[Bugfix] Fix FULL cudagraph bypass on partial cache hits for layerwise/hybrid connectors (#1130)
Purpose
When a cache hit falls just a few tokens short of the full request
length (e.g. request 8193 tokens, 8192 cached), vLLM dispatches the
prefill as uniform decode into FULL cudagraph. FULL cudagraph replay
skips wait_for_layer_load / save_kv_layer, breaking layerwise KV
load/save. The previous workaround only fired on full hit
(num_total_hit_tokens == request.num_tokens), missing the partial-hit
case. For hybrid models, the old external_hit_tokens -= 1 also left
total_hit_block_num unchanged, causing the mamba-align state block to
be loaded at a position beyond num_computed_tokens – feeding the SSM
a “future” state that already includes the recomputed positions.
Additionally, the dump wrote mamba state at every LCM boundary within a
chunk, but only the last boundary (chunk end) has a valid running state
in HBM, leaving intermediate state entries as garbage in the store.
Modifications

UCMDirectConnector._get_full_hit_recompute_tokens: Consolidate
the layerwise/non-layerwise logic here (check self.use_layerwise,
return 2 + num_speculative_tokens for layerwise, 1 otherwise).
Remove the UCMLayerWiseConnector override – both
UCMLayerWiseConnector and the hybrid connectors now inherit directly.
UCMDirectConnector.get_num_new_matched_tokens (non-hybrid):
Change the workaround condition from num_total_hit_tokens == request.num_tokens (full hit only) to actual_recompute_tokens < recompute_tokens, so partial hits that would leave too few recompute
tokens also trigger the workaround. Only the deficit is subtracted
from external_hit_tokens; total_hit_block_num is unchanged.
UCMHybridLinearAttentionConnector.get_num_new_matched_tokens
(hybrid): Move the workaround BEFORE the lookup – compute
max_hit_lcm_blocks = (request.num_tokens - recompute_tokens) // lcm_block_size and truncate lookup_block_ids before querying the
store. This ensures the last block(s) whose mamba state may not be valid
are never looked up; if the state is absent, stage 2 of the lookup
naturally degrades to a shorter hit instead of returning stale data. No
post-lookup reduction is needed.
_generate_hla_dispatch_meta (hybrid dump): Only dump the
mamba-align state at last_lcm_b (the last LCM boundary in the chunk),
not at every LCM boundary. Intermediate states within a chunk are not
valid in HBM because the mamba kernel only keeps the running state at
the current position. Dumping them wrote garbage data with valid keys
into the store, which the lookup would later find and return as a hit.

Test
Non-hybrid: QwQ-32B (block_size=128, spec=3, recompute=5,
max_num_batched_tokens=8192)
Send a first request of 8192 tokens to populate 64 blocks in the cache,
then send a second request sharing the same 8192-token prefix:
| Case | Request B length | Hit tokens | actual_recompute | Workaround?
| deficit | ext_hit_tokens returned | vLLM recompute |
|—|—|—|—|—|—|—|—|
| Full hit | 8192 | 8192 | 0 | Yes (0<5) | 5 | 8187 | 5 |
| 1 token short | 8193 | 8192 | 1 | Yes (1<5) | 4 | 8188 | 5 |
| 4 tokens short | 8196 | 8192 | 4 | Yes (4<5) | 1 | 8191 | 5 |
| Borderline (no workaround) | 8197 | 8192 | 5 | No (5<5 is False) | –
| 8192 | 5 |
| Enough new tokens | 8200 | 8192 | 8 | No | – | 8192 | 8 |
| No hit | Different prompt | 0 | – | No | – | 0 | Full |
Hybrid: Qwen3.5-35B-A3B (block_size=2048, spec=3, recompute=5)
With MTP (EAGLE pruning active)
EAGLE forces the last block into a separate chunk (last_cache_position = num_tokens - block_size), so an 8192-token request is chunked into
6144 + 2048. Each chunk boundary has a valid running state.
| Case | Request B length | max_hit_lcm_blocks | Truncated? | Blocks
looked up | ext_hit_tokens | vLLM recompute |
|—|—|—|—|—|—|—|
| Full hit | 8192 | (8192-5)//2048=3 | Yes (4->3) | 3 | 6144 | 2048 |
| 1 token short | 8193 | (8193-5)//2048=3 | Yes (4->3) | 3 | 6144 | 2049
|
| 4 tokens short | 8196 | (8196-5)//2048=3 | Yes (4->3) | 3 | 6144 |
2052 |
| Borderline (no truncation) | 8197 | (8197-5)//2048=4 | No (4==4) | 4 |
8192 | 5 |
| Enough new tokens | 8200 | (8200-5)//2048=4 | No | 4 | 8192 | 8 |
| No hit | Different prompt | 0 | No | 0 | 0 | Full |
Without MTP (no EAGLE pruning)
Without MTP, recompute_tokens = 2 (layerwise baseline). The request is
a single chunk, so only the last block boundary has a valid mamba state.
Pre-lookup truncation removes the last block; the backward scan finds no
valid intermediate states and degrades to 0 hit. This is correct but
sacrifices cache benefit.
| Case | Request B length | max_hit_lcm_blocks | Truncated? |
ext_hit_tokens | vLLM recompute |
|—|—|—|—|—|—|
| Full hit | 8192 | (8192-2)//2048=3 | Yes (4->3) | 0 (graceful
degradation) | 8192 |
| Borderline (no truncation) | 8194 | (8194-2)//2048=4 | No (4==4) |
8192 | 2 |
| Enough new tokens | 8200 | (8200-2)//2048=4 | No | 8192 | 8 |

To maximize cache hit without MTP, set max_num_batched_tokens = lcm_block_size (2048) so each chunk = 1 block and every block boundary
has a valid state.