[Bugfix] Fix /wake_up crash on hybrid models (Mamba/DeltaNet) (#41602)
Signed-off-by: Kevin Glynn <kevglynn@gmail.com>
Kevin Glynn committed
Jul 28, 2026 at 20:17 UTC
56f31af62afe6553369b14af225bfafb788e99a0
2 files changed
+66
-3
tests/v1/worker/test_gpu_model_runner.py
+56
@@ -1669,3 +1669,59 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
1669
1670
with pytest.raises(ValueError, match="max_num_seqs"):
1671
runner.initialize_kv_cache(kv_cache_config)
1672
+
1673
+
1674
+class TestInitFp8KvScalesHybridModels:
1675
+ """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries.
1676
+
1677
+ Hybrid models (Mamba, DeltaNet) store per-layer state as a list of tensors
1678
+ rather than a single tensor. init_fp8_kv_scales must iterate both forms.
1679
+ """
1680
+
1681
+ @staticmethod
1682
+ def _make_runner_stub(kv_caches):
1683
+ runner = Mock(spec=GPUModelRunner)
1684
+ runner.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3")
1685
+ runner.kv_caches = kv_caches
1686
+ runner.compilation_config = SimpleNamespace(static_forward_context={})
1687
+ runner.init_fp8_kv_scales = GPUModelRunner.init_fp8_kv_scales.__get__(
1688
+ runner, GPUModelRunner
1689
+ )
1690
+ return runner
1691
+
1692
+ def test_zeroes_both_tensor_and_list_entries(self):
1693
+ single_tensor = torch.ones(4, 8)
1694
+ list_tensors = [torch.ones(2, 4), torch.ones(3, 6)]
1695
+
1696
+ runner = self._make_runner_stub([single_tensor, list_tensors])
1697
+ runner.init_fp8_kv_scales()
1698
+
1699
+ assert (single_tensor == 0).all()
1700
+ assert all((t == 0).all() for t in list_tensors)
1701
+
1702
+ def test_skips_none_entries(self):
1703
+ tensor = torch.ones(4, 8)
1704
+ runner = self._make_runner_stub([None, tensor, None])
1705
+ runner.init_fp8_kv_scales()
1706
+
1707
+ assert (tensor == 0).all()
1708
+
1709
+ def test_noop_when_kv_cache_not_quantized(self):
1710
+ tensor = torch.ones(4, 8)
1711
+ runner = self._make_runner_stub([tensor])
1712
+ runner.cache_config.cache_dtype = "auto"
1713
+ runner.init_fp8_kv_scales()
1714
+
1715
+ assert (tensor == 1).all()
1716
+
1717
+ def test_mixed_none_tensor_and_list(self):
1718
+ t1 = torch.ones(2, 2)
1719
+ t2 = torch.ones(3, 3)
1720
+ list_entry = [torch.ones(1, 1), torch.ones(1, 1)]
1721
+
1722
+ runner = self._make_runner_stub([None, t1, list_entry, None, t2])
1723
+ runner.init_fp8_kv_scales()
1724
+
1725
+ assert (t1 == 0).all()
1726
+ assert (t2 == 0).all()
1727
+ assert all((t == 0).all() for t in list_entry)
vllm/v1/worker/gpu_model_runner.py
+10
-3
@@ -991,9 +991,16 @@ class GPUModelRunner(
991
return
992
993
kv_caches = getattr(self, "kv_caches", [])
994
- for cache_tensor in kv_caches:
995
- if cache_tensor is not None:
996
- cache_tensor.zero_()
994
+ for cache_entry in kv_caches:
995
+ if cache_entry is None:
996
+ continue
997
+ # Hybrid models (Mamba, DeltaNet) store per-layer state as a
998
+ # list of tensors rather than a single tensor.
999
+ if isinstance(cache_entry, list):
1000
+ for t in cache_entry:
1001
+ t.zero_()
1002
+ else:
1003
+ cache_entry.zero_()
1004
1005
k_attr_names = ("_k_scale", "k_scale")
1006
v_attr_names = ("_v_scale", "v_scale")