[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes by LyricZhao · Pull Request #304 · deepseek-ai/DeepGEMM

LyricZhao · 2026-04-16T07:57:12Z

New features

Mega MoE, fusing & overlapping dispatch/linear 1/SwiGLU/linear 2/combine into a single mega-kernel, overlapping NVLink communication and tensor core computation
- Performance number: Add various optimizations and Mega MoE benchmarks #316
- Only FP8 x FP4 MoE is supported
- Requires PyTorch >= 2.9
FP4 Indexer (MQA logits) with larger MTP support
FP8 x FP4 GEMM
PDL
Refactors on GEMM heuristics
Faster JIT compilation
GEMM optimizations (dynamic swap A/B, much faster MoE GEMM)
DeepEPv2 MoE GEMM layout

Bug fixes

JIT may crash on distributed FS
Some kernel hangs and IMA

Contributors

Mega MoE: @LyricZhao @zheanxu @bucket-xv @RayWang96 @interestingLSY @kurisu6912 @xay5421 @yukuai26
FP4 Indexer: @zheanxu @xay5421 @interestingLSY @kurisu6912
GEMM, PDL, JIT and bug fixes: @zheanxu @bucket-xv @xay5421 @yukuai26 @LyricZhao

Additional notes

Mega MoE is still under development and optimizations, stay tuned and optimization ideas are welcome!
Disclaimer: this release is only related to DeepGEMM's development, has nothing to do with internal model release.

yuho8818 · 2026-04-20T09:27:48Z

LGTM

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus the two NV-side feature commits: - Enable bf16 bias output for FP8 GEMM on SM100 - SM90 paged MQA logits: support kv_block=32 and next_n=4 The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh -- deepseek-ai#292 (d30fc36) already covers the same hazard via NamedBarrier.sync.

Extend the SM90 FP8 paged MQA logits kernel and host wrapper to cover block_kv=32 and next_n=4 on post-deepseek-ai#304 code. Kernel (deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh): - Keep MMA tile at kComputeBlockKV=64 and let block_kv vary (32 or 64). Each WGMMA iteration covers kNumBlocksPerMMA = 64 / block_kv physical blocks with one TMA copy each. - Load block_table entries with a loop of scalar __ldg's so block_table_stride no longer has to be a multiple of kNumBlocksPerMMA. - Add kNumKVMulticast template parameter for next_n=4; compute kNextNPerCTA = kNextN / kNumKVMulticast and split Q rows across two CTAs in a cluster via cta_rank_in_cluster, keeping smem and register budgets within SM90 limits (WGMMA N=256 otherwise exceeds kNumMathRegisters=104). - Use cute::cluster_id_in_grid().x for the scheduler index. Host wrapper (csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp): - Select num_kv_multicast = 2 when SM90+next_n=4 else 1. Build Q/weights TMA descriptors with next_n_per_cta * num_heads box, size smem from next_n_per_cta, launch with cluster_dim=num_kv_multicast. - Emit the extra kNumKVMulticast template arg only for SM90; SM100 instantiation is unchanged. Metadata (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh, csrc/apis/attention.hpp): - Pass num_next_n_atoms in as a runtime parameter instead of hard-coding it from next_n. SM90 multicast passes 1 (no atomization, one entry per cluster); SM100 keeps 2 for next_n=4. - Relax the schedule_meta size assertion to num_sms / num_kv_multicast + 1 to match the cluster grid. Other: - Add utils::shfl_sync helper in common/utils.cuh. - tests/test_attention.py: extend enumerate_paged_mqa_logits to cover block_kv=32 and next_n=4 on SM90; call get_paged_mqa_logits_metadata with num_clusters = num_sms / num_kv_multicast; initialize block_table with torch.zeros so unused slots hold a valid block id.

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus the two NV-side feature commits: - Enable bf16 bias output for FP8 GEMM on SM100 - SM90 paged MQA logits: support kv_block=32 and next_n=4 The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh -- deepseek-ai#292 (d30fc36) already covers the same hazard via NamedBarrier.sync.

Extend the SM90 FP8 paged MQA logits kernel and host wrapper to cover block_kv=32 and next_n=4 on post-deepseek-ai#304 code. Kernel (deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh): - Keep MMA tile at kComputeBlockKV=64 and let block_kv vary (32 or 64). Each WGMMA iteration covers kNumBlocksPerMMA = 64 / block_kv physical blocks with one TMA copy each. - Load block_table entries with a loop of scalar __ldg's so block_table_stride no longer has to be a multiple of kNumBlocksPerMMA. - Add kNumKVMulticast template parameter for next_n=4; compute kNextNPerCTA = kNextN / kNumKVMulticast and split Q rows across two CTAs in a cluster via cta_rank_in_cluster, keeping smem and register budgets within SM90 limits (WGMMA N=256 otherwise exceeds kNumMathRegisters=104). - Use cute::cluster_id_in_grid().x for the scheduler index. Host wrapper (csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp): - Select num_kv_multicast = 2 when SM90+next_n=4 else 1. Build Q/weights TMA descriptors with next_n_per_cta * num_heads box, size smem from next_n_per_cta, launch with cluster_dim=num_kv_multicast. - Emit the extra kNumKVMulticast template arg only for SM90; SM100 instantiation is unchanged. Metadata (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh, csrc/apis/attention.hpp): - Pass num_next_n_atoms in as a runtime parameter instead of hard-coding it from next_n. SM90 multicast passes 1 (no atomization, one entry per cluster); SM100 keeps 2 for next_n=4. - Relax the schedule_meta size assertion to num_sms / num_kv_multicast + 1 to match the cluster grid. Other: - Add utils::shfl_sync helper in common/utils.cuh. - tests/test_attention.py: extend enumerate_paged_mqa_logits to cover block_kv=32 and next_n=4 on SM90; call get_paged_mqa_logits_metadata with num_clusters = num_sms / num_kv_multicast; initialize block_table with torch.zeros so unused slots hold a valid block id.

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus the two NV-side feature commits: - Enable bf16 bias output for FP8 GEMM on SM100 - SM90 paged MQA logits: support kv_block=32 and next_n=4 The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh that was added by nv_dev's 8e866e4; deepseek-ai#304 removed the equivalent sync upstream (replaced with NamedBarrier.sync in sibling SM100 kernels) so keeping the nv_dev version would diverge from upstream's final sync design.

Extend the SM90 FP8 paged MQA logits kernel and host wrapper to cover block_kv=32 and next_n=4 on post-deepseek-ai#304 code. Kernel (deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh): - Keep MMA tile at kComputeBlockKV=64 and let block_kv vary (32 or 64). Each WGMMA iteration covers kNumBlocksPerMMA = 64 / block_kv physical blocks with one TMA copy each. - Load block_table entries with a loop of scalar __ldg's so block_table_stride no longer has to be a multiple of kNumBlocksPerMMA. - Add kNumKVMulticast template parameter for next_n=4; compute kNextNPerCTA = kNextN / kNumKVMulticast and split Q rows across two CTAs in a cluster via cta_rank_in_cluster, keeping smem and register budgets within SM90 limits (WGMMA N=256 otherwise exceeds kNumMathRegisters=104). - Use cute::cluster_id_in_grid().x for the scheduler index. Host wrapper (csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp): - Select num_kv_multicast = 2 when SM90+next_n=4 else 1. Build Q/weights TMA descriptors with next_n_per_cta * num_heads box, size smem from next_n_per_cta, launch with cluster_dim=num_kv_multicast. - Emit the extra kNumKVMulticast template arg only for SM90; SM100 instantiation is unchanged. Other: - Add utils::shfl_sync helper in common/utils.cuh. - tests/test_attention.py: extend enumerate_paged_mqa_logits to cover block_kv=32 and next_n=4 on SM90; call get_paged_mqa_logits_metadata with num_clusters = num_sms / num_kv_multicast; initialize block_table with torch.zeros so unused slots hold a valid block id.

…API mismatch The post-deepseek-ai#304 metadata kernel hard-codes num_next_n_atoms from next_n, which assumes the SM100 time-atomization scheme. On SM90 next_n=4 the cluster-multicast design processes one full q per cluster with no atomization, so the metadata ends up producing 2x the tasks the kernel expects and stepping past context_lens when initializing the scheduler (observed as illegal memory access in sm90_fp8_paged_mqa_logits). Fixes and hardening: - Pipe num_next_n_atoms as a runtime parameter into the metadata kernel (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh) and let the host decide. attention.hpp computes it based on arch + next_n + num_kv_multicast: SM90 multicast passes 1 (one entry per cluster), SM100 keeps 2 for next_n=4. - Relax the schedule_meta size assertion in attention.hpp to 'num_sms / num_kv_multicast + 1' so the check matches the cluster grid rather than the SM grid. - Extract 'get_paged_mqa_logits_num_kv_multicast' helper in the host wrapper so the formula lives in one place. - Add 'DG_HOST_ASSERT(num_sms % num_kv_multicast == 0)' guard in the cluster launch path for non-H200/H100-like num_sms. - Document the block_table unused-slot contract on the public API: callers MUST initialize unused block_table slots to a valid block id (e.g. 0) because the kernel may read slots up to num_kv * kNumBlocksPerMMA when context_len is not a multiple of kComputeBlockKV=64.

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus three NV-side commits: - Enable bf16 bias output for FP8 GEMM on SM100 (xueweil) - SM90 paged MQA logits: support kv_block=32 and next_n=4 (Ray) - Fix SM90 paged MQA logits for post-deepseek-ai#304 metadata/schedule API mismatch The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh that was added by nv_dev's 8e866e4; deepseek-ai#304 removed the equivalent sync upstream (replaced with NamedBarrier.sync in sibling SM100 kernels) so keeping the nv_dev version would diverge from upstream's final sync design.

Extend the SM90 FP8 paged MQA logits kernel and host wrapper to cover block_kv=32 and next_n=4 on post-deepseek-ai#304 code. Kernel (deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh): - Keep MMA tile at kComputeBlockKV=64 and let block_kv vary (32 or 64). Each WGMMA iteration covers kNumBlocksPerMMA = 64 / block_kv physical blocks with one TMA copy each. - Load block_table entries with a loop of scalar __ldg's so block_table_stride no longer has to be a multiple of kNumBlocksPerMMA. - Add kNumKVMulticast template parameter for next_n=4; compute kNextNPerCTA = kNextN / kNumKVMulticast and split Q rows across two CTAs in a cluster via cta_rank_in_cluster, keeping smem and register budgets within SM90 limits (WGMMA N=256 otherwise exceeds kNumMathRegisters=104). - Use cute::cluster_id_in_grid().x for the scheduler index. Host wrapper (csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp): - Select num_kv_multicast = 2 when SM90+next_n=4 else 1. Build Q/weights TMA descriptors with next_n_per_cta * num_heads box, size smem from next_n_per_cta, launch with cluster_dim=num_kv_multicast. - Emit the extra kNumKVMulticast template arg only for SM90; SM100 instantiation is unchanged. Other: - Add utils::shfl_sync helper in common/utils.cuh. - tests/test_attention.py: extend enumerate_paged_mqa_logits to cover block_kv=32 and next_n=4 on SM90; call get_paged_mqa_logits_metadata with num_clusters = num_sms / num_kv_multicast; initialize block_table with torch.zeros so unused slots hold a valid block id.

…API mismatch The post-deepseek-ai#304 metadata kernel hard-codes num_next_n_atoms from next_n, which assumes the SM100 time-atomization scheme. On SM90 next_n=4 the cluster-multicast design processes one full q per cluster with no atomization, so the metadata ends up producing 2x the tasks the kernel expects and stepping past context_lens when initializing the scheduler (observed as illegal memory access in sm90_fp8_paged_mqa_logits). Fixes and hardening: - Pipe num_next_n_atoms as a runtime parameter into the metadata kernel (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh) and let the host decide. attention.hpp computes it based on arch + next_n + num_kv_multicast: SM90 multicast passes 1 (one entry per cluster), SM100 keeps 2 for next_n=4. - Relax the schedule_meta size assertion in attention.hpp to 'num_sms / num_kv_multicast + 1' so the check matches the cluster grid rather than the SM grid. - Extract 'get_paged_mqa_logits_num_kv_multicast' helper in the host wrapper so the formula lives in one place. - Add 'DG_HOST_ASSERT(num_sms % num_kv_multicast == 0)' guard in the cluster launch path for non-H200/H100-like num_sms. - Document the block_table unused-slot contract on the public API: callers MUST initialize unused block_table slots to a valid block id (e.g. 0) because the kernel may read slots up to num_kv * kNumBlocksPerMMA when context_len is not a multiple of kComputeBlockKV=64.

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus three NV-side commits: - Enable bf16 bias output for FP8 GEMM on SM100 (xueweil) - SM90 paged MQA logits: support kv_block=32 and next_n=4 (Ray) - Fix SM90 paged MQA logits for post-deepseek-ai#304 metadata/schedule API mismatch The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh that was added by nv_dev's 8e866e4; deepseek-ai#304 removed the equivalent sync upstream (replaced with NamedBarrier.sync in sibling SM100 kernels) so keeping the nv_dev version would diverge from upstream's final sync design.

huangzhilin-hzl · 2026-04-21T09:27:43Z

Hi @LyricZhao tests/test_mega_moe.py currently depends on deep_ep.ElasticBuffer, but I could not find ElasticBuffer in the public deepseek-ai/DeepEP branches; was this test written against a private/internal DeepEP version?

LyricZhao · 2026-04-21T09:31:49Z

@huangzhilin-hzl deep_ep.ElasticBuffer (DeepEP V2) will be released in this week. And, you can still skip the legacy path test and just look at the perf numbers.

CyberQin · 2026-04-22T01:24:21Z

deep_ep基础设施还没更新完啊，那model release不是更佳悬了，yifan的爆料要失真了

…API mismatch The post-deepseek-ai#304 metadata kernel hard-codes num_next_n_atoms from next_n, which assumes the SM100 time-atomization scheme. On SM90 next_n=4 the cluster-multicast design processes one full q per cluster with no atomization, so the metadata ends up producing 2x the tasks the kernel expects and stepping past context_lens when initializing the scheduler (observed as illegal memory access in sm90_fp8_paged_mqa_logits). Fixes and hardening: - Pipe num_next_n_atoms as a runtime parameter into the metadata kernel (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh) and let the host decide. attention.hpp computes it based on arch + next_n + num_kv_multicast: SM90 multicast passes 1 (one entry per cluster), SM100 keeps 2 for next_n=4. - Relax the schedule_meta size assertion in attention.hpp to 'num_sms / num_kv_multicast + 1' so the check matches the cluster grid rather than the SM grid. - Extract 'get_paged_mqa_logits_num_kv_multicast' helper in the host wrapper so the formula lives in one place. - Add 'DG_HOST_ASSERT(num_sms % num_kv_multicast == 0)' guard. - Normalize the public 'get_paged_mqa_logits_metadata' helper so callers always pass the natural 'get_num_sms()' and physical 'block_kv': the helper now computes num_clusters internally, sizes the schedule as '{num_clusters + 1, 2}', and uses the compute-block size on SM90 regardless of the physical block_kv the caller supplies. - Tighten the SM90 block_table load guard to a per-physical-block 'phys_idx < ceil_div(context_len, BLOCK_KV)' check so BLOCK_KV=32 with odd physical block counts cannot read past the natural row end.

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus three NV-side commits: - Enable bf16 bias output for FP8 GEMM on SM100 (xueweil) - SM90 paged MQA logits: support kv_block=32 and next_n=4 (Ray) - Fix SM90 paged MQA logits for post-deepseek-ai#304 metadata/schedule API mismatch The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh that was added by nv_dev's 8e866e4; deepseek-ai#304 removed the equivalent sync upstream (replaced with NamedBarrier.sync in sibling SM100 kernels) so keeping the nv_dev version would diverge from upstream's final sync design.

Extend the SM90 FP8 paged MQA logits kernel and host wrapper to cover block_kv=32 and next_n=4 on post-deepseek-ai#304 code. Kernel (deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh): - Keep MMA tile at kComputeBlockKV=64 and let block_kv vary (32 or 64). Each WGMMA iteration covers kNumBlocksPerMMA = 64 / block_kv physical blocks with one TMA copy each. - Add kNumKVMulticast template parameter for next_n=4; compute kNextNPerCTA = kNextN / kNumKVMulticast and split Q rows across two CTAs in a cluster via cta_rank_in_cluster, keeping smem and register budgets within SM90 limits (WGMMA N=256 otherwise exceeds kNumMathRegisters=104). - Guard the block_table load per physical block so BLOCK_KV=32 with an odd physical block count cannot read past the natural row end. - Use cute::cluster_id_in_grid().x for the scheduler index. Post-deepseek-ai#304 metadata/schedule integration (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh, csrc/apis/attention.hpp): - Pipe num_next_n_atoms as a runtime parameter so SM90 cluster multicast (no atomization) passes 1 while SM100 time-atomization keeps 2 for next_n=4, matching the kernel's scheduling. - Relax the schedule_meta size assertion to num_sms / num_kv_multicast + 1 to match the cluster grid. - Normalize the public get_paged_mqa_logits_metadata helper so callers pass the natural get_num_sms() and physical block_kv; the helper now computes num_clusters internally, sizes the schedule as {num_clusters + 1, 2}, and uses the compute-block size on SM90 regardless of the physical block_kv the caller supplies. - Extract get_paged_mqa_logits_num_kv_multicast helper and add a num_sms % num_kv_multicast == 0 host-side guard. Host wrapper (csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp): - Select num_kv_multicast = 2 when SM90+next_n=4 else 1. Build Q/weights TMA descriptors with next_n_per_cta * num_heads box, size smem from next_n_per_cta, launch with cluster_dim=num_kv_multicast. - Emit the extra kNumKVMulticast template arg only for SM90; SM100 instantiation is unchanged. Other: - Add utils::shfl_sync helper in common/utils.cuh. - tests/test_attention.py: extend enumerate_paged_mqa_logits to cover block_kv=32 and next_n=4 on SM90.

The post-deepseek-ai#304 metadata kernel hard-codes num_next_n_atoms from next_n, which assumes the SM100 time-atomization scheme. On SM90 next_n=4 the cluster-multicast design processes one full q per cluster with no atomization, so the hard-coded value over-schedules by 2x; when the kernel scheduler (kNumNextNAtoms=1 in the SM90 template) decodes those entries it steps past context_lens and triggers an illegal memory access in sm90_fp8_paged_mqa_logits. Fixes: - Pipe num_next_n_atoms as a runtime parameter into the metadata kernel (deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh) and let the host decide. attention.hpp computes it based on arch + next_n + num_kv_multicast so SM90 multicast passes 1 while SM100 keeps 2 for next_n=4, matching the respective kernel template. - Relax the schedule_meta size assertion in attention.hpp to 'num_sms / num_kv_multicast + 1' so the check matches the cluster grid rather than the SM grid for SM90 next_n=4. - Relax the block_kv assertion in get_paged_mqa_logits_metadata to accept block_kv in {32, 64} and internally use the compute block size (64) on SM90, so callers can pass the physical block_kv the matching fp8_fp4_paged_mqa_logits call uses.

Brings in upstream commits d30fc36 (TMEM alloc/dealloc sync fix, deepseek-ai#292) and 7f2a703 (public release 26/04: Mega MoE, FP4 indexer, and related refactors, deepseek-ai#304), plus three NV-side commits: - Enable bf16 bias output for FP8 GEMM on SM100 (xueweil) - SM90 paged MQA logits: support kv_block=32 and next_n=4 (Ray) - Fix SM90 paged MQA logits for post-deepseek-ai#304 metadata/schedule API mismatch The upstream deepseek-ai#304 refactor reshapes the attention/MQA layer; conflicts with the pre-merge nv_dev code are resolved by taking the post-deepseek-ai#304 version. Also drops the redundant __syncthreads() in sm100_bmk_bnk_mn.cuh that was added by nv_dev's 8e866e4; deepseek-ai#304 removed the equivalent sync upstream (replaced with NamedBarrier.sync in sibling SM100 kernels) so keeping the nv_dev version would diverge from upstream's final sync design.

qinqinwo · 2026-04-28T04:48:00Z

请问有支持其他卡型的计划吗？比如H100。

yy16432 · 2026-05-01T10:49:03Z

Any plan to add support to RTX PRO 6000(SM 120)? Thanks!

Extend the SM90 FP8 paged MQA logits kernel and host wrapper to cover block_kv=32 and next_n=4 on post-deepseek-ai#304 code. Kernel (deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh): - Keep MMA tile at kComputeBlockKV=64 and let block_kv vary (32 or 64). Each WGMMA iteration covers kNumBlocksPerMMA = 64 / block_kv physical blocks with one TMA copy each. - Load block_table entries with a loop of scalar __ldg's so block_table_stride no longer has to be a multiple of kNumBlocksPerMMA. - Add kNumKVMulticast template parameter for next_n=4; compute kNextNPerCTA = kNextN / kNumKVMulticast and split Q rows across two CTAs in a cluster via cta_rank_in_cluster, keeping smem and register budgets within SM90 limits (WGMMA N=256 otherwise exceeds kNumMathRegisters=104). - Use cute::cluster_id_in_grid().x for the scheduler index. Host wrapper (csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp): - Select num_kv_multicast = 2 when SM90+next_n=4 else 1. Build Q/weights TMA descriptors with next_n_per_cta * num_heads box, size smem from next_n_per_cta, launch with cluster_dim=num_kv_multicast. - Emit the extra kNumKVMulticast template arg only for SM90; SM100 instantiation is unchanged. Other: - Add utils::shfl_sync helper in common/utils.cuh. - tests/test_attention.py: extend enumerate_paged_mqa_logits to cover block_kv=32 and next_n=4 on SM90; call get_paged_mqa_logits_metadata with num_clusters = num_sms / num_kv_multicast; initialize block_table with torch.zeros so unused slots hold a valid block id. (cherry picked from commit 12cb7c0) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

LyricZhao changed the title ~~[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other bug fixes~~ [Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes Apr 16, 2026

Merge with private repo

7dbe077

LyricZhao force-pushed the public-release-260416 branch from a050d09 to 7dbe077 Compare April 16, 2026 08:23

LyricZhao added 3 commits April 16, 2026 16:38

Update README

37ab81a

Update README

d01623b

Update README

0c72e11

LyricZhao requested a review from zheanxu April 16, 2026 08:48

zheanxu approved these changes Apr 16, 2026

View reviewed changes

LyricZhao added 3 commits April 16, 2026 17:37

Add PyTorch requirements

2d6fd1c

Fix sync scopes for MQA logits (#256)

9f4bed6

Update README

211d267

LyricZhao merged commit 7f2a703 into main Apr 17, 2026

YconquestY mentioned this pull request Apr 19, 2026

[feat] Enable DeepGEMM Mega MoE as MoE layer runner backend in NVLink domain sgl-project/sglang#23167

Closed

5 tasks

akoumpa mentioned this pull request Apr 20, 2026

Integrate Mega MoE NVIDIA-NeMo/Automodel#1918

Open

Barry-Delaney mentioned this pull request Apr 21, 2026

[Test] Merge & update Barry-Delaney/DeepGEMM#1

Closed

Barry-Delaney mentioned this pull request Apr 22, 2026

fix: SM90 paged MQA logits IMA on next_n=4 (nv_dev rebase) #314

Merged

b8zhong mentioned this pull request Apr 27, 2026

Use Torch torch.mm for Deepseek V3.2 Indexer GEMM sgl-project/sglang#23856

Merged

qinqinwo mentioned this pull request Apr 28, 2026

need MegaMOE for hopper style(H100) flashinfer-ai/flashinfer#3196

Closed

This was referenced May 28, 2026

fix(paged_mqa_logits): guard get_num_kv against out-of-batch q indices tirthasheshpatel/DeepGEMM#1

Closed

Fix an OOB access in the the paged MQA logits kernel #344

Closed

aganhui mentioned this pull request Sep 11, 2026

Does Mega MoE support RTX 50-series GPUs, and can it work in a PCIe-only setup without NVLink? #356

Open

Navigation Menu

Repository navigation

[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes - #304

[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes#304
LyricZhao merged 7 commits into
mainfrom
public-release-260416

LyricZhao commented Apr 16, 2026 •

edited

Loading

Uh oh!

yuho8818 commented Apr 20, 2026

Uh oh!

huangzhilin-hzl commented Apr 21, 2026

Uh oh!

LyricZhao commented Apr 21, 2026

Uh oh!

CyberQin commented Apr 22, 2026

Uh oh!

qinqinwo commented Apr 28, 2026

Uh oh!

yy16432 commented May 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

LyricZhao commented Apr 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

New features

Bug fixes

Contributors

Additional notes

Uh oh!

yuho8818 commented Apr 20, 2026

Uh oh!

huangzhilin-hzl commented Apr 21, 2026

Uh oh!

LyricZhao commented Apr 21, 2026

Uh oh!

CyberQin commented Apr 22, 2026

Uh oh!

qinqinwo commented Apr 28, 2026

Uh oh!

yy16432 commented May 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

LyricZhao commented Apr 16, 2026 •

edited

Loading