[Codoge]

Motivation

Kimi-K3 has 12 local MLA heads at TP8. _concat_and_cast_mha_k currently requires a power-of-two head count before calling the Triton helper, even though that helper already handles padded head counts. These shapes fall back to two separate tensor copies.

Modifications

Remove the redundant head-count restriction at the shared model dispatch point. Reuse the existing masked concat kernel and keep the feature-dimension guards and the 128-head specialized path.

For FA3 with FP8 KV storage, these head counts now follow the same output-dtype rule as the existing power-of-two Triton path. Other covered backends retain the input dtype.

Accuracy Tests

Upstream b146a8576de930cf3adc104e601feb2790cdf3f6 plus this PR only, NVIDIA B30Z (B300), PyTorch 2.14.1+cu130 and Triton 3.8.0:

python -m pytest test/registered/kernels/ops/attention/test_concat_mla.py -k model_mha_k_concat -q
# 14 passed, 94 deselected

The tests call the actual model helper and Triton implementation. Coverage includes BF16 / FP16, empty and 64K-token inputs, 12 / 24 heads, an 8-head control, strided inputs, and repeated CUDA Graph replay with changed inputs. Outputs match the reference exactly. The FA3 dtype case compares raw FP8 bytes for finite random inputs against the reference cast.

Full-model accuracy and serving throughput have not been rerun on this branch; the results below are isolated operator timings.

Upstream CUDA validation

The complete registered test/registered/kernels/ops/attention/test_concat_mla.py passed on the upstream H100 CI runner for head 81955274f6363c91979dd12d05e2662621bc4346: 108 passed, 0 skipped, including all 14 new model-helper cases and the existing JIT/AOT regressions. Passing CI job.

Speed Tests and Profiling

Isolated concat operator on one NVIDIA B30Z, 12 heads, nope_dim=128, rope_dim=64, BF16 output and strided KV / latent views. Baseline is the exact two-copy fallback at the base commit; patched runs the model helper. Both allocate their output. Eager execution, no profiler and no CUDA Graph; triton.testing.do_bench, 25 ms warmup and 100 ms measurement, median CUDA-event latency, L2 flush per sample. Baseline / patched order alternates across five rounds.

Tokens	Baseline, 5 rounds (μs)	Patched, 5 rounds (μs)	Median latency reduction
1	11.328, 11.264, 11.328, 11.264, 11.264	6.272, 6.304, 6.240, 6.240, 6.304	44.32%
8,192	50.112, 50.112, 50.112, 50.112, 50.144	16.448, 16.576, 16.576, 16.480, 16.480	67.11%
65,536	310.240, 310.240, 310.240, 310.240, 310.240	82.080, 82.112, 82.272, 82.560, 82.080	73.53%

These are operator timings, not an end-to-end throughput result. No full-model speedup is inferred from the percentages.

Checklist

• Format changed files with the repository's isort / Ruff rules; AST, whitespace and registered-test checks pass.
• Add regression coverage to registered test suites.
• Follow the SGLang code style guidance.
• Complete registered concat CUDA regression suite in upstream CI (108 passed, 0 skipped). Full-model measurements are outside this isolated-operator claim.

---

CI States

Latest PR Test (Base): ❌ Run #41299322000
Latest PR Test (Extra): ❌ Run #41299322057
Latest PR Test (AMD ROCm 10): ❌ Run #41299322005

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.

[Codoge]

@Fridge003 Ready for review. Could an authorized maintainer enable CI with /tag-and-rerun-ci and review this model-dispatch change?

The current red workflow results stopped at Block draft PR; the PR is now non-draft and satisfies the repository's MIN_BASE_SHA requirement.

At head 81955274f6, all 14 new model-helper cases pass on B300, including exact BF16/FP16 results, strided and empty inputs, 64K tokens, repeated CUDA Graph replay, and the FA3 FP8 output-dtype case. Five alternating baseline/patched operator measurements are in the description. This reuses an existing masked kernel; no end-to-end speedup or full-model accuracy result is claimed.

Focused CI coverage:

/rerun-test test/registered/kernels/ops/attention/test_concat_mla.py

This runs the existing regression cases together with the new model-dispatch coverage. The PR is independent of #47742 and #47744.

[BBuf]

[BBuf Review Agent] LGTM. Please wait for the registered concat CUDA tests to pass.

[BBuf]

/tag-and-rerun-ci

[Codoge]

@BBuf The registered concat CUDA suite has already passed in upstream CI for the current head 81955274f6363c91979dd12d05e2662621bc4346.

H100 CI job: jit-kernel-unit-test-1-gpu-large (1) checked out merge commit bab657f5422be3179be49c57b5abdf46736db294 (this head merged into c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0) and ran the complete registered file:

python3 test/registered/kernels/ops/attention/test_concat_mla.py
108 passed, 21 warnings in 8.97s

There were no failures or skips. The log explicitly includes all 14 new model-helper cases, including repeated CUDA Graph replay and the FA3 FP8 output-dtype case, along with the existing JIT/AOT regressions. The job is green, so the concat validation requested in your review is complete.

The overall Base result is red because CPU shard 16 hits the existing DCP fixture error (skips_prefill_metadata_on_idle_dp_rank), the same unrelated failure you confirmed on #47653. Neither affected DCP file is changed here. Current main now includes the missing fixture field. The remaining AMD disaggregation failure and optional-tier gates are separate from the passing concat suite.

Could you re-review for merge with this upstream test evidence? The production change is still the same dispatch guard removal; no new commit or performance claim has been added.

[BBuf]

[BBuf Review Agent] @Codoge The code is unchanged and the registered suite covers the affected dispatch, so I’m fine dropping the extra validation condition.