[clintg6]

Motivation

Kimi-K3 uses two attention representations during prefill at TP8 on MI355X:

• Fresh/zero-prefix MHA: 12 heads with QK/V dimensions 192/128.
• Cached-prefix absorbed MLA: 12 query heads, one KV head, and QK/V
  dimensions 576/512.

Current Triton prefill leaves performance on the table for both shapes:

1. The exact K3 kernels use one pipeline stage and natural-exponential
   online softmax.
2. Absorbed 576/512 prefill processes one full 512-wide value accumulator,
   increasing register pressure.
3. FP8 cached-prefix P·V can underflow small softmax probabilities.
4. Aiter supports native FP8 zero-prefix prefill, but Triton lacks an
   equivalent opt-in path.

Changes

K3 prefill tuning

• Use two pipeline stages for exact gfx950 K3 192/128 and N32 576/512
  prefill kernels.
• Fold log2(e) into the attention scale and use exp2 for unmodified K3
  logits.
• Convert stored LSE back to natural-log units.
• Preserve current compact ragged-grid scheduling and the generic
  Triton-3.7 fallback.

4×128 absorbed MLA kernel

Add a dedicated split-dimension kernel for cached-prefix K3 MLA:

Q/K latent: 4 × 128
Q/K tail:       64
V/output:   4 × 128

The kernel maintains one online-softmax state and four independent FP32
value accumulators before concatenating the normalized 128-wide outputs.

Dispatch is restricted to:

• gfx950;
• 12 query heads and one KV head;
• QK/V dimensions 576/512;
• BF16 current Q/K/V and output;
• BF16 cached prefix with unit scales, or opt-in FP8 cached prefix with
  supported descales;
• causal attention;
• page size 1;
• contiguous last dimensions;
• no custom mask, sliding window, logit cap, XAI temperature, sink, LSE,
  score modifier, auxiliary tensor, or skipped stage.

Unsupported calls immediately use the existing generic kernel.

FP8 prefill

Add:

SGLANG_TRITON_FP8_PREFILL_ATTN=1

This is enabled by default. Set SGLANG_TRITON_FP8_PREFILL_ATTN=0 to disable
native FP8 prefill.

On exact K3 gfx950 shapes:

• Zero-prefix 192/128 Q/K/V are cast from BF16 to FP8 before launch,
  matching Aiter behavior.
• Absorbed 576/512 prefill keeps current-chunk Q/K/V in BF16 while using
  the FP8 cached prefix.
• Softmax probabilities are multiplied by FP8_MAX before FP8 P·V and
  the accumulator contribution is divided by the same factor.
• Non-unit cached-prefix K/V descales are supported.

The default-on path achieved 98% GSM8K accuracy in model-level validation.

Performance

MI355X measurements used a device warmup and HIP-graph replay.

Against unmodified current main:

K3 prefill case	Current main	This change	Speedup
Zero-prefix BF16, 8K	0.726 ms	0.415 ms	1.75×
Zero-prefix FP8, 8K, casts included	0.726 ms	0.420 ms	1.73×
Absorbed BF16, 8K prefix + 8K extend	8.658 ms	4.133 ms	2.09×
Absorbed FP8, 8K prefix + 8K extend	6.770 ms	3.687 ms	1.84×

FP8 zero-prefix compared with the optimized BF16 path:

Length	BF16	FP8, casts included	Result
8K	0.414 ms	0.421 ms	No difference
16K	1.500 ms	1.400 ms	FP8 1.07× faster
32K	6.255 ms	5.145 ms	FP8 1.22× faster

The default-on policy intentionally follows Aiter and enables native FP8 at
every length despite the small 8K crossover loss.

End-to-end serving

Throughput uplift is computed as the geometric mean of the per-concurrency
ratios at concurrency 2, 4, 8, and 16.

Kimi-K3 TP8, ISL/OSL 8192/1024:

