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

test/registered/jit/benchmark/bench_spec_topk1.py

@@ -0,0 +1,165 @@
+"""Benchmark CUDA topk=1 speculative decoding helpers."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import (
+    DEFAULT_DEVICE,
+    get_benchmark_range,
+    run_benchmark,
+)
+from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+
+BATCH_SIZE_RANGE = get_benchmark_range(
+    full_range=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048],
+    ci_range=[1, 16, 256, 2048],
+)
+VOCAB_SIZES = {
+    "dsv4": 129280,
+    "glm5_2": 151552,
+}
+VOCAB_SIZE_RANGE = get_benchmark_range(
+    full_range=list(VOCAB_SIZES.values()),
+    ci_range=list(VOCAB_SIZES.values()),
+)
+NUM_STEPS = 3
+
+
+def make_logits(batch_size: int, vocab_size: int) -> torch.Tensor:
+    logits = torch.zeros(
+        (batch_size, vocab_size), dtype=torch.float32, device=DEFAULT_DEVICE
+    )
+    max_index = (
+        torch.arange(batch_size, dtype=torch.long, device=DEFAULT_DEVICE) * 9973 + 17
+    ) % vocab_size
+    logits.scatter_(1, max_index[:, None], 1000.0)
+    return logits
+
+
+def make_draft_case(batch_size: int, vocab_size: int):
+    logits = make_logits(batch_size, vocab_size)
+    positions = torch.zeros(batch_size, dtype=torch.long, device=DEFAULT_DEVICE)
+    return logits, positions
+
+
+def make_chain_case(batch_size: int, vocab_size: int):
+    seed_topk_index = torch.randint(
+        0, vocab_size, (batch_size, 1), dtype=torch.long, device=DEFAULT_DEVICE
+    )
+    logits = [make_logits(batch_size, vocab_size) for _ in range(NUM_STEPS - 1)]
+    positions = torch.zeros(batch_size, dtype=torch.long, device=DEFAULT_DEVICE)
+    return seed_topk_index, logits, positions
+
+
+def eager_draft_topk1_postprocess(logits: torch.Tensor, positions: torch.Tensor):
+    topk_index = torch.argmax(logits, dim=-1, keepdim=True)
+    topk_p = torch.ones_like(topk_index, dtype=torch.float32)
+    positions.add_(1)
+    return topk_p, topk_index
+
+
+def fused_draft_topk1_postprocess(logits: torch.Tensor, positions: torch.Tensor):
+    return draft_topk1_postprocess(logits, positions)
+
+
+def eager_chain_materialize(
+    seed_topk_index: torch.Tensor,
+    logits: list[torch.Tensor],
+    positions: torch.Tensor,
+):
+    token_list = [seed_topk_index]
+    for step_logits in logits:
+        _, topk_index = eager_draft_topk1_postprocess(step_logits, positions)
+        token_list.append(topk_index)
+    return torch.cat(token_list, dim=1)
+
+
+def fused_chain_materialize(
+    seed_topk_index: torch.Tensor,
+    logits: list[torch.Tensor],
+    positions: torch.Tensor,
+):
+    draft_tokens = torch.empty(
+        (seed_topk_index.shape[0], NUM_STEPS),
+        dtype=torch.long,
+        device=DEFAULT_DEVICE,
+    )
+    draft_tokens[:, :1].copy_(seed_topk_index)
+    for i, step_logits in enumerate(logits, start=1):
+        draft_topk1_postprocess(
+            step_logits,
+            positions,
+            draft_tokens,
+            draft_token_column=i,
+        )
+    return draft_tokens
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "vocab_size"],
+        x_vals=[(bs, vocab) for bs in BATCH_SIZE_RANGE for vocab in VOCAB_SIZE_RANGE],
+        line_arg="provider",
+        line_vals=["fused", "eager"],
+        line_names=["Fused Triton", "Eager torch"],
+        styles=[("blue", "-"), ("orange", "--")],
+        ylabel="us",
+        plot_name="spec-topk1-draft-postprocess",
+        args={},
+    )
+)
+def benchmark_draft_postprocess(
+    batch_size: int, vocab_size: int, provider: str
+) -> tuple[float, float, float]:
+    logits, positions = make_draft_case(batch_size, vocab_size)
+    if provider == "fused":
+        fn = lambda: fused_draft_topk1_postprocess(logits, positions)
+    elif provider == "eager":
+        fn = lambda: eager_draft_topk1_postprocess(logits, positions)
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+    fn()
+    torch.cuda.synchronize()
+    return run_benchmark(fn)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "vocab_size"],
+        x_vals=[(bs, vocab) for bs in BATCH_SIZE_RANGE for vocab in VOCAB_SIZE_RANGE],
+        line_arg="provider",
+        line_vals=["fused", "eager"],
+        line_names=["Fused Triton", "Eager argmax + cat"],
+        styles=[("blue", "-"), ("orange", "--")],
+        ylabel="us",
+        plot_name="spec-topk1-chain-materialize",
+        args={},
+    )
+)
+def benchmark_chain_materialize(
+    batch_size: int, vocab_size: int, provider: str
+) -> tuple[float, float, float]:
+    seed_topk_index, logits, positions = make_chain_case(batch_size, vocab_size)
+    if provider == "fused":
+        fn = lambda: fused_chain_materialize(seed_topk_index, logits, positions)
+    elif provider == "eager":
+        fn = lambda: eager_chain_materialize(seed_topk_index, logits, positions)
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+    fn()
+    torch.cuda.synchronize()
+    return run_benchmark(fn)
+
+
+if __name__ == "__main__":
+    benchmark_draft_postprocess.run(print_data=True)
+    benchmark_chain_materialize.run(print_data=True)

test/registered/kernels/test_kernels_namespace.py

@@ -69,6 +69,7 @@ EXPECTED_OPS = {
     "memory.alloc_extend_kernel": {"triton"},
     "attention.decode_attention_fwd": {"triton"},
     "kvcache.create_flashinfer_kv_indices_triton": {"triton"},
+    "speculative.draft_topk1_postprocess": {"triton"},
     "speculative.gather_spec_extras": {"triton"},
 }
 

test/registered/kernels/test_spec_topk1.py

@@ -0,0 +1,163 @@
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess
+from sglang.test.test_utils import CustomTestCase
+
+
+def _make_logits_with_unique_argmax(
+    batch_size: int,
+    vocab_size: int,
+    *,
+    dtype: torch.dtype,
+    device: torch.device,
+    seed: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    g = torch.Generator(device=device).manual_seed(seed)
+    logits = torch.randn(
+        (batch_size, vocab_size), dtype=dtype, device=device, generator=g
+    )
+    expected_index = (
+        torch.arange(batch_size, dtype=torch.long, device=device) * 9973 + 17
+    ) % vocab_size
+    logits.scatter_(1, expected_index[:, None], 1000.0)
+    return logits, expected_index[:, None]
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
+class TestSpecTopk1Triton(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.device = torch.device("cuda")
+
+    def test_draft_topk1_postprocess_matches_argmax_and_position_add(self):
+        configs = [
+            (1, 127, torch.float32),
+            (4, 8192, torch.float16),
+            (7, 8193, torch.bfloat16),
+            (3, 50000, torch.float32),
+        ]
+        for batch_size, vocab_size, dtype in configs:
+            with self.subTest(
+                batch_size=batch_size, vocab_size=vocab_size, dtype=dtype
+            ):
+                logits, expected_index = _make_logits_with_unique_argmax(
+                    batch_size,
+                    vocab_size,
+                    dtype=dtype,
+                    device=self.device,
+                    seed=vocab_size,
+                )
+                positions = torch.arange(
+                    batch_size, dtype=torch.long, device=self.device
+                )
+                expected_positions = positions + 1
+
+                topk_p, topk_index = draft_topk1_postprocess(logits, positions)
+
+                torch.testing.assert_close(topk_index, expected_index, rtol=0, atol=0)
+                torch.testing.assert_close(
+                    topk_p,
+                    torch.ones(
+                        (batch_size, 1), dtype=torch.float32, device=self.device
+                    ),
+                    rtol=0,
+                    atol=0,
+                )
+                torch.testing.assert_close(
+                    positions, expected_positions, rtol=0, atol=0
+                )
+
+    def test_draft_topk1_postprocess_can_write_draft_token_column(self):
+        batch_size = 17
+        # Multi-split vocab so the fused write composes with the split reduction.
+        vocab_size = 50000
+        logits, expected_index = _make_logits_with_unique_argmax(
+            batch_size,
+            vocab_size,
+            dtype=torch.float32,
+            device=self.device,
+            seed=0,
+        )
+        positions = torch.zeros(batch_size, dtype=torch.long, device=self.device)
+        backing = torch.full((batch_size, 5), -1, dtype=torch.long, device=self.device)
+        draft_tokens = backing[:, 1:4]
+
+        topk_p, topk_index = draft_topk1_postprocess(
+            logits, positions, draft_tokens, draft_token_column=2
+        )
+
+        torch.testing.assert_close(topk_index, expected_index, rtol=0, atol=0)
+        torch.testing.assert_close(topk_p, torch.ones_like(topk_p), rtol=0, atol=0)
+        # Exactly one backing column is written; both neighbors stay untouched.
+        expected_backing = torch.full_like(backing, -1)
+        expected_backing[:, 3] = expected_index[:, 0]
+        torch.testing.assert_close(backing, expected_backing, rtol=0, atol=0)
+        torch.testing.assert_close(
+            positions, torch.ones_like(positions), rtol=0, atol=0
+        )
+
+    def test_row_strided_logits_view_matches_argmax(self):
+        batch_size = 5
+        vocab_size = 8193
+        # Poison the padding columns: if the kernel used the dense vocab width
+        # as the row stride it would read them and pick the wrong index.
+        backing = torch.full(
+            (batch_size, vocab_size + 64),
+            2000.0,
+            dtype=torch.float32,
+            device=self.device,
+        )
+        logits, expected_index = _make_logits_with_unique_argmax(
+            batch_size,
+            vocab_size,
+            dtype=torch.float32,
+            device=self.device,
+            seed=1,
+        )
+        backing[:, :vocab_size] = logits
+        strided_logits = backing[:, :vocab_size]
+        self.assertFalse(strided_logits.is_contiguous())
+        positions = torch.zeros(batch_size, dtype=torch.long, device=self.device)
+
+        topk_p, topk_index = draft_topk1_postprocess(strided_logits, positions)
+
+        torch.testing.assert_close(topk_index, expected_index, rtol=0, atol=0)
+        torch.testing.assert_close(topk_p, torch.ones_like(topk_p), rtol=0, atol=0)
+        torch.testing.assert_close(
+            positions, torch.ones_like(positions), rtol=0, atol=0
+        )
+
+    def test_empty_batch(self):
+        logits = torch.empty((0, 1024), dtype=torch.float32, device=self.device)
+        positions = torch.empty((0,), dtype=torch.long, device=self.device)
+        draft_tokens = torch.empty((0, 3), dtype=torch.long, device=self.device)
+
+        topk_p, topk_index = draft_topk1_postprocess(
+            logits, positions, draft_tokens, draft_token_column=1
+        )
+
+        self.assertEqual(topk_p.shape, (0, 1))
+        self.assertEqual(topk_index.shape, (0, 1))
+        self.assertEqual(draft_tokens.numel(), 0)
+
+    def test_non_contiguous_inputs_raise(self):
+        logits = torch.empty((16, 4), dtype=torch.float32, device=self.device).t()
+        positions = torch.arange(8, dtype=torch.long, device=self.device)[::2]
+
+        with self.assertRaises(AssertionError):
+            draft_topk1_postprocess(
+                logits, torch.empty(4, dtype=torch.long, device=self.device)
+            )
+        with self.assertRaises(AssertionError):
+            draft_topk1_postprocess(torch.empty((4, 16), device=self.device), positions)
+
+
+if __name__ == "__main__":
+    unittest.main()