[BugFix] Mamba sparse retention keeps a state below each boundary under EAGLE/MTP - #556
Conversation
…er EAGLE/MTP The full-attention EAGLE lookup drops one block below what it matched, so until a request decodes past the block boundary after its prompt, the only candidate the coordinator can offer a Mamba group is one block below the replay boundary. Latest-only retention (prefix_cache_retention_interval=0, the default) kept exactly the boundary state, leaving every retained state one block above every reachable candidate: the reconciled prefix-cache hit was always zero on hybrid models running MTP/EAGLE spec decode. Measured live on GLM-5.3-Flash (MTP k=3, TP=4): 0 hits across 16,897 queries; per-group lookups on an identical resent 8,901-token prompt returned (6912, 6912, 6912) for the Mamba groups against 4608 for the eagle-dropped full-attention group, reconciling to 0. Keep the state one block below each reachable boundary as well when the group runs under EAGLE. The extra state only materializes where a block boundary coincides with an aligned prefill chunk end (the running-state block), which is how the align-mode scheduler chunks prompts.
|
Warning Review limit reachedNext included review available in 30 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (3)
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe Mamba reachable-block mask now retains the boundary state and the preceding state for EAGLE. A hybrid prefix-caching test verifies a two-block hit after replaying a 127-token prompt. ChangesMamba EAGLE retention
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to The PR restores prefix-cache hits for hybrid EAGLE/MTP requests by retaining one additional predecessor state while leaving non-speculative behavior unchanged. No actionable merge-blocking risk remains after normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…group's drop Two defects found reviewing the first cut, both caught by new tests: 1. The back-off was hard-coded as one Mamba block, but the drop the full-attention finder applies is min(alignment_tokens, its block_size) followed by a re-floor to the alignment -- which lands exactly one ALIGNMENT unit below the boundary either way. When a group's block size differs from the alignment (their LCM), one block is the wrong step: the retained state sits at an offset Mamba's own finder rejects, so the hit stays 0 and the extra block is dead weight. Now computed in tokens via reachable_hit_positions(), which is exact for alignment >, == and < the block size. 2. Retention keyed off the group's OWN use_eagle bit, which is 'this group holds draft layers'. The right predicate is 'some group's lookup shortens the candidate offered to me', since the coordinator reconciles all groups to one hit length. It worked only via the coordinator's flag-all fallback (no annotator exists for glm5_next); the day one lands, the zero-hit bug returns silently. The coordinator now sets lookup_drops_eagle_block on every manager from bool(self.eagle_group_ids). Tests: parametrize the MTP test over the three annotation routes (fallback, full-only, both) and give it real speculative blocks + lookahead; add a differing-block-size test pinning the alignment-unit back-off. Both fail on the previous cut (full_only: 'mamba hash 1 should be cached'; backoff: 'reachable state missing; cached=[5]'). 95/95 in test_prefix_caching.py.
|
Updated after review: two defects in the first cut, both now pinned by tests. The back-off is one alignment unit, not one block. The full-attention finder subtracts Retention must key off "any group drops", not "this group is eagle". The coordinator reconciles every group to one hit length, so a drop anywhere shortens the candidate offered here. The first cut read the group's own Tests: the MTP test is parametrized over all three annotation routes (flag-all fallback, full-attention-only, both) and now runs with real speculative blocks and lookahead; a new differing-block-size test pins the alignment-unit back-off. Both fail on the previous cut ( Two related defects found in review and deliberately not fixed here, since each is pre-existing and independent:
|
|
Independent validation of
This is a local vLLM automatic-prefix-cache correctness fix, not an LMCache #589 separately classifies DFlash/DSpark as preserving target KV cache, so its |
|
Validation against 1. Please fold the SWA
|
3 participants
Prefix caching never hits on hybrid models running MTP/EAGLE spec decode. Measured on GLM-5.3-Flash in production: 0 hit tokens against 16,897 queried, silently.
The full-attention EAGLE lookup drops one block below what it matched, but latest-only Mamba retention (
prefix_cache_retention_interval=0, the default) keeps only the state at the boundary itself, one block above anything a lookup can ever request. The reconciled hit is always 0.Fix: when the group runs under EAGLE, also retain the state one block below each reachable boundary. Costs at most one extra state page per request; non-speculative behavior unchanged.
Verified live on a 4x DGX Spark GLM-5.3-Flash deployment (MTP k=3, TP=4): an identical 8,901-token resend now hits 4,608 tokens, warm TTFT 2.66 s vs 5.88 s cold, warm output byte-identical at temperature 0. Regression test included (fails without the fix);
tests/v1/core/test_prefix_caching.pypasses 92/92.Summary by CodeRabbit