sampling: write master rows with a scalar fill_, not through a reused pinned row

mstar/engine/resources/sampler/utils.py

@@ -651,32 +651,31 @@ class Buffer:
     - ``buf``     ``[max_bs]``   per-step tensor, sliced to ``padded_bs`` and read
       by ``CudaGraphableSampler`` (its address must stay stable across replays).
     - ``master``  ``[capacity]`` slot-indexed cache, one row per active request.
-    - ``row_cpu`` ``[1]`` pinned staging for a single async H2D master-row write.
     """
     # ``[cg_slots, max_bs]``: one per-step row per double-buffer slot, so a
     # pre-plan gathering into one slot doesn't clobber the replay reading another
     buf: torch.Tensor
     master: torch.Tensor
-    row_cpu: torch.Tensor
     default: float
     dtype: torch.dtype
 
     @classmethod
     def allocate(
         cls, max_bs: int, capacity: int, device: torch.device,
-        dtype: torch.dtype, default: float, pinned: bool, cg_slots: int = 1,
+        dtype: torch.dtype, default: float, cg_slots: int = 1,
     ) -> "Buffer":
         return cls(
             buf=torch.full((cg_slots, max_bs), default, dtype=dtype, device=device),
             master=torch.full((capacity,), default, dtype=dtype, device=device),
-            row_cpu=torch.zeros(1, dtype=dtype, pin_memory=pinned),
             default=default,
             dtype=dtype,
         )
 
     def write_master_row(self, slot: int, value) -> None:
-        self.row_cpu[0] = value
-        self.master[slot:slot + 1].copy_(self.row_cpu, non_blocking=True)
+        # Staging through a shared pinned row +
+        # non_blocking H2D raced: the copy runs behind step N-1's kernels, so a
+        # second rid registered in the same step overwrote the row first.
+        self.master[slot:slot + 1].fill_(value)
 
     def grow_master(self, new_capacity: int) -> None:
         new = torch.full(
@@ -708,7 +707,7 @@ class MaskBuffer:
     """Three-tier storage for the per-request seen-token mask ``[*, V]`` (bool).
 
     Mirrors ``Buffer`` but 2-D and sourced from on-device ``SeenTokenMask``
-    tensors, so the master-row write is a GPU->GPU copy (no pinned staging).
+    tensors, so the master-row write is a GPU->GPU copy.
     """
     buf: torch.Tensor       # [cg_slots, max_bs, V] bool
     master: torch.Tensor    # [capacity, V] bool
@@ -753,8 +752,8 @@ class SamplerBuffers:
     """Pre-allocated static buffers for graph-safe sampling.
 
     Each per-request scalar parameter (temperature, top_k, top_p, seed,
-    repetition_penalty) is a ``Buffer`` owning a per-step slice, a slot-indexed
-    master cache, and pinned row staging. The RNG ``offset`` is also a ``Buffer``
+    repetition_penalty) is a ``Buffer`` owning a per-step slice and a slot-indexed
+    master cache. The RNG ``offset`` is also a ``Buffer``
     but round-trips on the GPU with no CPU middleman: gathered from its slot
     master before sampling, advanced in graph (``offset_buf += 1`` per sample),
     then scattered back to the master after replay — so a request's RNG stream
@@ -851,7 +850,7 @@ def allocate(
 
         def mk(dtype: torch.dtype, default: float, slots=cg_slots) -> Buffer:
             return Buffer.allocate(
-                max_batch_size, cap, device, dtype, default, pinned, slots
+                max_batch_size, cap, device, dtype, default, slots
             )
 
         # seen token mask is not double-buffered, as it depends on the GPU
@@ -899,9 +898,10 @@ def slice_for_bs(self, bs: int, cg_slot: int = 0) -> dict[str, Any]:
     # ------------------------------------------------------------------
 
     def _write_master_row(self, slot: int, cfg: SamplingConfig) -> None:
-        """Push one config row into each scalar master buffer via pinned H2D.
+        """Push one config row into each scalar master buffer.
 
-        Cheap async copies; only runs on register or actual config change
+        Five scalar fills, queued in stream order ahead of the gather that
+        reads them; only runs on register or actual config change
         (change-detection lives in ``update_request_config``). The seen-token
         mask is NOT written here (it changes every step — see
         ``update_request_config``).

test/modular/test_sampler_master_row_reuse.py

@@ -0,0 +1,50 @@
+"""Two requests registered in the same step must each get their own sampling row.
+
+``Buffer.write_master_row`` used to stage the value in one pinned ``[1]`` row
+per Buffer and queue a non_blocking H2D from it. The copy engine reads that row
+only when the stream reaches the copy -- behind whatever the GPU is still
+running (step N-1 under the worker's 2-step launch bound) -- while the host has
+long moved on to the next request's ``write_master_row`` on the same row. Rid A
+then landed rid B's temperature / top_k / top_p / seed / penalty. With ~20 ms of
+kernels queued ahead this failed 20/20 on an H100; a scalar ``fill_`` bakes the
+value into the launch and has no host row to reuse.
+
+GPU-only: the hazard needs a real asynchronous stream.
+"""
+
+from __future__ import annotations
+
+import sys
+
+sys.path.insert(0, ".")
+
+import pytest
+import torch
+
+from mstar.engine.resources.sampler.utils import Buffer
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="needs an asynchronous CUDA stream"
+)
+
+
+def test_back_to_back_master_row_writes_keep_their_values():
+    dev = torch.device("cuda")
+    buf = Buffer.allocate(
+        max_bs=4, capacity=8, device=dev, dtype=torch.float32, default=-1.0,
+    )
+    a = torch.randn(4096, 4096, device=dev)
+    b = torch.empty_like(a)
+
+    bad = 0
+    for _ in range(20):
+        buf.master.fill_(-1.0)
+        torch.cuda.synchronize()
+        for _ in range(10):  # ~20 ms of GPU work ahead of the writes: step N-1
+            torch.mm(a, a, out=b)
+        buf.write_master_row(0, 1.0)  # rid A registered
+        buf.write_master_row(1, 2.0)  # rid B, a few microseconds later
+        torch.cuda.synchronize()
+        if buf.master[0].item() != 1.0 or buf.master[1].item() != 2.0:
+            bad += 1
+    assert bad == 0, f"{bad}/20 trials landed the second value in the first slot"