Integrate main APIs and SM120 kernels while preserving nv_dev features - #447
lucifer1004 wants to merge 11 commits into
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>
4 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.