diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py b/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py
index ed75bce..b62c59f 100644
--- a/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py
+++ b/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py
@@ -78,3 +78,116 @@ def load_kv_index_tl(
         )
 
     return launch
+
+
+@triton.jit
+def _get_fused_dflash_get_input_kernel(
+    SLOT_MAPPING_SWA,
+    POSITIONS,
+    FTP,
+    LAST_VALID_OFFSET,
+    NEWX_SLOT_MAPPING_OUT,
+    POSITIONS_OUT,
+    INPUT_IDS_OUT,
+    VALID_TOKEN_MASK_OUT,
+    GATE_FORCE_RANDOM_OUT,
+    KV_INDEX_OUT,
+    kv_len,
+    window_size,
+    swa_block_size: tl.constexpr,
+    padded_kv_len: tl.constexpr,
+    has_valid_token_mask: tl.constexpr,
+    has_gate_force_random: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    req = tl.program_id(0)
+    bx = tl.program_id(1)
+    # The row stride of the packed output is the grid's second dimension, so it does not need to be
+    # a compile-time constant.
+    dflash_block_size = tl.num_programs(1)
+
+    ftp = tl.load(FTP + req)
+    # A padding request keeps its loads in bounds but writes -1 / 0 everywhere.
+    is_valid = ftp >= 0
+    idx = tl.maximum(ftp, 0)
+    so = tl.load(SLOT_MAPPING_SWA + idx) - tl.load(POSITIONS + idx) % swa_block_size
+    ap = tl.load(POSITIONS + idx) + tl.load(LAST_VALID_OFFSET + req)
+
+    out_offset = req * dflash_block_size + bx
+    tl.store(
+        NEWX_SLOT_MAPPING_OUT + out_offset,
+        tl.where(is_valid, (ap + 1 + bx) % swa_block_size + so, -1),
+    )
+    tl.store(POSITIONS_OUT + out_offset, tl.where(is_valid, ap + 1 + bx, 0))
+    tl.store(INPUT_IDS_OUT + out_offset, 0)
+    if has_valid_token_mask:
+        tl.store(VALID_TOKEN_MASK_OUT + out_offset, is_valid)
+    if has_gate_force_random:
+        tl.store(GATE_FORCE_RANDOM_OUT + out_offset, 0)
+
+    # Every draft slot of a request fills the same kv_index columns; the row differs.
+    start_pos = ap - (window_size - 1)
+    row = KV_INDEX_OUT + out_offset * padded_kv_len
+    # Both operands are compile-time constants, so this bounds a static loop. It must round up:
+    # `padded_kv_len` is a multiple of 64 but not necessarily of BLOCK, and floor division would
+    # leave the tail of the row unwritten.
+    for tile in range(0, (padded_kv_len + BLOCK - 1) // BLOCK):
+        cols = tile * BLOCK + tl.arange(0, BLOCK)
+        in_range = cols < padded_kv_len
+        value = start_pos + cols
+        keep = is_valid & in_range & (cols < kv_len) & (value >= 0)
+        tl.store(row + cols, tl.where(keep, value % swa_block_size + so, -1), mask=in_range)
+
+
+def get_fused_dflash_get_input_tl(
+    swa_block_size: int,
+    padded_kv_len: int,
+    has_valid_token_mask: bool,
+    has_gate_force_random: bool,
+) -> Any:
+    """Build the kernel that packs each request's draft-token inputs, one CTA per draft slot."""
+    block = 128
+
+    def launch(
+        slot_mapping_swa: torch.Tensor,
+        positions: torch.Tensor,
+        ftp: torch.Tensor,
+        last_valid_offset: torch.Tensor,
+        newx_slot_mapping_out: torch.Tensor,
+        positions_out: torch.Tensor,
+        input_ids_out: torch.Tensor,
+        valid_token_mask_out: torch.Tensor | None,
+        gate_force_random_out: torch.Tensor | None,
+        kv_index_out: torch.Tensor,
+        kv_len: int,
+        window_size: int,
+    ) -> None:
+        real_cnt, dflash_block_size = newx_slot_mapping_out.shape
+        if real_cnt == 0 or dflash_block_size == 0:
+            return
+        # The interface passes None for a disabled mask; Triton needs a real pointer even when the
+        # branch reading it is compiled out.
+        valid_mask = valid_token_mask_out if valid_token_mask_out is not None else input_ids_out
+        gate_out = gate_force_random_out if gate_force_random_out is not None else input_ids_out
+        _get_fused_dflash_get_input_kernel[(real_cnt, dflash_block_size)](
+            slot_mapping_swa,
+            positions,
+            ftp,
+            last_valid_offset,
+            newx_slot_mapping_out,
+            positions_out,
+            input_ids_out,
+            valid_mask,
+            gate_out,
+            kv_index_out,
+            kv_len,
+            window_size,
+            swa_block_size=swa_block_size,
+            padded_kv_len=padded_kv_len,
+            has_valid_token_mask=has_valid_token_mask,
+            has_gate_force_random=has_gate_force_random,
+            BLOCK=block,
+            num_warps=4,
+        )
+
+    return launch
diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py
index f0a2cfd..c00b47a 100644
--- a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py
+++ b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py
@@ -368,3 +368,457 @@ def copy_block_tables_tl(max_seqlen_q: int, max_num_blocks: int, BLOCK_SIZE: int
         )
 
     return launch
+
+
+@triton.jit
+def _get_accept_result_kernel(
+    RAW_ACCEPT,
+    GOLDEN,
+    DRAFT_POOL,
+    READ_ROWS,
+    CANDIDATE_TOKEN_IDS,
+    IS_VERIFY,
+    FTP,
+    ROW_TO_REQ,
+    TOKEN_IDS_OUT,
+    ACCEPT_OUT,
+    IS_FINAL_OUT,
+    total_q,
+    HAS_GOLDEN: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    rows = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    live = rows < total_q
+
+    accept = tl.load(RAW_ACCEPT + rows, mask=live, other=0) != 0
+    token = tl.load(CANDIDATE_TOKEN_IDS + rows, mask=live, other=0)
+    if HAS_GOLDEN:
+        # A golden token replaces the draft's verdict: the request already knows the token it
+        # wants there, so acceptance becomes "did the draft propose that same token".
+        golden = tl.load(GOLDEN + rows, mask=live, other=-1)
+        draft_row = tl.load(READ_ROWS + rows, mask=live, other=0)
+        draft_token = tl.load(DRAFT_POOL + draft_row, mask=live, other=0)
+        has_golden = golden != -1
+        accept = tl.where(has_golden, golden == draft_token, accept)
+        token = tl.where(has_golden, golden, token)
+
+    is_verify = tl.load(IS_VERIFY + rows, mask=live, other=0) != 0
+    req = tl.load(ROW_TO_REQ + rows, mask=live, other=0)
+    start = tl.load(FTP + req, mask=live, other=0)
+    span = tl.where(live, rows - start, 0)
+
+    # Prefix acceptance: a row survives only if every earlier draft row in its request both was a
+    # verify row and was accepted. Walking `start..row-1` per row is what the TileLang kernel does
+    # with a while loop; here the trip count is the widest span in the block, and lanes whose span
+    # is shorter hold their value on the extra iterations. MTP packs at most `num_steps + 1` rows
+    # per request, so that bound stays small.
+    alive = tl.full([BLOCK], 1, tl.int1)
+    max_span = tl.max(span, axis=0)
+    for offset in range(0, max_span):
+        prev = start + offset
+        active = live & (offset < span) & (prev < total_q)
+        prev_accept = tl.load(RAW_ACCEPT + prev, mask=active, other=0) != 0
+        if HAS_GOLDEN:
+            prev_golden = tl.load(GOLDEN + prev, mask=active, other=-1)
+            prev_draft_row = tl.load(READ_ROWS + prev, mask=active, other=0)
+            prev_draft_token = tl.load(DRAFT_POOL + prev_draft_row, mask=active, other=0)
+            prev_accept = tl.where(prev_golden != -1, prev_golden == prev_draft_token, prev_accept)
+        prev_verify = tl.load(IS_VERIFY + prev, mask=active, other=0) != 0
+        alive = alive & (~active | (prev_verify & prev_accept))
+
+    tl.store(TOKEN_IDS_OUT + rows, token, mask=live)
+    tl.store(ACCEPT_OUT + rows, alive & is_verify & accept, mask=live)
+    tl.store(IS_FINAL_OUT + rows, alive & (~is_verify | ~accept), mask=live)
+
+
+def get_accept_result_tl(has_golden: bool) -> Any:
+    """Build the kernel that turns raw draft verdicts into prefix-accepted tokens."""
+
+    def launch(
+        raw_accept: torch.Tensor,
+        golden: torch.Tensor,
+        draft_token_ids_pool: torch.Tensor,
+        read_rows: torch.Tensor,
+        candidate_token_ids: torch.Tensor,
+        is_verify: torch.Tensor,
+        ftp: torch.Tensor,
+        row_to_req: torch.Tensor,
+        token_ids_out: torch.Tensor,
+        accept_out: torch.Tensor,
+        is_final_out: torch.Tensor,
+    ) -> None:
+        total_q = raw_accept.shape[0]
+        if total_q == 0:
+            return
+        block = 128
+        _get_accept_result_kernel[(cdiv(total_q, block),)](
+            raw_accept,
+            golden,
+            draft_token_ids_pool,
+            read_rows,
+            candidate_token_ids,
+            is_verify,
+            ftp,
+            row_to_req,
+            token_ids_out,
+            accept_out,
+            is_final_out,
+            total_q,
+            HAS_GOLDEN=has_golden,
+            BLOCK=block,
+            num_warps=4,
+        )
+
+    return launch
+
+
+@triton.jit
+def _finalize_verification_zero_kernel(
+    NUM_GENERATED,
+    NUM_VERIFIED,
+    BATCH_CONFIDENCE_ERROR,
+    num_steps: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # Counters are accumulated with atomics in the main kernel, so they are zeroed in a separate
+    # launch rather than by a grid barrier that would disable the overlap PDL gives us.
+    steps = tl.arange(0, BLOCK)
+    if num_steps > 0:
+        live = steps < num_steps
+        tl.store(NUM_GENERATED + steps, tl.zeros([BLOCK], tl.int32), mask=live)
+        tl.store(NUM_VERIFIED + steps, tl.zeros([BLOCK], tl.int32), mask=live)
+    tl.store(BATCH_CONFIDENCE_ERROR, 0.0)
+
+
+@triton.jit
+def _finalize_verification_kernel(
+    EFFECTIVE_ACCEPT,
+    IS_FINAL_IN,
+    EFFECTIVE_TOKEN_IDS,
+    TOKEN_LOGPROBS,
+    SUM_SMALLER_PROBS,
+    TOP_LOGPROBS,
+    TOP_TOKEN_IDS,
+    FTP,
+    SEQLEN_Q,
+    ROW_TO_REQ,
+    IS_VERIFY,
+    ACCEPT_PROBS,
+    CONFIDENCE_POOL,
+    MTP_INFO_READ_ROWS,
+    STOP_TOKEN_IDS,
+    LAST_VALID_OFFSET,
+    NUM_GENERATED,
+    NUM_VERIFIED,
+    BATCH_CONFIDENCE_ERROR,
+    EMITTED_STOP_TOKEN,
+    SAMPLED_TOKEN_IDS,
+    SAMPLED_TOKEN_LOGPROBS,
+    SAMPLED_SUM_SMALLER_PROBS,
+    SAMPLED_TOP_LOGPROBS,
+    SAMPLED_TOP_TOKEN_IDS,
+    VERIFIED_TOKEN_IDS,
+    VERIFIED_TOKEN_LOGPROBS,
+    VERIFIED_SUM_SMALLER_PROBS,
+    VERIFIED_TOP_LOGPROBS,
+    VERIFIED_TOP_TOKEN_IDS,
+    total_q,
+    topk_n,
+    top_lp_row_stride,
+    top_ids_row_stride,
+    HAS_SUM_SMALLER_PROBS: tl.constexpr,
+    HAS_CONFIDENCE: tl.constexpr,
+    NUM_STEPS: tl.constexpr,
+    NUM_STOP_TOKENS: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+    live = rows < total_q
+
+    accepted = tl.load(EFFECTIVE_ACCEPT + rows, mask=live, other=0) != 0
+    is_final = tl.load(IS_FINAL_IN + rows, mask=live, other=0) != 0
+    is_verify = tl.load(IS_VERIFY + rows, mask=live, other=0) != 0
+    token = tl.load(EFFECTIVE_TOKEN_IDS + rows, mask=live, other=0)
+    logprob = tl.load(TOKEN_LOGPROBS + rows, mask=live, other=0.0)
+    req = tl.load(ROW_TO_REQ + rows, mask=live, other=0)
+    start = tl.load(FTP + req, mask=live, other=0)
+    step = rows - start
+
+    # Payload for rows that survive verification; rejected rows are recorded as -1 / 0 so a
+    # consumer can tell "no token here" from a genuine token 0.
+    tl.store(VERIFIED_TOKEN_IDS + rows, tl.where(accepted, token, -1), mask=live)
+    tl.store(VERIFIED_TOKEN_LOGPROBS + rows, tl.where(accepted, logprob, 0.0), mask=live)
+    if HAS_SUM_SMALLER_PROBS:
+        ssp = tl.load(SUM_SMALLER_PROBS + rows, mask=live, other=0.0)
+        tl.store(VERIFIED_SUM_SMALLER_PROBS + rows, tl.where(accepted, ssp, 0.0), mask=live)
+
+    if NUM_STEPS > 0:
+        tl.atomic_add(NUM_GENERATED + step, 1, mask=live & accepted & (step < NUM_STEPS))
+        verified_slot = tl.maximum(tl.minimum(step, NUM_STEPS - 1), 0)
+        tl.atomic_add(NUM_VERIFIED + verified_slot, 1, mask=live & is_verify)
+
+    if HAS_CONFIDENCE:
+        accept_prob = tl.load(ACCEPT_PROBS + rows, mask=live, other=0.0)
+        read_row = tl.load(MTP_INFO_READ_ROWS + rows, mask=live, other=0)
+        confidence = tl.load(CONFIDENCE_POOL + read_row, mask=live, other=0.0)
+        delta = accept_prob - confidence
+        delta = tl.where(delta < 0, -delta, delta)
+        # The counter is one scalar, so reduce within the program and issue a single atomic.
+        # Summing in a different order than the TileLang kernel does changes the last bits, which
+        # is why the comparison tolerance for float32 exists.
+        tl.atomic_add(BATCH_CONFIDENCE_ERROR, tl.sum(tl.where(live & is_verify, delta, 0.0)))
+
+    # Per-request payload comes from the request's final row, of which there is exactly one.
+    if HAS_CONFIDENCE:
+        # Did the emitted stream hit any stop token? Scan the rows this request will actually emit,
+        # which is its prefix up to and including the final row.
+        stop_hit = tl.zeros([BLOCK_ROWS], tl.int1)
+        for scan in tl.static_range(NUM_STEPS + 1):
+            scan_row = start + scan
+            in_prefix = scan <= step
+            scan_token = tl.load(EFFECTIVE_TOKEN_IDS + scan_row, mask=live & is_final & in_prefix, other=0)
+            for stop_idx in tl.static_range(NUM_STOP_TOKENS):
+                stop_token = tl.load(
+                    STOP_TOKEN_IDS + req * NUM_STOP_TOKENS + stop_idx,
+                    mask=live & is_final & in_prefix,
+                    other=-1,
+                )
+                stop_hit = stop_hit | (in_prefix & (scan_token == stop_token))
+        tl.store(EMITTED_STOP_TOKEN + req, stop_hit, mask=live & is_final)
+
+    tl.store(LAST_VALID_OFFSET + req, step, mask=live & is_final)
+    tl.store(SAMPLED_TOKEN_IDS + req, token, mask=live & is_final)
+    tl.store(SAMPLED_TOKEN_LOGPROBS + req, logprob, mask=live & is_final)
+    if HAS_SUM_SMALLER_PROBS:
+        tl.store(SAMPLED_SUM_SMALLER_PROBS + req, ssp, mask=live & is_final)
+
+    # Top-k payload: accepted rows keep theirs, the final row also publishes to the request slot.
+    for tile in range(0, tl.cdiv(topk_n, BLOCK_K)):
+        k = tile * BLOCK_K + tl.arange(0, BLOCK_K)
+        in_range = live[:, None] & (k[None, :] < topk_n)
+        keep = accepted[:, None] | is_final[:, None]
+        lp = tl.load(
+            TOP_LOGPROBS + rows[:, None] * top_lp_row_stride + k[None, :],
+            mask=in_range & keep,
+            other=0.0,
+        )
+        tok = tl.load(
+            TOP_TOKEN_IDS + rows[:, None] * top_ids_row_stride + k[None, :],
+            mask=in_range & keep,
+            other=0,
+        )
+        acc = accepted[:, None]
+        tl.store(
+            VERIFIED_TOP_LOGPROBS + rows[:, None] * topk_n + k[None, :],
+            tl.where(acc, lp, 0.0),
+            mask=in_range,
+        )
+        tl.store(
+            VERIFIED_TOP_TOKEN_IDS + rows[:, None] * topk_n + k[None, :],
+            tl.where(acc, tok, 0),
+            mask=in_range,
+        )
+        tl.store(
+            SAMPLED_TOP_LOGPROBS + req[:, None] * topk_n + k[None, :],
+            lp,
+            mask=in_range & is_final[:, None],
+        )
+        tl.store(
+            SAMPLED_TOP_TOKEN_IDS + req[:, None] * topk_n + k[None, :],
+            tok,
+            mask=in_range & is_final[:, None],
+        )
+
+
+def finalize_verification_tl(
+    has_sum_smaller_probs: bool,
+    has_confidence: bool,
+    num_steps: int,
+    num_stop_tokens: int,
+) -> Any:
+    """Build the kernel that finalizes accepted tokens and packs the per-request payload."""
+    block_rows = 16
+    block_k = 128
+
+    def launch(
+        effective_accept: torch.Tensor,
+        is_final_input: torch.Tensor,
+        effective_token_ids: torch.Tensor,
+        token_logprobs: torch.Tensor,
+        sum_smaller_probs: torch.Tensor,
+        top_logprobs: torch.Tensor,
+        top_token_ids: torch.Tensor,
+        ftp: torch.Tensor,
+        seqlen_q: torch.Tensor,
+        row_to_req: torch.Tensor,
+        is_verify: torch.Tensor,
+        accept_probs: torch.Tensor,
+        confidence_pool: torch.Tensor,
+        mtp_info_read_rows: torch.Tensor,
+        stop_token_ids: torch.Tensor,
+        last_valid_offset: torch.Tensor,
+        num_generated_tokens: torch.Tensor,
+        num_verified_tokens: torch.Tensor,
+        batch_confidence_error: torch.Tensor,
+        emitted_stop_token: torch.Tensor,
+        sampled_token_ids: torch.Tensor,
+        sampled_token_logprobs: torch.Tensor,
+        sampled_sum_smaller_probs: torch.Tensor,
+        sampled_top_logprobs: torch.Tensor,
+        sampled_top_token_ids: torch.Tensor,
+        verified_token_ids: torch.Tensor,
+        verified_token_logprobs: torch.Tensor,
+        verified_sum_smaller_probs: torch.Tensor,
+        verified_top_logprobs: torch.Tensor,
+        verified_top_token_ids: torch.Tensor,
+    ) -> None:
+        total_q = effective_accept.shape[0]
+        topk_n = top_logprobs.shape[-1]
+        _finalize_verification_zero_kernel[(1,)](
+            num_generated_tokens,
+            num_verified_tokens,
+            batch_confidence_error,
+            num_steps=num_steps,
+            BLOCK=next_power_of_2(max(num_steps, 1)),
+        )
+        if total_q == 0:
+            return
+        _finalize_verification_kernel[(cdiv(total_q, block_rows),)](
+            effective_accept,
+            is_final_input,
+            effective_token_ids,
+            token_logprobs,
+            sum_smaller_probs,
+            top_logprobs,
+            top_token_ids,
+            ftp,
+            seqlen_q,
+            row_to_req,
+            is_verify,
+            accept_probs,
+            confidence_pool,
+            mtp_info_read_rows,
+            stop_token_ids,
+            last_valid_offset,
+            num_generated_tokens,
+            num_verified_tokens,
+            batch_confidence_error,
+            emitted_stop_token,
+            sampled_token_ids,
+            sampled_token_logprobs,
+            sampled_sum_smaller_probs,
+            sampled_top_logprobs,
+            sampled_top_token_ids,
+            verified_token_ids,
+            verified_token_logprobs,
+            verified_sum_smaller_probs,
+            verified_top_logprobs,
+            verified_top_token_ids,
+            total_q,
+            topk_n,
+            top_logprobs.stride(0),
+            top_token_ids.stride(0),
+            HAS_SUM_SMALLER_PROBS=has_sum_smaller_probs,
+            HAS_CONFIDENCE=has_confidence,
+            NUM_STEPS=num_steps,
+            NUM_STOP_TOKENS=num_stop_tokens,
+            BLOCK_ROWS=block_rows,
+            BLOCK_K=block_k,
+            num_warps=8,
+        )
+
+    return launch
+
+
+@triton.jit
+def _remap_draft_topk_kernel(
+    DRAFT_TOPK_INDEX,
+    DRAFT_TOPK_LOGPROBS,
+    DRAFT_TOKEN_IDS_POOL,
+    TARGET_TOPK_INDEX,
+    READ_ROWS,
+    ALIGNED_LOGPROBS,
+    DRAFT_LOCAL_IDS,
+    IDENTITY_ROWS,
+    target_stride,
+    aligned_stride,
+    force_topk: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tblock = tl.program_id(1)
+    source_row = tl.load(READ_ROWS + row)
+    draft_token_id = tl.load(DRAFT_TOKEN_IDS_POOL + source_row).to(tl.uint32)
+
+    # One program owns a slice of this row's target indices and scans the whole draft slice.
+    tj = tblock * BLOCK_T + tl.arange(0, BLOCK_T)
+    t_valid = tj < force_topk
+    q_index = tl.load(TARGET_TOPK_INDEX + row * target_stride + tj, mask=t_valid, other=0).to(tl.uint32)
+
+    if tblock == 0:
+        tl.store(IDENTITY_ROWS + row, row)
+
+    # Where in the target slice the draft token itself sits. This is independent of the lookup
+    # below: it locates the draft token, not some other token's draft logprob.
+    is_draft = t_valid & (q_index == draft_token_id)
+    tl.atomic_max(DRAFT_LOCAL_IDS + row, tl.max(tl.where(is_draft, tj, -1), axis=0))
+
+    # The TileLang kernel builds a shared-memory lookup table over vocabulary-sized windows and
+    # reads it back. The same answer comes from matching each target index against the draft
+    # indices directly: a draft entry is unique per token id, so at most one candidate matches and
+    # taking the maximum over that block is well defined. This trades a table build for matched
+    # scans, which is only worth it because `force_topk` is at most a few thousand.
+    acc = tl.full([BLOCK_T], float("-inf"), tl.float32)
+    for dstart in range(0, force_topk, BLOCK_D):
+        dj = dstart + tl.arange(0, BLOCK_D)
+        d_valid = dj < force_topk
+        d_index = tl.load(DRAFT_TOPK_INDEX + source_row * force_topk + dj, mask=d_valid, other=0).to(tl.uint32)
+        d_value = tl.load(DRAFT_TOPK_LOGPROBS + source_row * force_topk + dj, mask=d_valid, other=0.0)
+        match = t_valid[:, None] & d_valid[None, :] & (q_index[:, None] == d_index[None, :])
+        acc = tl.maximum(acc, tl.max(tl.where(match, d_value, float("-inf")), axis=1))
+
+    tl.store(ALIGNED_LOGPROBS + row * aligned_stride + tj, acc, mask=t_valid)
+
+
+def remap_draft_topk_tl(force_topk: int, vocab_size: int) -> Any:
+    """Build the kernel that aligns the previous draft's top-k to the current target top-k."""
+    block_t = 128
+    block_d = 128
+
+    def launch(
+        draft_topk_index: torch.Tensor,
+        draft_topk_logprobs: torch.Tensor,
+        draft_token_ids_pool: torch.Tensor,
+        target_topk_index: torch.Tensor,
+        read_rows: torch.Tensor,
+        aligned_logprobs: torch.Tensor,
+        draft_local_ids: torch.Tensor,
+        identity_rows: torch.Tensor,
+    ) -> None:
+        total_q = target_topk_index.shape[0]
+        if total_q == 0 or force_topk == 0:
+            return
+        # `atomic_max` needs a starting value below every real position, and "not found" is -1.
+        draft_local_ids.fill_(-1)
+        _remap_draft_topk_kernel[(total_q, cdiv(force_topk, block_t))](
+            draft_topk_index,
+            draft_topk_logprobs,
+            draft_token_ids_pool,
+            target_topk_index,
+            read_rows,
+            aligned_logprobs,
+            draft_local_ids,
+            identity_rows,
+            target_topk_index.stride(0),
+            aligned_logprobs.stride(0),
+            force_topk=force_topk,
+            BLOCK_T=block_t,
+            BLOCK_D=block_d,
+            num_warps=4,
+        )
+
+    return launch
+
+    return launch