[jit_kernel] Add JIT ngram_utils kernel

python/sglang/jit_kernel/benchmark/bench_ngram_utils.py

@@ -0,0 +1,172 @@
+"""
+Benchmark: reconstruct_indices_from_tree_mask JIT vs AOT (sgl_kernel)
+
+Measures throughput (µs) across typical batch sizes and tree sizes.
+
+Run:
+    python python/sglang/jit_kernel/benchmark/bench_ngram_utils.py
+"""
+
+import itertools
+
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
+from sglang.jit_kernel.ngram_utils import (
+    reconstruct_indices_from_tree_mask as reconstruct_jit,
+)
+
+try:
+    from sgl_kernel import reconstruct_indices_from_tree_mask as reconstruct_aot
+
+    AOT_AVAILABLE = True
+except ImportError:
+    reconstruct_aot = None
+    AOT_AVAILABLE = False
+
+DEVICE = "cuda"
+
+LINE_VALS = ["jit", "aot"] if AOT_AVAILABLE else ["jit"]
+LINE_NAMES = ["JIT (new)", "AOT sgl_kernel"] if AOT_AVAILABLE else ["JIT (new)"]
+
+# ---------------------------------------------------------------------------
+# Benchmark configuration
+# ---------------------------------------------------------------------------
+
+BATCH_SIZE_RANGE = get_benchmark_range(
+    full_range=[1, 4, 8, 16],
+    ci_range=[4],
+)
+
+DRAFT_TOKEN_RANGE = get_benchmark_range(
+    full_range=[8, 16, 32, 64],
+    ci_range=[16],
+)
+
+
+# ---------------------------------------------------------------------------
+# Input helpers
+# ---------------------------------------------------------------------------
+
+
+def make_inputs(bs, draft_token_num):
+    tree_mask = torch.zeros(
+        bs * draft_token_num * draft_token_num, dtype=torch.bool, device=DEVICE
+    )
+    # Linear-chain tree mask
+    base = draft_token_num * draft_token_num
+    for b in range(bs):
+        for i in range(draft_token_num):
+            for j in range(i):
+                tree_mask[b * base + i * draft_token_num + j] = True
+
+    verified_seq_len = torch.full((bs,), 128, dtype=torch.int64, device=DEVICE)
+    positions = torch.zeros(bs * draft_token_num, dtype=torch.int64, device=DEVICE)
+    retrive_index = torch.zeros(bs, draft_token_num, dtype=torch.int64, device=DEVICE)
+    retrive_next_token = torch.full(
+        (bs, draft_token_num), -1, dtype=torch.int64, device=DEVICE
+    )
+    retrive_next_sibling = torch.full(
+        (bs, draft_token_num), -1, dtype=torch.int64, device=DEVICE
+    )
+
+    return dict(
+        tree_mask=tree_mask,
+        verified_seq_len=verified_seq_len,
+        positions=positions,
+        retrive_index=retrive_index,
+        retrive_next_token=retrive_next_token,
+        retrive_next_sibling=retrive_next_sibling,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["bs", "draft_token_num"],
+        x_vals=list(itertools.product(BATCH_SIZE_RANGE, DRAFT_TOKEN_RANGE)),
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=[("blue", "--"), ("orange", "-")][: len(LINE_VALS)],
+        ylabel="us",
+        plot_name="reconstruct-indices-from-tree-mask-performance",
+        args={},
+    )
+)
+def bench_reconstruct(bs: int, draft_token_num: int, provider: str):
+    inputs = make_inputs(bs, draft_token_num)
+
+    mutated_keys = {
+        "positions",
+        "retrive_index",
+        "retrive_next_token",
+        "retrive_next_sibling",
+    }
+    backups = {k: inputs[k].clone() for k in mutated_keys}
+
+    if provider == "jit":
+
+        def fn():
+            for k in mutated_keys:
+                inputs[k].copy_(backups[k])
+            reconstruct_jit(**inputs, batch_size=bs, draft_token_num=draft_token_num)
+
+    elif provider == "aot":
+
+        def fn():
+            for k in mutated_keys:
+                inputs[k].copy_(backups[k])
+            reconstruct_aot(**inputs, batch_size=bs, draft_token_num=draft_token_num)
+
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    return run_benchmark(fn)
+
+
+# ---------------------------------------------------------------------------
+# Quick correctness diff
+# ---------------------------------------------------------------------------
+
+
+def calculate_diff():
+    if not AOT_AVAILABLE:
+        print("sgl_kernel not available — skipping AOT diff check")
+        return
+
+    print("Correctness diff — reconstruct_indices_from_tree_mask (JIT vs AOT):")
+    for bs, draft_token_num in [(1, 8), (2, 16), (4, 32)]:
+        inp_jit = make_inputs(bs, draft_token_num)
+        inp_aot = {
+            k: v.clone() if isinstance(v, torch.Tensor) else v
+            for k, v in inp_jit.items()
+        }
+
+        reconstruct_jit(**inp_jit, batch_size=bs, draft_token_num=draft_token_num)
+        reconstruct_aot(**inp_aot, batch_size=bs, draft_token_num=draft_token_num)
+
+        match_pos = torch.equal(inp_jit["positions"], inp_aot["positions"])
+        match_idx = torch.equal(inp_jit["retrive_index"], inp_aot["retrive_index"])
+        match_next = torch.equal(
+            inp_jit["retrive_next_token"], inp_aot["retrive_next_token"]
+        )
+        match_sib = torch.equal(
+            inp_jit["retrive_next_sibling"], inp_aot["retrive_next_sibling"]
+        )
+        status = (
+            "OK" if all([match_pos, match_idx, match_next, match_sib]) else "MISMATCH"
+        )
+        print(f"  bs={bs:2d} draft_token_num={draft_token_num:2d}  [{status}]")
+
+
+if __name__ == "__main__":
+    calculate_diff()
+    print()
+    bench_reconstruct.run(print_data=True)

python/sglang/jit_kernel/csrc/speculative/ngram_utils.cuh

@@ -0,0 +1,184 @@
+/*
+ * Copyright (c) 2025 by SGLang team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// Adapted from
+// https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/speculative/ngram_utils.cu
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+// tree_mask: [bs * draft_token_num * draft_token_num]
+// verified_seq_len: [bs]
+// positions: [bs * draft_token_num]
+// retrive_index: [bs, draft_token_num]
+// retrive_next_token: [bs, draft_token_num]
+// retrive_next_sibling: [bs, draft_token_num]
+__global__ void reconstructIndicesFromTreeMask(
+    uint8_t* tree_mask,
+    int64_t* verified_seq_len,
+    int64_t* positions,
+    int64_t* retrive_index,
+    int64_t* retrive_next_token,
+    int64_t* retrive_next_sibling,
+    int batch_size,
+    int draft_token_num) {
+  int bid = blockIdx.x;
+  int tid = threadIdx.x;
+
+  if (bid >= batch_size || tid >= draft_token_num) {
+    return;
+  }
+  int base_offset = draft_token_num * draft_token_num;
+  // token_idx: [bid * draft_token_num, (bid + 1) * draft_token_num)
+  int token_idx = bid * draft_token_num;
+  // tree_mask_idx: [bid * base_offset, (bid + 1) * base_offset)
+  int tree_mask_offset = bid * base_offset;
+
+  int depth = 0;
+  int parent_idx = -1;
+
+  for (int i = tid - 1, start_idx = tree_mask_offset + tid * draft_token_num; i >= 0; i--) {
+    if (tree_mask[start_idx + i]) {
+      depth++;
+      if (parent_idx == -1) {
+        parent_idx = i;
+      }
+    }
+  }
+  retrive_index[token_idx + tid] = token_idx + tid;
+  positions[token_idx + tid] = depth + verified_seq_len[bid];
+
+  int next_token_idx = -1;
+  for (int i = tid + 1; i < draft_token_num; i++) {
+    if (tree_mask[tree_mask_offset + i * draft_token_num + tid]) {
+      next_token_idx = i;
+      break;
+    }
+  }
+  retrive_next_token[token_idx + tid] = next_token_idx;
+
+  int next_sibling_idx = -1;
+  if (parent_idx != -1) {
+    for (int i = tid + 1; i < draft_token_num; i++) {
+      int start_idx = tree_mask_offset + i * draft_token_num + parent_idx;
+      if (tree_mask[start_idx]) {
+        bool is_sibling = true;
+        int end_idx = tree_mask_offset + i * draft_token_num + i;
+        for (int j = start_idx + 1; j < end_idx; ++j) {
+          if (tree_mask[j]) {
+            is_sibling = false;
+            break;
+          }
+        }
+        if (is_sibling) {
+          next_sibling_idx = i;
+          break;
+        }
+      }
+    }
+  }
+  retrive_next_sibling[token_idx + tid] = next_sibling_idx;
+}
+
+namespace {
+
+// ---------------------------------------------------------------------------
+// tvm-ffi entry point
+// ---------------------------------------------------------------------------
+
+// tree_mask:            [bs * draft_token_num * draft_token_num] bool
+// verified_seq_len:     [bs] int64
+// positions:            [bs * draft_token_num] int64, mutable
+// retrive_index:        [bs, draft_token_num] int64, mutable
+// retrive_next_token:   [bs, draft_token_num] int64, mutable
+// retrive_next_sibling: [bs, draft_token_num] int64, mutable
+// batch_size, draft_token_num: scalars
+void reconstruct_indices_from_tree_mask(
+    tvm::ffi::TensorView tree_mask,
+    tvm::ffi::TensorView verified_seq_len,
+    tvm::ffi::TensorView positions,
+    tvm::ffi::TensorView retrive_index,
+    tvm::ffi::TensorView retrive_next_token,
+    tvm::ffi::TensorView retrive_next_sibling,
+    int64_t batch_size,
+    int64_t draft_token_num) {
+  using namespace host;
+
+  RuntimeCheck(tree_mask.device().device_type == kDLCUDA, "tree_mask must be a CUDA tensor");
+  RuntimeCheck(tree_mask.ndim() == 1, "tree_mask must be 1D: [bs * draft_token_num * draft_token_num]");
+  RuntimeCheck(tree_mask.is_contiguous(), "tree_mask must be contiguous");
+  RuntimeCheck(host::dtype_bytes(tree_mask.dtype()) == 1, "tree_mask element size must be 1 byte (bool or uint8)");
+  RuntimeCheck(
+      tree_mask.size(0) == batch_size * draft_token_num * draft_token_num,
+      "tree_mask size must equal batch_size * draft_token_num * draft_token_num");
+
+  RuntimeCheck(verified_seq_len.ndim() == 1, "verified_seq_len must be 1D: [bs]");
+  RuntimeCheck(verified_seq_len.is_contiguous(), "verified_seq_len must be contiguous");
+  RuntimeCheck(
+      verified_seq_len.dtype().code == kDLInt && verified_seq_len.dtype().bits == 64, "verified_seq_len must be int64");
+  RuntimeCheck(verified_seq_len.size(0) == batch_size, "verified_seq_len size must equal batch_size");
+
+  RuntimeCheck(positions.ndim() == 1, "positions must be 1D: [bs * draft_token_num]");
+  RuntimeCheck(positions.is_contiguous(), "positions must be contiguous");
+  RuntimeCheck(positions.dtype().code == kDLInt && positions.dtype().bits == 64, "positions must be int64");
+  RuntimeCheck(
+      positions.size(0) == batch_size * draft_token_num, "positions size must equal batch_size * draft_token_num");
+
+  RuntimeCheck(retrive_index.ndim() == 2, "retrive_index must be 2D: [bs, draft_token_num]");
+  RuntimeCheck(retrive_index.is_contiguous(), "retrive_index must be contiguous");
+  RuntimeCheck(retrive_index.dtype().code == kDLInt && retrive_index.dtype().bits == 64, "retrive_index must be int64");
+  RuntimeCheck(
+      retrive_index.size(0) == batch_size && retrive_index.size(1) == draft_token_num,
+      "retrive_index shape must be [batch_size, draft_token_num]");
+
+  RuntimeCheck(retrive_next_token.ndim() == 2, "retrive_next_token must be 2D: [bs, draft_token_num]");
+  RuntimeCheck(retrive_next_token.is_contiguous(), "retrive_next_token must be contiguous");
+  RuntimeCheck(
+      retrive_next_token.dtype().code == kDLInt && retrive_next_token.dtype().bits == 64,
+      "retrive_next_token must be int64");
+  RuntimeCheck(
+      retrive_next_token.size(0) == batch_size && retrive_next_token.size(1) == draft_token_num,
+      "retrive_next_token shape must be [batch_size, draft_token_num]");
+
+  RuntimeCheck(retrive_next_sibling.ndim() == 2, "retrive_next_sibling must be 2D: [bs, draft_token_num]");
+  RuntimeCheck(retrive_next_sibling.is_contiguous(), "retrive_next_sibling must be contiguous");
+  RuntimeCheck(
+      retrive_next_sibling.dtype().code == kDLInt && retrive_next_sibling.dtype().bits == 64,
+      "retrive_next_sibling must be int64");
+  RuntimeCheck(
+      retrive_next_sibling.size(0) == batch_size && retrive_next_sibling.size(1) == draft_token_num,
+      "retrive_next_sibling shape must be [batch_size, draft_token_num]");
+
+  cudaStream_t stream = LaunchKernel::resolve_device(tree_mask.device());
+  dim3 grid(static_cast<unsigned>(batch_size));
+  dim3 block(static_cast<unsigned>(draft_token_num));
+
+  LaunchKernel(grid, block, stream)(
+      reconstructIndicesFromTreeMask,
+      static_cast<uint8_t*>(tree_mask.data_ptr()),
+      static_cast<int64_t*>(verified_seq_len.data_ptr()),
+      static_cast<int64_t*>(positions.data_ptr()),
+      static_cast<int64_t*>(retrive_index.data_ptr()),
+      static_cast<int64_t*>(retrive_next_token.data_ptr()),
+      static_cast<int64_t*>(retrive_next_sibling.data_ptr()),
+      static_cast<int>(batch_size),
+      static_cast<int>(draft_token_num));
+}
+
+}  // namespace