DeepSeek-V4 Online Compress support MTP (#26471)

2026-06-15 19:56:07 -07:00
parent b5bcd76a41
commit 063ab89ac1
12 changed files with 1276 additions and 49 deletions
@@ -0,0 +1,166 @@
+"""Benchmark online c128 MTP write-prefix kernel."""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass
+
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import (
+    DEFAULT_DEVICE,
+    get_benchmark_range,
+    run_benchmark,
+    run_benchmark_no_cudagraph,
+)
+from sglang.jit_kernel.dsv4.online_c128_mtp import _jit_online_c128_mtp_module
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large")
+
+HEAD_DIM = 512
+STATE_DIM = HEAD_DIM * 3
+SWA_PAGE_SIZE = 128
+
+BATCH_SIZE_RANGE = get_benchmark_range(
+    full_range=[1, 2, 4, 8, 16, 32, 64, 128, 256],
+    ci_range=[8, 64],
+)
+NUM_VERIFY_TOKENS_RANGE = get_benchmark_range(
+    full_range=[1, 4, 8],
+    ci_range=[8],
+)
+BENCHMARK_CONFIGS = list(itertools.product(BATCH_SIZE_RANGE, NUM_VERIFY_TOKENS_RANGE))
+
+
+@dataclass
+class BenchmarkCase:
+    kv_score_input: torch.Tensor
+    seq_lens: torch.Tensor
+    req_pool_indices: torch.Tensor
+    req_to_token: torch.Tensor
+    full_to_swa: torch.Tensor
+    ape: torch.Tensor
+    state: torch.Tensor
+    layer_bs: int
+    num_verify_tokens: int
+    state_slot_stride: int
+
+
+def round_up_div(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+
+def make_seq_lens(batch_size: int, num_verify_tokens: int) -> torch.Tensor:
+    # Cover chunk positions around the interesting boundaries. This exercises
+    # both the has-partial path and the final_seq % 128 == 0 skip-write path.
+    offsets = torch.tensor([0, 1, 2, 63, 120, 126, 127], dtype=torch.int64)
+    seq_offsets = offsets[torch.arange(batch_size, dtype=torch.int64) % offsets.numel()]
+    base = 8 * SWA_PAGE_SIZE
+    seq_lens = base + seq_offsets
+    assert int(seq_lens.max()) + num_verify_tokens < base + 2 * SWA_PAGE_SIZE
+    return seq_lens.to(device=DEFAULT_DEVICE)
+
+
+def make_req_to_token(
+    batch_size: int, max_seq_len: int, num_chunks: int
+) -> torch.Tensor:
+    chunk_ids = torch.arange(max_seq_len, dtype=torch.int32) // SWA_PAGE_SIZE
+    req_offsets = torch.arange(batch_size, dtype=torch.int32).unsqueeze(1) * num_chunks
+    req_to_token = req_offsets + chunk_ids.unsqueeze(0)
+    return req_to_token.contiguous().to(device=DEFAULT_DEVICE)
+
+
+def make_case(batch_size: int, num_verify_tokens: int) -> BenchmarkCase:
+    seq_lens = make_seq_lens(batch_size, num_verify_tokens)
+    req_pool_indices = torch.arange(
+        batch_size, dtype=torch.int64, device=DEFAULT_DEVICE
+    )
+
+    max_seq_len = int(seq_lens.max().item()) + num_verify_tokens + SWA_PAGE_SIZE
+    num_chunks = round_up_div(max_seq_len, SWA_PAGE_SIZE)
+    req_to_token = make_req_to_token(batch_size, max_seq_len, num_chunks)
+
+    num_full_locs = batch_size * num_chunks
+    full_to_swa = (
+        torch.arange(num_full_locs, dtype=torch.int64, device=DEFAULT_DEVICE)
+        * SWA_PAGE_SIZE
+    )
+
+    state_slot_stride = num_full_locs
+    state = torch.empty(
+        (state_slot_stride * (1 + num_verify_tokens), STATE_DIM),
+        dtype=torch.float32,
+        device=DEFAULT_DEVICE,
+    )
+    state.normal_(mean=0.0, std=0.01)
+
+    kv_score_input = torch.randn(
+        batch_size * num_verify_tokens,
+        HEAD_DIM * 2,
+        dtype=torch.float32,
+        device=DEFAULT_DEVICE,
+    )
+    ape = torch.randn(128, HEAD_DIM, dtype=torch.float32, device=DEFAULT_DEVICE)
+
+    return BenchmarkCase(
+        kv_score_input=kv_score_input,
+        seq_lens=seq_lens,
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        full_to_swa=full_to_swa,
+        ape=ape,
+        state=state,
+        layer_bs=batch_size,
+        num_verify_tokens=num_verify_tokens,
+        state_slot_stride=state_slot_stride,
+    )
+
+
+def call_write_prefix(module, case: BenchmarkCase) -> None:
+    module.write_prefix_states(
+        case.kv_score_input,
+        case.seq_lens,
+        case.req_pool_indices,
+        case.req_to_token,
+        case.full_to_swa,
+        case.ape,
+        case.state,
+        case.layer_bs,
+        SWA_PAGE_SIZE,
+        case.num_verify_tokens,
+        case.state_slot_stride,
+    )
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "num_verify_tokens"],
+        x_vals=BENCHMARK_CONFIGS,
+        line_arg="launch_mode",
+        line_vals=["cuda_graph", "eager"],
+        line_names=["CUDA graph", "Eager launch"],
+        styles=[("blue", "-"), ("orange", "--")],
+        ylabel="us",
+        plot_name="online-c128-mtp-write-prefix-performance",
+        args={},
+    )
+)
+def benchmark(
+    batch_size: int, num_verify_tokens: int, launch_mode: str
+) -> tuple[float, float, float]:
+    module = _jit_online_c128_mtp_module(HEAD_DIM)
+    case = make_case(batch_size, num_verify_tokens)
+    fn = lambda: call_write_prefix(module, case)
+
+    if launch_mode == "cuda_graph":
+        return run_benchmark(fn)
+    if launch_mode == "eager":
+        return run_benchmark_no_cudagraph(fn)
+    raise ValueError(f"Unknown launch_mode: {launch_mode}")
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)