RFC: route GDN prefill through tokenspeed-kernel registry

## Summary This RFC documents the design we chose for making GDN chunk prefill registry-selected in `tokenspeed-kernel`, and how runtime handles the different checkpoint layouts exposed by the Triton/FLA and FlashInfer implementations. The high-level direction is the same pattern used by MHA: runtime calls a public `tokenspeed-kernel` op, while concrete implementations register under a shared operator family/mode and are selected by platform, dtype, and traits. ```text MambaAttnBackend / HybridLinearAttnBackend -> tokenspeed_kernel.ops.attention.gdn_chunk_prefill(...) -> select_kernel("attention", "gdn_chunk_prefill", traits=...) -> triton_gdn_chunk_prefill or flashinfer_gdn_chunk_prefill ``` ## Motivation Before this change, `MambaAttnBackend` owned too much backend choice logic for GDN prefill. It directly knew about the FlashInfer Blackwell fast path and the portable Triton/FLA fallback. That made GDN unlike MHA, where runtime calls a stable public op and `tokenspeed-kernel` owns solution selection. Moving GDN chunk prefill into the registry gives us: - vendor-neutral runtime wiring; - a portable Triton fallback for AMD/non-Blackwell; - a FlashInfer fast path on supported NVIDIA Blackwell shapes; - a future extension point for Gluon/CuteDSL without touching runtime call sites. ## Public API Shape We add a public attention op: ```python gdn_chunk_prefill( q, k, v, g, beta, *, scale, initial_state, cu_seqlens, qk_l2norm=False, output_final_state=True, output_h=False, override=None, solution=None, ) ``` The wrapper derives selection traits from the tensors and options: ```python traits = { "head_dim": q.shape[-1], "head_v_dim": v.shape[-1], "head_v_eq_head_k": v.shape[-1] == k.shape[-1], "num_v_gte_num_q": v.shape[-2] >= q.shape[-2], "qk_l2norm": qk_l2norm, "output_h": output_h, } ``` The selected backend is then invoked with the same runtime-facing arguments. ## Registered Implementations ### Triton / FLA The Triton implementation is registered as: ```text family="attention" mode="gdn_chunk_prefill" name="triton_gdn_chunk_prefill" solution="triton" ``` It wraps the existing FLA-derived `chunk_gated_delta_rule` implementation and is portable across AMD/NVIDIA. It advertises broad support for `qk_l2norm` and `output_h`. When `output_h=True`, Triton returns native FLA layout: ```text (out, final_state, h) ``` where: ```text h.shape = [1, total_ceil_chunks, H, K, V] h[i] = state before chunk i ``` ### FlashInfer The FlashInfer implementation is registered only when its runtime/library prerequisites are available. It is gated by: - NVIDIA Blackwell arch range `10.0..10.3`; - CUDA major >= 13; - FlashInfer GDN prefill kernel presence; - bf16 / head-dim / head-count traits. It registers as: ```text family="attention" mode="gdn_chunk_prefill" name="flashinfer_gdn_chunk_prefill" solution="flashinfer" ``` When `output_h=True`, FlashInfer returns compact full-chunk checkpoints: ```text (out, final_state, h_checkpoints, h_cu_starts) ``` where: ```text h_checkpoints.shape = [total_full_chunks_across_batch, H, K, V] h_checkpoints[i] = state after a complete chunk ``` ## Layout Decision We considered normalizing FLA output into FlashInfer's compact checkpoint layout inside the Triton wrapper. That would have made the public op return one checkpoint shape for all backends, but it would also add nontrivial copy/allocation overhead: each checkpoint row can be large (`H * K * V`, around megabytes for Qwen3.5 GDN TP=1). After comparing SGLang's GDN design, we chose the less invasive and lower-overhead route: - let Triton return native FLA `h`; - let FlashInfer return its native compact checkpoint layout; - make `MambaAttnBackend` understand which returned layout it received. This mirrors SGLang's broader shape: the runtime has a loose contract of `(out, final_state, h)` when a backend exposes intermediate states, and backend-specific details are handled at the boundary rather than forcing every implementation through one layout. ## Runtime Handling `MambaForwardMetadata` carries two h-source index views: ```python track_ssm_h_src # FlashInfer compact checkpoint index track_ssm_h_src_fla # FLA native h index ``` For FlashInfer layout, the source index is based on full chunks: ```text num_fi_ckpts = extend_seq_lens // 64 track_ssm_h_src = offset + (track_lens // 64 - 1) ``` For FLA layout, the source index is based on ceil chunks and state-before-chunk semantics: ```text num_fla_states = (extend_seq_lens - 1) // 64 + 1 track_ssm_h_src_fla = fla_offset + (track_lens // 64) ``` At runtime: ```python gdn_out = gdn_chunk_prefill(..., output_h=True) if len(gdn_out) == 4: # FlashInfer: (out, final_state, h, h_cu_starts) h_src = track_ssm_h_src elif len(gdn_out) == 3: # Triton/FLA: (out, final_state, h) h_src = track_ssm_h_src_fla ``` The actual copy remains shared: ```python ssm_states[track_ssm_h_dst] = h[h_src].to(ssm_states.dtype, copy=False) ``` ## Why Not Force One Layout? A single normalized output layout is cleaner as an API, but in this case it makes the Triton path pay for conversion solely to match FlashInfer. Because h-track is cache-policy metadata rather than core recurrence math, preserving native layouts avoids avoidable memory traffic and keeps the selected kernels closer to their natural representation. The tradeoff is that `MambaAttnBackend` now has to know how to interpret the two return forms. We accept that because it is localized and avoids widening the public kernel API with extra metadata objects or forced layout conversions. ## Current Assumptions - GDN chunk size is 64. - Qwen3.5 GDN uses `linear_key_head_dim=128` and `linear_value_head_dim=128`, even though full-attention MHA has `head_dim=256`. - FlashInfer GDN prefill is a specialized Blackwell fast path. - Triton/FLA remains the portable fallback. - Future Gluon/CuteDSL GDN implementations should register under the same `attention.gdn_chunk_prefill` mode and either return one of the existing layouts or introduce a clearly handled third layout. ## Validation Local validation performed on this branch includes: - `pre-commit run --all-files` - `python -m pytest tokenspeed-kernel/test/ops/test_attention_gdn.py tokenspeed-kernel/test/test_kernel_api_selection.py -q` - `python -m pytest test/runtime/layers/test_mamba_checkpoint_metadata.py test/runtime/layers/test_gdn_qkv_split_fused.py test/runtime/layers/test_gdn_flashinfer_fastpath.py -q` - TP=1 non-offline smoke for `amd-Qwen3.5-35B-A3B-MXFP4` via `tokenspeed serve` and `test_coherence.sh`.