[ModelRunnerV2] Fix scalar Mamba state update with int32 mappings (#50327)
Signed-off-by: Varun Shenoy <varun.vinayak.shenoy@oracle.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Varun Vinayak Shenoy committed
Aug 3, 2026 at 15:05 UTC
e578de311c2416e66c6a8b9819586bb91de10f9d
2 files changed
+58
-16
tests/v1/worker/test_mamba_hybrid_model_state.py
new
+29
@@ -0,0 +1,29 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
9
+
10
+
11
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
12
+@pytest.mark.parametrize(("num_sampled", "expected_value"), [(0, 1), (3, 3)])
13
+def test_postprocess_state_scalar_with_int32_mapping(
14
+ num_sampled: int, expected_value: int
15
+) -> None:
16
+ state = object.__new__(MambaHybridModelState)
17
+ state.num_accepted_tokens_gpu = torch.full(
18
+ (4,), 9, dtype=torch.int32, device="cuda"
19
+ )
20
+ state._align_mode = False
21
+ state._mamba_ctx = None
22
+ idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda")
23
+
24
+ state.postprocess_state(idx_mapping, num_sampled)
25
+
26
+ expected = torch.tensor(
27
+ [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda"
28
+ )
29
+ torch.testing.assert_close(state.num_accepted_tokens_gpu, expected)
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+29
-16
@@ -294,18 +294,20 @@ class MambaHybridModelState(DefaultModelState):
294
) -> None:
295
# Chunked prefill does not sample a token, so num_sampled can be 0.
296
# Mamba treats num_accepted_tokens=1 as the neutral non-spec value.
297
+ num_reqs = idx_mapping.shape[0]
298
+ if not num_reqs:
299
+ return
300
+
301
if not isinstance(num_sampled, int):
302
# idx_mapping may contain -1 sentinels (filtered rows) under PP; the
303
# kernel skips them rather than scattering with a host-side gather.
300
- n = idx_mapping.shape[0]
301
- if n:
302
- _scatter_num_accepted_kernel[(n,)](
303
- idx_mapping, num_sampled, self.num_accepted_tokens_gpu
304
- )
304
+ _scatter_num_accepted_kernel[(num_reqs,)](
305
+ idx_mapping, num_sampled, self.num_accepted_tokens_gpu
306
+ )
307
else:
308
# Fill with single value.
307
- self.num_accepted_tokens_gpu.index_fill_(
308
- 0, idx_mapping, max(num_sampled, 1)
309
+ _fill_num_accepted_kernel[(num_reqs,)](
310
+ idx_mapping, self.num_accepted_tokens_gpu, max(num_sampled, 1)
311
)
312
313
# Align: save the running state to the block-aligned position when
@@ -317,15 +319,13 @@ class MambaHybridModelState(DefaultModelState):
319
and num_computed_tokens is not None
320
and self._mamba_ctx is not None
321
):
320
- num_reqs = idx_mapping.shape[0]
321
- if num_reqs:
322
- self._mamba_ctx.run_fused_postprocess_align(
323
- num_reqs,
324
- self.num_accepted_tokens_gpu,
325
- self._mamba_state_idx_gpu,
326
- num_computed_tokens,
327
- idx_mapping,
328
- )
322
+ self._mamba_ctx.run_fused_postprocess_align(
323
+ num_reqs,
324
+ self.num_accepted_tokens_gpu,
325
+ self._mamba_state_idx_gpu,
326
+ num_computed_tokens,
327
+ idx_mapping,
328
+ )
329
330
331
@triton.jit
@@ -340,3 +340,16 @@ def _scatter_num_accepted_kernel(
340
return
341
num_sampled = tl.load(num_sampled_ptr + row)
342
tl.store(num_accepted_ptr + req_state_idx, tl.maximum(num_sampled, 1))
343
+
344
+
345
+@triton.jit
346
+def _fill_num_accepted_kernel(
347
+ idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip)
348
+ num_accepted_ptr, # [max_num_reqs]
349
+ num_sampled,
350
+):
351
+ row = tl.program_id(0)
352
+ req_state_idx = tl.load(idx_mapping_ptr + row)
353
+ if req_state_idx < 0:
354
+ return
355
+ tl.store(num_accepted_ptr + req_state_idx, num_sampled)