[ShangmingCai]

Motivation

With --disaggregation-enable-kv-checksum on, a mismatch on one TP rank hangs the decode engine, and healthy traffic is aborted on several supported configurations. This PR fixes those paths so the flag can be turned on.

Modifications

All-reduce the mismatch decision. Corruption is per rank, so a rank-local drop left the TP peers building different batches and hanging on the next collective. The drop set is MAX-reduced over attn_tp_cpu_group with all_reduce_flags before any request is aborted.

Verify once. The digest covers whole pages, and after a retraction the restored row's last page holds tokens the decode wrote, so re-verifying on resume was a false mismatch at page_size > 1. A request is digested when it first enters a batch and its expectation is then cleared.

Ship a layout signature with the digest. kv_layout_signature(kv_args) hashes the page size and the per-page KV and state item lengths; the prefill writes it into slot 1 of the kv_checksum row. A prefill with another TP width or a PP stage is admitted unverified with one warning instead of aborting every request. A prefill that wrote no digest is skipped silently.

Refuse the decode-side radix cache. The digest covers [0, seq_len) while the prefill then sends only [decode_prefix_len, seq_len), so every prefix hit looks like corruption. The decode logs once at startup and runs without the checksum.

Survive a metadata buffer count mismatch. The aux buffers pair positionally, and the Mooncake and NIXL send loops walked the decode's list (Mori's TCP loop the prefill's), so an engine with one more buffer than its peer was an IndexError in a transfer thread. The loops stop at the shorter list and the prefill logs an error at registration naming the optional buffers that can differ. The decode resets the checksum row together with the bootstrap room when a metadata slot is freed, so a later occupant whose prefill writes no digest cannot read the previous one.

Test hook. SGLANG_TEST_DISAGG_KV_CHECKSUM_MISMATCH_PROB flips a decode rank's computed digest with that probability, drawn from a per-rank generator (the global random is seeded identically on every TP rank, so the ranks would otherwise inject the same requests). The in-CI raise is suppressed while it is set.

Accuracy Tests

• test/registered/unit/disaggregation/test_kv_checksum.py: new cases for the signature, the slot reset, verify-once, the reduced drop set, the signature in the prefill's metadata write, and the injector not tripping the in-CI raise. The CPU-only cases pass on a CPU box; the CUDA cases need a runner.
• test/registered/disaggregation/test_disaggregation_kv_checksum.py (new, base-c / 4-gpu-h100): TP2 on both sides, because at TP1 the reduce is a no-op. Mismatches injected per decode rank at p=0.3 over 16 requests with a multi-chunk prompt: some must abort, some must complete, and both engines must stay up.

No GPU here; the CUDA unit cases and the e2e test need a runner.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): 🚫 Run #40182804723
Latest PR Test (Extra): ❌ Run #40182804764
Latest PR Test (AMD ROCm 10): 🚫 Run #40182804728

[ShangmingCai]

Review round 2

All four points were right; I verified each against the code before fixing.

1. _dsa_tail fed a descriptor into the index path — confirmed, and it was the same bug class this PR set out to remove. get_dsa_tail_state_indices returns [req_pool_idx, start_phys, first_n, 0, second_n, tail_size] for the transfer backends to decode, not rows. I wrapped it in _as_index_tensor and compute() used all six as row indices into _compress_tail_k[i], which is (req_pool_size, ...). Both failure modes hold: req_pool_idx differs between engines, and the other five land on rows owned by other in-flight requests whose tail state is mutating. Your correction is right too — in bounds, just wrong rows. DSA_TAIL now has no index mapping and takes the documented exclusion path.

I checked the other two helpers for the same shape: get_qsa_pending_state_indices and get_dsv4_c128_state_indices both return genuine one-element row indices, so the bug was isolated to _dsa_tail.

2. No true-positive e2e — confirmed. Added SGLANG_TEST_DISAGG_KV_CORRUPT_PROB (next to the existing SGLANG_TEST_DISAGG_FAILURE_PROB, per the env-var conventions) and corrupt_one_kv_row_for_test, which clobbers a landed row through the pool's own buffers. New TestDisaggregationKVChecksumDetectsCorruption asserts the request fails and both engines stay up.

3. The all-reduce never ran multi-rank — confirmed. PDDisaggregationServerBase defaults prefill_tp_size = decode_tp_size = 1, so the collective took its world_size == 1 early return every time. Added test_disaggregation_kv_checksum_tp2.py on a 4-gpu-h100 runner (prefill 0-1, decode 2-3) with a healthy-traffic case and a per-rank-disagreement case that asserts the engines keep serving rather than hang.

4. Unconditional collective per iteration — confirmed, and fixed with a rank-invariant gate. A new Req.kv_checksum_pending is set from the prefill's layout signature, which is identical on every rank, rather than from the digest value, which is not. get_new_prebuilt_batch now skips _verify_kv_checksums entirely when nothing is pending, so a queue of already-checked requests costs no collective. I did not fold it into poll_and_all_reduce — that is the better end state and would subsume the hisparse cleanup too, but it moves verification to commit time and I cannot run a single test in this environment, so I kept the change reviewable instead.

Your four "dropped" items match what I found as well; the two-page-size one in particular — getattr(pool, "page_size", page_size) is belt-and-braces, not a claim they differ.

Still unverified by me: everything that needs a GPU. Lint is green; the e2e and CUDA unit tests need a runner.

[ShangmingCai]

Review round 3

All three confirmed and fixed. Two of them are the same failure mode as rounds 1-2, which is worth naming: three separate times now I have added a test that could not fail.

1. The in-CI raise contradicted both injection tests — confirmed. is_in_ci() reads envs.SGLANG_IS_IN_CI; the workflows set SGLANG_IS_IN_CI: true at job level; and test_utils.py:905 merges rather than replaces (env = {**os.environ, **env}), so the decode engine inherits it. A mismatch would raise inside the event loop and kill the scheduler that both new tests assert is healthy.

Took your first option: suppress the raise only while the injector is active (is_in_ci() and SGLANG_TEST_DISAGG_KV_CORRUPT_PROB.get() == 0.0). A real mismatch in a normal CI run still fails loudly, which was the point of the raise. Added a unit test pinning the suppression so it cannot silently regress.

2. base-b-test-4-gpu-h100 does not exist — confirmed. base-b owns 1-gpu-large, 1-gpu-small, 2-gpu-large, 4-gpu-b200. The 4-gpu-h100 suites are base-c-test-4-gpu-h100 and extra-b-test-4-gpu-h100. Moved to base-c and left a comment saying why, since register_cuda_ci is AST-parsed and a bad pair fails silently. The file written specifically to make _all_reduce_kv_checksum_mismatches execute was registered where nothing would run it.

3. Page id used as a token slot — confirmed. The K/V buffers are [size, head_num, head_dim] (memory_pool.py:3190), indexed by token slot; kv_page_indices_gpu holds page ids. They coincide only at page_size == 1. Now multiplied through allocator.page_size. Also took the fp8 point: the clobber is now an XOR over a uint8 view rather than += 1 on the KV dtype, so it is guaranteed to change bytes on any dtype.

Still unverified by me: anything needing a GPU. Lint is green.

[ShangmingCai]

/tag-and-rerun-ci

[ShangmingCai]

/rerun-test test/registered/disaggregation/test_disaggregation_kv_checksum.py

[github-actions]

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_kv_checksum.py