[1/3] [EAGLE] perf: Fuse topk=1 draft postprocess (#30947)

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

@@ -18,6 +18,7 @@ _TRITON_KERNELS = [
     ("multi_layer_eagle", "rotate_input_ids_triton"),
     ("spec_tree", "sgl_build_tree_kernel_efficient_triton"),
     ("spec_tree", "verify_tree_greedy_kernel_triton"),
+    ("topk1", "draft_topk1_postprocess"),
 ]
 for _mod, _fn in _TRITON_KERNELS:
     register_kernel(

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

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_DRAFT_TOPK1_BLOCK = 8192
+
+
+@triton.jit
+def _draft_topk1_partial_argmax_kernel(
+    logits,
+    partial_vals,
+    partial_indices,
+    logits_row_stride,
+    vocab_size: tl.constexpr,
+    num_splits: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # int64 row base: row * stride overflows int32 once bs * vocab reaches 2^31.
+    row = tl.program_id(0).to(tl.int64)
+    split = tl.program_id(1)
+    offsets = split * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < vocab_size
+    vals = tl.load(
+        logits + row * logits_row_stride + offsets,
+        mask=mask,
+        other=-float("inf"),
+    ).to(tl.float32)
+
+    max_val = tl.max(vals, axis=0)
+    local_index = tl.argmax(vals, axis=0)
+    out_offset = row * num_splits + split
+    tl.store(partial_vals + out_offset, max_val)
+    tl.store(partial_indices + out_offset, split * BLOCK + local_index)
+
+
+@triton.jit
+def _draft_topk1_finalize_kernel(
+    partial_vals,
+    partial_indices,
+    topk_p,
+    topk_index,
+    positions,
+    draft_tokens,
+    draft_tokens_stride,
+    draft_token_column,
+    num_splits: tl.constexpr,
+    WRITE_DRAFT_TOKEN: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK)
+    mask = offsets < num_splits
+    vals = tl.load(
+        partial_vals + row * num_splits + offsets,
+        mask=mask,
+        other=-float("inf"),
+    )
+
+    split = tl.argmax(vals, axis=0)
+    index = tl.load(partial_indices + row * num_splits + split).to(tl.int64)
+    tl.store(topk_index + row, index)
+    tl.store(topk_p + row, 1.0)
+    if WRITE_DRAFT_TOKEN:
+        tl.store(draft_tokens + row * draft_tokens_stride + draft_token_column, index)
+
+    position = tl.load(positions + row)
+    tl.store(positions + row, position + 1)
+
+
+def draft_topk1_postprocess(
+    next_token_logits: torch.Tensor,
+    positions: torch.Tensor,
+    draft_tokens: torch.Tensor | None = None,
+    draft_token_column: int = 0,
+):
+    """Argmax draft logits for topk=1 and advance positions.
+
+    PyTorch eager argmax reduces each row with too little parallelism for the
+    GLM/DSV4 vocab widths in CUDA graph replay. This split reduction exposes
+    the vocab dimension across CTAs, then finalizes one token per row.
+
+    If ``draft_tokens`` is given, the finalize kernel also stores the argmax
+    into ``draft_tokens[:, draft_token_column]``, mutating the caller-owned
+    buffer in place. ``topk_p`` is returned as constant 1.0: topk=1 drafting
+    is greedy and the chain probabilities are unused downstream.
+    """
+    assert next_token_logits.ndim == 2
+    assert next_token_logits.stride(1) == 1
+    assert positions.ndim == 1
+    assert positions.is_contiguous()
+    assert positions.shape[0] == next_token_logits.shape[0]
+    assert positions.device == next_token_logits.device
+    write_draft_token = draft_tokens is not None
+    if write_draft_token:
+        assert draft_tokens.ndim == 2
+        assert draft_tokens.dtype == torch.long
+        assert draft_tokens.device == next_token_logits.device
+        assert draft_tokens.shape[0] == next_token_logits.shape[0]
+        assert draft_tokens.stride(1) == 1
+        assert 0 <= draft_token_column < draft_tokens.shape[1]
+
+    bs, vocab_size = next_token_logits.shape
+    topk_p = torch.empty((bs, 1), dtype=torch.float32, device=next_token_logits.device)
+    topk_index = torch.empty(
+        (bs, 1), dtype=torch.int64, device=next_token_logits.device
+    )
+    if bs == 0:
+        return topk_p, topk_index
+
+    block = _DRAFT_TOPK1_BLOCK
+    num_splits = triton.cdiv(vocab_size, block)
+    partial_vals = torch.empty(
+        (bs, num_splits), dtype=torch.float32, device=next_token_logits.device
+    )
+    partial_indices = torch.empty(
+        (bs, num_splits), dtype=torch.int32, device=next_token_logits.device
+    )
+
+    _draft_topk1_partial_argmax_kernel[(bs, num_splits)](
+        next_token_logits,
+        partial_vals,
+        partial_indices,
+        next_token_logits.stride(0),
+        vocab_size,
+        num_splits,
+        BLOCK=block,
+        num_warps=8,
+    )
+    # Dummy operand for the disabled draft-token slot: the pointer must be
+    # valid even though the kernel never dereferences it (gated off by
+    # WRITE_DRAFT_TOKEN).
+    _draft_topk1_finalize_kernel[(bs,)](
+        partial_vals,
+        partial_indices,
+        topk_p,
+        topk_index,
+        positions,
+        draft_tokens if write_draft_token else topk_index,
+        draft_tokens.stride(0) if write_draft_token else 0,
+        draft_token_column,
+        num_splits,
+        WRITE_DRAFT_TOKEN=write_draft_token,
+        BLOCK=triton.next_power_of_2(num_splits),
+        num_warps=1,
+    )
+    return topk_p, topk_index

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -6,6 +6,7 @@ from typing import List, Optional
 
 import torch
 
+from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import (
@@ -582,6 +583,27 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         if self.server_args.speculative_use_rejection_sampling:
             draft_probs_list: List[torch.Tensor] = [spec_info.draft_probs]
 
+        topk1_chain_fits = (
+            self.topk == 1
+            and topk_index.shape[0] <= self._topk1_parents_prealloc.shape[0]
+        )
+        # Materialize the chain directly only when the CUDA kernel can write
+        # every subsequent column. Other topk=1 paths retain the token list and
+        # assemble it with one final cat instead of launching a copy per step.
+        draft_tokens_topk1 = None
+        if (
+            topk1_chain_fits
+            and _is_cuda
+            and self.hot_token_id is None
+            and not self.server_args.speculative_use_rejection_sampling
+        ):
+            draft_tokens_topk1 = torch.empty(
+                (topk_index.shape[0], self.speculative_num_steps),
+                dtype=topk_index.dtype,
+                device=topk_index.device,
+            )
+            draft_tokens_topk1[:, :1].copy_(topk_index)
+
         # Forward multiple steps
         scores = None
         if self.index_share_for_mtp_iteration:
@@ -593,12 +615,15 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             ):
                 spec_info.dsa_topk_indices = None
         for i in range(self.speculative_num_steps):
-            input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
-                i, topk_p, topk_index, hidden_states, scores, self.topk
-            )
-            score_list.append(tree_info[0])
-            token_list.append(tree_info[1])
-            parents_list.append(tree_info[2])
+            if draft_tokens_topk1 is not None:
+                input_ids = topk_index.flatten()
+            else:
+                input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
+                    i, topk_p, topk_index, hidden_states, scores, self.topk
+                )
+                score_list.append(tree_info[0])
+                token_list.append(tree_info[1])
+                parents_list.append(tree_info[2])
 
             # We don't need to run the last forward. we get 1 token from draft prefill and (#spec steps - 1) tokens here
             if i == self.speculative_num_steps - 1:
@@ -641,11 +666,22 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                     forward_batch.sampling_info.temperatures,
                 )
                 draft_probs_list.append(probs)
