[qwen 3.8 next] Fuse exact greedy chain verification with argmax finalization

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

@@ -69,7 +69,7 @@ def _medium_argmax_partial_kernel(
     SPLITS: tl.constexpr,
     BLOCK: tl.constexpr,
 ):
-    row, part = tl.program_id(0), tl.program_id(1)
+    row, part = tl.program_id(0).to(tl.int64), tl.program_id(1)
     ix = part * BLOCK + tl.arange(0, BLOCK)
     valid = ix < N
     v = tl.load(X + row * SX + ix, valid, float("-inf"))
@@ -80,13 +80,23 @@ def _medium_argmax_partial_kernel(
 
 
 @triton.jit
-def _medium_argmax_final_kernel(PV, PI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexpr):
+def _medium_argmax_final_kernel(
+    PV,
+    PI,
+    OUT,
+    SPLITS: tl.constexpr,
+    BLOCK: tl.constexpr,
+    PROBS=None,
+    WRITE_PROBS: tl.constexpr = False,
+):
     row = tl.program_id(0)
     part = tl.arange(0, BLOCK)
     v = tl.load(PV + row * SPLITS + part, part < SPLITS, float("-inf"))
     i = tl.load(PI + row * SPLITS + part, part < SPLITS, 0x7FFFFFFF)
     _, index = tl.reduce((v, i), 0, _argmax_pair)
     tl.store(OUT + row, index.to(tl.int64))
+    if WRITE_PROBS:
+        tl.store(PROBS + row, 1.0)
 
 
 def row_argmax(x: torch.Tensor) -> torch.Tensor:
@@ -115,3 +125,104 @@ def row_argmax(x: torch.Tensor) -> torch.Tensor:
     )
     _argmax_final_kernel[(rows,)](pv, pi, out, SPLITS=_SPLITS, BLOCK=64, num_warps=2)
     return out
+
+
+def speculative_argmax(x: torch.Tensor, with_probs: bool = False):
+    assert x.ndim == 2 and x.dtype == torch.float32 and x.stride(1) == 1
+    rows, n = x.shape
+    assert n > 0
+    out = torch.empty((rows,), dtype=torch.int64, device=x.device)
+    probs = (
+        torch.empty((rows, 1), dtype=torch.float32, device=x.device)
+        if with_probs
+        else None
+    )
+    if rows:
+        block = 2048 if rows <= 64 else 4096 if rows <= 256 else 8192
+        splits = triton.cdiv(n, block)
+        pv = torch.empty((rows, splits), dtype=torch.float32, device=x.device)
+        pi = torch.empty((rows, splits), dtype=torch.int32, device=x.device)
+        _medium_argmax_partial_kernel[(rows, splits)](
+            x, pv, pi, n, x.stride(0), splits, block, num_warps=4
+        )
+        _medium_argmax_final_kernel[(rows,)](
+            pv,
+            pi,
+            out,
+            splits,
+            triton.next_power_of_2(splits),
+            PROBS=probs,
+            WRITE_PROBS=with_probs,
+            num_warps=1,
+        )
+    return (probs, out.view(rows, 1)) if with_probs else out
+
+
+@triton.jit
+def _argmax_verify_chain_kernel(
+    PV,
+    PI,
+    CANDIDATES,
+    TARGET,
+    PREDICT,
+    ACCEPT,
+    CORRECT,
+    WIDTH: tl.constexpr,
+    CSTRIDE: tl.constexpr,
+    SPLITS: tl.constexpr,
+    TOKENS: tl.constexpr,
+    PARTS: tl.constexpr,
+):
+    req = tl.program_id(0).to(tl.int64)
+    col = tl.arange(0, TOKENS)
+    part = tl.arange(0, PARTS)
+    row = req * WIDTH + col
+    valid = (col[:, None] < WIDTH) & (part[None, :] < SPLITS)
+    v = tl.load(PV + row[:, None] * SPLITS + part[None, :], valid, float("-inf"))
+    i = tl.load(PI + row[:, None] * SPLITS + part[None, :], valid, 0x7FFFFFFF)
+    _, target = tl.reduce((v, i), 1, _argmax_pair)
+    candidate = tl.load(CANDIDATES + req * CSTRIDE + col + 1, col + 1 < WIDTH, other=-1)
+    first = tl.min(
+        tl.where((col + 1 < WIDTH) & (candidate != target), col, WIDTH - 1), 0
+    )
+    tl.store(TARGET + row, target.to(tl.int64), col < WIDTH)
+    tl.store(PREDICT + row, tl.where(col <= first, target, 0), col < WIDTH)
+    tl.store(ACCEPT + row, tl.where(col <= first, row, -1), col < WIDTH)
+    tl.store(CORRECT + req, first)
+
+
+def greedy_verify_chain(logits, candidates):
+    bs, width = candidates.shape
+    rows, n = logits.shape
+    assert rows == bs * width and n > 0 and width > 0
+    assert (
+        logits.dtype == torch.float32 and logits.stride(1) == candidates.stride(1) == 1
+    )
+    predict = torch.empty((rows,), dtype=torch.int32, device=logits.device)
+    accept = torch.empty((bs, width), dtype=torch.int32, device=logits.device)
+    correct = torch.empty((bs,), dtype=torch.int32, device=logits.device)
+    target = torch.empty((bs, width), dtype=torch.int64, device=logits.device)
+    if bs:
+        block = 2048 if rows <= 64 else 4096 if rows <= 256 else 8192
+        splits = triton.cdiv(n, block)
+        pv = torch.empty((rows, splits), dtype=torch.float32, device=logits.device)
+        pi = torch.empty((rows, splits), dtype=torch.int32, device=logits.device)
+        _medium_argmax_partial_kernel[(rows, splits)](
+            logits, pv, pi, n, logits.stride(0), splits, block, num_warps=4
+        )
+        _argmax_verify_chain_kernel[(bs,)](
+            pv,
+            pi,
+            candidates,
+            target,
+            predict,
+            accept,
+            correct,
+            width,
+            candidates.stride(0),
+            splits,
+            triton.next_power_of_2(width),
+            triton.next_power_of_2(splits),
+            num_warps=4,
+        )
+    return predict, accept, correct, target

