[PR #447] [MERGED] Integrate main APIs and SM120 kernels while preserving nv_dev features #455

New issue

Closed

opened 2026-10-26 21:18:03 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:03 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/447
Author: @lucifer1004
Created: 9/15/2026
Status: ✅ Merged
Merged: 9/23/2026
Merged by: @RayWang96

Base: nv_dev ← Head: feature/sm120-main-api-migration

📝 Commits (10+)

66081d4 Public Release 26/09 (#432)
b945d55 feat(sm120): add sparse MQA support on upstream main APIs
40b9d2d feat(sm120): restore native GEMM paths on main APIs
794783a feat(sm120): preserve nv_dev features alongside main API SM120 kernels
5500e69 chore(sync): merge nv_dev ancestry into main API integration
3a04a69 fix(kernels): address migration review and restore grouped TMA pipeline
48e5de7 fix(sm120): harden GEMM execution and complete HeadSplits output mapping
0aa7482 test(sm120): collect SM120 tests into dedicated files
70a84f8 feat(sm120): vendor device layer from DeepGEMM-sm120 v0.1.3
f7bea76 fix(sm120): gate ptxas -O2 for sparse MQA logits on CUDA 13.3+

📊 Changes

193 files changed (+22685 additions, -7023 deletions)

View changed files

📝 .gitmodules (+3 -3)
📝 CMakeLists.txt (+4 -2)
📝 README.md (+58 -61)
📝 csrc/apis/attention.hpp (+376 -128)
📝 csrc/apis/config.hpp (+17 -21)
📝 csrc/apis/einsum.hpp (+284 -90)
📝 csrc/apis/gemm.hpp (+438 -214)
📝 csrc/apis/hyperconnection.hpp (+52 -11)
📝 csrc/apis/layout.hpp (+38 -25)
➕ csrc/apis/mega_gate.hpp (+234 -0)
➕ csrc/apis/mega_mhc.hpp (+287 -0)
📝 csrc/apis/mega_moe.hpp (+42 -248)
➕ csrc/apis/nvfp4_mega_moe.hpp (+276 -0)
➕ csrc/apis/sm120_gemm.hpp (+86 -0)
📝 csrc/apis/sm90_mega.hpp (+26 -12)
📝 csrc/indexing/main.cu (+1 -1)
➖ csrc/jit/cache.hpp (+0 -31)
➖ csrc/jit/compiler.hpp (+0 -369)
➖ csrc/jit/device_runtime.hpp (+0 -153)
➖ csrc/jit/handle.hpp (+0 -222)

...and 80 more files

📄 Description

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

Native BF16, FP8, FP4 and mixed FP8/FP4 dense GEMM; M-grouped contiguous/masked paths; BF16/FP8 K-grouped paths.
FP8 einsum, native BF16 forward einsum and batch reduction, and HC prenorm.
Dense/paged FP8 and MXFP4 MQA, plus DSv4.1 MXFP8/MXFP4 sparse and paged sparse MQA, using the current metadata APIs.
Current scaling-factor/layout interfaces, dispatch validation, focused regression tests and support documentation.

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

SM90 FP8 MegaMoE, including split L1/L2 execution and natural FP32 block-128 weight scales (#383).
NVFP4 MegaMoE with BF16 shared experts, per-expert scales and routed scaling (#409).
SiTU for both routed and shared experts (#396).
SM100 dense MQA with FP16 weights, preserving its dedicated MMA path rather than replacing it with BF16. FP16 paged MQA is not added.
SM90 paged MQA with physical page sizes 32/64, next_n=1/2/4, and zero-work/scheduler safety fixes.

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

Restore removed math/shuffle helpers privately and adapt host launches to DeepJIT without replacing kernel arithmetic.
Validate SM90 output/device/rank inputs and reject incompatible buffer protocols/layouts.
Restrict NVFP4 base= reuse to identical supported configurations, including the configured SM count; main-compatible buffer reuse remains available.
Fix an inherited NVFP4 combine double-buffer WAR hazard with one warp synchronization before issuing a TMA overwrite. PC/source correlation and a one-token, non-graph reproduction established missing consumer completion before stage reuse. The pre-fix combine region matched nv_dev; the fix changes no arithmetic or rounding order.

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.

SM120, final artifact: 349 focused tests, 2 packing tests and 1 HC test passed. All 310 SM120-exclusive tests passed memcheck (0 errors) and racecheck (0 hazards/errors/warnings); the additional 39 focused cases were not sanitizer-tested. With GPUs hidden, 310 tests skipped without CUDA initialization. Intentional empty-graph PyTorch warnings are separate from sanitizer findings.
H200: 73 corrected paged MQA tests and host guards passed. SM90 MegaMoE smoke/stateful checks passed at 1/2/4/8 ranks, including unequal token counts, an empty rank, live/masked/live reuse, CUDA Graph replay and exact cumulative statistics. Larger 1/2/4-rank sweeps also passed. Selected original BF16/FP8 GEMM tests passed numerical checks and memcheck/racecheck. These runs preceded the final NVFP4-only synchronization change; the H200 code is unchanged.
B300, single rank: 54 attention tests, SiTU and main MegaMoE contract tests passed before the final NVFP4-only change. The final artifact passed 82 protocol tests, 27 NVFP4 matrix invocations and additional one-token controls. Coverage includes shared counts 0/1/2, alpha/A2 combinations, routed scaling, empty/masked/tail inputs, graph/state reuse and exact statistics.
NVFP4 repair: previously failing shared1/shared2 racechecks now report zero hazards/errors/warnings. Representative stateful memcheck passed. Actual ring wrap exercised 24,576 routed tokens against a 10,752-token ring; ordinary stateful execution and a separate single-large-call racecheck passed.

Cautions

The extended 10-call large-ring stateful racecheck exceeded its 900-second execution bound without a final sanitizer summary. It is not counted as passed; the successful small graph/stateful and single-large-call checks are distinct coverage.
SM100 multirank/NVLink qualification remains outstanding; the available B300 allocation is single-GPU.
H200 MoE memcheck is not clean because NCCL 2.30.7 probes Blackwell-only symmetric kernels during initialization. A minimal NCCL-only program reproduces the handled CUDA API diagnostics without importing DeepGEMM. This is not evidence of a DeepGEMM memory-access bug, but no suppression or clean sanitizer claim is made. No demonstrated release/configuration fix was found in the inspected upstream versions.
The complete inherited main-release diff still needs upstream review. This is bounded kernel qualification, not a full-repository, installation, model-serving, SM121 or universal performance-parity claim.

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/DeepGEMM#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 drop block_scale), tensor_map_replace_global_dim_in_smem relocated into common/sm120_utils.cuh (main-lineage bases lack it in ptx/tma.cuh), and the FP4 smem pack factor inlined into layout/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_layout put a group boundary inside a BLOCK_M tile 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 -O3 with 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 (DeepJIT CompilerOptions::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/DeepGEMM#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 align and revert only ceil_div to / 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 the align, which is load-bearing; the reviewer's actual suggestion stands. Landed as 3aca951 + 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-up 4b3571e (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 — 14 fence_view_async_shared() insertions across 8 SM120 kernel headers. The SM120 kernels read TMA-written smem through ldmatrix/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.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/447 **Author:** [@lucifer1004](https://github.com/lucifer1004) **Created:** 9/15/2026 **Status:** ✅ Merged **Merged:** 9/23/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `feature/sm120-main-api-migration` --- ### 📝 Commits (10+) - [`66081d4`](https://github.com/deepseek-ai/DeepGEMM/commit/66081d4c9c7d7c44f13fea402e5b622aa0f409c2) Public Release 26/09 (#432) - [`b945d55`](https://github.com/deepseek-ai/DeepGEMM/commit/b945d5523850961dacdc1d411ebaea1600176bdb) feat(sm120): add sparse MQA support on upstream main APIs - [`40b9d2d`](https://github.com/deepseek-ai/DeepGEMM/commit/40b9d2d8206439b93f8f606543066ed3b1df58f2) feat(sm120): restore native GEMM paths on main APIs - [`794783a`](https://github.com/deepseek-ai/DeepGEMM/commit/794783af91d1d81dd26d64ccceef23c7978a9885) feat(sm120): preserve nv_dev features alongside main API SM120 kernels - [`5500e69`](https://github.com/deepseek-ai/DeepGEMM/commit/5500e69410d0efba4daa30420cc35e12213374af) chore(sync): merge nv_dev ancestry into main API integration - [`3a04a69`](https://github.com/deepseek-ai/DeepGEMM/commit/3a04a69435187162c223fb0ec918b65a4f596dd0) fix(kernels): address migration review and restore grouped TMA pipeline - [`48e5de7`](https://github.com/deepseek-ai/DeepGEMM/commit/48e5de772758dec1a86a1e0248e494ae4bb6388d) fix(sm120): harden GEMM execution and complete HeadSplits output mapping - [`0aa7482`](https://github.com/deepseek-ai/DeepGEMM/commit/0aa74827f462d4c30d658873319e77f84bf1b454) test(sm120): collect SM120 tests into dedicated files - [`70a84f8`](https://github.com/deepseek-ai/DeepGEMM/commit/70a84f8adf2d05f90c3298e01e1b42c4ead7eac0) feat(sm120): vendor device layer from DeepGEMM-sm120 v0.1.3 - [`f7bea76`](https://github.com/deepseek-ai/DeepGEMM/commit/f7bea767527f4c526de1e5d628ce88d96fdc7e33) fix(sm120): gate ptxas -O2 for sparse MQA logits on CUDA 13.3+ ### 📊 Changes **193 files changed** (+22685 additions, -7023 deletions) <details> <summary>View changed files</summary> 📝 `.gitmodules` (+3 -3) 📝 `CMakeLists.txt` (+4 -2) 📝 `README.md` (+58 -61) 📝 `csrc/apis/attention.hpp` (+376 -128) 📝 `csrc/apis/config.hpp` (+17 -21) 📝 `csrc/apis/einsum.hpp` (+284 -90) 📝 `csrc/apis/gemm.hpp` (+438 -214) 📝 `csrc/apis/hyperconnection.hpp` (+52 -11) 📝 `csrc/apis/layout.hpp` (+38 -25) ➕ `csrc/apis/mega_gate.hpp` (+234 -0) ➕ `csrc/apis/mega_mhc.hpp` (+287 -0) 📝 `csrc/apis/mega_moe.hpp` (+42 -248) ➕ `csrc/apis/nvfp4_mega_moe.hpp` (+276 -0) ➕ `csrc/apis/sm120_gemm.hpp` (+86 -0) 📝 `csrc/apis/sm90_mega.hpp` (+26 -12) 📝 `csrc/indexing/main.cu` (+1 -1) ➖ `csrc/jit/cache.hpp` (+0 -31) ➖ `csrc/jit/compiler.hpp` (+0 -369) ➖ `csrc/jit/device_runtime.hpp` (+0 -153) ➖ `csrc/jit/handle.hpp` (+0 -222) _...and 80 more files_ </details> ### 📄 Description ## 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 - Native BF16, FP8, FP4 and mixed FP8/FP4 dense GEMM; M-grouped contiguous/masked paths; BF16/FP8 K-grouped paths. - FP8 einsum, native BF16 forward einsum and batch reduction, and HC prenorm. - Dense/paged FP8 and MXFP4 MQA, plus **DSv4.1 MXFP8/MXFP4 sparse and paged sparse MQA**, using the current metadata APIs. - Current scaling-factor/layout interfaces, dispatch validation, focused regression tests and support documentation. 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 - SM90 FP8 MegaMoE, including split L1/L2 execution and natural FP32 block-128 weight scales (#383). - NVFP4 MegaMoE with BF16 shared experts, per-expert scales and routed scaling (#409). - SiTU for both routed and shared experts (#396). - SM100 **dense** MQA with FP16 weights, preserving its dedicated MMA path rather than replacing it with BF16. FP16 paged MQA is not added. - SM90 paged MQA with physical page sizes 32/64, next_n=1/2/4, and zero-work/scheduler safety fixes. 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 - Restore removed math/shuffle helpers privately and adapt host launches to DeepJIT without replacing kernel arithmetic. - Validate SM90 output/device/rank inputs and reject incompatible buffer protocols/layouts. - Restrict NVFP4 `base=` reuse to identical supported configurations, including the configured SM count; main-compatible buffer reuse remains available. - Fix an **inherited NVFP4 combine double-buffer WAR hazard** with one warp synchronization before issuing a TMA overwrite. PC/source correlation and a one-token, non-graph reproduction established missing consumer completion before stage reuse. The pre-fix combine region matched nv_dev; the fix changes no arithmetic or rounding order. ## 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. - **SM120, final artifact:** 349 focused tests, 2 packing tests and 1 HC test passed. All 310 SM120-exclusive tests passed memcheck (0 errors) and racecheck (0 hazards/errors/warnings); the additional 39 focused cases were not sanitizer-tested. With GPUs hidden, 310 tests skipped without CUDA initialization. Intentional empty-graph PyTorch warnings are separate from sanitizer findings. - **H200:** 73 corrected paged MQA tests and host guards passed. SM90 MegaMoE smoke/stateful checks passed at 1/2/4/8 ranks, including unequal token counts, an empty rank, live/masked/live reuse, CUDA Graph replay and exact cumulative statistics. Larger 1/2/4-rank sweeps also passed. Selected original BF16/FP8 GEMM tests passed numerical checks and memcheck/racecheck. These runs preceded the final NVFP4-only synchronization change; the H200 code is unchanged. - **B300, single rank:** 54 attention tests, SiTU and main MegaMoE contract tests passed before the final NVFP4-only change. The final artifact passed 82 protocol tests, 27 NVFP4 matrix invocations and additional one-token controls. Coverage includes shared counts 0/1/2, alpha/A2 combinations, routed scaling, empty/masked/tail inputs, graph/state reuse and exact statistics. - **NVFP4 repair:** previously failing shared1/shared2 racechecks now report zero hazards/errors/warnings. Representative stateful memcheck passed. Actual ring wrap exercised 24,576 routed tokens against a 10,752-token ring; ordinary stateful execution and a separate single-large-call racecheck passed. ## Cautions - The extended 10-call large-ring stateful racecheck exceeded its 900-second execution bound without a final sanitizer summary. It is **not** counted as passed; the successful small graph/stateful and single-large-call checks are distinct coverage. - SM100 multirank/NVLink qualification remains outstanding; the available B300 allocation is single-GPU. - H200 MoE memcheck is not clean because NCCL 2.30.7 probes Blackwell-only symmetric kernels during initialization. A minimal NCCL-only program reproduces the handled CUDA API diagnostics without importing DeepGEMM. This is not evidence of a DeepGEMM memory-access bug, but no suppression or clean sanitizer claim is made. No demonstrated release/configuration fix was found in the inspected upstream versions. - The complete inherited main-release diff still needs upstream review. This is bounded kernel qualification, not a full-repository, installation, model-serving, SM121 or universal performance-parity claim. ## 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](https://github.com/lucifer1004/DeepGEMM-sm120) **v0.1.4**, the single source of truth shared with the vllm-project fork's SM120 layer (vllm-project/DeepGEMM#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 drop `block_scale`), `tensor_map_replace_global_dim_in_smem` relocated into `common/sm120_utils.cuh` (main-lineage bases lack it in `ptx/tma.cuh`), and the FP4 smem pack factor inlined into `layout/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_layout` put a group boundary inside a `BLOCK_M` tile 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 `-O3` with any `--register-usage-level` — a compiler bug, not a kernel defect (root-caused in [DeepGEMM-sm120#2](https://github.com/lucifer1004/DeepGEMM-sm120/issues/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 (DeepJIT `CompilerOptions::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/DeepGEMM#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 `align` and revert only `ceil_div` to `/` 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 the `align`, which *is* load-bearing; the reviewer's actual suggestion stands. Landed as `3aca951` + `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-up `4b3571e` (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](https://github.com/lucifer1004/DeepGEMM-sm120) **v0.1.5** (`b31a688`): ports #453 (thanks @Sunt-ing) and extends it to every site with the same hazard — 14 `fence_view_async_shared()` insertions across 8 SM120 kernel headers. The SM120 kernels read TMA-written smem through `ldmatrix`/`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. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:18:03 +00:00