+                forward_batch.positions.add_(1)
             elif self.topk == 1 and not _is_hip:
-                topk_index = torch.argmax(
-                    logits_output.next_token_logits, dim=-1, keepdim=True
-                )
-                topk_p = torch.ones_like(topk_index, dtype=torch.float32)
+                if _is_cuda:
+                    # The positions advance is fused into the kernel.
+                    topk_p, topk_index = draft_topk1_postprocess(
+                        logits_output.next_token_logits,
+                        forward_batch.positions,
+                        draft_tokens_topk1,
+                        i + 1,
+                    )
+                else:
+                    topk_index = torch.argmax(
+                        logits_output.next_token_logits, dim=-1, keepdim=True
+                    )
+                    topk_p = torch.ones_like(topk_index, dtype=torch.float32)
+                    forward_batch.positions.add_(1)
             else:
                 probs = renorm_draft_probs(
                     logits_output.next_token_logits,
@@ -653,6 +689,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                     self.server_args.speculative_use_rejection_sampling,
                 )
                 topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
+                forward_batch.positions.add_(1)
             maybe_detect_oob(
                 topk_index,
                 0,
@@ -662,42 +699,35 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             if self.hot_token_id is not None:
                 topk_index = self.hot_token_id[topk_index]
             hidden_states = logits_output.hidden_states
-            forward_batch.positions.add_(1)
 
         if self.index_share_for_mtp_iteration:
             spec_info.dsa_topk_indices = None
             forward_batch.reuse_dsa_topk_indices = False
 
-        # Organize the results
-        if (
-            self.topk == 1
-            and token_list[0].shape[0] <= self._topk1_parents_prealloc.shape[0]
-        ):
-            # Chain topology: draft_tokens = concat of per-step tokens; the
-            # full-length topk/sort/gather over score_list collapses to an
-            # identity. parent_list and top_scores_index are runtime-invariant
-            # constants pre-allocated on the worker. Oversized batches (rare,
-            # would silently truncate the slice) fall through to the slow path.
-            bs = token_list[0].shape[0]
-            draft_tokens = torch.cat(token_list, dim=1)
-            top_scores_index = self._topk1_score_indices_prealloc[:bs]
-            parent_list = self._topk1_parents_prealloc[:bs]
-            draft_probs = (
-                torch.stack(draft_probs_list, dim=1)
-                if self.server_args.speculative_use_rejection_sampling
-                else None
-            )
-            return parent_list, top_scores_index, draft_tokens, draft_probs
-
-        parent_list, top_scores_index, draft_tokens = organize_draft_results(
-            score_list, token_list, parents_list, self.speculative_num_draft_tokens
-        )
-
         draft_probs = (
             torch.stack(draft_probs_list, dim=1)
             if self.server_args.speculative_use_rejection_sampling
             else None
         )
+
+        # Organize the results
+        if draft_tokens_topk1 is not None:
+            bs = draft_tokens_topk1.shape[0]
+            top_scores_index = self._topk1_score_indices_prealloc[:bs]
+            parent_list = self._topk1_parents_prealloc[:bs]
+            return parent_list, top_scores_index, draft_tokens_topk1, draft_probs
+
+        if topk1_chain_fits:
+            bs = token_list[0].shape[0]
+            draft_tokens = torch.cat(token_list, dim=1)
+            top_scores_index = self._topk1_score_indices_prealloc[:bs]
+            parent_list = self._topk1_parents_prealloc[:bs]
+            return parent_list, top_scores_index, draft_tokens, draft_probs
+
+        parent_list, top_scores_index, draft_tokens = organize_draft_results(
+            score_list, token_list, parents_list, self.speculative_num_draft_tokens
+        )
+
         return parent_list, top_scores_index, draft_tokens, draft_probs
 
     def draft_extend(self):