Multiple updates and refactorings - #347
Merged
Merged
Conversation
RayWang96
commented
Jun 1, 2026
RayWang96
commented
Collaborator
- Optimize FP8/FP4 Mega MoE dispatch, scheduling, and shared memory layout
- Add BF16 accumulation/output support updates for GEMM paths
- Improve paged MQA scheduler and attention coverage
- Minor fixes and test updates
* Optimize FP8/FP4 Mega MoE dispatch, scheduling, and shared memory layout * Add BF16 accumulation/output support updates for GEMM paths * Improve paged MQA scheduler and attention coverage * Minor fixes and test updates
LyricZhao
approved these changes
Jun 1, 2026
This was referenced Jun 1, 2026
Closed
RayWang96
pushed a commit
that referenced
this pull request
Jun 2, 2026
…s) into nv_dev Brings deepseek-ai PR #347 (88965b0, also pulls 714dd1a 'Update test_mega_moe.py') onto nv_dev (currently at ac1f285 = #316 merge #328 + PR #342 IMA guard). Merge base is 891d57b. Conflicts resolved (5 files): * scheduler/paged_mqa_logits.cuh: take #347's refresh_num_kv_and_advance + reversed metadata allocation; drop nv_dev's PR #342 exist_q_atom_idx guard and get_atom_advance (both subsumed — reversed alloc keeps current_q_atom_idx in-bounds, refresh_num_kv_and_advance reproduces the varlen 1-or-2 advance). NV's host-passed num_next_n_atoms survives in the auto-merged metadata kernel (verified: single param, no inline recompute, coherent with reversed alloc). * csrc/apis/gemm.hpp: drop redundant C/D dtype asserts (#347 cleanup; the pre-existing d-dtype check + d==c check already cover them). * sm100_fp8_fp4_gemm_1d1d.cuh: take #347's unconditional BF16/FP32 C/D assert (supersedes nv_dev's accumulation-only assert). * tests/generators.py: take #347's broadened BF16-accumulation enumeration (out_dtype==bf16 and (dtype==bf16 or arch==10)) — strict superset of nv_dev's fp8+SM100 case. * tests/test_attention.py: merge both — keep NV's SM90 next_n=4 multicast + num_clusters metadata sizing; add #347's pool-size limits, batch_size=4096, and block_table/context_lens sanity asserts. Auto-merged metadata kernel + smxx_fp8_fp4_paged_mqa_logits.hpp + attention.hpp verified coherent across all three layers (kernel signature / launch / host call all carry num_next_n_atoms + is_varlen + indices). Validated on B-card (SM100) and H200 (SM90) incl next_n=4 multicast. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
RayWang96
added a commit
that referenced
this pull request
Jun 2, 2026
Sync nv_dev with upstream #347
leavelet
pushed a commit
to leavelet/DeepGEMM
that referenced
this pull request
Jun 5, 2026
Brings in 9 upstream nv_dev commits: paged-MQA metadata OOB fix (deepseek-ai#353), IMA guard in the paged-MQA scheduler (deepseek-ai#342), the FP16-weights FP8 MQA logits kernel, the mega-MoE refactor (deepseek-ai#347), and assorted updates. Conflicts resolved (3 files): - csrc/apis/attention.hpp (dense MQA dispatch): union of both sides — keep SM120 FP4 (sm120_fp4_mqa_logits) and SM120 FP8 (arch 12) arms, and add nv_dev's FP16-weights arm (sm100_fp8_mqa_logits_f16_weights, SM100-only, already guarded by the arch==10 host assert). - deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh: take nv_dev's refactored scheduler. SM120's one-line OOB guard is subsumed by nv_dev's reversed-metadata allocation (the deepseek-ai#342/deepseek-ai#353 fix). The refactor replaces get_num_kv/get_atom_advance with refresh_num_kv_and_advance + get_last_advance. - tests/test_attention.py: keep SM120 small-S split-KV MQA shapes and add the upstream large-KV/32-head regime; next_n = (1..6) on SM120, (1,2,4,5,6) on SM100, (1,2,4) on SM90; take nv_dev's paged test body (memory caps + peak-mem reset). Semantic fixes required by the merge (no textual conflict, would not compile/run otherwise): - sm120_fp8/fp4_paged_mqa_logits.cuh: migrate the scheduler call get_atom_advance(next_q_idx, batch_size) -> get_last_advance() to match nv_dev's refactored scheduler API (same migration nv_dev applied to the SM100 paged kernel). - tests/test_attention.py: gate the FP16-weights path to arch 10. Upstream set FP16 weights unconditionally, which dispatches to the SM100-only kernel and trips the arch==10 host assert when test_attention runs on SM90/SM120. Validated on H200 (sm90): clean build; test_bf16, test_fp8_fp4, test_attention all pass. SM120 (arch 12) kernels are JIT and arch-gated, so they are neither compiled nor run on H200 — they require an RTX 6000 run before the final PR.
zyongye
pushed a commit
to zyongye/DeepGEMM
that referenced
this pull request
Jun 9, 2026
…wave scheduling) The squash-rewrite at 4ec3457 reverted three deepseek-ai/DeepGEMM deepseek-ai#347 mega-MoE optimizations, regressing token-64 (block_m=16) from ~304us to ~321us at the DeepSeek-V4-Pro rack shape (H=7168, E=384, K=6) on 8x B300: 1. block_k=256 for <=8.5 tokens/expert ("larger BLOCK_K for less synchronization") - halves the K-loop pipeline barrier syncs, and makes the block_m=16 mxf4 A-tile naturally 1024-aligned. 2. get_num_experts_per_wave: L2-reuse early-return + tail-ratio wave search. 3. scheduler get_wave_expert_end_idx partial-wave clamp. Restore all three by reverting the 6 mega-only files to their 89b5a63 (pre-rewrite, deepseek-ai#347-complete) versions, which carry the BLOCK_K-generic UMMA_BLOCK_K two-level K-loop. The only shared-header drift -- 4ec3457 dropped the optional smem_outer_dim param from make_tma_sf_desc -- is bridged mega-locally via make_tma_sf_desc_mega (calls the unmodified shared make_tma_2d_desc; block_k=256 needs smem_outer_dim=block_k/128=2). No other kernels touched. Validated numerics at block_m=16 (block_k=256 path), H in {2048,7168}: mxfp4 ~0.043, mxfp4+fp8-combine ~0.044 vs bf16 reference. token-64: 321 -> 302us (matches 89b5a63); larger tiers (block_k=128) unchanged. Also adds --act-format {fp8,mxfp4} to tests/test_mega_moe.py so the rack harness can benchmark MXFP4 activations. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
b8zhong
pushed a commit
to bzhng-development/DeepGEMM
that referenced
this pull request
Jun 22, 2026
Co-authored-by: Ray Wang <raywang@nvidia.com>
This was referenced Jun 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.