Concurrency	Baseline tok/s	This change tok/s	Uplift
2	874.18	879.54	+0.6%
4	1529.69	1586.17	+3.7%
8	2554.27	2586.82	+1.3%
16	3892.75	3940.95	+1.2%
Geomean	—	—	+1.70%

Kimi-K3 TP8, ISL/OSL 68K/350:

Concurrency	Baseline tok/s	This change tok/s	Uplift
2	6092.96	6910.34	+13.4%
4	7060.28	8185.23	+15.9%
8	7655.63	9017.19	+17.8%
16	7973.53	9540.03	+19.6%
Geomean	—	—	+16.67%

The long-context gain is predominantly improved prefill/TTFT; median ITL is
effectively unchanged in both comparisons.

Numerics

• BF16 absorbed split versus current-main generic output:
  relative L2 approximately 9.5e-5.
• FP8 absorbed split versus the corrected generic FP8 path:
  relative L2 approximately 6.6e-5.
• Zero-prefix FP8 versus BF16:
  relative L2 approximately 3.4-3.5%.
• Kimi-K3 model-level GSM8K accuracy with FP8 prefill enabled: 98%.

Additional long-context logits and workload validation remain recommended.

Platform scope

All specialized execution paths require gfx950 and exact K3 shapes.

On NVIDIA and other AMD architectures:

• the split-dimension dispatcher returns false;
• FP8 Q/K/V casting is disabled;
• FP8 probability scaling is disabled;
• K3 exp2 and two-stage specializations are disabled;
• existing generic attention behavior remains active.

Validation

• Ruff 0.15.1 passed with the repository's configured rule selection.
• Black 26.1.0 passed.
• Isort 7.0.0 passed.
• git diff --check passed.
• Kimi-K3 GSM8K accuracy with FP8 prefill enabled: 98%.
• Four focused K3 BF16/FP8 parity and dispatch tests pass on gfx950.
• The focused kernel tests are registered in the per-PR
  stage-b-test-1-gpu-small-amd-mi35x suite.
• Generic Triton-3.7 fallback tests passed after the final gfx950-only gate.

Recommended E2E matrix:

• ISL/OSL 8192/1024 and 68000/350;
• concurrency 2, 4, 8, 16, 32, and 64;
• BF16 and FP8 KV cache;
• flag off/on with identical prompts, seeds, request counts, and cache state;
• model-level accuracy and logits comparison.

Files changed

python/sglang/kernels/ops/attention/extend_attention.py
python/sglang/kernels/ops/attention/extend_attention_split_dim.py
python/sglang/srt/environ.py
test/registered/unit/layers/attention/test_triton_mla_prefill_gfx950.py

---

CI States

Latest PR Test (Base): ✅ Run #33842442737
Latest PR Test (Extra): ❌ Run #33842442510
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33842442551

[1am9trash]

Hi, @clintg6

In long-context cases (chunk-prefill or radix-cache on), the prefix sweep dominates split_dim_absorbed_extend_attention_fwd. It steps through the cached context BLOCK_N at a time, so a larger block halves its iterations. Since our prefix sweep uses fp8 matmul while the extend sweep uses bf16, we can decouple their block sizes and set the former to twice the latter, without increasing LDS usage or causing an occupancy regression.

Based on your PR, I added the change here. Both the radix-cache-on and the chunked-prefill paths benefit.
Branch: decouple block size commit

Test Config: radix-cache off, chunked-prefill-size 16384, conc 2/4/8, isl/osl 68k/350.
Kernel: 1.2–1.35× depending on prefix length. E2E:

conc	TTFT PR (ms)	TTFT new (ms)	gain	TTT PR (tok/s)	TTT new (tok/s)	gain
2	9378	8657	7.7%	7178.28	7585.71	5.7%
4	15991	14703	8.1%	8566.11	9178.24	7.1%
8	31193	28516	8.6%	9215.54	9994.40	8.5%

[HaiShaw]

NV CI is all good, ROCm pre-existing condition