Skip to content

Performance: BlockTile 256x128 optimizations enable 1500+ TF FP8 - #81

Merged
LyricZhao merged 9 commits into
mainfrom
blocktile-256x128
Apr 9, 2025
Merged

Performance: BlockTile 256x128 optimizations enable 1500+ TF FP8#81
LyricZhao merged 9 commits into
mainfrom
blocktile-256x128

Conversation

@sazczmh

@sazczmh sazczmh commented Apr 8, 2025

Copy link
Copy Markdown
Collaborator

By resuing the Accumulator registers of Tensor Cores to implement a 256x128 BlockTile structure, this approach significantly increases data reuse, reduces the demand for L2 Cache and HBM memory accesses, and enhances the SM's computational frequency, ultimately achieving FP8 performance exceeding 1,500+ TFLOPS.

M N K Base BMxBN Computation Opti BMxBN Computation Speedup
4096 24576 1536 128x160 1162 TF 256x128 1204 TF 3.61%
4096 32768 512 128x160 801 TF 256x128 777 TF -3.00%
4096 7168 16384 128x160 1451 TF 256x128 1500 TF 3.38%
4096 4096 7168 128x160 1304 TF 256x128 1377 TF 5.60%
4096 7168 2048 128x160 1185 TF 256x128 1159 TF -2.19%

Test on “H800”-SXM && CUDA 12.8.1

