[Codoge]

Motivation

Kimi-K3's token-major prefill convolution leaves strided Q/K/V views, and its merged gate projection leaves a strided beta view. FlashInfer's KDA wrapper materializes these inputs with separate copies. Batch those required copies into one Triton launch without changing KDA arithmetic, gates, state precision, or checkpoint handling.

This integrates the measured PP8 prototype into upstream main. Full-model serving validation now passes for these production changes on the common upstream main base 592754fc5e, with FlashInfer 0.7.2; the standalone and repeated combined results are reported below. Real GPU operator tests pass on both FlashInfer 0.7.1 and the pinned 0.7.2.

Relationship to #47451

The optimizations overlap on the Q/K/V materialization cost, but use different approaches:

• [Test] kda: Cover contiguous QKV graph inputs and asymmetric fallback #47451 changes the convolution producer to write contiguous Q/K/V groups, eliminating those three copies. It also covers the FlashInfer prefill-graph producer path.
• This PR leaves the producer layout alone and combines outstanding consumer-side copies, including beta. It covers both FlashInferKDAPrefillKernel.extend and FlashInferKDAPrefillGraph.run through the same helper.
• Contiguous inputs are returned unchanged. When fewer than two inputs require copying (e.g. grouped Q/K/V plus only strided beta), the ordinary contiguous path is retained. There is no redundant Q/K/V copy or extra fused-copy launch in that case.
• Their standalone speedups must not be added. Prefer the producer-layout solution where applicable; this is a separately reviewable alternative for remaining token-major inputs, not a second claimed gain on top of [Test] kda: Cover contiguous QKV graph inputs and asymmetric fallback #47451.

Modifications

• Add a bit-preserving multi-input copy kernel under kernels/ops/attention and use it at the FlashInfer KDA input boundary.
• Keep contiguous tensors aliased, validate device/dtype/indexing bounds, and fall back to PyTorch for unsupported inputs or autograd.
• Preserve input slicing and all existing KDA computation arguments. No new environment variable or CLI option.
• Add registered GPU coverage for the real 16K/96-head K3 layout, short/tail/empty shapes, BF16/FP32 bit patterns, already-grouped QKV, CUDA Graph replay, and real FlashInfer outputs/state.

Accuracy Tests

On an idle B300/B30Z (SM103), PyTorch 2.14.1+cu130, separately using FlashInfer 0.7.1 and 0.7.2:

• 5 GPU tests passed on each version, including actual Triton execution at [1, 16384, 96, 128] and CUDA Graph replay with changed inputs.
• Raw-bit comparisons cover NaN payloads, infinities, subnormals, and signed zero.
• Real FlashInfer KDA calls on strided inputs match contiguous-reference outputs and recurrent state exactly, in eager mode and through the production graph adapter replayed with updated inputs. Warmup and capture use the same stream, as required by the FlashInfer workspace.
• The candidate kernel and wrapper were loaded from this branch into the existing GPU environment; no installed source was modified. FlashInfer 0.7.2 was tested in an isolated wheel overlay without bypassing its version checks.
• Ruff, isort, AST parsing, registered-test lint, and git diff --check passed.

The current six full-service runs each pass 12 fixed greedy output checks covering input lengths 1,024 / 16,383 / 16,384 / 16,385 / 32,768 / 65,536, each repeated twice. All 60 comparisons against A1 match output token IDs and input SHA256 hashes. An additional eight C4 concurrent fixed-input checks match between A2 and B2. These are finite output-regression checks, not logits-level equality or a full task-accuracy benchmark.

Blackwell review follow-up (2026-11-01)

Reran both complete files on B300/B30Z (SM103), PyTorch 2.14.1+cu130 and FlashInfer 0.7.2, using the complete Python tree from PR head 7933473906e291d9a14d26ba78aca0284820ae1e:

• test/registered/kernels/ops/attention/test_kda_input_pack.py: 5 tests passed, 0 skipped.
• test/registered/kernels/ops/attention/test_kda_prefill_graph.py: 9 tests passed, 0 skipped, including the FlashInfer/Blackwell cases that were skipped on the H100 CI runner.

Speed Tests and Profiling

Full serving validation of the upstream integration (2026-11-01)

Tested production changes from #47651 (7933473906e) and #47653 (34238f31a5d) on the same upstream main base 592754fc5e, using FlashInfer 0.7.2 for every arm. Kimi-K3, one 8-GPU B300/B30Z, TP1/PP8/EP1, chunk 16,384, pipeline async depth 1, client C4, ISL 65,536 / OSL 1. CUDA Graphs and radix caching disabled; model precision and all serving settings held fixed. Each formal run contains 128 requests after a separate 8-request warmup. Formal measurements exclude profiler/diagnostic requests and prototype hooks.

Each run uses an independent autotune cache cloned from the same baseline startup cache. Baseline/KDA caches contain no 16K MoE seed; MoE/combined caches contain the FlashInfer 0.7.2 (128, 4) record on all eight PP ranks. Separate post-benchmark Torch GPU traces of all eight PP ranks confirm the actual production paths: 1,104 KDA input-copy kernels become 276 pack kernels, and all 736 MoE GEMMs change token tile 192 to 128. These C1 diagnostic traces verify execution paths; their timings are not used as C4 serving measurements.

