[kernel] Wire MLA split-KV decode

[rohash123]

Why this PR exists

This is an implementation for the Phase 3 item in #360: MLA-specific long-context partition/reduce decode.

#360 says the regular paged-attention split-KV work in #437 is useful reference material, but it did not cover MLA. This PR adds the MLA version.

In plain terms: when decoding one token from a very long prompt, MLA has to read a lot of cached context. Before this PR, one small Metal launch did that read mostly as one long scan. This can leave the GPU under-used at low concurrency. This PR lets MLA split that long scan into several independent pieces, run those pieces in parallel, and then combine the partial answers into the same final answer.

What changed

• paged_mla_attention can now run over a slice of the context instead of always scanning the full context in one pass.
• A new paged_mla_attention_reduce Metal kernel combines those slices into the final output.
• The reduce step uses each slice's softmax summary value (LSE, log-sum-exp) so the result matches normal attention math instead of doing a separate softmax per slice.
• Partition 0 writes directly into the final output buffer, so we avoid one extra temporary output slice.
• Host dispatch in paged_ops.cpp now chooses the split path only for conservative cases:
  • pure decode only
  • long enough context
  • small decode launch
  • 2 to 8 partitions
  • bounded temporary memory use

This keeps the existing MLA cache layout and runtime contract. It does not change scheduling, prefill, prefix cache behavior, or non-MLA attention.

Correctness

This PR adds targeted coverage to the existing MLA kernel test suite. The new split-KV cases compare partitioned MLA decode against the dense MLX reference and cover:

• 32K context, one request, two heads per threadgroup
• 32K context, two decode requests in the same batch
• 81K context, one request, several partitions

I also reran the existing MLA kernel and MLA paged-backend tests, so the old single-pass MLA path and wrapper routing coverage are still exercised.

Local validation on the clean PR branch:

python -m vllm_metal.metal.build
python -m pytest tests/test_mla_kernel.py tests/test_mla_paged_backend.py -q

75 passed, 1 warning

Wrapper benchmark

#360 asks for performance claims that include wrapper overhead, not only isolated kernel time. I re-ran the comparison through MLAPagedAttentionWrapper.__call__ rather than calling the Metal kernel directly.

This benchmark includes the wrapper work around the kernel: query projection, KV projection for the new token, RoPE, cache write, block-table/context metadata conversion, Metal attention call, unembed, and output projection.

Hardware: Apple M5 Pro, 20-core GPU, 24 GB, macOS 26.5.2.

Setup: synthetic absorbed-MLA decode wrapper, one decoded token, fp16, kv_lora_rank=512, qk_rope_head_dim=64, block size 32. The benchmark compared the new automatic split path with the old single-pass path in the same Python process, using the same tensors, shuffled order, median latency, and explicit mx.synchronize(). For the "old single-pass" numbers, I temporarily disabled the split decision in the local benchmark checkout so everything except the dispatch choice stayed identical.

The table reports the median across five shuffled trials. Each trial measured the median of 50 timed wrapper calls per mode after warmup.

Context	Selected split	New split path	Old single-pass path	Change
32K	2 partitions	0.973 ms	1.387 ms	+42.6% faster
49K	2 partitions	1.455 ms	1.873 ms	+28.8% faster
65K	2 partitions	1.437 ms	2.268 ms	+57.8% faster
81K	5 partitions	1.605 ms	2.870 ms	+78.8% faster
98K	6 partitions	2.247 ms	3.436 ms	+52.9% faster
131K	4 partitions	3.019 ms	4.780 ms	+58.3% faster

The exact percentages still move around across runs on my local machine, especially in this synthetic harness. The robust signal is that the split path is faster across the long-context, low-concurrency cases I tested. The chart below shows absolute latency with trial-to-trial min/max whiskers instead of a single-run speedup chart.

Raw benchmark summary: https://gist.github.com/rohash123/4ecf29d05091ec9053dbfd0fba83ea18

Serving sanity check

I also ran one small real-server check through vllm serve + vllm bench serve to make sure this path works outside the wrapper harness.

This is intentionally a sanity check, not the main performance claim: longer serving runs at 32K+ context put too much memory pressure on my local Mac. The useful signal here is that an MLA model loads, allocates paged MLA KV cache, serves through the OpenAI-compatible endpoint, and completes generation successfully with this PR.

Server:

