feat(triton): start on the MTP path with three draft kernels

compute_seqlens_k_zos, confidence_head_fused and dflash_input_store, all identical to the baseline:
15 cases across the three.

`compute_seqlens_k_zos` masks the remap lookup rather than clamping it: `-1` marks a fresh request,
and an unmasked `PREV_ACCEPTED + (-1)` would read the element before the buffer.

`confidence_head_fused` reads the weight as one flat vector and splits it at `n_embed`, so the
hidden and Markov halves are one accumulation instead of two passes. Without a Markov embedding the
second pointer is filled with the first tensor rather than left null, since the read is compiled out
but the argument still has to be a pointer. The scaled logit is clamped to +-60 before the sigmoid,
as the baseline does, so a small temperature cannot produce inf through `exp`.

`dflash_input_store` accumulates in fp32 even though both sides are bf16: the baseline sums into an
fp32 fragment, and a bf16 accumulator would round differently for hc_mult > 1.

One trap: `tl.next_power_of_2` inside a kernel body is not a compile-time constant, so the Markov
block size is resolved in the factory and passed as a constexpr.

deepinfer/model_executor/layers/tilelang_ops/triton/compute_seqlens_k_zos.py

@@ -0,0 +1,64 @@
+"""Triton version of the MTP sequence-length fixup.
+
+One kernel, one array: fold the current step's query length into the key length, and for a request
+that continues from the previous batch, replace the extra speculative slots with the number of
+tokens that batch actually accepted.
+
+The `prev_idx` lookup is masked rather than clamped: `-1` means "fresh request", and an unmasked
+`PREV_ACCEPTED + (-1)` would read the element before the buffer.
+"""
+
+from typing import Any
+
+import triton
+import triton.language as tl
+
+_BLOCK = 128
+
+
+@triton.jit
+def _compute_seqlens_k_zos_kernel(
+    SEQLENS_K,
+    SEQLENS_Q,
+    REMAP,
+    PREV_ACCEPTED,
+    n_seq,
+    nextn: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    idx = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    live = idx < n_seq
+
+    sk = tl.load(SEQLENS_K + idx, mask=live, other=0) + tl.load(SEQLENS_Q + idx, mask=live, other=0)
+    prev_idx = tl.load(REMAP + idx, mask=live, other=-1)
+    has_prev = live & (prev_idx >= 0)
+    # Only the rows that continue from the previous batch take the correction; the rest keep `sk`.
+    accepted = tl.load(PREV_ACCEPTED + tl.maximum(prev_idx, 0), mask=has_prev, other=0)
+    sk = tl.where(has_prev, sk - nextn + accepted, sk)
+    tl.store(SEQLENS_K + idx, sk, mask=live)
+
+
+def compute_seqlens_k_zos_tl(nextn: int) -> Any:
+    """Build the kernel; `nextn` is the number of speculative positions per step."""
+
+    def launch(
+        seqlens_k,
+        seqlens_q,
+        remap,
+        prev_num_accepted_tokens,
+    ) -> None:
+        n_seq = seqlens_k.shape[0]
+        if n_seq == 0:
+            return
+        _compute_seqlens_k_zos_kernel[(triton.cdiv(n_seq, _BLOCK),)](
+            seqlens_k,
+            seqlens_q,
+            remap,
+            prev_num_accepted_tokens,
+            n_seq,
+            nextn=nextn,
+            BLOCK=_BLOCK,
+            num_warps=4,
+        )
+
+    return launch

deepinfer/model_executor/layers/tilelang_ops/triton/dflash.py

@@ -0,0 +1,140 @@
+"""Triton versions of the two DFlash (MTP draft) kernels.
+
+`confidence_head_fused` is a dot product over `cat([hidden, markov_emb], -1)` followed by a scaled,
+clamped sigmoid: the MTP draft's confidence score. It reads the weight as one flat vector and splits
+it at `n_embed`, which keeps the two halves in one accumulation instead of two passes.
+
+`dflash_input_store` reduces the `hc_mult` copies of a row down to one: a plain copy when there is a
+single copy, the mean otherwise. The accumulation is in fp32 even though the inputs and the output
+are bf16, because the baseline sums into an fp32 fragment and a bf16 accumulator would differ.
+"""
+
+from typing import Any
+
+import triton
+import triton.language as tl
+
+_CONF_BLOCK = 1024
+_STORE_BLOCK = 1024
+
+
+@triton.jit
+def _confidence_head_fused_kernel(
+    X0,
+    X1,
+    WEIGHT,
+    OUT,
+    temperature,
+    hidden_stride,
+    out_stride,
+    N_EMBED: tl.constexpr,
+    MARKOV_RANK: tl.constexpr,
+    MARKOV_BLOCK: tl.constexpr,
+    HAS_MARKOV: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+
+    # The hidden part. `N_EMBED` is a compile-time multiple of the block by the factory's choice, so
+    # the tail is masked rather than assumed away.
+    cols = tl.arange(0, BLOCK)
+    acc = tl.zeros((), tl.float32)
+    for start in range(0, N_EMBED, BLOCK):
+        live = start + cols < N_EMBED
+        xv = tl.load(X0 + row * hidden_stride + start + cols, mask=live, other=0.0).to(tl.float32)
+        wv = tl.load(WEIGHT + start + cols, mask=live, other=0.0).to(tl.float32)
+        acc += tl.sum(xv * wv, axis=0)
+
+    if HAS_MARKOV:
+        # The Markov embedding continues the same flat weight vector, so it is one more dot product
+        # with an offset rather than a separate accumulation.
+        mcols = tl.arange(0, MARKOV_BLOCK)
+        mlive = mcols < MARKOV_RANK
+        mv = tl.load(X1 + row * MARKOV_RANK + mcols, mask=mlive, other=0.0).to(tl.float32)
+        mw = tl.load(WEIGHT + N_EMBED + mcols, mask=mlive, other=0.0).to(tl.float32)
+        acc += tl.sum(mv * mw, axis=0)
+
+    # The baseline clamps the scaled logit to +-60 before the sigmoid, so an extreme temperature
+    # cannot produce inf/nan through exp.
+    scaled = tl.minimum(tl.maximum(acc / temperature, -60.0), 60.0)
+    tl.store(OUT + row * out_stride, 1.0 / (1.0 + tl.exp(-scaled)))
+
+
+@triton.jit
+def _dflash_input_store_kernel(
+    X,
+    OUT,
+    out_stride,
+    N_EMBED: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    cols = tile * BLOCK + tl.arange(0, BLOCK)
+    live = cols < N_EMBED
+
+    if HC_MULT == 1:
+        value = tl.load(X + row * N_EMBED + cols, mask=live, other=0.0)
+    else:
+        total = tl.zeros((BLOCK,), tl.float32)
+        for j in tl.static_range(0, HC_MULT):
+            total += tl.load(X + row * HC_MULT * N_EMBED + j * N_EMBED + cols, mask=live, other=0.0).to(tl.float32)
+        value = (total * (1.0 / HC_MULT)).to(tl.bfloat16)
+
+    tl.store(OUT + row * out_stride + cols, value, mask=live)
+
+
+def confidence_head_fused_tl(n_embed: int, markov_rank: int, has_markov_emb: bool) -> Any:
+    """Build the confidence head for one (n_embed, markov_rank) shape."""
+
+    def launch(
+        x0: Any,
+        x1: Any,
+        weight: Any,
+        out: Any,
+        temperature: float,
+    ) -> None:
+        n = x0.shape[0]
+        if n == 0:
+            return
+        _confidence_head_fused_kernel[(n,)](
+            x0,
+            # Without a Markov embedding the kernel must not dereference the second input; passing
+            # `x0` keeps the pointer valid while `HAS_MARKOV` compiles the read out.
+            x1 if x1 is not None else x0,
+            weight,
+            out,
+            temperature,
+            x0.stride(0),
+            out.stride(0),
+            N_EMBED=n_embed,
+            MARKOV_RANK=markov_rank,
+            # Resolved here: a helper call inside the kernel body is not a compile-time constant.
+            MARKOV_BLOCK=triton.next_power_of_2(max(markov_rank, 1)),
+            HAS_MARKOV=has_markov_emb,
+            BLOCK=_CONF_BLOCK,
+            num_warps=4,
+        )
+
+    return launch
+
+
+def dflash_input_store_tl(n_embed: int, hc_mult: int) -> Any:
+    """Build the hidden-state reduction for one (n_embed, hc_mult) shape."""
+
+    def launch(x: Any, out: Any) -> None:
+        n = x.shape[0]
+        if n == 0:
+            return
+        _dflash_input_store_kernel[(n, triton.cdiv(n_embed, _STORE_BLOCK))](
+            x,
+            out,
+            out.stride(0),
+            N_EMBED=n_embed,
+            HC_MULT=hc_mult,
+            BLOCK=_STORE_BLOCK,
+            num_warps=4,
+        )
+
+    return launch

scripts/ab_cases.py

@@ -24,6 +24,8 @@ from deepinfer.model_executor.layers.tilelang_ops.compress import (
     compress_and_update_state_decode,
     compress_and_update_state_prefill,
 )
+from deepinfer.model_executor.layers.tilelang_ops.compute_seqlens_k_zos import compute_seqlens_k_zos
+from deepinfer.model_executor.layers.tilelang_ops.dflash import confidence_head_fused, dflash_input_store
 from deepinfer.model_executor.layers.tilelang_ops.flash_mla_utils import convert_kv_index
 from deepinfer.model_executor.layers.tilelang_ops.fused_sample import fused_sample
 from deepinfer.model_executor.layers.tilelang_ops.fused_transform import (
@@ -949,6 +951,76 @@ def _cases_fused_sample() -> list[tuple[str, bool, Callable[[], Any]]]:
     ]
 
 
+def _cases_compute_seqlens_k() -> list[tuple[str, bool, Callable[[], Any]]]:
+    """MTP key-length fixup: fresh rows keep `sk`, continuing rows take the accepted count."""
+
+    def make(n_seq: int, nextn: int, fresh_frac: float, n_prev: int | None = None):
+        def run() -> Any:
+            torch.manual_seed(0)
+            prev_n = n_prev if n_prev is not None else n_seq
+            seqlens_k = torch.randint(1, 200, (n_seq,), device="cuda", dtype=torch.int32)
+            seqlens_q = torch.randint(1, nextn + 2, (n_seq,), device="cuda", dtype=torch.int32)
+            # -1 marks a fresh request; anything else indexes the previous batch.
+            remap = torch.randint(0, prev_n, (n_seq,), device="cuda", dtype=torch.int32)
+            fresh = torch.rand(n_seq, device="cuda") < fresh_frac
+            remap = torch.where(fresh, torch.full_like(remap, -1), remap)
+            prev_accepted = torch.randint(0, nextn + 1, (prev_n,), device="cuda", dtype=torch.int32)
+            compute_seqlens_k_zos(seqlens_k, seqlens_q, remap, prev_accepted, nextn=nextn)
+            return seqlens_k
+
+        return run
+
+    return [
+        ("n256_nextn5_mixed", True, make(256, 5, 0.5)),
+        ("n256_nextn5_all_fresh", True, make(256, 5, 1.0)),
+        ("n256_nextn5_all_prev", True, make(256, 5, 0.0)),
+        ("n37_nextn3_mixed", True, make(37, 3, 0.5)),
+        ("n1_nextn1_fresh", True, make(1, 1, 1.0)),
+        ("n300_nextn5_small_prev", True, make(300, 5, 0.3, n_prev=7)),
+    ]
+
+
+def _cases_confidence_head() -> list[tuple[str, bool, Callable[[], Any]]]:
+    """MTP draft confidence: a dot product, a scaled clamp and a sigmoid."""
+
+    def make(n: int, n_embed: int, rank: int, temperature: float):
+        def run() -> Any:
+            torch.manual_seed(0)
+            hidden = torch.randn((n, n_embed), device="cuda", dtype=torch.bfloat16)
+            markov = torch.randn((n, rank), device="cuda", dtype=torch.bfloat16) if rank else None
+            weight = torch.randn((n_embed + rank, 1), device="cuda", dtype=torch.bfloat16)
+            return confidence_head_fused(hidden, weight, temperature, markov)
+
+        return run
+
+    return [
+        ("n64_e5120", False, make(64, 5120, 0, 1.0)),
+        ("n64_e5120_markov8", False, make(64, 5120, 8, 1.0)),
+        ("n64_e5120_markov8_cold", False, make(64, 5120, 8, 0.05)),
+        ("n1_e7168", False, make(1, 7168, 0, 1.0)),
+        ("n128_e512_markov4", False, make(128, 512, 4, 0.5)),
+    ]
+
+
+def _cases_dflash_input_store() -> list[tuple[str, bool, Callable[[], Any]]]:
+    """MTP draft input: reduce the hc_mult copies of each row to one."""
+
+    def make(n: int, hc_mult: int, n_embed: int):
+        def run() -> Any:
+            torch.manual_seed(0)
+            x = torch.randn((n, hc_mult, n_embed), device="cuda", dtype=torch.bfloat16)
+            return dflash_input_store(x)
+
+        return run
+
+    return [
+        ("n64_hc1_e5120", True, make(64, 1, 5120)),
+        ("n64_hc2_e5120", False, make(64, 2, 5120)),
+        ("n64_hc4_e512", False, make(64, 4, 512)),
+        ("n1_hc2_e7168", False, make(1, 2, 7168)),
+    ]
+
+
 CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
     "store_kv_cache_tl": _cases_store_kv_cache,
     "load_paged_kv_cache_tl": _cases_load_paged_kv_cache,
@@ -961,6 +1033,9 @@ CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
     "compress_and_update_state_prefill_tl": _cases_compress_prefill,
     "compress_and_update_state_decode_tl": _cases_compress_decode,
     "fused_sample_tl": _cases_fused_sample,
+    "compute_seqlens_k_zos_tl": _cases_compute_seqlens_k,
+    "confidence_head_fused_tl": _cases_confidence_head,
+    "dflash_input_store_tl": _cases_dflash_input_store,
     "copy_block_tables_tl": _cases_copy_block_tables,
     "gather_tl": _cases_gather,
     "gather_input_ids_zos_tl": _cases_gather_input_ids,