diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py
index 3151233d545b..751757a6d5f4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py
@@ -131,6 +131,13 @@ class FastPath:
         "multi-request decode FMHA over an over-subscribed KV split, 72..512 rows",
         "construct",
     ),
+    FastPath(
+        "mixed_step_gen_rows",
+        "full-attention generation rows of mixed steps with 9..64 generation requests on "
+        "multi_request_split_kv, context rows on the thop op",
+        "construct",
+        ("multi_request_split_kv",),
+    ),
     # Dense projections and prefetch.
     FastPath("dense_gemms", "dense projections on one named kernel per shape", "construct"),
     FastPath(
@@ -305,6 +312,13 @@ class FastPath:
         ("persistent_experts", "moe_prologue"),
     ),
     FastPath("replicated_embedding", "token embedding gathered from a replicated table", "load"),
+    FastPath(
+        "model_boundary",
+        "decode input boundary (gather, embed norm, residual, layer 0 pre-QKV norm) in one launch; "
+        "last-layer join storing BF16",
+        "load",
+        ("replicated_embedding",),
+    ),
     # Drafter.
     FastPath("drafter_block_conv", "drafter block convolutions as one fused kernel", "draft"),
     FastPath(
@@ -358,6 +372,13 @@ class FastPath:
         "draft",
         ("drafter_block_decode",),
     ),
+    FastPath(
+        "wide_drafter_hops",
+        "the drafter all-reduce hops of 9..56-row (sm107: 9..64-row) block decodes on a wide "
+        "one-shot transport",
+        "draft",
+        ("drafter_hops",),
+    ),
     FastPath(
         "drafter_pdl_chain",
         "drafter and draft head launches releasing their GEMMs early",
@@ -383,6 +404,12 @@ class FastPath:
         "draft",
         ("split_sampler",),
     ),
+    FastPath(
+        "vocab_parallel_head",
+        "the verify head of 16..1024-row decode steps on each rank's own vocabulary shard (no logits all-gather)",
+        "draft",
+        ("split_sampler",),
+    ),
     # Executor.
     FastPath("deferred_input_prep", "a decode step's input writes as one launch", "executor"),
     FastPath("sampler_store", "the sampler's store update and snapshots as one launch", "executor"),
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py
index ac97cd478efa..f159f2459851 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py
@@ -71,6 +71,8 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None:
         from .kernels.drafter_hops import hops as _drafter_hops
 
         _drafter_hops.install(getattr(shell.llm, "draft_model", None), shell._mapping)
+        if plan.enabled("wide_drafter_hops"):
+            _drafter_hops.install_wide(getattr(shell.llm, "draft_model", None), shell._mapping)
     if plan.enabled("drafter_pdl_chain"):
         # After the drafter_fixed_gemms GEMM methods and the drafter_hops hops:
         # the entry triggers rely on both.
@@ -82,7 +84,12 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None:
     if plan.enabled("split_sampler"):
         from .kernels.split_sampler import target_head as _target_head
 
-        _target_head.install(shell.llm, shell._mapping, plan.enabled("lm_head_inbox"))
+        _target_head.install(
+            shell.llm,
+            shell._mapping,
+            plan.enabled("lm_head_inbox"),
+            plan.enabled("vocab_parallel_head"),
+        )
         if plan.enabled("wide_split_sampler"):
             _target_head.widen(shell.llm)
     # Last: the arena compaction derives its worker class from whatever class every
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py
new file mode 100644
index 000000000000..75d2eeec934c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py
@@ -0,0 +1,571 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""A sequence-parallel layer's attention input all-gathered by the copy engines, QKV in row blocks.
+
+What it replaces. On a sequence-parallel context step every layer all-gathers its
+attention input (``[shard, 8192]`` BF16 per rank) with an NCCL ring kernel and then
+runs the QKV projection over all rows. The ring kernel needs SMs, so it runs next to
+nothing and the projection waits for its last byte.
+
+What runs instead, on steps of at least ``MIN_ROWS`` rows. Rank ``r`` gathers on a
+second stream ("copy stream"):
+
+* after the attention input is produced, the copy engines copy it into rows
+  ``[r * shard, (r + 1) * shard)`` of each peer's receive buffer, peers ``r+1, r+2, r+3``
+  in turn, and after each copy ``cuStreamWriteValue32`` sets the flag word that peer
+  keeps for source ``r`` (default flags: a fence before the write);
+* the QKV projection runs in row blocks on the main stream: rank ``r``'s own rows from
+  its own attention input first, then each peer's rows from the receive buffer after a
+  ``cuStreamWaitValue32`` on that peer's flag, peers in the order their copies arrive
+  (``r-1, r-2, r-3``: every sender copies to its next rank first). After the last
+  wait the flags are reset (fenced writes) and the main stream joins the copy stream.
+
+The receive buffers and flag words are one ``McastGPUBuffer`` per rank (the allocation
+of the projection-publish transport), used through its unicast peer mappings only.
+
+Ownership and progress (cuFlow):
+
+* One writer per region and layer: rank ``p``'s copy writes ``recv_r[p * shard...]``;
+  one reader: rank ``r``'s block GEMM for ``p``.
+* Publication: the copy completes, then the fenced flag write; acquire: the stream-side
+  wait, then the GEMM. The reader resets its flags after its last wait.
+* Visibility: the waits carry no ``CU_STREAM_WAIT_VALUE_FLUSH`` (the device reports
+  ``CAN_FLUSH_REMOTE_WRITES`` = 0). The copied rows are visible to the reader's GEMM
+  because the sender orders the copy before its fenced flag write on one stream and
+  the NVLink peer writes land in the reader's memory before the flag word does. This
+  rests on that ordering of the peer-write path, not on a documented acquire of the
+  wait; the unit and in-engine checks under graph replay are bit exact.
+* Consumers: on a step this takes, the view ``gather`` returns never holds this
+  rank's own rows (``project`` reads them from the attention input itself), so
+  ``project`` must stay its only reader; ``settle`` completes the view in place for
+  any other reader.
+* Reuse (one buffer for every layer): rank ``p`` writes ``recv_r`` again for layer
+  ``L + 1`` only after its own layer ``L`` passed the o_proj and experts
+  reduce-scatters, which need rank ``r``'s contributions, issued after ``r``'s layer
+  ``L`` block GEMMs and flag resets; between steps the forward's final all-gather does
+  the same. This needs one collective of all ranks between two gathers; every layer
+  has two.
+* Progress: a wait targets a copy its peer enqueued right after its own producer of the
+  same layer, which depends only on collectives this rank already passed. Copy
+  engines take no SM and the waits hold none, so nothing assumes co-residency. A peer
+  that never copies leaves the stream waiting without an error: ``start_watch``'s
+  thread sees a forward that began and did not end for ``stall_s`` seconds
+  (``first_stall_s`` before this rank ended its first one), reads this rank's flag
+  words on another stream, logs the sources that never arrived and aborts.
+
+Numerics: bitwise. The gathered rows are the same bytes, and each row block's GEMM is
+``torch.matmul`` with the projection's own weight, which computes every row as the
+whole-step GEMM does on the cuBLAS this was certified with (the in-engine checksum
+arm compares the hidden states).
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+from typing import Callable, List, Optional, Tuple
+
+import torch
+from cuda.bindings import driver as _cu
+from cuda.bindings import runtime as _cudart
+
+from tensorrt_llm._torch.distributed import allgather
+from tensorrt_llm._utils import prefer_pinned
+from tensorrt_llm.logger import logger
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:  # pragma: no cover - the image ships Triton
+    triton = None
+
+#: Seconds a sequence-parallel forward may stay begun and not ended before the watch
+#: aborts, once this rank has ended one (a served context step takes well under a second;
+#: a later step may still load or compile a kernel).
+STALL_S = 120.0
+#: The same bound while this rank has not yet ended any sequence-parallel forward: the
+#: first one (the warmup) compiles and loads the step's kernels, which on a cold cache
+#: takes longer than ``STALL_S``.
+FIRST_STALL_S = 900.0
+#: The fewest rows of a step whose attention input is gathered here: below it the
+#: flags' fixed per-flag cost eats the overlap (measured: next to no gain per layer at
+#: 2048 rows, a clear and growing gain from 4096 rows up).
+MIN_ROWS = 4096
+#: int32 words between two flag words (one 128-byte line each).
+FLAG_STRIDE = 32
+#: Seconds between two reads of the forward counters by the watch thread.
+WATCH_INTERVAL_S = 1.0
+
+_ALIGN = 2 << 20
+
+
+def _check(result):
+    err = result[0] if isinstance(result, tuple) else result
+    if err != _cu.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"[musespark_open] ce_gather: CUDA driver call failed: {err}")
+
+
+def _write_value(address: int, value: int) -> None:
+    """``cuStreamWriteValue32`` on the current stream with the default (fenced) flags."""
+    _check(
+        _cu.cuStreamWriteValue32(
+            _cu.CUstream(torch.cuda.current_stream().cuda_stream),
+            _cu.CUdeviceptr(address),
+            value,
+            _cu.CUstreamWriteValue_flags.CU_STREAM_WRITE_VALUE_DEFAULT,
+        )
+    )
+
+
+def _wait_value(address: int, value: int) -> None:
+    """``cuStreamWaitValue32`` (equal) on the current stream."""
+    _check(
+        _cu.cuStreamWaitValue32(
+            _cu.CUstream(torch.cuda.current_stream().cuda_stream),
+            _cu.CUdeviceptr(address),
+            value,
+            _cu.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ,
+        )
+    )
+
+
+class _HostWords:
+    """Pinned host int32 words a kernel updates at system scope and a host thread reads unsynchronized.
+
+    Triton reads ``data_ptr()`` and ``dtype`` of a pointer argument; the words' device
+    address is their host address (unified addressing).
+    """
+
+    dtype = torch.int32
+
+    def __init__(self, count: int) -> None:
+        # Without pinned memory there is no mapped device pointer: the error below makes
+        # ``CopyEngineGather.build`` refuse, and the NCCL all-gather runs instead.
+        self.host = torch.zeros(count, dtype=torch.int32, pin_memory=prefer_pinned())
+        status, pointer = _cudart.cudaHostGetDevicePointer(self.host.data_ptr(), 0)
+        if status != _cudart.cudaError_t.cudaSuccess:
+            raise RuntimeError(f"cudaHostGetDevicePointer of the forward counters: {status}")
+        self._pointer = int(pointer)
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+
+if triton is not None:
+
+    @triton.jit
+    def _count_kernel(words_ptr, INDEX: tl.constexpr):
+        """Add one to ``words[INDEX]`` (pinned host memory) at system scope."""
+        tl.atomic_add(words_ptr + INDEX, 1, sem="release", scope="sys")
+
+
+#: The forward counters' slots.
+_BEGUN, _ENDED = 0, 1
+
+
+def projection_refusal(linear) -> Optional[str]:
+    """Why ``linear`` (a layer's ``qkv_proj``) cannot run in row blocks here, or None."""
+    from tensorrt_llm._torch.modules.linear import UnquantizedLinearMethod
+
+    weight = getattr(linear, "weight", None)
+    if not isinstance(getattr(linear, "quant_method", None), UnquantizedLinearMethod):
+        return "the QKV projection is quantized"
+    if weight is None or weight.dtype != torch.bfloat16 or weight.dim() != 2:
+        return "the QKV projection's weight is not a 2-D BF16 matrix"
+    if getattr(linear, "bias", None) is not None:
+        return "the QKV projection has a bias"
+    if getattr(linear, "use_cute_dsl_bf16_gemm", False) or getattr(
+        linear, "use_custom_cublas_mm", False
+    ):
+        return "the QKV projection runs a custom GEMM"
+    if getattr(linear, "lora", None) is not None:
+        return "the QKV projection carries LoRA"
+    return None
+
+
+class CopyEngineGather:
+    """The copy-engine gather of every sequence-parallel layer of one model, and its row-block QKV.
+
+    Built (``build``, collective over the TP group) on the first eager prefill step, before
+    any CUDA graph is captured; until then, and on steps it does not take, ``gather``
+    is the NCCL all-gather the built-in layer runs.
+    """
+
+    def __init__(self, mapping, max_rows: int, hidden: int, comm=None) -> None:
+        self.mapping = mapping
+        self.rank = mapping.tp_rank
+        self.world = mapping.tp_size
+        self.hidden = int(hidden)
+        # A step of ``rows`` rows gathers ``world * ceil(rows / world)`` rows.
+        self.max_rows = -(-int(max_rows) // self.world) * self.world
+        self._comm = comm
+        self.built = False
+        self.refused: Optional[str] = None
+        self.allocated = 0
+        self._pending: Optional[Tuple[torch.Tensor, torch.Tensor, int, int]] = None
+        self._watch: Optional[threading.Thread] = None
+        self._watch_stop: Optional[threading.Event] = None
+        #: Set (instead of aborting) when the watch was started with ``abort=False`` (tests).
+        self.stalled: Optional[dict] = None
+        self._flag_snap: Optional[torch.Tensor] = None
+        self.gathers = 0
+
+    # ------------------------------------------------------------------ build
+
+    def build(self) -> None:
+        """Allocate the receive buffers and flags on every rank (collective); refuse alike on failure."""
+        from tensorrt_llm._torch.distributed.ops import (
+            _get_mnnvl_workspace_comm,
+            _make_mnnvl_mcast_buffer,
+            _mnnvl_workspace_all_succeeded,
+        )
+
+        if self.built:
+            return
+        self.built = True
+        comm = self._comm if self._comm is not None else _get_mnnvl_workspace_comm(self.mapping)
+        refusal = None
+        if torch.cuda.is_current_stream_capturing():
+            refusal = "the first prefill step is captured"
+        elif triton is None:
+            refusal = "Triton is not available"
+        if not _mnnvl_workspace_all_succeeded(comm, refusal is None):
+            self.refused = refusal or "another rank refused"
+            return
+        error: Optional[BaseException] = None
+        try:
+            recv_bytes = -(-self.max_rows * self.hidden * 2 // _ALIGN) * _ALIGN
+            flag_words = self.world * FLAG_STRIDE
+            total = recv_bytes + flag_words * 4
+            torch.cuda.synchronize()
+            free_before = torch.cuda.mem_get_info()[0]
+            self._handle = _make_mnnvl_mcast_buffer(
+                comm, total, self.mapping, self.mapping.is_multi_node()
+            )
+            self.recv: List[torch.Tensor] = [
+                self._handle.get_uc_buffer(p, (self.max_rows, self.hidden), torch.bfloat16, 0)
+                for p in range(self.world)
+            ]
+            self.flags = self._handle.get_uc_buffer(
+                self.rank, (flag_words,), torch.int32, recv_bytes // 4
+            )
+            self._flag_base = [
+                self._handle.get_uc_buffer(
+                    p, (flag_words,), torch.int32, recv_bytes // 4
+                ).data_ptr()
+                for p in range(self.world)
+            ]
+            self.flags.zero_()
+            self.copy_stream = torch.cuda.Stream()
+            self.watch_stream = torch.cuda.Stream()
+            self.counters = _HostWords(2)
+            # Launch both counter kernels and the flag read once now, while no stream waits:
+            # a kernel's first launch loads its module, and that load does not return while
+            # any stream of the device is blocked on a flag a peer has not yet written.
+            self._count(_BEGUN)
+            self._count(_ENDED)
+            self.flag_words()
+            torch.cuda.synchronize()
+            self.counters.host.zero_()
+            self.allocated = max(0, free_before - torch.cuda.mem_get_info()[0])
+        except Exception as exc:  # noqa: BLE001 - every failure must reach the agreement
+            error = exc
+        if not _mnnvl_workspace_all_succeeded(comm, error is None):
+            self.refused = f"the buffers could not be allocated: {error!r}"
+            return
+        logger.info(
+            f"[musespark_open] ce_gather: sequence-parallel steps of {MIN_ROWS}..{self.max_rows} rows "
+            f"gather the attention input on the copy engines ({total / 2**20:.1f} MiB requested, "
+            f"{self.allocated / 2**20:.1f} MiB of device memory on this rank)"
+        )
+
+    @property
+    def armed(self) -> bool:
+        return self.built and self.refused is None
+
+    def takes(self, rows: int) -> bool:
+        """Whether a step of ``rows`` rows gathers here (host facts, alike on every rank)."""
+        return self.armed and MIN_ROWS <= rows <= self.max_rows
+
+    # --------------------------------------------------------------- protocol
+
+    def _flag_address(self, owner: int, source: int) -> int:
+        return self._flag_base[owner] + 4 * source * FLAG_STRIDE
+
+    def arrival_order(self) -> List[int]:
+        """The sources in the order their copies reach this rank."""
+        return [(self.rank - k) % self.world for k in range(1, self.world)]
+
+    def gather(self, attn_input: torch.Tensor, rows: int) -> torch.Tensor:
+        """The step's attention input, ``[rows, hidden]``: the hook the built-in layer calls.
+
+        On a step this takes, the returned view's peer rows are complete only after
+        the row-block projection's waits, so it is handed to that projection only
+        (``project``); every other step gets the NCCL all-gather.
+        """
+        if not self.takes(rows):
+            return allgather(attn_input, self.mapping, dim=0)[:rows]
+        if self._pending is not None:
+            raise RuntimeError(
+                "[musespark_open] ce_gather: a gathered attention input was not projected"
+            )
+        shard = int(attn_input.shape[0])
+        if shard * self.world < rows or not attn_input.is_contiguous():
+            raise ValueError(
+                f"[musespark_open] ce_gather: a [{shard}, ...] shard cannot hold {rows} rows"
+            )
+        rank = self.rank
+        mine = slice(rank * shard, (rank + 1) * shard)
+        main = torch.cuda.current_stream()
+        self.copy_stream.wait_stream(main)
+        with torch.cuda.stream(self.copy_stream):
+            for k in range(1, self.world):
+                peer = (rank + k) % self.world
+                self.recv[peer][mine].copy_(attn_input)
+                _write_value(self._flag_address(peer, rank), 1)
+        view = self.recv[rank][:rows]
+        self._pending = (view, attn_input, shard, rows)
+        self.gathers += 1
+        return view
+
+    def is_pending(self, tensor: torch.Tensor) -> bool:
+        """Whether ``tensor`` is the view ``gather`` returned and nothing has projected yet."""
+        pending = self._pending
+        return (
+            pending is not None
+            and tensor.data_ptr() == pending[0].data_ptr()
+            and tensor.shape == pending[0].shape
+        )
+
+    def project(self, weight: torch.Tensor) -> torch.Tensor:
+        """``view @ weight.T`` for the pending view, in row blocks behind the flags."""
+        view, own, shard, rows = self._pending
+        self._pending = None
+        out = torch.empty((rows, weight.shape[0]), dtype=view.dtype, device=view.device)
+        w_t = weight.t()
+        rank = self.rank
+
+        def block(source: int) -> Optional[slice]:
+            lo, hi = source * shard, min((source + 1) * shard, rows)
+            return slice(lo, hi) if hi > lo else None
+
+        own_rows = block(rank)
+        if own_rows is not None:
+            torch.matmul(own[: own_rows.stop - own_rows.start], w_t, out=out[own_rows])
+        order = self.arrival_order()
+        for index, source in enumerate(order):
+            _wait_value(self._flag_address(rank, source), 1)
+            if index == len(order) - 1:
+                for done in order:
+                    _write_value(self._flag_address(rank, done), 0)
+            rows_of = block(source)
+            if rows_of is not None:
+                torch.matmul(self.recv[rank][rows_of], w_t, out=out[rows_of])
+        torch.cuda.current_stream().wait_stream(self.copy_stream)
+        return out
+
+    def settle(self) -> None:
+        """Complete the pending view in place (every flag waited and reset, own rows copied) and join."""
+        view, own, shard, rows = self._pending
+        self._pending = None
+        order = self.arrival_order()
+        for source in order:
+            _wait_value(self._flag_address(self.rank, source), 1)
+        for source in order:
+            _write_value(self._flag_address(self.rank, source), 0)
+        lo, hi = self.rank * shard, min((self.rank + 1) * shard, rows)
+        if hi > lo:
+            view[lo:hi].copy_(own[: hi - lo])
+        torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+    # ------------------------------------------------------------- the watch
+
+    def _count(self, index: int) -> None:
+        _count_kernel[(1,)](self.counters, INDEX=index, num_warps=1)
+
+    def begin(self) -> None:
+        """Count a sequence-parallel forward as begun (a one-thread kernel; captured with the step)."""
+        if self.armed:
+            self._count(_BEGUN)
+
+    def end(self) -> None:
+        """Count a sequence-parallel forward as ended."""
+        if self.armed:
+            self._count(_ENDED)
+
+    def flag_words(self, bound_s: float = 5.0) -> Optional[List[int]]:
+        """This rank's flag word per source, or None if not read within ``bound_s``.
+
+        One contiguous device-to-host copy of the flag block on the watch stream: a
+        copy-engine transfer into a pinned buffer kept from the first call, with no kernel
+        and no allocation, so it completes while the compute streams wait on a flag.
+        """
+        # Inference mode on every thread: the buffer is first made by ``build`` inside the
+        # engine's inference mode and later refilled from the watch's helper thread.
+        with torch.inference_mode():
+            snap = self._flag_snap
+            if snap is None:
+                snap = torch.empty(
+                    self.flags.numel(), dtype=torch.int32, pin_memory=prefer_pinned()
+                )
+                self._flag_snap = snap
+            with torch.cuda.stream(self.watch_stream):
+                snap.copy_(self.flags, non_blocking=True)
+                done = torch.cuda.Event()
+                done.record()
+        start = time.monotonic()
+        while not done.query():
+            if time.monotonic() - start > bound_s:
+                return None
+            time.sleep(0.001)
+        return [int(v) for v in snap[::FLAG_STRIDE]]
+
+    def _flag_words_bounded(self, bound_s: float = 5.0) -> Optional[List[int]]:
+        """``flag_words`` from a helper thread, abandoned after ``bound_s`` even if a CUDA call blocks."""
+        result: List[Optional[List[int]]] = []
+        reader = threading.Thread(
+            target=lambda: result.append(self.flag_words(bound_s)),
+            name="musespark_open.ce_gather_flags",
+            daemon=True,
+        )
+        reader.start()
+        reader.join(bound_s + 1.0)
+        return result[0] if result else None
+
+    def start_watch(
+        self,
+        stall_s: Optional[float] = None,
+        interval_s: float = WATCH_INTERVAL_S,
+        first_stall_s: Optional[float] = None,
+        abort: bool = True,
+        on_stall: Optional[Callable[[dict], None]] = None,
+    ) -> bool:
+        """Watch the forward counters from a host thread; returns whether a thread runs.
+
+        A forward that began and has not ended for ``stall_s`` seconds (default
+        ``STALL_S``) is waiting on a flag that
+        no peer set: the thread logs this rank's flag words (the sources still at 0 never
+        arrived) and aborts the process, or with ``abort=False`` records the report in
+        ``stalled`` and calls ``on_stall`` (tests). Until this rank has ended one forward
+        the bound is ``first_stall_s`` (default ``FIRST_STALL_S``), at least ``stall_s``.
+        """
+        if not self.armed or self._watch is not None:
+            return self._watch is not None
+        if stall_s is None:
+            stall_s = STALL_S
+        if first_stall_s is None:
+            first_stall_s = FIRST_STALL_S
+        first_stall_s = max(first_stall_s, stall_s)
+        stop = threading.Event()
+        device = torch.cuda.current_device()
+
+        def watch():
+            torch.cuda.set_device(device)
+            seen_end, since = None, None
+            while not stop.wait(interval_s):
+                begun = int(self.counters.host[_BEGUN])
+                ended = int(self.counters.host[_ENDED])
+                now = time.monotonic()
+                if begun <= ended or ended != seen_end:
+                    seen_end, since = ended, now
+                    continue
+                if now - since < (stall_s if ended > 0 else first_stall_s):
+                    continue
+                logger.error(
+                    f"[musespark_open] ce_gather: rank {self.rank} has waited {now - since:.1f} s in a "
+                    f"sequence-parallel forward (begun {begun}, ended {ended}); reading its flag words"
+                )
+                words = self._flag_words_bounded()
+                missing = (
+                    None if words is None else [s for s in self.arrival_order() if words[s] == 0]
+                )
+                report = dict(
+                    rank=self.rank,
+                    begun=begun,
+                    ended=ended,
+                    stalled_s=round(now - since, 3),
+                    flags=words,
+                    missing=missing,
+                )
+                logger.error(
+                    f"[musespark_open] ce_gather: rank {self.rank} has waited {now - since:.1f} s in a "
+                    f"sequence-parallel forward (begun {begun}, ended {ended}); flag words per source "
+                    f"{words}, sources that never arrived {missing}"
+                )
+                if abort:
+                    os.abort()
+                self.stalled = report
+                if on_stall is not None:
+                    on_stall(report)
+                return
+
+        self._watch_stop = stop
+        self._watch = threading.Thread(
+            target=watch, name="musespark_open.ce_gather_watch", daemon=True
+        )
+        self._watch.start()
+        return True
+
+    def stop_watch(self) -> None:
+        """Stop the watch thread, if one runs (tests)."""
+        if self._watch is not None:
+            self._watch_stop.set()
+            self._watch.join()
+            self._watch = None
+
+
+class CopyEngineGatherQKV:
+    """Mixin for a layer's ``self_attn.qkv_proj``: the pending gathered view runs in row blocks."""
+
+    _sp_ce_gather: Optional[CopyEngineGather] = None
+
+    def forward(self, input, *args, **kwargs):
+        gather = self._sp_ce_gather
+        if gather is not None and isinstance(input, torch.Tensor) and gather.is_pending(input):
+            if not args and not kwargs.get("lora_params"):
+                return gather.project(self.weight)
+            gather.settle()
+        return super().forward(input, *args, **kwargs)
+
+
+_CLASSES = {}
+
+
+def _with_mixin(base: type) -> type:
+    if issubclass(base, CopyEngineGatherQKV):
+        return base
+    if base not in _CLASSES:
+        _CLASSES[base] = type(
+            "CopyEngineGatherQKVOver" + base.__name__,
+            (CopyEngineGatherQKV, base),
+            {"__module__": __name__},
+        )
+    return _CLASSES[base]
+
+
+def install(model) -> Optional[CopyEngineGather]:
+    """Hook every layer's attention-input gather and return the gather, or None.
+
+    Changes nothing and returns None when a layer's projection cannot run in row blocks
+    (logged); the buffers are built later (``CopyEngineGather.build``).
+    """
+    layers = list(model.layers)
+    reasons = {layer.layer_idx: projection_refusal(layer.self_attn.qkv_proj) for layer in layers}
+    refused = {idx: why for idx, why in reasons.items() if why is not None}
+    if refused:
+        logger.warning(
+            f"[musespark_open] ce_gather: not armed: {sorted(set(refused.values()))} "
+            f"(layers {sorted(refused)}); the attention input is all-gathered by NCCL"
+        )
+        return None
+    config = model.model_config
+    gather = CopyEngineGather(
+        model._sp_mapping, config.max_num_tokens, config.pretrained_config.hidden_size
+    )
+    for layer in layers:
+        qkv = layer.self_attn.qkv_proj
+        qkv.__class__ = _with_mixin(type(qkv))
+        qkv._sp_ce_gather = gather
+        layer._sp_attn_input_gather = gather.gather
+    return gather
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py
index fbb1d93b1528..e3b1a3d083d0 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py
@@ -22,7 +22,14 @@
 bit for bit) and reduce-scatters the experts' partial sums in place of their
 all-reduce. Attention and experts still see every row of the step, so the generation
 rows of a mixed step are served as the context rows are. After the last layer the
-final norm runs on the shard and the step's rows are all-gathered.
+final norm runs on the shard and the step's rows are all-gathered. On sm103
+(``SHARD_PARTS_BY_SM``) each rank also routes its own rows and gathers their top-8 ids
+and weights instead of the router logits (``moe.shard_routing``), and the attention-side
+residual join writes the router's FP32 input next to the pre-FFN norm output
+(``router_input``); both are bitwise. There the attention input of steps of at least
+``ce_gather.MIN_ROWS`` rows is gathered by the copy engines and the QKV projection runs
+in row blocks behind per-peer flags (``ce_gather``); also bitwise. On sm107 the router
+logits and the attention input are all-gathered.
 
 DFlash taps the residual stream after its capture layers, and that stream is
 sharded here. For the forward the step's capture buffer
@@ -139,23 +146,52 @@ def refusal(model: MuseSparkOpenModel) -> Optional[str]:
     return None
 
 
-def install(model: MuseSparkOpenModel, sm: Optional[int] = None) -> "ContextSequenceParallel":
+#: The SM versions ``install`` wires, each with its own set of the path's parts
+#: (``SHARD_PARTS_BY_SM``).
+SHARD_PARTS_BY_SM = {
+    # sm103: the MoE re-class with routing on the shard, the router's FP32 input from the
+    # attention-side join, and the copy-engine gather of the attention input.
+    103: ("moe", "shard_routing", "router_input", "ce_gather"),
+    # sm107: the MoE re-class only; the router logits are gathered and every row of the
+    # step is routed, as the built-in arm does, and the attention input is all-gathered.
+    107: ("moe",),
+}
+
+
+def install(model: MuseSparkOpenModel, sm: Optional[int]) -> "ContextSequenceParallel":
     """The path's forward for ``model``, with every layer's MoE re-classed (``moe.install``).
 
-    ``sm`` is the SM version the model's composition was built for; it picks the eager
-    threshold (``EAGER_MIN_ROWS_BY_SM``).
+    ``sm`` is the SM version the model's composition was built for
+    (``dispatch.sm_of``), required: the parts differ per SM, so none is assumed. It
+    picks the eager threshold (``EAGER_MIN_ROWS_BY_SM``) and the parts installed
+    beside the MoE re-class (``SHARD_PARTS_BY_SM``).
 
     The MNNVL all-reduce workspaces are grown on the first prefill step
     (``ContextSequenceParallel.reserve``).
 
     Raises:
         ValueError: when ``refusal`` names a reason, a layer's MoE is not the built-in one,
-            or a layer's fused finalize all-reduce serves ``MIN_ROWS`` rows.
+            a layer's fused finalize all-reduce serves ``MIN_ROWS`` rows, or ``sm`` is
+            neither 103 nor 107 (None included).
     """
     from . import moe as _moe
 
+    parts = SHARD_PARTS_BY_SM.get(sm)
+    if parts is None:
+        got = "no SM version" if sm is None else f"sm{sm}"
+        raise ValueError(
+            f"[musespark_open] context_sequence_parallel is written for sm103 and sm107, got {got}"
+        )
     path = ContextSequenceParallel(model, sm)
-    _moe.install(model, min_rows=MIN_ROWS)
+    _moe.install(model, min_rows=MIN_ROWS, shard_routing="shard_routing" in parts)
+    if "router_input" in parts:
+        from . import router_input as _router_input
+
+        _router_input.install(model, path.step)
+    if "ce_gather" in parts:
+        from . import ce_gather as _ce_gather
+
+        path.ce_gather = _ce_gather.install(model)
     return path
 
 
@@ -223,11 +259,19 @@ class ContextSequenceParallel:
     """
 
     def __init__(self, model: MuseSparkOpenModel, sm: Optional[int] = None) -> None:
+        from .router_input import SequenceParallelStep
+
         reason = refusal(model)
         if reason is not None:
             raise ValueError(f"[musespark_open] context_sequence_parallel: {reason}")
         self.mapping = model._sp_mapping
         self.reserved = False
+        # The copy-engine gather of the attention input (``ce_gather.install``); None when
+        # a layer's QKV projection cannot run in row blocks.
+        self.ce_gather = None
+        # Set while ``forward`` runs the layers: the attention-side joins then also write
+        # the router's FP32 input (``router_input``).
+        self.step = SequenceParallelStep()
         #: The fewest rows of an eager step that runs sequence-parallel on this SM.
         self.eager_min_rows = EAGER_MIN_ROWS_BY_SM.get(sm, EAGER_MIN_ROWS)
 
@@ -247,6 +291,9 @@ def reserve(self, model: MuseSparkOpenModel) -> None:
                 "on an eager prefill step before any capture"
             )
         _reserve_all_reduce_workspaces(model)
+        if self.ce_gather is not None:
+            self.ce_gather.build()
+            self.ce_gather.start_watch()
         self.reserved = True
 
     def takes(self, attn_metadata: Any, rows: int, kwargs: Dict[str, Any]) -> bool:
@@ -293,24 +340,39 @@ def forward(
             embeds = embeds * core.m_emb
         # The FP32 residual stream of this rank's rows, entered once.
         hidden_states = embeds.to(torch.float32)
+        ce_gather = (
+            self.ce_gather if self.ce_gather is not None and self.ce_gather.takes(rows) else None
+        )
         with _shard_capture(kwargs.get("spec_metadata"), shard) as (capture, step_capture):
             layers = core.layers
             precomputed_attn_input = None
-            for index, decoder_layer in enumerate(layers):
-                next_norm = layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None
-                result = decoder_layer(
-                    position_ids=position_ids,
-                    hidden_states=hidden_states,
-                    attn_metadata=attn_metadata,
-                    precomputed_attn_input=precomputed_attn_input,
-                    next_layer_norm=next_norm,
-                    sp_num_tokens=rows,
-                    **kwargs,
-                )
-                if next_norm is None:
-                    hidden_states = result
-                else:
-                    hidden_states, precomputed_attn_input = result
+            self.step.active = True
+            if ce_gather is not None:
+                ce_gather.begin()
+            try:
+                for index, decoder_layer in enumerate(layers):
+                    next_norm = (
+                        layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None
+                    )
+                    result = decoder_layer(
+                        position_ids=position_ids,
+                        hidden_states=hidden_states,
+                        attn_metadata=attn_metadata,
+                        precomputed_attn_input=precomputed_attn_input,
+                        next_layer_norm=next_norm,
+                        sp_num_tokens=rows,
+                        **kwargs,
+                    )
+                    if next_norm is None:
+                        hidden_states = result
+                    else:
+                        hidden_states, precomputed_attn_input = result
+            finally:
+                self.step.active = False
+                # Also counted when a layer raised: the watch then sees the forward end once
+                # its queued work does, rather than a forward that never ends.
+                if ce_gather is not None:
+                    ce_gather.end()
         output = core.norm(hidden_states.to(core.dtype))
         if capture is None:
             return allgather(output, mapping, dim=0)[:rows]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py
index 70df6ccbd953..e586044a8e0e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py
@@ -30,6 +30,24 @@
 store rounds each FP32 sum to BF16 to nearest even, as the cast does, so the partials
 are the built-in arm's bit for bit and the FP32 partials are never written or read back.
 The built-in finalize still decides which kernel runs.
+
+Routing on the shard (``shard_routing``). The built-in arm gathers every rank's FP32
+router logits (1 KiB per row) and then routes every row of the step on every rank. The
+routing is per row, so where the built-in routing method would run its fused top-8
+kernel on the gathered logits, each rank runs that kernel on its own rows -- launched as
+the call over all the step's rows would be (``TORCH_ORDER`` and warps chosen from the
+step's row count) -- and the top-8 ids and FP32 weights (64 B per row) are gathered with
+the NVFP4 payload instead of the logits. The routing method's one-slot memo then holds
+the gathered ids and weights, keyed on a placeholder that stands for the logits (no
+values); the experts and the finalize read the memo, so both see the ids and weights
+the built-in arm computes, bit for bit.
+Every consumer of a shard-routed call's ``router_logits`` must read through that memo
+(``routing_method.apply``): the placeholder's values are uninitialized.
+
+The router's FP32 input. When the layer's attention-side residual join already wrote
+the FP32 widening of this call's input (``router_input.SequenceParallelAttnJoin``,
+``sp_router_input``), the router GEMM reads it instead of a fresh BF16 -> FP32 copy:
+the same values, the same GEMM.
 """
 
 from __future__ import annotations
@@ -43,13 +61,99 @@
 from tensorrt_llm._torch.models import modeling_musespark_open as _builtin
 from tensorrt_llm._torch.models.modeling_musespark_open import (
     MuseSparkOpenMoE,
+    MuseSparkOpenRoutingMethod,
     _pad_sequence_parallel_tokens,
+    _version_or_none,
 )
 from tensorrt_llm._torch.utils import Fp4QuantizedTensor
 from tensorrt_llm.logger import logger
 
 _CLASSES: Dict[Tuple[type, type], type] = {}
 
+#: The routed experts and the top-k the fused routing kernel serves.
+_EXPERTS, _TOP_K = 256, 8
+
+
+def shard_routing(
+    routing, logits: torch.Tensor, rows: int
+) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+    """The fused routing kernel's ``(selected, weights)`` for the shard's ``logits``, or None.
+
+    Launched with the arguments ``MuseSparkOpenRoutingMethod.apply`` chooses for a call
+    over all ``rows`` rows of the step (``TORCH_ORDER`` above 128 rows when the method
+    takes it, one warp from 512 rows), one program per row, so every row's ids and
+    weights are that call's. None where ``apply`` on the step's logits would not run the
+    fused kernel, or where its memo is off (the experts then could not read the
+    gathered routing).
+    """
+    bias = getattr(routing, "e_score_correction_bias", None)
+    if not (
+        isinstance(routing, MuseSparkOpenRoutingMethod)
+        and routing._memo_enabled
+        and rows > 0
+        and (rows <= 128 or routing._full_routing_all_batches)
+        and routing.top_k == _TOP_K
+        and routing.output_dtype == torch.float32
+        and logits.is_cuda
+        and logits.dtype == torch.float32
+        and logits.ndim == 2
+        and logits.shape[0] > 0
+        and logits.shape[1] == _EXPERTS
+        and logits.is_contiguous()
+        and bias is not None
+        and bias.is_cuda
+        and bias.dtype == torch.float32
+        and bias.shape == (_EXPERTS,)
+        and bias.is_contiguous()
+        and not torch.compiler.is_compiling()
+    ):
+        return None
+    shard = logits.shape[0]
+    selected = torch.empty((shard, _TOP_K), device=logits.device, dtype=torch.int32)
+    weights = torch.empty_like(selected, dtype=torch.float32)
+    _builtin._full_routing_kernel[(shard,)](
+        logits,
+        bias,
+        selected,
+        weights,
+        routing.renorm_eps,
+        EXPERTS=_EXPERTS,
+        TORCH_ORDER=routing._full_routing_torch_order and rows > 128,
+        num_warps=1 if rows >= 512 else 4,
+    )
+    return selected, weights
+
+
+def routed_placeholder(
+    routing, selected: torch.Tensor, weights: torch.Tensor, rows: int
+) -> torch.Tensor:
+    """A ``[rows, 256]`` FP32 stand-in for the step's logits, memoized to ``(selected, weights)``.
+
+    The stand-in is a zero-stride view of one element: it carries no values and is never
+    read, because every routing call on it returns the memo.
+    """
+    placeholder = selected.new_empty((1, 1), dtype=torch.float32).expand(rows, _EXPERTS)
+    routing._routing_memo = (placeholder, _version_or_none(placeholder), selected, weights)
+    return placeholder
+
+
+def _fp32_router_input_takes(moe, hidden_states: torch.Tensor) -> bool:
+    """Whether the built-in ``router_logits`` on ``hidden_states`` is ``router(hidden.float()) * m``.
+
+    That is its last arm: more than 128 rows (the small-batch kernels end there), no
+    BF16x3 router and no armed ``router_gemm`` call.
+    """
+    call = getattr(moe, "_router_gemm_call", None)
+    weight = getattr(moe.router, "weight", None)
+    return (
+        hidden_states.ndim == 2
+        and hidden_states.shape[0] > 128
+        and not getattr(moe, "_bf16x3_router_enabled", False)
+        and (call is None or not call.is_armed)
+        and weight is not None
+        and weight.dtype == torch.float32
+    )
+
 
 def _rows_of(tensor, rows: int, shape_tail: Tuple[int, ...]):
     """The first ``rows`` rows of a gathered ``[*, width]`` tensor, in the shard's own layout."""
@@ -144,6 +248,25 @@ class ContextSequenceParallelMoE:
     Acts on calls with ``sp_num_tokens`` only; every other call is the class behind it.
     """
 
+    #: ``(norm output, its FP32 widening)`` left by the layer's attention-side join for
+    #: the next sequence-parallel call (``router_input``); consumed or dropped by it.
+    sp_router_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    #: Whether a sequence-parallel call routes the shard's rows and gathers their ids and
+    #: weights (``shard_routing``); otherwise it gathers the router logits (``install``).
+    _sp_shard_routing: bool = True
+
+    def _sp_router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """The router logits of the shard's rows, from the join's FP32 rows when it left them."""
+        handed = self.sp_router_input
+        self.sp_router_input = None
+        if (
+            handed is not None
+            and handed[0] is hidden_states
+            and _fp32_router_input_takes(self, hidden_states)
+        ):
+            return self.router(handed[1]) * self.metap_multiplier
+        return self.router_logits(hidden_states)
+
     def forward(self, hidden_states, attn_metadata, *args, **kwargs):
         rows: Optional[int] = kwargs.get("sp_num_tokens")
         experts = self.experts
@@ -153,11 +276,13 @@ def forward(self, hidden_states, attn_metadata, *args, **kwargs):
             or getattr(experts, "comm", None) is not None
             or rows > getattr(experts, "moe_max_num_tokens", rows)
         ):
+            self.sp_router_input = None
             return super().forward(hidden_states, attn_metadata, *args, **kwargs)
         mapping = self.mapping
         shard = hidden_states.shape[0]
+        routing_method = experts.routing_method
         with self.nvtx("router"):
-            router_logits = self.router_logits(hidden_states)
+            router_logits = self._sp_router_logits(hidden_states)
         with self.nvtx("pool"):
             prepared = self.prepare_expert_input(self.pooler(hidden_states))
             if isinstance(prepared, Fp4QuantizedTensor):
@@ -169,21 +294,32 @@ def forward(self, hidden_states, attn_metadata, *args, **kwargs):
                         f"{tuple(scale.shape)} scales, swizzled {prepared.is_sf_swizzled}) "
                         f"for {shard} rows"
                     )
-                logits, fp4_rows, scale_rows = allgather(
-                    [router_logits, fp4.reshape(shard, -1), scale.reshape(shard, -1)],
-                    mapping,
-                    dim=0,
+                routed = (
+                    shard_routing(routing_method, router_logits, rows)
+                    if self._sp_shard_routing
+                    else None
                 )
+                head = [router_logits] if routed is None else list(routed)
+                gathered = allgather(
+                    [*head, fp4.reshape(shard, -1), scale.reshape(shard, -1)], mapping, dim=0
+                )
+                fp4_rows, scale_rows = gathered[-2:]
                 scale_tail = tuple(scale.shape[1:]) if scale.dim() > 1 else ()
                 pooled = Fp4QuantizedTensor(
                     _rows_of(fp4_rows, rows, tuple(fp4.shape[1:])),
                     scale_rows[:rows].reshape((rows, *scale_tail) if scale_tail else (-1,)),
                     is_sf_swizzled=False,
                 )
+                if routed is None:
+                    router_logits = gathered[0][:rows]
+                else:
+                    router_logits = routed_placeholder(
+                        routing_method, gathered[0][:rows], gathered[1][:rows], rows
+                    )
             else:
                 logits, pooled = allgather([router_logits, prepared], mapping, dim=0)
                 pooled = pooled[:rows]
-            router_logits = logits[:rows]
+                router_logits = logits[:rows]
         with self.nvtx("experts"):
             outputs = experts(
                 pooled,
@@ -201,7 +337,6 @@ def forward(self, hidden_states, attn_metadata, *args, **kwargs):
                 f"expanded_idx_to_permuted_idx) triple, got {type(outputs).__name__}"
             )
         gemm2_output, _, expanded_idx_to_permuted_idx = outputs
-        routing_method = experts.routing_method
         with self.nvtx("router"):
             _, token_final_scales = routing_method.apply(router_logits)
         routing_method.clear_routing_memo()
@@ -235,10 +370,13 @@ def _with_mixin(base: type) -> type:
     return _CLASSES[key]
 
 
-def install(model, min_rows: int) -> int:
+def install(model, min_rows: int, shard_routing: bool = True) -> int:
     """Re-class every layer's MoE with the mixin in front of its class; returns how many changed.
 
-    ``min_rows`` is the fewest rows of a sequence-parallel step. Runs after every
+    ``min_rows`` is the fewest rows of a sequence-parallel step. ``shard_routing`` has a
+    sequence-parallel call route the shard's rows and gather their top-8 ids and weights
+    (``shard_routing``); without it the call gathers the router logits and routes every
+    row of the step, as the built-in arm does. Runs after every
     construct-time re-class of the MoE (``PrefillTarget`` is built last in
     ``MuseSparkOpenTepModelingV2.adopt``), so no later fixed re-class drops the mixin.
 
@@ -265,4 +403,5 @@ def install(model, min_rows: int) -> int:
         if type(moe) is not mixed:
             moe.__class__ = mixed
             changed += 1
+        moe._sp_shard_routing = bool(shard_routing)
     return changed
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py
new file mode 100644
index 000000000000..2e7639f2ce33
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py
@@ -0,0 +1,226 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The attention-side residual join of a sequence-parallel layer also writes the router's FP32 input.
+
+On a sequence-parallel context step each rank's MoE routes its own rows: the built-in
+router reads the pre-FFN norm output (BF16, ``[shard, 8192]``) as FP32, so the step
+runs a standalone BF16 -> FP32 copy of it in front of the router GEMM (``router_logits``
+-> ``self.router(hidden_states.to(torch.float32))``). That norm output is written by
+the attention-side residual join (``MuseSparkOpenGatedResidual._add_after_rmsnorm_and_norm``,
+one Triton kernel), which holds every value in registers when it rounds it to BF16.
+
+``SequenceParallelAttnJoin`` re-classes a layer's ``add_attn`` so that, while a
+sequence-parallel forward runs (``SequenceParallelStep.active``), the join is this
+module's copy of the built-in kernel with one more store: the BF16-rounded norm output
+widened back to FP32, which is ``ffn_input.to(torch.float32)`` bit for bit. The pair is
+handed to the layer's MoE (``moe.sp_router_input``), whose sequence-parallel arm
+(``moe.ContextSequenceParallelMoE``) feeds the FP32 rows to the router GEMM when it is
+called with that very norm output. Every other call is the built-in join.
+
+Numerics: bitwise. The kernel body is the built-in one (same block, warps, loads,
+roundings and reduction), and the extra output is an exact widening of the BF16 value
+the built-in stores.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch.models import modeling_musespark_open as _builtin
+from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenGatedResidual
+
+#: Rows at or below which the built-in join may take its cluster boundary kernel instead
+#: of the Triton one; a sequence-parallel shard always has more.
+_CLUSTER_BOUNDARY_MAX_ROWS = 32
+
+
+@triton.jit
+def _join_with_router_input_kernel(
+    out_ptr,  # [T, HIDDEN] fp32, the new residual stream
+    residual_ptr,  # [T, HIDDEN] fp32, the old one
+    branch_ptr,  # [T, HIDDEN] branch output, activation dtype
+    gain_ptr,  # [HIDDEN] (1 + w) when HAS_GAIN, else unread
+    alpha_ptr,  # [HIDDEN] fp32
+    beta_ptr,  # [HIDDEN] fp32
+    eps,
+    next_out_ptr,  # [T, HIDDEN] activation dtype, the next norm's output
+    next_weight_ptr,
+    next_eps,
+    router_ptr,  # [T, HIDDEN] fp32, the next norm's output widened back from the activation dtype
+    HIDDEN: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    HAS_GAIN: tl.constexpr,
+    NORM_FP32_RESIDUAL: tl.constexpr,
+):
+    """The built-in ``_gated_residual_after_rmsnorm_kernel`` with ``FUSE_NEXT_NORM`` set, plus one store.
+
+    Statement for statement the built-in body up to the next norm's store; that store
+    rounds to the activation dtype explicitly (round to nearest even, the conversion the
+    built-in store applies) and the rounded value is also stored widened to FP32.
+    """
+    token = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_H)
+    mask = offs < HIDDEN
+    base = token * HIDDEN + offs
+
+    branch = tl.load(branch_ptr + base, mask=mask, other=0.0)
+    scaled = branch.to(tl.float32)
+    if HAS_GAIN:
+        gain = tl.load(gain_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        scaled = (scaled * gain).to(branch.dtype).to(tl.float32)
+    variance = tl.sum(scaled * scaled, axis=0) / HIDDEN
+    normed = (scaled * tl.rsqrt(variance + eps)).to(branch.dtype).to(tl.float32)
+
+    residual = tl.load(residual_ptr + base, mask=mask, other=0.0).to(tl.float32)
+    alpha = tl.load(alpha_ptr + offs, mask=mask, other=0.0)
+    beta = tl.load(beta_ptr + offs, mask=mask, other=0.0)
+    updated = alpha * residual + beta * normed
+    tl.store(out_ptr + base, updated, mask=mask)
+    if NORM_FP32_RESIDUAL:
+        next_input = updated
+    else:
+        next_input = updated.to(branch.dtype).to(tl.float32)
+    next_variance = tl.sum(next_input * next_input, axis=0) / HIDDEN
+    next_weight = tl.load(next_weight_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    next_output = next_input * tl.rsqrt(next_variance + next_eps)
+    next_output = next_output * (next_weight + 1.0)
+    rounded = next_output.to(next_out_ptr.dtype.element_ty)
+    tl.store(next_out_ptr + base, rounded, mask=mask)
+    tl.store(router_ptr + base, rounded.to(tl.float32), mask=mask)
+
+
+def join_with_router_input(
+    join: MuseSparkOpenGatedResidual,
+    residual: torch.Tensor,
+    branch: torch.Tensor,
+    eps: float,
+    next_norm,
+    gain: Optional[torch.Tensor] = None,
+) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """``(updated, next_input, next_input.float())`` from one launch, or None.
+
+    None wherever the built-in ``_add_after_rmsnorm_and_norm`` would not launch its
+    Triton kernel with this call's operands (its unfused or cluster-boundary arms); the
+    caller then runs the built-in join.
+    """
+    if not (
+        join._fused_enabled
+        and residual.is_cuda
+        and residual.dtype == torch.float32
+        and branch.dtype == torch.bfloat16
+        and residual.ndim == 2
+        and residual.shape == branch.shape
+        and residual.shape[0] > _CLUSTER_BOUNDARY_MAX_ROWS
+        and residual.is_contiguous()
+        and branch.is_contiguous()
+        and next_norm.use_gemma
+        and next_norm.weight.dtype == branch.dtype
+        and next_norm.weight.is_contiguous()
+        and (gain is None or (gain.dtype == branch.dtype and gain.is_contiguous()))
+        and not torch.compiler.is_compiling()
+    ):
+        return None
+    alpha, beta = join.coefficients()
+    updated = torch.empty_like(residual)
+    next_input = torch.empty_like(branch)
+    router_input = torch.empty_like(residual)
+    hidden_size = residual.shape[-1]
+    _join_with_router_input_kernel[(residual.shape[0],)](
+        updated,
+        residual,
+        branch,
+        gain if gain is not None else alpha,
+        alpha,
+        beta,
+        eps,
+        next_input,
+        next_norm.weight,
+        next_norm.variance_epsilon,
+        router_input,
+        HIDDEN=hidden_size,
+        BLOCK_H=triton.next_power_of_2(hidden_size),
+        HAS_GAIN=gain is not None,
+        NORM_FP32_RESIDUAL=join._preqkv_norm_fp32,
+        num_warps=_builtin._RESIDUAL_NUM_WARPS,
+    )
+    return updated, next_input, router_input
+
+
+class SequenceParallelStep:
+    """Whether a sequence-parallel forward is running (set by ``forward.ContextSequenceParallel``)."""
+
+    def __init__(self) -> None:
+        self.active = False
+
+
+class SequenceParallelAttnJoin:
+    """``MuseSparkOpenGatedResidual`` mixin for a layer's ``add_attn``.
+
+    While ``_sp_step.active``, a join that the built-in would run as its Triton kernel
+    runs ``join_with_router_input`` and hands ``(next_input, its FP32 widening)`` to the
+    layer's MoE (``sp_router_input``). Any other call is the built-in join.
+    """
+
+    _sp_step: Optional[SequenceParallelStep] = None
+    _sp_moe = None
+
+    def _add_after_rmsnorm_and_norm(self, residual, branch, eps, next_norm, gain=None):
+        step, moe = self._sp_step, self._sp_moe
+        if step is not None and step.active and moe is not None:
+            joined = join_with_router_input(self, residual, branch, eps, next_norm, gain)
+            if joined is not None:
+                updated, next_input, router_input = joined
+                moe.sp_router_input = (next_input, router_input)
+                return updated, next_input
+        return super()._add_after_rmsnorm_and_norm(residual, branch, eps, next_norm, gain)
+
+
+_CLASSES = {}
+
+
+def _with_mixin(base: type) -> type:
+    """``base`` with ``SequenceParallelAttnJoin`` in front, one class per base."""
+    if issubclass(base, SequenceParallelAttnJoin):
+        return base
+    if base not in _CLASSES:
+        _CLASSES[base] = type(
+            "SequenceParallelAttnJoinOver" + base.__name__,
+            (SequenceParallelAttnJoin, base),
+            {"__module__": __name__},
+        )
+    return _CLASSES[base]
+
+
+def install(model, step: SequenceParallelStep) -> int:
+    """Re-class every layer's ``add_attn`` and point it at ``step`` and the layer's MoE; returns how many.
+
+    Raises:
+        ValueError: if a layer's ``add_attn`` is not a built-in ``MuseSparkOpenGatedResidual``.
+    """
+    joins = [layer.add_attn for layer in model.layers]
+    wrong = [
+        layer.layer_idx
+        for layer, join in zip(model.layers, joins)
+        if not isinstance(join, MuseSparkOpenGatedResidual)
+    ]
+    if wrong:
+        raise ValueError(
+            "[musespark_open] context_sequence_parallel re-classes the built-in gated residual "
+            f"of every layer; layers {wrong} carry another class"
+        )
+    changed = 0
+    for layer, join in zip(model.layers, joins):
+        mixed = _with_mixin(type(join))
+        if type(join) is not mixed:
+            join.__class__ = mixed
+            changed += 1
+        # Plain attributes, not registered submodules: ``nn.Module.__setattr__`` would file
+        # the MoE under the join's children, where attribute lookup finds the class's
+        # ``_sp_moe = None`` first, and the module tree would hold the MoE twice.
+        object.__setattr__(join, "_sp_step", step)
+        object.__setattr__(join, "_sp_moe", layer.feed_forward)
+    return changed
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py
new file mode 100644
index 000000000000..d77863d03d36
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py
@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""decode_kv_splits on the generation rows of mixed steps.
+
+A mixed step (context rows and generation rows in one batch) runs ``PrefillTarget``,
+the built-in text-model forward. Each layer's built-in attention runs the context
+phase (flashinfer's trtllm-gen context FMHA) and then the generation phase of the
+batch's generation sequences (``run_generation`` with ``seq_offset = num_contexts``),
+which launches flashinfer's trtllm-gen decode FMHA over those sequences. flashinfer
+picks that call's KV split from the SM count exactly as on a decode step (``splits``'s
+module docstring), so a mixed step of 20 or more generation requests on a TEP4 rank
+runs the persistent one-CTA-per-request kernel on every full-attention layer, and its
+time is set by the longest of those requests.
+
+This mixin, re-classed in front of the full-attention layers' attention, arms the
+model's ``DecodeKvSplits`` for the duration of one attention call when the step has
+context rows and ``takes_generation`` accepts its generation requests (9 to 64). The
+split rule then acts on the generation phase's decode call exactly as on the decode
+steps ``DecodeTarget`` arms: same SM-count rule, same own counter buffer, same
+workspace check, same calls left alone. The context phase calls a different flashinfer
+entry point (``flashinfer.prefill``), which the stand-in passes through unchanged.
+
+Why at the attention call and not around ``PrefillTarget``'s forward: under the
+breakable prefill CUDA graph a replay never runs the model's forward; it replays the
+captured segments and re-runs the eager attention op between them, which reaches this
+module's ``_attn_impl`` with the step's metadata. Arming here applies the rule on every
+mixed step, captured or eager. Every read is host-side metadata of the step.
+
+Decode steps (no context rows) are left to ``DecodeTarget``'s own arming; steps of one
+to 8 generation requests and pure context steps run what they ran before.
+
+Numerics: those of ``decode_kv_splits`` on decode steps (the same trtllm-gen kernel
+family, multi-CTA KV with the in-kernel fp32 merge; only the KV partition of each
+request changes).
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Optional
+
+import torch
+
+from tensorrt_llm.logger import logger
+
+from .splits import MAX_ROWS, MIN_ROWS, ROWS_PER_REQUEST, DecodeKvSplits
+
+#: Generation requests of a mixed step the rule is armed for.
+MIN_REQUESTS = MIN_ROWS // ROWS_PER_REQUEST
+MAX_REQUESTS = MAX_ROWS // ROWS_PER_REQUEST
+
+
+def takes_generation(attn_metadata) -> bool:
+    """Whether a step is a mixed step whose generation rows the rule is armed for.
+
+    Context rows, and ``MIN_REQUESTS``..``MAX_REQUESTS`` generation sequences with beam
+    width 1. Host-side reads only.
+    """
+    num_contexts = getattr(attn_metadata, "num_contexts", 0)
+    num_generations = getattr(attn_metadata, "num_generations", 0)
+    return (
+        num_contexts > 0
+        and MIN_REQUESTS <= num_generations <= MAX_REQUESTS
+        and getattr(attn_metadata, "beam_width", 1) == 1
+    )
+
+
+class MixedStepKvSplitsAttention:
+    """Mixin re-classed in front of a full-attention layer's attention class.
+
+    A call on a step ``takes_generation`` accepts runs the class behind it with the model's
+    ``DecodeKvSplits`` armed; any other call is exactly the class behind it. Adds one plain
+    attribute (``_mixed_kv_splits``).
+    """
+
+    _mixed_kv_splits: Optional[DecodeKvSplits] = None
+
+    def _attn_impl(self, q, k, v, attn_metadata, *args, **kwargs):
+        rule = self._mixed_kv_splits
+        if rule is None or not takes_generation(attn_metadata):
+            return super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs)
+        split, kept = rule.split_calls, rule.kept_calls
+        with rule.armed():
+            result = super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs)
+        if rule.split_calls > split:
+            outcome = "split"
+        elif rule.kept_calls > kept:
+            outcome = "kept flashinfer's split"
+        else:
+            outcome = "no trtllm-gen decode call reached (another FMHA library ran)"
+        logger.info_once(
+            f"[musespark_open] decode_kv_splits: a mixed step's generation rows "
+            f"({attn_metadata.num_generations} requests, layer {self.layer_idx}): {outcome}",
+            key=f"musespark_open.decode_kv_splits.mixed.{outcome}",
+        )
+        return result
+
+
+@functools.lru_cache(maxsize=None)
+def _mixed_class(base: type) -> type:
+    """``base`` with ``MixedStepKvSplitsAttention`` in front, one class per base."""
+    return type(f"{base.__name__}MixedKv", (MixedStepKvSplitsAttention, base), {})
+
+
+def install(model: torch.nn.Module, rule: DecodeKvSplits) -> int:
+    """Re-class the full-attention layers' attention in front of whatever class they carry,
+    so their mixed-step attention calls arm ``rule``. Returns the number of layers re-classed.
+
+    Run after every other attention re-class: the mixin goes in front of all of them and
+    acts only on mixed steps, which none of the others intercept.
+    """
+    chosen = [layer.self_attn for layer in model.layers if not layer.self_attn.is_sliding]
+    for attention in chosen:
+        if not isinstance(attention, MixedStepKvSplitsAttention):
+            attention.__class__ = _mixed_class(type(attention))
+        attention._mixed_kv_splits = rule
+    logger.info(
+        f"[musespark_open] decode_kv_splits: the generation rows of mixed steps with "
+        f"{MIN_REQUESTS}..{MAX_REQUESTS} generation requests run the split rule on "
+        f"{len(chosen)} full-attention layers"
+    )
+    return len(chosen)
+
+
+__all__ = ["MixedStepKvSplitsAttention", "install", "takes_generation"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py
index c82710716294..306171d084b2 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py
@@ -164,6 +164,10 @@ def __init__(
         self.kv_heads = kv_heads
         self.q_heads = q_heads
         self.fallback: Optional[str] = None
+        #: Decode calls that reached the armed stand-in: how many were split, how many kept
+        #: flashinfer's own split.
+        self.split_calls = 0
+        self.kept_calls = 0
         #: The multi-CTA counters of every call the rule splits (zeroed once; the
         #: kernel resets what it used at the end of each launch).
         self.counter = torch.zeros(
@@ -250,7 +254,9 @@ def __getattr__(self, name: str):
     def trtllm_batch_decode_with_kv_cache(self, *args, **kwargs):
         presented = None if args else self._splits.presented_sm_count(kwargs)
         if presented is None:
+            self._splits.kept_calls += 1
             return self._real.trtllm_batch_decode_with_kv_cache(*args, **kwargs)
+        self._splits.split_calls += 1
         kwargs["multi_ctas_kv_counter_buffer"] = self._splits.counter
         decode = self._real
         device_sm_count = decode.get_device_sm_count
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py
index 0d9cc4b80b97..41b0068da411 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py
@@ -30,22 +30,48 @@
 full-vocab logits; for those the worker all-gathers the shard logits itself
 and runs the built-in selector, which is the built-in computation.
 
+A drafter that declares a logit transform (``DFlashForCausalLM.
+apply_draft_logit_transform``: a multiplier and a tanh softcap, in fp32) is
+transformed after the top-k instead of before it (``logit_transform.py``):
+stages 2 to 4 run on the bf16 shard with canonical keys, which rank as the
+transformed logits do, and the merge reads each candidate's transformed value
+from a table the drafter's own method computed. The built-in worker calls the
+transform right after the drafter's logits processor; the install puts a
+``_DeferredTransform`` on the drafter instance in its place, which hands back
+the very shard tensor ``ShardDraftLogitsProcessor`` just returned, untouched,
+and records it. ``ShardedDraftHeadDFlashWorker._refine_block_logits``, the
+next worker method the built-in forward calls on those logits, keeps them in
+bf16 only when this step takes the sharded selector, and applies the
+drafter's transform otherwise, so every other path (sampled drafts, d2t,
+another layout, any other tensor) gets the built-in transform bit for bit.
+With a transform installed, bf16 draft logits are untransformed ones: the
+transform always returns fp32. The tables hold only while the drafter's
+``logit_multiplier`` and ``logit_softcap`` are the ones they were built from;
+a step on which they differ takes the built-in path with the drafter's
+current transform. Each reason a step with a transform installed takes the
+built-in path is logged once, as a warning (``_warn_builtin``).
+
 Installed from the shell's ``load_draft_weights`` (the drafter's selector and
 the borrowed head exist only after the drafter loads), or from the shell's
 post-load hooks on load paths that never call it (``LoadFormat.DUMMY``, which
 never loads the drafter, so there the install refuses: no DFlash 2 drafter):
 the built-in DFlash worker is re-classed in place to
-``ShardedDraftHeadDFlashWorker`` and the drafter's logits processor is
-wrapped. Nothing else changes.
+``ShardedDraftHeadDFlashWorker``, the drafter's logits processor is wrapped
+and, with a logit transform, the drafter instance gets the deferring
+transform. Nothing else changes.
 
 The merge reproduces the tie order of the ``torch.topk`` build it was written
 against, which torch does not document. Before installing, ``tie_order_reason``
 runs ``torch.topk`` and the sharded path on tie-heavy logits of the deployed
-shape and refuses the install if they disagree.
+shape (with a transform, also on logits whose transformed values tie) and
+refuses the install if they disagree.
 """
 
 from __future__ import annotations
 
+import weakref
+from typing import Optional
+
 import torch
 from torch import nn
 
@@ -56,8 +82,10 @@
 from tensorrt_llm.logger import logger
 
 from ...ops.spec.draft_path_walk import DraftPathWalk
-from ...ops.spec.draft_shard_topk import CHUNK, NUM_WARPS, DraftShardTopK
+from ...ops.spec.draft_shard_topk import CHUNK, FINITE_KEYS, NUM_WARPS, DraftShardTopK
 from ...ops.spec.draft_topk_merge import SORT_SLOTS, DraftTopKMerge
+from . import logit_transform as _logit_transform
+from .logit_transform import DraftLogitTransform
 
 draft_shard_topk = DraftShardTopK()
 draft_topk_merge = DraftTopKMerge()
@@ -66,6 +94,24 @@
 #: Attribute of the zero-copy block-logits stand-in that carries the draft tokens.
 TOKENS_ATTR = "_draft_head_tokens"
 
+#: Worker attribute holding the drafter's ``DraftLogitTransform`` (absent without a transform).
+TRANSFORM_ATTR = "_draft_head_transform"
+#: Worker attribute: this step's logits were kept untransformed for the sharded selector.
+OWED_ATTR = "_draft_head_owed"
+#: Processor attribute: a weak reference to the shard tensor its last call returned.
+SHARD_ATTR = "_draft_head_shard"
+#: Drafter attribute: ``(data_ptr, numel)`` of the logits the deferring transform handed back.
+DEFERRED_ATTR = "_draft_head_deferred"
+#: Drafter attribute: why the deferring transform's last call applied the transform in full.
+DECLINED_ATTR = "_draft_head_declined"
+#: The drafter method the built-in worker calls on the draft logits.
+TRANSFORM_METHOD = "apply_draft_logit_transform"
+
+#: Why a step takes the built-in path: the drafter's transform attributes moved off the tables.
+STALE_TABLES = "the drafter's logit_multiplier or logit_softcap differs from the install's"
+#: Why a step takes the built-in path: its bf16 logits are not the shard the transform deferred.
+NOT_DEFERRED = "the draft logits are not the shard the deferring transform handed back"
+
 #: Largest shard the 16-bit column field of a shard key holds.
 MAX_SHARD = 65536
 
@@ -74,10 +120,37 @@
 #: of this target does (one request, K + 1 = 8 block positions).
 PROBE_ROWS = 8
 
+#: Rows of the install-time probe of many requests (24 requests of K = 7
+#: positions), run with a drafter logit transform.
+WIDE_PROBE_ROWS = 168
+
 #: Seed of the probe's logits, drawn on the host so no device RNG state moves.
 PROBE_SEED = 1234
 
 
+class DraftTransformStateError(RuntimeError):
+    """A step kept its draft logits untransformed and no selector applied the transform."""
+
+
+def _warn_builtin(reason: str) -> None:
+    """Warn, once per ``reason``, that a draft step with a transform installed takes the built-in path."""
+    logger.warning_once(
+        "[musespark_open] draft_head: a draft step takes the built-in path, with the drafter's "
+        f"logit transform applied before the top-k ({reason})",
+        key=f"musespark_open.draft_head.builtin.{reason}",
+    )
+
+
+def _builtin_transform(draft_model, logits: torch.Tensor) -> torch.Tensor:
+    """The drafter's class ``apply_draft_logit_transform`` on ``logits``, with its current attributes.
+
+    The call the built-in worker makes on a drafter without the deferring
+    transform: fp32 out when the drafter declares a transform, ``logits``
+    itself when it declares none.
+    """
+    return getattr(type(draft_model), TRANSFORM_METHOD)(draft_model, logits)
+
+
 def gathers_column_shards(lm_head: nn.Module) -> bool:
     """Whether ``lm_head`` is a column-parallel head whose forward is shard GEMM + all-gather.
 
@@ -112,6 +185,18 @@ def key_allreduce(mapping) -> nn.Module | None:
     return allreduce if allreduce.mnnvl_allreduce is not None else None
 
 
+def canon_kwargs(transform: Optional[DraftLogitTransform]) -> dict:
+    """``draft_shard_topk`` / ``draft_rank_topk`` keywords for ``transform`` (none without one)."""
+    if transform is None:
+        return {}
+    return {"canon": transform.canon, "canon_band": transform.band}
+
+
+def values_kwargs(transform: Optional[DraftLogitTransform]) -> dict:
+    """``draft_topk_merge`` keywords for ``transform`` (none without one)."""
+    return {} if transform is None else {"values": transform.values}
+
+
 def selector_draft_tokens(
     local_logits: torch.Tensor,
     block_hidden_states: torch.Tensor,
@@ -120,6 +205,8 @@ def selector_draft_tokens(
     mapping,
     allreduce: nn.Module | None = None,
     walk=None,
+    *,
+    transform: Optional[DraftLogitTransform],
 ):
     """Greedy DFlash 2 draft tokens ``[B, K]`` int32 from this rank's ``[B, K, shard]`` logits.
 
@@ -130,23 +217,32 @@ def selector_draft_tokens(
     (``merge_launch``) on the steps it takes. ``walk`` stands in for
     ``draft_path_walk`` (same arguments; its result is returned as is):
     the ``path_walk_tail`` fast path passes a walk that also writes the step's tail.
+
+    ``transform`` is the drafter's logit transform when ``local_logits`` are
+    the untransformed bf16 shard of a drafter that declares one, else None:
+    the keys are then canonical and the unary values the transformed ones, so
+    the tokens are the built-in ones on the transformed logits. Keyword-only
+    and required, so no call site drops it by accident.
     """
     requests, positions, shard = local_logits.shape
     local = local_logits.reshape(requests * positions, shard)
+    canon = canon_kwargs(transform)
     merge_launch = {}
     if allreduce is not None:
         takes = getattr(allreduce, "takes", None)
         if takes is not None and takes(local.shape[0]):
-            packed = allreduce.keys(local, selector.top_k, mapping.tp_rank, mapping.tp_size)
+            packed = allreduce.keys(
+                local, selector.top_k, mapping.tp_rank, mapping.tp_size, **canon
+            )
             merge_launch = allreduce.merge_launch(local.shape[0])
         else:
             packed = draft_shard_topk(
-                local, selector.top_k, CHUNK, NUM_WARPS, mapping.tp_rank, mapping.tp_size
+                local, selector.top_k, CHUNK, NUM_WARPS, mapping.tp_rank, mapping.tp_size, **canon
             )
         gathered = allreduce(packed)
     else:
         gathered = allgather(
-            draft_shard_topk(local, selector.top_k, CHUNK, NUM_WARPS), mapping, dim=0
+            draft_shard_topk(local, selector.top_k, CHUNK, NUM_WARPS, **canon), mapping, dim=0
         )
     # The built-in selector's projection call, operand for operand.
     gate = selector.hidden_projection(
@@ -164,6 +260,7 @@ def selector_draft_tokens(
         selector.top_k,
         mapping.tp_size,
         **merge_launch,
+        **values_kwargs(transform),
     )
     # dflash2_score_edges' einsum on tensors of its shapes and strides: the same kernel.
     edges = torch.einsum("blpr,blcr->blpc", predecessors, successors)
@@ -175,7 +272,9 @@ class ShardDraftLogitsProcessor(nn.Module):
 
     Only the call the DFlash worker makes (``return_context_logits=True`` on
     a head ``gathers_column_shards`` admits) is changed; anything else goes to
-    the wrapped processor.
+    the wrapped processor. Each call records a weak reference to the shard it
+    returned (``SHARD_ATTR``; cleared by every other call), which the
+    deferring drafter transform matches.
     """
 
     def __init__(self, inner: nn.Module):
@@ -183,19 +282,117 @@ def __init__(self, inner: nn.Module):
         self.inner = inner
 
     def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits=False):
+        self.__dict__.pop(SHARD_ATTR, None)
         if return_context_logits and gathers_column_shards(lm_head):
             fork = self.__dict__.get("_selector_fork")
             if fork is not None:
                 # The selector fork of the PDL chain (drafter_pdl_chain/chain.py): the
                 # selector's projection of these rows on a side stream, under the LM-head GEMM.
                 fork.launch(hidden_states)
-            return lm_head.apply_linear(hidden_states, lm_head.bias)
+            out = lm_head.apply_linear(hidden_states, lm_head.bias)
+            self.__dict__[SHARD_ATTR] = weakref.ref(out)
+            return out
         return self.inner(hidden_states, lm_head, attn_metadata, return_context_logits)
 
 
+class _DeferredTransform:
+    """The drafter's ``apply_draft_logit_transform`` with the draft head's transform installed.
+
+    Set on the drafter instance (``__dict__``), where it shadows the class
+    method. A call on the shard tensor the drafter's ``ShardDraftLogitsProcessor``
+    has just returned, while the drafter still declares the transform the
+    tables hold and the worker class still routes those logits through
+    ``ShardedDraftHeadDFlashWorker._refine_block_logits`` (``_defers_transform``),
+    returns that tensor untouched and records its ``(data_ptr, numel)`` on the
+    drafter (``DEFERRED_ATTR``). Every other call returns the drafter's own
+    transform with its current attributes, the built-in computation, records
+    why on the drafter (``DECLINED_ATTR``) and warns once per reason
+    (``_warn_builtin``). Holds weak references only.
+    """
+
+    def __init__(self, worker, draft_model):
+        self._worker = weakref.ref(worker)
+        self._model = weakref.ref(draft_model)
+
+    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
+        draft_model = self._model()
+        worker = self._worker()
+        if draft_model is None:
+            raise DraftTransformStateError("[musespark_open] draft_head: the drafter is gone")
+        transform = None if worker is None else worker.__dict__.get(TRANSFORM_ATTR)
+        draft_model.__dict__[DEFERRED_ATTR] = None
+        draft_model.__dict__[DECLINED_ATTR] = None
+        # The processor in place now, not the one of the install: a wrapper put
+        # around it later records nothing, and the transform then runs here.
+        processor = getattr(draft_model, "logits_processor", None)
+        shard = None
+        if isinstance(processor, ShardDraftLogitsProcessor):
+            shard = processor.__dict__.pop(SHARD_ATTR, None)
+        if transform is None:
+            return _builtin_transform(draft_model, logits)
+        if not isinstance(processor, ShardDraftLogitsProcessor):
+            reason = f"the drafter's logits processor is {type(processor).__name__}"
+        elif shard is None or shard() is not logits:
+            reason = "the transformed tensor is not the shard the logits processor returned"
+        elif logits.dtype != torch.bfloat16:
+            reason = f"the drafter's logits processor returned {logits.dtype} logits"
+        elif not transform.matches(draft_model):
+            reason = STALE_TABLES
+        elif not _defers_transform(type(worker)):
+            reason = (
+                f"the worker class {type(worker).__name__} overrides the draft head's "
+                "refinement or selector"
+            )
+        else:
+            draft_model.__dict__[DEFERRED_ATTR] = (logits.data_ptr(), logits.numel())
+            return logits
+        draft_model.__dict__[DECLINED_ATTR] = reason
+        _warn_builtin(reason)
+        return _builtin_transform(draft_model, logits)
+
+
+def _owner(cls: type, name: str, after: Optional[type] = None) -> Optional[type]:
+    """The class of ``cls``'s MRO (past ``after``, when given) that defines ``name``."""
+    mro = cls.__mro__
+    if after is not None:
+        mro = mro[mro.index(after) + 1 :]
+    return next((klass for klass in mro if name in vars(klass)), None)
+
+
+#: ``_defers_transform`` by worker class.
+_DEFERS: dict[type, bool] = {}
+
+
+def _defers_transform(cls: type) -> bool:
+    """Whether a worker of class ``cls`` takes deferred draft logits as this module expects.
+
+    ``ShardedDraftHeadDFlashWorker._refine_block_logits`` must be the one the
+    built-in forward calls, followed in the MRO by the built-in identity, and the
+    ``_apply_dflash2_selector`` it reaches must be one that applies the
+    transform after the top-k (a class marked ``_draft_transform_aware``).
+    """
+    known = _DEFERS.get(cls)
+    if known is None:
+        known = False
+        if issubclass(cls, ShardedDraftHeadDFlashWorker):
+            name = "_refine_block_logits"
+            selector = _owner(cls, "_apply_dflash2_selector")
+            known = (
+                _owner(cls, name) is ShardedDraftHeadDFlashWorker
+                and _owner(cls, name, after=ShardedDraftHeadDFlashWorker) is DFlashWorker
+                and selector is not None
+                and bool(vars(selector).get("_draft_transform_aware", False))
+            )
+        _DEFERS[cls] = known
+    return known
+
+
 class ShardedDraftHeadDFlashWorker(DFlashWorker):
     """The built-in DFlash worker with the draft head on vocab shards (fast path ``draft_head``)."""
 
+    #: Its ``_apply_dflash2_selector`` applies a deferred drafter transform (``_defers_transform``).
+    _draft_transform_aware = True
+
     def _draft_slot_ids(self, draft_model, num_gens, block_size, num_draft_tokens):
         """The built-in slot ids, built once per shape outside capture and reused.
 
@@ -226,27 +423,126 @@ def _pdl_head_rows(self, key, ids):
         rows = ranges[key]
         return ids if rows is None else _drafter_pdl.DraftRowRange(ids, *rows)
 
+    def _shard_reason(self, gen_logits, spec_metadata, full_vocab) -> Optional[str]:
+        """Why the sharded selector cannot take ``gen_logits`` on this step, or None."""
+        if gen_logits.shape[-1] == full_vocab:
+            return "full-vocab draft logits"
+        if gen_logits.dtype != torch.bfloat16:
+            return f"{gen_logits.dtype} draft logits"
+        if (
+            gen_logits.dim() != 3
+            or gen_logits.stride(-1) != 1
+            or gen_logits.shape[-1] * self.mapping.tp_size != full_vocab
+        ):
+            return "draft logits outside the [B, K, shard] layout"
+        if spec_metadata.wants_advanced_draft_sampling:
+            return "sampled drafts"
+        if self._d2t is not None:
+            return "a d2t draft vocab map"
+        return None
+
     def _shard_selector_applies(self, gen_logits, spec_metadata, full_vocab) -> bool:
-        mapping = self.mapping
-        return (
-            gen_logits.dim() == 3
-            and gen_logits.dtype == torch.bfloat16
-            and gen_logits.stride(-1) == 1
-            and gen_logits.shape[-1] * mapping.tp_size == full_vocab
-            and not spec_metadata.wants_advanced_draft_sampling
-            and self._d2t is None
-        )
+        return self._shard_reason(gen_logits, spec_metadata, full_vocab) is None
+
+    def _draft_head_takes(self, draft_model, gen_logits, spec_metadata, full_vocab) -> bool:
+        """Whether the sharded selector drafts this step.
+
+        ``_shard_selector_applies``, and, with a transform installed, tables
+        that still match ``draft_model``. A step it turns down with a transform
+        installed takes the built-in path, where bf16 logits get
+        ``_draft_head_transformed``, and warns once per reason.
+        """
+        reason = self._shard_reason(gen_logits, spec_metadata, full_vocab)
+        transform = self.__dict__.get(TRANSFORM_ATTR)
+        if transform is None:
+            return reason is None
+        if reason is None and not transform.matches(draft_model):
+            reason = STALE_TABLES
+        if reason is None:
+            return True
+        if gen_logits.dtype == torch.bfloat16:
+            # Logits transformed upstream (fp32) were warned for where that happened.
+            _warn_builtin(reason)
+        return False
+
+    def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata):
+        """The drafter's logit transform, unless this step's sharded selector applies it.
+
+        The deferring transform handed back the processor's bf16 shard
+        untransformed (``DEFERRED_ATTR``). These logits stay bf16 only when they
+        are that tensor (the reshape keeps its address and size), this step
+        takes the sharded selector and the drafter still declares the
+        transform the tables hold; the selector then owes the transform
+        (``OWED_ATTR``). Any other bf16 logits get the drafter's current
+        transform here, with a one-time warning per reason. Then the next
+        class's refinement, the built-in identity.
+        """
+        deferred = draft_model.__dict__.pop(DEFERRED_ATTR, None)
+        declined = draft_model.__dict__.pop(DECLINED_ATTR, None)
+        self.__dict__[OWED_ATTR] = False
+        transform = self.__dict__.get(TRANSFORM_ATTR)
+        if transform is not None and gen_logits.dtype == torch.bfloat16:
+            full_vocab = draft_model.candidate_selector.predecessor_codebook.shape[0]
+            # A declined call has applied the drafter's current transform already,
+            # which keeps bf16 logits only when that transform is the identity.
+            reason = declined
+            if reason is None and deferred != (gen_logits.data_ptr(), gen_logits.numel()):
+                reason = NOT_DEFERRED
+            if reason is None:
+                reason = self._shard_reason(gen_logits, spec_metadata, full_vocab)
+            if reason is None and not transform.matches(draft_model):
+                reason = STALE_TABLES
+            if reason is None:
+                self.__dict__[OWED_ATTR] = True
+            else:
+                _warn_builtin(reason)
+                gen_logits = self._draft_head_transformed(draft_model, gen_logits)
+        return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata)
+
+    def _draft_head_transform_for(self, draft_model, gen_logits) -> Optional[DraftLogitTransform]:
+        """The transform the sharded selector applies to ``gen_logits`` (bf16 means untransformed).
+
+        For a step ``_draft_head_takes`` admits, which checked that the tables
+        still match ``draft_model``.
+        """
+        self.__dict__[OWED_ATTR] = False
+        transform = self.__dict__.get(TRANSFORM_ATTR)
+        if transform is None or gen_logits.dtype != torch.bfloat16:
+            return None
+        if not transform.matches(draft_model):
+            raise DraftTransformStateError(
+                "[musespark_open] draft_head: the sharded selector was asked for tables that no "
+                f"longer hold ({STALE_TABLES})"
+            )
+        return transform
+
+    def _draft_head_transformed(self, draft_model, gen_logits) -> torch.Tensor:
+        """``gen_logits`` with the drafter's current transform applied once, for a built-in path.
+
+        With a transform installed, bf16 logits are untransformed or were
+        transformed by an identity: both get the drafter's class method with
+        its current attributes, which is what the built-in worker computes.
+        """
+        self.__dict__[OWED_ATTR] = False
+        if self.__dict__.get(TRANSFORM_ATTR) is None or gen_logits.dtype != torch.bfloat16:
+            return gen_logits
+        return _builtin_transform(draft_model, gen_logits)
 
     def _apply_dflash2_selector(
         self, draft_model, gen_logits, block_hidden_states, anchor_tokens, spec_metadata
     ):
         selector = draft_model.candidate_selector
         full_vocab = selector.predecessor_codebook.shape[0]
+        takes = self._draft_head_takes(draft_model, gen_logits, spec_metadata, full_vocab)
         if gen_logits.shape[-1] == full_vocab:
             return super()._apply_dflash2_selector(
-                draft_model, gen_logits, block_hidden_states, anchor_tokens, spec_metadata
+                draft_model,
+                self._draft_head_transformed(draft_model, gen_logits),
+                block_hidden_states,
+                anchor_tokens,
+                spec_metadata,
             )
-        if self._shard_selector_applies(gen_logits, spec_metadata, full_vocab):
+        if takes:
             tokens = selector_draft_tokens(
                 gen_logits,
                 block_hidden_states,
@@ -254,17 +550,27 @@ def _apply_dflash2_selector(
                 selector,
                 self.mapping,
                 self.__dict__.get("_draft_head_allreduce"),
+                transform=self._draft_head_transform_for(draft_model, gen_logits),
             )
             stand_in = tokens.unsqueeze(-1).expand(*tokens.shape, full_vocab)
             setattr(stand_in, TOKENS_ATTR, tokens)
             return stand_in
-        # The built-in path: the head's own all-gather and .float(), then the built-in selector.
-        full_logits = allgather(gen_logits, self.mapping, dim=-1).float()
+        # The built-in path: the transform, the head's own all-gather and .float(),
+        # then the built-in selector.
+        full_logits = allgather(
+            self._draft_head_transformed(draft_model, gen_logits), self.mapping, dim=-1
+        ).float()
         return super()._apply_dflash2_selector(
             draft_model, full_logits, block_hidden_states, anchor_tokens, spec_metadata
         )
 
     def sample_draft_tokens(self, logits, spec_metadata, batch_size, **kwargs):
+        if self.__dict__.get(OWED_ATTR):
+            self.__dict__[OWED_ATTR] = False
+            raise DraftTransformStateError(
+                "[musespark_open] draft_head: this step kept its draft logits untransformed for "
+                "the sharded selector, but no selector applied the drafter's logit transform"
+            )
         tokens = getattr(logits, TOKENS_ATTR, None)
         if tokens is not None:
             return tokens
@@ -300,6 +606,9 @@ def unsupported_reason(worker, draft_model) -> str | None:
             return f"selector {name} is not bf16"
     if selector.hidden_projection.weight.dtype != torch.bfloat16:
         return "selector hidden_projection is not bf16"
+    shadow = getattr(draft_model, "__dict__", {}).get(TRANSFORM_METHOD)
+    if shadow is not None and not isinstance(shadow, _DeferredTransform):
+        return f"the drafter's {TRANSFORM_METHOD} is replaced on the instance"
     return None
 
 
@@ -340,73 +649,214 @@ def place(row, values):
     return torch.stack([probe[r % len(probe)] for r in range(rows)]).to(torch.bfloat16)
 
 
+def transform_probe_logits(
+    vocab: int, shard: int, transform: DraftLogitTransform, rows: int = PROBE_ROWS
+) -> torch.Tensor:
+    """Host bf16 ``[rows, vocab]`` logits whose top-k under ``transform`` is decided by its ties.
+
+    Built from the transform's own tie groups (``logit_transform.tie_groups``),
+    so no threshold is assumed: the largest group above +0.0 (``A < B`` its
+    smallest and largest finite members) and the largest below -0.0 (``C <
+    D``). Rows, in turn: ten ``A`` at low ids and ten ``B`` at high ids above a
+    small grid; the same reversed; a row of the lower group's members with ten
+    ``C`` low and ten ``D`` high; five +inf among fifteen finite members of the
+    upper group; consecutive keys across the upper edge ``hi`` of the identity
+    band; members of the upper group on every shard's first and last columns;
+    a row of upper-group members only; consecutive keys across the lower edge
+    ``lo`` above lower-group members. A transform without ties gets the
+    largest finite values in their place, which tie nowhere.
+    """
+    g = torch.Generator().manual_seed(PROBE_SEED + 1)
+    lowest, highest = FINITE_KEYS
+    groups = _logit_transform.tie_groups(transform)
+
+    def largest(candidates, fallback):
+        return max(candidates, key=lambda grp: grp[1] - grp[0], default=fallback)
+
+    upper = largest([grp for grp in groups if grp[0] >= 0x8000], (highest - 255, highest))
+    lower = largest([grp for grp in groups if grp[1] <= 0x7FFF], (lowest, lowest + 255))
+
+    def finite(group):
+        return torch.arange(max(group[0], lowest + 1), min(group[1], highest - 1) + 1)
+
+    def members(group, n, with_inf=False):
+        keys = torch.arange(group[0], group[1] + 1) if with_inf else finite(group)
+        return keys[torch.randint(0, keys.numel(), (n,), generator=g)]
+
+    def grid(levels):
+        return (torch.randint(0, levels, (vocab,), generator=g).float() * 0.25 - 3.0).to(
+            torch.bfloat16
+        )
+
+    def place(row, columns, keys):
+        row[columns] = _logit_transform.bf16_of_key16(keys)
+        return row
+
+    def ends(group):
+        keys = finite(group)
+        return int(keys[0]), int(keys[-1])
+
+    low = torch.randperm(vocab // 4, generator=g)[:10]
+    high = vocab - 1 - torch.randperm(vocab // 4, generator=g)[:10]
+    a, b = ends(upper)
+    c, d = ends(lower)
+    ten = torch.ones(10, dtype=torch.int64)
+    top = min(transform.hi, highest)
+    bottom = max(transform.lo, lowest)
+    upper_ladder = torch.arange(max(top - 23, 0x8000), min(top + 8, highest) + 1)
+    lower_ladder = torch.arange(max(bottom - 8, lowest), min(bottom + 23, 0x7FFF) + 1)
+    edges = torch.tensor(
+        [c for r in range(vocab // shard) for c in (r * shard, (r + 1) * shard - 1)]
+    )
+    mixed = torch.randperm(vocab, generator=g)[:20]
+    probe = [
+        place(place(grid(8), low, a * ten), high, b * ten),
+        place(place(grid(8), low, b * ten), high, a * ten),
+        place(
+            place(place(grid(8), torch.arange(vocab), members(lower, vocab, True)), low, c * ten),
+            high,
+            d * ten,
+        ),
+        place(place(grid(8), mixed[:5], torch.full((5,), highest)), mixed[5:], members(upper, 15)),
+        place(grid(8), torch.randperm(vocab, generator=g)[: upper_ladder.numel()], upper_ladder),
+        place(grid(4), edges, members(upper, edges.numel())),
+        place(grid(8), torch.arange(vocab), members(upper, vocab, True)),
+        place(
+            place(grid(8), torch.arange(vocab), members(lower, vocab, True)),
+            torch.randperm(vocab, generator=g)[: lower_ladder.numel()],
+            lower_ladder,
+        ),
+    ]
+    return torch.stack([probe[r % len(probe)] for r in range(rows)])
+
+
+def tie_order_probes(
+    vocab: int, shard: int, transform: Optional[DraftLogitTransform], device
+) -> list[torch.Tensor]:
+    """The install's probe logits on ``device``, bf16 ``[rows, vocab]`` each.
+
+    ``tie_order_probe_logits``; with a transform also ``transform_probe_logits``
+    and, at ``WIDE_PROBE_ROWS`` rows, both repeated.
+    """
+    probes = [tie_order_probe_logits(vocab, shard).to(device)]
+    if transform is not None:
+        ties = transform_probe_logits(vocab, shard, transform).to(device)
+        both = torch.cat([probes[0], ties])
+        repeats = -(-WIDE_PROBE_ROWS // both.shape[0])
+        probes += [ties, both.repeat(repeats, 1)[:WIDE_PROBE_ROWS].contiguous()]
+    return probes
+
+
+def probe_reference(full: torch.Tensor, top_k: int, transform, topk=torch.topk):
+    """``topk`` as the built-in selector runs it on ``full`` (bf16 ``[rows, vocab]``).
+
+    On ``[1, rows, vocab]`` fp32: ``full.float()``, or the drafter's transform of it.
+    """
+    rows, vocab = full.shape
+    logits = full.float() if transform is None else transform.real(full)
+    return topk(logits.view(1, rows, vocab), top_k, dim=-1)
+
+
 def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool:
     """Bitwise equality of fp32 tensors, NaN payloads aside."""
-    nan = a.isnan()
-    if not torch.equal(nan, b.isnan()):
-        return False
-    return torch.equal(
-        torch.where(nan, 0.0, a).view(torch.int32), torch.where(nan, 0.0, b).view(torch.int32)
-    )
+    return _logit_transform.same_bits(a, b)
 
 
-def tie_order_reason(selector: nn.Module, num_ranks: int, topk=torch.topk) -> str | None:
+def tie_order_reason(
+    selector: nn.Module,
+    num_ranks: int,
+    topk=torch.topk,
+    transform: Optional[DraftLogitTransform] = None,
+) -> str | None:
     """Why the sharded candidates would differ from this torch's ``torch.topk``, or None.
 
     The merge copies the tie order of ``torch.topk`` (selection, arrangement
     and bitonic network), which torch does not document. This runs ``topk``
     as the built-in selector does -- ``[1, rows, vocab]`` fp32 -- on
-    ``tie_order_probe_logits``, and the sharded path on the same logits split
+    ``tie_order_probes``, and the sharded path on the same logits split
     into ``num_ranks`` shards with the keys exchanged in both layouts
     (concatenated int32 keys and summed fp32 halves, which is what the
     all-gather and the all-reduce return), and compares ids and values bit for
-    bit. ``topk`` is a parameter so a test can stand in another order.
+    bit. With ``transform`` the reference transforms the logits first and the
+    sharded path runs on canonical keys with the merge's values table. ``topk``
+    is a parameter so a test can stand in another order.
     """
     predecessor_codebook = selector.predecessor_codebook
     vocab, rank_dim = predecessor_codebook.shape
     shard = vocab // num_ranks
     top_k = selector.top_k
     device = predecessor_codebook.device
+    canon = canon_kwargs(transform)
     with torch.cuda.device(device):
-        full = tie_order_probe_logits(vocab, shard).to(device)
-        rows = full.shape[0]
-        ref_values, ref_ids = topk(full.float().view(1, rows, vocab), top_k, dim=-1)
-        shards = [full[:, r * shard : (r + 1) * shard].contiguous() for r in range(num_ranks)]
-        layouts = {
-            "int32 all-gather": torch.cat(
-                [draft_shard_topk(s, top_k, CHUNK, NUM_WARPS) for s in shards], dim=0
-            ),
-            "fp32 all-reduce": torch.stack(
-                [
-                    draft_shard_topk(s, top_k, CHUNK, NUM_WARPS, r, num_ranks)
-                    for r, s in enumerate(shards)
-                ]
-            ).sum(dim=0),
-        }
-        gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device)
-        anchor = torch.zeros(1, dtype=torch.int64, device=device)
-        for layout, keys in layouts.items():
-            ids, values, _, _ = draft_topk_merge(
-                keys,
-                shard,
-                1,
-                rows,
-                anchor,
-                gate,
-                predecessor_codebook,
-                selector.successor_codebook,
-                top_k,
-                num_ranks,
-            )
-            if not (torch.equal(ids, ref_ids) and _same_bits(values, ref_values)):
-                differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist()
-                return (
-                    f"torch {torch.__version__} orders top-k ties differently from the "
-                    f"merge ({layout} keys; probe rows {differ} differ)"
+        for probe, full in enumerate(tie_order_probes(vocab, shard, transform, device)):
+            rows = full.shape[0]
+            ref_values, ref_ids = probe_reference(full, top_k, transform, topk)
+            shards = [full[:, r * shard : (r + 1) * shard].contiguous() for r in range(num_ranks)]
+            layouts = {
+                "int32 all-gather": torch.cat(
+                    [draft_shard_topk(s, top_k, CHUNK, NUM_WARPS, **canon) for s in shards], dim=0
+                ),
+                "fp32 all-reduce": torch.stack(
+                    [
+                        draft_shard_topk(s, top_k, CHUNK, NUM_WARPS, r, num_ranks, **canon)
+                        for r, s in enumerate(shards)
+                    ]
+                ).sum(dim=0),
+            }
+            gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device)
+            anchor = torch.zeros(1, dtype=torch.int64, device=device)
+            for layout, keys in layouts.items():
+                ids, values, _, _ = draft_topk_merge(
+                    keys,
+                    shard,
+                    1,
+                    rows,
+                    anchor,
+                    gate,
+                    predecessor_codebook,
+                    selector.successor_codebook,
+                    top_k,
+                    num_ranks,
+                    **values_kwargs(transform),
                 )
+                if not (torch.equal(ids, ref_ids) and _same_bits(values, ref_values)):
+                    differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist()
+                    return (
+                        f"torch {torch.__version__} orders top-k ties differently from the "
+                        f"merge ({layout} keys; probe {probe} rows {differ} differ)"
+                    )
     return None
 
 
+def agreed_reason(mapping, reason: str | None, transform: Optional[DraftLogitTransform]):
+    """``reason`` as every rank of ``mapping``'s TP group agrees, tables compared.
+
+    Collective over the TP group (``dispatch.agreed_union`` and
+    ``dispatch.agreed_refusal``): every rank calls it at the same point of the
+    install, whatever its own reason. Ranks whose transform tables differ
+    (``logit_transform.digest``) refuse, since each rank's merge reads its own
+    tables and differing ones would draft different tokens per rank.
+    """
+    from ... import dispatch as _tp_agreement
+
+    mine = _logit_transform.digest(transform)
+    seen = _tp_agreement.agreed_union(mapping, [mine])
+    if reason is None and seen != [mine]:
+        reason = f"the TP ranks differ in their drafter logit transform ({'; '.join(seen)})"
+    return _tp_agreement.agreed_refusal(mapping, reason)
+
+
+def agreed_refusal(mapping, reason: str | None) -> str | None:
+    """``reason`` as every rank of ``mapping``'s TP group agrees (``dispatch.agreed_refusal``).
+
+    Collective over the TP group: every rank calls it at the same point of the
+    install, whatever its own reason.
+    """
+    from ... import dispatch as _tp_agreement
+
+    return _tp_agreement.agreed_refusal(mapping, reason)
+
+
 def installed(llm: nn.Module) -> bool:
     """Whether ``install`` has put the draft head on ``llm``'s DFlash worker."""
     return isinstance(getattr(llm, "spec_worker", None), ShardedDraftHeadDFlashWorker)
@@ -417,15 +867,22 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool:
 
     ``make_allreduce(mapping)`` builds the key exchange (collective: every rank
     installs at the same point of the load). A second call on an installed
-    ``llm`` changes nothing.
+    ``llm`` changes nothing. Everything is built and checked before anything
+    changes, and each refusal is agreed over the TP group (``agreed_reason``
+    for the checks, ``agreed_refusal`` for the key exchange ``make_allreduce``
+    built), so a refusing install leaves the worker and the drafter as they
+    were on every rank.
 
     Raises:
         ValueError: when there is no DFlash 2 drafter (no spec worker, or a
             drafter without a candidate selector), or its head layout is not one
-            the feature supports, or this torch's ``torch.topk`` orders ties
-            differently from the merge (``tie_order_reason``) -- running the
-            built-in path, or a path whose drafts differ, under the fast
-            path's name would misreport what runs.
+            the feature supports, or its logit transform is not one the draft
+            head can apply after the top-k (``logit_transform.unsupported_reason``),
+            or this torch's ``torch.topk`` orders ties differently from the merge
+            (``tie_order_reason``), or the key exchange's MNNVL all-reduce could
+            not be built on some rank -- running the built-in path, or a path
+            whose drafts differ, under the fast path's name would misreport what
+            runs.
     """
     if installed(llm):
         return True
@@ -435,29 +892,53 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool:
         raise ValueError(
             "[musespark_open] draft_head cannot run here: there is no DFlash 2 drafter"
         )
-    reason = unsupported_reason(worker, draft_model) or tie_order_reason(
-        draft_model.candidate_selector, worker.mapping.tp_size
-    )
+    selector = draft_model.candidate_selector
+    transform = None
+    reason = unsupported_reason(worker, draft_model)
+    if reason is None:
+        transform = _logit_transform.build(draft_model, selector.predecessor_codebook.device)
+        if transform is not None:
+            reason = _logit_transform.unsupported_reason(
+                transform, draft_model.lm_head.out_features
+            )
+    if reason is None:
+        reason = tie_order_reason(selector, worker.mapping.tp_size, transform=transform)
+    reason = agreed_reason(worker.mapping, reason, transform)
     if reason is not None:
         raise ValueError(f"[musespark_open] draft_head cannot run here: {reason}")
-    worker.__class__ = ShardedDraftHeadDFlashWorker
-    if type(draft_model.logits_processor) is LogitsProcessor:
-        draft_model.logits_processor = ShardDraftLogitsProcessor(draft_model.logits_processor)
-    if "_draft_head_allreduce" not in worker.__dict__:
+    if "_draft_head_allreduce" in worker.__dict__:
+        allreduce = worker.__dict__["_draft_head_allreduce"]
+    else:
         allreduce = make_allreduce(worker.mapping)
         if allreduce is None:
             # The composition's precondition checked the group takes an MNNVL all-reduce.
-            raise ValueError(
-                "the vocab-sharded draft head: the MNNVL FP32 all-reduce of the key exchange "
-                "could not be built on this TP group"
-            )
-        # Not registered as a submodule: it owns no parameters, only a fabric workspace.
-        worker.__dict__["_draft_head_allreduce"] = allreduce
-    exchange = "MNNVL all-reduce" if worker.__dict__["_draft_head_allreduce"] else "NCCL all-gather"
+            reason = "the MNNVL FP32 all-reduce of the key exchange could not be built"
+    # Building the all-reduce can fail on some ranks only; every rank refuses alike.
+    reason = agreed_refusal(worker.mapping, reason)
+    if reason is not None:
+        raise ValueError(f"[musespark_open] draft_head cannot run here: {reason}")
+    worker.__class__ = ShardedDraftHeadDFlashWorker
+    if type(draft_model.logits_processor) is LogitsProcessor:
+        draft_model.logits_processor = ShardDraftLogitsProcessor(draft_model.logits_processor)
+    # Not registered as a submodule: it owns no parameters, only a fabric workspace.
+    worker.__dict__["_draft_head_allreduce"] = allreduce
+    if transform is not None:
+        # Tables built outside capture, on the drafter's device; plain attributes, not buffers.
+        worker.__dict__[TRANSFORM_ATTR] = transform
+        draft_model.__dict__[TRANSFORM_METHOD] = _DeferredTransform(worker, draft_model)
+    exchange = "MNNVL all-reduce" if allreduce else "NCCL all-gather"
+    applied = (
+        ""
+        if transform is None
+        else (
+            f", drafter logit transform applied after the top-k (canonical keys outside "
+            f"[{transform.lo:#06x}, {transform.hi:#06x}])"
+        )
+    )
     logger.info(
         "[musespark_open] draft_head: DFlash 2 draft head on vocab shards "
         f"(shard {draft_model.lm_head.out_features}, tp {worker.mapping.tp_size}, "
         f"top_k {draft_model.candidate_selector.top_k}, keys by {exchange}, "
-        f"tie order checked against torch {torch.__version__})"
+        f"tie order checked against torch {torch.__version__}{applied})"
     )
     return True
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py
new file mode 100644
index 000000000000..9c58afd615fe
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py
@@ -0,0 +1,303 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The DFlash drafter's logit transform on the vocab-sharded draft head (``draft_head``).
+
+A DFlash drafter whose config declares ``output_multiplier`` or
+``final_logit_softcapping`` is trained on transformed logits, and the built-in
+worker applies ``DFlashForCausalLM.apply_draft_logit_transform`` to the draft
+logits before the selector: in fp32 and in this order, ``.float()``, ``*
+multiplier``, ``/ softcap``, ``tanh``, ``* softcap``. Call it ``t``. The
+selector then takes ``torch.topk(t(logits), top_k)``: its ids, their order and
+its values (the unary scores it adds to its pairwise ones).
+
+``t`` never decreases, so a top-k of the bf16 logits has the same candidates in
+the same order -- except where ``t`` maps distinct bf16 values to the same fp32
+value. tanh saturates: with the Muse Spark Open drafter's multiplier and softcap
+every logit of large enough magnitude transforms to exactly +-softcap. Those
+compare equal for ``torch.topk``, which then keeps and orders them by vocab id,
+while their bf16 values are distinct. So the draft head ranks canonical keys
+and reads the unary values from a table, both built once per drafter here:
+
+    canon   [65536] int32: 16-bit radix key of a bf16 logit -> the key of its tie
+            group's representative. A tie group is the bf16 values whose
+            transformed values have the same bits; ``t`` never decreases, so a
+            group is a run of consecutive keys, and canonical keys compare as the
+            transformed values' radix keys do. The representative is the group's
+            smallest key, or its largest below +0.0 so that the -0.0 group keeps
+            -0.0 (``-0.0 == +0.0`` under the selector sort's comparison, as for
+            the transformed values). NaN (0xFFFF) maps to itself.
+    values  [65536] fp32: key -> ``t`` of that bf16 value, computed by the
+            drafter's own method on the device, so it is the built-in transform
+            of every logit bit for bit (the transform is elementwise).
+    lo, hi  the key interval around 0 on which ``canon`` is the identity: logits
+            there need no table load.
+
+``unsupported_reason`` checks all of it exhaustively over the 65536 keys before
+the draft head installs it.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import weakref
+from typing import Callable, Optional
+
+import torch
+
+from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM
+
+from ...ops.spec.draft_shard_topk import CANON_ENTRIES, FINITE_KEYS, canon_reason, radix_key16
+
+__all__ = [
+    "DraftLogitTransform",
+    "bf16_of_key16",
+    "build",
+    "digest",
+    "ensure_transformed",
+    "same_bits",
+    "tie_groups",
+    "unsupported_reason",
+]
+
+#: Rows of the install-time layout check: eight requests of K = 7 positions.
+LAYOUT_ROWS = 56
+
+#: Seed of the layout check's logits, drawn on the host so no device RNG state moves.
+LAYOUT_SEED = 2718
+
+#: Key16 of +0.0 and -0.0.
+POSITIVE_ZERO, NEGATIVE_ZERO = 0x8000, 0x7FFF
+
+
+def bf16_of_key16(keys: torch.Tensor) -> torch.Tensor:
+    """The bf16 values of 16-bit radix keys (any int dtype, values in [0, 0xFFFF])."""
+    key = keys.to(torch.int32)
+    bits = torch.where(key >= 0x8000, key ^ 0x8000, key ^ 0xFFFF)
+    bits = torch.where(bits >= 0x8000, bits - 0x10000, bits).to(torch.int16)
+    return bits.view(torch.bfloat16)
+
+
+def _radix32(values: torch.Tensor) -> torch.Tensor:
+    """torch's fp32 radix key of ``values`` as int64 in [0, 2**32); every NaN is the top key."""
+    word = values.contiguous().view(torch.int32).to(torch.int64) & 0xFFFFFFFF
+    key = torch.where(word >= 0x80000000, word ^ 0xFFFFFFFF, word ^ 0x80000000)
+    return torch.where(values.isnan(), torch.full_like(key, 0xFFFFFFFF), key)
+
+
+def same_bits(a: torch.Tensor, b: torch.Tensor) -> bool:
+    """Bitwise equality of fp32 tensors, NaN payloads aside."""
+    nan = a.isnan()
+    if not torch.equal(nan, b.isnan()):
+        return False
+    return torch.equal(
+        torch.where(nan, 0.0, a).view(torch.int32), torch.where(nan, 0.0, b).view(torch.int32)
+    )
+
+
+class DraftLogitTransform:
+    """A drafter's logit transform as the draft head's tables (see the module docstring).
+
+    ``canon`` [65536] int32 and ``values`` [65536] fp32 live on the drafter's
+    device; ``lo`` and ``hi`` are plain ints; ``multiplier`` and ``softcap`` are
+    the drafter attributes the tables were built from. ``real`` is the drafter's
+    own transform, the built-in computation.
+    """
+
+    def __init__(
+        self,
+        draft_model,
+        function: Callable,
+        canon: torch.Tensor,
+        values: torch.Tensor,
+        lo: int,
+        hi: int,
+    ):
+        self._model = weakref.ref(draft_model)
+        self.function = function
+        self.canon = canon
+        self.values = values
+        self.lo = lo
+        self.hi = hi
+        self.multiplier = getattr(draft_model, "logit_multiplier", None)
+        self.softcap = getattr(draft_model, "logit_softcap", None)
+
+    @property
+    def band(self) -> tuple[int, int]:
+        """``(lo, hi)``: ``draft_shard_topk``'s ``canon_band``."""
+        return self.lo, self.hi
+
+    def matches(self, draft_model) -> bool:
+        """Whether ``draft_model`` still declares the transform these tables hold."""
+        return (
+            getattr(draft_model, "logit_multiplier", None),
+            getattr(draft_model, "logit_softcap", None),
+        ) == (self.multiplier, self.softcap)
+
+    def real(self, logits: torch.Tensor) -> torch.Tensor:
+        """The drafter's transform of ``logits``: the built-in computation, fp32 out."""
+        draft_model = self._model()
+        if draft_model is None:
+            raise RuntimeError("[musespark_open] draft_head: the drafter of this transform is gone")
+        return self.function(draft_model, logits)
+
+
+def build(draft_model, device) -> Optional[DraftLogitTransform]:
+    """The tables of ``draft_model``'s logit transform on ``device``, or None for no transform.
+
+    None when the drafter's class has no ``apply_draft_logit_transform`` or it
+    returns a bf16 probe unchanged (multiplier 1, no softcap). The probe is
+    bf16 because on an fp32 input the built-in transform runs in place and
+    returns the input itself even when it is on. Runs the drafter's method on
+    all 65536 bf16 values, so call it outside CUDA graph capture.
+    ``unsupported_reason`` says whether the result can be used.
+    """
+    function = getattr(type(draft_model), "apply_draft_logit_transform", None)
+    if function is None:
+        return None
+    probe = torch.zeros(4, dtype=torch.bfloat16, device=device)
+    if function(draft_model, probe) is probe:
+        return None
+    keys = torch.arange(CANON_ENTRIES, dtype=torch.int32, device=device)
+    values = function(draft_model, bf16_of_key16(keys))
+    canon = torch.arange(CANON_ENTRIES, dtype=torch.int32)
+    lo, hi = 0, CANON_ENTRIES - 1
+    if values.dtype == torch.float32 and tuple(values.shape) == (CANON_ENTRIES,):
+        first_key, last_key = FINITE_KEYS
+        radix = _radix32(values.cpu())[first_key : last_key + 1]
+        finite = torch.arange(first_key, last_key + 1, dtype=torch.int64)
+        start = torch.ones_like(radix, dtype=torch.bool)
+        start[1:] = radix[1:] != radix[:-1]
+        end = torch.ones_like(start)
+        end[:-1] = start[1:]
+        group = torch.cumsum(start.to(torch.int64), dim=0) - 1
+        first, last = finite[start], finite[end]
+        representative = torch.where(first >= POSITIVE_ZERO, first, last)
+        canon[first_key : last_key + 1] = representative[group].to(torch.int32)
+        identity = canon == torch.arange(CANON_ENTRIES, dtype=torch.int32)
+        above = (~identity[POSITIVE_ZERO:]).nonzero()
+        below = (~identity[:POSITIVE_ZERO]).nonzero()
+        hi = CANON_ENTRIES - 1 if above.numel() == 0 else POSITIVE_ZERO + int(above[0]) - 1
+        lo = 0 if below.numel() == 0 else int(below[-1]) + 1
+    return DraftLogitTransform(draft_model, function, canon.to(device), values.contiguous(), lo, hi)
+
+
+def tie_groups(transform: DraftLogitTransform) -> list[tuple[int, int]]:
+    """``(first, last)`` key16 of every tie group of two or more values, in key order."""
+    first_key, last_key = FINITE_KEYS
+    canon = transform.canon.cpu()[first_key : last_key + 1].to(torch.int64)
+    start = torch.ones_like(canon, dtype=torch.bool)
+    start[1:] = canon[1:] != canon[:-1]
+    begins = (start.nonzero().view(-1) + first_key).tolist()
+    ends = [b - 1 for b in begins[1:]] + [last_key]
+    return [(b, e) for b, e in zip(begins, ends) if e > b]
+
+
+def _layout_probe(rows: int, shard: int) -> torch.Tensor:
+    """Host bf16 ``[rows, shard]`` logits: every bf16 value once, then values of many scales.
+
+    The 65536 bf16 bit patterns (NaNs included) fill the first columns in a
+    random order; each remaining row is normal logits of one scale up to deep
+    saturation, with +-inf, +-0.0 and NaN scattered in.
+    """
+    g = torch.Generator().manual_seed(LAYOUT_SEED)
+    scales = (1.0, 10.0, 100.0, 400.0, 1000.0, 3000.0, 1e5, 1e30)
+    x = torch.randn(rows, shard, generator=g)
+    x *= torch.tensor([scales[r % len(scales)] for r in range(rows)])[:, None]
+    special = torch.tensor([float("inf"), float("-inf"), 0.0, -0.0, float("nan")])
+    picks = torch.randint(0, shard, (rows, 40), generator=g)
+    x.scatter_(1, picks, special[torch.randint(0, len(special), (rows, 40), generator=g)])
+    x = x.to(torch.bfloat16)
+    every = torch.arange(CANON_ENTRIES, dtype=torch.int32)[
+        torch.randperm(CANON_ENTRIES, generator=g)
+    ]
+    flat = x.view(-1)
+    count = min(CANON_ENTRIES, flat.numel())
+    flat[:count] = bf16_of_key16(every[:count])
+    return x
+
+
+def unsupported_reason(
+    transform: DraftLogitTransform, shard: int, rows: int = LAYOUT_ROWS
+) -> Optional[str]:
+    """Why the draft head cannot apply ``transform`` after its top-k, or None.
+
+    Exhaustive over the 65536 bf16 values: the transform is the built-in
+    ``DFlashForCausalLM`` method, returns fp32, keeps NaN, +0.0 and -0.0, never
+    decreases in torch's radix order (which makes the tie groups runs of
+    consecutive keys), and the canonical values compare under ``>`` as the
+    transformed values do. ``canon`` is a table ``draft_shard_topk`` reads
+    correctly with band ``(lo, hi)``. Last, the table gives the drafter's own
+    transform of a ``[rows, shard]`` bf16 slice bit for bit (a host-seeded
+    probe holding every bf16 value).
+    """
+    function = transform.function
+    if function is not DFlashForCausalLM.apply_draft_logit_transform:
+        return (
+            f"the drafter's apply_draft_logit_transform is {function.__qualname__}, "
+            "not DFlashForCausalLM's"
+        )
+    values = transform.values
+    if values.dtype != torch.float32 or tuple(values.shape) != (CANON_ENTRIES,):
+        return f"the drafter's logit transform returns {values.dtype}, not fp32"
+    table = values.cpu()
+    inputs = bf16_of_key16(torch.arange(CANON_ENTRIES, dtype=torch.int32))
+    if not torch.equal(inputs.isnan(), table.isnan()):
+        return "the drafter's logit transform does not map exactly the NaN logits to NaN"
+    words = table.view(torch.int32)
+    if int(words[POSITIVE_ZERO]) != 0 or int(words[NEGATIVE_ZERO]) != -(2**31):
+        return "the drafter's logit transform does not keep +0.0 and -0.0"
+    first_key, last_key = FINITE_KEYS
+    radix = _radix32(table)[first_key : last_key + 1]
+    drops = (radix[1:] < radix[:-1]).nonzero().view(-1)
+    if drops.numel():
+        at = float(inputs[first_key + int(drops[0]) + 1].float())
+        return f"the drafter's logit transform decreases at the bf16 logit {at}"
+    canon = transform.canon.cpu()
+    reason = canon_reason(canon, transform.band, torch.device("cpu"))
+    if reason is not None:
+        return reason
+    finite = canon[first_key : last_key + 1]
+    start = torch.ones_like(finite, dtype=torch.bool)
+    start[1:] = finite[1:] != finite[:-1]
+    members = torch.arange(first_key, last_key + 1)[start]
+    if bool((radix[1:] == radix[:-1])[start[1:]].any()):
+        return "canon splits a tie group of the transform"
+    if not torch.equal(_radix32(table[finite.long()]), _radix32(table[first_key : last_key + 1])):
+        return "canon maps a key outside its tie group"
+    shown = bf16_of_key16(finite[start]).float()
+    shown_values = table[members]
+    if not torch.equal(shown[1:] > shown[:-1], shown_values[1:] > shown_values[:-1]):
+        return "canonical values compare differently from the transformed values"
+    probe = _layout_probe(rows, shard).to(values.device)
+    got = values[radix_key16(probe.view(torch.int16)).long()]
+    if not same_bits(got, transform.real(probe)):
+        return (
+            f"the drafter's logit transform of a [{rows}, {shard}] slice differs from its "
+            "transform of the 65536 bf16 values"
+        )
+    return None
+
+
+def digest(transform: Optional[DraftLogitTransform]) -> str:
+    """A short digest of the tables, equal on ranks whose draft heads would agree."""
+    if transform is None:
+        return "no drafter logit transform"
+    h = hashlib.sha256()
+    h.update(transform.canon.cpu().numpy().tobytes())
+    h.update(transform.values.cpu().view(torch.int32).numpy().tobytes())
+    h.update(repr((transform.lo, transform.hi, transform.multiplier, transform.softcap)).encode())
+    return f"drafter logit transform tables {h.hexdigest()[:16]}"
+
+
+def ensure_transformed(
+    transform: Optional[DraftLogitTransform], logits: torch.Tensor
+) -> torch.Tensor:
+    """``logits`` with the drafter's transform applied once.
+
+    With a transform installed, bf16 draft logits are untransformed (the
+    transform always returns fp32, which ``unsupported_reason`` checks), so a
+    bf16 ``logits`` gets ``transform.real`` and anything else is returned as is.
+    """
+    if transform is None or logits.dtype != torch.bfloat16:
+        return logits
+    return transform.real(logits)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py
index 292e17633680..358c8e76a726 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py
@@ -68,9 +68,11 @@
 #: Operand dtypes the fused convolution takes; anything else runs the built-in.
 SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
 
-#: Block rows on which the PDL chain's conv / merge entry triggers apply (``chain.ROWS``):
-#: one request of K + 1 = 8 rows, whose successor GEMMs run the ``drafter_fixed_gemms`` tactic.
-_ENTRY_ROWS = 8
+#: Block rows on which the PDL chain's conv / merge entry triggers apply when the chain's
+#: install did not record their widths (``_pdl_entry_rows`` / ``_pdl_merge_rows``,
+#: ``chain.entry_rows`` / ``chain.MERGE_ENTRY_ROWS_BY_SM``): one request of K + 1 = 8 rows, whose
+#: successor GEMMs run the ``drafter_fixed_gemms`` tactic.
+_ENTRY_ROWS = (8,)
 
 
 class FusedDFlash2BlockConv(DFlash2BlockConv):
@@ -95,9 +97,13 @@ def _convolve(
             or hidden_states.dtype not in SUPPORTED_DTYPES
         ):
             return super()._convolve(hidden_states, delta, side, block_size)
-        if side == 0 and hidden_states.shape[0] == _ENTRY_ROWS and self.__dict__.get("_pdl_conv"):
-            # PDL chain conv trigger (drafter_pdl_chain/chain.py): the prepare side of an 8-row block, whose
-            # successor GEMM loads only its weight before its own wait.
+        if (
+            side == 0
+            and self.__dict__.get("_pdl_conv")
+            and hidden_states.shape[0] in self.__dict__.get("_pdl_entry_rows", _ENTRY_ROWS)
+        ):
+            # PDL chain conv trigger (drafter_pdl_chain/chain.py): the prepare side of a block whose
+            # successor GEMM runs a fixed tactic, which loads only its weight before its own wait.
             return dflash2_block_conv(
                 hidden_states, delta, base_kernel, block_size, self.group_size, True
             )
@@ -318,7 +324,9 @@ def windowed_block_decode(
     layers = self.model.layers
     prenormed = None
     # PDL chain merge trigger (drafter_pdl_chain/chain.py): the merge releases o_proj's GEMM at its entry.
-    merge_entry = bool(self.__dict__.get("_pdl_merge")) and B * block_size == _ENTRY_ROWS
+    merge_entry = bool(self.__dict__.get("_pdl_merge")) and B * block_size in self.__dict__.get(
+        "_pdl_merge_rows", _ENTRY_ROWS
+    )
     # wide_drafter_window_attention: a block decode of several requests on the wide op.
     wide = self.__dict__.get("_wide_window_attention")
     if wide is not None and not wide.takes(B * block_size):
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py
index eabd009108a8..d6970d7f8eb1 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py
@@ -18,6 +18,11 @@
 ``down_proj``) is the same on every rank and every forward; a forward takes
 the fused hops at all its sites or at none, decided on the block's row count
 (``DrafterHopCall.takes``), which every rank shares.
+
+With ``wide_drafter_hops`` (``install_wide``) a forward of ``MAX_ROWS + 1`` ..
+``wide_served_rows()`` rows (the blocks of several requests at K = 7, per device in
+``WIDE_SERVED_ROWS_BY_SM``) takes the same hops on a second, wide workspace; a forward of at
+most ``MAX_ROWS`` rows keeps the first one.
 """
 
 from __future__ import annotations
@@ -46,18 +51,29 @@ class DrafterHopCall:
 
     def __init__(self, workspace: "_hop.DrafterHopWorkspace") -> None:
         self.workspace = workspace
+        #: The wide workspace (``wide_drafter_hops``), or None.
+        self.wide: Optional["_hop.DrafterHopWorkspace"] = None
+        #: The workspace of the forward ``takes`` admitted last.
+        self.active = workspace
         #: Hops issued by this rank (a test counts them to see a forward took the fused path).
         self.calls = 0
 
     def takes(self, rows: int, block_size: int, dtype: torch.dtype) -> bool:
-        """Whether a forward of ``rows`` rows runs every site on the fused hop. Rank-uniform."""
-        if dtype != torch.bfloat16 or not 1 <= rows <= _hop.MAX_ROWS or block_size < 1:
+        """Whether a forward of ``rows`` rows runs every site on the fused hop, and on which
+        workspace. Rank-uniform."""
+        if dtype != torch.bfloat16 or block_size < 1 or rows < 1:
             return False
-        return True
+        if rows <= _hop.MAX_ROWS:
+            self.active = self.workspace
+            return True
+        if self.wide is not None and rows <= self.wide.max_rows:
+            self.active = self.wide
+            return True
+        return False
 
     def begin_forward(self) -> None:
         """A drafter forward starts; its hops name the sites from 0 (see the workspace's)."""
-        self.workspace.begin_forward()
+        self.active.begin_forward()
 
     def _site(self, site, partial, residual, coefficients, conv, norm, block_size):
         # ``coefficients`` is what ``prepare`` returned for ``finish``: [rows, taps, groups].
@@ -80,7 +96,7 @@ def _site(self, site, partial, residual, coefficients, conv, norm, block_size):
             norm.weight,
             norm.variance_epsilon,
             block_size,
-            self.workspace,
+            self.active,
         )
 
     def attention(
@@ -353,3 +369,42 @@ def install(draft_model: Optional[nn.Module], mapping=None) -> None:
     draft_model._drafter_hop = DrafterHopCall(workspace)
     _hop.log_once_enabled(2 * len(layers), False, early)
     logger.info("[musespark_open] drafter_hops: fused drafter hops installed")
+
+
+def install_wide(draft_model: Optional[nn.Module], mapping=None) -> None:
+    """``wide_drafter_hops``: the fused hops also on drafter forwards of ``MAX_ROWS + 1`` ..
+    ``wide_served_rows()`` rows of this device, on a wide workspace of their own. Collective;
+    after ``install``.
+
+    The calls the wide hops replace are the same three kernels per site as at 8 rows (the
+    MNNVL one-shot all-reduce ``request_mnnvl`` gave every row projection on one tray, the
+    finish convolution, flashinfer's fused add + RMSNorm); the hop reproduces them per row,
+    and the tap-1 mask per block (``row % block``). Idempotent.
+
+    Raises:
+        ValueError: when ``drafter_hops`` is not installed on this drafter.
+    """
+    if not installed(draft_model):
+        raise ValueError("[musespark_open] wide_drafter_hops needs drafter_hops installed first")
+    call: DrafterHopCall = draft_model._drafter_hop
+    if call.wide is not None:
+        return
+    narrow = call.workspace
+    layers = draft_model.model.layers
+    norms = [layer.post_attention_layernorm for layer in layers]
+    norms += [layer.input_layernorm for layer in layers[1:]] + [draft_model.model.norm]
+    eps = sorted({float(norm.variance_epsilon) for norm in norms})
+    rows = _hop.wide_served_rows()
+    call.wide = _hop.DrafterHopWorkspace(
+        layers[0].self_attn.o_proj.mapping,
+        sites=narrow.sites,
+        eps=eps,
+        early_trigger=narrow.early_trigger,
+        pre_wait=[tensor for _, tensor in pre_wait_operands(draft_model)],
+        max_rows=rows,
+    )
+    logger.info(
+        f"[musespark_open] wide_drafter_hops: fused drafter hops on forwards of {_hop.MAX_ROWS + 1}.."
+        f"{rows} rows ({call.wide.areas} areas, "
+        f"{call.wide.nbytes() / 2**20:.1f} MiB of fabric memory on this rank)"
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py
index 57c002131551..9a6e6aba94f4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py
@@ -37,10 +37,19 @@
   right before the draft LM-head GEMM and joined where the projection is
   called, after the key all-reduce, instead of on the main stream after it.
 
-Scope: the one-request step, whose successors are known. ``pdl_conv`` /
-``pdl_merge`` / ``pdl_silu`` act on the block decode's 8 rows (K + 1,
+Scope: block decodes whose successors are known. ``pdl_conv`` /
+``pdl_merge`` / ``pdl_silu`` act on the block decode's 8 rows (one request, K + 1,
 K = 7), where ``drafter_fixed_gemms`` has a fixed tactic for ``qkv_proj`` / ``gate_up_proj`` /
-``o_proj`` / ``down_proj`` / the first convolution's ``kernel_projection``.
+``o_proj`` / ``down_proj`` / the first convolution's ``kernel_projection``. ``pdl_conv``
+also acts on every other block width of ``ENTRY_ROW_CANDIDATES`` (2..4 requests: 16, 24
+and 32 rows) at which the drafter's table lists a fixed tactic for all four of those
+successors (``entry_rows``); a width without one runs the base launches. ``pdl_silu``
+and ``pdl_merge`` act on the widths their device's table lists among those
+(``SILU_ENTRY_ROWS_BY_SM`` / ``MERGE_ENTRY_ROWS_BY_SM``), each measured in the in-model
+drafter chain on that GPU. On GB300 all three triggers together are the fastest chain at
+1..4 requests. On Rubin an entry-released ``o_proj`` competes with the multi-request merge
+it follows (the chain is slower with the merge trigger at 2..4 requests), and the silu
+trigger on top of the conv trigger gains nothing at two requests and loses at three.
 ``pdl_head`` acts on a one-request step whose K slots are one contiguous row
 range (any K); its successor is the low-M dispatcher's runner for the LM
 head (cuBLAS through ``torch.mm``, which has no programmatic attribute, or
@@ -93,12 +102,50 @@
 from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod
 from tensorrt_llm.logger import logger
 
+from ...ops.gemm.fixed_low_m_gemm import device_capability
 from ...ops.spec.drafter_silu_and_mul import drafter_silu_and_mul
 
-#: The drafter block rows these launches act on: one request of K + 1 = 8 rows,
-#: the only width ``drafter_fixed_gemms`` lists a tactic for.
+#: The drafter block rows of one request, K + 1 = 8: the width ``pdl_conv`` / ``pdl_merge`` /
+#: ``pdl_silu`` require a fixed tactic at, and the ``pdl_head`` / ``selector_fork`` step.
 ROWS = 8
 
+#: The block widths the conv / merge / silu entry triggers may act on: 1..4 requests at
+#: K = 7, at most the fixed GEMM's ``MAX_M`` of 32 rows. A width is taken when the
+#: drafter's tactic table covers every successor at it (``entry_rows``).
+ENTRY_ROW_CANDIDATES = (8, 16, 24, 32)
+
+#: The block widths of the merge trigger per compute capability (``trigger_rows``), among
+#: those of ``entry_rows``. sm103: one to four requests. sm107: the one-request block only.
+MERGE_ENTRY_ROWS_BY_SM: dict[tuple[int, int], frozenset] = {
+    (10, 3): frozenset(ENTRY_ROW_CANDIDATES),
+    (10, 7): frozenset({ROWS}),
+}
+
+#: The block widths of the silu trigger per compute capability (``trigger_rows``), among
+#: those of ``entry_rows``. sm103: one to four requests. sm107: one and four requests.
+SILU_ENTRY_ROWS_BY_SM: dict[tuple[int, int], frozenset] = {
+    (10, 3): frozenset(ENTRY_ROW_CANDIDATES),
+    (10, 7): frozenset({ROWS, 4 * ROWS}),
+}
+
+
+def trigger_rows(table: dict, capability: Optional[tuple] = None) -> frozenset:
+    """The block widths ``table`` (``MERGE_ENTRY_ROWS_BY_SM`` / ``SILU_ENTRY_ROWS_BY_SM``)
+    lists for ``capability`` (the drafter GEMMs' device, ``fixed_low_m_gemm.device_capability``,
+    when None).
+
+    Raises:
+        ValueError: on a device the table does not list.
+    """
+    key = tuple(device_capability() if capability is None else capability)
+    rows = table.get(key)
+    if rows is None:
+        raise ValueError(
+            f"drafter_pdl_chain: no trigger widths for SM {key} (supported: {sorted(table)})"
+        )
+    return rows
+
+
 #: The draft head's rows for that request: its K = 7 draft slots. ``selector_fork``
 #: forks the selector projection only for this many rows.
 HEAD_ROWS = ROWS - 1
@@ -160,8 +207,29 @@ def _merge_reason(draft_model: nn.Module) -> Optional[str]:
     return None
 
 
+def entry_rows(draft_model: nn.Module) -> frozenset:
+    """The widths of ``ENTRY_ROW_CANDIDATES`` at which every successor of the conv / merge /
+    silu triggers (``qkv_proj``, ``gate_up_proj``, ``o_proj``, ``down_proj`` of every layer)
+    has a fixed tactic. A function of the drafter's shapes and TP table: rank-uniform."""
+    rows = []
+    for width in ENTRY_ROW_CANDIDATES:
+        if all(
+            _fixed_tactic_reason(linear, width, "") is None
+            for layer in draft_model.model.layers
+            for linear in (
+                layer.self_attn.qkv_proj,
+                layer.mlp.gate_up_proj,
+                layer.self_attn.o_proj,
+                layer.mlp.down_proj,
+            )
+        ):
+            rows.append(width)
+    return frozenset(rows)
+
+
 class EntryTriggerGatedMLP(GatedMLP):
-    """The built-in ``GatedMLP`` whose 8-row ``silu(gate) * up`` releases ``down_proj`` at entry.
+    """The built-in ``GatedMLP`` whose ``silu(gate) * up`` releases ``down_proj`` at entry on the
+    widths of ``_pdl_entry_rows`` (``entry_rows``; 8 rows when unset).
 
     Never constructed: ``install`` re-classes the drafter's ``GatedMLP``
     modules in place. Every other call is the built-in method's.
@@ -172,7 +240,7 @@ def _apply_activation(self, x, *, has_lora: bool = False):
             not has_lora
             and isinstance(x, torch.Tensor)
             and x.dim() == 2
-            and x.shape[0] == ROWS
+            and x.shape[0] in self.__dict__.get("_pdl_entry_rows", (ROWS,))
             and x.dtype == torch.bfloat16
             and x.is_contiguous()
         ):
@@ -385,6 +453,9 @@ def unsupported_reason(llm: nn.Module) -> Optional[str]:
     bound = draft_model.__dict__.get("dflash_forward")
     if getattr(bound, "__func__", None) is not _draft.windowed_block_decode:
         return "the drafter does not run the windowed block decode"
+    capability = device_capability()
+    if capability not in MERGE_ENTRY_ROWS_BY_SM or capability not in SILU_ENTRY_ROWS_BY_SM:
+        return f"no merge / silu trigger widths for SM {capability}"
     reason = _conv_reason(draft_model)
     if reason is not None:
         return f"pdl_conv: {reason}"
@@ -435,6 +506,10 @@ def installed(llm: nn.Module) -> frozenset:
     return frozenset(found)
 
 
+def _widths(widths) -> str:
+    return ",".join(str(w) for w in sorted(widths)) or "none"
+
+
 def install(llm: nn.Module, mapping=None) -> frozenset:
     """Install the five PDL chain launches; return the names this call installed.
 
@@ -458,12 +533,18 @@ def install(llm: nn.Module, mapping=None) -> frozenset:
         raise ValueError(reason)
     before = installed(llm)
     draft_model = llm.draft_model
+    widths = entry_rows(draft_model)
     for conv in list(draft_model.attention_convs) + list(draft_model.mlp_convs):
         conv._pdl_conv = True
+        conv._pdl_entry_rows = widths
+    merge_rows = widths & trigger_rows(MERGE_ENTRY_ROWS_BY_SM)
+    silu_rows = widths & trigger_rows(SILU_ENTRY_ROWS_BY_SM)
     draft_model._pdl_merge = True
+    draft_model._pdl_merge_rows = merge_rows
     for layer in draft_model.model.layers:
         if type(layer.mlp) is GatedMLP:
             layer.mlp.__class__ = EntryTriggerGatedMLP
+        layer.mlp._pdl_entry_rows = silu_rows
     llm.spec_worker._pdl_head = True
     projection = draft_model.candidate_selector.hidden_projection
     if type(projection) is nn.Linear:
@@ -476,6 +557,8 @@ def install(llm: nn.Module, mapping=None) -> frozenset:
     if changed:
         logger.info(
             f"[musespark_open] {','.join(sorted(changed))}: drafter PDL chain installed "
-            f"(in place: {','.join(sorted(now))})"
+            f"(in place: {','.join(sorted(now))}; block rows with the conv trigger "
+            f"{_widths(widths)}, the silu trigger {_widths(silu_rows)}, "
+            f"the merge trigger {_widths(merge_rows)})"
         )
     return changed
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py
new file mode 100644
index 000000000000..e2541f51f4f7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""mixed_step_gen_rows: the generation rows of mixed steps on multi_request_split_kv (sm_107)."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py
new file mode 100644
index 000000000000..baaf9cc37f21
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py
@@ -0,0 +1,257 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The generation rows of mixed steps on ``multi_request_split_kv`` (sm_107 full-attention layers).
+
+A mixed step (context rows and generation rows in one batch) runs ``PrefillTarget``, the
+built-in text-model forward. On sm_107 the public FMHA library list leaves flashinfer's
+trtllm-gen FMHA out; ``attention.select_decode_fmha`` puts a ``DecodeTrtllmGenFmha`` in front
+of every layer's libraries, which offers the generation phase only. With nothing in front of
+the thop fallback offering the context phase, the FMHA manager hands a mixed step's whole
+batch to the thop op, so its generation rows run thop's generation FMHA, not the decode-step
+paths (``DecodeTrtllmGenFmha`` and, in front of it, ``multi_request_split_kv``).
+
+This path changes that on the full-attention layers in two parts:
+
+* ``ThopContextPhaseFmha``, inserted in each full-attention layer's FMHA manager just before
+  the thop fallback, offers the context phase of a mixed step that ``takes_generation``
+  accepts (context rows, ``MIN_GENERATIONS``..``MAX_GENERATIONS`` generation sequences, beam
+  width 1). The manager then forms ``CombinedFmha(context=ThopContextPhaseFmha,
+  generation=DecodeTrtllmGenFmha)``: the context rows run the thop op with
+  ``attention_input_type=context_only`` (the same op, the same context launch as before) and the
+  generation rows run ``FlashInferTrtllmGenFmha.run_generation`` with ``seq_offset =
+  num_contexts``, as on sm_103.
+* ``MixedStepGenRowsAttention``, re-classed in front of the full-attention layers' attention,
+  runs such a step's attention call with the decode FMHA call intercepted by a
+  ``MixedStepCall`` (``attention_forward.intercepting``), so the generation phase launches
+  ``multi_request_split_kv`` as a decode step's does. ``MixedStepCall`` is
+  ``multi_request_split_kv.MultiRequestCall`` with the sequence lengths re-aligned: the
+  generation phase's view of them starts at ``num_contexts``. A call the stand-in declines runs
+  flashinfer's trtllm-gen decode FMHA.
+
+The band is read on the host from the step's metadata. Every bound is a point of the FMHA
+manager's batch-size cache grid, so a cached selection never crosses it. Sliding-window
+layers, pure context steps, decode steps and mixed steps outside the band run what they ran
+before. The arming sits at ``_attn_impl``, which the breakable prefill CUDA graph re-runs
+eagerly on every replay (``eager_on_graph``), so captured and eager mixed steps both take it.
+
+Numerics: the context rows are those of the thop op. The generation rows are those of
+``multi_request_split_kv`` on decode steps (not bit-identical to thop's generation FMHA;
+deterministic).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha
+from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase
+from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams, PhasedFmha
+from tensorrt_llm._torch.attention.backends.interface import AttentionInputType
+from tensorrt_llm.logger import logger
+
+from ...attention import DecodeTrtllmGenFmha
+from ...attention_forward import intercepting
+from ..multi_request_split_kv import attention as _mr_attention
+
+#: Generation sequences of a mixed step this path takes. Both bounds and their predecessors
+#: are points of the FMHA cache's batch-size grid (``fmha.manager``).
+MIN_GENERATIONS = 9
+MAX_GENERATIONS = 64
+
+
+def takes_generation(attn_metadata) -> bool:
+    """Whether a step is a mixed step whose generation rows this path takes.
+
+    Context rows, ``MIN_GENERATIONS``..``MAX_GENERATIONS`` generation sequences, beam width 1.
+    Host-side reads only.
+    """
+    return (
+        getattr(attn_metadata, "num_contexts", 0) > 0
+        and MIN_GENERATIONS <= getattr(attn_metadata, "num_generations", 0) <= MAX_GENERATIONS
+        and getattr(attn_metadata, "beam_width", 1) == 1
+    )
+
+
+class ThopContextPhaseFmha(PhasedFmha):
+    """The thop op's context phase, offered on the mixed steps ``takes_generation`` accepts.
+
+    It never takes a whole batch or a generation phase, so the manager pairs it with the
+    generation library in front of it. ``run_context`` calls the thop fallback with the
+    context rows and ``attention_input_type=context_only``.
+    """
+
+    supports_workspace_reclamation = True
+
+    def __init__(self, attn, fallback: FallbackFmha):
+        super().__init__(attn)
+        self._fallback = fallback
+        #: Context phases run, read by the served-path log.
+        self.context_calls = 0
+
+    def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool:
+        if phase != FmhaPhase.CONTEXT or not takes_generation(metadata):
+            return False
+        if forward_args.attention_input_type != AttentionInputType.mixed:
+            return False
+        if forward_args.output is None or forward_args.output_sf is not None:
+            return False
+        return self._fallback.is_supported(q, k, v, metadata, forward_args)
+
+    def run_context(self, params: FmhaParams) -> None:
+        rows = params.num_tokens
+        attention_input = params.qkv_input if params.qkv_input is not None else params.query_input
+        forward_args = dataclasses.replace(
+            params.fwd,
+            output=params.fwd.output[:rows],
+            attention_input_type=AttentionInputType.context_only,
+        )
+        self._fallback.forward(
+            attention_input, params.key_input, params.value_input, params.meta, forward_args
+        )
+        self.context_calls += 1
+
+    def run_generation(self, params: FmhaParams) -> None:
+        raise RuntimeError("ThopContextPhaseFmha serves the context phase only")
+
+
+def context_library(attn) -> Optional[ThopContextPhaseFmha]:
+    """The ``ThopContextPhaseFmha`` in ``attn``'s FMHA manager, or None."""
+    for fmha in attn._fmha_manager.fmha_libs:
+        if isinstance(fmha, ThopContextPhaseFmha):
+            return fmha
+    return None
+
+
+#: Tensors of the decode call ``multi_request_split_kv`` reads at a 16-byte aligned address.
+_ALIGNED_KWARGS = ("query", "out", "block_tables")
+
+
+class MixedStepCall(_mr_attention.MultiRequestCall):
+    """A mixed step's generation-phase decode call for ``multi_request_split_kv``.
+
+    The generation phase reads the sequence lengths from ``seq_offset = num_contexts`` on, so
+    ``seq_lens`` is a view that starts 4 bytes per context sequence into the step's lengths.
+    The kernel reads it at a 16-byte aligned address; ``launch`` hands it an aligned copy of
+    those ``B`` lengths when the view is not aligned. The other tensors it reads are declined
+    when not aligned.
+    """
+
+    def decline(self, kwargs: dict) -> Optional[str]:
+        reason = super().decline(kwargs)
+        if reason is not None:
+            return reason
+        for name in _ALIGNED_KWARGS:
+            if kwargs[name].data_ptr() % 16:
+                return f"a {name} not 16-byte aligned"
+        return None
+
+    def launch(self, kwargs: dict) -> None:
+        seq_lens = kwargs["seq_lens"]
+        if seq_lens.data_ptr() % 16:
+            kwargs = dict(kwargs, seq_lens=seq_lens.clone())
+        super().launch(kwargs)
+
+
+class MixedStepGenRowsAttention:
+    """Mixin re-classed in front of a full-attention layer's attention class.
+
+    A call on a step ``takes_generation`` accepts runs the class behind it with the decode
+    FMHA call intercepted by a ``MixedStepCall`` of the layer's KV heads; any other call is
+    exactly the class behind it. Adds plain attributes only.
+    """
+
+    _mixed_kv_heads: int = 4
+
+    def _attn_impl(self, q, k, v, attn_metadata, *args, **kwargs):
+        if not takes_generation(attn_metadata):
+            return super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs)
+        call = MixedStepCall(kv_heads=self._mixed_kv_heads, window=0)
+        with intercepting(call):
+            result = super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs)
+        if call.launched:
+            outcome = "multi_request_split_kv"
+        elif call.fallback is not None:
+            outcome = f"flashinfer's trtllm-gen decode FMHA ({call.fallback})"
+        else:
+            outcome = "no trtllm-gen decode call reached (the thop op ran the whole batch)"
+        logger.info_once(
+            f"[musespark_open] mixed_step_gen_rows: a mixed step's generation rows "
+            f"({attn_metadata.num_generations} requests, layer {self.layer_idx}): {outcome}",
+            key=f"musespark_open.mixed_step_gen_rows.{outcome}",
+        )
+        return result
+
+
+@functools.lru_cache(maxsize=None)
+def _mixed_class(base: type) -> type:
+    """``base`` with ``MixedStepGenRowsAttention`` in front, one class per base."""
+    return type(f"{base.__name__}MixedGen", (MixedStepGenRowsAttention, base), {})
+
+
+def _refusal(attention) -> Optional[str]:
+    libs = attention.attn._fmha_manager.fmha_libs
+    if not libs or not isinstance(libs[0], DecodeTrtllmGenFmha):
+        return "no DecodeTrtllmGenFmha in front of its FMHA libraries"
+    if not any(isinstance(fmha, FallbackFmha) for fmha in libs):
+        return "no thop fallback among its FMHA libraries"
+    if not isinstance(attention, _mr_attention.MultiRequestSplitKvAttention):
+        return "its attention is not on multi_request_split_kv"
+    return None
+
+
+def install(model: torch.nn.Module) -> int:
+    """Put a ``ThopContextPhaseFmha`` before the thop fallback of every full-attention layer's
+    FMHA manager and re-class those layers' attention with ``MixedStepGenRowsAttention`` in
+    front of whatever class they carry. Returns the number of layers changed.
+
+    Run after ``attention.select_decode_fmha`` and the ``multi_request_split_kv`` install, and
+    after every other attention re-class: the mixin acts only on mixed steps, which none of the
+    others intercept. Every layer is checked before any changes; idempotent.
+
+    Raises:
+        ValueError: if a full-attention layer's FMHA manager does not start with
+            ``DecodeTrtllmGenFmha``, holds no thop fallback, or its attention is not on
+            ``multi_request_split_kv``.
+    """
+    chosen = [layer.self_attn for layer in model.layers if not layer.self_attn.is_sliding]
+    refused = [
+        f"layer {attention.layer_idx}: {reason}"
+        for attention in chosen
+        if (reason := _refusal(attention)) is not None
+    ]
+    if refused:
+        raise ValueError(
+            "[musespark_open] mixed_step_gen_rows cannot route mixed steps' generation rows: "
+            + "; ".join(refused)
+        )
+    for attention in chosen:
+        attn = attention.attn
+        if context_library(attn) is None:
+            libs = attn._fmha_manager.fmha_libs
+            index = next(i for i, fmha in enumerate(libs) if isinstance(fmha, FallbackFmha))
+            libs.insert(index, ThopContextPhaseFmha(attn, libs[index]))
+        if not isinstance(attention, MixedStepGenRowsAttention):
+            attention.__class__ = _mixed_class(type(attention))
+        attention._mixed_kv_heads = attention._mr_kv_heads
+    logger.info(
+        f"[musespark_open] mixed_step_gen_rows: the generation rows of mixed steps with "
+        f"{MIN_GENERATIONS}..{MAX_GENERATIONS} generation requests run multi_request_split_kv "
+        f"on {len(chosen)} full-attention layers (context rows on the thop op)"
+    )
+    return len(chosen)
+
+
+__all__ = [
+    "MAX_GENERATIONS",
+    "MIN_GENERATIONS",
+    "MixedStepCall",
+    "MixedStepGenRowsAttention",
+    "ThopContextPhaseFmha",
+    "context_library",
+    "install",
+    "takes_generation",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py
new file mode 100644
index 000000000000..a0a97f10d1ff
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""model_boundary: the decode step's input boundary in one kernel, and a BF16 last-layer join."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py
new file mode 100644
index 000000000000..31e557cc3c3f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py
@@ -0,0 +1,224 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Fast path ``model_boundary``: the decode step's model input and output boundaries, fewer launches.
+
+Input. ``DecodeTarget`` runs, before layer 0's QKV projection, the replicated embedding gather,
+``embed_norm``, the FP32 residual copy, its BF16 copy and layer 0's ``pre_qkv_proj_norm``: five
+launches. With this fast path a decode step of at most ``replicated_embedding.MAX_ROWS`` ids
+runs ``input_boundary`` instead -- one launch, the same bits -- and hands its two outputs to
+layer 0 as the residual stream and ``precomputed_attn_input``, as every later layer receives
+them from the previous layer's fused boundary.
+
+Output. Layer 61 has no next norm, so its post-FFN join is the built-in
+``add_after_rmsnorm`` (the Triton ``_gated_residual_after_rmsnorm_kernel``, FP32 out), then
+``DecodeTarget`` casts the stream to BF16 for the final norm: a separate copy kernel. While
+``ModelBoundary.armed`` holds, that join's kernel stores BF16 directly -- the same kernel,
+its output pointer BF16, so Triton's store rounds the FP32 value to nearest even exactly as
+``.to(torch.bfloat16)`` does -- and the cast is a no-op. The FP32 stream after layer 61 has
+no reader: it is not a DFlash capture layer (``install`` refuses otherwise) and the final
+norm reads its BF16 cast only.
+
+What this module does:
+
+* ``install`` (after ``replicated_embedding``, from ``load_weights``) checks the model it acts
+  on -- the replicated table, ``m_emb == 1``, the historical pre-QKV ordering, flashinfer's
+  RMSNorm geometry, layer 61 not captured -- builds the kernel (agreed over the TP group: a
+  refusal on one rank makes every rank raise) and re-classes layer 61's post-FFN join.
+* ``ModelBoundary.takes`` / ``input`` serve the input boundary per call: shapes and dtypes only,
+  so every rank takes the same branch; ``inputs_embeds`` calls and other row counts keep the
+  built-in chain.
+* ``ModelBoundary.armed`` arms layer 61's BF16 join for one decode forward and disarms it on
+  any exit; prefill never arms it.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from typing import FrozenSet, Optional
+
+import torch
+
+from tensorrt_llm._torch.models import modeling_musespark_open as _builtin
+from tensorrt_llm.logger import logger
+
+from ... import dispatch as _tp_agreement
+from ..replicated_embedding import embedding as _replicated
+from . import input_boundary as _input
+
+
+class Bf16OutputJoin:
+    """Mixin for layer 61's post-FFN join: ``add_after_rmsnorm`` stores BF16 while armed.
+
+    ``_bf16_out`` is set by ``ModelBoundary.armed`` for one decode forward. Unarmed, the
+    join is the class it was before ``install``.
+    """
+
+    _bf16_out: bool = False
+
+    def add_after_rmsnorm(
+        self,
+        residual: torch.Tensor,
+        branch: torch.Tensor,
+        eps: float,
+        gain: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if not self._bf16_out:
+            return super().add_after_rmsnorm(residual, branch, eps, gain)
+        num_tokens = residual.shape[0]
+        hidden_size = residual.shape[-1]
+        if not (
+            self._fused_enabled
+            and residual.is_cuda
+            and num_tokens > 0
+            and residual.is_contiguous()
+            and branch.is_contiguous()
+            and branch.dtype == torch.bfloat16
+            and (gain is None or gain.dtype == branch.dtype)
+        ):
+            # The built-in composition, then the cast DecodeTarget would have made.
+            return super().add_after_rmsnorm(residual, branch, eps, gain).to(torch.bfloat16)
+        alpha, beta = self.coefficients()
+        out = torch.empty(residual.shape, dtype=torch.bfloat16, device=residual.device)
+        # The built-in launch (MuseSparkOpenGatedResidual.add_after_rmsnorm), with a BF16 output.
+        _builtin._gated_residual_after_rmsnorm_kernel[(num_tokens,)](
+            out,
+            residual,
+            branch,
+            gain if gain is not None else alpha,
+            alpha,
+            beta,
+            eps,
+            out,
+            alpha,
+            eps,
+            HIDDEN=hidden_size,
+            BLOCK_H=_builtin.triton.next_power_of_2(hidden_size),
+            HAS_GAIN=gain is not None,
+            FUSE_NEXT_NORM=False,
+            NORM_FP32_RESIDUAL=False,
+            num_warps=_builtin._RESIDUAL_NUM_WARPS,
+        )
+        return out
+
+
+_JOIN_CLASSES: dict = {}
+
+
+def _bf16_join_class(base: type) -> type:
+    cls = _JOIN_CLASSES.get(base)
+    if cls is None:
+        cls = type(f"ModelBoundaryBf16{base.__name__}", (Bf16OutputJoin, base), {})
+        _JOIN_CLASSES[base] = cls
+    return cls
+
+
+class ModelBoundary:
+    """What ``DecodeTarget`` holds under ``model_boundary`` (``core._model_boundary``)."""
+
+    def __init__(self, model) -> None:
+        self.table = model.embed_tokens.__dict__[_replicated.TABLE_ATTR]
+        self.embed_norm = model.embed_norm
+        self.qkv_norm = model.layers[0].pre_qkv_proj_norm
+        self.last_join = model.layers[-1].add_ffn
+
+    def takes(
+        self, input_ids: Optional[torch.Tensor], inputs_embeds: Optional[torch.Tensor]
+    ) -> bool:
+        """Whether this call's input boundary runs fused (shapes and dtypes only)."""
+        return (
+            inputs_embeds is None
+            and input_ids is not None
+            and input_ids.dim() == 1
+            and 1 <= input_ids.shape[0] <= _replicated.MAX_ROWS
+            and _input.takes(self.table, input_ids, self.embed_norm.weight, self.qkv_norm.weight)
+        )
+
+    def input(self, input_ids: torch.Tensor):
+        """``(FP32 residual, layer 0's pre-QKV norm output)`` of the step, one launch."""
+        return _input.input_boundary(
+            self.table,
+            input_ids,
+            self.embed_norm.weight,
+            self.embed_norm.variance_epsilon,
+            self.qkv_norm.weight,
+            self.qkv_norm.variance_epsilon,
+        )
+
+    @contextlib.contextmanager
+    def armed(self):
+        """Layer 61's join stores BF16 for one decode forward; disarmed on any exit."""
+        self.last_join._bf16_out = True
+        try:
+            yield
+        finally:
+            self.last_join._bf16_out = False
+
+
+def _check(model, captured: FrozenSet[int]) -> None:
+    """``install``'s checks. Raises ValueError naming the first that fails."""
+    reason = _input.available()
+    if reason is not None:
+        raise ValueError(f"[musespark_open] model_boundary: {reason}")
+    if _replicated.TABLE_ATTR not in model.embed_tokens.__dict__:
+        raise ValueError(
+            "[musespark_open] model_boundary gathers from the replicated table; replicated_embedding "
+            "is not installed"
+        )
+    if model.m_emb != 1.0:
+        raise ValueError(
+            f"[musespark_open] model_boundary is built for m_emb == 1, got {model.m_emb}"
+        )
+    layers = list(model.layers)
+    if layers[0].add_attn._preqkv_norm_fp32:
+        raise ValueError(
+            "[musespark_open] model_boundary computes layer 0's pre-QKV norm on the BF16 residual; "
+            "MUSESPARK_OPEN_PREQKV_NORM_FP32_RESIDUAL selects the FP32 ordering"
+        )
+    for name, norm in (
+        ("embed_norm", model.embed_norm),
+        ("pre_qkv_proj_norm", layers[0].pre_qkv_proj_norm),
+    ):
+        if norm.use_cuda_tile or norm.is_nvfp4 or norm.weight.dtype != torch.bfloat16:
+            raise ValueError(
+                f"[musespark_open] model_boundary: {name} is not flashinfer's BF16 RMSNorm here"
+            )
+    if model.embed_norm.use_gemma or not layers[0].pre_qkv_proj_norm.use_gemma:
+        raise ValueError(
+            "[musespark_open] model_boundary expects a plain embed_norm and a gemma pre_qkv_proj_norm"
+        )
+    last = len(layers) - 1
+    if last in captured:
+        raise ValueError(
+            f"[musespark_open] model_boundary: layer {last} is a DFlash capture layer; its FP32 "
+            "residual stream has a reader"
+        )
+    join = layers[-1].add_ffn
+    if not isinstance(join, _builtin.MuseSparkOpenGatedResidual):
+        raise ValueError(
+            f"[musespark_open] model_boundary: layer {last}'s post-FFN join is {type(join).__name__}"
+        )
+
+
+def install(model, captured: FrozenSet[int], mapping=None) -> ModelBoundary:
+    """Build the input-boundary kernel and re-class layer 61's join (agreed over ``mapping``'s TP group)."""
+    reason = None
+    try:
+        _check(model, captured)
+        # Compile now, with the step's id dtype (int32), so a compile failure is agreed here.
+        _input._compile(_input.HIDDEN, _input.EXIT, "int32", torch.cuda.current_device())
+    except Exception as error:  # noqa: BLE001 - every failure must reach the agreement
+        reason = str(error)
+    reason = _tp_agreement.agreed_refusal(mapping, reason)
+    if reason is not None:
+        raise ValueError(reason)
+    join = model.layers[-1].add_ffn
+    join.__class__ = _bf16_join_class(type(join))
+    join._bf16_out = False
+    boundary = ModelBoundary(model)
+    model._model_boundary = boundary
+    logger.info_once(
+        "[musespark_open] model_boundary: decode input boundary in one launch, layer "
+        f"{len(model.layers) - 1} join stores BF16",
+        key="musespark_open_model_boundary",
+    )
+    return boundary
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py
new file mode 100644
index 000000000000..4b60dea0e363
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py
@@ -0,0 +1,309 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The decode step's input boundary as one kernel: embedding gather, ``embed_norm``, residual, layer 0's pre-QKV norm.
+
+On a decode step ``DecodeTarget`` runs, before layer 0's QKV projection:
+
+1. ``replicated_embedding_gather`` (``ops/embedding/replicated_gather.py``): one row of the
+   replicated table per id, -0.0 written as +0.0, a +0.0 row for an id outside the table;
+2. ``embed_norm``: flashinfer's CuTe DSL ``RMSNormKernel`` (``weight_bias = 0``) on the BF16 rows;
+3. ``m_emb`` (1.0 for this checkpoint, so no kernel) and the FP32 copy, the residual stream;
+4. layer 0's ``pre_qkv_proj_norm`` on ``residual.to(bf16)`` -- the BF16 rows of step 2 again,
+   since they are exactly representable -- flashinfer's ``RMSNormKernel`` with
+   ``weight_bias = 1`` (gemma).
+
+That is five launches (a gather, two norms, two copies) of 8..128 CTAs each. This kernel does all
+of it in one launch of one 128-thread CTA per row: it gathers the row as step 1 does, then runs
+flashinfer's ``RMSNormKernel`` body twice on the same register fragment -- the same thread-value
+layout (``_make_tv_layout``), the same per-thread ``TensorSSA.reduce``, the same
+``row_reduce_sum_multirow`` (imported from flashinfer, not copied), the same ``rsqrt`` and the same
+product order ``x * rstd * (w + weight_bias)`` -- and stores the FP32 residual between the two.
+So every output is the five-launch chain's, bit for bit, for every table value the
+replicated-embedding install admits.
+
+PDL: every thread waits (``griddepcontrol.wait``) before its first global read (the ids are the
+step's input, written by the kernels before this one). The served trigger is ``"exit"``,
+flashinfer's placement after the last store, so every dependent the replaced chain's last norm
+had sees the same launch order. ``"entry"`` releases the dependent at once; it is correct only
+when every kernel that may launch behind this one waits before reading its outputs.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Optional
+
+import torch
+
+try:
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import BFloat16, Float32, Int32, Int64
+    from flashinfer.norm.kernels.rmsnorm import RMSNormKernel
+    from flashinfer.norm.utils import row_reduce_sum_multirow
+
+    HAS_CUTE_DSL = True
+except ImportError:  # pragma: no cover - the served image has the CuTe DSL and flashinfer
+    HAS_CUTE_DSL = False
+
+#: Width of the residual stream this kernel is compiled for.
+HIDDEN = 8192
+#: Where the grid lets its dependent launch.
+EXIT, ENTRY = 0, 1
+TRIGGERS = {"exit": EXIT, "entry": ENTRY}
+_ID_DTYPES = {torch.int32: "int32", torch.int64: "int64"}
+
+
+if HAS_CUTE_DSL:
+
+    class _InputBoundary:
+        """One CTA per row; flashinfer ``RMSNormKernel``'s geometry at ``HIDDEN`` columns."""
+
+        def __init__(self, hidden: int, trigger: int, id_dtype):
+            # The geometry flashinfer picks for a BF16 row of this width on this device.
+            ref = RMSNormKernel(BFloat16, hidden, 0.0)
+            if ref.cluster_n != 1 or ref.rows_per_block != 1 or ref.cols_per_tile != hidden:
+                raise ValueError(
+                    f"[musespark_open] model_boundary: flashinfer's RMSNormKernel at H={hidden} uses "
+                    f"cluster_n={ref.cluster_n}, rows_per_block={ref.rows_per_block}, "
+                    f"cols_per_tile={ref.cols_per_tile}; the input boundary is built for one CTA "
+                    "per row whose tile is the whole row"
+                )
+            self.hidden = hidden
+            self.trigger = trigger
+            self.id_dtype = id_dtype
+            self.vec_size = ref.vec_size
+            self.copy_bits = ref.copy_bits
+            self.threads_per_row = ref.threads_per_row
+            self.num_threads = ref.num_threads
+            self.rows_per_block = ref.rows_per_block
+            self.warps_per_row = ref.warps_per_row
+            self.num_vec_blocks = ref.num_vec_blocks
+            self.cols_per_tile = ref.cols_per_tile
+
+        @cute.jit
+        def __call__(self, mT, mIds, mWe, mWq, mR, mY, eps_e: Float32, eps_q: Float32, stream):
+            tv_shape, tv_stride = RMSNormKernel._make_tv_layout(
+                self.threads_per_row, self.rows_per_block, self.vec_size, self.num_vec_blocks
+            )
+            tv_layout = cute.make_layout(tv_shape, stride=tv_stride)
+            tiler_mn = (self.rows_per_block, self.cols_per_tile)
+            self.kernel(mT, mIds, mWe, mWq, mR, mY, eps_e, eps_q, tv_layout, tiler_mn).launch(
+                grid=[mIds.shape[0], 1, 1],
+                block=[self.num_threads, 1, 1],
+                smem=self.rows_per_block * self.warps_per_row * 4,
+                stream=stream,
+                use_pdl=True,
+            )
+
+        @cute.kernel
+        def kernel(
+            self,
+            mT: cute.Tensor,
+            mIds: cute.Tensor,
+            mWe: cute.Tensor,
+            mWq: cute.Tensor,
+            mR: cute.Tensor,
+            mY: cute.Tensor,
+            eps_e: Float32,
+            eps_q: Float32,
+            tv_layout: cute.Layout,
+            tiler_mn: cute.Shape,
+        ):
+            tidx, _, _ = cute.arch.thread_idx()
+            row, _, _ = cute.arch.block_idx()
+            if cutlass.const_expr(self.trigger == ENTRY):
+                cute.arch.griddepcontrol_launch_dependents()
+            # Before any global read: the ids are written by the kernels before this one.
+            cute.arch.griddepcontrol_wait()
+
+            H = self.hidden
+            threads_per_row = tv_layout.shape[0][0]
+            rows_per_block = tiler_mn[0]
+            warps_per_row = max(threads_per_row // 32, 1)
+
+            smem = cutlass.utils.SmemAllocator()
+            reduction_buffer = smem.allocate_tensor(
+                Float32, cute.make_layout((rows_per_block, warps_per_row)), byte_alignment=4
+            )
+
+            # The gather's row choice (ops/embedding/replicated_gather.py): an id outside the
+            # table reads nothing.
+            token = Int64(mIds[row])
+            valid = (token >= Int64(0)) & (token < Int64(mT.shape[0]))
+            source = Int64(0)
+            if valid:
+                source = token
+
+            gX = cute.local_tile(mT, tiler_mn, (source, 0))
+            gR = cute.local_tile(mR, tiler_mn, (row, 0))
+            gY = cute.local_tile(mY, tiler_mn, (row, 0))
+            we_layout = cute.prepend(mWe.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+            wq_layout = cute.prepend(mWq.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+            gWe = cute.local_tile(cute.make_tensor(mWe.iterator, we_layout), tiler_mn, (0, 0))
+            gWq = cute.local_tile(cute.make_tensor(mWq.iterator, wq_layout), tiler_mn, (0, 0))
+
+            copy_atom = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=self.copy_bits
+            )
+            tiled_copy = cute.make_tiled_copy(copy_atom, tv_layout, tiler_mn)
+            thr_copy = tiled_copy.get_slice(tidx)
+
+            tXgX = thr_copy.partition_S(gX)
+            tXrX = cute.make_fragment_like(tXgX)
+            tWgWe = thr_copy.partition_S(gWe)
+            tWrWe = cute.make_fragment_like(tWgWe)
+            tWgWq = thr_copy.partition_S(gWq)
+            tWrWq = cute.make_fragment_like(tWgWq)
+            tXgR = thr_copy.partition_D(gR)
+            tXgY = thr_copy.partition_D(gY)
+            tXrY = cute.make_fragment_like(tXgY)
+            tXrR = cute.make_fragment_like(tXgY, Float32)
+            tXrS = cute.make_fragment_like(tXgY, Float32)
+
+            tXrX.store(cute.zeros_like(tXrX, dtype=BFloat16))
+            if valid:
+                cute.copy(copy_atom, tXgX, tXrX)
+            cute.copy(copy_atom, tWgWe, tWrWe)
+            cute.copy(copy_atom, tWgWq, tWrWq)
+
+            # The gather's sign rule: a zero of either sign is +0.0.
+            tXrS.store(tXrX.load().to(Float32))
+            for i in cutlass.range_constexpr(cute.size(tXrS)):
+                if tXrS[i] == Float32(0.0):
+                    tXrS[i] = Float32(0.0)
+
+            # embed_norm: flashinfer RMSNormKernel, weight_bias 0.
+            x = tXrS.load()
+            x_sq = x * x
+            sum_sq = row_reduce_sum_multirow(x_sq, threads_per_row, reduction_buffer, None, 1)
+            mean_sq = sum_sq / Float32(H)
+            rstd = cute.math.rsqrt(mean_sq + eps_e, fastmath=True)
+            cute.arch.barrier()
+            w = tWrWe.load().to(Float32)
+            embeds = (x * rstd * (w + Float32(0.0))).to(BFloat16)
+
+            # The FP32 residual stream (m_emb == 1).
+            tXrR.store(embeds.to(Float32))
+            cute.autovec_copy(tXrR, tXgR)
+
+            # Layer 0's pre_qkv_proj_norm on residual.to(bf16) == embeds: RMSNormKernel, weight_bias 1.
+            x2 = embeds.to(Float32)
+            x2_sq = x2 * x2
+            sum_sq2 = row_reduce_sum_multirow(x2_sq, threads_per_row, reduction_buffer, None, 1)
+            mean_sq2 = sum_sq2 / Float32(H)
+            rstd2 = cute.math.rsqrt(mean_sq2 + eps_q, fastmath=True)
+            wq = tWrWq.load().to(Float32)
+            tXrY.store((x2 * rstd2 * (wq + Float32(1.0))).to(BFloat16))
+            cute.copy(copy_atom, tXrY, tXgY)
+
+            if cutlass.const_expr(self.trigger == EXIT):
+                cute.arch.griddepcontrol_launch_dependents()
+
+    @functools.lru_cache(maxsize=None)
+    def _compile(hidden: int, trigger: int, id_dtype: str, device_index: int):
+        ids_type = Int32 if id_dtype == "int32" else Int64
+        rows = cute.sym_int64()
+        vocab = cute.sym_int64()
+        table = cute.runtime.make_fake_compact_tensor(
+            BFloat16, (vocab, hidden), stride_order=(1, 0), assumed_align=16
+        )
+        ids = cute.runtime.make_fake_compact_tensor(ids_type, (rows,), assumed_align=4)
+        weight_e = cute.runtime.make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16)
+        weight_q = cute.runtime.make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16)
+        residual = cute.runtime.make_fake_compact_tensor(
+            Float32, (rows, hidden), stride_order=(1, 0), assumed_align=16
+        )
+        out = cute.runtime.make_fake_compact_tensor(
+            BFloat16, (rows, hidden), stride_order=(1, 0), assumed_align=16
+        )
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        with torch.cuda.device(device_index):
+            return cute.compile(
+                _InputBoundary(hidden, trigger, ids_type),
+                table,
+                ids,
+                weight_e,
+                weight_q,
+                residual,
+                out,
+                Float32(1e-6),
+                Float32(1e-6),
+                stream,
+                options="--enable-tvm-ffi",
+            )
+
+
+def available() -> Optional[str]:
+    """Why this kernel cannot run here, or ``None``."""
+    return (
+        None if HAS_CUTE_DSL else "the CuTe DSL or flashinfer's CuTe DSL RMSNorm is not importable"
+    )
+
+
+def takes(
+    table: torch.Tensor, ids: torch.Tensor, embed_weight: torch.Tensor, qkv_weight: torch.Tensor
+) -> bool:
+    """Whether a call fits the kernel (shapes, dtypes, contiguity; no pointer values)."""
+    return (
+        table.is_cuda
+        and table.dtype == torch.bfloat16
+        and table.dim() == 2
+        and table.shape[1] == HIDDEN
+        and table.is_contiguous()
+        and ids.device == table.device
+        and ids.dtype in _ID_DTYPES
+        and ids.dim() == 1
+        and ids.shape[0] > 0
+        and ids.is_contiguous()
+        and all(
+            w.device == table.device
+            and w.dtype == torch.bfloat16
+            and tuple(w.shape) == (HIDDEN,)
+            and w.is_contiguous()
+            for w in (embed_weight, qkv_weight)
+        )
+    )
+
+
+def input_boundary(
+    table: torch.Tensor,
+    ids: torch.Tensor,
+    embed_weight: torch.Tensor,
+    embed_eps: float,
+    qkv_weight: torch.Tensor,
+    qkv_eps: float,
+    trigger: str = "exit",
+) -> tuple:
+    """``(residual FP32 [rows, H], layer 0's pre-QKV norm output BF16 [rows, H])`` in one launch.
+
+    ``table`` is the whole replicated ``[vocab, H]`` BF16 embedding, ``ids`` the step's token
+    ids, ``embed_weight`` ``embed_norm``'s (unit) weight and ``qkv_weight`` layer 0's gemma
+    ``pre_qkv_proj_norm`` weight. The caller checks ``takes`` first.
+    """
+    rows = ids.shape[0]
+    residual = torch.empty((rows, HIDDEN), dtype=torch.float32, device=table.device)
+    out = torch.empty((rows, HIDDEN), dtype=torch.bfloat16, device=table.device)
+    _compile(HIDDEN, TRIGGERS[trigger], _ID_DTYPES[ids.dtype], table.device.index)(
+        table, ids, embed_weight, qkv_weight, residual, out, float(embed_eps), float(qkv_eps)
+    )
+    return residual, out
+
+
+def reference(
+    table: torch.Tensor,
+    ids: torch.Tensor,
+    embed_weight: torch.Tensor,
+    embed_eps: float,
+    qkv_weight: torch.Tensor,
+    qkv_eps: float,
+) -> tuple:
+    """The five-launch chain this kernel replaces, through the served ops."""
+    from tensorrt_llm._torch.custom_ops import flashinfer_gemma_rmsnorm, flashinfer_rmsnorm
+
+    from ...ops.embedding.replicated_gather import replicated_embedding_gather
+
+    embeds = replicated_embedding_gather(table, ids)
+    embeds = flashinfer_rmsnorm(embeds, embed_weight, embed_eps)
+    residual = embeds.to(torch.float32)
+    return residual, flashinfer_gemma_rmsnorm(residual.to(torch.bfloat16), qkv_weight, qkv_eps)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py
index a807e54158b0..9683a844e172 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py
@@ -83,6 +83,8 @@
 from ...ops.attention.multi_request_prep_append import (
     MAX_ROWS_PER_SEQ,
     MAX_SEQS,
+    PACKED_GEOMETRY,
+    launch_kind,
     multi_request_prep_append,
 )
 
@@ -690,10 +692,17 @@ class each layer carries by then, and acts only on decode steps of ``min_seqs``
         attention._mr_append_folded = None
         attention._mr_append_census = census
         attention._mr_append_min_seqs = min_seqs
+    grid = (
+        "packed grid (last rows, items per program): "
+        f"{[(last, items) for last, items, _, _ in PACKED_GEOMETRY]}, one item per program on "
+        "o_proj-prefetching steps"
+        if not torch.cuda.is_available() or launch_kind(torch.device("cuda")) == "packed"
+        else "grid of one program per (row, head)"
+    )
     logger.info(
         f"[musespark_open] multi_request_prep_append: the K/V append, E4M3 query and FMHA scales "
         f"of {len(attentions)} layers run in the decode prep on steps of {min_seqs}..{MAX_SEQS} sequences "
         f"of up to {MAX_ROWS_PER_SEQ} rows each; their trtllm-gen generation preprocess launch "
-        "(decoder-info + append) is skipped"
+        f"(decoder-info + append) is skipped; {grid}"
     )
     return len(attentions)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py
new file mode 100644
index 000000000000..b1ea75c49827
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""prefill_cutedsl_experts: large prefill expert calls on the Rubin CuTe DSL grouped GEMMs."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py
new file mode 100644
index 000000000000..2e258d81509c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py
@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""``prefill_cutedsl_experts``: the NVFP4 layers' large expert calls on the Rubin CuTe DSL GEMMs (sm107).
+
+Installed on Rubin (sm107) at tp 4 only (``serves``); every other GPU and tp size keeps
+the TRTLLM-gen expert calls. Every NVFP4 MoE layer whose leaf is
+the native TRTLLM-gen one on this rank's contiguous expert block
+(``ring_buddy.eligible``) gets a wrapper around its expert forward: an
+unfinalized call (``do_finalize=False``) of ``MIN_ROWS..MAX_ROWS`` NVFP4 rows --
+the prefill and mixed steps of at least 4096 tokens -- runs
+``ops/moe/prefill_cutedsl_experts`` on the leaf's own tensors and returns the
+same ``(gemm2_output, weights, expanded_idx_to_permuted_idx)`` triple the
+TRTLLM-gen call returns, which the model's per-expert post-norm finalize then
+consumes. Every other call (decode steps, smaller prefill chunks, the BF16 layers
+0 and 61) keeps the expert forward it had.
+
+The routing is the routing method's own ``apply`` (memoized on the logits, so the
+finalize reads the same weights). Rounding-different from the TRTLLM-gen call;
+expert selection is unchanged.
+
+Install after the weights are final and before any CUDA graph capture: the
+kernels are compiled here (``PrefillCuteDslExperts.prepare``), so no capture
+compiles them. Agreed over the TP group (``agreed_install``).
+"""
+
+from __future__ import annotations
+
+import types
+from typing import List, Optional, Tuple
+
+import torch
+
+from tensorrt_llm._torch.utils import Fp4QuantizedTensor
+from tensorrt_llm.logger import logger
+
+from ... import dispatch as _tp_agreement
+from ...ops.moe import prefill_cutedsl_experts as _op
+from ...ops.moe.ring_buddy_assign import EXPERTS, TOP_K
+from ..replicated_experts import ring_buddy as _ring_buddy
+
+#: The SM version and tp size this path is installed at: Rubin (sm107) tep4.
+SERVED_SM = 107
+SERVED_TP_SIZE = 4
+
+#: Rows of the expert calls this path serves.
+MIN_ROWS = 4096
+MAX_ROWS = 8192
+
+
+def serves(compute_capability: Optional[Tuple[int, int]], tp_size: int) -> bool:
+    """Whether a model on ``compute_capability`` at ``tp_size`` installs this path."""
+    return _tp_agreement.sm_of(compute_capability) == SERVED_SM and tp_size == SERVED_TP_SIZE
+
+
+def _refusal(moe, mapping) -> Optional[str]:
+    reason = _ring_buddy.eligible(moe, mapping)
+    if reason is not None:
+        return reason
+    if _ring_buddy.expert_format(moe) != "nvfp4":
+        return "the experts are not the native NVFP4 TRTLLM-gen leaf"
+    backend = _ring_buddy._backend(moe)
+    if not _ring_buddy._is_plain_swiglu(backend):
+        return "the experts are not plain SwiGLU"
+    if getattr(backend, "fc2_input_scale", None) is None:
+        return "the leaf has no FC2 input scale"
+    return None
+
+
+def _wrap(moe, op: "_op.PrefillCuteDslExperts") -> None:
+    experts = moe.experts
+    prior = experts.forward
+    routing = experts.routing_method
+
+    def forward(self, x, router_logits, *args, do_finalize=True, **kwargs):
+        if (
+            do_finalize
+            or args
+            or not isinstance(x, Fp4QuantizedTensor)
+            or x.is_sf_swizzled
+            or not MIN_ROWS <= x.fp4_tensor.shape[0] <= MAX_ROWS
+            or torch.compiler.is_compiling()
+        ):
+            return prior(x, router_logits, *args, do_finalize=do_finalize, **kwargs)
+        selected, weights = routing.apply(router_logits)
+        gemm2_output, expanded_idx_to_permuted_idx = op(
+            x.fp4_tensor, x.scaling_factor, selected, weights
+        )
+        return gemm2_output, weights, expanded_idx_to_permuted_idx
+
+    experts.forward = types.MethodType(forward, experts)
+
+
+def install(text_model, mapping) -> List[int]:
+    """Wrap every eligible NVFP4 layer's expert forward; returns the layer indices.
+
+    Collective over the TP group (the outcome is agreed). Raises ValueError on every
+    rank alike when some rank refuses (not sm107, no eligible layer) or cannot
+    compile the kernels.
+    """
+
+    def build() -> List[int]:
+        if torch.cuda.get_device_capability() != (10, 7):
+            raise ValueError(
+                "[musespark_open] prefill_cutedsl_experts is written for sm107 (Rubin) only"
+            )
+        installed, skipped, ops = [], {}, []
+        for index, layer in enumerate(text_model.layers):
+            moe = layer.feed_forward
+            reason = _refusal(moe, mapping)
+            if reason is not None:
+                skipped[index] = reason
+                continue
+            backend = _ring_buddy._backend(moe)
+            op = _op.PrefillCuteDslExperts(backend, EXPERTS, TOP_K, backend.slot_start)
+            ops.append((moe, op))
+            installed.append(index)
+        if not installed:
+            raise ValueError(
+                "[musespark_open] prefill_cutedsl_experts found no eligible NVFP4 layer"
+            )
+        try:
+            # One compile serves every layer (the kernels' shapes are runtime arguments).
+            ops[0][1].prepare(MAX_ROWS)
+            torch.cuda.synchronize()
+            # Hand the compile call's transient buffers back before the KV cache is sized.
+            torch.cuda.empty_cache()
+        except Exception as error:  # noqa: BLE001 - the compile's failure must reach the agreement
+            raise RuntimeError(
+                f"[musespark_open] prefill_cutedsl_experts could not compile: {error!r}"
+            ) from error
+        for moe, op in ops:
+            _wrap(moe, op)
+            moe._prefill_cutedsl_experts = op
+        logger.info(
+            f"[musespark_open] prefill_cutedsl_experts: {len(installed)} NVFP4 MoE layers run "
+            f"their {MIN_ROWS}..{MAX_ROWS}-row expert calls on the Rubin CuTe DSL GEMMs "
+            f"(tile {_op.TILE}); built-in expert call kept for layers {sorted(skipped)}: "
+            f"{sorted(set(skipped.values()))}"
+        )
+        return installed
+
+    return _tp_agreement.agreed_install(mapping, build, (ValueError, RuntimeError))
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py
index bd29c483e9ac..0535bd2926ab 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py
@@ -458,6 +458,9 @@ def _nvfp4_call(backend, routing, persistent: bool = False, stacks=None, persist
     from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import (  # noqa: E501
         fp4_block_scale_moe_runner,
     )
+    from tensorrt_llm._torch.autotuner import AutoTuner
+
+    from ...ops.moe import decode_expert_tactics as _decode_tactics
 
     intermediate = backend.w3_w1_weight.shape[-2] // 2
     tune_max_num_tokens = backend.max_num_tokens
@@ -521,6 +524,48 @@ def call(stacks, x, local_ids, weights_bf16):
                 stacks["fc31_alpha"],
                 stacks["fc2_alpha"],
             )
+        tactic = None
+        if not AutoTuner.get().is_tuning_mode:
+            tactic = _decode_tactics.fixed_tactic(
+                x.fp4_tensor.device,
+                act_type,
+                TOP_K,
+                2 * x.fp4_tensor.shape[1],
+                intermediate,
+                window,
+                x.fp4_tensor.shape[0],
+            )
+        if tactic is not None:
+            return _decode_tactics.run_fixed(
+                tactic,
+                [
+                    None,
+                    None,
+                    x.fp4_tensor,
+                    x_sf,
+                    stacks["w3_w1_weight"],
+                    stacks["w3_w1_weight_scale"].view(torch.float8_e4m3fn),
+                    None,
+                    None,
+                    None,
+                    None,
+                    stacks["w2_weight"],
+                    stacks["w2_weight_scale"].view(torch.float8_e4m3fn),
+                    None,
+                    stacks["fc31_scale_c"],
+                    stacks["fc31_alpha"],
+                    stacks["fc2_alpha"],
+                    weights_bf16,
+                    local_ids,
+                ],
+                EXPERTS,
+                TOP_K,
+                intermediate,
+                window,
+                routing.routing_method_type,
+                act_type,
+                tune_max_num_tokens,
+            )
         return fp4_block_scale_moe_runner(
             None,
             None,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py
index 95be6fda9d1c..db66937ca604 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py
@@ -16,6 +16,14 @@
   ``[8, shard]`` output straight into this rank's slot of the multicast inbox
   (``HeadGatherWorkspace``) and the sampler's first kernel gathers the row
   from the inbox (``gather_split_sample``), instead of the NCCL all-gather.
+* ``vocab_parallel_head`` -- on decode steps of ``head_sampler.VP_MIN_ROWS`` (16) rows or more,
+  no gather at all: the processor runs only the local LM-head GEMM and leaves a
+  ``[rows, vocab]`` BF16 placeholder pending (allocated, not written); the worker samples
+  the step from the ranks' shards (``head_sampler.vocab_parallel_sample``: each rank folds
+  its own shard, the ranks exchange per-virtual-thread partial states and crossings over a
+  multicast buffer) and argmaxes greedy batches the same way
+  (``vocab_parallel_argmax``, ``torch.argmax``'s token exactly). Every other consumer
+  gathers the shards with the built-in all-gather first (``materialize``).
 
 The processor instance is re-classed as ``SplitSamplerLogitsProcessor`` (only the
 target's: the drafter has its own processor), and the worker as a class
@@ -191,6 +199,11 @@ def context_kv_head_fork(lm_head) -> tuple:
 class TargetHead:
     """The state the re-classed processor and worker share: the step's pending BF16 row."""
 
+    #: ``vocab_parallel_head``'s fabric buffers (None: the head gathers every row).
+    vp: Optional[_hs.VocabParallelWorkspace] = None
+    #: This rank's ``[rows, shard]`` LM-head output while a vocab-parallel row is pending.
+    pending_local: Optional[torch.Tensor] = None
+
     def __init__(
         self,
         processor: MuseSparkOpenLogitsProcessor,
@@ -198,6 +211,7 @@ def __init__(
         worker: DFlashWorker,
         mapping,
         gather: bool,
+        vocab_parallel: bool = False,
     ) -> None:
         self.multiplier = float(processor.logit_multiplier)
         self.soft_cap = float(processor.soft_cap)
@@ -219,9 +233,19 @@ def __init__(
         self.max_rows = _hs.MAX_ROWS
         self.states = _hs.new_states(self.max_rows, self.device)
         self.wide_states: Optional[torch.Tensor] = None
+        # vocab_parallel_head: the fabric buffers for steps of VP_MIN_ROWS..max_rows rows
+        # (collective; rebuilt for the wider steps by ``widen``).
+        self.mapping = mapping
+        self.lm_head = lm_head
+        self.vp = None
+        if vocab_parallel:
+            self.vp = _hs.VocabParallelWorkspace(
+                mapping, int(lm_head.weight.shape[0]), self.max_rows
+            )
         self.load_seconds = time.perf_counter() - start
         self.pending: Optional[torch.Tensor] = None
         self.pending_inbox = False
+        self.pending_local = None
         self._bypass_logged = False
         self._wide_capture_logged = False
         self.hook_handle = None
@@ -298,6 +322,17 @@ def produce(self, hidden_states: torch.Tensor, lm_head: Linear) -> torch.Tensor:
                 "drained within its step"
             )
         rows = hidden_states.shape[0]
+        if self.vp is not None and _hs.VP_MIN_ROWS <= rows <= self.vp.max_rows:
+            # The built-in Linear.forward's local GEMM (column-parallel: apply_linear, then the
+            # all-gather this path leaves out), through the instance's apply_linear.
+            local = lm_head.apply_linear(hidden_states, lm_head.bias, None, None)
+            logits = torch.empty(
+                (rows, self.vp.vocab), dtype=torch.bfloat16, device=hidden_states.device
+            )
+            self.pending_local = local
+            self.pending_inbox = False
+            self.pending = logits
+            return logits
         if self.workspace is not None and rows == GATHER_ROWS:
             admitted, fork = context_kv_head_fork(lm_head)
             if not admitted:
@@ -347,9 +382,26 @@ def sample(
         offset: Optional[torch.Tensor],
     ) -> torch.Tensor:
         """The fused sampler's tokens for the pending row; consumes it."""
+        local = self.pending_local
         inbox = self.pending_inbox
         self.pending = None
         self.pending_inbox = False
+        self.pending_local = None
+        if local is not None:
+            return _hs.vocab_parallel_sample(
+                self.vp.view,
+                logits,
+                local,
+                temperatures,
+                top_ks,
+                top_ps,
+                min_ps,
+                seed,
+                offset,
+                self.multiplier,
+                self.soft_cap,
+                self.pdl,
+            )
         if inbox:
             ws = self.workspace
             return _hs.gather_split_sample(
@@ -384,7 +436,29 @@ def sample(
             self.pdl,
         )
 
-    def _gather_now(self, logits: torch.Tensor) -> None:
+    def vocab_parallel(self, logits) -> bool:
+        """Whether ``logits`` is the pending row of a vocab-parallel step."""
+        return self.owns(logits) and self.pending_local is not None
+
+    def argmax(self, logits: torch.Tensor) -> torch.Tensor:
+        """``torch.argmax`` of the built-in capped row of a vocab-parallel pending row (int64); consumes it."""
+        local = self.pending_local
+        self.pending = None
+        self.pending_local = None
+        return _hs.vocab_parallel_argmax(
+            self.vp.view, local, self.multiplier, self.soft_cap, self.pdl
+        )
+
+    def _gather_now(self, logits: torch.Tensor) -> Optional[torch.Tensor]:
+        """Consume the pending row; returns the gathered BF16 row of a vocab-parallel one."""
+        if self.pending_local is not None:
+            from tensorrt_llm._torch.distributed import allgather
+
+            local = self.pending_local
+            self.pending = None
+            self.pending_local = None
+            # The built-in Linear.forward's gather of the column-parallel output.
+            return allgather(local, self.lm_head.mapping, sizes=self.lm_head.gather_output_sizes)
         if self.pending_inbox:
             ws = self.workspace
             empty = ws.status  # unread by the gather-only kernel
@@ -408,15 +482,22 @@ def _gather_now(self, logits: torch.Tensor) -> None:
             )
         self.pending = None
         self.pending_inbox = False
+        return None
 
     def materialize(self, logits: torch.Tensor) -> torch.Tensor:
         """The built-in processor's FP32 capped logits for the pending row; consumes it."""
-        self._gather_now(logits)
-        return _hs.reference_capped(logits, self.multiplier, self.soft_cap)
+        gathered = self._gather_now(logits)
+        return _hs.reference_capped(
+            logits if gathered is None else gathered, self.multiplier, self.soft_cap
+        )
 
     def drain(self) -> None:
-        """Consume a pending row nothing sampled (resets its inbox words)."""
-        if self.pending is not None:
+        """Consume a pending row nothing sampled (resets its inbox words; a vocab-parallel row
+        has none, and is dropped without a gather)."""
+        if self.pending_local is not None:
+            self.pending = None
+            self.pending_local = None
+        elif self.pending is not None:
             self._gather_now(self.pending)
 
 
@@ -475,6 +556,13 @@ def _sample_tokens_for_batch(self, logits, spec_metadata, num_contexts, batch_si
         head = self._split_sampler_head
         head.check_routed(logits, "_sample_tokens_for_batch")
         if head.owns(logits):
+            if (
+                spec_metadata is not None
+                and spec_metadata.is_all_greedy_sample
+                and head.vocab_parallel(logits)
+            ):
+                # The built-in greedy branch: torch.argmax of the capped row.
+                return head.argmax(logits)
             if (
                 spec_metadata is not None
                 and not spec_metadata.is_all_greedy_sample
@@ -635,6 +723,13 @@ def widen(llm) -> bool:
     if head.pending is not None:
         raise RuntimeError("[musespark_open] wide_split_sampler: a head row is pending")
     head.max_rows = _hs.WIDE_MAX_ROWS
+    if head.vp is not None:
+        # Collective: every rank widens at the same point of the load.
+        head.vp = _hs.VocabParallelWorkspace(head.mapping, head.vp.layout.shard, head.max_rows)
+        logger.info(
+            f"[musespark_open] vocab_parallel_head: decode steps of {_hs.VP_MIN_ROWS}..{head.max_rows} "
+            f"rows; {head.vp.nbytes() / 2**20:.1f} MiB of fabric buffers on this rank"
+        )
     logger.info(
         f"[musespark_open] wide_split_sampler: the split sampler takes decode steps of up to "
         f"{head.max_rows} rows (the shipped 512-thread block above {_hs.MAX_ROWS}); its state "
@@ -643,7 +738,7 @@ def widen(llm) -> bool:
     return True
 
 
-def install(llm, mapping, gather: bool) -> None:
+def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None:
     """Install the split sampler on the target's processor and DFlash worker. Idempotent.
 
     With ``gather`` (``lm_head_inbox``) the LM-head GEMM also writes into the multicast
@@ -671,7 +766,7 @@ def install(llm, mapping, gather: bool) -> None:
     if reason is not None:
         raise ValueError(reason)
     processor, worker, lm_head = llm.logits_processor, llm.spec_worker, llm.lm_head
-    head = TargetHead(processor, lm_head, worker, mapping, gather)
+    head = TargetHead(processor, lm_head, worker, mapping, gather, vocab_parallel)
     processor.__class__ = SplitSamplerLogitsProcessor
     processor._split_head = head
     composed = attach_worker(worker, head)
@@ -683,6 +778,13 @@ def install(llm, mapping, gather: bool) -> None:
             if gather
             else ""
         )
+        + (
+            f"; vocab_parallel_head: decode steps of {_hs.VP_MIN_ROWS}..{head.max_rows} rows sample "
+            f"each rank's own shard (no logits all-gather), {head.vp.nbytes() / 2**20:.1f} MiB of "
+            "fabric buffers on this rank"
+            if head.vp is not None
+            else ""
+        )
         + f"; library {head.library} (load {head.load_seconds:.2f} s); "
         + f"worker class {composed.__name__}"
     )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py
index cbb858ad7d0a..1634b8830a70 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py
@@ -180,6 +180,9 @@ class SpecLoopFoldDFlashWorker(DFlashWorker):
 
     _fold_verify: bool = False
     _fold_path_walk_tail: bool = False
+    #: Its ``_apply_dflash2_selector`` applies the draft head's deferred drafter
+    #: logit transform (``draft_head.head._defers_transform``).
+    _draft_transform_aware = True
 
     # ------------------------------------------------------------------
     # The per-step state
@@ -736,7 +739,7 @@ def _apply_dflash2_selector(
         if (
             self._fold_path_walk_tail
             and gen_logits.shape[-1] != full_vocab
-            and self._shard_selector_applies(gen_logits, spec_metadata, full_vocab)
+            and self._draft_head_takes(draft_model, gen_logits, spec_metadata, full_vocab)
         ):
             args = self._path_walk_tail_args(step, gen_logits, spec_metadata)
         if args is None:
@@ -759,6 +762,7 @@ def walk(edges, unary, ids):
             self.mapping,
             self.__dict__.get("_draft_head_allreduce"),
             walk=walk,
+            transform=self._draft_head_transform_for(draft_model, gen_logits),
         )
         step.tail = _Tail(tokens, next_tokens, kv_lens is not None, batch_indices)
         if kv_lens is not None:
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py
index 85b29dcce163..8025d17ddf5f 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py
@@ -27,8 +27,10 @@
 the row count only, the same on every rank). Before wrapping it,
 ``tie_order_reason`` runs the rank top-k layout with every launch of this
 fast path against the running ``torch.topk`` on the draft head's tie-heavy
-probe (``draft_head.tie_order_probe_logits``) and refuses the install if they
-disagree.
+probes (``draft_head.tie_order_probes``) and refuses the install if they
+disagree. With the drafter's logit transform installed on the draft head
+(``draft_head/logit_transform.py``) ``keys`` takes its canonical keys, and the
+probes and their reference are the transformed ones.
 """
 
 from __future__ import annotations
@@ -42,8 +44,9 @@
 from tensorrt_llm.logger import logger
 
 from ...ops.spec.draft_rank_topk import MAX_KEYS, draft_rank_topk
-from ...ops.spec.draft_shard_topk import CHUNK
+from ...ops.spec.draft_shard_topk import CHUNK, FULL_BAND
 from ..draft_head import head as _draft_head
+from ..draft_head.logit_transform import DraftLogitTransform
 
 #: Fewest draft rows (requests x K) the exchange takes: five requests at K = 7. Fewer
 #: rows are the one-request path and the 9..32-row band.
@@ -67,9 +70,23 @@ def takes(self, rows: int) -> bool:
         """Whether a draft step of ``rows`` rows exchanges the rank top-k (``keys``)."""
         return rows >= MIN_ROWS
 
-    def keys(self, local: torch.Tensor, top_k: int, rank: int, num_ranks: int) -> torch.Tensor:
-        """This rank's own top ``top_k`` keys of ``local`` in the all-reduce layout."""
-        return draft_rank_topk(local, top_k, CHUNK, CHUNK_WARPS, rank, num_ranks)
+    def keys(
+        self,
+        local: torch.Tensor,
+        top_k: int,
+        rank: int,
+        num_ranks: int,
+        canon: Optional[torch.Tensor] = None,
+        canon_band: tuple[int, int] = FULL_BAND,
+    ) -> torch.Tensor:
+        """This rank's own top ``top_k`` keys of ``local`` in the all-reduce layout.
+
+        ``canon`` / ``canon_band``: a drafter logit transform's canonical keys
+        (``draft_shard_topk``'s), passed by ``draft_head.selector_draft_tokens``.
+        """
+        return draft_rank_topk(
+            local, top_k, CHUNK, CHUNK_WARPS, rank, num_ranks, canon=canon, canon_band=canon_band
+        )
 
     def merge_launch(self, rows: int) -> dict:
         """``draft_topk_merge``'s launch geometry for a step of ``rows`` rows."""
@@ -82,14 +99,21 @@ def __call__(self, packed: torch.Tensor) -> torch.Tensor:
         return self.allreduce(packed)
 
 
-def tie_order_reason(selector: nn.Module, num_ranks: int, topk=torch.topk) -> Optional[str]:
+def tie_order_reason(
+    selector: nn.Module,
+    num_ranks: int,
+    topk=torch.topk,
+    transform: Optional[DraftLogitTransform] = None,
+) -> Optional[str]:
     """Why the rank top-k exchange would differ from this torch's ``torch.topk``, or None.
 
-    ``draft_head.tie_order_reason``'s probe for the rank top-k layout: the
+    ``draft_head.tie_order_reason``'s probes for the rank top-k layout: the
     probe logits split into ``num_ranks`` shards, each shard's ``keys``, their
     sum (what the all-reduce returns) merged with every ``MERGE_LAUNCH``
     geometry, ids and values compared bit for bit with ``topk`` as the built-in
-    selector runs it. ``topk`` is a parameter so a test can stand in another order.
+    selector runs it. With ``transform`` (the draft head's) the keys are
+    canonical, the merge reads its values table and the reference transforms
+    the logits first. ``topk`` is a parameter so a test can stand in another order.
     """
     predecessor_codebook = selector.predecessor_codebook
     vocab, rank_dim = predecessor_codebook.shape
@@ -97,42 +121,48 @@ def tie_order_reason(selector: nn.Module, num_ranks: int, topk=torch.topk) -> Op
     top_k = selector.top_k
     device = predecessor_codebook.device
     exchange = RankTopKExchange(None)
+    canon = _draft_head.canon_kwargs(transform)
     with torch.cuda.device(device):
-        full = _draft_head.tie_order_probe_logits(vocab, shard).to(device)
-        rows = full.shape[0]
-        ref_values, ref_ids = topk(full.float().view(1, rows, vocab), top_k, dim=-1)
-        keys = torch.stack(
-            [
-                exchange.keys(
-                    full[:, r * shard : (r + 1) * shard].contiguous(), top_k, r, num_ranks
-                )
-                for r in range(num_ranks)
-            ]
-        ).sum(dim=0)
-        gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device)
-        anchor = torch.zeros(1, dtype=torch.int64, device=device)
-        for _, rank_block, num_warps in MERGE_LAUNCH:
-            ids, values, _, _ = _draft_head.draft_topk_merge(
-                keys,
-                shard,
-                1,
-                rows,
-                anchor,
-                gate,
-                predecessor_codebook,
-                selector.successor_codebook,
-                top_k,
-                num_ranks,
-                rank_block=rank_block,
-                num_warps=num_warps,
-            )
-            if not (torch.equal(ids, ref_ids) and _draft_head._same_bits(values, ref_values)):
-                differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist()
-                return (
-                    f"torch {torch.__version__} orders top-k ties differently from the merge over "
-                    f"rank top-k keys (rank_block {rank_block}, {num_warps} warps; probe rows "
-                    f"{differ} differ)"
+        for probe, full in enumerate(_draft_head.tie_order_probes(vocab, shard, transform, device)):
+            rows = full.shape[0]
+            ref_values, ref_ids = _draft_head.probe_reference(full, top_k, transform, topk)
+            keys = torch.stack(
+                [
+                    exchange.keys(
+                        full[:, r * shard : (r + 1) * shard].contiguous(),
+                        top_k,
+                        r,
+                        num_ranks,
+                        **canon,
+                    )
+                    for r in range(num_ranks)
+                ]
+            ).sum(dim=0)
+            gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device)
+            anchor = torch.zeros(1, dtype=torch.int64, device=device)
+            for _, rank_block, num_warps in MERGE_LAUNCH:
+                ids, values, _, _ = _draft_head.draft_topk_merge(
+                    keys,
+                    shard,
+                    1,
+                    rows,
+                    anchor,
+                    gate,
+                    predecessor_codebook,
+                    selector.successor_codebook,
+                    top_k,
+                    num_ranks,
+                    rank_block=rank_block,
+                    num_warps=num_warps,
+                    **_draft_head.values_kwargs(transform),
                 )
+                if not (torch.equal(ids, ref_ids) and _draft_head._same_bits(values, ref_values)):
+                    differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist()
+                    return (
+                        f"torch {torch.__version__} orders top-k ties differently from the merge "
+                        f"over rank top-k keys (rank_block {rank_block}, {num_warps} warps; probe "
+                        f"{probe} rows {differ} differ)"
+                    )
     return None
 
 
@@ -168,7 +198,11 @@ def install(llm: nn.Module) -> bool:
         return False
     reason = unsupported_reason(llm)
     if reason is None:
-        reason = tie_order_reason(llm.draft_model.candidate_selector, worker.mapping.tp_size)
+        reason = tie_order_reason(
+            llm.draft_model.candidate_selector,
+            worker.mapping.tp_size,
+            transform=worker.__dict__.get(_draft_head.TRANSFORM_ATTR),
+        )
     if reason is not None:
         raise ValueError(f"[musespark_open] wide_draft_topk cannot run here: {reason}")
     worker.__dict__[ATTRIBUTE] = RankTopKExchange(worker.__dict__[ATTRIBUTE])
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py
index 6644e4f977d0..8989f463b03c 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py
@@ -80,6 +80,34 @@
     "pool": ((WIDE_ROWS, "splitk:64,16,4,10"),),
     "unpool": ((16, "splitk:64,16,4,10"), (WIDE_ROWS, "splitk:64,16,2,11")),
 }
+
+
+def wide_kernel_steps(
+    capability: Optional[Tuple[int, int]],
+) -> Mapping[str, Tuple[Tuple[int, str], ...]]:
+    """The 9..32-row table of the device of ``capability`` (None: the target's own sm103,
+    as ``dense_gemms.device_capability`` reads a host without CUDA).
+
+    Raises:
+        ValueError: on sm107, whose TEP4 composition serves its 9..32-row projections on
+            ``wide_dense_gemms``' R200 table instead (``dense_gemms`` ``SM107_TP4_WIDE_KERNELS``),
+            and on any other device.
+    """
+    capability = (10, 3) if capability is None else tuple(capability)
+    if capability == (10, 3):
+        return TP4_WIDE_KERNELS
+    elif capability == (10, 7):
+        raise ValueError(
+            "[musespark_open] wide_projection_gathers has no sm107 table: sm107 TEP4 runs its "
+            "9..32-row projections on wide_dense_gemms"
+        )
+    else:
+        raise ValueError(
+            f"[musespark_open] wide_projection_gathers has no table for SM {capability} "
+            "(measured on sm103 only)"
+        )
+
+
 #: The per-rank ``(K, N)`` of each site at TP4.
 TP4_SHAPES: Mapping[str, Tuple[int, int]] = {
     "qkv": (8192, 4608),
@@ -166,7 +194,8 @@ def check(model: torch.nn.Module, plan: Optional[DenseGemmPlan]) -> None:
     """Whether ``install`` may run on ``model`` with ``dense_gemms``' ``plan``; changes nothing.
 
     Raises:
-        ValueError: if the model is not TP4, a layer's MoE or post-FFN join is not
+        ValueError: if the plan's device has no 9..32-row table (``wide_kernel_steps``: sm103
+            only), the model is not TP4, a layer's MoE or post-FFN join is not
             ``projection_gathers``' class with both sites fused, ``dense_gemms`` is not
             installed, or its plan prefetches (``dense_gemm_prefetch``): that prefetch's
             row exemptions (``skip_prefetch_rows``) are derived from the 1..8-row table, so
@@ -200,7 +229,7 @@ def check(model: torch.nn.Module, plan: Optional[DenseGemmPlan]) -> None:
                 f"[musespark_open] wide_projection_gathers widens projection_gathers' unpool site, "
                 f"which layer {layer.layer_idx} does not run ({type(join).__name__})"
             )
-    for site, by_rows in wide_kernels().items():
+    for site, by_rows in wide_kernels(wide_kernel_steps(plan.capability)).items():
         k, n = TP4_SHAPES[site]
         for rows, kernel in enumerate(by_rows, start=WIDE_MIN_ROWS):
             error = kernel_error(kernel, rows, k, n)
@@ -226,7 +255,8 @@ def install(model: torch.nn.Module, plan: DenseGemmPlan) -> WideProjections:
     pool_inbox = ProjectionInbox(
         mapping, row_words=_pool.local_columns(world) // 2, max_rows=WIDE_ROWS
     )
-    plan.set_wide_kernels(wide_kernels())
+    kernels = wide_kernels(wide_kernel_steps(plan.capability))
+    plan.set_wide_kernels(kernels)
     wide = WideProjections(plan)
     for layer in layers:
         moe, join = layer.feed_forward, layer.add_ffn
@@ -239,7 +269,7 @@ def install(model: torch.nn.Module, plan: DenseGemmPlan) -> WideProjections:
         f"in {len(layers)} layers: pool/unpool slices with their gathers fused into the "
         "pre-expert quantization and the post-FFN boundary (32-row inboxes); split-K for "
         + ", ".join(
-            f"{site} {'|'.join(dict.fromkeys(by_rows))}" for site, by_rows in wide_kernels().items()
+            f"{site} {'|'.join(dict.fromkeys(by_rows))}" for site, by_rows in kernels.items()
         )
     )
     return wide
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py
index 41cf773307e6..bdedd1c71ca7 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py
@@ -1053,8 +1053,9 @@ def adopt_model(
         full_attention.install_multi(model)
     decode_kv_splits = None
     if plan.enabled("decode_kv_splits"):
-        # Re-classes nothing: DecodeTarget arms it around a forward of 72..512 rows.
-        # Its refusal is agreed over the TP group, as the other attention installs.
+        # DecodeTarget arms it around a decode forward of 72..512 rows; the mixed-step
+        # mixin is installed after every other attention re-class (below). Its refusal is
+        # agreed over the TP group, as the other attention installs.
         from .kernels.decode_kv_splits import splits as _decode_kv_splits
 
         decode_kv_splits = _tp_agreement.agreed_install(
@@ -1094,6 +1095,13 @@ def adopt_model(
         _tp_agreement.agreed_install(
             model.layers[0].mapping, lambda: multi_request_append.install(model, min_seqs)
         )
+    if decode_kv_splits is not None:
+        # After every other attention re-class: the mixed-step mixin goes in front of all of
+        # them on the full-attention layers and arms the split rule only on mixed steps, which
+        # none of the others intercept.
+        from .kernels.decode_kv_splits import mixed as _mixed_kv_splits
+
+        _mixed_kv_splits.install(model, decode_kv_splits)
     if preattn_prep_on:
         logger.info(
             f"[musespark_open] preattn_prep: fused pre-attention prep on the decode path of all "
@@ -1109,6 +1117,15 @@ def adopt_model(
         swa_attention_block.log_installed(model.layers, swa_block)
         if plan.enabled("wide_swa_block"):
             swa_attention_block.install_multi(model)
+    if plan.enabled("mixed_step_gen_rows"):
+        # After every other attention re-class and after select_decode_fmha: the mixin goes
+        # in front of all of them and acts only on mixed steps, which none of the others
+        # intercept. Agreed like multi_request_split_kv's.
+        from .kernels.mixed_step_gen_rows import attention as mixed_step_gen_rows
+
+        _tp_agreement.agreed_install(
+            model.layers[0].mapping, lambda: mixed_step_gen_rows.install(model)
+        )
     layer_kwargs = {}
     if plan.enabled("projection_gathers"):
         # Imported only when planned: with no fast path the target imports
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py
index 63bcc3295d14..3544e9f35e61 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py
@@ -116,16 +116,30 @@ class acts only on sliding-window layers): their decode prep also writes the
   decode steps run it, while prefill and mixed steps keep the built-in
   selection (the ``thop`` attention). The drafter's attention is not a target
   layer and keeps the built-in selection.
-* ``decode_kv_splits`` (tep4 on sm_103 only) re-classes nothing: ``DecodeTarget``
+* ``decode_kv_splits`` (tep4 on sm_103 only): ``DecodeTarget``
   arms it for a forward of 72..512 rows (9 to 64 requests), and while armed the
   full-attention layers' trtllm-gen decode FMHA calls present flashinfer's
   KV-split selection with the SM count of a grid of about
   ``TARGET_CTAS_PER_SM`` CTAs per SM (``kernels/decode_kv_splits``), so those
-  steps keep multi-CTA KV at every batch size. Decode steps of up to 8
-  requests and prefill steps run what they ran before. Where
+  steps keep multi-CTA KV at every batch size. It also re-classes the
+  full-attention layers' attention with ``MixedStepKvSplitsAttention`` in front
+  of all of the above (``kernels/decode_kv_splits/mixed``): an attention call of
+  a mixed step with 9 to 64 generation requests runs with the same rule armed,
+  so the generation phase's decode FMHA of those steps keeps multi-CTA KV too.
+  Decode steps of up to 8 requests, mixed steps of up to 8 generation requests
+  and pure prefill steps run what they ran before. Where
   ``multi_request_split_kv`` serves a full-attention layer, its stand-in sits
   inside this one and launches its own kernel, so this split only acts on the
   calls that path declines.
+* ``mixed_step_gen_rows`` (tep4 on sm_107 only) takes the generation rows of
+  mixed steps with 9..64 generation requests off the thop attention on the
+  full-attention layers: a ``ThopContextPhaseFmha`` before each such layer's
+  thop fallback offers the context phase, so those steps run
+  ``CombinedFmha(context=thop, generation=DecodeTrtllmGenFmha)``, and
+  ``MixedStepGenRowsAttention`` in front of the layer's attention intercepts
+  the generation phase's decode call with ``multi_request_split_kv``'s
+  stand-in (``kernels/mixed_step_gen_rows``). Sliding-window layers, pure
+  context steps and mixed steps outside the band keep the thop attention.
 
 Dense projections and prefetch:
 
@@ -305,11 +319,13 @@ class acts only on sliding-window layers): their decode prep also writes the
   ``ring_buddy.install(windows=1)``), with nothing replicated.
 * ``three_window_placement`` stacks a third ring window on them.
 * ``half_window_experts`` (tp 4) stacks the ring predecessor's upper expert half on
-  them, which the prologue's routed kernel places onto; only up to the
+  them, which the prologue's routed kernel places onto (and, on calls of 9..32 rows
+  that reach the routing method, ``ring_buddy_assign``'s half flow); only up to the
   ``max_batch_size`` of its capacity gate (``dispatch.CAPACITY_GATES``).
 * ``bf16_routed_experts`` gives the two BF16 layers' stacked calls unpacked
   routing (on the owner-only window as well).
-* ``hot_replicas`` (tep4 on sm_103 only) hooks the NVFP4 layers' expert and routing
+* ``hot_replicas`` (tep4, in a plan without ``half_window_experts``:
+  ``dispatch.YIELDS``) hooks the NVFP4 layers' expert and routing
   calls after the owner-window install; the copies themselves are gathered on the
   first forward of a deployment whose largest decode step (``max_batch_size`` times
   the verify width) reaches 33 rows (``HotReplicas.decide``), and ``DecodeTarget``
@@ -380,7 +396,9 @@ class acts only on sliding-window layers): their decode prep also writes the
 write into a multicast inbox that the sampler gathers from (and fires the
 context-K/V fork's head fork point right after that GEMM); ``wide_split_sampler``
 (tep4 on sm_103 only) widens the head to decode steps of 129..1024 rows, which the
-sampler takes on the shipped sampler's 512-thread block (``target_head.widen``).
+sampler takes on the shipped sampler's 512-thread block (``target_head.widen``);
+``vocab_parallel_head`` (tep4 on sm_103 only) samples and argmaxes decode steps of
+16 rows or more from each rank's own vocabulary shard, without the logits all-gather.
 """
 
 from __future__ import annotations
@@ -734,12 +752,18 @@ def forward(
             self.qkv_prefetch.graph_start(
                 int((input_ids if input_ids is not None else inputs_embeds).shape[0])
             )
-        if inputs_embeds is None:
-            inputs_embeds = core.embed_tokens(input_ids)
-        embeds = core.embed_norm(inputs_embeds.to(core.dtype))
-        if core.m_emb != 1.0:
-            embeds = embeds * core.m_emb
-        hidden_states = embeds.to(torch.float32)
+        # Installed at weight load, after this target was built.
+        model_boundary = getattr(core, "_model_boundary", None)
+        first_attn_input = None
+        if model_boundary is not None and model_boundary.takes(input_ids, inputs_embeds):
+            hidden_states, first_attn_input = model_boundary.input(input_ids)
+        else:
+            if inputs_embeds is None:
+                inputs_embeds = core.embed_tokens(input_ids)
+            embeds = core.embed_norm(inputs_embeds.to(core.dtype))
+            if core.m_emb != 1.0:
+                embeds = embeds * core.m_emb
+            hidden_states = embeds.to(torch.float32)
         if self.preattn_prep is not None:
             kwargs = dict(kwargs, preattn_prep=self.preattn_prep)
         if self.decode_attention is not None:
@@ -787,7 +811,7 @@ def forward(
         batch_armed = batch is not None and batch.takes(rows)
 
         layers = core.layers
-        precomputed_attn_input = None
+        precomputed_attn_input = first_attn_input
         if hop_armed:
             arming = hop.armed(hidden_states.shape[0], hidden_states)
         elif wide_hop_armed:
@@ -806,6 +830,7 @@ def forward(
             two_shot.armed() if two_shot_armed else contextlib.nullcontext(),
             hot.armed() if hot_armed else contextlib.nullcontext(),
             batch.armed() if batch_armed else contextlib.nullcontext(),
+            model_boundary.armed() if model_boundary is not None else contextlib.nullcontext(),
         ):
             for index, decoder_layer in enumerate(layers):
                 next_norm = layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None
@@ -858,7 +883,10 @@ def forward(
 #: prefetch and PDL norm-gate of 33..512-row decode steps, the drafter window attention of
 #: 33..1024-row block decodes on a multi-request split geometry, the draft head's rank top-k
 #: key exchange on draft steps of 33 or more rows, the split verify sampler on
-#: 129..1024-row decode steps, and the context steps' pre-attention prep in one launch.
+#: 129..1024-row decode steps, the context steps' pre-attention prep in one launch, the
+#: verify head of 16..1024-row decode steps on each rank's own vocabulary shard over the
+#: tray's multicast, and the drafter hops of 9..56-row block decodes on a wide one-shot
+#: transport.
 _SM103_TEP4_ONLY_PATHS = frozenset(
     {
         "hot_replicas",
@@ -878,10 +906,18 @@ def forward(
         "wide_draft_topk",
         "wide_split_sampler",
         "context_prep",
+        "vocab_parallel_head",
+        "wide_drafter_hops",
     }
 )
+#: The fast paths validated at a tep4 rank's shape on sm107 only: the generation rows of
+#: mixed steps on the multi-request decode FMHA (on sm103 flashinfer's trtllm-gen FMHA
+#: already serves both phases of a mixed step).
+_SM107_TEP4_ONLY_PATHS = frozenset({"mixed_step_gen_rows"})
 #: TEP8: every fast path but the TP4-only ones of either architecture.
-_TEP8_PATHS = _dispatch.FAST_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS
+_TEP8_PATHS = (
+    _dispatch.FAST_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS - _SM107_TEP4_ONLY_PATHS
+)
 #: The fast paths written for any rank count whose output is bitwise the built-in
 #: model's.
 _TEP4_BITWISE_PATHS = frozenset(
@@ -892,6 +928,7 @@ def forward(
         "deferred_input_prep",
         "sampler_store",
         "replicated_embedding",
+        "model_boundary",
         "fused_accept",
         "verify_fold",
         "draft_head",
@@ -917,11 +954,20 @@ def forward(
 #: decode steps on the tray's MNNVL two-shot; on decode steps of 9..1024 rows, the
 #: post-attention hop on its own one-shot transport; large context steps
 #: sequence-parallel over the tray; and the context steps' pre-attention prep in one launch.
+#: Its NVFP4 expert layers also hold their ring predecessor's upper expert half
+#: (``half_window_experts``): decode steps of at most 8 rows are placed by the
+#: prologue's routed kernel, those of 9..32 rows (two to four requests, behind the
+#: multi-request router GEMM and projection gathers) by ``ring_buddy_assign``'s half
+#: flow, both balancing the selected experts over the four ranks. As on sm107 the
+#: window costs 45 GiB of KV cache per rank and is carried only up to
+#: ``max_batch_size`` 8 (``dispatch.CAPACITY_GATES``); a plan carrying it leaves
+#: ``hot_replicas`` out (``dispatch.YIELDS``).
 _SM103_TEP4_PATHS = (
     _TEP4_PATHS
     | _SM103_TEP4_ONLY_PATHS
     | frozenset(
         {
+            "half_window_experts",
             "swa_paged_decode",
             "swa_block",
             "full_attention_split_kv",
@@ -1023,7 +1069,7 @@ def forward(
 #: decode FMHA of steps of five or more requests (``multi_request_split_kv``; steps of two to
 #: four requests stay on ``wide_full_attention``).
 _SM107_TEP4_ATTENTION_PATHS = frozenset({"multi_request_split_kv"})
-_SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_ATTENTION_PATHS
+_SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_ATTENTION_PATHS | _SM107_TEP4_ONLY_PATHS
 _SM107_TEP4_PATHS |= _SM107_TEP4_COMM_PATHS
 #: The GB300 TEP4 GEMM and expert-balance paths also served at sm107 TEP4, tuned for
 #: R200: the router GEMM of decode steps of five or more requests (33..512 rows; 16..32
@@ -1040,13 +1086,15 @@ def forward(
 #: the two-to-four-request steps' append themselves), the drafter window attention of
 #: 33..1024-row block decodes on the multi-request split geometry (sized by the device's
 #: SM count), the draft head's rank top-k key exchange on draft steps of 33 or more rows,
-#: and the split verify sampler on 129..1024-row decode steps.
+#: the split verify sampler on 129..1024-row decode steps, and the drafter hops of
+#: 9..64-row block decodes on a wide one-shot transport.
 _SM107_TEP4_DRAFTER_AT_BATCH_PATHS = frozenset(
     {
         "multi_request_prep_append",
         "wide_drafter_window_attention",
         "wide_draft_topk",
         "wide_split_sampler",
+        "wide_drafter_hops",
     }
 )
 _SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_DRAFTER_AT_BATCH_PATHS
@@ -1162,7 +1210,9 @@ def load_weights(self, weights, weight_mapper=None):
         plan = self._plan
         if not plan:
             return result
-        _weights.install_loaded_features(self.llm.model, self.model_config.mapping, plan)
+        _weights.install_loaded_features(
+            self.llm.model, self.model_config.mapping, plan, device_compute_capability()
+        )
         # The token embedding after the experts, so the expert replication's
         # all-gather and this one's never interleave.
         _weights.install_embedding_features(self.llm.model, self.model_config.mapping, plan)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py
index 27494e01484c..21de86a7dee9 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py
@@ -72,6 +72,31 @@
 calls interleave freely, and a GEMM captured into a graph can store into a
 fixed area.
 
+Wide workspace (``max_rows`` above 8, fast path ``wide_drafter_hops``): the same kernel on
+drafter blocks of up to ``WIDE_MAX_ROWS`` rows (several requests' blocks; the convolution's
+tap-1 mask is per block, ``row % block``), with one slab per rank of the group and ``areas``
+receive areas used in site order modulo ``areas`` -- a divisor of ``sites`` of at least 3, so
+every window of ``areas`` consecutive calls (across forwards too) names distinct areas and
+the reset argument above holds with ``areas`` in place of ``sites``. A call resets the rows of
+the previous area that the previous call can have written: all ``max_rows`` at site 0 (the
+previous forward may have had any row count), at every other site the larger of its own row
+count and the previous call's (a forward's calls share one row count when served).
+
+Progress: a wide call launches ``CLUSTER * rows`` CTAs of 128 threads. A cluster publishes its
+row before it polls, and polls only the slots of rows ``t`` and ``t - 1`` of its column slice,
+which only the clusters of rows ``t`` and ``t - 1`` of the peers write. No other compute
+stream of the step is live during the drafter forward (the ``context_kv_fork`` side stream
+joins the main stream before the drafter's first hop), so no kernel of this rank holds an SM
+the hop's clusters wait for. Up to 32 rows (32 clusters = 256 CTAs) every cluster of the call
+is resident together on a GB300 (at most two 128-thread CTAs per SM), so progress needs no
+dispatch order. Above that (sm103's served 56 rows are 448 CTAs on 152 SMs, sm107's served 64
+rows are 512 CTAs on 212 SMs, and the qualification tests' 64..256 rows) not every cluster
+need be resident at once, and progress
+also relies on CTAs dispatched in grid order on every rank (the assumption the built-in MNNVL
+all-reduce kernels and ``post_attn_hop_wide`` make, which CUDA does not guarantee): the
+lowest-indexed unfinished cluster is then resident on every rank once every lower one has
+finished, so it publishes and finishes.
+
 Ordering contract: calls on one workspace must be issued in site order
 ``0, 1, ..., sites - 1, 0, ...`` on every rank, serialized on one stream (graph
 replay included), and every kernel launched with programmatic dependent launch
@@ -109,6 +134,33 @@
 HIDDEN = 8192
 #: Rows of a drafter block the transport is sized for (K + 1 at K = 7).
 MAX_ROWS = 8
+#: Rows of a wide workspace: the drafter blocks of up to 32 requests at K = 7.
+WIDE_MAX_ROWS = 256
+#: Rows of the served wide workspace (``wide_drafter_hops``) per compute capability
+#: (``wide_served_rows``). sm103: the drafter blocks of 2..7 requests at K = 7, where the
+#: one-shot hop is faster than the built-in all-reduce, convolution and add-RMSNorm on
+#: GB300 (at 8 requests the two tie; above, the built-in two-shot all-reduce wins). sm107:
+#: the blocks of 2..8 requests; on Rubin the one-shot hop is still the faster one at 64 rows
+#: and the slower one from 128 rows on.
+WIDE_SERVED_ROWS_BY_SM: dict[tuple[int, int], int] = {(10, 3): 56, (10, 7): 64}
+
+
+def wide_served_rows(capability: Optional[Tuple[int, int]] = None) -> int:
+    """The served wide workspace's rows on ``capability`` (this process's device when None).
+
+    Raises:
+        ValueError: on a device ``WIDE_SERVED_ROWS_BY_SM`` does not list.
+    """
+    key = tuple(torch.cuda.get_device_capability() if capability is None else capability)
+    rows = WIDE_SERVED_ROWS_BY_SM.get(key)
+    if rows is None:
+        raise ValueError(
+            f"drafter_hop: no served wide rows for SM {key} "
+            f"(supported: {sorted(WIDE_SERVED_ROWS_BY_SM)})"
+        )
+    return rows
+
+
 #: Largest tensor-parallel group; slots are laid out for it whatever the group size.
 MAX_RANKS = 8
 #: BF16 values per thread.
@@ -293,6 +345,24 @@ def _stale_bit(w0, w1, w2, w3, bit, *, loc=None, ip=None):
         )
         return Uint32(result)
 
+    @dsl_user_op
+    def _globaltimer(*, loc=None, ip=None):
+        """The global nanosecond timer (bounded-poll variant only)."""
+        result = _asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc, ip)
+        return Int64(result)
+
+    @dsl_user_op
+    def _trap_if(expired, *, loc=None, ip=None):
+        """Abort this rank's context when ``expired`` is non-zero (bounded-poll variant only)."""
+        _asm(
+            None,
+            [_u32(expired, loc, ip)],
+            "{\n.reg .pred p;\nsetp.ne.u32 p, $0, 0;\n@p trap;\n}",
+            "r",
+            loc,
+            ip,
+        )
+
     # ---- arithmetic: one PTX instruction each, as the replaced kernels run them ----
 
     def _f32_asm(text, operands, loc, ip):
@@ -424,12 +494,27 @@ def _pack_bf16x2(low, high, *, loc=None, ip=None):
     class _DrafterHopKernel:
         """All-reduce + DFlash 2 finish convolution + flashinfer add + RMSNorm, one row per cluster."""
 
-        def __init__(self, ranks: int, eps: float, pdl: bool, early_trigger: bool, direct: bool):
+        def __init__(
+            self,
+            ranks: int,
+            eps: float,
+            pdl: bool,
+            early_trigger: bool,
+            direct: bool,
+            max_rows: int = MAX_ROWS,
+            timeout_ns: int = 0,
+        ):
             self.ranks = ranks
             self.eps = float(eps)
             self.pdl = pdl
             self.early_trigger = early_trigger
             self.direct = direct
+            # Rows of one slab (one source rank's rows of an area) and its bytes.
+            self.max_rows = int(max_rows)
+            self.slab_bytes = self.max_rows * ROW_BYTES
+            # 0: the polls spin without a bound (served); else a poll still stale after this
+            # many nanoseconds executes ``trap`` (qualification builds).
+            self.timeout_ns = int(timeout_ns)
 
         @cute.jit
         def __call__(
@@ -449,6 +534,7 @@ def __call__(
             delta_row: Int32,
             delta_tap: Int32,
             dry: Int32,
+            reset_rows: Int32,
             stream: cuda.CUstream,
         ):
             self.kernel(
@@ -467,6 +553,7 @@ def __call__(
                 delta_row,
                 delta_tap,
                 dry,
+                reset_rows,
             ).launch(
                 grid=[rows * CLUSTER, 1, 1],
                 block=[THREADS, 1, 1],
@@ -494,6 +581,7 @@ def kernel(
             delta_row: Int32,
             delta_tap: Int32,
             dry: Int32,
+            reset_rows: Int32,
         ):
             tid = cute.arch.thread_idx()[0]
             lane = tid % 32
@@ -548,7 +636,7 @@ def kernel(
                     # Publish this rank's partial row into its slot on every rank
                     # (one store to the multicast address reaches every member).
                     _store_x4(
-                        mcast + Int64(rank) * SLAB_BYTES + slot_offset,
+                        mcast + Int64(rank) * self.slab_bytes + slot_offset,
                         own[0],
                         own[1],
                         own[2],
@@ -556,20 +644,36 @@ def kernel(
                     )
                 # Retire the previous site's area: every row, every slot this
                 # rank polls (rows t, t + rows, ... of this CTA's columns).
-                for i in cutlass.range_constexpr(MAX_ROWS):
-                    reset_row = row + i * rows
-                    in_area = Uint32(reset_row < MAX_ROWS)
-                    for p in cutlass.range_constexpr(self.ranks):
-                        reset_on = Uint32(1)
-                        if cutlass.const_expr(not self.direct):
-                            reset_on = Uint32(rank != p)
-                        _reset_x4(
-                            previous
-                            + Int64(p) * SLAB_BYTES
-                            + Int64(reset_row) * ROW_BYTES
-                            + slot_bytes,
-                            in_area & reset_on,
-                        )
+                if cutlass.const_expr(self.max_rows == MAX_ROWS):
+                    for i in cutlass.range_constexpr(MAX_ROWS):
+                        reset_row = row + i * rows
+                        in_area = Uint32(reset_row < MAX_ROWS)
+                        for p in cutlass.range_constexpr(self.ranks):
+                            reset_on = Uint32(1)
+                            if cutlass.const_expr(not self.direct):
+                                reset_on = Uint32(rank != p)
+                            _reset_x4(
+                                previous
+                                + Int64(p) * self.slab_bytes
+                                + Int64(reset_row) * ROW_BYTES
+                                + slot_bytes,
+                                in_area & reset_on,
+                            )
+                else:
+                    # A wide area: the rows the previous call can have written
+                    # (``reset_rows``), rows t, t + rows, ... of this CTA's columns.
+                    for reset_row in cutlass.range(row, reset_rows, rows, unroll=1):
+                        for p in cutlass.range_constexpr(self.ranks):
+                            reset_on = Uint32(1)
+                            if cutlass.const_expr(not self.direct):
+                                reset_on = Uint32(rank != p)
+                            _reset_x4(
+                                previous
+                                + Int64(p) * self.slab_bytes
+                                + Int64(reset_row) * ROW_BYTES
+                                + slot_bytes,
+                                reset_on,
+                            )
                 r = _load_x4(residual + row_offset)
                 group = Int64(column // GROUP_SIZE)
                 delta_address = delta + (Int64(row) * Int64(delta_row) + group) * 2
@@ -592,7 +696,7 @@ def kernel(
                             on = on & has_prev
                         mine = own if half == 0 else own_prev
                         v0, v1, v2, v3 = _poll_x4(
-                            receive_row + Int64(p) * SLAB_BYTES - half * ROW_BYTES,
+                            receive_row + Int64(p) * self.slab_bytes - half * ROW_BYTES,
                             on,
                             mine[0],
                             mine[1],
@@ -605,13 +709,16 @@ def kernel(
                         values[slot, 3] = v3
                         enabled[slot] = on
                         stale = stale | _stale_bit(v0, v1, v2, v3, on << Uint32(slot))
+                start = Int64(0)
+                if cutlass.const_expr(self.timeout_ns > 0):
+                    start = _globaltimer()
                 while stale != Uint32(0):
                     polled = stale
                     for half in cutlass.range_constexpr(2):
                         for p in cutlass.range_constexpr(self.ranks):
                             slot = half * self.ranks + p
                             v0, v1, v2, v3 = _poll_x4(
-                                receive_row + Int64(p) * SLAB_BYTES - half * ROW_BYTES,
+                                receive_row + Int64(p) * self.slab_bytes - half * ROW_BYTES,
                                 (polled >> Uint32(slot)) & Uint32(1),
                                 values[slot, 0],
                                 values[slot, 1],
@@ -631,6 +738,9 @@ def kernel(
                             values[slot, 3],
                             enabled[slot] << Uint32(slot),
                         )
+                    if cutlass.const_expr(self.timeout_ns > 0):
+                        elapsed = _globaltimer() - start
+                        _trap_if(Uint32((stale != Uint32(0)) & (elapsed > Int64(self.timeout_ns))))
                 # Early trigger: past the last transport access.
                 if cutlass.const_expr(self.pdl and self.early_trigger):
                     cute.arch.griddepcontrol_launch_dependents()
@@ -718,9 +828,16 @@ def kernel(
 
     @functools.lru_cache(maxsize=64)
     def _compiled_kernel(
-        device_index: int, ranks: int, eps: float, pdl: bool, early_trigger: bool, direct: bool
+        device_index: int,
+        ranks: int,
+        eps: float,
+        pdl: bool,
+        early_trigger: bool,
+        direct: bool,
+        max_rows: int = MAX_ROWS,
+        timeout_ns: int = 0,
     ):
-        kernel = _DrafterHopKernel(ranks, eps, pdl, early_trigger, direct)
+        kernel = _DrafterHopKernel(ranks, eps, pdl, early_trigger, direct, max_rows, timeout_ns)
         with torch.cuda.device(device_index):
             stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
             return cute.compile(
@@ -732,6 +849,7 @@ def _compiled_kernel(
                 Int32(0),
                 Int32(0),
                 Int32(1),
+                Int32(max_rows),
                 stream,
                 options="--enable-tvm-ffi",
             )
@@ -759,6 +877,13 @@ class DrafterHopWorkspace:
     use; construction checks that each is on this device, contiguous and
     16-byte aligned.
 
+    ``max_rows`` above ``MAX_ROWS`` builds a wide workspace (module docstring): one slab per
+    rank of the group, ``areas`` receive areas (default: the smallest divisor of ``sites`` of
+    at least 3), multicast publication only.
+
+    ``timeout_ns`` > 0 builds the bounded-poll variant for qualification: a poll still stale
+    after that long executes ``trap`` (aborting this rank's context) instead of spinning.
+
     Construction is collective over the group (the ``pre_wait`` check, fabric
     allocation, sentinel fill, compilation and a dry launch per ``eps``) and
     either succeeds on every rank or raises on every rank.
@@ -773,6 +898,9 @@ def __init__(
         early_trigger: bool = False,
         direct: bool = False,
         pre_wait: Sequence[torch.Tensor] = (),
+        max_rows: int = MAX_ROWS,
+        areas: Optional[int] = None,
+        timeout_ns: int = 0,
     ):
         from tensorrt_llm._torch.distributed.ops import (
             _get_mnnvl_workspace_comm,
@@ -791,10 +919,33 @@ def __init__(
             raise ValueError(f"the fused drafter hop serves 4 or 8 ranks, got {self.ranks}")
         if self.sites < MIN_SITES:
             raise ValueError(f"the fused drafter hop needs at least {MIN_SITES} sites")
+        self.max_rows = int(max_rows)
+        self.wide = self.max_rows > MAX_ROWS
+        self.timeout_ns = int(timeout_ns)
+        #: Rows of the last call issued (the reset bound of the next one inside a forward).
+        self._last_rows = self.max_rows
+        if self.wide:
+            if self.direct:
+                raise ValueError(
+                    "a wide drafter hop workspace publishes through the multicast only"
+                )
+            if not MAX_ROWS < self.max_rows <= WIDE_MAX_ROWS:
+                raise ValueError(f"a wide workspace takes {MAX_ROWS + 1}..{WIDE_MAX_ROWS} rows")
+            if areas is None:
+                areas = next(a for a in range(MIN_SITES, self.sites + 1) if self.sites % a == 0)
+            if areas < MIN_SITES or self.sites % areas:
+                raise ValueError(
+                    f"areas must divide the {self.sites} sites and be at least {MIN_SITES}"
+                )
+            self.areas = int(areas)
+            self.area_bytes = self.ranks * self.max_rows * ROW_BYTES
+        else:
+            self.areas = self.sites
+            self.area_bytes = AREA_BYTES
         self.device = torch.device("cuda", _mnnvl_device_index(mapping))
         self._kernels = {}
         self._comm = _get_mnnvl_workspace_comm(mapping)
-        size = self.sites * AREA_BYTES
+        size = self.areas * self.area_bytes
         error: Optional[BaseException] = None
         try:
             if not HAS_CUTE_DSL:
@@ -879,10 +1030,17 @@ def _kernel(self, eps: float):
         if torch.cuda.is_current_stream_capturing():
             raise RuntimeError(f"drafter hop variant eps={eps} was not built before graph capture")
         compiled = _compiled_kernel(
-            self.device.index, self.ranks, eps, self.pdl, self.early_trigger, self.direct
+            self.device.index,
+            self.ranks,
+            eps,
+            self.pdl,
+            self.early_trigger,
+            self.direct,
+            self.max_rows,
+            self.timeout_ns,
         )
         dummy = self._uc_base
-        compiled(*([dummy] * 9), 1, self.rank, 8, 0, 0, 1, self._stream())
+        compiled(*([dummy] * 9), 1, self.rank, 8, 0, 0, 1, self.max_rows, self._stream())
         self._kernels[eps] = compiled
         return compiled
 
@@ -890,7 +1048,11 @@ def _stream(self):
         return cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream)
 
     def _area(self, base: int, site: int) -> int:
-        return base + (site % self.sites) * AREA_BYTES
+        return base + (site % self.areas) * self.area_bytes
+
+    def nbytes(self) -> int:
+        """Fabric memory this rank holds for the transport."""
+        return self.areas * self.area_bytes
 
     def direct_output(self, site: int, rows: int) -> torch.Tensor:
         """This rank's slot of ``site``'s area through the multicast address, ``[rows, HIDDEN]`` BF16.
@@ -899,7 +1061,7 @@ def direct_output(self, site: int, rows: int) -> torch.Tensor:
         output here (every element once, its final value), and the hop then
         reads it from every rank's receive area.
         """
-        if not self.direct:
+        if not self.direct or self.wide:
             raise RuntimeError("direct_output is only for a direct (poll-only) workspace")
         start = (site % self.sites) * AREA_BYTES + self.rank * SLAB_BYTES
         flat = self._mc.view(torch.bfloat16)
@@ -928,6 +1090,11 @@ def launch(
             )
         compiled = self._kernel(float(eps))
         output = torch.empty_like(residual)
+        # The rows of the previous area its call can have written: any row count before the
+        # first site of a forward (the last forward that ran may not be the last one issued
+        # here: a capture that failed, a graph replayed after eager forwards), else the
+        # previous call's (wide areas only).
+        reset_rows = self.max_rows if site % self.sites == 0 else max(rows, self._last_rows)
         compiled(
             0 if partial is None else partial.data_ptr(),
             residual.data_ptr(),
@@ -944,9 +1111,11 @@ def launch(
             delta.stride(0),
             delta.stride(1),
             0,
+            reset_rows,
             self._stream(),
         )
         self.next_site = (site + 1) % self.sites
+        self._last_rows = rows
         return output
 
 
@@ -957,6 +1126,7 @@ def hop_takes(
     base: torch.Tensor,
     weight: torch.Tensor,
     block_size: int,
+    max_rows: int = MAX_ROWS,
 ) -> bool:
     """Whether the hop computes this call: rows, dtypes, shapes and devices only.
 
@@ -964,7 +1134,7 @@ def hop_takes(
     (``rows`` is the drafter block's row count), which is what lets a caller
     choose between this collective and the built-in all-reduce.
     """
-    if not 1 <= rows <= MAX_ROWS or block_size < 1:
+    if not 1 <= rows <= max_rows or block_size < 1:
         return False
     if any(t.dtype != torch.bfloat16 for t in (residual, delta, base, weight)):
         return False
@@ -1038,11 +1208,11 @@ def __call__(
         workspace: DrafterHopWorkspace,
     ) -> torch.Tensor:
         rows = residual.shape[0] if residual.dim() == 2 else 0
-        if not hop_takes(rows, residual, delta, base, weight, block_size) or not hop_layout_ok(
-            partial, residual, delta, base, weight
-        ):
+        if not hop_takes(
+            rows, residual, delta, base, weight, block_size, workspace.max_rows
+        ) or not hop_layout_ok(partial, residual, delta, base, weight):
             raise ValueError(
-                "drafter_hop takes BF16 [1..8, 8192] partial / residual, [rows, 2, 512] "
+                f"drafter_hop takes BF16 [1..{workspace.max_rows}, 8192] partial / residual, [rows, 2, 512] "
                 "coefficients with unit inner stride, a contiguous [2, 8192] base kernel and an "
                 "8192 weight, 16-byte aligned"
             )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py
index eb503cac93b3..e9fd45808a1f 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py
@@ -51,13 +51,24 @@
 ``prefetch`` (optional): a tensor whose bytes the kernel asks the L2 to fetch --
 the layer's attention output projection weight, read after the FMHA. Every
 program issues one ``cp.async.bulk.prefetch.L2`` of an equal slice right after
-its wait, so the requests leave every SM's TMA unit within the prep, before the
-FMHA's own TMA loads of K/V start. It reads nothing back and writes nothing: the
-values computed are the same with and without it.
+its wait, and a prefetching call runs one program per (row, head) item -- the
+built-in grid's program count, so the requests and their slice sizes are the
+built-in grid's -- and they leave every SM's TMA unit within the prep, before the
+FMHA's own TMA loads of K/V start. It reads nothing back and writes nothing: the values computed are the same
+with and without it.
+
+Launch geometry, per SM version (``LAUNCH_KIND_BY_SM``). On sm103 one program
+takes several consecutive (row, head) items of the row-major ``rows x (H + G)``
+grid (``PACKED_GEOMETRY``), each item running the per-head code of a
+one-program-per-(row, head) grid element for element (the same loads, the same
+additions in the same order, the same rounding points and stores), so the bytes
+written do not depend on the geometry. On sm107 the grid is one program per
+(row, head), one warp each (``_multi_request_prep_append_grid_kernel``).
 """
 
 from __future__ import annotations
 
+import functools
 from typing import Optional, Tuple
 
 import torch
@@ -74,7 +85,9 @@
     _load8,
     _normalize,
     _rotate,
+    _rsqrt_approx_ftz,
     is_rope_table,
+    libdevice,
     preattn_prep,
 )
 
@@ -83,16 +96,340 @@
 #: Query rows per sequence the entry takes: one accepted token and up to 7 drafts.
 MAX_ROWS_PER_SEQ = 8
 #: Sequences one call takes: 128 verify blocks of 8 rows (``preattn_prep``'s 1024-row cap), the
-#: largest decode step of a generation worker of ``max_batch_size`` 128. The grid is one program
-#: per (row, head), so a call's rows do not depend on how many sequences it carries.
+#: largest decode step of a generation worker of ``max_batch_size`` 128. The grid covers the
+#: (row, head) items, so a call's rows do not depend on how many sequences it carries.
 MAX_SEQS = 128
 #: Rows one call takes.
 MAX_ROWS = MAX_SEQS * MAX_ROWS_PER_SEQ
+#: The launch geometry by row count: ``(last row count, (row, head) items per program, warps per
+#: program, vector lanes)``, the first entry whose row count reaches the call's. An item is one
+#: head of one row; a program's items are consecutive in row-major (row, head) order. With vector
+#: lanes a lane loads and stores its 8 consecutive values as one vector (the same bytes). The
+#: entries above 32 rows are the fastest geometries of a GB300 sweep at 16..1024 rows; up to 32
+#: rows one item per program (the built-in grid) is faster in a served decode step on GB300.
+PACKED_GEOMETRY = (
+    (32, 1, 1, False),
+    (128, 2, 1, False),
+    (256, 4, 1, True),
+    (768, 8, 2, True),
+    (MAX_ROWS, 4, 1, True),
+)
+
+
+def packed_geometry(rows: int, heads: int, programs_at_least: int = 0) -> Tuple[int, int, bool]:
+    """``(items per program, warps, vector lanes)`` of a call of ``rows`` rows of ``heads`` heads.
+
+    ``programs_at_least`` lowers the items per program (never below one) until the grid has at
+    least that many programs -- the prefetching call's one bulk request per program then leaves
+    from at least that many programs.
+    """
+    items, warps, vec = next((i, w, v) for last, i, w, v in PACKED_GEOMETRY if rows <= last)
+    while items > 1 and -(-rows * heads // items) < programs_at_least:
+        items //= 2
+    warps = max(1, min(warps, items * 8 // 32))
+    return items, warps, vec
+
+
+@triton.jit
+def _fp8_store8(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16, mask):
+    """Eight BF16 values one norm thread owns, as E4M3 at ``ptr + 0..7`` where ``mask``."""
+    tl.store(ptr + 0, _to_cache_fp8(y0, scale_bf16), mask=mask)
+    tl.store(ptr + 1, _to_cache_fp8(y1, scale_bf16), mask=mask)
+    tl.store(ptr + 2, _to_cache_fp8(y2, scale_bf16), mask=mask)
+    tl.store(ptr + 3, _to_cache_fp8(y3, scale_bf16), mask=mask)
+    tl.store(ptr + 4, _to_cache_fp8(y4, scale_bf16), mask=mask)
+    tl.store(ptr + 5, _to_cache_fp8(y5, scale_bf16), mask=mask)
+    tl.store(ptr + 6, _to_cache_fp8(y6, scale_bf16), mask=mask)
+    tl.store(ptr + 7, _to_cache_fp8(y7, scale_bf16), mask=mask)
+
+
+@triton.jit
+def _split8(v):
+    """The 8 elements of ``v``'s last axis (shape ``[ITEMS, 2, 2, 2, 8]``), in order, each
+    ``[ITEMS, 2, 2, 2]``."""
+    v = tl.reshape(v, (v.shape[0], 2, 2, 2, 2, 2, 2))
+    # Element k = 4 a + 2 b + c of the last three axes (a, b, c): split c, then b, then a.
+    c0, c1 = tl.split(v)
+    c0b0, c0b1 = tl.split(c0)
+    c1b0, c1b1 = tl.split(c1)
+    x0, x4 = tl.split(c0b0)
+    x2, x6 = tl.split(c0b1)
+    x1, x5 = tl.split(c1b0)
+    x3, x7 = tl.split(c1b1)
+    return x0, x1, x2, x3, x4, x5, x6, x7
+
+
+@triton.jit
+def _join8(x0, x1, x2, x3, x4, x5, x6, x7):
+    """``_split8``'s inverse: ``[ITEMS, 2, 2, 2, 8]`` with element k = ``xk``."""
+    c0 = tl.join(tl.join(x0, x4), tl.join(x2, x6))
+    c1 = tl.join(tl.join(x1, x5), tl.join(x3, x7))
+    v = tl.join(c0, c1)
+    return tl.reshape(v, (v.shape[0], 2, 2, 2, 8))
+
+
+@triton.jit
+def _inverse_rms_per_item(x0, x1, x2, x3, x4, x5, x6, x7, eps):
+    """``_inverse_rms`` of each item of ``[ITEMS, 2, 2, 2]`` lane-shaped values, as ``[ITEMS]``.
+
+    The same additions in the same order: in order within a lane (``add.rn``), then lane
+    bit 0, 1, 2, each a single commutative add of a size-2 axis.
+    """
+    s = libdevice.add_rn(x0 * x0, x1 * x1)
+    s = libdevice.add_rn(s, x2 * x2)
+    s = libdevice.add_rn(s, x3 * x3)
+    s = libdevice.add_rn(s, x4 * x4)
+    s = libdevice.add_rn(s, x5 * x5)
+    s = libdevice.add_rn(s, x6 * x6)
+    s = libdevice.add_rn(s, x7 * x7)
+    total = tl.sum(tl.sum(tl.sum(s, axis=3), axis=2), axis=1)
+    return _rsqrt_approx_ftz(total / 64.0 + eps)
+
+
+@triton.jit
+def _multi_request_prep_append_kernel(
+    qkv_ptr,
+    qkv_out_ptr,
+    gate_out_ptr,
+    cos_sin_ptr,
+    positions_ptr,
+    q8_ptr,
+    pool_ptr,
+    bt_ptr,
+    bt_stride_seq,
+    bt_stride_kv,
+    seq_len_ptr,
+    orig_quant_ptr,
+    quant_orig_ptr,
+    bmm1_ptr,
+    bmm2_ptr,
+    cu_seqlens_ptr,
+    bmm1_host,
+    o_scale,
+    qkv_stride,
+    qkv_out_stride,
+    gate_out_stride,
+    eps,
+    q_scale,
+    num_seqs,
+    num_items,
+    pf_base,
+    pf_bytes,
+    pf_chunk,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_KV_HEADS: tl.constexpr,
+    QKV_PERMUTED: tl.constexpr,
+    APPLY_ROPE: tl.constexpr,
+    ROWS_PER_SEQ: tl.constexpr,
+    P: tl.constexpr,
+    LOG2_P: tl.constexpr,
+    SEQS_BLOCK: tl.constexpr,
+    LAUNCH_WITH_PDL: tl.constexpr,
+    ITEMS: tl.constexpr = 1,
+    LANES: tl.constexpr = 32,
+    VEC: tl.constexpr = False,
+    PREFETCH: tl.constexpr = False,
+):
+    """Program ``p`` takes the ``ITEMS`` items ``p * ITEMS ..`` of the ``rows * (H + G)``
+    (row, head) items in row-major order; every item runs the one-head code of the
+    one-program-per-(row, head) grid it replaces, element for element."""
+    D: tl.constexpr = 64
+    HALF: tl.constexpr = 32
+    HEADS: tl.constexpr = NUM_Q_HEADS + NUM_KV_HEADS
+    PAGE: tl.constexpr = NUM_KV_HEADS * P * D
+    if LAUNCH_WITH_PDL:
+        # Before any global read: the projection is the previous kernel's output.
+        # No early trigger: a dependent (the FMHA) launches when this grid ends.
+        tl.extra.cuda.gdc_wait()
+    program = tl.program_id(0)
+    if PREFETCH:
+        # One bulk L2 prefetch of this program's slice of ``pf_bytes`` (element 0 issues it;
+        # ``LANES`` is the program's thread count, so no element is held twice). The asm
+        # yields 0, so the store never happens; it keeps the request live.
+        lane = tl.arange(0, LANES)
+        start = program.to(tl.int64) * pf_chunk
+        size = tl.minimum(pf_bytes - start, pf_chunk)
+        size = tl.where((lane == 0) & (size > 0), size, 0).to(tl.int32)
+        issued = _bulk_prefetch_l2(tl.full([LANES], 0, tl.int64) + pf_base + start, size)
+        tl.store(cu_seqlens_ptr + lane, issued, mask=issued != 0)
+
+    # The program's items; a tail item past the last one computes item 0's addresses
+    # (every load stays in bounds) and stores nothing.
+    item = program * ITEMS + tl.arange(0, ITEMS)
+    valid = item < num_items
+    item = tl.where(valid, item, 0)
+    row = item // HEADS
+    head = item % HEADS
+    is_q = head < NUM_Q_HEADS
+    is_kv = head >= NUM_Q_HEADS
+    kv_head = head - NUM_Q_HEADS
+    row64 = row.to(tl.int64)
+    scale_bf16 = tl.load(orig_quant_ptr).to(tl.bfloat16)
+
+    # The row's cache position: the sequence's last ROWS_PER_SEQ positions
+    # (seq_lens counts the step's rows), addressed as KVBlockArray does.
+    seq = row // ROWS_PER_SEQ
+    position = tl.load(seq_len_ptr + seq) - ROWS_PER_SEQ + row % ROWS_PER_SEQ
+    table = bt_ptr + seq.to(tl.int64) * bt_stride_seq + (position >> LOG2_P)
+    slot = (kv_head * P + (position & (P - 1))) * D
+
+    # The pass-through half first, so its load is in flight during the norm:
+    # a query item copies its gate head, a key item its value head (and stores it
+    # into the cache's V page).
+    d = tl.arange(0, D)[None, :]
+    if QKV_PERMUTED:
+        gate_in = (NUM_Q_HEADS + 2 * NUM_KV_HEADS) * D + head * D
+        v_in = NUM_Q_HEADS * D + NUM_KV_HEADS * D + kv_head * D
+    else:
+        gate_in = head * 2 * D + D
+        v_in = 2 * NUM_Q_HEADS * D + NUM_KV_HEADS * D + kv_head * D
+    pass_in = tl.where(is_q, gate_in, v_in)
+    passed = tl.load(qkv_ptr + (row64 * qkv_stride + pass_in)[:, None] + d)
+    q_store = (valid & is_q)[:, None]
+    kv_store = (valid & is_kv)[:, None]
+    tl.store(gate_out_ptr + (row64 * gate_out_stride + head * D)[:, None] + d, passed, mask=q_store)
+    tl.store(
+        qkv_out_ptr
+        + (row64 * qkv_out_stride + (NUM_Q_HEADS + NUM_KV_HEADS) * D + kv_head * D)[:, None]
+        + d,
+        passed,
+        mask=kv_store,
+    )
+    page_v = tl.load(table + bt_stride_kv, mask=valid & is_kv, other=0).to(tl.int64)
+    tl.store(
+        pool_ptr + (page_v * PAGE + slot)[:, None] + d,
+        _to_cache_fp8(passed, scale_bf16),
+        mask=kv_store,
+    )
+
+    if QKV_PERMUTED:
+        in_off = tl.where(is_q, head * D, NUM_Q_HEADS * D + kv_head * D)
+    else:
+        in_off = tl.where(is_q, head * 2 * D, 2 * NUM_Q_HEADS * D + kv_head * D)
+    out_off = tl.where(is_q, head * D, NUM_Q_HEADS * D + kv_head * D)
+    scale = tl.where(is_q, q_scale, 1.0)[:, None, None, None]
+
+    lane8 = 8 * (
+        4 * tl.arange(0, 2)[None, :, None, None]
+        + 2 * tl.arange(0, 2)[None, None, :, None]
+        + tl.arange(0, 2)[None, None, None, :]
+    )
+    in_lane = qkv_ptr + (row64 * qkv_stride + in_off)[:, None, None, None] + lane8
+    if VEC:
+        # One 16-byte load per lane: the same 8 values.
+        x0, x1, x2, x3, x4, x5, x6, x7 = _split8(
+            tl.load(in_lane[:, :, :, :, None] + tl.arange(0, 8)[None, None, None, None, :]).to(
+                tl.float32
+            )
+        )
+    else:
+        x0, x1, x2, x3, x4, x5, x6, x7 = _load8(in_lane)
+    inverse_rms = _inverse_rms_per_item(x0, x1, x2, x3, x4, x5, x6, x7, eps)[:, None, None, None]
+    y0 = _normalize(x0, inverse_rms, scale)
+    y1 = _normalize(x1, inverse_rms, scale)
+    y2 = _normalize(x2, inverse_rms, scale)
+    y3 = _normalize(x3, inverse_rms, scale)
+    y4 = _normalize(x4, inverse_rms, scale)
+    y5 = _normalize(x5, inverse_rms, scale)
+    y6 = _normalize(x6, inverse_rms, scale)
+    y7 = _normalize(x7, inverse_rms, scale)
+
+    if APPLY_ROPE:
+        # Pair p = 4 * lane + i rotates elements 2p, 2p + 1 by table row
+        # [cos(32) | sin(32)] of this token's RoPE position.
+        rope_position = tl.load(positions_ptr + row).to(tl.int64)
+        rope = cos_sin_ptr + (rope_position * D)[:, None, None, None] + lane8 // 2
+        y0, y1 = _rotate(y0, y1, tl.load(rope + 0), tl.load(rope + HALF + 0))
+        y2, y3 = _rotate(y2, y3, tl.load(rope + 1), tl.load(rope + HALF + 1))
+        y4, y5 = _rotate(y4, y5, tl.load(rope + 2), tl.load(rope + HALF + 2))
+        y6, y7 = _rotate(y6, y7, tl.load(rope + 3), tl.load(rope + HALF + 3))
+
+    # The BF16 operand values: what the prep stores and what the append converts.
+    b0 = y0.to(tl.bfloat16)
+    b1 = y1.to(tl.bfloat16)
+    b2 = y2.to(tl.bfloat16)
+    b3 = y3.to(tl.bfloat16)
+    b4 = y4.to(tl.bfloat16)
+    b5 = y5.to(tl.bfloat16)
+    b6 = y6.to(tl.bfloat16)
+    b7 = y7.to(tl.bfloat16)
+    store = valid[:, None, None, None]
+    out = qkv_out_ptr + (row64 * qkv_out_stride + out_off)[:, None, None, None] + lane8
+    q8 = q8_ptr + ((row64 * NUM_Q_HEADS + head) * D)[:, None, None, None] + lane8
+    page_k = tl.load(table, mask=valid & is_kv, other=0).to(tl.int64)
+    k8 = pool_ptr + (page_k * PAGE + slot)[:, None, None, None] + lane8
+    if VEC:
+        # One vector store per lane each: the same bytes.
+        e = tl.arange(0, 8)[None, None, None, None, :]
+        b = _join8(b0, b1, b2, b3, b4, b5, b6, b7)
+        tl.store(out[:, :, :, :, None] + e, b, mask=store[:, :, :, :, None])
+        b8 = _to_cache_fp8(b, scale_bf16)
+        tl.store(q8[:, :, :, :, None] + e, b8, mask=(valid & is_q)[:, None, None, None, None])
+        tl.store(k8[:, :, :, :, None] + e, b8, mask=(valid & is_kv)[:, None, None, None, None])
+    else:
+        tl.store(out + 0, b0, mask=store)
+        tl.store(out + 1, b1, mask=store)
+        tl.store(out + 2, b2, mask=store)
+        tl.store(out + 3, b3, mask=store)
+        tl.store(out + 4, b4, mask=store)
+        tl.store(out + 5, b5, mask=store)
+        tl.store(out + 6, b6, mask=store)
+        tl.store(out + 7, b7, mask=store)
+        # Their E4M3 copies: Q into the FMHA's query buffer, K into the cache's K page.
+        _fp8_store8(
+            q8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16, (valid & is_q)[:, None, None, None]
+        )
+        _fp8_store8(
+            k8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16, (valid & is_kv)[:, None, None, None]
+        )
+
+    if program == 0:
+        # The FMHA's scale words, in the built-in kernel's operation order:
+        # (s_q * s_k) * bmm1_host, then that * log2(e); o_scale * s_v.
+        s = tl.load(quant_orig_ptr)
+        bmm1 = s * s * bmm1_host
+        tl.store(bmm1_ptr, bmm1)
+        tl.store(bmm1_ptr + 1, bmm1 * 1.4426950408889634)
+        tl.store(bmm2_ptr, o_scale * s)
+        if ROWS_PER_SEQ > 1:
+            # The decoder-info kernel's cumulative query lengths for fixed query lengths
+            # (it runs only for more than one query row per sequence).
+            i = tl.arange(0, SEQS_BLOCK)
+            tl.store(cu_seqlens_ptr + i, i * ROWS_PER_SEQ, mask=i <= num_seqs)
+
+
+#: The kernel ``multi_request_prep_append`` launches per SM version (``launch_kind``): on sm103
+#: the packed kernel above on ``PACKED_GEOMETRY``; on sm107 the grid of one program per
+#: (row, head), one warp each (``_multi_request_prep_append_grid_kernel``). Both compute every
+#: item with the same per-head code, so the bytes written are the same.
+LAUNCH_KIND_BY_SM = {103: "packed", 107: "grid"}
+
+
+@functools.lru_cache(maxsize=None)
+def _sm_of(index: int) -> int:
+    major, minor = torch.cuda.get_device_capability(index)
+    return 10 * major + minor
+
+
+def launch_kind(device: torch.device) -> str:
+    """``LAUNCH_KIND_BY_SM`` for ``device``.
+
+    Raises:
+        ValueError: on a device ``LAUNCH_KIND_BY_SM`` does not list.
+    """
+    sm = _sm_of(device.index if device.index is not None else torch.cuda.current_device())
+    kind = LAUNCH_KIND_BY_SM.get(sm)
+    if kind is None:
+        raise ValueError(
+            f"multi_request_prep_append: no launch for sm{sm} "
+            f"(supported: {sorted(LAUNCH_KIND_BY_SM)})"
+        )
+    return kind
 
 
 @triton.jit
-def _fp8_store8(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16):
-    """Eight BF16 values one norm thread owns, as E4M3 at ``ptr + 0..7``."""
+def _fp8_store8_all(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16):
+    """Eight BF16 values one norm thread owns, as E4M3 at ``ptr + 0..7`` (no mask)."""
     tl.store(ptr + 0, _to_cache_fp8(y0, scale_bf16))
     tl.store(ptr + 1, _to_cache_fp8(y1, scale_bf16))
     tl.store(ptr + 2, _to_cache_fp8(y2, scale_bf16))
@@ -104,7 +441,7 @@ def _fp8_store8(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16):
 
 
 @triton.jit
-def _multi_request_prep_append_kernel(
+def _multi_request_prep_append_grid_kernel(
     qkv_ptr,
     qkv_out_ptr,
     gate_out_ptr,
@@ -252,11 +589,11 @@ def _multi_request_prep_append_kernel(
     # Their E4M3 copies: Q into the FMHA's query buffer, K into the cache's K page.
     if is_q:
         q8 = q8_ptr + (row64 * NUM_Q_HEADS + head) * D + lane * 8
-        _fp8_store8(q8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16)
+        _fp8_store8_all(q8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16)
     else:
         page_k = tl.load(table).to(tl.int64)
         k8 = pool_ptr + page_k * PAGE + slot + lane * 8
-        _fp8_store8(k8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16)
+        _fp8_store8_all(k8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16)
 
     if (row == 0) & (head == 0):
         # The FMHA's scale words, in the built-in kernel's operation order:
@@ -321,7 +658,10 @@ class MultiRequestPrepAppend(OpWrapper):
     every target layer (46 RoPE + 16 NoPE) on decode steps of 2..128 sequences
     of the same number of rows (fast path multi_request_prep_append), whose
     trtllm-gen generation preprocess launch (decoder-info + append kernels)
-    the target then skips. Grid (rows, H + G), one warp per program.
+    the target then skips. Grid on sm103: the rows * (H + G) (row, head) items in
+    row-major order, PACKED_GEOMETRY's items, warps and lanes per program by row count
+    (with a prefetch, one program per item); on sm107: (rows, H + G), one warp per
+    program.
 
     Bit-identical, gated by its own tests: qkv_out and gate_out against
     preattn_prep; the cache bytes, the E4M3 query bytes, both scale words and
@@ -428,7 +768,11 @@ def __call__(
                 )
         elif position_ids is not None:
             raise ValueError("multi_request_prep_append: position_ids without a cos/sin table")
+        heads = num_q_heads + num_kv_heads
+        num_items = rows * heads
+        kind = launch_kind(qkv.device)
         pf_base, pf_bytes, pf_chunk = 0, 0, 16
+        programs_at_least = 0
         if prefetch is not None:
             pf_bytes = prefetch.numel() * prefetch.element_size()
             if not (
@@ -442,14 +786,66 @@ def __call__(
                     f"of a multiple of 16 bytes, got {tuple(prefetch.shape)} {prefetch.dtype} at "
                     f"{prefetch.data_ptr():#x}"
                 )
-            programs = rows * (num_q_heads + num_kv_heads)
+            # One bulk request per program, one program per (row, head) item: the o_proj
+            # weight is requested in the one-program-per-(row, head) grid's slices, from
+            # every SM, whatever the table's packing at these rows.
+            programs_at_least = num_items
+        if kind == "grid":
+            items, warps, vec = 1, 1, False
+        else:
+            items, warps, vec = packed_geometry(rows, heads, programs_at_least)
+        programs = -(-num_items // items)
+        if prefetch is not None:
             pf_base = prefetch.data_ptr()
             pf_chunk = ((pf_bytes + programs - 1) // programs + 15) // 16 * 16
         qkv_out = torch.empty(
             (rows, (num_q_heads + 2 * num_kv_heads) * HEAD_DIM), dtype=qkv.dtype, device=qkv.device
         )
         gate_out = torch.empty((rows, num_q_heads * HEAD_DIM), dtype=qkv.dtype, device=qkv.device)
-        _multi_request_prep_append_kernel[(rows, num_q_heads + num_kv_heads)](
+        if kind == "grid":
+            _multi_request_prep_append_grid_kernel[(rows, heads)](
+                qkv,
+                qkv_out,
+                gate_out,
+                cos_sin,
+                positions,
+                q_fp8.view(torch.uint8).view(torch.float8_e4m3fn),
+                kv_pool,
+                block_tables,
+                block_tables.stride(0),
+                block_tables.stride(1),
+                seq_lens,
+                kv_scale_orig_quant,
+                kv_scale_quant_orig,
+                bmm1_scale,
+                bmm2_scale,
+                cu_seqlens,
+                bmm1_host_scale,
+                o_scale,
+                qkv.stride(0),
+                qkv_out.stride(0),
+                gate_out.stride(0),
+                eps,
+                q_scale,
+                num_seqs,
+                pf_base,
+                pf_bytes,
+                pf_chunk,
+                NUM_Q_HEADS=num_q_heads,
+                NUM_KV_HEADS=num_kv_heads,
+                QKV_PERMUTED=qkv_permuted,
+                APPLY_ROPE=cos_sin is not None,
+                ROWS_PER_SEQ=rows_per_seq,
+                P=page,
+                LOG2_P=page.bit_length() - 1,
+                SEQS_BLOCK=triton.next_power_of_2(num_seqs + 1),
+                LAUNCH_WITH_PDL=enable_pdl,
+                PREFETCH=prefetch is not None,
+                num_warps=1,
+                launch_pdl=enable_pdl,
+            )
+            return qkv_out, gate_out
+        _multi_request_prep_append_kernel[(programs,)](
             qkv,
             qkv_out,
             gate_out,
@@ -474,6 +870,7 @@ def __call__(
             eps,
             q_scale,
             num_seqs,
+            num_items,
             pf_base,
             pf_bytes,
             pf_chunk,
@@ -486,8 +883,11 @@ def __call__(
             LOG2_P=page.bit_length() - 1,
             SEQS_BLOCK=triton.next_power_of_2(num_seqs + 1),
             LAUNCH_WITH_PDL=enable_pdl,
+            ITEMS=items,
+            LANES=32 * warps,
+            VEC=vec,
             PREFETCH=prefetch is not None,
-            num_warps=1,
+            num_warps=warps,
             launch_pdl=enable_pdl,
         )
         return qkv_out, gate_out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py
index a18b928bb409..1f514a082e42 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py
@@ -93,6 +93,30 @@
 #: The build the target launches: O rescale in 16-column chunks (spill-free), two Q terms.
 RESCALE_COLS = 16
 NQ = 2
+#: Pages (8 KB of K and V each) a CTA's ring holds in flight, by the number of E4M3 terms Q enters
+#: as, per SM version (``ring_stages``). At deep contexts the split kernel streams the cache and
+#: the ring depth sets its rate. sm103: each build takes the deepest ring one CTA's shared memory
+#: holds next to its Q buffers, 20 pages with a one-term Q (the served E4M3 Q) and 16 with a
+#: two-term Q (24 and 20 do not launch). sm107: 16 pages for either Q.
+RING_STAGES_BY_SM = {103: {1: 20, 2: 16}, 107: {1: 16, 2: 16}}
+#: sm103's ring depths.
+RING_STAGES = RING_STAGES_BY_SM[103]
+
+
+def ring_stages(device_index: int, q_terms: int) -> int:
+    """The ring depth of a build whose Q enters as ``q_terms`` E4M3 terms on the device.
+
+    Raises:
+        ValueError: on a device ``RING_STAGES_BY_SM`` does not list.
+    """
+    major, minor = torch.cuda.get_device_capability(device_index)
+    table = RING_STAGES_BY_SM.get(10 * major + minor)
+    if table is None:
+        raise ValueError(
+            f"multi_request_split_kv: no ring depth for sm{10 * major + minor} "
+            f"(supported: {sorted(RING_STAGES_BY_SM)})"
+        )
+    return table[q_terms]
 
 
 def default_num_work(kv_heads: int, sm_count: int) -> int:
@@ -163,10 +187,11 @@ def _compiled_kernel(
     rescale_cols: int = RESCALE_COLS,
     q_fp8: bool = False,
     device_scale: bool = False,
-    stages: int = 16,
+    stages: Optional[int] = None,
     max_requests: int = REQUEST_BUILDS[0],
 ):
-    """The compiled split kernel of one build."""
+    """The compiled split kernel of one build (its ring depth from ``ring_stages`` unless
+    ``stages`` is given)."""
     import cuda.bindings.driver as cuda
     import cutlass
     import cutlass.cute as cute
@@ -175,6 +200,8 @@ def _compiled_kernel(
     from . import multi_request_split_kv_kernel as mk
 
     assert (mk.MAX_ROWS, mk.REQUEST_BUILDS) == (MAX_ROWS, REQUEST_BUILDS)
+    if stages is None:
+        stages = ring_stages(device_index, 1 if q_fp8 else nq)
     kernel = mk.MultiRequestSplitKvKernel(
         enable_pdl=enable_pdl,
         kv_heads=kv_heads,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py
index 7ebefe6e9481..a96e07990486 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py
@@ -9,7 +9,7 @@
 The per-tile machinery is ``full_attention_split_kv_kernel``'s (the one-sequence kernel): a CTA
 serves one pair of KV heads (``blockIdx.y``) with 384 threads -- two softmax warpgroups that
 take alternating 128-key tiles, one MMA warp, one TMA loader warp, two other warps -- over a
-16-stage ring of 32-key pages (K in two 2 KB boxes, one per head, V in one 4 KB box). Q' is
+ring of stages 32-key pages (K in two 2 KB boxes, one per head, V in one 4 KB box). Q' is
 block-diagonal over the pair's two heads (row = head * 64 + token * 8 + q-head-in-group), so
 one M=128 MMA per Q term and K block scores a 128-key tile; P goes to TMEM in E4M3 in two
 halves; each warpgroup keeps its own running max / sum and TMEM O accumulator.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py
index 577e8ffb338c..f5a9c02a4de5 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py
@@ -58,7 +58,10 @@
 #: coefficient projection is the tep8 one. The fastest split-K tactic of each shape in
 #: chains of cold weights with PDL on; it beats cuBLAS on the qkv, gate_up and
 #: convolution shapes, ties it on down_proj and trails it slightly on o_proj in
-#: isolation, where ``drafter_pdl_chain``'s early weight load is what it is for.
+#: isolation, where ``drafter_pdl_chain``'s early weight load is what it is for. Besides the
+#: one-request block (8 rows) the table lists the blocks of 2..4 requests (16, 24 and 32
+#: rows), each the fastest split-K tactic at that row count on GB300's 152 SMs over cold
+#: weights with PDL on; every one is a single wave of at most 128 CTAs.
 SPLITK_TACTICS_TP4: dict[tuple[int, int, int], tuple[int, int, int, int]] = {
     # DFlash 2 convolution coefficient projection (replicated).
     (8, 8192, 2048): (64, 8, 4, 10),
@@ -70,16 +73,48 @@
     (8, 7168, 8192): (64, 8, 1, 11),
     # drafter gate_up_proj per rank.
     (8, 8192, 14336): (128, 8, 1, 6),
+    (16, 8192, 2048): (64, 16, 4, 8),
+    (16, 8192, 2560): (64, 16, 2, 10),
+    (16, 2048, 8192): (64, 16, 1, 10),
+    (16, 7168, 8192): (64, 16, 1, 10),
+    (16, 8192, 14336): (128, 16, 1, 6),
+    (24, 8192, 2048): (64, 16, 2, 11),
+    (24, 8192, 2560): (64, 32, 2, 9),
+    (24, 2048, 8192): (64, 32, 1, 9),
+    (24, 7168, 8192): (64, 32, 1, 9),
+    (24, 8192, 14336): (128, 32, 1, 5),
+    (32, 8192, 2048): (64, 16, 2, 11),
+    (32, 8192, 2560): (64, 32, 2, 9),
+    (32, 2048, 8192): (64, 32, 1, 9),
+    (32, 7168, 8192): (64, 32, 1, 9),
+    (32, 8192, 14336): (128, 32, 1, 5),
 }
 
 #: The tep4 drafter shapes on sm107 (212 SMs), each the fastest split-K tactic there in
-#: the same chains; every one beats cuBLAS at its shape.
+#: the same chains; every one beats cuBLAS at its shape. Besides the one-request block
+#: (8 rows) the table lists the blocks of 2..4 requests (16, 24 and 32 rows), each the
+#: fastest split-K tactic at that row count over cold weights with PDL on.
 SM107_SPLITK_TACTICS_TP4: dict[tuple[int, int, int], tuple[int, int, int, int]] = {
     (8, 8192, 2048): (64, 8, 4, 10),
     (8, 8192, 2560): (64, 8, 4, 12),
     (8, 2048, 8192): (128, 16, 2, 6),
     (8, 7168, 8192): (64, 8, 2, 10),
     (8, 8192, 14336): (64, 8, 2, 11),
+    (16, 8192, 2048): (64, 16, 4, 10),
+    (16, 8192, 2560): (64, 16, 4, 10),
+    (16, 2048, 8192): (64, 16, 1, 11),
+    (16, 7168, 8192): (64, 32, 1, 8),
+    (16, 8192, 14336): (128, 16, 1, 6),
+    (24, 8192, 2048): (64, 32, 4, 8),
+    (24, 8192, 2560): (64, 32, 4, 8),
+    (24, 2048, 8192): (64, 32, 1, 9),
+    (24, 7168, 8192): (64, 32, 1, 9),
+    (24, 8192, 14336): (128, 32, 1, 5),
+    (32, 8192, 2048): (64, 32, 4, 8),
+    (32, 8192, 2560): (64, 32, 4, 8),
+    (32, 2048, 8192): (64, 32, 1, 9),
+    (32, 7168, 8192): (64, 32, 2, 9),
+    (32, 8192, 14336): (64, 32, 2, 8),
 }
 
 #: The measured table of each ``(compute capability, drafter TP size)``.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py
new file mode 100644
index 000000000000..b35712085d19
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py
@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Fixed TRTLLM-gen tactics for the NVFP4 expert call of tep4 decode steps (sm103, sm107).
+
+``trtllm::fp4_block_scale_moe_runner`` takes its ``[tile_n, config]`` tactic from the
+autotuner, which profiles one tactic per power-of-two row bucket at warmup, on a dummy
+routing, separately on every rank: every row count of 256 to 511 runs the tactic
+profiled at 256 rows, where it settles on 64-token tiles.
+
+This table pins the tactic of the stacked expert call (``ring_buddy._nvfp4_call``) per
+``(expert window, rows)`` for listed decode row counts of tep4 (8 rows per request at the
+DFlash K = 7 verify width), one table per GPU:
+
+* sm107: the 32-token tile the autotuner serves at 128 to 255 rows, for 256 to 384 rows.
+  It computes every expert row bitwise as the autotuned tactics whose FC2 has the same
+  MMA K; tactics whose FC2 MMA K differs can round a rare row differently in its last
+  place (the autotuner picks from both).
+* sm103: the 128-token tile the autotuner serves at 512 rows, for 256 to 504 rows. Every
+  NVFP4 expert kernel sm103 loads has an MMA K of 64, and every tactic valid at these row
+  counts computes every expert row bitwise as the autotuned 64-token tactics.
+
+Row counts and windows not listed, and tuning passes, keep the autotuner. A GPU without a
+table is not supported.
+
+The config index is an index into the TRTLLM-gen MoE runner's list of (FC1, FC2) kernel
+pairs for one tile N, built from the kernels the running GPU loads, so an entry is only
+meaningful for the GPU and kernel library it was measured on; ``fixed_tactic`` also
+requires the entry to be a valid config of the call's shape (``get_valid_configs``) and
+returns None otherwise.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+#: The 32-token tile the autotuner serves at 128..255 rows (FC1
+#: ``t128x32x512u2_s5 … m256x32x128``, FC2 ``t128x32x512_s5 … m256x32x128``), pinned for the
+#: 256..384-row decode steps (32 to 48 requests) of the 68-expert ``hot_replicas`` window.
+#: Its FC2's MMA K is 128, as is that of the tile-64 tactics the autotuner picks there
+#: most, so every expert row is bitwise theirs.
+_TILE32_DECODE = (32, 236)
+
+#: ``(local_num_experts, rows) -> (tile_n, config)`` on sm107. Window 68 = the 64 own
+#: experts plus 4 hot copies (``hot_replicas``, 33..512 rows).
+SM107_TACTICS: Dict[Tuple[int, int], Tuple[int, int]] = {
+    (68, rows): _TILE32_DECODE for rows in range(256, 385, 8)
+}
+
+#: The 128-token tile the autotuner serves at 512 rows on sm103 (FC1
+#: ``t128x128x512_s3x3x3x3x1x3 … m256x128x64``, FC2 ``t128x128x256_s6 … m256x128x64``), pinned
+#: for the 256..504-row decode steps (32 to 63 requests) of the 68-expert window.
+_SM103_TILE128_DECODE = (128, 4)
+
+#: ``(local_num_experts, rows) -> (tile_n, config)`` on sm103.
+SM103_TACTICS: Dict[Tuple[int, int], Tuple[int, int]] = {
+    (68, rows): _SM103_TILE128_DECODE for rows in range(256, 505, 8)
+}
+
+_TABLES = {(10, 3): SM103_TACTICS, (10, 7): SM107_TACTICS}
+
+
+def table_for(capability: Tuple[int, int]) -> Dict[Tuple[int, int], Tuple[int, int]]:
+    """The table of a device capability: sm103 or sm107; any other GPU is not supported."""
+    table = _TABLES.get(tuple(capability))
+    if table is None:
+        raise NotImplementedError(
+            f"decode expert tactics: no table for compute capability {tuple(capability)} "
+            "(measured for sm103 and sm107 only)"
+        )
+    return table
+
+
+@lru_cache(maxsize=None)
+def _valid(act_type: int, top_k: int, hidden: int, intermediate: int, window: int, rows: int):
+    from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner
+
+    runner = FP4BlockScaleMoERunner.runner_dict.get((act_type,))
+    if runner is None:
+        runner = torch.classes.trtllm.FP4BlockScaleMoERunner(act_type)
+        FP4BlockScaleMoERunner.runner_dict[(act_type,)] = runner
+    configs = runner.get_valid_configs(top_k, hidden, intermediate, window, rows)
+    return frozenset(tuple(int(v) for v in c) for c in configs)
+
+
+def fixed_tactic(
+    device: torch.device,
+    act_type: int,
+    top_k: int,
+    hidden: int,
+    intermediate: int,
+    window: int,
+    rows: int,
+) -> Optional[List[int]]:
+    """The pinned ``[tile_n, config]`` of an expert call, or None to leave it to the autotuner.
+
+    None when the ``(window, rows)`` is not listed, the call is not the SwiGLU expert
+    (``act_type`` 0), or the listed tactic is not a valid config of the call's shape; a
+    device without a table raises (``table_for``).
+    """
+    if act_type != 0:
+        return None
+    tactic = table_for(torch.cuda.get_device_capability(device)).get((window, rows))
+    if tactic is None:
+        return None
+    if tactic not in _valid(act_type, top_k, hidden, intermediate, window, rows):
+        return None
+    return list(tactic)
+
+
+def run_fixed(
+    tactic: List[int],
+    inputs: List[Optional[torch.Tensor]],
+    num_experts: int,
+    top_k: int,
+    intermediate: int,
+    window: int,
+    routing_method_type: int,
+    act_type: int,
+    tune_max_num_tokens: int,
+) -> List[torch.Tensor]:
+    """The unfinalized pre-routed expert call on ``tactic``, without the autotuner.
+
+    ``inputs`` are ``FP4BlockScaleMoEInputs``'s 18 tensors in order (``routing_logits``
+    None, ``topk_weights`` / ``topk_ids`` the caller's). Returns what
+    ``trtllm::fp4_block_scale_moe_runner`` returns for ``do_finalize=False``.
+    """
+    from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner
+
+    runner = FP4BlockScaleMoERunner(
+        num_experts,
+        top_k,
+        None,
+        None,
+        intermediate,
+        0,
+        window,
+        None,
+        routing_method_type,
+        False,
+        act_type,
+        tune_max_num_tokens=tune_max_num_tokens,
+    )
+    return runner(inputs, tactic=tactic)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py
new file mode 100644
index 000000000000..f29b5cb9b895
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py
@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Prefill routed NVFP4 experts on the Rubin CuTe DSL grouped GEMMs, over TRTLLM-gen's tensors.
+
+The unfinalized expert call of a large (prefill or mixed) step: the same
+``(gemm2_output, expanded_idx_to_permuted_idx)`` contract as TRTLLM-gen's
+``fp4_block_scale_moe_runner`` with ``do_finalize=False``, computed by
+
+* ``torch.ops.trtllm.moe_sort`` (tile ``TILE`` rows),
+* the Rubin gather + grouped GEMM + SwiGLU + NVFP4 quantization kernel (FC1),
+  in its ``trtllm_gen_gated_layout`` mode,
+* the Rubin grouped GEMM (FC2) in its ``store_permuted`` +
+  ``trtllm_gen_n_shuffle`` + ``mma_k64`` mode: one BF16 row per permuted row,
+  no finalize, accumulated with the same 64-deep MMA as TRTLLM-gen's FC2.
+
+Both kernels read the served expert tensors as they are (TRTLLM-gen's shuffled
+rows and 128x4 swizzled block scales), so no second weight copy exists. The
+FC1 epilogue takes the served leaf's two FC1 scales as TRTLLM-gen's gated
+GEMM does (gate scale ``fc31_alpha``, output scale ``fc31_scale_c``) and uses
+TRTLLM-gen's SwiGLU and NVFP4 quantization arithmetic, so the intermediate
+codes and block scales equal TRTLLM-gen's; the model's per-expert post-norm
+then runs on the rows exactly as on TRTLLM-gen's.
+
+Graph capture: every kernel is compiled by ``prepare`` (outside any capture);
+the tactics are fixed per row bucket (``TACTICS``), so no autotuning happens in
+the served process.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Tuple
+
+import torch
+
+#: Routing tile (``moe_sort`` tile_tokens_dim, the kernels' M tile).
+TILE = 256
+
+#: Static tactics: (FC1, FC2) per smallest row count they serve, largest first.
+#: FC1 = (mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m, a_path);
+#: FC2 = (mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m).
+TACTICS: Tuple[Tuple[int, tuple, tuple], ...] = (
+    (
+        0,
+        ((256, 256, 256), (256, 256, 128), (2, 1), False, "cpasync"),
+        ((256, 256, 256), (256, 256, 128), (2, 1), False),
+    ),
+)
+
+
+def tactics_for(rows: int, table=None) -> Tuple[tuple, tuple]:
+    """The (FC1, FC2) tactics of a call of ``rows`` routed tokens."""
+    for min_rows, fc1, fc2 in table or TACTICS:
+        if rows >= min_rows:
+            return fc1, fc2
+    raise ValueError(f"no tactic for {rows} rows")
+
+
+class PrefillCuteDslExperts:
+    """One layer's prefill expert call over a TRTLLM-gen NVFP4 leaf's tensors.
+
+    ``backend`` is the leaf (or any object with its attributes: ``w3_w1_weight``,
+    ``w3_w1_weight_scale``, ``w2_weight``, ``w2_weight_scale``, ``fc31_alpha``,
+    ``fc31_scale_c``, ``fc2_alpha``, ``fc2_input_scale``). The tensors are read
+    on every call, as views, so a later re-binding of the leaf's parameters
+    (e.g. as views of a larger expert stack) is followed and nothing is copied.
+    The FC1 alpha pair ``[fc31_alpha | fc31_scale_c]`` is built once per set
+    of scale tensors.
+    """
+
+    def __init__(self, backend, num_experts: int, top_k: int, local_expert_offset: int):
+        from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as _ops
+
+        self.backend = backend
+        self.local_num_experts = int(backend.w3_w1_weight.shape[0])
+        self.num_experts = int(num_experts)
+        self.top_k = int(top_k)
+        self.local_expert_offset = int(local_expert_offset)
+        self.hidden = int(backend.w2_weight.shape[1])
+        e = self.local_num_experts
+        self._fc1 = _ops.Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner(
+            self.num_experts,
+            self.top_k,
+            e,
+            self.local_expert_offset,
+            TILE,
+            trtllm_gen_gated_layout=True,
+        )
+        self._fc2 = _ops.Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner(
+            self.num_experts,
+            self.top_k,
+            e,
+            self.local_expert_offset,
+            TILE,
+            torch.bfloat16,
+            store_permuted=True,
+            trtllm_gen_n_shuffle=True,
+            mma_k64=True,
+        )
+        self._scales_key = None
+        self._scales = None
+        self._dummy_scales: Dict[torch.device, torch.Tensor] = {}
+
+    def operands(self):
+        """``(w1, s1, w2, s2, gate_up_alpha, fc2_alpha, global_sf)`` in the kernels' dtypes (views)."""
+        be = self.backend
+        e = self.local_num_experts
+        w31 = be.w3_w1_weight
+        n1 = int(w31.shape[1])
+        w1 = w31.view(torch.uint8).view(torch.float4_e2m1fn_x2)
+        s1 = be.w3_w1_weight_scale.view(torch.uint8).reshape(e, n1, -1)
+        w2 = be.w2_weight.view(torch.uint8).view(torch.float4_e2m1fn_x2)
+        s2 = be.w2_weight_scale.view(torch.uint8).reshape(e, self.hidden, -1)
+        key = tuple(
+            (t.data_ptr(), t._version if not t.is_inference() else None)
+            for t in (be.fc31_alpha, be.fc31_scale_c, be.fc2_alpha, be.fc2_input_scale)
+        )
+        if key != self._scales_key:
+            # The FC1 kernel takes the output scale from the alpha pair; global_sf
+            # is a required operand it does not read in this mode.
+            g2 = be.fc2_input_scale.reshape(1).float()
+            gate_up_alpha = torch.cat(
+                (be.fc31_alpha.reshape(e).float(), be.fc31_scale_c.reshape(e).float())
+            ).contiguous()
+            self._scales = (
+                gate_up_alpha,
+                be.fc2_alpha.reshape(e).float().contiguous(),
+                g2.contiguous(),
+            )
+            self._scales_key = key
+        return (w1, s1, w2, s2) + self._scales
+
+    def __call__(
+        self,
+        x_fp4: torch.Tensor,
+        x_sf: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        tactics: Optional[Tuple[tuple, tuple]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """``(gemm2_output [permuted rows, hidden] bf16, expanded_idx_to_permuted_idx)``.
+
+        ``x_fp4`` is ``[T, hidden / 2]`` packed e2m1, ``x_sf`` the linear e4m3
+        block scales (``T * hidden / 16`` bytes), ``topk_ids`` global int32 ids
+        ``[T, top_k]``, ``topk_weights`` their combine weights (only sorted with
+        the ids; the rows carry no weight).
+        """
+        rows = int(x_fp4.shape[0])
+        fc1_tactic, fc2_tactic = tactics if tactics is not None else tactics_for(rows)
+        (
+            tile_idx_to_expert_idx,
+            tile_idx_to_mn_limit,
+            expanded_idx_to_permuted_idx,
+            permuted_idx_to_expanded_idx,
+            _total_num_padded_tokens,
+            num_non_exiting_tiles,
+        ) = torch.ops.trtllm.moe_sort(
+            token_selected_experts=topk_ids,
+            token_final_scales=topk_weights.float(),
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            local_expert_offset=self.local_expert_offset,
+            local_num_experts=self.local_num_experts,
+            tile_tokens_dim=TILE,
+        )
+        w1, s1, w2, s2, gate_up_alpha, fc2_alpha, global_sf = self.operands()
+        x = x_fp4.view(torch.uint8).view(torch.float4_e2m1fn_x2)
+        xs = x_sf.view(torch.uint8).reshape(rows, -1)
+        h, h_sf = self._fc1(
+            [
+                x,
+                w1,
+                xs,
+                s1,
+                gate_up_alpha,
+                tile_idx_to_expert_idx,
+                tile_idx_to_mn_limit,
+                permuted_idx_to_expanded_idx,
+                num_non_exiting_tiles,
+                global_sf,
+                None,
+                None,
+            ],
+            tactic=fc1_tactic,
+        )
+        m = int(permuted_idx_to_expanded_idx.shape[0])
+        out = torch.empty(m, self.hidden, dtype=torch.bfloat16, device=x_fp4.device)
+        scales = self._dummy_scales.get(x_fp4.device)
+        if scales is None:
+            # Unread in store_permuted mode; any valid pointer of the right dtype.
+            scales = torch.zeros(1, self.top_k, dtype=torch.float32, device=x_fp4.device)
+            self._dummy_scales[x_fp4.device] = scales
+        self._fc2(
+            [
+                h,
+                w2,
+                h_sf,
+                s2,
+                fc2_alpha,
+                out,
+                tile_idx_to_expert_idx,
+                tile_idx_to_mn_limit,
+                permuted_idx_to_expanded_idx,
+                num_non_exiting_tiles,
+                scales,
+            ],
+            tactic=fc2_tactic,
+        )
+        return out, expanded_idx_to_permuted_idx
+
+    def prepare(self, rows: int, table=None) -> None:
+        """Compile both kernels for ``rows`` (outside any graph capture) with an all-local routing."""
+        dev = self.backend.w3_w1_weight.device
+        ids = (
+            torch.arange(rows * self.top_k, device=dev, dtype=torch.int32) % self.local_num_experts
+            + self.local_expert_offset
+        ).reshape(rows, self.top_k)
+        x = torch.zeros(rows, self.hidden // 2, dtype=torch.uint8, device=dev)
+        xs = torch.zeros(rows * self.hidden // 16, dtype=torch.uint8, device=dev)
+        w = torch.full((rows, self.top_k), 1.0 / self.top_k, device=dev)
+        self(x, xs, ids, w, tactics=tactics_for(rows, table))
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu
index 3b251d689eab..b7aed7e11434 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu
@@ -60,6 +60,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <torch/library.h>
 
+#include <climits>
 #include <cstdint>
 #include <optional>
 
@@ -578,6 +579,569 @@ __global__ void splitWatchKernel(uint32_t const* buf, uint32_t* firstSeen, int64
     }
 }
 
+// ----------------------------------------------------------------------------------------
+// The vocabulary-parallel head (vp_*): each rank folds only its own LM-head shard.
+//
+// Virtual thread vt of the shipped VB-thread block owns the vectors v = vt (mod VB) of the
+// whole vocabulary; rank r holds the vectors [r * shardVecs, (r + 1) * shardVecs). vp_fold
+// folds, per (row, vt), the vectors of vt that lie in this rank's shard, in increasing
+// order, with the shipped per-element statements, and publishes that partial state to
+// every rank. vp_draw combines the R partials of its vt in rank order (the vocabulary
+// order of vt's vectors), then runs the shipped first round (block reduce, scan, Philox
+// draw, claimants); every rank walks its own portion of each claimant's elements, from
+// the claimant's base plus the masses of the earlier ranks' portions, and publishes the
+// crossing it finds (or none). vp_decide takes, per claimant, the first rank in rank order
+// with a crossing, and kernel B's rule across claimants. Greedy batches: vp_fold reduces
+// torch.argmax's comparator over each CTA and vp_argmax over the ranks' CTAs.
+//
+// Exchange buffer (one copy per rank, every 32-bit word holds kSplitSentinel between
+// calls): states [ranks][stateStride] x 16 B, greedy words [ranks][greedyStride] x 8 B at
+// greedyOffset, candidates [ranks][candStride] x 4 B at candOffset. Each value is stored
+// through every publish target (the multicast address in service; every rank's copy in
+// the single-GPU emulation); every rank reads and resets its own copy.
+
+//! Ranks the vocabulary-parallel head takes.
+constexpr int kVpMaxRanks = 8;
+//! Claimants of a row whose crossings the head exchanges (in thread order).
+constexpr int kVpClaimants = 4;
+//! A candidate word: no crossing in this rank's portion.
+constexpr uint32_t kVpNone = 0x7FFFFFFEu;
+//! Most vectors of one virtual thread in one shard (the walk's shared-memory stage).
+constexpr int kVpWalkVectors = 64;
+//! Threads per CTA of vp_fold.
+constexpr int kVpFoldThreads = 128;
+//! Greedy words per (rank, row): one per vp_fold CTA, at most kSplitBlock / kVpFoldThreads.
+constexpr int kVpGreedySlots = kSplitBlock / kVpFoldThreads;
+//! vp_draw's per-row record for vp_decide: [claimants or -1 (filtered), argmax, claimant tids].
+constexpr int kVpInfoWords = 2 + kVpClaimants;
+
+struct VpParams
+{
+    //! Row parameters; vocabSize is the whole vocabulary; logits is the pending
+    //! [numRows, vocabSize] BF16 row (written by vp_draw for filtered rows only).
+    FusedSamplingParams fs;
+    //! This rank's [numRows][shardVecs] local LM-head output.
+    uint2 const* local{nullptr};
+    //! This rank's staging rows [rows][shardVecs] (fabric memory; filtered rows only).
+    uint2* staging{nullptr};
+    //! Every rank's staging buffer, in rank order (peer addresses).
+    uint2 const* stagingOf[kVpMaxRanks]{};
+    //! Exchange-buffer bases the values are stored through.
+    char* publish[kVpMaxRanks]{};
+    int numPublish{0};
+    //! This rank's copy of the exchange buffer.
+    char* copy{nullptr};
+    int* info{nullptr};
+    int* status{nullptr};
+    int64_t* greedyTokens{nullptr};
+    SplitCap cap{};
+    int rank{0};
+    int ranks{0};
+    int shardVecs{0};
+    long long stateStride{0};
+    long long greedyOffset{0};
+    long long greedyStride{0};
+    long long candOffset{0};
+    long long candStride{0};
+    unsigned long long timeoutNs{0};
+};
+
+__device__ __forceinline__ uint4* vpState(char* base, VpParams const& p, int rank, int row, int vt, int vb)
+{
+    return reinterpret_cast<uint4*>(base) + rank * p.stateStride + static_cast<long long>(row) * vb + vt;
+}
+
+__device__ __forceinline__ uint2* vpGreedy(char* base, VpParams const& p, int rank, int row, int slot)
+{
+    return reinterpret_cast<uint2*>(base + p.greedyOffset) + rank * p.greedyStride
+        + static_cast<long long>(row) * kVpGreedySlots + slot;
+}
+
+__device__ __forceinline__ uint32_t* vpCand(char* base, VpParams const& p, int rank, int row, int c)
+{
+    return reinterpret_cast<uint32_t*>(base + p.candOffset) + rank * p.candStride
+        + static_cast<long long>(row) * kVpClaimants + c;
+}
+
+__device__ __forceinline__ uint4 loadRelaxedSys4(uint4 const* address)
+{
+    uint4 v;
+    asm volatile("ld.relaxed.sys.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(address)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ void storeRelaxedSys4(uint4* address, uint4 v)
+{
+    asm volatile(
+        "st.relaxed.sys.global.v4.b32 [%0], {%1, %2, %3, %4};" ::"l"(address), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)
+        : "memory");
+}
+
+__device__ __forceinline__ uint32_t loadRelaxedSys1(uint32_t const* address)
+{
+    uint32_t v;
+    asm volatile("ld.relaxed.sys.global.b32 %0, [%1];" : "=r"(v) : "l"(address) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ void storeRelaxedSys1(uint32_t* address, uint32_t v)
+{
+    asm volatile("st.relaxed.sys.global.b32 [%0], %1;" ::"l"(address), "r"(v) : "memory");
+}
+
+//! Records a poll that waited past the timeout and traps (a peer that never published).
+__device__ __noinline__ void vpPollTimedOut(int* status)
+{
+    atomicOr(status, kSplitPollTimeout);
+    __threadfence_system();
+    __trap();
+}
+
+//! Spin until none of the first `words` 32-bit words at `address` is the sentinel.
+template <int WORDS>
+__device__ __forceinline__ uint4 vpAwait(void const* address, int* status, unsigned long long timeoutNs)
+{
+    unsigned long long start = 0;
+    for (int spins = 0;; ++spins)
+    {
+        uint4 v{kSplitSentinel, kSplitSentinel, kSplitSentinel, kSplitSentinel};
+        if constexpr (WORDS == 1)
+        {
+            v.x = loadRelaxedSys1(static_cast<uint32_t const*>(address));
+        }
+        else if constexpr (WORDS == 2)
+        {
+            uint2 const w = loadRelaxedSys(static_cast<uint2 const*>(address));
+            v.x = w.x;
+            v.y = w.y;
+        }
+        else
+        {
+            v = loadRelaxedSys4(static_cast<uint4 const*>(address));
+        }
+        bool ready = v.x != kSplitSentinel;
+        if constexpr (WORDS >= 2)
+        {
+            ready = ready && v.y != kSplitSentinel;
+        }
+        if constexpr (WORDS >= 3)
+        {
+            ready = ready && v.z != kSplitSentinel;
+        }
+        if (ready)
+        {
+            return v;
+        }
+        if (spins == 64)
+        {
+            start = globalTimerNs();
+        }
+        if (spins >= 64)
+        {
+            __nanosleep(64);
+            if (globalTimerNs() - start > timeoutNs)
+            {
+                vpPollTimedOut(status);
+            }
+        }
+    }
+}
+
+//! torch.argmax's comparator (GreaterOrNan): whether (a, ia) is chosen over (b, ib).
+__device__ __forceinline__ bool vpArgmaxBetter(float a, int ia, float b, int ib)
+{
+    if (isnan(a))
+    {
+        return isnan(b) ? ia < ib : true;
+    }
+    return a == b ? ia < ib : a > b;
+}
+
+struct VpMaxArg
+{
+    float value;
+    int index;
+};
+
+struct VpMaxArgOp
+{
+    __device__ __forceinline__ VpMaxArg operator()(VpMaxArg const& a, VpMaxArg const& b) const
+    {
+        return vpArgmaxBetter(a.value, a.index, b.value, b.index) ? a : b;
+    }
+};
+
+//! The first vector of virtual thread vt in the shard starting at global vector `first`, as
+//! an offset into the shard.
+__device__ __forceinline__ int vpLead(int vt, int first, int vb)
+{
+    return ((vt - first) % vb + vb) % vb;
+}
+
+//! vp_fold: grid (numRows, VB / kVpFoldThreads), kVpFoldThreads threads. PDL: waits on
+//! the LM-head GEMM before reading anything; releases its dependent after publishing.
+template <bool GREEDY, int VB>
+__global__ void __launch_bounds__(kVpFoldThreads) vpFoldKernel(VpParams p)
+{
+    gridDependencyWait();
+    int const row = blockIdx.x;
+    int const vt = blockIdx.y * kVpFoldThreads + threadIdx.x;
+    int const sv = p.shardVecs;
+    int const first = p.rank * sv;
+    int const lead = vpLead(vt, first, VB);
+    uint2 const* localRow = p.local + static_cast<size_t>(row) * sv;
+    float tempInv = 1.0f;
+    bool stage = false;
+    if constexpr (!GREEDY)
+    {
+        RowParams const rp = loadRowParams(p.fs, row);
+        tempInv = rp.tempInv;
+        stage = rp.needTopK || rp.needTopP || rp.needMinP;
+    }
+    uint2* stagingRow = p.staging + static_cast<size_t>(row) * sv;
+    OnlineSoftmax st{-FLT_MAX, 0.0f, 0};
+    VpMaxArg best{-INFINITY, INT_MAX};
+    for (int l0 = lead; l0 < sv; l0 += kSplitPrefetch * VB)
+    {
+        uint2 buf[kSplitPrefetch];
+#pragma unroll
+        for (int u = 0; u < kSplitPrefetch; ++u)
+        {
+            int const l = l0 + u * VB;
+            if (l < sv)
+            {
+                buf[u] = __ldg(localRow + l);
+            }
+        }
+#pragma unroll
+        for (int u = 0; u < kSplitPrefetch; ++u)
+        {
+            int const l = l0 + u * VB;
+            if (l < sv)
+            {
+                if (stage)
+                {
+                    stagingRow[l] = buf[u];
+                }
+                __nv_bfloat16 const* elems = reinterpret_cast<__nv_bfloat16 const*>(&buf[u]);
+#pragma unroll
+                for (int j = 0; j < 4; ++j)
+                {
+                    float const logit = capOf(__bfloat162float(elems[j]), p.cap);
+                    int const index = (first + l) * 4 + j;
+                    if constexpr (GREEDY)
+                    {
+                        if (vpArgmaxBetter(logit, index, best.value, best.index))
+                        {
+                            best = VpMaxArg{logit, index};
+                        }
+                    }
+                    else
+                    {
+                        st = combineOnlineSoftmax(st, OnlineSoftmax{__fmul_rn(logit, tempInv), 1.0f, index});
+                    }
+                }
+            }
+        }
+    }
+    if constexpr (GREEDY)
+    {
+        using BlockReduceMaxArg = cub::BlockReduce<VpMaxArg, kVpFoldThreads>;
+        __shared__ typename BlockReduceMaxArg::TempStorage temp;
+        VpMaxArg const cta = BlockReduceMaxArg(temp).Reduce(best, VpMaxArgOp());
+        if (threadIdx.x == 0)
+        {
+            uint2 const word = make_uint2(__float_as_uint(cta.value), static_cast<uint32_t>(cta.index));
+            for (int i = 0; i < p.numPublish; ++i)
+            {
+                *vpGreedy(p.publish[i], p, p.rank, row, blockIdx.y) = word;
+            }
+        }
+    }
+    else
+    {
+        if (stage)
+        {
+            // The staged row reaches the peers before this thread's state does.
+            __threadfence_system();
+        }
+        uint4 const word
+            = make_uint4(__float_as_uint(st.max), __float_as_uint(st.sum), static_cast<uint32_t>(st.argmax), 0u);
+        for (int i = 0; i < p.numPublish; ++i)
+        {
+            *vpState(p.publish[i], p, p.rank, row, vt, VB) = word;
+        }
+    }
+    gridDependencyTrigger();
+}
+
+//! vp_draw: one VB-thread CTA per row. PDL: waits on vp_fold before reading anything.
+template <int VB>
+__global__ void __launch_bounds__(VB, VB == kSplitNarrowBlock ? 2 : 1) vpDrawKernel(VpParams p)
+{
+    __shared__ FusedSamplingShared<VB> shared;
+    __shared__ typename cub::BlockScan<int, VB>::TempStorage scanClaims;
+    __shared__ __align__(16) float wbuf[4 * kVpWalkVectors];
+    __shared__ int sClaimTid[kVpClaimants];
+    __shared__ float sClaimStart[kVpClaimants];
+    __shared__ uint32_t sCand[kVpClaimants];
+    int const row = blockIdx.x;
+    int const tid = threadIdx.x;
+    if (tid == 0)
+    {
+        sSplitCap = p.cap;
+    }
+    __syncthreads();
+    gridDependencyWait();
+    RowParams const rp = loadRowParams(p.fs, row);
+    bool const filtered = rp.needTopK || rp.needTopP || rp.needMinP;
+    // This thread's R partial states, combined in rank order. They stay in the copy until
+    // the end of the call (a claimant reads the earlier ranks' again for its walk offset).
+    OnlineSoftmax st{-FLT_MAX, 0.0f, 0};
+    for (int r = 0; r < p.ranks; ++r)
+    {
+        uint4 const w = vpAwait<3>(vpState(p.copy, p, r, row, tid, VB), p.status, p.timeoutNs);
+        OnlineSoftmax const part{__uint_as_float(w.x), __uint_as_float(w.y), static_cast<int>(w.z)};
+        st = r == 0 ? part : combineOnlineSoftmax(st, part);
+    }
+    uint4 const sentinel4 = make_uint4(kSplitSentinel, kSplitSentinel, kSplitSentinel, kSplitSentinel);
+    int const sv = p.shardVecs;
+    int const vecCount = p.fs.vocabSize / 4;
+    int* const info = p.info + static_cast<size_t>(row) * kVpInfoWords;
+    if (filtered)
+    {
+        // Every rank staged this row before publishing its states (vp_fold); gather it
+        // into the pending row and run the shipped body on it, as kernel B does.
+        __threadfence_system();
+        __syncthreads();
+        uint2* const pendingRow
+            = reinterpret_cast<uint2*>(const_cast<void*>(p.fs.logits)) + static_cast<size_t>(row) * vecCount;
+        for (int v = tid; v < vecCount; v += VB)
+        {
+            int const r = v / sv;
+            pendingRow[v] = loadRelaxedSys(p.stagingOf[r] + static_cast<size_t>(row) * sv + (v - r * sv));
+        }
+        __syncthreads();
+        fusedSamplingBody<CapBf16, VB, true, false>(p.fs, shared);
+        for (int r = 0; r < p.ranks; ++r)
+        {
+            storeRelaxedSys4(vpState(p.copy, p, r, row, tid, VB), sentinel4);
+        }
+        if (tid < kVpClaimants)
+        {
+            for (int i = 0; i < p.numPublish; ++i)
+            {
+                *vpCand(p.publish[i], p, p.rank, row, tid) = kVpNone;
+            }
+        }
+        if (tid == 0)
+        {
+            info[0] = -1;
+        }
+        gridDependencyTrigger();
+        return;
+    }
+    using BlockReduceOnline = cub::BlockReduce<OnlineSoftmax, VB>;
+    using BlockScanF = cub::BlockScan<float, VB>;
+    OnlineSoftmax const rowStats = BlockReduceOnline(shared.temp.reduceOnline).Reduce(st, OnlineSoftmaxOp());
+    if (tid == 0)
+    {
+        shared.maxScaled = rowStats.max;
+        shared.argmax = rowStats.argmax;
+    }
+    __syncthreads();
+    float const maxScaled = shared.maxScaled;
+    float const localMass = st.sum * __expf(st.max - maxScaled);
+    if (tid == 0)
+    {
+        initRowRng(p.fs, row, &shared.rejectRng);
+    }
+    __syncthreads();
+    float base = 0.0f;
+    float supportMass = 0.0f;
+    BlockScanF(shared.temp.scanF).ExclusiveSum(localMass, base, supportMass);
+    __syncthreads();
+    if (tid == 0)
+    {
+        shared.target = curand_uniform(&shared.rejectRng) * supportMass;
+    }
+    __syncthreads();
+    float const target = shared.target;
+    bool const claims = target >= base && target < base + localMass;
+    int claimIdx = 0;
+    int claimants = 0;
+    cub::BlockScan<int, VB>(scanClaims).ExclusiveSum(claims ? 1 : 0, claimIdx, claimants);
+    if (claims && claimIdx < kVpClaimants)
+    {
+        // The masses of this thread's portions on the ranks before this one.
+        float earlier = 0.0f;
+        for (int r = 0; r < p.rank; ++r)
+        {
+            uint4 const w = loadRelaxedSys4(vpState(p.copy, p, r, row, tid, VB));
+            earlier += __uint_as_float(w.y) * __expf(__uint_as_float(w.x) - maxScaled);
+        }
+        sClaimTid[claimIdx] = tid;
+        sClaimStart[claimIdx] = base + earlier;
+    }
+    for (int r = 0; r < p.ranks; ++r)
+    {
+        storeRelaxedSys4(vpState(p.copy, p, r, row, tid, VB), sentinel4);
+    }
+    __syncthreads();
+    int const first = p.rank * sv;
+    uint2 const* localRow = p.local + static_cast<size_t>(row) * sv;
+    int const walkers = min(claimants, kVpClaimants);
+    for (int c = 0; c < walkers; ++c)
+    {
+        int const walker = sClaimTid[c];
+        int const lead = vpLead(walker, first, VB);
+        int const walkVectors = lead < sv ? (sv - lead + VB - 1) / VB : 0;
+        if (tid < walkVectors)
+        {
+            uint2 const packed = __ldg(localRow + lead + tid * VB);
+            __nv_bfloat16 const* elems = reinterpret_cast<__nv_bfloat16 const*>(&packed);
+#pragma unroll
+            for (int j = 0; j < 4; ++j)
+            {
+                wbuf[tid * 4 + j] = __expf(__fmul_rn(capOf(__bfloat162float(elems[j]), p.cap), rp.tempInv) - maxScaled);
+            }
+        }
+        __syncthreads();
+        if (tid == 0)
+        {
+            float running = sClaimStart[c];
+            uint32_t cand = kVpNone;
+            float4 const* w4 = reinterpret_cast<float4 const*>(wbuf);
+            for (int q = 0; q < walkVectors && cand == kVpNone; ++q)
+            {
+                float4 const quad = w4[q];
+                float const ws[4] = {quad.x, quad.y, quad.z, quad.w};
+#pragma unroll
+                for (int j = 0; j < 4; ++j)
+                {
+                    float const w = ws[j];
+                    if (cand == kVpNone && w >= 0.0f && w > -1.0f)
+                    {
+                        running += w;
+                        if (running > target)
+                        {
+                            cand = static_cast<uint32_t>((first + lead + q * VB) * 4 + j);
+                        }
+                    }
+                }
+            }
+            sCand[c] = cand;
+        }
+        __syncthreads();
+    }
+    if (tid < kVpClaimants)
+    {
+        uint32_t const cand = tid < walkers ? sCand[tid] : kVpNone;
+        for (int i = 0; i < p.numPublish; ++i)
+        {
+            *vpCand(p.publish[i], p, p.rank, row, tid) = cand;
+        }
+    }
+    if (tid == 0)
+    {
+        info[0] = claimants;
+        info[1] = shared.argmax;
+        for (int c = 0; c < kVpClaimants; ++c)
+        {
+            info[2 + c] = c < walkers ? sClaimTid[c] : -1;
+        }
+    }
+    gridDependencyTrigger();
+}
+
+//! vp_decide: one warp per row. PDL: waits on vp_draw before reading anything.
+__global__ void __launch_bounds__(32) vpDecideKernel(VpParams p, int vb)
+{
+    gridDependencyWait();
+    int const row = blockIdx.x;
+    int const lane = threadIdx.x;
+    uint32_t mine = kVpNone;
+    if (lane < p.ranks * kVpClaimants)
+    {
+        uint32_t* const slot = vpCand(p.copy, p, lane / kVpClaimants, row, lane % kVpClaimants);
+        mine = vpAwait<1>(slot, p.status, p.timeoutNs).x;
+        storeRelaxedSys1(slot, kSplitSentinel);
+    }
+    uint32_t words[kVpMaxRanks * kVpClaimants];
+#pragma unroll
+    for (int i = 0; i < kVpMaxRanks * kVpClaimants; ++i)
+    {
+        words[i] = __shfl_sync(0xffffffffu, mine, i);
+    }
+    if (lane == 0)
+    {
+        int const* info = p.info + static_cast<size_t>(row) * kVpInfoWords;
+        int const claimants = info[0];
+        if (claimants >= 0)
+        {
+            int token = info[1];
+            unsigned long long bestKey = 0ull;
+            int const walkers = min(claimants, kVpClaimants);
+            for (int c = 0; c < walkers; ++c)
+            {
+                uint32_t cand = kVpNone;
+                for (int r = 0; r < p.ranks && cand == kVpNone; ++r)
+                {
+                    cand = words[r * kVpClaimants + c];
+                }
+                if (cand == kVpNone)
+                {
+                    continue;
+                }
+                int const walker = info[2 + c];
+                int const v = static_cast<int>(cand) / 4;
+                int const position = ((v - walker) / vb) * 4 + static_cast<int>(cand) % 4 + 1;
+                unsigned long long const key
+                    = (static_cast<unsigned long long>(position) << 32) | static_cast<unsigned long long>(walker + 1);
+                if (walkers == 1 || key > bestKey)
+                {
+                    bestKey = key;
+                    token = static_cast<int>(cand);
+                }
+            }
+            p.fs.outputTokens[row] = token;
+        }
+    }
+    gridDependencyTrigger();
+}
+
+//! vp_argmax: one warp per row; the greedy token as int64. PDL: waits on vp_fold first.
+__global__ void __launch_bounds__(32) vpArgmaxKernel(VpParams p, int slots)
+{
+    gridDependencyWait();
+    int const row = blockIdx.x;
+    int const lane = threadIdx.x;
+    VpMaxArg best{-INFINITY, INT_MAX};
+    for (int i = lane; i < p.ranks * slots; i += 32)
+    {
+        uint2* const slot = vpGreedy(p.copy, p, i / slots, row, i % slots);
+        uint4 const w = vpAwait<2>(slot, p.status, p.timeoutNs);
+        storeRelaxedSys(slot, make_uint2(kSplitSentinel, kSplitSentinel));
+        VpMaxArg const cand{__uint_as_float(w.x), static_cast<int>(w.y)};
+        best = VpMaxArgOp()(best, cand);
+    }
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+    {
+        VpMaxArg const other{
+            __shfl_xor_sync(0xffffffffu, best.value, offset), __shfl_xor_sync(0xffffffffu, best.index, offset)};
+        best = VpMaxArgOp()(best, other);
+    }
+    if (lane == 0)
+    {
+        p.greedyTokens[row] = static_cast<int64_t>(best.index);
+    }
+    gridDependencyTrigger();
+}
+
 } // namespace fusedSampling
 } // namespace kernels
 } // namespace tensorrt_llm
@@ -837,6 +1401,193 @@ void watch(at::Tensor const& buf, at::Tensor const& firstSeen, int64_t blocks, i
         firstSeen.numel(), static_cast<unsigned long long>(timeoutMs) * 1000000ull);
 }
 
+//! The vocabulary-parallel head's common parameters: rank geometry and the exchange buffer.
+fs::VpParams vpLayout(int64_t rank, int64_t ranks, int64_t shardVecs, int64_t copy, at::IntArrayRef publish,
+    int64_t stateStride, int64_t greedyOffset, int64_t greedyStride, int64_t candOffset, int64_t candStride,
+    int64_t timeoutS)
+{
+    TORCH_CHECK(ranks >= 1 && ranks <= fs::kVpMaxRanks && rank >= 0 && rank < ranks, "rank / ranks out of range");
+    TORCH_CHECK(publish.size() >= 1 && publish.size() <= static_cast<size_t>(fs::kVpMaxRanks), "1..8 publish targets");
+    TORCH_CHECK(shardVecs >= 1, "empty shard");
+    fs::VpParams p;
+    p.rank = static_cast<int>(rank);
+    p.ranks = static_cast<int>(ranks);
+    p.shardVecs = static_cast<int>(shardVecs);
+    p.copy = reinterpret_cast<char*>(copy);
+    for (size_t i = 0; i < publish.size(); ++i)
+    {
+        TORCH_CHECK(publish[i] % 16 == 0, "publish targets must be 16-byte aligned");
+        p.publish[i] = reinterpret_cast<char*>(publish[i]);
+    }
+    p.numPublish = static_cast<int>(publish.size());
+    TORCH_CHECK(copy % 16 == 0 && greedyOffset % 16 == 0 && candOffset % 16 == 0, "exchange regions 16-byte aligned");
+    p.stateStride = stateStride;
+    p.greedyOffset = greedyOffset;
+    p.greedyStride = greedyStride;
+    p.candOffset = candOffset;
+    p.candStride = candStride;
+    p.timeoutNs = static_cast<unsigned long long>(timeoutS) * 1000000000ull;
+    return p;
+}
+
+void checkVpRows(
+    int64_t rows, int64_t shardVecs, int64_t ranks, int64_t stateStride, int64_t greedyStride, int64_t candStride)
+{
+    TORCH_CHECK(rows >= 1 && rows <= fs::kSplitMaxRows, "1..", fs::kSplitMaxRows, " rows");
+    int const vb = fs::splitVirtualBlock(static_cast<int>(rows));
+    TORCH_CHECK((shardVecs + vb - 1) / vb <= fs::kVpWalkVectors, "shard too wide for the walk stage");
+    TORCH_CHECK(rows * vb <= stateStride, "rows exceed the state slots");
+    TORCH_CHECK(rows * fs::kVpGreedySlots <= greedyStride, "rows exceed the greedy slots");
+    TORCH_CHECK(rows * fs::kVpClaimants <= candStride, "rows exceed the candidate slots");
+    TORCH_CHECK(ranks * shardVecs * 4 <= INT32_MAX, "vocabulary too large");
+}
+
+//! vp_fold: this rank's shard folded per (row, virtual thread), published to every rank.
+void vpFold(at::Tensor const& local, at::Tensor const& temperatures, at::Tensor const& topKs, at::Tensor const& topPs,
+    at::Tensor const& minPs, int64_t staging, at::IntArrayRef publish, int64_t rank, int64_t ranks, int64_t stateStride,
+    int64_t greedyOffset, int64_t greedyStride, int64_t candStride, bool greedy, double multiplier, double invCap,
+    double cap, bool pdl)
+{
+    TORCH_CHECK(local.scalar_type() == at::kBFloat16 && local.dim() == 2 && local.is_contiguous() && local.is_cuda(),
+        "local must be a contiguous CUDA BF16 [rows, shard] tensor");
+    TORCH_CHECK(local.size(1) % 4 == 0, "the shard must be a whole number of 4-element vectors");
+    TORCH_CHECK(reinterpret_cast<uintptr_t>(local.data_ptr()) % 8 == 0, "local must be 8-byte aligned");
+    auto const rows = local.size(0);
+    auto const shardVecs = local.size(1) / 4;
+    checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride);
+    at::cuda::CUDAGuard guard(local.device());
+    auto p = vpLayout(
+        rank, ranks, shardVecs, publish[0], publish, stateStride, greedyOffset, greedyStride, 0, candStride, 0);
+    p.fs.numRows = static_cast<int32_t>(rows);
+    p.fs.vocabSize = static_cast<int32_t>(ranks * shardVecs * 4);
+    if (!greedy)
+    {
+        checkRowVector(temperatures, "temperatures", at::kFloat, rows);
+        checkRowVector(topKs, "top_ks", at::kInt, rows);
+        checkRowVector(topPs, "top_ps", at::kFloat, rows);
+        checkRowVector(minPs, "min_ps", at::kFloat, rows);
+        p.fs.temperatures = temperatures.data_ptr<float>();
+        p.fs.topKs = topKs.data_ptr<int32_t>();
+        p.fs.topPs = topPs.data_ptr<float>();
+        p.fs.minPs = minPs.data_ptr<float>();
+        TORCH_CHECK(staging % 8 == 0, "staging must be 8-byte aligned");
+    }
+    p.local = reinterpret_cast<uint2 const*>(local.data_ptr());
+    p.staging = reinterpret_cast<uint2*>(staging);
+    p.cap = softCap(multiplier, invCap, cap);
+    int const vb = fs::splitVirtualBlock(static_cast<int>(rows));
+    dim3 const grid(static_cast<unsigned>(rows), vb / fs::kVpFoldThreads);
+    if (vb == fs::kSplitBlock)
+    {
+        if (greedy)
+        {
+            launch(fs::vpFoldKernel<true, fs::kSplitBlock>, grid, dim3(fs::kVpFoldThreads), pdl, p);
+        }
+        else
+        {
+            launch(fs::vpFoldKernel<false, fs::kSplitBlock>, grid, dim3(fs::kVpFoldThreads), pdl, p);
+        }
+    }
+    else if (greedy)
+    {
+        launch(fs::vpFoldKernel<true, fs::kSplitNarrowBlock>, grid, dim3(fs::kVpFoldThreads), pdl, p);
+    }
+    else
+    {
+        launch(fs::vpFoldKernel<false, fs::kSplitNarrowBlock>, grid, dim3(fs::kVpFoldThreads), pdl, p);
+    }
+}
+
+//! vp_draw: combine, reduce, scan, draw, walk this rank's portions; publishes candidates.
+//! Returns the token tensor vp_decide completes (filtered rows are written here).
+at::Tensor vpDraw(at::Tensor const& pending, at::Tensor const& local, at::Tensor const& temperatures,
+    at::Tensor const& topKs, at::Tensor const& topPs, at::Tensor const& minPs, std::optional<at::Tensor> const& seed,
+    std::optional<at::Tensor> const& offset, at::Tensor const& info, at::Tensor const& status, int64_t copy,
+    at::IntArrayRef stagingOf, at::IntArrayRef publish, int64_t rank, int64_t ranks, int64_t stateStride,
+    int64_t greedyStride, int64_t candOffset, int64_t candStride, int64_t timeoutS, double multiplier, double invCap,
+    double cap, bool pdl)
+{
+    TORCH_CHECK(pending.scalar_type() == at::kBFloat16, "pending must be BF16");
+    TORCH_CHECK(
+        local.scalar_type() == at::kBFloat16 && local.dim() == 2 && local.is_contiguous(), "local [rows, shard]");
+    at::cuda::CUDAGuard guard(pending.device());
+    auto fsParams = rowParams(pending, temperatures, topKs, topPs, minPs, seed, offset);
+    auto const rows = pending.size(0);
+    auto const shardVecs = local.size(1) / 4;
+    TORCH_CHECK(local.size(0) == rows && local.size(1) % 4 == 0 && pending.size(1) == ranks * local.size(1),
+        "pending is [rows, ranks * shard] of local [rows, shard]");
+    checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride);
+    TORCH_CHECK(static_cast<int64_t>(stagingOf.size()) == ranks, "one staging address per rank");
+    TORCH_CHECK(info.is_cuda() && info.scalar_type() == at::kInt && info.is_contiguous()
+            && info.numel() >= rows * fs::kVpInfoWords,
+        "info holds kVpInfoWords int32 per row");
+    TORCH_CHECK(status.is_cuda() && status.scalar_type() == at::kInt && status.numel() >= 1, "status int32[1]");
+    auto p = vpLayout(
+        rank, ranks, shardVecs, copy, publish, stateStride, 0, greedyStride, candOffset, candStride, timeoutS);
+    p.fs = fsParams;
+    for (int64_t r = 0; r < ranks; ++r)
+    {
+        p.stagingOf[r] = reinterpret_cast<uint2 const*>(stagingOf[r]);
+    }
+    p.local = reinterpret_cast<uint2 const*>(local.data_ptr());
+    p.info = info.data_ptr<int32_t>();
+    p.status = status.data_ptr<int32_t>();
+    p.cap = softCap(multiplier, invCap, cap);
+    auto tokens = at::empty({rows}, pending.options().dtype(at::kInt));
+    p.fs.outputTokens = tokens.data_ptr<int32_t>();
+    if (fs::splitVirtualBlock(static_cast<int>(rows)) == fs::kSplitBlock)
+    {
+        launch(fs::vpDrawKernel<fs::kSplitBlock>, dim3(static_cast<unsigned>(rows)), dim3(fs::kSplitBlock), pdl, p);
+    }
+    else
+    {
+        launch(fs::vpDrawKernel<fs::kSplitNarrowBlock>, dim3(static_cast<unsigned>(rows)), dim3(fs::kSplitNarrowBlock),
+            pdl, p);
+    }
+    return tokens;
+}
+
+//! vp_decide: every rank's candidates -> the token of each unfiltered row, into `tokens`.
+void vpDecide(at::Tensor const& tokens, at::Tensor const& info, at::Tensor const& status, int64_t copy, int64_t rank,
+    int64_t ranks, int64_t shardVecs, int64_t stateStride, int64_t greedyStride, int64_t candOffset, int64_t candStride,
+    int64_t timeoutS, bool pdl)
+{
+    TORCH_CHECK(tokens.is_cuda() && tokens.scalar_type() == at::kInt && tokens.dim() == 1, "tokens int32 [rows]");
+    auto const rows = tokens.size(0);
+    checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride);
+    TORCH_CHECK(info.is_cuda() && info.scalar_type() == at::kInt && info.numel() >= rows * fs::kVpInfoWords, "info");
+    TORCH_CHECK(status.is_cuda() && status.scalar_type() == at::kInt && status.numel() >= 1, "status int32[1]");
+    at::cuda::CUDAGuard guard(tokens.device());
+    int64_t const self[1] = {copy};
+    auto p = vpLayout(rank, ranks, shardVecs, copy, at::IntArrayRef(self, 1), stateStride, 0, greedyStride, candOffset,
+        candStride, timeoutS);
+    p.fs.numRows = static_cast<int32_t>(rows);
+    p.fs.outputTokens = tokens.data_ptr<int32_t>();
+    p.info = info.data_ptr<int32_t>();
+    p.status = status.data_ptr<int32_t>();
+    launch(fs::vpDecideKernel, dim3(static_cast<unsigned>(rows)), dim3(32), pdl, p,
+        fs::splitVirtualBlock(static_cast<int>(rows)));
+}
+
+//! vp_argmax: every rank's greedy words -> torch.argmax's int64 token per row.
+at::Tensor vpArgmax(at::Tensor const& status, int64_t rows, int64_t copy, int64_t rank, int64_t ranks,
+    int64_t shardVecs, int64_t stateStride, int64_t greedyOffset, int64_t greedyStride, int64_t candStride,
+    int64_t timeoutS, bool pdl)
+{
+    checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride);
+    TORCH_CHECK(status.is_cuda() && status.scalar_type() == at::kInt && status.numel() >= 1, "status int32[1]");
+    at::cuda::CUDAGuard guard(status.device());
+    int64_t const self[1] = {copy};
+    auto p = vpLayout(rank, ranks, shardVecs, copy, at::IntArrayRef(self, 1), stateStride, greedyOffset, greedyStride,
+        0, candStride, timeoutS);
+    p.status = status.data_ptr<int32_t>();
+    auto tokens = at::empty({rows}, status.options().dtype(at::kLong));
+    p.greedyTokens = tokens.data_ptr<int64_t>();
+    launch(fs::vpArgmaxKernel, dim3(static_cast<unsigned>(rows)), dim3(32), pdl, p,
+        fs::splitVirtualBlock(static_cast<int>(rows)) / fs::kVpFoldThreads);
+    return tokens;
+}
+
 } // namespace
 
 TORCH_LIBRARY(musespark_open_head_sampler, m)
@@ -854,6 +1605,23 @@ TORCH_LIBRARY(musespark_open_head_sampler, m)
     m.def("ref_fold(Tensor capped, Tensor temperatures, Tensor(a!) states) -> ()");
     m.def("cap(Tensor x, float multiplier, float inv_cap, float cap) -> Tensor");
     m.def("watch(Tensor buf, Tensor(a!) first_seen, int blocks, int timeout_ms) -> ()");
+    m.def(
+        "vp_fold(Tensor local, Tensor temperatures, Tensor top_ks, Tensor top_ps, Tensor min_ps, int staging, "
+        "int[] publish, int rank, int ranks, int state_stride, int greedy_offset, int greedy_stride, "
+        "int cand_stride, bool greedy, float multiplier, float inv_cap, float cap, bool pdl) -> ()");
+    m.def(
+        "vp_draw(Tensor(a!) pending, Tensor local, Tensor temperatures, Tensor top_ks, Tensor top_ps, "
+        "Tensor min_ps, Tensor? seed, Tensor? offset, Tensor(b!) info, Tensor(c!) status, int copy, "
+        "int[] staging_of, int[] publish, int rank, int ranks, int state_stride, int greedy_stride, "
+        "int cand_offset, int cand_stride, int timeout_s, float multiplier, float inv_cap, float cap, "
+        "bool pdl) -> Tensor");
+    m.def(
+        "vp_decide(Tensor(a!) tokens, Tensor info, Tensor(b!) status, int copy, int rank, int ranks, "
+        "int shard_vecs, int state_stride, int greedy_stride, int cand_offset, int cand_stride, int timeout_s, "
+        "bool pdl) -> ()");
+    m.def(
+        "vp_argmax(Tensor(a!) status, int rows, int copy, int rank, int ranks, int shard_vecs, int state_stride, "
+        "int greedy_offset, int greedy_stride, int cand_stride, int timeout_s, bool pdl) -> Tensor");
 }
 
 TORCH_LIBRARY_IMPL(musespark_open_head_sampler, CUDA, m)
@@ -864,4 +1632,8 @@ TORCH_LIBRARY_IMPL(musespark_open_head_sampler, CUDA, m)
     m.impl("ref_fold", &refFold);
     m.impl("cap", &capRows);
     m.impl("watch", &watch);
+    m.impl("vp_fold", &vpFold);
+    m.impl("vp_draw", &vpDraw);
+    m.impl("vp_decide", &vpDecide);
+    m.impl("vp_argmax", &vpArgmax);
 }
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py
index 05a2c2ef7448..5a0db30bcdbe 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py
@@ -15,6 +15,11 @@
 ``draft_topk_merge`` over the sum of all ranks' outputs gives the result it
 gives over the per-chunk keys, bit for bit (selection, tie order and the
 selector's operands), from 1/13 of the bytes at TP4.
+
+``canon`` and ``canon_band`` (a drafter logit transform's canonical keys) go to
+the per-chunk stage unchanged; the rank's top ``top_k`` of canonical keys holds
+every key of the global top ``top_k`` under the transformed order for the same
+reason.
 """
 
 from __future__ import annotations
@@ -23,9 +28,9 @@
 import triton
 import triton.language as tl
 
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
 
-from .draft_shard_topk import CHUNK, NUM_WARPS, TOP_K, DraftShardTopK
+from .draft_shard_topk import CHUNK, FULL_BAND, NUM_WARPS, SHARD, TOP_K, DraftShardTopK
 
 #: Per-chunk keys one row program reduces.
 MAX_KEYS = 4096
@@ -78,6 +83,8 @@ class DraftRankTopK(OpWrapper):
         chunk      columns per program of the per-chunk stage (a power of two >= top_k)
         num_warps  warps per program of the per-chunk stage
         rank, num_ranks
+        canon, canon_band
+                   ``draft_shard_topk``'s canonical-key table and its identity band
 
     Returns ``[M, num_ranks * top_k * 2]`` fp32: rank ``rank``'s keys at columns
     ``[rank * top_k * 2, (rank + 1) * top_k * 2)`` as (high, low) 16-bit halves,
@@ -87,6 +94,20 @@ class DraftRankTopK(OpWrapper):
     """
 
     ARCHS = frozenset({Arch.SM_103, Arch.SM_107})
+    CELLS: tuple[Cell, ...] = (
+        Cell(
+            "five requests at TP8 (35 rows, the fewest wide_draft_topk takes)",
+            {"rows": 35, "shard": SHARD, "transform": None},
+        ),
+        Cell(
+            "five requests at TP8, the Muse Spark Open drafter transform's canon",
+            {"rows": 35, "shard": SHARD, "transform": "musespark"},
+        ),
+        Cell(
+            "24 requests at TP4 (a 50512-column shard), the drafter transform",
+            {"rows": 168, "shard": 2 * SHARD, "transform": "musespark"},
+        ),
+    )
     note = (
         "Used with draft_topk_merge after the draft head's MNNVL fp32 all-reduce "
         "(wide_draft_topk's RankTopKExchange, from 33 rows): the sum of "
@@ -103,9 +124,13 @@ def __call__(
         num_warps: int = NUM_WARPS,
         rank: int = 0,
         num_ranks: int = 1,
+        canon: torch.Tensor | None = None,
+        canon_band: tuple[int, int] = FULL_BAND,
     ) -> torch.Tensor:
         rows = logits.shape[0]
-        keys = _draft_shard_topk(logits, top_k, chunk, num_warps)
+        keys = _draft_shard_topk(
+            logits, top_k, chunk, num_warps, canon=canon, canon_band=canon_band
+        )
         out = torch.empty((rows, num_ranks * top_k * 2), dtype=torch.float32, device=logits.device)
         _draft_rank_topk_kernel[(rows,)](
             keys,
@@ -129,11 +154,15 @@ def reference(
         num_warps: int = NUM_WARPS,
         rank: int = 0,
         num_ranks: int = 1,
+        canon: torch.Tensor | None = None,
+        canon_band: tuple[int, int] = FULL_BAND,
     ) -> torch.Tensor:
         """The rank's top ``top_k`` keys over the whole shard, from torch ops; geometry ignored."""
         rows = logits.shape[0]
         # One chunk spanning the shard: draft_shard_topk's encoding of the rank's top_k.
-        keys = _draft_shard_topk.reference(logits, top_k, 1 << 16).to(torch.int64)
+        keys = _draft_shard_topk.reference(
+            logits, top_k, 1 << 16, canon=canon, canon_band=canon_band
+        ).to(torch.int64)
         unsigned = torch.where(keys < 0, keys + 2**32, keys)
         halves = torch.stack(((unsigned >> 16) & 0xFFFF, unsigned & 0xFFFF), dim=-1).float()
         packed = torch.zeros((rows, num_ranks, top_k, 2), device=logits.device)
@@ -148,9 +177,13 @@ def is_valid(
         num_warps: int = NUM_WARPS,
         rank: int = 0,
         num_ranks: int = 1,
+        canon: torch.Tensor | None = None,
+        canon_band: tuple[int, int] = FULL_BAND,
         **unused,
     ) -> None:
-        _draft_shard_topk.is_valid(logits, top_k, chunk, num_warps)
+        _draft_shard_topk.is_valid(
+            logits, top_k, chunk, num_warps, canon=canon, canon_band=canon_band
+        )
         if not 0 <= rank < num_ranks:
             raise ValueError(f"rank {rank} outside [0, {num_ranks})")
         if triton.cdiv(logits.shape[1], chunk) * top_k > MAX_KEYS:
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py
index 1a7c336c10c6..65a206fc5c32 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py
@@ -20,6 +20,17 @@
 
 whose signed order is (value desc, column asc) and which is unique per element.
 ``INT32_MIN`` marks an empty slot; no element encodes to it.
+
+With ``canon`` (the canonical-key table of a drafter logit transform ``t``,
+``kernels/draft_head/logit_transform.py``) the keys order the slice as
+``torch.topk`` orders the transformed fp32 logits ``t(x)``. Every key16 outside
+the band ``canon_band = (lo, hi)`` is replaced by ``canon[key16]``, the key16 of
+its tie group's representative (a tie group is the bf16 values whose
+transformed values have the same bits). ``t`` never decreases, so a tie group
+is a run of consecutive key16 values and ``canon`` is non-decreasing: two
+canonical keys compare as the radix keys of the transformed values do, and
+equal ones tie on the column as torch's do. Inside the band every group is a
+single value and ``canon`` is the identity, so those elements load nothing.
 """
 
 from __future__ import annotations
@@ -28,7 +39,7 @@
 import triton
 import triton.language as tl
 
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
 
 #: The draft head this target runs: 202048 / 8 vocabulary rows per rank, K + 1
 #: = 8 block positions with K <= 7 drafted, top 16 candidates.
@@ -42,20 +53,33 @@
 #: The empty-slot key.
 EMPTY_KEY = -(2**31)
 
+#: Entries of a canonical-key table: one per 16-bit radix key.
+CANON_ENTRIES = 1 << 16
+
+#: The band of a call without a table: every key16 lies inside it.
+FULL_BAND = (0, CANON_ENTRIES - 1)
+
+#: The non-NaN 16-bit radix keys, -inf (0x007F) to +inf (0xFF80).
+FINITE_KEYS = (0x007F, 0xFF80)
+
 
 @triton.jit
 def _draft_shard_topk_kernel(
     bits_ptr,
     keys_ptr,
+    canon_ptr,
     shard,
     bits_row_stride,
     keys_row_stride,
     rank,
     num_ranks,
     keys_per_rank,
+    canon_lo,
+    canon_hi,
     CHUNK: tl.constexpr,
     TOP_K: tl.constexpr,
     HALVES: tl.constexpr,
+    CANON: tl.constexpr,
 ):
     row = tl.program_id(0)
     chunk = tl.program_id(1)
@@ -66,6 +90,10 @@ def _draft_shard_topk_kernel(
     key16 = tl.where((bits & 0x8000) != 0, bits ^ 0xFFFF, bits ^ 0x8000)
     is_nan = ((bits & 0x7F80) == 0x7F80) & ((bits & 0x7F) != 0)
     key16 = tl.where(is_nan, 0xFFFF, key16)
+    if CANON:
+        # Only the keys outside the identity band read the table.
+        far = ((key16 < canon_lo) | (key16 > canon_hi)) & in_shard
+        key16 = tl.where(far, tl.load(canon_ptr + key16, mask=far, other=0), key16)
     key = ((key16 ^ 0x8000) << 16) | (0xFFFF - local)
     key = tl.where(in_shard, key, -2147483648)
 
@@ -91,34 +119,91 @@ def _draft_shard_topk_kernel(
         tl.store(keys_ptr + row.to(tl.int64) * keys_row_stride + chunk * TOP_K + slots, out)
 
 
-def _key16(bits: torch.Tensor) -> torch.Tensor:
-    """torch's radix key of bf16 bit patterns, as int32 in [0, 0xFFFF]."""
+def radix_key16(bits: torch.Tensor) -> torch.Tensor:
+    """torch's radix key of bf16 bit patterns, as int32 in [0, 0xFFFF]; every NaN is 0xFFFF."""
     b = bits.to(torch.int32) & 0xFFFF
     key = torch.where((b & 0x8000) != 0, b ^ 0xFFFF, b ^ 0x8000)
     is_nan = ((b & 0x7F80) == 0x7F80) & ((b & 0x7F) != 0)
     return torch.where(is_nan, torch.full_like(key, 0xFFFF), key)
 
 
+def canon_reason(canon: torch.Tensor, canon_band: tuple[int, int], device) -> str | None:
+    """Why ``canon`` is not a canonical-key table the kernels read correctly, or None.
+
+    A table is ``[65536]`` int32, contiguous, on ``device``; holds 16-bit keys;
+    is non-decreasing over the non-NaN keys (``FINITE_KEYS``), so it keeps
+    their order; is the identity on ``canon_band`` (the kernels skip the load
+    there); and maps 0xFFFF (every NaN) to itself.
+    """
+    if canon.dtype != torch.int32 or tuple(canon.shape) != (CANON_ENTRIES,):
+        return f"canon must be [{CANON_ENTRIES}] int32, got {canon.dtype} {tuple(canon.shape)}"
+    if not canon.is_contiguous() or canon.device != torch.device(device):
+        return "canon must be contiguous and on the logits' device"
+    lo, hi = canon_band
+    if not 0 <= lo <= hi < CANON_ENTRIES:
+        return f"canon_band {canon_band} is not an interval of 16-bit keys"
+    table = canon.cpu()
+    if int(table.min()) < 0 or int(table.max()) >= CANON_ENTRIES:
+        return "canon holds a value outside the 16-bit keys"
+    ordered = table[FINITE_KEYS[0] : FINITE_KEYS[1] + 1]
+    if bool((ordered[1:] < ordered[:-1]).any()):
+        return "canon is not non-decreasing over the non-NaN keys"
+    if not torch.equal(table[lo : hi + 1], torch.arange(lo, hi + 1, dtype=torch.int32)):
+        return f"canon is not the identity on its band {tuple(canon_band)}"
+    if int(table[CANON_ENTRIES - 1]) != CANON_ENTRIES - 1:
+        return "canon must map 0xFFFF (NaN) to itself"
+    return None
+
+
 class DraftShardTopK(OpWrapper):
     """Top ``top_k`` sort keys of every ``chunk`` columns of a bf16 logits slice.
 
-        logits   [M, shard] bf16, stride(-1) == 1, row stride free; shard <= 65536
-        top_k    candidates kept per chunk (a power of two)
-        chunk    columns per program (a power of two >= top_k)
+        logits      [M, shard] bf16, stride(-1) == 1, row stride free; shard <= 65536
+        top_k       candidates kept per chunk (a power of two)
+        chunk       columns per program (a power of two >= top_k)
         num_warps
+        rank, num_ranks
+                    with ``rank`` set, the all-reduce layout (see below)
+        canon       None, or a canonical-key table ``[65536]`` int32 (``canon_reason``)
+        canon_band  ``(lo, hi)``: the key16 interval on which ``canon`` is the identity
 
     Returns ``[M, ceil(shard / chunk) * top_k]`` int32: chunk ``c``'s keys at
     columns ``[c * top_k, (c + 1) * top_k)``, largest first, ``EMPTY_KEY`` where
-    the chunk has fewer than ``top_k`` columns. Integer selection: the output
-    is exact, not approximate.
+    the chunk has fewer than ``top_k`` columns. With ``rank`` set, ``[M,
+    num_ranks * ceil(shard / chunk) * top_k * 2]`` fp32: this rank's keys as
+    (high, low) 16-bit halves at its slot, +0.0 in every other rank's. Integer
+    selection: the output is exact, not approximate.
     """
 
     ARCHS = frozenset({Arch.SM_103})
+    CELLS: tuple[Cell, ...] = (
+        Cell(
+            "one request at TP8: 7 rows of a 25256-column shard, int32 keys for the all-gather",
+            {"rows": 7, "shard": SHARD, "halves": False, "transform": None},
+        ),
+        Cell(
+            "one request at TP8, fp32 halves for the MNNVL all-reduce",
+            {"rows": 7, "shard": SHARD, "halves": True, "transform": None},
+        ),
+        Cell(
+            "one request at TP8, fp32 halves, the Muse Spark Open drafter transform's canon",
+            {"rows": 7, "shard": SHARD, "halves": True, "transform": "musespark"},
+        ),
+        Cell(
+            "eight requests at TP4: a 50512-column shard, int32 keys, the drafter transform",
+            {"rows": 56, "shard": 2 * SHARD, "halves": False, "transform": "musespark"},
+        ),
+        Cell(
+            "24 requests at TP8, fp32 halves, the drafter transform",
+            {"rows": 168, "shard": SHARD, "halves": True, "transform": "musespark"},
+        ),
+    )
     note = (
         "Used with draft_topk_merge, which decodes the keys; the pair is exact "
-        "against torch.topk over the gathered fp32 logits. NaN payloads are not "
-        "carried (a NaN logit ranks first, as in torch, and decodes to the "
-        "canonical NaN)."
+        "against torch.topk over the gathered fp32 logits, and with canon (and the "
+        "merge's values table) against torch.topk over the drafter's transformed "
+        "logits. NaN payloads are not carried (a NaN logit ranks first, as in "
+        "torch, and decodes to the canonical NaN)."
     )
 
     def __call__(
@@ -129,6 +214,8 @@ def __call__(
         num_warps: int = NUM_WARPS,
         rank: int | None = None,
         num_ranks: int = 1,
+        canon: torch.Tensor | None = None,
+        canon_band: tuple[int, int] = FULL_BAND,
     ) -> torch.Tensor:
         rows, shard = logits.shape
         chunks = triton.cdiv(shard, chunk)
@@ -141,18 +228,23 @@ def __call__(
         else:
             keys = torch.empty((rows, keys_per_rank), dtype=torch.int32, device=logits.device)
         bits = logits.view(torch.int16)
+        lo, hi = canon_band
         _draft_shard_topk_kernel[(rows, chunks)](
             bits,
             keys,
+            bits if canon is None else canon,
             shard,
             bits.stride(0),
             keys.stride(0),
             rank if halves else 0,
             num_ranks,
             keys_per_rank,
+            lo,
+            hi,
             CHUNK=chunk,
             TOP_K=top_k,
             HALVES=halves,
+            CANON=canon is not None,
             num_warps=num_warps,
         )
         return keys
@@ -165,11 +257,18 @@ def reference(
         num_warps: int = NUM_WARPS,
         rank: int | None = None,
         num_ranks: int = 1,
+        canon: torch.Tensor | None = None,
+        canon_band: tuple[int, int] = FULL_BAND,
     ) -> torch.Tensor:
         """The same keys from torch ops; ``num_warps`` is ignored."""
         rows, shard = logits.shape
         local = torch.arange(shard, device=logits.device, dtype=torch.int64)
-        key16 = _key16(logits.contiguous().view(torch.int16)).to(torch.int64)
+        key16 = radix_key16(logits.contiguous().view(torch.int16))
+        if canon is not None:
+            lo, hi = canon_band
+            far = (key16 < lo) | (key16 > hi)
+            key16 = torch.where(far, canon.to(key16.device)[key16.long()], key16)
+        key16 = key16.to(torch.int64)
         # The unsigned 32-bit key: (value desc, column asc) as plain integer order.
         key32 = (key16 << 16) | (0xFFFF - local)
         out = []
@@ -203,6 +302,8 @@ def is_valid(
         num_warps: int = NUM_WARPS,
         rank: int | None = None,
         num_ranks: int = 1,
+        canon: torch.Tensor | None = None,
+        canon_band: tuple[int, int] = FULL_BAND,
         **unused,
     ) -> None:
         if logits.dtype != torch.bfloat16 or logits.dim() != 2:
@@ -217,3 +318,7 @@ def is_valid(
             )
         if rank is not None and not 0 <= rank < num_ranks:
             raise ValueError(f"rank {rank} outside [0, {num_ranks})")
+        if canon is not None:
+            reason = canon_reason(canon, canon_band, logits.device)
+            if reason is not None:
+                raise ValueError(reason)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py
index 97117280ab01..9dbca5a955a7 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py
@@ -25,6 +25,16 @@
   Both are exact: a bf16 widens to fp32 exactly and a bf16 * bf16 product
   fits fp32's significand, so these equal the built-in tensors bit for bit,
   and the built-in einsum over them is unchanged.
+
+With a drafter logit transform ``t`` the keys are canonical
+(``draft_shard_topk``'s ``canon``) and ``values`` is the transform's table
+(``kernels/draft_head/logit_transform.py``): ``values[key16]`` is ``t`` of the
+bf16 value of ``key16``, computed by the drafter's own method. The selection,
+the arrangement and the network above run on canonical keys and their decoded
+bf16 values unchanged, which compare as the transformed values do, so they give
+``torch.topk(t(logits))``'s ids and order; the unary values are then
+``values[key16]``, bitwise ``t`` of every candidate, since all members of a tie
+group transform to the same bits.
 """
 
 from __future__ import annotations
@@ -33,7 +43,9 @@
 import triton
 import triton.language as tl
 
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+from .draft_shard_topk import CANON_ENTRIES, SHARD, radix_key16
 
 #: ``torch.topk`` sorts k <= 32 results with this many bitonic slots.
 SORT_SLOTS = 32
@@ -93,6 +105,7 @@ def _draft_topk_merge_kernel(
     unary_ptr,
     pred_out_ptr,
     succ_out_ptr,
+    values_ptr,
     rows_per_rank,
     num_ranks,
     num_keys,
@@ -110,6 +123,7 @@ def _draft_topk_merge_kernel(
     SLOTS: tl.constexpr,
     RANK_BLOCK: tl.constexpr,
     HALVES: tl.constexpr,
+    TRANSFORM: tl.constexpr,
 ):
     # Every program of a row runs the (small) merge; each gathers one RANK_BLOCK
     # slice of the codebook rows, so the gathers spread over more SMs.
@@ -194,7 +208,11 @@ def _draft_topk_merge_kernel(
     ids = result >> 16
     if part == 0:
         tl.store(ids_ptr + row * TOP_K + kk, ids)
-        tl.store(unary_ptr + row * TOP_K + kk, _decode_value(result & 0xFFFF))
+        if TRANSFORM:
+            # The drafter transform of each candidate: one table entry per tie group.
+            tl.store(unary_ptr + row * TOP_K + kk, tl.load(values_ptr + (result & 0xFFFF)))
+        else:
+            tl.store(unary_ptr + row * TOP_K + kk, _decode_value(result & 0xFFFF))
 
     # Selector operands for this row (successors) and the next one (predecessors).
     out_rows = (row * TOP_K + kk[:, None]).to(tl.int64) * rank_dim
@@ -252,15 +270,77 @@ class DraftTopKMerge(OpWrapper):
         successor_codebook    [vocab, rank] bf16, stride(-1) == 1
         top_k                 <= 32, a power of two
         num_ranks             ranks in the gather (the tensor-parallel size)
+        rank_block, num_warps launch geometry
+        values                None, or a drafter logit transform's table [65536]
+                              fp32 (key16 -> transformed value); the keys are then
+                              ``draft_shard_topk``'s canonical keys of that transform
 
     Returns ``(candidate_ids [B, K, top_k] int64, unary [B, K, top_k] fp32,
     predecessors [B, K, top_k, rank] fp32, successors [B, K, top_k, rank] fp32)``:
     ``candidate_ids`` / ``unary`` are ``torch.topk(full_logits.float(), top_k)``'s
-    indices / values, and the operands are ``dflash2_score_edges``'s
+    indices / values -- with ``values``, ``torch.topk(t(full_logits))``'s, ``t``
+    the transform -- and the operands are ``dflash2_score_edges``'s
     ``predecessors * gate.float()[:, :, None]`` and ``successors``. Exact.
     """
 
     ARCHS = frozenset({Arch.SM_103})
+    CELLS: tuple[Cell, ...] = (
+        Cell(
+            "one request at TP8, int32 all-gather keys",
+            {
+                "requests": 1,
+                "positions": 7,
+                "shard": SHARD,
+                "ranks": 8,
+                "layout": "int32",
+                "transform": None,
+            },
+        ),
+        Cell(
+            "one request at TP8, fp32 all-reduce halves",
+            {
+                "requests": 1,
+                "positions": 7,
+                "shard": SHARD,
+                "ranks": 8,
+                "layout": "halves",
+                "transform": None,
+            },
+        ),
+        Cell(
+            "one request at TP8, fp32 halves, the Muse Spark Open drafter transform",
+            {
+                "requests": 1,
+                "positions": 7,
+                "shard": SHARD,
+                "ranks": 8,
+                "layout": "halves",
+                "transform": "musespark",
+            },
+        ),
+        Cell(
+            "two requests at TP4, int32 keys, the drafter transform",
+            {
+                "requests": 2,
+                "positions": 7,
+                "shard": 2 * SHARD,
+                "ranks": 4,
+                "layout": "int32",
+                "transform": "musespark",
+            },
+        ),
+        Cell(
+            "24 requests at TP8, rank top-k halves, the drafter transform",
+            {
+                "requests": 24,
+                "positions": 7,
+                "shard": SHARD,
+                "ranks": 8,
+                "layout": "rank",
+                "transform": "musespark",
+            },
+        ),
+    )
     note = (
         "Tie order is torch 2.14's (SortUtils.cuh bitonicSort, k <= 32 -> 32 "
         "slots), which torch does not document. draft_head.install checks it "
@@ -282,6 +362,7 @@ def __call__(
         num_ranks: int,
         rank_block: int = RANK_BLOCK,
         num_warps: int = NUM_WARPS,
+        values: torch.Tensor | None = None,
     ):
         rows = num_requests * num_positions
         halves = gathered_keys.dtype == torch.float32
@@ -307,6 +388,7 @@ def __call__(
             unary,
             preds,
             succs,
+            unary if values is None else values,
             rows,
             num_ranks,
             num_keys,
@@ -324,6 +406,7 @@ def __call__(
             SLOTS=SORT_SLOTS,
             RANK_BLOCK=rank_block,
             HALVES=halves,
+            TRANSFORM=values is not None,
             num_warps=num_warps,
         )
         return ids, unary, preds, succs
@@ -342,6 +425,7 @@ def reference(
         num_ranks: int,
         rank_block: int = RANK_BLOCK,
         num_warps: int = NUM_WARPS,
+        values: torch.Tensor | None = None,
     ):
         """``torch.topk`` over the candidates scattered into a -inf vocabulary row.
 
@@ -349,21 +433,28 @@ def reference(
 
         Every element outside the candidate lists is below the k-th value or
         tied with it at a higher id, so replacing it with -inf changes neither
-        the selection nor the arrangement torch sorts.
+        the selection nor the arrangement torch sorts. With ``values`` the
+        candidates carry the decoded canonical values, which order and tie as
+        the transformed values do; each selected value is then replaced by
+        ``values`` at its key16.
         """
         rows = num_requests * num_positions
         ranks = num_ranks
         if gathered_keys.dtype == torch.float32:
             gathered_keys = unpack_halves(gathered_keys, ranks)
         vocab_size = ranks * shard
-        vocab, values = _decode_keys(gathered_keys, ranks, shard)
+        vocab, decoded = _decode_keys(gathered_keys, ranks, shard)
         dense = torch.full((rows, vocab_size), float("-inf"), device=gathered_keys.device)
         row_index = torch.arange(rows, device=gathered_keys.device)[:, None].expand_as(vocab[:rows])
         for r in range(ranks):
-            ids_r, vals_r = vocab[r * rows : (r + 1) * rows], values[r * rows : (r + 1) * rows]
+            ids_r, vals_r = vocab[r * rows : (r + 1) * rows], decoded[r * rows : (r + 1) * rows]
             keep = ids_r >= 0
             dense[row_index[keep], ids_r[keep]] = vals_r[keep]
         unary, ids = torch.topk(dense, top_k, dim=-1)
+        if values is not None:
+            # Decoded canonical values are bf16-exact: their key16 indexes the table.
+            key16 = radix_key16(unary.to(torch.bfloat16).view(torch.int16))
+            unary = values.to(unary.device)[key16.long()]
         ids = ids.view(num_requests, num_positions, top_k)
         unary = unary.view(num_requests, num_positions, top_k)
         successors = successor_codebook[ids].float()
@@ -385,8 +476,20 @@ def is_valid(
         successor_codebook: torch.Tensor,
         top_k: int,
         num_ranks: int,
+        rank_block: int = RANK_BLOCK,
+        num_warps: int = NUM_WARPS,
+        values: torch.Tensor | None = None,
         **unused,
     ) -> None:
+        if values is not None and (
+            values.dtype != torch.float32
+            or tuple(values.shape) != (CANON_ENTRIES,)
+            or not values.is_contiguous()
+            or values.device != gathered_keys.device
+        ):
+            raise ValueError(
+                f"values must be a contiguous [{CANON_ENTRIES}] fp32 table on the keys' device"
+            )
         rows = num_requests * num_positions
         ranks = num_ranks
         if gathered_keys.dtype == torch.int32:
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py
index f694783969ff..08a01afdab61 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py
@@ -598,3 +598,468 @@ def shared_gather_workspace(mapping: Mapping, shard: int) -> HeadGatherWorkspace
         workspace = HeadGatherWorkspace(mapping, shard)
         _SHARED[key] = workspace
     return workspace
+
+
+# ---------------------------------------------------------------------------------------
+# The vocabulary-parallel head (fast path ``vocab_parallel_head``): no logits all-gather.
+
+#: Fewest rows the vocabulary-parallel head takes (two K=7 verify blocks); 8-row steps keep
+#: the ``lm_head_inbox`` gather.
+VP_MIN_ROWS = 16
+#: Claimants of a row whose crossings the ranks exchange (``kVpClaimants``).
+VP_CLAIMANTS = 4
+#: Greedy words per (rank, row): one per ``vp_fold`` CTA of the 1024-thread block.
+VP_GREEDY_SLOTS = VIRTUAL_THREADS // FOLD_THREADS
+#: ``vp_draw``'s per-row record for ``vp_decide`` (int32 words).
+VP_INFO_WORDS = 2 + VP_CLAIMANTS
+#: Most ranks (``kVpMaxRanks``).
+VP_MAX_RANKS = 8
+
+
+class VocabParallelLayout:
+    """The exchange buffer of the vocabulary-parallel head for ``ranks`` shards of ``shard``
+    elements and calls of at most ``max_rows`` rows.
+
+    One copy per rank, every 32-bit word the sentinel ``0xFFFFFFFF`` between calls:
+    ``[ranks][state_stride]`` 16-byte states (``vp_fold``'s per-virtual-thread partials), then
+    ``[ranks][greedy_stride]`` 8-byte greedy words at ``greedy_offset``, then
+    ``[ranks][cand_stride]`` 4-byte candidates at ``cand_offset``; ``nbytes`` in all.
+    """
+
+    def __init__(self, ranks: int, shard: int, max_rows: int):
+        if not 1 <= ranks <= VP_MAX_RANKS:
+            raise ValueError(f"1..{VP_MAX_RANKS} ranks, got {ranks}")
+        if shard % 4:
+            raise ValueError(f"shard {shard} is not a whole number of 4-element vectors")
+        if not VP_MIN_ROWS <= max_rows <= WIDE_MAX_ROWS:
+            raise ValueError(f"max_rows must be in {VP_MIN_ROWS}..{WIDE_MAX_ROWS}, got {max_rows}")
+        self.ranks = int(ranks)
+        self.shard = int(shard)
+        self.shard_vecs = self.shard // 4
+        self.vocab = self.ranks * self.shard
+        self.max_rows = int(max_rows)
+        widest = max(
+            -(-self.shard_vecs // virtual_threads(r)) for r in {min(max_rows, MAX_ROWS), max_rows}
+        )
+        if widest > 64:
+            raise ValueError(f"shard {shard} is too wide for the walk stage")
+        self.state_stride = max(r * virtual_threads(r) for r in {min(max_rows, MAX_ROWS), max_rows})
+        self.greedy_stride = self.max_rows * VP_GREEDY_SLOTS
+        self.cand_stride = self.max_rows * VP_CLAIMANTS
+        self.greedy_offset = self.ranks * self.state_stride * 16
+        self.cand_offset = self.greedy_offset + self.ranks * self.greedy_stride * 8
+        self.nbytes = -(-(self.cand_offset + self.ranks * self.cand_stride * 4) // 16) * 16
+        self.staging_bytes = self.max_rows * self.shard * 2
+
+    def strides(self) -> dict:
+        return dict(
+            state_stride=self.state_stride,
+            greedy_stride=self.greedy_stride,
+            cand_stride=self.cand_stride,
+        )
+
+
+class _VocabParallelRank:
+    """One rank's view: its exchange copy, the publish targets, every rank's staging rows."""
+
+    def __init__(
+        self,
+        layout: VocabParallelLayout,
+        rank: int,
+        copy: int,
+        publish: list,
+        staging: int,
+        staging_of: list,
+        info: torch.Tensor,
+        status: torch.Tensor,
+    ):
+        self.layout = layout
+        self.rank = int(rank)
+        self.copy = int(copy)
+        self.publish = [int(x) for x in publish]
+        self.staging = int(staging)
+        self.staging_of = [int(x) for x in staging_of]
+        self.info = info
+        self.status = status
+
+    def fold(self, local, temperatures, top_ks, top_ps, min_ps, multiplier, soft_cap, greedy, pdl):
+        m, inv, cap = cap_constants(multiplier, soft_cap)
+        lay = self.layout
+        torch.ops.musespark_open_head_sampler.vp_fold(
+            local,
+            temperatures,
+            top_ks,
+            top_ps,
+            min_ps,
+            self.staging,
+            self.publish,
+            self.rank,
+            lay.ranks,
+            lay.state_stride,
+            lay.greedy_offset,
+            lay.greedy_stride,
+            lay.cand_stride,
+            greedy,
+            m,
+            inv,
+            cap,
+            pdl,
+        )
+
+    def draw(
+        self,
+        pending,
+        local,
+        temperatures,
+        top_ks,
+        top_ps,
+        min_ps,
+        seed,
+        offset,
+        multiplier,
+        soft_cap,
+        pdl,
+    ):
+        m, inv, cap = cap_constants(multiplier, soft_cap)
+        lay = self.layout
+        return torch.ops.musespark_open_head_sampler.vp_draw(
+            pending,
+            local,
+            temperatures,
+            top_ks,
+            top_ps,
+            min_ps,
+            seed,
+            offset,
+            self.info,
+            self.status,
+            self.copy,
+            self.staging_of,
+            self.publish,
+            self.rank,
+            lay.ranks,
+            lay.state_stride,
+            lay.greedy_stride,
+            lay.cand_offset,
+            lay.cand_stride,
+            POLL_TIMEOUT_S,
+            m,
+            inv,
+            cap,
+            pdl,
+        )
+
+    def decide(self, tokens, pdl):
+        lay = self.layout
+        torch.ops.musespark_open_head_sampler.vp_decide(
+            tokens,
+            self.info,
+            self.status,
+            self.copy,
+            self.rank,
+            lay.ranks,
+            lay.shard_vecs,
+            lay.state_stride,
+            lay.greedy_stride,
+            lay.cand_offset,
+            lay.cand_stride,
+            POLL_TIMEOUT_S,
+            pdl,
+        )
+
+    def argmax(self, rows, pdl):
+        lay = self.layout
+        return torch.ops.musespark_open_head_sampler.vp_argmax(
+            self.status,
+            rows,
+            self.copy,
+            self.rank,
+            lay.ranks,
+            lay.shard_vecs,
+            lay.state_stride,
+            lay.greedy_offset,
+            lay.greedy_stride,
+            lay.cand_stride,
+            POLL_TIMEOUT_S,
+            pdl,
+        )
+
+
+def vocab_parallel_sample(
+    view: _VocabParallelRank,
+    pending: torch.Tensor,
+    local: torch.Tensor,
+    temperatures: torch.Tensor,
+    top_ks: torch.Tensor,
+    top_ps: torch.Tensor,
+    min_ps: torch.Tensor,
+    seed: Optional[torch.Tensor],
+    offset: Optional[torch.Tensor],
+    multiplier: float,
+    soft_cap: float,
+    pdl: bool = False,
+) -> torch.Tensor:
+    """One rank's sampled tokens (int32 ``[rows]``) of the step whose shard ``local`` is.
+
+    ``vp_fold``, ``vp_draw`` and ``vp_decide`` in order; every rank of the group makes the same
+    call with its own shard. ``pending`` is the ``[rows, vocab]`` BF16 row: written for rows with
+    a top-k / top-p / min-p filter only (gathered from the ranks' staging rows), left as it is
+    otherwise.
+    """
+    rows = local.shape[0]
+    view.fold(local, temperatures, top_ks, top_ps, min_ps, multiplier, soft_cap, False, pdl)
+    tokens = view.draw(
+        pending,
+        local,
+        temperatures,
+        top_ks,
+        top_ps,
+        min_ps,
+        seed,
+        offset,
+        multiplier,
+        soft_cap,
+        pdl,
+    )
+    view.decide(tokens[:rows], pdl)
+    return tokens
+
+
+def vocab_parallel_argmax(
+    view: _VocabParallelRank,
+    local: torch.Tensor,
+    multiplier: float,
+    soft_cap: float,
+    pdl: bool = False,
+) -> torch.Tensor:
+    """One rank's greedy tokens (int64 ``[rows]``): ``torch.argmax`` of the soft-capped FP32 row."""
+    view.fold(
+        local, view.status, view.status, view.status, view.status, multiplier, soft_cap, True, pdl
+    )
+    return view.argmax(local.shape[0], pdl)
+
+
+def vocab_parallel_reference(
+    locals_: list,
+    temperatures: torch.Tensor,
+    top_ks: torch.Tensor,
+    top_ps: torch.Tensor,
+    min_ps: torch.Tensor,
+    seed: Optional[torch.Tensor],
+    offset: Optional[torch.Tensor],
+    multiplier: float,
+    soft_cap: float,
+    greedy: bool,
+) -> torch.Tensor:
+    """The built-in head on the rank-order concatenation of the shards: ``torch.argmax`` of the
+    soft-capped FP32 row (greedy) or ``fused_sample_from_logits`` on it."""
+    capped = reference_capped(torch.cat(list(locals_), dim=1), multiplier, soft_cap)
+    if greedy:
+        return torch.argmax(capped, dim=-1)
+    return torch.ops.trtllm.fused_sample_from_logits(
+        capped, temperatures, top_ks, top_ps, min_ps, seed, offset
+    )
+
+
+class VocabParallelEmulation:
+    """``ranks`` virtual ranks on one GPU: every value is stored into each rank's copy (instead
+    of one multicast store), every rank's staging rows are plain tensors. The kernels are the
+    served ones; a call runs every rank's ``vp_fold``, then every ``vp_draw``, then every
+    ``vp_decide`` (so no poll waits)."""
+
+    def __init__(self, ranks: int, shard: int, max_rows: int, device):
+        load_library()
+        self.layout = VocabParallelLayout(ranks, shard, max_rows)
+        lay = self.layout
+        self.device = torch.device(device)
+        self.copies = [
+            torch.full((lay.nbytes // 4,), -1, dtype=torch.int32, device=self.device)
+            for _ in range(ranks)
+        ]
+        self.staging = [
+            torch.empty((max_rows, shard), dtype=torch.bfloat16, device=self.device)
+            for _ in range(ranks)
+        ]
+        publish = [c.data_ptr() for c in self.copies]
+        staging_of = [s.data_ptr() for s in self.staging]
+        self.status = torch.zeros(1, dtype=torch.int32, device=self.device)
+        self.views = [
+            _VocabParallelRank(
+                lay,
+                r,
+                self.copies[r].data_ptr(),
+                publish,
+                staging_of[r],
+                staging_of,
+                torch.empty(max_rows * VP_INFO_WORDS, dtype=torch.int32, device=self.device),
+                self.status,
+            )
+            for r in range(ranks)
+        ]
+
+    def sample(
+        self,
+        locals_,
+        temperatures,
+        top_ks,
+        top_ps,
+        min_ps,
+        seed,
+        offset,
+        multiplier,
+        soft_cap,
+        pdl=False,
+    ):
+        """Every rank's tokens (list of int32 ``[rows]``) and pending rows."""
+        rows = locals_[0].shape[0]
+        pendings = [
+            torch.zeros((rows, self.layout.vocab), dtype=torch.bfloat16, device=self.device)
+            for _ in self.views
+        ]
+        for view, local in zip(self.views, locals_):
+            view.fold(local, temperatures, top_ks, top_ps, min_ps, multiplier, soft_cap, False, pdl)
+        tokens = [
+            view.draw(
+                pending,
+                local,
+                temperatures,
+                top_ks,
+                top_ps,
+                min_ps,
+                seed,
+                offset,
+                multiplier,
+                soft_cap,
+                pdl,
+            )
+            for view, pending, local in zip(self.views, pendings, locals_)
+        ]
+        for view, t in zip(self.views, tokens):
+            view.decide(t, pdl)
+        return tokens, pendings
+
+    def argmax(self, locals_, multiplier, soft_cap, pdl=False):
+        for view, local in zip(self.views, locals_):
+            view.fold(
+                local,
+                self.status,
+                self.status,
+                self.status,
+                self.status,
+                multiplier,
+                soft_cap,
+                True,
+                pdl,
+            )
+        return [view.argmax(locals_[0].shape[0], pdl) for view in self.views]
+
+    def all_sentinel(self) -> bool:
+        return all(bool((c == -1).all()) for c in self.copies)
+
+
+class VocabParallelWorkspace:
+    """The vocabulary-parallel head's fabric buffers for one tensor-parallel group.
+
+    ``exchange``: the multicast exchange buffer (``VocabParallelLayout``), each rank's copy filled
+    with the sentinel; values are stored through the multicast address, read and reset in the
+    own copy. ``staging``: ``[max_rows, shard]`` BF16 per rank (unicast; peers read it), written
+    by ``vp_fold`` for rows with a filter only.
+
+    Protocol: (1) every exchange slot has one writer per call (``vp_fold`` per (rank, row,
+    virtual thread) or (rank, row, CTA); ``vp_draw`` per (rank, row, claimant)); every rank
+    reads each slot of its own copy once per call and resets it. (2) At the start of every call
+    every word of every copy is the sentinel: after construction (fill, synchronize, agreement)
+    and restored by the call's readers. (3) Reuse across calls: a peer stores into this copy (and
+    into its staging rows) again only in a later call, which it reaches after a collective of
+    the group issued between the two calls (the decode step's all-reduces and hops, the
+    drafter's hops). This rank publishes into that collective only after this call's last
+    kernel has completed: every kernel between them is a plain launch or waits for its
+    predecessor (``griddepcontrol.wait``) before it completes, and the collective publishes
+    after its own wait -- so an all-PDL chain (the Lamport hops) orders it as an NCCL call
+    does. Two calls with no collective between them are not ordered: a peer's next fold can
+    overwrite a word this rank has not read. (4) Every rank makes the same calls: the path
+    depends only on rank-uniform facts.
+
+    Construction is collective over the group and succeeds or raises on every rank.
+    """
+
+    def __init__(self, mapping: Mapping, shard: int, max_rows: int):
+        from tensorrt_llm._torch.distributed.ops import (
+            _get_mnnvl_workspace_comm,
+            _make_mnnvl_mcast_buffer,
+            _mnnvl_device_index,
+            _mnnvl_workspace_all_succeeded,
+        )
+
+        self.ranks = mapping.tp_size
+        self.rank = mapping.tp_rank
+        self.device = torch.device("cuda", _mnnvl_device_index(mapping))
+        self._comm = _get_mnnvl_workspace_comm(mapping)
+        self.library = ""
+        error: Optional[BaseException] = None
+        try:
+            self.layout = VocabParallelLayout(self.ranks, shard, max_rows)
+            if torch.cuda.get_device_capability(self.device) not in ((10, 3), (10, 7)):
+                raise RuntimeError("the vocabulary-parallel head runs on SM103 or SM107 only")
+            self.library = load_library()
+        except Exception as exc:  # noqa: BLE001 - every failure must reach the agreement
+            error = exc
+        if not _mnnvl_workspace_all_succeeded(self._comm, error is None):
+            raise RuntimeError(
+                "[musespark_open] vocab_parallel_head: the head sampler library failed to load on at "
+                "least one rank"
+                f"{f' (this rank: {error!r})' if error is not None else ''}"
+            ) from error
+        lay = self.layout
+        try:
+            multi_node = mapping.is_multi_node()
+            self._exchange = _make_mnnvl_mcast_buffer(self._comm, lay.nbytes, mapping, multi_node)
+            self._staging = _make_mnnvl_mcast_buffer(
+                self._comm, lay.staging_bytes, mapping, multi_node
+            )
+            self.copy = self._exchange.get_uc_buffer(self.rank, (lay.nbytes // 4,), torch.int32, 0)
+            self._mc = self._exchange.get_mc_buffer((lay.nbytes // 4,), torch.int32, 0)
+            self.staging_of = [
+                self._staging.get_uc_buffer(peer, (lay.staging_bytes // 2,), torch.bfloat16, 0)
+                for peer in range(self.ranks)
+            ]
+            with torch.cuda.device(self.device):
+                self.copy.fill_(-1)
+                self.info = torch.empty(
+                    lay.max_rows * VP_INFO_WORDS, dtype=torch.int32, device=self.device
+                )
+                self.status = torch.zeros(1, dtype=torch.int32, device=self.device)
+                torch.cuda.synchronize(self.device)
+        except Exception as exc:  # noqa: BLE001 - every failure must reach the agreement
+            error = exc
+        if not _mnnvl_workspace_all_succeeded(self._comm, error is None):
+            raise RuntimeError(
+                "[musespark_open] vocab_parallel_head: workspace failed on at least one rank"
+                f"{f' (this rank: {error!r})' if error is not None else ''}"
+            ) from error
+        self.view = _VocabParallelRank(
+            lay,
+            self.rank,
+            self.copy.data_ptr(),
+            [self._mc.data_ptr()],
+            self.staging_of[self.rank].data_ptr(),
+            [t.data_ptr() for t in self.staging_of],
+            self.info,
+            self.status,
+        )
+
+    @property
+    def max_rows(self) -> int:
+        return self.layout.max_rows
+
+    @property
+    def vocab(self) -> int:
+        return self.layout.vocab
+
+    def nbytes(self) -> int:
+        """Device memory this rank holds for the head (exchange copy + staging rows)."""
+        return self.layout.nbytes + self.layout.staging_bytes
+
+    def errors(self) -> int:
+        """The poll status word (synchronizes)."""
+        return int(self.status.item())
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py
index d9bb985aaaf1..9181f8d25191 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py
@@ -11,7 +11,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenModel
 from tensorrt_llm._torch.models.modeling_musespark_open_omni import (
@@ -38,7 +38,10 @@ def load(
 
 
 def install_loaded_features(
-    text_model: MuseSparkOpenModel, mapping, plan: "_dispatch.Plan"
+    text_model: MuseSparkOpenModel,
+    mapping,
+    plan: "_dispatch.Plan",
+    compute_capability: Optional[Tuple[int, int]] = None,
 ) -> None:
     """Install the plan's expert paths on a text model whose weights are final.
 
@@ -60,6 +63,10 @@ def install_loaded_features(
 
     Then ``hot_replicas`` (``hot_replicas.install``: hooks only, at tp 4; the copies are
     gathered on the first forward of a deployment that serves 33..512-row decode steps),
+    then, on Rubin (sm107) at tp 4 (``prefill_cutedsl_experts.serves`` of
+    ``compute_capability``, the device's; None installs it nowhere),
+    ``prefill_cutedsl_experts`` (``prefill_cutedsl_experts.install``: the
+    4096..8192-row expert calls of the NVFP4 layers on the Rubin CuTe DSL GEMMs),
     then ``early_expert_ids`` (``early_ids.install``, on the layers just installed),
     then ``pool_into_route`` (``DirectPool.install_sites``, whose routed kernel
     variants carry the placement and the early expert ids), each collective over the
@@ -76,6 +83,13 @@ def install_loaded_features(
 
         # After the owner-window install, whose expert and routing hooks it wraps.
         _hot_replicas.install(text_model, mapping)
+    if compute_capability is not None:
+        from .kernels.prefill_cutedsl_experts import install as _prefill_cutedsl
+
+        if _prefill_cutedsl.serves(compute_capability, mapping.tp_size):
+            # After the expert paths above, whose expert forwards it wraps for the large
+            # prefill calls only.
+            _prefill_cutedsl.install(text_model, mapping)
     if plan.enabled("early_expert_ids"):
         _install_early_expert_ids(text_model, mapping)
     if plan.enabled("pool_into_route"):
@@ -180,3 +194,9 @@ def install_embedding_features(
         from .kernels.replicated_embedding import embedding as _replicated_embedding
 
         _replicated_embedding.install(text_model, mapping)
+    if plan.enabled("model_boundary"):
+        # After replicated_embedding, whose table it gathers from. Collective (agreed refusal).
+        from .kernels.boundary_early_release import triggers as _pdl_triggers
+        from .kernels.model_boundary import boundary as _model_boundary
+
+        _model_boundary.install(text_model, _pdl_triggers.capture_layers(text_model), mapping)
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index e641111adb4e..b0e365826d1f 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -14784,8 +14784,13 @@ def __init__(
                     scaling_vector_size: int = 16,
                     activation_type: ActivationType = ActivationType.Swiglu,
                     situ_beta: Optional[float] = None,
-                    situ_linear_beta: Optional[float] = None):
+                    situ_linear_beta: Optional[float] = None,
+                    trtllm_gen_gated_layout: bool = False):
                 super().__init__()
+                # trtllm_gen_gated_layout: the weights are TRTLLM-gen's gated FC1
+                # tensors, alpha is [gate scales | output scales] (2 * L values)
+                # and the epilogue follows TRTLLM-gen's arithmetic.
+                self.trtllm_gen_gated_layout = trtllm_gen_gated_layout
                 self.num_experts = num_experts
                 self.top_k = top_k
                 self.num_local_experts = num_local_experts
@@ -14819,6 +14824,7 @@ def unique_id(self):
                     int(self.activation_type),
                     self.situ_beta,
                     self.situ_linear_beta,
+                    self.trtllm_gen_gated_layout,
                 )
 
             def get_valid_tactics(
@@ -14998,7 +15004,8 @@ def forward(self, inputs: List[torch.Tensor],
                 assert b_sf.size(0) == l
                 assert b_sf.size(1) == n
                 assert b_sf.size(2) == scale_k
-                assert alpha.size(0) == l
+                assert alpha.size(0) == l * (2 if self.trtllm_gen_gated_layout
+                                             else 1)
 
                 num_tiles = m // self.tile_size
                 assert tile_idx_to_group_idx.dtype == torch.int32
@@ -15122,7 +15129,8 @@ def forward(self, inputs: List[torch.Tensor],
                              cluster_shape_mn, raster_along_m,
                              locality_domain_half_gemm, a_path,
                              int(self.activation_type), self.situ_beta,
-                             self.situ_linear_beta, max_active_clusters)
+                             self.situ_linear_beta, max_active_clusters,
+                             self.trtllm_gen_gated_layout)
                 if cache_key not in self.__class__.kernel_cache:
                     gemm = self.__class__.kernel_class(
                         sf_vec_size=self.scaling_vector_size,
@@ -15137,6 +15145,7 @@ def forward(self, inputs: List[torch.Tensor],
                         activation_type=self.activation_type,
                         situ_beta=self.situ_beta,
                         situ_linear_beta=self.situ_linear_beta,
+                        trtllm_gen_gated_layout=self.trtllm_gen_gated_layout,
                     )
                     compiled_gemm = cute.compile(
                         gemm.wrapper,
@@ -16161,8 +16170,20 @@ def __init__(self,
                          local_expert_offset: int,
                          tile_size: int,
                          output_dtype: torch.dtype,
-                         scaling_vector_size: int = 16):
+                         scaling_vector_size: int = 16,
+                         store_permuted: bool = False,
+                         trtllm_gen_n_shuffle: bool = False,
+                         mma_k64: bool = False):
                 super().__init__()
+                # store_permuted: C gets one row per permuted row (alpha * acc,
+                # no token scale, no scatter); see the kernel class.
+                # trtllm_gen_n_shuffle: the weights are TRTLLM-gen's FC2 tensors
+                # (shuffle_matrix_a row order); C columns come out in logical order.
+                # mma_k64: the 64-deep SM100-family MMA (fp32 accumulators
+                # equal to SM100-family NVFP4 GEMMs'); see the kernel class.
+                self.store_permuted = store_permuted
+                self.trtllm_gen_n_shuffle = trtllm_gen_n_shuffle
+                self.mma_k64 = mma_k64
                 self.num_experts = num_experts
                 self.top_k = top_k
                 self.num_local_experts = num_local_experts
@@ -16192,6 +16213,9 @@ def unique_id(self):
                     self.tile_size,
                     self.output_dtype,
                     self.scaling_vector_size,
+                    self.store_permuted,
+                    self.trtllm_gen_n_shuffle,
+                    self.mma_k64,
                 )
 
             @staticmethod
@@ -16291,13 +16315,17 @@ def get_tuning_config(self) -> TuningConfig:
                                                      self.num_local_experts,
                                                      self.local_expert_offset,
                                                      self.tile_size)
+                    # Stored permuted rows: C has as many rows as A.
+                    c_rows = ((lambda shapes: shapes[0][0])
+                              if self.store_permuted else
+                              helper.infer_shape_num_tokens)
                     self.__class__.tuning_config_cache[key] = TuningConfig(
                         dynamic_tensor_specs=(DynamicTensorSpec(
                             0, 0, helper.gen_tuning_buckets,
                             helper.map_to_tuning_buckets), ),
                         constraint_specs=(
                             ConstraintSpec(2, 0, fp4_scale_infer_shape),
-                            ConstraintSpec(5, 0, helper.infer_shape_num_tokens),
+                            ConstraintSpec(5, 0, c_rows),
                             ConstraintSpec(6, 0,
                                            helper.infer_shape_max_num_tiles),
                             ConstraintSpec(7, 0,
@@ -16342,6 +16370,9 @@ def forward(self, inputs: List[torch.Tensor],
                 assert c.dim() == 2
                 num_tokens = c.size(0)
                 assert c.size(1) == n or c.size(1) == n * 2
+                if self.store_permuted:
+                    assert num_tokens == m, (
+                        "store_permuted needs one C row per permuted row")
 
                 num_tiles = m // self.tile_size
                 assert tile_idx_to_group_idx.dtype == torch.int32
@@ -16354,7 +16385,8 @@ def forward(self, inputs: List[torch.Tensor],
                 assert num_non_exiting_tiles.numel() == 1
                 assert token_final_scales.dtype == torch.float32
                 assert token_final_scales.dim() == 2
-                assert token_final_scales.size() == (num_tokens, self.top_k)
+                if not self.store_permuted:
+                    assert token_final_scales.size() == (num_tokens, self.top_k)
 
                 locality_domain_id = get_current_locality_domain()
                 locality_domain_half_gemm = locality_domain_id is not None
@@ -16434,7 +16466,9 @@ def forward(self, inputs: List[torch.Tensor],
                 cache_key = (self.scaling_vector_size, self.tile_size,
                              self.top_k, mma_tiler, mma_inst_shape,
                              cluster_shape_mn, raster_along_m,
-                             locality_domain_half_gemm, max_active_clusters)
+                             locality_domain_half_gemm, max_active_clusters,
+                             self.store_permuted, self.trtllm_gen_n_shuffle,
+                             self.mma_k64)
                 if cache_key not in self.__class__.kernel_cache:
                     gemm = self.__class__.kernel_class(
                         sf_vec_size=self.scaling_vector_size,
@@ -16443,6 +16477,9 @@ def forward(self, inputs: List[torch.Tensor],
                         cluster_shape_mn=cluster_shape_mn,
                         raster_along_m=raster_along_m,
                         topK=self.top_k,
+                        store_permuted=self.store_permuted,
+                        trtllm_gen_n_shuffle=self.trtllm_gen_n_shuffle,
+                        mma_k64=self.mma_k64,
                     )
                     compiled_gemm = cute.compile(
                         gemm.wrapper,
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py
index 89351bf6c0d1..f38a1a4497b5 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py
@@ -203,8 +203,17 @@ def __init__(
         activation_type: ActivationType = ActivationType.Swiglu,
         situ_beta: Optional[float] = None,
         situ_linear_beta: Optional[float] = None,
+        trtllm_gen_gated_layout: bool = False,
     ):
         self.a_path = a_path
+        # trtllm_gen_gated_layout: B (and SFB) rows are in TRTLLM-gen's gated
+        # FC1 layout (see _trtllm_gen_gated_split) instead of the granularity-64
+        # [up, gate] interleave, alpha holds two values per group (the gate
+        # scale at [g], the output scale at [l + g], see
+        # _apply_trtllm_gen_swiglu_epilogue), and the SwiGLU + NVFP4
+        # quantization use TRTLLM-gen's arithmetic, so the output codes and
+        # block scales equal TRTLLM-gen's gated-activation GEMM bit for bit.
+        self.trtllm_gen_gated_layout = trtllm_gen_gated_layout
         # locality domain half-GEMM: two partitions write their N-half into a shared
         # full-width C/SFC buffer at a column offset (see wrapper/__call__).
         self.locality_domain_half_gemm = locality_domain_half_gemm
@@ -231,6 +240,12 @@ def __init__(
         self.is_gated = is_gated_activation(self.activation_type)
         if locality_domain_half_gemm and not self.is_gated:
             raise ValueError("Rubin locality domain half-GEMM currently supports SwiGLU only")
+        if trtllm_gen_gated_layout and (
+            self.activation_type != ActivationType.Swiglu or locality_domain_half_gemm
+        ):
+            raise ValueError(
+                "trtllm_gen_gated_layout supports SwiGLU without the locality domain half-GEMM only"
+            )
         self.acc_dtype = cutlass.Float32
         self.mma_inst_shape = mma_inst_shape
         self.mma_tiler = mma_tiler
@@ -2934,7 +2949,12 @@ def kernel(
                 #
 
                 expert_idx = mma_tile_coord_mnl[2]
-                alpha_val = alpha[expert_idx]
+                if cutlass.const_expr(self.trtllm_gen_gated_layout):
+                    # gate scale; output scale (applied by the quantization)
+                    alpha_val = alpha[(expert_idx, 0)]
+                    scale_c_val = alpha[(expert_idx, 1)]
+                else:
+                    alpha_val = alpha[expert_idx]
 
                 #
                 # Slice to per mma tile index
@@ -3046,7 +3066,12 @@ def kernel(
 
                         acc_vec_up = tTR_rAcc_up.load()
                         tCompute = cute.make_rmem_tensor(acc_vec_up.shape, self.acc_dtype)
-                        if cutlass.const_expr(self.activation_type == ActivationType.Swiglu):
+                        if cutlass.const_expr(self.trtllm_gen_gated_layout):
+                            tUp, tGate = self._trtllm_gen_gated_split(tTR_rAcc_up, tTR_rAcc_gate)
+                            self._apply_trtllm_gen_swiglu_epilogue(
+                                tUp.load(), tGate.load(), alpha_val, tCompute
+                            )
+                        elif cutlass.const_expr(self.activation_type == ActivationType.Swiglu):
                             acc_vec_gate = tTR_rAcc_gate.load()
                             self._apply_swiglu_epilogue(
                                 acc_vec_up, acc_vec_gate, alpha_val, tCompute
@@ -3094,7 +3119,31 @@ def kernel(
                                 abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype
                             )
 
-                            if cutlass.const_expr(self.vectorized_f32):
+                            if cutlass.const_expr(self.trtllm_gen_gated_layout):
+                                # TRTLLM-gen: SF = e4m3((amax * scale_c) * (1/6))
+                                for vi in cutlass.range_constexpr(abs_acc_frg.shape[1]):
+                                    tCrSFC_pvscale[vi] = abs_acc_frg[None, vi].reduce(
+                                        cute.ReductionOp.MAX,
+                                        cutlass.Float32(0.0),
+                                        0,  # Use 0.0 as init for abs values
+                                    )
+                                for vi in cutlass.range_constexpr(0, abs_acc_frg.shape[1], 2):
+                                    tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = (
+                                        cute.arch.mul_packed_f32x2(
+                                            (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]),
+                                            (scale_c_val, scale_c_val),
+                                        )
+                                    )
+                                    tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = (
+                                        cute.arch.mul_packed_f32x2(
+                                            (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]),
+                                            (
+                                                self.get_dtype_rcp_limits(self.c_dtype),
+                                                self.get_dtype_rcp_limits(self.c_dtype),
+                                            ),
+                                        )
+                                    )
+                            elif cutlass.const_expr(self.vectorized_f32):
                                 for vi in cutlass.range_constexpr(abs_acc_frg.shape[1]):
                                     tCrSFC_pvscale[vi] = abs_acc_frg[None, vi].reduce(
                                         cute.ReductionOp.MAX,
@@ -3146,7 +3195,31 @@ def kernel(
                             # TODO: need to add f8x2 -> f32x2 conversion
                             tCrSFC_qpvscale_up = tCrSFC.load().to(cutlass.Float32)
                             fp32_max = cutlass.Float32(3.40282346638528859812e38)
-                            if cutlass.const_expr(self.vectorized_f32):
+                            if cutlass.const_expr(self.trtllm_gen_gated_layout):
+                                # TRTLLM-gen: q = (x * rcp(SF)) * scale_c
+                                for vi in cutlass.range_constexpr(0, cute.size(tCrSFC), 2):
+                                    acc_scale_min0 = fmin(
+                                        cute.arch.rcp_approx(tCrSFC_qpvscale_up[vi]),
+                                        fp32_max,
+                                        nan=True,
+                                    )
+                                    acc_scale_min1 = fmin(
+                                        cute.arch.rcp_approx(tCrSFC_qpvscale_up[vi + 1]),
+                                        fp32_max,
+                                        nan=True,
+                                    )
+                                    vec0 = tTR_rAcc_frg[None, vi]
+                                    vec1 = tTR_rAcc_frg[None, vi + 1]
+                                    for ei in cutlass.range_constexpr(self.sf_vec_size):
+                                        vec0[ei], vec1[ei] = cute.arch.mul_packed_f32x2(
+                                            (vec0[ei], vec1[ei]),
+                                            (acc_scale_min0, acc_scale_min1),
+                                        )
+                                        vec0[ei], vec1[ei] = cute.arch.mul_packed_f32x2(
+                                            (vec0[ei], vec1[ei]),
+                                            (scale_c_val, scale_c_val),
+                                        )
+                            elif cutlass.const_expr(self.vectorized_f32):
                                 for vi in cutlass.range_constexpr(0, cute.size(tCrSFC), 2):
                                     acc_scale = cute.arch.mul_packed_f32x2(
                                         (
@@ -3253,6 +3326,76 @@ def kernel(
 
         cute.arch.griddepcontrol_launch_dependents()
 
+    @cute.jit
+    def _trtllm_gen_gated_split(self, tAccLo: cute.Tensor, tAccHi: cute.Tensor):
+        """Gather up / gate for 64 output columns from TRTLLM-gen gated FC1 accumulators.
+
+        ``tAccLo`` / ``tAccHi`` are one thread's two consecutive 64-column
+        accumulator subtiles (128 B rows). In TRTLLM-gen's gated layout every
+        32-row block of B covers 16 intermediate columns: rows 0-7 are up of the
+        even columns, 8-15 gate of the even columns, 16-23 up of the odd
+        columns and 24-31 gate of the odd columns. Output column ``j`` of the
+        64-column output subtile therefore takes, in subtile ``j // 32`` with
+        ``r = j % 32 = 16 g + 2 kk + o``, up from column ``32 g + 16 o + kk``
+        and gate from column ``32 g + 16 o + 8 + kk``. Only register indices
+        move (compile-time constants).
+        """
+        tUp = cute.make_rmem_tensor(tAccLo.shape, self.acc_dtype)
+        tGate = cute.make_rmem_tensor(tAccLo.shape, self.acc_dtype)
+        for j in cutlass.range_constexpr(64):
+            r = j % 32
+            g = r // 16
+            o = r % 2
+            kk = (r % 16) // 2
+            src_up = 32 * g + 16 * o + kk
+            src = tAccLo if j < 32 else tAccHi
+            tUp[j] = src[src_up]
+            tGate[j] = src[src_up + 8]
+        return tUp, tGate
+
+    @cute.jit
+    def _apply_trtllm_gen_swiglu_epilogue(
+        self,
+        acc_vec_up: cute.Tensor,
+        acc_vec_gate: cute.Tensor,
+        scale_gate,
+        tCompute: cute.Tensor,
+    ):
+        """SwiGLU in TRTLLM-gen's gated-activation GEMM arithmetic (fp32, in this order)::
+
+            k = log2(e) * scale_gate
+            s = rcp(1 + exp2(-(gate[i] * k)))
+            tCompute[i] = (up[i] * scale_gate) * (gate[i] * s)
+
+        i.e. ``scale_gate * up * silu(scale_gate * gate)``, with the gate scale
+        folded into the exponent constant and the up half. The output scale is
+        applied by the quantization (``SF = e4m3((amax * scale_c) / 6)``,
+        ``q = (x * rcp(SF)) * scale_c``).
+        """
+        k = cutlass.Float32(1.4426950408889634) * cutlass.Float32(scale_gate)
+        for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2):
+            arg = cute.arch.mul_packed_f32x2(
+                (acc_vec_gate[i], acc_vec_gate[i + 1]),
+                (k, k),
+            )
+            arg = cute.arch.mul_packed_f32x2(arg, (-1.0, -1.0))
+            den = cute.arch.add_packed_f32x2(
+                (
+                    cute.math.exp2(arg[0], fastmath=True),
+                    cute.math.exp2(arg[1], fastmath=True),
+                ),
+                (1.0, 1.0),
+            )
+            act = cute.arch.mul_packed_f32x2(
+                (acc_vec_gate[i], acc_vec_gate[i + 1]),
+                (cute.arch.rcp_approx(den[0]), cute.arch.rcp_approx(den[1])),
+            )
+            lin = cute.arch.mul_packed_f32x2(
+                (acc_vec_up[i], acc_vec_up[i + 1]),
+                (cutlass.Float32(scale_gate), cutlass.Float32(scale_gate)),
+            )
+            tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(lin, act)
+
     @cute.jit
     def _apply_swiglu_epilogue(
         self,
@@ -4145,7 +4288,10 @@ def wrapper(
                 order=(2, 1, 4, 0, 3, 5),
             ),
         )
-        alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,)))
+        if cutlass.const_expr(self.trtllm_gen_gated_layout):
+            alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l, 2), stride=(1, l)))
+        else:
+            alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,)))
 
         tile_idx_to_group_idx = cute.make_tensor(
             tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,))
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py
index 7cf15e17f995..73412e40f528 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py
@@ -117,6 +117,24 @@ def blk_reduce_fp32(dst_gemm, src_smem, size, loc=None, ip=None):
     )
 
 
+@dsl_user_op
+def blk_store(dst_gemm, src_smem, size, loc=None, ip=None):
+    """Bulk copy of ``size`` bytes from shared to global memory (cp.async.bulk)."""
+    llvm.inline_asm(
+        None,
+        [
+            dst_gemm.iterator.llvm_ptr,
+            src_smem.iterator.llvm_ptr,
+            size.ir_value(),
+        ],
+        "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;",
+        "l,l,r",
+        has_side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
 @dsl_user_op
 def blk_reduce_fp16(dst_gemm, src_smem, size, loc=None, ip=None):
     """Block reduce for FP16 using cp.reduce.async.bulk."""
@@ -171,9 +189,30 @@ def __init__(
         cluster_shape_mn: Tuple[int, int],
         raster_along_m: bool = False,
         topK: int = 1,
+        store_permuted: bool = False,
+        trtllm_gen_n_shuffle: bool = False,
+        mma_k64: bool = False,
     ):
+        # store_permuted: the epilogue stores each valid permuted row
+        # (alpha * acc, no token scale) to row ``permuted_row`` of C, which then
+        # has one row per permuted row, instead of scatter-adding the
+        # token-scaled row into its token's row.
+        # trtllm_gen_n_shuffle: B (and SFB) rows are in TRTLLM-gen's
+        # shuffle_matrix_a order (epilogue_tile_m 128: within every 32 rows,
+        # stored row p is logical row 4 * (p % 8) + p // 8); the epilogue writes
+        # C columns back in logical order.
+        # mma_k64: issue the 64-deep tcgen05 kind::mxf4nvf4 MMA (the SM100 family
+        # instruction) instead of the 128-deep SM107 one. The fp32 accumulators
+        # then equal those of SM100-family NVFP4 GEMMs (e.g. TRTLLM-gen's) bit
+        # for bit; the K=128 instruction rounds its partial sums differently.
+        # mma_tiler[2] stays 256 (four MMAs per k-tile).
+        self.store_permuted = store_permuted
+        self.trtllm_gen_n_shuffle = trtllm_gen_n_shuffle
+        self.mma_k64 = mma_k64
         self.sf_vec_size = sf_vec_size
         self.acc_dtype = cutlass.Float32
+        if mma_k64:
+            mma_inst_shape = (mma_inst_shape[0], mma_inst_shape[1], 64)
         self.mma_inst_shape = mma_inst_shape
         self.mma_tiler = mma_tiler
         self.cluster_shape_mn = cluster_shape_mn
@@ -231,6 +270,49 @@ def __init__(
 
         # B-reuse pattern control
         self.enable_breuse = True if mma_tiler[0] // mma_inst_shape[0] == 2 else False
+        if mma_k64 and self.enable_breuse:
+            raise ValueError("mma_k64 does not support the B-reuse pattern")
+
+    def _make_tiled_mma(self, cta_group, mma_inst_shape, permutation_mnk=None):
+        """The block-scaled tiled MMA (SM107 K=128, or the SM100-family K=64 one with mma_k64)."""
+        if cutlass.const_expr(self.mma_k64):
+            return sm100_utils.make_blockscaled_trivial_tiled_mma(
+                self.a_dtype,
+                self.b_dtype,
+                self.a_major_mode,
+                self.b_major_mode,
+                self.sf_dtype,
+                self.sf_vec_size,
+                cta_group,
+                mma_inst_shape[:2],
+            )
+        if cutlass.const_expr(permutation_mnk is None):
+            return sm107_utils.make_blockscaled_trivial_tiled_mma(
+                self.a_dtype,
+                self.b_dtype,
+                self.a_major_mode,
+                self.b_major_mode,
+                self.sf_dtype,
+                self.sf_vec_size,
+                cta_group,
+                mma_inst_shape,
+                a_collector_op=CollectorOp.DISCARD,
+                b_collector_op=CollectorOp.DISCARD,
+            )
+        return sm107_utils.make_blockscaled_trivial_tiled_mma(
+            self.a_dtype,
+            self.b_dtype,
+            self.a_major_mode,
+            self.b_major_mode,
+            self.sf_dtype,
+            self.sf_vec_size,
+            cta_group,
+            mma_inst_shape,
+            a_collector_op=CollectorOp.DISCARD,
+            b_collector_op=CollectorOp.DISCARD,
+            atom_layout_mnk=(1, 1, 1),
+            permutation_mnk=permutation_mnk,
+        )
 
     def _get_mma_permutation_mnk(self):
         if cutlass.const_expr(self.use_2cta_instrs and self.enable_breuse):
@@ -312,32 +394,12 @@ def _setup_attributes(self):
             self.mma_inst_shape[2],
         )
 
-        tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma(
-            self.a_dtype,
-            self.b_dtype,
-            self.a_major_mode,
-            self.b_major_mode,
-            self.sf_dtype,
-            self.sf_vec_size,
-            self.cta_group,
-            self.mma_inst_shape,
-            a_collector_op=CollectorOp.DISCARD,
-            b_collector_op=CollectorOp.DISCARD,
-            atom_layout_mnk=(1, 1, 1),
-            permutation_mnk=self._get_mma_permutation_mnk(),
+        tiled_mma = self._make_tiled_mma(
+            self.cta_group, self.mma_inst_shape, self._get_mma_permutation_mnk()
         )
 
-        tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma(
-            self.a_dtype,
-            self.b_dtype,
-            self.a_major_mode,
-            self.b_major_mode,
-            self.sf_dtype,
-            self.sf_vec_size,
-            cute.nvgpu.tcgen05.CtaGroup.ONE,
-            self.mma_inst_shape_sfb,
-            a_collector_op=CollectorOp.DISCARD,
-            b_collector_op=CollectorOp.DISCARD,
+        tiled_mma_sfb = self._make_tiled_mma(
+            cute.nvgpu.tcgen05.CtaGroup.ONE, self.mma_inst_shape_sfb
         )
 
         # Compute mma/cluster/tile shapes
@@ -633,35 +695,14 @@ def __call__(
         atom_layout_mnk = (1, 1, 1)
         permutation_mnk = self._get_mma_permutation_mnk()
 
-        tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma(
-            self.a_dtype,
-            self.b_dtype,
-            self.a_major_mode,
-            self.b_major_mode,
-            self.sf_dtype,
-            self.sf_vec_size,
-            self.cta_group,
-            self.mma_inst_shape,
-            a_collector_op=CollectorOp.DISCARD,
-            b_collector_op=CollectorOp.DISCARD,
-            atom_layout_mnk=atom_layout_mnk,
-            permutation_mnk=permutation_mnk,
-        )
+        assert atom_layout_mnk == (1, 1, 1)
+        tiled_mma = self._make_tiled_mma(self.cta_group, self.mma_inst_shape, permutation_mnk)
 
         tiled_mma.set(tcgen05.Field.NEGATE_A, False)
         tiled_mma.set(tcgen05.Field.NEGATE_B, False)
 
-        tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma(
-            self.a_dtype,
-            self.b_dtype,
-            self.a_major_mode,
-            self.b_major_mode,
-            self.sf_dtype,
-            self.sf_vec_size,
-            cute.nvgpu.tcgen05.CtaGroup.ONE,
-            self.mma_inst_shape_sfb,
-            a_collector_op=CollectorOp.DISCARD,
-            b_collector_op=CollectorOp.DISCARD,
+        tiled_mma_sfb = self._make_tiled_mma(
+            cute.nvgpu.tcgen05.CtaGroup.ONE, self.mma_inst_shape_sfb
         )
         tiled_mma_sfb.set(tcgen05.Field.NEGATE_A, False)
         tiled_mma_sfb.set(tcgen05.Field.NEGATE_B, False)
@@ -1587,11 +1628,12 @@ def kernel(
                     # Compute token info and scaled alpha for this row
                     token_idx = cutlass.Int32(0)
                     alpha_val_iter = alpha_val
-                    if is_valid_row:
-                        token_idx = expanded_idx // self.topK
-                        topk_idx = expanded_idx % self.topK
-                        token_scale = token_final_scales[(token_idx, topk_idx)]
-                        alpha_val_iter = alpha_val * token_scale
+                    if cutlass.const_expr(not self.store_permuted):
+                        if is_valid_row:
+                            token_idx = expanded_idx // self.topK
+                            topk_idx = expanded_idx % self.topK
+                            token_scale = token_final_scales[(token_idx, topk_idx)]
+                            alpha_val_iter = alpha_val * token_scale
 
                     for n_iter_idx in cutlass.range(n_subtile_cnt):
                         # Load accumulator from tensor memory buffer to register
@@ -1600,7 +1642,12 @@ def kernel(
                         cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
 
                         # For block reduce: retile then load, compute, store
-                        acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
+                        if cutlass.const_expr(self.trtllm_gen_n_shuffle):
+                            acc_vec = tiled_copy_r2s.retile(
+                                self._trtllm_gen_n_unshuffle(tTR_rAcc)
+                            ).load()
+                        else:
+                            acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
                         acc_vec_final = (alpha_val_iter * acc_vec).to(self.c_dtype)
                         tRS_rC.store(acc_vec_final)
                         if is_valid_row:
@@ -1628,32 +1675,38 @@ def kernel(
                     is_valid_row = permuted_row < tile_info[4]
 
                     if is_valid_row:
-                        expanded_idx = permuted_idx_to_expanded_idx[permuted_row]
-                        token_idx = expanded_idx // self.topK
                         coord_n = mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk[1]
-                        scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), c)
-
                         # sC row index: m_iter_idx * 128 + epi_tidx
                         sC_row = m_iter_idx * 128 + epi_tidx
-
-                        if cutlass.const_expr(self.c_dtype == cutlass.BFloat16):
-                            blk_reduce_bf16(
-                                scatter_out_offset,
-                                sC[sC_row, None, 0],
-                                cutlass.Int32(self.copy_size),
-                            )
-                        elif cutlass.const_expr(self.c_dtype == cutlass.Float32):
-                            blk_reduce_fp32(
-                                scatter_out_offset,
-                                sC[sC_row, None, 0],
-                                cutlass.Int32(self.copy_size),
-                            )
-                        elif cutlass.const_expr(self.c_dtype == cutlass.Float16):
-                            blk_reduce_fp16(
-                                scatter_out_offset,
+                        if cutlass.const_expr(self.store_permuted):
+                            store_out_offset = cute.domain_offset((permuted_row, coord_n, 0), c)
+                            blk_store(
+                                store_out_offset,
                                 sC[sC_row, None, 0],
                                 cutlass.Int32(self.copy_size),
                             )
+                        else:
+                            expanded_idx = permuted_idx_to_expanded_idx[permuted_row]
+                            token_idx = expanded_idx // self.topK
+                            scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), c)
+                            if cutlass.const_expr(self.c_dtype == cutlass.BFloat16):
+                                blk_reduce_bf16(
+                                    scatter_out_offset,
+                                    sC[sC_row, None, 0],
+                                    cutlass.Int32(self.copy_size),
+                                )
+                            elif cutlass.const_expr(self.c_dtype == cutlass.Float32):
+                                blk_reduce_fp32(
+                                    scatter_out_offset,
+                                    sC[sC_row, None, 0],
+                                    cutlass.Int32(self.copy_size),
+                                )
+                            elif cutlass.const_expr(self.c_dtype == cutlass.Float16):
+                                blk_reduce_fp16(
+                                    scatter_out_offset,
+                                    sC[sC_row, None, 0],
+                                    cutlass.Int32(self.copy_size),
+                                )
 
                 cute.arch.cp_async_bulk_commit_group()
                 cute.arch.cp_async_bulk_wait_group(0, read=True)
@@ -1676,6 +1729,20 @@ def kernel(
             self.epilog_sync_barrier.arrive_and_wait()
             tmem.free(acc_tmem_ptr)
 
+    @cute.jit
+    def _trtllm_gen_n_unshuffle(self, tAcc: cute.Tensor) -> cute.Tensor:
+        """One thread's accumulator columns back in logical order (trtllm_gen_n_shuffle).
+
+        Column q of every 32-column group takes accumulator column
+        ``8 * (q % 4) + q // 4``. Only register indices move.
+        """
+        tOut = cute.make_rmem_tensor(tAcc.shape, self.acc_dtype)
+        for q in cutlass.range_constexpr(cute.size(tAcc)):
+            base = q - q % 32
+            r = q % 32
+            tOut[q] = tAcc[base + 8 * (r % 4) + r // 4]
+        return tOut
+
     @staticmethod
     def _compute_grid(
         gemm_shape: Tuple[int, int, int],
diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open.py b/tensorrt_llm/_torch/models/modeling_musespark_open.py
index babc872064d3..0c2dd676ddec 100644
--- a/tensorrt_llm/_torch/models/modeling_musespark_open.py
+++ b/tensorrt_llm/_torch/models/modeling_musespark_open.py
@@ -2524,6 +2524,11 @@ def forward(
 
 
 class MuseSparkOpenDecoderLayer(DecoderLayer):
+    #: The sequence-parallel arm's attention-input all-gather, ``(attn_input, sp_num_tokens) ->
+    #: rows``; None (the default) runs the NCCL all-gather. A model-specific path may set it
+    #: on a layer, together with a QKV projection that consumes what it returns.
+    _sp_attn_input_gather = None
+
     def __init__(
         self,
         model_config: ModelConfig[MuseSparkOpenConfig],
@@ -2627,7 +2632,11 @@ def forward(
                 and not torch.cuda.is_current_stream_capturing()
             )
             if sp_num_tokens is not None:
-                attn_input = allgather(attn_input, self.mapping, dim=0)[:sp_num_tokens]
+                gather = self._sp_attn_input_gather
+                if gather is None:
+                    attn_input = allgather(attn_input, self.mapping, dim=0)[:sp_num_tokens]
+                else:
+                    attn_input = gather(attn_input, sp_num_tokens)
                 kwargs = dict(kwargs, all_reduce_params=AllReduceParams(enable_allreduce=False))
             attn_out = self.self_attn(
                 position_ids=position_ids,
diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py b/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py
index 378a0f33bf14..f4201ade494c 100644
--- a/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py
+++ b/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py
@@ -165,12 +165,15 @@ class MuseSparkOpenOmniInputProcessor(
     semantics rather than imported from a shim.
     """
 
+    supports_tokenization_cache = True
+
     def __init__(
         self,
         model_path: str,
         config: PretrainedConfig,
         tokenizer,
         trust_remote_code: bool = True,
+        enable_tokenization_cache: bool = False,
         **kwargs,
     ):
         super().__init__(
@@ -183,6 +186,12 @@ def __init__(
         self._config = config
         self._tokenizer = tokenizer
         self._model_path = model_path
+        # A media-free prompt is ``tokenizer.encode(text, add_special_tokens=False)``
+        # (``_encode``), which is exactly what the prefix-tokenization cache
+        # computes, so its spliced ids are the ids this processor would produce.
+        # A processor built without a tokenizer (the engine's geometry copy)
+        # never tokenizes a served prompt and gets no cache.
+        self._init_tokenization_cache(enable_tokenization_cache, tokenizer)
         # `config.json` is not the whole preprocessing config on newer exports:
         # they drop `vision_config.navit_config` and declare the same knobs in
         # `processor_config.json` under HuggingFace spellings. Without this the
@@ -541,7 +550,12 @@ def call_with_text_prompt(
         mm_data = inputs.get("multi_modal_data") or {}
 
         if not mm_data:
-            return self._encode(text_prompt), None
+            ids = (
+                self._encode_with_tokenization_cache(text_prompt, sampling_params)
+                if text_prompt
+                else None
+            )
+            return (self._encode(text_prompt) if ids is None else ids), None
 
         # Every modality key must be consumed below. Anything else -- including
         # "audio" against a checkpoint whose config does not enable the
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 37fc9e0e3a6d..57560fe32204 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -12,6 +12,7 @@
 # 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.
+import contextlib
 import hashlib
 import math
 import os
@@ -24,6 +25,7 @@
     Callable,
     Dict,
     Iterable,
+    Iterator,
     List,
     NamedTuple,
     Optional,
@@ -1230,6 +1232,13 @@ class KVCacheManagerV2(BaseResourceManager):
     # Read by KvCacheCreator: a subclass that overrides the context
     # commit/history protocol opts out of generic reuse-match backoff.
     _supports_reuse_match_backoff = True
+    # Deferred close (opt-in, TRTLLM_KV_DEFER_CLOSE=1): see arm_deferred_close().
+    # Declared on the class so free_resources and shutdown also work on
+    # instances built without running __init__.
+    _defer_close_enabled: bool = False
+    _defer_close_armed: bool = False
+    _defer_close_scope: bool = False
+    _deferred_closes: Tuple[Tuple[int, _KVCache], ...] = ()
 
     def __init__(
         self,
@@ -1928,6 +1937,9 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
             index_mapper_capacity, max_beam_width, self.max_copy_beam_width
         )
         self._early_freed_index_requests: set[int] = set()
+        self._defer_close_enabled = (
+            os.environ.get("TRTLLM_KV_DEFER_CLOSE", "0") == "1" and not self.is_draft
+        )
         # Requests whose pages a connector is still reading from, so the
         # release half of `preempt_request` has to wait.
         self._pending_preemption: Dict[int, LlmRequest] = {}
@@ -5635,15 +5647,71 @@ def release_index_slot(self, request_id: int) -> None:
             # mirrored, and the target may release the same request twice.
             return
         if kv_cache is not None:
-            if self.sparse_metadata_batch is not None:
-                self.sparse_metadata_batch.record_read(self._stream.cuda_stream)
-                self.sparse_metadata_batch.remove(kv_cache)
-            for beam_idx in range(int(kv_cache.beam_width)):
-                for pool_idx in range(self.num_pools):
-                    kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, None)
+            self._detach_index_slot(kv_cache)
         self.index_mapper.remove_sequence(request_id)
         self._early_freed_index_requests.add(request_id)
 
+    def _detach_index_slot(self, kv_cache: _KVCache) -> None:
+        """Unbind a cache from its IndexMapper slot's page-index buffers (and
+        the sparse metadata batch) so the slot can be reused while the cache
+        stays open."""
+        if self.sparse_metadata_batch is not None:
+            self.sparse_metadata_batch.record_read(self._stream.cuda_stream)
+            self.sparse_metadata_batch.remove(kv_cache)
+        for beam_idx in range(int(kv_cache.beam_width)):
+            for pool_idx in range(self.num_pools):
+                kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, None)
+
+    def arm_deferred_close(self, armed: bool) -> bool:
+        """Arm or disarm deferred closing of finished requests' caches.
+
+        While armed, free_resources called inside deferring_closes() releases
+        the (non-dummy) request's IndexMapper slot at once (as
+        release_index_slot does) and queues the cache; its close(), which
+        commits the last block and releases the pages, runs in
+        flush_deferred_closes(). Every other release (preemption, admission
+        rollbacks, reservation drops) closes at once, since its caller counts
+        on the pages being free when it returns. The overlap executor arms this
+        only for its loop, releases terminated requests inside
+        deferring_closes(), and flushes right after launching each forward
+        step, where the GPU still has queued work to cover the close, before it
+        waits on an empty request queue with no request left, and before it
+        fires a control request. Takes effect only with TRTLLM_KV_DEFER_CLOSE=1
+        on a target-model manager. Disarming flushes. Returns whether deferred
+        closing is armed.
+        """
+        was_armed = self._defer_close_armed
+        self._defer_close_armed = bool(armed) and self._defer_close_enabled
+        if self._defer_close_armed and not was_armed:
+            logger.info(
+                "KVCacheManagerV2: deferred close armed (TRTLLM_KV_DEFER_CLOSE=1): finished "
+                "requests' caches close after the next forward step is launched"
+            )
+        if not self._defer_close_armed:
+            self.flush_deferred_closes()
+        return self._defer_close_armed
+
+    @contextlib.contextmanager
+    def deferring_closes(self) -> Iterator[None]:
+        """Let free_resources defer the close of the requests it releases for
+        the duration, while deferred closing is armed (arm_deferred_close).
+        For the release of requests that have terminated only."""
+        outer = self._defer_close_scope
+        self._defer_close_scope = True
+        try:
+            yield
+        finally:
+            self._defer_close_scope = outer
+
+    def flush_deferred_closes(self) -> int:
+        """Close the caches whose close() free_resources deferred, in the order
+        their requests finished. Returns how many were closed."""
+        pending, self._deferred_closes = self._deferred_closes, ()
+        for request_id, kv_cache in pending:
+            kv_cache.close()
+            self.impl.clear_stats_excluded(request_id)
+        return len(pending)
+
     def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
         # A request awaiting preemption can still be cancelled or fail while
         # its saves drain, and a dead entry here would block every later
@@ -5680,6 +5748,16 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
             self.impl.clear_stats_excluded(request.py_request_id)
             return
         kv_cache.discard_pending_stats()
+        # Dummy requests reuse their ids, which the stats-exclusion bookkeeping
+        # done after close() is keyed on, so they always close at once.
+        if self._defer_close_armed and self._defer_close_scope and not request.is_dummy:
+            if request.py_request_id in self._early_freed_index_requests:
+                self._early_freed_index_requests.discard(request.py_request_id)
+            else:
+                self._detach_index_slot(kv_cache)
+                self.index_mapper.remove_sequence(request.py_request_id)
+            self._deferred_closes = (*self._deferred_closes, (request.py_request_id, kv_cache))
+            return
         kv_cache.close()
         self.impl.clear_stats_excluded(request.py_request_id)
         if request.py_request_id in self._early_freed_index_requests:
@@ -5931,6 +6009,8 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool
         return bool(has_invalid_values)
 
     def shutdown(self):
+        self._defer_close_armed = False
+        self.flush_deferred_closes()
         if self.sparse_metadata_batch is not None:
             self.sparse_metadata_batch.record_read(self._stream.cuda_stream)
             self.sparse_metadata_batch.close()
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 9711f72f6f11..3d18f15bbaec 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -400,6 +400,10 @@ class PyExecutor:
     # If the number of micro batches is too large, the executor will spend too much host memory (No additional GPU memory is required).
     # 1024 in-flight micro batches can avoid synchronization in most cases and keep host memory usage low.
     MIN_ASYNC_MICRO_BATCH_NUM = 1024
+    # True while the overlap loop has armed KVCacheManagerV2's opt-in deferred
+    # close (TRTLLM_KV_DEFER_CLOSE=1). Declared on the class so instances built
+    # without running __init__ read it too.
+    _kv_defer_close_armed = False
 
     def __init__(
             self,
@@ -4902,6 +4906,10 @@ def _handle_control_request(self):
                      f"drain={pending.control_requires_drain} "
                      f"active_requests={len(self.active_requests)} "
                      f"waiting_queue={len(self.waiting_queue)}")
+        if self._kv_defer_close_armed:
+            # The action must not run while a finished request's cache is
+            # still open: its close() commits the last block and frees pages.
+            self.kv_cache_manager.flush_deferred_closes()
         # Quiesce the device before the action mutates GPU state. Under the
         # overlap scheduler a previous batch's forward/sample kernels may still
         # be in flight, so an in-place update_weights reload (or sleep/wakeup
@@ -5222,6 +5230,8 @@ def _maybe_rebalance_kv_pools(self) -> None:
 
         torch.cuda.current_stream().synchronize()
         self._consume_previous_batch_for_rebalance()
+        if self._kv_defer_close_armed:
+            self.kv_cache_manager.flush_deferred_closes()
         self._rebalance_kv_pools_now()
 
     def _rebalance_kv_pools_now(self) -> None:
@@ -5422,6 +5432,9 @@ def _executor_loop_overlap(self):
         torch.cuda.set_device(self.device_id)
         # ensure the context is created, otherwise, some MPI calls will fail.
         CUASSERT(cudart.cudaSetDevice(self.device_id))
+        self._kv_defer_close_armed = (
+            self._is_kv_manager_v2
+            and self.kv_cache_manager.arm_deferred_close(True))
         with self._profiler() as profile_step, self.hang_detector:
             iter_start_time = time.time()
             iter_stats = None
@@ -5639,6 +5652,12 @@ def _executor_loop_overlap(self):
                         self._maybe_prefetch_next_iter_mm_encoders(
                             scheduled_batch)
 
+                if self._kv_defer_close_armed:
+                    # Close the caches of requests that finished in the
+                    # previous iteration here, after this iteration's forward
+                    # step is queued and before waiting for the previous one.
+                    self.kv_cache_manager.flush_deferred_closes()
+
                 if self.previous_batch is not None and should_process_previous_batch:
                     self._update_requests(self.previous_batch.sample_state)
                     # Turning off speculative decoding when Acceptance Rate is low.
@@ -5773,6 +5792,9 @@ def _executor_loop_overlap(self):
                     self.disagg.pace_idle()
 
                 self.iter_counter += 1
+        if self._kv_defer_close_armed:
+            self._kv_defer_close_armed = False
+            self.kv_cache_manager.arm_deferred_close(False)
 
     @nvtx_range("_accept_draft_tokens")
     def _accept_draft_tokens(
@@ -6097,6 +6119,11 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue,
         idle = (total_num_live_requests == 0 and len(waiting_queue) == 0
                 and not self.is_shutdown
                 and not self._has_pending_connector_transfers())
+        if idle and self._kv_defer_close_armed:
+            # No forward step is coming to flush the caches of the requests
+            # that just finished, so close them before waiting: the next
+            # request then sees their pages and last blocks when scheduled.
+            self.kv_cache_manager.flush_deferred_closes()
         if idle:
             # In Ray path (TLLM_DISABLE_MPI=1), use a periodic heartbeat timeout so rank 0
             # reaches the broadcast path regularly to prevent trtllm-serve timeout when idle.
@@ -8436,7 +8463,11 @@ def _free_request_resources(self, request: LlmRequest) -> None:
         self.disagg.forget_request(request.py_request_id)
 
     def _do_terminate_request(self, request: LlmRequest) -> None:
-        self._free_request_resources(request)
+        # Only a terminated request's cache may close late; preemption and
+        # recompute pauses release theirs through other paths, at once.
+        with (self.kv_cache_manager.deferring_closes()
+              if self._kv_defer_close_armed else nullcontext()):
+            self._free_request_resources(request)
         # Cancellation and request-scoped failures can terminate before the
         # normal post-prefill release point, including with a partial buffer.
         _strip_py_multimodal_data_post_prefill(request)
diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
index 01f338aa3725..95539c89adf4 100644
--- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
@@ -15,6 +15,7 @@
 
 import enum
 import os
+import time
 from collections import Counter
 from typing import Callable, Optional
 
@@ -249,6 +250,138 @@ def withdraw_if_not_run(
         return True
 
 
+class ContextPageBackoff:
+    """Stop retrying context admission that the KV pool keeps refusing.
+
+    Admitting a first context chunk is not a cheap check. ``prepare_context``
+    creates the request's cache and claims its reusable prefix, the resume
+    brings that prefix onto the GPU (onboarding pages that live on a lower
+    tier and evicting others to make room), and when the chunk's own pages
+    then do not fit, the whole claim is rolled back and freed. The cost grows
+    with the prompt, not with the chunk. With the pool near full, every
+    pending request repeats this on every pass and is refused every time, so
+    host time per step grows with the number of pending requests while no
+    context work is admitted.
+
+    With ``passes`` N, once a pass that schedules generation sees a context
+    admission refused for pages, the rest of that pass tries no more first
+    chunks, and the next N passes that schedule generation try none either.
+    One pass in N + 1 then probes whether pages have come back. Chunk
+    continuations of started prefills are always tried, and passes without
+    generation are never held: with nothing decoding, nothing else can use
+    the step.
+
+    The state is a pure function of the sequence of scheduling decisions, so
+    ranks that schedule the same requests stay in lockstep.
+    """
+
+    def __init__(self, passes: int):
+        if passes <= 0:
+            raise ValueError(f"passes must be positive, got {passes}")
+        self.passes = passes
+        self._remaining = 0
+
+    @property
+    def remaining(self) -> int:
+        """Generation passes left before first chunks are tried again."""
+        return self._remaining
+
+    def holds_first_chunks(self, has_generation: bool, refused_this_pass: bool) -> bool:
+        """Whether this pass must skip first-chunk context admission."""
+        return has_generation and (self._remaining > 0 or refused_this_pass)
+
+    def record_pass(self, has_generation: bool, refused_this_pass: bool) -> None:
+        """Advance the backoff by one scheduling pass."""
+        if not has_generation:
+            return
+        if refused_this_pass:
+            self._remaining = self.passes
+        elif self._remaining > 0:
+            self._remaining -= 1
+
+
+class _AdmissionDiagnostics:
+    """Opt-in host-time accounting of phase-2 context admission.
+
+    Logs, every ``log_passes`` scheduling passes, how many context admissions
+    were attempted and refused for pages, and where the host time of the
+    admission phase went. A pass slower than ``slow_pass_ms`` is also logged
+    on its own, at most ``max_slow_logs`` times per window.
+    """
+
+    _SECTIONS = ("prepare", "allocate", "release")
+
+    def __init__(self, log_passes: int = 1024, slow_pass_ms: float = 100.0, max_slow_logs: int = 4):
+        self.log_passes = log_passes
+        self.slow_pass_ms = slow_pass_ms
+        self.max_slow_logs = max_slow_logs
+        self._window: Counter = Counter()
+        self._pass: Counter = Counter()
+        self._pass_start = 0.0
+        self._slow_logged = 0
+
+    def wrap(self, name: str, fn: Callable) -> Callable:
+        """Return *fn* timed into section *name* of the current pass."""
+
+        def timed(*args, **kwargs):
+            start = time.perf_counter()
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                self._pass[f"{name}_ms"] += (time.perf_counter() - start) * 1e3
+                self._pass[f"{name}_calls"] += 1
+
+        return timed
+
+    def begin_pass(self) -> None:
+        self._pass.clear()
+        self._pass_start = time.perf_counter()
+
+    def count(self, key: str, n: int = 1) -> None:
+        self._pass[key] += n
+
+    def end_pass(self, num_pending: int, num_scheduled_ctx: int) -> None:
+        pass_ms = (time.perf_counter() - self._pass_start) * 1e3
+        stats = self._pass
+        stats["phase2_ms"] = pass_ms
+        window = self._window
+        window["passes"] += 1
+        window.update(stats)
+        window["max_phase2_ms"] = max(window["max_phase2_ms"], pass_ms)
+        if stats["page_refused"]:
+            window["refused_passes"] += 1
+            window["refused_pass_ms"] += pass_ms
+        if pass_ms > self.slow_pass_ms:
+            window["slow_passes"] += 1
+            window["slow_pass_ms"] += pass_ms
+            if self._slow_logged < self.max_slow_logs:
+                self._slow_logged += 1
+                logger.info(
+                    f"[V2Scheduler] slow admission pass {pass_ms:.1f} ms: pending {num_pending}, "
+                    f"attempts {stats['attempts']}, page refusals {stats['page_refused']}, "
+                    f"admitted {num_scheduled_ctx}, held by backoff {stats['held']}, "
+                    + ", ".join(
+                        f"{s} {stats[s + '_calls']} calls {stats[s + '_ms']:.1f} ms"
+                        for s in self._SECTIONS
+                    )
+                )
+        if window["passes"] >= self.log_passes:
+            logger.info(
+                f"[V2Scheduler] admission diagnostics over {window['passes']} passes: "
+                f"attempts {window['attempts']}, page refusals {window['page_refused']} "
+                f"in {window['refused_passes']} passes ({window['refused_pass_ms']:.0f} ms), "
+                f"held by backoff {window['held']}, phase-2 total {window['phase2_ms']:.0f} ms, "
+                f"max {window['max_phase2_ms']:.1f} ms, passes > {self.slow_pass_ms:.0f} ms "
+                f"{window['slow_passes']} ({window['slow_pass_ms']:.0f} ms), "
+                + ", ".join(
+                    f"{s} {window[s + '_calls']} calls {window[s + '_ms']:.0f} ms"
+                    for s in self._SECTIONS
+                )
+            )
+            window.clear()
+            self._slow_logged = 0
+
+
 class ShortestPrefillFirstOrder:
     """Order pending context requests by the prefill work they still need.
 
@@ -480,6 +613,31 @@ def __init__(
                 _read_non_negative_int_env("TRTLLM_SCHED_SPF_MAX_WAIT_PASSES", 128)
             )
         self._context_admission_stats: Counter = Counter()
+
+        # Opt-in (default off): back off first-chunk context admission after
+        # the KV pool refuses one; see ContextPageBackoff.
+        backoff_passes = _read_non_negative_int_env("TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES", 0)
+        self._ctx_page_backoff = ContextPageBackoff(backoff_passes) if backoff_passes > 0 else None
+        # Set by the context admission paths when the KV pool refused pages
+        # during the current pass.
+        self._ctx_page_refused = False
+        # Opt-in (default off) host-time accounting of context admission.
+        self._admission_diag = None
+        if os.environ.get("TRTLLM_SCHED_ADMISSION_DIAG", "0") == "1":
+            diag = _AdmissionDiagnostics(
+                log_passes=_read_non_negative_int_env("TRTLLM_SCHED_ADMISSION_DIAG_PASSES", 1024)
+                or 1024
+            )
+            self._prepare_context_pair = diag.wrap("prepare", self._prepare_context_pair)
+            self._try_allocate_context = diag.wrap("allocate", self._try_allocate_context)
+            self.release_context_admission = diag.wrap("release", self.release_context_admission)
+            self._admission_diag = diag
+        if self._ctx_page_backoff is not None or self._admission_diag is not None:
+            logger.info(
+                "KVCacheV2Scheduler context page backoff: "
+                f"passes={backoff_passes}, admission_diagnostics="
+                f"{self._admission_diag is not None}"
+            )
         if self._prefill_decode_interval is not None or self._shortest_prefill_first is not None:
             logger.info(
                 "KVCacheV2Scheduler context admission: "
@@ -722,6 +880,11 @@ def preempt_for_pages(req: LlmRequest) -> bool:
         ctx_candidates, deferred_by_interval = self._admissible_pending_ctx(
             pending_ctx, requests_list, has_generation
         )
+        self._ctx_page_refused = False
+        backoff = self._ctx_page_backoff
+        diag = self._admission_diag
+        if diag is not None:
+            diag.begin_pass()
 
         # Prefix-aware skip: when several first-chunk context requests would
         # contribute the same not-yet-cached block, admitting them together
@@ -754,6 +917,15 @@ def preempt_for_pages(req: LlmRequest) -> bool:
                 continue
             if req.py_request_id in preempted_ids:
                 continue
+            if (
+                backoff is not None
+                and req.state_value == self._context_init_state_value
+                and req.is_first_context_chunk
+                and backoff.holds_first_chunks(has_generation, self._ctx_page_refused)
+            ):
+                if diag is not None:
+                    diag.count("held")
+                continue
             # Fast path: a pending context request whose draft-mirror admission
             # is certainly refused this iteration cannot be scheduled, so the
             # prefix-reuse probe below (two O(prompt_len / tokens_per_block)
@@ -796,6 +968,8 @@ def preempt_for_pages(req: LlmRequest) -> bool:
                 scheduled_encoder.append(req)
                 budget.commit(req, tokens, peft_pages)
             else:
+                if diag is not None:
+                    diag.count("attempts")
                 action, tokens, chunking_flag = self._try_schedule_context(
                     req, budget, preempt_for_pages
                 )
@@ -820,6 +994,10 @@ def preempt_for_pages(req: LlmRequest) -> bool:
                 if first_new_block is not None:
                     contributed_blocks.add(first_new_block)
 
+        if backoff is not None:
+            backoff.record_pass(has_generation, self._ctx_page_refused)
+        if diag is not None:
+            diag.end_pass(len(ctx_candidates), len(scheduled_ctx))
         if self._prefill_decode_interval is not None:
             self._record_context_admission(scheduled_ctx, has_generation, deferred_by_interval)
         if self._shortest_prefill_first is not None:
@@ -1137,6 +1315,7 @@ def _try_schedule_context_full(
         # before budget check.
         if not self._prepare_context_pair(req):
             logger.debug(f"prepare_context failed for context request {req.py_request_id}")
+            self._note_page_refusal()
             return ScheduleAction.STOP, 0, False
 
         context_tokens = req.context_remaining_length
@@ -1159,6 +1338,7 @@ def _try_schedule_context_full(
             # prepare_context to resume it.
             if preempt_for_pages(req):
                 return ScheduleAction.STOP, 0, False
+            self._note_page_refusal()
             return ScheduleAction.SKIP, 0, False
 
         cross_action = self._try_schedule_cross_context(req)
@@ -1168,6 +1348,12 @@ def _try_schedule_context_full(
 
         return ScheduleAction.SCHEDULED, req_tokens, False
 
+    def _note_page_refusal(self) -> None:
+        """Record that the KV pool refused a context admission this pass."""
+        self._ctx_page_refused = True
+        if self._admission_diag is not None:
+            self._admission_diag.count("page_refused")
+
     def _has_context_chunk_budget(self, budget: BudgetTracker) -> bool:
         remaining = budget.remaining_tokens
         return remaining is None or (
@@ -1202,6 +1388,7 @@ def _try_schedule_context_chunked(
         # Prepare context (create _KVCache, block reuse, resume — no resize)
         if not self._prepare_context_pair(req):
             logger.debug(f"prepare_context failed for chunked context request {req.py_request_id}")
+            self._note_page_refusal()
             return ScheduleAction.SKIP, 0, False
 
         # Calculate chunk size from remaining budget
@@ -1273,6 +1460,7 @@ def _try_schedule_context_chunked(
             # Out of pages, as in _try_schedule_context_full.
             if preempt_for_pages(req):
                 return ScheduleAction.STOP, 0, False
+            self._note_page_refusal()
             return ScheduleAction.SKIP, 0, False
 
         cross_action = self._try_schedule_cross_context(req)
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py
new file mode 100644
index 000000000000..5320e156e7fe
--- /dev/null
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py
@@ -0,0 +1,489 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""KVCacheManagerV2's opt-in deferred close (TRTLLM_KV_DEFER_CLOSE=1).
+
+While armed, free_resources inside deferring_closes() releases a terminated
+request's IndexMapper slot at once and leaves the cache's close() to
+flush_deferred_closes(); every other release closes at once. The overlap
+executor arms it for its loop, releases terminated requests inside the scope,
+and flushes after launching each forward step, before waiting for requests with
+none left, and before firing a control request.
+"""
+
+import types
+from unittest.mock import MagicMock, Mock
+
+from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+
+def _manager(events, *, enabled=True, num_pools=2, sparse=False):
+    mgr = object.__new__(KVCacheManagerV2)
+    mgr._defer_close_enabled = enabled
+    mgr._pending_preemption = {}
+    mgr.kv_connector_manager = None
+    mgr.conversation_manager = None
+    mgr._allocated_draft_lens = {}
+    mgr._request_stats_enabled_ids = set()
+    mgr._fresh_pages_filled = {}
+    mgr._disagg_receive_ready = {}
+    mgr.kv_cache_map = {}
+    mgr._early_freed_index_requests = set()
+    mgr.num_pools = num_pools
+    mgr.is_draft = False
+    mgr._stream = types.SimpleNamespace(cuda_stream=0)
+    mgr.impl = Mock()
+    mgr.impl.clear_stats_excluded.side_effect = lambda rid: events.append(("clear_stats", rid))
+    mgr.index_mapper = Mock()
+    mgr.index_mapper.remove_sequence.side_effect = lambda rid: events.append(("remove_seq", rid))
+    if sparse:
+        mgr.sparse_metadata_batch = Mock()
+        mgr.sparse_metadata_batch.remove.side_effect = lambda kv: events.append(
+            ("sparse_remove", kv.rid)
+        )
+        mgr.sparse_metadata_batch.close.side_effect = lambda: events.append(("sparse_close",))
+    return mgr
+
+
+def _cache(mgr, events, rid):
+    kv = Mock()
+    kv.rid = rid
+    kv.beam_width = 1
+    kv.close.side_effect = lambda: events.append(("close", rid))
+    kv.set_base_page_index_buf.side_effect = lambda beam, pool, buf: events.append(
+        ("unbind", rid, pool, buf)
+    )
+    mgr.kv_cache_map[rid] = kv
+    return kv
+
+
+def _request(rid, *, dummy=False):
+    return types.SimpleNamespace(py_request_id=rid, is_dummy=dummy)
+
+
+def _finish(mgr, request):
+    """Release *request* as the executor releases a terminated one."""
+    with mgr.deferring_closes():
+        mgr.free_resources(request)
+
+
+def test_off_unless_enabled_closes_at_once():
+    events = []
+    mgr = _manager(events, enabled=False)
+    _cache(mgr, events, 7)
+    assert mgr.arm_deferred_close(True) is False
+    _finish(mgr, _request(7))
+    assert events == [("close", 7), ("clear_stats", 7), ("remove_seq", 7)]
+    assert mgr._deferred_closes == ()
+
+
+def test_class_defaults_keep_managers_built_without_init_unchanged():
+    events = []
+    mgr = _manager(events)
+    del mgr._defer_close_enabled
+    _cache(mgr, events, 3)
+    assert mgr.arm_deferred_close(True) is False
+    _finish(mgr, _request(3))
+    assert events == [("close", 3), ("clear_stats", 3), ("remove_seq", 3)]
+
+
+def test_armed_releases_slot_now_and_closes_on_flush():
+    events = []
+    mgr = _manager(events, num_pools=2)
+    kv = _cache(mgr, events, 7)
+    assert mgr.arm_deferred_close(True) is True
+    _finish(mgr, _request(7))
+    # Slot released at once, after the cache is unbound from it; no close yet.
+    assert events == [("unbind", 7, 0, None), ("unbind", 7, 1, None), ("remove_seq", 7)]
+    kv.close.assert_not_called()
+    assert 7 not in mgr.kv_cache_map
+    events.clear()
+    assert mgr.flush_deferred_closes() == 1
+    assert events == [("close", 7), ("clear_stats", 7)]
+    assert mgr.flush_deferred_closes() == 0
+
+
+def test_flush_closes_in_finish_order():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    for rid in (5, 2, 9):
+        _cache(mgr, events, rid)
+    mgr.arm_deferred_close(True)
+    for rid in (5, 2, 9):
+        _finish(mgr, _request(rid))
+    events.clear()
+    assert mgr.flush_deferred_closes() == 3
+    assert [e for e in events if e[0] == "close"] == [("close", 5), ("close", 2), ("close", 9)]
+    assert [e for e in events if e[0] == "clear_stats"] == [
+        ("clear_stats", 5),
+        ("clear_stats", 2),
+        ("clear_stats", 9),
+    ]
+
+
+def test_each_stats_exclusion_is_cleared_after_its_close():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    for rid in (1, 4):
+        _cache(mgr, events, rid)
+    mgr.arm_deferred_close(True)
+    for rid in (1, 4):
+        _finish(mgr, _request(rid))
+    events.clear()
+    mgr.flush_deferred_closes()
+    assert events == [("close", 1), ("clear_stats", 1), ("close", 4), ("clear_stats", 4)]
+
+
+def test_dummy_requests_close_at_once_while_armed():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 11)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(11, dummy=True))
+    assert events == [("close", 11), ("clear_stats", 11), ("remove_seq", 11)]
+    assert mgr._deferred_closes == ()
+
+
+def test_early_released_slot_is_not_released_twice():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 6)
+    mgr._early_freed_index_requests.add(6)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(6))
+    assert events == []
+    assert 6 not in mgr._early_freed_index_requests
+    mgr.flush_deferred_closes()
+    assert events == [("close", 6), ("clear_stats", 6)]
+
+
+def test_sparse_batch_is_left_before_the_slot_is_released():
+    events = []
+    mgr = _manager(events, num_pools=1, sparse=True)
+    _cache(mgr, events, 8)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(8))
+    assert events == [("sparse_remove", 8), ("unbind", 8, 0, None), ("remove_seq", 8)]
+    mgr.sparse_metadata_batch.record_read.assert_called_once_with(0)
+
+
+def test_disarm_flushes_and_stops_deferring():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 1)
+    _cache(mgr, events, 2)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(1))
+    events.clear()
+    assert mgr.arm_deferred_close(False) is False
+    assert events == [("close", 1), ("clear_stats", 1)]
+    events.clear()
+    _finish(mgr, _request(2))
+    assert events == [("close", 2), ("clear_stats", 2), ("remove_seq", 2)]
+
+
+def test_shutdown_closes_deferred_caches_before_the_sparse_batch():
+    events = []
+    mgr = _manager(events, num_pools=1, sparse=True)
+    _cache(mgr, events, 3)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(3))
+    events.clear()
+    mgr.event_manager = None
+    mgr.shutdown()
+    assert events.index(("close", 3)) < events.index(("sparse_close",))
+    assert mgr._defer_close_armed is False
+    assert mgr._deferred_closes == ()
+
+
+def _overlap_loop_stub(kv_cache_manager):
+    executor = object.__new__(PyExecutor)
+    executor.device_id = 0
+    profiler = MagicMock()
+    profiler.__enter__.return_value = Mock()
+    executor._profiler = Mock(return_value=profiler)
+    executor.hang_detector = MagicMock()
+    executor.enable_iter_perf_stats = False
+    executor._resource_governor_enabled = False
+    executor._is_kv_manager_v2 = True
+    executor.kv_cache_manager = kv_cache_manager
+    executor._can_pause_for_rebalance = Mock(return_value=False)
+    executor._wait_for_model_engine_input_copy = Mock()
+    executor._prepare_and_schedule_batch = Mock(return_value=(None, None))
+    executor._mm_encoder_item_scheduling_enabled = False
+    executor.is_benchmark_disagg = False
+    executor._flush_pending_transfer_responses = Mock()
+    return executor
+
+
+def _patch_loop_cuda(monkeypatch):
+    for name in ("torch.cuda.set_device", "cudart.cudaSetDevice", "CUASSERT"):
+        monkeypatch.setattr(f"tensorrt_llm._torch.pyexecutor.py_executor.{name}", Mock())
+
+
+def test_overlap_loop_arms_for_its_run_and_disarms_on_exit(monkeypatch):
+    mgr = Mock()
+    mgr.arm_deferred_close = Mock(side_effect=lambda armed: armed)
+    executor = _overlap_loop_stub(mgr)
+    _patch_loop_cuda(monkeypatch)
+
+    PyExecutor._executor_loop_overlap(executor)
+
+    assert [c.args for c in mgr.arm_deferred_close.call_args_list] == [(True,), (False,)]
+    assert executor._kv_defer_close_armed is False
+
+
+def test_overlap_loop_leaves_the_manager_alone_when_not_enabled(monkeypatch):
+    mgr = Mock()
+    mgr.arm_deferred_close = Mock(return_value=False)
+    executor = _overlap_loop_stub(mgr)
+    _patch_loop_cuda(monkeypatch)
+
+    PyExecutor._executor_loop_overlap(executor)
+
+    mgr.arm_deferred_close.assert_called_once_with(True)
+    mgr.flush_deferred_closes.assert_not_called()
+    assert executor._kv_defer_close_armed is False
+
+
+class _QueueWait(Exception):
+    pass
+
+
+def _fetch_stub(kv_cache_manager, events, *, armed=True):
+    executor = object.__new__(PyExecutor)
+    executor.kv_cache_manager = kv_cache_manager
+    executor._kv_defer_close_armed = armed
+    executor.kv_connector_manager = None
+    executor.control_requests = []
+    executor.is_shutdown = False
+    executor._has_pending_connector_transfers = Mock(return_value=False)
+    executor._disable_mpi = False
+    executor.dist = types.SimpleNamespace(rank=0)
+    executor.request_accumulated = []
+    executor.hang_detector = MagicMock()
+    executor.executor_request_queue = Mock()
+
+    def wait(timeout):
+        events.append(("queue_wait", timeout))
+        raise _QueueWait
+
+    executor.executor_request_queue.get_from_request_queue.side_effect = wait
+    return executor
+
+
+def _fetch(executor, total_num_live_requests, waiting=()):
+    try:
+        PyExecutor._fetch_and_enqueue_requests(executor, list(waiting), total_num_live_requests)
+    except _QueueWait:
+        pass
+
+
+def test_last_request_closes_before_the_idle_queue_wait():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 4)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(4))
+    events.clear()
+    executor = _fetch_stub(mgr, events)
+
+    _fetch(executor, total_num_live_requests=0)
+
+    # The blocking wait for the next request starts only after the close.
+    assert events == [("close", 4), ("clear_stats", 4), ("queue_wait", None)]
+    assert mgr._deferred_closes == ()
+
+
+def test_busy_fetch_leaves_deferred_closes_to_the_loop():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 4)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(4))
+    events.clear()
+    executor = _fetch_stub(mgr, events)
+
+    _fetch(executor, total_num_live_requests=3)
+    _fetch(executor, total_num_live_requests=0, waiting=[object()])
+
+    assert [e[0] for e in events] == ["queue_wait", "queue_wait"]
+    assert len(mgr._deferred_closes) == 1
+
+
+def test_idle_fetch_does_not_flush_when_not_armed():
+    events = []
+    mgr = Mock()
+    executor = _fetch_stub(mgr, events, armed=False)
+
+    _fetch(executor, total_num_live_requests=0)
+
+    mgr.flush_deferred_closes.assert_not_called()
+    assert events == [("queue_wait", None)]
+
+
+def test_rebalance_flushes_deferred_closes_before_adjusting(monkeypatch):
+    order = []
+    executor = object.__new__(PyExecutor)
+    executor.kv_cache_manager = Mock()
+    executor.kv_cache_manager.flush_deferred_closes.side_effect = lambda: order.append("flush")
+    executor._agreed_need_adjustment = Mock(return_value=True)
+    executor._consume_previous_batch_for_rebalance = Mock(
+        side_effect=lambda: order.append("consume")
+    )
+    executor._rebalance_kv_pools_now = Mock(side_effect=lambda: order.append("adjust"))
+    executor._kv_defer_close_armed = True
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.current_stream", Mock()
+    )
+
+    PyExecutor._maybe_rebalance_kv_pools(executor)
+
+    assert order == ["consume", "flush", "adjust"]
+
+
+def test_armed_release_outside_the_terminate_scope_closes_at_once():
+    # Admission rollbacks and reservation drops call free_resources directly
+    # and count on the pages being free when it returns.
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 7)
+    mgr.arm_deferred_close(True)
+    mgr.free_resources(_request(7))
+    assert events == [("close", 7), ("clear_stats", 7), ("remove_seq", 7)]
+    assert mgr._deferred_closes == ()
+
+
+def test_preemption_releases_pages_at_once_while_armed():
+    # preempt_request reports the pages released; the scheduler retries the
+    # allocation right after each victim, so the close cannot wait for a flush.
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 9)
+    mgr.arm_deferred_close(True)
+    assert mgr.preempt_request(_request(9)) is True
+    assert events == [("close", 9), ("clear_stats", 9), ("remove_seq", 9)]
+    assert mgr._deferred_closes == ()
+
+
+def test_preemption_inside_the_terminate_scope_still_closes_at_once():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 2)
+    _cache(mgr, events, 3)
+    mgr.arm_deferred_close(True)
+    with mgr.deferring_closes():
+        mgr.free_resources(_request(2))
+        with mgr.deferring_closes():
+            pass
+        assert mgr._defer_close_scope is True
+    assert mgr._defer_close_scope is False
+    mgr.preempt_request(_request(3))
+    assert ("close", 3) in events and ("close", 2) not in events
+    assert [rid for rid, _ in mgr._deferred_closes] == [2]
+
+
+def _terminate_stub(kv_cache_manager, *, armed=True):
+    executor = object.__new__(PyExecutor)
+    executor.kv_cache_manager = kv_cache_manager
+    executor._kv_defer_close_armed = armed
+    executor.resource_manager = types.SimpleNamespace(
+        free_resources=kv_cache_manager.free_resources
+    )
+    executor._prefetched_request_ids = set()
+    executor._disagg_coordinator = Mock()
+    executor.gather_all_responses = False
+    executor.dist = types.SimpleNamespace(rank=1)
+    executor.result_wait_queues = {}
+    return executor
+
+
+def test_terminated_request_defers_and_recompute_pause_does_not():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 4)
+    _cache(mgr, events, 5)
+    mgr.arm_deferred_close(True)
+    executor = _terminate_stub(mgr)
+
+    PyExecutor._do_terminate_request(executor, _request(4))
+    assert ("close", 4) not in events
+    assert [rid for rid, _ in mgr._deferred_closes] == [4]
+
+    # The recompute-pause teardown releases through _free_request_resources.
+    PyExecutor._free_request_resources(executor, _request(5))
+    assert ("close", 5) in events
+    assert [rid for rid, _ in mgr._deferred_closes] == [4]
+
+
+def test_terminate_closes_at_once_when_not_armed():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 6)
+    executor = _terminate_stub(mgr, armed=False)
+
+    PyExecutor._do_terminate_request(executor, _request(6))
+
+    assert events == [("close", 6), ("clear_stats", 6), ("remove_seq", 6)]
+
+
+def _control_stub(kv_cache_manager, events, *, active=()):
+    executor = object.__new__(PyExecutor)
+    executor.kv_cache_manager = kv_cache_manager
+    executor._kv_defer_close_armed = True
+    executor.control_requests = [
+        types.SimpleNamespace(control_requires_drain=True, control_id=None)
+    ]
+    executor.active_requests = list(active)
+    executor.waiting_queue = []
+    executor._has_pending_connector_transfers = Mock(return_value=False)
+    executor.control_request_barrier = Mock()
+    executor.control_request_barrier.set.side_effect = lambda: events.append(("fire",))
+    executor.control_action_done = Mock()
+    return executor
+
+
+def test_drained_control_request_fires_after_the_deferred_closes(monkeypatch):
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 8)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(8))
+    events.clear()
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.synchronize",
+        Mock(side_effect=lambda: events.append(("sync",))),
+    )
+    executor = _control_stub(mgr, events)
+
+    PyExecutor._handle_control_request(executor)
+
+    assert events == [("close", 8), ("clear_stats", 8), ("sync",), ("fire",)]
+    assert mgr._deferred_closes == ()
+
+
+def test_parked_control_request_leaves_deferred_closes_alone():
+    events = []
+    mgr = _manager(events, num_pools=1)
+    _cache(mgr, events, 8)
+    mgr.arm_deferred_close(True)
+    _finish(mgr, _request(8))
+    events.clear()
+    executor = _control_stub(mgr, events, active=[object()])
+
+    PyExecutor._handle_control_request(executor)
+
+    assert events == []
+    assert len(mgr._deferred_closes) == 1
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py
index c3c45c0515ba..ef973b720d0d 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py
@@ -5366,3 +5366,199 @@ def test_started_prefill_may_still_preempt_under_spf(self, admission_env):
         # The protection applies to requests that have not started prefilling;
         # between started requests the usual newest-first preemption holds.
         assert ids(out.recompute_paused_requests) == [2]
+
+
+# ===========================================================================
+# Opt-in context page backoff and admission diagnostics
+# ===========================================================================
+
+
+_BACKOFF_ENV = "TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES"
+_DIAG_ENV = "TRTLLM_SCHED_ADMISSION_DIAG"
+_DIAG_PASSES_ENV = "TRTLLM_SCHED_ADMISSION_DIAG_PASSES"
+
+
+@pytest.fixture
+def backoff_env(monkeypatch):
+    for name in (
+        _INTERVAL_ENV,
+        _SPF_ENV,
+        _SPF_WAIT_ENV,
+        _BACKOFF_ENV,
+        _DIAG_ENV,
+        _DIAG_PASSES_ENV,
+    ):
+        monkeypatch.delenv(name, raising=False)
+    return monkeypatch
+
+
+def _refusing_manager(refused_ids):
+    """Manager whose pool refuses the context pages of every id in *refused_ids*."""
+    return make_kv_cache_manager(
+        resize_context_fn=lambda req, n: req.py_request_id not in refused_ids
+    )
+
+
+def _resize_attempts(mgr):
+    return [c.args[0].py_request_id for c in mgr.resize_context.call_args_list]
+
+
+class TestContextPageBackoff:
+    def test_off_by_default_retries_every_pass(self, backoff_env):
+        mgr = _refusing_manager({1, 2})
+        sched = make_scheduler(mgr)
+        assert sched._ctx_page_backoff is None
+        assert sched._admission_diag is None
+        gen = make_gen_request(0)
+        pending = [make_ctx_request(1, 100), make_ctx_request(2, 100)]
+        for _ in range(3):
+            out = sched.schedule_request([gen, *pending], set())
+            assert ids(out.generation_requests) == [0]
+            assert ids(out.context_requests) == []
+        assert _resize_attempts(mgr) == [1, 2] * 3
+
+    @pytest.mark.parametrize("value", ["-1", "x", "2.5"])
+    def test_invalid_value_rejected(self, backoff_env, value):
+        backoff_env.setenv(_BACKOFF_ENV, value)
+        with pytest.raises(ValueError, match=_BACKOFF_ENV):
+            make_scheduler(make_kv_cache_manager())
+
+    def test_zero_disables(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "0")
+        assert make_scheduler(make_kv_cache_manager())._ctx_page_backoff is None
+
+    @pytest.mark.parametrize("passes", [1, 3])
+    def test_refusal_holds_first_chunks_for_n_generation_passes(self, backoff_env, passes):
+        backoff_env.setenv(_BACKOFF_ENV, str(passes))
+        mgr = _refusing_manager({1, 2})
+        sched = make_scheduler(mgr)
+        gen = make_gen_request(0)
+        pending = [make_ctx_request(1, 100), make_ctx_request(2, 100)]
+        attempted_passes = []
+        for step in range(2 * (passes + 1)):
+            before = len(mgr.resize_context.call_args_list)
+            out = sched.schedule_request([gen, *pending], set())
+            # Decoding never waits on the backoff.
+            assert ids(out.generation_requests) == [0]
+            assert ids(out.context_requests) == []
+            attempts = _resize_attempts(mgr)[before:]
+            if attempts:
+                # The first refusal ends first-chunk admission for the pass.
+                assert attempts == [1]
+                attempted_passes.append(step)
+        assert attempted_passes == [0, passes + 1]
+
+    def test_admission_resumes_once_pages_return(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "2")
+        refused = {1}
+        mgr = _refusing_manager(refused)
+        sched = make_scheduler(mgr)
+        gen = make_gen_request(0)
+        ctx = make_ctx_request(1, 100)
+        assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == []
+        refused.clear()
+        assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == []
+        assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == []
+        assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == [1]
+        assert sched._ctx_page_backoff.remaining == 0
+
+    def test_successful_admission_does_not_arm(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "4")
+        sched = make_scheduler(_refusing_manager(set()))
+        gen = make_gen_request(0)
+        for request_id in (1, 2, 3):
+            out = sched.schedule_request([gen, make_ctx_request(request_id, 100)], set())
+            assert ids(out.context_requests) == [request_id]
+        assert sched._ctx_page_backoff.remaining == 0
+
+    def test_chunk_continuations_are_not_held(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "4")
+        mgr = _refusing_manager({2})
+        sched = make_scheduler(mgr, max_num_tokens=1000, ctx_chunk_config=(None, 64))
+        gen = make_gen_request(0)
+        refused = make_ctx_request(2, 100)
+        sched.schedule_request([gen, refused], set())
+        assert sched._ctx_page_backoff.remaining == 4
+        started = make_ctx_request(
+            1, context_remaining_length=500, prompt_len=5000, is_first_context_chunk=False
+        )
+        fresh = make_ctx_request(3, 100)
+        out = sched.schedule_request([gen, started, fresh], set())
+        assert ids(out.context_requests) == [1]
+        assert 3 not in _resize_attempts(mgr)
+
+    def test_not_held_without_generation(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "4")
+        mgr = _refusing_manager({1})
+        sched = make_scheduler(mgr)
+        gen = make_gen_request(0)
+        sched.schedule_request([gen, make_ctx_request(1, 100)], set())
+        assert sched._ctx_page_backoff.remaining == 4
+        # Nothing decodes: the step has no other use, so admission is tried.
+        out = sched.schedule_request([make_ctx_request(2, 100)], set())
+        assert ids(out.context_requests) == [2]
+        assert sched._ctx_page_backoff.remaining == 4
+
+    def test_held_context_is_not_a_stall(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "64")
+        sched = make_scheduler(_refusing_manager({1}))
+        gen = make_gen_request(0)
+        for _ in range(64):
+            sched.schedule_request([gen, make_ctx_request(1, 100)], set())
+        assert sched._stalled_schedules == 0
+
+    def test_composes_with_prefill_decode_interval(self, backoff_env):
+        backoff_env.setenv(_BACKOFF_ENV, "2")
+        backoff_env.setenv(_INTERVAL_ENV, "1")
+        mgr = _refusing_manager({1})
+        sched = make_scheduler(mgr)
+        gen = make_gen_request(0)
+        ctx = make_ctx_request(1, 100)
+        attempted = []
+        for step in range(6):
+            before = len(mgr.resize_context.call_args_list)
+            sched.schedule_request([gen, ctx], set())
+            if _resize_attempts(mgr)[before:]:
+                attempted.append(step)
+        # The interval never arms (nothing is admitted); the backoff alone
+        # spaces the refused attempts.
+        assert attempted == [0, 3]
+
+
+class TestAdmissionDiagnostics:
+    def test_window_summary_counts_attempts_and_refusals(self, backoff_env):
+        backoff_env.setenv(_DIAG_ENV, "1")
+        backoff_env.setenv(_DIAG_PASSES_ENV, "3")
+        mgr = _refusing_manager({1})
+        sched = make_scheduler(mgr)
+        gen = make_gen_request(0)
+        with patch("tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2.logger") as log:
+            for request_id in (2, 3, 4):
+                sched.schedule_request(
+                    [gen, make_ctx_request(1, 100), make_ctx_request(request_id, 10)], set()
+                )
+        summaries = [
+            c.args[0] for c in log.info.call_args_list if "admission diagnostics" in c.args[0]
+        ]
+        assert len(summaries) == 1
+        assert "over 3 passes: attempts 6, page refusals 3 in 3 passes" in summaries[0]
+        assert "allocate 6 calls" in summaries[0]
+        assert "prepare 6 calls" in summaries[0]
+
+    def test_diagnostics_do_not_change_decisions(self, backoff_env):
+        def run():
+            mgr = _refusing_manager({1})
+            sched = make_scheduler(mgr)
+            gen = make_gen_request(0)
+            return [
+                ids(
+                    sched.schedule_request(
+                        [gen, make_ctx_request(1, 100), make_ctx_request(r, 10)], set()
+                    ).context_requests
+                )
+                for r in (2, 3, 4)
+            ]
+
+        plain = run()
+        backoff_env.setenv(_DIAG_ENV, "1")
+        assert run() == plain
diff --git a/tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py b/tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py
new file mode 100644
index 000000000000..4bf20fe277b4
--- /dev/null
+++ b/tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py
@@ -0,0 +1,222 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The MuseSparkOpen input processor with the prefix-tokenization cache.
+
+With ``enable_tokenization_cache`` the processor's text-only path splices the
+token ids of a cached earlier prompt with a freshly tokenized tail. The ids
+must equal what the uncached path produces, turn by turn, for the prompt
+shapes an agentic conversation has under this checkpoint's chat template:
+multi-turn growth, consecutive assistant tool-call messages (whose end token
+changes once another message follows), reasoning blocks, tool results, and the
+``Current date`` line at the top of the system message.
+
+Runs against the real checkpoint's tokenizer and template, like
+``tests/unittest/llmapi/test_musespark_open_tokenizer.py``; it is skipped when
+``MUSESPARK_REFERENCE_CKPT`` is unset.
+"""
+
+import os
+
+import pytest
+
+from tensorrt_llm.sampling_params import SamplingParams
+
+pytestmark = pytest.mark.cpu_only
+
+CHECKPOINT = os.environ.get("MUSESPARK_REFERENCE_CKPT", "")
+
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "Bash",
+            "description": "Run a shell command and return its output. " * 20,
+            "parameters": {
+                "type": "object",
+                "properties": {"command": {"type": "string"}, "timeout": {"type": "number"}},
+                "required": ["command"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "Read",
+            "description": "Read a file from the local filesystem. " * 20,
+            "parameters": {
+                "type": "object",
+                "properties": {"file_path": {"type": "string"}},
+                "required": ["file_path"],
+            },
+        },
+    },
+]
+
+FILLER = (
+    "Line {i}: the build log reports 1234567 bytes copied in 0.25s; "
+    "naïve café ünïcode, tabs\tand  double  spaces, and a URL "
+    "https://example.com/a?b=c&d=e#f.\n"
+)
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    if not CHECKPOINT or not os.path.isdir(CHECKPOINT):
+        pytest.skip("checkpoint not present; set MUSESPARK_REFERENCE_CKPT")
+    from tensorrt_llm.tokenizer import load_custom_tokenizer
+
+    return load_custom_tokenizer("musespark_open", CHECKPOINT)
+
+
+def _processor(tokenizer, enable):
+    from tensorrt_llm.inputs.registry import create_input_processor
+
+    return create_input_processor(
+        CHECKPOINT, tokenizer, trust_remote_code=True, enable_tokenization_cache=enable
+    )
+
+
+def _render(tokenizer, messages, current_date="2026-10-20"):
+    return tokenizer.tokenizer.apply_chat_template(
+        messages,
+        tools=TOOLS,
+        add_generation_prompt=True,
+        tokenize=False,
+        current_date=current_date,
+    )
+
+
+def _filler(start, count):
+    return "".join(FILLER.format(i=i) for i in range(start, start + count))
+
+
+def _tool_call(call_id, name, **arguments):
+    return {"id": call_id, "type": "function", "function": {"name": name, "arguments": arguments}}
+
+
+def _tool_result(call_id, step):
+    return {"role": "tool", "tool_call_id": call_id, "content": _filler(step * 7, 30)}
+
+
+def _agentic_turns():
+    """Message lists of one growing agentic conversation, one per request."""
+    messages = [
+        {"role": "system", "content": "You are a coding agent. " + "Be careful. " * 400},
+        {"role": "user", "content": "Fix the failing test.\n" + _filler(0, 200)},
+    ]
+    turns = [list(messages)]
+    for step in range(12):
+        if step % 4 == 1:
+            # Two tool calls emitted as consecutive assistant messages. The
+            # first one ends with <|eot|> while it is the last message and with
+            # <|eom|> once the second follows, so the next prompt does not
+            # extend this one character for character.
+            messages.append(
+                {
+                    "role": "assistant",
+                    "reasoning_content": f"I should list files (step {step}).",
+                    "content": "",
+                    "tool_calls": [_tool_call(f"a{step}", "Bash", command=f"ls -la dir{step}")],
+                }
+            )
+            turns.append(list(messages))
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [_tool_call(f"b{step}", "Read", file_path=f"/src/f{step}.py")],
+                }
+            )
+            messages.append(_tool_result(f"a{step}", step))
+            messages.append(_tool_result(f"b{step}", step))
+        elif step % 4 == 3:
+            # A reasoning block plus a final answer, then a new user turn.
+            messages.append(
+                {
+                    "role": "assistant",
+                    "reasoning_content": "Thinking " * 50,
+                    "content": f"Done with step {step}.",
+                }
+            )
+            messages.append({"role": "user", "content": f"Next: step{step + 1}" + "x" * step})
+        else:
+            messages.append(
+                {
+                    "role": "assistant",
+                    "reasoning_content": f"Run the tests ({step}).",
+                    "content": "Running the tests now.",
+                    "tool_calls": [
+                        _tool_call(f"c{step}", "Bash", command="pytest -x", timeout=600)
+                    ],
+                }
+            )
+            messages.append(_tool_result(f"c{step}", step))
+        turns.append(list(messages))
+    return turns
+
+
+def test_cache_is_wired_only_when_enabled(tokenizer):
+    from tensorrt_llm._torch.models.modeling_musespark_open_omni import (
+        MuseSparkOpenOmniInputProcessor,
+    )
+
+    on = _processor(tokenizer, True)
+    off = _processor(tokenizer, False)
+    assert isinstance(on, MuseSparkOpenOmniInputProcessor)
+    assert on._prefix_token_cache is not None
+    assert off._prefix_token_cache is None
+
+
+def test_agentic_turns_are_id_identical(tokenizer):
+    on = _processor(tokenizer, True)
+    off = _processor(tokenizer, False)
+    params = SamplingParams(add_special_tokens=False)
+    turns = _agentic_turns()
+    for turn, messages in enumerate(turns):
+        prompt = _render(tokenizer, messages)
+        expected = tokenizer.encode(prompt, add_special_tokens=False)
+        cached, extra = on({"prompt": prompt}, params)
+        uncached, _ = off({"prompt": prompt}, params)
+        assert extra is None
+        assert uncached == expected, f"turn {turn}: uncached path changed"
+        assert cached == expected, f"turn {turn}: cached ids differ"
+    cache = on._prefix_token_cache
+    assert not cache.disabled
+    # Every turn but the first and the three that follow a lone assistant
+    # tool call extends the previous prompt, so the cache must splice them.
+    assert cache.hits >= len(turns) - 4, cache.summary()
+
+
+def test_date_change_misses_but_stays_identical(tokenizer):
+    on = _processor(tokenizer, True)
+    params = SamplingParams(add_special_tokens=False)
+    turns = _agentic_turns()
+    for date, messages in (("2026-10-20", turns[3]), ("2026-10-21", turns[4])):
+        prompt = _render(tokenizer, messages, current_date=date)
+        cached, _ = on({"prompt": prompt}, params)
+        assert cached == tokenizer.encode(prompt, add_special_tokens=False)
+    # The date sits at the top of the system message, so the second prompt
+    # does not extend the first one.
+    assert on._prefix_token_cache.hits == 0
+
+
+@pytest.mark.parametrize("stride_offset", [1, 7, 31, 97, 1009])
+def test_appends_starting_mid_word_are_identical(tokenizer, stride_offset):
+    on = _processor(tokenizer, True)
+    params = SamplingParams(add_special_tokens=False)
+    full = _render(tokenizer, _agentic_turns()[-1])
+    for end in range(len(full) // 2, len(full), len(full) // 40 + stride_offset):
+        prompt = full[:end]
+        cached, _ = on({"prompt": prompt}, params)
+        assert cached == tokenizer.encode(prompt, add_special_tokens=False), end
+    assert on._prefix_token_cache.hits > 0
+
+
+def test_cache_is_bypassed_when_special_tokens_are_requested(tokenizer):
+    on = _processor(tokenizer, True)
+    prompt = _render(tokenizer, _agentic_turns()[2])
+    ids, _ = on({"prompt": prompt}, SamplingParams(add_special_tokens=True))
+    # The processor never adds special tokens itself; the template carries BOS.
+    assert ids == tokenizer.encode(prompt, add_special_tokens=False)
+    cache = on._prefix_token_cache
+    assert cache.hits + cache.misses == 0
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py
new file mode 100644
index 000000000000..1930de55d98b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py
@@ -0,0 +1,328 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Inputs and built-in references shared by the draft-head tests (fast path ``draft_head``).
+
+Not a test module: the entry tests under ``ops/spec`` and the draft-head tests
+here import it. Every reference transform is the drafter's own
+``DFlashForCausalLM.apply_draft_logit_transform``, never a formula written
+here: on the device that method divides by the softcap as a multiplication by
+its fp32 reciprocal, which a hand-written ``x / softcap`` does not reproduce.
+"""
+
+from __future__ import annotations
+
+import types
+import zlib
+
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.draft_head import (  # noqa: E501
+    head as draft_head,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.draft_head import (  # noqa: E501
+    logit_transform,
+)
+from tensorrt_llm._torch.models.modeling_dflash import DFlash2CandidateSelector, DFlashForCausalLM
+
+draft_shard_topk = draft_head.draft_shard_topk
+draft_topk_merge = draft_head.draft_topk_merge
+draft_path_walk = draft_head.draft_path_walk
+
+VOCAB, HIDDEN, RANK, TOP_K, TP = 202048, 8192, 256, 16, 8
+
+#: The Muse Spark Open DFlash 2 drafter's declared transform: (output_multiplier, softcap).
+MUSESPARK = (0.1767766952966369, 20.0)
+
+#: Transform configurations the table tests cover besides MUSESPARK.
+OTHER_TRANSFORMS = {
+    "multiplier only": (0.5, None),
+    "softcap only": (1.0, 20.0),
+    "softcap 30": (1.0, 30.0),
+}
+
+#: Logit families of the end-to-end and merge matrices.
+FAMILIES = (
+    "gemm1",
+    "gemm10",
+    "gemm100",
+    "gemm400",
+    "grid",
+    "special_nan",
+    "straddle_up",
+    "straddle_down",
+    "ladder",
+    "saturated",
+    "edges",
+)
+
+
+def seed_of(text: str) -> int:
+    return zlib.crc32(text.encode())
+
+
+# --------------------------------------------------------------------------- drafters
+
+
+_DRAFTERS: dict = {}
+_TRANSFORMS: dict = {}
+
+
+def drafter(multiplier: float = 1.0, softcap: float | None = None) -> DFlashForCausalLM:
+    """A ``DFlashForCausalLM`` holding only the logit transform attributes (kept alive here)."""
+    key = (multiplier, softcap)
+    if key not in _DRAFTERS:
+        model = DFlashForCausalLM.__new__(DFlashForCausalLM)
+        torch.nn.Module.__init__(model)
+        model.logit_multiplier, model.logit_softcap = multiplier, softcap
+        _DRAFTERS[key] = model
+    return _DRAFTERS[key]
+
+
+def transform(config=MUSESPARK) -> logit_transform.DraftLogitTransform:
+    """The draft head's tables of ``config``'s transform on the current device (built once)."""
+    if config not in _TRANSFORMS:
+        built = logit_transform.build(drafter(*config), torch.device("cuda"))
+        assert built is not None
+        _TRANSFORMS[config] = built
+    return _TRANSFORMS[config]
+
+
+def builtin_transform(config, logits: torch.Tensor) -> torch.Tensor:
+    """The built-in transform: the drafter's own method, the call the DFlash worker makes."""
+    model = drafter(*config)
+    return DFlashForCausalLM.apply_draft_logit_transform(model, logits)
+
+
+# --------------------------------------------------------------------------- logits
+
+
+def _keys_values(keys) -> torch.Tensor:
+    return logit_transform.bf16_of_key16(torch.as_tensor(keys, device="cuda"))
+
+
+def _groups(config):
+    """(upper, lower): the largest tie group above +0.0 and below -0.0 of ``config``'s transform."""
+    groups = logit_transform.tie_groups(transform(config))
+    upper = max((g for g in groups if g[0] >= 0x8000), key=lambda g: g[1] - g[0])
+    lower = max((g for g in groups if g[1] <= 0x7FFF), key=lambda g: g[1] - g[0])
+    return upper, lower
+
+
+def logits(kind: str, rows: int, vocab: int, seed: int, shard: int | None = None, config=MUSESPARK):
+    """bf16 draft logits ``[rows, vocab]`` of one value family.
+
+    The tie families take their values from ``config``'s tie groups:
+    ``straddle_up`` ten copies each of two distinct members of the upper
+    saturating group at low and at high ids (both orders across rows),
+    ``straddle_down`` the same for the lower group above a row of lower-group
+    members, ``ladder`` consecutive bf16 values across the upper edge of the
+    identity band, ``saturated`` whole rows of upper-group members with +inf,
+    ``edges`` upper-group members on every shard's first and last columns.
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+
+    def pick(n):
+        return torch.randint(0, vocab, (rows, n), generator=g, device="cuda")
+
+    def grid(levels=8):
+        return (
+            torch.randint(0, levels, (rows, vocab), generator=g, device="cuda").float() * 0.25 - 3.0
+        )
+
+    def members(group, shape, with_inf=False):
+        first = group[0] if with_inf else max(group[0], 0x0080)
+        last = group[1] if with_inf else min(group[1], 0xFF7F)
+        keys = torch.randint(first, last + 1, shape, generator=g, device="cuda")
+        return _keys_values(keys)
+
+    if kind.startswith("gemm"):
+        sigma = float(kind[4:])
+        x = torch.randn(rows, vocab, generator=g, device="cuda") * sigma
+    elif kind == "logits":
+        x = torch.randn(rows, vocab, generator=g, device="cuda") * 3.0
+    elif kind == "grid":
+        x = torch.randint(0, 40, (rows, vocab), generator=g, device="cuda").float() * 0.125
+    elif kind == "zeros":
+        x = torch.zeros(rows, vocab, device="cuda")
+    elif kind in ("special", "special_nan"):
+        x = -(torch.rand(rows, vocab, generator=g, device="cuda") + 1.0)
+        for value, n in ((0.0, 12), (-0.0, 12), (float("inf"), 3), (float("-inf"), 2000)):
+            idx = pick(n)
+            x.scatter_(-1, idx, torch.full(idx.shape, value, device="cuda"))
+        if kind == "special_nan":
+            idx = pick(4)
+            x.scatter_(-1, idx, torch.full(idx.shape, float("nan"), device="cuda"))
+    elif kind in ("straddle_up", "straddle_down", "ladder", "saturated", "edges"):
+        upper, lower = _groups(config)
+        x = grid().to(torch.bfloat16)
+        quarter = vocab // 4
+        low = torch.randperm(quarter, generator=g, device="cuda")[:10]
+        high = vocab - 1 - torch.randperm(quarter, generator=g, device="cuda")[:10]
+        if kind in ("straddle_up", "straddle_down"):
+            group = upper if kind == "straddle_up" else lower
+            if kind == "straddle_down":
+                x = members(lower, (rows, vocab), with_inf=True)
+            a = _keys_values([max(group[0], 0x0080)])
+            b = _keys_values([min(group[1], 0xFF7F)])
+            for r in range(rows):
+                first, second = (a, b) if r % 2 == 0 else (b, a)
+                x[r, low] = first
+                x[r, high] = second
+        elif kind == "ladder":
+            hi = transform(config).hi
+            ladder = _keys_values(torch.arange(hi - 23, min(hi + 9, 0xFF80)))
+            for r in range(rows):
+                x[r, torch.randperm(vocab, generator=g, device="cuda")[: ladder.numel()]] = ladder
+        elif kind == "saturated":
+            x = members(upper, (rows, vocab), with_inf=True)
+        else:
+            assert shard is not None, "the edges family places its ties on shard edges"
+            cols = torch.tensor(
+                [c for r in range(vocab // shard) for c in (r * shard, (r + 1) * shard - 1)],
+                device="cuda",
+            )
+            x[:, cols] = members(upper, (rows, cols.numel()))
+        return x
+    else:
+        raise ValueError(kind)
+    return x.to(torch.bfloat16)
+
+
+# --------------------------------------------------------------------------- key exchange
+
+
+def split(full: torch.Tensor, ranks: int) -> list:
+    shard = full.shape[1] // ranks
+    return [full[:, r * shard : (r + 1) * shard].contiguous() for r in range(ranks)]
+
+
+def shard_keys(full: torch.Tensor, ranks: int, t=None) -> torch.Tensor:
+    """Every rank's shard keys, concatenated rank-major: what the all-gather returns."""
+    return torch.cat(
+        [draft_shard_topk(s, **draft_head.canon_kwargs(t)) for s in split(full, ranks)], dim=0
+    )
+
+
+def packed_keys(full: torch.Tensor, ranks: int, t=None) -> torch.Tensor:
+    """Every rank's fp32 key halves summed: what the all-reduce returns."""
+    packed = [
+        draft_shard_topk(s, rank=r, num_ranks=ranks, **draft_head.canon_kwargs(t))
+        for r, s in enumerate(split(full, ranks))
+    ]
+    return torch.stack(packed).sum(dim=0)
+
+
+def rank_keys(full: torch.Tensor, ranks: int, t=None) -> torch.Tensor:
+    """Every rank's own top-k halves (``wide_draft_topk``) summed: what its all-reduce returns."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import (  # noqa: E501
+        exchange,
+    )
+
+    ex = exchange.RankTopKExchange(None)
+    return torch.stack(
+        [
+            ex.keys(s, TOP_K, r, ranks, **draft_head.canon_kwargs(t))
+            for r, s in enumerate(split(full, ranks))
+        ]
+    ).sum(dim=0)
+
+
+def layout_keys(layout: str, full: torch.Tensor, ranks: int, t=None) -> torch.Tensor:
+    return {"int32": shard_keys, "halves": packed_keys, "rank": rank_keys}[layout](full, ranks, t)
+
+
+def merge_launch(layout: str, rows: int) -> dict:
+    """The merge geometry a step of ``rows`` rows runs with in ``layout``."""
+    if layout != "rank":
+        return {}
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import (  # noqa: E501
+        exchange,
+    )
+
+    return exchange.RankTopKExchange(None).merge_launch(rows)
+
+
+def assert_same_floats(a: torch.Tensor, b: torch.Tensor) -> None:
+    """Bitwise equality, NaN payloads aside (both NaN counts as equal)."""
+    assert a.shape == b.shape and a.dtype == b.dtype, (a.shape, b.shape, a.dtype, b.dtype)
+    assert logit_transform.same_bits(a, b), "values differ in their bits"
+
+
+# --------------------------------------------------------------------------- selector
+
+
+_SELECTOR = None
+
+
+def selector() -> DFlash2CandidateSelector:
+    """A random selector of the published drafter's shape."""
+    global _SELECTOR
+    if _SELECTOR is None:
+        g = torch.Generator(device="cuda").manual_seed(7)
+        pred = (0.05 * torch.randn(VOCAB, RANK, generator=g, device="cuda")).bfloat16()
+        succ = (0.05 * torch.randn(VOCAB, RANK, generator=g, device="cuda")).bfloat16()
+        proj = (0.01 * torch.randn(RANK, HIDDEN, generator=g, device="cuda")).bfloat16()
+        _SELECTOR = DFlash2CandidateSelector(pred, succ, proj, TOP_K, VOCAB, RANK).cuda()
+    return _SELECTOR
+
+
+def draft_step(requests, positions, kind, seed, ranks=TP):
+    """Per-rank bf16 shard logits of one draft step, plus its selector inputs."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    rows = requests * positions
+    hidden = torch.randn(rows, HIDDEN, generator=g, device="cuda").bfloat16()
+    shard = VOCAB // ranks
+    if kind.startswith("gemm") and requests * positions <= 56:
+        # The head's GEMM on a random weight of the given logit scale.
+        sigma = float(kind[4:])
+        shards = []
+        for r in range(ranks):
+            weight = (
+                torch.randn(shard, HIDDEN, generator=g, device="cuda") * (sigma / HIDDEN**0.5)
+            ).bfloat16()
+            shards.append(F.linear(hidden, weight))
+            del weight
+    else:
+        full = logits(kind, rows, VOCAB, seed, shard=shard)
+        shards = split(full, ranks)
+    anchor = torch.randint(0, VOCAB, (requests,), generator=g, device="cuda")
+    return shards, hidden.view(requests, positions, HIDDEN), anchor
+
+
+class DraftStub:
+    """What the worker's selector path reads off the drafter."""
+
+    select_candidate_path = DFlashForCausalLM.select_candidate_path
+
+    def __init__(self, candidate_selector):
+        self.candidate_selector = candidate_selector
+
+
+def builtin_full_logits(shards, requests, positions, config=None) -> torch.Tensor:
+    """The built-in ``[B, K, vocab]`` fp32 logits: each shard transformed, then all-gathered."""
+    if config is None:
+        parts = [s.float() for s in shards]
+    else:
+        parts = [builtin_transform(config, s) for s in shards]
+    return torch.cat(parts, dim=-1).float().view(requests, positions, VOCAB)
+
+
+def builtin_candidates(shards, requests, positions, config=None):
+    """``(unary, ids)``: ``torch.topk`` as the built-in selector runs it."""
+    return torch.topk(builtin_full_logits(shards, requests, positions, config), TOP_K, dim=-1)
+
+
+def builtin_draft_tokens(shards, block_hidden, anchor, config=None):
+    """The built-in path on the gathered logits: topk, selector, walk, scatter, argmax."""
+    requests, positions, _ = block_hidden.shape
+    full = builtin_full_logits(shards, requests, positions, config)
+    unary, ids = torch.topk(full, TOP_K, dim=-1)
+    block = DFlashForCausalLM.select_candidate_path(
+        DraftStub(selector()), full, ids, unary, block_hidden, anchor.long()
+    )
+    return torch.argmax(block.reshape(-1, VOCAB), dim=-1).to(torch.int32).view(requests, positions)
+
+
+def mapping(rank: int, ranks: int = TP):
+    return types.SimpleNamespace(tp_size=ranks, tp_rank=rank)
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py
new file mode 100644
index 000000000000..86a2af50f802
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py
@@ -0,0 +1,1099 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: the vocab-sharded draft head with the DFlash drafter's logit transform on.
+
+The claim is that with a drafter that declares a logit transform (the Muse
+Spark Open drafter's multiplier and softcap) the draft head's fast path still
+runs -- the deferring transform hands the bf16 shard through, no full-vocab
+all-gather and no full-vocab ``torch.topk`` -- and drafts the built-in tokens
+bit for bit: the built-in path transforms each shard with the drafter's own
+method, all-gathers, runs ``torch.topk`` and the selector. Then every path
+that keeps the built-in computation (sampled drafts, d2t, another layout,
+another tensor, a stale record) gets the built-in transform bit for bit, the
+fold's tail walk takes the same inputs, CUDA-graph replay follows new logits,
+and the install is atomic, idempotent and checked. A drafter whose transform
+attributes change after the install drafts the built-in tokens of its new
+transform through the built-in path, at every point of the step, and each
+reason a step takes the built-in path is warned once.
+
+TP ranks are emulated in one process: each rank's drafter, worker and head
+shard are built here, and the key exchange returns what every rank's kernels
+give for the same step.
+"""
+
+import types
+
+import pytest
+import torch
+
+__extra_import_path__ = ["."]
+from _draft_head_cases import (  # noqa: E402
+    FAMILIES,
+    HIDDEN,
+    MUSESPARK,
+    TOP_K,
+    TP,
+    VOCAB,
+    DraftStub,
+    assert_same_floats,
+    builtin_candidates,
+    builtin_draft_tokens,
+    builtin_full_logits,
+    builtin_transform,
+    draft_head,
+    draft_step,
+    logit_transform,
+    packed_keys,
+    rank_keys,
+    seed_of,
+    selector,
+    shard_keys,
+)
+
+from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM  # noqa: E402
+from tensorrt_llm._torch.modules.linear import TensorParallelMode  # noqa: E402
+from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor  # noqa: E402
+
+assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device"
+
+K = 7
+
+
+# --------------------------------------------------------------------------- the emulated group
+
+
+class StubHead:
+    """A gathered bf16 column-parallel head whose shard GEMM returns ``source`` (a copy)."""
+
+    tp_mode = TensorParallelMode.COLUMN
+    gather_output = True
+    padding_size = 0
+    gather_output_sizes = None
+    lora = None
+    enable_lm_head_tp_in_adp = False
+    bias = None
+
+    def __init__(self, source: torch.Tensor, ranks: int):
+        self.source = source
+        self.tp_size = ranks
+        self.out_features = source.shape[-1]
+        self.weight = torch.zeros(1, dtype=torch.bfloat16, device="cuda")
+
+    def apply_linear(self, hidden_states, bias):
+        return self.source.clone()
+
+
+class Exchange:
+    """The key all-reduce, emulated: every rank's keys of ``full`` (read at call time) summed."""
+
+    def __init__(self, full, ranks, t):
+        self.full, self.ranks, self.t = full, ranks, t
+        self.calls = 0
+
+    def __call__(self, packed):
+        self.calls += 1
+        chunk = packed_keys(self.full, self.ranks, self.t)
+        if packed.shape == chunk.shape:
+            return chunk
+        return rank_keys(self.full, self.ranks, self.t)
+
+
+def make_drafter(config, lm_head):
+    model = DFlashForCausalLM.__new__(DFlashForCausalLM)
+    torch.nn.Module.__init__(model)
+    model.logit_multiplier, model.logit_softcap = config
+    model.candidate_selector = selector()
+    model.logits_processor = LogitsProcessor()
+    model.lm_head = lm_head
+    return model
+
+
+def make_worker(rank, ranks=TP):
+    from tensorrt_llm._torch.speculative.dflash import DFlashWorker
+    from tensorrt_llm.mapping import Mapping
+
+    config = types.SimpleNamespace(max_draft_len=K, attention_backend="VANILLA")
+    return DFlashWorker(config, Mapping(world_size=ranks, tp_size=ranks, rank=rank))
+
+
+def spec_metadata(requests=1, advanced=False):
+    return types.SimpleNamespace(
+        wants_advanced_draft_sampling=advanced,
+        is_all_greedy_sample=not advanced,
+        is_cuda_graph=False,
+        draft_vocab_size=VOCAB,
+        vocab_size=VOCAB,
+        batch_indices_cuda=torch.arange(requests, dtype=torch.int32, device="cuda"),
+    )
+
+
+@pytest.fixture
+def single_process(monkeypatch):
+    """The install's TP agreements in one process: each emulated rank agrees with itself."""
+    monkeypatch.setattr(draft_head, "agreed_reason", lambda mapping, reason, transform: reason)
+    monkeypatch.setattr(draft_head, "agreed_refusal", lambda mapping, reason: reason)
+
+
+def install_rank(rank, local, full, ranks=TP, config=MUSESPARK, layout="halves"):
+    """One emulated rank: its worker and drafter with the draft head installed."""
+    worker = make_worker(rank, ranks)
+    drafter = make_drafter(config, StubHead(local, ranks))
+    llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter)
+    exchange = Exchange(full, ranks, None)
+    assert draft_head.install(llm, make_allreduce=lambda mapping: exchange)
+    exchange.t = worker.__dict__.get(draft_head.TRANSFORM_ATTR)
+    if layout == "int32":
+        worker.__dict__["_draft_head_allreduce"] = None
+    elif layout == "rank":
+        from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import (  # noqa: E501
+            exchange as wide,
+        )
+
+        assert wide.install(llm)
+    return llm
+
+
+def patch_allgather(monkeypatch, full, ranks, t, gathered=None):
+    """The head's NCCL all-gather: keys along dim 0; ``gathered`` (or a raise) along dim -1."""
+
+    def allgather(tensor, mapping, dim=-1, sizes=None):
+        if dim == 0:
+            return shard_keys(full, ranks, t)
+        if gathered is None:
+            raise AssertionError("the fast path must not all-gather the vocabulary")
+        return gathered.clone()
+
+    monkeypatch.setattr(draft_head, "allgather", allgather)
+
+
+def drafter_step(worker, drafter, hidden, anchor, md):
+    """``DFlashWorker._forward_impl`` from the drafter's logits processor to the draft tokens.
+
+    The statements the pin test (``test_modeling_v2_draft_transform_pin.py``)
+    holds the built-in forward to.
+    """
+    num_gens = hidden.shape[0] // K
+    gen_logits = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    apply_logit_transform = getattr(drafter, "apply_draft_logit_transform", None)
+    if apply_logit_transform is not None:
+        gen_logits = apply_logit_transform(gen_logits)
+    vocab_size = gen_logits.shape[-1]
+    gen_logits = gen_logits.reshape(num_gens, K, vocab_size)
+    gen_logits = worker._refine_block_logits(drafter, gen_logits, {}, md)
+    gen_logits = worker._apply_dflash2_selector(
+        drafter, gen_logits, hidden.reshape(num_gens, K, -1), anchor, md
+    )
+    return worker.sample_draft_tokens(gen_logits, md, num_gens, num_contexts=0)
+
+
+# --------------------------------------------------------------------------- end to end
+
+
+@pytest.mark.parametrize("kind", FAMILIES)
+@pytest.mark.parametrize("layout", ["int32", "halves", "rank"])
+@pytest.mark.parametrize("ranks", [4, 8])
+@pytest.mark.parametrize("requests", [1, 2, 8, 24])
+def test_selector_draft_tokens_equal_builtin(requests, ranks, layout, kind, monkeypatch):
+    """``selector_draft_tokens`` on every rank's bf16 shard with the transform = the built-in drafts."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import (  # noqa: E501
+        exchange as wide,
+    )
+
+    keep = make_drafter(MUSESPARK, None)
+    t = logit_transform.build(keep, torch.device("cuda"))
+    shards, block_hidden, anchor = draft_step(
+        requests, K, kind, seed_of(f"e2e{requests}{ranks}{kind}"), ranks=ranks
+    )
+    full = torch.cat(shards, dim=-1)
+    expected = builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK)
+    ref_unary, ref_ids = builtin_candidates(shards, requests, K, MUSESPARK)
+    patch_allgather(monkeypatch, full, ranks, t)
+    allreduce = {
+        "int32": None,
+        "halves": Exchange(full, ranks, t),
+        "rank": wide.RankTopKExchange(Exchange(full, ranks, t)),
+    }[layout]
+    for rank in range(ranks):
+        seen = {}
+
+        def walk(edges, unary, ids, seen=seen):
+            seen["unary"], seen["ids"] = unary.clone(), ids.clone()
+            return draft_head.draft_path_walk(edges, unary, ids)
+
+        got = draft_head.selector_draft_tokens(
+            shards[rank].view(requests, K, -1),
+            block_hidden,
+            anchor,
+            selector(),
+            types.SimpleNamespace(tp_size=ranks, tp_rank=rank),
+            allreduce,
+            walk=walk,
+            transform=t,
+        )
+        assert torch.equal(got, expected), f"rank {rank}: {got.tolist()} != {expected.tolist()}"
+        assert torch.equal(seen["ids"], ref_ids)
+        assert_same_floats(seen["unary"], ref_unary)
+
+
+def test_transform_is_a_required_keyword():
+    with pytest.raises(TypeError):
+        draft_head.selector_draft_tokens(None, None, None, None, None)
+
+
+@pytest.mark.parametrize("kind", ["gemm100", "straddle_up", "saturated"])
+@pytest.mark.parametrize("requests", [1, 8])
+def test_path_walk_tail_takes_the_builtin_inputs(requests, kind, monkeypatch):
+    """The fold's tail walk on the fast path's unary values = the tail on the built-in candidates."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec.draft_path_walk_tail import (  # noqa: E501
+        draft_path_walk_tail,
+    )
+
+    keep = make_drafter(MUSESPARK, None)
+    t = logit_transform.build(keep, torch.device("cuda"))
+    shards, block_hidden, anchor = draft_step(requests, K, kind, seed_of(f"tail{requests}{kind}"))
+    full = torch.cat(shards, dim=-1)
+    g = torch.Generator(device="cuda").manual_seed(requests)
+    accepted = torch.randint(0, VOCAB, (requests, K + 1), generator=g, device="cuda").int()
+    num_accepted = torch.randint(1, K + 2, (requests,), generator=g, device="cuda").int()
+    batch_indices = torch.arange(requests, dtype=torch.int32, device="cuda")
+    kv_lens = torch.randint(10, 100, (requests,), generator=g, device="cuda").int()
+    rewind = (1 - num_accepted).int()
+
+    # The built-in candidates and einsum operands, then the tail's reference.
+    sel = selector()
+    ref_unary, ref_ids = builtin_candidates(shards, requests, K, MUSESPARK)
+    gate = sel.hidden_projection(block_hidden.to(sel.hidden_projection.weight.dtype))
+    successors = sel.successor_codebook[ref_ids].float()
+    predecessor_ids = torch.cat(
+        (anchor.long()[:, None, None].expand(-1, 1, TOP_K), ref_ids[:, :-1]), 1
+    )
+    predecessors = sel.predecessor_codebook[predecessor_ids].float() * gate.float()[:, :, None]
+    edges = torch.einsum("blpr,blcr->blpc", predecessors, successors)
+    ref_kv = kv_lens.clone()
+    ref_tokens, ref_next = draft_path_walk_tail.reference(
+        edges, ref_unary, ref_ids, accepted, num_accepted, batch_indices, ref_kv, rewind
+    )
+    patch_allgather(monkeypatch, full, TP, t)
+    ours_kv = kv_lens.clone()
+
+    def walk(edges, unary, ids):
+        return draft_path_walk_tail(
+            edges, unary, ids, accepted, num_accepted, batch_indices, ours_kv, rewind
+        )
+
+    tokens, next_tokens = draft_head.selector_draft_tokens(
+        shards[5].view(requests, K, -1),
+        block_hidden,
+        anchor,
+        sel,
+        types.SimpleNamespace(tp_size=TP, tp_rank=5),
+        Exchange(full, TP, t),
+        walk=walk,
+        transform=t,
+    )
+    assert torch.equal(tokens, ref_tokens) and torch.equal(next_tokens, ref_next)
+    assert torch.equal(ours_kv, ref_kv)
+    assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK))
+
+
+# --------------------------------------------------------------------------- the gate
+
+
+def _forbid_full_vocab_topk(monkeypatch):
+    real_topk = torch.topk
+
+    def topk(x, *args, **kwargs):
+        if x.shape[-1] == VOCAB:
+            raise AssertionError("the fast path must not run torch.topk over the full vocabulary")
+        return real_topk(x, *args, **kwargs)
+
+    monkeypatch.setattr(torch, "topk", topk)
+
+
+@pytest.mark.parametrize("layout", ["int32", "halves", "rank"])
+@pytest.mark.parametrize("requests", [1, 5])
+def test_the_fast_path_runs_with_the_transform_on(requests, layout, monkeypatch, single_process):
+    shards, block_hidden, anchor = draft_step(
+        requests, K, "gemm100", seed_of(f"gate{layout}{requests}")
+    )
+    full = torch.cat(shards, dim=-1)
+    expected = builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK)
+    hidden = block_hidden.view(requests * K, HIDDEN)
+    llms = [install_rank(rank, shards[rank], full, layout=layout) for rank in range(TP)]
+    t = llms[0].spec_worker.__dict__[draft_head.TRANSFORM_ATTR]
+    for llm in llms:
+        # The drafter instance carries the deferring transform; the class is unchanged.
+        assert isinstance(
+            vars(llm.draft_model)[draft_head.TRANSFORM_METHOD], draft_head._DeferredTransform
+        )
+        assert DFlashForCausalLM.apply_draft_logit_transform is not None
+    patch_allgather(monkeypatch, full, TP, t)
+    _forbid_full_vocab_topk(monkeypatch)
+    calls = []
+    real = draft_head.selector_draft_tokens
+
+    def spy(*args, **kwargs):
+        calls.append(kwargs["transform"])
+        return real(*args, **kwargs)
+
+    monkeypatch.setattr(draft_head, "selector_draft_tokens", spy)
+    md = spec_metadata(requests)
+    for rank, llm in enumerate(llms):
+        worker, drafter = llm.spec_worker, llm.draft_model
+        out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+        assert out.dtype == torch.bfloat16
+        handed = drafter.apply_draft_logit_transform(out)
+        assert handed is out, "the deferring transform must hand the shard through"
+        block = handed.reshape(requests, K, -1)
+        refined = worker._refine_block_logits(drafter, block, {}, md)
+        assert refined.dtype == torch.bfloat16 and refined.data_ptr() == out.data_ptr()
+        assert worker._shard_selector_applies(refined, md, VOCAB)
+        stand_in = worker._apply_dflash2_selector(drafter, refined, block_hidden, anchor, md)
+        tokens = worker.sample_draft_tokens(stand_in, md, requests, num_contexts=0)
+        assert torch.equal(tokens, expected), f"rank {rank}"
+        assert calls[-1] is worker.__dict__[draft_head.TRANSFORM_ATTR]
+        # The whole sequence once more, as the built-in forward runs it.
+        assert torch.equal(drafter_step(worker, drafter, hidden, anchor, md), expected)
+
+
+def test_the_fast_path_runs_under_the_fold_tail(monkeypatch, single_process):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.verify_fold import (  # noqa: E501
+        fold,
+    )
+
+    requests = 2
+    shards, block_hidden, anchor = draft_step(requests, K, "straddle_up", seed_of("fold"))
+    full = torch.cat(shards, dim=-1)
+    expected = builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK)
+    hidden = block_hidden.view(requests * K, HIDDEN)
+    g = torch.Generator(device="cuda").manual_seed(3)
+    accepted = torch.randint(0, VOCAB, (requests, K + 1), generator=g, device="cuda").int()
+    num_accepted = torch.randint(1, K + 2, (requests,), generator=g, device="cuda").int()
+    llms = {rank: install_rank(rank, shards[rank], full) for rank in (0, 6)}
+    for llm in llms.values():
+        assert fold.install(llm.spec_worker, verify_fold=False, path_walk_tail=True)
+        assert draft_head._defers_transform(type(llm.spec_worker))
+    t = llms[0].spec_worker.__dict__[draft_head.TRANSFORM_ATTR]
+    patch_allgather(monkeypatch, full, TP, t)
+    _forbid_full_vocab_topk(monkeypatch)
+    for rank, llm in llms.items():
+        worker, drafter = llm.spec_worker, llm.draft_model
+        md = spec_metadata(requests)
+        kv_lens = torch.full((requests,), 50, dtype=torch.int32, device="cuda")
+        attn = types.SimpleNamespace(num_contexts=0, num_seqs=requests, kv_lens_cuda=kv_lens)
+        worker._kv_rewind_amount = (1 - num_accepted).int()
+        worker._kv_rewind_pending, worker._kv_rewind_nc, worker._kv_rewind_bs = True, 0, requests
+        step = fold._Step(attn_metadata=attn, spec_metadata=md, draft_model=drafter, warmup=False)
+        step.accepted, step.num_accepted = accepted, num_accepted
+        worker.__dict__["_fold_step"] = step
+        try:
+            tokens = drafter_step(worker, drafter, hidden, anchor, md)
+        finally:
+            worker.__dict__["_fold_step"] = None
+        assert step.tail is not None and torch.equal(tokens, expected), f"rank {rank}"
+        assert torch.equal(kv_lens, (50 - (1 - num_accepted)).clamp(min=0).int())
+
+
+def test_the_real_lm_head_hands_its_shard_through(single_process):
+    """With the gathered column-parallel ``LMHead`` the processor's shard reaches the selector untouched."""
+    from tensorrt_llm._torch.modules.embedding import LMHead
+    from tensorrt_llm.mapping import Mapping
+
+    rank = 3
+    head = LMHead(
+        VOCAB,
+        HIDDEN,
+        dtype=torch.bfloat16,
+        mapping=Mapping(world_size=TP, tp_size=TP, rank=rank),
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
+        gather_output=True,
+        reduce_output=False,
+    ).cuda()
+    with torch.no_grad():
+        head.weight.copy_((torch.randn_like(head.weight, dtype=torch.float32) * 0.05).bfloat16())
+    worker = make_worker(rank)
+    drafter = make_drafter(MUSESPARK, head)
+    llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter)
+    assert draft_head.install(llm, make_allreduce=lambda mapping: object())
+    hidden = torch.randn(K, HIDDEN, device="cuda").bfloat16()
+    out = drafter.logits_processor(hidden, head, None, True)
+    assert out.dtype == torch.bfloat16 and out.shape == (K, VOCAB // TP)
+    assert drafter.apply_draft_logit_transform(out) is out
+    refined = worker._refine_block_logits(drafter, out.reshape(1, K, -1), {}, spec_metadata())
+    assert refined.dtype == torch.bfloat16 and worker.__dict__[draft_head.OWED_ATTR]
+
+
+# --------------------------------------------------------------------------- built-in paths
+
+
+def _one_rank(monkeypatch, kind="gemm400", requests=1, rank=2):
+    shards, block_hidden, anchor = draft_step(requests, K, kind, seed_of(f"fallback{kind}{rank}"))
+    full = torch.cat(shards, dim=-1)
+    llm = install_rank(rank, shards[rank], full)
+    worker, drafter = llm.spec_worker, llm.draft_model
+    hidden = block_hidden.view(requests * K, HIDDEN)
+    return shards, block_hidden, anchor, full, worker, drafter, hidden
+
+
+def _builtin_shard(shard, requests):
+    return builtin_transform(MUSESPARK, shard).view(requests, K, -1)
+
+
+def test_sampled_drafts_get_the_builtin_transform(monkeypatch, single_process):
+    shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch)
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    handed = drafter.apply_draft_logit_transform(out)
+    refined = worker._refine_block_logits(
+        drafter, handed.reshape(1, K, -1), {}, spec_metadata(advanced=True)
+    )
+    # What the inherited _sample_dflash2_block receives: the built-in fp32 shard, bit for bit.
+    assert refined.dtype == torch.float32
+    assert_same_floats(refined, _builtin_shard(shards[2], 1))
+    assert not worker.__dict__[draft_head.OWED_ATTR]
+
+
+@pytest.mark.parametrize("why", ["d2t", "stride"])
+def test_layouts_the_selector_cannot_take_get_the_builtin_path(why, monkeypatch, single_process):
+    shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(monkeypatch)
+    md = spec_metadata()
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    handed = drafter.apply_draft_logit_transform(out)
+    block = handed.reshape(1, K, -1)
+    if why == "d2t":
+        worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda")
+    else:
+        block = block.as_strided(block.shape, (block.shape[-1] * K, 1, K), 0)
+        assert block.stride(-1) != 1 and block.data_ptr() == out.data_ptr()
+    refined = worker._refine_block_logits(drafter, block, {}, md)
+    assert refined.dtype == torch.float32
+    assert_same_floats(refined, builtin_transform(MUSESPARK, block))
+    if why == "d2t":
+        gathered = builtin_full_logits(shards, 1, K, MUSESPARK)
+        patch_allgather(monkeypatch, full, TP, None, gathered=gathered)
+        got = worker._apply_dflash2_selector(drafter, refined, block_hidden, anchor, md)
+        ref = DFlashForCausalLM.select_candidate_path(
+            DraftStub(selector()),
+            gathered.clone(),
+            *reversed(torch.topk(gathered, TOP_K, dim=-1)),
+            block_hidden,
+            anchor.long(),
+        )
+        assert torch.equal(got, ref)
+
+
+def test_any_other_tensor_gets_the_builtin_transform(monkeypatch, single_process):
+    shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch)
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    other = out.clone()
+    assert_same_floats(
+        drafter.apply_draft_logit_transform(other), builtin_transform(MUSESPARK, other)
+    )
+    # The shard itself once the record is spent.
+    assert_same_floats(drafter.apply_draft_logit_transform(out), builtin_transform(MUSESPARK, out))
+    # A full-vocab bf16 output of the wrapped processor is never deferred.
+    full_vocab = torch.randn(K, VOCAB, device="cuda").bfloat16()
+
+    class Inner(torch.nn.Module):
+        def forward(self, *args):
+            return full_vocab
+
+    drafter.logits_processor.inner = Inner()
+    assert drafter.logits_processor(hidden, None, None, True) is full_vocab
+    assert draft_head.SHARD_ATTR not in vars(drafter.logits_processor)
+    assert_same_floats(
+        drafter.apply_draft_logit_transform(full_vocab), builtin_transform(MUSESPARK, full_vocab)
+    )
+    # A drafter whose attributes moved away from the tables' transform is not deferred.
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    drafter.logit_softcap = 30.0
+    handed = drafter.apply_draft_logit_transform(out)
+    assert handed is not out and handed.dtype == torch.float32
+    assert_same_floats(handed, builtin_transform((MUSESPARK[0], 30.0), out))
+
+
+def test_a_stale_record_changes_nothing(monkeypatch, single_process):
+    shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch)
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    drafter.__dict__[draft_head.DEFERRED_ATTR] = (out.data_ptr() + 2, out.numel())
+    refined = worker._refine_block_logits(drafter, out.reshape(1, K, -1), {}, spec_metadata())
+    assert refined.dtype == torch.float32
+    assert_same_floats(refined, _builtin_shard(shards[2], 1))
+
+
+def test_the_selector_applies_the_transform_without_refine(monkeypatch, single_process):
+    """Bypassing refine: bf16 logits at the selector are untransformed, so it still drafts exactly."""
+    shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(
+        monkeypatch, "saturated"
+    )
+    t = worker.__dict__[draft_head.TRANSFORM_ATTR]
+    patch_allgather(monkeypatch, full, TP, t)
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    handed = drafter.apply_draft_logit_transform(out)
+    md = spec_metadata()
+    stand_in = worker._apply_dflash2_selector(
+        drafter, handed.reshape(1, K, -1), block_hidden, anchor, md
+    )
+    tokens = worker.sample_draft_tokens(stand_in, md, 1, num_contexts=0)
+    assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK))
+
+
+def test_a_step_that_owes_the_transform_refuses_to_sample(monkeypatch, single_process):
+    shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch)
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    handed = drafter.apply_draft_logit_transform(out)
+    refined = worker._refine_block_logits(drafter, handed.reshape(1, K, -1), {}, spec_metadata())
+    assert refined.dtype == torch.bfloat16
+    with pytest.raises(draft_head.DraftTransformStateError):
+        worker.sample_draft_tokens(refined, spec_metadata(), 1, num_contexts=0)
+
+
+def test_a_class_composed_above_with_its_own_refine_is_not_deferred(monkeypatch, single_process):
+    shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch)
+
+    class Above(type(worker)):
+        def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata):
+            return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata)
+
+    worker.__class__ = Above
+    assert not draft_head._defers_transform(Above)
+    out = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    handed = drafter.apply_draft_logit_transform(out)
+    assert handed.dtype == torch.float32
+    assert_same_floats(handed, builtin_transform(MUSESPARK, out))
+
+
+# --------------------------------------------------------------------------- CUDA graph
+
+
+@pytest.mark.parametrize("tail", [False, True])
+def test_cuda_graph_replay_follows_new_logits(tail, monkeypatch, single_process):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.verify_fold import (  # noqa: E501
+        fold,
+    )
+
+    requests, rank = 1, 2
+    shards, block_hidden, anchor = draft_step(requests, K, "gemm100", seed_of(f"graph{tail}"))
+    full = torch.cat(shards, dim=-1)
+    shard = VOCAB // TP
+    local = full[:, rank * shard : (rank + 1) * shard]
+    llm = install_rank(rank, local, full)
+    worker, drafter = llm.spec_worker, llm.draft_model
+    t = worker.__dict__[draft_head.TRANSFORM_ATTR]
+    tables = (t.canon.data_ptr(), t.values.data_ptr())
+    patch_allgather(monkeypatch, full, TP, t)
+    md = spec_metadata(requests)
+    hidden = block_hidden.view(requests * K, HIDDEN)
+    kv_lens = torch.full((requests,), 50, dtype=torch.int32, device="cuda")
+    accepted = torch.zeros(requests, K + 1, dtype=torch.int32, device="cuda")
+    num_accepted = torch.ones(requests, dtype=torch.int32, device="cuda")
+    if tail:
+        assert fold.install(worker, verify_fold=False, path_walk_tail=True)
+
+    def step():
+        if tail:
+            attn = types.SimpleNamespace(num_contexts=0, num_seqs=requests, kv_lens_cuda=kv_lens)
+            worker._kv_rewind_amount = (1 - num_accepted).int()
+            worker._kv_rewind_pending, worker._kv_rewind_nc, worker._kv_rewind_bs = (
+                True,
+                0,
+                requests,
+            )
+            s = fold._Step(attn_metadata=attn, spec_metadata=md, draft_model=drafter, warmup=False)
+            s.accepted, s.num_accepted = accepted, num_accepted
+            worker.__dict__["_fold_step"] = s
+        try:
+            return drafter_step(worker, drafter, hidden, anchor, md)
+        finally:
+            worker.__dict__["_fold_step"] = None
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        step()
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = step()
+    assert worker.__dict__[draft_head.TRANSFORM_ATTR] is t
+    assert (t.canon.data_ptr(), t.values.data_ptr()) == tables
+    for replay, kind in enumerate(("gemm400", "straddle_up", "saturated", "ladder", "special_nan")):
+        new_shards, new_hidden, new_anchor = draft_step(requests, K, kind, 60 + replay)
+        full.copy_(torch.cat(new_shards, dim=-1))
+        block_hidden.copy_(new_hidden)
+        anchor.copy_(new_anchor)
+        graph.replay()
+        expected = builtin_draft_tokens(new_shards, new_hidden, new_anchor, MUSESPARK)
+        assert torch.equal(out, expected), f"replay {replay} ({kind})"
+
+
+# --------------------------------------------------------------------------- install
+
+
+def test_install_is_idempotent(single_process):
+    shards, _, _ = draft_step(1, K, "gemm3", 5)
+    full = torch.cat(shards, dim=-1)
+    llm = install_rank(0, shards[0], full)
+    t = llm.spec_worker.__dict__[draft_head.TRANSFORM_ATTR]
+    shadow = vars(llm.draft_model)[draft_head.TRANSFORM_METHOD]
+    assert draft_head.install(llm, make_allreduce=lambda mapping: 1 / 0)
+    assert llm.spec_worker.__dict__[draft_head.TRANSFORM_ATTR] is t
+    assert vars(llm.draft_model)[draft_head.TRANSFORM_METHOD] is shadow
+
+
+def test_an_identity_transform_installs_nothing_new(monkeypatch, single_process):
+    shards, block_hidden, anchor = draft_step(1, K, "gemm3", 6)
+    full = torch.cat(shards, dim=-1)
+    llm = install_rank(1, shards[1], full, config=(1.0, None))
+    worker, drafter = llm.spec_worker, llm.draft_model
+    assert draft_head.TRANSFORM_ATTR not in worker.__dict__
+    assert draft_head.TRANSFORM_METHOD not in vars(drafter)
+    seen = []
+    real = draft_head.draft_shard_topk
+
+    class Spy:
+        def __call__(self, *args, **kwargs):
+            seen.append(sorted(kwargs))
+            return real(*args, **kwargs)
+
+    monkeypatch.setattr(draft_head, "draft_shard_topk", Spy())
+    patch_allgather(monkeypatch, full, TP, None)
+    md = spec_metadata()
+    tokens = drafter_step(worker, drafter, block_hidden.view(K, HIDDEN), anchor, md)
+    assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor))
+    assert seen and all("canon" not in kw for kw in seen)
+
+
+def test_tie_order_probe_agrees_with_this_torch_under_the_transform():
+    keep = make_drafter(MUSESPARK, None)
+    t = logit_transform.build(keep, torch.device("cuda"))
+    assert draft_head.tie_order_reason(selector(), TP, transform=t) is None
+    assert draft_head.tie_order_reason(selector(), 4, transform=t) is None
+
+
+def _stable_topk(x, k, dim=-1):
+    """A top-k that breaks ties by ascending index: torch's selection, another order."""
+    values, ids = torch.sort(x, dim=dim, descending=True, stable=True)
+    return values.narrow(dim, 0, k), ids.narrow(dim, 0, k)
+
+
+def test_tie_order_probe_has_teeth_under_the_transform(monkeypatch):
+    keep = make_drafter(MUSESPARK, None)
+    t = logit_transform.build(keep, torch.device("cuda"))
+    reason = draft_head.tie_order_reason(selector(), TP, topk=_stable_topk, transform=t)
+    assert reason is not None and "orders top-k ties differently" in reason
+    # Without canonical keys the transform probe's saturated ties come out wrong.
+    monkeypatch.setattr(draft_head, "canon_kwargs", lambda transform: {})
+    reason = draft_head.tie_order_reason(selector(), TP, transform=t)
+    assert reason is not None and "probe 1" in reason
+
+
+def test_transform_probe_logits_hold_the_tie_families():
+    keep = make_drafter(MUSESPARK, None)
+    t = logit_transform.build(keep, torch.device("cuda"))
+    probe = draft_head.transform_probe_logits(VOCAB, VOCAB // TP, t)
+    assert probe.dtype == torch.bfloat16 and probe.shape == (draft_head.PROBE_ROWS, VOCAB)
+    again = draft_head.transform_probe_logits(VOCAB, VOCAB // TP, t)
+    assert torch.equal(probe.view(torch.int16), again.view(torch.int16)), "not deterministic"
+    transformed = t.real(probe.cuda())
+    # Rows 0, 1, 2, 3, 6: more than top_k candidates share the row's largest transformed value.
+    for row in (0, 1, 2, 3, 6):
+        top = transformed[row].max()
+        assert (transformed[row] == top).sum() > TOP_K, row
+        tied = probe[row][transformed[row].cpu() == top.cpu()]
+        assert tied.view(torch.int16).unique().numel() > 1, row
+
+
+def test_a_refused_install_changes_nothing(monkeypatch, single_process):
+    from tensorrt_llm._torch.speculative.dflash import DFlashWorker
+
+    monkeypatch.setattr(
+        draft_head, "tie_order_reason", lambda *args, **kwargs: "probe rows [3] differ"
+    )
+    worker = make_worker(0)
+    shards, _, _ = draft_step(1, K, "gemm3", 8)
+    drafter = make_drafter(MUSESPARK, StubHead(shards[0], TP))
+    inner = drafter.logits_processor
+    llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter)
+    built = []
+    with pytest.raises(ValueError, match="probe rows"):
+        draft_head.install(llm, make_allreduce=lambda mapping: built.append(mapping))
+    assert type(worker) is DFlashWorker and drafter.logits_processor is inner and not built
+    assert draft_head.TRANSFORM_ATTR not in worker.__dict__
+    assert draft_head.TRANSFORM_METHOD not in vars(drafter)
+
+
+def test_ranks_whose_tables_differ_refuse(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep import (  # noqa: E501
+        dispatch,
+    )
+
+    keep = make_drafter(MUSESPARK, None)
+    t = logit_transform.build(keep, torch.device("cuda"))
+    mine = logit_transform.digest(t)
+    monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: [mine, "another"])
+    monkeypatch.setattr(dispatch, "agreed_refusal", lambda mapping, local: local)
+    reason = draft_head.agreed_reason(object(), None, t)
+    assert reason is not None and "differ in their drafter logit transform" in reason
+    monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: list(local))
+    assert draft_head.agreed_reason(object(), None, t) is None
+
+
+def test_an_exchange_one_rank_cannot_build_is_refused_on_every_rank(monkeypatch, single_process):
+    """``make_allreduce``'s refusal is agreed over the TP group: a rank that built it refuses too."""
+    from tensorrt_llm._torch.speculative.dflash import DFlashWorker
+
+    agreed = []
+    elsewhere = (
+        "the MNNVL FP32 all-reduce of the key exchange could not be built (on TP ranks [3] of 8)"
+    )
+
+    def agreed_refusal(mapping, reason):
+        agreed.append(reason)
+        return elsewhere
+
+    monkeypatch.setattr(draft_head, "agreed_refusal", agreed_refusal)
+    shards, _, _ = draft_step(1, K, "gemm3", 9)
+    for built in (object(), None):
+        worker = make_worker(0)
+        drafter = make_drafter(MUSESPARK, StubHead(shards[0], TP))
+        inner = drafter.logits_processor
+        llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter)
+        with pytest.raises(ValueError, match="could not be built"):
+            draft_head.install(llm, make_allreduce=lambda mapping, built=built: built)
+        assert type(worker) is DFlashWorker and drafter.logits_processor is inner
+        assert "_draft_head_allreduce" not in worker.__dict__
+        assert draft_head.TRANSFORM_ATTR not in worker.__dict__
+        assert draft_head.TRANSFORM_METHOD not in vars(drafter)
+    # This rank's own reason went into the agreement: none, then its own refusal.
+    assert agreed[0] is None and "could not be built" in agreed[1]
+
+
+def test_agreed_refusal_is_the_dispatch_agreement(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep import (  # noqa: E501
+        dispatch,
+    )
+
+    seen = []
+
+    def agreed_refusal(mapping, local):
+        seen.append((mapping, local))
+        return "agreed"
+
+    monkeypatch.setattr(dispatch, "agreed_refusal", agreed_refusal)
+    mapping = object()
+    assert draft_head.agreed_refusal(mapping, "mine") == "agreed"
+    assert seen == [(mapping, "mine")]
+
+
+# --------------------------------------------------------------------------- the tables' drafter
+
+
+#: Transforms a drafter may declare after the install, the identity among them.
+CHANGED = {
+    "identity": (1.0, None),
+    "softcap 30": (MUSESPARK[0], 30.0),
+    "multiplier only": (0.5, None),
+    "softcap only": (1.0, MUSESPARK[1]),
+}
+
+
+def staged_step(worker, drafter, hidden, anchor, md, change=None, when=None):
+    """``drafter_step`` up to the selector, with ``change()`` run at ``when``.
+
+    ``when`` is ``before`` (the logits processor), ``after_defer`` (the
+    drafter's transform) or ``after_refine`` (``_refine_block_logits``).
+    Returns what the selector returns.
+    """
+    num_gens = hidden.shape[0] // K
+    if when == "before":
+        change()
+    gen_logits = drafter.logits_processor(hidden, drafter.lm_head, None, True)
+    apply_logit_transform = getattr(drafter, "apply_draft_logit_transform", None)
+    if apply_logit_transform is not None:
+        gen_logits = apply_logit_transform(gen_logits)
+    if when == "after_defer":
+        change()
+    gen_logits = gen_logits.reshape(num_gens, K, gen_logits.shape[-1])
+    gen_logits = worker._refine_block_logits(drafter, gen_logits, {}, md)
+    if when == "after_refine":
+        change()
+    return worker._apply_dflash2_selector(
+        drafter, gen_logits, hidden.reshape(num_gens, K, -1), anchor, md
+    )
+
+
+def patch_builtin_gather(monkeypatch, shards, rank, config, requests):
+    """The head's vocabulary all-gather, checking this rank hands it the built-in transform.
+
+    ``config``'s built-in transform of this rank's shard must reach the
+    all-gather bit for bit (bf16 for the identity, fp32 otherwise); it returns
+    every rank's built-in transform, concatenated. A key exchange fails.
+    """
+    local = builtin_transform(config, shards[rank]).view(requests, K, -1)
+    gathered = torch.cat([builtin_transform(config, s) for s in shards], dim=-1)
+    gathered = gathered.view(requests, K, VOCAB)
+    calls = []
+
+    def allgather(tensor, mapping, dim=-1, sizes=None):
+        assert dim == -1, "a step whose tables do not hold must not exchange shard keys"
+        assert tensor.dtype == local.dtype, (tensor.dtype, local.dtype)
+        assert_same_floats(tensor.float(), local.float())
+        calls.append(dim)
+        return gathered.clone()
+
+    monkeypatch.setattr(draft_head, "allgather", allgather)
+    return calls
+
+
+def _spy_sharded_selector(monkeypatch):
+    calls = []
+    real = draft_head.selector_draft_tokens
+
+    def spy(*args, **kwargs):
+        calls.append(kwargs.get("transform"))
+        return real(*args, **kwargs)
+
+    monkeypatch.setattr(draft_head, "selector_draft_tokens", spy)
+    return calls
+
+
+def _tokens_of(worker, out):
+    """The step's draft tokens off the built-in selector's full-vocab block logits."""
+    assert getattr(out, draft_head.TOKENS_ATTR, None) is None, "the sharded selector drafted"
+    assert not worker.__dict__[draft_head.OWED_ATTR]
+    requests, positions, vocab = out.shape
+    return torch.argmax(out.reshape(-1, vocab), dim=-1).to(torch.int32).view(requests, positions)
+
+
+@pytest.mark.parametrize("when", ["before", "after_defer", "after_refine"])
+@pytest.mark.parametrize("name", list(CHANGED))
+@pytest.mark.parametrize("requests", [1, 3])
+def test_a_changed_transform_drafts_the_builtin_tokens(
+    requests, name, when, monkeypatch, single_process
+):
+    """The drafter's multiplier / softcap changed after the install: no stale tables, built-in drafts."""
+    config = CHANGED[name]
+    rank = 5
+    shards, block_hidden, anchor = draft_step(
+        requests, K, "saturated", seed_of(f"changed{requests}{name}{when}")
+    )
+    full = torch.cat(shards, dim=-1)
+    llm = install_rank(rank, shards[rank], full)
+    worker, drafter = llm.spec_worker, llm.draft_model
+    t = worker.__dict__[draft_head.TRANSFORM_ATTR]
+    expected = builtin_draft_tokens(shards, block_hidden, anchor, config)
+    gathers = patch_builtin_gather(monkeypatch, shards, rank, config, requests)
+    sharded = _spy_sharded_selector(monkeypatch)
+    md = spec_metadata(requests)
+    hidden = block_hidden.view(requests * K, HIDDEN)
+
+    def change():
+        drafter.logit_multiplier, drafter.logit_softcap = config
+
+    out = staged_step(worker, drafter, hidden, anchor, md, change, when)
+    assert not sharded and gathers == [-1]
+    assert torch.equal(_tokens_of(worker, out), expected)
+    # The tables stay installed and serve again once the drafter declares their transform.
+    assert worker.__dict__[draft_head.TRANSFORM_ATTR] is t
+    drafter.logit_multiplier, drafter.logit_softcap = MUSESPARK
+    patch_allgather(monkeypatch, full, TP, t)
+    tokens = worker.sample_draft_tokens(
+        staged_step(worker, drafter, hidden, anchor, md), md, requests, num_contexts=0
+    )
+    assert sharded == [t]
+    assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK))
+
+
+@pytest.mark.parametrize("when", ["before", "after_refine"])
+@pytest.mark.parametrize("name", ["identity", "softcap 30"])
+def test_a_changed_transform_under_the_fold_tail(name, when, monkeypatch, single_process):
+    """The fold's tail walk turns the step down too: the built-in drafts of the new transform."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.verify_fold import (  # noqa: E501
+        fold,
+    )
+
+    config = CHANGED[name]
+    requests, rank = 2, 6
+    shards, block_hidden, anchor = draft_step(
+        requests, K, "straddle_up", seed_of(f"fold{name}{when}")
+    )
+    full = torch.cat(shards, dim=-1)
+    llm = install_rank(rank, shards[rank], full)
+    worker, drafter = llm.spec_worker, llm.draft_model
+    assert fold.install(worker, verify_fold=False, path_walk_tail=True)
+    expected = builtin_draft_tokens(shards, block_hidden, anchor, config)
+    gathers = patch_builtin_gather(monkeypatch, shards, rank, config, requests)
+    sharded = _spy_sharded_selector(monkeypatch)
+    md = spec_metadata(requests)
+    hidden = block_hidden.view(requests * K, HIDDEN)
+    g = torch.Generator(device="cuda").manual_seed(4)
+    accepted = torch.randint(0, VOCAB, (requests, K + 1), generator=g, device="cuda").int()
+    num_accepted = torch.randint(1, K + 2, (requests,), generator=g, device="cuda").int()
+    kv_lens = torch.full((requests,), 50, dtype=torch.int32, device="cuda")
+    attn = types.SimpleNamespace(num_contexts=0, num_seqs=requests, kv_lens_cuda=kv_lens)
+    worker._kv_rewind_amount = (1 - num_accepted).int()
+    worker._kv_rewind_pending, worker._kv_rewind_nc, worker._kv_rewind_bs = True, 0, requests
+    step = fold._Step(attn_metadata=attn, spec_metadata=md, draft_model=drafter, warmup=False)
+    step.accepted, step.num_accepted = accepted, num_accepted
+    worker.__dict__["_fold_step"] = step
+
+    def change():
+        drafter.logit_multiplier, drafter.logit_softcap = config
+
+    try:
+        out = staged_step(worker, drafter, hidden, anchor, md, change, when)
+    finally:
+        worker.__dict__["_fold_step"] = None
+    assert not sharded and gathers == [-1] and step.tail is None
+    assert torch.equal(_tokens_of(worker, out), expected)
+
+
+def test_the_selector_refuses_tables_that_do_not_hold(monkeypatch, single_process):
+    shards, _, _, _, worker, drafter, _ = _one_rank(monkeypatch)
+    logits = shards[2].view(1, K, -1)
+    assert worker._draft_head_transform_for(drafter, logits) is not None
+    drafter.logit_softcap = 30.0
+    assert not worker._draft_head_takes(drafter, logits, spec_metadata(), VOCAB)
+    with pytest.raises(draft_head.DraftTransformStateError):
+        worker._draft_head_transform_for(drafter, logits)
+
+
+# --------------------------------------------------------------------------- warnings
+
+
+@pytest.fixture
+def logged_warnings(monkeypatch):
+    """The draft head's warnings logged from here on, with a fresh once-only record."""
+    logged = []
+
+    def log(severity, *msg):
+        message = " ".join(map(str, msg))
+        if severity == draft_head.logger.WARNING and "] draft_head:" in message:
+            logged.append(message)
+
+    monkeypatch.setattr(draft_head.logger, "_appeared_keys", set())
+    monkeypatch.setattr(draft_head.logger, "log", log)
+    return logged
+
+
+def _sampled(worker, drafter, md):
+    md.wants_advanced_draft_sampling = True
+
+
+def _d2t(worker, drafter, md):
+    worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda")
+
+
+def _softcap_30(worker, drafter, md):
+    drafter.logit_softcap = 30.0
+
+
+def _identity(worker, drafter, md):
+    drafter.logit_multiplier, drafter.logit_softcap = 1.0, None
+
+
+def _wrapped_processor(worker, drafter, md):
+    class Wrapper(torch.nn.Module):
+        def __init__(self, inner):
+            super().__init__()
+            self.inner = inner
+
+        def forward(self, *args):
+            return self.inner(*args)
+
+    drafter.logits_processor = Wrapper(drafter.logits_processor)
+
+
+def _fp32_head(worker, drafter, md):
+    drafter.lm_head.source = drafter.lm_head.source.float()
+
+
+def _class_above(worker, drafter, md):
+    class Above(type(worker)):
+        def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata):
+            return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata)
+
+    worker.__class__ = Above
+
+
+#: Each way a step with the transform installed takes the built-in path: (set-up, message part).
+FALLBACKS = {
+    "sampled": (_sampled, "sampled drafts"),
+    "d2t": (_d2t, "a d2t draft vocab map"),
+    "attributes": (_softcap_30, draft_head.STALE_TABLES),
+    "identity": (_identity, draft_head.STALE_TABLES),
+    "processor": (_wrapped_processor, "the drafter's logits processor is Wrapper"),
+    "fp32 logits": (_fp32_head, "returned torch.float32 logits"),
+    "class above": (_class_above, "the worker class Above overrides"),
+    "not deferred": (None, draft_head.NOT_DEFERRED),
+}
+
+
+@pytest.mark.parametrize("why", list(FALLBACKS))
+def test_a_fallback_warns_once(why, logged_warnings, monkeypatch, single_process):
+    setup, part = FALLBACKS[why]
+    shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(monkeypatch)
+    md = spec_metadata()
+    if setup is not None:
+        setup(worker, drafter, md)
+    gathered = builtin_full_logits(shards, 1, K, MUSESPARK)
+    patch_allgather(monkeypatch, full, TP, None, gathered=gathered)
+    for _ in range(3):
+        if why == "sampled":
+            # Sampled drafts leave the step after the refinement, for the built-in sampler.
+            out = drafter.apply_draft_logit_transform(
+                drafter.logits_processor(hidden, drafter.lm_head, None, True)
+            )
+            worker._refine_block_logits(drafter, out.reshape(1, K, -1), {}, md)
+        elif why == "not deferred":
+            # bf16 logits the deferring transform never saw.
+            out = drafter.logits_processor(hidden, drafter.lm_head, None, True).reshape(1, K, -1)
+            refined = worker._refine_block_logits(drafter, out, {}, md)
+            assert_same_floats(refined, builtin_transform(MUSESPARK, out))
+        else:
+            out = staged_step(worker, drafter, hidden, anchor, md)
+            assert getattr(out, draft_head.TOKENS_ATTR, None) is None
+    assert len(logged_warnings) == 1, logged_warnings
+    assert logged_warnings[0].startswith(
+        "[musespark_open] draft_head: a draft step takes the built-in path"
+    )
+    assert part in logged_warnings[0], logged_warnings[0]
+
+
+def test_each_reason_warns_once_and_the_fast_path_never(
+    logged_warnings, monkeypatch, single_process
+):
+    shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(monkeypatch)
+    t = worker.__dict__[draft_head.TRANSFORM_ATTR]
+    patch_allgather(monkeypatch, full, TP, t)
+    md = spec_metadata()
+    for _ in range(2):
+        out = staged_step(worker, drafter, hidden, anchor, md)
+        assert getattr(out, draft_head.TOKENS_ATTR, None) is not None
+    assert logged_warnings == []
+    softcap_30 = builtin_full_logits(shards, 1, K, (MUSESPARK[0], 30.0))
+    patch_allgather(monkeypatch, full, TP, None, gathered=softcap_30)
+    drafter.logit_softcap = 30.0
+    for _ in range(2):
+        staged_step(worker, drafter, hidden, anchor, md)
+    worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda")
+    drafter.logit_softcap = MUSESPARK[1]
+    musespark = builtin_full_logits(shards, 1, K, MUSESPARK)
+    patch_allgather(monkeypatch, full, TP, None, gathered=musespark)
+    for _ in range(2):
+        staged_step(worker, drafter, hidden, anchor, md)
+    assert len(logged_warnings) == 2, logged_warnings
+    assert draft_head.STALE_TABLES in logged_warnings[0] and "d2t" in logged_warnings[1]
+
+
+def test_without_a_transform_nothing_warns(logged_warnings, monkeypatch, single_process):
+    shards, block_hidden, anchor = draft_step(1, K, "gemm3", 11)
+    full = torch.cat(shards, dim=-1)
+    llm = install_rank(1, shards[1], full, config=(1.0, None))
+    worker, drafter = llm.spec_worker, llm.draft_model
+    worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda")
+    gathered = builtin_full_logits(shards, 1, K)
+    patch_allgather(monkeypatch, full, TP, None, gathered=gathered)
+    out = staged_step(worker, drafter, block_hidden.view(K, HIDDEN), anchor, spec_metadata())
+    assert getattr(out, draft_head.TOKENS_ATTR, None) is None
+    assert logged_warnings == []
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py
new file mode 100644
index 000000000000..a070c92d7fae
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py
@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the draft head's drafter logit transform tables (``kernels/draft_head/logit_transform.py``).
+
+The claim is that ``values`` is the drafter's own transform of every bf16 value
+bit for bit, in any layout, and that ``canon`` groups exactly the bf16 values
+whose transformed values have the same bits, in runs of consecutive keys, with
+the identity on ``[lo, hi]``. Checked exhaustively over the 65536 bf16 values
+on this device, for the Muse Spark Open drafter's transform and three others.
+"""
+
+import pytest
+import torch
+
+__extra_import_path__ = ["."]
+from _draft_head_cases import (  # noqa: E402
+    MUSESPARK,
+    OTHER_TRANSFORMS,
+    builtin_transform,
+    drafter,
+    logit_transform,
+    transform,
+)
+
+assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device"
+
+CONFIGS = {"musespark": MUSESPARK, **OTHER_TRANSFORMS}
+ALL = torch.arange(1 << 16, dtype=torch.int32)
+
+
+def _radix32(v: torch.Tensor) -> torch.Tensor:
+    word = v.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
+    key = torch.where(word >= 0x80000000, word ^ 0xFFFFFFFF, word ^ 0x80000000)
+    return torch.where(v.isnan(), torch.full_like(key, 0xFFFFFFFF), key)
+
+
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_values_are_the_builtin_transform_of_every_bf16_value(name):
+    config = CONFIGS[name]
+    t = transform(config)
+    x = logit_transform.bf16_of_key16(ALL.cuda())
+    assert t.values.dtype == torch.float32 and t.values.shape == (1 << 16,)
+    assert logit_transform.same_bits(t.values, builtin_transform(config, x))
+
+
+@pytest.mark.parametrize("requests", [1, 2, 8, 24])
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_values_hold_in_the_runtime_layout(name, requests):
+    """``values[key16(x)]`` equals the transform of x laid out as ``[B * 7, shard]``, shuffled."""
+    config = CONFIGS[name]
+    t = transform(config)
+    rows, shard = requests * 7, 25256
+    g = torch.Generator(device="cuda").manual_seed(requests)
+    keys = torch.randint(0, 1 << 16, (rows, shard), generator=g, device="cuda", dtype=torch.int32)
+    keys.view(-1)[: 1 << 16] = ALL.cuda()[torch.randperm(1 << 16, generator=g, device="cuda")]
+    x = logit_transform.bf16_of_key16(keys)
+    from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec.draft_shard_topk import (  # noqa: E501
+        radix_key16,
+    )
+
+    got = t.values[radix_key16(x.view(torch.int16)).long()]
+    assert logit_transform.same_bits(got, builtin_transform(config, x))
+
+
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_canon_groups_exactly_the_ties(name):
+    t = transform(CONFIGS[name])
+    canon, values = t.canon.cpu().long(), t.values.cpu()
+    radix = _radix32(values)
+    finite = torch.arange(0x007F, 0xFF81)
+    # Non-decreasing over the non-NaN keys, so the tie groups are runs of consecutive keys.
+    assert (radix[finite][1:] >= radix[finite][:-1]).all()
+    # Same canonical key exactly when the transformed values have the same bits.
+    assert torch.equal(radix[canon[finite]], radix[finite])
+    split = canon[finite][1:] != canon[finite][:-1]
+    assert torch.equal(split, radix[finite][1:] != radix[finite][:-1])
+    # The band: canon is the identity on [lo, hi] and on no larger interval around zero.
+    assert torch.equal(canon[t.lo : t.hi + 1], torch.arange(t.lo, t.hi + 1))
+    assert t.lo == 0 or canon[t.lo - 1] != t.lo - 1
+    assert t.hi == 0xFFFF or canon[t.hi + 1] != t.hi + 1
+    # NaN maps to itself; +0.0 and -0.0 stay apart and keep their sign.
+    assert int(canon[0xFFFF]) == 0xFFFF
+    assert int(canon[0x8000]) == 0x8000 and int(canon[0x7FFF]) == 0x7FFF
+    assert values[0x8000].item() == 0.0 and not torch.signbit(values[0x8000])
+    assert torch.signbit(values[0x7FFF]) and values[0x7FFF].item() == 0.0
+    assert logit_transform.unsupported_reason(t, 25256) is None
+
+
+def test_musespark_band_and_saturation_on_this_device():
+    """The identity band holds every logit of magnitude up to 512; +-inf saturate to +-softcap."""
+    t = transform(MUSESPARK)
+    lo_value = float(logit_transform.bf16_of_key16(torch.tensor([t.lo])).float())
+    hi_value = float(logit_transform.bf16_of_key16(torch.tensor([t.hi])).float())
+    groups = logit_transform.tie_groups(t)
+    first_up = min(g[0] for g in groups if g[0] >= 0x8000)
+    print(
+        f"identity band [{lo_value}, {hi_value}], first tie above zero at bf16 "
+        f"{float(logit_transform.bf16_of_key16(torch.tensor([first_up])).float())}, "
+        f"{len(groups)} tie groups"
+    )
+    assert lo_value <= -512.0 and hi_value >= 512.0
+    softcap = MUSESPARK[1]
+    assert t.values[0xFF80].item() == softcap and t.values[0x007F].item() == -softcap
+
+
+def test_identity_transform_builds_nothing():
+    assert logit_transform.build(drafter(1.0, None), torch.device("cuda")) is None
+
+
+def test_an_fp32_probe_would_not_have_found_the_transform():
+    """Why build probes in bf16: on fp32 the built-in transform runs in place and returns its input."""
+    x = torch.ones(4, device="cuda")
+    assert builtin_transform(MUSESPARK, x) is x
+    assert logit_transform.build(drafter(*MUSESPARK), torch.device("cuda")) is not None
+
+
+def test_ensure_transformed_applies_once():
+    t = transform(MUSESPARK)
+    x = torch.randn(3, 9, device="cuda").bfloat16()
+    y = logit_transform.ensure_transformed(t, x)
+    assert y.dtype == torch.float32 and logit_transform.same_bits(
+        y, builtin_transform(MUSESPARK, x)
+    )
+    assert logit_transform.ensure_transformed(t, y) is y
+    assert logit_transform.ensure_transformed(None, x) is x
+
+
+def test_the_tables_follow_the_drafter_attributes():
+    t = transform(MUSESPARK)
+    model = drafter(*MUSESPARK)
+    assert t.matches(model)
+    model.logit_softcap = 30.0
+    try:
+        assert not t.matches(model)
+    finally:
+        model.logit_softcap = MUSESPARK[1]
+
+
+def test_a_transform_that_decreases_is_refused():
+    t = transform(MUSESPARK)
+    bad = logit_transform.DraftLogitTransform(
+        drafter(*MUSESPARK), t.function, t.canon, -t.values, t.lo, t.hi
+    )
+    reason = logit_transform.unsupported_reason(bad, 25256)
+    assert reason is not None
+
+
+def test_another_method_is_refused():
+    t = transform(MUSESPARK)
+
+    def other(model, logits):
+        return logits.float()
+
+    bad = logit_transform.DraftLogitTransform(
+        drafter(*MUSESPARK), other, t.canon, t.values, t.lo, t.hi
+    )
+    assert "not DFlashForCausalLM's" in logit_transform.unsupported_reason(bad, 25256)
+
+
+def test_digest_tells_tables_apart():
+    a = logit_transform.digest(transform(MUSESPARK))
+    assert a == logit_transform.digest(transform(MUSESPARK))
+    assert a != logit_transform.digest(transform(OTHER_TRANSFORMS["softcap 30"]))
+    assert logit_transform.digest(None) != a
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py
new file mode 100644
index 000000000000..77c819887bde
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py
@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The built-in DFlash forward's statements the draft head's deferred logit transform relies on.
+
+The draft head (``kernels/draft_head/head.py``) hands the drafter's bf16 shard
+logits through the drafter's ``apply_draft_logit_transform`` untransformed and
+applies the transform after its top-k, in ``ShardedDraftHeadDFlashWorker``'s
+``_refine_block_logits`` and selector. That is the built-in computation only
+while nothing between the transform call and ``_refine_block_logits`` reads the
+logits, and only while ``_refine_block_logits`` runs before the selector and
+the sampler. This test reads ``DFlashWorker._forward_impl``'s source (no
+import, no device) and fails when that sequence changes, so a change to it is
+reviewed against the draft head instead of silently reaching a reader with
+untransformed logits.
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytestmark = pytest.mark.cpu_only
+
+#: The statements from the drafter's logits processor to the selector, as ``ast.unparse`` prints them.
+SEQUENCE = (
+    "gen_logits = draft_model.logits_processor(gen_hidden_states, draft_model.lm_head, "
+    "attn_metadata, True)",
+    "apply_logit_transform = getattr(draft_model, 'apply_draft_logit_transform', None)",
+    "if apply_logit_transform is not None:\n    gen_logits = apply_logit_transform(gen_logits)",
+    "vocab_size = gen_logits.shape[-1]",
+    "gen_logits = gen_logits.reshape(num_gens, K, vocab_size)",
+    "gen_logits = self._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata)",
+    "has_selector = getattr(draft_model, 'has_candidate_selector', False)",
+)
+
+
+def _source(relative: str) -> str:
+    spec = importlib.util.find_spec("tensorrt_llm")
+    assert spec is not None and spec.submodule_search_locations, "tensorrt_llm is not importable"
+    return (Path(next(iter(spec.submodule_search_locations))) / relative).read_text()
+
+
+def _forward_impl() -> ast.FunctionDef:
+    tree = ast.parse(_source("_torch/speculative/dflash.py"))
+    worker = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "DFlashWorker")
+    return next(
+        n for n in worker.body if isinstance(n, ast.FunctionDef) and n.name == "_forward_impl"
+    )
+
+
+def _bodies(node: ast.AST):
+    for child in ast.walk(node):
+        for name in ("body", "orelse"):
+            body = getattr(child, name, None)
+            if isinstance(body, list) and body and isinstance(body[0], ast.stmt):
+                yield body
+
+
+def test_nothing_reads_the_draft_logits_between_the_transform_and_refine():
+    forward = _forward_impl()
+    found = []
+    for body in _bodies(forward):
+        printed = [ast.unparse(s) for s in body]
+        if SEQUENCE[0] in printed:
+            start = printed.index(SEQUENCE[0])
+            found.append(tuple(printed[start : start + len(SEQUENCE)]))
+    assert len(found) == 1, (
+        "the drafter's logits processor call is not where the draft head expects it"
+    )
+    assert found[0] == SEQUENCE, (
+        "DFlashWorker._forward_impl changed between the drafter's logits processor and "
+        "_refine_block_logits; check the draft head's deferred transform against it:\n"
+        + "\n".join(found[0])
+    )
+
+
+def test_refine_runs_before_the_selector_and_the_samplers():
+    forward = _forward_impl()
+    calls = [
+        (node.lineno, node.func.attr)
+        for node in ast.walk(forward)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "self"
+        and node.func.attr
+        in (
+            "_refine_block_logits",
+            "_apply_dflash2_selector",
+            "_sample_dflash2_block",
+            "sample_draft_tokens",
+        )
+    ]
+    names = [name for _, name in sorted(calls)]
+    assert names.count("_refine_block_logits") == 1, names
+    assert names[0] == "_refine_block_logits", names
+
+
+def test_the_transform_has_one_caller():
+    """The deferring transform is consulted only by that call."""
+    root = Path(next(iter(importlib.util.find_spec("tensorrt_llm").submodule_search_locations)))
+    callers = sorted(
+        str(path.relative_to(root))
+        for path in (root / "_torch").rglob("*.py")
+        if "apply_draft_logit_transform" in path.read_text(encoding="utf-8", errors="ignore")
+        and "_experimental" not in path.parts
+    )
+    assert callers == ["_torch/models/modeling_dflash.py", "_torch/speculative/dflash.py"], callers
+
+
+#: This target's own sources, relative to the ``tensorrt_llm`` package.
+TARGET = (
+    "_torch/_experimental/modeling_v2/models/musespark_open/"
+    "muse_spark_open_816b_a42b_nvfp4__sm_103__tep"
+)
+
+
+def test_the_target_composes_no_other_reader_of_deferred_logits():
+    """Only the draft head and the fold take deferred logits; only the draft head wraps the processor.
+
+    A worker class of another fast path composed above the draft head that
+    defined ``_refine_block_logits`` or ``_apply_dflash2_selector`` would see the
+    deferred bf16 logits first (the draft head then stops deferring for it:
+    ``head._defers_transform``), and a fast path that replaced the drafter's
+    logits processor would end the deferral without a word. Read off the
+    target's sources, so a new one of either is reviewed against the draft head.
+    """
+    root = Path(next(iter(importlib.util.find_spec("tensorrt_llm").submodule_search_locations)))
+    owners = {"_refine_block_logits": [], "_apply_dflash2_selector": [], "processor": []}
+    for path in sorted((root / TARGET).rglob("*.py")):
+        source = path.read_text(encoding="utf-8", errors="ignore")
+        rel = str(path.relative_to(root / TARGET))
+        for node in ast.walk(ast.parse(source)):
+            if isinstance(node, ast.FunctionDef) and node.name in owners:
+                owners[node.name].append(rel)
+            if isinstance(node, ast.Assign) and any(
+                isinstance(target, ast.Attribute) and target.attr == "logits_processor"
+                for target in node.targets
+            ):
+                owners["processor"].append(rel)
+    assert owners["_refine_block_logits"] == ["kernels/draft_head/head.py"], owners
+    assert sorted(owners["_apply_dflash2_selector"]) == [
+        "kernels/draft_head/head.py",
+        "kernels/verify_fold/fold.py",
+    ], owners
+    assert owners["processor"] == ["kernels/draft_head/head.py"], owners
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py
new file mode 100644
index 000000000000..67920d43b88d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py
@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for ``draft_rank_topk`` (fast path ``wide_draft_topk``).
+
+Gated on its cells with ``torch.equal`` against its torch reference, with and
+without a drafter logit transform's canonical keys, and the merge over its
+output against the merge over ``draft_shard_topk``'s per-chunk keys.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import (  # noqa: E501
+    draft_rank_topk as rank_mod,
+)
+
+__extra_import_path__ = ["../../../../..", "../../kernels/draft_head"]
+from _draft_head_cases import (  # noqa: E402
+    MUSESPARK,
+    TOP_K,
+    VOCAB,
+    draft_head,
+    logits,
+    packed_keys,
+    rank_keys,
+    seed_of,
+    selector,
+    transform,
+)
+from _validating import certifying  # noqa: E402,F401
+
+assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device"
+
+draft_rank_topk = rank_mod.draft_rank_topk
+CELLS = rank_mod.DraftRankTopK.CELLS
+
+
+@pytest.mark.parametrize("kind", ["gemm3", "special_nan", "straddle_up", "saturated"])
+@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:48] for c in CELLS])
+def test_rank_topk_cells(cell, kind, certifying):  # noqa: F811
+    certifying(draft_rank_topk)
+    spec = cell.spec
+    x = logits(kind, spec["rows"], spec["shard"], seed_of(cell.why + kind))
+    kwargs = {} if spec["transform"] is None else draft_head.canon_kwargs(transform(MUSESPARK))
+    ranks = VOCAB // spec["shard"]
+    out = draft_rank_topk(x, TOP_K, 4096, 4, 1, ranks, **kwargs)
+    ref = draft_rank_topk.reference(x, TOP_K, 4096, 4, 1, ranks, **kwargs)
+    assert torch.equal(out, ref), f"{(out != ref).sum().item()} key halves differ"
+
+
+@pytest.mark.parametrize("kind", ["gemm400", "straddle_down", "ladder"])
+def test_merge_over_rank_keys_is_the_merge_over_chunk_keys(kind):
+    t = transform(MUSESPARK)
+    rows, ranks = 35, 8
+    full = logits(kind, rows, VOCAB, seed_of(kind))
+    sel = selector()
+    gate = torch.zeros(
+        1, rows, sel.predecessor_codebook.shape[1], dtype=torch.bfloat16, device="cuda"
+    )
+    anchor = torch.zeros(1, dtype=torch.int64, device="cuda")
+    outs = []
+    for keys in (packed_keys(full, ranks, t), rank_keys(full, ranks, t)):
+        outs.append(
+            draft_head.draft_topk_merge(
+                keys,
+                VOCAB // ranks,
+                1,
+                rows,
+                anchor,
+                gate,
+                sel.predecessor_codebook,
+                sel.successor_codebook,
+                TOP_K,
+                ranks,
+                values=t.values,
+            )
+        )
+    for a, b in zip(*outs):
+        assert torch.equal(
+            a.view(torch.int32) if a.is_floating_point() else a,
+            b.view(torch.int32) if b.is_floating_point() else b,
+        )
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py
new file mode 100644
index 000000000000..703fa48b68b8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py
@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for ``draft_shard_topk`` (fast path ``draft_head``).
+
+Gated on its cells with ``torch.equal`` against its torch reference, in both
+output layouts (the int32 keys the all-gather takes and the fp32 halves the
+fabric all-reduce takes), with and without a drafter logit transform's
+canonical keys.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import (  # noqa: E501
+    draft_shard_topk as shard_mod,
+)
+
+__extra_import_path__ = ["../../../../..", "../../kernels/draft_head"]
+from _draft_head_cases import (  # noqa: E402
+    MUSESPARK,
+    draft_head,
+    draft_shard_topk,
+    logits,
+    seed_of,
+    transform,
+)
+from _validating import certifying, validating  # noqa: E402,F401
+
+assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device"
+
+CELLS = shard_mod.DraftShardTopK.CELLS
+KINDS = ("gemm3", "special_nan", "straddle_up", "straddle_down", "ladder", "saturated")
+
+
+def _call_args(spec, x, rank=3, num_ranks=8):
+    kwargs = {}
+    if spec["halves"]:
+        kwargs.update(rank=rank, num_ranks=num_ranks)
+    if spec["transform"] is not None:
+        kwargs.update(draft_head.canon_kwargs(transform(MUSESPARK)))
+    return kwargs
+
+
+@pytest.mark.parametrize("kind", KINDS)
+@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:48] for c in CELLS])
+def test_shard_topk_cells(cell, kind, certifying):  # noqa: F811
+    certifying(draft_shard_topk)
+    spec = cell.spec
+    x = logits(
+        kind, spec["rows"], spec["shard"], seed_of(cell.why + kind), shard=spec["shard"] // 2
+    )
+    kwargs = _call_args(spec, x)
+    out = draft_shard_topk(x, **kwargs)
+    ref = draft_shard_topk.reference(x, **kwargs)
+    assert out.dtype == ref.dtype and out.shape == ref.shape
+    assert torch.equal(out, ref), f"{(out != ref).sum().item()} keys differ"
+
+
+def test_without_a_table_the_keys_are_the_plain_ones():
+    """``canon=None`` is the table-free kernel; a table changes only keys outside its band."""
+    t = transform(MUSESPARK)
+    x = logits("gemm100", 7, 25256, 11)
+    plain = draft_shard_topk(x)
+    assert torch.equal(plain, draft_shard_topk.reference(x))
+    inside = x.clamp(-500, 500)
+    assert torch.equal(
+        draft_shard_topk(inside, **draft_head.canon_kwargs(t)), draft_shard_topk(inside)
+    )
+    straddle = logits("straddle_up", 7, 25256, 12)
+    assert not torch.equal(
+        draft_shard_topk(straddle, **draft_head.canon_kwargs(t)), draft_shard_topk(straddle)
+    )
+
+
+def test_shard_topk_reads_a_strided_row_view():
+    t = transform(MUSESPARK)
+    full = logits("gemm400", 7, 202048, 5)
+    view = full[:, 3 * 25256 : 4 * 25256]
+    assert view.stride(0) == 202048
+    kwargs = draft_head.canon_kwargs(t)
+    assert torch.equal(
+        draft_shard_topk(view, **kwargs), draft_shard_topk(view.contiguous(), **kwargs)
+    )
+
+
+def test_shard_topk_guards_refuse_what_the_kernel_cannot_take():
+    t = transform(MUSESPARK)
+    x = torch.zeros(2, 1024, device="cuda", dtype=torch.bfloat16)
+    with validating(draft_shard_topk):
+        with pytest.raises(ValueError):  # wider than the 16-bit column field
+            draft_shard_topk(torch.zeros(2, 70000, device="cuda", dtype=torch.bfloat16))
+        with pytest.raises(ValueError):  # not bf16
+            draft_shard_topk(torch.zeros(2, 1024, device="cuda", dtype=torch.float32))
+        with pytest.raises(ValueError):  # a table that is not the identity on its band
+            draft_shard_topk(x, canon=t.canon, canon_band=(0, 0xFFFF))
+        with pytest.raises(ValueError):  # a table of the wrong dtype
+            draft_shard_topk(x, canon=t.canon.long(), canon_band=t.band)
+        with pytest.raises(ValueError):  # a table that does not keep the order
+            draft_shard_topk(x, canon=t.canon.flip(0).contiguous(), canon_band=(0x8000, 0x8000))
+        draft_shard_topk(x, **draft_head.canon_kwargs(t))
diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py
new file mode 100644
index 000000000000..fb562330f2b0
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for ``draft_topk_merge`` (fast path ``draft_head``) with and without a drafter logit transform.
+
+Gated with ``torch.equal`` on the ids and bitwise on the values against what
+the built-in path computes: ``torch.topk`` over the gathered fp32 logits, or,
+with the transform, over the drafter's transform of them (each shard
+transformed by the drafter's own method, then gathered). Then against the
+entry's own reference, and the einsum operands against ``dflash2_score_edges``.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import (  # noqa: E501
+    draft_topk_merge as merge_mod,
+)
+from tensorrt_llm._torch.models.modeling_dflash import dflash2_score_edges
+
+__extra_import_path__ = ["../../../../..", "../../kernels/draft_head"]
+from _draft_head_cases import (  # noqa: E402
+    FAMILIES,
+    MUSESPARK,
+    RANK,
+    TOP_K,
+    VOCAB,
+    assert_same_floats,
+    builtin_candidates,
+    draft_topk_merge,
+    layout_keys,
+    logits,
+    merge_launch,
+    seed_of,
+    selector,
+    split,
+    transform,
+)
+from _validating import certifying  # noqa: E402,F401
+
+assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device"
+
+CELLS = merge_mod.DraftTopKMerge.CELLS
+
+
+def _run(requests, positions, ranks, layout, kind, config, seed, canon=True):
+    """(merge outputs, built-in (unary, ids), gate, anchor) for one draft step."""
+    rows = requests * positions
+    t = None if config is None else transform(config)
+    full = logits(kind, rows, VOCAB, seed, shard=VOCAB // ranks)
+    g = torch.Generator(device="cuda").manual_seed(seed + 1)
+    gate = torch.randn(requests, positions, RANK, generator=g, device="cuda").bfloat16()
+    anchor = torch.randint(0, VOCAB, (requests,), generator=g, device="cuda")
+    keys = layout_keys(layout, full, ranks, t if canon else None)
+    sel = selector()
+    args = (
+        keys,
+        VOCAB // ranks,
+        requests,
+        positions,
+        anchor,
+        gate,
+        sel.predecessor_codebook,
+        sel.successor_codebook,
+        TOP_K,
+        ranks,
+    )
+    kwargs = dict(merge_launch(layout, rows))
+    if t is not None:
+        kwargs["values"] = t.values
+    out = draft_topk_merge(*args, **kwargs)
+    ref = builtin_candidates(split(full, ranks), requests, positions, config)
+    return out, ref, args, kwargs, gate, anchor
+
+
+@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:48] for c in CELLS])
+def test_topk_merge_cells(cell, certifying):  # noqa: F811
+    certifying(draft_topk_merge)
+    spec = cell.spec
+    config = None if spec["transform"] is None else MUSESPARK
+    requests, positions, ranks = spec["requests"], spec["positions"], spec["ranks"]
+    for kind in ("gemm3", "special_nan", "straddle_up", "saturated"):
+        (ids, unary, preds, succs), (ref_unary, ref_ids), args, kwargs, gate, anchor = _run(
+            requests, positions, ranks, spec["layout"], kind, config, seed_of(cell.why + kind)
+        )
+        assert torch.equal(ids, ref_ids), f"{kind}: candidate ids or their order differ"
+        assert_same_floats(unary, ref_unary)
+        r_ids, r_unary, r_preds, r_succs = draft_topk_merge.reference(*args, **kwargs)
+        assert torch.equal(ids, r_ids)
+        assert_same_floats(unary, r_unary)
+        assert torch.equal(preds, r_preds) and torch.equal(succs, r_succs)
+        sel = selector()
+        builtin = dflash2_score_edges(
+            sel.predecessor_codebook, sel.successor_codebook, ref_ids, ref_unary, gate, anchor
+        )
+        ours = unary.float()[:, :, None] + torch.einsum("blpr,blcr->blpc", preds, succs)
+        assert_same_floats(ours, builtin)
+
+
+@pytest.mark.parametrize("kind", FAMILIES)
+@pytest.mark.parametrize("layout", ["int32", "halves", "rank"])
+@pytest.mark.parametrize("ranks", [4, 8])
+@pytest.mark.parametrize("requests", [1, 2, 8, 24])
+def test_transformed_candidates_are_the_builtin_ones(requests, ranks, layout, kind):
+    (ids, unary, _, _), (ref_unary, ref_ids), *_ = _run(
+        requests, 7, ranks, layout, kind, MUSESPARK, seed_of(f"{requests}{ranks}{layout}{kind}")
+    )
+    differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist()
+    assert not differ, f"rows {differ} differ from torch.topk over the transformed logits"
+    assert_same_floats(unary, ref_unary)
+
+
+@pytest.mark.parametrize("kind", ["gemm3", "grid", "special_nan", "edges"])
+@pytest.mark.parametrize("layout", ["int32", "halves", "rank"])
+@pytest.mark.parametrize("requests", [8, 24])
+def test_untransformed_control(requests, layout, kind):
+    """The same matrix without a transform: a failure here is the merge's, not the transform's."""
+    (ids, unary, _, _), (ref_unary, ref_ids), *_ = _run(
+        requests, 7, 8, layout, kind, None, seed_of(f"control{requests}{layout}{kind}")
+    )
+    assert torch.equal(ids, ref_ids)
+    assert_same_floats(unary, ref_unary)
+
+
+@pytest.mark.parametrize("kind", ["straddle_up", "straddle_down", "saturated"])
+def test_without_canonical_keys_the_ties_go_wrong(kind):
+    """The negative control: bf16 keys order saturated logits by their bf16 values, not by id."""
+    (ids, _, _, _), (_, ref_ids), *_ = _run(1, 7, 8, "halves", kind, MUSESPARK, 3, canon=False)
+    assert not torch.equal(ids, ref_ids)