[ZLkanyo009]

Ready for review. No blocking issues.

Since #2519, Qwen3.5 GDN prefix checkpoints are copied into MHA KV blocks as PAGE images. aiter's gluon paged-attention decode reads KV rows no sequence owns: padded block-table lanes load physical block 0, and a block's unwritten token slots load too. Those lanes get zero probability, but at fp8 an image byte is often a NaN encoding, and 0 * NaN in the PV MFMA wipes the output row. With an image on block 0 the whole decode batch is hit. Kimi K3, whose KDA checkpoints take the same PAGE path into its MLA pool, reads physical page 0 the same way.

Dev record

• Found: a NaN-poison probe fails all 48 shapes of aiter main pa_decode_gluon. aiter mla_decode_fwd (K3, page size 1) also reads physical page 0. On K3 the damage is quieter: the fp8 V up-projection quantizes per 128-element group with a NaN-ignoring max and a clamp, so a few NaN latent columns become finite but wrong values instead of NaN rows.
• Decided: guard it in ATOM. A pool holding state images never hands out block 0. Units a checkpoint releases are drained into the next batch (StateMaintenanceOps.checkpoint_clears) and zeroed before its copies and forward, one index_fill_ per MHA field. Check the drain order in take_state_maintenance_ops and CommonAttentionBuilder.build first.
• Surprised: greedy output is not reproducible across restarts (per-rank Triton autotune, MoE bf16 atomics), so every A/B below pins autotune and alternates restarts.
• Left undone: clears for the K3 MLA pool. The hook is a no-op there, because page size 1 has no unwritten tails; the null block alone covers K3. An in-kernel aiter fix: zeroing dead V rows costs 7-23% decode-attention time.

Named result

TP8 GSM8K, original aiter kernel: 0.8673 / 0.8704 / 0.8696 with the guard; 4 of 6 passes at 0.824-0.833 without it. Kimi K3 GSM8K: 0.9553 / 0.9545 with the guard; 0.9303 on main when block 0 held an image the whole pass. Serving throughput unchanged within restart noise.

Accuracy

GSM8K, original aiter kernel, Triton autotune pinned (TRITON_CACHE_AUTOTUNING=1, one cache).

Qwen3.5-397B-A17B-FP8. TP8, --gpu-memory-utilization 0.26, --state-checkpoint-interval-tokens 256, block 16, max_model_len 2048, prefix caching, c65, 3-shot, LM Cache off. "Long thinking" counts answers that open a Thinking Process block and get cut at 256 tokens; it is where the lost points go.

