Public release 26/07 - #377
Conversation
| num_shared_experts: int = 0, | ||
| use_fp8_dispatch: Union[bool, None] = None, |
There was a problem hiding this comment.
🟡 warning: Preserve the existing positional argument order: Existing positional calls now bind use_fp8_dispatch to num_shared_experts. For example, get_symm_buffer_for_mega_moe(..., False, "bf16xbf16") previously selected BF16, but now leaves mma_type as FP8 and fails the compatibility assertion. Append the new parameter after the existing optional parameters or make it keyword-only.
🤖 v6
| from deep_gemm.utils import align, get_mk_alignment_for_contiguous_layout | ||
|
|
||
|
|
||
| def assert_psum_zero_padding(a: torch.Tensor | tuple, d: torch.Tensor, grouped_layout: torch.Tensor, dtype_label: str) -> None: |
There was a problem hiding this comment.
🟡 warning: Use Python 3.8-compatible annotation syntax: The project documents Python 3.8+ support, but the torch.Tensor | tuple syntax requires Python 3.10. On Python 3.8 or 3.9, importing this helper raises a syntax error and prevents the BF16 and FP8/FP4 tests from being collected; use typing.Union instead.
🤖 v6
| @@ -69,32 +70,13 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, | |||
| num_experts: int, | |||
| num_max_tokens_per_rank: int, num_topk: int, | |||
| hidden: int, intermediate_hidden: int, | |||
There was a problem hiding this comment.
🟡 warning: Buffer is now sized to the WORST-CASE full ring over all candidate block sizes (get_symm_buffer_size_for_mega_moe) and the entire old per-call ring-budget/experts-per-wave host heuristic is deleted. The removed code deliberately chose a much smaller ring for prefill vs. decode and balanced L2/cache reuse by selecting experts-per-wave; the new code uses a maximum ring capacity + in-kernel task pool. This is simpler/correct but can substantially increase symmetric-memory footprint and changes runtime scheduling behavior. Please confirm the memory/scheduling regression (and the loss of the cache-reuse/tail-wave tuning) is intended and benchmarked.
🤖 v4
| get_num_wave_pool_tokens(num_ranks, num_topk, num_max_tokens_per_rank, 1, layout::kLCMCandidateBlockM), | ||
| get_num_wave_pool_tokens(num_ranks, num_topk, num_max_tokens_per_rank, num_experts_per_rank, layout::kLCMCandidateBlockM) | ||
| }; | ||
| static int get_block_m_for_mega_moe( |
There was a problem hiding this comment.
🟡 warning: The new host-side get_block_m_for_mega_moe recomputes and returns only one BLOCK_M by re-running a copy of the (now heavily rewritten) config path with runtime num_sms/num_tokens, while the actual BLOCK_M/schedule is decided inside the in-kernel scheduler/MegaMoEBuffer. The host value is only advisory, so any divergence (e.g. different SM count or num_tokens between the sizing/build path and runtime/PDL) gives a wrong schedule. Consider returning the full selected config, or unifying the host and kernel scheduling so they cannot mismatch.
🤖 v4
| input_topk_idx_buffer.get_end_ptr()); | ||
|
|
||
| // Shared expert buffers | ||
| shared_l1_token_buffer = input_token_buffer; |
There was a problem hiding this comment.
🔴 critical: shared_l1_token_buffer = input_token_buffer (and the host slice sets shared_l1_acts = x), i.e. the shared-expert L1 activation aliases the per-rank input token buffer, and the shared/ring buffers use a different alignment/capacity basis than the routed ring (kMinCandidateBlockM/gran alignment). If any shared expert reads input tokens that a routed expert has not yet consumed, or runs before/after routed pull in the in-kernel scheduler, the aliasing silently corrupts data. The in-place/ordering invariant must be enforced (asserted) in the scheduler/kernel, otherwise this is a data-race/corruption hazard.
🤖 v4
| const int num_warmup_waves_for_interleave_schedule = math::constexpr_ceil_div( | ||
| num_l1_n_clusters + (num_total_m_blocks - 1) * num_interleave_cluster_diff_per_m_block, | ||
| num_clusters) + 1; | ||
|
|
There was a problem hiding this comment.
🔵 suggestion: get_num_max_live_pool_blocks hardcodes kMegaMoEBlockN=128 and kNumCTAsPerCluster=2 as local constexprs and asserts (intermediate_hidden2) % (2128) and hidden % (2*128). These duplicate the BLOCK_N / cluster assumptions that also live in the scheduler template params (BLOCK_N, kNumL1Clusters=kNumL1BlockNs/2). If BLOCK_N ever changes for a config, this helper would silently disagree with the actual kernel geometry. Derive these from a shared constant or add a static_assert coupling them to the scheduler's BLOCK_N.
🤖 v3
| } | ||
|
|
||
| static std::tuple<int64_t, std::function<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(const torch::Tensor&)>> | ||
| static std::tuple<int64_t, std::function<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, |
There was a problem hiding this comment.
🟡 warning: get_symm_buffer_size_for_mega_moe sizes the ring buffer from device_runtime->get_num_sms() with the note 'we temporarily assume the SM count is consistent with the runtime value'. Both the returned buffer size and the kernel warmup-wave logic depend on kNumSMs, so a buffer allocated in a context with a different SM count than the one that later launches the kernel could under-provision the ring — a correctness (not just perf) hazard. Add an assert at launch that the num_sms used for sizing matches the kernel launch's SM count, or otherwise pin/validate this value.
🤖 v3
| @@ -83,62 +83,42 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::opt | |||
| const at::ScalarType& logits_dtype) { | |||
There was a problem hiding this comment.
🔵 suggestion: is_fp4 is now derived from the tensor dtype (qk_dtype==kPackedFP4) while is_mx_sf is derived from q_sf.has_value(). A plain-FP8 call that mistakenly supplies q_sf would now silently take the is_mx_sf branch (int32 SF, arch_major==10 required) instead of the legacy float-SF FP8 path, with no explicit rejection. Consider asserting that q_sf is only present for MX-scaled inputs (e.g. is_mx_sf implies the caller intended an MX path), so a stray scale tensor fails loudly rather than changing the dispatch semantics.
🤖 v3
| # Must with FP32 accumulation and 1D1D kernels | ||
| for num_groups, m, n, expected_k_per_group in (( 4, 4096, 7168, 8192), ( 4, 7168, 2048, 8192), # EP64 | ||
| ( 8, 768, 2048, 128), ( 8, 4096, 7168, 4096), ( 8, 7168, 2048, 4096), # EP32 | ||
| # NOTES: the first shape has many small groups, for stressing the SM90 in-place tensor map update |
There was a problem hiding this comment.
🔵 suggestion: The reordered k-grouped shape (8, 768, 2048, 128) placed first is intended to stress the SM90 in-place tensormap update paired with the sm90_fp8_gemm_1d1d.cuh change (kIsMulticastOnA false->true and the __syncwarp fix). Confirm this test is actually exercised on an SM90 target in CI (it is arch-gated); otherwise the tensormap race fix ships untested.
🤖 v3
| mma_type: str = 'fp8xfp4', | ||
| activation: str = 'swiglu') -> SymmBuffer: | ||
| # Align token count | ||
| num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) |
There was a problem hiding this comment.
🟡 warning: SymmBuffer / get_symm_buffer_for_mega_moe drop the num_ring_tokens parameter and the >=6144-token heuristic, replacing them with num_shared_experts and delegating ring sizing to C++. This is a breaking public-API change (num_ring_tokens argument and self.num_ring_tokens attribute removed; new num_shared_experts arg; fp8_fp4_mega_moe/bf16_mega_moe gain shared_l1/l2 weight args). Update all downstream callers and add a note to the README News section documenting the breaking mega API change.
🤖 v3
🤖 ds-review-bot Code Reviewv6The patch introduces a backward-incompatible positional API change and makes the test suite unparsable on documented supported Python versions. v4This is a large, squashed v3Public release 26/07 (commit 61ca7f1) for DeepGEMM: a 30-file, +1909/-1142 release focused on the Mega MoE path and the MQA logits (lightning indexer) kernels. Main themes: (1) Mega MoE scheduler rewrite from a host-side 'experts per wave' heuristic to an in-kernel, dynamically scheduled, atomic-counter-based producer/consumer scheme with a warmup-wave-based L1->L2 deadlock-avoidance analysis and a host-callable ring-capacity estimator (get_num_max_live_pool_blocks / get_block_m_for_mega_moe); (2) shared-experts support added end-to-end for both bf16 and fp8xfp4 Mega MoE (new optional weight args, new symmetric-buffer slices, new SharedLinear1/2 block phases, MegaMoEBuffer layout refactor); (3) MQA logits generalized from a binary kIsFP4 switch to a (qk_dtype_t, kIsMXSF) parameterization so it supports FP8 / MXFP4 / MXFP8 uniformly, with a shared attention.hpp validation/dispatch path and block_kv=128 support; (4) infrastructure: packed-FP4 UMMA descriptor pack-factor handling in mma/sm100.cuh, new PTX helpers (st_async_cluster/mapa_shared/atomic_add(u32)/ld_volatile(u32)), DG_IN_CUDA_COMPILATION guards so device-only PTX headers can be included from host TUs, and per_token_cast_to_fp8 SF padding for small head dims. Also a targeted SM90 fix (kIsMulticastOnA flag flip + __syncwarp fix + test shape reordering to stress in-place tensormap updates) and a new tests/utils.py helper. Verified as correct / no action needed: get_umma_desc_pack_factor / packed-FP4 make_umma_desc generalization is sound (byte-pointer arithmetic with kPackFactor, K-major-only compile-time guard) and correctly lets MQA logits drop the hand-rolled kFP4Layout descriptor; the smem_weights alignment change (dropping explicit weights padding) is safe because BLOCK_Q=128/kNumHeads makes BLOCK_Q*kNumHeads==128, so each per-stage weight row is 128 elements and remains 128B aligned; get_logical_shape plus the new t.is_cuda() assert cleanly replaces scattered 'head_dim *= 2' logic; the attention.hpp FP4/FP8 branch dedup keyed on (is_fp4 from dtype, is_mx_sf from q_sf presence) is internally consistent and additionally enables MXFP8; per_token_cast_to_fp8's 1.0-valued SF padding covers no real elements and is never read by the kernel. Overall the release is internally consistent and the alignment/deadlock reasoning checks out; the remaining comments are hardening/maintainability suggestions rather than blocking defects. Files reviewed: 30 |
Brings deepseek-ai PR #377 (559d79f, including #364 and #343) onto nv_dev (a6b593d = #369 merge). Merge base is 88965b0. Conflict resolution highlights: * MQA: adopt #377's unified SM100 implementation and common MMA/PTX refactor while preserving nv_dev's SM90 next_n=4 multicast and block_kv=32 support, dedicated SM120 paths, FP16/BF16 weights, and 16-head indexer coverage. * Scheduler safety: retain zero-context and empty-varlen handling; drive producer prefetch from the scheduler's actual next query; bound the SM90 block_kv=32 odd-page tail before reading block tables. * FP16 weights: port the implementation to the new headers and preserve padded host offsets, per-warpgroup row indexing, and per-row compressed window bounds. * GEMM/layout: combine #377's grouped scale-factor, PSUM, zero-padding, and Mega-MoE changes with nv_dev's SM120 dense/grouped dispatch. * Tests: merge upstream coverage with nv_dev's SM90/SM100/SM120 cases and add regressions for empty metadata, skipped requests, FP16 compressed windows, and the SM90 odd physical-page tail. Validation: * All Git-tracked Python files pass py_compile. * No unresolved conflicts, conflict markers, missing project includes, or references to removed headers were found. * CUDA build/JIT and GPU numerical tests were not run because this host has no nvcc, PyTorch, or visible GPU tooling.
Sync nv_dev with upstream #377
3 participants
Uh oh!
There was an error while loading. Please reload this page.