Run	Requests	Input tok/s	P50 TTFT s	P99 TTFT s
baseline A1	128	42,203.78	6.128	7.046
KDA only	128	42,899.91	6.022	6.937
MoE only	128	42,813.05	6.034	6.926
combined B1	128	43,569.94	5.925	6.835
baseline A2	128	42,130.58	6.137	7.050
combined B2	128	43,524.13	5.940	6.823
baseline pooled	256	42,167.15	6.132	7.359
combined pooled	256	43,547.02	5.930	7.125

KDA-only throughput changes by +1.65% against A1. This is one standalone screening pair. The repeated combined KDA+MoE result changes throughput by +3.27% (paired runs +3.24% / +3.31%); it must not be attributed to either PR alone. Pooled throughput is total input tokens / total formal duration; P50/P99 are recomputed from all formal request TTFTs. All 768 formal requests succeed with zero cached tokens and retractions. The combined pooled P50/P99 TTFT decrease by 3.28% / 3.17%; no formal request reaches 15 seconds. OSL is 1, so TPOT/ITL are not reported.

Earlier prototype evidence

Historical prototype results, not a new end-to-end benchmark of this branch. Kimi-K3 / SGLang 0.5.22 / FlashInfer 0.7.1, one 8-GPU B300, TP1/PP8, chunk 16,384, client concurrency 4, ISL 65,536 / OSL 1. CUDA Graphs and radix caching disabled; all model precision settings unchanged. Warmups and profiling requests excluded.

Run	Requests	Input tok/s	P50 TTFT s	P99 TTFT s
baseline A1	128	41,924.66	6.161	7.135
pack B1	128	42,610.43	6.062	7.022
baseline A2	128	41,811.10	6.176	7.110
pack B2	128	42,526.45	6.078	6.996
baseline pooled	256	41,867.80	6.164	7.438
pack pooled	256	42,568.40	6.071	7.320

Pooled throughput changes by +1.67%. Throughput uses total tokens / total duration; percentiles are recomputed from all formal requests, not averaged between runs. All requests succeeded with no cache hits or retractions. Only one output token is generated, so TPOT/ITL are not reported.

The combined prototype's NSYS trace showed KDA copy-related time changing from 1.5405 to 0.3566 ms per matched layer invocation. That includes remaining copies and is not the complete KDA layer time. Required element traffic is unchanged. The combined MoE+KDA serving gain is not attributed to this PR alone.

Checklist

• Keep arithmetic and precision unchanged; add no configuration switch.
• Run actual GPU copy, graph replay, and FlashInfer output/state checks.
• Document overlap with [Test] kda: Cover contiguous QKV graph inputs and asymmetric fallback #47451 and avoid redundant work with grouped QKV.
• Separate prototype performance evidence from exact-branch validation.
• Complete common-base full-model serving validation, including repeated baseline/combined A/B and fixed output checks.

---

CI States

Latest PR Test (Base): ✅ Run #41201160895
Latest PR Test (Extra): ❌ Run #41201160980
Latest PR Test (AMD ROCm 10): ❌ Run #41201160985

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

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_kda_input_pack.py test/registered/kernels/ops/attention/test_kda_prefill_graph.py

[github-actions]

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_kda_input_pack.py
cd test/ && python3 registered/kernels/ops/attention/test_kda_prefill_graph.py

[BBuf]

[BBuf Review Agent] LGTM. Please run the copy and graph tests on B200/B300 before merging.

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_kda_input_pack.py test/registered/kernels/ops/attention/test_kda_prefill_graph.py

[github-actions]

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_kda_input_pack.py
cd test/ && python3 registered/kernels/ops/attention/test_kda_prefill_graph.py

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[Codoge]

@BBuf Thanks for the review. Both requested test files have now been rerun against PR head 7933473906e291d9a14d26ba78aca0284820ae1e on B300/B30Z (SM103), PyTorch 2.14.1+cu130, FlashInfer 0.7.2:

python3 test/registered/kernels/ops/attention/test_kda_input_pack.py -v
Ran 5 tests in 7.085s
OK

python3 test/registered/kernels/ops/attention/test_kda_prefill_graph.py -v
Ran 9 tests in 82.067s
OK

There were zero skips in either file. This covers the Blackwell-only FlashInfer cases skipped by the earlier H100 CI run. A complete PR Python tree and a version-matched FlashInfer overlay were used; version checks remained enabled.

The PR description now also contains the completed common-base full-service comparison: KDA-only +1.65% (one screening pair); combined KDA+MoE +3.27% over two A/B pairs (+3.24% / +3.31%). The combined gain is not attributed to this PR alone. All 768 formal requests and finite fixed-output regressions passed.

Marked ready for review. Please re-review when convenient, including the documented overlap/choice with #47451. My account cannot dispatch targeted /rerun-test; if another upstream run is needed, the exact command is:

/rerun-test test/registered/kernels/ops/attention/test_kda_input_pack.py test/registered/kernels/ops/attention/test_kda_prefill_graph.py

[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.

[Codoge]

/rerun-failed-ci

[BBuf]

[BBuf Review Agent] @Codoge The zero-skip B300 runs on 7933473906e address my Blackwell test request. This covers the remaining validation I asked for.

[BBuf]

/rerun-failed-ci