[Test] e2e hybrid-Mamba prefix-cache corruption regression tests (#43559)

[puririshi98]

Purpose

End-to-end output-correctness regression tests for #43559 — hybrid-Mamba
prefix caching silently corrupting output under mamba_cache_mode="align"

• MTP/EAGLE speculative decoding. These are the e2e tests split out of
  [Test] Regression test for hybrid-Mamba eagle cache-peek in Mooncake connector (#43559) #48361 (per review) and supersede the test-only [Test] Deterministic e2e regression tests (xfail) for hybrid-Mamba prefix cache corruption (#43559) #48198.

Two graded arms, each compared against an enable_prefix_caching=False
control using deterministic greedy needle recall:

• cold-race (write-side): concurrent cold prefills fragment a prefill
  chunk mid Mamba block; the chunk-end state can be cached as the boundary
  snapshot.
• multi-turn (read-side): reuse of blocks written during (speculative)
  decode, where the eagle-lookahead lookup can let the Mamba hit length
  overrun the attention-verified hit.

Parametrized on Nemotron-Super-120B-A12B-BF16 (TP4) and
Qwen3.6-27B-FP8. Corruption raises CorruptionDetected; geometry /
control-quality / cache-liveness problems hard-fail with dedicated
exceptions (never pytest.skip, since the fork runner reports in-body
skips as PASS). An anti-vacuity liveness gate hard-fails a "clean" run
whose prefix cache is not demonstrably live (and now also hard-fails if the
≥3-block liveness probe cannot fit max_model_len, rather than silently
disabling itself). Registered as a non-gating optional Buildkite step.

What these tests actually guard (scope)

These are in-process (LLM) tests; they exercise the
HybridKVCacheCoordinator read path and the align-mode scheduler, not
the Mooncake KV-connector. Concretely:

• The corruption arms are a regression guard for the already-merged
  in-process [Bug]: Accuracy drops ~20% when --enable-prefix-caching is used together with MTP speculative decoding (Qwen3.6 35B-A3B) #43559 fixes ([2/N][Core] support partial prefix cache hit for hybrid model #46384 kv_cache_coordinator eagle-margin
  gate + [Core] Preserve Marconi caching with selective hybrid cache retention #47782 selective retention). On current main these arms are
  green because those fixes are already in; the RED baseline in the table
  below is against a pre-[2/N][Core] support partial prefix cache hit for hybrid model #46384/[Core] Preserve Marconi caching with selective hybrid cache retention #47782 tree. If a future change
  regresses the in-process fix, the corruption arms go red.
• The cold-race liveness gate additionally requires [V1][Metrics] Preserve prefix-cache stats on zero-output steps #48668 — it
  reads vllm:prefix_cache_hits, which the offline LLMEngine.step()
  path under-reports to 0 (hits applied on non-output-bearing prefill
  steps are dropped) without [V1][Metrics] Preserve prefix-cache stats on zero-output steps #48668, which would make the gate fire even
  though the cache is live.

Dependency: requires #48668 to pass. It does not exercise
#48361's Mooncake connector path (_find_hit_blocks) — #48361 is covered
by its own unit test, not here. (Kept in the #43559 stack for context;
merge order relative to #48361 does not matter for these tests.)

Why not a duplicate

Supersedes #48198 (test-only, draft). The production fixes live in #48361
(Mooncake connector) and #48668 (stats); this PR is the in-process e2e
regression coverage for #46384/#47782 + the #48668 liveness gate.

Test Plan / Test Result

pytest -v -s -rA tests/v1/e2e/test_hybrid_mamba_prefix_cache_correctness.py

Validated on GB200 (Nemotron-Super-120B-A12B-BF16, TP4, current main +
#48668; #48361 present but not exercised):

Arm	Result	Engagement
cold-race	1 passed	liveness_probe_hits=4224, no corruption
multi-turn	1 passed	wave-2 hit ratio 0.477 (67584 hits), no corruption

RED baseline (pre-#46384/#47782 main) reproduced corruption on both arms
(documented in issue #43559).

AI assistance disclosure

Developed with AI assistance (Claude Code), including GB200 validation and
an adversarial self-review that corrected the dependency scope above;
every line reviewed by the human submitter.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[puririshi98]

@TomerBN-Nvidia @benchislett

[mgoin]

Please reduce the amount of comment slop in this PR. There must be close to 200 lines of comments here with dense blocks. This is simply not needed to this degree. Reduce to match the style of the codebase and long-lived concerns/information

[puririshi98]

Thanks — fair, and done in 050adedf.

Prose is down from 286 lines to 137 (38% of the file to 23%): 252 deletions, 101 insertions. Specifically:

• module docstring 64 lines -> 23
• the three exception-class docstrings (~10 lines each) -> 1-4 lines
• dropped a 5-line paragraph on which exception type means what that was duplicated verbatim in both test docstrings; it is stated once at module level now
• ~10 inline blocks explaining calibration numbers, metric-flush mechanics and re-derivations of already-stated rationale
• the buildkite step comment 5 lines -> 1 (it was the only commented step in that file, which rather made your point)

What I kept is the non-obvious, long-lived rationale: why grading is relative to an APC-off control rather than zero-tolerance (vLLM forces mamba_cache_mode="none" when prefix caching is off, so the arms differ numerically by construction), why these paths hard-fail instead of calling pytest.skip (the fork-based per-test runner reports an in-body skip as PASS and would silently green the step), and why ControlQualityFailure is a plain Exception rather than pytest.fail (Failed derives from BaseException and escapes the fork wrapper). Happy to trim further if any of those still read as slop.

Executable code is unchanged — verified by comparing the ASTs with all docstrings stripped — and pytest --collect-only still collects all 3 parametrizations on a GB200 node.

[mgoin]

This is an extremely expensive test to run, can we make this more efficient? I think we should be able to run this on a single B200, which would be preferred due to the amount we have

[mgoin]

Is there a reason why you can't use smaller models, or more heavily quantized versions? For instance why not Qwen3.5 4B and nemotron 3 nano, and NVFP4 versions?

[puririshi98]

Thanks — both points were right, and chasing the second one turned up a defect in my own test.

Cost. 240 min was a guess. Measured on this branch: the Qwen arm is ~130 s and the Nemotron arm ~330 s. I've split the step by model — the Qwen param was already tp=1, so it now runs on a single B200 at a 45-min budget; only the 120B TP4 param needs 4 GPUs, kept optional at 90. The remaining budget is headroom for weight download, not runtime.

Nemotron 3 Nano. No variant ships an MTP head — Nano-30B-A3B BF16/FP8, 4B and Nano-Omni have no num_nextn_predict_layers, and the 30B safetensors index has 0 of 6243 tensors matching mtp|nextn — so method="mtp" raises NotImplementedError at construction. Worth flagging against my own case: this PR's paths gate on use_eagle(), which covers eagle3 as well as mtp, and NemotronH implements SupportsEagle3 — so an eagle3 draft on Nano would exercise the same read-side hit-length drop and write-side boundary backoff. I haven't validated any Nano eagle3 head, so I can't claim that configuration works, only that "Nano is impossible here" would be wrong.

NVFP4. nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 does retain a real MTP module (1040 mtp.* tensors, left unquantized in BF16) and the drafter genuinely engages — I measured 88 drafts / 172 accepted at tp=1. What blocks it is geometry, not the model: at tp=1 the resolved Mamba block jumps to 8400 while the weights cap gpu_memory_utilization near 0.59, leaving ~30 GiB of KV. Every prompt in this test is sized in blocks, so the working set grows as the cache shrinks and the arms fail. It's a different test point rather than the same test on cheaper hardware, so I'd rather not fold it in here.

I also removed the cold-prefill arm, because your push on cost made me run the checked-in config end to end and it does not grade its own trigger.

It never armed. Under align mode with EAGLE/MTP the last cacheable boundary is pulled back one block (scheduler._mamba_block_aligned_split), so the sub-2-block prompt it builds commits no boundary snapshot at all — measured zero prefix-cache hits across the whole arm. Its cache-liveness probe then asserted on vllm:prefix_cache_hits, which excludes preempted and unadmitted requests (kv_cache_manager.record_prefix_cache_stats → preempted_*, never exported); I instrumented it and the probe's own query delta was 0, so that counter never observed the request. The assertion "no corruption + zero probe hits ⇒ cache is dead" was therefore unsound as written.

I rebuilt it at a >2-block geometry as a same-engine cached-vs-uncached A/B — same engine, same config, single-request batches, so cache state is the only variable. That version did fire when I reverted the #45477/#47861 boundary backoff. But on a fixed tree it reported 12/40 recall regressions versus 3/40 with the fix reverted: anti-correlated, because resuming from a Mamba prefix cache legitimately differs from full recomputation by considerably more than the defect does. An 8-probe variant detected the reverted fix but flaked red 1-in-3 on correct code, with signal and noise both at ~1 probe. Rather than ship something that flags correct code, I've dropped the arm — #45477/#47861 cover that mechanism at the scheduler level, and the module docstring now says so.

The multi-turn arm is unaffected and passes on both models (hit ratios 0.444 and 0.477, 4147 and 3993 MTP drafts).

[puririshi98]

@mgoin let me know if this alleviates your concerns or if you would me to make any further changes

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @puririshi98.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

Replaced by #53189