arm	flexible-extract	long thinking
fork checkpoints (pre-#2519 path), 2 restarts	0.8650-0.8779	2-5
PAGE, original aiter kernel, 2 restarts	0.8241 / 0.8332 / 0.8271 / 0.8264 / 0.8643 / 0.8711	83 / 55 / 60 / 73 / 13 / 6
PAGE + this guard (earlier env-gated version), 2 restarts	0.8582-0.8741	1-6
PAGE + this branch	0.8673 / 0.8704 / 0.8696	5 / 4 / 8

Kimi K3. TP8, agentic recipe: block 128, MLA page size 1, FP8 KV, prefix caching, --state-checkpoint-interval-tokens -1, ATOM_STATE_CHECKPOINT_DEMAND=0, util 0.93, 64 seqs. 5-shot, completions, c64. main is 53e4fb524 plus env-gated debug code that logs block 0's checkpoint reservations and, in passes 2-4, checks its first row each decode step; this branch has the same base. The recipe range is 0.9538-0.9591.

pass (in run order)	server	flexible-extract	strict-match	wrong	garbled answers	block 0 holding an image
main 1	restart 1	0.9477	0.9477	69	0	last 30% of the pass
main 2	restart 2	0.9522	0.9522	63	0	last 36%
this branch 1	restart 3	0.9553	0.9560	59	0	never
this branch 2	same server	0.9545	0.9545	60	0	never
main 3	restart 4	0.9515	0.9515	64	0	last 36%
main 4	same server	0.9303	0.9303	92	7	the whole pass

Standard error is 0.006 per pass (0.007 for 0.9303). Main averages 0.9454 over four passes (0.9505 over the first three); this branch averages 0.9549. Only one of the 7 garbled answers lost its point; the others had one stray token mid-answer and still ended on the right number.

Block 0 is the lowest id in the free heap, so once the pool fills (about 4 minutes in) each checkpoint that releases it hands it straight to the next one. Garbled answers are stray CJK tokens mid-answer, e.g. <<4*3=蚕豆>>12. K3 answers vary a lot from pass to pass (about 23% identical text between two passes on one server), so compare totals.

Performance

Qwen3.5-397B-A17B-FP8, TP8, FP8 KV, block 16, prefix caching, benchmark_serving random dataset, 640 prompts at concurrency 64. Two restarts per arm, alternating main / branch, Triton autotune pinned. Measured at e439bd8; 7abd832 only rewords comments and drops a flag.

scenario	arm	total tok/s	mean TTFT (ms)	mean TPOT (ms)
util 0.9, interval 8192, ISL/OSL 1024/1024, no checkpoint released	main	6632 / 6646	260.7 / 260.0	18.66 / 18.62
	this PR	6616 / 6629	263.2 / 264.1	18.70 / 18.66
	delta of means	-0.25%	+1.3%	+0.2%
util 0.30, interval 256, ISL/OSL 1024/256, ~546 checkpoints released per run	main	8413 / 8362	368.9 / 366.2	36.24 / 36.47
	this PR	8542 / 8342	369.0 / 369.6	35.65 / 36.55
	delta of means	+0.65%	+0.5%	-0.7%

Main against itself across a restart moves 0.2% (first scenario) and 0.6% (second) in throughput.

Gates

branch @ 7abd832: CPU suite failures identical to main (38, environment); the 24 GPU-only test files 573 passed; 8 new tests; GSM8K smoke (200 questions, 169 checkpoints released) 0.880. black, ruff clean on changed files.

Evidence

Probe: fill KV rows no sequence references (and unwritten tails) with zeros, then NaN; the outputs must be finite and bit-identical. pa_decode_gluon on aiter 84077b3c0: 48 of 48 shapes fail (ps / partitioned, 1-2 KV heads, bf16 / fp8 KV, block 16 / 64 / 1024). mla_decode_fwd, page size 1: poisoning all of physical page 0 wipes every row, and a single NaN byte there turns that latent column NaN in every row; poisoning other unlisted pages changes nothing.

K3 mechanism, measured: block 0's page 0 held 2-3 NaN latent columns per MLA layer. In the kernel alone each NaN byte there turns that column NaN for every row and head. batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant then maps the NaN to finite values: 1 NaN column gives 13% relative error in the V up-projection, 3 give 24%. Only an all-NaN page 0 zeroes the output, and injecting that into a block 0 no sequence owns drops GSM8K to 0.005 (200 questions).

The second performance scenario first ran at util 0.26, but KV block counts are sized from free memory at startup and two restarts differed by 19% there, so it moved to 0.30 (blocks within 3.5%).

🤖 Generated with Claude Code

[github-actions]

🏷️ 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 2635 --add-label <label>

[valarLip]

@yhl-amd

[ZLkanyo009]

The CI failures here (GLM-5.3-Flash, Kimi-K3, Kimi-K3 DSpark in run 39380615429) all stop at triton_mla.py:7, which imports aiter.ops.triton.attention.mla_decode. That module was removed from aiter, and main moved off it in #2669. A rerun would reuse the old merge ref and fail the same way, so the fix is a main merge.

I have that merge ready locally (origin/main at 5fea029, clean, c62a145). On its own it checks out:

• black/ruff clean.
• CPU unit tests: 3 failed / 9728 passed, the same 3 failures as main.
• GPU: test_gdn_page_checkpoint.py, test_mha_kv_pool.py, test_block_pool.py, test_state_checkpoint.py and test_lmcache_mp_native_scheduler.py: 448 passed.

I have not pushed it, because this PR's clear-on-release causes a regression when it is combined with #2636. With Qwen3.8 in-process LMCache offload and interval 256, GSM8K pass 2 (retrieved from CPU) drops to 0.927-0.973, against 0.977-0.986 for pass 1. The pass-2 answers solve a different problem or run to the token limit. Without this PR, or with Qwen3.8's clear_page_units made a no-op, pass 2 is back at 0.977-0.984. One measured cause is that the connector's start_load_kv writes a re-used raw unit before the batch's build clears it, which breaks the BlockPool assumption that "the batch that allocation serves drains it". Numbers and diagnostics are in #2636.

The GDN clear_page_units here probably has the same exposure with Qwen3.5 offload (#2562), though I haven't tested that. Once it's decided whether the fix goes here or in #2636, I'll push the merge (or someone can push c62a145's equivalent) and ci:atom will rerun.

[ZLkanyo009]

Clear-on-release zeroes KV a connector load just wrote: fixed in cfa9b4c1

With this branch, a request fetched back from the LMCache CPU tier can lose part of its context. The released PAGE units are cleared after a connector load has already filled them. This hits Qwen3.5 on main's LMCache path (#2519) as well as Qwen3.8 with #2636. The fix is a new commit cfa9b4c1, pushed together with the main merge c62a1450 that was already prepared (main 5fea029d). There was no force push. The branch still merges cleanly with main 0ba93b24 and with #2636 484dc30e.

Root cause

release_units makes a unit allocatable again immediately and queues it in _released_units. The batch drains that queue into checkpoint_clears, and CommonAttentionBuilder.build zeroes the units before the forward. The comment on take_released_units assumes that whatever reuses a unit writes it only after that batch's clear. That holds for the forward. It does not hold for a connector load:

• EngineCore dispatches the connector metadata (start_load_kv) before the forward.
• The LMCache load fills its blocks from its own threads and streams. Nothing orders it after the compute stream.
• Under pool pressure, a load's target blocks are exactly the given-back units. Allocating them evicts a checkpoint, and that checkpoint's units come straight back out.

So the clear lands after the load in two ways:

1. A pass with a forward. The load and the build-time clear race.
2. A pass with no forward. This is the deterministic case. If every scheduled request is parked on a load (WAITING_FOR_REMOTE_KVS), which is how the second pass of an offload run starts, the decode branch builds an empty batch. That batch drains nothing, but its metadata still goes out and the loads start. The next non-empty batch then zeroes the units after the loads have finished. _dispatch_idle_offload_work and PP's _dispatch_connector_only_batch also dispatch loads with no batch at all.

With bf16 KV, a zeroed block removes those tokens from every full-attention layer. The GDN state is loaded into the slot separately, so it is unaffected. The model loses part of the prompt. It then answers a different few-shot question or runs to the length limit.

A diagnostic build (#2636 + this branch without the fix) recorded which blocks each dispatched load wrote and counted how many of them build later cleared. On one 440-question fetched pass:

• 948 blocks were written by loads dispatched in the 6 forward-less passes at the start of the pass.
• 1107 blocks were hit by same-pass races.
• No other path showed up.

The other writers onto these units are not affected:

• State store to CPU: the source is released only after gather and D2H.
• State load: writes the slot, not PAGE units.
• HBM restore and store: run in build after the clear, and are pinned.
• KV save: reads live blocks only.

lmcache_mp KV loads and reserve_transfer_units are connector writes of the same kind, and the fix covers them. They were not measured on a GPU.

Fix (cfa9b4c1)

Move the clear ahead of the loads rather than holding units back from allocation. KV allocation evicts a checkpoint and reuses its units within the same pass (_ensure_page_units), and can_allocate already counts on that.

• Scheduler.carry_checkpoint_clears(meta): whenever connector metadata that has work is about to be dispatched, every unit given back so far moves onto it as ConnectorMetadata.checkpoint_clears. It is called:

  • at both scheduling branches, including the empty pass;
  • in _dispatch_idle_offload_work;
  • in PP's _dispatch_connector_only_batch.

  Metadata with no work is left alone. The engine may drop it, and since it has no load, the batch clears those units as before.

• ModelRunner.process_kvconnector_output: zeroes meta.checkpoint_clears on the compute stream, then waits for every clear issued so far, including the previous forward's build, before start_load_kv. The builder methods are clear_released_units and wait_released_unit_clears, and the latter is one event sync. It costs nothing once the clears have run.

• The NaN guard is unchanged. The same units are zeroed with the same index_fill_, block 0 is still never handed out, and the clear still runs before any forward that could read them. Only the case where a load would fill a unit first now clears it earlier.

Data

MI355X, TP1, bf16 KV, in-process lmcache_offload, --state-checkpoint-interval-tokens 256 unless noted.

The method is the per-segment GSM8K replay: questions 0–439, 440–879 and 880–1318 each run twice. Pass 1 stores. Pass 2 is served from the CPU tier. pass2 below is the fetched pass. "Empty / truncated" counts answers that are empty and hit the 12288-token limit.

Tree	Model	Spec	Questions	pass1	pass2 (fetched)	Empty / truncated	Fetched
main + this branch + #2562, before	Qwen3.5-35B-A3B-FP8 (interval 1024)	0	440	0.9636	0.9136	27	373/440
main + this branch + #2562, after	Qwen3.5-35B-A3B-FP8 (interval 1024)	0	440	0.9682	0.9750	2	373/440
#2636 + this branch, before	Qwen3.8 Flash Next MXFP4	0	1319	0.9727	0.9325	42	1318/1319
#2636 + this branch, before	Qwen3.8 Flash Next MXFP4	MTP 3	440	0.9841	0.9545	7	440/440
#2636 + this branch, after	Qwen3.8 Flash Next MXFP4	0	1319	0.9780	0.9780	0	1318/1319
#2636 + this branch, after	Qwen3.8 Flash Next MXFP4	MTP 3	1319	0.9773	0.9780	1	1319/1319
#2636 alone (reference)	Qwen3.8 Flash Next MXFP4	MTP 3	440	0.9773	0.9773	1	440/440

• After the fix, the per-question pass1-vs-pass2 differences are noise. With MTP 3, 7 questions are right only in pass 1 and 8 only in pass 2; with spec 0 it is 7 and 7. Qwen3.5 after the fix: 4 vs 7. Before the fix the split was 65 vs 12 (Qwen3.8, spec 0, 1319 questions), 14 vs 1 (Qwen3.8, MTP 3, 440 questions) and 23 vs 1 (Qwen3.5).
• Before the fix, MTP 3 tokens per forward rose from 3.16 in pass 1 to 3.39 in pass 2. The output degenerates into repetition, which the draft predicts more easily. After the fix it is 3.16 in both passes.
• Qwen3.5-35B-A3B-FP8 runs on one GPU with thinking on, about 6.4k tokens per question. At interval 256 its checkpoints overflow the CPU state pool, so the run uses interval 1024 and --gpu-memory-utilization 0.22, which pushes KV out of HBM so pass 2 has to load it. Its pass 1 already has 2–5 answers that think up to the length limit. The 397B was not run: its weights are not mounted in the available containers.
• Full GSM8K on the fixed [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 tree, MTP 3, no offload: 0.978 strict. Acceptance 0.721, 3.16 tokens per forward.
• MTP 3 pass-2 wall time per 440 questions: 114 s after the fix, 122 s before it, 119 s for [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 alone. No measurable cost.

Tests

New tests:

• test_state_checkpoint.py: the scheduler carries the units on the metadata of a pass where every request is parked on a load, and of a pass with a forward. It leaves them to the batch when the metadata has no work.
• test_lmcache_offload_connector.py: the idle dispatch and PP's connector-only dispatch carry them.
• test_gdn_page_checkpoint.py:
  • the worker clears and then waits before start_load_kv (needs AITER);
  • the wait actually waits for the device (needs a GPU).

Without the fix, all 8 new tests fail. With only the carry call sites removed, the 4 carrying tests fail.

Results with the fix:

• CPU suite: 3 failed / 9734 passed. The 3 failures are identical to main's (test_graph_marker ×2, test_sparse_indexer_side_effect_op).
• GPU files of this PR (test_gdn_page_checkpoint, test_mha_kv_pool, test_block_pool, test_state_checkpoint, test_lmcache_mp_native_scheduler, test_lmcache_offload_connector): 713 passed, 0 skipped. Merged with [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636, plus its offload test files: 815 passed.
• black 26.10.1 and ruff 0.17.2 are clean.

The NaN-poison probe was not rerun: the fix does not change what is cleared or the block-0 rule.

[ZLkanyo009]

@valarLip could you take a look at this one?

Updated today: main merge c62a1450, plus cfa9b4c1, which zeroes given-back PAGE units before a connector load fills them. Without it, a request fetched back from the LMCache CPU tier could lose part of its context. This also hits Qwen3.5 on main's LMCache path, not only Qwen3.8 with #2636. With the fix, the CPU-fetched GSM8K pass is back to normal: Qwen3.5-35B-A3B-FP8 goes from 0.9136 to 0.9750, and Qwen3.8 with #2636 gives 0.9780 at both spec 0 and MTP 3. Root cause and full table: comment.

CI: Pre Checkin is green. The one Accuracy failure (GLM-5.3-Flash kpool, run 40284323652) ran on runner pit2-p03-g40. Like the three on #2636, the step was killed (exit 137) mid-run and the container disappeared. The job passes on main, so I've rerun it.

[zhuyuhua-v]

will fix in kernel level

[github-actions]

🏷️ 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 2635 --add-label <label>

[ZLkanyo009]

Kernel-level fix for the FlyDSL path: ROCm/aiter#7092 (FlyDSL pa_decode only, gfx950).

On the boundary tile, it zeroes V bytes for tokens at or past context_len before PV, and it zeroes per-token scales there. That covers both padded lanes pinned to block 0 and the unwritten tail of the last page. Clean-input output is bit-identical to main. Kernel time: -0.6% geomean, +5% worst, over 513 decode shapes.

E2E on Qwen3.5-35B-A3B-FP8 with ATOM_PA_FLYDSL=1: ATOM main without this PR gets GSM8K 0.063 on aiter main and 0.964 with aiter#7092. This PR on aiter main also gets 0.964.

With that in, all three parts here (block 0 never allocated, released units zeroed, zeroing before connector loads) are redundant on the FlyDSL path. They are still needed for gluon (default, not fixed there; peeling the last partition is the cheap way, details in the aiter PR), for calls that fall back from FlyDSL, and for K3's MLA pool. So nothing changes here yet. Once gluon gets the same mask, the MHA-pool parts can be dropped.

[ZLkanyo009]

@valarLip Merged latest main (bdec6812) into this branch to clear the conflict; new merge commit d5a4e6bf. No rebase, no force push. This only merges main: the PR's own changes are unchanged (cfa9b4c1), and it stays a draft.

Conflict. Only the import block of tests/test_state_checkpoint.py conflicted. #2562 added AdmissionProbe, and this PR added ConnectorMetadata. Both imports are kept:

from atom.kv_transfer.disaggregation.types import ConnectorMetadata
from atom.model_engine.block_manager import AdmissionProbe, BlockManager

Verification

• black 26.10.1 and ruff are clean on the touched file.
• CPU unit suite (.github/scripts/run_unit_tests.sh): 10093 passed, 4 failed.
  • Three of the failures also fail on plain main bdec6812 on the same machine (10081 passed, 3 failed): test_graph_marker::test_compiled_marker_preserves_inplace_aliases[False/True] and test_sparse_indexer_side_effect_op::test_side_effect_op_stays_ordered_without_a_synthetic_output, all inductor g++ compile errors.
  • The fourth was a timing-based check, test_stream_marker_properties::...::test_a_long_call_costs_what_it_is_and_not_its_square[qwen], which this PR doesn't touch. It passed when rerun (14/14 in that class).
• GPU unit tests on MI355 for the files this PR touches plus the scheduler / block manager suites (test_state_checkpoint, test_block_pool, test_gdn_page_checkpoint, test_lmcache_mp_native_scheduler, test_lmcache_offload_connector, test_mha_kv_pool, test_block_manager, test_scheduler): 1003 passed. Plain main gives 987 on the same files.

[ZLkanyo009]

Closing this in favor of the kernel-level fix. ROCm/aiter#7092 is merged (2db75fc6c): FlyDSL pa_decode now zeroes V rows past the context before PV (mask_unowned_kv, on by default). On the FlyDSL path (ATOM_PA_FLYDSL=1), Qwen3.5 no longer needs this ATOM-side guard: in #7092's end-to-end run, ATOM main without this PR went from GSM8K 0.063 to 0.964.

Two decode kernels still read KV rows that no sequence owns. They are left to the kernel owners (cc @zhuyuhua-v @fsx950223 @coderfeli):

• gluon pa_decode_gluon: ATOM's default MHA decode, and also the fallback for calls FlyDSL does not cover, such as bf16 KV or a sliding window.
  • A NaN-injection probe fails all 432 shapes on aiter main: bf16 outputs NaN or Inf, fp8 outputs finite but wrong values.
  • Costs measured while preparing #7092: an unconditional gl.where(token < context_len, v, 0) costs +13.4% geomean / +66% worst. Masking only the boundary partition behind a branch costs +1.2% / +24% worst, because the loop body splits.
  • The cheap version needs the last partition peeled out of the loop in each of the four decode kernels.
• mla_decode_fwd (Kimi-K3, page size 1): padded block-table lanes read physical page 0.
  • With a state image there, each NaN byte turns that latent column NaN for every row.
  • The fp8 V up-projection then maps it to finite but wrong values, about 13% relative error for one NaN column.

Until those land, Qwen3.5 / Qwen3-Next on the default gluon path and Kimi-K3 with PAGE checkpoints can still pick up a state image's bytes in a KV row the decode reads. FP8 KV makes this most visible. The measurements and the guard's design remain in this PR's description for reference.