[Kimi-K3 MLA TP8] MLA performance

[zejunchen-zejun]

Kimi-K3 MLA decode kernel — input specification

Measured on a live vLLM-ATOM server on 8x MI355X (gfx950): every argument of
aiter.mla.mla_decode_fwd on the decode path, with shape, dtype, stride and
size. Values are from the tail of a 350-token generation (kv_len ≈ 68,43x),
not from the first decode steps.

1. Environment

item	value
model	KimiK3ForConditionalGeneration, 93 layers
GPUs	8x MI355X (288 GiB HBM3E), --tensor-parallel-size 8
execution	eager, no CUDA graphs, no torch.compile
online weight quant	off
KV cache dtype	fp8 → float8_e4m3fn
vLLM block size	128
max model len	73728
concurrency	16 and 24
workload	63911-token shared prefix + 4089 unique input tokens, 350 output tokens, prefix caching on
captured layers	83, 87, 91, 92 (last four full-attention layers), TP rank 0

component	version	commit
aiter	amd-aiter 0.1.20.dev41+g334fd72ad	334fd72ad6e69d16785efc97cb8b6b4aec3c2a92 (main, 2026-08-03)
vLLM	0.25.2.dev0+g752a3a504.rocm724	752a3a504485790a2e8491cacbb35c137339ad34
ATOM	0.1.6rc1.dev219	452a0aca281b0317fde12f98e55a803e8c88e753 (k3-rmsnorm-fusion)
PyTorch	2.10.0+rocm7.2.4	3d3aa833
Triton	3.7.0+amd.rocm7.2.0	89002410
HIP / ROCm	7.2.53211	—
GPU	gfx950:sramecc+:xnack-	—

2. K3 MLA geometry

K3 is a hybrid stack: 69 KDA (gated-delta) layers and 24 full-attention MLA
layers. MLA runs in absorbed (latent) form, so the kernel attends over a
single 576-dim latent entry per token rather than per-head K/V.

item	value
MLA layers, 0-based	3, 7, 11, ... , 79, 83, 87, 91, 92 (24 total; last two adjacent)
num_attention_heads (global)	96
heads per GPU at TP8	12
kv_lora_rank	512
qk_nope_head_dim / qk_rope_head_dim	128 / 64
qk_head_dim	192
v_head_dim	128
KV entry width	576 = 512 latent + 64 rope
sm_scale	0.07216878364870322 = 1/sqrt(192)

All four captured layers produced byte-identical argument signatures, so one
table describes all 24 MLA layers.

3. Entry point

aiter.mla.mla_decode_fwd, aiter/mla.py:270:

mla_decode_fwd(q, kv_buffer, o, qo_indptr, kv_indptr, kv_indices,
               kv_last_page_lens, max_seqlen_q, page_size=1, nhead_kv=1,
               sm_scale=None, logit_cap=0.0, num_kv_splits=None,
               num_kv_splits_indptr=None, work_meta_data=None, work_indptr=None,
               work_info_set=None, reduce_indptr=None, reduce_final_map=None,
               reduce_partial_map=None, q_scale=None, kv_scale=None,
               intra_batch_mode=False, return_logits=False, return_lse=False,
               g_kv_indptr=None, cp_world_size=1, cp_rank=0)

Called from AttentionForVllmMLA._forward_decode. Two defaults are not
overridden and matter: page_size stays 1 (see 4.2) and nhead_kv stays
1 (single latent KV head).

The binary dispatched on this configuration:

aiter::mla_a8w8_qh16_qseqlen1_gqaratio16_ps
aiter_meta/hsa/gfx950/mla/mla_a8w8_qh16_qseqlen1_gqaratio16_ps.co

a8w8 = fp8 Q and fp8 KV, qh16 = 16 query heads, qseqlen1 = single-token
decode, ps = persistent. The qh16 specialization is why 12 heads get padded
to 16 (see 4.1).

The persistent (split-KV) path is always taken: the concurrency gate in
_use_persistent_mla_decode (aiter/mla.py:245) only downgrades to the
non-persistent kernel for the bf16 Q + bf16 KV 16-head profile, and this
deployment is fp8/fp8.

q is produced by aiter.fused_qk_rope_concat_and_cache_mla, which also writes
the step's new KV entry into the same cache the kernel reads. The 512-dim latent
o is then mapped back to v_head_dim by a separate BMM. Those surrounding ops
are out of scope here.

4. Arguments

contiguous flags a strided view; size is the tensor's full byte footprint.
Both q and o are contiguous.

