[KDA] Support BF16 recurrent state in the CuTeDSL MTP verify kernel

[milo-modal]

Motivation

The fused CuTeDSL KDA MTP verify kernel (python/sglang/kernels/ops/attention/kda_decode_mtp.py) only accepts FP32 recurrent state. With a BF16 SSM state pool (mamba_ssm_dtype=bfloat16), KDAAttnBackend._can_run_dspark_cutedsl_mtp rejects the call, and target verify falls back to the slower Triton path.

Modifications

• Kernel (kda_decode_mtp.py):
  • Stage recurrent-state tiles through shared memory in the storage dtype. cp.async copies 4 FP32 or 8 BF16 values per 128-bit transfer.
  • Widen the tiles to FP32 as they load into registers. The token recurrence stays FP32 and only the per-token snapshot stores (intermediate_ssm) round to BF16.
  • Keep the shared-memory bank-conflict padding constant in bytes rather than elements, so BF16 rows land on the same conflict-free banks as FP32 rows.
  • The FP32 kernel compiles to the same configuration as before.
• Wrapper (fused_kda_decode_mtp_dspark): accept FP32 or BF16 state. Check cp.async alignment in bytes (16-byte slot stride and base pointer). Require intermediate_ssm to be present, with the same dtype as the state, when the ReplaySSM ring is not used.
• Dispatch (kda_backend.py): take the fused path for FP32 or BF16 state when the intermediate snapshot cache has the same dtype; otherwise fall back as before.

Accuracy Tests

Kernel level, test/registered/kernels/ops/attention/test_kda_mtp_cutedsl_replayssm_ring.py on B200 (SM100): 9/9 pass. On current main, the three BF16 cases fail and the six FP32 cases pass.

• The ReplaySSM ring parity and CUDA-graph padding-slot tests now also run with BF16 state.
• New test_cutedsl_bf16_state_matches_fp32_register_math_block8: with a BF16-representable initial state, BF16 storage produces output bitwise equal to FP32 storage, and BF16 snapshots equal the FP32 snapshots rounded to BF16. This shows the arithmetic stays FP32.

Speed Tests and Profiling

One KDA target-verify layer on B200

num_spec	batch	Triton BF16 (main)	CuTeDSL BF16 (this PR)	Speedup
3	1	13.1	11.6	1.13x
3	8	14.5	12.0	1.21x
3	32	31.7	30.7	1.03x
3	48	35.9	39.0	0.92x
3	128	85.7	104.2	0.82x
7	1	23.3	20.3	1.15x
7	8	25.2	20.7	1.22x
7	32	56.3	54.5	1.03x
7	48	62.8	69.5	0.90x
7	128	154.8	184.2	0.84x

BF16 state on the selected nv_cutedsl backend is faster up to about 32 requests per verify step and slower from about 48. Unlike Triton, the CuTeDSL kernel gains little from halving the state size. For FP32 state, CuTeDSL is faster than Triton at every batch size measured (1–256), so FP32 behavior is unchanged. High-concurrency BF16 deployments can keep the previous behavior with --linear-attn-verify-backend triton. This PR honors the selected backend rather than silently falling back.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. N/A: no user-facing option changed.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. Kernel level; see above.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #36767611565
Latest PR Test (Extra): ❌ Run #36767611149
Latest PR Test (AMD ROCm 10): ❌ Run #36767611801

[rodamani]

/tag-and-rerun-ci

[zcnrex]

LGTM