[Feature][KVOffload][P/D] Recompute offload supports linear hybrid model (vllm-project#12881)

### What this PR does / why we need it?
Recompute offload supports linear hybrid model, such as Qwen3.5. When
the MTP accept length changes, the offload Mamba cache block
corresponding to the accept length needs to be placed at the start
position of mamba cache groups, and this logic is implemented in this
PR.

Specifically, when preemption occurs, the last `1+num_spec_tokens` mamba
blocks are offloaded. When the blocks are reloaded, the index of the
mamba block that needs to be reloaded is calculated based on the actual
length of the last scheduling, that is, `accept_token_idx =
self.num_spec_tokens - (state.num_computed_tokens - request.num_tokens +
1)`. In this case, we only need to reload the block to the 0th mamba
block index of the resumed request.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
In local self-verification, I use Qwen3.6 35B-A3B for test. The GPU
memory usage of D node is set to 0.75, and the DP1TP2 parallel
configuration is used to trigger recompute offload at a high frequency.
The verification result shows that the score of the GPQA dataset meets
the expectation.
```
| dataset | version | metric | mode | vllm-api-general-chat |
|----- | ----- | ----- | ----- | -----|
| GPQA_diamond | b1ed2c | accuracy | gen | 85.86 |
```


- vLLM version: v0.25.1
- vLLM main:
vllm-project/vllm@fe784ff

---------

Signed-off-by: nwpu-zxr <zhouxuerong2@huawei.com>

Author: nwpu-zxr

tests/ut/kv_offload/test_recompute_cpu_offload.py

@@ -197,6 +197,7 @@ def test_recompute_cpu_offload_scheduler_aligns_sliding_window_blocks():
 def test_recompute_cpu_offload_scheduler_d2h_keeps_sliding_window_offsets():
     scheduler = RecomputeCPUOffloadScheduler.__new__(RecomputeCPUOffloadScheduler)
     scheduler._group_is_sliding_window = [True]
+    scheduler._group_is_mamba = [False]
     scheduler.cpu_kv_cache_config = SimpleNamespace(
         kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))]
     )
@@ -231,6 +232,7 @@ def test_recompute_cpu_offload_scheduler_d2h_keeps_sliding_window_offsets():
 def test_recompute_cpu_offload_scheduler_h2d_skips_sliding_window_null_blocks():
     scheduler = RecomputeCPUOffloadScheduler.__new__(RecomputeCPUOffloadScheduler)
     scheduler._group_is_sliding_window = [True]
+    scheduler._group_is_mamba = [False]
     scheduler.cpu_kv_cache_config = SimpleNamespace(
         kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))]
     )
@@ -262,6 +264,7 @@ def test_recompute_cpu_offload_scheduler_h2d_skips_sliding_window_null_blocks():
 def test_recompute_cpu_offload_scheduler_h2d_clips_mtp_tail_blocks():
     scheduler = RecomputeCPUOffloadScheduler.__new__(RecomputeCPUOffloadScheduler)
     scheduler._group_is_sliding_window = [False]
+    scheduler._group_is_mamba = [False]
     scheduler.cpu_kv_cache_config = SimpleNamespace(
         kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))]
     )

vllm_ascend/core/recompute_scheduler.py

@@ -993,7 +993,7 @@ def update_from_output(
         for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
             assert num_tokens_scheduled > 0
             request = self.requests.get(req_id)
-            if request is not None:
+            if not vllm_version_is("0.25.1") and request is not None:
                 request.num_in_flight_tokens -= num_tokens_scheduled
             if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
                 # skip failed or rescheduled requests from KV load failure

vllm_ascend/distributed/kv_transfer/kv_pool/recompute_cpu_offload/manager.py

@@ -18,7 +18,7 @@
 )
 from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.kv_cache_interface import SlidingWindowSpec, UniformTypeKVCacheSpecs
