Route MLA RoPE KV writes through embedding API

[qedawkins]

Summary

Add an MLA-specific fused RoPE KV-cache descriptor and route it through the existing embedding.rope API. The normal fused KV path assumes MHA-style K and V buffers with the same per-head width, but DeepSeek MLA stores one combined cache row shaped as k_nope / latent width kv_lora_rank concatenated with rotated k_rope width qk_rope_head_dim. The new fused_mla_set_kv_buffer_arg makes that layout explicit, letting the RoPE kernel copy k_nope and write rotated k_rope directly into the combined row. The normal fused-KV path for DeepSeekV3 was unused and replaced here.

Testing Plan

• test/runtime/cache/test_mla_kv_buffer.py covers MLA KV set/get against torch references across bf16/fp8, cache locations, and round trips.
• test_mla_rope_set_kv_buffer_matches_reference covers the fused RoPE+MLA KV write path used by DeepSeek decode.

[borontion]

do you have a timeline/perf comparison before and after this change?

[qedawkins]

do you have a timeline/perf comparison before and after this change?

Oops forgot to add that part. From Codex:

Fused MLA RoPE KV Benchmark Results

Setup

• Before: /home/qdawkins/tokenspeed-main-clean at 6a226130
• After: /home/qdawkins/tokenspeed-mla-fused-set-kv-pr at 4065dca1
• Model: amd/Kimi-K2.5-MXFP4
• Runtime: TP4, --attention-backend mla, --quantization mxfp4, --kv-cache-dtype auto, prefix cache disabled
• GPUs: HIP_VISIBLE_DEVICES=4,5,6,7
• Note: the local tokenspeed-scheduler extension was reinstalled from the matching checkout before each before/after run because the two checkouts use different scheduler ABI fields.

Isolated Kernel Benchmark

Command shape: synthetic BF16 MLA RoPE + MLA KV write, Kimi/DeepSeek dimensions (nope_dim=512, rope_dim=64), 200 warmup iterations and 3000 measured iterations per row.

The fused path is compared against both the pre-PR Torch RoPE + KV write sequence and the generic Triton RoPE + KV write sequence. Values below are median milliseconds.

tokens	q heads	pre-PR unfused ms	generic Triton unfused ms	fused ms	speedup vs pre-PR	speedup vs generic
1	8	0.11296	0.06444	0.03168	3.57x	2.03x
1	16	0.11360	0.06548	0.03100	3.66x	2.11x
4	8	0.11676	0.06548	0.03124	3.74x	2.10x
4	16	0.11444	0.06492	0.03120	3.67x	2.08x
16	8	0.11680	0.06576	0.03160	3.70x	2.08x
16	16	0.11480	0.06524	0.03108	3.69x	2.10x
128	8	0.12572	0.06540	0.03120	4.03x	2.10x
128	16	0.12448	0.06480	0.03092	4.03x	2.10x
512	8	0.11396	0.06588	0.03020	3.77x	2.18x
512	16	0.11500	0.06524	0.03012	3.82x	2.17x

Raw data: /home/qdawkins/fused_kv_kernel_benchmark_after.json

E2E Decode Benchmark

Command shape: in-process Engine, synthetic token IDs, input_len=64, output_len=16, 5 warmup requests and 100 measured requests.

The controlled after run disables prefill graph to match the before behavior. The default after run is included for reference only because the branch default also engages prefill graph.

run	prefill graph	tok/s	avg latency s	p95 latency s	change vs before
before origin/main	disabled/effectively off	37.62	0.4252	0.4281	baseline
after controlled	disabled	49.09	0.3259	0.3302	+30.5% tok/s, -23.4% avg latency
after default	enabled	55.05	0.2906	0.2917	+46.3% tok/s, -31.7% avg latency

Trace Comparison

Trace shape: one warmup request, then one profiled request with input_len=64, output_len=4. The after trace uses --disable-prefill-graph. The intent here is to estimate the wall clock time saved during real execution of Kimi TP4 on mi355.

Single-request measured latency from the profiled runs:

run	tok/s	request latency s
before trace	20.55	0.1945
after trace	24.70	0.1619

Rank-0 kernel-launch changes in the profiled request:

kernel/category	before count	after count	note
set_mla_kv_buffer_kernel	305	61	decode-side standalone KV write removed; remaining 61 are prefill-side
_mla_rope_set_kv_buffer_kernel	0	244	new decode-side fused RoPE + MLA KV write, 61 layers x 4 decode tokens
_rope_apply_kernel	0	61	remaining prefill-side generic RoPE
Torch cat copy kernels	610	0	eliminated from the profiled request
Torch BF16 copy kernels	670	60	mostly eliminated
Torch direct copy kernels	1890	975	reduced, but not eliminated
Torch binary elementwise kernels	1280	60	decode-side RoPE math eliminated
Torch add kernels	971	361	reduced
Torch neg kernels	610	0	eliminated from the profiled request

Immediate pre-MLA decode window, rank 0:

window	before median us	after median us	note
last 4 kernels before _mla_decode_kernel	28.10	22.88	conservative local span; includes neighboring copy/producer work
old unfused helper sequence visible before _mla_decode_kernel	about 98 us	replaced by fused path	before used Torch RoPE math, cat/copy kernels, and standalone KV write

Interpretation: the isolated benchmark is the cleanest measurement of this specific fused operation. The trace confirms that the decode-side standalone RoPE/cat/copy/KV-write sequence no longer appears: 244 decode-side standalone KV writes become 244 fused RoPE+KV kernels, and the large Torch RoPE helper launch count disappears from the profiled request.