[yyihuang]

📌 Description

This PR adds an optimized SM100a BF16 recurrent-KDA prefill backend and
dispatches to it through the existing recurrent_kda API.

The new backend coexists with the current CuTe-DSL implementation. Decode
(T=1), speculative decode, GQA, state-pool/checkpoint features, unsupported
layouts and dtypes, and non-B200 devices continue to use the existing backend
unchanged.

Dispatch contract

The optimized path is selected only when all of the following hold:

• the device is exactly compute capability 10.0 (B200 / SM100a);
• the call is ordinary multi-token prefill: fixed T > 1, or packed input
  with more tokens than sequences;
• Q, K, V, and G are contiguous BF16 [B, T, H, 128] tensors with a shared
  head count, and beta is contiguous BF16 [B, T, H];
• A_log is contiguous FP32 [H], while dt_bias is contiguous FP32
  [H, 128] or [H * 128];
• use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True,
  beta_is_logit=True, and lower_bound is finite and negative;
• speculative decode, GQA, state indices, committed-state sources, and
  accepted-token/checkpoint features are disabled.

Packed input uses B=1 and contiguous CUDA int32/int64 cu_seqlens. The
binding consumes int64 offsets; CUDA graph capture therefore requires callers
to supply int64 offsets directly. An optional contiguous CUDA int32
seq_order can order packed sequences for better tail utilization.

Calls that do not exactly match this contract fall back to the existing
CuTe-DSL path.

Physical schedules and true in-place state update

• M64 is used only for fixed-layout B=1, H=64.
• M128 handles every other eligible fixed-layout call and all eligible packed
  calls.
• State layout remains BF16 [N, H, V, K] with V=K=128.
• A supplied initial_state is passed as both the initial- and final-state
  kernel pointer. The kernel updates that allocation directly; there is no
  state scratch allocation and no post-kernel copy_/memcpy.
• M128 assigns one CTA to each (sequence, head). M64 assigns two CTAs whose
  64-row value partitions are disjoint. Each CTA loads all initial-state rows
  it owns before storing those same final-state rows, making exact pointer
  aliasing safe.
• The binding accepts either disjoint state buffers or an exact full-range
  alias and rejects partial overlap.
• When initial_state=None, final-state storage is created only if
  output_final_state=True.

The generated source emits __restrict__ for every pointer and has no
per-argument alias escape hatch, so the removal of __restrict__ from only the
two state parameters is an explicitly delimited FlashInfer ABI integration
patch. The frozen-source test restores those two qualifiers, removes the
existing tensor-map acquire integration prologue, and verifies the remaining
generated source against its normalized SHA256.

CUDA graph and tensor-map safety

RecurrentKDAPrefillWorkspace owns optional final-state scratch for calls
without an initial state, beta padding, and separate 768-byte M64/M128
tensor-map descriptor blocks. It binds to one stream, is warmed eagerly
against the exact tensor signature, and is used by one captured
recurrent_kda invocation. A workspace that has participated in capture is
rejected by later Python calls; graph replay remains valid because it does not
re-enter Python.

Tensor maps are prepared outside capture and published to stable global
storage on the caller stream. Every consumer CTA executes
fence.proxy.tensormap::generic.acquire.gpu for all six maps, followed by a
CTA barrier, before any TMA instruction can consume them.

Other integration

• Adds JIT and AOT registration for separate M64 and M128 translation units.
• Updates tracing and API documentation for seq_order and
  prefill_workspace.
• Adds correctness, routing/fallback, ABI, stream, overlap, workspace, and
  CUDA graph tests.
• Adds a CUPTI cold-L2 benchmark for the six fixed/packed contract shapes,
  including a commit- and binary-verified MoonshotAI/FlashKDA peer.

🧪 Validation

Final validation was run on pushed head
f6c3b0c486787e64f76b2be91ddb239ea3ba66ce, on an NVIDIA B200
(compute capability 10.0), PyTorch 2.13.0.dev20260503+cu132, CUDA 13.2,
and Python 3.13.13.

Gate	Result
Frozen-source/JIT specification tests	8 passed
Full targeted recurrent-KDA prefill suite, including fixed/packed correctness, M64/M128, non-default streams, and CUDA graph capture/replay	28 passed
Focused M64/M128 compute-sanitizer --tool memcheck	5 passed; ERROR SUMMARY: 0 errors
Focused M64/M128 compute-sanitizer --tool synccheck	5 passed; ERROR SUMMARY: 0 errors
Pre-commit hooks on the change	Passed

The full repository test suite was not run.

Fresh B200 CUPTI benchmark

The candidate is invoked only through the public
flashinfer.kda_decode.recurrent_kda entry point and performs its state update
in place inside the kernel. Its timed region contains no state memcpy.

The comparison uses six fixed/packed shapes with deterministic seeds, matching
tensor distributions, BF16 state, and scale=1/sqrt(128). Preinitialized
rotating state buffers ensure each timed candidate invocation receives the
same initial state. Allocation, metadata creation, sequence ordering,
state-pool reset, and JIT/cache warmup are outside the measured region.

The peer is MoonshotAI/FlashKDA._fwd_raw, from source commit
d2ff19a6
and CUTLASS
5c149f52.
The loaded extension SHA256 is
997c3a1d1338f8bf9dba3c1a01386b1b74448214c294d64409454cc11141c04c.
The benchmark records the source revision and independently computed extension
digest.

Timing uses CUPTI activity tracing, cold-L2 flushing, no CUDA graph, 20 ms
warmup, and 100 ms measurement per block. Each value is the median of two
independent block medians; the reported scopes occupy symmetric
PR/raw/.../raw/PR positions. Speedup is FlashKDA / this PR.

Case	Layout / sequence lengths	Variant	This PR public in-place API (ms)	FlashKDA raw peer (ms)	Speedup
H=96 fixed	[8192]	M128	0.506095	1.085277	2.1444x
H=96 mixed	[1300, 547, 2048, 963, 271, 3063]	M128	0.391135	0.894702	2.2875x
H=96 uniform	[1024] × 8	M128	0.434992	0.735535	1.6909x
H=64 fixed	[8192]	M64	0.469711	0.992638	2.1133x
H=64 mixed	[1300, 547, 2048, 963, 271, 3063]	M128	0.269903	0.676574	2.5067x
H=64 uniform	[1024] × 8	M128	0.292567	0.495871	1.6949x
Geometric mean					2.0512x

All six output/state comparisons against FlashKDA passed at
atol=rtol=1e-2; the maximum observed absolute errors were 0.0009765625
for output and 0.015625 for state. The benchmark JSON SHA256 is
c34b27f9e8fb4e4dfbc500e6976e63668fcd0bff2bb00086cd516de59491f259.

Upstream recurrent-KDA comparison

Pinned upstream main does not implement ordinary multi-token recurrent-KDA
prefill through its public API: fixed T != 1 is rejected, and packed standard
decode processes one token per sequence. Speculative decode has different
checkpoint/state semantics and is not an equivalent workload. The upstream
prefill comparison is therefore unsupported / N/A, rather than timing a
different operation.

🚀 Pull Request Checklist

• Tests were added or updated.
• Targeted tests and pre-commit hooks pass.
• Full repository test suite was not run.

Reviewer Notes

The main review surfaces are:

1. strict eligibility and unchanged fallback behavior;
2. exact-alias state ABI and direct in-kernel update;
3. workspace lifetime, stream binding, and CUDA graph rules;
4. tensor-map publication/acquire ordering;
5. the two explicitly delimited frozen-CUDA integration patches and their
   integrity normalization.