From ae64d7f04d9d152e6fe9e4245d69288cefe1095a Mon Sep 17 00:00:00 2001
From: Mika Senghaas <mika@primeintellect.ai>
Date: Wed, 2 Sep 2026 01:37:55 +0000
Subject: [PATCH] Keep bounded sampling masks sparse

vLLM expands every sampling mask to a full-vocabulary bitset before the
asynchronous host copy. The engine then unpacks that bitset on the CPU for
every decode step. This stalls sampling replay even though training caps
each kept set at 512 tokens.

Capture bounded masks as fixed-width token IDs instead. Preserve the
existing scheduler and response formats, and use vLLM's native bitset for
wider masks.
---
 src/prime_rl/inference/patches.py | 71 +++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py
index e6c5ec3a1f..7dc3334a2f 100644
--- a/src/prime_rl/inference/patches.py
+++ b/src/prime_rl/inference/patches.py
@@ -17,6 +17,77 @@ def apply_shared_vllm_patches():
     monkey_patch_return_routed_experts_with_nixl_connector()
     monkey_patch_kv_xfer_finished_tolerate_freed()
     monkey_patch_online_fp8_parameter_cast()
+    monkey_patch_bounded_sampling_mask()
+
+
+def monkey_patch_bounded_sampling_mask():
+    """Keep bounded sampling masks sparse across the device-to-host copy."""
+    from contextvars import ContextVar
+    from typing import NamedTuple
+
+    import numpy as np
+    from vllm.v1.outputs import SamplingMaskLists
+    from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
+    from vllm.v1.worker.gpu.sample.sampler import Sampler
+
+    original_call = Sampler.__call__
+    original_from_logits = SamplingMaskTensors.from_logits
+    if getattr(original_call, "_prime_rl_uses_sparse_sampling_masks", False):
+        return
+
+    max_sparse_width = 512
+    use_sparse: ContextVar[bool] = ContextVar("use_sparse_sampling_mask", default=False)
+
+    class SparseSamplingMaskTensors(NamedTuple):
+        token_ids: torch.Tensor
+        counts: torch.Tensor
+
+        def to_cpu_nonblocking(self):
+            if self.token_ids.device.type == "cpu":
+                return self
+            return SparseSamplingMaskTensors(
+                self.token_ids.to("cpu", non_blocking=True),
+                self.counts.to("cpu", non_blocking=True),
+            )
+
+        def tolists(self, num_sampled_tokens: np.ndarray) -> SamplingMaskLists:
+            sampled_rows = np.flatnonzero(num_sampled_tokens)
+            counts = self.counts.numpy()[sampled_rows]
+            offsets = np.empty(len(counts) + 1, dtype=np.int64)
+            offsets[0] = 0
+            np.cumsum(counts, dtype=np.int64, out=offsets[1:])
+            token_ids = self.token_ids.numpy()[sampled_rows]
+            token_ids = token_ids[token_ids >= 0]
+            return SamplingMaskLists(
+                token_ids=token_ids,
+                offsets=offsets,
+                cu_num_generated_tokens=np.cumsum(np.concatenate(([0], num_sampled_tokens))).tolist(),
+            )
+
+    def _call(self, logits, input_batch):
+        top_k = self.sampling_states.top_k.np[input_batch.idx_mapping_np]
+        token = use_sparse.set(self.return_sampling_mask and int(top_k.max()) <= max_sparse_width)
+        try:
+            return original_call(self, logits, input_batch)
+        finally:
+            use_sparse.reset(token)
+
+    def _from_logits(cls, logits, num_sampled_tokens):
+        if not use_sparse.get():
+            return original_from_logits(logits, num_sampled_tokens)
+
+        width = min(max_sparse_width + 1, logits.shape[1])
+        values, token_ids = logits.topk(width, dim=-1)
+        finite = torch.isfinite(values)
+        overflow = finite[:, -1:] if width > max_sparse_width else torch.zeros_like(finite[:, :1])
+        keep = finite & ~overflow & (num_sampled_tokens[:, None] > 0)
+        token_ids = token_ids.to(torch.int32).masked_fill_(~keep, -1)
+        counts = keep.sum(dim=-1, dtype=torch.int32)
+        return SparseSamplingMaskTensors(token_ids, counts)
+
+    _call._prime_rl_uses_sparse_sampling_masks = True
+    Sampler.__call__ = _call
+    SamplingMaskTensors.from_logits = classmethod(_from_logits)
 
 
 def monkey_patch_online_fp8_parameter_cast():