flashmap 0.1.0__py3-none-any.whl

flashmap/sglang/sampler.py

@@ -0,0 +1,188 @@
+# coding=utf-8
+"""DSPARK's Markov head and draft sampler, specialized for FlashMap.
+
+sglang's DSPARK worker requires a Markov head (a low-rank bias on each draft
+position's logits from the previous token) and samples the block one position
+at a time through it. FlashMap has no such head: its block positions are
+independent given the noise. :class:`IdentityMarkov` stands in for the head
+(rank 1, never loaded, contributes nothing) so DSPARK samples the map's logits
+unchanged and verifies them with its usual rejection sampling.
+
+With no bias, the positions are independent, yet DSPARK's graph-folded sampler
+would still walk them one at a time (an exponential-noise fill and two
+online-softmax kernels per position, or an argmax per position for greedy
+batches). :class:`FlashMapDraftSampler` samples the whole block at once inside
+the draft's CUDA graph: the logits written into DSPARK's buffer by one GEMM
+against the folded head, then one ``SampleStepTokens`` pass with fresh
+exponential noise per row and each request's temperature (exactly DSPARK's
+per-position sampling, since the rows are independent; greedy rows take their
+argmax). Drafted indices of a restricted vocabulary map back to token ids in
+the graph too, and the logits rejection sampling reads stay over the full
+vocabulary (-inf outside the draft vocabulary).
+
+``FLASHMAP_FUSED_SAMPLER=1`` (opt-in, measured on Qwen3-8B only) samples with
+one kernel instead: each greedy row's
+argmax, each other row Gumbel-max-sampled at its temperature with in-kernel
+Philox noise, so there is no noise buffer, and the full-vocabulary logits are
+scattered only when DSPARK reads them for a sampling batch, so greedy batches
+never pay for them.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from sglang.kernels.ops.speculative.dspark.dspark_draft_model import SampleStepTokens
+from sglang.srt.models.dspark import VanillaMarkov
+from sglang.srt.speculative.dspark_components.dspark_draft_sampler import (
+    DsparkDraftSampler,
+)
+from sglang.srt.speculative.spec_tp_sync import SpecTpSyncSite
+
+from .kernels import sample_rows
+
+
+class IdentityMarkov(VanillaMarkov):
+    """A Markov head that leaves the draft's logits unchanged.
+
+    With a restricted draft vocabulary (``hot_ids``, set by the draft) its
+    eager block sampler also maps drafted indices back to token ids and
+    scatters the logits back over the full vocabulary (-inf elsewhere, so
+    rejection sampling stays exact)."""
+
+    def __init__(self, *, vocab_size: int) -> None:
+        super().__init__(vocab_size=vocab_size, markov_rank=1)
+        for weight in self.parameters():
+            weight.requires_grad_(False)
+            torch.nn.init.zeros_(weight)
+        self.hot_ids = None
+
+    def apply_step_logits(self, logits, *args, **kwargs):
+        return logits
+
+    def apply_block_logits(self, base_logits, *args, **kwargs):
+        return base_logits
+
+    def sample_block_greedy_fused(self, *args, **kwargs):
+        return None  # the caller falls back to sample_block
+
+    def sample_block(
+        self,
+        base_logits,
+        *,
+        first_prev_tokens,
+        hidden_states,
+        sampler,
+        collect_corrected: bool = True,
+    ):
+        del first_prev_tokens, hidden_states
+        tokens = torch.stack(
+            [sampler(base_logits[:, i], i) for i in range(base_logits.shape[1])], dim=1
+        )
+        if self.hot_ids is None:
+            return tokens, base_logits if collect_corrected else None
+        logits = None
+        if collect_corrected:
+            logits = base_logits.new_full(
+                (*base_logits.shape[:2], self.vocab_size), float("-inf")
+            )
+            logits.index_copy_(2, self.hot_ids, base_logits)
+        return self.hot_ids[tokens], logits
+
+
+class _FullLogits:
+    """DSPARK's ``corrected_out`` for a restricted draft vocabulary: slicing
+    the first ``rows`` scatters those rows' draft-vocabulary logits into the
+    full-vocabulary buffer (-inf elsewhere) and returns them."""
+
+    def __init__(self, full: torch.Tensor, hot: torch.Tensor, hot_ids: torch.Tensor):
+        self.full, self.hot, self.hot_ids = full, hot, hot_ids
+        self.full.fill_(float("-inf"))
+
+    def __getitem__(self, index: slice) -> torch.Tensor:
+        assert index.start is None and index.step is None, index
+        rows = index.stop
+        return self.full[:rows].index_copy_(1, self.hot_ids, self.hot[:rows])
+
+
+class FlashMapDraftSampler(DsparkDraftSampler):
+    """DSPARK's graph-folded draft sampler for a FlashMap draft: the whole block
+    in one pass (see the module docstring). Same contract."""
+
+    # DsparkDraftSampler state this sampler takes over (sglang 0.5.20).
+    _INHERITED = (
+        "model", "gamma", "sample_from_anchor", "query_token_num", "out",
+        "confidence_out", "folded_sampling", "temperatures", "greedy_mask",
+        "exp_noise", "corrected_out", "_tp_sync",
+    )  # fmt: skip
+
+    def __init__(self, sampler: DsparkDraftSampler, *, max_bs: int, device):
+        missing = [name for name in self._INHERITED if not hasattr(sampler, name)]
+        assert not missing, f"DsparkDraftSampler lacks {missing}: unsupported sglang"
+        self.__dict__.update(sampler.__dict__)
+        assert self.sample_from_anchor and self.confidence_out is None
+        self.fused = os.environ.get("FLASHMAP_FUSED_SAMPLER", "0") == "1"
+        self.hot_logits = None
+        if not self.folded_sampling:
+            return
+        hot_ids = self.model.hot_ids
+        rows, draft_vocab = int(max_bs) * self.gamma, self.model.folded_head[0].shape[0]
+        self.exp_noise = None
+        if not self.fused:
+            # Exponential noise for every row of the block, not one position's.
+            self.row_noise = torch.empty(
+                (rows, draft_vocab), dtype=torch.float32, device=device
+            )
+            if hot_ids is not None:
+                self.corrected_out.fill_(float("-inf"))
+            return
+        self.seed = torch.zeros((), dtype=torch.int64, device=device)
+        if hot_ids is not None:
+            # The draft-vocabulary logits outlive the graph replay until DSPARK
+            # reads them (``_FullLogits``).
+            self.hot_logits = torch.empty(
+                (rows, draft_vocab), dtype=self.corrected_out.dtype, device=device
+            )
+            self.corrected_out = _FullLogits(
+                self.corrected_out, self.hot_logits, hot_ids
+            )
+
+    def __call__(self, hidden_states, input_ids):
+        del input_ids
+        bs = hidden_states.shape[0] // self.query_token_num
+        rows = bs * self.gamma
+        hot_ids = self.model.hot_ids
+        if not self.folded_sampling:
+            logits = self.model.head_logits(hidden_states)
+            tokens = self._tp_sync.sync(
+                SpecTpSyncSite.DSPARK_GRAPH_GREEDY, logits.argmax(dim=-1)
+            )
+        else:
+            temperatures = self.temperatures[:bs].repeat_interleave(self.gamma)
+            greedy = self.greedy_mask[:bs].repeat_interleave(self.gamma)
+            if self.fused:
+                out = self.hot_logits if hot_ids is not None else self.corrected_out
+                logits = out[:rows]
+                logits.copy_(self.model.head_logits(hidden_states))
+                self.seed.add_(1)
+                tokens = sample_rows(logits, temperatures, greedy, self.seed)
+            else:
+                logits = self.model.head_logits(hidden_states)
+                if hot_ids is None:
+                    self.corrected_out[:rows].copy_(logits)
+                else:
+                    self.corrected_out[:rows].index_copy_(1, hot_ids, logits)
+                tokens = SampleStepTokens.execute(
+                    step_logits=logits,
+                    temperatures=temperatures,
+                    greedy_mask=greedy,
+                    exp_noise=self.row_noise[:rows].exponential_(),
+                )
+            tokens = self._tp_sync.sync(SpecTpSyncSite.DSPARK_GRAPH_SAMPLE, tokens)
+        if hot_ids is not None:
+            tokens = hot_ids[tokens]
+        self.out[:rows].copy_(tokens.reshape(-1))
+
+
+__all__ = ["FlashMapDraftSampler", "IdentityMarkov"]

flashmap/sglang/tree.py

@@ -0,0 +1,78 @@
+# coding=utf-8
+"""The verify tree of a multi-candidate FlashMap draft (pure torch, no sglang
+imports)."""
+
+import torch
+import torch.nn.functional as F
+
+
+def build_tree_abi(candidates: torch.Tensor, width: int):
+    """The verify tree from N candidate blocks per request, in EAGLE's encoding.
+
+    ``candidates (bs, N, K)`` are the drafts after each request's anchor (the
+    tree's root). Their trie has a node per distinct prefix; a node's vote is
+    how many candidates pass through it, a Monte Carlo estimate of the prefix's
+    probability under the map that never grows with depth. The ``width - 1``
+    nodes with the most votes (ties: shallower, then lower candidate) are kept,
+    so every kept node's parent is kept too, and listed parents first (by depth,
+    then candidate). A trie with fewer nodes is padded with copies of the first
+    node under the root, listed after it: verification takes the first matching
+    child, so a copy is never accepted.
+
+    Returns ``tokens (bs, width - 1)``; ``ids (bs, width - 1)``, node n's id
+    ``parent(n) * width + n`` (root = 0), which build_tree_kernel_efficient
+    decodes with topk = width and parent_list[:, n] = ids[:, n - 1]; and
+    ``stats (bs, 2)``: the trie's node count before the cut and the number of
+    distinct candidates.
+    All on device, fixed shapes, no host synchronization.
+    """
+    bs, n, k = candidates.shape
+    m = width - 1
+    device = candidates.device
+    # eq[b, i, j, d]: candidates i and j agree on positions 0..d.
+    eq = (
+        (candidates.unsqueeze(2) == candidates.unsqueeze(1))
+        .to(torch.int32)
+        .cumprod(dim=-1)
+    )
+    votes = eq.sum(dim=2)  # (bs, N, K)
+    first = eq.argmax(dim=2)  # the lowest candidate with the same prefix
+    cand = torch.arange(n, device=device).view(1, n, 1)
+    depth = torch.arange(k, device=device).view(1, 1, k)
+    real = first == cand  # one representative slot per trie node
+    key = torch.where(
+        real, votes * (n * k) + (k - 1 - depth) * n + (n - 1 - cand), -1
+    ).view(bs, n * k)
+    if n * k < m:
+        key = F.pad(key, (0, m - n * k), value=-1)
+    # Keys of real nodes are distinct, so a sort ranks exactly as topk would;
+    # topk's radix select costs ~15 us per step here against ~4 us for a sort.
+    top, slot = key.sort(dim=1, descending=True)
+    top, slot = top[:, :m], slot[:, :m]
+    pad = top < 0
+    slot = slot.clamp(max=n * k - 1)
+    i, d = slot // k, slot % k
+    order = torch.where(
+        pad, n * k + torch.arange(m, device=device).view(1, m), d * n + i
+    )
+    order, perm = order.sort(dim=1)
+    slot, pad, i, d = (
+        slot.gather(1, perm),
+        pad.gather(1, perm),
+        i.gather(1, perm),
+        d.gather(1, perm),
+    )
+    index = torch.arange(1, m + 1, device=device).view(1, m).expand(bs, m)
+    # Node index of each kept slot; padding goes to a spare last column.
+    node_of = torch.zeros(bs, n * k + 1, dtype=torch.int64, device=device)
+    node_of.scatter_(1, torch.where(pad, n * k, slot), index)
+    # A node's parent is the representative of its prefix one shorter.
+    up = first.view(bs, n * k).gather(1, (slot - 1).clamp(min=0)) * k + (d - 1).clamp(
+        min=0
+    )
+    parent = torch.where((d == 0) | pad, 0, node_of.gather(1, up))
+    tokens = candidates.reshape(bs, n * k).gather(1, slot)
+    tokens = torch.where(pad, tokens[:, :1], tokens)
+    distinct = real[:, :, k - 1].sum(dim=1)  # candidates that are their own last node
+    stats = torch.stack([real.view(bs, n * k).sum(dim=1), distinct], dim=1)
+    return tokens, parent * width + index, stats

flashmap/sglang/tree_worker.py

@@ -0,0 +1,396 @@
+# coding=utf-8
+"""Tree verification of multi-candidate FlashMap drafts inside sglang's DSPARK
+worker.
+
+Opt-in with ``FLASHMAP_TREE_N=<N>`` on a FlashMap server, with
+``SGLANG_DSPARK_FOLDED_PROPOSAL=0`` (the tree path does its own proposal), a
+target attention backend that honours a tree mask (``fa3``, ``flashinfer`` or
+``triton``) and ``--page-size 1``. Greedy decoding only: the scheduler refuses
+a request that samples (temperature > 0), asks for logprobs or constrains its
+output (JSON schema, regex, EBNF, structural tag) when it arrives, with an error
+response saying why (``tree_refusal``); the server keeps serving.
+
+Each decode step, replacing DSPARK's chain step:
+
+1. The draft (one CUDA-graph replay) drafts up to N blocks per request from
+   independent noise in one forward, merges them into a trie and keeps the
+   block_size nodes the most candidates pass through, in EAGLE's tree encoding
+   (``tree.build_tree_abi``). With the root that is DSPARK's own
+   verify width, so the target's verify graphs and KV reservation are unchanged.
+2. ``build_tree_kernel_efficient`` turns it into the tree mask, positions and
+   child/sibling links; the target verifies the tree in one forward over the
+   request's window ``req_to_token[prefix : prefix + W]``.
+3. ``verify_tree_greedy`` walks the tree along the target's argmax. Instead of
+   moving the accepted nodes' KV, the window's ``req_to_token`` entries are
+   reordered so the accepted path comes first (``tree_commit``, one launch);
+   the path's hidden states become draft context there.
+
+``FLASHMAP_SPLITKV_VERIFY=1`` (with ``--attention-backend triton``) runs the
+target's verify attention with ``verify_attention.SplitKVVerify``: GQA
+rows packed per KV head, the prefix split across programs, the tree mask only
+on each request's W x W block.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+TREE_ENABLED = int(os.environ.get("FLASHMAP_TREE_N", "0")) >= 1
+SPLITKV_ENABLED = os.environ.get("FLASHMAP_SPLITKV_VERIFY", "0") == "1"
+
+if TREE_ENABLED or SPLITKV_ENABLED:
+    from . import check_sglang_version
+
+    check_sglang_version()
+
+    from sglang.srt.layers.attention import flashattention_backend, triton_backend
+    from sglang.srt.lora.layers import unwrap_lora_layer
+    from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler
+    from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+    from sglang.srt.speculative import dflash_utils
+    from sglang.srt.speculative.dflash_info import DFlashVerifyInput
+    from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+    from sglang.srt.speculative.dflash_utils import (
+        apply_dflash_verify_logits_adjustments,
+        resolve_dflash_verify_mask_policy,
+    )
+    from sglang.srt.speculative.dspark_components import dspark_worker_v2
+    from sglang.srt.speculative.dspark_components.dspark_draft import (
+        make_next_draft_input,
+    )
+    from sglang.srt.speculative.dspark_components.dspark_planner import (
+        alloc_verify_window,
+    )
+    from sglang.srt.speculative.eagle_utils import (
+        TreeMaskMode,
+        build_tree_kernel_efficient,
+        verify_tree_greedy_func,
+    )
+
+    from .kernels import tree_commit
+    from .verify_attention import SplitKVVerify
+
+
+def _install_splitkv() -> None:
+    """The split-KV verify on the target's Triton backend (sglang routes a
+    verify through ``verify_splitkv_fwd`` when ``use_verify_splitkv`` is set,
+    upstream on AMD only), installed before the target's graphs are captured."""
+    triton_init = triton_backend.TritonAttnBackend.__init__
+
+    def init(self, model_runner, *args, **kwargs):
+        triton_init(self, model_runner, *args, **kwargs)
+        if getattr(model_runner, "is_draft_worker", False) or self.use_mla:
+            return
+        self.use_verify_shared_kv = False
+        self.use_verify_splitkv = True
+        self._flashmap_splitkv = SplitKVVerify()
+        # The tree mask as each request's W x W block (QLEN_ONLY).
+        self._flashmap_splitkv.qlen_mask = TREE_ENABLED
+        self.verify_splitkv_fwd = torch.compiler.disable(
+            self._flashmap_splitkv.__call__
+        )
+        logger.info("FlashMap split-KV verify attention on the target's Triton backend")
+
+    triton_backend.TritonAttnBackend.__init__ = init
+
+
+_CONSTRAINTS = ("json_schema", "regex", "ebnf", "structural_tag")
+
+
+def tree_refusal(req):
+    """Why a tree-mode server refuses ``req`` (an sglang ``Req``), or None.
+
+    The tree verify walks the target's argmax, so it serves greedy requests
+    only; a request that samples, asks for logprobs or constrains its output is
+    refused when it arrives rather than failing the batch it would join."""
+    params = req.sampling_params
+    asks = []
+    if params.top_k > 1:  # sglang's greedy is top_k 1 (temperature 0)
+        asks.append("sampling (temperature > 0)")
+    if req.return_logprob:
+        asks.append("logprobs")
+    if any(getattr(params, name, None) for name in _CONSTRAINTS):
+        asks.append("constrained decoding")
+    if not asks:
+        return None
+    return (
+        "this server drafts FlashMap trees (FLASHMAP_TREE_N), which serve greedy "
+        f"requests (temperature 0) only; the request asks for {' and '.join(asks)}. "
+        "Send it with temperature 0, or serve without FLASHMAP_TREE_N."
+    )
+
+
+def _install_tree() -> None:
+    # Capture the target's verify graphs with a custom-mask buffer on the
+    # backends that read one (DSPARK otherwise captures them mask-free).
+    dflash_utils._DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS = (
+        dflash_utils._DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS
+        - {"FlashInferAttnBackend", "TritonAttnBackend"}
+    )
+    fa3_init = flashattention_backend.FlashAttentionBackend.__init__
+
+    def fa3_tree_init(self, model_runner, *args, **kwargs):
+        """FA3 verifies a tree, as for EAGLE with top-k > 1, by cascade
+        attention; every such branch keys on ``topk > 1``, which DSPARK pins to
+        1, so any top-k above 1 selects it on the target."""
+        fa3_init(self, model_runner, *args, **kwargs)
+        if not model_runner.is_draft_worker and self.topk <= 1:
+            self.topk = 2
+
+    flashattention_backend.FlashAttentionBackend.__init__ = fa3_tree_init
+    forward_decode = dspark_worker_v2.DSparkWorkerV2._forward_decode
+
+    def _forward_decode(self, batch, on_publish, grammar_barrier=None):
+        if not getattr(self.draft_model, "tree_mode", False) or (
+            batch.forward_mode.is_idle()
+        ):
+            return forward_decode(self, batch, on_publish, grammar_barrier)
+        return _forward_decode_tree(self, batch, on_publish)
+
+    dspark_worker_v2.DSparkWorkerV2._forward_decode = _forward_decode
+
+    # Refuse what the tree verify cannot serve when it arrives: every request
+    # (new, or out of the grammar queue) enters the waiting queue here. An
+    # aborted request is still run once (a one-token prefill, and with overlap
+    # scheduling possibly one decode step) before its error is returned, so it
+    # is also made greedy (aborting already drops its grammar and logprobs):
+    # every batch it joins stays one the tree step serves.
+    add_request_to_queue = Scheduler._add_request_to_queue
+
+    def _add_request_to_queue(self, req, is_retracted: bool = False):
+        if not is_retracted and not req.finished():
+            refusal = tree_refusal(req)
+            if refusal is not None:
+                req.set_finish_with_abort(refusal)
+                req.sampling_params.top_k = 1
+        return add_request_to_queue(self, req, is_retracted)
+
+    Scheduler._add_request_to_queue = _add_request_to_queue
+
+
+def _tree_state(self):
+    """Per-worker constants of the tree path, checked once."""
+    state = getattr(self, "_flashmap_tree", None)
+    if state is not None:
+        return state
+    backend, reads_mask = resolve_dflash_verify_mask_policy(
+        self.model_runner.attn_backend
+    )
+    attn = self.model_runner.attn_backend
+    while getattr(attn, "full_attn_backend", None) is not None:
+        attn = attn.full_attn_backend
+    if not (reads_mask or (backend == "FlashAttentionBackend" and attn.topk > 1)):
+        raise RuntimeError(
+            "FlashMap tree verify needs a target attention backend that reads a "
+            f"tree mask (fa3, flashinfer or triton), got {backend}"
+        )
+    if self._draft_sampler is not None:
+        raise RuntimeError(
+            "FlashMap tree verify does its own proposal: set "
+            "SGLANG_DSPARK_FOLDED_PROPOSAL=0"
+        )
+    width = int(self.draft_model.tree_width)
+    assert width == self.verify_num_draft_tokens == self.gamma + 1, (
+        width,
+        self.verify_num_draft_tokens,
+    )
+    # Reordering a window's req_to_token entries keeps each slot's KV where it
+    # is only if every token is its own page.
+    if self.model_runner.page_size != 1:
+        raise RuntimeError("FlashMap tree verify needs --page-size 1")
+    splitkv = getattr(attn, "_flashmap_splitkv", None)
+    state = {
+        "width": width,
+        "mask_mode": (
+            TreeMaskMode.QLEN_ONLY
+            if getattr(splitkv, "qlen_mask", False)
+            else TreeMaskMode.FULL_MASK
+        ),
+        "offsets": torch.arange(width, device=self.device, dtype=torch.int64),
+        "req_to_token": self.model_runner.req_to_token_pool.req_to_token,
+    }
+    self._flashmap_tree = state
+    logger.info(
+        "FlashMap tree verify: up to %d candidates per request (budget %d per "
+        "step), %d verify tokens, target backend %s",
+        self.draft_model.tree_max_n,
+        self.draft_model.tree_budget,
+        width,
+        backend,
+    )
+    return state
+
+
+def _forward_decode_tree(self, batch, on_publish):
+    if batch.spec_info is None:
+        batch.spec_info = DFlashDraftInputV2.create_idle_input(device=self.device)
+    draft_input = batch.spec_info
+    sampling_info = batch.sampling_info
+    if (
+        batch.has_grammar
+        or batch.return_logprob
+        or not (sampling_info is None or sampling_info.is_all_greedy)
+    ):
+        # The scheduler refuses such requests on arrival (tree_refusal), so
+        # this is unreachable. Falling back to DSPARK's chain step is not safe:
+        # tried with greedy logprob requests, it broke their outputs on a
+        # tree-mode server.
+        raise NotImplementedError(
+            "FlashMap tree verify (FLASHMAP_TREE_N) serves greedy requests "
+            "(temperature 0) without grammar or logprobs only"
+        )
+    state = _tree_state(self)
+    width, depth = state["width"], self.gamma
+
+    batch.seq_lens.record_stream(torch.get_device_module(self.device).current_stream())
+    bs, device = len(batch.seq_lens), self.device
+    prefix_lens = batch.seq_lens
+    target_model = self.target_worker.model_runner.model
+    seq_lens_cpu = (
+        batch.seq_lens_cpu
+        if batch.seq_lens_cpu is not None
+        else draft_input.nxt_kv_lens_cpu
+    )
+    window = alloc_verify_window(
+        batch=batch,
+        bs=bs,
+        device=device,
+        verify_num_draft_tokens=width,
+        block_pos_offsets=state["offsets"],
+        model_runner=self.model_runner,
+    )
+
+    # 1. Draft: one graph replay fills the draft's tree buffers.
+    with self._draft_context():
+        self._proposer._run_forward(
+            batch=batch,
+            draft_input=draft_input,
+            verify_window=window,
+            bs=bs,
+            device=device,
+            embed_module=unwrap_lora_layer(target_model.get_input_embeddings()),
+            draft_sampler=None,
+            sampling_info=sampling_info,
+        )
+    node_tokens, node_ids, parent_list, _ = self.draft_model.tree_outputs(bs)
+
+    # 2. The tree's mask, positions and links; verify it.
+    (
+        tree_mask,
+        positions,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        verify_ids,
+    ) = build_tree_kernel_efficient(
+        draft_input.bonus_tokens.view(-1),
+        parent_list,
+        node_ids,
+        node_tokens,
+        prefix_lens,
+        int(seq_lens_cpu.sum()) + bs * width,
+        topk=width,
+        spec_steps=depth,
+        num_verify_tokens=width,
+        tree_mask_mode=state["mask_mode"],
+    )
+    verify_input = DFlashVerifyInput(
+        draft_token=verify_ids,
+        positions=positions,
+        draft_token_num=width,
+        custom_mask=tree_mask,
+        capture_hidden_mode=CaptureHiddenMode.FULL,
+        live_seq_lens_cpu=batch.seq_lens_cpu,
+    )
+    batch.out_cache_loc = window.verify_cache_loc
+    executor = self._verify_executor
+    seq_lens_cpu_backup, seq_lens_sum_backup = batch.seq_lens_cpu, batch.seq_lens_sum
+    if not executor._verify_backend_self_adds_seq_lens():
+        if seq_lens_cpu_backup is not None:
+            batch.seq_lens_cpu = seq_lens_cpu_backup + width
+            batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
+        elif draft_input.nxt_kv_lens_cpu is not None:
+            batch.seq_lens_cpu = draft_input.nxt_kv_lens_cpu
+            batch.seq_lens_sum = int(draft_input.nxt_kv_lens_sum)
+    verified = executor._forward_prepared_verify(
+        batch=batch,
+        verify_input=verify_input,
+        seq_lens_cpu_backup=seq_lens_cpu_backup,
+        seq_lens_sum_backup=seq_lens_sum_backup,
+    )
+    logits_output = verified.logits_output
+    logits = logits_output.next_token_logits
+    if sampling_info is not None:
+        apply_dflash_verify_logits_adjustments(
+            next_token_logits=logits, sampling_info=sampling_info, draft_token_num=width
+        )
+
+    # 3. Accept the longest path the target agrees with.
+    predict = torch.zeros(bs * width, dtype=torch.int32, device=device)
+    accept_index = torch.full((bs, depth + 1), -1, dtype=torch.int32, device=device)
+    num_correct = torch.empty((bs,), dtype=torch.int32, device=device)
+    verify_tree_greedy_func(
+        predicts=predict,
+        accept_index=accept_index,
+        accept_token_num=num_correct,
+        candidates=verify_ids.view(bs, width),
+        retrieve_index=retrieve_index,
+        retrieve_next_token=retrieve_next_token,
+        retrieve_next_sibling=retrieve_next_sibling,
+        target_predict=logits.argmax(dim=-1).view(bs, width),
+    )
+    # 4. Commit without moving KV: reorder each request's window so the
+    # accepted path (root first) is its first commit_lens slots, where its KV
+    # already is; the path's hidden states become draft context there (rows off
+    # the path go to the pool's dummy slot 0).
+    out_tokens, bonus, commit_lens, new_seq_lens, inject_locs = tree_commit(
+        accept_index,
+        num_correct,
+        window.verify_cache_loc_2d,
+        predict,
+        batch.req_pool_indices,
+        prefix_lens,
+        state["req_to_token"],
+    )
+    # sglang's scheduler assumes a forward only reads req_to_token: its
+    # write-after-read barrier must wait for this write too.
+    reordered = torch.get_device_module(self.device).Event()
+    reordered.record()
+    self.model_runner.shared_read_done_event = reordered
+    if on_publish is not None:
+        on_publish(new_seq_lens)
+    executor.kv_injector.inject_target_hidden(
+        target_hidden=logits_output.hidden_states,
+        cache_loc=inject_locs,
+        positions=positions,
+    )
+    logits_output.hidden_states = None
+    return GenerationBatchResult(
+        logits_output=logits_output,
+        next_token_ids=out_tokens.reshape(-1),
+        accept_lens=commit_lens,
+        block_accept_lens=commit_lens,
+        cap_lens=None,
+        can_run_cuda_graph=verified.can_run_cuda_graph,
+        next_draft_input=make_next_draft_input(
+            bonus_tokens=bonus, new_seq_lens=new_seq_lens
+        ),
+        speculative_num_draft_tokens=int(self.verify_num_draft_tokens),
+        new_seq_lens=new_seq_lens,
+    )
+
+
+# Installed once, however often the module is imported.
+if SPLITKV_ENABLED and not getattr(
+    triton_backend, "_flashmap_sglang_splitkv", False
+):
+    _install_splitkv()
+    triton_backend._flashmap_sglang_splitkv = True
+if TREE_ENABLED and not getattr(dspark_worker_v2, "_flashmap_sglang_tree", False):
+    _install_tree()
+    dspark_worker_v2._flashmap_sglang_tree = True