Repository navigation
Integrate main APIs and SM120 kernels while preserving nv_dev features - #447
Conversation
* Public release 26/09 * Update News
Port contiguous and paged sparse kernels to DeepJIT while preserving SM100 dispatch, stream-local metadata, entry-balanced scheduling and register guards. Restore histogram synchronization and capture-safe UE8M0 validation, with independent SM120 numerical and CUDA graph regression fixtures.
Restore the original BF16/FP8/FP4 kernels, scheduler, heuristics, and split-K paths from 139f504 through DeepJIT adapters. Support alpha, independent C/D strides, GPU-only KPSUM, and bounded padding cleanup. Separately repair inherited mixed-tail, producer handoff, grouped row ownership, and TMA output-width defects. Preserve the original TMA/MMA pipelines and dispatch choices except for the documented safety eligibility guard. This is a bounded restoration checkpoint, not full migration or universal performance parity. BF16 odd-N, FP4 K-grouped, and remaining public integrations are unfinished.
Resolve the nv_dev sync using an ordinary merge, following the ancestry-preserving approach of prior sync PRs deepseek-ai#314, deepseek-ai#328, deepseek-ai#349 and deepseek-ai#384. Audit all 75 target-side changed paths, including 59 conflicts and seven automatically changed paths, against the prepared main-API implementation. Retain DeepJIT, native SM120 kernels, restored nv_dev attention and MegaMoE contracts, and the validated NVFP4 WAR fix. Preserve the target's FP16 MQA indexing include; this is the only tree change relative to the validated integration head. Standard standalone build, delivered-package import, 81 host protocol tests, 349 SM120 focused tests, packing and HC passed. Runtime sources are unchanged; prior bounded hardware and sanitizer evidence applies. Existing combined IDE indexing build limitations and previously documented qualification limits remain.
|
Resolved the target-branch conflicts in 5500e69 using an ordinary ancestry-preserving merge, following the prior nv_dev sync patterns in #314, #328, #349 and #384. No rebase, force-push or blanket ours strategy was used. Audited all 75 target-side changed paths, including 59 conflicts and seven automatically changed paths. The resolved tree retains the main APIs/DeepJIT and the prepared SM120/nv_dev feature integration. The only tree delta from the previously validated head is the target’s FP16 MQA header include in the developer indexing translation unit; runtime implementations, Python, device headers and tests are unchanged. A fresh standard standalone build, external import, 81 CPU protocol tests, 349 SM120 focused tests, packing and HC passed against the resolved merge artifact. The added header also compiled independently. The optional combined IDE indexing target retains pre-existing limitations (missing pybind11 CMake package and a shared-memory declaration conflict also reproduced without this added include); these are not reported as passing. The PR remains Draft for the qualification limits described above; conflict resolution does not clear the outstanding SM100 multirank, extended ring sanitizer or NCCL initialization diagnostics. |
| } else if (arch_major == 12) { | ||
|
|
||
| const auto b_data = sm120_to_k_major(b.first, major_b, n); | ||
| const bool is_mixed_fp4 = (a.first.scalar_type() != b_data.scalar_type()) and | ||
| (a.first.scalar_type() == kPackedFP4 or b_data.scalar_type() == kPackedFP4); | ||
| DG_HOST_ASSERT(!is_mixed_fp4 or k % 128 == 0); | ||
| sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d(a.first, sfa, b_data, sfb, d, grouped_layout, |
There was a problem hiding this comment.
🟡 warning: 在 SM120 分组 GEMM 分派前拒绝 FP32 缩放因子: 在 SM120 上传入 FP32 scales、recipe=(1,1,128) 和 disable_ue8m0_cast=True 时，SF 转换会保留 FP32 格式。这里移除 sfa.scalar_type() == torch::kInt 校验后，会将其直接交给仅支持 packed UE8M0 的内核，导致静默数值错误：例如 FP32 的 1.0 会被拆成指数字节 [0,0,128,63]，而非单位缩放。contiguous 和 masked 两个分支都应校验转换后的 scale 类型，或像 dense 路径一样拒绝该参数组合。
🤖 v6
There was a problem hiding this comment.
Fixed in 3a04a69: both SM120 grouped entry points now require INT32 scales on both A and B when disable_ue8m0_cast=True, before transforms or launches. Automatic FP32 conversion and standalone-transform compatibility remain unchanged. The regression covers 128 parameter combinations, including independently packed/FP32 A/B scales; the final SM120 artifact passed the 226-case memcheck/racecheck selection with zero errors and zero race hazards/warnings.
| } else if (task_info.block_phase == sched::nvfp4::BlockPhase::Linear1) { | ||
| // Wait L2 block empty | ||
| const auto l2_empty_ptr = workspace.get_l2_empty_count_ptr(ring_block_idx); | ||
| const auto num_expected_blocks = (L2_SHAPE_N / BLOCK_N) * (pool_block_idx / kNumRingBlocks); |
There was a problem hiding this comment.
🟡 warning: The NVFP4 amax reduction still uses math::ReduceMax<float>(), which in main is now specialized to fmaxf (NaN-dropping). nv_dev's original — and the private copy you restored for SM90 in sm90_mega_moe_math.cuh — is a > b ? a : b (NaN-propagating depending on operand order). Results are identical for finite inputs, but a NaN activation may now be quantized with a finite SF instead of poisoning amax. This contradicts the PR statement that helpers were restored privately "without replacing kernel arithmetic" for the NVFP4 path. Either point NVFP4 at the private functor (rename sm90_moe_math → nv_moe_math) or explicitly document the intentional change.
🤖 v5
There was a problem hiding this comment.
Fixed in 3a04a69: the two NVFP4 warp-amax reductions now use a private copy of the original a > b ? a : b functor; main and SM90 math helpers are unchanged. The earlier frozen B300 artifact passed a GPU regression of the extracted functor, including operand-order-sensitive NaN cases and finite/Inf inputs. This establishes helper preservation, not end-to-end NaN poisoning by the quantization kernel; downstream maxima/clamping can still discard NaNs.
| @@ -210,7 +232,7 @@ void sm90_fp8_paged_mqa_logits(const uint32_t batch_size, | |||
| : 0; | |||
There was a problem hiding this comment.
🔵 suggestion: "Multicast" is a misnomer here: both CTAs in the cluster issue their own tma::copy for KV with the default num_tma_multicast = 1, so no TMA multicast happens — the 2-CTA cluster only splits next_n across CTAs and KV traffic is doubled relative to real multicast. Suggest renaming kNumKVMulticast → kNumCTAsPerCluster (host side too) and fixing this comment, or tracking true multicast as a follow-up, so readers don't assume the bandwidth benefit exists.
🤖 v5
There was a problem hiding this comment.
Agreed; corrected in 3a04a69. Renamed the device/host counts to kNumCTAsPerCluster / num_ctas_per_cluster and clarified that the CTAs split next_n while issuing independent KV copies. Cluster size, template positions, launch arithmetic and numerical operations are unchanged; no TMA multicast or associated bandwidth benefit is claimed.
| // Create metadata tensor. `num_sms` here is actually the scheduler slot count | ||
| // (= num_clusters on SM90 next_n=4 multicast, = num_sms elsewhere); callers | ||
| // pre-divide. | ||
| // num_sms is the scheduler slot count: callers pre-divide by 2 for SM90 next_n=4. |
There was a problem hiding this comment.
🔵 suggestion: Inherited from nv_dev, not a blocker: requiring callers to pass get_num_sms() // 2 for SM90 next_n=4 is a leaky contract (a wrong value only fails later in fp8_fp4_paged_mqa_logits via the _schedule_meta_size assert). This function already has context_lens (→ next_n) and the arch, so it can compute the slot count itself and keep num_sms semantics uniform across SM90/SM100/SM120. Worth a follow-up issue.
🤖 v5
There was a problem hiding this comment.
Agreed that a next_n-aware helper would be cleaner, but deliberately deferred in 3a04a69. The existing metadata argument counts scheduler slots: SM90 next_n=4 callers already pass get_num_sms() // 2. Dividing again inside this API would break migrated callers and metadata shapes. The documented slot-count contract and shape guards remain; a future explicit/versioned convenience API needs migration guidance rather than a silent semantic change.
| l1.shape != (2 * width, buffer.hidden) or l2.size(0) != buffer.hidden): | ||
| raise ValueError('Shared-expert weights do not match the symmetric buffer layout') | ||
| count = width // buffer.intermediate_hidden | ||
| if buffer.num_shared_experts != count: |
There was a problem hiding this comment.
🔵 suggestion: Main-visible behavior tightening not mentioned in the description: fp8_fp4_mega_moe / bf16_mega_moe now raise when buffer.num_shared_experts != <shared experts in the call>. On main this was allowed and safe (input views precede the shared-expert regions in MegaMoEBuffer, and byte capacity suffices), so existing callers that reuse one buffer for layers with and without shared experts will start failing. The base= path still covers the use case, but please call this out in the description/README as an intentional main-API change.
🤖 v5
There was a problem hiding this comment.
Fixed rather than documenting a new restriction in 3a04a69: exact shared-expert-count equality now applies only to NVFP4. Main paths retain capacity checks and call-time C++ re-slicing, allowing a sufficiently sized buffer to serve different counts. The earlier frozen B300 artifact passed the 2 -> 0 -> 1 -> 2 reuse regression for BF16, FP8xFP8 and FP8xFP4, eager and graph, with the existing numerical oracle. NVFP4 strict layout/protocol checks remain unchanged.
| if (runtime->device.get_arch_major() == 12) { | ||
| runtime->default_compiler_options.arch = "120f"; | ||
| runtime->default_compiler_options.extra_nvcc_flags.emplace_back( | ||
| "-gencode=arch=compute_120f,code=sm_120f"); |
There was a problem hiding this comment.
🔵 suggestion: DeepJIT's Device::get_arch() already yields "120f" for major 12 and CompilerOptions::get_flags() emits --gpu-architecture=sm_120f, so both the explicit arch = "120f" override and this extra -gencode=arch=compute_120f,code=sm_120f are redundant (duplicate codegen for the same target). It evidently works since the SM120 suite passed, but it can be dropped.
🤖 v5
There was a problem hiding this comment.
Partly applied in 3a04a69: removed the duplicate -gencode and the redundant arch override for CC 12.0. The pinned DeepJIT returns 120f for 12.0 but 121f for 12.1, so the 120f override is retained for nonzero 12.x minors to avoid silently changing the existing target. This is conservative target preservation, not an SM121 runtime-validation or support claim.
| DG_HOST_ASSERT(get_major_type_ab(a) == cute::UMMA::Major::K); | ||
| DG_HOST_ASSERT(get_major_type_ab(b) == cute::UMMA::Major::K); | ||
| check_major_type_cd(d); | ||
| if (cached.major != 12 or a.size(0) != 0) |
There was a problem hiding this comment.
🔵 suggestion: Skipping check_major_type_cd(d) only when major == 12 && m == 0 needs a comment. If the motivation is a degenerate stride on a zero-row d, the condition should be a.size(0) == 0 for all arches rather than SM120-specific. Also note the new device-property equality assert above applies to SM90/SM100 as well (reasonable, but it is a main-visible change on heterogeneous multi-GPU hosts).
🤖 v5
There was a problem hiding this comment.
Clarified in 3a04a69: ordinary PyTorch empty split outputs have a nonzero batch stride, which fails the legacy layout guard. SM120 accepts these as a no-op; SM90/SM100 retain their existing acceptance/rejection boundary rather than receiving an unrelated API extension. Older paths already return early after validation, so this is not a claim that their kernels cannot handle an empty request. Same-device, shape and dtype checks remain before return. The earlier H200/B300 artifacts each passed six empty-layout contract cases.
| return task_info.is_valid(); | ||
| } | ||
|
|
||
| CUTLASS_DEVICE void release_task_info() const { |
There was a problem hiding this comment.
🟡 warning: NVFP4 release_task_info() arrives on the task-info empty barrier without first fencing the consumer's metadata reads. This is a new copy of the exact pattern that main later fixed in #441 (fence.acq_rel.cta before the arrive) to prevent the producer from overwriting task_infos[...] while a consumer is still reading it. Because this scheduler is introduced by this PR, add ptx::fence_acq_rel_cta(); before task_info_empty_barriers[sched_stage_idx ^ 1].arrive(0u); (and the helper in ptx/ld_st.cuh). This changes no arithmetic or rounding.
🤖 v4
There was a problem hiding this comment.
Applied in 3a04a69: ported the #441 ptx::fence_acq_rel_cta() helper, including its compiler memory clobber, and placed it before the NVFP4 task-info empty-barrier arrival. This orders consumer metadata reads before slot reuse without changing arithmetic. The earlier frozen B300 artifact passed canonical stateful NVFP4 tests and the shared-one, global-PDL-enabled memcheck/racecheck run with zero errors and zero race hazards/warnings; extended ring-wrap qualification is still outstanding.
| DG_STATIC_ASSERT(SHARED_L1_SHAPE_K % SHARED_BLOCK_K == 0, "Invalid shared shape"); | ||
| DG_STATIC_ASSERT(SHARED_L2_SHAPE_K % SHARED_BLOCK_K == 0, "Invalid shared shape"); | ||
| DG_STATIC_ASSERT(SHARED_L1_SHAPE_K % BLOCK_K == 0, "Invalid shared shape"); | ||
| DG_STATIC_ASSERT(SHARED_L2_SHAPE_K % BLOCK_K == 0, "Invalid shared shape"); | ||
|
|
||
| // NOTES: N block counts must be even so that 2 adjacent CTAs in a cluster | ||
| // always land on the same m_block_idx with n_block_idx differing by 1 |
There was a problem hiding this comment.
🟡 warning: The preserved main MegaMoE scheduler is the #432 baseline and also lacks the #441 release-ordering fence, so sm100_bf16_mega_moe and sm100_fp8_fp4_mega_moe inherit the ordering window. This PR is the vehicle that brings #432 into nv_dev; either include #441 here or explicitly record it as an outstanding follow-up in the PR description.
🤖 v4
There was a problem hiding this comment.
Applied in 3a04a69: the main scheduler now uses the same #441 fence-before-empty-arrive ordering. Also added pre-election __syncwarp() barriers before combine-stage refill in the main BF16 and FP8/FP4 kernels. On the earlier frozen B300 artifact, the unchanged 12-case combine reproduction went from racecheck failures to zero hazards/errors/warnings; the three-dtype production stress also passed memcheck/racecheck without filters or suppression. These are ordering fixes, not arithmetic changes.
| @@ -1,41 +1,56 @@ | |||
| # DeepGEMM | |||
There was a problem hiding this comment.
🔵 suggestion: The documented SM120 coverage and scaling-factor contract should stay synchronized with the actual dispatch guards (e.g. K-grouped FP4 NT is rejected, paged clean_logits is unsupported, SM100-only scheduling metadata is unsupported). The text currently matches the host checks I inspected, but it is worth re-checking after any further conflict resolution so the docs do not overstate support.
🤖 v4
There was a problem hiding this comment.
Rechecked against 3a04a69. The README and scaling-factor contract still match the dispatch exclusions: no SM120 K-grouped FP4 NT, paged cleaning, SM100-only dense scheduling metadata, or dense/paged MXFP8 mode. Standalone FP32-scale transform compatibility does not imply an FP32-scale SM120 MMA path. The final change restores the grouped TMA pipeline and removes obsolete cp.async-specific code/comments; it does not expand the documented architecture coverage.
🤖 ds-review-bot Code Reviewv6SM120 分组 GEMM 存在缩放格式误解释，可产生静默数值错误。当前环境缺少 PyTorch/CUDA，未复跑 GPU 测试；讨论中 SM100 多卡及扩展环形缓冲区 sanitizer 验证仍未完成。 v5Reviewed #447 at merge commit 5500e69 (parents 794783a + nv_dev 572557e). Verified: (1) the merge commit's only tree delta vs. the previously validated head is the single Findings (none blocking for a draft, but the first two are cheap and touch paths this PR claims to harden): a hang-instead-of-trap risk in the new SM90 paged scheduler zero-context skip loop; a subtle NaN-semantics change in the NVFP4 amax reduction because it still uses main's now- The draft caveats (single-GPU B300, timed-out 10-call large-ring racecheck, NCCL 2.30.7 init diagnostics) are stated honestly and nothing I found contradicts them. Before undrafting, re-run the long-ring racecheck with a larger bound so it can actually be counted, and address the first two findings. v4The integration is broad but appears faithful. I verified statically that: (1) the SM90/NVFP4 paths keep a private workspace (LCM 384, candidate block Ms up to 192) separate from main's (LCM 1920, with 240); (2) the restored shfl helper in sm90_fp8_paged_mqa_logits.cuh is identical to the helper main removed, so no arithmetic changed; (3) the NVFP4 combine WAR fix places __syncwarp() before the elected-lane TMA overwrite; (4) the NVFP4 base= key includes the configured SM count and weight validation enforces INT8/INT32, per-rank shapes, MN-major packed-SF strides and device, while main-compatible reuse remains allowed; and (5) the SM90 output/device/rank/ptr/SF checks and paged page-size/next_n guards are present, with the metadata call passing block_kv=64 only as a 256-divisor. The one substantive concern is that the integration is pinned at #432 and therefore omits the later main fix #441 (fence.acq_rel.cta before the task-info empty arrive). This matters most for the newly added NVFP4 scheduler, which is a copy of the unfenced pattern, and it also leaves the preserved main MegaMoE paths exposed. This is a correctness ordering fix with no arithmetic impact and should be ported or explicitly deferred. Hardware/sanitizer/NCCL claims in the description were not reproducible in this environment and remain unverified. Files reviewed: 172 📍 未定位到 diff 的评论🟡 warning while (current_num_kv == 0 and exist_q_atom_idx(current_q_atom_idx)) { ... }
...
if (not exist_q_atom_idx(current_q_atom_idx))
return false;This keeps the intended skip semantics but makes the loop bounded. 🤖 v5 |
| @@ -348,7 +344,7 @@ sm100_fp4_fp4_mega_moe_impl(void* y, | |||
| comm::cluster_sync_with_relaxed_arrive(); | |||
There was a problem hiding this comment.
Main's cudaGridDependencySynchronize() is imported into sm100_fp8_fp4_mega_moe.cuh only. This kernel and sm90_fp8_mega_moe.cuh are still launched with PDL enabled by default but never wait, so under set_pdl(True) their dispatch loads can race the producer kernel. Since this PR touches all three, could the same wait be added at the equivalent point here and in SM90 (or .enable_pdl = false in their launchers)?
There was a problem hiding this comment.
Thanks, the global-enabled case is a real ordering risk. Fixed conservatively in 3a04a69 with explicit .enable_pdl = false in the NVFP4 launcher and the shared SM90 launcher used by both L1/L2 phases. The pinned DeepJIT default is actually false, but set_pdl(True) previously propagated to these launches without device waits. The override preserves ordinary stream/graph ordering rather than introducing unqualified overlap. Earlier H200/B300 artifacts passed their SM90/NVFP4 stateful tests under global PDL 0 and 1; these are not a dedicated early-release-producer stress proof, nor final-artifact cross-architecture runtime qualification.
| for fmt in fmts: | ||
| is_mxfp4 = fmt == 'mxfp4' | ||
| for logits_dtype in (torch.bfloat16, torch.float): | ||
| weights_dtypes = (torch.float, torch.bfloat16, torch.float16) if arch_major == 10 else (torch.float, ) | ||
| for weights_dtype in weights_dtypes: | ||
| for weights_dtype in ((torch.float, torch.bfloat16) if get_arch_major() == 10 else (torch.float, )): |
There was a problem hiding this comment.
The SM100 float16 weights sweep (shapes up to (8192, 65536)) was dropped with main's sweep; test_nv_attention_restoration.py only covers fp16 at seq_len<=516. The f16-weights kernel is still dispatched, so please re-add float16 here or add large shapes to the restoration test.
There was a problem hiding this comment.
Restored in 3a04a69 as an independent 96-case large-FP16 matrix in test_nv_attention_restoration.py, up to (8192, 65536): 32/64 heads, dimensions 32/64/128, FP32/BF16 logits, full+clean and compressed output. It checks every output region with a bounded-memory independent reference, preserves the original tolerances, and checks bitwise repeatability. The earlier frozen B300 artifact passed all 150 SM100 restoration cases, including these 96. No FP16 production arithmetic was changed, and this historical result is not relabeled as a final-artifact B300 run.
Barry-Delaney
left a comment
Barry-Delaney
left a comment
There was a problem hiding this comment.
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh:2239: This combine loop has the same double-buffer structure as the NVFP4 one you fixed, but has neither the new __syncwarp() before elect_one_sync() nor main's fence_view_async_shared() after the reads (fp8_fp4 :1474). Same WAR window between the lanes' ld.shared of combine_load_buffer[s] and the next iteration's TMA into it; suggest applying the same fix here.
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
Review follow-up in 3a04a69; replies are attached to all 12 inline comments. Correctness / compatibility. Both grouped SM120 APIs reject disabled conversion unless both scales are INT32. NVFP4 restores the original private ternary ReduceMax (helper semantics, not an end-to-end NaN guarantee). Main buffer reuse with differing shared counts is restored; NVFP4 remains strict. Both MoE schedulers include #441 release fences, and combine-stage pre-refill warp barriers cover main BF16/FP8 and SM90. NVFP4 and both SM90 phases explicitly disable PDL, including when global PDL is enabled. CTA-cluster naming and HC empty-layout comments are corrected; duplicate SM120 codegen is removed without changing the existing nonzero-minor 12.x target. SM120 exclusions were rechecked. Metadata slot-count convenience remains deferred for compatibility. The general SM90 paged-scheduler finding is addressed with physical-batch bounds, lexicographic exclusive endpoints and guarded initial/final refreshes, plus metadata-prefix warp synchronization. This bounds stale/invalid traversal, not arbitrary malformed-input numerical behavior. Earlier frozen H200 evidence: 105 attention cases passed functionally and under racecheck; 16 bounds cases passed memcheck. The SM90 combine WAR fix is included, but MoE sanitizer qualification remains limited by NCCL 2.30.7 initialization diagnostics. Earlier frozen B300 evidence includes 150 attention cases (96 restored large FP16 cases up to 8192 x 65536), shared-count reuse, and clean main/NVFP4 targeted sanitizer runs. These H200/B300 results belong to their original artifacts, not final-artifact cross-architecture runtime runs. Final SM120 artifact (attempt 02). Standard build/import and CPU protocol checks passed: 94 passed, one GPU opt-in skipped. GPU selections passed: grouped-boundary 30 functional + 30 memcheck + 30 racecheck; exact reproduction 1; stage coverage 80; memcheck/racecheck 226 each; focused 459; packing/HC 3; expanded memcheck/racecheck 310 each. Every sanitizer run exited successfully: memcheck zero errors, racecheck zero hazards/errors/warnings, with no kernel filters or suppression. Pytest reported four non-sanitizer warnings in the broader selections. These selections overlap and must not be added together. Frozen source/artifact identity checks passed. Audited performance. Six paired rounds x 20 cases: final, V5 and TMA/direct-D control accepted 120 case-rounds each; upstream accepted 108 because mixed K=384/640 failed correctness, so no speedup is claimed for those failures. Timing is graph-replay GPU public-call latency after per-invocation L2 eviction, including copies/cleanup, not CPU/Python/allocator time. Displayed latency is the median of six round medians; percentages use the median of six paired ratios, not ratios of displayed medians. Negative percentages mean lower latency:
Both W2 variants remain about 3.9% slower than TMA/direct-D control: an explicit trade-off, not a blanket speedup. Six adapter cases / 36 rounds were accepted. Staged-copy medians are 1.376–3.713 us; nonzero cleanup medians 0.864–1.952 us. These are components, not the complete incremental adapter cost; component medians cannot be summed. Profiler attribution relies on capture/correlation consistency; calibration does not prove all benchmark dependency edges, and observed timestamp overlaps provide no absolute error bound. SM121 runtime, SM100 multirank and extended ring sanitizer qualification remain outstanding. This update does not change the PR body, draft/state, or resolve threads. |
- Wait for in-flight TMA descriptor reads before overwriting the per-CTA K-grouped BF16 descriptor, and add full-SM randomized alias/separate-C regressions. - Require packed INT32 SFB on the SM120 skip-head path instead of silently reading FP32 scales as UE8M0. - Apply HeadSplits index mapping with logical-column bounds in scalar stores and split-K reduction, replacing the former capability guards while keeping alpha/C/stride contracts. - Store BF16 batched einsum and HC prenorm directly to supported strided outputs, keeping the existing fallback for unsupported layouts. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Restore tests/test_{attention,bf16,fp8_fp4,hyperconnection}.py to their
upstream byte state and consolidate all SM120-specific coverage into
tests/test_sm120_*.py plus the shared tests/sm120_exercise.py helper.
This eliminates per-release merge conflicts with upstream's wholesale
rewrites of the shared test files.
Switch gating from collection-time string skipif (which initialized
CUDA during pytest collection on non-SM120 machines and broke the
sanitizer runner's direct-call path) to upstream's call-time
test_filter convention, and give each standalone file its own
__main__ runner matching upstream script style. Moved test bodies are
AST-identical; test counts and results are unchanged (741 passed both
before and after on SM120, memcheck/racecheck subsets clean).
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
d5e3bbf to
5a0d912
Compare
Adopt the standalone DeepGEMM-sm120 repository as the single source of truth for the SM120 device headers: add the CUDA>=13 compile-time guard (silent block_scale drop on pre-13 ptxas), relocate tensor_map_replace_global_dim_in_smem into common/sm120_utils.cuh so the kernels no longer depend on ptx/tma.cuh providing it, and inline the FP4 smem pack factor in layout/sparse_mqa_logits.cuh. Marker-free vendoring: the files carry no canonical-source comments or other foreign markers; provenance and drift control live on the canonical repo's side (per-tag manifest, check_vendor.py --fork, scheduled downstream watcher). Verified: 741 SM120 tests pass on sm_120a from a fresh JIT cache. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
0bf7f15 to
70a84f8
Compare
|
Force-pushed ( |
CUDA 13.3's ptxas crashes (SIGSEGV; DeepJIT compile subprocess exit 139) when compiling the non-warp-specialized MXFP4 sparse MQA logits variants at -O3 with any --register-usage-level (lucifer1004/DeepGEMM-sm120#2); CUDA 13.2, the validated toolchain, is unaffected. Add a toolchain-gated per-kernel override via DeepJIT CompilerOptions::extra_nvcc_flags: when nvcc reports >= 13.3, the non-warp-specialized instantiations compile with ptxas -O2, which avoids the crash and reproduces the 13.2 register allocation (REG:101 for the reference instantiation). The warp-specialized variant is not affected by the crash and keeps the default flags (its compiled reg==64 contract also verified under -O2). Validated on sm_120a: sparse MQA suite 9/9 under CUDA 13.3.73 with the gate active (previously 3 failing tests), and 9/9 under CUDA 13.2.86 with the gate inactive (unchanged default path). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Vendor DeepGEMM-sm120 v0.1.4, which derives BLOCK_KV = min(PAGE_KV, 64) in the FP8 paged kernel (mirroring the FP4 sibling): a 32-row page gets a 32-row compute tile and 4 KV groups at SPLIT_KV=128, instead of a 64-row tile straddling two non-contiguous physical pages. DSv4.1 indexer caches mix 64- and 32-state pages (vllm-project#14). Host glue: relax the FP8 launcher and fused-cache API gates to admit block_kv=32 for arch 12, derive tile_kv = min(block_kv, 64) like the FP4 launcher, and drop to two KV stages for page32 + 64 heads + paired queries, where three stages exceed the 99 KiB SMEM budget by 4 bytes. Tests: the paged-MQA contract matrix gains (fp8, page32), and a focused case covers the two-stage fallback (page32, 64 heads, paired/varlen). Validated on sm_120a from a fresh JIT cache: paged MQA suite 20/20 passed, incl. graph replay and legacy-API cross-checks. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
The stale/OOB-metadata protection added during the migration lived in the per-task hot path (compound has_next_task/exist_q_atom_idx checks per fetch), regressing paged MQA logits on SM90: ~1% on upstream shapes and 13-16% at next_n=4. Move the protection to the constructor: clamp end_q_atom_idx to the in-batch sentinel (zeroing end_kv_idx at the sentinel) and collapse out-of-range starts, then make the per-task end test a single ordered 64-bit compare on (q_atom, kv) packed keys. The ordered compare also keeps the hang fix: an advance overshooting the end is caught by >=, never by exact equality. context_lens/indices are only dereferenced when current_q_atom_idx is provably in batch (end is clamped at or below the sentinel, and the sentinel range always has end_kv == 0). Patch from local review. Validated on H200: test_paged_mqa_logits 24/24 cases pass; next_n=4 (H=64, D=128, page 64) 140.6us vs 154.6us pre-patch at L=8192 and 1093.3us vs 1233.1us at L=65536 (-9%/-11%), upstream-shape sweep at parity or better. Cross-compiled for sm_90a. Reviewer's full 105-case SM90 suite passes on their side. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
pack_fp32_into_ue8m0's per-group loop ran three runtime integer divisions per group per thread on the psum path (align of previous end, align of group K, ceil_div to SF rows). gran_k and k_alignment are powers of two at every call site (host-asserted), so detect that once per block and use shift forms instead; the generic division path is kept for any future non-pow2 alignment. The earlier suggestion to drop the align entirely on SM90/100 was unsafe: psum tail K sizes are not k_alignment-multiples, and without the align the row count comes up short (B300 fp8 psum accumulate regression, diff 0.024 vs 0.008 tolerance). Validated: pack output bitwise identical to the division form on a 128-group psum shape (B300); kernel time 38.7 -> 22.1 us/call (-43%) on the same shape. SM100 k-grouped fp8/bf16 contiguous suites pass on B300 (fresh JIT cache), SM120 k-grouped regressions pass on sm_120a, and the TU cross-compiles for sm_103a/sm_120a. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
psum end offsets are k_alignment-multiples and gran_k divides k_alignment there, so aligned_group_k % gran_k == 0 and the rounding shift is exact division; the ceil semantics matter only on SM120. Comment-only change. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
3aca951 (SF pack pow2 shift path) — the runtime pow2_layout branch only covers pow2 alignments. With k_alignment = 384 (a third of the k-grouped cases) the kernel is now 25 % slower than main (128 groups: 91 → 124 µs), reproduced in both run orders on B300; pow2 alignments are 18 % faster. Numerics are bit-exact. An alternative that covers both: pass gran_k / k_alignment into the JIT instantiation as template constants and keep the plain math::align / math::ceil_div; the compiler turns every division into mul-shift for any alignment, no branch. |
Follow-up to 3aca951: the runtime pow2 branch regressed non-pow2 k_alignment (reported by RayWang96: k_alignment=384, a third of their k-grouped cases, 91 -> 124 us vs main on B300). Instead of branching at runtime, pass gran_k/k_alignment as JIT template parameters: the per-group align/ceil_div fold to shifts (pow2) or multiply-high sequences (non-pow2) at compile time, covering both alignment classes with no runtime branch. Runtime kernel args drop accordingly. Validated on B300: pack output bitwise identical to the pre-change form at k_alignment 128 and 384; pack kernel 38.8 -> 14.0 us/call at 128 (-64%) and 65.2 -> 32.3 us/call at 384 (-50%); k-grouped fp8/bf16 contiguous suites pass from a fresh JIT cache. SM120 k-grouped regressions pass on sm_120a; the TU cross-compiles for sm_103a/sm_120a. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
@RayWang96 Good catch — adopted your suggestion in 4b3571e: gran_k/k_alignment are now JIT template parameters, so the per-group align/ceil_div fold at compile time (shifts for pow2, multiply-high for non-pow2) and the runtime branch is gone entirely. Reproduced your k_alignment=384 case on B300: the runtime-branch version was indeed slower there (65.2 us vs 38.8 us at 128 on my 128-group psum shape), and the template-constant form fixes it: 38.8 -> 14.0 us (-64%) at 128, 65.2 -> 32.3 us (-50%) at 384, bitwise-identical output at both, k-grouped fp8/bf16 suites pass from a fresh cache. |
Vendor lucifer1004/DeepGEMM-sm120 v0.1.5 (b31a688): 14 fence_view_async_shared() sites across 8 SM120 kernel headers, ordering generic-proxy smem reads (ldmatrix/ld_shared) before the empty-barrier arrive that hands each stage back to the TMA producer. Ports deepseek-ai#453 (bf16 GEMM) and fixes the same hazard found by audit in bmk_bnk_mn, tf32_hc_prenorm, fp8_fp4_gemm_1d1d, and the fp8/fp4 dense+paged MQA logits kernels. Sparse MQA is exempt (legacy cp.async producer). Perf A/B on RTX PRO 6000 (bf16 GEMM, fp8 dense/paged MQA, 3 runs): all deltas within run-to-run noise. Fork suite: 235 passed (GPU4, fresh JIT cache). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
No further issue, merging |
5 participants
Summary
Integrate the current main APIs / DeepJIT runtime and native SM120 kernels into
nv_dev, while preserving its branch-specific functionality. This is a substantial release integration, not a small SM120-only patch. The API/runtime baseline is main's Public Release 26/09 (#432); the intended merge target is nv_dev, not main.SM120 support
The existing SM120 kernel algorithms, heuristics and pipelines are retained through the API migration. DeepSelect is a separate project and is not included here. README and scaling-factor documentation describe unsupported combinations; this does not add SM120 MegaMoE, Mega Gate or Mega mHC support.
Preserve nv_dev features
The restored SM90/NVFP4 paths retain a private dynamic-counter workspace (LCM 384), separate from main's workspace (LCM 1920). NVFP4 uses E4M3 scales per 16 elements and is not aliased to main's MXFP4/UE8M0 enum or layout.
Correctness changes distinguished from mechanical migration
base=reuse to identical supported configurations, including the configured SM count; main-compatible buffer reuse remains available.Validation
Standard standalone builds and external package imports passed. Tests select the built package, not an injected source-tree import. Numerical thresholds were not relaxed.
Cautions
Attribution
The SM120 work builds on #324, with credit to @leavelet, adapted to the current main APIs and DeepJIT. Branch-specific functionality is preserved from nv_dev, including #383, #409 and #396. The test-helper discovery cleanup is tracked separately in #446.
SM120 device-layer provenance
The SM120 device headers are byte-identical vendored copies of lucifer1004/DeepGEMM-sm120 v0.1.4, the single source of truth shared with the vllm-project fork's SM120 layer (vllm-project#10). The vendored files carry no manifests, tooling, or header markers — they are plain DeepGEMM-style headers; provenance and drift control live entirely on the canonical repo's side (per-tag sha256 manifest,
check_vendor.py --fork, and a scheduled downstream watcher). Device-side deltas vs the pre-vendoring head of this PR: a CUDA >= 13 compile-time guard for SM120 device passes (a pre-13 ptxas was reported to silently dropblock_scale),tensor_map_replace_global_dim_in_smemrelocated intocommon/sm120_utils.cuh(main-lineage bases lack it inptx/tma.cuh), and the FP4 smem pack factor inlined intolayout/sparse_mqa_logits.cuh. Re-verified after vendoring: the full SM120 suite (741 tests at that head) passed on sm_120a from a fresh JIT cache against a freshly rebuilt package, and the vendored kernels are byte-identical at the current head. The labels-contract regression cases were later deferred to upstream #450 (see the note below); the directly affected regression file was re-run at the current head and passes.Note: m-grouped contiguous labels contract
A downstream field report ("middle empty group corrupts later groups") was root-caused to a caller-contract violation, not a kernel bug: labels built at a finer granularity than the runtime
mk_alignment_for_contiguous_layoutput a group boundary inside aBLOCK_Mtile whose B/SFB is selected from the tile's first-row label (verified with the reporter's exact shapes at alignments 32/64/128 — contract-respecting layouts are always safe). The opt-in host checker and regression coverage are proposed arch-generically to main in #450; this branch intentionally carries no fork-local copy and will pick the checker up at the next sync once #450 lands.CUDA 13.3 toolchain note (f7bea76)
CUDA 13.3's ptxas crashes (SIGSEGV) compiling the non-warp-specialized MXFP4 sparse MQA logits kernel at
-O3with any--register-usage-level— a compiler bug, not a kernel defect (root-caused in DeepGEMM-sm120#2; CUDA 13.2 remains the validated toolchain). So downstream users on 13.3 are not broken, this kernel now carries a toolchain-gated per-compile override (DeepJITCompilerOptions::extra_nvcc_flags): on nvcc >= 13.3 the affected instantiations build with ptxas-O2, which reproduces the 13.2 register allocation exactly (REG:101 reference). The warp-specialized variant is unaffected and untouched. Validated on sm_120a: sparse MQA suite 9/9 under both CUDA 13.3.73 (gate active; previously 3 failures) and 13.2.86 (gate inactive, default path unchanged).Follow-up folded in (b6acafe): FP8 paged MQA PAGE_KV=32
Vendored DeepGEMM-sm120 v0.1.4 plus host enablement: FP8 paged MQA logits admits block_kv=32 (API + launcher gates),
tile_kv = min(block_kv, 64), and a two-KV-stage fallback for page32 + 64 heads + paired queries (three stages exceed the 99 KiB SMEM budget by 4 bytes). Motivation: DSv4.1 indexer caches mix 64- and 32-state pages (vllm-project#14; kernel design mirrors the FP4 sibling). Validated on sm_120a from a fresh JIT cache: paged MQA suite 20/20 passed, including new (fp8, page32) contract cases, a focused two-stage-fallback case (page32 x 64 heads, paired/varlen), graph replay and legacy-API cross-checks.Follow-up folded in (d9aa49d): SM90 paged MQA scheduler perf
Local review found the migration's stale/OOB-metadata protection sat in the per-task hot path (SM90 paged MQA logits: ~+1% upstream shapes, +13-16% at next_n=4). The protection now runs once at scheduler construction (clamp end to the in-batch sentinel, collapse out-of-range starts), and the per-task end test is a single ordered 64-bit compare on packed (q_atom, kv) keys — which also preserves the hang fix (overshooting advances are caught by >=, never exact equality). Validated on H200: 24/24 paged cases pass; next_n=4 is 9-11% faster than the pre-patch head, upstream-shape sweep at parity or better.
Second review item (k-grouped SF pack, +1.7% on SM100): the suggestion was to keep the psum
alignand revert onlyceil_divto/on SM90/100. That is provably exact there (psum ends are k_alignment-multiples, gran_k | k_alignment, so aligned_group_k % gran_k == 0) — an earlier counter-test of mine that failed on B300 had additionally dropped thealign, which is load-bearing; the reviewer's actual suggestion stands. Landed as3aca951+1f78b85: a power-of-two shift fast path (pow2 holds at every call site, host-asserted) that subsumes the revert — shift forms equal exact division on SM90/100 and ceil on SM120 — with the generic division path retained for non-pow2 alignments. Follow-up4b3571e(review: non-pow2 k_alignment regressed under the runtime pow2 branch): gran_k/k_alignment moved into the JIT template, folding the align/ceil_div at compile time for pow2 AND non-pow2 alike; measured 38.8 -> 14.0 us at k_alignment=128 and 65.2 -> 32.3 us at 384 (B300, bitwise-identical). Pack output bitwise identical (128-group psum, B300); pack kernel 38.7 -> 22.1 us/call (-43%) on that shape; SM100 fp8/bf16 k-grouped suites pass on B300, SM120 k-grouped regressions pass on sm_120a.Follow-up folded in (381e2d7): v0.1.5 generic-proxy WAR fences
Vendored DeepGEMM-sm120 v0.1.5 (
b31a688): ports #453 (thanks @Sunt-ing) and extends it to every site with the same hazard — 14fence_view_async_shared()insertions across 8 SM120 kernel headers. The SM120 kernels read TMA-written smem throughldmatrix/ld_shared(generic proxy), so the empty-barrier arrive that hands a stage back to the TMA producer must be ordered after those reads; without the fence the next TMA can overlap the consumer's reads (generic/async proxy WAR hazard). Upstream's SM100 kernels carry the same fence for the analogous UMMA case ("Release KV smem only after UMMA commits TMEM"); the SM120 ports had inherited the barrier protocol but not the fences. Sites: bf16_gemm (the #453 report), bmk_bnk_mn, tf32_hc_prenorm_gemm, fp8_fp4_gemm_1d1d (×3), fp8/fp4 dense MQA logits (empty_kv + empty_q each), fp8/fp4 paged MQA logits (empty_q stage release + empty_kv each). Sparse MQA is exempt (legacy cp.async producer). Evidence: #453 author reproduces 2/6000 failures on RTX 5090 without the fence, 0/6000 with it; local RTX PRO 6000 30k-iteration soak does not reproduce either way, so the fence rests on the memory model plus the author's A/B/A, not on local observation. Perf A/B (bf16 GEMM, fp8 dense/paged MQA, 3 runs each): all deltas within run-to-run noise. SM120 suite re-run at this head from a fresh JIT cache: 235 passed.