[jianc99]

This is component 1 of the DFlash adaptive-verification series. Original implementation and component-only review: jianc99#1. Integration and full-system results: #46182.

Updated onto upstream main (6a319e7). All applicable pre-commit checks passed on this branch. The GPU tests and benchmark results below were run previously on the original source branch; they have not been rerun after this upstream update. Submission preparation used AI assistance.

---

Motivation

On hybrid models, linear-chain speculative verify (EAGLE chain, MTP, DFLASH, DSPARK, standalone) calls causal_conv1d_update with an intermediate conv window. The window lets the conv state of the accepted prefix be restored after verify.

On CUDA and ROCm, MambaPool stores this window deduplicated. Each slot holds one [dim, T + K - 2] row, exposed as an overlapping [T, dim, K - 1] view whose step stride equals its window stride.

The kernel still writes the window per token, with three lane-strided element stores per draft token. Each column is therefore written up to three times by narrow scattered stores. On GB300 (dim 10240, bf16, T = 8, bs 128), the verify conv takes 106 us with the deduplicated window, against 39 us with a dense window.

Modifications

• _causal_conv1d_update_kernel: adds a DEDUP_CONV_WINDOW path.

  • Step t's window is columns t..t+2 of the row [state[1], state[2], x_0, ..., x_{T-1}].
  • The kernel writes that row once, as a [BLOCK_N, next_pow2(T + 2)] tile before the token loop, and skips the per-token window stores.
  • The tile mask covers exactly the T + 2 columns the per-token path writes.
  • Padded rows still exit before any store.
  • Two tl.static_asserts keep the branch from being compiled for a conv width other than 4 or together with a tree mask, so the kernel is safe without relying on its caller.

• causal_conv1d_update: takes the tile path only when all three hold:

  • the window view's step stride equals its window stride;
  • the verify is a chain (no tree mask);
  • the conv width is 4.

  With step stride == window stride == s, the per-token stores write column t + w of the row at (t + w) * s, and the tile writes column c at c * s: the same addresses and the same values, for any layout with those strides. MambaPool's deduplicated view is one such layout. There are no new arguments, and callers are unchanged. Kernels that do not take the tile get NP2_WINDOW_COLS=1, so no extra variants are compiled.

• Test: test_causal_conv1d_update_dedup_window_matches_per_token.

  • It runs the kernel on a dense window (per-token path) and on the deduplicated view (tile path).
  • It checks bitwise that the shared rows, the windows read through the view and the conv state match a reference row built in torch.
  • It checks that slots of padded rows are untouched.
  • It compares the conv output with the file's update tolerances.
  • Coverage: T in {2, 3, 8, 16} (T + 2 pads to tiles of 4, 8, 16 and 32), a num_accept_tokens state offset, MambaPool's axis-0 window layout seen transposed, fp32 and bf16, a dim that leaves a partial channel block, with and without padded rows.

• A comment in MambaPool is updated to describe the new write.

Scope

This path is used by every chain verify (topk <= 1) whose conv window uses MambaPool's deduplicated layout. Examples are GDN models (Qwen3.5/3.8) and Mamba2, on CUDA and ROCm.

These are unchanged and keep the per-token stores:

• dense windows: UnifiedMambaPool, KDA / Kimi-Linear, tree verify;
• deduplicated windows of a conv width other than 4 (for example LFM2, width 3).

Accuracy Tests

• The new test passes, and so does the full test file (270 tests at 8eb01f2, which adds the tree-mask assert and the accept-offset, axis-0 and odd-T cases).
• The test is sensitive: a tile that ignores num_accept_tokens fails exactly the accept-offset cases, and forcing the tile together with a tree mask fails at compile time on the new assert.
• compute-sanitizer memcheck, racecheck and initcheck are clean on the tile test as of 8177311 (T in {2, 4, 8, 16}); besides comments and tests, 8eb01f2 adds only a compile-time assert and NP2_WINDOW_COLS=1 for launches without the tile. The extra 72-case check below, which covers accepted-token offsets, transposed views and odd T, is also memcheck-clean.
• An extra 72-case check also passes, comparing the tile path bitwise with the per-token path. It covers T in {1, 3, 9}, mixed input and state dtypes, transposed views and accepted-token offsets.
• End to end, greedy, c1: acceptance length and output tokens are identical to main.

Benchmarking and Profiling

Kernel time per call in microseconds (GB300, dim 10240, K 4, bf16, CUDA graph):

bs	T	before	after	dense window
1	8	5.35	2.49	3.71
64	8	57.39	9.96	21.08
128	8	106.45	17.55	38.58
128	16	322.63	39.10	74.12

End-to-end setup:

• Qwen3.8-27B + DFlash2 (block 8), TP1, one GB300;
• chain verify, bf16 mamba state, overlap scheduling on, --max-running-requests 128;
• SPEED-Bench first turns, T = 1, top_p 0.95, top_k 20, up to 4096 output tokens;
• mean of 3 interleaved runs.

Decode tok/s = per-request decode rate (output tokens after the first / time after the first token), times concurrency; c1 is the single-request decode rate.

concurrency	main	this PR	change
64	6048	6361	+5.2%
128	6992	7329	+4.8%

Acceptance length is 3.59-3.63 in all six runs.

At c128 the running cap equals the concurrency, so about one request waits in the queue. #46182 shows that this setup reads 7-15% higher at that level than a larger cap, on both sides, so compare c128 only within this table.

Single runs at lower concurrency: c8 2031 -> 2104 (+3.6%), c32 4723 -> 4973 (+5.3%). At c1 the change is within noise: 332 -> 321, with acceptance length 3.70 vs 3.60 from sampling; per accepted token the rate differs by under 0.5%. Greedy (single runs): c1 357.0 -> 357.1 with identical acceptance length and output tokens, c64 6746 -> 7145 (+5.9%).

The end-to-end runs used the first commit; the two later commits add compile-time asserts, comments and tests, and pass NP2_WINDOW_COLS=1 to kernels that skip the tile, with no runtime change on the tile path.

Related upstream work

• [Fix] Honor DSpark SPS profiler budget pins and handle ragged mRoPE positions #38897 adds a separate packed variable-length kernel, causal_conv1d_update_varlen, to the same file. It writes the intermediate window per token with three stores per step, the pattern this PR replaces. There is no textual conflict (it is a new function), but if both land, packed verify through that kernel does not get the tile store. [Feat] gdn: run target verify on packed ragged tokens (ReplaySSM) #46189 takes the other route: it adds packed rows to this kernel, so the tile store covers them too.
• This PR does not depend on any open PR.

Notes

• The tile store is SubSir's work from the dynamic-verify branch. This PR ports it to the current kernel, gates it on the view's strides and adds the tests.
• Only GB300 was tested; ROCm takes the same path but has not been run.
• Commits: 7526326 (tile store and stride gate), 8177311 (width assert, test trim, MambaPool comment), 8eb01f2 (tree-mask assert, stride-based wrapper comment, NP2_WINDOW_COLS=1 when the tile is off, more test layouts). Squash them when merging upstream.

---

CI States

Latest PR Test (Base): ✅ Run #40568730789
Latest PR Test (Extra): ❌ Run #40568730891
Latest PR Test (AMD ROCm 10): ❌ Run #40568730701