[Kimi-K3] Enable KDA projection fusion for DP attention

[ajit283]

Motivation

Kimi-K3's full-rank KDA projection fusion was disabled whenever attention TP differed from global TP. That excludes DP-attention deployments such as TP16 + DP-attention16, even though the full-rank merged QKVG projection accepts explicit attention-TP rank/size and can safely use the replicated attention weights.

This enables the existing fusion for that layout. No new GPU kernel is introduced.

Modifications

• Allow full-rank KDA projection fusion when attn_tp_size != tp_size, including mixed block-FP8 attention projections.
• Compute fused QKVG output splits using attn_tp_size.
• Preserve the existing restriction for the low-rank repeated/batched layout, which still requires matching attention/global TP and no quantization config.
• Add unit coverage for the full-rank DP-attention and low-rank policies, direct fused/unfused projection math, and CUDA-graph side-stream replay.
• Add a CUDA-graph projection microbenchmark and an Nsight slice comparison utility under benchmark/kernels/kimi_k3/.

For DEP16, this changes the projection prologue from five logical GEMMs:

QKV + beta + f_a + f_b + gate

to three:

QKVG + BFA + f_b

Accuracy Tests

Tested with the real Kimi-K3 checkpoint on 16 GB300 GPUs, TP16 + DP-attention16 + EP16:

• GPU unit job 630507: 7 passed, including five parameterized subtests.
  • Fused projections match the corresponding unfused BF16 math at M={1,32}.
  • Side-stream eager and CUDA-graph replay outputs match the serial path.
• Deterministic full-model A/B job 631695:
  • 32 fixed prompts, 64 generated tokens each.
  • 32/32 generated texts identical.
  • 2048/2048 output token IDs identical.
  • 2048/2048 output token logprobs identical; maximum absolute difference 0.0.
• Candidate job 523990 loaded the real checkpoint, captured all CUDA graphs, and served 2048/2048 requests.
• Post-rebase pre-commit, Python compilation, and git diff --check pass.

A preliminary normal-mode comparison produced different outputs, but a baseline-versus-baseline control also diverged despite fixed server seeds and sequential requests. The normal serving configuration is not batch-invariant and therefore is not a valid exact-equality oracle. Enabling SGLang deterministic inference made the baseline/candidate comparison bit-identical as reported above.

GSM8K

Full 1,319-question GSM8K test split, standard 5-shot completion evaluation, greedy decoding, and three independent server starts per variant:

Variant	Scores	Mean	Median
Unfused	95.814%, 95.358%, 95.510%	95.561%	95.510%
Fused	95.282%, 95.586%, 94.673%	95.180%	95.282%

The fused-minus-unfused difference is -0.381 percentage points by mean and -0.228 points by median, within the observed normal-mode run-to-run spread. The deterministic full-model comparison above remains bit-identical.

Speed Tests and Profiling

Measured on NVIDIA GB300 (SM103), TP16 + DP-attention16 + EP16, global concurrency 512 (local decode M=32), ISL/OSL 8192/1024, FP8 KV cache, BF16 SSM state, and CUDA graphs.

Repeated unprofiled serving A/B

Job 630711, three repetitions per variant on the same four-node NVL72 allocation:

Metric	Unfused median	Fused median	Change
Output throughput	3055.28 tok/s	3122.06 tok/s	+2.19%
Request throughput	2.9837 req/s	3.0489 req/s	+2.19%
Mean TPOT	95.160 ms	93.123 ms	-2.14%
Mean E2E latency	161.229 s	157.630 s	-2.23%

Per-run output throughput:

Unfused: 3038.33, 3055.28, 3061.53 tok/s
Fused:   3106.03, 3133.26, 3122.06 tok/s

Nsight profiling

Metric	Unfused	Fused	Change
Individual KDA projection prologue	140.064 us	117.281 us	-16.27%
Three projection prologues in aligned slice	420.304 us	351.697 us	-16.32%
Aligned four-layer wall time	1977.695 us	1951.731 us	-1.31%
Full graph median	45.262 ms	44.807 ms	-1.00%
Kernels per four-layer slice	114	105	-7.89%
CUDA graph nodes	2651	2444	-7.81%

Summed kernel duration increases from 2190.766 us to 2363.971 us because the side-stream BFA/f_b work overlaps QKVG and runs longer under SM contention. Wall time and full-graph latency improve, so those are the relevant end-to-end metrics.

The branch was subsequently rebased to current main; the tested kimi_k3.py remained byte-identical across that rebase.

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.

---

CI States

Latest PR Test (Base): ⏳ Run #35610810492
Latest PR Test (Extra): ❌ Run #35610809019
Latest PR Test (AMD ROCm 10): ⏳ Run #35610809647

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

All NV pipelines have passed

[nvpohanh]

cc @DarkSharpness since this is related to K3

[kpham-sgl]

Can you also test with GSM8K and/or AIME26? Thanks!

[nvpohanh]

@ajit283 please address the comments. Thanks!

[ajit283]

@kpham-sgl I ran the full 1,319-question GSM8K eval three times per variant. Baseline averaged 95.56% and this PR averaged 95.18% (medians 95.51% vs 95.28%). That is a -0.38 pp mean / -0.23 pp median difference, within the run-to-run spread we saw in normal serving. The deterministic full-model A/B is still bit-identical. I added the full numbers to the PR description.

[nvpohanh]

@ajit283 please fix the conflicts. thanks

[kpham-sgl]

@ajit283 @nvpohanh seems like #39589 already landed most of this PR's functionality.... We can close this one out