+from vllm.v1.kv_cache_interface import MambaSpec, SlidingWindowSpec, UniformTypeKVCacheSpecs
 from vllm.v1.outputs import KVConnectorOutput
 
 from vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.metadata import (
@@ -70,9 +70,13 @@ def __init__(
         assert kv_cache_config is not None
         self.vllm_config = vllm_config
         self.enable_offload_prefix_caching = enable_offload_prefix_caching
+        self.num_spec_tokens = (
+            vllm_config.speculative_config.num_speculative_tokens if vllm_config.speculative_config else 0
+        )
         self.cpu_kv_cache_config = self._derive_cpu_config(kv_cache_config, cpu_capacity_bytes)
         self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
         self._group_is_sliding_window = self._get_group_is_sliding_window(kv_cache_config)
+        self._group_is_mamba = self._get_group_is_mamba(kv_cache_config)
         self.enable_kv_cache_events = (
             vllm_config.kv_events_config is not None and vllm_config.kv_events_config.enable_kv_cache_events
         )
@@ -129,6 +133,18 @@ def _get_group_is_sliding_window(kv_cache_config: "KVCacheConfig") -> list[bool]
                 group_is_sliding_window.append(isinstance(group.kv_cache_spec, SlidingWindowSpec))
         return group_is_sliding_window
 
+    @staticmethod
+    def _get_group_is_mamba(kv_cache_config: "KVCacheConfig") -> list[bool]:
+        group_is_mamba: list[bool] = []
+        for group in kv_cache_config.kv_cache_groups:
+            if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs):
+                group_is_mamba.append(
+                    any(isinstance(spec, MambaSpec) for spec in group.kv_cache_spec.kv_cache_specs.values())
+                )
+            else:
+                group_is_mamba.append(isinstance(group.kv_cache_spec, MambaSpec))
+        return group_is_mamba
+
     @staticmethod
     def _derive_cpu_config(gpu_config: "KVCacheConfig", cpu_capacity_bytes: int) -> "KVCacheConfig":
         from vllm.v1.kv_cache_interface import KVCacheConfig as KVCacheConfigCls
@@ -247,53 +263,49 @@ def _create_preempt_state(
         group_gpu_hashes: list[list[BlockHashWithGroupId | None]] = []
         missing_hashes: set[BlockHashWithGroupId] = set()
         num_unhashed = 0
+        num_mamba_blocks = 0
 
         for g, group_gpu_ids in enumerate(block_ids_by_group):
-            group_block_size = kv_cache_groups[g].kv_cache_spec.block_size
-            logical_num_blocks = cdiv(num_computed_tokens, group_block_size)
-            aligned_group_gpu_ids = self._align_group_block_ids(g, group_gpu_ids, logical_num_blocks)
-            eviction_group_gpu_ids = self._align_group_block_ids(
-                g,
-                group_gpu_ids,
-                max(logical_num_blocks, len(group_gpu_ids)),
-            )
             gpu_blocks: list[KVCacheBlock | None] = []
             effective_hashes: list[BlockHashWithGroupId | None] = []
-
-            for block_idx, block_id in enumerate(eviction_group_gpu_ids):
-                if block_id <= 0:
-                    continue
-                gpu_block = self._gpu_block_pool.blocks[block_id]
-                block_is_computed = (block_idx + 1) * group_block_size <= num_computed_tokens
-                if not block_is_computed and gpu_block.block_hash is not None:
-                    # allocate_slots() may assign a hash using tokens planned
-                    # for this scheduling step. If the request is then
-                    # preempted before forward, that block does not contain the
-                    # hashed KV and must not remain in the GPU prefix cache.
-                    self._gpu_block_pool._maybe_evict_cached_block(gpu_block)
-
-            for block_idx, block_id in enumerate(aligned_group_gpu_ids):
-                if block_id <= 0:
-                    gpu_blocks.append(None)
-                    effective_hashes.append(None)
-                    continue
-
-                gpu_block = self._gpu_block_pool.blocks[block_id]
-                block_is_computed = (block_idx + 1) * group_block_size <= num_computed_tokens
-                block_hash = gpu_block.block_hash if block_is_computed and self.enable_offload_prefix_caching else None
-                gpu_blocks.append(gpu_block)
-                effective_hashes.append(block_hash)
-                if block_hash is None:
-                    num_unhashed += 1
-                elif (
-                    self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash) is None
-                    and block_hash not in self._pending_hash_blocks
-                ):
-                    missing_hashes.add(block_hash)
+            if self._group_is_mamba[g]:
+                # For Mamba cache, only the last `1 + num_spec_tokens` blocks need to offload
+                offload_start_idx = len(group_gpu_ids) - self.num_spec_tokens - 1
+                for block_idx, block_id in enumerate(group_gpu_ids):
+                    if block_idx >= offload_start_idx:
+                        num_mamba_blocks += 1
+                        gpu_block = self._gpu_block_pool.blocks[block_id]
+                        gpu_blocks.append(gpu_block)
+                        effective_hashes.append(None)
+            else:
+                group_block_size = kv_cache_groups[g].kv_cache_spec.block_size
+                logical_num_blocks = cdiv(num_computed_tokens, group_block_size)
+                aligned_group_gpu_ids = self._align_group_block_ids(g, group_gpu_ids, logical_num_blocks)
+
+                for block_idx, block_id in enumerate(aligned_group_gpu_ids):
+                    if block_id <= 0:
+                        gpu_blocks.append(None)
+                        effective_hashes.append(None)
+                        continue
+
+                    gpu_block = self._gpu_block_pool.blocks[block_id]
+                    block_is_computed = (block_idx + 1) * group_block_size <= num_computed_tokens
+                    block_hash = (
+                        gpu_block.block_hash if block_is_computed and self.enable_offload_prefix_caching else None
+                    )
+                    gpu_blocks.append(gpu_block)
+                    effective_hashes.append(block_hash)
+                    if block_hash is None:
+                        num_unhashed += 1
+                    elif (
+                        self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash) is None
+                        and block_hash not in self._pending_hash_blocks
+                    ):
+                        missing_hashes.add(block_hash)
             group_gpu_blocks.append(gpu_blocks)
             group_gpu_hashes.append(effective_hashes)
 