@sazczmh sazczmh added the perf label Apr 8, 2025
@sazczmh sazczmh self-assigned this Apr 8, 2025
@LyricZhao
LyricZhao force-pushed the blocktile-256x128 branch from 1eeb98a to 48a5f07 Compare April 9, 2025 02:01
@LyricZhao
LyricZhao requested a review from zheanxu April 9, 2025 03:10
@LyricZhao
LyricZhao merged commit fed3e4d into main Apr 9, 2025
@LyricZhao
LyricZhao deleted the blocktile-256x128 branch April 11, 2025 03:35
@YouJiacheng YouJiacheng mentioned this pull request Jan 12, 2026
Merged
LyricZhao added a commit that referenced this pull request Apr 16, 2026
@LyricZhao
Share pipeline stages between blocks (#81)
87e44ab 
* Shared stages between blocks

* Optimize `elect_one_sync`

* Minor fix

* Enable more stages

* Better block N selections

* Optimize cluster sync

* Less unrolling

* Better performance

* Share TMA store pipeline between blocks

* Fix cluster sync

* Better performance

* Minor fix

* Fix unrolling

* Use sync-free UMMA impl

* Remove redundant `__syncwarp`

* Remove comments

* Lint
qiushixiaoyu pushed a commit to qiushixiaoyu/DeepGEMM that referenced this pull request Sep 4, 2026
@zcnrex @Fridge003
Merge pull request deepseek-ai#81 from zcnrex/megamoe-nvfp4-fix
b7fb616 
fix and add cache hint
@vipulSharma18 vipulSharma18 mentioned this pull request Oct 13, 2026
Merged
vipulSharma18 added a commit to fw-ai-external/DeepGEMM that referenced this pull request Oct 13, 2026
@vipulSharma18 @claude
Add sgl-project's NVFP4 MegaMoE as an additive deep_gemm.mega_nvfp4…
ef210f2 
… module

Port the SM100 NVFP4 x NVFP4 MegaMoE decode kernel and the fused
`mega_moe_pre_dispatch` kernel from sgl-project/DeepGEMM@ec226bc (#27, 593e376,
deepseek-ai#80, deepseek-ai#81) so Fireworks can drop the separate `sgl_deep_gemm` wheel.

Purely additive: the only change to an existing file is the registration hook in
csrc/python_api.cpp. No existing header changes, so no existing kernel's JIT
cache key changes.

- Device code: sgl's headers copied into deep_gemm/include/deep_gemm/mega_nvfp4/
  (kernel, pre-dispatch, layout, scheduler, barrier, math/utils/tma_copy, ptx
  ld_st/tcgen05/tma/utils, sym_buffer, and the `MmaKind` block of types.cuh),
  moved into namespace `deep_gemm::mega_nvfp4` with includes redirected and
  nothing else changed (line numbers kept, so even assert-path SASS matches).
- Host code: csrc/apis/mega_nvfp4.hpp plus heuristics/nvfp4_mega_moe.hpp and
  impls/sm100_nvfp4_mega_moe{,_pre_dispatch}.hpp, in `deep_gemm::mega_nvfp4`
  with `nvfp4`-marked names. Like the separate sgl library, these kernels keep
  their own SM count, compile for sgl's `sm_100f` target and launch without
  PDL, so `deep_gemm.set_num_sms`/`set_pdl` and these kernels never affect each
  other. Not ported: TRT-LLM weight layout, FP8 combine, BF16 and SM90 paths.
- Python: deep_gemm/mega_nvfp4 exposes the sgl_deep_gemm subset Fireworks
  uses (get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe,
  mega_moe_pre_dispatch, nvfp4_mega_moe, set_num_sms, get_block_m_for_mega_moe,
  transform_ue4m3_sf_into_required_layout).
- tests/test_mega_moe_nvfp4.py ports sgl's pre-dispatch and NVFP4 alphas tests
  and adds a bytewise check of the NVFP4 pre-dispatch quantization.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
vipulSharma18 added a commit to fw-ai-external/DeepGEMM that referenced this pull request Oct 14, 2026
@vipulSharma18 @claude
Add sgl-project's NVFP4 MegaMoE as an additive deep_gemm.mega_nvfp4…
9b56b6f 
… module

Port the SM100 NVFP4 x NVFP4 MegaMoE decode kernel and the fused
`mega_moe_pre_dispatch` kernel from sgl-project/DeepGEMM@ec226bc (#27, 593e376,
deepseek-ai#80, deepseek-ai#81) so Fireworks can drop the separate `sgl_deep_gemm` wheel.

Purely additive: the only change to an existing file is the registration hook in
csrc/python_api.cpp. No existing header changes, so no existing kernel's JIT
cache key changes.

- Device code: sgl's headers copied into deep_gemm/include/deep_gemm/mega_nvfp4/
  (kernel, pre-dispatch, layout, scheduler, barrier, math/utils/tma_copy, ptx
  ld_st/tcgen05/tma/utils, sym_buffer, and the `MmaKind` block of types.cuh),
  moved into namespace `deep_gemm::mega_nvfp4` with includes redirected and
  nothing else changed (line numbers kept, so even assert-path SASS matches).
- Host code: csrc/apis/mega_nvfp4.hpp plus heuristics/nvfp4_mega_moe.hpp and
  impls/sm100_nvfp4_mega_moe{,_pre_dispatch}.hpp, in `deep_gemm::mega_nvfp4`
  with `nvfp4`-marked names. Like the separate sgl library, these kernels keep
  their own SM count, compile for sgl's `sm_100f` target and launch without
  PDL, so `deep_gemm.set_num_sms`/`set_pdl` and these kernels never affect each
  other. Not ported: TRT-LLM weight layout, FP8 combine, BF16 and SM90 paths.
- Python: deep_gemm/mega_nvfp4 exposes the sgl_deep_gemm subset Fireworks
  uses (get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe,
  mega_moe_pre_dispatch, nvfp4_mega_moe, set_num_sms, get_block_m_for_mega_moe,
  transform_ue4m3_sf_into_required_layout).
- tests/test_mega_moe_nvfp4.py ports sgl's pre-dispatch and NVFP4 alphas tests
  and adds a bytewise check of the NVFP4 pre-dispatch quantization.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
vipulSharma18 added a commit to fw-ai-external/DeepGEMM that referenced this pull request Oct 14, 2026
@vipulSharma18 @claude
Add sgl-project's NVFP4 MegaMoE as an additive deep_gemm.mega_nvfp4…
32dd03c 
… module

Port the SM100 NVFP4 x NVFP4 MegaMoE decode kernel and the fused
`mega_moe_pre_dispatch` kernel from sgl-project/DeepGEMM@ec226bc (#27, 593e376,
deepseek-ai#80, deepseek-ai#81) so Fireworks can drop the separate `sgl_deep_gemm` wheel.

Purely additive: the only change to an existing file is the registration hook in
csrc/python_api.cpp. No existing header changes, so no existing kernel's JIT
cache key changes.

- Device code: sgl's headers copied into deep_gemm/include/deep_gemm/mega_nvfp4/
  (kernel, pre-dispatch, layout, scheduler, barrier, math/utils/tma_copy, ptx
  ld_st/tcgen05/tma/utils, sym_buffer, and the `MmaKind` block of types.cuh),
  moved into namespace `deep_gemm::mega_nvfp4` with includes redirected and
  nothing else changed (line numbers kept, so even assert-path SASS matches).
- Host code: csrc/apis/mega_nvfp4.hpp plus heuristics/nvfp4_mega_moe.hpp and
  impls/sm100_nvfp4_mega_moe{,_pre_dispatch}.hpp, in `deep_gemm::mega_nvfp4`
  with `nvfp4`-marked names. Like the separate sgl library, these kernels keep
  their own SM count, compile for sgl's `sm_100f` target and launch without
  PDL, so `deep_gemm.set_num_sms`/`set_pdl` and these kernels never affect each
  other. Not ported: TRT-LLM weight layout, FP8 combine, BF16 and SM90 paths.
- Python: deep_gemm/mega_nvfp4 exposes the sgl_deep_gemm subset Fireworks
  uses (get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe,
  mega_moe_pre_dispatch, nvfp4_mega_moe, set_num_sms, get_block_m_for_mega_moe,
  transform_ue4m3_sf_into_required_layout).
- tests/test_mega_moe_nvfp4.py ports sgl's pre-dispatch and NVFP4 alphas tests
  and adds a bytewise check of the NVFP4 pre-dispatch quantization.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
vipulSharma18 added a commit to fw-ai-external/DeepGEMM that referenced this pull request Oct 14, 2026
@vipulSharma18 @claude
Add sgl-project's NVFP4 MegaMoE as an additive deep_gemm.mega_nvfp4 m…
0ed7cc8 
…odule (#18)

* Add sgl-project's NVFP4 MegaMoE as an additive `deep_gemm.mega_nvfp4` module

Port the SM100 NVFP4 x NVFP4 MegaMoE decode kernel and the fused
`mega_moe_pre_dispatch` kernel from sgl-project/DeepGEMM@ec226bc (#27, 593e376,
deepseek-ai#80, deepseek-ai#81) so Fireworks can drop the separate `sgl_deep_gemm` wheel.

Purely additive: the only change to an existing file is the registration hook in
csrc/python_api.cpp. No existing header changes, so no existing kernel's JIT
cache key changes.

- Device code: sgl's headers copied into deep_gemm/include/deep_gemm/mega_nvfp4/
  (kernel, pre-dispatch, layout, scheduler, barrier, math/utils/tma_copy, ptx
  ld_st/tcgen05/tma/utils, sym_buffer, and the `MmaKind` block of types.cuh),
  moved into namespace `deep_gemm::mega_nvfp4` with includes redirected and
  nothing else changed (line numbers kept, so even assert-path SASS matches).
- Host code: csrc/apis/mega_nvfp4.hpp plus heuristics/nvfp4_mega_moe.hpp and
  impls/sm100_nvfp4_mega_moe{,_pre_dispatch}.hpp, in `deep_gemm::mega_nvfp4`
  with `nvfp4`-marked names. Like the separate sgl library, these kernels keep
  their own SM count, compile for sgl's `sm_100f` target and launch without
  PDL, so `deep_gemm.set_num_sms`/`set_pdl` and these kernels never affect each
  other. Not ported: TRT-LLM weight layout, FP8 combine, BF16 and SM90 paths.
- Python: deep_gemm/mega_nvfp4 exposes the sgl_deep_gemm subset Fireworks
  uses (get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe,
  mega_moe_pre_dispatch, nvfp4_mega_moe, set_num_sms, get_block_m_for_mega_moe,
  transform_ue4m3_sf_into_required_layout).
- tests/test_mega_moe_nvfp4.py ports sgl's pre-dispatch and NVFP4 alphas tests
  and adds a bytewise check of the NVFP4 pre-dispatch quantization.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(mega_nvfp4): fence task-info metadata reads before slot release (upstream deepseek-ai#441)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zheanxu zheanxu zheanxu approved these changes
+1 more reviewer
@soundOfDestiny soundOfDestiny soundOfDestiny approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

@sazczmh sazczmh

Labels

perf

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@sazczmh @soundOfDestiny @zheanxu @LyricZhao