[qwen 3.8 next] Fuse speculative relay payload stores and avoid single-request gathers

python/sglang/kernels/ops/speculative/gather_spec_extras.py

@@ -115,3 +115,42 @@ def gather_spec_extras(
         BLOCK=block,
     )
     return topk_p, topk_index, bonus_tokens, hidden_states
+
+
+@triton.jit
+def _scatter_spec_extras_kernel(
+    indices,
+    sources,
+    destinations,
+    INDEX_STRIDE: tl.constexpr,
+    SOURCE_STRIDES: tl.constexpr,
+    WIDTHS: tl.constexpr,
+    SLOTS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    dst = tl.load(indices + row * INDEX_STRIDE).to(tl.int64)
+    dst = tl.where(dst < 0, dst + SLOTS, dst)
+    col = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    for i in tl.static_range(len(WIDTHS)):
+        value = tl.load(sources[i] + row * SOURCE_STRIDES[i] + col, col < WIDTHS[i])
+        tl.store(destinations[i] + dst * WIDTHS[i] + col, value, col < WIDTHS[i])
+
+
+def scatter_spec_extras(indices, pairs):
+    if indices.numel() == 0 or not pairs:
+        return
+    sources = tuple(src.reshape(src.shape[0], -1) for _, src in pairs)
+    sources = tuple(src if src.stride(1) == 1 else src.contiguous() for src in sources)
+    widths = tuple(src.shape[1] for src in sources)
+    block = min(1024, triton.next_power_of_2(max(widths)))
+    _scatter_spec_extras_kernel[(indices.numel(), triton.cdiv(max(widths), block))](
+        indices,
+        sources,
+        tuple(dst for dst, _ in pairs),
+        indices.stride(0),
+        tuple(src.stride(0) for src in sources),
+        widths,
+        pairs[0][0].shape[0],
+        BLOCK=block,
+    )

python/sglang/srt/managers/overlap_utils.py

@@ -6,7 +6,10 @@
 import msgspec
 import torch
 
-from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras
+from sglang.kernels.ops.speculative.gather_spec_extras import (
+    gather_spec_extras,
+    scatter_spec_extras,
+)
 from sglang.srt.environ import envs
 from sglang.srt.runtime_context import (
     get_exec,
@@ -421,6 +424,33 @@ def _resolve_spec_extras(self, batch: ScheduleBatch) -> None:
         indices = draft_input.future_indices
         if indices.shape[0] == 0:
             return
+        if (
+            indices.is_cuda
+            and indices.shape[0] == 1
+            and indices is batch.req_pool_indices
+            and len(batch.reqs) == 1
+            and not _DEBUG_ASSERT
+        ):
+            row = batch.reqs[0].kv.req_pool_idx
+            draft_input.bonus_tokens = self.output_tokens_buf.narrow(0, row, 1)
+            if self.need_topk:
+                draft_input.topk_p = self.topk_p_buf.narrow(0, row, 1)
+                draft_input.topk_index = self.topk_index_buf.narrow(0, row, 1)
+                if (
+                    self.draft_probs_buf is not None
+                    and draft_input.draft_probs is not None
+                ):
+                    draft_input.draft_probs = self.draft_probs_buf.narrow(0, row, 1)
+            if self.need_hidden_states:
+                draft_input.hidden_states = self.hidden_states_buf.narrow(0, row, 1)
+            if draft_input.future_dsa_topk_indices_available:
+                assert self.dsa_topk_indices_buf is not None
+                draft_input.dsa_topk_indices = self.dsa_topk_indices_buf.narrow(
+                    0, row, 1
+                )
+            else:
+                draft_input.dsa_topk_indices = None
+            return
         # FIXME: indices = batch.req_pool_indices, pinned 2 iters via
         # record_batch_in_overlap; record_stream here is redundant.
         indices.record_stream(torch.get_device_module(self.device).current_stream())
@@ -605,6 +635,26 @@ def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None:
             return
         self._maybe_init_forward_bufs(payload)
         self._maybe_init_dsa_topk_indices_buf(payload)
+        if indices.is_cuda:
+            pairs = [(self.output_tokens_buf, payload.bonus_tokens)]
+            if self.need_topk:
+                pairs.extend(
+                    [
+                        (self.topk_p_buf, payload.topk_p),
+                        (self.topk_index_buf, payload.topk_index),
+                    ]
+                )
+            if self.need_hidden_states:
+                pairs.append((self.hidden_states_buf, payload.hidden_states))
+            if self.draft_probs_buf is not None and payload.draft_probs is not None:
+                pairs.append((self.draft_probs_buf, payload.draft_probs))
+            if (
+                self.dsa_topk_indices_buf is not None
+                and payload.dsa_topk_indices is not None
+            ):
+                pairs.append((self.dsa_topk_indices_buf, payload.dsa_topk_indices))
+            scatter_spec_extras(indices, pairs)
+            return
         self.output_tokens_buf[indices] = payload.bonus_tokens.to(
             self.output_tokens_buf.dtype
         )

test/registered/unit/layers/test_next_relay.py

@@ -0,0 +1,48 @@
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+
+
+@pytest.mark.parametrize("bs", [0, 1, 129])
+@pytest.mark.parametrize("width", [1, 17, 2560])
+@pytest.mark.parametrize("strided", [False, True])
+def test_relay_scatter_graph(bs, width, strided):
+    from sglang.kernels.ops.speculative.gather_spec_extras import scatter_spec_extras
+
+    slots = bs + 3
+    indices = torch.randperm(slots, device="cuda")[:bs].repeat_interleave(2)[::2]
+    if bs:
+        indices[0] -= slots
+    pairs = []
+    for n, src_type, dst_type in [
+        (1, torch.int32, torch.int64),
+        (3, torch.float32, torch.float32),
+        (width, torch.float32, torch.bfloat16),
+        (7, torch.int64, torch.int64),
+    ]:
+        stride = 2 if strided else 1
+        src = (
+            torch.arange(bs * n * stride, device="cuda")
+            .reshape(bs, n * stride)[:, ::stride]
+            .to(src_type)
+        )
+        dst = torch.full((slots, n), -7, dtype=dst_type, device="cuda")
+        pairs.append((dst, src))
+    scatter_spec_extras(indices, pairs)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        scatter_spec_extras(indices, pairs)
+    for replay in range(3):
+        expected = []
+        for dst, src in pairs:
+            dst.fill_(-7)
+            src.add_(replay)
+            ref = dst.clone()
+            ref[indices] = src.to(dst.dtype)
+            expected.append(ref)
+        graph.replay()
+        for (actual, _), ref in zip(pairs, expected):
+            torch.testing.assert_close(actual, ref, rtol=0, atol=0)

test/registered/unit/managers/test_single_request_relay.py

@@ -0,0 +1,71 @@
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+
+from array import array
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.managers.overlap_utils import FutureMap
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.sampling.sampling_params import SamplingParams
+
+
+@pytest.mark.parametrize("hidden", [False, True])
+@pytest.mark.parametrize("topk", [False, True])
+@pytest.mark.parametrize("dsa", [False, True])
+@pytest.mark.parametrize("same", [False, True])
+def test_single_relay(hidden, topk, dsa, same):
+    f = object.__new__(FutureMap)
+    f.spec_algo = SimpleNamespace(is_ngram=lambda: False)
+    f.device = "cuda"
+    f.need_topk = topk
+    f.need_hidden_states = hidden
+    f.output_tokens_buf = torch.arange(9, device="cuda", dtype=torch.int64)
+    f.topk_p_buf = torch.randn(9, 3, device="cuda")
+    f.topk_index_buf = torch.randint(0, 100, (9, 3), device="cuda")
+    f.hidden_states_buf = torch.randn(9, 17, device="cuda", dtype=torch.bfloat16)
+    f.draft_probs_buf = torch.randn(9, 23, device="cuda")
+    f.dsa_topk_indices_buf = torch.randint(0, 100, (9, 4), device="cuda")
+    ids = torch.tensor([3], device="cuda", dtype=torch.int64)
+    d = SimpleNamespace(
+        future_indices=ids if same else ids.clone(),
+        draft_probs=torch.empty(0, device="cuda"),
+        future_dsa_topk_indices_available=dsa,
+    )
+    request = Req(
+        rid="relay",
+        origin_input_text="",
+        origin_input_ids=array("q", [1]),
+        sampling_params=SamplingParams(temperature=0, max_new_tokens=16),
+    )
+    request.kv.req_pool_idx = 3
+    b = SimpleNamespace(spec_info=d, req_pool_indices=ids, reqs=[request])
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        f._resolve_spec_extras(b)
+        fields = [("bonus_tokens", "output_tokens_buf")]
+        if topk:
+            fields.extend(
+                [
+                    ("topk_p", "topk_p_buf"),
+                    ("topk_index", "topk_index_buf"),
+                    ("draft_probs", "draft_probs_buf"),
+                ]
+            )
+        if hidden:
+            fields.append(("hidden_states", "hidden_states_buf"))
+        if dsa:
+            fields.append(("dsa_topk_indices", "dsa_topk_indices_buf"))
+        for k, v in fields:
+            assert (getattr(d, k).data_ptr() == getattr(f, v)[3:4].data_ptr()) == same
+        copies = [getattr(d, k).clone() for k, v in fields]
+        references = [getattr(f, v)[3:4].clone() for k, v in fields]
+        for k, v in fields:
+            getattr(f, v).fill_(-1)
+    stream.synchronize()
+    for x, y in zip(copies, references):
+        assert torch.equal(x, y)