-        num_needed = num_unhashed + len(missing_hashes)
+        num_needed = num_unhashed + len(missing_hashes) + num_mamba_blocks
         if not any(any(gpu_block is not None for gpu_block in group) for group in group_gpu_blocks):
             return False
         if num_needed > self.cpu_block_pool.get_num_free_blocks():
@@ -410,45 +422,50 @@ def _prepare_preempt_load_after_alloc(
         gpu_block_ids: list[int] = []
         cpu_block_ids: list[int] = []
         for g, group_cpu_ids in enumerate(state.cpu_block_ids):
-            group_block_size = self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size
-            start_block = load_start_tokens // group_block_size
-            end_block = min(
-                len(group_cpu_ids),
-                len(
-                    self._align_group_block_ids(
-                        g,
-                        block_ids_by_group[g],
-                        max(
-                            cdiv(load_end_tokens, group_block_size),
-                            len(block_ids_by_group[g]),
-                        ),
-                    )
-                ),
-                cdiv(load_end_tokens, group_block_size),
-            )
-            if end_block == start_block:
-                continue
-            if end_block < start_block:
-                raise RuntimeError(
-                    "Recompute H2D produced an empty block range: "
-                    f"req_id={request.request_id}, group={g}, "
-                    f"start_block={start_block}, end_block={end_block}, "
-                    f"gpu_blocks={len(block_ids_by_group[g])}, "
-                    f"cpu_blocks={len(group_cpu_ids)}"
+            if self._group_is_mamba[g]:
+                accept_token_idx = self.num_spec_tokens - (state.num_computed_tokens - request.num_tokens + 1)
+                cpu_block_ids.append(group_cpu_ids[accept_token_idx])
+                gpu_block_ids.append(block_ids_by_group[g][0])
+            else:
+                group_block_size = self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size
+                start_block = load_start_tokens // group_block_size
+                end_block = min(
+                    len(group_cpu_ids),
+                    len(
+                        self._align_group_block_ids(
+                            g,
+                            block_ids_by_group[g],
+                            max(
+                                cdiv(load_end_tokens, group_block_size),
+                                len(block_ids_by_group[g]),
+                            ),
+                        )
+                    ),
+                    cdiv(load_end_tokens, group_block_size),
                 )
-
-            aligned_group_gpu_ids = self._align_group_block_ids(
-                g,
-                block_ids_by_group[g],
-                end_block,
-            )
-            for block_idx in range(start_block, end_block):
-                cpu_block_id = group_cpu_ids[block_idx]
-                gpu_block_id = aligned_group_gpu_ids[block_idx]
-                if cpu_block_id <= 0 or gpu_block_id <= 0:
+                if end_block == start_block:
                     continue
-                cpu_block_ids.append(cpu_block_id)
-                gpu_block_ids.append(gpu_block_id)
+                if end_block < start_block:
+                    raise RuntimeError(
+                        "Recompute H2D produced an empty block range: "
+                        f"req_id={request.request_id}, group={g}, "
+                        f"start_block={start_block}, end_block={end_block}, "
+                        f"gpu_blocks={len(block_ids_by_group[g])}, "
+                        f"cpu_blocks={len(group_cpu_ids)}"
+                    )
+
+                aligned_group_gpu_ids = self._align_group_block_ids(
+                    g,
+                    block_ids_by_group[g],
+                    end_block,
+                )
+                for block_idx in range(start_block, end_block):
+                    cpu_block_id = group_cpu_ids[block_idx]
+                    gpu_block_id = aligned_group_gpu_ids[block_idx]
+                    if cpu_block_id <= 0 or gpu_block_id <= 0:
+                        continue
+                    cpu_block_ids.append(cpu_block_id)
+                    gpu_block_ids.append(gpu_block_id)
 
         if not cpu_block_ids or len(cpu_block_ids) != len(gpu_block_ids):
             raise RuntimeError(