test: allow bounded checkpoint lag in Mamba HiCache tests

[alphabetc1]

Motivation

TestUnifiedMambaHiCache.test_multiturn_logprobs_match can reject a valid checkpoint after KV-memory retraction. This CI job expected at least 4864 cached tokens but received 4839.

Retraction discards the in-flight decode checkpoint. Resumed prefill saves states relative to the recovered prefix: 4391 + floor(491 / 64) * 64 = 4839. If generation ends before the next absolute decode checkpoint, that prefix-relative checkpoint is the one available for reuse. The same assertion failure reproduced twice on unmodified main with cached_tokens=4820.

Modifications

Keep the shared Mamba decode-cache assertion strict by default. Add an explicit max_checkpoint_lag option and set it to 63 tokens (less than one 64-token prefill chunk) only for TestUnifiedMambaHiCache, whose 12,000-token KV pool and four concurrent requests can trigger retraction. Other callers and zero-output assertions keep their original bounds. The existing KL replay and threshold are unchanged. No retraction metadata plumbing is needed.

Two files change: the assertion helper and the Mamba HiCache test configuration. No inference implementation or callback interface changes.

Accuracy Tests

• Focused checks confirm that the default remains strict and only the opt-in assertion accepts the bounded lag: both observed values (4839 and 4820) pass; the minimum allowed value (4801) passes; 4800 and a complete cache miss fail; zero-output and short-input bounds are preserved.
• Applied the old and new assertions to the four recorded responses from the reproduction: the old assertion fails at 4864 -> 4820; all four pass with the fix.
• Applicable pre-commit checks passed.
• 4×H100 / FA3, page size 1: rerun-test on this head passed all four TestUnifiedMambaHiCache tests in 313 seconds. GSM8K accuracy: 0.950. The three KL averages were 0.003134, 0.002697, and 0.003543, all below the unchanged 0.005 threshold. The log confirms a natural retraction during the prefill-branching test.
• Extended 4×B300 / FlashInfer validation: two matrices covered 624 generated-request replays and 284 retractions, including up to six retractions in a single request. Each matrix included natural pressure, a larger-pool no-retraction control, forced retraction, and cold single-request generation with cache flushing. Output lengths 129, 257, and 512 were exercised. The original 4820-token cache hit recurred in all three long-request repetitions after aligning DeepGEMM/DeepEP with the source requirements; cache-length checks passed throughout.
• Known numerical limitation on B300: the extended matrices were not fully green under the unchanged KL threshold. Failures also occurred without retraction. Even with cached_tokens=0 and num_retractions=0, 11 of 24 requests exceeded 0.005 in the aligned-dependency cold control. This per-request screening is stricter than the registered batch-average assertion; registered batch-average KL failures also occurred in the B300 matrices. These results do not establish retraction-specific state corruption, nor do they rule out subtler numerical issues. No KL threshold was relaxed. The constructed B300 pressure inputs differ from the original CI inputs.

Speed Tests and Profiling

Not applicable: test-only change.

---

CI States

Latest PR Test (Base): ❌ Run #37409487107
Latest PR Test (Extra): ❌ Run #37409486896
Latest PR Test (AMD ROCm 10): ❌ Run #37409487157

[alphabetc1]

/rerun-test -c

[github-actions]

⛔ --changed: this PR adds or modifies no runnable test files under test/registered/ or python/sglang/multimodal_gen/test/.

[alphabetc1]

/rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit fea3c088a141.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[alphabetc1]

/rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

[github-actions]

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

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py