vllm serve <local DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx> \
  --served-model-name deepseek-coder-v2-lite-mlx \
  --host 127.0.0.1 \
  --port 8011 \
  --max-model-len 8448 \
  --max-num-seqs 1 \
  --max-num-batched-tokens 1024 \
  --no-enable-prefix-caching

Environment:

VLLM_PLUGINS=metal
VLLM_METAL_USE_PAGED_ATTENTION=1
VLLM_METAL_MLA_KERNEL=1
VLLM_METAL_MEMORY_FRACTION=0.55

Client:

vllm bench serve \
  --backend openai \
  --base-url http://127.0.0.1:8011 \
  --model deepseek-coder-v2-lite-mlx \
  --dataset-name random \
  --random-input-len 8192 \
  --random-output-len 8 \
  --random-range-ratio 0.0 \
  --num-prompts 1 \
  --max-concurrency 1 \
  --request-rate inf \
  --ignore-eos \
  --temperature 0

Result:

Metric	Value
Successful requests	1
Failed requests	0
Input tokens	8192
Generated tokens	8
Mean TTFT	18628.88 ms
Mean TPOT	270.93 ms
Mean E2E latency	20525.39 ms
Total token throughput	399.47 tok/s

Scope notes

MLX is still the default path for MLA. This PR only affects users who opt into the Metal MLA kernel with VLLM_METAL_MLA_KERNEL=1.

When that opt-in is enabled, the C++ side chooses between the existing one-shot Metal path and the new split path based on the request shape.

This PR does not change the MLA cache format, request scheduling, or upstream vLLM MLA behavior.

I also ran the wrapper-level grid from the MLA kernel RFC comment: heads {16, 64, 128}, batch {1, 8, 32}, context {128, 2048, 8192}, and both fp16/bf16. This compares the MLX default path with the existing opt-in Metal MLA path (VLLM_METAL_MLA_KERNEL=1).

Important: this grid stops at 8K context, while this PR's new split path only starts at 32K+ context. So these numbers do not measure the new split/reduce optimization. They are included as a broader default-readiness check for Metal MLA overall.

Result: Metal wins 24/54 cells overall, but not the full grid. This supports keeping MLX as the default.

Slice	Metal wins	Median Metal vs MLX
fp16	13/27	-1.1%
bf16	11/27	-5.5%
H=16	14/18	+25.8%
H=64	7/18	-4.3%
H=128	3/18	-10.3%
ctx=128	11/18	+4.4%
ctx=2048	8/18	-2.3%
ctx=8192	5/18	-19.5%

Raw grid CSV: https://gist.githubusercontent.com/rohash123/4ecf29d05091ec9053dbfd0fba83ea18/raw/ea5953f827b1a86cbb3371429117a1949125aad1/mla_production_grid_wrapper.csv

Progresses #360.

[LxYuan0420]

This policy is too hard-coded for shared dispatch. Before landing, prove 32K+ split Metal beats the default MLX path, then put the thresholds behind a small device-aware split-plan helper.

[LxYuan0420]

These parity tests can still pass if split dispatch never runs. Assert the selected split plan, cover bf16/block16, and add one mixed long/short batch.

[LxYuan0420]

This scratch guard looks unreachable under the existing bounds. Delete the dead calculation and keep the occupancy and partition-count gates.

[LxYuan0420]

docs/configuration.md:15 :: This is stale now. The Metal MLA path is no longer only single-pass; document the 32K+ split/reduce path and its eligibility.

[LxYuan0420]

Requesting changes for now. After the policy/test cleanup, please share rerunnable 32K+ proof from the rebased head: MLX vs single-pass Metal vs split Metal, with device, dtype, heads, batch, context, and repeated latency numbers.

[reneleonhardt]

@rohash123 If M4 Max 128 GB would help you get that 32K+ proof, just prepare a bench script for me, happy to help.

[LxYuan0420]

Closing this for now and keeping #360 open.

The problem is real, but this branch is too large and still does not prove a win over default MLX. A restart should preserve the existing one-pass Metal path and include rerunnable 32K+ MLX vs single-pass vs split evidence.

[reneleonhardt]

@rohash123 If you eventually need me to run your 32K+ benchmark, just let me know.
If the win is real, maybe the solution can be trimmed into optimal shape, thank you for all your MLX optimization work!

Just curious, there are amazing MLX optimizations in youssofal/MTPLX , unsloth/ and jundot/omlx, could any of them help with MLA split-KV decode?