Kimi-K3 plugin: LMCache MP offload with the KDA state riding the PAGE object by PerryZhang01 · Pull Request #3117 · ROCm/ATOM

PerryZhang01 · 2026-11-09T11:06:40Z

Ready for review. No blocking issues.

Kimi-K3 on the vLLM plugin could offload only in-process; the MP backend refused it. This follows #2598: the KDA state at a boundary rides the PAGE object as a native-state image, so the KV and the state are one LMCache object.

Dev record

Found: vLLM lays every KV cache group out layer-outer from one allocation, so mamba layer j aliases the j-th attention layer it allocates. On K3 the three 23-layer KDA groups share 23 of the 29 MLA layers, so a boundary's state is one PAGE unit per group (61,046,784 B per rank) and needs no copy.
Decided: the worker checks that aliasing on the registered tensors and refuses to start when it does not hold; the aliased layers publish first, since PAGE planes are otherwise in name order. Check this first.
Decided: only Kimi-K3 rides, gated on the kimi_k3 offload layout (its KDA layers report GDN_ATTN, as Qwen's GDN does). V41StateMPScheduler is now StateRideMPScheduler, shared by both models.
Surprised: on GPU the name order and vLLM's allocation order differ (DSpark); the check refused the first boot.
Left undone: K3 under DSpark on main also needs vllm-atom: GLM-5.3-Flash prefix caching and MTP on the plugin backend #3137's spec_state_indices_in_tensor default to start, with or without this PR.

Named result

Kimi-K3 offloads over the LMCache MP server, given --enable-prefix-caching, --mamba-cache-mode align, kv_load_failure_policy: recompute and an LMCache chunk that is a multiple of 1536. The in-process K3 path is unchanged.

Gates

branch @ ceabea5: K3, V4.1 and planner plugin tests 109 passed, 1 failed (fails the same on main). ruff, black clean.
K3 TP8 over mp: tier-resumed probe 736/736 vs recomputed 703/704; gsm8k later passes 399/400 vs 320/320.

Evidence

The unit-test failure is test_decode_keeps_are_stored_by_the_next_build_before_it_stages: HIP error: no ROCm-capable device is detected in the GPU-less test container.
test_the_state_rides_as_the_native_layout_registers_it builds the layout with build_native_state_mp_layout over CPU storage aliased the way vLLM allocates it (MLA only, MLA with a DSA index plane, and a layer that sorts first without being aliased) and checks byte by byte that image unit j holds mamba group j's state.
GPU runs and their analysis are in the round 3 comment.

🤖 Generated with Claude Code

…ache MP server Kimi-K3 on the MP backend was refused ("recurrent (mamba) state offloads in-process only"). It now takes the path DeepSeek-V4.1's checkpoints take (#2598): the KDA state at a boundary rides the PAGE object as a native-state image, so KV and state commit and restore as one LMCache object. No copy is needed. vLLM lays every KV cache group out layer-outer from byte 0 of one allocation, so mamba group layer j aliases MLA plane j. On K3 the three mamba groups of 23 KDA layers sit on the first 23 of the 29 MLA planes, and the state a group holds in block b is PAGE unit b of those planes. A boundary's image is one unit per mamba group (3 x 23 x 884,736 B). - kda_ride.py: KdaRidePlanner (caps hits to the chunk grid, names the load destination row (n - 1) // block in every mamba group or none), kda_ride_geometry (the checkpoint spec both halves derive from vLLM's KV cache config), publish_kda_ride_state (checks on the registered tensors that every KDA layer starts at, and strides with, one of the leading PAGE planes before declaring the image), and the startup refusals (align mode, prefix caching, kv_load_failure_policy=recompute). - V41StateMPScheduler had nothing V4.1-specific; renamed StateRideMPScheduler (ride_mp_scheduler.py) and shared by both models. _v41_request_of is now _request_of. - The KDA state is per rank (TP-sharded heads); the in-process K3 leg is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions · 2026-11-09T11:06:54Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 3117 --add-label <label>

valar-cop

some changes needed..

k3 tp8 mp boot log, and a two-pass or gsm8k showing a tier-resumed prefix matches the cold run?

ci green.

valar-cop · 2026-11-09T11:34:00Z

+    return span * tensor.element_size()
+
+
+def publish_kda_ride_state(transfer_tensors, tensors_by_group, geometry) -> None:


never run, so the aliasing check and the in-place kda restore have no execution behind them.. please boot k3 tp8 over mp and post a two-pass run where a tier-resumed prefix matches the cold answer

Booted K3 TP8 over mp and ran the two-pass and a probe that reads the restored region: #3117 (comment). The aliasing check refused the first boot, which ceabea5 fixes.

valar-cop · 2026-11-09T11:34:01Z

+        len(groups[group_id].layer_names) for group_id in group_ids
+    )
+    plane_bytes = _layer_page_bytes(groups[attention_group_id])
+    state_planes = max(layers_per_group)


state_planes is per mamba layer but used as attention page plane count, a dsa layer adds an index_cache plane and this breaks.. size it from the layer's movable planes instead?

Fixed in ea0074a. The spec now sums each leading attention layer's movable bytes from vLLM's specs, folding an index cache into the layer that owns it as build_kv_cache_tensors does, and the worker counts the image's planes off the published regions' layer names. The end-to-end test also runs with a DSA index plane per layer.

valar-cop · 2026-11-09T11:34:01Z

+            return False
+        if self._is_deepseek_v41:
+            return bool(self._v41_checkpoint_group_ids)
+        return bool(self._mamba_groups)


this now returns True for any mamba model on mp, but only kimi-k3 is named and tested.. intended? if so say the scope in the docstring, else gate on the kda mamba type

Gated in ea0074a. K3's KDA layers report GDN_ATTN, the same mamba type as Qwen's GDN, so the gate is the kimi_k3 offload layout (text config model_type kimi_linear). Any other mamba model on mp is refused again, and the docstring names the two models that ride.

Review of #3117: - The image covered the leading N PAGE planes with N = mamba layers per group, which counts layers, not planes. An attention layer with a DSA index cache publishes two planes, so the count fell short. The spec now sums each leading attention layer's movable bytes from vLLM's specs (an index cache folded into the layer that owns it, as build_kv_cache_tensors publishes it), and the worker counts the region planes off the published regions' layer names before it checks the aliasing. - Only Kimi-K3 rides on MP. Its KDA layers report GDN_ATTN, as Qwen's GDN layers do, so the gate is the kimi_k3 offload layout; any other mamba model on MP is refused as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valar-cop

one thing still open..

k3 tp8 mp boot log + a two-pass or gsm8k run where a tier-resumed prefix matches the cold run?

ci green.

valar-cop · 2026-11-10T09:21:40Z

+    spec = geometry.spec
+    layers = [_published_layer(page.region) for page in pages]
+    state_planes = sum(1 for layer in layers if layer < state_layers)
+    if layers[:state_planes] != sorted(layers[:state_planes]) or any(


this arithmetic only ran on the hand-built cpu pool, never a real vllm layout. also the leading-run check accepts any leading plane, so 3 distinct groups vs 3 groups aliasing the same 23 planes look the same..

…E order The first K3 TP8 boot over mp refused at registration: mamba group 0 layer 18 started on PAGE plane 23, outside the image's leading 23 planes. ATOM publishes PAGE planes by layer name, while vLLM aliases mamba layer j onto the j-th attention layer it allocates, and with DSpark the two orders differ (29 attention layers, five sort ahead of the first aliased one). kda_ride_page_order moves the attention layers the mamba layers alias to the front, in mamba-layer order, matched on addresses, and the scheduler sizes the image from vLLM's allocation order. The end-to-end test gains a layer that sorts first without being aliased; without the reorder it raises the same refusal the boot did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valar-cop

some changes needed..

ci green.

valar-cop · 2026-11-10T12:06:10Z

+    return span * tensor.element_size()
+
+
+def publish_kda_ride_state(transfer_tensors, tensors_by_group, geometry) -> None:


still open from both reviews: this has never run on a real gpu/vllm layout, tests build the aliasing by hand.. need the k3 tp8 mp boot log and a two-pass run where the resumed prefix matches cold.

PerryZhang01 · 2026-11-10T14:02:41Z

Round 3 pushed ceabea5. K3 TP8 boots over mp; tier-resumed prefixes match the cold run.

First boot refused registration: PAGE planes were in name order, vLLM aliases in allocation order. Fixed in ceabea5.
Probe: tier-resumed 736/736, recomputed 703/704.
gsm8k later passes: on 399/400, off 320/320. The miss resembles off's cold miss.

Why, and what ran

The first boot stopped with mamba group 0 layer 18 starting on PAGE plane 23. vLLM aliases mamba layer j onto the j-th attention layer it allocates, while PAGE planes were published by layer name, and with DSpark 5 of the 29 attention layers sort ahead of the aliased ones. ceabea5 publishes the aliased layers first and sizes the image in vLLM's order; the end-to-end test now has a layer that sorts first without being aliased.

Runs: this head plus #3137's one-line spec_state_indices_in_tensor default (main's K3 plugin does not start under DSpark without it); TP8, DSpark x2, FULL_AND_PIECEWISE; HBM pool of 200 to 300 blocks so every later pass resumes from the tier. off is the same server without the connector. The probe asks for a line inside the first 7680 of 9.4k tokens, the part a tier hit restores.

	on, resumed from the tier	off, recomputed
probe answers correct	736/736	703/704
probe diverging from the cold pass before token 12	0	3
probe answer tokens, max abs logprob delta	0.36	1.33
gsm8k 64-shot, later passes	399/400	320/320

The gsm8k miss is degenerate from its first token (run 1, question 10); the same question resumed from the tier is correct in the four other passes. The off arm has a comparable event with no tier involved: a cold probe answers a wrong code, its right token falling from -0.03 to -1.86. Both read as the K3 plugin's base rate at concurrency 16. No load failure, withdrawn promise or tier error in any run.

Evidence

Boot (every on run, 8 ranks each):

ATOM LMCache offload: KDA state rides the PAGE object as 3 unit(s) over the 23 planes of the leading 23 attention layers (mamba groups of 23,23,23 layers), 57443328 B of state per boundary in a 61046784 B image
ATOM LMCache offload: publishing PAGE layout (29 planes, 25657344 B/block, 200 blocks; checkpoint image 61046784 B over 3 units), tp_replication_factor=8
LMCache MP native state registered rank=R images=True native_image=61046784 units=3 groups=4 chunk=1536
ATOM LMCache offload: KDA state rides the PAGE object on group(s) 0,1,2 (mamba_block=1536, hash_block=1536, chunk=1536, prefill_tail=2)

Where the cached prefix came from, per later pass (vLLM metrics, deltas): on, probe 737,280 of 903,856 prompt tokens from the tier and 0 from HBM; on, gsm8k 837,120 of 1,007,813 from the tier and 0 from HBM; off, 0 from either.

Runs (probe passes at concurrency 16 except the first two, 32 prompts one at a time; gsm8k 80 questions at concurrency 16):

run	arm	probe	gsm8k pass 1 / later
1	on	32/32, tokens identical to cold	80/80 / 79/80
2	on	32/32, identical	80/80 / 80/80
3	off	32/32	80/80 / 80/80
4	on	2 x 96/96, identical	none
5	off	96/96, 95/96	none
6	on, 64 tokens	5 x 96/96, first divergence at token 12 (the next made-up name)	80/80 / 3 x 80/80
7	off, 64 tokens	5 x 96/96, first divergence at token 12, twice at 3	80/80 / 3 x 80/80

valar-cop · 2026-11-10T14:08:34Z

ci triage for ceabea5:

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/42830150240

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/42830150130

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/42830150357

#2365 made the shared K3 KDA forward read `kda_metadata.spec_state_indices_in_tensor` on its speculative-decode branch. vLLM's KimiK3KDAMetadata has no such field, and the plugin only defaulted its non-spec twin, so K3 on the plugin dies in profiling under DSpark (AttributeError on the first verify-shaped forward). None keeps the forward on the write slot, as the non-spec default does. The same hunk as #3137's, so the two apply cleanly in either order. The K3 TP8 LMCache MP runs on this PR ran with it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

main renamed V41StateMPScheduler to StateRideMPScheduler (#3117); the plugin test fixture keeps the token shipper it needs and uses the new class.

valar-cop reviewed Nov 9, 2026

View reviewed changes

zufayu requested a review from valarLip November 10, 2026 01:20

zufayu assigned valarLip Nov 10, 2026

valar-cop reviewed Nov 10, 2026

View reviewed changes

valarLip approved these changes Nov 11, 2026

View reviewed changes

valarLip merged commit eec6c47 into main Nov 11, 2026
21 of 23 checks passed

valarLip deleted the perzhang/k3-lmcache-mp branch November 11, 2026 03:17

ZLkanyo009 mentioned this pull request Nov 11, 2026

fix(offload): key multimodal prompts by media, one LMCache engine id per server; a native MP restore has its own deadline and waits for a free slot before dispatch #3027

Merged

		return span * tensor.element_size()


		def publish_kda_ride_state(transfer_tensors, tensors_by_group, geometry) -> None:

Conversation

PerryZhang01 commented Nov 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev record

Named result

Gates

Uh oh!

github-actions Bot commented Nov 9, 2026

🏷️ CI Guide

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PerryZhang01 commented Nov 10, 2026

Uh oh!

valar-cop commented Nov 10, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

PerryZhang01 commented Nov 9, 2026 •

edited

Loading