4.1 Concurrency 16

#	argument	shape	dtype	stride	contiguous	size
0	q	(16, 16, 576)	float8_e4m3fn	(9216, 576, 1)	yes	144.0 KiB
1	kv_buffer	(5283840, 1, 1, 576)	float8_e4m3fn	(576, 576, 576, 1)	yes	2.8 GiB
2	o	(16, 16, 512)	bfloat16	(8192, 512, 1)	yes	256.0 KiB
3	qo_indptr	(17)	int32	(1,)	yes	68 B
4	kv_indptr	(17)	int32	(1,)	yes	68 B
5	kv_indices	(4718592)	int32	(1,)	yes	18.0 MiB
6	kv_last_page_lens	(16)	int32	(1,)	yes	64 B
7	max_seqlen_q	non-tensor				1
8	sm_scale	non-tensor				0.07216878364870322
9	q_scale	scalar	float32	()	yes	4 B
10	kv_scale	scalar	float32	()	yes	4 B
11	work_meta_data	(2)	uint64	(1,)	yes	16 B
12	work_indptr	(257)	int32	(1,)	yes	1.0 KiB
13	work_info_set	(319, 8)	int32	(8, 1)	yes	10.0 KiB
14	reduce_indptr	(65)	int32	(1,)	yes	260 B
15	reduce_final_map	(64, 2)	int32	(2, 1)	yes	512 B
16	reduce_partial_map	(319)	int32	(1,)	yes	1.2 KiB
17	return_lse	non-tensor				False

Batch: num_decodes = 16, num_decode_tokens = 16, num_prefills = 0 (pure
decode). num_heads 12 before padding, 16 after. fold_factor = None,
use_persistent_mode = True, dcp_world_size = 1.

Per-request KV length over the run: 68,090 → 68,437 tokens
(kv_len_sum 1,089,472 → 1,094,960), q_len always 1.

Late decode steps (step counts MLA-layer forwards, so consecutive model steps
are 24 apart):

step	kv_len min	kv_len max	kv_len sum
16677	68,430	68,434	1,094,912
16701	68,431	68,435	1,094,928
16725	68,432	68,436	1,094,944
16749	68,433	68,437	1,094,960

4.2 Concurrency 24

#	argument	shape	dtype	stride	contiguous	size
0	q	(24, 16, 576)	float8_e4m3fn	(9216, 576, 1)	yes	216.0 KiB
1	kv_buffer	(5283840, 1, 1, 576)	float8_e4m3fn	(576, 576, 576, 1)	yes	2.8 GiB
2	o	(24, 16, 512)	bfloat16	(8192, 512, 1)	yes	384.0 KiB
3	qo_indptr	(25)	int32	(1,)	yes	100 B
4	kv_indptr	(25)	int32	(1,)	yes	100 B
5	kv_indices	(4718592)	int32	(1,)	yes	18.0 MiB
6	kv_last_page_lens	(24)	int32	(1,)	yes	96 B
7	max_seqlen_q	non-tensor				1
8	sm_scale	non-tensor				0.07216878364870322
9	q_scale	scalar	float32	()	yes	4 B
10	kv_scale	scalar	float32	()	yes	4 B
11	work_meta_data	(2)	uint64	(1,)	yes	16 B
12	work_indptr	(257)	int32	(1,)	yes	1.0 KiB
13	work_info_set	(319, 8)	int32	(8, 1)	yes	10.0 KiB
14	reduce_indptr	(65)	int32	(1,)	yes	260 B
15	reduce_final_map	(64, 2)	int32	(2, 1)	yes	512 B
16	reduce_partial_map	(319)	int32	(1,)	yes	1.2 KiB
17	return_lse	non-tensor				False

Batch: num_decodes = 24, num_decode_tokens = 24, num_prefills = 0. Same
head padding, persistent mode and fold settings as c16.

Per-request KV length over the run: 68,089 → 68,438 tokens
(kv_len_sum 1,634,250 → 1,642,410), q_len always 1.

Late decode steps:

step	kv_len min	kv_len max	kv_len sum
33549	68,426	68,435	1,642,338
33573	68,427	68,436	1,642,362
33597	68,428	68,437	1,642,386
33621	68,429	68,438	1,642,410

5. Notes for optimization

5.1 12 real heads padded to 16 — 25% of the Q tile is zeros

