perf(mla): chunk the non-FP4 gather_kv_b_proj over KV

[zejunchen-zejun]

The MLA cached-prefix gather expands the paged latent into full per-head K/V, once per MLA layer, over the WHOLE cached context -- its cost tracks the context length and not the new tokens the forward is actually for. That makes it the kernel a high prefix-cache hit rate leans on hardest: a hit is exactly the case where the context is huge and the new-token count is small.

The FP4 impl already partitions (batch, head, KV chunk). The non-FP4 impl, which is what every bf16 and per-token-FP8 kv_b_proj lands on -- i.e. everything that is not an FP4 checkpoint -- still launched (batch x head) programs and walked the context as a serial loop inside each one. On Kimi-K3 at TP8 that is 12 workgroups on a 256-CU part, 4.7% of the device, and it measured 0.18-0.21 TB/s against a ~8 TB/s peak.

Lifting the loop onto the grid is the whole change; the body is unmodified apart from the indent, since each chunk already initialised its own accumulators and stored its own slice with no cross-chunk state. A program past the chunk count of the sequence it belongs to returns immediately, the same guard the FP4 path uses, so a grid sized from the longest sequence stays correct for a ragged batch.

Kimi-K3, MI355X, TP8, 12 local heads, fp8 paged cache, page_size 1:

ctx	before	after	speedup
4096	0.180 ms	0.038 ms	4.7x
32000	1.247 ms	0.227 ms	5.5x
65536	3.015 ms	0.442 ms	6.8x
109440	5.003 ms	0.711 ms	7.0x

0.21 -> 1.16 TB/s at 32k. Output is bit-identical across seven shapes, including context lengths that are not a multiple of the chunk: max abs diff is exactly 0.

End to end on the SemiAnalysis cc-traces agentic replay, Kimi-K3 + DSpark(2), concurrency 8, 3600s, everything else held fixed:

metric	before	after	delta
throughput/chip	3354.82 tok/s	3565.44 tok/s	+6.3%
P90 E2E normalized	25.97 tok/s/user	30.14 tok/s/user	+16.1%
TTFT mean	4107 ms	2464 ms	-40.0%
TTFT p90	63166 ms	48290 ms	-23.6%

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4776 --add-label <label>

[Copilot]

Pull request overview

This PR increases GPU occupancy and throughput for MLA cached-prefix gather_kv_b_proj by chunking the non-FP4 path over the KV dimension, matching the existing FP4 partitioning strategy. This shifts work from an in-kernel serial loop to a larger Triton launch grid so cost scales better with long cached contexts when only a small number of new tokens are processed.

Changes:

• Change the host-side launch grid to include a KV-chunk axis sized from k_prefix’s row count (valid output tokens), not kv_indices capacity.
• Update the non-FP4 Triton kernel to use one program per (batch, head, KV chunk) and add an early-exit guard for out-of-range chunks.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
aiter/ops/triton/gather_kv_b_proj.py	Updates the launch grid to include a KV chunk dimension for both FP4 and non-FP4 paths, using k_prefix token count as the upper bound.
aiter/ops/triton/_triton_kernels/gather_kv_b_proj.py	Refactors non-FP4 implementation to map KV chunking onto the Triton grid (one program per batch/head/chunk) and removes the per-program serial loop over chunks.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.