[Core] Fix gpu<->cpu syncs in MRV2 mamba_hybrid.py (#49736)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Nick Hill committed
Jul 27, 2026 at 03:41 UTC
439f336212227833e126526d3c5f3ef3968dfbf5
1 file changed
+4
-4
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+4
-4
@@ -101,12 +101,12 @@ class MambaHybridModelState(DefaultModelState):
101
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
102
super().add_request(req_index, new_req_data)
103
# Must reset the speculative acceptance count in this idx which could be stale.
104
- self.num_accepted_tokens_gpu[req_index] = 1
104
+ self.num_accepted_tokens_gpu[req_index].fill_(1)
105
if self._align_mode:
106
# Seed the running state block from the resumed/prefilled position.
107
- self._mamba_state_idx_gpu[req_index] = (
108
- new_req_data.num_computed_tokens - 1
109
- ) // self.cache_config.block_size
107
+ self._mamba_state_idx_gpu[req_index].fill_(
108
+ (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size
109
+ )
110
111
def _get_mamba_group_info(
112
self, kv_cache_config: KVCacheConfig