python/sglang/srt/speculative/eagle_utils.py

@@ -797,36 +797,66 @@ def eagle_sample(
         grammar_mask.apply(next_token_logits)
 
     candidates = verify_input.draft_token.reshape(bs, verify_input.draft_token_num)
-    predict_shape = list(next_token_logits.shape)[:-1]
-    predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
-    accept_index = torch.full(
-        (bs, verify_input.max_tree_depth), -1, dtype=torch.int32, device=device
-    )
-    num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device)
-
-    # Sample tokens
-    target_predict = None
     use_rejection_sampling = get_spec().speculative_use_rejection_sampling
-    if _verify_uses_greedy(
+    use_greedy = _verify_uses_greedy(
         is_all_greedy=sampling_info.is_all_greedy,
         is_cpu=_is_cpu,
         is_hip=_is_hip,
         is_xpu=_is_xpu,
         use_rejection_sampling=use_rejection_sampling,
-    ):
-        target_predict = torch.argmax(next_token_logits, dim=-1)
-        target_predict = target_predict.reshape(bs, verify_input.draft_token_num)
-        predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
-            predicts=predict,  # mutable
-            accept_index=accept_index,  # mutable
-            accept_token_num=num_correct_drafts,  # mutable
-            candidates=candidates,
-            retrieve_index=verify_input.retrieve_index,
-            retrieve_next_token=verify_input.retrieve_next_token,
-            retrieve_next_sibling=verify_input.retrieve_next_sibling,
-            target_predict=target_predict,
-            topk=verify_input.tree_topk,
+    )
+    fused_chain_greedy = (
+        use_greedy
+        and _is_cuda
+        and verify_input.tree_topk == 1
+        and verify_input.draft_token_num == verify_input.max_tree_depth
+        and next_token_logits.dtype == torch.float32
+        and next_token_logits.stride(-1) == 1
+        and next_token_logits.shape[-1] >= 131072
+        and candidates.stride(-1) == 1
+        and verify_input.draft_token_num <= 16
+    )
+    target_predict = None
+    if fused_chain_greedy:
+        from sglang.kernels.ops.speculative.row_argmax import greedy_verify_chain
+
+        predict, accept_index, num_correct_drafts, target_predict = greedy_verify_chain(
+            next_token_logits, candidates
+        )
+    else:
+        predict_shape = list(next_token_logits.shape)[:-1]
+        predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
+        accept_index = torch.full(
+            (bs, verify_input.max_tree_depth), -1, dtype=torch.int32, device=device
         )
