[Bugfix][Model Runner V2] Preserve sampling masks in batch-sharded sampling

tests/test_config.py

@@ -753,9 +753,7 @@ def test_v1_model_runner_rejects_v2_only_features():
         VllmConfig._validate_v1_model_runner(config)
 
 
-def test_batch_sharded_sampling_rejects_return_sampling_mask():
-    """The batch-sharded gather drops sampling masks, so the combination must
-    fail loudly instead of returning ``sampling_mask=None``."""
+def test_batch_sharded_sampling_allows_return_sampling_mask():
     config = SimpleNamespace(
         parallel_config=SimpleNamespace(
             enable_batch_sharded_sampling=True, tensor_parallel_size=2
@@ -765,8 +763,8 @@ def test_batch_sharded_sampling_rejects_return_sampling_mask():
         speculative_config=None,
     )
 
-    with pytest.raises(ValueError, match="sampling masks"):
-        VllmConfig._validate_batch_sharded_sampling(config)
+    VllmConfig._validate_batch_sharded_sampling(config)
+    assert config.parallel_config.enable_batch_sharded_sampling
 
 
 @pytest.mark.skip_global_cleanup

tests/v1/worker/test_gpu_batch_shard.py

@@ -461,6 +461,119 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
         assert torch.equal(out.num_rejected.cpu(), expected_rejected.to(torch.int32))
 
 
+@requires_cuda
+def test_gather_sampler_output_sampling_masks():
+    """Sampling masks survive mixed-owner sharding with an empty rank."""
+    from dataclasses import replace as dc_replace
+
+    from vllm.v1.worker.gpu.sample.output import (
+        SamplerOutput,
+        SamplingMaskTensors,
+    )
+
+    tp_size = 3
+    rng = np.random.default_rng(0)
+    slot_pool = np.array([1, 2, 4, 5, 7, 8], dtype=np.int32)
+    batch = _make_batch(
+        rng,
+        num_reqs=len(slot_pool),
+        max_num_reqs=16,
+        max_spec=0,
+        slot_pool=slot_pool,
+    )
+    results = _shard_all_ranks(batch, max_num_reqs=16, tp_size=tp_size)
+    assert [result[2].num_local_reqs for result in results] == [0, 3, 3]
+    owners = batch.idx_mapping_np % tp_size
+    assert not np.array_equal(owners, np.sort(owners))
+
+    device = torch.device(DEVICE)
+    vocab_size = 32
+    compact_width = 4
+    packed_width = (vocab_size + 7) // 8
+    expected_token_ids = torch.empty(batch.num_reqs, compact_width, dtype=torch.int32)
+    expected_packed_masks = torch.empty(batch.num_reqs, packed_width, dtype=torch.uint8)
+    expected_counts = torch.empty(batch.num_reqs, dtype=torch.int32)
+    local_outputs: list[SamplerOutput | None] = []
+
+    for local, _, metadata in results:
+        if metadata.num_local_reqs == 0:
+            local_outputs.append(None)
+            continue
+        owned = torch.from_numpy(_owned_batch_indices(local))
+        token_ids = owned[:, None] * compact_width + torch.arange(compact_width)
+        packed_masks = owned[:, None] + torch.arange(packed_width)
+        counts = owned % compact_width + 1
+        expected_token_ids[owned] = token_ids.to(torch.int32)
+        expected_packed_masks[owned] = packed_masks.to(torch.uint8)
+        expected_counts[owned] = counts.to(torch.int32)
+        local_outputs.append(
+            SamplerOutput(
+                sampled_token_ids=torch.zeros(
+                    metadata.num_local_reqs, 1, dtype=torch.int64, device=device
+                ),
+                logprobs_tensors=None,
+                num_nans=None,
+                num_sampled=torch.ones(
+                    metadata.num_local_reqs, dtype=torch.int32, device=device
+                ),
+                num_rejected=torch.zeros(
+                    metadata.num_local_reqs, dtype=torch.int32, device=device
+                ),
+                sampling_mask_tensors=SamplingMaskTensors(
+                    token_ids=token_ids.to(device=device, dtype=torch.int32),
+                    packed_mask=packed_masks.to(device=device, dtype=torch.uint8),
+                    counts=counts.to(device=device, dtype=torch.int32),
+                    vocab_size=vocab_size,
+                ),
+            )
+        )
+
+    recorded: dict[int, list[torch.Tensor]] = {}
+    gathered: list[torch.Tensor] = []
+
+    def run(rank: int) -> SamplerOutput:
+        metadata = dc_replace(
+            results[rank][2],
+            gathered_src_indices=results[rank][2].gathered_src_indices.to(device),
+        )
+        call = 0
+
+        def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
+            nonlocal call
+            result = gathered[call] if gathered else torch.cat([x] * tp_size)
+            call += 1
+            if not gathered:
+                recorded.setdefault(rank, []).append(x.clone())
+            return result
+
+        with mock.patch.object(
+            batch_shard, "tensor_model_parallel_all_gather", fake_all_gather
+        ):
+            return batch_shard.gather_sampler_output(
+                local_outputs[rank],
+                metadata,
+                device,
+                global_batch=batch,
+                local_batch=results[rank][0],
+                sampling_mask_dims=(vocab_size, compact_width),
+            )
+
+    for rank in range(tp_size):
+        run(rank)
+    gathered.extend(
+        torch.cat([recorded[rank][i] for rank in range(tp_size)]) for i in range(4)
+    )
+
+    for rank in range(tp_size):
+        output = run(rank)
+        assert output.sampling_mask_tensors is not None
+        mask = output.sampling_mask_tensors
+        assert mask.vocab_size == vocab_size
+        assert torch.equal(mask.token_ids.cpu(), expected_token_ids)
+        assert torch.equal(mask.packed_mask.cpu(), expected_packed_masks)
+        assert torch.equal(mask.counts.cpu(), expected_counts)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_gather_sampler_output_logprobs_and_nans():
     """num_nans is reduced per request and gathered; LogprobsTensors are

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -18,6 +18,62 @@
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+
+
+def test_batch_sharding_preserves_custom_sampler_call_interface(monkeypatch):
+    """Mask plumbing must not change custom samplers when masks are disabled."""
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.device = torch.device("cpu")
+    runner.vocab_size = 4
+    runner.model_config = SimpleNamespace(return_sampling_mask=False)
+    runner.rejection_sampler = None
+    runner.speculator = None
+
+    global_batch = SimpleNamespace(idx_mapping_np=[0], num_draft_tokens=0)
+    local_batch = SimpleNamespace(num_reqs=1, num_draft_tokens=0)
+    shard_metadata = object()
+    runner.batch_sharder = SimpleNamespace(
+        shard_sampler_inputs=lambda *_args: (
+            local_batch,
+            torch.tensor([0]),
+            None,
+            shard_metadata,
+        )
+    )
+    runner.model = SimpleNamespace(compute_logits_local=lambda hidden: hidden)
+
+    class CustomSampler:
+        compute_nans = False
+
+        def __init__(self):
+            self.called = False
+
+        def get_logprobs_dims(self, *_args, **_kwargs):
+            return None
+
+        def __call__(self, logits, input_batch):
+            self.called = True
+            one = torch.ones(1, dtype=torch.int32)
+            return SamplerOutput(
+                sampled_token_ids=torch.ones(1, 1, dtype=torch.int64),
+                logprobs_tensors=None,
+                num_nans=None,
+                num_sampled=one,
+                num_rejected=torch.zeros_like(one),
+            )
+
+    runner.sampler = CustomSampler()
+    monkeypatch.setattr(model_runner_module, "all_to_all_logits", lambda x, _: x)
+    monkeypatch.setattr(
+        model_runner_module,
+        "gather_sampler_output",
+        lambda output, *_args, **_kwargs: output,
+    )
+
+    GPUModelRunner.sample(runner, torch.ones(1, runner.vocab_size), global_batch, None)
+
+    assert runner.sampler.called
 
 
 def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):

tests/v1/worker/test_gpu_sampler_flags.py

@@ -23,7 +23,7 @@ class MockReasoningConfig:
     natural_reasoning_end_token_ids = [91]
 
 
-def _make_sampler() -> Sampler:
+def _make_sampler(return_sampling_mask: bool = False) -> Sampler:
     req_states = RequestState(
         max_num_reqs=4,
         max_model_len=64,
@@ -38,6 +38,7 @@ def _make_sampler() -> Sampler:
         device=DEVICE,
         req_states=req_states,
         reasoning_config=MockReasoningConfig(),
+        return_sampling_mask=return_sampling_mask,
     )
 
 
@@ -88,3 +89,12 @@ def test_logits_processing_cache_only_checks_active_requests():
 
     assert not np.any(sampler.needs_logits_processing[sampling_only])
     assert np.any(sampler.needs_logits_processing[with_processing])
+
+
+def test_sampling_mask_width_uses_all_active_requests():
+    sampler = _make_sampler(return_sampling_mask=True)
+    sampler.add_request(0, 1, SamplingParams(top_k=3))
+    sampler.add_request(2, 1, SamplingParams(top_k=7))
+
+    assert sampler.get_sampling_mask_width(np.array([0], dtype=np.int32)) == 3
+    assert sampler.get_sampling_mask_width(np.array([0, 2], dtype=np.int32)) == 7

vllm/config/vllm.py

@@ -2703,13 +2703,6 @@ def _validate_batch_sharded_sampling(self) -> None:
             # fixed-width logprobs gather cannot reasonably size for.
             blockers.append("max_logprobs is -1, allowing vocab-size logprob requests")
 
-        if self.model_config is not None and self.model_config.return_sampling_mask:
-            # gather_sampler_output() drops SamplingMaskTensors: masks come back None.
-            blockers.append(
-                "return_sampling_mask is set and the batch-sharded gather does "
-                "not forward sampling masks"
-            )
-
         if (
             self.speculative_config is not None
             and self.speculative_config.enable_adaptive_verification

vllm/v1/worker/gpu/model_runner.py

@@ -1477,13 +1477,26 @@ def sample(
             )
 
         sampler_output: SamplerOutput | None
+        sampling_mask_width = None
+        if shard_metadata is not None and self.model_config.return_sampling_mask:
+            assert self.sampler is not None
+            sampling_mask_width = self.sampler.get_sampling_mask_width(
+                global_input_batch.idx_mapping_np
+            )
         if input_batch.num_reqs == 0:
             # This rank owns no requests this step. It contributes an
             # all-padding block to the gather below.
             sampler_output = None
         elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
             assert self.sampler is not None
-            sampler_output = self.sampler(logits, input_batch)
+            if sampling_mask_width is None:
+                sampler_output = self.sampler(logits, input_batch)
+            else:
+                sampler_output = self.sampler(
+                    logits,
+                    input_batch,
+                    sampling_mask_width=sampling_mask_width,
+                )
         else:
             # Rejection sampling for spec decoding.
             assert self.rejection_sampler is not None
@@ -1514,6 +1527,11 @@ def sample(
                         or self.rejection_sampler is None
                     ),
                 ),
+                sampling_mask_dims=(
+                    (self.vocab_size, sampling_mask_width)
+                    if sampling_mask_width is not None
+                    else None
+                ),
             )
 
         assert sampler_output is not None

vllm/v1/worker/gpu/sample/batch_shard.py

@@ -13,7 +13,7 @@
 from vllm.v1.core.sched.output import GrammarOutput
 from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.worker.gpu.input_batch import InputBatch
-from vllm.v1.worker.gpu.sample.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
 
 
 @dataclass
@@ -591,6 +591,64 @@ def _gather_logprobs_tensors(
     )
 
 
+def _gather_sampling_mask_tensors(
+    local_output: SamplerOutput | None,
+    metadata: BatchShardMetadata,
+    device: torch.device,
+    vocab_size: int,
+    compact_width: int,
+) -> SamplingMaskTensors:
+    """Gather and restore sampling masks from the owner-sharded requests."""
+    packed_width = (vocab_size + 7) // 8
+    token_ids = torch.zeros(
+        metadata.max_num_reqs_per_rank,
+        compact_width,
+        dtype=torch.int32,
+        device=device,
+    )
+    packed_mask = torch.zeros(
+        metadata.max_num_reqs_per_rank,
+        packed_width,
+        dtype=torch.uint8,
+        device=device,
+    )
+    counts = torch.zeros(
+        metadata.max_num_reqs_per_rank,
+        dtype=torch.int32,
+        device=device,
+    )
+
+    if local_output is not None:
+        assert local_output.sampling_mask_tensors is not None
+        local_mask = local_output.sampling_mask_tensors
+        assert local_mask.vocab_size == vocab_size
+        assert local_mask.token_ids.shape == (
+            metadata.num_local_reqs,
+            compact_width,
+        )
+        assert local_mask.packed_mask.shape == (
+            metadata.num_local_reqs,
+            packed_width,
+        )
+        assert local_mask.counts.shape == (metadata.num_local_reqs,)
+        token_ids[: metadata.num_local_reqs].copy_(local_mask.token_ids)
+        packed_mask[: metadata.num_local_reqs].copy_(local_mask.packed_mask)
+        counts[: metadata.num_local_reqs].copy_(local_mask.counts)
+
+    # Every rank must participate even when it owns no requests. The zero-padded
+    # rows are never selected by gathered_src_indices.
+    gathered_token_ids = tensor_model_parallel_all_gather(token_ids, dim=0)
+    gathered_packed_mask = tensor_model_parallel_all_gather(packed_mask, dim=0)
+    gathered_counts = tensor_model_parallel_all_gather(counts, dim=0)
+    src_indices = metadata.gathered_src_indices
+    return SamplingMaskTensors(
+        token_ids=gathered_token_ids[src_indices],
+        packed_mask=gathered_packed_mask[src_indices],
+        counts=gathered_counts[src_indices],
+        vocab_size=vocab_size,
+    )
+
+
 def gather_sampler_output(
     local_output: SamplerOutput | None,
     metadata: BatchShardMetadata,
@@ -599,6 +657,7 @@ def gather_sampler_output(
     local_batch: InputBatch,
     gather_num_nans: bool = False,
     logprobs_dims: tuple[int, int] | None = None,
+    sampling_mask_dims: tuple[int, int] | None = None,
 ) -> SamplerOutput:
     max_num_logits_per_req = metadata.max_num_logits_per_req
     num_packed_cols = max_num_logits_per_req + 2 + (1 if gather_num_nans else 0)
@@ -642,6 +701,17 @@ def gather_sampler_output(
             local_output, metadata, global_batch, local_batch, logprobs_dims, device
         )
 
+    sampling_mask_tensors = None
+    if sampling_mask_dims is not None:
+        vocab_size, compact_width = sampling_mask_dims
+        sampling_mask_tensors = _gather_sampling_mask_tensors(
+            local_output,
+            metadata,
+            device,
+            vocab_size,
+            compact_width,
+        )
+
     # Unpack the gathered tensor into a sampler output object.
     num_reqs = metadata.gathered_src_indices.shape[0]
     sampled_token_ids = torch.empty(
@@ -671,4 +741,5 @@ def gather_sampler_output(
         num_nans=num_nans,
         num_sampled=num_sampled,
         num_rejected=num_rejected,
+        sampling_mask_tensors=sampling_mask_tensors,
     )