DSA indexer forward: add a lean SM100 fast path for the head_dim=128 / qhead_per_kv_head=64 regime (up to 1.46x) - #416
Open
zkyue wants to merge 4 commits into
Open
Conversation
…28 MQA regime Add IndexerForwardLean / indexer_forward_lean_wrapper, an additive APIBase fast path next to IndexerForward, specialized for the head_dim=128, qhead_per_kv_head=64, h_kv=1 indexer-forward regime (BF16 Q/K, BF16 or FP32 W ingested directly, FP32 scores, batched uniform-length BSHD). The lean kernel (indexer_fwd_sm100_lean.py) replaces the generic tile scheduler with a static single-wave persistent grid in reversed-LPT tile order built around the uniform triangular causal work distribution, sweeping dense 128-row-aligned KV tiles via TMA. The schedule is computed on the host and the kernel compiles per (S_q, S_k, W dtype, sm_scale variant, device) by design, targeting training-style workloads where shapes repeat; the batch size and ratio never reach codegen, so one compiled kernel serves every B of a shape (execute() reads B from the runtime tensors). Compiled kernels, dispatch verdicts, and cached window tensors live in bounded thread-safe LRU caches (following graph.py's graph_cache bound), cached instances are immutable after compile (sm_scale is a per-call execute() argument), and the cached offset-free window tensors carry a CUDA event so consumers on other streams wait for their production. API-contract parity with the legacy wrapper is preserved: sm_scale and q_causal_offsets are fully supported (offsets are folded into per-row visibility windows on the host in int64 arithmetic; offset-free windows are cached per shape), scores are returned contiguous, and results match the legacy kernel to FP32 accumulation-order rounding with a deterministic run-to-run reduction order. check_support() validates the exact W shape, requires one common CUDA device across the runtime descriptors (meta descriptors are accepted for metadata-only support checks), resolves compute capability as a full (major, minor) pair plus sm_count from that target device, and requires a saturated persistent grid (num_m_tiles >= LEAN_MIN_WAVES * sm_count). Base-pointer TMA alignment (16 B, not implied by contiguity for storage-offset views) is validated at runtime in execute() and by the dispatcher. A configuration the gate accepts is expected to compile: JIT failures raise RuntimeError with context instead of silently falling back. The grid threshold was calibrated on B200 with nsys pure-kernel medians, where the lean kernel is 1.27-1.48x faster than the generic SM100 kernel across S_q 256-8192 at DeepSeek-V4-like shapes (H=64, D=128, causal ratio 2 and 4), 1.43x at S_q=8192, S_kv=4096. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
… the lean kernel Route indexer_forward_wrapper calls that fall inside the lean specialization (uniform-length BSHD, default tuning parameters, IndexerForwardLean.check_support() true) to the lean fast-path kernel on SM100. The dispatch is strictly additive: THD/varlen, qhead_per_kv_head=32, non-default tuning knobs, SM90/SM110+, unsaturated grids, and non-contiguous or non-16-byte-aligned tensors all keep using the legacy path unchanged, and malformed inputs fall through so the legacy error surface is preserved. The dispatch exception boundary is explicit: cheap non-throwing eligibility checks (rank, device, dtype, shape consistency, contiguity, alignment, finite sm_scale) run before any JIT, offsets metadata is screened before compilation, only validation ValueErrors fall back to legacy, and a lean compile failure for an accepted configuration raises RuntimeError rather than silently masking a real bug. Because the legacy kernel evaluates (q_causal_offset + i + 1) in int32 on device while the lean host path uses int64, offsets with offset + S_q + 1 > INT32_MAX stay on the legacy path so dispatched behavior is unchanged for extreme-but-valid offsets. Setting the CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN environment variable (to any non-empty value) forces the legacy kernel for every configuration as an escape hatch. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Add test_DSA_indexer_forward_lean.py covering the lean wrapper and
the transparent dispatch: score parity against an FP64 reference and
the legacy kernel, bitwise equivalence of wrapper-dispatched vs
directly-invoked lean calls, BF16- and FP32-W ingestion (bitwise
identical) plus run-to-run determinism, sm_scale and
q_causal_offsets handling, batched B>1 shapes with compile sharing
across batch sizes (a new B must not recompile), ragged S_k tails
whose visibility windows actually reach the physical tail columns
(S_k in {1, 127, 129, 999}, sweeping the 128-column KV-tile
boundary, via ratio=1 and via offsets), all-masked windows
(untouched -inf pre-fill), the int32 q_causal_offsets dispatch
bound (boundary-accepted offsets stay lean, one-past-the-bound
routes legacy), check_support() gate rejections isolated per cause
(H=32 at a saturating S_q, THD/varlen, wrong head counts/dims,
non-default tuning knobs on an otherwise eligible shape, grid and
tile-multiple limits derived from the device SM count), and the
CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN escape hatch.
Document the lean fast path and its support gate in
docs/fe-oss-apis/dsa.md — including the alignment requirement, the
compile-granularity and bounded-cache behavior, the int32 offsets
dispatch bound, and the loud-compile-failure policy — and link the
DSA page from the FE OSS API overview.
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a specialized SM100 lean Indexer Forward kernel, dispatch eligibility and fallback logic, public exports, documentation, caching, causal-window handling, and CUDA tests covering BSHD and THD correctness and routing. ChangesSM100 Lean Indexer Forward
Estimated code review effort: 5 (Critical) | ~90+ minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant indexer_forward_wrapper
participant IndexerForwardLean
participant IndexerForwardSm100Lean
Caller->>indexer_forward_wrapper: submit SM100 indexer-forward tensors
indexer_forward_wrapper->>IndexerForwardLean: check support and compile/cache
indexer_forward_wrapper->>IndexerForwardLean: execute with causal windows
IndexerForwardLean->>IndexerForwardSm100Lean: launch persistent kernel
IndexerForwardSm100Lean-->>IndexerForwardLean: produce FP32 scores
IndexerForwardLean-->>Caller: return scores
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.py (1)
419-419: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUnused loop variable
iin the MMA loop.The MMA warp drives
pipe_K/pipe_Spurely through pipeline state, soiis never referenced. Rename to_to clear the flagged RuffB007.♻️ Suggested change
- for i in cutlass.range(n_iters): + for _ in cutlass.range(n_iters):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.py` at line 419, In the MMA loop around `cutlass.range(n_iters)`, rename the unused loop variable `i` to `_`; preserve the loop body and pipeline-state behavior unchanged.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In
`@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.py`:
- Line 419: In the MMA loop around `cutlass.range(n_iters)`, rename the unused
loop variable `i` to `_`; preserve the loop body and pipeline-state behavior
unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cb3faaf6-d700-434f-aaef-78c7c86fccbf
📒 Files selected for processing (8)
docs/fe-oss-apis/dsa.mddocs/fe-oss-apis/overview.mdpython/cudnn/deepseek_sparse_attention/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_forward/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_forward/api.pypython/cudnn/deepseek_sparse_attention/indexer_forward/api_lean.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.pytest/python/fe_api/dsa/test_DSA_indexer_forward_lean.py
The lean H=64/D=128 score kernel is a pure per-row-window sweep with no batch/segment axis, so THD/varlen (ragged packed) support needs no kernel change: absolute compressed-KV windows ks=cu_seqlens_k[seg], ke=ks+visible encode segment isolation AND the ratio-causal mask together. This is exactly the gtp csa_host_index_math integer contract. api_lean additions (BSHD path byte-identical, guarded by cu_seqlens is None): - IndexerForwardLean gains cu_seqlens_q/k + max_seqlen_q/k; check_support early-branches to _check_support_thd (packed ranks, cu_seqlens validity/monotonicity/extent-consistency, H=64/D=128/MQA, T_q%4, saturated grid); execute early-branches to _execute_thd (one packed launch, B folds away). - _thd_ratio_causal_windows builds the absolute (T_q,) int32 windows; indexer_forward_lean_wrapper accepts cu_seqlens and returns global-column (T_q, m_total) scores via _indexer_forward_lean_thd. THD is exposed only through the explicit wrapper: the global-column layout differs from the legacy (total_q, max_seqlen_k) local-column output, so indexer_forward_wrapper still routes every THD call to the (unchanged) legacy kernel. q_causal_offsets on the THD path is gated to legacy for now. Verified on B200 (SM100): 42/42 lean tests pass (13 new THD cases); BSHD bitwise-identical to 04e2fec across 4 prod-ish configs; THD rel-L2 vs fp64 oracle 1.42e-7 <= legacy 1.45e-7 with exact segment isolation; nsys pure-kernel median lean-THD 57.3us vs legacy-THD 78.2us = 1.36x on a ragged B=3 (T_q=8192, ratio=4, m_total=2046) shape. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds an opt-in "lean" SM100 kernel for the DSA indexer forward (lightning-indexer scoring), plus a transparent dispatch hook in
indexer_forward_wrapper: eligible configs route to the lean kernel; everything else falls through to the existing kernel unchanged. Escape hatch:CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN.Why
On B200 (SM100), at DeepSeek-V4-like shapes (
head_dim=128,qhead_per_kv_head=64, bf16-in/fp32-out, ratio-causal), the existing indexer-forward spends ~2x the machine instructions of a leaner schedule for the same math (ncu: 41.6M vs 20.3Msm__inst_executed, ~80% FMA+ALU -- CLC per-tile address replay and the dual-stage epilogue account for it). A static single-wave persistent grid (min(sm_count x CTAs/SM, num_tiles)CTAs, runtime-queried, grid-stride, reversed-LPT order) with an N=256 query tile and register-resident head-sum removes that overhead in the grid-saturated regime:(nsys pure-kernel per-launch median, 50 warmup/100 reps, fp64-oracle-gated both kernels on identical inputs. Sweep across S_q 256-8192, ratio in {2,4}: 1.27-1.48x, no losing point measured. Calibration data is B200/148-SM; the wave gate is runtime-
sm_count-based.)Precision / determinism
Same contract as the existing kernel (bf16 tensor-core, fp32 accumulate, fp32 scores,
-infoutside the causal window). fp64-oracle rel-L2: lean 1.41e-7 vs existing 1.45e-7. Deterministic run-to-run (no atomics, fixed reduction order). Dispatched output is bitwise-identical to direct lean invocation; bf16-W ingest is bitwise-identical to fp32-W.Support surface (checked in
check_support; everything else -> existing kernel)SM100 only; uniform-length BSHD (any B);
head_dim=128,qhead_per_kv_head=64,h_kv=1; bf16 Q/K, bf16-or-fp32 W, fp32 out; contiguous, 16B-aligned base pointers;S_q % 4 == 0, anyS_kv >= 1;sm_scaleandq_causal_offsetssupported (offsets restricted to the int32-safe range); grid-saturation gatenum_m_tiles >= sm_count. THD/varlen (cu_seqlens) intentionally not routed -- the static grid cannot honor per-batch KV offsets; those stay on the existing path.Testing
test/python/fe_api/dsa/test_DSA_indexer_forward_lean.py: check_support accept/reject matrix, numerics vs reference (B in {1,2,3}, causal offsets incl. int32 boundary, sm_scale, fp32/bf16 W bitwise-equal, ragged tails S_kv in {1,127,129,999} with full-visibility offsets, all-masked windows, run-to-run determinism, dispatch-routing assertions incl. env hatch and non-default-knob fallback). Existing indexer-forward tests pass unchanged.black -l 160clean.Summary by CodeRabbit
New Features
indexer_forward.IndexerForwardLeanandindexer_forward_lean_wrapper) and updated theindexer_forwardmodule exports.Documentation
Tests