Repository navigation
Fix Mega MoE task info slot release ordering - #441
Conversation
|
|
||
| CUTLASS_DEVICE void release_task_info() const { | ||
| // Complete metadata reads before the scheduler can overwrite this slot. | ||
| ptx::fence_acq_rel_cta(); |
There was a problem hiding this comment.
🔵 suggestion: 非阻塞建议：非 leader CTA 的消费者与位于 CTA 0 的调度 warp 是跨 CTA 同步，严格按 PTX 内存模型应使用 .cluster scope（fence.acq_rel.cluster 或 mbarrier.arrive.release.cluster.shared::cluster）。实际硬件上 fence.acq_rel.cta 已足以让本线程未完成的 shared 加载在 arrive 前完成，因此行为安全且开销更低；建议在注释中补充说明这一取舍（为何选 .cta 而非 .cluster），或在 ld_st.cuh 中同时提供 fence_acq_rel_cluster() 变体以备后续需要。
🤖 v5
🤖 ds-review-bot Code Reviewv6v5该提交为 Mega MoE 调度器的 task info 槽位释放路径补上了内存序保证。原实现中，消费者线程在 v4The MR fixes a real producer/consumer race on the double-buffered Files reviewed: 2 |
…ring (#16) * Fix a race condition where in-flight tensormaps are updated in-place (deepseek-ai#343) * Public release 26/07 (deepseek-ai#377) * Public Release 26/09 (deepseek-ai#432) * Public release 26/09 * Update News * Fix Mega MoE task info slot release ordering (deepseek-ai#441) --------- Co-authored-by: Ivan Komarov <Ivan.Komarov@dfyz.info> Co-authored-by: Pengle Zhang <53843527+guyan364@users.noreply.github.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com>
…upstream deepseek-ai#441) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…upstream deepseek-ai#441) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…upstream deepseek-ai#441) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…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>
…-ai#471) into dev Two-parent merge of upstream 1fc9c66 (deepseek-ai#441, deepseek-ai#462, deepseek-ai#471) into sgl dev b296ce3, keeping the SGL wheel, the TORCH_LIBRARY binding and the downstream kernels. - Mega MoE: upstream locality-domain scheduling, 4D weight TMA, L1/L2 store blocks and the TMEM prefetch epilogue, with NVFP4/MXFP4, TRT-LLM weights, alphas, SiTU and FP8 combine carried over. TRT weights stay unlocalized. - GEMM epilogue classes threaded through every dispatch; SM90/SM120 assert that no custom epilogue is set. SM120 kernels drop the head-split epilogue parameter, as upstream did for SM90. - Attention takes upstream's API as-is: no clean_logits/logits_dtype, no legacy fp8_mqa_logits/fp8_paged_mqa_logits, no fp8_gemm_nt_skip_head_mid. The SM120 MQA kernels that only served the old API are removed. - Mega gate is sqrtsoftplus-only; bf16_mega_gate defaults to it and get_bf16_mega_gate_config is gone. - TORCH_LIBRARY binding: epilogue classes (Identity, Alpha, FP8CastWithClamp, BF16StochasticRounding) cross as (kind, scalar, sfd, count) into every GEMM/einsum op; shutdown and the locality-domain ops are registered; MLOPart reaches Python through ctypes callbacks installed by deep_gemm/locality_domain.py, which re-raises the callback's exception from the allocating call. - sgl_deep_gemm tests pick up upstream's test changes and the new test_locality_domain.py. - Dev's per-expert post-norm (#104) sits on upstream's restructured epilogue and combine. Its store block of 48 at block_m 192 becomes the L1 store block, upstream's FP8xFP4 L1/L2 blocks at 240 are kept, and DG_MEGA_MOE_TILES sets both L1 and L2. - deep_jit stays at dev's 3732a3b, a descendant of upstream's 2efdab4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ence Root cause: deepseek-ai#441 put `fence.acq_rel.cta` (MEMBAR.ALL.CTA) before every task-info slot release in the epilogue. Besides the slot reads, the fence waits for all of the thread's outstanding memory operations, including the previous block's remote combine stores, which costs ~6 us per fp8xfp4 Mega MoE call at a 32-rank decode shape. Fix: copy the slot after the full-barrier wait with two `ld.acquire.cta.shared::cta.v4.u32` loads and drop the fence. The acquire orders every later operation of the thread, including the release arrive, after the slot reads, at the same scope as deepseek-ai#441. ptxas emits no MEMBAR or CCTL for it, and nothing waits on unrelated stores. Validation (32 ranks over 8 GB300 nodes, fp8xfp4, 416 experts, top-6, 108 tokens/rank, BLOCK_M 64, CUDA graph, interleaved rounds): slowest-rank median 89.6-89.8 -> 83.5-83.8 us/call, equal to dropping the fence; outputs bitwise identical to the fenced kernel and across 200 graph replays x 10 calls on every rank.
…ad of a fence (#145) * Order Mega MoE task-info slot reads with acquire loads instead of a fence Root cause: deepseek-ai#441 put `fence.acq_rel.cta` (MEMBAR.ALL.CTA) before every task-info slot release in the epilogue. Besides the slot reads, the fence waits for all of the thread's outstanding memory operations, including the previous block's remote combine stores, which costs ~6 us per fp8xfp4 Mega MoE call at a 32-rank decode shape. Fix: copy the slot after the full-barrier wait with two `ld.acquire.cta.shared::cta.v4.u32` loads and drop the fence. The acquire orders every later operation of the thread, including the release arrive, after the slot reads, at the same scope as deepseek-ai#441. ptxas emits no MEMBAR or CCTL for it, and nothing waits on unrelated stores. Validation (32 ranks over 8 GB300 nodes, fp8xfp4, 416 experts, top-6, 108 tokens/rank, BLOCK_M 64, CUDA graph, interleaved rounds): slowest-rank median 89.6-89.8 -> 83.5-83.8 us/call, equal to dropping the fence; outputs bitwise identical to the fenced kernel and across 200 graph replays x 10 calls on every rank. * Make the task-info slot load an overload of ld_acquire_cta_shared --------- Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
No description provided.