_MLA_MIN_HEADS = 16 in atom/model_ops/attention_mla.py:109 ("AITER MLA
kernels require at least 16 attention heads"). K3 at TP8 has 96/8 = 12
heads. 16 is not a multiple of 12, so head-repeat cannot be used and the code
zero-pads (head_pad = 4):

# atom/model_ops/attention_mla.py:342
if self.head_pad > 0:
    return torch.nn.functional.pad(q, (0, 0, 0, self.head_pad))

So q arrives with 16 heads of which only the first 12 carry data, and o
comes back with 16 of which 4 are discarded. Every QK^T and PV MFMA tile wastes
25% of its issue slots. A kernel accepting 12 heads natively removes this
directly.

5.2 The kernel runs with page_size = 1

Although vLLM is configured with --block-size 128, the decode side passes a
token-major view — kv_cache is (5283840, 1, 576), one entry per token
slot, not per block:

kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
mla_decode_fwd(q, kv_buffer.view(-1, 1, 1, q.shape[-1]), o, ...)

Consequently kv_indices is a per-token index list, kv_last_page_lens is
all ones, and kv_indptr[i+1] - kv_indptr[i] is request i's KV length in
tokens (~68,400), not its block count. At c24 the kernel gathers 1,642,410
individual token indices per layer per step (6.3 MiB of int32 index traffic
per layer, ~150 MiB per decode step across 24 layers).

Block-granular addressing (page_size 128, matching the allocator) would cut
index traffic by 128x and allow contiguous 128-token loads instead of a gather.

Worth confirming on your side: the persistent work partitioner is built by
get_mla_metadata_v1 with page_size=self.block_size and
kv_granularity=max(self.block_size, 16) from kv_cache_spec.block_size, while
the decode call itself uses page_size=1. The kernel side is unambiguously
page_size 1 (token-count indptr deltas, all-ones last-page lens); whether the
work-splitting metadata used the same granularity is unchecked.

5.3 Strongly memory-bound

Per MLA layer, per GPU, at the captured late steps:

	c16	c24
gathered KV token slots	1,094,960	1,642,410
KV bytes read (fp8, 576 B/entry)	601 MiB	902 MiB
index bytes read (int32)	4.2 MiB	6.3 MiB
Q in / O out	144 KiB / 256 KiB	216 KiB / 384 KiB
FLOPs (padded 16 heads)	38.1 GFLOP	57.2 GFLOP
FLOPs actually useful (12 heads)	28.6 GFLOP	42.9 GFLOP
arithmetic intensity	~60 FLOP/byte	~60 FLOP/byte
per decode step, all 24 MLA layers	15.2 GB	22.9 GB
HBM-bound floor @ 8 TB/s	1.91 ms/step	2.86 ms/step

MI355X machine balance is roughly 625 FLOP/byte for dense FP8, so at ~60
FLOP/byte this kernel is an order of magnitude away from compute bound.
Effort should go to bytes moved and load efficiency, not MFMA utilization.

5.4 Possible cross-request KV reuse (unverified)

All requests share a 63,911-token prefix with prefix caching on, so the block
tables point at the same physical KV slots for the shared portion. Derived
from the configuration (499 shared 128-token blocks plus a ~4.4k-token unique
tail per request):

• c16: ~137k distinct slots vs 1,095k gathered → ~8.0x redundancy
• c24: ~173k distinct slots vs 1,642k gathered → ~9.5x redundancy

At 576 B/slot that is 75–95 MiB of distinct data, within MI355X's 256 MiB LLC.
If the kernel gathers per-request independently it re-reads the shared prefix
once per request; blocking the KV loop across requests (cascade / shared-prefix
attention) could cut HBM traffic substantially.

This ratio is derived from configuration, not measured — treat it as a
hypothesis to confirm.

5.5 Other fixed properties

• max_seqlen_q = 1 — pure decode, no speculative decoding or MTP.
• return_lse = False, dcp_world_size = 1 — no decode context parallelism, so
  no LSE correction path.
• Persistent split-KV scheduler active: work_indptr(257) = 256 work
  partitions, work_info_set(319, 8), reduce metadata reduce_indptr(65) /
  reduce_final_map(64, 2) / reduce_partial_map(319). The 64 is
  max_num_seqs, not the live batch.
• fold_factor = None — head-folding path not taken.
• KV cache allocation is 5,283,840 slots x 576 B = 3.04 GB per layer per
  GPU, ~73 GB across the 24 MLA layers.
• KV lengths are near-uniform across requests (within ~10 tokens), so
  per-request work distribution is naturally balanced.
• Eager mode was used for clean attribution, so no latency figures are quoted;
  they would not represent production.