From 47cba233404494890f321f1df8c5c9bab4e5cb82 Mon Sep 17 00:00:00 2001 From: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:03:17 -0700 Subject: [PATCH 1/2] sampling: write master rows with a scalar fill_, not through a reused pinned row --- mstar/engine/resources/sampler/utils.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index dd9d10649d..fba80031b8 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/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``). From 5f0b7c84e9b294d6b056bbd195b6276bb25d1be5 Mon Sep 17 00:00:00 2001 From: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:03:17 -0700 Subject: [PATCH 2/2] test: back-to-back master-row writes keep their values (GPU) --- test/modular/test_sampler_master_row_reuse.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 test/modular/test_sampler_master_row_reuse.py diff --git a/test/modular/test_sampler_master_row_reuse.py b/test/modular/test_sampler_master_row_reuse.py new file mode 100644 index 0000000000..3f5fb54835 --- /dev/null +++ b/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"