+        num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device)
+
+    # Sample tokens
+    if use_greedy:
+        if not fused_chain_greedy:
+            if (
+                _is_cuda
+                and next_token_logits.dtype == torch.float32
+                and next_token_logits.stride(-1) == 1
+                and next_token_logits.shape[-1] >= 131072
+            ):
+                from sglang.kernels.ops.speculative.row_argmax import speculative_argmax
+
+                target_predict = speculative_argmax(next_token_logits)
+            else:
+                target_predict = torch.argmax(next_token_logits, dim=-1)
+            target_predict = target_predict.reshape(bs, verify_input.draft_token_num)
+            predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
+                predicts=predict,  # mutable
+                accept_index=accept_index,  # mutable
+                accept_token_num=num_correct_drafts,  # mutable
+                candidates=candidates,
+                retrieve_index=verify_input.retrieve_index,
+                retrieve_next_token=verify_input.retrieve_next_token,
+                retrieve_next_sibling=verify_input.retrieve_next_sibling,
+                target_predict=target_predict,
+                topk=verify_input.tree_topk,
+            )
 
         if _is_hip:
             # On ROCm, the per-rank draft tokens can differ, so ranks accept a

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -1253,10 +1253,19 @@ def _draft_extend_for_decode(
         elif self.topk == 1 and not _is_hip:
             # Gated to CUDA: see #26358 — ROCm's argmax tie-break corrupts
             # MTP draft selection on FP8 logits.
-            ret_topk_index = torch.argmax(
-                draft_logits_output.next_token_logits, dim=-1, keepdim=True
-            )
-            ret_topk_p = torch.ones_like(ret_topk_index, dtype=torch.float32)
+            logits = draft_logits_output.next_token_logits
+            if (
+                _is_cuda
+                and logits.dtype == torch.float32
+                and logits.stride(-1) == 1
+                and logits.shape[-1] >= 131072
+            ):
+                from sglang.kernels.ops.speculative.row_argmax import speculative_argmax
+
+                ret_topk_p, ret_topk_index = speculative_argmax(logits, with_probs=True)
+            else:
+                ret_topk_index = torch.argmax(logits, dim=-1, keepdim=True)
+                ret_topk_p = torch.ones_like(ret_topk_index, dtype=torch.float32)
             ret_draft_probs = None
         else:
             probs = renorm_draft_probs(

test/registered/unit/layers/test_next_sampling.py

@@ -0,0 +1,85 @@
+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("rows", [0, 1, 4, 128])
+@pytest.mark.parametrize("cols", [32771, 151936, 248320])
+def test_speculative_argmax_graph(rows, cols):
+    from sglang.kernels.ops.speculative.row_argmax import speculative_argmax
+
+    logits = torch.randn((rows * 2, cols), device="cuda")[::2]
+    speculative_argmax(logits)
+    speculative_argmax(logits, with_probs=True)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        ids = speculative_argmax(logits)
+        probabilities, ids2 = speculative_argmax(logits, with_probs=True)
+    for case in range(4):
+        if case == 1:
+            logits.fill_(-float("inf"))
+        elif case == 2:
+            logits[:, 2049] = float("inf")
+            logits[:, -1] = float("inf")
+        elif case == 3:
+            logits[:, 15] = float("nan")
+            logits[:, -17] = float("nan")
+        graph.replay()
+        expected = logits.argmax(dim=-1)
+        torch.testing.assert_close(ids, expected, rtol=0, atol=0)
+        torch.testing.assert_close(ids2[:, 0], expected, rtol=0, atol=0)
+        torch.testing.assert_close(
+            probabilities, torch.ones_like(probabilities), rtol=0, atol=0
+        )
+
+
+@pytest.mark.parametrize("bs", [1, 3, 128])
+@pytest.mark.parametrize("width", [1, 2, 4, 7])
+@pytest.mark.parametrize("vocab", [131073, 248320])
+def test_greedy_verify_chain_graph(bs, width, vocab):
+    from sglang.kernels.ops.speculative.row_argmax import greedy_verify_chain
+    from sglang.srt.speculative.eagle_utils import verify_tree_greedy_func
+
+    logits = torch.randn((bs * width, vocab), device="cuda")
+    candidates = torch.empty((bs * 2, width), device="cuda", dtype=torch.int64)[::2]
+    indices = torch.arange(bs * width, device="cuda").reshape(bs, width)
+    next_tokens = torch.arange(1, width + 1, device="cuda").repeat(bs, 1)
+    next_tokens[:, -1] = -1
+    siblings = torch.full_like(indices, -1)
+    greedy_verify_chain(logits, candidates)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = greedy_verify_chain(logits, candidates)
+    for case in range(3):
+        if case == 1:
+            logits.fill_(-float("inf"))
+            logits[:, -2:] = float("inf")
+        elif case == 2:
+            logits[:, 3] = float("nan")
+            logits[:, -1] = float("nan")
+        target = logits.argmax(-1).reshape(bs, width)
+        for accepted in range(width):
+            candidates[:, 0] = 17
+            candidates[:, 1:] = target[:, :-1]
+            if accepted < width - 1:
+                candidates[:, accepted + 1].add_(1).remainder_(vocab)
+            graph.replay()
+            expected = verify_tree_greedy_func(
+                predicts=torch.zeros(bs * width, device="cuda", dtype=torch.int32),
+                accept_index=torch.full(
+                    (bs, width), -1, device="cuda", dtype=torch.int32
+                ),
+                accept_token_num=torch.empty(bs, device="cuda", dtype=torch.int32),
+                candidates=candidates.contiguous(),
+                retrieve_index=indices,
+                retrieve_next_token=next_tokens,
+                retrieve_next_sibling=siblings,
+                target_predict=target,
+                topk=1,
+            )
+            for a, e in zip(actual[:3], expected):
+                torch.testing.assert_close(a, e, rtol=0, atol=0)
+            torch.testing.assert_close(actual[3], target, rtol=0, atol=0)