Conversation
ameynaik-hub
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 22, 2026 01:21
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdded SM90+ GDN ucache decode and verify-plus-flush kernels with bf16/fp16 variants, Python launch APIs, fp32-reference tests, strided-path and commit-guard regression coverage, and graph-captured CUDA benchmarks. ChangesGDN ucache execution
Estimated code review effort: 5 (Critical) | ~120 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant FlushAPI as gated_delta_rule_mtp_ucache_flush
participant Kernel as GdnDecodeUCacheFlushKernel
participant Ring as HistoryRing
participant State as CheckpointState
Caller->>FlushAPI: submit decode inputs and hist_len
FlushAPI->>Kernel: compile and launch with flush_min
Kernel->>Ring: read history and append entries
Kernel->>State: fold ring entries when flush threshold is reached
FlushAPI->>Ring: reset flushed history lengths
FlushAPI-->>Caller: return output tensor
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
ameynaik-hub
force-pushed
the
ameyn/gdn-ucache-replayssm-kernels
branch
from
July 22, 2026 01:26
61909af to
c7c55fa
Compare
…-only (SM100)
Two CuTe-DSL kernels that make speculative decoding cheap for GDN
(Gated DeltaNet) models, plus tests and a benchmark.
The problem: during speculative decoding, the engine must be able to
roll back rejected draft tokens. For GDN models the naive way is to
save a full copy of the recurrent state for every draft position — but
GDN states are megabytes per layer per request, so those copies eat
both memory capacity and bandwidth.
The idea: do not save states at all. Save only the small per-token
ingredients (update vector u, normalized key k, decay g) in a 16-slot
ring per request. One kernel launch per layer per decode step then does
everything: computes the verify output for the draft tokens, appends
the new ingredients to the ring, and — only when a request's ring is
full, roughly once every 4-5 steps — folds the ring into the single
checkpoint state and restarts it. Rejected tokens cost nothing (just a
cursor that does not advance).
Included:
- gdn_decode_bf16_wy_ucache_flush.py — the main fused verify+flush kernel
- gdn_decode_bf16_wy_ucache.py — a verify-only variant
- tests/gdn/test_decode_ucache.py — correctness vs a simple fp32 reference
- benchmarks/bench_gdn_ucache_flush.py — latency benchmark
Supported precisions (decay ring is always fp32):
mode | inputs | u/k rings | state pool
bf16 (default) | bf16 | bf16 | bf16
fp16-state (GDN_UCACHE_STATE_DTYPE) | bf16 | bf16 | fp16
fp16-IO (GDN_UCACHE_IO_DTYPE) | fp16 | fp16 | fp16
The fp16-state ("mixed") mode is the analogue of
mamba_ssm_cache_dtype=float16 in serving frameworks: only the
checkpoint gains fp16's extra mantissa bits; everything else stays bf16.
How to run:
pytest tests/gdn/test_decode_ucache.py -v
python benchmarks/bench_gdn_ucache_flush.py --arm bf16 --iters 1000
python benchmarks/bench_gdn_ucache_flush.py --arm fp16_state --iters 1000
Benchmark — us per layer per decode step, 1x B200, Qwen3.5-122B GDN
geometry at TP1 (H=16 key heads, HV=64 value heads, K=V=128, T=4 draft
window, ring W=16); columns = fraction of requests folding their ring
that step, scattered at exact counts (CUDA-graph replay, CUPTI cold-L2):
batch | 0% | 20% | 40% | 80%
8 | 13.2 | 17.1 | 17.8 | 18.2
32 | 26.3 | 31.1 | 34.1 | 40.4
64 | 43.7 | 51.1 | 58.2 | 70.4
128 | 73.2 | 88.4 | 102.2 | 126.6
256 | 133.0 | 162.4 | 188.5 | 242.6
Scope/limits: validated on B200 (SM100), head dims K=V=128, draft
windows T in {4, 8}. Dtype selection via env vars and a couple of API
stubs are known follow-ups.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py (2)
840-852: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead no-op expressions.
Lines 840, 844, 848, 852 are bare
cutlass.Int32(1)expressions whose results are discarded — likely leftover element-stride (sq_e/sk_e/…) assignments that were dropped. They are harmless (element stride 1 is never referenced), but read as dead code and can mislead. Consider removing them.🤖 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 `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py` around lines 840 - 852, Remove the standalone cutlass.Int32(1) expressions in the stride setup around sq_h, sk_h, and sv_hv, including the trailing occurrence, since their results are discarded. Leave the existing sq_*, sk_*, and sv_* stride assignments unchanged.
2388-2415: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd an explicit compute-capability gate on the public API.
The kernel requires SM90+ (TMA + mbarrier, per Line 2326) but nothing enforces it, so a mis-routed call on older hardware fails obscurely. Per coding guidelines, capability-dependent APIs should be gated; based on retrieved learnings, a single-backend, architecture-gated API should use
supported_compute_capability([...])rather thanbackend_requirement. PR objectives already list "an explicit SM100 gate" as a follow-up — worth wiring in here.🤖 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 `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py` around lines 2388 - 2415, Add an explicit compute-capability gate to the public gated_delta_rule_mtp_ucache API using supported_compute_capability([...]), rather than backend_requirement. Configure the gate for the kernel’s required SM90+ capability, so calls on older architectures are rejected before dispatch while supported hardware behavior remains unchanged.Sources: Coding guidelines, Learnings
🤖 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.
Inline comments:
In `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py`:
- Around line 3346-3347: Update the restart logic around restart_hist_on_flush
and hist_len so masked_fill_ always modifies the caller-visible tensor: require
hist_len to be contiguous before this operation, or apply the mask directly to
the original strided view instead of a contiguous temporary. Preserve resetting
entries where hist_len is at least flush_min.
In `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py`:
- Around line 2790-2796: Update flashinfer/__init__.py to import and re-export
gated_delta_rule_mtp_ucache from gdn_decode_bf16_wy_ucache, making the new
operation available through the package-level API while preserving the existing
module-level export.
- Around line 864-870: Update the stride setup near sa_t and sb_t to use _ab_t
for the b-view temporal and batch strides: set sb_t from _ab_t and derive sb_b
from _ab_rows multiplied by _ab_t. Preserve sb_hv as the unit stride and keep
the existing a-stride handling unchanged.
In `@tests/gdn/test_decode_ucache.py`:
- Around line 268-273: Apply Ruff formatting to all affected files: indent the
hist_len argument consistently in tests/gdn/test_decode_ucache.py (lines
268-273), and run ruff-format on
flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py (line 1) and
benchmarks/bench_gdn_ucache_flush.py (line 1). Run pre-commit across all files
and commit the resulting formatting changes.
---
Nitpick comments:
In `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py`:
- Around line 840-852: Remove the standalone cutlass.Int32(1) expressions in the
stride setup around sq_h, sk_h, and sv_hv, including the trailing occurrence,
since their results are discarded. Leave the existing sq_*, sk_*, and sv_*
stride assignments unchanged.
- Around line 2388-2415: Add an explicit compute-capability gate to the public
gated_delta_rule_mtp_ucache API using supported_compute_capability([...]),
rather than backend_requirement. Configure the gate for the kernel’s required
SM90+ capability, so calls on older architectures are rejected before dispatch
while supported hardware behavior remains unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 700fb765-0866-4960-a759-99c5435f0b7c
📥 Commits
Reviewing files that changed from the base of the PR and between 7f786be and 61909af25186ddef71f9a186271e7473a4e8a277.
📒 Files selected for processing (4)
benchmarks/bench_gdn_ucache_flush.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.pytests/gdn/test_decode_ucache.py
ameynaik-hub
force-pushed
the
ameyn/gdn-ucache-replayssm-kernels
branch
from
July 22, 2026 01:31
c7c55fa to
fe89ffd
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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.
Inline comments:
In `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py`:
- Around line 1087-1089: Update the warp-3 b-view stride setup near sb_hv, sb_t,
and sb_b so sb_t uses the computed _ab_t stride rather than hardcoded HV,
matching the a-view indexing and preserving correct behavior for packed or
strided a/b inputs.
- Around line 3356-3362: Update flashinfer/__init__.py to import and re-export
gated_delta_rule_mtp_ucache_flush from gdn_decode_bf16_wy_ucache_flush, making
it available as a package-level flashinfer API while preserving the existing
module-level export.
In `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py`:
- Line 1: Both affected kernel files are not formatted according to Ruff. Run
ruff format (or pre-commit run --all-files) on
flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py:1-1 and
flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py:1-1, then commit the
resulting formatting changes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ef0f8948-b62f-41a8-bb39-791718e77f42
📥 Commits
Reviewing files that changed from the base of the PR and between 61909af25186ddef71f9a186271e7473a4e8a277 and fe89ffd.
📒 Files selected for processing (4)
benchmarks/bench_gdn_ucache_flush.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.pytests/gdn/test_decode_ucache.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/gdn/test_decode_ucache.py
- benchmarks/bench_gdn_ucache_flush.py
…he flush kernel Adds the ring_fp16 and fp16_state_cache arms to the test + benchmark. Run accuracy (vs the fp32 reference, all 5 arms): pytest tests/gdn/test_decode_ucache.py -v # single arm: pytest tests/gdn/test_decode_ucache.py -v -k fp16_state_cache Run benchmark (CUDA-graph replay, CUPTI cold-L2; pass --arm to pick a mode): python benchmarks/bench_gdn_ucache_flush.py --arm bf16 python benchmarks/bench_gdn_ucache_flush.py --arm fp16_state_cache Supported precision modes (g_cache is fp32 in every mode): ┌─────┬───────────────────┬────────────────────────────────────┬────────────────┬───────────┬───────────┬──────────┐ │ # │ mode / arm name │ env knobs │ q/k/v + output │ SSM state │ u/k cache │ g (gate) │ ├─────┼───────────────────┼────────────────────────────────────┼────────────────┼───────────┼───────────┼──────────┤ │ 1 │ bf16 (default) │ (none) │ bf16 │ bf16 │ bf16 │ fp32 │ ├─────┼───────────────────┼────────────────────────────────────┼────────────────┼───────────┼───────────┼──────────┤ │ 2 │ fp16_io │ IO_DTYPE=fp16 │ fp16 │ fp16 │ fp16 │ fp32 │ ├─────┼───────────────────┼────────────────────────────────────┼────────────────┼───────────┼───────────┼──────────┤ │ 3 │ fp16_state │ STATE_DTYPE=fp16 │ bf16 │ fp16 │ bf16 │ fp32 │ ├─────┼───────────────────┼────────────────────────────────────┼────────────────┼───────────┼───────────┼──────────┤ │ 4 │ ring_fp16 (cache) │ RING_DTYPE=fp16 │ bf16 │ bf16 │ fp16 │ fp32 │ ├─────┼───────────────────┼────────────────────────────────────┼────────────────┼───────────┼───────────┼──────────┤ │ 5 │ fp16_state_cache │ STATE_DTYPE=fp16 + RING_DTYPE=fp16 │ bf16 │ fp16 │ fp16 │ fp32 │ └─────┴───────────────────┴────────────────────────────────────┴────────────────┴───────────┴───────────┴──────────┘ (env knobs carry the GDN_UCACHE_ prefix.)
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
benchmarks/bench_gdn_ucache_flush.py (1)
145-153: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReset every mutated kernel input before each timed replay.
The flush kernel updates
pooland restarts/appends ring state, but Line 150 restores onlyhist_len. Subsequent graph replays therefore benchmark evolved checkpoint/cache contents rather than the same workload. Capture a reset graph from pristine device-resident copies forpool,kc,uc, andgc, replay it before each run graph, and time only the run graph.Based on learnings, in-place benchmarks need separate reset and timed replay graphs.
🤖 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 `@benchmarks/bench_gdn_ucache_flush.py` around lines 145 - 153, The benchmark’s timed replay currently restores only hist_len, so mutated pool, kc, uc, and gc state carries across iterations. Update fn and the surrounding graphed/bench_gpu_time setup to create a reset graph from pristine device-resident copies, replay that reset graph before each timed run graph, and ensure only the run graph is included in timing.Source: Learnings
🤖 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.
Inline comments:
In `@benchmarks/bench_gdn_ucache_flush.py`:
- Around line 60-66: Update the --arm CLI help text near the argument definition
to list all entries in ARMS, including ring_fp16 and fp16_state_cache alongside
the existing configurations. Keep the documented names synchronized with the
supported ARMS keys.
---
Outside diff comments:
In `@benchmarks/bench_gdn_ucache_flush.py`:
- Around line 145-153: The benchmark’s timed replay currently restores only
hist_len, so mutated pool, kc, uc, and gc state carries across iterations.
Update fn and the surrounding graphed/bench_gpu_time setup to create a reset
graph from pristine device-resident copies, replay that reset graph before each
timed run graph, and ensure only the run graph is included in timing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 950c1009-0212-447d-93a4-937ad17ef147
📒 Files selected for processing (3)
benchmarks/bench_gdn_ucache_flush.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.pytests/gdn/test_decode_ucache.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/gdn/test_decode_ucache.py
- flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py
Comment on lines
+60
to
+66
| ARMS = { | ||
| "bf16": (None, None, None, torch.bfloat16, torch.bfloat16, torch.bfloat16), | ||
| "fp16_state": (None, "fp16", None, torch.bfloat16, torch.float16, torch.bfloat16), | ||
| "fp16_io": ("fp16", None, None, torch.float16, torch.float16, torch.float16), | ||
| "ring_fp16": (None, None, "fp16", torch.bfloat16, torch.bfloat16, torch.float16), | ||
| "fp16_state_cache": (None, "fp16", "fp16", | ||
| torch.bfloat16, torch.float16, torch.float16), |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document all supported dtype arms in the CLI help.
--arm accepts ring_fp16 and fp16_state_cache, but Line 165 still lists only three configurations.
Proposed fix
- help="dtype config: bf16 | fp16_state | fp16_io")
+ help=("dtype config: bf16 | fp16_state | fp16_io | ring_fp16 | "
+ "fp16_state_cache"))As per coding guidelines, keep documentation in sync with code changes.
🤖 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 `@benchmarks/bench_gdn_ucache_flush.py` around lines 60 - 66, Update the --arm
CLI help text near the argument definition to list all entries in ARMS,
including ring_fp16 and fp16_state_cache alongside the existing configurations.
Keep the documented names synchronized with the supported ARMS keys.
Source: Coding guidelines
Contributor
Author
The u-cache decode kernels (verify + flush) computed b's per-token and per-batch global-memory strides as HV, while a correctly used _ab_t. The wrapper only takes the native-a/b path when tuple(b.stride()) == tuple(a.stride()), so b carries the SAME non-HV token stride as a on the packed/chunk-view path (vLLM's strided-QKV setup, where a/b are HV-wide slices of a wider tensor). With sb_t=HV the packed b was read from the wrong rows -> wrong gate/beta -> wrong output. _ab_t == HV in the compact case, so the default/compact path is byte-identical. Verified on B200: a strided (chunk-view) a/b run vs a compact run with the SAME values went from max|y|=2.2e-2 (corrupted) to 0.0 (bit-identical). pytest tests/gdn/test_decode_ucache.py: 21/21 unchanged, plus a new test_strided_ab_matches_compact guard (fails on the old code, passes now).
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/gdn/test_decode_ucache.py (1)
302-303: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRun
ruff formatbefore re-submitting.
ruff format --check tests/gdn/test_decode_ucache.pyreports the file would be reformatted, so the repository will flag this as a formatting failure.🤖 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 `@tests/gdn/test_decode_ucache.py` around lines 302 - 303, Run ruff format on tests/gdn/test_decode_ucache.py, ensuring the importlib.util.spec_from_file_location call and surrounding file match the repository’s formatting requirements.Source: Coding guidelines
🤖 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 `@tests/gdn/test_decode_ucache.py`:
- Around line 302-303: Run ruff format on tests/gdn/test_decode_ucache.py,
ensuring the importlib.util.spec_from_file_location call and surrounding file
match the repository’s formatting requirements.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: d0b6c776-40f2-466d-b19b-ff782998bea6
📒 Files selected for processing (3)
flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.pytests/gdn/test_decode_ucache.py
🚧 Files skipped from review as they are similar to previous changes (2)
- flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py
- flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py
…st the live window; racing tail restart deleted Live window [base, base+P) mod 32 with caller-owned cursor commits (commit_gdn_replayssm_spec semantics). History loads ring-rotate at load time into logically-ordered smem tiles, so every downstream tile/GEMM/ mask is unchanged (order is carried by per-row g). Appends (k leader, u, g) land at (base+P+s)&31 for flush and verify rows alike — past every sibling's fold-source window — eliminating the same-launch inter-CTA K-ring race that the old single-buffer restart mitigated by launch-skew timing only. New per-request cache_base input; W_RING=16 stays as the max-window/tile constant; RING_SLOTS/RING_MASK exported.
…. wrapped windows, logical-gather oracle, flush no-overwrite property test
…ore, --base flag for wrapped-window measurement
The wrapper's standalone convenience cursor-commit (restart_hist_on_flush=True) plus the bench's own per-iteration cursor restores add ~4-6us of elementwise ops to the timed graph -- fine for wrapper-level A/Bs, misleading for kernel-level ones (vLLM serving passes False and owns commits in the builder). --no-commit disables both; legal since the kernel never mutates cursors.
…mmit path Review findings 1-4 on the ring conversion. All host-side wrapper guards; the device kernel is untouched (pure-kernel bench at B=256 matches the documented anchors, wrapped window == base 0). - restart_hist_on_flush=True now REQUIRES a caller-owned cache_base: the in-place commit previously slid the base of an internal zeros_like temp, silently corrupting any legacy stateful caller after its first flush (the next call re-defaulted base to 0 and read a stale window). - restart_hist_on_flush=True asserts hist_len/cache_base contiguity: .contiguous() on a strided cursor view silently copies, and the commit would mutate the copy (base never advances -> cursor desync). - The standalone path validates hist_len in [0,16] and cache_base in [0,32), skipped during CUDA-graph capture (the eager warmups validate the same tensors). The ring's &31 masking makes oversized windows corrupt SILENTLY: hist_len>16 breaks the W_RING=16 tile math, and hist_len>=29 wraps appends back into the live window -- the write-after-read race the ring removes. The caller-owned-commit serving path (vLLM) stays sync-free and unvalidated as before. - verify-only kernel: the 16-deep shape assert now spells out the ring format incompatibility with the flush kernel (do not share one ring allocation between the kernels; do not pad these rings to 32). tests: +test_standalone_commit_guards_reject_silent_corruption_patterns (suite: 43/43).
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/gdn/test_decode_ucache.py (1)
347-360: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRing rows passed to the oracle without
_logical_rings.This is the only remaining call site that hands physical ring rows straight to
_ref_fp32. It happens to be correct because_make_caseis called withoutbases(all zeros,P=13 < RING), but it silently depends on that. Passing through_logical_ringskeeps the test robust if the case ever gains non-zero bases.♻️ Suggested consistency fix
for r in range(B): + kc_l, uc_l, gc_l = _logical_rings(kc[r], uc[r], gc[r], 0) _, S_ref = _ref_fp32(q[r], k[r], v[r], a[r], b[r], A_log, - dt_bias, pool_before[r], kc[r], uc[r], - gc[r], 13) + dt_bias, pool_before[r], kc_l, uc_l, + gc_l, 13)🤖 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 `@tests/gdn/test_decode_ucache.py` around lines 347 - 360, Update the `_ref_fp32` call in the `gated_delta_rule_mtp_ucache_flush` test to pass the ring rows transformed through `_logical_rings`, matching the other oracle call sites. Preserve the existing test setup and expected-reference calculations while ensuring non-zero ring bases are handled correctly.
🤖 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.
Inline comments:
In `@tests/gdn/test_decode_ucache.py`:
- Line 479: Update the test statements around hl_bad in the relevant decode
ucache test so each semicolon-joined assignment is placed on its own line,
including the matching case at the other flagged location, while preserving the
existing test behavior.
---
Nitpick comments:
In `@tests/gdn/test_decode_ucache.py`:
- Around line 347-360: Update the `_ref_fp32` call in the
`gated_delta_rule_mtp_ucache_flush` test to pass the ring rows transformed
through `_logical_rings`, matching the other oracle call sites. Preserve the
existing test setup and expected-reference calculations while ensuring non-zero
ring bases are handled correctly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: fcdf84ff-5c67-4944-bb73-8d33db672515
📒 Files selected for processing (4)
benchmarks/bench_gdn_ucache_flush.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.pytests/gdn/test_decode_ucache.py
🚧 Files skipped from review as they are similar to previous changes (3)
- benchmarks/bench_gdn_ucache_flush.py
- flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py
- flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache.py
4 tasks
ameynaik-hub
added a commit
to ameynaik-hub/flashinfer
that referenced
this pull request
Jul 28, 2026
…tate mode Ring semantics (parity with the u-cache kernels, PR flashinfer-ai#4081): the four raw rings become 32-slot circular buffers; the live window of request b is [cache_base[b], cache_base[b]+hist_len[b]) mod 32 and the T new entries are appended at (cache_base+hist_len+s) & 31 — PAST the window, for flush and verify rows alike, so an append can never overwrite rows a sibling CTA still reads. The racing single-buffer tail restart is deleted, which also lets the shared k ring append move IN-KERNEL (from SMEM, pre-norm raw k) and removes the k-append micro-kernel launch entirely. Cursor commits are caller-owned (flush: base' = (base+len) & 31, len' = accepted; verify: len' += accepted); restart_hist_on_flush=True applies the commit in the wrapper for standalone use (graph-capturable), with cursor-safety validation (caller-owned contiguous cache_base required; range checks skipped during graph capture). fp16 SSM-state mode (PR flashinfer-ai#4081's fp16_state): GDN_VCACHE_STATE_DTYPE=fp16 stores the checkpoint fp16 while q/k/v/a/b, the rings, and the output stay bf16. State-touching MMAs (H GEMM + Step-A KH piggyback) convert their four shared A-fragments bf16->f16 in registers (exact in range) and issue .f16.f16; the Step-E fold unpacks fp16 state pairs through f32 FMAs and repacks fp16. Default bf16 path is bit-identical to before. Validated on B200 in BOTH state modes: correctness sweep vs the fp32 per-request oracle incl. wrapped windows (base near 32), ring property checks (append placement, flush no-overwrite, cursor commits), verify-only semantics, CUDA-graph replay bit-equality, and a new test_vcache_fp16_state_commits_more_precisely_than_bf16. Perf, 33% flush, B=8/64/256: 14.3/63.1/215.0 us (bf16; was 17.3/67.9/217.4 — the micro-kernel launch is gone) vs u-cache 13.0/53.0/178.2; fp16 state 220.6 us at B=256 (~2.6% for the A-fragment cvts). Wrapped windows are free (--base 28: 214.9). Run additions: GDN_VCACHE_STATE_DTYPE=fp16 pytest tests/gdn/test_decode_vcache.py -v GDN_VCACHE_STATE_DTYPE=fp16 python benchmarks/bench_gdn_vcache_flush.py python benchmarks/bench_gdn_vcache_flush.py --base 28 # wrapped windows
kahyunnam
reviewed
Jul 30, 2026
kahyunnam
left a comment
kahyunnam
left a comment
Member
There was a problem hiding this comment.
The pre-commit is red (https://github.com/flashinfer-ai/flashinfer/actions/runs/30186413244/job/89751881862?pr=4081) + we probably need to export this from https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/gdn_kernels/__init__.py#L94
Comment on lines
+68
to
+70
| k_cache [pool, H, 16, K] IO in/out ring: L2-normalized keys | ||
| u_cache [pool, HV, 16, V] IO in/out ring: correction vectors | ||
| g_cache [pool, HV, 16] f32 in/out ring: cumulative log-decay |
Contributor
Author
There was a problem hiding this comment.
during one kernel launch, two things happen to the same request's cache at the same time:
(1) if the row is flushing, the kernel is reading the history rows (up to 16 of them) to fold them into the state checkpoint, and
(2) the kernel is WRITING the current step's new tokens (T=4) into the cache.
If the buffer were exactly 16 slots, there'd be nowhere to put the new tokens except by wrapping around to the front. Example: history is 13 rows (the flush threshold), so the new 4 tokens go to positions 13, 14, 15, 16 but 16 wraps to slot 0 on a 16-slot ring, which is a history row the flush is still reading. That's a read/write race, and there's no synchronization between the thread blocks to prevent it.
So the physical ring is bigger than the logical window: new tokens are always written PAST the live window into fresh slots, never on top of it. Worst case is a full 16-row window plus up to 8 new tokens = 24 slots needed at once → round up to the next power of two = 32.
it also makes rollback trivial. The new tokens sit outside the committed window, and a cursor advances by the number of accepted tokens after the launch if some tokens are rejected, nothing in the live history was ever touched.
| @@ -0,0 +1,499 @@ | |||
| """ | |||
Member
There was a problem hiding this comment.
all cases load the flush module; verify-only has no oracle coverage, can we have a thin test for verify-only?
| V_DIM_C = 128 # full V tile per CTA | ||
| BK_H = 16 # K-tile for the H GEMM (multiple of 16 for mma.k=16) | ||
| EPS = 1e-6 | ||
| io = cutlass.BFloat16 |
Member
There was a problem hiding this comment.
let's assert q/k/v/a/b are bf16 (or convert), and/or add dtype to the verify cache key once body can specialize.
Context: trying to avoid GDN-C1 and GDN-C3 items from #4214
- tests: thin fp32-oracle test for the verify-only kernel (+ negative dtype test); pins read-only pool and ring-append contracts - verify wrapper: assert q/k/v/a/b bf16 + io dtype in compile cache key (flashinfer-ai#4214 GDN-C1/C3); parity with the flush wrapper - export both ucache entry points from gdn_kernels/__init__ - ruff-format the PR files + two SIM300 fixes (pre-commit green)
Contributor
Author
|
/bot run tests/gdn |
Collaborator
Collaborator
|
[FAILED] Pipeline #61519802 — 16/18 executed test jobs passed Compared with nightly #61367193. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsPR-related regressions
New relative to nightly (attribution uncertain)
|
kahyunnam
approved these changes
Aug 11, 2026
kahyunnam
enabled auto-merge (squash)
August 11, 2026 17:23
Vinnie6167
added a commit
to Vinnie6167/flashinfer
that referenced
this pull request
Aug 21, 2026
… target the GPU" Reverts the skip guard in tests/gdn/test_decode_ucache.py. The analysis behind it still holds -- those kernels compile through @cute.experimental.jit, which resolves the device arch inside the DSL, so `KeyError: 'sm_107a'` on DSL 4.7 has no flashinfer-side site to intercept and CUTE_DSL_ARCH=sm_100f does not reach that path -- but the 45 failures are being left visible for now rather than skipped. Keeping the PR to the dispatch-layer changes also keeps it out of a file owned by the kernel author (flashinfer-ai#4081), who is better placed to decide whether the requirement is inherent or the kernel can be made family-portable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
bkryu
pushed a commit
that referenced
this pull request
Aug 27, 2026
…4649) ## Problem On SM107 (Rubin) with a CuTe DSL older than 4.8, FlashInfer fails with a bare `KeyError: 'sm_107a'` raised from `enum.py` inside `cute.compile` — no FlashInfer frame in the traceback, and no warning. The root cause is that `supported_compute_capability` gates on **hardware capability alone**. #4122 widened the cute-dsl lists to `[100, 103, 107]`, which tells the dispatcher Rubin is supported regardless of which DSL is installed. Public CuTe DSL tops out at 4.7.0 on PyPI, and that release has no `sm_107a` in its `Arch` enum. This is reachable without asking for cute-dsl explicitly — two `auto` heuristics route SM107 *toward* it: - `_heuristic_func_mm_fp4`: `elif is_sm107: candidate_backends = ("cudnn", "cutlass", "cute-dsl")` - `_heuristic_func_bmm_fp8`: appends `"cute-dsl_sm107"` when `is_sm107_supported` ## Changes **1. Decline the cute-dsl backend when the installed DSL cannot emit for the device** (`fix(gemm)`) The three cute-dsl requirement functions now call `_check_cute_dsl_arch(...)`, which sits beside the existing `_check_cute_dsl_availability()` and delegates to `require_cute_dsl_arch()` — the helper added in #4122, which owns both the predicate and the message (it derives the family arch and names the exact `CUTE_DSL_ARCH=sm_100f` to export). Only the exception type is adapted, and that part is load bearing: `require_cute_dsl_arch` raises `NotImplementedError`, while `suitable_auto_backends` catches `ValueError` to mean "backend not suitable" and keeps searching. Left unadapted, an unsupported DSL would propagate out of the auto path and fail the call instead of falling back to cutlass/cudnn. Returning `False` instead of raising was also rejected: on the explicit-backend path that surfaces as `ValueError: Problem size is not supported`, which is misleading. **No capability lists change.** This is deliberately an *availability* check, not a capability one. The kernels do exist for sm_107, so `@supported_compute_capability([100, 103, 107])` stays as-is and `is_backend_supported("cute-dsl", 107)` keeps answering `True` — it is a public method on the wrapper, called with no tensors by e.g. `flashinfer/trace/templates/gemm.py:707`, and making it vary with an installed pip package would also have made the skip reason in `tests/grouped_mm/conftest.py` environment-dependent. This mirrors how the codebase already separates the two axes: `_cudnn_mm_mxfp8_requirement` lists its capabilities statically while `CUDNN_AVAILABLE` handles presence, and `_is_cudnn_override_shape_available` handles a dependency that is present but too old. **2. GDN CP delta rule resolves the arch instead of formatting it** (`fix(gdn)`) `_blackwell_compile_options` guards on the major only, then builds `f"sm_{major}{minor}a"`. Rubin is 10.7, so it passes a check written when "compute 10.x" meant Blackwell 10.0/10.3. This is the only place FlashInfer names the arch for a compute-10 device; everywhere else the DSL derives it internally. `cute_dsl_compile_arch()` returns the device's own arch when the DSL has it, the family arch when the DSL is targeting `sm_100f`, and otherwise raises `NotImplementedError` naming `CUTE_DSL_ARCH`. Same rule as the capability gate, so the two cannot disagree. ## Testing Rubin CI, `TEST_PATH="tests/gemm tests/gdn"`, against `release-v0.6.18`, with `CUTE_DSL_ARCH=sm_100f` exported and public CuTe DSL 4.7.0: | | before | after | |---|---|---| | passed | 8,141 | **12,341** | | failed | 4,230 | **1** | | `KeyError: 'sm_107a'` | 4,482 | **0** | Identical results on **both** VR200 (`hecate`, 4 workers, 2,078s) and GR100 (8 workers, 3,424s); `suite_complete=true` on both, well inside the 13,500s deadline. Per-file, verified independently on both boards: | File | before | after | |---|---|---| | `tests/gdn/test_prefill_delta_rule.py` | 2,678 | **0** | | `tests/gemm/test_mm_mxfp8.py` | 501 | **0** | | `tests/gdn/test_decode_delta_rule.py` | 417 | **0** | | `tests/gdn/test_prefill_cp_delta_rule.py` | 232 | **0** | The remaining failures are `tests/gdn/test_decode_ucache.py` and `tests/gemm/test_bmm_fp8.py` — see below. `BackendSupportedError` count is **0**, so the cute-dsl backends are being selected and compiling successfully against `sm_100f` — not silently skipped. The node accounting reconciles exactly: the plan drops 44,275 → 44,229 nodes and 37 → 36 units, i.e. the 46 tests in the ucache module leave collection entirely (a module-level skip is taken during collection, so those nodes are not counted as `skipped`). `passed` moves +287 = +288 Triton tests now compiling, −1 ucache test that previously passed; `failed` moves −333 = −288 Triton −45 ucache. Also unit-tested away from hardware: the new decorator resolves conditional 107 as False on DSL 4.7, True on 4.8+/`CUTE_DSL_ARCH`, False when the predicate raises, and yields a plain `set` when no conditional is given. `cute_dsl_compile_arch` was verified against a stubbed `Arch` enum for native / family / unsupported / Blackwell-unchanged, and the skip predicate for all four DSL-vs-arch combinations. ### Caveats - **The numbers above do not reflect this branch.** They were measured at `93143db2`, which carried a skip guard for `tests/gdn/test_decode_ucache.py` that has since been reverted, so 45 of those tests now fail again rather than skipping. - **The gemm mechanism changed after that measurement.** The two commits after it moved the check out of the decorator and into the requirement functions; that mechanism is unit-tested (adapter pass-through, `NotImplementedError` → `ValueError`, silent when the probe cannot be imported) but has not been re-run on hardware. - The measurement runs also carry `CUTE_DSL_ARCH=sm_100f` from the CI side. With it set the DSL *can* target sm_107, so `_check_cute_dsl_arch` passes and the gemm change is a no-op; only a run without that variable exercises the deselect-and-fall-back path. - `cute_dsl_compile_arch` changes `gdn_cp_prefill.py` for **all** compute-10 devices, not just Rubin. Blackwell resolution (`sm_100a` / `sm_103a`) is verified against a stubbed `Arch` enum, not on B200/GB200 hardware. ## Not addressed - **45 `KeyError: 'sm_107a'`** in `tests/gdn/test_decode_ucache.py`. Not fixable from FlashInfer: those kernels compile through `@cute.experimental.jit` / `@cute.experimental.kernel`, passing no arch and no compile options, so the DSL resolves the device arch itself and looks up `sm_107a` in its own enum. There is no FlashInfer-side site to guard, the traceback bottoms out at `enum.py:813` with no FlashInfer frame, and `CUTE_DSL_ARCH=sm_100f` does not help because that path never consults it — which points at a genuine **CuTe DSL 4.8** requirement. Left visible rather than skipped; the kernel author (#4081) is better placed to say whether it is inherent. - **1 `No valid cute-dsl SM107 bmm_fp8 config`** in `tests/gemm/test_bmm_fp8.py` — pre-existing, and present on the internal DSL 4.8 stack too (18 vs 20 occurrences across stacks), so it is independent of the DSL version question. The 288 Triton `PTXASError` failures previously seen in `tests/gemm/test_group_gemm.py` were a CI-side issue, not a FlashInfer one: Triton resolves ptxas through its own knobs (`TRITON_PTXAS_PATH`, and `TRITON_PTXAS_BLACKWELL_PATH` for arch >= 100, which is the one Rubin selects) and otherwise falls back to `$CUDA_HOME/bin/ptxas`. Fixed in flashinfer-ci!354; this run confirms 0 remaining. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added architecture detection for CuTe DSL compilation, including native and family-compatible GPU architectures. * Added clear guidance when the installed DSL cannot compile for a target GPU. * **Bug Fixes** * Improved Blackwell architecture handling, including support for devices with nonstandard architecture identifiers. * Prevented unsuitable CuTe DSL backends from being selected automatically when architecture support is unavailable. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
PetersonGuo
pushed a commit
to PetersonGuo/flashinfer
that referenced
this pull request
Aug 28, 2026
…lashinfer-ai#4649) ## Problem On SM107 (Rubin) with a CuTe DSL older than 4.8, FlashInfer fails with a bare `KeyError: 'sm_107a'` raised from `enum.py` inside `cute.compile` — no FlashInfer frame in the traceback, and no warning. The root cause is that `supported_compute_capability` gates on **hardware capability alone**. flashinfer-ai#4122 widened the cute-dsl lists to `[100, 103, 107]`, which tells the dispatcher Rubin is supported regardless of which DSL is installed. Public CuTe DSL tops out at 4.7.0 on PyPI, and that release has no `sm_107a` in its `Arch` enum. This is reachable without asking for cute-dsl explicitly — two `auto` heuristics route SM107 *toward* it: - `_heuristic_func_mm_fp4`: `elif is_sm107: candidate_backends = ("cudnn", "cutlass", "cute-dsl")` - `_heuristic_func_bmm_fp8`: appends `"cute-dsl_sm107"` when `is_sm107_supported` ## Changes **1. Decline the cute-dsl backend when the installed DSL cannot emit for the device** (`fix(gemm)`) The three cute-dsl requirement functions now call `_check_cute_dsl_arch(...)`, which sits beside the existing `_check_cute_dsl_availability()` and delegates to `require_cute_dsl_arch()` — the helper added in flashinfer-ai#4122, which owns both the predicate and the message (it derives the family arch and names the exact `CUTE_DSL_ARCH=sm_100f` to export). Only the exception type is adapted, and that part is load bearing: `require_cute_dsl_arch` raises `NotImplementedError`, while `suitable_auto_backends` catches `ValueError` to mean "backend not suitable" and keeps searching. Left unadapted, an unsupported DSL would propagate out of the auto path and fail the call instead of falling back to cutlass/cudnn. Returning `False` instead of raising was also rejected: on the explicit-backend path that surfaces as `ValueError: Problem size is not supported`, which is misleading. **No capability lists change.** This is deliberately an *availability* check, not a capability one. The kernels do exist for sm_107, so `@supported_compute_capability([100, 103, 107])` stays as-is and `is_backend_supported("cute-dsl", 107)` keeps answering `True` — it is a public method on the wrapper, called with no tensors by e.g. `flashinfer/trace/templates/gemm.py:707`, and making it vary with an installed pip package would also have made the skip reason in `tests/grouped_mm/conftest.py` environment-dependent. This mirrors how the codebase already separates the two axes: `_cudnn_mm_mxfp8_requirement` lists its capabilities statically while `CUDNN_AVAILABLE` handles presence, and `_is_cudnn_override_shape_available` handles a dependency that is present but too old. **2. GDN CP delta rule resolves the arch instead of formatting it** (`fix(gdn)`) `_blackwell_compile_options` guards on the major only, then builds `f"sm_{major}{minor}a"`. Rubin is 10.7, so it passes a check written when "compute 10.x" meant Blackwell 10.0/10.3. This is the only place FlashInfer names the arch for a compute-10 device; everywhere else the DSL derives it internally. `cute_dsl_compile_arch()` returns the device's own arch when the DSL has it, the family arch when the DSL is targeting `sm_100f`, and otherwise raises `NotImplementedError` naming `CUTE_DSL_ARCH`. Same rule as the capability gate, so the two cannot disagree. ## Testing Rubin CI, `TEST_PATH="tests/gemm tests/gdn"`, against `release-v0.6.18`, with `CUTE_DSL_ARCH=sm_100f` exported and public CuTe DSL 4.7.0: | | before | after | |---|---|---| | passed | 8,141 | **12,341** | | failed | 4,230 | **1** | | `KeyError: 'sm_107a'` | 4,482 | **0** | Identical results on **both** VR200 (`hecate`, 4 workers, 2,078s) and GR100 (8 workers, 3,424s); `suite_complete=true` on both, well inside the 13,500s deadline. Per-file, verified independently on both boards: | File | before | after | |---|---|---| | `tests/gdn/test_prefill_delta_rule.py` | 2,678 | **0** | | `tests/gemm/test_mm_mxfp8.py` | 501 | **0** | | `tests/gdn/test_decode_delta_rule.py` | 417 | **0** | | `tests/gdn/test_prefill_cp_delta_rule.py` | 232 | **0** | The remaining failures are `tests/gdn/test_decode_ucache.py` and `tests/gemm/test_bmm_fp8.py` — see below. `BackendSupportedError` count is **0**, so the cute-dsl backends are being selected and compiling successfully against `sm_100f` — not silently skipped. The node accounting reconciles exactly: the plan drops 44,275 → 44,229 nodes and 37 → 36 units, i.e. the 46 tests in the ucache module leave collection entirely (a module-level skip is taken during collection, so those nodes are not counted as `skipped`). `passed` moves +287 = +288 Triton tests now compiling, −1 ucache test that previously passed; `failed` moves −333 = −288 Triton −45 ucache. Also unit-tested away from hardware: the new decorator resolves conditional 107 as False on DSL 4.7, True on 4.8+/`CUTE_DSL_ARCH`, False when the predicate raises, and yields a plain `set` when no conditional is given. `cute_dsl_compile_arch` was verified against a stubbed `Arch` enum for native / family / unsupported / Blackwell-unchanged, and the skip predicate for all four DSL-vs-arch combinations. ### Caveats - **The numbers above do not reflect this branch.** They were measured at `93143db2`, which carried a skip guard for `tests/gdn/test_decode_ucache.py` that has since been reverted, so 45 of those tests now fail again rather than skipping. - **The gemm mechanism changed after that measurement.** The two commits after it moved the check out of the decorator and into the requirement functions; that mechanism is unit-tested (adapter pass-through, `NotImplementedError` → `ValueError`, silent when the probe cannot be imported) but has not been re-run on hardware. - The measurement runs also carry `CUTE_DSL_ARCH=sm_100f` from the CI side. With it set the DSL *can* target sm_107, so `_check_cute_dsl_arch` passes and the gemm change is a no-op; only a run without that variable exercises the deselect-and-fall-back path. - `cute_dsl_compile_arch` changes `gdn_cp_prefill.py` for **all** compute-10 devices, not just Rubin. Blackwell resolution (`sm_100a` / `sm_103a`) is verified against a stubbed `Arch` enum, not on B200/GB200 hardware. ## Not addressed - **45 `KeyError: 'sm_107a'`** in `tests/gdn/test_decode_ucache.py`. Not fixable from FlashInfer: those kernels compile through `@cute.experimental.jit` / `@cute.experimental.kernel`, passing no arch and no compile options, so the DSL resolves the device arch itself and looks up `sm_107a` in its own enum. There is no FlashInfer-side site to guard, the traceback bottoms out at `enum.py:813` with no FlashInfer frame, and `CUTE_DSL_ARCH=sm_100f` does not help because that path never consults it — which points at a genuine **CuTe DSL 4.8** requirement. Left visible rather than skipped; the kernel author (flashinfer-ai#4081) is better placed to say whether it is inherent. - **1 `No valid cute-dsl SM107 bmm_fp8 config`** in `tests/gemm/test_bmm_fp8.py` — pre-existing, and present on the internal DSL 4.8 stack too (18 vs 20 occurrences across stacks), so it is independent of the DSL version question. The 288 Triton `PTXASError` failures previously seen in `tests/gemm/test_group_gemm.py` were a CI-side issue, not a FlashInfer one: Triton resolves ptxas through its own knobs (`TRITON_PTXAS_PATH`, and `TRITON_PTXAS_BLACKWELL_PATH` for arch >= 100, which is the one Rubin selects) and otherwise falls back to `$CUDA_HOME/bin/ptxas`. Fixed in flashinfer-ci!354; this run confirms 0 remaining. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added architecture detection for CuTe DSL compilation, including native and family-compatible GPU architectures. * Added clear guidance when the installed DSL cannot compile for a target GPU. * **Bug Fixes** * Improved Blackwell architecture handling, including support for devices with nonstandard architecture identifiers. * Prevented unsuitable CuTe DSL backends from being selected automatically when architecture support is unavailable. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
PetersonGuo
pushed a commit
to PetersonGuo/flashinfer
that referenced
this pull request
Aug 28, 2026
…lashinfer-ai#4649) ## Problem On SM107 (Rubin) with a CuTe DSL older than 4.8, FlashInfer fails with a bare `KeyError: 'sm_107a'` raised from `enum.py` inside `cute.compile` — no FlashInfer frame in the traceback, and no warning. The root cause is that `supported_compute_capability` gates on **hardware capability alone**. flashinfer-ai#4122 widened the cute-dsl lists to `[100, 103, 107]`, which tells the dispatcher Rubin is supported regardless of which DSL is installed. Public CuTe DSL tops out at 4.7.0 on PyPI, and that release has no `sm_107a` in its `Arch` enum. This is reachable without asking for cute-dsl explicitly — two `auto` heuristics route SM107 *toward* it: - `_heuristic_func_mm_fp4`: `elif is_sm107: candidate_backends = ("cudnn", "cutlass", "cute-dsl")` - `_heuristic_func_bmm_fp8`: appends `"cute-dsl_sm107"` when `is_sm107_supported` ## Changes **1. Decline the cute-dsl backend when the installed DSL cannot emit for the device** (`fix(gemm)`) The three cute-dsl requirement functions now call `_check_cute_dsl_arch(...)`, which sits beside the existing `_check_cute_dsl_availability()` and delegates to `require_cute_dsl_arch()` — the helper added in flashinfer-ai#4122, which owns both the predicate and the message (it derives the family arch and names the exact `CUTE_DSL_ARCH=sm_100f` to export). Only the exception type is adapted, and that part is load bearing: `require_cute_dsl_arch` raises `NotImplementedError`, while `suitable_auto_backends` catches `ValueError` to mean "backend not suitable" and keeps searching. Left unadapted, an unsupported DSL would propagate out of the auto path and fail the call instead of falling back to cutlass/cudnn. Returning `False` instead of raising was also rejected: on the explicit-backend path that surfaces as `ValueError: Problem size is not supported`, which is misleading. **No capability lists change.** This is deliberately an *availability* check, not a capability one. The kernels do exist for sm_107, so `@supported_compute_capability([100, 103, 107])` stays as-is and `is_backend_supported("cute-dsl", 107)` keeps answering `True` — it is a public method on the wrapper, called with no tensors by e.g. `flashinfer/trace/templates/gemm.py:707`, and making it vary with an installed pip package would also have made the skip reason in `tests/grouped_mm/conftest.py` environment-dependent. This mirrors how the codebase already separates the two axes: `_cudnn_mm_mxfp8_requirement` lists its capabilities statically while `CUDNN_AVAILABLE` handles presence, and `_is_cudnn_override_shape_available` handles a dependency that is present but too old. **2. GDN CP delta rule resolves the arch instead of formatting it** (`fix(gdn)`) `_blackwell_compile_options` guards on the major only, then builds `f"sm_{major}{minor}a"`. Rubin is 10.7, so it passes a check written when "compute 10.x" meant Blackwell 10.0/10.3. This is the only place FlashInfer names the arch for a compute-10 device; everywhere else the DSL derives it internally. `cute_dsl_compile_arch()` returns the device's own arch when the DSL has it, the family arch when the DSL is targeting `sm_100f`, and otherwise raises `NotImplementedError` naming `CUTE_DSL_ARCH`. Same rule as the capability gate, so the two cannot disagree. ## Testing Rubin CI, `TEST_PATH="tests/gemm tests/gdn"`, against `release-v0.6.18`, with `CUTE_DSL_ARCH=sm_100f` exported and public CuTe DSL 4.7.0: | | before | after | |---|---|---| | passed | 8,141 | **12,341** | | failed | 4,230 | **1** | | `KeyError: 'sm_107a'` | 4,482 | **0** | Identical results on **both** VR200 (`hecate`, 4 workers, 2,078s) and GR100 (8 workers, 3,424s); `suite_complete=true` on both, well inside the 13,500s deadline. Per-file, verified independently on both boards: | File | before | after | |---|---|---| | `tests/gdn/test_prefill_delta_rule.py` | 2,678 | **0** | | `tests/gemm/test_mm_mxfp8.py` | 501 | **0** | | `tests/gdn/test_decode_delta_rule.py` | 417 | **0** | | `tests/gdn/test_prefill_cp_delta_rule.py` | 232 | **0** | The remaining failures are `tests/gdn/test_decode_ucache.py` and `tests/gemm/test_bmm_fp8.py` — see below. `BackendSupportedError` count is **0**, so the cute-dsl backends are being selected and compiling successfully against `sm_100f` — not silently skipped. The node accounting reconciles exactly: the plan drops 44,275 → 44,229 nodes and 37 → 36 units, i.e. the 46 tests in the ucache module leave collection entirely (a module-level skip is taken during collection, so those nodes are not counted as `skipped`). `passed` moves +287 = +288 Triton tests now compiling, −1 ucache test that previously passed; `failed` moves −333 = −288 Triton −45 ucache. Also unit-tested away from hardware: the new decorator resolves conditional 107 as False on DSL 4.7, True on 4.8+/`CUTE_DSL_ARCH`, False when the predicate raises, and yields a plain `set` when no conditional is given. `cute_dsl_compile_arch` was verified against a stubbed `Arch` enum for native / family / unsupported / Blackwell-unchanged, and the skip predicate for all four DSL-vs-arch combinations. ### Caveats - **The numbers above do not reflect this branch.** They were measured at `93143db2`, which carried a skip guard for `tests/gdn/test_decode_ucache.py` that has since been reverted, so 45 of those tests now fail again rather than skipping. - **The gemm mechanism changed after that measurement.** The two commits after it moved the check out of the decorator and into the requirement functions; that mechanism is unit-tested (adapter pass-through, `NotImplementedError` → `ValueError`, silent when the probe cannot be imported) but has not been re-run on hardware. - The measurement runs also carry `CUTE_DSL_ARCH=sm_100f` from the CI side. With it set the DSL *can* target sm_107, so `_check_cute_dsl_arch` passes and the gemm change is a no-op; only a run without that variable exercises the deselect-and-fall-back path. - `cute_dsl_compile_arch` changes `gdn_cp_prefill.py` for **all** compute-10 devices, not just Rubin. Blackwell resolution (`sm_100a` / `sm_103a`) is verified against a stubbed `Arch` enum, not on B200/GB200 hardware. ## Not addressed - **45 `KeyError: 'sm_107a'`** in `tests/gdn/test_decode_ucache.py`. Not fixable from FlashInfer: those kernels compile through `@cute.experimental.jit` / `@cute.experimental.kernel`, passing no arch and no compile options, so the DSL resolves the device arch itself and looks up `sm_107a` in its own enum. There is no FlashInfer-side site to guard, the traceback bottoms out at `enum.py:813` with no FlashInfer frame, and `CUTE_DSL_ARCH=sm_100f` does not help because that path never consults it — which points at a genuine **CuTe DSL 4.8** requirement. Left visible rather than skipped; the kernel author (flashinfer-ai#4081) is better placed to say whether it is inherent. - **1 `No valid cute-dsl SM107 bmm_fp8 config`** in `tests/gemm/test_bmm_fp8.py` — pre-existing, and present on the internal DSL 4.8 stack too (18 vs 20 occurrences across stacks), so it is independent of the DSL version question. The 288 Triton `PTXASError` failures previously seen in `tests/gemm/test_group_gemm.py` were a CI-side issue, not a FlashInfer one: Triton resolves ptxas through its own knobs (`TRITON_PTXAS_PATH`, and `TRITON_PTXAS_BLACKWELL_PATH` for arch >= 100, which is the one Rubin selects) and otherwise falls back to `$CUDA_HOME/bin/ptxas`. Fixed in flashinfer-ci!354; this run confirms 0 remaining. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added architecture detection for CuTe DSL compilation, including native and family-compatible GPU architectures. * Added clear guidance when the installed DSL cannot compile for a target GPU. * **Bug Fixes** * Improved Blackwell architecture handling, including support for devices with nonstandard architecture identifiers. * Prevented unsuitable CuTe DSL backends from being selected automatically when architecture support is unavailable. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
5 tasks
kahyunnam
added
model: qwen3.5 / 3.6 / 3.8
op: linear attention
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
3 participants
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.
📌 Description
Two CuTe-DSL kernels that make speculative decoding cheap for GDN
(Gated DeltaNet) models, plus tests and a benchmark.
The problem: during speculative decoding, the engine must be able to
roll back rejected draft tokens. For GDN models the naive way is to
save a full copy of the recurrent state for every draft position — but
GDN states are megabytes per layer per request, so those copies eat
both memory capacity and bandwidth.
The idea: do not save states at all. Save only the small per-token
ingredients (update vector u, normalized key k, decay g) in a small
ring per request. One kernel launch per layer per decode step then does
everything: computes the verify output for the draft tokens, appends
the new ingredients to the ring, and — only when a request's live
window is full, roughly once every 3-5 steps — folds the window into
the single checkpoint state. Rejected tokens cost nothing (just a
cursor that does not advance).
Ring layout and cursors: the ring has 32 physical slots per request
(
RING_SLOTS), of which at most 16 are ever live (W_RING, the maxhistory window). Two per-request int32 cursors describe the window:
cache_base(where it starts) andhist_len(how long it is); row jlives at physical slot
(cache_base + j) & 31. The kernel treats bothcursors as READ-ONLY: history is read through the rotation, and new
tokens are always appended past the live window at
(cache_base + hist_len + s) & 31— for folding and non-foldingrequests alike. Since window (≤16) plus appends (≤8) always fit in 32
slots, an append can never touch a slot any of the request's sibling
CTAs is still reading — the fold/append overlap is impossible by
construction, with no inter-CTA synchronization needed.
Cursor commits happen OUTSIDE the launch, after acceptance is known:
fold →
cache_base = (cache_base + hist_len) & 31; hist_len = n_accepted; no fold →hist_len += n_accepted. These are exactly thecursor semantics of vLLM's Triton ReplaySSM spec backend, so a serving
integration can share one cursor set between that backend and this
kernel. For standalone use the wrapper can apply the fold commit for
you (
restart_hist_on_flush=True); this path requires a caller-ownedcache_basetensor and validateshist_len ∈ [0, 16]/cache_base ∈ [0, 32)(validation is skipped while a CUDA graph iscapturing; the eager warmup calls cover it). The serving path
(
restart_hist_on_flush=False) stays sync-free.Included:
kernel (32-slot ring as above)
variant still uses the legacy 16-deep flat layout (history at rows
[0, hist_len), no cache_base) and is NOT ring-format compatible with
the flush kernel — do not share one ring allocation between them.
reference, including wrapped-window cases (base+len crossing the ring
boundary), a fold-never-overwrites-the-live-window property test, the
strided (chunk-view) q/k/v + a/b path, and negative tests for the
cursor-misuse guards
(
--no-committimes the pure kernel;--basemeasures wrappedwindows)
Supported precision modes (
g_cacheis fp32 in every mode):IO_DTYPE=fp16STATE_DTYPE=fp16RING_DTYPE=fp16STATE_DTYPE=fp16 + RING_DTYPE=fp16The fp16-state ("mixed") mode is the analogue of
mamba_ssm_cache_dtype=float16 in serving frameworks: only the
checkpoint gains fp16's extra mantissa bits; everything else stays bf16.
How to run:
pytest tests/gdn/test_decode_ucache.py -v
python benchmarks/bench_gdn_ucache_flush.py --arm bf16 --iters 1000 --no-commit
python benchmarks/bench_gdn_ucache_flush.py --arm fp16_state --iters 1000 --no-commit
Benchmark — us per layer per decode step, 1x B200, Qwen3.5-122B GDN
geometry at TP1 (H=16 key heads, HV=64 value heads, K=V=128, T=4 draft
window, 32-slot ring / 16-row window); columns = fraction of requests
folding their window that step, scattered at exact counts (CUDA-graph
replay, CUPTI cold-L2, median of 1000,
--no-commit= pure kerneltime; the cursor commit is caller-owned in serving and excluded here).
Wrapped windows (base near the ring boundary) measure the same as
base=0.
Scope/limits: validated on B200 (SM100), head dims K=V=128, draft
windows T in {4, 8}. Ring depth is fixed at 32 slots with a 16-row max
window (sized for T ≤ 8). The verify-only variant has not been
converted to the ring layout. Dtype selection via env vars and a couple
of API stubs are known follow-ups.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Summary