[SunskyXH]

Motivation

_get_compiled_kernel in cutedsl_kda.py compiles the CuTe DSL KDA decode kernels against a bf16 dt_bias placeholder, and the dtype is not part of the compile-cache key. CuTe launches do not check tensor dtypes, so when a caller passes an fp32 dt_bias, the kernel reads the fp32 bytes as bf16. Kimi-Linear keeps dt_bias in fp32 (kimi_linear.py), and kda_cutedsl.py forwards it without a cast. So with --linear-attn-decode-backend cutedsl the forget gates are wrong, and NaN shows up in the decode output and in every state slot the batch touches.

Modifications

• cutedsl_kda.py: compile with the caller's dt_bias dtype, and add the dtype to the compile-cache key. The kernel bodies are unchanged; they already convert dt_bias to fp32 on load.
• New test/registered/kernels/ops/attention/test_cutedsl_kda.py, modeled on test_cutedsl_gdn.py. For batch sizes 16 and 128, with dense and varlen inputs, it runs bf16 and then fp32 dt_bias at the same shape. Over 3 decode steps it compares the outputs and the whole state pool against Triton fused_sigmoid_gating_delta_rule_update(is_kda=True), with one padded -1 slot. Running both dtypes at one shape means a kernel compiled for one dtype cannot quietly serve the other.

Accuracy Tests

The new test on one H200 (H = HV = 32, K = V = 128):

	bf16 dt_bias	fp32 dt_bias (as in Kimi-Linear)	test
main	matches Triton	NaN, 97% of outputs mismatched	4/4 fail
placeholder dtype fixed, cache key unchanged	matches Triton	NaN (reuses the bf16 kernel)	4/4 fail
this PR	matches Triton	matches Triton	4/4 pass

test_cutedsl_gdn.py still passes.

Speed Tests and Profiling

The kernels are unchanged. A model compiles one variant for its dt_bias dtype, and nothing is added per call.

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.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #41751237296
Latest PR Test (Extra): ❌ Run #41751237093
Latest PR Test (AMD ROCm 10): 🚫 Run #41751237232

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[BBuf]

LGTM. Compiling with the caller's dt_bias dtype and including it in the cache key addresses both the incorrect loads and cross-dtype kernel reuse. The regression test exercises both dtypes at the same shape and checks outputs and the full state pool over multiple decode steps. Reviewed the code; GPU tests were not rerun locally.

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[BBuf]

[BBuf Review Agent] @SunskyXH these still fail after a rerun, and I can't tell from the logs whether it's this PR. Could you take a look?

• base-b-test-1-gpu-large (2): {\"file\": \"test/registered/dllm/test_dllm_batching_fdfo.py\", \"passed\": false, \"elapsed\": 82}" identifies a failed test file, but the tail contains no error or traceback establishing its cause o

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci