diff --git a/python/sglang/jit_kernel/benchmark/bench_ngram_utils.py b/python/sglang/jit_kernel/benchmark/bench_ngram_utils.py
new file mode 100644
index 000000000000..8a56cbbffccd
--- /dev/null
+++ b/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)
diff --git a/python/sglang/jit_kernel/csrc/speculative/ngram_utils.cuh b/python/sglang/jit_kernel/csrc/speculative/ngram_utils.cuh
new file mode 100644
index 000000000000..edb59325c37c
--- /dev/null
+++ b/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
diff --git a/python/sglang/jit_kernel/ngram_utils.py b/python/sglang/jit_kernel/ngram_utils.py
new file mode 100644
index 000000000000..685217b95348
--- /dev/null
+++ b/python/sglang/jit_kernel/ngram_utils.py
@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_ngram_utils_module() -> Module:
+    return load_jit(
+        "ngram_utils",
+        cuda_files=["speculative/ngram_utils.cuh"],
+        cuda_wrappers=[
+            (
+                "reconstruct_indices_from_tree_mask",
+                "reconstruct_indices_from_tree_mask",
+            ),
+        ],
+    )
+
+
+@register_custom_op(
+    op_name="reconstruct_indices_from_tree_mask_out",
+    mutates_args=[
+        "positions",
+        "retrive_index",
+        "retrive_next_token",
+        "retrive_next_sibling",
+    ],
+)
+def reconstruct_indices_from_tree_mask(
+    tree_mask: torch.Tensor,
+    verified_seq_len: torch.Tensor,
+    positions: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    batch_size: int,
+    draft_token_num: int,
+) -> None:
+    """
+    Reconstruct tree indices from a flat boolean tree mask.
+
+    Args:
+        tree_mask:            [bs * draft_token_num * draft_token_num] bool —
+                              attention mask encoding the tree structure
+        verified_seq_len:     [bs] int64 — verified sequence lengths
+        positions:            [bs * draft_token_num] int64 — filled with token positions
+        retrive_index:        [bs, draft_token_num] int64 — filled with retrieval indices
+        retrive_next_token:   [bs, draft_token_num] int64 — filled with next token links
+        retrive_next_sibling: [bs, draft_token_num] int64 — filled with sibling links
+        batch_size:           number of sequences in the batch
+        draft_token_num:      number of draft tokens per sequence
+    """
+    module = _jit_ngram_utils_module()
+    module.reconstruct_indices_from_tree_mask(
+        tree_mask,
+        verified_seq_len,
+        positions,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        batch_size,
+        draft_token_num,
+    )
diff --git a/python/sglang/jit_kernel/tests/test_ngram_utils.py b/python/sglang/jit_kernel/tests/test_ngram_utils.py
new file mode 100644
index 000000000000..80b5e2580d01
--- /dev/null
+++ b/python/sglang/jit_kernel/tests/test_ngram_utils.py
@@ -0,0 +1,296 @@
+"""
+Tests for the JIT reconstruct_indices_from_tree_mask kernel.
+
+Correctness is validated by:
+1. Smoke tests across batch sizes and tree sizes.
+2. Known-answer tests: linear chain and branching tree with hand-crafted
+   tree masks and expected positions / next_token / next_sibling values.
+3. JIT vs AOT cross-validation (when sgl_kernel is available).
+"""
+
+import pytest
+import torch
+
+DEVICE = "cuda"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def make_inputs(bs, draft_token_num, device=DEVICE):
+    """Allocate output tensors; tree_mask and verified_seq_len filled by caller."""
+    tree_mask = torch.zeros(
+        bs * draft_token_num * draft_token_num, dtype=torch.bool, device=device
+    )
+    verified_seq_len = torch.zeros(bs, 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,
+    )
+
+
+def build_linear_chain_mask(bs, draft_token_num, device=DEVICE):
+    """
+    Build a tree_mask for a linear chain: 0 → 1 → 2 → ... → draft_token_num-1.
+
+    tree_mask[b, i, j] = True if token j is an ancestor of token i (j < i).
+    For a linear chain, token i has ancestors {0, 1, ..., i-1}.
+    """
+    tree_mask = torch.zeros(
+        bs * draft_token_num * draft_token_num, dtype=torch.bool, device=device
+    )
+    base = draft_token_num * draft_token_num
+    for b in range(bs):
+        for i in range(draft_token_num):
+            for j in range(i):  # all predecessors are ancestors in a chain
+                tree_mask[b * base + i * draft_token_num + j] = True
+    return tree_mask
+
+
+# ---------------------------------------------------------------------------
+# Smoke tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("bs", [1, 2, 4])
+@pytest.mark.parametrize("draft_token_num", [4, 8, 16])
+def test_smoke(bs, draft_token_num):
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_linear_chain_mask(bs, draft_token_num)
+    inputs["verified_seq_len"].fill_(10)
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    # retrive_index[b, i] must equal b * draft_token_num + i
+    for b in range(bs):
+        for i in range(draft_token_num):
+            assert inputs["retrive_index"][b, i].item() == b * draft_token_num + i
+
+
+# ---------------------------------------------------------------------------
+# Known-answer: linear chain
+# ---------------------------------------------------------------------------
+
+
+def test_linear_chain_positions():
+    """positions[b*N + i] = verified_seq_len[b] + depth(i)  (depth = i for chain)."""
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    bs, draft_token_num, seq_len = 2, 4, 5
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_linear_chain_mask(bs, draft_token_num)
+    inputs["verified_seq_len"].fill_(seq_len)
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    for b in range(bs):
+        for i in range(draft_token_num):
+            expected = seq_len + i  # depth of token i in a chain = i
+            assert (
+                inputs["positions"][b * draft_token_num + i].item() == expected
+            ), f"positions[{b},{i}] = {inputs['positions'][b*draft_token_num+i].item()}, expected {expected}"
+
+
+def test_linear_chain_next_token():
+    """In a linear chain, retrive_next_token[b, i] = i+1 (last = -1)."""
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    bs, draft_token_num = 1, 5
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_linear_chain_mask(bs, draft_token_num)
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    for i in range(draft_token_num - 1):
+        assert inputs["retrive_next_token"][0, i].item() == i + 1
+    assert inputs["retrive_next_token"][0, draft_token_num - 1].item() == -1
+
+
+def test_linear_chain_no_siblings():
+    """In a linear chain, there are no siblings."""
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    bs, draft_token_num = 1, 5
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_linear_chain_mask(bs, draft_token_num)
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    assert (inputs["retrive_next_sibling"] == -1).all()
+
+
+# ---------------------------------------------------------------------------
+# Known-answer: branching tree
+#
+# Tree structure (draft_token_num=5, bs=1):
+#
+#         0  (root, depth 0)
+#        / \
+#       1   2  (depth 1, siblings: next_sibling[1]=2)
+#      / \
+#     3   4  (depth 2, siblings: next_sibling[3]=4, children of 1)
+#
+# Ancestor sets:
+#   token 0: {}
+#   token 1: {0}
+#   token 2: {0}
+#   token 3: {0, 1}
+#   token 4: {0, 1}
+#
+# Expected outputs (seq_len=3):
+#   positions:            [3, 4, 4, 5, 5]
+#   retrive_next_token:   [1, 3, -1, -1, -1]
+#   retrive_next_sibling: [-1, 2, -1, 4, -1]
+# ---------------------------------------------------------------------------
+
+
+def build_branching_tree_mask(device=DEVICE):
+    """
+    Build tree_mask for the 5-token branching tree described above.
+    tree_mask[i*N + j] = True if token j is an ancestor of token i.
+    """
+    N = 5
+    mask = torch.zeros(N * N, dtype=torch.bool, device=device)
+    # token 1: ancestor 0
+    mask[1 * N + 0] = True
+    # token 2: ancestor 0
+    mask[2 * N + 0] = True
+    # token 3: ancestors 0, 1
+    mask[3 * N + 0] = True
+    mask[3 * N + 1] = True
+    # token 4: ancestors 0, 1
+    mask[4 * N + 0] = True
+    mask[4 * N + 1] = True
+    return mask
+
+
+def test_branching_tree_positions():
+    """positions reflect actual tree depth, not sequential order."""
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    bs, draft_token_num, seq_len = 1, 5, 3
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_branching_tree_mask()
+    inputs["verified_seq_len"].fill_(seq_len)
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    expected_positions = [3, 4, 4, 5, 5]
+    for i, exp in enumerate(expected_positions):
+        got = inputs["positions"][i].item()
+        assert got == exp, f"positions[{i}]={got}, expected {exp}"
+
+
+def test_branching_tree_next_token():
+    """retrive_next_token points to the first child of each node."""
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    bs, draft_token_num = 1, 5
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_branching_tree_mask()
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    expected_next = [1, 3, -1, -1, -1]
+    for i, exp in enumerate(expected_next):
+        got = inputs["retrive_next_token"][0, i].item()
+        assert got == exp, f"retrive_next_token[0,{i}]={got}, expected {exp}"
+
+
+def test_branching_tree_next_sibling():
+    """retrive_next_sibling links tokens that share the same parent."""
+    from sglang.jit_kernel.ngram_utils import reconstruct_indices_from_tree_mask
+
+    bs, draft_token_num = 1, 5
+    inputs = make_inputs(bs, draft_token_num)
+    inputs["tree_mask"] = build_branching_tree_mask()
+
+    reconstruct_indices_from_tree_mask(
+        **inputs, batch_size=bs, draft_token_num=draft_token_num
+    )
+
+    expected_sibling = [-1, 2, -1, 4, -1]
+    for i, exp in enumerate(expected_sibling):
+        got = inputs["retrive_next_sibling"][0, i].item()
+        assert got == exp, f"retrive_next_sibling[0,{i}]={got}, expected {exp}"
+
+
+# ---------------------------------------------------------------------------
+# JIT vs AOT cross-validation
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "bs,draft_token_num",
+    [
+        (1, 4),
+        (2, 8),
+        (4, 16),
+    ],
+)
+def test_vs_aot(bs, draft_token_num):
+    try:
+        from sgl_kernel import reconstruct_indices_from_tree_mask as reconstruct_aot
+    except ImportError:
+        pytest.skip("sgl_kernel not available")
+
+    from sglang.jit_kernel.ngram_utils import (
+        reconstruct_indices_from_tree_mask as reconstruct_jit,
+    )
+
+    inputs_jit = make_inputs(bs, draft_token_num)
+    inputs_jit["tree_mask"] = build_linear_chain_mask(bs, draft_token_num)
+    inputs_jit["verified_seq_len"].fill_(7)
+
+    inputs_aot = {
+        k: v.clone() if isinstance(v, torch.Tensor) else v
+        for k, v in inputs_jit.items()
+    }
+
+    reconstruct_jit(**inputs_jit, batch_size=bs, draft_token_num=draft_token_num)
+    reconstruct_aot(**inputs_aot, batch_size=bs, draft_token_num=draft_token_num)
+
+    assert torch.equal(
+        inputs_jit["positions"], inputs_aot["positions"]
+    ), "positions mismatch"
+    assert torch.equal(
+        inputs_jit["retrive_index"], inputs_aot["retrive_index"]
+    ), "retrive_index mismatch"
+    assert torch.equal(
+        inputs_jit["retrive_next_token"], inputs_aot["retrive_next_token"]
+    ), "retrive_next_token mismatch"
+    assert torch.equal(
+        inputs_jit["retrive_next_sibling"], inputs_aot["retrive_next_sibling"]
+    ), "retrive_next_sibling mismatch"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])