From 2d06973be6b32d59de5fc14d176fb5c44f76ad93 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 09:10:37 -0700
Subject: [PATCH 01/19] [None][feat] PyTorch executor: runtime hooks for
 model-owned decode work

Add RuntimeHooks (pyexecutor/runtime_hooks.py), a per-model seam through
which a modeling_v2 target can take over parts of the decode step without
changing the engine's behaviour for other models:

- step-input staging, so a target can write the next step's inputs from its
  own buffers;
- the speculative sampler's store writer;
- executor events and an execution-stream loop;
- a decode prologue that runs before the model forward.

The executor, CUDA graph runner, decoder runner, KV cache manager V2,
attention backends and the speculative modules read the hooks at their call
sites. Every hook is off unless a model installs it, and the default
behaviour is unchanged.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_torch/attention/backends/fmha/manager.py |   5 +-
 .../attention/backends/sparse/dsa/indexer.py  | 136 ++-
 .../attention/backends/sparse/dsa/metadata.py |  98 +-
 .../_torch/attention/backends/trtllm.py       |  61 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   8 +-
 .../_torch/pyexecutor/cuda_graph_runner.py    |  14 +-
 .../engine/runners/decoder/runner.py          | 189 +++-
 .../kv_cache/kv_cache_manager_v2.py           |  13 +
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  65 +-
 .../_torch/pyexecutor/runtime_hooks.py        | 216 +++++
 tensorrt_llm/_torch/speculative/eagle3.py     |   9 +-
 tensorrt_llm/_torch/speculative/interface.py  |  10 +-
 .../_torch/speculative/spec_sampler_base.py   |  82 +-
 .../attention/sparse/dsa/test_dsa_indexer.py  | 100 ++
 .../_torch/attention/test_fmha_manager.py     |  57 ++
 .../_torch/executor/test_gpu_keepalive.py     |   1 +
 .../_torch/executor/test_py_executor.py       |   3 +
 .../_torch/executor/test_runtime_hooks.py     | 918 ++++++++++++++++++
 18 files changed, 1876 insertions(+), 109 deletions(-)
 create mode 100644 tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
 create mode 100644 tests/unittest/_torch/executor/test_runtime_hooks.py

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py
index c9afff66e1d8..a4770ad3dbba 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py
@@ -308,7 +308,10 @@ def __init__(self, attn: TrtllmAttention) -> None:
         # normal forwarding never pays for an environment lookup.
         self._cache_sanity_check_enabled = _is_fmha_cache_sanity_check_enabled()
         self.fmha_libs: list[Fmha] = []
-        for fmha_cls in get_enabled_fmha_lib_classes():
+        for fmha_cls in (
+            *getattr(attn, "prepended_fmha_lib_classes", ()),
+            *get_enabled_fmha_lib_classes(),
+        ):
             if fmha_cls.is_available(attn):
                 self.fmha_libs.append(fmha_cls(attn))
 
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
index abe993e31438..61884b628af4 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
@@ -1050,6 +1050,18 @@ def recompute_slot_mappings(
         metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices
         metadata.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices
 
+        stager = getattr(metadata, "step_input_stager", None)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                metadata.slot_mapping_fp8[:total_new_kv_tokens],
+                metadata.host_slot_mapping_fp8[:total_new_kv_tokens],
+            )
+            stager.copy(
+                metadata.slot_mapping_scale[:total_new_kv_tokens],
+                metadata.host_slot_mapping_scale[:total_new_kv_tokens],
+            )
+            return
         metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_(
             metadata.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True
         )
@@ -1308,7 +1320,13 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None:
         if num_generations > 0:
             # Prepare schedule metadata for fp8_paged_mqa_logits
             # This is a preprocessing step that computes scheduling information for the kernel
-            Indexer.prepare_scheduler_metadata(metadata)
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager = getattr(metadata, "step_input_stager", None)
+            if stager is None:
+                Indexer.prepare_scheduler_metadata(metadata)
+            elif not stager.scheduler_metadata(metadata):
+                # Reads staged KV lengths: run right after the staged launch.
+                stager.defer(lambda: Indexer.prepare_scheduler_metadata(metadata))
 
     def _update_k_cache(
         self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata
@@ -1771,9 +1789,24 @@ def sparse_attn_indexer(
         )
 
         if has_decode and not metadata.skip_indexer_for_gen_reqs and reuse_topk:
-            topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = (
-                metadata.shared_topk_indices[:num_generations, :]
+            self._copy_reused_topk(
+                metadata,
+                topk_indices_buffer,
+                token_offset,
+                num_gen_tokens,
+                num_generations,
+                has_prefill,
+                gvr_prior_indices,
             )
+        elif (
+            has_decode
+            and not metadata.skip_indexer_for_gen_reqs
+            and self._fill_decode_topk(
+                metadata, q_fp8, weights, topk_indices_buffer, token_offset, num_gen_tokens
+            )
+        ):
+            # A subclass filled the generation rows.
+            pass
         elif has_decode and not metadata.skip_indexer_for_gen_reqs:
             # Get decode lengths per request (from seq_lens) for validation
             gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
@@ -2034,30 +2067,83 @@ def sparse_attn_indexer(
                 prev_topk_dst.copy_(last_mtp_topk)
 
         if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk:
-            rows = None
-            if has_decode:
-                next_n = num_gen_tokens // num_generations
-                gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
-                rows = self._mtp_last_accepted_rows(
-                    gen_topk, metadata, num_contexts, num_generations, next_n
-                )
-            if has_prefill:
-                ctx_last = (
-                    torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
-                )
-                ctx_rows = topk_indices_buffer[ctx_last]
-                rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
-            if rows is not None:
-                shared_topk_indices = metadata.shared_topk_indices
-                if shared_topk_indices is None:
-                    metadata.shared_topk_indices = rows.contiguous()
-                else:
-                    row_start = num_contexts if is_generation else 0
-                    shared_topk_indices[
-                        row_start : row_start + rows.shape[0], : rows.shape[1]
-                    ].copy_(rows)
+            self._store_mtp_shared_rows(
+                topk_indices_buffer,
+                metadata,
+                num_contexts,
+                num_generations,
+                num_gen_tokens,
+                token_offset,
+                has_prefill,
+                has_decode,
+                is_generation,
+            )
         return topk_indices_buffer
 
+    def _copy_reused_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+        num_generations: int,
+        has_prefill: bool,
+        gvr_prior_indices: Optional[torch.Tensor],
+    ) -> None:
+        """Copy the shared TopK rows an MTP draft step reuses into its generation rows."""
+        topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = (
+            metadata.shared_topk_indices[:num_generations, :]
+        )
+
+    def _fill_decode_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        q_fp8: torch.Tensor,
+        weights: torch.Tensor,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+    ) -> bool:
+        """Fill the generation TopK rows ahead of the decode branch; True if it filled them.
+
+        This class leaves them to the decode branch of ``sparse_attn_indexer``.
+        """
+        return False
+
+    def _store_mtp_shared_rows(
+        self,
+        topk_indices_buffer: torch.Tensor,
+        metadata: DSAtrtllmAttentionMetadata,
+        num_contexts: int,
+        num_generations: int,
+        num_gen_tokens: int,
+        token_offset: int,
+        has_prefill: bool,
+        has_decode: bool,
+        is_generation: Optional[bool],
+    ) -> None:
+        """Store each request's last accepted TopK row for the next MTP draft step to reuse."""
+        rows = None
+        if has_decode:
+            next_n = num_gen_tokens // num_generations
+            gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+            rows = self._mtp_last_accepted_rows(
+                gen_topk, metadata, num_contexts, num_generations, next_n
+            )
+        if has_prefill:
+            ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
+            ctx_rows = topk_indices_buffer[ctx_last]
+            rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
+        if rows is not None:
+            shared_topk_indices = metadata.shared_topk_indices
+            if shared_topk_indices is None:
+                metadata.shared_topk_indices = rows.contiguous()
+            else:
+                row_start = num_contexts if is_generation else 0
+                shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(
+                    rows
+                )
+
     def _mtp_last_accepted_rows(
         self,
         gen_topk: torch.Tensor,
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
index c6af4efd7dba..610a996f7e5c 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
@@ -1454,6 +1454,15 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale
         num_blocks = page_indices.shape[1]
         self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices)
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.indexer_k_cache_block_offsets[: self.num_seqs],
+                self.host_indexer_k_cache_block_offsets[: self.num_seqs],
+                clamp_min=0,
+            )
+            self.step_input_stager.clamp_min(self.indexer_k_cache_block_offsets[self.num_seqs :], 0)
+            return pool_indices
         self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
             self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
         )
@@ -1670,9 +1679,16 @@ def prepare_dense_topk_indices(self, kv_lens, device=False):  # device=False mea
                     kv_lens[self.num_contexts : self.num_seqs],
                     self.num_tokens - self.num_ctx_tokens,
                 )
-                self.topk_indices_buffer[gen_range, :].copy_(
-                    self.host_topk_indices_buffer[gen_range, :], non_blocking=True
-                )
+                if self.step_input_stager is not None:
+                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                    self.step_input_stager.copy(
+                        self.topk_indices_buffer[gen_range, :],
+                        self.host_topk_indices_buffer[gen_range, :],
+                    )
+                else:
+                    self.topk_indices_buffer[gen_range, :].copy_(
+                        self.host_topk_indices_buffer[gen_range, :], non_blocking=True
+                    )
 
     def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
         # The DeepGEMM paged-MQA kernel (fp8_paged_mqa_logits) runs a native
@@ -1753,9 +1769,16 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 num_tokens = self.num_generations * expand_factor
                 gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(expand_factor)
                 self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded)
-                self.kv_lens_expanded_cuda[:num_tokens].copy_(
-                    self.kv_lens_expanded_host[:num_tokens], non_blocking=True
-                )
+                if self.step_input_stager is not None:
+                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                    self.step_input_stager.copy(
+                        self.kv_lens_expanded_cuda[:num_tokens],
+                        self.kv_lens_expanded_host[:num_tokens],
+                    )
+                else:
+                    self.kv_lens_expanded_cuda[:num_tokens].copy_(
+                        self.kv_lens_expanded_host[:num_tokens], non_blocking=True
+                    )
                 self._refresh_expanded_block_table(expand_factor)
         else:
             self.dsl_expand_factor = 1
@@ -1984,6 +2007,15 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None):
         self.host_block_table_expanded[:num_tokens, :max_len].copy_(
             expanded_blocks, non_blocking=True
         )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.block_table_expanded[:num_tokens],
+                self.host_block_table_expanded[:num_tokens],
+                clamp_min=0,
+            )
+            self.step_input_stager.clamp_min(self.block_table_expanded[num_tokens:], 0)
+            return
         self.block_table_expanded[:num_tokens].copy_(
             self.host_block_table_expanded[:num_tokens], non_blocking=True
         )
@@ -2024,9 +2056,15 @@ def prepare_for_indexer_k_cache(self):
             pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1
         )
         # Copy to GPU
-        self.block_table[: self.num_seqs, :max_blocks_used].copy_(
-            host_block_table, non_blocking=True
-        )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.block_table[: self.num_seqs, :max_blocks_used], host_block_table
+            )
+        else:
+            self.block_table[: self.num_seqs, :max_blocks_used].copy_(
+                host_block_table, non_blocking=True
+            )
 
     def prepare_for_skip_indexer(self, kv_lens: torch.Tensor):
         num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens
@@ -2092,9 +2130,17 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_cached_token_indptr[1 : self.num_generations + 1],
             )
-            self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
-                self.host_gen_cached_token_indptr[: self.num_generations + 1], non_blocking=True
-            )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self.gen_cached_token_indptr[: self.num_generations + 1],
+                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
+                )
+            else:
+                self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
+                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
+                    non_blocking=True,
+                )
             # generation kv indptr
             torch.cumsum(
                 kv_lens[self.num_contexts : self.num_seqs],
@@ -2102,9 +2148,16 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_kv_indptr[1 : self.num_generations + 1],
             )
-            self.gen_kv_indptr[: self.num_generations + 1].copy_(
-                self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
-            )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self.gen_kv_indptr[: self.num_generations + 1],
+                    self.host_gen_kv_indptr[: self.num_generations + 1],
+                )
+            else:
+                self.gen_kv_indptr[: self.num_generations + 1].copy_(
+                    self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
+                )
         else:
             self.max_gen_seq_len = 0
 
@@ -2116,7 +2169,14 @@ def prepare_for_indices_conversion(self):
             self.seq_lens,
             dim=0,
         )
-        self.req_idx_per_token[: self.num_tokens].copy_(
-            self.host_req_idx_per_token[: self.num_tokens],
-            non_blocking=True,
-        )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.req_idx_per_token[: self.num_tokens],
+                self.host_req_idx_per_token[: self.num_tokens],
+            )
+        else:
+            self.req_idx_per_token[: self.num_tokens].copy_(
+                self.host_req_idx_per_token[: self.num_tokens],
+                non_blocking=True,
+            )
diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py
index 048138a48582..9a6634434d90 100644
--- a/tensorrt_llm/_torch/attention/backends/trtllm.py
+++ b/tensorrt_llm/_torch/attention/backends/trtllm.py
@@ -28,6 +28,7 @@
     from tensorrt_llm.mapping import Mapping
 
     from ...model_config import ModelConfig
+    from ...pyexecutor.runtime_hooks import StepInputStager
     from ...speculative.interface import SpecMetadata
     from ...speculative.spec_tree_manager import SpecTreeManager
     from .sparse.dsa.metadata import TokenMajorGenView
@@ -123,6 +124,8 @@ def generate_spec_decoding_packed_mask(max_num_requests: int,
 
 @dataclass(kw_only=True)
 class TrtllmAttentionMetadata(AttentionMetadata):
+    # A step-input stager while it records this metadata's writes, else None.
+    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     workspace: Optional[torch.Tensor] = None
     cuda_graph_workspace: Optional[torch.Tensor] = None
     workspace_reclaimable: bool = field(default=True, init=False)
@@ -997,8 +1000,14 @@ def prepare(self) -> None:
             device='cpu',
         )
         self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens)
-        self.prompt_lens_cuda[:self.num_seqs].copy_(
-            self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
+        stager = self.step_input_stager
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.prompt_lens_cuda[:self.num_seqs],
+                        self.prompt_lens_cpu[:self.num_seqs])
+        else:
+            self.prompt_lens_cuda[:self.num_seqs].copy_(
+                self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
 
         # number of tokens in the kv cache for each sequence in the batch
         cached_token_lens = torch.tensor(
@@ -1051,9 +1060,14 @@ def prepare(self) -> None:
         # the sequence length including the cached tokens and the input tokens.
         self.kv_lens[:self.num_seqs].copy_(
             kv_lens + self.kv_cache_params.num_extra_kv_tokens)
-        self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
-            kv_lens[:self.num_seqs]),
-                                                non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.kv_lens_cuda[:self.num_seqs],
+                        kv_lens[:self.num_seqs])
+        else:
+            self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
+                kv_lens[:self.num_seqs]),
+                                                    non_blocking=True)
         # total kv lens for context requests and generation requests, without extra tokens
         self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item()
         self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self.
@@ -1420,23 +1434,42 @@ def prepare_context_mla_with_cached_kv(self,
                      dtype=torch.int64,
                      out=self.host_ctx_cached_token_indptr[1:self.num_contexts +
                                                            1])
-        self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
-            non_blocking=True)
+        stager = self.step_input_stager
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.ctx_cached_token_indptr[:self.num_contexts + 1],
+                self.host_ctx_cached_token_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
+                non_blocking=True)
         torch.cumsum(
             self.seq_lens[:self.num_contexts],
             dim=0,
             dtype=torch.int64,
             out=self.host_ctx_uncached_token_indptr[1:self.num_contexts + 1])
-        self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
-            non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.ctx_uncached_token_indptr[:self.num_contexts + 1],
+                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
+                non_blocking=True)
         torch.cumsum(kv_lens[:self.num_contexts],
                      dim=0,
                      dtype=torch.int64,
                      out=self.host_ctx_kv_indptr[1:self.num_contexts + 1])
-        self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.ctx_kv_indptr[:self.num_contexts + 1],
+                        self.host_ctx_kv_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_kv_indptr[:self.num_contexts + 1],
+                non_blocking=True)
 
     def compute_max_num_custom_mask_tiles_kv_upper_bound(
             self, max_seq_len_kv, min_first_sparse_mask_offset_kv,
@@ -1678,6 +1711,8 @@ def is_sm_version_trtllm_gen_kernel(self, sm):
 class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]):
 
     Metadata = TrtllmAttentionMetadata
+    # FMHA library classes this instance tries before the registry's libraries.
+    prepended_fmha_lib_classes: tuple = ()
 
     @staticmethod
     def is_sm_version_trtllm_gen_kernel(sm):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index bcfb7529311b..26940ca8bb50 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -53,6 +53,7 @@
 from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens,
                            get_num_spec_layers, get_spec_decoder,
                            should_use_separate_draft_kv_cache)
+from ..speculative.spec_sampler_base import SpecSampler
 from ..utils import is_gdn_replay_enabled
 from . import config_utils
 from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer,
@@ -77,6 +78,7 @@
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
+from .runtime_hooks import runtime_hooks_of
 from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
                         KVCacheV2Scheduler, MultimodalEagerEncoderScheduler,
@@ -4451,7 +4453,11 @@ def instantiate_sampler(
     )
     if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder(
     ):
-        return get_spec_decoder(sampler_args, engine.spec_config)
+        sampler = get_spec_decoder(sampler_args, engine.spec_config)
+        store_writer = runtime_hooks_of(engine.model).sampler_store_writer
+        if store_writer is not None and isinstance(sampler, SpecSampler):
+            sampler.store_writer = store_writer
+        return sampler
 
     if mm_encoder_only:
         # NOTE: handle model outputs specially for mm encoder executor/engine.
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index 19ed990f33a1..73aa638f740d 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -162,6 +162,10 @@ class CUDAGraphRunner:
     multiple graphs, keyed by batch shape and execution-path specializations.
     """
     WARMUP_STEPS = 1
+    # Tokens of input_ids / position_ids that a step-input stager already
+    # wrote into the static inputs for the next replay (0: the replay copies
+    # them).
+    staged_static_input_tokens = 0
 
     def __init__(self, config: CUDAGraphRunnerConfig):
         self.config = config
@@ -816,7 +820,12 @@ def replay(
                 f"replay() got {seqlen} tokens for key {key}, but the graph "
                 f"was captured for {expected_num_tokens} tokens. A shorter "
                 "input_ids leaves the tail of the static input buffer stale.")
-        static_tensors["input_ids"][:seqlen].copy_(input_ids)
+        # A step-input stager already wrote input_ids / position_ids[:seqlen]
+        # into the static buffers.
+        staged = self.staged_static_input_tokens == seqlen
+        self.staged_static_input_tokens = 0
+        if not staged:
+            static_tensors["input_ids"][:seqlen].copy_(input_ids)
 
         position_ids = current_inputs["position_ids"]
         if self.config.use_mrope:
@@ -863,7 +872,8 @@ def replay(
                     f"for key {key}, but expected {expected_position_ids_shape}. "
                     "torch.Tensor.copy_() silently broadcasts mismatched shapes, "
                     "which would corrupt the static input buffer.")
-            static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
+            if not staged:
+                static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
 
         num_encoder_tokens = key.num_encoder_tokens
         if num_encoder_tokens:
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index c74a0a8d9e7e..986d05081590 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -54,6 +54,12 @@
     ResourceManager,
     ResourceManagerType,
 )
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import (
+    DEFAULT_RUNTIME_HOOKS,
+    RuntimeHooks,
+    StepInputStager,
+    runtime_hooks_of,
+)
 from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import warmup_sampling_module
 from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType
@@ -315,6 +321,11 @@ def __getattr__(self, name: str) -> Any:
 class DecoderRunner(ScheduledModelRunner):
     """Run scheduled decoder generation with its own buffers, graphs and metadata."""
 
+    # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise.
+    runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS
+    # The model's step-input stager while it records this runner's writes, else None.
+    step_input_stager: StepInputStager | None = None
+
     # Optional execution paths that a subclass can turn off.
     _context_warmups_supported = True
     _context_graph_promotion_supported = True
@@ -535,6 +546,7 @@ def _init_decoder_state(self) -> None:
             cuda_graph_manager=None,
         )
         self._prepare_inputs_event: torch.cuda.Event | None = None
+        self.runtime_hooks = runtime_hooks_of(self.model)
         # Let the first CUDA graph capture create its private pool. Piecewise
         # CUDA graphs use a separate pool owned by their runners, so sharing a
         # pre-created pool handle with the outer graph runner is unnecessary.
@@ -2856,6 +2868,16 @@ def _preprocess_inputs(
         Make some changes to the device inputs and avoid blocking the async data transfer
         """
         attn_meta = inputs.get("attn_metadata")
+        # A model's decode prologue may replace the device work below with one
+        # launch; it decides on the host before any device work and returns
+        # False to keep the code below.
+        prologue = self.runtime_hooks.decode_prologue
+        if prologue is not None and prologue(
+            self, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=runtime_draft_len
+        ):
+            if self.guided_decoder is not None:
+                self.guided_decoder.token_event.record()
+            return inputs
         # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step.
         if attn_meta is not None:
             attn_meta.on_update_kv_lens()
@@ -4076,21 +4098,51 @@ def previous_seq_slots_device():
             f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self._config.max_num_tokens})"  # noqa: E501
         )
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
+        stager = self.step_input_stager
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids, dtype=torch.int, pin_memory=prefer_pinned())
-            self.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(self.input_ids_cuda[:num_tokens], input_ids)
+            else:
+                self.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
 
         if num_draft_tokens > 0:
             draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned())
-            self.draft_tokens_cuda[: len(draft_tokens)].copy_(draft_tokens, non_blocking=True)
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(self.draft_tokens_cuda[: len(draft_tokens)], draft_tokens)
+            else:
+                self.draft_tokens_cuda[: len(draft_tokens)].copy_(draft_tokens, non_blocking=True)
         if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0:
             num_accepted_draft_tokens = torch.tensor(
                 num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()
             )
-            self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_(
-                num_accepted_draft_tokens, non_blocking=True
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(
+                    self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)],
+                    num_accepted_draft_tokens,
+                )
+            else:
+                self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_(
+                    num_accepted_draft_tokens, non_blocking=True
+                )
+        if next_draft_tokens_device is not None and stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.overlap_inputs(
+                self,
+                previous_batch_indices,
+                previous_pos_indices,
+                num_tokens,
+                num_draft_tokens,
+                self.get_runtime_tokens_per_gen_step(runtime_draft_len),
+                len(extend_requests) - len(extend_dummy_requests),
+                new_tokens_device,
+                next_draft_tokens_device,
+                new_tokens_lens_device,
             )
-        if next_draft_tokens_device is not None:
+        elif next_draft_tokens_device is not None:
             # Initialize these two values to zeros
             self.previous_pos_id_offsets_cuda *= 0
             self.previous_kv_lens_offsets_cuda *= 0
@@ -4188,6 +4240,10 @@ def previous_seq_slots_device():
         # cover the previous-batch rows; padding-request rows are covered by
         # neither, so without this the embedding would gather stale contents
         # of the uninitialized buffer, which may exceed the vocabulary size.
+        if stager is not None and unwritten_input_id_rows:
+            # A native write to input_ids_cuda inside the window needs a flush
+            # first; see the window banner in DecoderRunner._forward_decoder.
+            stager.flush()
         _zero_fill_row_spans(self.input_ids_cuda, unwritten_input_id_rows)
 
         position_ids = apply_position_id_offset(position_ids, model=self.model)
@@ -4229,11 +4285,21 @@ def previous_seq_slots_device():
                     segment[:, :, : end_idx - start_idx], non_blocking=True
                 )
             final_position_ids = self.mrope_position_ids_cuda[:, :, :total_num_tokens]
+        elif stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.position_ids(self, self.position_ids_cuda[:total_num_tokens], host_position_ids)
+            final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0)
         else:
             self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, non_blocking=True)
             final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0)
 
-        if enable_spec_decode:
+        if enable_spec_decode and stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.gather_ids_cuda[: len(gather_ids)],
+                torch.tensor(gather_ids, dtype=torch.int, pin_memory=prefer_pinned()),
+            )
+        elif enable_spec_decode:
             self.gather_ids_cuda[: len(gather_ids)].copy_(
                 torch.tensor(gather_ids, dtype=torch.int, pin_memory=prefer_pinned()),
                 non_blocking=True,
@@ -4423,7 +4489,13 @@ def previous_seq_slots_device():
         )
 
         virtual_num_tokens = total_num_tokens
-        if attn_metadata.padded_num_tokens is not None:
+        if attn_metadata.padded_num_tokens is not None and stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.zero(self.input_ids_cuda[total_num_tokens:padded_num_tokens])
+            stager.zero(self.position_ids_cuda[total_num_tokens:padded_num_tokens])
+            virtual_num_tokens = padded_num_tokens
+            final_position_ids = self.position_ids_cuda[:virtual_num_tokens].unsqueeze(0)
+        elif attn_metadata.padded_num_tokens is not None:
             self.input_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0)
             virtual_num_tokens = padded_num_tokens
             # Match the rank of the unpadded branch: MRoPE models always use
@@ -4856,29 +4928,69 @@ def _forward_decoder(
                     execution_requests.generation_requests,
                 )
 
-            with self._cuda_graph_metadata_prepare_scope(
-                attn_metadata,
-                key,
-                execution_requests,
-                execution_promoted_context_ids,
-                enable_spec_decode=enable_spec_decode,
-            ):
-                inputs, gather_ids, runtime_draft_len = self._prepare_inputs(
+            # ==== Step-input window opens here ====
+            # When the model declares a step-input stager and begin_if_eligible
+            # returns it, every device write that _prepare_inputs (and the
+            # metadata and KV-cache-manager code it calls) routes through
+            # ``step_input_stager`` (call sites marked "Staged write") is only
+            # recorded; stager.end() below, right before the graph replay,
+            # issues them all. Until then, do not read or write those buffers
+            # with any other device op: a read sees stale bytes and a write is
+            # overwritten. Route new writes through the stager or call
+            # stager.flush() first. Exception: input_ids_cuda and
+            # position_ids_cuda are also staged into the graph's static
+            # inputs, which the replay then does not refresh from them; to
+            # change them here, call stager.flush() and then write them
+            # natively. Nothing checks this contract in a production run.
+            # Contract: StepInputStager in pyexecutor/runtime_hooks.py.
+            stager = None
+            stager_source = self.runtime_hooks.step_input_stager
+            if stager_source is not None:
+                stager = stager_source.begin_if_eligible(
+                    self,
                     execution_requests,
-                    kv_cache_manager,
                     attn_metadata,
                     spec_metadata,
                     new_tensors_device,
-                    cache_indirection_buffer,
-                    resource_manager,
-                    can_run_graph,
-                    execution_promoted_context_ids,
-                    use_lora_graph=use_lora_graph,
-                    enable_spec_decode=enable_spec_decode,
-                    runtime_draft_len=runtime_draft_len,
+                    graph_replay=(
+                        can_run_graph
+                        and not self.cuda_graph_runner.needs_capture(key)
+                        and not self.cuda_graph_runner.is_warmup_only
+                    ),
                     is_dummy=is_dummy,
-                    gather_context_logits=gather_context_logits,
+                    enable_spec_decode=enable_spec_decode,
+                    promoted=bool(execution_promoted_context_ids),
                 )
+            try:
+                # The window spans the scope: the scope is a no-op whenever
+                # speculative decoding is on, which every stager requires.
+                with self._cuda_graph_metadata_prepare_scope(
+                    attn_metadata,
+                    key,
+                    execution_requests,
+                    execution_promoted_context_ids,
+                    enable_spec_decode=enable_spec_decode,
+                ):
+                    inputs, gather_ids, runtime_draft_len = self._prepare_inputs(
+                        execution_requests,
+                        kv_cache_manager,
+                        attn_metadata,
+                        spec_metadata,
+                        new_tensors_device,
+                        cache_indirection_buffer,
+                        resource_manager,
+                        can_run_graph,
+                        execution_promoted_context_ids,
+                        use_lora_graph=use_lora_graph,
+                        enable_spec_decode=enable_spec_decode,
+                        runtime_draft_len=runtime_draft_len,
+                        is_dummy=is_dummy,
+                        gather_context_logits=gather_context_logits,
+                    )
+            except BaseException:
+                if stager is not None:
+                    stager.discard()
+                raise
             if execution_promoted_context_ids:
                 self.iter_states["num_ctx_requests"] = scheduled_requests.num_context_requests
                 self.iter_states["num_ctx_tokens"] = sum(
@@ -4888,7 +5000,8 @@ def _forward_decoder(
                     scheduled_requests.num_generation_requests
                 )
             self._prepare_inputs_event = torch.cuda.Event()
-            self._prepare_inputs_event.record()
+            if stager is None:
+                self._prepare_inputs_event.record()
 
             breakable_runner = self.breakable_cuda_graph_runner
 
@@ -4976,10 +5089,30 @@ def capture_postprocess_fn(inputs: dict[str, Any]):
                         finally:
                             restore_attn_metadata_after_draft_replay(attn_metadata, saved_draft)
                     else:
-                        saved_draft = prepare_attn_metadata_for_draft_replay(
-                            attn_metadata, draft_kv_cache_manager
-                        )
                         try:
+                            saved_draft = prepare_attn_metadata_for_draft_replay(
+                                attn_metadata, draft_kv_cache_manager
+                            )
+                        except BaseException:
+                            if stager is not None:
+                                stager.discard()
+                            raise
+                        try:
+                            if stager is not None:
+                                # ==== Step-input window closes here ====
+                                # end() issues every staged write; nothing between
+                                # begin_if_eligible above and this call may have
+                                # touched the staged buffers with other device ops.
+                                # The input-copy event must follow the launch (host
+                                # page-table reads happen at execution time). A
+                                # nonzero static_inputs_tokens() makes the replay
+                                # skip its input_ids / position_ids copy.
+                                stager.end(record_event=False)
+                                self._prepare_inputs_event.record()
+                                stager.set_completion_event(self._prepare_inputs_event)
+                                self.cuda_graph_runner.staged_static_input_tokens = (
+                                    stager.static_inputs_tokens()
+                                )
                             with MoeLoadBalancerIterContext(moe_load_balancer):
                                 outputs = self.cuda_graph_runner.replay(
                                     key, inputs, position_id_offsets=position_id_offsets
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 5ac8934c6d85..f1f60ee84448 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
@@ -135,6 +135,7 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS
@@ -1231,6 +1232,8 @@ def _spec_recompute_claim_limit(
 class KVCacheManagerV2(BaseResourceManager):
     # Sparse managers attach a metadata batch during initialization.
     sparse_metadata_batch: Batch | None = None
+    # A step-input stager while it records this manager's writes, else None.
+    step_input_stager: Optional["StepInputStager"] = None
     # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction.
     # Declared on the class so it is present even when an instance is built without running __init__.
     _cold_pool_group_membership_cache: Optional[tuple[tuple[int, frozenset[int]], ...]] = None
@@ -6565,6 +6568,16 @@ def copy_batch_block_offsets(
             )
             return
 
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.block_offsets(
+                self.host_kv_cache_block_offsets,
+                dst_tensor,
+                copy_idx.tolist(),
+                self._index_scale_ints,
+                self.kv_offset.tolist(),
+            )
+            return
         copy_batch_block_offsets_to_device(
             self.host_kv_cache_block_offsets,
             dst_tensor,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index c4b14391194b..2bb83d55b49e 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -114,6 +114,7 @@
                             get_from_waiting_queue, merge_requests)
 from .resource_manager import (NoFreeSlotsError, ResourceManager,
                                ResourceManagerType, request_context)
+from .runtime_hooks import ExecutorEvents, runtime_hooks_of
 from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState,
                       SampleStateTensors)
 from .scheduler import (RequestScheduler, ScheduledRequests,
@@ -405,6 +406,16 @@ class PendingEncoderStep:
 
 
 class PyExecutor:
+    # True while the overlap loop runs with the execution stream current
+    # (_overlap_loop_stream_scope).
+    _loop_on_execution_stream = False
+    # Whether the model opted into running the overlap loop on the execution
+    # stream (RuntimeHooks.loop_on_execution_stream); read once at construction.
+    _execution_stream_loop_armed = False
+    # The model's executor events (RuntimeHooks.executor_events); read once at
+    # construction.
+    _executor_events: Optional[ExecutorEvents] = None
+
     # Minimum number of async micro batches for async PP execution.
     # This is a trade-off between memory usage and performance.
     # If the number of micro batches is too small, the executor will spend too much time in synchronization.
@@ -522,6 +533,9 @@ def __init__(
         # related modules
         self.resource_manager = resource_manager
         self.model_engine = model_engine
+        hooks = runtime_hooks_of(model_engine.model)
+        self._execution_stream_loop_armed = hooks.loop_on_execution_stream
+        self._executor_events = hooks.executor_events
         self._enable_adp_dummy_fixes = getattr(model_engine,
                                                "_enable_adp_dummy_fixes", False)
         self._enable_scheduler_aware_adp_dummy = getattr(
@@ -916,6 +930,10 @@ def __init__(
         # Ensure the default stream waits for execution_stream to complete
         # before subsequent operations.
         torch.cuda.current_stream().wait_stream(self.execution_stream)
+        if self._executor_events is not None:
+            # The model's actions after warmup and graph capture, on every rank
+            # (some are collective, so every rank reaches this point).
+            self._executor_events.after_warmup()
         self.encoder_launch_executor = (ThreadPoolExecutor(
             max_workers=1, thread_name_prefix="encoder-launch")
                                         if self.is_encoder_decoder else None)
@@ -3021,6 +3039,9 @@ def _executor_loop_cleanup(self):
                     f"Error waiting on PP send handles during cleanup: "
                     f"{traceback.format_exc()}")
 
+        if self._executor_events is not None and self._event_loop_completed:
+            self._executor_events.at_shutdown()
+
     def _pp_schedule_and_propagate(self, microbatch_id: int):
         """The first PP rank schedules the requests and propagates the result to all other PP ranks."""
 
@@ -5714,11 +5735,44 @@ def _wait_for_model_engine_input_copy(self):
         if wait_for_input_copy is not None:
             wait_for_input_copy()
 
+    @contextmanager
+    def _overlap_loop_stream_scope(self):
+        """Run the overlap loop with the execution stream current, when the model opted in.
+
+        _forward_step makes the execution stream wait for the current stream
+        before every forward and the current stream wait for the execution
+        stream after it. On the thread's default stream, which in steady decode
+        carries only these waits and timing events, step N+1's inputs wait for
+        decode graph N through two cross-stream event hops. With the execution
+        stream current the loop's stream work is queued in the same order on one
+        stream and _forward_step skips both waits.
+
+        Not taken with a two-model drafter, a guided decoder, a KV connector,
+        an encoder-decoder model or multimodal encoder item scheduling, whose
+        side streams synchronize with the current stream.
+        """
+        if not (self._execution_stream_loop_armed and self.drafter is None and
+                self.draft_model_engine is None and self.guided_decoder is None
+                and self.kv_connector_manager is None
+                and not self.is_encoder_decoder
+                and not self._mm_encoder_item_scheduling_enabled):
+            yield
+            return
+        with torch.cuda.stream(self.execution_stream):
+            self._loop_on_execution_stream = True
+            try:
+                yield
+            finally:
+                self._loop_on_execution_stream = False
+
     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))
-        with self._profiler() as profile_step, self.hang_detector:
+        if self._executor_events is not None:
+            self._executor_events.before_overlap_loop()
+        with self._overlap_loop_stream_scope(), self._profiler() as profile_step, \
+                self.hang_detector:
             iter_start_time = time.time()
             iter_stats = None
             target_inputs = None
@@ -8423,8 +8477,10 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
             cache_indirection_buffer = self.sampler.get_cache_indirection()
 
             # Run model forward on the execution stream for proper synchronization
-            # with KVCacheTransferManager's onboard/offload operations.
-            self.execution_stream.wait_stream(torch.cuda.current_stream())
+            # with KVCacheTransferManager's onboard/offload operations. With the
+            # loop on the execution stream both waits are program order.
+            if not self._loop_on_execution_stream:
+                self.execution_stream.wait_stream(torch.cuda.current_stream())
             self._attach_encoder_output_to_execution_stream(scheduled_requests)
             with torch.cuda.stream(self.execution_stream):
                 outputs = forward(scheduled_requests, self.resource_manager,
@@ -8438,7 +8494,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
 
             # Ensure the default stream waits for execution_stream to complete
             # before downstream operations use the outputs.
-            torch.cuda.current_stream().wait_stream(self.execution_stream)
+            if not self._loop_on_execution_stream:
+                torch.cuda.current_stream().wait_stream(self.execution_stream)
 
             self._kv_connector_wait_for_save()
             if route_capture is not None:
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
new file mode 100644
index 000000000000..f8c6717b3909
--- /dev/null
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -0,0 +1,216 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Step-loop behaviour a model may opt into, read through ``runtime_hooks_of``.
+
+A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class
+attribute or a property). Every field defaults to off, and a model that declares
+nothing gets ``RuntimeHooks()``: the step loop, the samplers and the attention
+metadata then run exactly as they do without this module. Every engine call site
+of a hook runs its original statements when the hook is off.
+
+The hooks change *where* and *how* the step loop moves data, never what a
+forward computes. The protocols below state only the methods an engine call
+site needs; the objects that implement them are owned by the model.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, Sequence
+
+if TYPE_CHECKING:
+    import torch
+
+
+class StepInputStager(Protocol):
+    """Records a decode step's input preparation and issues it as one launch.
+
+    ``begin_if_eligible`` opens a deferred-write window on a step it supports
+    and returns the stager, or returns None and the engine prepares the step
+    natively. While the window is open, the engine's call sites route their
+    device writes to the staging methods instead of issuing them. ``end``
+    issues every staged write in staging order, then the deferred native
+    operations; ``discard`` drops them after a failed preparation.
+
+    While its window is open the stager sets itself as ``step_input_stager`` on
+    the runner, the attention metadata, the speculative metadata and the KV
+    cache managers, and resets that attribute to None at ``end`` and
+    ``discard``. Every call site reads the attribute and runs its original
+    statement when it is None.
+    """
+
+    def begin_if_eligible(
+        self,
+        runner: Any,
+        requests: Any,
+        attn_metadata: Any,
+        spec_metadata: Any,
+        new_tensors_device: Any,
+        *,
+        graph_replay: bool,
+        is_dummy: bool,
+        enable_spec_decode: bool,
+        promoted: bool,
+    ) -> Optional["StepInputStager"]:
+        """Open the window for this step, or return None to keep the native preparation."""
+
+    def copy(
+        self, dst: "torch.Tensor", src: "torch.Tensor", clamp_min: Optional[int] = None
+    ) -> None:
+        """Stage ``dst.copy_(src, non_blocking=True)`` for a host ``src``."""
+
+    def zero(self, dst: "torch.Tensor") -> None:
+        """Stage ``dst.zero_()``."""
+
+    def clamp_min(self, dst: "torch.Tensor", lo: int) -> None:
+        """Stage ``dst.clamp_(min=lo)``."""
+
+    def block_offsets(
+        self,
+        host_table: "torch.Tensor",
+        dst: "torch.Tensor",
+        copy_index: Sequence[int],
+        index_scales: Sequence[int],
+        kv_offsets: Sequence[int],
+    ) -> None:
+        """Stage the KV cache manager's batch copy of block offsets to the device."""
+
+    def position_ids(
+        self, runner: Any, dst: "torch.Tensor", host_position_ids: "torch.Tensor"
+    ) -> None:
+        """Stage the position-id upload and the CUDA graph's static copy of it."""
+
+    def overlap_inputs(
+        self,
+        runner: Any,
+        previous_batch_indices: Sequence[int],
+        previous_pos_indices: Sequence[int],
+        num_tokens: int,
+        num_draft_tokens: int,
+        tokens_per_gen_step: int,
+        num_extend_requests: int,
+        new_tokens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+    ) -> None:
+        """Stage the overlap scheduler's gathers of the previous step's outputs."""
+
+    def scheduler_metadata(self, metadata: Any) -> bool:
+        """Stage the sparse-attention indexer's scheduler metadata.
+
+        Returns False, having staged nothing, when the caller must defer the
+        native computation instead.
+        """
+
+    def defer(self, fn: Callable[[], None]) -> None:
+        """Run ``fn`` (native device work that reads staged values) after the launch."""
+
+    def flush(self) -> None:
+        """Issue every write staged so far, before the window closes."""
+
+    def end(self, record_event: bool = True) -> None:
+        """Close the window and issue the staged preparation."""
+
+    def discard(self) -> None:
+        """Close the window after a failed preparation; the staged writes are dropped."""
+
+    def set_completion_event(self, event: "torch.cuda.Event") -> None:
+        """Event the caller recorded right after ``end(record_event=False)``."""
+
+    def static_inputs_tokens(self) -> int:
+        """Tokens of the CUDA graph's static inputs the last launch wrote (0: none)."""
+
+
+class SamplerStoreWriter(Protocol):
+    """Stores a speculative sampler step's outputs and copies the store to the host in one launch."""
+
+    def supported(
+        self,
+        store: Any,
+        new_tokens: "torch.Tensor",
+        next_new_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+    ) -> bool:
+        """Whether this step's output tensors fit the writer; False keeps the native stores."""
+
+    def store_and_copy_out(
+        self,
+        store: Any,
+        slots: Sequence[int],
+        new_tokens: "torch.Tensor",
+        next_new_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+        host_new_tokens: "torch.Tensor",
+        host_new_tokens_lens: "torch.Tensor",
+        host_next_draft_tokens: "torch.Tensor",
+        record_event: bool = True,
+    ) -> None:
+        """Scatter the outputs into the store rows ``slots`` and copy the store to the host tensors."""
+
+    def set_completion_event(self, event: "torch.cuda.Event") -> None:
+        """Event the caller recorded right after ``store_and_copy_out(record_event=False)``."""
+
+
+class DecodePrologue(Protocol):
+    """Replaces the decode step's input prologue with one launch.
+
+    Called first in the step's input preprocessing, inside CUDA graph capture
+    as well. It decides on the host before any device work; returning False
+    keeps the native prologue, and returning True means every engine-visible
+    tensor the native prologue writes holds the value it would have written.
+    """
+
+    def __call__(
+        self, runner: Any, inputs: Any, *, enable_spec_decode: bool, runtime_draft_len: int
+    ) -> bool: ...
+
+
+class ExecutorEvents(Protocol):
+    """Points in the executor's life a model may act at."""
+
+    def after_warmup(self) -> None:
+        """After warmup and CUDA graph capture, before the first request is served."""
+
+    def before_overlap_loop(self) -> None:
+        """Right before the overlap executor loop starts, on every rank; must not raise."""
+
+    def at_shutdown(self) -> None:
+        """When the executor loop has completed, during cleanup; must not raise."""
+
+
+@dataclass(frozen=True)
+class RuntimeHooks:
+    """The step-loop hooks a model opts into. All off by default."""
+
+    # Run the overlap executor loop with the execution stream current, so the
+    # loop's own stream work and the forward are ordered on one stream.
+    loop_on_execution_stream: bool = False
+    # Decode input preparation recorded and issued as one launch (see StepInputStager).
+    step_input_stager: Optional[StepInputStager] = None
+    # The speculative sampler's store update and host copy in one launch (see SamplerStoreWriter).
+    sampler_store_writer: Optional[SamplerStoreWriter] = None
+    # The decode step's input prologue in one launch (see DecodePrologue).
+    decode_prologue: Optional[DecodePrologue] = None
+    # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents).
+    executor_events: Optional[ExecutorEvents] = None
+
+
+DEFAULT_RUNTIME_HOOKS = RuntimeHooks()
+
+
+def runtime_hooks_of(model: object) -> RuntimeHooks:
+    """The hooks ``model`` declares, or the all-off default when it declares none.
+
+    Declaring hooks is optional for every model class, so absence is the common
+    case and means the default.
+    """
+    hooks = getattr(model, "runtime_hooks", None)
+    if hooks is None:
+        return DEFAULT_RUNTIME_HOOKS
+    if not isinstance(hooks, RuntimeHooks):
+        raise TypeError(
+            f"{type(model).__name__}.runtime_hooks must be a RuntimeHooks, got {type(hooks).__name__}"
+        )
+    return hooks
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index fc29a0633f41..72561a7642fe 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -352,8 +352,13 @@ def prepare(self):
                                      dtype=torch.int,
                                      device='cpu',
                                      pin_memory=prefer_pinned())
-        self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
-                                                 non_blocking=True)
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(self.batch_indices_cuda[:num_seqs],
+                                        batch_indices)
+        else:
+            self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
+                                                     non_blocking=True)
         # Snapshot the real (unpadded) scheduled token count before the decrement
         # below. maybe_capture_hidden_states runs during the target forward and
         # must bound the capture by this count; self.num_tokens is about to be
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 8237c09ee8b7..030e6fc9a82b 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import TYPE_CHECKING, List, Optional, Protocol, Type
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Protocol, Type
 
 import torch
 from packaging.version import Version
@@ -37,6 +37,7 @@
 if TYPE_CHECKING:
     from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
     from ..pyexecutor.llm_request import LlmRequest
+    from ..pyexecutor.runtime_hooks import StepInputStager
 
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer
@@ -472,6 +473,8 @@ class SpecMetadata:
     """
     Metadata for speculative decoding.
     """
+    # A step-input stager while it records this metadata's writes, else None.
+    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     # The max number of requests in a single batch.
     max_num_requests: int
     # The number of draft layers. (Also the number of draft tokens for the linear tree.)
@@ -894,6 +897,11 @@ def _populate_request_rng_state(
                                        device='cuda')
 
         def _upload(dst: torch.Tensor, values: list[int]) -> None:
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    dst[:len(values)], torch.tensor(values, dtype=torch.int64))
+                return
             dst[:len(values)].copy_(torch.tensor(values,
                                                  dtype=torch.int64,
                                                  pin_memory=prefer_pinned()),
diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
index 5fb5b185dc0b..dc40cca3a7c5 100644
--- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py
+++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
@@ -27,9 +27,11 @@
 
 import torch
 
+from ..._utils import prefer_pinned
 from ...sampling_params import SamplingParams
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from ..pyexecutor.resource_manager import BaseResourceManager
+from ..pyexecutor.runtime_hooks import SamplerStoreWriter
 from ..pyexecutor.sampler import (
     DEFAULT_BEAM_IDX,
     AsyncWorkerMixin,
@@ -167,6 +169,10 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin):
 
     SampleState = SampleStateSpec
 
+    # The model's one-launch store writer, bound once by instantiate_sampler; None keeps
+    # the native stores and host copies.
+    store_writer: Optional[SamplerStoreWriter] = None
+
     def is_generation_model(self) -> bool:
         return True
 
@@ -442,8 +448,11 @@ def sample_async(
             for r in sampling_requests
         ]
 
-        slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
-        slots = slots.to(device="cuda", non_blocking=True)
+        store_writer = self.store_writer
+        slots = None
+        if store_writer is None:
+            slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
+            slots = slots.to(device="cuda", non_blocking=True)
 
         o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests]
         o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests]
@@ -514,11 +523,28 @@ def sample_async(
         elif o_next_new_tokens.shape[1] > next_new_tokens_width:
             o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width]
 
-        # Use index_copy_ for efficient copying (slots are unique)
-        self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
-        self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
-        self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
-        self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
+        # The writer stores the four outputs and copies the three store tensors
+        # only; a step that publishes verify_lens keeps the native path.
+        staged_store = (
+            store_writer is not None
+            and o_verify_lens is None
+            and prefer_pinned()
+            and not self._async_worker_active()
+            and store_writer.supported(
+                self.store, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens
+            )
+        )
+        if not staged_store:
+            if slots is None:
+                slots = torch.as_tensor(
+                    [r.py_seq_slot for r in sampling_requests], dtype=torch.long
+                )
+                slots = slots.to(device="cuda", non_blocking=True)
+            # Use index_copy_ for efficient copying (slots are unique)
+            self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
+            self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
+            self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
+            self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
 
         # Create sample state with async D2H copy
         device_tensors = SampleStateTensorsSpec(
@@ -527,13 +553,43 @@ def sample_async(
             next_draft_tokens=self.store.next_draft_tokens,
         )
 
-        host_tensors = SampleStateTensorsSpec(
-            new_tokens=self._copy_to_host(self.store.new_tokens),
-            new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
-            next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
-            verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None),
-        )
+        if staged_store:
+            # staged_store implies prefer_pinned(), so these buffers are pinned.
+            host_tensors = SampleStateTensorsSpec(
+                new_tokens=torch.empty_like(
+                    self.store.new_tokens, device="cpu", pin_memory=prefer_pinned()
+                ),
+                new_tokens_lens=torch.empty_like(
+                    self.store.new_tokens_lens, device="cpu", pin_memory=prefer_pinned()
+                ),
+                next_draft_tokens=torch.empty_like(
+                    self.store.next_draft_tokens, device="cpu", pin_memory=prefer_pinned()
+                ),
+            )
+            store_writer.store_and_copy_out(
+                self.store,
+                [r.py_seq_slot for r in sampling_requests],
+                o_new_tokens,
+                o_next_new_tokens,
+                o_new_tokens_lens,
+                o_next_draft_tokens,
+                host_tensors.new_tokens,
+                host_tensors.new_tokens_lens,
+                host_tensors.next_draft_tokens,
+                record_event=False,
+            )
+        else:
+            host_tensors = SampleStateTensorsSpec(
+                new_tokens=self._copy_to_host(self.store.new_tokens),
+                new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
+                next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
+                verify_lens=(
+                    self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None
+                ),
+            )
         sampler_event = self._record_sampler_event()
+        if staged_store:
+            store_writer.set_completion_event(sampler_event.cuda_event)
 
         # Add dummy draft tokens to context requests for KV cache preparation
         for request in finished_context_requests:
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index 23d90f830a49..b73dc418ef74 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -4880,3 +4880,103 @@ def _run_indexer():
     assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, (
         "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena"
     )
+
+
+# --- The subclass seams of Indexer.sparse_attn_indexer ----------------------------------------
+
+
+def _seam_indexer():
+    indexer = Indexer.__new__(Indexer)
+    torch.nn.Module.__init__(indexer)
+    return indexer
+
+
+def test_reused_topk_copy_is_the_row_copy():
+    gen = torch.Generator().manual_seed(0)
+    shared = torch.randint(0, 1000, (4, 16), generator=gen, dtype=torch.int32)
+    buffer = torch.full((12, 16), -1, dtype=torch.int32)
+    expected = buffer.clone()
+    expected[2:5, :] = shared[:3, :]
+
+    _seam_indexer()._copy_reused_topk(
+        SimpleNamespace(shared_topk_indices=shared), buffer, 2, 3, 3, False, None
+    )
+
+    assert torch.equal(buffer, expected)
+
+
+def test_fill_decode_topk_leaves_the_rows_to_the_decode_branch():
+    buffer = torch.full((6, 8), -1, dtype=torch.int32)
+
+    filled = _seam_indexer()._fill_decode_topk(SimpleNamespace(), None, None, buffer, 0, 6)
+
+    assert filled is False
+    assert bool((buffer == -1).all())
+
+
+def _stored_rows_inline(
+    indexer,
+    topk_indices_buffer,
+    metadata,
+    num_contexts,
+    num_generations,
+    num_gen_tokens,
+    token_offset,
+    has_prefill,
+    has_decode,
+    is_generation,
+):
+    """The statements ``sparse_attn_indexer`` ran in place of ``_store_mtp_shared_rows``."""
+    rows = None
+    if has_decode:
+        next_n = num_gen_tokens // num_generations
+        gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+        rows = indexer._mtp_last_accepted_rows(
+            gen_topk, metadata, num_contexts, num_generations, next_n
+        )
+    if has_prefill:
+        ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
+        ctx_rows = topk_indices_buffer[ctx_last]
+        rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
+    if rows is not None:
+        shared_topk_indices = metadata.shared_topk_indices
+        if shared_topk_indices is None:
+            metadata.shared_topk_indices = rows.contiguous()
+        else:
+            row_start = num_contexts if is_generation else 0
+            shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(rows)
+
+
+@pytest.mark.parametrize("has_prefill, has_decode", [(False, True), (True, True), (True, False)])
+@pytest.mark.parametrize("is_generation", [True, False])
+@pytest.mark.parametrize("preset", [False, True], ids=["unset", "preset"])
+@pytest.mark.parametrize("accepted", [None, (0, 0, 1, 4, 2)], ids=["no_counts", "counts"])
+def test_mtp_shared_rows_are_the_inline_statements(
+    has_prefill, has_decode, is_generation, preset, accepted
+):
+    num_contexts, num_generations, next_n = 2, 3, 4
+    seq_lens = torch.tensor([5, 3] + [next_n] * num_generations, dtype=torch.int32)
+    token_offset = int(seq_lens[:num_contexts].sum())
+    num_gen_tokens = num_generations * next_n
+    gen = torch.Generator().manual_seed(1)
+    buffer = torch.randint(
+        0, 1000, (token_offset + num_gen_tokens, 6), generator=gen, dtype=torch.int32
+    )
+
+    def metadata():
+        return SimpleNamespace(
+            seq_lens_cuda=seq_lens.clone(),
+            mtp_num_accepted=None if accepted is None else torch.tensor(accepted),
+            shared_topk_indices=(torch.full((8, 8), -7, dtype=torch.int32) if preset else None),
+        )
+
+    args = (num_contexts, num_generations, num_gen_tokens, token_offset)
+    flags = (has_prefill, has_decode, is_generation)
+    expected, got = metadata(), metadata()
+    indexer = _seam_indexer()
+    _stored_rows_inline(indexer, buffer.clone(), expected, *args, *flags)
+    seam_buffer = buffer.clone()
+    indexer._store_mtp_shared_rows(seam_buffer, got, *args, *flags)
+
+    assert torch.equal(seam_buffer, buffer)
+    assert torch.equal(got.shared_topk_indices, expected.shared_topk_indices)
diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py
index 7cfb59eee093..397bfd93a4b5 100644
--- a/tests/unittest/_torch/attention/test_fmha_manager.py
+++ b/tests/unittest/_torch/attention/test_fmha_manager.py
@@ -974,3 +974,60 @@ def __init__(self, attn: TrtllmAttention) -> None:
     assert isinstance(attn._fmha_manager.fmha_libs[0], _NewFmha)
     assert quant_states_during_construction == [True]
     assert events == [("support", "old", None)]
+
+
+def _listed_fmha(name: str, available: bool, checked: list) -> type:
+    """An FMHA library class that records its availability check."""
+
+    def is_available(cls, attn) -> bool:
+        checked.append((name, attn))
+        return available
+
+    return type(
+        name,
+        (),
+        {"is_available": classmethod(is_available), "__init__": lambda self, attn: None},
+    )
+
+
+@pytest.mark.parametrize("declared", [False, True], ids=["undeclared", "empty"])
+def test_manager_without_prepended_libraries_builds_the_registry_list(declared: bool) -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)]
+    attn = FakeAttention()
+    if declared:
+        attn.prepended_fmha_lib_classes = ()
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == registry
+    assert TrtllmAttention.prepended_fmha_lib_classes == ()
+
+
+def test_manager_tries_a_prepended_library_first() -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)]
+    prepended = _listed_fmha("Prepended", True, checked)
+    attn = FakeAttention()
+    attn.prepended_fmha_lib_classes = (prepended,)
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == [prepended, *registry]
+    assert [name for name, _ in checked] == ["Prepended", "First", "Second"]
+    assert all(checked_attn is attn for _, checked_attn in checked)
+
+
+def test_manager_skips_an_unavailable_prepended_library() -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked)]
+    attn = FakeAttention()
+    attn.prepended_fmha_lib_classes = (_listed_fmha("Prepended", False, checked),)
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == registry
+    assert [name for name, _ in checked] == ["Prepended", "First"]
diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive.py b/tests/unittest/_torch/executor/test_gpu_keepalive.py
index 5c3318c00690..51d9b63b277e 100644
--- a/tests/unittest/_torch/executor/test_gpu_keepalive.py
+++ b/tests/unittest/_torch/executor/test_gpu_keepalive.py
@@ -315,6 +315,7 @@ class Stub:
     # notified first, then the keepalive is closed, then the PP handles are awaited.
     class CleanupStub:
         _executor_loop_cleanup = pe.PyExecutor._executor_loop_cleanup
+        _executor_events = None
 
         def __init__(self, keepalive):
             self.events = []
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index b571678d3c5c..c24923d0b935 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -102,6 +102,7 @@ def test_forward_step_carries_context_logits_request_to_runner(request_flags, ex
         iter_counter=0,
         sampler=Mock(get_cache_indirection=Mock(return_value=cache_indirection)),
         execution_stream=Mock(),
+        _loop_on_execution_stream=False,
         _iter_adp_dummy_ctx_tokens=0,
         _iter_adp_dummy_gen_tokens=0,
         _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens,
@@ -2243,6 +2244,8 @@ class _CleanupStub:
     """Stub for _executor_loop_cleanup: records the order in which the
     shutdown notification and PP-handle wait happen."""
 
+    _executor_events = None
+
     def __init__(self, pp_handles_raise=False):
         self.response_lock = threading.Lock()
         self.response_cv = threading.Condition(self.response_lock)
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
new file mode 100644
index 000000000000..1f3e63148340
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -0,0 +1,918 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).
+
+The engine's call sites run their original statements when a hook is off: every
+object a hook attaches to declares it off at class level, every call a call
+site makes on a hook is a method of the hook's protocol, every read of a hook in
+the engine files is decided by its all-off value toward the original statements
+(checked on the source), and call sites of each hook run on stand-ins issue
+their original work with the hook off and hand it to the hook when it is on.
+"""
+
+import ast
+import contextlib
+import dataclasses
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+import tensorrt_llm
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import (
+    DEFAULT_RUNTIME_HOOKS,
+    ExecutorEvents,
+    RuntimeHooks,
+    SamplerStoreWriter,
+    StepInputStager,
+    runtime_hooks_of,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+
+class _PlainModel(torch.nn.Module):
+    """A model that declares no hooks, like every model outside the opt-in."""
+
+
+class _ClassHooksModel(torch.nn.Module):
+    runtime_hooks = RuntimeHooks(loop_on_execution_stream=True)
+
+
+class _Events:
+    def after_warmup(self) -> None:
+        pass
+
+    def before_overlap_loop(self) -> None:
+        pass
+
+    def at_shutdown(self) -> None:
+        pass
+
+
+@pytest.mark.parametrize("model", [_PlainModel(), object()], ids=["module", "object"])
+def test_a_model_without_hooks_gets_the_all_off_default(model):
+    hooks = runtime_hooks_of(model)
+
+    assert hooks is DEFAULT_RUNTIME_HOOKS
+    assert hooks == RuntimeHooks()
+    assert not hooks.loop_on_execution_stream
+    assert hooks.step_input_stager is None
+    assert hooks.sampler_store_writer is None
+    assert hooks.decode_prologue is None
+    assert hooks.executor_events is None
+
+
+def test_every_field_defaults_to_off():
+    for field in dataclasses.fields(RuntimeHooks):
+        assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) in (False, None), field.name
+
+
+def test_hooks_declared_as_none_mean_the_default():
+    model = _PlainModel()
+    model.runtime_hooks = None
+
+    assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS
+
+
+def test_class_attribute_hooks_are_returned_as_declared():
+    assert runtime_hooks_of(_ClassHooksModel()) is _ClassHooksModel.runtime_hooks
+
+
+def test_instance_attribute_hooks_are_returned_as_declared():
+    hooks = RuntimeHooks(executor_events=_Events())
+    model = _PlainModel()
+    model.runtime_hooks = hooks
+
+    assert runtime_hooks_of(model) is hooks
+
+
+def test_instance_attribute_overrides_class_attribute():
+    hooks = RuntimeHooks()
+    model = _ClassHooksModel()
+    model.runtime_hooks = hooks
+
+    assert runtime_hooks_of(model) is hooks
+
+
+@pytest.mark.parametrize(
+    "declared",
+    [True, {"loop_on_execution_stream": True}, "RuntimeHooks()"],
+    ids=["bool", "dict", "str"],
+)
+def test_hooks_of_the_wrong_type_raise(declared):
+    model = _PlainModel()
+    model.runtime_hooks = declared
+
+    with pytest.raises(TypeError, match=r"_PlainModel\.runtime_hooks must be a RuntimeHooks"):
+        runtime_hooks_of(model)
+
+
+def test_hooks_are_immutable_once_declared():
+    hooks = RuntimeHooks()
+
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        hooks.loop_on_execution_stream = True
+
+
+def test_engine_classes_declare_the_stager_off():
+    """Every object a hook attaches to declares the attribute off at class level."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+    )
+    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+    from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata
+    from tensorrt_llm._torch.speculative.interface import SpecMetadata
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    for cls in (
+        TrtllmAttentionMetadata,
+        DSAtrtllmAttentionMetadata,
+        SpecMetadata,
+        Eagle3OneModelSpecMetadata,
+    ):
+        assert cls.step_input_stager is None, cls
+        # A class attribute, not a dataclass field: the constructors do not take it.
+        assert "step_input_stager" not in {f.name for f in dataclasses.fields(cls)}, cls
+    assert DecoderRunner.step_input_stager is None
+    assert DecoderRunner.runtime_hooks is DEFAULT_RUNTIME_HOOKS
+    assert KVCacheManagerV2.step_input_stager is None
+    assert CUDAGraphRunner.staged_static_input_tokens == 0
+    assert SpecSampler.store_writer is None
+    assert PyExecutor._execution_stream_loop_armed is False
+    assert PyExecutor._loop_on_execution_stream is False
+    assert PyExecutor._executor_events is None
+
+
+# Per hook protocol: the local names and attribute names a call site holds the hook in, and
+# the engine files with call sites.
+_HOOK_CONSUMERS = {
+    StepInputStager: (
+        ("stager", "stager_source"),
+        ("step_input_stager",),
+        (
+            "pyexecutor/engine/runners/decoder/runner.py",
+            "pyexecutor/kv_cache/kv_cache_manager_v2.py",
+            "attention/backends/trtllm.py",
+            "attention/backends/sparse/dsa/metadata.py",
+            "attention/backends/sparse/dsa/indexer.py",
+            "speculative/eagle3.py",
+            "speculative/interface.py",
+        ),
+    ),
+    SamplerStoreWriter: (
+        ("store_writer",),
+        ("store_writer",),
+        ("speculative/spec_sampler_base.py",),
+    ),
+    ExecutorEvents: ((), ("_executor_events",), ("pyexecutor/py_executor.py",)),
+}
+
+
+def _protocol_methods(protocol) -> set:
+    return {
+        name
+        for name, value in vars(protocol).items()
+        if inspect.isfunction(value) and not name.startswith("_")
+    }
+
+
+def _hook_calls(tree: ast.AST, local_names, attr_names):
+    """``(method, line)`` of every call on a hook held in one of the locals or attributes."""
+    for node in ast.walk(tree):
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
+            continue
+        receiver = node.func.value
+        if (isinstance(receiver, ast.Name) and receiver.id in local_names) or (
+            isinstance(receiver, ast.Attribute) and receiver.attr in attr_names
+        ):
+            yield node.func.attr, node.lineno
+
+
+@pytest.mark.parametrize("protocol", list(_HOOK_CONSUMERS), ids=lambda p: p.__name__)
+def test_every_hook_call_site_uses_its_protocol(protocol):
+    """A call site calls only its hook's protocol methods, and the protocol has no unused method."""
+    local_names, attr_names, consumers = _HOOK_CONSUMERS[protocol]
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    methods = _protocol_methods(protocol)
+    used, unknown = set(), []
+    for rel in consumers:
+        tree = ast.parse((root / rel).read_text())
+        calls = list(_hook_calls(tree, local_names, attr_names))
+        assert calls, f"{rel} has no {protocol.__name__} call site"
+        for method, line in calls:
+            used.add(method)
+            if method not in methods:
+                unknown.append(f"{rel}:{line}: {method}")
+    assert not unknown, f"calls outside the {protocol.__name__} protocol:\n" + "\n".join(unknown)
+    assert used == methods
+
+
+class _StoreWriterModel(torch.nn.Module):
+    def __init__(self, writer) -> None:
+        super().__init__()
+        self.runtime_hooks = RuntimeHooks(sampler_store_writer=writer)
+
+
+def _instantiate_sampler(monkeypatch, model, sampler):
+    """``instantiate_sampler`` for a speculative engine whose spec decoder is ``sampler``."""
+    from tensorrt_llm._torch.pyexecutor import _util
+
+    monkeypatch.setattr(_util, "confidential_compute_enabled", lambda: False)
+    monkeypatch.setattr(_util, "resolve_max_num_sequences", lambda *args, **kwargs: 8)
+    monkeypatch.setattr(_util, "create_torch_sampler_args", lambda **kwargs: "sampler args")
+    monkeypatch.setattr(_util, "get_spec_decoder", lambda args, spec_config: sampler)
+    engine = types.SimpleNamespace(
+        max_seq_len=64,
+        spec_config=types.SimpleNamespace(
+            spec_dec_mode=types.SimpleNamespace(has_spec_decoder=lambda: True)
+        ),
+        model=model,
+    )
+    llm_args = types.SimpleNamespace(
+        sampler_force_async_worker=False,
+        disable_overlap_scheduler=False,
+        enable_speculative_beam_history_d2h=False,
+    )
+    return _util.instantiate_sampler(
+        engine,
+        llm_args,
+        None,
+        max_batch_size=8,
+        max_beam_width=1,
+        mm_encoder_only=False,
+        speculative_config=None,
+    )
+
+
+def test_instantiate_sampler_binds_no_writer_without_hooks(monkeypatch):
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    sampler = SpecSampler.__new__(SpecSampler)
+
+    assert _instantiate_sampler(monkeypatch, _PlainModel(), sampler) is sampler
+    assert "store_writer" not in vars(sampler)
+    assert sampler.store_writer is None
+
+
+def test_instantiate_sampler_binds_the_declared_writer(monkeypatch):
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    writer = object()
+    sampler = SpecSampler.__new__(SpecSampler)
+
+    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(writer), sampler) is sampler
+    assert sampler.store_writer is writer
+
+
+def test_executor_events_protocol():
+    assert _protocol_methods(ExecutorEvents) == {
+        "after_warmup",
+        "before_overlap_loop",
+        "at_shutdown",
+    }
+    events = _Events()
+    assert RuntimeHooks(executor_events=events).executor_events is events
+
+
+class _StreamRecorder:
+    """Stands in for ``torch.cuda.stream``: records the streams it is entered with."""
+
+    def __init__(self) -> None:
+        self.entered = []
+
+    @contextlib.contextmanager
+    def __call__(self, stream):
+        self.entered.append(stream)
+        yield
+
+
+def _scope_executor(armed: bool, **overrides):
+    """A ``PyExecutor`` with only the attributes ``_overlap_loop_stream_scope`` reads."""
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor._execution_stream_loop_armed = armed
+    executor.execution_stream = object()
+    attrs = dict(
+        drafter=None,
+        draft_model_engine=None,
+        guided_decoder=None,
+        kv_connector_manager=None,
+        is_encoder_decoder=False,
+        _mm_encoder_item_scheduling_enabled=False,
+    )
+    attrs.update(overrides)
+    for name, value in attrs.items():
+        setattr(executor, name, value)
+    return executor
+
+
+def test_overlap_loop_scope_is_a_no_op_when_not_armed(monkeypatch):
+    executor = _scope_executor(armed=False)
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is False
+    assert recorder.entered == []
+    assert "_loop_on_execution_stream" not in vars(executor)
+
+
+def test_overlap_loop_scope_engages_when_armed(monkeypatch):
+    executor = _scope_executor(armed=True)
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is True
+    assert recorder.entered == [executor.execution_stream]
+    assert executor._loop_on_execution_stream is False
+
+
+def test_overlap_loop_scope_resets_the_flag_when_the_loop_raises(monkeypatch):
+    executor = _scope_executor(armed=True)
+    monkeypatch.setattr(torch.cuda, "stream", _StreamRecorder())
+
+    with pytest.raises(RuntimeError, match="loop failed"):
+        with executor._overlap_loop_stream_scope():
+            raise RuntimeError("loop failed")
+    assert executor._loop_on_execution_stream is False
+
+
+@pytest.mark.parametrize(
+    "name, value",
+    [
+        ("drafter", object()),
+        ("draft_model_engine", object()),
+        ("guided_decoder", object()),
+        ("kv_connector_manager", object()),
+        ("is_encoder_decoder", True),
+        ("_mm_encoder_item_scheduling_enabled", True),
+    ],
+)
+def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value):
+    executor = _scope_executor(armed=True, **{name: value})
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is False
+    assert recorder.entered == []
+
+
+def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
+    sampler = types.SimpleNamespace()
+
+    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(object()), sampler) is sampler
+    assert "store_writer" not in vars(sampler)
+
+
+# ---------------------------------------------------------------------------------------------
+# With every hook off, every call site runs its original statements.
+#
+# The check below reads the engine files that consume a hook and evaluates each read of a hook
+# at its all-off value. An ``if`` that reads a hook must be decided by that value; the branch it
+# takes is checked like any other code and the other branch is not run, so it may use the hook
+# freely. Outside such tests, live code may only bind a hook value to a local (or to an
+# attribute that declares the hook) at its all-off value. Any other read is reported with its
+# file and line.
+# ---------------------------------------------------------------------------------------------
+
+_SEAM_CONSUMERS = (
+    "pyexecutor/engine/runners/decoder/runner.py",
+    "pyexecutor/cuda_graph_runner.py",
+    "pyexecutor/kv_cache/kv_cache_manager_v2.py",
+    "pyexecutor/py_executor.py",
+    "pyexecutor/_util.py",
+    "attention/backends/trtllm.py",
+    "attention/backends/sparse/dsa/metadata.py",
+    "attention/backends/sparse/dsa/indexer.py",
+    "speculative/eagle3.py",
+    "speculative/interface.py",
+    "speculative/spec_sampler_base.py",
+)
+
+# Attributes that hold a hook, and their value when every hook is off.
+_SEAM_ATTRS = {
+    "step_input_stager": None,
+    "store_writer": None,
+    "_executor_events": None,
+    "_loop_on_execution_stream": False,
+    "_execution_stream_loop_armed": False,
+    "staged_static_input_tokens": 0,
+    "runtime_hooks": DEFAULT_RUNTIME_HOOKS,
+}
+# Locals the call sites hold a hook (or a decision on one) in.
+_SEAM_LOCALS = ("stager", "stager_source", "store_writer", "prologue", "staged", "staged_store")
+_HOOK_FIELDS = {field.name for field in dataclasses.fields(RuntimeHooks)}
+
+
+class _Unknown:
+    """A value the all-off evaluation cannot decide."""
+
+
+_UNKNOWN = _Unknown()
+
+
+class _AllOffWalk:
+    """Walks one engine file with every hook at its all-off value."""
+
+    def __init__(self, rel: str) -> None:
+        self.rel = rel
+        self.problems = []
+        self.decided = 0
+
+    def report(self, node: ast.AST, what: str) -> None:
+        self.problems.append(f"{self.rel}:{node.lineno}: {what}")
+
+    # -- reads of a hook ------------------------------------------------------------------
+    @staticmethod
+    def _is_seam_call(node: ast.AST) -> bool:
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)):
+            return False
+        if node.func.id == "runtime_hooks_of":
+            return True
+        return (
+            node.func.id == "getattr"
+            and len(node.args) >= 2
+            and isinstance(node.args[1], ast.Constant)
+            and node.args[1].value == "step_input_stager"
+        )
+
+    def reads_seam(self, node: ast.AST, env: dict) -> bool:
+        for sub in ast.walk(node):
+            if isinstance(sub, ast.Attribute) and sub.attr in _SEAM_ATTRS:
+                return True
+            if isinstance(sub, ast.Name) and sub.id in env:
+                return True
+            if self._is_seam_call(sub):
+                return True
+        return False
+
+    def value(self, node: ast.AST, env: dict):
+        """The all-off value of ``node``, or ``_UNKNOWN``."""
+        if isinstance(node, ast.Constant):
+            return node.value
+        if isinstance(node, ast.Name):
+            return env.get(node.id, _UNKNOWN)
+        if isinstance(node, ast.Attribute):
+            receiver = self.value(node.value, env)
+            if receiver is DEFAULT_RUNTIME_HOOKS and node.attr in _HOOK_FIELDS:
+                return getattr(DEFAULT_RUNTIME_HOOKS, node.attr)
+            if node.attr in _SEAM_ATTRS:
+                return _SEAM_ATTRS[node.attr]
+            return _UNKNOWN
+        if self._is_seam_call(node):
+            return DEFAULT_RUNTIME_HOOKS if node.func.id == "runtime_hooks_of" else None
+        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
+            operand = self.value(node.operand, env)
+            return _UNKNOWN if operand is _UNKNOWN else not operand
+        if isinstance(node, ast.BoolOp):
+            values = [self.value(v, env) for v in node.values]
+            known = [v for v in values if v is not _UNKNOWN]
+            if isinstance(node.op, ast.And):
+                if any(not v for v in known):
+                    return False
+                return True if len(known) == len(values) else _UNKNOWN
+            if any(known):
+                return True
+            return False if len(known) == len(values) else _UNKNOWN
+        if isinstance(node, ast.Compare) and len(node.ops) == 1:
+            left, right = self.value(node.left, env), self.value(node.comparators[0], env)
+            op = node.ops[0]
+            if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN:
+                if isinstance(node.comparators[0], ast.Constant) and right is None:
+                    return (left is None) == isinstance(op, ast.Is)
+            if (
+                isinstance(op, ast.Eq)
+                and isinstance(node.left, ast.Attribute)
+                and node.left.attr == "staged_static_input_tokens"
+            ):
+                # No tokens were staged; a graph replay has at least one token.
+                return False
+        return _UNKNOWN
+
+    # -- statements -------------------------------------------------------------------------
+    def block(self, body, env: dict) -> bool:
+        """Walk live statements; True when the block always ends in return or raise."""
+        for stmt in body:
+            if self.stmt(stmt, env):
+                return True
+        return False
+
+    def bind(self, stmt, targets, value, env: dict) -> None:
+        reads = self.reads_seam(value, env)
+        names = [t.id for t in targets if isinstance(t, ast.Name)]
+        attrs = [t for t in targets if isinstance(t, ast.Attribute) and t.attr in _SEAM_ATTRS]
+        if not (reads or attrs or any(name in env or name in _SEAM_LOCALS for name in names)):
+            return
+        if len(names) + len(attrs) != len(targets):
+            self.report(stmt, "a hook value is bound to something other than a local")
+            return
+        off = self.value(value, env)
+        if off is _UNKNOWN:
+            if reads or attrs:
+                self.report(stmt, "a hook value with every hook off is not decided")
+            for name in names:
+                env.pop(name, None)
+            return
+        for name in names:
+            env[name] = off
+        for attr in attrs:
+            if off != _SEAM_ATTRS[attr.attr]:
+                self.report(stmt, f"{attr.attr} is set to a value other than its all-off value")
+
+    def stmt(self, stmt: ast.stmt, env: dict) -> bool:
+        if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            self.block(stmt.body, dict(env))
+            return False
+        if isinstance(stmt, ast.ClassDef):
+            self.block(stmt.body, {})
+            return False
+        if isinstance(stmt, (ast.If, ast.While)):
+            if not self.reads_seam(stmt.test, env):
+                body_ends = self.block(stmt.body, env)
+                else_ends = self.block(stmt.orelse, env)
+                return isinstance(stmt, ast.If) and body_ends and else_ends
+            taken = self.value(stmt.test, env)
+            if taken is _UNKNOWN:
+                self.report(stmt, "a test of a hook is not decided with every hook off")
+                return False
+            self.decided += 1
+            if isinstance(stmt, ast.While) and taken:
+                self.report(stmt, "a loop on a hook runs with every hook off")
+                return False
+            return self.block(stmt.body if taken else stmt.orelse, env)
+        if isinstance(stmt, ast.Assign):
+            self.bind(stmt, stmt.targets, stmt.value, env)
+            return False
+        if isinstance(stmt, ast.AnnAssign):
+            if stmt.value is not None:
+                self.bind(stmt, [stmt.target], stmt.value, env)
+            return False
+        if isinstance(stmt, (ast.With, ast.AsyncWith)):
+            for item in stmt.items:
+                if self.reads_seam(item.context_expr, env):
+                    self.report(stmt, "a hook is entered as a context")
+            return self.block(stmt.body, env)
+        if isinstance(stmt, ast.Try):
+            ends = self.block(stmt.body, env)
+            for handler in stmt.handlers:
+                self.block(handler.body, env)
+            self.block(stmt.orelse, env)
+            return self.block(stmt.finalbody, env) or (ends and not stmt.handlers)
+        if isinstance(stmt, (ast.For, ast.AsyncFor)):
+            if self.reads_seam(stmt.iter, env):
+                self.report(stmt, "a loop iterates over a hook")
+            self.block(stmt.body, env)
+            self.block(stmt.orelse, env)
+            return False
+        if self.reads_seam(stmt, env):
+            self.report(stmt, "a hook is used outside a test decided with every hook off")
+        return isinstance(stmt, (ast.Return, ast.Raise))
+
+
+@pytest.mark.parametrize("rel", _SEAM_CONSUMERS)
+def test_every_seam_read_has_a_native_branch(rel):
+    """With every hook off, each call site takes the branch with its original statements."""
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    walk = _AllOffWalk(rel)
+    walk.block(ast.parse((root / rel).read_text()).body, {})
+    assert not walk.problems, "\n".join(walk.problems)
+    assert walk.decided, f"{rel} has no call site of a hook"
+
+
+def test_the_seam_check_reports_undecided_reads():
+    source = """
+def step(self, x):
+    stager = self.step_input_stager
+    if stager is not None:
+        stager.copy(x)
+    else:
+        x.copy_()
+    stager.flush()
+    if x or stager:
+        pass
+    self.step_input_stager = x
+"""
+    walk = _AllOffWalk("sample.py")
+    walk.block(ast.parse(source).body, {})
+    assert walk.decided == 1
+    assert walk.problems == [
+        "sample.py:8: a hook is used outside a test decided with every hook off",
+        "sample.py:9: a test of a hook is not decided with every hook off",
+        "sample.py:11: a hook value with every hook off is not decided",
+    ]
+
+
+def test_the_decode_prologue_call_site_uses_its_protocol():
+    """The runner calls the prologue with the protocol's positional and keyword parameters."""
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue
+
+    signature = inspect.signature(DecodePrologue.__call__)
+    positional = [
+        p.name
+        for p in signature.parameters.values()
+        if p.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD and p.name != "self"
+    ]
+    keywords = {
+        p.name for p in signature.parameters.values() if p.kind is inspect.Parameter.KEYWORD_ONLY
+    }
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    tree = ast.parse((root / "pyexecutor/engine/runners/decoder/runner.py").read_text())
+    calls = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == "prologue"
+    ]
+    assert len(calls) == 1
+    assert len(calls[0].args) == len(positional)
+    assert {keyword.arg for keyword in calls[0].keywords} == keywords
+
+
+# -- the decode prologue ------------------------------------------------------------------------
+
+
+class _AttnMetadata:
+    def __init__(self) -> None:
+        self.kv_lens_updates = 0
+
+    def on_update_kv_lens(self) -> None:
+        self.kv_lens_updates += 1
+
+
+def _prologue_runner(prologue):
+    """A ``DecoderRunner`` with only what ``_preprocess_inputs`` reads without spec decoding."""
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+
+    runner = DecoderRunner.__new__(DecoderRunner)
+    if prologue is not None:
+        runner.runtime_hooks = RuntimeHooks(decode_prologue=prologue)
+    runner.guided_decoder = None
+    runner._config = types.SimpleNamespace(disable_overlap_scheduler=True)
+    runner.mapping = types.SimpleNamespace(has_cp_helix=lambda: False)
+    return runner
+
+
+def test_preprocess_inputs_runs_the_native_prologue_without_a_hook():
+    runner = _prologue_runner(None)
+    inputs = {"attn_metadata": _AttnMetadata()}
+
+    assert runner.runtime_hooks is DEFAULT_RUNTIME_HOOKS
+    assert runner._preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=0) is (
+        inputs
+    )
+    assert inputs["attn_metadata"].kv_lens_updates == 1
+
+
+@pytest.mark.parametrize("engaged", [False, True], ids=["declined", "engaged"])
+def test_preprocess_inputs_keeps_the_native_prologue_when_the_hook_declines(engaged):
+    calls = []
+
+    def prologue(runner, inputs, *, enable_spec_decode, runtime_draft_len):
+        calls.append((runner, inputs, enable_spec_decode, runtime_draft_len))
+        return engaged
+
+    runner = _prologue_runner(prologue)
+    inputs = {"attn_metadata": _AttnMetadata()}
+
+    assert runner._preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=3) is (
+        inputs
+    )
+    assert calls == [(runner, inputs, False, 3)]
+    assert inputs["attn_metadata"].kv_lens_updates == (0 if engaged else 1)
+
+
+# -- the execution-stream loop --------------------------------------------------------------------
+
+
+class _Stream:
+    def __init__(self, name: str, log: list) -> None:
+        self.name, self.log = name, log
+
+    def wait_stream(self, other: "_Stream") -> None:
+        self.log.append(f"{self.name} waits for {other.name}")
+
+
+def _forward_step_executor(monkeypatch, log: list):
+    """A ``PyExecutor`` with only what ``_forward_step`` reads, logging its stream work."""
+    from tensorrt_llm._torch.pyexecutor import py_executor
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    current = _Stream("current", log)
+
+    @contextlib.contextmanager
+    def stream(s):
+        log.append(f"enter {s.name}")
+        yield
+        log.append(f"leave {s.name}")
+
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda: current)
+    monkeypatch.setattr(torch.cuda, "stream", stream)
+    monkeypatch.setattr(py_executor, "nvtx_range", lambda message: (lambda fn: fn))
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor.iter_counter = 0
+    executor.execution_stream = _Stream("execution", log)
+    executor.resource_manager = object()
+    executor.sampler = types.SimpleNamespace(get_cache_indirection=lambda: None)
+    executor.model_engine = types.SimpleNamespace(
+        forward=lambda *args, **kwargs: log.append("forward") or "outputs"
+    )
+    executor._iter_adp_dummy_ctx_tokens = 0
+    executor._iter_adp_dummy_gen_tokens = 0
+    executor._compute_adp_dummy_tokens = lambda requests: (0, 0)
+    executor._maybe_record_hang_diagnostic_phase = lambda *args, **kwargs: None
+    executor._attach_encoder_output_to_execution_stream = lambda requests: None
+    executor._mark_cross_kv_projection_consumed = lambda requests: None
+    executor._kv_connector_wait_for_save = lambda: None
+    return executor
+
+
+_NO_REQUESTS = types.SimpleNamespace(
+    context_requests=[], num_context_requests=0, num_generation_requests=0
+)
+
+
+def test_forward_step_orders_the_streams_off_the_execution_stream_loop(monkeypatch):
+    log = []
+    executor = _forward_step_executor(monkeypatch, log)
+
+    assert executor._forward_step(_NO_REQUESTS) == "outputs"
+    assert log == [
+        "execution waits for current",
+        "enter execution",
+        "forward",
+        "leave execution",
+        "current waits for execution",
+    ]
+
+
+def test_forward_step_skips_the_waits_on_the_execution_stream_loop(monkeypatch):
+    log = []
+    executor = _forward_step_executor(monkeypatch, log)
+    executor._loop_on_execution_stream = True
+
+    assert executor._forward_step(_NO_REQUESTS) == "outputs"
+    assert log == ["enter execution", "forward", "leave execution"]
+
+
+# -- executor events --------------------------------------------------------------------------
+
+
+def _cleanup_executor(events, completed: bool):
+    """A ``PyExecutor`` with only what ``_executor_loop_cleanup`` reads, without pipeline stages."""
+    import threading
+
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor.response_cv = threading.Condition()
+    executor.shutdown_event = threading.Event()
+    executor.num_micro_batches = 0
+    executor._event_loop_completed = completed
+    if events is not None:
+        executor._executor_events = events
+    return executor
+
+
+class _RecordingEvents(_Events):
+    def __init__(self) -> None:
+        self.calls = []
+
+    def at_shutdown(self) -> None:
+        self.calls.append("at_shutdown")
+
+
+@pytest.mark.parametrize("completed", [False, True])
+def test_executor_loop_cleanup_without_events(completed):
+    executor = _cleanup_executor(None, completed)
+
+    executor._executor_loop_cleanup()
+    assert executor.is_shutdown and executor.shutdown_event.is_set()
+
+
+@pytest.mark.parametrize("completed", [False, True])
+def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(completed):
+    events = _RecordingEvents()
+    executor = _cleanup_executor(events, completed)
+
+    executor._executor_loop_cleanup()
+    assert executor.is_shutdown
+    assert events.calls == (["at_shutdown"] if completed else [])
+
+
+# -- the step-input stager ------------------------------------------------------------------
+
+
+class _RecordingStager:
+    def __init__(self) -> None:
+        self.calls = []
+
+    def __getattr__(self, name):
+        return lambda *args, **kwargs: self.calls.append((name, args, kwargs))
+
+
+def _block_offsets_manager(monkeypatch, log: list):
+    """A ``KVCacheManagerV2`` with only what ``copy_batch_block_offsets`` reads (one page table)."""
+    from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+
+    monkeypatch.setattr(
+        kv_cache_manager_v2,
+        "copy_batch_block_offsets_to_device",
+        lambda *args: log.append(("native", args)),
+    )
+    manager = KVCacheManagerV2.__new__(KVCacheManagerV2)
+    manager._publish_sparse_metadata = lambda: None
+    manager.max_copy_beam_width = 1
+    manager.kv_cache_type = None
+    manager._use_per_layer_page_tables = False
+    manager.index_mapper = types.SimpleNamespace(
+        get_copy_index=lambda request_ids, num_contexts, beam_width, cross: torch.tensor([1, 0])
+    )
+    manager.host_kv_cache_block_offsets = torch.zeros(2, 4, dtype=torch.int32)
+    manager.index_scales = torch.ones(1, dtype=torch.int32)
+    manager._index_scale_ints = [1]
+    manager.kv_offset = torch.zeros(1, dtype=torch.int32)
+    manager._stream = types.SimpleNamespace(cuda_stream=7)
+    return manager
+
+
+def test_kv_cache_block_offsets_copy_natively_without_a_stager(monkeypatch):
+    log = []
+    manager = _block_offsets_manager(monkeypatch, log)
+    dst = torch.empty(2, 4, dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
+    assert len(log) == 1
+    kind, args = log[0]
+    assert kind == "native"
+    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
+    assert args[2].tolist() == [1, 0]
+    assert args[3] is manager.index_scales and args[4] is manager.kv_offset and args[5] == 7
+
+
+def test_kv_cache_block_offsets_are_staged_with_a_stager(monkeypatch):
+    log = []
+    manager = _block_offsets_manager(monkeypatch, log)
+    manager.step_input_stager = _RecordingStager()
+    dst = torch.empty(2, 4, dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
+    assert log == []
+    ((name, args, _),) = manager.step_input_stager.calls
+    assert name == "block_offsets"
+    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
+    assert args[2:] == ([1, 0], [1], [0])
+
+
+def _slot_mapping_inputs(tokens: int):
+    metadata = types.SimpleNamespace(
+        host_indexer_k_cache_block_offsets=torch.tensor([[0, 1]], dtype=torch.int32),
+        host_slot_mapping_fp8=torch.zeros(tokens, dtype=torch.int64),
+        host_slot_mapping_scale=torch.zeros(tokens, dtype=torch.int64),
+        slot_mapping_fp8=torch.full((tokens,), -1, dtype=torch.int64),
+        slot_mapping_scale=torch.full((tokens,), -1, dtype=torch.int64),
+    )
+    params = types.SimpleNamespace(
+        batch_size=1,
+        tokens_per_block=4,
+        head_dim=8,
+        new_kv_tokens=torch.tensor([tokens]),
+        data_bytes_per_token=8,
+        cached_kv_tokens=torch.tensor([1]),
+        quant_block_size=8,
+    )
+    return metadata, params
+
+
+def test_indexer_slot_mappings_copy_natively_without_a_stager():
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+    metadata, params = _slot_mapping_inputs(3)
+
+    Indexer.recompute_slot_mappings(metadata, params)
+    assert torch.equal(metadata.slot_mapping_fp8, metadata.host_slot_mapping_fp8)
+    assert torch.equal(metadata.slot_mapping_scale, metadata.host_slot_mapping_scale)
+
+
+def test_indexer_slot_mappings_are_staged_with_a_stager():
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+    metadata, params = _slot_mapping_inputs(3)
+    metadata.step_input_stager = _RecordingStager()
+
+    Indexer.recompute_slot_mappings(metadata, params)
+    assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
+    assert bool((metadata.slot_mapping_fp8 == -1).all())
+    assert bool((metadata.slot_mapping_scale == -1).all())

From 23324b77981d1b88d4e87657b77815cb82974217 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 09:10:38 -0700
Subject: [PATCH 02/19] [None][feat] GLM-MoE-DSA: CUDA kernels and torch ops
 for GLM-5.2 decode

Add model-specific CUDA kernels under cpp/tensorrt_llm/kernels/glmMoeDsa/
and their torch ops under cpp/tensorrt_llm/thop/glmMoeDsa/, one
<stem>Op.cpp per kernel <stem>.cu, registered as
trtllm::glm_moe_dsa_<what>. They cover the decode path of GLM-5.2: the
post-attention hop and its routing, the RoPE KV append, the sparse MLA
decode, the cross-rank indexer top-k, the dense MLP, and the draft model's
routed MoE, live-row hop and tail.

The kernels build as one OBJECT library for the SM100 family and are used
only by the GLM-5.2 modeling_v2 targets.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 cpp/tensorrt_llm/CMakeLists.txt               |    1 +
 cpp/tensorrt_llm/kernels/CMakeLists.txt       |    3 +
 .../kernels/glmMoeDsa/CMakeLists.txt          |   42 +
 .../kernels/glmMoeDsa/crossRankTopk.cu        | 1557 ++++++++++
 .../kernels/glmMoeDsa/crossRankTopk.h         |  162 +
 .../glmMoeDsa/crossRankTopkClusterMerge.cu    | 2420 +++++++++++++++
 .../glmMoeDsa/crossRankTopkClusterMerge.h     |  191 ++
 .../glmMoeDsa/crossRankTopkGlobalEmit.cu      | 2757 +++++++++++++++++
 .../glmMoeDsa/crossRankTopkGlobalEmit.h       |  252 ++
 .../kernels/glmMoeDsa/denseMlp.cu             |  453 +++
 cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h |   62 +
 .../kernels/glmMoeDsa/draftLiveRowHop.cu      |  586 ++++
 .../kernels/glmMoeDsa/draftLiveRowHop.h       |   88 +
 .../kernels/glmMoeDsa/draftRoutedMoe.cu       | 2628 ++++++++++++++++
 .../kernels/glmMoeDsa/draftRoutedMoe.h        |  205 ++
 .../kernels/glmMoeDsa/draftTail.cu            | 1568 ++++++++++
 .../kernels/glmMoeDsa/draftTail.h             |  144 +
 .../glmMoeDsa/indexerLayerRopeKvAppend.cu     |  439 +++
 .../glmMoeDsa/indexerLayerRopeKvAppend.h      |   47 +
 .../kernels/glmMoeDsa/postAttentionHop.cu     |  680 ++++
 .../kernels/glmMoeDsa/postAttentionHop.h      |   88 +
 .../glmMoeDsa/postAttentionHopFromPartials.cu |  707 +++++
 .../glmMoeDsa/postAttentionHopFromPartials.h  |   92 +
 .../glmMoeDsa/postAttentionHopRouting.cu      |  889 ++++++
 .../glmMoeDsa/postAttentionHopRouting.h       |  116 +
 .../kernels/glmMoeDsa/ropeKvAppend.cu         |  515 +++
 .../kernels/glmMoeDsa/ropeKvAppend.h          |  136 +
 .../kernels/glmMoeDsa/sparseMlaDecode.cu      |  714 +++++
 .../kernels/glmMoeDsa/sparseMlaDecode.h       |  108 +
 cpp/tensorrt_llm/thop/CMakeLists.txt          |   10 +-
 .../thop/glmMoeDsa/CMakeLists.txt             |   19 +
 .../glmMoeDsa/crossRankTopkClusterMergeOp.cpp |  108 +
 .../glmMoeDsa/crossRankTopkGlobalEmitOp.cpp   |  187 ++
 .../thop/glmMoeDsa/crossRankTopkOp.cpp        |  104 +
 .../thop/glmMoeDsa/denseMlpOp.cpp             |  131 +
 .../thop/glmMoeDsa/draftLiveRowHopOp.cpp      |  150 +
 .../thop/glmMoeDsa/draftRoutedMoeOp.cpp       |  246 ++
 .../thop/glmMoeDsa/draftTailOp.cpp            |  168 +
 .../glmMoeDsa/indexerLayerRopeKvAppendOp.cpp  |  242 ++
 .../postAttentionHopFromPartialsOp.cpp        |  150 +
 .../thop/glmMoeDsa/postAttentionHopOp.cpp     |  170 +
 .../glmMoeDsa/postAttentionHopRoutingOp.cpp   |  189 ++
 .../thop/glmMoeDsa/ropeKvAppendOp.cpp         |  330 ++
 .../thop/glmMoeDsa/sparseMlaDecodeOp.cpp      |  121 +
 44 files changed, 19974 insertions(+), 1 deletion(-)
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu
 create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp

diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt
index 25d1d5784c91..b7d4594f53fe 100644
--- a/cpp/tensorrt_llm/CMakeLists.txt
+++ b/cpp/tensorrt_llm/CMakeLists.txt
@@ -194,6 +194,7 @@ set(TRTLLM_LINK_LIBS
     mamba2_mtp_ssm_cache_src
     kda_decode_src
     kimi_k3_attn_res_src
+    glm_moe_dsa_kernels_src
     ws_layernorm_src
     fusedGatedRMSNormQuant_src
     fpA_intB_gemm_src
diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt
index 0e7fc9390dc0..7b386929da95 100644
--- a/cpp/tensorrt_llm/kernels/CMakeLists.txt
+++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt
@@ -32,6 +32,7 @@ add_subdirectory(fusedGatedRMSNormQuant)
 add_subdirectory(mamba2MTPSSMCache)
 add_subdirectory(kdaDecode)
 add_subdirectory(kimiK3AttnRes)
+add_subdirectory(glmMoeDsa)
 add_subdirectory(mhcKernels)
 add_subdirectory(compressorKernels)
 add_subdirectory(fusedSampling)
@@ -68,6 +69,8 @@ list(FILTER SRC_CPP EXCLUDE REGEX "kdaDecode/.*")
 list(FILTER SRC_CU EXCLUDE REGEX "kdaDecode/.*")
 list(FILTER SRC_CPP EXCLUDE REGEX "kimiK3AttnRes/.*")
 list(FILTER SRC_CU EXCLUDE REGEX "kimiK3AttnRes/.*")
+list(FILTER SRC_CPP EXCLUDE REGEX "glmMoeDsa/.*")
+list(FILTER SRC_CU EXCLUDE REGEX "glmMoeDsa/.*")
 list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*")
 list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*")
 list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*")
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt b/cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt
new file mode 100644
index 000000000000..db750686e8db
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt
@@ -0,0 +1,42 @@
+# 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.
+
+add_library(
+  glm_moe_dsa_kernels_src OBJECT
+  crossRankTopk.cu
+  crossRankTopkClusterMerge.cu
+  crossRankTopkGlobalEmit.cu
+  denseMlp.cu
+  draftLiveRowHop.cu
+  draftRoutedMoe.cu
+  draftTail.cu
+  indexerLayerRopeKvAppend.cu
+  postAttentionHop.cu
+  postAttentionHopFromPartials.cu
+  postAttentionHopRouting.cu
+  ropeKvAppend.cu
+  sparseMlaDecode.cu)
+
+# The GLM-MoE-DSA kernels are written for the SM100 (datacenter Blackwell)
+# family (for example 32-byte global loads), so pin the target to that family
+# instead of inheriting the global architecture list. On builds without any
+# sm_100-family architecture the target still compiles: device code that needs
+# newer instructions is guarded by __CUDA_ARCH__ and traps.
+set_cuda_architectures(glm_moe_dsa_kernels_src 100f)
+
+set_property(TARGET glm_moe_dsa_kernels_src PROPERTY POSITION_INDEPENDENT_CODE
+                                                     ON)
+set_property(TARGET glm_moe_dsa_kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS
+                                                     ON)
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu
new file mode 100644
index 000000000000..d859f89960b4
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu
@@ -0,0 +1,1557 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdlib>
+#include <stdexcept>
+#include <string>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::cross_rank_topk
+{
+namespace
+{
+
+constexpr int32_t kBins = 2048;
+static_assert(kBins >= kTopK, "the output row is staged in the histogram");
+constexpr int32_t kWarps = kThreads / 32;
+//! Per-warp key statistics kept in shared memory (see KeyStats).
+constexpr int32_t kNumKeyStats = 5;
+//! Tokens of each page one source holds.
+constexpr int32_t kBlockTokens = kPageTokens / kSourcesPerRank;
+//! Static shared memory of the kernel (round-2 segment tables), kept out of the dynamic budget.
+constexpr size_t kStaticSmemBytes = 1024;
+
+// Scalar slots in shared memory.
+enum ScalarSlot : int32_t
+{
+    kSlotBin = 0,
+    kSlotAbove,
+    kSlotBinCount,
+    kSlotTotal,
+    kSlotCounter2,
+    kSlotSrcCount,                               // [kMaxSources]
+    kSlotSrcBound = kSlotSrcCount + kMaxSources, // [kMaxSources]
+    kSlotSrcLen = kSlotSrcBound + kMaxSources,   // [kMaxSources]
+    kSlotSrcOffset = kSlotSrcLen + kMaxSources,  // [kMaxSources + 1]
+    kNumScalarSlots = kSlotSrcOffset + kMaxSources + 1
+};
+
+__host__ __device__ inline int32_t numSources(int32_t world)
+{
+    return world * kSourcesPerRank;
+}
+
+// Even share of kTopK per source.
+__host__ __device__ inline int32_t evenShare(int32_t world)
+{
+    return (kTopK + numSources(world) - 1) / numSources(world);
+}
+
+// Most keys a source sends in round 1 (an even number).
+__host__ __device__ inline int32_t roundCap(int32_t world)
+{
+    return min(kSourceCapacity, 4 * evenShare(world));
+}
+
+// Keys one source can hold: one block of every local page.
+__host__ __device__ inline int32_t chunkCapacity(int32_t localCapacity)
+{
+    int32_t const pages = (localCapacity + kPageTokens - 1) / kPageTokens;
+    return pages * kBlockTokens;
+}
+
+__host__ __device__ inline int32_t mergeCapacity(int32_t world)
+{
+    return numSources(world) * roundCap(world);
+}
+
+__host__ __device__ inline size_t smemBytes(int32_t localCapacity, int32_t world)
+{
+    size_t words = 2 * size_t{kBins} + 2 * size_t{kWarps} + size_t{kNumKeyStats} * kWarps + size_t{kNumScalarSlots};
+    words = (words + 3) / 4 * 4;
+    words += static_cast<size_t>(chunkCapacity(localCapacity)) + 2 * static_cast<size_t>(mergeCapacity(world));
+    return words * sizeof(uint32_t);
+}
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900)
+// The kernel targets SM90+ (multicast Lamport exchange).
+__global__ void __launch_bounds__(kThreads, 1) cpTopkKernel(CpTopkParams p)
+{
+    __trap();
+}
+#else
+
+//! Bits of one radix-select digit (kBins bins).
+constexpr int32_t kDigitBits = 11;
+//! Order-preserving key of +0.0: the keys of non-negative values are at least this, those of negative values below.
+//! Real indexer logits are mostly negative, and the cut of a row (and of many sources) often lies among small
+//! negative values while the largest keys are positive: one key range across zero spans ~2^31 key units, where a
+//! radix digit has no resolution. The selects therefore work on one sign at a time: with fewer than `need` keys at
+//! or above kPositiveKey these are all taken and the cut is searched among the negative keys only, else among the
+//! non-negative keys only; each range is narrow enough for one or two passes.
+constexpr uint32_t kPositiveKey = 0x80000000u;
+//! Independent 16-byte candidate loads one merge thread issues before it waits on any of them.
+constexpr int32_t kPollBatch = 4;
+//! Polls of a still-empty Lamport word between two reads of the global timer.
+constexpr int32_t kPollsPerTimerRead = 256;
+
+// Keys a saturated source aims to send in round 1 (kSourceBudget at 8 ranks).
+__device__ __forceinline__ int32_t roundNeed(int32_t world)
+{
+    return min(kSourceCapacity, 3 * evenShare(world));
+}
+
+// Tokens of `rank` among the global tokens [0, n).
+__host__ __device__ __forceinline__ int32_t localLength(int32_t n, int32_t rank, int32_t world)
+{
+    if (n <= 0)
+    {
+        return 0;
+    }
+    int32_t const fullPages = n / kPageTokens;
+    int32_t const rem = n % kPageTokens;
+    int32_t const ownedFull = (fullPages - rank + world - 1) / world;
+    int32_t len = ownedFull * kPageTokens;
+    if (rem > 0 && fullPages % world == rank)
+    {
+        len += rem;
+    }
+    return len;
+}
+
+// Source `src` of a rank holds one block of kBlockTokens consecutive tokens of every local page: block
+// (src + page) % kSourcesPerRank. A fully selected page is thus spread over all sources of its rank.
+// Keys of source `src` among `localLen` local columns (only the last local page can be partial).
+__host__ __device__ __forceinline__ int32_t sourceLength(int32_t localLen, int32_t src)
+{
+    int32_t const pages = (localLen + kPageTokens - 1) / kPageTokens;
+    if (pages == 0)
+    {
+        return 0;
+    }
+    int32_t const last = pages - 1;
+    int32_t const tail = localLen - last * kPageTokens - ((src + last) % kSourcesPerRank) * kBlockTokens;
+    return last * kBlockTokens + min(max(tail, 0), kBlockTokens);
+}
+
+// Global token of a local column under the page-interleaved partition.
+__device__ __forceinline__ uint32_t localToGlobal(int32_t col, int32_t rank, int32_t world)
+{
+    int32_t const page = col / kPageTokens;
+    return static_cast<uint32_t>((page * world + rank) * kPageTokens + (col % kPageTokens));
+}
+
+struct SmemLayout
+{
+    uint32_t* hist;     // [2 * kBins] radix histograms of alternate passes; the first one stages the output row
+    uint32_t* warpSums; // [2 * kWarps]
+    uint32_t* stats;    // [kNumKeyStats * kWarps] per-warp key statistics
+    uint32_t* scal;     // [kNumScalarSlots]
+    uint32_t* chunk;    // [chunkCapacity] order-preserving keys of this source
+    uint32_t* mkey;     // [mergeCapacity] staged round-1 pairs, then the merge candidates (merge CTA)
+    uint32_t* midx;     // [mergeCapacity]
+};
+
+__device__ inline SmemLayout carveSmem(uint32_t* base, int32_t localCapacity, int32_t world)
+{
+    SmemLayout s;
+    uint32_t* p = base;
+    s.hist = p;
+    p += 2 * kBins;
+    s.warpSums = p;
+    p += 2 * kWarps;
+    s.stats = p;
+    p += kNumKeyStats * kWarps;
+    s.scal = p;
+    p += kNumScalarSlots;
+    p = base + ((p - base) + 3) / 4 * 4;
+    s.chunk = p;
+    p += chunkCapacity(localCapacity);
+    s.mkey = p;
+    p += mergeCapacity(world);
+    s.midx = p;
+    return s;
+}
+
+__device__ __forceinline__ uint32_t orderedKey(float v)
+{
+    uint32_t const b = __float_as_uint(v);
+    uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u);
+    // NaN keys above +inf are clamped so that no key, and no exclusive bound key + 1, equals the empty marker.
+    return k < kEmptyWord - 1u ? k : kEmptyWord - 2u;
+}
+
+// Bits of x (0 for 0): the digit range of keys in [base, base + x].
+__device__ __forceinline__ int32_t bitLength(uint32_t x)
+{
+    return x == 0u ? 0 : 32 - __clz(x);
+}
+
+// Statistics of a set of keys: extremes overall and per sign, and the count of non-negative keys.
+struct KeyStats
+{
+    uint32_t maxKey = 0u;         // largest key (0 when none)
+    uint32_t minKey = kEmptyWord; // smallest key (kEmptyWord when none)
+    uint32_t positives = 0u;      // keys >= kPositiveKey
+    uint32_t maxNeg = 0u;         // largest key below kPositiveKey (0 when none)
+    uint32_t minPos = kEmptyWord; // smallest key >= kPositiveKey (kEmptyWord when none)
+
+    __device__ __forceinline__ void add(uint32_t k)
+    {
+        bool const pos = k >= kPositiveKey;
+        maxKey = max(maxKey, k);
+        minKey = min(minKey, k);
+        positives += pos ? 1u : 0u;
+        maxNeg = pos ? maxNeg : max(maxNeg, k);
+        minPos = pos ? min(minPos, k) : minPos;
+    }
+};
+
+// Reduce a thread's statistics over its warp; lane 0 stores them in the per-warp array (read with readStats after
+// the next block barrier).
+__device__ __forceinline__ void storeWarpStats(KeyStats const& t, uint32_t* stats)
+{
+    uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, t.maxKey);
+    uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, t.minKey);
+    uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, t.positives);
+    uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg);
+    uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, t.minPos);
+    if (threadIdx.x % 32 == 0)
+    {
+        int32_t const w = threadIdx.x / 32;
+        stats[w] = v0;
+        stats[kWarps + w] = v1;
+        stats[2 * kWarps + w] = v2;
+        stats[3 * kWarps + w] = v3;
+        stats[4 * kWarps + w] = v4;
+    }
+}
+
+// Block statistics from the per-warp array (every warp reduces it; no barrier).
+__device__ __forceinline__ KeyStats readStats(uint32_t const* stats)
+{
+    int32_t const lane = threadIdx.x % 32;
+    KeyStats b;
+    b.maxKey = __reduce_max_sync(0xFFFFFFFFu, stats[lane]);
+    b.minKey = __reduce_min_sync(0xFFFFFFFFu, stats[kWarps + lane]);
+    b.positives = __reduce_add_sync(0xFFFFFFFFu, stats[2 * kWarps + lane]);
+    b.maxNeg = __reduce_max_sync(0xFFFFFFFFu, stats[3 * kWarps + lane]);
+    b.minPos = __reduce_min_sync(0xFFFFFFFFu, stats[4 * kWarps + lane]);
+    return b;
+}
+
+// Digit setup of a select of the need-th largest key: over the non-negative keys when there are at least `need` of
+// them, else over the negative keys with every non-negative key counted as taken (`aboveHi`).
+struct SignRange
+{
+    uint32_t lo;      // smallest key of the range (also the digit base)
+    uint64_t hi;      // exclusive end of the histogrammed range
+    int32_t bits;     // digit range of the range
+    uint32_t aboveHi; // keys at or above hi, all taken
+    bool negative;    // the cut is among the negative keys
+};
+
+__device__ __forceinline__ SignRange signRange(KeyStats const& st, uint32_t need)
+{
+    if (st.positives >= need || st.positives == 0u)
+    {
+        uint32_t const lo = st.positives >= need ? st.minPos : st.minKey;
+        return SignRange{lo, uint64_t{kEmptyWord}, bitLength(st.maxKey - lo), 0u, st.positives == 0u};
+    }
+    return SignRange{st.minKey, uint64_t{kPositiveKey}, bitLength(st.maxNeg - st.minKey), st.positives, true};
+}
+
+__device__ __forceinline__ int64_t round1Offset(int32_t buf, int32_t src, int32_t row)
+{
+    return int64_t{buf} * kBufferWords + (int64_t{src} * kMaxRows + row) * kRound1SlotWords;
+}
+
+__device__ __forceinline__ int64_t round2Offset(int32_t buf, int32_t src, int32_t row)
+{
+    return int64_t{buf} * kBufferWords + int64_t{kMaxSources} * kMaxRows * kRound1SlotWords
+        + (int64_t{src} * kMaxRows + row) * kRound2SlotWords;
+}
+
+//! Word offset of the candidates of a slot (after the 16-byte header).
+constexpr int32_t kSlotHeaderWords = 4;
+
+__device__ __forceinline__ uint2 loadVolatile2(uint32_t const* ptr)
+{
+    uint2 v;
+    asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(ptr) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadVolatile4(uint32_t const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ void store2(uint32_t* ptr, uint32_t a, uint32_t b)
+{
+    asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(ptr), "r"(a), "r"(b) : "memory");
+}
+
+__device__ __forceinline__ void store4(uint32_t* ptr, uint4 v)
+{
+    asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)
+                 : "memory");
+}
+
+__device__ __forceinline__ bool isEmpty(uint2 v)
+{
+    return v.x == kEmptyWord || v.y == kEmptyWord;
+}
+
+// Whether the first pair of a 16-byte candidate vector, and the second one when `second`, were written.
+__device__ __forceinline__ bool vectorReady(uint4 v, bool second)
+{
+    return v.x != kEmptyWord && v.y != kEmptyWord && (!second || (v.z != kEmptyWord && v.w != kEmptyWord));
+}
+
+__device__ __forceinline__ uint64_t globalTimerNs()
+{
+    uint64_t t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return t;
+}
+
+// What a polling thread needs to give up: the error word, the timeout and this rank (for the report).
+struct WaitCtx
+{
+    int32_t* state;
+    uint64_t timeoutNs;
+    int32_t rank;
+};
+
+// A peer did not deliver within the timeout. The selection cannot be completed correctly any more, and the peer's
+// late writes could land in a buffer this rank resets or reuses, so record which wait failed and abort the kernel:
+// the resulting CUDA error fails the run instead of letting any rank continue with a partial selection.
+__device__ __noinline__ void abortOnTimeout(WaitCtx const& w, int32_t what)
+{
+    bool const first = atomicExch(&w.state[3], what) == 0;
+    __threadfence_system();
+    if (first)
+    {
+        printf(
+            "[glm_moe_dsa] crossRankTopk: rank %d, row %d: wait %d for a peer exceeded %llu ms; aborting the kernel\n",
+            w.rank, static_cast<int32_t>(blockIdx.y), what, static_cast<unsigned long long>(w.timeoutNs / 1000000ull));
+    }
+    __trap();
+}
+
+// Spin until both words of a Lamport pair were written; aborts the kernel after the timeout.
+__device__ uint2 pollPair(uint32_t const* ptr, WaitCtx const& w, int32_t what)
+{
+    uint2 v = loadVolatile2(ptr);
+    uint64_t start = 0;
+    for (int32_t polls = 1; isEmpty(v); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, what);
+            }
+        }
+        v = loadVolatile2(ptr);
+    }
+    return v;
+}
+
+// Spin until a 16-byte candidate vector was written (see vectorReady); aborts the kernel after the timeout.
+__device__ uint4 pollVector(uint32_t const* ptr, bool second, WaitCtx const& w, int32_t what)
+{
+    uint4 v = loadVolatile4(ptr);
+    uint64_t start = 0;
+    for (int32_t polls = 1; !vectorReady(v, second); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, what);
+            }
+        }
+        v = loadVolatile4(ptr);
+    }
+    return v;
+}
+
+// Block-wide inclusive scan of one value per thread. Returns the inclusive prefix; `total` gets the block sum.
+__device__ uint32_t blockInclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    uint32_t x = v;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    if (lane == 31)
+    {
+        warpSums[warp] = x;
+    }
+    __syncthreads();
+    if (warp == 0)
+    {
+        uint32_t w = warpSums[lane];
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, w, o);
+            if (lane >= o)
+            {
+                w += y;
+            }
+        }
+        warpSums[kWarps + lane] = w;
+    }
+    __syncthreads();
+    total = warpSums[2 * kWarps - 1];
+    uint32_t const before = warp > 0 ? warpSums[kWarps + warp - 1] : 0u;
+    return x + before;
+}
+
+// Exclusive block-wide scan with a single barrier: every warp scans the per-warp sums itself. Returns this
+// thread's exclusive prefix; `total` gets the block sum. warpSums must not be rewritten before the next barrier.
+__device__ uint32_t blockExclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    uint32_t x = v;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    if (lane == 31)
+    {
+        warpSums[warp] = x;
+    }
+    __syncthreads();
+    uint32_t const w = warpSums[lane];
+    uint32_t y = w;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const z = __shfl_up_sync(0xFFFFFFFFu, y, o);
+        if (lane >= o)
+        {
+            y += z;
+        }
+    }
+    total = __shfl_sync(0xFFFFFFFFu, y, 31);
+    return __shfl_sync(0xFFFFFFFFu, y - w, warp) + x - v;
+}
+
+// Find, counting from the top bin, the bin of `hist` that holds the target-th largest key, with two barriers:
+// per-warp sums, then each warp scans them and only the warp holding the target scans its own bins. Returns the
+// total count (every thread); when target <= total, the bin, the count above it and its own count go to the
+// scalar slots.
+__device__ uint32_t findBin(SmemLayout const& s, uint32_t const* hist, int32_t nbins, uint32_t target)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    int32_t const perThread = (nbins + kThreads - 1) / kThreads;
+    int32_t const top = nbins - 1 - perThread * static_cast<int32_t>(threadIdx.x);
+    uint32_t mine = 0;
+    for (int32_t i = 0; i < perThread && top - i >= 0; ++i)
+    {
+        mine += hist[top - i];
+    }
+    uint32_t const warpTotal = __reduce_add_sync(0xFFFFFFFFu, mine);
+    if (lane == 0)
+    {
+        s.warpSums[warp] = warpTotal;
+    }
+    __syncthreads();
+    uint32_t const w = s.warpSums[lane];
+    uint32_t x = w;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    uint32_t const total = __shfl_sync(0xFFFFFFFFu, x, 31);
+    uint32_t const warpBefore = __shfl_sync(0xFFFFFFFFu, x - w, warp);
+    if (warpBefore < target && target <= warpBefore + warpTotal)
+    {
+        uint32_t t = mine;
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, t, o);
+            if (lane >= o)
+            {
+                t += y;
+            }
+        }
+        uint32_t acc = warpBefore + t - mine;
+        if (acc < target && target <= acc + mine)
+        {
+            for (int32_t i = 0; i < perThread && top - i >= 0; ++i)
+            {
+                uint32_t const c = hist[top - i];
+                if (acc + c >= target)
+                {
+                    s.scal[kSlotBin] = static_cast<uint32_t>(top - i);
+                    s.scal[kSlotAbove] = acc;
+                    s.scal[kSlotBinCount] = c;
+                    break;
+                }
+                acc += c;
+            }
+        }
+    }
+    __syncthreads();
+    return total;
+}
+
+struct SelectOut
+{
+    uint32_t edge;     // lower edge of the selected bin (the exact key after the last digit)
+    uint32_t above;    // predicate keys strictly above the selected bin
+    uint32_t binCount; // predicate keys in the selected bin
+    uint32_t total;    // predicate keys
+    bool all;          // take every predicate key (fewer than needed, or all fit the capacity)
+    bool full;         // exact mode: the keys at or above `edge` are exactly the `need` largest
+    int32_t passes;    // radix passes run
+};
+
+// Candidate keys and token indices held in shared memory (16-byte aligned, readable in groups of four up to the
+// next multiple of four).
+struct SmemKeys
+{
+    static constexpr bool kVector = true;
+    uint32_t const* keys;
+    uint32_t const* idx;
+
+    __device__ uint4 keys4(int32_t j) const
+    {
+        return *reinterpret_cast<uint4 const*>(keys + j);
+    }
+
+    __device__ uint4 index4(int32_t j) const
+    {
+        return *reinterpret_cast<uint4 const*>(idx + j);
+    }
+};
+
+// Merge candidates in shared memory followed by round-2 candidates read from their global Lamport slots.
+struct MergeKeys
+{
+    static constexpr bool kVector = false;
+    uint32_t const* smemKeys;
+    uint32_t const* smemIdx;
+    int32_t smemCount;
+    uint32_t const* const* slots; // round-2 slot of each segment
+    uint32_t const* segBegin;     // [numSegments + 1] candidate offsets
+    int32_t numSegments;
+
+    __device__ uint32_t const* pair(int32_t k) const
+    {
+        // Last segment whose begin is <= k.
+        int32_t lo = 0;
+        int32_t hi = numSegments - 1;
+        while (lo < hi)
+        {
+            int32_t const mid = (lo + hi + 1) / 2;
+            if (static_cast<int32_t>(segBegin[mid]) <= k)
+            {
+                lo = mid;
+            }
+            else
+            {
+                hi = mid - 1;
+            }
+        }
+        return slots[lo] + kSlotHeaderWords + 2 * (k - static_cast<int32_t>(segBegin[lo]));
+    }
+
+    __device__ uint32_t operator()(int32_t j) const
+    {
+        return j < smemCount ? smemKeys[j] : pair(j - smemCount)[0];
+    }
+
+    __device__ uint32_t index(int32_t j) const
+    {
+        return j < smemCount ? smemIdx[j] : pair(j - smemCount)[1];
+    }
+};
+
+// Debug slots of this CTA starting at `offset` (nullptr without a debug buffer).
+__device__ __forceinline__ int64_t* debugSlots(CpTopkParams const& p, int32_t offset)
+{
+    if (p.debugTimes == nullptr)
+    {
+        return nullptr;
+    }
+    return p.debugTimes + (static_cast<int64_t>(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots + offset;
+}
+
+// clock64 of thread 0 into a debug slot (no-op without a debug buffer).
+__device__ __forceinline__ void stamp(int64_t* dbg, int32_t slot)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[slot] = clock64();
+    }
+}
+
+// A value into a debug slot (thread 0; no-op without a debug buffer).
+__device__ __forceinline__ void putDebug(int64_t* dbg, int32_t slot, int64_t value)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[slot] = value;
+    }
+}
+
+// Phase timestamp for diagnosis (thread 0 of the CTA; no-op without a debug buffer).
+__device__ __forceinline__ void markPhase(CpTopkParams const& p, int32_t phase)
+{
+    if (p.debugTimes != nullptr && threadIdx.x == 0)
+    {
+        uint64_t t;
+        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+        int64_t* d = p.debugTimes + (static_cast<int64_t>(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots;
+        d[phase] = static_cast<int64_t>(t);
+        d[16 + phase] = static_cast<int64_t>(clock64());
+    }
+}
+
+// Histogram update of one key of a radix pass.
+__device__ __forceinline__ void histogramKey(uint32_t* hist, bool inRange, uint32_t k, uint32_t lo, uint64_t hi,
+    uint32_t base, uint32_t prefMask, uint32_t prefVal, int32_t shift, uint32_t binMask)
+{
+    uint32_t const d = k - base;
+    if (inRange && k >= lo && static_cast<uint64_t>(k) < hi && (d & prefMask) == prefVal)
+    {
+        atomicAdd(&hist[(d >> shift) & binMask], 1u);
+    }
+}
+
+// Block radix select of the need-th largest of the n keys at or above lo, of which `aboveHi` (< need) lie at or
+// above hi and are counted as taken without being histogrammed. Every key in [lo, hi) is at least `base`, and
+// key - base < 2^topBits; digits of up to kDigitBits bits of key - base are refined from the top. With cap > 0 the
+// select stops at the first bin edge whose keys at or above it number at most `cap`. With cap = 0 it runs to the
+// exact key, or stops early when the selected bin is wholly taken (`full`). Alternate passes use alternate
+// histograms, so the next one is cleared while the current one is searched. Counts in the result include aboveHi.
+template <typename KeyAt>
+__device__ SelectOut blockSelect(SmemLayout const& s, KeyAt const& keyAt, int32_t n, uint32_t lo, uint64_t hi,
+    uint32_t need, uint32_t cap, uint32_t base, int32_t topBits, int64_t* dbg = nullptr, uint32_t aboveHi = 0u)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[0] = n;
+    }
+    SelectOut out{lo, 0u, 0u, 0u, false, false, 0};
+    uint32_t prefMask = 0;
+    uint32_t prefVal = 0;
+    uint32_t above = aboveHi;
+    int32_t remaining = max(topBits, 1);
+    uint32_t* hist = s.hist;
+    uint32_t* histNext = s.hist + kBins;
+    for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads)
+    {
+        hist[b] = 0;
+    }
+    __syncthreads();
+#pragma unroll 1
+    for (int32_t pass = 0; remaining > 0; ++pass)
+    {
+        int32_t const width = min(kDigitBits, remaining);
+        int32_t const shift = remaining - width;
+        int32_t const nbins = 1 << width;
+        uint32_t const binMask = static_cast<uint32_t>(nbins - 1);
+        out.passes = pass + 1;
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[1 + 3 * pass] = clock64();
+        }
+        if constexpr (KeyAt::kVector)
+        {
+            for (int32_t j0 = 4 * static_cast<int32_t>(threadIdx.x); j0 < n; j0 += 4 * kThreads)
+            {
+                uint4 const q = keyAt.keys4(j0);
+                histogramKey(hist, true, q.x, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 1 < n, q.y, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 2 < n, q.z, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 3 < n, q.w, lo, hi, base, prefMask, prefVal, shift, binMask);
+            }
+        }
+        else
+        {
+            for (int32_t j = threadIdx.x; j < n; j += kThreads)
+            {
+                histogramKey(hist, true, keyAt(j), lo, hi, base, prefMask, prefVal, shift, binMask);
+            }
+        }
+        __syncthreads();
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[2 + 3 * pass] = clock64();
+        }
+        if (shift > 0)
+        {
+            for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads)
+            {
+                histNext[b] = 0;
+            }
+        }
+        // Find the bin holding the (need - above)-th largest key, counting from the top bin.
+        uint32_t const target = need - above;
+        uint32_t const total = findBin(s, hist, nbins, target);
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[3 + 3 * pass] = clock64();
+        }
+        if (pass == 0)
+        {
+            out.total = total + aboveHi;
+            if (total + aboveHi <= cap || target > total)
+            {
+                out.all = true;
+                out.above = total + aboveHi;
+                return out;
+            }
+        }
+        // The scalars are next written after two more barriers (the next histogram and scan).
+        uint32_t const bin = s.scal[kSlotBin];
+        above += s.scal[kSlotAbove];
+        out.binCount = s.scal[kSlotBinCount];
+        prefVal |= bin << shift;
+        prefMask |= binMask << shift;
+        out.edge = base + prefVal;
+        out.above = above;
+        remaining = shift;
+        uint32_t* const t = hist;
+        hist = histNext;
+        histNext = t;
+        if (cap > 0 && above + out.binCount <= cap)
+        {
+            return out;
+        }
+        if (cap == 0 && remaining > 0 && above + out.binCount == need)
+        {
+            out.full = true;
+            return out;
+        }
+    }
+    return out;
+}
+
+// Send one source's selected keys (keys in [lo, hi) of the chunk) to its slot; returns the count sent.
+// Takes every predicate key at or above the edge when they fit `cap`. Otherwise (more than `cap` keys tie at the
+// exact edge) it takes every key above the edge, plus, with kPartialTies, the first keys equal to it in source
+// order up to `cap` (any subset of equal keys is a valid tie-aware choice). Round 1 sends no partial ties, so the
+// keys a second round requests from a source (at most the edge) were never sent before.
+// The header carries the count and an exclusive upper bound of the keys left unsent (0 when none).
+// With `stage` (round 1, at most kSourceCapacity keys), the pairs are gathered in shared memory first and written
+// with coalesced 16-byte stores; otherwise each pair is stored where it is found.
+template <bool kPartialTies>
+__device__ uint32_t sendSource(SmemLayout const& s, SelectOut const& sel, int32_t n, uint32_t lo, uint64_t hi,
+    uint32_t cap, uint32_t* mcSlot, int32_t src, int32_t rank, int32_t world, uint2* stage, int64_t* dbg = nullptr)
+{
+    bool const strict = !sel.all && sel.above + sel.binCount > cap;
+    uint32_t const threshold = sel.all ? lo : sel.edge;
+    // Keys equal to the edge that fit after the keys above it (sel.above < need <= cap when strict).
+    uint32_t const tieBudget = kPartialTies && strict ? cap - sel.above : 0u;
+    uint32_t const sent = sel.all ? sel.total : (strict ? sel.above + tieBudget : sel.above + sel.binCount);
+    uint32_t const bound = sel.all ? lo : (strict ? sel.edge + 1u : sel.edge);
+    stamp(dbg, 42);
+    // One page block (kBlockTokens keys) per thread; positions from a block scan, so the slot holds the keys in
+    // source order: first those taken outright, then (strict only) the admitted keys equal to the edge.
+    int32_t const groups = (n + kBlockTokens - 1) / kBlockTokens;
+    uint32_t takenMain = 0;
+    uint32_t takenTies = 0;
+    uint32_t const tieBase = strict ? sel.above : 0u;
+    for (int32_t g0 = 0; g0 < groups; g0 += kThreads)
+    {
+        int32_t const g = g0 + static_cast<int32_t>(threadIdx.x);
+        uint32_t k[kBlockTokens];
+        uint32_t mask = 0;
+        uint32_t tieMask = 0;
+        if (g < groups)
+        {
+            uint4 const a = *reinterpret_cast<uint4 const*>(s.chunk + kBlockTokens * g);
+            uint4 const b = *reinterpret_cast<uint4 const*>(s.chunk + kBlockTokens * g + 4);
+            k[0] = a.x;
+            k[1] = a.y;
+            k[2] = a.z;
+            k[3] = a.w;
+            k[4] = b.x;
+            k[5] = b.y;
+            k[6] = b.z;
+            k[7] = b.w;
+#pragma unroll
+            for (int32_t i = 0; i < kBlockTokens; ++i)
+            {
+                bool const in = kBlockTokens * g + i < n && k[i] >= lo && static_cast<uint64_t>(k[i]) < hi;
+                bool const take = in && (strict ? k[i] > threshold : k[i] >= threshold);
+                bool const tie = kPartialTies && in && strict && k[i] == threshold;
+                mask |= (take ? 1u : 0u) << i;
+                tieMask |= (tie ? 1u : 0u) << i;
+            }
+        }
+        // At most kThreads * kBlockTokens keys per pass: both counts fit 16 bits.
+        uint32_t const c = static_cast<uint32_t>(__popc(mask)) | (static_cast<uint32_t>(__popc(tieMask)) << 16);
+        uint32_t chunkTaken;
+        uint32_t const before = blockExclusiveScan(c, s.warpSums, chunkTaken);
+        uint32_t pos = takenMain + (before & 0xFFFFu);
+        uint32_t tiePos = takenTies + (before >> 16);
+        takenMain += chunkTaken & 0xFFFFu;
+        takenTies += chunkTaken >> 16;
+        if ((mask | tieMask) != 0u)
+        {
+            int32_t const colBase = g * kPageTokens + ((src + g) % kSourcesPerRank) * kBlockTokens;
+#pragma unroll
+            for (int32_t i = 0; i < kBlockTokens; ++i)
+            {
+                uint32_t slot = kEmptyWord;
+                if ((mask >> i) & 1u)
+                {
+                    slot = pos++;
+                }
+                else if constexpr (kPartialTies)
+                {
+                    if ((tieMask >> i) & 1u)
+                    {
+                        slot = tiePos < tieBudget ? tieBase + tiePos : kEmptyWord;
+                        ++tiePos;
+                    }
+                }
+                if (slot != kEmptyWord)
+                {
+                    uint32_t const token = localToGlobal(colBase + i, rank, world);
+                    if (stage != nullptr)
+                    {
+                        stage[slot] = make_uint2(k[i], token);
+                    }
+                    else
+                    {
+                        store2(mcSlot + kSlotHeaderWords + 2 * static_cast<int64_t>(slot), k[i], token);
+                    }
+                }
+            }
+        }
+        // The staged pairs are complete, and the scan's warp sums may be rewritten.
+        __syncthreads();
+    }
+    stamp(dbg, 43);
+    if (stage != nullptr)
+    {
+        // Two pairs per 16-byte store; an odd count leaves the last half of the last vector empty.
+        int32_t const vectors = static_cast<int32_t>(sent + 1u) / 2;
+        for (int32_t v = threadIdx.x; v < vectors; v += kThreads)
+        {
+            uint2 const a = stage[2 * v];
+            uint2 const b
+                = 2 * v + 1 < static_cast<int32_t>(sent) ? stage[2 * v + 1] : make_uint2(kEmptyWord, kEmptyWord);
+            store4(mcSlot + kSlotHeaderWords + 4 * static_cast<int64_t>(v), make_uint4(a.x, a.y, b.x, b.y));
+        }
+    }
+    // Receivers check every candidate word, so the header needs no ordering against the candidates.
+    if (threadIdx.x == kThreads - 1)
+    {
+        store2(mcSlot, sent, bound);
+    }
+    stamp(dbg, 44);
+    return sent;
+}
+
+// Load one source's keys (one block of kBlockTokens tokens of every local page) from the local logits into the
+// chunk; returns the number of keys. Consecutive threads read consecutive keys (whole 32-byte blocks per eight
+// threads), and each thread issues kLoadBatch loads before it uses any. With `withStats`, the per-warp key
+// statistics are stored for readStats after the next block barrier.
+__device__ int32_t loadSource(
+    SmemLayout const& s, float const* rowLogits, int32_t localLen, int32_t src, bool withStats)
+{
+    constexpr int32_t kLoadBatch = 8;
+    int32_t const n = sourceLength(localLen, src);
+    KeyStats st;
+    for (int32_t j0 = threadIdx.x; j0 < n; j0 += kLoadBatch * kThreads)
+    {
+        float v[kLoadBatch];
+#pragma unroll
+        for (int32_t i = 0; i < kLoadBatch; ++i)
+        {
+            int32_t const j = j0 + i * kThreads;
+            int32_t const page = j / kBlockTokens;
+            int32_t const col = page * kPageTokens + ((src + page) % kSourcesPerRank) * kBlockTokens + j % kBlockTokens;
+            v[i] = j < n ? rowLogits[col] : 0.F;
+        }
+#pragma unroll
+        for (int32_t i = 0; i < kLoadBatch; ++i)
+        {
+            int32_t const j = j0 + i * kThreads;
+            if (j < n)
+            {
+                uint32_t const k = orderedKey(v[i]);
+                s.chunk[j] = k;
+                st.add(k);
+            }
+        }
+    }
+    if (withStats)
+    {
+        storeWarpStats(st, s.stats);
+    }
+    return n;
+}
+
+// Whether source q withheld keys that may lie above the merged k-th key.
+__device__ __forceinline__ bool sourceOverflows(SmemLayout const& s, int32_t q, uint32_t kth)
+{
+    uint32_t const cnt = s.scal[kSlotSrcCount + q];
+    uint64_t const bound = s.scal[kSlotSrcBound + q];
+    return cnt < s.scal[kSlotSrcLen + q] && bound > uint64_t{kth} + 1u;
+}
+
+// Emit the top-kTopK of the merge candidates. Keys above `kth` are all taken; ties at `kth` are taken in candidate
+// order when only some of them fit (deterministic on every rank), otherwise all of them. With `takeAll`, every
+// candidate is taken and the row is padded with -1. The row is assembled in shared memory (the first histogram)
+// and written with coalesced stores.
+template <typename KeyAt>
+__device__ bool emitSelection(SmemLayout const& s, KeyAt const& keys, int32_t total, uint32_t kth, bool takeAll,
+    int32_t* outRow, int32_t overflowSources, int64_t* dbg = nullptr)
+{
+    uint32_t const want = static_cast<uint32_t>(kTopK);
+    // Positions from one block scan over contiguous candidate ranges (no contended counter); the order is the
+    // candidate order, identical on every rank. Shared-memory candidates are read four at a time.
+    int32_t per = (total + kThreads - 1) / kThreads;
+    if constexpr (KeyAt::kVector)
+    {
+        per = (per + 3) / 4 * 4;
+    }
+    int32_t const j0 = min(total, static_cast<int32_t>(threadIdx.x) * per);
+    int32_t const j1 = min(total, j0 + per);
+    uint32_t cGt = 0;
+    uint32_t cEq = 0;
+    if constexpr (KeyAt::kVector)
+    {
+        for (int32_t j = j0; j < j1; j += 4)
+        {
+            uint4 const q = keys.keys4(j);
+            uint32_t const kk[4] = {q.x, q.y, q.z, q.w};
+#pragma unroll
+            for (int32_t i = 0; i < 4; ++i)
+            {
+                bool const in = j + i < j1;
+                cGt += (in && (takeAll || kk[i] > kth)) ? 1u : 0u;
+                cEq += (in && !takeAll && kk[i] == kth) ? 1u : 0u;
+            }
+        }
+    }
+    else
+    {
+        for (int32_t j = j0; j < j1; ++j)
+        {
+            uint32_t const k = keys(j);
+            cGt += (takeAll || k > kth) ? 1u : 0u;
+            cEq += (!takeAll && k == kth) ? 1u : 0u;
+        }
+    }
+    // The first `overflowSources` threads flag a source that withheld keys above the threshold (the flag slot is
+    // zero on entry); with a flag set, nothing is written and the caller runs the second round.
+    if (static_cast<int32_t>(threadIdx.x) < overflowSources && sourceOverflows(s, threadIdx.x, kth))
+    {
+        atomicOr(&s.scal[kSlotCounter2], 1u);
+    }
+    uint32_t allGt;
+    uint32_t allEq;
+    uint32_t posGt;
+    uint32_t posEq;
+    if (total < 0x10000)
+    {
+        // Both counts fit 16 bits: one scan of the packed pair.
+        uint32_t packedTotal;
+        uint32_t const packed = cGt | (cEq << 16);
+        uint32_t const before = blockExclusiveScan(packed, s.warpSums, packedTotal);
+        allGt = packedTotal & 0xFFFFu;
+        allEq = packedTotal >> 16;
+        posGt = before & 0xFFFFu;
+        posEq = allGt + (before >> 16);
+    }
+    else
+    {
+        posGt = blockInclusiveScan(cGt, s.warpSums, allGt) - cGt;
+        __syncthreads();
+        posEq = blockInclusiveScan(cEq, s.warpSums, allEq) - cEq;
+        posEq += allGt;
+    }
+    stamp(dbg, 59);
+    if (overflowSources > 0 && s.scal[kSlotCounter2] != 0u)
+    {
+        return true;
+    }
+    uint32_t* row = s.hist;
+    auto place = [&](uint32_t k, uint32_t idx)
+    {
+        if (takeAll || k > kth)
+        {
+            if (posGt < want)
+            {
+                row[posGt] = idx;
+            }
+            ++posGt;
+        }
+        else if (k == kth)
+        {
+            if (posEq < want)
+            {
+                row[posEq] = idx;
+            }
+            ++posEq;
+        }
+    };
+    if constexpr (KeyAt::kVector)
+    {
+        for (int32_t j = j0; j < j1; j += 4)
+        {
+            uint4 const q = keys.keys4(j);
+            uint4 const x = keys.index4(j);
+            uint32_t const kk[4] = {q.x, q.y, q.z, q.w};
+            uint32_t const xx[4] = {x.x, x.y, x.z, x.w};
+#pragma unroll
+            for (int32_t i = 0; i < 4; ++i)
+            {
+                if (j + i < j1)
+                {
+                    place(kk[i], xx[i]);
+                }
+            }
+        }
+    }
+    else
+    {
+        for (int32_t j = j0; j < j1; ++j)
+        {
+            place(keys(j), keys.index(j));
+        }
+    }
+    __syncthreads();
+    stamp(dbg, 60);
+    uint32_t const selected = min(allGt + allEq, want);
+    for (uint32_t j = threadIdx.x; j < want; j += kThreads)
+    {
+        outRow[j] = j < selected ? static_cast<int32_t>(row[j]) : -1;
+    }
+    stamp(dbg, 61);
+    return false;
+}
+
+// Threshold of an exact (cap = 0) merge select: keys above `kth` are taken, ties at `kth` while they fit.
+__device__ __forceinline__ uint32_t mergeThreshold(SelectOut const& m, bool& takeAll)
+{
+    takeAll = m.all || (m.full && m.edge == 0u);
+    if (takeAll)
+    {
+        return 0u;
+    }
+    // A wholly taken bin: the keys at or above its edge are exactly kTopK.
+    return m.full ? m.edge - 1u : m.edge;
+}
+
+// Reset the slots of Lamport buffer `buf` on this rank's copy, using one warp per slot. A slot was written iff its
+// header was; the header's count bounds the candidates to reset.
+__device__ void clearBuffer(CpTopkParams const& p, int32_t buf, int32_t clearer, int32_t numClearers)
+{
+    int32_t const warp = threadIdx.x / 32;
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const round1Slots = kMaxSources * kMaxRows;
+    uint4 const empty = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord);
+    for (int32_t slot = clearer * kWarps + warp; slot < 2 * round1Slots; slot += numClearers * kWarps)
+    {
+        bool const second = slot >= round1Slots;
+        int32_t const qr = second ? slot - round1Slots : slot;
+        uint32_t* base = p.ucWorkspace
+            + (second ? round2Offset(buf, qr / kMaxRows, qr % kMaxRows)
+                      : round1Offset(buf, qr / kMaxRows, qr % kMaxRows));
+        uint2 h = make_uint2(kEmptyWord, kEmptyWord);
+        if (lane == 0)
+        {
+            h = loadVolatile2(base);
+        }
+        h.x = __shfl_sync(0xFFFFFFFFu, h.x, 0);
+        h.y = __shfl_sync(0xFFFFFFFFu, h.y, 0);
+        if (h.x == kEmptyWord && h.y == kEmptyWord)
+        {
+            continue;
+        }
+        uint32_t const capacity = second ? kRound2Capacity : kSourceCapacity;
+        uint32_t const count = h.x == kEmptyWord ? capacity : min(h.x, capacity);
+        int32_t const vectors = static_cast<int32_t>(count + 1u) / 2;
+        for (int32_t v = lane; v < vectors; v += 32)
+        {
+            store4(base + kSlotHeaderWords + 4 * static_cast<int64_t>(v), empty);
+        }
+        if (lane == 0)
+        {
+            store2(base, kEmptyWord, kEmptyWord);
+        }
+    }
+}
+
+// Merge of one row on this rank (one CTA): gathers every source's round-1 candidates, selects the exact
+// top-kTopK, runs the second round when a source withheld keys above the threshold, and writes the row.
+__device__ void mergeRow(CpTopkParams const& p, SmemLayout const& s, uint32_t const** segSlots, uint32_t* segBegin,
+    int32_t buf, int32_t row, int32_t rowLen, int32_t localLen, float const* rowLogits, int32_t* outRow,
+    WaitCtx const& wait)
+{
+    int32_t const world = p.worldSize;
+    int32_t const rank = p.rank;
+    int32_t const sources = numSources(world);
+    uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+    __syncthreads();
+    if (threadIdx.x == 0)
+    {
+        s.scal[kSlotCounter2] = 0u;
+    }
+    if (static_cast<int32_t>(threadIdx.x) < sources)
+    {
+        // Keys each source holds, computed while its header is in flight.
+        int32_t const q = threadIdx.x;
+        s.scal[kSlotSrcLen + q] = static_cast<uint32_t>(
+            sourceLength(min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank));
+        uint2 const h = pollPair(p.ucWorkspace + round1Offset(buf, threadIdx.x, row), wait, kWaitRound1Header);
+        s.scal[kSlotSrcCount + threadIdx.x] = min(h.x, cap);
+        s.scal[kSlotSrcBound + threadIdx.x] = h.y;
+    }
+    __syncthreads();
+    if (threadIdx.x < 32)
+    {
+        // Exclusive prefix of the source counts (at most 64 sources: two per lane).
+        uint32_t const a = static_cast<int32_t>(threadIdx.x) < sources ? s.scal[kSlotSrcCount + threadIdx.x] : 0u;
+        uint32_t const b
+            = static_cast<int32_t>(threadIdx.x) + 32 < sources ? s.scal[kSlotSrcCount + threadIdx.x + 32] : 0u;
+        uint32_t x = a;
+        uint32_t z = b;
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+            uint32_t const w = __shfl_up_sync(0xFFFFFFFFu, z, o);
+            if (static_cast<int32_t>(threadIdx.x) >= o)
+            {
+                x += y;
+                z += w;
+            }
+        }
+        uint32_t const totalA = __shfl_sync(0xFFFFFFFFu, x, 31);
+        s.scal[kSlotSrcOffset + threadIdx.x] = x - a;
+        s.scal[kSlotSrcOffset + 32 + threadIdx.x] = totalA + z - b;
+        if (threadIdx.x == 31)
+        {
+            s.scal[kSlotSrcOffset + kMaxSources] = totalA + z;
+        }
+    }
+    __syncthreads();
+    int32_t const total = static_cast<int32_t>(s.scal[kSlotSrcOffset + kMaxSources]);
+    markPhase(p, 4);
+
+    // Candidates: padded (source, vector) positions; issue a batch of independent 16-byte loads, then re-poll the
+    // ones not yet written. The key statistics are tracked for the merge select.
+    int32_t const vectorsPerSource = static_cast<int32_t>(cap) / 2;
+    int32_t const padded = sources * vectorsPerSource;
+    KeyStats st;
+#pragma unroll 1
+    for (int32_t v0 = 0; v0 < padded; v0 += kThreads * kPollBatch)
+    {
+        uint4 v[kPollBatch];
+#pragma unroll
+        for (int32_t i = 0; i < kPollBatch; ++i)
+        {
+            int32_t const jv = v0 + i * kThreads + static_cast<int32_t>(threadIdx.x);
+            int32_t const q = jv / vectorsPerSource;
+            int32_t const kv = jv % vectorsPerSource;
+            v[i] = make_uint4(0u, 0u, 0u, 0u);
+            if (jv < padded && 2 * kv < static_cast<int32_t>(s.scal[kSlotSrcCount + q]))
+            {
+                v[i] = loadVolatile4(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv);
+            }
+        }
+#pragma unroll
+        for (int32_t i = 0; i < kPollBatch; ++i)
+        {
+            int32_t const jv = v0 + i * kThreads + static_cast<int32_t>(threadIdx.x);
+            int32_t const q = jv / vectorsPerSource;
+            int32_t const kv = jv % vectorsPerSource;
+            if (jv < padded && 2 * kv < static_cast<int32_t>(s.scal[kSlotSrcCount + q]))
+            {
+                bool const second = 2 * kv + 1 < static_cast<int32_t>(s.scal[kSlotSrcCount + q]);
+                if (!vectorReady(v[i], second))
+                {
+                    v[i] = pollVector(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv, second,
+                        wait, kWaitRound1Candidates);
+                }
+                int32_t const j = static_cast<int32_t>(s.scal[kSlotSrcOffset + q]) + 2 * kv;
+                s.mkey[j] = v[i].x;
+                s.midx[j] = v[i].y;
+                st.add(v[i].x);
+                if (second)
+                {
+                    s.mkey[j + 1] = v[i].z;
+                    s.midx[j + 1] = v[i].w;
+                    st.add(v[i].z);
+                }
+            }
+        }
+    }
+    stamp(debugSlots(p, 0), 45);
+    storeWarpStats(st, s.stats);
+    __syncthreads();
+    markPhase(p, 5);
+
+    // Digits of the merge select span the candidates of the sign the cut lies in (see kPositiveKey).
+    SignRange const r = total > 0 ? signRange(readStats(s.stats), static_cast<uint32_t>(kTopK))
+                                  : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false};
+    SmemKeys const merged{s.mkey, s.midx};
+    SelectOut const m = blockSelect(
+        s, merged, total, r.lo, r.hi, static_cast<uint32_t>(kTopK), 0u, r.lo, r.bits, debugSlots(p, 48), r.aboveHi);
+    putDebug(debugSlots(p, 0), 62, total);
+    putDebug(debugSlots(p, 0), 63, m.passes + (r.negative ? 10 : 0));
+    bool takeAll = false;
+    uint32_t const kth = mergeThreshold(m, takeAll);
+    markPhase(p, 6);
+    stamp(debugSlots(p, 0), 58);
+    if (!emitSelection(s, merged, total, kth, takeAll, outRow, sources, debugSlots(p, 0)))
+    {
+        return;
+    }
+
+    // ---- Round 2 (rare; every rank takes it): withheld keys in [kth, bound) of overflowing sources ----------
+    __syncthreads();
+    // 1. This rank's overflowing sources send their top-kRound2Capacity withheld keys.
+    for (int32_t c = 0; c < kSourcesPerRank; ++c)
+    {
+        int32_t const q = rank * kSourcesPerRank + c;
+        if (!sourceOverflows(s, q, kth))
+        {
+            continue;
+        }
+        uint64_t const b = s.scal[kSlotSrcBound + q];
+        int32_t const nq = loadSource(s, rowLogits, localLen, c, false);
+        __syncthreads();
+        SelectOut const sel2 = blockSelect(
+            s, SmemKeys{s.chunk, nullptr}, nq, kth, b, static_cast<uint32_t>(kRound2Capacity), kRound2Capacity, 0u, 32);
+        (void) sendSource<true>(
+            s, sel2, nq, kth, b, kRound2Capacity, p.mcWorkspace + round2Offset(buf, q, row), c, rank, world, nullptr);
+        __syncthreads();
+    }
+    // 2. Gather the round-2 segments of every overflowing source (all ranks).
+    if (threadIdx.x == 0)
+    {
+        int32_t segs = 0;
+        uint32_t off = 0;
+        for (int32_t q = 0; q < sources; ++q)
+        {
+            if (!sourceOverflows(s, q, kth))
+            {
+                continue;
+            }
+            uint32_t const* slot = p.ucWorkspace + round2Offset(buf, q, row);
+            uint2 const h = pollPair(slot, wait, kWaitRound2Header);
+            segSlots[segs] = slot;
+            segBegin[segs] = off;
+            off += min(h.x, static_cast<uint32_t>(kRound2Capacity));
+            ++segs;
+        }
+        segBegin[segs] = off;
+        s.scal[kSlotTotal] = static_cast<uint32_t>(segs);
+    }
+    __syncthreads();
+    int32_t const segs = static_cast<int32_t>(s.scal[kSlotTotal]);
+    int32_t const total2 = static_cast<int32_t>(segBegin[segs]);
+    MergeKeys const keys{s.mkey, s.midx, total, segSlots, segBegin, segs};
+    for (int32_t k = threadIdx.x; k < total2; k += kThreads)
+    {
+        (void) pollPair(keys.pair(k), wait, kWaitRound2Candidates);
+    }
+    __syncthreads();
+    SelectOut const m2
+        = blockSelect(s, keys, total + total2, 0u, uint64_t{kEmptyWord}, static_cast<uint32_t>(kTopK), 0u, 0u, 32);
+    bool takeAll2 = false;
+    uint32_t const kth2 = mergeThreshold(m2, takeAll2);
+    (void) emitSelection(s, keys, total + total2, kth2, takeAll2, outRow, 0);
+    if (threadIdx.x == 0)
+    {
+        atomicAdd(&p.state[2], 1);
+    }
+}
+
+__global__ void __launch_bounds__(kThreads, 1) cpTopkKernel(CpTopkParams p)
+{
+    markPhase(p, 0);
+    extern __shared__ __align__(16) uint32_t smemRaw[];
+    __shared__ uint32_t const* segSlots[kMaxSources];
+    __shared__ uint32_t segBegin[kMaxSources + 1];
+    int32_t const world = p.worldSize;
+    int32_t const rank = p.rank;
+    int32_t const src = blockIdx.x;
+    int32_t const row = blockIdx.y;
+    SmemLayout const s = carveSmem(smemRaw, p.localCapacity, world);
+    int32_t* outRow = p.out + static_cast<int64_t>(row) * kTopK;
+    // Every global read below (lengths, epoch, logits) comes after the wait, so a programmatic early launch is safe
+    // whatever the preceding kernel writes.
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    int32_t const kvLen = p.kvLens[0];
+    int32_t const epoch = *reinterpret_cast<int32_t volatile*>(&p.state[0]);
+    int32_t const buf = epoch % kNumLamportBuffers;
+    int32_t const rowLen = max(0, kvLen - p.nextN + row + 1);
+    // Identity launch (every row has at most kTopK tokens; the same decision on every rank): no exchange, and
+    // neither the workspace nor the epoch is touched, so identity launches never shorten the reuse chain of the
+    // Lamport buffers (see crossRankTopk.h).
+    bool const exchange = kvLen - p.nextN + p.numRows > kTopK;
+    __syncthreads();
+
+    if (rowLen <= kTopK)
+    {
+        // Every rank agrees on this branch: identity selection, no exchange for this row.
+        if (src == 0)
+        {
+            for (int32_t j = threadIdx.x; j < kTopK; j += kThreads)
+            {
+                outRow[j] = j < rowLen ? j : -1;
+            }
+        }
+    }
+    else
+    {
+        int32_t const localLen = min(localLength(rowLen, rank, world), p.localCapacity);
+        float const* rowLogits = p.logits + static_cast<int64_t>(row) * p.logitsStride;
+
+        // ---- Round 1: this source's local top keys ---------------------------------------------------------
+        int32_t const n = loadSource(s, rowLogits, localLen, src, true);
+        __syncthreads();
+        markPhase(p, 1);
+        uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+        uint32_t const need = static_cast<uint32_t>(roundNeed(world));
+        SelectOut sel{0u, 0u, 0u, static_cast<uint32_t>(n), true, false, 0};
+        uint32_t selLo = 0u;
+        int32_t passInfo = 0;
+        if (n > static_cast<int32_t>(cap))
+        {
+            // The cut lies among the keys of one sign (see kPositiveKey); the digits span that sign's key range.
+            SignRange const r = signRange(readStats(s.stats), need);
+            selLo = r.lo;
+            SmemKeys const keys{s.chunk, nullptr};
+            sel = blockSelect(s, keys, n, r.lo, r.hi, need, cap, r.lo, r.bits, debugSlots(p, 32), r.aboveHi);
+            passInfo = sel.passes + (r.negative ? 10 : 0);
+        }
+        markPhase(p, 2);
+        uint32_t const sent = sendSource<false>(s, sel, n, selLo, uint64_t{kEmptyWord}, cap,
+            p.mcWorkspace + round1Offset(buf, rank * kSourcesPerRank + src, row), src, rank, world,
+            reinterpret_cast<uint2*>(s.mkey), debugSlots(p, 0));
+        putDebug(debugSlots(p, 0), 46, sent);
+        putDebug(debugSlots(p, 0), 47, passInfo);
+        markPhase(p, 3);
+
+        // ---- Merge: one CTA per row and rank ---------------------------------------------------------------
+        if (src == 0)
+        {
+            WaitCtx const wait{p.state, p.timeoutNs, rank};
+            mergeRow(p, s, segSlots, segBegin, buf, row, rowLen, localLen, rowLogits, outRow, wait);
+            markPhase(p, 7);
+        }
+    }
+    if (!exchange)
+    {
+        return;
+    }
+
+    // The other CTAs reset the buffer of the previous exchange launch, whose candidates every rank has consumed: a
+    // peer writes it again only two exchange launches later, after this launch completed on this rank.
+    if (src != 0)
+    {
+        clearBuffer(p, (epoch + kNumLamportBuffers - 1) % kNumLamportBuffers, row * (kSourcesPerRank - 1) + src - 1,
+            static_cast<int32_t>(gridDim.y) * (kSourcesPerRank - 1));
+    }
+
+    // Advance the epoch once every CTA of this launch has read it.
+    __syncthreads();
+    markPhase(p, 8);
+    if (threadIdx.x == 0)
+    {
+        __threadfence();
+        int32_t const numCtas = static_cast<int32_t>(gridDim.x * gridDim.y);
+        int32_t const arrived = atomicAdd(&p.state[1], 1);
+        if (arrived == numCtas - 1)
+        {
+            p.state[1] = 0;
+            __threadfence();
+            atomicAdd(&p.state[0], 1);
+        }
+    }
+}
+
+#endif
+
+} // namespace
+
+size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize)
+{
+    return smemBytes(localCapacity, worldSize);
+}
+
+int32_t cpTopkMaxLocalCapacity(int32_t worldSize)
+{
+    static int32_t cached[kMaxRanks + 1] = {};
+    if (worldSize < 1 || worldSize > kMaxRanks)
+    {
+        return 0;
+    }
+    if (cached[worldSize] == 0)
+    {
+        int device = 0;
+        cudaGetDevice(&device);
+        int maxSmemOptin = 0;
+        cudaDeviceGetAttribute(&maxSmemOptin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+        int32_t cap = 0;
+        for (int32_t c = 1 << 22; c >= kPageTokens * kSourcesPerRank; c -= kPageTokens * kSourcesPerRank)
+        {
+            if (smemBytes(c, worldSize) + kStaticSmemBytes <= static_cast<size_t>(maxSmemOptin))
+            {
+                cap = c;
+                break;
+            }
+        }
+        cached[worldSize] = cap;
+    }
+    return cached[worldSize];
+}
+
+int64_t cpTopkDefaultTimeoutMs()
+{
+    static int64_t const timeoutMs = []()
+    {
+        char const* env = std::getenv("TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS");
+        if (env == nullptr || env[0] == '\0')
+        {
+            return kDefaultTimeoutMs;
+        }
+        char* end = nullptr;
+        long long const v = std::strtoll(env, &end, 10);
+        if (end == env || *end != '\0' || v <= 0)
+        {
+            throw std::invalid_argument(
+                std::string("glm_moe_dsa_cross_rank_topk: bad TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS=")
+                + env);
+        }
+        return static_cast<int64_t>(v);
+    }();
+    return timeoutMs;
+}
+
+void invokeCpTopk(CpTopkParams const& inParams, cudaStream_t stream)
+{
+    if (inParams.numRows <= 0)
+    {
+        return;
+    }
+    CpTopkParams params = inParams;
+    if (params.timeoutNs == 0)
+    {
+        params.timeoutNs = static_cast<uint64_t>(cpTopkDefaultTimeoutMs()) * 1000000ull;
+    }
+    if (params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks || params.rank < 0
+        || params.rank >= params.worldSize)
+    {
+        throw std::invalid_argument("glm_moe_dsa_cross_rank_topk: unsupported rows/world/rank");
+    }
+    size_t const smem = smemBytes(params.localCapacity, params.worldSize);
+    // The attribute is per function and device; raise it only when a larger footprint is requested.
+    static thread_local size_t configuredSmem[64] = {};
+    int device = 0;
+    cudaGetDevice(&device);
+    cudaError_t err = cudaSuccess;
+    if (device >= 64 || configuredSmem[device] < smem)
+    {
+        err = cudaFuncSetAttribute(cpTopkKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem));
+        if (err != cudaSuccess)
+        {
+            throw std::runtime_error(
+                std::string("glm_moe_dsa_cross_rank_topk: smem attribute: ") + cudaGetErrorString(err));
+        }
+        if (device < 64)
+        {
+            configuredSmem[device] = smem;
+        }
+    }
+    // Programmatic dependent launch as for the other TensorRT-LLM kernels (on unless TRTLLM_ENABLE_PDL=0).
+    static bool const enablePdl = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && env[0] == '0' && env[1] == '\0');
+    }();
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSourcesPerRank, params.numRows);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    err = cudaLaunchKernelEx(&config, cpTopkKernel, params);
+    if (err == cudaSuccess)
+    {
+        err = cudaGetLastError();
+    }
+    if (err != cudaSuccess)
+    {
+        throw std::runtime_error(std::string("glm_moe_dsa_cross_rank_topk: launch: ") + cudaGetErrorString(err));
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::cross_rank_topk
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h
new file mode 100644
index 000000000000..2ff17f85e48e
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h
@@ -0,0 +1,162 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
+#include <cuda_runtime_api.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::cross_rank_topk
+{
+
+// Context-partitioned DSA indexer top-k.
+//
+// The indexer context of one generation request is partitioned over the ranks of a tensor-parallel group by whole
+// 64-token pages: global page p belongs to rank p % worldSize, and a rank's local column c maps to global token
+// ((c / 64) * worldSize + rank) * 64 + c % 64. Every rank scores only its own pages (the caller supplies the local
+// logits). Within a rank, each verify row is handled by kSourcesPerRank CTAs; CTA s owns one 8-token block of every
+// local page (block (s + page) % kSourcesPerRank), so the (rank, CTA) "sources" interleave the context finely and a
+// fully selected page is spread over all sources of its rank. Each source
+//   1. selects its local top keys (kSourceBudget at 8 ranks, see below) with a block radix select,
+//   2. broadcasts them as (order-preserving key, global token) pairs, plus a header (count, exclusive bound of the
+//      keys it withheld), through a multicast Lamport workspace.
+// One CTA per row and rank then merges the candidates of all sources into the exact top-kTopK of the row,
+// identically on every rank. A source that withheld a key above the merged threshold makes every rank run a
+// second round in which that source's top keys between the threshold and its bound are sent (at most
+// kRound2Capacity, by the merge CTA of the source's rank, from the local logits), so the result is the tie-aware
+// exact top-kTopK for any input. In the second round, a source whose cut falls inside a run of equal keys sends the
+// keys above the cut and the first equal ones in source order up to its capacity (round 1 sends only the keys above
+// such a cut, so the two rounds never send the same key).
+// Rows of at most kTopK valid tokens produce the identity selection followed by -1, without an exchange.
+// A launch whose rows all have at most kTopK valid tokens (the same on every rank: the lengths come from the shared
+// KV length) is an identity launch: it touches neither the workspace nor the epoch. Every other launch is an
+// exchange launch. Exchange launches rotate kNumLamportBuffers buffers by a device epoch that counts exchange
+// launches; in exchange launch x the CTAs that do not merge reset, on their own rank, the buffer of exchange launch
+// x - 1. Every rank's merges of x - 1 have read their copy of it, and a peer writes it again only in exchange
+// launch x + 2, which starts after the peer's merge of x + 1 received this rank's sources of x + 1, i.e. after this
+// rank completed launch x (identity launches in between neither write nor wait, so they do not shorten this chain).
+//
+// A merge that waits longer than the timeout (CpTopkParams::timeoutNs, %globaltimer) for a peer's words records the
+// wait in the state's error word and aborts the kernel with a trap: the exchange cannot complete correctly any more,
+// and a late peer write could land in a buffer this rank reuses, so the run fails instead of continuing.
+
+//! Selected indices per verify row (DSA index_topk).
+constexpr int32_t kTopK = 2048;
+//! Ranks of the context-parallel group.
+constexpr int32_t kMaxRanks = 8;
+//! Verify rows per call.
+constexpr int32_t kMaxRows = 8;
+//! Tokens per indexer K page; the context is partitioned by whole pages.
+constexpr int32_t kPageTokens = 64;
+//! CTAs (exchange sources) per rank and verify row.
+constexpr int32_t kSourcesPerRank = 8;
+//! Exchange sources per row.
+constexpr int32_t kMaxSources = kMaxRanks * kSourcesPerRank;
+//! Round-1 slots per (source, row).
+constexpr int32_t kSourceCapacity = 256;
+//! Keys each saturated source aims to send in round 1 at 8 ranks (64 sources): three times the even share of kTopK
+//! (the largest share of the exact top-2048 that one source of this layout held on real GLM-5.2 indexer logits was
+//! 69). A source stops refining as soon as its keys at or above a bin edge number at most four times the even
+//! share, so it sends between 96 and 128 keys at 64 sources. Other group sizes scale both with ceil(kTopK / sources).
+constexpr int32_t kSourceBudget = 96;
+//! Round-2 slots per (source, row).
+constexpr int32_t kRound2Capacity = 2048;
+//! Threads per CTA.
+constexpr int32_t kThreads = 1024;
+//! Lamport buffers rotated by the device epoch.
+constexpr int32_t kNumLamportBuffers = 3;
+//! 32-bit word that marks an empty Lamport slot; never a valid key, token index, count or bound.
+constexpr uint32_t kEmptyWord = 0xFFFFFFFFu;
+
+//! 32-bit words of one (source, row) slot: a 16-byte header (count, bound, two unused words) then two words per
+//! candidate, so candidates can be written and read with 16-byte vector accesses.
+constexpr int64_t kRound1SlotWords = 4 + 2 * int64_t{kSourceCapacity};
+constexpr int64_t kRound2SlotWords = 4 + 2 * int64_t{kRound2Capacity};
+//! 32-bit words of one Lamport buffer (round-1 region, then round-2 region).
+constexpr int64_t kBufferWords = int64_t{kMaxSources} * kMaxRows * (kRound1SlotWords + kRound2SlotWords);
+//! 32-bit words of the whole exchange workspace; every word must start as kEmptyWord.
+constexpr int64_t kWorkspaceWords = kNumLamportBuffers * kBufferWords;
+//! int32 words of the per-rank device state: epoch (exchange launches), CTA arrivals, second rounds taken, error
+//! word (nonzero = the wait that timed out, written right before the kernel aborts).
+constexpr int32_t kStateWords = 4;
+//! Default wait for a peer's Lamport words before the kernel aborts
+//! (TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS overrides it).
+constexpr int64_t kDefaultTimeoutMs = 120000;
+
+//! Values of the error word.
+enum CpTopkWait : int32_t
+{
+    kWaitRound1Header = 1,
+    kWaitRound1Candidates = 2,
+    kWaitRound2Header = 3,
+    kWaitRound2Candidates = 4
+};
+
+struct CpTopkParams
+{
+    //! [numRows, logitsStride] fp32 logits of this rank's local columns.
+    float const* logits;
+    int64_t logitsStride;
+    //! Local columns the logits view holds (each row's local length is clamped to it).
+    int32_t localCapacity;
+    //! [1] int32 global indexer KV length of the request; row r has kvLen - nextN + r + 1 valid tokens.
+    int32_t const* kvLens;
+    //! [numRows, kTopK] int32 output: global token indices, -1 for unused slots.
+    int32_t* out;
+    int32_t numRows;
+    int32_t nextN;
+    int32_t rank;
+    int32_t worldSize;
+    //! Multicast address of the exchange workspace (kWorkspaceWords words, identical layout on every rank).
+    uint32_t* mcWorkspace;
+    //! This rank's unicast address of the same workspace.
+    uint32_t* ucWorkspace;
+    //! [kStateWords] int32 per-rank device state, zero-initialized.
+    int32_t* state;
+    //! Optional [numRows, kSourcesPerRank, kDebugSlots] int64 phase timestamps (diagnosis only; nullptr = off).
+    int64_t* debugTimes = nullptr;
+    //! Longest wait for a peer's words before the kernel aborts (ns); 0 = the default (kDefaultTimeoutMs, or
+    //! TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS).
+    uint64_t timeoutNs = 0;
+};
+
+//! Timestamps per CTA in CpTopkParams::debugTimes: [0, 16) %globaltimer (ns), [16, 32) clock64 per phase,
+//! [32, 42) local select (keys, then clock64 at the start of / after the histogram of / after the bin search of
+//! each pass), [42, 45) send (counter ready, keys taken, stores issued), 45 candidates polled before the key-range
+//! reduction, 46 round-1 keys sent by this source, 47 local select passes (+10 when the cut was searched among the
+//! negative keys), [48, 58) merge select (as the local select), [58, 62) emission (threshold decided, positions
+//! scanned, row assembled, row written), 62 merge candidates, 63 merge select passes (+10 likewise); clock64 and
+//! values of thread 0.
+constexpr int32_t kDebugSlots = 64;
+
+//! Effective wait timeout of a launch with params.timeoutNs == 0 (ms).
+int64_t cpTopkDefaultTimeoutMs();
+
+//! Dynamic shared memory the kernel needs (bytes).
+size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize);
+
+//! Largest local capacity the kernel accepts for a world size.
+int32_t cpTopkMaxLocalCapacity(int32_t worldSize);
+
+void invokeCpTopk(CpTopkParams const& params, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::cross_rank_topk
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu
new file mode 100644
index 000000000000..a0b9671537c5
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu
@@ -0,0 +1,2420 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdlib>
+#include <stdexcept>
+#include <string>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge
+{
+namespace
+{
+
+constexpr int32_t kBins = 2048;
+static_assert(kBins >= kTopK, "the output row is staged in the histogram");
+constexpr int32_t kWarps = kThreads / 32;
+//! Per-warp key statistics kept in shared memory (see KeyStats).
+constexpr int32_t kNumKeyStats = 5;
+//! Tokens of each page one source holds.
+constexpr int32_t kBlockTokens = kPageTokens / kSourcesPerRank;
+//! CTAs of one row's cluster.
+constexpr int32_t kClusterCtas = kSourcesPerRank;
+//! Histogram bins per group of the cluster bin search (one warp lane per bin).
+constexpr int32_t kGroupBins = 32;
+constexpr int32_t kMaxGroups = kBins / kGroupBins;
+//! Static shared memory of the kernel (round-2 segment tables and the cluster exchange records), kept out of the
+//! dynamic budget.
+constexpr size_t kStaticSmemBytes = 8192;
+
+// Records the CTAs of a row's cluster send each other with st.async (static, so every CTA has them at the same
+// address). Exchange x of a row completes on barrier bar[x % 2], phase x / 2; the record sets alternate by pass
+// parity and are written again only after every CTA has completed the exchange in between.
+struct ClusterRecords
+{
+    //! Transaction barriers of the exchanges (arrival count 1: this CTA's expect-tx).
+    uint64_t bar[2];
+    //! [pass parity][source CTA][group] histogram totals per group of kGroupBins bins.
+    uint32_t groups[2][kClusterCtas][kMaxGroups];
+    //! [pass parity][source CTA][bin] counts of the selected group's bins, ascending.
+    uint32_t bins[2][kClusterCtas][kGroupBins];
+};
+
+// Scalar slots in shared memory.
+enum ScalarSlot : int32_t
+{
+    kSlotBin = 0,
+    kSlotAbove,
+    kSlotBinCount,
+    kSlotTotal,
+    kSlotCounter2,
+    kSlotSrcCount,                                       // [kMaxSources]
+    kSlotSrcBound = kSlotSrcCount + kMaxSources,         // [kMaxSources]
+    kSlotSrcLen = kSlotSrcBound + kMaxSources,           // [kMaxSources]
+    kSlotSrcOffset = kSlotSrcLen + kMaxSources,          // [kMaxSources + 1]
+    kSlotCountBefore = kSlotSrcOffset + kMaxSources + 1, // candidates of the cluster CTAs before this one
+    kSlotCountMine,                                      // candidates of this cluster CTA
+    kSlotCountTotal,                                     // candidates of the row
+    kSlotPositivesBefore,                                // candidates >= +0.0 of the cluster CTAs before this one
+    kNumScalarSlots
+};
+
+__host__ __device__ inline int32_t numSources(int32_t world)
+{
+    return world * kSourcesPerRank;
+}
+
+// Even share of kTopK per source.
+__host__ __device__ inline int32_t evenShare(int32_t world)
+{
+    return (kTopK + numSources(world) - 1) / numSources(world);
+}
+
+// Most keys a source sends in round 1 (an even number).
+__host__ __device__ inline int32_t roundCap(int32_t world)
+{
+    return min(kSourceCapacity, 4 * evenShare(world));
+}
+
+// Keys one source can hold: one block of every local page.
+__host__ __device__ inline int32_t chunkCapacity(int32_t localCapacity)
+{
+    int32_t const pages = (localCapacity + kPageTokens - 1) / kPageTokens;
+    return pages * kBlockTokens;
+}
+
+__host__ __device__ inline int32_t mergeCapacity(int32_t world)
+{
+    return numSources(world) * roundCap(world);
+}
+
+__host__ __device__ inline size_t smemBytes(int32_t localCapacity, int32_t world)
+{
+    size_t words = 2 * size_t{kBins} + 2 * size_t{kWarps} + size_t{kNumKeyStats} * kWarps + size_t{kNumScalarSlots};
+    words = (words + 3) / 4 * 4;
+    words += static_cast<size_t>(chunkCapacity(localCapacity)) + 2 * static_cast<size_t>(mergeCapacity(world));
+    return words * sizeof(uint32_t);
+}
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900)
+// The kernel targets SM90+ (multicast Lamport exchange, thread-block clusters).
+__global__ void __launch_bounds__(kThreads, 1) cpTopkZ9Kernel(CpTopkParams p)
+{
+    __trap();
+}
+#else
+
+//! Bits of one radix-select digit (kBins bins).
+constexpr int32_t kDigitBits = 11;
+//! Order-preserving key of +0.0: the keys of non-negative values are at least this, those of negative values below.
+//! Real indexer logits are mostly negative, and the cut of a row (and of many sources) often lies among small
+//! negative values while the largest keys are positive: one key range across zero spans ~2^31 key units, where a
+//! radix digit has no resolution. The selects therefore work on one sign at a time: with fewer than `need` keys at
+//! or above kPositiveKey these are all taken and the cut is searched among the negative keys only, else among the
+//! non-negative keys only; each range is narrow enough for one or two passes.
+constexpr uint32_t kPositiveKey = 0x80000000u;
+//! Independent 16-byte candidate loads one merge thread issues before it waits on any of them.
+constexpr int32_t kPollBatch = 4;
+//! Polls of a still-empty Lamport word between two reads of the global timer.
+constexpr int32_t kPollsPerTimerRead = 256;
+//! Warps of the cluster merge: warps [0, kPollWarps) poll the round-1 candidates of this CTA's sources (one 16-byte
+//! vector per thread: at most 520 vectors per CTA, at world sizes 5 to 7), the next kHeaderWarps warps poll the
+//! round-1 headers of every source (one per thread), and the remaining kResetWarps warps reset the previous exchange
+//! launch's buffer meanwhile.
+constexpr int32_t kPollWarps = 17;
+constexpr int32_t kHeaderWarps = 2;
+constexpr int32_t kFirstResetWarp = kPollWarps + kHeaderWarps;
+constexpr int32_t kResetWarps = kWarps - kFirstResetWarp;
+static_assert(kHeaderWarps * 32 >= kMaxSources, "one header per header-warp thread");
+
+// Keys a saturated source aims to send in round 1 (kSourceBudget at 8 ranks).
+__device__ __forceinline__ int32_t roundNeed(int32_t world)
+{
+    return min(kSourceCapacity, 3 * evenShare(world));
+}
+
+// Tokens of `rank` among the global tokens [0, n).
+__host__ __device__ __forceinline__ int32_t localLength(int32_t n, int32_t rank, int32_t world)
+{
+    if (n <= 0)
+    {
+        return 0;
+    }
+    int32_t const fullPages = n / kPageTokens;
+    int32_t const rem = n % kPageTokens;
+    int32_t const ownedFull = (fullPages - rank + world - 1) / world;
+    int32_t len = ownedFull * kPageTokens;
+    if (rem > 0 && fullPages % world == rank)
+    {
+        len += rem;
+    }
+    return len;
+}
+
+// Source `src` of a rank holds one block of kBlockTokens consecutive tokens of every local page: block
+// (src + page) % kSourcesPerRank. A fully selected page is thus spread over all sources of its rank.
+// Keys of source `src` among `localLen` local columns (only the last local page can be partial).
+__host__ __device__ __forceinline__ int32_t sourceLength(int32_t localLen, int32_t src)
+{
+    int32_t const pages = (localLen + kPageTokens - 1) / kPageTokens;
+    if (pages == 0)
+    {
+        return 0;
+    }
+    int32_t const last = pages - 1;
+    int32_t const tail = localLen - last * kPageTokens - ((src + last) % kSourcesPerRank) * kBlockTokens;
+    return last * kBlockTokens + min(max(tail, 0), kBlockTokens);
+}
+
+// Global token of a local column under the page-interleaved partition.
+__device__ __forceinline__ uint32_t localToGlobal(int32_t col, int32_t rank, int32_t world)
+{
+    int32_t const page = col / kPageTokens;
+    return static_cast<uint32_t>((page * world + rank) * kPageTokens + (col % kPageTokens));
+}
+
+struct SmemLayout
+{
+    uint32_t* hist;     // [2 * kBins] radix histograms of alternate passes; the first one stages the output row
+    uint32_t* warpSums; // [2 * kWarps]
+    uint32_t* stats;    // [kNumKeyStats * kWarps] per-warp key statistics
+    uint32_t* scal;     // [kNumScalarSlots]
+    uint32_t* chunk;    // [chunkCapacity] order-preserving keys of this source
+    uint32_t* mkey;     // [mergeCapacity] staged round-1 pairs, then the merge candidates (merge CTA)
+    uint32_t* midx;     // [mergeCapacity]
+};
+
+__device__ inline SmemLayout carveSmem(uint32_t* base, int32_t localCapacity, int32_t world)
+{
+    SmemLayout s;
+    uint32_t* p = base;
+    s.hist = p;
+    p += 2 * kBins;
+    s.warpSums = p;
+    p += 2 * kWarps;
+    s.stats = p;
+    p += kNumKeyStats * kWarps;
+    s.scal = p;
+    p += kNumScalarSlots;
+    p = base + ((p - base) + 3) / 4 * 4;
+    s.chunk = p;
+    p += chunkCapacity(localCapacity);
+    s.mkey = p;
+    p += mergeCapacity(world);
+    s.midx = p;
+    return s;
+}
+
+__device__ __forceinline__ uint32_t orderedKey(float v)
+{
+    uint32_t const b = __float_as_uint(v);
+    uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u);
+    // NaN keys above +inf are clamped so that no key, and no exclusive bound key + 1, equals the empty marker.
+    return k < kEmptyWord - 1u ? k : kEmptyWord - 2u;
+}
+
+// Bits of x (0 for 0): the digit range of keys in [base, base + x].
+__device__ __forceinline__ int32_t bitLength(uint32_t x)
+{
+    return x == 0u ? 0 : 32 - __clz(x);
+}
+
+// Statistics of a set of keys: extremes overall and per sign, and the count of non-negative keys.
+struct KeyStats
+{
+    uint32_t maxKey = 0u;         // largest key (0 when none)
+    uint32_t minKey = kEmptyWord; // smallest key (kEmptyWord when none)
+    uint32_t positives = 0u;      // keys >= kPositiveKey
+    uint32_t maxNeg = 0u;         // largest key below kPositiveKey (0 when none)
+    uint32_t minPos = kEmptyWord; // smallest key >= kPositiveKey (kEmptyWord when none)
+
+    __device__ __forceinline__ void add(uint32_t k)
+    {
+        bool const pos = k >= kPositiveKey;
+        maxKey = max(maxKey, k);
+        minKey = min(minKey, k);
+        positives += pos ? 1u : 0u;
+        maxNeg = pos ? maxNeg : max(maxNeg, k);
+        minPos = pos ? min(minPos, k) : minPos;
+    }
+};
+
+// Reduce a thread's statistics over its warp; lane 0 stores them in the per-warp array (read with readStats after
+// the next block barrier).
+__device__ __forceinline__ void storeWarpStats(KeyStats const& t, uint32_t* stats)
+{
+    uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, t.maxKey);
+    uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, t.minKey);
+    uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, t.positives);
+    uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg);
+    uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, t.minPos);
+    if (threadIdx.x % 32 == 0)
+    {
+        int32_t const w = threadIdx.x / 32;
+        stats[w] = v0;
+        stats[kWarps + w] = v1;
+        stats[2 * kWarps + w] = v2;
+        stats[3 * kWarps + w] = v3;
+        stats[4 * kWarps + w] = v4;
+    }
+}
+
+// Block statistics from the per-warp array (every warp reduces it; no barrier).
+__device__ __forceinline__ KeyStats readStats(uint32_t const* stats)
+{
+    int32_t const lane = threadIdx.x % 32;
+    KeyStats b;
+    b.maxKey = __reduce_max_sync(0xFFFFFFFFu, stats[lane]);
+    b.minKey = __reduce_min_sync(0xFFFFFFFFu, stats[kWarps + lane]);
+    b.positives = __reduce_add_sync(0xFFFFFFFFu, stats[2 * kWarps + lane]);
+    b.maxNeg = __reduce_max_sync(0xFFFFFFFFu, stats[3 * kWarps + lane]);
+    b.minPos = __reduce_min_sync(0xFFFFFFFFu, stats[4 * kWarps + lane]);
+    return b;
+}
+
+// Digit setup of a select of the need-th largest key: over the non-negative keys when there are at least `need` of
+// them, else over the negative keys with every non-negative key counted as taken (`aboveHi`).
+struct SignRange
+{
+    uint32_t lo;      // smallest key of the range (also the digit base)
+    uint64_t hi;      // exclusive end of the histogrammed range
+    int32_t bits;     // digit range of the range
+    uint32_t aboveHi; // keys at or above hi, all taken
+    bool negative;    // the cut is among the negative keys
+};
+
+__device__ __forceinline__ SignRange signRange(KeyStats const& st, uint32_t need)
+{
+    if (st.positives >= need || st.positives == 0u)
+    {
+        uint32_t const lo = st.positives >= need ? st.minPos : st.minKey;
+        return SignRange{lo, uint64_t{kEmptyWord}, bitLength(st.maxKey - lo), 0u, st.positives == 0u};
+    }
+    return SignRange{st.minKey, uint64_t{kPositiveKey}, bitLength(st.maxNeg - st.minKey), st.positives, true};
+}
+
+__device__ __forceinline__ int64_t round1Offset(int32_t buf, int32_t src, int32_t row)
+{
+    return int64_t{buf} * kBufferWords + (int64_t{src} * kMaxRows + row) * kRound1SlotWords;
+}
+
+__device__ __forceinline__ int64_t round2Offset(int32_t buf, int32_t src, int32_t row)
+{
+    return int64_t{buf} * kBufferWords + int64_t{kMaxSources} * kMaxRows * kRound1SlotWords
+        + (int64_t{src} * kMaxRows + row) * kRound2SlotWords;
+}
+
+__device__ __forceinline__ uint2 loadVolatile2(uint32_t const* ptr)
+{
+    uint2 v;
+    asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(ptr) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadVolatile4(uint32_t const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ void store2(uint32_t* ptr, uint32_t a, uint32_t b)
+{
+    asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(ptr), "r"(a), "r"(b) : "memory");
+}
+
+__device__ __forceinline__ void store4(uint32_t* ptr, uint4 v)
+{
+    asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)
+                 : "memory");
+}
+
+__device__ __forceinline__ bool isEmpty(uint2 v)
+{
+    return v.x == kEmptyWord || v.y == kEmptyWord;
+}
+
+// Whether the first pair of a 16-byte candidate vector, and the second one when `second`, were written.
+__device__ __forceinline__ bool vectorReady(uint4 v, bool second)
+{
+    return v.x != kEmptyWord && v.y != kEmptyWord && (!second || (v.z != kEmptyWord && v.w != kEmptyWord));
+}
+
+__device__ __forceinline__ uint64_t globalTimerNs()
+{
+    uint64_t t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return t;
+}
+
+// What a polling thread needs to give up: the error word, the timeout and this rank (for the report).
+struct WaitCtx
+{
+    int32_t* state;
+    uint64_t timeoutNs;
+    int32_t rank;
+};
+
+// A peer did not deliver within the timeout. The selection cannot be completed correctly any more, and the peer's
+// late writes could land in a buffer this rank resets or reuses, so record which wait failed and abort the kernel:
+// the resulting CUDA error fails the run instead of letting any rank continue with a partial selection.
+__device__ __noinline__ void abortOnTimeout(WaitCtx const& w, int32_t what)
+{
+    bool const first = atomicExch(&w.state[3], what) == 0;
+    __threadfence_system();
+    if (first)
+    {
+        printf(
+            "[glm_moe_dsa] crossRankTopkClusterMerge: rank %d, row %d: wait %d for a peer exceeded %llu ms; aborting "
+            "the kernel\n",
+            w.rank, static_cast<int32_t>(blockIdx.y), what, static_cast<unsigned long long>(w.timeoutNs / 1000000ull));
+    }
+    __trap();
+}
+
+// Spin until both words of a Lamport pair were written; aborts the kernel after the timeout.
+__device__ uint2 pollPair(uint32_t const* ptr, WaitCtx const& w, int32_t what)
+{
+    uint2 v = loadVolatile2(ptr);
+    uint64_t start = 0;
+    for (int32_t polls = 1; isEmpty(v); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, what);
+            }
+        }
+        v = loadVolatile2(ptr);
+    }
+    return v;
+}
+
+// Spin until a 16-byte candidate vector was written (see vectorReady); aborts the kernel after the timeout.
+__device__ uint4 pollVector(uint32_t const* ptr, bool second, WaitCtx const& w, int32_t what)
+{
+    uint4 v = loadVolatile4(ptr);
+    uint64_t start = 0;
+    for (int32_t polls = 1; !vectorReady(v, second); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, what);
+            }
+        }
+        v = loadVolatile4(ptr);
+    }
+    return v;
+}
+
+// Block-wide inclusive scan of one value per thread. Returns the inclusive prefix; `total` gets the block sum.
+__device__ uint32_t blockInclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    uint32_t x = v;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    if (lane == 31)
+    {
+        warpSums[warp] = x;
+    }
+    __syncthreads();
+    if (warp == 0)
+    {
+        uint32_t w = warpSums[lane];
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, w, o);
+            if (lane >= o)
+            {
+                w += y;
+            }
+        }
+        warpSums[kWarps + lane] = w;
+    }
+    __syncthreads();
+    total = warpSums[2 * kWarps - 1];
+    uint32_t const before = warp > 0 ? warpSums[kWarps + warp - 1] : 0u;
+    return x + before;
+}
+
+// Exclusive block-wide scan with a single barrier: every warp scans the per-warp sums itself. Returns this
+// thread's exclusive prefix; `total` gets the block sum. warpSums must not be rewritten before the next barrier.
+__device__ uint32_t blockExclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    uint32_t x = v;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    if (lane == 31)
+    {
+        warpSums[warp] = x;
+    }
+    __syncthreads();
+    uint32_t const w = warpSums[lane];
+    uint32_t y = w;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const z = __shfl_up_sync(0xFFFFFFFFu, y, o);
+        if (lane >= o)
+        {
+            y += z;
+        }
+    }
+    total = __shfl_sync(0xFFFFFFFFu, y, 31);
+    return __shfl_sync(0xFFFFFFFFu, y - w, warp) + x - v;
+}
+
+// Find, counting from the top bin, the bin of `hist` that holds the target-th largest key, with two barriers:
+// per-warp sums, then each warp scans them and only the warp holding the target scans its own bins. Returns the
+// total count (every thread); when target <= total, the bin, the count above it and its own count go to the
+// scalar slots.
+__device__ uint32_t findBin(SmemLayout const& s, uint32_t const* hist, int32_t nbins, uint32_t target)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    int32_t const perThread = (nbins + kThreads - 1) / kThreads;
+    int32_t const top = nbins - 1 - perThread * static_cast<int32_t>(threadIdx.x);
+    uint32_t mine = 0;
+    for (int32_t i = 0; i < perThread && top - i >= 0; ++i)
+    {
+        mine += hist[top - i];
+    }
+    uint32_t const warpTotal = __reduce_add_sync(0xFFFFFFFFu, mine);
+    if (lane == 0)
+    {
+        s.warpSums[warp] = warpTotal;
+    }
+    __syncthreads();
+    uint32_t const w = s.warpSums[lane];
+    uint32_t x = w;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    uint32_t const total = __shfl_sync(0xFFFFFFFFu, x, 31);
+    uint32_t const warpBefore = __shfl_sync(0xFFFFFFFFu, x - w, warp);
+    if (warpBefore < target && target <= warpBefore + warpTotal)
+    {
+        uint32_t t = mine;
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, t, o);
+            if (lane >= o)
+            {
+                t += y;
+            }
+        }
+        uint32_t acc = warpBefore + t - mine;
+        if (acc < target && target <= acc + mine)
+        {
+            for (int32_t i = 0; i < perThread && top - i >= 0; ++i)
+            {
+                uint32_t const c = hist[top - i];
+                if (acc + c >= target)
+                {
+                    s.scal[kSlotBin] = static_cast<uint32_t>(top - i);
+                    s.scal[kSlotAbove] = acc;
+                    s.scal[kSlotBinCount] = c;
+                    break;
+                }
+                acc += c;
+            }
+        }
+    }
+    __syncthreads();
+    return total;
+}
+
+struct SelectOut
+{
+    uint32_t edge;     // lower edge of the selected bin (the exact key after the last digit)
+    uint32_t above;    // predicate keys strictly above the selected bin
+    uint32_t binCount; // predicate keys in the selected bin
+    uint32_t total;    // predicate keys
+    bool all;          // take every predicate key (fewer than needed, or all fit the capacity)
+    bool full;         // exact mode: the keys at or above `edge` are exactly the `need` largest
+    int32_t passes;    // radix passes run
+};
+
+// Candidate keys and token indices held in shared memory (16-byte aligned, readable in groups of four up to the
+// next multiple of four).
+struct SmemKeys
+{
+    static constexpr bool kVector = true;
+    uint32_t const* keys;
+    uint32_t const* idx;
+
+    __device__ uint4 keys4(int32_t j) const
+    {
+        return *reinterpret_cast<uint4 const*>(keys + j);
+    }
+
+    __device__ uint4 index4(int32_t j) const
+    {
+        return *reinterpret_cast<uint4 const*>(idx + j);
+    }
+};
+
+// Merge candidates in shared memory followed by round-2 candidates read from their global Lamport slots.
+struct MergeKeys
+{
+    static constexpr bool kVector = false;
+    uint32_t const* smemKeys;
+    uint32_t const* smemIdx;
+    int32_t smemCount;
+    uint32_t const* const* slots; // round-2 slot of each segment
+    uint32_t const* segBegin;     // [numSegments + 1] candidate offsets
+    int32_t numSegments;
+
+    __device__ uint32_t const* pair(int32_t k) const
+    {
+        // Last segment whose begin is <= k.
+        int32_t lo = 0;
+        int32_t hi = numSegments - 1;
+        while (lo < hi)
+        {
+            int32_t const mid = (lo + hi + 1) / 2;
+            if (static_cast<int32_t>(segBegin[mid]) <= k)
+            {
+                lo = mid;
+            }
+            else
+            {
+                hi = mid - 1;
+            }
+        }
+        return slots[lo] + kSlotHeaderWords + 2 * (k - static_cast<int32_t>(segBegin[lo]));
+    }
+
+    __device__ uint32_t operator()(int32_t j) const
+    {
+        return j < smemCount ? smemKeys[j] : pair(j - smemCount)[0];
+    }
+
+    __device__ uint32_t index(int32_t j) const
+    {
+        return j < smemCount ? smemIdx[j] : pair(j - smemCount)[1];
+    }
+};
+
+// Debug slots of this CTA starting at `offset` (nullptr without a debug buffer).
+__device__ __forceinline__ int64_t* debugSlots(CpTopkParams const& p, int32_t offset)
+{
+    if (p.debugTimes == nullptr)
+    {
+        return nullptr;
+    }
+    return p.debugTimes + (static_cast<int64_t>(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots + offset;
+}
+
+// clock64 of thread 0 into a debug slot (no-op without a debug buffer).
+__device__ __forceinline__ void stamp(int64_t* dbg, int32_t slot)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[slot] = clock64();
+    }
+}
+
+// A value into a debug slot (thread 0; no-op without a debug buffer).
+__device__ __forceinline__ void putDebug(int64_t* dbg, int32_t slot, int64_t value)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[slot] = value;
+    }
+}
+
+// Phase timestamp for diagnosis (thread 0 of the CTA; no-op without a debug buffer).
+__device__ __forceinline__ void markPhase(CpTopkParams const& p, int32_t phase)
+{
+    if (p.debugTimes != nullptr && threadIdx.x == 0)
+    {
+        uint64_t t;
+        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+        int64_t* d = p.debugTimes + (static_cast<int64_t>(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots;
+        d[phase] = static_cast<int64_t>(t);
+        d[16 + phase] = static_cast<int64_t>(clock64());
+    }
+}
+
+// Histogram update of one key of a radix pass.
+__device__ __forceinline__ void histogramKey(uint32_t* hist, bool inRange, uint32_t k, uint32_t lo, uint64_t hi,
+    uint32_t base, uint32_t prefMask, uint32_t prefVal, int32_t shift, uint32_t binMask)
+{
+    uint32_t const d = k - base;
+    if (inRange && k >= lo && static_cast<uint64_t>(k) < hi && (d & prefMask) == prefVal)
+    {
+        atomicAdd(&hist[(d >> shift) & binMask], 1u);
+    }
+}
+
+// Block radix select of the need-th largest of the n keys at or above lo, of which `aboveHi` (< need) lie at or
+// above hi and are counted as taken without being histogrammed. Every key in [lo, hi) is at least `base`, and
+// key - base < 2^topBits; digits of up to kDigitBits bits of key - base are refined from the top. With cap > 0 the
+// select stops at the first bin edge whose keys at or above it number at most `cap`. With cap = 0 it runs to the
+// exact key, or stops early when the selected bin is wholly taken (`full`). Alternate passes use alternate
+// histograms, so the next one is cleared while the current one is searched. Counts in the result include aboveHi.
+template <typename KeyAt>
+__device__ SelectOut blockSelect(SmemLayout const& s, KeyAt const& keyAt, int32_t n, uint32_t lo, uint64_t hi,
+    uint32_t need, uint32_t cap, uint32_t base, int32_t topBits, int64_t* dbg = nullptr, uint32_t aboveHi = 0u)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[0] = n;
+    }
+    SelectOut out{lo, 0u, 0u, 0u, false, false, 0};
+    uint32_t prefMask = 0;
+    uint32_t prefVal = 0;
+    uint32_t above = aboveHi;
+    int32_t remaining = max(topBits, 1);
+    uint32_t* hist = s.hist;
+    uint32_t* histNext = s.hist + kBins;
+    for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads)
+    {
+        hist[b] = 0;
+    }
+    __syncthreads();
+#pragma unroll 1
+    for (int32_t pass = 0; remaining > 0; ++pass)
+    {
+        int32_t const width = min(kDigitBits, remaining);
+        int32_t const shift = remaining - width;
+        int32_t const nbins = 1 << width;
+        uint32_t const binMask = static_cast<uint32_t>(nbins - 1);
+        out.passes = pass + 1;
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[1 + 3 * pass] = clock64();
+        }
+        if constexpr (KeyAt::kVector)
+        {
+            for (int32_t j0 = 4 * static_cast<int32_t>(threadIdx.x); j0 < n; j0 += 4 * kThreads)
+            {
+                uint4 const q = keyAt.keys4(j0);
+                histogramKey(hist, true, q.x, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 1 < n, q.y, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 2 < n, q.z, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 3 < n, q.w, lo, hi, base, prefMask, prefVal, shift, binMask);
+            }
+        }
+        else
+        {
+            for (int32_t j = threadIdx.x; j < n; j += kThreads)
+            {
+                histogramKey(hist, true, keyAt(j), lo, hi, base, prefMask, prefVal, shift, binMask);
+            }
+        }
+        __syncthreads();
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[2 + 3 * pass] = clock64();
+        }
+        if (shift > 0)
+        {
+            for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads)
+            {
+                histNext[b] = 0;
+            }
+        }
+        // Find the bin holding the (need - above)-th largest key, counting from the top bin.
+        uint32_t const target = need - above;
+        uint32_t const total = findBin(s, hist, nbins, target);
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[3 + 3 * pass] = clock64();
+        }
+        if (pass == 0)
+        {
+            out.total = total + aboveHi;
+            if (total + aboveHi <= cap || target > total)
+            {
+                out.all = true;
+                out.above = total + aboveHi;
+                return out;
+            }
+        }
+        // The scalars are next written after two more barriers (the next histogram and scan).
+        uint32_t const bin = s.scal[kSlotBin];
+        above += s.scal[kSlotAbove];
+        out.binCount = s.scal[kSlotBinCount];
+        prefVal |= bin << shift;
+        prefMask |= binMask << shift;
+        out.edge = base + prefVal;
+        out.above = above;
+        remaining = shift;
+        uint32_t* const t = hist;
+        hist = histNext;
+        histNext = t;
+        if (cap > 0 && above + out.binCount <= cap)
+        {
+            return out;
+        }
+        if (cap == 0 && remaining > 0 && above + out.binCount == need)
+        {
+            out.full = true;
+            return out;
+        }
+    }
+    return out;
+}
+
+// Send one source's selected keys (keys in [lo, hi) of the chunk) to its slot; returns the count sent.
+// Takes every predicate key at or above the edge when they fit `cap`. Otherwise (more than `cap` keys tie at the
+// exact edge) it takes every key above the edge, plus, with kPartialTies, the first keys equal to it in source
+// order up to `cap` (any subset of equal keys is a valid tie-aware choice). Round 1 sends no partial ties, so the
+// keys a second round requests from a source (at most the edge) were never sent before.
+// The header carries the count and an exclusive upper bound of the keys left unsent (0 when none); in round 1
+// (!kPartialTies) also the key statistics of the keys sent (see kSlotHeaderWords).
+// With `stage` (round 1, at most kSourceCapacity keys), the pairs are gathered in shared memory first and written
+// with coalesced 16-byte stores; otherwise each pair is stored where it is found.
+template <bool kPartialTies>
+__device__ uint32_t sendSource(SmemLayout const& s, SelectOut const& sel, int32_t n, uint32_t lo, uint64_t hi,
+    uint32_t cap, uint32_t* mcSlot, int32_t src, int32_t rank, int32_t world, uint2* stage, int64_t* dbg = nullptr)
+{
+    bool const strict = !sel.all && sel.above + sel.binCount > cap;
+    uint32_t const threshold = sel.all ? lo : sel.edge;
+    // Keys equal to the edge that fit after the keys above it (sel.above < need <= cap when strict).
+    uint32_t const tieBudget = kPartialTies && strict ? cap - sel.above : 0u;
+    uint32_t const sent = sel.all ? sel.total : (strict ? sel.above + tieBudget : sel.above + sel.binCount);
+    uint32_t const bound = sel.all ? lo : (strict ? sel.edge + 1u : sel.edge);
+    stamp(dbg, 42);
+    // One page block (kBlockTokens keys) per thread; positions from a block scan, so the slot holds the keys in
+    // source order: first those taken outright, then (strict only) the admitted keys equal to the edge.
+    int32_t const groups = (n + kBlockTokens - 1) / kBlockTokens;
+    uint32_t takenMain = 0;
+    uint32_t takenTies = 0;
+    uint32_t const tieBase = strict ? sel.above : 0u;
+    for (int32_t g0 = 0; g0 < groups; g0 += kThreads)
+    {
+        int32_t const g = g0 + static_cast<int32_t>(threadIdx.x);
+        uint32_t k[kBlockTokens];
+        uint32_t mask = 0;
+        uint32_t tieMask = 0;
+        if (g < groups)
+        {
+            uint4 const a = *reinterpret_cast<uint4 const*>(s.chunk + kBlockTokens * g);
+            uint4 const b = *reinterpret_cast<uint4 const*>(s.chunk + kBlockTokens * g + 4);
+            k[0] = a.x;
+            k[1] = a.y;
+            k[2] = a.z;
+            k[3] = a.w;
+            k[4] = b.x;
+            k[5] = b.y;
+            k[6] = b.z;
+            k[7] = b.w;
+#pragma unroll
+            for (int32_t i = 0; i < kBlockTokens; ++i)
+            {
+                bool const in = kBlockTokens * g + i < n && k[i] >= lo && static_cast<uint64_t>(k[i]) < hi;
+                bool const take = in && (strict ? k[i] > threshold : k[i] >= threshold);
+                bool const tie = kPartialTies && in && strict && k[i] == threshold;
+                mask |= (take ? 1u : 0u) << i;
+                tieMask |= (tie ? 1u : 0u) << i;
+            }
+        }
+        // At most kThreads * kBlockTokens keys per pass: both counts fit 16 bits.
+        uint32_t const c = static_cast<uint32_t>(__popc(mask)) | (static_cast<uint32_t>(__popc(tieMask)) << 16);
+        uint32_t chunkTaken;
+        uint32_t const before = blockExclusiveScan(c, s.warpSums, chunkTaken);
+        uint32_t pos = takenMain + (before & 0xFFFFu);
+        uint32_t tiePos = takenTies + (before >> 16);
+        takenMain += chunkTaken & 0xFFFFu;
+        takenTies += chunkTaken >> 16;
+        if ((mask | tieMask) != 0u)
+        {
+            int32_t const colBase = g * kPageTokens + ((src + g) % kSourcesPerRank) * kBlockTokens;
+#pragma unroll
+            for (int32_t i = 0; i < kBlockTokens; ++i)
+            {
+                uint32_t slot = kEmptyWord;
+                if ((mask >> i) & 1u)
+                {
+                    slot = pos++;
+                }
+                else if constexpr (kPartialTies)
+                {
+                    if ((tieMask >> i) & 1u)
+                    {
+                        slot = tiePos < tieBudget ? tieBase + tiePos : kEmptyWord;
+                        ++tiePos;
+                    }
+                }
+                if (slot != kEmptyWord)
+                {
+                    uint32_t const token = localToGlobal(colBase + i, rank, world);
+                    if (stage != nullptr)
+                    {
+                        stage[slot] = make_uint2(k[i], token);
+                    }
+                    else
+                    {
+                        store2(mcSlot + kSlotHeaderWords + 2 * static_cast<int64_t>(slot), k[i], token);
+                    }
+                }
+            }
+        }
+        // The staged pairs are complete, and the scan's warp sums may be rewritten.
+        __syncthreads();
+    }
+    stamp(dbg, 43);
+    if (stage != nullptr)
+    {
+        // Two pairs per 16-byte store; an odd count leaves the last half of the last vector empty.
+        int32_t const vectors = static_cast<int32_t>(sent + 1u) / 2;
+        for (int32_t v = threadIdx.x; v < vectors; v += kThreads)
+        {
+            uint2 const a = stage[2 * v];
+            uint2 const b
+                = 2 * v + 1 < static_cast<int32_t>(sent) ? stage[2 * v + 1] : make_uint2(kEmptyWord, kEmptyWord);
+            store4(mcSlot + kSlotHeaderWords + 4 * static_cast<int64_t>(v), make_uint4(a.x, a.y, b.x, b.y));
+        }
+    }
+    // Receivers check every candidate word, so the header needs no ordering against the candidates.
+    if constexpr (kPartialTies)
+    {
+        if (threadIdx.x == kThreads - 1)
+        {
+            store2(mcSlot, sent, bound);
+        }
+    }
+    else
+    {
+        // The statistics of the keys sent, from the staged pairs (round 1 always stages), by the last warp, which
+        // stores no vector (at most kSourceCapacity / 2), while the first warps store the vectors.
+        static_assert(kSourceCapacity / 2 <= kThreads - 32, "the last warp stores no candidate vector");
+        if (threadIdx.x >= kThreads - 32)
+        {
+            int32_t const lane = threadIdx.x % 32;
+            KeyStats t;
+#pragma unroll
+            for (int32_t j = 0; j < kSourceCapacity / 32; ++j)
+            {
+                int32_t const i = lane + 32 * j;
+                if (i < static_cast<int32_t>(sent))
+                {
+                    t.add(stage[i].x);
+                }
+            }
+            uint32_t const maxKey = __reduce_max_sync(0xFFFFFFFFu, t.maxKey);
+            uint32_t const minKey = __reduce_min_sync(0xFFFFFFFFu, t.minKey);
+            uint32_t const positives = __reduce_add_sync(0xFFFFFFFFu, t.positives);
+            uint32_t const maxNeg = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg);
+            uint32_t const minPos = __reduce_min_sync(0xFFFFFFFFu, t.minPos);
+            if (lane == 31)
+            {
+                store4(mcSlot, make_uint4(sent, bound, maxKey, minKey + 1u));
+                store4(mcSlot + 4, make_uint4(positives, maxNeg, minPos + 1u, 0u));
+            }
+        }
+    }
+    stamp(dbg, 44);
+    return sent;
+}
+
+// Load one source's keys (one block of kBlockTokens tokens of every local page) from the local logits into the
+// chunk; returns the number of keys. Consecutive threads read consecutive keys (whole 32-byte blocks per eight
+// threads), and each thread issues kLoadBatch loads before it uses any. With `withStats`, the per-warp key
+// statistics are stored for readStats after the next block barrier.
+__device__ int32_t loadSource(
+    SmemLayout const& s, float const* rowLogits, int32_t localLen, int32_t src, bool withStats)
+{
+    constexpr int32_t kLoadBatch = 8;
+    int32_t const n = sourceLength(localLen, src);
+    KeyStats st;
+    for (int32_t j0 = threadIdx.x; j0 < n; j0 += kLoadBatch * kThreads)
+    {
+        float v[kLoadBatch];
+#pragma unroll
+        for (int32_t i = 0; i < kLoadBatch; ++i)
+        {
+            int32_t const j = j0 + i * kThreads;
+            int32_t const page = j / kBlockTokens;
+            int32_t const col = page * kPageTokens + ((src + page) % kSourcesPerRank) * kBlockTokens + j % kBlockTokens;
+            v[i] = j < n ? rowLogits[col] : 0.F;
+        }
+#pragma unroll
+        for (int32_t i = 0; i < kLoadBatch; ++i)
+        {
+            int32_t const j = j0 + i * kThreads;
+            if (j < n)
+            {
+                uint32_t const k = orderedKey(v[i]);
+                s.chunk[j] = k;
+                st.add(k);
+            }
+        }
+    }
+    if (withStats)
+    {
+        storeWarpStats(st, s.stats);
+    }
+    return n;
+}
+
+// Whether source q withheld keys that may lie above the merged k-th key.
+__device__ __forceinline__ bool sourceOverflows(SmemLayout const& s, int32_t q, uint32_t kth)
+{
+    uint32_t const cnt = s.scal[kSlotSrcCount + q];
+    uint64_t const bound = s.scal[kSlotSrcBound + q];
+    return cnt < s.scal[kSlotSrcLen + q] && bound > uint64_t{kth} + 1u;
+}
+
+// Emit the top-kTopK of the merge candidates. Keys above `kth` are all taken; ties at `kth` are taken in candidate
+// order when only some of them fit (deterministic on every rank), otherwise all of them. With `takeAll`, every
+// candidate is taken and the row is padded with -1. The row is assembled in shared memory (the first histogram)
+// and written with coalesced stores.
+template <typename KeyAt>
+__device__ bool emitSelection(SmemLayout const& s, KeyAt const& keys, int32_t total, uint32_t kth, bool takeAll,
+    int32_t* outRow, int32_t overflowSources, int64_t* dbg = nullptr)
+{
+    uint32_t const want = static_cast<uint32_t>(kTopK);
+    // Positions from one block scan over contiguous candidate ranges (no contended counter); the order is the
+    // candidate order, identical on every rank. Shared-memory candidates are read four at a time.
+    int32_t per = (total + kThreads - 1) / kThreads;
+    if constexpr (KeyAt::kVector)
+    {
+        per = (per + 3) / 4 * 4;
+    }
+    int32_t const j0 = min(total, static_cast<int32_t>(threadIdx.x) * per);
+    int32_t const j1 = min(total, j0 + per);
+    uint32_t cGt = 0;
+    uint32_t cEq = 0;
+    if constexpr (KeyAt::kVector)
+    {
+        for (int32_t j = j0; j < j1; j += 4)
+        {
+            uint4 const q = keys.keys4(j);
+            uint32_t const kk[4] = {q.x, q.y, q.z, q.w};
+#pragma unroll
+            for (int32_t i = 0; i < 4; ++i)
+            {
+                bool const in = j + i < j1;
+                cGt += (in && (takeAll || kk[i] > kth)) ? 1u : 0u;
+                cEq += (in && !takeAll && kk[i] == kth) ? 1u : 0u;
+            }
+        }
+    }
+    else
+    {
+        for (int32_t j = j0; j < j1; ++j)
+        {
+            uint32_t const k = keys(j);
+            cGt += (takeAll || k > kth) ? 1u : 0u;
+            cEq += (!takeAll && k == kth) ? 1u : 0u;
+        }
+    }
+    // The first `overflowSources` threads flag a source that withheld keys above the threshold (the flag slot is
+    // zero on entry); with a flag set, nothing is written and the caller runs the second round.
+    if (static_cast<int32_t>(threadIdx.x) < overflowSources && sourceOverflows(s, threadIdx.x, kth))
+    {
+        atomicOr(&s.scal[kSlotCounter2], 1u);
+    }
+    uint32_t allGt;
+    uint32_t allEq;
+    uint32_t posGt;
+    uint32_t posEq;
+    if (total < 0x10000)
+    {
+        // Both counts fit 16 bits: one scan of the packed pair.
+        uint32_t packedTotal;
+        uint32_t const packed = cGt | (cEq << 16);
+        uint32_t const before = blockExclusiveScan(packed, s.warpSums, packedTotal);
+        allGt = packedTotal & 0xFFFFu;
+        allEq = packedTotal >> 16;
+        posGt = before & 0xFFFFu;
+        posEq = allGt + (before >> 16);
+    }
+    else
+    {
+        posGt = blockInclusiveScan(cGt, s.warpSums, allGt) - cGt;
+        __syncthreads();
+        posEq = blockInclusiveScan(cEq, s.warpSums, allEq) - cEq;
+        posEq += allGt;
+    }
+    stamp(dbg, 59);
+    if (overflowSources > 0 && s.scal[kSlotCounter2] != 0u)
+    {
+        return true;
+    }
+    uint32_t* row = s.hist;
+    auto place = [&](uint32_t k, uint32_t idx)
+    {
+        if (takeAll || k > kth)
+        {
+            if (posGt < want)
+            {
+                row[posGt] = idx;
+            }
+            ++posGt;
+        }
+        else if (k == kth)
+        {
+            if (posEq < want)
+            {
+                row[posEq] = idx;
+            }
+            ++posEq;
+        }
+    };
+    if constexpr (KeyAt::kVector)
+    {
+        for (int32_t j = j0; j < j1; j += 4)
+        {
+            uint4 const q = keys.keys4(j);
+            uint4 const x = keys.index4(j);
+            uint32_t const kk[4] = {q.x, q.y, q.z, q.w};
+            uint32_t const xx[4] = {x.x, x.y, x.z, x.w};
+#pragma unroll
+            for (int32_t i = 0; i < 4; ++i)
+            {
+                if (j + i < j1)
+                {
+                    place(kk[i], xx[i]);
+                }
+            }
+        }
+    }
+    else
+    {
+        for (int32_t j = j0; j < j1; ++j)
+        {
+            place(keys(j), keys.index(j));
+        }
+    }
+    __syncthreads();
+    stamp(dbg, 60);
+    uint32_t const selected = min(allGt + allEq, want);
+    for (uint32_t j = threadIdx.x; j < want; j += kThreads)
+    {
+        outRow[j] = j < selected ? static_cast<int32_t>(row[j]) : -1;
+    }
+    stamp(dbg, 61);
+    return false;
+}
+
+// Threshold of an exact (cap = 0) merge select: keys above `kth` are taken, ties at `kth` while they fit.
+__device__ __forceinline__ uint32_t mergeThreshold(SelectOut const& m, bool& takeAll)
+{
+    takeAll = m.all || (m.full && m.edge == 0u);
+    if (takeAll)
+    {
+        return 0u;
+    }
+    // A wholly taken bin: the keys at or above its edge are exactly kTopK.
+    return m.full ? m.edge - 1u : m.edge;
+}
+
+// Reset of the previous exchange launch's buffer on this rank's copy. Its 2 * kMaxSources * kMaxRows slots (round 1,
+// then round 2) are spread over the kResetWarps reset warps (the warps from kFirstResetWarp on) of every CTA of the
+// launch, at most ten consecutive slots per warp, one per lane: each lane fetches its slot's header once the epoch is
+// known (fetchResetHeaders, after the grid dependency wait and the first block barrier; an asynchronous copy into
+// shared memory that completes during the local select and the send) and the warp resets the written slots while the
+// poll warps wait for the candidates (resetSlots). A slot was written iff its header was; the header's count bounds
+// the candidates to reset.
+constexpr int32_t kResetSlots = 2 * kMaxSources * kMaxRows;
+
+// This CTA's index among the launch's clearers, and their number (from the grid; recomputed where used).
+__device__ __forceinline__ int32_t resetClearer()
+{
+    return static_cast<int32_t>(blockIdx.y) * kSourcesPerRank + static_cast<int32_t>(blockIdx.x);
+}
+
+__device__ __forceinline__ int32_t resetClearers()
+{
+    return static_cast<int32_t>(gridDim.y) * kSourcesPerRank;
+}
+
+__device__ __forceinline__ int32_t resetSlotsPerWarp(int32_t numClearers)
+{
+    int32_t const warps = numClearers * kResetWarps;
+    return (kResetSlots + warps - 1) / warps;
+}
+
+// First slot of this (reset) warp.
+__device__ __forceinline__ int32_t resetFirstSlot(int32_t clearer, int32_t numClearers)
+{
+    int32_t const warp = static_cast<int32_t>(threadIdx.x) / 32 - kFirstResetWarp;
+    return (clearer * kResetWarps + warp) * resetSlotsPerWarp(numClearers);
+}
+
+__device__ __forceinline__ int64_t resetSlotOffset(int32_t buf, int32_t slot)
+{
+    int32_t const round1Slots = kMaxSources * kMaxRows;
+    bool const second = slot >= round1Slots;
+    int32_t const qr = second ? slot - round1Slots : slot;
+    return second ? round2Offset(buf, qr / kMaxRows, qr % kMaxRows) : round1Offset(buf, qr / kMaxRows, qr % kMaxRows);
+}
+
+__device__ __forceinline__ uint32_t* resetSlotBase(CpTopkParams const& p, int32_t buf, int32_t slot)
+{
+    return p.ucWorkspace + resetSlotOffset(buf, slot);
+}
+
+//! Most slots one reset warp resets (one row: the fewest reset warps).
+constexpr int32_t kMaxResetSlotsPerWarp
+    = (kResetSlots + kSourcesPerRank * kResetWarps - 1) / (kSourcesPerRank * kResetWarps);
+static_assert(kMaxResetSlotsPerWarp <= 32, "one slot header per lane");
+
+// Shared-memory home of the reset warps' slot headers (16 bytes per slot: count, bound and two further header words),
+// in the merge candidate area beyond the round-1 staging of the send; it is free from the first block barrier until
+// the reset.
+constexpr int32_t kResetHeaderOffset = 2 * kSourceCapacity;
+static_assert(kResetHeaderOffset + 4 * kMaxResetSlotsPerWarp * kResetWarps <= kSourcesPerRank * kSourceCapacity,
+    "the reset headers fit the smallest merge candidate area");
+
+__device__ __forceinline__ uint4* resetHeaderHome(SmemLayout const& s)
+{
+    int32_t const warp = static_cast<int32_t>(threadIdx.x) / 32 - kFirstResetWarp;
+    return reinterpret_cast<uint4*>(s.mkey + kResetHeaderOffset) + warp * kMaxResetSlotsPerWarp;
+}
+
+// Word offset, within one Lamport buffer, of the slot whose header this lane fetches for the reset; -1 for the lanes
+// and warps without one. Arithmetic on the grid only (no memory access).
+__device__ __forceinline__ int32_t resetHeaderOffset()
+{
+    int32_t const numClearers = resetClearers();
+    int32_t const lane = threadIdx.x % 32;
+    if (static_cast<int32_t>(threadIdx.x) / 32 < kFirstResetWarp || lane >= resetSlotsPerWarp(numClearers))
+    {
+        return -1;
+    }
+    int32_t const slot = resetFirstSlot(resetClearer(), numClearers) + lane;
+    if (slot >= kResetSlots)
+    {
+        return -1;
+    }
+    static_assert(kBufferWords <= int64_t{1} << 31, "slot offsets within a buffer fit int32");
+    return static_cast<int32_t>(resetSlotOffset(0, slot));
+}
+
+// Starts the copy of this lane's reset slot header of buffer `buf` (offset from resetHeaderOffset) into the warp's
+// home with an asynchronous copy (cp.async, generic proxy, completed in resetSlots): a pending copy holds no register
+// and no block barrier waits for it. Lanes without a slot do nothing.
+__device__ __forceinline__ void fetchResetHeader(
+    CpTopkParams const& p, SmemLayout const& s, int32_t buf, int32_t offset)
+{
+    if (offset >= 0)
+    {
+        uint32_t const dst = static_cast<uint32_t>(
+            __cvta_generic_to_shared(resetHeaderHome(s) + static_cast<int32_t>(threadIdx.x % 32)));
+        uint32_t const* const src = p.ucWorkspace + int64_t{buf} * kBufferWords + offset;
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(src) : "memory");
+    }
+}
+
+// Reset this reset warp's slots of buffer `buf` (headers fetched by fetchResetHeaders); other warps return.
+__device__ void resetSlots(CpTopkParams const& p, SmemLayout const& s, int32_t buf)
+{
+    if (static_cast<int32_t>(threadIdx.x) / 32 < kFirstResetWarp)
+    {
+        return;
+    }
+    int32_t const clearer = resetClearer();
+    int32_t const numClearers = resetClearers();
+    // Each lane completes its own header copy; the warp barrier then makes every lane's header visible to the warp.
+    asm volatile("cp.async.wait_all;" ::: "memory");
+    __syncwarp();
+    int32_t const per = resetSlotsPerWarp(numClearers);
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const first = resetFirstSlot(clearer, numClearers);
+    uint4 const* const headers = resetHeaderHome(s);
+    uint4 const empty = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord);
+    for (int32_t i = 0; i < per && first + i < kResetSlots; ++i)
+    {
+        uint4 const h = headers[i];
+        if (h.x == kEmptyWord && h.y == kEmptyWord)
+        {
+            continue;
+        }
+        int32_t const slot = first + i;
+        uint32_t* base = resetSlotBase(p, buf, slot);
+        uint32_t const capacity = slot >= kMaxSources * kMaxRows ? kRound2Capacity : kSourceCapacity;
+        uint32_t const count = h.x == kEmptyWord ? capacity : min(h.x, capacity);
+        int32_t const vectors = static_cast<int32_t>(count + 1u) / 2;
+        for (int32_t v = lane; v < vectors; v += 32)
+        {
+            store4(base + kSlotHeaderWords + 4 * static_cast<int64_t>(v), empty);
+        }
+        if (lane < 2)
+        {
+            store4(base + 4 * lane, empty);
+        }
+    }
+}
+
+// Merge of one row on this rank (one CTA): gathers every source's round-1 candidates, selects the exact
+// top-kTopK, runs the second round when a source withheld keys above the threshold, and writes the row.
+__device__ void mergeRow(CpTopkParams const& p, SmemLayout const& s, uint32_t const** segSlots, uint32_t* segBegin,
+    int32_t buf, int32_t row, int32_t rowLen, int32_t localLen, float const* rowLogits, int32_t* outRow,
+    WaitCtx const& wait)
+{
+    int32_t const world = p.worldSize;
+    int32_t const rank = p.rank;
+    int32_t const sources = numSources(world);
+    uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+    __syncthreads();
+    if (threadIdx.x == 0)
+    {
+        s.scal[kSlotCounter2] = 0u;
+    }
+    if (static_cast<int32_t>(threadIdx.x) < sources)
+    {
+        // Keys each source holds, computed while its header is in flight.
+        int32_t const q = threadIdx.x;
+        s.scal[kSlotSrcLen + q] = static_cast<uint32_t>(
+            sourceLength(min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank));
+        uint2 const h = pollPair(p.ucWorkspace + round1Offset(buf, threadIdx.x, row), wait, kWaitRound1Header);
+        s.scal[kSlotSrcCount + threadIdx.x] = min(h.x, cap);
+        s.scal[kSlotSrcBound + threadIdx.x] = h.y;
+    }
+    __syncthreads();
+    if (threadIdx.x < 32)
+    {
+        // Exclusive prefix of the source counts (at most 64 sources: two per lane).
+        uint32_t const a = static_cast<int32_t>(threadIdx.x) < sources ? s.scal[kSlotSrcCount + threadIdx.x] : 0u;
+        uint32_t const b
+            = static_cast<int32_t>(threadIdx.x) + 32 < sources ? s.scal[kSlotSrcCount + threadIdx.x + 32] : 0u;
+        uint32_t x = a;
+        uint32_t z = b;
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+            uint32_t const w = __shfl_up_sync(0xFFFFFFFFu, z, o);
+            if (static_cast<int32_t>(threadIdx.x) >= o)
+            {
+                x += y;
+                z += w;
+            }
+        }
+        uint32_t const totalA = __shfl_sync(0xFFFFFFFFu, x, 31);
+        s.scal[kSlotSrcOffset + threadIdx.x] = x - a;
+        s.scal[kSlotSrcOffset + 32 + threadIdx.x] = totalA + z - b;
+        if (threadIdx.x == 31)
+        {
+            s.scal[kSlotSrcOffset + kMaxSources] = totalA + z;
+        }
+    }
+    __syncthreads();
+    int32_t const total = static_cast<int32_t>(s.scal[kSlotSrcOffset + kMaxSources]);
+    markPhase(p, 4);
+
+    // Candidates: padded (source, vector) positions; issue a batch of independent 16-byte loads, then re-poll the
+    // ones not yet written. The key statistics are tracked for the merge select.
+    int32_t const vectorsPerSource = static_cast<int32_t>(cap) / 2;
+    int32_t const padded = sources * vectorsPerSource;
+    KeyStats st;
+#pragma unroll 1
+    for (int32_t v0 = 0; v0 < padded; v0 += kThreads * kPollBatch)
+    {
+        uint4 v[kPollBatch];
+#pragma unroll
+        for (int32_t i = 0; i < kPollBatch; ++i)
+        {
+            int32_t const jv = v0 + i * kThreads + static_cast<int32_t>(threadIdx.x);
+            int32_t const q = jv / vectorsPerSource;
+            int32_t const kv = jv % vectorsPerSource;
+            v[i] = make_uint4(0u, 0u, 0u, 0u);
+            if (jv < padded && 2 * kv < static_cast<int32_t>(s.scal[kSlotSrcCount + q]))
+            {
+                v[i] = loadVolatile4(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv);
+            }
+        }
+#pragma unroll
+        for (int32_t i = 0; i < kPollBatch; ++i)
+        {
+            int32_t const jv = v0 + i * kThreads + static_cast<int32_t>(threadIdx.x);
+            int32_t const q = jv / vectorsPerSource;
+            int32_t const kv = jv % vectorsPerSource;
+            if (jv < padded && 2 * kv < static_cast<int32_t>(s.scal[kSlotSrcCount + q]))
+            {
+                bool const second = 2 * kv + 1 < static_cast<int32_t>(s.scal[kSlotSrcCount + q]);
+                if (!vectorReady(v[i], second))
+                {
+                    v[i] = pollVector(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv, second,
+                        wait, kWaitRound1Candidates);
+                }
+                int32_t const j = static_cast<int32_t>(s.scal[kSlotSrcOffset + q]) + 2 * kv;
+                s.mkey[j] = v[i].x;
+                s.midx[j] = v[i].y;
+                st.add(v[i].x);
+                if (second)
+                {
+                    s.mkey[j + 1] = v[i].z;
+                    s.midx[j + 1] = v[i].w;
+                    st.add(v[i].z);
+                }
+            }
+        }
+    }
+    stamp(debugSlots(p, 0), 45);
+    storeWarpStats(st, s.stats);
+    __syncthreads();
+    markPhase(p, 5);
+
+    // Digits of the merge select span the candidates of the sign the cut lies in (see kPositiveKey).
+    SignRange const r = total > 0 ? signRange(readStats(s.stats), static_cast<uint32_t>(kTopK))
+                                  : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false};
+    SmemKeys const merged{s.mkey, s.midx};
+    SelectOut const m = blockSelect(
+        s, merged, total, r.lo, r.hi, static_cast<uint32_t>(kTopK), 0u, r.lo, r.bits, debugSlots(p, 48), r.aboveHi);
+    putDebug(debugSlots(p, 0), 62, total);
+    putDebug(debugSlots(p, 0), 63, m.passes + (r.negative ? 10 : 0));
+    bool takeAll = false;
+    uint32_t const kth = mergeThreshold(m, takeAll);
+    markPhase(p, 6);
+    stamp(debugSlots(p, 0), 58);
+    if (!emitSelection(s, merged, total, kth, takeAll, outRow, sources, debugSlots(p, 0)))
+    {
+        return;
+    }
+
+    // ---- Round 2 (rare; every rank takes it): withheld keys in [kth, bound) of overflowing sources ----------
+    __syncthreads();
+    // 1. This rank's overflowing sources send their top-kRound2Capacity withheld keys.
+    for (int32_t c = 0; c < kSourcesPerRank; ++c)
+    {
+        int32_t const q = rank * kSourcesPerRank + c;
+        if (!sourceOverflows(s, q, kth))
+        {
+            continue;
+        }
+        uint64_t const b = s.scal[kSlotSrcBound + q];
+        int32_t const nq = loadSource(s, rowLogits, localLen, c, false);
+        __syncthreads();
+        SelectOut const sel2 = blockSelect(
+            s, SmemKeys{s.chunk, nullptr}, nq, kth, b, static_cast<uint32_t>(kRound2Capacity), kRound2Capacity, 0u, 32);
+        (void) sendSource<true>(
+            s, sel2, nq, kth, b, kRound2Capacity, p.mcWorkspace + round2Offset(buf, q, row), c, rank, world, nullptr);
+        __syncthreads();
+    }
+    // 2. Gather the round-2 segments of every overflowing source (all ranks).
+    if (threadIdx.x == 0)
+    {
+        int32_t segs = 0;
+        uint32_t off = 0;
+        for (int32_t q = 0; q < sources; ++q)
+        {
+            if (!sourceOverflows(s, q, kth))
+            {
+                continue;
+            }
+            uint32_t const* slot = p.ucWorkspace + round2Offset(buf, q, row);
+            uint2 const h = pollPair(slot, wait, kWaitRound2Header);
+            segSlots[segs] = slot;
+            segBegin[segs] = off;
+            off += min(h.x, static_cast<uint32_t>(kRound2Capacity));
+            ++segs;
+        }
+        segBegin[segs] = off;
+        s.scal[kSlotTotal] = static_cast<uint32_t>(segs);
+    }
+    __syncthreads();
+    int32_t const segs = static_cast<int32_t>(s.scal[kSlotTotal]);
+    int32_t const total2 = static_cast<int32_t>(segBegin[segs]);
+    MergeKeys const keys{s.mkey, s.midx, total, segSlots, segBegin, segs};
+    for (int32_t k = threadIdx.x; k < total2; k += kThreads)
+    {
+        (void) pollPair(keys.pair(k), wait, kWaitRound2Candidates);
+    }
+    __syncthreads();
+    SelectOut const m2
+        = blockSelect(s, keys, total + total2, 0u, uint64_t{kEmptyWord}, static_cast<uint32_t>(kTopK), 0u, 0u, 32);
+    bool takeAll2 = false;
+    uint32_t const kth2 = mergeThreshold(m2, takeAll2);
+    (void) emitSelection(s, keys, total + total2, kth2, takeAll2, outRow, 0);
+    if (threadIdx.x == 0)
+    {
+        atomicAdd(&p.state[2], 1);
+    }
+}
+
+// ---- Thread-block cluster primitives (one cluster of kClusterCtas CTAs per row; cluster rank = blockIdx.x) ----------
+
+__device__ __forceinline__ uint32_t smemAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// Shared-memory window address of `ptr` (this CTA's shared memory) in CTA `rank` of the cluster.
+__device__ __forceinline__ uint32_t clusterAddr(void const* ptr, int32_t rank)
+{
+    uint32_t remote;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smemAddr(ptr)), "r"(rank));
+    return remote;
+}
+
+// Initialize the exchange barriers and announce it to the cluster (relaxed arrive; clusterInitWait completes it
+// before the first exchange). Touches shared memory only.
+__device__ __forceinline__ void clusterInit(ClusterRecords& rec)
+{
+    if (threadIdx.x == 0)
+    {
+        asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[0])) : "memory");
+        asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[1])) : "memory");
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+    asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory");
+}
+
+__device__ __forceinline__ void clusterInitWait()
+{
+    asm volatile("barrier.cluster.wait.aligned;" ::: "memory");
+}
+
+// This CTA's single arrival on exchange barrier `bar`, expecting `bytes` from the cluster's st.async stores.
+__device__ __forceinline__ void expectBytes(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile(
+        "{\n\t.reg .b64 st;\n\tmbarrier.arrive.expect_tx.shared::cta.b64 st, [%0], %1;\n\t}" ::"r"(smemAddr(bar)),
+        "r"(bytes)
+        : "memory");
+}
+
+// Wait until the phase of exchange barrier `bar` with parity `parity` completed (every expected byte arrived).
+__device__ __forceinline__ void waitExchange(uint64_t* bar, uint32_t parity)
+{
+    uint32_t done;
+    do
+    {
+        asm volatile(
+            "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], %2;\n\t"
+            "selp.u32 %0, 1, 0, p;\n\t}"
+            : "=r"(done)
+            : "r"(smemAddr(bar)), "r"(parity)
+            : "memory");
+    } while (done == 0u);
+}
+
+// Store 16 / 8 bytes at `ptr` (a record of this CTA's layout) in CTA `rank`, completing them on its barrier `bar`.
+__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint4 v)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(
+                     clusterAddr(ptr, rank)),
+                 "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(clusterAddr(bar, rank))
+                 : "memory");
+}
+
+__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint2 v)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.b32 [%0], {%1, %2}, [%3];" ::"r"(
+                     clusterAddr(ptr, rank)),
+                 "r"(v.x), "r"(v.y), "r"(clusterAddr(bar, rank))
+                 : "memory");
+}
+
+// Exchange sequence of a row: per pass x = 2 * pass group totals and x = 2 * pass + 1 the selected group's bins.
+// Every CTA sends the records of exchange x + 1 only after its own exchange x completed, so a barrier phase (and an
+// alternating record set) is reused only after every CTA consumed the previous one.
+__device__ __forceinline__ uint64_t* exchangeBar(ClusterRecords& rec, uint32_t x)
+{
+    return &rec.bar[x & 1u];
+}
+
+__device__ __forceinline__ uint32_t exchangeParity(uint32_t x)
+{
+    return (x >> 1) & 1u;
+}
+
+// Inclusive scan over the lanes of a warp.
+__device__ __forceinline__ uint32_t warpInclusiveScan(uint32_t v)
+{
+    int32_t const lane = threadIdx.x % 32;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, v, o);
+        if (lane >= o)
+        {
+            v += y;
+        }
+    }
+    return v;
+}
+
+// Tallies of the cluster select over the candidates of the cluster CTAs before this one (in candidate order), kept
+// by warp 0 and published in shared memory: those strictly above the selected bins of all passes so far, and those
+// in the selected bin of the last pass.
+struct CtaTally
+{
+    uint32_t aboveBefore;
+    uint32_t binBefore;
+};
+
+// Cluster radix select of the need-th largest key among the candidates of every CTA of the cluster (this thread
+// holds k0 and k1; kEmptyWord is no candidate, as it lies above every range). It computes exactly what blockSelect
+// with cap = 0 computes over the union of the candidates: the same digits (from `base` and topBits), the same bins
+// and counts from the cluster-summed histograms, and the same stop rules. Per pass: every CTA histograms its own
+// candidates and sends its totals per group of kGroupBins bins to every CTA (st.async, two groups per warp); warp 0
+// of every CTA finds, from the top, the group holding the target, sends its own counts of that group's bins to every
+// CTA, and finds the bin from the eight CTAs' counts. From the same records, warp 0 also tallies the candidates of
+// the CTAs before this one above the selected bins and in the last selected bin (`tally`), which give this CTA's
+// emission offsets without another exchange.
+__device__ SelectOut clusterSelect(SmemLayout const& s, ClusterRecords& rec, CtaTally& tally, uint32_t k0, uint32_t k1,
+    uint32_t lo, uint64_t hi, uint32_t need, uint32_t base, int32_t topBits, uint32_t aboveHi, int64_t* dbg)
+{
+    int32_t const me = blockIdx.x;
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    SelectOut out{lo, 0u, 0u, 0u, false, false, 0};
+    uint32_t prefMask = 0;
+    uint32_t prefVal = 0;
+    uint32_t above = aboveHi;
+    int32_t remaining = max(topBits, 1);
+    uint32_t* hist = s.hist;
+    uint32_t* histNext = s.hist + kBins;
+    // Warp 0: the running tallies of the CTAs before this one.
+    uint32_t aboveBefore = 0u;
+    uint32_t binBefore = 0u;
+    for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads)
+    {
+        hist[b] = 0;
+    }
+    __syncthreads();
+#pragma unroll 1
+    for (int32_t pass = 0; remaining > 0; ++pass)
+    {
+        int32_t const width = min(kDigitBits, remaining);
+        int32_t const shift = remaining - width;
+        int32_t const nbins = 1 << width;
+        uint32_t const binMask = static_cast<uint32_t>(nbins - 1);
+        int32_t const groups = (nbins + kGroupBins - 1) / kGroupBins;
+        int32_t const pairs = (groups + 1) / 2;
+        uint32_t const xg = 2u * pass;
+        uint32_t const xb = xg + 1u;
+        out.passes = pass + 1;
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[1 + 3 * pass] = clock64();
+        }
+        if (threadIdx.x == 0)
+        {
+            expectBytes(exchangeBar(rec, xg), kClusterCtas * pairs * 8);
+        }
+        // The next pass's histogram: its last reads (the previous pass's bins sent by warp 0) are complete.
+        if (shift > 0)
+        {
+            for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads)
+            {
+                histNext[b] = 0;
+            }
+        }
+        histogramKey(hist, true, k0, lo, hi, base, prefMask, prefVal, shift, binMask);
+        histogramKey(hist, true, k1, lo, hi, base, prefMask, prefVal, shift, binMask);
+        __syncthreads();
+        uint32_t(*const groupRec)[kMaxGroups] = rec.groups[pass & 1];
+        uint32_t(*const binRec)[kGroupBins] = rec.bins[pass & 1];
+        if (warp < pairs)
+        {
+            int32_t const b0 = 2 * warp * kGroupBins + lane;
+            int32_t const b1 = b0 + kGroupBins;
+            uint32_t const t0 = __reduce_add_sync(0xFFFFFFFFu, b0 < nbins ? hist[b0] : 0u);
+            uint32_t const t1 = __reduce_add_sync(0xFFFFFFFFu, b1 < nbins ? hist[b1] : 0u);
+            if (lane < kClusterCtas)
+            {
+                sendTo(lane, &groupRec[me][2 * warp], exchangeBar(rec, xg), make_uint2(t0, t1));
+            }
+        }
+        uint32_t const target = need - above;
+        if (warp == 0)
+        {
+            waitExchange(exchangeBar(rec, xg), exchangeParity(xg));
+            if (dbg != nullptr && lane == 0 && pass < 3)
+            {
+                dbg[2 + 3 * pass] = clock64();
+            }
+            // Lane l sums the groups at positions 2l and 2l + 1 from the top over the cluster.
+            int32_t const ga = groups - 1 - 2 * lane;
+            int32_t const gb = ga - 1;
+            uint32_t a = 0;
+            uint32_t b = 0;
+            uint32_t beforeA = 0;
+            uint32_t beforeB = 0;
+#pragma unroll
+            for (int32_t c = 0; c < kClusterCtas; ++c)
+            {
+                uint32_t const x = ga >= 0 ? groupRec[c][ga] : 0u;
+                uint32_t const y = gb >= 0 ? groupRec[c][gb] : 0u;
+                a += x;
+                b += y;
+                beforeA += c < me ? x : 0u;
+                beforeB += c < me ? y : 0u;
+            }
+            uint32_t const incl = warpInclusiveScan(a + b);
+            uint32_t const before = incl - a - b;
+            uint32_t const total = __shfl_sync(0xFFFFFFFFu, incl, 31);
+            bool const holds = before < target && target <= incl;
+            uint32_t const holder = __ballot_sync(0xFFFFFFFFu, holds);
+            if (lane == 0)
+            {
+                s.scal[kSlotTotal] = total;
+            }
+            if (holder != 0u)
+            {
+                int32_t const hl = __ffs(holder) - 1;
+                bool const first = target <= before + a;
+                int32_t const gSel = __shfl_sync(0xFFFFFFFFu, first ? ga : gb, hl);
+                uint32_t const gAbove = __shfl_sync(0xFFFFFFFFu, first ? before : before + a, hl);
+                if (lane == 0)
+                {
+                    expectBytes(exchangeBar(rec, xb), kClusterCtas * kGroupBins * 4);
+                }
+                // This CTA's counts of the selected group's bins to every CTA: lane l sends vector l % 8 to CTAs
+                // l / 8 and l / 8 + 4.
+                int32_t const v = lane % 8;
+                int32_t const b0 = gSel * kGroupBins + 4 * v;
+                uint4 const q = make_uint4(b0 < nbins ? hist[b0] : 0u, b0 + 1 < nbins ? hist[b0 + 1] : 0u,
+                    b0 + 2 < nbins ? hist[b0 + 2] : 0u, b0 + 3 < nbins ? hist[b0 + 3] : 0u);
+                sendTo(lane / 8, &binRec[me][4 * v], exchangeBar(rec, xb), q);
+                sendTo(lane / 8 + 4, &binRec[me][4 * v], exchangeBar(rec, xb), q);
+                // While the bins travel: the earlier CTAs' candidates in the groups above the selected one (both
+                // groups of the lanes before the holder, and the holder's first group when it selected the second).
+                bool const selFirst = __shfl_sync(0xFFFFFFFFu, first, hl);
+                aboveBefore += __reduce_add_sync(
+                    0xFFFFFFFFu, lane < hl ? beforeA + beforeB : (lane == hl && !selFirst ? beforeA : 0u));
+                waitExchange(exchangeBar(rec, xb), exchangeParity(xb));
+                // Bins of the selected group from the top: lane l holds bin 32 * gSel + 31 - l.
+                int32_t const bin = gSel * kGroupBins + kGroupBins - 1 - lane;
+                uint32_t c = 0;
+                uint32_t cBefore = 0;
+#pragma unroll
+                for (int32_t d = 0; d < kClusterCtas; ++d)
+                {
+                    uint32_t const x = binRec[d][kGroupBins - 1 - lane];
+                    c += x;
+                    cBefore += d < me ? x : 0u;
+                }
+                uint32_t const y = warpInclusiveScan(c);
+                uint32_t const acc = gAbove + y - c;
+                bool const found = acc < target && target <= acc + c;
+                if (found)
+                {
+                    s.scal[kSlotBin] = static_cast<uint32_t>(bin);
+                    s.scal[kSlotAbove] = acc;
+                    s.scal[kSlotBinCount] = c;
+                }
+                // The earlier CTAs' candidates in the bins above the selected one (lanes before the finder) and in
+                // the selected bin (the finder's).
+                int32_t const finder = __ffs(__ballot_sync(0xFFFFFFFFu, found)) - 1;
+                aboveBefore += __reduce_add_sync(0xFFFFFFFFu, lane < finder ? cBefore : 0u);
+                binBefore = __shfl_sync(0xFFFFFFFFu, cBefore, finder);
+            }
+            else if (lane == 0)
+            {
+                // Every key is taken (pass 0 only, the same decision in every CTA): the bin exchange still
+                // completes, with no bytes, so that the exchange sequence stays aligned.
+                expectBytes(exchangeBar(rec, xb), 0u);
+                waitExchange(exchangeBar(rec, xb), exchangeParity(xb));
+            }
+            if (lane == 0)
+            {
+                tally.aboveBefore = aboveBefore;
+                tally.binBefore = binBefore;
+            }
+        }
+        __syncthreads();
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[3 + 3 * pass] = clock64();
+        }
+        uint32_t const total = s.scal[kSlotTotal];
+        if (pass == 0)
+        {
+            out.total = total + aboveHi;
+            if (total + aboveHi == 0u || target > total)
+            {
+                out.all = true;
+                out.above = total + aboveHi;
+                return out;
+            }
+        }
+        // The scalars are next written by warp 0 after the next histogram's block barrier.
+        uint32_t const bin = s.scal[kSlotBin];
+        above += s.scal[kSlotAbove];
+        out.binCount = s.scal[kSlotBinCount];
+        prefVal |= bin << shift;
+        prefMask |= binMask << shift;
+        out.edge = base + prefVal;
+        out.above = above;
+        remaining = shift;
+        uint32_t* const t = hist;
+        hist = histNext;
+        histNext = t;
+        if (remaining > 0 && above + out.binCount == need)
+        {
+            out.full = true;
+            return out;
+        }
+    }
+    return out;
+}
+
+// Spin until the header of a round-1 slot was written and, when the slot's count covers candidate vector kv, that
+// vector too; returns the header's count (clamped to `cap`) and bound. Both are polled together, so neither waits
+// for the other; aborts the kernel after the timeout.
+__device__ uint2 pollSlotVector(uint32_t const* slot, int32_t kv, uint32_t cap, uint4& v, WaitCtx const& w)
+{
+    uint32_t const* vec = slot + kSlotHeaderWords + 4 * kv;
+    uint2 h = loadVolatile2(slot);
+    v = loadVolatile4(vec);
+    uint64_t start = 0;
+    for (int32_t polls = 1;; ++polls)
+    {
+        bool const headerReady = !isEmpty(h);
+        uint32_t const cnt = headerReady ? min(h.x, cap) : 0u;
+        if (headerReady && (2u * kv >= cnt || vectorReady(v, 2u * kv + 1u < cnt)))
+        {
+            return make_uint2(cnt, h.y);
+        }
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, headerReady ? kWaitRound1Candidates : kWaitRound1Header);
+            }
+        }
+        if (!headerReady)
+        {
+            h = loadVolatile2(slot);
+        }
+        v = loadVolatile4(vec);
+    }
+}
+
+// Spin until the whole 32-byte round-1 header of a slot was written; aborts the kernel after the timeout.
+__device__ void pollHeader(uint32_t const* slot, uint4& a, uint4& b, WaitCtx const& w)
+{
+    a = loadVolatile4(slot);
+    b = loadVolatile4(slot + 4);
+    uint64_t start = 0;
+    for (int32_t polls = 1; !(vectorReady(a, true) && vectorReady(b, true)); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, kWaitRound1Header);
+            }
+        }
+        a = loadVolatile4(slot);
+        b = loadVolatile4(slot + 4);
+    }
+}
+
+// Merge of one row by its cluster. CTA c gathers the round-1 candidates of sources [c * world, (c + 1) * world) (one
+// 16-byte vector, two candidate slots, per thread of the poll warps, kept in registers: thread order is crossRankTopk's
+// candidate order) while its header warps read every source's header (counts, bounds and the key statistics of the keys
+// sent, so every CTA has the row's statistics and can decide every source's overflow without an exchange). The cluster
+// runs crossRankTopk's merge select on the summed histograms (clusterSelect), whose records also give every CTA's
+// number of keys above the threshold and of ties, and every CTA writes its selected candidates at crossRankTopk's
+// positions: keys above the threshold in candidate order, then ties in candidate order while they fit, then -1. Returns
+// true, writing nothing, when a source withheld keys that may lie above the threshold (the same decision on every CTA
+// and rank); the row is then merged by mergeRow.
+__device__ bool clusterMergeRow(CpTopkParams const& p, SmemLayout const& s, ClusterRecords& rec, CtaTally& tally,
+    int32_t buf, int32_t row, int32_t rowLen, int32_t* outRow, WaitCtx const& wait)
+{
+    int32_t const world = p.worldSize;
+    int32_t const sources = numSources(world);
+    int32_t const me = blockIdx.x;
+    int32_t const tid = threadIdx.x;
+    int32_t const lane = tid % 32;
+    int32_t const warp = tid / 32;
+    uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+    int32_t const vectorsPerSource = static_cast<int32_t>(cap) / 2;
+    int32_t const q0 = me * world;
+    int64_t* const dbg = debugSlots(p, 0);
+    // This thread's candidates (kEmptyWord: none; never a key).
+    uint32_t k0 = kEmptyWord;
+    uint32_t k1 = kEmptyWord;
+    if (tid == 0)
+    {
+        s.scal[kSlotCounter2] = 0u;
+    }
+    if (tid < world * vectorsPerSource)
+    {
+        int32_t const q = q0 + tid / vectorsPerSource;
+        int32_t const kv = tid % vectorsPerSource;
+        uint4 v;
+        uint2 const h = pollSlotVector(p.ucWorkspace + round1Offset(buf, q, row), kv, cap, v, wait);
+        if (2u * kv < h.x)
+        {
+            k0 = v.x;
+            if (2u * kv + 1u < h.x)
+            {
+                k1 = v.z;
+            }
+        }
+        // The tokens wait in shared memory until the emission (s.midx is free in the cluster merge).
+        *reinterpret_cast<uint2*>(s.midx + 2 * tid) = make_uint2(v.y, v.w);
+    }
+    else if (warp >= kPollWarps && warp < kFirstResetWarp)
+    {
+        // Header warps: thread h reads source h's header; per source its count (clamped to the round-1 capacity),
+        // bound and length, and per header warp the reduced key statistics.
+        int32_t const q = tid - kPollWarps * 32;
+        KeyStats st;
+        if (q < sources)
+        {
+            uint4 a;
+            uint4 b;
+            pollHeader(p.ucWorkspace + round1Offset(buf, q, row), a, b, wait);
+            s.scal[kSlotSrcCount + q] = min(a.x, cap);
+            s.scal[kSlotSrcBound + q] = a.y;
+            s.scal[kSlotSrcLen + q] = static_cast<uint32_t>(sourceLength(
+                min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank));
+            st.maxKey = a.z;
+            st.minKey = a.w - 1u;
+            st.positives = b.x;
+            st.maxNeg = b.y;
+            st.minPos = b.z - 1u;
+        }
+        uint32_t const cnt = q < sources ? s.scal[kSlotSrcCount + q] : 0u;
+        // Candidates and keys >= +0.0 of the cluster CTAs before this one (CTA c merges sources
+        // [c * world, (c + 1) * world)), of this one and of the row.
+        int32_t const c = q < sources ? q / world : kClusterCtas;
+        uint32_t const nBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? cnt : 0u);
+        uint32_t const nMine = __reduce_add_sync(0xFFFFFFFFu, c == me ? cnt : 0u);
+        uint32_t const nAll = __reduce_add_sync(0xFFFFFFFFu, cnt);
+        uint32_t const pBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? st.positives : 0u);
+        if (lane == 0)
+        {
+            atomicAdd(&s.scal[kSlotCountBefore], nBefore);
+            atomicAdd(&s.scal[kSlotCountMine], nMine);
+            atomicAdd(&s.scal[kSlotCountTotal], nAll);
+            atomicAdd(&s.scal[kSlotPositivesBefore], pBefore);
+        }
+        uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, st.maxKey);
+        uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, st.minKey);
+        uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, st.positives);
+        uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, st.maxNeg);
+        uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, st.minPos);
+        if (lane == 0)
+        {
+            // In the scan scratch (s.stats may still be read by the send's last warp; s.warpSums is next written
+            // by the emission scan, after the reads below).
+            uint32_t* const out = s.warpSums + kNumKeyStats * (warp - kPollWarps);
+            out[0] = v0;
+            out[1] = v1;
+            out[2] = v2;
+            out[3] = v3;
+            out[4] = v4;
+        }
+    }
+    stamp(dbg, 45);
+    __syncthreads();
+    markPhase(p, 4);
+
+    // Statistics of the row's candidates from every source's header (identical in every CTA and rank).
+    KeyStats g;
+#pragma unroll
+    for (int32_t w = 0; w < kHeaderWarps; ++w)
+    {
+        uint32_t const* const in = s.warpSums + kNumKeyStats * w;
+        g.maxKey = max(g.maxKey, in[0]);
+        g.minKey = min(g.minKey, in[1]);
+        g.positives += in[2];
+        g.maxNeg = max(g.maxNeg, in[3]);
+        g.minPos = min(g.minPos, in[4]);
+    }
+    uint32_t const total = s.scal[kSlotCountTotal];
+    stamp(dbg, 61);
+    markPhase(p, 5);
+
+    SignRange const r
+        = total > 0u ? signRange(g, static_cast<uint32_t>(kTopK)) : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false};
+    clusterInitWait();
+    SelectOut const m = clusterSelect(
+        s, rec, tally, k0, k1, r.lo, r.hi, static_cast<uint32_t>(kTopK), r.lo, r.bits, r.aboveHi, debugSlots(p, 48));
+    putDebug(debugSlots(p, 48), 0, s.scal[kSlotCountMine]);
+    putDebug(dbg, 62, total);
+    putDebug(dbg, 63, m.passes + (r.negative ? 10 : 0));
+    bool takeAll = false;
+    uint32_t const kth = mergeThreshold(m, takeAll);
+    markPhase(p, 6);
+
+    // Keys above the threshold and ties of the cluster CTAs before this one and of the row (identical in every CTA),
+    // from the select's tallies; with the threshold below a wholly taken bin the ties fall beyond kTopK and are taken
+    // as keys above it.
+    uint32_t gtBefore;
+    uint32_t eqBefore;
+    uint32_t allGt;
+    uint32_t allEq;
+    if (takeAll)
+    {
+        gtBefore = s.scal[kSlotCountBefore];
+        eqBefore = 0u;
+        allGt = total;
+        allEq = 0u;
+    }
+    else
+    {
+        uint32_t const aboveBefore = tally.aboveBefore + (r.aboveHi > 0u ? s.scal[kSlotPositivesBefore] : 0u);
+        gtBefore = m.full ? aboveBefore + tally.binBefore : aboveBefore;
+        eqBefore = m.full ? 0u : tally.binBefore;
+        allGt = m.full ? m.above + m.binCount : m.above;
+        allEq = m.full ? 0u : m.binCount;
+    }
+    // This CTA's positions, scanned in thread order; a source of the row that withheld keys above the threshold
+    // sends the row to the second round.
+    bool const gt0 = k0 != kEmptyWord && (takeAll || k0 > kth);
+    bool const eq0 = k0 != kEmptyWord && !takeAll && k0 == kth;
+    bool const gt1 = k1 != kEmptyWord && (takeAll || k1 > kth);
+    bool const eq1 = k1 != kEmptyWord && !takeAll && k1 == kth;
+    if (tid < sources && sourceOverflows(s, tid, kth))
+    {
+        atomicOr(&s.scal[kSlotCounter2], 1u);
+    }
+    // Exclusive prefix in thread order (k0 before k1): within the warp from ballots, across warps from the warp
+    // totals (one block barrier).
+    uint32_t const below = (1u << lane) - 1u;
+    uint32_t const bGt0 = __ballot_sync(0xFFFFFFFFu, gt0);
+    uint32_t const bGt1 = __ballot_sync(0xFFFFFFFFu, gt1);
+    uint32_t const bEq0 = __ballot_sync(0xFFFFFFFFu, eq0);
+    uint32_t const bEq1 = __ballot_sync(0xFFFFFFFFu, eq1);
+    if (lane == 0)
+    {
+        s.warpSums[warp] = static_cast<uint32_t>(__popc(bGt0) + __popc(bGt1))
+            + (static_cast<uint32_t>(__popc(bEq0) + __popc(bEq1)) << 16);
+    }
+    __syncthreads();
+    uint32_t const warpsBefore = __reduce_add_sync(0xFFFFFFFFu, lane < warp ? s.warpSums[lane] : 0u);
+    uint32_t const before = warpsBefore + static_cast<uint32_t>(__popc(bGt0 & below) + __popc(bGt1 & below))
+        + (static_cast<uint32_t>(__popc(bEq0 & below) + __popc(bEq1 & below)) << 16);
+    stamp(dbg, 58);
+    if (s.scal[kSlotCounter2] != 0u)
+    {
+        return true;
+    }
+    uint32_t const want = static_cast<uint32_t>(kTopK);
+    uint32_t posGt = gtBefore + (before & 0xFFFFu);
+    uint32_t posEq = allGt + eqBefore + (before >> 16);
+    uint32_t x0 = 0u;
+    uint32_t x1 = 0u;
+    if (k0 != kEmptyWord)
+    {
+        uint2 const x = *reinterpret_cast<uint2 const*>(s.midx + 2 * tid);
+        x0 = x.x;
+        x1 = x.y;
+    }
+    if (gt0 && posGt < want)
+    {
+        outRow[posGt] = static_cast<int32_t>(x0);
+    }
+    posGt += gt0 ? 1u : 0u;
+    if (eq0 && posEq < want)
+    {
+        outRow[posEq] = static_cast<int32_t>(x0);
+    }
+    posEq += eq0 ? 1u : 0u;
+    if (gt1 && posGt < want)
+    {
+        outRow[posGt] = static_cast<int32_t>(x1);
+    }
+    if (eq1 && posEq < want)
+    {
+        outRow[posEq] = static_cast<int32_t>(x1);
+    }
+    uint32_t const selected = min(allGt + allEq, want);
+    for (uint32_t j = selected + static_cast<uint32_t>(me * kThreads + tid); j < want; j += kClusterCtas * kThreads)
+    {
+        outRow[j] = -1;
+    }
+    stamp(dbg, 60);
+    return false;
+}
+
+__global__ void __launch_bounds__(kThreads, 1) cpTopkZ9Kernel(CpTopkParams p)
+{
+    markPhase(p, 0);
+    extern __shared__ __align__(16) uint32_t smemRaw[];
+    __shared__ uint32_t const* segSlots[kMaxSources];
+    __shared__ uint32_t segBegin[kMaxSources + 1];
+    __shared__ __align__(16) ClusterRecords rec;
+    __shared__ CtaTally tally;
+    __shared__ int32_t epochShared;
+    int32_t const world = p.worldSize;
+    int32_t const rank = p.rank;
+    int32_t const src = blockIdx.x;
+    int32_t const row = blockIdx.y;
+    SmemLayout const s = carveSmem(smemRaw, p.localCapacity, world);
+    int32_t* outRow = p.out + static_cast<int64_t>(row) * kTopK;
+    clusterInit(rec);
+    if (threadIdx.x <= kSlotPositivesBefore - kSlotCountBefore)
+    {
+        s.scal[kSlotCountBefore + threadIdx.x] = 0u;
+    }
+    int32_t const resetOffset = resetHeaderOffset();
+    // Every global read (lengths, logits, the epoch word, the workspace) comes after the grid dependency wait, so a
+    // programmatic early launch is safe whatever the preceding kernels write.
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    markPhase(p, 10);
+    // Thread 0 and the reset warps read this CTA's epoch word, written only by the previous exchange launch of this
+    // op, before the key loads; nothing waits for the value until the key loads are done. Thread 0 publishes it to the
+    // CTA at the block barrier that follows the key loads, and the reset warps start the copies of the previous
+    // buffer's headers right before that barrier (slot offsets computed before the wait, from the grid only).
+    int32_t* const epochWord = p.state + kEpochBase + row * kSourcesPerRank + src;
+    int32_t epochRead = 0;
+    if (threadIdx.x == 0 || static_cast<int32_t>(threadIdx.x) >= kFirstResetWarp * 32)
+    {
+        epochRead = *reinterpret_cast<int32_t volatile*>(epochWord);
+    }
+    int32_t const prevBufRead = (epochRead + kNumLamportBuffers - 1) % kNumLamportBuffers;
+    int32_t const kvLen = p.kvLens[0];
+    int32_t const rowLen = max(0, kvLen - p.nextN + row + 1);
+    // Identity launch (every row has at most kTopK tokens; the same decision on every rank): no exchange, and
+    // neither the workspace nor the epoch words are touched, so identity launches never shorten the reuse chain of
+    // the Lamport buffers (see crossRankTopkClusterMerge.h).
+    bool const exchange = kvLen - p.nextN + p.numRows > kTopK;
+    // Thread 0 advances this CTA's epoch word (the CTAs of the last row also those of the rows this launch does not
+    // have, so every word counts every exchange launch), after a block barrier that follows every read of the word in
+    // this CTA. Nothing else of this launch reads the epoch words, and the next launch reads them only after its grid
+    // dependency wait.
+    auto advanceEpoch = [&]()
+    {
+        if (exchange && threadIdx.x == 0)
+        {
+            int32_t const rows = row == static_cast<int32_t>(gridDim.y) - 1 ? kMaxRows - row : 1;
+            for (int32_t r = 0; r < rows; ++r)
+            {
+                epochWord[r * kSourcesPerRank] = epochRead + 1;
+            }
+        }
+    };
+
+    if (rowLen <= kTopK)
+    {
+        // Every rank agrees on this branch: identity selection, no exchange for this row.
+        if (src == 0)
+        {
+            for (int32_t j = threadIdx.x; j < kTopK; j += kThreads)
+            {
+                outRow[j] = j < rowLen ? j : -1;
+            }
+        }
+        if (exchange)
+        {
+            fetchResetHeader(p, s, prevBufRead, resetOffset);
+            resetSlots(p, s, prevBufRead);
+            markPhase(p, 9);
+        }
+        __syncthreads();
+        advanceEpoch();
+        clusterInitWait();
+    }
+    else
+    {
+        int32_t const localLen = min(localLength(rowLen, rank, world), p.localCapacity);
+        float const* rowLogits = p.logits + static_cast<int64_t>(row) * p.logitsStride;
+
+        // ---- Round 1: this source's local top keys ---------------------------------------------------------
+        int32_t const n = loadSource(s, rowLogits, localLen, src, true);
+        if (threadIdx.x == 0)
+        {
+            epochShared = epochRead;
+        }
+        // The reset warps start fetching the headers of their share of the previous exchange launch's buffer (an
+        // exchange row implies an exchange launch); the values are first used by the reset after the send.
+        fetchResetHeader(p, s, prevBufRead, resetOffset);
+        __syncthreads();
+        markPhase(p, 1);
+        advanceEpoch();
+        uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+        uint32_t const need = static_cast<uint32_t>(roundNeed(world));
+        SelectOut sel{0u, 0u, 0u, static_cast<uint32_t>(n), true, false, 0};
+        uint32_t selLo = 0u;
+        int32_t passInfo = 0;
+        if (n > static_cast<int32_t>(cap))
+        {
+            // The cut lies among the keys of one sign (see kPositiveKey); the digits span that sign's key range.
+            SignRange const r = signRange(readStats(s.stats), need);
+            selLo = r.lo;
+            SmemKeys const keys{s.chunk, nullptr};
+            sel = blockSelect(s, keys, n, r.lo, r.hi, need, cap, r.lo, r.bits, debugSlots(p, 32), r.aboveHi);
+            passInfo = sel.passes + (r.negative ? 10 : 0);
+        }
+        markPhase(p, 2);
+        int32_t const sendBuf = epochShared % kNumLamportBuffers;
+        uint32_t const sent = sendSource<false>(s, sel, n, selLo, uint64_t{kEmptyWord}, cap,
+            p.mcWorkspace + round1Offset(sendBuf, rank * kSourcesPerRank + src, row), src, rank, world,
+            reinterpret_cast<uint2*>(s.mkey), debugSlots(p, 0));
+        putDebug(debugSlots(p, 0), 46, sent);
+        putDebug(debugSlots(p, 0), 47, passInfo);
+        markPhase(p, 3);
+
+        // The reset warps reset the previous exchange launch's buffer, whose candidates every rank has consumed (a
+        // peer writes it again only two exchange launches later, after this launch completed on this rank), while
+        // the poll and header warps wait for this launch's candidates.
+        resetSlots(p, s, (epochShared + kNumLamportBuffers - 1) % kNumLamportBuffers);
+        int32_t const buf = epochShared % kNumLamportBuffers;
+        if (threadIdx.x == kFirstResetWarp * 32)
+        {
+            int64_t* const d = debugSlots(p, 0);
+            if (d != nullptr)
+            {
+                d[9] = static_cast<int64_t>(globalTimerNs());
+            }
+        }
+
+        // ---- Merge: the row's cluster -----------------------------------------------------------------------
+        WaitCtx const wait{p.state, p.timeoutNs, rank};
+        if (clusterMergeRow(p, s, rec, tally, buf, row, rowLen, outRow, wait) && src == 0)
+        {
+            // A source withheld keys above the threshold: crossRankTopk's merge, with its second round, on CTA 0.
+            mergeRow(p, s, segSlots, segBegin, buf, row, rowLen, localLen, rowLogits, outRow, wait);
+            int64_t* const dbg = debugSlots(p, 0);
+            putDebug(dbg, 63, dbg != nullptr ? dbg[63] + 100 : 0);
+        }
+        markPhase(p, 7);
+    }
+    markPhase(p, 8);
+}
+
+#endif
+
+} // namespace
+
+size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize)
+{
+    return smemBytes(localCapacity, worldSize);
+}
+
+int32_t cpTopkMaxLocalCapacity(int32_t worldSize)
+{
+    static int32_t cached[kMaxRanks + 1] = {};
+    if (worldSize < 1 || worldSize > kMaxRanks)
+    {
+        return 0;
+    }
+    if (cached[worldSize] == 0)
+    {
+        int device = 0;
+        cudaGetDevice(&device);
+        int maxSmemOptin = 0;
+        cudaDeviceGetAttribute(&maxSmemOptin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+        int32_t cap = 0;
+        for (int32_t c = 1 << 22; c >= kPageTokens * kSourcesPerRank; c -= kPageTokens * kSourcesPerRank)
+        {
+            if (smemBytes(c, worldSize) + kStaticSmemBytes <= static_cast<size_t>(maxSmemOptin))
+            {
+                cap = c;
+                break;
+            }
+        }
+        cached[worldSize] = cap;
+    }
+    return cached[worldSize];
+}
+
+int32_t cpTopkMaxActiveClusters(int32_t worldSize)
+{
+    // Queried once per world size and device (the value + 1; 0 = not queried yet).
+    constexpr int32_t kMaxDevices = 64;
+    static int32_t cached[kMaxDevices][kMaxRanks + 1] = {};
+    if (worldSize < 1 || worldSize > kMaxRanks)
+    {
+        return 0;
+    }
+    int device = 0;
+    cudaGetDevice(&device);
+    if (device >= 0 && device < kMaxDevices && cached[device][worldSize] > 0)
+    {
+        return cached[device][worldSize] - 1;
+    }
+    int32_t const capacity = cpTopkMaxLocalCapacity(worldSize);
+    if (capacity == 0)
+    {
+        return 0;
+    }
+    size_t const smem = smemBytes(capacity, worldSize);
+    if (cudaFuncSetAttribute(cpTopkZ9Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem))
+        != cudaSuccess)
+    {
+        (void) cudaGetLastError();
+        return 0;
+    }
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSourcesPerRank, kMaxRows);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = smem;
+    cudaLaunchAttribute attr{};
+    attr.id = cudaLaunchAttributeClusterDimension;
+    attr.val.clusterDim.x = kSourcesPerRank;
+    attr.val.clusterDim.y = 1;
+    attr.val.clusterDim.z = 1;
+    config.attrs = &attr;
+    config.numAttrs = 1;
+    int clusters = 0;
+    if (cudaOccupancyMaxActiveClusters(&clusters, cpTopkZ9Kernel, &config) != cudaSuccess)
+    {
+        (void) cudaGetLastError();
+        return 0;
+    }
+    if (device >= 0 && device < kMaxDevices)
+    {
+        cached[device][worldSize] = clusters + 1;
+    }
+    return clusters;
+}
+
+int64_t cpTopkDefaultTimeoutMs()
+{
+    static int64_t const timeoutMs = []()
+    {
+        char const* env = std::getenv("TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS");
+        if (env == nullptr || env[0] == '\0')
+        {
+            return kDefaultTimeoutMs;
+        }
+        char* end = nullptr;
+        long long const v = std::strtoll(env, &end, 10);
+        if (end == env || *end != '\0' || v <= 0)
+        {
+            throw std::invalid_argument(std::string("glm_moe_dsa_cross_rank_topk_cluster_merge: bad "
+                                                    "TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS=")
+                + env);
+        }
+        return static_cast<int64_t>(v);
+    }();
+    return timeoutMs;
+}
+
+void invokeCpTopk(CpTopkParams const& inParams, cudaStream_t stream)
+{
+    if (inParams.numRows <= 0)
+    {
+        return;
+    }
+    CpTopkParams params = inParams;
+    if (params.timeoutNs == 0)
+    {
+        params.timeoutNs = static_cast<uint64_t>(cpTopkDefaultTimeoutMs()) * 1000000ull;
+    }
+    if (params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks || params.rank < 0
+        || params.rank >= params.worldSize)
+    {
+        throw std::invalid_argument("glm_moe_dsa_cross_rank_topk_cluster_merge: unsupported rows/world/rank");
+    }
+    size_t const smem = smemBytes(params.localCapacity, params.worldSize);
+    // The attribute is per function and device; raise it only when a larger footprint is requested.
+    static thread_local size_t configuredSmem[64] = {};
+    int device = 0;
+    cudaGetDevice(&device);
+    cudaError_t err = cudaSuccess;
+    if (device >= 64 || configuredSmem[device] < smem)
+    {
+        err = cudaFuncSetAttribute(cpTopkZ9Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem));
+        if (err != cudaSuccess)
+        {
+            throw std::runtime_error(
+                std::string("glm_moe_dsa_cross_rank_topk_cluster_merge: smem attribute: ") + cudaGetErrorString(err));
+        }
+        if (device < 64)
+        {
+            configuredSmem[device] = smem;
+        }
+    }
+    // Programmatic dependent launch as for the other TensorRT-LLM kernels (on unless TRTLLM_ENABLE_PDL=0).
+    static bool const enablePdl = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && env[0] == '0' && env[1] == '\0');
+    }();
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSourcesPerRank, params.numRows);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[2];
+    attrs[0].id = cudaLaunchAttributeClusterDimension;
+    attrs[0].val.clusterDim.x = kSourcesPerRank;
+    attrs[0].val.clusterDim.y = 1;
+    attrs[0].val.clusterDim.z = 1;
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    err = cudaLaunchKernelEx(&config, cpTopkZ9Kernel, params);
+    if (err == cudaSuccess)
+    {
+        err = cudaGetLastError();
+    }
+    if (err != cudaSuccess)
+    {
+        throw std::runtime_error(
+            std::string("glm_moe_dsa_cross_rank_topk_cluster_merge: launch: ") + cudaGetErrorString(err));
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h
new file mode 100644
index 000000000000..0f6d53a52e92
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h
@@ -0,0 +1,191 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
+#include <cuda_runtime_api.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge
+{
+
+// Context-partitioned DSA indexer top-k with a cluster merge.
+//
+// It selects exactly what the kernel of glmMoeDsa/crossRankTopk.h selects, in the same output order, with the
+// same context partition, exchange sources, round-1 / round-2 protocol and Lamport workspace layout:
+//   - global 64-token page p belongs to rank p % worldSize; within a rank, each verify row is handled by
+//     kSourcesPerRank CTAs, CTA s owning the 8-token block (s + page) % kSourcesPerRank of every local page
+//     (exchange source q = rank * kSourcesPerRank + s);
+//   - each source selects its local top keys with crossRankTopk's rule and broadcasts them as (order-preserving key,
+//     global token) pairs in source order, plus a header (count, exclusive bound of the withheld keys; this kernel adds
+//     the key statistics of the keys sent), through the multicast Lamport workspace.
+// What differs is how the row is merged and how a launch ends:
+//   - the kSourcesPerRank CTAs of one row form a thread-block cluster. CTA s merges the candidates of sources
+//     [s * worldSize, (s + 1) * worldSize) (at 8 ranks: the 8 sources of rank s). The cluster runs crossRankTopk's
+//     merge select (the same digits and early-stop rules) on the cluster-summed histograms, and every CTA writes its
+//     selected candidates directly into the output row at crossRankTopk's positions (keys above the threshold in
+//     candidate order, then ties in candidate order, then -1). Each round-1 slot header carries its source's key
+//     statistics, so every CTA derives the row's statistics from the headers it polls; the CTAs exchange only the
+//     select's histogram totals, with st.async stores into each other's shared memory completed on transaction
+//     barriers, and derive every CTA's output counts from the per-CTA tallies of the select. The only cluster-wide
+//     barrier orders the barriers' initialization before the first exchange. A row in which a source withheld keys
+//     above the threshold (identical decision on every rank) is merged instead by CTA 0 with crossRankTopk's single-CTA
+//     merge, including its second round;
+//   - every CTA keeps its own epoch word (one per (row, source CTA) slot): its thread 0 reads the word after the grid
+//     dependency wait (only the previous exchange launch of this op writes it), the value is published to the CTA at
+//     the block barrier that follows the key loads, and the CTA advances the word with a plain store after that
+//     barrier (nothing else in the launch reads it; the next launch reads it after its own grid dependency wait); the
+//     CTAs of the last row also advance the words of the rows the launch does not have, so every word counts the
+//     exchange launches. There is no arrival counter and no end-of-kernel fence.
+//
+// Rows of at most kTopK valid tokens produce the identity selection followed by -1, without an exchange. A launch whose
+// rows all have at most kTopK valid tokens (the same on every rank: the lengths come from the shared KV length) is an
+// identity launch: it writes neither the workspace nor the epoch words. Every other launch is an exchange launch.
+// Exchange launches rotate kNumLamportBuffers buffers by the epoch, which counts exchange launches; in exchange launch
+// x every CTA resets its share of its own rank's copy of the buffer of exchange launch x - 1 between its send and its
+// merge (the headers of its share are fetched once the epoch is known, after the grid dependency wait). Every rank's
+// merges of x - 1 have read their copy of it, and a peer writes it again only in exchange launch x + 2, which starts
+// after the peer's merges of x + 1 received this rank's sources of x + 1, i.e. after this rank completed launch x
+// (identity launches in between neither write nor wait, so they do not shorten this chain). Every global memory
+// access of the kernel follows its grid dependency wait, which orders it after the completion (and memory
+// visibility) of the preceding kernels of the stream, including the previous launch of this op; the kernel never
+// triggers its dependents early.
+//
+// A merge that waits longer than the timeout (CpTopkParams::timeoutNs, %globaltimer) for a peer's words records the
+// wait in the state's error word and aborts the kernel with a trap: the exchange cannot complete correctly any more,
+// and a late peer write could land in a buffer this rank reuses, so the run fails instead of continuing.
+
+//! Selected indices per verify row (DSA index_topk).
+constexpr int32_t kTopK = 2048;
+//! Ranks of the context-parallel group.
+constexpr int32_t kMaxRanks = 8;
+//! Verify rows per call.
+constexpr int32_t kMaxRows = 8;
+//! Tokens per indexer K page; the context is partitioned by whole pages.
+constexpr int32_t kPageTokens = 64;
+//! CTAs (exchange sources) per rank and verify row; also the cluster size (one cluster per row).
+constexpr int32_t kSourcesPerRank = 8;
+//! Exchange sources per row.
+constexpr int32_t kMaxSources = kMaxRanks * kSourcesPerRank;
+//! Round-1 slots per (source, row).
+constexpr int32_t kSourceCapacity = 256;
+//! Keys each saturated source aims to send in round 1 at 8 ranks (crossRankTopk's rule: three times the even share of
+//! kTopK; refining stops at the first bin edge with at most four times the even share of keys at or above it).
+constexpr int32_t kSourceBudget = 96;
+//! Round-2 slots per (source, row).
+constexpr int32_t kRound2Capacity = 2048;
+//! Threads per CTA.
+constexpr int32_t kThreads = 1024;
+//! Lamport buffers rotated by the epoch.
+constexpr int32_t kNumLamportBuffers = 3;
+//! 32-bit word that marks an empty Lamport slot; never a valid key, token index, count or bound.
+constexpr uint32_t kEmptyWord = 0xFFFFFFFFu;
+
+//! 32-bit words of one (source, row) slot: a 32-byte header then two words per candidate, so candidates can be written
+//! and read with 16-byte vector accesses. Round-1 header: count, exclusive bound of the withheld keys, then the key
+//! statistics of the keys sent (largest key, smallest key + 1, keys >= +0.0, largest key below +0.0, smallest key
+//! >= +0.0 plus 1; the "+ 1" maps the "none" value 0xFFFFFFFF to 0), then 0. Round-2 header: count, bound.
+constexpr int64_t kSlotHeaderWords = 8;
+constexpr int64_t kRound1SlotWords = kSlotHeaderWords + 2 * int64_t{kSourceCapacity};
+constexpr int64_t kRound2SlotWords = kSlotHeaderWords + 2 * int64_t{kRound2Capacity};
+//! 32-bit words of one Lamport buffer (round-1 region, then round-2 region).
+constexpr int64_t kBufferWords = int64_t{kMaxSources} * kMaxRows * (kRound1SlotWords + kRound2SlotWords);
+//! 32-bit words of the whole exchange workspace; every word must start as kEmptyWord.
+constexpr int64_t kWorkspaceWords = kNumLamportBuffers * kBufferWords;
+//! int32 words of the per-rank device state, zero-initialized: [0, 2) unused, [2] second rounds taken, [3] error word
+//! (nonzero = the wait that timed out, written right before the kernel aborts), then one epoch word per
+//! (row, source CTA) slot: state[kEpochBase + row * kSourcesPerRank + s].
+constexpr int32_t kEpochBase = 4;
+constexpr int32_t kStateWords = kEpochBase + kMaxRows * kSourcesPerRank;
+//! Default wait for a peer's Lamport words before the kernel aborts
+//! (TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS overrides it).
+constexpr int64_t kDefaultTimeoutMs = 120000;
+
+//! Values of the error word.
+enum CpTopkWait : int32_t
+{
+    kWaitRound1Header = 1,
+    kWaitRound1Candidates = 2,
+    kWaitRound2Header = 3,
+    kWaitRound2Candidates = 4
+};
+
+struct CpTopkParams
+{
+    //! [numRows, logitsStride] fp32 logits of this rank's local columns.
+    float const* logits;
+    int64_t logitsStride;
+    //! Local columns the logits view holds (each row's local length is clamped to it).
+    int32_t localCapacity;
+    //! [1] int32 global indexer KV length of the request; row r has kvLen - nextN + r + 1 valid tokens.
+    int32_t const* kvLens;
+    //! [numRows, kTopK] int32 output: global token indices, -1 for unused slots.
+    int32_t* out;
+    int32_t numRows;
+    int32_t nextN;
+    int32_t rank;
+    int32_t worldSize;
+    //! Multicast address of the exchange workspace (kWorkspaceWords words, identical layout on every rank).
+    uint32_t* mcWorkspace;
+    //! This rank's unicast address of the same workspace.
+    uint32_t* ucWorkspace;
+    //! [kStateWords] int32 per-rank device state, zero-initialized.
+    int32_t* state;
+    //! Optional [numRows, kSourcesPerRank, kDebugSlots] int64 phase timestamps (diagnosis only; nullptr = off).
+    int64_t* debugTimes = nullptr;
+    //! Longest wait for a peer's words before the kernel aborts (ns); 0 = the default (kDefaultTimeoutMs, or
+    //! TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS).
+    uint64_t timeoutNs = 0;
+};
+
+//! Timestamps per CTA in CpTopkParams::debugTimes: [0, 16) %globaltimer (ns) and [16, 32) clock64 per phase (0 start,
+//! 10 grid dependency wait returned, 1 keys loaded, 2 local select, 3 sent, 9 previous buffer reset (%globaltimer of
+//! the first reset warp only), 4 candidates polled, 5 cluster statistics agreed, 6 merge select, 7 row written, 8 end),
+//! [32, 42) local select (keys, then clock64 at the start of / after the histogram of / after the bin search of each
+//! pass), [42, 45) send (start, keys taken, stores issued), 46 round-1 keys sent by this source, 47 local select passes
+//! (+10 when the cut was searched among the negative keys), [48, 58) merge select (keys = this CTA's candidates; per
+//! pass clock64 at its start / after the group totals arrived / after the bin was decided), clock64 of the merge: 45
+//! candidates polled, 61 statistics combined, 58 emission positions scanned, 60 row written; 62 candidates of the row
+//! (all CTAs), 63 merge select passes (+10 likewise; +100 when the row took crossRankTopk's single-CTA merge, which
+//! then reuses 45 and 58-61 as crossRankTopk does); clock64 and values of thread 0.
+constexpr int32_t kDebugSlots = 64;
+
+//! Effective wait timeout of a launch with params.timeoutNs == 0 (ms).
+int64_t cpTopkDefaultTimeoutMs();
+
+//! Dynamic shared memory the kernel needs (bytes).
+size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize);
+
+//! Largest local capacity the kernel accepts for a world size.
+int32_t cpTopkMaxLocalCapacity(int32_t worldSize);
+
+//! Clusters of the kernel's launch configuration (largest local capacity) that can be resident at once on this
+//! device (cudaOccupancyMaxActiveClusters, queried once per world size and device); 0 when the query fails. No row
+//! waits for another row of its rank, but a row's cluster waits for the same row's sources on every peer rank, and
+//! the ranks may make their rows resident in different orders; the caller therefore requires kMaxRows co-resident
+//! clusters, so every cluster of a launch can be resident at once (other kernels sharing the SMs are finite and do
+//! not wait for this one).
+int32_t cpTopkMaxActiveClusters(int32_t worldSize);
+
+void invokeCpTopk(CpTopkParams const& params, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu
new file mode 100644
index 000000000000..685a2da67a97
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu
@@ -0,0 +1,2757 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdlib>
+#include <stdexcept>
+#include <string>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit
+{
+namespace
+{
+
+constexpr int32_t kBins = 2048;
+static_assert(kBins >= kTopK, "the output row is staged in the histogram");
+constexpr int32_t kWarps = kThreads / 32;
+//! Per-warp key statistics kept in shared memory (see KeyStats).
+constexpr int32_t kNumKeyStats = 5;
+//! Tokens of each page one source holds.
+constexpr int32_t kBlockTokens = kPageTokens / kSourcesPerRank;
+//! CTAs of one row's cluster.
+constexpr int32_t kClusterCtas = kSourcesPerRank;
+//! Histogram bins per group of the cluster bin search (one warp lane per bin).
+constexpr int32_t kGroupBins = 32;
+constexpr int32_t kMaxGroups = kBins / kGroupBins;
+//! Static shared memory of the kernel (round-2 segment tables and the cluster exchange records), kept out of the
+//! dynamic budget.
+constexpr size_t kStaticSmemBytes = 8192;
+
+// Records the CTAs of a row's cluster send each other with st.async (static, so every CTA has them at the same
+// address). Exchange x of a row completes on barrier bar[x % 2], phase x / 2; the record sets alternate by pass
+// parity and are written again only after every CTA has completed the exchange in between.
+struct ClusterRecords
+{
+    //! Transaction barriers of the exchanges (arrival count 1: this CTA's expect-tx).
+    uint64_t bar[2];
+    //! [pass parity][source CTA][group] histogram totals per group of kGroupBins bins.
+    uint32_t groups[2][kClusterCtas][kMaxGroups];
+    //! [pass parity][source CTA][bin] counts of the selected group's bins, ascending.
+    uint32_t bins[2][kClusterCtas][kGroupBins];
+};
+
+// Scalar slots in shared memory.
+enum ScalarSlot : int32_t
+{
+    kSlotBin = 0,
+    kSlotAbove,
+    kSlotBinCount,
+    kSlotTotal,
+    kSlotCounter2,
+    kSlotSrcCount,                                       // [kMaxSources]
+    kSlotSrcBound = kSlotSrcCount + kMaxSources,         // [kMaxSources]
+    kSlotSrcLen = kSlotSrcBound + kMaxSources,           // [kMaxSources]
+    kSlotSrcOffset = kSlotSrcLen + kMaxSources,          // [kMaxSources + 1]
+    kSlotCountBefore = kSlotSrcOffset + kMaxSources + 1, // candidates of the cluster CTAs before this one
+    kSlotCountMine,                                      // candidates of this cluster CTA
+    kSlotCountTotal,                                     // candidates of the row
+    kSlotPositivesBefore,                                // candidates >= +0.0 of the cluster CTAs before this one
+    kNumScalarSlots
+};
+
+__host__ __device__ inline int32_t numSources(int32_t world)
+{
+    return world * kSourcesPerRank;
+}
+
+// Even share of kTopK per source.
+__host__ __device__ inline int32_t evenShare(int32_t world)
+{
+    return (kTopK + numSources(world) - 1) / numSources(world);
+}
+
+// Most keys a source sends in round 1 (an even number).
+__host__ __device__ inline int32_t roundCap(int32_t world)
+{
+    return min(kSourceCapacity, 4 * evenShare(world));
+}
+
+// Keys one source can hold: one block of every local page.
+__host__ __device__ inline int32_t chunkCapacity(int32_t localCapacity)
+{
+    int32_t const pages = (localCapacity + kPageTokens - 1) / kPageTokens;
+    return pages * kBlockTokens;
+}
+
+__host__ __device__ inline int32_t mergeCapacity(int32_t world)
+{
+    return numSources(world) * roundCap(world);
+}
+
+__host__ __device__ inline size_t smemBytes(int32_t localCapacity, int32_t world)
+{
+    size_t words = 2 * size_t{kBins} + 2 * size_t{kWarps} + size_t{kNumKeyStats} * kWarps + size_t{kNumScalarSlots};
+    words = (words + 3) / 4 * 4;
+    words += static_cast<size_t>(chunkCapacity(localCapacity)) + 2 * static_cast<size_t>(mergeCapacity(world));
+    return words * sizeof(uint32_t);
+}
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900)
+// The kernel targets SM90+ (multicast Lamport exchange, thread-block clusters).
+__global__ void __launch_bounds__(kThreads, 1) cpTopkZ11Kernel(CpTopkParams p)
+{
+    __trap();
+}
+#else
+
+//! Bits of one radix-select digit (kBins bins).
+constexpr int32_t kDigitBits = 11;
+//! Order-preserving key of +0.0: the keys of non-negative values are at least this, those of negative values below.
+//! Real indexer logits are mostly negative, and the cut of a row (and of many sources) often lies among small
+//! negative values while the largest keys are positive: one key range across zero spans ~2^31 key units, where a
+//! radix digit has no resolution. The selects therefore work on one sign at a time: with fewer than `need` keys at
+//! or above kPositiveKey these are all taken and the cut is searched among the negative keys only, else among the
+//! non-negative keys only; each range is narrow enough for one or two passes.
+constexpr uint32_t kPositiveKey = 0x80000000u;
+//! Independent 16-byte candidate loads one merge thread issues before it waits on any of them.
+constexpr int32_t kPollBatch = 4;
+//! Polls of a still-empty Lamport word between two reads of the global timer.
+constexpr int32_t kPollsPerTimerRead = 256;
+//! Warps of the cluster merge: warps [0, kPollWarps) poll the round-1 candidates of this CTA's sources (one 16-byte
+//! vector per thread: at most 520 vectors per CTA, at world sizes 5 to 7), the next kHeaderWarps warps poll the
+//! round-1 headers of every source (one per thread), and the remaining kResetWarps warps reset the previous exchange
+//! launch's buffer meanwhile.
+constexpr int32_t kPollWarps = 17;
+constexpr int32_t kHeaderWarps = 2;
+constexpr int32_t kFirstResetWarp = kPollWarps + kHeaderWarps;
+constexpr int32_t kResetWarps = kWarps - kFirstResetWarp;
+static_assert(kHeaderWarps * 32 >= kMaxSources, "one header per header-warp thread");
+
+// Keys a saturated source aims to send in round 1 (kSourceBudget at 8 ranks).
+__device__ __forceinline__ int32_t roundNeed(int32_t world)
+{
+    return min(kSourceCapacity, 3 * evenShare(world));
+}
+
+// Tokens of `rank` among the global tokens [0, n).
+__host__ __device__ __forceinline__ int32_t localLength(int32_t n, int32_t rank, int32_t world)
+{
+    if (n <= 0)
+    {
+        return 0;
+    }
+    int32_t const fullPages = n / kPageTokens;
+    int32_t const rem = n % kPageTokens;
+    int32_t const ownedFull = (fullPages - rank + world - 1) / world;
+    int32_t len = ownedFull * kPageTokens;
+    if (rem > 0 && fullPages % world == rank)
+    {
+        len += rem;
+    }
+    return len;
+}
+
+// Source `src` of a rank holds one block of kBlockTokens consecutive tokens of every local page: block
+// (src + page) % kSourcesPerRank. A fully selected page is thus spread over all sources of its rank.
+// Keys of source `src` among `localLen` local columns (only the last local page can be partial).
+__host__ __device__ __forceinline__ int32_t sourceLength(int32_t localLen, int32_t src)
+{
+    int32_t const pages = (localLen + kPageTokens - 1) / kPageTokens;
+    if (pages == 0)
+    {
+        return 0;
+    }
+    int32_t const last = pages - 1;
+    int32_t const tail = localLen - last * kPageTokens - ((src + last) % kSourcesPerRank) * kBlockTokens;
+    return last * kBlockTokens + min(max(tail, 0), kBlockTokens);
+}
+
+// Global token of a local column under the page-interleaved partition.
+__device__ __forceinline__ uint32_t localToGlobal(int32_t col, int32_t rank, int32_t world)
+{
+    int32_t const page = col / kPageTokens;
+    return static_cast<uint32_t>((page * world + rank) * kPageTokens + (col % kPageTokens));
+}
+
+struct SmemLayout
+{
+    uint32_t* hist;     // [2 * kBins] radix histograms of alternate passes; the first one stages the output row
+    uint32_t* warpSums; // [2 * kWarps]
+    uint32_t* stats;    // [kNumKeyStats * kWarps] per-warp key statistics
+    uint32_t* scal;     // [kNumScalarSlots]
+    uint32_t* chunk;    // [chunkCapacity] order-preserving keys of this source
+    uint32_t* mkey;     // [mergeCapacity] staged round-1 pairs, then the merge candidates (merge CTA)
+    uint32_t* midx;     // [mergeCapacity]
+};
+
+__device__ inline SmemLayout carveSmem(uint32_t* base, int32_t localCapacity, int32_t world)
+{
+    SmemLayout s;
+    uint32_t* p = base;
+    s.hist = p;
+    p += 2 * kBins;
+    s.warpSums = p;
+    p += 2 * kWarps;
+    s.stats = p;
+    p += kNumKeyStats * kWarps;
+    s.scal = p;
+    p += kNumScalarSlots;
+    p = base + ((p - base) + 3) / 4 * 4;
+    s.chunk = p;
+    p += chunkCapacity(localCapacity);
+    s.mkey = p;
+    p += mergeCapacity(world);
+    s.midx = p;
+    return s;
+}
+
+__device__ __forceinline__ uint32_t orderedKey(float v)
+{
+    uint32_t const b = __float_as_uint(v);
+    uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u);
+    // NaN keys above +inf are clamped so that no key, and no exclusive bound key + 1, equals the empty marker.
+    return k < kEmptyWord - 1u ? k : kEmptyWord - 2u;
+}
+
+// Bits of x (0 for 0): the digit range of keys in [base, base + x].
+__device__ __forceinline__ int32_t bitLength(uint32_t x)
+{
+    return x == 0u ? 0 : 32 - __clz(x);
+}
+
+// Statistics of a set of keys: extremes overall and per sign, and the count of non-negative keys.
+struct KeyStats
+{
+    uint32_t maxKey = 0u;         // largest key (0 when none)
+    uint32_t minKey = kEmptyWord; // smallest key (kEmptyWord when none)
+    uint32_t positives = 0u;      // keys >= kPositiveKey
+    uint32_t maxNeg = 0u;         // largest key below kPositiveKey (0 when none)
+    uint32_t minPos = kEmptyWord; // smallest key >= kPositiveKey (kEmptyWord when none)
+
+    __device__ __forceinline__ void add(uint32_t k)
+    {
+        bool const pos = k >= kPositiveKey;
+        maxKey = max(maxKey, k);
+        minKey = min(minKey, k);
+        positives += pos ? 1u : 0u;
+        maxNeg = pos ? maxNeg : max(maxNeg, k);
+        minPos = pos ? min(minPos, k) : minPos;
+    }
+};
+
+// Reduce a thread's statistics over its warp; lane 0 stores them in the per-warp array (read with readStats after
+// the next block barrier).
+__device__ __forceinline__ void storeWarpStats(KeyStats const& t, uint32_t* stats)
+{
+    uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, t.maxKey);
+    uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, t.minKey);
+    uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, t.positives);
+    uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg);
+    uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, t.minPos);
+    if (threadIdx.x % 32 == 0)
+    {
+        int32_t const w = threadIdx.x / 32;
+        stats[w] = v0;
+        stats[kWarps + w] = v1;
+        stats[2 * kWarps + w] = v2;
+        stats[3 * kWarps + w] = v3;
+        stats[4 * kWarps + w] = v4;
+    }
+}
+
+// Block statistics from the per-warp array (every warp reduces it; no barrier).
+__device__ __forceinline__ KeyStats readStats(uint32_t const* stats)
+{
+    int32_t const lane = threadIdx.x % 32;
+    KeyStats b;
+    b.maxKey = __reduce_max_sync(0xFFFFFFFFu, stats[lane]);
+    b.minKey = __reduce_min_sync(0xFFFFFFFFu, stats[kWarps + lane]);
+    b.positives = __reduce_add_sync(0xFFFFFFFFu, stats[2 * kWarps + lane]);
+    b.maxNeg = __reduce_max_sync(0xFFFFFFFFu, stats[3 * kWarps + lane]);
+    b.minPos = __reduce_min_sync(0xFFFFFFFFu, stats[4 * kWarps + lane]);
+    return b;
+}
+
+// Digit setup of a select of the need-th largest key: over the non-negative keys when there are at least `need` of
+// them, else over the negative keys with every non-negative key counted as taken (`aboveHi`).
+struct SignRange
+{
+    uint32_t lo;      // smallest key of the range (also the digit base)
+    uint64_t hi;      // exclusive end of the histogrammed range
+    int32_t bits;     // digit range of the range
+    uint32_t aboveHi; // keys at or above hi, all taken
+    bool negative;    // the cut is among the negative keys
+};
+
+__device__ __forceinline__ SignRange signRange(KeyStats const& st, uint32_t need)
+{
+    if (st.positives >= need || st.positives == 0u)
+    {
+        uint32_t const lo = st.positives >= need ? st.minPos : st.minKey;
+        return SignRange{lo, uint64_t{kEmptyWord}, bitLength(st.maxKey - lo), 0u, st.positives == 0u};
+    }
+    return SignRange{st.minKey, uint64_t{kPositiveKey}, bitLength(st.maxNeg - st.minKey), st.positives, true};
+}
+
+__device__ __forceinline__ int64_t round1Offset(int32_t buf, int32_t src, int32_t row)
+{
+    return int64_t{buf} * kBufferWords + (int64_t{src} * kMaxRows + row) * kRound1SlotWords;
+}
+
+__device__ __forceinline__ int64_t round2Offset(int32_t buf, int32_t src, int32_t row)
+{
+    return int64_t{buf} * kBufferWords + int64_t{kMaxSources} * kMaxRows * kRound1SlotWords
+        + (int64_t{src} * kMaxRows + row) * kRound2SlotWords;
+}
+
+__device__ __forceinline__ uint2 loadVolatile2(uint32_t const* ptr)
+{
+    uint2 v;
+    asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(ptr) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadVolatile4(uint32_t const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ void store2(uint32_t* ptr, uint32_t a, uint32_t b)
+{
+    asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(ptr), "r"(a), "r"(b) : "memory");
+}
+
+__device__ __forceinline__ void store4(uint32_t* ptr, uint4 v)
+{
+    asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)
+                 : "memory");
+}
+
+__device__ __forceinline__ bool isEmpty(uint2 v)
+{
+    return v.x == kEmptyWord || v.y == kEmptyWord;
+}
+
+// Whether the first pair of a 16-byte candidate vector, and the second one when `second`, were written.
+__device__ __forceinline__ bool vectorReady(uint4 v, bool second)
+{
+    return v.x != kEmptyWord && v.y != kEmptyWord && (!second || (v.z != kEmptyWord && v.w != kEmptyWord));
+}
+
+__device__ __forceinline__ uint64_t globalTimerNs()
+{
+    uint64_t t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return t;
+}
+
+// What a polling thread needs to give up: the error word, the timeout and this rank (for the report).
+struct WaitCtx
+{
+    int32_t* state;
+    uint64_t timeoutNs;
+    int32_t rank;
+};
+
+// A peer did not deliver within the timeout. The selection cannot be completed correctly any more, and the peer's
+// late writes could land in a buffer this rank resets or reuses, so record which wait failed and abort the kernel:
+// the resulting CUDA error fails the run instead of letting any rank continue with a partial selection.
+__device__ __noinline__ void abortOnTimeout(WaitCtx const& w, int32_t what)
+{
+    bool const first = atomicExch(&w.state[3], what) == 0;
+    __threadfence_system();
+    if (first)
+    {
+        printf(
+            "[glm_moe_dsa] crossRankTopkGlobalEmit: rank %d, row %d: wait %d for a peer exceeded %llu ms; aborting the "
+            "kernel\n",
+            w.rank, static_cast<int32_t>(blockIdx.y), what, static_cast<unsigned long long>(w.timeoutNs / 1000000ull));
+    }
+    __trap();
+}
+
+// Spin until both words of a Lamport pair were written; aborts the kernel after the timeout.
+__device__ uint2 pollPair(uint32_t const* ptr, WaitCtx const& w, int32_t what)
+{
+    uint2 v = loadVolatile2(ptr);
+    uint64_t start = 0;
+    for (int32_t polls = 1; isEmpty(v); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, what);
+            }
+        }
+        v = loadVolatile2(ptr);
+    }
+    return v;
+}
+
+// Spin until a 16-byte candidate vector was written (see vectorReady); aborts the kernel after the timeout.
+__device__ uint4 pollVector(uint32_t const* ptr, bool second, WaitCtx const& w, int32_t what)
+{
+    uint4 v = loadVolatile4(ptr);
+    uint64_t start = 0;
+    for (int32_t polls = 1; !vectorReady(v, second); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, what);
+            }
+        }
+        v = loadVolatile4(ptr);
+    }
+    return v;
+}
+
+// Block-wide inclusive scan of one value per thread. Returns the inclusive prefix; `total` gets the block sum.
+__device__ uint32_t blockInclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    uint32_t x = v;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    if (lane == 31)
+    {
+        warpSums[warp] = x;
+    }
+    __syncthreads();
+    if (warp == 0)
+    {
+        uint32_t w = warpSums[lane];
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, w, o);
+            if (lane >= o)
+            {
+                w += y;
+            }
+        }
+        warpSums[kWarps + lane] = w;
+    }
+    __syncthreads();
+    total = warpSums[2 * kWarps - 1];
+    uint32_t const before = warp > 0 ? warpSums[kWarps + warp - 1] : 0u;
+    return x + before;
+}
+
+// Exclusive block-wide scan with a single barrier: every warp scans the per-warp sums itself. Returns this
+// thread's exclusive prefix; `total` gets the block sum. warpSums must not be rewritten before the next barrier.
+__device__ uint32_t blockExclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    uint32_t x = v;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    if (lane == 31)
+    {
+        warpSums[warp] = x;
+    }
+    __syncthreads();
+    uint32_t const w = warpSums[lane];
+    uint32_t y = w;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const z = __shfl_up_sync(0xFFFFFFFFu, y, o);
+        if (lane >= o)
+        {
+            y += z;
+        }
+    }
+    total = __shfl_sync(0xFFFFFFFFu, y, 31);
+    return __shfl_sync(0xFFFFFFFFu, y - w, warp) + x - v;
+}
+
+// Find, counting from the top bin, the bin of `hist` that holds the target-th largest key, with two barriers:
+// per-warp sums, then each warp scans them and only the warp holding the target scans its own bins. Returns the
+// total count (every thread); when target <= total, the bin, the count above it and its own count go to the
+// scalar slots.
+__device__ uint32_t findBin(SmemLayout const& s, uint32_t const* hist, int32_t nbins, uint32_t target)
+{
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    int32_t const perThread = (nbins + kThreads - 1) / kThreads;
+    int32_t const top = nbins - 1 - perThread * static_cast<int32_t>(threadIdx.x);
+    uint32_t mine = 0;
+    for (int32_t i = 0; i < perThread && top - i >= 0; ++i)
+    {
+        mine += hist[top - i];
+    }
+    uint32_t const warpTotal = __reduce_add_sync(0xFFFFFFFFu, mine);
+    if (lane == 0)
+    {
+        s.warpSums[warp] = warpTotal;
+    }
+    __syncthreads();
+    uint32_t const w = s.warpSums[lane];
+    uint32_t x = w;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+        if (lane >= o)
+        {
+            x += y;
+        }
+    }
+    uint32_t const total = __shfl_sync(0xFFFFFFFFu, x, 31);
+    uint32_t const warpBefore = __shfl_sync(0xFFFFFFFFu, x - w, warp);
+    if (warpBefore < target && target <= warpBefore + warpTotal)
+    {
+        uint32_t t = mine;
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, t, o);
+            if (lane >= o)
+            {
+                t += y;
+            }
+        }
+        uint32_t acc = warpBefore + t - mine;
+        if (acc < target && target <= acc + mine)
+        {
+            for (int32_t i = 0; i < perThread && top - i >= 0; ++i)
+            {
+                uint32_t const c = hist[top - i];
+                if (acc + c >= target)
+                {
+                    s.scal[kSlotBin] = static_cast<uint32_t>(top - i);
+                    s.scal[kSlotAbove] = acc;
+                    s.scal[kSlotBinCount] = c;
+                    break;
+                }
+                acc += c;
+            }
+        }
+    }
+    __syncthreads();
+    return total;
+}
+
+struct SelectOut
+{
+    uint32_t edge;     // lower edge of the selected bin (the exact key after the last digit)
+    uint32_t above;    // predicate keys strictly above the selected bin
+    uint32_t binCount; // predicate keys in the selected bin
+    uint32_t total;    // predicate keys
+    bool all;          // take every predicate key (fewer than needed, or all fit the capacity)
+    bool full;         // exact mode: the keys at or above `edge` are exactly the `need` largest
+    int32_t passes;    // radix passes run
+};
+
+// Candidate keys and token indices held in shared memory (16-byte aligned, readable in groups of four up to the
+// next multiple of four).
+struct SmemKeys
+{
+    static constexpr bool kVector = true;
+    uint32_t const* keys;
+    uint32_t const* idx;
+
+    __device__ uint4 keys4(int32_t j) const
+    {
+        return *reinterpret_cast<uint4 const*>(keys + j);
+    }
+
+    __device__ uint4 index4(int32_t j) const
+    {
+        return *reinterpret_cast<uint4 const*>(idx + j);
+    }
+};
+
+// Merge candidates in shared memory followed by round-2 candidates read from their global Lamport slots.
+struct MergeKeys
+{
+    static constexpr bool kVector = false;
+    uint32_t const* smemKeys;
+    uint32_t const* smemIdx;
+    int32_t smemCount;
+    uint32_t const* const* slots; // round-2 slot of each segment
+    uint32_t const* segBegin;     // [numSegments + 1] candidate offsets
+    int32_t numSegments;
+
+    __device__ uint32_t const* pair(int32_t k) const
+    {
+        // Last segment whose begin is <= k.
+        int32_t lo = 0;
+        int32_t hi = numSegments - 1;
+        while (lo < hi)
+        {
+            int32_t const mid = (lo + hi + 1) / 2;
+            if (static_cast<int32_t>(segBegin[mid]) <= k)
+            {
+                lo = mid;
+            }
+            else
+            {
+                hi = mid - 1;
+            }
+        }
+        return slots[lo] + kSlotHeaderWords + 2 * (k - static_cast<int32_t>(segBegin[lo]));
+    }
+
+    __device__ uint32_t operator()(int32_t j) const
+    {
+        return j < smemCount ? smemKeys[j] : pair(j - smemCount)[0];
+    }
+
+    __device__ uint32_t index(int32_t j) const
+    {
+        return j < smemCount ? smemIdx[j] : pair(j - smemCount)[1];
+    }
+};
+
+// Debug slots of this CTA starting at `offset` (nullptr without a debug buffer).
+__device__ __forceinline__ int64_t* debugSlots(CpTopkParams const& p, int32_t offset)
+{
+    if (p.debugTimes == nullptr)
+    {
+        return nullptr;
+    }
+    return p.debugTimes + (static_cast<int64_t>(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots + offset;
+}
+
+// clock64 of thread 0 into a debug slot (no-op without a debug buffer).
+__device__ __forceinline__ void stamp(int64_t* dbg, int32_t slot)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[slot] = clock64();
+    }
+}
+
+// A value into a debug slot (thread 0; no-op without a debug buffer).
+__device__ __forceinline__ void putDebug(int64_t* dbg, int32_t slot, int64_t value)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[slot] = value;
+    }
+}
+
+// Phase timestamp for diagnosis (thread 0 of the CTA; no-op without a debug buffer).
+__device__ __forceinline__ void markPhase(CpTopkParams const& p, int32_t phase)
+{
+    if (p.debugTimes != nullptr && threadIdx.x == 0)
+    {
+        uint64_t t;
+        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+        int64_t* d = p.debugTimes + (static_cast<int64_t>(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots;
+        d[phase] = static_cast<int64_t>(t);
+        d[16 + phase] = static_cast<int64_t>(clock64());
+    }
+}
+
+// ---- Bridge outputs (topkBridgeKernel's values, written at the emission) --------------------------------------------
+
+// Row scalars of the bridge in an exchange row, in static shared memory, from warp kFirstResetWarp:
+//   - the row's request index before the bias (lane kReqLane's asynchronous copy, started right after the grid
+//     dependency wait and completed before the block barrier that follows the key loads);
+//   - the layer ids and the accepted count (asynchronous copies started right after the grid dependency wait and
+//     completed with the warp's cp.async.wait_all in resetSlots);
+//   - the row's pool row and copy row (stored by lane 0 after resetSlots);
+// the last two groups are read after the block barrier that follows the candidate poll.
+struct BridgeShared
+{
+    int64_t acc64;
+    int32_t* poolRow;
+    int32_t* copyRow;
+    int32_t acc32;
+    int32_t reqRaw;
+    int32_t layer[kMaxGroupSize];
+};
+
+// Lane of warp kFirstResetWarp that copies the request index.
+constexpr int32_t kReqLane = 30;
+
+// Request of an exchange row from its staged index.
+__device__ __forceinline__ int32_t sharedRequest(CpTopkBridgeParams const& b, BridgeShared const& bs)
+{
+    return bs.reqRaw - b.reqIdxBias;
+}
+
+// Pool base of token `tok` read from the block table (the bridge's rule: -1 for tok < 0 or a page beyond the table).
+__device__ __forceinline__ int32_t directBase(CpTopkBridgeParams const& b, int32_t req, int32_t tok)
+{
+    if (tok < 0)
+    {
+        return -1;
+    }
+    int32_t const page = tok / kPageTokens;
+    if (page >= b.maxBlocksPerReq)
+    {
+        return -1;
+    }
+    return b.blockTable[req * b.btStride0 + page * b.btStride1];
+}
+
+// Copy row of verify row `row` (-1: the row is not copied). Copy mode 0: rows [0, copyRows) to copyRowStart + row.
+// Copy mode 1: request i = row / copyNextN < copyRows copies its last accepted row i * copyNextN + clamp(acc - 1, 0,
+// copyNextN - 1) (copyNextN - 1 without accepted counts) to copyRowStart + i; that row is the only one of request i
+// that can match, so this is the bridge's selection seen from the source row.
+__device__ __forceinline__ int32_t bridgeCopyRow(CpTopkBridgeParams const& b, int32_t row, int64_t acc)
+{
+    if (b.copyDst == nullptr)
+    {
+        return -1;
+    }
+    if (b.copyMode == 0)
+    {
+        return row < b.copyRows ? b.copyRowStart + row : -1;
+    }
+    int32_t const i = row / b.copyNextN;
+    if (i >= b.copyRows)
+    {
+        return -1;
+    }
+    int64_t off = b.copyNextN - 1;
+    if (b.numAccepted != nullptr)
+    {
+        off = acc - 1;
+        off = off < 0 ? 0 : (off > b.copyNextN - 1 ? b.copyNextN - 1 : off);
+    }
+    int32_t const srcRow = static_cast<int32_t>(static_cast<int64_t>(i) * b.copyNextN + off);
+    return srcRow == row ? b.copyRowStart + i : -1;
+}
+
+// Whether verify row `row` needs an accepted count (copy mode 1 with counts, a request below copyRows).
+__device__ __forceinline__ bool needsAccepted(CpTopkBridgeParams const& b, int32_t row)
+{
+    return b.copyDst != nullptr && b.copyMode == 1 && b.numAccepted != nullptr && row / b.copyNextN < b.copyRows;
+}
+
+// Accepted count staged in shared memory (BridgeShared) in the dtype of the counts.
+__device__ __forceinline__ int64_t sharedAccepted(CpTopkBridgeParams const& b, BridgeShared const& bs)
+{
+    return b.numAcceptedIs64 ? bs.acc64 : static_cast<int64_t>(bs.acc32);
+}
+
+// Accepted count of the request of verify row `row`, read from global memory (0 when not needed).
+__device__ __forceinline__ int64_t loadAccepted(CpTopkBridgeParams const& b, int32_t row)
+{
+    if (!needsAccepted(b, row))
+    {
+        return 0;
+    }
+    int32_t const i = row / b.copyNextN;
+    return b.numAcceptedIs64 ? static_cast<int64_t const*>(b.numAccepted)[i]
+                             : static_cast<int64_t>(static_cast<int32_t const*>(b.numAccepted)[i]);
+}
+
+// Destinations of one verify row's bridge outputs: the pool row of group member 0 (member g at g * strideG further)
+// and the copy row (nullptr: not copied).
+struct BridgeRowOut
+{
+    int32_t* pool;
+    int64_t strideG;
+    int32_t* copy;
+};
+
+__device__ __forceinline__ BridgeRowOut bridgeRowOut(CpTopkBridgeParams const& b, int32_t row, int32_t copyRow)
+{
+    return BridgeRowOut{
+        b.out + row * b.outStride0, b.outStrideG, copyRow >= 0 ? b.copyDst + copyRow * b.copyStride0 : nullptr};
+}
+
+// The bridge outputs of output position `pos`: for every group member g, the bridge's remapOne(tok, layer g) =
+// base * strideFactor + (tok % kPageTokens + layer[g] * kPageTokens), or -1 for tok < 0 or base < 0, with its int32
+// wraparound arithmetic (evaluated in uint32, where the grouping of the sum does not change the bits); then the copy.
+template <typename LayerAt>
+__device__ __forceinline__ void writeBridge(
+    CpTopkBridgeParams const& b, BridgeRowOut const& o, uint32_t pos, int32_t tok, int32_t base, LayerAt const& layerAt)
+{
+    bool const valid = tok >= 0 && base >= 0;
+    uint32_t const common = static_cast<uint32_t>(base) * static_cast<uint32_t>(b.strideFactor)
+        + (static_cast<uint32_t>(tok) & static_cast<uint32_t>(kPageTokens - 1));
+    int32_t* dst = o.pool + pos;
+    for (int32_t g = 0; g < b.groupSize; ++g)
+    {
+        *dst = valid ? static_cast<int32_t>(common + static_cast<uint32_t>(layerAt(g)) * kPageTokens) : -1;
+        dst += o.strideG;
+    }
+    if (o.copy != nullptr)
+    {
+        o.copy[pos] = tok;
+    }
+}
+
+// Warp kFirstResetWarp of an exchange row: start the asynchronous copies of the row's request index, layer ids (or
+// store the single layer id) and its request's accepted count into `bs` (cp.async, generic proxy; see BridgeShared).
+__device__ __forceinline__ void fetchRowScalars(CpTopkParams const& p, BridgeShared& bs, int32_t row)
+{
+    CpTopkBridgeParams const& b = p.bridge;
+    int32_t const lane = static_cast<int32_t>(threadIdx.x) % 32;
+    if (lane < b.groupSize)
+    {
+        if (b.layerIds != nullptr)
+        {
+            uint32_t const dst = static_cast<uint32_t>(__cvta_generic_to_shared(&bs.layer[lane]));
+            asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(dst), "l"(b.layerIds + lane) : "memory");
+        }
+        else
+        {
+            bs.layer[lane] = b.layerId;
+        }
+    }
+    static_assert(kMaxGroupSize < kReqLane && kReqLane < 31, "the request and accepted-count lanes hold no layer id");
+    if (lane == kReqLane)
+    {
+        // A group of its own, completed (cp.async.wait_group 0) before the block barrier that follows the key loads.
+        uint32_t const dst = static_cast<uint32_t>(__cvta_generic_to_shared(&bs.reqRaw));
+        asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(dst), "l"(b.reqIdx + row) : "memory");
+        asm volatile("cp.async.commit_group;" ::: "memory");
+    }
+    if (lane == 31 && needsAccepted(b, row))
+    {
+        int32_t const i = row / b.copyNextN;
+        if (b.numAcceptedIs64)
+        {
+            uint32_t const dst = static_cast<uint32_t>(__cvta_generic_to_shared(&bs.acc64));
+            asm volatile("cp.async.ca.shared.global [%0], [%1], 8;" ::"r"(dst),
+                         "l"(static_cast<int64_t const*>(b.numAccepted) + i)
+                         : "memory");
+        }
+        else
+        {
+            uint32_t const dst = static_cast<uint32_t>(__cvta_generic_to_shared(&bs.acc32));
+            asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(dst),
+                         "l"(static_cast<int32_t const*>(b.numAccepted) + i)
+                         : "memory");
+        }
+    }
+}
+
+// Start the copy of candidate token `tok`'s pool base (its block-table entry, 4-byte cp.async) into `dst`, or store
+// -1 for tok < 0 or a page beyond the table (the bridge's rule). Only the calling thread reads `dst`, after its
+// cp.async.wait_all.
+__device__ __forceinline__ void fetchBase(CpTopkBridgeParams const& b, int32_t req, int32_t tok, uint32_t* dst)
+{
+    int32_t const page = tok / kPageTokens;
+    if (tok < 0 || page >= b.maxBlocksPerReq)
+    {
+        *dst = 0xFFFFFFFFu;
+        return;
+    }
+    uint32_t const addr = static_cast<uint32_t>(__cvta_generic_to_shared(dst));
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(addr),
+                 "l"(b.blockTable + req * b.btStride0 + page * b.btStride1)
+                 : "memory");
+}
+
+// Cluster-merge emission of one thread: its (at most two) candidates at pos0 / pos1 and its padding position posPad
+// (kTopK: none), ids and bridge outputs (same values as writeBridge): the id row outRow, the pool row of group member
+// g at poolRow + g * strideG (kGroup members; 0: groupSize, at most kMaxGroupSize) and the copy row (nullptr: none),
+// with the pool bases fetched at the poll and the layer ids in shared memory.
+//
+// It is a called function (not inlined) so that its instructions are at one address: warp kFirstResetWarp's lane 0
+// runs it once into shared scratch while it waits for the poll (warmEmission), so the emission does not fetch them
+// from memory on the critical path. The emission runs once per launch, at its end; when the L2 has been streamed since
+// the previous launch (as by a layer's weights in the served step), its instructions would otherwise come from DRAM.
+template <int32_t kGroup>
+__device__ __noinline__ void emitCandidates(int32_t* outRow, int32_t* poolRow, int32_t* copyRow, int64_t strideG,
+    uint32_t strideFactor, int32_t groupSize, int32_t const* layer, uint32_t pos0, uint32_t pos1, uint32_t posPad,
+    uint2 tokens, uint2 bases)
+{
+    uint32_t const want = static_cast<uint32_t>(kTopK);
+    auto emit = [&](uint32_t pos, int32_t tok, int32_t base)
+    {
+        outRow[pos] = tok;
+        bool const valid = tok >= 0 && base >= 0;
+        uint32_t const common = static_cast<uint32_t>(base) * strideFactor
+            + (static_cast<uint32_t>(tok) & static_cast<uint32_t>(kPageTokens - 1));
+        int32_t* const dst = poolRow + pos;
+#pragma unroll
+        for (int32_t g = 0; g < (kGroup > 0 ? kGroup : kMaxGroupSize); ++g)
+        {
+            if (kGroup > 0 || g < groupSize)
+            {
+                dst[g * strideG]
+                    = valid ? static_cast<int32_t>(common + static_cast<uint32_t>(layer[g]) * kPageTokens) : -1;
+            }
+        }
+        if (copyRow != nullptr)
+        {
+            copyRow[pos] = tok;
+        }
+    };
+    if (pos0 < want)
+    {
+        emit(pos0, static_cast<int32_t>(tokens.x), static_cast<int32_t>(bases.x));
+    }
+    if (pos1 < want)
+    {
+        emit(pos1, static_cast<int32_t>(tokens.y), static_cast<int32_t>(bases.y));
+    }
+    if (posPad < want)
+    {
+        emit(posPad, -1, -1);
+    }
+}
+
+// emitCandidates for a run-time group size; the served sizes (per-layer remaps and groups of four) are unrolled.
+__device__ __forceinline__ void emitCandidatesFor(int32_t groupSize, int32_t* outRow, int32_t* poolRow,
+    int32_t* copyRow, int64_t strideG, uint32_t strideFactor, int32_t const* layer, uint32_t pos0, uint32_t pos1,
+    uint32_t posPad, uint2 tokens, uint2 bases)
+{
+    switch (groupSize)
+    {
+    case 1:
+        emitCandidates<1>(
+            outRow, poolRow, copyRow, strideG, strideFactor, groupSize, layer, pos0, pos1, posPad, tokens, bases);
+        break;
+    case 4:
+        emitCandidates<4>(
+            outRow, poolRow, copyRow, strideG, strideFactor, groupSize, layer, pos0, pos1, posPad, tokens, bases);
+        break;
+    default:
+        emitCandidates<0>(
+            outRow, poolRow, copyRow, strideG, strideFactor, groupSize, layer, pos0, pos1, posPad, tokens, bases);
+        break;
+    }
+}
+
+//! Words of the shared scratch of warmEmission: ids [0, 4), pool rows [4, 36) (stride 4), copy row [36, 40).
+constexpr int32_t kWarmScratchWords = 40;
+
+// One thread runs the emission of this launch's group size once into `scratch` (three positions, every store taken
+// once, no global access), so that its instructions are cached on the SM before the cluster emission runs.
+__device__ __forceinline__ void warmEmission(int32_t groupSize, int32_t* scratch, int32_t const* layer)
+{
+    emitCandidatesFor(groupSize, scratch, scratch + 4, scratch + 36, 4, 0u, layer, 0u, 1u, 2u, make_uint2(0u, 0u),
+        make_uint2(~0u, ~0u));
+}
+
+// Histogram update of one key of a radix pass.
+__device__ __forceinline__ void histogramKey(uint32_t* hist, bool inRange, uint32_t k, uint32_t lo, uint64_t hi,
+    uint32_t base, uint32_t prefMask, uint32_t prefVal, int32_t shift, uint32_t binMask)
+{
+    uint32_t const d = k - base;
+    if (inRange && k >= lo && static_cast<uint64_t>(k) < hi && (d & prefMask) == prefVal)
+    {
+        atomicAdd(&hist[(d >> shift) & binMask], 1u);
+    }
+}
+
+// Block radix select of the need-th largest of the n keys at or above lo, of which `aboveHi` (< need) lie at or
+// above hi and are counted as taken without being histogrammed. Every key in [lo, hi) is at least `base`, and
+// key - base < 2^topBits; digits of up to kDigitBits bits of key - base are refined from the top. With cap > 0 the
+// select stops at the first bin edge whose keys at or above it number at most `cap`. With cap = 0 it runs to the
+// exact key, or stops early when the selected bin is wholly taken (`full`). Alternate passes use alternate
+// histograms, so the next one is cleared while the current one is searched. Counts in the result include aboveHi.
+template <typename KeyAt>
+__device__ SelectOut blockSelect(SmemLayout const& s, KeyAt const& keyAt, int32_t n, uint32_t lo, uint64_t hi,
+    uint32_t need, uint32_t cap, uint32_t base, int32_t topBits, int64_t* dbg = nullptr, uint32_t aboveHi = 0u)
+{
+    if (dbg != nullptr && threadIdx.x == 0)
+    {
+        dbg[0] = n;
+    }
+    SelectOut out{lo, 0u, 0u, 0u, false, false, 0};
+    uint32_t prefMask = 0;
+    uint32_t prefVal = 0;
+    uint32_t above = aboveHi;
+    int32_t remaining = max(topBits, 1);
+    uint32_t* hist = s.hist;
+    uint32_t* histNext = s.hist + kBins;
+    for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads)
+    {
+        hist[b] = 0;
+    }
+    __syncthreads();
+#pragma unroll 1
+    for (int32_t pass = 0; remaining > 0; ++pass)
+    {
+        int32_t const width = min(kDigitBits, remaining);
+        int32_t const shift = remaining - width;
+        int32_t const nbins = 1 << width;
+        uint32_t const binMask = static_cast<uint32_t>(nbins - 1);
+        out.passes = pass + 1;
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[1 + 3 * pass] = clock64();
+        }
+        if constexpr (KeyAt::kVector)
+        {
+            for (int32_t j0 = 4 * static_cast<int32_t>(threadIdx.x); j0 < n; j0 += 4 * kThreads)
+            {
+                uint4 const q = keyAt.keys4(j0);
+                histogramKey(hist, true, q.x, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 1 < n, q.y, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 2 < n, q.z, lo, hi, base, prefMask, prefVal, shift, binMask);
+                histogramKey(hist, j0 + 3 < n, q.w, lo, hi, base, prefMask, prefVal, shift, binMask);
+            }
+        }
+        else
+        {
+            for (int32_t j = threadIdx.x; j < n; j += kThreads)
+            {
+                histogramKey(hist, true, keyAt(j), lo, hi, base, prefMask, prefVal, shift, binMask);
+            }
+        }
+        __syncthreads();
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[2 + 3 * pass] = clock64();
+        }
+        if (shift > 0)
+        {
+            for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads)
+            {
+                histNext[b] = 0;
+            }
+        }
+        // Find the bin holding the (need - above)-th largest key, counting from the top bin.
+        uint32_t const target = need - above;
+        uint32_t const total = findBin(s, hist, nbins, target);
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[3 + 3 * pass] = clock64();
+        }
+        if (pass == 0)
+        {
+            out.total = total + aboveHi;
+            if (total + aboveHi <= cap || target > total)
+            {
+                out.all = true;
+                out.above = total + aboveHi;
+                return out;
+            }
+        }
+        // The scalars are next written after two more barriers (the next histogram and scan).
+        uint32_t const bin = s.scal[kSlotBin];
+        above += s.scal[kSlotAbove];
+        out.binCount = s.scal[kSlotBinCount];
+        prefVal |= bin << shift;
+        prefMask |= binMask << shift;
+        out.edge = base + prefVal;
+        out.above = above;
+        remaining = shift;
+        uint32_t* const t = hist;
+        hist = histNext;
+        histNext = t;
+        if (cap > 0 && above + out.binCount <= cap)
+        {
+            return out;
+        }
+        if (cap == 0 && remaining > 0 && above + out.binCount == need)
+        {
+            out.full = true;
+            return out;
+        }
+    }
+    return out;
+}
+
+// Send one source's selected keys (keys in [lo, hi) of the chunk) to its slot; returns the count sent.
+// Takes every predicate key at or above the edge when they fit `cap`. Otherwise (more than `cap` keys tie at the
+// exact edge) it takes every key above the edge, plus, with kPartialTies, the first keys equal to it in source
+// order up to `cap` (any subset of equal keys is a valid tie-aware choice). Round 1 sends no partial ties, so the
+// keys a second round requests from a source (at most the edge) were never sent before.
+// The header carries the count and an exclusive upper bound of the keys left unsent (0 when none); in round 1
+// (!kPartialTies) also the key statistics of the keys sent (see kSlotHeaderWords).
+// With `stage` (round 1, at most kSourceCapacity keys), the pairs are gathered in shared memory first and written
+// with coalesced 16-byte stores; otherwise each pair is stored where it is found.
+template <bool kPartialTies>
+__device__ uint32_t sendSource(SmemLayout const& s, SelectOut const& sel, int32_t n, uint32_t lo, uint64_t hi,
+    uint32_t cap, uint32_t* mcSlot, int32_t src, int32_t rank, int32_t world, uint2* stage, int64_t* dbg = nullptr)
+{
+    bool const strict = !sel.all && sel.above + sel.binCount > cap;
+    uint32_t const threshold = sel.all ? lo : sel.edge;
+    // Keys equal to the edge that fit after the keys above it (sel.above < need <= cap when strict).
+    uint32_t const tieBudget = kPartialTies && strict ? cap - sel.above : 0u;
+    uint32_t const sent = sel.all ? sel.total : (strict ? sel.above + tieBudget : sel.above + sel.binCount);
+    uint32_t const bound = sel.all ? lo : (strict ? sel.edge + 1u : sel.edge);
+    stamp(dbg, 42);
+    // One page block (kBlockTokens keys) per thread; positions from a block scan, so the slot holds the keys in
+    // source order: first those taken outright, then (strict only) the admitted keys equal to the edge.
+    int32_t const groups = (n + kBlockTokens - 1) / kBlockTokens;
+    uint32_t takenMain = 0;
+    uint32_t takenTies = 0;
+    uint32_t const tieBase = strict ? sel.above : 0u;
+    for (int32_t g0 = 0; g0 < groups; g0 += kThreads)
+    {
+        int32_t const g = g0 + static_cast<int32_t>(threadIdx.x);
+        uint32_t k[kBlockTokens];
+        uint32_t mask = 0;
+        uint32_t tieMask = 0;
+        if (g < groups)
+        {
+            uint4 const a = *reinterpret_cast<uint4 const*>(s.chunk + kBlockTokens * g);
+            uint4 const b = *reinterpret_cast<uint4 const*>(s.chunk + kBlockTokens * g + 4);
+            k[0] = a.x;
+            k[1] = a.y;
+            k[2] = a.z;
+            k[3] = a.w;
+            k[4] = b.x;
+            k[5] = b.y;
+            k[6] = b.z;
+            k[7] = b.w;
+#pragma unroll
+            for (int32_t i = 0; i < kBlockTokens; ++i)
+            {
+                bool const in = kBlockTokens * g + i < n && k[i] >= lo && static_cast<uint64_t>(k[i]) < hi;
+                bool const take = in && (strict ? k[i] > threshold : k[i] >= threshold);
+                bool const tie = kPartialTies && in && strict && k[i] == threshold;
+                mask |= (take ? 1u : 0u) << i;
+                tieMask |= (tie ? 1u : 0u) << i;
+            }
+        }
+        // At most kThreads * kBlockTokens keys per pass: both counts fit 16 bits.
+        uint32_t const c = static_cast<uint32_t>(__popc(mask)) | (static_cast<uint32_t>(__popc(tieMask)) << 16);
+        uint32_t chunkTaken;
+        uint32_t const before = blockExclusiveScan(c, s.warpSums, chunkTaken);
+        uint32_t pos = takenMain + (before & 0xFFFFu);
+        uint32_t tiePos = takenTies + (before >> 16);
+        takenMain += chunkTaken & 0xFFFFu;
+        takenTies += chunkTaken >> 16;
+        if ((mask | tieMask) != 0u)
+        {
+            int32_t const colBase = g * kPageTokens + ((src + g) % kSourcesPerRank) * kBlockTokens;
+#pragma unroll
+            for (int32_t i = 0; i < kBlockTokens; ++i)
+            {
+                uint32_t slot = kEmptyWord;
+                if ((mask >> i) & 1u)
+                {
+                    slot = pos++;
+                }
+                else if constexpr (kPartialTies)
+                {
+                    if ((tieMask >> i) & 1u)
+                    {
+                        slot = tiePos < tieBudget ? tieBase + tiePos : kEmptyWord;
+                        ++tiePos;
+                    }
+                }
+                if (slot != kEmptyWord)
+                {
+                    uint32_t const token = localToGlobal(colBase + i, rank, world);
+                    if (stage != nullptr)
+                    {
+                        stage[slot] = make_uint2(k[i], token);
+                    }
+                    else
+                    {
+                        store2(mcSlot + kSlotHeaderWords + 2 * static_cast<int64_t>(slot), k[i], token);
+                    }
+                }
+            }
+        }
+        // The staged pairs are complete, and the scan's warp sums may be rewritten.
+        __syncthreads();
+    }
+    stamp(dbg, 43);
+    if (stage != nullptr)
+    {
+        // Two pairs per 16-byte store; an odd count leaves the last half of the last vector empty.
+        int32_t const vectors = static_cast<int32_t>(sent + 1u) / 2;
+        for (int32_t v = threadIdx.x; v < vectors; v += kThreads)
+        {
+            uint2 const a = stage[2 * v];
+            uint2 const b
+                = 2 * v + 1 < static_cast<int32_t>(sent) ? stage[2 * v + 1] : make_uint2(kEmptyWord, kEmptyWord);
+            store4(mcSlot + kSlotHeaderWords + 4 * static_cast<int64_t>(v), make_uint4(a.x, a.y, b.x, b.y));
+        }
+    }
+    // Receivers check every candidate word, so the header needs no ordering against the candidates.
+    if constexpr (kPartialTies)
+    {
+        if (threadIdx.x == kThreads - 1)
+        {
+            store2(mcSlot, sent, bound);
+        }
+    }
+    else
+    {
+        // The statistics of the keys sent, from the staged pairs (round 1 always stages), by the last warp, which
+        // stores no vector (at most kSourceCapacity / 2), while the first warps store the vectors.
+        static_assert(kSourceCapacity / 2 <= kThreads - 32, "the last warp stores no candidate vector");
+        if (threadIdx.x >= kThreads - 32)
+        {
+            int32_t const lane = threadIdx.x % 32;
+            KeyStats t;
+#pragma unroll
+            for (int32_t j = 0; j < kSourceCapacity / 32; ++j)
+            {
+                int32_t const i = lane + 32 * j;
+                if (i < static_cast<int32_t>(sent))
+                {
+                    t.add(stage[i].x);
+                }
+            }
+            uint32_t const maxKey = __reduce_max_sync(0xFFFFFFFFu, t.maxKey);
+            uint32_t const minKey = __reduce_min_sync(0xFFFFFFFFu, t.minKey);
+            uint32_t const positives = __reduce_add_sync(0xFFFFFFFFu, t.positives);
+            uint32_t const maxNeg = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg);
+            uint32_t const minPos = __reduce_min_sync(0xFFFFFFFFu, t.minPos);
+            if (lane == 31)
+            {
+                store4(mcSlot, make_uint4(sent, bound, maxKey, minKey + 1u));
+                store4(mcSlot + 4, make_uint4(positives, maxNeg, minPos + 1u, 0u));
+            }
+        }
+    }
+    stamp(dbg, 44);
+    return sent;
+}
+
+// Load one source's keys (one block of kBlockTokens tokens of every local page) from the local logits into the
+// chunk; returns the number of keys. Consecutive threads read consecutive keys (whole 32-byte blocks per eight
+// threads), and each thread issues kLoadBatch loads before it uses any. With `withStats`, the per-warp key
+// statistics are stored for readStats after the next block barrier.
+__device__ int32_t loadSource(
+    SmemLayout const& s, float const* rowLogits, int32_t localLen, int32_t src, bool withStats)
+{
+    constexpr int32_t kLoadBatch = 8;
+    int32_t const n = sourceLength(localLen, src);
+    KeyStats st;
+    for (int32_t j0 = threadIdx.x; j0 < n; j0 += kLoadBatch * kThreads)
+    {
+        float v[kLoadBatch];
+#pragma unroll
+        for (int32_t i = 0; i < kLoadBatch; ++i)
+        {
+            int32_t const j = j0 + i * kThreads;
+            int32_t const page = j / kBlockTokens;
+            int32_t const col = page * kPageTokens + ((src + page) % kSourcesPerRank) * kBlockTokens + j % kBlockTokens;
+            v[i] = j < n ? rowLogits[col] : 0.F;
+        }
+#pragma unroll
+        for (int32_t i = 0; i < kLoadBatch; ++i)
+        {
+            int32_t const j = j0 + i * kThreads;
+            if (j < n)
+            {
+                uint32_t const k = orderedKey(v[i]);
+                s.chunk[j] = k;
+                st.add(k);
+            }
+        }
+    }
+    if (withStats)
+    {
+        storeWarpStats(st, s.stats);
+    }
+    return n;
+}
+
+// Whether source q withheld keys that may lie above the merged k-th key.
+__device__ __forceinline__ bool sourceOverflows(SmemLayout const& s, int32_t q, uint32_t kth)
+{
+    uint32_t const cnt = s.scal[kSlotSrcCount + q];
+    uint64_t const bound = s.scal[kSlotSrcBound + q];
+    return cnt < s.scal[kSlotSrcLen + q] && bound > uint64_t{kth} + 1u;
+}
+
+// Emit the top-kTopK of the merge candidates. Keys above `kth` are all taken; ties at `kth` are taken in candidate
+// order when only some of them fit (deterministic on every rank), otherwise all of them. With `takeAll`, every
+// candidate is taken and the row is padded with -1. The row is assembled in shared memory (the first histogram)
+// and written with coalesced stores, together with the bridge outputs of verify row `verifyRow` (pool bases read from
+// the block table; the row scalars from `bs`).
+template <typename KeyAt>
+__device__ bool emitSelection(SmemLayout const& s, KeyAt const& keys, int32_t total, uint32_t kth, bool takeAll,
+    int32_t* outRow, int32_t overflowSources, CpTopkParams const& p, BridgeShared const& bs, int32_t verifyRow,
+    int64_t* dbg = nullptr)
+{
+    uint32_t const want = static_cast<uint32_t>(kTopK);
+    // Positions from one block scan over contiguous candidate ranges (no contended counter); the order is the
+    // candidate order, identical on every rank. Shared-memory candidates are read four at a time.
+    int32_t per = (total + kThreads - 1) / kThreads;
+    if constexpr (KeyAt::kVector)
+    {
+        per = (per + 3) / 4 * 4;
+    }
+    int32_t const j0 = min(total, static_cast<int32_t>(threadIdx.x) * per);
+    int32_t const j1 = min(total, j0 + per);
+    uint32_t cGt = 0;
+    uint32_t cEq = 0;
+    if constexpr (KeyAt::kVector)
+    {
+        for (int32_t j = j0; j < j1; j += 4)
+        {
+            uint4 const q = keys.keys4(j);
+            uint32_t const kk[4] = {q.x, q.y, q.z, q.w};
+#pragma unroll
+            for (int32_t i = 0; i < 4; ++i)
+            {
+                bool const in = j + i < j1;
+                cGt += (in && (takeAll || kk[i] > kth)) ? 1u : 0u;
+                cEq += (in && !takeAll && kk[i] == kth) ? 1u : 0u;
+            }
+        }
+    }
+    else
+    {
+        for (int32_t j = j0; j < j1; ++j)
+        {
+            uint32_t const k = keys(j);
+            cGt += (takeAll || k > kth) ? 1u : 0u;
+            cEq += (!takeAll && k == kth) ? 1u : 0u;
+        }
+    }
+    // The first `overflowSources` threads flag a source that withheld keys above the threshold (the flag slot is
+    // zero on entry); with a flag set, nothing is written and the caller runs the second round.
+    if (static_cast<int32_t>(threadIdx.x) < overflowSources && sourceOverflows(s, threadIdx.x, kth))
+    {
+        atomicOr(&s.scal[kSlotCounter2], 1u);
+    }
+    uint32_t allGt;
+    uint32_t allEq;
+    uint32_t posGt;
+    uint32_t posEq;
+    if (total < 0x10000)
+    {
+        // Both counts fit 16 bits: one scan of the packed pair.
+        uint32_t packedTotal;
+        uint32_t const packed = cGt | (cEq << 16);
+        uint32_t const before = blockExclusiveScan(packed, s.warpSums, packedTotal);
+        allGt = packedTotal & 0xFFFFu;
+        allEq = packedTotal >> 16;
+        posGt = before & 0xFFFFu;
+        posEq = allGt + (before >> 16);
+    }
+    else
+    {
+        posGt = blockInclusiveScan(cGt, s.warpSums, allGt) - cGt;
+        __syncthreads();
+        posEq = blockInclusiveScan(cEq, s.warpSums, allEq) - cEq;
+        posEq += allGt;
+    }
+    stamp(dbg, 59);
+    if (overflowSources > 0 && s.scal[kSlotCounter2] != 0u)
+    {
+        return true;
+    }
+    uint32_t* row = s.hist;
+    auto place = [&](uint32_t k, uint32_t idx)
+    {
+        if (takeAll || k > kth)
+        {
+            if (posGt < want)
+            {
+                row[posGt] = idx;
+            }
+            ++posGt;
+        }
+        else if (k == kth)
+        {
+            if (posEq < want)
+            {
+                row[posEq] = idx;
+            }
+            ++posEq;
+        }
+    };
+    if constexpr (KeyAt::kVector)
+    {
+        for (int32_t j = j0; j < j1; j += 4)
+        {
+            uint4 const q = keys.keys4(j);
+            uint4 const x = keys.index4(j);
+            uint32_t const kk[4] = {q.x, q.y, q.z, q.w};
+            uint32_t const xx[4] = {x.x, x.y, x.z, x.w};
+#pragma unroll
+            for (int32_t i = 0; i < 4; ++i)
+            {
+                if (j + i < j1)
+                {
+                    place(kk[i], xx[i]);
+                }
+            }
+        }
+    }
+    else
+    {
+        for (int32_t j = j0; j < j1; ++j)
+        {
+            place(keys(j), keys.index(j));
+        }
+    }
+    __syncthreads();
+    stamp(dbg, 60);
+    uint32_t const selected = min(allGt + allEq, want);
+    BridgeRowOut const bridgeOut{bs.poolRow, p.bridge.outStrideG, bs.copyRow};
+    int32_t const req = sharedRequest(p.bridge, bs);
+    for (uint32_t j = threadIdx.x; j < want; j += kThreads)
+    {
+        int32_t const tok = j < selected ? static_cast<int32_t>(row[j]) : -1;
+        outRow[j] = tok;
+        writeBridge(
+            p.bridge, bridgeOut, j, tok, directBase(p.bridge, req, tok), [&](int32_t g) { return bs.layer[g]; });
+    }
+    stamp(dbg, 61);
+    return false;
+}
+
+// Threshold of an exact (cap = 0) merge select: keys above `kth` are taken, ties at `kth` while they fit.
+__device__ __forceinline__ uint32_t mergeThreshold(SelectOut const& m, bool& takeAll)
+{
+    takeAll = m.all || (m.full && m.edge == 0u);
+    if (takeAll)
+    {
+        return 0u;
+    }
+    // A wholly taken bin: the keys at or above its edge are exactly kTopK.
+    return m.full ? m.edge - 1u : m.edge;
+}
+
+// Reset of the previous exchange launch's buffer on this rank's copy. Its 2 * kMaxSources * kMaxRows slots (round 1,
+// then round 2) are spread over the kResetWarps reset warps (the warps from kFirstResetWarp on) of every CTA of the
+// launch, at most ten consecutive slots per warp, one per lane: each lane fetches its slot's header once the epoch is
+// known (fetchResetHeaders, after the grid dependency wait and the first block barrier; an asynchronous copy into
+// shared memory that completes during the local select and the send) and the warp resets the written slots while the
+// poll warps wait for the candidates (resetSlots). A slot was written iff its header was; the header's count bounds
+// the candidates to reset.
+constexpr int32_t kResetSlots = 2 * kMaxSources * kMaxRows;
+
+// This CTA's index among the launch's clearers, and their number (from the grid; recomputed where used).
+__device__ __forceinline__ int32_t resetClearer()
+{
+    return static_cast<int32_t>(blockIdx.y) * kSourcesPerRank + static_cast<int32_t>(blockIdx.x);
+}
+
+__device__ __forceinline__ int32_t resetClearers()
+{
+    return static_cast<int32_t>(gridDim.y) * kSourcesPerRank;
+}
+
+__device__ __forceinline__ int32_t resetSlotsPerWarp(int32_t numClearers)
+{
+    int32_t const warps = numClearers * kResetWarps;
+    return (kResetSlots + warps - 1) / warps;
+}
+
+// First slot of this (reset) warp.
+__device__ __forceinline__ int32_t resetFirstSlot(int32_t clearer, int32_t numClearers)
+{
+    int32_t const warp = static_cast<int32_t>(threadIdx.x) / 32 - kFirstResetWarp;
+    return (clearer * kResetWarps + warp) * resetSlotsPerWarp(numClearers);
+}
+
+__device__ __forceinline__ int64_t resetSlotOffset(int32_t buf, int32_t slot)
+{
+    int32_t const round1Slots = kMaxSources * kMaxRows;
+    bool const second = slot >= round1Slots;
+    int32_t const qr = second ? slot - round1Slots : slot;
+    return second ? round2Offset(buf, qr / kMaxRows, qr % kMaxRows) : round1Offset(buf, qr / kMaxRows, qr % kMaxRows);
+}
+
+__device__ __forceinline__ uint32_t* resetSlotBase(CpTopkParams const& p, int32_t buf, int32_t slot)
+{
+    return p.ucWorkspace + resetSlotOffset(buf, slot);
+}
+
+//! Most slots one reset warp resets (one row: the fewest reset warps).
+constexpr int32_t kMaxResetSlotsPerWarp
+    = (kResetSlots + kSourcesPerRank * kResetWarps - 1) / (kSourcesPerRank * kResetWarps);
+static_assert(kMaxResetSlotsPerWarp <= 32, "one slot header per lane");
+
+// Shared-memory home of the reset warps' slot headers (16 bytes per slot: count, bound and two further header words),
+// in the merge candidate area beyond the round-1 staging of the send; it is free from the first block barrier until
+// the reset.
+constexpr int32_t kResetHeaderOffset = 2 * kSourceCapacity;
+static_assert(kResetHeaderOffset + 4 * kMaxResetSlotsPerWarp * kResetWarps <= kSourcesPerRank * kSourceCapacity,
+    "the reset headers fit the smallest merge candidate area");
+
+__device__ __forceinline__ uint4* resetHeaderHome(SmemLayout const& s)
+{
+    int32_t const warp = static_cast<int32_t>(threadIdx.x) / 32 - kFirstResetWarp;
+    return reinterpret_cast<uint4*>(s.mkey + kResetHeaderOffset) + warp * kMaxResetSlotsPerWarp;
+}
+
+// Word offset, within one Lamport buffer, of the slot whose header this lane fetches for the reset; -1 for the lanes
+// and warps without one. Arithmetic on the grid only (no memory access).
+__device__ __forceinline__ int32_t resetHeaderOffset()
+{
+    int32_t const numClearers = resetClearers();
+    int32_t const lane = threadIdx.x % 32;
+    if (static_cast<int32_t>(threadIdx.x) / 32 < kFirstResetWarp || lane >= resetSlotsPerWarp(numClearers))
+    {
+        return -1;
+    }
+    int32_t const slot = resetFirstSlot(resetClearer(), numClearers) + lane;
+    if (slot >= kResetSlots)
+    {
+        return -1;
+    }
+    static_assert(kBufferWords <= int64_t{1} << 31, "slot offsets within a buffer fit int32");
+    return static_cast<int32_t>(resetSlotOffset(0, slot));
+}
+
+// Starts the copy of this lane's reset slot header of buffer `buf` (offset from resetHeaderOffset) into the warp's
+// home with an asynchronous copy (cp.async, generic proxy, completed in resetSlots): a pending copy holds no register
+// and no block barrier waits for it. Lanes without a slot do nothing.
+__device__ __forceinline__ void fetchResetHeader(
+    CpTopkParams const& p, SmemLayout const& s, int32_t buf, int32_t offset)
+{
+    if (offset >= 0)
+    {
+        uint32_t const dst = static_cast<uint32_t>(
+            __cvta_generic_to_shared(resetHeaderHome(s) + static_cast<int32_t>(threadIdx.x % 32)));
+        uint32_t const* const src = p.ucWorkspace + int64_t{buf} * kBufferWords + offset;
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(src) : "memory");
+    }
+}
+
+// Reset this reset warp's slots of buffer `buf` (headers fetched by fetchResetHeaders); other warps return.
+__device__ void resetSlots(CpTopkParams const& p, SmemLayout const& s, int32_t buf)
+{
+    if (static_cast<int32_t>(threadIdx.x) / 32 < kFirstResetWarp)
+    {
+        return;
+    }
+    int32_t const clearer = resetClearer();
+    int32_t const numClearers = resetClearers();
+    // Each lane completes its own header copy; the warp barrier then makes every lane's header visible to the warp.
+    asm volatile("cp.async.wait_all;" ::: "memory");
+    __syncwarp();
+    int32_t const per = resetSlotsPerWarp(numClearers);
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const first = resetFirstSlot(clearer, numClearers);
+    uint4 const* const headers = resetHeaderHome(s);
+    uint4 const empty = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord);
+    for (int32_t i = 0; i < per && first + i < kResetSlots; ++i)
+    {
+        uint4 const h = headers[i];
+        if (h.x == kEmptyWord && h.y == kEmptyWord)
+        {
+            continue;
+        }
+        int32_t const slot = first + i;
+        uint32_t* base = resetSlotBase(p, buf, slot);
+        uint32_t const capacity = slot >= kMaxSources * kMaxRows ? kRound2Capacity : kSourceCapacity;
+        uint32_t const count = h.x == kEmptyWord ? capacity : min(h.x, capacity);
+        int32_t const vectors = static_cast<int32_t>(count + 1u) / 2;
+        for (int32_t v = lane; v < vectors; v += 32)
+        {
+            store4(base + kSlotHeaderWords + 4 * static_cast<int64_t>(v), empty);
+        }
+        if (lane < 2)
+        {
+            store4(base + 4 * lane, empty);
+        }
+    }
+}
+
+// Merge of one row on this rank (one CTA): gathers every source's round-1 candidates, selects the exact
+// top-kTopK, runs the second round when a source withheld keys above the threshold, and writes the row and its bridge
+// outputs.
+__device__ void mergeRow(CpTopkParams const& p, SmemLayout const& s, uint32_t const** segSlots, uint32_t* segBegin,
+    int32_t buf, int32_t row, int32_t rowLen, int32_t localLen, float const* rowLogits, int32_t* outRow,
+    WaitCtx const& wait, BridgeShared const& bs)
+{
+    int32_t const world = p.worldSize;
+    int32_t const rank = p.rank;
+    int32_t const sources = numSources(world);
+    uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+    __syncthreads();
+    if (threadIdx.x == 0)
+    {
+        s.scal[kSlotCounter2] = 0u;
+    }
+    if (static_cast<int32_t>(threadIdx.x) < sources)
+    {
+        // Keys each source holds, computed while its header is in flight.
+        int32_t const q = threadIdx.x;
+        s.scal[kSlotSrcLen + q] = static_cast<uint32_t>(
+            sourceLength(min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank));
+        uint2 const h = pollPair(p.ucWorkspace + round1Offset(buf, threadIdx.x, row), wait, kWaitRound1Header);
+        s.scal[kSlotSrcCount + threadIdx.x] = min(h.x, cap);
+        s.scal[kSlotSrcBound + threadIdx.x] = h.y;
+    }
+    __syncthreads();
+    if (threadIdx.x < 32)
+    {
+        // Exclusive prefix of the source counts (at most 64 sources: two per lane).
+        uint32_t const a = static_cast<int32_t>(threadIdx.x) < sources ? s.scal[kSlotSrcCount + threadIdx.x] : 0u;
+        uint32_t const b
+            = static_cast<int32_t>(threadIdx.x) + 32 < sources ? s.scal[kSlotSrcCount + threadIdx.x + 32] : 0u;
+        uint32_t x = a;
+        uint32_t z = b;
+#pragma unroll
+        for (int32_t o = 1; o < 32; o <<= 1)
+        {
+            uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o);
+            uint32_t const w = __shfl_up_sync(0xFFFFFFFFu, z, o);
+            if (static_cast<int32_t>(threadIdx.x) >= o)
+            {
+                x += y;
+                z += w;
+            }
+        }
+        uint32_t const totalA = __shfl_sync(0xFFFFFFFFu, x, 31);
+        s.scal[kSlotSrcOffset + threadIdx.x] = x - a;
+        s.scal[kSlotSrcOffset + 32 + threadIdx.x] = totalA + z - b;
+        if (threadIdx.x == 31)
+        {
+            s.scal[kSlotSrcOffset + kMaxSources] = totalA + z;
+        }
+    }
+    __syncthreads();
+    int32_t const total = static_cast<int32_t>(s.scal[kSlotSrcOffset + kMaxSources]);
+    markPhase(p, 4);
+
+    // Candidates: padded (source, vector) positions; issue a batch of independent 16-byte loads, then re-poll the
+    // ones not yet written. The key statistics are tracked for the merge select.
+    int32_t const vectorsPerSource = static_cast<int32_t>(cap) / 2;
+    int32_t const padded = sources * vectorsPerSource;
+    KeyStats st;
+#pragma unroll 1
+    for (int32_t v0 = 0; v0 < padded; v0 += kThreads * kPollBatch)
+    {
+        uint4 v[kPollBatch];
+#pragma unroll
+        for (int32_t i = 0; i < kPollBatch; ++i)
+        {
+            int32_t const jv = v0 + i * kThreads + static_cast<int32_t>(threadIdx.x);
+            int32_t const q = jv / vectorsPerSource;
+            int32_t const kv = jv % vectorsPerSource;
+            v[i] = make_uint4(0u, 0u, 0u, 0u);
+            if (jv < padded && 2 * kv < static_cast<int32_t>(s.scal[kSlotSrcCount + q]))
+            {
+                v[i] = loadVolatile4(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv);
+            }
+        }
+#pragma unroll
+        for (int32_t i = 0; i < kPollBatch; ++i)
+        {
+            int32_t const jv = v0 + i * kThreads + static_cast<int32_t>(threadIdx.x);
+            int32_t const q = jv / vectorsPerSource;
+            int32_t const kv = jv % vectorsPerSource;
+            if (jv < padded && 2 * kv < static_cast<int32_t>(s.scal[kSlotSrcCount + q]))
+            {
+                bool const second = 2 * kv + 1 < static_cast<int32_t>(s.scal[kSlotSrcCount + q]);
+                if (!vectorReady(v[i], second))
+                {
+                    v[i] = pollVector(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv, second,
+                        wait, kWaitRound1Candidates);
+                }
+                int32_t const j = static_cast<int32_t>(s.scal[kSlotSrcOffset + q]) + 2 * kv;
+                s.mkey[j] = v[i].x;
+                s.midx[j] = v[i].y;
+                st.add(v[i].x);
+                if (second)
+                {
+                    s.mkey[j + 1] = v[i].z;
+                    s.midx[j + 1] = v[i].w;
+                    st.add(v[i].z);
+                }
+            }
+        }
+    }
+    stamp(debugSlots(p, 0), 45);
+    storeWarpStats(st, s.stats);
+    __syncthreads();
+    markPhase(p, 5);
+
+    // Digits of the merge select span the candidates of the sign the cut lies in (see kPositiveKey).
+    SignRange const r = total > 0 ? signRange(readStats(s.stats), static_cast<uint32_t>(kTopK))
+                                  : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false};
+    SmemKeys const merged{s.mkey, s.midx};
+    SelectOut const m = blockSelect(
+        s, merged, total, r.lo, r.hi, static_cast<uint32_t>(kTopK), 0u, r.lo, r.bits, debugSlots(p, 48), r.aboveHi);
+    putDebug(debugSlots(p, 0), 62, total);
+    putDebug(debugSlots(p, 0), 63, m.passes + (r.negative ? 10 : 0));
+    bool takeAll = false;
+    uint32_t const kth = mergeThreshold(m, takeAll);
+    markPhase(p, 6);
+    stamp(debugSlots(p, 0), 58);
+    if (!emitSelection(s, merged, total, kth, takeAll, outRow, sources, p, bs, row, debugSlots(p, 0)))
+    {
+        return;
+    }
+
+    // ---- Round 2 (rare; every rank takes it): withheld keys in [kth, bound) of overflowing sources ----------
+    __syncthreads();
+    // 1. This rank's overflowing sources send their top-kRound2Capacity withheld keys.
+    for (int32_t c = 0; c < kSourcesPerRank; ++c)
+    {
+        int32_t const q = rank * kSourcesPerRank + c;
+        if (!sourceOverflows(s, q, kth))
+        {
+            continue;
+        }
+        uint64_t const b = s.scal[kSlotSrcBound + q];
+        int32_t const nq = loadSource(s, rowLogits, localLen, c, false);
+        __syncthreads();
+        SelectOut const sel2 = blockSelect(
+            s, SmemKeys{s.chunk, nullptr}, nq, kth, b, static_cast<uint32_t>(kRound2Capacity), kRound2Capacity, 0u, 32);
+        (void) sendSource<true>(
+            s, sel2, nq, kth, b, kRound2Capacity, p.mcWorkspace + round2Offset(buf, q, row), c, rank, world, nullptr);
+        __syncthreads();
+    }
+    // 2. Gather the round-2 segments of every overflowing source (all ranks).
+    if (threadIdx.x == 0)
+    {
+        int32_t segs = 0;
+        uint32_t off = 0;
+        for (int32_t q = 0; q < sources; ++q)
+        {
+            if (!sourceOverflows(s, q, kth))
+            {
+                continue;
+            }
+            uint32_t const* slot = p.ucWorkspace + round2Offset(buf, q, row);
+            uint2 const h = pollPair(slot, wait, kWaitRound2Header);
+            segSlots[segs] = slot;
+            segBegin[segs] = off;
+            off += min(h.x, static_cast<uint32_t>(kRound2Capacity));
+            ++segs;
+        }
+        segBegin[segs] = off;
+        s.scal[kSlotTotal] = static_cast<uint32_t>(segs);
+    }
+    __syncthreads();
+    int32_t const segs = static_cast<int32_t>(s.scal[kSlotTotal]);
+    int32_t const total2 = static_cast<int32_t>(segBegin[segs]);
+    MergeKeys const keys{s.mkey, s.midx, total, segSlots, segBegin, segs};
+    for (int32_t k = threadIdx.x; k < total2; k += kThreads)
+    {
+        (void) pollPair(keys.pair(k), wait, kWaitRound2Candidates);
+    }
+    __syncthreads();
+    SelectOut const m2
+        = blockSelect(s, keys, total + total2, 0u, uint64_t{kEmptyWord}, static_cast<uint32_t>(kTopK), 0u, 0u, 32);
+    bool takeAll2 = false;
+    uint32_t const kth2 = mergeThreshold(m2, takeAll2);
+    (void) emitSelection(s, keys, total + total2, kth2, takeAll2, outRow, 0, p, bs, row);
+    if (threadIdx.x == 0)
+    {
+        atomicAdd(&p.state[2], 1);
+    }
+}
+
+// ---- Thread-block cluster primitives (one cluster of kClusterCtas CTAs per row; cluster rank = blockIdx.x) ----------
+
+__device__ __forceinline__ uint32_t smemAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// Shared-memory window address of `ptr` (this CTA's shared memory) in CTA `rank` of the cluster.
+__device__ __forceinline__ uint32_t clusterAddr(void const* ptr, int32_t rank)
+{
+    uint32_t remote;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smemAddr(ptr)), "r"(rank));
+    return remote;
+}
+
+// Initialize the exchange barriers and announce it to the cluster (relaxed arrive; clusterInitWait completes it
+// before the first exchange). Touches shared memory only.
+__device__ __forceinline__ void clusterInit(ClusterRecords& rec)
+{
+    if (threadIdx.x == 0)
+    {
+        asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[0])) : "memory");
+        asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[1])) : "memory");
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+    asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory");
+}
+
+__device__ __forceinline__ void clusterInitWait()
+{
+    asm volatile("barrier.cluster.wait.aligned;" ::: "memory");
+}
+
+// This CTA's single arrival on exchange barrier `bar`, expecting `bytes` from the cluster's st.async stores.
+__device__ __forceinline__ void expectBytes(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile(
+        "{\n\t.reg .b64 st;\n\tmbarrier.arrive.expect_tx.shared::cta.b64 st, [%0], %1;\n\t}" ::"r"(smemAddr(bar)),
+        "r"(bytes)
+        : "memory");
+}
+
+// Wait until the phase of exchange barrier `bar` with parity `parity` completed (every expected byte arrived).
+__device__ __forceinline__ void waitExchange(uint64_t* bar, uint32_t parity)
+{
+    uint32_t done;
+    do
+    {
+        asm volatile(
+            "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], %2;\n\t"
+            "selp.u32 %0, 1, 0, p;\n\t}"
+            : "=r"(done)
+            : "r"(smemAddr(bar)), "r"(parity)
+            : "memory");
+    } while (done == 0u);
+}
+
+// Store 16 / 8 bytes at `ptr` (a record of this CTA's layout) in CTA `rank`, completing them on its barrier `bar`.
+__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint4 v)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(
+                     clusterAddr(ptr, rank)),
+                 "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(clusterAddr(bar, rank))
+                 : "memory");
+}
+
+__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint2 v)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.b32 [%0], {%1, %2}, [%3];" ::"r"(
+                     clusterAddr(ptr, rank)),
+                 "r"(v.x), "r"(v.y), "r"(clusterAddr(bar, rank))
+                 : "memory");
+}
+
+// Exchange sequence of a row: per pass x = 2 * pass group totals and x = 2 * pass + 1 the selected group's bins.
+// Every CTA sends the records of exchange x + 1 only after its own exchange x completed, so a barrier phase (and an
+// alternating record set) is reused only after every CTA consumed the previous one.
+__device__ __forceinline__ uint64_t* exchangeBar(ClusterRecords& rec, uint32_t x)
+{
+    return &rec.bar[x & 1u];
+}
+
+__device__ __forceinline__ uint32_t exchangeParity(uint32_t x)
+{
+    return (x >> 1) & 1u;
+}
+
+// Inclusive scan over the lanes of a warp.
+__device__ __forceinline__ uint32_t warpInclusiveScan(uint32_t v)
+{
+    int32_t const lane = threadIdx.x % 32;
+#pragma unroll
+    for (int32_t o = 1; o < 32; o <<= 1)
+    {
+        uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, v, o);
+        if (lane >= o)
+        {
+            v += y;
+        }
+    }
+    return v;
+}
+
+// Tallies of the cluster select over the candidates of the cluster CTAs before this one (in candidate order), kept
+// by warp 0 and published in shared memory: those strictly above the selected bins of all passes so far, and those
+// in the selected bin of the last pass.
+struct CtaTally
+{
+    uint32_t aboveBefore;
+    uint32_t binBefore;
+};
+
+// Cluster radix select of the need-th largest key among the candidates of every CTA of the cluster (this thread
+// holds k0 and k1; kEmptyWord is no candidate, as it lies above every range). It computes exactly what blockSelect
+// with cap = 0 computes over the union of the candidates: the same digits (from `base` and topBits), the same bins
+// and counts from the cluster-summed histograms, and the same stop rules. Per pass: every CTA histograms its own
+// candidates and sends its totals per group of kGroupBins bins to every CTA (st.async, two groups per warp); warp 0
+// of every CTA finds, from the top, the group holding the target, sends its own counts of that group's bins to every
+// CTA, and finds the bin from the eight CTAs' counts. From the same records, warp 0 also tallies the candidates of
+// the CTAs before this one above the selected bins and in the last selected bin (`tally`), which give this CTA's
+// emission offsets without another exchange.
+__device__ SelectOut clusterSelect(SmemLayout const& s, ClusterRecords& rec, CtaTally& tally, uint32_t k0, uint32_t k1,
+    uint32_t lo, uint64_t hi, uint32_t need, uint32_t base, int32_t topBits, uint32_t aboveHi, int64_t* dbg)
+{
+    int32_t const me = blockIdx.x;
+    int32_t const lane = threadIdx.x % 32;
+    int32_t const warp = threadIdx.x / 32;
+    SelectOut out{lo, 0u, 0u, 0u, false, false, 0};
+    uint32_t prefMask = 0;
+    uint32_t prefVal = 0;
+    uint32_t above = aboveHi;
+    int32_t remaining = max(topBits, 1);
+    uint32_t* hist = s.hist;
+    uint32_t* histNext = s.hist + kBins;
+    // Warp 0: the running tallies of the CTAs before this one.
+    uint32_t aboveBefore = 0u;
+    uint32_t binBefore = 0u;
+    for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads)
+    {
+        hist[b] = 0;
+    }
+    __syncthreads();
+#pragma unroll 1
+    for (int32_t pass = 0; remaining > 0; ++pass)
+    {
+        int32_t const width = min(kDigitBits, remaining);
+        int32_t const shift = remaining - width;
+        int32_t const nbins = 1 << width;
+        uint32_t const binMask = static_cast<uint32_t>(nbins - 1);
+        int32_t const groups = (nbins + kGroupBins - 1) / kGroupBins;
+        int32_t const pairs = (groups + 1) / 2;
+        uint32_t const xg = 2u * pass;
+        uint32_t const xb = xg + 1u;
+        out.passes = pass + 1;
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[1 + 3 * pass] = clock64();
+        }
+        if (threadIdx.x == 0)
+        {
+            expectBytes(exchangeBar(rec, xg), kClusterCtas * pairs * 8);
+        }
+        // The next pass's histogram: its last reads (the previous pass's bins sent by warp 0) are complete.
+        if (shift > 0)
+        {
+            for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads)
+            {
+                histNext[b] = 0;
+            }
+        }
+        histogramKey(hist, true, k0, lo, hi, base, prefMask, prefVal, shift, binMask);
+        histogramKey(hist, true, k1, lo, hi, base, prefMask, prefVal, shift, binMask);
+        __syncthreads();
+        uint32_t(*const groupRec)[kMaxGroups] = rec.groups[pass & 1];
+        uint32_t(*const binRec)[kGroupBins] = rec.bins[pass & 1];
+        if (warp < pairs)
+        {
+            int32_t const b0 = 2 * warp * kGroupBins + lane;
+            int32_t const b1 = b0 + kGroupBins;
+            uint32_t const t0 = __reduce_add_sync(0xFFFFFFFFu, b0 < nbins ? hist[b0] : 0u);
+            uint32_t const t1 = __reduce_add_sync(0xFFFFFFFFu, b1 < nbins ? hist[b1] : 0u);
+            if (lane < kClusterCtas)
+            {
+                sendTo(lane, &groupRec[me][2 * warp], exchangeBar(rec, xg), make_uint2(t0, t1));
+            }
+        }
+        uint32_t const target = need - above;
+        if (warp == 0)
+        {
+            waitExchange(exchangeBar(rec, xg), exchangeParity(xg));
+            if (dbg != nullptr && lane == 0 && pass < 3)
+            {
+                dbg[2 + 3 * pass] = clock64();
+            }
+            // Lane l sums the groups at positions 2l and 2l + 1 from the top over the cluster.
+            int32_t const ga = groups - 1 - 2 * lane;
+            int32_t const gb = ga - 1;
+            uint32_t a = 0;
+            uint32_t b = 0;
+            uint32_t beforeA = 0;
+            uint32_t beforeB = 0;
+#pragma unroll
+            for (int32_t c = 0; c < kClusterCtas; ++c)
+            {
+                uint32_t const x = ga >= 0 ? groupRec[c][ga] : 0u;
+                uint32_t const y = gb >= 0 ? groupRec[c][gb] : 0u;
+                a += x;
+                b += y;
+                beforeA += c < me ? x : 0u;
+                beforeB += c < me ? y : 0u;
+            }
+            uint32_t const incl = warpInclusiveScan(a + b);
+            uint32_t const before = incl - a - b;
+            uint32_t const total = __shfl_sync(0xFFFFFFFFu, incl, 31);
+            bool const holds = before < target && target <= incl;
+            uint32_t const holder = __ballot_sync(0xFFFFFFFFu, holds);
+            if (lane == 0)
+            {
+                s.scal[kSlotTotal] = total;
+            }
+            if (holder != 0u)
+            {
+                int32_t const hl = __ffs(holder) - 1;
+                bool const first = target <= before + a;
+                int32_t const gSel = __shfl_sync(0xFFFFFFFFu, first ? ga : gb, hl);
+                uint32_t const gAbove = __shfl_sync(0xFFFFFFFFu, first ? before : before + a, hl);
+                if (lane == 0)
+                {
+                    expectBytes(exchangeBar(rec, xb), kClusterCtas * kGroupBins * 4);
+                }
+                // This CTA's counts of the selected group's bins to every CTA: lane l sends vector l % 8 to CTAs
+                // l / 8 and l / 8 + 4.
+                int32_t const v = lane % 8;
+                int32_t const b0 = gSel * kGroupBins + 4 * v;
+                uint4 const q = make_uint4(b0 < nbins ? hist[b0] : 0u, b0 + 1 < nbins ? hist[b0 + 1] : 0u,
+                    b0 + 2 < nbins ? hist[b0 + 2] : 0u, b0 + 3 < nbins ? hist[b0 + 3] : 0u);
+                sendTo(lane / 8, &binRec[me][4 * v], exchangeBar(rec, xb), q);
+                sendTo(lane / 8 + 4, &binRec[me][4 * v], exchangeBar(rec, xb), q);
+                // While the bins travel: the earlier CTAs' candidates in the groups above the selected one (both
+                // groups of the lanes before the holder, and the holder's first group when it selected the second).
+                bool const selFirst = __shfl_sync(0xFFFFFFFFu, first, hl);
+                aboveBefore += __reduce_add_sync(
+                    0xFFFFFFFFu, lane < hl ? beforeA + beforeB : (lane == hl && !selFirst ? beforeA : 0u));
+                waitExchange(exchangeBar(rec, xb), exchangeParity(xb));
+                // Bins of the selected group from the top: lane l holds bin 32 * gSel + 31 - l.
+                int32_t const bin = gSel * kGroupBins + kGroupBins - 1 - lane;
+                uint32_t c = 0;
+                uint32_t cBefore = 0;
+#pragma unroll
+                for (int32_t d = 0; d < kClusterCtas; ++d)
+                {
+                    uint32_t const x = binRec[d][kGroupBins - 1 - lane];
+                    c += x;
+                    cBefore += d < me ? x : 0u;
+                }
+                uint32_t const y = warpInclusiveScan(c);
+                uint32_t const acc = gAbove + y - c;
+                bool const found = acc < target && target <= acc + c;
+                if (found)
+                {
+                    s.scal[kSlotBin] = static_cast<uint32_t>(bin);
+                    s.scal[kSlotAbove] = acc;
+                    s.scal[kSlotBinCount] = c;
+                }
+                // The earlier CTAs' candidates in the bins above the selected one (lanes before the finder) and in
+                // the selected bin (the finder's).
+                int32_t const finder = __ffs(__ballot_sync(0xFFFFFFFFu, found)) - 1;
+                aboveBefore += __reduce_add_sync(0xFFFFFFFFu, lane < finder ? cBefore : 0u);
+                binBefore = __shfl_sync(0xFFFFFFFFu, cBefore, finder);
+            }
+            else if (lane == 0)
+            {
+                // Every key is taken (pass 0 only, the same decision in every CTA): the bin exchange still
+                // completes, with no bytes, so that the exchange sequence stays aligned.
+                expectBytes(exchangeBar(rec, xb), 0u);
+                waitExchange(exchangeBar(rec, xb), exchangeParity(xb));
+            }
+            if (lane == 0)
+            {
+                tally.aboveBefore = aboveBefore;
+                tally.binBefore = binBefore;
+            }
+        }
+        __syncthreads();
+        if (dbg != nullptr && threadIdx.x == 0 && pass < 3)
+        {
+            dbg[3 + 3 * pass] = clock64();
+        }
+        uint32_t const total = s.scal[kSlotTotal];
+        if (pass == 0)
+        {
+            out.total = total + aboveHi;
+            if (total + aboveHi == 0u || target > total)
+            {
+                out.all = true;
+                out.above = total + aboveHi;
+                return out;
+            }
+        }
+        // The scalars are next written by warp 0 after the next histogram's block barrier.
+        uint32_t const bin = s.scal[kSlotBin];
+        above += s.scal[kSlotAbove];
+        out.binCount = s.scal[kSlotBinCount];
+        prefVal |= bin << shift;
+        prefMask |= binMask << shift;
+        out.edge = base + prefVal;
+        out.above = above;
+        remaining = shift;
+        uint32_t* const t = hist;
+        hist = histNext;
+        histNext = t;
+        if (remaining > 0 && above + out.binCount == need)
+        {
+            out.full = true;
+            return out;
+        }
+    }
+    return out;
+}
+
+// Spin until the header of a round-1 slot was written and, when the slot's count covers candidate vector kv, that
+// vector too; returns the header's count (clamped to `cap`) and bound. Both are polled together, so neither waits
+// for the other; aborts the kernel after the timeout.
+__device__ uint2 pollSlotVector(uint32_t const* slot, int32_t kv, uint32_t cap, uint4& v, WaitCtx const& w)
+{
+    uint32_t const* vec = slot + kSlotHeaderWords + 4 * kv;
+    uint2 h = loadVolatile2(slot);
+    v = loadVolatile4(vec);
+    uint64_t start = 0;
+    for (int32_t polls = 1;; ++polls)
+    {
+        bool const headerReady = !isEmpty(h);
+        uint32_t const cnt = headerReady ? min(h.x, cap) : 0u;
+        if (headerReady && (2u * kv >= cnt || vectorReady(v, 2u * kv + 1u < cnt)))
+        {
+            return make_uint2(cnt, h.y);
+        }
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, headerReady ? kWaitRound1Candidates : kWaitRound1Header);
+            }
+        }
+        if (!headerReady)
+        {
+            h = loadVolatile2(slot);
+        }
+        v = loadVolatile4(vec);
+    }
+}
+
+// Spin until the whole 32-byte round-1 header of a slot was written; aborts the kernel after the timeout.
+__device__ void pollHeader(uint32_t const* slot, uint4& a, uint4& b, WaitCtx const& w)
+{
+    a = loadVolatile4(slot);
+    b = loadVolatile4(slot + 4);
+    uint64_t start = 0;
+    for (int32_t polls = 1; !(vectorReady(a, true) && vectorReady(b, true)); ++polls)
+    {
+        if (polls % kPollsPerTimerRead == 0)
+        {
+            uint64_t const now = globalTimerNs();
+            if (start == 0)
+            {
+                start = now;
+            }
+            else if (now - start > w.timeoutNs)
+            {
+                abortOnTimeout(w, kWaitRound1Header);
+            }
+        }
+        a = loadVolatile4(slot);
+        b = loadVolatile4(slot + 4);
+    }
+}
+
+// Merge of one row by its cluster. CTA c gathers the round-1 candidates of sources [c * world, (c + 1) * world) (one
+// 16-byte vector, two candidate slots, per thread of the poll warps, kept in registers: thread order is crossRankTopk's
+// candidate order) while its header warps read every source's header (counts, bounds and the key statistics of the keys
+// sent, so every CTA has the row's statistics and can decide every source's overflow without an exchange). The cluster
+// runs crossRankTopk's merge select on the summed histograms (clusterSelect), whose records also give every CTA's
+// number of keys above the threshold and of ties, and every CTA writes its selected candidates at crossRankTopk's
+// positions: keys above the threshold in candidate order, then ties in candidate order while they fit, then -1, each
+// position with its bridge outputs (pool bases fetched per candidate during the cluster select). Returns true, writing
+// nothing, when a source withheld keys that may lie above the threshold (the same decision on every CTA and rank); the
+// row is then merged by mergeRow.
+__device__ bool clusterMergeRow(CpTopkParams const& p, SmemLayout const& s, ClusterRecords& rec, CtaTally& tally,
+    int32_t buf, int32_t row, int32_t rowLen, int32_t* outRow, WaitCtx const& wait, BridgeShared const& bs)
+{
+    int32_t const world = p.worldSize;
+    int32_t const sources = numSources(world);
+    int32_t const me = blockIdx.x;
+    int32_t const tid = threadIdx.x;
+    int32_t const lane = tid % 32;
+    int32_t const warp = tid / 32;
+    uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+    int32_t const vectorsPerSource = static_cast<int32_t>(cap) / 2;
+    int32_t const q0 = me * world;
+    int64_t* const dbg = debugSlots(p, 0);
+    // This thread's candidates (kEmptyWord: none; never a key).
+    uint32_t k0 = kEmptyWord;
+    uint32_t k1 = kEmptyWord;
+    if (tid == 0)
+    {
+        s.scal[kSlotCounter2] = 0u;
+    }
+    if (tid < world * vectorsPerSource)
+    {
+        int32_t const q = q0 + tid / vectorsPerSource;
+        int32_t const kv = tid % vectorsPerSource;
+        uint4 v;
+        uint2 const h = pollSlotVector(p.ucWorkspace + round1Offset(buf, q, row), kv, cap, v, wait);
+        if (2u * kv < h.x)
+        {
+            k0 = v.x;
+            if (2u * kv + 1u < h.x)
+            {
+                k1 = v.z;
+            }
+        }
+        // The tokens wait in shared memory until the emission (s.midx is free in the cluster merge). Their pool bases
+        // are fetched now, while the cluster selects: copies of the candidates' block-table entries into this
+        // thread's words beyond kBaseOffsetWords, completed before the block barrier of the emission prefix.
+        *reinterpret_cast<uint2*>(s.midx + 2 * tid) = make_uint2(v.y, v.w);
+        if (k0 != kEmptyWord)
+        {
+            int32_t const req = sharedRequest(p.bridge, bs);
+            uint32_t* const bases = s.midx + kBaseOffsetWords + 2 * tid;
+            fetchBase(p.bridge, req, static_cast<int32_t>(v.y), bases);
+            if (k1 != kEmptyWord)
+            {
+                fetchBase(p.bridge, req, static_cast<int32_t>(v.w), bases + 1);
+            }
+        }
+    }
+    else if (warp >= kPollWarps && warp < kFirstResetWarp)
+    {
+        // Header warps: thread h reads source h's header; per source its count (clamped to the round-1 capacity),
+        // bound and length, and per header warp the reduced key statistics.
+        int32_t const q = tid - kPollWarps * 32;
+        KeyStats st;
+        if (q < sources)
+        {
+            uint4 a;
+            uint4 b;
+            pollHeader(p.ucWorkspace + round1Offset(buf, q, row), a, b, wait);
+            s.scal[kSlotSrcCount + q] = min(a.x, cap);
+            s.scal[kSlotSrcBound + q] = a.y;
+            s.scal[kSlotSrcLen + q] = static_cast<uint32_t>(sourceLength(
+                min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank));
+            st.maxKey = a.z;
+            st.minKey = a.w - 1u;
+            st.positives = b.x;
+            st.maxNeg = b.y;
+            st.minPos = b.z - 1u;
+        }
+        uint32_t const cnt = q < sources ? s.scal[kSlotSrcCount + q] : 0u;
+        // Candidates and keys >= +0.0 of the cluster CTAs before this one (CTA c merges sources
+        // [c * world, (c + 1) * world)), of this one and of the row.
+        int32_t const c = q < sources ? q / world : kClusterCtas;
+        uint32_t const nBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? cnt : 0u);
+        uint32_t const nMine = __reduce_add_sync(0xFFFFFFFFu, c == me ? cnt : 0u);
+        uint32_t const nAll = __reduce_add_sync(0xFFFFFFFFu, cnt);
+        uint32_t const pBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? st.positives : 0u);
+        if (lane == 0)
+        {
+            atomicAdd(&s.scal[kSlotCountBefore], nBefore);
+            atomicAdd(&s.scal[kSlotCountMine], nMine);
+            atomicAdd(&s.scal[kSlotCountTotal], nAll);
+            atomicAdd(&s.scal[kSlotPositivesBefore], pBefore);
+        }
+        uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, st.maxKey);
+        uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, st.minKey);
+        uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, st.positives);
+        uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, st.maxNeg);
+        uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, st.minPos);
+        if (lane == 0)
+        {
+            // In the scan scratch (s.stats may still be read by the send's last warp; s.warpSums is next written
+            // by the emission scan, after the reads below).
+            uint32_t* const out = s.warpSums + kNumKeyStats * (warp - kPollWarps);
+            out[0] = v0;
+            out[1] = v1;
+            out[2] = v2;
+            out[3] = v3;
+            out[4] = v4;
+        }
+    }
+    stamp(dbg, 45);
+    __syncthreads();
+    markPhase(p, 4);
+
+    // Statistics of the row's candidates from every source's header (identical in every CTA and rank).
+    KeyStats g;
+#pragma unroll
+    for (int32_t w = 0; w < kHeaderWarps; ++w)
+    {
+        uint32_t const* const in = s.warpSums + kNumKeyStats * w;
+        g.maxKey = max(g.maxKey, in[0]);
+        g.minKey = min(g.minKey, in[1]);
+        g.positives += in[2];
+        g.maxNeg = max(g.maxNeg, in[3]);
+        g.minPos = min(g.minPos, in[4]);
+    }
+    uint32_t const total = s.scal[kSlotCountTotal];
+    stamp(dbg, 61);
+    markPhase(p, 5);
+
+    SignRange const r
+        = total > 0u ? signRange(g, static_cast<uint32_t>(kTopK)) : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false};
+    clusterInitWait();
+    SelectOut const m = clusterSelect(
+        s, rec, tally, k0, k1, r.lo, r.hi, static_cast<uint32_t>(kTopK), r.lo, r.bits, r.aboveHi, debugSlots(p, 48));
+    putDebug(debugSlots(p, 48), 0, s.scal[kSlotCountMine]);
+    putDebug(dbg, 62, total);
+    putDebug(dbg, 63, m.passes + (r.negative ? 10 : 0));
+    bool takeAll = false;
+    uint32_t const kth = mergeThreshold(m, takeAll);
+    markPhase(p, 6);
+
+    // Keys above the threshold and ties of the cluster CTAs before this one and of the row (identical in every CTA),
+    // from the select's tallies; with the threshold below a wholly taken bin the ties fall beyond kTopK and are taken
+    // as keys above it.
+    uint32_t gtBefore;
+    uint32_t eqBefore;
+    uint32_t allGt;
+    uint32_t allEq;
+    if (takeAll)
+    {
+        gtBefore = s.scal[kSlotCountBefore];
+        eqBefore = 0u;
+        allGt = total;
+        allEq = 0u;
+    }
+    else
+    {
+        uint32_t const aboveBefore = tally.aboveBefore + (r.aboveHi > 0u ? s.scal[kSlotPositivesBefore] : 0u);
+        gtBefore = m.full ? aboveBefore + tally.binBefore : aboveBefore;
+        eqBefore = m.full ? 0u : tally.binBefore;
+        allGt = m.full ? m.above + m.binCount : m.above;
+        allEq = m.full ? 0u : m.binCount;
+    }
+    // This CTA's positions, scanned in thread order; a source of the row that withheld keys above the threshold
+    // sends the row to the second round.
+    bool const gt0 = k0 != kEmptyWord && (takeAll || k0 > kth);
+    bool const eq0 = k0 != kEmptyWord && !takeAll && k0 == kth;
+    bool const gt1 = k1 != kEmptyWord && (takeAll || k1 > kth);
+    bool const eq1 = k1 != kEmptyWord && !takeAll && k1 == kth;
+    if (tid < sources && sourceOverflows(s, tid, kth))
+    {
+        atomicOr(&s.scal[kSlotCounter2], 1u);
+    }
+    // Exclusive prefix in thread order (k0 before k1): within the warp from ballots, across warps from the warp
+    // totals (one block barrier).
+    uint32_t const below = (1u << lane) - 1u;
+    uint32_t const bGt0 = __ballot_sync(0xFFFFFFFFu, gt0);
+    uint32_t const bGt1 = __ballot_sync(0xFFFFFFFFu, gt1);
+    uint32_t const bEq0 = __ballot_sync(0xFFFFFFFFu, eq0);
+    uint32_t const bEq1 = __ballot_sync(0xFFFFFFFFu, eq1);
+    if (lane == 0)
+    {
+        s.warpSums[warp] = static_cast<uint32_t>(__popc(bGt0) + __popc(bGt1))
+            + (static_cast<uint32_t>(__popc(bEq0) + __popc(bEq1)) << 16);
+    }
+    // The pool-base copies are complete before the barrier (a second-round row reuses s.midx).
+    asm volatile("cp.async.wait_all;" ::: "memory");
+    __syncthreads();
+    uint32_t const warpsBefore = __reduce_add_sync(0xFFFFFFFFu, lane < warp ? s.warpSums[lane] : 0u);
+    uint32_t const before = warpsBefore + static_cast<uint32_t>(__popc(bGt0 & below) + __popc(bGt1 & below))
+        + (static_cast<uint32_t>(__popc(bEq0 & below) + __popc(bEq1 & below)) << 16);
+    stamp(dbg, 58);
+    if (s.scal[kSlotCounter2] != 0u)
+    {
+        return true;
+    }
+    uint32_t const want = static_cast<uint32_t>(kTopK);
+    uint32_t posGt = gtBefore + (before & 0xFFFFu);
+    uint32_t posEq = allGt + eqBefore + (before >> 16);
+    uint2 tokens = make_uint2(0u, 0u);
+    uint2 tokenBases = make_uint2(0u, 0u);
+    if (k0 != kEmptyWord)
+    {
+        tokens = *reinterpret_cast<uint2 const*>(s.midx + 2 * tid);
+        tokenBases = *reinterpret_cast<uint2 const*>(s.midx + kBaseOffsetWords + 2 * tid);
+    }
+    // Output positions of this thread's two candidates (want: not written) and of its share of the -1 padding (at
+    // most one position per thread of the cluster).
+    static_assert(kClusterCtas * kThreads >= kTopK, "one padding position per thread");
+    uint32_t const pos0 = gt0 ? posGt : (eq0 ? posEq : want);
+    posGt += gt0 ? 1u : 0u;
+    posEq += eq0 ? 1u : 0u;
+    uint32_t const pos1 = gt1 ? posGt : (eq1 ? posEq : want);
+    uint32_t const selected = min(allGt + allEq, want);
+    uint32_t const posPad = selected + static_cast<uint32_t>(me * kThreads + tid);
+    emitCandidatesFor(p.bridge.groupSize, outRow, bs.poolRow, bs.copyRow, p.bridge.outStrideG,
+        static_cast<uint32_t>(p.bridge.strideFactor), bs.layer, pos0, pos1, posPad, tokens, tokenBases);
+    stamp(dbg, 60);
+    return false;
+}
+
+__global__ void __launch_bounds__(kThreads, 1) cpTopkZ11Kernel(CpTopkParams p)
+{
+    markPhase(p, 0);
+    extern __shared__ __align__(16) uint32_t smemRaw[];
+    __shared__ uint32_t const* segSlots[kMaxSources];
+    __shared__ uint32_t segBegin[kMaxSources + 1];
+    __shared__ __align__(16) ClusterRecords rec;
+    __shared__ CtaTally tally;
+    __shared__ int32_t epochShared;
+    __shared__ __align__(16) BridgeShared bridgeRow;
+    __shared__ int32_t emitScratch[kWarmScratchWords];
+    int32_t const world = p.worldSize;
+    int32_t const rank = p.rank;
+    int32_t const src = blockIdx.x;
+    int32_t const row = blockIdx.y;
+    SmemLayout const s = carveSmem(smemRaw, p.localCapacity, world);
+    int32_t* outRow = p.out + static_cast<int64_t>(row) * kTopK;
+    clusterInit(rec);
+    if (threadIdx.x <= kSlotPositivesBefore - kSlotCountBefore)
+    {
+        s.scal[kSlotCountBefore + threadIdx.x] = 0u;
+    }
+    int32_t const resetOffset = resetHeaderOffset();
+    // Every global read (lengths, logits, the epoch word, the workspace, the bridge inputs) comes after the grid
+    // dependency wait, so a programmatic early launch is safe whatever the preceding kernels write.
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    markPhase(p, 10);
+    // Thread 0 and the reset warps read this CTA's epoch word, written only by the previous exchange launch of this
+    // op, before the key loads; nothing waits for the value until the key loads are done. Thread 0 publishes it to the
+    // CTA at the block barrier that follows the key loads, and the reset warps start the copies of the previous
+    // buffer's headers right before that barrier (slot offsets computed before the wait, from the grid only).
+    int32_t* const epochWord = p.state + kEpochBase + row * kSourcesPerRank + src;
+    int32_t epochRead = 0;
+    if (threadIdx.x == 0 || static_cast<int32_t>(threadIdx.x) >= kFirstResetWarp * 32)
+    {
+        epochRead = *reinterpret_cast<int32_t volatile*>(epochWord);
+    }
+    int32_t const prevBufRead = (epochRead + kNumLamportBuffers - 1) % kNumLamportBuffers;
+    int32_t const kvLen = p.kvLens[0];
+    int32_t const rowLen = max(0, kvLen - p.nextN + row + 1);
+    // Identity launch (every row has at most kTopK tokens; the same decision on every rank): no exchange, and
+    // neither the workspace nor the epoch words are touched, so identity launches never shorten the reuse chain of
+    // the Lamport buffers (see crossRankTopkGlobalEmit.h).
+    bool const exchange = kvLen - p.nextN + p.numRows > kTopK;
+    // Thread 0 advances this CTA's epoch word (the CTAs of the last row also those of the rows this launch does not
+    // have, so every word counts every exchange launch), after a block barrier that follows every read of the word in
+    // this CTA. Nothing else of this launch reads the epoch words, and the next launch reads them only after its grid
+    // dependency wait.
+    auto advanceEpoch = [&]()
+    {
+        if (exchange && threadIdx.x == 0)
+        {
+            int32_t const rows = row == static_cast<int32_t>(gridDim.y) - 1 ? kMaxRows - row : 1;
+            for (int32_t r = 0; r < rows; ++r)
+            {
+                epochWord[r * kSourcesPerRank] = epochRead + 1;
+            }
+        }
+    };
+
+    if (rowLen <= kTopK)
+    {
+        // Every rank agrees on this branch: identity selection, no exchange for this row. Each CTA of the row writes
+        // kTopK / kClusterCtas positions with their bridge outputs (bridge inputs read directly).
+        constexpr int32_t kIdentityPerCta = kTopK / kClusterCtas;
+        if (static_cast<int32_t>(threadIdx.x) < kIdentityPerCta)
+        {
+            CpTopkBridgeParams const& b = p.bridge;
+            int32_t const pos = src * kIdentityPerCta + static_cast<int32_t>(threadIdx.x);
+            int32_t const tok = pos < rowLen ? pos : -1;
+            outRow[pos] = tok;
+            int32_t const req = b.reqIdx[row] - b.reqIdxBias;
+            BridgeRowOut const bridgeOut = bridgeRowOut(b, row, bridgeCopyRow(b, row, loadAccepted(b, row)));
+            writeBridge(b, bridgeOut, static_cast<uint32_t>(pos), tok, directBase(b, req, tok),
+                [&](int32_t g) { return b.layerIds != nullptr ? b.layerIds[g] : b.layerId; });
+        }
+        if (exchange)
+        {
+            fetchResetHeader(p, s, prevBufRead, resetOffset);
+            resetSlots(p, s, prevBufRead);
+            markPhase(p, 9);
+        }
+        __syncthreads();
+        advanceEpoch();
+        clusterInitWait();
+    }
+    else
+    {
+        int32_t const localLen = min(localLength(rowLen, rank, world), p.localCapacity);
+        float const* rowLogits = p.logits + static_cast<int64_t>(row) * p.logitsStride;
+        // The first reset warp starts the copies of the row's request index, layer ids and accepted count (completed
+        // with the reset).
+        if (static_cast<int32_t>(threadIdx.x) / 32 == kFirstResetWarp)
+        {
+            fetchRowScalars(p, bridgeRow, row);
+        }
+
+        // ---- Round 1: this source's local top keys ---------------------------------------------------------
+        int32_t const n = loadSource(s, rowLogits, localLen, src, true);
+        if (threadIdx.x == 0)
+        {
+            epochShared = epochRead;
+        }
+        // The request index copy is complete before the barrier (only its group is committed at this point).
+        if (static_cast<int32_t>(threadIdx.x) == kFirstResetWarp * 32 + kReqLane)
+        {
+            asm volatile("cp.async.wait_group 0;" ::: "memory");
+        }
+        // The reset warps start fetching the headers of their share of the previous exchange launch's buffer (an
+        // exchange row implies an exchange launch); the values are first used by the reset after the send.
+        fetchResetHeader(p, s, prevBufRead, resetOffset);
+        __syncthreads();
+        markPhase(p, 1);
+        advanceEpoch();
+        uint32_t const cap = static_cast<uint32_t>(roundCap(world));
+        uint32_t const need = static_cast<uint32_t>(roundNeed(world));
+        SelectOut sel{0u, 0u, 0u, static_cast<uint32_t>(n), true, false, 0};
+        uint32_t selLo = 0u;
+        int32_t passInfo = 0;
+        if (n > static_cast<int32_t>(cap))
+        {
+            // The cut lies among the keys of one sign (see kPositiveKey); the digits span that sign's key range.
+            SignRange const r = signRange(readStats(s.stats), need);
+            selLo = r.lo;
+            SmemKeys const keys{s.chunk, nullptr};
+            sel = blockSelect(s, keys, n, r.lo, r.hi, need, cap, r.lo, r.bits, debugSlots(p, 32), r.aboveHi);
+            passInfo = sel.passes + (r.negative ? 10 : 0);
+        }
+        markPhase(p, 2);
+        int32_t const sendBuf = epochShared % kNumLamportBuffers;
+        uint32_t const sent = sendSource<false>(s, sel, n, selLo, uint64_t{kEmptyWord}, cap,
+            p.mcWorkspace + round1Offset(sendBuf, rank * kSourcesPerRank + src, row), src, rank, world,
+            reinterpret_cast<uint2*>(s.mkey), debugSlots(p, 0));
+        putDebug(debugSlots(p, 0), 46, sent);
+        putDebug(debugSlots(p, 0), 47, passInfo);
+        markPhase(p, 3);
+
+        // The reset warps reset the previous exchange launch's buffer, whose candidates every rank has consumed (a
+        // peer writes it again only two exchange launches later, after this launch completed on this rank), while
+        // the poll and header warps wait for this launch's candidates.
+        resetSlots(p, s, (epochShared + kNumLamportBuffers - 1) % kNumLamportBuffers);
+        int32_t const buf = epochShared % kNumLamportBuffers;
+        if (threadIdx.x == kFirstResetWarp * 32)
+        {
+            int64_t* const d = debugSlots(p, 0);
+            if (d != nullptr)
+            {
+                d[9] = static_cast<int64_t>(globalTimerNs());
+            }
+            // The row's pool and copy rows (the accepted count is visible after resetSlots' warp barrier), read by
+            // the emission after the block barrier that follows the candidate poll.
+            CpTopkBridgeParams const& b = p.bridge;
+            int32_t const copyRow = bridgeCopyRow(b, row, sharedAccepted(b, bridgeRow));
+            bridgeRow.poolRow = b.out + row * b.outStride0;
+            bridgeRow.copyRow = copyRow >= 0 ? b.copyDst + copyRow * b.copyStride0 : nullptr;
+            // The emission's instructions, run once now while this warp waits for the poll (see emitCandidates);
+            // the layer ids are complete after resetSlots.
+            warmEmission(b.groupSize, emitScratch, bridgeRow.layer);
+        }
+
+        // ---- Merge: the row's cluster -----------------------------------------------------------------------
+        WaitCtx const wait{p.state, p.timeoutNs, rank};
+        if (clusterMergeRow(p, s, rec, tally, buf, row, rowLen, outRow, wait, bridgeRow) && src == 0)
+        {
+            // A source withheld keys above the threshold: crossRankTopk's merge, with its second round, on CTA 0.
+            mergeRow(p, s, segSlots, segBegin, buf, row, rowLen, localLen, rowLogits, outRow, wait, bridgeRow);
+            int64_t* const dbg = debugSlots(p, 0);
+            putDebug(dbg, 63, dbg != nullptr ? dbg[63] + 100 : 0);
+        }
+        markPhase(p, 7);
+    }
+    markPhase(p, 8);
+}
+
+#endif
+
+} // namespace
+
+size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize)
+{
+    return smemBytes(localCapacity, worldSize);
+}
+
+int32_t cpTopkMaxLocalCapacity(int32_t worldSize)
+{
+    static int32_t cached[kMaxRanks + 1] = {};
+    if (worldSize < 1 || worldSize > kMaxRanks)
+    {
+        return 0;
+    }
+    if (cached[worldSize] == 0)
+    {
+        int device = 0;
+        cudaGetDevice(&device);
+        int maxSmemOptin = 0;
+        cudaDeviceGetAttribute(&maxSmemOptin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+        int32_t cap = 0;
+        for (int32_t c = 1 << 22; c >= kPageTokens * kSourcesPerRank; c -= kPageTokens * kSourcesPerRank)
+        {
+            if (smemBytes(c, worldSize) + kStaticSmemBytes <= static_cast<size_t>(maxSmemOptin))
+            {
+                cap = c;
+                break;
+            }
+        }
+        cached[worldSize] = cap;
+    }
+    return cached[worldSize];
+}
+
+int32_t cpTopkMaxActiveClusters(int32_t worldSize)
+{
+    // Queried once per world size and device (the value + 1; 0 = not queried yet).
+    constexpr int32_t kMaxDevices = 64;
+    static int32_t cached[kMaxDevices][kMaxRanks + 1] = {};
+    if (worldSize < 1 || worldSize > kMaxRanks)
+    {
+        return 0;
+    }
+    int device = 0;
+    cudaGetDevice(&device);
+    if (device >= 0 && device < kMaxDevices && cached[device][worldSize] > 0)
+    {
+        return cached[device][worldSize] - 1;
+    }
+    int32_t const capacity = cpTopkMaxLocalCapacity(worldSize);
+    if (capacity == 0)
+    {
+        return 0;
+    }
+    size_t const smem = smemBytes(capacity, worldSize);
+    if (cudaFuncSetAttribute(cpTopkZ11Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem))
+        != cudaSuccess)
+    {
+        (void) cudaGetLastError();
+        return 0;
+    }
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSourcesPerRank, kMaxRows);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = smem;
+    cudaLaunchAttribute attr{};
+    attr.id = cudaLaunchAttributeClusterDimension;
+    attr.val.clusterDim.x = kSourcesPerRank;
+    attr.val.clusterDim.y = 1;
+    attr.val.clusterDim.z = 1;
+    config.attrs = &attr;
+    config.numAttrs = 1;
+    int clusters = 0;
+    if (cudaOccupancyMaxActiveClusters(&clusters, cpTopkZ11Kernel, &config) != cudaSuccess)
+    {
+        (void) cudaGetLastError();
+        return 0;
+    }
+    if (device >= 0 && device < kMaxDevices)
+    {
+        cached[device][worldSize] = clusters + 1;
+    }
+    return clusters;
+}
+
+int64_t cpTopkDefaultTimeoutMs()
+{
+    static int64_t const timeoutMs = []()
+    {
+        char const* env = std::getenv("TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS");
+        if (env == nullptr || env[0] == '\0')
+        {
+            return kDefaultTimeoutMs;
+        }
+        char* end = nullptr;
+        long long const v = std::strtoll(env, &end, 10);
+        if (end == env || *end != '\0' || v <= 0)
+        {
+            throw std::invalid_argument(
+                std::string(
+                    "glm_moe_dsa_cross_rank_topk_global_emit: bad TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS=")
+                + env);
+        }
+        return static_cast<int64_t>(v);
+    }();
+    return timeoutMs;
+}
+
+void invokeCpTopk(CpTopkParams const& inParams, cudaStream_t stream)
+{
+    if (inParams.numRows <= 0)
+    {
+        return;
+    }
+    CpTopkParams params = inParams;
+    if (params.timeoutNs == 0)
+    {
+        params.timeoutNs = static_cast<uint64_t>(cpTopkDefaultTimeoutMs()) * 1000000ull;
+    }
+    if (params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks || params.rank < 0
+        || params.rank >= params.worldSize)
+    {
+        throw std::invalid_argument("glm_moe_dsa_cross_rank_topk_global_emit: unsupported rows/world/rank");
+    }
+    CpTopkBridgeParams const& b = params.bridge;
+    if (b.out == nullptr || b.reqIdx == nullptr || b.blockTable == nullptr || b.groupSize < 1
+        || b.groupSize > kMaxGroupSize || b.maxBlocksPerReq < 0 || (b.copyMode != 0 && b.copyMode != 1)
+        || b.copyNextN < 1 || b.copyRows < 0
+        || (b.copyDst != nullptr
+            && (b.copyMode == 0 ? b.copyRows > params.numRows
+                                : static_cast<int64_t>(b.copyRows) * b.copyNextN > params.numRows)))
+    {
+        throw std::invalid_argument("glm_moe_dsa_cross_rank_topk_global_emit: bad bridge parameters");
+    }
+    // The cluster merge stages two tokens per polled candidate vector below kBaseOffsetWords and their pool bases
+    // beyond it.
+    int32_t const pollThreads = params.worldSize * (roundCap(params.worldSize) / 2);
+    if (2 * pollThreads > kBaseOffsetWords || kBaseOffsetWords + 2 * pollThreads > mergeCapacity(params.worldSize))
+    {
+        throw std::invalid_argument(
+            "glm_moe_dsa_cross_rank_topk_global_emit: the candidate staging exceeds the merge area");
+    }
+    size_t const smem = smemBytes(params.localCapacity, params.worldSize);
+    // The attribute is per function and device; raise it only when a larger footprint is requested.
+    static thread_local size_t configuredSmem[64] = {};
+    int device = 0;
+    cudaGetDevice(&device);
+    cudaError_t err = cudaSuccess;
+    if (device >= 64 || configuredSmem[device] < smem)
+    {
+        err = cudaFuncSetAttribute(
+            cpTopkZ11Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem));
+        if (err != cudaSuccess)
+        {
+            throw std::runtime_error(
+                std::string("glm_moe_dsa_cross_rank_topk_global_emit: smem attribute: ") + cudaGetErrorString(err));
+        }
+        if (device < 64)
+        {
+            configuredSmem[device] = smem;
+        }
+    }
+    // Programmatic dependent launch as for the other TensorRT-LLM kernels (on unless TRTLLM_ENABLE_PDL=0).
+    static bool const enablePdl = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && env[0] == '0' && env[1] == '\0');
+    }();
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSourcesPerRank, params.numRows);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[2];
+    attrs[0].id = cudaLaunchAttributeClusterDimension;
+    attrs[0].val.clusterDim.x = kSourcesPerRank;
+    attrs[0].val.clusterDim.y = 1;
+    attrs[0].val.clusterDim.z = 1;
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    err = cudaLaunchKernelEx(&config, cpTopkZ11Kernel, params);
+    if (err == cudaSuccess)
+    {
+        err = cudaGetLastError();
+    }
+    if (err != cudaSuccess)
+    {
+        throw std::runtime_error(
+            std::string("glm_moe_dsa_cross_rank_topk_global_emit: launch: ") + cudaGetErrorString(err));
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h
new file mode 100644
index 000000000000..c51b6fff00fe
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h
@@ -0,0 +1,252 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
+#include <cuda_runtime_api.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit
+{
+
+// Context-partitioned DSA indexer top-k with a cluster merge that also writes the top-k bridge's outputs.
+//
+// It is the cluster-merge kernel (glmMoeDsa/crossRankTopkClusterMerge.h; the selection, exchange, workspace, epoch,
+// reset and timeout below are that kernel's, unchanged) whose emission also writes, at every output position, what the
+// top-k index bridge (glmMoeDsa/ropeKvAppend.h, topkBridgeKernel) writes from the selected row:
+//   - the global KV-pool rows of every layer of the remap group: pool[g][row][pos] = remapOne(tok, layer g), i.e.
+//     -1 for tok < 0 or a page beyond the request's block table, else
+//     blockTable[req, tok / 64] * strideFactor + tok % 64 + layerId[g] * 64 (int32 arithmetic, as the bridge);
+//   - the native top-k copy row of this row, if any: copy mode 0 copies rows [0, copyRows) to copy rows
+//     copyRowStart + row; copy mode 1 copies, for request i < copyRows, its last accepted row
+//     i * copyNextN + clamp(numAccepted[i] - 1, 0, copyNextN - 1) (i * copyNextN + copyNextN - 1 without
+//     numAccepted) to copy row copyRowStart + i.
+// In the cluster merge, the pool bases come off the critical path: right after the candidate poll, every poll thread
+// starts 4-byte cp.async copies of its (at most two) candidates' block-table entries blockTable[req, tok / 64] into
+// its own shared-memory words, which complete while the cluster selects the threshold; the emission then only adds
+// and stores. The emission is a called function that one thread of the CTA first runs into shared scratch while it
+// waits for the poll, so its instructions are cached on the SM when the emission runs. The row's request index, layer
+// ids and accepted count are copied into shared memory by one reset warp right after the grid dependency wait.
+// Identity rows and the CTA-0 second round read the block table directly. Every bridge input (request indices, block
+// table, layer ids, accepted counts) is read after the grid dependency wait, like every other global access of the
+// kernel.
+//
+// The merge selects exactly what the kernel of glmMoeDsa/crossRankTopk.h selects, in the same output order, with the
+// same context partition, exchange sources, round-1 / round-2 protocol and Lamport workspace layout:
+//   - global 64-token page p belongs to rank p % worldSize; within a rank, each verify row is handled by
+//     kSourcesPerRank CTAs, CTA s owning the 8-token block (s + page) % kSourcesPerRank of every local page
+//     (exchange source q = rank * kSourcesPerRank + s);
+//   - each source selects its local top keys with crossRankTopk's rule and broadcasts them as (order-preserving key,
+//     global token) pairs in source order, plus a header (count, exclusive bound of the withheld keys; the cluster
+//     merge adds the key statistics of the keys sent), through the multicast Lamport workspace.
+// What differs is how the row is merged and how a launch ends:
+//   - the kSourcesPerRank CTAs of one row form a thread-block cluster. CTA s merges the candidates of sources
+//     [s * worldSize, (s + 1) * worldSize) (at 8 ranks: the 8 sources of rank s). The cluster runs crossRankTopk's
+//     merge select (the same digits and early-stop rules) on the cluster-summed histograms, and every CTA writes its
+//     selected candidates directly into the output row at crossRankTopk's positions (keys above the threshold in
+//     candidate order, then ties in candidate order, then -1). Each round-1 slot header carries its source's key
+//     statistics, so every CTA derives the row's statistics from the headers it polls; the CTAs exchange only the
+//     select's histogram totals, with st.async stores into each other's shared memory completed on transaction
+//     barriers, and derive every CTA's output counts from the per-CTA tallies of the select. The only cluster-wide
+//     barrier orders the barriers' initialization before the first exchange. A row in which a source withheld keys
+//     above the threshold (identical decision on every rank) is merged instead by CTA 0 with crossRankTopk's single-CTA
+//     merge, including its second round;
+//   - every CTA keeps its own epoch word (one per (row, source CTA) slot): its thread 0 reads the word after the grid
+//     dependency wait (only the previous exchange launch of this op writes it), the value is published to the CTA at
+//     the block barrier that follows the key loads, and the CTA advances the word with a plain store after that
+//     barrier (nothing else in the launch reads it; the next launch reads it after its own grid dependency wait); the
+//     CTAs of the last row also advance the words of the rows the launch does not have, so every word counts the
+//     exchange launches. There is no arrival counter and no end-of-kernel fence.
+//
+// Rows of at most kTopK valid tokens produce the identity selection followed by -1, without an exchange. A launch whose
+// rows all have at most kTopK valid tokens (the same on every rank: the lengths come from the shared KV length) is an
+// identity launch: it writes neither the workspace nor the epoch words. Every other launch is an exchange launch.
+// Exchange launches rotate kNumLamportBuffers buffers by the epoch, which counts exchange launches; in exchange launch
+// x every CTA resets its share of its own rank's copy of the buffer of exchange launch x - 1 between its send and its
+// merge (the headers of its share are fetched once the epoch is known, after the grid dependency wait). Every rank's
+// merges of x - 1 have read their copy of it, and a peer writes it again only in exchange launch x + 2, which starts
+// after the peer's merges of x + 1 received this rank's sources of x + 1, i.e. after this rank completed launch x
+// (identity launches in between neither write nor wait, so they do not shorten this chain). Every global memory
+// access of the kernel follows its grid dependency wait, which orders it after the completion (and memory
+// visibility) of the preceding kernels of the stream, including the previous launch of this op; the kernel never
+// triggers its dependents early.
+//
+// A merge that waits longer than the timeout (CpTopkParams::timeoutNs, %globaltimer) for a peer's words records the
+// wait in the state's error word and aborts the kernel with a trap: the exchange cannot complete correctly any more,
+// and a late peer write could land in a buffer this rank reuses, so the run fails instead of continuing.
+
+//! Selected indices per verify row (DSA index_topk).
+constexpr int32_t kTopK = 2048;
+//! Ranks of the context-parallel group.
+constexpr int32_t kMaxRanks = 8;
+//! Verify rows per call.
+constexpr int32_t kMaxRows = 8;
+//! Tokens per indexer K page; the context is partitioned by whole pages.
+constexpr int32_t kPageTokens = 64;
+//! CTAs (exchange sources) per rank and verify row; also the cluster size (one cluster per row).
+constexpr int32_t kSourcesPerRank = 8;
+//! Exchange sources per row.
+constexpr int32_t kMaxSources = kMaxRanks * kSourcesPerRank;
+//! Round-1 slots per (source, row).
+constexpr int32_t kSourceCapacity = 256;
+//! Keys each saturated source aims to send in round 1 at 8 ranks (crossRankTopk's rule: three times the even share of
+//! kTopK; refining stops at the first bin edge with at most four times the even share of keys at or above it).
+constexpr int32_t kSourceBudget = 96;
+//! Round-2 slots per (source, row).
+constexpr int32_t kRound2Capacity = 2048;
+//! Threads per CTA.
+constexpr int32_t kThreads = 1024;
+//! Lamport buffers rotated by the epoch.
+constexpr int32_t kNumLamportBuffers = 3;
+//! 32-bit word that marks an empty Lamport slot; never a valid key, token index, count or bound.
+constexpr uint32_t kEmptyWord = 0xFFFFFFFFu;
+
+//! 32-bit words of one (source, row) slot: a 32-byte header then two words per candidate, so candidates can be written
+//! and read with 16-byte vector accesses. Round-1 header: count, exclusive bound of the withheld keys, then the key
+//! statistics of the keys sent (largest key, smallest key + 1, keys >= +0.0, largest key below +0.0, smallest key
+//! >= +0.0 plus 1; the "+ 1" maps the "none" value 0xFFFFFFFF to 0), then 0. Round-2 header: count, bound.
+constexpr int64_t kSlotHeaderWords = 8;
+constexpr int64_t kRound1SlotWords = kSlotHeaderWords + 2 * int64_t{kSourceCapacity};
+constexpr int64_t kRound2SlotWords = kSlotHeaderWords + 2 * int64_t{kRound2Capacity};
+//! 32-bit words of one Lamport buffer (round-1 region, then round-2 region).
+constexpr int64_t kBufferWords = int64_t{kMaxSources} * kMaxRows * (kRound1SlotWords + kRound2SlotWords);
+//! 32-bit words of the whole exchange workspace; every word must start as kEmptyWord.
+constexpr int64_t kWorkspaceWords = kNumLamportBuffers * kBufferWords;
+//! int32 words of the per-rank device state, zero-initialized: [0, 2) unused, [2] second rounds taken, [3] error word
+//! (nonzero = the wait that timed out, written right before the kernel aborts), then one epoch word per
+//! (row, source CTA) slot: state[kEpochBase + row * kSourcesPerRank + s].
+constexpr int32_t kEpochBase = 4;
+constexpr int32_t kStateWords = kEpochBase + kMaxRows * kSourcesPerRank;
+//! Default wait for a peer's Lamport words before the kernel aborts
+//! (TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS overrides it).
+constexpr int64_t kDefaultTimeoutMs = 120000;
+
+//! Values of the error word.
+enum CpTopkWait : int32_t
+{
+    kWaitRound1Header = 1,
+    kWaitRound1Candidates = 2,
+    kWaitRound2Header = 3,
+    kWaitRound2Candidates = 4
+};
+
+//! Most layers of one remap group (the bridge's kMaxGroupSize).
+constexpr int32_t kMaxGroupSize = 8;
+//! Words of the merge area (the token staging of the cluster merge) its candidates' tokens use: at most 520 16-byte
+//! candidate vectors per CTA (world sizes 5 to 7), two tokens each. The candidates' pool bases follow them, at the
+//! same per-thread offsets.
+constexpr int32_t kBaseOffsetWords = 1040;
+
+//! Outputs of the top-k bridge written at the emission (see the comment at the top).
+struct CpTopkBridgeParams
+{
+    //! [groupSize, numRows, kTopK] int32 pool rows with strides (outStrideG, outStride0, 1).
+    int32_t* out = nullptr;
+    int64_t outStrideG = 0;
+    int64_t outStride0 = 0;
+    int32_t groupSize = 1;
+    //! Optional [groupSize] int32 layer ids; when null every member uses layerId.
+    int32_t const* layerIds = nullptr;
+    int32_t layerId = 0;
+    //! [numRows] int32 request index per row, before subtracting reqIdxBias.
+    int32_t const* reqIdx = nullptr;
+    int32_t reqIdxBias = 0;
+    //! int32 block table [requests, maxBlocksPerReq] with strides (btStride0, btStride1), block size kPageTokens.
+    int32_t const* blockTable = nullptr;
+    int64_t btStride0 = 0;
+    int64_t btStride1 = 1;
+    int32_t maxBlocksPerReq = 0;
+    int32_t strideFactor = 0;
+    //! Optional copy destination [>= copyRowStart + copyRows, >= kTopK] int32 with row stride copyStride0.
+    int32_t* copyDst = nullptr;
+    int64_t copyStride0 = 0;
+    int32_t copyRows = 0;
+    int32_t copyRowStart = 0;
+    int32_t copyMode = 0;
+    int32_t copyNextN = 1;
+    //! Optional [copyRows] accepted counts of copy mode 1, int32 or int64.
+    void const* numAccepted = nullptr;
+    bool numAcceptedIs64 = false;
+};
+
+struct CpTopkParams
+{
+    //! [numRows, logitsStride] fp32 logits of this rank's local columns.
+    float const* logits;
+    int64_t logitsStride;
+    //! Local columns the logits view holds (each row's local length is clamped to it).
+    int32_t localCapacity;
+    //! [1] int32 global indexer KV length of the request; row r has kvLen - nextN + r + 1 valid tokens.
+    int32_t const* kvLens;
+    //! [numRows, kTopK] int32 output: global token indices, -1 for unused slots.
+    int32_t* out;
+    int32_t numRows;
+    int32_t nextN;
+    int32_t rank;
+    int32_t worldSize;
+    //! Multicast address of the exchange workspace (kWorkspaceWords words, identical layout on every rank).
+    uint32_t* mcWorkspace;
+    //! This rank's unicast address of the same workspace.
+    uint32_t* ucWorkspace;
+    //! [kStateWords] int32 per-rank device state, zero-initialized.
+    int32_t* state;
+    //! Optional [numRows, kSourcesPerRank, kDebugSlots] int64 phase timestamps (diagnosis only; nullptr = off).
+    int64_t* debugTimes = nullptr;
+    //! Longest wait for a peer's words before the kernel aborts (ns); 0 = the default (kDefaultTimeoutMs, or
+    //! TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS).
+    uint64_t timeoutNs = 0;
+    //! The bridge outputs (out, reqIdx and blockTable are required).
+    CpTopkBridgeParams bridge;
+};
+
+//! Timestamps per CTA in CpTopkParams::debugTimes: [0, 16) %globaltimer (ns) and [16, 32) clock64 per phase (0 start,
+//! 10 grid dependency wait returned, 1 keys loaded, 2 local select, 3 sent, 9 previous buffer reset (%globaltimer of
+//! the first reset warp only), 4 candidates polled, 5 cluster statistics agreed, 6 merge select, 7 row written, 8 end),
+//! [32, 42) local select (keys, then clock64 at the start of / after the histogram of / after the bin search of each
+//! pass), [42, 45) send (start, keys taken, stores issued), 46 round-1 keys sent by this source, 47 local select passes
+//! (+10 when the cut was searched among the negative keys), [48, 58) merge select (keys = this CTA's candidates; per
+//! pass clock64 at its start / after the group totals arrived / after the bin was decided), clock64 of the merge: 45
+//! candidates polled, 61 statistics combined, 58 emission positions scanned, 60 row written; 62 candidates of the row
+//! (all CTAs), 63 merge select passes (+10 likewise; +100 when the row took crossRankTopk's single-CTA merge, which
+//! then reuses 45 and 58-61 as crossRankTopk does); clock64 and values of thread 0.
+constexpr int32_t kDebugSlots = 64;
+
+//! Effective wait timeout of a launch with params.timeoutNs == 0 (ms).
+int64_t cpTopkDefaultTimeoutMs();
+
+//! Dynamic shared memory the kernel needs (bytes).
+size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize);
+
+//! Largest local capacity the kernel accepts for a world size (shared memory; the cluster merge's limit).
+int32_t cpTopkMaxLocalCapacity(int32_t worldSize);
+
+//! Clusters of the kernel's launch configuration (largest local capacity) that can be resident at once on this
+//! device (cudaOccupancyMaxActiveClusters, queried once per world size and device); 0 when the query fails. No row
+//! waits for another row of its rank, but a row's cluster waits for the same row's sources on every peer rank, and
+//! the ranks may make their rows resident in different orders; the caller therefore requires kMaxRows co-resident
+//! clusters, so every cluster of a launch can be resident at once (other kernels sharing the SMs are finite and do
+//! not wait for this one).
+int32_t cpTopkMaxActiveClusters(int32_t worldSize);
+
+void invokeCpTopk(CpTopkParams const& params, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu
new file mode 100644
index 000000000000..d2fe7ad69100
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu
@@ -0,0 +1,453 @@
+/*
+ * 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.
+ */
+
+// Weight-streaming BF16 kernels for the dense (non-MoE) MLP at decode row counts (M <= 8).
+//
+// Both projections are bound by reading the BF16 weight matrix once from HBM. Each CTA owns a
+// 16-row weight tile and its warps split the contraction dimension. Every thread loads 16 bytes of
+// two weight rows and of one activation row per step and feeds them to two m16n8k16 tensor-core
+// MMAs with the weight on the MMA M side and the (at most 8) activation rows on the MMA N side.
+// The contraction index assigned to an MMA k-slot depends only on the thread's position inside its
+// quad, so weight and activation fragments use the same permutation of k and the sum is exact up to
+// FP32 summation order. Each quad therefore reads 64 contiguous bytes of one weight row per step.
+//
+// The gate/up kernel places 8 gate rows and the matching 8 up rows in one tile, so each thread ends
+// with gate and up of the same output column and applies SwiGLU in registers.
+//
+// Bytes in flight: a thread holds one chunk of steps in registers, and one lane per quad prefetches
+// the quad's row segment of the next chunk into L2 with cp.async.bulk.prefetch, which keeps the HBM
+// queue full without spending registers.
+//
+// Programmatic dependent launch: weights never depend on the preceding kernel, so the first chunk
+// (registers) and the L2 prefetch of the next one are issued before griddepcontrol.wait; dependents
+// are released right after the wait. The tilings keep each grid within one wave of resident CTAs,
+// so early-launched dependent CTAs that wait in griddepcontrol.wait cannot starve the primary grid.
+
+#include "tensorrt_llm/kernels/glmMoeDsa/denseMlp.h"
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::dense_mlp
+{
+
+namespace
+{
+
+using Bf16 = __nv_bfloat16;
+
+constexpr int kWarpSize = 32;
+//! Weight rows per CTA tile (MMA M).
+constexpr int kTileRows = 16;
+//! Output columns per gate/up CTA: 8 gate rows + 8 up rows form one tile.
+constexpr int kGateUpCols = kTileRows / 2;
+//! BF16 values per 16-byte vector load.
+constexpr int kVecElems = 8;
+//! Threads per MMA quad; a quad covers kQuadThreads * kVecElems contraction elements per step.
+constexpr int kQuadThreads = 4;
+//! Contraction elements consumed by one warp step (two m16n8k16 MMAs).
+constexpr int kKPerStep = kQuadThreads * kVecElems;
+
+//! Compile-time tiling of one kernel instance.
+//! \tparam kWarps warps per CTA; each warp owns a contiguous 1/kWarps slice of the contraction.
+//! \tparam kSteps steps per register chunk.
+//! \tparam kMinBlocks minimum resident CTAs per SM requested from the compiler (register budget).
+//! \tparam kCachedX load activations with ld.global.ca instead of the non-coherent path.
+//! \tparam kL2Ahead chunks of each weight row segment prefetched into L2 ahead of the register loads.
+template <int kWarps_, int kSteps_, int kMinBlocks_, bool kCachedX_, int kL2Ahead_>
+struct Tiling
+{
+    static constexpr int kWarps = kWarps_;
+    static constexpr int kSteps = kSteps_;
+    static constexpr int kMinBlocks = kMinBlocks_;
+    static constexpr bool kCachedX = kCachedX_;
+    static constexpr int kL2Ahead = kL2Ahead_;
+    static constexpr int kKPerChunk = kKPerStep * kSteps;
+};
+
+__device__ __forceinline__ void pdlWait()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+#endif
+}
+
+__device__ __forceinline__ void pdlLaunchDependents()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+#endif
+}
+
+//! Streaming weight load: not kept in L1, 256-byte L2 fetch granularity.
+__device__ __forceinline__ uint4 loadWeight(Bf16 const* ptr)
+{
+    uint4 r;
+    asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w)
+                 : "l"(ptr));
+    return r;
+}
+
+//! Activation load: every CTA reads the same few rows, so keep them in L1.
+template <bool kCached>
+__device__ __forceinline__ uint4 loadActivation(Bf16 const* ptr)
+{
+    if constexpr (kCached)
+    {
+        uint4 r;
+        asm volatile("ld.global.ca.v4.u32 {%0, %1, %2, %3}, [%4];"
+                     : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w)
+                     : "l"(ptr));
+        return r;
+    }
+    else
+    {
+        return __ldg(reinterpret_cast<uint4 const*>(ptr));
+    }
+}
+
+__device__ __forceinline__ void prefetchL2(void const* ptr, uint32_t bytes)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(ptr), "r"(bytes) : "memory");
+#endif
+}
+
+__device__ __forceinline__ void mmaBf16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+//! Registers of one chunk: kSteps steps of two weight rows and one activation row.
+template <int kSteps>
+struct Chunk
+{
+    uint4 w0[kSteps];
+    uint4 w1[kSteps];
+    uint4 x[kSteps];
+};
+
+template <int kSteps>
+__device__ __forceinline__ void loadChunkWeights(Chunk<kSteps>& chunk, Bf16 const* w0, Bf16 const* w1, int chunkIdx)
+{
+    int const offset = chunkIdx * kKPerStep * kSteps;
+#pragma unroll
+    for (int u = 0; u < kSteps; ++u)
+    {
+        chunk.w0[u] = loadWeight(w0 + offset + u * kKPerStep);
+        chunk.w1[u] = loadWeight(w1 + offset + u * kKPerStep);
+    }
+}
+
+template <bool kCached, int kSteps>
+__device__ __forceinline__ void loadChunkActivations(Chunk<kSteps>& chunk, Bf16 const* x, bool hasRow, int chunkIdx)
+{
+    int const offset = chunkIdx * kKPerStep * kSteps;
+#pragma unroll
+    for (int u = 0; u < kSteps; ++u)
+    {
+        chunk.x[u] = hasRow ? loadActivation<kCached>(x + offset + u * kKPerStep) : make_uint4(0U, 0U, 0U, 0U);
+    }
+}
+
+//! Two MMAs per step. MMA A uses contraction elements 0..3 of the thread's 8-element vector, MMA B
+//! elements 4..7. Fragment registers: a0/a2 = weight row "group", a1/a3 = weight row "group + 8",
+//! b0/b1 = activation row "group".
+template <int kSteps>
+__device__ __forceinline__ void computeChunk(float (&accA)[4], float (&accB)[4], Chunk<kSteps> const& chunk)
+{
+#pragma unroll
+    for (int u = 0; u < kSteps; ++u)
+    {
+        mmaBf16(accA, chunk.w0[u].x, chunk.w1[u].x, chunk.w0[u].y, chunk.w1[u].y, chunk.x[u].x, chunk.x[u].y);
+        mmaBf16(accB, chunk.w0[u].z, chunk.w1[u].z, chunk.w0[u].w, chunk.w1[u].w, chunk.x[u].z, chunk.x[u].w);
+    }
+}
+
+//! Streams the CTA's two weight rows per thread over the warp's contraction slice and returns the
+//! CTA-reduced accumulator fragment in warp 0 (other warps return false).
+//! acc[0..1] = tile row "group", activation rows 2 * tig + {0, 1}; acc[2..3] = tile row "group + 8".
+template <typename T>
+__device__ __forceinline__ bool streamTile(float (&acc)[4], Bf16 const* __restrict__ x,
+    Bf16 const* __restrict__ weight, int numRows, int kDim, int row0, int row1)
+{
+    constexpr int kWarps = T::kWarps;
+    constexpr int kSteps = T::kSteps;
+    __shared__ float reduceBuf[kWarps][4][kWarpSize];
+
+    int const lane = threadIdx.x % kWarpSize;
+    int const warp = threadIdx.x / kWarpSize;
+    int const group = lane / kQuadThreads;
+    int const tig = lane % kQuadThreads;
+
+    int const kPerWarp = kDim / kWarps;
+    int const kBase = warp * kPerWarp + tig * kVecElems;
+    Bf16 const* w0 = weight + static_cast<int64_t>(row0) * kDim + kBase;
+    Bf16 const* w1 = weight + static_cast<int64_t>(row1) * kDim + kBase;
+    bool const hasRow = group < numRows;
+    Bf16 const* xRow = x + static_cast<int64_t>(hasRow ? group : 0) * kDim + kBase;
+    int const numChunks = kPerWarp / T::kKPerChunk;
+
+    // One lane per quad prefetches the quad's contiguous row segments of a chunk into L2.
+    auto prefetchChunkL2 = [&](int chunkIdx)
+    {
+        if constexpr (T::kL2Ahead > 0)
+        {
+            if (tig == 0 && chunkIdx < numChunks)
+            {
+                uint32_t const bytes = T::kKPerChunk * static_cast<uint32_t>(sizeof(Bf16));
+                prefetchL2(w0 + chunkIdx * T::kKPerChunk, bytes);
+                prefetchL2(w1 + chunkIdx * T::kKPerChunk, bytes);
+            }
+        }
+    };
+
+    float accA[4] = {0.F, 0.F, 0.F, 0.F};
+    float accB[4] = {0.F, 0.F, 0.F, 0.F};
+
+    // Weights do not depend on the preceding kernel: request them before the dependency wait.
+    Chunk<kSteps> chunk;
+    loadChunkWeights(chunk, w0, w1, 0);
+#pragma unroll
+    for (int a = 1; a <= T::kL2Ahead; ++a)
+    {
+        prefetchChunkL2(a);
+    }
+    pdlWait();
+    pdlLaunchDependents();
+    loadChunkActivations<T::kCachedX>(chunk, xRow, hasRow, 0);
+
+#pragma unroll 1
+    for (int c = 0; c < numChunks; ++c)
+    {
+        if (c > 0)
+        {
+            loadChunkWeights(chunk, w0, w1, c);
+            loadChunkActivations<T::kCachedX>(chunk, xRow, hasRow, c);
+        }
+        prefetchChunkL2(c + 1 + T::kL2Ahead);
+        computeChunk(accA, accB, chunk);
+    }
+
+#pragma unroll
+    for (int e = 0; e < 4; ++e)
+    {
+        acc[e] = accA[e] + accB[e];
+    }
+    if constexpr (kWarps > 1)
+    {
+#pragma unroll
+        for (int e = 0; e < 4; ++e)
+        {
+            reduceBuf[warp][e][lane] = acc[e];
+        }
+        __syncthreads();
+        if (warp != 0)
+        {
+            return false;
+        }
+        // Fixed order over warps keeps the result deterministic.
+#pragma unroll
+        for (int w = 1; w < kWarps; ++w)
+        {
+#pragma unroll
+            for (int e = 0; e < 4; ++e)
+            {
+                acc[e] += reduceBuf[w][e][lane];
+            }
+        }
+    }
+    return true;
+}
+
+__device__ __forceinline__ float roundToBf16(float v)
+{
+    return __bfloat162float(__float2bfloat16_rn(v));
+}
+
+template <typename T>
+__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) gateUpSwigluKernel(
+    Bf16 const* __restrict__ x, Bf16 const* __restrict__ weight, Bf16* __restrict__ h, int numRows, int hidden, int inter)
+{
+    int const lane = threadIdx.x % kWarpSize;
+    int const group = lane / kQuadThreads;
+    int const tig = lane % kQuadThreads;
+    int const col = blockIdx.x * kGateUpCols + group;
+
+    float acc[4];
+    if (!streamTile<T>(acc, x, weight, numRows, hidden, col, inter + col))
+    {
+        return;
+    }
+#pragma unroll
+    for (int b = 0; b < 2; ++b)
+    {
+        int const row = 2 * tig + b;
+        if (row < numRows)
+        {
+            // gate and up are rounded to BF16 first, matching a BF16 GEMM output feeding SwiGLU.
+            float const gate = roundToBf16(acc[b]);
+            float const up = roundToBf16(acc[2 + b]);
+            float const sig = 1.F / (1.F + __expf(-gate));
+            h[static_cast<int64_t>(row) * inter + col] = __float2bfloat16_rn(gate * sig * up);
+        }
+    }
+}
+
+template <typename T>
+__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) downKernel(
+    Bf16 const* __restrict__ h, Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int inter, int hidden)
+{
+    int const lane = threadIdx.x % kWarpSize;
+    int const group = lane / kQuadThreads;
+    int const tig = lane % kQuadThreads;
+    int const n0 = blockIdx.x * kTileRows + group;
+
+    float acc[4];
+    if (!streamTile<T>(acc, h, weight, numRows, inter, n0, n0 + kTileRows / 2))
+    {
+        return;
+    }
+#pragma unroll
+    for (int b = 0; b < 2; ++b)
+    {
+        int const row = 2 * tig + b;
+        if (row < numRows)
+        {
+            out[static_cast<int64_t>(row) * hidden + n0] = __float2bfloat16_rn(acc[b]);
+            out[static_cast<int64_t>(row) * hidden + n0 + kTileRows / 2] = __float2bfloat16_rn(acc[2 + b]);
+        }
+    }
+}
+
+template <typename T>
+cudaError_t launchKernel(void (*kernel)(Bf16 const*, Bf16 const*, Bf16*, int, int, int), int gridSize, Bf16 const* in,
+    Bf16 const* weight, Bf16* out, int numRows, int kDim, int outDim, bool enablePdl, cudaStream_t stream)
+{
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(gridSize);
+    config.blockDim = dim3(T::kWarps * kWarpSize);
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    return cudaLaunchKernelEx(&config, kernel, in, weight, out, numRows, kDim, outDim);
+}
+
+//! Variant table (warps, steps); the template arguments of each index are in launchVariant.
+struct VariantDesc
+{
+    int warps;
+    int steps;
+};
+
+constexpr VariantDesc kVariants[] = {
+    {4, 8}, // 0: down default. L2 prefetch one chunk ahead
+    {8, 4}, // 1: gate/up default. Three CTAs/SM register budget, L2 prefetch one chunk ahead
+    {8, 4}, // 2: as 1 with ld.global.ca activation loads
+    {4, 8}, // 3: as 0 without L2 prefetch
+    {8, 8}, // 4: two CTAs/SM, no L2 prefetch
+};
+constexpr int kNumVariants = sizeof(kVariants) / sizeof(kVariants[0]);
+
+bool isContractionSupported(int kDim, int variant)
+{
+    VariantDesc const& v = kVariants[variant];
+    return kDim > 0 && kDim % (v.warps * kKPerStep * v.steps) == 0;
+}
+
+//! Dispatches a variant index to its kernel instance; kGateUp selects the gate/up or the down kernel.
+template <bool kGateUp>
+cudaError_t launchVariant(int variant, Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int kDim,
+    int outDim, bool enablePdl, cudaStream_t stream)
+{
+    auto launch = [&](auto tiling) -> cudaError_t
+    {
+        using T = decltype(tiling);
+        if constexpr (kGateUp)
+        {
+            return launchKernel<T>(gateUpSwigluKernel<T>, outDim / kGateUpCols, in, weight, out, numRows, kDim, outDim,
+                enablePdl, stream);
+        }
+        else
+        {
+            return launchKernel<T>(
+                downKernel<T>, outDim / kTileRows, in, weight, out, numRows, kDim, outDim, enablePdl, stream);
+        }
+    };
+    switch (variant)
+    {
+    case 0: return launch(Tiling<4, 8, 1, false, 1>{});
+    case 1: return launch(Tiling<8, 4, 3, false, 1>{});
+    case 2: return launch(Tiling<8, 4, 3, true, 1>{});
+    case 3: return launch(Tiling<4, 8, 1, false, 0>{});
+    case 4: return launch(Tiling<8, 8, 2, false, 0>{});
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace
+
+int getNumVariants()
+{
+    return kNumVariants;
+}
+
+bool isGateUpSupported(int numRows, int hidden, int inter, int variant)
+{
+    return variant >= 0 && variant < kNumVariants && numRows >= 1 && numRows <= kMaxRows && inter > 0
+        && inter % kGateUpCols == 0 && isContractionSupported(hidden, variant);
+}
+
+bool isDownSupported(int numRows, int inter, int hidden, int variant)
+{
+    return variant >= 0 && variant < kNumVariants && numRows >= 1 && numRows <= kMaxRows && hidden > 0
+        && hidden % kTileRows == 0 && isContractionSupported(inter, variant);
+}
+
+cudaError_t invokeGateUpSwiglu(Bf16 const* x, Bf16 const* gateUpWeight, Bf16* h, int numRows, int hidden, int inter,
+    int variant, bool enablePdl, cudaStream_t stream)
+{
+    if (!isGateUpSupported(numRows, hidden, inter, variant))
+    {
+        return cudaErrorInvalidValue;
+    }
+    return launchVariant<true>(variant, x, gateUpWeight, h, numRows, hidden, inter, enablePdl, stream);
+}
+
+cudaError_t invokeDown(Bf16 const* h, Bf16 const* downWeight, Bf16* out, int numRows, int inter, int hidden,
+    int variant, bool enablePdl, cudaStream_t stream)
+{
+    if (!isDownSupported(numRows, inter, hidden, variant))
+    {
+        return cudaErrorInvalidValue;
+    }
+    return launchVariant<false>(variant, h, downWeight, out, numRows, inter, hidden, enablePdl, stream);
+}
+
+} // namespace kernels::glm_moe_dsa::dense_mlp
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h
new file mode 100644
index 000000000000..3d82801ffca0
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h
@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::dense_mlp
+{
+
+//! Largest number of activation rows (tokens) the small-M dense MLP kernels accept.
+constexpr int kMaxRows = 8;
+
+//! Number of tiling variants compiled for each kernel. Variant 0 is the default.
+int getNumVariants();
+
+//! Host-side eligibility checks (no device work).
+//! \param numRows activation rows M.
+//! \param hidden model hidden size (gate/up contraction size, down output size).
+//! \param inter rank-local intermediate size (gate/up output size, down contraction size).
+bool isGateUpSupported(int numRows, int hidden, int inter, int variant);
+bool isDownSupported(int numRows, int inter, int hidden, int variant);
+
+//! Gate/up projection fused with SwiGLU.
+//! For every row m < numRows and column j < inter:
+//!   gate = bf16(sum_k x[m, k] * W[j, k]), up = bf16(sum_k x[m, k] * W[inter + j, k]),
+//!   h[m, j] = bf16((gate * sigmoid(gate)) * up), all sums in FP32.
+//! \param x activations [numRows, hidden], row-major, BF16.
+//! \param gateUpWeight [2 * inter, hidden], row-major, BF16 (rows [0, inter) gate, [inter, 2 * inter) up).
+//! \param h output [numRows, inter], row-major, BF16.
+cudaError_t invokeGateUpSwiglu(__nv_bfloat16 const* x, __nv_bfloat16 const* gateUpWeight, __nv_bfloat16* h, int numRows,
+    int hidden, int inter, int variant, bool enablePdl, cudaStream_t stream);
+
+//! Down projection: out[m, n] = bf16(sum_k h[m, k] * W[n, k]) with FP32 sums.
+//! \param h activations [numRows, inter], row-major, BF16.
+//! \param downWeight [hidden, inter], row-major, BF16.
+//! \param out output [numRows, hidden], row-major, BF16.
+cudaError_t invokeDown(__nv_bfloat16 const* h, __nv_bfloat16 const* downWeight, __nv_bfloat16* out, int numRows,
+    int inter, int hidden, int variant, bool enablePdl, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::dense_mlp
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu
new file mode 100644
index 000000000000..33bfd92f7cef
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu
@@ -0,0 +1,586 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h"
+
+#include <cooperative_groups.h>
+#include <cuda_bf16.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::draft_live_rows
+{
+namespace
+{
+namespace cg = cooperative_groups;
+
+// Work decomposition (the post-attention router hop's, on the compact rows).
+//
+// Every cluster has kSlices CTAs; CTA c owns slice [c * kSlice, (c + 1) * kSlice) of the hidden dimension (96 threads
+// x 8 elements), the slice a CTA of the native one-shot kernel owns, so the RMSNorm reduction reproduces the native
+// tree. There are kG x kRowSplit clusters (kG = 4 expert groups, kRowSplit = min(R, 3) row groups, R = live rows).
+// Cluster (g, h) collects and normalizes the compact rows r with r % kRowSplit == h and computes their logits for
+// experts [64 g, 64 g + 64). Its router weights live in registers (loaded before the PDL wait). Each CTA multiplies
+// its rows of slice c by the experts' weights of slice c with BF16 MMAs (experts on m16, rows on n8) and pushes the
+// partial logits to the CTA that owns the expert; owners add the 8 slices in a fixed order, so the logits are
+// identical on every rank and do not depend on R or on the row's position. Clusters never synchronize with each
+// other; inside a cluster, slice sums and partial logits travel as st.async stores that complete transactions on the
+// consumer's mbarrier.
+constexpr int kSlices = 8;
+constexpr int kSlice = kHiddenDim / kSlices;
+constexpr int kEltsPerVec = 8;
+constexpr int kVecPerSlice = kSlice / kEltsPerVec;
+constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32;
+// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two
+// rows in one quarter-warp phase on disjoint banks.
+constexpr int kNormedStride = kVecPerSlice + 4;
+constexpr int kMmaChains = 4;
+constexpr int kExpertGroups = 4;
+constexpr uint32_t kLamportSentinel = 0x80000000U;
+
+static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0);
+static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32);
+static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension");
+
+struct HopConfig
+{
+    static constexpr int kThreads = 384;
+    static constexpr int kWarps = kThreads / 32;
+    static constexpr int kExpertsPerCluster = kNumExperts / kExpertGroups;
+    static constexpr int kExpertTiles = kExpertsPerCluster / 16;
+    static constexpr int kKGroups = kWarps / kExpertTiles;
+    static constexpr int kKPerGroup = kSlice / kKGroups;
+    static constexpr int kPairsPerGroup = kKPerGroup / 32;
+    static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices;
+    static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0);
+    static_assert(kPairsPerGroup % (kMmaChains / 2) == 0);
+};
+
+constexpr int kMaxRowsPerCtaAll = kMaxRows / 2;
+
+struct alignas(16) HopStorage
+{
+    uint4 gamma[kVecPerSlice];
+    float vecSq[kMaxRowsPerCtaAll][kVecPerSlice];
+    float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice];
+    float blockSq[kMaxRowsPerCtaAll][kSlices];
+    // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx).
+    alignas(8) uint64_t normBar;
+    // Completes when all partial logits of the experts this CTA owns have arrived.
+    alignas(8) uint64_t routerBar;
+};
+
+struct alignas(16) RouterStorage
+{
+    uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero
+    float partial[kSlices][HopConfig::kKGroups][HopConfig::kExpertsPerOwner][kMaxRows];
+};
+
+union BfVec
+{
+    uint4 packed;
+    __nv_bfloat16 elts[kEltsPerVec];
+};
+
+__device__ __forceinline__ uint4 loadVolatile(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadConstant(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr));
+    return v;
+}
+
+__device__ __forceinline__ bool isDirty(uint4 const& v)
+{
+    return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel;
+}
+
+// Lamport payloads must never contain the sentinel word: map bf16 -0.0 to +0.0 (the reduction starts from +0.0, so
+// the sum is unchanged).
+__device__ __forceinline__ uint4 sanitize(uint4 v)
+{
+    BfVec b;
+    b.packed = v;
+#pragma unroll
+    for (int i = 0; i < kEltsPerVec; ++i)
+    {
+        if (__bfloat16_as_ushort(b.elts[i]) == 0x8000U)
+        {
+            b.elts[i] = __ushort_as_bfloat16(0);
+        }
+    }
+    return b.packed;
+}
+
+__device__ __forceinline__ uint32_t sharedAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// Address of the same shared-memory object in CTA `rank` of the cluster.
+__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank)
+{
+    uint32_t out;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank));
+    return out;
+}
+
+__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory");
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar)
+{
+    asm volatile(
+        "{\n"
+        ".reg .pred done;\n"
+        "LIVE_ROW_MBAR_WAIT_%=:\n"
+        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n"
+        "@!done bra LIVE_ROW_MBAR_WAIT_%=;\n"
+        "}\n" ::"r"(sharedAddr(bar))
+        : "memory");
+}
+
+// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier.
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v),
+                 "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr),
+                 "f"(a), "f"(b), "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mmaBf16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+// R = live (compact) rows; kRowSplit = row groups.
+template <int R, int kRowSplit>
+__global__ void __launch_bounds__(HopConfig::kThreads, 1) liveRowHopKernel(LiveRowHopParams const params)
+{
+    static_assert(R >= 1 && R <= kMaxRows);
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    using Cfg = HopConfig;
+    constexpr int kClusters = kExpertGroups * kRowSplit;
+    constexpr int kGridCtas = kClusters * kSlices;
+    // Rows one CTA collects, normalizes and multiplies.
+    constexpr int kRowsPerCta = (R + kRowSplit - 1) / kRowSplit;
+    static_assert(kRowsPerCta <= kMaxRowsPerCtaAll);
+    static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps);
+    static_assert(Cfg::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads);
+
+    __shared__ HopStorage hop;
+    __shared__ RouterStorage rs;
+
+    cg::cluster_group cluster = cg::this_cluster();
+    int const slice = static_cast<int>(cluster.block_rank());
+    int const clusterId = blockIdx.y;
+    int const expertGroup = clusterId % kExpertGroups;
+    int const rowGroup = clusterId / kExpertGroups;
+    int const tid = threadIdx.x;
+    int const lane = tid & 31;
+    int const warp = tid >> 5;
+    int const sliceBase = slice * kSlice;
+
+    // Compact rows this CTA collects and normalizes: local row j is compact row rowGroup + j * kRowSplit.
+    int const numRowsCta = (R - rowGroup + kRowSplit - 1) / kRowSplit;
+    // This thread's item: one 16-byte vector of one row slice.
+    int const itemRow = tid / kVecPerSlice;
+    int const itemVec = tid % kVecPerSlice;
+    bool const hasItem = itemRow < numRowsCta;
+    int const row = rowGroup + itemRow * kRowSplit;
+    // The rank's own partial of a row slice is published by exactly one CTA of the grid.
+    bool const sends = hasItem && itemRow % kExpertGroups == expertGroup;
+    size_t const colOff = static_cast<size_t>(sliceBase) + itemVec * kEltsPerVec;
+    size_t const rowOff = static_cast<size_t>(row) * kHiddenDim + colOff;
+    // The cluster that writes the normalized rows and the residual.
+    bool const writesOutputs = expertGroup == 0;
+
+    // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ----
+    if (tid < kVecPerSlice)
+    {
+        hop.gamma[tid] = loadConstant(reinterpret_cast<uint4 const*>(params.gamma) + slice * kVecPerSlice + tid);
+    }
+    int const gid = lane >> 2;
+    int const tig = lane & 3;
+    uint4 wLo[Cfg::kPairsPerGroup];
+    uint4 wHi[Cfg::kPairsPerGroup];
+    {
+        int const tile = warp % Cfg::kExpertTiles;
+        int const kGroup = warp / Cfg::kExpertTiles;
+        auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight);
+        int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid;
+#pragma unroll
+        for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+        {
+            int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec;
+            wLo[p] = loadConstant(w + static_cast<size_t>(expert) * kHiddenDim + k);
+            wHi[p] = loadConstant(w + static_cast<size_t>(expert + 8) * kHiddenDim + k);
+        }
+        // MMA columns of rows this CTA does not own stay zero.
+        for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads)
+        {
+            rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0);
+        }
+    }
+
+    // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums
+    // of its rows and the 8 x kKGroups partial tiles of the experts it owns. A split cluster barrier publishes the
+    // initialization; its wait precedes the first remote store, long after all CTAs arrived.
+    if (tid == 0)
+    {
+        mbarrierInitExpectTx(&hop.normBar, static_cast<uint32_t>(numRowsCta * kSlices * sizeof(float)));
+        mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast<uint32_t>(sizeof(rs.partial)) : 0U);
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+    asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory");
+
+    // ---- Wait for the producer of `input` (and transitively of `residualIn`, `rowIds` and the protocol state). ----
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    // Dependents only start their own prologue; they still wait for this grid to complete.
+    asm volatile("griddepcontrol.launch_dependents;");
+
+    // The live row's source row (written by an earlier kernel of the stream: a plain load after the wait).
+    int64_t srcRow = 0;
+    if (hasItem)
+    {
+        srcRow = params.rowIds[row];
+    }
+    uint4 const flags = loadVolatile(params.flags);
+    srcRow = srcRow < 0 ? 0 : (srcRow >= params.numInputRows ? params.numInputRows - 1 : srcRow);
+    size_t const srcOff = static_cast<size_t>(srcRow) * kHiddenDim + colOff;
+    auto const* input = reinterpret_cast<__nv_bfloat16 const*>(params.input);
+    auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn);
+    uint4 res = make_uint4(0, 0, 0, 0);
+    // This rank's own partial, as published (sanitized); it is used from registers, not from the multicast copy.
+    uint4 own = make_uint4(0, 0, 0, 0);
+    if (hasItem)
+    {
+        res = *reinterpret_cast<uint4 const*>(residualIn + srcOff);
+        own = sanitize(*reinterpret_cast<uint4 const*>(input + srcOff));
+    }
+    uint32_t const cur = flags.x;
+    uint32_t const clearBytes = flags.y;
+    char* ucBase = reinterpret_cast<char*>(params.ucBuffer);
+    auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes);
+    auto* curMc
+        = reinterpret_cast<__nv_bfloat16*>(reinterpret_cast<char*>(params.mcBuffer) + cur * kLamportBufferBytes);
+    // Lamport layout of one buffer: [compact row][source rank][kHiddenDim].
+    size_t const slotBase = static_cast<size_t>(row) * kNumRanks * kHiddenDim + colOff;
+
+    // ---- Publish this rank's partial to every rank (one multicast store per vector). ----
+    if (sends)
+    {
+        *reinterpret_cast<uint4*>(curMc + slotBase + static_cast<size_t>(params.rank) * kHiddenDim) = own;
+    }
+
+    // ---- Collect the other ranks' partials. ----
+    uint4 vals[kNumRanks];
+    if (hasItem)
+    {
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            vals[s] = s == params.rank ? own : loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+        }
+    }
+
+    // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ----
+    {
+        uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers;
+        auto* dirtyBuf = reinterpret_cast<uint4*>(ucBase + dirty * kLamportBufferBytes);
+        uint32_t const numVec = clearBytes / sizeof(uint4);
+        uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel);
+        uint32_t const ctaLinear = clusterId * kSlices + slice;
+        for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads)
+        {
+            dirtyBuf[j] = sentinel;
+        }
+    }
+
+    if (hasItem)
+    {
+        bool pending = true;
+        while (pending)
+        {
+            pending = false;
+#pragma unroll
+            for (int s = 0; s < kNumRanks; ++s)
+            {
+                if (isDirty(vals[s]))
+                {
+                    pending = true;
+                    vals[s] = loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+                }
+            }
+        }
+    }
+
+    // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ----
+    uint4 resOut = make_uint4(0, 0, 0, 0);
+    if (hasItem)
+    {
+        float acc[kEltsPerVec];
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            acc[e] = 0.F;
+        }
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            BfVec x;
+            x.packed = vals[s];
+#pragma unroll
+            for (int e = 0; e < kEltsPerVec; ++e)
+            {
+                acc[e] += __bfloat162float(x.elts[e]);
+            }
+        }
+        BfVec rv;
+        rv.packed = res;
+        BfVec y;
+        float sq = 0.F;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]);
+            sq += __bfloat162float(__hmul(y.elts[e], y.elts[e]));
+        }
+        resOut = y.packed;
+        hop.vecSq[itemRow][itemVec] = sq;
+    }
+    __syncthreads();
+    if (tid == 0)
+    {
+        // Every thread of this CTA has read the protocol state; CTA 0 advances it only after all CTAs report.
+        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory");
+    }
+
+    // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ----
+    if (warp < numRowsCta * kNativeWarpsPerSlice)
+    {
+        int const j = warp / kNativeWarpsPerSlice;
+        int const q = warp % kNativeWarpsPerSlice;
+        float val = hop.vecSq[j][q * 32 + lane];
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+        {
+            val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32);
+        }
+        if (lane == 0)
+        {
+            hop.warpSq[j][q] = val;
+        }
+    }
+    __syncthreads();
+    asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
+    if (tid < numRowsCta * kSlices)
+    {
+        int const j = tid / kSlices;
+        int const dst = tid % kSlices;
+        float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1];
+        storeAsync(
+            clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst));
+    }
+    mbarrierWaitPhase0(&hop.normBar);
+
+    if (hasItem)
+    {
+        float fullSum = 0.F;
+#pragma unroll
+        for (int c = 0; c < kSlices; ++c)
+        {
+            fullSum += hop.blockSq[itemRow][c];
+        }
+        float const rcpRms = rsqrtf(fullSum / static_cast<float>(kHiddenDim) + params.eps);
+        BfVec y;
+        y.packed = resOut;
+        BfVec g;
+        g.packed = hop.gamma[itemVec];
+        BfVec n;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e]));
+        }
+        if (writesOutputs)
+        {
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed;
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed;
+        }
+        rs.normed[itemRow][itemVec] = n.packed;
+    }
+
+    if (numRowsCta > 0)
+    {
+        __syncthreads();
+        int const tile = warp % Cfg::kExpertTiles;
+        int const kGroup = warp / Cfg::kExpertTiles;
+        // Independent accumulator chains keep the MMA latency off a single dependency chain.
+        float acc[kMmaChains][4] = {};
+#pragma unroll
+        for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+        {
+            // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands.
+            uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig];
+            int const chain = (2 * p) % kMmaChains;
+            mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y);
+            mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w);
+        }
+        float sum4[4];
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+        {
+            sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]);
+        }
+        // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner.
+        int const e0 = tile * 16 + gid;
+        int const e1 = e0 + 8;
+        uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner;
+        uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner;
+        storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0),
+            sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0));
+        storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1),
+            sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1));
+        // Also keeps this CTA resident until every partial addressed to it has landed.
+        mbarrierWaitPhase0(&hop.routerBar);
+
+        // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane groups
+        // of 8 sum over the 8 slices.
+        if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices)
+        {
+            int const c = tid % kSlices;
+            int const ij = tid / kSlices;
+            int const ei = ij / numRowsCta;
+            int const j = ij % numRowsCta;
+            float sum = 0.F;
+#pragma unroll
+            for (int kg = 0; kg < Cfg::kKGroups; ++kg)
+            {
+                sum += rs.partial[c][kg][ei][j];
+            }
+#pragma unroll
+            for (int mask = kSlices / 2; mask > 0; mask >>= 1)
+            {
+                sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices);
+            }
+            if (c == 0)
+            {
+                int const r = rowGroup + j * kRowSplit;
+                params.routerLogits[r * kNumExperts + expertGroup * Cfg::kExpertsPerCluster
+                    + slice * Cfg::kExpertsPerOwner + ei]
+                    = sum;
+            }
+        }
+    }
+
+    // ---- CTA 0 advances the protocol state once every CTA has read it. ----
+    if (tid == 0 && slice == 0 && clusterId == 0)
+    {
+        uint32_t arrived = 0;
+        do
+        {
+            asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory");
+        } while (arrived < static_cast<uint32_t>(kGridCtas));
+        // The next call clears exactly the slots this call wrote: compact rows 0..R-1 of every rank.
+        uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers,
+            static_cast<uint32_t>(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U);
+        asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x),
+                     "r"(next.y), "r"(next.z), "r"(next.w)
+                     : "memory");
+    }
+#else
+    // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices.
+    __trap();
+#endif
+}
+
+template <int R, int kRowSplit>
+cudaError_t launchTyped(LiveRowHopParams const& params)
+{
+    cudaLaunchAttribute attrs[2];
+    attrs[0].id = cudaLaunchAttributeClusterDimension;
+    attrs[0].val.clusterDim.x = kSlices;
+    attrs[0].val.clusterDim.y = 1;
+    attrs[0].val.clusterDim.z = 1;
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0;
+
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSlices, kExpertGroups * kRowSplit, 1);
+    config.blockDim = dim3(HopConfig::kThreads, 1, 1);
+    config.dynamicSmemBytes = 0;
+    config.stream = params.stream;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    return cudaLaunchKernelEx(&config, liveRowHopKernel<R, kRowSplit>, params);
+}
+
+} // namespace
+
+cudaError_t invokeDraftLiveRowHop(LiveRowHopParams const& params)
+{
+    if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr
+        || params.mcBuffer == nullptr || params.rowIds == nullptr || params.routerWeight == nullptr
+        || params.routerLogits == nullptr || params.numInputRows < 1 || params.numInputRows > kMaxRows
+        || params.numRows < 1 || params.numRows > params.numInputRows)
+    {
+        return cudaErrorInvalidValue;
+    }
+    // 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices).
+    switch (params.numRows)
+    {
+    case 1: return launchTyped<1, 1>(params);
+    case 2: return launchTyped<2, 2>(params);
+    case 3: return launchTyped<3, 3>(params);
+    case 4: return launchTyped<4, 3>(params);
+    case 5: return launchTyped<5, 3>(params);
+    case 6: return launchTyped<6, 3>(params);
+    case 7: return launchTyped<7, 3>(params);
+    case 8: return launchTyped<8, 3>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::draft_live_rows
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h
new file mode 100644
index 000000000000..a6c6f40b40f1
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h
@@ -0,0 +1,88 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::draft_live_rows
+{
+
+//! Live-row post-attention hop (MTP draft step 0).
+//!
+//! A fork of the router hop of postAttentionHop.h that serves only the ``numRows`` live rows of a call with
+//! ``numInputRows`` rows. Compact row ``r`` (0 <= r < numRows) is source row ``rowIds[r]`` of ``input`` and
+//! ``residualIn``; for it the kernel computes, on every rank and with identical bits on every rank:
+//!   sum      = bf16(sum over ranks 0..7 of float(partial_rank))          (one-shot Lamport exchange, fixed order)
+//!   residual = sum + residualIn                                         (bf16 add)
+//!   normed   = bf16(float(residual) * rsqrt(mean(bf16(residual^2)) + eps) * float(gamma))
+//!   logits   = normed @ routerWeight^T in FP32                          ([numRows, kNumExperts])
+//! Every quantity of a row is computed by the same per-row arithmetic as the post-attention hop, so compact row ``r``
+//! equals row ``rowIds[r]`` of a post-attention hop over all ``numInputRows`` rows.
+//!
+//! ``rowIds`` is read on the device after the kernel's programmatic dependency wait, never before it; it may be
+//! written by any earlier kernel of the stream. Ids outside ``[0, numInputRows)`` are clamped into that range.
+//!
+//! The exchange runs on the post-attention hop's multicast workspace and protocol state, with the same layout
+//! (``[row][source rank][kHiddenDim]`` bf16 per Lamport buffer, ``kNumLamportBuffers`` rotating buffers, flags
+//! ``{current buffer, bytes to clear in the previous buffer, arrival counter, unused}``), so calls of both kernels
+//! may be issued in any order on the stream that issues the group's hops.
+struct LiveRowHopParams
+{
+    void const* input = nullptr;        //!< [numInputRows, kHiddenDim] bf16, rank-local partial (o_proj output)
+    void const* residualIn = nullptr;   //!< [numInputRows, kHiddenDim] bf16
+    int64_t const* rowIds = nullptr;    //!< [numRows] device ids of the live rows in [0, numInputRows)
+    void const* gamma = nullptr;        //!< [kHiddenDim] bf16 RMSNorm weight
+    void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16
+    void* normOut = nullptr;            //!< [numRows, kHiddenDim] bf16
+    void* residualOut = nullptr;        //!< [numRows, kHiddenDim] bf16
+    float* routerLogits = nullptr;      //!< [numRows, kNumExperts] fp32
+    void* ucBuffer = nullptr;           //!< local unicast view of the multicast workspace
+    void* mcBuffer = nullptr;           //!< multicast address of the same workspace
+    uint32_t* flags = nullptr;          //!< [4] device-side protocol state of this workspace
+    int numRows = 0;                    //!< live rows, 1 <= numRows <= numInputRows
+    int numInputRows = 0;               //!< rows of ``input`` / ``residualIn``, <= kMaxRows
+    int rank = 0;
+    float eps = 1e-5F;
+    bool enablePdl = true;
+    cudaStream_t stream = nullptr;
+};
+
+// The post-attention hop's workspace geometry (the kernel runs on that workspace).
+constexpr int kHiddenDim = 6144;
+constexpr int kNumRanks = 8;
+constexpr int kNumExperts = 256;
+constexpr int kMaxRows = 8;
+constexpr int kNumLamportBuffers = 3;
+
+//! Bytes of one Lamport buffer: [kMaxRows][kNumRanks][kHiddenDim] bf16.
+constexpr size_t kLamportBufferBytes = size_t(kMaxRows) * kNumRanks * kHiddenDim * 2;
+//! Bytes of the whole multicast workspace.
+constexpr size_t kWorkspaceBytes = kLamportBufferBytes * kNumLamportBuffers;
+
+//! Launch the live-row hop on ``params.stream``. Returns the launch status.
+cudaError_t invokeDraftLiveRowHop(LiveRowHopParams const& params);
+
+} // namespace kernels::glm_moe_dsa::draft_live_rows
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu
new file mode 100644
index 000000000000..a4ca68a4859d
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu
@@ -0,0 +1,2628 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h"
+
+#include <cfloat>
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::draft_routed_moe
+{
+
+namespace
+{
+
+// Ring depths (2 KB chunks per warp; the deeper pair with kFlagDeepRing) and the FC2 intermediate ring (one entry
+// per expert of a tile).
+constexpr int kFc1Stages = 12;
+constexpr int kFc1StagesDeep = 20;
+constexpr int kFc2Stages = 8;
+constexpr int kFc2StagesDeep = 12;
+constexpr int kFc2HStages = 6;
+constexpr int kChunkBytes = 16 * kKChunk * 2;       // 16 rows x 128 bytes
+constexpr int kXRowBytes = kHidden * 2;             // 12288
+constexpr int kHEntryBytes = kMaxRows * kInter * 2; // 4096
+constexpr int kStageBytes = kMaxRows * kTopK * 16 * 2;
+constexpr float kSumEpsilon = 1e-20f;
+constexpr float kBf16MaxFinite = 3.38953139e38f;
+// Bytes of one FC2 weight block of an expert (one K chunk of all 6144 rows) and of one L2 prefetch piece.
+constexpr int kFc2BlockBytes = kHidden * kKChunk * 2;
+constexpr int kPrefetchPieceBytes = 16384;
+constexpr int kPrefetchPiecesPerBlock = kFc2BlockBytes / kPrefetchPieceBytes;
+static_assert(kFc2BlockBytes % kPrefetchPieceBytes == 0, "prefetch pieces must tile an FC2 block");
+// FC2 intermediate ring: expert s (in the warp's order over its items) has its K chunks at stream positions
+// 4s..4s+3 (alone or as the first of a pair) or 4s-3, 4s-1, 4s+1, 4s+3 (second of a pair), so it is last read at
+// iteration 4s+7 or earlier, and expert s + kFc2HStages loads into the same entry together with its first chunk, at
+// position >= 4(s + kFc2HStages) - 3, issued at the end of iteration (position - kStages + 1). The entry is free when
+// that iteration is not earlier than the last read: kStages <= 4 * kFc2HStages - 9.
+static_assert(kFc2Stages <= 4 * kFc2HStages - 9 && kFc2StagesDeep <= 4 * kFc2HStages - 9,
+    "an FC2 intermediate entry is refilled before its last read");
+
+// ---------------------------------------------------------------------------------------------------------------
+// Arithmetic with the exact instructions of the native kernels.
+// ---------------------------------------------------------------------------------------------------------------
+
+__device__ __forceinline__ float mulFtz(float a, float b)
+{
+    float r;
+    asm("mul.ftz.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+    return r;
+}
+
+__device__ __forceinline__ float addFtz(float a, float b)
+{
+    float r;
+    asm("add.ftz.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+    return r;
+}
+
+__device__ __forceinline__ float fmaFtz(float a, float b, float c)
+{
+    float r;
+    asm("fma.rn.ftz.f32 %0, %1, %2, %3;" : "=f"(r) : "f"(a), "f"(b), "f"(c));
+    return r;
+}
+
+__device__ __forceinline__ float tanhApprox(float x)
+{
+    float r;
+    asm("tanh.approx.f32 %0, %1;" : "=f"(r) : "f"(x));
+    return r;
+}
+
+__device__ __forceinline__ float divApproxFtz(float a, float b)
+{
+    float r;
+    asm("div.approx.ftz.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+    return r;
+}
+
+__device__ __forceinline__ float ex2ApproxFtz(float x)
+{
+    float r;
+    asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
+    return r;
+}
+
+__device__ __forceinline__ float rcpApproxFtz(float x)
+{
+    float r;
+    asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
+    return r;
+}
+
+__device__ __forceinline__ float mulRn(float a, float b)
+{
+    float r;
+    asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+    return r;
+}
+
+__device__ __forceinline__ float addRn(float a, float b)
+{
+    float r;
+    asm("add.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+    return r;
+}
+
+//! FlashInfer ``sigmoid_accurate`` as compiled with ``-use_fast_math``.
+__device__ __forceinline__ float routingSigmoid(float x)
+{
+    return fmaFtz(tanhApprox(mulFtz(x, 0.5f)), 0.5f, 0.5f);
+}
+
+//! Round to BF16 like ``F2FP.SATFINITE.BF16.F32``: round to nearest even, overflow saturates, NaN stays NaN.
+__device__ __forceinline__ __nv_bfloat16 bf16SatFinite(float x)
+{
+    __nv_bfloat16 r = __float2bfloat16_rn(x);
+    float back = __bfloat162float(r);
+    if (back == INFINITY || back == -INFINITY)
+    {
+        r = __float2bfloat16_rn(x > 0.f ? kBf16MaxFinite : -kBf16MaxFinite);
+    }
+    return r;
+}
+
+//! The native gated epilogue (``swiGlu``, no clamp limit, alpha 1, beta 0): ``(up + 0) * (gate * sigmoid(gate))``
+//! with ``sigmoid(g) = rcp.approx(1 + ex2.approx(-(g * log2 e)))``; both inputs are clamped to +-FLT_MAX first.
+__device__ __forceinline__ float swiGlu(float up, float gate)
+{
+    float const g = gate > FLT_MAX ? FLT_MAX : (gate < -FLT_MAX ? -FLT_MAX : gate);
+    float const u = up > FLT_MAX ? FLT_MAX : (up < -FLT_MAX ? -FLT_MAX : up);
+    float const t = mulRn(mulRn(g, 1.4426950216293334961f), -1.f);
+    float const s = rcpApproxFtz(addRn(ex2ApproxFtz(t), 1.f));
+    return mulRn(addRn(u, 0.f), mulRn(g, s));
+}
+
+// ---------------------------------------------------------------------------------------------------------------
+// Data movement and MMA.
+// ---------------------------------------------------------------------------------------------------------------
+
+__device__ __forceinline__ uint32_t smemAddr(void const* p)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
+}
+
+__device__ __forceinline__ void cpAsync16(uint32_t dst, void const* src)
+{
+    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(__cvta_generic_to_global(src)) : "memory");
+}
+
+__device__ __forceinline__ void cpAsyncCommit()
+{
+    asm volatile("cp.async.commit_group;" ::: "memory");
+}
+
+template <int N>
+__device__ __forceinline__ void cpAsyncWait()
+{
+    asm volatile("cp.async.wait_group %0;" ::"n"(N) : "memory");
+}
+
+__device__ __forceinline__ void ldmatrixX4(uint32_t (&r)[4], uint32_t addr)
+{
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+                 : "r"(addr));
+}
+
+__device__ __forceinline__ void mmaBf16(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};"
+        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
+}
+
+__device__ __forceinline__ void gridDependencyWait()
+{
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+}
+
+__device__ __forceinline__ void gridDependencyTrigger()
+{
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+}
+
+//! Byte offset of 16-byte unit ``u`` of a 128-byte row segment for row ``r`` (XOR swizzle, conflict-free ldmatrix).
+__device__ __forceinline__ uint32_t swizzled(int r, int u)
+{
+    return static_cast<uint32_t>(((u & ~7) | ((u & 7) ^ (r & 7))) << 4);
+}
+
+//! Copies one 16-row x 64-element weight chunk (2 KB, contiguous in the BlockMajorK layout) into a ring slot.
+__device__ __forceinline__ void issueWeightChunk(uint32_t slot, __nv_bfloat16 const* src, int lane)
+{
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        int const idx = lane + 32 * j;
+        int const r = idx >> 3;
+        int const u = idx & 7;
+        cpAsync16(slot + r * 128 + swizzled(r, u), src + r * kKChunk + u * 8);
+    }
+}
+
+//! A fragments of k-step ``q`` (16 elements) of a chunk: rows 0..15, natural K order (m16n8k16 row-major A).
+__device__ __forceinline__ void loadA(uint32_t (&a)[4], uint32_t slot, int q, int lane)
+{
+    int const m = lane >> 3;
+    int const r = (lane & 7) | ((m & 1) << 3);
+    int const u = 2 * q + (m >> 1);
+    ldmatrixX4(a, slot + r * 128 + swizzled(r, u));
+}
+
+//! B fragments of two consecutive k-steps from a [token][k] BF16 tile (rows ``rowBytes`` apart, units swizzled
+//! by token); tokens at or above ``rows`` read token 0 (their columns are discarded).
+__device__ __forceinline__ void loadB2(uint32_t (&b)[4], uint32_t base, int rowBytes, int unit, int rows, int lane)
+{
+    int tok = lane & 7;
+    tok = tok < rows ? tok : 0;
+    int const u = unit + (lane >> 3);
+    ldmatrixX4(b, base + tok * rowBytes + swizzled(tok, u));
+}
+
+// ---------------------------------------------------------------------------------------------------------------
+// Routing (FlashInfer routingCustom: SigmoidBias + ScaledSumNormalize, warp per token).
+// ---------------------------------------------------------------------------------------------------------------
+
+//! Selection key: the score's bits mapped to an unsigned order (a NaN score ranks above every number), then the
+//! expert id so that ties go to the lower index.
+__device__ __forceinline__ uint64_t packKey(float key, int expert)
+{
+    uint32_t bits = __float_as_uint(key);
+    bits = (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+    bits = key != key ? 0xFFFFFFFFu : bits;
+    return (static_cast<uint64_t>(bits) << 32) | static_cast<uint64_t>(0xFFFF & (65535 - expert));
+}
+
+//! Lexicographic warp maximum of 64-bit keys with two ``redux.sync.max.u32`` (as the native top-k reduction).
+__device__ __forceinline__ uint64_t warpMax64(uint64_t v)
+{
+    uint32_t const hi = static_cast<uint32_t>(v >> 32);
+    uint32_t const lo = static_cast<uint32_t>(v);
+    uint32_t maxHi;
+    asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(maxHi) : "r"(hi));
+    uint32_t const loContrib = hi == maxHi ? lo : 0u;
+    uint32_t maxLo;
+    asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(maxLo) : "r"(loContrib));
+    return (static_cast<uint64_t>(maxHi) << 32) | maxLo;
+}
+
+struct RouteSmem
+{
+    int topExpert[kMaxRows][kTopK];
+    float topWeight[kMaxRows][kTopK]; //!< Value of the BF16 weight.
+    int slotDistinct[kMaxRows][kTopK];
+    int distinct[kMaxDistinct];
+    uint32_t mask[kNumExperts / 32]; //!< Experts selected by any row: bit e % 32 of word e / 32.
+    int numDistinct;
+};
+
+//! A lane's routing inputs of one token: logits of experts 8l .. 8l + 7 (two 16-byte loads).
+struct RouteLogits
+{
+    float4 l0;
+    float4 l1;
+};
+
+__device__ __forceinline__ RouteLogits loadRouteLogits(DraftRoutedMoeParams const& p, int token, int lane)
+{
+    float4 const* logits = reinterpret_cast<float4 const*>(p.logits + token * kNumExperts + lane * (kNumExperts / 32));
+    return RouteLogits{logits[0], logits[1]};
+}
+
+//! Compare-exchange for a descending order: a[i] >= a[j] afterwards.
+__device__ __forceinline__ void keyExchange(uint64_t& hi, uint64_t& lo)
+{
+    uint64_t const a = hi;
+    uint64_t const b = lo;
+    hi = a > b ? a : b;
+    lo = a > b ? b : a;
+}
+
+//! Sorts 8 keys in descending order (Batcher's odd-even merge sort, 19 compare-exchanges).
+__device__ __forceinline__ void sortKeysDescending(uint64_t (&k)[8])
+{
+    keyExchange(k[0], k[1]);
+    keyExchange(k[2], k[3]);
+    keyExchange(k[4], k[5]);
+    keyExchange(k[6], k[7]);
+    keyExchange(k[0], k[2]);
+    keyExchange(k[1], k[3]);
+    keyExchange(k[4], k[6]);
+    keyExchange(k[5], k[7]);
+    keyExchange(k[1], k[2]);
+    keyExchange(k[5], k[6]);
+    keyExchange(k[0], k[4]);
+    keyExchange(k[1], k[5]);
+    keyExchange(k[2], k[6]);
+    keyExchange(k[3], k[7]);
+    keyExchange(k[2], k[4]);
+    keyExchange(k[3], k[5]);
+    keyExchange(k[1], k[2]);
+    keyExchange(k[3], k[4]);
+    keyExchange(k[5], k[6]);
+}
+
+//! Top-8 of one token by one warp; lane k < 8 receives slot k (expert, weight). Lane l holds experts 8l .. 8l + 7
+//! (its logits and its 16-byte slice of the BF16 bias, ``braw``); the selection key carries the expert id, so the
+//! result does not depend on this placement.
+__device__ void routeToken(
+    DraftRoutedMoeParams const& p, int token, int lane, RouteLogits const& in, uint4 braw, RouteSmem& rs)
+{
+    constexpr int kPerLane = kNumExperts / 32;
+    float const logit[kPerLane] = {in.l0.x, in.l0.y, in.l0.z, in.l0.w, in.l1.x, in.l1.y, in.l1.z, in.l1.w};
+    uint32_t const biasPairs[kPerLane / 2] = {braw.x, braw.y, braw.z, braw.w};
+    float sig[kPerLane];
+    uint64_t key[kPerLane];
+#pragma unroll
+    for (int i = 0; i < kPerLane; ++i)
+    {
+        uint32_t const pair = biasPairs[i / 2];
+        float const bias = __uint_as_float((i % 2 == 0 ? pair << 16 : pair & 0xFFFF0000u));
+        sig[i] = routingSigmoid(logit[i]);
+        key[i] = packKey(addFtz(sig[i], bias), lane * kPerLane + i);
+    }
+    // The lane's keys in descending order; then every round takes the warp maximum of the lanes' leading keys and
+    // the owning lane drops its leading key (keys are unique: they carry the expert id). Every round thus selects
+    // the largest key not selected yet, as the native top-k reduction does.
+    sortKeysDescending(key);
+    float mySig = 0.f;
+    int myExpert = 0;
+#pragma unroll
+    for (int k = 0; k < kTopK; ++k)
+    {
+        uint64_t const best = warpMax64(key[0]);
+        bool const pop = key[0] == best;
+#pragma unroll
+        for (int i = 0; i + 1 < kPerLane; ++i)
+        {
+            key[i] = pop ? key[i + 1] : key[i];
+        }
+        key[kPerLane - 1] = pop ? 0 : key[kPerLane - 1];
+        int const e = 65535 - static_cast<int>(best & 0xFFFF);
+        int const owner = e / kPerLane;
+        int const slot = e % kPerLane;
+        float s = 0.f;
+#pragma unroll
+        for (int i = 0; i < kPerLane; ++i)
+        {
+            s = i == slot ? sig[i] : s;
+        }
+        s = __shfl_sync(0xffffffffu, s, owner);
+        if (lane == k)
+        {
+            mySig = s;
+            myExpert = e;
+        }
+    }
+    // cg::reduce(warp, sigmoid, plus): XOR butterfly 16, 8, 4, 2, 1; lanes >= 8 hold 0.
+    float sum = lane < kTopK ? mySig : 0.f;
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+    {
+        sum = addFtz(sum, __shfl_xor_sync(0xffffffffu, sum, off));
+    }
+    if (lane < kTopK)
+    {
+        float const w = divApproxFtz(mulFtz(mySig, p.routedScaling), addFtz(sum, kSumEpsilon));
+        rs.topExpert[token][lane] = myExpert;
+        rs.topWeight[token][lane] = __bfloat162float(__float2bfloat16_rn(w));
+        atomicOr(&rs.mask[myExpert >> 5], 1u << (myExpert & 31));
+    }
+}
+
+//! Routing of all rows plus the distinct-expert list (ascending ids), by the whole CTA. ``tokensDone`` (optional)
+//! receives clock64 when this thread's warp has routed its tokens.
+template <int kWarps>
+__device__ void routeAll(
+    DraftRoutedMoeParams const& p, RouteSmem& rs, int warp, int lane, int64_t* tokensDone = nullptr)
+{
+    if (threadIdx.x < kNumExperts / 32)
+    {
+        rs.mask[threadIdx.x] = 0;
+    }
+    // A warp routes rows warp, warp + kWarps, ...: their logit loads are all issued before the first selection.
+    constexpr int kPerWarp = (kMaxRows + kWarps - 1) / kWarps;
+    uint4 const braw = *reinterpret_cast<uint4 const*>(p.bias + lane * (kNumExperts / 32));
+    RouteLogits in[kPerWarp];
+#pragma unroll
+    for (int i = 0; i < kPerWarp; ++i)
+    {
+        if (warp + i * kWarps < p.rows)
+        {
+            in[i] = loadRouteLogits(p, warp + i * kWarps, lane);
+        }
+    }
+    __syncthreads(); // the mask is cleared before any row sets its bits
+#pragma unroll
+    for (int i = 0; i < kPerWarp; ++i)
+    {
+        if (warp + i * kWarps < p.rows)
+        {
+            routeToken(p, warp + i * kWarps, lane, in[i], braw, rs);
+        }
+    }
+    if (tokensDone != nullptr)
+    {
+        *tokensDone = clock64();
+    }
+    __syncthreads();
+    // Distinct experts in ascending order: the index of expert e is the number of selected experts below it.
+    if (warp == 0 && lane < kNumExperts / 32)
+    {
+        int idx = 0;
+        for (int w = 0; w < lane; ++w)
+        {
+            idx += __popc(rs.mask[w]);
+        }
+        uint32_t m = rs.mask[lane];
+        if (lane == kNumExperts / 32 - 1)
+        {
+            rs.numDistinct = idx + __popc(m);
+        }
+        while (m != 0)
+        {
+            rs.distinct[idx++] = lane * 32 + __ffs(m) - 1;
+            m &= m - 1;
+        }
+    }
+    for (int i = threadIdx.x; i < p.rows * kTopK; i += kWarps * 32)
+    {
+        int const e = rs.topExpert[i / kTopK][i % kTopK];
+        int idx = __popc(rs.mask[e >> 5] & ((1u << (e & 31)) - 1u));
+        for (int w = 0; w < (e >> 5); ++w)
+        {
+            idx += __popc(rs.mask[w]);
+        }
+        rs.slotDistinct[i / kTopK][i % kTopK] = idx;
+    }
+    __syncthreads();
+}
+
+// ---------------------------------------------------------------------------------------------------------------
+// FC1: routing + gate/up GEMM + SwiGLU.
+// ---------------------------------------------------------------------------------------------------------------
+
+template <int kWarps, bool kPairs, int kStages>
+__global__ void __launch_bounds__(kWarps * 32) mtpMoeFc1Kernel(DraftRoutedMoeParams const p)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    extern __shared__ __align__(128) uint8_t smem[];
+    __shared__ RouteSmem rs;
+    int const warp = threadIdx.x >> 5;
+    int const lane = threadIdx.x & 31;
+    int const rows = p.rows;
+    uint8_t* xs = smem;
+    uint8_t* ring = smem + rows * kXRowBytes + warp * kStages * kChunkBytes;
+    uint32_t const xsAddr = smemAddr(xs);
+    uint32_t const ringAddr = smemAddr(ring);
+
+    // Every input below (logits, activation) is produced by the predecessor grid.
+    gridDependencyWait();
+
+    // The activation rows, swizzled per 128-byte segment by token.
+    for (int i = threadIdx.x; i < rows * (kXRowBytes / 16); i += kWarps * 32)
+    {
+        int const t = i / (kXRowBytes / 16);
+        int const u = i % (kXRowBytes / 16);
+        cpAsync16(xsAddr + t * kXRowBytes + swizzled(t, u), p.x + t * kHidden + u * 8);
+    }
+    cpAsyncCommit();
+
+    routeAll<kWarps>(p, rs, warp, lane);
+    int const numDistinct = rs.numDistinct;
+    if (blockIdx.x == 0)
+    {
+        for (int i = threadIdx.x; i < kRouteWords; i += kWarps * 32)
+        {
+            int v = 0;
+            if (i == 0)
+            {
+                v = numDistinct;
+            }
+            else if (i < kRouteSlotOffset)
+            {
+                v = i - kRouteDistinctOffset < numDistinct ? rs.distinct[i - kRouteDistinctOffset] : -1;
+            }
+            else if (i < kRouteWeightOffset)
+            {
+                int const j = i - kRouteSlotOffset;
+                v = j / kTopK < rows ? rs.slotDistinct[j / kTopK][j % kTopK] : -1;
+            }
+            else
+            {
+                int const j = i - kRouteWeightOffset;
+                v = j / kTopK < rows ? __float_as_int(rs.topWeight[j / kTopK][j % kTopK]) : 0;
+            }
+            p.route[i] = v;
+        }
+    }
+    if (p.flags & kFlagEarlyTrigger)
+    {
+        gridDependencyTrigger();
+    }
+    cpAsyncWait<0>();
+    __syncthreads();
+
+    int const totalWarps = gridDim.x * kWarps;
+    int const gw = blockIdx.x * kWarps + warp;
+    int const items = numDistinct * kFc1Tiles;
+    int const myItems = items > gw ? (items - gw + totalWarps - 1) / totalWarps : 0;
+    int const total = myItems * kFc1Chunks;
+    // With kPairs, the warp's tiles are taken two at a time and their chunks interleave (A0 B0 A1 B1 ...): two
+    // independent accumulation chains per warp, each still in natural K order. An odd last tile runs alone.
+    int const pairChunks = kPairs ? (myItems / 2) * 2 * kFc1Chunks : 0;
+
+    // Chunk n of the warp's stream -> (local tile index, K chunk).
+    auto chunkOf = [&](int n, int& tile, int& kc)
+    {
+        if (n < pairChunks)
+        {
+            int const r = n % (2 * kFc1Chunks);
+            tile = 2 * (n / (2 * kFc1Chunks)) + (r & 1);
+            kc = r >> 1;
+        }
+        else
+        {
+            tile = (pairChunks / kFc1Chunks) + (n - pairChunks) / kFc1Chunks;
+            kc = (n - pairChunks) % kFc1Chunks;
+        }
+    };
+
+    auto issue = [&](int n)
+    {
+        int tile, kc;
+        chunkOf(n, tile, kc);
+        int const item = gw + tile * totalWarps;
+        int const e = rs.distinct[item / kFc1Tiles];
+        int const a = item % kFc1Tiles;
+        __nv_bfloat16 const* src = p.w13 + ((static_cast<int64_t>(e) * kFc1Chunks + kc) * kFc1Rows + a * 16) * kKChunk;
+        issueWeightChunk(ringAddr + (n % kStages) * kChunkBytes, src, lane);
+    };
+
+    auto epilogue = [&](int tile, float const(&acc)[4])
+    {
+        int const item = gw + tile * totalWarps;
+        int const d = item / kFc1Tiles;
+        int const a = item % kFc1Tiles;
+        int const g = lane >> 2;
+        int const t = lane & 3;
+        // Rows 0..7 of the tile are up rows, rows 8..15 the gate rows of the same intermediate columns.
+        int const j = 16 * (a >> 1) + 2 * g + (a & 1);
+#pragma unroll
+        for (int c = 0; c < 2; ++c)
+        {
+            int const tok = 2 * t + c;
+            if (tok < rows)
+            {
+                p.h[(static_cast<int64_t>(d) * kMaxRows + tok) * kInter + j]
+                    = bf16SatFinite(swiGlu(acc[c], acc[2 + c]));
+            }
+        }
+    };
+
+#pragma unroll 1
+    for (int n = 0; n < kStages - 1; ++n)
+    {
+        if (n < total)
+        {
+            issue(n);
+        }
+        cpAsyncCommit();
+    }
+
+    float acc[4] = {-0.f, -0.f, -0.f, -0.f};
+    float acc2[4] = {-0.f, -0.f, -0.f, -0.f};
+#pragma unroll 1
+    for (int n = 0; n < total; ++n)
+    {
+        // Chunk n has landed once at most kStages - 2 newer groups are pending.
+        cpAsyncWait<kStages - 2>();
+        __syncwarp();
+        int tile, kc;
+        chunkOf(n, tile, kc);
+        uint32_t const slot = ringAddr + (n % kStages) * kChunkBytes;
+        if (n >= pairChunks)
+        {
+#pragma unroll
+            for (int q = 0; q < 4; q += 2)
+            {
+                uint32_t b[4];
+                loadB2(b, xsAddr, kXRowBytes, kc * 8 + 2 * q, rows, lane);
+                uint32_t a[4];
+                loadA(a, slot, q, lane);
+                mmaBf16(acc, a, b[0], b[1]);
+                loadA(a, slot, q + 1, lane);
+                mmaBf16(acc, a, b[2], b[3]);
+            }
+            if (kc == kFc1Chunks - 1)
+            {
+                epilogue(tile, acc);
+                acc[0] = acc[1] = acc[2] = acc[3] = -0.f;
+            }
+        }
+        else if (n & 1)
+        {
+            // Both chunks of this K step of the pair are in: chunk n - 1 (first tile) and n (second tile).
+            uint32_t const slot0 = ringAddr + ((n - 1) % kStages) * kChunkBytes;
+#pragma unroll
+            for (int q = 0; q < 4; q += 2)
+            {
+                uint32_t b[4];
+                loadB2(b, xsAddr, kXRowBytes, kc * 8 + 2 * q, rows, lane);
+                uint32_t a0[4];
+                uint32_t a1[4];
+                loadA(a0, slot0, q, lane);
+                loadA(a1, slot, q, lane);
+                mmaBf16(acc, a0, b[0], b[1]);
+                mmaBf16(acc2, a1, b[0], b[1]);
+                loadA(a0, slot0, q + 1, lane);
+                loadA(a1, slot, q + 1, lane);
+                mmaBf16(acc, a0, b[2], b[3]);
+                mmaBf16(acc2, a1, b[2], b[3]);
+            }
+            if (kc == kFc1Chunks - 1)
+            {
+                epilogue(tile - 1, acc);
+                epilogue(tile, acc2);
+                acc[0] = acc[1] = acc[2] = acc[3] = -0.f;
+                acc2[0] = acc2[1] = acc2[2] = acc2[3] = -0.f;
+            }
+        }
+        __syncwarp();
+        // The slot refilled now held chunk n - 1, consumed in this or the previous iteration.
+        if (n + kStages - 1 < total)
+        {
+            issue(n + kStages - 1);
+        }
+        cpAsyncCommit();
+    }
+    cpAsyncWait<0>();
+
+    if (p.flags & kFlagPrefetchFc2)
+    {
+        int const pieces = numDistinct * kFc2Chunks * kPrefetchPiecesPerBlock;
+        for (int piece = gw * 32 + lane; piece < pieces; piece += totalWarps * 32)
+        {
+            int const d = piece / (kFc2Chunks * kPrefetchPiecesPerBlock);
+            int const rem = piece % (kFc2Chunks * kPrefetchPiecesPerBlock);
+            int const kc = rem / kPrefetchPiecesPerBlock;
+            int const off = rem % kPrefetchPiecesPerBlock;
+            uint8_t const* addr = reinterpret_cast<uint8_t const*>(p.w2
+                                      + (static_cast<int64_t>(rs.distinct[d]) * kFc2Chunks + kc) * kHidden * kKChunk)
+                + static_cast<int64_t>(off) * kPrefetchPieceBytes;
+            asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(__cvta_generic_to_global(addr)),
+                         "r"(kPrefetchPieceBytes)
+                         : "memory");
+        }
+    }
+    if (!(p.flags & kFlagEarlyTrigger))
+    {
+        gridDependencyTrigger();
+    }
+#endif
+}
+
+// ---------------------------------------------------------------------------------------------------------------
+// FC2: down GEMM per distinct expert + native finalize, 16 hidden rows per work item.
+// ---------------------------------------------------------------------------------------------------------------
+
+struct Fc2Smem
+{
+    float weight[kMaxRows][kTopK];
+    int distinct[kMaxDistinct];
+    int8_t slotOf[kMaxRows][kMaxDistinct];
+    int numDistinct;
+};
+
+__device__ __forceinline__ int hiddenOfStoredRow(int s)
+{
+    return (s & ~31) | ((s & 7) << 2) | ((s & 31) >> 3);
+}
+
+template <int kWarps, bool kPairs, int kStages>
+__global__ void __launch_bounds__(kWarps * 32) mtpMoeFc2Kernel(DraftRoutedMoeParams const p)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    extern __shared__ __align__(128) uint8_t smem[];
+    __shared__ Fc2Smem fs;
+    int const warp = threadIdx.x >> 5;
+    int const lane = threadIdx.x & 31;
+    int const rows = p.rows;
+    constexpr int kWarpBytes = kStages * kChunkBytes + kFc2HStages * kHEntryBytes + kStageBytes;
+    uint8_t* mine = smem + warp * kWarpBytes;
+    uint32_t const ringAddr = smemAddr(mine);
+    uint32_t const hAddr = ringAddr + kStages * kChunkBytes;
+    __nv_bfloat16* stage = reinterpret_cast<__nv_bfloat16*>(mine + kStages * kChunkBytes + kFc2HStages * kHEntryBytes);
+
+    // The routing record and the intermediate are written by the FC1 grid.
+    gridDependencyWait();
+
+    int const numDistinct = p.route[0];
+    for (int i = threadIdx.x; i < kMaxRows * kMaxDistinct; i += kWarps * 32)
+    {
+        fs.slotOf[i / kMaxDistinct][i % kMaxDistinct] = -1;
+    }
+    for (int i = threadIdx.x; i < kMaxDistinct; i += kWarps * 32)
+    {
+        fs.distinct[i] = p.route[kRouteDistinctOffset + i];
+    }
+    __syncthreads();
+    for (int i = threadIdx.x; i < rows * kTopK; i += kWarps * 32)
+    {
+        int const t = i / kTopK;
+        int const k = i % kTopK;
+        fs.slotOf[t][p.route[kRouteSlotOffset + i]] = static_cast<int8_t>(k);
+        fs.weight[t][k] = __int_as_float(p.route[kRouteWeightOffset + i]);
+    }
+    __syncthreads();
+
+    int const totalWarps = gridDim.x * kWarps;
+    int const gw = blockIdx.x * kWarps + warp;
+    int const perItem = numDistinct * kFc2Chunks;
+    int const myItems = kFc2Tiles > gw ? (kFc2Tiles - gw + totalWarps - 1) / totalWarps : 0;
+    int const total = myItems * perItem;
+    // With kPairs, the experts of an item are taken two at a time and their chunks interleave (two independent
+    // accumulation chains, each in natural K order); an odd last expert runs alone.
+    int const pairChunks = kPairs ? (numDistinct / 2) * 2 * kFc2Chunks : 0;
+
+    // Chunk m of an item -> (distinct expert index, K chunk).
+    auto chunkOf = [&](int m, int& d, int& kc)
+    {
+        if (m < pairChunks)
+        {
+            int const r = m % (2 * kFc2Chunks);
+            d = 2 * (m / (2 * kFc2Chunks)) + (r & 1);
+            kc = r >> 1;
+        }
+        else
+        {
+            d = (pairChunks / kFc2Chunks) + (m - pairChunks) / kFc2Chunks;
+            kc = (m - pairChunks) % kFc2Chunks;
+        }
+    };
+    // Intermediate entry of expert d of local item il (entries are consumed in this order).
+    auto hEntry = [&](int il, int d) { return hAddr + ((il * numDistinct + d) % kFc2HStages) * kHEntryBytes; };
+
+    auto issue = [&](int n)
+    {
+        int const il = n / perItem;
+        int d, kc;
+        chunkOf(n % perItem, d, kc);
+        int const tile = gw + il * totalWarps;
+        int const e = fs.distinct[d];
+        __nv_bfloat16 const* src = p.w2 + ((static_cast<int64_t>(e) * kFc2Chunks + kc) * kHidden + tile * 16) * kKChunk;
+        issueWeightChunk(ringAddr + (n % kStages) * kChunkBytes, src, lane);
+        if (kc == 0)
+        {
+            uint32_t const entry = hEntry(il, d);
+            for (int u = lane; u < rows * (kInter / 8); u += 32)
+            {
+                int const tok = u / (kInter / 8);
+                int const unit = u % (kInter / 8);
+                cpAsync16(entry + tok * (kInter * 2) + swizzled(tok, unit),
+                    p.h + (static_cast<int64_t>(d) * kMaxRows + tok) * kInter + unit * 8);
+            }
+        }
+    };
+
+    // FC2 rows of expert d (C fragment) -> BF16, into the token's slot of the staging area.
+    auto scatter = [&](int tile, int d, float const(&acc)[4])
+    {
+        int const g = lane >> 2;
+        int const t = lane & 3;
+#pragma unroll
+        for (int c = 0; c < 2; ++c)
+        {
+            int const tok = 2 * t + c;
+            int const k = tok < rows ? fs.slotOf[tok][d] : -1;
+            if (k >= 0)
+            {
+                __nv_bfloat16 const lo = bf16SatFinite(acc[c]);
+                __nv_bfloat16 const hi = bf16SatFinite(acc[2 + c]);
+                stage[(tok * kTopK + k) * 16 + g] = lo;
+                stage[(tok * kTopK + k) * 16 + g + 8] = hi;
+                if (p.debugFc2 != nullptr)
+                {
+                    int64_t const base = (static_cast<int64_t>(tok) * kTopK + k) * kHidden;
+                    p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g)] = lo;
+                    p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g + 8)] = hi;
+                }
+            }
+        }
+    };
+
+    // Native finalize of the item's 16 hidden rows for every token: FP32 fma chain from 0 in slot order.
+    auto finalize = [&](int tile)
+    {
+        __syncwarp();
+        for (int pr = lane; pr < rows * 16; pr += 32)
+        {
+            int const tok = pr >> 4;
+            int const row = pr & 15;
+            float sum = 0.f;
+#pragma unroll
+            for (int k = 0; k < kTopK; ++k)
+            {
+                sum = fmaFtz(fs.weight[tok][k], __bfloat162float(stage[(tok * kTopK + k) * 16 + row]), sum);
+            }
+            p.out[static_cast<int64_t>(tok) * kHidden + hiddenOfStoredRow(tile * 16 + row)] = __float2bfloat16_rn(sum);
+        }
+    };
+
+#pragma unroll 1
+    for (int n = 0; n < kStages - 1; ++n)
+    {
+        if (n < total)
+        {
+            issue(n);
+        }
+        cpAsyncCommit();
+    }
+
+    float acc[4] = {-0.f, -0.f, -0.f, -0.f};
+    float acc2[4] = {-0.f, -0.f, -0.f, -0.f};
+#pragma unroll 1
+    for (int n = 0; n < total; ++n)
+    {
+        cpAsyncWait<kStages - 2>();
+        __syncwarp();
+        int const il = n / perItem;
+        int const m = n % perItem;
+        int const tile = gw + il * totalWarps;
+        int d, kc;
+        chunkOf(m, d, kc);
+        uint32_t const slot = ringAddr + (n % kStages) * kChunkBytes;
+        if (m >= pairChunks)
+        {
+            uint32_t const entry = hEntry(il, d);
+#pragma unroll
+            for (int q = 0; q < 4; q += 2)
+            {
+                uint32_t b[4];
+                loadB2(b, entry, kInter * 2, kc * 8 + 2 * q, rows, lane);
+                uint32_t a[4];
+                loadA(a, slot, q, lane);
+                mmaBf16(acc, a, b[0], b[1]);
+                loadA(a, slot, q + 1, lane);
+                mmaBf16(acc, a, b[2], b[3]);
+            }
+            if (kc == kFc2Chunks - 1)
+            {
+                scatter(tile, d, acc);
+                acc[0] = acc[1] = acc[2] = acc[3] = -0.f;
+            }
+        }
+        else if (m & 1)
+        {
+            uint32_t const slot0 = ringAddr + ((n - 1) % kStages) * kChunkBytes;
+            uint32_t const entry0 = hEntry(il, d - 1);
+            uint32_t const entry1 = hEntry(il, d);
+#pragma unroll
+            for (int q = 0; q < 4; q += 2)
+            {
+                uint32_t b0[4];
+                uint32_t b1[4];
+                loadB2(b0, entry0, kInter * 2, kc * 8 + 2 * q, rows, lane);
+                loadB2(b1, entry1, kInter * 2, kc * 8 + 2 * q, rows, lane);
+                uint32_t a0[4];
+                uint32_t a1[4];
+                loadA(a0, slot0, q, lane);
+                loadA(a1, slot, q, lane);
+                mmaBf16(acc, a0, b0[0], b0[1]);
+                mmaBf16(acc2, a1, b1[0], b1[1]);
+                loadA(a0, slot0, q + 1, lane);
+                loadA(a1, slot, q + 1, lane);
+                mmaBf16(acc, a0, b0[2], b0[3]);
+                mmaBf16(acc2, a1, b1[2], b1[3]);
+            }
+            if (kc == kFc2Chunks - 1)
+            {
+                scatter(tile, d - 1, acc);
+                scatter(tile, d, acc2);
+                acc[0] = acc[1] = acc[2] = acc[3] = -0.f;
+                acc2[0] = acc2[1] = acc2[2] = acc2[3] = -0.f;
+            }
+        }
+        if (m == perItem - 1)
+        {
+            finalize(tile);
+        }
+        __syncwarp();
+        if (n + kStages - 1 < total)
+        {
+            issue(n + kStages - 1);
+        }
+        cpAsyncCommit();
+    }
+    cpAsyncWait<0>();
+#endif
+}
+
+// ---------------------------------------------------------------------------------------------------------------
+// TMA-fed kernels (kFlagTma). Per CTA one producer thread (last warp, lane 0) streams 3D weight boxes (16 rows per
+// chain x 64 elements x several K chunks of the BlockMajorK layout, 128-byte swizzle = the layout loadA reads; 32 KB
+// per FC1 box, one expert's K = 256 per FC2 box) into a ring of S stages; W MMA warps each own C 16-row accumulation
+// chains per tile and consume every stage. Full barriers complete on the transaction count; each MMA warp releases
+// a stage (empty barrier, count W) after the MMAs that read it. The arithmetic of every chain (fragments, k-step
+// order, accumulator init, epilogues) is that of the kernels above.
+// ---------------------------------------------------------------------------------------------------------------
+
+#ifndef GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE
+#define GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE 0
+#endif
+
+//! Probe builds (GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1) record per-CTA events into DraftRoutedMoeParams::timeline;
+//! production builds compile none of it.
+constexpr bool kTimeline = GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE != 0;
+
+//! Columns of a CTA's event record: %globaltimer at entry and when MMA warp 0 is done (ns), the SM id, clock64
+//! cycles after entry, and counts.
+enum TimelineSlot : int
+{
+    kTlEntryNs = 0,
+    kTlDoneNs = 1,
+    kTlSmId = 2,
+    kTlWait = 3,          //!< griddepcontrol.wait returned
+    kTlSetup = 4,         //!< routing (FC1) / routing record (FC2) in shared memory, roles split
+    kTlFirst = 5,         //!< MMA warp 0 found its first stage full
+    kTlLast = 6,          //!< MMA warp 0 found its last stage full
+    kTlDone = 7,          //!< MMA warp 0 done (epilogue / finalize stored)
+    kTlProdFirst = 8,     //!< producer issued its first weight copy
+    kTlProdLast = 9,      //!< producer issued its last weight copy
+    kTlFullBlocked = 10,  //!< MMA warp 0: stage waits whose first poll failed
+    kTlEmptyBlocked = 11, //!< producer: re-arm waits whose first poll failed
+    kTlChunks = 12,       //!< MMA warp 0: K chunks consumed
+    kTlRouteTokens = 13,  //!< FC1 warp 0 has routed its tokens
+    kTlRouted = 14,       //!< FC1 routing and distinct-expert list done (before the activation wait)
+};
+
+//! K chunks (2 KB per 16-row chain) per stage of the TMA-fed FC1 kernel: one 32 KB 3D box per stage.
+static_assert(kTmaStageChunks * kChunkBytes == 32 * 1024, "32 KB FC1 stages");
+
+__device__ __forceinline__ int64_t globalTimerNs()
+{
+    uint64_t t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return static_cast<int64_t>(t);
+}
+
+__device__ __forceinline__ int64_t smId()
+{
+    uint32_t id;
+    asm volatile("mov.u32 %0, %%smid;" : "=r"(id));
+    return id;
+}
+
+//! This CTA's event record of kernel ``kernel`` (0: FC1, 1: FC2), or null.
+__device__ __forceinline__ int64_t* timelineRow(DraftRoutedMoeParams const& p, int kernel)
+{
+    if (!kTimeline || p.timeline == nullptr || blockIdx.x >= kTimelineCtas)
+    {
+        return nullptr;
+    }
+    return p.timeline + (static_cast<int64_t>(kernel) * kTimelineCtas + blockIdx.x) * kTimelineSlots;
+}
+
+__device__ __forceinline__ void mbarInit(uint32_t bar, uint32_t count)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(bar), "r"(count) : "memory");
+}
+
+__device__ __forceinline__ void mbarFenceInit()
+{
+    asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+}
+
+__device__ __forceinline__ void mbarArriveExpectTx(uint32_t bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory");
+}
+
+__device__ __forceinline__ void mbarArrive(uint32_t bar)
+{
+    asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(bar) : "memory");
+}
+
+__device__ __forceinline__ bool mbarTryWait(uint32_t bar, uint32_t parity)
+{
+    uint32_t done;
+    asm volatile(
+        "{\n\t"
+        ".reg .pred p;\n\t"
+        "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n\t"
+        "selp.u32 %0, 1, 0, p;\n\t"
+        "}"
+        : "=r"(done)
+        : "r"(bar), "r"(parity)
+        : "memory");
+    return done != 0;
+}
+
+//! Non-blocking probe of the phase of ``bar`` with parity ``parity`` (acquire when it returns true).
+__device__ __forceinline__ bool mbarTestWait(uint32_t bar, uint32_t parity)
+{
+    uint32_t done;
+    asm volatile(
+        "{\n\t"
+        ".reg .pred p;\n\t"
+        "mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n\t"
+        "selp.u32 %0, 1, 0, p;\n\t"
+        "}"
+        : "=r"(done)
+        : "r"(bar), "r"(parity)
+        : "memory");
+    return done != 0;
+}
+
+//! Blocks until the phase of ``bar`` with parity ``parity`` has completed. Every wait is between threads of one CTA
+//! (never on another grid), so one that does not complete within 2^28 polls is a protocol error: trap, do not hang.
+__device__ __forceinline__ void mbarWait(uint32_t bar, uint32_t parity)
+{
+    for (uint32_t i = 0; !mbarTryWait(bar, parity); ++i)
+    {
+        if (i == (1u << 28))
+        {
+            __trap();
+        }
+    }
+}
+
+//! mbarWait; returns 1 when the first poll failed (probe builds count these), else 0.
+[[maybe_unused]] __device__ __forceinline__ int mbarWaitCounted(uint32_t bar, uint32_t parity)
+{
+    if (mbarTryWait(bar, parity))
+    {
+        return 0;
+    }
+    mbarWait(bar, parity);
+    return 1;
+}
+
+//! Barrier among the W MMA warps (threads 0 .. 32 W - 1) of a warp-specialized kernel; the producer warp does not
+//! take part. Named barrier 1 (``__syncthreads`` uses 0).
+template <int W>
+__device__ __forceinline__ void mmaWarpsSync()
+{
+    asm volatile("bar.sync 1, %0;" ::"n"(W * 32) : "memory");
+}
+
+__device__ __forceinline__ void tmaLoad3d(uint32_t dst, CUtensorMap const* map, int c0, int c1, int c2, uint32_t bar)
+{
+    asm volatile(
+        "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4}], "
+        "[%5];" ::"r"(dst),
+        "l"(reinterpret_cast<uint64_t>(map)), "r"(c0), "r"(c1), "r"(c2), "r"(bar)
+        : "memory");
+}
+
+//! tmaLoad3d with an L2 cache policy (createpolicy).
+__device__ __forceinline__ void tmaLoad3dHint(
+    uint32_t dst, CUtensorMap const* map, int c0, int c1, int c2, uint32_t bar, uint64_t policy)
+{
+    asm volatile(
+        "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1, {%2, "
+        "%3, %4}], [%5], %6;" ::"r"(dst),
+        "l"(reinterpret_cast<uint64_t>(map)), "r"(c0), "r"(c1), "r"(c2), "r"(bar), "l"(policy)
+        : "memory");
+}
+
+//! L2 policy that evicts the lines it is attached to first.
+__device__ __forceinline__ uint64_t l2EvictFirstPolicy()
+{
+    uint64_t policy;
+    asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy));
+    return policy;
+}
+
+__device__ __forceinline__ void tmaPrefetchL2(CUtensorMap const* map, int c0, int c1, int c2)
+{
+    asm volatile(
+        "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [%0, {%1, %2, %3}];" ::"l"(reinterpret_cast<uint64_t>(map)),
+        "r"(c0), "r"(c1), "r"(c2)
+        : "memory");
+}
+
+__device__ __forceinline__ void bulkLoad(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar)
+{
+    asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst),
+                 "l"(__cvta_generic_to_global(src)), "r"(bytes), "r"(bar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void prefetchTensorMap(CUtensorMap const* map)
+{
+    asm volatile("prefetch.tensormap [%0];" ::"l"(reinterpret_cast<uint64_t>(map)) : "memory");
+}
+
+//! L2 prefetch of this CTA's share of the distinct experts' FC2 weights (16 KB pieces, expert-major, CTA-strided, so
+//! the experts the FC2 grid reads first are fetched first). Weights only: legal at any point of the FC1 grid.
+__device__ __forceinline__ void prefetchFc2Share(DraftRoutedMoeParams const& p, int const* distinct, int numDistinct)
+{
+    constexpr int kPiecesPerExpert = kFc2Chunks * kPrefetchPiecesPerBlock;
+    int const pieces = numDistinct * kPiecesPerExpert;
+    for (int piece = blockIdx.x; piece < pieces; piece += gridDim.x)
+    {
+        uint8_t const* addr
+            = reinterpret_cast<uint8_t const*>(
+                  p.w2 + static_cast<int64_t>(distinct[piece / kPiecesPerExpert]) * kFc2Chunks * kHidden * kKChunk)
+            + static_cast<int64_t>(piece % kPiecesPerExpert) * kPrefetchPieceBytes;
+        asm volatile(
+            "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(__cvta_generic_to_global(addr)), "r"(kPrefetchPieceBytes)
+            : "memory");
+    }
+}
+
+//! Fragments of one 64-element K chunk for C 16-row chains: A of k-steps 0..3 per chain, B of k-steps (0,1), (2,3).
+template <int C>
+struct ChunkFrags
+{
+    uint32_t a[C][4][4];
+    uint32_t b[2][4];
+};
+
+//! The four k-steps of a chunk in natural K order into each chain's own accumulator (the sequence of the kernels
+//! above; the chains are interleaved per k-step and never mixed).
+template <int C>
+__device__ __forceinline__ void mmaChunk(float (&acc)[C][4], ChunkFrags<C> const& f)
+{
+#pragma unroll
+    for (int q = 0; q < 4; ++q)
+    {
+#pragma unroll
+        for (int c = 0; c < C; ++c)
+        {
+            mmaBf16(acc[c], f.a[c][q], f.b[q >> 1][2 * (q & 1)], f.b[q >> 1][2 * (q & 1) + 1]);
+        }
+    }
+}
+
+//! Per-lane ldmatrix offsets inside one K chunk. ``a[q]``: the lane's row of k-step q in a 2 KB 16-row chunk (as
+//! loadA). ``b[h]``: the lane's 16-byte unit of k-steps (2h, 2h + 1) in the first 128-byte K chunk of a [token][k]
+//! tile with ``bRowBytes`` per token (as loadB2). The swizzle only permutes the 16-byte units inside a 128-byte K
+//! chunk, so K chunk kc of a token row is the same pattern ``kc * 128`` bytes further.
+struct LaneOffsets
+{
+    uint32_t a[4];
+    uint32_t b[2];
+};
+
+__device__ __forceinline__ LaneOffsets laneOffsets(int rows, int bRowBytes, int lane)
+{
+    LaneOffsets o;
+    int const m = lane >> 3;
+    int const r = (lane & 7) | ((m & 1) << 3);
+#pragma unroll
+    for (int q = 0; q < 4; ++q)
+    {
+        o.a[q] = r * 128 + swizzled(r, 2 * q + (m >> 1));
+    }
+    int tok = lane & 7;
+    tok = tok < rows ? tok : 0;
+#pragma unroll
+    for (int h = 0; h < 2; ++h)
+    {
+        o.b[h] = tok * bRowBytes + swizzled(tok, 4 * h + m);
+    }
+    return o;
+}
+
+//! Per-lane ldmatrix offsets of the B fragments in a streamed activation slice (kFlagStreamX): one TMA box of
+//! ``DraftRoutedMoeTensorMaps::x`` is [8 tokens][kXBoxChunks K chunks][128 bytes] with the 128-byte swizzle, so K chunk
+//! p of token t is line 2t + p of the box. ``b[p][h]``: the lane's 16-byte unit of k-steps (2h, 2h + 1) of K chunk p.
+struct LaneOffsetsX
+{
+    uint32_t b[2][2];
+};
+
+__device__ __forceinline__ LaneOffsetsX laneOffsetsX(int rows, int lane)
+{
+    static_assert(kXBoxChunks == 2, "two K chunks per activation box");
+    LaneOffsetsX o;
+    int const m = lane >> 3;
+    int tok = lane & 7;
+    tok = tok < rows ? tok : 0;
+#pragma unroll
+    for (int q = 0; q < 2; ++q)
+    {
+        int const line = kXBoxChunks * tok + q;
+#pragma unroll
+        for (int h = 0; h < 2; ++h)
+        {
+            o.b[q][h] = line * 128 + swizzled(line, 4 * h + m);
+        }
+    }
+    return o;
+}
+
+//! The fragments of one K chunk for C 16-row chains (chain c's 2 KB chunk follows chain 0's at c * kChunkBytes)
+//! from per-stage bases and constant chunk offsets, so the loads of a stage need no address arithmetic.
+//! ``aBase[q]``: chain 0 of the stage's first chunk + LaneOffsets::a[q]; ``bBase[h]``: the activation K chunk of the
+//! stage's first chunk + LaneOffsets::b[h].
+template <int C>
+__device__ __forceinline__ void loadFragsAt(
+    ChunkFrags<C>& f, uint32_t const (&aBase)[4], uint32_t aOff, uint32_t const (&bBase)[2], uint32_t bOff)
+{
+#pragma unroll
+    for (int c = 0; c < C; ++c)
+    {
+#pragma unroll
+        for (int q = 0; q < 4; ++q)
+        {
+            ldmatrixX4(f.a[c][q], aBase[q] + aOff + c * kChunkBytes);
+        }
+    }
+#pragma unroll
+    for (int h = 0; h < 2; ++h)
+    {
+        ldmatrixX4(f.b[h], bBase[h] + bOff);
+    }
+}
+
+//! loadFragsAt for chains ``kChainStride`` bytes apart (the sub-tile boxes of the triples path); chains ``c >= valid``
+//! (absent sub-tiles of the grid's last CTA) are not loaded.
+template <int C, int kChainStride>
+__device__ __forceinline__ void loadFragsStrided(
+    ChunkFrags<C>& f, uint32_t const (&aBase)[4], uint32_t aOff, uint32_t const (&bBase)[2], uint32_t bOff, int valid)
+{
+#pragma unroll
+    for (int c = 0; c < C; ++c)
+    {
+        if (c < valid)
+        {
+#pragma unroll
+            for (int q = 0; q < 4; ++q)
+            {
+                ldmatrixX4(f.a[c][q], aBase[q] + aOff + c * kChainStride);
+            }
+        }
+    }
+#pragma unroll
+    for (int h = 0; h < 2; ++h)
+    {
+        ldmatrixX4(f.b[h], bBase[h] + bOff);
+    }
+}
+
+//! mmaChunk for the chains ``c < valid`` only (each chain's k-steps in natural order, as mmaChunk).
+template <int C>
+__device__ __forceinline__ void mmaChunkValid(float (&acc)[C][4], ChunkFrags<C> const& f, int valid)
+{
+#pragma unroll
+    for (int q = 0; q < 4; ++q)
+    {
+#pragma unroll
+        for (int c = 0; c < C; ++c)
+        {
+            if (c < valid)
+            {
+                mmaBf16(acc[c], f.a[c][q], f.b[q >> 1][2 * (q & 1)], f.b[q >> 1][2 * (q & 1) + 1]);
+            }
+        }
+    }
+}
+
+//! Fragments of one K chunk for E FC2 experts of one 16-row tile: per expert its A (weights) and B (intermediate).
+template <int E>
+struct ExpertFrags
+{
+    uint32_t a[E][4][4];
+    uint32_t b[E][2][4];
+};
+
+//! Byte offset of element ``j`` of token ``tok`` in an intermediate entry stored as its shared-memory image (512-byte
+//! token rows, 16-byte units swizzled by token as loadB2 reads them).
+__device__ __forceinline__ int hImageElem(int tok, int j)
+{
+    int const u = j >> 3;
+    return tok * kInter + (((u & ~7) | ((u & 7) ^ (tok & 7))) << 3) + (j & 7);
+}
+
+//! TMA-fed FC1 kernel: W MMA warps with C chains each, a ring of S stages of 32 KB (one box of ``map``). With S3 > 0
+//! (4 MMA warps, kFlagFc1Triples) a call whose tiles exceed the grid but whose 64-row sub-tiles fit three per CTA
+//! runs the triples path instead (decided per call from the routing, identically in every CTA): CTA b takes sub-tiles
+//! 3b .. 3b + 2 (in distinct-expert order), MMA warp w owns rows 16w .. 16w + 15 of each (three chains), and a stage
+//! of its ring of S3 is three boxes of ``map3`` (64 rows x kTripleStageChunks K chunks), one per sub-tile.
+//! With X (kFlagStreamX) the activation rows are not kept in shared memory: every stage also carries the activation
+//! K chunks it needs (boxes of ``mapX``, kXBoxBytes each, after the weight boxes), so the ring can be deeper.
+template <int W, int C, int S, int S3, bool X = false>
+__global__ void __launch_bounds__((W + 1) * 32)
+    mtpMoeFc1TmaKernel(DraftRoutedMoeParams const p, __grid_constant__ CUtensorMap const map,
+        __grid_constant__ CUtensorMap const map3, __grid_constant__ CUtensorMap const mapX)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    constexpr int kChains = W * C;
+    constexpr int kKs = kTmaStageChunks / kChains; // K chunks per stage (one TMA box)
+    constexpr int kChunkSet = kChains * kChunkBytes;
+    constexpr int kStage = kKs * kChunkSet;
+    constexpr int kTileRows = 16 * kChains;
+    constexpr int kTilesPerExpert = kFc1Rows / kTileRows;
+    constexpr int kStagesPerTile = kFc1Chunks / kKs;
+    constexpr int kL2Ahead = 2 * S;
+    static_assert(kFc1Chunks % kKs == 0, "a tile must be a whole number of stages");
+    static_assert(kKs % 2 == 0, "the fragment double buffer restarts at f[0] on every stage");
+    // Triples path (S3 > 0): sub-tiles of 16 W rows, three per CTA; a stage holds one box per sub-tile.
+    constexpr int kSubRows = 16 * W;
+    constexpr int kSubPerExpert = kFc1Rows / kSubRows;
+    constexpr int kChunkStride3 = kSubRows * kKChunk * 2; // next K chunk inside a sub-tile box
+    constexpr int kBox3 = kTripleStageChunks * kChunkStride3;
+    constexpr int kStage3 = kTripleSubTiles * kBox3;
+    constexpr int kStagesPerTile3 = kFc1Chunks / kTripleStageChunks;
+    static_assert(S3 == 0 || (W == 4 && kSubRows == kTripleRows), "the triples path has 4 MMA warps (64-row boxes)");
+    static_assert(kTripleStageChunks == 2, "the triples consumer walks two K chunks per stage");
+    // Ring slot strides: the weight boxes, then (X) the activation boxes of the stage's K chunks.
+    constexpr int kXBoxes = X ? kKs / kXBoxChunks : 0;
+    constexpr int kSlot = kStage + kXBoxes * kXBoxBytes;
+    constexpr int kSlot3 = kStage3 + (X ? kXBoxBytes : 0);
+    static_assert(!X || (kKs % kXBoxChunks == 0 && kTripleStageChunks == kXBoxChunks), "whole activation boxes");
+    static_assert(kSlot % 1024 == 0 && kSlot3 % 1024 == 0, "slots stay aligned to the 128-byte swizzle pattern");
+    constexpr int kRingBytes = S * kSlot > S3 * kSlot3 ? S * kSlot : S3 * kSlot3;
+    constexpr int kBars = S > S3 ? S : S3;
+    extern __shared__ uint8_t smemRaw[];
+    __shared__ RouteSmem rs;
+    __shared__ __align__(8) uint64_t bars[2 * kBars];
+    int64_t const t0 = kTimeline ? clock64() : 0;
+    int64_t const entryNs = kTimeline ? globalTimerNs() : 0;
+    int const warp = threadIdx.x >> 5;
+    int const lane = threadIdx.x & 31;
+    int const rows = p.rows;
+    uint32_t const ringAddr = (smemAddr(smemRaw) + 1023u) & ~1023u;
+    uint32_t const xsAddr = ringAddr + kRingBytes;
+    uint32_t const fullBar = smemAddr(&bars[0]);
+    uint32_t const emptyBar = smemAddr(&bars[kBars]);
+    if (threadIdx.x == 0)
+    {
+        for (int s = 0; s < kBars; ++s)
+        {
+            mbarInit(fullBar + 8 * s, 1);
+            mbarInit(emptyBar + 8 * s, W * 32); // every MMA thread releases its own reads
+        }
+        mbarFenceInit();
+        prefetchTensorMap(&map);
+        if constexpr (S3 > 0)
+        {
+            prefetchTensorMap(&map3);
+        }
+        if constexpr (X)
+        {
+            prefetchTensorMap(&mapX);
+        }
+    }
+
+    // Every input below (logits, activation) is produced by the predecessor grid.
+    gridDependencyWait();
+    int64_t const tWait = kTimeline ? clock64() - t0 : 0;
+
+    // The MMA warps stage the activation rows (they alone read them); the producer warp does not wait for them.
+    for (int i = threadIdx.x; i < (X || warp == W ? 0 : rows * (kXRowBytes / 16)); i += W * 32)
+    {
+        int const t = i / (kXRowBytes / 16);
+        int const u = i % (kXRowBytes / 16);
+        cpAsync16(xsAddr + t * kXRowBytes + swizzled(t, u), p.x + t * kHidden + u * 8);
+    }
+    cpAsyncCommit();
+
+    int64_t tTokens = 0;
+    // Ends with __syncthreads (also publishes the barrier initialization).
+    routeAll<W + 1>(p, rs, warp, lane, kTimeline ? &tTokens : nullptr);
+    int64_t const tRouted = kTimeline ? clock64() - t0 : 0;
+    int const numDistinct = rs.numDistinct;
+    if (blockIdx.x == 0)
+    {
+        for (int i = threadIdx.x; i < kRouteWords; i += (W + 1) * 32)
+        {
+            int v = 0;
+            if (i == 0)
+            {
+                v = numDistinct;
+            }
+            else if (i < kRouteSlotOffset)
+            {
+                v = i - kRouteDistinctOffset < numDistinct ? rs.distinct[i - kRouteDistinctOffset] : -1;
+            }
+            else if (i < kRouteWeightOffset)
+            {
+                int const j = i - kRouteSlotOffset;
+                v = j / kTopK < rows ? rs.slotDistinct[j / kTopK][j % kTopK] : -1;
+            }
+            else
+            {
+                int const j = i - kRouteWeightOffset;
+                v = j / kTopK < rows ? __float_as_int(rs.topWeight[j / kTopK][j % kTopK]) : 0;
+            }
+            p.route[i] = v;
+        }
+    }
+    int const tiles = numDistinct * kTilesPerExpert;
+    bool const noMma = kTimeline && (p.flags & kFlagProbeNoMma) != 0;
+    bool const resident = kTimeline && (p.flags & kFlagProbeResident) != 0;
+    // The path depends only on numDistinct and the launch, so every CTA takes the same one.
+    int const numSub = numDistinct * kSubPerExpert;
+    int const sub0 = kTripleSubTiles * static_cast<int>(blockIdx.x);
+    bool const triples = S3 > 0 && (p.flags & kFlagFc1Triples) != 0 && !resident && tiles > static_cast<int>(gridDim.x)
+        && numSub <= kTripleSubTiles * static_cast<int>(gridDim.x);
+    int const valid3 = min(kTripleSubTiles, numSub - sub0); // sub-tiles of this CTA (<= 0: none)
+    // Early trigger (after routing): always with kFlagEarlyTrigger, or for calls on the triples path only.
+    bool const early = (p.flags & kFlagEarlyTrigger) != 0 || (triples && (p.flags & kFlagEarlyTriggerTriples) != 0);
+    if (early)
+    {
+        gridDependencyTrigger();
+    }
+    bool const evictFirst = (p.flags & kFlagFc1EvictFirst) != 0;
+
+    if (warp == W)
+    {
+        if constexpr (S3 > 0)
+        {
+            if (lane == 0 && triples)
+            {
+                uint64_t const policy = evictFirst ? l2EvictFirstPolicy() : 0;
+                if (p.flags & kFlagPrefetchFc2Early)
+                {
+                    prefetchFc2Share(p, rs.distinct, numDistinct);
+                }
+                int emptyBlocked = 0;
+                int64_t tFirst = 0;
+                for (int u = 0; u < (valid3 > 0 ? kStagesPerTile3 : 0); ++u)
+                {
+                    int const s = u % S3;
+                    if (u >= S3)
+                    {
+                        if constexpr (kTimeline)
+                        {
+                            emptyBlocked += mbarWaitCounted(emptyBar + 8 * s, ((u / S3) - 1) & 1);
+                        }
+                        else
+                        {
+                            mbarWait(emptyBar + 8 * s, ((u / S3) - 1) & 1);
+                        }
+                    }
+                    mbarArriveExpectTx(fullBar + 8 * s, valid3 * kBox3 + (X ? kXBoxBytes : 0));
+                    if constexpr (X)
+                    {
+                        tmaLoad3d(
+                            ringAddr + s * kSlot3 + kStage3, &mapX, 0, u * kTripleStageChunks, 0, fullBar + 8 * s);
+                    }
+                    for (int c = 0; c < valid3; ++c)
+                    {
+                        int const sub = sub0 + c;
+                        uint32_t const dst = ringAddr + s * kSlot3 + c * kBox3;
+                        int const row = (sub % kSubPerExpert) * kSubRows;
+                        int const slab = rs.distinct[sub / kSubPerExpert] * kFc1Chunks + u * kTripleStageChunks;
+                        if (evictFirst)
+                        {
+                            tmaLoad3dHint(dst, &map3, 0, row, slab, fullBar + 8 * s, policy);
+                        }
+                        else
+                        {
+                            tmaLoad3d(dst, &map3, 0, row, slab, fullBar + 8 * s);
+                        }
+                    }
+                    if (kTimeline && u == 0)
+                    {
+                        tFirst = clock64() - t0;
+                    }
+                }
+                int64_t const tLast = kTimeline ? clock64() - t0 : 0;
+                if ((p.flags & kFlagPrefetchFc2) && !(p.flags & kFlagPrefetchFc2Early))
+                {
+                    prefetchFc2Share(p, rs.distinct, numDistinct);
+                }
+                if (int64_t* row = timelineRow(p, 0))
+                {
+                    row[kTlProdFirst] = tFirst;
+                    row[kTlProdLast] = tLast;
+                    row[kTlEmptyBlocked] = emptyBlocked;
+                }
+            }
+        }
+        if (lane == 0 && !triples)
+        {
+            int const myTiles
+                = static_cast<int>(blockIdx.x) < tiles ? (tiles - 1 - static_cast<int>(blockIdx.x)) / gridDim.x + 1 : 0;
+            int const total = myTiles * kStagesPerTile;
+            // Box coordinates of this CTA's u-th stage: row of the tile, (expert, first K chunk) slab.
+            auto rowOf = [&](int u)
+            {
+                int const t = blockIdx.x + (u / kStagesPerTile) * gridDim.x;
+                return (t % kTilesPerExpert) * kTileRows;
+            };
+            auto slabOf = [&](int u)
+            {
+                int const t = blockIdx.x + (u / kStagesPerTile) * gridDim.x;
+                return rs.distinct[t / kTilesPerExpert] * kFc1Chunks + (u % kStagesPerTile) * kKs;
+            };
+            bool const l2 = (p.flags & kFlagL2Prefetch) != 0;
+            uint64_t const policy = evictFirst ? l2EvictFirstPolicy() : 0;
+            if (p.flags & kFlagPrefetchFc2Early)
+            {
+                prefetchFc2Share(p, rs.distinct, numDistinct);
+            }
+            if (l2)
+            {
+                for (int u = 0; u < kL2Ahead && u < total; ++u)
+                {
+                    tmaPrefetchL2(&map, 0, rowOf(u), slabOf(u));
+                }
+            }
+            int emptyBlocked = 0;
+            int64_t tFirst = 0;
+            for (int u = 0; u < total && !(resident && u == S); ++u)
+            {
+                if (l2 && u + kL2Ahead < total)
+                {
+                    tmaPrefetchL2(&map, 0, rowOf(u + kL2Ahead), slabOf(u + kL2Ahead));
+                }
+                int const s = u % S;
+                if (u >= S)
+                {
+                    if constexpr (kTimeline)
+                    {
+                        emptyBlocked += mbarWaitCounted(emptyBar + 8 * s, ((u / S) - 1) & 1);
+                    }
+                    else
+                    {
+                        mbarWait(emptyBar + 8 * s, ((u / S) - 1) & 1);
+                    }
+                }
+                mbarArriveExpectTx(fullBar + 8 * s, kSlot);
+                if constexpr (X)
+                {
+#pragma unroll
+                    for (int b = 0; b < kXBoxes; ++b)
+                    {
+                        tmaLoad3d(ringAddr + s * kSlot + kStage + b * kXBoxBytes, &mapX, 0,
+                            (u % kStagesPerTile) * kKs + b * kXBoxChunks, 0, fullBar + 8 * s);
+                    }
+                }
+                if (evictFirst)
+                {
+                    tmaLoad3dHint(ringAddr + s * kSlot, &map, 0, rowOf(u), slabOf(u), fullBar + 8 * s, policy);
+                }
+                else
+                {
+                    tmaLoad3d(ringAddr + s * kSlot, &map, 0, rowOf(u), slabOf(u), fullBar + 8 * s);
+                }
+                if (kTimeline && u == 0)
+                {
+                    tFirst = clock64() - t0;
+                }
+            }
+            int64_t const tLast = kTimeline ? clock64() - t0 : 0;
+            if ((p.flags & kFlagPrefetchFc2) && !(p.flags & kFlagPrefetchFc2Early))
+            {
+                prefetchFc2Share(p, rs.distinct, numDistinct);
+            }
+            if (int64_t* row = timelineRow(p, 0))
+            {
+                row[kTlProdFirst] = tFirst;
+                row[kTlProdLast] = tLast;
+                row[kTlEmptyBlocked] = emptyBlocked;
+            }
+        }
+    }
+    else
+    {
+        cpAsyncWait<0>();
+        mmaWarpsSync<W>(); // the activation rows of every MMA thread are in shared memory
+        int64_t const tSetup = kTimeline ? clock64() - t0 : 0;
+        int n = 0;         // K chunks consumed by this warp (all tiles)
+        int fullBlocked = 0;
+        int64_t tFirst = 0;
+        int64_t tLast = 0;
+        // Stages are used in ring order; a tile starts on a stage boundary (kFc1Chunks % kKs == 0). The fragment
+        // bases point at the stage of ring slot ``fetchSlot`` (``aBase``) and at the K chunk of the activation rows
+        // the current stage starts with (``bBase``); both move by constant steps, so no lane offset stays live.
+        int fetchSlot = 0;
+        uint32_t fetchPhase = 0;
+        int releaseSlot = 0;
+        int fetchUses = 0;
+        uint32_t aBase[4];
+        uint32_t bBase[2];
+        uint32_t bX[2][2]; // X: K chunk parity p of the activation boxes of ring slot ``fetchSlot``
+        {
+            LaneOffsets const lo = laneOffsets(rows, kXRowBytes, lane);
+            LaneOffsetsX const lx = laneOffsetsX(rows, lane);
+#pragma unroll
+            for (int q = 0; q < 4; ++q)
+            {
+                aBase[q] = ringAddr + warp * C * kChunkBytes + lo.a[q];
+            }
+#pragma unroll
+            for (int h = 0; h < 2; ++h)
+            {
+                bBase[h] = xsAddr + lo.b[h];
+                bX[0][h] = ringAddr + kStage + lx.b[0][h];
+                bX[1][h] = ringAddr + kStage + lx.b[1][h];
+            }
+        }
+        // Waits for the stage of ring slot ``fetchSlot`` unless an earlier probe found it complete (``ready``).
+        auto waitStage = [&](bool ready)
+        {
+            if constexpr (kTimeline)
+            {
+                bool const wait = !resident || fetchUses++ < S;
+                if (wait && !ready)
+                {
+                    fullBlocked += mbarWaitCounted(fullBar + 8 * fetchSlot, fetchPhase);
+                }
+                tLast = clock64() - t0;
+                tFirst = tFirst == 0 ? tLast : tFirst;
+            }
+            else if (!ready)
+            {
+                mbarWait(fullBar + 8 * fetchSlot, fetchPhase);
+            }
+        };
+        // Moves the A bases (X: and the activation bases) to the next ring slot.
+        auto nextSlot = [&]()
+        {
+            int const step = fetchSlot + 1 == S ? -(S - 1) * kSlot : kSlot;
+#pragma unroll
+            for (int q = 0; q < 4; ++q)
+            {
+                aBase[q] += step;
+            }
+            if constexpr (X)
+            {
+#pragma unroll
+                for (int h = 0; h < 2; ++h)
+                {
+                    bX[0][h] += step;
+                    bX[1][h] += step;
+                }
+            }
+            if (++fetchSlot == S)
+            {
+                fetchSlot = 0;
+                fetchPhase ^= 1;
+            }
+        };
+        if constexpr (S3 > 0)
+        {
+            if (triples && valid3 > 0)
+            {
+                // Chain c of this warp: rows 16 warp .. 16 warp + 15 of sub-tile sub0 + c, in box c of each stage.
+                uint32_t a3[4];
+                uint32_t b3[2];
+                uint32_t b3X[2][2]; // X: K chunk 0 / 1 of the stage's activation box
+                {
+                    LaneOffsets const lo = laneOffsets(rows, kXRowBytes, lane);
+                    LaneOffsetsX const lx = laneOffsetsX(rows, lane);
+#pragma unroll
+                    for (int q = 0; q < 4; ++q)
+                    {
+                        a3[q] = ringAddr + warp * kChunkBytes + lo.a[q];
+                    }
+#pragma unroll
+                    for (int h = 0; h < 2; ++h)
+                    {
+                        b3[h] = xsAddr + lo.b[h];
+                        b3X[0][h] = ringAddr + kStage3 + lx.b[0][h];
+                        b3X[1][h] = ringAddr + kStage3 + lx.b[1][h];
+                    }
+                }
+                // The fragments of K chunk j (0 / 1) of the current stage.
+                auto frags3 = [&](ChunkFrags<kTripleSubTiles>& f, int j)
+                {
+                    if constexpr (X)
+                    {
+                        loadFragsStrided<kTripleSubTiles, kBox3>(f, a3, j * kChunkStride3, b3X[j], 0, valid3);
+                    }
+                    else
+                    {
+                        loadFragsStrided<kTripleSubTiles, kBox3>(f, a3, j * kChunkStride3, b3, j * kKChunk * 2, valid3);
+                    }
+                };
+                float acc[kTripleSubTiles][4];
+#pragma unroll
+                for (int c = 0; c < kTripleSubTiles; ++c)
+                {
+                    acc[c][0] = acc[c][1] = acc[c][2] = acc[c][3] = -0.f;
+                }
+                ChunkFrags<kTripleSubTiles> f[2];
+                waitStage(false);
+                if (!noMma)
+                {
+                    frags3(f[0], 0);
+                }
+#pragma unroll 1
+                for (int st = 0; st < kStagesPerTile3; ++st)
+                {
+                    // K chunk 0 of the stage: fetch chunk 1, MMAs of chunk 0.
+                    if (!noMma)
+                    {
+                        frags3(f[1], 1);
+                        mmaChunkValid<kTripleSubTiles>(acc, f[0], valid3);
+                    }
+                    if (st + 1 < kStagesPerTile3)
+                    {
+                        // K chunk 1: move to the next stage, probe it, MMAs of chunk 1, then wait only if the probe
+                        // failed and fetch the next stage's chunk 0.
+                        int const step = fetchSlot + 1 == S3 ? -(S3 - 1) * kSlot3 : kSlot3;
+#pragma unroll
+                        for (int q = 0; q < 4; ++q)
+                        {
+                            a3[q] += step;
+                        }
+                        if (++fetchSlot == S3)
+                        {
+                            fetchSlot = 0;
+                            fetchPhase ^= 1;
+                        }
+#pragma unroll
+                        for (int h = 0; h < 2; ++h)
+                        {
+                            b3[h] += kTripleStageChunks * kKChunk * 2;
+                            b3X[0][h] += step;
+                            b3X[1][h] += step;
+                        }
+                        bool const ready = mbarTestWait(fullBar + 8 * fetchSlot, fetchPhase);
+                        if (!noMma)
+                        {
+                            mmaChunkValid<kTripleSubTiles>(acc, f[1], valid3);
+                        }
+                        waitStage(ready);
+                        if (!noMma)
+                        {
+                            frags3(f[0], 0);
+                        }
+                    }
+                    else if (!noMma)
+                    {
+                        mmaChunkValid<kTripleSubTiles>(acc, f[1], valid3);
+                    }
+                    // Every thread releases its own reads of the stage (their fragments fed the MMAs above).
+                    mbarArrive(emptyBar + 8 * releaseSlot);
+                    releaseSlot = releaseSlot + 1 == S3 ? 0 : releaseSlot + 1;
+                    n += kTripleStageChunks;
+                }
+                int const g = lane >> 2;
+                int const tq = lane & 3;
+#pragma unroll
+                for (int c = 0; c < kTripleSubTiles; ++c)
+                {
+                    if (c < valid3)
+                    {
+                        int const sub = sub0 + c;
+                        int const d = sub / kSubPerExpert;
+                        int const a = (sub % kSubPerExpert) * W + warp;
+                        int const j = 16 * (a >> 1) + 2 * g + (a & 1);
+#pragma unroll
+                        for (int h = 0; h < 2; ++h)
+                        {
+                            int const tok = 2 * tq + h;
+                            if (tok < rows)
+                            {
+                                p.h[static_cast<int64_t>(d) * kMaxRows * kInter + hImageElem(tok, j)]
+                                    = bf16SatFinite(swiGlu(acc[c][h], acc[c][2 + h]));
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        // The regular path (no tile when the triples path ran).
+        for (int t = triples ? tiles : static_cast<int>(blockIdx.x); t < tiles; t += gridDim.x)
+        {
+            float acc[C][4];
+#pragma unroll
+            for (int c = 0; c < C; ++c)
+            {
+                acc[c][0] = acc[c][1] = acc[c][2] = acc[c][3] = -0.f;
+            }
+            // The fragments of chunk j + 1 are loaded before the MMAs of chunk j, across stage boundaries too (the
+            // next stage is waited for at the last chunk of the current one); f[0] holds a stage's first chunk.
+            ChunkFrags<C> f[2];
+            // The fragments of K chunk j of the current stage (X: chunk j is in activation box j / 2, parity j % 2).
+            auto frags = [&](ChunkFrags<C>& fr, int j)
+            {
+                if constexpr (X)
+                {
+                    loadFragsAt<C>(fr, aBase, j * kChunkSet, bX[j & 1], (j >> 1) * kXBoxBytes);
+                }
+                else
+                {
+                    loadFragsAt<C>(fr, aBase, j * kChunkSet, bBase, j * 128);
+                }
+            };
+            if (t != static_cast<int>(blockIdx.x))
+            {
+                nextSlot();
+#pragma unroll
+                for (int h = 0; h < 2; ++h)
+                {
+                    bBase[h] -= (kStagesPerTile - 1) * kKs * 128;
+                }
+            }
+            waitStage(false);
+            if (!noMma)
+            {
+                frags(f[0], 0);
+            }
+#pragma unroll 1
+            for (int st = 0; st < kStagesPerTile; ++st)
+            {
+#pragma unroll
+                for (int j = 0; j < kKs; ++j)
+                {
+                    if (j + 1 < kKs)
+                    {
+                        if (!noMma)
+                        {
+                            frags(f[(j + 1) & 1], j + 1);
+                            mmaChunk<C>(acc, f[j & 1]);
+                        }
+                    }
+                    else if (st + 1 < kStagesPerTile)
+                    {
+                        // Last chunk of the stage: probe the next stage, issue this chunk's MMAs, then wait only if
+                        // the probe failed, so the barrier round trip overlaps the MMAs.
+                        nextSlot();
+#pragma unroll
+                        for (int h = 0; h < 2; ++h)
+                        {
+                            bBase[h] += kKs * 128;
+                        }
+                        bool const ready = mbarTestWait(fullBar + 8 * fetchSlot, fetchPhase);
+                        if (!noMma)
+                        {
+                            mmaChunk<C>(acc, f[j & 1]);
+                        }
+                        waitStage(ready);
+                        if (!noMma)
+                        {
+                            frags(f[0], 0);
+                        }
+                    }
+                    else if (!noMma)
+                    {
+                        mmaChunk<C>(acc, f[j & 1]);
+                    }
+                }
+                // Every thread releases its own reads of the stage (their fragments fed the MMAs above).
+                mbarArrive(emptyBar + 8 * releaseSlot);
+                releaseSlot = releaseSlot + 1 == S ? 0 : releaseSlot + 1;
+                n += kKs;
+            }
+            int const d = t / kTilesPerExpert;
+            int const g = lane >> 2;
+            int const tq = lane & 3;
+#pragma unroll
+            for (int c = 0; c < C; ++c)
+            {
+                int const a = (t % kTilesPerExpert) * kChains + warp * C + c;
+                int const j = 16 * (a >> 1) + 2 * g + (a & 1);
+#pragma unroll
+                for (int h = 0; h < 2; ++h)
+                {
+                    int const tok = 2 * tq + h;
+                    if (tok < rows)
+                    {
+                        p.h[static_cast<int64_t>(d) * kMaxRows * kInter + hImageElem(tok, j)]
+                            = bf16SatFinite(swiGlu(acc[c][h], acc[c][2 + h]));
+                    }
+                }
+            }
+        }
+        if (warp == 0 && lane == 0)
+        {
+            if (int64_t* row = timelineRow(p, 0))
+            {
+                int64_t const cycles = clock64() - t0;
+                row[kTlEntryNs] = entryNs;
+                row[kTlDoneNs] = globalTimerNs();
+                row[kTlSmId] = smId();
+                row[kTlWait] = tWait;
+                row[kTlSetup] = tSetup;
+                row[kTlFirst] = tFirst;
+                row[kTlLast] = tLast;
+                row[kTlDone] = cycles;
+                row[kTlFullBlocked] = fullBlocked;
+                row[kTlChunks] = n;
+                row[kTlRouteTokens] = tTokens != 0 ? tTokens - t0 : 0;
+                row[kTlRouted] = tRouted;
+            }
+        }
+    }
+    if (!early)
+    {
+        // A CTA counts as triggered as soon as one of its threads triggers: wait for the MMA warps, or the idle
+        // producer lanes would release the dependent grid onto free SMs right after routing.
+        __syncthreads();
+        gridDependencyTrigger();
+    }
+#endif
+}
+
+//! FC2 ring depth for W MMA warps (stages of W * 8 KB weights + one intermediate entry, about 140 KB in all).
+template <int W>
+constexpr int fc2TmaStages()
+{
+    return W == 4 ? 4 : (W == 3 ? 5 : (W == 2 ? 7 : 12));
+}
+
+template <int W, int S, int E>
+__global__ void __launch_bounds__((W + 1) * 32)
+    mtpMoeFc2TmaKernel(DraftRoutedMoeParams const p, __grid_constant__ CUtensorMap const map)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    static_assert(E == 1 || E == 2, "one or two experts per step");
+    static_assert(S >= 2 * E, "the ring must hold the step being consumed and the next one");
+    // One stage per (group, expert): the expert's W * 16 rows x K = 256 weight box (one 3D TMA) and its
+    // intermediate entry (one bulk copy), both counted on the stage's full barrier.
+    constexpr int kWBytes = W * kFc2Chunks * kChunkBytes;
+    constexpr int kStage = kWBytes + kHEntryBytes;
+    constexpr int kGroups = kFc2Tiles / W;
+    static_assert(kWBytes % 1024 == 0 && kHEntryBytes % 1024 == 0, "stages must stay 1 KB aligned");
+    extern __shared__ uint8_t smemRaw[];
+    __shared__ Fc2Smem fs;
+    __shared__ __align__(8) uint64_t bars[2 * S];
+    int64_t const t0 = kTimeline ? clock64() : 0;
+    int64_t const entryNs = kTimeline ? globalTimerNs() : 0;
+    int const warp = threadIdx.x >> 5;
+    int const lane = threadIdx.x & 31;
+    int const rows = p.rows;
+    uint32_t const ringAddr = (smemAddr(smemRaw) + 1023u) & ~1023u;
+    uint8_t* const stageBase = smemRaw + (ringAddr - smemAddr(smemRaw)) + S * kStage;
+    uint32_t const fullBar = smemAddr(&bars[0]);
+    uint32_t const emptyBar = smemAddr(&bars[S]);
+    if (threadIdx.x == 0)
+    {
+        for (int s = 0; s < S; ++s)
+        {
+            mbarInit(fullBar + 8 * s, 1);
+            mbarInit(emptyBar + 8 * s, W * 32); // every MMA thread releases its own reads
+        }
+        mbarFenceInit();
+        prefetchTensorMap(&map);
+    }
+
+    // The routing record and the intermediate are written by the FC1 grid.
+    gridDependencyWait();
+    int64_t const tWait = kTimeline ? clock64() - t0 : 0;
+    bool const noMma = kTimeline && (p.flags & kFlagProbeNoMma) != 0;
+    // Publishes the barrier initialization. Past this point the producer only needs the distinct experts and the MMA
+    // warps only the slot and weight tables, so each role reads its part of the routing record by itself.
+    __syncthreads();
+    int const numDistinct = p.route[0];
+
+    if (warp == W)
+    {
+        for (int i = lane; i < kMaxDistinct; i += 32)
+        {
+            fs.distinct[i] = p.route[kRouteDistinctOffset + i];
+        }
+        __syncwarp();
+        if (lane == 0)
+        {
+            // The intermediate was written with generic stores by the previous grid; bulk copies read it.
+            asm volatile("fence.proxy.async.global;" ::: "memory");
+            int u = 0;
+            int emptyBlocked = 0;
+            int64_t tFirst = 0;
+            for (int gi = blockIdx.x; gi < kGroups; gi += gridDim.x)
+            {
+                for (int d = 0; d < numDistinct; ++d, ++u)
+                {
+                    int const s = u % S;
+                    if (u >= S)
+                    {
+                        if constexpr (kTimeline)
+                        {
+                            emptyBlocked += mbarWaitCounted(emptyBar + 8 * s, ((u / S) - 1) & 1);
+                        }
+                        else
+                        {
+                            mbarWait(emptyBar + 8 * s, ((u / S) - 1) & 1);
+                        }
+                    }
+                    mbarArriveExpectTx(fullBar + 8 * s, kWBytes + rows * kInter * 2);
+                    tmaLoad3d(
+                        ringAddr + s * kStage, &map, 0, gi * 16 * W, fs.distinct[d] * kFc2Chunks, fullBar + 8 * s);
+                    bulkLoad(ringAddr + s * kStage + kWBytes, p.h + static_cast<int64_t>(d) * kMaxRows * kInter,
+                        rows * kInter * 2, fullBar + 8 * s);
+                    if (kTimeline && u == 0)
+                    {
+                        tFirst = clock64() - t0;
+                    }
+                }
+            }
+            if (int64_t* row = timelineRow(p, 1))
+            {
+                row[kTlProdFirst] = tFirst;
+                row[kTlProdLast] = clock64() - t0;
+                row[kTlEmptyBlocked] = emptyBlocked;
+            }
+        }
+    }
+    else
+    {
+        // Slot and weight tables, shared by the W MMA warps only (named barrier 1); first read at the end of a step.
+        for (int i = threadIdx.x; i < kMaxRows * kMaxDistinct; i += W * 32)
+        {
+            fs.slotOf[i / kMaxDistinct][i % kMaxDistinct] = -1;
+        }
+        mmaWarpsSync<W>();
+        for (int i = threadIdx.x; i < rows * kTopK; i += W * 32)
+        {
+            int const t = i / kTopK;
+            int const k = i % kTopK;
+            fs.slotOf[t][p.route[kRouteSlotOffset + i]] = static_cast<int8_t>(k);
+            fs.weight[t][k] = __int_as_float(p.route[kRouteWeightOffset + i]);
+        }
+        mmaWarpsSync<W>();
+        int64_t const tSetup = kTimeline ? clock64() - t0 : 0;
+        __nv_bfloat16* stage = reinterpret_cast<__nv_bfloat16*>(stageBase + warp * kStageBytes);
+        int u0 = 0; // stage uses of the previous groups
+        int fullBlocked = 0;
+        int64_t tFirst = 0;
+        int64_t tLast = 0;
+        // A group's experts are taken E at a time (a step); step j holds experts E*j .. E*j + E - 1, the last step
+        // may hold fewer. Each expert has its own stage and accumulator; their MMAs are interleaved per k-step.
+        int const steps = (numDistinct + E - 1) / E;
+        int fetchSlot = 0; // stages are used in ring order, one per (group, expert)
+        uint32_t fetchPhase = 0;
+        int releaseSlot = 0;
+        // Fragment bases of the next ring slot to be fetched; they move by constant steps.
+        uint32_t slotA[4];
+        uint32_t slotB[2];
+        {
+            LaneOffsets const lo = laneOffsets(rows, kInter * 2, lane);
+#pragma unroll
+            for (int q = 0; q < 4; ++q)
+            {
+                slotA[q] = ringAddr + warp * kChunkBytes + lo.a[q];
+            }
+#pragma unroll
+            for (int h = 0; h < 2; ++h)
+            {
+                slotB[h] = ringAddr + kWBytes + lo.b[h];
+            }
+        }
+        uint32_t aBase[E][4];
+        uint32_t bBase[E][2];
+        // Points the fragment bases at the stages of step j (one per expert of the step, in ring order) and probes
+        // them; finishStep waits for those the probe did not find complete. A step past the last expert reads the
+        // step's first expert again (results discarded).
+        uint32_t pendBar[E];
+        uint32_t pendParity[E];
+        bool pendWait[E];
+        auto enterStep = [&](int j)
+        {
+#pragma unroll
+            for (int e = 0; e < E; ++e)
+            {
+                pendWait[e] = false;
+                if (E * j + e < numDistinct)
+                {
+                    pendBar[e] = fullBar + 8 * fetchSlot;
+                    pendParity[e] = fetchPhase;
+                    pendWait[e] = !mbarTestWait(pendBar[e], pendParity[e]);
+#pragma unroll
+                    for (int q = 0; q < 4; ++q)
+                    {
+                        aBase[e][q] = slotA[q];
+                    }
+#pragma unroll
+                    for (int h = 0; h < 2; ++h)
+                    {
+                        bBase[e][h] = slotB[h];
+                    }
+                    int const step = fetchSlot + 1 == S ? -(S - 1) * kStage : kStage;
+#pragma unroll
+                    for (int q = 0; q < 4; ++q)
+                    {
+                        slotA[q] += step;
+                    }
+#pragma unroll
+                    for (int h = 0; h < 2; ++h)
+                    {
+                        slotB[h] += step;
+                    }
+                    if (++fetchSlot == S)
+                    {
+                        fetchSlot = 0;
+                        fetchPhase ^= 1;
+                    }
+                }
+                else
+                {
+#pragma unroll
+                    for (int q = 0; q < 4; ++q)
+                    {
+                        aBase[e][q] = aBase[0][q];
+                    }
+#pragma unroll
+                    for (int h = 0; h < 2; ++h)
+                    {
+                        bBase[e][h] = bBase[0][h];
+                    }
+                }
+            }
+        };
+        auto finishStep = [&]()
+        {
+#pragma unroll
+            for (int e = 0; e < E; ++e)
+            {
+                if (pendWait[e])
+                {
+                    if constexpr (kTimeline)
+                    {
+                        ++fullBlocked;
+                    }
+                    mbarWait(pendBar[e], pendParity[e]);
+                }
+            }
+            if constexpr (kTimeline)
+            {
+                tLast = clock64() - t0;
+                tFirst = tFirst == 0 ? tLast : tFirst;
+            }
+        };
+        // K chunk kc of every expert of the current step.
+        auto load = [&](ExpertFrags<E>& f, int kc)
+        {
+#pragma unroll
+            for (int e = 0; e < E; ++e)
+            {
+#pragma unroll
+                for (int q = 0; q < 4; ++q)
+                {
+                    ldmatrixX4(f.a[e][q], aBase[e][q] + kc * W * kChunkBytes);
+                }
+#pragma unroll
+                for (int h = 0; h < 2; ++h)
+                {
+                    ldmatrixX4(f.b[e][h], bBase[e][h] + kc * 128);
+                }
+            }
+        };
+        for (int gi = blockIdx.x; gi < kGroups; gi += gridDim.x)
+        {
+            int const tile = gi * W + warp;
+            float acc[E][4];
+#pragma unroll
+            for (int e = 0; e < E; ++e)
+            {
+                acc[e][0] = acc[e][1] = acc[e][2] = acc[e][3] = -0.f;
+            }
+            // The fragments of chunk kc + 1 are loaded before the MMAs of chunk kc, across steps too (the next
+            // step's accumulators are independent); f[0] holds a step's first chunk.
+            ExpertFrags<E> f[2];
+            enterStep(0);
+            finishStep();
+            if (!noMma)
+            {
+                load(f[0], 0);
+            }
+#pragma unroll 1
+            for (int j = 0; j < steps; ++j)
+            {
+#pragma unroll
+                for (int kc = 0; kc < kFc2Chunks; ++kc)
+                {
+                    auto mmas = [&]()
+                    {
+                        if (!noMma)
+                        {
+#pragma unroll
+                            for (int q = 0; q < 4; ++q)
+                            {
+#pragma unroll
+                                for (int e = 0; e < E; ++e)
+                                {
+                                    ExpertFrags<E> const& fr = f[kc & 1];
+                                    mmaBf16(acc[e], fr.a[e][q], fr.b[e][q >> 1][2 * (q & 1)],
+                                        fr.b[e][q >> 1][2 * (q & 1) + 1]);
+                                }
+                            }
+                        }
+                    };
+                    if (kc + 1 < kFc2Chunks)
+                    {
+                        if (!noMma)
+                        {
+                            load(f[(kc + 1) & 1], kc + 1);
+                        }
+                        mmas();
+                    }
+                    else if (j + 1 < steps)
+                    {
+                        // Last chunk of the step: probe the next step's stages, issue this chunk's MMAs, then wait
+                        // only for stages the probe did not find complete.
+                        enterStep(j + 1);
+                        mmas();
+                        finishStep();
+                        if (!noMma)
+                        {
+                            load(f[0], 0);
+                        }
+                    }
+                    else
+                    {
+                        mmas();
+                    }
+                }
+                // Every thread releases its own reads of the step's stages (their fragments fed the MMAs above).
+#pragma unroll
+                for (int e = 0; e < E; ++e)
+                {
+                    if (E * j + e < numDistinct)
+                    {
+                        mbarArrive(emptyBar + 8 * releaseSlot);
+                        releaseSlot = releaseSlot + 1 == S ? 0 : releaseSlot + 1;
+                    }
+                }
+                // FC2 rows of each expert -> BF16, into the token's slot of the warp's staging area.
+                int const g = lane >> 2;
+                int const tq = lane & 3;
+#pragma unroll
+                for (int e = 0; e < E; ++e)
+                {
+                    int const d = E * j + e;
+#pragma unroll
+                    for (int c = 0; c < 2; ++c)
+                    {
+                        int const tok = 2 * tq + c;
+                        int const k = (d < numDistinct && tok < rows) ? fs.slotOf[tok][d] : -1;
+                        if (k >= 0)
+                        {
+                            __nv_bfloat16 const lo16 = bf16SatFinite(acc[e][c]);
+                            __nv_bfloat16 const hi16 = bf16SatFinite(acc[e][2 + c]);
+                            stage[(tok * kTopK + k) * 16 + g] = lo16;
+                            stage[(tok * kTopK + k) * 16 + g + 8] = hi16;
+                            if (p.debugFc2 != nullptr)
+                            {
+                                int64_t const base = (static_cast<int64_t>(tok) * kTopK + k) * kHidden;
+                                p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g)] = lo16;
+                                p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g + 8)] = hi16;
+                            }
+                        }
+                    }
+                    acc[e][0] = acc[e][1] = acc[e][2] = acc[e][3] = -0.f;
+                }
+            }
+            u0 += numDistinct;
+            // Native finalize of the tile's 16 hidden rows for every token: FP32 fma chain from 0 in slot order.
+            __syncwarp();
+            for (int pr = lane; pr < rows * 16; pr += 32)
+            {
+                int const tok = pr >> 4;
+                int const row = pr & 15;
+                float sum = 0.f;
+#pragma unroll
+                for (int k = 0; k < kTopK; ++k)
+                {
+                    sum = fmaFtz(fs.weight[tok][k], __bfloat162float(stage[(tok * kTopK + k) * 16 + row]), sum);
+                }
+                p.out[static_cast<int64_t>(tok) * kHidden + hiddenOfStoredRow(tile * 16 + row)]
+                    = __float2bfloat16_rn(sum);
+            }
+            __syncwarp();
+        }
+        if (warp == 0 && lane == 0)
+        {
+            if (int64_t* row = timelineRow(p, 1))
+            {
+                int64_t const cycles = clock64() - t0;
+                row[kTlEntryNs] = entryNs;
+                row[kTlDoneNs] = globalTimerNs();
+                row[kTlSmId] = smId();
+                row[kTlWait] = tWait;
+                row[kTlSetup] = tSetup;
+                row[kTlFirst] = tFirst;
+                row[kTlLast] = tLast;
+                row[kTlDone] = cycles;
+                row[kTlFullBlocked] = fullBlocked;
+                row[kTlChunks] = u0 * kFc2Chunks;
+            }
+        }
+    }
+#endif
+}
+
+template <typename Kernel>
+cudaError_t launchOne(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl, DraftRoutedMoeParams const& params,
+    cudaStream_t stream)
+{
+    cudaError_t status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes);
+    if (status != cudaSuccess)
+    {
+        return status;
+    }
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(ctas);
+    config.blockDim = dim3(threads);
+    config.dynamicSmemBytes = smemBytes;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    return cudaLaunchKernelEx(&config, kernel, params);
+}
+
+template <typename Kernel>
+cudaError_t launchTma(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl, DraftRoutedMoeParams const& params,
+    CUtensorMap const& map, cudaStream_t stream)
+{
+    cudaError_t status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes);
+    if (status != cudaSuccess)
+    {
+        return status;
+    }
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(ctas);
+    config.blockDim = dim3(threads);
+    config.dynamicSmemBytes = smemBytes;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    return cudaLaunchKernelEx(&config, kernel, params, map);
+}
+
+template <typename Kernel>
+cudaError_t launchTma3(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl,
+    DraftRoutedMoeParams const& params, CUtensorMap const& map, CUtensorMap const& map3, CUtensorMap const& mapX,
+    cudaStream_t stream)
+{
+    cudaError_t status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes);
+    if (status != cudaSuccess)
+    {
+        return status;
+    }
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(ctas);
+    config.blockDim = dim3(threads);
+    config.dynamicSmemBytes = smemBytes;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    return cudaLaunchKernelEx(&config, kernel, params, map, map3, mapX);
+}
+
+template <typename Kernel>
+cudaError_t launchTma2(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl,
+    DraftRoutedMoeParams const& params, CUtensorMap const& map, CUtensorMap const& map3, cudaStream_t stream)
+{
+    return launchTma3(kernel, ctas, threads, smemBytes, pdl, params, map, map3, map, stream);
+}
+
+} // namespace
+
+int fc1SmemBytes(int rows, int warps, bool deep)
+{
+    return rows * kXRowBytes + warps * (deep ? kFc1StagesDeep : kFc1Stages) * kChunkBytes;
+}
+
+int fc2SmemBytes(int warps, bool deep)
+{
+    return warps * ((deep ? kFc2StagesDeep : kFc2Stages) * kChunkBytes + kFc2HStages * kHEntryBytes + kStageBytes);
+}
+
+namespace
+{
+
+template <int W, int S>
+cudaError_t launchFc1(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream)
+{
+    return pairs ? launchOne(mtpMoeFc1Kernel<W, true, S>, params.fc1Ctas, 32 * W, smem, pdl, params, stream)
+                 : launchOne(mtpMoeFc1Kernel<W, false, S>, params.fc1Ctas, 32 * W, smem, pdl, params, stream);
+}
+
+template <int W, int S>
+cudaError_t launchFc2(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream)
+{
+    return pairs ? launchOne(mtpMoeFc2Kernel<W, true, S>, params.fc2Ctas, 32 * W, smem, pdl, params, stream)
+                 : launchOne(mtpMoeFc2Kernel<W, false, S>, params.fc2Ctas, 32 * W, smem, pdl, params, stream);
+}
+
+template <int S>
+cudaError_t launchFc1Warps(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream)
+{
+    switch (params.fc1Warps)
+    {
+    case 1: return launchFc1<1, S>(pairs, smem, pdl, params, stream);
+    case 2: return launchFc1<2, S>(pairs, smem, pdl, params, stream);
+    case 4: return launchFc1<4, S>(pairs, smem, pdl, params, stream);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+template <int S>
+cudaError_t launchFc2Warps(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream)
+{
+    switch (params.fc2Warps)
+    {
+    case 1: return launchFc2<1, S>(pairs, smem, pdl, params, stream);
+    case 2: return launchFc2<2, S>(pairs, smem, pdl, params, stream);
+    case 4: return launchFc2<4, S>(pairs, smem, pdl, params, stream);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace
+
+int fc1TmaSmemBytes(int rows, int stages)
+{
+    return 1024 + stages * kTmaStageChunks * kChunkBytes + rows * kXRowBytes;
+}
+
+int fc1TmaTriplesSmemBytes(int rows, int stages, int tripleStages)
+{
+    int const ring = stages * kTmaStageChunks * kChunkBytes;
+    int const ring3 = tripleStages * kTripleSubTiles * kTripleStageChunks * kTripleRows * kKChunk * 2;
+    return 1024 + (ring > ring3 ? ring : ring3) + rows * kXRowBytes;
+}
+
+int fc1TmaStreamXSmemBytes()
+{
+    int const ring = 6 * (kTmaStageChunks * kChunkBytes + kXBoxBytes);
+    int const ring3 = 4 * (kTripleSubTiles * kTripleRows * kTripleStageChunks * kKChunk * 2 + kXBoxBytes);
+    return 1024 + (ring > ring3 ? ring : ring3);
+}
+
+int fc2TmaSmemBytes(int warps, int stages)
+{
+    return 1024 + stages * (warps * kFc2Chunks * kChunkBytes + kHEntryBytes) + warps * kStageBytes;
+}
+
+namespace
+{
+
+//! Dynamic shared memory budget (the static routing state and barriers take about 3.5 KB of the 227 KB).
+constexpr int kMaxDynamicSmem = 220 * 1024;
+
+//! FC1 kernel with ring depth S (and S3 for the triples path, 4 MMA warps only; 0: no triples path). The instantiated
+//! (S, S3) pairs are those launchFc1TmaC chooses: (6, 4), (4, 3), (4, 2), (3, 2).
+template <int C, int S>
+cudaError_t launchFc1TmaCS(int smem, int tripleStages, bool pdl, DraftRoutedMoeParams const& params,
+    DraftRoutedMoeTensorMaps const& maps, cudaStream_t stream)
+{
+    int const n = params.fc1Ctas;
+    CUtensorMap const& w = maps.w13;
+    switch (params.fc1Warps)
+    {
+    case 1: return launchTma2(mtpMoeFc1TmaKernel<1, C, S, 0>, n, 64, smem, pdl, params, w, w, stream);
+    case 2: return launchTma2(mtpMoeFc1TmaKernel<2, C, S, 0>, n, 96, smem, pdl, params, w, w, stream);
+    case 4:
+    {
+        CUtensorMap const& w3 = maps.w13Triple;
+        if (tripleStages == 0)
+        {
+            return launchTma2(mtpMoeFc1TmaKernel<4, C, S, 0>, n, 160, smem, pdl, params, w, w, stream);
+        }
+        if constexpr (S == 6)
+        {
+            if (tripleStages == 4)
+            {
+                return launchTma2(mtpMoeFc1TmaKernel<4, C, 6, 4>, n, 160, smem, pdl, params, w, w3, stream);
+            }
+        }
+        if constexpr (S == 4)
+        {
+            if (tripleStages == 3)
+            {
+                return launchTma2(mtpMoeFc1TmaKernel<4, C, 4, 3>, n, 160, smem, pdl, params, w, w3, stream);
+            }
+            if (tripleStages == 2)
+            {
+                return launchTma2(mtpMoeFc1TmaKernel<4, C, 4, 2>, n, 160, smem, pdl, params, w, w3, stream);
+            }
+        }
+        if constexpr (S == 3)
+        {
+            if (tripleStages == 2)
+            {
+                return launchTma2(mtpMoeFc1TmaKernel<4, C, 3, 2>, n, 160, smem, pdl, params, w, w3, stream);
+            }
+        }
+        return cudaErrorInvalidValue;
+    }
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+//! FC1 ring depth: the deepest of 6 / 4 / 3 stages of 32 KB that fits next to the activation rows; with
+//! kFlagFc1Triples and 4 MMA warps also the deepest triples ring (48 KB stages) of the instantiated pairs that fits.
+//! kFlagStreamX (4 MMA warps, tile pairs; ignored otherwise): 6 stages of 32 + 2 KB and 4 triples stages of 48 + 2 KB
+//! for any row count.
+template <int C>
+cudaError_t launchFc1TmaC(
+    bool pdl, DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const& maps, cudaStream_t stream)
+{
+    if constexpr (C == 2)
+    {
+        if ((params.flags & kFlagStreamX) != 0 && params.fc1Warps == 4)
+        {
+            int const smem = fc1TmaStreamXSmemBytes();
+            return launchTma3(mtpMoeFc1TmaKernel<4, 2, 6, 4, true>, params.fc1Ctas, 160, smem, pdl, params, maps.w13,
+                maps.w13Triple, maps.x, stream);
+        }
+    }
+    bool const triples = (params.flags & kFlagFc1Triples) != 0 && params.fc1Warps == 4;
+    for (int stages : {6, 4, 3})
+    {
+        int smem = fc1TmaSmemBytes(params.rows, stages);
+        if (smem > kMaxDynamicSmem)
+        {
+            continue;
+        }
+        int tripleStages = 0;
+        if (triples)
+        {
+            for (int t : {4, 3, 2})
+            {
+                bool const instantiated
+                    = (stages == 6 && t == 4) || (stages == 4 && (t == 3 || t == 2)) || (stages == 3 && t == 2);
+                if (instantiated && fc1TmaTriplesSmemBytes(params.rows, stages, t) <= kMaxDynamicSmem)
+                {
+                    tripleStages = t;
+                    smem = fc1TmaTriplesSmemBytes(params.rows, stages, t);
+                    break;
+                }
+            }
+            if (tripleStages == 0)
+            {
+                return cudaErrorInvalidValue;
+            }
+        }
+        switch (stages)
+        {
+        case 6: return launchFc1TmaCS<C, 6>(smem, tripleStages, pdl, params, maps, stream);
+        case 4: return launchFc1TmaCS<C, 4>(smem, tripleStages, pdl, params, maps, stream);
+        default: return launchFc1TmaCS<C, 3>(smem, tripleStages, pdl, params, maps, stream);
+        }
+    }
+    return cudaErrorInvalidValue;
+}
+
+template <int W>
+cudaError_t launchFc2TmaW(bool pdl, DraftRoutedMoeParams const& params, CUtensorMap const& map, cudaStream_t stream)
+{
+    constexpr int kS = fc2TmaStages<W>();
+    int const smem = fc2TmaSmemBytes(W, kS);
+    return (params.flags & kFlagFc2Pairs)
+        ? launchTma(mtpMoeFc2TmaKernel<W, kS, 2>, params.fc2Ctas, (W + 1) * 32, smem, pdl, params, map, stream)
+        : launchTma(mtpMoeFc2TmaKernel<W, kS, 1>, params.fc2Ctas, (W + 1) * 32, smem, pdl, params, map, stream);
+}
+
+cudaError_t launchTmaPair(DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const& maps, cudaStream_t stream)
+{
+    bool const pdl = (params.flags & kFlagPdl) != 0;
+    cudaError_t status = (params.flags & kFlagFc1Pairs) ? launchFc1TmaC<2>(pdl, params, maps, stream)
+                                                        : launchFc1TmaC<1>(pdl, params, maps, stream);
+    if (status != cudaSuccess)
+    {
+        return status;
+    }
+    switch (params.fc2Warps)
+    {
+    case 1: return launchFc2TmaW<1>(pdl, params, maps.w2, stream);
+    case 2: return launchFc2TmaW<2>(pdl, params, maps.w2, stream);
+    case 3: return launchFc2TmaW<3>(pdl, params, maps.w2, stream);
+    case 4: return launchFc2TmaW<4>(pdl, params, maps.w2, stream);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace
+
+cudaError_t invokeDraftRoutedMoe(
+    DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const* maps, cudaStream_t stream)
+{
+    if (params.rows < 1 || params.rows > kMaxRows || params.fc1Ctas < 1 || params.fc2Ctas < 1)
+    {
+        return cudaErrorInvalidValue;
+    }
+    if (params.flags & kFlagTma)
+    {
+        return maps != nullptr ? launchTmaPair(params, *maps, stream) : cudaErrorInvalidValue;
+    }
+    bool const pdl = (params.flags & kFlagPdl) != 0;
+    bool const pairs1 = (params.flags & kFlagFc1Pairs) != 0;
+    bool const pairs2 = (params.flags & kFlagFc2Pairs) != 0;
+    bool const deep = (params.flags & kFlagDeepRing) != 0;
+    int const smem1 = fc1SmemBytes(params.rows, params.fc1Warps, deep);
+    cudaError_t status = deep ? launchFc1Warps<kFc1StagesDeep>(pairs1, smem1, pdl, params, stream)
+                              : launchFc1Warps<kFc1Stages>(pairs1, smem1, pdl, params, stream);
+    if (status != cudaSuccess)
+    {
+        return status;
+    }
+    int const smem2 = fc2SmemBytes(params.fc2Warps, deep);
+    return deep ? launchFc2Warps<kFc2StagesDeep>(pairs2, smem2, pdl, params, stream)
+                : launchFc2Warps<kFc2Stages>(pairs2, smem2, pdl, params, stream);
+}
+
+} // namespace kernels::glm_moe_dsa::draft_routed_moe
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h
new file mode 100644
index 000000000000..19d8adada74b
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h
@@ -0,0 +1,205 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::draft_routed_moe
+{
+
+//! The routed experts of the GLM-5.2 MTP draft layer (BF16 weights, TP8/EP1 rank shard) on the
+//! FlashInfer TRT-LLM-Gen BF16 layouts, bit-identical to ``trtllm_bf16_moe`` (DeepSeek-V3 routing with
+//! ``n_group <= 1``, SwiGLU, ``do_finalize=True``).
+//!
+//! Two kernels, chained with programmatic dependent launch:
+//!
+//! * ``fc1``: every CTA computes the native top-8 routing of all rows from the FP32 router logits (sigmoid via
+//!   ``tanh.approx``, selection key ``sigmoid + bias``, ties to the lower expert index, NaN keys first, weights
+//!   ``bf16(sigmoid * scale / (sum + 1e-20))`` with the native FP32 sum order and approximate division), then its
+//!   warps stream 16-row tiles of the distinct experts' gate/up weights (BlockMajorK, rows shuffled for the gated
+//!   epilogue) and run ``mma.sync.m16n8k16`` over K = 6144 in natural order with one FP32 accumulator per tile,
+//!   followed by the native SwiGLU epilogue. The BF16 intermediate goes to a scratch buffer.
+//! * ``fc2``: after the FC1 grid has completed, each warp computes 16 down-projection rows for every distinct
+//!   expert (K = 256 in natural order), rounds each row to BF16 like the native FC2 epilogue, and finalizes the
+//!   rows of every token in top-k slot order (``fma`` from 0, one BF16 rounding), as ``finalizeKernel`` does.
+//!
+//! The weights are read in place: ``w13`` is ``[256][96][512][64]`` and ``w2`` is ``[256][4][6144][64]`` (BF16).
+
+constexpr int kHidden = 6144;
+constexpr int kInter = 256;   //!< Per-rank intermediate size (TP8).
+constexpr int kFc1Rows = 512; //!< Gate and up rows per expert per rank.
+constexpr int kNumExperts = 256;
+constexpr int kTopK = 8;
+constexpr int kMaxRows = 8;
+constexpr int kMaxDistinct = kMaxRows * kTopK;
+constexpr int kKChunk = 64; //!< BlockMajorK block (elements, 128 bytes).
+constexpr int kFc1Chunks = kHidden / kKChunk;
+constexpr int kFc2Chunks = kInter / kKChunk;
+constexpr int kFc1Tiles = kFc1Rows / 16;
+constexpr int kFc2Tiles = kHidden / 16;
+
+//! Routing record written by the FC1 grid and read by the FC2 grid (int32 words).
+constexpr int kRouteDistinctOffset = 1;                                 //!< [kMaxDistinct] expert ids, ascending.
+constexpr int kRouteSlotOffset = kRouteDistinctOffset + kMaxDistinct;   //!< [kMaxRows][kTopK] distinct index.
+constexpr int kRouteWeightOffset = kRouteSlotOffset + kMaxRows * kTopK; //!< [kMaxRows][kTopK] float bits of bf16.
+constexpr int kRouteWords = kRouteWeightOffset + kMaxRows * kTopK;
+//! Intermediate scratch: [kMaxDistinct][kMaxRows][kInter] BF16.
+constexpr int64_t kScratchHElems = int64_t{kMaxDistinct} * kMaxRows * kInter;
+
+//! Optional behaviour (``flags``).
+enum DraftRoutedMoeFlags : int
+{
+    kFlagPdl = 1,          //!< Programmatic dependent launch for both kernels (as the native kernels).
+    kFlagPrefetchFc2 = 2,  //!< FC1 warps prefetch the distinct experts' FC2 weights into L2 after their tiles.
+    kFlagEarlyTrigger = 4, //!< FC1 triggers its dependents after routing instead of after its tiles.
+    kFlagFc1Pairs = 8,     //!< FC1 warps interleave two tiles (two independent accumulation chains).
+    kFlagFc2Pairs = 16,    //!< FC2 warps interleave two experts (two independent accumulation chains).
+    kFlagDeepRing = 32,    //!< Deeper weight rings (FC1 20, FC2 12 chunks of 2 KB per warp instead of 12 and 8).
+    kFlagTma = 64,         //!< TMA-fed kernels: one producer thread streams 16W-row weight boxes into a ring that W
+                   //!< MMA warps consume (same per-chain arithmetic). Needs the tensor maps of invokeDraftRoutedMoe.
+                   //!< With it, kFlagFc1Pairs gives every FC1 MMA warp two 16-row chains (boxes of 32W
+                   //!< rows), kFlagFc2Pairs gives every FC2 MMA warp two experts per step (two stages, two
+                   //!< accumulators), kFlagPrefetchFc2 is issued by the FC1 producer after its last weight
+                   //!< copy, and kFlagDeepRing does not apply (the rings are sized to the shared memory).
+    kFlagL2Prefetch = 128, //!< TMA FC1: the producer also prefetches its weight boxes into L2, 2 x ring depth ahead.
+    kFlagPrefetchFc2Early = 256, //!< TMA FC1: the producer prefetches its share of the FC2 weights right after routing.
+    //! GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE builds only: MMA warps wait for and release stages without reading them
+    //! (memory schedule alone; the output is meaningless).
+    kFlagProbeNoMma = 512,
+    //! GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE builds only: the FC1 producer fills the ring once and the MMA warps cycle
+    //! over the resident stages without waiting (MMA warps alone; meaningless output).
+    kFlagProbeResident = 1024,
+    //! TMA FC1 with 4 MMA warps: when the call's tiles (decided on the device from the number of distinct experts)
+    //! would need more than one round of the grid, every CTA instead takes three consecutive 64-row sub-tiles (three
+    //! chains per MMA warp, three TMA boxes per stage), if those fit in one round. Needs
+    //! DraftRoutedMoeTensorMaps::w13Triple.
+    kFlagFc1Triples = 4096,
+    //! TMA FC1: the early trigger of kFlagEarlyTrigger for calls on the kFlagFc1Triples path only (other calls trigger
+    //! after their tiles).
+    kFlagEarlyTriggerTriples = 8192,
+    //! TMA FC1: the weight boxes are loaded with an L2 evict-first policy (each is read once per call), so that they
+    //! do not displace FC2 weights prefetched into L2.
+    kFlagFc1EvictFirst = 32768,
+    //! TMA FC1 with 4 MMA warps and tile pairs (ignored otherwise): the activation rows are not kept in shared
+    //! memory; every ring stage also carries the activation K chunks it needs (boxes of DraftRoutedMoeTensorMaps::x),
+    //! so the weight ring is 6 stages deep (4 on the triples path) for every row count. Needs
+    //! DraftRoutedMoeTensorMaps::x.
+    kFlagStreamX = 65536
+};
+
+//! Event record of the probe builds (compiled with ``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1``): int64
+//! ``[2][kTimelineCtas][kTimelineSlots]`` (FC1, FC2), one row per CTA; see TimelineSlot in draftRoutedMoe.cu.
+constexpr int kTimelineCtas = 1024;
+constexpr int kTimelineSlots = 16;
+
+//! K chunks per FC1 stage of the TMA-fed kernel for all its 16-row chains together (32 KB per stage and box).
+constexpr int kTmaStageChunks = 16;
+
+//! Sub-tiles per CTA, rows per sub-tile and K chunks per stage of the kFlagFc1Triples path (4 MMA warps): a stage is
+//! three boxes of 64 rows x 2 K chunks (16 KB each).
+constexpr int kTripleSubTiles = 3;
+constexpr int kTripleRows = 64;
+constexpr int kTripleStageChunks = 2;
+
+//! Activation box of kFlagStreamX: 8 tokens (rows beyond the call's are zero-filled) x kXBoxChunks K chunks.
+constexpr int kXBoxChunks = 2;
+constexpr int kXBoxBytes = 8 * kXBoxChunks * kKChunk * 2;
+
+//! Tensor maps of the TMA-fed kernels (kFlagTma): 3D views [slab][row][64] BF16 of ``w13`` (256 * 96 slabs of 512
+//! rows) and ``w2`` (256 * 4 slabs of 6144 rows), 128-byte swizzle; boxes of 64 x fc1TmaBoxRows x fc1TmaBoxSlabs
+//! (FC1) and 64 x 16 * fc2Warps x 4 (FC2).
+struct DraftRoutedMoeTensorMaps
+{
+    CUtensorMap w13;
+    CUtensorMap w2;
+    CUtensorMap w13Triple; //!< kFlagFc1Triples: view of ``w13`` with boxes of 64 x kTripleRows x kTripleStageChunks.
+    //! kFlagStreamX: 3D view [token][K chunk][64] of the call's ``x`` (dims 64, 96, rows), boxes of 64 x kXBoxChunks x
+    //! 8, 128-byte swizzle.
+    CUtensorMap x;
+};
+
+struct DraftRoutedMoeParams
+{
+    float const* logits;       //!< [rows, 256] FP32.
+    __nv_bfloat16 const* bias; //!< [256] router bias.
+    __nv_bfloat16 const* x;    //!< [rows, 6144] activation (16-byte aligned).
+    __nv_bfloat16 const* w13;  //!< [256][96][512][64].
+    __nv_bfloat16 const* w2;   //!< [256][4][6144][64].
+    __nv_bfloat16* out;        //!< [rows, 6144] finalized routed output.
+    __nv_bfloat16* h;          //!< Scratch, kScratchHElems.
+    int32_t* route;            //!< Scratch, kRouteWords.
+    __nv_bfloat16* debugFc2;   //!< Optional [rows, 8, 6144]: FC2 rows per (token, slot) in hidden order.
+    int64_t* timeline; //!< Optional event record (GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE builds; ignored otherwise).
+    int rows;
+    float routedScaling;
+    int fc1Ctas;
+    int fc1Warps; //!< 1, 2 or 4 warps per FC1 CTA.
+    int fc2Ctas;
+    int fc2Warps; //!< 1, 2 or 4 warps per FC2 CTA (also 3 with kFlagTma: 128 groups of 48 rows).
+    int flags;
+};
+
+//! Dynamic shared memory of the FC1 kernel for ``rows`` and ``warps`` (bytes); ``deep``: kFlagDeepRing.
+int fc1SmemBytes(int rows, int warps, bool deep);
+
+//! Dynamic shared memory of the FC2 kernel for ``warps`` (bytes); ``deep``: kFlagDeepRing.
+int fc2SmemBytes(int warps, bool deep);
+
+//! Dynamic shared memory of the TMA-fed FC1 kernel (bytes, incl. 1 KB alignment slack); stages: ring depth.
+int fc1TmaSmemBytes(int rows, int stages);
+
+//! Dynamic shared memory of the TMA-fed FC1 kernel with the kFlagFc1Triples path (bytes): the larger of the two
+//! rings (``stages`` of 32 KB, ``tripleStages`` of three sub-tile boxes) next to the activation rows.
+int fc1TmaTriplesSmemBytes(int rows, int stages, int tripleStages);
+
+//! Dynamic shared memory of the TMA-fed FC1 kernel with kFlagStreamX (bytes): the larger of the 6-stage regular and
+//! the 4-stage triples rings, activation boxes included.
+int fc1TmaStreamXSmemBytes();
+
+//! Dynamic shared memory of the TMA-fed FC2 kernel (bytes, incl. 1 KB alignment slack); stages: ring depth.
+int fc2TmaSmemBytes(int warps, int stages);
+
+//! Launches the FC1 and FC2 kernels on ``stream``; ``fc1Warps``/``fc2Warps`` must be 1, 2 or 4 (``fc2Warps`` also 3
+//! with kFlagTma). ``maps`` is required
+//! with kFlagTma (see DraftRoutedMoeTensorMaps) and ignored otherwise.
+cudaError_t invokeDraftRoutedMoe(
+    DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const* maps, cudaStream_t stream);
+
+//! Rows of the FC1 weight box of the TMA-fed kernel (16 per chain).
+inline int fc1TmaBoxRows(int flags, int fc1Warps)
+{
+    return 16 * fc1Warps * ((flags & kFlagFc1Pairs) ? 2 : 1);
+}
+
+//! K chunks (slabs) of the FC1 weight box of the TMA-fed kernel.
+inline int fc1TmaBoxSlabs(int flags, int fc1Warps)
+{
+    return kTmaStageChunks * 16 / fc1TmaBoxRows(flags, fc1Warps);
+}
+
+} // namespace kernels::glm_moe_dsa::draft_routed_moe
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu
new file mode 100644
index 000000000000..a22ae0d29395
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu
@@ -0,0 +1,1568 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/draftTail.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::draft_tail
+{
+
+namespace
+{
+
+// ---------------------------------------------------------------------------
+// Candidate keys
+// ---------------------------------------------------------------------------
+
+// Monotone map of BF16 bits to 16 bits: -0 and +0 compare equal, NaN of either
+// sign ranks above +inf (as the native torch.max reduction does).
+__device__ __forceinline__ uint32_t orderedBf16(uint32_t bits)
+{
+    uint32_t const mag = bits & 0x7FFFu;
+    if (mag > 0x7F80u)
+    {
+        return 0xFFFFu;
+    }
+    if (mag == 0u)
+    {
+        bits = 0u;
+    }
+    return (bits & 0x8000u) ? (~bits & 0xFFFFu) : (bits | 0x8000u);
+}
+
+// Larger key == larger value, then smaller global id. Never zero.
+__device__ __forceinline__ uint64_t makeKey(uint32_t bf16Bits, int64_t globalId)
+{
+    return (static_cast<uint64_t>(orderedBf16(bf16Bits)) << 32)
+        | static_cast<uint64_t>(0xFFFFFFFFu - static_cast<uint32_t>(globalId));
+}
+
+__device__ __forceinline__ int32_t keyToToken(uint64_t key)
+{
+    return static_cast<int32_t>(0xFFFFFFFFu - static_cast<uint32_t>(key & 0xFFFFFFFFull));
+}
+
+__device__ __forceinline__ uint64_t umax64(uint64_t a, uint64_t b)
+{
+    return a > b ? a : b;
+}
+
+__device__ __forceinline__ uint64_t warpMax64(uint64_t v)
+{
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+    {
+        v = umax64(v, __shfl_xor_sync(0xFFFFFFFFu, v, off));
+    }
+    return v;
+}
+
+// ---------------------------------------------------------------------------
+// Memory helpers
+// ---------------------------------------------------------------------------
+
+__device__ __forceinline__ void griddepWait()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+#endif
+}
+
+__device__ __forceinline__ uint64_t ldRelaxedSys(uint64_t const* addr)
+{
+    uint64_t v;
+    asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(addr) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ void stRelaxedSys(uint64_t* addr, uint64_t v)
+{
+    asm volatile("st.relaxed.sys.global.u64 [%0], %1;" ::"l"(addr), "l"(v) : "memory");
+}
+
+__device__ __forceinline__ void mcStoreRelaxedSys(uint64_t* mcAddr, uint64_t v)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("multimem.st.relaxed.sys.global.b64 [%0], %1;" ::"l"(mcAddr), "l"(v) : "memory");
+#endif
+}
+
+__device__ __forceinline__ uint64_t evictFirstPolicy()
+{
+    uint64_t policy;
+    asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy));
+    return policy;
+}
+
+__device__ __forceinline__ uint64_t evictLastPolicy()
+{
+    uint64_t policy;
+    asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;" : "=l"(policy));
+    return policy;
+}
+
+// One 32-byte chunk = 16 BF16 values.
+struct alignas(16) Chunk
+{
+    uint64_t v[4];
+};
+
+enum LoadKind : int
+{
+    kLoad32Hint = 0, // 32-byte load, L2 evict_first policy
+    kLoad32 = 1,     // 32-byte load
+    kLoad16x2Pf = 2, // two 16-byte loads with 256-byte L2 prefetch
+};
+
+template <int kKind>
+__device__ __forceinline__ void loadChunk(Chunk& c, void const* ptr, uint64_t policy)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000)
+    // 32-byte global loads need SM100+.
+    __trap();
+#else
+    if constexpr (kKind == kLoad32Hint)
+    {
+        asm volatile("ld.global.nc.L1::no_allocate.L2::cache_hint.v4.b64 {%0, %1, %2, %3}, [%4], %5;"
+                     : "=l"(c.v[0]), "=l"(c.v[1]), "=l"(c.v[2]), "=l"(c.v[3])
+                     : "l"(ptr), "l"(policy));
+    }
+    else if constexpr (kKind == kLoad32)
+    {
+        asm volatile("ld.global.nc.L1::no_allocate.v4.b64 {%0, %1, %2, %3}, [%4];"
+                     : "=l"(c.v[0]), "=l"(c.v[1]), "=l"(c.v[2]), "=l"(c.v[3])
+                     : "l"(ptr));
+    }
+    else
+    {
+        int4 a;
+        int4 b;
+        asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.s32 {%0, %1, %2, %3}, [%4];"
+                     : "=r"(a.x), "=r"(a.y), "=r"(a.z), "=r"(a.w)
+                     : "l"(ptr));
+        asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.s32 {%0, %1, %2, %3}, [%4];"
+                     : "=r"(b.x), "=r"(b.y), "=r"(b.z), "=r"(b.w)
+                     : "l"(static_cast<char const*>(ptr) + 16));
+        c.v[0] = (static_cast<uint64_t>(static_cast<uint32_t>(a.y)) << 32) | static_cast<uint32_t>(a.x);
+        c.v[1] = (static_cast<uint64_t>(static_cast<uint32_t>(a.w)) << 32) | static_cast<uint32_t>(a.z);
+        c.v[2] = (static_cast<uint64_t>(static_cast<uint32_t>(b.y)) << 32) | static_cast<uint32_t>(b.x);
+        c.v[3] = (static_cast<uint64_t>(static_cast<uint32_t>(b.w)) << 32) | static_cast<uint32_t>(b.z);
+    }
+#endif
+}
+
+// acc += dot(w, h) over 16 BF16 pairs, in a fixed order.
+__device__ __forceinline__ float dotChunk(Chunk const& w, Chunk const& h, float acc)
+{
+#pragma unroll
+    for (int i = 0; i < 4; ++i)
+    {
+        uint32_t const wl = static_cast<uint32_t>(w.v[i]);
+        uint32_t const wh = static_cast<uint32_t>(w.v[i] >> 32);
+        uint32_t const hl = static_cast<uint32_t>(h.v[i]);
+        uint32_t const hh = static_cast<uint32_t>(h.v[i] >> 32);
+        acc = fmaf(__uint_as_float(wl << 16), __uint_as_float(hl << 16), acc);
+        acc = fmaf(__uint_as_float(wl & 0xFFFF0000u), __uint_as_float(hl & 0xFFFF0000u), acc);
+        acc = fmaf(__uint_as_float(wh << 16), __uint_as_float(hh << 16), acc);
+        acc = fmaf(__uint_as_float(wh & 0xFFFF0000u), __uint_as_float(hh & 0xFFFF0000u), acc);
+    }
+    return acc;
+}
+
+constexpr int kHidden6144ChunksDev = 6144 / 512; // 32-byte chunks per lane per row
+
+__device__ __forceinline__ uint32_t* arrivalCounter(uint64_t* scratch)
+{
+    return reinterpret_cast<uint32_t*>(scratch + kMaxRows);
+}
+
+__device__ __forceinline__ uint32_t* epochCounter(uint64_t* scratch)
+{
+    return reinterpret_cast<uint32_t*>(scratch + kMaxRows) + 1;
+}
+
+// ---------------------------------------------------------------------------
+// Cross-rank exchange (one warp)
+//
+// Slot s of the exchange buffer holds kMaxRanks x kMaxRows keys; rank q writes
+// only its own row block, every rank polls its local copy until all tpSize
+// entries of each row are non-zero. Epoch e uses slot e % 3. Before publishing
+// epoch e a rank clears slot (e + 1) % 3, whose last use was epoch e - 2: every
+// key of that epoch was already observed here, and no peer can publish epoch
+// e + 1 before it has observed this rank's epoch-e key, which the release fence
+// orders after the clear (the peer's own fence before its next publish is the
+// matching acquire). All ranks run the same sequence of epochs.
+// ---------------------------------------------------------------------------
+__device__ void exchangeAndWrite(
+    DraftTailParams const& p, uint64_t const* localKeys /* smem [rows] */, uint64_t (*gathered)[kMaxRows] /* smem */)
+{
+    int const lane = threadIdx.x & 31;
+    int const rows = p.rows;
+    bool const exchange = p.mode != ExchangeMode::kNone && p.tpSize > 1;
+    uint64_t result = 0;
+    if (!exchange)
+    {
+        if (lane < rows)
+        {
+            result = localKeys[lane];
+        }
+    }
+    else
+    {
+        uint32_t* epochPtr = epochCounter(p.scratch);
+        uint32_t const epoch = *reinterpret_cast<uint32_t volatile*>(epochPtr);
+        int const cur = static_cast<int>(epoch % kNumSlots);
+        int const nxt = static_cast<int>((epoch + 1) % kNumSlots);
+        for (int i = lane; i < kSlotWords; i += 32)
+        {
+            p.ucSlotsLocal[nxt * kSlotWords + i] = 0ull;
+        }
+        // Release: the clear above is ordered before this rank's publish below;
+        // a peer's next write to the cleared slot is causally after that publish.
+        asm volatile("fence.acq_rel.sys;" ::: "memory");
+        __syncwarp();
+        int const mine = cur * kSlotWords + p.tpRank * kMaxRows;
+        if (p.mode == ExchangeMode::kMulticast)
+        {
+            if (lane < rows)
+            {
+                mcStoreRelaxedSys(p.mcSlots + mine + lane, localKeys[lane]);
+            }
+        }
+        else
+        {
+            for (int i = lane; i < p.tpSize * rows; i += 32)
+            {
+                int const q = i / rows;
+                int const r = i - q * rows;
+                stRelaxedSys(p.ucSlotsPeers[q] + mine + r, localKeys[r]);
+            }
+        }
+        long long const start = clock64();
+        for (int i = lane; i < p.tpSize * rows; i += 32)
+        {
+            int const q = i / rows;
+            int const r = i - q * rows;
+            uint64_t const* addr = p.ucSlotsLocal + cur * kSlotWords + q * kMaxRows + r;
+            uint64_t v = ldRelaxedSys(addr);
+            while (v == 0ull)
+            {
+                if (p.timeoutCycles > 0 && clock64() - start > p.timeoutCycles)
+                {
+                    // A peer never published: the ranks disagree on the
+                    // collective sequence. Record it and abort loudly.
+                    atomicExch(reinterpret_cast<unsigned long long*>(p.scratch + kMaxRows + 1),
+                        0x5733000000000000ull | (static_cast<unsigned long long>(epoch) << 8)
+                            | static_cast<unsigned long long>(q));
+                    __threadfence_system();
+                    __trap();
+                }
+                v = ldRelaxedSys(addr);
+            }
+            gathered[q][r] = v;
+        }
+        __syncwarp();
+        if (lane < rows)
+        {
+            for (int q = 0; q < p.tpSize; ++q)
+            {
+                result = umax64(result, gathered[q][lane]);
+            }
+        }
+        if (lane == 0)
+        {
+            *epochPtr = epoch + 1;
+        }
+    }
+    if (lane < rows)
+    {
+        p.tokens[lane] = keyToToken(result);
+        if (p.keysOut != nullptr)
+        {
+            p.keysOut[lane] = result;
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Fused kernel: GEMV + rank-local max + exchange
+//
+// CTA c owns a contiguous block of vocab rows; warp w of it takes rows
+// begin + w, begin + w + kWarps, ... Each row is split into kSplit units of
+// kUnit 32-byte chunks per lane, double-buffered in registers so a warp always
+// has the next unit in flight while it computes the current one. The first
+// unit is requested before griddepcontrol.wait, i.e. while the producer of
+// the hidden rows may still be running under programmatic dependent launch.
+// ---------------------------------------------------------------------------
+// ---------------------------------------------------------------------------
+// Direct exchange: every CTA publishes its partial key for each row into
+// entry [slot][row][rank][cta] of every rank's buffer; CTA 0 of each rank polls
+// all tpSize x numCtas entries of each row, writes the tokens, then clears
+// exactly the entries it polled and advances the epoch. A slot's entries are
+// written again only three epochs later, i.e. after this rank's next two
+// publishes, which follow this launch; so no peer can write a cleared entry
+// before the clear has completed. CTAs without vocab rows publish key 1, below
+// every real key.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ bool isDirectMode(ExchangeMode m)
+{
+    return m == ExchangeMode::kMulticastDirect || m == ExchangeMode::kUnicastDirect;
+}
+
+__device__ __forceinline__ size_t directEntry(int slot, int row, int rank, int cta)
+{
+    return kDirectOffsetWords + static_cast<size_t>(slot) * kDirectSlotWords
+        + (static_cast<size_t>(row) * kMaxRanks + rank) * kMaxDirectCtas + cta;
+}
+
+// All threads of the CTA call this once ctaKeys[0..rows) (shared) is final.
+__device__ void directExchange(DraftTailParams const& p, uint64_t const* ctaKeys, int ctaIdx, int numCtas)
+{
+    __shared__ uint32_t sEpoch;
+    __shared__ uint64_t red[32][kMaxRows];
+    int const tid = threadIdx.x;
+    int const lane = tid & 31;
+    int const warp = tid >> 5;
+    int const rows = p.rows;
+    if (tid == 0)
+    {
+        sEpoch = *reinterpret_cast<uint32_t volatile*>(epochCounter(p.scratch));
+    }
+    __syncthreads();
+    uint32_t const epoch = sEpoch;
+    int const cur = static_cast<int>(epoch % kNumSlots);
+    if (p.mode == ExchangeMode::kMulticastDirect)
+    {
+        if (tid < rows)
+        {
+            uint64_t const key = ctaKeys[tid] != 0ull ? ctaKeys[tid] : 1ull;
+            mcStoreRelaxedSys(p.mcSlots + directEntry(cur, tid, p.tpRank, ctaIdx), key);
+        }
+    }
+    else
+    {
+        for (int i = tid; i < p.tpSize * rows; i += blockDim.x)
+        {
+            int const q = i / rows;
+            int const r = i - q * rows;
+            uint64_t const key = ctaKeys[r] != 0ull ? ctaKeys[r] : 1ull;
+            stRelaxedSys(p.ucSlotsPeers[q] + directEntry(cur, r, p.tpRank, ctaIdx), key);
+        }
+    }
+    if (ctaIdx != 0)
+    {
+        return;
+    }
+    long long const start = clock64();
+    int const perRow = p.tpSize * numCtas;
+    // Each thread issues up to kPoll independent loads per round and re-polls
+    // only the entries still empty, so a row costs about one memory round
+    // trip after the last entry lands, however many CTAs published.
+    constexpr int kPoll = 32;
+    for (int r = 0; r < rows; ++r)
+    {
+        uint64_t m = 0ull;
+        for (int base = tid; base < perRow; base += kPoll * blockDim.x)
+        {
+            uint64_t v[kPoll];
+#pragma unroll
+            for (int b = 0; b < kPoll; ++b)
+            {
+                int const i = base + b * blockDim.x;
+                v[b] = i < perRow ? ldRelaxedSys(p.ucSlotsLocal + directEntry(cur, r, i / numCtas, i % numCtas)) : 1ull;
+            }
+            bool done = false;
+            while (!done)
+            {
+                done = true;
+#pragma unroll
+                for (int b = 0; b < kPoll; ++b)
+                {
+                    if (v[b] == 0ull)
+                    {
+                        int const i = base + b * blockDim.x;
+                        done = false;
+                        v[b] = ldRelaxedSys(p.ucSlotsLocal + directEntry(cur, r, i / numCtas, i % numCtas));
+                    }
+                }
+                if (!done && p.timeoutCycles > 0 && clock64() - start > p.timeoutCycles)
+                {
+                    atomicExch(reinterpret_cast<unsigned long long*>(p.scratch + kMaxRows + 1),
+                        0x5734000000000000ull | (static_cast<unsigned long long>(epoch) << 16)
+                            | static_cast<unsigned long long>(base));
+                    __threadfence_system();
+                    __trap();
+                }
+            }
+#pragma unroll
+            for (int b = 0; b < kPoll; ++b)
+            {
+                m = umax64(m, v[b]);
+            }
+        }
+        m = warpMax64(m);
+        if (lane == 0)
+        {
+            red[warp][r] = m;
+        }
+    }
+    __syncthreads();
+    if (tid < rows)
+    {
+        uint64_t result = 0ull;
+        int const numWarps = (blockDim.x + 31) >> 5;
+        for (int w = 0; w < numWarps; ++w)
+        {
+            result = umax64(result, red[w][tid]);
+        }
+        p.tokens[tid] = keyToToken(result);
+        if (p.keysOut != nullptr)
+        {
+            p.keysOut[tid] = result;
+        }
+    }
+    for (int i = tid; i < rows * perRow; i += blockDim.x)
+    {
+        int const r = i / perRow;
+        int const j = i - r * perRow;
+        int const q = j / numCtas;
+        int const c = j - q * numCtas;
+        p.ucSlotsLocal[directEntry(cur, r, q, c)] = 0ull;
+    }
+    if (tid == 0)
+    {
+        *epochCounter(p.scratch) = epoch + 1;
+    }
+}
+
+// Per-CTA epilogue of the fused kernels: CTA max of the warps' keys, one
+// atomicMax per row into rank-local scratch, then the last CTA to arrive takes
+// (and resets) the rank-local keys and runs the exchange.
+template <int R, int kWarps>
+__device__ __forceinline__ void ctaFinish(DraftTailParams const& p, uint64_t const (&best)[R])
+{
+    __shared__ uint64_t warpBest[kWarps][R];
+    __shared__ uint64_t ctaKeys[kMaxRows];
+    __shared__ uint64_t gathered[kMaxRanks][kMaxRows];
+    __shared__ int isLast;
+    int const lane = threadIdx.x & 31;
+    int const warp = threadIdx.x >> 5;
+    if (lane == 0)
+    {
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            warpBest[warp][r] = best[r];
+        }
+    }
+    __syncthreads();
+    if (isDirectMode(p.mode))
+    {
+        if (threadIdx.x < R)
+        {
+            uint64_t m = 0ull;
+#pragma unroll
+            for (int w = 0; w < kWarps; ++w)
+            {
+                m = umax64(m, warpBest[w][threadIdx.x]);
+            }
+            ctaKeys[threadIdx.x] = m;
+        }
+        __syncthreads();
+        directExchange(p, ctaKeys, blockIdx.x, gridDim.x);
+        return;
+    }
+    if (threadIdx.x < R)
+    {
+        uint64_t m = 0ull;
+#pragma unroll
+        for (int w = 0; w < kWarps; ++w)
+        {
+            m = umax64(m, warpBest[w][threadIdx.x]);
+        }
+        if (m != 0ull)
+        {
+            atomicMax(reinterpret_cast<unsigned long long*>(p.scratch + threadIdx.x), m);
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0)
+    {
+        __threadfence();
+        uint32_t const prev = atomicAdd(arrivalCounter(p.scratch), 1u);
+        isLast = prev == gridDim.x - 1;
+    }
+    __syncthreads();
+    if (!isLast || warp != 0)
+    {
+        return;
+    }
+    __threadfence();
+    if (lane < R)
+    {
+        ctaKeys[lane] = atomicExch(reinterpret_cast<unsigned long long*>(p.scratch + lane), 0ull);
+    }
+    if (lane == 0)
+    {
+        atomicExch(arrivalCounter(p.scratch), 0u);
+    }
+    __syncwarp();
+    exchangeAndWrite(p, ctaKeys, gathered);
+}
+
+template <int kUnit, int kLoad>
+__device__ __forceinline__ void loadUnit(Chunk (&buf)[kUnit], char const* base, uint64_t policy)
+{
+#pragma unroll
+    for (int j = 0; j < kUnit; ++j)
+    {
+        loadChunk<kLoad>(buf[j], base + j * 32 * 32, policy);
+    }
+}
+
+// Byte address of unit u of this lane: row firstRow + (u / kSplit) * kWarps,
+// part u % kSplit; chunk c of a row lives at bytes [32 c, 32 c + 32) and lane l
+// owns chunks part * kUnit * 32 + j * 32 + l.
+template <int kWarps, int kChunks, int kSplit>
+__device__ __forceinline__ char const* unitAddress(__nv_bfloat16 const* weight, int firstRow, int u, int lane)
+{
+    constexpr int kUnit = kChunks / kSplit;
+    int const rowIdx = u / kSplit;
+    int const part = u - rowIdx * kSplit;
+    int64_t const row = firstRow + static_cast<int64_t>(rowIdx) * kWarps;
+    return reinterpret_cast<char const*>(weight + row * (kChunks * 32 * 16)) + (part * kUnit * 32 + lane) * 32;
+}
+
+template <int R, int kWarps, int kChunks, int kSplit>
+__device__ __forceinline__ void computeUnit(Chunk const (&buf)[kChunks / kSplit], Chunk const* smemH, int u,
+    int firstRow, int lane, int64_t vocabOffset, float (&acc)[R], uint64_t (&best)[R])
+{
+    constexpr int kUnit = kChunks / kSplit;
+    int const rowIdx = u / kSplit;
+    int const part = u - rowIdx * kSplit;
+#pragma unroll
+    for (int j = 0; j < kUnit; ++j)
+    {
+        int const c = (part * kUnit + j) * 32 + lane;
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            acc[r] = dotChunk(buf[j], smemH[r * kChunks * 32 + c], acc[r]);
+        }
+    }
+    if (part == kSplit - 1)
+    {
+        int64_t const row = firstRow + static_cast<int64_t>(rowIdx) * kWarps;
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            float s = acc[r];
+#pragma unroll
+            for (int off = 16; off > 0; off >>= 1)
+            {
+                s += __shfl_xor_sync(0xFFFFFFFFu, s, off);
+            }
+            uint32_t const bits = static_cast<uint32_t>(__bfloat16_as_ushort(__float2bfloat16_rn(s)));
+            best[r] = umax64(best[r], makeKey(bits, vocabOffset + row));
+            acc[r] = 0.f;
+        }
+    }
+}
+
+template <int R, int kWarps, int kChunks, int kSplit, int kLoad, int kMinBlocks>
+__global__ void __launch_bounds__(kWarps * 32, kMinBlocks) fusedDraftTailKernel(DraftTailParams p)
+{
+    constexpr int kUnit = kChunks / kSplit;
+    static_assert(kChunks % kSplit == 0, "row must split evenly");
+
+    extern __shared__ uint4 smemDyn[];
+    Chunk* smemH = reinterpret_cast<Chunk*>(smemDyn); // [R][kChunks * 32]
+
+    int const lane = threadIdx.x & 31;
+    int const warp = threadIdx.x >> 5;
+    int const numRows = p.vocabLocal;
+    int const rowsPerCta = (numRows + gridDim.x - 1) / gridDim.x;
+    int const rowBegin = blockIdx.x * rowsPerCta;
+    int const rowEnd = min(numRows, rowBegin + rowsPerCta);
+    int const firstRow = rowBegin + warp;
+    int const myRows = firstRow < rowEnd ? (rowEnd - firstRow + kWarps - 1) / kWarps : 0;
+    int const numUnits = myRows * kSplit;
+
+    uint64_t const policyFirst = kLoad == kLoad32Hint ? evictFirstPolicy() : 0ull;
+    uint64_t const policyLast = kLoad == kLoad32Hint ? evictLastPolicy() : 0ull;
+    int const keepRows = static_cast<int>((static_cast<int64_t>(myRows) * p.keepPermille) / 1000);
+    __nv_bfloat16 const* weight = p.weight;
+
+    Chunk bufA[kUnit];
+    Chunk bufB[kUnit];
+    if (numUnits > 0)
+    {
+        loadUnit<kUnit, kLoad>(bufA, unitAddress<kWarps, kChunks, kSplit>(weight, firstRow, 0, lane),
+            keepRows > 0 ? policyLast : policyFirst);
+    }
+
+    griddepWait();
+
+    {
+        int const total = R * kChunks * 32;
+        for (int i = threadIdx.x; i < total; i += blockDim.x)
+        {
+            int const r = i / (kChunks * 32);
+            int const c = i - r * (kChunks * 32);
+            uint4 const* src = reinterpret_cast<uint4 const*>(p.hidden + r * p.hiddenStride) + 2 * c;
+            uint4 const lo = src[0];
+            uint4 const hi = src[1];
+            Chunk ch;
+            ch.v[0] = (static_cast<uint64_t>(lo.y) << 32) | lo.x;
+            ch.v[1] = (static_cast<uint64_t>(lo.w) << 32) | lo.z;
+            ch.v[2] = (static_cast<uint64_t>(hi.y) << 32) | hi.x;
+            ch.v[3] = (static_cast<uint64_t>(hi.w) << 32) | hi.z;
+            smemH[i] = ch;
+        }
+    }
+    __syncthreads();
+
+    float acc[R];
+    uint64_t best[R];
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+    {
+        acc[r] = 0.f;
+        best[r] = 0ull;
+    }
+
+    for (int u = 0; u < numUnits; u += 2)
+    {
+        if (u + 1 < numUnits)
+        {
+            loadUnit<kUnit, kLoad>(bufB, unitAddress<kWarps, kChunks, kSplit>(weight, firstRow, u + 1, lane),
+                (u + 1) / kSplit < keepRows ? policyLast : policyFirst);
+        }
+        computeUnit<R, kWarps, kChunks, kSplit>(bufA, smemH, u, firstRow, lane, p.vocabOffset, acc, best);
+        if (u + 1 < numUnits)
+        {
+            if (u + 2 < numUnits)
+            {
+                loadUnit<kUnit, kLoad>(bufA, unitAddress<kWarps, kChunks, kSplit>(weight, firstRow, u + 2, lane),
+                    (u + 2) / kSplit < keepRows ? policyLast : policyFirst);
+            }
+            computeUnit<R, kWarps, kChunks, kSplit>(bufB, smemH, u + 1, firstRow, lane, p.vocabOffset, acc, best);
+        }
+    }
+
+    ctaFinish<R, kWarps>(p, best);
+}
+
+// ---------------------------------------------------------------------------
+// Fused kernel, bulk-copy flavour: each warp streams its rows through its own
+// ring of kSlots shared-memory row buffers filled by cp.async.bulk (one 12 KB
+// copy per row, completion on a per-slot mbarrier), so the bytes in flight do
+// not depend on registers. Lane l reads chunks l + 32 j of a row in the same
+// order as the register kernel, so both produce bitwise-identical keys.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ uint32_t smemAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+__device__ __forceinline__ void mbarInit(uint64_t* bar, uint32_t count)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smemAddr(bar)), "r"(count) : "memory");
+}
+
+__device__ __forceinline__ void mbarExpectTx(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemAddr(bar)), "r"(bytes) : "memory");
+}
+
+__device__ __forceinline__ void mbarWait(uint64_t* bar, uint32_t phase)
+{
+    uint32_t const addr = smemAddr(bar);
+    asm volatile(
+        "{\n"
+        ".reg .pred done;\n"
+        "DRAFT_TAIL_WAIT_%=:\n"
+        "mbarrier.try_wait.parity.shared::cta.b64 done, [%0], %1;\n"
+        "@!done bra DRAFT_TAIL_WAIT_%=;\n"
+        "}\n" ::"r"(addr),
+        "r"(phase)
+        : "memory");
+}
+
+__device__ __forceinline__ void bulkLoad(void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t policy)
+{
+    asm volatile(
+        "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" ::
+            "r"(smemAddr(dst)),
+        "l"(src), "r"(bytes), "r"(smemAddr(bar)), "l"(policy)
+        : "memory");
+}
+
+template <int R, int kWarps, int kSlots>
+__global__ void __launch_bounds__(kWarps * 32, 1) fusedBulkDraftTailKernel(DraftTailParams p)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900)
+    // Bulk copies and transaction mbarriers need SM90+.
+    __trap();
+#else
+    constexpr int kChunks = kHidden6144ChunksDev;
+    constexpr int kRowBytes = kChunks * 32 * 32;
+    extern __shared__ uint4 smemDyn[];
+    Chunk* smemH = reinterpret_cast<Chunk*>(smemDyn); // [R][kChunks * 32]
+    Chunk* ring = smemH + R * kChunks * 32;           // [kWarps][kSlots][kChunks * 32]
+    __shared__ __align__(8) uint64_t bars[kWarps][kSlots];
+
+    int const lane = threadIdx.x & 31;
+    int const warp = threadIdx.x >> 5;
+    int const numRows = p.vocabLocal;
+    int const rowsPerCta = (numRows + gridDim.x - 1) / gridDim.x;
+    int const rowBegin = blockIdx.x * rowsPerCta;
+    int const rowEnd = min(numRows, rowBegin + rowsPerCta);
+    int const firstRow = rowBegin + warp;
+    int const myRows = firstRow < rowEnd ? (rowEnd - firstRow + kWarps - 1) / kWarps : 0;
+    int const keepRows = static_cast<int>((static_cast<int64_t>(myRows) * p.keepPermille) / 1000);
+    Chunk* myRing = ring + warp * kSlots * kChunks * 32;
+
+    uint64_t const policyFirst = evictFirstPolicy();
+    uint64_t const policyLast = evictLastPolicy();
+    auto rowPtr = [&](int i) -> void const*
+    { return p.weight + (static_cast<int64_t>(firstRow) + static_cast<int64_t>(i) * kWarps) * (kChunks * 32 * 16); };
+
+    if (lane == 0)
+    {
+#pragma unroll
+        for (int s = 0; s < kSlots; ++s)
+        {
+            mbarInit(&bars[warp][s], 1);
+        }
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+#pragma unroll
+        for (int s = 0; s < kSlots; ++s)
+        {
+            if (s < myRows)
+            {
+                mbarExpectTx(&bars[warp][s], kRowBytes);
+                bulkLoad(myRing + s * kChunks * 32, rowPtr(s), kRowBytes, &bars[warp][s],
+                    s < keepRows ? policyLast : policyFirst);
+            }
+        }
+    }
+
+    griddepWait();
+    {
+        int const total = R * kChunks * 32;
+        for (int i = threadIdx.x; i < total; i += blockDim.x)
+        {
+            int const r = i / (kChunks * 32);
+            int const c = i - r * (kChunks * 32);
+            uint4 const* src = reinterpret_cast<uint4 const*>(p.hidden + r * p.hiddenStride) + 2 * c;
+            uint4 const lo = src[0];
+            uint4 const hi = src[1];
+            Chunk ch;
+            ch.v[0] = (static_cast<uint64_t>(lo.y) << 32) | lo.x;
+            ch.v[1] = (static_cast<uint64_t>(lo.w) << 32) | lo.z;
+            ch.v[2] = (static_cast<uint64_t>(hi.y) << 32) | hi.x;
+            ch.v[3] = (static_cast<uint64_t>(hi.w) << 32) | hi.z;
+            smemH[i] = ch;
+        }
+    }
+    __syncthreads();
+
+    uint64_t best[R];
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+    {
+        best[r] = 0ull;
+    }
+    for (int i = 0; i < myRows; ++i)
+    {
+        int const slot = i % kSlots;
+        mbarWait(&bars[warp][slot], static_cast<uint32_t>((i / kSlots) & 1));
+        Chunk const* row = myRing + slot * kChunks * 32;
+        float acc[R];
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            acc[r] = 0.f;
+        }
+#pragma unroll
+        for (int j = 0; j < kChunks; ++j)
+        {
+            int const c = j * 32 + lane;
+            Chunk const w = row[c];
+#pragma unroll
+            for (int r = 0; r < R; ++r)
+            {
+                acc[r] = dotChunk(w, smemH[r * kChunks * 32 + c], acc[r]);
+            }
+        }
+        __syncwarp();
+        if (lane == 0 && i + kSlots < myRows)
+        {
+            mbarExpectTx(&bars[warp][slot], kRowBytes);
+            bulkLoad(myRing + slot * kChunks * 32, rowPtr(i + kSlots), kRowBytes, &bars[warp][slot],
+                i + kSlots < keepRows ? policyLast : policyFirst);
+        }
+        int64_t const rowId = firstRow + static_cast<int64_t>(i) * kWarps;
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            float sum = acc[r];
+#pragma unroll
+            for (int off = 16; off > 0; off >>= 1)
+            {
+                sum += __shfl_xor_sync(0xFFFFFFFFu, sum, off);
+            }
+            uint32_t const bits = static_cast<uint32_t>(__bfloat16_as_ushort(__float2bfloat16_rn(sum)));
+            best[r] = umax64(best[r], makeKey(bits, p.vocabOffset + rowId));
+        }
+    }
+    ctaFinish<R, kWarps>(p, best);
+#endif
+}
+
+// ---------------------------------------------------------------------------
+// Fused kernel, tensor-core flavour: each warp owns 16-row tiles of the vocab
+// shard and computes them with mma.sync m16n8k16 (BF16 in, FP32 accumulate),
+// the R <= 8 hidden rows being the n columns. Lane (g = lane / 4, t = lane % 4)
+// loads 32 contiguous bytes (16 values) of rows g and g + 8 per 64-value K
+// block; inside a block the K order is permuted consistently for A and B
+// (values 16t + 4s .. 16t + 4s + 3 feed k-step s), so a lane's two 16-byte
+// halves map directly onto the A/B fragments without shuffles. This replaces
+// the CUDA-core unpack + FMA loop (instruction bound) with 4 HMMA per 2 KB
+// of weight per warp.
+//
+// Native-order flavours (MmaOrder below): k-step q of a 64-value block takes
+// the 16 consecutive values 16q .. 16q + 15, so every MMA sums the same K
+// values, in the same k-step order, as a GEMM that walks K in 16-value steps.
+// They differ only in how the values reach the lanes.
+// ---------------------------------------------------------------------------
+__device__ __forceinline__ void mmaBf16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};\n"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+constexpr int kMmaBlockElems = 64;                // K values per lane-group block (4 lanes x 16)
+constexpr int kMmaBlocks = 6144 / kMmaBlockElems; // 96
+constexpr int kMmaHStride = 6144 + 16;            // padded smem row (elements) against bank conflicts
+
+// Which K values lane t (= lane % 4) of a quad feeds into the 4 mma k-steps
+// q = 0..3 of each 64-value block, and how they are loaded. Every lane feeds
+// 4 values (one 8-byte piece) per k-step and row.
+enum MmaOrder : int
+{
+    // Lane t: values 16 t .. 16 t + 15 (one 32-byte load), piece q into k-step
+    // q; each k-step sums a K subset spread over the block (not native order).
+    kOrderPermuted = 0,
+    // Native orders: k-step q sums exactly values 16 q .. 16 q + 15.
+    // Lane t loads piece t of every k-step: four 8-byte loads.
+    kOrderNat8 = 1,
+    // Lane t loads k-step t whole (one 32-byte load); three in-quad shuffles
+    // hand piece (t - q) mod 4 of k-step q to lane t.
+    kOrderRot32 = 2,
+    // Lane t loads half (t >> 1) of k-steps 2 (t & 1) and 2 (t & 1) + 1 (two
+    // 16-byte loads), keeps the first piece of each half and swaps the second
+    // with lane t ^ 1.
+    kOrderPair16 = 3,
+};
+
+// Byte offset of lane t's first load inside a 128-byte block row.
+template <int kOrder>
+__device__ __forceinline__ int mmaLaneOffset(int t)
+{
+    if constexpr (kOrder == kOrderNat8)
+    {
+        return 8 * t;
+    }
+    else if constexpr (kOrder == kOrderPair16)
+    {
+        return 64 * (t & 1) + 16 * (t >> 1);
+    }
+    else
+    {
+        return 32 * t;
+    }
+}
+
+// Rot32 / Pair16: which 8-byte piece (0..3) of k-step q lane t feeds.
+template <int kOrder>
+__device__ __forceinline__ int mmaNativePiece(int t, int q)
+{
+    if constexpr (kOrder == kOrderRot32)
+    {
+        return (t - q) & 3;
+    }
+    else
+    {
+        return q < 2 ? t : (t ^ 1);
+    }
+}
+
+template <int kLoad>
+__device__ __forceinline__ void load16(uint64_t& a, uint64_t& b, void const* ptr, uint64_t policy)
+{
+    if constexpr (kLoad == kLoad32Hint)
+    {
+        asm volatile("ld.global.nc.L1::no_allocate.L2::cache_hint.v2.b64 {%0, %1}, [%2], %3;"
+                     : "=l"(a), "=l"(b)
+                     : "l"(ptr), "l"(policy));
+    }
+    else
+    {
+        asm volatile("ld.global.nc.L1::no_allocate.v2.b64 {%0, %1}, [%2];" : "=l"(a), "=l"(b) : "l"(ptr));
+    }
+}
+
+template <int kLoad>
+__device__ __forceinline__ uint64_t load8(void const* ptr, uint64_t policy)
+{
+    uint64_t v;
+    if constexpr (kLoad == kLoad32Hint)
+    {
+        asm volatile("ld.global.nc.L1::no_allocate.L2::cache_hint.b64 %0, [%1], %2;" : "=l"(v) : "l"(ptr), "l"(policy));
+    }
+    else
+    {
+        asm volatile("ld.global.nc.L1::no_allocate.b64 %0, [%1];" : "=l"(v) : "l"(ptr));
+    }
+    return v;
+}
+
+// row0 / row1 already include the lane offset (mmaLaneOffset).
+template <int kStage, int kLoad, int kOrder>
+__device__ __forceinline__ void mmaLoadStage(
+    Chunk (&x)[kStage], Chunk (&y)[kStage], char const* row0, char const* row1, int block, uint64_t policy)
+{
+#pragma unroll
+    for (int j = 0; j < kStage; ++j)
+    {
+        int const off = (block + j) * kMmaBlockElems * 2;
+        if constexpr (kOrder == kOrderNat8)
+        {
+#pragma unroll
+            for (int q = 0; q < 4; ++q)
+            {
+                x[j].v[q] = load8<kLoad>(row0 + off + 32 * q, policy);
+                y[j].v[q] = load8<kLoad>(row1 + off + 32 * q, policy);
+            }
+        }
+        else if constexpr (kOrder == kOrderPair16)
+        {
+            // v[0], v[1]: k-step 2 (t & 1); v[2], v[3]: k-step 2 (t & 1) + 1.
+            load16<kLoad>(x[j].v[0], x[j].v[1], row0 + off, policy);
+            load16<kLoad>(x[j].v[2], x[j].v[3], row0 + off + 32, policy);
+            load16<kLoad>(y[j].v[0], y[j].v[1], row1 + off, policy);
+            load16<kLoad>(y[j].v[2], y[j].v[3], row1 + off + 32, policy);
+        }
+        else
+        {
+            loadChunk<kLoad>(x[j], row0 + off, policy);
+            loadChunk<kLoad>(y[j], row1 + off, policy);
+        }
+    }
+}
+
+// a[q]: the weight piece lane t feeds into k-step q, from what it loaded.
+// Every lane of the warp must call this (in-quad shuffles).
+template <int kOrder>
+__device__ __forceinline__ void mmaWeightPieces(uint64_t (&a)[4], Chunk const& c, int t)
+{
+    if constexpr (kOrder == kOrderRot32)
+    {
+        // c = k-step t. Round r takes piece r of k-step (t - r) mod 4 from the
+        // lane that loaded it, so k-step q arrives in round (t - q) mod 4.
+        uint64_t s[4];
+        s[0] = c.v[0];
+        int const quadBase = (threadIdx.x & 31) & ~3;
+#pragma unroll
+        for (int r = 1; r < 4; ++r)
+        {
+            s[r] = __shfl_sync(0xFFFFFFFFu, c.v[r], quadBase | ((t - r) & 3));
+        }
+#pragma unroll
+        for (int q = 0; q < 4; ++q)
+        {
+            int const idx = (t - q) & 3;
+            uint64_t const lo = (idx & 1) ? s[1] : s[0];
+            uint64_t const hi = (idx & 1) ? s[3] : s[2];
+            a[q] = (idx & 2) ? hi : lo;
+        }
+    }
+    else if constexpr (kOrder == kOrderPair16)
+    {
+        // Even lanes hold k-steps 0, 1 and odd lanes k-steps 2, 3 (pieces
+        // 2 (t >> 1) and 2 (t >> 1) + 1 of each); swap the second pieces.
+        bool const odd = (t & 1) != 0;
+        uint64_t const r0 = __shfl_xor_sync(0xFFFFFFFFu, c.v[1], 1);
+        uint64_t const r1 = __shfl_xor_sync(0xFFFFFFFFu, c.v[3], 1);
+        a[0] = odd ? r0 : c.v[0];
+        a[1] = odd ? r1 : c.v[2];
+        a[2] = odd ? c.v[0] : r0;
+        a[3] = odd ? c.v[2] : r1;
+    }
+    else
+    {
+#pragma unroll
+        for (int q = 0; q < 4; ++q)
+        {
+            a[q] = c.v[q];
+        }
+    }
+}
+
+template <int R, int kStage, int kOrder>
+__device__ __forceinline__ void mmaComputeStage(float (&d)[4], Chunk const (&x)[kStage], Chunk const (&y)[kStage],
+    __nv_bfloat16 const* hRow, bool hLive, int block, int t)
+{
+#pragma unroll
+    for (int j = 0; j < kStage; ++j)
+    {
+        uint4 h0 = make_uint4(0, 0, 0, 0);
+        uint4 h1 = make_uint4(0, 0, 0, 0);
+        if (hLive)
+        {
+            if constexpr (kOrder == kOrderNat8)
+            {
+                // Values 16 q + 4 t .. + 3 for q = 0..3: 8 bytes every 32.
+                uint2 const* src = reinterpret_cast<uint2 const*>(hRow + (block + j) * kMmaBlockElems + 4 * t);
+                uint2 const q0 = src[0];
+                uint2 const q1 = src[4];
+                uint2 const q2 = src[8];
+                uint2 const q3 = src[12];
+                h0 = make_uint4(q0.x, q0.y, q1.x, q1.y);
+                h1 = make_uint4(q2.x, q2.y, q3.x, q3.y);
+            }
+            else
+            {
+                // Permuted: values 16 t .. 16 t + 15 of the block. Rot32 /
+                // Pair16: the row is staged in lane order, so the same 32
+                // bytes are lane t's pieces for k-steps 0..3.
+                uint4 const* src = reinterpret_cast<uint4 const*>(hRow + (block + j) * kMmaBlockElems + 16 * t);
+                h0 = src[0];
+                h1 = src[1];
+            }
+        }
+        uint32_t const hb[8] = {h0.x, h0.y, h0.z, h0.w, h1.x, h1.y, h1.z, h1.w};
+        uint64_t xp[4];
+        uint64_t yp[4];
+        mmaWeightPieces<kOrder>(xp, x[j], t);
+        mmaWeightPieces<kOrder>(yp, y[j], t);
+#pragma unroll
+        for (int sIdx = 0; sIdx < 4; ++sIdx)
+        {
+            // Piece s holds the pairs mapped to k positions (2t, 2t + 1) and
+            // (2t + 8, 2t + 9) of k-step s, for A and B alike.
+            uint64_t const xw = xp[sIdx];
+            uint64_t const yw = yp[sIdx];
+            mmaBf16(d, static_cast<uint32_t>(xw), static_cast<uint32_t>(yw), static_cast<uint32_t>(xw >> 32),
+                static_cast<uint32_t>(yw >> 32), hb[2 * sIdx], hb[2 * sIdx + 1]);
+        }
+    }
+}
+
+template <int R, int kWarps, int kStage, int kLoad, int kMinBlocks, int kOrder>
+__global__ void __launch_bounds__(kWarps * 32, kMinBlocks) fusedMmaDraftTailKernel(DraftTailParams p)
+{
+    static_assert(R >= 1 && R <= 8, "n columns of m16n8k16");
+    static_assert(kMmaBlocks % (2 * kStage) == 0, "double-buffered stages must tile K");
+    extern __shared__ uint4 smemDyn[];
+    __nv_bfloat16* smemH = reinterpret_cast<__nv_bfloat16*>(smemDyn); // [R][kMmaHStride]
+
+    int const lane = threadIdx.x & 31;
+    int const warp = threadIdx.x >> 5;
+    int const g = lane >> 2;
+    int const t = lane & 3;
+    int const numRows = p.vocabLocal;
+    int const numTiles = (numRows + 15) / 16;
+    int const warpStride = gridDim.x * kWarps;
+    int tile = blockIdx.x * kWarps + warp;
+
+    uint64_t const policy = kLoad == kLoad32Hint ? evictFirstPolicy() : 0ull;
+    char const* weight = reinterpret_cast<char const*>(p.weight);
+    // Rows beyond the shard read row 0 (valid memory); their results are dropped.
+    auto rowPtr = [&](int row) -> char const*
+    { return weight + static_cast<int64_t>(row < numRows ? row : 0) * (6144 * 2) + mmaLaneOffset<kOrder>(t); };
+
+    Chunk xa[kStage], ya[kStage], xb[kStage], yb[kStage];
+    if (tile < numTiles)
+    {
+        mmaLoadStage<kStage, kLoad, kOrder>(xa, ya, rowPtr(tile * 16 + g), rowPtr(tile * 16 + g + 8), 0, policy);
+    }
+
+    griddepWait();
+    if constexpr (kOrder == kOrderRot32 || kOrder == kOrderPair16)
+    {
+        // Lane order: in each 64-value block, lane t's pieces for k-steps
+        // 0..3 are contiguous at values 16 t .. 16 t + 15.
+        for (int i = threadIdx.x; i < R * kMmaBlocks * 16; i += blockDim.x)
+        {
+            int const r = i / (kMmaBlocks * 16);
+            int const rem = i - r * (kMmaBlocks * 16);
+            int const block = rem >> 4;
+            int const lt = (rem >> 2) & 3;
+            int const q = rem & 3;
+            reinterpret_cast<uint64_t*>(smemH + r * kMmaHStride)[rem] = reinterpret_cast<uint64_t const*>(
+                p.hidden + r * p.hiddenStride)[block * 16 + q * 4 + mmaNativePiece<kOrder>(lt, q)];
+        }
+    }
+    else
+    {
+        for (int i = threadIdx.x; i < R * 6144 / 8; i += blockDim.x)
+        {
+            int const r = i / (6144 / 8);
+            int const c = i - r * (6144 / 8);
+            reinterpret_cast<uint4*>(smemH + r * kMmaHStride)[c]
+                = reinterpret_cast<uint4 const*>(p.hidden + r * p.hiddenStride)[c];
+        }
+    }
+    __syncthreads();
+
+    bool const hLive = g < R;
+    __nv_bfloat16 const* hRow = smemH + (hLive ? g : 0) * kMmaHStride;
+    uint64_t bestJ[2] = {0ull, 0ull};
+
+    for (; tile < numTiles; tile += warpStride)
+    {
+        int const row0 = tile * 16 + g;
+        int const row1 = row0 + 8;
+        char const* p0 = rowPtr(row0);
+        char const* p1 = rowPtr(row1);
+        float d[4] = {0.f, 0.f, 0.f, 0.f};
+        for (int b = 0; b < kMmaBlocks; b += 2 * kStage)
+        {
+            mmaLoadStage<kStage, kLoad, kOrder>(xb, yb, p0, p1, b + kStage, policy);
+            mmaComputeStage<R, kStage, kOrder>(d, xa, ya, hRow, hLive, b, t);
+            if (b + 2 * kStage < kMmaBlocks)
+            {
+                mmaLoadStage<kStage, kLoad, kOrder>(xa, ya, p0, p1, b + 2 * kStage, policy);
+            }
+            else if (tile + warpStride < numTiles)
+            {
+                int const next = tile + warpStride;
+                mmaLoadStage<kStage, kLoad, kOrder>(
+                    xa, ya, rowPtr(next * 16 + g), rowPtr(next * 16 + g + 8), 0, policy);
+            }
+            mmaComputeStage<R, kStage, kOrder>(d, xb, yb, hRow, hLive, b + kStage, t);
+        }
+        // d[j] = (row0, column 2t + j), d[2 + j] = (row1, column 2t + j).
+#pragma unroll
+        for (int j = 0; j < 2; ++j)
+        {
+            int const n = 2 * t + j;
+            if (n < R)
+            {
+                if (row0 < numRows)
+                {
+                    __nv_bfloat16 const v = __float2bfloat16_rn(d[j]);
+                    bestJ[j] = umax64(bestJ[j], makeKey(__bfloat16_as_ushort(v), p.vocabOffset + row0));
+                    if (p.logitsOut != nullptr)
+                    {
+                        p.logitsOut[n * p.logitsOutStride + row0] = v;
+                    }
+                }
+                if (row1 < numRows)
+                {
+                    __nv_bfloat16 const v = __float2bfloat16_rn(d[2 + j]);
+                    bestJ[j] = umax64(bestJ[j], makeKey(__bfloat16_as_ushort(v), p.vocabOffset + row1));
+                    if (p.logitsOut != nullptr)
+                    {
+                        p.logitsOut[n * p.logitsOutStride + row1] = v;
+                    }
+                }
+            }
+        }
+    }
+    // Reduce over g (lane bits 2..4), then lane 0 collects column n from lane n / 2.
+#pragma unroll
+    for (int j = 0; j < 2; ++j)
+    {
+#pragma unroll
+        for (int off = 4; off < 32; off <<= 1)
+        {
+            bestJ[j] = umax64(bestJ[j], __shfl_xor_sync(0xFFFFFFFFu, bestJ[j], off));
+        }
+    }
+    uint64_t best[R];
+#pragma unroll
+    for (int n = 0; n < R; ++n)
+    {
+        best[n] = __shfl_sync(0xFFFFFFFFu, bestJ[n & 1], n >> 1);
+    }
+    ctaFinish<R, kWarps>(p, best);
+}
+
+// ---------------------------------------------------------------------------
+// Logits kernel: rank-local max over precomputed BF16 logits + exchange.
+// ---------------------------------------------------------------------------
+template <int R>
+__global__ void __launch_bounds__(512, 1) fromLogitsDraftTailKernel(DraftTailParams p)
+{
+    __shared__ uint64_t warpBest[32][R];
+    __shared__ uint64_t ctaKeys[kMaxRows];
+    __shared__ uint64_t gathered[kMaxRanks][kMaxRows];
+
+    griddepWait();
+    int const lane = threadIdx.x & 31;
+    int const warp = threadIdx.x >> 5;
+    int const numWarps = blockDim.x >> 5;
+    uint64_t best[R];
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+    {
+        best[r] = 0ull;
+    }
+    // Flat (row, 8-element vector) index space; each thread issues kBatch
+    // independent 16-byte loads before reducing them, so a row costs about one
+    // memory round trip instead of one per loop iteration.
+    constexpr int kBatch = 8;
+    bool vec = (p.logitsStride % 8) == 0 && (reinterpret_cast<uintptr_t>(p.logits) & 15) == 0;
+    int const nVec = vec ? p.vocabLocal / 8 : 0;
+    int const total = R * nVec;
+    for (int base = threadIdx.x; base < total; base += kBatch * blockDim.x)
+    {
+        uint4 x[kBatch];
+#pragma unroll
+        for (int b = 0; b < kBatch; ++b)
+        {
+            int const i = base + b * blockDim.x;
+            if (i < total)
+            {
+                int const r = i / nVec;
+                int const v = i - r * nVec;
+                x[b] = __ldg(reinterpret_cast<uint4 const*>(p.logits + r * p.logitsStride) + v);
+            }
+        }
+#pragma unroll
+        for (int b = 0; b < kBatch; ++b)
+        {
+            int const i = base + b * blockDim.x;
+            if (i < total)
+            {
+                int const r = i / nVec;
+                int const v = i - r * nVec;
+                uint32_t const w[4] = {x[b].x, x[b].y, x[b].z, x[b].w};
+                uint64_t m = 0ull;
+#pragma unroll
+                for (int k = 0; k < 4; ++k)
+                {
+                    int64_t const id = p.vocabOffset + 8 * v + 2 * k;
+                    m = umax64(m, makeKey(w[k] & 0xFFFFu, id));
+                    m = umax64(m, makeKey(w[k] >> 16, id + 1));
+                }
+#pragma unroll
+                for (int rr = 0; rr < R; ++rr)
+                {
+                    if (rr == r)
+                    {
+                        best[rr] = umax64(best[rr], m);
+                    }
+                }
+            }
+        }
+    }
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+    {
+        __nv_bfloat16 const* row = p.logits + r * p.logitsStride;
+        for (int i = nVec * 8 + threadIdx.x; i < p.vocabLocal; i += blockDim.x)
+        {
+            uint32_t const bits = __bfloat16_as_ushort(row[i]);
+            best[r] = umax64(best[r], makeKey(bits, p.vocabOffset + i));
+        }
+        best[r] = warpMax64(best[r]);
+    }
+    if (lane == 0)
+    {
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            warpBest[warp][r] = best[r];
+        }
+    }
+    __syncthreads();
+    if (warp == 0)
+    {
+#pragma unroll
+        for (int r = 0; r < R; ++r)
+        {
+            uint64_t m = lane < numWarps ? warpBest[lane][r] : 0ull;
+            m = warpMax64(m);
+            if (lane == 0)
+            {
+                ctaKeys[r] = m;
+            }
+        }
+    }
+    __syncthreads();
+    if (isDirectMode(p.mode))
+    {
+        directExchange(p, ctaKeys, 0, 1);
+        return;
+    }
+    if (warp != 0)
+    {
+        return;
+    }
+    exchangeAndWrite(p, ctaKeys, gathered);
+}
+
+// ---------------------------------------------------------------------------
+// Launch
+// ---------------------------------------------------------------------------
+
+constexpr int kHidden6144Chunks = 6144 / 512; // 32-byte chunks per lane per row
+
+struct VariantDesc
+{
+    int warps;
+    int split;
+    int load;
+    int minBlocks;
+};
+
+// Kernel shape and load flavour, selectable at run time for tuning.
+constexpr VariantDesc kVariants[] = {
+    {8, 2, kLoad32Hint, 1},  // 0: default
+    {8, 2, kLoad32, 1},      // 1: no L2 eviction hint
+    {8, 2, kLoad16x2Pf, 1},  // 2: 16-byte loads with L2 prefetch
+    {16, 3, kLoad32Hint, 1}, // 3: more warps, shorter units
+    {4, 2, kLoad32Hint, 2},  // 4: two CTAs per SM
+    {8, 1, kLoad32Hint, 1},  // 5: whole-row units
+    {8, 0, kLoad32Hint, 1},  // 6: per-warp cp.async.bulk row ring (R <= 2; else variant 0)
+    {8, 0, kLoad32Hint, 1},  // 7: tensor-core tiles, 8 warps, 2-block stages
+    {4, 0, kLoad32Hint, 2},  // 8: tensor-core tiles, 4 warps, 2 CTAs/SM
+    {8, 0, kLoad32, 1},      // 9: tensor-core tiles, no L2 hint
+    {8, 0, kLoad32Hint, 1},  // 10: tensor-core tiles, 8 warps, 4-block stages
+    {8, 0, kLoad32Hint, 1},  // 11: as 10, native K order (kOrderNat8: 8-byte lane loads)
+    {4, 0, kLoad32Hint, 2},  // 12: native K order (kOrderNat8), 4 warps, 2 CTAs/SM
+    {8, 0, kLoad32Hint, 1},  // 13: as 10, native K order (kOrderRot32: 32-byte loads, quad shuffles)
+    {8, 0, kLoad32Hint, 1},  // 14: as 10, native K order (kOrderPair16: 16-byte loads, pair shuffles)
+};
+constexpr int kNumVariants = sizeof(kVariants) / sizeof(kVariants[0]);
+// Variants from here on are tensor-core kernels; their default grid covers
+// one 16-row tile per warp.
+constexpr int kFirstMmaVariant = 7;
+
+template <int R, int kWarps, int kSplit, int kLoad, int kMinBlocks>
+void launchFusedShape(DraftTailParams const& p, int numCtas, bool pdl, cudaStream_t stream)
+{
+    auto kernel = fusedDraftTailKernel<R, kWarps, kHidden6144Chunks, kSplit, kLoad, kMinBlocks>;
+    size_t const smem = static_cast<size_t>(R) * 6144 * sizeof(__nv_bfloat16);
+    static bool const attrSet = [&]()
+    {
+        TLLM_CUDA_CHECK(
+            cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem)));
+        return true;
+    }();
+    (void) attrSet;
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(numCtas);
+    config.blockDim = dim3(kWarps * 32);
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p));
+}
+
+template <int R, int kWarps, int kSlots>
+void launchBulkShape(DraftTailParams const& p, int numCtas, bool pdl, cudaStream_t stream)
+{
+    auto kernel = fusedBulkDraftTailKernel<R, kWarps, kSlots>;
+    size_t const smem = static_cast<size_t>(R + kWarps * kSlots) * 6144 * sizeof(__nv_bfloat16);
+    static bool const attrSet = [&]()
+    {
+        TLLM_CUDA_CHECK(
+            cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem)));
+        return true;
+    }();
+    (void) attrSet;
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(numCtas);
+    config.blockDim = dim3(kWarps * 32);
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p));
+}
+
+template <int R, int kWarps, int kStage, int kLoad, int kMinBlocks, int kOrder = kOrderPermuted>
+void launchMmaShape(DraftTailParams const& p, int numCtas, bool pdl, cudaStream_t stream)
+{
+    auto kernel = fusedMmaDraftTailKernel<R, kWarps, kStage, kLoad, kMinBlocks, kOrder>;
+    size_t const smem = static_cast<size_t>(R) * kMmaHStride * sizeof(__nv_bfloat16);
+    static bool const attrSet = [&]()
+    {
+        TLLM_CUDA_CHECK(
+            cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem)));
+        return true;
+    }();
+    (void) attrSet;
+    if (numCtas <= 0)
+    {
+        int const numTiles = (p.vocabLocal + 15) / 16;
+        numCtas = (numTiles + kWarps - 1) / kWarps;
+    }
+    TLLM_CHECK_WITH_INFO((p.mode != ExchangeMode::kMulticastDirect && p.mode != ExchangeMode::kUnicastDirect)
+            || numCtas <= kMaxDirectCtas,
+        "glm_moe_dsa_draft_tail: direct exchange supports at most %d CTAs", kMaxDirectCtas);
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(numCtas);
+    config.blockDim = dim3(kWarps * 32);
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p));
+}
+
+template <int R>
+void launchFusedRows(DraftTailParams const& p, int variant, int numCtas, bool pdl, cudaStream_t stream)
+{
+    switch (variant)
+    {
+    case 0: launchFusedShape<R, 8, 2, kLoad32Hint, 1>(p, numCtas, pdl, stream); break;
+    case 1: launchFusedShape<R, 8, 2, kLoad32, 1>(p, numCtas, pdl, stream); break;
+    case 2: launchFusedShape<R, 8, 2, kLoad16x2Pf, 1>(p, numCtas, pdl, stream); break;
+    case 3: launchFusedShape<R, 16, 3, kLoad32Hint, 1>(p, numCtas, pdl, stream); break;
+    case 4: launchFusedShape<R, 4, 2, kLoad32Hint, 2>(p, numCtas, pdl, stream); break;
+    case 5: launchFusedShape<R, 8, 1, kLoad32Hint, 1>(p, numCtas, pdl, stream); break;
+    case 6:
+        if constexpr (R <= 2)
+        {
+            launchBulkShape<R, 8, 2>(p, numCtas, pdl, stream);
+        }
+        else
+        {
+            launchFusedShape<R, 8, 2, kLoad32Hint, 1>(p, numCtas, pdl, stream);
+        }
+        break;
+    case 7: launchMmaShape<R, 8, 2, kLoad32Hint, 1>(p, numCtas, pdl, stream); break;
+    case 8: launchMmaShape<R, 4, 2, kLoad32Hint, 2>(p, numCtas, pdl, stream); break;
+    case 9: launchMmaShape<R, 8, 2, kLoad32, 1>(p, numCtas, pdl, stream); break;
+    case 10: launchMmaShape<R, 8, 4, kLoad32Hint, 1>(p, numCtas, pdl, stream); break;
+    case 11: launchMmaShape<R, 8, 4, kLoad32Hint, 1, kOrderNat8>(p, numCtas, pdl, stream); break;
+    case 12: launchMmaShape<R, 4, 4, kLoad32Hint, 2, kOrderNat8>(p, numCtas, pdl, stream); break;
+    case 13: launchMmaShape<R, 8, 4, kLoad32Hint, 1, kOrderRot32>(p, numCtas, pdl, stream); break;
+    case 14: launchMmaShape<R, 8, 4, kLoad32Hint, 1, kOrderPair16>(p, numCtas, pdl, stream); break;
+    default: TLLM_THROW("glm_moe_dsa_draft_tail: unknown fused variant %d", variant);
+    }
+}
+
+int smCount()
+{
+    static int const count = []()
+    {
+        int dev = 0;
+        TLLM_CUDA_CHECK(cudaGetDevice(&dev));
+        int n = 0;
+        TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&n, cudaDevAttrMultiProcessorCount, dev));
+        return n;
+    }();
+    return count;
+}
+
+} // namespace
+
+int numFusedVariants()
+{
+    return kNumVariants;
+}
+
+void invokeDraftTailFused(DraftTailParams const& p, FusedLaunchConfig const& config, cudaStream_t stream)
+{
+    TLLM_CHECK_WITH_INFO(p.hiddenDim == 6144, "glm_moe_dsa_draft_tail: fused kernel supports hidden size 6144 only");
+    TLLM_CHECK_WITH_INFO(p.rows >= 1 && p.rows <= kMaxRows, "glm_moe_dsa_draft_tail: rows out of range");
+    TLLM_CHECK_WITH_INFO(config.variant >= 0 && config.variant < kNumVariants, "glm_moe_dsa_draft_tail: bad variant");
+    TLLM_CHECK_WITH_INFO(p.logitsOut == nullptr || config.variant >= kFirstMmaVariant,
+        "glm_moe_dsa_draft_tail: logits output needs a tensor-core variant");
+    int numCtas = config.numCtas;
+    if (numCtas <= 0 && config.variant < kFirstMmaVariant)
+    {
+        numCtas = smCount() * kVariants[config.variant].minBlocks;
+    }
+    TLLM_CHECK_WITH_INFO((p.mode != ExchangeMode::kMulticastDirect && p.mode != ExchangeMode::kUnicastDirect)
+            || numCtas <= kMaxDirectCtas,
+        "glm_moe_dsa_draft_tail: direct exchange supports at most %d CTAs", kMaxDirectCtas);
+    switch (p.rows)
+    {
+    case 1: launchFusedRows<1>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 2: launchFusedRows<2>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 3: launchFusedRows<3>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 4: launchFusedRows<4>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 5: launchFusedRows<5>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 6: launchFusedRows<6>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 7: launchFusedRows<7>(p, config.variant, numCtas, config.pdl, stream); break;
+    case 8: launchFusedRows<8>(p, config.variant, numCtas, config.pdl, stream); break;
+    default: break;
+    }
+}
+
+template <int R>
+void launchFromLogitsRows(DraftTailParams const& p, bool pdl, cudaStream_t stream)
+{
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(1);
+    config.blockDim = dim3(512);
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, fromLogitsDraftTailKernel<R>, p));
+}
+
+void invokeDraftTailFromLogits(DraftTailParams const& p, bool pdl, cudaStream_t stream)
+{
+    TLLM_CHECK_WITH_INFO(p.rows >= 1 && p.rows <= kMaxRows, "glm_moe_dsa_draft_tail: rows out of range");
+    switch (p.rows)
+    {
+    case 1: launchFromLogitsRows<1>(p, pdl, stream); break;
+    case 2: launchFromLogitsRows<2>(p, pdl, stream); break;
+    case 3: launchFromLogitsRows<3>(p, pdl, stream); break;
+    case 4: launchFromLogitsRows<4>(p, pdl, stream); break;
+    case 5: launchFromLogitsRows<5>(p, pdl, stream); break;
+    case 6: launchFromLogitsRows<6>(p, pdl, stream); break;
+    case 7: launchFromLogitsRows<7>(p, pdl, stream); break;
+    case 8: launchFromLogitsRows<8>(p, pdl, stream); break;
+    default: break;
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::draft_tail
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h
new file mode 100644
index 000000000000..82e13a7f084a
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h
@@ -0,0 +1,144 @@
+/*
+ * 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.
+ */
+
+// Greedy draft-token selection over a vocab-sharded LM head.
+//
+// For R hidden rows and this rank's vocab shard W [V_local, H] the kernels
+// produce, on every rank, the token id of the global maximum of the BF16
+// logits h @ W^T over all tensor-parallel shards; ties go to the lowest global
+// token id and NaN ranks above every number. That is the selection rule of
+// SpecWorkerBase.greedy_sample_draft_with_tp_gather.
+//
+// The logits kernel selects from BF16 logits computed by the caller (the
+// native LM-head GEMM), so its tokens are the native ones. The fused kernels
+// compute the logits themselves (FP32 accumulation, then BF16 rounding);
+// where the largest logits are within one BF16 rounding step of each other,
+// the selected token depends on the exact FP32 sums, i.e. on the accumulation
+// order. Fused variants 11 to 14 accumulate like the native cuBLAS LM-head
+// GEMM selected for this shape on SM103 (K in 16-value tensor-core steps from
+// 0 to H into one FP32 accumulator), so their BF16 logits and tokens equal the
+// native ones bit for bit (checked by the draft_tail unit tests against F.linear;
+// re-check whenever cuBLAS or the GPU changes). The other fused variants
+// accumulate in other orders and are decision-level variants at BF16
+// near-ties.
+//
+// Each candidate is packed into one 64-bit key, (ordered BF16 bits << 32) |
+// (0xFFFFFFFF - global id), so the selection is a plain unsigned max that is
+// associative, commutative and independent of reduction order. The keys of
+// all ranks are exchanged through a multicast (or unicast peer) buffer with a
+// 0-sentinel protocol over three rotating slots; see draftTail.cu.
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::draft_tail
+{
+
+constexpr int kMaxRows = 8;
+constexpr int kMaxRanks = 8;
+constexpr int kNumSlots = 3;
+// Exchange buffer: kNumSlots x kMaxRanks x kMaxRows 64-bit keys.
+constexpr int kSlotWords = kMaxRanks * kMaxRows;
+constexpr int kExchangeWords = kNumSlots * kSlotWords;
+// Rank-local scratch, in 64-bit words: kMaxRows partial keys, then the CTA
+// arrival counter and the exchange epoch (32 bits each in one word), then an
+// error word.
+constexpr int kScratchWords = kMaxRows + 2;
+
+enum class ExchangeMode : int
+{
+    kNone = 0,      // tpSize == 1, or debug: tokens from the rank-local maximum
+    kMulticast = 1, // one multimem.st per row into the multicast slot
+    kUnicast = 2,   // one store per (peer, row) into each peer's slot
+    // Direct variants: every CTA publishes its own partial key (no rank-local
+    // atomics or last-CTA election); CTA 0 of each rank polls all
+    // tpSize x numCtas entries, writes the tokens and recycles the slot.
+    kMulticastDirect = 3,
+    kUnicastDirect = 4,
+};
+
+constexpr int kMaxDirectCtas = 512;
+constexpr int kDirectSlotWords = kMaxRows * kMaxRanks * kMaxDirectCtas;
+// The direct region follows the last-CTA region in the same buffer.
+constexpr int kDirectOffsetWords = kExchangeWords;
+constexpr int kExchangeBufferWords = kExchangeWords + kNumSlots * kDirectSlotWords;
+
+struct DraftTailParams
+{
+    // Fused mode: hidden [rows, hiddenDim] (row stride hiddenStride elements)
+    // and the vocab shard weight [vocabLocal, hiddenDim], row-major.
+    __nv_bfloat16 const* hidden = nullptr;
+    int64_t hiddenStride = 0;
+    __nv_bfloat16 const* weight = nullptr;
+    // Logits mode: BF16 logits [rows, vocabLocal] (row stride logitsStride).
+    __nv_bfloat16 const* logits = nullptr;
+    int64_t logitsStride = 0;
+
+    int rows = 0;
+    int vocabLocal = 0;
+    int hiddenDim = 0;
+    int64_t vocabOffset = 0;     // global id of local column 0
+
+    int32_t* tokens = nullptr;   // [rows] output
+    uint64_t* keysOut = nullptr; // optional [rows] output: the winning key
+
+    uint64_t* scratch = nullptr; // [kScratchWords], zero-initialized once
+
+    int tpRank = 0;
+    int tpSize = 1;
+    ExchangeMode mode = ExchangeMode::kNone;
+    uint64_t* mcSlots = nullptr;             // multicast address of the exchange buffer
+    uint64_t* ucSlotsLocal = nullptr;        // this rank's unicast address of the buffer
+    uint64_t* const* ucSlotsPeers = nullptr; // device array [tpSize] of unicast addresses
+    long long timeoutCycles = 0;             // 0: wait forever
+    // Optional BF16 logits [rows, vocabLocal] (row stride logitsOutStride),
+    // written by the tensor-core fused kernels only (tests compare them with
+    // the native GEMM bit for bit); nullptr in production.
+    __nv_bfloat16* logitsOut = nullptr;
+    int64_t logitsOutStride = 0;
+    // Fused kernel, L2-hinted variants: each warp loads the first
+    // keepPermille / 1000 of its rows with an L2 evict_last policy (the rest
+    // evict_first), so that part of the shard tends to stay in L2 across the
+    // draft steps that re-read the same weight.
+    int keepPermille = 0;
+};
+
+struct FusedLaunchConfig
+{
+    int numCtas = 0; // 0: one per SM
+    int variant = 0; // see draftTail.cu: kernel shape and load flavour
+    bool pdl = true;
+};
+
+int numFusedVariants();
+
+// GEMV + rank-local max + exchange in one launch.
+void invokeDraftTailFused(DraftTailParams const& params, FusedLaunchConfig const& config, cudaStream_t stream);
+
+// Rank-local max over precomputed BF16 logits + exchange in one launch.
+void invokeDraftTailFromLogits(DraftTailParams const& params, bool pdl, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::draft_tail
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu
new file mode 100644
index 000000000000..8d21b5092336
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu
@@ -0,0 +1,439 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h"
+
+#include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h"
+
+#include <cuda_fp8.h>
+
+#include <algorithm>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::indexer_layer_front
+{
+
+using glm_moe_dsa::rope_kv_append::kKvLoraRank;
+using glm_moe_dsa::rope_kv_append::kLatentDim;
+using glm_moe_dsa::rope_kv_append::kRopeDim;
+using glm_moe_dsa::rope_kv_append::RopeAppendParams;
+
+namespace
+{
+
+constexpr int32_t kRopeThreads = 256;
+constexpr int32_t kEltsPerVec = 8;                           // bf16 elements per 16-byte load
+constexpr int32_t kNopeVecs = kKvLoraRank / kEltsPerVec;     // 64
+constexpr int32_t kRopeVecs = kRopeDim / kEltsPerVec;        // 8
+constexpr int32_t kTotalVecsPerHead = kNopeVecs + kRopeVecs; // 72
+constexpr int32_t kIdxBytesPerItem = 4;
+
+__device__ __forceinline__ void pdlWait()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+#endif
+}
+
+__device__ __forceinline__ void pdlLaunchDependents()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+#endif
+}
+
+//! Same arithmetic as quantCopy<__nv_bfloat16, 8> in mlaKernels.cu: bf16 -> fp32, scale, saturating E4M3.
+__device__ __forceinline__ uint2 quantizeBf16x8(uint4 const& src, float scale)
+{
+    __nv_bfloat162 const* in = reinterpret_cast<__nv_bfloat162 const*>(&src);
+    uint2 out;
+    __nv_fp8x2_e4m3* dst = reinterpret_cast<__nv_fp8x2_e4m3*>(&out);
+#pragma unroll
+    for (int32_t j = 0; j < 4; ++j)
+    {
+        float2 v = __bfloat1622float2(in[j]);
+        v.x *= scale;
+        v.y *= scale;
+        dst[j] = __nv_fp8x2_e4m3(v);
+    }
+    return out;
+}
+
+//! GPT-J RoPE on 8 bf16 values (4 adjacent pairs), identical to the per-pair transform of the native kernel.
+__device__ __forceinline__ void rotateBf16x8(uint4& data, float2 const (&coef)[4])
+{
+    __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&data);
+#pragma unroll
+    for (int32_t j = 0; j < 4; ++j)
+    {
+        pairs[j] = mmha::rotary_embedding_transform(pairs[j], coef[j]);
+    }
+}
+
+//! Native unravel of a flat byte index over the strided 4-D indexer cache view (indexerKCacheScatter.cu);
+//! 32-bit arithmetic when the index fits, same result.
+__device__ __forceinline__ int64_t idxFlatToOffset(int64_t flat, int32_t const (&dims)[4], int64_t const (&strides)[4])
+{
+    int32_t i0, i1, i2, i3;
+    if (flat >= 0 && flat <= 0x7fffffff)
+    {
+        uint32_t f = static_cast<uint32_t>(flat);
+        i3 = static_cast<int32_t>(f % static_cast<uint32_t>(dims[3]));
+        f /= static_cast<uint32_t>(dims[3]);
+        i2 = static_cast<int32_t>(f % static_cast<uint32_t>(dims[2]));
+        f /= static_cast<uint32_t>(dims[2]);
+        i1 = static_cast<int32_t>(f % static_cast<uint32_t>(dims[1]));
+        f /= static_cast<uint32_t>(dims[1]);
+        i0 = static_cast<int32_t>(f);
+    }
+    else
+    {
+        i3 = flat % dims[3];
+        flat /= dims[3];
+        i2 = flat % dims[2];
+        flat /= dims[2];
+        i1 = flat % dims[1];
+        flat /= dims[1];
+        i0 = flat;
+    }
+    return i0 * strides[0] + i1 * strides[1] + i2 * strides[2] + i3 * strides[3];
+}
+
+enum class RopeItem : int32_t
+{
+    kNone = 0,
+    kQNope,
+    kQRope,
+    kKRope,
+    kKNope,
+    kIdxData,
+    kIdxScale,
+};
+
+//! One thread per 16-byte work item of one token. Per token the item space is
+//!   [0, H*64)            q_nope FP8 quantization      (head = i / 64, vec = i % 64)
+//!   [H*64, H*72)         q_pe RoPE + FP8               (head, vec)
+//!   [H*72, H*72+8)       k_pe RoPE + FP8 cache write   (vec)
+//!   [H*72+8, H*72+72)    compressed-KV FP8 cache write (vec)
+//!   then idxHeadBytes/4 indexer K data items (4 bytes each) and one indexer scale item.
+//! Per-step metadata (cache lengths, KV block offsets, indexer slot mappings) and the item's source data are read only
+//! after griddepcontrol.wait (see the header). Identical to ropeAppendScatterKernel except for the trigger position.
+template <bool kEarlyTrigger>
+__global__ void __launch_bounds__(kRopeThreads) ropeAppendScatterZ14Kernel(RopeAppendParams const p)
+{
+    int32_t const token = blockIdx.x;
+    int32_t const item = blockIdx.y * kRopeThreads + threadIdx.x;
+    int32_t const qNopeEnd = p.numHeads * kNopeVecs;
+    int32_t const qRopeEnd = qNopeEnd + p.numHeads * kRopeVecs;
+    int32_t const kRopeEnd = qRopeEnd + kRopeVecs;
+    int32_t const kNopeEnd = kRopeEnd + kNopeVecs;
+    int32_t const idxDataEnd = kNopeEnd + p.idxHeadBytes / kIdxBytesPerItem;
+
+    bool const ropeToken = token < p.numTokens;
+    bool const idxToken = token < p.idxNumTokens;
+
+    RopeItem kind = RopeItem::kNone;
+    int32_t head = 0;
+    int32_t vec = 0;
+    if (item < qNopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kQNope : RopeItem::kNone;
+        head = item / kNopeVecs;
+        vec = item % kNopeVecs;
+    }
+    else if (item < qRopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kQRope : RopeItem::kNone;
+        head = (item - qNopeEnd) / kRopeVecs;
+        vec = (item - qNopeEnd) % kRopeVecs;
+    }
+    else if (item < kRopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kKRope : RopeItem::kNone;
+        vec = item - qRopeEnd;
+    }
+    else if (item < kNopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kKNope : RopeItem::kNone;
+        vec = item - kRopeEnd;
+    }
+    else if (item < idxDataEnd)
+    {
+        kind = idxToken ? RopeItem::kIdxData : RopeItem::kNone;
+        vec = item - kNopeEnd;
+    }
+    else if (item == idxDataEnd)
+    {
+        kind = idxToken ? RopeItem::kIdxScale : RopeItem::kNone;
+    }
+
+    int32_t const batchIdx = token / p.seqLen;
+    int32_t const localIdx = token % p.seqLen;
+    bool const needPos = kind == RopeItem::kQRope || kind == RopeItem::kKRope || kind == RopeItem::kKNope;
+    bool const needCoef = kind == RopeItem::kQRope || kind == RopeItem::kKRope;
+    bool const needBlock = kind == RopeItem::kKRope || kind == RopeItem::kKNope;
+    bool const isIdx = kind == RopeItem::kIdxData || kind == RopeItem::kIdxScale;
+
+    // Values that never change while kernels of the step run (the quantization scales) may be read before the wait;
+    // per-step metadata and producer outputs may not.
+    float scale = 1.F;
+    if (needBlock)
+    {
+        scale = p.quantScaleKv ? p.quantScaleKv[0] : 1.F;
+    }
+    else if (kind == RopeItem::kQNope || kind == RopeItem::kQRope)
+    {
+        scale = p.quantScaleQ ? p.quantScaleQ[0] : 1.F;
+    }
+    auto loadCoef = [&](int32_t pos, float2(&c)[4])
+    {
+        float2 const* src = p.cosSin + static_cast<size_t>(kRopeDim) * pos + (vec * kEltsPerVec) / 2;
+#pragma unroll
+        for (int32_t j = 0; j < 4; ++j)
+        {
+            c[j] = src[j];
+        }
+    };
+    auto idxOffset = [&](int64_t slotData, int64_t slotScale) -> int64_t
+    {
+        if (slotData < 0 || slotScale < 0)
+        {
+            return -1;
+        }
+        return kind == RopeItem::kIdxData ? idxFlatToOffset(slotData + vec * kIdxBytesPerItem, p.idxDims, p.idxStrides)
+                                          : idxFlatToOffset(slotScale, p.idxDims, p.idxStrides);
+    };
+    // Source data of this item; written by producers, so it is read after the wait.
+    uint4 vdata = make_uint4(0, 0, 0, 0);
+    uint32_t wdata = 0;
+    auto loadSource = [&]()
+    {
+        switch (kind)
+        {
+        case RopeItem::kQNope:
+            vdata = *reinterpret_cast<uint4 const*>(
+                p.fusedQ + (static_cast<size_t>(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec);
+            break;
+        case RopeItem::kQRope:
+            vdata = *reinterpret_cast<uint4 const*>(p.qPe + static_cast<size_t>(token) * p.qPeStride
+                + static_cast<size_t>(p.qPeLd) * head + vec * kEltsPerVec);
+            break;
+        case RopeItem::kKRope:
+            vdata = *reinterpret_cast<uint4 const*>(
+                p.latent + static_cast<size_t>(token) * kLatentDim + kKvLoraRank + vec * kEltsPerVec);
+            break;
+        case RopeItem::kKNope:
+            vdata = *reinterpret_cast<uint4 const*>(
+                p.latent + static_cast<size_t>(token) * kLatentDim + vec * kEltsPerVec);
+            break;
+        case RopeItem::kIdxData:
+            wdata = *reinterpret_cast<uint32_t const*>(
+                p.idxK + static_cast<int64_t>(token) * p.idxHeadBytes + vec * kIdxBytesPerItem);
+            break;
+        case RopeItem::kIdxScale:
+            wdata = *reinterpret_cast<uint32_t const*>(p.idxScale + static_cast<int64_t>(token) * kIdxBytesPerItem);
+            break;
+        default: break;
+        }
+    };
+
+    int32_t kvPos = 0;
+    float2 coef[4];
+    __nv_fp8_e4m3* kBlock = nullptr;
+    int64_t idxDst = -1;
+    pdlWait();
+    // Issue the source-data load first so it overlaps the metadata loads and the address math below.
+    loadSource();
+    if (needPos)
+    {
+        kvPos = p.kvLens[batchIdx] - p.seqLen + localIdx;
+    }
+    if (needCoef)
+    {
+        loadCoef(kvPos, coef);
+    }
+    if (needBlock)
+    {
+        kBlock = reinterpret_cast<__nv_fp8_e4m3*>(p.kvCache.getKBlockPtr(batchIdx, kvPos));
+    }
+    if (isIdx)
+    {
+        idxDst = idxOffset(p.idxSlotData[token], p.idxSlotScale[token]);
+    }
+    if constexpr (kEarlyTrigger)
+    {
+        // Every global read of this thread's item has been issued. The block counts as triggered at the first of
+        // its threads to get here (later invocations in the block have no further effect), so a dependent may be
+        // resident before this block's stores below and before block (0, 0)'s scheduler-buffer stores and their
+        // cache-length reads. Dependents read this kernel's outputs only after their own griddepcontrol.wait, and
+        // write no global memory before it, so neither the stores nor the reads below can race with them.
+        pdlLaunchDependents();
+    }
+
+    switch (kind)
+    {
+    case RopeItem::kQNope:
+    {
+        size_t const off = (static_cast<size_t>(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec;
+        *reinterpret_cast<uint2*>(p.quantQ + off) = quantizeBf16x8(vdata, scale);
+        break;
+    }
+    case RopeItem::kQRope:
+    {
+        uint4 data = vdata;
+        rotateBf16x8(data, coef);
+        size_t const dst
+            = (static_cast<size_t>(token) * p.numHeads + head) * kLatentDim + kKvLoraRank + vec * kEltsPerVec;
+        *reinterpret_cast<uint2*>(p.quantQ + dst) = quantizeBf16x8(data, scale);
+        break;
+    }
+    case RopeItem::kKRope:
+    {
+        uint4 data = vdata;
+        rotateBf16x8(data, coef);
+        int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, kNopeVecs + vec);
+        *reinterpret_cast<uint2*>(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(data, scale);
+        break;
+    }
+    case RopeItem::kKNope:
+    {
+        int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, vec);
+        *reinterpret_cast<uint2*>(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(vdata, scale);
+        if (vec == 0 && localIdx == 0)
+        {
+            p.cuQSeqlens[batchIdx + 1] = p.numHeads * p.seqLen * (batchIdx + 1);
+        }
+        break;
+    }
+    case RopeItem::kIdxData:
+    case RopeItem::kIdxScale:
+    {
+        if (idxDst >= 0)
+        {
+            *reinterpret_cast<uint32_t*>(p.idxCache + idxDst) = wdata;
+        }
+        break;
+    }
+    default: break;
+    }
+
+    if (blockIdx.x == 0 && blockIdx.y == 0)
+    {
+        if (threadIdx.x == 0)
+        {
+            p.fmhaTileCounter[0] = 0;
+            p.cuQSeqlens[0] = 0;
+            float const dequantQ = p.dequantScaleQ ? p.dequantScaleQ[0] : 1.F;
+            float const dequantKv = p.dequantScaleKv ? p.dequantScaleKv[0] : 1.F;
+            float const quantO = p.quantScaleO ? p.quantScaleO[0] : 1.F;
+            if (p.bmm1Scale)
+            {
+                constexpr float kLog2e = 1.4426950408889634074F;
+                float const bmm1 = dequantQ * dequantKv * p.hostBmm1Scale;
+                p.bmm1Scale[0] = bmm1;
+                p.bmm1Scale[1] = bmm1 * kLog2e;
+            }
+            if (p.bmm2Scale)
+            {
+                p.bmm2Scale[0] = quantO * dequantKv;
+            }
+        }
+        if (threadIdx.x < 32)
+        {
+            // Exclusive prefix sum of the cache lengths -> cuKvSeqlens[0..batch].
+            int32_t const lane = threadIdx.x;
+            int32_t const batch = p.numTokens / p.seqLen;
+            int32_t carry = 0;
+            for (int32_t base = 0; base <= batch; base += 32)
+            {
+                int32_t const idx = base + lane;
+                int32_t const len = idx < batch ? p.kvLens[idx] : 0;
+                int32_t incl = len;
+#pragma unroll
+                for (int32_t off = 1; off < 32; off <<= 1)
+                {
+                    int32_t const n = __shfl_up_sync(0xffffffffU, incl, off);
+                    if (lane >= off)
+                    {
+                        incl += n;
+                    }
+                }
+                if (idx <= batch)
+                {
+                    p.cuKvSeqlens[idx] = incl - len + carry;
+                }
+                carry += __shfl_sync(0xffffffffU, incl, 31);
+            }
+        }
+    }
+
+    if constexpr (!kEarlyTrigger)
+    {
+        // ropeAppendScatterKernel's position.
+        pdlLaunchDependents();
+    }
+}
+
+template <typename Kernel, typename Params>
+void launchWithOptionalPdl(Kernel kernel, dim3 grid, dim3 block, bool enablePdl, cudaStream_t stream, Params const& p)
+{
+    cudaLaunchConfig_t config{};
+    config.gridDim = grid;
+    config.blockDim = block;
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0;
+    config.numAttrs = 1;
+    config.attrs = attrs;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p));
+}
+
+} // namespace
+
+void invokeIndexerLayerRopeKvAppend(
+    RopeAppendParams const& params, bool enablePdl, bool earlyTrigger, cudaStream_t stream)
+{
+    int32_t const gridTokens = std::max(params.numTokens, params.idxNumTokens);
+    if (gridTokens == 0)
+    {
+        return;
+    }
+    TLLM_CHECK_WITH_INFO(params.seqLen > 0, "seqLen must be positive");
+    TLLM_CHECK_WITH_INFO(params.cosSinRows > 0, "the RoPE table must have at least one position");
+    TLLM_CHECK_WITH_INFO(params.idxHeadBytes % kIdxBytesPerItem == 0, "indexer head bytes must be a multiple of 4");
+    int32_t const itemsPerToken
+        = params.numHeads * kTotalVecsPerHead + kTotalVecsPerHead + params.idxHeadBytes / kIdxBytesPerItem + 1;
+    dim3 const grid(gridTokens, (itemsPerToken + kRopeThreads - 1) / kRopeThreads);
+    dim3 const block(kRopeThreads);
+    if (earlyTrigger)
+    {
+        launchWithOptionalPdl(ropeAppendScatterZ14Kernel<true>, grid, block, enablePdl, stream, params);
+    }
+    else
+    {
+        launchWithOptionalPdl(ropeAppendScatterZ14Kernel<false>, grid, block, enablePdl, stream, params);
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::indexer_layer_front
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h
new file mode 100644
index 000000000000..20a8e8b530de
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h
@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h"
+
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::indexer_layer_front
+{
+
+//! Fork of the generation rope of ropeKvAppend.h (ropeAppendScatterKernel): the same bytes for the same parameters
+//! (glm_moe_dsa::rope_kv_append::RopeAppendParams), with the programmatic-launch trigger moved.
+//!
+//! Memory order: like ropeAppendScatterKernel, only the quantization scales are read before griddepcontrol.wait; every
+//! producer output and every per-step metadata word is read after it. With earlyTrigger every thread executes
+//! griddepcontrol.launch_dependents right after the global reads of its work item (source item, cache length, rotary
+//! coefficients, KV block offset and indexer slot mapping), before its stores. A block counts as triggered at its first
+//! such invocation, so dependents may become resident before any store of the kernel, including block (0, 0)'s
+//! scheduler buffers (cuQSeqlens, cuKvSeqlens, the tile counter, the bmm scales) and the cache-length reads of its
+//! prefix sum. Without earlyTrigger the trigger is at the end of the kernel, as there. A dependent must not read
+//! anything this kernel writes, nor write anything it reads, before its own griddepcontrol.wait; the rope's dependents
+//! already require this, since a trigger at the end of the code does not make the stores visible either.
+void invokeIndexerLayerRopeKvAppend(glm_moe_dsa::rope_kv_append::RopeAppendParams const& params, bool enablePdl,
+    bool earlyTrigger, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::indexer_layer_front
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu
new file mode 100644
index 000000000000..5f85a965e1d7
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu
@@ -0,0 +1,680 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h"
+
+#include <cooperative_groups.h>
+#include <cuda_bf16.h>
+#include <type_traits>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::post_attention_hop
+{
+namespace
+{
+namespace cg = cooperative_groups;
+
+// Work decomposition.
+//
+// Every cluster has kSlices CTAs; CTA c owns slice [c * kSlice, (c + 1) * kSlice), the slice a CTA of the native
+// one-shot kernel owns for numRows <= 8 (96 threads x 8 elements, 8 CTAs per row), so the RMSNorm reduction below
+// reproduces the native tree exactly.
+//
+// Without router logits (dense layers) the layout is the native one: cluster y handles row y, and every CTA collects
+// one row slice (8 ranks x 768 elements).
+//
+// With router logits there are kG x kRowSplit clusters (kG = 4, kRowSplit = min(R, 3): at most 12 clusters, which stay
+// co-resident). Cluster (g, h) collects and normalizes the rows r with r % kRowSplit == h, so a CTA reads at most
+// ceil(R / kRowSplit) row slices, and computes the logits of those rows for experts [64 g, 64 g + 64). Its router
+// weights live in registers (loaded before the PDL wait). Each CTA multiplies its rows of slice c by the experts'
+// weights of slice c with BF16 MMAs (experts on m16, rows on n8) and pushes the partial logits to the CTA that owns the
+// expert; owners add the 8 slices in a fixed order, so the logits are identical on every rank. Clusters never
+// synchronize with each other; inside a cluster, slice sums and partial logits travel as st.async stores that complete
+// transactions on the consumer's mbarrier.
+constexpr int kSlices = 8;
+constexpr int kSlice = kHiddenDim / kSlices;
+constexpr int kEltsPerVec = 8;
+constexpr int kVecPerSlice = kSlice / kEltsPerVec;
+constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32;
+// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two
+// rows in one quarter-warp phase on disjoint banks.
+constexpr int kNormedStride = kVecPerSlice + 4;
+constexpr int kMmaChains = 4;
+constexpr uint32_t kLamportSentinel = 0x80000000U;
+
+static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0);
+static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32);
+static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension");
+
+// kG = number of expert groups (0: hop without router logits); kRowSplit = number of row groups.
+template <int kG, int kRowSplit>
+struct HopConfig
+{
+    static constexpr int kThreads = 384;
+    static constexpr int kWarps = kThreads / 32;
+    static constexpr int kExpertsPerCluster = kNumExperts / kG;
+    static constexpr int kExpertTiles = kExpertsPerCluster / 16;
+    static constexpr int kKGroups = kWarps / kExpertTiles;
+    static constexpr int kKPerGroup = kSlice / kKGroups;
+    static constexpr int kPairsPerGroup = kKPerGroup / 32;
+    static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices;
+    static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0);
+    static_assert(kPairsPerGroup % (kMmaChains / 2) == 0);
+};
+
+template <int kRowSplit>
+struct HopConfig<0, kRowSplit>
+{
+    static constexpr int kThreads = 128;
+    static constexpr int kWarps = kThreads / 32;
+    static constexpr int kMaxRowsPerCta = 1;
+    static_assert(kMaxRowsPerCta * kVecPerSlice <= kThreads && kMaxRowsPerCta * kNativeWarpsPerSlice <= kWarps);
+};
+
+constexpr int kMaxRowsPerCtaAll = kMaxRows / 2;
+
+struct alignas(16) HopStorage
+{
+    uint4 gamma[kVecPerSlice];
+    float vecSq[kMaxRowsPerCtaAll][kVecPerSlice];
+    float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice];
+    float blockSq[kMaxRowsPerCtaAll][kSlices];
+    // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx).
+    alignas(8) uint64_t normBar;
+    // Completes when all partial logits of the experts this CTA owns have arrived.
+    alignas(8) uint64_t routerBar;
+};
+
+template <int kG, int kRowSplit>
+struct alignas(16) RouterStorage
+{
+    using Cfg = HopConfig<kG, kRowSplit>;
+    uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero
+    float partial[kSlices][Cfg::kKGroups][Cfg::kExpertsPerOwner][kMaxRows];
+};
+
+struct EmptyStorage
+{
+};
+
+#ifdef TLLM_GLM_LL_W2_TIMING
+// Development-only phase timestamps (globaltimer, ns) of thread 0 of every CTA for the most recent launch.
+__device__ unsigned long long gW2Timing[kTimingMaxCtas][kTimingStamps];
+#define W2_STAMP(i)                                                                                                    \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (threadIdx.x == 0)                                                                                          \
+        {                                                                                                              \
+            unsigned long long t;                                                                                      \
+            asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));                                                      \
+            gW2Timing[blockIdx.y * gridDim.x + blockIdx.x][i] = t;                                                     \
+        }                                                                                                              \
+    } while (0)
+#else
+#define W2_STAMP(i)                                                                                                    \
+    do                                                                                                                 \
+    {                                                                                                                  \
+    } while (0)
+#endif
+
+union BfVec
+{
+    uint4 packed;
+    __nv_bfloat16 elts[kEltsPerVec];
+};
+
+__device__ __forceinline__ uint4 loadVolatile(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadConstant(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr));
+    return v;
+}
+
+__device__ __forceinline__ bool isDirty(uint4 const& v)
+{
+    return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel;
+}
+
+// Lamport payloads must never contain the sentinel word: map bf16 -0.0 to +0.0 (the reduction starts from +0.0, so
+// the sum is unchanged).
+__device__ __forceinline__ uint4 sanitize(uint4 v)
+{
+    BfVec b;
+    b.packed = v;
+#pragma unroll
+    for (int i = 0; i < kEltsPerVec; ++i)
+    {
+        if (__bfloat16_as_ushort(b.elts[i]) == 0x8000U)
+        {
+            b.elts[i] = __ushort_as_bfloat16(0);
+        }
+    }
+    return b.packed;
+}
+
+__device__ __forceinline__ uint32_t sharedAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// Address of the same shared-memory object in CTA `rank` of the cluster.
+__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank)
+{
+    uint32_t out;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank));
+    return out;
+}
+
+__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory");
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar)
+{
+    asm volatile(
+        "{\n"
+        ".reg .pred done;\n"
+        "HOP_MBAR_WAIT_%=:\n"
+        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n"
+        "@!done bra HOP_MBAR_WAIT_%=;\n"
+        "}\n" ::"r"(sharedAddr(bar))
+        : "memory");
+}
+
+// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier.
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v),
+                 "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr),
+                 "f"(a), "f"(b), "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mmaBf16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+template <int R, int kG, int kRowSplit>
+__global__ void __launch_bounds__(HopConfig<kG, kRowSplit>::kThreads, 1)
+    postAttnHopKernel(PostAttnHopParams const params)
+{
+    static_assert(R >= 1 && R <= kMaxRows);
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    constexpr bool kRouter = kG > 0;
+    using Cfg = HopConfig<kG, kRowSplit>;
+    constexpr int kClusters = kRouter ? kG * kRowSplit : R;
+    constexpr int kGridCtas = kClusters * kSlices;
+    // Rows one CTA collects, normalizes and multiplies.
+    constexpr int kRowsPerCta = kRouter ? (R + kRowSplit - 1) / kRowSplit : 1;
+    static_assert(kRowsPerCta <= kMaxRowsPerCtaAll);
+    static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps);
+    if constexpr (kRouter)
+    {
+        static_assert(
+            HopConfig<(kG > 0 ? kG : 4), kRowSplit>::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads);
+    }
+
+    __shared__ HopStorage hop;
+    __shared__ std::conditional_t<kRouter, RouterStorage<(kG > 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>, EmptyStorage> rs;
+
+    cg::cluster_group cluster = cg::this_cluster();
+    int const slice = static_cast<int>(cluster.block_rank());
+    int const clusterId = blockIdx.y;
+    int const expertGroup = kRouter ? clusterId % (kG > 0 ? kG : 1) : 0;
+    int const rowGroup = kRouter ? clusterId / (kG > 0 ? kG : 1) : clusterId;
+    int const tid = threadIdx.x;
+    int const lane = tid & 31;
+    int const warp = tid >> 5;
+    int const sliceBase = slice * kSlice;
+
+    // Rows this CTA collects and normalizes: local row j is global row rowGroup + j * rowStride.
+    constexpr int kRowStride = kRouter ? kRowSplit : 1;
+    int const numRowsCta = kRouter ? (R - rowGroup + kRowSplit - 1) / kRowSplit : 1;
+    // This thread's item: one 16-byte vector of one row slice.
+    int const itemRow = tid / kVecPerSlice;
+    int const itemVec = tid % kVecPerSlice;
+    bool const hasItem = itemRow < numRowsCta;
+    int const row = rowGroup + itemRow * kRowStride;
+    // The rank's own partial of a row slice is published by exactly one CTA of the grid.
+    bool const sends = hasItem && (!kRouter || itemRow % (kG > 0 ? kG : 1) == expertGroup);
+    size_t const rowOff = static_cast<size_t>(row) * kHiddenDim + sliceBase + itemVec * kEltsPerVec;
+    // The cluster that writes the normalized rows and the residual.
+    bool const writesOutputs = !kRouter || expertGroup == 0;
+
+    // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ----
+    if (tid < kVecPerSlice)
+    {
+        hop.gamma[tid] = loadConstant(reinterpret_cast<uint4 const*>(params.gamma) + slice * kVecPerSlice + tid);
+    }
+    [[maybe_unused]] int const gid = lane >> 2;
+    [[maybe_unused]] int const tig = lane & 3;
+    [[maybe_unused]] uint4 wLo[kRouter ? HopConfig<(kG > 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>::kPairsPerGroup : 1];
+    [[maybe_unused]] uint4 wHi[kRouter ? HopConfig<(kG > 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>::kPairsPerGroup : 1];
+    if constexpr (kRouter)
+    {
+        int const tile = warp % Cfg::kExpertTiles;
+        int const kGroup = warp / Cfg::kExpertTiles;
+        auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight);
+        int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid;
+#pragma unroll
+        for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+        {
+            int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec;
+            wLo[p] = loadConstant(w + static_cast<size_t>(expert) * kHiddenDim + k);
+            wHi[p] = loadConstant(w + static_cast<size_t>(expert + 8) * kHiddenDim + k);
+        }
+        // MMA columns of rows this CTA does not own stay zero.
+        for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads)
+        {
+            rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0);
+        }
+    }
+
+    // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums
+    // of its rows and, with the router, the 8 x kKGroups partial tiles of the experts it owns. A split cluster
+    // barrier publishes the initialization; its wait precedes the first remote store, long after all CTAs arrived.
+    if (tid == 0)
+    {
+        mbarrierInitExpectTx(&hop.normBar, static_cast<uint32_t>(numRowsCta * kSlices * sizeof(float)));
+        if constexpr (kRouter)
+        {
+            mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast<uint32_t>(sizeof(rs.partial)) : 0U);
+        }
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+    asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory");
+
+    // ---- Wait for the producer of `input` (and transitively of `residualIn` and the protocol state). ----
+    W2_STAMP(0);
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    W2_STAMP(1);
+    // Dependents only start their own prologue; they still wait for this grid to complete.
+    asm volatile("griddepcontrol.launch_dependents;");
+
+    uint4 const flags = loadVolatile(params.flags);
+    auto const* input = reinterpret_cast<__nv_bfloat16 const*>(params.input);
+    auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn);
+    uint4 res = make_uint4(0, 0, 0, 0);
+    // This rank's own partial, as published (sanitized); it is used from registers, not from the multicast copy.
+    uint4 own = make_uint4(0, 0, 0, 0);
+    if (hasItem)
+    {
+        res = *reinterpret_cast<uint4 const*>(residualIn + rowOff);
+        own = sanitize(*reinterpret_cast<uint4 const*>(input + rowOff));
+    }
+    uint32_t const cur = flags.x;
+    uint32_t const clearBytes = flags.y;
+    char* ucBase = reinterpret_cast<char*>(params.ucBuffer);
+    auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes);
+    auto* curMc
+        = reinterpret_cast<__nv_bfloat16*>(reinterpret_cast<char*>(params.mcBuffer) + cur * kLamportBufferBytes);
+    // Lamport layout of one buffer: [row][source rank][kHiddenDim].
+    size_t const slotBase = static_cast<size_t>(row) * kNumRanks * kHiddenDim + sliceBase + itemVec * kEltsPerVec;
+
+    // ---- Publish this rank's partial to every rank (one multicast store per vector). ----
+    if (sends)
+    {
+        *reinterpret_cast<uint4*>(curMc + slotBase + static_cast<size_t>(params.rank) * kHiddenDim) = own;
+    }
+
+    W2_STAMP(2);
+    // ---- Collect the other ranks' partials. ----
+    uint4 vals[kNumRanks];
+    if (hasItem)
+    {
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            vals[s] = s == params.rank ? own : loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+        }
+    }
+
+    // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ----
+    {
+        uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers;
+        auto* dirtyBuf = reinterpret_cast<uint4*>(ucBase + dirty * kLamportBufferBytes);
+        uint32_t const numVec = clearBytes / sizeof(uint4);
+        uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel);
+        uint32_t const ctaLinear = clusterId * kSlices + slice;
+        for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads)
+        {
+            dirtyBuf[j] = sentinel;
+        }
+    }
+
+    if (hasItem)
+    {
+        bool pending = true;
+        while (pending)
+        {
+            pending = false;
+#pragma unroll
+            for (int s = 0; s < kNumRanks; ++s)
+            {
+                if (isDirty(vals[s]))
+                {
+                    pending = true;
+                    vals[s] = loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+                }
+            }
+        }
+    }
+
+    W2_STAMP(3);
+    // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ----
+    uint4 resOut = make_uint4(0, 0, 0, 0);
+    if (hasItem)
+    {
+        float acc[kEltsPerVec];
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            acc[e] = 0.F;
+        }
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            BfVec x;
+            x.packed = vals[s];
+#pragma unroll
+            for (int e = 0; e < kEltsPerVec; ++e)
+            {
+                acc[e] += __bfloat162float(x.elts[e]);
+            }
+        }
+        BfVec rv;
+        rv.packed = res;
+        BfVec y;
+        float sq = 0.F;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]);
+            sq += __bfloat162float(__hmul(y.elts[e], y.elts[e]));
+        }
+        resOut = y.packed;
+        hop.vecSq[itemRow][itemVec] = sq;
+    }
+    __syncthreads();
+    if (tid == 0)
+    {
+        // Every thread of this CTA has read the protocol state; CTA 0 advances it only after all CTAs report.
+        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory");
+    }
+
+    // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ----
+    if (warp < numRowsCta * kNativeWarpsPerSlice)
+    {
+        int const j = warp / kNativeWarpsPerSlice;
+        int const q = warp % kNativeWarpsPerSlice;
+        float val = hop.vecSq[j][q * 32 + lane];
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+        {
+            val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32);
+        }
+        if (lane == 0)
+        {
+            hop.warpSq[j][q] = val;
+        }
+    }
+    __syncthreads();
+    asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
+    if (tid < numRowsCta * kSlices)
+    {
+        int const j = tid / kSlices;
+        int const dst = tid % kSlices;
+        float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1];
+        storeAsync(
+            clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst));
+    }
+    W2_STAMP(4);
+    mbarrierWaitPhase0(&hop.normBar);
+    W2_STAMP(5);
+
+    if (hasItem)
+    {
+        float fullSum = 0.F;
+#pragma unroll
+        for (int c = 0; c < kSlices; ++c)
+        {
+            fullSum += hop.blockSq[itemRow][c];
+        }
+        float const rcpRms = rsqrtf(fullSum / static_cast<float>(kHiddenDim) + params.eps);
+        BfVec y;
+        y.packed = resOut;
+        BfVec g;
+        g.packed = hop.gamma[itemVec];
+        BfVec n;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e]));
+        }
+        if (writesOutputs)
+        {
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed;
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed;
+        }
+        if constexpr (kRouter)
+        {
+            rs.normed[itemRow][itemVec] = n.packed;
+        }
+    }
+
+    if constexpr (kRouter)
+    {
+        if (numRowsCta > 0)
+        {
+            __syncthreads();
+            W2_STAMP(6);
+            int const tile = warp % Cfg::kExpertTiles;
+            int const kGroup = warp / Cfg::kExpertTiles;
+            // Independent accumulator chains keep the MMA latency off a single dependency chain.
+            float acc[kMmaChains][4] = {};
+#pragma unroll
+            for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+            {
+                // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands.
+                uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig];
+                int const chain = (2 * p) % kMmaChains;
+                mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y);
+                mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w);
+            }
+            float sum4[4];
+#pragma unroll
+            for (int i = 0; i < 4; ++i)
+            {
+                sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]);
+            }
+            // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner.
+            int const e0 = tile * 16 + gid;
+            int const e1 = e0 + 8;
+            uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner;
+            uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner;
+            storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0),
+                sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0));
+            storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1),
+                sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1));
+            W2_STAMP(7);
+            // Also keeps this CTA resident until every partial addressed to it has landed.
+            mbarrierWaitPhase0(&hop.routerBar);
+            W2_STAMP(8);
+
+            // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane
+            // groups of 8 sum over the 8 slices.
+            if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices)
+            {
+                int const c = tid % kSlices;
+                int const ij = tid / kSlices;
+                int const ei = ij / numRowsCta;
+                int const j = ij % numRowsCta;
+                float sum = 0.F;
+#pragma unroll
+                for (int kg = 0; kg < Cfg::kKGroups; ++kg)
+                {
+                    sum += rs.partial[c][kg][ei][j];
+                }
+#pragma unroll
+                for (int mask = kSlices / 2; mask > 0; mask >>= 1)
+                {
+                    sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices);
+                }
+                if (c == 0)
+                {
+                    int const r = rowGroup + j * kRowStride;
+                    params.routerLogits[r * kNumExperts + expertGroup * Cfg::kExpertsPerCluster
+                        + slice * Cfg::kExpertsPerOwner + ei]
+                        = sum;
+                }
+            }
+        }
+    }
+
+    W2_STAMP(9);
+    // ---- CTA 0 advances the protocol state once every CTA has read it. ----
+    if (tid == 0 && slice == 0 && clusterId == 0)
+    {
+        uint32_t arrived = 0;
+        do
+        {
+            asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory");
+        } while (arrived < static_cast<uint32_t>(kGridCtas));
+        uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers,
+            static_cast<uint32_t>(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U);
+        asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x),
+                     "r"(next.y), "r"(next.z), "r"(next.w)
+                     : "memory");
+    }
+#else
+    // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices.
+    __trap();
+#endif
+}
+
+template <int R, int kG, int kRowSplit>
+cudaError_t launchTyped(PostAttnHopParams const& params)
+{
+    using Cfg = HopConfig<kG, kRowSplit>;
+    cudaLaunchAttribute attrs[2];
+    attrs[0].id = cudaLaunchAttributeClusterDimension;
+    attrs[0].val.clusterDim.x = kSlices;
+    attrs[0].val.clusterDim.y = 1;
+    attrs[0].val.clusterDim.z = 1;
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0;
+
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSlices, kG > 0 ? kG * kRowSplit : R, 1);
+    config.blockDim = dim3(Cfg::kThreads, 1, 1);
+    config.dynamicSmemBytes = 0;
+    config.stream = params.stream;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    return cudaLaunchKernelEx(&config, postAttnHopKernel<R, kG, kRowSplit>, params);
+}
+
+template <int kG, int kRowSplit>
+cudaError_t launchRows(PostAttnHopParams const& params)
+{
+    switch (params.numRows)
+    {
+    case 1: return launchTyped<1, kG, kRowSplit>(params);
+    case 2: return launchTyped<2, kG, kRowSplit>(params);
+    case 3: return launchTyped<3, kG, kRowSplit>(params);
+    case 4: return launchTyped<4, kG, kRowSplit>(params);
+    case 5: return launchTyped<5, kG, kRowSplit>(params);
+    case 6: return launchTyped<6, kG, kRowSplit>(params);
+    case 7: return launchTyped<7, kG, kRowSplit>(params);
+    case 8: return launchTyped<8, kG, kRowSplit>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+// Router hop: 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices).
+cudaError_t launchRouter(PostAttnHopParams const& params)
+{
+    switch (params.numRows)
+    {
+    case 1: return launchTyped<1, 4, 1>(params);
+    case 2: return launchTyped<2, 4, 2>(params);
+    case 3: return launchTyped<3, 4, 3>(params);
+    case 4: return launchTyped<4, 4, 3>(params);
+    case 5: return launchTyped<5, 4, 3>(params);
+    case 6: return launchTyped<6, 4, 3>(params);
+    case 7: return launchTyped<7, 4, 3>(params);
+    case 8: return launchTyped<8, 4, 3>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace
+
+#ifdef TLLM_GLM_LL_W2_TIMING
+cudaError_t copyPostAttnHopTiming(unsigned long long* dst, cudaStream_t stream)
+{
+    return cudaMemcpyFromSymbolAsync(dst, gW2Timing, sizeof(gW2Timing), 0, cudaMemcpyDeviceToHost, stream);
+}
+#endif
+
+cudaError_t invokePostAttentionHop(PostAttnHopParams const& params)
+{
+    if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr
+        || params.mcBuffer == nullptr)
+    {
+        return cudaErrorInvalidValue;
+    }
+    if (params.routerWeight != nullptr)
+    {
+        if (params.routerLogits == nullptr)
+        {
+            return cudaErrorInvalidValue;
+        }
+        return launchRouter(params);
+    }
+    return launchRows<0, 1>(params);
+}
+
+} // namespace kernels::glm_moe_dsa::post_attention_hop
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h
new file mode 100644
index 000000000000..ee7c62202f5b
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h
@@ -0,0 +1,88 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::post_attention_hop
+{
+
+//! Fused post-attention hop.
+//!
+//! For ``numRows`` rows of hidden size ``kHiddenDim`` on a TP group of ``kNumRanks`` ranks it computes, on every
+//! rank and with identical bits on every rank:
+//!   sum      = bf16(sum over ranks 0..7 of float(partial_rank))          (one-shot Lamport exchange, fixed order)
+//!   residual = sum + residualIn                                         (bf16 add)
+//!   normed   = bf16(float(residual) * rsqrt(mean(bf16(residual^2)) + eps) * float(gamma))
+//!   logits   = normed @ routerWeight^T in FP32                          (optional, [numRows, kNumExperts])
+//! ``sum``, ``residual`` and ``normed`` follow the arithmetic and reduction tree of the native MNNVL one-shot
+//! AR + residual + RMSNorm kernel, so they are bit-identical to it.
+//!
+//! The exchange uses a private multicast workspace of ``kNumLamportBuffers`` rotating buffers. ``flags`` holds the
+//! device-side protocol state ``{current buffer, bytes to clear in the previous buffer, arrival counter, unused}``
+//! and must be zero-initialized together with the buffers filled with the Lamport sentinel (all 32-bit words equal
+//! to 0x80000000) on every rank before the first call.
+struct PostAttnHopParams
+{
+    void const* input = nullptr;        //!< [numRows, kHiddenDim] bf16, rank-local partial (o_proj output)
+    void const* residualIn = nullptr;   //!< [numRows, kHiddenDim] bf16
+    void const* gamma = nullptr;        //!< [kHiddenDim] bf16 RMSNorm weight
+    void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16 or nullptr (no router logits)
+    void* normOut = nullptr;            //!< [numRows, kHiddenDim] bf16
+    void* residualOut = nullptr;        //!< [numRows, kHiddenDim] bf16
+    float* routerLogits = nullptr;      //!< [numRows, kNumExperts] fp32, required iff routerWeight != nullptr
+    void* ucBuffer = nullptr;           //!< local unicast view of the multicast workspace
+    void* mcBuffer = nullptr;           //!< multicast address of the same workspace
+    uint32_t* flags = nullptr;          //!< [4] device-side protocol state of this workspace
+    int numRows = 0;
+    int rank = 0;
+    float eps = 1e-5F;
+    bool enablePdl = true;
+    cudaStream_t stream = nullptr;
+};
+
+constexpr int kHiddenDim = 6144;
+constexpr int kNumRanks = 8;
+constexpr int kNumExperts = 256;
+constexpr int kMaxRows = 8;
+constexpr int kNumLamportBuffers = 3;
+
+//! Bytes of one Lamport buffer: [kMaxRows][kNumRanks][kHiddenDim] bf16.
+constexpr size_t kLamportBufferBytes = size_t(kMaxRows) * kNumRanks * kHiddenDim * 2;
+//! Bytes of the whole multicast workspace.
+constexpr size_t kWorkspaceBytes = kLamportBufferBytes * kNumLamportBuffers;
+
+//! Launch the fused hop on ``params.stream``. Returns the launch status.
+cudaError_t invokePostAttentionHop(PostAttnHopParams const& params);
+
+#ifdef TLLM_GLM_LL_W2_TIMING
+//! Development builds only: phase timestamps of the most recent launch, [kTimingMaxCtas][kTimingStamps] ns.
+constexpr int kTimingMaxCtas = 128;
+constexpr int kTimingStamps = 10;
+cudaError_t copyPostAttnHopTiming(unsigned long long* dst, cudaStream_t stream);
+#endif
+
+} // namespace kernels::glm_moe_dsa::post_attention_hop
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu
new file mode 100644
index 000000000000..17c81e24877a
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu
@@ -0,0 +1,707 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h"
+
+#include <cooperative_groups.h>
+#include <cuda_bf16.h>
+#include <type_traits>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::attention_output_into_hop
+{
+namespace
+{
+namespace cg = cooperative_groups;
+
+// The post-attention hop (glmMoeDsa/postAttentionHop.cu) with its input taken from the transport: the only differences
+// are marked "Transport" below; the work decomposition, the arithmetic and the protocol are that hop's.
+//
+// Work decomposition.
+//
+// Every cluster has kSlices CTAs; CTA c owns slice [c * kSlice, (c + 1) * kSlice), the slice a CTA of the native
+// one-shot kernel owns for numRows <= 8 (96 threads x 8 elements, 8 CTAs per row), so the RMSNorm reduction below
+// reproduces the native tree exactly.
+//
+// Without router logits (dense layers) the layout is the native one: cluster y handles row y, and every CTA collects
+// one row slice (8 ranks x 768 elements).
+//
+// With router logits there are kG x kRowSplit clusters (kG = 4, kRowSplit = min(R, 3): at most 12 clusters, which stay
+// co-resident). Cluster (g, h) collects and normalizes the rows r with r % kRowSplit == h, so a CTA reads at most
+// ceil(R / kRowSplit) row slices, and computes the logits of those rows for experts [64 g, 64 g + 64). Its router
+// weights live in registers (loaded before the PDL wait). Each CTA multiplies its rows of slice c by the experts'
+// weights of slice c with BF16 MMAs (experts on m16, rows on n8) and pushes the partial logits to the CTA that owns the
+// expert; owners add the 8 slices in a fixed order, so the logits are identical on every rank. Clusters never
+// synchronize with each other; inside a cluster, slice sums and partial logits travel as st.async stores that complete
+// transactions on the consumer's mbarrier.
+constexpr int kSlices = 8;
+constexpr int kSlice = kHiddenDim / kSlices;
+constexpr int kEltsPerVec = 8;
+constexpr int kVecPerSlice = kSlice / kEltsPerVec;
+constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32;
+// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two
+// rows in one quarter-warp phase on disjoint banks.
+constexpr int kNormedStride = kVecPerSlice + 4;
+constexpr int kMmaChains = 4;
+constexpr uint32_t kLamportSentinel = 0x80000000U;
+
+static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0);
+static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32);
+static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension");
+
+// kG = number of expert groups (0: hop without router logits); kRowSplit = number of row groups.
+template <int kG, int kRowSplit>
+struct HopConfig
+{
+    static constexpr int kThreads = 384;
+    static constexpr int kWarps = kThreads / 32;
+    static constexpr int kExpertsPerCluster = kNumExperts / kG;
+    static constexpr int kExpertTiles = kExpertsPerCluster / 16;
+    static constexpr int kKGroups = kWarps / kExpertTiles;
+    static constexpr int kKPerGroup = kSlice / kKGroups;
+    static constexpr int kPairsPerGroup = kKPerGroup / 32;
+    static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices;
+    static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0);
+    static_assert(kPairsPerGroup % (kMmaChains / 2) == 0);
+};
+
+template <int kRowSplit>
+struct HopConfig<0, kRowSplit>
+{
+    static constexpr int kThreads = 128;
+    static constexpr int kWarps = kThreads / 32;
+    static constexpr int kMaxRowsPerCta = 1;
+    static_assert(kMaxRowsPerCta * kVecPerSlice <= kThreads && kMaxRowsPerCta * kNativeWarpsPerSlice <= kWarps);
+};
+
+constexpr int kMaxRowsPerCtaAll = kMaxRows / 2;
+
+struct alignas(16) HopStorage
+{
+    uint4 gamma[kVecPerSlice];
+    float vecSq[kMaxRowsPerCtaAll][kVecPerSlice];
+    float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice];
+    float blockSq[kMaxRowsPerCtaAll][kSlices];
+    // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx).
+    alignas(8) uint64_t normBar;
+    // Completes when all partial logits of the experts this CTA owns have arrived.
+    alignas(8) uint64_t routerBar;
+};
+
+template <int kG, int kRowSplit>
+struct alignas(16) RouterStorage
+{
+    using Cfg = HopConfig<kG, kRowSplit>;
+    uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero
+    float partial[kSlices][Cfg::kKGroups][Cfg::kExpertsPerOwner][kMaxRows];
+};
+
+struct EmptyStorage
+{
+};
+
+union BfVec
+{
+    uint4 packed;
+    __nv_bfloat16 elts[kEltsPerVec];
+};
+
+__device__ __forceinline__ uint4 loadVolatile(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadConstant(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr));
+    return v;
+}
+
+__device__ __forceinline__ void prefetchL2(void const* ptr)
+{
+    asm volatile("prefetch.global.L2 [%0];" ::"l"(ptr));
+}
+
+__device__ __forceinline__ bool isDirty(uint4 const& v)
+{
+    return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel;
+}
+
+// Transport: the producer writes the Lamport payloads already sanitized (bf16 -0.0 mapped to +0.0), so they never
+// contain the sentinel word.
+
+__device__ __forceinline__ uint32_t sharedAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// Address of the same shared-memory object in CTA `rank` of the cluster.
+__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank)
+{
+    uint32_t out;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank));
+    return out;
+}
+
+__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory");
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar)
+{
+    asm volatile(
+        "{\n"
+        ".reg .pred done;\n"
+        "HOP_MBAR_WAIT_%=:\n"
+        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n"
+        "@!done bra HOP_MBAR_WAIT_%=;\n"
+        "}\n" ::"r"(sharedAddr(bar))
+        : "memory");
+}
+
+// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier.
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v),
+                 "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr),
+                 "f"(a), "f"(b), "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mmaBf16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+// Transport: one warp's router-weight MMA fragments (the post-attention hop's register layout): experts
+// tile * 16 + gid (+ 8) of this cluster's expert group, k pairs of this warp's k group within the CTA's slice.
+template <class Cfg>
+__device__ __forceinline__ void loadRouterFragments(uint4 (&wLo)[Cfg::kPairsPerGroup],
+    uint4 (&wHi)[Cfg::kPairsPerGroup], void const* routerWeight, int expertGroup, int warp, int gid, int tig,
+    int sliceBase)
+{
+    int const tile = warp % Cfg::kExpertTiles;
+    int const kGroup = warp / Cfg::kExpertTiles;
+    auto const* w = reinterpret_cast<__nv_bfloat16 const*>(routerWeight);
+    int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid;
+#pragma unroll
+    for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+    {
+        int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec;
+        wLo[p] = loadConstant(w + static_cast<size_t>(expert) * kHiddenDim + k);
+        wHi[p] = loadConstant(w + static_cast<size_t>(expert + 8) * kHiddenDim + k);
+    }
+}
+
+// Transport: where the router-weight fragments are loaded (all orders give the same values and arithmetic):
+//   kRouterAfterBarrier: after the barrier that ends the rank sum, overlapping the RMS tree and cluster exchange;
+//   kRouterAfterSum:     right after the rank sum, before that barrier (the default; the loads go through
+//                        loadRouterFragments so the rank sum is not sunk below them and the kernel stays at
+//                        few enough registers to share an SM with a producer CTA);
+//   kRouterBeforeMma:    right before the MMAs (fewest registers, the load latency follows the RMSNorm).
+constexpr int kRouterAfterBarrier = 0;
+constexpr int kRouterAfterSum = 1;
+constexpr int kRouterBeforeMma = 2;
+
+template <int R, int kG, int kRowSplit, int kRouterLoad>
+__global__ void __launch_bounds__(HopConfig<kG, kRowSplit>::kThreads, 1)
+    postAttnHopX1Kernel(PostAttentionHopFromPartialsParams const params)
+{
+    static_assert(R >= 1 && R <= kMaxRows);
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    constexpr bool kRouter = kG > 0;
+    using Cfg = HopConfig<kG, kRowSplit>;
+    constexpr int kClusters = kRouter ? kG * kRowSplit : R;
+    constexpr int kGridCtas = kClusters * kSlices;
+    // Rows one CTA collects, normalizes and multiplies.
+    constexpr int kRowsPerCta = kRouter ? (R + kRowSplit - 1) / kRowSplit : 1;
+    static_assert(kRowsPerCta <= kMaxRowsPerCtaAll);
+    static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps);
+    if constexpr (kRouter)
+    {
+        static_assert(
+            HopConfig<(kG > 0 ? kG : 4), kRowSplit>::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads);
+    }
+
+    __shared__ HopStorage hop;
+    __shared__ std::conditional_t<kRouter, RouterStorage<(kG > 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>, EmptyStorage> rs;
+
+    cg::cluster_group cluster = cg::this_cluster();
+    int const slice = static_cast<int>(cluster.block_rank());
+    int const clusterId = blockIdx.y;
+    int const expertGroup = kRouter ? clusterId % (kG > 0 ? kG : 1) : 0;
+    int const rowGroup = kRouter ? clusterId / (kG > 0 ? kG : 1) : clusterId;
+    int const tid = threadIdx.x;
+    int const lane = tid & 31;
+    int const warp = tid >> 5;
+    int const sliceBase = slice * kSlice;
+
+    // Rows this CTA collects and normalizes: local row j is global row rowGroup + j * rowStride.
+    constexpr int kRowStride = kRouter ? kRowSplit : 1;
+    int const numRowsCta = kRouter ? (R - rowGroup + kRowSplit - 1) / kRowSplit : 1;
+    // This thread's item: one 16-byte vector of one row slice.
+    int const itemRow = tid / kVecPerSlice;
+    int const itemVec = tid % kVecPerSlice;
+    bool const hasItem = itemRow < numRowsCta;
+    int const row = rowGroup + itemRow * kRowStride;
+    size_t const rowOff = static_cast<size_t>(row) * kHiddenDim + sliceBase + itemVec * kEltsPerVec;
+    // The cluster that writes the normalized rows and the residual.
+    bool const writesOutputs = !kRouter || expertGroup == 0;
+
+    // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ----
+    if (tid < kVecPerSlice)
+    {
+        hop.gamma[tid] = loadConstant(reinterpret_cast<uint4 const*>(params.gamma) + slice * kVecPerSlice + tid);
+    }
+    [[maybe_unused]] int const gid = lane >> 2;
+    [[maybe_unused]] int const tig = lane & 3;
+    // Transport: the router-weight fragments are loaded after the poll (not here), so they are not live across it; with
+    // fewer registers a CTA of this kernel can be resident next to a CTA of the producer. The prologue only asks L2
+    // for the same lines (no registers held, no architectural effect).
+    if constexpr (kRouter)
+    {
+        {
+            int const tile = warp % Cfg::kExpertTiles;
+            int const kGroup = warp / Cfg::kExpertTiles;
+            auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight);
+            int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid;
+#pragma unroll
+            for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+            {
+                int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec;
+                prefetchL2(w + static_cast<size_t>(expert) * kHiddenDim + k);
+                prefetchL2(w + static_cast<size_t>(expert + 8) * kHiddenDim + k);
+            }
+        }
+        // MMA columns of rows this CTA does not own stay zero.
+        for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads)
+        {
+            rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0);
+        }
+    }
+
+    // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums
+    // of its rows and, with the router, the 8 x kKGroups partial tiles of the experts it owns. A split cluster
+    // barrier publishes the initialization; its wait precedes the first remote store, long after all CTAs arrived.
+    if (tid == 0)
+    {
+        mbarrierInitExpectTx(&hop.normBar, static_cast<uint32_t>(numRowsCta * kSlices * sizeof(float)));
+        if constexpr (kRouter)
+        {
+            mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast<uint32_t>(sizeof(rs.partial)) : 0U);
+        }
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+    asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory");
+
+    // Transport: with earlyPoll no wait here. The producer read the protocol state and wrote the transport after its
+    // own griddepcontrol.wait and triggered this grid only after that wait, so every grid before the producer has
+    // completed and the protocol state written by the previous hop is final (see postAttentionHopFromPartials.h).
+    auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn);
+    uint4 res = make_uint4(0, 0, 0, 0);
+    if (!params.earlyPoll)
+    {
+        asm volatile("griddepcontrol.wait;" ::: "memory");
+        // As in the post-attention hop: dependents may start their prologue; they still wait for this grid to complete.
+        asm volatile("griddepcontrol.launch_dependents;");
+        // Transport: after the wait the residual is final; its load overlaps the poll instead of following it.
+        if (hasItem)
+        {
+            res = *reinterpret_cast<uint4 const*>(residualIn + rowOff);
+        }
+    }
+    uint4 const flags = loadVolatile(params.flags);
+    uint32_t const cur = flags.x;
+    uint32_t const clearBytes = flags.y;
+    char* ucBase = reinterpret_cast<char*>(params.ucBuffer);
+    auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes);
+    // Lamport layout of one buffer: [row][source rank][kHiddenDim].
+    size_t const slotBase = static_cast<size_t>(row) * kNumRanks * kHiddenDim + sliceBase + itemVec * kEltsPerVec;
+
+    // ---- Transport: collect every rank's partial, this rank's own included, from the transport. ----
+    uint4 vals[kNumRanks];
+    if (hasItem)
+    {
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            vals[s] = loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+        }
+    }
+
+    // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ----
+    {
+        uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers;
+        auto* dirtyBuf = reinterpret_cast<uint4*>(ucBase + dirty * kLamportBufferBytes);
+        uint32_t const numVec = clearBytes / sizeof(uint4);
+        uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel);
+        uint32_t const ctaLinear = clusterId * kSlices + slice;
+        for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads)
+        {
+            dirtyBuf[j] = sentinel;
+        }
+    }
+
+    if (hasItem)
+    {
+        bool pending = true;
+        while (pending)
+        {
+            pending = false;
+#pragma unroll
+            for (int s = 0; s < kNumRanks; ++s)
+            {
+                if (isDirty(vals[s]))
+                {
+                    pending = true;
+                    vals[s] = loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+                }
+            }
+        }
+    }
+
+    // ---- Transport (earlyPoll): wait for the producer grid before reading the residual or writing any output. ----
+    if (params.earlyPoll)
+    {
+        asm volatile("griddepcontrol.wait;" ::: "memory");
+        // Dependents only start their own prologue; they still wait for this grid to complete.
+        asm volatile("griddepcontrol.launch_dependents;");
+        if (hasItem)
+        {
+            res = *reinterpret_cast<uint4 const*>(residualIn + rowOff);
+        }
+    }
+
+    // Transport: router-weight fragments (see kRouterLoad).
+    constexpr int kWPairs = kRouter ? HopConfig<(kG > 0 ? kG : 4), kRowSplit>::kPairsPerGroup : 1;
+    [[maybe_unused]] uint4 wLo[kWPairs];
+    [[maybe_unused]] uint4 wHi[kWPairs];
+
+    // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ----
+    uint4 resOut = make_uint4(0, 0, 0, 0);
+    if (hasItem)
+    {
+        float acc[kEltsPerVec];
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            acc[e] = 0.F;
+        }
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            BfVec x;
+            x.packed = vals[s];
+#pragma unroll
+            for (int e = 0; e < kEltsPerVec; ++e)
+            {
+                acc[e] += __bfloat162float(x.elts[e]);
+            }
+        }
+        BfVec rv;
+        rv.packed = res;
+        BfVec y;
+        float sq = 0.F;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]);
+            sq += __bfloat162float(__hmul(y.elts[e], y.elts[e]));
+        }
+        resOut = y.packed;
+        hop.vecSq[itemRow][itemVec] = sq;
+    }
+    if constexpr (kRouter && kRouterLoad == kRouterAfterSum)
+    {
+        if (numRowsCta > 0)
+        {
+            loadRouterFragments<Cfg>(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase);
+        }
+    }
+    __syncthreads();
+    if (tid == 0)
+    {
+        // Every thread of this CTA has read the protocol state; CTA 0 advances it only after all CTAs report.
+        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory");
+    }
+    if constexpr (kRouter && kRouterLoad == kRouterAfterBarrier)
+    {
+        if (numRowsCta > 0)
+        {
+            loadRouterFragments<Cfg>(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase);
+        }
+    }
+
+    // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ----
+    if (warp < numRowsCta * kNativeWarpsPerSlice)
+    {
+        int const j = warp / kNativeWarpsPerSlice;
+        int const q = warp % kNativeWarpsPerSlice;
+        float val = hop.vecSq[j][q * 32 + lane];
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+        {
+            val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32);
+        }
+        if (lane == 0)
+        {
+            hop.warpSq[j][q] = val;
+        }
+    }
+    __syncthreads();
+    asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
+    if (tid < numRowsCta * kSlices)
+    {
+        int const j = tid / kSlices;
+        int const dst = tid % kSlices;
+        float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1];
+        storeAsync(
+            clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst));
+    }
+    mbarrierWaitPhase0(&hop.normBar);
+
+    if (hasItem)
+    {
+        float fullSum = 0.F;
+#pragma unroll
+        for (int c = 0; c < kSlices; ++c)
+        {
+            fullSum += hop.blockSq[itemRow][c];
+        }
+        float const rcpRms = rsqrtf(fullSum / static_cast<float>(kHiddenDim) + params.eps);
+        BfVec y;
+        y.packed = resOut;
+        BfVec g;
+        g.packed = hop.gamma[itemVec];
+        BfVec n;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e]));
+        }
+        if (writesOutputs)
+        {
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed;
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed;
+        }
+        if constexpr (kRouter)
+        {
+            rs.normed[itemRow][itemVec] = n.packed;
+        }
+    }
+
+    if constexpr (kRouter)
+    {
+        if (numRowsCta > 0)
+        {
+            __syncthreads();
+            int const tile = warp % Cfg::kExpertTiles;
+            int const kGroup = warp / Cfg::kExpertTiles;
+            if constexpr (kRouterLoad == kRouterBeforeMma)
+            {
+                loadRouterFragments<Cfg>(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase);
+            }
+            // Independent accumulator chains keep the MMA latency off a single dependency chain.
+            float acc[kMmaChains][4] = {};
+#pragma unroll
+            for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+            {
+                // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands.
+                uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig];
+                int const chain = (2 * p) % kMmaChains;
+                mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y);
+                mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w);
+            }
+            float sum4[4];
+#pragma unroll
+            for (int i = 0; i < 4; ++i)
+            {
+                sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]);
+            }
+            // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner.
+            int const e0 = tile * 16 + gid;
+            int const e1 = e0 + 8;
+            uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner;
+            uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner;
+            storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0),
+                sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0));
+            storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1),
+                sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1));
+            // Also keeps this CTA resident until every partial addressed to it has landed.
+            mbarrierWaitPhase0(&hop.routerBar);
+
+            // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane
+            // groups of 8 sum over the 8 slices.
+            if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices)
+            {
+                int const c = tid % kSlices;
+                int const ij = tid / kSlices;
+                int const ei = ij / numRowsCta;
+                int const j = ij % numRowsCta;
+                float sum = 0.F;
+#pragma unroll
+                for (int kg = 0; kg < Cfg::kKGroups; ++kg)
+                {
+                    sum += rs.partial[c][kg][ei][j];
+                }
+#pragma unroll
+                for (int mask = kSlices / 2; mask > 0; mask >>= 1)
+                {
+                    sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices);
+                }
+                if (c == 0)
+                {
+                    int const r = rowGroup + j * kRowStride;
+                    params.routerLogits[r * kNumExperts + expertGroup * Cfg::kExpertsPerCluster
+                        + slice * Cfg::kExpertsPerOwner + ei]
+                        = sum;
+                }
+            }
+        }
+    }
+
+    // ---- CTA 0 advances the protocol state once every CTA has read it. ----
+    if (tid == 0 && slice == 0 && clusterId == 0)
+    {
+        uint32_t arrived = 0;
+        do
+        {
+            asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory");
+        } while (arrived < static_cast<uint32_t>(kGridCtas));
+        uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers,
+            static_cast<uint32_t>(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U);
+        asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x),
+                     "r"(next.y), "r"(next.z), "r"(next.w)
+                     : "memory");
+    }
+#else
+    // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices.
+    __trap();
+#endif
+}
+
+template <int R, int kG, int kRowSplit, int kRouterLoad = kRouterAfterSum>
+cudaError_t launchTyped(PostAttentionHopFromPartialsParams const& params)
+{
+    using Cfg = HopConfig<kG, kRowSplit>;
+    cudaLaunchAttribute attrs[2];
+    attrs[0].id = cudaLaunchAttributeClusterDimension;
+    attrs[0].val.clusterDim.x = kSlices;
+    attrs[0].val.clusterDim.y = 1;
+    attrs[0].val.clusterDim.z = 1;
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0;
+
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSlices, kG > 0 ? kG * kRowSplit : R, 1);
+    config.blockDim = dim3(Cfg::kThreads, 1, 1);
+    config.dynamicSmemBytes = 0;
+    config.stream = params.stream;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    return cudaLaunchKernelEx(&config, postAttnHopX1Kernel<R, kG, kRowSplit, kRouterLoad>, params);
+}
+
+template <int kG, int kRowSplit>
+cudaError_t launchRows(PostAttentionHopFromPartialsParams const& params)
+{
+    switch (params.numRows)
+    {
+    case 1: return launchTyped<1, kG, kRowSplit>(params);
+    case 2: return launchTyped<2, kG, kRowSplit>(params);
+    case 3: return launchTyped<3, kG, kRowSplit>(params);
+    case 4: return launchTyped<4, kG, kRowSplit>(params);
+    case 5: return launchTyped<5, kG, kRowSplit>(params);
+    case 6: return launchTyped<6, kG, kRowSplit>(params);
+    case 7: return launchTyped<7, kG, kRowSplit>(params);
+    case 8: return launchTyped<8, kG, kRowSplit>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+// Router hop: 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices).
+template <int kRouterLoad>
+cudaError_t launchRouterAt(PostAttentionHopFromPartialsParams const& params)
+{
+    switch (params.numRows)
+    {
+    case 1: return launchTyped<1, 4, 1, kRouterLoad>(params);
+    case 2: return launchTyped<2, 4, 2, kRouterLoad>(params);
+    case 3: return launchTyped<3, 4, 3, kRouterLoad>(params);
+    case 4: return launchTyped<4, 4, 3, kRouterLoad>(params);
+    case 5: return launchTyped<5, 4, 3, kRouterLoad>(params);
+    case 6: return launchTyped<6, 4, 3, kRouterLoad>(params);
+    case 7: return launchTyped<7, 4, 3, kRouterLoad>(params);
+    case 8: return launchTyped<8, 4, 3, kRouterLoad>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+cudaError_t launchRouter(PostAttentionHopFromPartialsParams const& params)
+{
+    switch (params.routerLoad)
+    {
+    case kRouterAfterBarrier: return launchRouterAt<kRouterAfterBarrier>(params);
+    case kRouterAfterSum: return launchRouterAt<kRouterAfterSum>(params);
+    case kRouterBeforeMma: return launchRouterAt<kRouterBeforeMma>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace
+
+cudaError_t invokePostAttentionHopFromPartials(PostAttentionHopFromPartialsParams const& params)
+{
+    if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr)
+    {
+        return cudaErrorInvalidValue;
+    }
+    if (params.routerWeight != nullptr)
+    {
+        if (params.routerLogits == nullptr)
+        {
+            return cudaErrorInvalidValue;
+        }
+        return launchRouter(params);
+    }
+    return launchRows<0, 1>(params);
+}
+
+} // namespace kernels::glm_moe_dsa::attention_output_into_hop
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h
new file mode 100644
index 000000000000..362fa1da20fa
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h
@@ -0,0 +1,92 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::attention_output_into_hop
+{
+
+//! Post-attention hop that takes its input from the transport: the hop of postAttentionHop.h with its input taken
+//! from the transport instead of from a partial tensor.
+//!
+//! The producer of this call (the attention output tail's o_proj kernel, publishing into the hop) has already
+//! written every rank's sanitized rank-local partial (bf16 -0.0 mapped to +0.0) into the current Lamport buffer of
+//! the post-attention hop's workspace, slot ``[row][rank][kHiddenDim]``, with multicast stores. This kernel polls
+//! all ``kNumRanks`` slots of its rows (its own included) and then performs exactly the post-attention hop's
+//! arithmetic: the rank sum in rank order 0..7, the bf16 residual add, RMSNorm with the native reduction tree and
+//! the optional router logits with the post-attention hop's MMA decomposition and summation order. Given the same
+//! partials, ``normed``, ``residualOut`` and ``routerLogits`` are bit-identical to ``postAttnHop`` consuming the
+//! ordinary partial (whose own contribution is the same sanitized value, taken from registers there).
+//!
+//! The workspace, buffer rotation and protocol state are the post-attention hop's
+//! (``kernels::glm_moe_dsa::post_attention_hop``), so calls of either kind can follow each other on one workspace.
+//! Memory-order contract with the producer (programmatic dependent launch):
+//!   * the producer reads the protocol state and writes the transport only after its own ``griddepcontrol.wait`` and
+//!     triggers its dependents only after that wait, so when this kernel starts, every grid before the producer has
+//!     completed: the protocol state written by the previous hop and ``residualIn`` are final;
+//!   * with ``earlyPoll`` this kernel reads the protocol state and polls the transport before its
+//!     ``griddepcontrol.wait`` (the transport is private to the package and Lamport-validated); it reads
+//!     ``residualIn`` and writes every output only after the wait, i.e. after the producer grid has completed.
+//!     Without ``earlyPoll`` it waits first. The early read relies on every kernel between two hops waiting on its
+//!     predecessor before it completes, which holds for the served graph's kernels (each consumes its predecessor's
+//!     output).
+struct PostAttentionHopFromPartialsParams
+{
+    void const* residualIn = nullptr;   //!< [numRows, kHiddenDim] bf16
+    void const* gamma = nullptr;        //!< [kHiddenDim] bf16 RMSNorm weight
+    void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16 or nullptr (no router logits)
+    void* normOut = nullptr;            //!< [numRows, kHiddenDim] bf16
+    void* residualOut = nullptr;        //!< [numRows, kHiddenDim] bf16
+    float* routerLogits = nullptr;      //!< [numRows, kNumExperts] fp32, required iff routerWeight != nullptr
+    void* ucBuffer = nullptr;           //!< local unicast view of the post-attention hop's multicast workspace
+    uint32_t* flags = nullptr;          //!< [4] protocol state of that workspace
+    int numRows = 0;
+    int rank = 0;
+    float eps = 1e-5F;
+    bool enablePdl = true;
+    //! true: read the protocol state and poll the transport before ``griddepcontrol.wait`` (relies on the producer
+    //! contract above); false: wait first, as postAttentionHop does, then read the state and poll.
+    bool earlyPoll = true;
+    //! Where the router-weight MMA fragments are loaded (same values and arithmetic in every case): 0 after the
+    //! barrier that ends the rank sum, 1 right after the rank sum (default), 2 right before the MMAs.
+    int routerLoad = 1;
+    cudaStream_t stream = nullptr;
+};
+
+using glm_moe_dsa::post_attention_hop::kHiddenDim;
+using glm_moe_dsa::post_attention_hop::kLamportBufferBytes;
+using glm_moe_dsa::post_attention_hop::kMaxRows;
+using glm_moe_dsa::post_attention_hop::kNumExperts;
+using glm_moe_dsa::post_attention_hop::kNumLamportBuffers;
+using glm_moe_dsa::post_attention_hop::kNumRanks;
+using glm_moe_dsa::post_attention_hop::kWorkspaceBytes;
+
+//! Launch the transport-input hop on ``params.stream``. Returns the launch status.
+cudaError_t invokePostAttentionHopFromPartials(PostAttentionHopFromPartialsParams const& params);
+
+} // namespace kernels::glm_moe_dsa::attention_output_into_hop
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu
new file mode 100644
index 000000000000..fe8cc14ea457
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu
@@ -0,0 +1,889 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h"
+
+#include <cooperative_groups.h>
+#include <cuda_bf16.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::hop_routing
+{
+namespace
+{
+namespace cg = cooperative_groups;
+
+// The partials hop (glmMoeDsa/postAttentionHopFromPartials.cu, router variant) and top-8 routing. Everything up to and
+// including the router-logit stores is the partials hop's code; the additions are marked "Routing".
+//
+// Work decomposition (the partials hop's). Every cluster has kSlices CTAs; CTA c owns hidden slice
+// [c * kSlice, (c + 1) * kSlice). There are kG x kRowSplit clusters (kG = 4, kRowSplit = min(R, 3)). Cluster (g, h)
+// collects and normalizes the rows r with r % kRowSplit == h and computes their logits for experts [64 g, 64 g + 64);
+// CTA c of the cluster owns the 8 experts 64 g + 8 c + [0, 8) and writes their logits.
+//
+// Routing. The owner of logit (r, e) also writes the score bits of (r, e) into the gather buffer. Row
+// r = h + j * kRowSplit (local row j of row group h) is finished by CTA 0 of cluster (g = j, h): its 256 threads poll
+// the row's 256 scores, each warp ranks its 32 keys and keeps its top 8 (level 1), the 256 threads rank the 64
+// candidates (level 2, four lanes per candidate), and lanes 0..7 write the row's ids and weights.
+constexpr int kSlices = 8;
+constexpr int kSlice = kHiddenDim / kSlices;
+constexpr int kEltsPerVec = 8;
+constexpr int kVecPerSlice = kSlice / kEltsPerVec;
+constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32;
+// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two
+// rows in one quarter-warp phase on disjoint banks.
+constexpr int kNormedStride = kVecPerSlice + 4;
+constexpr int kMmaChains = 4;
+constexpr uint32_t kLamportSentinel = 0x80000000U;
+// The empty gather slot: a routing score is never -0.0f (see postAttentionHopRouting.h).
+constexpr uint32_t kGatherSentinel = 0x80000000U;
+constexpr int kG = 4;
+constexpr int kExpertWarps = kNumExperts / 32;
+constexpr int kCandidates = kExpertWarps * kTopK;
+
+static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0);
+static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32);
+static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension");
+static_assert(kNumExperts == 256 && kTopK == 8, "the finisher maps one expert per thread and 8 candidates per warp");
+static_assert(kCandidates * 4 == kNumExperts, "level 2 ranks each candidate on four lanes");
+
+template <int kRowSplit>
+struct HopConfig
+{
+    static constexpr int kThreads = 384;
+    static constexpr int kWarps = kThreads / 32;
+    static constexpr int kExpertsPerCluster = kNumExperts / kG;
+    static constexpr int kExpertTiles = kExpertsPerCluster / 16;
+    static constexpr int kKGroups = kWarps / kExpertTiles;
+    static constexpr int kKPerGroup = kSlice / kKGroups;
+    static constexpr int kPairsPerGroup = kKPerGroup / 32;
+    static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices;
+    static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0);
+    static_assert(kPairsPerGroup % (kMmaChains / 2) == 0);
+    static_assert(kThreads >= kNumExperts);
+};
+
+constexpr int kMaxRowsPerCtaAll = kMaxRows / 2;
+
+struct alignas(16) HopStorage
+{
+    uint4 gamma[kVecPerSlice];
+    float vecSq[kMaxRowsPerCtaAll][kVecPerSlice];
+    float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice];
+    float blockSq[kMaxRowsPerCtaAll][kSlices];
+    // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx).
+    alignas(8) uint64_t normBar;
+    // Completes when all partial logits of the experts this CTA owns have arrived.
+    alignas(8) uint64_t routerBar;
+};
+
+template <int kRowSplit>
+struct alignas(16) RouterStorage
+{
+    using Cfg = HopConfig<kRowSplit>;
+    uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero
+    float partial[kSlices][Cfg::kKGroups][Cfg::kExpertsPerOwner][kMaxRows];
+};
+
+// Routing: state of one CTA.
+struct alignas(16) RouteStorage
+{
+    unsigned long long candKey[kCandidates]; // level-1 survivors: native 64-bit keys
+    float candSg[kCandidates];               // their score - bias
+    unsigned long long selKey[kTopK];        // level-2 result in rank order
+    float selSg[kTopK];
+    float biasOwn[8];                        // bias of the experts this CTA owns
+    uint32_t parity;                         // gather parity of this call
+    float biasAll[kNumExperts];              // finisher CTAs: bias of every expert
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+    unsigned long long stamp[kNumStamps];
+#endif
+};
+
+union BfVec
+{
+    uint4 packed;
+    __nv_bfloat16 elts[kEltsPerVec];
+};
+
+__device__ __forceinline__ uint4 loadVolatile(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+                 : "l"(ptr)
+                 : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint32_t loadVolatileU32(void const* ptr)
+{
+    uint32_t v;
+    asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ uint4 loadConstant(void const* ptr)
+{
+    uint4 v;
+    asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr));
+    return v;
+}
+
+__device__ __forceinline__ float loadBiasConstant(void const* bias, int expert)
+{
+    unsigned short raw;
+    asm volatile("ld.global.nc.u16 %0, [%1];" : "=h"(raw) : "l"(reinterpret_cast<__nv_bfloat16 const*>(bias) + expert));
+    return __bfloat162float(__ushort_as_bfloat16(raw));
+}
+
+__device__ __forceinline__ void prefetchL2(void const* ptr)
+{
+    asm volatile("prefetch.global.L2 [%0];" ::"l"(ptr));
+}
+
+__device__ __forceinline__ void storeRelaxed(uint32_t* ptr, uint32_t v)
+{
+    asm volatile("st.relaxed.gpu.global.u32 [%0], %1;" ::"l"(ptr), "r"(v) : "memory");
+}
+
+__device__ __forceinline__ uint32_t loadRelaxed(uint32_t const* ptr)
+{
+    uint32_t v;
+    asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory");
+    return v;
+}
+
+__device__ __forceinline__ bool isDirty(uint4 const& v)
+{
+    return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel;
+}
+
+__device__ __forceinline__ uint32_t sharedAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// Address of the same shared-memory object in CTA `rank` of the cluster.
+__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank)
+{
+    uint32_t out;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank));
+    return out;
+}
+
+__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory");
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar)
+{
+    asm volatile(
+        "{\n"
+        ".reg .pred done;\n"
+        "ROUTING_MBAR_WAIT_%=:\n"
+        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n"
+        "@!done bra ROUTING_MBAR_WAIT_%=;\n"
+        "}\n" ::"r"(sharedAddr(bar))
+        : "memory");
+}
+
+// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier.
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v),
+                 "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr),
+                 "f"(a), "f"(b), "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void mmaBf16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+// One warp's router-weight MMA fragments (the post-attention hop's register layout): experts tile * 16 + gid (+ 8)
+// of this cluster's expert group, k pairs of this warp's k group within the CTA's slice.
+template <class Cfg>
+__device__ __forceinline__ void loadRouterFragments(uint4 (&wLo)[Cfg::kPairsPerGroup],
+    uint4 (&wHi)[Cfg::kPairsPerGroup], void const* routerWeight, int expertGroup, int warp, int gid, int tig,
+    int sliceBase)
+{
+    int const tile = warp % Cfg::kExpertTiles;
+    int const kGroup = warp / Cfg::kExpertTiles;
+    auto const* w = reinterpret_cast<__nv_bfloat16 const*>(routerWeight);
+    int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid;
+#pragma unroll
+    for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+    {
+        int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec;
+        wLo[p] = loadConstant(w + static_cast<size_t>(expert) * kHiddenDim + k);
+        wHi[p] = loadConstant(w + static_cast<size_t>(expert + 8) * kHiddenDim + k);
+    }
+}
+
+// Router-weight fragment load points, as in the partials hop (all orders give the same values and arithmetic).
+constexpr int kRouterAfterBarrier = 0;
+constexpr int kRouterAfterSum = 1;
+constexpr int kRouterBeforeMma = 2;
+
+// Routing: the native routing score (RoutingKernel.cuh sigmoid_accurate plus the SigmoidBias preprocess).
+__device__ __forceinline__ float routingScore(float logit, float bias)
+{
+    float const s = 0.5f * tanhf(0.5f * logit) + 0.5f;
+    return s + bias;
+}
+
+// Routing: cub::Traits<float>::TwiddleIn, the order-preserving map of FP32 bits onto unsigned integers.
+__device__ __forceinline__ uint32_t twiddleIn(uint32_t bits)
+{
+    return bits ^ ((bits & 0x80000000U) != 0U ? 0xFFFFFFFFU : 0x80000000U);
+}
+
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+__device__ __forceinline__ unsigned long long globalTimer()
+{
+    unsigned long long t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return t;
+}
+
+#define ROUTING_STAMP(field) atomicMax(&route.stamp[field], globalTimer())
+#else
+#define ROUTING_STAMP(field)                                                                                           \
+    do                                                                                                                 \
+    {                                                                                                                  \
+    } while (0)
+#endif
+
+// Stamp fields (diagnostic builds).
+constexpr int kStampEntry = 0;
+constexpr int kStampWait = 1;
+constexpr int kStampLogits = 2;
+constexpr int kStampGather = 3;
+constexpr int kStampPolled = 4;
+constexpr int kStampLevel1 = 5;
+constexpr int kStampRouted = 6;
+constexpr int kStampEnd = 7;
+static_assert(kStampEnd < kNumStamps);
+
+// The register cap of the partials hop, under which one CTA of it fits on an SM next to a CTA of the publishing
+// o_proj. The routing code after the MMA phase would raise the allocation above it; the cap keeps the partials hop's
+// allocation without spills.
+constexpr int kMaxRegisters = 112;
+
+template <int R, int kRowSplit, int kRouterLoad, bool kLogitsIn>
+__global__ void __maxnreg__(kMaxRegisters) postAttnHopZ2Kernel(PostAttentionHopRoutingParams const params)
+{
+    static_assert(R >= 1 && R <= kMaxRows);
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    using Cfg = HopConfig<kRowSplit>;
+    constexpr int kClusters = kG * kRowSplit;
+    constexpr int kGridCtas = kClusters * kSlices;
+    static_assert(kGridCtas <= kMaxGridCtas);
+    // Rows one CTA collects, normalizes and multiplies.
+    constexpr int kRowsPerCta = (R + kRowSplit - 1) / kRowSplit;
+    static_assert(kRowsPerCta <= kMaxRowsPerCtaAll && kRowsPerCta <= kG);
+    static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps);
+    static_assert(Cfg::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads);
+
+    __shared__ HopStorage hop;
+    __shared__ RouterStorage<kRowSplit> rs;
+    __shared__ RouteStorage route;
+
+    cg::cluster_group cluster = cg::this_cluster();
+    int const slice = static_cast<int>(cluster.block_rank());
+    int const clusterId = blockIdx.y;
+    int const expertGroup = clusterId % kG;
+    int const rowGroup = clusterId / kG;
+    int const tid = threadIdx.x;
+    int const lane = tid & 31;
+    int const warp = tid >> 5;
+    int const sliceBase = slice * kSlice;
+    uint32_t const ctaLinear = clusterId * kSlices + slice;
+
+    // Rows this CTA collects and normalizes: local row j is global row rowGroup + j * kRowSplit.
+    constexpr int kRowStride = kRowSplit;
+    int const numRowsCta = (R - rowGroup + kRowSplit - 1) / kRowSplit;
+    // This thread's item: one 16-byte vector of one row slice.
+    int const itemRow = tid / kVecPerSlice;
+    int const itemVec = tid % kVecPerSlice;
+    bool const hasItem = itemRow < numRowsCta;
+    int const row = rowGroup + itemRow * kRowStride;
+    size_t const rowOff = static_cast<size_t>(row) * kHiddenDim + sliceBase + itemVec * kEltsPerVec;
+    // The cluster that writes the normalized rows and the residual.
+    bool const writesOutputs = expertGroup == 0;
+    // Routing: this CTA finishes local row `expertGroup` of its row group, if that row exists.
+    bool const finisher = slice == 0 && expertGroup < numRowsCta;
+
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+    if (tid < kNumStamps)
+    {
+        route.stamp[tid] = 0ULL;
+    }
+    __syncthreads();
+    if (tid == 0)
+    {
+        ROUTING_STAMP(kStampEntry);
+    }
+#endif
+
+    // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ----
+    if (tid < kVecPerSlice)
+    {
+        hop.gamma[tid] = loadConstant(reinterpret_cast<uint4 const*>(params.gamma) + slice * kVecPerSlice + tid);
+    }
+    // Routing: router bias of the owned experts and, on finishers, of every expert (a model parameter).
+    if (tid < Cfg::kExpertsPerOwner)
+    {
+        route.biasOwn[tid] = loadBiasConstant(
+            params.routingBias, expertGroup * Cfg::kExpertsPerCluster + slice * Cfg::kExpertsPerOwner + tid);
+    }
+    if (finisher && tid < kNumExperts)
+    {
+        route.biasAll[tid] = loadBiasConstant(params.routingBias, tid);
+    }
+    int const gid = lane >> 2;
+    int const tig = lane & 3;
+    // The router-weight fragments are loaded after the poll (not here), so they are not live across it; the prologue
+    // only asks L2 for the same lines (no registers held, no architectural effect).
+    {
+        int const tile = warp % Cfg::kExpertTiles;
+        int const kGroup = warp / Cfg::kExpertTiles;
+        auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight);
+        int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid;
+#pragma unroll
+        for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+        {
+            int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec;
+            prefetchL2(w + static_cast<size_t>(expert) * kHiddenDim + k);
+            prefetchL2(w + static_cast<size_t>(expert + 8) * kHiddenDim + k);
+        }
+    }
+    // MMA columns of rows this CTA does not own stay zero.
+    for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads)
+    {
+        rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0);
+    }
+
+    // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums
+    // of its rows and the 8 x kKGroups partial tiles of the experts it owns. A split cluster barrier publishes the
+    // initialization; its wait precedes the first remote store, long after all CTAs arrived.
+    if (tid == 0)
+    {
+        mbarrierInitExpectTx(&hop.normBar, static_cast<uint32_t>(numRowsCta * kSlices * sizeof(float)));
+        mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast<uint32_t>(sizeof(rs.partial)) : 0U);
+        asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+    asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory");
+
+    // With earlyPoll no wait here (the protocol state and the transport may be read early, see
+    // postAttentionHopFromPartials.h).
+    auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn);
+    uint4 res = make_uint4(0, 0, 0, 0);
+    if (!params.earlyPoll)
+    {
+        asm volatile("griddepcontrol.wait;" ::: "memory");
+        // Dependents may start their prologue; they still wait for this grid to complete.
+        asm volatile("griddepcontrol.launch_dependents;");
+        if (hasItem)
+        {
+            res = *reinterpret_cast<uint4 const*>(residualIn + rowOff);
+        }
+    }
+    uint4 const flags = loadVolatile(params.flags);
+    uint32_t const cur = flags.x;
+    uint32_t const clearBytes = flags.y;
+    char* ucBase = reinterpret_cast<char*>(params.ucBuffer);
+    auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes);
+    // Lamport layout of one buffer: [row][source rank][kHiddenDim].
+    size_t const slotBase = static_cast<size_t>(row) * kNumRanks * kHiddenDim + sliceBase + itemVec * kEltsPerVec;
+
+    // ---- Collect every rank's partial, this rank's own included, from the transport. ----
+    uint4 vals[kNumRanks];
+    if (hasItem)
+    {
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            vals[s] = loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+        }
+    }
+
+    // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ----
+    {
+        uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers;
+        auto* dirtyBuf = reinterpret_cast<uint4*>(ucBase + dirty * kLamportBufferBytes);
+        uint32_t const numVec = clearBytes / sizeof(uint4);
+        uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel);
+        for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads)
+        {
+            dirtyBuf[j] = sentinel;
+        }
+    }
+
+    if (hasItem)
+    {
+        bool pending = true;
+        while (pending)
+        {
+            pending = false;
+#pragma unroll
+            for (int s = 0; s < kNumRanks; ++s)
+            {
+                if (isDirty(vals[s]))
+                {
+                    pending = true;
+                    vals[s] = loadVolatile(curUc + slotBase + static_cast<size_t>(s) * kHiddenDim);
+                }
+            }
+        }
+    }
+
+    // ---- earlyPoll: wait for the producer grid before reading the residual or writing any output. ----
+    if (params.earlyPoll)
+    {
+        asm volatile("griddepcontrol.wait;" ::: "memory");
+        // Dependents only start their own prologue; they still wait for this grid to complete.
+        asm volatile("griddepcontrol.launch_dependents;");
+        if (hasItem)
+        {
+            res = *reinterpret_cast<uint4 const*>(residualIn + rowOff);
+        }
+    }
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+    if (tid == 0)
+    {
+        ROUTING_STAMP(kStampWait);
+    }
+#endif
+
+    // Router-weight fragments (see kRouterLoad).
+    uint4 wLo[Cfg::kPairsPerGroup];
+    uint4 wHi[Cfg::kPairsPerGroup];
+
+    // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ----
+    uint4 resOut = make_uint4(0, 0, 0, 0);
+    if (hasItem)
+    {
+        float acc[kEltsPerVec];
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            acc[e] = 0.F;
+        }
+#pragma unroll
+        for (int s = 0; s < kNumRanks; ++s)
+        {
+            BfVec x;
+            x.packed = vals[s];
+#pragma unroll
+            for (int e = 0; e < kEltsPerVec; ++e)
+            {
+                acc[e] += __bfloat162float(x.elts[e]);
+            }
+        }
+        BfVec rv;
+        rv.packed = res;
+        BfVec y;
+        float sq = 0.F;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]);
+            sq += __bfloat162float(__hmul(y.elts[e], y.elts[e]));
+        }
+        resOut = y.packed;
+        hop.vecSq[itemRow][itemVec] = sq;
+    }
+    if constexpr (kRouterLoad == kRouterAfterSum)
+    {
+        if (numRowsCta > 0)
+        {
+            loadRouterFragments<Cfg>(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase);
+        }
+    }
+    if (tid == 0)
+    {
+        // Routing: the gather parity of this call (the routing state is read only after the dependency wait). It is
+        // kept in shared memory, published by the barrier below, so no register holds it across the MMA phase.
+        route.parity = loadVolatileU32(params.zstate) & 1U;
+    }
+    __syncthreads();
+    if (tid == 0)
+    {
+        // Every thread of this CTA has read the protocol state and thread 0 the gather parity; CTA 0 advances both only
+        // after all CTAs report.
+        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory");
+    }
+    // ---- Routing: reset the previous call's gather buffer (its readers belong to a completed grid), off the routing
+    // path. ----
+    {
+        uint4* const other = reinterpret_cast<uint4*>(params.gather + (route.parity ^ 1U) * kGatherBufferWords);
+        uint4 const sentinel = make_uint4(kGatherSentinel, kGatherSentinel, kGatherSentinel, kGatherSentinel);
+        for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < kGatherBufferWords / 4; j += kGridCtas * Cfg::kThreads)
+        {
+            other[j] = sentinel;
+        }
+    }
+    if constexpr (kRouterLoad == kRouterAfterBarrier)
+    {
+        if (numRowsCta > 0)
+        {
+            loadRouterFragments<Cfg>(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase);
+        }
+    }
+
+    // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ----
+    if (warp < numRowsCta * kNativeWarpsPerSlice)
+    {
+        int const j = warp / kNativeWarpsPerSlice;
+        int const q = warp % kNativeWarpsPerSlice;
+        float val = hop.vecSq[j][q * 32 + lane];
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+        {
+            val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32);
+        }
+        if (lane == 0)
+        {
+            hop.warpSq[j][q] = val;
+        }
+    }
+    __syncthreads();
+    asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
+    if (tid < numRowsCta * kSlices)
+    {
+        int const j = tid / kSlices;
+        int const dst = tid % kSlices;
+        float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1];
+        storeAsync(
+            clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst));
+    }
+    mbarrierWaitPhase0(&hop.normBar);
+
+    if (hasItem)
+    {
+        float fullSum = 0.F;
+#pragma unroll
+        for (int c = 0; c < kSlices; ++c)
+        {
+            fullSum += hop.blockSq[itemRow][c];
+        }
+        float const rcpRms = rsqrtf(fullSum / static_cast<float>(kHiddenDim) + params.eps);
+        BfVec y;
+        y.packed = resOut;
+        BfVec g;
+        g.packed = hop.gamma[itemVec];
+        BfVec n;
+#pragma unroll
+        for (int e = 0; e < kEltsPerVec; ++e)
+        {
+            n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e]));
+        }
+        if (writesOutputs)
+        {
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed;
+            *reinterpret_cast<uint4*>(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed;
+        }
+        rs.normed[itemRow][itemVec] = n.packed;
+    }
+
+    if (numRowsCta > 0)
+    {
+        __syncthreads();
+        int const tile = warp % Cfg::kExpertTiles;
+        int const kGroup = warp / Cfg::kExpertTiles;
+        if constexpr (kRouterLoad == kRouterBeforeMma)
+        {
+            loadRouterFragments<Cfg>(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase);
+        }
+        // Independent accumulator chains keep the MMA latency off a single dependency chain.
+        float acc[kMmaChains][4] = {};
+#pragma unroll
+        for (int p = 0; p < Cfg::kPairsPerGroup; ++p)
+        {
+            // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands.
+            uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig];
+            int const chain = (2 * p) % kMmaChains;
+            mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y);
+            mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w);
+        }
+        float sum4[4];
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+        {
+            sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]);
+        }
+        // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner.
+        int const e0 = tile * 16 + gid;
+        int const e1 = e0 + 8;
+        uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner;
+        uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner;
+        storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0),
+            sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0));
+        storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1),
+            sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1));
+        // Also keeps this CTA resident until every partial addressed to it has landed.
+        mbarrierWaitPhase0(&hop.routerBar);
+
+        // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane
+        // groups of 8 sum over the 8 slices.
+        if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices)
+        {
+            int const c = tid % kSlices;
+            int const ij = tid / kSlices;
+            int const ei = ij / numRowsCta;
+            int const j = ij % numRowsCta;
+            float sum = 0.F;
+#pragma unroll
+            for (int kg = 0; kg < Cfg::kKGroups; ++kg)
+            {
+                sum += rs.partial[c][kg][ei][j];
+            }
+#pragma unroll
+            for (int mask = kSlices / 2; mask > 0; mask >>= 1)
+            {
+                sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices);
+            }
+            if (c == 0)
+            {
+                int const r = rowGroup + j * kRowStride;
+                int const e = expertGroup * Cfg::kExpertsPerCluster + slice * Cfg::kExpertsPerOwner + ei;
+                params.routerLogits[r * kNumExperts + e] = sum;
+                ROUTING_STAMP(kStampLogits);
+                // Routing: the score of (r, e) into this call's gather buffer.
+                float const logit = kLogitsIn ? params.logitsIn[r * kNumExperts + e] : sum;
+                storeRelaxed(params.gather + route.parity * kGatherBufferWords + r * kNumExperts + e,
+                    __float_as_uint(routingScore(logit, route.biasOwn[ei])));
+                ROUTING_STAMP(kStampGather);
+            }
+        }
+    }
+
+    // ---- Routing: finisher CTAs select the top 8 of their row. ----
+    if (finisher)
+    {
+        int const rf = rowGroup + expertGroup * kRowStride;
+        if (tid < kNumExperts)
+        {
+            uint32_t const* slot = params.gather + route.parity * kGatherBufferWords + rf * kNumExperts + tid;
+            uint32_t bits = loadRelaxed(slot);
+            while (bits == kGatherSentinel)
+            {
+                bits = loadRelaxed(slot);
+            }
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+            __syncwarp();
+            if (lane == 0)
+            {
+                ROUTING_STAMP(kStampPolled);
+            }
+#endif
+            float const sg = __uint_as_float(bits) - route.biasAll[tid];
+            uint32_t const hi = twiddleIn(bits);
+            // Level 1: rank of this key among the warp's 32 (expert 32 * warp + i sits in lane i, so equal scores
+            // rank the lower lane first, as the native key does).
+            int rank = 0;
+#pragma unroll
+            for (int i = 0; i < 32; ++i)
+            {
+                uint32_t const other = __shfl_sync(0xffffffffU, hi, i);
+                rank += (other > hi || (other == hi && i < lane)) ? 1 : 0;
+            }
+            if (rank < kTopK)
+            {
+                route.candKey[warp * kTopK + rank]
+                    = (static_cast<unsigned long long>(hi) << 32) | static_cast<unsigned long long>(65535 - tid);
+                route.candSg[warp * kTopK + rank] = sg;
+            }
+        }
+        __syncthreads();
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+        if (tid == 0)
+        {
+            ROUTING_STAMP(kStampLevel1);
+        }
+#endif
+        if (tid < kNumExperts)
+        {
+            // Level 2: rank of the 64 candidates by the full native key (all keys are distinct). Four lanes share
+            // candidate tid / 4; lane q of them counts the larger keys among candidates q, q + 4, ..., q + 60.
+            int const ci = tid >> 2;
+            int const q = tid & 3;
+            unsigned long long const key = route.candKey[ci];
+            int rank = 0;
+#pragma unroll
+            for (int i = 0; i < kCandidates / 4; ++i)
+            {
+                rank += route.candKey[i * 4 + q] > key ? 1 : 0;
+            }
+            rank += __shfl_xor_sync(0xffffffffU, rank, 1);
+            rank += __shfl_xor_sync(0xffffffffU, rank, 2);
+            if (q == 0 && rank < kTopK)
+            {
+                route.selKey[rank] = key;
+                route.selSg[rank] = route.candSg[ci];
+            }
+        }
+        __syncthreads();
+        if (warp == 0)
+        {
+            if (lane < kTopK)
+            {
+                float s[kTopK];
+#pragma unroll
+                for (int k = 0; k < kTopK; ++k)
+                {
+                    s[k] = route.selSg[k];
+                }
+                // The native warp sum over lanes 0..7 (zeros above): butterfly offsets 16, 8, 4, 2, 1.
+                float const total = ((s[0] + s[4]) + (s[2] + s[6])) + ((s[1] + s[5]) + (s[3] + s[7]));
+                float const weight = route.selSg[lane] * params.routedScaling / total;
+                int const expert = 65535 - static_cast<int>(route.selKey[lane] & 0xFFFFULL);
+                params.topkIds[rf * kTopK + lane] = expert;
+                params.topkWeights[rf * kTopK + lane] = __bfloat162float(__float2bfloat16(weight));
+                ROUTING_STAMP(kStampRouted);
+            }
+        }
+    }
+
+    // ---- CTA 0 advances the protocol state and the gather parity once every CTA has read them. ----
+    if (tid == 0 && slice == 0 && clusterId == 0)
+    {
+        uint32_t arrived = 0;
+        do
+        {
+            asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory");
+        } while (arrived < static_cast<uint32_t>(kGridCtas));
+        uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers,
+            static_cast<uint32_t>(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U);
+        asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x),
+                     "r"(next.y), "r"(next.z), "r"(next.w)
+                     : "memory");
+        storeRelaxed(params.zstate, route.parity ^ 1U);
+    }
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+    __syncthreads();
+    if (tid == 0)
+    {
+        ROUTING_STAMP(kStampEnd);
+    }
+    __syncthreads();
+    if (tid < kNumStamps && params.stamps != nullptr)
+    {
+        params.stamps[ctaLinear * kNumStamps + tid] = route.stamp[tid];
+    }
+#endif
+#else
+    // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices.
+    __trap();
+#endif
+}
+
+template <int R, int kRowSplit, int kRouterLoad, bool kLogitsIn>
+cudaError_t launchTyped(PostAttentionHopRoutingParams const& params)
+{
+    using Cfg = HopConfig<kRowSplit>;
+    cudaLaunchAttribute attrs[2];
+    attrs[0].id = cudaLaunchAttributeClusterDimension;
+    attrs[0].val.clusterDim.x = kSlices;
+    attrs[0].val.clusterDim.y = 1;
+    attrs[0].val.clusterDim.z = 1;
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0;
+
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kSlices, kG * kRowSplit, 1);
+    config.blockDim = dim3(Cfg::kThreads, 1, 1);
+    config.dynamicSmemBytes = 0;
+    config.stream = params.stream;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    return cudaLaunchKernelEx(&config, postAttnHopZ2Kernel<R, kRowSplit, kRouterLoad, kLogitsIn>, params);
+}
+
+// 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices).
+template <int kRouterLoad, bool kLogitsIn>
+cudaError_t launchAt(PostAttentionHopRoutingParams const& params)
+{
+    switch (params.numRows)
+    {
+    case 1: return launchTyped<1, 1, kRouterLoad, kLogitsIn>(params);
+    case 2: return launchTyped<2, 2, kRouterLoad, kLogitsIn>(params);
+    case 3: return launchTyped<3, 3, kRouterLoad, kLogitsIn>(params);
+    case 4: return launchTyped<4, 3, kRouterLoad, kLogitsIn>(params);
+    case 5: return launchTyped<5, 3, kRouterLoad, kLogitsIn>(params);
+    case 6: return launchTyped<6, 3, kRouterLoad, kLogitsIn>(params);
+    case 7: return launchTyped<7, 3, kRouterLoad, kLogitsIn>(params);
+    case 8: return launchTyped<8, 3, kRouterLoad, kLogitsIn>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace
+
+bool stampsCompiled()
+{
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+    return true;
+#else
+    return false;
+#endif
+}
+
+cudaError_t invokePostAttentionHopRouting(PostAttentionHopRoutingParams const& params)
+{
+    if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr
+        || params.routerWeight == nullptr || params.routingBias == nullptr || params.routerLogits == nullptr
+        || params.topkIds == nullptr || params.topkWeights == nullptr || params.gather == nullptr
+        || params.zstate == nullptr)
+    {
+        return cudaErrorInvalidValue;
+    }
+    if (params.logitsIn != nullptr)
+    {
+        // Test instance: the default router-weight load point only.
+        return params.routerLoad == kRouterAfterSum ? launchAt<kRouterAfterSum, true>(params) : cudaErrorInvalidValue;
+    }
+    switch (params.routerLoad)
+    {
+    case kRouterAfterBarrier: return launchAt<kRouterAfterBarrier, false>(params);
+    case kRouterAfterSum: return launchAt<kRouterAfterSum, false>(params);
+    case kRouterBeforeMma: return launchAt<kRouterBeforeMma, false>(params);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace kernels::glm_moe_dsa::hop_routing
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h
new file mode 100644
index 000000000000..cb19a3f7f947
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h
@@ -0,0 +1,116 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::hop_routing
+{
+
+//! Pre-MoE hop with routing: the transport-input hop (postAttentionHopFromPartials.h) with router logits, plus
+//! the native DeepSeek-V3 top-8 routing of every row computed from those logits.
+//!
+//! Normalized hidden, residual and router logits are produced exactly as by
+//! ``glm_moe_dsa::attention_output_into_hop::invokePostAttentionHopFromPartials`` (same work decomposition, arithmetic
+//! and protocol on the same post-attention hop workspace), so they are bit-identical to that hop's.
+//!
+//! Routing (the native ``SigmoidBias`` + ``ScaledSumNormalize`` routing of the TRT-LLM-Gen routing runner with one
+//! expert group): every logit owner computes ``score = 0.5f * tanhf(0.5f * logit) + 0.5f + float(bias[e])`` and stores
+//! its bits into the gather buffer; one finisher CTA per row collects the row's 256 scores and selects the 8 largest
+//! native keys ``(TwiddleIn(score) << 32) | (65535 - e)`` (ties go to the lower expert) in descending order, and writes
+//! ``topkIds[r][k]`` and ``topkWeights[r][k] = bf16((score_k - bias_k) * routedScaling / sum)`` (stored as FP32), where
+//! ``sum`` adds the eight ``score_k - bias_k`` in the native warp-butterfly order ((s0+s4)+(s2+s6))+((s1+s5)+(s3+s7)).
+//!
+//! Gather protocol (private to this kernel): ``gather`` holds two buffers of ``kMaxRows x kNumExperts`` 32-bit words
+//! and ``zstate[0]`` the parity ``p`` of the buffer this call uses. A score is never -0.0f, so the word 0x80000000
+//! marks an empty slot. Each call resets buffer ``p ^ 1`` (the previous call's) to that sentinel and CTA 0 flips the
+//! parity together with the hop's state advance, after every CTA has read both states. Calls of other hop kinds on the
+//! same hop workspace do not touch the routing state.
+//!
+//! Memory-order contract with the producer (programmatic dependent launch): as
+//! ``glm_moe_dsa::attention_output_into_hop``. Before ``griddepcontrol.wait`` the kernel reads only the RMSNorm
+//! weight, the router weight (L2 prefetch) and the router bias, which no kernel writes while serving; with
+//! ``earlyPoll`` it also reads the protocol state of the post-attention hop workspace and polls the transport early
+//! (see postAttentionHopFromPartials.h). The routing state, the gather buffers and every output are read or written
+//! only after the wait. Finishers spin on gather words that every CTA writes before it waits on anything outside its
+//! own cluster, so the wait-for graph is acyclic.
+struct PostAttentionHopRoutingParams
+{
+    void const* residualIn = nullptr;   //!< [numRows, kHiddenDim] bf16
+    void const* gamma = nullptr;        //!< [kHiddenDim] bf16 RMSNorm weight
+    void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16
+    void const* routingBias = nullptr;  //!< [kNumExperts] bf16 score-correction bias
+    //! Test only: when set, the routing scores are computed from these [numRows, kNumExperts] FP32 logits instead of
+    //! the computed router logits (which are still written to ``routerLogits``). Read after the dependency wait.
+    float const* logitsIn = nullptr;
+    void* normOut = nullptr;       //!< [numRows, kHiddenDim] bf16
+    void* residualOut = nullptr;   //!< [numRows, kHiddenDim] bf16
+    float* routerLogits = nullptr; //!< [numRows, kNumExperts] fp32
+    int32_t* topkIds = nullptr;    //!< [numRows, kTopK] int32, descending native key order
+    float* topkWeights = nullptr;  //!< [numRows, kTopK] fp32 holding bf16-rounded weights
+    void* ucBuffer = nullptr;      //!< local unicast view of the post-attention hop's multicast workspace
+    uint32_t* flags = nullptr;     //!< [4] protocol state of that workspace
+    uint32_t* gather = nullptr;    //!< [kGatherWords] routing gather buffers
+    uint32_t* zstate = nullptr;    //!< [4] routing state; word 0 is the parity
+    //! Diagnostic builds only (``TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS``): [gridCtas, kNumStamps] globaltimer values.
+    unsigned long long* stamps = nullptr;
+    int numRows = 0;
+    int rank = 0;
+    float eps = 1e-5F;
+    float routedScaling = 2.5F;
+    bool enablePdl = true;
+    bool earlyPoll = false;
+    int routerLoad = 1;
+    cudaStream_t stream = nullptr;
+};
+
+using glm_moe_dsa::post_attention_hop::kHiddenDim;
+using glm_moe_dsa::post_attention_hop::kLamportBufferBytes;
+using glm_moe_dsa::post_attention_hop::kMaxRows;
+using glm_moe_dsa::post_attention_hop::kNumExperts;
+using glm_moe_dsa::post_attention_hop::kNumLamportBuffers;
+using glm_moe_dsa::post_attention_hop::kNumRanks;
+using glm_moe_dsa::post_attention_hop::kWorkspaceBytes;
+
+constexpr int kTopK = 8;
+//! Words of one gather buffer and of both.
+constexpr int kGatherBufferWords = kMaxRows * kNumExperts;
+constexpr int kGatherWords = 2 * kGatherBufferWords;
+constexpr int kStateWords = 4;
+//! Diagnostic stamp fields per CTA: entry, wait, last logit store, last gather store, poll complete, level-1 done,
+//! routing stored, end.
+constexpr int kNumStamps = 8;
+//! Upper bound of the grid size (4 expert groups x 3 row groups x 8 CTAs).
+constexpr int kMaxGridCtas = 96;
+
+//! Whether this build records diagnostic stamps.
+bool stampsCompiled();
+
+//! Launch the routing hop on ``params.stream``. Returns the launch status.
+cudaError_t invokePostAttentionHopRouting(PostAttentionHopRoutingParams const& params);
+
+} // namespace kernels::glm_moe_dsa::hop_routing
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu
new file mode 100644
index 000000000000..efe8a1e21e09
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu
@@ -0,0 +1,515 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h"
+
+#include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h"
+
+#include <cuda_fp8.h>
+
+#include <algorithm>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::rope_kv_append
+{
+
+namespace
+{
+
+constexpr int32_t kRopeThreads = 256;
+constexpr int32_t kBridgeThreads = 128;
+constexpr int32_t kMaxGroupSize = 8;
+constexpr int32_t kEltsPerVec = 8;                           // bf16 elements per 16-byte load
+constexpr int32_t kNopeVecs = kKvLoraRank / kEltsPerVec;     // 64
+constexpr int32_t kRopeVecs = kRopeDim / kEltsPerVec;        // 8
+constexpr int32_t kTotalVecsPerHead = kNopeVecs + kRopeVecs; // 72
+constexpr int32_t kIdxBytesPerItem = 4;
+
+__device__ __forceinline__ void pdlWait()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+#endif
+}
+
+__device__ __forceinline__ void pdlLaunchDependents()
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+#endif
+}
+
+//! Same arithmetic as quantCopy<__nv_bfloat16, 8> in mlaKernels.cu: bf16 -> fp32, scale, saturating E4M3.
+__device__ __forceinline__ uint2 quantizeBf16x8(uint4 const& src, float scale)
+{
+    __nv_bfloat162 const* in = reinterpret_cast<__nv_bfloat162 const*>(&src);
+    uint2 out;
+    __nv_fp8x2_e4m3* dst = reinterpret_cast<__nv_fp8x2_e4m3*>(&out);
+#pragma unroll
+    for (int32_t j = 0; j < 4; ++j)
+    {
+        float2 v = __bfloat1622float2(in[j]);
+        v.x *= scale;
+        v.y *= scale;
+        dst[j] = __nv_fp8x2_e4m3(v);
+    }
+    return out;
+}
+
+//! GPT-J RoPE on 8 bf16 values (4 adjacent pairs), identical to the per-pair transform of the native kernel.
+__device__ __forceinline__ void rotateBf16x8(uint4& data, float2 const (&coef)[4])
+{
+    __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&data);
+#pragma unroll
+    for (int32_t j = 0; j < 4; ++j)
+    {
+        pairs[j] = mmha::rotary_embedding_transform(pairs[j], coef[j]);
+    }
+}
+
+//! Native unravel of a flat byte index over the strided 4-D indexer cache view (indexerKCacheScatter.cu);
+//! 32-bit arithmetic when the index fits, same result.
+__device__ __forceinline__ int64_t idxFlatToOffset(int64_t flat, int32_t const (&dims)[4], int64_t const (&strides)[4])
+{
+    int32_t i0, i1, i2, i3;
+    if (flat >= 0 && flat <= 0x7fffffff)
+    {
+        uint32_t f = static_cast<uint32_t>(flat);
+        i3 = static_cast<int32_t>(f % static_cast<uint32_t>(dims[3]));
+        f /= static_cast<uint32_t>(dims[3]);
+        i2 = static_cast<int32_t>(f % static_cast<uint32_t>(dims[2]));
+        f /= static_cast<uint32_t>(dims[2]);
+        i1 = static_cast<int32_t>(f % static_cast<uint32_t>(dims[1]));
+        f /= static_cast<uint32_t>(dims[1]);
+        i0 = static_cast<int32_t>(f);
+    }
+    else
+    {
+        i3 = flat % dims[3];
+        flat /= dims[3];
+        i2 = flat % dims[2];
+        flat /= dims[2];
+        i1 = flat % dims[1];
+        flat /= dims[1];
+        i0 = flat;
+    }
+    return i0 * strides[0] + i1 * strides[1] + i2 * strides[2] + i3 * strides[3];
+}
+
+enum class RopeItem : int32_t
+{
+    kNone = 0,
+    kQNope,
+    kQRope,
+    kKRope,
+    kKNope,
+    kIdxData,
+    kIdxScale,
+};
+
+//! One thread per 16-byte work item of one token. Per token the item space is
+//!   [0, H*64)            q_nope FP8 quantization      (head = i / 64, vec = i % 64)
+//!   [H*64, H*72)         q_pe RoPE + FP8               (head, vec)
+//!   [H*72, H*72+8)       k_pe RoPE + FP8 cache write   (vec)
+//!   [H*72+8, H*72+72)    compressed-KV FP8 cache write (vec)
+//!   then idxHeadBytes/4 indexer K data items (4 bytes each) and one indexer scale item.
+//! Per-step metadata (cache lengths, KV block offsets, indexer slot mappings) and the item's source data are read
+//! only after griddepcontrol.wait (see the header).
+__global__ void __launch_bounds__(kRopeThreads) ropeAppendScatterKernel(RopeAppendParams const p)
+{
+    int32_t const token = blockIdx.x;
+    int32_t const item = blockIdx.y * kRopeThreads + threadIdx.x;
+    int32_t const qNopeEnd = p.numHeads * kNopeVecs;
+    int32_t const qRopeEnd = qNopeEnd + p.numHeads * kRopeVecs;
+    int32_t const kRopeEnd = qRopeEnd + kRopeVecs;
+    int32_t const kNopeEnd = kRopeEnd + kNopeVecs;
+    int32_t const idxDataEnd = kNopeEnd + p.idxHeadBytes / kIdxBytesPerItem;
+
+    bool const ropeToken = token < p.numTokens;
+    bool const idxToken = token < p.idxNumTokens;
+
+    RopeItem kind = RopeItem::kNone;
+    int32_t head = 0;
+    int32_t vec = 0;
+    if (item < qNopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kQNope : RopeItem::kNone;
+        head = item / kNopeVecs;
+        vec = item % kNopeVecs;
+    }
+    else if (item < qRopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kQRope : RopeItem::kNone;
+        head = (item - qNopeEnd) / kRopeVecs;
+        vec = (item - qNopeEnd) % kRopeVecs;
+    }
+    else if (item < kRopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kKRope : RopeItem::kNone;
+        vec = item - qRopeEnd;
+    }
+    else if (item < kNopeEnd)
+    {
+        kind = ropeToken ? RopeItem::kKNope : RopeItem::kNone;
+        vec = item - kRopeEnd;
+    }
+    else if (item < idxDataEnd)
+    {
+        kind = idxToken ? RopeItem::kIdxData : RopeItem::kNone;
+        vec = item - kNopeEnd;
+    }
+    else if (item == idxDataEnd)
+    {
+        kind = idxToken ? RopeItem::kIdxScale : RopeItem::kNone;
+    }
+
+    int32_t const batchIdx = token / p.seqLen;
+    int32_t const localIdx = token % p.seqLen;
+    bool const needPos = kind == RopeItem::kQRope || kind == RopeItem::kKRope || kind == RopeItem::kKNope;
+    bool const needCoef = kind == RopeItem::kQRope || kind == RopeItem::kKRope;
+    bool const needBlock = kind == RopeItem::kKRope || kind == RopeItem::kKNope;
+    bool const isIdx = kind == RopeItem::kIdxData || kind == RopeItem::kIdxScale;
+
+    // Values that never change while kernels of the step run (the quantization scales) may be read before the wait;
+    // per-step metadata and producer outputs may not.
+    float scale = 1.F;
+    if (needBlock)
+    {
+        scale = p.quantScaleKv ? p.quantScaleKv[0] : 1.F;
+    }
+    else if (kind == RopeItem::kQNope || kind == RopeItem::kQRope)
+    {
+        scale = p.quantScaleQ ? p.quantScaleQ[0] : 1.F;
+    }
+    auto loadCoef = [&](int32_t pos, float2 (&c)[4])
+    {
+        float2 const* src = p.cosSin + static_cast<size_t>(kRopeDim) * pos + (vec * kEltsPerVec) / 2;
+#pragma unroll
+        for (int32_t j = 0; j < 4; ++j)
+        {
+            c[j] = src[j];
+        }
+    };
+    auto idxOffset = [&](int64_t slotData, int64_t slotScale) -> int64_t
+    {
+        if (slotData < 0 || slotScale < 0)
+        {
+            return -1;
+        }
+        return kind == RopeItem::kIdxData ? idxFlatToOffset(slotData + vec * kIdxBytesPerItem, p.idxDims, p.idxStrides)
+                                          : idxFlatToOffset(slotScale, p.idxDims, p.idxStrides);
+    };
+    // Source data of this item; written by producers, so it is read after the wait.
+    uint4 vdata = make_uint4(0, 0, 0, 0);
+    uint32_t wdata = 0;
+    auto loadSource = [&]()
+    {
+        switch (kind)
+        {
+        case RopeItem::kQNope:
+            vdata = *reinterpret_cast<uint4 const*>(
+                p.fusedQ + (static_cast<size_t>(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec);
+            break;
+        case RopeItem::kQRope:
+            vdata = *reinterpret_cast<uint4 const*>(p.qPe + static_cast<size_t>(token) * p.qPeStride
+                + static_cast<size_t>(p.qPeLd) * head + vec * kEltsPerVec);
+            break;
+        case RopeItem::kKRope:
+            vdata = *reinterpret_cast<uint4 const*>(
+                p.latent + static_cast<size_t>(token) * kLatentDim + kKvLoraRank + vec * kEltsPerVec);
+            break;
+        case RopeItem::kKNope:
+            vdata = *reinterpret_cast<uint4 const*>(
+                p.latent + static_cast<size_t>(token) * kLatentDim + vec * kEltsPerVec);
+            break;
+        case RopeItem::kIdxData:
+            wdata = *reinterpret_cast<uint32_t const*>(
+                p.idxK + static_cast<int64_t>(token) * p.idxHeadBytes + vec * kIdxBytesPerItem);
+            break;
+        case RopeItem::kIdxScale:
+            wdata = *reinterpret_cast<uint32_t const*>(p.idxScale + static_cast<int64_t>(token) * kIdxBytesPerItem);
+            break;
+        default: break;
+        }
+    };
+
+    int32_t kvPos = 0;
+    float2 coef[4];
+    __nv_fp8_e4m3* kBlock = nullptr;
+    int64_t idxDst = -1;
+    pdlWait();
+    // Issue the source-data load first so it overlaps the metadata loads and the address math below.
+    loadSource();
+    if (needPos)
+    {
+        kvPos = p.kvLens[batchIdx] - p.seqLen + localIdx;
+    }
+    if (needCoef)
+    {
+        loadCoef(kvPos, coef);
+    }
+    if (needBlock)
+    {
+        kBlock = reinterpret_cast<__nv_fp8_e4m3*>(p.kvCache.getKBlockPtr(batchIdx, kvPos));
+    }
+    if (isIdx)
+    {
+        idxDst = idxOffset(p.idxSlotData[token], p.idxSlotScale[token]);
+    }
+
+    switch (kind)
+    {
+    case RopeItem::kQNope:
+    {
+        size_t const off = (static_cast<size_t>(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec;
+        *reinterpret_cast<uint2*>(p.quantQ + off) = quantizeBf16x8(vdata, scale);
+        break;
+    }
+    case RopeItem::kQRope:
+    {
+        uint4 data = vdata;
+        rotateBf16x8(data, coef);
+        size_t const dst
+            = (static_cast<size_t>(token) * p.numHeads + head) * kLatentDim + kKvLoraRank + vec * kEltsPerVec;
+        *reinterpret_cast<uint2*>(p.quantQ + dst) = quantizeBf16x8(data, scale);
+        break;
+    }
+    case RopeItem::kKRope:
+    {
+        uint4 data = vdata;
+        rotateBf16x8(data, coef);
+        int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, kNopeVecs + vec);
+        *reinterpret_cast<uint2*>(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(data, scale);
+        break;
+    }
+    case RopeItem::kKNope:
+    {
+        int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, vec);
+        *reinterpret_cast<uint2*>(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(vdata, scale);
+        if (vec == 0 && localIdx == 0)
+        {
+            p.cuQSeqlens[batchIdx + 1] = p.numHeads * p.seqLen * (batchIdx + 1);
+        }
+        break;
+    }
+    case RopeItem::kIdxData:
+    case RopeItem::kIdxScale:
+    {
+        if (idxDst >= 0)
+        {
+            *reinterpret_cast<uint32_t*>(p.idxCache + idxDst) = wdata;
+        }
+        break;
+    }
+    default: break;
+    }
+
+    if (blockIdx.x == 0 && blockIdx.y == 0)
+    {
+        if (threadIdx.x == 0)
+        {
+            p.fmhaTileCounter[0] = 0;
+            p.cuQSeqlens[0] = 0;
+            float const dequantQ = p.dequantScaleQ ? p.dequantScaleQ[0] : 1.F;
+            float const dequantKv = p.dequantScaleKv ? p.dequantScaleKv[0] : 1.F;
+            float const quantO = p.quantScaleO ? p.quantScaleO[0] : 1.F;
+            if (p.bmm1Scale)
+            {
+                constexpr float kLog2e = 1.4426950408889634074F;
+                float const bmm1 = dequantQ * dequantKv * p.hostBmm1Scale;
+                p.bmm1Scale[0] = bmm1;
+                p.bmm1Scale[1] = bmm1 * kLog2e;
+            }
+            if (p.bmm2Scale)
+            {
+                p.bmm2Scale[0] = quantO * dequantKv;
+            }
+        }
+        if (threadIdx.x < 32)
+        {
+            // Exclusive prefix sum of the cache lengths -> cuKvSeqlens[0..batch].
+            int32_t const lane = threadIdx.x;
+            int32_t const batch = p.numTokens / p.seqLen;
+            int32_t carry = 0;
+            for (int32_t base = 0; base <= batch; base += 32)
+            {
+                int32_t const idx = base + lane;
+                int32_t const len = idx < batch ? p.kvLens[idx] : 0;
+                int32_t incl = len;
+#pragma unroll
+                for (int32_t off = 1; off < 32; off <<= 1)
+                {
+                    int32_t const n = __shfl_up_sync(0xffffffffU, incl, off);
+                    if (lane >= off)
+                    {
+                        incl += n;
+                    }
+                }
+                if (idx <= batch)
+                {
+                    p.cuKvSeqlens[idx] = incl - len + carry;
+                }
+                carry += __shfl_sync(0xffffffffU, incl, 31);
+            }
+        }
+    }
+
+    pdlLaunchDependents();
+}
+
+__device__ __forceinline__ int32_t remapOne(
+    int32_t tok, int32_t layerId, int32_t blockSize, int32_t strideFactor, int32_t base)
+{
+    if (tok < 0 || base < 0)
+    {
+        return -1;
+    }
+    int32_t const inblockOff = tok % blockSize + layerId * blockSize;
+    return base * strideFactor + inblockOff;
+}
+
+//! One thread per (row, column) element, like the native remap, so the scattered block-table reads spread over
+//! many SMs; the group members' layer offsets are hoisted out of the store loop. Everything is read after the wait.
+__global__ void __launch_bounds__(kBridgeThreads) topkBridgeKernel(TopkBridgeParams const p)
+{
+    int32_t const rowSlot = blockIdx.x;
+    int32_t const col = blockIdx.y * kBridgeThreads + threadIdx.x;
+    bool const isOutRow = rowSlot < p.numRows;
+
+    int32_t req = 0;
+    int32_t srcRow = rowSlot;
+    int32_t layerIds[kMaxGroupSize];
+    pdlWait();
+    if (isOutRow)
+    {
+        req = p.reqIdx[rowSlot] - p.reqIdxBias;
+#pragma unroll
+        for (int32_t g = 0; g < kMaxGroupSize; ++g)
+        {
+            layerIds[g] = g < p.groupSize ? (p.layerIds ? p.layerIds[g] : p.layerId) : 0;
+        }
+    }
+    else
+    {
+        int32_t const i = rowSlot - p.numRows;
+        if (p.numAccepted == nullptr)
+        {
+            srcRow = i * p.nextN + p.nextN - 1;
+        }
+        else
+        {
+            int64_t const acc = p.numAcceptedIs64
+                ? static_cast<int64_t const*>(p.numAccepted)[i]
+                : static_cast<int64_t>(static_cast<int32_t const*>(p.numAccepted)[i]);
+            int64_t off = acc - 1;
+            off = off < 0 ? 0 : (off > p.nextN - 1 ? p.nextN - 1 : off);
+            srcRow = static_cast<int32_t>(static_cast<int64_t>(i) * p.nextN + off);
+        }
+    }
+
+    if (col < p.numTopk)
+    {
+        int32_t const tok = p.topk[static_cast<int64_t>(srcRow) * p.topkStride0 + col];
+        if (isOutRow)
+        {
+            int32_t base = -1;
+            if (tok >= 0)
+            {
+                int32_t const blockId = tok / p.blockSize;
+                if (blockId < p.maxBlocksPerReq)
+                {
+                    base = p.blockTable[req * p.btStride0 + blockId * p.btStride1];
+                }
+            }
+            int32_t* dst = p.out + static_cast<int64_t>(rowSlot) * p.outStride0 + col;
+#pragma unroll
+            for (int32_t g = 0; g < kMaxGroupSize; ++g)
+            {
+                if (g < p.groupSize)
+                {
+                    dst[g * p.outStrideG] = remapOne(tok, layerIds[g], p.blockSize, p.strideFactor, base);
+                }
+            }
+        }
+        bool const doCopy = p.copyDst != nullptr
+            && ((p.copyMode == 0 && isOutRow && rowSlot < p.copyRows) || (p.copyMode == 1 && !isOutRow));
+        if (doCopy)
+        {
+            int32_t const copyRow = p.copyRowStart + (isOutRow ? rowSlot : rowSlot - p.numRows);
+            p.copyDst[static_cast<int64_t>(copyRow) * p.copyStride0 + col] = tok;
+        }
+    }
+
+    pdlLaunchDependents();
+}
+
+template <typename Kernel, typename Params>
+void launchWithOptionalPdl(Kernel kernel, dim3 grid, dim3 block, bool enablePdl, cudaStream_t stream, Params const& p)
+{
+    cudaLaunchConfig_t config{};
+    config.gridDim = grid;
+    config.blockDim = block;
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0;
+    config.numAttrs = 1;
+    config.attrs = attrs;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p));
+}
+
+} // namespace
+
+void invokeRopeAppendScatter(RopeAppendParams const& params, bool enablePdl, cudaStream_t stream)
+{
+    int32_t const gridTokens = std::max(params.numTokens, params.idxNumTokens);
+    if (gridTokens == 0)
+    {
+        return;
+    }
+    TLLM_CHECK_WITH_INFO(params.seqLen > 0, "seqLen must be positive");
+    TLLM_CHECK_WITH_INFO(params.cosSinRows > 0, "the RoPE table must have at least one position");
+    TLLM_CHECK_WITH_INFO(params.idxHeadBytes % kIdxBytesPerItem == 0, "indexer head bytes must be a multiple of 4");
+    int32_t const itemsPerToken
+        = params.numHeads * kTotalVecsPerHead + kTotalVecsPerHead + params.idxHeadBytes / kIdxBytesPerItem + 1;
+    dim3 const grid(gridTokens, (itemsPerToken + kRopeThreads - 1) / kRopeThreads);
+    dim3 const block(kRopeThreads);
+    launchWithOptionalPdl(ropeAppendScatterKernel, grid, block, enablePdl, stream, params);
+}
+
+void invokeTopkBridge(TopkBridgeParams const& params, bool enablePdl, cudaStream_t stream)
+{
+    int32_t const extraRows = params.copyMode == 1 ? params.copyRows : 0;
+    int32_t const rows = params.numRows + extraRows;
+    if (rows == 0 || params.numTopk == 0)
+    {
+        return;
+    }
+    TLLM_CHECK_WITH_INFO(params.groupSize >= 1 && params.groupSize <= kMaxGroupSize,
+        "glm_moe_dsa_topk_index_bridge: group sizes must be 1..%d, got %d", kMaxGroupSize, params.groupSize);
+    dim3 const grid(rows, (params.numTopk + kBridgeThreads - 1) / kBridgeThreads);
+    dim3 const block(kBridgeThreads);
+    launchWithOptionalPdl(topkBridgeKernel, grid, block, enablePdl, stream, params);
+}
+
+} // namespace kernels::glm_moe_dsa::rope_kv_append
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h
new file mode 100644
index 000000000000..8f12c8b6cd61
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h
@@ -0,0 +1,136 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/kvCacheUtils.h"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::rope_kv_append
+{
+
+//! MLA absorption-mode latent width and RoPE width supported by the fused kernels.
+constexpr int32_t kKvLoraRank = 512;
+constexpr int32_t kRopeDim = 64;
+constexpr int32_t kLatentDim = kKvLoraRank + kRopeDim;
+
+//! Memory order of both kernels: they may start early as programmatic dependents (PDL) and read everything a
+//! producer writes, including per-step metadata (cache lengths, KV block offsets, indexer slot mappings, request
+//! indices, MTP accepted counts), only after griddepcontrol.wait. Only the quantization scales, which no kernel of the
+//! step writes, are read before the wait.
+
+//! Generation-phase MLA RoPE + FP8 query quantization + FP8 latent KV append, optionally fused with the DSA
+//! indexer K-cache scatter. With an FP8 MLA cache the outputs are byte-identical to
+//! applyMLARopeAndAssignQKVKernelGeneration (absorption mode, rope_append, no helix) followed by
+//! indexerKCacheScatterUnifiedKernel.
+struct RopeAppendParams
+{
+    // Query side.
+    __nv_bfloat16 const* fusedQ{nullptr}; //!< [numTokens, numHeads, kLatentDim]; the first kKvLoraRank are read.
+    __nv_bfloat16 const* qPe{nullptr};    //!< [numTokens, numHeads, kRopeDim] with strides (qPeStride, qPeLd, 1).
+    int32_t qPeLd{0};
+    int32_t qPeStride{0};
+    uint8_t* quantQ{nullptr};          //!< FP8 E4M3 [numTokens, numHeads, kLatentDim].
+    float const* quantScaleQ{nullptr}; //!< Optional; 1.0 when null.
+
+    // KV side.
+    __nv_bfloat16 const* latent{nullptr}; //!< [numTokens, kLatentDim] rows (packed).
+    float2 const* cosSin{nullptr};        //!< Rotary table addressed as cosSin[kRopeDim * position + d / 2].
+    int32_t cosSinRows{0};                //!< Positions in the rotary table.
+    int32_t const* kvLens{nullptr};       //!< [batch] cache lengths after this step's append.
+    float const* quantScaleKv{nullptr};   //!< Optional; 1.0 when null.
+    KVBlockArray kvCache;
+
+    // Scheduler buffers consumed by the FP8 MLA generation FMHA.
+    int32_t* cuQSeqlens{nullptr};
+    int32_t* cuKvSeqlens{nullptr};
+    uint32_t* fmhaTileCounter{nullptr};
+    float* bmm1Scale{nullptr};            //!< Optional [2].
+    float* bmm2Scale{nullptr};            //!< Optional [1].
+    float const* dequantScaleQ{nullptr};  //!< Optional; 1.0 when null.
+    float const* dequantScaleKv{nullptr}; //!< Optional; 1.0 when null.
+    float const* quantScaleO{nullptr};    //!< Optional; 1.0 when null.
+    float hostBmm1Scale{1.F};
+
+    int32_t numTokens{0};
+    int32_t seqLen{1};
+    int32_t numHeads{0};
+
+    // Optional DSA indexer K-cache scatter (skipped when idxNumTokens == 0).
+    uint8_t const* idxK{nullptr};     //!< [idxNumTokens, idxHeadBytes] bytes.
+    uint8_t const* idxScale{nullptr}; //!< [idxNumTokens, 4] bytes.
+    uint8_t* idxCache{nullptr};       //!< Strided 4-D byte view of one layer's indexer pool.
+    int64_t const* idxSlotData{nullptr};
+    int64_t const* idxSlotScale{nullptr};
+    int32_t idxNumTokens{0};
+    int32_t idxHeadBytes{128};
+    int32_t idxDims[4]{0, 0, 0, 0};
+    int64_t idxStrides[4]{0, 0, 0, 0};
+};
+
+//! Fused bridge from per-row local top-k token ids to global KV-pool indices, optionally writing the native
+//! top-k reuse copy in the same pass. Output values are identical to convertReqIndexToGlobal[Grouped]Kernel.
+struct TopkBridgeParams
+{
+    int32_t const* topk{nullptr}; //!< Source rows [srcRows, numTopk] with row stride topkStride0, unit column stride.
+    int64_t topkStride0{0};
+    int32_t numRows{0}; //!< Rows remapped to global indices.
+    int32_t numTopk{0};
+
+    int32_t const* reqIdx{nullptr}; //!< [numRows] request index per row, before subtracting reqIdxBias.
+    int32_t reqIdxBias{0};
+    int32_t const* blockTable{nullptr};
+    int64_t btStride0{0};
+    int64_t btStride1{1};
+    int32_t maxBlocksPerReq{0};
+    int32_t blockSize{0};
+    int32_t strideFactor{0};
+
+    int32_t const* layerIds{nullptr}; //!< Optional [groupSize]; when null the single layerId is used.
+    int32_t layerId{0};
+    int32_t groupSize{1};
+    int32_t* out{nullptr}; //!< [groupSize, numRows, numTopk] with strides (outStrideG, outStride0, 1).
+    int64_t outStrideG{0};
+    int64_t outStride0{0};
+
+    //! Copy rows: dst[copyRowStart + i, :numTopk] = topk[srcRow(i), :] for i in [0, copyRows).
+    //! srcRow(i) = i (copyMode 0) or the MTP last-accepted row i * nextN + clamp(numAccepted[i] - 1, 0, nextN - 1)
+    //! (copyMode 1; i * nextN + nextN - 1 when numAccepted is null).
+    int32_t* copyDst{nullptr};
+    int64_t copyStride0{0};
+    int32_t copyRows{0};
+    int32_t copyRowStart{0};
+    int32_t copyMode{0};
+    int32_t nextN{1};
+    void const* numAccepted{nullptr}; //!< Already offset to the first generation request.
+    bool numAcceptedIs64{false};
+};
+
+void invokeRopeAppendScatter(RopeAppendParams const& params, bool enablePdl, cudaStream_t stream);
+
+void invokeTopkBridge(TopkBridgeParams const& params, bool enablePdl, cudaStream_t stream);
+
+} // namespace kernels::glm_moe_dsa::rope_kv_append
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu
new file mode 100644
index 000000000000..b9c0a82f04ec
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu
@@ -0,0 +1,714 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/envUtils.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h"
+
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+
+#include <cmath>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::sparse_mla_decode
+{
+namespace
+{
+
+// Work split: one CTA handles one query row and 128 consecutive positions of its index list.
+// The 16 CTAs of a row form one thread-block cluster. Each CTA pushes its unnormalized partial
+// output and softmax statistics into the shared memory of the CTA that finalizes those output
+// dims with asynchronous stores that complete on that CTA's mbarrier.
+int constexpr kKeysPerCta = 128;
+int constexpr kNumSplits = kTopK / kKeysPerCta;
+int constexpr kNumWarps = 8;
+int constexpr kThreads = kNumWarps * 32;
+int constexpr kKeysPerWarp = kKeysPerCta / kNumWarps;
+int constexpr kDimsPerWarp = kHeadDimV / kNumWarps;
+int constexpr kDimsPerCtaOut = kHeadDimV / kNumSplits;
+[[maybe_unused]] int constexpr kQkSteps = kHeadDimQk / 32;
+[[maybe_unused]] int constexpr kPvSteps = kKeysPerCta / 32;
+[[maybe_unused]] int constexpr kPvTilesPerWarp = kDimsPerWarp / 16;
+[[maybe_unused]] int constexpr kQChunks = kNumHeads * kHeadDimQk / 16;
+// Probabilities are scaled to the e4m3 range before rounding, as in the native kernel.
+[[maybe_unused]] float constexpr kProbScale = 448.F;
+
+// Shared-memory pitches, padded so that ldmatrix and fragment loads are bank-conflict free: K rows
+// shift by 20 banks (eight ldmatrix rows hit disjoint 16-byte bank groups); the 64-bit query
+// fragment loads of rows g = 0-3 cover banks 8g..8g+7 (pitch = 8 banks mod 32); the 32-bit
+// probability fragment loads of rows g = 0-7 cover banks 4g..4g+3.
+int constexpr kKRowBytes = kHeadDimQk + 16;
+int constexpr kQRowHalves = kHeadDimQk + 16;
+int constexpr kPRowHalves = kKeysPerCta + 8;
+
+int constexpr kKTileOffset = 0;
+int constexpr kQTileOffset = kKTileOffset + kKeysPerCta * kKRowBytes;
+int constexpr kPTileOffset = kQTileOffset + kNumHeads * kQRowHalves * 2;
+// [source split][head][dim within this CTA's 32 output dims] FP32 partial outputs.
+int constexpr kRecvOffset = kPTileOffset + kNumHeads * kPRowHalves * 2;
+// [max of heads 0-7, sum of heads 0-7][source split] FP32.
+int constexpr kRecvStatsOffset = kRecvOffset + kNumSplits * kNumHeads * kDimsPerCtaOut * 4;
+int constexpr kRedMaxOffset = kRecvStatsOffset + kNumSplits * 2 * kNumHeads * 4;
+int constexpr kRedSumOffset = kRedMaxOffset + kNumWarps * kNumHeads * 4;
+// mbarriers: one per warp for its row gather, one for the pushed partials of this CTA.
+int constexpr kBarrierOffset = kRedSumOffset + kNumWarps * kNumHeads * 4;
+int constexpr kSmemBytes = kBarrierOffset + (kNumWarps + 1) * 8;
+// Bytes each CTA receives: 16 sources x (8 heads x 32 dims + 16 statistics) FP32 values.
+[[maybe_unused]] uint32_t constexpr kRecvBytes = kNumSplits * (kNumHeads * kDimsPerCtaOut + 2 * kNumHeads) * 4;
+
+static_assert(kTopK % kKeysPerCta == 0, "the index list must split evenly");
+static_assert(kNumSplits <= 16, "the merge cluster holds at most 16 CTAs");
+static_assert(kKeysPerWarp == 16, "one m16 tile of keys per warp");
+static_assert(kNumHeads == 8, "the eight heads form the n8 side of the MMAs");
+static_assert(kDimsPerCtaOut == 32 && kDimsPerWarp == 2 * kDimsPerCtaOut, "each warp feeds two merge CTAs");
+static_assert(kDimsPerCtaOut * kNumHeads == kThreads, "one merged output element per thread");
+static_assert(kKRowBytes % 16 == 0 && kQRowHalves % 8 == 0 && kPRowHalves % 8 == 0, "16-byte aligned shared rows");
+static_assert(
+    kKRowBytes / 4 % 32 == 20 && kQRowHalves / 2 % 32 == 8 && kPRowHalves / 2 % 32 == 4, "bank-conflict-free pitches");
+static_assert(kBarrierOffset % 8 == 0, "mbarrier alignment");
+static_assert(kQChunks <= 2 * kThreads, "two query chunks per thread at most");
+static_assert(2 * kNumHeads * kNumSplits == kThreads, "one statistics store per thread");
+static_assert(
+    kNumWarps == kNumHeads && (kNumSplits & (kNumSplits - 1)) == 0, "merge: one warp per head, one lane per split");
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+
+__device__ __forceinline__ uint32_t smemAddr(void const* ptr)
+{
+    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+__device__ __forceinline__ void mbarrierInit(uint32_t bar, uint32_t count)
+{
+    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory");
+}
+
+__device__ __forceinline__ void fenceBarrierInit()
+{
+    asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void mbarrierArriveExpectTx(uint32_t bar, uint32_t bytes)
+{
+    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory");
+}
+
+__device__ __forceinline__ void mbarrierArrive(uint32_t bar)
+{
+    asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" ::"r"(bar) : "memory");
+}
+
+__device__ __forceinline__ void mbarrierWait(uint32_t bar, uint32_t parity)
+{
+    uint32_t done = 0;
+    while (!done)
+    {
+        asm volatile(
+            "{\n"
+            ".reg .pred p;\n"
+            "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n"
+            "selp.u32 %0, 1, 0, p;\n"
+            "}\n"
+            : "=r"(done)
+            : "r"(bar), "r"(parity)
+            : "memory");
+    }
+}
+
+// The same wait with cluster-scope acquire: the arrivals and transaction bytes it waits for come
+// from asynchronous stores issued by other CTAs of the cluster.
+__device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar, uint32_t parity)
+{
+    uint32_t done = 0;
+    while (!done)
+    {
+        asm volatile(
+            "{\n"
+            ".reg .pred p;\n"
+            "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], %2;\n"
+            "selp.u32 %0, 1, 0, p;\n"
+            "}\n"
+            : "=r"(done)
+            : "r"(bar), "r"(parity)
+            : "memory");
+    }
+}
+
+__device__ __forceinline__ void bulkCopyG2S(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar)
+{
+    asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"(dst),
+                 "l"(src), "r"(bytes), "r"(bar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void gridDependencyWait()
+{
+    asm volatile("griddepcontrol.wait;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void gridDependencyLaunchDependents()
+{
+    asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void clusterArriveRelaxed()
+{
+    asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void clusterWait()
+{
+    asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory");
+}
+
+__device__ __forceinline__ uint32_t mapToRank(uint32_t localAddr, uint32_t rank)
+{
+    uint32_t remote;
+    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(remote) : "r"(localAddr), "r"(rank));
+    return remote;
+}
+
+// Asynchronous stores into a peer CTA's shared memory that complete-tx on the peer's mbarrier.
+__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float value, uint32_t remoteBar)
+{
+    asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];\n" ::"r"(remoteAddr),
+                 "f"(value), "r"(remoteBar)
+                 : "memory");
+}
+
+__device__ __forceinline__ void storeAsync2(uint32_t remoteAddr, float a, float b, uint32_t remoteBar)
+{
+    asm volatile(
+        "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];\n" ::"r"(remoteAddr),
+        "f"(a), "f"(b), "r"(remoteBar)
+        : "memory");
+}
+
+__device__ __forceinline__ void ldmatrixX4(uint32_t (&r)[4], uint32_t addr)
+{
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+                 : "r"(addr));
+}
+
+__device__ __forceinline__ void ldmatrixX4Trans(uint32_t (&r)[4], uint32_t addr)
+{
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+                 : "r"(addr));
+}
+
+// D += A * B with A 16x16 f16 (row), B 16x8 f16 (col), FP32 accumulation.
+__device__ __forceinline__ void mmaF16(
+    float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
+        "{%0, %1, %2, %3};\n"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+
+__device__ __forceinline__ float fastExp2(float x)
+{
+    float y;
+    asm volatile("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(y) : "f"(x));
+    return y;
+}
+
+// Two e4m3 values (low byte first) to two FP16 values; exact.
+__device__ __forceinline__ uint32_t e4m3x2ToHalf2(uint32_t pair)
+{
+    __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(pair & 0xffffU), __NV_E4M3);
+    return static_cast<uint32_t>(h.x) | (static_cast<uint32_t>(h.y) << 16);
+}
+
+// e4m3(448 * p) as an exact FP16 value.
+__device__ __forceinline__ __half roundProbability(float p)
+{
+    __nv_fp8_storage_t const q = __nv_cvt_float_to_fp8(p * kProbScale, __NV_SATFINITE, __NV_E4M3);
+    return __half(__nv_cvt_fp8_to_halfraw(q, __NV_E4M3));
+}
+
+#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+
+#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING
+// Development timeline: thread 0 of each CTA records the SM clock at phase boundaries and the
+// global timer at entry and exit.
+__device__ unsigned long long gTimeline[kTimelineMaxCtas][kTimelineSlots];
+#endif
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING
+__device__ __forceinline__ void markTimeline(int slot)
+{
+    if (threadIdx.x == 0 && blockIdx.x < kTimelineMaxCtas)
+    {
+        gTimeline[blockIdx.x][slot] = clock64();
+    }
+}
+
+__device__ __forceinline__ void markGlobalTime(int slot)
+{
+    if (threadIdx.x == 0 && blockIdx.x < kTimelineMaxCtas)
+    {
+        unsigned long long t;
+        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+        gTimeline[blockIdx.x][slot] = t;
+    }
+}
+#else
+__device__ __forceinline__ void markTimeline(int) {}
+
+__device__ __forceinline__ void markGlobalTime(int) {}
+#endif
+#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+
+__global__ void __launch_bounds__(kThreads) sparseMlaDecodeKernel(SparseMlaDecodeParams const params)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    extern __shared__ __align__(128) uint8_t smem[];
+
+    int const tid = threadIdx.x;
+    int const warp = tid >> 5;
+    int const lane = tid & 31;
+    int const groupId = lane >> 2;
+    int const tig = lane & 3;
+    int const row = blockIdx.x / kNumSplits;
+    int const split = blockIdx.x % kNumSplits;
+
+    uint8_t* kTile = smem + kKTileOffset;
+    __half* qTile = reinterpret_cast<__half*>(smem + kQTileOffset);
+    __half* pTile = reinterpret_cast<__half*>(smem + kPTileOffset);
+    float* recv = reinterpret_cast<float*>(smem + kRecvOffset);
+    float* recvStats = reinterpret_cast<float*>(smem + kRecvStatsOffset);
+    float* redMax = reinterpret_cast<float*>(smem + kRedMaxOffset);
+    float* redSum = reinterpret_cast<float*>(smem + kRedSumOffset);
+    uint32_t const gatherBar = smemAddr(smem + kBarrierOffset + warp * 8);
+    uint32_t const recvBar = smemAddr(smem + kBarrierOffset + kNumWarps * 8);
+
+    markGlobalTime(0);
+    markTimeline(1);
+    if (lane == 0)
+    {
+        mbarrierInit(gatherBar, kKeysPerWarp);
+    }
+    if (tid == 0)
+    {
+        mbarrierInit(recvBar, 1);
+    }
+    // The initializations must be visible to the bulk-copy engine and to the peers' async stores.
+    fenceBarrierInit();
+    __syncthreads();
+    clusterArriveRelaxed();
+
+    // Everything below reads tensors written by preceding kernels in the stream. The kernel can be
+    // resident before those kernels finish, so these tensors are read with coherent loads, not
+    // through the non-coherent read-only path.
+    gridDependencyWait();
+    markTimeline(2);
+
+    // The query does not depend on the index list: start its loads first.
+    uint8_t const* qRow = static_cast<uint8_t const*>(params.q) + static_cast<int64_t>(row) * kNumHeads * kHeadDimQk;
+    uint4 qRaw[2];
+#pragma unroll
+    for (int i = 0; i < 2; ++i)
+    {
+        int const chunk = tid + i * kThreads;
+        if (chunk < kQChunks)
+        {
+            qRaw[i] = reinterpret_cast<uint4 const*>(qRow)[chunk];
+        }
+    }
+
+    int32_t poolRow = -1;
+    if (lane < kKeysPerWarp)
+    {
+        int const seq = row / params.seqLenQ;
+        int const rowInSeq = row - seq * params.seqLenQ;
+        int const key = split * kKeysPerCta + warp * kKeysPerWarp + lane;
+        int32_t const listed = params.indices[static_cast<int64_t>(row) * params.indicesRowStride + key];
+        int const usable = params.seqLensKv[seq] - (params.seqLenQ - 1 - rowInSeq);
+        poolRow = key < usable ? listed : -1;
+    }
+    uint32_t const validBits = __ballot_sync(0xffffffffU, poolRow >= 0);
+    markTimeline(3);
+
+    // Gather this warp's 16 rows with one 576-byte bulk copy each; unusable rows are zeroed so that
+    // zero probabilities never meet non-finite bytes.
+    if (lane < kKeysPerWarp)
+    {
+        uint8_t* dst = kTile + (warp * kKeysPerWarp + lane) * kKRowBytes;
+        if (poolRow >= 0)
+        {
+            mbarrierArriveExpectTx(gatherBar, kHeadDimQk);
+            bulkCopyG2S(smemAddr(dst),
+                static_cast<uint8_t const*>(params.kvPool) + static_cast<int64_t>(poolRow) * kHeadDimQk, kHeadDimQk,
+                gatherBar);
+        }
+        else
+        {
+#pragma unroll 4
+            for (int c = 0; c < kHeadDimQk / 16; ++c)
+            {
+                *reinterpret_cast<uint4*>(dst + c * 16) = make_uint4(0U, 0U, 0U, 0U);
+            }
+            mbarrierArrive(gatherBar);
+        }
+    }
+
+    // FP16 copy of the query, shared by all warps.
+#pragma unroll
+    for (int i = 0; i < 2; ++i)
+    {
+        int const chunk = tid + i * kThreads;
+        if (chunk < kQChunks)
+        {
+            int const head = chunk / (kHeadDimQk / 16);
+            int const col = (chunk - head * (kHeadDimQk / 16)) * 16;
+            uint4 lo;
+            uint4 hi;
+            lo.x = e4m3x2ToHalf2(qRaw[i].x);
+            lo.y = e4m3x2ToHalf2(qRaw[i].x >> 16);
+            lo.z = e4m3x2ToHalf2(qRaw[i].y);
+            lo.w = e4m3x2ToHalf2(qRaw[i].y >> 16);
+            hi.x = e4m3x2ToHalf2(qRaw[i].z);
+            hi.y = e4m3x2ToHalf2(qRaw[i].z >> 16);
+            hi.z = e4m3x2ToHalf2(qRaw[i].w);
+            hi.w = e4m3x2ToHalf2(qRaw[i].w >> 16);
+            uint4* dst = reinterpret_cast<uint4*>(qTile + head * kQRowHalves + col);
+            dst[0] = lo;
+            dst[1] = hi;
+        }
+    }
+    float const scaleLog2 = params.bmm1Scale[1];
+    float const outScale = params.bmm2Scale[0];
+
+    gridDependencyLaunchDependents();
+    __syncthreads();
+
+    // Every peer has initialized its barriers once phase 0 of the cluster barrier completes; waiting
+    // here overlaps that with the gather. Pushes into peers start after the softmax.
+    clusterWait();
+    mbarrierWait(gatherBar, 0);
+    __syncwarp();
+    markTimeline(4);
+
+    // S^T = K * Q^T for this warp's 16 keys: keys on the M side, heads on the N side, FP16 MMAs on
+    // converted e4m3 keys. Within each 16-value k chunk the MMA k-slot order is permuted so that
+    // the two halves of an ldmatrix register map straight onto the A fragment: slots 2t, 2t+1 hold
+    // k = 4t, 4t+1 and slots 2t+8, 2t+9 hold k = 4t+2, 4t+3. The FP16 query is read in that order.
+    float sAcc[4];
+    {
+        float acc[2][4] = {};
+        uint32_t const aAddr = smemAddr(kTile + (warp * kKeysPerWarp + (lane & 15)) * kKRowBytes + (lane >> 4) * 16);
+        __half const* qFrag = qTile + groupId * kQRowHalves + 4 * tig;
+#pragma unroll
+        for (int kk = 0; kk < kQkSteps; ++kk)
+        {
+            uint32_t r[4];
+            ldmatrixX4(r, aAddr + kk * 32);
+            uint2 const b0 = *reinterpret_cast<uint2 const*>(qFrag + kk * 32);
+            uint2 const b1 = *reinterpret_cast<uint2 const*>(qFrag + kk * 32 + 16);
+            mmaF16(acc[0], e4m3x2ToHalf2(r[0]), e4m3x2ToHalf2(r[1]), e4m3x2ToHalf2(r[0] >> 16),
+                e4m3x2ToHalf2(r[1] >> 16), b0.x, b0.y);
+            mmaF16(acc[1], e4m3x2ToHalf2(r[2]), e4m3x2ToHalf2(r[3]), e4m3x2ToHalf2(r[2] >> 16),
+                e4m3x2ToHalf2(r[3] >> 16), b1.x, b1.y);
+        }
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+        {
+            sAcc[i] = acc[0][i] + acc[1][i];
+        }
+    }
+    markTimeline(5);
+
+    // Fragment element (key, head): 0 -> (groupId, 2*tig), 1 -> (groupId, 2*tig+1), 2/3 -> key + 8.
+    bool const valid0 = (validBits >> groupId) & 1U;
+    bool const valid1 = (validBits >> (groupId + 8)) & 1U;
+    float s[4];
+    s[0] = valid0 ? sAcc[0] * scaleLog2 : -INFINITY;
+    s[1] = valid0 ? sAcc[1] * scaleLog2 : -INFINITY;
+    s[2] = valid1 ? sAcc[2] * scaleLog2 : -INFINITY;
+    s[3] = valid1 ? sAcc[3] * scaleLog2 : -INFINITY;
+
+    float maxA = fmaxf(s[0], s[2]);
+    float maxB = fmaxf(s[1], s[3]);
+#pragma unroll
+    for (int offset = 4; offset < 32; offset <<= 1)
+    {
+        maxA = fmaxf(maxA, __shfl_xor_sync(0xffffffffU, maxA, offset));
+        maxB = fmaxf(maxB, __shfl_xor_sync(0xffffffffU, maxB, offset));
+    }
+    if (lane < 4)
+    {
+        redMax[warp * kNumHeads + 2 * tig] = maxA;
+        redMax[warp * kNumHeads + 2 * tig + 1] = maxB;
+    }
+    __syncthreads();
+
+    float headMaxA = -INFINITY;
+    float headMaxB = -INFINITY;
+#pragma unroll
+    for (int w = 0; w < kNumWarps; ++w)
+    {
+        headMaxA = fmaxf(headMaxA, redMax[w * kNumHeads + 2 * tig]);
+        headMaxB = fmaxf(headMaxB, redMax[w * kNumHeads + 2 * tig + 1]);
+    }
+    float p[4];
+    p[0] = s[0] == -INFINITY ? 0.F : fastExp2(s[0] - headMaxA);
+    p[1] = s[1] == -INFINITY ? 0.F : fastExp2(s[1] - headMaxB);
+    p[2] = s[2] == -INFINITY ? 0.F : fastExp2(s[2] - headMaxA);
+    p[3] = s[3] == -INFINITY ? 0.F : fastExp2(s[3] - headMaxB);
+
+    float sumA = p[0] + p[2];
+    float sumB = p[1] + p[3];
+#pragma unroll
+    for (int offset = 4; offset < 32; offset <<= 1)
+    {
+        sumA += __shfl_xor_sync(0xffffffffU, sumA, offset);
+        sumB += __shfl_xor_sync(0xffffffffU, sumB, offset);
+    }
+    if (lane < 4)
+    {
+        redSum[warp * kNumHeads + 2 * tig] = sumA;
+        redSum[warp * kNumHeads + 2 * tig + 1] = sumB;
+    }
+    {
+        int const key = warp * kKeysPerWarp + groupId;
+        pTile[(2 * tig) * kPRowHalves + key] = roundProbability(p[0]);
+        pTile[(2 * tig + 1) * kPRowHalves + key] = roundProbability(p[1]);
+        pTile[(2 * tig) * kPRowHalves + key + 8] = roundProbability(p[2]);
+        pTile[(2 * tig + 1) * kPRowHalves + key + 8] = roundProbability(p[3]);
+    }
+    __syncthreads();
+
+    // This CTA's softmax statistics to every CTA of the cluster: thread t sends value t / 16
+    // (maxima of heads 0-7, then sums) to rank t % 16.
+    {
+        int const index = tid / kNumSplits;
+        int const dest = tid % kNumSplits;
+        int const head = index & (kNumHeads - 1);
+        bool const isMax = index < kNumHeads;
+        float value = isMax ? -INFINITY : 0.F;
+#pragma unroll
+        for (int w = 0; w < kNumWarps; ++w)
+        {
+            value = isMax ? fmaxf(value, redMax[w * kNumHeads + head]) : value + redSum[w * kNumHeads + head];
+        }
+        storeAsync(mapToRank(smemAddr(recvStats + index * kNumSplits + split), dest), value, mapToRank(recvBar, dest));
+    }
+    markTimeline(6);
+
+    // O^T = V^T * P^T for this warp's 64 output dims, one 16-dim tile at a time; each finished tile
+    // is pushed to the CTA that finalizes its dims while the next tile is computed. Transposing the
+    // e4m3 V rows with b16 ldmatrix yields keys (2t, 2t+1) x dims (2g, 2g+1) per register; a byte
+    // permute splits it into the two dims, so M row g of a tile is dim 2g and M row g+8 is dim 2g+1.
+    {
+        uint32_t pFrag[kPvSteps][2][2];
+        __half const* pRow = pTile + groupId * kPRowHalves + 2 * tig;
+#pragma unroll
+        for (int ks = 0; ks < kPvSteps; ++ks)
+        {
+#pragma unroll
+            for (int half = 0; half < 2; ++half)
+            {
+                pFrag[ks][half][0] = *reinterpret_cast<uint32_t const*>(pRow + ks * 32 + half * 16);
+                pFrag[ks][half][1] = *reinterpret_cast<uint32_t const*>(pRow + ks * 32 + half * 16 + 8);
+            }
+        }
+        uint32_t const vAddr = smemAddr(kTile + lane * kKRowBytes + warp * kDimsPerWarp);
+        uint32_t const recvBase = smemAddr(recv);
+        float const invProbScale = 1.F / kProbScale;
+#pragma unroll
+        for (int mt = 0; mt < kPvTilesPerWarp; ++mt)
+        {
+            // Independent accumulators per MMA keep the tensor pipe busy within one tile.
+            float part[kPvSteps][2][4] = {};
+#pragma unroll
+            for (int ks = 0; ks < kPvSteps; ++ks)
+            {
+                uint32_t r[4];
+                ldmatrixX4Trans(r, vAddr + ks * 32 * kKRowBytes + mt * 16);
+                uint32_t dimEven[4];
+                uint32_t dimOdd[4];
+#pragma unroll
+                for (int j = 0; j < 4; ++j)
+                {
+                    uint32_t const pairs = __byte_perm(r[j], 0, 0x3120);
+                    dimEven[j] = e4m3x2ToHalf2(pairs);
+                    dimOdd[j] = e4m3x2ToHalf2(pairs >> 16);
+                }
+                mmaF16(part[ks][0], dimEven[0], dimOdd[0], dimEven[1], dimOdd[1], pFrag[ks][0][0], pFrag[ks][0][1]);
+                mmaF16(part[ks][1], dimEven[2], dimOdd[2], dimEven[3], dimOdd[3], pFrag[ks][1][0], pFrag[ks][1][1]);
+            }
+            float o[4];
+#pragma unroll
+            for (int i = 0; i < 4; ++i)
+            {
+                o[i] = ((part[0][0][i] + part[0][1][i]) + (part[1][0][i] + part[1][1][i]))
+                    + ((part[2][0][i] + part[2][1][i]) + (part[3][0][i] + part[3][1][i]));
+            }
+            int const dim = warp * kDimsPerWarp + mt * 16 + 2 * groupId;
+            uint32_t const dest = dim / kDimsPerCtaOut;
+            int const localDim = dim - static_cast<int>(dest) * kDimsPerCtaOut;
+            uint32_t const addrA
+                = recvBase + ((split * kNumHeads + 2 * tig) * kDimsPerCtaOut + localDim) * sizeof(float);
+            uint32_t const addrB = addrA + kDimsPerCtaOut * sizeof(float);
+            uint32_t const remoteBar = mapToRank(recvBar, dest);
+            storeAsync2(mapToRank(addrA, dest), o[0] * invProbScale, o[2] * invProbScale, remoteBar);
+            storeAsync2(mapToRank(addrB, dest), o[1] * invProbScale, o[3] * invProbScale, remoteBar);
+        }
+    }
+    markTimeline(7);
+    markTimeline(8);
+    if (tid == 0)
+    {
+        mbarrierArriveExpectTx(recvBar, kRecvBytes);
+    }
+    markTimeline(9);
+    mbarrierWaitCluster(recvBar, 0);
+    markTimeline(10);
+
+    // Finalize 32 of the 512 output dims of all eight heads from the pushed partials: warp h
+    // handles head h, lane i < 16 the statistics of source split i, every lane one output dim.
+    {
+        int const head = warp;
+        float const peerMax = lane < kNumSplits ? recvStats[head * kNumSplits + lane] : -INFINITY;
+        float const peerSum = lane < kNumSplits ? recvStats[(kNumHeads + head) * kNumSplits + lane] : 0.F;
+        float globalMax = peerMax;
+#pragma unroll
+        for (int offset = kNumSplits / 2; offset > 0; offset >>= 1)
+        {
+            globalMax = fmaxf(globalMax, __shfl_xor_sync(0xffffffffU, globalMax, offset));
+        }
+        float const weight = peerMax == -INFINITY ? 0.F : fastExp2(peerMax - globalMax);
+        float sum = weight * peerSum;
+#pragma unroll
+        for (int offset = kNumSplits / 2; offset > 0; offset >>= 1)
+        {
+            sum += __shfl_xor_sync(0xffffffffU, sum, offset);
+        }
+        sum = __shfl_sync(0xffffffffU, sum, 0);
+        float acc = 0.F;
+#pragma unroll
+        for (int i = 0; i < kNumSplits; ++i)
+        {
+            acc += __shfl_sync(0xffffffffU, weight, i) * recv[(i * kNumHeads + head) * kDimsPerCtaOut + lane];
+        }
+        float const result = sum > 0.F ? acc * (outScale / sum) : 0.F;
+        params.out[static_cast<int64_t>(row) * kNumHeads * kHeadDimV + head * kHeadDimV + split * kDimsPerCtaOut + lane]
+            = __float2bfloat16_rn(result);
+    }
+    markTimeline(11);
+    markGlobalTime(12);
+#else
+    __trap();
+#endif
+}
+
+void prepareKernelAttributes()
+{
+    TLLM_CUDA_CHECK(
+        cudaFuncSetAttribute(sparseMlaDecodeKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes));
+    TLLM_CUDA_CHECK(cudaFuncSetAttribute(sparseMlaDecodeKernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
+}
+
+cudaLaunchAttribute clusterAttribute()
+{
+    cudaLaunchAttribute attr;
+    attr.id = cudaLaunchAttributeClusterDimension;
+    attr.val.clusterDim.x = kNumSplits;
+    attr.val.clusterDim.y = 1;
+    attr.val.clusterDim.z = 1;
+    return attr;
+}
+
+} // namespace
+
+#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING
+void copySparseMlaDecodeTimeline(unsigned long long* host)
+{
+    TLLM_CUDA_CHECK(cudaMemcpyFromSymbol(host, gTimeline, sizeof(gTimeline)));
+}
+#endif
+
+size_t getSparseMlaDecodeSmemBytes()
+{
+    return static_cast<size_t>(kSmemBytes);
+}
+
+bool isSparseMlaDecodeDeviceSupported(int device)
+{
+    int major = 0;
+    int maxSmem = 0;
+    TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device));
+    TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device));
+    int constexpr kBlackwellDatacenterMajor = 10;
+    if (major != kBlackwellDatacenterMajor || maxSmem < kSmemBytes)
+    {
+        return false;
+    }
+    int currentDevice = 0;
+    TLLM_CUDA_CHECK(cudaGetDevice(&currentDevice));
+    TLLM_CUDA_CHECK(cudaSetDevice(device));
+    prepareKernelAttributes();
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(kNumSplits);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = kSmemBytes;
+    cudaLaunchAttribute attr = clusterAttribute();
+    config.attrs = &attr;
+    config.numAttrs = 1;
+    int numClusters = 0;
+    cudaError_t const status = cudaOccupancyMaxActiveClusters(&numClusters, sparseMlaDecodeKernel, &config);
+    TLLM_CUDA_CHECK(cudaSetDevice(currentDevice));
+    if (status != cudaSuccess)
+    {
+        // Clear the error recorded by the failed query.
+        cudaGetLastError();
+        return false;
+    }
+    return numClusters > 0;
+}
+
+void invokeSparseMlaDecode(SparseMlaDecodeParams const& params, cudaStream_t stream)
+{
+    if (params.numRows <= 0)
+    {
+        return;
+    }
+    prepareKernelAttributes();
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(params.numRows * kNumSplits);
+    config.blockDim = dim3(kThreads);
+    config.dynamicSmemBytes = kSmemBytes;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[2];
+    attrs[0] = clusterAttribute();
+    attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0;
+    config.attrs = attrs;
+    config.numAttrs = 2;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, sparseMlaDecodeKernel, params));
+}
+
+} // namespace kernels::glm_moe_dsa::sparse_mla_decode
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h
new file mode 100644
index 000000000000..18dadfa6a556
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h
@@ -0,0 +1,108 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::glm_moe_dsa::sparse_mla_decode
+{
+
+// Sparse MLA generation attention over an FP8 (e4m3) latent KV pool.
+//
+// Each query row attends to the pool rows listed in its `kTopK`-entry index list. A pool row
+// holds the 512-dim latent followed by the 64-dim rotary part, 576 e4m3 values in total; the
+// value vector is the first 512 values of the same row. Eight query heads share the row.
+//
+// Rules of the TRT-LLM generation FMHA for this layout, with the one difference marked below:
+// * Index-list position j of query row r is used only when j < min(kTopK, kvLen - (seqLenQ - 1 - r % seqLenQ)),
+//   where kvLen counts every token of the sequence including this step's rows. Repeated pool rows
+//   count once per occurrence.
+// * Difference: a negative entry at a used position is skipped. The native kernel does not check
+//   it and reads the pool row that -1 addresses. The native producers (indexer top-k, local-to-global
+//   remap) write -1 only at positions past the causal length, so the outputs agree on their inputs.
+// * Scores are FP32 dot products of the e4m3 query and key times bmm1Scale[1] (log2 domain).
+// * Probabilities relative to a running maximum are multiplied by 448 and rounded to e4m3 before
+//   the P*V product; row sums use the unrounded FP32 probabilities.
+// * A row without a usable entry produces zeros.
+//
+// Programmatic dependent launch (PDL), active when TRTLLM_ENABLE_PDL allows it:
+// * The kernel is launched with programmatic stream serialization. Before griddepcontrol.wait it
+//   only initializes its own shared-memory barriers and arrives on its cluster barrier; it accesses
+//   no global memory there, so it may start while the preceding kernel still writes its inputs.
+// * After the wait it reads q, indices, seqLensKv, bmm1Scale and bmm2Scale and issues the
+//   asynchronous gathers of the pool rows. It then executes griddepcontrol.launch_dependents, before
+//   those gathers complete and before any compute, so the next kernel of the stream can become
+//   resident while this one still runs.
+// * Its only global write is out, stored at the end of each CTA.
+// * Obligation of a PDL dependent: execute griddepcontrol.wait (cudaGridDependencySynchronize)
+//   before reading out and before writing any buffer read here (q, indices, seqLensKv, the scales,
+//   the pool rows). Work before that wait may access only memory this kernel neither reads nor
+//   writes, such as weights streamed into shared memory or L2; such work competes with this
+//   kernel's gathers for memory bandwidth and SM resources.
+
+int constexpr kNumHeads = 8;
+int constexpr kHeadDimQk = 576;
+int constexpr kHeadDimV = 512;
+int constexpr kTopK = 2048;
+
+struct SparseMlaDecodeParams
+{
+    // [numRows, kNumHeads, kHeadDimQk] e4m3 query, rows contiguous.
+    void const* q;
+    // Base of the FP8 latent pool; pool row i starts at kvPool + i * kHeadDimQk bytes.
+    void const* kvPool;
+    // [numRows, indicesRowStride] int32 pool rows.
+    int32_t const* indices;
+    int64_t indicesRowStride;
+    // [2] device scales; element 1 is the dequantization scale times the softmax scale times log2(e).
+    float const* bmm1Scale;
+    // [1] device scale applied to the normalized P*V output.
+    float const* bmm2Scale;
+    // [numRows / seqLenQ] KV lengths including this step's rows.
+    int32_t const* seqLensKv;
+    int32_t seqLenQ;
+    // [numRows, kNumHeads, kHeadDimV] BF16 output.
+    __nv_bfloat16* out;
+    int32_t numRows;
+};
+
+// Dynamic shared memory needed by one CTA.
+size_t getSparseMlaDecodeSmemBytes();
+
+// Whether the current device can launch the kernel (architecture and 16-CTA clusters).
+bool isSparseMlaDecodeDeviceSupported(int device);
+
+void invokeSparseMlaDecode(SparseMlaDecodeParams const& params, cudaStream_t stream);
+
+#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING
+// Development builds only: per-CTA phase timeline of the most recent launches.
+int constexpr kTimelineMaxCtas = 4096;
+int constexpr kTimelineSlots = 13;
+void copySparseMlaDecodeTimeline(unsigned long long* host);
+#endif
+
+} // namespace kernels::glm_moe_dsa::sparse_mla_decode
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt
index ab424318c862..228d096f7022 100644
--- a/cpp/tensorrt_llm/thop/CMakeLists.txt
+++ b/cpp/tensorrt_llm/thop/CMakeLists.txt
@@ -161,6 +161,13 @@ add_th_common_family(
   mhcOp.cpp
   compressorOp.cpp)
 
+add_th_common_family(th_common_glm_moe_dsa)
+
+# GLM-MoE-DSA torch ops live in glmMoeDsa/ (kernels in
+# cpp/tensorrt_llm/kernels/glmMoeDsa/); that directory adds them to
+# th_common_glm_moe_dsa with target_sources().
+add_subdirectory(glmMoeDsa)
+
 add_library(
   th_common SHARED
   $<TARGET_OBJECTS:th_common_attention>
@@ -169,7 +176,8 @@ add_library(
   $<TARGET_OBJECTS:th_common_distributed>
   $<TARGET_OBJECTS:th_common_sequence>
   $<TARGET_OBJECTS:th_common_visual_gen>
-  $<TARGET_OBJECTS:th_common_decoding>)
+  $<TARGET_OBJECTS:th_common_decoding>
+  $<TARGET_OBJECTS:th_common_glm_moe_dsa>)
 set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_target_properties(th_common PROPERTIES CXX_STANDARD 20 CXX_STANDARD_REQUIRED
                                                            ON)
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt b/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt
new file mode 100644
index 000000000000..7f8bc2d1593c
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt
@@ -0,0 +1,19 @@
+# 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.
+
+# torch.ops.trtllm.glm_moe_dsa_* ops, one file per kernel of
+# cpp/tensorrt_llm/kernels/glmMoeDsa/. The parent defines th_common_glm_moe_dsa.
+file(GLOB_RECURSE SRC_CPP CONFIGURE_DEPENDS *.cpp)
+target_sources(th_common_glm_moe_dsa PRIVATE ${SRC_CPP})
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp
new file mode 100644
index 000000000000..d53886b3d03f
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp
@@ -0,0 +1,108 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
+
+#include <optional>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+namespace
+{
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::cross_rank_topk_cluster_merge;
+
+void crossRankTopkClusterMerge(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, at::Tensor& state,
+    int64_t mcWorkspace, int64_t ucWorkspace, int64_t rank, int64_t worldSize, int64_t nextN,
+    std::optional<at::Tensor> const& debugTimes)
+{
+    TORCH_CHECK(logits.is_cuda() && kvLens.is_cuda() && out.is_cuda() && state.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(logits.scalar_type() == at::kFloat && logits.dim() == 2 && logits.stride(1) == 1,
+        "logits must be a row-major [rows, local_capacity] float32 view");
+    TORCH_CHECK(kvLens.scalar_type() == at::kInt && kvLens.numel() >= 1, "kv_lens must be int32 with >= 1 element");
+    TORCH_CHECK(out.scalar_type() == at::kInt && out.dim() == 2 && out.is_contiguous() && out.size(1) == tk::kTopK,
+        "out must be a contiguous int32 [rows, ", tk::kTopK, "] tensor");
+    TORCH_CHECK(out.size(0) == logits.size(0), "out and logits must have the same rows");
+    TORCH_CHECK(logits.size(0) <= tk::kMaxRows, "at most ", tk::kMaxRows, " rows");
+    TORCH_CHECK(state.scalar_type() == at::kInt && state.numel() >= tk::kStateWords && state.is_contiguous(),
+        "state must be a contiguous int32 tensor of >= ", tk::kStateWords, " words");
+    TORCH_CHECK(worldSize >= 1 && worldSize <= tk::kMaxRanks && rank >= 0 && rank < worldSize, "bad rank/world");
+    TORCH_CHECK(nextN >= 1 && nextN == logits.size(0), "next_n must equal the number of rows (one request)");
+    TORCH_CHECK(logits.size(1) <= tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)), "local capacity ",
+        logits.size(1), " exceeds the kernel limit ", tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)));
+    TORCH_CHECK(mcWorkspace != 0 && ucWorkspace != 0 && (mcWorkspace % 16) == 0 && (ucWorkspace % 16) == 0,
+        "workspace addresses must be non-null and 16-byte aligned");
+
+    tk::CpTopkParams params{};
+    params.logits = logits.data_ptr<float>();
+    params.logitsStride = logits.stride(0);
+    params.localCapacity = static_cast<int32_t>(logits.size(1));
+    params.kvLens = kvLens.data_ptr<int32_t>();
+    params.out = out.data_ptr<int32_t>();
+    params.numRows = static_cast<int32_t>(logits.size(0));
+    params.nextN = static_cast<int32_t>(nextN);
+    params.rank = static_cast<int32_t>(rank);
+    params.worldSize = static_cast<int32_t>(worldSize);
+    params.mcWorkspace = reinterpret_cast<uint32_t*>(mcWorkspace);
+    params.ucWorkspace = reinterpret_cast<uint32_t*>(ucWorkspace);
+    params.state = state.data_ptr<int32_t>();
+    if (debugTimes.has_value())
+    {
+        TORCH_CHECK(debugTimes->is_cuda() && debugTimes->scalar_type() == at::kLong && debugTimes->is_contiguous()
+                && debugTimes->numel() >= logits.size(0) * tk::kSourcesPerRank * tk::kDebugSlots,
+            "debug_times must be a contiguous int64 CUDA tensor of rows x ", tk::kSourcesPerRank, " x ",
+            tk::kDebugSlots);
+        params.debugTimes = debugTimes->data_ptr<int64_t>();
+    }
+    tk::invokeCpTopk(params, at::cuda::getCurrentCUDAStream(logits.get_device()));
+}
+
+//! [workspace words, top-k, max ranks, max rows, page tokens, sources per rank, first-round slot capacity, max local
+//! capacity, state words, Lamport buffers, default timeout (ms), resident clusters of kMaxRows-row launches]
+std::vector<int64_t> crossRankTopkClusterMergeWorkspaceInfo(int64_t worldSize)
+{
+    return {tk::kWorkspaceWords, tk::kTopK, tk::kMaxRanks, tk::kMaxRows, tk::kPageTokens, tk::kSourcesPerRank,
+        tk::kSourceCapacity, tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)), tk::kStateWords,
+        tk::kNumLamportBuffers, tk::cpTopkDefaultTimeoutMs(),
+        tk::cpTopkMaxActiveClusters(static_cast<int32_t>(worldSize))};
+}
+
+} // namespace
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_cross_rank_topk_cluster_merge(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) state, "
+        "int mc_workspace, "
+        "int uc_workspace, int rank, int world_size, int next_n, Tensor(c!)? debug_times=None) -> ()");
+    m.def("glm_moe_dsa_cross_rank_topk_cluster_merge_workspace_info(int world_size) -> int[]",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkClusterMergeWorkspaceInfo);
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl(
+        "glm_moe_dsa_cross_rank_topk_cluster_merge", &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkClusterMerge);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp
new file mode 100644
index 000000000000..0796f5bd96ee
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp
@@ -0,0 +1,187 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
+
+#include <optional>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+namespace
+{
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::cross_rank_topk_global_emit;
+
+//! The cluster-merge cross-rank top-k (same arguments as glm_moe_dsa_cross_rank_topk_cluster_merge) that also writes
+//! the outputs of the top-k index bridge (trtllm::glm_moe_dsa_topk_index_bridge with the same bridge arguments, applied
+//! to `out`): returns the pool rows, [G, rows, top-k] with layer_ids (G = its length) or [rows, top-k] without, and
+//! writes the copy rows into copy_dst.
+at::Tensor crossRankTopkGlobalEmit(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out,
+    at::Tensor& state, int64_t mcWorkspace, int64_t ucWorkspace, int64_t rank, int64_t worldSize, int64_t nextN,
+    at::Tensor const& reqIdx, int64_t reqIdxBias, at::Tensor const& blockTable, int64_t blockSize, int64_t strideFactor,
+    std::optional<at::Tensor> const& layerIds, int64_t layerId, std::optional<at::Tensor> const& copyDst,
+    int64_t copyRowStart, int64_t copyRows, int64_t copyMode, int64_t copyNextN,
+    std::optional<at::Tensor> const& numAccepted, std::optional<at::Tensor> const& debugTimes)
+{
+    TORCH_CHECK(logits.is_cuda() && kvLens.is_cuda() && out.is_cuda() && state.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(logits.scalar_type() == at::kFloat && logits.dim() == 2 && logits.stride(1) == 1,
+        "logits must be a row-major [rows, local_capacity] float32 view");
+    TORCH_CHECK(kvLens.scalar_type() == at::kInt && kvLens.numel() >= 1, "kv_lens must be int32 with >= 1 element");
+    TORCH_CHECK(out.scalar_type() == at::kInt && out.dim() == 2 && out.is_contiguous() && out.size(1) == tk::kTopK,
+        "out must be a contiguous int32 [rows, ", tk::kTopK, "] tensor");
+    TORCH_CHECK(out.size(0) == logits.size(0), "out and logits must have the same rows");
+    TORCH_CHECK(logits.size(0) <= tk::kMaxRows, "at most ", tk::kMaxRows, " rows");
+    TORCH_CHECK(state.scalar_type() == at::kInt && state.numel() >= tk::kStateWords && state.is_contiguous(),
+        "state must be a contiguous int32 tensor of >= ", tk::kStateWords, " words");
+    TORCH_CHECK(worldSize >= 1 && worldSize <= tk::kMaxRanks && rank >= 0 && rank < worldSize, "bad rank/world");
+    TORCH_CHECK(nextN >= 1 && nextN == logits.size(0), "next_n must equal the number of rows (one request)");
+    TORCH_CHECK(logits.size(1) <= tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)), "local capacity ",
+        logits.size(1), " exceeds the kernel limit ", tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)));
+    TORCH_CHECK(mcWorkspace != 0 && ucWorkspace != 0 && (mcWorkspace % 16) == 0 && (ucWorkspace % 16) == 0,
+        "workspace addresses must be non-null and 16-byte aligned");
+
+    // The bridge inputs, with the rules of trtllm::glm_moe_dsa_topk_index_bridge (topk = out, num_rows = rows).
+    int64_t const rows = out.size(0);
+    TORCH_CHECK(reqIdx.is_cuda() && blockTable.is_cuda(), "req_idx and block_table must be CUDA tensors");
+    TORCH_CHECK(reqIdx.scalar_type() == at::kInt && reqIdx.dim() == 1 && reqIdx.is_contiguous(),
+        "req_idx must be packed int32");
+    TORCH_CHECK(reqIdx.size(0) >= rows, "req_idx shorter than the rows");
+    TORCH_CHECK(blockTable.scalar_type() == at::kInt && blockTable.dim() == 2, "block_table must be int32 2-D");
+    TORCH_CHECK(blockSize == tk::kPageTokens, "block_size must be ", tk::kPageTokens, " (the CP page)");
+    TORCH_CHECK(copyMode == 0 || copyMode == 1, "copy_mode must be 0 or 1");
+    TORCH_CHECK(copyNextN >= 1 && copyRows >= 0 && copyRowStart >= 0, "bad copy rows");
+    TORCH_CHECK(copyMode == 1 || copyRows <= rows, "copy_rows exceeds the rows");
+    TORCH_CHECK(copyMode == 0 || copyRows * copyNextN <= rows, "copy_rows * copy_next_n exceeds the rows");
+
+    int32_t groupSize = 1;
+    int32_t const* layerIdsPtr = nullptr;
+    if (layerIds.has_value())
+    {
+        TORCH_CHECK(layerIds->is_cuda() && layerIds->scalar_type() == at::kInt && layerIds->dim() == 1
+                && layerIds->is_contiguous(),
+            "layer_ids must be packed int32 on the GPU");
+        groupSize = static_cast<int32_t>(layerIds->size(0));
+        TORCH_CHECK(groupSize >= 1 && groupSize <= tk::kMaxGroupSize,
+            "glm_moe_dsa_cross_rank_topk_global_emit: remap group sizes must be 1..", tk::kMaxGroupSize, ", got ",
+            groupSize);
+        layerIdsPtr = layerIds->data_ptr<int32_t>();
+    }
+    int64_t const topK = tk::kTopK;
+    at::Tensor pool = layerIds.has_value() ? at::empty({static_cast<int64_t>(groupSize), rows, topK}, out.options())
+                                           : at::empty({rows, topK}, out.options());
+
+    tk::CpTopkParams params{};
+    params.logits = logits.data_ptr<float>();
+    params.logitsStride = logits.stride(0);
+    params.localCapacity = static_cast<int32_t>(logits.size(1));
+    params.kvLens = kvLens.data_ptr<int32_t>();
+    params.out = out.data_ptr<int32_t>();
+    params.numRows = static_cast<int32_t>(rows);
+    params.nextN = static_cast<int32_t>(nextN);
+    params.rank = static_cast<int32_t>(rank);
+    params.worldSize = static_cast<int32_t>(worldSize);
+    params.mcWorkspace = reinterpret_cast<uint32_t*>(mcWorkspace);
+    params.ucWorkspace = reinterpret_cast<uint32_t*>(ucWorkspace);
+    params.state = state.data_ptr<int32_t>();
+
+    tk::CpTopkBridgeParams& b = params.bridge;
+    b.out = pool.data_ptr<int32_t>();
+    b.outStrideG = layerIds.has_value() ? pool.stride(0) : 0;
+    b.outStride0 = layerIds.has_value() ? pool.stride(1) : pool.stride(0);
+    b.groupSize = groupSize;
+    b.layerIds = layerIdsPtr;
+    b.layerId = static_cast<int32_t>(layerId);
+    b.reqIdx = reqIdx.data_ptr<int32_t>();
+    b.reqIdxBias = static_cast<int32_t>(reqIdxBias);
+    b.blockTable = blockTable.data_ptr<int32_t>();
+    b.btStride0 = blockTable.stride(0);
+    b.btStride1 = blockTable.stride(1);
+    b.maxBlocksPerReq = static_cast<int32_t>(blockTable.size(1));
+    b.strideFactor = static_cast<int32_t>(strideFactor);
+    if (copyDst.has_value() && copyRows > 0)
+    {
+        auto const& dst = copyDst.value();
+        TORCH_CHECK(dst.is_cuda() && dst.scalar_type() == at::kInt && dst.dim() == 2 && dst.stride(1) == 1
+                && dst.size(1) >= tk::kTopK && dst.size(0) >= copyRowStart + copyRows,
+            "copy_dst must be int32 [>= start + rows, >= top-k] with unit column stride");
+        b.copyDst = dst.data_ptr<int32_t>();
+        b.copyStride0 = dst.stride(0);
+        b.copyRows = static_cast<int32_t>(copyRows);
+        b.copyRowStart = static_cast<int32_t>(copyRowStart);
+        b.copyMode = static_cast<int32_t>(copyMode);
+        b.copyNextN = static_cast<int32_t>(copyNextN);
+        if (copyMode == 1 && numAccepted.has_value())
+        {
+            auto const& acc = numAccepted.value();
+            TORCH_CHECK(acc.is_cuda() && (acc.scalar_type() == at::kInt || acc.scalar_type() == at::kLong)
+                    && acc.dim() == 1 && acc.is_contiguous() && acc.size(0) >= copyRows,
+                "num_accepted must be packed int32/int64 with at least copy_rows entries");
+            b.numAccepted = acc.data_ptr();
+            b.numAcceptedIs64 = acc.scalar_type() == at::kLong;
+        }
+    }
+    if (debugTimes.has_value())
+    {
+        TORCH_CHECK(debugTimes->is_cuda() && debugTimes->scalar_type() == at::kLong && debugTimes->is_contiguous()
+                && debugTimes->numel() >= logits.size(0) * tk::kSourcesPerRank * tk::kDebugSlots,
+            "debug_times must be a contiguous int64 CUDA tensor of rows x ", tk::kSourcesPerRank, " x ",
+            tk::kDebugSlots);
+        params.debugTimes = debugTimes->data_ptr<int64_t>();
+    }
+    tk::invokeCpTopk(params, at::cuda::getCurrentCUDAStream(logits.get_device()));
+    return pool;
+}
+
+//! [workspace words, top-k, max ranks, max rows, page tokens, sources per rank, first-round slot capacity, max local
+//! capacity, state words, Lamport buffers, default timeout (ms), resident clusters of kMaxRows-row launches, max remap
+//! group size]
+std::vector<int64_t> crossRankTopkGlobalEmitWorkspaceInfo(int64_t worldSize)
+{
+    return {tk::kWorkspaceWords, tk::kTopK, tk::kMaxRanks, tk::kMaxRows, tk::kPageTokens, tk::kSourcesPerRank,
+        tk::kSourceCapacity, tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)), tk::kStateWords,
+        tk::kNumLamportBuffers, tk::cpTopkDefaultTimeoutMs(),
+        tk::cpTopkMaxActiveClusters(static_cast<int32_t>(worldSize)), tk::kMaxGroupSize};
+}
+
+} // namespace
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_cross_rank_topk_global_emit(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) state, int "
+        "mc_workspace, "
+        "int uc_workspace, int rank, int world_size, int next_n, Tensor req_idx, int req_idx_bias, "
+        "Tensor block_table, int block_size, int stride_factor, Tensor? layer_ids, int layer_id, "
+        "Tensor(c!)? copy_dst, int copy_row_start, int copy_rows, int copy_mode, int copy_next_n, "
+        "Tensor? num_accepted, Tensor(d!)? debug_times=None) -> Tensor");
+    m.def("glm_moe_dsa_cross_rank_topk_global_emit_workspace_info(int world_size) -> int[]",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkGlobalEmitWorkspaceInfo);
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_cross_rank_topk_global_emit", &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkGlobalEmit);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp
new file mode 100644
index 000000000000..710b0e716b88
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp
@@ -0,0 +1,104 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
+
+#include <optional>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+namespace
+{
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::cross_rank_topk;
+
+void crossRankTopk(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, at::Tensor& state,
+    int64_t mcWorkspace, int64_t ucWorkspace, int64_t rank, int64_t worldSize, int64_t nextN,
+    std::optional<at::Tensor> const& debugTimes)
+{
+    TORCH_CHECK(logits.is_cuda() && kvLens.is_cuda() && out.is_cuda() && state.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(logits.scalar_type() == at::kFloat && logits.dim() == 2 && logits.stride(1) == 1,
+        "logits must be a row-major [rows, local_capacity] float32 view");
+    TORCH_CHECK(kvLens.scalar_type() == at::kInt && kvLens.numel() >= 1, "kv_lens must be int32 with >= 1 element");
+    TORCH_CHECK(out.scalar_type() == at::kInt && out.dim() == 2 && out.is_contiguous() && out.size(1) == tk::kTopK,
+        "out must be a contiguous int32 [rows, ", tk::kTopK, "] tensor");
+    TORCH_CHECK(out.size(0) == logits.size(0), "out and logits must have the same rows");
+    TORCH_CHECK(logits.size(0) <= tk::kMaxRows, "at most ", tk::kMaxRows, " rows");
+    TORCH_CHECK(state.scalar_type() == at::kInt && state.numel() >= tk::kStateWords && state.is_contiguous(),
+        "state must be a contiguous int32 tensor of >= ", tk::kStateWords, " words");
+    TORCH_CHECK(worldSize >= 1 && worldSize <= tk::kMaxRanks && rank >= 0 && rank < worldSize, "bad rank/world");
+    TORCH_CHECK(nextN >= 1 && nextN == logits.size(0), "next_n must equal the number of rows (one request)");
+    TORCH_CHECK(logits.size(1) <= tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)), "local capacity ",
+        logits.size(1), " exceeds the kernel limit ", tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)));
+    TORCH_CHECK(mcWorkspace != 0 && ucWorkspace != 0 && (mcWorkspace % 16) == 0 && (ucWorkspace % 16) == 0,
+        "workspace addresses must be non-null and 16-byte aligned");
+
+    tk::CpTopkParams params{};
+    params.logits = logits.data_ptr<float>();
+    params.logitsStride = logits.stride(0);
+    params.localCapacity = static_cast<int32_t>(logits.size(1));
+    params.kvLens = kvLens.data_ptr<int32_t>();
+    params.out = out.data_ptr<int32_t>();
+    params.numRows = static_cast<int32_t>(logits.size(0));
+    params.nextN = static_cast<int32_t>(nextN);
+    params.rank = static_cast<int32_t>(rank);
+    params.worldSize = static_cast<int32_t>(worldSize);
+    params.mcWorkspace = reinterpret_cast<uint32_t*>(mcWorkspace);
+    params.ucWorkspace = reinterpret_cast<uint32_t*>(ucWorkspace);
+    params.state = state.data_ptr<int32_t>();
+    if (debugTimes.has_value())
+    {
+        TORCH_CHECK(debugTimes->is_cuda() && debugTimes->scalar_type() == at::kLong && debugTimes->is_contiguous()
+                && debugTimes->numel() >= logits.size(0) * tk::kSourcesPerRank * tk::kDebugSlots,
+            "debug_times must be a contiguous int64 CUDA tensor of rows x ", tk::kSourcesPerRank, " x ",
+            tk::kDebugSlots);
+        params.debugTimes = debugTimes->data_ptr<int64_t>();
+    }
+    tk::invokeCpTopk(params, at::cuda::getCurrentCUDAStream(logits.get_device()));
+}
+
+std::vector<int64_t> crossRankTopkWorkspaceInfo(int64_t worldSize)
+{
+    return {tk::kWorkspaceWords, tk::kTopK, tk::kMaxRanks, tk::kMaxRows, tk::kPageTokens, tk::kSourcesPerRank,
+        tk::kSourceCapacity, tk::cpTopkMaxLocalCapacity(static_cast<int32_t>(worldSize)), tk::kStateWords,
+        tk::kNumLamportBuffers, tk::cpTopkDefaultTimeoutMs()};
+}
+
+} // namespace
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_cross_rank_topk(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) state, int "
+        "mc_workspace, "
+        "int uc_workspace, int rank, int world_size, int next_n, Tensor(c!)? debug_times=None) -> ()");
+    m.def("glm_moe_dsa_cross_rank_topk_workspace_info(int world_size) -> int[]",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkWorkspaceInfo);
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_cross_rank_topk", &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopk);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp
new file mode 100644
index 000000000000..dd335256942d
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp
@@ -0,0 +1,131 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/envUtils.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/denseMlp.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
+
+#include <cstdint>
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::dense_mlp;
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+constexpr int64_t kVecAlignBytes = 16;
+
+void checkBf16Matrix(at::Tensor const& t, char const* name)
+{
+    TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+    TORCH_CHECK(t.scalar_type() == at::kBFloat16, name, " must be BF16");
+    TORCH_CHECK(t.dim() == 2, name, " must be 2-D");
+    TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+    TORCH_CHECK(reinterpret_cast<uintptr_t>(t.data_ptr()) % kVecAlignBytes == 0, name, " must be 16-byte aligned");
+}
+
+void launchGateUp(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tensor& h, int64_t variant, bool enablePdl)
+{
+    auto const stream = at::cuda::getCurrentCUDAStream(x.get_device());
+    TLLM_CUDA_CHECK(tk::invokeGateUpSwiglu(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()),
+        reinterpret_cast<__nv_bfloat16 const*>(gateUpWeight.data_ptr()), reinterpret_cast<__nv_bfloat16*>(h.data_ptr()),
+        static_cast<int>(x.size(0)), static_cast<int>(x.size(1)), static_cast<int>(h.size(1)),
+        static_cast<int>(variant), enablePdl, stream));
+}
+
+void launchDownProj(
+    at::Tensor const& h, at::Tensor const& downWeight, at::Tensor& out, int64_t variant, bool enablePdl)
+{
+    auto const stream = at::cuda::getCurrentCUDAStream(h.get_device());
+    TLLM_CUDA_CHECK(tk::invokeDown(reinterpret_cast<__nv_bfloat16 const*>(h.data_ptr()),
+        reinterpret_cast<__nv_bfloat16 const*>(downWeight.data_ptr()), reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
+        static_cast<int>(h.size(0)), static_cast<int>(h.size(1)), static_cast<int>(out.size(1)),
+        static_cast<int>(variant), enablePdl, stream));
+}
+
+} // namespace
+
+//! SwiGLU of the fused gate/up projection: x [M, H], gate_up_weight [2I, H] -> h [M, I].
+at::Tensor denseMlpGateUpSwiglu(at::Tensor const& x, at::Tensor const& gateUpWeight, int64_t variant)
+{
+    checkBf16Matrix(x, "x");
+    checkBf16Matrix(gateUpWeight, "gate_up_weight");
+    int64_t const numRows = x.size(0);
+    int64_t const hidden = x.size(1);
+    TORCH_CHECK(gateUpWeight.size(1) == hidden, "gate_up_weight must be [2 * inter, hidden]");
+    TORCH_CHECK(gateUpWeight.size(0) % 2 == 0, "gate_up_weight rows must be even");
+    int64_t const inter = gateUpWeight.size(0) / 2;
+    TORCH_CHECK(tk::isGateUpSupported(static_cast<int>(numRows), static_cast<int>(hidden), static_cast<int>(inter),
+                    static_cast<int>(variant)),
+        "glm_moe_dsa_dense_mlp_gate_up_swiglu: unsupported shape or variant");
+    at::Tensor h = at::empty({numRows, inter}, x.options());
+    launchGateUp(x, gateUpWeight, h, variant, tensorrt_llm::common::getEnvEnablePDL());
+    return h;
+}
+
+//! Down projection: h [M, I], down_weight [H, I] -> out [M, H] (rank-local partial, no reduction).
+at::Tensor denseMlpDown(at::Tensor const& h, at::Tensor const& downWeight, int64_t variant)
+{
+    checkBf16Matrix(h, "h");
+    checkBf16Matrix(downWeight, "down_weight");
+    int64_t const numRows = h.size(0);
+    int64_t const inter = h.size(1);
+    TORCH_CHECK(downWeight.size(1) == inter, "down_weight must be [hidden, inter]");
+    int64_t const hidden = downWeight.size(0);
+    TORCH_CHECK(tk::isDownSupported(static_cast<int>(numRows), static_cast<int>(inter), static_cast<int>(hidden),
+                    static_cast<int>(variant)),
+        "glm_moe_dsa_dense_mlp_down: unsupported shape or variant");
+    at::Tensor out = at::empty({numRows, hidden}, h.options());
+    launchDownProj(h, downWeight, out, variant, tensorrt_llm::common::getEnvEnablePDL());
+    return out;
+}
+
+//! Dense gated MLP without the output all-reduce: x [M, H] -> rank-local partial [M, H].
+at::Tensor denseMlp(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tensor const& downWeight,
+    int64_t gateUpVariant, int64_t downVariant)
+{
+    at::Tensor const h = denseMlpGateUpSwiglu(x, gateUpWeight, gateUpVariant);
+    return denseMlpDown(h, downWeight, downVariant);
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def("glm_moe_dsa_dense_mlp_gate_up_swiglu(Tensor x, Tensor gate_up_weight, int variant=1) -> Tensor");
+    m.def("glm_moe_dsa_dense_mlp_down(Tensor h, Tensor down_weight, int variant=0) -> Tensor");
+    m.def(
+        "glm_moe_dsa_dense_mlp(Tensor x, Tensor gate_up_weight, Tensor down_weight, int gate_up_variant=1, "
+        "int down_variant=0) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_dense_mlp_gate_up_swiglu", &tensorrt_llm::torch_ext::glm_moe_dsa::denseMlpGateUpSwiglu);
+    m.impl("glm_moe_dsa_dense_mlp_down", &tensorrt_llm::torch_ext::glm_moe_dsa::denseMlpDown);
+    m.impl("glm_moe_dsa_dense_mlp", &tensorrt_llm::torch_ext::glm_moe_dsa::denseMlp);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp
new file mode 100644
index 000000000000..b439f5d003fd
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp
@@ -0,0 +1,150 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/library.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::draft_live_rows;
+
+// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0.
+bool pdlEnabled()
+{
+    static bool const enabled = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && std::strcmp(env, "0") == 0);
+    }();
+    return enabled;
+}
+
+void checkRows(at::Tensor const& t, char const* name)
+{
+    TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_draft_live_row_hop] ", name,
+        " must be a CUDA bf16 tensor");
+    TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_draft_live_row_hop] ", name, " must be contiguous");
+    TORCH_CHECK(t.dim() >= 1 && t.size(-1) == tk::kHiddenDim, "[glm_moe_dsa_draft_live_row_hop] ", name,
+        " must have hidden size ", tk::kHiddenDim);
+}
+
+} // namespace
+
+//! Live-row post-attention hop: for the rows ``row_ids`` of ``input`` / ``residual`` only, the one-shot cross-rank
+//! sum + residual + RMSNorm + router logits of the post-attention hop, on that hop's workspace.
+//! Returns ``[normed, residual_out, router_logits]`` with ``row_ids.numel()`` rows (compact, in ``row_ids`` order).
+//! See kernels/glmMoeDsa/draftLiveRowHop.h for the arithmetic and the workspace protocol.
+std::vector<at::Tensor> draftLiveRowHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& rowIds,
+    at::Tensor const& gamma, double eps, at::Tensor const& routerWeight, at::Tensor const& ucBuffer, int64_t mcPtr,
+    at::Tensor const& flags, int64_t rank, int64_t nRanks)
+{
+    checkRows(input, "input");
+    checkRows(residual, "residual");
+    TORCH_CHECK(
+        residual.sizes() == input.sizes(), "[glm_moe_dsa_draft_live_row_hop] residual must have the shape of input");
+    checkRows(gamma, "gamma");
+    TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_draft_live_row_hop] gamma must be 1-D");
+    int64_t const numInputRows = input.numel() / tk::kHiddenDim;
+    TORCH_CHECK(numInputRows >= 1 && numInputRows <= tk::kMaxRows,
+        "[glm_moe_dsa_draft_live_row_hop] number of input rows ", numInputRows, " is outside the supported range [1, ",
+        tk::kMaxRows, "]");
+    TORCH_CHECK(rowIds.is_cuda() && rowIds.scalar_type() == at::kLong && rowIds.dim() == 1 && rowIds.is_contiguous(),
+        "[glm_moe_dsa_draft_live_row_hop] row_ids must be a contiguous 1-D CUDA int64 tensor");
+    int64_t const numRows = rowIds.numel();
+    TORCH_CHECK(numRows >= 1 && numRows <= numInputRows, "[glm_moe_dsa_draft_live_row_hop] number of live rows ",
+        numRows, " is outside [1, ", numInputRows, "]");
+    TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_draft_live_row_hop] only TP", tk::kNumRanks,
+        " is supported, got ", nRanks);
+    TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_draft_live_row_hop] invalid rank ", rank);
+    TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous()
+            && static_cast<size_t>(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes,
+        "[glm_moe_dsa_draft_live_row_hop] uc_buffer must be a contiguous CUDA tensor of at least ", tk::kWorkspaceBytes,
+        " bytes");
+    TORCH_CHECK(
+        mcPtr != 0, "[glm_moe_dsa_draft_live_row_hop] mc_ptr must be the multicast address of uc_buffer's workspace");
+    TORCH_CHECK(flags.is_cuda() && flags.is_contiguous() && flags.scalar_type() == at::kInt && flags.numel() >= 4,
+        "[glm_moe_dsa_draft_live_row_hop] flags must be a contiguous CUDA int32 tensor with at least 4 elements");
+    TORCH_CHECK(routerWeight.is_cuda() && routerWeight.scalar_type() == at::kBFloat16 && routerWeight.is_contiguous()
+            && routerWeight.dim() == 2 && routerWeight.size(0) == tk::kNumExperts
+            && routerWeight.size(1) == tk::kHiddenDim,
+        "[glm_moe_dsa_draft_live_row_hop] router_weight must be a contiguous CUDA bf16 tensor of shape [",
+        tk::kNumExperts, ", ", tk::kHiddenDim, "]");
+    TORCH_CHECK(input.get_device() == ucBuffer.get_device() && input.get_device() == flags.get_device()
+            && input.get_device() == rowIds.get_device(),
+        "[glm_moe_dsa_draft_live_row_hop] all tensors must be on the same device");
+
+    c10::cuda::CUDAGuard const deviceGuard(input.device());
+    auto const rowsOpts = input.options();
+    at::Tensor normOut = at::empty({numRows, tk::kHiddenDim}, rowsOpts);
+    at::Tensor residualOut = at::empty({numRows, tk::kHiddenDim}, rowsOpts);
+    at::Tensor routerLogits = at::empty({numRows, tk::kNumExperts}, rowsOpts.dtype(at::kFloat));
+
+    tk::LiveRowHopParams params;
+    params.input = input.const_data_ptr();
+    params.residualIn = residual.const_data_ptr();
+    params.rowIds = rowIds.const_data_ptr<int64_t>();
+    params.gamma = gamma.const_data_ptr();
+    params.routerWeight = routerWeight.const_data_ptr();
+    params.normOut = normOut.mutable_data_ptr();
+    params.residualOut = residualOut.mutable_data_ptr();
+    params.routerLogits = routerLogits.mutable_data_ptr<float>();
+    params.ucBuffer = ucBuffer.mutable_data_ptr();
+    params.mcBuffer = reinterpret_cast<void*>(mcPtr);
+    params.flags = reinterpret_cast<uint32_t*>(flags.mutable_data_ptr<int32_t>());
+    params.numRows = static_cast<int>(numRows);
+    params.numInputRows = static_cast<int>(numInputRows);
+    params.rank = static_cast<int>(rank);
+    params.eps = static_cast<float>(eps);
+    params.enablePdl = pdlEnabled();
+    params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
+
+    cudaError_t const status = tk::invokeDraftLiveRowHop(params);
+    TORCH_CHECK(
+        status == cudaSuccess, "[glm_moe_dsa_draft_live_row_hop] kernel launch failed: ", cudaGetErrorString(status));
+    return {normOut, residualOut, routerLogits};
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_draft_live_row_hop(Tensor input, Tensor residual, Tensor row_ids, Tensor gamma, float eps, "
+        "Tensor router_weight, Tensor uc_buffer, int mc_ptr, Tensor flags, int rank, int nranks) -> Tensor[]");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_draft_live_row_hop", &tensorrt_llm::torch_ext::glm_moe_dsa::draftLiveRowHop);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp
new file mode 100644
index 000000000000..15d02bced53a
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp
@@ -0,0 +1,246 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/library.h>
+
+#include <cuda.h>
+#include <cudaTypedefs.h>
+
+#include <cstdint>
+#include <initializer_list>
+#include <map>
+#include <mutex>
+#include <optional>
+#include <tuple>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::draft_routed_moe;
+
+#ifndef GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE
+#define GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE 0
+#endif
+//! Probe builds compile the kernels with their event record (see draftRoutedMoe.cu).
+constexpr bool kTimelineBuild = GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE != 0;
+
+void checkTensor(at::Tensor const& t, char const* name, at::ScalarType dtype, std::initializer_list<int64_t> shape,
+    c10::DeviceIndex device)
+{
+    TORCH_CHECK(t.is_cuda() && t.get_device() == device, "[glm_moe_dsa_draft_routed_moe] ", name,
+        " must be on the activation's device");
+    TORCH_CHECK(t.scalar_type() == dtype, "[glm_moe_dsa_draft_routed_moe] ", name, " has dtype ", t.scalar_type());
+    TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_draft_routed_moe] ", name, " must be contiguous");
+    TORCH_CHECK(t.sizes().vec() == std::vector<int64_t>(shape), "[glm_moe_dsa_draft_routed_moe] ", name, " has shape ",
+        t.sizes());
+}
+
+//! ``cuTensorMapEncodeTiled`` through the runtime's driver entry point (no link-time driver dependency).
+PFN_cuTensorMapEncodeTiled_v12000 encodeTiledFn()
+{
+    static PFN_cuTensorMapEncodeTiled_v12000 fn = []
+    {
+        void* ptr = nullptr;
+        cudaDriverEntryPointQueryResult query{};
+        cudaError_t const status
+            = cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &ptr, 12000, cudaEnableDefault, &query);
+        TORCH_CHECK(status == cudaSuccess && query == cudaDriverEntryPointSuccess && ptr != nullptr,
+            "[glm_moe_dsa_draft_routed_moe] cuTensorMapEncodeTiled is not available");
+        return reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(ptr);
+    }();
+    return fn;
+}
+
+//! Tensor map of a BlockMajorK weight buffer viewed as [slabs][slabRows][64] BF16 (a slab: one K chunk of one
+//! expert) with boxes of 64 x ``boxRows`` x ``boxSlabs`` and the 128-byte swizzle. Cached per (device, address,
+//! shape, box): the expert buffers are fixed model parameters.
+CUtensorMap const& weightMap(
+    int device, void const* base, uint64_t slabRows, uint64_t slabs, uint32_t boxRows, uint32_t boxSlabs)
+{
+    static std::mutex mutex;
+    static std::map<std::tuple<int, void const*, uint64_t, uint64_t, uint32_t, uint32_t>, CUtensorMap> cache;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto const key = std::make_tuple(device, base, slabRows, slabs, boxRows, boxSlabs);
+    auto it = cache.find(key);
+    if (it != cache.end())
+    {
+        return it->second;
+    }
+    CUtensorMap map{};
+    cuuint64_t const rowBytes = static_cast<cuuint64_t>(tk::kKChunk) * 2;
+    cuuint64_t const dims[3] = {static_cast<cuuint64_t>(tk::kKChunk), slabRows, slabs};
+    cuuint64_t const strides[2] = {rowBytes, rowBytes * slabRows};
+    cuuint32_t const box[3] = {static_cast<cuuint32_t>(tk::kKChunk), boxRows, boxSlabs};
+    cuuint32_t const elementStrides[3] = {1, 1, 1};
+    CUresult const result = encodeTiledFn()(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast<void*>(base), dims,
+        strides, box, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+        CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
+    TORCH_CHECK(result == CUDA_SUCCESS,
+        "[glm_moe_dsa_draft_routed_moe] cuTensorMapEncodeTiled failed: ", static_cast<int>(result));
+    return cache.emplace(key, map).first->second;
+}
+
+} // namespace
+
+//! Routed experts of the MTP draft layer (see kernels/glmMoeDsa/draftRoutedMoe.h): native DeepSeek-V3
+//! routing of ``router_logits`` with ``routing_bias``, FC1 + SwiGLU + FC2 + finalize on the FlashInfer TRT-LLM-Gen
+//! BF16 layouts of ``w3_w1`` / ``w2``. Returns the finalized routed output ``[R, 6144]`` (BF16). ``h_scratch`` and
+//! ``route_scratch`` are written (declared mutated in the schema) and reused by every call; calls sharing them must
+//! be ordered on one stream. ``debug_fc2`` (optional, ``[R, 8, 6144]`` BF16)
+//! receives the FC2 row of every (token, slot). ``timeline`` (optional, int64 ``[2, kTimelineCtas, kTimelineSlots]``)
+//! receives the per-CTA event record of the TMA-fed kernels; only probe builds
+//! (``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1``) accept it.
+at::Tensor draftRoutedMoe(at::Tensor const& x, at::Tensor const& routerLogits, at::Tensor const& routingBias,
+    at::Tensor const& w3w1, at::Tensor const& w2, at::Tensor const& hScratch, at::Tensor const& routeScratch,
+    double routedScaling, int64_t fc1Ctas, int64_t fc1Warps, int64_t fc2Ctas, int64_t fc2Warps, int64_t flags,
+    std::optional<at::Tensor> const& debugFc2, std::optional<at::Tensor> const& timeline)
+{
+    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2
+            && x.size(1) == tk::kHidden,
+        "[glm_moe_dsa_draft_routed_moe] x must be a contiguous CUDA bf16 tensor [R, ", tk::kHidden, "]");
+    int64_t const rows = x.size(0);
+    TORCH_CHECK(rows >= 1 && rows <= tk::kMaxRows, "[glm_moe_dsa_draft_routed_moe] rows ", rows, " outside [1, ",
+        tk::kMaxRows, "]");
+    TORCH_CHECK(reinterpret_cast<uintptr_t>(x.const_data_ptr()) % 16 == 0,
+        "[glm_moe_dsa_draft_routed_moe] x must be 16-byte aligned");
+    auto const device = x.get_device();
+    checkTensor(routerLogits, "router_logits", at::kFloat, {rows, tk::kNumExperts}, device);
+    checkTensor(routingBias, "routing_bias", at::kBFloat16, {tk::kNumExperts}, device);
+    TORCH_CHECK(reinterpret_cast<uintptr_t>(routerLogits.const_data_ptr()) % 16 == 0
+            && reinterpret_cast<uintptr_t>(routingBias.const_data_ptr()) % 16 == 0,
+        "[glm_moe_dsa_draft_routed_moe] router_logits and routing_bias must be 16-byte aligned");
+    checkTensor(w3w1, "w3_w1", at::kBFloat16, {tk::kNumExperts, tk::kFc1Chunks, tk::kFc1Rows, tk::kKChunk}, device);
+    checkTensor(w2, "w2", at::kBFloat16, {tk::kNumExperts, tk::kFc2Chunks, tk::kHidden, tk::kKChunk}, device);
+    TORCH_CHECK(hScratch.is_cuda() && hScratch.get_device() == device && hScratch.scalar_type() == at::kBFloat16
+            && hScratch.is_contiguous() && hScratch.numel() >= tk::kScratchHElems,
+        "[glm_moe_dsa_draft_routed_moe] h_scratch must be a contiguous CUDA bf16 tensor with at least ",
+        tk::kScratchHElems, " elements");
+    TORCH_CHECK(routeScratch.is_cuda() && routeScratch.get_device() == device && routeScratch.scalar_type() == at::kInt
+            && routeScratch.is_contiguous() && routeScratch.numel() >= tk::kRouteWords,
+        "[glm_moe_dsa_draft_routed_moe] route_scratch must be a contiguous CUDA int32 tensor with at least ",
+        tk::kRouteWords, " elements");
+    TORCH_CHECK(
+        fc1Warps == 1 || fc1Warps == 2 || fc1Warps == 4, "[glm_moe_dsa_draft_routed_moe] fc1_warps must be 1, 2 or 4");
+    TORCH_CHECK(fc2Warps == 1 || fc2Warps == 2 || fc2Warps == 4 || (fc2Warps == 3 && (flags & tk::kFlagTma)),
+        "[glm_moe_dsa_draft_routed_moe] fc2_warps must be 1, 2 or 4 (or 3 with the TMA flag)");
+    TORCH_CHECK(fc1Ctas >= 1 && fc1Ctas <= 1024 && fc2Ctas >= 1 && fc2Ctas <= 1024,
+        "[glm_moe_dsa_draft_routed_moe] grid sizes");
+    if (debugFc2.has_value())
+    {
+        checkTensor(debugFc2.value(), "debug_fc2", at::kBFloat16, {rows, tk::kTopK, tk::kHidden}, device);
+    }
+    if (timeline.has_value())
+    {
+        TORCH_CHECK(kTimelineBuild,
+            "[glm_moe_dsa_draft_routed_moe] timeline needs a build with GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1");
+        checkTensor(timeline.value(), "timeline", at::kLong, {2, tk::kTimelineCtas, tk::kTimelineSlots}, device);
+    }
+
+    c10::cuda::CUDAGuard const deviceGuard(x.device());
+    at::Tensor out = at::empty({rows, tk::kHidden}, x.options());
+
+    tk::DraftRoutedMoeParams params{};
+    params.logits = routerLogits.const_data_ptr<float>();
+    params.bias = reinterpret_cast<__nv_bfloat16 const*>(routingBias.const_data_ptr());
+    params.x = reinterpret_cast<__nv_bfloat16 const*>(x.const_data_ptr());
+    params.w13 = reinterpret_cast<__nv_bfloat16 const*>(w3w1.const_data_ptr());
+    params.w2 = reinterpret_cast<__nv_bfloat16 const*>(w2.const_data_ptr());
+    params.out = reinterpret_cast<__nv_bfloat16*>(out.mutable_data_ptr());
+    params.h = reinterpret_cast<__nv_bfloat16*>(hScratch.mutable_data_ptr());
+    params.route = routeScratch.mutable_data_ptr<int32_t>();
+    params.debugFc2 = debugFc2.has_value() ? reinterpret_cast<__nv_bfloat16*>(debugFc2->mutable_data_ptr()) : nullptr;
+    params.timeline = timeline.has_value() ? timeline->mutable_data_ptr<int64_t>() : nullptr;
+    params.rows = static_cast<int>(rows);
+    params.routedScaling = static_cast<float>(routedScaling);
+    params.fc1Ctas = static_cast<int>(fc1Ctas);
+    params.fc1Warps = static_cast<int>(fc1Warps);
+    params.fc2Ctas = static_cast<int>(fc2Ctas);
+    params.fc2Warps = static_cast<int>(fc2Warps);
+    params.flags = static_cast<int>(flags);
+
+    tk::DraftRoutedMoeTensorMaps maps{};
+    bool const tma = (flags & tk::kFlagTma) != 0;
+    if (tma)
+    {
+        int const f = static_cast<int>(flags);
+        int const w1 = static_cast<int>(fc1Warps);
+        maps.w13 = weightMap(device, w3w1.const_data_ptr(), tk::kFc1Rows,
+            static_cast<uint64_t>(tk::kNumExperts) * tk::kFc1Chunks, static_cast<uint32_t>(tk::fc1TmaBoxRows(f, w1)),
+            static_cast<uint32_t>(tk::fc1TmaBoxSlabs(f, w1)));
+        maps.w2 = weightMap(device, w2.const_data_ptr(), tk::kHidden,
+            static_cast<uint64_t>(tk::kNumExperts) * tk::kFc2Chunks, static_cast<uint32_t>(16 * fc2Warps),
+            static_cast<uint32_t>(tk::kFc2Chunks));
+        // The triples path of the FC1 kernel (4 MMA warps) reads 64-row boxes; otherwise the map is not read.
+        maps.w13Triple = ((f & tk::kFlagFc1Triples) && w1 == 4) ? weightMap(device, w3w1.const_data_ptr(), tk::kFc1Rows,
+                             static_cast<uint64_t>(tk::kNumExperts) * tk::kFc1Chunks,
+                             static_cast<uint32_t>(tk::kTripleRows), static_cast<uint32_t>(tk::kTripleStageChunks))
+                                                                : maps.w13;
+        // kFlagStreamX: the call's activation as [token][K chunk][64] (boxes of 8 tokens; rows beyond the call's are
+        // zero-filled and only feed discarded MMA columns). Encoded per call (``x`` changes between calls).
+        maps.x = maps.w13;
+        if (f & tk::kFlagStreamX)
+        {
+            cuuint64_t const dims[3] = {static_cast<cuuint64_t>(tk::kKChunk), static_cast<cuuint64_t>(tk::kFc1Chunks),
+                static_cast<cuuint64_t>(rows)};
+            cuuint64_t const strides[2]
+                = {static_cast<cuuint64_t>(tk::kKChunk) * 2, static_cast<cuuint64_t>(tk::kHidden) * 2};
+            cuuint32_t const box[3] = {static_cast<cuuint32_t>(tk::kKChunk), static_cast<cuuint32_t>(tk::kXBoxChunks),
+                static_cast<cuuint32_t>(tk::kMaxRows)};
+            cuuint32_t const elementStrides[3] = {1, 1, 1};
+            CUresult const result
+                = encodeTiledFn()(&maps.x, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast<void*>(x.const_data_ptr()),
+                    dims, strides, box, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                    CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
+            TORCH_CHECK(result == CUDA_SUCCESS,
+                "[glm_moe_dsa_draft_routed_moe] cuTensorMapEncodeTiled (x) failed: ", static_cast<int>(result));
+        }
+    }
+    cudaError_t const status
+        = tk::invokeDraftRoutedMoe(params, tma ? &maps : nullptr, at::cuda::getCurrentCUDAStream(device));
+    TORCH_CHECK(
+        status == cudaSuccess, "[glm_moe_dsa_draft_routed_moe] kernel launch failed: ", cudaGetErrorString(status));
+    return out;
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_draft_routed_moe(Tensor x, Tensor router_logits, Tensor routing_bias, Tensor w3_w1, Tensor w2, "
+        "Tensor(a!) h_scratch, Tensor(b!) route_scratch, float routed_scaling, int fc1_ctas, int fc1_warps, "
+        "int fc2_ctas, int fc2_warps, int flags, Tensor(c!)? debug_fc2=None, Tensor(d!)? timeline=None) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_draft_routed_moe", &tensorrt_llm::torch_ext::glm_moe_dsa::draftRoutedMoe);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp
new file mode 100644
index 000000000000..62c53823f6ba
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp
@@ -0,0 +1,168 @@
+/*
+ * 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.
+ */
+
+// Torch ops of the draft tail (greedy draft-token selection over a
+// vocab-sharded LM head). See kernels/glmMoeDsa/draftTail.h.
+
+#include "tensorrt_llm/kernels/glmMoeDsa/draftTail.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include <optional>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::draft_tail;
+
+void fillCommon(tk::DraftTailParams& p, int64_t rows, int64_t vocabLocal, at::Tensor& tokens, at::Tensor& scratch,
+    std::optional<at::Tensor> const& keysOut, int64_t vocabOffset, int64_t tpRank, int64_t tpSize, int64_t mode,
+    int64_t mcPtr, int64_t ucPtr, std::optional<at::Tensor> const& peerPtrs, int64_t timeoutCycles, c10::Device device)
+{
+    TORCH_CHECK(rows >= 1 && rows <= tk::kMaxRows, "glm_moe_dsa_draft_tail: rows must be in [1, ", tk::kMaxRows, "]");
+    TORCH_CHECK(tokens.device() == device && tokens.scalar_type() == at::kInt && tokens.is_contiguous()
+            && tokens.numel() >= rows,
+        "glm_moe_dsa_draft_tail: tokens must be a contiguous int32 tensor with at least rows elements");
+    TORCH_CHECK(scratch.device() == device && scratch.scalar_type() == at::kLong && scratch.is_contiguous()
+            && scratch.numel() >= tk::kScratchWords,
+        "glm_moe_dsa_draft_tail: scratch must be a contiguous int64 tensor with ", tk::kScratchWords, " elements");
+    TORCH_CHECK(vocabOffset >= 0 && vocabOffset + vocabLocal <= 0xFFFFFFFFll,
+        "glm_moe_dsa_draft_tail: vocab range out of bounds");
+    TORCH_CHECK(tpSize >= 1 && tpSize <= tk::kMaxRanks && tpRank >= 0 && tpRank < tpSize,
+        "glm_moe_dsa_draft_tail: bad tp rank");
+    TORCH_CHECK(mode >= 0 && mode <= 4, "glm_moe_dsa_draft_tail: bad exchange mode");
+    p.rows = static_cast<int>(rows);
+    p.vocabLocal = static_cast<int>(vocabLocal);
+    p.vocabOffset = vocabOffset;
+    p.tokens = tokens.data_ptr<int32_t>();
+    p.scratch = reinterpret_cast<uint64_t*>(scratch.data_ptr<int64_t>());
+    if (keysOut.has_value())
+    {
+        TORCH_CHECK(keysOut->device() == device && keysOut->scalar_type() == at::kLong && keysOut->is_contiguous()
+                && keysOut->numel() >= rows,
+            "glm_moe_dsa_draft_tail: keys_out must be a contiguous int64 tensor with at least rows elements");
+        p.keysOut = reinterpret_cast<uint64_t*>(keysOut->data_ptr<int64_t>());
+    }
+    p.tpRank = static_cast<int>(tpRank);
+    p.tpSize = static_cast<int>(tpSize);
+    p.mode = static_cast<tk::ExchangeMode>(mode);
+    p.timeoutCycles = timeoutCycles;
+    if (p.mode != tk::ExchangeMode::kNone && tpSize > 1)
+    {
+        TORCH_CHECK(ucPtr != 0 && ucPtr % 16 == 0, "glm_moe_dsa_draft_tail: exchange needs the local unicast buffer");
+        p.ucSlotsLocal = reinterpret_cast<uint64_t*>(ucPtr);
+        if (p.mode == tk::ExchangeMode::kMulticast || p.mode == tk::ExchangeMode::kMulticastDirect)
+        {
+            TORCH_CHECK(
+                mcPtr != 0 && mcPtr % 16 == 0, "glm_moe_dsa_draft_tail: multicast exchange needs the multicast buffer");
+            p.mcSlots = reinterpret_cast<uint64_t*>(mcPtr);
+        }
+        else
+        {
+            TORCH_CHECK(peerPtrs.has_value() && peerPtrs->device() == device && peerPtrs->scalar_type() == at::kLong
+                    && peerPtrs->is_contiguous() && peerPtrs->numel() >= tpSize,
+                "glm_moe_dsa_draft_tail: unicast exchange needs a device int64 tensor of tp_size peer addresses");
+            p.ucSlotsPeers = reinterpret_cast<uint64_t* const*>(peerPtrs->data_ptr<int64_t>());
+        }
+    }
+}
+
+} // namespace
+
+void draftTail(at::Tensor const& hidden, at::Tensor const& weight, at::Tensor& tokens, at::Tensor& scratch,
+    std::optional<at::Tensor> const& keysOut, int64_t vocabOffset, int64_t tpRank, int64_t tpSize, int64_t mode,
+    int64_t mcPtr, int64_t ucPtr, std::optional<at::Tensor> const& peerPtrs, int64_t variant, int64_t numCtas,
+    int64_t keepPermille, bool pdl, int64_t timeoutCycles, std::optional<at::Tensor> const& logitsOut)
+{
+    TORCH_CHECK(hidden.is_cuda() && hidden.dim() == 2 && hidden.scalar_type() == at::kBFloat16,
+        "glm_moe_dsa_draft_tail: hidden must be a 2D BF16 CUDA tensor");
+    TORCH_CHECK(weight.device() == hidden.device() && weight.dim() == 2 && weight.scalar_type() == at::kBFloat16
+            && weight.is_contiguous(),
+        "glm_moe_dsa_draft_tail: weight must be a contiguous 2D BF16 tensor on the hidden device");
+    TORCH_CHECK(hidden.size(1) == 6144 && weight.size(1) == 6144, "glm_moe_dsa_draft_tail: hidden size must be 6144");
+    TORCH_CHECK(hidden.stride(1) == 1 && hidden.stride(0) % 8 == 0
+            && reinterpret_cast<uintptr_t>(hidden.data_ptr()) % 16 == 0
+            && reinterpret_cast<uintptr_t>(weight.data_ptr()) % 32 == 0,
+        "glm_moe_dsa_draft_tail: hidden rows must be 16-byte aligned and the weight 32-byte aligned");
+    tk::DraftTailParams p;
+    fillCommon(p, hidden.size(0), weight.size(0), tokens, scratch, keysOut, vocabOffset, tpRank, tpSize, mode, mcPtr,
+        ucPtr, peerPtrs, timeoutCycles, hidden.device());
+    p.hidden = reinterpret_cast<__nv_bfloat16 const*>(hidden.data_ptr());
+    p.hiddenStride = hidden.stride(0);
+    p.weight = reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr());
+    p.hiddenDim = static_cast<int>(hidden.size(1));
+    TORCH_CHECK(
+        keepPermille >= 0 && keepPermille <= 1000, "glm_moe_dsa_draft_tail: keep_permille must be in [0, 1000]");
+    p.keepPermille = static_cast<int>(keepPermille);
+    if (logitsOut.has_value())
+    {
+        TORCH_CHECK(logitsOut->device() == hidden.device() && logitsOut->scalar_type() == at::kBFloat16
+                && logitsOut->dim() == 2 && logitsOut->size(0) >= hidden.size(0)
+                && logitsOut->size(1) == weight.size(0) && logitsOut->stride(1) == 1,
+            "glm_moe_dsa_draft_tail: logits_out must be a BF16 [rows, vocab_local] tensor with unit last stride");
+        p.logitsOut = reinterpret_cast<__nv_bfloat16*>(logitsOut->data_ptr());
+        p.logitsOutStride = logitsOut->stride(0);
+    }
+    tk::FusedLaunchConfig config;
+    config.numCtas = static_cast<int>(numCtas);
+    config.variant = static_cast<int>(variant);
+    config.pdl = pdl;
+    tk::invokeDraftTailFused(p, config, at::cuda::getCurrentCUDAStream(hidden.get_device()));
+}
+
+void draftTailFromLogits(at::Tensor const& logits, at::Tensor& tokens, at::Tensor& scratch,
+    std::optional<at::Tensor> const& keysOut, int64_t vocabOffset, int64_t tpRank, int64_t tpSize, int64_t mode,
+    int64_t mcPtr, int64_t ucPtr, std::optional<at::Tensor> const& peerPtrs, bool pdl, int64_t timeoutCycles)
+{
+    TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.scalar_type() == at::kBFloat16 && logits.stride(1) == 1,
+        "glm_moe_dsa_draft_tail: logits must be a 2D BF16 CUDA tensor with unit last stride");
+    tk::DraftTailParams p;
+    fillCommon(p, logits.size(0), logits.size(1), tokens, scratch, keysOut, vocabOffset, tpRank, tpSize, mode, mcPtr,
+        ucPtr, peerPtrs, timeoutCycles, logits.device());
+    p.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr());
+    p.logitsStride = logits.stride(0);
+    tk::invokeDraftTailFromLogits(p, pdl, at::cuda::getCurrentCUDAStream(logits.get_device()));
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_draft_tail(Tensor hidden, Tensor weight, Tensor(a!) tokens, Tensor(b!) scratch, "
+        "Tensor(c!)? keys_out, int vocab_offset, int tp_rank, int tp_size, int mode, int mc_ptr, int uc_ptr, "
+        "Tensor? peer_ptrs, int variant, int num_ctas, int keep_permille, bool pdl, int timeout_cycles, "
+        "Tensor(d!)? logits_out=None) -> ()");
+    m.def(
+        "glm_moe_dsa_draft_tail_from_logits(Tensor logits, Tensor(a!) tokens, Tensor(b!) scratch, Tensor(c!)? "
+        "keys_out, "
+        "int vocab_offset, int tp_rank, int tp_size, int mode, int mc_ptr, int uc_ptr, Tensor? peer_ptrs, bool pdl, "
+        "int timeout_cycles) -> ()");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_draft_tail", &tensorrt_llm::torch_ext::glm_moe_dsa::draftTail);
+    m.impl("glm_moe_dsa_draft_tail_from_logits", &tensorrt_llm::torch_ext::glm_moe_dsa::draftTailFromLogits);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp
new file mode 100644
index 000000000000..d51ca99cf61d
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp
@@ -0,0 +1,242 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/envUtils.h"
+#include "tensorrt_llm/common/quantization.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h"
+#include "tensorrt_llm/kernels/kvCacheUtils.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include <climits>
+#include <cmath>
+#include <cstdint>
+#include <optional>
+
+namespace tk = tensorrt_llm::kernels;
+namespace tc = tensorrt_llm::common;
+namespace rope_kv_append = tensorrt_llm::kernels::glm_moe_dsa::rope_kv_append;
+namespace indexer_layer_front = tensorrt_llm::kernels::glm_moe_dsa::indexer_layer_front;
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+void checkCuda(torch::Tensor const& t, char const* name)
+{
+    TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+}
+
+//! Paged MLA latent cache of one layer: the KVBlockArray that buildPagedKvCacheBuffers (thop/attentionOp.cpp)
+//! builds for an FP8 MLA cache (kvFactor 1, one KV head, 8-bit elements, beam width 1, no FP4 scale pool).
+tk::KVBlockArray buildFp8MlaKvBlockArray(torch::Tensor const& kvCacheBlockOffsets,
+    torch::Tensor const& hostKvCachePoolPointers, torch::Tensor const& hostKvCachePoolMapping, int64_t layerIdx,
+    int32_t batchSize, int64_t tokensPerBlock, int64_t attentionWindowSize, int64_t seqOffset)
+{
+    TORCH_CHECK(hostKvCachePoolMapping.device().is_cpu() && hostKvCachePoolMapping.dim() == 2
+            && hostKvCachePoolMapping.size(1) >= 2 && hostKvCachePoolMapping.scalar_type() == torch::kInt32,
+        "host_kv_cache_pool_mapping must be a CPU int32 [layers, >=2] tensor");
+    TORCH_CHECK(layerIdx >= 0 && layerIdx < hostKvCachePoolMapping.size(0), "layer index out of the pool mapping");
+    auto const* mapping = hostKvCachePoolMapping.data_ptr<int32_t>();
+    int32_t const poolIndex = mapping[layerIdx * hostKvCachePoolMapping.stride(0)];
+    int32_t const layerIdxInCachePool
+        = mapping[layerIdx * hostKvCachePoolMapping.stride(0) + hostKvCachePoolMapping.stride(1)];
+
+    TORCH_CHECK(hostKvCachePoolPointers.device().is_cpu() && hostKvCachePoolPointers.dim() == 2
+            && hostKvCachePoolPointers.scalar_type() == torch::kInt64 && poolIndex >= 0
+            && poolIndex < hostKvCachePoolPointers.size(0) && hostKvCachePoolPointers.size(1) >= 2,
+        "host_kv_cache_pool_pointers must be a CPU int64 [pools, 2] tensor");
+    TORCH_CHECK(kvCacheBlockOffsets.is_cuda() && kvCacheBlockOffsets.scalar_type() == torch::kInt32
+            && kvCacheBlockOffsets.dim() >= 2 && poolIndex < kvCacheBlockOffsets.size(0)
+            && seqOffset < kvCacheBlockOffsets.size(1),
+        "kv_cache_block_offsets must be CUDA int32 [pools, seqs, 2, max_blocks]");
+
+    int64_t constexpr kCacheElemBits = 8;
+    int64_t const blockSize = tokensPerBlock * /*kv_head_num=*/1 * rope_kv_append::kLatentDim;
+    int64_t const bytesPerBlock = blockSize * kCacheElemBits / CHAR_BIT;
+    int64_t const intraPoolOffset = layerIdxInCachePool * /*kvFactor=*/1 * bytesPerBlock;
+    auto const sizePerToken = static_cast<int32_t>(rope_kv_append::kLatentDim * kCacheElemBits / 8);
+    auto const* pointers = hostKvCachePoolPointers.data_ptr<int64_t>();
+    auto const rowOffset = poolIndex * hostKvCachePoolPointers.stride(0);
+    void* primaryPoolPtr = reinterpret_cast<char*>(pointers[rowOffset]) + intraPoolOffset;
+    void* secondaryPoolPtr
+        = reinterpret_cast<char*>(pointers[rowOffset + hostKvCachePoolPointers.stride(1)]) + intraPoolOffset;
+    auto* blockOffsets
+        = reinterpret_cast<tk::KVBlockArray::DataType*>(static_cast<int32_t*>(kvCacheBlockOffsets.data_ptr())
+            + poolIndex * kvCacheBlockOffsets.stride(0) + seqOffset * kvCacheBlockOffsets.stride(1));
+    auto const maxBlocksPerSequence = static_cast<int32_t>(kvCacheBlockOffsets.size(-1));
+    auto const window = static_cast<int32_t>(attentionWindowSize);
+    return tk::KVBlockArray(batchSize, maxBlocksPerSequence, static_cast<int32_t>(tokensPerBlock), sizePerToken, window,
+        window, /*sinkTokenLen=*/0, /*canUseOneMoreBlock=*/false, primaryPoolPtr, secondaryPoolPtr, blockOffsets);
+}
+
+} // namespace
+
+//! Variant of trtllm::glm_moe_dsa_rope_kv_append: same arguments, checks and written bytes; with early_trigger each
+//! kernel thread triggers the programmatic dependents right after the global reads of its work item (otherwise at the
+//! end of the kernel, as in trtllm::glm_moe_dsa_rope_kv_append).
+void indexerLayerRopeKvAppend(torch::Tensor const& fused_q, torch::Tensor const& q_pe,
+    torch::Tensor const& latent_cache, torch::Tensor const& rotary_cos_sin, torch::Tensor const& cu_q_seqlens,
+    torch::Tensor const& cu_kv_seqlens, torch::Tensor const& fmha_scheduler_counter,
+    torch::Tensor const& mla_bmm1_scale, torch::Tensor const& mla_bmm2_scale, torch::Tensor const& quant_q_buffer,
+    torch::Tensor const& sequence_length, torch::Tensor const& host_context_lengths, int64_t num_contexts,
+    torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers,
+    torch::Tensor const& host_kv_cache_pool_mapping, std::optional<torch::Tensor> const& kv_cache_scale_orig_quant,
+    int64_t layer_idx, int64_t num_heads, int64_t tokens_per_block, int64_t attention_window_size, int64_t quant_mode,
+    double q_scaling, int64_t qk_nope_head_dim, int64_t qk_rope_head_dim, int64_t kv_lora_rank,
+    std::optional<torch::Tensor> const& idx_k_fp8, std::optional<torch::Tensor> const& idx_k_scale,
+    std::optional<torch::Tensor> const& idx_k_cache, std::optional<torch::Tensor> const& idx_slot_mapping_fp8,
+    std::optional<torch::Tensor> const& idx_slot_mapping_scale, int64_t idx_num_tokens, bool early_trigger)
+{
+    TORCH_CHECK(kv_lora_rank == rope_kv_append::kKvLoraRank && qk_rope_head_dim == rope_kv_append::kRopeDim,
+        "glm_moe_dsa_indexer_layer_rope_kv_append: supports kv_lora_rank 512 and qk_rope_head_dim 64 only");
+    auto const kvQuantMode = tc::QuantMode(static_cast<uint32_t>(quant_mode));
+    TORCH_CHECK(kvQuantMode.hasFp8KvCache() && !kvQuantMode.hasFp4KvCache(),
+        "glm_moe_dsa_indexer_layer_rope_kv_append: needs an FP8 MLA cache");
+    checkCuda(fused_q, "fused_q");
+    checkCuda(q_pe, "q_pe");
+    checkCuda(latent_cache, "latent_cache");
+    TORCH_CHECK(fused_q.scalar_type() == torch::kBFloat16 && q_pe.scalar_type() == torch::kBFloat16
+            && latent_cache.scalar_type() == torch::kBFloat16,
+        "glm_moe_dsa_indexer_layer_rope_kv_append: needs BF16 fused_q, q_pe and latent_cache");
+    TORCH_CHECK(rotary_cos_sin.scalar_type() == torch::kFloat32 && rotary_cos_sin.is_contiguous(),
+        "rotary_cos_sin must be packed float32");
+    int64_t const cosSinRows = rotary_cos_sin.numel() / (2 * rope_kv_append::kRopeDim);
+    TORCH_CHECK(cosSinRows > 0 && cosSinRows <= INT32_MAX, "rotary_cos_sin has no complete position row");
+    TORCH_CHECK(q_pe.dim() == 3 && q_pe.stride(2) == 1, "q_pe must be [tokens, heads, rope] with unit inner stride");
+    TORCH_CHECK(latent_cache.dim() == 2 && latent_cache.size(1) == rope_kv_append::kLatentDim
+            && latent_cache.stride(1) == 1 && latent_cache.stride(0) == rope_kv_append::kLatentDim,
+        "latent_cache must be packed [tokens, 576]");
+    TORCH_CHECK(fused_q.is_contiguous() && quant_q_buffer.is_contiguous(), "fused_q and quant_q_buffer must be packed");
+    TORCH_CHECK(host_kv_cache_pool_mapping.dim() >= 1, "host_kv_cache_pool_mapping is required");
+
+    int32_t const numTokens = static_cast<int32_t>(latent_cache.size(0));
+    int32_t const numSeqs = static_cast<int32_t>(host_context_lengths.size(0));
+    int32_t const numGenerations = numSeqs - static_cast<int32_t>(num_contexts);
+    TORCH_CHECK(numGenerations > 0, "glm_moe_dsa_indexer_layer_rope_kv_append: needs at least one generation request");
+    TORCH_CHECK(numTokens % numGenerations == 0, "MLA can only support input sequences with the same sequence length.");
+    TORCH_CHECK(fused_q.numel() >= static_cast<int64_t>(numTokens) * num_heads * rope_kv_append::kLatentDim,
+        "fused_q is smaller than [tokens, heads, 576]");
+    TORCH_CHECK(quant_q_buffer.numel() * quant_q_buffer.element_size()
+            >= static_cast<int64_t>(numTokens) * num_heads * rope_kv_append::kLatentDim,
+        "quant_q_buffer is smaller than [tokens, heads, 576] bytes");
+    TORCH_CHECK(mla_bmm1_scale.numel() >= 2 && mla_bmm2_scale.numel() >= 1, "bmm scale buffers too small");
+    TORCH_CHECK(cu_q_seqlens.numel() >= numGenerations + 1 && cu_kv_seqlens.numel() >= numGenerations + 1,
+        "cu_seqlens buffers too small");
+
+    auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device());
+    rope_kv_append::RopeAppendParams params{};
+    params.fusedQ = reinterpret_cast<__nv_bfloat16 const*>(fused_q.data_ptr());
+    params.qPe = reinterpret_cast<__nv_bfloat16 const*>(q_pe.data_ptr());
+    params.qPeLd = static_cast<int32_t>(q_pe.stride(1));
+    params.qPeStride = static_cast<int32_t>(q_pe.stride(0));
+    params.quantQ = reinterpret_cast<uint8_t*>(quant_q_buffer.data_ptr());
+    params.quantScaleQ = nullptr;
+    params.latent = reinterpret_cast<__nv_bfloat16 const*>(latent_cache.data_ptr());
+    params.cosSin = reinterpret_cast<float2 const*>(rotary_cos_sin.data_ptr());
+    params.cosSinRows = static_cast<int32_t>(cosSinRows);
+    params.kvLens = sequence_length.slice(0, num_contexts).data_ptr<int32_t>();
+    params.quantScaleKv
+        = kv_cache_scale_orig_quant.has_value() ? kv_cache_scale_orig_quant->data_ptr<float>() : nullptr;
+    params.kvCache = buildFp8MlaKvBlockArray(kv_cache_block_offsets, host_kv_cache_pool_pointers,
+        host_kv_cache_pool_mapping, layer_idx, numGenerations, tokens_per_block, attention_window_size, num_contexts);
+    params.cuQSeqlens = cu_q_seqlens.data_ptr<int32_t>();
+    params.cuKvSeqlens = cu_kv_seqlens.data_ptr<int32_t>();
+    params.fmhaTileCounter = reinterpret_cast<uint32_t*>(fmha_scheduler_counter.data_ptr());
+    params.bmm1Scale = mla_bmm1_scale.data_ptr<float>();
+    params.bmm2Scale = mla_bmm2_scale.data_ptr<float>();
+    params.dequantScaleQ = nullptr;
+    params.dequantScaleKv = nullptr;
+    params.quantScaleO = nullptr;
+    // Same expression as the native op so the folded FMHA scales match bit for bit.
+    params.hostBmm1Scale = 1.F / (q_scaling * std::sqrt(static_cast<float>(qk_nope_head_dim + qk_rope_head_dim)));
+    params.numTokens = numTokens;
+    params.seqLen = numTokens / numGenerations;
+    params.numHeads = static_cast<int32_t>(num_heads);
+
+    if (idx_num_tokens > 0)
+    {
+        TORCH_CHECK(idx_k_fp8.has_value() && idx_k_scale.has_value() && idx_k_cache.has_value()
+                && idx_slot_mapping_fp8.has_value() && idx_slot_mapping_scale.has_value(),
+            "indexer scatter needs k_fp8, k_scale, k_cache and both slot mappings");
+        auto const& kFp8 = idx_k_fp8.value();
+        auto const& kScale = idx_k_scale.value();
+        auto const& kCache = idx_k_cache.value();
+        auto const& slotFp8 = idx_slot_mapping_fp8.value();
+        auto const& slotScale = idx_slot_mapping_scale.value();
+        checkCuda(kFp8, "k_fp8");
+        checkCuda(kScale, "k_scale");
+        checkCuda(kCache, "k_cache");
+        TORCH_CHECK(kFp8.dim() == 2 && kFp8.element_size() == 1 && kFp8.is_contiguous(), "k_fp8 must be packed bytes");
+        TORCH_CHECK(kScale.dim() == 2 && kScale.element_size() == 4 && kScale.is_contiguous()
+                && kScale.size(1) * kScale.element_size() == 4,
+            "k_scale must be one packed 4-byte word per token");
+        TORCH_CHECK(kCache.dim() == 4 && kCache.size(2) == 1, "k_cache must be [blocks, block_size, 1, bytes]");
+        TORCH_CHECK(slotFp8.scalar_type() == torch::kInt64 && slotScale.scalar_type() == torch::kInt64
+                && slotFp8.is_contiguous() && slotScale.is_contiguous(),
+            "slot mappings must be packed int64");
+        int32_t const headBytes = static_cast<int32_t>(kFp8.size(1));
+        TORCH_CHECK(headBytes == 128 || headBytes == 64, "indexer k head bytes must be 128 or 64");
+        TORCH_CHECK(kFp8.size(0) >= idx_num_tokens && kScale.size(0) >= idx_num_tokens
+                && slotFp8.numel() >= idx_num_tokens && slotScale.numel() >= idx_num_tokens,
+            "indexer inputs are shorter than idx_num_tokens");
+        params.idxK = reinterpret_cast<uint8_t const*>(kFp8.data_ptr());
+        params.idxScale = reinterpret_cast<uint8_t const*>(kScale.data_ptr());
+        params.idxCache = reinterpret_cast<uint8_t*>(kCache.data_ptr());
+        params.idxSlotData = slotFp8.data_ptr<int64_t>();
+        params.idxSlotScale = slotScale.data_ptr<int64_t>();
+        params.idxNumTokens = static_cast<int32_t>(idx_num_tokens);
+        params.idxHeadBytes = headBytes;
+        for (int32_t d = 0; d < 4; ++d)
+        {
+            params.idxDims[d] = static_cast<int32_t>(kCache.size(d));
+            params.idxStrides[d] = kCache.stride(d);
+        }
+    }
+
+    indexer_layer_front::invokeIndexerLayerRopeKvAppend(params, tc::getEnvEnablePDL(), early_trigger, stream);
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_indexer_layer_rope_kv_append(Tensor fused_q, Tensor q_pe, Tensor latent_cache, Tensor "
+        "rotary_cos_sin, "
+        "Tensor(a!) cu_q_seqlens, Tensor(b!) cu_kv_seqlens, Tensor(c!) fmha_scheduler_counter, "
+        "Tensor(d!) mla_bmm1_scale, Tensor(e!) mla_bmm2_scale, Tensor(f!) quant_q_buffer, Tensor sequence_length, "
+        "Tensor host_context_lengths, int num_contexts, Tensor kv_cache_block_offsets, "
+        "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_cache_scale_orig_quant, "
+        "int layer_idx, int num_heads, int tokens_per_block, int attention_window_size, int quant_mode, "
+        "float q_scaling, int qk_nope_head_dim, int qk_rope_head_dim, int kv_lora_rank, Tensor? idx_k_fp8, "
+        "Tensor? idx_k_scale, Tensor(g!)? idx_k_cache, Tensor? idx_slot_mapping_fp8, Tensor? idx_slot_mapping_scale, "
+        "int idx_num_tokens, bool early_trigger) -> ()");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_indexer_layer_rope_kv_append", &tensorrt_llm::torch_ext::glm_moe_dsa::indexerLayerRopeKvAppend);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp
new file mode 100644
index 000000000000..7d94a64871ff
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp
@@ -0,0 +1,150 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/library.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <optional>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::attention_output_into_hop;
+
+// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0.
+bool pdlEnabled()
+{
+    static bool const enabled = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && std::strcmp(env, "0") == 0);
+    }();
+    return enabled;
+}
+
+void checkRows(at::Tensor const& t, char const* name)
+{
+    TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_post_attention_hop_from_partials] ",
+        name, " must be a CUDA bf16 tensor");
+    TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_post_attention_hop_from_partials] ", name, " must be contiguous");
+    TORCH_CHECK(t.dim() >= 1 && t.size(-1) == tk::kHiddenDim, "[glm_moe_dsa_post_attention_hop_from_partials] ", name,
+        " must have hidden size ", tk::kHiddenDim);
+}
+
+} // namespace
+
+//! Post-attention hop whose rank-local partials were already published into the hop's transport by the attention
+//! output tail. ``residual`` gives the row count. Returns ``[normed, residual_out, router_logits]``; ``router_logits``
+//! is an empty tensor when ``router_weight`` is not given. See kernels/glmMoeDsa/postAttentionHopFromPartials.h.
+std::vector<at::Tensor> postAttentionHopFromPartials(at::Tensor const& residual, at::Tensor const& gamma, double eps,
+    std::optional<at::Tensor> const& routerWeight, at::Tensor const& ucBuffer, at::Tensor const& flags, int64_t rank,
+    int64_t nRanks, bool earlyPoll, int64_t routerLoad)
+{
+    checkRows(residual, "residual");
+    checkRows(gamma, "gamma");
+    TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_post_attention_hop_from_partials] gamma must be 1-D");
+    int64_t const numRows = residual.numel() / tk::kHiddenDim;
+    TORCH_CHECK(numRows >= 1 && numRows <= tk::kMaxRows,
+        "[glm_moe_dsa_post_attention_hop_from_partials] number of rows ", numRows,
+        " is outside the supported range [1, ", tk::kMaxRows, "]");
+    TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_post_attention_hop_from_partials] only TP", tk::kNumRanks,
+        " is supported, got ", nRanks);
+    TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_post_attention_hop_from_partials] invalid rank ", rank);
+    TORCH_CHECK(routerLoad >= 0 && routerLoad <= 2,
+        "[glm_moe_dsa_post_attention_hop_from_partials] router_load must be 0, 1 or 2, got ", routerLoad);
+    TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous()
+            && static_cast<size_t>(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes,
+        "[glm_moe_dsa_post_attention_hop_from_partials] uc_buffer must be a contiguous CUDA tensor of at least ",
+        tk::kWorkspaceBytes, " bytes");
+    TORCH_CHECK(flags.is_cuda() && flags.is_contiguous() && flags.scalar_type() == at::kInt && flags.numel() >= 4,
+        "[glm_moe_dsa_post_attention_hop_from_partials] flags must be a contiguous CUDA int32 tensor with at least 4 "
+        "elements");
+    TORCH_CHECK(residual.get_device() == ucBuffer.get_device() && residual.get_device() == flags.get_device(),
+        "[glm_moe_dsa_post_attention_hop_from_partials] all tensors must be on the same device");
+
+    at::Tensor routerLogits;
+    if (routerWeight.has_value())
+    {
+        auto const& w = routerWeight.value();
+        TORCH_CHECK(w.is_cuda() && w.scalar_type() == at::kBFloat16 && w.is_contiguous() && w.dim() == 2
+                && w.size(0) == tk::kNumExperts && w.size(1) == tk::kHiddenDim,
+            "[glm_moe_dsa_post_attention_hop_from_partials] router_weight must be a contiguous CUDA bf16 tensor of "
+            "shape [",
+            tk::kNumExperts, ", ", tk::kHiddenDim, "]");
+        routerLogits = at::empty({numRows, tk::kNumExperts}, residual.options().dtype(at::kFloat));
+    }
+    else
+    {
+        routerLogits = at::empty({0}, residual.options().dtype(at::kFloat));
+    }
+
+    c10::cuda::CUDAGuard const deviceGuard(residual.device());
+    at::Tensor normOut = at::empty_like(residual);
+    at::Tensor residualOut = at::empty_like(residual);
+
+    tk::PostAttentionHopFromPartialsParams params;
+    params.residualIn = residual.const_data_ptr();
+    params.gamma = gamma.const_data_ptr();
+    params.routerWeight = routerWeight.has_value() ? routerWeight->const_data_ptr() : nullptr;
+    params.normOut = normOut.mutable_data_ptr();
+    params.residualOut = residualOut.mutable_data_ptr();
+    params.routerLogits = routerWeight.has_value() ? routerLogits.mutable_data_ptr<float>() : nullptr;
+    params.ucBuffer = ucBuffer.mutable_data_ptr();
+    params.flags = reinterpret_cast<uint32_t*>(flags.mutable_data_ptr<int32_t>());
+    params.numRows = static_cast<int>(numRows);
+    params.rank = static_cast<int>(rank);
+    params.eps = static_cast<float>(eps);
+    params.enablePdl = pdlEnabled();
+    params.earlyPoll = earlyPoll;
+    params.routerLoad = static_cast<int>(routerLoad);
+    params.stream = at::cuda::getCurrentCUDAStream(residual.get_device());
+
+    cudaError_t const status = tk::invokePostAttentionHopFromPartials(params);
+    TORCH_CHECK(status == cudaSuccess,
+        "[glm_moe_dsa_post_attention_hop_from_partials] kernel launch failed: ", cudaGetErrorString(status));
+    return {normOut, residualOut, routerLogits};
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_post_attention_hop_from_partials(Tensor residual, Tensor gamma, float eps, Tensor? router_weight, "
+        "Tensor uc_buffer, Tensor flags, int rank, int nranks, bool early_poll, int router_load=1) -> Tensor[]");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_post_attention_hop_from_partials",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHopFromPartials);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp
new file mode 100644
index 000000000000..de99e3968c6b
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp
@@ -0,0 +1,170 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/library.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <optional>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::post_attention_hop;
+
+// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0.
+bool pdlEnabled()
+{
+    static bool const enabled = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && std::strcmp(env, "0") == 0);
+    }();
+    return enabled;
+}
+
+void checkRows(at::Tensor const& t, char const* name)
+{
+    TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_post_attention_hop] ", name,
+        " must be a CUDA bf16 tensor");
+    TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_post_attention_hop] ", name, " must be contiguous");
+    TORCH_CHECK(t.dim() >= 1 && t.size(-1) == tk::kHiddenDim, "[glm_moe_dsa_post_attention_hop] ", name,
+        " must have hidden size ", tk::kHiddenDim);
+}
+
+} // namespace
+
+//! Fused post-attention hop: one-shot cross-rank sum of ``input`` + residual + RMSNorm (+ router logits).
+//! Returns ``[normed, residual_out, router_logits]``; ``router_logits`` is an empty tensor when ``router_weight`` is
+//! not given. See kernels/glmMoeDsa/postAttentionHop.h for the arithmetic and the workspace protocol.
+std::vector<at::Tensor> postAttentionHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma,
+    double eps, std::optional<at::Tensor> const& routerWeight, at::Tensor const& ucBuffer, int64_t mcPtr,
+    at::Tensor const& flags, int64_t rank, int64_t nRanks)
+{
+    checkRows(input, "input");
+    checkRows(residual, "residual");
+    TORCH_CHECK(
+        residual.sizes() == input.sizes(), "[glm_moe_dsa_post_attention_hop] residual must have the shape of input");
+    checkRows(gamma, "gamma");
+    TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_post_attention_hop] gamma must be 1-D");
+    int64_t const numRows = input.numel() / tk::kHiddenDim;
+    TORCH_CHECK(numRows >= 1 && numRows <= tk::kMaxRows, "[glm_moe_dsa_post_attention_hop] number of rows ", numRows,
+        " is outside the supported range [1, ", tk::kMaxRows, "]");
+    TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_post_attention_hop] only TP", tk::kNumRanks,
+        " is supported, got ", nRanks);
+    TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_post_attention_hop] invalid rank ", rank);
+    TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous()
+            && static_cast<size_t>(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes,
+        "[glm_moe_dsa_post_attention_hop] uc_buffer must be a contiguous CUDA tensor of at least ", tk::kWorkspaceBytes,
+        " bytes");
+    TORCH_CHECK(
+        mcPtr != 0, "[glm_moe_dsa_post_attention_hop] mc_ptr must be the multicast address of uc_buffer's workspace");
+    TORCH_CHECK(flags.is_cuda() && flags.is_contiguous() && flags.scalar_type() == at::kInt && flags.numel() >= 4,
+        "[glm_moe_dsa_post_attention_hop] flags must be a contiguous CUDA int32 tensor with at least 4 elements");
+    TORCH_CHECK(input.get_device() == ucBuffer.get_device() && input.get_device() == flags.get_device(),
+        "[glm_moe_dsa_post_attention_hop] all tensors must be on the same device");
+
+    at::Tensor routerLogits;
+    if (routerWeight.has_value())
+    {
+        auto const& w = routerWeight.value();
+        TORCH_CHECK(w.is_cuda() && w.scalar_type() == at::kBFloat16 && w.is_contiguous() && w.dim() == 2
+                && w.size(0) == tk::kNumExperts && w.size(1) == tk::kHiddenDim,
+            "[glm_moe_dsa_post_attention_hop] router_weight must be a contiguous CUDA bf16 tensor of shape [",
+            tk::kNumExperts, ", ", tk::kHiddenDim, "]");
+        routerLogits = at::empty({numRows, tk::kNumExperts}, input.options().dtype(at::kFloat));
+    }
+    else
+    {
+        routerLogits = at::empty({0}, input.options().dtype(at::kFloat));
+    }
+
+    c10::cuda::CUDAGuard const deviceGuard(input.device());
+    at::Tensor normOut = at::empty_like(input);
+    at::Tensor residualOut = at::empty_like(residual);
+
+    tk::PostAttnHopParams params;
+    params.input = input.const_data_ptr();
+    params.residualIn = residual.const_data_ptr();
+    params.gamma = gamma.const_data_ptr();
+    params.routerWeight = routerWeight.has_value() ? routerWeight->const_data_ptr() : nullptr;
+    params.normOut = normOut.mutable_data_ptr();
+    params.residualOut = residualOut.mutable_data_ptr();
+    params.routerLogits = routerWeight.has_value() ? routerLogits.mutable_data_ptr<float>() : nullptr;
+    params.ucBuffer = ucBuffer.mutable_data_ptr();
+    params.mcBuffer = reinterpret_cast<void*>(mcPtr);
+    params.flags = reinterpret_cast<uint32_t*>(flags.mutable_data_ptr<int32_t>());
+    params.numRows = static_cast<int>(numRows);
+    params.rank = static_cast<int>(rank);
+    params.eps = static_cast<float>(eps);
+    params.enablePdl = pdlEnabled();
+    params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
+
+    cudaError_t const status = tk::invokePostAttentionHop(params);
+    TORCH_CHECK(
+        status == cudaSuccess, "[glm_moe_dsa_post_attention_hop] kernel launch failed: ", cudaGetErrorString(status));
+    return {normOut, residualOut, routerLogits};
+}
+
+#ifdef TLLM_GLM_LL_W2_TIMING
+//! Development builds only: phase timestamps of the most recent hop launch.
+at::Tensor postAttentionHopTiming()
+{
+    at::Tensor out = at::empty({tk::kTimingMaxCtas, tk::kTimingStamps}, at::TensorOptions().dtype(at::kLong));
+    auto stream = at::cuda::getCurrentCUDAStream();
+    TORCH_CHECK(tk::copyPostAttnHopTiming(reinterpret_cast<unsigned long long*>(out.data_ptr<int64_t>()), stream)
+        == cudaSuccess);
+    TORCH_CHECK(cudaStreamSynchronize(stream) == cudaSuccess);
+    return out;
+}
+#endif
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_post_attention_hop(Tensor input, Tensor residual, Tensor gamma, float eps, Tensor? router_weight, "
+        "Tensor uc_buffer, int mc_ptr, Tensor flags, int rank, int nranks) -> Tensor[]");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_post_attention_hop", &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHop);
+}
+
+#ifdef TLLM_GLM_LL_W2_TIMING
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def("glm_moe_dsa_post_attention_hop_timing() -> Tensor",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHopTiming);
+}
+#endif
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp
new file mode 100644
index 000000000000..0d71f2b34496
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp
@@ -0,0 +1,189 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/library.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <optional>
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::hop_routing;
+
+// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0.
+bool routingPdlEnabled()
+{
+    static bool const enabled = []()
+    {
+        char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+        return !(env != nullptr && std::strcmp(env, "0") == 0);
+    }();
+    return enabled;
+}
+
+void checkBf16(at::Tensor const& t, char const* name, int64_t lastDim)
+{
+    TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_post_attention_hop_routing] ", name,
+        " must be a CUDA bf16 tensor");
+    TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_post_attention_hop_routing] ", name, " must be contiguous");
+    TORCH_CHECK(t.dim() >= 1 && t.size(-1) == lastDim, "[glm_moe_dsa_post_attention_hop_routing] ", name,
+        " must have last dimension ", lastDim);
+}
+
+void checkWords(at::Tensor const& t, char const* name, int64_t minWords)
+{
+    TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == at::kInt && t.numel() >= minWords,
+        "[glm_moe_dsa_post_attention_hop_routing] ", name, " must be a contiguous CUDA int32 tensor with at least ",
+        minWords, " elements");
+}
+
+} // namespace
+
+//! Pre-MoE hop with top-8 routing: the transport-input hop with router logits, plus the native
+//! DeepSeek-V3 routing of every row. Returns ``[normed, residual_out, router_logits, routing]`` where ``routing`` is
+//! an int32 ``[2, R, 8]`` tensor: ``routing[0]`` holds the expert ids and ``routing[1]`` the FP32 bits of the
+//! bf16-rounded weights, both in descending native key order. See kernels/glmMoeDsa/postAttentionHopRouting.h.
+std::vector<at::Tensor> postAttentionHopRouting(at::Tensor const& residual, at::Tensor const& gamma, double eps,
+    at::Tensor const& routerWeight, at::Tensor const& routingBias, at::Tensor const& ucBuffer, at::Tensor const& flags,
+    at::Tensor const& gather, at::Tensor const& zstate, int64_t rank, int64_t nRanks, bool earlyPoll,
+    int64_t routerLoad, double routedScaling, std::optional<at::Tensor> const& logitsIn,
+    std::optional<at::Tensor> const& stamps)
+{
+    checkBf16(residual, "residual", tk::kHiddenDim);
+    checkBf16(gamma, "gamma", tk::kHiddenDim);
+    TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_post_attention_hop_routing] gamma must be 1-D");
+    checkBf16(routerWeight, "router_weight", tk::kHiddenDim);
+    TORCH_CHECK(routerWeight.dim() == 2 && routerWeight.size(0) == tk::kNumExperts,
+        "[glm_moe_dsa_post_attention_hop_routing] router_weight must have shape [", tk::kNumExperts, ", ",
+        tk::kHiddenDim, "]");
+    checkBf16(routingBias, "routing_bias", tk::kNumExperts);
+    TORCH_CHECK(routingBias.dim() == 1, "[glm_moe_dsa_post_attention_hop_routing] routing_bias must be 1-D");
+    int64_t const numRows = residual.numel() / tk::kHiddenDim;
+    TORCH_CHECK(numRows >= 1 && numRows <= tk::kMaxRows, "[glm_moe_dsa_post_attention_hop_routing] number of rows ",
+        numRows, " is outside the supported range [1, ", tk::kMaxRows, "]");
+    TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_post_attention_hop_routing] only TP", tk::kNumRanks,
+        " is supported, got ", nRanks);
+    TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_post_attention_hop_routing] invalid rank ", rank);
+    TORCH_CHECK(routerLoad >= 0 && routerLoad <= 2,
+        "[glm_moe_dsa_post_attention_hop_routing] router_load must be 0, 1 or 2, got ", routerLoad);
+    TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous()
+            && static_cast<size_t>(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes,
+        "[glm_moe_dsa_post_attention_hop_routing] uc_buffer must be a contiguous CUDA tensor of at least ",
+        tk::kWorkspaceBytes, " bytes");
+    checkWords(flags, "flags", 4);
+    checkWords(gather, "gather", tk::kGatherWords);
+    checkWords(zstate, "zstate", tk::kStateWords);
+    auto const device = residual.get_device();
+    TORCH_CHECK(gamma.get_device() == device && routerWeight.get_device() == device
+            && routingBias.get_device() == device && ucBuffer.get_device() == device && flags.get_device() == device
+            && gather.get_device() == device && zstate.get_device() == device,
+        "[glm_moe_dsa_post_attention_hop_routing] all tensors must be on the same device");
+    if (logitsIn.has_value())
+    {
+        auto const& l = logitsIn.value();
+        TORCH_CHECK(l.is_cuda() && l.get_device() == device && l.scalar_type() == at::kFloat && l.is_contiguous()
+                && l.numel() == numRows * tk::kNumExperts,
+            "[glm_moe_dsa_post_attention_hop_routing] logits_in must be a contiguous CUDA fp32 tensor of shape [",
+            numRows, ", ", tk::kNumExperts, "]");
+        TORCH_CHECK(routerLoad == 1, "[glm_moe_dsa_post_attention_hop_routing] logits_in needs router_load=1");
+    }
+    if (stamps.has_value())
+    {
+        auto const& s = stamps.value();
+        TORCH_CHECK(tk::stampsCompiled(),
+            "[glm_moe_dsa_post_attention_hop_routing] stamps need a TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS build");
+        TORCH_CHECK(s.is_cuda() && s.get_device() == device && s.scalar_type() == at::kLong && s.is_contiguous()
+                && s.numel() >= tk::kMaxGridCtas * tk::kNumStamps,
+            "[glm_moe_dsa_post_attention_hop_routing] stamps must be a contiguous CUDA int64 tensor with at least ",
+            tk::kMaxGridCtas * tk::kNumStamps, " elements");
+    }
+
+    c10::cuda::CUDAGuard const deviceGuard(residual.device());
+    at::Tensor normOut = at::empty_like(residual);
+    at::Tensor residualOut = at::empty_like(residual);
+    at::Tensor routerLogits = at::empty({numRows, tk::kNumExperts}, residual.options().dtype(at::kFloat));
+    at::Tensor routing = at::empty({2, numRows, tk::kTopK}, residual.options().dtype(at::kInt));
+
+    tk::PostAttentionHopRoutingParams params;
+    params.residualIn = residual.const_data_ptr();
+    params.gamma = gamma.const_data_ptr();
+    params.routerWeight = routerWeight.const_data_ptr();
+    params.routingBias = routingBias.const_data_ptr();
+    params.logitsIn = logitsIn.has_value() ? logitsIn->const_data_ptr<float>() : nullptr;
+    params.normOut = normOut.mutable_data_ptr();
+    params.residualOut = residualOut.mutable_data_ptr();
+    params.routerLogits = routerLogits.mutable_data_ptr<float>();
+    params.topkIds = routing.mutable_data_ptr<int32_t>();
+    params.topkWeights = reinterpret_cast<float*>(routing.mutable_data_ptr<int32_t>() + numRows * tk::kTopK);
+    params.ucBuffer = ucBuffer.mutable_data_ptr();
+    params.flags = reinterpret_cast<uint32_t*>(flags.mutable_data_ptr<int32_t>());
+    params.gather = reinterpret_cast<uint32_t*>(gather.mutable_data_ptr<int32_t>());
+    params.zstate = reinterpret_cast<uint32_t*>(zstate.mutable_data_ptr<int32_t>());
+    params.stamps
+        = stamps.has_value() ? reinterpret_cast<unsigned long long*>(stamps->mutable_data_ptr<int64_t>()) : nullptr;
+    params.numRows = static_cast<int>(numRows);
+    params.rank = static_cast<int>(rank);
+    params.eps = static_cast<float>(eps);
+    params.routedScaling = static_cast<float>(routedScaling);
+    params.enablePdl = routingPdlEnabled();
+    params.earlyPoll = earlyPoll;
+    params.routerLoad = static_cast<int>(routerLoad);
+    params.stream = at::cuda::getCurrentCUDAStream(device);
+
+    cudaError_t const status = tk::invokePostAttentionHopRouting(params);
+    TORCH_CHECK(status == cudaSuccess,
+        "[glm_moe_dsa_post_attention_hop_routing] kernel launch failed: ", cudaGetErrorString(status));
+    return {normOut, residualOut, routerLogits, routing};
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+// A diagnostic build (stamps) registers under its own name, so it can be loaded next to a build that has the op.
+#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS)
+#define GLM_MOE_DSA_HOP_ROUTING_OP_NAME "glm_moe_dsa_post_attention_hop_routing_diag"
+#else
+#define GLM_MOE_DSA_HOP_ROUTING_OP_NAME "glm_moe_dsa_post_attention_hop_routing"
+#endif
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(GLM_MOE_DSA_HOP_ROUTING_OP_NAME
+        "(Tensor residual, Tensor gamma, float eps, Tensor router_weight, Tensor routing_bias, "
+        "Tensor uc_buffer, Tensor flags, Tensor gather, Tensor zstate, int rank, int nranks, bool early_poll, "
+        "int router_load, float routed_scaling, Tensor? logits_in=None, Tensor? stamps=None) -> Tensor[]");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl(GLM_MOE_DSA_HOP_ROUTING_OP_NAME, &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHopRouting);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp
new file mode 100644
index 000000000000..4411b6f88af8
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp
@@ -0,0 +1,330 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/common/envUtils.h"
+#include "tensorrt_llm/common/quantization.h"
+#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h"
+#include "tensorrt_llm/kernels/kvCacheUtils.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include <climits>
+#include <cmath>
+#include <cstdint>
+#include <optional>
+
+namespace tk = tensorrt_llm::kernels;
+namespace tc = tensorrt_llm::common;
+namespace rope_kv_append = tensorrt_llm::kernels::glm_moe_dsa::rope_kv_append;
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+void checkCuda(torch::Tensor const& t, char const* name)
+{
+    TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+}
+
+//! Paged MLA latent cache of one layer: the KVBlockArray that buildPagedKvCacheBuffers (thop/attentionOp.cpp)
+//! builds for an FP8 MLA cache (kvFactor 1, one KV head, 8-bit elements, beam width 1, no FP4 scale pool).
+tk::KVBlockArray buildFp8MlaKvBlockArray(torch::Tensor const& kvCacheBlockOffsets,
+    torch::Tensor const& hostKvCachePoolPointers, torch::Tensor const& hostKvCachePoolMapping, int64_t layerIdx,
+    int32_t batchSize, int64_t tokensPerBlock, int64_t attentionWindowSize, int64_t seqOffset)
+{
+    TORCH_CHECK(hostKvCachePoolMapping.device().is_cpu() && hostKvCachePoolMapping.dim() == 2
+            && hostKvCachePoolMapping.size(1) >= 2 && hostKvCachePoolMapping.scalar_type() == torch::kInt32,
+        "host_kv_cache_pool_mapping must be a CPU int32 [layers, >=2] tensor");
+    TORCH_CHECK(layerIdx >= 0 && layerIdx < hostKvCachePoolMapping.size(0), "layer index out of the pool mapping");
+    auto const* mapping = hostKvCachePoolMapping.data_ptr<int32_t>();
+    int32_t const poolIndex = mapping[layerIdx * hostKvCachePoolMapping.stride(0)];
+    int32_t const layerIdxInCachePool
+        = mapping[layerIdx * hostKvCachePoolMapping.stride(0) + hostKvCachePoolMapping.stride(1)];
+
+    TORCH_CHECK(hostKvCachePoolPointers.device().is_cpu() && hostKvCachePoolPointers.dim() == 2
+            && hostKvCachePoolPointers.scalar_type() == torch::kInt64 && poolIndex >= 0
+            && poolIndex < hostKvCachePoolPointers.size(0) && hostKvCachePoolPointers.size(1) >= 2,
+        "host_kv_cache_pool_pointers must be a CPU int64 [pools, 2] tensor");
+    TORCH_CHECK(kvCacheBlockOffsets.is_cuda() && kvCacheBlockOffsets.scalar_type() == torch::kInt32
+            && kvCacheBlockOffsets.dim() >= 2 && poolIndex < kvCacheBlockOffsets.size(0)
+            && seqOffset < kvCacheBlockOffsets.size(1),
+        "kv_cache_block_offsets must be CUDA int32 [pools, seqs, 2, max_blocks]");
+
+    int64_t constexpr kCacheElemBits = 8;
+    int64_t const blockSize = tokensPerBlock * /*kv_head_num=*/1 * rope_kv_append::kLatentDim;
+    int64_t const bytesPerBlock = blockSize * kCacheElemBits / CHAR_BIT;
+    int64_t const intraPoolOffset = layerIdxInCachePool * /*kvFactor=*/1 * bytesPerBlock;
+    auto const sizePerToken = static_cast<int32_t>(rope_kv_append::kLatentDim * kCacheElemBits / 8);
+    auto const* pointers = hostKvCachePoolPointers.data_ptr<int64_t>();
+    auto const rowOffset = poolIndex * hostKvCachePoolPointers.stride(0);
+    void* primaryPoolPtr = reinterpret_cast<char*>(pointers[rowOffset]) + intraPoolOffset;
+    void* secondaryPoolPtr
+        = reinterpret_cast<char*>(pointers[rowOffset + hostKvCachePoolPointers.stride(1)]) + intraPoolOffset;
+    auto* blockOffsets
+        = reinterpret_cast<tk::KVBlockArray::DataType*>(static_cast<int32_t*>(kvCacheBlockOffsets.data_ptr())
+            + poolIndex * kvCacheBlockOffsets.stride(0) + seqOffset * kvCacheBlockOffsets.stride(1));
+    auto const maxBlocksPerSequence = static_cast<int32_t>(kvCacheBlockOffsets.size(-1));
+    auto const window = static_cast<int32_t>(attentionWindowSize);
+    return tk::KVBlockArray(batchSize, maxBlocksPerSequence, static_cast<int32_t>(tokensPerBlock), sizePerToken, window,
+        window, /*sinkTokenLen=*/0, /*canUseOneMoreBlock=*/false, primaryPoolPtr, secondaryPoolPtr, blockOffsets);
+}
+
+} // namespace
+
+//! Generation MLA RoPE + FP8 query quantization + FP8 latent-cache append (the FP8-cache, absorption-mode,
+//! non-helix subset of trtllm::mla_rope_generation), optionally fused with the DSA indexer K-cache scatter of
+//! trtllm::indexer_k_cache_scatter_op. Argument meanings follow those two native ops.
+void ropeKvAppend(torch::Tensor const& fused_q, torch::Tensor const& q_pe, torch::Tensor const& latent_cache,
+    torch::Tensor const& rotary_cos_sin, torch::Tensor const& cu_q_seqlens, torch::Tensor const& cu_kv_seqlens,
+    torch::Tensor const& fmha_scheduler_counter, torch::Tensor const& mla_bmm1_scale,
+    torch::Tensor const& mla_bmm2_scale, torch::Tensor const& quant_q_buffer, torch::Tensor const& sequence_length,
+    torch::Tensor const& host_context_lengths, int64_t num_contexts, torch::Tensor const& kv_cache_block_offsets,
+    torch::Tensor const& host_kv_cache_pool_pointers, torch::Tensor const& host_kv_cache_pool_mapping,
+    std::optional<torch::Tensor> const& kv_cache_scale_orig_quant, int64_t layer_idx, int64_t num_heads,
+    int64_t tokens_per_block, int64_t attention_window_size, int64_t quant_mode, double q_scaling,
+    int64_t qk_nope_head_dim, int64_t qk_rope_head_dim, int64_t kv_lora_rank,
+    std::optional<torch::Tensor> const& idx_k_fp8, std::optional<torch::Tensor> const& idx_k_scale,
+    std::optional<torch::Tensor> const& idx_k_cache, std::optional<torch::Tensor> const& idx_slot_mapping_fp8,
+    std::optional<torch::Tensor> const& idx_slot_mapping_scale, int64_t idx_num_tokens)
+{
+    TORCH_CHECK(kv_lora_rank == rope_kv_append::kKvLoraRank && qk_rope_head_dim == rope_kv_append::kRopeDim,
+        "glm_moe_dsa_rope_kv_append: supports kv_lora_rank 512 and qk_rope_head_dim 64 only");
+    auto const kvQuantMode = tc::QuantMode(static_cast<uint32_t>(quant_mode));
+    TORCH_CHECK(kvQuantMode.hasFp8KvCache() && !kvQuantMode.hasFp4KvCache(),
+        "glm_moe_dsa_rope_kv_append: needs an FP8 MLA cache");
+    checkCuda(fused_q, "fused_q");
+    checkCuda(q_pe, "q_pe");
+    checkCuda(latent_cache, "latent_cache");
+    TORCH_CHECK(fused_q.scalar_type() == torch::kBFloat16 && q_pe.scalar_type() == torch::kBFloat16
+            && latent_cache.scalar_type() == torch::kBFloat16,
+        "glm_moe_dsa_rope_kv_append: needs BF16 fused_q, q_pe and latent_cache");
+    TORCH_CHECK(rotary_cos_sin.scalar_type() == torch::kFloat32 && rotary_cos_sin.is_contiguous(),
+        "rotary_cos_sin must be packed float32");
+    int64_t const cosSinRows = rotary_cos_sin.numel() / (2 * rope_kv_append::kRopeDim);
+    TORCH_CHECK(cosSinRows > 0 && cosSinRows <= INT32_MAX, "rotary_cos_sin has no complete position row");
+    TORCH_CHECK(q_pe.dim() == 3 && q_pe.stride(2) == 1, "q_pe must be [tokens, heads, rope] with unit inner stride");
+    TORCH_CHECK(latent_cache.dim() == 2 && latent_cache.size(1) == rope_kv_append::kLatentDim
+            && latent_cache.stride(1) == 1 && latent_cache.stride(0) == rope_kv_append::kLatentDim,
+        "latent_cache must be packed [tokens, 576]");
+    TORCH_CHECK(fused_q.is_contiguous() && quant_q_buffer.is_contiguous(), "fused_q and quant_q_buffer must be packed");
+    TORCH_CHECK(host_kv_cache_pool_mapping.dim() >= 1, "host_kv_cache_pool_mapping is required");
+
+    int32_t const numTokens = static_cast<int32_t>(latent_cache.size(0));
+    int32_t const numSeqs = static_cast<int32_t>(host_context_lengths.size(0));
+    int32_t const numGenerations = numSeqs - static_cast<int32_t>(num_contexts);
+    TORCH_CHECK(numGenerations > 0, "glm_moe_dsa_rope_kv_append: needs at least one generation request");
+    TORCH_CHECK(numTokens % numGenerations == 0, "MLA can only support input sequences with the same sequence length.");
+    TORCH_CHECK(fused_q.numel() >= static_cast<int64_t>(numTokens) * num_heads * rope_kv_append::kLatentDim,
+        "fused_q is smaller than [tokens, heads, 576]");
+    TORCH_CHECK(quant_q_buffer.numel() * quant_q_buffer.element_size()
+            >= static_cast<int64_t>(numTokens) * num_heads * rope_kv_append::kLatentDim,
+        "quant_q_buffer is smaller than [tokens, heads, 576] bytes");
+    TORCH_CHECK(mla_bmm1_scale.numel() >= 2 && mla_bmm2_scale.numel() >= 1, "bmm scale buffers too small");
+    TORCH_CHECK(cu_q_seqlens.numel() >= numGenerations + 1 && cu_kv_seqlens.numel() >= numGenerations + 1,
+        "cu_seqlens buffers too small");
+
+    auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device());
+    rope_kv_append::RopeAppendParams params{};
+    params.fusedQ = reinterpret_cast<__nv_bfloat16 const*>(fused_q.data_ptr());
+    params.qPe = reinterpret_cast<__nv_bfloat16 const*>(q_pe.data_ptr());
+    params.qPeLd = static_cast<int32_t>(q_pe.stride(1));
+    params.qPeStride = static_cast<int32_t>(q_pe.stride(0));
+    params.quantQ = reinterpret_cast<uint8_t*>(quant_q_buffer.data_ptr());
+    params.quantScaleQ = nullptr;
+    params.latent = reinterpret_cast<__nv_bfloat16 const*>(latent_cache.data_ptr());
+    params.cosSin = reinterpret_cast<float2 const*>(rotary_cos_sin.data_ptr());
+    params.cosSinRows = static_cast<int32_t>(cosSinRows);
+    params.kvLens = sequence_length.slice(0, num_contexts).data_ptr<int32_t>();
+    params.quantScaleKv
+        = kv_cache_scale_orig_quant.has_value() ? kv_cache_scale_orig_quant->data_ptr<float>() : nullptr;
+    params.kvCache = buildFp8MlaKvBlockArray(kv_cache_block_offsets, host_kv_cache_pool_pointers,
+        host_kv_cache_pool_mapping, layer_idx, numGenerations, tokens_per_block, attention_window_size, num_contexts);
+    params.cuQSeqlens = cu_q_seqlens.data_ptr<int32_t>();
+    params.cuKvSeqlens = cu_kv_seqlens.data_ptr<int32_t>();
+    params.fmhaTileCounter = reinterpret_cast<uint32_t*>(fmha_scheduler_counter.data_ptr());
+    params.bmm1Scale = mla_bmm1_scale.data_ptr<float>();
+    params.bmm2Scale = mla_bmm2_scale.data_ptr<float>();
+    params.dequantScaleQ = nullptr;
+    params.dequantScaleKv = nullptr;
+    params.quantScaleO = nullptr;
+    // Same expression as the native op so the folded FMHA scales match bit for bit.
+    params.hostBmm1Scale = 1.F / (q_scaling * std::sqrt(static_cast<float>(qk_nope_head_dim + qk_rope_head_dim)));
+    params.numTokens = numTokens;
+    params.seqLen = numTokens / numGenerations;
+    params.numHeads = static_cast<int32_t>(num_heads);
+
+    if (idx_num_tokens > 0)
+    {
+        TORCH_CHECK(idx_k_fp8.has_value() && idx_k_scale.has_value() && idx_k_cache.has_value()
+                && idx_slot_mapping_fp8.has_value() && idx_slot_mapping_scale.has_value(),
+            "indexer scatter needs k_fp8, k_scale, k_cache and both slot mappings");
+        auto const& kFp8 = idx_k_fp8.value();
+        auto const& kScale = idx_k_scale.value();
+        auto const& kCache = idx_k_cache.value();
+        auto const& slotFp8 = idx_slot_mapping_fp8.value();
+        auto const& slotScale = idx_slot_mapping_scale.value();
+        checkCuda(kFp8, "k_fp8");
+        checkCuda(kScale, "k_scale");
+        checkCuda(kCache, "k_cache");
+        TORCH_CHECK(kFp8.dim() == 2 && kFp8.element_size() == 1 && kFp8.is_contiguous(), "k_fp8 must be packed bytes");
+        TORCH_CHECK(kScale.dim() == 2 && kScale.element_size() == 4 && kScale.is_contiguous()
+                && kScale.size(1) * kScale.element_size() == 4,
+            "k_scale must be one packed 4-byte word per token");
+        TORCH_CHECK(kCache.dim() == 4 && kCache.size(2) == 1, "k_cache must be [blocks, block_size, 1, bytes]");
+        TORCH_CHECK(slotFp8.scalar_type() == torch::kInt64 && slotScale.scalar_type() == torch::kInt64
+                && slotFp8.is_contiguous() && slotScale.is_contiguous(),
+            "slot mappings must be packed int64");
+        int32_t const headBytes = static_cast<int32_t>(kFp8.size(1));
+        TORCH_CHECK(headBytes == 128 || headBytes == 64, "indexer k head bytes must be 128 or 64");
+        TORCH_CHECK(kFp8.size(0) >= idx_num_tokens && kScale.size(0) >= idx_num_tokens
+                && slotFp8.numel() >= idx_num_tokens && slotScale.numel() >= idx_num_tokens,
+            "indexer inputs are shorter than idx_num_tokens");
+        params.idxK = reinterpret_cast<uint8_t const*>(kFp8.data_ptr());
+        params.idxScale = reinterpret_cast<uint8_t const*>(kScale.data_ptr());
+        params.idxCache = reinterpret_cast<uint8_t*>(kCache.data_ptr());
+        params.idxSlotData = slotFp8.data_ptr<int64_t>();
+        params.idxSlotScale = slotScale.data_ptr<int64_t>();
+        params.idxNumTokens = static_cast<int32_t>(idx_num_tokens);
+        params.idxHeadBytes = headBytes;
+        for (int32_t d = 0; d < 4; ++d)
+        {
+            params.idxDims[d] = static_cast<int32_t>(kCache.size(d));
+            params.idxStrides[d] = kCache.stride(d);
+        }
+    }
+
+    rope_kv_append::invokeRopeAppendScatter(params, tc::getEnvEnablePDL(), stream);
+}
+
+//! Local top-k token ids -> global KV-pool indices (values of trtllm::convert_req_index_to_global and its grouped
+//! variant), optionally fused with the native top-k reuse copy. Returns [group, rows, topk] when layer_ids is given,
+//! otherwise [rows, topk].
+torch::Tensor topkIndexBridge(torch::Tensor const& topk, int64_t num_rows, torch::Tensor const& req_idx,
+    int64_t req_idx_bias, torch::Tensor const& block_table, int64_t block_size, int64_t stride_factor,
+    std::optional<torch::Tensor> const& layer_ids, int64_t layer_id, std::optional<torch::Tensor> const& copy_dst,
+    int64_t copy_row_start, int64_t copy_rows, int64_t copy_mode, int64_t next_n,
+    std::optional<torch::Tensor> const& num_accepted)
+{
+    checkCuda(topk, "topk");
+    checkCuda(req_idx, "req_idx");
+    checkCuda(block_table, "block_table");
+    TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.stride(1) == 1,
+        "topk must be int32 [rows, topk] with unit column stride");
+    TORCH_CHECK(req_idx.scalar_type() == torch::kInt32 && req_idx.dim() == 1 && req_idx.is_contiguous(),
+        "req_idx must be packed int32");
+    TORCH_CHECK(block_table.scalar_type() == torch::kInt32 && block_table.dim() == 2, "block_table must be int32 2-D");
+    TORCH_CHECK(copy_mode == 0 || copy_mode == 1, "copy_mode must be 0 or 1");
+    int64_t const numTopk = topk.size(1);
+    TORCH_CHECK(req_idx.size(0) >= num_rows, "req_idx shorter than num_rows");
+    TORCH_CHECK(copy_mode == 1 || topk.size(0) >= num_rows, "topk shorter than num_rows");
+    TORCH_CHECK(copy_mode == 0 || topk.size(0) >= copy_rows * next_n, "topk shorter than copy_rows * next_n");
+    TORCH_CHECK(copy_mode == 1 || copy_rows <= num_rows, "copy_rows exceeds num_rows");
+
+    int32_t groupSize = 1;
+    int32_t const* layerIdsPtr = nullptr;
+    if (layer_ids.has_value())
+    {
+        TORCH_CHECK(layer_ids->scalar_type() == torch::kInt32 && layer_ids->dim() == 1 && layer_ids->is_contiguous(),
+            "layer_ids must be packed int32");
+        groupSize = static_cast<int32_t>(layer_ids->size(0));
+        layerIdsPtr = layer_ids->data_ptr<int32_t>();
+    }
+    torch::Tensor out = layer_ids.has_value() ? torch::empty({groupSize, num_rows, numTopk}, topk.options())
+                                              : torch::empty({num_rows, numTopk}, topk.options());
+
+    rope_kv_append::TopkBridgeParams params{};
+    params.topk = topk.data_ptr<int32_t>();
+    params.topkStride0 = topk.stride(0);
+    params.numRows = static_cast<int32_t>(num_rows);
+    params.numTopk = static_cast<int32_t>(numTopk);
+    params.reqIdx = req_idx.data_ptr<int32_t>();
+    params.reqIdxBias = static_cast<int32_t>(req_idx_bias);
+    params.blockTable = block_table.data_ptr<int32_t>();
+    params.btStride0 = block_table.stride(0);
+    params.btStride1 = block_table.stride(1);
+    params.maxBlocksPerReq = static_cast<int32_t>(block_table.size(1));
+    params.blockSize = static_cast<int32_t>(block_size);
+    params.strideFactor = static_cast<int32_t>(stride_factor);
+    params.layerIds = layerIdsPtr;
+    params.layerId = static_cast<int32_t>(layer_id);
+    params.groupSize = groupSize;
+    params.out = out.data_ptr<int32_t>();
+    params.outStrideG = layer_ids.has_value() ? out.stride(0) : 0;
+    params.outStride0 = layer_ids.has_value() ? out.stride(1) : out.stride(0);
+    if (copy_dst.has_value() && copy_rows > 0)
+    {
+        auto const& dst = copy_dst.value();
+        checkCuda(dst, "copy_dst");
+        TORCH_CHECK(dst.scalar_type() == torch::kInt32 && dst.dim() == 2 && dst.stride(1) == 1 && dst.size(1) >= numTopk
+                && dst.size(0) >= copy_row_start + copy_rows,
+            "copy_dst must be int32 [>= start + rows, >= topk] with unit column stride");
+        params.copyDst = dst.data_ptr<int32_t>();
+        params.copyStride0 = dst.stride(0);
+        params.copyRows = static_cast<int32_t>(copy_rows);
+        params.copyRowStart = static_cast<int32_t>(copy_row_start);
+        params.copyMode = static_cast<int32_t>(copy_mode);
+        params.nextN = static_cast<int32_t>(next_n);
+        if (copy_mode == 1 && num_accepted.has_value())
+        {
+            auto const& acc = num_accepted.value();
+            checkCuda(acc, "num_accepted");
+            TORCH_CHECK((acc.scalar_type() == torch::kInt32 || acc.scalar_type() == torch::kInt64) && acc.dim() == 1
+                    && acc.is_contiguous() && acc.size(0) >= copy_rows,
+                "num_accepted must be packed int32/int64 with at least copy_rows entries");
+            params.numAccepted = acc.data_ptr();
+            params.numAcceptedIs64 = acc.scalar_type() == torch::kInt64;
+        }
+    }
+
+    auto stream = at::cuda::getCurrentCUDAStream(topk.get_device());
+    rope_kv_append::invokeTopkBridge(params, tc::getEnvEnablePDL(), stream);
+    return out;
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_rope_kv_append(Tensor fused_q, Tensor q_pe, Tensor latent_cache, Tensor rotary_cos_sin, "
+        "Tensor(a!) cu_q_seqlens, Tensor(b!) cu_kv_seqlens, Tensor(c!) fmha_scheduler_counter, "
+        "Tensor(d!) mla_bmm1_scale, Tensor(e!) mla_bmm2_scale, Tensor(f!) quant_q_buffer, Tensor sequence_length, "
+        "Tensor host_context_lengths, int num_contexts, Tensor kv_cache_block_offsets, "
+        "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_cache_scale_orig_quant, "
+        "int layer_idx, int num_heads, int tokens_per_block, int attention_window_size, int quant_mode, "
+        "float q_scaling, int qk_nope_head_dim, int qk_rope_head_dim, int kv_lora_rank, Tensor? idx_k_fp8, "
+        "Tensor? idx_k_scale, Tensor(g!)? idx_k_cache, Tensor? idx_slot_mapping_fp8, Tensor? idx_slot_mapping_scale, "
+        "int idx_num_tokens) -> ()");
+    m.def(
+        "glm_moe_dsa_topk_index_bridge(Tensor topk, int num_rows, Tensor req_idx, int req_idx_bias, Tensor "
+        "block_table, "
+        "int block_size, int stride_factor, Tensor? layer_ids, int layer_id, Tensor(a!)? copy_dst, "
+        "int copy_row_start, int copy_rows, int copy_mode, int next_n, Tensor? num_accepted) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_rope_kv_append", &tensorrt_llm::torch_ext::glm_moe_dsa::ropeKvAppend);
+    m.impl("glm_moe_dsa_topk_index_bridge", &tensorrt_llm::torch_ext::glm_moe_dsa::topkIndexBridge);
+}
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp
new file mode 100644
index 000000000000..a7dba3610644
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp
@@ -0,0 +1,121 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
+#include <torch/types.h>
+
+#include <optional>
+
+namespace tk = tensorrt_llm::kernels::glm_moe_dsa::sparse_mla_decode;
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext::glm_moe_dsa
+{
+
+namespace
+{
+
+bool isFp8ByteTensor(torch::Tensor const& t)
+{
+    return t.scalar_type() == torch::kUInt8 || t.scalar_type() == torch::kFloat8_e4m3fn;
+}
+
+} // namespace
+
+// Sparse MLA generation attention; see sparseMlaDecode.h.
+void sparseMlaDecode(torch::Tensor const& q, int64_t kvPoolPtr, torch::Tensor const& indices,
+    torch::Tensor const& bmm1Scale, torch::Tensor const& bmm2Scale, torch::Tensor const& seqLensKv, int64_t seqLenQ,
+    torch::Tensor& out)
+{
+    TORCH_CHECK(q.is_cuda() && indices.is_cuda() && bmm1Scale.is_cuda() && bmm2Scale.is_cuda() && out.is_cuda(),
+        "glm_moe_dsa_sparse_mla_decode expects CUDA tensors");
+    TORCH_CHECK(isFp8ByteTensor(q), "q must hold e4m3 bytes (uint8 or float8_e4m3fn)");
+    TORCH_CHECK(q.is_contiguous(), "q must be contiguous");
+    int64_t const numRows = q.size(0);
+    TORCH_CHECK(q.numel() == numRows * tk::kNumHeads * tk::kHeadDimQk, "q must be [rows, ", tk::kNumHeads, ", ",
+        tk::kHeadDimQk, "]");
+    TORCH_CHECK(indices.scalar_type() == torch::kInt32 && indices.dim() == 2, "indices must be a 2-D int32 tensor");
+    TORCH_CHECK(indices.size(0) == numRows, "indices must have one row per query row");
+    TORCH_CHECK(indices.size(1) == tk::kTopK, "indices must hold ", tk::kTopK, " entries per row");
+    TORCH_CHECK(indices.stride(1) == 1, "indices rows must be contiguous");
+    TORCH_CHECK(bmm1Scale.scalar_type() == torch::kFloat32 && bmm1Scale.numel() >= 2, "bmm1_scale must be float32[2]");
+    TORCH_CHECK(bmm2Scale.scalar_type() == torch::kFloat32 && bmm2Scale.numel() >= 1, "bmm2_scale must be float32[1]");
+    TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.is_contiguous(), "out must be contiguous bfloat16");
+    TORCH_CHECK(out.numel() == numRows * tk::kNumHeads * tk::kHeadDimV, "out must be [rows, ",
+        tk::kNumHeads * tk::kHeadDimV, "]");
+    TORCH_CHECK(kvPoolPtr != 0, "kv_pool_ptr must be a valid device address");
+
+    tk::SparseMlaDecodeParams params{};
+    params.q = q.data_ptr();
+    params.kvPool = reinterpret_cast<void const*>(kvPoolPtr);
+    params.indices = indices.data_ptr<int32_t>();
+    params.indicesRowStride = indices.stride(0);
+    params.bmm1Scale = bmm1Scale.data_ptr<float>();
+    params.bmm2Scale = bmm2Scale.data_ptr<float>();
+    TORCH_CHECK(seqLensKv.is_cuda() && seqLensKv.scalar_type() == torch::kInt32 && seqLensKv.is_contiguous(),
+        "seq_lens_kv must be a contiguous int32 CUDA tensor");
+    TORCH_CHECK(seqLenQ > 0 && numRows % seqLenQ == 0, "rows must be a multiple of seq_len_q");
+    TORCH_CHECK(seqLensKv.numel() >= numRows / seqLenQ, "seq_lens_kv must cover every sequence");
+    params.seqLensKv = seqLensKv.data_ptr<int32_t>();
+    params.seqLenQ = static_cast<int32_t>(seqLenQ);
+    params.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr());
+    params.numRows = static_cast<int32_t>(numRows);
+
+    auto const stream = at::cuda::getCurrentCUDAStream(q.get_device());
+    tk::invokeSparseMlaDecode(params, stream);
+}
+
+#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING
+torch::Tensor sparseMlaDecodeDebugTimeline()
+{
+    auto out = torch::empty({tk::kTimelineMaxCtas, tk::kTimelineSlots}, torch::dtype(torch::kInt64));
+    tk::copySparseMlaDecodeTimeline(reinterpret_cast<unsigned long long*>(out.data_ptr<int64_t>()));
+    return out;
+}
+#endif
+
+bool sparseMlaDecodeSupported(int64_t device)
+{
+    return tk::isSparseMlaDecodeDeviceSupported(static_cast<int>(device));
+}
+
+} // namespace torch_ext::glm_moe_dsa
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "glm_moe_dsa_sparse_mla_decode(Tensor q, int kv_pool_ptr, Tensor indices, Tensor bmm1_scale, "
+        "Tensor bmm2_scale, Tensor seq_lens_kv, int seq_len_q, Tensor(a!) out) -> ()");
+    m.def("glm_moe_dsa_sparse_mla_decode_supported(int device) -> bool",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::sparseMlaDecodeSupported);
+#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING
+    m.def("glm_moe_dsa_sparse_mla_decode_debug_timeline() -> Tensor",
+        &tensorrt_llm::torch_ext::glm_moe_dsa::sparseMlaDecodeDebugTimeline);
+#endif
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("glm_moe_dsa_sparse_mla_decode", &tensorrt_llm::torch_ext::glm_moe_dsa::sparseMlaDecode);
+}

From e10663b9b4b76ae002a42f6c944bcbc8ebde6849 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 09:10:38 -0700
Subject: [PATCH 03/19] [None][feat] modeling_v2: GLM-5.2 NVFP4 targets for
 SM103

Add the glm_moe_dsa family to modeling_v2 with three targets, each claiming
its own parallel topology through the family's route():

- glm_5_2_nvfp4__sm_103__tp8: tensor-parallel attention and MoE on 8 GPUs;
- glm_5_2_nvfp4__sm_103__tep: tensor-parallel attention with
  expert-parallel experts on 4 or 8 GPUs;
- glm_5_2_nvfp4__sm_103__dep: attention data parallelism with
  expert-parallel experts on 8 or 16 GPUs.

Code the targets share (kernel packages, dispatch, attention, the layer
composition, the speculative worker and the runtime-hook wiring) lives in
models/glm_moe_dsa/shared/. Each target selects the fast paths it serves
through its own plan; the paths default on inside the target and can be
turned off with TRTLLM_STAIRCASE_GLM_LL_DISABLE or per-path
TRTLLM_STAIRCASE_GLM_LL_* variables. Any configuration that no target
claims runs the built-in model.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_experimental/modeling_v2/__init__.py     |    2 +
 .../modeling_v2/_router_index.py              |   22 +-
 .../models/glm_moe_dsa/__init__.py            |    3 +
 .../glm_5_2_nvfp4__sm_103__dep/__init__.py    |    3 +
 .../glm_5_2_nvfp4__sm_103__dep/modeling.py    |  443 ++++
 .../glm_5_2_nvfp4__sm_103__dep/weights.py     |   21 +
 .../glm_5_2_nvfp4__sm_103__tep/__init__.py    |    3 +
 .../glm_5_2_nvfp4__sm_103__tep/modeling.py    |  357 +++
 .../glm_5_2_nvfp4__sm_103__tep/weights.py     |   21 +
 .../glm_5_2_nvfp4__sm_103__tp8/__init__.py    |    3 +
 .../kernels/__init__.py                       |    3 +
 .../cross_layer_expert_prefetch/__init__.py   |   29 +
 .../cross_layer_expert_prefetch/kernels.py    |  405 ++++
 .../spec_experts.py                           |  508 +++++
 .../kernels/draft_routed_moe/__init__.py      |   30 +
 .../kernels/draft_routed_moe/routed_moe.py    |  408 ++++
 .../kernels/hop_routing/__init__.py           |   46 +
 .../kernels/hop_routing/hop.py                |  251 +++
 .../kernels/hop_routing/routed_moe.py         |  316 +++
 .../kernels/routed_moe/__init__.py            |   39 +
 .../kernels/routed_moe/kernel.py              | 1987 +++++++++++++++++
 .../kernels/routed_moe/kernel_xc.py           | 1775 +++++++++++++++
 .../kernels/routed_moe/ptx.py                 |  844 +++++++
 .../kernels/routed_moe/routed_moe.py          |  880 ++++++++
 .../shared_expert_prefetch/__init__.py        |   36 +
 .../shared_expert_prefetch/kernel_xc.py       |  745 ++++++
 .../shared_expert_prefetch/routed_moe.py      |  308 +++
 .../glm_5_2_nvfp4__sm_103__tp8/modeling.py    |  318 +++
 .../glm_5_2_nvfp4__sm_103__tp8/weights.py     |   21 +
 .../modeling_v2/models/glm_moe_dsa/routing.py |  290 +++
 .../models/glm_moe_dsa/shared/__init__.py     |    4 +
 .../models/glm_moe_dsa/shared/attention.py    |  569 +++++
 .../glm_moe_dsa/shared/attention_forward.py   |  744 ++++++
 .../models/glm_moe_dsa/shared/dispatch.py     | 1209 ++++++++++
 .../glm_moe_dsa/shared/kernels/__init__.py    |    3 +
 .../shared/kernels/absorbed_query/__init__.py |   48 +
 .../shared/kernels/absorbed_query/kernels.py  | 1764 +++++++++++++++
 .../kernels/absorbed_query/query_proj.py      |  886 ++++++++
 .../kernels/absorbed_query/shared_indexer.py  | 1348 +++++++++++
 .../absorbed_query/shared_indexer_cute.py     |  945 ++++++++
 .../kernels/attention_output_tail/__init__.py |   46 +
 .../attention_output_tail/cute_fused.py       |  585 +++++
 .../attention_output_tail/cute_oproj.py       |  390 ++++
 .../attention_output_tail/cute_oproj_ring.py  |  461 ++++
 .../kernels/attention_output_tail/into_hop.py |  566 +++++
 .../kernels/attention_output_tail/kernels.py  |  612 +++++
 .../kernels/attention_output_tail/prefetch.py |  217 ++
 .../kernels/attention_output_tail/tail.py     |  628 ++++++
 .../glm_moe_dsa/shared/kernels/contracts.py   |  244 ++
 .../cross_rank_topk_cluster_merge/__init__.py |   39 +
 .../cp_indexer_decode.py                      |  264 +++
 .../cp_workspace.py                           |  247 ++
 .../cross_rank_topk_global_emit/__init__.py   |   60 +
 .../cross_rank_topk_global_emit/bridge.py     |  286 +++
 .../cp_indexer_decode.py                      |  332 +++
 .../cp_workspace.py                           |  246 ++
 .../kernels/decode_prologue/__init__.py       |   35 +
 .../kernels/decode_prologue/prologue.py       |  596 +++++
 .../kernels/deferred_input_prep/__init__.py   |   30 +
 .../kernels/deferred_input_prep/compat.py     |   84 +
 .../kernels/deferred_input_prep/guard.py      |  233 ++
 .../kernels/deferred_input_prep/packet.py     |  398 ++++
 .../deferred_input_prep/paged_mqa_schedule.py |   70 +
 .../deferred_input_prep/sampler_store.py      |  248 ++
 .../kernels/deferred_input_prep/stager.py     |  703 ++++++
 .../shared/kernels/deferred_norm/__init__.py  |   72 +
 .../shared/kernels/deferred_norm/hop.py       |  666 ++++++
 .../kernels/deferred_norm/hop_kernel.py       |  758 +++++++
 .../shared/kernels/deferred_norm/packet.py    |  105 +
 .../kernels/deferred_norm/proj_kernel.py      |  889 ++++++++
 .../kernels/deferred_norm/projection.py       |  305 +++
 .../shared/kernels/deferred_norm/ptx.py       |  619 +++++
 .../kernels/deferred_norm/query_entry.py      |  224 ++
 .../shared/kernels/deferred_norm/scope.py     |  134 ++
 .../deferred_norm_evict_hints/__init__.py     |   62 +
 .../deferred_norm_evict_hints/entry.py        |  122 +
 .../kernels/deferred_norm_evict_hints/hop.py  |  239 ++
 .../deferred_norm_evict_hints/hop_kernel.py   |  814 +++++++
 .../deferred_norm_evict_hints/proj_kernel.py  |  703 ++++++
 .../deferred_norm_evict_hints/projection.py   |  227 ++
 .../kernels/deferred_norm_evict_hints/ptx.py  |  166 ++
 .../shared/kernels/dense_mlp/__init__.py      |   26 +
 .../shared/kernels/dense_mlp/dense_mlp.py     |  152 ++
 .../shared/kernels/draft_front/__init__.py    |   42 +
 .../shared/kernels/draft_front/draft_front.py |  223 ++
 .../shared/kernels/draft_front/kernels.py     |  445 ++++
 .../shared/kernels/draft_front/native_norm.py |  137 ++
 .../kernels/draft_live_rows/__init__.py       |   50 +
 .../shared/kernels/draft_live_rows/hop.py     |  277 +++
 .../shared/kernels/draft_live_rows/step0.py   |  103 +
 .../shared/kernels/draft_loop/__init__.py     |   46 +
 .../shared/kernels/draft_loop/draft_loop.py   |  450 ++++
 .../shared/kernels/draft_loop/draft_step.py   |  497 +++++
 .../shared/kernels/draft_loop/embedding.py    |  175 ++
 .../shared/kernels/draft_loop/verify.py       |  245 ++
 .../shared/kernels/draft_tail/__init__.py     |   31 +
 .../shared/kernels/draft_tail/draft_tail.py   |  564 +++++
 .../fused_kv_indexer_projection/__init__.py   |   38 +
 .../fused_kv_indexer_projection/entry.py      |  208 ++
 .../proj_kernel.py                            |  906 ++++++++
 .../fused_kv_indexer_projection/projection.py |  316 +++
 .../fused_kv_indexer_projection/ptx.py        |  475 ++++
 .../kernels/host_page_retention/__init__.py   |   31 +
 .../host_page_retention/host_policy.py        |  288 +++
 .../kernels/indexer_layer_front/__init__.py   |   55 +
 .../indexer_layer_front/cp_indexer_decode.py  |  248 ++
 .../indexer_layer_front/cp_logits_runner.py   |  217 ++
 .../fp8_paged_mqa_logits_cp.py                | 1844 +++++++++++++++
 .../kernels/indexer_layer_front/kernels.py    |  732 ++++++
 .../kernels/indexer_layer_front/query_proj.py |  403 ++++
 .../kernels/indexer_layer_front/rope.py       |  110 +
 .../shared/kernels/indexer_topk/__init__.py   |   45 +
 .../indexer_topk/cross_rank_indexer_decode.py |  266 +++
 .../indexer_topk/cross_rank_logits_runner.py  |  219 ++
 .../kernels/indexer_topk/cross_rank_policy.py |  143 ++
 .../indexer_topk/cross_rank_workspace.py      |  212 ++
 .../fp8_paged_mqa_logits_cross_rank.py        | 1806 +++++++++++++++
 .../indexer_topk/fp8_paged_mqa_logits_rows.py | 1747 +++++++++++++++
 .../fp8_paged_mqa_logits_rowsplit.py          |  773 +++++++
 .../kernels/indexer_topk/indexer_decode.py    |  296 +++
 .../kernels/indexer_topk/logits_runner.py     |  237 ++
 .../shared/kernels/indexer_topk/policy.py     |  240 ++
 .../shared/kernels/indexer_topk/region.py     |   99 +
 .../kernels/indexer_topk/topk_select.py       |  475 ++++
 .../kernels/oproj_weight_staging/__init__.py  |   48 +
 .../oproj_weight_staging/cute_oproj_stage.py  |  841 +++++++
 .../kernels/oproj_weight_staging/kernels.py   |  389 ++++
 .../kernels/oproj_weight_staging/tail.py      |  208 ++
 .../kernels/post_attention_hop/__init__.py    |   74 +
 .../post_attention_hop/from_partials.py       |  207 ++
 .../shared/kernels/post_attention_hop/hop.py  |  430 ++++
 .../shared/kernels/post_moe_hop/__init__.py   |   45 +
 .../kernels/post_moe_hop/cute_kernel.py       | 1078 +++++++++
 .../shared/kernels/post_moe_hop/health.py     |  119 +
 .../shared/kernels/post_moe_hop/hop.py        |  785 +++++++
 .../shared/kernels/post_moe_hop/kernels.py    |  446 ++++
 .../shared/kernels/post_moe_hop/workspace.py  |  200 ++
 .../shared/kernels/rope_kv_append/__init__.py |   70 +
 .../shared/kernels/rope_kv_append/ops.py      |  517 +++++
 .../shared/kernels/router_gemm/__init__.py    |   40 +
 .../shared/kernels/router_gemm/hop.py         |  146 ++
 .../kernels/router_gemm/router_kernel.py      |  571 +++++
 .../kernels/sparse_mla_decode/__init__.py     |   31 +
 .../kernels/sparse_mla_decode/sparse_mla.py   |  335 +++
 .../models/glm_moe_dsa/shared/layers.py       | 1189 ++++++++++
 .../models/glm_moe_dsa/shared/runtime.py      |  103 +
 .../models/glm_moe_dsa/shared/spec_worker.py  |  621 ++++++
 147 files changed, 55745 insertions(+), 2 deletions(-)
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/spec_experts.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/routed_moe.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/routed_moe.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel_xc.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/ptx.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/kernel_xc.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/routed_moe.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention_forward.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/dispatch.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/query_proj.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer_cute.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/prefetch.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_indexer_decode.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_workspace.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/bridge.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_indexer_decode.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_workspace.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/prologue.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/guard.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/paged_mqa_schedule.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/sampler_store.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/stager.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/packet.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/proj_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/projection.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/ptx.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/query_entry.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/scope.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/entry.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/proj_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/projection.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/ptx.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/draft_front.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/native_norm.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/step0.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_loop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_step.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/embedding.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/verify.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/entry.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/projection.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/ptx.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/host_policy.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_indexer_decode.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_logits_runner.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/query_proj.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/rope.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_indexer_decode.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_logits_runner.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_policy.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_workspace.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/indexer_decode.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/logits_runner.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/policy.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/region.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/topk_select.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/tail.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/from_partials.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/cute_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/health.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/kernels.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/workspace.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/ops.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/hop.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/router_kernel.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/__init__.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/layers.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/runtime.py
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/spec_worker.py

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py
index c5140e7c912b..7dc4b711fd91 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py
@@ -32,6 +32,7 @@
     ModelingV2Mode,
     NoTarget,
     modeling_v2_resolve,
+    resolved_mode,
     step_contract_enabled,
 )
 
@@ -41,5 +42,6 @@
     "ModelingV2Mode",
     "NoTarget",
     "modeling_v2_resolve",
+    "resolved_mode",
     "step_contract_enabled",
 ]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py
index b0b6e78ef8d8..4d6bb76f1ef4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py
@@ -37,7 +37,7 @@
 import enum
 import os
 from importlib import import_module
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
 from tensorrt_llm.logger import logger
 
@@ -51,6 +51,7 @@
 MODELING_V2_ROUTERS = {
     "GptOssForCausalLM": "models.gpt_oss.routing",
     "DeepseekV3ForCausalLM": "models.deepseek_v3.routing",
+    "GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing",
     # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the
     # config loader's composite branch writes the same name back.
     "Qwen4ExpForConditionalGeneration": "models.qwen4_exp.routing",
@@ -104,6 +105,20 @@ class NoTarget(Exception):
     """
 
 
+# The mode each target class was last resolved under (the loader builds the class right after resolving it).
+_RESOLVED_MODES: Dict[type, ModelingV2Mode] = {}
+
+
+def resolved_mode(target: type) -> ModelingV2Mode:
+    """The mode ``modeling_v2_resolve`` returned ``target`` under; ``auto`` for a class built without it.
+
+    A target reads it at construction: under ``require`` its in-target dispatch
+    raises where it would otherwise fall back to a built-in path inside the
+    shapes it claims.
+    """
+    return _RESOLVED_MODES.get(target, ModelingV2Mode.AUTO)
+
+
 def routing_module(arch: str):
     """Import and return the routing module for ``arch``, or None."""
     name = MODELING_V2_ROUTERS.get(arch)
@@ -146,7 +161,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti
             raise NoTarget(
                 f"no routing module for this architecture; routed architectures: {known}"
             )
-        return routing.route(
+        target = routing.route(
             pretrained_config, config.mapping, _device_compute_capability(), config.spec_config
         )
     except NoTarget as why:
@@ -156,3 +171,6 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti
             f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation"
         )
         return None
+    if target is not None:
+        _RESOLVED_MODES[target] = mode
+    return target
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py
new file mode 100644
index 000000000000..638d9884388f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GlmMoeDsaForCausalLM targets."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py
new file mode 100644
index 000000000000..dd69779ecd5e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-5.2-NVFP4 / sm_103 / dep."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py
new file mode 100644
index 000000000000..5b2625da55cb
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py
@@ -0,0 +1,443 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / dep.
+
+GLM-5.2, modelopt NVFP4 export: 78 layers, hidden 6144, 64 query heads, MLA with
+DSA sparse attention (index top-k 2048), layers 0-2 dense and 3-77 MoE (256
+routed experts at top-8 plus a shared expert). Layer 78 is the checkpoint's bf16
+MTP module, driven by the one-model MTP-Eagle loop. `dep` is attention data
+parallelism over eight or sixteen ranks (each rank runs the whole attention on its
+own requests) and the experts expert-parallel over the same ranks.
+
+``route`` claims three deployments by the facts it reads: the rank count, whether
+the LM head is tensor-parallel in attention DP, and the draft tokens per request.
+Two more facts fix what the fast paths can serve there, the KV-cache dtype and the
+MoE backend, and only the model config carries them, so this target checks them
+first, at construction (``CELLS``, ``target_paths``). Under ``require`` a
+deployment outside the cells raises; under ``auto`` it serves no fast path and
+runs the built-in forward.
+
+Under attention DP the fast paths that take this target's shapes are rank-local:
+the attention's rope, KV append and top-k bridge with an FP8 KV cache
+(``rope_kv_append``), the rope of the full-indexer layers' query front
+(``indexer_layer_front``), the draft loop (``draft_loop``), the decode step's
+prologue (``decode_prologue``) and the host page retention policy
+(``host_page_retention``). The plan of a cell serves those and the bases they
+extend; the bases decline every call under attention DP. The other fast paths of
+the family take tensor-parallel attention, so this target serves none of them: its
+decoder layers, MoE modules and draft layer stay the built-in classes, which run
+the built-in model's attention-DP communication on every step.
+
+The shell is the built-in ``GlmMoeDsaForCausalLM``: the same module tree, the
+same weights and the same engine preferences (KV cache manager, KV transceiver,
+sparse-attention defaults), so a built-in context server and this target's
+generation server exchange the same KV and indexer-K layout. Construction runs,
+in order:
+
+1. the cell check (``target_paths``), from the model config alone, so every rank
+   decides alike before any collective;
+2. the rank-uniform check of the knobs (``dispatch.check_rank_uniform``), before
+   every collective the construction below makes;
+3. the dispatch plan (``dispatch.resolve``) of the paths the cell serves, before
+   the built-in model is built, so a knob value no rank can parse fails early;
+4. the plan is installed (``dispatch.install``): from here on the fast paths read
+   it;
+5. the built-in model, then the core: its text model adopted as
+   ``GlmMoeDsaDepModelingV2``, whose forward hands each step to
+   ``PrefillTarget`` or ``DecodeTarget``;
+6. ``adopt_model``, which re-classes the attention modules of every layer
+   (``shared.attention``) when the plan runs an attention path, and the
+   speculative worker (``shared.spec_worker``) when it runs the draft loop;
+7. the rank-uniform check of the attached fast-path state
+   (``dispatch.check_rank_uniform_adopted``, ``shared.layers.adopted_state``).
+
+The step-loop hooks the target declares to the engine (``runtime_hooks``,
+``shared.runtime.make_runtime_hooks``) are built later, on the engine's first
+read, after the model is loaded.
+
+The rank-uniform check of step 2 also covers the shared switches the kernels
+read and the modeling_v2 mode (``dispatch.rank_uniform_inputs``).
+
+Every step runs under the phase of the target that takes it
+(``dispatch.step``): the shell's forward enters it around the built-in forward,
+so the model, the speculative worker and its draft loop all run in it. A
+prefill step runs the built-in decoder forward and the built-in draft loop; a
+decode step runs ``DecodeTarget``'s loop over the built-in layers, the attention
+modules' decode composition and the worker's draft loop. Each rank picks its
+phase from its own requests; the layers' communication is the built-in one in
+both phases.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Mapping, Optional, Tuple
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
+    ModelingV2Mode,
+    resolved_mode,
+)
+from tensorrt_llm._torch.attention.attention import maybe_allgather_for_helix_cp
+from tensorrt_llm._torch.attention.backends import AttentionMetadata
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
+from tensorrt_llm._torch.speculative import SpecMetadata
+from tensorrt_llm.logger import logger
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+from ..shared import dispatch, runtime
+from ..shared.attention import GlmMLA
+from ..shared.layers import adopted_state
+from ..shared.spec_worker import GlmMtpWorker
+from . import weights as _weights
+
+if TYPE_CHECKING:
+    from ..shared.kernels.decode_prologue import DecodePrologueHook
+
+
+# The paths that engage under attention DP with an FP8 KV cache: the attention's rope, KV append
+# and top-k bridge, the full-indexer layers' query-front rope, the draft loop, the decode
+# prologue and the host page retention policy.
+_FP8_KV_ENGAGED = frozenset(
+    {
+        "rope_kv_append",
+        "indexer_layer_front",
+        "draft_loop",
+        "decode_prologue",
+        "host_page_retention",
+    }
+)
+# The bases these paths extend. Under attention DP their kernels decline every call: the query
+# fronts on every layer, the top-k selection on every step, and the post-attention hop is never
+# built (the built-in layers run no post-attention all-reduce under attention DP).
+_FP8_KV_BASES = frozenset(
+    {
+        "absorbed_query",
+        "draft_absorbed_query",
+        "indexer_topk",
+        "cross_rank_topk",
+        "post_attention_hop",
+    }
+)
+# With an NVFP4 KV cache the attention paths are outside the shapes their kernels take; the
+# draft loop, the decode prologue and the host page retention policy engage.
+_NVFP4_KV_ENGAGED = frozenset({"draft_loop", "decode_prologue", "host_page_retention"})
+_NVFP4_KV_BASES = frozenset({"post_attention_hop"})
+_NO_LAYERS: frozenset = frozenset()
+
+#: What this target serves with an FP8 KV cache. The query fronts decline attention DP on every
+#: layer, which is outside their shapes here.
+FP8_KV_PATHS = dispatch.TargetPaths(
+    served=_FP8_KV_ENGAGED | _FP8_KV_BASES,
+    layer_cells={
+        **dispatch.LAYER_CELLS,
+        "absorbed_query": {"layer": _NO_LAYERS},
+        "draft_absorbed_query": {"layer": _NO_LAYERS},
+        "indexer_layer_front": {"target_layer": _NO_LAYERS, "draft_layer": _NO_LAYERS},
+    },
+)
+#: What this target serves with an NVFP4 KV cache.
+NVFP4_KV_PATHS = dispatch.TargetPaths(served=_NVFP4_KV_ENGAGED | _NVFP4_KV_BASES)
+#: What a deployment outside the cells serves under ``auto``: no fast path.
+NO_PATHS = dispatch.TargetPaths(served=frozenset())
+
+
+@dataclass(frozen=True)
+class Cell:
+    """One deployment this target serves: the facts ``route`` does not read, and the plan."""
+
+    kv_cache: QuantAlgo
+    moe_backend: str
+    paths: dispatch.TargetPaths
+
+
+#: The deployments this target serves, per (ranks, LM head tensor-parallel in attention DP,
+#: draft tokens per request), the facts ``route`` claims them by.
+CELLS: Mapping[Tuple[int, bool, int], Cell] = {
+    (16, False, 3): Cell(QuantAlgo.FP8, "TRTLLM", FP8_KV_PATHS),
+    (8, False, 3): Cell(QuantAlgo.FP8, "CUTEDSL", FP8_KV_PATHS),
+    (8, True, 5): Cell(QuantAlgo.NVFP4, "CUTEDSL", NVFP4_KV_PATHS),
+}
+
+
+def target_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.TargetPaths:
+    """What this target serves on ``model_config``'s deployment: its cell's paths, or none.
+
+    Reads only the model config, so every rank decides alike.
+
+    Raises:
+        ValueError: under ``require``, if the deployment is not one of ``CELLS`` (its KV-cache
+            dtype or MoE backend is not its cell's).
+    """
+    mapping = model_config.mapping
+    spec_config = model_config.spec_config
+    key = (
+        mapping.world_size,
+        bool(mapping.enable_lm_head_tp_in_adp),
+        getattr(spec_config, "max_draft_len", None),
+    )
+    kv_cache = model_config.quant_config.kv_cache_quant_algo
+    cell = CELLS.get(key)
+    if cell is not None and (kv_cache, model_config.moe_backend) == (
+        cell.kv_cache,
+        cell.moe_backend,
+    ):
+        return cell.paths
+    served = "; ".join(
+        f"{ranks} ranks, LM head tensor-parallel in attention DP {lm_head_tp}, {draft} draft "
+        f"tokens: KV cache {cell.kv_cache.value}, MoE backend {cell.moe_backend}"
+        for (ranks, lm_head_tp, draft), cell in CELLS.items()
+    )
+    words = (
+        f"{key[0]} ranks, LM head tensor-parallel in attention DP {key[1]}, {key[2]} draft "
+        f"tokens, KV cache {kv_cache}, MoE backend {model_config.moe_backend} is not a "
+        f"deployment this target serves ({served})"
+    )
+    if mode is ModelingV2Mode.REQUIRE:
+        raise ValueError(f"modeling_v2='require': {words}")
+    logger.warning(f"[glm_moe_dsa] dispatch: {words}; no fast path runs")
+    return NO_PATHS
+
+
+def adopt_model(model: "ModelingV2Glm52Nvfp4Sm103Dep", plan: dispatch.Plan) -> None:
+    """Re-class the attention modules of every layer and the speculative worker, as far as
+    ``plan`` runs their fast paths.
+
+    The attention modules of the decoder layers and of the draft layer, in index order, when
+    the plan runs the rope and KV append (the attention paths that engage under attention DP
+    extend it); then the speculative worker, which the built-in model builds after the draft
+    layer, when the plan runs the draft loop. The layers, their MoE modules and the draft
+    layer stay the built-in classes. The fast paths read the installed plan, which is ``plan``.
+
+    Raises:
+        TypeError: if an attention module or the speculative worker is not exactly the
+            built-in class.
+    """
+    if dispatch.is_enabled("rope_kv_append"):
+        for layer in model.model.layers:
+            GlmMLA.adopt(layer.self_attn)
+    if model.spec_worker is not None and dispatch.is_enabled("draft_loop"):
+        GlmMtpWorker.adopt(model.spec_worker)
+
+
+class GlmMoeDsaDepModelingV2(DeepseekV3Model):
+    """The text model: the built-in ``DeepseekV3Model``, whose forward picks a target per step.
+
+    Never constructed directly: ``adopt`` re-classes the built-in model the shell
+    has already built, so the module tree, the parameters and their names stay
+    the built-in model's. The class adds methods and plain attributes only.
+    """
+
+    _prefill: "PrefillTarget"
+    _decode: "DecodeTarget"
+    # The decode prologue the step loop runs, set when the target's runtime hooks are built;
+    # the decode target takes the embedding rows of its launch.
+    decode_prologue: Optional["DecodePrologueHook"] = None
+
+    @classmethod
+    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaDepModelingV2":
+        """Re-class the built-in text model as this core and build its two targets.
+
+        Raises:
+            TypeError: if ``model`` is not exactly the built-in ``DeepseekV3Model``.
+        """
+        if type(model) is not DeepseekV3Model:
+            raise TypeError(
+                f"{cls.__name__} adopts the built-in DeepseekV3Model, got {type(model).__name__}"
+            )
+        model.__class__ = cls
+        model._prefill = PrefillTarget(model)
+        model._decode = DecodeTarget(model)
+        return model
+
+    def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
+        """Which target runs this step.
+
+        One host-side integer read and nothing else: a decode-only CUDA graph is a
+        per-capture constant environment. A mixed batch -- context and generation
+        rows in one step -- goes to prefill, so prefill is the general case and
+        decode the specialization that may assume it has no context rows.
+        """
+        return self._decode if attn_metadata.num_contexts == 0 else self._prefill
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self._select_target(attn_metadata).forward(
+            attn_metadata,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+
+class PrefillTarget:
+    """The general case: context rows, and possibly generation rows beside them.
+
+    Runs the built-in decoder forward.
+    """
+
+    phase = dispatch.Phase.PREFILL
+
+    def __init__(self, core: GlmMoeDsaDepModelingV2) -> None:
+        self.core = core
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return DeepseekV3Model.forward(
+            self.core,
+            attn_metadata,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+
+class DecodeTarget:
+    """The specialization: routed to only when there are no context rows.
+
+    Runs this target's decoder loop.
+    """
+
+    phase = dispatch.Phase.DECODE
+
+    def __init__(self, core: GlmMoeDsaDepModelingV2) -> None:
+        self.core = core
+
+    def forward(
+        self,
+        attn_metadata: AttentionMetadata,
+        input_ids: Optional[torch.IntTensor] = None,
+        position_ids: Optional[torch.IntTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """The decoder layers of the text model on a decode step.
+
+        The built-in loop with the decode prologue's embedding. Each layer is the built-in
+        decoder layer, called as a module, so its forward hooks run; its attention runs the
+        attention modules' decode composition under the step's phase. The draft layer is not
+        run here; the last decoder layer applies the final norm.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        prologue = self.core.decode_prologue
+        if inputs_embeds is None and prologue is not None:
+            # The rows the decode prologue launch embedded (None unless that launch ran for
+            # exactly these ids and this table).
+            inputs_embeds = prologue.take_embedding(
+                attn_metadata, input_ids, self.core.embed_tokens
+            )
+        if inputs_embeds is None:
+            inputs_embeds = self.core.embed_tokens(input_ids)
+
+        hidden_states = inputs_embeds
+        residual = None
+
+        for idx, decoder_layer in enumerate(self.core.layers[: self.core.num_hidden_layers]):
+            hidden_states, residual = decoder_layer(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                residual=residual,
+                spec_metadata=spec_metadata,
+            )
+
+        hidden_states = maybe_allgather_for_helix_cp(
+            hidden_states, attn_metadata, self.core.mapping_with_cp
+        )
+        return hidden_states
+
+
+class ModelingV2Glm52Nvfp4Sm103Dep(GlmMoeDsaForCausalLM):
+    """GLM-5.2-NVFP4 on sm_103 at dep: the built-in causal LM with this target's core and dispatch."""
+
+    def __init__(self, model_config: ModelConfig):
+        world_size = model_config.mapping.world_size
+        mode = resolved_mode(type(self))
+        paths = target_paths(model_config, mode)
+        knobs = dispatch.read_knobs()
+        dispatch.check_rank_uniform(dispatch.rank_uniform_inputs(knobs, mode), world_size)
+        plan = dispatch.resolve(mode, dict(knobs), paths)
+        dispatch.install(plan)
+        super().__init__(model_config)
+        GlmMoeDsaDepModelingV2.adopt(self.model)
+        self.plan = plan
+        adopt_model(self, self.plan)
+        dispatch.check_rank_uniform_adopted(adopted_state(self), world_size)
+
+    @property
+    def runtime_hooks(self) -> RuntimeHooks:
+        """The step-loop hooks this target declares, built on the engine's first read.
+
+        Built after the model is loaded: a hook may hold pinned host memory, which model
+        construction (under meta-device initialization) must not allocate.
+
+        Raises:
+            RuntimeError: if building the hooks raises ``AttributeError``, which
+                ``runtime_hooks_of`` would otherwise read as "no hooks declared".
+        """
+        hooks = self.__dict__.get("_runtime_hooks")
+        if hooks is None:
+            try:
+                hooks = runtime.make_runtime_hooks(self)
+            except AttributeError as exc:
+                raise RuntimeError(
+                    f"building the step-loop hooks of {type(self).__name__} failed: {exc}"
+                ) from exc
+            self._runtime_hooks = hooks
+        return hooks
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        return_context_logits: bool = False,
+        spec_metadata=None,
+        resource_manager=None,
+        **kwargs,
+    ):
+        """The built-in forward under the phase of the target that takes this step.
+
+        The scope covers the text model, the speculative worker and its draft loop.
+        """
+        with dispatch.step(self.model._select_target(attn_metadata).phase):
+            return super().forward(
+                attn_metadata,
+                input_ids=input_ids,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                return_context_logits=return_context_logits,
+                spec_metadata=spec_metadata,
+                resource_manager=resource_manager,
+                **kwargs,
+            )
+
+    def load_weights(self, weights):
+        _weights.load(self, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py
new file mode 100644
index 000000000000..b9046fba42dd
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py
@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The weights this target expects: the GLM-5.2 NVFP4 checkpoint as the built-in model loads it.
+
+The target's module tree is the built-in model's, so the built-in loader maps
+every checkpoint tensor, including the bf16 MTP layer and the NVFP4 scales.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+
+if TYPE_CHECKING:
+    from .modeling import ModelingV2Glm52Nvfp4Sm103Dep
+
+
+def load(model: "ModelingV2Glm52Nvfp4Sm103Dep", weights: Any) -> None:
+    """Load the checkpoint ``weights`` into ``model`` with the built-in loader."""
+    GlmMoeDsaForCausalLM.load_weights(model, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py
new file mode 100644
index 000000000000..bfe54ffe7404
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-5.2-NVFP4 / sm_103 / tep."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py
new file mode 100644
index 000000000000..d4c83cbaec3c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py
@@ -0,0 +1,357 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / tep.
+
+GLM-5.2, modelopt NVFP4 export: 78 layers, hidden 6144, 64 query heads, MLA with
+DSA sparse attention (index top-k 2048), layers 0-2 dense and 3-77 MoE (256
+routed experts at top-8 plus a shared expert). Layer 78 is the checkpoint's bf16
+MTP module, driven by the one-model MTP-Eagle loop at five draft tokens. `tep`
+is the attention heads tensor-parallel over four or eight ranks and the experts
+expert-parallel over the same ranks: each rank holds whole experts and runs the
+built-in expert-parallel MoE on them.
+
+The fast paths this target serves depend on its tensor-parallel size
+(``TARGET_PATHS``): at eight ranks the attention runs at the same shapes as on
+the tp8 target, so it serves every fast path of the family except the routed-MoE
+ones, which take only tensor-parallel experts; at four ranks it serves the paths
+whose kernels take four ranks (the indexer, the top-k selection, the router GEMM,
+the dense MLP, the draft loop and the step-loop paths) and the bases they extend.
+The layers run the family's composition (``shared.layers``), in which the block of
+a path this target does not serve never runs.
+
+The shell is the built-in ``GlmMoeDsaForCausalLM``: the same module tree, the
+same weights and the same engine preferences (KV cache manager, KV transceiver,
+sparse-attention defaults), so a built-in context server and this target's
+generation server exchange the same KV and indexer-K layout. Construction runs,
+in order:
+
+1. the rank-uniform check of the knobs (``dispatch.check_rank_uniform``), first,
+   so it runs before every collective the construction below makes;
+2. the dispatch plan (``dispatch.resolve``) of the paths this target serves at its
+   tensor-parallel size, before the built-in model is built, so a knob value no
+   rank can parse fails early;
+3. the plan is installed (``dispatch.install``): from here on the fast paths read
+   it;
+4. the built-in model, then the core: its text model adopted as
+   ``GlmMoeDsaTepModelingV2``, whose forward hands each step to
+   ``PrefillTarget`` or ``DecodeTarget``;
+5. ``adopt_model`` (``shared.layers``), which re-classes the decoder layers,
+   the draft layer, their attention modules (``shared.attention``) and their MoE
+   modules as the family's layer classes and attaches the fast-path objects the
+   plan enables to them, then re-classes the speculative worker
+   (``shared.spec_worker``). The attention backends build their FMHA libraries
+   after construction, in ``create_weights``, each with the libraries adoption
+   declared on it first;
+6. the rank-uniform check of that state (``dispatch.check_rank_uniform_adopted``).
+
+The step-loop hooks the target declares to the engine (``runtime_hooks``,
+``shared.runtime.make_runtime_hooks``) are built later, on the engine's first
+read, after the model is loaded.
+
+The rank-uniform check of step 1 also covers the shared switches the kernels
+read and the modeling_v2 mode (``dispatch.rank_uniform_inputs``).
+
+Every step runs under the phase of the target that takes it
+(``dispatch.step``): the shell's forward enters it around the built-in forward,
+so the model, the speculative worker and its draft loop all run in it. A
+prefill step runs the built-in decoder forward, the built-in layers and the
+built-in draft loop; a decode step runs ``DecodeTarget``'s loop, the layers'
+decode composition and the worker's draft loop.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import resolved_mode
+from tensorrt_llm._torch.attention.attention import maybe_allgather_for_helix_cp
+from tensorrt_llm._torch.attention.backends import AttentionMetadata
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
+from tensorrt_llm._torch.speculative import SpecMetadata
+
+from ..shared import dispatch, runtime
+from ..shared.layers import adopt_model, adopted_state
+from . import weights as _weights
+
+if TYPE_CHECKING:
+    from ..shared.kernels.decode_prologue import DecodePrologueHook
+
+
+# The routed-MoE paths take tensor-parallel experts (every expert on every rank); with the experts
+# expert-parallel their kernels decline every call, so this target serves none of them.
+_ROUTED_MOE_PATHS = frozenset(
+    {
+        "routed_moe",
+        "routed_moe_into_hop",
+        "hop_routing",
+        "shared_expert_prefetch",
+        "draft_routed_moe",
+        "cross_layer_expert_prefetch",
+    }
+)
+# At four ranks the attention operands, the attention output and the post-MoE hop are outside
+# the kernels' shapes (sixteen local heads, a four-rank output projection, a single-node group);
+# these paths take four ranks.
+_TEP4_PATHS = frozenset(
+    {
+        "indexer_layer_front",
+        "rope_kv_append",
+        "indexer_topk",
+        "cross_rank_topk",
+        "cross_rank_topk_cluster_merge",
+        "cross_rank_topk_global_emit",
+        "router_gemm",
+        "dense_mlp",
+        "draft_front",
+        "draft_tail",
+        "draft_loop",
+        "decode_prologue",
+        "deferred_input_prep",
+        "sampler_store",
+        "execution_stream_loop",
+        "host_page_retention",
+    }
+)
+# The bases these paths extend at four ranks. Their kernels decline every layer there: the query
+# fronts on their layer cells, the post-attention hop on its native all-reduce, which a
+# single-node group runs without multi-node NVLink.
+_TEP4_BASES = frozenset({"absorbed_query", "draft_absorbed_query", "post_attention_hop"})
+_NO_LAYERS: frozenset = frozenset()
+
+#: What this target serves at each tensor-parallel size.
+TARGET_PATHS = {
+    4: dispatch.TargetPaths(
+        served=_TEP4_PATHS | _TEP4_BASES,
+        layer_cells={
+            **dispatch.LAYER_CELLS,
+            "absorbed_query": {"layer": _NO_LAYERS},
+            "draft_absorbed_query": {"layer": _NO_LAYERS},
+            "indexer_layer_front": {"target_layer": _NO_LAYERS, "draft_layer": _NO_LAYERS},
+        },
+        uncovered={"post_attention_hop": ("does not use MNNVL",)},
+    ),
+    8: dispatch.TargetPaths(served=frozenset(dispatch.FAST_PATH_NAMES) - _ROUTED_MOE_PATHS),
+}
+
+
+class GlmMoeDsaTepModelingV2(DeepseekV3Model):
+    """The text model: the built-in ``DeepseekV3Model``, whose forward picks a target per step.
+
+    Never constructed directly: ``adopt`` re-classes the built-in model the shell
+    has already built, so the module tree, the parameters and their names stay
+    the built-in model's. The class adds methods and plain attributes only.
+    """
+
+    _prefill: "PrefillTarget"
+    _decode: "DecodeTarget"
+    # The decode prologue the step loop runs, set when the target's runtime hooks are built;
+    # the decode target takes the embedding rows of its launch.
+    decode_prologue: Optional["DecodePrologueHook"] = None
+
+    @classmethod
+    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaTepModelingV2":
+        """Re-class the built-in text model as this core and build its two targets.
+
+        Raises:
+            TypeError: if ``model`` is not exactly the built-in ``DeepseekV3Model``.
+        """
+        if type(model) is not DeepseekV3Model:
+            raise TypeError(
+                f"{cls.__name__} adopts the built-in DeepseekV3Model, got {type(model).__name__}"
+            )
+        model.__class__ = cls
+        model._prefill = PrefillTarget(model)
+        model._decode = DecodeTarget(model)
+        return model
+
+    def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
+        """Which target runs this step.
+
+        One host-side integer read and nothing else: a decode-only CUDA graph is a
+        per-capture constant environment. A mixed batch -- context and generation
+        rows in one step -- goes to prefill, so prefill is the general case and
+        decode the specialization that may assume it has no context rows.
+        """
+        return self._decode if attn_metadata.num_contexts == 0 else self._prefill
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self._select_target(attn_metadata).forward(
+            attn_metadata,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+
+class PrefillTarget:
+    """The general case: context rows, and possibly generation rows beside them.
+
+    Runs the built-in decoder forward.
+    """
+
+    phase = dispatch.Phase.PREFILL
+
+    def __init__(self, core: GlmMoeDsaTepModelingV2) -> None:
+        self.core = core
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return DeepseekV3Model.forward(
+            self.core,
+            attn_metadata,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+
+class DecodeTarget:
+    """The specialization: routed to only when there are no context rows.
+
+    Runs this target's decoder loop.
+    """
+
+    phase = dispatch.Phase.DECODE
+
+    def __init__(self, core: GlmMoeDsaTepModelingV2) -> None:
+        self.core = core
+
+    def forward(
+        self,
+        attn_metadata: AttentionMetadata,
+        input_ids: Optional[torch.IntTensor] = None,
+        position_ids: Optional[torch.IntTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """The decoder layers of the text model on a decode step.
+
+        The built-in loop with the decode prologue's embedding.
+        Each layer is called as a module, so its forward hooks run, and runs its decode
+        composition under the step's phase. The draft layer is not run here; the last decoder
+        layer applies the final norm.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        prologue = self.core.decode_prologue
+        if inputs_embeds is None and prologue is not None:
+            # The rows the decode prologue launch embedded (None unless that launch ran for
+            # exactly these ids and this table).
+            inputs_embeds = prologue.take_embedding(
+                attn_metadata, input_ids, self.core.embed_tokens
+            )
+        if inputs_embeds is None:
+            inputs_embeds = self.core.embed_tokens(input_ids)
+
+        hidden_states = inputs_embeds
+        residual = None
+
+        for idx, decoder_layer in enumerate(self.core.layers[: self.core.num_hidden_layers]):
+            hidden_states, residual = decoder_layer(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                residual=residual,
+                spec_metadata=spec_metadata,
+            )
+
+        hidden_states = maybe_allgather_for_helix_cp(
+            hidden_states, attn_metadata, self.core.mapping_with_cp
+        )
+        return hidden_states
+
+
+class ModelingV2Glm52Nvfp4Sm103Tep(GlmMoeDsaForCausalLM):
+    """GLM-5.2-NVFP4 on sm_103 at tep: the built-in causal LM with this target's core and dispatch."""
+
+    def __init__(self, model_config: ModelConfig):
+        world_size = model_config.mapping.world_size
+        knobs = dispatch.read_knobs()
+        mode = resolved_mode(type(self))
+        dispatch.check_rank_uniform(dispatch.rank_uniform_inputs(knobs, mode), world_size)
+        plan = dispatch.resolve(mode, dict(knobs), TARGET_PATHS[model_config.mapping.tp_size])
+        dispatch.install(plan)
+        super().__init__(model_config)
+        GlmMoeDsaTepModelingV2.adopt(self.model)
+        self.plan = plan
+        adopt_model(self, self.plan)
+        dispatch.check_rank_uniform_adopted(adopted_state(self), world_size)
+
+    @property
+    def runtime_hooks(self) -> RuntimeHooks:
+        """The step-loop hooks this target declares, built on the engine's first read.
+
+        Built after the model is loaded: a hook may hold pinned host memory, which model
+        construction (under meta-device initialization) must not allocate.
+
+        Raises:
+            RuntimeError: if building the hooks raises ``AttributeError``, which
+                ``runtime_hooks_of`` would otherwise read as "no hooks declared".
+        """
+        hooks = self.__dict__.get("_runtime_hooks")
+        if hooks is None:
+            try:
+                hooks = runtime.make_runtime_hooks(self)
+            except AttributeError as exc:
+                raise RuntimeError(
+                    f"building the step-loop hooks of {type(self).__name__} failed: {exc}"
+                ) from exc
+            self._runtime_hooks = hooks
+        return hooks
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        return_context_logits: bool = False,
+        spec_metadata=None,
+        resource_manager=None,
+        **kwargs,
+    ):
+        """The built-in forward under the phase of the target that takes this step.
+
+        The scope covers the text model, the speculative worker and its draft loop.
+        """
+        with dispatch.step(self.model._select_target(attn_metadata).phase):
+            return super().forward(
+                attn_metadata,
+                input_ids=input_ids,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                return_context_logits=return_context_logits,
+                spec_metadata=spec_metadata,
+                resource_manager=resource_manager,
+                **kwargs,
+            )
+
+    def load_weights(self, weights):
+        _weights.load(self, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py
new file mode 100644
index 000000000000..09b1849fe5f4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py
@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The weights this target expects: the GLM-5.2 NVFP4 checkpoint as the built-in model loads it.
+
+The target's module tree is the built-in model's, so the built-in loader maps
+every checkpoint tensor, including the bf16 MTP layer and the NVFP4 scales.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+
+if TYPE_CHECKING:
+    from .modeling import ModelingV2Glm52Nvfp4Sm103Tep
+
+
+def load(model: "ModelingV2Glm52Nvfp4Sm103Tep", weights: Any) -> None:
+    """Load the checkpoint ``weights`` into ``model`` with the built-in loader."""
+    GlmMoeDsaForCausalLM.load_weights(model, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py
new file mode 100644
index 000000000000..87156adf61ff
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-5.2-NVFP4 / sm_103 / tp8."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py
new file mode 100644
index 000000000000..94b44234a92b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Kernels of the GLM-5.2-NVFP4 sm_103 tp8 target, one package per fast path."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/__init__.py
new file mode 100644
index 000000000000..555ee2664d03
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/__init__.py
@@ -0,0 +1,29 @@
+# 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.
+"""cross_layer_expert_prefetch: cross-layer speculative routed-expert L2 prefetch.
+
+Registered entry points: ``link_layers(layers, num_hidden_layers)`` (from the
+model's ``setup_aliases``) and ``join()`` (end of the target model's forward).
+Per-layer hooks are methods of the ``ExpertPrefetchLink`` bound to each MoE layer
+as ``_expert_prefetch_link``.
+"""
+
+from ....shared import dispatch
+from .spec_experts import ExpertPrefetchLink, config, join, link_layers
+
+dispatch.register("cross_layer_expert_prefetch", "link_layers", link_layers)
+dispatch.register("cross_layer_expert_prefetch", "join", join)
+
+__all__ = ["ExpertPrefetchLink", "config", "join", "link_layers"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/kernels.py
new file mode 100644
index 000000000000..ab6a84c51599
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/kernels.py
@@ -0,0 +1,405 @@
+# 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.
+"""Triton kernels of cross_layer_expert_prefetch.
+
+``rank_and_stats``: one launch over the 256 experts of one MoE layer, with no
+sequential loop over experts.
+
+* Prediction (``NE / BE`` programs, one block of ``BE`` experts each): from
+  approximate router logits of the next layer (``h_L`` times the next layer's
+  gate), the score ``sigmoid(logit) + bias`` selects each live row's top-8
+  experts; experts are ranked by how many rows select them, then by their
+  maximum score over the rows, then by ID. Each program computes the rank of
+  its own experts by pairwise comparison against all 256 and writes the IDs
+  ranked below ``KMAX`` to their rank slot.
+* Statistics (one more program): from this layer's router logits (the ones
+  its routed MoE uses) and the ranking predicted for this layer by the
+  previous layer, the program adds to this layer's counters: calls, distinct
+  experts D, and the number of the D experts among the first K predicted, for
+  each K of ``STAT_KS``.
+
+``gemm_rank``: the prediction in one launch, without the router GEMM: each of
+``NE / 16`` programs computes the logits of 16 experts (``tl.dot`` over the
+hidden size) and their scores, publishes them, and the last program to arrive
+ranks all experts in the same order as above, with one sort of a unique key.
+
+``delay``: one warp that sleeps for a given time, to place the work queued after
+it on the same stream later in the decode step.
+
+``prefetch_experts``: fire-and-forget ``cp.async.bulk.prefetch.L2`` of the
+local FC1 / FC2 weights and block scales of the first ``K`` predicted experts,
+expert by expert in rank order, FC1 bytes before FC2 bytes.
+"""
+
+import triton
+import triton.language as tl
+
+NUM_EXPERTS = 256
+TOP_K = 8
+MAX_ROWS = 8
+KMAX = 32
+# Experts ranked per prediction program.
+RANK_BLOCK = 32
+# Experts per program of the fused GEMM + rank kernel.
+GEMM_BLOCK = 16
+# Columns of the per-layer counter rows (int64).
+STAT_KS = (8, 12, 18, 24, 32)
+C_CALLS = 0
+C_DISTINCT = 1
+C_HITS = 2  # C_HITS + i for STAT_KS[i]
+NUM_COUNTERS = C_HITS + len(STAT_KS)
+
+
+@triton.jit
+def _scores_topk_union(s, live, NE: tl.constexpr, TOPK: tl.constexpr):
+    """For scores ``s`` [BR, NE] (dead rows -inf), the experts in each live
+    row's top-``TOPK``; returns (count over rows [NE], max score [NE])."""
+    e = tl.arange(0, NE)
+    # TOPK-th largest score of every row: TOPK passes, each removing one
+    # maximum (the lowest index among equal maxima).
+    work = s
+    for _ in tl.static_range(TOPK):
+        thr, idx = tl.max(work, axis=1, return_indices=True)
+        work = tl.where(e[None, :] == idx[:, None], float("-inf"), work)
+    sel = (s >= thr[:, None]) & live[:, None]
+    cnt = tl.sum(sel.to(tl.int32), axis=0)
+    mx = tl.max(s, axis=0)
+    return cnt, mx
+
+
+@triton.jit
+def _row_topk_union(
+    logits_ptr, bias_ptr, rows, NE: tl.constexpr, BR: tl.constexpr, TOPK: tl.constexpr
+):
+    """Per live row, the experts in that row's top-``TOPK`` of
+    ``sigmoid(logit) + bias``; returns (count over rows [NE], max score [NE])."""
+    r = tl.arange(0, BR)
+    e = tl.arange(0, NE)
+    live = r < rows
+    x = tl.load(logits_ptr + r[:, None] * NE + e[None, :], mask=live[:, None], other=0.0).to(
+        tl.float32
+    )
+    b = tl.load(bias_ptr + e).to(tl.float32)
+    s = 1.0 / (1.0 + tl.exp(-x)) + b[None, :]
+    s = tl.where(live[:, None], s, float("-inf"))
+    return _scores_topk_union(s, live, NE, TOPK)
+
+
+@triton.jit
+def _store_ranked(cnt, mx, eb, out_ids_ptr, NE: tl.constexpr, KM: tl.constexpr):
+    """Rank the experts ``eb`` against all NE by (rows selecting, max score,
+    ID) and store each one ranked below ``KM`` to its rank slot."""
+    e = tl.arange(0, NE)
+    same = e[None, :] == eb[:, None]
+    cb = tl.sum(tl.where(same, cnt[None, :], 0), axis=1)
+    mb = tl.max(tl.where(same, mx[None, :], float("-inf")), axis=1)
+    cj = cnt[None, :]
+    mj = mx[None, :]
+    better = (cj > cb[:, None]) | (
+        (cj == cb[:, None])
+        & ((mj > mb[:, None]) | ((mj == mb[:, None]) & (e[None, :] < eb[:, None])))
+    )
+    rank = tl.sum(better.to(tl.int32), axis=1)
+    tl.store(out_ids_ptr + rank, eb.to(tl.int32), mask=rank < KM)
+
+
+@triton.jit
+def _rank_and_stats_kernel(
+    next_logits_ptr,
+    next_bias_ptr,
+    out_ids_ptr,
+    cur_logits_ptr,
+    cur_bias_ptr,
+    cur_ids_ptr,
+    counters_ptr,
+    rows,
+    NPRED: tl.constexpr,
+    HAS_STATS: tl.constexpr,
+    NE: tl.constexpr,
+    BR: tl.constexpr,
+    BE: tl.constexpr,
+    TOPK: tl.constexpr,
+    KM: tl.constexpr,
+    K0: tl.constexpr,
+    K1: tl.constexpr,
+    K2: tl.constexpr,
+    K3: tl.constexpr,
+    K4: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    e = tl.arange(0, NE)
+    if pid < NPRED:
+        cnt, mx = _row_topk_union(next_logits_ptr, next_bias_ptr, rows, NE, BR, TOPK)
+        # this program's experts, ranked with the values of the full vectors
+        _store_ranked(cnt, mx, pid * BE + tl.arange(0, BE), out_ids_ptr, NE, KM)
+    elif HAS_STATS:
+        cnt, _ = _row_topk_union(cur_logits_ptr, cur_bias_ptr, rows, NE, BR, TOPK)
+        actual = cnt > 0
+        # rank position of every expert in the ranking predicted for this
+        # layer (KM = not among the first KM)
+        j = tl.arange(0, KM)
+        ids = tl.load(cur_ids_ptr + j)
+        pos = tl.min(tl.where(ids[:, None] == e[None, :], j[:, None], KM), axis=0)
+        tl.atomic_add(counters_ptr + 0, tl.full([], 1, tl.int64))
+        tl.atomic_add(counters_ptr + 1, tl.sum(actual.to(tl.int64), axis=0))
+        tl.atomic_add(counters_ptr + 2, tl.sum((actual & (pos < K0)).to(tl.int64), axis=0))
+        tl.atomic_add(counters_ptr + 3, tl.sum((actual & (pos < K1)).to(tl.int64), axis=0))
+        tl.atomic_add(counters_ptr + 4, tl.sum((actual & (pos < K2)).to(tl.int64), axis=0))
+        tl.atomic_add(counters_ptr + 5, tl.sum((actual & (pos < K3)).to(tl.int64), axis=0))
+        tl.atomic_add(counters_ptr + 6, tl.sum((actual & (pos < K4)).to(tl.int64), axis=0))
+
+
+def rank_and_stats(
+    next_logits,
+    next_bias,
+    out_ids,
+    cur_logits,
+    cur_bias,
+    cur_ids,
+    counters,
+    rows: int,
+    has_pred: bool,
+    has_stats: bool,
+) -> None:
+    """Launch the prediction and / or statistics programs on the current
+    stream. Unused pointer arguments may be any tensor of the right device."""
+    npred = NUM_EXPERTS // RANK_BLOCK if has_pred else 0
+    grid = npred + (1 if has_stats else 0)
+    if grid == 0:
+        return
+    _rank_and_stats_kernel[(grid,)](
+        next_logits,
+        next_bias,
+        out_ids,
+        cur_logits,
+        cur_bias,
+        cur_ids,
+        counters,
+        rows,
+        NPRED=npred,
+        HAS_STATS=has_stats,
+        NE=NUM_EXPERTS,
+        BR=MAX_ROWS,
+        BE=RANK_BLOCK,
+        TOPK=TOP_K,
+        KM=KMAX,
+        K0=STAT_KS[0],
+        K1=STAT_KS[1],
+        K2=STAT_KS[2],
+        K3=STAT_KS[3],
+        K4=STAT_KS[4],
+        num_warps=4,
+    )
+
+
+@triton.jit
+def _gemm_rank_kernel(
+    h_ptr,
+    w_ptr,
+    bias_ptr,
+    scores_ptr,
+    ticket_ptr,
+    out_ids_ptr,
+    rows,
+    H: tl.constexpr,
+    NE: tl.constexpr,
+    BR: tl.constexpr,
+    BM: tl.constexpr,
+    BE: tl.constexpr,
+    BK: tl.constexpr,
+    TOPK: tl.constexpr,
+    KM: tl.constexpr,
+):
+    # Logits and scores of this program's BE experts for the rows of h; the
+    # gate rows are read contiguously along the hidden dimension.
+    pid = tl.program_id(0)
+    NPROG: tl.constexpr = NE // BE
+    m = tl.arange(0, BM)
+    eb = pid * BE + tl.arange(0, BE)
+    kk = tl.arange(0, BK)
+    acc = tl.zeros([BM, BE], dtype=tl.float32)
+    for k0 in range(0, H, BK):
+        a = tl.load(h_ptr + m[:, None] * H + k0 + kk[None, :], mask=(m < rows)[:, None], other=0.0)
+        b = tl.load(w_ptr + eb[:, None] * H + k0 + kk[None, :])
+        acc += tl.dot(a, tl.trans(b))
+    bias = tl.load(bias_ptr + eb).to(tl.float32)
+    s = 1.0 / (1.0 + tl.exp(-acc)) + bias[None, :]
+    # Publish the scores at L2 and take a ticket; the last program to arrive
+    # ranks all NE experts and resets the ticket for the next launch.
+    tl.atomic_xchg(
+        scores_ptr + m[:, None] * NE + eb[None, :], s, mask=(m < BR)[:, None], sem="relaxed"
+    )
+    tl.debug_barrier()
+    ticket = tl.atomic_add(ticket_ptr, 1, sem="acq_rel")
+    if ticket == NPROG - 1:
+        r = tl.arange(0, BR)
+        e = tl.arange(0, NE)
+        live = r < rows
+        sc = tl.load(scores_ptr + r[:, None] * NE + e[None, :], cache_modifier=".cg")
+        sc = tl.where(live[:, None], sc, float("-inf"))
+        cnt, mx = _scores_topk_union(sc, live, NE, TOPK)
+        # One sort of a unique key per expert: rows selecting, then the max
+        # score (order-preserving bits), then the lower ID first.
+        bits = mx.to(tl.int32, bitcast=True)
+        ordered = tl.where(bits < 0, bits ^ 0x7FFFFFFF, bits).to(tl.int64) + 2147483648
+        key = (cnt.to(tl.int64) << 41) + (ordered << 8) + (NE - 1 - e).to(tl.int64)
+        key = tl.sort(key, dim=0, descending=True)
+        ids = (NE - 1 - (key & (NE - 1))).to(tl.int32)
+        tl.store(out_ids_ptr + e, ids, mask=e < KM)
+        tl.atomic_xchg(ticket_ptr, 0, sem="relaxed")
+
+
+def gemm_rank(hidden, weight, bias, scores, ticket, out_ids, rows: int) -> None:
+    """Prediction in one launch on the current stream: ``hidden`` [rows, H]
+    BF16 times the next layer's gate ``weight`` [256, H] BF16, scores, per-row
+    top-8 union and ranking into ``out_ids`` (as ``rank_and_stats``).
+    ``scores`` is an FP32 [MAX_ROWS, 256] scratch and ``ticket`` an int32 [1]
+    zero; both are reused by every launch on the same stream."""
+    _gemm_rank_kernel[(NUM_EXPERTS // GEMM_BLOCK,)](
+        hidden,
+        weight,
+        bias,
+        scores,
+        ticket,
+        out_ids,
+        rows,
+        H=hidden.shape[1],
+        NE=NUM_EXPERTS,
+        BR=MAX_ROWS,
+        BM=16,
+        BE=GEMM_BLOCK,
+        BK=512,
+        TOPK=TOP_K,
+        KM=KMAX,
+        num_warps=4,
+        num_stages=4,
+    )
+
+
+@triton.jit
+def _globaltimer(z):
+    return tl.inline_asm_elementwise(
+        "mov.u64 $0, %globaltimer; // $1", "=l,r", [z], dtype=tl.int64, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _nanosleep(ns):
+    return tl.inline_asm_elementwise(
+        "nanosleep.u32 $1; mov.u32 $0, 0;", "=r,r", [ns], dtype=tl.int32, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _delay_kernel(delay_ns):
+    # One warp sleeps until delay_ns of global time have passed.
+    z = tl.zeros([32], dtype=tl.int32)
+    t0 = tl.max(_globaltimer(z), axis=0)
+    t = t0
+    while t - t0 < delay_ns:
+        _nanosleep(z + 256)
+        t = tl.max(_globaltimer(z), axis=0)
+
+
+def delay(delay_ns: int) -> None:
+    """Hold the current stream for about ``delay_ns`` nanoseconds with one
+    warp (the work queued after it starts that much later)."""
+    _delay_kernel[(1,)](delay_ns, num_warps=1)
+
+
+@triton.jit
+def _bulk_prefetch_l2(addr, nbytes):
+    tl.inline_asm_elementwise(
+        "cp.async.bulk.prefetch.L2.global [$1], $2; // dummy $0",
+        "=r,l,r",
+        [addr, nbytes],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _prefetch_experts_kernel(
+    ids_ptr,
+    w13_addr,
+    w13_sf_addr,
+    w2_addr,
+    w2_sf_addr,
+    K: tl.constexpr,
+    CHUNK: tl.constexpr,
+    C13: tl.constexpr,
+    C13SF: tl.constexpr,
+    C2: tl.constexpr,
+    C2SF: tl.constexpr,
+    NPROG: tl.constexpr,
+    NT: tl.constexpr,
+):
+    # Chunks of one expert, in the order the routed kernel reads them:
+    # FC1 weight, FC1 scales, FC2 weight, FC2 scales. Addresses are int64
+    # byte addresses (the four tensors have different element types).
+    CPE: tl.constexpr = C13 + C13SF + C2 + C2SF
+    TOTAL: tl.constexpr = K * CPE
+    ITERS: tl.constexpr = (TOTAL + NPROG * NT - 1) // (NPROG * NT)
+    pid = tl.program_id(0).to(tl.int64)
+    lanes = tl.arange(0, NT).to(tl.int64)
+    for it in range(ITERS):
+        c = (it * NPROG + pid) * NT + lanes
+        live = c < TOTAL
+        c = tl.where(live, c, 0)
+        slot = c // CPE
+        off = c - slot * CPE
+        ex = tl.load(ids_ptr + slot, mask=live, other=0).to(tl.int64)
+        a13 = w13_addr + ex * (C13 * CHUNK) + off * CHUNK
+        o1 = off - C13
+        a13sf = w13_sf_addr + ex * (C13SF * CHUNK) + o1 * CHUNK
+        o2 = o1 - C13SF
+        a2 = w2_addr + ex * (C2 * CHUNK) + o2 * CHUNK
+        o3 = o2 - C2
+        a2sf = w2_sf_addr + ex * (C2SF * CHUNK) + o3 * CHUNK
+        addr = tl.where(off < C13, a13, tl.where(o1 < C13SF, a13sf, tl.where(o2 < C2, a2, a2sf)))
+        # Lanes past the end re-issue the first chunk, which is harmless.
+        _bulk_prefetch_l2(addr, tl.full([NT], CHUNK, tl.int32))
+
+
+def prefetch_experts(
+    ids, w13, w13_sf, w2, w2_sf, k: int, chunk: int, ctas: int, threads: int = 32
+) -> None:
+    """Prefetch the first ``k`` experts of ``ids`` into L2 on the current
+    stream. Every expert slice of the four tensors must be a whole number of
+    ``chunk`` bytes."""
+
+    def per_expert_chunks(t):
+        nbytes = t[0].numel() * t.element_size()
+        assert nbytes % chunk == 0, (tuple(t.shape), chunk)
+        return nbytes // chunk
+
+    _prefetch_experts_kernel[(ctas,)](
+        ids,
+        w13.data_ptr(),
+        w13_sf.data_ptr(),
+        w2.data_ptr(),
+        w2_sf.data_ptr(),
+        K=k,
+        CHUNK=chunk,
+        C13=per_expert_chunks(w13),
+        C13SF=per_expert_chunks(w13_sf),
+        C2=per_expert_chunks(w2),
+        C2SF=per_expert_chunks(w2_sf),
+        NPROG=ctas,
+        NT=threads,
+        num_warps=max(1, threads // 32),
+        num_stages=1,
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/spec_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/spec_experts.py
new file mode 100644
index 000000000000..82dec07518b5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/spec_experts.py
@@ -0,0 +1,508 @@
+# 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.
+"""Cross-layer speculative routed-expert L2 prefetch (``cross_layer_expert_prefetch``).
+
+At a decode call of MoE layer ``L`` (1 to 8 rows), the post-attention hop has
+produced this layer's normalized hidden ``h_L`` and router logits. On a
+private auxiliary stream, forked from the caller's stream at that point:
+
+1. ``h_L`` times the next MoE layer's gate gives approximate router logits of
+   layer ``L + 1``; every live row's top-8 of ``sigmoid(logit) + bias`` is
+   taken, and the experts are ranked by how many rows select them (then by
+   their best score). The first ``KMAX`` expert IDs are kept for layer
+   ``L + 1``.
+2. Layer ``L``'s own routing (from its router logits) is compared with the
+   ranking predicted for it at layer ``L - 1``: per-layer device counters of
+   calls, distinct experts and hits among the first K predicted experts.
+3. After layer ``L``'s routed MoE (``start="moe"``),
+   or right after the prediction (``hop``), or after layer ``L``'s post-MoE
+   hop (``post``), a kernel issues ``cp.async.bulk.prefetch.L2`` over the
+   local FC1 / FC2 weights and block scales of the first ``K`` predicted
+   experts of layer ``L + 1`` (``k``, 0 = predict and count only).
+
+With ``issue="xc"`` no prefetch kernel is launched: the routed_moe
+expert-cluster kernel of layer ``L + 1`` issues the prefetch itself, before
+its dependency wait, i.e. while the pre-MoE hop of layer ``L + 1`` still runs
+(``routed_moe/kernel_xc.py``, ``spec_k``). With ``pred="self"`` there
+is no auxiliary work at all: the expert-cluster kernel of layer ``L`` writes
+its own experts, ranked by row count, and its next call (the next decode
+step) prefetches them.
+
+A prefetch has no architectural effect, and the auxiliary work writes only
+this package's buffers, so model outputs are unchanged. The caller's stream
+never waits for the auxiliary stream inside a layer: the hidden state and
+router logits the auxiliary work reads are recorded on the auxiliary stream
+(``Tensor.record_stream``), so the caching allocator does not reuse them
+before that work is done (under CUDA-graph capture, not before the capture
+ends). The caller's stream waits for all auxiliary work once, at the end of
+the target model's forward (``join``), so a CUDA-graph capture always sees the
+auxiliary stream rejoin.
+
+Configuration (``ExpertPrefetchConfig``; this target runs ``SERVED_CONFIG``):
+
+* ``k``: experts prefetched per layer, 0..32 (8).
+* ``start``: ``moe``, ``hop`` or ``post`` (``moe``).
+* ``pred``: where the prediction (and counters) is forked:
+  ``hop`` (right after the pre-MoE hop, concurrent with layer L's
+  routed MoE) or ``moe`` (right after layer L's routed MoE, together with a
+  ``start=moe`` prefetch); ``moe`` excludes ``start=hop``. ``self``: no
+  cross-layer prediction (see above; needs ``issue=xc``, no counters) (``self``).
+* ``issue``: ``aux`` (a prefetch kernel on the auxiliary stream, placed by
+  ``start`` and ``delay_ns``) or ``xc`` (the routed MoE kernel of the
+  predicted layer) (``xc``).
+* ``xc_ctas``: with ``issue=xc``, the routed-kernel CTAs that issue the
+  prefetch (0 = all) (32).
+* ``fused``: predict with one fused Triton launch (``kernels.gemm_rank``: gate
+  GEMM, scores and ranking) when the next gate is a BF16 [256, 6144] weight;
+  otherwise the next gate module plus the ranking kernel (True).
+* ``delay_ns``: nanoseconds the auxiliary stream waits (one sleeping warp)
+  right before each prefetch, to move the prefetch traffic later into layer
+  L+1 (0).
+* ``chunk``: bytes per prefetch instruction (16384).
+* ``ctas``: CTAs of the prefetch kernel (16).
+
+``start``, ``fused``, ``delay_ns``, ``chunk`` and ``ctas`` act only with the
+``hop`` / ``moe`` predictors and the ``aux`` prefetch kernel; under ``self``
+they appear only in the capability key.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+
+from ....shared import dispatch
+from . import kernels
+
+HIDDEN = 6144
+STARTS = ("moe", "hop", "post")
+PREDS = ("hop", "moe", "self")
+ISSUES = ("aux", "xc")
+
+
+@dataclass(frozen=True)
+class ExpertPrefetchConfig:
+    k: int
+    start: str
+    pred: str
+    fused: bool
+    delay_ns: int
+    chunk: int
+    ctas: int
+    issue: str = "xc"
+    xc_ctas: int = 32
+
+
+#: The expert-prefetch configuration of this target: the routed MoE kernel of each layer writes
+#: its own experts (``self``) and its next call prefetches the first 8 on 32 of its CTAs.
+SERVED_CONFIG = ExpertPrefetchConfig(
+    k=8,
+    start="moe",
+    pred="self",
+    fused=True,
+    delay_ns=0,
+    chunk=16384,
+    ctas=16,
+    issue="xc",
+    xc_ctas=32,
+)
+
+
+def config() -> ExpertPrefetchConfig:
+    """The expert-prefetch configuration of this target, ``SERVED_CONFIG``."""
+    return SERVED_CONFIG
+
+
+class _DeviceState:
+    """Per-device stream, events and buffers shared by every linked layer.
+
+    Created on the first eager call of a device, never replaced (captured
+    graphs keep the buffer addresses).
+    """
+
+    def __init__(self, device: torch.device, num_layers: int):
+        self.device = device
+        self.aux = torch.cuda.Stream(device=device)
+        self.ev_fork = torch.cuda.Event()
+        self.ev_tail = torch.cuda.Event()
+        self.counters = torch.zeros(
+            num_layers, kernels.NUM_COUNTERS, dtype=torch.int64, device=device
+        )
+        self.pred_ids = torch.zeros(num_layers + 1, kernels.KMAX, dtype=torch.int32, device=device)
+        self.dummy_f32 = torch.zeros(
+            kernels.MAX_ROWS, kernels.NUM_EXPERTS, dtype=torch.float32, device=device
+        )
+        self.dummy_bias = torch.zeros(kernels.NUM_EXPERTS, dtype=torch.float32, device=device)
+        # scratch of the fused prediction (kernels.gemm_rank)
+        self.scores = torch.zeros(
+            kernels.MAX_ROWS, kernels.NUM_EXPERTS, dtype=torch.float32, device=device
+        )
+        self.ticket = torch.zeros(1, dtype=torch.int32, device=device)
+        self.forked = False
+        # prefetch launches issued from Python (eager calls and captures, not
+        # graph replays)
+        self.prefetch_calls = 0
+
+
+_STATES: Dict[int, _DeviceState] = {}
+_NUM_LAYERS = {"n": 0}
+
+
+def _state(device: torch.device) -> _DeviceState:
+    idx = device.index if device.index is not None else torch.cuda.current_device()
+    st = _STATES.get(idx)
+    if st is None:
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "[glm_moe_dsa] cross_layer_expert_prefetch device state must be created by an "
+                "eager call before CUDA-graph capture"
+            )
+        try:
+            from torch.utils._python_dispatch import _disable_current_modes
+        except ImportError:
+            import contextlib
+
+            _disable_current_modes = contextlib.nullcontext
+        with _disable_current_modes():
+            st = _DeviceState(torch.device("cuda", idx), _NUM_LAYERS["n"])
+        _STATES[idx] = st
+    return st
+
+
+def _expert_tensors(experts) -> Optional[List[torch.Tensor]]:
+    """The four per-expert weight tensors of a native NVFP4 MoE leaf, or None."""
+    from ..routed_moe.routed_moe import resolve_expert_owner
+
+    owner, _ = resolve_expert_owner(experts)
+    if owner is None:
+        return None
+    names = ("w3_w1_weight", "w3_w1_weight_scale", "w2_weight", "w2_weight_scale")
+    ts = [owner._parameters.get(n) for n in names]
+    if any(t is None for t in ts):
+        return None
+    return ts
+
+
+def _prefetch_reason(ts: Optional[List[torch.Tensor]], chunk: int) -> str:
+    if ts is None:
+        return "next layer has no native NVFP4 expert leaf"
+    for t in ts:
+        if not t.is_cuda or not t.is_contiguous():
+            return "expert tensor is not contiguous on CUDA"
+        if t.shape[0] != kernels.NUM_EXPERTS:
+            return f"expert tensor {tuple(t.shape)} is not [256, ...]"
+        if t.data_ptr() % 16:
+            return "expert tensor is not 16-byte aligned"
+        if (t[0].numel() * t.element_size()) % chunk:
+            return f"expert slice of {tuple(t.shape)} is not a multiple of {chunk} bytes"
+    return ""
+
+
+class ExpertPrefetchLink:
+    """cross_layer_expert_prefetch state of one MoE decoder layer ``L`` (bound by ``link_layers``)."""
+
+    def __init__(self, layer_idx: int, cur_gate, next_gate, next_experts, has_prev: bool):
+        self.layer_idx = layer_idx
+        self.cur_gate = cur_gate
+        self.next_gate = next_gate
+        self.next_experts = next_experts
+        self.has_prev = has_prev
+        self._armed = False
+        self._pending_prefetch = False
+        self._pending_predict: Optional[tuple] = None
+        self._next_tensors: Optional[List[torch.Tensor]] = None
+        self._prefetch_ok: Optional[bool] = None
+        self._fused_ok: Optional[bool] = None
+
+    @staticmethod
+    def _eligible(hidden, router_logits) -> str:
+        if not isinstance(hidden, torch.Tensor) or not hidden.is_cuda:
+            return "hidden is not a CUDA tensor"
+        if hidden.dtype != torch.bfloat16 or hidden.dim() != 2 or hidden.shape[1] != HIDDEN:
+            return f"hidden {tuple(hidden.shape)} {hidden.dtype}"
+        rows = hidden.shape[0]
+        if not 1 <= rows <= kernels.MAX_ROWS:
+            return f"{rows} rows outside [1, {kernels.MAX_ROWS}]"
+        if not hidden.is_contiguous():
+            return "hidden is not contiguous"
+        if router_logits is not None and (
+            router_logits.dtype != torch.float32
+            or tuple(router_logits.shape) != (rows, kernels.NUM_EXPERTS)
+            or not router_logits.is_contiguous()
+        ):
+            return f"router logits {tuple(router_logits.shape)} {router_logits.dtype}"
+        return ""
+
+    def _prefetch_ready(self, cfg: ExpertPrefetchConfig) -> bool:
+        if self._prefetch_ok is None:
+            ts = _expert_tensors(self.next_experts) if self.next_experts is not None else None
+            reason = _prefetch_reason(ts, cfg.chunk)
+            self._prefetch_ok = not reason
+            self._next_tensors = ts if not reason else None
+            dispatch.log_capability(
+                "cross_layer_expert_prefetch", "expert_prefetch", self._prefetch_ok, reason
+            )
+        return self._prefetch_ok
+
+    def _launch_prefetch(self, st: _DeviceState, cfg: ExpertPrefetchConfig) -> None:
+        if cfg.delay_ns:
+            kernels.delay(cfg.delay_ns)
+        w13, w13_sf, w2, w2_sf = self._next_tensors
+        kernels.prefetch_experts(
+            st.pred_ids[self.layer_idx + 1], w13, w13_sf, w2, w2_sf, cfg.k, cfg.chunk, cfg.ctas
+        )
+        st.prefetch_calls += 1
+
+    def _fused_ready(self, cfg: ExpertPrefetchConfig) -> bool:
+        if self._fused_ok is None:
+            reason = ""
+            if not cfg.fused:
+                reason = "fused=0"
+            else:
+                w = getattr(self.next_gate, "weight", None)
+                if (
+                    not isinstance(w, torch.Tensor)
+                    or w.dtype != torch.bfloat16
+                    or tuple(w.shape) != (kernels.NUM_EXPERTS, HIDDEN)
+                    or not w.is_contiguous()
+                ):
+                    reason = "next gate weight is not a contiguous BF16 [256, 6144]"
+            self._fused_ok = not reason
+            dispatch.log_capability(
+                "cross_layer_expert_prefetch", "fused_predict", self._fused_ok, reason
+            )
+        return self._fused_ok
+
+    def _predict(
+        self,
+        st: _DeviceState,
+        cfg: ExpertPrefetchConfig,
+        hidden,
+        router_logits,
+        has_pred: bool,
+        has_stats: bool,
+    ) -> None:
+        """Prediction for layer L+1 and counters of layer L (current stream)."""
+        rows = hidden.shape[0]
+        out_ids = st.pred_ids[self.layer_idx + 1]
+        fused = has_pred and self._fused_ready(cfg)
+        if fused:
+            kernels.gemm_rank(
+                hidden,
+                self.next_gate.weight,
+                self.next_gate.e_score_correction_bias,
+                st.scores,
+                st.ticket,
+                out_ids,
+                rows,
+            )
+        if fused and not has_stats:
+            return
+        has_pred = has_pred and not fused
+        next_logits = self.next_gate(hidden) if has_pred else st.dummy_f32
+        next_bias = self.next_gate.e_score_correction_bias if has_pred else st.dummy_bias
+        kernels.rank_and_stats(
+            next_logits,
+            next_bias,
+            out_ids,
+            router_logits if has_stats else st.dummy_f32,
+            self.cur_gate.e_score_correction_bias if has_stats else st.dummy_bias,
+            st.pred_ids[self.layer_idx],
+            st.counters[self.layer_idx],
+            rows,
+            has_pred,
+            has_stats,
+        )
+
+    def _fork(
+        self, st: _DeviceState, cfg: ExpertPrefetchConfig, predict: Optional[tuple], prefetch: bool
+    ) -> None:
+        """Fork auxiliary work at this point of the caller's stream: the
+        prediction (``predict`` = arguments of ``_predict``) and / or the
+        prefetch of layer L+1, in that order."""
+        caller = torch.cuda.current_stream()
+        st.ev_fork.record(caller)
+        st.aux.wait_event(st.ev_fork)
+        with torch.cuda.stream(st.aux):
+            if predict is not None:
+                self._predict(st, cfg, *predict)
+            if prefetch:
+                self._launch_prefetch(st, cfg)
+            st.ev_tail.record(st.aux)
+        st.forked = True
+
+    def on_hop(self, hidden, router_logits) -> None:
+        """After layer L's pre-MoE hop: fork the prediction (``pred=hop``) or
+        keep its inputs for ``after_moe`` (``pred=moe``)."""
+        self._armed = False
+        self._pending_prefetch = False
+        self._pending_predict = None
+        reason = self._eligible(hidden, router_logits)
+        if reason:
+            # calls outside the decode shape (e.g. warmup with many rows)
+            dispatch.log_capability(
+                "cross_layer_expert_prefetch", "spec_experts_declined", False, reason
+            )
+            return
+        dispatch.log_capability("cross_layer_expert_prefetch", "spec_experts", True)
+        cfg = config()
+        # created eagerly here: the routed kernel's prefetch hook reads it
+        st = _state(hidden.device)
+        if cfg.pred == "self":
+            return
+        has_pred = self.next_gate is not None
+        # The hit counters are not reported in this target.
+        has_stats = False
+        if not has_pred and not has_stats:
+            return
+        hidden.record_stream(st.aux)
+        if has_stats:
+            router_logits.record_stream(st.aux)
+        predict = (hidden, router_logits, has_pred, has_stats)
+        if cfg.pred == "hop":
+            prefetch = (
+                has_pred
+                and cfg.k > 0
+                and cfg.issue == "aux"
+                and cfg.start == "hop"
+                and self._prefetch_ready(cfg)
+            )
+            self._fork(st, cfg, predict, prefetch)
+        else:
+            self._pending_predict = predict
+        self._armed = True
+        self._pending_prefetch = (
+            has_pred and cfg.k > 0 and cfg.issue == "aux" and cfg.start != "hop"
+        )
+
+    def after_moe(self) -> None:
+        """After layer L's routed MoE: fork the pending prediction
+        (``pred=moe``) and / or the prefetch of layer L+1 (``start=moe``)."""
+        if not self._armed:
+            return
+        cfg = config()
+        prefetch = self._pending_prefetch and cfg.start == "moe" and self._prefetch_ready(cfg)
+        if prefetch:
+            self._pending_prefetch = False
+        predict, self._pending_predict = self._pending_predict, None
+        if predict is not None or prefetch:
+            st = _state(torch.device("cuda", torch.cuda.current_device()))
+            self._fork(st, cfg, predict, prefetch)
+
+    def after_post_moe(self) -> None:
+        """After layer L's post-MoE hop (``start=post``)."""
+        if not self._armed:
+            return
+        self._armed = False
+        cfg = config()
+        if self._pending_prefetch and cfg.start == "post" and self._prefetch_ready(cfg):
+            self._pending_prefetch = False
+            st = _state(torch.device("cuda", torch.cuda.current_device()))
+            self._fork(st, cfg, None, True)
+
+
+def _is_moe_with_gate(layer) -> bool:
+    mlp = getattr(layer, "mlp", None)
+    gate = getattr(mlp, "gate", None)
+    return (
+        mlp is not None
+        and getattr(mlp, "experts", None) is not None
+        and gate is not None
+        and hasattr(gate, "e_score_correction_bias")
+        and hasattr(gate, "weight")
+    )
+
+
+def _xc_hook(row: int, cfg: ExpertPrefetchConfig):
+    """Prefetch hook of a routed-MoE leaf (read by ``routed_moe.routed_moe``): the
+    ``pred_ids`` row its expert-cluster kernel prefetches, or ``None`` before
+    the device state exists."""
+    self_store = cfg.pred == "self"
+
+    def hook():
+        st = _STATES.get(torch.cuda.current_device())
+        if st is None:
+            return None
+        return (st.pred_ids[row], cfg.k, cfg.xc_ctas, self_store)
+
+    return hook
+
+
+def _xc_issue_reason(cfg: ExpertPrefetchConfig) -> str:
+    """Why the routed kernel cannot issue the prefetch ("" when it can)."""
+    if not dispatch.is_enabled("routed_moe"):
+        return "issue=xc needs routed_moe"
+    return ""
+
+
+def _bind_xc(experts, row: int, cfg: ExpertPrefetchConfig) -> bool:
+    from ..routed_moe.routed_moe import resolve_expert_owner
+
+    owner, _ = resolve_expert_owner(experts)
+    if owner is None:
+        return False
+    owner._expert_prefetch_spec = _xc_hook(row, cfg)
+    return True
+
+
+def link_layers(layers, num_hidden_layers: int) -> int:
+    """Bind a ``ExpertPrefetchLink`` to every target MoE layer (structural only)."""
+    cfg = config()
+    _NUM_LAYERS["n"] = max(_NUM_LAYERS["n"], num_hidden_layers)
+    linked = 0
+    xc_bound = 0
+    xc_on = cfg.issue == "xc" and cfg.k > 0
+    if xc_on:
+        reason = _xc_issue_reason(cfg)
+        xc_on = dispatch.log_capability(
+            "cross_layer_expert_prefetch", "xc_prefetch", not reason, reason
+        )
+    for idx in range(num_hidden_layers):
+        layer = layers[idx]
+        if not _is_moe_with_gate(layer):
+            continue
+        nxt = layers[idx + 1] if idx + 1 < num_hidden_layers else None
+        nxt_moe = nxt is not None and _is_moe_with_gate(nxt)
+        prev_moe = idx > 0 and _is_moe_with_gate(layers[idx - 1])
+        layer._expert_prefetch_link = ExpertPrefetchLink(
+            idx,
+            layer.mlp.gate,
+            nxt.mlp.gate if nxt_moe else None,
+            nxt.mlp.experts if nxt_moe else None,
+            prev_moe,
+        )
+        linked += 1
+        if xc_on:
+            if cfg.pred == "self":
+                xc_bound += _bind_xc(layer.mlp.experts, idx, cfg)
+            elif nxt_moe:
+                xc_bound += _bind_xc(nxt.mlp.experts, idx + 1, cfg)
+    dispatch.log_capability(
+        "cross_layer_expert_prefetch",
+        f"link[k={cfg.k},start={cfg.start},pred={cfg.pred},fused={int(cfg.fused)},"
+        f"delay_ns={cfg.delay_ns},chunk={cfg.chunk},"
+        f"ctas={cfg.ctas},issue={cfg.issue},xc_ctas={cfg.xc_ctas}]",
+        linked > 0,
+        f"{linked} MoE layers, {xc_bound} routed-kernel prefetch hooks",
+    )
+    return linked
+
+
+def join() -> None:
+    """Make the current stream wait for all forked auxiliary work."""
+    for st in _STATES.values():
+        if st.forked and st.device.index == torch.cuda.current_device():
+            torch.cuda.current_stream().wait_event(st.ev_tail)
+            st.forked = False
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/__init__.py
new file mode 100644
index 000000000000..2156cea0dd74
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/__init__.py
@@ -0,0 +1,30 @@
+# 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.
+"""Package draft_routed_moe: the MTP draft layer's routed experts (BF16) at the DRAM rate.
+
+Entry point registered with ``dispatch``:
+
+* ``routed_output(module, routing_bias, x, router_logits, do_finalize)``:
+  the finalized routed output ``[R, 6144]`` of the BF16 FlashInfer
+  TRT-LLM-Gen leaf, bit-identical to ``trtllm_bf16_moe``, or ``None`` (the
+  native path; decided before any device work).
+"""
+
+from ....shared.dispatch import register
+from .routed_moe import eligibility, launch_config, op_available, routed_output
+
+register("draft_routed_moe", "routed_output", routed_output)
+
+__all__ = ["eligibility", "launch_config", "op_available", "routed_output"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/routed_moe.py
new file mode 100644
index 000000000000..66fc869f73d3
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/routed_moe.py
@@ -0,0 +1,408 @@
+# 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.
+"""Host side of package draft_routed_moe: the routed experts of the MTP draft layer.
+
+The MTP layer of the GLM-5.2 NVFP4 checkpoint is unquantized, so its routed
+experts run FlashInfer's TRT-LLM-Gen BF16 MoE (``FlashinferTrtllmGenBf16Impl``,
+``trtllm_bf16_moe``). ``routed_output`` replaces that call with
+``torch.ops.trtllm.glm_moe_dsa_draft_routed_moe`` on the leaf's own weight buffers (no
+copy, no re-layout) and returns the same finalized ``[R, 6144]`` BF16 tensor,
+bit for bit:
+
+* routing: FlashInfer's ``routingCustom`` with ``SigmoidBias`` preprocess and
+  ``ScaledSumNormalize`` postprocess (DeepSeek-V3 with ``n_group <= 1``),
+  including its fast-math instructions, tie order and FP32 sum order;
+* FC1/FC2: ``mma.sync`` over K in natural order with one FP32 accumulator per
+  output element (the batched-GEMM cubins accumulate K sequentially and have
+  no split-K), the cubin's SwiGLU epilogue and BF16 rounding;
+* finalize: ``finalizeKernel``'s FP32 ``fma`` chain in top-k slot order.
+
+The eligibility decision uses host-side state only (module types, routing
+configuration, shapes, dtypes, devices, the op's availability and whether the
+scratch buffers exist before a CUDA graph capture) and is taken before any
+device work; ``None`` leaves the native path with no side effect.
+
+Programmatic dependent launch (memory order and progress):
+
+* both kernels are PDL secondaries; neither reads anything before its
+  ``griddepcontrol.wait`` (the FC1 kernel's first reads, the router logits
+  and the activation, are the predecessor's outputs);
+* the FC1 grid writes the routing record and the BF16 intermediate into the
+  per-device scratch after its wait; the FC2 grid reads them after its own
+  wait, which returns once the FC1 grid has completed and its writes are
+  visible; the FC1 grid triggers its dependents after its tiles (or, with
+  flag bit 4 ``FLAG_EARLY_TRIGGER``, after the routing), which only lets the
+  FC2 CTAs be scheduled earlier;
+* the scratch of call ``i`` is rewritten by the FC1 grid of call ``i + 1``
+  only after that grid's wait, i.e. after its predecessor has completed; every
+  kernel between the two calls on the stream either is launched without
+  programmatic overlap or waits for its own predecessor before completing, so
+  that predecessor's completion implies the completion of call ``i``'s FC2
+  grid;
+* the optional L2 prefetch of FC2 weights reads immutable model parameters;
+* no CTA waits for another CTA of its grid (no spin, no grid barrier), so
+  progress does not depend on co-residency with the shared expert's kernels.
+"""
+
+import os
+from typing import Dict, Optional, Tuple
+
+import torch
+
+HIDDEN = 6144
+INTER = 256
+NUM_EXPERTS = 256
+TOP_K = 8
+MAX_ROWS = 8
+K_CHUNK = 64
+SCRATCH_H_ELEMS = 64 * MAX_ROWS * INTER
+ROUTE_WORDS = 1 + 64 + 64 + 64
+
+FLAG_PDL = 1
+FLAG_PREFETCH_FC2 = 2
+FLAG_EARLY_TRIGGER = 4
+FLAG_FC1_PAIRS = 8
+FLAG_FC2_PAIRS = 16
+FLAG_DEEP_RING = 32
+FLAG_TMA = 64
+FLAG_L2_PREFETCH = 128
+FLAG_PREFETCH_FC2_EARLY = 256
+FLAG_FC1_TRIPLES = 4096
+FLAG_EARLY_TRIGGER_TRIPLES = 8192
+FLAG_FC1_EVICT_FIRST = 32768
+FLAG_STREAM_X = 65536
+FLAG_OPTIONS = (
+    FLAG_PREFETCH_FC2
+    | FLAG_EARLY_TRIGGER
+    | FLAG_FC1_PAIRS
+    | FLAG_FC2_PAIRS
+    | FLAG_DEEP_RING
+    | FLAG_TMA
+    | FLAG_L2_PREFETCH
+    | FLAG_PREFETCH_FC2_EARLY
+    | FLAG_FC1_TRIPLES
+    | FLAG_EARLY_TRIGGER_TRIPLES
+    | FLAG_FC1_EVICT_FIRST
+    | FLAG_STREAM_X
+)
+# Shared memory of the FC1 kernel: the activation rows plus a weight ring per warp (2 KB chunks); the deep ring
+# applies to a call only when it fits this budget (the static routing state takes about 3 KB of the 227 KB).
+FC1_RING_CHUNKS_DEEP = 20
+FC1_DYNAMIC_SMEM_BUDGET = 220 * 1024
+
+OP_NAME = "glm_moe_dsa_draft_routed_moe"
+
+# Launch configuration (``launch_config``): (ctas, warps) of each kernel and extra flag bits (2: FC2 weights prefetched
+# into L2 by the FC1 warps, 4: early dependent trigger, 8: FC1 tile pairs, 16: FC2 expert pairs, 32: deeper weight
+# rings, dropped for a call whose FC1 shared memory would not fit, 64: TMA-fed kernels, where "ctas" is the persistent
+# grid size and "warps" the number of MMA warps (FC2 also 3), each owning a 16-row chain (two with bit 8), next to one
+# producer warp; with them bit 2 is issued by the FC1 producer after its last weight copy, 16 gives every FC2 warp two
+# experts per step, 128 adds L2 prefetches of the FC1 weights ahead of the ring, 256 moves the FC2 prefetch to right
+# after routing, and bit 32 does not apply; 4096 (4 FC1 MMA warps): a call whose FC1 tiles exceed the grid runs in one
+# round of three 64-row sub-tiles per CTA when they fit, chosen on the device from the routing; 8192: the early
+# trigger for calls on that path only; 32768: FC1 weight loads with an L2 evict-first policy; 65536 (4 FC1 MMA warps
+# with bit 8): the activation streamed through the FC1 ring with the weights instead of kept in shared memory, so the
+# ring is deeper).
+# Each value is a pair (calls with one row, calls with more rows). None of them changes a bit of the output.
+# (one row, more rows): TMA-fed kernels. One row: 64 FC1 CTAs (64-row tiles, one per CTA), which leaves SMs to the
+# shared expert on the auxiliary stream, and the FC2 weights prefetched into L2 behind the last FC1 copy. More rows:
+# 128-row FC1 tiles (two chains per MMA warp) on up to 152 CTAs, one round for up to 38 distinct experts; above that
+# the triples path (three 64-row sub-tiles per CTA, one round up to 57 distinct experts) with the early trigger, so
+# the FC2 grid takes the SMs FC1 leaves and is resident when FC1 ends (calls on the regular path trigger after their
+# tiles, which leaves those SMs to the shared expert); FC1 weights (read once per call) are loaded L2 evict-first.
+DEFAULT_FC1 = ((64, 4), (152, 4))
+DEFAULT_FC2 = ((128, 3), (128, 3))
+DEFAULT_FLAGS = (
+    FLAG_TMA | FLAG_PREFETCH_FC2,
+    FLAG_TMA
+    | FLAG_FC1_PAIRS
+    | FLAG_FC1_TRIPLES
+    | FLAG_EARLY_TRIGGER_TRIPLES
+    | FLAG_FC1_EVICT_FIRST,
+)
+
+
+def op_available() -> bool:
+    return hasattr(torch.ops.trtllm, OP_NAME)
+
+
+def _pdl_enabled() -> bool:
+    return os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0"
+
+
+def launch_config(rows: int = 1) -> Tuple[Tuple[int, int], Tuple[int, int], int]:
+    """(FC1 ctas, warps), (FC2 ctas, warps) and the flag bits for calls with ``rows`` rows
+    (``DEFAULT_FC1``, ``DEFAULT_FC2``, ``DEFAULT_FLAGS``)."""
+    pick = 0 if rows == 1 else 1
+    flags = DEFAULT_FLAGS[pick] & FLAG_OPTIONS
+    if _pdl_enabled():
+        flags |= FLAG_PDL
+    fc1 = DEFAULT_FC1[pick]
+    # (The TMA-fed kernels choose their ring depth in the launcher.)
+    if (
+        flags & FLAG_DEEP_RING
+        and not flags & FLAG_TMA
+        and rows * HIDDEN * 2 + fc1[1] * FC1_RING_CHUNKS_DEEP * 2048 > FC1_DYNAMIC_SMEM_BUDGET
+    ):
+        flags &= ~FLAG_DEEP_RING
+    fc2 = DEFAULT_FC2[pick]
+    if fc1[1] == 3 or (fc2[1] == 3 and not flags & FLAG_TMA):
+        raise ValueError("3 warps per CTA exist only for the TMA-fed FC2 kernel")
+    return fc1, fc2, flags
+
+
+class _Scratch:
+    """Per-device scratch of the two kernels (routing record and intermediate)."""
+
+    def __init__(self, device: torch.device):
+        self.h = torch.empty(SCRATCH_H_ELEMS, dtype=torch.bfloat16, device=device)
+        self.route = torch.empty(ROUTE_WORDS, dtype=torch.int32, device=device)
+
+
+# One scratch per device, shared by every draft_routed_moe call on it: correct only while those
+# calls are ordered on one stream (the MTP layer's draft steps run in sequence on the model's main
+# stream, eagerly and in every captured graph, and graphs replay one at a time). Calls that may
+# overlap on other streams would each need their own scratch.
+_SCRATCH: Dict[int, _Scratch] = {}
+
+
+def _scratch(device: torch.device, create: bool) -> Optional[_Scratch]:
+    key = device.index if device.index is not None else torch.cuda.current_device()
+    if key not in _SCRATCH and create:
+        _SCRATCH[key] = _Scratch(torch.device("cuda", key))
+    return _SCRATCH.get(key)
+
+
+def resolve_owner(module) -> Tuple[Optional[torch.nn.Module], str]:
+    """The leaf owning the BF16 expert parameters (``ConfigurableMoE.backend`` or the leaf)."""
+    backend = getattr(module, "backend", None)
+    owner = backend if isinstance(backend, torch.nn.Module) else module
+    if not isinstance(owner, torch.nn.Module):
+        return None, f"{type(module).__name__} is not a module"
+    if "w3_w1_weight" not in owner._parameters or "w2_weight" not in owner._parameters:
+        return None, f"{type(owner).__name__} owns no expert parameters"
+    return owner, ""
+
+
+def _wrapper_reason(wrapper) -> str:
+    """Wrapper-level behaviour the draft_routed_moe call would skip; empty when none."""
+    if getattr(wrapper, "comm", None) is not None:
+        return "MoE communication strategy"
+    if getattr(wrapper, "layer_load_balancer", None) is not None:
+        return "EPLB load balancer"
+    if getattr(wrapper, "enable_dwdp", False):
+        return "DWDP"
+    if getattr(wrapper, "unfinalized_combine_fn", None) is not None:
+        return "unfinalized_combine_fn"
+    if getattr(wrapper, "apply_router_weight_on_input", False):
+        return "router weight on input"
+    if getattr(wrapper, "_enable_perfect_router", False):
+        return "perfect router"
+    if getattr(wrapper, "reduce_results", False) and getattr(wrapper, "parallel_size", 1) > 1:
+        return "in-op all-reduce"
+    return ""
+
+
+def _routing_reason(rm) -> str:
+    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
+
+    if not isinstance(rm, DeepSeekV3MoeRoutingMethod):
+        return f"routing method {type(rm).__name__}"
+    impl = rm.routing_impl
+    groups = (impl.top_k, impl.n_group, impl.topk_group)
+    if groups != (TOP_K, 1, 1):
+        return f"routing {groups}"
+    scale = impl.routed_scaling_factor
+    if (
+        scale is None
+        or not (float(scale) == float(scale))
+        or float(scale) in (float("inf"), float("-inf"))
+    ):
+        return f"routed scaling {scale}"
+    return ""
+
+
+def _owner_reason(owner, device: torch.device) -> str:
+    """Leaf-level eligibility: FlashInfer TRT-LLM-Gen BF16 leaf, TP shard with every expert local."""
+    from tensorrt_llm._torch.moe.fused_moe.quantization import BF16TRTLLMGenFusedMoEMethod
+    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.flashinfer_bf16 import (
+        FlashinferTrtllmGenBf16Impl,
+    )
+    from tensorrt_llm._torch.utils import ActivationType
+
+    if not isinstance(owner, FlashinferTrtllmGenBf16Impl):
+        return f"{type(owner).__name__} is not the FlashInfer TRT-LLM-Gen BF16 leaf"
+    if getattr(owner, "use_dp", False):
+        return "attention DP"
+    routes_outside = getattr(owner, "_routes_outside_the_kernel", None)
+    if routes_outside is None or routes_outside():
+        return "separated routing"
+    if getattr(owner, "activation_type", None) != ActivationType.Swiglu:
+        return f"activation {getattr(owner, 'activation_type', None)}"
+    if getattr(owner, "bias", False):
+        return "expert bias"
+    for name in ("act_alpha", "act_beta", "act_clamp"):
+        if getattr(owner, name, None) is not None:
+            return f"{name} set"
+    if (
+        getattr(owner, "num_slots", None) != NUM_EXPERTS
+        or getattr(owner, "slot_start", None) != 0
+        or getattr(owner, "expert_size_per_partition", None) != NUM_EXPERTS
+    ):
+        return "not all 256 experts local (EP or EPLB slots)"
+    if getattr(owner, "hidden_size", None) != HIDDEN:
+        return f"hidden {getattr(owner, 'hidden_size', None)}"
+    if getattr(owner, "intermediate_size_per_partition", None) != INTER:
+        return f"intermediate {getattr(owner, 'intermediate_size_per_partition', None)}"
+    qm = getattr(owner, "quant_method", None)
+    if not isinstance(qm, BF16TRTLLMGenFusedMoEMethod):
+        return f"quant method {type(qm).__name__}"
+    if not getattr(qm, "use_shuffled_weight", False) or getattr(qm, "block_k", None) != K_CHUNK:
+        return "weights not shuffled into 64-element BlockMajorK blocks"
+    from tensorrt_llm._torch.moe.fused_moe.quantization import (
+        TRTLLM_GEN_WEIGHT_LAYOUT_BLOCK_MAJOR_K,
+    )
+
+    if getattr(qm, "weight_layout", None) != TRTLLM_GEN_WEIGHT_LAYOUT_BLOCK_MAJOR_K:
+        return f"weight layout {getattr(qm, 'weight_layout', None)}"
+    if getattr(owner, "_trtllm_gen_layout_transform_pending", False):
+        return "weight layout transform pending"
+    w13 = owner._parameters["w3_w1_weight"]
+    w2 = owner._parameters["w2_weight"]
+    if (
+        tuple(w13.shape) != (NUM_EXPERTS, HIDDEN // K_CHUNK, 2 * INTER, K_CHUNK)
+        or w13.dtype != torch.bfloat16
+    ):
+        return f"w3_w1 {tuple(w13.shape)} {w13.dtype}"
+    if (
+        tuple(w2.shape) != (NUM_EXPERTS, INTER // K_CHUNK, HIDDEN, K_CHUNK)
+        or w2.dtype != torch.bfloat16
+    ):
+        return f"w2 {tuple(w2.shape)} {w2.dtype}"
+    for t in (w13, w2):
+        if not t.is_contiguous() or t.device != device:
+            return "expert buffers not contiguous on the activation's device"
+    return _routing_reason(getattr(owner, "routing_method", None))
+
+
+def eligibility(
+    module, routing_bias, x: torch.Tensor, router_logits: torch.Tensor, do_finalize: bool
+) -> Tuple[bool, str]:
+    """Decide draft_routed_moe support for this call from host-side state only."""
+    if not do_finalize:
+        return False, "unfinalized output requested"
+    if not torch.cuda.is_available() or x.device.type != "cuda":
+        return False, "not on CUDA"
+    major, _ = torch.cuda.get_device_capability(x.device)
+    if major != 10:
+        return False, f"SM{major}x is not the SM100 family"
+    if x.dtype != torch.bfloat16 or x.dim() != 2 or x.shape[1] != HIDDEN:
+        return False, f"activation {tuple(x.shape)} {x.dtype}"
+    rows = x.shape[0]
+    if rows < 1 or rows > MAX_ROWS:
+        return False, f"{rows} rows outside [1, {MAX_ROWS}]"
+    if not x.is_contiguous() or router_logits is None or not router_logits.is_contiguous():
+        return False, "non-contiguous input"
+    if (
+        router_logits.dtype != torch.float32
+        or tuple(router_logits.shape) != (rows, NUM_EXPERTS)
+        or router_logits.device != x.device
+    ):
+        return False, f"logits {tuple(router_logits.shape)} {router_logits.dtype}"
+    if torch.compiler.is_compiling():
+        return False, "torch.compile tracing"
+    if not op_available():
+        return False, f"torch.ops.trtllm.{OP_NAME} is not built"
+    owner, reason = resolve_owner(module)
+    if owner is None:
+        return False, reason
+    if owner is not module:
+        reason = _wrapper_reason(module)
+        if reason:
+            return False, reason
+    reason = _owner_reason(owner, x.device)
+    if reason:
+        return False, reason
+    native_bias = getattr(getattr(owner, "routing_method", None), "e_score_correction_bias", None)
+    if (
+        routing_bias is None
+        or routing_bias.dtype != torch.bfloat16
+        or tuple(routing_bias.shape) != (NUM_EXPERTS,)
+        or not routing_bias.is_contiguous()
+        or routing_bias.device != x.device
+    ):
+        return False, "router bias is not a contiguous BF16 [256] on the device"
+    if native_bias is None or native_bias.data_ptr() != routing_bias.data_ptr():
+        return False, "router bias is not the routing method's bias"
+    if torch.cuda.is_current_stream_capturing() and _scratch(x.device, create=False) is None:
+        # Allocating inside a capture would record the allocation into the graph; the model's eager warmup
+        # forwards create the scratch before capture.
+        return False, "not warmed up before CUDA graph capture"
+    return True, ""
+
+
+def run(owner, routing_bias, x, router_logits, debug_fc2=None, timeline=None) -> torch.Tensor:
+    """Launch the two kernels (no eligibility check; ``owner`` is the BF16 leaf)."""
+    (c1, w1), (c2, w2), flags = launch_config(x.shape[0])
+    scratch = _scratch(x.device, create=True)
+    # The kernels read the activation, the logits and the bias with 16-byte loads.
+    if x.data_ptr() % 16 != 0:
+        x = x.clone(memory_format=torch.contiguous_format)
+    if router_logits.data_ptr() % 16 != 0:
+        router_logits = router_logits.clone(memory_format=torch.contiguous_format)
+    if routing_bias.data_ptr() % 16 != 0:
+        routing_bias = routing_bias.clone(memory_format=torch.contiguous_format)
+    rm = owner.routing_method
+    return getattr(torch.ops.trtllm, OP_NAME)(
+        x,
+        router_logits,
+        routing_bias,
+        owner._parameters["w3_w1_weight"],
+        owner._parameters["w2_weight"],
+        scratch.h,
+        scratch.route,
+        float(rm.routing_impl.routed_scaling_factor),
+        c1,
+        w1,
+        c2,
+        w2,
+        flags,
+        debug_fc2,
+        timeline,
+    )
+
+
+def routed_output(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    do_finalize: bool,
+) -> Optional[torch.Tensor]:
+    """Entry point of the MTP routed-MoE call site; ``None`` means "use the native path".
+
+    ``module`` is the served ``ConfigurableMoE`` wrapper (or the bare BF16 leaf), ``routing_bias`` the gate's
+    ``e_score_correction_bias``. The decision is made before any device work.
+    """
+    from ....shared import dispatch
+
+    ok, reason = eligibility(module, routing_bias, x, router_logits, do_finalize)
+    dispatch.log_capability(
+        "draft_routed_moe", f"rows={x.shape[0]} {reason or 'engaged'}", ok, reason
+    )
+    if not ok:
+        return None
+    owner, _ = resolve_owner(module)
+    return run(owner, routing_bias, x, router_logits)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/__init__.py
new file mode 100644
index 000000000000..17274028d3bb
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/__init__.py
@@ -0,0 +1,46 @@
+# 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.
+"""Package hop_routing: native top-8 routing in the pre-MoE into-hop entry.
+
+Entry points registered with ``dispatch``:
+
+* ``make_transport(into_hop)`` returns the :class:`~.hop.HopRoutingTransport` of an
+  ``attention_output_into_hop`` transport (or None, the into-hop path). Call it at the same point
+  of model construction on every rank.
+* ``engagement_counts()`` returns the host-side routing-hop call counters.
+* ``routed_output_into_hop(module, routing_bias, x, router_logits, routing,
+  hop)`` and ``routed_output(module, routing_bias, x, router_logits, routing,
+  do_finalize)``: routed_moe's entries of the same names with the hop's routing
+  (``routing``, int32 ``[2, R, 8]``) in place of routed_moe's own; ``None`` means
+  "run routed_moe on ``router_logits``" (decided before any device work).
+"""
+
+from ....shared.dispatch import register
+from .hop import HopRoutingTransport, engagement_counts, op_available, split_routing
+from .routed_moe import routed_output, routed_output_into_hop
+
+register("hop_routing", "make_transport", HopRoutingTransport.create)
+register("hop_routing", "engagement_counts", engagement_counts)
+register("hop_routing", "routed_output", routed_output)
+register("hop_routing", "routed_output_into_hop", routed_output_into_hop)
+
+__all__ = [
+    "HopRoutingTransport",
+    "engagement_counts",
+    "op_available",
+    "routed_output",
+    "routed_output_into_hop",
+    "split_routing",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop.py
new file mode 100644
index 000000000000..4f9e12391d26
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop.py
@@ -0,0 +1,251 @@
+# 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.
+"""Routing hop of hop_routing: the into-hop entry plus native top-8 routing.
+
+``HopRoutingTransport`` wraps the ``attention_output_into_hop`` :class:`IntoHopTransport` of a
+``post_attention_hop``. The producer side is the into-hop entry's unchanged (the attention tail
+publishes into the same transport); the hop runs ``trtllm::glm_moe_dsa_post_attention_hop_routing``,
+which returns the into-hop entry's normalized hidden, residual and router logits bit for bit plus
+the native DeepSeek-V3 routing of every row (expert ids int32 ``[R, 8]`` and
+weights FP32 ``[R, 8]`` holding the bf16-rounded native weights, descending
+native key order). See
+``cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h`` for the
+arithmetic, the gather protocol and the memory-order contract.
+
+The gather buffers and the parity word are private to hop_routing and shared by every
+layer that uses the same ``post_attention_hop`` (one TP group, one workspace); they are
+allocated once, at model construction.
+"""
+
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from ....shared.dispatch import log_capability
+
+FEATURE = "hop_routing"
+_OP_NAME = "glm_moe_dsa_post_attention_hop_routing"
+HIDDEN_SIZE = 6144
+NUM_EXPERTS = 256
+TOP_K = 8
+MAX_ROWS = 8
+NUM_RANKS = 8
+ROUTED_SCALING = 2.5
+GATHER_WORDS = 2 * MAX_ROWS * NUM_EXPERTS
+STATE_WORDS = 4
+NUM_STAMPS = 8
+MAX_GRID_CTAS = 96
+# The empty gather slot (0x80000000 as int32).
+GATHER_SENTINEL = -(1 << 31)
+
+_counts: Dict[str, int] = {}
+
+
+def _count(key: str) -> None:
+    _counts[key] = _counts.get(key, 0) + 1
+
+
+def engagement_counts() -> Dict[str, int]:
+    """Copy of the host-side call counters (eager calls and CUDA graph
+    captures, not replays): ``hop_routing_hop`` counts routing-hop calls."""
+    return dict(_counts)
+
+
+def reset_engagement_counts() -> None:
+    _counts.clear()
+
+
+def op_available(op_name: str = _OP_NAME) -> bool:
+    """Whether ``trtllm::<op_name>`` (default: the routing hop) is loaded."""
+    return hasattr(torch.ops.trtllm, op_name)
+
+
+def _agreed_op_reason(hop, op_name: str = _OP_NAME) -> str:
+    """Why no rank of ``hop``'s group may use the routing hop ("" when every
+    rank may). Agreed once per hop and op over the hop's TP communicator."""
+    cache = hop.__dict__.setdefault("_hop_routing_op_reason", {})
+    if op_name in cache:
+        return cache[op_name]
+    local = "" if op_available(op_name) else f"trtllm::{op_name} is not loaded"
+    comm = getattr(hop, "comm", None)
+    agreed = comm is None or comm.allreduce(0 if local else 1) == comm.Get_size()
+    reason = local or ("" if agreed else f"trtllm::{op_name} is not loaded on another rank")
+    cache[op_name] = reason
+    return reason
+
+
+class _GatherState:
+    """Gather buffers and parity of one ``post_attention_hop`` workspace."""
+
+    def __init__(self, device: torch.device):
+        self.gather = torch.full((GATHER_WORDS,), GATHER_SENTINEL, dtype=torch.int32, device=device)
+        self.zstate = torch.zeros(STATE_WORDS, dtype=torch.int32, device=device)
+
+
+def _gather_state(hop) -> _GatherState:
+    state = getattr(hop, "_hop_routing_gather_state", None)
+    if state is None:
+        state = _GatherState(hop.flags.device)
+        hop._hop_routing_gather_state = state
+    return state
+
+
+class HopRoutingTransport:
+    """into-hop producer view plus the routing hop."""
+
+    def __init__(self, into_hop, op_name: str = _OP_NAME):
+        self.into_hop = into_hop
+        self.hop = into_hop.hop
+        self.rank = into_hop.rank
+        self.flags = into_hop.flags
+        self.mc_words = into_hop.mc_words
+        state = _gather_state(self.hop)
+        self.gather = state.gather
+        self.zstate = state.zstate
+        # Same schedule knobs as the into-hop entry (the hop arithmetic and protocol are the
+        # into-hop entry's).
+        self.early_poll = into_hop.early_poll
+        self.router_load = into_hop.router_load
+        # Tests may use the diagnostic build (``..._diag``, stamps).
+        self.op_name = op_name
+
+    @classmethod
+    def create(cls, into_hop, op_name: str = _OP_NAME) -> Optional["HopRoutingTransport"]:
+        """The routing hop of ``into_hop``'s hop, or None (into-hop path).
+
+        Rank-uniform: ``into_hop`` exists on all ranks of the group or on none, and
+        the op's availability is agreed over the hop's TP group the first
+        time, at the same point of model construction on every rank.
+        """
+        reason = (
+            "attention_output_into_hop transport unavailable"
+            if into_hop is None
+            else _agreed_op_reason(into_hop.hop, op_name)
+        )
+        log_capability(FEATURE, "routing_hop", not reason, reason)
+        return None if reason else cls(into_hop, op_name)
+
+    def can_run(
+        self,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        router_weight: Optional[torch.Tensor],
+        routing_bias: Optional[torch.Tensor],
+    ) -> bool:
+        """Host-only, rank-uniform eligibility of one call: the into-hop entry's, with router
+        logits, and a contiguous bf16 ``[256]`` router bias on the device."""
+        if router_weight is None or routing_bias is None:
+            return False
+        if (
+            routing_bias.dtype != torch.bfloat16
+            or tuple(routing_bias.shape) != (NUM_EXPERTS,)
+            or not routing_bias.is_contiguous()
+            or routing_bias.device != residual.device
+        ):
+            return False
+        return self.into_hop.can_run(residual, norm_weight, router_weight)
+
+    def __call__(
+        self,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+        router_weight: torch.Tensor,
+        routing_bias: torch.Tensor,
+        logits_in: Optional[torch.Tensor] = None,
+        stamps: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Run the routing hop on the partials the producer already published.
+
+        Returns ``(normed, residual_out, router_logits, routing)`` where
+        ``routing`` is int32 ``[2, R, 8]``: ``routing[0]`` the expert ids and
+        ``routing[1]`` the FP32 bits of the weights (see :func:`split_routing`).
+        ``logits_in`` (tests only) replaces the logits the routing reads.
+        """
+        _count("hop_routing_hop")
+        normed, residual_out, logits, routing = getattr(torch.ops.trtllm, self.op_name)(
+            residual.contiguous(),
+            norm_weight.contiguous(),
+            float(eps),
+            router_weight,
+            routing_bias,
+            self.hop.uc_buffer,
+            self.flags,
+            self.gather,
+            self.zstate,
+            self.rank,
+            NUM_RANKS,
+            self.early_poll,
+            self.router_load,
+            ROUTED_SCALING,
+            logits_in,
+            stamps,
+        )
+        return normed, residual_out, logits, routing
+
+
+def split_routing(routing: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """``(ids int32 [R, 8], weights fp32 [R, 8])`` views of a routing tensor."""
+    return routing[0], routing[1].view(torch.float32)
+
+
+def _register_fake() -> None:
+    if not op_available():
+        return
+    try:
+
+        @torch.library.register_fake(f"trtllm::{_OP_NAME}")
+        def _(
+            residual,
+            gamma,
+            eps,
+            router_weight,
+            routing_bias,
+            uc_buffer,
+            flags,
+            gather,
+            zstate,
+            rank,
+            nranks,
+            early_poll,
+            router_load,
+            routed_scaling,
+            logits_in=None,
+            stamps=None,
+        ):
+            rows = residual.numel() // HIDDEN_SIZE
+            return [
+                torch.empty_like(residual),
+                torch.empty_like(residual),
+                residual.new_empty((rows, NUM_EXPERTS), dtype=torch.float32),
+                residual.new_empty((2, rows, TOP_K), dtype=torch.int32),
+            ]
+    except RuntimeError:
+        # Already registered by an earlier import in this process.
+        pass
+
+
+_register_fake()
+
+__all__ = [
+    "GATHER_SENTINEL",
+    "NUM_STAMPS",
+    "MAX_GRID_CTAS",
+    "HopRoutingTransport",
+    "engagement_counts",
+    "op_available",
+    "reset_engagement_counts",
+    "split_routing",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/routed_moe.py
new file mode 100644
index 000000000000..391112f8835d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/routed_moe.py
@@ -0,0 +1,316 @@
+# 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.
+"""routed_moe fork of package hop_routing: the routed MoE consuming the hop's routing.
+
+The expert-cluster kernel of routed_moe (``routed_moe/kernel_xc.py``) with its routing
+phase replaced: instead of loading the router logits and selecting the top 8
+(one warp per token), warp ``t`` reads token ``t``'s eight expert ids and
+weights that the hop_routing hop wrote, fills the same shared-memory tables
+(``s_topk_id``, ``s_topk_w`` and the expert bitmap) and continues exactly as
+routed_moe. Given the same routing, every later step (slot tables, schedule, FC1,
+SwiGLU, quantization, FC2, the outputs of each mode) is routed_moe's code, so the
+numerics are routed_moe's.
+
+The routing operand is the hop's int32 ``[2, R, 8]`` tensor (ids, then the
+FP32 bits of the weights; see ``hop.split_routing``). It is passed in routed_moe's
+router-logits address slot; the weights start ``R * 8`` words after the ids.
+It is produced on the stream by the hop_routing hop and read only after
+``griddepcontrol.wait`` (routed_moe's own memory-order contract, see
+``routed_moe/routed_moe.py``).
+
+Before the wait nothing else changes, except under package cross_layer_expert_prefetch's speculative
+prefetch, which the fork compiles in from the runner key exactly as routed_moe's
+runner does (``spec_k``, ``spec_ctas``, ``spec_self`` of ``key[6]``). Then the
+TMA warp of each issuing CTA also reads the layer's row of predicted expert
+IDs (cross_layer_expert_prefetch's int32 row of ``SPEC_KMAX`` entries) with relaxed loads and issues
+L2 bulk prefetches of those experts' weights, which are model parameters
+that nothing writes while serving. With cross_layer_expert_prefetch's predictor ``self``
+(``spec_experts.SERVED_CONFIG``) the row has one writer: CTA 0's MMA warp of the
+same layer's fork, which stores the call's distinct experts ranked by row count
+once the routing tables are
+built, i.e. after that call's own dependency wait. A call's pre-wait read
+therefore sees the row its layer's previous call stored (that call completed
+before this launch, in stream order behind the intervening dependency
+waits), or, on a CTA that becomes resident after this call's CTA 0 has
+stored, this call's own experts. Either value only selects which immutable
+bytes are prefetched (each ID is masked into [0, 256)), never what the
+kernel reads or computes.
+
+Entry points (``None`` means "use the routed_moe path with the logits"; the
+decision is host-only and precedes any device work):
+
+* ``routed_output_into_hop(module, bias, x, logits, routing, hop)``: routed_moe_into_hop,
+  every routed-input mode of the post_moe_hop;
+* ``routed_output(module, bias, x, logits, routing, do_finalize)``.
+
+Only the expert-cluster schedule is forked; a call that routed_moe would run with
+its grid-wide schedule (``v6``) declines.
+"""
+
+import threading
+from typing import Dict, List, Optional
+
+import torch
+
+from ..routed_moe import routed_moe as routed_moe_mod
+
+FEATURE = "hop_routing"
+TOP_K = routed_moe_mod.TOP_K
+MAX_ROWS = routed_moe_mod.MAX_ROWS
+
+_LOCK = threading.Lock()
+_COMPILED: Dict[tuple, object] = {}
+_KERNEL_CLASS = None
+
+
+def _kernel_class():
+    """``HopRoutedExpertClusterKernel``, defined on first use (imports the CuTe DSL)."""
+    global _KERNEL_CLASS
+    if _KERNEL_CLASS is not None:
+        return _KERNEL_CLASS
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cutlass_dsl import Int32, Int64, T, Uint32, dsl_user_op
+
+    from ..routed_moe import ptx
+    from ..routed_moe.kernel_xc import RoutedMoeExpertClusterKernel
+
+    @dsl_user_op
+    def ld_nc_u32(addr: Int64, *, loc=None, ip=None) -> Uint32:
+        return Uint32(
+            llvm.inline_asm(
+                T.i32(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.global.nc.u32 $0, [$1];",
+                "=r,l",
+                has_side_effects=True,
+                is_align_stack=False,
+                asm_dialect=llvm.AsmDialect.AD_ATT,
+                loc=loc,
+                ip=ip,
+            )
+        )
+
+    class HopRoutedExpertClusterKernel(RoutedMoeExpertClusterKernel):
+        """routed_moe's expert-cluster kernel reading the routing instead of computing it."""
+
+        @cute.jit
+        def _route_tokens(
+            self, w, lane, num_tokens, logits_addr, bias_r, s_topk_id, s_topk_w, storage
+        ):
+            """Warp ``w`` takes token ``w``'s ids and weights from the hop's
+            routing (``logits_addr``: ids ``[R, 8]`` int32, then weights
+            ``[R, 8]`` FP32) and marks the experts in the bitmap."""
+            if w < num_tokens:
+                if lane < TOP_K:
+                    idx = w * TOP_K + lane
+                    e = Int32(ld_nc_u32(logits_addr + idx * 4))
+                    wt = ptx.ld_nc_f32(logits_addr + (num_tokens * TOP_K + idx) * 4)
+                    s_topk_id[idx] = e
+                    s_topk_w[idx] = wt
+                    ptx.atom_or_shared_u32(
+                        Int32(storage.bitmap.data_ptr().toint()) + (e // 32) * 4,
+                        Uint32(1) << Uint32(e % 32),
+                    )
+
+    del cutlass
+    _KERNEL_CLASS = HopRoutedExpertClusterKernel
+    return HopRoutedExpertClusterKernel
+
+
+class HopRoutedMoe(routed_moe_mod.RoutedMoe):
+    """routed_moe runner whose expert-cluster kernel consumes the hop's routing."""
+
+    def ready_for_capture(
+        self, finalize: bool, debug: int = 0, rows: int = MAX_ROWS, red: bool = False
+    ) -> bool:
+        key = self._key(finalize, debug, rows, red)
+        dev = self.device.index if self.device.index is not None else torch.cuda.current_device()
+        ws = routed_moe_mod._WORKSPACES.get(dev)
+        return key in _COMPILED and ws is not None and ws.num_ctas >= self.num_ctas
+
+    def _compiled(self, finalize: bool, debug: int, rows: int = MAX_ROWS, red: bool = False):
+        """routed_moe's compiled launchers with ``HopRoutedExpertClusterKernel`` (expert clusters only)."""
+        key = self._key(finalize, debug, rows, red)
+        if key[3] != "xc":
+            raise ValueError("the hop_routing fork runs the expert-cluster schedule only")
+        with _LOCK:
+            fns = _COMPILED.get(key)
+            if fns is not None:
+                return fns
+            import cutlass
+            import cutlass.cute as cute
+            from cuda.bindings import driver as cuda
+            from cutlass.cute.runtime import make_ptr
+
+            from ..routed_moe.kernel_xc import DELAY_FINALIZE, RoutedMoeFinalizeKernel
+
+            ptrs = self._pointers(make_ptr, cute, cutlass)
+            stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+            inkernel = finalize and self.inkernel_finalize()
+            # Package cross_layer_expert_prefetch's speculative expert prefetch
+            # (``RoutedMoe._key``).
+            spec = key[6] or (0, 0, False)
+            kern = _kernel_class()(
+                self._xc_ctas(rows, finalize, red),
+                debug=debug,
+                finalize=inkernel,
+                red_out=red,
+                **routed_moe_mod.xc_kernel_args(),
+                spec_k=spec[0],
+                spec_ctas=spec[1],
+                spec_self=spec[2],
+            )
+            ints = [cutlass.Int64(0)] * routed_moe_mod._NUM_XC_ADDRS
+            routed = cute.compile(kern.run, *ptrs, *ints, cutlass.Int32(1), stream)
+            if spec[0]:
+                from ....shared import dispatch
+
+                dispatch.log_capability(
+                    FEATURE,
+                    f"expert_prefetch[k={spec[0]},ctas={spec[1]},self={int(spec[2])}]",
+                    True,
+                )
+            fin = None
+            if finalize and not inkernel and not red:
+                delay, delay_ns = routed_moe_mod.DELAY_INJECTION
+                fin = cute.compile(
+                    RoutedMoeFinalizeKernel(delay_ns=delay_ns if delay & DELAY_FINALIZE else 0).run,
+                    cutlass.Int64(0),
+                    cutlass.Int64(0),
+                    cutlass.Int64(0),
+                    cutlass.Int32(1),
+                    stream,
+                )
+            fns = (routed, fin)
+            _COMPILED[key] = fns
+            return fns
+
+
+def routing_reason(routing: Optional[torch.Tensor], x: torch.Tensor) -> str:
+    """Why ``routing`` cannot feed the fork for activation ``x`` ("" if it can)."""
+    if routing is None:
+        return "no routing from the hop"
+    rows = x.shape[0]
+    if (
+        routing.dtype != torch.int32
+        or tuple(routing.shape) != (2, rows, TOP_K)
+        or not routing.is_contiguous()
+        or routing.device != x.device
+    ):
+        return f"routing {tuple(routing.shape)} {routing.dtype}"
+    return ""
+
+
+def _engaged_runner(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    routing: Optional[torch.Tensor],
+    do_finalize: bool,
+    red: bool = False,
+) -> Optional[HopRoutedMoe]:
+    """The fork's runner for this call, or ``None`` (host-only decision)."""
+    from ....shared import dispatch
+
+    reason = routing_reason(routing, x)
+    ok = not reason
+    if ok:
+        # routed_moe's own eligibility (shapes, dtypes, device, routing method, expert owner).
+        ok, reason = routed_moe_mod.eligibility(module, x, router_logits)
+    owner = routed_moe_mod.resolve_expert_owner(module)[0] if ok else None
+    if ok:
+        rm = getattr(owner, "routing_method", None)
+        native_bias = getattr(rm, "e_score_correction_bias", None)
+        if (
+            routing_bias is None
+            or routing_bias.dtype != torch.bfloat16
+            or tuple(routing_bias.shape) != (routed_moe_mod.NUM_EXPERTS,)
+            or not routing_bias.is_contiguous()
+            or routing_bias.device != x.device
+        ):
+            ok, reason = False, "router bias is not a contiguous BF16 [256] on the device"
+        elif native_bias is not None and native_bias.data_ptr() != routing_bias.data_ptr():
+            ok, reason = False, "router bias is not the routing method's bias"
+    runner = getattr(owner, "_hop_routing_runner", None) if ok else None
+    if ok and runner is not None and runner.bias.data_ptr() != routing_bias.data_ptr():
+        runner = None
+    if ok and runner is None:
+        runner = HopRoutedMoe(owner, routing_bias)
+    rows = x.shape[0]
+    if ok and runner.schedule(rows, do_finalize, red) != "xc":
+        ok, reason = False, "routed_moe would run its grid-wide schedule"
+    if ok and torch.cuda.is_current_stream_capturing():
+        if getattr(owner, "_hop_routing_runner", None) is None or not runner.ready_for_capture(
+            do_finalize, rows=rows, red=red
+        ):
+            ok, reason = False, "not warmed up before CUDA graph capture"
+    dispatch.log_capability(
+        FEATURE, f"routed_moe_fork rows={rows} {reason or 'engaged'}", ok, reason
+    )
+    if not ok:
+        return None
+    owner._hop_routing_runner = runner
+    return runner
+
+
+def routed_output(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    routing: Optional[torch.Tensor],
+    do_finalize: bool,
+) -> Optional[List[torch.Tensor]]:
+    """routed_moe's ``routed_output`` with the hop's routing; ``None``: use routed_moe."""
+    runner = _engaged_runner(module, routing_bias, x, router_logits, routing, do_finalize)
+    if runner is None:
+        return None
+    return runner(x, routing, do_finalize=do_finalize)
+
+
+def routed_output_into_hop(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    routing: Optional[torch.Tensor],
+    hop,
+) -> Optional[List[torch.Tensor]]:
+    """routed_moe's ``routed_output_into_hop`` (routed_moe_into_hop) with the hop's routing; same
+    modes and returns. ``None``: use routed_moe."""
+    if hop is None or not getattr(hop, "_routed_enabled", False):
+        return None
+    rows = x.shape[0]
+    if not hop.can_run(rows):
+        return None
+    mode = getattr(hop, "routed_mode", "row")
+    if mode == "row":
+        runner = _engaged_runner(module, routing_bias, x, router_logits, routing, True)
+        if runner is None:
+            return None
+        target = hop.routed_input_buffer(rows)
+        return [runner(x, routing, do_finalize=True, out=target)]
+    red = mode == "fp32"
+    runner = _engaged_runner(module, routing_bias, x, router_logits, routing, False, red)
+    if runner is None:
+        return None
+    if red:
+        return runner(x, routing, red_out=hop.routed_acc_buffer(rows))
+    rows_buf, scales_buf = hop.routed_raw_buffers(rows)
+    runner(x, routing, do_finalize=False, raw_out=(rows_buf, scales_buf))
+    return [rows_buf]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/__init__.py
new file mode 100644
index 000000000000..f06123668cf5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/__init__.py
@@ -0,0 +1,39 @@
+# 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.
+"""routed_moe: routed NVFP4 MoE (the routed_moe contract) as one CuTe DSL launch.
+
+Registered entry points (``dispatch.get_impl("routed_moe", name)``):
+
+* ``routed_output(module, routing_bias, x, router_logits, do_finalize)`` ->
+  finalized ``[R, 6144]`` BF16, the native unfinalized triple, or ``None``
+  when the call is ineligible (the caller then runs the native path; the
+  decision is taken before any device work). ``module`` is the served
+  ``ConfigurableMoE`` (``Deepseekv3MoE.experts``) or its TRT-LLM-Gen NVFP4
+  leaf.
+* ``eligibility(module, x, router_logits)`` -> ``(bool, reason)``.
+* ``routed_output_into_hop(module, routing_bias, x, router_logits, hop)``
+  (used only when the ``routed_moe_into_hop`` feature is selected): the finalized routed
+  output written into the post_moe_hop's routed-input buffer, as ``[routed]``, or
+  ``None``.
+"""
+
+from ....shared import dispatch
+from .routed_moe import eligibility, routed_output, routed_output_into_hop
+
+dispatch.register("routed_moe", "routed_output", routed_output)
+dispatch.register("routed_moe", "eligibility", eligibility)
+dispatch.register("routed_moe", "routed_output_into_hop", routed_output_into_hop)
+
+__all__ = ["eligibility", "routed_output", "routed_output_into_hop"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel.py
new file mode 100644
index 000000000000..f252c60a3f2f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel.py
@@ -0,0 +1,1987 @@
+# 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.
+"""Single-launch routed NVFP4 MoE for small-M decode (CuTe DSL, SM100 family).
+
+One persistent kernel computes, for R <= 8 rows of one rank's TP shard of a
+256-expert / top-8 / intermediate-256 / hidden-6144 NVFP4 MoE:
+
+1. routing (sigmoid + bias top-8, scaled-sum normalization), recomputed by
+   every CTA from the FP32 router logits, so no CTA waits for another;
+2. NVFP4 quantization of the activation rows (bulk-copied to shared memory
+   before any weight traffic) into the MMA B operand, all R rows: unused
+   columns are masked in the epilogues;
+3. FC1 (gate/up) as a split-K partition of all (expert, 128-row tile,
+   256-element k-block) work over the grid: TMA streams the stock weight and
+   scale buffers through a 4-stage ring straight into tcgen05 block-scaled
+   MMAs (M=128 weight rows, N=8 rows of tokens, K=64 per instruction);
+4. the FC1 epilogue (split-K reduction, SwiGLU, NVFP4 quantization of the
+   intermediate) publishes a per-expert FC2 operand;
+5. FC2 as a split partition of (128-row output tile, expert) units; its
+   weight TMA starts as soon as the CTA's FC1 loads are issued;
+6. the top-8 combine (finalize) in the FC2 epilogue, or the native
+   unfinalized tuple.
+
+Warp roles: 0-3 quantize, then run the epilogues; 4 issues MMAs; 5 issues
+TMA; 6-7 quantize, then fetch the FC2 operands.
+
+Cross-CTA hand-offs are kept to one round trip on the consumer side. A
+finished FC1 tile is stored straight into its expert's FC2 operand image
+(the exact shared-memory layout of one B / SFB stage) and published with an
+epoch flag; each consumer CTA waits for the flags of its experts and fetches
+every image with two bulk copies that complete on a per-expert mbarrier.
+Split-tile partials move as LL words (8 bytes ``{payload, epoch}`` per
+16-byte store, polled by the reader). A split tile is reduced by a fixed
+owner (the CTA holding its first k-block or expert, which it processes last)
+in CTA order, so results are deterministic. Each CTA keeps its own launch
+counter as the epoch; all CTAs of a launch hold the same value without
+synchronizing. The weight stream is loaded with an evict-first L2 policy.
+
+The weight buffers are consumed in the stock TRT-LLM-Gen layout: gate/up rows
+interleaved and shuffled in 32-row blocks, scale factors in the 128x4
+interleaved layout. The epilogues undo the row shuffle with index arithmetic.
+"""
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass.cute.nvgpu import tcgen05
+from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32
+
+from . import ptx
+
+HIDDEN = 6144
+INTER = 256
+NUM_EXPERTS = 256
+TOP_K = 8
+NTOK = 8
+KTILE = 256
+FC1_KB = HIDDEN // KTILE  # 24 k-blocks per FC1 tile
+FC1_MT = 2 * INTER // 128  # 4 FC1 tiles per expert
+FC2_MT = HIDDEN // 128  # 48 FC2 tiles
+MAX_SLOTS = NTOK * TOP_K  # 64 distinct experts at most
+MAX_LOCAL_SLOTS = 24  # FC2 operand cache slots per CTA
+ROUTED_SCALING = 2.5
+
+# Stock weight geometry in bytes (L2 prefetch address arithmetic).
+W13_ROW_BYTES = HIDDEN // 2
+W13_EXPERT_BYTES = 2 * INTER * W13_ROW_BYTES
+W13_SF_TILE_BYTES = 128 * HIDDEN // 16
+W13_SF_EXPERT_BYTES = FC1_MT * W13_SF_TILE_BYTES
+W13_SF_KB_BYTES = 128 * KTILE // 16
+KB_ROW_BYTES = KTILE // 2
+W2_TILE_BYTES = 128 * INTER // 2
+W2_EXPERT_BYTES = HIDDEN * INTER // 2
+W2_SF_TILE_BYTES = 128 * INTER // 16
+W2_SF_EXPERT_BYTES = HIDDEN * INTER // 16
+
+# Workspace: one launch-counter (epoch) word per CTA.
+WS_WORDS = 1024
+# LL exchange buffers (8-byte words {payload u32, epoch u32}).
+LL_PART_BYTES = 128 * NTOK * 8  # one 128x8 FP32 partial
+# Per-expert FC2 operand, stored as the exact shared-memory image of one B
+# stage (8 rows x 128 B FP4, 16-byte chunks swizzled by row) followed by one
+# compact SFB stage, so a consumer fetches it with two bulk copies.
+HIMG_B_BYTES = NTOK * INTER // 2
+HIMG_SF_BYTES = 4 * 128
+HIMG_BYTES = HIMG_B_BYTES + HIMG_SF_BYTES
+# LL form of the same operand (hand-off mode 1): 8-byte {4 payload bytes,
+# epoch} words; words n*32 + j hold FP4 bytes 4j..4j+3 of row n, words
+# 256 + n*4 + mt hold tile mt's four scale bytes of row n.
+H_LL_DATA_WORDS = NTOK * INTER // 8
+H_LL_WORDS = H_LL_DATA_WORDS + NTOK * FC1_MT
+H_LL_BYTES = H_LL_WORDS * 8
+H_LL_ITEMS = H_LL_BYTES // 16
+H_LL_DATA_ITEMS = H_LL_DATA_WORDS // 2
+HBUF_BYTES = max(HIMG_BYTES, H_LL_BYTES)  # per-expert hand-off buffer
+# Shared staging of one FC1 tile's output: 8 rows x 32 FP4 bytes, then
+# 8 rows x 4 scale bytes.
+HSTAGE_SF = NTOK * 32
+HSTAGE_BYTES = HSTAGE_SF + NTOK * 4
+MAX_PEERS = 7  # other CTAs contributing to one split tile
+GATHER_PEERS = 4  # peers' partials fetched per round trip
+
+# Warp roles
+EPI_WARPS = 4
+MMA_WARP = 4
+TMA_WARP = 5
+NUM_WARPS = 8
+NUM_THREADS = 32 * NUM_WARPS
+QUANT_WARPS = 6  # warps 0-3 and 6-7 quantize
+GATHER_WARP0 = 6  # warps 6-7 fetch the FC2 operands
+GATHER_THREADS = 64
+
+BAR_QUANT = 3  # quantizing warps (debug dump only)
+BAR_EPI = 4  # warps 0-3
+BAR_GATHER = 5  # warps 6-7
+
+ACC_STAGES = 4
+TMEM_COLS = 128
+TMEM_SFA_COL = 64
+TMEM_SFB_COL = 80
+# Compact SFB: only token rows 0..7 of each 128-row scale chunk are read by
+# the N=8 MMA, so consecutive 64-K chunks are stored 128 B apart (overlapping
+# the unused rows) instead of 512 B.
+SFB_CHUNK_STRIDE = 128
+SFB_STAGE_STRIDE = 4 * SFB_CHUNK_STRIDE
+SFB_BYTES = (MAX_LOCAL_SLOTS - 1) * SFB_STAGE_STRIDE + 3 * SFB_CHUNK_STRIDE + 512
+B2_BAR = FC1_KB  # first of MAX_LOCAL_SLOTS FC2 operand barriers
+X_BAR = B2_BAR + MAX_LOCAL_SLOTS  # the staged activation rows
+FC1_DONE_BAR = X_BAR + 1  # all FC1 MMAs of this CTA have completed
+STG_BAR = FC1_DONE_BAR + 1  # LL operand staged in sX (mode 1), per slot
+NUM_BAR = STG_BAR + MAX_LOCAL_SLOTS
+A_STAGE_BYTES = 128 * KTILE // 2
+SFA_STAGE_BYTES = 128 * KTILE // 16
+# L2 eviction policy for the weight stream (read once): evict first, so the
+# kernel's code, activations and exchange buffers stay in L2.
+L2_EVICT_FIRST = 0x12F0000000000000
+X_ROW_BYTES = HIDDEN * 2
+
+# L2 prefetch modes (bit mask, ``RoutedMoeKernel(prefetch=...)``; off in this target)
+PF_FC1 = 1  # rest of this CTA's FC1 rows, once the ring is primed
+PF_FC2_EARLY = 2  # this CTA's FC2 tiles, right after the FC1 prefetch
+PF_FC2_LATE = 4  # this CTA's FC2 tiles, after the last FC1 TMA
+
+# Shared misc words
+MISC_NUM_SLOTS = 0
+
+DEBUG_ROUTING = 1
+DEBUG_QUANT = 2
+DEBUG_FC1 = 4
+DEBUG_TIMELINE = 8
+
+# Stress-test delay injection sites (``delay`` bits; off in production). A
+# pseudo-random third of the CTAs, re-drawn every launch from the epoch,
+# spins 1..6 units of ``delay_ns`` on the global timer at each enabled site.
+DELAY_ENTRY = 1  # whole CTA, right after the launch epoch is read
+DELAY_PRODUCER = 2  # before publishing an FC1 operand tile or a split partial
+DELAY_OWNER = 4  # split-tile owner, before it fetches its peers' partials
+DELAY_CONSUMER = 8  # FC2 operand fetch, before polling the FC1 flags
+TL_BYTE_BASE = 1600000
+TL_SLOTS = 48
+# Stamps are SM clock cycles; slot TL_GT0 holds the global timer read next
+# to the entry stamp so the host can place every CTA on one time axis.
+TL_GT0 = TL_SLOTS - 1
+(
+    TL_ENTRY,
+    TL_PDL,
+    TL_ROUTED,
+    TL_Q_START,
+    TL_Q_G0_COMPUTE,
+    TL_Q_G0_FENCE,
+    TL_Q_G1_FENCE,
+    TL_Q_G2_FENCE,
+    TL_Q_DONE,
+    TL_MMA_FC1_START,
+    TL_TMA_FC1,
+    TL_PF_ISSUED,
+    TL_MMA_FC1_DONE,
+    TL_EPI_FIRST_ACC,
+    TL_EPI_FC1_GATHERED,
+    TL_EPI_FC1_DONE,
+    TL_FC2_LL_ISSUED,
+    TL_FC2_LL_READY,
+    TL_B2_READY,
+    TL_TMA_FC2,
+    TL_MMA_FC2_DONE,
+    TL_FC2_FIRST_ACC,
+    TL_FC2_GATHERED,
+    TL_FC2_DONE,
+    TL_EPI_START,
+    TL_FC1_MMA_COMPLETE,
+    TL_FC2_SENTINELS,
+    TL_EPI_T_COMPUTED,
+    TL_EPI_T_STORED,
+    TL_EPI_T_TMEM,
+    TL_EPI_T_PARAMS,
+    TL_G_STS_DONE,
+    TL_G_FENCE_DONE,
+    TL_G_BAR0_DONE,
+    TL_FC2_LAST_SLOT,
+) = range(35)
+
+
+def _stamp(enabled, tl_s, idx):
+    # Timeline stamps go to shared memory (a global store here would wait in
+    # the same queues the stamps are meant to measure) and are flushed once
+    # at the end of the kernel.
+    if enabled:
+        ptx.sts_u64(tl_s + idx * 8, ptx.clock64())
+
+
+class RoutedMoeKernel:
+    """CuTe DSL kernel object; one compiled instance per configuration."""
+
+    def __init__(
+        self,
+        num_ctas: int,
+        finalize: bool = True,
+        num_stages: int = 4,
+        prefetch: int = 0,
+        min_fc1_kb: int = 12,
+        acc_stages: int = 4,
+        handoff: int = 0,
+        delay: int = 0,
+        delay_ns: int = 2000,
+        debug: int = 0,
+    ):
+        assert num_ctas <= WS_WORDS
+        # A split FC1 tile spans at most ceil(24 / min_fc1_kb) + 1 CTAs.
+        assert (FC1_KB + min_fc1_kb - 1) // min_fc1_kb <= MAX_PEERS
+        self.num_ctas = num_ctas
+        self.finalize = finalize
+        self.num_stages = num_stages
+        self.prefetch = prefetch
+        self.min_fc1_kb = min_fc1_kb
+        # TMEM: acc_stages x 8 accumulator columns, then the scale columns
+        assert acc_stages * NTOK <= TMEM_SFA_COL
+        self.acc_stages = acc_stages
+        # FC1 -> FC2 operand hand-off: 0 = operand image + per-tile release
+        # flag, bulk-copied into place; 1 = LL words (no producer fence),
+        # bulk-copied into the freed activation buffer and validated there
+        self.handoff = handoff
+        # Stress testing only: DELAY_* site bits and the delay unit
+        self.delay = delay
+        self.delay_ns = delay_ns
+        self.debug = debug
+        self.sf_vec_size = 16
+        self.acc_dtype = Float32
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.mma_inst_mn = (128, NTOK)
+        self.mma_tiler = (128, NTOK, KTILE)
+
+    def _setup(self):
+        self.tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+            cutlass.Float4E2M1FN,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.K,
+            cutlass.Float8E4M3FN,
+            self.sf_vec_size,
+            self.cta_group,
+            self.mma_inst_mn,
+        )
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((1, 1, 1)), (self.tiled_mma.thr_id.shape,)
+        )
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            self.tiled_mma, self.mma_tiler, cutlass.Float4E2M1FN, self.num_stages
+        )
+        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+            self.tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_stages
+        )
+        # The B operand is resident: one "stage" per FC1 k-block (FC1) and
+        # one per local expert slot (FC2); both need 24.
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            self.tiled_mma, self.mma_tiler, cutlass.Float4E2M1FN, MAX_LOCAL_SLOTS
+        )
+        # Standard SFB layout (defines the TMEM layout) and the compact SMEM
+        # layout actually used for storage (see SFB_CHUNK_STRIDE).
+        self.sfb_smem_layout_std = blockscaled_utils.make_smem_layout_sfb(
+            self.tiled_mma, self.mma_tiler, self.sf_vec_size, MAX_LOCAL_SLOTS
+        )
+        self.sfb_smem_layout_staged = cute.make_layout(
+            self.sfb_smem_layout_std.shape,
+            stride=((((16, 4), 0), (0, 1)), 0, SFB_CHUNK_STRIDE, SFB_STAGE_STRIDE),
+        )
+        self.num_tma_load_bytes = A_STAGE_BYTES + SFA_STAGE_BYTES
+
+    @cute.jit
+    def run(
+        self,
+        w13_ptr: cute.Pointer,
+        w13_sf_ptr: cute.Pointer,
+        w2_ptr: cute.Pointer,
+        w2_sf_ptr: cute.Pointer,
+        x_addr: Int64,
+        logits_addr: Int64,
+        bias_addr: Int64,
+        alpha1_addr: Int64,
+        scale_c1_addr: Int64,
+        alpha2_addr: Int64,
+        xgs_addr: Int64,
+        out_addr: Int64,
+        raw_scale_addr: Int64,
+        raw_idx_addr: Int64,
+        ws_addr: Int64,
+        part_addr: Int64,
+        himg_addr: Int64,
+        hflag_addr: Int64,
+        dbg_addr: Int64,
+        w13_addr: Int64,
+        w13_sf_addr: Int64,
+        w2_addr: Int64,
+        w2_sf_addr: Int64,
+        num_tokens: Int32,
+        stream,
+    ):
+        self._setup()
+        print(f"[routed_moe] a_smem {self.a_smem_layout_staged}")
+        print(f"[routed_moe] b_smem {self.b_smem_layout_staged}")
+        print(f"[routed_moe] sfa_smem {self.sfa_smem_layout_staged}")
+        print(f"[routed_moe] sfb_smem {self.sfb_smem_layout_std} -> {self.sfb_smem_layout_staged}")
+        # A operands: [rows, K, experts], K-major FP4.
+        w13 = cute.make_tensor(
+            w13_ptr, cute.make_ordered_layout((2 * INTER, HIDDEN, NUM_EXPERTS), order=(1, 0, 2))
+        )
+        w2 = cute.make_tensor(
+            w2_ptr, cute.make_ordered_layout((HIDDEN, INTER, NUM_EXPERTS), order=(1, 0, 2))
+        )
+        w13_sf = cute.make_tensor(
+            w13_sf_ptr, blockscaled_utils.tile_atom_to_shape_SF(w13.shape, self.sf_vec_size)
+        )
+        w2_sf = cute.make_tensor(
+            w2_sf_ptr, blockscaled_utils.tile_atom_to_shape_SF(w2.shape, self.sf_vec_size)
+        )
+
+        a_op = sm100_utils.cluster_shape_to_tma_atom_A((1, 1), self.tiled_mma.thr_id)
+        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+        sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
+        tma_atom_a1, tma_a1 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op, w13, a_smem_layout, self.mma_tiler, self.tiled_mma, self.cluster_layout_vmnk.shape
+        )
+        tma_atom_sfa1, tma_sfa1 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            w13_sf,
+            sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+        tma_atom_a2, tma_a2 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op, w2, a_smem_layout, self.mma_tiler, self.tiled_mma, self.cluster_layout_vmnk.shape
+        )
+        tma_atom_sfa2, tma_sfa2 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            w2_sf,
+            sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            ring_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.num_stages]
+            acc_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.acc_stages]
+            b_mbar: cute.struct.MemRange[cutlass.Int64, NUM_BAR]
+            tmem_holding: cutlass.Int32
+            tl_buf: cute.struct.MemRange[cutlass.Int64, TL_SLOTS]
+            hstage: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, HSTAGE_BYTES // 8], 16]
+            misc: cute.struct.MemRange[cutlass.Int32, 16]
+            bitmap: cute.struct.MemRange[cutlass.Uint32, 8]
+            topk_id: cute.struct.MemRange[cutlass.Int32, NTOK * TOP_K]
+            topk_w: cute.struct.MemRange[cutlass.Float32, NTOK * TOP_K]
+            slot_expert: cute.struct.MemRange[cutlass.Int32, MAX_SLOTS]
+            slot_mask: cute.struct.MemRange[cutlass.Uint32, MAX_SLOTS]
+            slot_w: cute.struct.MemRange[cutlass.Float32, MAX_SLOTS * NTOK]
+            slot_k: cute.struct.MemRange[cutlass.Int32, MAX_SLOTS * NTOK]
+            slot_alpha1: cute.struct.MemRange[cutlass.Float32, MAX_SLOTS]
+            slot_scale_c1: cute.struct.MemRange[cutlass.Float32, MAX_SLOTS]
+            slot_alpha2: cute.struct.MemRange[cutlass.Float32, MAX_SLOTS]
+            expert_slot: cute.struct.MemRange[cutlass.Int32, NUM_EXPERTS]
+            sA: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.a_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFA: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float8E4M3FN, cute.cosize(self.sfa_smem_layout_staged)
+                ],
+                1024,
+            ]
+            sX: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, NTOK * X_ROW_BYTES], 1024]
+            sB: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.b_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFB: cute.struct.Align[cute.struct.MemRange[cutlass.Float8E4M3FN, SFB_BYTES], 1024]
+
+        self.shared_storage = SharedStorage
+
+        self.kernel(
+            self.tiled_mma,
+            tma_atom_a1,
+            tma_a1,
+            tma_atom_sfa1,
+            tma_sfa1,
+            tma_atom_a2,
+            tma_a2,
+            tma_atom_sfa2,
+            tma_sfa2,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.sfb_smem_layout_std,
+            x_addr,
+            logits_addr,
+            bias_addr,
+            alpha1_addr,
+            scale_c1_addr,
+            alpha2_addr,
+            xgs_addr,
+            out_addr,
+            raw_scale_addr,
+            raw_idx_addr,
+            ws_addr,
+            part_addr,
+            himg_addr,
+            hflag_addr,
+            dbg_addr,
+            w13_addr,
+            w13_sf_addr,
+            w2_addr,
+            w2_sf_addr,
+            num_tokens,
+        ).launch(
+            grid=[self.num_ctas, 1, 1],
+            block=[NUM_THREADS, 1, 1],
+            cluster=[1, 1, 1],
+            smem=self.shared_storage.size_in_bytes(),
+            min_blocks_per_mp=1,
+            stream=stream,
+            use_pdl=True,
+        )
+
+    # ------------------------------------------------------------------
+    # Device code
+    # ------------------------------------------------------------------
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a1: cute.CopyAtom,
+        mA1: cute.Tensor,
+        tma_atom_sfa1: cute.CopyAtom,
+        mSFA1: cute.Tensor,
+        tma_atom_a2: cute.CopyAtom,
+        mA2: cute.Tensor,
+        tma_atom_sfa2: cute.CopyAtom,
+        mSFA2: cute.Tensor,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfb_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_std: cute.Layout,
+        x_addr: Int64,
+        logits_addr: Int64,
+        bias_addr: Int64,
+        alpha1_addr: Int64,
+        scale_c1_addr: Int64,
+        alpha2_addr: Int64,
+        xgs_addr: Int64,
+        out_addr: Int64,
+        raw_scale_addr: Int64,
+        raw_idx_addr: Int64,
+        ws_addr: Int64,
+        part_addr: Int64,
+        himg_addr: Int64,
+        hflag_addr: Int64,
+        dbg_addr: Int64,
+        w13_addr: Int64,
+        w13_sf_addr: Int64,
+        w2_addr: Int64,
+        w2_sf_addr: Int64,
+        num_tokens: Int32,
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        tidx, _, _ = cute.arch.thread_idx()
+        lane = tidx % 32
+        cta, _, _ = cute.arch.block_idx()
+        grid, _, _ = cute.arch.grid_dim()
+        cta = Int32(cta)
+        grid = Int32(grid)
+
+        if warp_idx == TMA_WARP:
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a1)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa1)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a2)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa2)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        ring = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.ring_mbar.data_ptr(),
+            num_stages=self.num_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+            tx_count=self.num_tma_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+        acc_pipe = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_mbar.data_ptr(),
+            num_stages=self.acc_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, EPI_WARPS),
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+        b_mbar = storage.b_mbar.data_ptr()
+        if warp_idx == MMA_WARP:
+            with cute.arch.elect_one():
+                # one arrival per quantized k-block (all rows)
+                for kb in cutlass.range_constexpr(FC1_KB):
+                    cute.arch.mbarrier_init(b_mbar + kb, 1)
+                for ls in cutlass.range_constexpr(MAX_LOCAL_SLOTS):
+                    cute.arch.mbarrier_init(b_mbar + B2_BAR + ls, 1)
+                cute.arch.mbarrier_init(b_mbar + X_BAR, 1)
+                cute.arch.mbarrier_init(b_mbar + FC1_DONE_BAR, 1)
+                for ls in cutlass.range_constexpr(MAX_LOCAL_SLOTS):
+                    cute.arch.mbarrier_init(b_mbar + STG_BAR + ls, 1)
+
+        cute.arch.mbarrier_init_fence()
+
+        if warp_idx == 0:
+            cute.arch.alloc_tmem(TMEM_COLS, storage.tmem_holding)
+
+        # Shared routing tables, zeroed before routing.
+        s_misc = storage.misc.get_tensor(cute.make_layout(16))
+        s_bitmap = storage.bitmap.get_tensor(cute.make_layout(8))
+        s_topk_id = storage.topk_id.get_tensor(cute.make_layout(NTOK * TOP_K))
+        s_topk_w = storage.topk_w.get_tensor(cute.make_layout(NTOK * TOP_K))
+        s_slot_expert = storage.slot_expert.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_slot_mask = storage.slot_mask.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_slot_w = storage.slot_w.get_tensor(cute.make_layout(MAX_SLOTS * NTOK))
+        s_slot_k = storage.slot_k.get_tensor(cute.make_layout(MAX_SLOTS * NTOK))
+        s_alpha1 = storage.slot_alpha1.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_scale_c1 = storage.slot_scale_c1.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_alpha2 = storage.slot_alpha2.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_expert_slot = storage.expert_slot.get_tensor(cute.make_layout(NUM_EXPERTS))
+
+        s_expert_slot[tidx] = Int32(-1)
+        s_slot_w[tidx] = Float32(0.0)
+        s_slot_w[tidx + NUM_THREADS] = Float32(0.0)
+        if tidx < MAX_SLOTS:
+            s_slot_mask[tidx] = Uint32(0)
+        if tidx < 16:
+            s_misc[tidx] = Int32(0)
+        if cutlass.const_expr(self.debug & DEBUG_TIMELINE):
+            if tidx < TL_SLOTS:
+                ptx.sts_u64(Int32(storage.tl_buf.data_ptr().toint()) + tidx * 8, Int64(0))
+        if tidx < 8:
+            s_bitmap[tidx] = Uint32(0)
+        # The router bias is a parameter, not an upstream output: read it
+        # before waiting on the previous grid.
+        bias_r = cute.make_rmem_tensor((8,), Float32)
+        for j in cutlass.range_constexpr(8):
+            bias_r[j] = ptx.bf16x2_to_f32(ptx.ld_nc_u16(bias_addr + (lane + 32 * j) * 2), False)
+        cute.arch.sync_threads()
+
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_ENTRY)
+            if cutlass.const_expr(tl):
+                ptx.sts_u64(tl_s + TL_GT0 * 8, ptx.globaltimer_u64())
+        # Everything below reads tensors produced by earlier kernels or
+        # touches the workspace shared with the previous launch.
+        cute.arch.griddepcontrol_wait()
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_PDL)
+        # Every CTA counts its own launches, so all CTAs of a launch read the
+        # same value without synchronizing; this launch stamps LL words with
+        # the incremented value (stale words carry older epochs).
+        epoch_next = ptx.ld_relaxed_gpu_u32(ws_addr + cta * 4) + Uint32(1)
+        if cutlass.const_expr(self.delay & DELAY_ENTRY):
+            self._delay_spin(DELAY_ENTRY * 53, cta, epoch_next)
+        # Stage the activation rows in shared memory with bulk copies before
+        # any weight streaming starts, so they never queue behind it.
+        if warp_idx == TMA_WARP:
+            with cute.arch.elect_one():
+                xbar = Int32((b_mbar + X_BAR).toint())
+                sx = Int32(storage.sX.data_ptr().toint())
+                ptx.mbar_arrive_expect_tx(xbar, num_tokens * X_ROW_BYTES)
+                r = Int32(0)
+                while r < num_tokens:
+                    ptx.bulk_g2s(
+                        sx + r * X_ROW_BYTES, x_addr + r * X_ROW_BYTES, Int32(X_ROW_BYTES), xbar
+                    )
+                    r = r + 1
+
+        # -------- routing (warp t handles token t) -> slot tables --------
+        self._route_tokens(
+            warp_idx, lane, num_tokens, logits_addr, bias_r, s_topk_id, s_topk_w, storage
+        )
+        cute.arch.sync_threads()
+        # Every thread has read the epoch word (barrier above): advance it.
+        if tidx == 0:
+            ptx.st_relaxed_gpu_u32(ws_addr + cta * 4, epoch_next)
+        e = tidx
+        bw = s_bitmap[e // 32]
+        prefix = Int32(0)
+        for w in cutlass.range_constexpr(8):
+            word = s_bitmap[w]
+            if w < e // 32:
+                prefix = prefix + ptx.popc_u32(word)
+        below = bw & ((Uint32(1) << Uint32(e % 32)) - Uint32(1))
+        prefix = prefix + ptx.popc_u32(below)
+        if ((bw >> Uint32(e % 32)) & Uint32(1)) != Uint32(0):
+            s_slot_expert[prefix] = e
+            s_expert_slot[e] = prefix
+        if tidx == 0:
+            total = Int32(0)
+            for w in cutlass.range_constexpr(8):
+                total = total + ptx.popc_u32(s_bitmap[w])
+            s_misc[MISC_NUM_SLOTS] = total
+        cute.arch.sync_threads()
+        if tidx < NTOK * TOP_K:
+            t = tidx // TOP_K
+            k = tidx % TOP_K
+            if t < num_tokens:
+                e = s_topk_id[tidx]
+                sl = s_expert_slot[e]
+                s_slot_w[sl * NTOK + t] = s_topk_w[tidx]
+                s_slot_k[sl * NTOK + t] = k
+                ptx.atom_or_shared_u32(
+                    Int32(storage.slot_mask.data_ptr().toint()) + sl * 4, Uint32(1) << Uint32(t)
+                )
+        cute.arch.sync_threads()
+        num_slots = s_misc[MISC_NUM_SLOTS]
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_ROUTED)
+
+        # ---------------- static work partition ----------------
+        # FC1 runs on p1 <= grid CTAs (at least min_fc1_kb k-blocks each);
+        # FC2 runs on the whole grid.
+        t1 = num_slots * (FC1_MT * FC1_KB)
+        p1 = cutlass.min(grid, t1 // self.min_fc1_kb)
+        g0 = t1
+        g1 = t1
+        if cta < p1:
+            g0 = (cta * t1) // p1
+            g1 = ((cta + 1) * t1) // p1
+        t2 = num_slots * FC2_MT
+        u0 = (cta * t2) // grid
+        u1 = ((cta + 1) * t2) // grid
+
+        if warp_idx == TMA_WARP:
+            self._tma_producer(
+                tma_atom_a1,
+                mA1,
+                tma_atom_sfa1,
+                mSFA1,
+                tma_atom_a2,
+                mA2,
+                tma_atom_sfa2,
+                mSFA2,
+                a_smem_layout_staged,
+                sfa_smem_layout_staged,
+                tiled_mma,
+                cluster_layout_vmnk,
+                storage,
+                ring,
+                s_slot_expert,
+                g0,
+                g1,
+                u0,
+                u1,
+                num_slots,
+                w13_addr,
+                w13_sf_addr,
+                w2_addr,
+                w2_sf_addr,
+            )
+        if warp_idx == MMA_WARP:
+            self._mma_issuer(
+                tiled_mma,
+                storage,
+                ring,
+                acc_pipe,
+                b_mbar,
+                a_smem_layout_staged,
+                sfa_smem_layout_staged,
+                b_smem_layout_staged,
+                sfb_smem_layout_staged,
+                sfb_smem_layout_std,
+                g0,
+                g1,
+                u0,
+                u1,
+                num_slots,
+            )
+        if (warp_idx != TMA_WARP) & (warp_idx != MMA_WARP):
+            self._quantize_x(
+                warp_idx, lane, xgs_addr, num_tokens, g0, g1, storage, dbg_addr, b_mbar
+            )
+        if warp_idx >= GATHER_WARP0:
+            self._fetch_fc2_operands(
+                tidx - 32 * GATHER_WARP0,
+                u0,
+                u1,
+                num_slots,
+                epoch_next,
+                storage,
+                b_mbar,
+                himg_addr,
+                hflag_addr,
+            )
+
+        if cutlass.const_expr(not self.finalize):
+            if cta == 0:
+                if tidx < NTOK * TOP_K:
+                    t = tidx // TOP_K
+                    if t < num_tokens:
+                        ptx.st_u16(raw_scale_addr + tidx * 2, ptx.bf16_bits(s_topk_w[tidx]))
+                        ptx.st_u32(raw_idx_addr + tidx * 4, Uint32(tidx))
+
+        if cutlass.const_expr(self.debug & DEBUG_ROUTING):
+            if cta == 0:
+                if tidx < NTOK * TOP_K:
+                    ptx.st_u32(dbg_addr + tidx * 4, Uint32(s_topk_id[tidx]))
+                    ptx.st_f32(dbg_addr + 256 + tidx * 4, s_topk_w[tidx])
+                if tidx < MAX_SLOTS:
+                    ptx.st_u32(dbg_addr + 512 + tidx * 4, Uint32(s_slot_expert[tidx]))
+                    ptx.st_u32(dbg_addr + 768 + tidx * 4, s_slot_mask[tidx])
+                if tidx == 0:
+                    ptx.st_u32(dbg_addr + 1024, Uint32(num_slots))
+
+        if warp_idx < EPI_WARPS:
+            if tidx < num_slots:
+                e = s_slot_expert[tidx]
+                s_alpha1[tidx] = ptx.ld_nc_f32(alpha1_addr + e * 4)
+                s_scale_c1[tidx] = ptx.ld_nc_f32(scale_c1_addr + e * 4)
+                s_alpha2[tidx] = ptx.ld_nc_f32(alpha2_addr + e * 4)
+            ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+            if tidx == 0:
+                _stamp(tl, tl_s, TL_EPI_START)
+            self._epilogue(
+                tiled_mma,
+                tidx,
+                warp_idx,
+                lane,
+                cta,
+                grid,
+                p1,
+                storage,
+                acc_pipe,
+                b_mbar,
+                s_slot_mask,
+                s_slot_w,
+                s_slot_k,
+                s_alpha1,
+                s_scale_c1,
+                s_alpha2,
+                num_slots,
+                num_tokens,
+                epoch_next,
+                g0,
+                g1,
+                t1,
+                u0,
+                u1,
+                t2,
+                out_addr,
+                part_addr,
+                himg_addr,
+                hflag_addr,
+                dbg_addr,
+            )
+            ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+            if cutlass.const_expr(self.debug & DEBUG_TIMELINE):
+                if tidx < TL_SLOTS:
+                    ptx.st_u64(
+                        dbg_addr + TL_BYTE_BASE + (cta * TL_SLOTS + tidx) * 8,
+                        ptx.lds_u64(tl_s + tidx * 8),
+                    )
+            if warp_idx == 0:
+                tmem_ptr = cute.arch.retrieve_tmem_ptr(
+                    Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+                )
+                cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _route_tokens(self, w, lane, num_tokens, logits_addr, bias_r, s_topk_id, s_topk_w, storage):
+        """Native SigmoidBias + ScaledSumNormalize top-8 for token ``w``.
+
+        Lane ``l`` owns experts ``l + 32 j``. Selection key: twiddled FP32
+        ``sigmoid + bias`` bits above ``65535 - expert`` (ties -> lower id),
+        as the native 64-bit key.
+        """
+        if w < num_tokens:
+            lg = cute.make_rmem_tensor((8,), Float32)
+            for j in cutlass.range_constexpr(8):
+                lg[j] = ptx.ld_nc_f32(logits_addr + (w * NUM_EXPERTS + lane + 32 * j) * 4)
+            self._route_one(w, lane, lg, 0, bias_r, s_topk_id, s_topk_w, storage)
+
+    @cute.jit
+    def _route_one(self, t, lane, lg, h, bias_r, s_topk_id, s_topk_w, storage):
+        hi = cute.make_rmem_tensor((8,), Uint32)
+        sg = cute.make_rmem_tensor((8,), Float32)
+        for j in cutlass.range_constexpr(8):
+            s = Float32(0.5) * cute.math.tanh(Float32(0.5) * lg[8 * h + j]) + Float32(0.5)
+            score = s + bias_r[j]
+            sg[j] = score - bias_r[j]
+            bits = ptx.f32_bits(score)
+            tw = bits ^ Uint32(0x80000000)
+            if (bits >> Uint32(31)) != Uint32(0):
+                tw = bits ^ Uint32(0xFFFFFFFF)
+            hi[j] = tw
+        taken = Uint32(0)
+        sel_e = cute.make_rmem_tensor((8,), Int32)
+        sig = cute.make_rmem_tensor((8,), Float32)
+        for k in cutlass.range_constexpr(8):
+            bh = Uint32(0)
+            bl = Uint32(0)
+            bs = Float32(0.0)
+            for j in cutlass.range_constexpr(8):
+                lo_j = Uint32(65535 - (lane + 32 * j))
+                avail = ((taken >> Uint32(j)) & Uint32(1)) == Uint32(0)
+                better = (hi[j] > bh) | ((hi[j] == bh) & (lo_j > bl))
+                if avail & better:
+                    bh = hi[j]
+                    bl = lo_j
+                    bs = sg[j]
+            mh = ptx.redux_max_u32(bh)
+            lo_c = Uint32(0)
+            if bh == mh:
+                lo_c = bl
+            ml = ptx.redux_max_u32(lo_c)
+            win = Int32(65535) - Int32(ml)
+            if (win % 32) == lane:
+                taken = taken | (Uint32(1) << Uint32(win // 32))
+            sel_e[k] = win
+            sig[k] = cute.arch.shuffle_sync(bs, win % 32)
+        # Butterfly order of the native warp reduction over lanes 0..7.
+        total = ((sig[0] + sig[4]) + (sig[2] + sig[6])) + ((sig[1] + sig[5]) + (sig[3] + sig[7]))
+        if lane < TOP_K:
+            my_sig = sig[0]
+            my_e = sel_e[0]
+            for k in cutlass.range_constexpr(1, 8):
+                if lane == k:
+                    my_sig = sig[k]
+                    my_e = sel_e[k]
+            wgt = ptx.bf16_round((my_sig * Float32(ROUTED_SCALING)) / total)
+            s_topk_id[t * TOP_K + lane] = my_e
+            s_topk_w[t * TOP_K + lane] = wgt
+            ptx.atom_or_shared_u32(
+                Int32(storage.bitmap.data_ptr().toint()) + (my_e // 32) * 4,
+                Uint32(1) << Uint32(my_e % 32),
+            )
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _prefetch_fc1(self, lane, g_begin, g1, s_slot_expert, w13_addr, w13_sf_addr):
+        """L2 prefetch of FC1 units [g_begin, g1): per tile segment, each of
+        the 128 rows' contiguous k-range and the segment's scale bytes."""
+        g = g_begin
+        while g < g1:
+            tile = g // FC1_KB
+            tile_base = tile * FC1_KB
+            seg_end = cutlass.min(g1, tile_base + FC1_KB)
+            slot = tile // FC1_MT
+            mt = tile - slot * FC1_MT
+            expert = Int64(s_slot_expert[slot])
+            ka = g - tile_base
+            nk = seg_end - g
+            row0 = (
+                w13_addr
+                + expert * W13_EXPERT_BYTES
+                + Int64(mt * 128) * W13_ROW_BYTES
+                + Int64(ka * KB_ROW_BYTES)
+            )
+            for i in cutlass.range_constexpr(4):
+                ptx.prefetch_l2_bulk(row0 + Int64(lane + 32 * i) * W13_ROW_BYTES, nk * KB_ROW_BYTES)
+            if lane == 0:
+                ptx.prefetch_l2_bulk(
+                    w13_sf_addr
+                    + expert * W13_SF_EXPERT_BYTES
+                    + Int64(mt * W13_SF_TILE_BYTES + ka * W13_SF_KB_BYTES),
+                    nk * W13_SF_KB_BYTES,
+                )
+            g = seg_end
+
+    @cute.jit
+    def _prefetch_fc2(self, lane, u0, u1, num_slots, s_slot_expert, w2_addr, w2_sf_addr):
+        """L2 prefetch of FC2 units [u0, u1): 16 KB weights + 2 KB scales."""
+        u = u0 + lane
+        while u < u1:
+            mt2 = u // num_slots
+            slot = u - mt2 * num_slots
+            expert = Int64(s_slot_expert[slot])
+            ptx.prefetch_l2_bulk(
+                w2_addr + expert * W2_EXPERT_BYTES + Int64(mt2 * W2_TILE_BYTES),
+                Int32(W2_TILE_BYTES),
+            )
+            ptx.prefetch_l2_bulk(
+                w2_sf_addr + expert * W2_SF_EXPERT_BYTES + Int64(mt2 * W2_SF_TILE_BYTES),
+                Int32(W2_SF_TILE_BYTES),
+            )
+            u = u + 32
+
+    @cute.jit
+    def _tma_producer(
+        self,
+        tma_atom_a1,
+        mA1,
+        tma_atom_sfa1,
+        mSFA1,
+        tma_atom_a2,
+        mA2,
+        tma_atom_sfa2,
+        mSFA2,
+        a_smem_layout_staged,
+        sfa_smem_layout_staged,
+        tiled_mma,
+        cluster_layout_vmnk,
+        storage,
+        ring,
+        s_slot_expert,
+        g0,
+        g1,
+        u0,
+        u1,
+        num_slots,
+        w13_addr,
+        w13_sf_addr,
+        w2_addr,
+        w2_sf_addr,
+    ):
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        lane = cute.arch.lane_idx()
+        lane0 = lane == 0
+        sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
+        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+        thr_mma = tiled_mma.get_slice(0)
+        a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+
+        gA1 = cute.local_tile(mA1, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
+        gSFA1 = cute.local_tile(
+            mSFA1, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        gA2 = cute.local_tile(mA2, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
+        gSFA2 = cute.local_tile(
+            mSFA2, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        tAsA, tAgA1 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_a1,
+            0,
+            a_cta_layout,
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gA1), 0, 3),
+        )
+        tAsSFA, tAgSFA1 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_sfa1,
+            0,
+            a_cta_layout,
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gSFA1), 0, 3),
+        )
+        tAsSFA = cute.filter_zeros(tAsSFA)
+        tAgSFA1 = cute.filter_zeros(tAgSFA1)
+        _, tAgA2 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_a2,
+            0,
+            a_cta_layout,
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gA2), 0, 3),
+        )
+        _, tAgSFA2 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_sfa2,
+            0,
+            a_cta_layout,
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gSFA2), 0, 3),
+        )
+        tAgSFA2 = cute.filter_zeros(tAgSFA2)
+
+        state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_stages)
+        # Prime the ring first, then hand the rest of this CTA's weight
+        # bytes to L2 prefetches so DRAM sees long sequential streams.
+        n_prime = cutlass.min(Int32(self.num_stages), g1 - g0)
+        g = g0
+        while g < g1:
+            if g - g0 == n_prime:
+                self._issue_prefetches(
+                    lane,
+                    g0 + n_prime,
+                    g1,
+                    u0,
+                    u1,
+                    num_slots,
+                    s_slot_expert,
+                    w13_addr,
+                    w13_sf_addr,
+                    w2_addr,
+                    w2_sf_addr,
+                )
+                if lane0:
+                    _stamp(tl, tl_s, TL_PF_ISSUED)
+            tile = g // FC1_KB
+            kb = g - tile * FC1_KB
+            slot = tile // FC1_MT
+            mt = tile - slot * FC1_MT
+            expert = s_slot_expert[slot]
+            ring.producer_acquire(state)
+            bar = ring.producer_get_barrier(state)
+            cute.copy(
+                tma_atom_a1,
+                tAgA1[(None, mt, kb, expert)],
+                tAsA[(None, state.index)],
+                tma_bar_ptr=bar,
+                cache_policy=Int64(L2_EVICT_FIRST),
+            )
+            cute.copy(
+                tma_atom_sfa1,
+                tAgSFA1[(None, mt, kb, expert)],
+                tAsSFA[(None, state.index)],
+                tma_bar_ptr=bar,
+                cache_policy=Int64(L2_EVICT_FIRST),
+            )
+            state.advance()
+            g = g + 1
+        if g1 - g0 == n_prime:
+            # Fewer FC1 units than ring stages: nothing left of FC1.
+            self._issue_prefetches(
+                lane,
+                g1,
+                g1,
+                u0,
+                u1,
+                num_slots,
+                s_slot_expert,
+                w13_addr,
+                w13_sf_addr,
+                w2_addr,
+                w2_sf_addr,
+            )
+            if lane0:
+                _stamp(tl, tl_s, TL_PF_ISSUED)
+        if lane0:
+            _stamp(tl, tl_s, TL_TMA_FC1)
+        if cutlass.const_expr(self.prefetch & PF_FC2_LATE):
+            self._prefetch_fc2(lane, u0, u1, num_slots, s_slot_expert, w2_addr, w2_sf_addr)
+        u = u0
+        while u < u1:
+            mt2 = u // num_slots
+            slot = u - mt2 * num_slots
+            expert = s_slot_expert[slot]
+            ring.producer_acquire(state)
+            bar = ring.producer_get_barrier(state)
+            cute.copy(
+                tma_atom_a2,
+                tAgA2[(None, mt2, 0, expert)],
+                tAsA[(None, state.index)],
+                tma_bar_ptr=bar,
+                cache_policy=Int64(L2_EVICT_FIRST),
+            )
+            cute.copy(
+                tma_atom_sfa2,
+                tAgSFA2[(None, mt2, 0, expert)],
+                tAsSFA[(None, state.index)],
+                tma_bar_ptr=bar,
+                cache_policy=Int64(L2_EVICT_FIRST),
+            )
+            state.advance()
+            u = u + 1
+        if lane0:
+            _stamp(tl, tl_s, TL_TMA_FC2)
+        ring.producer_tail(state)
+
+    @cute.jit
+    def _issue_prefetches(
+        self,
+        lane,
+        g_begin,
+        g1,
+        u0,
+        u1,
+        num_slots,
+        s_slot_expert,
+        w13_addr,
+        w13_sf_addr,
+        w2_addr,
+        w2_sf_addr,
+    ):
+        if cutlass.const_expr(self.prefetch & PF_FC1):
+            self._prefetch_fc1(lane, g_begin, g1, s_slot_expert, w13_addr, w13_sf_addr)
+        if cutlass.const_expr(self.prefetch & PF_FC2_EARLY):
+            self._prefetch_fc2(lane, u0, u1, num_slots, s_slot_expert, w2_addr, w2_sf_addr)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _mma_issuer(
+        self,
+        tiled_mma,
+        storage,
+        ring,
+        acc_pipe,
+        b_mbar,
+        a_smem_layout_staged,
+        sfa_smem_layout_staged,
+        b_smem_layout_staged,
+        sfb_smem_layout_staged,
+        sfb_smem_layout_std,
+        g0,
+        g1,
+        u0,
+        u1,
+        num_slots,
+    ):
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        lane0 = cute.arch.lane_idx() == 0
+        sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
+        sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
+        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+        sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
+        tCrA = tiled_mma.make_fragment_A(sA)
+        tCrB = tiled_mma.make_fragment_B(sB)
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.acc_stages))
+
+        acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
+            Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+        )
+        tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+        sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + TMEM_SFA_COL, dtype=cutlass.Float8E4M3FN)
+        tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            self.sf_vec_size,
+            cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+        )
+        tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+        sfb_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + TMEM_SFB_COL, dtype=cutlass.Float8E4M3FN)
+        tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            self.sf_vec_size,
+            cute.slice_(sfb_smem_layout_std, (None, None, None, 0)),
+        )
+        tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+
+        copy_s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = self._s2t_partition(sSFA, tCtSFA)
+        copy_s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = self._s2t_partition(sSFB, tCtSFB)
+
+        ring_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_stages
+        )
+        acc_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.acc_stages
+        )
+        num_kphases = cute.size(tCrA, mode=[2])
+
+        # FC1: wait for the quantized activations, then one accumulator per
+        # split-K segment.
+        if lane0:
+            _stamp(tl, tl_s, TL_MMA_FC1_START)
+        g = g0
+        while g < g1:
+            tile = g // FC1_KB
+            seg_end = cutlass.min(g1, (tile + 1) * FC1_KB)
+            acc_pipe.producer_acquire(acc_state)
+            tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)]
+            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+            while g < seg_end:
+                kb = g - (g // FC1_KB) * FC1_KB
+                cute.arch.mbarrier_wait(b_mbar + kb, 0)
+                ring.consumer_wait(ring_state)
+                cute.copy(
+                    copy_s2t_sfa, tCsSFA_s2t[(None, None, None, None, ring_state.index)], tCtSFA_s2t
+                )
+                cute.copy(copy_s2t_sfb, tCsSFB_s2t[(None, None, None, None, kb)], tCtSFB_s2t)
+                for kp in cutlass.range_constexpr(num_kphases):
+                    tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator)
+                    tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator)
+                    cute.gemm(
+                        tiled_mma,
+                        tCtAcc,
+                        tCrA[(None, None, kp, ring_state.index)],
+                        tCrB[(None, None, kp, kb)],
+                        tCtAcc,
+                    )
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                ring.consumer_release(ring_state)
+                ring_state.advance()
+                g = g + 1
+            acc_pipe.producer_commit(acc_state)
+            acc_state.advance()
+
+        # The FC2 operands overwrite the FC1 B operand: tell the gathering
+        # warps once every FC1 MMA (and scale copy) has completed.
+        with cute.arch.elect_one():
+            tcgen05.commit(b_mbar + FC1_DONE_BAR)
+        # FC2: one accumulator per (output tile, expert) unit.
+        if lane0:
+            _stamp(tl, tl_s, TL_MMA_FC1_DONE)
+        s_first = u0 - (u0 // num_slots) * num_slots
+        u = u0
+        while u < u1:
+            mt2 = u // num_slots
+            slot = u - mt2 * num_slots
+            ls = slot - s_first
+            if ls < 0:
+                ls = ls + num_slots
+            # operands of local slot ls are in shared memory
+            cute.arch.mbarrier_wait(b_mbar + B2_BAR + ls, 0)
+            acc_pipe.producer_acquire(acc_state)
+            tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)]
+            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+            ring_state = self._fc2_mma_ring(
+                tiled_mma,
+                ring,
+                ring_state,
+                tCtAcc,
+                copy_s2t_sfa,
+                tCsSFA_s2t,
+                tCtSFA_s2t,
+                copy_s2t_sfb,
+                tCsSFB_s2t,
+                tCtSFB_s2t,
+                tCtSFA,
+                tCtSFB,
+                tCrA,
+                tCrB,
+                ls,
+                num_kphases,
+            )
+            acc_pipe.producer_commit(acc_state)
+            acc_state.advance()
+            u = u + 1
+        if lane0:
+            _stamp(tl, tl_s, TL_MMA_FC2_DONE)
+        acc_pipe.producer_tail(acc_state)
+
+    @cute.jit
+    def _fc2_mma_ring(
+        self,
+        tiled_mma,
+        ring,
+        ring_state,
+        tCtAcc,
+        copy_s2t_sfa,
+        tCsSFA_s2t,
+        tCtSFA_s2t,
+        copy_s2t_sfb,
+        tCsSFB_s2t,
+        tCtSFB_s2t,
+        tCtSFA,
+        tCtSFB,
+        tCrA,
+        tCrB,
+        ls,
+        num_kphases,
+    ):
+        """One FC2 unit from the weight ring (A) and local slot ``ls`` (B)."""
+        ring.consumer_wait(ring_state)
+        cute.copy(copy_s2t_sfa, tCsSFA_s2t[(None, None, None, None, ring_state.index)], tCtSFA_s2t)
+        cute.copy(copy_s2t_sfb, tCsSFB_s2t[(None, None, None, None, ls)], tCtSFB_s2t)
+        for kp in cutlass.range_constexpr(num_kphases):
+            tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator)
+            tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator)
+            cute.gemm(
+                tiled_mma,
+                tCtAcc,
+                tCrA[(None, None, kp, ring_state.index)],
+                tCrB[(None, None, kp, ls)],
+                tCtAcc,
+            )
+            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+        ring.consumer_release(ring_state)
+        ring_state.advance()
+        return ring_state
+
+    def _s2t_partition(self, sSF, tSF):
+        tCsSF_compact = cute.filter_zeros(sSF)
+        tCtSF_compact = cute.filter_zeros(tSF)
+        copy_atom_s2t = cute.make_copy_atom(
+            tcgen05.Cp4x32x128bOp(self.cta_group), cutlass.Float8E4M3FN
+        )
+        tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
+        thr_copy_s2t = tiled_copy_s2t.get_slice(0)
+        tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
+        tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_)
+        tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
+        return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _quantize_x(self, warp_idx, lane, xgs_addr, num_tokens, g0, g1, storage, dbg_addr, b_mbar):
+        """NVFP4-quantize the activation k-blocks this CTA's FC1 range uses.
+
+        One warp task is one k-block for all rows: 256 BF16 elements per row,
+        8 per lane, two lanes per 16-element scale block, read from the
+        staged rows in shared memory; the rows are independent, so their
+        dependency chains interleave. Same arithmetic as the stock
+        ``fp4_quantize`` (E4M3 scale ``gs * amax * rcp(6)``, RN e2m1). Warps
+        take k-blocks in the MMA's order and each finished k-block arrives on
+        its mbarrier, so the first FC1 MMAs start while later k-blocks are
+        still being quantized. Padded rows (>= num_tokens) are quantized from
+        zeros: zero data and zero scales.
+        """
+        sb_base = Int32(storage.sB.data_ptr().toint())
+        sfb_base = Int32(storage.sSFB.data_ptr().toint())
+        gs = ptx.ld_nc_f32(xgs_addr)
+        rgs = ptx.rcp_approx_ftz(gs)
+        rcp6 = ptx.rcp_approx_ftz(Float32(6.0))
+        nkb = cutlass.min(g1 - g0, FC1_KB)
+        kb_first = g0 - (g0 // FC1_KB) * FC1_KB
+        if cutlass.const_expr(self.debug & DEBUG_QUANT):
+            nkb = Int32(FC1_KB)
+            kb_first = Int32(0)
+        qw = warp_idx
+        if warp_idx > TMA_WARP:
+            qw = warp_idx - 2
+        sx = Int32(storage.sX.data_ptr().toint())
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        stamp_q = (warp_idx == 0) & (lane == 0)
+        blk = lane // 2
+        sf_off = (blk // 4) * SFB_CHUNK_STRIDE + blk % 4
+        cute.arch.mbarrier_wait(b_mbar + X_BAR, 0)
+        if stamp_q:
+            _stamp(tl, tl_s, TL_Q_START)
+        j = Int32(0)
+        while qw + QUANT_WARPS * j < nkb:
+            kb = kb_first + qw + QUANT_WARPS * j
+            if kb >= FC1_KB:
+                kb = kb - FC1_KB
+            self._quant_rows(0, kb, lane, num_tokens, sx, sb_base, sfb_base, sf_off, gs, rgs, rcp6)
+            if num_tokens > NTOK // 2:
+                self._quant_rows(
+                    1, kb, lane, num_tokens, sx, sb_base, sfb_base, sf_off, gs, rgs, rcp6
+                )
+            else:
+                # rows 4..7 are padding: zero their scales
+                if lane < NTOK // 2 * 4:
+                    n = NTOK // 2 + lane // 4
+                    ptx.sts_u32(
+                        sfb_base + kb * SFB_STAGE_STRIDE + (lane % 4) * SFB_CHUNK_STRIDE + n * 16,
+                        Uint32(0),
+                    )
+            if stamp_q:
+                if j == 0:
+                    _stamp(tl, tl_s, TL_Q_G0_COMPUTE)
+            ptx.fence_proxy_async_shared()
+            cute.arch.sync_warp()
+            if lane == 0:
+                cute.arch.mbarrier_arrive(b_mbar + kb)
+            if stamp_q:
+                if j < 3:
+                    _stamp(tl, tl_s, TL_Q_G0_FENCE + j)
+            j = j + 1
+        if stamp_q:
+            _stamp(tl, tl_s, TL_Q_DONE)
+        if cutlass.const_expr(self.debug & DEBUG_QUANT):
+            # debug only: wait for every k-block, then CTA 0 dumps B / SFB
+            qt = warp_idx * 32 + lane
+            if warp_idx > TMA_WARP:
+                qt = qt - 64
+            ptx.bar_sync(BAR_QUANT, 32 * QUANT_WARPS)
+            cta_q, _, _ = cute.arch.block_idx()
+            if cta_q == 0:
+                w = qt
+                while w < (FC1_KB * 1024 + SFB_BYTES) // 4:
+                    if w < FC1_KB * 256:
+                        ptx.st_u32(dbg_addr + 4096 + w * 4, ptx.lds_u32(sb_base + w * 4))
+                    else:
+                        r2 = w - FC1_KB * 256
+                        ptx.st_u32(dbg_addr + 4096 + w * 4, ptx.lds_u32(sfb_base + r2 * 4))
+                    w = w + 32 * QUANT_WARPS
+
+    @cute.jit
+    def _quant_rows(self, half, kb, lane, num_tokens, sx, sb_base, sfb_base, sf_off, gs, rgs, rcp6):
+        """Quantize rows ``4*half .. 4*half+3`` of k-block ``kb``; rows at or
+        beyond ``num_tokens`` get zero data and zero scales."""
+        rows = NTOK // 2
+        # All four rows' loads first (independent chains); rows past
+        # num_tokens read stale shared memory and are masked to zero below.
+        d = cute.make_rmem_tensor((rows * 4,), Uint32)
+        for r in cutlass.range_constexpr(rows):
+            n = half * rows + r
+            x4 = ptx.lds_v4_u32(sx + n * X_ROW_BYTES + kb * 512 + lane * 16)
+            for c in cutlass.range_constexpr(4):
+                d[4 * r + c] = x4[c]
+        for r in cutlass.range_constexpr(rows):
+            n = half * rows + r
+            vals = cute.make_rmem_tensor((8,), Float32)
+            for i in cutlass.range_constexpr(4):
+                vals[2 * i] = ptx.bf16x2_to_f32(d[4 * r + i], False)
+                vals[2 * i + 1] = ptx.bf16x2_to_f32(d[4 * r + i], True)
+            amax = Float32(0.0)
+            for i in cutlass.range_constexpr(8):
+                amax = cute.arch.fmax(amax, cute.math.absf(vals[i]))
+            amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, 1))
+            sf = ptx.cvt_e4m3_rn_sat(gs * (amax * rcp6))
+            sf_f = ptx.e4m3_to_f32(sf)
+            oscale = Float32(0.0)
+            if amax != Float32(0.0):
+                oscale = ptx.rcp_approx_ftz(sf_f * rgs)
+            word = ptx.cvt_e2m1x8(
+                vals[0] * oscale,
+                vals[1] * oscale,
+                vals[2] * oscale,
+                vals[3] * oscale,
+                vals[4] * oscale,
+                vals[5] * oscale,
+                vals[6] * oscale,
+                vals[7] * oscale,
+            )
+            if n >= num_tokens:
+                word = Uint32(0)
+                sf = Uint32(0)
+            ptx.sts_u32(
+                sb_base + kb * 1024 + n * 128 + (((lane // 4) ^ n) * 16) + (lane % 4) * 4, word
+            )
+            if lane % 2 == 0:
+                ptx.sts_u8(sfb_base + kb * SFB_STAGE_STRIDE + n * 16 + sf_off, sf)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _fetch_fc2_operands(
+        self, gt, u0, u1, num_slots, epoch_next, storage, b_mbar, himg_addr, hflag_addr
+    ):
+        """Warps 6-7: fetch this CTA's FC2 operands into shared memory.
+
+        Thread ``ls`` (< local slots) waits for the four FC1 tile flags of
+        expert slot ``s_first + ls``, then bulk-copies the slot's operand
+        image into B stage ``ls`` / compact SFB stage ``ls``, completing on
+        that slot's mbarrier, so the FC2 MMAs of early slots start while
+        later ones are in flight. Shared memory is written only after every
+        FC1 MMA of this CTA has completed.
+        """
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        nls = cutlass.min(num_slots, u1 - u0)
+        s_first = u0 - (u0 // num_slots) * num_slots
+        cta, _, _ = cute.arch.block_idx()
+        if cutlass.const_expr(self.delay & DELAY_CONSUMER):
+            self._delay_spin(DELAY_CONSUMER * 53, Int32(cta), epoch_next)
+        if cutlass.const_expr(self.handoff == 1):
+            self._fetch_fc2_operands_ll(
+                gt, nls, s_first, num_slots, epoch_next, storage, b_mbar, himg_addr, tl, tl_s
+            )
+        if cutlass.const_expr(self.handoff == 0):
+            if gt < nls:
+                slot = s_first + gt
+                if slot >= num_slots:
+                    slot = slot - num_slots
+                fa = hflag_addr + slot * (FC1_MT * 4)
+                f0 = ptx.ld_acquire_gpu_u32(fa)
+                f1 = ptx.ld_acquire_gpu_u32(fa + 4)
+                f2 = ptx.ld_acquire_gpu_u32(fa + 8)
+                f3 = ptx.ld_acquire_gpu_u32(fa + 12)
+                while (
+                    (f0 != epoch_next)
+                    | (f1 != epoch_next)
+                    | (f2 != epoch_next)
+                    | (f3 != epoch_next)
+                ):
+                    f0 = ptx.ld_acquire_gpu_u32(fa)
+                    f1 = ptx.ld_acquire_gpu_u32(fa + 4)
+                    f2 = ptx.ld_acquire_gpu_u32(fa + 8)
+                    f3 = ptx.ld_acquire_gpu_u32(fa + 12)
+                if gt == 0:
+                    _stamp(tl, tl_s, TL_FC2_SENTINELS)
+                if gt == nls - 1:
+                    _stamp(tl, tl_s, TL_FC2_LAST_SLOT)
+                # the images were written by generic stores of other CTAs; the
+                # bulk copy reads them through the async proxy
+                ptx.fence_proxy_async_global()
+                cute.arch.mbarrier_wait(b_mbar + FC1_DONE_BAR, 0)
+                if gt == 0:
+                    _stamp(tl, tl_s, TL_FC1_MMA_COMPLETE)
+                bar = Int32((b_mbar + B2_BAR).toint()) + gt * 8
+                ptx.mbar_arrive_expect_tx(bar, Int32(HIMG_BYTES))
+                src = himg_addr + slot * HBUF_BYTES
+                ptx.bulk_g2s(
+                    Int32(storage.sB.data_ptr().toint()) + gt * HIMG_B_BYTES,
+                    src,
+                    Int32(HIMG_B_BYTES),
+                    bar,
+                )
+                ptx.bulk_g2s(
+                    Int32(storage.sSFB.data_ptr().toint()) + gt * SFB_STAGE_STRIDE,
+                    src + HIMG_B_BYTES,
+                    Int32(HIMG_SF_BYTES),
+                    bar,
+                )
+                if gt == 0:
+                    _stamp(tl, tl_s, TL_B2_READY)
+
+    @cute.jit
+    def _fetch_fc2_operands_ll(
+        self, gt, nls, s_first, num_slots, epoch_next, storage, b_mbar, hll_addr, tl, tl_s
+    ):
+        """Hand-off mode 1: LL operands staged through the activation buffer.
+
+        Thread ``ls`` waits for one LL word of each of its expert's four FC1
+        tiles (relaxed polls, no fence), then bulk-copies the expert's LL
+        words into sX; all 64 threads then validate every word (re-polling
+        global memory for the rare stale one) and scatter the payload into
+        B / compact SFB stage ``ls``, expert by expert, releasing each to the
+        MMA warp with its own mbarrier.
+        """
+        sx = Int32(storage.sX.data_ptr().toint())
+        if gt < nls:
+            slot = s_first + gt
+            if slot >= num_slots:
+                slot = slot - num_slots
+            base = hll_addr + slot * HBUF_BYTES
+            # epoch half of the row-7 scale word of each tile
+            sa = base + (H_LL_DATA_WORDS + (NTOK - 1) * FC1_MT) * 8 + 4
+            f0 = ptx.ld_relaxed_gpu_u32(sa)
+            f1 = ptx.ld_relaxed_gpu_u32(sa + 8)
+            f2 = ptx.ld_relaxed_gpu_u32(sa + 16)
+            f3 = ptx.ld_relaxed_gpu_u32(sa + 24)
+            while (f0 != epoch_next) | (f1 != epoch_next) | (f2 != epoch_next) | (f3 != epoch_next):
+                f0 = ptx.ld_relaxed_gpu_u32(sa)
+                f1 = ptx.ld_relaxed_gpu_u32(sa + 8)
+                f2 = ptx.ld_relaxed_gpu_u32(sa + 16)
+                f3 = ptx.ld_relaxed_gpu_u32(sa + 24)
+            if gt == 0:
+                _stamp(tl, tl_s, TL_FC2_SENTINELS)
+            if gt == nls - 1:
+                _stamp(tl, tl_s, TL_FC2_LAST_SLOT)
+            # sX is free: quantization is complete once all FC1 MMAs are
+            cute.arch.mbarrier_wait(b_mbar + FC1_DONE_BAR, 0)
+            bar = Int32((b_mbar + STG_BAR).toint()) + gt * 8
+            ptx.mbar_arrive_expect_tx(bar, Int32(H_LL_BYTES))
+            ptx.bulk_g2s(sx + gt * H_LL_BYTES, base, Int32(H_LL_BYTES), bar)
+        sb_base = Int32(storage.sB.data_ptr().toint())
+        sfb_base = Int32(storage.sSFB.data_ptr().toint())
+        ls = Int32(0)
+        while ls < nls:
+            cute.arch.mbarrier_wait(b_mbar + STG_BAR + ls, 0)
+            if gt == 0:
+                if ls == 0:
+                    _stamp(tl, tl_s, TL_FC1_MMA_COMPLETE)
+            i = gt
+            while i < H_LL_ITEMS:
+                d = ptx.lds_v4_u32(sx + ls * H_LL_BYTES + i * 16)
+                if (d[1] != epoch_next) | (d[3] != epoch_next):
+                    slot = s_first + ls
+                    if slot >= num_slots:
+                        slot = slot - num_slots
+                    a = hll_addr + slot * HBUF_BYTES + i * 16
+                    d = ptx.ld_ll2(a)
+                    while (d[1] != epoch_next) | (d[3] != epoch_next):
+                        d = ptx.ld_ll2(a)
+                if i < H_LL_DATA_ITEMS:
+                    # words (n, w), (n, w + 1) of the FP4 row n
+                    n = i // 16
+                    w = 2 * (i - n * 16)
+                    ptx.sts_v2_u32(
+                        sb_base + ls * HIMG_B_BYTES + n * 128 + (((w // 4) ^ n) * 16) + (w % 4) * 4,
+                        d[0],
+                        d[2],
+                    )
+                else:
+                    # scale words (n, mt), (n, mt + 1)
+                    rr = i - H_LL_DATA_ITEMS
+                    n = rr // 2
+                    mt0 = 2 * (rr - n * 2)
+                    ptx.sts_u32(
+                        sfb_base + ls * SFB_STAGE_STRIDE + mt0 * SFB_CHUNK_STRIDE + n * 16, d[0]
+                    )
+                    ptx.sts_u32(
+                        sfb_base + ls * SFB_STAGE_STRIDE + (mt0 + 1) * SFB_CHUNK_STRIDE + n * 16,
+                        d[2],
+                    )
+                i = i + GATHER_THREADS
+            ptx.fence_proxy_async_shared()
+            ptx.bar_sync(BAR_GATHER, GATHER_THREADS)
+            if gt == 0:
+                cute.arch.mbarrier_arrive(b_mbar + B2_BAR + ls)
+                if ls == 0:
+                    _stamp(tl, tl_s, TL_B2_READY)
+            ls = ls + 1
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _delay_spin(self, salt, cta, epoch):
+        """Stress testing: spin on the global timer.
+
+        The delayed CTAs and the delay length are a hash of (CTA, launch
+        epoch, site salt), so every launch perturbs a different set of
+        producers and consumers. Call sites are compiled only when their
+        ``DELAY_*`` bit is set.
+        """
+        h = (
+            Uint32(cta) * Uint32(37) + (epoch & Uint32(0xFFFF)) * Uint32(101) + Uint32(salt)
+        ) % Uint32(17)
+        if h < Uint32(6):
+            d = Int64(h + Uint32(1)) * Int64(self.delay_ns)
+            t0 = ptx.globaltimer_u64()
+            t = t0
+            while t - t0 < d:
+                t = ptx.globaltimer_u64()
+
+    @cute.jit
+    def _cta_of(self, g, total, parts):
+        """CTA whose range [c*total//parts, (c+1)*total//parts) holds g."""
+        return ((g + 1) * parts - 1) // total
+
+    @cute.jit
+    def _ll_publish(self, v, tidx, cta, base, flag):
+        """Fire-and-forget LL store of this thread's 8 partial values."""
+        a = base + cta * LL_PART_BYTES + tidx * (NTOK * 8)
+        for j in cutlass.range_constexpr(NTOK // 2):
+            ptx.st_ll2(a + j * 16, ptx.f32_bits(v[2 * j]), ptx.f32_bits(v[2 * j + 1]), flag)
+
+    @cute.jit
+    def _ll_issue(self, buf, tidx, c_begin, c_end, base):
+        """Start loading the LL partials of CTAs c_begin.. (at most
+        GATHER_PEERS of them) into ``buf``; the owner issues this before its
+        own last segment, so the round trip overlaps that work."""
+        for p in cutlass.range_constexpr(GATHER_PEERS):
+            if c_begin + p <= c_end:
+                a = base + (c_begin + p) * LL_PART_BYTES + tidx * (NTOK * 8)
+                for j in cutlass.range_constexpr(NTOK // 2):
+                    d = ptx.ld_ll2(a + j * 16)
+                    for k in cutlass.range_constexpr(4):
+                        buf[p * 2 * NTOK + j * 4 + k] = d[k]
+
+    @cute.jit
+    def _ll_finish(self, v, buf, tidx, c_begin, c_end, base, flag):
+        """v += the LL partials of CTAs c_begin..c_end, summed in CTA order.
+
+        The first GATHER_PEERS peers' words are already in flight in
+        ``buf``; stale words are re-polled in rounds. Further peers (only
+        with very small splits) are fetched GATHER_PEERS at a time.
+        """
+        c0 = c_begin
+        while c0 <= c_end:
+            if c0 != c_begin:
+                self._ll_issue(buf, tidx, c0, c_end, base)
+            pending = Int32(1)
+            while pending != 0:
+                pending = Int32(0)
+                for p in cutlass.range_constexpr(GATHER_PEERS):
+                    if c0 + p <= c_end:
+                        a = base + (c0 + p) * LL_PART_BYTES + tidx * (NTOK * 8)
+                        for j in cutlass.range_constexpr(NTOK // 2):
+                            o = p * 2 * NTOK + j * 4
+                            if (buf[o + 1] != flag) | (buf[o + 3] != flag):
+                                d = ptx.ld_ll2(a + j * 16)
+                                for k in cutlass.range_constexpr(4):
+                                    buf[o + k] = d[k]
+                                pending = Int32(1)
+            for p in cutlass.range_constexpr(GATHER_PEERS):
+                if c0 + p <= c_end:
+                    for j in cutlass.range_constexpr(NTOK // 2):
+                        o = p * 2 * NTOK + j * 4
+                        v[2 * j] = v[2 * j] + ptx.bits_f32(buf[o])
+                        v[2 * j + 1] = v[2 * j + 1] + ptx.bits_f32(buf[o + 2])
+            c0 = c0 + GATHER_PEERS
+
+    @cute.jit
+    def _epilogue(
+        self,
+        tiled_mma,
+        tidx,
+        warp_idx,
+        lane,
+        cta,
+        grid,
+        p1,
+        storage,
+        acc_pipe,
+        b_mbar,
+        s_slot_mask,
+        s_slot_w,
+        s_slot_k,
+        s_alpha1,
+        s_scale_c1,
+        s_alpha2,
+        num_slots,
+        num_tokens,
+        epoch_next,
+        g0,
+        g1,
+        t1,
+        u0,
+        u1,
+        t2,
+        out_addr,
+        part_addr,
+        himg_addr,
+        hflag_addr,
+        dbg_addr,
+    ):
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        hst = Int32(storage.hstage.data_ptr().toint())
+        acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
+            Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+        )
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.acc_stages))
+        tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+        lane_off = (warp_idx * 32) << 16
+        acc_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.acc_stages
+        )
+        v = cute.make_rmem_tensor((NTOK,), Float32)
+        part1 = part_addr
+        part2 = part_addr + grid * LL_PART_BYTES
+
+        # ---------------- FC1 epilogue ----------------
+        # A tile split over CTAs is finished by the CTA holding its k-block
+        # 0 (its last segment); the others publish their partial (their
+        # first segment) as LL words.
+        pbuf = cute.make_rmem_tensor((GATHER_PEERS * 2 * NTOK,), Uint32)
+        g = g0
+        while g < g1:
+            tile = g // FC1_KB
+            tile_base = tile * FC1_KB
+            seg_end = cutlass.min(g1, tile_base + FC1_KB)
+            split_head = (g == tile_base) & (seg_end != tile_base + FC1_KB)
+            c_last = self._cta_of(tile_base + FC1_KB - 1, t1, p1)
+            if split_head:
+                if cutlass.const_expr(self.delay & DELAY_OWNER):
+                    self._delay_spin(DELAY_OWNER * 53, cta, epoch_next)
+                self._ll_issue(pbuf, tidx, cta + 1, c_last, part1)
+            acc_pipe.consumer_wait(acc_state)
+            if tidx == 0:
+                if g == g0:
+                    _stamp(tl, tl_s, TL_EPI_FIRST_ACC)
+            vals = ptx.tmem_ld_32x32b_x8(
+                Int32(tCtAcc_base[(None, None, None, acc_state.index)].iterator.toint()) + lane_off
+            )
+            for n in cutlass.range_constexpr(NTOK):
+                v[n] = vals[n]
+            cute.arch.fence_view_async_tmem_load()
+            if tidx == 0:
+                _stamp(tl, tl_s, TL_EPI_T_TMEM)
+            with cute.arch.elect_one():
+                acc_pipe.consumer_release(acc_state)
+            acc_state.advance()
+            if g == tile_base:
+                if split_head:
+                    self._ll_finish(v, pbuf, tidx, cta + 1, c_last, part1, epoch_next)
+                    if tidx == 0:
+                        _stamp(tl, tl_s, TL_EPI_FC1_GATHERED)
+                self._fc1_tile_epilogue(
+                    v,
+                    tile,
+                    warp_idx,
+                    lane,
+                    tidx,
+                    s_slot_mask,
+                    s_alpha1,
+                    s_scale_c1,
+                    num_tokens,
+                    epoch_next,
+                    himg_addr,
+                    hflag_addr,
+                    hst,
+                    tl_s,
+                    dbg_addr,
+                )
+            else:
+                if cutlass.const_expr(self.delay & DELAY_PRODUCER):
+                    self._delay_spin(DELAY_PRODUCER * 53, cta, epoch_next)
+                self._ll_publish(v, tidx, cta, part1, epoch_next)
+            g = seg_end
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_EPI_FC1_DONE)
+
+        # ---------------- FC2 epilogue ----------------
+        o = cute.make_rmem_tensor((NTOK,), Float32)
+        hid = warp_idx * 32 + (lane % 8) * 4 + lane // 8
+        u = u0
+        while u < u1:
+            mt2 = u // num_slots
+            seg_start = u
+            seg_end = cutlass.min(u1, (mt2 + 1) * num_slots)
+            tile_base = mt2 * num_slots
+            c_last = self._cta_of(tile_base + num_slots - 1, t2, grid)
+            split_head = (seg_start == tile_base) & (seg_end != tile_base + num_slots)
+            if cutlass.const_expr(self.finalize):
+                if split_head:
+                    if cutlass.const_expr(self.delay & DELAY_OWNER):
+                        self._delay_spin(DELAY_OWNER * 53, cta, epoch_next)
+                    self._ll_issue(pbuf, tidx, cta + 1, c_last, part2)
+            for n in cutlass.range_constexpr(NTOK):
+                o[n] = Float32(0.0)
+            while u < seg_end:
+                slot = u - mt2 * num_slots
+                mask = s_slot_mask[slot]
+                a2 = s_alpha2[slot]
+                acc_pipe.consumer_wait(acc_state)
+                if tidx == 0:
+                    if u == u0:
+                        _stamp(tl, tl_s, TL_FC2_FIRST_ACC)
+                vals = ptx.tmem_ld_32x32b_x8(
+                    Int32(tCtAcc_base[(None, None, None, acc_state.index)].iterator.toint())
+                    + lane_off
+                )
+                cute.arch.fence_view_async_tmem_load()
+                with cute.arch.elect_one():
+                    acc_pipe.consumer_release(acc_state)
+                acc_state.advance()
+                for n in cutlass.range_constexpr(NTOK):
+                    y = ptx.bf16_round(vals[n] * a2)
+                    if ((mask >> Uint32(n)) & Uint32(1)) != Uint32(0):
+                        if cutlass.const_expr(self.finalize):
+                            o[n] = o[n] + s_slot_w[slot * NTOK + n] * y
+                        else:
+                            k = s_slot_k[slot * NTOK + n]
+                            row = n * TOP_K + k
+                            ptx.st_u16(
+                                out_addr + (row * HIDDEN + mt2 * 128 + hid) * 2, ptx.bf16_bits(y)
+                            )
+                u = u + 1
+            if cutlass.const_expr(self.finalize):
+                if seg_start == tile_base:
+                    if split_head:
+                        self._ll_finish(o, pbuf, tidx, cta + 1, c_last, part2, epoch_next)
+                        if tidx == 0:
+                            _stamp(tl, tl_s, TL_FC2_GATHERED)
+                    for n in cutlass.range_constexpr(NTOK):
+                        if n < num_tokens:
+                            ptx.st_u16(
+                                out_addr + (n * HIDDEN + mt2 * 128 + hid) * 2, ptx.bf16_bits(o[n])
+                            )
+                else:
+                    if cutlass.const_expr(self.delay & DELAY_PRODUCER):
+                        self._delay_spin(DELAY_PRODUCER * 53, cta, epoch_next)
+                    self._ll_publish(o, tidx, cta, part2, epoch_next)
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_FC2_DONE)
+
+    @cute.jit
+    def _fc1_tile_epilogue(
+        self,
+        v,
+        tile,
+        warp_idx,
+        lane,
+        tidx,
+        s_slot_mask,
+        s_alpha1,
+        s_scale_c1,
+        num_tokens,
+        epoch_next,
+        himg_addr,
+        hflag_addr,
+        hst,
+        tl_s,
+        dbg_addr,
+    ):
+        """SwiGLU + NVFP4 quantization of one 128-row FC1 tile.
+
+        Row ``32w + q`` of the stock layout holds intermediate index
+        ``64*mt + 16*w + j`` with ``j = 2*(q % 8) + (q // 16)`` as the up row
+        when ``q % 16 < 8`` and its gate row at lane ``q + 8``, so one warp
+        owns exactly one 16-element scale block of the intermediate. The
+        tile's 256 FP4 bytes and 32 scale bytes are gathered in shared memory,
+        stored into the expert's operand image and published with the
+        tile's epoch flag.
+        """
+        slot = tile // FC1_MT
+        mt = tile - slot * FC1_MT
+        mask = s_slot_mask[slot]
+        a_gate = s_alpha1[slot]
+        s_up = s_scale_c1[slot]
+        is_up = (lane % 16) < 8
+        rcp6 = ptx.rcp_approx_ftz(Float32(6.0))
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_EPI_T_PARAMS)
+        if cutlass.const_expr(self.debug & DEBUG_FC1):
+            for n in cutlass.range_constexpr(NTOK):
+                ptx.st_f32(
+                    dbg_addr + 4 * (65536 + ((slot * FC1_MT + mt) * 128 + tidx) * NTOK + n), v[n]
+                )
+        for n in cutlass.range_constexpr(NTOK):
+            gate_raw = cute.arch.shuffle_sync_down(v[n], 8)
+            gr = gate_raw * a_gate
+            up = v[n] * s_up
+            sig = ptx.rcp_approx_ftz(
+                Float32(1.0) + ptx.ex2_approx(gr * Float32(-1.4426950408889634))
+            )
+            h = Float32(0.0)
+            if is_up:
+                h = up * (gr * sig)
+            am = cute.math.absf(h)
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 1))
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 2))
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 4))
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 16))
+            sf = ptx.cvt_e4m3_rn_sat(am * rcp6)
+            sf_f = ptx.e4m3_to_f32(sf)
+            osc = Float32(0.0)
+            if am != Float32(0.0):
+                osc = ptx.rcp_approx_ftz(sf_f)
+            partner = cute.arch.shuffle_sync_down(h, 16)
+            byte = ptx.cvt_e2m1x2(h * osc, partner * osc)
+            routed = (((mask >> Uint32(n)) & Uint32(1)) != Uint32(0)) & (n < num_tokens)
+            if not routed:
+                byte = Uint32(0)
+                sf = Uint32(0)
+            if lane < 8:
+                ptx.sts_u8(hst + n * 32 + warp_idx * 8 + lane, byte)
+            if lane == 0:
+                ptx.sts_u8(hst + HSTAGE_SF + n * 4 + warp_idx, sf)
+        if tidx == 0:
+            _stamp(tl, tl_s, TL_EPI_T_COMPUTED)
+        ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+        cta, _, _ = cute.arch.block_idx()
+        if cutlass.const_expr(self.delay & DELAY_PRODUCER):
+            self._delay_spin(DELAY_PRODUCER * 53, Int32(cta), epoch_next)
+        if cutlass.const_expr(self.handoff == 1):
+            # LL words: payload and epoch in one store, no fence
+            base = himg_addr + slot * HBUF_BYTES
+            if tidx < 32:
+                r = tidx // 4
+                p = tidx - r * 4
+                w = ptx.lds_v2_u32(hst + r * 32 + p * 8)
+                ptx.st_ll2(base + (r * 32 + mt * 8 + 2 * p) * 8, w[0], w[1], epoch_next)
+            if (tidx >= 32) & (tidx < 32 + NTOK):
+                r = tidx - 32
+                w1 = ptx.lds_u32(hst + HSTAGE_SF + r * 4)
+                ptx.st_ll1(base + (H_LL_DATA_WORDS + r * FC1_MT + mt) * 8, w1, epoch_next)
+            if tidx == 0:
+                _stamp(tl, tl_s, TL_EPI_T_STORED)
+            # the staging buffer is reused by the next tile
+            ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+        else:
+            img = himg_addr + slot * HBUF_BYTES
+            if tidx < 2 * NTOK:
+                # 16-byte chunk c16 = 2*mt + half of row r, swizzled as in sB
+                r = tidx // 2
+                half = tidx - r * 2
+                c16 = 2 * mt + half
+                d = ptx.lds_v4_u32(hst + r * 32 + half * 16)
+                ptx.st_v4_u32(img + r * 128 + ((c16 ^ r) * 16), d[0], d[1], d[2], d[3])
+                ptx.fence_proxy_async_global()
+            if (tidx >= 2 * NTOK) & (tidx < 3 * NTOK):
+                # the tile's four scale bytes of row r, compact SFB chunk mt
+                r = tidx - 2 * NTOK
+                w1 = ptx.lds_u32(hst + HSTAGE_SF + r * 4)
+                ptx.st_u32(img + HIMG_B_BYTES + mt * SFB_CHUNK_STRIDE + r * 16, w1)
+                ptx.fence_proxy_async_global()
+            if tidx == 0:
+                _stamp(tl, tl_s, TL_EPI_T_STORED)
+            # all image stores (and staging reads) precede the flag; the staging
+            # buffer is free for the next tile after this barrier
+            ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+            if tidx == 0:
+                ptx.fence_acq_rel_gpu()
+                ptx.st_relaxed_gpu_u32(hflag_addr + (slot * FC1_MT + mt) * 4, epoch_next)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel_xc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel_xc.py
new file mode 100644
index 000000000000..6d905fed2d14
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel_xc.py
@@ -0,0 +1,1775 @@
+# 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.
+"""Expert-cluster schedule of the routed_moe routed NVFP4 MoE (CuTe DSL, SM100 family).
+
+The grid is a set of 4-CTA thread block clusters. After every CTA has
+recomputed the routing, each cluster owns one expert per round (one round when
+the number of distinct experts D is at most the number of clusters, two
+otherwise): cluster ``c`` of round ``rho`` serves expert slot ``v % D`` with
+``v = rho * clusters + c``, so an expert is served by ``r`` clusters when there
+are more clusters than experts, and each of those clusters computes the whole
+FC1 of the expert (the repeated weight reads are L2 hits) and its share of the
+48 FC2 output tiles.
+
+Inside a cluster, CTA ``k`` stages and NVFP4-quantizes activation k-blocks
+``[6k, 6k + 6)`` of every row and bulk-copies each finished k-block into the
+three peers (distributed shared memory, completing on their k-block
+mbarriers), so every CTA holds the whole FC1 B operand after a quarter of the
+quantization work. CTA ``k`` then computes FC1 tile ``k`` (128 of the 512
+stock gate/up rows, all 24 k-blocks, so no split-K) through an 8-stage TMA
+ring, applies SwiGLU and the NVFP4 quantization of the intermediate, and
+stores its 64 intermediate columns of every token straight into the FC2
+operand stage of all four CTAs, then release-arrives on each CTA's operand
+mbarrier. Every CTA then runs its FC2 tiles from that operand. Nothing crosses
+a cluster: there are no global flags, split partials or epochs, and a cluster
+that starts late (for example behind a co-running kernel) delays only its own
+expert. The host sizes the grid so that every cluster is co-resident and
+SMs remain for the concurrently running shared expert (on sm_103 32 clusters,
+26 for the raw tuple consumed by the post-MoE hop when the router logits come
+from the pre-MoE hop; see ``routed_moe.RoutedMoe._xc_ctas``); with
+``min_clusters`` set, calls with
+few distinct experts use only ``clamp(D, min_clusters, clusters)`` clusters
+and the rest exit after routing.
+
+The kernel writes the native unfinalized tuple (FC2 rows ``t * 8 + k``,
+routing weights, identity permutation). The finalized output is produced by
+``RoutedMoeFinalizeKernel``, launched right after with programmatic dependent
+launch: the native finalize arithmetic (FP32 FMA over the top-k in order, one
+BF16 rounding).
+"""
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass.cute.nvgpu import tcgen05
+from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32
+
+from . import ptx
+from .kernel import (
+    A_STAGE_BYTES,
+    BAR_EPI,
+    BAR_QUANT,
+    DEBUG_FC1,
+    DEBUG_QUANT,
+    DEBUG_ROUTING,
+    DEBUG_TIMELINE,
+    DELAY_ENTRY,
+    DELAY_PRODUCER,
+    EPI_WARPS,
+    FC1_KB,
+    FC1_MT,
+    FC2_MT,
+    HIDDEN,
+    HSTAGE_BYTES,
+    HSTAGE_SF,
+    L2_EVICT_FIRST,
+    MAX_SLOTS,
+    MMA_WARP,
+    NTOK,
+    NUM_EXPERTS,
+    NUM_THREADS,
+    QUANT_WARPS,
+    SFA_STAGE_BYTES,
+    SFB_BYTES,
+    SFB_CHUNK_STRIDE,
+    SFB_STAGE_STRIDE,
+    TL_BYTE_BASE,
+    TL_SLOTS,
+    TMA_WARP,
+    TMEM_COLS,
+    TMEM_SFA_COL,
+    TMEM_SFB_COL,
+    TOP_K,
+    W2_EXPERT_BYTES,
+    W2_SF_EXPERT_BYTES,
+    W2_SF_TILE_BYTES,
+    W2_TILE_BYTES,
+    W13_EXPERT_BYTES,
+    W13_ROW_BYTES,
+    W13_SF_EXPERT_BYTES,
+    W13_SF_TILE_BYTES,
+    X_ROW_BYTES,
+    RoutedMoeKernel,
+    _stamp,
+)
+
+DELAY_FINALIZE = 16  # RoutedMoeFinalizeKernel CTAs, before their output stores
+# L2 prefetch (``l2pf`` bits; 0 = off): the TMA warp issues
+# ``cp.async.bulk.prefetch.L2`` for weight bytes its ring reaches later, so
+# that those ring loads hit L2 instead of waiting for DRAM.
+L2PF_FC1_R1 = 1  # round-1 FC1 tile (by the first cluster of the round-1 slot)
+L2PF_FC2 = 2  # this CTA's FC2 tiles of both rounds
+L2PF_EARLY = 4  # issue them before the first ring load (else after the first ring fill)
+FC1_TILE_BYTES = 128 * W13_ROW_BYTES
+PF_FC1_CHUNK = FC1_TILE_BYTES // 32  # one bulk prefetch per lane
+CLUSTER = 4  # CTAs per cluster = FC1 tiles per expert
+# Speculative prefetch of predicted experts (``spec_k`` > 0, package cross_layer_expert_prefetch):
+# 16 KB chunks of one expert's FC1 weight, FC1 scales, FC2 weight and FC2
+# scales, in that order, for ``SPEC_KMAX``-wide rows of ranked expert IDs.
+SPEC_CHUNK = 16384
+SPEC_C13 = W13_EXPERT_BYTES // SPEC_CHUNK
+SPEC_C13SF = W13_SF_EXPERT_BYTES // SPEC_CHUNK
+SPEC_C2 = W2_EXPERT_BYTES // SPEC_CHUNK
+SPEC_C2SF = W2_SF_EXPERT_BYTES // SPEC_CHUNK
+SPEC_CPE = SPEC_C13 + SPEC_C13SF + SPEC_C2 + SPEC_C2SF
+SPEC_KMAX = 32
+assert SPEC_CPE * SPEC_CHUNK == (
+    W13_EXPERT_BYTES + W13_SF_EXPERT_BYTES + W2_EXPERT_BYTES + W2_SF_EXPERT_BYTES
+)
+MAX_ROUNDS = 2  # expert rounds per cluster (D <= 2 * clusters)
+L2_EVICT_NORMAL = 0x1000000000000000
+HIMG_B_BYTES = NTOK * 128  # one FC2 B operand stage (8 rows x 128 B FP4)
+SFB2_BYTES = (MAX_ROUNDS - 1) * SFB_STAGE_STRIDE + 3 * SFB_CHUNK_STRIDE + 512
+
+# mbarriers: one per quantized k-block, the staged activation rows, one FC2
+# operand barrier per round (4 arrivals: one per cluster CTA)
+X_BAR = FC1_KB
+B2_BAR = X_BAR + 1
+NUM_BAR = B2_BAR + MAX_ROUNDS
+# Activation quantization is split over the cluster: CTA k stages and
+# quantizes k-blocks [6k, 6k + 6) of every row and bulk-copies each finished
+# k-block (1 KB B stage + 512 B compact SFB stage) into the three peers.
+KB_PER_CTA = FC1_KB // CLUSTER
+XS_ROW_BYTES = KB_PER_CTA * 512  # staged columns of one activation row
+KB_TX_BYTES = 1024 + SFB_STAGE_STRIDE  # one k-block's B + SFB stage
+
+# Timeline slots (debug builds only)
+(
+    XT_ENTRY,
+    XT_PDL,
+    XT_ROUTED,
+    XT_Q_DONE,
+    XT_TMA_FC1,
+    XT_MMA_FC1,
+    XT_EPI_FC1_ACC,
+    XT_EPI_FC1_PUB,
+    XT_B2_READY,
+    XT_TMA_FC2,
+    XT_FC2_FIRST_ACC,
+    XT_MMA_FC2,
+    XT_FC2_DONE,
+    XT_EXIT,
+    XT_MMA_FC1_R0,
+    XT_EPI_FC1_PUB_R1,
+    XT_B2_READY_R1,
+    XT_FC2_R0_DONE,
+    XT_TMA_FC1_R0,
+    XT_PF_ISSUED,
+    XT_EPI_FC1_ACC_R1,
+    XT_EPI_FC1_SWIGLU,
+    XT_EPI_FC1_SWIGLU_R1,
+    XT_EPI_FC1_TMEM,  # 23, 24: accumulator loaded from TMEM (rounds 0, 1)
+    XT_EPI_FC1_TMEM_R1,
+    XT_EPI_FC1_SWQ,  # 25, 26: SwiGLU + quantization done, before the barrier
+    XT_EPI_FC1_SWQ_R1,
+    XT_EPI_FC1_SWQ_W_R1,  # 27..30: the same, round 1, epilogue warp w (lane 0)
+) = range(28)
+# slots 28..30 are XT_EPI_FC1_SWQ_W_R1 + 1..3
+XT_GT0 = TL_SLOTS - 1
+
+
+class RoutedMoeExpertClusterKernel(RoutedMoeKernel):
+    """Expert-cluster kernel; raw output only (see module docstring)."""
+
+    def __init__(
+        self,
+        num_ctas: int,
+        num_stages: int = 8,
+        debug: int = 0,
+        delay: int = 0,
+        delay_ns: int = 2000,
+        finalize: bool = False,
+        trigger: str = "late",
+        min_clusters: int = 0,
+        red_out: bool = False,
+        l2pf: int = 0,
+        two_round_clusters: int = 0,
+        acc_stages: int = 4,
+        spec_k: int = 0,
+        spec_ctas: int = 0,
+        spec_self: bool = False,
+    ):
+        assert num_ctas % CLUSTER == 0
+        assert 0 <= spec_k <= SPEC_KMAX
+        # Speculative L2 prefetch (cross_layer_expert_prefetch): before its dependency wait, the TMA
+        # warp of each of the first ``spec_ctas`` CTAs (0 = all) issues
+        # prefetches of the first ``spec_k`` experts listed at
+        # ``spec_ids_addr`` (int32, ``SPEC_KMAX`` wide); 0 compiles it out.
+        # ``spec_self``: after routing, CTA 0 also writes this call's distinct
+        # experts there, ranked by row count, for the next call of the layer.
+        self.spec_k = spec_k
+        self.spec_ctas = spec_ctas
+        self.spec_self = spec_self
+        # 0: a two-round call (D > clusters) uses every cluster. Otherwise it
+        # uses this many (the rest exit after routing and leave their SMs to
+        # the shared expert), as long as two rounds of them cover D.
+        self.two_round_clusters = two_round_clusters
+        # L2 prefetch bits (L2PF_*); 0 compiles the prefetch out
+        self.l2pf = l2pf
+        # routed_moe_into_hop "fp32" entry: instead of the raw rows, every FC2 output adds
+        # its weighted contribution w[t, k] * bf16(y) (exact in FP32) to an
+        # FP32 [R, H] accumulator; the post-MoE hop rounds and clears it.
+        self.red_out = red_out
+        # programmatic-dependent-launch trigger: "late" after each CTA's last
+        # FC2 tile, "early" right after the dependency wait (the successor's
+        # CTAs may then wait on the free SMs), "none" (the successor launches
+        # when the grid completes)
+        assert trigger in ("late", "early", "none")
+        self.trigger = trigger
+        # 0: every cluster works. Otherwise, when D <= clusters, only
+        # clamp(D, min_clusters, clusters) clusters work and the others exit
+        # right after routing, leaving their SMs to concurrent kernels (the
+        # shared expert on the aux stream).
+        self.min_clusters = min_clusters
+        # Two rounds cover at most 2 x clusters distinct experts; the host
+        # launches this kernel only for calls with rows x 8 within that bound
+        # (``RoutedMoe.schedule``).
+        super().__init__(
+            num_ctas,
+            finalize=finalize,
+            num_stages=num_stages,
+            debug=debug,
+            delay=delay,
+            delay_ns=delay_ns,
+            acc_stages=acc_stages,
+        )
+
+    def _setup(self):
+        super()._setup()
+        self.b2_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            self.tiled_mma, self.mma_tiler, cutlass.Float4E2M1FN, MAX_ROUNDS
+        )
+        sfb2_std = blockscaled_utils.make_smem_layout_sfb(
+            self.tiled_mma, self.mma_tiler, self.sf_vec_size, MAX_ROUNDS
+        )
+        self.sfb2_smem_layout_staged = cute.make_layout(
+            sfb2_std.shape,
+            stride=((((16, 4), 0), (0, 1)), 0, SFB_CHUNK_STRIDE, SFB_STAGE_STRIDE),
+        )
+
+    @cute.jit
+    def run(
+        self,
+        w13_ptr: cute.Pointer,
+        w13_sf_ptr: cute.Pointer,
+        w2_ptr: cute.Pointer,
+        w2_sf_ptr: cute.Pointer,
+        x_addr: Int64,
+        logits_addr: Int64,
+        bias_addr: Int64,
+        alpha1_addr: Int64,
+        scale_c1_addr: Int64,
+        alpha2_addr: Int64,
+        xgs_addr: Int64,
+        rows_addr: Int64,
+        raw_scale_addr: Int64,
+        raw_idx_addr: Int64,
+        out_addr: Int64,
+        counter_addr: Int64,
+        dbg_addr: Int64,
+        w13_addr: Int64,
+        w13_sf_addr: Int64,
+        w2_addr: Int64,
+        w2_sf_addr: Int64,
+        spec_ids_addr: Int64,
+        num_tokens: Int32,
+        stream,
+    ):
+        self._setup()
+        w13 = cute.make_tensor(
+            w13_ptr, cute.make_ordered_layout((2 * 256, HIDDEN, NUM_EXPERTS), order=(1, 0, 2))
+        )
+        w2 = cute.make_tensor(
+            w2_ptr, cute.make_ordered_layout((HIDDEN, 256, NUM_EXPERTS), order=(1, 0, 2))
+        )
+        w13_sf = cute.make_tensor(
+            w13_sf_ptr, blockscaled_utils.tile_atom_to_shape_SF(w13.shape, self.sf_vec_size)
+        )
+        w2_sf = cute.make_tensor(
+            w2_sf_ptr, blockscaled_utils.tile_atom_to_shape_SF(w2.shape, self.sf_vec_size)
+        )
+        a_op = sm100_utils.cluster_shape_to_tma_atom_A((1, 1), self.tiled_mma.thr_id)
+        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+        sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
+        tma_atom_a1, tma_a1 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op, w13, a_smem_layout, self.mma_tiler, self.tiled_mma, self.cluster_layout_vmnk.shape
+        )
+        tma_atom_sfa1, tma_sfa1 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            w13_sf,
+            sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+        tma_atom_a2, tma_a2 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op, w2, a_smem_layout, self.mma_tiler, self.tiled_mma, self.cluster_layout_vmnk.shape
+        )
+        tma_atom_sfa2, tma_sfa2 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            w2_sf,
+            sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            ring_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.num_stages]
+            acc_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.acc_stages]
+            b_mbar: cute.struct.MemRange[cutlass.Int64, NUM_BAR]
+            tmem_holding: cutlass.Int32
+            tl_buf: cute.struct.MemRange[cutlass.Int64, TL_SLOTS]
+            hstage: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, HSTAGE_BYTES // 8], 16]
+            misc: cute.struct.MemRange[cutlass.Int32, 16]
+            closed: cute.struct.MemRange[cutlass.Int32, MAX_ROUNDS * FC2_MT]
+            bitmap: cute.struct.MemRange[cutlass.Uint32, 8]
+            topk_id: cute.struct.MemRange[cutlass.Int32, NTOK * TOP_K]
+            topk_w: cute.struct.MemRange[cutlass.Float32, NTOK * TOP_K]
+            slot_expert: cute.struct.MemRange[cutlass.Int32, MAX_SLOTS]
+            slot_mask: cute.struct.MemRange[cutlass.Uint32, MAX_SLOTS]
+            slot_w: cute.struct.MemRange[cutlass.Float32, MAX_SLOTS * NTOK]
+            slot_k: cute.struct.MemRange[cutlass.Int32, MAX_SLOTS * NTOK]
+            expert_slot: cute.struct.MemRange[cutlass.Int32, NUM_EXPERTS]
+            sA: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.a_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFA: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float8E4M3FN, cute.cosize(self.sfa_smem_layout_staged)
+                ],
+                1024,
+            ]
+            # this CTA's share of the activation rows (NTOK x 6 k-blocks)
+            sX: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, NTOK * XS_ROW_BYTES], 1024]
+            sB: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.b_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFB: cute.struct.Align[cute.struct.MemRange[cutlass.Float8E4M3FN, SFB_BYTES], 1024]
+            sB2: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.b2_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFB2: cute.struct.Align[cute.struct.MemRange[cutlass.Float8E4M3FN, SFB2_BYTES], 1024]
+
+        self.shared_storage = SharedStorage
+
+        self.kernel(
+            self.tiled_mma,
+            tma_atom_a1,
+            tma_a1,
+            tma_atom_sfa1,
+            tma_sfa1,
+            tma_atom_a2,
+            tma_a2,
+            tma_atom_sfa2,
+            tma_sfa2,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.sfb_smem_layout_std,
+            self.b2_smem_layout_staged,
+            self.sfb2_smem_layout_staged,
+            x_addr,
+            logits_addr,
+            bias_addr,
+            alpha1_addr,
+            scale_c1_addr,
+            alpha2_addr,
+            xgs_addr,
+            rows_addr,
+            raw_scale_addr,
+            raw_idx_addr,
+            out_addr,
+            counter_addr,
+            dbg_addr,
+            w13_addr,
+            w13_sf_addr,
+            w2_addr,
+            w2_sf_addr,
+            spec_ids_addr,
+            num_tokens,
+        ).launch(
+            grid=[self.num_ctas, 1, 1],
+            block=[NUM_THREADS, 1, 1],
+            cluster=[CLUSTER, 1, 1],
+            smem=self.shared_storage.size_in_bytes(),
+            min_blocks_per_mp=1,
+            stream=stream,
+            use_pdl=True,
+        )
+
+    # ------------------------------------------------------------------
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a1: cute.CopyAtom,
+        mA1: cute.Tensor,
+        tma_atom_sfa1: cute.CopyAtom,
+        mSFA1: cute.Tensor,
+        tma_atom_a2: cute.CopyAtom,
+        mA2: cute.Tensor,
+        tma_atom_sfa2: cute.CopyAtom,
+        mSFA2: cute.Tensor,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfb_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_std: cute.Layout,
+        b2_smem_layout_staged: cute.ComposedLayout,
+        sfb2_smem_layout_staged: cute.Layout,
+        x_addr: Int64,
+        logits_addr: Int64,
+        bias_addr: Int64,
+        alpha1_addr: Int64,
+        scale_c1_addr: Int64,
+        alpha2_addr: Int64,
+        xgs_addr: Int64,
+        rows_addr: Int64,
+        raw_scale_addr: Int64,
+        raw_idx_addr: Int64,
+        out_addr: Int64,
+        counter_addr: Int64,
+        dbg_addr: Int64,
+        w13_addr: Int64,
+        w13_sf_addr: Int64,
+        w2_addr: Int64,
+        w2_sf_addr: Int64,
+        spec_ids_addr: Int64,
+        num_tokens: Int32,
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        tidx, _, _ = cute.arch.thread_idx()
+        lane = tidx % 32
+        bidx, _, _ = cute.arch.block_idx()
+        grid, _, _ = cute.arch.grid_dim()
+        bidx = Int32(bidx)
+        rk = ptx.cluster_ctarank()
+        cl = bidx // CLUSTER
+        ncl = Int32(grid) // CLUSTER
+
+        if warp_idx == TMA_WARP:
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a1)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa1)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a2)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa2)
+            if cutlass.const_expr(self.spec_k > 0):
+                self._xc_spec_prefetch(
+                    bidx,
+                    Int32(grid),
+                    lane,
+                    spec_ids_addr,
+                    w13_addr,
+                    w13_sf_addr,
+                    w2_addr,
+                    w2_sf_addr,
+                )
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        ring = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.ring_mbar.data_ptr(),
+            num_stages=self.num_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+            tx_count=A_STAGE_BYTES + SFA_STAGE_BYTES,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+        acc_pipe = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_mbar.data_ptr(),
+            num_stages=self.acc_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, EPI_WARPS),
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+        b_mbar = storage.b_mbar.data_ptr()
+        if warp_idx == MMA_WARP:
+            with cute.arch.elect_one():
+                for kb in cutlass.range_constexpr(FC1_KB):
+                    cute.arch.mbarrier_init(b_mbar + kb, 1)
+                cute.arch.mbarrier_init(b_mbar + X_BAR, 1)
+                for rho in cutlass.range_constexpr(MAX_ROUNDS):
+                    cute.arch.mbarrier_init(b_mbar + B2_BAR + rho, CLUSTER)
+                # k-blocks quantized by a peer complete on its bulk copies
+                for kb in cutlass.range_constexpr(FC1_KB):
+                    if Int32(kb // KB_PER_CTA) != rk:
+                        ptx.mbar_arrive_expect_tx(Int32((b_mbar + kb).toint()), Int32(KB_TX_BYTES))
+        cute.arch.mbarrier_init_fence()
+        # Peers arrive on this CTA's operand barriers: their initialization
+        # must be visible cluster-wide first (waited before the first remote
+        # arrive, below).
+        cute.arch.cluster_arrive_relaxed()
+
+        if warp_idx == 0:
+            cute.arch.alloc_tmem(TMEM_COLS, storage.tmem_holding)
+
+        s_misc = storage.misc.get_tensor(cute.make_layout(16))
+        s_bitmap = storage.bitmap.get_tensor(cute.make_layout(8))
+        s_topk_id = storage.topk_id.get_tensor(cute.make_layout(NTOK * TOP_K))
+        s_topk_w = storage.topk_w.get_tensor(cute.make_layout(NTOK * TOP_K))
+        s_slot_expert = storage.slot_expert.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_slot_mask = storage.slot_mask.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_slot_w = storage.slot_w.get_tensor(cute.make_layout(MAX_SLOTS * NTOK))
+        s_slot_k = storage.slot_k.get_tensor(cute.make_layout(MAX_SLOTS * NTOK))
+        s_expert_slot = storage.expert_slot.get_tensor(cute.make_layout(NUM_EXPERTS))
+
+        s_expert_slot[tidx] = Int32(-1)
+        s_slot_w[tidx] = Float32(0.0)
+        s_slot_w[tidx + NUM_THREADS] = Float32(0.0)
+        if tidx < MAX_SLOTS:
+            s_slot_mask[tidx] = Uint32(0)
+        if tidx < 16:
+            s_misc[tidx] = Int32(0)
+        if cutlass.const_expr(self.debug & DEBUG_TIMELINE):
+            if tidx < TL_SLOTS:
+                ptx.sts_u64(Int32(storage.tl_buf.data_ptr().toint()) + tidx * 8, Int64(0))
+        if tidx < 8:
+            s_bitmap[tidx] = Uint32(0)
+        bias_r = cute.make_rmem_tensor((8,), Float32)
+        for j in cutlass.range_constexpr(8):
+            bias_r[j] = ptx.bf16x2_to_f32(ptx.ld_nc_u16(bias_addr + (lane + 32 * j) * 2), False)
+        cute.arch.sync_threads()
+
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_ENTRY)
+            if cutlass.const_expr(tl):
+                ptx.sts_u64(tl_s + XT_GT0 * 8, ptx.globaltimer_u64())
+        cute.arch.griddepcontrol_wait()
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_PDL)
+        if cutlass.const_expr(self.delay & DELAY_ENTRY):
+            self._xc_delay(DELAY_ENTRY * 53, bidx)
+        if warp_idx == TMA_WARP:
+            with cute.arch.elect_one():
+                xbar = Int32((b_mbar + X_BAR).toint())
+                sx = Int32(storage.sX.data_ptr().toint())
+                ptx.mbar_arrive_expect_tx(xbar, num_tokens * XS_ROW_BYTES)
+                r = Int32(0)
+                while r < num_tokens:
+                    ptx.bulk_g2s(
+                        sx + r * XS_ROW_BYTES,
+                        x_addr + r * X_ROW_BYTES + rk * XS_ROW_BYTES,
+                        Int32(XS_ROW_BYTES),
+                        xbar,
+                    )
+                    r = r + 1
+        # Every CTA of the cluster has initialized its barriers: remote
+        # copies and arrivals may start.
+        cute.arch.cluster_wait()
+        if cutlass.const_expr(self.trigger == "early"):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        # -------- routing (warp t handles token t) -> slot tables --------
+        self._route_tokens(
+            warp_idx, lane, num_tokens, logits_addr, bias_r, s_topk_id, s_topk_w, storage
+        )
+        cute.arch.sync_threads()
+        e = tidx
+        bw = s_bitmap[e // 32]
+        prefix = Int32(0)
+        for w in cutlass.range_constexpr(8):
+            word = s_bitmap[w]
+            if w < e // 32:
+                prefix = prefix + ptx.popc_u32(word)
+        below = bw & ((Uint32(1) << Uint32(e % 32)) - Uint32(1))
+        prefix = prefix + ptx.popc_u32(below)
+        if ((bw >> Uint32(e % 32)) & Uint32(1)) != Uint32(0):
+            s_slot_expert[prefix] = e
+            s_expert_slot[e] = prefix
+        if tidx == 0:
+            total = Int32(0)
+            for w in cutlass.range_constexpr(8):
+                total = total + ptx.popc_u32(s_bitmap[w])
+            s_misc[0] = total
+        cute.arch.sync_threads()
+        if tidx < NTOK * TOP_K:
+            t = tidx // TOP_K
+            k = tidx % TOP_K
+            if t < num_tokens:
+                e = s_topk_id[tidx]
+                sl = s_expert_slot[e]
+                s_slot_w[sl * NTOK + t] = s_topk_w[tidx]
+                s_slot_k[sl * NTOK + t] = k
+                ptx.atom_or_shared_u32(
+                    Int32(storage.slot_mask.data_ptr().toint()) + sl * 4, Uint32(1) << Uint32(t)
+                )
+        cute.arch.sync_threads()
+        num_slots = s_misc[0]
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_ROUTED)
+
+        # -------- expert rounds of this cluster --------
+        # Round 0 spreads min(D, clusters) experts over all clusters (an
+        # expert served by r clusters splits its FC2 tiles r ways); when
+        # D > clusters, round 1 spreads the remaining D - clusters experts
+        # over all clusters the same way. Every cluster then streams two FC1
+        # tiles per CTA and an even share of the FC2 tiles.
+        nact = ncl
+        if cutlass.const_expr(self.min_clusters > 0):
+            if num_slots <= ncl:
+                nact = cutlass.max(num_slots, Int32(self.min_clusters))
+                nact = cutlass.min(nact, ncl)
+        if cutlass.const_expr(self.two_round_clusters > 0):
+            if num_slots > ncl:
+                if num_slots <= Int32(MAX_ROUNDS * self.two_round_clusters):
+                    nact = cutlass.min(Int32(self.two_round_clusters), ncl)
+        rounds = (num_slots + nact - 1) // nact
+        if cl >= nact:
+            rounds = Int32(0)
+        slots = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        shares = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        firsts = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        t_lo = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        t_hi = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            base = Int32(0)
+            count = cutlass.min(num_slots, nact)
+            if cutlass.const_expr(rho == 1):
+                base = nact
+                count = cutlass.max(num_slots - nact, Int32(1))
+            i = cl % count
+            j = cl // count
+            nshare = (nact - 1 - i) // count + 1
+            a = (j * FC2_MT) // nshare
+            n = ((j + 1) * FC2_MT) // nshare - a
+            slots[rho] = base + i
+            shares[rho] = nshare
+            firsts[rho] = Int32(0)
+            if j == 0:
+                firsts[rho] = Int32(1)
+            t_lo[rho] = a + (rk * n) // CLUSTER
+            t_hi[rho] = a + ((rk + 1) * n) // CLUSTER
+
+        if warp_idx == TMA_WARP:
+            self._xc_tma_producer(
+                tma_atom_a1,
+                mA1,
+                tma_atom_sfa1,
+                mSFA1,
+                tma_atom_a2,
+                mA2,
+                tma_atom_sfa2,
+                mSFA2,
+                a_smem_layout_staged,
+                sfa_smem_layout_staged,
+                tiled_mma,
+                cluster_layout_vmnk,
+                storage,
+                ring,
+                b_mbar,
+                s_slot_expert,
+                rk,
+                rounds,
+                slots,
+                shares,
+                firsts,
+                t_lo,
+                t_hi,
+                w13_addr,
+                w13_sf_addr,
+                w2_addr,
+                w2_sf_addr,
+            )
+        if cutlass.const_expr(self.spec_self):
+            if bidx == 0:
+                if warp_idx == MMA_WARP:
+                    self._xc_spec_store(lane, num_slots, spec_ids_addr, storage, s_slot_expert)
+        if warp_idx == MMA_WARP:
+            self._xc_mma_issuer(
+                tiled_mma,
+                storage,
+                ring,
+                acc_pipe,
+                b_mbar,
+                a_smem_layout_staged,
+                sfa_smem_layout_staged,
+                b_smem_layout_staged,
+                sfb_smem_layout_staged,
+                sfb_smem_layout_std,
+                b2_smem_layout_staged,
+                sfb2_smem_layout_staged,
+                rounds,
+                t_lo,
+                t_hi,
+            )
+        if (warp_idx != TMA_WARP) & (warp_idx != MMA_WARP):
+            if rounds > 0:
+                self._xc_quantize_x(
+                    warp_idx, lane, xgs_addr, num_tokens, storage, b_mbar, dbg_addr, rk
+                )
+            else:
+                # an idle cluster still lets its activation copy land
+                cute.arch.mbarrier_wait(b_mbar + X_BAR, 0)
+
+        if bidx == 0:
+            if tidx < NTOK * TOP_K:
+                t = tidx // TOP_K
+                if t < num_tokens:
+                    ptx.st_u16(raw_scale_addr + tidx * 2, ptx.bf16_bits(s_topk_w[tidx]))
+                    ptx.st_u32(raw_idx_addr + tidx * 4, Uint32(tidx))
+        if cutlass.const_expr(self.debug & DEBUG_ROUTING):
+            if bidx == 0:
+                if tidx < NTOK * TOP_K:
+                    ptx.st_u32(dbg_addr + tidx * 4, Uint32(s_topk_id[tidx]))
+                    ptx.st_f32(dbg_addr + 256 + tidx * 4, s_topk_w[tidx])
+                if tidx < MAX_SLOTS:
+                    ptx.st_u32(dbg_addr + 512 + tidx * 4, Uint32(s_slot_expert[tidx]))
+                    ptx.st_u32(dbg_addr + 768 + tidx * 4, s_slot_mask[tidx])
+                if tidx == 0:
+                    ptx.st_u32(dbg_addr + 1024, Uint32(num_slots))
+
+        if warp_idx < EPI_WARPS:
+            self._xc_epilogue(
+                tiled_mma,
+                tidx,
+                warp_idx,
+                lane,
+                storage,
+                acc_pipe,
+                b_mbar,
+                s_slot_expert,
+                s_slot_mask,
+                s_slot_k,
+                rk,
+                rounds,
+                slots,
+                firsts,
+                t_lo,
+                t_hi,
+                num_tokens,
+                alpha1_addr,
+                scale_c1_addr,
+                alpha2_addr,
+                rows_addr,
+                out_addr,
+                counter_addr,
+                s_topk_w,
+                num_slots,
+                bidx,
+                dbg_addr,
+            )
+        # No CTA leaves while a peer may still store into its shared memory.
+        cute.arch.cluster_arrive()
+        cute.arch.cluster_wait()
+        if warp_idx < EPI_WARPS:
+            if cutlass.const_expr(self.debug & DEBUG_TIMELINE):
+                ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+                if tidx == 0:
+                    _stamp(tl, tl_s, XT_EXIT)
+                ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+                if tidx < TL_SLOTS:
+                    ptx.st_u64(
+                        dbg_addr + TL_BYTE_BASE + (bidx * TL_SLOTS + tidx) * 8,
+                        ptx.lds_u64(tl_s + tidx * 8),
+                    )
+            if warp_idx == 0:
+                tmem_ptr = cute.arch.retrieve_tmem_ptr(
+                    Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+                )
+                cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_tma_producer(
+        self,
+        tma_atom_a1,
+        mA1,
+        tma_atom_sfa1,
+        mSFA1,
+        tma_atom_a2,
+        mA2,
+        tma_atom_sfa2,
+        mSFA2,
+        a_smem_layout_staged,
+        sfa_smem_layout_staged,
+        tiled_mma,
+        cluster_layout_vmnk,
+        storage,
+        ring,
+        b_mbar,
+        s_slot_expert,
+        rk,
+        rounds,
+        slots,
+        shares,
+        firsts,
+        t_lo,
+        t_hi,
+        w13_addr,
+        w13_sf_addr,
+        w2_addr,
+        w2_sf_addr,
+    ):
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        lane0 = cute.arch.lane_idx() == 0
+        sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
+        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+        thr_mma = tiled_mma.get_slice(0)
+        a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+        gA1 = cute.local_tile(mA1, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
+        gSFA1 = cute.local_tile(
+            mSFA1, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        gA2 = cute.local_tile(mA2, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
+        gSFA2 = cute.local_tile(
+            mSFA2, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        tAsA, tAgA1 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_a1,
+            0,
+            a_cta_layout,
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gA1), 0, 3),
+        )
+        tAsSFA, tAgSFA1 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_sfa1,
+            0,
+            a_cta_layout,
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gSFA1), 0, 3),
+        )
+        tAsSFA = cute.filter_zeros(tAsSFA)
+        tAgSFA1 = cute.filter_zeros(tAgSFA1)
+        _, tAgA2 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_a2,
+            0,
+            a_cta_layout,
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gA2), 0, 3),
+        )
+        _, tAgSFA2 = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_sfa2,
+            0,
+            a_cta_layout,
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(thr_mma.partition_A(gSFA2), 0, 3),
+        )
+        tAgSFA2 = cute.filter_zeros(tAgSFA2)
+
+        state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_stages)
+        if cutlass.const_expr(self.l2pf & L2PF_EARLY):
+            self._xc_l2_prefetch(
+                tl,
+                tl_s,
+                rk,
+                rounds,
+                slots,
+                firsts,
+                t_lo,
+                t_hi,
+                s_slot_expert,
+                w13_addr,
+                w13_sf_addr,
+                w2_addr,
+                w2_sf_addr,
+            )
+        # FC1: tile rk of each round's expert, all k-blocks. A weight read by
+        # several clusters stays in L2 for the others (evict-normal).
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                expert = s_slot_expert[slots[rho]]
+                policy = Int64(L2_EVICT_FIRST)
+                if shares[rho] > 1:
+                    policy = Int64(L2_EVICT_NORMAL)
+                kb = Int32(0)
+                while kb < FC1_KB:
+                    if cutlass.const_expr(
+                        rho == 0 and self.l2pf != 0 and not (self.l2pf & L2PF_EARLY)
+                    ):
+                        # after the first ring fill of round 0
+                        if kb == Int32(self.num_stages):
+                            self._xc_l2_prefetch(
+                                tl,
+                                tl_s,
+                                rk,
+                                rounds,
+                                slots,
+                                firsts,
+                                t_lo,
+                                t_hi,
+                                s_slot_expert,
+                                w13_addr,
+                                w13_sf_addr,
+                                w2_addr,
+                                w2_sf_addr,
+                            )
+                    ring.producer_acquire(state)
+                    bar = ring.producer_get_barrier(state)
+                    cute.copy(
+                        tma_atom_a1,
+                        tAgA1[(None, rk, kb, expert)],
+                        tAsA[(None, state.index)],
+                        tma_bar_ptr=bar,
+                        cache_policy=policy,
+                    )
+                    cute.copy(
+                        tma_atom_sfa1,
+                        tAgSFA1[(None, rk, kb, expert)],
+                        tAsSFA[(None, state.index)],
+                        tma_bar_ptr=bar,
+                        cache_policy=policy,
+                    )
+                    state.advance()
+                    kb = kb + 1
+                if cutlass.const_expr(rho == 0):
+                    if lane0:
+                        _stamp(tl, tl_s, XT_TMA_FC1_R0)
+        if lane0:
+            _stamp(tl, tl_s, XT_TMA_FC1)
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                expert = s_slot_expert[slots[rho]]
+                t = t_lo[rho]
+                while t < t_hi[rho]:
+                    ring.producer_acquire(state)
+                    bar = ring.producer_get_barrier(state)
+                    cute.copy(
+                        tma_atom_a2,
+                        tAgA2[(None, t, 0, expert)],
+                        tAsA[(None, state.index)],
+                        tma_bar_ptr=bar,
+                        cache_policy=Int64(L2_EVICT_FIRST),
+                    )
+                    cute.copy(
+                        tma_atom_sfa2,
+                        tAgSFA2[(None, t, 0, expert)],
+                        tAsSFA[(None, state.index)],
+                        tma_bar_ptr=bar,
+                        cache_policy=Int64(L2_EVICT_FIRST),
+                    )
+                    state.advance()
+                    t = t + 1
+        if lane0:
+            _stamp(tl, tl_s, XT_TMA_FC2)
+        ring.producer_tail(state)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_l2_prefetch(
+        self,
+        tl,
+        tl_s,
+        rk,
+        rounds,
+        slots,
+        firsts,
+        t_lo,
+        t_hi,
+        s_slot_expert,
+        w13_addr,
+        w13_sf_addr,
+        w2_addr,
+        w2_sf_addr,
+    ):
+        """Bulk L2 prefetch of weight bytes this CTA's ring reaches later.
+
+        ``L2PF_FC1_R1``: FC1 tile ``rk`` of the round-1 expert, by the first
+        cluster serving that slot only (the others read it from L2 anyway);
+        one 12 KB chunk per lane, lane 0 also the tile's scales.
+        ``L2PF_FC2``: this CTA's FC2 tiles of both rounds, one tile (16 KB +
+        2 KB scales) per lane. Prefetches change no data, only where the
+        later TMA loads find it.
+        """
+        lane = cute.arch.lane_idx()
+        if cutlass.const_expr(self.l2pf & L2PF_FC1_R1):
+            if rounds > 1:
+                if firsts[1] != 0:
+                    e1 = Int64(s_slot_expert[slots[1]])
+                    base = w13_addr + e1 * W13_EXPERT_BYTES + Int64(rk) * Int64(FC1_TILE_BYTES)
+                    ptx.prefetch_l2_bulk(
+                        base + Int64(lane) * Int64(PF_FC1_CHUNK), Int32(PF_FC1_CHUNK)
+                    )
+                    if lane == 0:
+                        ptx.prefetch_l2_bulk(
+                            w13_sf_addr
+                            + e1 * W13_SF_EXPERT_BYTES
+                            + Int64(rk) * Int64(W13_SF_TILE_BYTES),
+                            Int32(W13_SF_TILE_BYTES),
+                        )
+        if cutlass.const_expr(self.l2pf & L2PF_FC2):
+            for rho in cutlass.range_constexpr(MAX_ROUNDS):
+                if rho < rounds:
+                    e2 = Int64(s_slot_expert[slots[rho]])
+                    t = t_lo[rho] + lane
+                    if t < t_hi[rho]:
+                        ptx.prefetch_l2_bulk(
+                            w2_addr + e2 * W2_EXPERT_BYTES + Int64(t) * Int64(W2_TILE_BYTES),
+                            Int32(W2_TILE_BYTES),
+                        )
+                        ptx.prefetch_l2_bulk(
+                            w2_sf_addr
+                            + e2 * W2_SF_EXPERT_BYTES
+                            + Int64(t) * Int64(W2_SF_TILE_BYTES),
+                            Int32(W2_SF_TILE_BYTES),
+                        )
+        if lane == 0:
+            _stamp(tl, tl_s, XT_PF_ISSUED)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_spec_prefetch(
+        self, bidx, grid, lane, ids_addr, w13_addr, w13_sf_addr, w2_addr, w2_sf_addr
+    ):
+        """Bulk L2 prefetch of the first ``spec_k`` experts at ``ids_addr``.
+
+        Runs before the dependency wait: the IDs are written on another stream
+        (cross_layer_expert_prefetch) and only choose which weight bytes are prefetched, so
+        a stale or concurrently updated ID costs bandwidth, never correctness
+        (each ID is masked into [0, 256)). Chunk ``c`` of the ``spec_k`` x
+        ``SPEC_CPE`` chunks (expert-rank major) is issued by lane ``c %
+        (32 * n)`` of the first ``n`` CTAs, so the best-ranked experts go
+        first.
+        """
+        nsp = grid
+        if cutlass.const_expr(self.spec_ctas > 0):
+            nsp = cutlass.min(grid, Int32(self.spec_ctas))
+        if bidx < nsp:
+            c = bidx * 32 + lane
+            while c < Int32(self.spec_k * SPEC_CPE):
+                slot = c // SPEC_CPE
+                off = c - slot * SPEC_CPE
+                ex = Int64(
+                    Int32(
+                        ptx.ld_relaxed_gpu_u32(ids_addr + Int64(slot) * 4) & Uint32(NUM_EXPERTS - 1)
+                    )
+                )
+                o = Int64(off) * SPEC_CHUNK
+                addr = w13_addr + ex * W13_EXPERT_BYTES + o
+                if off >= SPEC_C13:
+                    addr = w13_sf_addr + ex * W13_SF_EXPERT_BYTES + o - SPEC_C13 * SPEC_CHUNK
+                if off >= SPEC_C13 + SPEC_C13SF:
+                    addr = w2_addr + ex * W2_EXPERT_BYTES + o - (SPEC_C13 + SPEC_C13SF) * SPEC_CHUNK
+                if off >= SPEC_C13 + SPEC_C13SF + SPEC_C2:
+                    addr = (
+                        w2_sf_addr
+                        + ex * W2_SF_EXPERT_BYTES
+                        + o
+                        - (SPEC_C13 + SPEC_C13SF + SPEC_C2) * SPEC_CHUNK
+                    )
+                ptx.prefetch_l2_bulk(addr, Int32(SPEC_CHUNK))
+                c = c + nsp * 32
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_spec_store(self, lane, num_slots, ids_addr, storage, s_slot_expert):
+        """Write this call's distinct experts to ``ids_addr``, ranked by the
+        number of rows that select them (ties: lower expert ID first); slots
+        past ``num_slots`` keep their previous IDs."""
+        mask_s = Int32(storage.slot_mask.data_ptr().toint())
+        for h in cutlass.range_constexpr(MAX_SLOTS // 32):
+            i = lane + 32 * h
+            if i < num_slots:
+                ci = ptx.popc_u32(ptx.lds_u32(mask_s + i * 4))
+                rank = Int32(0)
+                j = Int32(0)
+                while j < num_slots:
+                    cj = ptx.popc_u32(ptx.lds_u32(mask_s + j * 4))
+                    if cj > ci:
+                        rank = rank + 1
+                    if cj == ci:
+                        if j < i:
+                            rank = rank + 1
+                    j = j + 1
+                if rank < SPEC_KMAX:
+                    ptx.st_relaxed_gpu_u32(ids_addr + Int64(rank) * 4, Uint32(s_slot_expert[i]))
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_mma_issuer(
+        self,
+        tiled_mma,
+        storage,
+        ring,
+        acc_pipe,
+        b_mbar,
+        a_smem_layout_staged,
+        sfa_smem_layout_staged,
+        b_smem_layout_staged,
+        sfb_smem_layout_staged,
+        sfb_smem_layout_std,
+        b2_smem_layout_staged,
+        sfb2_smem_layout_staged,
+        rounds,
+        t_lo,
+        t_hi,
+    ):
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        lane0 = cute.arch.lane_idx() == 0
+        sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
+        sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
+        sB2 = storage.sB2.get_tensor(
+            b2_smem_layout_staged.outer, swizzle=b2_smem_layout_staged.inner
+        )
+        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+        sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
+        sSFB2 = storage.sSFB2.get_tensor(sfb2_smem_layout_staged)
+        tCrA = tiled_mma.make_fragment_A(sA)
+        tCrB = tiled_mma.make_fragment_B(sB)
+        tCrB2 = tiled_mma.make_fragment_B(sB2)
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.acc_stages))
+        acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
+            Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+        )
+        tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+        sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + TMEM_SFA_COL, dtype=cutlass.Float8E4M3FN)
+        tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            self.sf_vec_size,
+            cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+        )
+        tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+        sfb_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + TMEM_SFB_COL, dtype=cutlass.Float8E4M3FN)
+        tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            self.sf_vec_size,
+            cute.slice_(sfb_smem_layout_std, (None, None, None, 0)),
+        )
+        tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+        copy_s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = self._s2t_partition(sSFA, tCtSFA)
+        copy_s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = self._s2t_partition(sSFB, tCtSFB)
+        _, tCsSFB2_s2t, _ = self._s2t_partition(sSFB2, tCtSFB)
+
+        ring_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_stages
+        )
+        acc_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.acc_stages
+        )
+        num_kphases = cute.size(tCrA, mode=[2])
+
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                acc_pipe.producer_acquire(acc_state)
+                tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)]
+                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                kb = Int32(0)
+                while kb < FC1_KB:
+                    cute.arch.mbarrier_wait(b_mbar + kb, 0)
+                    ring.consumer_wait(ring_state)
+                    cute.copy(
+                        copy_s2t_sfa,
+                        tCsSFA_s2t[(None, None, None, None, ring_state.index)],
+                        tCtSFA_s2t,
+                    )
+                    cute.copy(copy_s2t_sfb, tCsSFB_s2t[(None, None, None, None, kb)], tCtSFB_s2t)
+                    for kp in cutlass.range_constexpr(num_kphases):
+                        tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator)
+                        tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator)
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc,
+                            tCrA[(None, None, kp, ring_state.index)],
+                            tCrB[(None, None, kp, kb)],
+                            tCtAcc,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    ring.consumer_release(ring_state)
+                    ring_state.advance()
+                    kb = kb + 1
+                acc_pipe.producer_commit(acc_state)
+                acc_state.advance()
+                if cutlass.const_expr(rho == 0):
+                    if lane0:
+                        _stamp(tl, tl_s, XT_MMA_FC1_R0)
+        if lane0:
+            _stamp(tl, tl_s, XT_MMA_FC1)
+
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                # this round's FC2 operand: the four FC1 tiles of the cluster
+                bar = Int32((b_mbar + B2_BAR + rho).toint())
+                done = ptx.mbar_try_wait_cluster(bar, Int32(0))
+                while done == 0:
+                    done = ptx.mbar_try_wait_cluster(bar, Int32(0))
+                ptx.fence_proxy_async_shared()
+                if cutlass.const_expr(rho == 0):
+                    if lane0:
+                        _stamp(tl, tl_s, XT_B2_READY)
+                else:
+                    if lane0:
+                        _stamp(tl, tl_s, XT_B2_READY_R1)
+                t = t_lo[rho]
+                while t < t_hi[rho]:
+                    acc_pipe.producer_acquire(acc_state)
+                    tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)]
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    ring.consumer_wait(ring_state)
+                    cute.copy(
+                        copy_s2t_sfa,
+                        tCsSFA_s2t[(None, None, None, None, ring_state.index)],
+                        tCtSFA_s2t,
+                    )
+                    cute.copy(copy_s2t_sfb, tCsSFB2_s2t[(None, None, None, None, rho)], tCtSFB_s2t)
+                    for kp in cutlass.range_constexpr(num_kphases):
+                        tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator)
+                        tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator)
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc,
+                            tCrA[(None, None, kp, ring_state.index)],
+                            tCrB2[(None, None, kp, rho)],
+                            tCtAcc,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    ring.consumer_release(ring_state)
+                    ring_state.advance()
+                    acc_pipe.producer_commit(acc_state)
+                    acc_state.advance()
+                    t = t + 1
+        if lane0:
+            _stamp(tl, tl_s, XT_MMA_FC2)
+        acc_pipe.producer_tail(acc_state)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_quantize_x(self, warp_idx, lane, xgs_addr, num_tokens, storage, b_mbar, dbg_addr, rk):
+        """NVFP4-quantize this CTA's six activation k-blocks (same arithmetic
+        as the stock ``fp4_quantize``), one k-block per quantizing warp, into
+        the FC1 B operand of every CTA of the cluster.
+
+        The k-block is written to the local B / SFB stage and released on the
+        local mbarrier, then bulk-copied into the same stage of the three
+        peers, completing on their k-block mbarriers (which expect the bytes).
+        """
+        sb_base = Int32(storage.sB.data_ptr().toint())
+        sfb_base = Int32(storage.sSFB.data_ptr().toint())
+        gs = ptx.ld_nc_f32(xgs_addr)
+        rgs = ptx.rcp_approx_ftz(gs)
+        rcp6 = ptx.rcp_approx_ftz(Float32(6.0))
+        qw = warp_idx
+        if warp_idx > TMA_WARP:
+            qw = warp_idx - 2
+        sx = Int32(storage.sX.data_ptr().toint())
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        blk = lane // 2
+        sf_off = (blk // 4) * SFB_CHUNK_STRIDE + blk % 4
+        cute.arch.mbarrier_wait(b_mbar + X_BAR, 0)
+        kb = rk * KB_PER_CTA + qw
+        self._quant_rows_xs(
+            0, kb, qw, lane, num_tokens, sx, sb_base, sfb_base, sf_off, gs, rgs, rcp6
+        )
+        if num_tokens > NTOK // 2:
+            self._quant_rows_xs(
+                1, kb, qw, lane, num_tokens, sx, sb_base, sfb_base, sf_off, gs, rgs, rcp6
+            )
+        else:
+            if lane < NTOK // 2 * 4:
+                n = NTOK // 2 + lane // 4
+                ptx.sts_u32(
+                    sfb_base + kb * SFB_STAGE_STRIDE + (lane % 4) * SFB_CHUNK_STRIDE + n * 16,
+                    Uint32(0),
+                )
+        ptx.fence_proxy_async_shared()
+        cute.arch.sync_warp()
+        kbar = Int32((b_mbar + kb).toint())
+        if lane == 0:
+            cute.arch.mbarrier_arrive(b_mbar + kb)
+        # lanes 1..6: (peer, B or SFB) bulk copies into the other CTAs
+        if (lane >= 1) & (lane <= 2 * (CLUSTER - 1)):
+            peer = (lane - 1) // 2
+            if peer >= rk:
+                peer = peer + 1
+            src = sb_base + kb * 1024
+            size = Int32(1024)
+            if ((lane - 1) % 2) != 0:
+                src = sfb_base + kb * SFB_STAGE_STRIDE
+                size = Int32(SFB_STAGE_STRIDE)
+            ptx.bulk_s2cluster(ptx.mapa_u32(src, peer), src, size, ptx.mapa_u32(kbar, peer))
+        if (warp_idx == 0) & (lane == 0):
+            _stamp(tl, tl_s, XT_Q_DONE)
+        if cutlass.const_expr(self.debug & DEBUG_QUANT):
+            # debug only: wait for every k-block (local and peers'), then CTA 0
+            # dumps B / SFB
+            for k2 in cutlass.range_constexpr(FC1_KB):
+                cute.arch.mbarrier_wait(b_mbar + k2, 0)
+            qt = warp_idx * 32 + lane
+            if warp_idx > TMA_WARP:
+                qt = qt - 64
+            ptx.bar_sync(BAR_QUANT, 32 * QUANT_WARPS)
+            cta_q, _, _ = cute.arch.block_idx()
+            if cta_q == 0:
+                w = qt
+                while w < (FC1_KB * 1024 + SFB_BYTES) // 4:
+                    if w < FC1_KB * 256:
+                        ptx.st_u32(dbg_addr + 4096 + w * 4, ptx.lds_u32(sb_base + w * 4))
+                    else:
+                        r2 = w - FC1_KB * 256
+                        ptx.st_u32(dbg_addr + 4096 + w * 4, ptx.lds_u32(sfb_base + r2 * 4))
+                    w = w + 32 * QUANT_WARPS
+
+    @cute.jit
+    def _quant_rows_xs(
+        self, half, kb, kb_local, lane, num_tokens, sx, sb_base, sfb_base, sf_off, gs, rgs, rcp6
+    ):
+        """Rows ``4*half .. 4*half+3`` of k-block ``kb``, read from this CTA's
+        staged share of the activation rows (``kb_local`` within it); rows at
+        or beyond ``num_tokens`` get zero data and zero scales."""
+        rows = NTOK // 2
+        d = cute.make_rmem_tensor((rows * 4,), Uint32)
+        for r in cutlass.range_constexpr(rows):
+            n = half * rows + r
+            x4 = ptx.lds_v4_u32(sx + n * XS_ROW_BYTES + kb_local * 512 + lane * 16)
+            for c in cutlass.range_constexpr(4):
+                d[4 * r + c] = x4[c]
+        for r in cutlass.range_constexpr(rows):
+            n = half * rows + r
+            vals = cute.make_rmem_tensor((8,), Float32)
+            for i in cutlass.range_constexpr(4):
+                vals[2 * i] = ptx.bf16x2_to_f32(d[4 * r + i], False)
+                vals[2 * i + 1] = ptx.bf16x2_to_f32(d[4 * r + i], True)
+            amax = Float32(0.0)
+            for i in cutlass.range_constexpr(8):
+                amax = cute.arch.fmax(amax, cute.math.absf(vals[i]))
+            amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, 1))
+            sf = ptx.cvt_e4m3_rn_sat(gs * (amax * rcp6))
+            sf_f = ptx.e4m3_to_f32(sf)
+            oscale = Float32(0.0)
+            if amax != Float32(0.0):
+                oscale = ptx.rcp_approx_ftz(sf_f * rgs)
+            word = ptx.cvt_e2m1x8(
+                vals[0] * oscale,
+                vals[1] * oscale,
+                vals[2] * oscale,
+                vals[3] * oscale,
+                vals[4] * oscale,
+                vals[5] * oscale,
+                vals[6] * oscale,
+                vals[7] * oscale,
+            )
+            if n >= num_tokens:
+                word = Uint32(0)
+                sf = Uint32(0)
+            ptx.sts_u32(
+                sb_base + kb * 1024 + n * 128 + (((lane // 4) ^ n) * 16) + (lane % 4) * 4, word
+            )
+            if lane % 2 == 0:
+                ptx.sts_u8(sfb_base + kb * SFB_STAGE_STRIDE + n * 16 + sf_off, sf)
+
+    # ------------------------------------------------------------------
+    @cute.jit
+    def _xc_epilogue(
+        self,
+        tiled_mma,
+        tidx,
+        warp_idx,
+        lane,
+        storage,
+        acc_pipe,
+        b_mbar,
+        s_slot_expert,
+        s_slot_mask,
+        s_slot_k,
+        rk,
+        rounds,
+        slots,
+        firsts,
+        t_lo,
+        t_hi,
+        num_tokens,
+        alpha1_addr,
+        scale_c1_addr,
+        alpha2_addr,
+        rows_addr,
+        out_addr,
+        counter_addr,
+        s_topk_w,
+        num_slots,
+        bidx,
+        dbg_addr,
+    ):
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        hst = Int32(storage.hstage.data_ptr().toint())
+        acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
+            Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+        )
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.acc_stages))
+        tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+        lane_off = (warp_idx * 32) << 16
+        acc_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.acc_stages
+        )
+        v = cute.make_rmem_tensor((NTOK,), Float32)
+        sb2 = Int32(storage.sB2.data_ptr().toint())
+        ssfb2 = Int32(storage.sSFB2.data_ptr().toint())
+        b2bar = Int32((b_mbar + B2_BAR).toint())
+
+        # ---------------- FC1 tile epilogue -> cluster operand ----------------
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                slot = slots[rho]
+                expert = s_slot_expert[slot]
+                mask = s_slot_mask[slot]
+                a_gate = ptx.ld_nc_f32(alpha1_addr + expert * 4)
+                s_up = ptx.ld_nc_f32(scale_c1_addr + expert * 4)
+                acc_pipe.consumer_wait(acc_state)
+                if tidx == 0:
+                    _stamp(tl, tl_s, XT_EPI_FC1_ACC if rho == 0 else XT_EPI_FC1_ACC_R1)
+                vals = ptx.tmem_ld_32x32b_x8(
+                    Int32(tCtAcc_base[(None, None, None, acc_state.index)].iterator.toint())
+                    + lane_off
+                )
+                for n in cutlass.range_constexpr(NTOK):
+                    v[n] = vals[n]
+                cute.arch.fence_view_async_tmem_load()
+                if tidx == 0:
+                    _stamp(tl, tl_s, XT_EPI_FC1_TMEM if rho == 0 else XT_EPI_FC1_TMEM_R1)
+                with cute.arch.elect_one():
+                    acc_pipe.consumer_release(acc_state)
+                acc_state.advance()
+                if cutlass.const_expr(self.debug & DEBUG_FC1):
+                    # the first cluster of a slot dumps its tile's accumulators
+                    if firsts[rho] != 0:
+                        for n in cutlass.range_constexpr(NTOK):
+                            ptx.st_f32(
+                                dbg_addr
+                                + 4 * (65536 + ((slot * FC1_MT + rk) * 128 + tidx) * NTOK + n),
+                                v[n],
+                            )
+                self._swiglu_quant_to_stage(v, warp_idx, lane, mask, a_gate, s_up, num_tokens, hst)
+                if tidx == 0:
+                    _stamp(tl, tl_s, XT_EPI_FC1_SWQ if rho == 0 else XT_EPI_FC1_SWQ_R1)
+                if cutlass.const_expr(rho == 1):
+                    if lane == 0:
+                        _stamp(tl, tl_s, XT_EPI_FC1_SWQ_W_R1 + warp_idx)
+                if cutlass.const_expr(self.delay & DELAY_PRODUCER):
+                    self._xc_delay(DELAY_PRODUCER * 53, bidx)
+                ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+                if tidx == 0:
+                    _stamp(tl, tl_s, XT_EPI_FC1_SWIGLU if rho == 0 else XT_EPI_FC1_SWIGLU_R1)
+                # 64 threads: 16-byte FP4 chunks (row r, half h) for CTA d;
+                # 32 threads: the tile's four scale bytes of row r for CTA d.
+                if tidx < 4 * 2 * NTOK:
+                    d = tidx // (2 * NTOK)
+                    q = tidx - d * (2 * NTOK)
+                    r = q // 2
+                    h = q - r * 2
+                    c16 = 2 * rk + h
+                    w = ptx.lds_v4_u32(hst + r * 32 + h * 16)
+                    dst = ptx.mapa_u32(sb2 + rho * HIMG_B_BYTES + r * 128 + ((c16 ^ r) * 16), d)
+                    ptx.st_cluster_v4_u32(dst, w[0], w[1], w[2], w[3])
+                    ptx.fence_proxy_async_shared_cluster()
+                if (tidx >= 4 * 2 * NTOK) & (tidx < 4 * 3 * NTOK):
+                    q = tidx - 4 * 2 * NTOK
+                    d = q // NTOK
+                    r = q - d * NTOK
+                    w1 = ptx.lds_u32(hst + HSTAGE_SF + r * 4)
+                    dst = ptx.mapa_u32(
+                        ssfb2 + rho * SFB_STAGE_STRIDE + rk * SFB_CHUNK_STRIDE + r * 16, d
+                    )
+                    ptx.st_cluster_u32(dst, w1)
+                    ptx.fence_proxy_async_shared_cluster()
+                # all operand stores (and staging reads) precede the arrivals;
+                # the staging buffer is free for the next round afterwards
+                ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+                if tidx < CLUSTER:
+                    ptx.mbar_arrive_remote(ptx.mapa_u32(b2bar + rho * 8, tidx))
+                if cutlass.const_expr(rho == 0):
+                    if tidx == 0:
+                        _stamp(tl, tl_s, XT_EPI_FC1_PUB)
+                else:
+                    if tidx == 0:
+                        _stamp(tl, tl_s, XT_EPI_FC1_PUB_R1)
+
+        # ---------------- FC2 epilogue: raw rows t * 8 + k ----------------
+        hid = warp_idx * 32 + (lane % 8) * 4 + lane // 8
+        first = Int32(1)
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                slot = slots[rho]
+                expert = s_slot_expert[slot]
+                mask = s_slot_mask[slot]
+                a2 = ptx.ld_nc_f32(alpha2_addr + expert * 4)
+                kk = cute.make_rmem_tensor((NTOK,), Int32)
+                for n in cutlass.range_constexpr(NTOK):
+                    kk[n] = s_slot_k[slot * NTOK + n]
+                t = t_lo[rho]
+                while t < t_hi[rho]:
+                    acc_pipe.consumer_wait(acc_state)
+                    if tidx == 0:
+                        if first != 0:
+                            _stamp(tl, tl_s, XT_FC2_FIRST_ACC)
+                    first = Int32(0)
+                    vals = ptx.tmem_ld_32x32b_x8(
+                        Int32(tCtAcc_base[(None, None, None, acc_state.index)].iterator.toint())
+                        + lane_off
+                    )
+                    cute.arch.fence_view_async_tmem_load()
+                    with cute.arch.elect_one():
+                        acc_pipe.consumer_release(acc_state)
+                    acc_state.advance()
+                    for n in cutlass.range_constexpr(NTOK):
+                        if ((mask >> Uint32(n)) & Uint32(1)) != Uint32(0):
+                            y = ptx.bf16_round(vals[n] * a2)
+                            if cutlass.const_expr(self.red_out):
+                                w = s_topk_w[n * TOP_K + kk[n]]
+                                ptx.red_add_f32(out_addr + (n * HIDDEN + t * 128 + hid) * 4, w * y)
+                            else:
+                                row = n * TOP_K + kk[n]
+                                ptx.st_u16(
+                                    rows_addr + (row * HIDDEN + t * 128 + hid) * 2,
+                                    ptx.bf16_bits(y),
+                                )
+                    t = t + 1
+                if cutlass.const_expr(rho == 0):
+                    if tidx == 0:
+                        _stamp(tl, tl_s, XT_FC2_R0_DONE)
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_FC2_DONE)
+        if cutlass.const_expr(self.finalize):
+            self._xc_finalize_last_arriver(
+                tidx,
+                rounds,
+                t_lo,
+                t_hi,
+                num_slots,
+                num_tokens,
+                s_topk_w,
+                rows_addr,
+                out_addr,
+                counter_addr,
+                storage,
+            )
+        # the successor (finalize or the post-MoE hop) may start launching
+        if cutlass.const_expr(self.trigger == "late"):
+            cute.arch.griddepcontrol_launch_dependents()
+
+    @cute.jit
+    def _xc_finalize_last_arriver(
+        self,
+        tidx,
+        rounds,
+        t_lo,
+        t_hi,
+        num_slots,
+        num_tokens,
+        s_topk_w,
+        rows_addr,
+        out_addr,
+        counter_addr,
+        storage,
+    ):
+        """In-kernel native finalize, by the last contributor of each tile.
+
+        Every (expert, output tile) unit is computed by exactly one CTA, so
+        an output tile is complete once ``num_slots`` units have counted in.
+        After its raw rows are stored, a CTA releases them (one GPU-scope
+        fence), counts each of its units into the tile's counter, and
+        finalizes every tile for which it was the last arrival: FP32 FMA over
+        the top-k in order and one BF16 rounding (the native finalize), then
+        resets the counter for the next launch. Nobody waits for anyone.
+        """
+        # every thread releases its own row stores, then the units count in
+        ptx.fence_acq_rel_gpu()
+        ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+        closed_base = Int32(storage.closed.data_ptr().toint())
+        u = Int32(0)
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                n = t_hi[rho] - t_lo[rho]
+                if (tidx >= u) & (tidx < u + n):
+                    t = t_lo[rho] + (tidx - u)
+                    old = ptx.atom_add_relaxed_gpu_u32(counter_addr + t * 4, Uint32(1))
+                    closed = Int32(0)
+                    if old + Uint32(1) == Uint32(num_slots):
+                        closed = Int32(1)
+                        ptx.st_relaxed_gpu_u32(counter_addr + t * 4, Uint32(0))
+                        ptx.fence_acq_rel_gpu()
+                    ptx.sts_u32(closed_base + tidx * 4, Uint32(closed))
+                u = u + n
+        ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+        # one (token, 8 columns) task per thread: 16 x num_tokens <= 128
+        n_tok = tidx // 16
+        c8 = (tidx % 16) * 8
+        u = Int32(0)
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            if rho < rounds:
+                n = t_hi[rho] - t_lo[rho]
+                i = Int32(0)
+                while i < n:
+                    if ptx.lds_u32(closed_base + (u + i) * 4) != Uint32(0):
+                        if n_tok < num_tokens:
+                            col = (t_lo[rho] + i) * 128 + c8
+                            acc = cute.make_rmem_tensor((8,), Float32)
+                            for j in cutlass.range_constexpr(8):
+                                acc[j] = Float32(0.0)
+                            for k in cutlass.range_constexpr(TOP_K):
+                                w = s_topk_w[n_tok * TOP_K + k]
+                                y = ptx.ld_ll2(rows_addr + ((n_tok * TOP_K + k) * HIDDEN + col) * 2)
+                                for j in cutlass.range_constexpr(4):
+                                    acc[2 * j] = ptx.fma_rn(
+                                        w, ptx.bf16x2_to_f32(y[j], False), acc[2 * j]
+                                    )
+                                    acc[2 * j + 1] = ptx.fma_rn(
+                                        w, ptx.bf16x2_to_f32(y[j], True), acc[2 * j + 1]
+                                    )
+                            ptx.st_v4_u32(
+                                out_addr + (n_tok * HIDDEN + col) * 2,
+                                ptx.pack_bf16x2(acc[0], acc[1]),
+                                ptx.pack_bf16x2(acc[2], acc[3]),
+                                ptx.pack_bf16x2(acc[4], acc[5]),
+                                ptx.pack_bf16x2(acc[6], acc[7]),
+                            )
+                    i = i + 1
+                u = u + n
+
+    @cute.jit
+    def _xc_delay(self, salt, bidx):
+        """Stress testing: a pseudo-random third of the CTAs (re-drawn per
+        launch from the global timer) spins 1..6 units of ``delay_ns``."""
+        t0 = ptx.globaltimer_u64()
+        seed = Uint32(Int32(t0 >> Int64(12)) & Int32(0xFFFF))
+        h = (Uint32(bidx) * Uint32(37) + seed * Uint32(101) + Uint32(salt)) % Uint32(17)
+        if h < Uint32(6):
+            d = Int64(h + Uint32(1)) * Int64(self.delay_ns)
+            t = t0
+            while t - t0 < d:
+                t = ptx.globaltimer_u64()
+
+    @cute.jit
+    def _swiglu_quant_to_stage(self, v, warp_idx, lane, mask, a_gate, s_up, num_tokens, hst):
+        """SwiGLU + NVFP4 quantization of one 128-row FC1 tile into the
+        staging buffer: 8 rows x 32 FP4 bytes, then 8 rows x 4 scale bytes.
+
+        Row ``32w + q`` of the stock layout holds intermediate index
+        ``64*mt + 16*w + j`` with ``j = 2*(q % 8) + (q // 16)`` as the up row
+        when ``q % 16 < 8`` and its gate row at lane ``q + 8``, so one warp
+        owns exactly one 16-element scale block of the intermediate. Tokens
+        not routed to the expert get zero bytes and zero scales.
+        """
+        is_up = (lane % 16) < 8
+        rcp6 = ptx.rcp_approx_ftz(Float32(6.0))
+        for n in cutlass.range_constexpr(NTOK):
+            gate_raw = cute.arch.shuffle_sync_down(v[n], 8)
+            gr = gate_raw * a_gate
+            up = v[n] * s_up
+            sig = ptx.rcp_approx_ftz(
+                Float32(1.0) + ptx.ex2_approx(gr * Float32(-1.4426950408889634))
+            )
+            h = Float32(0.0)
+            if is_up:
+                h = up * (gr * sig)
+            am = cute.math.absf(h)
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 1))
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 2))
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 4))
+            am = cute.arch.fmax(am, cute.arch.shuffle_sync_bfly(am, 16))
+            sf = ptx.cvt_e4m3_rn_sat(am * rcp6)
+            sf_f = ptx.e4m3_to_f32(sf)
+            osc = Float32(0.0)
+            if am != Float32(0.0):
+                osc = ptx.rcp_approx_ftz(sf_f)
+            partner = cute.arch.shuffle_sync_down(h, 16)
+            byte = ptx.cvt_e2m1x2(h * osc, partner * osc)
+            routed = (((mask >> Uint32(n)) & Uint32(1)) != Uint32(0)) & (n < num_tokens)
+            if not routed:
+                byte = Uint32(0)
+                sf = Uint32(0)
+            if lane < 8:
+                ptx.sts_u8(hst + n * 32 + warp_idx * 8 + lane, byte)
+            if lane == 0:
+                ptx.sts_u8(hst + HSTAGE_SF + n * 4 + warp_idx, sf)
+
+
+class RoutedMoeFinalizeKernel:
+    """Native finalize of the raw tuple: ``out[t] = bf16(sum_k fma(w[t, k],
+    rows[t * 8 + k]))`` in top-k order with FP32 accumulation; each thread owns
+    8 hidden columns of one token. Launched with programmatic dependent launch
+    right after the routed kernel.
+
+    ``delay_ns`` (stress testing only, 0 compiles it out): a pseudo-random
+    third of the CTAs, re-drawn every launch from the global timer, spins 1..6
+    units before storing, so a consumer that read the output before this
+    grid completed would see stale rows.
+    """
+
+    THREADS = 128
+
+    def __init__(self, delay_ns: int = 0):
+        self.delay_ns = delay_ns
+
+    @cute.jit
+    def run(self, rows_addr: Int64, scale_addr: Int64, out_addr: Int64, num_tokens: Int32, stream):
+        self.kernel(rows_addr, scale_addr, out_addr).launch(
+            grid=[HIDDEN // (8 * self.THREADS), num_tokens, 1],
+            block=[self.THREADS, 1, 1],
+            stream=stream,
+            use_pdl=True,
+        )
+
+    @cute.kernel
+    def kernel(self, rows_addr: Int64, scale_addr: Int64, out_addr: Int64):
+        tidx, _, _ = cute.arch.thread_idx()
+        bx, t, _ = cute.arch.block_idx()
+        col = (Int32(bx) * self.THREADS + Int32(tidx)) * 8
+        t = Int32(t)
+        cute.arch.griddepcontrol_wait()
+        acc = cute.make_rmem_tensor((8,), Float32)
+        for i in cutlass.range_constexpr(8):
+            acc[i] = Float32(0.0)
+        # written by the routed kernel: plain (coherent) loads after the wait
+        ws = ptx.ld_v4_u32(scale_addr + t * (TOP_K * 2))
+        for k in cutlass.range_constexpr(TOP_K):
+            w = ptx.bf16x2_to_f32(ws[k // 2], k % 2 == 1)
+            y = ptx.ld_v4_u32(rows_addr + ((t * TOP_K + k) * HIDDEN + col) * 2)
+            for i in cutlass.range_constexpr(4):
+                acc[2 * i] = ptx.fma_rn(w, ptx.bf16x2_to_f32(y[i], False), acc[2 * i])
+                acc[2 * i + 1] = ptx.fma_rn(w, ptx.bf16x2_to_f32(y[i], True), acc[2 * i + 1])
+        if cutlass.const_expr(self.delay_ns > 0):
+            t0 = ptx.globaltimer_u64()
+            seed = Uint32(Int32(t0 >> Int64(12)) & Int32(0xFFFF))
+            h = (
+                Uint32(bx + t * 8) * Uint32(37) + seed * Uint32(101) + Uint32(DELAY_FINALIZE)
+            ) % Uint32(17)
+            if h < Uint32(6):
+                d = Int64(h + Uint32(1)) * Int64(self.delay_ns)
+                tn = t0
+                while tn - t0 < d:
+                    tn = ptx.globaltimer_u64()
+        ptx.st_v4_u32(
+            out_addr + (t * HIDDEN + col) * 2,
+            ptx.pack_bf16x2(acc[0], acc[1]),
+            ptx.pack_bf16x2(acc[2], acc[3]),
+            ptx.pack_bf16x2(acc[4], acc[5]),
+            ptx.pack_bf16x2(acc[6], acc[7]),
+        )
+        cute.arch.griddepcontrol_launch_dependents()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/ptx.py
new file mode 100644
index 000000000000..556b138389c9
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/ptx.py
@@ -0,0 +1,844 @@
+# 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.
+"""Inline-PTX primitives used by the routed-MoE kernel."""
+
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import Float32, Int32, Int64, T, Uint32, dsl_user_op
+
+
+def _asm(res, operands, text, constraints, *, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        operands,
+        text,
+        constraints,
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def f32_bits(x: Float32, *, loc=None, ip=None) -> Uint32:
+    return Uint32(llvm.bitcast(T.i32(), Float32(x).ir_value(loc=loc, ip=ip), loc=loc, ip=ip))
+
+
+@dsl_user_op
+def bits_f32(x: Uint32, *, loc=None, ip=None) -> Float32:
+    return Float32(llvm.bitcast(T.f32(), Uint32(x).ir_value(loc=loc, ip=ip), loc=loc, ip=ip))
+
+
+@dsl_user_op
+def redux_max_u32(x: Uint32, *, loc=None, ip=None) -> Uint32:
+    """Warp-wide unsigned max (all 32 lanes must participate)."""
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Uint32(x).ir_value(loc=loc, ip=ip)],
+            "redux.sync.max.u32 $0, $1, 0xffffffff;",
+            "=r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def cvt_e4m3_rn_sat(x: Float32, *, loc=None, ip=None) -> Uint32:
+    """FP32 -> E4M3 (RN, satfinite), returned in the low byte."""
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Float32(x).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .b16 h;\n.reg .f32 z;\n"
+            "mov.f32 z, 0f00000000;\n"
+            "cvt.rn.satfinite.e4m3x2.f32 h, z, $1;\n"
+            "cvt.u32.u16 $0, h;\n}",
+            "=r,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def e4m3_to_f32(byte: Uint32, *, loc=None, ip=None) -> Float32:
+    """E4M3 in the low byte -> FP32 (exact)."""
+    return Float32(
+        _asm(
+            T.f32(),
+            [Uint32(byte).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .b16 lo;\n.reg .b32 hx2;\n.reg .b16 h0, h1;\n"
+            "cvt.u16.u32 lo, $1;\n"
+            "cvt.rn.f16x2.e4m3x2 hx2, lo;\n"
+            "mov.b32 {h0, h1}, hx2;\n"
+            "cvt.f32.f16 $0, h0;\n}",
+            "=f,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def cvt_e2m1x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Uint32:
+    """Two FP32 -> one byte of E2M1 (RN, satfinite); ``lo`` in bits 0..3."""
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Float32(hi).ir_value(loc=loc, ip=ip), Float32(lo).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .b8 b;\n.reg .b32 t;\n"
+            "cvt.rn.satfinite.e2m1x2.f32 b, $1, $2;\n"
+            "mov.b32 t, {b, b, b, b};\n"
+            "and.b32 $0, t, 255;\n}",
+            "=r,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def rcp_approx_ftz(x: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [Float32(x).ir_value(loc=loc, ip=ip)],
+            "rcp.approx.ftz.f32 $0, $1;",
+            "=f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def ex2_approx(x: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [Float32(x).ir_value(loc=loc, ip=ip)],
+            "ex2.approx.ftz.f32 $0, $1;",
+            "=f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def bf16_round(x: Float32, *, loc=None, ip=None) -> Float32:
+    """Round FP32 to BF16 (RN) and back to FP32."""
+    return Float32(
+        _asm(
+            T.f32(),
+            [Float32(x).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\ncvt.f32.bf16 $0, h;\n}",
+            "=f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def bf16_bits(x: Float32, *, loc=None, ip=None) -> Uint32:
+    """FP32 -> BF16 (RN) bits in the low half."""
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Float32(x).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\ncvt.u32.u16 $0, h;\n}",
+            "=r,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def bf16x2_to_f32(packed: Uint32, upper: bool, *, loc=None, ip=None) -> Float32:
+    """One BF16 lane of a packed pair -> FP32 (exact)."""
+    text = (
+        "{\n.reg .b32 t;\nand.b32 t, $1, 0xFFFF0000;\nmov.b32 $0, t;\n}"
+        if upper
+        else "{\n.reg .b32 t;\nshl.b32 t, $1, 16;\nmov.b32 $0, t;\n}"
+    )
+    return Float32(
+        _asm(T.f32(), [Uint32(packed).ir_value(loc=loc, ip=ip)], text, "=f,r", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def tmem_ld_32x32b_x8(taddr: Int32, *, loc=None, ip=None):
+    """tcgen05.ld 32x32b.x8: this thread's lane, 8 consecutive columns."""
+    res = _asm(
+        llvm.StructType.get_literal([T.f32()] * 8),
+        [Int32(taddr).ir_value(loc=loc, ip=ip)],
+        "tcgen05.ld.sync.aligned.32x32b.x8.b32 "
+        "{$0, $1, $2, $3, $4, $5, $6, $7}, [$8];\n"
+        "tcgen05.wait::ld.sync.aligned;",
+        "=f,=f,=f,=f,=f,=f,=f,=f,r",
+        loc=loc,
+        ip=ip,
+    )
+    return [Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(8)]
+
+
+@dsl_user_op
+def st_relaxed_gpu_u32(addr: Int64, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "st.relaxed.gpu.global.u32 [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def ld_relaxed_gpu_u32(addr: Int64, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Int64(addr).ir_value(loc=loc, ip=ip)],
+            "ld.relaxed.gpu.global.u32 $0, [$1];",
+            "=r,l",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def st_v4_u32(addr: Int64, a: Uint32, b: Uint32, c: Uint32, d: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [
+            Int64(addr).ir_value(loc=loc, ip=ip),
+            Uint32(a).ir_value(loc=loc, ip=ip),
+            Uint32(b).ir_value(loc=loc, ip=ip),
+            Uint32(c).ir_value(loc=loc, ip=ip),
+            Uint32(d).ir_value(loc=loc, ip=ip),
+        ],
+        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_u32(addr: Int64, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "st.global.u32 [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_u16(addr: Int64, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.global.u16 [$0], h;\n}",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_add_f32(addr: Int64, value: Float32, *, loc=None, ip=None):
+    """Relaxed GPU-scope FP32 atomic add without a return value."""
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Float32(value).ir_value(loc=loc, ip=ip)],
+        "red.relaxed.gpu.global.add.f32 [$0], $1;",
+        "l,f",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_f32(addr: Int64, value: Float32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Float32(value).ir_value(loc=loc, ip=ip)],
+        "st.global.f32 [$0], $1;",
+        "l,f",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def sts_u32(smem_addr: Int32, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(smem_addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "st.shared.u32 [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def sts_u8(smem_addr: Int32, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(smem_addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.shared.u8 [$0], h;\n}",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def lds_u32(smem_addr: Int32, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
+            "ld.shared.u32 $0, [$1];",
+            "=r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def atom_or_shared_u32(smem_addr: Int32, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(smem_addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "red.shared.or.b32 [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def popc_u32(x: Uint32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [Uint32(x).ir_value(loc=loc, ip=ip)],
+            "popc.b32 $0, $1;",
+            "=r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def ld_nc_f32(addr: Int64, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [Int64(addr).ir_value(loc=loc, ip=ip)],
+            "ld.global.nc.f32 $0, [$1];",
+            "=f,l",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def ld_nc_u16(addr: Int64, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Int64(addr).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .b16 h;\nld.global.nc.u16 h, [$1];\ncvt.u32.u16 $0, h;\n}",
+            "=r,l",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def lds_v4_u32(smem_addr: Int32, *, loc=None, ip=None):
+    res = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
+        "ld.shared.v4.u32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return [Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def fence_acq_rel_gpu(*, loc=None, ip=None):
+    _asm(None, [], "fence.acq_rel.gpu;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def fence_proxy_async_shared(*, loc=None, ip=None):
+    _asm(None, [], "fence.proxy.async.shared::cta;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def bar_sync(bar_id: int, num_threads: int, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(bar_id).ir_value(loc=loc, ip=ip), Int32(num_threads).ir_value(loc=loc, ip=ip)],
+        "bar.sync $0, $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def globaltimer_u64(*, loc=None, ip=None) -> Int64:
+    return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def clock64(*, loc=None, ip=None) -> Int64:
+    return Int64(_asm(T.i64(), [], "mov.u64 $0, %clock64;", "=l", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def st_u64(addr: Int64, value: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Int64(value).ir_value(loc=loc, ip=ip)],
+        "st.global.u64 [$0], $1;",
+        "l,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def bulk_g2s(smem_dst: Int32, gmem_src: Int64, nbytes: Int32, mbar: Int32, *, loc=None, ip=None):
+    """cp.async.bulk global -> shared, completing on ``mbar`` (tx bytes)."""
+    _asm(
+        None,
+        [
+            Int32(smem_dst).ir_value(loc=loc, ip=ip),
+            Int64(gmem_src).ir_value(loc=loc, ip=ip),
+            Int32(nbytes).ir_value(loc=loc, ip=ip),
+            Int32(mbar).ir_value(loc=loc, ip=ip),
+        ],
+        "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];",
+        "r,l,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mbar_arrive_expect_tx(mbar: Int32, nbytes: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(mbar).ir_value(loc=loc, ip=ip), Int32(nbytes).ir_value(loc=loc, ip=ip)],
+        "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def sts_u64(smem_addr: Int32, value: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(smem_addr).ir_value(loc=loc, ip=ip), Int64(value).ir_value(loc=loc, ip=ip)],
+        "st.shared.u64 [$0], $1;",
+        "r,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def lds_u64(smem_addr: Int32, *, loc=None, ip=None) -> Int64:
+    return Int64(
+        _asm(
+            T.i64(),
+            [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
+            "ld.shared.u64 $0, [$1];",
+            "=l,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def cvt_e2m1x8(
+    v0: Float32,
+    v1: Float32,
+    v2: Float32,
+    v3: Float32,
+    v4: Float32,
+    v5: Float32,
+    v6: Float32,
+    v7: Float32,
+    *,
+    loc=None,
+    ip=None,
+) -> Uint32:
+    """Eight FP32 -> eight E2M1 (RN, satfinite); ``v0`` in bits 0..3."""
+    ops = [Float32(v).ir_value(loc=loc, ip=ip) for v in (v1, v0, v3, v2, v5, v4, v7, v6)]
+    return Uint32(
+        _asm(
+            T.i32(),
+            ops,
+            "{\n.reg .b8 b0, b1, b2, b3;\n"
+            "cvt.rn.satfinite.e2m1x2.f32 b0, $1, $2;\n"
+            "cvt.rn.satfinite.e2m1x2.f32 b1, $3, $4;\n"
+            "cvt.rn.satfinite.e2m1x2.f32 b2, $5, $6;\n"
+            "cvt.rn.satfinite.e2m1x2.f32 b3, $7, $8;\n"
+            "mov.b32 $0, {b0, b1, b2, b3};\n}",
+            "=r,f,f,f,f,f,f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def st_ll2(addr: Int64, d0: Uint32, d1: Uint32, flag: Uint32, *, loc=None, ip=None):
+    """Two LL words ``{d0, flag}, {d1, flag}`` in one 16-byte store.
+
+    LL protocol (as NCCL's LL): every 4-byte payload travels with the epoch
+    of the launch that wrote it, so a reader that sees the expected epoch in
+    an 8-byte half also sees that half's payload; no fence or flag is needed.
+    GPU-scope relaxed (not volatile, which is system scope) so the stores
+    and polls stay within the GPU's L2.
+    """
+    _asm(
+        None,
+        [
+            Int64(addr).ir_value(loc=loc, ip=ip),
+            Uint32(d0).ir_value(loc=loc, ip=ip),
+            Uint32(flag).ir_value(loc=loc, ip=ip),
+            Uint32(d1).ir_value(loc=loc, ip=ip),
+        ],
+        "st.relaxed.gpu.global.v4.u32 [$0], {$1, $2, $3, $2};",
+        "l,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def ld_ll2(addr: Int64, *, loc=None, ip=None):
+    """16-byte poll of two LL words: ``[d0, flag0, d1, flag1]``."""
+    res = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [Int64(addr).ir_value(loc=loc, ip=ip)],
+        "ld.relaxed.gpu.global.v4.u32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,l",
+        loc=loc,
+        ip=ip,
+    )
+    return [Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def prefetch_l2_bulk(addr: Int64, nbytes: Int32, *, loc=None, ip=None):
+    """Bulk L2 prefetch of ``nbytes`` (multiple of 16) contiguous bytes."""
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Int32(nbytes).ir_value(loc=loc, ip=ip)],
+        "cp.async.bulk.prefetch.L2.global [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def fence_proxy_async_global(*, loc=None, ip=None):
+    """Order generic-proxy global accesses with async-proxy (bulk copy)
+    accesses of this thread."""
+    _asm(None, [], "fence.proxy.async.global;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def ld_acquire_gpu_u32(addr: Int64, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Int64(addr).ir_value(loc=loc, ip=ip)],
+            "ld.acquire.gpu.global.u32 $0, [$1];",
+            "=r,l",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def st_ll1(addr: Int64, d0: Uint32, flag: Uint32, *, loc=None, ip=None):
+    """One LL word ``{d0, flag}`` (8-byte GPU-scope relaxed store)."""
+    _asm(
+        None,
+        [
+            Int64(addr).ir_value(loc=loc, ip=ip),
+            Uint32(d0).ir_value(loc=loc, ip=ip),
+            Uint32(flag).ir_value(loc=loc, ip=ip),
+        ],
+        "st.relaxed.gpu.global.v2.u32 [$0], {$1, $2};",
+        "l,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def lds_v2_u32(smem_addr: Int32, *, loc=None, ip=None):
+    res = _asm(
+        llvm.StructType.get_literal([T.i32()] * 2),
+        [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
+        "ld.shared.v2.u32 {$0, $1}, [$2];",
+        "=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return [Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(2)]
+
+
+@dsl_user_op
+def sts_v2_u32(smem_addr: Int32, a: Uint32, b: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [
+            Int32(smem_addr).ir_value(loc=loc, ip=ip),
+            Uint32(a).ir_value(loc=loc, ip=ip),
+            Uint32(b).ir_value(loc=loc, ip=ip),
+        ],
+        "st.shared.v2.u32 [$0], {$1, $2};",
+        "r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Thread block cluster primitives (expert-cluster schedule)
+# ---------------------------------------------------------------------------
+@dsl_user_op
+def cluster_ctarank(*, loc=None, ip=None) -> Int32:
+    return Int32(_asm(T.i32(), [], "mov.u32 $0, %cluster_ctarank;", "=r", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def mapa_u32(smem_addr: Int32, cta_rank: Int32, *, loc=None, ip=None) -> Int32:
+    """Shared::cluster address of ``smem_addr`` in cluster CTA ``cta_rank``."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [Int32(smem_addr).ir_value(loc=loc, ip=ip), Int32(cta_rank).ir_value(loc=loc, ip=ip)],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def st_cluster_v4_u32(
+    cluster_addr: Int32, a: Uint32, b: Uint32, c: Uint32, d: Uint32, *, loc=None, ip=None
+):
+    _asm(
+        None,
+        [Int32(cluster_addr).ir_value(loc=loc, ip=ip)]
+        + [Uint32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c, d)],
+        "st.shared::cluster.v4.u32 [$0], {$1, $2, $3, $4};",
+        "r,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_cluster_u32(cluster_addr: Int32, value: Uint32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(cluster_addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "st.shared::cluster.u32 [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def fence_proxy_async_shared_cluster(*, loc=None, ip=None):
+    """Order this thread's generic-proxy shared::cluster stores before later
+    async-proxy (UMMA, tcgen05.cp) reads of the same bytes."""
+    _asm(None, [], "fence.proxy.async.shared::cluster;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def mbar_arrive_remote(cluster_addr: Int32, *, loc=None, ip=None):
+    """Release-arrive (cluster scope) on an mbarrier of a cluster CTA."""
+    _asm(
+        None,
+        [Int32(cluster_addr).ir_value(loc=loc, ip=ip)],
+        "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [$0];",
+        "r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mbar_try_wait_cluster(mbar: Int32, phase: Int32, *, loc=None, ip=None) -> Int32:
+    """One acquire (cluster scope) try-wait on a local mbarrier phase; 1 when done."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [Int32(mbar).ir_value(loc=loc, ip=ip), Int32(phase).ir_value(loc=loc, ip=ip)],
+            "{\n.reg .pred p;\n"
+            "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n"
+            "selp.u32 $0, 1, 0, p;\n}",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def ld_nc_v4_u32(addr: Int64, *, loc=None, ip=None):
+    res = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [Int64(addr).ir_value(loc=loc, ip=ip)],
+        "ld.global.nc.v4.u32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,l",
+        loc=loc,
+        ip=ip,
+    )
+    return [Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def ld_v4_u32(addr: Int64, *, loc=None, ip=None):
+    """Plain (coherent) 16-byte global load."""
+    res = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [Int64(addr).ir_value(loc=loc, ip=ip)],
+        "ld.global.v4.u32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,l",
+        loc=loc,
+        ip=ip,
+    )
+    return [Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def fma_rn(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [Float32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c)],
+            "fma.rn.f32 $0, $1, $2, $3;",
+            "=f,f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def pack_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Uint32:
+    """Two FP32 values -> one word of two BF16 (RN), ``lo`` in the low half."""
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Float32(lo).ir_value(loc=loc, ip=ip), Float32(hi).ir_value(loc=loc, ip=ip)],
+            "cvt.rn.bf16x2.f32 $0, $2, $1;",
+            "=r,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def atom_add_relaxed_gpu_u32(addr: Int64, value: Uint32, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+            "atom.relaxed.gpu.global.add.u32 $0, [$1], $2;",
+            "=r,l,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def smid_u32(*, loc=None, ip=None) -> Int32:
+    return Int32(_asm(T.i32(), [], "mov.u32 $0, %smid;", "=r", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def bar_arrive(bar_id: int, num_threads: int, *, loc=None, ip=None):
+    _asm(
+        None,
+        [Int32(bar_id).ir_value(loc=loc, ip=ip), Int32(num_threads).ir_value(loc=loc, ip=ip)],
+        "bar.arrive $0, $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def bulk_s2cluster(
+    dst_cluster: Int32, src_cta: Int32, nbytes: Int32, mbar_cluster: Int32, *, loc=None, ip=None
+):
+    """cp.async.bulk shared::cta -> shared::cluster (a peer CTA), completing
+    on the peer's mbarrier (tx bytes)."""
+    _asm(
+        None,
+        [
+            Int32(dst_cluster).ir_value(loc=loc, ip=ip),
+            Int32(src_cta).ir_value(loc=loc, ip=ip),
+            Int32(nbytes).ir_value(loc=loc, ip=ip),
+            Int32(mbar_cluster).ir_value(loc=loc, ip=ip),
+        ],
+        "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];",
+        "r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe.py
new file mode 100644
index 000000000000..c1a8096cf8ed
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe.py
@@ -0,0 +1,880 @@
+# 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.
+"""Host side of the routed_moe routed NVFP4 MoE (the routed_moe contract).
+
+``RoutedMoe`` wraps one stock TRT-LLM-Gen NVFP4 MoE leaf (one rank's TP
+shard) and runs the single-launch CuTe DSL kernel on its weight buffers in
+place: no packed copies, no layout conversion. The call site may pass the
+served ``ConfigurableMoE`` wrapper or the bare leaf; the leaf that owns the
+expert parameters is ``wrapper.backend``. The schedule is the
+expert-cluster kernel (``kernel_xc.py``, raw tuple, plus a native-order
+finalize kernel for finalized output); a call whose rows need more
+co-resident clusters than the grid holds runs the grid-wide kernel
+(``kernel.py``, ``RoutedMoe.schedule``). ``eligibility`` is decided from
+host-side state (types, shapes, dtypes, routing configuration and the device)
+before any device work, so an ineligible call site keeps the native path with
+no side effect.
+
+Outputs follow the routed_moe contract: the finalized routed output ``[R, 6144]`` BF16, or
+the native unfinalized tuple ``(fc2_output [R*8, 6144] BF16, expert scales
+[R, 8] BF16, expanded_idx_to_permuted_idx [R, 8] int32)`` whose permutation is
+the identity ``t*8 + k`` (every row valid; a stock finalizer sums
+``scale[t, k] * fc2_output[idx[t, k]]``).
+
+Physical layout of the raw tuple. ``fc2_output`` is token-major: row
+``t * 8 + k`` holds the FC2 output of token ``t``'s ``k``-th expert, with a row
+pitch of 6144 elements, and the permutation maps ``(t, k)`` to that row. The
+native runner groups its rows by expert and pads them; both layouts are read
+through the permutation, so the native unfinalized consumers (the single-node
+``MoEAllReduce`` finalize + all-reduce and the post_moe_hop) take either
+one. The expert scales carry the native dtype, values and top-k order.
+
+Programmatic dependent launch (PDL). Both kernels are launched as PDL
+secondaries of whatever kernel precedes them on the stream (the router GEMM,
+router_gemm's router or a pre-MoE hop):
+
+* before ``griddepcontrol.wait`` they only set up shared memory, barriers and
+  tensor memory and read the router bias, which no kernel writes during
+  serving (it is a model parameter); every input produced on the stream
+  (activation, router logits) and every output store comes after the wait.
+  With package cross_layer_expert_prefetch's speculative prefetch, the expert-cluster kernel also
+  reads the predicted expert IDs there and issues L2 prefetches of those
+  experts' weights; the IDs are written on another stream (or by the
+  previous call of the same layer) and only steer the prefetch, so any value
+  read is harmless;
+* the expert-cluster kernel issues ``griddepcontrol.launch_dependents`` after
+  each CTA's last FC2 tile; that only lets the successor's CTAs be scheduled
+  earlier, and the successor (the finalize kernel or the post-MoE hop) reads
+  the routed output only after its own ``griddepcontrol.wait``, which returns
+  when this grid has completed and its memory operations are visible;
+* the finalize kernel reads the raw rows after its ``griddepcontrol.wait``
+  and triggers its dependents after its stores;
+* outputs written into a consumer's buffer (routed_moe_into_hop) are not reused by this
+  package; the consumer owns their lifetime and clears the FP32 accumulator
+  after reading it.
+
+Work on other streams (the shared expert on the aux stream) joins through
+events as in the native model, not through PDL.
+
+Numerics. The expert-cluster schedule computes every token's FC1, SwiGLU,
+intermediate quantization and FC2 independently of the other rows of the
+call (no split-K), and the finalize kernel sums each token's top-k in order,
+so the raw rows and the finalized output of a live row do not depend on the
+content of padded rows (bit for bit). The FP32 accumulator output (routed_moe_into_hop
+"fp32") adds the eight weighted rows of a token with atomic FP32 adds whose
+order depends on the schedule, so it is run-to-run nondeterministic within
+the FP32 reorder bound. The grid-wide ``v6`` schedule splits FC1 over the
+distinct experts of the whole call and is not padding-invariant.
+"""
+
+import contextlib
+import threading
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+HIDDEN = 6144
+INTER = 256
+NUM_EXPERTS = 256
+TOP_K = 8
+MAX_ROWS = 8
+ROUTED_SCALING = 2.5
+
+_NUM_ADDRS = 19
+# Expert-cluster schedule defaults for the served neighbours (attention
+# weight stream, pre-MoE hop, the BF16 shared expert on the aux stream, post_moe_hop /
+# MNNVL all-reduce hop). The shared expert runs concurrently with the routed
+# kernel and needs SMs and DRAM latency headroom; how much slack it has
+# depends on the output mode:
+# * finalized: the finalize kernel, the add and the all-reduce follow on the
+#   main stream, so at most 32 four-CTA clusters work;
+# * raw tuple (into the post_moe_hop, which joins the shared branch right after the
+#   routed kernel) with the router logits supplied by the pre-MoE hop (post_attention_hop or
+#   router_gemm selected: the MoE's router GEMM is skipped, so the shared expert's
+#   first GEMM starts together with the routed kernel instead of during the
+#   router GEMM): 26 clusters, raised to cover rows x top-8 in two rounds for
+#   R > 6;
+# * the same with routed_moe_into_hop's FP32 accumulator output: the routed kernel writes R
+#   rows instead of R x 8 and finishes earlier, so the shared expert is the
+#   critical path at 26 clusters; 24 clusters (raised the same way for
+#   R > 6);
+# * raw tuple after the MoE's own router GEMM: 32 clusters.
+# For D <= clusters distinct experts only clamp(D, 24, clusters) clusters
+# work; the others exit after routing and leave their SMs to the shared
+# expert.
+XC_MAX_CLUSTERS = 32
+XC_RAW_CLUSTERS = 26
+XC_RED_CLUSTERS = 24
+XC_MIN_CLUSTERS = 24
+# selector IDs of the pre-MoE hops that emit the router logits (the router contract)
+_LOGIT_HOPS = ("post_attention_hop", "router_gemm")
+
+
+#: Stress-test delay injection of the routed kernels: ``(site bits, unit in ns)``, the sites being
+#: ``kernel.DELAY_*`` / ``kernel_xc.DELAY_*``. Serving delays no site; tests set it to hold CTAs at
+#: those sites. It is part of every runner's compile key.
+DELAY_INJECTION = (0, 2000)
+
+
+def xc_kernel_args() -> dict:
+    """Launch parameters of the expert-cluster kernel (``kernel_xc``) and of its forks.
+
+    Ring depth 8, TMEM accumulator stages 4, PDL trigger ``late``, at least
+    ``XC_MIN_CLUSTERS`` working clusters, no expert-cluster L2 prefetch
+    (``l2pf`` 0), every cluster working in a two-round call
+    (``two_round_clusters`` 0) and ``DELAY_INJECTION``.
+    """
+    delay, delay_ns = DELAY_INJECTION
+    return dict(
+        num_stages=8,
+        delay=delay,
+        delay_ns=delay_ns,
+        trigger="late",
+        min_clusters=XC_MIN_CLUSTERS,
+        l2pf=0,
+        two_round_clusters=0,
+        acc_stages=4,
+    )
+
+
+def _logits_from_hop() -> bool:
+    """Whether a pre-MoE hop supplies the router logits (rank-identical)."""
+    from ....shared import dispatch
+
+    return any(dispatch.is_enabled(f) for f in _LOGIT_HOPS)
+
+
+_NUM_XC_ADDRS = 18
+
+_LOCK = threading.Lock()
+_COMPILED: Dict[tuple, object] = {}
+_WORKSPACES: Dict[int, "_Workspace"] = {}
+
+
+def _sm_count(device: torch.device) -> int:
+    return torch.cuda.get_device_properties(device).multi_processor_count
+
+
+_MAX_CLUSTERS: Dict[tuple, int] = {}
+
+
+def _max_active_clusters(device: torch.device, cluster: int) -> int:
+    """Clusters of ``cluster`` CTAs that are co-resident at one CTA per SM
+    (a cluster must fit in one GPC, so this can be below sm_count / cluster)."""
+    idx = device.index if device.index is not None else torch.cuda.current_device()
+    key = (idx, cluster)
+    if key not in _MAX_CLUSTERS:
+        import cutlass.utils
+
+        _MAX_CLUSTERS[key] = int(
+            cutlass.utils.HardwareInfo(device_id=idx).get_max_active_clusters(cluster)
+        )
+    return _MAX_CLUSTERS[key]
+
+
+class _Workspace:
+    """Graph-stable per-device scratch shared by every routed_moe call of the process.
+
+    The scratch is private to routed_moe. Calls that use it must be ordered on the
+    device: one workspace serves all layers, both output modes and every
+    captured CUDA graph. ``ws`` holds one launch counter per CTA; flags and LL
+    words carry the epoch of the launch that wrote them, so nothing needs
+    clearing between calls or graph replays.
+
+    Ordering is enforced for eager calls: a call issued from another stream
+    than the previous eager call first makes its stream wait for that
+    stream's submitted work (``order_on``). Captured calls are ordered by the
+    stream the graph is replayed on, which must be the stream of the eager
+    calls (or synchronized with it), as in the model executor.
+
+    It is created once per device, outside any dispatch mode (a model built
+    under the loader's meta-device mode still gets device memory), and never
+    replaced, because captured graphs keep its addresses.
+    """
+
+    def __init__(self, device: torch.device, num_ctas: int):
+        from .kernel import FC1_MT, HBUF_BYTES, LL_PART_BYTES, MAX_SLOTS, WS_WORDS
+
+        self.num_ctas = num_ctas
+        self.ws = torch.zeros(WS_WORDS, dtype=torch.int32, device=device)
+        # FC1 and FC2 split partials, one LL slot per CTA each
+        self.partials = torch.zeros(2 * num_ctas * LL_PART_BYTES, dtype=torch.uint8, device=device)
+        # per-expert FC2 operand hand-off buffers and per-tile epoch flags
+        self.himg = torch.zeros(MAX_SLOTS * HBUF_BYTES, dtype=torch.uint8, device=device)
+        self.hflag = torch.zeros(MAX_SLOTS * FC1_MT, dtype=torch.int32, device=device)
+        self.dummy = torch.zeros(64, dtype=torch.int32, device=device)
+        # expert-cluster schedule, finalized mode: raw rows and routing
+        # weights consumed by the finalize kernel
+        self.rows = torch.empty(MAX_ROWS * TOP_K, HIDDEN, dtype=torch.bfloat16, device=device)
+        self.scales = torch.empty(MAX_ROWS, TOP_K, dtype=torch.bfloat16, device=device)
+        # per-output-tile arrival counters of the in-kernel finalize; the
+        # closing CTA resets its tile's counter
+        self.tile_count = torch.zeros(64, dtype=torch.int32, device=device)
+        # stream of the last eager call (see order_on)
+        self.stream: Optional[torch.cuda.Stream] = None
+
+    def tensors(self) -> Dict[str, torch.Tensor]:
+        return {name: t for name, t in vars(self).items() if isinstance(t, torch.Tensor)}
+
+    def nbytes(self) -> int:
+        return sum(t.numel() * t.element_size() for t in self.tensors().values())
+
+    def order_on(self, stream: torch.cuda.Stream) -> None:
+        """Order an eager call on ``stream`` after the previous eager call.
+
+        A no-op while ``stream`` is being captured: graph nodes run when the
+        graph is replayed, in the order of the replay stream.
+        """
+        if torch.cuda.is_current_stream_capturing():
+            return
+        prev = self.stream
+        if prev is not None and prev != stream:
+            stream.wait_stream(prev)
+        self.stream = stream
+
+
+def _workspace(device: torch.device, num_ctas: int) -> _Workspace:
+    """The device's workspace, sized for every SM; created on first use."""
+    key = device.index if device.index is not None else torch.cuda.current_device()
+    ws = _WORKSPACES.get(key)
+    if ws is None:
+        try:
+            from torch.utils._python_dispatch import _disable_current_modes
+        except ImportError:  # the buffers are still checked below
+            _disable_current_modes = contextlib.nullcontext
+        dev = torch.device("cuda", key)
+        with _disable_current_modes():
+            ws = _Workspace(dev, max(num_ctas, _sm_count(dev)))
+        bad = [n for n, t in ws.tensors().items() if t.is_meta or t.device != dev]
+        if bad:
+            raise RuntimeError(f"routed_moe workspace buffers {bad} are not device memory on {dev}")
+        _WORKSPACES[key] = ws
+    elif ws.num_ctas < num_ctas:
+        raise RuntimeError(
+            f"routed_moe workspace of cuda:{key} holds {ws.num_ctas} CTA slots, a launch needs "
+            f"{num_ctas}; it cannot be replaced once CUDA graphs may reference it"
+        )
+    return ws
+
+
+def resolve_expert_owner(module) -> Tuple[Optional[torch.nn.Module], str]:
+    """The leaf that owns the NVFP4 expert parameters of ``module``.
+
+    The served model passes a ``ConfigurableMoE`` wrapper, whose expert
+    tensors live on ``wrapper.backend`` (the wrapper only has delegating
+    properties); a bare TRT-LLM-Gen leaf owns them itself. Ownership is
+    decided by registered parameters, not by attribute probes.
+    """
+    backend = getattr(module, "backend", None)
+    owner = backend if isinstance(backend, torch.nn.Module) else module
+    if not isinstance(owner, torch.nn.Module):
+        return None, f"{type(module).__name__} is not a module"
+    if "w3_w1_weight" not in owner._parameters or "w2_weight" not in owner._parameters:
+        return None, f"{type(owner).__name__} owns no NVFP4 expert parameters"
+    return owner, ""
+
+
+def _param(owner, name):
+    return owner._parameters.get(name) if isinstance(owner, torch.nn.Module) else None
+
+
+def _check_vector(t, numel: int) -> bool:
+    return t is not None and t.dtype == torch.float32 and t.numel() == numel and t.is_contiguous()
+
+
+def _wrapper_reason(wrapper) -> str:
+    """Wrapper-level behaviour the routed_moe call would skip; empty when none."""
+    if getattr(wrapper, "comm", None) is not None:
+        return "MoE communication strategy"
+    if getattr(wrapper, "layer_load_balancer", None) is not None:
+        return "EPLB load balancer"
+    if getattr(wrapper, "enable_dwdp", False):
+        return "DWDP"
+    if getattr(wrapper, "unfinalized_combine_fn", None) is not None:
+        return "unfinalized_combine_fn"
+    if getattr(wrapper, "apply_router_weight_on_input", False):
+        return "router weight on input"
+    if getattr(wrapper, "_enable_perfect_router", False):
+        return "perfect router"
+    if getattr(wrapper, "reduce_results", False) and getattr(wrapper, "parallel_size", 1) > 1:
+        return "in-op all-reduce"
+    return ""
+
+
+def _routing_reason(rm) -> str:
+    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
+
+    if not isinstance(rm, DeepSeekV3MoeRoutingMethod):
+        return f"routing method {type(rm).__name__}"
+    impl = rm.routing_impl
+    groups = (impl.top_k, impl.n_group, impl.topk_group)
+    if groups != (TOP_K, 1, 1) or float(impl.routed_scaling_factor) != ROUTED_SCALING:
+        return f"routing {groups} scale {impl.routed_scaling_factor}"
+    return ""
+
+
+def _owner_reason(owner, device: torch.device) -> str:
+    """Leaf-level eligibility: stock TRT-LLM-Gen NVFP4 TP8/EP1 shard."""
+    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.fp4_block_scale import TRTLLMGenNvfp4Base
+    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.identity import PROVIDER_TRTLLM
+    from tensorrt_llm._torch.utils import ActivationType
+
+    if not isinstance(owner, TRTLLMGenNvfp4Base):
+        return f"{type(owner).__name__} is not a TRT-LLM-Gen NVFP4 leaf"
+    if getattr(type(owner), "provider", None) != PROVIDER_TRTLLM:
+        return f"provider {getattr(type(owner), 'provider', None)}"
+    if getattr(owner, "use_dp", False):
+        return "attention DP"
+    routes_outside = getattr(owner, "_routes_outside_the_kernel", None)
+    if routes_outside is None or routes_outside():
+        return "separated routing"
+    if getattr(owner, "activation_type", None) != ActivationType.Swiglu:
+        return f"activation {getattr(owner, 'activation_type', None)}"
+    if getattr(owner, "bias", False):
+        return "expert bias"
+    for name in ("act_alpha", "act_beta", "act_clamp"):
+        if getattr(owner, name, None) is not None:
+            return f"{name} set"
+    if (
+        getattr(owner, "num_slots", None) != NUM_EXPERTS
+        or getattr(owner, "slot_start", None) != 0
+        or getattr(owner, "expert_size_per_partition", None) != NUM_EXPERTS
+    ):
+        return "not all 256 experts local (EP or EPLB slots)"
+    if getattr(owner, "hidden_size", None) != HIDDEN:
+        return f"hidden {getattr(owner, 'hidden_size', None)}"
+    quant_method = getattr(owner, "quant_method", None)
+    lean = getattr(quant_method, "intermediate_size_per_partition_lean", None)
+    if lean not in (None, INTER):
+        return f"lean intermediate {lean}"
+    w13 = _param(owner, "w3_w1_weight")
+    w2 = _param(owner, "w2_weight")
+    w13_sf = _param(owner, "w3_w1_weight_scale")
+    w2_sf = _param(owner, "w2_weight_scale")
+    if tuple(w13.shape) != (NUM_EXPERTS, 2 * INTER, HIDDEN // 2) or w13.dtype != torch.uint8:
+        return f"w3_w1 {tuple(w13.shape)} {w13.dtype}"
+    if tuple(w2.shape) != (NUM_EXPERTS, HIDDEN, INTER // 2) or w2.dtype != torch.uint8:
+        return f"w2 {tuple(w2.shape)} {w2.dtype}"
+    if (
+        w13_sf is None
+        or w2_sf is None
+        or tuple(w13_sf.shape) != (NUM_EXPERTS, 2 * INTER, HIDDEN // 16)
+        or tuple(w2_sf.shape) != (NUM_EXPERTS, HIDDEN, INTER // 16)
+    ):
+        return "block-scale shapes"
+    for t in (w13, w2, w13_sf, w2_sf):
+        if not t.is_contiguous() or t.device != device:
+            return "expert buffers not contiguous on the activation's device"
+    for name in ("fc31_alpha", "fc31_scale_c", "fc2_alpha"):
+        t = _param(owner, name)
+        if not _check_vector(t, NUM_EXPERTS) or t.device != device:
+            return f"{name} is not FP32 [256] on the device"
+    t = _param(owner, "fc31_input_scale")
+    if not _check_vector(t, 1) or t.device != device:
+        return "fc31_input_scale is not one FP32 scalar on the device"
+    return _routing_reason(getattr(owner, "routing_method", None))
+
+
+# Both schedules copy each activation row into shared memory with
+# ``cp.async.bulk``, whose global source must be 16-byte aligned.
+_BULK_ALIGN = 16
+_REALIGNED_LOGGED = []
+
+
+def _bulk_copy_ready(x: torch.Tensor) -> torch.Tensor:
+    """``x`` itself, or an aligned copy when its base is not 16-byte aligned.
+
+    The base address comes from this rank's allocator, so it never decides
+    between routed_moe and the native path (with routed_moe_into_hop the choice also selects the
+    post_moe_hop hop variant): a misaligned activation costs one device copy instead.
+    """
+    if x.data_ptr() % _BULK_ALIGN == 0:
+        return x
+    if not _REALIGNED_LOGGED:
+        _REALIGNED_LOGGED.append(True)
+        from tensorrt_llm.logger import logger
+
+        logger.warning(
+            "[glm_moe_dsa] routed_moe: activation base is not 16-byte aligned "
+            f"(data_ptr % 16 = {x.data_ptr() % _BULK_ALIGN}); copying it before every such call"
+        )
+    return x.clone(memory_format=torch.contiguous_format)
+
+
+def eligibility(module, x: torch.Tensor, router_logits: torch.Tensor) -> Tuple[bool, str]:
+    """Decide routed_moe support for this call from host-side state only.
+
+    ``module`` is the served ``ConfigurableMoE`` wrapper or a bare
+    TRT-LLM-Gen NVFP4 leaf.
+    """
+    if not torch.cuda.is_available() or x.device.type != "cuda":
+        return False, "not on CUDA"
+    major, _ = torch.cuda.get_device_capability(x.device)
+    if major != 10:
+        return False, f"SM{major}x is not the SM100 family"
+    if x.dtype != torch.bfloat16 or x.dim() != 2 or x.shape[1] != HIDDEN:
+        return False, f"activation {tuple(x.shape)} {x.dtype}"
+    rows = x.shape[0]
+    if rows < 1 or rows > MAX_ROWS:
+        return False, f"{rows} rows outside [1, {MAX_ROWS}]"
+    if not x.is_contiguous() or not router_logits.is_contiguous():
+        return False, "non-contiguous input"
+    if (
+        router_logits.dtype != torch.float32
+        or tuple(router_logits.shape) != (rows, NUM_EXPERTS)
+        or router_logits.device != x.device
+    ):
+        return False, f"logits {tuple(router_logits.shape)} {router_logits.dtype}"
+    if torch.compiler.is_compiling():
+        return False, "torch.compile tracing"
+    owner, reason = resolve_expert_owner(module)
+    if owner is None:
+        return False, reason
+    if owner is not module:
+        reason = _wrapper_reason(module) or _routing_reason(getattr(module, "routing_method", None))
+        if reason:
+            return False, reason
+    reason = _owner_reason(owner, x.device)
+    if reason:
+        return False, reason
+    return True, ""
+
+
+class RoutedMoe:
+    """Callable bound to one native NVFP4 MoE leaf and its router bias.
+
+    ``module`` may be the ``ConfigurableMoE`` wrapper or the leaf itself; the
+    kernel always reads the leaf's parameters.
+    """
+
+    def __init__(self, module, routing_bias: torch.Tensor):
+        owner, reason = resolve_expert_owner(module)
+        if owner is None:
+            raise ValueError(f"routed_moe needs an NVFP4 expert leaf: {reason}")
+        self.module = owner
+        self.bias = routing_bias
+        self.device = owner.w3_w1_weight.device
+        self.num_ctas = _sm_count(self.device)
+
+    @classmethod
+    def from_native(cls, module, routing_bias: torch.Tensor) -> "RoutedMoe":
+        return cls(module, routing_bias)
+
+    def _spec(self):
+        """Speculative prefetch of this layer's experts (cross_layer_expert_prefetch):
+        ``(ids, k, ctas, self_store)`` or ``None``."""
+        hook = getattr(self.module, "_expert_prefetch_spec", None)
+        return hook() if hook is not None else None
+
+    def _key(self, finalize: bool, debug: int, rows: int = MAX_ROWS, red: bool = False):
+        """``(num_ctas, finalize, debug, schedule, grid, red, prefetch spec, delay injection)`` of a
+        launch."""
+        sched = self.schedule(rows, finalize, red)
+        grid = self._xc_ctas(rows, finalize, red) if sched == "xc" else self.num_ctas
+        spec = self._spec() if sched == "xc" else None
+        spec_key = spec[1:] if spec is not None else None
+        return (self.num_ctas, finalize, debug, sched, grid, red, spec_key, DELAY_INJECTION)
+
+    def schedule(self, rows: int = MAX_ROWS, finalize: bool = True, red: bool = False) -> str:
+        """``xc`` (expert clusters) or ``v6`` (grid-wide hand-off).
+
+        ``xc`` needs enough co-resident clusters for two rounds to cover the
+        largest distinct-expert count of a ``rows``-row call (8 per row);
+        otherwise the grid-wide schedule runs.
+        """
+        from .kernel_xc import CLUSTER, MAX_ROUNDS
+
+        if (self._xc_ctas(rows, finalize, red) // CLUSTER) * MAX_ROUNDS < rows * TOP_K:
+            return "v6"
+        return "xc"
+
+    @staticmethod
+    def inkernel_finalize() -> bool:
+        """Whether ``xc`` finalizes inside the routed kernel: no, a second, PDL-launched kernel
+        finalizes."""
+        return False
+
+    def _xc_ctas(self, rows: int = MAX_ROWS, finalize: bool = True, red: bool = False) -> int:
+        """Grid of the expert-cluster schedule for a ``rows``-row call.
+
+        Every cluster must be co-resident, and in the served graph the BF16
+        shared expert runs concurrently on the aux stream, so the grid leaves
+        SMs free: ``XC_MAX_CLUSTERS`` clusters (32 on sm_103 = 128 of 152 SMs),
+        or, when a pre-MoE hop supplies the router logits, ``XC_RAW_CLUSTERS``
+        (26 = 104 SMs) for the raw tuple and ``XC_RED_CLUSTERS`` (24 = 96 SMs)
+        for the FP32 accumulator output (``red``), raised to
+        ``ceil(rows * 8 / 2)`` so that two rounds cover every distinct expert
+        of the call.
+        """
+        from .kernel_xc import CLUSTER, MAX_ROUNDS
+
+        co_resident = min(XC_MAX_CLUSTERS, _max_active_clusters(self.device, CLUSTER))
+        clusters = co_resident
+        if not finalize and _logits_from_hop():
+            need = -(-rows * TOP_K // MAX_ROUNDS)
+            base = XC_RED_CLUSTERS if red else XC_RAW_CLUSTERS
+            clusters = min(co_resident, max(base, need))
+        n = CLUSTER * clusters
+        n = min(n, self.num_ctas)
+        return n - n % CLUSTER
+
+    def ready_for_capture(
+        self, finalize: bool, debug: int = 0, rows: int = MAX_ROWS, red: bool = False
+    ) -> bool:
+        """True when a launch needs no compilation and no workspace
+        allocation, i.e. it can be recorded into a CUDA graph."""
+        key = self._key(finalize, debug, rows, red)
+        dev = self.device.index if self.device.index is not None else torch.cuda.current_device()
+        ws = _WORKSPACES.get(dev)
+        return key in _COMPILED and ws is not None and ws.num_ctas >= self.num_ctas
+
+    def _compiled(self, finalize: bool, debug: int, rows: int = MAX_ROWS, red: bool = False):
+        """Compiled launcher(s): ``(routed, finalize_or_None)``.
+
+        ``red``: the expert-cluster kernel adds the weighted FC2 outputs into
+        an FP32 accumulator instead of writing the raw rows (routed_moe_into_hop "fp32").
+        """
+        key = self._key(finalize, debug, rows, red)
+        with _LOCK:
+            fns = _COMPILED.get(key)
+            if fns is not None:
+                return fns
+            import cutlass
+            import cutlass.cute as cute
+            from cuda.bindings import driver as cuda
+            from cutlass.cute.runtime import make_ptr
+
+            ptrs = self._pointers(make_ptr, cute, cutlass)
+            stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+            if key[3] == "xc":
+                from .kernel_xc import RoutedMoeExpertClusterKernel, RoutedMoeFinalizeKernel
+
+                inkernel = finalize and self.inkernel_finalize()
+                spec = key[6] or (0, 0, False)
+                kern = RoutedMoeExpertClusterKernel(
+                    self._xc_ctas(rows, finalize, red),
+                    debug=debug,
+                    finalize=inkernel,
+                    red_out=red,
+                    **xc_kernel_args(),
+                    spec_k=spec[0],
+                    spec_ctas=spec[1],
+                    spec_self=spec[2],
+                )
+                ints = [cutlass.Int64(0)] * _NUM_XC_ADDRS
+                routed = cute.compile(kern.run, *ptrs, *ints, cutlass.Int32(1), stream)
+                if spec[0]:
+                    from ....shared import dispatch
+
+                    dispatch.log_capability(
+                        "routed_moe",
+                        f"expert_prefetch[k={spec[0]},ctas={spec[1]},self={int(spec[2])}]",
+                        True,
+                    )
+                fin = None
+                if finalize and not inkernel and not red:
+                    from .kernel_xc import DELAY_FINALIZE
+
+                    delay, delay_ns = DELAY_INJECTION
+                    fin = cute.compile(
+                        RoutedMoeFinalizeKernel(
+                            delay_ns=delay_ns if delay & DELAY_FINALIZE else 0
+                        ).run,
+                        cutlass.Int64(0),
+                        cutlass.Int64(0),
+                        cutlass.Int64(0),
+                        cutlass.Int32(1),
+                        stream,
+                    )
+                fns = (routed, fin)
+            else:
+                from .kernel import RoutedMoeKernel
+
+                delay, delay_ns = DELAY_INJECTION
+                kern = RoutedMoeKernel(
+                    self.num_ctas, finalize=finalize, debug=debug, delay=delay, delay_ns=delay_ns
+                )
+                ints = [cutlass.Int64(0)] * _NUM_ADDRS
+                fns = (cute.compile(kern.run, *ptrs, *ints, cutlass.Int32(1), stream), None)
+            _COMPILED[key] = fns
+            return fns
+
+    def _pointers(self, make_ptr, cute, cutlass):
+        m = self.module
+        return [
+            make_ptr(
+                cutlass.Float4E2M1FN,
+                m.w3_w1_weight.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=32,
+            ),
+            make_ptr(
+                cutlass.Float8E4M3FN,
+                m.w3_w1_weight_scale.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            ),
+            make_ptr(
+                cutlass.Float4E2M1FN,
+                m.w2_weight.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=32,
+            ),
+            make_ptr(
+                cutlass.Float8E4M3FN,
+                m.w2_weight_scale.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            ),
+        ]
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        do_finalize: bool = True,
+        debug: int = 0,
+        debug_buffer: Optional[torch.Tensor] = None,
+        out: Optional[torch.Tensor] = None,
+        raw_out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        red_out: Optional[torch.Tensor] = None,
+    ):
+        """Run the routed MoE.
+
+        ``out``: finalized output buffer. ``raw_out``: ``(rows [R*8, H] bf16,
+        scales [R, 8] bf16)`` buffers for the raw tuple (identity layout; the
+        permutation is not written). ``red_out``: FP32 ``[R, H]`` accumulator
+        that receives ``sum_k w[t, k] * fc2[t, k]`` (zero on entry; expert
+        cluster schedule only).
+        """
+        import cutlass
+        import cutlass.cute as cute
+        from cuda.bindings import driver as cuda
+        from cutlass.cute.runtime import make_ptr
+
+        m = self.module
+        rows = x.shape[0]
+        x = _bulk_copy_ready(x)
+        red = red_out is not None
+        if red:
+            do_finalize = False
+            if self.schedule(rows, False, True) != "xc":
+                raise ValueError("FP32 accumulator output needs the expert-cluster schedule")
+        routed_fn, fin_fn = self._compiled(do_finalize, debug, rows, red)
+        ws = _workspace(self.device, self.num_ctas)
+        dbg = debug_buffer if debug_buffer is not None else ws.dummy
+        current = torch.cuda.current_stream()
+        ws.order_on(current)
+        stream = cuda.CUstream(current.cuda_stream)
+        ptrs = self._pointers(make_ptr, cute, cutlass)
+        params = [
+            x.data_ptr(),
+            router_logits.data_ptr(),
+            self.bias.data_ptr(),
+            m.fc31_alpha.data_ptr(),
+            m.fc31_scale_c.data_ptr(),
+            m.fc2_alpha.data_ptr(),
+            m.fc31_input_scale.data_ptr(),
+        ]
+        if self.schedule(rows, do_finalize, red) == "xc":
+            if red:
+                out = red_out
+                raw_rows, raw_scale, raw_idx = ws.rows, ws.scales, ws.dummy
+            elif do_finalize:
+                raw_rows, raw_scale, raw_idx = ws.rows, ws.scales, ws.dummy
+                if out is None:
+                    out = torch.empty(rows, HIDDEN, dtype=torch.bfloat16, device=x.device)
+            elif raw_out is not None:
+                out = ws.dummy
+                raw_rows, raw_scale = raw_out
+                raw_idx = ws.dummy
+            else:
+                out = ws.dummy
+                raw_rows = torch.empty(rows * TOP_K, HIDDEN, dtype=torch.bfloat16, device=x.device)
+                raw_scale = torch.empty(rows, TOP_K, dtype=torch.bfloat16, device=x.device)
+                raw_idx = torch.empty(rows, TOP_K, dtype=torch.int32, device=x.device)
+            addrs = params + [
+                raw_rows.data_ptr(),
+                raw_scale.data_ptr(),
+                raw_idx.data_ptr(),
+                out.data_ptr(),
+                ws.tile_count.data_ptr(),
+                dbg.data_ptr(),
+                m.w3_w1_weight.data_ptr(),
+                m.w3_w1_weight_scale.data_ptr(),
+                m.w2_weight.data_ptr(),
+                m.w2_weight_scale.data_ptr(),
+            ]
+            spec = self._spec()
+            addrs.append(spec[0].data_ptr() if spec is not None else ws.dummy.data_ptr())
+            assert len(addrs) == _NUM_XC_ADDRS
+            routed_fn(*ptrs, *[cutlass.Int64(a) for a in addrs], cutlass.Int32(rows), stream)
+            if red:
+                return [red_out]
+            if not do_finalize:
+                return [raw_rows, raw_scale, raw_idx]
+            if fin_fn is None:
+                return out
+            fin_fn(
+                cutlass.Int64(raw_rows.data_ptr()),
+                cutlass.Int64(raw_scale.data_ptr()),
+                cutlass.Int64(out.data_ptr()),
+                cutlass.Int32(rows),
+                stream,
+            )
+            return out
+        if do_finalize:
+            if out is None:
+                out = torch.empty(rows, HIDDEN, dtype=torch.bfloat16, device=x.device)
+            raw_scale = ws.dummy
+            raw_idx = ws.dummy
+        else:
+            out = torch.empty(rows * TOP_K, HIDDEN, dtype=torch.bfloat16, device=x.device)
+            raw_scale = torch.empty(rows, TOP_K, dtype=torch.bfloat16, device=x.device)
+            raw_idx = torch.empty(rows, TOP_K, dtype=torch.int32, device=x.device)
+        addrs = params + [
+            out.data_ptr(),
+            raw_scale.data_ptr(),
+            raw_idx.data_ptr(),
+            ws.ws.data_ptr(),
+            ws.partials.data_ptr(),
+            ws.himg.data_ptr(),
+            ws.hflag.data_ptr(),
+            dbg.data_ptr(),
+            m.w3_w1_weight.data_ptr(),
+            m.w3_w1_weight_scale.data_ptr(),
+            m.w2_weight.data_ptr(),
+            m.w2_weight_scale.data_ptr(),
+        ]
+        assert len(addrs) == _NUM_ADDRS
+        routed_fn(*ptrs, *[cutlass.Int64(a) for a in addrs], cutlass.Int32(rows), stream)
+        if do_finalize:
+            return out
+        return [out, raw_scale, raw_idx]
+
+
+def _engaged_runner(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    do_finalize: bool,
+    feature: str = "routed_moe",
+    red: bool = False,
+) -> Optional["RoutedMoe"]:
+    """The runner for this call, or ``None`` (decided before any device work)."""
+    from ....shared import dispatch
+
+    ok, reason = eligibility(module, x, router_logits)
+    owner = resolve_expert_owner(module)[0] if ok else None
+    if ok:
+        rm = getattr(owner, "routing_method", None)
+        native_bias = getattr(rm, "e_score_correction_bias", None)
+        if (
+            routing_bias is None
+            or routing_bias.dtype != torch.bfloat16
+            or tuple(routing_bias.shape) != (NUM_EXPERTS,)
+            or not routing_bias.is_contiguous()
+            or routing_bias.device != x.device
+        ):
+            ok, reason = False, "router bias is not a contiguous BF16 [256] on the device"
+        elif native_bias is not None and native_bias.data_ptr() != routing_bias.data_ptr():
+            ok, reason = False, "router bias is not the routing method's bias"
+    runner = getattr(owner, "_routed_moe_runner", None) if ok else None
+    if ok and runner is not None and runner.bias.data_ptr() != routing_bias.data_ptr():
+        runner = None
+    if ok and torch.cuda.is_current_stream_capturing():
+        # Compiling or allocating the workspace inside a capture would record
+        # the workspace initialization into the graph; the model's eager
+        # warmup forwards make both ready before capture.
+        if runner is None or not runner.ready_for_capture(do_finalize, rows=x.shape[0], red=red):
+            ok, reason = False, "not warmed up before CUDA graph capture"
+    dispatch.log_capability(feature, f"rows={x.shape[0]} {reason or 'engaged'}", ok, reason)
+    if not ok:
+        return None
+    if runner is None:
+        runner = RoutedMoe(owner, routing_bias)
+        owner._routed_moe_runner = runner
+    return runner
+
+
+def routed_output(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    do_finalize: bool,
+) -> Optional[List[torch.Tensor]]:
+    """Entry point for the routed-MoE call site; ``None`` means "use the native path".
+
+    The decision is made before any device work, so returning ``None`` leaves
+    no side effect behind.
+    """
+    runner = _engaged_runner(module, routing_bias, x, router_logits, do_finalize)
+    if runner is None:
+        return None
+    return runner(x, router_logits, do_finalize=do_finalize)
+
+
+def routed_output_into_hop(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    hop,
+) -> Optional[List[torch.Tensor]]:
+    """routed_moe_into_hop: routed_moe writes its per-rank routed result into the post_moe_hop
+    post-MoE hop's routed-input buffer; post_moe_hop then adds the shared output,
+    reduces across ranks and applies residual and RMSNorm (``hop.hop_routed``).
+
+    The hop's routed-input mode (``hop.routed_mode``) decides what is written:
+
+    * "row": the finalized routed output ``[R, H]`` bf16 (finalize kernel);
+    * "identity": the raw FC2 rows ``[R * 8, H]`` in the identity layout plus
+      the routing weights in the hop's scale buffer; post_moe_hop finalizes without the
+      permutation round trip (bit-identical to the raw tuple);
+    * "fp32": an FP32 ``[R, H]`` accumulator of ``w[t, k] * fc2[t, k]``, added
+      by the routed kernel itself (no finalize kernel, no raw rows); the sum
+      order differs from the native finalize (FP32 reorder class).
+
+    Returns ``[buffer]`` (a view of the hop's buffer) or ``None`` when routed_moe or
+    the hop's routed input is not eligible; the decision is taken before any
+    device work and depends only on rank-identical state.
+    """
+    if hop is None or not getattr(hop, "_routed_enabled", False):
+        return None
+    rows = x.shape[0]
+    if not hop.can_run(rows):
+        return None
+    mode = getattr(hop, "routed_mode", "row")
+    if mode == "row":
+        runner = _engaged_runner(
+            module, routing_bias, x, router_logits, True, feature="routed_moe_into_hop"
+        )
+        if runner is None:
+            return None
+        target = hop.routed_input_buffer(rows)
+        return [runner(x, router_logits, do_finalize=True, out=target)]
+    red = mode == "fp32"
+    runner = _engaged_runner(
+        module, routing_bias, x, router_logits, False, "routed_moe_into_hop", red
+    )
+    if runner is None or runner.schedule(rows, False, red) != "xc":
+        return None
+    if red:
+        return runner(x, router_logits, red_out=hop.routed_acc_buffer(rows))
+    rows_buf, scales_buf = hop.routed_raw_buffers(rows)
+    runner(x, router_logits, do_finalize=False, raw_out=(rows_buf, scales_buf))
+    return [rows_buf]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/__init__.py
new file mode 100644
index 000000000000..bbdb33993a2c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/__init__.py
@@ -0,0 +1,36 @@
+# 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.
+"""Package shared_expert_prefetch: the MoE region END (routed END and shared END).
+
+Entry point registered with ``dispatch``:
+
+* ``routed_output_into_hop(module, bias, x, logits, routing, hop, shared)``:
+  hop_routing's routed_moe fork as ``SharedPrefetchExpertClusterKernel`` (``routed_moe.py``,
+  ``kernel_xc.py``). Before its dependency wait, while the pre-MoE hop runs, it prefetches the
+  shared expert's weights into L2, so the native shared expert on the auxiliary stream
+  reads them from L2 instead of competing with the routed kernel for DRAM.
+  ``None`` means "run hop_routing's fork" (decided before any device work).
+
+The shared expert itself stays native (``GatedMLP``); shared_expert_prefetch has no numerical
+variant: given the same routing its expert rows are routed_moe's and its FP32
+accumulator is in routed_moe_into_hop's ``fp32`` class.
+"""
+
+from ....shared.dispatch import register
+from .routed_moe import routed_output_into_hop
+
+register("shared_expert_prefetch", "routed_output_into_hop", routed_output_into_hop)
+
+__all__ = ["routed_output_into_hop"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/kernel_xc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/kernel_xc.py
new file mode 100644
index 000000000000..cb6e29f6ef7c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/kernel_xc.py
@@ -0,0 +1,745 @@
+# 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.
+"""Routed-MoE kernel of package shared_expert_prefetch: hop_routing's expert-cluster fork that
+stages the shared expert's weights in L2.
+
+``SharedPrefetchExpertClusterKernel`` is ``HopRoutedExpertClusterKernel`` (routed_moe's
+expert-cluster kernel reading the pre-MoE hop's routing, ``hop_routing/routed_moe.py``) with one
+addition. Before its ``griddepcontrol.wait``, each CTA issues ``cp.async.bulk.prefetch.L2`` for
+its ``1 / grid`` slice of the layer's BF16 shared-expert weights
+(``gate_up`` ``[512, 6144]`` and ``down`` ``[6144, 256]``, one TP8 rank shard,
+9,437,184 bytes). The lanes of the TMA warp each prefetch one chunk.
+
+The routed kernel is launched with programmatic dependent launch during the
+pre-MoE hop and waits there while DRAM is nearly idle. The prefetch uses
+that window, so the shared expert, which runs beside the routed kernel on
+the auxiliary stream right after the hop, reads its weights from L2 instead
+of competing with the routed weight stream for DRAM.
+
+``run`` and ``kernel`` are routed_moe's (``routed_moe/kernel_xc.py``, including package
+cross_layer_expert_prefetch's speculative expert prefetch) with the two shared-expert addresses
+added to their parameters (after cross_layer_expert_prefetch's ``spec_ids_addr``) and the
+shared-expert prefetch added before the wait. Every other line, every helper method and
+``_route_tokens`` (hop_routing's) are unchanged, so the schedule, the arithmetic and
+the outputs are those of hop_routing's fork.
+
+With package cross_layer_expert_prefetch (``spec_k`` > 0), the TMA warp runs one pre-wait prologue:
+each lane issues its shared-expert chunks first (the shared expert reads all
+of them right after the hop), then its chunks of cross_layer_expert_prefetch's predicted experts in
+rank order (``_xc_spec_prefetch``). When cross_layer_expert_prefetch restricts its issuers to the
+first ``spec_ctas`` CTAs (the CTAs resident before the dependency release),
+the shared-expert chunks are spread over the same CTAs (``se_ctas``);
+otherwise over the whole grid, as without cross_layer_expert_prefetch. With ``spec_self``, CTA 0's
+MMA warp stores the call's distinct experts, ranked by row count, into cross_layer_expert_prefetch's
+row after routing (``_xc_spec_store``), as in routed_moe's kernel.
+
+Memory order and progress:
+
+* Before ``griddepcontrol.wait`` the addition reads only the shared-expert
+  weights, through the L2 prefetch. They are model parameters that nothing
+  writes while serving. It writes nothing and touches no flag, buffer or
+  counter.
+* Everything a producer owns (activation, routing, the FP32 accumulator) is
+  read and written only after the wait, exactly as in hop_routing's fork.
+* The prefetch is an async-proxy read of memory that the generic proxy never
+  writes while serving, so it needs no fence. Whether the lines are in L2 or
+  not, the shared expert's own loads (ordered by its own graph edges) read the
+  same bytes.
+* No wait, spin or cross-CTA dependency is added.
+* cross_layer_expert_prefetch's prologue reads its row of predicted expert IDs with relaxed loads
+  before the wait. The row is written only by cross_layer_expert_prefetch (on another stream) or by
+  CTA 0 of the same layer's previous call, after that call's wait; any value
+  read only selects which immutable weight bytes are prefetched (each ID is
+  masked into [0, 256)), never what the kernel reads or computes (the memory
+  order of ``hop_routing/routed_moe.py`` and ``routed_moe/routed_moe.py``).
+"""
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32
+
+from ..hop_routing import routed_moe as hop_routed_mod
+from ..routed_moe import ptx
+from ..routed_moe.kernel_xc import (
+    A_STAGE_BYTES,
+    B2_BAR,
+    BAR_EPI,
+    CLUSTER,
+    DEBUG_ROUTING,
+    DEBUG_TIMELINE,
+    DELAY_ENTRY,
+    EPI_WARPS,
+    FC1_KB,
+    FC2_MT,
+    HIDDEN,
+    HSTAGE_BYTES,
+    KB_PER_CTA,
+    KB_TX_BYTES,
+    MAX_ROUNDS,
+    MAX_SLOTS,
+    MMA_WARP,
+    NTOK,
+    NUM_BAR,
+    NUM_EXPERTS,
+    NUM_THREADS,
+    SFA_STAGE_BYTES,
+    SFB2_BYTES,
+    SFB_BYTES,
+    TL_BYTE_BASE,
+    TL_SLOTS,
+    TMA_WARP,
+    TMEM_COLS,
+    TOP_K,
+    X_BAR,
+    X_ROW_BYTES,
+    XS_ROW_BYTES,
+    XT_ENTRY,
+    XT_EXIT,
+    XT_GT0,
+    XT_PDL,
+    XT_ROUTED,
+    _stamp,
+)
+
+SE_INTER = 256  # per-rank shared intermediate size (2048 / TP8)
+SE_W1_BYTES = 2 * SE_INTER * HIDDEN * 2  # gate_up [512, 6144] BF16
+SE_W2_BYTES = HIDDEN * SE_INTER * 2  # down [6144, 256] BF16
+PF_LANES = 32
+# Debug-timeline slots of the pre-wait prologue (routed_moe's ``XT_*`` slots end at
+# 30; ``XT_GT0`` is 47): SM clock at its start and after its prefetches are
+# issued, the global timer at its start, and the global timer when the
+# dependency wait returns (the release, next to ``XT_PDL``).
+XZ_PRO_CLK0 = 40
+XZ_PRO_CLK1 = 41
+XZ_PRO_GT0 = 42
+XZ_PDL_GT = 43
+
+
+def _chunk(total: int, ctas: int) -> int:
+    """Bytes per prefetching lane (a multiple of 16) so that ``ctas x 32`` lanes cover ``total``."""
+    per = -(-total // (ctas * PF_LANES))
+    return -(-per // 16) * 16
+
+
+class SharedPrefetchExpertClusterKernel(hop_routed_mod._kernel_class()):
+    """hop_routing's expert-cluster kernel plus the pre-wait L2 prefetch of the shared expert's weights."""
+
+    def __init__(self, num_ctas: int, *args, se_pf: bool = True, **kwargs):
+        super().__init__(num_ctas, *args, **kwargs)
+        self.se_pf = se_pf
+        # CTAs that issue the shared-expert prefetch: every CTA, or, when
+        # cross_layer_expert_prefetch's prefetch is compiled in for its first ``spec_ctas`` CTAs
+        # (the CTAs resident before the dependency release), the same CTAs.
+        self.se_ctas = num_ctas
+        if self.spec_k > 0 and self.spec_ctas > 0:
+            self.se_ctas = min(num_ctas, self.spec_ctas)
+        self.se_chunk1 = _chunk(SE_W1_BYTES, self.se_ctas)
+        self.se_chunk2 = _chunk(SE_W2_BYTES, self.se_ctas)
+
+    @cute.jit
+    def _se_prefetch(self, bidx, lane, se_w1_addr: Int64, se_w2_addr: Int64):
+        """Lane ``lane`` of CTA ``bidx`` prefetches chunk ``bidx * 32 + lane`` of ``gate_up`` and of ``down``.
+
+        The chunks are sized for the first ``se_ctas`` CTAs; CTAs past them own no chunk.
+        """
+        i = Int64(bidx * PF_LANES + lane)
+        s1 = i * Int64(self.se_chunk1)
+        if s1 < Int64(SE_W1_BYTES):
+            n1 = cutlass.min(Int64(self.se_chunk1), Int64(SE_W1_BYTES) - s1)
+            ptx.prefetch_l2_bulk(se_w1_addr + s1, Int32(n1))
+        s2 = i * Int64(self.se_chunk2)
+        if s2 < Int64(SE_W2_BYTES):
+            n2 = cutlass.min(Int64(self.se_chunk2), Int64(SE_W2_BYTES) - s2)
+            ptx.prefetch_l2_bulk(se_w2_addr + s2, Int32(n2))
+
+    @cute.jit
+    def run(
+        self,
+        w13_ptr: cute.Pointer,
+        w13_sf_ptr: cute.Pointer,
+        w2_ptr: cute.Pointer,
+        w2_sf_ptr: cute.Pointer,
+        x_addr: Int64,
+        logits_addr: Int64,
+        bias_addr: Int64,
+        alpha1_addr: Int64,
+        scale_c1_addr: Int64,
+        alpha2_addr: Int64,
+        xgs_addr: Int64,
+        rows_addr: Int64,
+        raw_scale_addr: Int64,
+        raw_idx_addr: Int64,
+        out_addr: Int64,
+        counter_addr: Int64,
+        dbg_addr: Int64,
+        w13_addr: Int64,
+        w13_sf_addr: Int64,
+        w2_addr: Int64,
+        w2_sf_addr: Int64,
+        spec_ids_addr: Int64,
+        se_w1_addr: Int64,
+        se_w2_addr: Int64,
+        num_tokens: Int32,
+        stream,
+    ):
+        self._setup()
+        w13 = cute.make_tensor(
+            w13_ptr, cute.make_ordered_layout((2 * 256, HIDDEN, NUM_EXPERTS), order=(1, 0, 2))
+        )
+        w2 = cute.make_tensor(
+            w2_ptr, cute.make_ordered_layout((HIDDEN, 256, NUM_EXPERTS), order=(1, 0, 2))
+        )
+        w13_sf = cute.make_tensor(
+            w13_sf_ptr, blockscaled_utils.tile_atom_to_shape_SF(w13.shape, self.sf_vec_size)
+        )
+        w2_sf = cute.make_tensor(
+            w2_sf_ptr, blockscaled_utils.tile_atom_to_shape_SF(w2.shape, self.sf_vec_size)
+        )
+        a_op = sm100_utils.cluster_shape_to_tma_atom_A((1, 1), self.tiled_mma.thr_id)
+        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+        sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
+        tma_atom_a1, tma_a1 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op, w13, a_smem_layout, self.mma_tiler, self.tiled_mma, self.cluster_layout_vmnk.shape
+        )
+        tma_atom_sfa1, tma_sfa1 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            w13_sf,
+            sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+        tma_atom_a2, tma_a2 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op, w2, a_smem_layout, self.mma_tiler, self.tiled_mma, self.cluster_layout_vmnk.shape
+        )
+        tma_atom_sfa2, tma_sfa2 = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            w2_sf,
+            sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            ring_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.num_stages]
+            acc_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.acc_stages]
+            b_mbar: cute.struct.MemRange[cutlass.Int64, NUM_BAR]
+            tmem_holding: cutlass.Int32
+            tl_buf: cute.struct.MemRange[cutlass.Int64, TL_SLOTS]
+            hstage: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, HSTAGE_BYTES // 8], 16]
+            misc: cute.struct.MemRange[cutlass.Int32, 16]
+            closed: cute.struct.MemRange[cutlass.Int32, MAX_ROUNDS * FC2_MT]
+            bitmap: cute.struct.MemRange[cutlass.Uint32, 8]
+            topk_id: cute.struct.MemRange[cutlass.Int32, NTOK * TOP_K]
+            topk_w: cute.struct.MemRange[cutlass.Float32, NTOK * TOP_K]
+            slot_expert: cute.struct.MemRange[cutlass.Int32, MAX_SLOTS]
+            slot_mask: cute.struct.MemRange[cutlass.Uint32, MAX_SLOTS]
+            slot_w: cute.struct.MemRange[cutlass.Float32, MAX_SLOTS * NTOK]
+            slot_k: cute.struct.MemRange[cutlass.Int32, MAX_SLOTS * NTOK]
+            expert_slot: cute.struct.MemRange[cutlass.Int32, NUM_EXPERTS]
+            sA: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.a_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFA: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float8E4M3FN, cute.cosize(self.sfa_smem_layout_staged)
+                ],
+                1024,
+            ]
+            # this CTA's share of the activation rows (NTOK x 6 k-blocks)
+            sX: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, NTOK * XS_ROW_BYTES], 1024]
+            sB: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.b_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFB: cute.struct.Align[cute.struct.MemRange[cutlass.Float8E4M3FN, SFB_BYTES], 1024]
+            sB2: cute.struct.Align[
+                cute.struct.MemRange[
+                    cutlass.Float4E2M1FN, cute.cosize(self.b2_smem_layout_staged.outer)
+                ],
+                1024,
+            ]
+            sSFB2: cute.struct.Align[cute.struct.MemRange[cutlass.Float8E4M3FN, SFB2_BYTES], 1024]
+
+        self.shared_storage = SharedStorage
+
+        self.kernel(
+            self.tiled_mma,
+            tma_atom_a1,
+            tma_a1,
+            tma_atom_sfa1,
+            tma_sfa1,
+            tma_atom_a2,
+            tma_a2,
+            tma_atom_sfa2,
+            tma_sfa2,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.sfb_smem_layout_std,
+            self.b2_smem_layout_staged,
+            self.sfb2_smem_layout_staged,
+            x_addr,
+            logits_addr,
+            bias_addr,
+            alpha1_addr,
+            scale_c1_addr,
+            alpha2_addr,
+            xgs_addr,
+            rows_addr,
+            raw_scale_addr,
+            raw_idx_addr,
+            out_addr,
+            counter_addr,
+            dbg_addr,
+            w13_addr,
+            w13_sf_addr,
+            w2_addr,
+            w2_sf_addr,
+            spec_ids_addr,
+            se_w1_addr,
+            se_w2_addr,
+            num_tokens,
+        ).launch(
+            grid=[self.num_ctas, 1, 1],
+            block=[NUM_THREADS, 1, 1],
+            cluster=[CLUSTER, 1, 1],
+            smem=self.shared_storage.size_in_bytes(),
+            min_blocks_per_mp=1,
+            stream=stream,
+            use_pdl=True,
+        )
+
+    # ------------------------------------------------------------------
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a1: cute.CopyAtom,
+        mA1: cute.Tensor,
+        tma_atom_sfa1: cute.CopyAtom,
+        mSFA1: cute.Tensor,
+        tma_atom_a2: cute.CopyAtom,
+        mA2: cute.Tensor,
+        tma_atom_sfa2: cute.CopyAtom,
+        mSFA2: cute.Tensor,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfb_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_std: cute.Layout,
+        b2_smem_layout_staged: cute.ComposedLayout,
+        sfb2_smem_layout_staged: cute.Layout,
+        x_addr: Int64,
+        logits_addr: Int64,
+        bias_addr: Int64,
+        alpha1_addr: Int64,
+        scale_c1_addr: Int64,
+        alpha2_addr: Int64,
+        xgs_addr: Int64,
+        rows_addr: Int64,
+        raw_scale_addr: Int64,
+        raw_idx_addr: Int64,
+        out_addr: Int64,
+        counter_addr: Int64,
+        dbg_addr: Int64,
+        w13_addr: Int64,
+        w13_sf_addr: Int64,
+        w2_addr: Int64,
+        w2_sf_addr: Int64,
+        spec_ids_addr: Int64,
+        se_w1_addr: Int64,
+        se_w2_addr: Int64,
+        num_tokens: Int32,
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        tidx, _, _ = cute.arch.thread_idx()
+        lane = tidx % 32
+        bidx, _, _ = cute.arch.block_idx()
+        grid, _, _ = cute.arch.grid_dim()
+        bidx = Int32(bidx)
+        rk = ptx.cluster_ctarank()
+        cl = bidx // CLUSTER
+        ncl = Int32(grid) // CLUSTER
+
+        # shared_expert_prefetch (+ cross_layer_expert_prefetch): one pre-wait prologue in the TMA
+        # warp. Its lanes first start the L2 prefetch of this CTA's slice of the layer's
+        # shared-expert weights, then (cross_layer_expert_prefetch) of their chunks of the
+        # predicted routed experts in rank order; both are immutable model
+        # parameters (see the module docstring). With cross_layer_expert_prefetch's issuing CTAs
+        # set, both lists are spread over those CTAs only.
+        # Debug timeline only: when the prologue starts (SM clock and global
+        # timer) and when its prefetches are issued (SM clock), recorded by the
+        # TMA warp's lane 0 and stored into slots XZ_PRO_* after the shared
+        # memory set-up below.
+        pro_on = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        if cutlass.const_expr(pro_on):
+            pro_c0 = Int64(0)
+            pro_g0 = Int64(0)
+            pro_c1 = Int64(0)
+        if warp_idx == TMA_WARP:
+            if cutlass.const_expr(pro_on):
+                pro_c0 = ptx.clock64()
+                pro_g0 = ptx.globaltimer_u64()
+            if cutlass.const_expr(self.se_pf):
+                self._se_prefetch(bidx, lane, se_w1_addr, se_w2_addr)
+            if cutlass.const_expr(self.spec_k > 0):
+                self._xc_spec_prefetch(
+                    bidx,
+                    Int32(grid),
+                    lane,
+                    spec_ids_addr,
+                    w13_addr,
+                    w13_sf_addr,
+                    w2_addr,
+                    w2_sf_addr,
+                )
+            if cutlass.const_expr(pro_on):
+                pro_c1 = ptx.clock64()
+
+        if warp_idx == TMA_WARP:
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a1)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa1)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a2)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa2)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        ring = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.ring_mbar.data_ptr(),
+            num_stages=self.num_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+            tx_count=A_STAGE_BYTES + SFA_STAGE_BYTES,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+        acc_pipe = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_mbar.data_ptr(),
+            num_stages=self.acc_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, EPI_WARPS),
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+        b_mbar = storage.b_mbar.data_ptr()
+        if warp_idx == MMA_WARP:
+            with cute.arch.elect_one():
+                for kb in cutlass.range_constexpr(FC1_KB):
+                    cute.arch.mbarrier_init(b_mbar + kb, 1)
+                cute.arch.mbarrier_init(b_mbar + X_BAR, 1)
+                for rho in cutlass.range_constexpr(MAX_ROUNDS):
+                    cute.arch.mbarrier_init(b_mbar + B2_BAR + rho, CLUSTER)
+                # k-blocks quantized by a peer complete on its bulk copies
+                for kb in cutlass.range_constexpr(FC1_KB):
+                    if Int32(kb // KB_PER_CTA) != rk:
+                        ptx.mbar_arrive_expect_tx(Int32((b_mbar + kb).toint()), Int32(KB_TX_BYTES))
+        cute.arch.mbarrier_init_fence()
+        # Peers arrive on this CTA's operand barriers: their initialization
+        # must be visible cluster-wide first (waited before the first remote
+        # arrive, below).
+        cute.arch.cluster_arrive_relaxed()
+
+        if warp_idx == 0:
+            cute.arch.alloc_tmem(TMEM_COLS, storage.tmem_holding)
+
+        s_misc = storage.misc.get_tensor(cute.make_layout(16))
+        s_bitmap = storage.bitmap.get_tensor(cute.make_layout(8))
+        s_topk_id = storage.topk_id.get_tensor(cute.make_layout(NTOK * TOP_K))
+        s_topk_w = storage.topk_w.get_tensor(cute.make_layout(NTOK * TOP_K))
+        s_slot_expert = storage.slot_expert.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_slot_mask = storage.slot_mask.get_tensor(cute.make_layout(MAX_SLOTS))
+        s_slot_w = storage.slot_w.get_tensor(cute.make_layout(MAX_SLOTS * NTOK))
+        s_slot_k = storage.slot_k.get_tensor(cute.make_layout(MAX_SLOTS * NTOK))
+        s_expert_slot = storage.expert_slot.get_tensor(cute.make_layout(NUM_EXPERTS))
+
+        s_expert_slot[tidx] = Int32(-1)
+        s_slot_w[tidx] = Float32(0.0)
+        s_slot_w[tidx + NUM_THREADS] = Float32(0.0)
+        if tidx < MAX_SLOTS:
+            s_slot_mask[tidx] = Uint32(0)
+        if tidx < 16:
+            s_misc[tidx] = Int32(0)
+        if cutlass.const_expr(self.debug & DEBUG_TIMELINE):
+            if tidx < TL_SLOTS:
+                ptx.sts_u64(Int32(storage.tl_buf.data_ptr().toint()) + tidx * 8, Int64(0))
+        if tidx < 8:
+            s_bitmap[tidx] = Uint32(0)
+        bias_r = cute.make_rmem_tensor((8,), Float32)
+        for j in cutlass.range_constexpr(8):
+            bias_r[j] = ptx.bf16x2_to_f32(ptx.ld_nc_u16(bias_addr + (lane + 32 * j) * 2), False)
+        cute.arch.sync_threads()
+
+        tl = cutlass.const_expr(self.debug & DEBUG_TIMELINE)
+        tl_s = Int32(storage.tl_buf.data_ptr().toint())
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_ENTRY)
+            if cutlass.const_expr(tl):
+                ptx.sts_u64(tl_s + XT_GT0 * 8, ptx.globaltimer_u64())
+        if cutlass.const_expr(pro_on):
+            if warp_idx == TMA_WARP:
+                if lane == 0:
+                    ptx.sts_u64(tl_s + XZ_PRO_CLK0 * 8, pro_c0)
+                    ptx.sts_u64(tl_s + XZ_PRO_GT0 * 8, pro_g0)
+                    ptx.sts_u64(tl_s + XZ_PRO_CLK1 * 8, pro_c1)
+        cute.arch.griddepcontrol_wait()
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_PDL)
+            if cutlass.const_expr(pro_on):
+                ptx.sts_u64(tl_s + XZ_PDL_GT * 8, ptx.globaltimer_u64())
+        if cutlass.const_expr(self.delay & DELAY_ENTRY):
+            self._xc_delay(DELAY_ENTRY * 53, bidx)
+        if warp_idx == TMA_WARP:
+            with cute.arch.elect_one():
+                xbar = Int32((b_mbar + X_BAR).toint())
+                sx = Int32(storage.sX.data_ptr().toint())
+                ptx.mbar_arrive_expect_tx(xbar, num_tokens * XS_ROW_BYTES)
+                r = Int32(0)
+                while r < num_tokens:
+                    ptx.bulk_g2s(
+                        sx + r * XS_ROW_BYTES,
+                        x_addr + r * X_ROW_BYTES + rk * XS_ROW_BYTES,
+                        Int32(XS_ROW_BYTES),
+                        xbar,
+                    )
+                    r = r + 1
+        # Every CTA of the cluster has initialized its barriers: remote
+        # copies and arrivals may start.
+        cute.arch.cluster_wait()
+        if cutlass.const_expr(self.trigger == "early"):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        # -------- routing (warp t handles token t) -> slot tables --------
+        self._route_tokens(
+            warp_idx, lane, num_tokens, logits_addr, bias_r, s_topk_id, s_topk_w, storage
+        )
+        cute.arch.sync_threads()
+        e = tidx
+        bw = s_bitmap[e // 32]
+        prefix = Int32(0)
+        for w in cutlass.range_constexpr(8):
+            word = s_bitmap[w]
+            if w < e // 32:
+                prefix = prefix + ptx.popc_u32(word)
+        below = bw & ((Uint32(1) << Uint32(e % 32)) - Uint32(1))
+        prefix = prefix + ptx.popc_u32(below)
+        if ((bw >> Uint32(e % 32)) & Uint32(1)) != Uint32(0):
+            s_slot_expert[prefix] = e
+            s_expert_slot[e] = prefix
+        if tidx == 0:
+            total = Int32(0)
+            for w in cutlass.range_constexpr(8):
+                total = total + ptx.popc_u32(s_bitmap[w])
+            s_misc[0] = total
+        cute.arch.sync_threads()
+        if tidx < NTOK * TOP_K:
+            t = tidx // TOP_K
+            k = tidx % TOP_K
+            if t < num_tokens:
+                e = s_topk_id[tidx]
+                sl = s_expert_slot[e]
+                s_slot_w[sl * NTOK + t] = s_topk_w[tidx]
+                s_slot_k[sl * NTOK + t] = k
+                ptx.atom_or_shared_u32(
+                    Int32(storage.slot_mask.data_ptr().toint()) + sl * 4, Uint32(1) << Uint32(t)
+                )
+        cute.arch.sync_threads()
+        num_slots = s_misc[0]
+        if tidx == 0:
+            _stamp(tl, tl_s, XT_ROUTED)
+
+        # -------- expert rounds of this cluster --------
+        # Round 0 spreads min(D, clusters) experts over all clusters (an
+        # expert served by r clusters splits its FC2 tiles r ways); when
+        # D > clusters, round 1 spreads the remaining D - clusters experts
+        # over all clusters the same way. Every cluster then streams two FC1
+        # tiles per CTA and an even share of the FC2 tiles.
+        nact = ncl
+        if cutlass.const_expr(self.min_clusters > 0):
+            if num_slots <= ncl:
+                nact = cutlass.max(num_slots, Int32(self.min_clusters))
+                nact = cutlass.min(nact, ncl)
+        if cutlass.const_expr(self.two_round_clusters > 0):
+            if num_slots > ncl:
+                if num_slots <= Int32(MAX_ROUNDS * self.two_round_clusters):
+                    nact = cutlass.min(Int32(self.two_round_clusters), ncl)
+        rounds = (num_slots + nact - 1) // nact
+        if cl >= nact:
+            rounds = Int32(0)
+        slots = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        shares = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        firsts = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        t_lo = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        t_hi = cute.make_rmem_tensor((MAX_ROUNDS,), Int32)
+        for rho in cutlass.range_constexpr(MAX_ROUNDS):
+            base = Int32(0)
+            count = cutlass.min(num_slots, nact)
+            if cutlass.const_expr(rho == 1):
+                base = nact
+                count = cutlass.max(num_slots - nact, Int32(1))
+            i = cl % count
+            j = cl // count
+            nshare = (nact - 1 - i) // count + 1
+            a = (j * FC2_MT) // nshare
+            n = ((j + 1) * FC2_MT) // nshare - a
+            slots[rho] = base + i
+            shares[rho] = nshare
+            firsts[rho] = Int32(0)
+            if j == 0:
+                firsts[rho] = Int32(1)
+            t_lo[rho] = a + (rk * n) // CLUSTER
+            t_hi[rho] = a + ((rk + 1) * n) // CLUSTER
+
+        if warp_idx == TMA_WARP:
+            self._xc_tma_producer(
+                tma_atom_a1,
+                mA1,
+                tma_atom_sfa1,
+                mSFA1,
+                tma_atom_a2,
+                mA2,
+                tma_atom_sfa2,
+                mSFA2,
+                a_smem_layout_staged,
+                sfa_smem_layout_staged,
+                tiled_mma,
+                cluster_layout_vmnk,
+                storage,
+                ring,
+                b_mbar,
+                s_slot_expert,
+                rk,
+                rounds,
+                slots,
+                shares,
+                firsts,
+                t_lo,
+                t_hi,
+                w13_addr,
+                w13_sf_addr,
+                w2_addr,
+                w2_sf_addr,
+            )
+        if cutlass.const_expr(self.spec_self):
+            if bidx == 0:
+                if warp_idx == MMA_WARP:
+                    self._xc_spec_store(lane, num_slots, spec_ids_addr, storage, s_slot_expert)
+        if warp_idx == MMA_WARP:
+            self._xc_mma_issuer(
+                tiled_mma,
+                storage,
+                ring,
+                acc_pipe,
+                b_mbar,
+                a_smem_layout_staged,
+                sfa_smem_layout_staged,
+                b_smem_layout_staged,
+                sfb_smem_layout_staged,
+                sfb_smem_layout_std,
+                b2_smem_layout_staged,
+                sfb2_smem_layout_staged,
+                rounds,
+                t_lo,
+                t_hi,
+            )
+        if (warp_idx != TMA_WARP) & (warp_idx != MMA_WARP):
+            if rounds > 0:
+                self._xc_quantize_x(
+                    warp_idx, lane, xgs_addr, num_tokens, storage, b_mbar, dbg_addr, rk
+                )
+            else:
+                # an idle cluster still lets its activation copy land
+                cute.arch.mbarrier_wait(b_mbar + X_BAR, 0)
+
+        if bidx == 0:
+            if tidx < NTOK * TOP_K:
+                t = tidx // TOP_K
+                if t < num_tokens:
+                    ptx.st_u16(raw_scale_addr + tidx * 2, ptx.bf16_bits(s_topk_w[tidx]))
+                    ptx.st_u32(raw_idx_addr + tidx * 4, Uint32(tidx))
+        if cutlass.const_expr(self.debug & DEBUG_ROUTING):
+            if bidx == 0:
+                if tidx < NTOK * TOP_K:
+                    ptx.st_u32(dbg_addr + tidx * 4, Uint32(s_topk_id[tidx]))
+                    ptx.st_f32(dbg_addr + 256 + tidx * 4, s_topk_w[tidx])
+                if tidx < MAX_SLOTS:
+                    ptx.st_u32(dbg_addr + 512 + tidx * 4, Uint32(s_slot_expert[tidx]))
+                    ptx.st_u32(dbg_addr + 768 + tidx * 4, s_slot_mask[tidx])
+                if tidx == 0:
+                    ptx.st_u32(dbg_addr + 1024, Uint32(num_slots))
+
+        if warp_idx < EPI_WARPS:
+            self._xc_epilogue(
+                tiled_mma,
+                tidx,
+                warp_idx,
+                lane,
+                storage,
+                acc_pipe,
+                b_mbar,
+                s_slot_expert,
+                s_slot_mask,
+                s_slot_k,
+                rk,
+                rounds,
+                slots,
+                firsts,
+                t_lo,
+                t_hi,
+                num_tokens,
+                alpha1_addr,
+                scale_c1_addr,
+                alpha2_addr,
+                rows_addr,
+                out_addr,
+                counter_addr,
+                s_topk_w,
+                num_slots,
+                bidx,
+                dbg_addr,
+            )
+        # No CTA leaves while a peer may still store into its shared memory.
+        cute.arch.cluster_arrive()
+        cute.arch.cluster_wait()
+        if warp_idx < EPI_WARPS:
+            if cutlass.const_expr(self.debug & DEBUG_TIMELINE):
+                ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+                if tidx == 0:
+                    _stamp(tl, tl_s, XT_EXIT)
+                ptx.bar_sync(BAR_EPI, 32 * EPI_WARPS)
+                if tidx < TL_SLOTS:
+                    ptx.st_u64(
+                        dbg_addr + TL_BYTE_BASE + (bidx * TL_SLOTS + tidx) * 8,
+                        ptx.lds_u64(tl_s + tidx * 8),
+                    )
+            if warp_idx == 0:
+                tmem_ptr = cute.arch.retrieve_tmem_ptr(
+                    Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding
+                )
+                cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/routed_moe.py
new file mode 100644
index 000000000000..ba398a5031cf
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/routed_moe.py
@@ -0,0 +1,308 @@
+# 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.
+"""Routed side of package shared_expert_prefetch: hop_routing's routed_moe fork that stages the
+shared expert in L2, on a chosen grid.
+
+The MoE region ends with the later of two branches. On the
+main stream, hop_routing's fork of routed_moe's expert-cluster kernel adds the routed experts
+into routed_moe_into_hop's FP32 accumulator. On the auxiliary stream, the native BF16 shared
+expert runs. Both stream their weights from DRAM at the same time, and the
+routed grid is kept at 24 clusters so that the shared expert has SMs.
+
+``SharedPrefetchRoutedMoe`` runs ``SharedPrefetchExpertClusterKernel`` (``kernel_xc.py``). That is
+hop_routing's kernel plus an L2 prefetch of the layer's shared-expert weights, issued before its
+dependency wait while the pre-MoE hop is still running. Its grid in the FP32
+accumulator mode is ``min(co-resident clusters, DEFAULT_CLUSTERS)`` clusters, at
+least ``ceil(R * 8 / 2)``. The schedule, its rules (routed_moe's
+``xc_kernel_args``), the memory order after the wait and the arithmetic are those
+of hop_routing's fork. The grid changes only the FP32 order in which top-k
+contributions are added into the accumulator (routed_moe_into_hop's ``fp32``
+class).
+
+With package cross_layer_expert_prefetch the runner key carries cross_layer_expert_prefetch's
+prefetch spec (``key[6]``, ``RoutedMoe._key``), which ``_compiled`` passes to
+``SharedPrefetchExpertClusterKernel`` exactly as hop_routing's and routed_moe's runners do
+(capability ``shared_expert_prefetch/expert_prefetch[k=..,ctas=.., self=..]``), and ``__call__``
+passes cross_layer_expert_prefetch's row of predicted experts as the 18th routed_moe address, before
+the two shared-expert addresses. The kernel then issues the shared-expert chunks and
+cross_layer_expert_prefetch's chunks from one pre-wait prologue (``kernel_xc.py``).
+
+``routed_output_into_hop(module, bias, x, logits, routing, hop, shared)``
+returns ``None`` before any device work unless the hop takes the FP32
+accumulator, hop_routing's routing is present, and ``shared`` is a BF16 shared expert
+whose weights the kernel can prefetch. The caller then runs hop_routing's fork.
+"""
+
+import threading
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from ..hop_routing import routed_moe as hop_routed_mod
+from ..routed_moe import routed_moe as routed_moe_mod
+
+FEATURE = "shared_expert_prefetch"
+DEFAULT_CLUSTERS = 24
+SE_HIDDEN = 6144
+SE_INTER = 256
+
+_LOCK = threading.Lock()
+_COMPILED: Dict[tuple, object] = {}
+
+
+def shared_weights(shared) -> Tuple[Optional[Tuple[torch.Tensor, torch.Tensor]], str]:
+    """``(gate_up, down)`` weights of a served BF16 shared expert, or ``(None, reason)`` (host-only)."""
+    if shared is None or type(shared).__name__ != "GatedMLP":
+        return None, "no BF16 GatedMLP shared expert"
+    w1 = getattr(getattr(shared, "gate_up_proj", None), "weight", None)
+    w2 = getattr(getattr(shared, "down_proj", None), "weight", None)
+    for name, w, shape in (
+        ("gate_up", w1, (2 * SE_INTER, SE_HIDDEN)),
+        ("down", w2, (SE_HIDDEN, SE_INTER)),
+    ):
+        if (
+            not isinstance(w, torch.Tensor)
+            or w.dtype != torch.bfloat16
+            or tuple(w.shape) != shape
+            or not w.is_contiguous()
+            or not w.is_cuda
+            or w.data_ptr() % 16
+        ):
+            return (
+                None,
+                f"shared {name} weight is not a contiguous 16-byte aligned BF16 {shape} device tensor",
+            )
+    return (w1, w2), ""
+
+
+class SharedPrefetchRoutedMoe(hop_routed_mod.HopRoutedMoe):
+    """hop_routing's fork running ``SharedPrefetchExpertClusterKernel`` on the
+    shared_expert_prefetch grid (FP32 accumulator output only)."""
+
+    def _xc_ctas(
+        self, rows: int = routed_moe_mod.MAX_ROWS, finalize: bool = True, red: bool = False
+    ) -> int:
+        from ..routed_moe.kernel_xc import CLUSTER, MAX_ROUNDS
+
+        if not red or finalize:
+            return super()._xc_ctas(rows, finalize, red)
+        co_resident = routed_moe_mod._max_active_clusters(self.device, CLUSTER)
+        want = DEFAULT_CLUSTERS
+        need = -(-rows * routed_moe_mod.TOP_K // MAX_ROUNDS)
+        n = CLUSTER * min(co_resident, max(want, need))
+        n = min(n, self.num_ctas)
+        return n - n % CLUSTER
+
+    def _key(
+        self, finalize: bool, debug: int, rows: int = routed_moe_mod.MAX_ROWS, red: bool = False
+    ):
+        return super()._key(finalize, debug, rows, red) + (("shared_expert_prefetch", True),)
+
+    def ready_for_capture(
+        self,
+        finalize: bool = False,
+        debug: int = 0,
+        rows: int = routed_moe_mod.MAX_ROWS,
+        red: bool = True,
+    ) -> bool:
+        key = self._key(finalize, debug, rows, red)
+        dev = self.device.index if self.device.index is not None else torch.cuda.current_device()
+        ws = routed_moe_mod._WORKSPACES.get(dev)
+        return key in _COMPILED and ws is not None and ws.num_ctas >= self.num_ctas
+
+    def _compiled(
+        self, finalize: bool, debug: int, rows: int = routed_moe_mod.MAX_ROWS, red: bool = False
+    ):
+        """The compiled ``SharedPrefetchExpertClusterKernel`` launcher (expert clusters, FP32 accumulator only)."""
+        key = self._key(finalize, debug, rows, red)
+        if key[3] != "xc" or finalize or not red:
+            raise ValueError(
+                "the shared_expert_prefetch fork runs the expert-cluster schedule into the FP32 accumulator only"
+            )
+        with _LOCK:
+            fn = _COMPILED.get(key)
+            if fn is not None:
+                return fn
+            import cutlass
+            import cutlass.cute as cute
+            from cuda.bindings import driver as cuda
+            from cutlass.cute.runtime import make_ptr
+
+            from .kernel_xc import SharedPrefetchExpertClusterKernel
+
+            ptrs = self._pointers(make_ptr, cute, cutlass)
+            stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+            # Package cross_layer_expert_prefetch's speculative expert prefetch
+            # (``RoutedMoe._key``).
+            spec = key[6] or (0, 0, False)
+            kern = SharedPrefetchExpertClusterKernel(
+                self._xc_ctas(rows, finalize, red),
+                debug=debug,
+                finalize=False,
+                red_out=True,
+                **routed_moe_mod.xc_kernel_args(),
+                spec_k=spec[0],
+                spec_ctas=spec[1],
+                spec_self=spec[2],
+                se_pf=True,
+            )
+            ints = [cutlass.Int64(0)] * (routed_moe_mod._NUM_XC_ADDRS + 2)
+            fn = cute.compile(kern.run, *ptrs, *ints, cutlass.Int32(1), stream)
+            if spec[0]:
+                from ....shared import dispatch
+
+                dispatch.log_capability(
+                    FEATURE,
+                    f"expert_prefetch[k={spec[0]},ctas={spec[1]},self={int(spec[2])}]",
+                    True,
+                )
+            _COMPILED[key] = fn
+            return fn
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        routing: torch.Tensor,
+        red_out: torch.Tensor,
+        se_weights: Tuple[torch.Tensor, torch.Tensor],
+        debug: int = 0,
+        debug_buffer: Optional[torch.Tensor] = None,
+    ) -> List[torch.Tensor]:
+        """Add the routed experts of ``x`` (routing from the hop) into ``red_out`` (FP32 ``[R, H]``, zero on
+        entry); ``se_weights`` = the shared expert's ``(gate_up, down)`` weights to prefetch."""
+        import cutlass
+        import cutlass.cute as cute
+        from cuda.bindings import driver as cuda
+        from cutlass.cute.runtime import make_ptr
+
+        m = self.module
+        rows = x.shape[0]
+        x = routed_moe_mod._bulk_copy_ready(x)
+        if self.schedule(rows, False, True) != "xc":
+            raise ValueError("FP32 accumulator output needs the expert-cluster schedule")
+        routed_fn = self._compiled(False, debug, rows, True)
+        ws = routed_moe_mod._workspace(self.device, self.num_ctas)
+        dbg = debug_buffer if debug_buffer is not None else ws.dummy
+        current = torch.cuda.current_stream()
+        ws.order_on(current)
+        stream = cuda.CUstream(current.cuda_stream)
+        ptrs = self._pointers(make_ptr, cute, cutlass)
+        se_w1, se_w2 = se_weights
+        addrs = [
+            x.data_ptr(),
+            routing.data_ptr(),
+            self.bias.data_ptr(),
+            m.fc31_alpha.data_ptr(),
+            m.fc31_scale_c.data_ptr(),
+            m.fc2_alpha.data_ptr(),
+            m.fc31_input_scale.data_ptr(),
+            ws.rows.data_ptr(),
+            ws.scales.data_ptr(),
+            ws.dummy.data_ptr(),
+            red_out.data_ptr(),
+            ws.tile_count.data_ptr(),
+            dbg.data_ptr(),
+            m.w3_w1_weight.data_ptr(),
+            m.w3_w1_weight_scale.data_ptr(),
+            m.w2_weight.data_ptr(),
+            m.w2_weight_scale.data_ptr(),
+        ]
+        # cross_layer_expert_prefetch's row of predicted experts (the kernel ignores it unless
+        # cross_layer_expert_prefetch is compiled in).
+        spec = self._spec()
+        addrs.append(spec[0].data_ptr() if spec is not None else ws.dummy.data_ptr())
+        addrs += [se_w1.data_ptr(), se_w2.data_ptr()]
+        assert len(addrs) == routed_moe_mod._NUM_XC_ADDRS + 2
+        routed_fn(*ptrs, *[cutlass.Int64(a) for a in addrs], cutlass.Int32(rows), stream)
+        return [red_out]
+
+
+def _engaged_runner(module, routing_bias, x, router_logits, routing, shared):
+    """``(runner, se_weights)`` for an FP32-accumulator call, or ``(None, None)`` (host-only decision)."""
+    from ....shared import dispatch
+
+    reason = hop_routed_mod.routing_reason(routing, x)
+    ok = not reason
+    se = None
+    if ok:
+        se, reason = shared_weights(shared)
+        ok = se is not None
+    if ok:
+        ok, reason = routed_moe_mod.eligibility(module, x, router_logits)
+    owner = routed_moe_mod.resolve_expert_owner(module)[0] if ok else None
+    if ok:
+        rm = getattr(owner, "routing_method", None)
+        native_bias = getattr(rm, "e_score_correction_bias", None)
+        if (
+            routing_bias is None
+            or routing_bias.dtype != torch.bfloat16
+            or tuple(routing_bias.shape) != (routed_moe_mod.NUM_EXPERTS,)
+            or not routing_bias.is_contiguous()
+            or routing_bias.device != x.device
+        ):
+            ok, reason = False, "router bias is not a contiguous BF16 [256] on the device"
+        elif native_bias is not None and native_bias.data_ptr() != routing_bias.data_ptr():
+            ok, reason = False, "router bias is not the routing method's bias"
+    runner = getattr(owner, "_shared_prefetch_runner", None) if ok else None
+    if ok and runner is not None and runner.bias.data_ptr() != routing_bias.data_ptr():
+        runner = None
+    if ok and runner is None:
+        runner = SharedPrefetchRoutedMoe(owner, routing_bias)
+    rows = x.shape[0]
+    if ok and runner.schedule(rows, False, True) != "xc":
+        ok, reason = False, "routed_moe would run its grid-wide schedule"
+    if ok and torch.cuda.is_current_stream_capturing():
+        if getattr(owner, "_shared_prefetch_runner", None) is None or not runner.ready_for_capture(
+            False, rows=rows, red=True
+        ):
+            ok, reason = False, "not warmed up before CUDA graph capture"
+    grid = runner._xc_ctas(rows, False, True) if ok else 0
+    dispatch.log_capability(
+        FEATURE,
+        f"routed rows={rows} grid={grid} sepf=1 {reason or 'engaged'}",
+        ok,
+        reason,
+    )
+    if not ok:
+        return None, None
+    owner._shared_prefetch_runner = runner
+    return runner, se
+
+
+def routed_output_into_hop(
+    module,
+    routing_bias: torch.Tensor,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    routing: Optional[torch.Tensor],
+    hop,
+    shared=None,
+) -> Optional[List[torch.Tensor]]:
+    """hop_routing's ``routed_output_into_hop`` (FP32 accumulator only) through
+    ``SharedPrefetchExpertClusterKernel``, which prefetches ``shared``'s weights into L2; ``None``:
+    use hop_routing's fork."""
+    if (
+        hop is None
+        or not getattr(hop, "_routed_enabled", False)
+        or getattr(hop, "routed_mode", "") != "fp32"
+    ):
+        return None
+    rows = x.shape[0]
+    if not hop.can_run(rows):
+        return None
+    runner, se = _engaged_runner(module, routing_bias, x, router_logits, routing, shared)
+    if runner is None:
+        return None
+    return runner(x, routing, hop.routed_acc_buffer(rows), se)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py
new file mode 100644
index 000000000000..7abbfb3c7d0d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py
@@ -0,0 +1,318 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / tp8.
+
+GLM-5.2, modelopt NVFP4 export: 78 layers, hidden 6144, 64 query heads, MLA with
+DSA sparse attention (index top-k 2048), layers 0-2 dense and 3-77 MoE (256
+routed experts at top-8 plus a shared expert). Layer 78 is the checkpoint's bf16
+MTP module, driven by the one-model MTP-Eagle loop at five draft tokens. `tp8`
+is the attention heads and the experts tensor-parallel over eight ranks.
+
+The shell is the built-in ``GlmMoeDsaForCausalLM``: the same module tree, the
+same weights and the same engine preferences (KV cache manager, KV transceiver,
+sparse-attention defaults), so a built-in context server and this target's
+generation server exchange the same KV and indexer-K layout. Construction runs,
+in order:
+
+1. the rank-uniform check of the knobs (``dispatch.check_rank_uniform``), first,
+   so it runs before every collective the construction below makes;
+2. the dispatch plan (``dispatch.resolve``), before the built-in model is built,
+   so a knob value no rank can parse fails early;
+3. the plan is installed (``dispatch.install``): from here on the fast paths read
+   it;
+4. the built-in model, then the core: its text model adopted as
+   ``GlmMoeDsaTp8ModelingV2``, whose forward hands each step to
+   ``PrefillTarget`` or ``DecodeTarget``;
+5. ``adopt_model`` (``shared.layers``), which re-classes the decoder layers,
+   the draft layer, their attention modules (``shared.attention``) and their MoE
+   modules as the family's layer classes and attaches the fast-path objects the
+   plan enables to them, then re-classes the speculative worker
+   (``shared.spec_worker``). The attention backends build their FMHA libraries
+   after construction, in ``create_weights``, each with the libraries adoption
+   declared on it first;
+6. the rank-uniform check of that state (``dispatch.check_rank_uniform_adopted``).
+
+The step-loop hooks the target declares to the engine (``runtime_hooks``,
+``shared.runtime.make_runtime_hooks``) are built later, on the engine's first
+read, after the model is loaded.
+
+The rank-uniform check of step 1 also covers the shared switches the kernels
+read and the modeling_v2 mode (``dispatch.rank_uniform_inputs``).
+
+Every step runs under the phase of the target that takes it
+(``dispatch.step``): the shell's forward enters it around the built-in forward,
+so the model, the speculative worker and its draft loop all run in it. A
+prefill step runs the built-in decoder forward, the built-in layers and the
+built-in draft loop; a decode step runs ``DecodeTarget``'s loop, the layers'
+decode composition and the worker's draft loop.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import resolved_mode
+from tensorrt_llm._torch.attention.attention import maybe_allgather_for_helix_cp
+from tensorrt_llm._torch.attention.backends import AttentionMetadata
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
+from tensorrt_llm._torch.speculative import SpecMetadata
+
+from ..shared import dispatch, runtime
+from ..shared.layers import adopt_model, adopted_state
+from . import weights as _weights
+
+if TYPE_CHECKING:
+    from ..shared.kernels.decode_prologue import DecodePrologueHook
+
+
+# Every fast path of the family, with the family's tables; the routed-MoE paths run from this
+# target's own kernel packages.
+TARGET_PATHS = dispatch.TargetPaths(
+    kernel_modules={
+        path: f"{__package__}.kernels.{package}"
+        for path, package in (
+            ("routed_moe", "routed_moe"),
+            ("routed_moe_into_hop", "routed_moe"),
+            ("hop_routing", "hop_routing"),
+            ("shared_expert_prefetch", "shared_expert_prefetch"),
+            ("draft_routed_moe", "draft_routed_moe"),
+            ("cross_layer_expert_prefetch", "cross_layer_expert_prefetch"),
+        )
+    }
+)
+
+
+class GlmMoeDsaTp8ModelingV2(DeepseekV3Model):
+    """The text model: the built-in ``DeepseekV3Model``, whose forward picks a target per step.
+
+    Never constructed directly: ``adopt`` re-classes the built-in model the shell
+    has already built, so the module tree, the parameters and their names stay
+    the built-in model's. The class adds methods and plain attributes only.
+    """
+
+    _prefill: "PrefillTarget"
+    _decode: "DecodeTarget"
+    # The decode prologue the step loop runs, set when the target's runtime hooks are built;
+    # the decode target takes the embedding rows of its launch.
+    decode_prologue: Optional["DecodePrologueHook"] = None
+
+    @classmethod
+    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaTp8ModelingV2":
+        """Re-class the built-in text model as this core and build its two targets.
+
+        Raises:
+            TypeError: if ``model`` is not exactly the built-in ``DeepseekV3Model``.
+        """
+        if type(model) is not DeepseekV3Model:
+            raise TypeError(
+                f"{cls.__name__} adopts the built-in DeepseekV3Model, got {type(model).__name__}"
+            )
+        model.__class__ = cls
+        model._prefill = PrefillTarget(model)
+        model._decode = DecodeTarget(model)
+        return model
+
+    def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
+        """Which target runs this step.
+
+        One host-side integer read and nothing else: a decode-only CUDA graph is a
+        per-capture constant environment. A mixed batch -- context and generation
+        rows in one step -- goes to prefill, so prefill is the general case and
+        decode the specialization that may assume it has no context rows.
+        """
+        return self._decode if attn_metadata.num_contexts == 0 else self._prefill
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self._select_target(attn_metadata).forward(
+            attn_metadata,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+
+class PrefillTarget:
+    """The general case: context rows, and possibly generation rows beside them.
+
+    Runs the built-in decoder forward.
+    """
+
+    phase = dispatch.Phase.PREFILL
+
+    def __init__(self, core: GlmMoeDsaTp8ModelingV2) -> None:
+        self.core = core
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return DeepseekV3Model.forward(
+            self.core,
+            attn_metadata,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+
+class DecodeTarget:
+    """The specialization: routed to only when there are no context rows.
+
+    Runs this target's decoder loop.
+    """
+
+    phase = dispatch.Phase.DECODE
+
+    def __init__(self, core: GlmMoeDsaTp8ModelingV2) -> None:
+        self.core = core
+
+    def forward(
+        self,
+        attn_metadata: AttentionMetadata,
+        input_ids: Optional[torch.IntTensor] = None,
+        position_ids: Optional[torch.IntTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """The decoder layers of the text model on a decode step.
+
+        The built-in loop with the decode prologue's embedding and the expert-prefetch join.
+        Each layer is called as a module, so its forward hooks run, and runs its decode
+        composition under the step's phase. The draft layer is not run here; the last decoder
+        layer applies the final norm.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        prologue = self.core.decode_prologue
+        if inputs_embeds is None and prologue is not None:
+            # The rows the decode prologue launch embedded (None unless that launch ran for
+            # exactly these ids and this table).
+            inputs_embeds = prologue.take_embedding(
+                attn_metadata, input_ids, self.core.embed_tokens
+            )
+        if inputs_embeds is None:
+            inputs_embeds = self.core.embed_tokens(input_ids)
+
+        hidden_states = inputs_embeds
+        residual = None
+
+        for idx, decoder_layer in enumerate(self.core.layers[: self.core.num_hidden_layers]):
+            hidden_states, residual = decoder_layer(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                residual=residual,
+                spec_metadata=spec_metadata,
+            )
+        if dispatch.is_enabled("cross_layer_expert_prefetch"):
+            # The expert-prefetch stream rejoins before the drafts.
+            dispatch.get_impl("cross_layer_expert_prefetch", "join")()
+
+        hidden_states = maybe_allgather_for_helix_cp(
+            hidden_states, attn_metadata, self.core.mapping_with_cp
+        )
+        return hidden_states
+
+
+class ModelingV2Glm52Nvfp4Sm103Tp8(GlmMoeDsaForCausalLM):
+    """GLM-5.2-NVFP4 on sm_103 at tp8: the built-in causal LM with this target's core and dispatch."""
+
+    def __init__(self, model_config: ModelConfig):
+        world_size = model_config.mapping.world_size
+        knobs = dispatch.read_knobs()
+        mode = resolved_mode(type(self))
+        dispatch.check_rank_uniform(dispatch.rank_uniform_inputs(knobs, mode), world_size)
+        plan = dispatch.resolve(mode, dict(knobs), TARGET_PATHS)
+        dispatch.install(plan)
+        super().__init__(model_config)
+        GlmMoeDsaTp8ModelingV2.adopt(self.model)
+        self.plan = plan
+        adopt_model(self, self.plan)
+        dispatch.check_rank_uniform_adopted(adopted_state(self), world_size)
+
+    @property
+    def runtime_hooks(self) -> RuntimeHooks:
+        """The step-loop hooks this target declares, built on the engine's first read.
+
+        Built after the model is loaded: a hook may hold pinned host memory, which model
+        construction (under meta-device initialization) must not allocate.
+
+        Raises:
+            RuntimeError: if building the hooks raises ``AttributeError``, which
+                ``runtime_hooks_of`` would otherwise read as "no hooks declared".
+        """
+        hooks = self.__dict__.get("_runtime_hooks")
+        if hooks is None:
+            try:
+                hooks = runtime.make_runtime_hooks(self)
+            except AttributeError as exc:
+                raise RuntimeError(
+                    f"building the step-loop hooks of {type(self).__name__} failed: {exc}"
+                ) from exc
+            self._runtime_hooks = hooks
+        return hooks
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.IntTensor | None = None,
+        position_ids: torch.IntTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        return_context_logits: bool = False,
+        spec_metadata=None,
+        resource_manager=None,
+        **kwargs,
+    ):
+        """The built-in forward under the phase of the target that takes this step.
+
+        The scope covers the text model, the speculative worker and its draft loop.
+        """
+        with dispatch.step(self.model._select_target(attn_metadata).phase):
+            return super().forward(
+                attn_metadata,
+                input_ids=input_ids,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                return_context_logits=return_context_logits,
+                spec_metadata=spec_metadata,
+                resource_manager=resource_manager,
+                **kwargs,
+            )
+
+    def setup_aliases(self) -> None:
+        """The built-in aliases, then the cross-layer expert prefetch links."""
+        super().setup_aliases()
+        # Bind each decoder MoE layer to the next MoE layer's gate and experts (structural
+        # references only).
+        if dispatch.is_enabled("cross_layer_expert_prefetch"):
+            dispatch.get_impl("cross_layer_expert_prefetch", "link_layers")(
+                self.model.layers, self.config.num_hidden_layers
+            )
+
+    def load_weights(self, weights):
+        _weights.load(self, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py
new file mode 100644
index 000000000000..b77cd3c04a62
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py
@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The weights this target expects: the GLM-5.2 NVFP4 checkpoint as the built-in model loads it.
+
+The target's module tree is the built-in model's, so the built-in loader maps
+every checkpoint tensor, including the bf16 MTP layer and the NVFP4 scales.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+
+if TYPE_CHECKING:
+    from .modeling import ModelingV2Glm52Nvfp4Sm103Tp8
+
+
+def load(model: "ModelingV2Glm52Nvfp4Sm103Tp8", weights: Any) -> None:
+    """Load the checkpoint ``weights`` into ``model`` with the built-in loader."""
+    GlmMoeDsaForCausalLM.load_weights(model, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py
new file mode 100644
index 000000000000..5faf1a030aa2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py
@@ -0,0 +1,290 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Where a GlmMoeDsaForCausalLM config lands. Read this file and you know.
+
+``route`` is one forward-reading decision: the facts that fix which code runs,
+compared against the deployment the target below was written for. When it does
+not fit, ``route`` raises ``NoTarget`` stating those facts. In ``auto`` the
+engine then uses the built-in GlmMoeDsa implementation; in ``require`` it raises
+that error.
+
+The three targets are decode specializations of one checkpoint on one GPU
+architecture, and differ in how the attention and the experts are split. Their
+fast paths are built for the checkpoint's operand shapes and layer pattern, their
+ranks and the rows per request (the target token and its draft tokens), so each of
+those facts is a routing criterion rather than a runtime branch:
+
+* ``glm_5_2_nvfp4__sm_103__tp8``: attention and experts tensor-parallel over
+  eight ranks, five draft tokens;
+* ``glm_5_2_nvfp4__sm_103__tep``: attention tensor-parallel over four or eight
+  ranks and the experts expert-parallel over the same ranks, five draft tokens;
+* ``glm_5_2_nvfp4__sm_103__dep``: attention data parallelism and the experts
+  expert-parallel over the same ranks, at the rank counts, LM-head splits and
+  draft lengths of ``_DEP``. Two more facts fix what its fast paths serve, the
+  KV-cache dtype and the MoE backend; ``route`` does not read them, so the
+  target checks them when it is built.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Tuple
+
+from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+from tensorrt_llm.quantization.modelopt_config import (
+    canonicalize_quant_algo,
+    is_modelopt_quant_config,
+    read_modelopt_quant_config,
+)
+
+from ..._router_index import NoTarget
+
+# The one GPU architecture the fast paths are built and validated for.
+_SM = (10, 3)
+
+# The checkpoint's config shape; see the gpt_oss routing module for what
+# shape-sniffing does and does not pin. The head count fixes the per-rank
+# attention operands at tp8, and index_topk fixes the sparse-attention selection
+# width the decode kernels are built for. The last four fields fix the layer
+# pattern: the dense layers before the first MoE layer, and which layers run
+# their own indexer and which reuse an earlier layer's selection (an explicit
+# index_topk_pattern overrides the frequency and offset, so it must be absent).
+_SHAPE_FIELDS = (
+    "num_hidden_layers",
+    "hidden_size",
+    "n_routed_experts",
+    "q_lora_rank",
+    "num_attention_heads",
+    "index_topk",
+    "first_k_dense_replace",
+    "index_topk_freq",
+    "index_skip_topk_offset",
+    "index_topk_pattern",
+)
+_SHAPE = (78, 6144, 256, 2048, 64, 2048, 3, 4, 3, None)
+
+# Ranks of tensor-parallel attention and tensor-parallel experts.
+_TP = 8
+
+# Ranks of tensor-parallel attention with the experts expert-parallel over the same ranks.
+_TEP = (4, 8)
+
+# Draft tokens per request. The decode kernels take the target token plus this
+# many draft tokens per request in one step.
+_DRAFT_LEN = 5
+
+# Attention data parallelism with the experts expert-parallel over the same ranks: per
+# (ranks, LM head tensor-parallel in attention DP), the draft tokens per request.
+_DEP = {(16, False): 3, (8, False): 3, (8, True): 5}
+
+
+def _nvfp4_with_bf16_mtp_layer(pretrained_config: Any) -> Tuple[bool, str]:
+    """Whether the checkpoint is NVFP4 with a bf16 MTP layer, and its quantization in words.
+
+    Reads the inline ``quantization_config`` of config.json, and parses only the
+    modelopt shape, with the modelopt normalizer (which raises only
+    ``ValueError``). The engine itself prefers ``hf_quant_config.json`` when the
+    checkpoint has one, which ``route`` does not see, and warns when the inline
+    copy diverges from it. Any other declaration is named and declined, never
+    parsed, so a config this target does not serve cannot fail here. The export
+    keeps the whole MTP layer (layer ``num_hidden_layers``) in bf16; the draft
+    fast paths read bf16 weights there.
+    """
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+
+    declared = getattr(pretrained_config, "quantization_config", None)
+    if declared is None:
+        return False, "none"
+    if not isinstance(declared, dict) and callable(getattr(declared, "to_dict", None)):
+        declared = declared.to_dict()
+    if not isinstance(declared, dict):
+        return False, f"declared as a {type(declared).__name__}, not a dict"
+    if not is_modelopt_quant_config(declared):
+        return False, f"quant_method {declared.get('quant_method')!r} (not a modelopt export)"
+    try:
+        normalized = read_modelopt_quant_config(declared)
+        quant_algo = QuantAlgo(canonicalize_quant_algo(normalized.get("quant_algo")))
+    except ValueError as why:
+        return False, f"a modelopt config without a readable global algorithm ({why})"
+    excluded = normalized.get("exclude_modules")
+    if excluded is not None and not (
+        isinstance(excluded, (list, tuple)) and all(isinstance(m, str) for m in excluded)
+    ):
+        return False, f"modelopt {quant_algo.name} with unreadable excluded modules"
+    mtp_layer = f"model.layers.{pretrained_config.num_hidden_layers}"
+    unquantized = QuantConfig(exclude_modules=excluded).is_module_excluded_from_quantization(
+        mtp_layer
+    )
+    words = f"{quant_algo.name} with the MTP layer {'unquantized' if unquantized else 'quantized'}"
+    return quant_algo == QuantAlgo.NVFP4 and unquantized, words
+
+
+def _tp8(m: Any) -> bool:
+    """Attention heads and experts both tensor-parallel over eight ranks, with every
+    expert on every rank and no attention data parallelism.
+
+    The split decides which weights each rank loads and into which shapes, so it
+    selects a target rather than a runtime branch.
+    """
+    return (
+        m.world_size == _TP
+        and m.tp_size == _TP
+        and m.pp_size == 1
+        and m.cp_size == 1
+        and m.moe_tp_size == _TP
+        and m.moe_ep_size == 1
+        and not m.enable_attention_dp
+    )
+
+
+def _tep(m: Any) -> bool:
+    """Attention heads tensor-parallel over four or eight ranks and the experts expert-parallel
+    over the same ranks (each rank holds whole experts), with no attention data parallelism.
+
+    The split decides which weights each rank loads and into which shapes, so it
+    selects a target rather than a runtime branch.
+    """
+    return (
+        m.world_size == m.tp_size
+        and m.tp_size in _TEP
+        and m.pp_size == 1
+        and m.cp_size == 1
+        and m.moe_tp_size == 1
+        and m.moe_ep_size == m.tp_size
+        and not m.enable_attention_dp
+    )
+
+
+def _dep(m: Any) -> bool:
+    """Attention data parallelism (each rank runs the whole attention on its own requests)
+    and the experts expert-parallel over the same ranks, at a rank count and LM-head split
+    of ``_DEP``.
+
+    The split decides which weights each rank loads and into which shapes, so it
+    selects a target rather than a runtime branch.
+    """
+    return (
+        m.world_size == m.tp_size
+        and m.pp_size == 1
+        and m.cp_size == 1
+        and m.moe_tp_size == 1
+        and m.moe_ep_size == m.world_size
+        and m.enable_attention_dp
+        and (m.world_size, bool(m.enable_lm_head_tp_in_adp)) in _DEP
+    )
+
+
+def _mtp_eagle_linear(spec_config: Optional[Any], draft_len: int) -> bool:
+    """The checkpoint's MTP layer driven by the one-model MTP-Eagle loop, as a
+    linear chain of exactly ``draft_len`` draft tokens.
+
+    No speculation and other draft lengths do not route here: the decode
+    kernels are built for the target's rows per request. Nor do the
+    suffix-automaton draft enhancer and relaxed acceptance: both are fixed for
+    the engine's lifetime, and the draft fast paths implement neither.
+    """
+    return (
+        isinstance(spec_config, MTPDecodingConfig)
+        and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+        and spec_config.max_draft_len == draft_len
+        and spec_config.max_total_draft_tokens == draft_len
+        and not spec_config.use_dynamic_tree
+        and spec_config.sa_config is None
+        and not spec_config.use_relaxed_acceptance_for_thinking
+    )
+
+
+def _speculation(spec_config: Optional[Any]) -> str:
+    if spec_config is None:
+        return "off"
+    words = (
+        f"{spec_config.spec_dec_mode.name} max_draft_len={spec_config.max_draft_len} "
+        f"max_total_draft_tokens={spec_config.max_total_draft_tokens}"
+    )
+    if getattr(spec_config, "use_dynamic_tree", False):
+        words += " dynamic tree"
+    if getattr(spec_config, "sa_config", None) is not None:
+        words += " with the suffix-automaton enhancer"
+    if getattr(spec_config, "use_relaxed_acceptance_for_thinking", False):
+        words += " with relaxed acceptance"
+    return words
+
+
+def route(
+    pretrained_config: Any,
+    mapping: Any,
+    compute_capability: Tuple[int, int],
+    spec_config: Optional[Any],
+) -> type:
+    """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
+
+    The target is imported where it is returned, so declining costs no import
+    of modeling code.
+    """
+    # A GlmMoeDsa config without one of these fields is another assembly: it reads as
+    # None and the shape does not match.
+    shape = tuple(getattr(pretrained_config, name, None) for name in _SHAPE_FIELDS)
+
+    # GLM-5.2-NVFP4 (bf16 MTP layer) on sm103, 8 ranks of tensor-parallel attention
+    # and experts, one-model MTP-Eagle at five draft tokens. The quantization is
+    # read last, and only for this deployment.
+    quantization = "not read"
+    if (
+        compute_capability == _SM
+        and shape == _SHAPE
+        and _tp8(mapping)
+        and _mtp_eagle_linear(spec_config, _DRAFT_LEN)
+    ):
+        served, quantization = _nvfp4_with_bf16_mtp_layer(pretrained_config)
+        if served:
+            from .glm_5_2_nvfp4__sm_103__tp8.modeling import ModelingV2Glm52Nvfp4Sm103Tp8
+
+            return ModelingV2Glm52Nvfp4Sm103Tp8
+
+    # The same checkpoint, architecture and speculation, with 4 or 8 ranks of
+    # tensor-parallel attention and the experts expert-parallel over them.
+    if (
+        compute_capability == _SM
+        and shape == _SHAPE
+        and _tep(mapping)
+        and _mtp_eagle_linear(spec_config, _DRAFT_LEN)
+    ):
+        served, quantization = _nvfp4_with_bf16_mtp_layer(pretrained_config)
+        if served:
+            from .glm_5_2_nvfp4__sm_103__tep.modeling import ModelingV2Glm52Nvfp4Sm103Tep
+
+            return ModelingV2Glm52Nvfp4Sm103Tep
+
+    # The same checkpoint and architecture with attention data parallelism and the experts
+    # expert-parallel over the same ranks, at the draft length of the rank count and LM-head
+    # split.
+    if (
+        compute_capability == _SM
+        and shape == _SHAPE
+        and _dep(mapping)
+        and _mtp_eagle_linear(
+            spec_config, _DEP[(mapping.world_size, bool(mapping.enable_lm_head_tp_in_adp))]
+        )
+    ):
+        served, quantization = _nvfp4_with_bf16_mtp_layer(pretrained_config)
+        if served:
+            from .glm_5_2_nvfp4__sm_103__dep.modeling import ModelingV2Glm52Nvfp4Sm103Dep
+
+            return ModelingV2Glm52Nvfp4Sm103Dep
+
+    raise NoTarget(
+        f"compute capability {compute_capability}, config shape {shape}, quantization "
+        f"{quantization}, parallel ws={mapping.world_size} tp={mapping.tp_size} "
+        f"pp={mapping.pp_size} cp={mapping.cp_size} moe_tp={mapping.moe_tp_size} "
+        f"moe_ep={mapping.moe_ep_size} attention_dp={mapping.enable_attention_dp} "
+        f"lm_head_tp_in_adp={mapping.enable_lm_head_tp_in_adp}, "
+        f"speculative decoding {_speculation(spec_config)}; the glm_moe_dsa targets are "
+        "the GLM-5.2 modelopt NVFP4 checkpoint (bf16 MTP layer) on sm103, with one-model "
+        "MTP-Eagle as a linear chain, without the suffix-automaton enhancer or relaxed "
+        "acceptance, at tp8 (8 ranks of tensor-parallel attention and experts) or tep (4 or 8 "
+        "ranks of tensor-parallel attention with the experts expert-parallel over the same "
+        "ranks), both without attention DP and at max_draft_len 5, or at dep (attention DP "
+        "with the experts expert-parallel over the same ranks: 16 ranks at max_draft_len 3, "
+        "8 ranks at max_draft_len 3, or 8 ranks with the LM head tensor-parallel in attention "
+        "DP at max_draft_len 5)"
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/__init__.py
new file mode 100644
index 000000000000..506daf40bc38
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/__init__.py
@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Code two or more glm_moe_dsa targets call: the fast-path plan, the attention adoption, the
+speculative worker, the step-loop hooks and the kernel packages they share."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention.py
new file mode 100644
index 000000000000..a61b18ea447a
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention.py
@@ -0,0 +1,569 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The attention modules of the glm_moe_dsa targets: the MLA, its DSA hooks, its DSA backend and
+its indexer.
+
+``GlmMLA.adopt`` re-classes a layer's built-in attention modules in place -- the
+``DeepseekV32Attention`` as ``GlmMLA``, its ``DSAHooks`` as ``GlmDsaHooks``, its
+``DSATrtllmAttention`` as ``GlmDsaBackend`` and its ``Indexer`` as ``GlmIndexer`` -- and attaches
+the attention fast paths the installed plan enables: the attention output tail and, with
+``sparse_mla_decode``, that FMHA library ahead of the registry's on the layer's attention
+backends. The backends build their FMHA libraries later, in ``create_weights``, so the library is
+created when and where the built-in model creates its own. The classes add methods and
+class-level defaults only, so the module tree and the weights stay the built-in model's.
+
+Each forward override runs the family's decode composition (``attention_forward``) when the
+step is a decode step (``dispatch.Phase.DECODE``) and the built-in method otherwise.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+
+import tensorrt_llm
+import tensorrt_llm.bindings
+from tensorrt_llm._torch.attention.attention import _helix_cp_allgather_input
+from tensorrt_llm._torch.attention.backends.interface import (
+    AttentionForwardArgs,
+    AttentionInputType,
+    AttentionMetadata,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata
+from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
+from tensorrt_llm._torch.distributed import AllReduceParams
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+from tensorrt_llm._torch.utils import Fp4QuantizedTensor
+
+from . import attention_forward, dispatch
+from .kernels.indexer_topk import region
+
+
+def reclass(module, builtin: type, cls: type) -> None:
+    """Re-class ``module``, which must be exactly the built-in class ``builtin``, as ``cls``.
+
+    Raises:
+        TypeError: if ``module`` is not exactly ``builtin`` (a subclass of it is refused too).
+    """
+    if type(module) is not builtin:
+        raise TypeError(
+            f"{cls.__name__} adopts the built-in {builtin.__name__}, got {type(module).__name__}"
+        )
+    module.__class__ = cls
+
+
+def _deferred_norm_view(hidden_states) -> Optional[torch.Tensor]:
+    """The tensor view of a deferred-norm packet, else None."""
+    if dispatch.is_enabled("deferred_norm") and dispatch.get_impl("deferred_norm", "is_deferred")(
+        hidden_states
+    ):
+        return hidden_states.residual
+    return None
+
+
+class GlmMLA(DeepseekV32Attention):
+    """A layer's MLA: the built-in attention, with the family's decode composition."""
+
+    # The attention output tail (``attention_output_tail``) and whether it also prefetches its
+    # weights into L2 (``attention_output_tail_prefetch``); None keeps the built-in output.
+    _attn_tail = None
+    _tail_prefetch = False
+
+    @classmethod
+    def adopt(cls, mla: DeepseekV32Attention) -> "GlmMLA":
+        """Re-class a layer's built-in attention modules and attach its attention fast paths.
+
+        Raises:
+            TypeError: if ``mla``, its DSA hooks, its DSA backend or its indexer is not exactly
+                the built-in class.
+        """
+        reclass(mla, DeepseekV32Attention, cls)
+        mla._adopt_attention_tail()
+        GlmDsaHooks.adopt(mla.sparse_attn_hooks)
+        GlmDsaBackend.adopt(mla.mqa)
+        if mla.mqa.indexer is not None:
+            GlmIndexer.adopt(mla.mqa.indexer)
+        # The sparse-MLA decode library goes ahead of the registry's libraries on both attention
+        # backends, which build their FMHA libraries in ``create_weights``.
+        if dispatch.is_enabled("sparse_mla_decode"):
+            libraries = (dispatch.get_impl("sparse_mla_decode", "fmha_lib"),)
+            for backend in (mla.mha, mla.mqa):
+                if backend is not None:
+                    backend.prepended_fmha_lib_classes = libraries
+        return mla
+
+    def _adopt_attention_tail(self) -> None:
+        """The attention output tail: the fused generation W_UV and the local o_proj."""
+        # attention_output_tail_prefetch uses the same tail object to stage the tail weights
+        # into L2 off the critical path while the attention runs.
+        self._tail_prefetch = dispatch.is_enabled("attention_output_tail_prefetch")
+        self._attn_tail = None
+        # attention_output_into_hop publishes the tail's partial into the post-attention hop's
+        # transport.
+        if (
+            dispatch.is_enabled("attention_output_tail")
+            or self._tail_prefetch
+            or dispatch.is_enabled("attention_output_into_hop")
+        ):
+            self._attn_tail = dispatch.get_impl("attention_output_tail", "make_attn_tail")(self)
+            # oproj_weight_staging: the tail's o_proj stages the W_o rows it streams after its
+            # programmatic-dependent-launch wait before that wait.
+            if dispatch.is_enabled("oproj_weight_staging"):
+                self._attn_tail = dispatch.get_impl("oproj_weight_staging", "wrap_attn_tail")(
+                    self, self._attn_tail
+                )
+
+    def forward_absorption_generation(
+        self,
+        q: torch.Tensor,
+        compressed_kv: torch.Tensor,
+        k_pe: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        output: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None,
+        latent_cache: Optional[torch.Tensor] = None,
+        absorbed_fused_q: Optional[torch.Tensor] = None,
+        deferred_indexer_k=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward_absorption_generation(
+                q=q,
+                compressed_kv=compressed_kv,
+                k_pe=k_pe,
+                attn_metadata=attn_metadata,
+                output=output,
+                position_ids=position_ids,
+                latent_cache=latent_cache,
+                **kwargs,
+            )
+        return attention_forward.forward_absorption_generation(
+            self,
+            q,
+            compressed_kv,
+            k_pe,
+            attn_metadata,
+            output,
+            position_ids=position_ids,
+            latent_cache=latent_cache,
+            absorbed_fused_q=absorbed_fused_q,
+            deferred_indexer_k=deferred_indexer_k,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        position_ids: Optional[torch.Tensor],
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        all_reduce_params: Optional[AllReduceParams] = None,
+        latent_cache_gen: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                all_reduce_params=all_reduce_params,
+                latent_cache_gen=latent_cache_gen,
+            )
+        hidden_states = _helix_cp_allgather_input(
+            hidden_states, attn_metadata, self.mapping, self.layer_idx
+        )
+
+        # Unquantized view of the module input, used by create_mla_outputs and
+        # by _apply_output_gate.
+        output_hidden_states = hidden_states
+        deferred_view = _deferred_norm_view(hidden_states)
+        if deferred_view is not None:
+            # A deferred-norm packet: only the DSA query path consumes it; the outputs are
+            # allocated on its tensor view.
+            output_hidden_states = deferred_view
+        elif isinstance(hidden_states, Fp4QuantizedTensor):
+            assert hidden_states.unquantized_hidden_states is not None, (
+                "MLA.forward received an Fp4QuantizedTensor without a "
+                "unquantized_hidden_states view"
+            )
+            output_hidden_states = hidden_states.unquantized_hidden_states
+
+        # The attention tail decides here, before the attention op, so the W_UV step inside the
+        # op and the projection below see one decision. arm() returns False under
+        # torch.compile and inside breakable CUDA graphs, where the op runs apart from the
+        # traced or captured projection.
+        attn_tail = getattr(self, "_attn_tail", None)
+        tail_armed = attn_tail is not None and attn_tail.arm(self, attn_metadata)
+
+        if self.register_to_config:
+            attn_output = [
+                torch.ops.trtllm.create_mla_outputs(output_hidden_states, self.layer_idx_str)
+            ]
+            self._forward_custom_op(
+                hidden_states,
+                position_ids,
+                attn_output,
+                latent_cache_gen,
+            )
+        else:
+            attn_output = self._create_outputs(
+                hidden_states if deferred_view is None else deferred_view, attn_metadata
+            )
+            self.forward_impl(
+                position_ids,
+                hidden_states,
+                attn_metadata,
+                attn_output=attn_output,
+                latent_cache_gen=latent_cache_gen,
+            )
+
+        attn_output[0] = self._apply_output_gate(output_hidden_states, attn_output[0])
+        if tail_armed:
+            partial = attn_tail.project(self, attn_output[0])
+            if partial is not None:
+                # The tail returns o_proj's rank-local GEMM output; reduce it as the
+                # ROW-parallel Linear.forward does for this o_proj (the tail takes only an
+                # o_proj without bias, LoRA or fused GEMM+all-reduce).
+                if self.o_proj.reduce_output:
+                    partial = self.o_proj.all_reduce(partial, all_reduce_params=all_reduce_params)
+                return partial
+        return self._project_output(attn_output, position_ids, attn_metadata, all_reduce_params)
+
+
+class GlmDsaHooks(DSAHooks):
+    """A layer's DSA hooks: the built-in sparse attention, with the family's decode forward."""
+
+    @classmethod
+    def adopt(cls, hooks: DSAHooks) -> "GlmDsaHooks":
+        """Re-class a layer's built-in DSA hooks.
+
+        Raises:
+            TypeError: if ``hooks`` is not exactly the built-in ``DSAHooks``.
+        """
+        reclass(hooks, DSAHooks, cls)
+        return hooks
+
+    def forward(
+        self,
+        mla,
+        position_ids: Optional[torch.Tensor],
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        attn_output: list[torch.Tensor],
+    ) -> None:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward(
+                mla=mla,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                attn_output=attn_output,
+            )
+        attention_forward.forward_sparse_attn(
+            mla, position_ids, hidden_states, attn_metadata, attn_output
+        )
+
+    def forward_custom_op(
+        self,
+        mla,
+        hidden_states: torch.Tensor,
+        position_ids: Optional[torch.Tensor],
+        attn_output: list[torch.Tensor],
+        latent_cache_gen: Optional[torch.Tensor],
+    ) -> bool:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward_custom_op(
+                mla=mla,
+                hidden_states=hidden_states,
+                position_ids=position_ids,
+                attn_output=attn_output,
+                latent_cache_gen=latent_cache_gen,
+            )
+        attention_forward.forward_sparse_attn_custom_op(
+            mla,
+            hidden_states,
+            position_ids,
+            attn_output,
+            latent_cache_gen,
+        )
+        return True
+
+
+class GlmDsaBackend(DSATrtllmAttention):
+    """A layer's DSA attention backend, with the family's top-k bridge on decode steps."""
+
+    @classmethod
+    def adopt(cls, backend: DSATrtllmAttention) -> "GlmDsaBackend":
+        """Re-class a layer's built-in DSA attention backend.
+
+        Raises:
+            TypeError: if ``backend`` is not exactly the built-in ``DSATrtllmAttention``.
+        """
+        reclass(backend, DSATrtllmAttention, cls)
+        return backend
+
+    def sparse_attn_predict(
+        self,
+        q: torch.Tensor,
+        k: Optional[torch.Tensor],
+        metadata: DSAtrtllmAttentionMetadata,
+        forward_args: AttentionForwardArgs,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """Run the DSA indexer and convert its local TopK to paged indices."""
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().sparse_attn_predict(
+                q=q, k=k, metadata=metadata, forward_args=forward_args
+            )
+        is_generation = forward_args.attention_input_type == AttentionInputType.generation_only
+        sparse_backend_args = forward_args.sparse_backend_args
+        phase_start = metadata.num_ctx_tokens if is_generation else 0
+        phase_end = metadata.num_tokens if is_generation else metadata.num_ctx_tokens
+
+        shared_topk_indices = metadata.shared_topk_indices
+        # rope_kv_append: one fused bridge kernel performs the native top-k copy of this call
+        # and the local -> global remap. With cross_rank_topk_global_emit the cross-rank top-k
+        # writes them and records it on its packet.
+        bridge = None
+        bridge_path = (
+            "cross_rank_topk_global_emit"
+            if dispatch.is_enabled("cross_rank_topk_global_emit")
+            else "rope_kv_append"
+        )
+        if dispatch.is_enabled("rope_kv_append"):
+            bridge = dispatch.get_impl(bridge_path, "new_bridge_packet")(metadata, is_generation)
+        if self.indexer is None:
+            topk_indices = shared_topk_indices[phase_start:phase_end]
+        else:
+            indexer_intermediates = sparse_backend_args.indexer_intermediates
+            # The indexer's top-k region reads the bridge of this call only.
+            self.indexer._topk_bridge = bridge
+            try:
+                topk_indices = self.indexer.forward_from_projected(
+                    metadata,
+                    q,
+                    indexer_intermediates,
+                    is_generation=is_generation,
+                )
+            finally:
+                self.indexer._topk_bridge = None
+            preserve_mtp_topk = metadata.in_mtp_draft_loop and self.indexer.mtp_index_share
+            if shared_topk_indices is not None and not preserve_mtp_topk:
+                if bridge is not None:
+                    dispatch.get_impl("rope_kv_append", "defer_shared_copy")(
+                        bridge, shared_topk_indices, phase_start, topk_indices
+                    )
+                else:
+                    shared_topk_indices[
+                        phase_start : phase_start + topk_indices.shape[0],
+                        : topk_indices.shape[1],
+                    ].copy_(topk_indices)
+
+        local_layer_idx = self.get_local_layer_idx(metadata)
+        kv_cache_dtype = getattr(metadata.kv_cache_manager, "dtype", None)
+        if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4 and not is_generation:
+            metadata._ensure_pool_view_cached()
+            page_index_scale, layer_offset = (
+                metadata.kv_cache_manager.get_primary_pool_page_index_params(local_layer_idx)
+            )
+            total_kv_tokens = metadata.num_ctx_mla_kv_tokens
+            if total_kv_tokens <= 0:
+                raise RuntimeError("NVFP4 DSA context has no active KV tokens")
+            head_dim = metadata.kv_cache_manager.head_dim
+            # Static sparse FMHA may prefetch up to TopK rows even when a
+            # short context has fewer valid KV tokens. The compaction itself
+            # can select at most min(total KV tokens, query/TopK pairs), so a
+            # late chunk with a long cached prefix does not need one fp8 row
+            # for every token in that prefix.
+            max_selected_tokens = min(total_kv_tokens, topk_indices.numel())
+            scratch_capacity = max(max_selected_tokens, topk_indices.shape[1])
+            # The previous layer's attention has already been enqueued on this
+            # stream. Drop its Python reference before allocating the next
+            # layer's scratch so the caching allocator can reuse that storage
+            # in stream order instead of transiently keeping two full buffers.
+            metadata.nvfp4_mla_context_fp8_scratch = None
+            scratch = torch.empty(
+                (scratch_capacity, 1, head_dim),
+                dtype=torch.float8_e4m3fn,
+                device=topk_indices.device,
+            )
+            compact_indices = torch.empty_like(topk_indices)
+            torch.ops.trtllm.nvfp4_mla_context_kv_cache_gather(
+                metadata.host_kv_cache_pool_pointers,
+                metadata.host_kv_cache_pool_mapping,
+                topk_indices,
+                metadata._cached_req_idx_ctx,
+                metadata._cached_block_table_ctx,
+                metadata.ctx_kv_indptr[: metadata.num_contexts + 1],
+                scratch,
+                compact_indices,
+                self._nvfp4_mla_kv_scale_quant_orig,
+                local_layer_idx,
+                total_kv_tokens,
+                metadata._cached_tokens_per_block,
+                page_index_scale * metadata._cached_tokens_per_block,
+                layer_offset,
+                metadata.kv_cache_manager.mla_kv_cache_residual_dim,
+                metadata._cached_num_pool_tokens,
+            )
+            metadata.nvfp4_mla_context_fp8_scratch = scratch
+            forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr()
+            return compact_indices, None
+
+        if bridge is not None:
+            topk_indices_global = dispatch.get_impl(bridge_path, "remap_topk_to_global")(
+                self, topk_indices, metadata, bridge
+            )
+        else:
+            topk_indices_global = self._remap_topk_to_global(topk_indices, metadata, is_generation)
+
+        if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4:
+            scratch = metadata.nvfp4_mla_fp8_scratch
+            compact_indices = metadata.nvfp4_mla_compact_indices
+            if scratch is None or compact_indices is None:
+                raise RuntimeError("NVFP4 MLA scratch buffers were not allocated")
+            num_rows = topk_indices_global.shape[0]
+            if num_rows > scratch.shape[0] or num_rows > compact_indices.shape[0]:
+                raise RuntimeError(
+                    "NVFP4 MLA generation rows exceed the allocated MTP "
+                    f"workspace: requested {num_rows}, capacity {scratch.shape[0]}"
+                )
+            scratch = scratch[:num_rows]
+            compact_indices = compact_indices[:num_rows]
+            torch.ops.trtllm.nvfp4_mla_kv_cache_gather(
+                metadata.host_kv_cache_pool_pointers,
+                metadata.host_kv_cache_pool_mapping,
+                topk_indices_global,
+                scratch,
+                compact_indices,
+                self._nvfp4_mla_kv_scale_quant_orig,
+                local_layer_idx,
+                metadata.kv_cache_manager.mla_kv_cache_residual_dim,
+                metadata._cached_num_pool_tokens,
+            )
+            # Static sparse MLA treats indices as offsets from kvPtr. Feed it
+            # compact offsets and the FP8 scratch base; TRTLLMGen is unchanged.
+            forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr()
+            return compact_indices, None
+
+        return topk_indices_global, None
+
+
+class GlmIndexer(Indexer):
+    """A full-indexer layer's DSA indexer, with the family's decode top-k region."""
+
+    # The top-k bridge of the current ``GlmDsaBackend.sparse_attn_predict`` call (set only for
+    # that call, on a decode step); None keeps the built-in copies.
+    _topk_bridge = None
+
+    @classmethod
+    def adopt(cls, indexer: Indexer) -> "GlmIndexer":
+        """Re-class a layer's built-in indexer.
+
+        Raises:
+            TypeError: if ``indexer`` is not exactly the built-in ``Indexer``.
+        """
+        reclass(indexer, Indexer, cls)
+        return indexer
+
+    def _copy_reused_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+        num_generations: int,
+        has_prefill: bool,
+        gvr_prior_indices: Optional[torch.Tensor],
+    ) -> None:
+        # rope_kv_append: the fused bridge reads the shared rows and writes this copy itself.
+        bridge = self._topk_bridge
+        if (
+            bridge is not None
+            and not has_prefill
+            and num_gen_tokens == num_generations
+            and dispatch.get_impl("rope_kv_append", "defer_reuse_copy")(
+                bridge,
+                topk_indices_buffer,
+                token_offset,
+                metadata.shared_topk_indices[:num_generations, :],
+                gvr_prior_indices is not None,
+            )
+        ):
+            return
+        super()._copy_reused_topk(
+            metadata,
+            topk_indices_buffer,
+            token_offset,
+            num_gen_tokens,
+            num_generations,
+            has_prefill,
+            gvr_prior_indices,
+        )
+
+    def _fill_decode_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        q_fp8: torch.Tensor,
+        weights: torch.Tensor,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+    ) -> bool:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super()._fill_decode_topk(
+                metadata=metadata,
+                q_fp8=q_fp8,
+                weights=weights,
+                topk_indices_buffer=topk_indices_buffer,
+                token_offset=token_offset,
+                num_gen_tokens=num_gen_tokens,
+            )
+        out = topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :]
+        # The context-partitioned top-k reproduces the default selection only, so a
+        # deterministic top-k goes to the indexer_topk entry.
+        if not self.top_k.gvr_deterministic and region.cross_rank_topk_decode(
+            self, metadata, q_fp8, weights, token_offset, num_gen_tokens, out, self._topk_bridge
+        ):
+            return True
+        return region.indexer_topk_decode(
+            self, metadata, q_fp8, weights, token_offset, num_gen_tokens, out
+        )
+
+    def _store_mtp_shared_rows(
+        self,
+        topk_indices_buffer: torch.Tensor,
+        metadata: DSAtrtllmAttentionMetadata,
+        num_contexts: int,
+        num_generations: int,
+        num_gen_tokens: int,
+        token_offset: int,
+        has_prefill: bool,
+        has_decode: bool,
+        is_generation: Optional[bool],
+    ) -> None:
+        # rope_kv_append: the fused bridge stashes each request's last accepted row (the same
+        # selection as ``_mtp_last_accepted_rows``).
+        bridge = self._topk_bridge
+        if bridge is not None and has_decode and not has_prefill and is_generation:
+            num_accepted = metadata.mtp_num_accepted
+            if num_accepted is not None:
+                num_accepted = num_accepted[num_contexts : num_contexts + num_generations]
+            if dispatch.get_impl("rope_kv_append", "defer_mtp_row_stash")(
+                bridge,
+                metadata.shared_topk_indices,
+                num_contexts,
+                num_generations,
+                num_gen_tokens // num_generations,
+                num_accepted,
+            ):
+                return
+        super()._store_mtp_shared_rows(
+            topk_indices_buffer,
+            metadata,
+            num_contexts,
+            num_generations,
+            num_gen_tokens,
+            token_offset,
+            has_prefill,
+            has_decode,
+            is_generation,
+        )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention_forward.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention_forward.py
new file mode 100644
index 000000000000..91133f458fca
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/attention_forward.py
@@ -0,0 +1,744 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The MLA generation and DSA sparse-attention forward of the glm_moe_dsa targets on a decode step.
+
+These are the built-in functions of ``attention/mla.py`` and ``attention/backends/sparse/dsa``
+with the family's fast paths composed in:
+
+- the absorbed-query packet: one producer runs the DSA projections and the absorbed query
+  (``q_nope @ W_UK``) for a layer, so the generation skips its absorption BMM;
+- the deferred-norm packet: the hidden state a post-MoE hop hands over before its RMSNorm,
+  normalized inside the next layer's projection (or materialized when no producer serves);
+- the fused generation rope, which also performs the deferred indexer K-cache scatter;
+- the attention output tail, which takes over the W_UV BMM when ``GlmMLA.forward`` armed it.
+
+``attention.GlmMLA`` and ``attention.GlmDsaHooks`` call them only on a decode step; every other
+step runs the built-in functions. ``trtllm::glm_moe_dsa_mla_dsa_attn_inplace`` is the family's
+copy of the built-in DSA op 2 (``trtllm::mla_dsa_attn_inplace``): the same op on this module's
+``_forward_dsa_attn``.
+"""
+
+from typing import List, Optional
+
+import torch
+
+from tensorrt_llm._torch.attention.attention import extract_extra_attrs
+from tensorrt_llm._torch.attention.backends.interface import AttentionInputType, AttentionMetadata
+from tensorrt_llm._torch.attention.backends.sparse.dsa.module import (
+    forward_context_sparse_attn,
+    forward_dsa_proj,
+    forward_sparse_mla_kvcache_bf16,
+    should_use_short_mha,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs
+from tensorrt_llm._torch.attention.mla import fp8_block_scaling_bmm_out
+from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel
+from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph
+from tensorrt_llm._torch.utils import Fp4QuantizedTensor
+from tensorrt_llm._utils import get_sm_version
+
+from . import dispatch
+
+
+def _absorbed_query_packet(self, position_ids, hidden_states, attn_metadata):
+    """Return the absorbed-query packet of this layer's DSA projections, or None.
+
+    The producers are tried in order; each decides its eligibility before any device work and
+    returns None outside its domain, in which case the next one, and finally the built-in
+    projection, runs:
+
+    - a deferred-norm packet goes to the deferred-norm projection (the evict-hints variant when
+      selected), which normalizes it inside the projection;
+    - ``fused_kv_indexer_projection``: one projection launch for the fused [q_a | kv_a] and the
+      indexer's wk | weights_proj outputs;
+    - ``indexer_layer_front``: the query kernel of the full-indexer target layers and the draft
+      layer (the same packet bytes);
+    - ``absorbed_query`` (with ``shared_indexer_absorbed_query`` first on the shared-indexer
+      layers);
+    - ``draft_absorbed_query`` on the draft layer.
+    """
+    if dispatch.is_enabled("deferred_norm") and dispatch.get_impl("deferred_norm", "is_deferred")(
+        hidden_states
+    ):
+        # The deferred-norm projection feeds the absorbed-query kernels; None makes the caller
+        # materialize the packet.
+        entry = (
+            "deferred_norm_evict_hints"
+            if dispatch.is_enabled("deferred_norm_evict_hints")
+            else "deferred_norm"
+        )
+        return dispatch.get_impl(entry, "dsa_proj")(
+            self, position_ids, hidden_states, attn_metadata
+        )
+    if dispatch.is_enabled("fused_kv_indexer_projection"):
+        packet = dispatch.get_impl("fused_kv_indexer_projection", "dsa_proj")(
+            self, position_ids, hidden_states, attn_metadata
+        )
+        if packet is not None:
+            return packet
+    if dispatch.is_enabled("indexer_layer_front"):
+        # The indexer-front query kernel on the full-indexer target layers and the draft layer
+        # (the same packet bytes); None runs the absorbed-query entries below.
+        packet = dispatch.get_impl("indexer_layer_front", "dsa_proj")(
+            self, position_ids, hidden_states, attn_metadata
+        )
+        if packet is not None:
+            return packet
+    if dispatch.is_enabled("absorbed_query"):
+        if dispatch.is_enabled("shared_indexer_absorbed_query"):
+            packet = dispatch.get_impl("absorbed_query", "shared_dsa_proj")(
+                self, position_ids, hidden_states, attn_metadata
+            )
+            if packet is not None:
+                return packet
+        packet = dispatch.get_impl("absorbed_query", "dsa_proj")(
+            self, position_ids, hidden_states, attn_metadata
+        )
+        if packet is not None:
+            return packet
+    if dispatch.is_enabled("draft_absorbed_query"):
+        return dispatch.get_impl("absorbed_query", "mtp_dsa_proj")(
+            self, position_ids, hidden_states, attn_metadata
+        )
+    return None
+
+
+def _deferred_norm_materialize(hidden_states):
+    """The normalized hidden state of a deferred-norm packet (bit-identical to the post-MoE
+    hop's normalized output); any other input passes through."""
+    if dispatch.is_enabled("deferred_norm") and dispatch.get_impl("deferred_norm", "is_deferred")(
+        hidden_states
+    ):
+        return hidden_states.materialize()
+    return hidden_states
+
+
+def forward_absorption_generation(
+    self,
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    k_pe: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+    output: torch.Tensor,
+    position_ids: Optional[torch.Tensor] = None,
+    latent_cache: Optional[torch.Tensor] = None,
+    absorbed_fused_q: Optional[torch.Tensor] = None,
+    deferred_indexer_k=None,
+    **kwargs,
+) -> torch.Tensor:
+    """Absorbed-MLA generation (``MLA.forward_absorption_generation``) on a decode step.
+
+    ``absorbed_fused_q`` is the absorbed-query packet: a
+    ``[num_tokens, num_heads_tp, kv_lora_rank + qk_rope_head_dim]`` buffer whose
+    ``[..., :kv_lora_rank]`` already holds ``q_nope @ W_UK``. When given, only the absorption
+    BMM is skipped; RoPE, cache append, FP8 query quantization and scheduler buffers run as
+    usual. ``deferred_indexer_k`` is the indexer K-cache scatter ``_forward_dsa_attn`` deferred
+    into this call's rope.
+    """
+    num_tokens = q.shape[0]
+    q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split(
+        [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+    )
+
+    # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank]
+    # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp
+    num_seqs = attn_metadata.num_seqs
+
+    # The rope kernel writes seqQOffset[batch_idx + 1] for every
+    # batch_idx < batch_size, and CUB-scans total_s_len / seq_len entries
+    # into seqKVOffsets (mlaKernels.cu:579, :650). `batch_size` there is the
+    # count of presented rows, which equals num_seqs only while every
+    # request contributes the same number of query tokens. Under the
+    # token-major presentation a row is a token, so size these by the token
+    # count. Uniform verification keeps the established request-major
+    # allocation rather than scaling these temporary buffers by K.
+    cu_seqlens_rows = (
+        max(num_seqs, num_tokens) + 1
+        if getattr(attn_metadata, "is_ragged_verify", False)
+        else num_seqs + 1
+    )
+    cu_q_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device)
+    cu_kv_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device)
+    fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=q.device)
+    use_fp8_mla = getattr(self.mqa, "has_fp8_kv_cache", False) or getattr(
+        self.mqa, "has_fp4_kv_cache", False
+    )
+
+    mla_bmm1_scale = None
+    mla_bmm2_scale = None
+    quant_q_buffer = None
+    if use_fp8_mla:
+        mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=q.device)
+        mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=q.device)
+        quant_q_buffer = torch.empty(
+            num_tokens,
+            self.num_heads_tp,
+            (self.kv_lora_rank + self.qk_rope_head_dim),
+            dtype=torch.uint8,
+            device=q.device,
+        )
+
+    if hasattr(self, "k_b_proj_trans"):
+        if absorbed_fused_q is not None:
+            fused_q = absorbed_fused_q
+        else:
+            fused_q = torch.empty(
+                [
+                    num_tokens,
+                    self.num_heads_tp,
+                    (self.kv_lora_rank + self.qk_rope_head_dim),
+                ],
+                dtype=q.dtype,
+                device=q.device,
+            )
+
+        fp4_mla = (
+            self.sparse_params is None
+            and self.mqa.support_fp4_kv_cache()
+            and self.quant_config is not None
+            and self.quant_config.layer_quant_mode.has_fp4_kv_cache()
+        )
+        if fp4_mla and latent_cache is None:
+            raise RuntimeError(
+                "FP4 MLA generation requires a latent cache for fused "
+                "Q quantization, RoPE, and cache update."
+            )
+        fuse_fp4_q_quant = fp4_mla
+        fp4_rope_kwargs = {"fuse_fp4_q_quant": True} if fuse_fp4_q_quant else {}
+
+        def _mla_gen_rope():
+            # rope_kv_append: one fused launch writes the bytes of mla_rope_generation plus the
+            # deferred indexer K-cache scatter. Every check runs before any device work;
+            # otherwise the deferred scatter runs natively first and the native path follows.
+            if deferred_indexer_k is not None or dispatch.is_enabled("rope_kv_append"):
+                rope_fused = (
+                    not self.apply_rotary_emb
+                    and self.kv_cache_dtype != "fp8_ds_mla"
+                    and not fp4_rope_kwargs
+                    and dispatch.get_impl("rope_kv_append", "rope_eligible")(
+                        self.mqa,
+                        attn_metadata,
+                        fused_q,
+                        q_pe,
+                        latent_cache,
+                        quant_q_buffer,
+                        mla_bmm1_scale,
+                        mla_bmm2_scale,
+                    )
+                )
+                if rope_fused:
+                    # indexer_layer_front: the same fused rope, forked to trigger its
+                    # dependents after its last read (the same bytes).
+                    rope_owner = "rope_kv_append"
+                    if dispatch.is_enabled("indexer_layer_front"):
+                        front_rope = dispatch.get_impl("indexer_layer_front", "rope_available")()
+                        dispatch.log_capability(
+                            "indexer_layer_front",
+                            "rope",
+                            front_rope,
+                            "" if front_rope else "indexer_layer_front rope op not loaded",
+                        )
+                        if front_rope:
+                            rope_owner = "indexer_layer_front"
+                    dispatch.get_impl(rope_owner, "mla_rope_generation")(
+                        self.mqa,
+                        fused_q,
+                        q_pe,
+                        latent_cache,
+                        attn_metadata,
+                        cu_q_seqlens,
+                        cu_kv_seqlens,
+                        fmha_scheduler_counter,
+                        mla_bmm1_scale,
+                        mla_bmm2_scale,
+                        quant_q_buffer,
+                        indexer_k=deferred_indexer_k,
+                    )
+                    return
+                if deferred_indexer_k is not None:
+                    dispatch.get_impl("rope_kv_append", "run_indexer_k_scatter_native")(
+                        deferred_indexer_k
+                    )
+            if self.apply_rotary_emb:
+                # Non-fused backends (Vanilla / FlashInfer) do not fuse RoPE
+                # in the attention kernel. Reuse apply_rope, which rotates
+                # q's q_pe slice in place, then copy rotated k_pe into the
+                # latent cache that the backend appends.
+                assert position_ids is not None
+                assert latent_cache is not None
+                gen_position_ids = position_ids[..., attn_metadata.num_ctx_tokens :]
+                k_pe_rope = self.apply_rope(q, k_pe, gen_position_ids)
+                fused_q[..., self.kv_lora_rank :] = q_pe
+                latent_cache[..., self.kv_lora_rank :] = k_pe_rope
+            else:
+                # Fused backend (TRTLLM): RoPE, latent-cache append and the
+                # trtllm-gen scheduler buffers are all produced in-kernel.
+                if self.kv_cache_dtype == "fp8_ds_mla":
+                    fused_q[..., self.kv_lora_rank :] = q_pe
+                else:
+                    self.mqa.mla_rope_generation(
+                        fused_q,
+                        q_pe,
+                        latent_cache,
+                        attn_metadata,
+                        cu_q_seqlens,
+                        cu_kv_seqlens,
+                        fmha_scheduler_counter,
+                        mla_bmm1_scale,
+                        mla_bmm2_scale,
+                        quant_q_buffer,
+                        **fp4_rope_kwargs,
+                    )
+
+        if fuse_fp4_q_quant or use_fp8_mla:
+            rope_stream = None
+        else:
+            rope_stream = self.aux_stream
+        if absorbed_fused_q is not None:
+            _mla_gen_rope()
+        elif self.k_b_proj_trans.dtype == torch.bfloat16:
+            # [num_heads, num_tokens, self.qk_nope_head_dim]
+            q_nope_t = q_nope.transpose(0, 1)
+            # [num_heads, num_tokens, self.kv_lora_rank]
+            q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1)
+
+            # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim]
+            # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank]
+            # The output of bmm is written directly into fused_q
+            maybe_execute_in_parallel(
+                lambda: self._bmm_bf16_out(
+                    q_nope_t,
+                    self.k_b_proj_trans,
+                    self.k_b_proj_trans.transpose(1, 2),
+                    q_nope_out,
+                    self._locality_domain_k_b_proj_trans_shards,
+                ),
+                _mla_gen_rope,
+                self.ln_events[0],
+                self.ln_events[1],
+                rope_stream,
+            )
+
+        elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn:
+            # [num_heads, num_tokens, self.kv_lora_rank]
+            q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1)
+
+            maybe_execute_in_parallel(
+                lambda: fp8_block_scaling_bmm_out(
+                    q_nope,
+                    self.k_b_proj_trans,
+                    self.k_b_proj_trans_scale,
+                    q_nope_out,
+                    self.k_b_proj_trans_dequant,
+                    self.use_cute_dsl_blockscaling_bmm,
+                ),
+                _mla_gen_rope,
+                self.ln_events[0],
+                self.ln_events[1],
+                rope_stream,
+            )
+        else:
+            raise NotImplementedError(f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.")
+
+        fused_q = fused_q.view(
+            [
+                num_tokens,
+                self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim),
+            ]
+        )
+    else:
+        raise RuntimeError("MLA absorption requires k_b_proj_trans")
+
+    # attention_output_tail_prefetch: on a call that GlmMLA.forward armed, fork the L2
+    # prefetch of the tail weights onto a private stream while the attention runs; the
+    # projection step joins it. No-op otherwise (decided before any device work).
+    attn_tail = getattr(self, "_attn_tail", None)
+    if getattr(self, "_tail_prefetch", False) and attn_tail is not None:
+        attn_tail.before_attention(self, point="attn")
+
+    attn_out_latent = self._attn_forward_gen(
+        self.mqa,
+        fused_q,
+        None,
+        None,
+        position_ids,
+        attn_metadata,
+        attention_input_type=AttentionInputType.generation_only,
+        out_scale=self.out_scale,
+        latent_cache=latent_cache,  # kvcache and k_pe
+        q_pe=q_pe,  # used by `invokeMLARopeGeneration`
+        cu_q_seqlens=cu_q_seqlens,  # used by `mlaGeneration`
+        cu_kv_seqlens=cu_kv_seqlens,  # used by `mlaGeneration`
+        fmha_scheduler_counter=fmha_scheduler_counter,  # used by `mlaGeneration`
+        mla_bmm1_scale=mla_bmm1_scale,  # used by `mlaGeneration`
+        mla_bmm2_scale=mla_bmm2_scale,  # used by `mlaGeneration`
+        quant_q_buffer=quant_q_buffer,  # used by `mlaGeneration`
+        **kwargs,
+    )
+    fused_q = None
+
+    # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp
+    assert (
+        attn_out_latent.shape[0] == q.shape[0]
+        and attn_out_latent.shape[1] == self.num_heads_tp_cp * self.kv_lora_rank
+    )
+
+    # [seq, num_heads, kv_lora_rank]
+    attn_out_latent = attn_out_latent.view([-1, self.num_heads_tp_cp, self.kv_lora_rank])
+
+    # Acts only on a call that GlmMLA.forward armed; when the attention tail takes the call,
+    # its projection step consumes the latent and the native W_UV BMM is skipped. The latent
+    # shape is checked before any device work.
+    attn_tail = getattr(self, "_attn_tail", None)
+    if attn_tail is not None and attn_tail.defer_w_uv(self, attn_metadata, attn_out_latent, output):
+        return output
+
+    attn_output = output.view([num_tokens, self.num_heads_tp_cp, self.v_head_dim])
+
+    if self.v_b_proj.dtype == torch.bfloat16:
+        # [num_heads, seq, kv_lora_rank] x [num_heads, kv_lora_rank, v_head_dim]
+        # -> [num_heads, seq, v_head_dim]
+        self._bmm_bf16_out(
+            attn_out_latent.transpose(0, 1),
+            self.v_b_proj,
+            self.v_b_proj.transpose(1, 2),
+            attn_output.transpose(0, 1),
+            self._locality_domain_v_b_proj_shards,
+        )
+    elif self.v_b_proj.dtype == torch.float8_e4m3fn:
+        fp8_block_scaling_bmm_out(
+            attn_out_latent,
+            self.v_b_proj,
+            self.v_b_proj_scale,
+            attn_output.transpose(0, 1),
+            self.v_b_proj_dequant,
+            self.use_cute_dsl_blockscaling_bmm,
+        )
+    else:
+        raise NotImplementedError(f"Missing bmm impl for dtype: {self.v_b_proj.dtype}.")
+
+    return output
+
+
+def forward_generation_sparse_attn(
+    self,
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    k_pe: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+    output: torch.Tensor,
+    position_ids: Optional[torch.Tensor] = None,
+    latent_cache: Optional[torch.Tensor] = None,
+    indexer_intermediates: Optional[List[torch.Tensor]] = None,
+    absorbed_fused_q: Optional[torch.Tensor] = None,
+    deferred_indexer_k=None,
+) -> torch.Tensor:
+    if get_sm_version() >= 100:
+        return self.forward_absorption_generation(
+            q,
+            compressed_kv,
+            k_pe,
+            attn_metadata,
+            output,
+            position_ids=position_ids,
+            latent_cache=latent_cache,
+            absorbed_fused_q=absorbed_fused_q,
+            deferred_indexer_k=deferred_indexer_k,
+            sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=indexer_intermediates),
+        )
+    assert absorbed_fused_q is None, "absorbed_query preabsorbed queries require SM100+"
+    assert deferred_indexer_k is None
+    return forward_sparse_mla_kvcache_bf16(
+        self,
+        q,
+        latent_cache,
+        attn_metadata,
+        output,
+        indexer_intermediates,
+        is_generation=True,
+    )
+
+
+def forward_sparse_attn(
+    self,
+    position_ids: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+    attn_output: list[torch.Tensor],
+) -> None:
+    """
+    Forward pass for the MLA module with DSA (always in MQA mode).
+    Writes result into output tensor in-place.
+
+    Delegates to forward_dsa_proj (token-wise projections) followed by
+    _forward_dsa_attn (batch-dependent attention dispatch). An absorbed-query packet replaces
+    forward_dsa_proj; a deferred-norm packet that no producer takes is materialized first.
+
+    Args:
+        position_ids (Optional[torch.IntTensor]): The position IDs.
+        hidden_states (torch.Tensor): The hidden states.
+        attn_metadata (AttentionMetadata): The attention metadata.
+        attn_output: Algorithm-defined output tensors. DSA writes to the
+            standard output at index zero.
+    """
+    output = attn_output[0]
+    query_packet = _absorbed_query_packet(self, position_ids, hidden_states, attn_metadata)
+    if query_packet is None:
+        hidden_states = _deferred_norm_materialize(hidden_states)
+        proj_outputs = forward_dsa_proj(self, position_ids, hidden_states, attn_metadata)
+        absorbed_fused_q = None
+    else:
+        proj_outputs, absorbed_fused_q = query_packet.proj_outputs, query_packet.fused_q
+    q, compressed_kv, k_pe, latent_cache = proj_outputs[:4]
+    indexer_intermediates = proj_outputs[4:]
+    _forward_dsa_attn(
+        self,
+        q,
+        compressed_kv,
+        k_pe,
+        latent_cache,
+        indexer_intermediates,
+        position_ids,
+        attn_metadata,
+        output,
+        absorbed_fused_q=absorbed_fused_q,
+    )
+
+
+def forward_sparse_attn_custom_op(
+    self,
+    hidden_states: torch.Tensor,
+    position_ids: Optional[torch.Tensor],
+    attn_output: list[torch.Tensor],
+    latent_cache_gen: Optional[torch.Tensor],
+) -> None:
+    """Run DSA's two-stage custom ops for piecewise CUDA graph capture.
+
+    The second op is the family's ``trtllm::glm_moe_dsa_mla_dsa_attn_inplace``. With an
+    absorbed-query packet the attention runs directly, without the ops.
+    """
+    del latent_cache_gen
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import custom_ops  # noqa: F401
+
+    if (
+        dispatch.is_enabled("absorbed_query") or dispatch.is_enabled("draft_absorbed_query")
+    ) and not isinstance(hidden_states, Fp4QuantizedTensor):
+        metadata, _ = extract_extra_attrs(self.layer_idx_str, "mla")
+        query_packet = _absorbed_query_packet(self, position_ids, hidden_states, metadata)
+        if query_packet is not None:
+            # The packet producers decline under torch.compile and breakable CUDA graphs, so
+            # the two-op split is not needed on this path.
+            proj_outputs = query_packet.proj_outputs
+            _forward_dsa_attn(
+                self,
+                *proj_outputs[:4],
+                proj_outputs[4:],
+                position_ids,
+                metadata,
+                attn_output[0],
+                absorbed_fused_q=query_packet.fused_q,
+            )
+            return
+        hidden_states = _deferred_norm_materialize(hidden_states)
+
+    if isinstance(hidden_states, Fp4QuantizedTensor):
+        proj_outputs = torch.ops.trtllm.mla_dsa_proj(
+            hidden_states.unquantized_hidden_states,
+            position_ids,
+            self.layer_idx_str,
+            hidden_states.fp4_tensor,
+            hidden_states.scaling_factor,
+        )
+    else:
+        proj_outputs = torch.ops.trtllm.mla_dsa_proj(
+            hidden_states, position_ids, self.layer_idx_str
+        )
+    q, compressed_kv, k_pe, latent_cache = proj_outputs[:4]
+    indexer_intermediates = proj_outputs[4:]
+    maybe_bcg_glm_moe_dsa_mla_dsa_attn_inplace(
+        q,
+        compressed_kv,
+        k_pe,
+        latent_cache,
+        indexer_intermediates,
+        position_ids,
+        self.layer_idx_str,
+        attn_output[0],
+    )
+
+
+def _forward_dsa_attn(
+    self,
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    k_pe: torch.Tensor,
+    latent_cache: torch.Tensor,
+    indexer_intermediates: List[torch.Tensor],
+    position_ids: Optional[torch.Tensor],
+    attn_metadata: AttentionMetadata,
+    output: torch.Tensor,
+    absorbed_fused_q: Optional[torch.Tensor] = None,
+) -> None:
+    """Batch-structure-dependent attention for DSA MLA (Op 2, not graph-captured).
+
+    indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale]
+    when the indexer ran in Op 1, or [] when short-MHA handled all tokens.
+
+    All num_tokens slicing happens here (not in Op 1) because
+    num_tokens comes from batch-specific metadata and must not be
+    baked into CUDA graph capture.
+    """
+    num_contexts = attn_metadata.num_contexts
+    num_generations = attn_metadata.num_generations
+    num_ctx_tokens = attn_metadata.num_ctx_tokens
+    num_tokens = attn_metadata.num_tokens
+
+    # Slice Op 1 outputs to actual num_tokens (Op 1 operates on the
+    # full padded tensor for CUDA graph compatibility).
+    q = q[:num_tokens, ...]
+    compressed_kv = compressed_kv[:num_tokens, ...]
+    k_pe = k_pe[:num_tokens, ...]
+    latent_cache = latent_cache[:num_tokens, ...]
+    if position_ids is not None:
+        position_ids = position_ids[..., :num_tokens]
+
+    use_short_mha_for_ctx = num_contexts > 0 and should_use_short_mha(
+        self, attn_metadata, position_ids
+    )
+
+    deferred_indexer_k = None
+    if not (use_short_mha_for_ctx and num_generations == 0) and self.mqa.indexer is not None:
+        q_fp8, k_fp8, k_scale, weights, q_scale = indexer_intermediates
+        # Slice indexer intermediates to actual num_tokens (they were
+        # computed on the full padded tensor in Op 1).
+        q_fp8 = q_fp8[:num_tokens, ...]
+        k_fp8 = k_fp8[:num_tokens, ...]
+        k_scale = k_scale[:num_tokens, ...]
+        weights = weights[:num_tokens, ...]
+        q_scale = q_scale[:num_tokens, ...]
+        indexer_intermediates = [q_fp8, k_fp8, k_scale, weights, q_scale]
+        # Update the full mixed batch once before context/generation prediction
+        # splits. Context chunk gathers must observe the newly appended keys.
+        # rope_kv_append: a pure-generation batch folds this scatter into the same layer's
+        # fused generation rope, which still runs before the indexer reads the K cache. The
+        # deferred scatter must reach the generation call below, which is checked right after
+        # that call.
+        if dispatch.is_enabled("rope_kv_append"):
+            deferred_indexer_k = dispatch.get_impl("rope_kv_append", "defer_indexer_k_scatter")(
+                self, k_fp8, k_scale, attn_metadata, forwards_packet=True
+            )
+        if deferred_indexer_k is None:
+            self.mqa.indexer._update_k_cache(k_fp8, k_scale, attn_metadata)
+    else:
+        indexer_intermediates = None
+
+    assert output is not None, "output must be provided"
+
+    if num_contexts > 0:
+        q_ctx = q[:num_ctx_tokens, ...]
+        compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...]
+        k_pe_ctx = k_pe[:num_ctx_tokens, ...]
+        latent_cache_ctx = latent_cache[:num_ctx_tokens, ...]
+        ctx_position_ids = position_ids[..., :num_ctx_tokens] if position_ids is not None else None
+        if self.apply_rotary_emb:
+            assert ctx_position_ids is not None
+            k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids)
+
+        forward_context_sparse_attn(
+            self,
+            q_ctx,
+            compressed_kv_ctx,
+            k_pe_ctx,
+            attn_metadata,
+            output[:num_ctx_tokens, :],
+            latent_cache_ctx,
+            position_ids=ctx_position_ids,
+            indexer_intermediates=indexer_intermediates,
+        )
+
+    if num_generations > 0:
+        q_gen = q[num_ctx_tokens:, ...]
+        compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...]
+        k_pe_gen = k_pe[num_ctx_tokens:, ...]
+        latent_cache_gen = latent_cache[num_ctx_tokens:, ...]
+        gen_position_ids = (
+            position_ids[..., num_ctx_tokens:num_tokens] if position_ids is not None else None
+        )
+        if self.apply_rotary_emb:
+            assert gen_position_ids is not None
+            k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids)
+
+        forward_generation_sparse_attn(
+            self,
+            q_gen,
+            compressed_kv_gen,
+            k_pe_gen,
+            attn_metadata,
+            output[num_ctx_tokens:num_tokens, :],
+            latent_cache=latent_cache_gen,
+            position_ids=gen_position_ids,
+            indexer_intermediates=indexer_intermediates,
+            absorbed_fused_q=(
+                absorbed_fused_q[num_ctx_tokens:num_tokens]
+                if absorbed_fused_q is not None
+                else None
+            ),
+            deferred_indexer_k=deferred_indexer_k,
+        )
+    if deferred_indexer_k is not None:
+        # Raises unless the generation rope (fused or the native fallback) executed the
+        # deferred indexer K-cache scatter.
+        dispatch.get_impl("rope_kv_append", "check_indexer_k_scatter_consumed")(deferred_indexer_k)
+
+    # Join the prev_topk copy forked in sparse_attn_indexer; CUDA graph
+    # capture requires the join within the same layer's forward.
+    if self.mqa.indexer is not None:
+        self.mqa.indexer.maybe_join_prev_topk_copy()
+
+
+@torch.library.custom_op("trtllm::glm_moe_dsa_mla_dsa_attn_inplace", mutates_args=("output",))
+def glm_moe_dsa_mla_dsa_attn_inplace(
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    k_pe: torch.Tensor,
+    latent_cache: torch.Tensor,
+    indexer_intermediates: List[torch.Tensor],
+    position_ids: Optional[torch.Tensor],
+    layer_idx: str,
+    output: torch.Tensor,
+) -> None:
+    """Batch-structure-dependent attention dispatch for DSA MLA, on ``_forward_dsa_attn`` above.
+
+    The built-in ``trtllm::mla_dsa_attn_inplace`` with the family's op-2 function:
+    indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] when the indexer ran in
+    Op 1, or [] when short-MHA handled all tokens. This op is excluded from CUDA graph capture.
+    """
+    metadata, mla_layer = extract_extra_attrs(layer_idx, "mla")
+    _forward_dsa_attn(
+        mla_layer,
+        q,
+        compressed_kv,
+        k_pe,
+        latent_cache,
+        indexer_intermediates,
+        position_ids,
+        metadata,
+        output,
+    )
+
+
+@glm_moe_dsa_mla_dsa_attn_inplace.register_fake
+def _glm_moe_dsa_mla_dsa_attn_inplace_fake(
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    k_pe: torch.Tensor,
+    latent_cache: torch.Tensor,
+    indexer_intermediates: List[torch.Tensor],
+    position_ids: Optional[torch.Tensor],
+    layer_idx: str,
+    output: torch.Tensor,
+) -> None:
+    """Model the in-place output mutation during fake-tensor propagation."""
+
+
+maybe_bcg_glm_moe_dsa_mla_dsa_attn_inplace = eager_on_graph(glm_moe_dsa_mla_dsa_attn_inplace)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/dispatch.py
new file mode 100644
index 000000000000..1cf9a9abb921
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/dispatch.py
@@ -0,0 +1,1209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The decode fast paths of the glm_moe_dsa targets, which of them a process runs, and what happens
+when one cannot.
+
+Composition. ``FAST_PATHS`` names every fast path of the family once, with the
+paths it extends (``depends_on``): a path that extends another has no effect
+without it. Each target serves some of them (``TargetPaths``), and every path
+it serves is on by default. ``TRTLLM_STAIRCASE_GLM_LL_DISABLE`` takes a
+comma-separated list of path names and marks those off in the plan together
+with every path that depends on them, transitively; a name that is not a path
+the target serves raises. ``resolve`` reads it once, at construction, into a
+``Plan``. The fast paths read the installed plan through ``is_enabled``;
+``DISABLE`` turns a path and its dependents off.
+
+Rank uniformity. Several fast paths are collectives across the model's ranks: a
+rank that runs one while a peer runs the generic path waits forever. So every
+rank must read the same knob values. ``check_rank_uniform`` allgathers them,
+with the shared switches the kernels read and the modeling_v2 mode
+(``rank_uniform_inputs``), once, at construction, before they are parsed and
+before any fast path builds a collective workspace, and raises on every rank
+with the per-rank table on a mismatch. The plan is a function of the knobs
+alone, so equal knobs give every rank the same plan, and a value no rank can
+parse fails on every rank alike. Once the model is built,
+``check_rank_uniform_adopted`` compares the fast-path state every rank
+attached to its modules the same way.
+
+Phases. ``step`` runs a model step under the ``Phase`` of the target that
+takes it (``PREFILL`` or ``DECODE``); ``phase`` reads it, and is None outside
+a step.
+
+Two kinds of not taking a fast path, kept apart:
+
+* **uncovered**: the step is outside the shapes the path takes (more rows than
+  its kernel holds, a KV length outside its tier, a prefill step). ``uncovered``
+  logs the reason once per (path, reason); the caller runs the next tier or the
+  generic path in every mode.
+* **missed**: the step is inside the shapes the path takes, but a precondition
+  is absent (a workspace, a compiled kernel, multi-node NVLink memory, a device
+  check). Under ``modeling_v2='require'`` ``missed`` raises
+  ``FastPathUnavailable`` naming the path and the precondition; otherwise it
+  logs once per (path, reason), counts the miss, and the caller falls back. A
+  miss decided per call never raises: it is counted and logged once.
+
+The kernels report their decisions through ``log_capability``, which also
+classifies every decline (``classify_decline``) by where it is decided
+(``DECLINE_KINDS``), the layers each path serves (``LAYER_CELLS``) and the
+reason: a path turned off is neither kind, a decline decided once inside the
+path's layers is a miss, and a decline decided per call is a miss only when
+its reason names a missing resource.
+
+Selection. A kernel or a shared hook asks ``is_enabled(path)`` whether a fast
+path runs, looks up its entry points with ``get_impl(path, name)`` (the kernel
+package of the path registers them with ``register`` when it is imported) and
+logs its per-rank decisions with ``log_capability``. ``is_enabled`` answers
+from the installed plan (``install``); without one, which is the case of every
+process that does not build a target of the family, no fast path runs. ``install`` logs a
+record line (``feature_record``): every enabled path with its numerical variant
+(``NUMERICAL_VARIANTS``), every set knob of the family and the shared
+TensorRT-LLM switches the kernels read.
+
+Validation. ``TRTLLM_MODELING_V2_VALIDATE`` (read once, in ``resolve``) turns on
+the diagnostics of the fast paths that have some (``validate_mode``): the
+step-input stager's guard and the draft loop's in-pipeline native oracle.
+
+Knobs. The knobs of the family are ``KNOBS``: ``DISABLE``, the numerical variant
+of ``routed_moe_into_hop`` and the cross-rank top-k wait timeout (read by the
+C++ kernels). A target refuses a knob that only a path it does not serve reads
+(``KNOB_PATHS``). The kernels run at fixed launch parameters; ``resolve`` refuses
+every other ``TRTLLM_STAIRCASE_GLM_LL_*`` variable that is set.
+"""
+
+from __future__ import annotations
+
+import enum
+import functools
+import hashlib
+import importlib
+import json
+import os
+import re
+import threading
+from dataclasses import dataclass, field
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    FrozenSet,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+)
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
+    MODELING_V2_VALIDATE_ENV,
+    ModelingV2Mode,
+    step_contract_enabled,
+)
+from tensorrt_llm.logger import logger
+
+#: Prefix of every knob of this target; every set variable with it is part of the rank-uniform
+#: check.
+KNOB_PREFIX = "TRTLLM_STAIRCASE_GLM_LL_"
+#: Comma-separated fast-path names to turn off, with every path that depends on them.
+DISABLE_ENV = KNOB_PREFIX + "DISABLE"
+#: The numerical variant of ``routed_moe_into_hop`` (``NUMERICAL_VARIANTS``).
+ACCUMULATION_ENV = KNOB_PREFIX + "ROUTED_MOE_INTO_HOP_ACCUMULATION"
+#: How long the cross-rank top-k kernels wait for their peers before they trap, in ms (C++).
+CROSS_RANK_TOPK_TIMEOUT_ENV = KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS"
+#: The knobs of this target; ``resolve`` refuses any other set ``KNOB_PREFIX`` variable.
+KNOBS: FrozenSet[str] = frozenset({DISABLE_ENV, ACCUMULATION_ENV, CROSS_RANK_TOPK_TIMEOUT_ENV})
+#: A fast-path selector this target does not read; ``resolve`` refuses it.
+FEATURE_ENV = KNOB_PREFIX + "FEATURES"
+#: A narrowing of ``deferred_input_prep`` to some of its parts; ``resolve`` refuses it.
+PARTS_ENV = KNOB_PREFIX + "DEFERRED_INPUT_PREP_PARTS"
+
+_VALIDATE_HINT = f"set {MODELING_V2_VALIDATE_ENV}=1 instead"
+#: Knobs this target does not read, each with what to use instead; ``resolve`` refuses them.
+REMOVED_KNOBS: Mapping[str, str] = {
+    FEATURE_ENV: f"every fast path of this target is on; {KNOB_PREFIX}DISABLE turns paths off",
+    PARTS_ENV: (
+        "each part is a fast path of its own (prep = deferred_input_prep, store = "
+        f"sampler_store, stream = execution_stream_loop); list the parts to drop in "
+        f"{KNOB_PREFIX}DISABLE"
+    ),
+    KNOB_PREFIX + "DEFERRED_INPUT_PREP_DEBUG": _VALIDATE_HINT,
+    KNOB_PREFIX + "DRAFT_LOOP_VERIFY": _VALIDATE_HINT,
+    KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_STATS": "the hit counters are not reported",
+    KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_REPORT_EVERY": "the hit counters are not reported",
+    KNOB_PREFIX + "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE": (
+        f"list shared_indexer_absorbed_query in {KNOB_PREFIX}DISABLE to run absorbed_query on the "
+        "shared-indexer layers"
+    ),
+}
+#: What ``resolve`` says about any other set variable with the prefix that is not in ``KNOBS``.
+_FIXED_HINT = "the fast paths of this target run at fixed launch parameters"
+
+_LOG_TAG = "[glm_moe_dsa]"
+
+
+@dataclass(frozen=True)
+class FastPath:
+    """One decode fast path: its name and the paths it extends."""
+
+    name: str
+    depends_on: Tuple[str, ...] = ()
+
+
+# The attention query front, the indexer and top-k selection, sparse attention,
+# the attention output and its cross-rank hop, the MoE region and its hop, the
+# dense layers, the draft (MTP) layer and loop, and the engine-side step work.
+FAST_PATHS: Tuple[FastPath, ...] = (
+    # The target layers' query front: the query projections and the absorbed query.
+    FastPath("absorbed_query"),
+    # The same query front on the draft layer.
+    FastPath("draft_absorbed_query"),
+    # The query front of layers that reuse an earlier layer's indexer selection.
+    FastPath("shared_indexer_absorbed_query", depends_on=("absorbed_query",)),
+    # One projection launch for the latent KV and the indexer key and weights.
+    FastPath("fused_kv_indexer_projection", depends_on=("absorbed_query",)),
+    # The query front of layers that run the full indexer, with its rope and
+    # top-k exchange started early.
+    FastPath(
+        "indexer_layer_front",
+        depends_on=(
+            "absorbed_query",
+            "draft_absorbed_query",
+            "indexer_topk",
+            "rope_kv_append",
+            "cross_rank_topk",
+        ),
+    ),
+    # The next layer's RMSNorm deferred across the post-MoE hop into its projection.
+    FastPath("deferred_norm", depends_on=("absorbed_query", "post_moe_hop")),
+    # The deferred-norm hop and projection chain with cache-eviction hints.
+    FastPath("deferred_norm_evict_hints", depends_on=("deferred_norm",)),
+    # Rope, FP8 query, latent-KV append, indexer-K scatter and the top-k to
+    # KV-index bridge in one launch per layer.
+    FastPath("rope_kv_append"),
+    # Indexer logits and top-k selection on decode steps.
+    FastPath("indexer_topk"),
+    # Top-k over the KV pages each rank owns, merged across the ranks.
+    FastPath("cross_rank_topk"),
+    # The cross-rank top-k with a cluster merge of the per-rank candidates.
+    FastPath("cross_rank_topk_cluster_merge", depends_on=("cross_rank_topk",)),
+    # The cluster merge also emits the top-k to KV-index bridge.
+    FastPath(
+        "cross_rank_topk_global_emit",
+        depends_on=("rope_kv_append", "cross_rank_topk_cluster_merge"),
+    ),
+    # Sparse MLA attention for the decode rows.
+    FastPath("sparse_mla_decode"),
+    # The attention output projection and its local tail.
+    FastPath("attention_output_tail"),
+    # The attention tail's prefetch schedule: the next weights are fetched while the current
+    # attention runs.
+    FastPath("attention_output_tail_prefetch", depends_on=("attention_output_tail",)),
+    # The attention output published straight into the post-attention hop.
+    FastPath(
+        "attention_output_into_hop", depends_on=("attention_output_tail", "post_attention_hop")
+    ),
+    # Output-projection weight staging inside an already-waiting kernel of the attention window.
+    FastPath("oproj_weight_staging", depends_on=("attention_output_into_hop", "sparse_mla_decode")),
+    # Cross-rank sum, residual add, RMSNorm and router after attention, in one hop.
+    FastPath("post_attention_hop"),
+    # The MoE routing (top-8) computed inside the post-attention hop.
+    FastPath("hop_routing", depends_on=("attention_output_into_hop", "routed_moe")),
+    # The router GEMM at the pre-MoE boundary.
+    FastPath("router_gemm"),
+    # The routed NVFP4 experts in one launch.
+    FastPath("routed_moe"),
+    # The routed-expert result written straight into the post-MoE hop.
+    FastPath("routed_moe_into_hop", depends_on=("post_moe_hop", "routed_moe")),
+    # The shared expert's weights prefetched into L2 by the routed-expert launch.
+    FastPath("shared_expert_prefetch", depends_on=("routed_moe_into_hop", "hop_routing")),
+    # Cross-rank finalize, reduction and RMSNorm after the MoE, in one hop.
+    FastPath("post_moe_hop"),
+    # The next MoE layer's predicted experts prefetched into L2.
+    FastPath("cross_layer_expert_prefetch", depends_on=("routed_moe",)),
+    # The dense layers' small-row MLP.
+    FastPath("dense_mlp"),
+    # The draft layer's input normalization and packing.
+    FastPath("draft_front"),
+    # The draft layer's bf16 routed experts.
+    FastPath("draft_routed_moe"),
+    # Only the accepted rows of the first draft step through the post-attention hop.
+    FastPath("draft_live_rows", depends_on=("post_attention_hop", "draft_loop")),
+    # Draft-token selection at the end of each draft step.
+    FastPath("draft_tail"),
+    # The draft loop's embedding, bookkeeping and sparse-attention metadata.
+    FastPath("draft_loop"),
+    # The decode step's metadata prologue and the target embedding in one launch.
+    FastPath("decode_prologue", depends_on=("draft_loop",)),
+    # The decode step's input preparation recorded and issued as one launch.
+    FastPath("deferred_input_prep"),
+    # The speculative sampler's store update and host copy in one launch.
+    FastPath("sampler_store"),
+    # The overlap executor loop on the execution stream.
+    FastPath("execution_stream_loop"),
+    # Freed host pages returned to the operating system lazily on decode ranks.
+    FastPath("host_page_retention", depends_on=("post_attention_hop",)),
+)
+
+FAST_PATH_NAMES: Tuple[str, ...] = tuple(p.name for p in FAST_PATHS)
+_BY_NAME: Mapping[str, FastPath] = {p.name: p for p in FAST_PATHS}
+
+
+def _dependents() -> Dict[str, FrozenSet[str]]:
+    """Every path that depends on each path, directly or transitively."""
+    direct: Dict[str, List[str]] = {name: [] for name in FAST_PATH_NAMES}
+    for path in FAST_PATHS:
+        for base in path.depends_on:
+            direct[base].append(path.name)
+    closure = {}
+    for name in FAST_PATH_NAMES:
+        seen, todo = set(), list(direct[name])
+        while todo:
+            dependent = todo.pop()
+            if dependent not in seen:
+                seen.add(dependent)
+                todo.extend(direct[dependent])
+        closure[name] = frozenset(seen)
+    return closure
+
+
+_DEPENDENTS = _dependents()
+
+
+def parse_disable(raw: Optional[str], served: Iterable[str] = FAST_PATH_NAMES) -> FrozenSet[str]:
+    """The fast-path names a ``DISABLE`` value lists (case and blanks ignored).
+
+    ``served`` is the fast paths of the target that reads the value (every fast path by default).
+
+    Raises:
+        ValueError: naming every entry that is not one of ``served``.
+    """
+    if not raw:
+        return frozenset()
+    names = {tok.strip().lower() for tok in raw.split(",") if tok.strip()}
+    served = frozenset(served)
+    unknown = sorted(names - served)
+    if unknown:
+        raise ValueError(
+            f"{DISABLE_ENV} names {unknown}, which are not fast paths of this target; "
+            f"the fast paths are {[name for name in FAST_PATH_NAMES if name in served]}"
+        )
+    return frozenset(names)
+
+
+def disable_closure(names: FrozenSet[str]) -> FrozenSet[str]:
+    """``names`` and every path that depends on one of them."""
+    closed = set(names)
+    for name in names:
+        closed |= _DEPENDENTS[name]
+    return frozenset(closed)
+
+
+#: Every knob of this target set in an environment, as (name, value) sorted by name.
+Knobs = Tuple[Tuple[str, str], ...]
+
+
+def read_knobs(environ: Optional[Mapping[str, str]] = None) -> Knobs:
+    """Every knob of the family set in ``environ`` (the process environment by default)."""
+    env = os.environ if environ is None else environ
+    return tuple(sorted((k, v) for k, v in env.items() if k.startswith(KNOB_PREFIX)))
+
+
+def knob_digest(knobs: Knobs) -> str:
+    """A short digest of ``knobs`` for the log line."""
+    return hashlib.sha256(repr(knobs).encode()).hexdigest()[:16]
+
+
+@dataclass(frozen=True)
+class Plan:
+    """The composition this process resolved: which fast paths run, and under which mode."""
+
+    mode: ModelingV2Mode
+    # Paths the DISABLE knob names.
+    disabled: FrozenSet[str]
+    # Paths turned off only because a path they depend on is disabled.
+    disabled_dependents: FrozenSet[str]
+    # Every knob of this target set in the environment, as (name, value).
+    knobs: Knobs
+    # Whether the fast paths run their diagnostics (TRTLLM_MODELING_V2_VALIDATE).
+    validate: bool = False
+    # What the target that resolved the plan serves (``TargetPaths``); None serves every path.
+    target: Optional["TargetPaths"] = None
+
+    @property
+    def require(self) -> bool:
+        return self.mode is ModelingV2Mode.REQUIRE
+
+    @property
+    def served(self) -> FrozenSet[str]:
+        """The fast paths the target serves; the others are off in every plan of the target."""
+        return frozenset(FAST_PATH_NAMES) if self.target is None else self.target.served
+
+    def enabled(self, name: str) -> bool:
+        """Whether fast path ``name`` runs in this process.
+
+        Raises:
+            KeyError: if ``name`` is not a fast path, so a misspelled call site fails.
+        """
+        if name not in _BY_NAME:
+            raise KeyError(f"{name!r} is not a fast path of this target")
+        return (
+            name in self.served
+            and name not in self.disabled
+            and name not in self.disabled_dependents
+        )
+
+    def describe(self) -> str:
+        """One line naming the mode, the number of fast paths that run and what is off."""
+        on = sum(self.enabled(name) for name in FAST_PATH_NAMES)
+        words = f"modeling_v2={self.mode.value}, {on}/{len(self.served)} fast paths"
+        if self.disabled:
+            words += f", disabled {sorted(self.disabled)}"
+        if self.disabled_dependents:
+            words += f" and their dependents {sorted(self.disabled_dependents)}"
+        return words + f", knobs {knob_digest(self.knobs)}"
+
+
+#: The knobs that only one fast path reads; a target that does not serve the path refuses them.
+KNOB_PATHS: Mapping[str, str] = {ACCUMULATION_ENV: "routed_moe_into_hop"}
+
+
+def _refuse_removed_knobs(knobs: Knobs, served: Iterable[str] = FAST_PATH_NAMES) -> None:
+    """Raise if a variable with the knob prefix that is not in ``KNOBS`` is set, or a knob of a
+    fast path outside ``served``, naming each with what to use instead."""
+    served = frozenset(served)
+    refused = [
+        f"{name} ({REMOVED_KNOBS.get(name, _FIXED_HINT)})" for name, _ in knobs if name not in KNOBS
+    ]
+    refused += [
+        f"{name} (this target does not serve {KNOB_PATHS[name]})"
+        for name, _ in knobs
+        if name in KNOB_PATHS and KNOB_PATHS[name] not in served
+    ]
+    if refused:
+        raise ValueError(f"this target does not read {'; '.join(refused)}; unset them")
+
+
+def resolve(
+    mode: ModelingV2Mode,
+    environ: Optional[Mapping[str, str]] = None,
+    target: Optional["TargetPaths"] = None,
+) -> Plan:
+    """Read the knobs once and resolve the composition of ``target``; arms ``missed`` for ``mode``.
+
+    ``target`` is what the target serves (every fast path, with the family's tables, when None).
+
+    Raises:
+        ValueError: if ``DISABLE`` names something that is not a fast path the target serves, or
+            if a variable with the knob prefix that is not in ``KNOBS`` (or a knob of a path the
+            target does not serve) is set.
+    """
+    knobs = read_knobs(environ)
+    served = FAST_PATH_NAMES if target is None else target.served
+    _refuse_removed_knobs(knobs, served)
+    disabled = parse_disable(dict(knobs).get(DISABLE_ENV), served)
+    plan = Plan(
+        mode=mode,
+        disabled=disabled,
+        disabled_dependents=disable_closure(disabled) - disabled,
+        knobs=knobs,
+        validate=step_contract_enabled(),
+        target=target,
+    )
+    set_require(plan.require)
+    logger.info(f"{_LOG_TAG} dispatch: {plan.describe()}")
+    return plan
+
+
+# --------------------------------------------------------------------------------------------------
+# Rank uniformity
+# --------------------------------------------------------------------------------------------------
+
+#: Gathers one rank's value from every rank of a ``world_size``-rank group, or returns None
+#: (having logged why) when no such group is available.
+Allgather = Callable[[Any, int], Optional[List[Any]]]
+
+#: The key under which ``rank_uniform_inputs`` records the modeling_v2 mode.
+MODE_KEY = "modeling_v2"
+
+
+def rank_uniform_inputs(
+    knobs: Knobs, mode: ModelingV2Mode, environ: Optional[Mapping[str, str]] = None
+) -> Knobs:
+    """What every rank must agree on before the model is built, as (name, value) sorted by name.
+
+    ``knobs``, the shared switches the kernels read to choose a collective or a protocol
+    (``SHARED_ENV``) and the modeling_v2 validation switch, where set in ``environ`` (the
+    process environment by default), and ``mode``.
+    """
+    env = os.environ if environ is None else environ
+    shared = [(name, env[name]) for name in SHARED_ENV + (MODELING_V2_VALIDATE_ENV,) if name in env]
+    return tuple(sorted(list(knobs) + shared + [(MODE_KEY, mode.value)]))
+
+
+def _world_allgather(value: Any, world_size: int) -> Optional[List[Any]]:
+    """``value`` from every rank of the session's MPI communicator.
+
+    Returns None, with a warning, without MPI or when the communicator is not
+    the ``world_size`` ranks of this model: gathering over any other group
+    would either prove nothing or wait on ranks that never build this model.
+    """
+    from tensorrt_llm._utils import mpi_comm, mpi_disabled
+
+    if mpi_disabled():
+        logger.warning(
+            f"{_LOG_TAG} dispatch: no MPI communicator, so the ranks are not checked for agreement"
+        )
+        return None
+    comm = mpi_comm()
+    if comm.Get_size() != world_size:
+        logger.warning(
+            f"{_LOG_TAG} dispatch: the MPI communicator has {comm.Get_size()} ranks, not the "
+            f"{world_size} of this model, so the ranks are not checked for agreement"
+        )
+        return None
+    return comm.allgather(value)
+
+
+def _rank_table(gathered: Sequence[Knobs]) -> str:
+    knob_names = sorted({name for knobs in gathered for name, _ in knobs})
+    lines = []
+    for rank, knobs in enumerate(gathered):
+        values = dict(knobs)
+        knob_words = ", ".join(f"{name}={values.get(name, '<unset>')!r}" for name in knob_names)
+        lines.append(f"  rank {rank}: knobs {knob_digest(knobs)} ({knob_words or 'none set'})")
+    return "\n".join(lines)
+
+
+def check_rank_uniform(
+    knobs: Knobs, world_size: int, allgather: Optional[Allgather] = None
+) -> None:
+    """Raise on every rank unless every rank read the same knob values.
+
+    Called once, at construction, on every rank, before the knobs are parsed:
+    the allgather is itself a collective, so a rank must not raise before it.
+    ``knobs`` is ``rank_uniform_inputs`` of the knobs. ``allgather`` defaults to
+    the session's MPI communicator.
+
+    Raises:
+        RuntimeError: with one line per rank, when the ranks disagree.
+    """
+    gathered = (allgather or _world_allgather)(knobs, world_size)
+    if gathered is None or all(peer == knobs for peer in gathered):
+        return
+    raise RuntimeError(
+        f"{_LOG_TAG} the ranks read different knob values, so they would resolve different "
+        "fast paths; a collective fast path would wait forever on the ranks that skip it. "
+        f"Set the {KNOB_PREFIX}* variables, {', '.join(SHARED_ENV)}, "
+        f"{MODELING_V2_VALIDATE_ENV} and {MODE_KEY} identically on every rank (before the "
+        f"ranks start):\n{_rank_table(gathered)}"
+    )
+
+
+#: The fast-path state of a built model, module by module in construction order: per module,
+#: its path and, per fast-path object attached to it, the attribute, the object's class and its
+#: rank-independent flags.
+AdoptedState = Tuple[Tuple[str, Tuple[Tuple[str, ...], ...]], ...]
+
+
+def check_rank_uniform_adopted(
+    state: AdoptedState, world_size: int, allgather: Optional[Allgather] = None
+) -> None:
+    """Raise on every rank unless every rank attached the same fast-path state to its modules.
+
+    Called once, on every rank, after the fast paths are attached; a fast path that engaged on
+    some ranks only would run its collective on those ranks alone.
+
+    Raises:
+        RuntimeError: naming the first module whose state differs, with one line per rank.
+    """
+    gathered = (allgather or _world_allgather)(state, world_size)
+    if gathered is None or all(peer == state for peer in gathered):
+        return
+    depth = max(len(peer) for peer in gathered)
+    first = next(
+        i
+        for i in range(depth)
+        if len({peer[i] if i < len(peer) else None for peer in gathered}) > 1
+    )
+    lines = []
+    for rank, peer in enumerate(gathered):
+        lines.append(f"  rank {rank}: {peer[first] if first < len(peer) else '<no module>'}")
+    raise RuntimeError(
+        f"{_LOG_TAG} the ranks attached different fast-path state to their modules; a "
+        "collective fast path would wait forever on the ranks that skip it. First module that "
+        "differs:\n" + "\n".join(lines)
+    )
+
+
+# --------------------------------------------------------------------------------------------------
+# Fallback policy
+# --------------------------------------------------------------------------------------------------
+
+
+class FastPathUnavailable(RuntimeError):
+    """A step inside a fast path's shapes could not take it under ``modeling_v2='require'``."""
+
+
+_lock = threading.Lock()
+_require = [False]
+_logged: set = set()
+_missed: Dict[str, int] = {}
+
+
+def set_require(require: bool) -> None:
+    """Whether ``missed`` raises (``modeling_v2='require'``); set by ``resolve``."""
+    _require[0] = bool(require)
+
+
+def require_mode() -> bool:
+    return _require[0]
+
+
+def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None:
+    key = (path, reason)
+    with _lock:
+        if key in _logged:
+            return
+        _logged.add(key)
+    level(message)
+
+
+def missed(path: str, rows: Optional[int], reason: str, *, per_call: bool = False) -> None:
+    """``path`` takes this step's shape but cannot run: raise under require, else log once and count.
+
+    A miss decided per call (``per_call``) never raises: a step that is already being served
+    falls back, and the miss is counted and logged once. The caller falls back after it returns.
+    """
+    at = "" if rows is None else f" at {rows} rows"
+    if _require[0] and not per_call:
+        raise FastPathUnavailable(
+            f"modeling_v2='require': the GLM-5.2 decode fast path {path!r} takes this step{at} "
+            f"but cannot run: {reason}"
+        )
+    with _lock:
+        _missed[path] = _missed.get(path, 0) + 1
+    _log_once(
+        logger.warning,
+        path,
+        reason,
+        f"{_LOG_TAG} {path}/missed{at}: {reason}; falling back (logged once per reason)",
+    )
+
+
+def uncovered(path: str, reason: str) -> None:
+    """The step is outside ``path``'s shapes: log once per (path, reason); the next tier runs."""
+    _log_once(logger.info, path, reason, f"{_LOG_TAG} {path}/uncovered: {reason} (logged once)")
+
+
+def missed_counts() -> Dict[str, int]:
+    """Misses per fast path in this process: per-call misses in every mode, the others under
+    ``auto`` (``require`` raises on them instead)."""
+    with _lock:
+        return dict(_missed)
+
+
+def reset_fallback_state() -> None:
+    """Forget the logged reasons, the miss counts and the installed plan."""
+    with _lock:
+        _logged.clear()
+        _missed.clear()
+    _INSTALLED[0] = None
+    _VALIDATE[0] = False
+    _TARGET[0] = DEFAULT_TARGET
+    classify_decline.cache_clear()
+
+
+# --------------------------------------------------------------------------------------------------
+# Phases
+# --------------------------------------------------------------------------------------------------
+
+
+class Phase(enum.Enum):
+    """Which target takes a model step."""
+
+    PREFILL = "prefill"
+    DECODE = "decode"
+
+
+_PHASE: List[Optional[Phase]] = [None]
+
+
+class _Step:
+    """The scope of one model step under a phase; restores the enclosing phase on exit."""
+
+    __slots__ = ("_phase", "_previous")
+
+    def __init__(self, phase: Phase) -> None:
+        self._phase = phase
+        self._previous: Optional[Phase] = None
+
+    def __enter__(self) -> None:
+        self._previous = _PHASE[0]
+        _PHASE[0] = self._phase
+
+    def __exit__(self, *exc_info: Any) -> None:
+        _PHASE[0] = self._previous
+
+
+def step(phase: Phase) -> _Step:
+    """A context manager that runs the enclosed model step under ``phase``; scopes nest."""
+    return _Step(phase)
+
+
+def phase() -> Optional[Phase]:
+    """The phase of the model step being run, or None outside a step."""
+    return _PHASE[0]
+
+
+# --------------------------------------------------------------------------------------------------
+# Decline classes
+# --------------------------------------------------------------------------------------------------
+
+
+class Decline(enum.Enum):
+    """What a fast path's decline means."""
+
+    #: The path is turned off: neither a miss nor outside its shapes.
+    DISABLED = "disabled"
+    #: Outside the path's shapes (a layer it does not serve, a row count, a step kind).
+    UNCOVERED = "uncovered"
+    #: Decided once, at construction or in the warm-up, inside the path's layers, with a
+    #: precondition absent: ``missed``, which raises under ``require``.
+    MISSED = "missed"
+    #: Decided per call, with a resource absent: counted and logged once, never raised.
+    MISSED_PER_CALL = "missed per call"
+
+
+_TARGET_LAYERS = frozenset(range(78))
+_DRAFT_LAYERS = frozenset({78})
+# The decoder layers that run the full indexer: the three dense layers and every fourth layer
+# from layer 6. The other target layers reuse an earlier layer's top-k selection.
+_FULL_INDEXER_LAYERS = frozenset((0, 1, 2) + tuple(range(6, 78, 4)))
+_SHARED_INDEXER_LAYERS = _TARGET_LAYERS - _FULL_INDEXER_LAYERS
+
+#: Per fast path and per layer-keyed capability key family (the key without its layer index),
+#: the layers on which the path engages on every rank when this target serves its checkpoint.
+#: A decline on a layer outside its family's cells is outside the path's shapes.
+LAYER_CELLS: Mapping[str, Mapping[str, FrozenSet[int]]] = {
+    "absorbed_query": {"layer": _TARGET_LAYERS},
+    "draft_absorbed_query": {"layer": _DRAFT_LAYERS},
+    "shared_indexer_absorbed_query": {"layer": _SHARED_INDEXER_LAYERS},
+    # Only the shared-indexer layer that the last dense layer feeds.
+    "fused_kv_indexer_projection": {"layer": frozenset({3})},
+    "indexer_layer_front": {"target_layer": _FULL_INDEXER_LAYERS, "draft_layer": _DRAFT_LAYERS},
+    # Boundaries into the shared-indexer layers that a MoE layer feeds.
+    "deferred_norm": {"scope[shared]_into_layer": _SHARED_INDEXER_LAYERS - {3}},
+    "deferred_norm_evict_hints": {"projection_layer": _SHARED_INDEXER_LAYERS - {3}},
+    "sparse_mla_decode": {
+        "layer": _TARGET_LAYERS | _DRAFT_LAYERS,
+        "call:layer": _TARGET_LAYERS | _DRAFT_LAYERS,
+    },
+    "attention_output_tail": {"layer": _TARGET_LAYERS},
+    "oproj_weight_staging": {"layer": _TARGET_LAYERS},
+}
+
+#: Reason fragments that turn a path off rather than decline it.
+DISABLED_REASONS: Tuple[str, ...] = ("DISABLE",)
+#: Reason fragments that name a layer kind or a composition the path does not serve.
+LAYER_KIND_REASONS: Tuple[str, ...] = (
+    "target layers only",
+    "draft (MTP) layers only",
+    "full indexer layer",
+    "shared-indexer layer",
+    "layers not selected (",
+    "draft or unindexed layer",
+    "not a target layer",
+    "is a full-indexer layer",
+    "not a DSA layer",
+    "not MLA",
+    "needs the post_attention_hop hop",
+    "no post_attention_hop hop",
+    "post_attention_hop without sparse_mla_decode",
+)
+#: Reason fragments (any case) of a per-call decline that name a missing resource.
+RESOURCE_REASONS: Tuple[str, ...] = (
+    "workspace",
+    "compile failed",
+    "not compiled",
+    "not loaded",
+    "not built",
+    "failed",
+    "unavailable",
+    "not warmed up",
+    "before CUDA graph capture",
+    "packet does not match",
+    "does not forward",
+    "another rank",
+    "MNNVL",
+    "fabric",
+)
+
+
+def _any_of(fragments: Iterable[str], flags: int = 0) -> "re.Pattern[str]":
+    return re.compile("|".join(map(re.escape, fragments)), flags)
+
+
+_DISABLED_RE = _any_of(DISABLED_REASONS)
+_LAYER_KIND_RE = _any_of(LAYER_KIND_REASONS)
+_RESOURCE_RE = _any_of(RESOURCE_REASONS, re.IGNORECASE)
+
+
+@dataclass(frozen=True)
+class DeclineSite:
+    """Capability keys of a fast path that are decided alike.
+
+    ``key`` matches a whole key; a key that names a layer captures its index as ``layer``, and
+    ``cells`` names the ``LAYER_CELLS`` family of those layers. A ``per_call`` site decides on
+    every call; the others decide once, at construction or in the warm-up. ``uncovered``
+    matches reasons that are outside the path's shapes at this site.
+    """
+
+    path: str
+    key: "re.Pattern[str]"
+    per_call: bool = False
+    cells: Optional[str] = None
+    uncovered: Optional["re.Pattern[str]"] = None
+
+
+def _site(
+    path: str,
+    key: str,
+    per_call: bool = False,
+    cells: Optional[str] = None,
+    uncovered: Optional[str] = None,
+) -> DeclineSite:
+    return DeclineSite(
+        path,
+        re.compile(key),
+        per_call,
+        cells,
+        None if uncovered is None else re.compile(re.escape(uncovered)),
+    )
+
+
+_LAYER = r"(?P<layer>\d+)"
+_EXPERT_PREFETCH = r"expert_prefetch\[.*\]"
+
+#: Every capability key the fast paths log, by where and when it is decided.
+DECLINE_KINDS: Tuple[DeclineSite, ...] = (
+    _site("absorbed_query", "form"),
+    _site("absorbed_query", f"layer{_LAYER}", cells="layer"),
+    _site("draft_absorbed_query", f"layer{_LAYER}", cells="layer"),
+    _site("shared_indexer_absorbed_query", f"layer{_LAYER}", cells="layer"),
+    _site("fused_kv_indexer_projection", f"layer{_LAYER}", cells="layer"),
+    _site("indexer_layer_front", f"target_layer{_LAYER}", cells="target_layer"),
+    _site("indexer_layer_front", f"draft_layer{_LAYER}", cells="draft_layer"),
+    _site("indexer_layer_front", "rope"),
+    _site("indexer_layer_front", r"cp_indexer_decode_.+", per_call=True),
+    _site("deferred_norm", "deferred_hop"),
+    _site(
+        "deferred_norm", rf"scope\[shared\]_into_layer{_LAYER}", cells="scope[shared]_into_layer"
+    ),
+    _site("deferred_norm", r"deferred_hop_routed\[.*\]"),
+    _site("deferred_norm", f"projection_layer{_LAYER}", cells="projection_layer"),
+    _site("deferred_norm_evict_hints", "deferred_hop"),
+    _site("deferred_norm_evict_hints", r"deferred_hop_routed\[.*\]"),
+    _site("deferred_norm_evict_hints", f"projection_layer{_LAYER}", cells="projection_layer"),
+    _site("rope_kv_append", "rope|scatter_defer|bridge", per_call=True),
+    _site("indexer_topk", "indexer_decode", per_call=True),
+    _site("cross_rank_topk", "cp_indexer_decode", per_call=True),
+    _site("cross_rank_topk_cluster_merge", "cp_indexer_decode", per_call=True),
+    _site("cross_rank_topk_global_emit", "cp_indexer_decode", per_call=True),
+    _site("sparse_mla_decode", f"layer{_LAYER}", cells="layer"),
+    _site("sparse_mla_decode", r"request:.*", per_call=True),
+    _site("sparse_mla_decode", f"call:layer{_LAYER}", per_call=True, cells="call:layer"),
+    _site("attention_output_tail", f"layer{_LAYER}", cells="layer"),
+    _site("attention_output_tail_prefetch", f"layer{_LAYER}", cells="layer"),
+    _site("attention_output_tail_prefetch", "schedule"),
+    _site("attention_output_into_hop", f"layer{_LAYER}", cells="layer"),
+    _site("attention_output_into_hop", "into_hop_transport"),
+    _site("oproj_weight_staging", f"layer{_LAYER}", cells="layer"),
+    _site("post_attention_hop", r"post_attn_hop(:.*)?"),
+    _site("hop_routing", "routing_hop"),
+    _site("hop_routing", r"routed_moe_fork rows=.*", per_call=True),
+    _site("hop_routing", _EXPERT_PREFETCH),
+    _site("router_gemm", r"router:.*"),
+    _site("routed_moe", r"rows=.*", per_call=True),
+    _site("routed_moe", _EXPERT_PREFETCH),
+    _site("routed_moe_into_hop", r"rows=.*", per_call=True),
+    _site("routed_moe_into_hop", r"post_moe_routed_input.*"),
+    _site("shared_expert_prefetch", r"routed rows=.*", per_call=True),
+    _site("shared_expert_prefetch", _EXPERT_PREFETCH),
+    _site("post_moe_hop", r"cluster_kernel(_compile)?|post_moe_hop\[.*\]"),
+    _site("cross_layer_expert_prefetch", r"expert_prefetch|fused_predict|xc_prefetch|link\[.*\]"),
+    _site("cross_layer_expert_prefetch", r"spec_experts(_declined)?", per_call=True),
+    _site("dense_mlp", r"dense_mlp:.*", per_call=True),
+    _site("draft_front", "draft_front"),
+    _site("draft_front", r"draft_front_call(_fallback)?", per_call=True),
+    _site("draft_routed_moe", r"rows=.*", per_call=True),
+    _site("draft_live_rows", r"live_row_hop(:.*)?"),
+    _site("draft_live_rows", r"live_row_hop_call(:.*)?|step0:.*", per_call=True),
+    _site("draft_tail", r"draft_tail\[tp\d+\]"),
+    _site("draft_tail", r"draft_tail\[(step: .*|graph|eager)\]", per_call=True),
+    _site("draft_loop", r"(rows|draft_step|finalize|draft_loop):.*", per_call=True),
+    _site("draft_loop", "draft_embedding"),
+    _site("decode_prologue", r"target_(prologue|embedding)(\[.*\])?|take_embedding", per_call=True),
+    _site("deferred_input_prep", r"prep(\[.*\])?", per_call=True),
+    # Before the overlap executor loop declares its domain every call declines; the decision
+    # is final only once the loop has declared it.
+    _site(
+        "deferred_input_prep",
+        "loop_domain",
+        uncovered="no overlap executor loop declared its domain",
+    ),
+    _site("deferred_input_prep", "engaged"),
+    _site("sampler_store", "engaged"),
+    _site("execution_stream_loop", "engaged"),
+    _site("host_page_retention", "host_policy"),
+)
+
+_SITES_BY_PATH: Mapping[str, Tuple[DeclineSite, ...]] = {
+    name: tuple(s for s in DECLINE_KINDS if s.path == name) for name in FAST_PATH_NAMES
+}
+
+
+@dataclass(frozen=True)
+class TargetPaths:
+    """What one target of the family serves of the fast paths, and where their code lives.
+
+    ``served`` is the fast paths the target runs: a path outside it is off in every plan of the
+    target, ``DISABLE`` may not name it and a knob only it reads is refused. ``layer_cells`` is
+    ``LAYER_CELLS`` for the target's layers and shapes; ``uncovered`` lists, per path, reason
+    fragments of declines that are outside the target's shapes at every site of the path.
+    ``kernel_modules`` names the module of each path whose kernel package is the target's own;
+    every other path's package is the family's (``KERNELS_PACKAGE``).
+
+    Raises:
+        ValueError: if a name is not a fast path, or a served path extends one the target does not
+            serve.
+    """
+
+    served: FrozenSet[str] = frozenset(FAST_PATH_NAMES)
+    layer_cells: Mapping[str, Mapping[str, FrozenSet[int]]] = field(
+        default_factory=lambda: LAYER_CELLS
+    )
+    uncovered: Mapping[str, Tuple[str, ...]] = field(default_factory=dict)
+    kernel_modules: Mapping[str, str] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        named = set(self.served) | set(self.layer_cells) | set(self.uncovered)
+        unknown = sorted((named | set(self.kernel_modules)) - set(FAST_PATH_NAMES))
+        if unknown:
+            raise ValueError(f"{unknown} are not fast paths of the family")
+        lacking = sorted(
+            f"{name} extends {base}"
+            for name in self.served
+            for base in _BY_NAME[name].depends_on
+            if base not in self.served
+        )
+        if lacking:
+            raise ValueError(f"a served fast path extends one the target does not serve: {lacking}")
+
+
+#: Every fast path, with the family's tables.
+DEFAULT_TARGET = TargetPaths()
+# The tables of the installed plan's target.
+_TARGET: List[TargetPaths] = [DEFAULT_TARGET]
+
+
+def _decline_site(path: str, key: str) -> Optional[Tuple[DeclineSite, "re.Match[str]"]]:
+    """The site of ``path`` whose key pattern matches all of ``key``, with the match."""
+    for site in _SITES_BY_PATH.get(path, ()):
+        match = site.key.fullmatch(key)
+        if match is not None:
+            return site, match
+    return None
+
+
+@functools.lru_cache(maxsize=4096)
+def classify_decline(path: str, key: str, reason: str) -> Decline:
+    """What a decline of fast path ``path`` at capability ``key`` for ``reason`` means.
+
+    In order: a reason that turns the path off is ``DISABLED``; a key no site of the path
+    matches is ``UNCOVERED``; a per-call decline is ``MISSED_PER_CALL`` when its reason names a
+    missing resource and ``UNCOVERED`` otherwise; a decline decided once is ``UNCOVERED`` when
+    its reason names a layer kind or a composition the path does not serve, or its layer is
+    outside the path's cells, and ``MISSED`` otherwise. The cells and the target's uncovered
+    reasons are the installed plan's target's (``TargetPaths``). Memoized: a repeated decline
+    costs one lookup.
+    """
+    if _DISABLED_RE.search(reason):
+        return Decline.DISABLED
+    found = _decline_site(path, key)
+    if found is None:
+        return Decline.UNCOVERED
+    site, match = found
+    if site.uncovered is not None and site.uncovered.search(reason):
+        return Decline.UNCOVERED
+    if any(fragment in reason for fragment in _TARGET[0].uncovered.get(path, ())):
+        return Decline.UNCOVERED
+    if site.per_call:
+        return Decline.MISSED_PER_CALL if _RESOURCE_RE.search(reason) else Decline.UNCOVERED
+    if _LAYER_KIND_RE.search(reason):
+        return Decline.UNCOVERED
+    if site.cells is not None:
+        cells = _TARGET[0].layer_cells.get(path, {}).get(site.cells, frozenset())
+        if int(match["layer"]) not in cells:
+            return Decline.UNCOVERED
+    return Decline.MISSED
+
+
+def _record_decline(path: str, key: str, reason: str) -> None:
+    kind = classify_decline(path, key, reason)
+    if kind is Decline.MISSED:
+        missed(path, None, f"{key}: {reason}" if reason else key)
+    elif kind is Decline.MISSED_PER_CALL:
+        missed(path, None, reason or key, per_call=True)
+
+
+# --------------------------------------------------------------------------------------------------
+# Selection
+# --------------------------------------------------------------------------------------------------
+
+#: The parts of the decode step's input preparation, each a fast path of its own: the step-input
+#: stager, the sampler's store and the overlap loop on the execution stream.
+PART_PATHS: Mapping[str, str] = {
+    "prep": "deferred_input_prep",
+    "store": "sampler_store",
+    "stream": "execution_stream_loop",
+}
+_PART_ONLY_PATHS = frozenset(PART_PATHS.values()) - {PART_PATHS["prep"]}
+
+#: Every fast path except the parts that only ``deferred_input_prep`` brings along.
+SELECTOR_PATHS: Tuple[str, ...] = tuple(n for n in FAST_PATH_NAMES if n not in _PART_ONLY_PATHS)
+
+#: The family's kernel packages, one directory per fast path under ``kernels/``.
+KERNELS_PACKAGE = f"{__package__}.kernels"
+#: The fast paths whose code lives in the kernel directory of the path they extend.
+KERNEL_DIRECTORY: Mapping[str, str] = {
+    "draft_absorbed_query": "absorbed_query",
+    "shared_indexer_absorbed_query": "absorbed_query",
+    "cross_rank_topk": "indexer_topk",
+    "attention_output_tail_prefetch": "attention_output_tail",
+    "attention_output_into_hop": "attention_output_tail",
+    "routed_moe_into_hop": "routed_moe",
+    "sampler_store": "deferred_input_prep",
+    "execution_stream_loop": "deferred_input_prep",
+}
+
+
+def _bases() -> Dict[str, FrozenSet[str]]:
+    """Every path each path extends, directly or transitively."""
+    closure: Dict[str, FrozenSet[str]] = {}
+
+    def visit(name: str) -> FrozenSet[str]:
+        if name not in closure:
+            bases = set()
+            for base in _BY_NAME[name].depends_on:
+                bases |= {base} | visit(base)
+            closure[name] = frozenset(bases)
+        return closure[name]
+
+    for name in FAST_PATH_NAMES:
+        visit(name)
+    return closure
+
+
+#: The paths each path has no effect without (the transitive closure of ``depends_on``).
+REQUIRED_FEATURES: Mapping[str, FrozenSet[str]] = {
+    name: bases for name, bases in _bases().items() if bases
+}
+
+# Numerical variants (a variant has its own name). Per fast path with variants: the knob that
+# selects one (None: a single variant) and the variant names, default first. A path without an
+# entry has one numerical form, equal to the generic path's or held to a frozen error budget.
+NUMERICAL_VARIANTS: Dict[str, Tuple[Optional[str], Tuple[str, ...]]] = {
+    # Exact top-k over each page owner's indexer K, identical on every rank.
+    "cross_rank_topk": (None, ("rank-consistent",)),
+    # The cluster merge's rows are byte-identical to the cross-rank top-k's (same set and order).
+    "cross_rank_topk_cluster_merge": (None, ("rank-consistent",)),
+    # Its id rows are the cluster merge's; its pool and copy rows are the values of the
+    # rope/KV-append top-k bridge.
+    "cross_rank_topk_global_emit": (None, ("rank-consistent",)),
+    # What the routed experts write into the post-MoE hop: fp32 sums the top-k contributions in
+    # another order (nondeterministic from run to run); identity is bit-exact.
+    "routed_moe_into_hop": (ACCUMULATION_ENV, ("fp32", "identity")),
+    # The next layer's RMSNorm deferred into its projection (changes rounding), only on boundaries
+    # into shared-indexer layers: full-indexer layers keep the post-MoE hop's normalization.
+    "deferred_norm": (None, ("shared",)),
+    # One projection launch for the fused-A GEMM and the indexer key/weights GEMM: kva keeps the
+    # fused-A order (bit-identical); kw (full-indexer and draft layers) takes the same mma.sync
+    # order instead of the cuBLAS split-K GEMM (rounds differently). This target runs it on the
+    # dense-fed shared-indexer layer, which computes kva only.
+    "fused_kv_indexer_projection": (None, ("kw-fused-a-order",)),
+}
+
+#: Prefixes of the set variables the record line lists (``side_selectors``).
+SIDE_SELECTOR_PREFIXES = (KNOB_PREFIX,)
+#: Shared TensorRT-LLM switches that the kernels also read.
+SHARED_ENV = ("TLLM_TEST_MNNVL", "TRTLLM_ENABLE_PDL", "TRTLLM_FORCE_MNNVL_AR")
+
+
+def feature_variants(
+    paths: Iterable[str], environ: Optional[Mapping[str, str]] = None
+) -> Dict[str, str]:
+    """The numerical variant of every path in ``paths`` that has some.
+
+    Raises:
+        ValueError: if a knob names an unknown variant of a path in ``paths``.
+    """
+    env = os.environ if environ is None else environ
+    variants = {}
+    for path in sorted(paths):
+        if path not in NUMERICAL_VARIANTS:
+            continue
+        name, choices = NUMERICAL_VARIANTS[path]
+        raw = env.get(name) if name else None
+        value = (raw or "").strip().lower() or choices[0]
+        if value not in choices:
+            raise ValueError(
+                f"{name}={raw!r} is not a {path} variant; known variants are {list(choices)} "
+                f"(unset selects {choices[0]!r})"
+            )
+        variants[path] = value
+    return variants
+
+
+def side_selectors(environ: Optional[Mapping[str, str]] = None) -> Dict[str, str]:
+    """Every set knob of the family."""
+    env = os.environ if environ is None else environ
+    return {name: env[name] for name in sorted(env) if name.startswith(SIDE_SELECTOR_PREFIXES)}
+
+
+def shared_env(environ: Optional[Mapping[str, str]] = None) -> Dict[str, str]:
+    """The set shared TensorRT-LLM switches that the kernels read."""
+    env = os.environ if environ is None else environ
+    return {name: env[name] for name in SHARED_ENV if name in env}
+
+
+def feature_record(paths: Iterable[str], environ: Optional[Mapping[str, str]] = None) -> str:
+    """The one-line record of a selection.
+
+    ``paths=`` lists every path as ``path`` or ``path[variant]``; ``side_selectors=`` and
+    ``shared_env=`` are JSON objects of the set variables.
+
+    Raises:
+        ValueError: as ``feature_variants``.
+    """
+    variants = feature_variants(paths, environ)
+    listed = ",".join(f"{p}[{variants[p]}]" if p in variants else p for p in sorted(paths))
+
+    def _json(values: Dict[str, str]) -> str:
+        return json.dumps(values, sort_keys=True, separators=(",", ":"))
+
+    return (
+        f"paths={listed or '<none>'} "
+        f"side_selectors={_json(side_selectors(environ))} "
+        f"shared_env={_json(shared_env(environ))}"
+    )
+
+
+def _log_info(message: str) -> None:
+    logger.info(message)
+
+
+# The fast paths of the installed plan; None while no plan is installed.
+_INSTALLED: List[Optional[FrozenSet[str]]] = [None]
+# Whether the installed plan runs the fast paths' diagnostics.
+_VALIDATE: List[bool] = [False]
+
+
+def install(plan: Plan, environ: Optional[Mapping[str, str]] = None) -> None:
+    """Answer ``is_enabled`` from ``plan`` from now on, and log its record line once.
+
+    Raises:
+        ValueError: as ``feature_variants``, for a knob that names an unknown variant of an
+            enabled path.
+    """
+    paths = frozenset(n for n in FAST_PATH_NAMES if plan.enabled(n))
+    record = feature_record(paths, environ)
+    _INSTALLED[0] = paths
+    _VALIDATE[0] = plan.validate
+    _TARGET[0] = DEFAULT_TARGET if plan.target is None else plan.target
+    classify_decline.cache_clear()
+    _log_info(f"{_LOG_TAG} dispatch: record {record}")
+
+
+def installed() -> Optional[FrozenSet[str]]:
+    """The fast paths of the installed plan, or None."""
+    return _INSTALLED[0]
+
+
+def enabled_features() -> FrozenSet[str]:
+    """The fast paths that run in this process: the installed plan's, else none."""
+    plan_paths = _INSTALLED[0]
+    return plan_paths if plan_paths is not None else frozenset()
+
+
+def validate_mode() -> bool:
+    """Whether the installed plan runs the fast paths' diagnostics (False without a plan)."""
+    return _VALIDATE[0]
+
+
+def is_enabled(path: str) -> bool:
+    """Whether fast path ``path`` runs in this process.
+
+    Raises:
+        ValueError: if ``path`` is not a fast path of the family.
+    """
+    if path not in _BY_NAME:
+        raise ValueError(f"{path!r} is not a fast path of this target")
+    return path in enabled_features()
+
+
+# Entry points per fast path. The kernel package of a path registers them when it is imported;
+# a call site looks one up with ``get_impl`` only after ``is_enabled`` returned True.
+_REGISTRY: Dict[str, Dict[str, Callable]] = {}
+
+# Per-rank capability decisions, logged once per (path, key).
+_CAPABILITY_LOGGED: Dict[Tuple[str, str], bool] = {}
+
+
+def register(path: str, name: str, fn: Callable) -> Callable:
+    """Register ``fn`` as entry point ``name`` of fast path ``path``.
+
+    Raises:
+        ValueError: if ``path`` is not a fast path of the family.
+    """
+    if path not in _BY_NAME:
+        raise ValueError(f"{path!r} is not a fast path of this target")
+    _REGISTRY.setdefault(path, {})[name] = fn
+    return fn
+
+
+def get_impl(path: str, name: str) -> Callable:
+    """Import the kernel package of ``path`` on first use and return its entry ``name``.
+
+    The package is the installed plan's target's own (``TargetPaths.kernel_modules``) or the
+    family's.
+
+    Raises:
+        KeyError: if the package does not register ``name`` for ``path``.
+    """
+    if path not in _REGISTRY:
+        module = _TARGET[0].kernel_modules.get(path)
+        importlib.import_module(module or f"{KERNELS_PACKAGE}.{KERNEL_DIRECTORY.get(path, path)}")
+    entries = _REGISTRY.get(path, {})
+    if name not in entries:
+        raise KeyError(
+            f"fast path {path!r} does not register {name!r}; registered: {sorted(entries)}"
+        )
+    return entries[name]
+
+
+def log_capability(path: str, key: str, eligible: bool, reason: str = "") -> bool:
+    """Log a per-rank capability decision once per (path, key) and return ``eligible``.
+
+    Every decline, logged or not, is also classified (``classify_decline``): a miss decided
+    once goes to ``missed``, which raises under ``require``; a miss decided per call is counted
+    and logged once; a decline outside the path's shapes or of a path turned off adds nothing
+    to the capability line.
+    """
+    if (path, key) not in _CAPABILITY_LOGGED:
+        _CAPABILITY_LOGGED[(path, key)] = eligible
+        state = "engaged" if eligible else "native fallback"
+        suffix = f" ({reason})" if reason else ""
+        _log_info(f"{_LOG_TAG} {path}/{key}: {state}{suffix}")
+    if not eligible:
+        _record_decline(path, key, reason)
+    return eligible
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/__init__.py
new file mode 100644
index 000000000000..61038c261bb1
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Kernels two or more glm_moe_dsa targets call, one package per fast path."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/__init__.py
new file mode 100644
index 000000000000..11a8ee3db664
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/__init__.py
@@ -0,0 +1,48 @@
+# 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.
+"""absorbed_query: target query/projection region (the query_front contract, token-wise part of the
+index_selection contract).
+
+Registered entry points (``dispatch.get_impl("absorbed_query", name)``):
+
+* ``dsa_proj(mla, position_ids, hidden_states, attn_metadata)`` ->
+  ``Optional[AbsorbedQueryPacket]`` for target layers (``absorbed_query``); ``None``
+  means "run the native path".
+* ``mtp_dsa_proj(mla, position_ids, hidden_states, attn_metadata)`` -> the
+  same packet for the one-model MTP draft layer at every draft step
+  (package ``draft_absorbed_query``, which extends absorbed_query to the draft layer); ``None`` for
+  target layers and outside the support domain.
+* ``shared_dsa_proj(mla, position_ids, hidden_states, attn_metadata)`` ->
+  the preabsorbed packet of a shared-indexer target layer from one kernel
+  (package ``shared_indexer_absorbed_query``, which extends absorbed_query's shared-layer path);
+  ``None`` for full-indexer and draft layers and outside the support domain.
+"""
+
+from ...dispatch import register
+from .query_proj import AbsorbedQueryPacket, AbsorbedQueryProjection, dsa_proj, mtp_dsa_proj
+from .shared_indexer import SharedIndexerProjection, shared_dsa_proj
+
+register("absorbed_query", "dsa_proj", dsa_proj)
+register("absorbed_query", "mtp_dsa_proj", mtp_dsa_proj)
+register("absorbed_query", "shared_dsa_proj", shared_dsa_proj)
+
+__all__ = [
+    "AbsorbedQueryPacket",
+    "AbsorbedQueryProjection",
+    "SharedIndexerProjection",
+    "dsa_proj",
+    "mtp_dsa_proj",
+    "shared_dsa_proj",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/kernels.py
new file mode 100644
index 000000000000..be6c61752cd6
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/kernels.py
@@ -0,0 +1,1764 @@
+# 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.
+"""Triton kernels of the absorbed_query/projection package.
+
+``qproj_kernel`` consumes the native fused ``kv_a_proj_with_mqa`` output
+``kva = [q_a | compressed_kv | k_pe]`` of ``R <= BLOCK_M`` token rows (and,
+with ``PRENORM``, the native ``q_a_layernorm`` output) and produces, in one
+launch:
+
+* ``q = q_b_proj(rmsnorm(q_a))`` ``[R, H * 256]`` (BF16, the native ``q``),
+* optionally ``qr = rmsnorm(q_a)`` ``[R, 2048]`` for the native indexer,
+* ``latent = [rmsnorm(compressed_kv) | k_pe]`` ``[R, 576]``,
+* optionally the absorbed query ``fused_q[:, h, :512] = q_nope[:, h] @ W_UK[h]``
+  written into the caller's ``[R, H, 576]`` buffer.
+
+``qproj_indexer_kernel`` additionally produces the token-wise indexer
+outputs of full indexer layers.
+
+Numerics follow the native chain bit for bit on the default path
+(``PRENORM``, the native kv_a GEMM, ``q_a_layernorm``, fused ``wk`` /
+``weights_proj`` GEMM and ``k_norm`` as bridges): ``q`` and the indexer
+``wq_b`` output are FP32-accumulated MMA chains in ascending K order (equal to
+the native GEMMs), the ``compressed_kv`` RMSNorm reproduces the native
+FlashInfer kernel's reduction order (``_fi_rmsnorm_rows_exact``), the
+indexer RoPE its FMA pattern (``_rope_pairs``) and the FP8 tail the native
+``fusedCatFp8Kernel``; ``q`` is rounded to BF16 before the absorption (native
+intermediate). The experimental paths (own hidden GEMV, prefetching kernel's
+in-kernel q_a norm, ``PRENORM=False``) are not bit-exact.
+
+Program roles, in program-id order (each kernel documents its own list):
+``q_b`` producer tiles ``[0, N_QB)`` (each ``BN`` output columns over the full
+``K``; ``TPH = HEAD_DIM // BN`` tiles per head, the first ``NOPE // BN`` of a
+head are its no-RoPE tiles), then indexer ``wq_b`` producer tiles, then the
+programs that only consume.
+
+Inter-program synchronization contract
+--------------------------------------
+* A producer stores its tile, executes a CTA barrier and then one thread
+  releases the tile with ``atom.add.release.gpu`` on a per-head counter. A
+  consumer spins with ``atom.add.acquire.gpu`` (adding 0) until the counter
+  reaches the number of producer tiles it needs, executes a CTA barrier and
+  only then reads the tiles, with ``ld.global.cg`` (L2, never a stale L1 line).
+* Every counter is per module (per layer) and is zero between launches. The
+  last consumer of a counter resets it (relaxed exchange) after every
+  consumer of that counter has read its tiles: for the absorption counters a
+  second ``done`` counter (``acq_rel``) counts the consumers of a head; an
+  indexer-head counter has a single consumer. Counters are touched only after
+  ``griddepcontrol.wait``, so under PDL a launch never observes the counters
+  of an earlier launch that has not completed.
+* Forward progress: every spin-wait waits only on counters released by
+  programs with a *lower* program id. Absorption job ``j`` (head
+  ``j // TPH``, ``BC`` latent columns) is run, with ``ABSORB_IN_PRODUCER``
+  and ``PRODUCER_JOB_SHIFT == TPH``, by producer ``j + TPH`` (a producer of
+  the next head) after it has published its own tile, and the last head's
+  jobs by ``TPH`` extra consumer programs placed after all producers. So the
+  lowest-numbered unfinished program never waits, and a launch completes
+  with a single resident program slot when programs are dispatched in id
+  order, and on any dispatch order when the whole grid is resident at once
+  (the default pipelined grids, <= 234 programs of <= 96 KB shared memory,
+  fit on an idle sm_103 GPU; the experimental prefetching kernel's do not).
+  ``PRODUCER_JOB_SHIFT == 0`` is the legacy layout (producer ``p`` runs job
+  ``p`` of its own head group), which needs ``NOPE // BN`` co-resident
+  programs of one head group; it is kept only as a stress-test reference.
+* ``DELAY_NS > 0`` (tests only) delays a deterministic subset of producers
+  before they publish and of consumers before they finish, by up to
+  ``DELAY_NS`` nanoseconds of ``%globaltimer``, to stress the protocol.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _rms_rcp(
+    src_ptr, stride_src, rows, rmask, N: tl.constexpr, BLOCK_M: tl.constexpr, BK: tl.constexpr, eps
+):
+    cols = tl.arange(0, BK)
+    ss = tl.zeros([BLOCK_M], dtype=tl.float32)
+    for k0 in tl.range(0, N, BK, num_stages=N // BK):
+        x = tl.load(
+            src_ptr + rows[:, None] * stride_src + (k0 + cols)[None, :],
+            mask=rmask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        ss += tl.sum(x * x, axis=1)
+    return tl.rsqrt(ss / N + eps)
+
+
+@triton.jit
+def _rms_rcp_partials(
+    ss_ptr,
+    group: tl.constexpr,
+    n_parts,
+    N: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+    eps,
+):
+    parts = tl.arange(0, MAX_PARTS)
+    rows = tl.arange(0, BLOCK_M)
+    ss = tl.load(
+        ss_ptr + parts[:, None] * 2 * BLOCK_M + group * BLOCK_M + rows[None, :],
+        mask=(parts < n_parts)[:, None],
+        other=0.0,
+    )
+    return tl.rsqrt(tl.sum(ss, axis=0) / N + eps)
+
+
+# --------------------------------------------------------------------------
+# Bit-exact reproduction of the native FlashInfer CuTe-DSL RMSNorm (the
+# ``kv_a_layernorm`` / ``q_a_layernorm`` kernels, H in {512, 2048}): one warp
+# per row, lane t owns columns 256 * b + 8 * t + v (b < H / 256, v < 8), a
+# sequential per-lane fold of the squares in column order (blocks outer,
+# columns inner), a butterfly over the 32 lanes (xor offsets 1, 2, 4, 8, 16),
+# rstd = rsqrt(sum / H + eps), y = bf16((x * rstd) * (w + 0)). Every add is
+# explicit (reshape/split, sums over axes of size 2), so the result does not
+# depend on the compiled layout or on the rows per program. The same
+# arithmetic as ``shared_indexer._fi_rmsnorm_rows_vec`` (shared_indexer_absorbed_query), checked bit
+# for bit against the native kernel on real and wide-range rows.
+# --------------------------------------------------------------------------
+
+
+@triton.jit
+def _fi_butterfly32(s, ROWS: tl.constexpr):
+    """Sum of ``s [ROWS, 32]`` over the lanes in shuffle-xor butterfly order
+    (offsets 1, 2, 4, 8, 16): a balanced pairwise tree in lane order."""
+    s = tl.sum(tl.reshape(s, [ROWS, 16, 2]), axis=2)
+    s = tl.sum(tl.reshape(s, [ROWS, 8, 2]), axis=2)
+    s = tl.sum(tl.reshape(s, [ROWS, 4, 2]), axis=2)
+    s = tl.sum(tl.reshape(s, [ROWS, 2, 2]), axis=2)
+    return tl.sum(s, axis=1)
+
+
+@triton.jit
+def _v8_leaves(x, ROWS: tl.constexpr):
+    """The 8 columns ``v`` of ``x [ROWS, 32, 8]`` as ``[ROWS, 32]`` tensors,
+    in ``v`` order (register-level reshape/split, no data movement)."""
+    e, o = tl.split(tl.reshape(x, [ROWS, 32, 4, 2]))
+    ee, eo = tl.split(tl.reshape(e, [ROWS, 32, 2, 2]))
+    oe, oo = tl.split(tl.reshape(o, [ROWS, 32, 2, 2]))
+    v0, v4 = tl.split(ee)
+    v2, v6 = tl.split(eo)
+    v1, v5 = tl.split(oe)
+    v3, v7 = tl.split(oo)
+    return v0, v1, v2, v3, v4, v5, v6, v7
+
+
+@triton.jit
+def _fi_rmsnorm_rows_exact(
+    x_ptr, stride_x, w_ptr, y_ptr, stride_y, rows, rmask, eps, H: tl.constexpr, ROWS: tl.constexpr
+):
+    """``y[rows, :H] = bf16(rmsnorm(x[rows, :H]) * w)`` bit for bit as the
+    native FlashInfer RMSNorm (see above)."""
+    lanes = tl.arange(0, 32)
+    vv = tl.arange(0, 8)
+    offs = rows[:, None, None] * stride_x + 8 * lanes[None, :, None] + vv[None, None, :]
+    s = tl.zeros([ROWS, 32], dtype=tl.float32)
+    for b in tl.static_range(H // 256):
+        x = tl.load(x_ptr + offs + 256 * b, mask=rmask[:, None, None], other=0.0).to(tl.float32)
+        v0, v1, v2, v3, v4, v5, v6, v7 = _v8_leaves(x, ROWS)
+        s = s + v0 * v0
+        s = s + v1 * v1
+        s = s + v2 * v2
+        s = s + v3 * v3
+        s = s + v4 * v4
+        s = s + v5 * v5
+        s = s + v6 * v6
+        s = s + v7 * v7
+    rstd = tl.rsqrt(_fi_butterfly32(s, ROWS) / H + eps)
+    cols = tl.arange(0, H)
+    x = tl.load(
+        x_ptr + rows[:, None] * stride_x + cols[None, :], mask=rmask[:, None], other=0.0
+    ).to(tl.float32)
+    w = tl.load(w_ptr + cols).to(tl.float32) + 0.0
+    y = (x * rstd[:, None]) * w[None, :]
+    tl.store(
+        y_ptr + rows[:, None] * stride_y + cols[None, :], y.to(tl.bfloat16), mask=rmask[:, None]
+    )
+
+
+@triton.jit
+def _stress_delay(pid, salt: tl.constexpr, DELAY_NS: tl.constexpr):
+    """Test-only skew: spin ``(h / 10) * DELAY_NS`` ns, ``h = (37 pid + salt)
+    mod 11`` (so some programs do not wait and others wait up to DELAY_NS)."""
+    if DELAY_NS > 0:
+        wait_ns = ((pid * 37 + salt) % 11) * (DELAY_NS // 10)
+        t0 = tl.extra.cuda.globaltimer()
+        while tl.extra.cuda.globaltimer() - t0 < wait_ns:
+            pass
+
+
+@triton.jit
+def _ld_wuk(a, valid, wuk_ptr, KV_LORA: tl.constexpr, NOPE: tl.constexpr, BC: tl.constexpr):
+    """``W_UK[h, c0:c0+BC, :NOPE]`` of absorption job ``a`` as [NOPE, BC]
+    (split 128 + NOPE-128 along NOPE); nothing is read unless ``valid``."""
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    h = a // chunks_per_head
+    c0 = (a % chunks_per_head) * BC
+    cc = tl.arange(0, BC)
+    w_base = wuk_ptr + h * (KV_LORA * NOPE) + (c0 + cc)[None, :] * NOPE
+    w_lo = tl.load(w_base + tl.arange(0, 128)[:, None], mask=valid, other=0.0)
+    w_hi = tl.load(w_base + 128 + tl.arange(0, NOPE - 128)[:, None], mask=valid, other=0.0)
+    return w_lo, w_hi
+
+
+@triton.jit
+def _qb_program(
+    pid,
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    wqb_ptr,
+    q_ptr,
+    qr_ptr,
+    sync_ptr,
+    ss_ptr,
+    n_parts,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    WRITE_QR: tl.constexpr,
+    DO_ABSORB: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+    PRENORM: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    if PRENORM:
+        # kva_ptr/stride_kva address the already normalized BF16 q_a rows.
+        _prenorm_gemm(
+            pid,
+            kva_ptr,
+            stride_kva,
+            wqb_ptr,
+            q_ptr,
+            H * HEAD_DIM,
+            rows,
+            rmask,
+            Q_LORA,
+            BLOCK_M,
+            BN,
+            BK,
+        )
+    else:
+        _norm_gemm(
+            pid,
+            kva_ptr,
+            stride_kva,
+            gq_ptr,
+            wqb_ptr,
+            q_ptr,
+            qr_ptr,
+            ss_ptr,
+            n_parts,
+            rows,
+            rmask,
+            eps,
+            Q_LORA,
+            HEAD_DIM,
+            H,
+            BLOCK_M,
+            BN,
+            BK,
+            WRITE_QR,
+            SS_PARTIALS,
+            MAX_PARTS,
+        )
+    if DO_ABSORB:
+        if (pid % (HEAD_DIM // BN)) < (NOPE // BN):
+            _stress_delay(pid, 1, DELAY_NS)
+            # Publish this no-RoPE tile of head pid // (HEAD_DIM // BN): every
+            # thread's store is ordered before the release by the barrier.
+            tl.debug_barrier()
+            tl.atomic_add(sync_ptr + pid // (HEAD_DIM // BN), 1, sem="release", scope="gpu")
+
+
+@triton.jit
+def _prenorm_gemm(
+    t,
+    x_ptr,
+    stride_x,
+    w_ptr,
+    out_ptr,
+    OUT_COLS: tl.constexpr,
+    rows,
+    rmask,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    """``out[:, t*BN:(t+1)*BN] = x @ W[t*BN:(t+1)*BN]^T`` for BF16 ``x``."""
+    cols = tl.arange(0, BK)
+    n_offs = t * BN + tl.arange(0, BN)
+    acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32)
+    for k0 in tl.range(0, K, BK):
+        x = tl.load(
+            x_ptr + rows[:, None] * stride_x + (k0 + cols)[None, :], mask=rmask[:, None], other=0.0
+        )
+        w = tl.load(w_ptr + n_offs[None, :] * K + (k0 + cols)[:, None])
+        acc = tl.dot(x, w, acc)
+    tl.store(
+        out_ptr + rows[:, None] * OUT_COLS + n_offs[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def _norm_gemm(
+    pid,
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    wqb_ptr,
+    q_ptr,
+    qr_ptr,
+    ss_ptr,
+    n_parts,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    H: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    WRITE_QR: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+):
+    if SS_PARTIALS:
+        rcp = _rms_rcp_partials(ss_ptr, 0, n_parts, Q_LORA, BLOCK_M, MAX_PARTS, eps)
+    else:
+        rcp = _rms_rcp(kva_ptr, stride_kva, rows, rmask, Q_LORA, BLOCK_M, BK, eps)
+    cols = tl.arange(0, BK)
+    n_offs = pid * BN + tl.arange(0, BN)
+    acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32)
+    for k0 in tl.range(0, Q_LORA, BK):
+        x = tl.load(
+            kva_ptr + rows[:, None] * stride_kva + (k0 + cols)[None, :],
+            mask=rmask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        g = tl.load(gq_ptr + k0 + cols).to(tl.float32)
+        xn = (x * rcp[:, None] * g[None, :]).to(tl.bfloat16)
+        if WRITE_QR:
+            if pid == 0:
+                tl.store(
+                    qr_ptr + rows[:, None] * Q_LORA + (k0 + cols)[None, :], xn, mask=rmask[:, None]
+                )
+        w = tl.load(wqb_ptr + n_offs[None, :] * Q_LORA + (k0 + cols)[:, None])
+        acc = tl.dot(xn, w, acc)
+    tl.store(
+        q_ptr + rows[:, None] * (H * HEAD_DIM) + n_offs[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def _absorb_program(
+    a,
+    wuk_ptr,
+    q_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    rows,
+    rmask,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BN: tl.constexpr,
+    BC: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    h = a // chunks_per_head
+    c0 = (a % chunks_per_head) * BC
+    cc = tl.arange(0, BC)
+    d_lo = tl.arange(0, 128)
+    d_hi = tl.arange(0, NOPE - 128)
+    # W_UK[h] is k_b_proj_trans[h] = [KV_LORA, NOPE]; load the [NOPE, BC]
+    # slice for this program before waiting so it overlaps the q_b GEMM.
+    w_base = wuk_ptr + h * (KV_LORA * NOPE) + (c0 + cc)[None, :] * NOPE
+    w_lo = tl.load(w_base + d_lo[:, None])
+    w_hi = tl.load(w_base + 128 + d_hi[:, None])
+    # Weights are never written by an earlier grid, so only the counter and
+    # buffer accesses below have to wait for the programmatic dependency.
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu")
+    while ready < NOPE // BN:
+        ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu")
+    tl.debug_barrier()
+
+    q_base = q_ptr + rows[:, None] * (H * HEAD_DIM) + h * HEAD_DIM
+    q_lo = tl.load(q_base + d_lo[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg")
+    q_hi = tl.load(
+        q_base + 128 + d_hi[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    )
+    acc = tl.dot(q_lo, w_lo)
+    acc = tl.dot(q_hi, w_hi, acc)
+    tl.store(
+        fq_ptr + rows[:, None] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+    # Every program of head h has read the q tiles once this count reaches
+    # chunks_per_head; the last one resets both counters for the next launch.
+    _stress_delay(a, 3, DELAY_NS)
+    tl.debug_barrier()
+    done = tl.atomic_add(sync_ptr + H + h, 1, sem="acq_rel", scope="gpu")
+    if done == chunks_per_head - 1:
+        tl.atomic_xchg(sync_ptr + h, 0, sem="relaxed", scope="gpu")
+        tl.atomic_xchg(sync_ptr + H + h, 0, sem="relaxed", scope="gpu")
+
+
+@triton.jit
+def _kv_program(
+    kva_ptr,
+    stride_kva,
+    gkv_ptr,
+    latent_ptr,
+    ss_ptr,
+    n_parts,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+):
+    if SS_PARTIALS:
+        # Experimental own hidden GEMV only (not bit-exact: its partial sums
+        # follow the GEMV's column split).
+        kc = tl.arange(0, KV_LORA)
+        x = tl.load(
+            kva_ptr + rows[:, None] * stride_kva + Q_LORA + kc[None, :],
+            mask=rmask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        rcp = _rms_rcp_partials(ss_ptr, 1, n_parts, KV_LORA, BLOCK_M, MAX_PARTS, eps)
+        g = tl.load(gkv_ptr + kc).to(tl.float32)
+        y = (x * rcp[:, None] * g[None, :]).to(tl.bfloat16)
+        tl.store(
+            latent_ptr + rows[:, None] * (KV_LORA + ROPE) + kc[None, :], y, mask=rmask[:, None]
+        )
+    else:
+        # The native kv_a_layernorm, bit for bit.
+        _fi_rmsnorm_rows_exact(
+            kva_ptr + Q_LORA,
+            stride_kva,
+            gkv_ptr,
+            latent_ptr,
+            KV_LORA + ROPE,
+            rows,
+            rmask,
+            eps,
+            KV_LORA,
+            BLOCK_M,
+        )
+    rc = tl.arange(0, ROPE)
+    k_pe = tl.load(
+        kva_ptr + rows[:, None] * stride_kva + Q_LORA + KV_LORA + rc[None, :],
+        mask=rmask[:, None],
+        other=0.0,
+    )
+    tl.store(
+        latent_ptr + rows[:, None] * (KV_LORA + ROPE) + KV_LORA + rc[None, :],
+        k_pe,
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def qproj_kernel(
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    gkv_ptr,
+    wqb_ptr,
+    wuk_ptr,
+    q_ptr,
+    qr_ptr,
+    latent_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    ss_ptr,
+    n_parts,
+    R,
+    eps_q,
+    eps_kv,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    BC: tl.constexpr,
+    N_QB: tl.constexpr,
+    N_ABS: tl.constexpr,
+    WRITE_QR: tl.constexpr,
+    DO_ABSORB: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    PRENORM: tl.constexpr = False,
+    qn_ptr=None,
+    stride_qn=0,
+    ABSORB_IN_PRODUCER: tl.constexpr = False,
+    PRODUCER_JOB_SHIFT: tl.constexpr = 0,
+    ABS_JOB0: tl.constexpr = 0,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """Programs: ``[0, N_QB)`` q_b producers (with ``ABSORB_IN_PRODUCER``,
+    producer ``p >= PRODUCER_JOB_SHIFT`` then runs absorption job
+    ``p - PRODUCER_JOB_SHIFT``); ``[N_QB, N_QB + N_ABS)`` absorption jobs
+    ``ABS_JOB0 + i``; the last program normalizes ``compressed_kv``."""
+    pid = tl.program_id(0)
+    rows = tl.arange(0, BLOCK_M)
+    rmask = rows < R
+    if pid < N_QB:
+        if ABSORB_IN_PRODUCER:
+            job = tl.maximum(pid - PRODUCER_JOB_SHIFT, 0)
+            w_lo, w_hi = _ld_wuk(job, pid >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC)
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        if PRENORM:
+            _qb_program(
+                pid,
+                qn_ptr,
+                stride_qn,
+                gq_ptr,
+                wqb_ptr,
+                q_ptr,
+                qr_ptr,
+                sync_ptr,
+                ss_ptr,
+                n_parts,
+                rows,
+                rmask,
+                eps_q,
+                Q_LORA,
+                HEAD_DIM,
+                NOPE,
+                H,
+                BLOCK_M,
+                BN,
+                BK,
+                False,
+                DO_ABSORB,
+                False,
+                MAX_PARTS,
+                True,
+                DELAY_NS,
+            )
+        else:
+            _qb_program(
+                pid,
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                wqb_ptr,
+                q_ptr,
+                qr_ptr,
+                sync_ptr,
+                ss_ptr,
+                n_parts,
+                rows,
+                rmask,
+                eps_q,
+                Q_LORA,
+                HEAD_DIM,
+                NOPE,
+                H,
+                BLOCK_M,
+                BN,
+                BK,
+                WRITE_QR,
+                DO_ABSORB,
+                SS_PARTIALS,
+                MAX_PARTS,
+                False,
+                DELAY_NS,
+            )
+        if ABSORB_IN_PRODUCER:
+            if pid >= PRODUCER_JOB_SHIFT:
+                _absorb_job(
+                    job,
+                    w_lo,
+                    w_hi,
+                    q_ptr,
+                    fq_ptr,
+                    stride_fq_row,
+                    sync_ptr,
+                    rows,
+                    rmask,
+                    KV_LORA,
+                    ROPE,
+                    HEAD_DIM,
+                    NOPE,
+                    H,
+                    BN,
+                    BC,
+                    DELAY_NS,
+                )
+    elif pid < N_QB + N_ABS:
+        _absorb_program(
+            ABS_JOB0 + pid - N_QB,
+            wuk_ptr,
+            q_ptr,
+            fq_ptr,
+            stride_fq_row,
+            sync_ptr,
+            rows,
+            rmask,
+            KV_LORA,
+            ROPE,
+            HEAD_DIM,
+            NOPE,
+            H,
+            BN,
+            BC,
+            USE_PDL,
+            DELAY_NS,
+        )
+    else:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        _kv_program(
+            kva_ptr,
+            stride_kva,
+            gkv_ptr,
+            latent_ptr,
+            ss_ptr,
+            n_parts,
+            rows,
+            rmask,
+            eps_kv,
+            Q_LORA,
+            KV_LORA,
+            ROPE,
+            BLOCK_M,
+            SS_PARTIALS,
+            MAX_PARTS,
+        )
+
+
+@triton.jit
+def _gemv_cols(
+    x_ptr,
+    stride_x,
+    w_ptr,
+    out_ptr,
+    stride_o,
+    ss_ptr,
+    R,
+    col0,
+    n_cols,
+    SS_SPLIT: tl.constexpr,
+    SS_END: tl.constexpr,
+    WRITE_SS: tl.constexpr,
+    CPP: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    rows = tl.arange(0, BLOCK_M)
+    rmask = rows < R
+    lane = tl.arange(0, BN)
+    cols = col0 + lane
+    cmask = (lane < CPP) & (cols < n_cols)
+    kk = tl.arange(0, BK)
+    acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32)
+    for k0 in tl.range(0, K, BK):
+        x = tl.load(
+            x_ptr + rows[:, None] * stride_x + (k0 + kk)[None, :], mask=rmask[:, None], other=0.0
+        )
+        w = tl.load(w_ptr + cols[None, :] * K + (k0 + kk)[:, None], mask=cmask[None, :], other=0.0)
+        acc = tl.dot(x, w, acc)
+    out = acc.to(tl.bfloat16)
+    tl.store(
+        out_ptr + rows[:, None] * stride_o + cols[None, :],
+        out,
+        mask=rmask[:, None] & cmask[None, :],
+    )
+    if WRITE_SS:
+        # Per-program partial sums of squares of the BF16 outputs over the
+        # RMSNorm groups [0, SS_SPLIT) and [SS_SPLIT, SS_END); the query
+        # kernel reduces them instead of re-reading the rows.
+        sq = out.to(tl.float32)
+        sq = sq * sq
+        in_a = cmask & (cols < SS_SPLIT)
+        in_b = cmask & (cols >= SS_SPLIT) & (cols < SS_END)
+        ss_a = tl.sum(tl.where(in_a[None, :], sq, 0.0), axis=1)
+        ss_b = tl.sum(tl.where(in_b[None, :], sq, 0.0), axis=1)
+        pid = tl.program_id(0)
+        tl.store(ss_ptr + pid * 2 * BLOCK_M + rows, ss_a)
+        tl.store(ss_ptr + pid * 2 * BLOCK_M + BLOCK_M + rows, ss_b)
+
+
+@triton.jit
+def hidden_gemv_kernel(
+    x_ptr,
+    stride_x,
+    w0_ptr,
+    w1_ptr,
+    out0_ptr,
+    stride_o0,
+    out1_ptr,
+    stride_o1,
+    ss_ptr,
+    R,
+    N0,
+    N1,
+    P0: tl.constexpr,
+    CPP0: tl.constexpr,
+    CPP1: tl.constexpr,
+    SS_SPLIT: tl.constexpr,
+    SS_END: tl.constexpr,
+    WRITE_SS: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    USE_PDL: tl.constexpr,
+):
+    """``out0 = x @ W0^T`` and ``out1 = x @ W1^T`` for ``R <= BLOCK_M`` rows.
+
+    ``W0 [N0, K]`` (the native fused ``kv_a_proj_with_mqa`` weight) and the
+    optional ``W1 [N1, K]`` (the indexer's fused ``wk``/``weights_proj``
+    weight) share the input. Programs ``[0, P0)`` own ``CPP0`` consecutive
+    columns of ``W0``, the rest ``CPP1`` columns of ``W1`` (``BN`` is the
+    next power of two), each over the full ``K``, so the work is split
+    evenly over about one wave. FP32 accumulation, BF16 outputs. With
+    ``WRITE_SS``, ``W0`` programs also emit RMSNorm partial sums (see
+    ``_gemv_cols``) into ``ss_ptr [P0, 2, BLOCK_M]``.
+    """
+    pid = tl.program_id(0)
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+    if pid < P0:
+        _gemv_cols(
+            x_ptr,
+            stride_x,
+            w0_ptr,
+            out0_ptr,
+            stride_o0,
+            ss_ptr,
+            R,
+            pid * CPP0,
+            N0,
+            SS_SPLIT,
+            SS_END,
+            WRITE_SS,
+            CPP0,
+            K,
+            BLOCK_M,
+            BN,
+            BK,
+        )
+    else:
+        _gemv_cols(
+            x_ptr,
+            stride_x,
+            w1_ptr,
+            out1_ptr,
+            stride_o1,
+            ss_ptr,
+            R,
+            (pid - P0) * CPP1,
+            N1,
+            0,
+            0,
+            False,
+            CPP1,
+            K,
+            BLOCK_M,
+            BN,
+            BK,
+        )
+
+
+# --------------------------------------------------------------------------
+# Full-indexer layers: the same query kernel plus the token-wise indexer
+# projection (the index_selection contract, token-wise part), reproducing the native
+# ``Indexer.pre_indexer_proj`` chain for BF16 projections and FP8 keys bit
+# for bit:
+#   q = wq_b(qr); rope(q[:, :, :64]); fused_cat_fp8(q)          (per head)
+#   k = k_norm(wk(hidden)); rope(k[:, :64]); fused_cat_fp8(k)
+#   weights = weights_proj(hidden) * q_scale * softmax_scale * n_heads^-0.5
+# The fused wk/weights_proj GEMM and ``k_norm`` (LayerNorm under
+# torch.compile) run natively as bridges; the kernel consumes their outputs.
+# --------------------------------------------------------------------------
+
+
+@triton.jit
+def _iq_program(
+    t,
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    wiq_ptr,
+    iq_ptr,
+    isync_ptr,
+    ss_ptr,
+    n_parts,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+    PRENORM: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    if PRENORM:
+        _prenorm_gemm(
+            t, kva_ptr, stride_kva, wiq_ptr, iq_ptr, IH * IHD, rows, rmask, Q_LORA, BLOCK_M, BN, BK
+        )
+    else:
+        _iq_norm_gemm(
+            t,
+            kva_ptr,
+            stride_kva,
+            gq_ptr,
+            wiq_ptr,
+            iq_ptr,
+            ss_ptr,
+            n_parts,
+            rows,
+            rmask,
+            eps,
+            Q_LORA,
+            IHD,
+            IH,
+            BLOCK_M,
+            BN,
+            BK,
+            SS_PARTIALS,
+            MAX_PARTS,
+        )
+    _stress_delay(t, 5, DELAY_NS)
+    tl.debug_barrier()
+    tl.atomic_add(isync_ptr + t // (IHD // BN), 1, sem="release", scope="gpu")
+
+
+@triton.jit
+def _iq_norm_gemm(
+    t,
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    wiq_ptr,
+    iq_ptr,
+    ss_ptr,
+    n_parts,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+):
+    if SS_PARTIALS:
+        rcp = _rms_rcp_partials(ss_ptr, 0, n_parts, Q_LORA, BLOCK_M, MAX_PARTS, eps)
+    else:
+        rcp = _rms_rcp(kva_ptr, stride_kva, rows, rmask, Q_LORA, BLOCK_M, BK, eps)
+    cols = tl.arange(0, BK)
+    n_offs = t * BN + tl.arange(0, BN)
+    acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32)
+    for k0 in tl.range(0, Q_LORA, BK):
+        x = tl.load(
+            kva_ptr + rows[:, None] * stride_kva + (k0 + cols)[None, :],
+            mask=rmask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        g = tl.load(gq_ptr + k0 + cols).to(tl.float32)
+        xn = (x * rcp[:, None] * g[None, :]).to(tl.bfloat16)
+        w = tl.load(wiq_ptr + n_offs[None, :] * Q_LORA + (k0 + cols)[:, None])
+        acc = tl.dot(xn, w, acc)
+    tl.store(
+        iq_ptr + rows[:, None] * (IH * IHD) + n_offs[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def _fp8_ue8m0_scale(amax):
+    """``2^ceil(log2(max(amax, 1e-12) / 448))`` and its exact inverse, as in
+    the native ``fusedCatFp8Kernel<UseUe8m0=true>``."""
+    ratio = tl.maximum(amax, 1.0e-12) * (1.0 / 448.0)
+    bits = ratio.to(tl.uint32, bitcast=True)
+    exp_bits = bits & 0x7F800000
+    exp_bits = tl.where((bits & 0x007FFFFF) != 0, exp_bits + 0x00800000, exp_bits)
+    scale = exp_bits.to(tl.float32, bitcast=True)
+    inv_scale = ((254 << 23) - exp_bits).to(tl.float32, bitcast=True)
+    return scale, inv_scale
+
+
+@triton.jit
+def _mul_rn(a, b):
+    """FP32 multiply rounded once (never contracted into an FMA)."""
+    return tl.inline_asm_elementwise(
+        "mul.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1
+    )
+
+
+@triton.jit
+def _fma_rn(a, b, c):
+    """``a * b + c`` with a single rounding."""
+    return tl.inline_asm_elementwise(
+        "fma.rn.f32 $0, $1, $2, $3;", "=r,r,r,r", [a, b, c], dtype=tl.float32, is_pure=True, pack=1
+    )
+
+
+@triton.jit
+def _rope_pairs(x_even, x_odd, cos, sin):
+    """Interleaved (GPT-J) RoPE on BF16 pairs, FP32 math, BF16 result, with
+    the exact arithmetic of the native indexer RoPE (FlashInfer
+    ``vec_apply_llama_rope_cos_sin_interleave_reuse_half``: ``x[i] * cos +
+    (+-x[i ^ 1]) * sin`` compiled as ``fma(x[i], cos, rn((+-x[i ^ 1]) *
+    sin))``), pinned with explicit ``fma.rn`` / ``mul.rn``. Checked bit for
+    bit against the native kernel on real and synthetic rows."""
+    o_even = _fma_rn(x_even, cos, -_mul_rn(x_odd, sin))
+    o_odd = _fma_rn(x_odd, cos, _mul_rn(x_even, sin))
+    return (o_even.to(tl.bfloat16).to(tl.float32), o_odd.to(tl.bfloat16).to(tl.float32))
+
+
+@triton.jit
+def _qpost_program(
+    i,
+    iq_ptr,
+    isync_ptr,
+    kw_ptr,
+    stride_kw,
+    pos_ptr,
+    cos_sin_ptr,
+    qfp8_ptr,
+    qscale_ptr,
+    wout_ptr,
+    w_scale,
+    rows,
+    rmask,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    BN: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    half: tl.constexpr = ROPE_I // 2
+    nope: tl.constexpr = IHD - ROPE_I
+    j = tl.arange(0, half)
+    c = tl.arange(0, nope)
+    pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64)
+    cs_row = cos_sin_ptr + pos[:, None] * ROPE_I
+    cos = tl.load(cs_row + j[None, :], mask=rmask[:, None], other=0.0)
+    sin = tl.load(cs_row + half + j[None, :], mask=rmask[:, None], other=0.0)
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+    w = tl.load(kw_ptr + rows * stride_kw + IHD + i, mask=rmask, other=0.0).to(tl.float32)
+
+    ready = tl.atomic_add(isync_ptr + i, 0, sem="acquire", scope="gpu")
+    while ready < IHD // BN:
+        ready = tl.atomic_add(isync_ptr + i, 0, sem="acquire", scope="gpu")
+    tl.debug_barrier()
+
+    base = iq_ptr + rows[:, None] * (IH * IHD) + i * IHD
+    xe = tl.load(base + 2 * j[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg").to(
+        tl.float32
+    )
+    xo = tl.load(
+        base + 2 * j[None, :] + 1, mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    ).to(tl.float32)
+    xn = tl.load(
+        base + ROPE_I + c[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    ).to(tl.float32)
+    oe, oo = _rope_pairs(xe, xo, cos, sin)
+    amax = tl.maximum(
+        tl.maximum(tl.max(tl.abs(oe), axis=1), tl.max(tl.abs(oo), axis=1)),
+        tl.max(tl.abs(xn), axis=1),
+    )
+    scale, inv = _fp8_ue8m0_scale(amax)
+    out = qfp8_ptr + rows[:, None] * (IH * IHD) + i * IHD
+    tl.store(out + 2 * j[None, :], (oe * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + 2 * j[None, :] + 1, (oo * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + ROPE_I + c[None, :], (xn * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(qscale_ptr + rows * IH + i, scale, mask=rmask)
+    tl.store(wout_ptr + rows * IH + i, w * scale * w_scale, mask=rmask)
+    # Single consumer per head: reset for the next launch.
+    _stress_delay(i, 7, DELAY_NS)
+    tl.atomic_xchg(isync_ptr + i, 0, sem="relaxed", scope="gpu")
+
+
+@triton.jit
+def _kpost_program(
+    kln_ptr,
+    stride_kln,
+    pos_ptr,
+    cos_sin_ptr,
+    kfp8_ptr,
+    kscale_ptr,
+    rows,
+    rmask,
+    IHD: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    USE_PDL: tl.constexpr,
+):
+    """Indexer key tail on the native ``k_norm`` output ``kln [R, IHD]``
+    (the LayerNorm itself stays native: its inductor reduction order
+    depends on the compiled configuration): native RoPE arithmetic
+    (``_rope_pairs``) and the native ``fusedCatFp8Kernel<UseUe8m0>``
+    quantization."""
+    half: tl.constexpr = ROPE_I // 2
+    nope: tl.constexpr = IHD - ROPE_I
+    j = tl.arange(0, half)
+    c = tl.arange(0, nope)
+    pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64)
+    cs_row = cos_sin_ptr + pos[:, None] * ROPE_I
+    cos = tl.load(cs_row + j[None, :], mask=rmask[:, None], other=0.0)
+    sin = tl.load(cs_row + half + j[None, :], mask=rmask[:, None], other=0.0)
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+    base = kln_ptr + rows[:, None] * stride_kln
+    ye = tl.load(base + 2 * j[None, :], mask=rmask[:, None], other=0.0).to(tl.float32)
+    yo = tl.load(base + 2 * j[None, :] + 1, mask=rmask[:, None], other=0.0).to(tl.float32)
+    yn = tl.load(base + ROPE_I + c[None, :], mask=rmask[:, None], other=0.0).to(tl.float32)
+    oe, oo = _rope_pairs(ye, yo, cos, sin)
+    amax = tl.maximum(
+        tl.maximum(tl.max(tl.abs(oe), axis=1), tl.max(tl.abs(oo), axis=1)),
+        tl.max(tl.abs(yn), axis=1),
+    )
+    scale, inv = _fp8_ue8m0_scale(amax)
+    out = kfp8_ptr + rows[:, None] * IHD
+    tl.store(out + 2 * j[None, :], (oe * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + 2 * j[None, :] + 1, (oo * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + ROPE_I + c[None, :], (yn * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(kscale_ptr + rows, scale, mask=rmask)
+
+
+@triton.jit
+def qproj_indexer_kernel(
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    gkv_ptr,
+    wqb_ptr,
+    wuk_ptr,
+    q_ptr,
+    latent_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    ss_ptr,
+    n_parts,
+    R,
+    eps_q,
+    eps_kv,
+    wiq_ptr,
+    iq_ptr,
+    kw_ptr,
+    stride_kw,
+    kln_ptr,
+    stride_kln,
+    pos_ptr,
+    cos_sin_ptr,
+    qfp8_ptr,
+    qscale_ptr,
+    kfp8_ptr,
+    kscale_ptr,
+    wout_ptr,
+    w_scale,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    BC: tl.constexpr,
+    N_QB: tl.constexpr,
+    N_IQ: tl.constexpr,
+    N_ABS: tl.constexpr,
+    DO_ABSORB: tl.constexpr,
+    SS_PARTIALS: tl.constexpr,
+    MAX_PARTS: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    PRENORM: tl.constexpr = False,
+    qn_ptr=None,
+    stride_qn=0,
+    ABSORB_IN_PRODUCER: tl.constexpr = False,
+    PRODUCER_JOB_SHIFT: tl.constexpr = 0,
+    ABS_JOB0: tl.constexpr = 0,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """``qproj_kernel`` plus the indexer's token-wise projection.
+
+    Program roles, in program-id order (every consumer after the producers
+    it waits on): ``q_b`` tiles (with ``ABSORB_IN_PRODUCER``, producer
+    ``p >= PRODUCER_JOB_SHIFT`` then runs absorption job
+    ``p - PRODUCER_JOB_SHIFT``), indexer ``wq_b`` tiles (publish per indexer
+    head), absorption jobs ``ABS_JOB0 + i``, indexer query post-processing
+    (RoPE, FP8, weights; one program per indexer head), ``compressed_kv``
+    norm, indexer key post-processing. ``sync_ptr`` holds ``2 * H``
+    absorption counters followed by ``IH`` indexer-head counters.
+    """
+    pid = tl.program_id(0)
+    rows = tl.arange(0, BLOCK_M)
+    rmask = rows < R
+    if pid < N_QB:
+        if ABSORB_IN_PRODUCER:
+            job = tl.maximum(pid - PRODUCER_JOB_SHIFT, 0)
+            w_lo, w_hi = _ld_wuk(job, pid >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC)
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        if PRENORM:
+            _qb_program(
+                pid,
+                qn_ptr,
+                stride_qn,
+                gq_ptr,
+                wqb_ptr,
+                q_ptr,
+                q_ptr,
+                sync_ptr,
+                ss_ptr,
+                n_parts,
+                rows,
+                rmask,
+                eps_q,
+                Q_LORA,
+                HEAD_DIM,
+                NOPE,
+                H,
+                BLOCK_M,
+                BN,
+                BK,
+                False,
+                DO_ABSORB,
+                False,
+                MAX_PARTS,
+                True,
+                DELAY_NS,
+            )
+        else:
+            _qb_program(
+                pid,
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                wqb_ptr,
+                q_ptr,
+                q_ptr,
+                sync_ptr,
+                ss_ptr,
+                n_parts,
+                rows,
+                rmask,
+                eps_q,
+                Q_LORA,
+                HEAD_DIM,
+                NOPE,
+                H,
+                BLOCK_M,
+                BN,
+                BK,
+                False,
+                DO_ABSORB,
+                SS_PARTIALS,
+                MAX_PARTS,
+                False,
+                DELAY_NS,
+            )
+        if ABSORB_IN_PRODUCER:
+            if pid >= PRODUCER_JOB_SHIFT:
+                _absorb_job(
+                    job,
+                    w_lo,
+                    w_hi,
+                    q_ptr,
+                    fq_ptr,
+                    stride_fq_row,
+                    sync_ptr,
+                    rows,
+                    rmask,
+                    KV_LORA,
+                    ROPE,
+                    HEAD_DIM,
+                    NOPE,
+                    H,
+                    BN,
+                    BC,
+                    DELAY_NS,
+                )
+    elif pid < N_QB + N_IQ:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        if PRENORM:
+            _iq_program(
+                pid - N_QB,
+                qn_ptr,
+                stride_qn,
+                gq_ptr,
+                wiq_ptr,
+                iq_ptr,
+                sync_ptr + 2 * H,
+                ss_ptr,
+                n_parts,
+                rows,
+                rmask,
+                eps_q,
+                Q_LORA,
+                IHD,
+                IH,
+                BLOCK_M,
+                BN,
+                BK,
+                False,
+                MAX_PARTS,
+                True,
+                DELAY_NS,
+            )
+        else:
+            _iq_program(
+                pid - N_QB,
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                wiq_ptr,
+                iq_ptr,
+                sync_ptr + 2 * H,
+                ss_ptr,
+                n_parts,
+                rows,
+                rmask,
+                eps_q,
+                Q_LORA,
+                IHD,
+                IH,
+                BLOCK_M,
+                BN,
+                BK,
+                SS_PARTIALS,
+                MAX_PARTS,
+                False,
+                DELAY_NS,
+            )
+    elif pid < N_QB + N_IQ + N_ABS:
+        _absorb_program(
+            ABS_JOB0 + pid - N_QB - N_IQ,
+            wuk_ptr,
+            q_ptr,
+            fq_ptr,
+            stride_fq_row,
+            sync_ptr,
+            rows,
+            rmask,
+            KV_LORA,
+            ROPE,
+            HEAD_DIM,
+            NOPE,
+            H,
+            BN,
+            BC,
+            USE_PDL,
+            DELAY_NS,
+        )
+    elif pid < N_QB + N_IQ + N_ABS + IH:
+        _qpost_program(
+            pid - N_QB - N_IQ - N_ABS,
+            iq_ptr,
+            sync_ptr + 2 * H,
+            kw_ptr,
+            stride_kw,
+            pos_ptr,
+            cos_sin_ptr,
+            qfp8_ptr,
+            qscale_ptr,
+            wout_ptr,
+            w_scale,
+            rows,
+            rmask,
+            IHD,
+            IH,
+            ROPE_I,
+            BN,
+            USE_PDL,
+            DELAY_NS,
+        )
+    elif pid == N_QB + N_IQ + N_ABS + IH:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        _kv_program(
+            kva_ptr,
+            stride_kva,
+            gkv_ptr,
+            latent_ptr,
+            ss_ptr,
+            n_parts,
+            rows,
+            rmask,
+            eps_kv,
+            Q_LORA,
+            KV_LORA,
+            ROPE,
+            BLOCK_M,
+            SS_PARTIALS,
+            MAX_PARTS,
+        )
+    else:
+        _kpost_program(
+            kln_ptr,
+            stride_kln,
+            pos_ptr,
+            cos_sin_ptr,
+            kfp8_ptr,
+            kscale_ptr,
+            rows,
+            rmask,
+            IHD,
+            ROPE_I,
+            USE_PDL,
+        )
+
+
+# --------------------------------------------------------------------------
+# Prefetching variant (shared and full indexer layers).
+#
+# The cp.async-pipelined producers above are latency bound: each CTA keeps
+# only (num_stages - 1) K-tiles in flight and starts its weight stream only
+# after the RMSNorm pre-pass. Here every producer issues its whole [K, BN]
+# weight tile (and its absorption W_UK slice) as independent loads at
+# program start -- before ``griddepcontrol.wait``, so under programmatic
+# dependent launch the weight stream overlaps the preceding native kv_a
+# GEMM -- then waits, loads the eight q_a row tiles at once, normalizes and
+# multiplies.
+# Each producer program of head h + 1 also performs one absorption job of
+# head h after publishing its own tile (produce-then-consume), and TPH extra
+# programs after all producers absorb the last head, so every wait targets
+# lower program ids (see the module docstring).
+# --------------------------------------------------------------------------
+
+
+@triton.jit
+def _ld_w(w_ptr, n_offs, j: tl.constexpr, K: tl.constexpr, BK: tl.constexpr):
+    k = j * BK + tl.arange(0, BK)
+    return tl.load(w_ptr + n_offs[None, :] * K + k[:, None])
+
+
+@triton.jit
+def _ld_x(kva_ptr, stride_kva, rows, rmask, j: tl.constexpr, BK: tl.constexpr):
+    k = j * BK + tl.arange(0, BK)
+    return tl.load(
+        kva_ptr + rows[:, None] * stride_kva + k[None, :], mask=rmask[:, None], other=0.0
+    )
+
+
+@triton.jit
+def _sq(x):
+    xf = x.to(tl.float32)
+    return tl.sum(xf * xf, axis=1)
+
+
+@triton.jit
+def _nx(x, rcp, gq_ptr, j: tl.constexpr, BK: tl.constexpr):
+    g = tl.load(gq_ptr + j * BK + tl.arange(0, BK)).to(tl.float32)
+    return (x.to(tl.float32) * rcp[:, None] * g[None, :]).to(tl.bfloat16)
+
+
+@triton.jit
+def _prefetch_tile_gemm(
+    t,
+    w_ptr,
+    out_ptr,
+    OUT_COLS: tl.constexpr,
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    qr_ptr,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    WRITE_QR: tl.constexpr,
+    USE_PDL: tl.constexpr,
+):
+    """One [R, BN] tile of ``rmsnorm(q_a) @ W[t*BN:(t+1)*BN]^T`` with the
+    whole weight tile requested before the grid-dependency wait."""
+    BK: tl.constexpr = Q_LORA // 8
+    n_offs = t * BN + tl.arange(0, BN)
+    w0 = _ld_w(w_ptr, n_offs, 0, Q_LORA, BK)
+    w1 = _ld_w(w_ptr, n_offs, 1, Q_LORA, BK)
+    w2 = _ld_w(w_ptr, n_offs, 2, Q_LORA, BK)
+    w3 = _ld_w(w_ptr, n_offs, 3, Q_LORA, BK)
+    w4 = _ld_w(w_ptr, n_offs, 4, Q_LORA, BK)
+    w5 = _ld_w(w_ptr, n_offs, 5, Q_LORA, BK)
+    w6 = _ld_w(w_ptr, n_offs, 6, Q_LORA, BK)
+    w7 = _ld_w(w_ptr, n_offs, 7, Q_LORA, BK)
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+    x0 = _ld_x(kva_ptr, stride_kva, rows, rmask, 0, BK)
+    x1 = _ld_x(kva_ptr, stride_kva, rows, rmask, 1, BK)
+    x2 = _ld_x(kva_ptr, stride_kva, rows, rmask, 2, BK)
+    x3 = _ld_x(kva_ptr, stride_kva, rows, rmask, 3, BK)
+    x4 = _ld_x(kva_ptr, stride_kva, rows, rmask, 4, BK)
+    x5 = _ld_x(kva_ptr, stride_kva, rows, rmask, 5, BK)
+    x6 = _ld_x(kva_ptr, stride_kva, rows, rmask, 6, BK)
+    x7 = _ld_x(kva_ptr, stride_kva, rows, rmask, 7, BK)
+    ss = _sq(x0) + _sq(x1) + _sq(x2) + _sq(x3) + _sq(x4) + _sq(x5) + _sq(x6) + _sq(x7)
+    rcp = tl.rsqrt(ss / Q_LORA + eps)
+    xn0 = _nx(x0, rcp, gq_ptr, 0, BK)
+    xn1 = _nx(x1, rcp, gq_ptr, 1, BK)
+    xn2 = _nx(x2, rcp, gq_ptr, 2, BK)
+    xn3 = _nx(x3, rcp, gq_ptr, 3, BK)
+    xn4 = _nx(x4, rcp, gq_ptr, 4, BK)
+    xn5 = _nx(x5, rcp, gq_ptr, 5, BK)
+    xn6 = _nx(x6, rcp, gq_ptr, 6, BK)
+    xn7 = _nx(x7, rcp, gq_ptr, 7, BK)
+    if WRITE_QR:
+        if t == 0:
+            kk = tl.arange(0, BK)
+            qr_row = qr_ptr + rows[:, None] * Q_LORA + kk[None, :]
+            tl.store(qr_row + 0 * BK, xn0, mask=rmask[:, None])
+            tl.store(qr_row + 1 * BK, xn1, mask=rmask[:, None])
+            tl.store(qr_row + 2 * BK, xn2, mask=rmask[:, None])
+            tl.store(qr_row + 3 * BK, xn3, mask=rmask[:, None])
+            tl.store(qr_row + 4 * BK, xn4, mask=rmask[:, None])
+            tl.store(qr_row + 5 * BK, xn5, mask=rmask[:, None])
+            tl.store(qr_row + 6 * BK, xn6, mask=rmask[:, None])
+            tl.store(qr_row + 7 * BK, xn7, mask=rmask[:, None])
+    acc = tl.dot(xn0, w0)
+    acc = tl.dot(xn1, w1, acc)
+    acc = tl.dot(xn2, w2, acc)
+    acc = tl.dot(xn3, w3, acc)
+    acc = tl.dot(xn4, w4, acc)
+    acc = tl.dot(xn5, w5, acc)
+    acc = tl.dot(xn6, w6, acc)
+    acc = tl.dot(xn7, w7, acc)
+    tl.store(
+        out_ptr + rows[:, None] * OUT_COLS + n_offs[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def _absorb_job(
+    a,
+    w_lo,
+    w_hi,
+    q_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    rows,
+    rmask,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BN: tl.constexpr,
+    BC: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    h = a // chunks_per_head
+    c0 = (a % chunks_per_head) * BC
+    cc = tl.arange(0, BC)
+    d_lo = tl.arange(0, 128)
+    d_hi = tl.arange(0, NOPE - 128)
+    ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu")
+    while ready < NOPE // BN:
+        ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu")
+    tl.debug_barrier()
+    q_base = q_ptr + rows[:, None] * (H * HEAD_DIM) + h * HEAD_DIM
+    q_lo = tl.load(q_base + d_lo[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg")
+    q_hi = tl.load(
+        q_base + 128 + d_hi[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    )
+    acc = tl.dot(q_lo, w_lo)
+    acc = tl.dot(q_hi, w_hi, acc)
+    tl.store(
+        fq_ptr + rows[:, None] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+    _stress_delay(a, 3, DELAY_NS)
+    tl.debug_barrier()
+    done = tl.atomic_add(sync_ptr + H + h, 1, sem="acq_rel", scope="gpu")
+    if done == chunks_per_head - 1:
+        tl.atomic_xchg(sync_ptr + h, 0, sem="relaxed", scope="gpu")
+        tl.atomic_xchg(sync_ptr + H + h, 0, sem="relaxed", scope="gpu")
+
+
+@triton.jit
+def qproj_prefetch_kernel(
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    gkv_ptr,
+    wqb_ptr,
+    wuk_ptr,
+    q_ptr,
+    qr_ptr,
+    latent_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    R,
+    eps_q,
+    eps_kv,
+    wiq_ptr,
+    iq_ptr,
+    kw_ptr,
+    stride_kw,
+    kln_ptr,
+    stride_kln,
+    pos_ptr,
+    cos_sin_ptr,
+    qfp8_ptr,
+    qscale_ptr,
+    kfp8_ptr,
+    kscale_ptr,
+    wout_ptr,
+    w_scale,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BC: tl.constexpr,
+    N_QB: tl.constexpr,
+    N_IQ: tl.constexpr,
+    WRITE_QR: tl.constexpr,
+    DO_ABSORB: tl.constexpr,
+    FUSE_INDEXER: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    N_ABS: tl.constexpr = 0,
+    PRODUCER_JOB_SHIFT: tl.constexpr = 0,
+    ABS_JOB0: tl.constexpr = 0,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """Prefetching query kernel. Program roles in program-id order:
+
+    * ``[0, N_QB)``: ``q_b`` tile ``pid`` (prefetched), then (``DO_ABSORB``,
+      ``pid >= PRODUCER_JOB_SHIFT``) absorption job ``pid -
+      PRODUCER_JOB_SHIFT`` (``N_QB == H * KV_LORA // BC``; its ``W_UK`` slice
+      is also prefetched before the wait);
+    * ``[N_QB, N_QB + N_IQ)`` (``FUSE_INDEXER``): indexer ``wq_b`` tiles;
+    * ``N_ABS`` absorption jobs ``ABS_JOB0 + i``;
+    * ``IH`` indexer-query post-processing programs (``FUSE_INDEXER``);
+    * the ``compressed_kv`` norm / latent program;
+    * the indexer-key program (``FUSE_INDEXER``).
+    """
+    pid = tl.program_id(0)
+    rows = tl.arange(0, BLOCK_M)
+    rmask = rows < R
+    if pid < N_QB:
+        if DO_ABSORB:
+            job = tl.maximum(pid - PRODUCER_JOB_SHIFT, 0)
+            w_lo, w_hi = _ld_wuk(job, pid >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC)
+        _prefetch_tile_gemm(
+            pid,
+            wqb_ptr,
+            q_ptr,
+            H * HEAD_DIM,
+            kva_ptr,
+            stride_kva,
+            gq_ptr,
+            qr_ptr,
+            rows,
+            rmask,
+            eps_q,
+            Q_LORA,
+            BLOCK_M,
+            BN,
+            WRITE_QR,
+            USE_PDL,
+        )
+        if DO_ABSORB:
+            if (pid % (HEAD_DIM // BN)) < (NOPE // BN):
+                _stress_delay(pid, 1, DELAY_NS)
+                tl.debug_barrier()
+                tl.atomic_add(sync_ptr + pid // (HEAD_DIM // BN), 1, sem="release", scope="gpu")
+            if pid >= PRODUCER_JOB_SHIFT:
+                _absorb_job(
+                    job,
+                    w_lo,
+                    w_hi,
+                    q_ptr,
+                    fq_ptr,
+                    stride_fq_row,
+                    sync_ptr,
+                    rows,
+                    rmask,
+                    KV_LORA,
+                    ROPE,
+                    HEAD_DIM,
+                    NOPE,
+                    H,
+                    BN,
+                    BC,
+                    DELAY_NS,
+                )
+    elif pid < N_QB + N_IQ:
+        t = pid - N_QB
+        _prefetch_tile_gemm(
+            t,
+            wiq_ptr,
+            iq_ptr,
+            IH * IHD,
+            kva_ptr,
+            stride_kva,
+            gq_ptr,
+            qr_ptr,
+            rows,
+            rmask,
+            eps_q,
+            Q_LORA,
+            BLOCK_M,
+            BN,
+            False,
+            USE_PDL,
+        )
+        _stress_delay(t, 5, DELAY_NS)
+        tl.debug_barrier()
+        tl.atomic_add(sync_ptr + 2 * H + t // (IHD // BN), 1, sem="release", scope="gpu")
+    elif pid < N_QB + N_IQ + N_ABS:
+        _absorb_program(
+            ABS_JOB0 + pid - N_QB - N_IQ,
+            wuk_ptr,
+            q_ptr,
+            fq_ptr,
+            stride_fq_row,
+            sync_ptr,
+            rows,
+            rmask,
+            KV_LORA,
+            ROPE,
+            HEAD_DIM,
+            NOPE,
+            H,
+            BN,
+            BC,
+            USE_PDL,
+            DELAY_NS,
+        )
+    elif pid < N_QB + N_IQ + N_ABS + IH:
+        _qpost_program(
+            pid - N_QB - N_IQ - N_ABS,
+            iq_ptr,
+            sync_ptr + 2 * H,
+            kw_ptr,
+            stride_kw,
+            pos_ptr,
+            cos_sin_ptr,
+            qfp8_ptr,
+            qscale_ptr,
+            wout_ptr,
+            w_scale,
+            rows,
+            rmask,
+            IHD,
+            IH,
+            ROPE_I,
+            BN,
+            USE_PDL,
+            DELAY_NS,
+        )
+    elif pid == N_QB + N_IQ + N_ABS + IH:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        _kv_program(
+            kva_ptr,
+            stride_kva,
+            gkv_ptr,
+            latent_ptr,
+            sync_ptr,
+            0,
+            rows,
+            rmask,
+            eps_kv,
+            Q_LORA,
+            KV_LORA,
+            ROPE,
+            BLOCK_M,
+            False,
+            16,
+        )
+    else:
+        _kpost_program(
+            kln_ptr,
+            stride_kln,
+            pos_ptr,
+            cos_sin_ptr,
+            kfp8_ptr,
+            kscale_ptr,
+            rows,
+            rmask,
+            IHD,
+            ROPE_I,
+            USE_PDL,
+        )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/query_proj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/query_proj.py
new file mode 100644
index 000000000000..ecf1e58de8cc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/query_proj.py
@@ -0,0 +1,886 @@
+# 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.
+"""absorbed_query: target query/projection region of GLM-5.x DSA MLA (the query_front contract).
+
+``dsa_proj`` is the package entry point. It replaces the token-wise work of
+the native ``forward_dsa_proj`` for one generation-only target layer and
+returns a :class:`AbsorbedQueryPacket`. ``mtp_dsa_proj`` (``draft_absorbed_query``) does
+the same for the one-model MTP draft layer at every draft step. The packet
+holds:
+
+* ``proj_outputs``: the native 4/9-tensor list ``[q, compressed_kv, k_pe,
+  latent_cache(, q_fp8, k_fp8, k_scale, weights, q_scale)]`` with native
+  shapes, dtypes and values (``compressed_kv``/``k_pe`` are views of
+  ``latent_cache``);
+* ``fused_q``: in the preabsorbed form, the ``[R, H, kv_lora + rope]`` buffer
+  whose ``[..., :kv_lora]`` already holds ``q_nope @ W_UK`` (BF16). The
+  caller hands it to the native ``forward_absorption_generation`` in place of
+  its own allocation and skips only the absorption BMM; native RoPE, cache
+  append, FP8 query quantization and scheduler buffers still write the rest.
+  ``None`` in the native-q form.
+
+The kv_a GEMM and, on full indexer layers, ``Indexer.pre_indexer_proj`` run
+natively around the absorbed_query kernel (adapter bridges). ``dsa_proj`` returns ``None``
+before any GPU work when the call is outside the support domain, so the
+caller runs the unchanged native path.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+
+from ...dispatch import log_capability
+from . import kernels
+
+#: How the form's capability line names it (``<label>=<form>``).
+FORM_LABEL = "TRTLLM_STAIRCASE_GLM_LL_ABSORBED_QUERY_FORM"
+# "auto": preabsorbed on full indexer layers, native-q on shared layers (the
+# default form for each layer type); the other two force one form. Both
+# forms are bit-identical to native; the form is logged once per process.
+FORMS = ("auto", "preabsorbed", "native_q")
+#: The form this target runs.
+SERVED_FORM = FORMS[0]
+
+# Support domain (GLM-5.2 DSA MLA, TP8, non-ADP).
+_Q_LORA = 2048
+_KV_LORA = 512
+_ROPE = 64
+_NOPE = 192
+_HEAD_DIM = _NOPE + _ROPE
+_LOCAL_HEADS = 8
+_MAX_ROWS = 16
+
+
+def selected_form() -> str:
+    """The form the family's targets run, ``SERVED_FORM``."""
+    return SERVED_FORM
+
+
+def check_package_selectors() -> str:
+    """Log (once per process) this package's form and return it.
+
+    Called by every entry point before its per-module decision, so it runs
+    before any GPU work.
+    """
+    form = selected_form()
+    log_capability("absorbed_query", "form", True, f"{FORM_LABEL}={form}")
+    return form
+
+
+def absorb_for(mla, form: Optional[str] = None) -> bool:
+    """Whether the call returns the preabsorbed form for this module."""
+    form = form or selected_form()
+    if form == "auto":
+        return mla.mqa.indexer is not None
+    return form == "preabsorbed"
+
+
+@dataclass
+class AbsorbedQueryPacket:
+    """Explicit query_front packet of one absorbed_query invocation."""
+
+    proj_outputs: List[torch.Tensor]
+    fused_q: Optional[torch.Tensor]
+
+
+@dataclass(frozen=True)
+class KernelConfig:
+    # True: kernels.qproj_prefetch_kernel (whole weight tiles requested
+    # before the grid-dependency wait); False: the cp.async-pipelined
+    # kernels.qproj_kernel / qproj_indexer_kernel.
+    prefetch: bool = False
+    # Normalize q_a once with the native q_a_layernorm (FlashInfer RMSNorm)
+    # and feed its BF16 output to the q_b / indexer wq_b MMA loops, instead of
+    # re-normalizing the padded q_a block inside every producer program.
+    prenorm: bool = True
+    # Preabsorbed form: q_b programs perform absorption jobs after
+    # publishing their own tile (instead of separate spinning programs).
+    absorb_in_producer: bool = True
+    # With absorb_in_producer: the producers of head h + 1 absorb head h and
+    # HEAD_DIM // BN extra programs after all producers absorb the last head,
+    # so every spin-wait targets lower program ids (forward progress with one
+    # resident program slot; see kernels.py). False is the legacy layout
+    # (producer p absorbs job p of its own head group), which needs
+    # NOPE // BN co-resident programs; kept only as a stress-test reference.
+    absorb_prev_head: bool = True
+    # Tests only: skew producers/consumers by up to this many nanoseconds
+    # (kernels._stress_delay). 0 compiles the delay out.
+    debug_delay_ns: int = 0
+    # Tests only: keep the fused indexer kernel's BF16 wq_b output
+    # (``self.debug["iq"]``) for comparison with the native GEMM.
+    debug_keep_intermediates: bool = False
+    prefetch_bn: int = 16
+    prefetch_warps: int = 8
+    # Pipelined query-projection kernel; tuned launch configuration: BN=32,
+    # 4 stages.
+    block_m: int = 16
+    bn: int = 32
+    bk: int = 256
+    bc: int = 64
+    num_warps: int = 4
+    num_stages: int = 4
+    # Hidden-state GEMV (kernels.hidden_gemv_kernel) replacing the native
+    # kv_a GEMM and, on full indexer layers, the indexer wk/weights_proj GEMM.
+    # Experimental, off: it would replace the native kv_a GEMM, and its
+    # results are outside the frozen error budget.
+    own_hidden_gemv: bool = False
+    # Full indexer layers: compute the indexer query/key/weights tail in the
+    # query kernel from the native fused wk/weights_proj output and the
+    # native k_norm output; otherwise the native Indexer.pre_indexer_proj
+    # runs on the absorbed_query qr.
+    fuse_indexer: bool = True
+    # Full indexer layers, multi-stream (CUDA-graph capture): run the native
+    # fused wk/weights_proj GEMM and k_norm (bridges) on the indexer's aux
+    # stream, concurrently with the native kv_a GEMM and q_a_layernorm, and
+    # join before the absorbed_query kernel. False (or no multi-stream): in stream order.
+    indexer_k_on_aux: bool = True
+    gemv_programs: int = 0  # 0: one program per SM
+    gemv_bk: int = 256
+    gemv_warps: int = 4
+    gemv_stages: int = 6
+
+
+# The DSA hooks of a DSA MLA module: the built-in class and the one this target adopts it as.
+_DSA_HOOKS = ("DSAHooks", "GlmDsaHooks")
+
+
+def module_eligibility(mla, draft: bool = False) -> Tuple[bool, str]:
+    """Static (per-module) support check; no GPU work.
+
+    ``draft=False`` (absorbed_query): target layers, ``layer_idx < num_hidden_layers``.
+    ``draft=True`` (draft_absorbed_query): the one-model MTP draft layer,
+    ``layer_idx >= num_hidden_layers``. Every other check is shared.
+    """
+    from tensorrt_llm._utils import get_sm_version
+
+    if get_sm_version() not in (100, 103):
+        return False, f"sm{get_sm_version()} not validated"
+    if (
+        getattr(mla, "sparse_attn_hooks", None) is None
+        or type(mla.sparse_attn_hooks).__name__ not in _DSA_HOOKS
+    ):
+        return False, "not a DSA MLA module"
+    if getattr(mla, "is_lite", True) or not getattr(mla, "fuse_qkv_a_proj", False):
+        return False, "needs the fused [q_a | kv_a] projection"
+    num_layers = getattr(mla, "num_hidden_layers", None)
+    if draft:
+        if num_layers is None or mla.layer_idx is None or mla.layer_idx < num_layers:
+            return False, "draft (MTP) layers only"
+    elif num_layers is None or mla.layer_idx is None or mla.layer_idx >= num_layers:
+        return False, "target layers only"
+    if mla.mapping.enable_attention_dp or mla.mapping.cp_size != 1:
+        return False, "attention DP / CP not supported"
+    if (
+        mla.q_lora_rank,
+        mla.kv_lora_rank,
+        mla.qk_rope_head_dim,
+        mla.qk_nope_head_dim,
+        mla.num_heads_tp,
+    ) != (_Q_LORA, _KV_LORA, _ROPE, _NOPE, _LOCAL_HEADS):
+        return False, "shape outside the GLM-5.2 TP8 domain"
+    if mla.apply_rotary_emb or mla.kv_cache_dtype == "fp8_ds_mla":
+        return False, "needs the fused-RoPE TRTLLM backend cache path"
+    if getattr(mla, "llama_4_scaling", False):
+        return False, "llama-4 query scaling not supported"
+    for name in ("kv_a_proj_with_mqa", "q_b_proj"):
+        lin = getattr(mla, name)
+        if lin.has_any_quant or lin.weight.dtype != torch.bfloat16 or (lin.bias is not None):
+            return False, f"{name} must be unquantized BF16 without bias"
+    if mla.k_b_proj_trans.dtype != torch.bfloat16 or tuple(mla.k_b_proj_trans.shape) != (
+        _LOCAL_HEADS,
+        _KV_LORA,
+        _NOPE,
+    ):
+        return False, "k_b_proj_trans must be BF16 [8, 512, 192]"
+    if mla.q_a_layernorm.use_gemma or mla.kv_a_layernorm.use_gemma:
+        return False, "Gemma-style RMSNorm not supported"
+    if getattr(mla.q_a_layernorm, "is_nvfp4", False):
+        return False, "q_a_layernorm fused with NVFP4 quantization"
+    return True, ""
+
+
+def call_eligibility(mla, hidden_states, attn_metadata) -> Tuple[bool, str]:
+    """Per-call support check; no GPU work."""
+    from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph
+    from tensorrt_llm._torch.utils import is_torch_compiling
+
+    if is_torch_compiling():
+        return False, "torch.compile graph (piecewise custom ops) stays native"
+    if is_in_breakable_cuda_graph():
+        return False, "breakable CUDA graph keeps the native custom-op break"
+    if not isinstance(hidden_states, torch.Tensor) or type(hidden_states) is not torch.Tensor:
+        return False, "quantized hidden-state input"
+    if hidden_states.dtype != torch.bfloat16 or hidden_states.dim() != 2:
+        return False, "hidden states must be 2-D BF16"
+    rows = hidden_states.shape[0]
+    if not 1 <= rows <= _MAX_ROWS:
+        return False, f"{rows} rows outside [1, {_MAX_ROWS}]"
+    if (
+        attn_metadata is None
+        or attn_metadata.num_contexts != 0
+        or (attn_metadata.num_generations == 0)
+    ):
+        return False, "generation-only batches only"
+    return True, ""
+
+
+class AbsorbedQueryProjection:
+    """query_front producer bound to one native DSA ``MLA`` module."""
+
+    def __init__(self, mla, config: KernelConfig = KernelConfig()):
+        self.mla = mla
+        self.config = config
+        device = mla.q_b_proj.weight.device
+        # Per-layer ready/done counters (see kernels.py): 2 * H absorption
+        # counters, then one per indexer head; always zero between launches.
+        indexer = mla.mqa.indexer
+        n_index_heads = indexer.n_heads if indexer is not None else 0
+        self.sync = torch.zeros(2 * _LOCAL_HEADS + n_index_heads, dtype=torch.int32, device=device)
+        from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+        self.use_pdl = get_env_enable_pdl()
+        self.num_sms = torch.cuda.get_device_properties(device).multi_processor_count
+        # Fork/join events of the indexer-key bridges (``native_bridges``).
+        self.events = [torch.cuda.Event(), torch.cuda.Event()]
+        self.debug = {}
+
+    def launch(
+        self,
+        kva: torch.Tensor,
+        need_qr: bool,
+        absorb: bool,
+        ss: Optional[torch.Tensor] = None,
+        n_parts: int = 0,
+    ):
+        """Run the query kernel on a kv_a output ``[R, 2624]``.
+
+        ``ss``/``n_parts`` are the hidden GEMV's RMSNorm partial sums; without
+        them the kernel recomputes the row norms from ``kva``.
+        """
+        mla = self.mla
+        cfg = self.config
+        rows = kva.shape[0]
+        assert kva.stride(1) == 1 and rows <= cfg.block_m
+        q = kva.new_empty(rows, _LOCAL_HEADS * _HEAD_DIM)
+        prenorm = cfg.prenorm and ss is None
+        qn = self.native_q_norm(kva) if prenorm else None
+        if prenorm:
+            qr = qn
+            need_qr = False
+        else:
+            qr = kva.new_empty(rows, _Q_LORA) if need_qr else q
+        latent = kva.new_empty(rows, _KV_LORA + _ROPE)
+        fused_q = kva.new_empty(rows, _LOCAL_HEADS, _KV_LORA + _ROPE) if absorb else None
+        n_qb = (_LOCAL_HEADS * _HEAD_DIM) // cfg.bn
+        in_producer = absorb and cfg.absorb_in_producer
+        layout = absorb_layout(n_qb, absorb, in_producer, cfg.absorb_prev_head, cfg.bc)
+        grid = (n_qb + layout["N_ABS"] + 1,)
+        kernels.qproj_kernel[grid](
+            kva,
+            kva.stride(0),
+            mla.q_a_layernorm.weight,
+            mla.kv_a_layernorm.weight,
+            mla.q_b_proj.weight,
+            mla.k_b_proj_trans,
+            q,
+            qr,
+            latent,
+            fused_q if absorb else q,
+            fused_q.stride(0) if absorb else 0,
+            self.sync,
+            ss if ss is not None else self.sync,
+            n_parts,
+            rows,
+            mla.q_a_layernorm.variance_epsilon,
+            mla.kv_a_layernorm.variance_epsilon,
+            Q_LORA=_Q_LORA,
+            KV_LORA=_KV_LORA,
+            ROPE=_ROPE,
+            HEAD_DIM=_HEAD_DIM,
+            NOPE=_NOPE,
+            H=_LOCAL_HEADS,
+            BLOCK_M=cfg.block_m,
+            BN=cfg.bn,
+            BK=cfg.bk,
+            N_QB=n_qb,
+            WRITE_QR=need_qr,
+            DO_ABSORB=absorb,
+            SS_PARTIALS=ss is not None,
+            MAX_PARTS=triton_next_pow2(max(n_parts, 1)),
+            USE_PDL=self.use_pdl,
+            PRENORM=prenorm,
+            qn_ptr=qn if prenorm else q,
+            stride_qn=qn.stride(0) if prenorm else 0,
+            ABSORB_IN_PRODUCER=in_producer,
+            DELAY_NS=cfg.debug_delay_ns,
+            **layout,
+            num_warps=cfg.num_warps,
+            num_stages=cfg.num_stages,
+            launch_pdl=self.use_pdl,
+        )
+        return q, (qr if (need_qr or prenorm) else None), latent, fused_q
+
+    def native_q_norm(self, kva: torch.Tensor) -> torch.Tensor:
+        """The native ``q_a_layernorm`` call on the q_a slice (bridge)."""
+        return self.mla._q_a_layernorm_maybe_fused(kva[:, :_Q_LORA])
+
+    def fuse_indexer_kw(self) -> bool:
+        """Whether the hidden GEMV also produces the indexer wk/weights_proj
+        output (BF16 indexer projection, FP8 indexer keys only)."""
+        indexer = self.mla.mqa.indexer
+        return (
+            indexer is not None
+            and indexer._indexer_bf16
+            and not indexer.use_fp4
+            and indexer._fused_wk_wp_weight is not None
+            and indexer._fused_wk_wp_weight.dtype == torch.bfloat16
+        )
+
+    def hidden_gemv(self, hidden_states: torch.Tensor, fuse_indexer_kw: bool):
+        """``hidden @ [W_kv_a; W_indexer_kw]^T`` in one balanced launch.
+
+        Returns ``(kva, kw, ss, n_parts)``: the BF16 kv_a output, the BF16
+        indexer wk/weights_proj output (or None) and the per-program RMSNorm
+        partial sums ``[n_parts, 2, block_m]`` of the q_a / compressed_kv
+        column groups.
+        """
+        mla = self.mla
+        cfg = self.config
+        rows = hidden_states.shape[0]
+        w0 = mla.kv_a_proj_with_mqa.weight
+        n0 = w0.shape[0]
+        kva = hidden_states.new_empty(rows, n0)
+        programs = cfg.gemv_programs or self.num_sms
+        if fuse_indexer_kw:
+            w1 = mla.mqa.indexer._fused_wk_wp_weight
+            n1 = w1.shape[0]
+            kw = hidden_states.new_empty(rows, n1)
+            p1 = max(1, round(programs * n1 / (n0 + n1)))
+            cpp1 = triton_cdiv(n1, p1)
+            p1 = triton_cdiv(n1, cpp1)
+        else:
+            w1, kw, n1, p1, cpp1 = w0, kva, 0, 0, 1
+        cpp0 = triton_cdiv(n0, programs - p1)
+        p0 = triton_cdiv(n0, cpp0)
+        ss = torch.empty(p0, 2, cfg.block_m, dtype=torch.float32, device=hidden_states.device)
+        kernels.hidden_gemv_kernel[(p0 + p1,)](
+            hidden_states,
+            hidden_states.stride(0),
+            w0,
+            w1,
+            kva,
+            kva.stride(0),
+            kw,
+            kw.stride(0),
+            ss,
+            rows,
+            n0,
+            n1,
+            P0=p0,
+            CPP0=cpp0,
+            CPP1=cpp1,
+            SS_SPLIT=_Q_LORA,
+            SS_END=_Q_LORA + _KV_LORA,
+            WRITE_SS=True,
+            K=hidden_states.shape[1],
+            BLOCK_M=cfg.block_m,
+            BN=triton_next_pow2(max(cpp0, cpp1)),
+            BK=cfg.gemv_bk,
+            USE_PDL=self.use_pdl,
+            num_warps=cfg.gemv_warps,
+            num_stages=cfg.gemv_stages,
+            launch_pdl=self.use_pdl,
+        )
+        return kva, (kw if fuse_indexer_kw else None), ss, p0
+
+    def indexer_fusable(self) -> bool:
+        """Whether the query kernel can produce the indexer tail itself."""
+        indexer = self.mla.mqa.indexer
+        if not (self.config.fuse_indexer and self.fuse_indexer_kw()):
+            return False
+        rope = indexer.rotary_emb
+        return (
+            indexer.scale_fmt == "ue8m0"
+            and indexer.head_dim == 128
+            and indexer.rope_dim == 64
+            and indexer.n_heads <= 64
+            and not rope.is_neox
+            and not rope.inverse
+            and rope.head_dim == indexer.rope_dim
+            and rope.rotary_cos_sin.dtype == torch.float32
+            and rope.rotary_cos_sin.is_contiguous()
+            and tuple(rope.rotary_cos_sin.shape[1:]) == (2, 32)
+            and indexer.wq_b.weight.dtype == torch.bfloat16
+            and not indexer.wq_b.has_any_quant
+        )
+
+    def launch_indexer(
+        self,
+        kva: torch.Tensor,
+        kw: torch.Tensor,
+        k_ln: torch.Tensor,
+        ss: Optional[torch.Tensor],
+        n_parts: int,
+        position_ids: torch.Tensor,
+        absorb: bool,
+        qn: Optional[torch.Tensor] = None,
+    ):
+        """Query kernel with the fused indexer tail (full indexer layers).
+
+        ``kw`` is the native fused wk/weights_proj output, ``k_ln`` the
+        native ``k_norm`` of its key columns and ``qn`` (optional) the native
+        ``q_a_layernorm`` output (computed here when absent).
+        """
+        mla = self.mla
+        indexer = mla.mqa.indexer
+        cfg = self.config
+        rows = kva.shape[0]
+        assert (
+            kva.stride(1) == 1 and kw.stride(1) == 1 and k_ln.stride(1) == 1 and rows <= cfg.block_m
+        )
+        ih, ihd = indexer.n_heads, indexer.head_dim
+        q = kva.new_empty(rows, _LOCAL_HEADS * _HEAD_DIM)
+        prenorm = cfg.prenorm and ss is None
+        if prenorm and qn is None:
+            qn = self.native_q_norm(kva)
+        elif not prenorm:
+            qn = q
+        latent = kva.new_empty(rows, _KV_LORA + _ROPE)
+        fused_q = kva.new_empty(rows, _LOCAL_HEADS, _KV_LORA + _ROPE) if absorb else None
+        iq = kva.new_empty(rows, ih * ihd)
+        q_fp8 = torch.empty(rows, ih, ihd, dtype=torch.float8_e4m3fn, device=kva.device)
+        q_scale = torch.empty(rows, ih, 1, dtype=torch.float32, device=kva.device)
+        k_fp8 = torch.empty(rows, ihd, dtype=torch.float8_e4m3fn, device=kva.device)
+        k_scale = torch.empty(rows, 1, dtype=torch.float32, device=kva.device)
+        weights = torch.empty(rows, ih, dtype=torch.float32, device=kva.device)
+        pos = position_ids.reshape(-1)
+        assert pos.numel() >= rows
+        n_qb = (_LOCAL_HEADS * _HEAD_DIM) // cfg.bn
+        n_iq = (ih * ihd) // cfg.bn
+        in_producer = absorb and cfg.absorb_in_producer
+        layout = absorb_layout(n_qb, absorb, in_producer, cfg.absorb_prev_head, cfg.bc)
+        grid = (n_qb + n_iq + layout["N_ABS"] + ih + 2,)
+        kernels.qproj_indexer_kernel[grid](
+            kva,
+            kva.stride(0),
+            mla.q_a_layernorm.weight,
+            mla.kv_a_layernorm.weight,
+            mla.q_b_proj.weight,
+            mla.k_b_proj_trans,
+            q,
+            latent,
+            fused_q if absorb else q,
+            fused_q.stride(0) if absorb else 0,
+            self.sync,
+            ss if ss is not None else self.sync,
+            n_parts,
+            rows,
+            mla.q_a_layernorm.variance_epsilon,
+            mla.kv_a_layernorm.variance_epsilon,
+            indexer.wq_b.weight,
+            iq,
+            kw,
+            kw.stride(0),
+            k_ln,
+            k_ln.stride(0),
+            pos,
+            indexer.rotary_emb.rotary_cos_sin,
+            q_fp8,
+            q_scale,
+            k_fp8,
+            k_scale,
+            weights,
+            indexer.weight_scale_factor,
+            Q_LORA=_Q_LORA,
+            KV_LORA=_KV_LORA,
+            ROPE=_ROPE,
+            HEAD_DIM=_HEAD_DIM,
+            NOPE=_NOPE,
+            H=_LOCAL_HEADS,
+            IHD=ihd,
+            IH=ih,
+            ROPE_I=indexer.rope_dim,
+            BLOCK_M=cfg.block_m,
+            BN=cfg.bn,
+            BK=cfg.bk,
+            N_QB=n_qb,
+            N_IQ=n_iq,
+            DO_ABSORB=absorb,
+            SS_PARTIALS=ss is not None,
+            MAX_PARTS=triton_next_pow2(max(n_parts, 1)),
+            USE_PDL=self.use_pdl,
+            PRENORM=prenorm,
+            qn_ptr=qn,
+            stride_qn=qn.stride(0) if prenorm else 0,
+            ABSORB_IN_PRODUCER=in_producer,
+            DELAY_NS=cfg.debug_delay_ns,
+            **layout,
+            num_warps=cfg.num_warps,
+            num_stages=cfg.num_stages,
+            launch_pdl=self.use_pdl,
+        )
+        if cfg.debug_keep_intermediates:
+            self.debug["iq"] = iq
+        outs = [
+            q,
+            latent[:, :_KV_LORA],
+            latent[:, _KV_LORA:],
+            latent,
+            q_fp8,
+            k_fp8,
+            k_scale,
+            weights,
+            q_scale,
+        ]
+        return AbsorbedQueryPacket(outs, fused_q)
+
+    def launch_prefetch(
+        self,
+        kva: torch.Tensor,
+        kw: Optional[torch.Tensor],
+        hidden_states: torch.Tensor,
+        position_ids: Optional[torch.Tensor],
+        absorb: bool,
+        k_ln: Optional[torch.Tensor] = None,
+    ) -> AbsorbedQueryPacket:
+        """Prefetching query kernel (both layer types; experimental).
+
+        ``kw`` is the native fused indexer wk/weights_proj output on full
+        indexer layers whose tail is fused (``k_ln``: its native ``k_norm``,
+        computed here when absent); with ``kw`` None on a full layer the
+        native indexer tail runs on the returned ``qr``.
+        """
+        mla = self.mla
+        indexer = mla.mqa.indexer
+        cfg = self.config
+        rows = kva.shape[0]
+        assert kva.stride(1) == 1 and rows <= cfg.block_m
+        bn = cfg.prefetch_bn
+        n_qb = (_LOCAL_HEADS * _HEAD_DIM) // bn
+        # The prefetching kernel always absorbs inside the producers.
+        layout = absorb_layout(
+            n_qb, absorb, True, cfg.absorb_prev_head, (_LOCAL_HEADS * _KV_LORA) // n_qb
+        )
+        fuse = kw is not None
+        need_qr = indexer is not None and not fuse
+        q = kva.new_empty(rows, _LOCAL_HEADS * _HEAD_DIM)
+        qr = kva.new_empty(rows, _Q_LORA) if need_qr else q
+        latent = kva.new_empty(rows, _KV_LORA + _ROPE)
+        fused_q = kva.new_empty(rows, _LOCAL_HEADS, _KV_LORA + _ROPE) if absorb else None
+        if fuse:
+            ih, ihd = indexer.n_heads, indexer.head_dim
+            iq = kva.new_empty(rows, ih * ihd)
+            q_fp8 = torch.empty(rows, ih, ihd, dtype=torch.float8_e4m3fn, device=kva.device)
+            q_scale = torch.empty(rows, ih, 1, dtype=torch.float32, device=kva.device)
+            k_fp8 = torch.empty(rows, ihd, dtype=torch.float8_e4m3fn, device=kva.device)
+            k_scale = torch.empty(rows, 1, dtype=torch.float32, device=kva.device)
+            weights = torch.empty(rows, ih, dtype=torch.float32, device=kva.device)
+            pos = position_ids.reshape(-1)
+            assert pos.numel() >= rows
+            n_iq = (ih * ihd) // bn
+            if k_ln is None:
+                k_ln = indexer.k_norm(kw[:, :ihd])
+            iargs = (
+                indexer.wq_b.weight,
+                iq,
+                kw,
+                kw.stride(0),
+                k_ln,
+                k_ln.stride(0),
+                pos,
+                indexer.rotary_emb.rotary_cos_sin,
+                q_fp8,
+                q_scale,
+                k_fp8,
+                k_scale,
+                weights,
+                indexer.weight_scale_factor,
+            )
+            iconst = dict(IHD=ihd, IH=ih, ROPE_I=indexer.rope_dim)
+        else:
+            ih, n_iq = 0, 0
+            dummy = q
+            iargs = (
+                dummy,
+                dummy,
+                dummy,
+                0,
+                dummy,
+                0,
+                self.sync,
+                self.sync,
+                dummy,
+                self.sync,
+                dummy,
+                self.sync,
+                self.sync,
+                0.0,
+            )
+            iconst = dict(IHD=128, IH=0, ROPE_I=64)
+        grid = (n_qb + n_iq + layout["N_ABS"] + ih + 1 + (1 if fuse else 0),)
+        kernels.qproj_prefetch_kernel[grid](
+            kva,
+            kva.stride(0),
+            mla.q_a_layernorm.weight,
+            mla.kv_a_layernorm.weight,
+            mla.q_b_proj.weight,
+            mla.k_b_proj_trans,
+            q,
+            qr,
+            latent,
+            fused_q if absorb else q,
+            fused_q.stride(0) if absorb else 0,
+            self.sync,
+            rows,
+            mla.q_a_layernorm.variance_epsilon,
+            mla.kv_a_layernorm.variance_epsilon,
+            *iargs,
+            Q_LORA=_Q_LORA,
+            KV_LORA=_KV_LORA,
+            ROPE=_ROPE,
+            HEAD_DIM=_HEAD_DIM,
+            NOPE=_NOPE,
+            H=_LOCAL_HEADS,
+            BLOCK_M=cfg.block_m,
+            BN=bn,
+            N_QB=n_qb,
+            N_IQ=n_iq,
+            WRITE_QR=need_qr,
+            DO_ABSORB=absorb,
+            FUSE_INDEXER=fuse,
+            USE_PDL=self.use_pdl,
+            DELAY_NS=cfg.debug_delay_ns,
+            **layout,
+            num_warps=cfg.prefetch_warps,
+            num_stages=1,
+            launch_pdl=self.use_pdl,
+            **iconst,
+        )
+        outs = [q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent]
+        if fuse:
+            outs += [q_fp8, k_fp8, k_scale, weights, q_scale]
+        elif indexer is not None:
+            outs.extend(indexer.pre_indexer_proj(qr, hidden_states, position_ids))
+        return AbsorbedQueryPacket(outs, fused_q)
+
+    def native_bridges(self, hidden_states: torch.Tensor):
+        """The native token-wise calls the fused full-layer kernel consumes.
+
+        Returns ``(kva, qn, kw, k_ln)``: the native kv_a GEMM output, the
+        native ``q_a_layernorm`` of its q_a columns (None without
+        ``prenorm``), the native BF16 fused wk/weights_proj GEMM output (the
+        call ``Indexer.pre_indexer_proj`` makes with
+        ``TRTLLM_DSA_INDEXER_BF16``) and the native ``k_norm`` of its key
+        columns. The last two depend only on ``hidden_states``; with
+        ``indexer_k_on_aux`` and multi-stream enabled (CUDA-graph capture)
+        they run on the indexer's aux stream concurrently with the first two
+        and are joined before returning, as the native region overlaps its
+        own aux-stream work.
+        """
+        mla = self.mla
+        indexer = mla.mqa.indexer
+        prenorm = self.config.prenorm
+
+        def main():
+            kva = mla.kv_a_proj_with_mqa(hidden_states)
+            return kva, (self.native_q_norm(kva) if prenorm else None)
+
+        def aux():
+            kw = torch.nn.functional.linear(hidden_states, indexer._fused_wk_wp_weight)
+            return kw, indexer.k_norm(kw[:, : indexer.head_dim])
+
+        if self.config.indexer_k_on_aux:
+            from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel
+
+            (kva, qn), (kw, k_ln) = maybe_execute_in_parallel(
+                main, aux, self.events[0], self.events[1], indexer.aux_stream
+            )
+        else:
+            kva, qn = main()
+            kw, k_ln = aux()
+        return kva, qn, kw, k_ln
+
+    def project(
+        self, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, absorb: bool
+    ) -> AbsorbedQueryPacket:
+        """Native bridges (or the experimental hidden GEMV), absorbed_query kernel,
+        indexer tail."""
+        mla = self.mla
+        indexer = mla.mqa.indexer
+        fusable = indexer is not None and self.indexer_fusable()
+        kw, ss, n_parts = None, None, 0
+        if self.config.own_hidden_gemv:
+            kva, kw, ss, n_parts = self.hidden_gemv(hidden_states, self.fuse_indexer_kw())
+        elif fusable:
+            kva, qn, kw, k_ln = self.native_bridges(hidden_states)
+            if self.config.prefetch:
+                return self.launch_prefetch(kva, kw, hidden_states, position_ids, absorb, k_ln)
+            return self.launch_indexer(kva, kw, k_ln, None, 0, position_ids, absorb, qn)
+        else:
+            kva = mla.kv_a_proj_with_mqa(hidden_states)
+        if self.config.prefetch and not self.config.own_hidden_gemv:
+            return self.launch_prefetch(kva, None, hidden_states, position_ids, absorb)
+        if fusable:
+            # Experimental own hidden GEMV: its kw, the native k_norm.
+            return self.launch_indexer(
+                kva,
+                kw,
+                indexer.k_norm(kw[:, : indexer.head_dim]),
+                ss,
+                n_parts,
+                position_ids,
+                absorb,
+            )
+        q, qr, latent, fused_q = self.launch(
+            kva, need_qr=indexer is not None, absorb=absorb, ss=ss, n_parts=n_parts
+        )
+        outs = [q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent]
+        if indexer is not None:
+            if kw is not None:
+                outs.extend(indexer_tail_from_projected(indexer, qr, kw, position_ids))
+            else:
+                outs.extend(indexer.pre_indexer_proj(qr, hidden_states, position_ids))
+        return AbsorbedQueryPacket(outs, fused_q)
+
+    def forward_preabsorbed(self, position_ids, hidden_states):
+        return self.project(position_ids, hidden_states, absorb=True)
+
+
+def absorb_layout(
+    n_qb: int, absorb: bool, in_producer: bool, prev_head: bool, bc_separate: int
+) -> dict:
+    """Absorption program layout of one query-kernel launch.
+
+    Returns the kernel's ``BC``, ``N_ABS`` (separate absorption programs),
+    ``ABS_JOB0`` (first job of those programs) and ``PRODUCER_JOB_SHIFT``
+    (producer ``p`` runs job ``p - shift`` when ``ABSORB_IN_PRODUCER``).
+    """
+    tiles_per_head = n_qb // _LOCAL_HEADS
+    if not absorb:
+        return dict(BC=bc_separate, N_ABS=0, ABS_JOB0=0, PRODUCER_JOB_SHIFT=0)
+    if not in_producer:
+        return dict(
+            BC=bc_separate,
+            N_ABS=_LOCAL_HEADS * (_KV_LORA // bc_separate),
+            ABS_JOB0=0,
+            PRODUCER_JOB_SHIFT=0,
+        )
+    # One job per producer: BC latent columns, KV_LORA // BC == tiles/head.
+    bc = (_LOCAL_HEADS * _KV_LORA) // n_qb
+    assert _KV_LORA // bc == tiles_per_head
+    if not prev_head:
+        return dict(BC=bc, N_ABS=0, ABS_JOB0=0, PRODUCER_JOB_SHIFT=0)
+    return dict(
+        BC=bc,
+        N_ABS=tiles_per_head,
+        ABS_JOB0=n_qb - tiles_per_head,
+        PRODUCER_JOB_SHIFT=tiles_per_head,
+    )
+
+
+def triton_cdiv(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+def triton_next_pow2(n: int) -> int:
+    return 1 << max(4, (n - 1).bit_length())
+
+
+def indexer_tail_from_projected(
+    indexer, qr: torch.Tensor, fused_kw: torch.Tensor, position_ids: Optional[torch.Tensor]
+):
+    """Native ``Indexer.pre_indexer_proj`` after its fused wk/weights_proj
+    GEMM, fed with that GEMM's output ``fused_kw [R, head_dim + n_heads]``.
+
+    Calls the native indexer methods in the native order (FP8 keys, BF16
+    projection), so only the source of ``fused_kw`` differs.
+    """
+    from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel
+
+    indexer_k, weights = fused_kw.split([indexer.head_dim, indexer.n_heads], dim=-1)
+    q_pe, q_nope, k_pe, k_nope = indexer._qk_projection_and_rope(qr, indexer_k, position_ids)
+    q, k = maybe_execute_in_parallel(
+        lambda: indexer._prep_q_or_k(q_pe, q_nope),
+        lambda: indexer._prep_q_or_k(k_pe, k_nope),
+        indexer.ln_events[0],
+        indexer.ln_events[1],
+        indexer.aux_stream,
+    )
+    q_fp8, q_scale = q
+    k_fp8, k_scale = k
+    q_fp8 = q_fp8.view(-1, indexer.n_heads, indexer.head_dim)
+    q_scale = q_scale.view(-1, indexer.n_heads, 1)
+    weights = indexer._weight_scale(weights, q_scale)
+    return q_fp8, k_fp8, k_scale, weights, q_scale
+
+
+def dsa_proj(
+    mla, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """absorbed_query replacement of the token-wise ``forward_dsa_proj`` work.
+
+    Returns ``None`` (before any GPU work) when the module or the call is
+    outside the support domain; the caller then runs the native path.
+    """
+    runner = getattr(mla, "_absorbed_query", None)
+    if runner is None:
+        check_package_selectors()
+        ok, reason = module_eligibility(mla)
+        log_capability("absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+        runner = AbsorbedQueryProjection(mla) if ok else False
+        mla._absorbed_query = runner
+    if runner is False:
+        return None
+    ok, _ = call_eligibility(mla, hidden_states, attn_metadata)
+    if not ok:
+        return None
+    return runner.project(position_ids, hidden_states, absorb=absorb_for(mla))
+
+
+def mtp_dsa_proj(
+    mla, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """draft_absorbed_query: the absorbed_query query_front producer on the one-model MTP draft layer.
+
+    The MTP layer (``layer_idx >= num_hidden_layers``, GLM-5.2: layer 78)
+    is called once per draft step of the MTP Eagle draft loop (``R = 6``
+    rows at step 0, one row per request at steps 1-4 for one request). Its
+    attention weights have the target-layer shapes and dtypes and it always
+    owns a full indexer, so the unchanged absorbed_query kernel path serves every draft
+    step. The native RoPE, cache appends (latent and indexer key), FP8 query
+    quantization, scheduler buffers and the indexer's draft-loop top-k reuse
+    (``mtp_index_share``) all stay in the native code after this call. Every
+    output of the native 9-tensor list is written at every draft step, also
+    on steps whose indexer reuses the step-0 top-k.
+
+    Returns ``None`` (before any GPU work) for target layers and outside the
+    absorbed_query support domain; the caller then runs the native path. The per-module
+    decision is cached in ``mla._draft_absorbed_query`` and logged once under
+    ``draft_absorbed_query``.
+    """
+    runner = getattr(mla, "_draft_absorbed_query", None)
+    if runner is None:
+        check_package_selectors()
+        ok, reason = module_eligibility(mla, draft=True)
+        log_capability("draft_absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+        runner = AbsorbedQueryProjection(mla) if ok else False
+        mla._draft_absorbed_query = runner
+    if runner is False:
+        return None
+    ok, _ = call_eligibility(mla, hidden_states, attn_metadata)
+    if not ok:
+        return None
+    return runner.project(position_ids, hidden_states, absorb=absorb_for(mla))
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer.py
new file mode 100644
index 000000000000..63641d50abca
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer.py
@@ -0,0 +1,1348 @@
+# 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.
+"""shared_indexer_absorbed_query: one-kernel query region of the shared-indexer layers (absorbed_query package).
+
+On a GLM-5.x DSA layer without its own indexer (the indexer reuses the
+previous full layer's top-k), the absorbed_query path runs the native ``fused_a`` kv_a
+GEMM, the native ``q_a_layernorm`` (FlashInfer RMSNorm), the absorbed_query kernel
+in its native-q form and then, inside ``MLA.forward_absorption_generation``,
+the native absorption BMM. With package ``shared_indexer_absorbed_query`` selected (on top of
+``absorbed_query``) ``shared_dsa_proj`` replaces everything after the kv_a GEMM by one kernel
+that returns the **preabsorbed** query_front packet (``q``, ``latent`` =
+``[rmsnorm(compressed_kv) | k_pe]`` and ``fused_q[..., :kv_lora]``):
+
+* R <= 8 rows (``SharedQueryKernelConfig.cute``, default): the CuTe-DSL cluster
+  kernel of ``shared_indexer_cute`` (one 8-CTA cluster per head, the
+  normalized rows recomputed per CTA, the no-RoPE ``q`` tiles exchanged
+  through distributed shared memory; see that module).
+* otherwise (R <= 16, or ``cute=False``): ``shared_query_kernel`` below.
+
+``shared_query_kernel`` (Triton):
+
+* norm programs reproduce the native ``q_a_layernorm`` (FlashInfer CuTe-DSL
+  ``RMSNormKernel``) bit for bit -- same lane-to-column map, same per-lane
+  fold order, same butterfly, ``rsqrt``, ``bf16((x * rstd) * w)`` -- write
+  the normalized rows to a scratch buffer and release them to the producers;
+  they also produce ``latent`` with the same reproduction of the native
+  ``kv_a_layernorm``;
+* ``q_b`` producer programs (the absorbed_query tiling: ``BN`` output columns over the
+  full ``K``) wait for the normalized rows, compute ``q`` with the same MMA
+  sequence as absorbed_query and publish each head's no-RoPE tiles;
+* absorption jobs (absorbed_query's shifted layout: producers of head ``h + 1`` absorb
+  head ``h``, ``HEAD_DIM // BN`` extra programs absorb the last head) write
+  ``fused_q[:, h, :kv_lora] = bf16(q_nope[:, h] @ W_UK[h]^T)``.
+
+Both kernels are launched with programmatic dependent launch right after the
+kv_a GEMM. Before ``griddepcontrol.wait`` a program only reads weights (the
+Triton kernel asks L2 to prefetch its ``q_b`` tile and ``W_UK`` slice, the
+CuTe kernel copies them into shared memory), so the weight stream overlaps
+the kv_a GEMM; every activation, output and counter access comes after it.
+With ``early_trigger`` (default; off: the implicit trigger at grid
+completion) every program issues ``griddepcontrol.launch_dependents`` right
+after its own wait, so the next kernel (its programmatic dependent) can be
+scheduled while this one runs.
+
+PDL contract with the dependent (whichever kernel follows: the native
+``mla_rope_generation`` kernel, or another package's replacement such as
+rope_kv_append's fused rope):
+
+* What the trigger guarantees. The dependent is released only once every
+  program of this grid has issued the trigger (or exited). At that point
+  every program is resident and past its own wait. So (1) everything this
+  kernel's wait covers has completed, as it has at an implicit trigger, and
+  (2) a waiting dependent program cannot keep a program of this grid, or a
+  counter it spins on, from being scheduled.
+* What the dependent must do:
+
+  - execute ``griddepcontrol.wait`` before reading any output of this kernel
+    (``q``, ``latent`` / ``compressed_kv`` / ``k_pe``, ``fused_q``) or
+    anything derived from them. This holds for every proxy, TMA included:
+    the outputs are generic-proxy stores of a grid that has completed when
+    that wait returns;
+  - write nothing before its wait that this kernel reads or writes: ``kva``,
+    the weights, the per-layer counters, the outputs, and any graph-pool
+    memory that may back them;
+  - read nothing before its wait that this kernel writes;
+  - never spin on this kernel's progress before its wait.
+
+These rules do not depend on ``early_trigger``: with the implicit trigger
+the dependent is released when this grid's programs exit, and the same
+rules apply. The served successors follow them. The native rope kernel
+reads its inputs after its wait. rope_kv_append's fused rope reads ``fused_q`` and
+``latent`` after its wait. Its per-step metadata reads happen after the
+wait by default; in its opt-in prefetch mode they happen before the wait,
+under rope_kv_append's own condition on the metadata writers. That condition is
+unaffected by this trigger, because the trigger comes after this kernel's
+own wait.
+
+Synchronization of the Triton kernel, per layer: with ``epoch_sync``
+(default) monotonic 64-bit counters ``epochs`` = [finished programs,
+released norm programs, published no-RoPE tiles per head] that are never
+reset; a program derives the launch index from the finished-program count
+(the grid is the same in every launch) and waits for ``(launch + 1) *
+expected`` releases. Without it, absorbed_query's int32 counters ``sync`` that the last
+consumer resets to zero. Every spin-wait targets counters released by lower
+program ids (norm programs first, then producers, then the extra
+absorption programs), so the progress argument of ``kernels.py`` holds. The
+CuTe kernel has no inter-cluster dependency.
+
+Numerics: ``q``, the absorbed query and ``latent`` are bit-identical to the
+native ``forward_dsa_proj`` + absorption BMM when the per-lane fold order
+(``SharedQueryKernelConfig.norm_order``) matches the native RMSNorm kernel; any
+other order is a different rounding and is rejected by the contract tests.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from ...dispatch import log_capability
+from . import kernels as aq_kernels
+from .query_proj import (
+    _HEAD_DIM,
+    _KV_LORA,
+    _LOCAL_HEADS,
+    _NOPE,
+    _Q_LORA,
+    _ROPE,
+    AbsorbedQueryPacket,
+    absorb_layout,
+    call_eligibility,
+    check_package_selectors,
+    module_eligibility,
+)
+
+# Rows of one shared_indexer_absorbed_query launch (the absorbed_query support domain).
+_SHARED_QUERY_MAX_ROWS = 16
+# Per-lane fold orders of the RMSNorm reproduction (see _fi_lane_sumsq).
+NORM_ORDERS = (0, 1, 2, 3, 4, 5, 6)
+
+
+@dataclass(frozen=True)
+class SharedQueryKernelConfig:
+    # The CuTe-DSL cluster kernel (``shared_indexer_cute``) for R <= 8 rows;
+    # the Triton kernel below otherwise.
+    cute: bool = True
+    # CuTe kernel: copy the weights into shared memory before the
+    # grid-dependency wait (overlapping the kv_a GEMM).
+    cute_weights_before_wait: bool = True
+    # CuTe kernel: st.async pushes + per-CTA mbarrier (True) or plain
+    # cluster stores + one cluster barrier (False).
+    cute_push_async: bool = True
+    block_m: int = 16
+    # Rows per norm program (the native FlashInfer kernel uses 4 per CTA).
+    norm_rows: int = 4
+    bn: int = 32
+    bk: int = 256
+    num_warps: int = 4
+    num_stages: int = 4
+    # Per-lane fold order of the native RMSNorm reproduction.
+    norm_order: int = 0
+    # Order 0 with 16-byte loads (one vector per lane and 256-column block).
+    norm_vectorized: bool = True
+    # q_b tile and absorption computed transposed (weights as the MMA M
+    # operand, the padded token rows as N); needs bn, bc in {64, 128}.
+    gemm_transposed: bool = False
+    absorb_transposed: bool = False
+    # L2 prefetch of the q_b tile and W_UK slice before griddepcontrol.wait.
+    prefetch_l2: bool = True
+    # griddepcontrol.launch_dependents right after each program's wait
+    # (see the PDL contract in the module docstring), so the successor (the
+    # native rope kernel or rope_kv_append's fused rope kernel) can be scheduled while
+    # this kernel runs; the successor still waits for this grid before it
+    # reads any output.
+    early_trigger: bool = True
+    # Load each producer's W_UK slice into registers before the wait.
+    preload_wuk: bool = True
+    # Counter poll: 0 atom.add 0 (absorbed_query), 1 ld.acquire.gpu, 2 nanosleep + ld.
+    poll_mode: int = 1
+    # Monotonic 64-bit counters with per-launch targets instead of absorbed_query's
+    # last-consumer resets (no program waits for an atomic's return value).
+    epoch_sync: bool = True
+    # Tests only: skew programs by up to this many ns (kernels._stress_delay).
+    debug_delay_ns: int = 0
+    # Tests only: keep the normalized q_a rows (self.debug["qn"]).
+    debug_keep_qn: bool = False
+    # Tests only: per-program %globaltimer stamps (self.debug["ts"]).
+    debug_timestamps: bool = False
+
+
+# --------------------------------------------------------------------------
+# Reproduction of the native FlashInfer CuTe-DSL RMSNorm (H in {512, 2048}):
+# one warp per row, lane t owns columns 256 * b + 8 * t + v (b < H / 256,
+# v < 8), per-lane fold of the squares, butterfly over the 32 lanes,
+# rstd = rsqrt(sum / H + eps), y = bf16((x * rstd) * (w + 0)).
+# --------------------------------------------------------------------------
+
+
+@triton.jit
+def _tree8(a0, a1, a2, a3, a4, a5, a6, a7):
+    return ((a0 + a1) + (a2 + a3)) + ((a4 + a5) + (a6 + a7))
+
+
+@triton.jit
+def _ld_sq(base, rmask, off: tl.constexpr):
+    x = tl.load(base + off, mask=rmask[:, None], other=0.0).to(tl.float32)
+    return x * x
+
+
+@triton.jit
+def _tree8_at(base, rmask, b: tl.constexpr):
+    o: tl.constexpr = 256 * b
+    return _tree8(
+        _ld_sq(base, rmask, o + 0),
+        _ld_sq(base, rmask, o + 1),
+        _ld_sq(base, rmask, o + 2),
+        _ld_sq(base, rmask, o + 3),
+        _ld_sq(base, rmask, o + 4),
+        _ld_sq(base, rmask, o + 5),
+        _ld_sq(base, rmask, o + 6),
+        _ld_sq(base, rmask, o + 7),
+    )
+
+
+@triton.jit
+def _seq8_at(base, rmask, b: tl.constexpr, ROWS: tl.constexpr):
+    o: tl.constexpr = 256 * b
+    s = tl.zeros([ROWS, 32], dtype=tl.float32)
+    for v in tl.static_range(8):
+        s = s + _ld_sq(base, rmask, o + v)
+    return s
+
+
+@triton.jit
+def _fi_lane_sumsq(base, rmask, NB: tl.constexpr, ORDER: tl.constexpr, ROWS: tl.constexpr):
+    """``[ROWS, 32]`` per-lane sums of squares, NB in {2, 8}.
+
+    ``base`` points at column ``8 * lane`` of each row. Orders: 0 sequential
+    (v fastest), 1 sequential (b fastest), 2 pairwise tree, 3 per-v
+    accumulators over b then sequential over v, 4 the same then a tree over
+    v, 5 per-b tree over v then sequential over b, 6 per-b sequential over v
+    then a tree over b.
+    """
+    s = tl.zeros([ROWS, 32], dtype=tl.float32)
+    if ORDER == 0:
+        for b in tl.static_range(NB):
+            for v in tl.static_range(8):
+                s = s + _ld_sq(base, rmask, 256 * b + v)
+    elif ORDER == 1:
+        for v in tl.static_range(8):
+            for b in tl.static_range(NB):
+                s = s + _ld_sq(base, rmask, 256 * b + v)
+    elif ORDER == 2:
+        if NB == 2:
+            s = _tree8_at(base, rmask, 0) + _tree8_at(base, rmask, 1)
+        else:
+            s = _tree8(
+                _tree8_at(base, rmask, 0),
+                _tree8_at(base, rmask, 1),
+                _tree8_at(base, rmask, 2),
+                _tree8_at(base, rmask, 3),
+                _tree8_at(base, rmask, 4),
+                _tree8_at(base, rmask, 5),
+                _tree8_at(base, rmask, 6),
+                _tree8_at(base, rmask, 7),
+            )
+    elif ORDER == 3 or ORDER == 4:
+        a0 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a1 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a2 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a3 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a4 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a5 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a6 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        a7 = tl.zeros([ROWS, 32], dtype=tl.float32)
+        for b in tl.static_range(NB):
+            a0 = a0 + _ld_sq(base, rmask, 256 * b + 0)
+            a1 = a1 + _ld_sq(base, rmask, 256 * b + 1)
+            a2 = a2 + _ld_sq(base, rmask, 256 * b + 2)
+            a3 = a3 + _ld_sq(base, rmask, 256 * b + 3)
+            a4 = a4 + _ld_sq(base, rmask, 256 * b + 4)
+            a5 = a5 + _ld_sq(base, rmask, 256 * b + 5)
+            a6 = a6 + _ld_sq(base, rmask, 256 * b + 6)
+            a7 = a7 + _ld_sq(base, rmask, 256 * b + 7)
+        if ORDER == 3:
+            s = ((((((a0 + a1) + a2) + a3) + a4) + a5) + a6) + a7
+        else:
+            s = _tree8(a0, a1, a2, a3, a4, a5, a6, a7)
+    elif ORDER == 5:
+        for b in tl.static_range(NB):
+            s = s + _tree8_at(base, rmask, b)
+    else:
+        if NB == 2:
+            s = _seq8_at(base, rmask, 0, ROWS) + _seq8_at(base, rmask, 1, ROWS)
+        else:
+            s = _tree8(
+                _seq8_at(base, rmask, 0, ROWS),
+                _seq8_at(base, rmask, 1, ROWS),
+                _seq8_at(base, rmask, 2, ROWS),
+                _seq8_at(base, rmask, 3, ROWS),
+                _seq8_at(base, rmask, 4, ROWS),
+                _seq8_at(base, rmask, 5, ROWS),
+                _seq8_at(base, rmask, 6, ROWS),
+                _seq8_at(base, rmask, 7, ROWS),
+            )
+    return s
+
+
+@triton.jit
+def _fi_butterfly32(s, ROWS: tl.constexpr):
+    """Sum over the 32 lanes in shuffle-xor butterfly order (offsets 1, 2,
+    4, 8, 16), i.e. a balanced pairwise tree in lane order."""
+    s = tl.sum(tl.reshape(s, [ROWS, 16, 2]), axis=2)
+    s = tl.sum(tl.reshape(s, [ROWS, 8, 2]), axis=2)
+    s = tl.sum(tl.reshape(s, [ROWS, 4, 2]), axis=2)
+    s = tl.sum(tl.reshape(s, [ROWS, 2, 2]), axis=2)
+    return tl.sum(s, axis=1)
+
+
+@triton.jit
+def _fi_rmsnorm_rows(
+    x_ptr,
+    stride_x,
+    w_ptr,
+    y_ptr,
+    stride_y,
+    rows,
+    rmask,
+    eps,
+    H: tl.constexpr,
+    ROWS: tl.constexpr,
+    ORDER: tl.constexpr,
+):
+    """``y[rows, :H] = rmsnorm(x[rows, :H]) * w`` as the native kernel."""
+    lanes = tl.arange(0, 32)
+    base = x_ptr + rows[:, None] * stride_x + 8 * lanes[None, :]
+    s = _fi_lane_sumsq(base, rmask, H // 256, ORDER, ROWS)
+    rstd = tl.rsqrt(_fi_butterfly32(s, ROWS) / H + eps)
+    for c in tl.static_range(H // 256):
+        cols = c * 256 + tl.arange(0, 256)
+        x = tl.load(
+            x_ptr + rows[:, None] * stride_x + cols[None, :], mask=rmask[:, None], other=0.0
+        ).to(tl.float32)
+        w = tl.load(w_ptr + cols).to(tl.float32) + 0.0
+        y = (x * rstd[:, None]) * w[None, :]
+        tl.store(
+            y_ptr + rows[:, None] * stride_y + cols[None, :], y.to(tl.bfloat16), mask=rmask[:, None]
+        )
+
+
+@triton.jit
+def _v8_leaves(x, ROWS: tl.constexpr):
+    """The 8 columns ``v`` of ``x [ROWS, 32, 8]`` as ``[ROWS, 32]`` tensors,
+    in ``v`` order (register-level reshape/split, no data movement)."""
+    e, o = tl.split(tl.reshape(x, [ROWS, 32, 4, 2]))
+    ee, eo = tl.split(tl.reshape(e, [ROWS, 32, 2, 2]))
+    oe, oo = tl.split(tl.reshape(o, [ROWS, 32, 2, 2]))
+    v0, v4 = tl.split(ee)
+    v2, v6 = tl.split(eo)
+    v1, v5 = tl.split(oe)
+    v3, v7 = tl.split(oo)
+    return v0, v1, v2, v3, v4, v5, v6, v7
+
+
+@triton.jit
+def _fi_rmsnorm_rows_vec(
+    x_ptr,
+    stride_x,
+    w_ptr,
+    y_ptr,
+    stride_y,
+    rows,
+    rmask,
+    eps,
+    H: tl.constexpr,
+    ROWS: tl.constexpr,
+    ts_ptr=None,
+    pid=0,
+    TS: tl.constexpr = False,
+):
+    """Fold order 0 of ``_fi_rmsnorm_rows`` with 16-byte loads: each lane
+    loads its 8 consecutive columns of a 256-column block at once and folds
+    them in column order (blocks outer, columns inner)."""
+    lanes = tl.arange(0, 32)
+    vv = tl.arange(0, 8)
+    offs = rows[:, None, None] * stride_x + 8 * lanes[None, :, None] + vv[None, None, :]
+    s = tl.zeros([ROWS, 32], dtype=tl.float32)
+    for b in tl.static_range(H // 256):
+        x = tl.load(x_ptr + offs + 256 * b, mask=rmask[:, None, None], other=0.0).to(tl.float32)
+        v0, v1, v2, v3, v4, v5, v6, v7 = _v8_leaves(x, ROWS)
+        s = s + v0 * v0
+        s = s + v1 * v1
+        s = s + v2 * v2
+        s = s + v3 * v3
+        s = s + v4 * v4
+        s = s + v5 * v5
+        s = s + v6 * v6
+        s = s + v7 * v7
+    rstd = tl.rsqrt(_fi_butterfly32(s, ROWS) / H + eps)
+    if TS:
+        tl.debug_barrier()
+        _ts(ts_ptr, pid, 3, TS)
+    # One load of the whole rows and gain, then one store: separate loads
+    # per block would each wait behind the previous block's stores.
+    cols = tl.arange(0, H)
+    x = tl.load(
+        x_ptr + rows[:, None] * stride_x + cols[None, :], mask=rmask[:, None], other=0.0
+    ).to(tl.float32)
+    w = tl.load(w_ptr + cols).to(tl.float32) + 0.0
+    y = (x * rstd[:, None]) * w[None, :]
+    tl.store(
+        y_ptr + rows[:, None] * stride_y + cols[None, :], y.to(tl.bfloat16), mask=rmask[:, None]
+    )
+    if TS:
+        tl.debug_barrier()
+        _ts(ts_ptr, pid, 4, TS)
+
+
+@triton.jit
+def _fi_rmsnorm_vec_kernel(
+    x_ptr, stride_x, w_ptr, y_ptr, stride_y, R, eps, H: tl.constexpr, ROWS: tl.constexpr
+):
+    rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS)
+    _fi_rmsnorm_rows_vec(x_ptr, stride_x, w_ptr, y_ptr, stride_y, rows, rows < R, eps, H, ROWS)
+
+
+@triton.jit
+def _fi_rmsnorm_kernel(
+    x_ptr,
+    stride_x,
+    w_ptr,
+    y_ptr,
+    stride_y,
+    R,
+    eps,
+    H: tl.constexpr,
+    ROWS: tl.constexpr,
+    ORDER: tl.constexpr,
+):
+    """Standalone launch of the reproduction (tests and probes)."""
+    rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS)
+    _fi_rmsnorm_rows(x_ptr, stride_x, w_ptr, y_ptr, stride_y, rows, rows < R, eps, H, ROWS, ORDER)
+
+
+def fi_rmsnorm_reproduction(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    order: int,
+    rows_per_program: int = 4,
+    vectorized: bool = False,
+) -> torch.Tensor:
+    """Run the RMSNorm reproduction on ``x [R, H]`` (row stride free);
+    ``vectorized`` selects the 16-byte-load implementation of order 0."""
+    rows, h = x.shape
+    assert h in (512, 2048) and x.stride(1) == 1
+    y = torch.empty(rows, h, dtype=torch.bfloat16, device=x.device)
+    grid = (triton.cdiv(rows, rows_per_program),)
+    if vectorized:
+        assert order == 0
+        _fi_rmsnorm_vec_kernel[grid](
+            x,
+            x.stride(0),
+            weight,
+            y,
+            y.stride(0),
+            rows,
+            eps,
+            H=h,
+            ROWS=rows_per_program,
+            num_warps=4,
+        )
+    else:
+        _fi_rmsnorm_kernel[grid](
+            x,
+            x.stride(0),
+            weight,
+            y,
+            y.stride(0),
+            rows,
+            eps,
+            H=h,
+            ROWS=rows_per_program,
+            ORDER=order,
+            num_warps=4,
+        )
+    return y
+
+
+# --------------------------------------------------------------------------
+# The shared_indexer_absorbed_query kernel.
+# --------------------------------------------------------------------------
+
+
+@triton.jit
+def _l2_prefetch_1k(base_ptr, N: tl.constexpr, n_valid):
+    """Ask L2 to prefetch ``n_valid`` (<= N) consecutive 1 KiB chunks."""
+    idx = tl.minimum(tl.arange(0, N), n_valid - 1)
+    ptrs = base_ptr.to(tl.pointer_type(tl.int8)) + idx * 1024
+    tl.inline_asm_elementwise(
+        "cp.async.bulk.prefetch.L2.global [$1], 1024; mov.u32 $0, 0;",
+        "=r,l",
+        [ptrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _shared_query_gemm(
+    t,
+    x_ptr,
+    stride_x,
+    w_ptr,
+    out_ptr,
+    OUT_COLS: tl.constexpr,
+    rows,
+    rmask,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    """``out[:, t*BN:(t+1)*BN] = x @ W[t*BN:(t+1)*BN]^T`` (absorbed_query's MMA sequence);
+    ``x`` was written by another program of this grid, so it is read with
+    ``.cg`` (L2)."""
+    cols = tl.arange(0, BK)
+    n_offs = t * BN + tl.arange(0, BN)
+    acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32)
+    for k0 in tl.range(0, K, BK):
+        x = tl.load(
+            x_ptr + rows[:, None] * stride_x + (k0 + cols)[None, :],
+            mask=rmask[:, None],
+            other=0.0,
+            cache_modifier=".cg",
+        )
+        w = tl.load(w_ptr + n_offs[None, :] * K + (k0 + cols)[:, None])
+        acc = tl.dot(x, w, acc)
+    tl.store(
+        out_ptr + rows[:, None] * OUT_COLS + n_offs[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def _shared_query_gemm_t(
+    t,
+    x_ptr,
+    stride_x,
+    w_ptr,
+    out_ptr,
+    OUT_COLS: tl.constexpr,
+    rows,
+    rmask,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    """The same tile as ``_shared_query_gemm``, computed transposed:
+    ``out^T [BN, BLOCK_M] = W[t*BN:(t+1)*BN] @ x^T``, so the MMA ``M`` is the
+    weight tile (64 or 128 rows, eligible for tcgen05) and ``N`` the padded
+    token rows. Same products, same ascending K order."""
+    kk = tl.arange(0, BK)
+    n_offs = t * BN + tl.arange(0, BN)
+    acc = tl.zeros([BN, BLOCK_M], dtype=tl.float32)
+    for k0 in tl.range(0, K, BK):
+        w = tl.load(w_ptr + n_offs[:, None] * K + (k0 + kk)[None, :])
+        x = tl.load(
+            x_ptr + rows[None, :] * stride_x + (k0 + kk)[:, None],
+            mask=rmask[None, :],
+            other=0.0,
+            cache_modifier=".cg",
+        )
+        acc = tl.dot(w, x, acc)
+    tl.store(
+        out_ptr + rows[None, :] * OUT_COLS + n_offs[:, None],
+        acc.to(tl.bfloat16),
+        mask=rmask[None, :],
+    )
+
+
+@triton.jit
+def _ld_wuk_t(a, wuk_ptr, KV_LORA: tl.constexpr, NOPE: tl.constexpr, BC: tl.constexpr):
+    """``W_UK[h, c0:c0+BC, :NOPE]`` of absorption job ``a`` as [BC, NOPE]
+    (split 128 + NOPE-128 along NOPE)."""
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    h = a // chunks_per_head
+    c0 = (a % chunks_per_head) * BC
+    cc = tl.arange(0, BC)
+    base = wuk_ptr + h * (KV_LORA * NOPE) + (c0 + cc)[:, None] * NOPE
+    w_lo = tl.load(base + tl.arange(0, 128)[None, :])
+    w_hi = tl.load(base + 128 + tl.arange(0, NOPE - 128)[None, :])
+    return w_lo, w_hi
+
+
+@triton.jit
+def _poll_once(ptr, POLL: tl.constexpr, W64: tl.constexpr = False):
+    """One acquire read of an int32 (``W64``: int64) counter. ``POLL`` 0:
+    ``atom.add 0`` (absorbed_query's protocol); 1: ``ld.acquire.gpu``; 2: ``nanosleep``
+    then ``ld.acquire.gpu``."""
+    if POLL == 0:
+        v = tl.atomic_add(ptr, 0, sem="acquire", scope="gpu")
+    else:
+        p = ptr + tl.zeros([1], dtype=tl.int32)
+        if W64:
+            if POLL == 1:
+                r = tl.inline_asm_elementwise(
+                    "ld.acquire.gpu.global.b64 $0, [$1];",
+                    "=l,l",
+                    [p],
+                    dtype=tl.int64,
+                    is_pure=False,
+                    pack=1,
+                )
+            else:
+                r = tl.inline_asm_elementwise(
+                    "nanosleep.u32 32; ld.acquire.gpu.global.b64 $0, [$1];",
+                    "=l,l",
+                    [p],
+                    dtype=tl.int64,
+                    is_pure=False,
+                    pack=1,
+                )
+        else:
+            if POLL == 1:
+                r = tl.inline_asm_elementwise(
+                    "ld.acquire.gpu.global.b32 $0, [$1];",
+                    "=r,l",
+                    [p],
+                    dtype=tl.int32,
+                    is_pure=False,
+                    pack=1,
+                )
+            else:
+                r = tl.inline_asm_elementwise(
+                    "nanosleep.u32 32; ld.acquire.gpu.global.b32 $0, [$1];",
+                    "=r,l",
+                    [p],
+                    dtype=tl.int32,
+                    is_pure=False,
+                    pack=1,
+                )
+        v = tl.max(r, axis=0)
+    return v
+
+
+@triton.jit
+def _wait_ge(ptr, target, POLL: tl.constexpr, W64: tl.constexpr = False):
+    """Spin until ``*ptr >= target`` (acquire), then a CTA barrier so every
+    thread's later reads are ordered after the acquire."""
+    v = _poll_once(ptr, POLL, W64)
+    while v < target:
+        v = _poll_once(ptr, POLL, W64)
+    tl.debug_barrier()
+
+
+@triton.jit
+def _ts(ts_ptr, pid, k: tl.constexpr, TS: tl.constexpr):
+    """Tests only: record ``%globaltimer`` (ns) at point ``k`` of ``pid``."""
+    if TS:
+        tl.store(ts_ptr + pid * 8 + k, tl.extra.cuda.globaltimer())
+
+
+@triton.jit
+def _shared_query_absorb_job(
+    a,
+    w_lo,
+    w_hi,
+    q_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    ready_ptr,
+    ready_target,
+    rows,
+    rmask,
+    ts_ptr,
+    pid,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BN: tl.constexpr,
+    BC: tl.constexpr,
+    POLL: tl.constexpr,
+    TS: tl.constexpr,
+    ABS_T: tl.constexpr,
+    EPOCH: tl.constexpr,
+    DELAY_NS: tl.constexpr,
+):
+    """Absorption job ``a`` (head ``a // (KV_LORA // BC)``, ``BC`` latent
+    columns) with its ``W_UK`` slice already loaded: absorbed_query's
+    ``kernels._absorb_job`` (same products, same d order) with a selectable
+    poll. Waits for ``ready_ptr[h] >= ready_target``; without ``EPOCH`` it
+    then runs absorbed_query's done/reset protocol on ``sync_ptr``. ``ABS_T`` computes
+    the transposed tile (``W_UK`` slice as the MMA ``M`` operand, from
+    ``_ld_wuk_t``)."""
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    h = a // chunks_per_head
+    c0 = (a % chunks_per_head) * BC
+    cc = tl.arange(0, BC)
+    d_lo = tl.arange(0, 128)
+    d_hi = tl.arange(0, NOPE - 128)
+    _wait_ge(ready_ptr + h, ready_target, POLL, EPOCH)
+    _ts(ts_ptr, pid, 5, TS)
+    if ABS_T:
+        q_base = q_ptr + rows[None, :] * (H * HEAD_DIM) + h * HEAD_DIM
+        q_lo = tl.load(q_base + d_lo[:, None], mask=rmask[None, :], other=0.0, cache_modifier=".cg")
+        q_hi = tl.load(
+            q_base + 128 + d_hi[:, None], mask=rmask[None, :], other=0.0, cache_modifier=".cg"
+        )
+        acc = tl.dot(w_lo, q_lo)
+        acc = tl.dot(w_hi, q_hi, acc)
+        tl.store(
+            fq_ptr + rows[None, :] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[:, None],
+            acc.to(tl.bfloat16),
+            mask=rmask[None, :],
+        )
+    else:
+        q_base = q_ptr + rows[:, None] * (H * HEAD_DIM) + h * HEAD_DIM
+        q_lo = tl.load(q_base + d_lo[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg")
+        q_hi = tl.load(
+            q_base + 128 + d_hi[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+        )
+        acc = tl.dot(q_lo, w_lo)
+        acc = tl.dot(q_hi, w_hi, acc)
+        tl.store(
+            fq_ptr + rows[:, None] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[None, :],
+            acc.to(tl.bfloat16),
+            mask=rmask[:, None],
+        )
+    aq_kernels._stress_delay(a, 3, DELAY_NS)
+    if not EPOCH:
+        tl.debug_barrier()
+        done = tl.atomic_add(sync_ptr + H + h, 1, sem="acq_rel", scope="gpu")
+        if done == chunks_per_head - 1:
+            tl.atomic_xchg(sync_ptr + h, 0, sem="relaxed", scope="gpu")
+            tl.atomic_xchg(sync_ptr + H + h, 0, sem="relaxed", scope="gpu")
+    _ts(ts_ptr, pid, 6, TS)
+
+
+@triton.jit
+def _shared_query_ld_wuk(
+    job,
+    valid,
+    wuk_ptr,
+    KV_LORA: tl.constexpr,
+    NOPE: tl.constexpr,
+    BC: tl.constexpr,
+    ABS_T: tl.constexpr,
+):
+    if ABS_T:
+        w_lo, w_hi = _ld_wuk_t(job, wuk_ptr, KV_LORA, NOPE, BC)
+    else:
+        w_lo, w_hi = aq_kernels._ld_wuk(job, valid, wuk_ptr, KV_LORA, NOPE, BC)
+    return w_lo, w_hi
+
+
+@triton.jit
+def shared_query_kernel(
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    gkv_ptr,
+    wqb_ptr,
+    wuk_ptr,
+    qn_ptr,
+    q_ptr,
+    latent_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    ep_ptr,
+    ts_ptr,
+    R,
+    n_norm,
+    n_programs,
+    eps_q,
+    eps_kv,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    NORM_ROWS: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    BC: tl.constexpr,
+    N_QB: tl.constexpr,
+    N_ABS: tl.constexpr,
+    ABS_JOB0: tl.constexpr,
+    PRODUCER_JOB_SHIFT: tl.constexpr,
+    NORM_ORDER: tl.constexpr,
+    NORM_VEC: tl.constexpr,
+    GEMM_T: tl.constexpr,
+    ABS_T: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    PREFETCH_L2: tl.constexpr,
+    PF_W: tl.constexpr,
+    PF_WUK: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+    PRELOAD_WUK: tl.constexpr,
+    POLL: tl.constexpr,
+    EPOCH: tl.constexpr = False,
+    TS: tl.constexpr = False,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """Programs: ``[0, n_norm)`` norm programs (``NORM_ROWS`` rows each),
+    then ``N_QB`` ``q_b`` producers (producer ``t >= PRODUCER_JOB_SHIFT``
+    then runs absorption job ``t - PRODUCER_JOB_SHIFT``), then ``N_ABS``
+    absorption programs (jobs ``ABS_JOB0 + i``).
+
+    Counter protocols: without ``EPOCH``, absorbed_query's (int32 ``sync_ptr``, reset by
+    the last consumer of every counter). With ``EPOCH`` the counters in the
+    int64 ``ep_ptr`` are never reset: ``ep_ptr[0]`` counts finished
+    programs over all launches, ``ep_ptr[1]`` released norm programs and
+    ``ep_ptr[2 + h]`` published no-RoPE tiles of head ``h``. The grid size
+    ``n_programs`` and ``n_norm`` are the same in every launch, and a program
+    reads ``ep_ptr[0]`` after ``griddepcontrol.wait`` (every earlier launch
+    has finished) and before it finishes, so ``launch = ep_ptr[0] //
+    n_programs`` and its wait targets are ``(launch + 1) * count``. No
+    program waits for an atomic's return value.
+
+    ``TS`` (tests only) records ``%globaltimer`` per program into
+    ``ts_ptr [grid, 8]``: 0 entry, 1 after the grid-dependency wait, 2 norm
+    rows released / norm flag acquired, 3 ``q`` tile stored, 4 tile
+    published, 5 absorption inputs acquired, 6 absorption done, 7 exit (norm
+    programs: 3 rstd ready, 4 rows stored, 5 kv norm done).
+    """
+    pid = tl.program_id(0)
+    _ts(ts_ptr, pid, 0, TS)
+    FLAG: tl.constexpr = 2 * H
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    if pid < n_norm:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+            if EARLY_TRIGGER:
+                tl.extra.cuda.gdc_launch_dependents()
+        _ts(ts_ptr, pid, 1, TS)
+        # (Names differ from the other roles' [BLOCK_M] row vectors: Triton
+        # merges same-named values across the branches of a runtime if.)
+        nrows = pid * NORM_ROWS + tl.arange(0, NORM_ROWS)
+        nmask = nrows < R
+        if NORM_VEC:
+            _fi_rmsnorm_rows_vec(
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                qn_ptr,
+                Q_LORA,
+                nrows,
+                nmask,
+                eps_q,
+                Q_LORA,
+                NORM_ROWS,
+                ts_ptr,
+                pid,
+                TS,
+            )
+        else:
+            _fi_rmsnorm_rows(
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                qn_ptr,
+                Q_LORA,
+                nrows,
+                nmask,
+                eps_q,
+                Q_LORA,
+                NORM_ROWS,
+                NORM_ORDER,
+            )
+        aq_kernels._stress_delay(pid, 11, DELAY_NS)
+        # Every thread's qn stores precede the release (CTA barrier).
+        tl.debug_barrier()
+        if EPOCH:
+            tl.atomic_add(ep_ptr + 1, 1, sem="release", scope="gpu")
+        else:
+            tl.atomic_add(sync_ptr + FLAG, 1, sem="release", scope="gpu")
+        _ts(ts_ptr, pid, 2, TS)
+        # Off the critical path: latent = [rmsnorm(compressed_kv) | k_pe].
+        if NORM_VEC:
+            _fi_rmsnorm_rows_vec(
+                kva_ptr + Q_LORA,
+                stride_kva,
+                gkv_ptr,
+                latent_ptr,
+                KV_LORA + ROPE,
+                nrows,
+                nmask,
+                eps_kv,
+                KV_LORA,
+                NORM_ROWS,
+            )
+        else:
+            _fi_rmsnorm_rows(
+                kva_ptr + Q_LORA,
+                stride_kva,
+                gkv_ptr,
+                latent_ptr,
+                KV_LORA + ROPE,
+                nrows,
+                nmask,
+                eps_kv,
+                KV_LORA,
+                NORM_ROWS,
+                NORM_ORDER,
+            )
+        rc = tl.arange(0, ROPE)
+        k_pe = tl.load(
+            kva_ptr + nrows[:, None] * stride_kva + Q_LORA + KV_LORA + rc[None, :],
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        tl.store(
+            latent_ptr + nrows[:, None] * (KV_LORA + ROPE) + KV_LORA + rc[None, :],
+            k_pe,
+            mask=nmask[:, None],
+        )
+        _ts(ts_ptr, pid, 5, TS)
+        if EPOCH:
+            tl.atomic_add(ep_ptr, 1, sem="relaxed", scope="gpu")
+    elif pid < n_norm + N_QB:
+        t = pid - n_norm
+        rows = tl.arange(0, BLOCK_M)
+        rmask = rows < R
+        job = tl.maximum(t - PRODUCER_JOB_SHIFT, 0)
+        wuk_job = (
+            wuk_ptr
+            + (job // chunks_per_head) * (KV_LORA * NOPE)
+            + (job % chunks_per_head) * BC * NOPE
+        )
+        if PREFETCH_L2:
+            # Weights only (never written by an earlier grid).
+            _l2_prefetch_1k(wqb_ptr + t * BN * Q_LORA, PF_W, (BN * Q_LORA * 2) // 1024)
+            if t >= PRODUCER_JOB_SHIFT:
+                _l2_prefetch_1k(wuk_job, PF_WUK, (BC * NOPE * 2) // 1024)
+        if PRELOAD_WUK:
+            w_lo, w_hi = _shared_query_ld_wuk(
+                job, t >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC, ABS_T
+            )
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+            if EARLY_TRIGGER:
+                tl.extra.cuda.gdc_launch_dependents()
+        _ts(ts_ptr, pid, 1, TS)
+        if EPOCH:
+            launch1 = tl.load(ep_ptr, volatile=True) // n_programs + 1
+            _wait_ge(ep_ptr + 1, launch1 * n_norm, POLL, True)
+            _ts(ts_ptr, pid, 2, TS)
+        else:
+            # Wait for every norm program, then count this producer out; the
+            # last producer resets the flag for the next launch.
+            _wait_ge(sync_ptr + FLAG, n_norm, POLL)
+            _ts(ts_ptr, pid, 2, TS)
+            done = tl.atomic_add(sync_ptr + FLAG + 1, 1, sem="acq_rel", scope="gpu")
+            if done == N_QB - 1:
+                tl.atomic_xchg(sync_ptr + FLAG, 0, sem="relaxed", scope="gpu")
+                tl.atomic_xchg(sync_ptr + FLAG + 1, 0, sem="relaxed", scope="gpu")
+        if GEMM_T:
+            _shared_query_gemm_t(
+                t,
+                qn_ptr,
+                Q_LORA,
+                wqb_ptr,
+                q_ptr,
+                H * HEAD_DIM,
+                rows,
+                rmask,
+                Q_LORA,
+                BLOCK_M,
+                BN,
+                BK,
+            )
+        else:
+            _shared_query_gemm(
+                t,
+                qn_ptr,
+                Q_LORA,
+                wqb_ptr,
+                q_ptr,
+                H * HEAD_DIM,
+                rows,
+                rmask,
+                Q_LORA,
+                BLOCK_M,
+                BN,
+                BK,
+            )
+        _ts(ts_ptr, pid, 3, TS)
+        if (t % (HEAD_DIM // BN)) < (NOPE // BN):
+            aq_kernels._stress_delay(t, 1, DELAY_NS)
+            tl.debug_barrier()
+            if EPOCH:
+                tl.atomic_add(ep_ptr + 2 + t // (HEAD_DIM // BN), 1, sem="release", scope="gpu")
+            else:
+                tl.atomic_add(sync_ptr + t // (HEAD_DIM // BN), 1, sem="release", scope="gpu")
+        _ts(ts_ptr, pid, 4, TS)
+        if t >= PRODUCER_JOB_SHIFT:
+            if not PRELOAD_WUK:
+                w_lo, w_hi = _shared_query_ld_wuk(job, True, wuk_ptr, KV_LORA, NOPE, BC, ABS_T)
+            if EPOCH:
+                _shared_query_absorb_job(
+                    job,
+                    w_lo,
+                    w_hi,
+                    q_ptr,
+                    fq_ptr,
+                    stride_fq_row,
+                    sync_ptr,
+                    ep_ptr + 2,
+                    launch1 * (NOPE // BN),
+                    rows,
+                    rmask,
+                    ts_ptr,
+                    pid,
+                    KV_LORA,
+                    ROPE,
+                    HEAD_DIM,
+                    NOPE,
+                    H,
+                    BN,
+                    BC,
+                    POLL,
+                    TS,
+                    ABS_T,
+                    True,
+                    DELAY_NS,
+                )
+            else:
+                _shared_query_absorb_job(
+                    job,
+                    w_lo,
+                    w_hi,
+                    q_ptr,
+                    fq_ptr,
+                    stride_fq_row,
+                    sync_ptr,
+                    sync_ptr,
+                    NOPE // BN,
+                    rows,
+                    rmask,
+                    ts_ptr,
+                    pid,
+                    KV_LORA,
+                    ROPE,
+                    HEAD_DIM,
+                    NOPE,
+                    H,
+                    BN,
+                    BC,
+                    POLL,
+                    TS,
+                    ABS_T,
+                    False,
+                    DELAY_NS,
+                )
+        if EPOCH:
+            tl.atomic_add(ep_ptr, 1, sem="relaxed", scope="gpu")
+    else:
+        rows = tl.arange(0, BLOCK_M)
+        rmask = rows < R
+        a = ABS_JOB0 + pid - n_norm - N_QB
+        w_lo, w_hi = _shared_query_ld_wuk(a, True, wuk_ptr, KV_LORA, NOPE, BC, ABS_T)
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+            if EARLY_TRIGGER:
+                tl.extra.cuda.gdc_launch_dependents()
+        _ts(ts_ptr, pid, 1, TS)
+        if EPOCH:
+            launch1 = tl.load(ep_ptr, volatile=True) // n_programs + 1
+            _shared_query_absorb_job(
+                a,
+                w_lo,
+                w_hi,
+                q_ptr,
+                fq_ptr,
+                stride_fq_row,
+                sync_ptr,
+                ep_ptr + 2,
+                launch1 * (NOPE // BN),
+                rows,
+                rmask,
+                ts_ptr,
+                pid,
+                KV_LORA,
+                ROPE,
+                HEAD_DIM,
+                NOPE,
+                H,
+                BN,
+                BC,
+                POLL,
+                TS,
+                ABS_T,
+                True,
+                DELAY_NS,
+            )
+            tl.atomic_add(ep_ptr, 1, sem="relaxed", scope="gpu")
+        else:
+            _shared_query_absorb_job(
+                a,
+                w_lo,
+                w_hi,
+                q_ptr,
+                fq_ptr,
+                stride_fq_row,
+                sync_ptr,
+                sync_ptr,
+                NOPE // BN,
+                rows,
+                rmask,
+                ts_ptr,
+                pid,
+                KV_LORA,
+                ROPE,
+                HEAD_DIM,
+                NOPE,
+                H,
+                BN,
+                BC,
+                POLL,
+                TS,
+                ABS_T,
+                False,
+                DELAY_NS,
+            )
+    _ts(ts_ptr, pid, 7, TS)
+
+
+class SharedIndexerProjection:
+    """query_front producer (preabsorbed form) for one shared-indexer layer."""
+
+    def __init__(self, mla, config: SharedQueryKernelConfig = SharedQueryKernelConfig()):
+        assert config.norm_order in NORM_ORDERS
+        self.mla = mla
+        self.config = config
+        device = mla.q_b_proj.weight.device
+        self.sync = torch.zeros(2 * _LOCAL_HEADS + 2, dtype=torch.int32, device=device)
+        # Epoch protocol: [finished programs, released norm programs,
+        # published no-RoPE tiles per head]; never reset.
+        self.epochs = torch.zeros(2 + _LOCAL_HEADS, dtype=torch.int64, device=device)
+        from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+        self.use_pdl = get_env_enable_pdl()
+        self.debug = {}
+
+    def launch(self, kva: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Run the shared_indexer_absorbed_query kernel on a native kv_a output ``[R, 2624]``.
+
+        Returns ``(q [R, H*256], latent [R, 576], fused_q [R, H, 576])``;
+        ``fused_q[..., :512]`` holds the absorbed query, ``fused_q[...,
+        512:]`` is left for the native RoPE kernel.
+        """
+        mla = self.mla
+        cfg = self.config
+        rows = kva.shape[0]
+        assert kva.stride(1) == 1 and 1 <= rows <= cfg.block_m
+        if cfg.cute and rows <= 8 and kva.is_contiguous():
+            from . import shared_indexer_cute
+
+            qn = kva.new_empty(rows, _Q_LORA) if cfg.debug_keep_qn else None
+            ts = self._ts_buffer(72, kva.device) if cfg.debug_timestamps else None
+            q, latent, fused_q = shared_indexer_cute.launch(
+                mla,
+                kva,
+                self.use_pdl,
+                cfg.early_trigger,
+                qn,
+                ts,
+                cfg.cute_weights_before_wait,
+                cfg.cute_push_async,
+            )
+            if cfg.debug_keep_qn:
+                self.debug["qn"] = qn
+            if cfg.debug_timestamps:
+                self.debug["ts"] = ts
+                self.debug["grid"] = dict(cute=True)
+            return q, latent, fused_q
+        q = kva.new_empty(rows, _LOCAL_HEADS * _HEAD_DIM)
+        qn = kva.new_empty(rows, _Q_LORA)
+        latent = kva.new_empty(rows, _KV_LORA + _ROPE)
+        fused_q = kva.new_empty(rows, _LOCAL_HEADS, _KV_LORA + _ROPE)
+        n_qb = (_LOCAL_HEADS * _HEAD_DIM) // cfg.bn
+        layout = absorb_layout(n_qb, True, True, True, 64)
+        pf_w = triton.next_power_of_2((cfg.bn * _Q_LORA * 2) // 1024)
+        pf_wuk = triton.next_power_of_2(max(1, (layout["BC"] * _NOPE * 2) // 1024))
+        # The epoch protocol needs the same grid in every launch.
+        n_norm = triton.cdiv(cfg.block_m if cfg.epoch_sync else rows, cfg.norm_rows)
+        grid = (n_norm + n_qb + layout["N_ABS"],)
+        ts = self._ts_buffer(grid[0], kva.device) if cfg.debug_timestamps else self.sync
+        compiled = shared_query_kernel[grid](
+            kva,
+            kva.stride(0),
+            mla.q_a_layernorm.weight,
+            mla.kv_a_layernorm.weight,
+            mla.q_b_proj.weight,
+            mla.k_b_proj_trans,
+            qn,
+            q,
+            latent,
+            fused_q,
+            fused_q.stride(0),
+            self.sync,
+            self.epochs,
+            ts,
+            rows,
+            n_norm,
+            grid[0],
+            mla.q_a_layernorm.variance_epsilon,
+            mla.kv_a_layernorm.variance_epsilon,
+            Q_LORA=_Q_LORA,
+            KV_LORA=_KV_LORA,
+            ROPE=_ROPE,
+            HEAD_DIM=_HEAD_DIM,
+            NOPE=_NOPE,
+            H=_LOCAL_HEADS,
+            BLOCK_M=cfg.block_m,
+            NORM_ROWS=cfg.norm_rows,
+            BN=cfg.bn,
+            BK=cfg.bk,
+            N_QB=n_qb,
+            NORM_ORDER=cfg.norm_order,
+            NORM_VEC=cfg.norm_vectorized and cfg.norm_order == 0,
+            GEMM_T=cfg.gemm_transposed,
+            ABS_T=cfg.absorb_transposed,
+            USE_PDL=self.use_pdl,
+            PREFETCH_L2=cfg.prefetch_l2,
+            PF_W=pf_w,
+            PF_WUK=pf_wuk,
+            EARLY_TRIGGER=cfg.early_trigger,
+            PRELOAD_WUK=cfg.preload_wuk,
+            POLL=cfg.poll_mode,
+            EPOCH=cfg.epoch_sync,
+            TS=cfg.debug_timestamps,
+            DELAY_NS=cfg.debug_delay_ns,
+            **layout,
+            num_warps=cfg.num_warps,
+            num_stages=cfg.num_stages,
+            launch_pdl=self.use_pdl,
+        )
+        self.launch_shape = (n_norm, grid[0])
+        if cfg.debug_keep_qn:
+            self.debug["qn"] = qn
+        if cfg.debug_timestamps:
+            self.debug["compiled"] = compiled
+            self.debug["ts"] = ts
+            self.debug["grid"] = dict(n_norm=n_norm, n_qb=n_qb, **layout)
+        return q, latent, fused_q
+
+    def _ts_buffer(self, programs: int, device) -> torch.Tensor:
+        """Tests: a persistent stamp buffer (no allocation kernel between
+        the kv_a GEMM and the shared_indexer_absorbed_query kernel inside a captured graph)."""
+        buf = self.debug.get("ts_buf")
+        if buf is None or buf.shape[0] != programs:
+            buf = torch.zeros(programs, 8, dtype=torch.int64, device=device)
+            self.debug["ts_buf"] = buf
+        return buf
+
+    def counters_quiescent(self) -> bool:
+        """Tests: the counters are consistent with completed launches only
+        (absorbed_query protocol: all zero; epoch protocol: every counter at the same
+        whole number of launches)."""
+        if int(self.sync.abs().sum()) != 0:
+            return False
+        if not self.config.epoch_sync:
+            return int(self.epochs.abs().sum()) == 0
+        e = self.epochs.tolist()
+        if not hasattr(self, "launch_shape"):
+            return not any(e)
+        n_norm, grid = self.launch_shape
+        launches, rem = divmod(e[0], grid)
+        tiles = _NOPE // self.config.bn
+        return rem == 0 and e[1] == launches * n_norm and all(v == launches * tiles for v in e[2:])
+
+    def project(self, hidden_states: torch.Tensor) -> AbsorbedQueryPacket:
+        """Native kv_a GEMM (bridge), then the shared_indexer_absorbed_query kernel."""
+        kva = self.mla.kv_a_proj_with_mqa(hidden_states)
+        q, latent, fused_q = self.launch(kva)
+        return AbsorbedQueryPacket([q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent], fused_q)
+
+
+def _native_rmsnorm_is_flashinfer_cute(norm, hidden: int) -> Tuple[bool, str]:
+    """Whether ``norm`` (an ``RMSNorm``) runs the FlashInfer CuTe-DSL kernel
+    this package reproduces."""
+    from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE
+
+    if not IS_FLASHINFER_AVAILABLE:
+        return False, "FlashInfer not available"
+    if (
+        getattr(norm, "use_cuda_tile", False)
+        or norm.use_gemma
+        or getattr(norm, "is_nvfp4", False)
+        or getattr(norm, "return_hp_output", False)
+    ):
+        return False, "RMSNorm variant not reproduced"
+    w = norm.weight
+    if w.dtype != torch.bfloat16 or tuple(w.shape) != (hidden,):
+        return False, "RMSNorm weight must be BF16"
+    try:
+        from flashinfer.norm import _use_cuda_norm
+
+        if _use_cuda_norm(w.device):
+            return False, "FlashInfer CUDA (non-CuTe) RMSNorm selected"
+    except ImportError:
+        return False, "FlashInfer norm dispatch not inspectable"
+    return True, ""
+
+
+def shared_module_eligibility(mla) -> Tuple[bool, str]:
+    """Static support check of shared_indexer_absorbed_query; no GPU work."""
+    ok, reason = module_eligibility(mla)
+    if not ok:
+        return ok, reason
+    if mla.mqa.indexer is not None:
+        return False, "full indexer layer (shared_indexer_absorbed_query serves shared layers)"
+    for name, hidden in (("q_a_layernorm", _Q_LORA), ("kv_a_layernorm", _KV_LORA)):
+        ok, reason = _native_rmsnorm_is_flashinfer_cute(getattr(mla, name), hidden)
+        if not ok:
+            return False, f"{name}: {reason}"
+    return True, ""
+
+
+def shared_dsa_proj(
+    mla, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """shared_indexer_absorbed_query: preabsorbed query_front packet for a shared-indexer target layer.
+
+    Returns ``None`` (before any GPU work) for full-indexer and draft layers
+    and outside the support domain; the caller then runs absorbed_query's ``dsa_proj``
+    (or the native path). The per-module decision is cached in
+    ``mla._shared_indexer_query`` and logged once under ``shared_indexer_absorbed_query``.
+    """
+    del position_ids  # shared layers have no indexer RoPE
+    runner = getattr(mla, "_shared_indexer_query", None)
+    if runner is None:
+        check_package_selectors()
+        ok, reason = shared_module_eligibility(mla)
+        log_capability("shared_indexer_absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+        runner = SharedIndexerProjection(mla) if ok else False
+        mla._shared_indexer_query = runner
+    if runner is False:
+        return None
+    ok, _ = call_eligibility(mla, hidden_states, attn_metadata)
+    if not ok or hidden_states.shape[0] > _SHARED_QUERY_MAX_ROWS:
+        return None
+    return runner.project(hidden_states)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer_cute.py
new file mode 100644
index 000000000000..c97fccd3aa13
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/shared_indexer_cute.py
@@ -0,0 +1,945 @@
+# 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.
+"""shared_indexer_absorbed_query: CuTe-DSL cluster kernel for the shared-indexer-layer query region.
+
+Same contract as ``shared_indexer.shared_query_kernel`` (the preabsorbed
+query_front packet of a GLM-5.x DSA layer without its own indexer, bit for
+bit equal to the native ``forward_dsa_proj`` + absorption BMM), computed
+without exchanging data through global memory:
+
+* grid of 9 thread-block clusters of 8 CTAs (256 threads each); cluster
+  ``h < 8`` serves local attention head ``h``, CTA ``j`` of it the 32
+  ``q_b`` output columns ``[h * 256 + 32 j, +32)`` (``j < 6``: no-RoPE,
+  ``j = 6, 7``: RoPE) and the 64 absorbed latent columns ``[64 j, +64)``;
+* before ``griddepcontrol.wait`` a CTA copies its ``q_b`` weight tile (32 x
+  2048 BF16), its ``W_UK`` slice (64 x 192 BF16) and the q_a norm gain
+  into shared memory with ``cp.async`` (weights only; under programmatic
+  dependent launch this overlaps the preceding kv_a GEMM);
+* after the wait every CTA reproduces the native ``q_a_layernorm``
+  (FlashInfer CuTe-DSL ``RMSNormKernel``: one warp per row, lane ``t`` owns
+  columns ``256 b + 8 t + v``, sequential fold over ``b`` then ``v``,
+  butterfly over lanes, ``rsqrt(sum / H + eps)`` (fast-math), ``bf16((x *
+  rstd) * w)``) for all rows into its own shared memory -- no cross-CTA
+  flag;
+* two warps run the ``q_b`` tile as ``mma.sync.m16n8k16`` chains over K =
+  2048 in ascending order (weights as the A operand, the token rows as the
+  8-wide B operand: the same products and order as the Triton kernels);
+* the no-RoPE CTAs push their ``q`` tile into every cluster peer's shared
+  memory with ``st.async`` completing on the receiver's mbarrier (each CTA
+  waits only for the bytes it receives; the barriers are initialized and
+  made cluster-visible before the dependency wait); every CTA then absorbs
+  its 64 latent columns with a 12-step ``mma.sync`` chain over d = 0..191
+  and stores them;
+* cluster 8 reproduces the native ``kv_a_layernorm`` (H = 512) into
+  ``latent[:, :512]`` and copies ``k_pe`` (one row per CTA).
+
+Support domain of this kernel: 1 <= R <= 8 token rows (one 8-wide MMA
+tile), the GLM-5.2 TP8 shapes, BF16, sm100 / sm103.
+"""
+
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import Float32, Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.utils.smem_allocator import SmemAllocator
+
+_Q_LORA = 2048
+_KV_LORA = 512
+_ROPE = 64
+_NOPE = 192
+_HEAD_DIM = 256
+_H = 8
+_KVA_COLS = _Q_LORA + _KV_LORA + _ROPE  # 2624
+_MAX_ROWS = 8
+
+# Shared-memory row strides (bytes) padded by 16 B so that the 8 rows an
+# ldmatrix / fragment load touches fall into distinct banks.
+_W_STRIDE = _Q_LORA * 2 + 16  # q_b weight tile rows and normalized q_a rows
+_UK_STRIDE = _NOPE * 2 + 16  # W_UK slice rows and gathered q_nope rows
+_TILE_COLS = 32  # q_b columns per CTA
+_ABS_COLS = 64  # latent columns per CTA
+
+# --------------------------------------------------------------------------
+# PTX helpers.
+# --------------------------------------------------------------------------
+
+
+def _asm(res, args, text, cons, side_effects=True, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=side_effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _gaddr(t: cute.Tensor, *, loc=None, ip=None) -> Int64:
+    """64-bit global address of a tensor's first element."""
+    return Int64(llvm.ptrtoint(T.i64(), t.iterator.llvm_ptr, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def _saddr(p: cute.Pointer, *, loc=None, ip=None) -> Int32:
+    """32-bit shared-window address of a shared-memory pointer."""
+    return Int32(p.toint(loc=loc, ip=ip))
+
+
+@dsl_user_op
+def _cp_async16(saddr: Int32, gaddr: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [saddr.ir_value(), gaddr.ir_value()],
+        "cp.async.cg.shared.global [$0], [$1], 16;",
+        "r,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _cp_async_commit(*, loc=None, ip=None):
+    _asm(None, [], "cp.async.commit_group;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _cp_async_wait_group(n: int, *, loc=None, ip=None):
+    """Wait until at most ``n`` committed cp.async groups are pending."""
+    _asm(None, [], f"cp.async.wait_group {n};", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _cp_async_commit_wait_all(*, loc=None, ip=None):
+    _asm(None, [], "cp.async.commit_group;\n\tcp.async.wait_group 0;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _ld_global_v4(gaddr: Int64, *, loc=None, ip=None):
+    out = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [gaddr.ir_value()],
+        "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,l",
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4))
+
+
+@dsl_user_op
+def _st_global_b16(gaddr: Int64, v: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [gaddr.ir_value(), v.ir_value()],
+        "{\n\t.reg .b16 lo;\n\tcvt.u16.u32 lo, $1;\n\tst.global.b16 [$0], lo;\n\t}",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _st_global_v4(gaddr: Int64, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [gaddr.ir_value(), a.ir_value(), b.ir_value(), c.ir_value(), d.ir_value()],
+        "st.global.v4.b32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _st_global_b64(gaddr: Int64, v: Int64, *, loc=None, ip=None):
+    _asm(None, [gaddr.ir_value(), v.ir_value()], "st.global.b64 [$0], $1;", "l,l", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _st_shared_v4(saddr: Int32, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [saddr.ir_value(), a.ir_value(), b.ir_value(), c.ir_value(), d.ir_value()],
+        "st.shared.v4.b32 [$0], {$1, $2, $3, $4};",
+        "r,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _st_shared_b16(saddr: Int32, v: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [saddr.ir_value(), v.ir_value()],
+        "{\n\t.reg .b16 lo;\n\tcvt.u16.u32 lo, $1;\n\tst.shared.b16 [$0], lo;\n\t}",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _ld_shared_b32(saddr: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(T.i32(), [saddr.ir_value()], "ld.shared.b32 $0, [$1];", "=r,r", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def _ld_shared_v4(saddr: Int32, *, loc=None, ip=None):
+    out = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [saddr.ir_value()],
+        "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4))
+
+
+@dsl_user_op
+def _st_cluster_v4(
+    saddr: Int32, rank: Int32, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None
+):
+    """Store 16 bytes to the same shared-memory offset of cluster CTA
+    ``rank``."""
+    _asm(
+        None,
+        [saddr.ir_value(), rank.ir_value(), a.ir_value(), b.ir_value(), c.ir_value(), d.ir_value()],
+        "{\n\t.reg .b32 ra;\n\tmapa.shared::cluster.u32 ra, $0, $1;\n\t"
+        "st.shared::cluster.v4.b32 [ra], {$2, $3, $4, $5};\n\t}",
+        "r,r,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _st_async_cluster_v4(
+    saddr: Int32,
+    mbar: Int32,
+    rank: Int32,
+    a: Int32,
+    b: Int32,
+    c: Int32,
+    d: Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    """``st.async`` of 16 bytes to the same shared-memory offset of cluster
+    CTA ``rank``, completing transaction bytes on that CTA's mbarrier at the
+    same offset as the local ``mbar``."""
+    _asm(
+        None,
+        [
+            saddr.ir_value(),
+            mbar.ir_value(),
+            rank.ir_value(),
+            a.ir_value(),
+            b.ir_value(),
+            c.ir_value(),
+            d.ir_value(),
+        ],
+        "{\n\t.reg .b32 ra, rb;\n\t"
+        "mapa.shared::cluster.u32 ra, $0, $2;\n\t"
+        "mapa.shared::cluster.u32 rb, $1, $2;\n\t"
+        "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 "
+        "[ra], {$3, $4, $5, $6}, [rb];\n\t}",
+        "r,r,r,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _mbar_wait_cluster(mbar: Int32, phase: Int32, *, loc=None, ip=None):
+    """Wait for phase ``phase`` of the local mbarrier with acquire semantics
+    at cluster scope (its transaction bytes come from peer CTAs)."""
+    _asm(
+        None,
+        [mbar.ir_value(), phase.ir_value()],
+        "{\n\t.reg .pred p;\n\tINDEXER_Q_MBAR_WAIT:\n\t"
+        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 "
+        "p, [$0], $1;\n\t@!p bra INDEXER_Q_MBAR_WAIT;\n\t}",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _ldsm_x4(saddr: Int32, *, loc=None, ip=None):
+    out = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [saddr.ir_value()],
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4))
+
+
+@dsl_user_op
+def _mma_bf16(
+    a0: Int32,
+    a1: Int32,
+    a2: Int32,
+    a3: Int32,
+    b0: Int32,
+    b1: Int32,
+    c0: Float32,
+    c1: Float32,
+    c2: Float32,
+    c3: Float32,
+    *,
+    loc=None,
+    ip=None,
+):
+    out = _asm(
+        llvm.StructType.get_literal([T.f32()] * 4),
+        [
+            a0.ir_value(),
+            a1.ir_value(),
+            a2.ir_value(),
+            a3.ir_value(),
+            b0.ir_value(),
+            b1.ir_value(),
+            c0.ir_value(),
+            c1.ir_value(),
+            c2.ir_value(),
+            c3.ir_value(),
+        ],
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, "
+        "{$10, $11, $12, $13};",
+        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+        side_effects=False,
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4))
+
+
+@dsl_user_op
+def _bf16x2_lo(u: Int32, *, loc=None, ip=None) -> Float32:
+    v = _asm(
+        T.i32(), [u.ir_value()], "shl.b32 $0, $1, 16;", "=r,r", side_effects=False, loc=loc, ip=ip
+    )
+    return Float32(llvm.bitcast(T.f32(), v, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def _bf16x2_hi(u: Int32, *, loc=None, ip=None) -> Float32:
+    v = _asm(
+        T.i32(),
+        [u.ir_value()],
+        "and.b32 $0, $1, 0xffff0000;",
+        "=r,r",
+        side_effects=False,
+        loc=loc,
+        ip=ip,
+    )
+    return Float32(llvm.bitcast(T.f32(), v, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def _pack_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Int32:
+    """Round-to-nearest-even of two FP32 values into one BF16x2 word
+    (``lo`` in the low half)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [hi.ir_value(), lo.ir_value()],
+            "cvt.rn.bf16x2.f32 $0, $1, $2;",
+            "=r,f,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _bf16_bits(x: Float32, *, loc=None, ip=None) -> Int32:
+    """Round-to-nearest-even of one FP32 value to BF16 (in the low half)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [x.ir_value()],
+            "{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tcvt.u32.u16 $0, h;\n\t}",
+            "=r,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _cluster_sync(*, loc=None, ip=None):
+    _asm(
+        None,
+        [],
+        "barrier.cluster.arrive.release.aligned;\n\tbarrier.cluster.wait.acquire.aligned;",
+        "",
+        loc=loc,
+        ip=ip,
+    )
+
+
+# --------------------------------------------------------------------------
+# Native RMSNorm reproduction (one warp per row).
+# --------------------------------------------------------------------------
+
+
+@cute.jit
+def _norm_row_to_smem(
+    src: Int64,
+    gw: Int64,
+    sgw: Int32,
+    dst_smem: Int32,
+    lane: Int32,
+    eps: Float32,
+    NB: cutlass.Constexpr[int],
+    H: cutlass.Constexpr[int],
+    GAIN_SMEM: cutlass.Constexpr[bool],
+):
+    """Normalize one BF16 row of ``H = 256 * NB`` columns at global address
+    ``src`` with gain ``gw`` (global) or ``sgw`` (shared memory, with
+    ``GAIN_SMEM``) and store the BF16 result to shared memory at
+    ``dst_smem`` (contiguous row), as the native FlashInfer kernel does:
+    lane ``t`` owns columns ``256 b + 8 t + v``, sequential fold over ``b``
+    then ``v``, butterfly over the lanes, ``rsqrt(sum / H + eps)``
+    (fast-math), ``bf16((x * rstd) * (w + 0))``. All loads (row and gain)
+    are issued before the arithmetic."""
+    xu = []
+    for b in cutlass.range_constexpr(NB):
+        xu.append(_ld_global_v4(src + Int64((256 * b + 8 * lane) * 2)))
+    wu = []
+    for b in cutlass.range_constexpr(NB):
+        if cutlass.const_expr(GAIN_SMEM):
+            wu.append(_ld_shared_v4(sgw + Int32((256 * b + 8 * lane) * 2)))
+        else:
+            wu.append(_ld_global_v4(gw + Int64((256 * b + 8 * lane) * 2)))
+    xs = cute.make_rmem_tensor((NB, 8), Float32)
+    for b in cutlass.range_constexpr(NB):
+        u = xu[b]
+        for i in cutlass.range_constexpr(4):
+            xs[b, 2 * i] = _bf16x2_lo(u[i])
+            xs[b, 2 * i + 1] = _bf16x2_hi(u[i])
+    s = Float32(0.0)
+    for b in cutlass.range_constexpr(NB):
+        for v in cutlass.range_constexpr(8):
+            s = s + xs[b, v] * xs[b, v]
+    for i in cutlass.range_constexpr(5):
+        s = s + cute.arch.shuffle_sync_bfly(s, offset=1 << i)
+    mean = s / Float32(H)
+    rstd = cute.math.rsqrt(mean + eps, fastmath=True)
+    for b in cutlass.range_constexpr(NB):
+        w = wu[b]
+        packed = cute.make_rmem_tensor(4, Int32)
+        for i in cutlass.range_constexpr(4):
+            y0 = (xs[b, 2 * i] * rstd) * (_bf16x2_lo(w[i]) + Float32(0.0))
+            y1 = (xs[b, 2 * i + 1] * rstd) * (_bf16x2_hi(w[i]) + Float32(0.0))
+            packed[i] = _pack_bf16x2(y0, y1)
+        _st_shared_v4(
+            dst_smem + Int32((256 * b + 8 * lane) * 2), packed[0], packed[1], packed[2], packed[3]
+        )
+
+
+@cute.jit
+def _stamp(
+    enabled: cutlass.Constexpr[bool], a_ts: Int64, bid: Int32, tid: Int32, k: cutlass.Constexpr[int]
+):
+    """Tests only: ``%globaltimer`` of CTA ``bid`` at point ``k``."""
+    if cutlass.const_expr(enabled):
+        if tid == 0:
+            _st_global_b64(a_ts + Int64((bid * 8 + k) * 8), cute.arch.globaltimer())
+
+
+@cute.jit
+def _mma_chain(
+    a_row: Int32, b_row: Int32, STEPS: cutlass.Constexpr[int], DIST: cutlass.Constexpr[int]
+):
+    """``sum_s A_s B_s`` over ``STEPS`` ascending 16-wide K steps of one
+    m16n8 tile (A via ldmatrix.x4 at ``a_row + 32 s``, B via two 32-bit
+    loads at ``b_row + 32 s`` and ``+ 16``), with the fragments of step
+    ``s + DIST`` requested before the MMA of step ``s``."""
+    frags = []
+    for s in cutlass.range_constexpr(DIST):
+        frags.append(
+            (
+                _ldsm_x4(a_row + s * 32),
+                _ld_shared_b32(b_row + s * 32),
+                _ld_shared_b32(b_row + s * 32 + 16),
+            )
+        )
+    c0 = Float32(0.0)
+    c1 = Float32(0.0)
+    c2 = Float32(0.0)
+    c3 = Float32(0.0)
+    for s in cutlass.range_constexpr(STEPS):
+        if cutlass.const_expr(s + DIST < STEPS):
+            n = s + DIST
+            frags.append(
+                (
+                    _ldsm_x4(a_row + n * 32),
+                    _ld_shared_b32(b_row + n * 32),
+                    _ld_shared_b32(b_row + n * 32 + 16),
+                )
+            )
+        a, b0, b1 = frags[s]
+        c0, c1, c2, c3 = _mma_bf16(a[0], a[1], a[2], a[3], b0, b1, c0, c1, c2, c3)
+    return c0, c1, c2, c3
+
+
+@cute.jit
+def _copy_weights(
+    sW: Int32,
+    sUK: Int32,
+    a_wqb: Int64,
+    a_wuk: Int64,
+    col0: Int32,
+    uk0: Int32,
+    tid: Int32,
+    NT: cutlass.Constexpr[int],
+):
+    """``cp.async`` of a CTA's q_b weight tile (32 x 2048) and W_UK slice
+    (64 x 192) into shared memory (padded row strides)."""
+    for i in cutlass.range_constexpr(32 * _Q_LORA // 8 // NT):
+        c = tid + NT * i
+        r = c // 256
+        k16 = c % 256
+        _cp_async16(
+            sW + r * _W_STRIDE + k16 * 16, a_wqb + Int64(((col0 + r) * _Q_LORA + k16 * 8) * 2)
+        )
+    for i in cutlass.range_constexpr(_ABS_COLS * _NOPE // 8 // NT):
+        c = tid + NT * i
+        r = c // 24
+        k16 = c % 24
+        _cp_async16(sUK + r * _UK_STRIDE + k16 * 16, a_wuk + Int64((uk0 + r * _NOPE + k16 * 8) * 2))
+
+
+class SharedQueryClusterKernel:
+    """The CuTe-DSL shared_indexer_absorbed_query kernel (see the module docstring)."""
+
+    # Fragment prefetch distance of the MMA chains (steps).
+    DIST = 4
+    # Threads per CTA (one warp per normalized row).
+    NT = 256
+
+    def __init__(
+        self,
+        use_pdl: bool,
+        early_trigger: bool = True,
+        debug_qn: bool = False,
+        debug_ts: bool = False,
+        weights_before_wait: bool = True,
+        push_async: bool = True,
+    ):
+        self.use_pdl = use_pdl
+        self.early_trigger = early_trigger
+        self.debug_qn = debug_qn
+        self.debug_ts = debug_ts
+        self.weights_before_wait = weights_before_wait
+        # True: q_nope tiles are pushed with st.async and every CTA waits on
+        # its own mbarrier for the bytes it receives; False: plain
+        # st.shared::cluster stores and one cluster barrier.
+        self.push_async = push_async
+
+    @cute.jit
+    def __call__(
+        self,
+        kva: cute.Tensor,
+        gq: cute.Tensor,
+        gkv: cute.Tensor,
+        wqb: cute.Tensor,
+        wuk: cute.Tensor,
+        q_out: cute.Tensor,
+        latent: cute.Tensor,
+        fq: cute.Tensor,
+        qn_dbg: cute.Tensor,
+        ts_dbg: cute.Tensor,
+        rows: Int32,
+        eps_q: Float32,
+        eps_kv: Float32,
+        stream,
+    ):
+        self.kernel(
+            kva, gq, gkv, wqb, wuk, q_out, latent, fq, qn_dbg, ts_dbg, rows, eps_q, eps_kv
+        ).launch(
+            grid=[(_H + 1) * 8, 1, 1],
+            block=[self.NT, 1, 1],
+            cluster=[8, 1, 1],
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        kva: cute.Tensor,
+        gq: cute.Tensor,
+        gkv: cute.Tensor,
+        wqb: cute.Tensor,
+        wuk: cute.Tensor,
+        q_out: cute.Tensor,
+        latent: cute.Tensor,
+        fq: cute.Tensor,
+        qn_dbg: cute.Tensor,
+        ts_dbg: cute.Tensor,
+        rows: Int32,
+        eps_q: Float32,
+        eps_kv: Float32,
+    ):
+        """Tests: ``debug_ts`` writes ``%globaltimer`` stamps per CTA into
+        ``ts_dbg [72, 8]`` (int64): 0 entry, 1 after the grid-dependency
+        wait, 2 normalized rows and weights in shared memory, 3 q_b tile
+        done, 4 after the q exchange, 5 absorbed query stored, 6 exit, 7
+        warp 0's normalized row in shared memory (the kv cluster stamps 0,
+        1, 2 and 6)."""
+        tid, _, _ = cute.arch.thread_idx()
+        bid, _, _ = cute.arch.block_idx()
+        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        lane = cute.arch.lane_idx()
+        head = bid // 8
+        j = bid % 8
+
+        smem = SmemAllocator()
+        sW = _saddr(smem.allocate(32 * _W_STRIDE, byte_alignment=128))
+        sUK = _saddr(smem.allocate(_ABS_COLS * _UK_STRIDE, byte_alignment=128))
+        sQn = _saddr(smem.allocate(_MAX_ROWS * _W_STRIDE, byte_alignment=128))
+        sStage = _saddr(smem.allocate(_MAX_ROWS * _TILE_COLS * 2, byte_alignment=128))
+        sNope = _saddr(smem.allocate(_MAX_ROWS * _UK_STRIDE, byte_alignment=128))
+        sGq = _saddr(smem.allocate(_Q_LORA * 2, byte_alignment=128))
+        mbar = smem.allocate_array(Int64, num_elems=1, byte_alignment=8)
+        sMbar = _saddr(mbar)
+
+        a_kva = _gaddr(kva)
+        a_gq = _gaddr(gq)
+        a_gkv = _gaddr(gkv)
+        a_wqb = _gaddr(wqb)
+        a_wuk = _gaddr(wuk)
+        a_q = _gaddr(q_out)
+        a_lat = _gaddr(latent)
+        a_fq = _gaddr(fq)
+        a_qn = _gaddr(qn_dbg)
+        a_ts = _gaddr(ts_dbg)
+        _stamp(self.debug_ts, a_ts, bid, tid, 0)
+
+        if head < _H:
+            col0 = head * _HEAD_DIM + j * _TILE_COLS
+            uk0 = (head * _KV_LORA + j * _ABS_COLS) * _NOPE
+            # 1. Weights -> shared memory: they do not depend on earlier
+            #    grids, so they are requested before the dependency wait.
+            #    The q_a norm gain (4 KB, one 16-byte chunk per thread) is
+            #    its own cp.async group and is complete and visible to all
+            #    warps before the wait.
+            _cp_async16(sGq + tid * 16, a_gq + Int64(tid * 16))
+            _cp_async_commit()
+            if cutlass.const_expr(self.weights_before_wait):
+                _copy_weights(sW, sUK, a_wqb, a_wuk, col0, uk0, tid, self.NT)
+                _cp_async_commit()
+                _cp_async_wait_group(1)
+            else:
+                _cp_async_wait_group(0)
+            cute.arch.sync_threads()
+            if cutlass.const_expr(self.push_async):
+                # Receive buffer barrier: expects the 6 no-RoPE tiles (R_max
+                # rows x 32 columns of BF16 each).
+                if tid == 0:
+                    cute.arch.mbarrier_init(mbar, 1)
+                    cute.arch.mbarrier_arrive_and_expect_tx(
+                        mbar, (_NOPE // _TILE_COLS) * _MAX_ROWS * _TILE_COLS * 2
+                    )
+                cute.arch.mbarrier_init_fence()
+            # Every CTA of the cluster has started (and initialized its
+            # barrier) before any CTA writes to a peer's shared memory.
+            cute.arch.cluster_arrive_relaxed()
+            cute.arch.cluster_wait()
+            if cutlass.const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_wait()
+                if cutlass.const_expr(self.early_trigger):
+                    cute.arch.griddepcontrol_launch_dependents()
+            if cutlass.const_expr(not self.weights_before_wait):
+                _copy_weights(sW, sUK, a_wqb, a_wuk, col0, uk0, tid, self.NT)
+            _stamp(self.debug_ts, a_ts, bid, tid, 1)
+
+            # 2. Normalized q_a rows -> shared memory, one warp per row;
+            #    rows >= R are zero.
+            r = warp
+            if r < rows:
+                _norm_row_to_smem(
+                    a_kva + Int64(r * _KVA_COLS * 2),
+                    a_gq,
+                    sGq,
+                    sQn + r * _W_STRIDE,
+                    lane,
+                    eps_q,
+                    8,
+                    _Q_LORA,
+                    True,
+                )
+            else:
+                for b in cutlass.range_constexpr(8):
+                    _st_shared_v4(
+                        sQn + r * _W_STRIDE + (256 * b + 8 * lane) * 2,
+                        Int32(0),
+                        Int32(0),
+                        Int32(0),
+                        Int32(0),
+                    )
+            _stamp(self.debug_ts, a_ts, bid, tid, 7)
+            _cp_async_commit_wait_all()
+            cute.arch.sync_threads()
+            _stamp(self.debug_ts, a_ts, bid, tid, 2)
+            if cutlass.const_expr(self.debug_qn):
+                # Tests: the normalized rows of CTA 0 -> qn_dbg [R, 2048].
+                if bid == 0:
+                    for i in cutlass.range_constexpr(_MAX_ROWS * 256 // self.NT):
+                        c = tid + self.NT * i
+                        rr = c // 256
+                        k16 = c % 256
+                        if rr < rows:
+                            u = _ld_shared_v4(sQn + rr * _W_STRIDE + k16 * 16)
+                            _st_global_v4(
+                                a_qn + Int64(rr * _Q_LORA * 2 + k16 * 16), u[0], u[1], u[2], u[3]
+                            )
+
+            # 3. q_b tile: warps 0 and 1, 16 weight rows each, K ascending.
+            g = lane // 4
+            t4 = lane % 4
+            if warp < 2:
+                acc = _mma_chain(
+                    sW + (16 * warp + lane % 16) * _W_STRIDE + (lane // 16) * 16,
+                    sQn + g * _W_STRIDE + t4 * 4,
+                    _Q_LORA // 16,
+                    self.DIST,
+                )
+                _stamp(self.debug_ts, a_ts, bid, tid, 3)
+                # C[m][n]: m = weight row (16 warp + g, +8), n = token
+                # (2 t4, 2 t4 + 1); staged as bf16 [token, 32 columns].
+                for hi in cutlass.range_constexpr(2):
+                    m = 16 * warp + g + 8 * hi
+                    for e in cutlass.range_constexpr(2):
+                        n = 2 * t4 + e
+                        _st_shared_b16(
+                            sStage + (n * _TILE_COLS + m) * 2, _bf16_bits(acc[2 * hi + e])
+                        )
+            cute.arch.sync_threads()
+
+            # 4. No-RoPE tiles -> every cluster peer's q_nope buffer
+            #    (32 chunks of 16 B to 8 peers: one per thread), the cluster
+            #    barrier, then the q tile to global memory (after the
+            #    barrier, so that its release does not wait for them).
+            if j < _NOPE // _TILE_COLS:
+                peer = tid // 32
+                ch = tid % 32
+                n = ch // 4
+                part = ch % 4
+                u = _ld_shared_v4(sStage + n * _TILE_COLS * 2 + part * 16)
+                dst = sNope + n * _UK_STRIDE + (j * _TILE_COLS) * 2 + part * 16
+                if cutlass.const_expr(self.push_async):
+                    _st_async_cluster_v4(dst, sMbar, peer, u[0], u[1], u[2], u[3])
+                else:
+                    _st_cluster_v4(dst, peer, u[0], u[1], u[2], u[3])
+            if cutlass.const_expr(self.push_async):
+                _mbar_wait_cluster(sMbar, Int32(0))
+            else:
+                _cluster_sync()
+            _stamp(self.debug_ts, a_ts, bid, tid, 4)
+            if warp >= 4:
+                if tid < 128 + _MAX_ROWS * 4:
+                    n = (tid - 128) // 4
+                    part = tid % 4
+                    if n < rows:
+                        u = _ld_shared_v4(sStage + n * _TILE_COLS * 2 + part * 16)
+                        _st_global_v4(
+                            a_q + Int64((n * _H * _HEAD_DIM + col0) * 2 + part * 16),
+                            u[0],
+                            u[1],
+                            u[2],
+                            u[3],
+                        )
+
+            # 5. Absorption: 4 warps x 16 latent columns, d ascending.
+            if warp < 4:
+                acc2 = _mma_chain(
+                    sUK + (16 * warp + lane % 16) * _UK_STRIDE + (lane // 16) * 16,
+                    sNope + g * _UK_STRIDE + t4 * 4,
+                    _NOPE // 16,
+                    self.DIST,
+                )
+                for hi in cutlass.range_constexpr(2):
+                    cl = j * _ABS_COLS + 16 * warp + g + 8 * hi
+                    for e in cutlass.range_constexpr(2):
+                        n = 2 * t4 + e
+                        if n < rows:
+                            _st_global_b16(
+                                a_fq + Int64(((n * _H + head) * (_KV_LORA + _ROPE) + cl) * 2),
+                                _bf16_bits(acc2[2 * hi + e]),
+                            )
+            _stamp(self.debug_ts, a_ts, bid, tid, 5)
+        else:
+            # kv_a RMSNorm -> latent[:, :512]; k_pe -> latent[:, 512:].
+            if cutlass.const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_wait()
+                if cutlass.const_expr(self.early_trigger):
+                    cute.arch.griddepcontrol_launch_dependents()
+            _stamp(self.debug_ts, a_ts, bid, tid, 1)
+            r = j
+            if r < rows:
+                if warp == 0:
+                    _norm_row_to_smem(
+                        a_kva + Int64((r * _KVA_COLS + _Q_LORA) * 2),
+                        a_gkv,
+                        Int32(0),
+                        sQn,
+                        lane,
+                        eps_kv,
+                        2,
+                        _KV_LORA,
+                        False,
+                    )
+                cute.arch.sync_threads()
+                _stamp(self.debug_ts, a_ts, bid, tid, 2)
+                if tid < 64:
+                    u = _ld_shared_v4(sQn + tid * 16)
+                    _st_global_v4(
+                        a_lat + Int64((r * (_KV_LORA + _ROPE) * 2) + tid * 16),
+                        u[0],
+                        u[1],
+                        u[2],
+                        u[3],
+                    )
+                elif tid < 72:
+                    k = tid - 64
+                    u = _ld_global_v4(
+                        a_kva + Int64((r * _KVA_COLS + _Q_LORA + _KV_LORA) * 2 + k * 16)
+                    )
+                    _st_global_v4(
+                        a_lat + Int64((r * (_KV_LORA + _ROPE) + _KV_LORA) * 2 + k * 16),
+                        u[0],
+                        u[1],
+                        u[2],
+                        u[3],
+                    )
+        _stamp(self.debug_ts, a_ts, bid, tid, 6)
+
+
+_COMPILED = {}
+
+
+def _fake_tensors():
+    sym_r = cute.sym_int()
+    bf16 = cutlass.BFloat16
+    mk = cute.runtime.make_fake_compact_tensor
+    return [
+        mk(bf16, (sym_r, _KVA_COLS), stride_order=(1, 0)),
+        mk(bf16, (_Q_LORA,), stride_order=(0,)),
+        mk(bf16, (_KV_LORA,), stride_order=(0,)),
+        mk(bf16, (_H * _HEAD_DIM, _Q_LORA), stride_order=(1, 0)),
+        mk(bf16, (_H, _KV_LORA, _NOPE), stride_order=(2, 1, 0)),
+        mk(bf16, (sym_r, _H * _HEAD_DIM), stride_order=(1, 0)),
+        mk(bf16, (sym_r, _KV_LORA + _ROPE), stride_order=(1, 0)),
+        mk(bf16, (sym_r, _H, _KV_LORA + _ROPE), stride_order=(2, 1, 0)),
+        mk(bf16, (cute.sym_int(), _Q_LORA), stride_order=(1, 0)),
+        mk(cutlass.Int64, (cute.sym_int(), 8), stride_order=(1, 0)),
+    ]
+
+
+def compiled_kernel(
+    use_pdl: bool,
+    early_trigger: bool = True,
+    debug_qn: bool = False,
+    debug_ts: bool = False,
+    weights_before_wait: bool = True,
+    push_async: bool = True,
+):
+    """Compile (once per process and flag set) and return the kernel. The
+    compile allocates, so the first call must not be inside a CUDA-graph
+    capture (the runner's eager warm-up forwards come first)."""
+    key = (use_pdl, early_trigger, debug_qn, debug_ts, weights_before_wait, push_async)
+    if key not in _COMPILED:
+        kernel = SharedQueryClusterKernel(
+            use_pdl, early_trigger, debug_qn, debug_ts, weights_before_wait, push_async
+        )
+        fakes = _fake_tensors()
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        _COMPILED[key] = cute.compile(
+            kernel,
+            *fakes,
+            Int32(1),
+            Float32(1e-5),
+            Float32(1e-5),
+            stream,
+            options="--enable-tvm-ffi",
+        )
+    return _COMPILED[key]
+
+
+def launch(
+    mla,
+    kva: torch.Tensor,
+    use_pdl: bool,
+    early_trigger: bool = True,
+    qn_dbg: torch.Tensor = None,
+    ts_dbg: torch.Tensor = None,
+    weights_before_wait: bool = True,
+    push_async: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Run the cluster kernel on a native kv_a output ``[R <= 8, 2624]``;
+    with ``qn_dbg [R, 2048]`` the normalized q_a rows are also written,
+    with ``ts_dbg [72, 8]`` (int64) the per-CTA timer stamps."""
+    rows = kva.shape[0]
+    assert 1 <= rows <= _MAX_ROWS and kva.is_contiguous()
+    q = kva.new_empty(rows, _H * _HEAD_DIM)
+    latent = kva.new_empty(rows, _KV_LORA + _ROPE)
+    fused_q = kva.new_empty(rows, _H, _KV_LORA + _ROPE)
+    debug_qn = qn_dbg is not None
+    debug_ts = ts_dbg is not None
+    if qn_dbg is None:
+        qn_dbg = kva.new_empty(1, _Q_LORA)
+    if ts_dbg is None:
+        ts_dbg = torch.empty(1, 8, dtype=torch.int64, device=kva.device)
+    compiled_kernel(use_pdl, early_trigger, debug_qn, debug_ts, weights_before_wait, push_async)(
+        kva,
+        mla.q_a_layernorm.weight,
+        mla.kv_a_layernorm.weight,
+        mla.q_b_proj.weight,
+        mla.k_b_proj_trans,
+        q,
+        latent,
+        fused_q,
+        qn_dbg,
+        ts_dbg,
+        rows,
+        mla.q_a_layernorm.variance_epsilon,
+        mla.kv_a_layernorm.variance_epsilon,
+    )
+    return q, latent, fused_q
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/__init__.py
new file mode 100644
index 000000000000..9b18ef89cfbd
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/__init__.py
@@ -0,0 +1,46 @@
+# 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.
+"""attention_output_tail: local attention output tail (the attention_output_tail contract).
+
+Replaces the GEN-phase W_UV BMM and the local ``o_proj`` GEMM of an MLA layer
+with attention_output_tail-owned kernels (a Triton W_UV kernel and a CuTe DSL o_proj kernel that
+streams its weights around its PDL wait) that keep the native BF16 rounding
+of the intermediate and emit the ordinary rank-local BF16 partial
+``[R, hidden]``. The native all-reduce consumes that partial unchanged; attention_output_tail
+hides no communication.
+
+Registered entry points (``dispatch.get_impl("attention_output_tail", name)``):
+
+* ``make_attn_tail(mla) -> AttnTail | None``: per-layer adapter used by the
+  MLA call sites (see ``tail.py``).
+* ``attn_tail(latent, w_uv, w_o, num_rows, out_rows=None)``: the functional
+  tail (see ``kernels.py``).
+"""
+
+from ... import dispatch
+from .kernels import ROW_TILE, attn_tail, check_tail_shapes, check_weights
+from .tail import AttnTail, make_attn_tail
+
+dispatch.register("attention_output_tail", "make_attn_tail", make_attn_tail)
+dispatch.register("attention_output_tail", "attn_tail", attn_tail)
+
+__all__ = [
+    "ROW_TILE",
+    "AttnTail",
+    "attn_tail",
+    "check_tail_shapes",
+    "check_weights",
+    "make_attn_tail",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py
new file mode 100644
index 000000000000..5102e384c037
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py
@@ -0,0 +1,585 @@
+# 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.
+"""Single-launch CuTe DSL attention_output_tail: W_UV and o_proj in one kernel (SM100/SM103).
+
+One CTA per SM; 8 compute warps plus 1 producer warp.
+
+* Producer warp: streams the CTA's 40-41 contiguous ``w_o`` rows into a
+  row-padded shared-memory slab with one 1-D ``cp.async.bulk`` per 4 KiB row,
+  each completing on its own mbarrier, keeping at most ``window`` rows in
+  flight. Bounding the in-flight weight traffic keeps the memory system's
+  queues short, so the latency-critical W_UV loads, ``v`` stores, flag
+  updates and ``v`` loads of the other phases are not stuck behind tens of
+  megabytes of queued copies; a window of a few rows per SM is still enough
+  to keep HBM saturated.
+* Phase 1, W_UV: CTAs ``0 .. K/16 - 1`` each own 16 columns of ``v`` (one
+  head, 16 of its V rows of ``w_uv``). The compute warps split the 512-wide
+  latent dimension, run ``mma.sync.m16n8k16`` (BF16 in, FP32 accumulate), add
+  their partials in fixed warp order and store ``v`` rounded to BF16 -- the
+  native intermediate rounding -- then publish the slice by incrementing a
+  per-layer ``done`` counter (``fence.acq_rel.gpu`` + ``red``).
+* Phase 2, o_proj: every CTA waits (acquire) until ``done`` reaches the target
+  of its invocation, loads ``v`` from L2, and runs ``mma.sync`` over its slab
+  rows as they land, with a fixed-order cross-warp reduction and one BF16
+  rounding of the partial. At exit every CTA adds 1 to ``finished``.
+
+With ``pdl_entry`` the kernel is a PDL dependent of the latent's producer:
+before ``griddepcontrol.wait`` it only streams immutable weights into its own
+shared memory and loads ``w_uv``; the latent, ``v``, the counters and ``out``
+are touched only after the wait, which guarantees the producer grid has
+completed and its writes are visible (the same contract as stream order).
+
+Invocation ``e`` of a state buffer is ``finished_at_entry // grid``: all CTAs
+of earlier invocations have exited (stream order) and fewer than ``grid`` CTAs
+of this one, so the value lies in ``[e * grid, (e + 1) * grid)``. Its slices
+are complete when ``done >= (e + 1) * K / 16`` (with ``pdl_entry``,
+``finished`` is read after the wait). Both counters are monotonic 64-bit
+integers and never reset. The grid is one CTA per SM and the kernel never
+triggers its dependents early, so a waiting CTA waits only for CTAs that
+either run or will run once concurrent, independent work leaves their SM. One
+state buffer must not be used by two concurrently running invocations (one per
+layer and one stream, as in the MLA adapter).
+"""
+
+import functools
+from typing import Optional
+
+import torch
+
+from .cute_oproj import _TILE_M, _WARPS, _sm_count, plan
+
+_SLICE = 16  # v columns per W_UV CTA
+_COMPUTE_THREADS = _WARPS * 32
+_THREADS = _COMPUTE_THREADS + 32  # plus one producer warp
+_SPIN_LIMIT = 1 << 26
+DEFAULT_WINDOW = 8
+"""Rows (4 KiB each) of ``w_o`` a CTA keeps in flight."""
+
+
+def check_fused_shapes(heads: int, v_dim: int, c_dim: int, n_dim: int, grid: int) -> Optional[str]:
+    """Return why the fused kernel cannot take these shapes, or ``None``."""
+    k_dim = heads * v_dim
+    if v_dim % _SLICE or c_dim % (_WARPS * 16):
+        return "W_UV shape not a multiple of the fused tiles"
+    if k_dim // _SLICE > grid:
+        return "more W_UV slices than CTAs"
+    if k_dim % (_WARPS * 16):
+        return "K is not a multiple of 128"
+    p = plan(n_dim, k_dim, grid)
+    if p["smem_bytes"] + _WARPS * 2 * 128 * 4 > 200 * 1024:
+        return "weight slab does not fit in shared memory"
+    return None
+
+
+@functools.lru_cache(maxsize=None)
+def _build_fused(
+    heads: int,
+    v_dim: int,
+    c_dim: int,
+    n_dim: int,
+    grid: int,
+    nt: int,
+    traced: bool,
+    window: int,
+    pdl: bool,
+):
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import BFloat16, Float32, Int32, Int64
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cute import runtime as cute_runtime
+    from cutlass.cutlass_dsl import T, dsl_user_op
+    from cutlass.utils.smem_allocator import SmemAllocator
+
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+
+    k_dim = heads * v_dim
+    geo = plan(n_dim, k_dim, grid)
+    rows_max = geo["rows_max"]
+    num_tiles = geo["num_tiles"]
+    pitch = geo["pitch"]
+    half_pitch = pitch // 2
+    k_per_warp = geo["k_per_warp"]
+    steps = k_per_warp // 16
+    row_bytes = k_dim * 2
+    red_stride = 128
+    num_slices = k_dim // _SLICE
+    slices_per_head = v_dim // _SLICE
+    c_per_warp = c_dim // _WARPS
+    c_steps = c_per_warp // 16
+    c_words = c_dim // 2
+    k_words = k_dim // 2
+    compute_bar = 1
+
+    def _asm(res, operands, text, constraints, loc, ip):
+        return llvm.inline_asm(
+            res, operands, text, constraints, has_side_effects=True, asm_dialect=0, loc=loc, ip=ip
+        )
+
+    @dsl_user_op
+    def bulk_g2s(dst_smem, src_gmem, num_bytes, mbar, *, loc=None, ip=None):
+        _asm(
+            None,
+            [
+                dst_smem.toint(loc=loc, ip=ip).ir_value(),
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+                mbar.toint(loc=loc, ip=ip).ir_value(),
+            ],
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            " [$0], [$1], $2, [$3];",
+            "r,l,r,r",
+            loc,
+            ip,
+        )
+
+    @dsl_user_op
+    def red_add_u64(addr, val, *, loc=None, ip=None):
+        _asm(
+            None,
+            [Int64(addr).ir_value(loc=loc, ip=ip), Int64(val).ir_value(loc=loc, ip=ip)],
+            "red.relaxed.gpu.global.add.u64 [$0], $1;",
+            "l,l",
+            loc,
+            ip,
+        )
+
+    @dsl_user_op
+    def ld_relaxed_u64(addr, *, loc=None, ip=None):
+        return Int64(
+            _asm(
+                T.i64(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.relaxed.gpu.global.u64 $0, [$1];",
+                "=l,l",
+                loc,
+                ip,
+            )
+        )
+
+    @dsl_user_op
+    def global_timer(*, loc=None, ip=None):
+        return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc, ip))
+
+    @dsl_user_op
+    def ld_acquire_u64(addr, *, loc=None, ip=None):
+        return Int64(
+            _asm(
+                T.i64(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.acquire.gpu.global.u64 $0, [$1];",
+                "=l,l",
+                loc,
+                ip,
+            )
+        )
+
+    @dsl_user_op
+    def ld_cg_b32(addr, *, loc=None, ip=None):
+        return Int32(
+            _asm(
+                T.i32(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.global.cg.b32 $0, [$1];",
+                "=r,l",
+                loc,
+                ip,
+            )
+        )
+
+    @dsl_user_op
+    def grid_dependency_wait(*, loc=None, ip=None):
+        _asm(None, [], "griddepcontrol.wait;", "", loc, ip)
+
+    @dsl_user_op
+    def fence_gpu(*, loc=None, ip=None):
+        _asm(None, [], "fence.acq_rel.gpu;", "", loc, ip)
+
+    @dsl_user_op
+    def sleep_ns(*, loc=None, ip=None):
+        _asm(None, [], "nanosleep.u32 32;", "", loc, ip)
+
+    @dsl_user_op
+    def trap(*, loc=None, ip=None):
+        _asm(None, [], "trap;", "", loc, ip)
+
+    @cute.jit
+    def compute_sync():
+        cute.arch.barrier(barrier_id=compute_bar, number_of_threads=_COMPUTE_THREADS)
+
+    class FusedTail:
+        @cute.kernel
+        def kernel(
+            self,
+            latent: cute.Tensor,
+            w_uv: cute.Tensor,
+            w: cute.Tensor,
+            v: cute.Tensor,
+            out: cute.Tensor,
+            state: cute.Tensor,
+            trace: cute.Tensor,
+        ):
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            warp = tidx // Int32(32)
+            lane = tidx % Int32(32)
+            grp = lane // Int32(4)
+            quad = lane % Int32(4)
+            t_entry = Int64(0)
+            if cutlass.const_expr(traced):
+                t_entry = global_timer()
+
+            num_rows = Int32(latent.shape[0])
+            last_row = num_rows - Int32(1)
+            out_rows = Int32(out.shape[0])
+            n0 = (bidx * Int32(n_dim)) // Int32(grid)
+            n1 = ((bidx + Int32(1)) * Int32(n_dim)) // Int32(grid)
+            nrows = n1 - n0
+
+            smem = SmemAllocator()
+            slab = smem.allocate_tensor(
+                BFloat16, cute.make_layout((rows_max * pitch,)), byte_alignment=128
+            )
+            red1 = smem.allocate_tensor(
+                Float32, cute.make_layout((_WARPS * nt * red_stride,)), byte_alignment=16
+            )
+            row_bar = smem.allocate_tensor(Int64, cute.make_layout((rows_max,)), byte_alignment=8)
+            slab32 = cute.recast_tensor(slab, Int32)
+            slab_f32 = cute.recast_tensor(slab, Float32)
+            lat32 = cute.recast_tensor(latent, Int32)
+            wuv32 = cute.recast_tensor(w_uv, Int32)
+            finished_addr = state.iterator.toint()
+            done_addr = finished_addr + Int64(8)
+            v_addr = v.iterator.toint()
+
+            # Thread 0 arms one mbarrier per row (4 KiB of expected bytes)
+            # and reads the invocation count.
+            finished = Int64(0)
+            if tidx == Int32(0):
+                if cutlass.const_expr(not pdl):
+                    finished = ld_relaxed_u64(finished_addr)
+                for i in cutlass.range_constexpr(rows_max):
+                    cute.arch.mbarrier_init(row_bar.iterator + i, 1)
+                for i in cutlass.range_constexpr(rows_max):
+                    if Int32(i) < nrows:
+                        cute.arch.mbarrier_arrive_and_expect_tx(
+                            row_bar.iterator + i, Int32(row_bytes)
+                        )
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+
+            t_issued = Int64(0)
+            t_published = Int64(0)
+            t_ready = Int64(0)
+            t_tile0 = Int64(0)
+            if warp == Int32(_WARPS):
+                # Producer: at most `window` rows in flight. The whole warp
+                # runs the loop so every mbarrier wait is warp-uniform; lane 0
+                # issues the copies.
+                i = Int32(0)
+                while i < nrows:
+                    if i >= Int32(window):
+                        cute.arch.mbarrier_wait(row_bar.iterator + (i - Int32(window)), 0)
+                    if lane == Int32(0):
+                        bulk_g2s(
+                            slab.iterator + i * Int32(pitch),
+                            w.iterator + (n0 + i) * Int32(k_dim),
+                            row_bytes,
+                            row_bar.iterator + i,
+                        )
+                    i = i + Int32(1)
+                if cutlass.const_expr(traced):
+                    t_issued = global_timer()
+            else:
+                # Phase 1: CTA b < num_slices computes 16 columns of v.
+                if bidx < Int32(num_slices):
+                    head = bidx // Int32(slices_per_head)
+                    col0 = (bidx % Int32(slices_per_head)) * Int32(_SLICE)
+                    wrow_lo = head * Int32(v_dim) + col0 + grp
+                    wrow_hi = wrow_lo + Int32(8)
+                    c_word0 = warp * Int32(c_per_warp // 2) + quad
+                    lat_word0 = head * Int32(c_words) + c_word0
+                    wa = cute.make_rmem_tensor((8, c_steps), BFloat16)
+                    wa32 = cute.recast_tensor(wa, Int32)
+                    la = cute.make_rmem_tensor((4, nt, c_steps), BFloat16)
+                    la32 = cute.recast_tensor(la, Int32)
+                    for s in cutlass.range_constexpr(c_steps):
+                        wa32[0, s] = wuv32[wrow_lo, c_word0 + Int32(8 * s)]
+                        wa32[1, s] = wuv32[wrow_hi, c_word0 + Int32(8 * s)]
+                        wa32[2, s] = wuv32[wrow_lo, c_word0 + Int32(8 * s + 4)]
+                        wa32[3, s] = wuv32[wrow_hi, c_word0 + Int32(8 * s + 4)]
+                    if cutlass.const_expr(pdl):
+                        grid_dependency_wait()
+                    for s in cutlass.range_constexpr(c_steps):
+                        for j in cutlass.range_constexpr(nt):
+                            r = Int32(8 * j) + grp
+                            rc = r
+                            if rc > last_row:
+                                rc = last_row
+                            lo = lat32[rc, lat_word0 + Int32(8 * s)]
+                            hi = lat32[rc, lat_word0 + Int32(8 * s + 4)]
+                            if r > last_row:
+                                lo = Int32(0)
+                                hi = Int32(0)
+                            la32[0, j, s] = lo
+                            la32[1, j, s] = hi
+                    acc1 = cute.make_rmem_tensor((4, nt), Float32)
+                    acc1.fill(0.0)
+                    for s in cutlass.range_constexpr(c_steps):
+                        for j in cutlass.range_constexpr(nt):
+                            acc1[None, j] = mma_sync(wa[None, s], la[None, j, s], acc1[None, j])
+                    for j in cutlass.range_constexpr(nt):
+                        for e in cutlass.range_constexpr(4):
+                            red1[
+                                (warp * Int32(nt) + Int32(j)) * Int32(red_stride)
+                                + lane * Int32(4)
+                                + Int32(e)
+                            ] = acc1[e, j]
+                    compute_sync()
+                    idx1 = tidx
+                    while idx1 < Int32(nt * red_stride):
+                        j_i = idx1 // Int32(red_stride)
+                        l_i = (idx1 // Int32(4)) % Int32(32)
+                        e_i = idx1 % Int32(4)
+                        vsum = Float32(0.0)
+                        for wi in cutlass.range_constexpr(_WARPS):
+                            vsum = vsum + red1[Int32(wi * nt * red_stride) + idx1]
+                        m = l_i // Int32(4) + Int32(8) * (e_i // Int32(2))
+                        r = j_i * Int32(8) + Int32(2) * (l_i % Int32(4)) + (e_i % Int32(2))
+                        if r < num_rows:
+                            v[r, head * Int32(v_dim) + col0 + m] = vsum.to(BFloat16)
+                        idx1 = idx1 + Int32(_COMPUTE_THREADS)
+                    compute_sync()
+                    if tidx == Int32(0):
+                        fence_gpu()
+                        red_add_u64(done_addr, Int64(1))
+                        if cutlass.const_expr(traced):
+                            t_published = global_timer()
+
+                else:
+                    if cutlass.const_expr(pdl):
+                        grid_dependency_wait()
+
+                # Phase 2: wait for all slices of this invocation.
+                if tidx == Int32(0):
+                    if cutlass.const_expr(pdl):
+                        # After griddepcontrol.wait: every earlier invocation
+                        # of this state has completed.
+                        finished = ld_relaxed_u64(finished_addr)
+                    target = (finished // Int64(grid) + Int64(1)) * Int64(num_slices)
+                    spins = Int32(0)
+                    seen = ld_acquire_u64(done_addr)
+                    while seen < target:
+                        sleep_ns()
+                        spins = spins + Int32(1)
+                        if spins > Int32(_SPIN_LIMIT):
+                            trap()
+                        seen = ld_acquire_u64(done_addr)
+                compute_sync()
+                if cutlass.const_expr(traced):
+                    t_ready = global_timer()
+
+                vb = cute.make_rmem_tensor((4, nt, steps), BFloat16)
+                vb32 = cute.recast_tensor(vb, Int32)
+                k_word0 = warp * Int32(k_per_warp // 2) + quad
+                for s in cutlass.range_constexpr(steps):
+                    for j in cutlass.range_constexpr(nt):
+                        r = Int32(8 * j) + grp
+                        rc = r
+                        if rc > last_row:
+                            rc = last_row
+                        word = rc * Int32(k_words) + k_word0 + Int32(8 * s)
+                        lo = ld_cg_b32(v_addr + Int64(word) * Int64(4))
+                        hi = ld_cg_b32(v_addr + Int64(word + Int32(4)) * Int64(4))
+                        if r > last_row:
+                            lo = Int32(0)
+                            hi = Int32(0)
+                        vb32[0, j, s] = lo
+                        vb32[1, j, s] = hi
+
+                acc = cute.make_rmem_tensor((4, nt, num_tiles), Float32)
+                acc.fill(0.0)
+                a2 = cute.make_rmem_tensor((8,), BFloat16)
+                a2_32 = cute.recast_tensor(a2, Int32)
+                last = nrows - Int32(1)
+                for t in cutlass.range_constexpr(num_tiles):
+                    # Warp-uniform waits on every row of the tile.
+                    for q in cutlass.range_constexpr(_TILE_M):
+                        wrow = Int32(t * _TILE_M + q)
+                        if wrow < nrows:
+                            cute.arch.mbarrier_wait(row_bar.iterator + wrow, 0)
+                    row_lo = Int32(t * _TILE_M) + grp
+                    row_hi = row_lo + Int32(8)
+                    if row_lo > last:
+                        row_lo = last
+                    if row_hi > last:
+                        row_hi = last
+                    if cutlass.const_expr(traced and t == 0):
+                        t_tile0 = global_timer()
+                    base_lo = row_lo * Int32(half_pitch) + k_word0
+                    base_hi = row_hi * Int32(half_pitch) + k_word0
+                    for s in cutlass.range_constexpr(steps):
+                        a2_32[0] = slab32[base_lo + Int32(8 * s)]
+                        a2_32[1] = slab32[base_hi + Int32(8 * s)]
+                        a2_32[2] = slab32[base_lo + Int32(8 * s + 4)]
+                        a2_32[3] = slab32[base_hi + Int32(8 * s + 4)]
+                        for j in cutlass.range_constexpr(nt):
+                            acc[None, j, t] = mma_sync(a2, vb[None, j, s], acc[None, j, t])
+
+                # Fixed-order cross-warp reduction through the idle slab.
+                compute_sync()
+                for t in cutlass.range_constexpr(num_tiles):
+                    for j in cutlass.range_constexpr(nt):
+                        base = (
+                            (warp * Int32(num_tiles) + Int32(t)) * Int32(nt) + Int32(j)
+                        ) * Int32(red_stride) + (lane * Int32(4))
+                        for e in cutlass.range_constexpr(4):
+                            slab_f32[base + Int32(e)] = acc[e, j, t]
+                compute_sync()
+
+                total = num_tiles * nt * red_stride
+                out_limit = out_rows
+                if out_limit > Int32(8 * nt):
+                    out_limit = Int32(8 * nt)
+                idx = tidx
+                while idx < Int32(total):
+                    t_i = idx // Int32(nt * red_stride)
+                    j_i = (idx // Int32(red_stride)) % Int32(nt)
+                    l_i = (idx // Int32(4)) % Int32(32)
+                    e_i = idx % Int32(4)
+                    acc_sum = Float32(0.0)
+                    for wi in cutlass.range_constexpr(_WARPS):
+                        acc_sum = acc_sum + slab_f32[Int32(wi * num_tiles * nt * red_stride) + idx]
+                    m = t_i * Int32(_TILE_M) + l_i // Int32(4) + Int32(8) * (e_i // Int32(2))
+                    r = j_i * Int32(8) + Int32(2) * (l_i % Int32(4)) + (e_i % Int32(2))
+                    if m < nrows:
+                        if r < out_limit:
+                            out[r, n0 + m] = acc_sum.to(BFloat16)
+                    idx = idx + Int32(_COMPUTE_THREADS)
+                compute_sync()
+                if tidx == Int32(0):
+                    red_add_u64(finished_addr, Int64(1))
+                    if cutlass.const_expr(traced):
+                        trace[bidx, 0] = t_entry
+                        trace[bidx, 2] = t_published
+                        trace[bidx, 3] = t_ready
+                        trace[bidx, 4] = t_tile0
+                        trace[bidx, 5] = global_timer()
+            if cutlass.const_expr(traced):
+                if warp == Int32(_WARPS):
+                    if lane == Int32(0):
+                        trace[bidx, 1] = t_issued
+
+        @cute.jit
+        def __call__(
+            self,
+            latent: cute.Tensor,
+            w_uv: cute.Tensor,
+            w: cute.Tensor,
+            v: cute.Tensor,
+            out: cute.Tensor,
+            state: cute.Tensor,
+            trace: cute.Tensor,
+            stream,
+        ):
+            self.kernel(latent, w_uv, w, v, out, state, trace).launch(
+                grid=(grid, 1, 1),
+                block=(_THREADS, 1, 1),
+                stream=stream,
+                use_pdl=pdl,
+                min_blocks_per_mp=1,
+            )
+
+    r_sym, v_sym, o_sym = cute.sym_int(), cute.sym_int(), cute.sym_int()
+
+    def fake(dtype, shape):
+        return cute_runtime.make_fake_compact_tensor(
+            dtype, shape, stride_order=tuple(range(len(shape) - 1, -1, -1)), assumed_align=16
+        )
+
+    return cute.compile(
+        FusedTail(),
+        fake(BFloat16, (r_sym, heads * c_dim)),
+        fake(BFloat16, (heads * v_dim, c_dim)),
+        fake(BFloat16, (n_dim, k_dim)),
+        fake(BFloat16, (v_sym, k_dim)),
+        fake(BFloat16, (o_sym, n_dim)),
+        cute_runtime.make_fake_compact_tensor(Int64, (2,), stride_order=(0,), assumed_align=16),
+        cute_runtime.make_fake_compact_tensor(
+            Int64, (grid, 8), stride_order=(1, 0), assumed_align=16
+        ),
+        stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def new_state(device: torch.device) -> torch.Tensor:
+    """Allocate the zeroed ``[finished, done]`` counters of one tail instance."""
+    return torch.zeros(2, dtype=torch.int64, device=device)
+
+
+_NO_TRACE = {}
+
+
+def fused_tail(
+    latent: torch.Tensor,
+    w_uv: torch.Tensor,
+    w_o: torch.Tensor,
+    out: torch.Tensor,
+    state: torch.Tensor,
+    v_scratch: Optional[torch.Tensor] = None,
+    trace: Optional[torch.Tensor] = None,
+    window: int = DEFAULT_WINDOW,
+    pdl_entry: bool = False,
+) -> None:
+    """Launch the single-kernel tail on the current stream.
+
+    Args:
+        latent: ``[num_rows, H * C]`` contiguous BF16, ``num_rows <= 16``.
+        w_uv: ``[H, V, C]`` contiguous BF16.
+        w_o: ``[N, H * V]`` contiguous BF16, 16-byte aligned.
+        out: ``[out_rows, N]`` contiguous BF16; rows ``[0, min(out_rows, 16))``
+            are written (rows past ``num_rows`` as zeros).
+        state: ``[2]`` int64 counters from ``new_state``, private to the
+            caller and never used by two invocations at once.
+        v_scratch: optional ``[num_rows, H * V]`` BF16 buffer for ``v``.
+        trace: optional ``[grid, 8]`` int64 buffer; when given, a traced
+            variant records ``%globaltimer`` at entry, after the last weight
+            copy is issued, after the v slice is published, after the wait,
+            when the first weight tile is ready, and at exit (columns 0-5).
+        window: ``w_o`` rows a CTA keeps in flight.
+        pdl_entry: launch as a PDL dependent of the latent's producer. The
+            producer warp then streams ``w_o`` and W_UV CTAs load ``w_uv``
+            before ``griddepcontrol.wait``; every latent read, every global
+            write and the read of ``finished`` follow the wait.
+    """
+    heads, v_dim, c_dim = w_uv.shape
+    n_dim, k_dim = w_o.shape
+    num_rows = latent.shape[0]
+    if v_scratch is None:
+        v_scratch = torch.empty((num_rows, k_dim), dtype=w_o.dtype, device=w_o.device)
+    grid = _sm_count(w_o.device.index or 0)
+    nt = 1 if max(num_rows, out.shape[0]) <= 8 else 2
+    traced = trace is not None
+    if trace is None:
+        key = (w_o.device.index or 0, grid)
+        if key not in _NO_TRACE:
+            _NO_TRACE[key] = torch.empty((grid, 8), dtype=torch.int64, device=w_o.device)
+        trace = _NO_TRACE[key]
+    compiled = _build_fused(heads, v_dim, c_dim, n_dim, grid, nt, traced, window, pdl_entry)
+    compiled(latent, w_uv.view(heads * v_dim, c_dim), w_o, v_scratch, out, state, trace)
+
+
+__all__ = ["DEFAULT_WINDOW", "check_fused_shapes", "fused_tail", "new_state"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py
new file mode 100644
index 000000000000..6ab70ca5ab18
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py
@@ -0,0 +1,390 @@
+# 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.
+"""CuTe DSL streaming o_proj kernel for the attention_output_tail (SM100/SM103).
+
+``out[r, n] = bf16(sum_k v[r, k] * w_o[n, k])`` with FP32 accumulation, for at
+most 16 live rows ``r``. The kernel is built to start streaming ``w_o`` before
+its input ``v`` exists:
+
+* One persistent CTA per SM owns a contiguous block of 40-41 ``w_o`` rows.
+  At entry, warp 0 issues a 1-D ``cp.async.bulk`` global→shared copy per row
+  (4 KiB each) into a row-padded shared-memory slab, completing on one
+  mbarrier per 16-row tile.
+* All threads then execute ``griddepcontrol.wait``. When the kernel is a PDL
+  dependent of the W_UV kernel, this is where it waits for ``v``; the first
+  ``early_rows`` rows per CTA are already in flight and the rest are issued
+  after the wait, so the early stream does not queue ~24 MiB of copies in
+  front of the W_UV kernel's own memory traffic. Without PDL every row is
+  issued at entry.
+* Eight warps split K into 256-wide ranges. Each warp keeps its ``v``
+  fragments in registers and runs ``mma.sync.m16n8k16`` (BF16 in, FP32
+  accumulate) over each 16-row tile as soon as that tile's mbarrier flips.
+* The per-warp partial sums are added in fixed warp order through shared
+  memory (reusing the slab), so the result is deterministic, and rounded once
+  to BF16. Rows ``[num_rows, out_rows)`` are written as zeros.
+
+The 16-byte row padding (pitch ``K + 8`` elements) makes the MMA operand loads
+bank-conflict free.
+"""
+
+import functools
+import math
+from typing import Optional
+
+import torch
+
+_THREADS = 256
+_WARPS = _THREADS // 32
+_ROW_PAD = 8  # BF16 elements of padding per shared-memory row (16 bytes)
+_TILE_M = 16
+
+
+@functools.lru_cache(maxsize=None)
+def _sm_count(device_index: int) -> int:
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def plan(n_dim: int, k_dim: int, grid: int) -> dict:
+    """Compile-time geometry of the kernel for one weight shape."""
+    rows_max = math.ceil(n_dim / grid)
+    pitch = k_dim + _ROW_PAD
+    return {
+        "rows_max": rows_max,
+        "num_tiles": math.ceil(rows_max / _TILE_M),
+        "pitch": pitch,
+        "smem_bytes": rows_max * pitch * 2,
+        "k_per_warp": k_dim // _WARPS,
+    }
+
+
+def check_shapes(n_dim: int, k_dim: int, grid: int) -> Optional[str]:
+    """Return why the kernel cannot take this ``[n_dim, k_dim]`` weight."""
+    if k_dim % (_WARPS * 16):
+        return "K is not a multiple of 128"
+    p = plan(n_dim, k_dim, grid)
+    red_bytes = _WARPS * p["num_tiles"] * 2 * 32 * 4 * 4
+    if p["smem_bytes"] < red_bytes:
+        return "slab smaller than the reduction buffer"
+    # 227 KB of opt-in shared memory per CTA on SM100/SM103, minus the
+    # tile mbarriers and alignment.
+    if p["smem_bytes"] > 226 * 1024:
+        return "weight slab does not fit in shared memory"
+    if (k_dim * 2) % 16:
+        return "row bytes not a 16-byte multiple"
+    return None
+
+
+def _build(
+    n_dim: int,
+    k_dim: int,
+    grid: int,
+    n_tiles_r: int,
+    use_pdl: bool,
+    early_rows: int,
+    trigger_dependents: bool,
+    trigger_late: bool,
+):
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import BFloat16, Float32, Int32
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cute import runtime as cute_runtime
+    from cutlass.cutlass_dsl import dsl_user_op
+    from cutlass.utils.smem_allocator import SmemAllocator
+
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+
+    geo = plan(n_dim, k_dim, grid)
+    rows_max = geo["rows_max"]
+    num_tiles = geo["num_tiles"]
+    pitch = geo["pitch"]
+    half_pitch = pitch // 2
+    k_per_warp = geo["k_per_warp"]
+    steps = k_per_warp // 16
+    row_bytes = k_dim * 2
+    nt = n_tiles_r
+    red_stride = 128  # 32 lanes x 4 accumulators
+
+    @dsl_user_op
+    def bulk_g2s(dst_smem, src_gmem, num_bytes, mbar, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                dst_smem.toint(loc=loc, ip=ip).ir_value(),
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+                mbar.toint(loc=loc, ip=ip).ir_value(),
+            ],
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            " [$0], [$1], $2, [$3];",
+            "r,l,r,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def grid_launch_dependents(*, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [],
+            "griddepcontrol.launch_dependents;",
+            "",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def grid_dependency_wait(*, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [],
+            "griddepcontrol.wait;",
+            "",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    class StreamOproj:
+        @cute.kernel
+        def kernel(self, v: cute.Tensor, w: cute.Tensor, out: cute.Tensor):
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            warp = tidx // Int32(32)
+            lane = tidx % Int32(32)
+            grp = lane // Int32(4)
+            quad = lane % Int32(4)
+
+            num_rows = Int32(v.shape[0])
+            out_rows = Int32(out.shape[0])
+            n0 = (bidx * Int32(n_dim)) // Int32(grid)
+            n1 = ((bidx + Int32(1)) * Int32(n_dim)) // Int32(grid)
+            nrows = n1 - n0
+
+            smem = SmemAllocator()
+            slab = smem.allocate_tensor(
+                BFloat16, cute.make_layout((rows_max * pitch,)), byte_alignment=128
+            )
+            mbar = smem.allocate_tensor(
+                cutlass.Int64, cute.make_layout((num_tiles,)), byte_alignment=8
+            )
+            slab32 = cute.recast_tensor(slab, Int32)
+            slab_f32 = cute.recast_tensor(slab, Float32)
+            v32 = cute.recast_tensor(v, Int32)
+
+            if tidx == Int32(0):
+                for t in cutlass.range_constexpr(num_tiles):
+                    cute.arch.mbarrier_init(mbar.iterator + t, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+
+            # Stream this CTA's rows of w into shared memory before waiting
+            # for the producer of v: lane 0 arms the tile mbarriers, then
+            # warp 0 issues one 1-D bulk copy per row.
+            if tidx == Int32(0):
+                for t in cutlass.range_constexpr(num_tiles):
+                    tile_rows = nrows - Int32(t * _TILE_M)
+                    if tile_rows > Int32(_TILE_M):
+                        tile_rows = Int32(_TILE_M)
+                    if tile_rows < Int32(0):
+                        tile_rows = Int32(0)
+                    cute.arch.mbarrier_arrive_and_expect_tx(
+                        mbar.iterator + t, tile_rows * Int32(row_bytes)
+                    )
+            cute.arch.sync_warp()
+            # Rows [0, early_rows) go out before the wait on v's producer and
+            # the rest after it: a bounded early stream overlaps the producer
+            # without queueing ~24 MiB of copies in front of its traffic.
+            if warp == Int32(0):
+                for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                    row = lane + Int32(32 * q)
+                    if row < nrows:
+                        if row < Int32(early_rows):
+                            bulk_g2s(
+                                slab.iterator + row * Int32(pitch),
+                                w.iterator + (n0 + row) * Int32(k_dim),
+                                row_bytes,
+                                mbar.iterator + row // Int32(_TILE_M),
+                            )
+
+            if cutlass.const_expr(use_pdl):
+                grid_dependency_wait()
+
+            if warp == Int32(0):
+                for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                    row = lane + Int32(32 * q)
+                    if row < nrows:
+                        if row >= Int32(early_rows):
+                            bulk_g2s(
+                                slab.iterator + row * Int32(pitch),
+                                w.iterator + (n0 + row) * Int32(k_dim),
+                                row_bytes,
+                                mbar.iterator + row // Int32(_TILE_M),
+                            )
+            if cutlass.const_expr(trigger_dependents and not trigger_late):
+                # Dependents still wait for this whole grid (and its memory)
+                # before they read anything it produces.
+                grid_launch_dependents()
+
+            # v fragments of this warp's K range (B operand, [n][k] storage).
+            vb = cute.make_rmem_tensor((4, nt, steps), BFloat16)
+            vb32 = cute.recast_tensor(vb, Int32)
+            k_word0 = warp * Int32(k_per_warp // 2) + quad
+            for s in cutlass.range_constexpr(steps):
+                for j in cutlass.range_constexpr(nt):
+                    r = Int32(8 * j) + grp
+                    lo = Int32(0)
+                    hi = Int32(0)
+                    if r < num_rows:
+                        lo = v32[r, k_word0 + Int32(8 * s)]
+                        hi = v32[r, k_word0 + Int32(8 * s + 4)]
+                    vb32[0, j, s] = lo
+                    vb32[1, j, s] = hi
+
+            acc = cute.make_rmem_tensor((4, nt, num_tiles), Float32)
+            acc.fill(0.0)
+            a = cute.make_rmem_tensor((8,), BFloat16)
+            a32 = cute.recast_tensor(a, Int32)
+            last = nrows - Int32(1)
+            for t in cutlass.range_constexpr(num_tiles):
+                cute.arch.mbarrier_wait(mbar.iterator + t, 0)
+                row_lo = Int32(t * _TILE_M) + grp
+                row_hi = row_lo + Int32(8)
+                if row_lo > last:
+                    row_lo = last
+                if row_hi > last:
+                    row_hi = last
+                base_lo = row_lo * Int32(half_pitch) + k_word0
+                base_hi = row_hi * Int32(half_pitch) + k_word0
+                for s in cutlass.range_constexpr(steps):
+                    a32[0] = slab32[base_lo + Int32(8 * s)]
+                    a32[1] = slab32[base_hi + Int32(8 * s)]
+                    a32[2] = slab32[base_lo + Int32(8 * s + 4)]
+                    a32[3] = slab32[base_hi + Int32(8 * s + 4)]
+                    for j in cutlass.range_constexpr(nt):
+                        acc[None, j, t] = mma_sync(a, vb[None, j, s], acc[None, j, t])
+
+            if cutlass.const_expr(trigger_dependents and trigger_late):
+                # Late trigger: only once this CTA's MMAs are done.
+                grid_launch_dependents()
+            # Fixed-order cross-warp reduction through the (now idle) slab.
+            cute.arch.sync_threads()
+            for t in cutlass.range_constexpr(num_tiles):
+                for j in cutlass.range_constexpr(nt):
+                    base = ((warp * Int32(num_tiles) + Int32(t)) * Int32(nt) + Int32(j)) * Int32(
+                        red_stride
+                    ) + lane * Int32(4)
+                    for e in cutlass.range_constexpr(4):
+                        slab_f32[base + Int32(e)] = acc[e, j, t]
+            cute.arch.sync_threads()
+
+            total = num_tiles * nt * red_stride
+            out_limit = out_rows
+            if out_limit > Int32(8 * nt):
+                out_limit = Int32(8 * nt)
+            idx = tidx
+            while idx < Int32(total):
+                t_i = idx // Int32(nt * red_stride)
+                j_i = (idx // Int32(red_stride)) % Int32(nt)
+                l_i = (idx // Int32(4)) % Int32(32)
+                e_i = idx % Int32(4)
+                acc_sum = Float32(0.0)
+                for wi in cutlass.range_constexpr(_WARPS):
+                    acc_sum = acc_sum + slab_f32[Int32(wi * num_tiles * nt * red_stride) + idx]
+                m = t_i * Int32(_TILE_M) + l_i // Int32(4) + Int32(8) * (e_i // Int32(2))
+                r = j_i * Int32(8) + Int32(2) * (l_i % Int32(4)) + e_i % Int32(2)
+                if m < nrows:
+                    if r < out_limit:
+                        out[r, n0 + m] = acc_sum.to(BFloat16)
+                idx = idx + Int32(_THREADS)
+
+        @cute.jit
+        def __call__(self, v: cute.Tensor, w: cute.Tensor, out: cute.Tensor, stream):
+            self.kernel(v, w, out).launch(
+                grid=(grid, 1, 1),
+                block=(_THREADS, 1, 1),
+                stream=stream,
+                use_pdl=use_pdl,
+                min_blocks_per_mp=1,
+            )
+
+    rows_sym, out_sym = cute.sym_int(), cute.sym_int()
+    v_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (rows_sym, k_dim), stride_order=(1, 0), assumed_align=16
+    )
+    w_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (n_dim, k_dim), stride_order=(1, 0), assumed_align=16
+    )
+    out_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (out_sym, n_dim), stride_order=(1, 0), assumed_align=16
+    )
+    stream = cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        StreamOproj(), v_fake, w_fake, out_fake, stream=stream, options="--enable-tvm-ffi"
+    )
+
+
+_COMPILED = {}
+
+
+DEFAULT_EARLY_ROWS = 16
+
+
+def stream_oproj(
+    v: torch.Tensor,
+    w_o: torch.Tensor,
+    out: torch.Tensor,
+    use_pdl: bool,
+    grid: Optional[int] = None,
+    early_rows: int = DEFAULT_EARLY_ROWS,
+    trigger_dependents: bool = False,
+    trigger_late: bool = False,
+) -> None:
+    """Launch the streaming o_proj kernel on the current stream.
+
+    Args:
+        v: ``[num_rows, K]`` contiguous BF16, ``num_rows <= 16``.
+        w_o: ``[N, K]`` contiguous BF16, 16-byte aligned.
+        out: ``[out_rows, N]`` contiguous BF16; the kernel writes rows
+            ``[0, min(out_rows, 16))`` (rows past ``num_rows`` as zeros).
+        use_pdl: launch as a PDL dependent of the previous kernel.
+        grid: CTA count (default: the device's SM count).
+        early_rows: with ``use_pdl``, rows of ``w_o`` per CTA copied before
+            ``griddepcontrol.wait`` (the rest follow the wait).
+        trigger_dependents: execute ``griddepcontrol.launch_dependents`` once
+            all row copies are issued, so a PDL-launched consumer can start
+            its prologue; it must still wait for this grid before reading.
+        trigger_late: with ``trigger_dependents``, trigger only after the
+            CTA's MMAs instead (experiments).
+    """
+    n_dim, k_dim = w_o.shape
+    if grid is None:
+        grid = _sm_count(w_o.device.index or 0)
+    n_tiles_r = 1 if out.shape[0] <= 8 else 2
+    if not use_pdl:
+        early_rows = 1 << 30
+    key = (n_dim, k_dim, grid, n_tiles_r, use_pdl, early_rows, trigger_dependents, trigger_late)
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        compiled = _build(
+            n_dim, k_dim, grid, n_tiles_r, use_pdl, early_rows, trigger_dependents, trigger_late
+        )
+        _COMPILED[key] = compiled
+    compiled(v, w_o, out)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py
new file mode 100644
index 000000000000..521fe4810d70
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py
@@ -0,0 +1,461 @@
+# 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.
+"""Small-footprint CuTe DSL o_proj for weights staged in L2 (SM100/SM103).
+
+``out[r, n] = bf16(sum_k v[r, k] * w_o[n, k])`` with FP32 accumulation, at
+most 16 live rows ``r``, like ``cute_oproj.stream_oproj`` but on ``grid``
+CTAs (default 96, the footprint of the native o_proj GEMM) instead of one per
+SM, so a register-heavy PDL successor (the post_attention_hop needs a whole SM per CTA)
+can become resident on the remaining SMs while this kernel runs.
+
+* CTA ``b`` owns ``rows = N / grid`` contiguous ``w_o`` rows, processed as
+  16-row tiles through a ring of ``slots`` row-padded shared-memory tiles
+  (pitch ``K + 8`` elements, bank-conflict-free operand loads). Each tile
+  completes on its own mbarrier; a slot is refilled with the tile
+  ``slots`` positions later once all warps have consumed it.
+* The first ``early`` tiles are requested before ``griddepcontrol.wait``
+  (only immutable weights are read before the wait); everything else, and
+  every read of ``v``, follows it.
+* Eight warps split K into 256-wide ranges, keep their ``v`` fragments in
+  registers and run ``mma.sync.m16n8k16`` (BF16 in, FP32 accumulate); the
+  per-warp partials are added in fixed warp order (deterministic) and rounded
+  once to BF16. Rows ``[num_rows, out_rows)`` are written as zeros.
+* The PDL successor is released once this CTA's MMAs are issued after the
+  wait (``trigger``); it still waits for the whole grid before reading.
+* ``publish`` (not supported; ``ring_oproj`` refuses it): the epilogue would
+  additionally write the live rows of the partial into slot ``[row][rank]``
+  of the current Lamport buffer of the post_attention_hop workspace. The into-hop entry polls, and
+  later resets, every row of the hop including padding, so a producer that
+  publishes only the live rows leaves a padded call spinning or leaves stale
+  words behind; attention_output_into_hop publishes through ``into_hop``, which writes every
+  hop row.
+"""
+
+import functools
+from typing import Optional
+
+import torch
+
+_THREADS = 256
+_WARPS = _THREADS // 32
+_ROW_PAD = 8
+_TILE_M = 16
+_HIDDEN = 6144
+_NUM_RANKS = 8
+_LAMPORT_BUFFER_WORDS = 8 * _NUM_RANKS * _HIDDEN * 2 // 4
+
+
+def plan(n_dim: int, k_dim: int, grid: int, slots: int) -> dict:
+    rows = n_dim // grid
+    pitch = k_dim + _ROW_PAD
+    return {
+        "rows": rows,
+        "tiles": rows // _TILE_M,
+        "pitch": pitch,
+        "smem_bytes": slots * _TILE_M * pitch * 2,
+        "k_per_warp": k_dim // _WARPS,
+    }
+
+
+def check_shapes(n_dim: int, k_dim: int, grid: int, slots: int) -> Optional[str]:
+    """Return why the kernel cannot take this ``[n_dim, k_dim]`` weight."""
+    if k_dim % (_WARPS * 16) or (k_dim * 2) % 16:
+        return "K is not a multiple of 128"
+    if n_dim % (grid * _TILE_M):
+        return "N is not a multiple of grid * 16"
+    p = plan(n_dim, k_dim, grid, slots)
+    if p["smem_bytes"] > 210 * 1024:
+        return "ring does not fit in shared memory"
+    if p["smem_bytes"] < _WARPS * p["tiles"] * 2 * 128 * 4:
+        return "ring smaller than the reduction buffer"
+    return None
+
+
+@functools.lru_cache(maxsize=None)
+def _build(
+    n_dim: int,
+    k_dim: int,
+    grid: int,
+    slots: int,
+    early: int,
+    nt: int,
+    use_pdl: bool,
+    trigger: bool,
+    publish_rank: int,
+):
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import BFloat16, Float32, Int32, Int64
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cute import runtime as cute_runtime
+    from cutlass.cutlass_dsl import T, dsl_user_op
+    from cutlass.utils.smem_allocator import SmemAllocator
+
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+
+    geo = plan(n_dim, k_dim, grid, slots)
+    rows = geo["rows"]
+    tiles = geo["tiles"]
+    pitch = geo["pitch"]
+    half_pitch = pitch // 2
+    k_per_warp = geo["k_per_warp"]
+    steps = k_per_warp // 16
+    row_bytes = k_dim * 2
+    red_stride = 128
+    publish = publish_rank >= 0
+    slots_used = min(slots, tiles)
+    early = min(early, slots_used)
+
+    def _asm(res, operands, text, constraints, loc, ip):
+        return llvm.inline_asm(
+            res, operands, text, constraints, has_side_effects=True, asm_dialect=0, loc=loc, ip=ip
+        )
+
+    @dsl_user_op
+    def bulk_g2s(dst_smem, src_gmem, num_bytes, mbar, *, loc=None, ip=None):
+        _asm(
+            None,
+            [
+                dst_smem.toint(loc=loc, ip=ip).ir_value(),
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+                mbar.toint(loc=loc, ip=ip).ir_value(),
+            ],
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            " [$0], [$1], $2, [$3];",
+            "r,l,r,r",
+            loc,
+            ip,
+        )
+
+    @dsl_user_op
+    def grid_dependency_wait(*, loc=None, ip=None):
+        _asm(None, [], "griddepcontrol.wait;", "", loc, ip)
+
+    @dsl_user_op
+    def grid_launch_dependents(*, loc=None, ip=None):
+        _asm(None, [], "griddepcontrol.launch_dependents;", "", loc, ip)
+
+    @dsl_user_op
+    def ld_volatile_u32(addr, *, loc=None, ip=None):
+        return Int32(
+            _asm(
+                T.i32(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.volatile.global.u32 $0, [$1];",
+                "=r,l",
+                loc,
+                ip,
+            )
+        )
+
+    @dsl_user_op
+    def st_global_b32(addr, val, *, loc=None, ip=None):
+        _asm(
+            None,
+            [Int64(addr).ir_value(loc=loc, ip=ip), Int32(val).ir_value(loc=loc, ip=ip)],
+            "st.global.b32 [$0], $1;",
+            "l,r",
+            loc,
+            ip,
+        )
+
+    class RingOproj:
+        @cute.kernel
+        def kernel(
+            self,
+            v: cute.Tensor,
+            w: cute.Tensor,
+            out: cute.Tensor,
+            mc: cute.Tensor,
+            flags: cute.Tensor,
+        ):
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            warp = tidx // Int32(32)
+            lane = tidx % Int32(32)
+            grp = lane // Int32(4)
+            quad = lane % Int32(4)
+
+            num_rows = Int32(v.shape[0])
+            out_rows = Int32(out.shape[0])
+            n0 = bidx * Int32(rows)
+
+            smem = SmemAllocator()
+            slab = smem.allocate_tensor(
+                BFloat16, cute.make_layout((slots_used * _TILE_M * pitch,)), byte_alignment=128
+            )
+            mbar = smem.allocate_tensor(
+                cutlass.Int64, cute.make_layout((slots_used,)), byte_alignment=8
+            )
+            slab32 = cute.recast_tensor(slab, Int32)
+            slab_f32 = cute.recast_tensor(slab, Float32)
+            v32 = cute.recast_tensor(v, Int32)
+
+            if tidx == Int32(0):
+                for s in cutlass.range_constexpr(slots_used):
+                    cute.arch.mbarrier_init(mbar.iterator + s, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+
+            # Tile t goes to slot t % slots: lane 0 of warp 0 arms the slot's
+            # mbarrier, lanes 0..15 copy one row each.
+            if warp == Int32(0):
+                for t in cutlass.range_constexpr(early):
+                    if lane == Int32(0):
+                        cute.arch.mbarrier_arrive_and_expect_tx(
+                            mbar.iterator + t, Int32(_TILE_M * row_bytes)
+                        )
+                    cute.arch.sync_warp()
+                    if lane < Int32(_TILE_M):
+                        bulk_g2s(
+                            slab.iterator + (Int32(t * _TILE_M) + lane) * Int32(pitch),
+                            w.iterator + (n0 + Int32(t * _TILE_M) + lane) * Int32(k_dim),
+                            row_bytes,
+                            mbar.iterator + t,
+                        )
+
+            if cutlass.const_expr(use_pdl):
+                grid_dependency_wait()
+
+            if warp == Int32(0):
+                for t in cutlass.range_constexpr(early, slots_used):
+                    if lane == Int32(0):
+                        cute.arch.mbarrier_arrive_and_expect_tx(
+                            mbar.iterator + t, Int32(_TILE_M * row_bytes)
+                        )
+                    cute.arch.sync_warp()
+                    if lane < Int32(_TILE_M):
+                        bulk_g2s(
+                            slab.iterator + (Int32(t * _TILE_M) + lane) * Int32(pitch),
+                            w.iterator + (n0 + Int32(t * _TILE_M) + lane) * Int32(k_dim),
+                            row_bytes,
+                            mbar.iterator + t,
+                        )
+
+            # v fragments of this warp's K range (B operand, [n][k] storage).
+            vb = cute.make_rmem_tensor((4, nt, steps), BFloat16)
+            vb32 = cute.recast_tensor(vb, Int32)
+            k_word0 = warp * Int32(k_per_warp // 2) + quad
+            for s in cutlass.range_constexpr(steps):
+                for j in cutlass.range_constexpr(nt):
+                    r = Int32(8 * j) + grp
+                    lo = Int32(0)
+                    hi = Int32(0)
+                    if r < num_rows:
+                        lo = v32[r, k_word0 + Int32(8 * s)]
+                        hi = v32[r, k_word0 + Int32(8 * s + 4)]
+                    vb32[0, j, s] = lo
+                    vb32[1, j, s] = hi
+
+            if cutlass.const_expr(trigger):
+                # After the wait: every grid before this one has completed.
+                grid_launch_dependents()
+
+            acc = cute.make_rmem_tensor((4, nt, tiles), Float32)
+            acc.fill(0.0)
+            a = cute.make_rmem_tensor((8,), BFloat16)
+            a32 = cute.recast_tensor(a, Int32)
+            for t in cutlass.range_constexpr(tiles):
+                s = t % slots_used
+                cute.arch.mbarrier_wait(mbar.iterator + s, (t // slots_used) % 2)
+                base_lo = (Int32(s * _TILE_M) + grp) * Int32(half_pitch) + (k_word0)
+                base_hi = base_lo + Int32(8 * half_pitch)
+                for st in cutlass.range_constexpr(steps):
+                    a32[0] = slab32[base_lo + Int32(8 * st)]
+                    a32[1] = slab32[base_hi + Int32(8 * st)]
+                    a32[2] = slab32[base_lo + Int32(8 * st + 4)]
+                    a32[3] = slab32[base_hi + Int32(8 * st + 4)]
+                    for j in cutlass.range_constexpr(nt):
+                        acc[None, j, t] = mma_sync(a, vb[None, j, st], acc[None, j, t])
+                if cutlass.const_expr(t + slots_used < tiles):
+                    # Every warp has read slot s: refill it with tile
+                    # t + slots.
+                    cute.arch.sync_threads()
+                    if warp == Int32(0):
+                        if lane == Int32(0):
+                            cute.arch.mbarrier_arrive_and_expect_tx(
+                                mbar.iterator + s, Int32(_TILE_M * row_bytes)
+                            )
+                        cute.arch.sync_warp()
+                        if lane < Int32(_TILE_M):
+                            bulk_g2s(
+                                slab.iterator + (Int32(s * _TILE_M) + lane) * Int32(pitch),
+                                w.iterator
+                                + (n0 + Int32((t + slots_used) * _TILE_M) + lane) * Int32(k_dim),
+                                row_bytes,
+                                mbar.iterator + s,
+                            )
+
+            # Fixed-order cross-warp reduction through the (now idle) ring.
+            cute.arch.sync_threads()
+            for t in cutlass.range_constexpr(tiles):
+                for j in cutlass.range_constexpr(nt):
+                    base = ((warp * Int32(tiles) + Int32(t)) * Int32(nt) + Int32(j)) * Int32(
+                        red_stride
+                    ) + lane * Int32(4)
+                    for e in cutlass.range_constexpr(4):
+                        slab_f32[base + Int32(e)] = acc[e, j, t]
+            cute.arch.sync_threads()
+
+            out_limit = out_rows
+            if out_limit > Int32(8 * nt):
+                out_limit = Int32(8 * nt)
+            buf_word0 = Int32(0)
+            if cutlass.const_expr(publish):
+                buf_word0 = ld_volatile_u32(flags.iterator.toint()) * Int32(_LAMPORT_BUFFER_WORDS)
+            mc_base = mc.iterator.toint()
+            # Pairs of adjacent output columns (m, m + 1) of one row: lane
+            # groups 2g and 2g + 1 of the MMA C fragment.
+            pk = cute.make_rmem_tensor((2,), BFloat16)
+            pk32 = cute.recast_tensor(pk, Int32)
+            total_pairs = tiles * nt * 64
+            p = tidx
+            while p < Int32(total_pairs):
+                t_i = p // Int32(nt * 64)
+                j_i = (p // Int32(64)) % Int32(nt)
+                q = p % Int32(64)
+                gh = q // Int32(16)
+                quad_i = (q // Int32(4)) % Int32(4)
+                e_i = q % Int32(4)
+                idx0 = (
+                    (t_i * Int32(nt) + j_i) * Int32(red_stride)
+                    + (gh * Int32(8) + quad_i) * Int32(4)
+                    + e_i
+                )
+                idx1 = idx0 + Int32(16)
+                s0 = Float32(0.0)
+                s1 = Float32(0.0)
+                for wi in cutlass.range_constexpr(_WARPS):
+                    s0 = s0 + slab_f32[Int32(wi * tiles * nt * red_stride) + idx0]
+                    s1 = s1 + slab_f32[Int32(wi * tiles * nt * red_stride) + idx1]
+                m0 = t_i * Int32(_TILE_M) + gh * Int32(2) + Int32(8) * (e_i // Int32(2))
+                r = j_i * Int32(8) + Int32(2) * quad_i + e_i % Int32(2)
+                if r < out_limit:
+                    pk[0] = s0.to(BFloat16)
+                    pk[1] = s1.to(BFloat16)
+                    out[r, n0 + m0] = pk[0]
+                    out[r, n0 + m0 + Int32(1)] = pk[1]
+                    if cutlass.const_expr(publish):
+                        if r < num_rows:
+                            word = pk32[0]
+                            lo = word & Int32(0xFFFF)
+                            hi = (word >> Int32(16)) & Int32(0xFFFF)
+                            if lo == Int32(0x8000):
+                                lo = Int32(0)
+                            if hi == Int32(0x8000):
+                                hi = Int32(0)
+                            word = lo | (hi << Int32(16))
+                            widx = buf_word0 + (
+                                (r * Int32(_NUM_RANKS) + Int32(publish_rank)) * Int32(_HIDDEN)
+                                + n0
+                                + m0
+                            ) // Int32(2)
+                            st_global_b32(mc_base + Int64(widx) * Int64(4), word)
+                p = p + Int32(_THREADS)
+
+        @cute.jit
+        def __call__(
+            self,
+            v: cute.Tensor,
+            w: cute.Tensor,
+            out: cute.Tensor,
+            mc: cute.Tensor,
+            flags: cute.Tensor,
+            stream,
+        ):
+            self.kernel(v, w, out, mc, flags).launch(
+                grid=(grid, 1, 1),
+                block=(_THREADS, 1, 1),
+                stream=stream,
+                use_pdl=use_pdl,
+                min_blocks_per_mp=1,
+            )
+
+    rows_sym, out_sym = cute.sym_int(), cute.sym_int()
+    mc_sym = cute.sym_int()
+
+    def fake(dtype, shape):
+        return cute_runtime.make_fake_compact_tensor(
+            dtype, shape, stride_order=tuple(range(len(shape) - 1, -1, -1)), assumed_align=16
+        )
+
+    return cute.compile(
+        RingOproj(),
+        fake(BFloat16, (rows_sym, k_dim)),
+        fake(BFloat16, (n_dim, k_dim)),
+        fake(BFloat16, (out_sym, n_dim)),
+        fake(Int32, (mc_sym,)),
+        fake(Int32, (4,)),
+        stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+_DUMMY = {}
+
+
+def _dummy(device: torch.device) -> torch.Tensor:
+    key = device.index or 0
+    if key not in _DUMMY:
+        _DUMMY[key] = torch.zeros(4, dtype=torch.int32, device=device)
+    return _DUMMY[key]
+
+
+def ring_oproj(
+    v: torch.Tensor,
+    w_o: torch.Tensor,
+    out: torch.Tensor,
+    use_pdl: bool,
+    grid: int = 96,
+    slots: int = 3,
+    early: int = 3,
+    trigger: bool = True,
+    mc_words: Optional[torch.Tensor] = None,
+    flags: Optional[torch.Tensor] = None,
+    publish_rank: int = -1,
+) -> None:
+    """Launch the small-footprint o_proj on the current stream.
+
+    Args:
+        v: ``[num_rows, K]`` contiguous BF16, ``num_rows <= 16``.
+        w_o: ``[N, K]`` contiguous BF16, 16-byte aligned.
+        out: ``[out_rows, N]`` contiguous BF16; rows ``[0, min(out_rows, 16))``
+            are written (rows past ``num_rows`` as zeros).
+        use_pdl: launch as a PDL dependent of the previous kernel.
+        grid: CTAs; ``N`` must be a multiple of ``16 * grid``.
+        slots: 16-row shared-memory tiles per CTA.
+        early: tiles requested before ``griddepcontrol.wait``.
+        trigger: release the PDL successor after the wait.
+        mc_words / flags / publish_rank: unsupported publish target; only
+            ``publish_rank < 0`` (no publishing) is accepted.
+
+    Raises:
+        ValueError: ``publish_rank >= 0`` (see the module docstring).
+    """
+    n_dim, k_dim = w_o.shape
+    nt = 1 if out.shape[0] <= 8 else 2
+    if publish_rank >= 0:
+        raise ValueError(
+            "ring_oproj publishing is withdrawn: it publishes only the live rows, "
+            "the attention_output_into_hop hop needs every hop row; publish with into_hop"
+        )
+    del mc_words, flags
+    mc_words = flags = _dummy(w_o.device)
+    compiled = _build(n_dim, k_dim, grid, slots, early, nt, use_pdl, trigger, publish_rank)
+    compiled(v, w_o, out, mc_words, flags)
+
+
+__all__ = ["check_shapes", "plan", "ring_oproj"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py
new file mode 100644
index 000000000000..7a6dec5e0681
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py
@@ -0,0 +1,566 @@
+# 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.
+"""attention_output_tail streaming o_proj that also publishes into the post_attention_hop transport
+(``attention_output_into_hop``).
+
+The kernel of ``cute_oproj.stream_oproj`` (one CTA per SM streams its
+``w_o`` rows into a padded shared-memory slab around its PDL wait, 8 K-split
+warps run ``mma.sync``, fixed-order cross-warp reduction, one BF16
+rounding), with two changes for ``attention_output_into_hop``:
+
+* rows are split at even boundaries (40 or 42 per CTA on 152 SMs), so each
+  thread of the epilogue owns two adjacent output columns of one row and
+  stores them as one 32-bit word;
+* besides the ordinary partial (the attention_output_tail contract stays available), the epilogue
+  writes the words of every row of ``out`` (live rows and the zero rows
+  past them, as the post_attention_hop sums them from the ordinary partial), with each
+  bf16 ``-0.0`` mapped to ``+0.0`` (the Lamport sentinel can never appear),
+  into slot ``[row][rank]`` of the current Lamport buffer of the post_attention_hop
+  workspace through its multicast view. The hop polls exactly the rows of
+  its residual and resets exactly those rows of the buffer later, so the
+  caller passes an ``out`` with the hop's row count. The buffer index is read from the protocol state after
+  ``griddepcontrol.wait``; the PDL successor is released only after that
+  wait (producer side of the contract in
+  ``cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h``).
+  When every CTA owns a multiple of 8 columns (e.g. 48 on 128 CTAs) the
+  words are first gathered in shared memory and each row chunk of 8 columns
+  goes out as one 16-byte multicast store, as the post_attention_hop's own publish does;
+  otherwise every word is its own 32-bit multicast store.
+
+The arithmetic of every output element (K split, MMA order, warp order,
+rounding) is the same as ``stream_oproj``'s, so the ordinary partial equals
+``stream_oproj``'s output bit for bit.
+"""
+
+import functools
+import math
+from typing import Optional
+
+import torch
+
+_THREADS = 256
+_HIDDEN = 6144
+_NUM_RANKS = 8
+_LAMPORT_BUFFER_WORDS = 8 * _NUM_RANKS * _HIDDEN * 2 // 4
+_WARPS = _THREADS // 32
+_ROW_PAD = 8  # BF16 elements of padding per shared-memory row (16 bytes)
+_TILE_M = 16
+
+
+@functools.lru_cache(maxsize=None)
+def _sm_count(device_index: int) -> int:
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def plan(n_dim: int, k_dim: int, grid: int) -> dict:
+    """Compile-time geometry of the kernel for one weight shape."""
+    rows_max = 2 * math.ceil(n_dim / 2 / grid)
+    pitch = k_dim + _ROW_PAD
+    return {
+        "rows_max": rows_max,
+        "num_tiles": math.ceil(rows_max / _TILE_M),
+        "pitch": pitch,
+        "smem_bytes": rows_max * pitch * 2,
+        "k_per_warp": k_dim // _WARPS,
+    }
+
+
+def check_shapes(n_dim: int, k_dim: int, grid: int) -> Optional[str]:
+    """Return why the kernel cannot take this ``[n_dim, k_dim]`` weight."""
+    if k_dim % (_WARPS * 16):
+        return "K is not a multiple of 128"
+    p = plan(n_dim, k_dim, grid)
+    red_bytes = _WARPS * p["num_tiles"] * 2 * 32 * 4 * 4
+    if p["smem_bytes"] < red_bytes:
+        return "slab smaller than the reduction buffer"
+    # 227 KB of opt-in shared memory per CTA on SM100/SM103, minus the
+    # tile mbarriers and alignment.
+    if p["smem_bytes"] > 226 * 1024:
+        return "weight slab does not fit in shared memory"
+    if (k_dim * 2) % 16:
+        return "row bytes not a 16-byte multiple"
+    return None
+
+
+def _build(
+    n_dim: int,
+    k_dim: int,
+    grid: int,
+    n_tiles_r: int,
+    use_pdl: bool,
+    early_rows: int,
+    trigger_dependents: bool,
+    trigger_late: bool,
+    publish_rank: int,
+    vector_publish: bool,
+):
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import BFloat16, Float32, Int32, Int64
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cute import runtime as cute_runtime
+    from cutlass.cutlass_dsl import T, dsl_user_op
+    from cutlass.utils.smem_allocator import SmemAllocator
+
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+
+    geo = plan(n_dim, k_dim, grid)
+    rows_max = geo["rows_max"]
+    num_tiles = geo["num_tiles"]
+    pitch = geo["pitch"]
+    half_pitch = pitch // 2
+    k_per_warp = geo["k_per_warp"]
+    steps = k_per_warp // 16
+    row_bytes = k_dim * 2
+    nt = n_tiles_r
+    red_stride = 128  # 32 lanes x 4 accumulators
+    # 16-byte publish: every CTA owns the same whole number of 8-column
+    # chunks; the words are staged right after the reduction buffer.
+    vec_publish = vector_publish and n_dim % (8 * grid) == 0
+    cols = n_dim // grid
+    chunks = cols // 8
+    stage_pitch = cols // 2
+    stage0 = _WARPS * num_tiles * nt * red_stride
+    if vec_publish:
+        assert stage0 + 8 * stage_pitch <= rows_max * pitch // 2
+
+    @dsl_user_op
+    def bulk_g2s(dst_smem, src_gmem, num_bytes, mbar, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                dst_smem.toint(loc=loc, ip=ip).ir_value(),
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+                mbar.toint(loc=loc, ip=ip).ir_value(),
+            ],
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            " [$0], [$1], $2, [$3];",
+            "r,l,r,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def grid_launch_dependents(*, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [],
+            "griddepcontrol.launch_dependents;",
+            "",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def ld_volatile_u32(addr, *, loc=None, ip=None):
+        return Int32(
+            llvm.inline_asm(
+                T.i32(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.volatile.global.u32 $0, [$1];",
+                "=r,l",
+                has_side_effects=True,
+                asm_dialect=0,
+                loc=loc,
+                ip=ip,
+            )
+        )
+
+    @dsl_user_op
+    def st_global_b32(addr, val, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [Int64(addr).ir_value(loc=loc, ip=ip), Int32(val).ir_value(loc=loc, ip=ip)],
+            "st.global.b32 [$0], $1;",
+            "l,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def st_global_v4_b32(addr, w0, w1, w2, w3, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                Int64(addr).ir_value(loc=loc, ip=ip),
+                Int32(w0).ir_value(loc=loc, ip=ip),
+                Int32(w1).ir_value(loc=loc, ip=ip),
+                Int32(w2).ir_value(loc=loc, ip=ip),
+                Int32(w3).ir_value(loc=loc, ip=ip),
+            ],
+            "st.global.v4.b32 [$0], {$1, $2, $3, $4};",
+            "l,r,r,r,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def grid_dependency_wait(*, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [],
+            "griddepcontrol.wait;",
+            "",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    class StreamOprojIntoHop:
+        @cute.kernel
+        def kernel(
+            self,
+            v: cute.Tensor,
+            w: cute.Tensor,
+            out: cute.Tensor,
+            mc: cute.Tensor,
+            flags: cute.Tensor,
+        ):
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            warp = tidx // Int32(32)
+            lane = tidx % Int32(32)
+            grp = lane // Int32(4)
+            quad = lane % Int32(4)
+
+            num_rows = Int32(v.shape[0])
+            out_rows = Int32(out.shape[0])
+            # Even row boundaries: each epilogue thread stores a 32-bit pair.
+            n0 = Int32(2) * ((bidx * Int32(n_dim // 2)) // Int32(grid))
+            n1 = Int32(2) * (((bidx + Int32(1)) * Int32(n_dim // 2)) // Int32(grid))
+            nrows = n1 - n0
+
+            smem = SmemAllocator()
+            slab = smem.allocate_tensor(
+                BFloat16, cute.make_layout((rows_max * pitch,)), byte_alignment=128
+            )
+            mbar = smem.allocate_tensor(
+                cutlass.Int64, cute.make_layout((num_tiles,)), byte_alignment=8
+            )
+            slab32 = cute.recast_tensor(slab, Int32)
+            slab_f32 = cute.recast_tensor(slab, Float32)
+            v32 = cute.recast_tensor(v, Int32)
+
+            if tidx == Int32(0):
+                for t in cutlass.range_constexpr(num_tiles):
+                    cute.arch.mbarrier_init(mbar.iterator + t, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+
+            # Stream this CTA's rows of w into shared memory before waiting
+            # for the producer of v: lane 0 arms the tile mbarriers, then
+            # warp 0 issues one 1-D bulk copy per row.
+            if tidx == Int32(0):
+                for t in cutlass.range_constexpr(num_tiles):
+                    tile_rows = nrows - Int32(t * _TILE_M)
+                    if tile_rows > Int32(_TILE_M):
+                        tile_rows = Int32(_TILE_M)
+                    if tile_rows < Int32(0):
+                        tile_rows = Int32(0)
+                    cute.arch.mbarrier_arrive_and_expect_tx(
+                        mbar.iterator + t, tile_rows * Int32(row_bytes)
+                    )
+            cute.arch.sync_warp()
+            # Rows [0, early_rows) go out before the wait on v's producer and
+            # the rest after it: a bounded early stream overlaps the producer
+            # without queueing ~24 MiB of copies in front of its traffic.
+            if warp == Int32(0):
+                for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                    row = lane + Int32(32 * q)
+                    if row < nrows:
+                        if row < Int32(early_rows):
+                            bulk_g2s(
+                                slab.iterator + row * Int32(pitch),
+                                w.iterator + (n0 + row) * Int32(k_dim),
+                                row_bytes,
+                                mbar.iterator + row // Int32(_TILE_M),
+                            )
+
+            if cutlass.const_expr(use_pdl):
+                grid_dependency_wait()
+
+            if warp == Int32(0):
+                for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                    row = lane + Int32(32 * q)
+                    if row < nrows:
+                        if row >= Int32(early_rows):
+                            bulk_g2s(
+                                slab.iterator + row * Int32(pitch),
+                                w.iterator + (n0 + row) * Int32(k_dim),
+                                row_bytes,
+                                mbar.iterator + row // Int32(_TILE_M),
+                            )
+            if cutlass.const_expr(trigger_dependents and not trigger_late):
+                # Dependents still wait for this whole grid (and its memory)
+                # before they read anything it produces.
+                grid_launch_dependents()
+
+            # v fragments of this warp's K range (B operand, [n][k] storage).
+            vb = cute.make_rmem_tensor((4, nt, steps), BFloat16)
+            vb32 = cute.recast_tensor(vb, Int32)
+            k_word0 = warp * Int32(k_per_warp // 2) + quad
+            for s in cutlass.range_constexpr(steps):
+                for j in cutlass.range_constexpr(nt):
+                    r = Int32(8 * j) + grp
+                    lo = Int32(0)
+                    hi = Int32(0)
+                    if r < num_rows:
+                        lo = v32[r, k_word0 + Int32(8 * s)]
+                        hi = v32[r, k_word0 + Int32(8 * s + 4)]
+                    vb32[0, j, s] = lo
+                    vb32[1, j, s] = hi
+
+            acc = cute.make_rmem_tensor((4, nt, num_tiles), Float32)
+            acc.fill(0.0)
+            a = cute.make_rmem_tensor((8,), BFloat16)
+            a32 = cute.recast_tensor(a, Int32)
+            last = nrows - Int32(1)
+            for t in cutlass.range_constexpr(num_tiles):
+                cute.arch.mbarrier_wait(mbar.iterator + t, 0)
+                row_lo = Int32(t * _TILE_M) + grp
+                row_hi = row_lo + Int32(8)
+                if row_lo > last:
+                    row_lo = last
+                if row_hi > last:
+                    row_hi = last
+                base_lo = row_lo * Int32(half_pitch) + k_word0
+                base_hi = row_hi * Int32(half_pitch) + k_word0
+                for s in cutlass.range_constexpr(steps):
+                    a32[0] = slab32[base_lo + Int32(8 * s)]
+                    a32[1] = slab32[base_hi + Int32(8 * s)]
+                    a32[2] = slab32[base_lo + Int32(8 * s + 4)]
+                    a32[3] = slab32[base_hi + Int32(8 * s + 4)]
+                    for j in cutlass.range_constexpr(nt):
+                        acc[None, j, t] = mma_sync(a, vb[None, j, s], acc[None, j, t])
+
+            if cutlass.const_expr(trigger_dependents and trigger_late):
+                # Late trigger: only once this CTA's MMAs are done.
+                grid_launch_dependents()
+            # Fixed-order cross-warp reduction through the (now idle) slab.
+            cute.arch.sync_threads()
+            for t in cutlass.range_constexpr(num_tiles):
+                for j in cutlass.range_constexpr(nt):
+                    base = ((warp * Int32(num_tiles) + Int32(t)) * Int32(nt) + Int32(j)) * Int32(
+                        red_stride
+                    ) + lane * Int32(4)
+                    for e in cutlass.range_constexpr(4):
+                        slab_f32[base + Int32(e)] = acc[e, j, t]
+            cute.arch.sync_threads()
+
+            out_limit = out_rows
+            if out_limit > Int32(8 * nt):
+                out_limit = Int32(8 * nt)
+            # After the wait: the protocol state written by the previous hop
+            # is final.
+            buf_word0 = ld_volatile_u32(flags.iterator.toint()) * Int32(_LAMPORT_BUFFER_WORDS)
+            mc_base = mc.iterator.toint()
+            pk = cute.make_rmem_tensor((2,), BFloat16)
+            pk32 = cute.recast_tensor(pk, Int32)
+            total_pairs = num_tiles * nt * 64
+            p = tidx
+            while p < Int32(total_pairs):
+                t_i = p // Int32(nt * 64)
+                j_i = (p // Int32(64)) % Int32(nt)
+                q = p % Int32(64)
+                gh = q // Int32(16)
+                quad_i = (q // Int32(4)) % Int32(4)
+                e_i = q % Int32(4)
+                idx0 = (
+                    (t_i * Int32(nt) + j_i) * Int32(red_stride)
+                    + (gh * Int32(8) + quad_i) * Int32(4)
+                    + e_i
+                )
+                idx1 = idx0 + Int32(16)
+                s0 = Float32(0.0)
+                s1 = Float32(0.0)
+                for wi in cutlass.range_constexpr(_WARPS):
+                    s0 = s0 + slab_f32[Int32(wi * num_tiles * nt * red_stride) + idx0]
+                    s1 = s1 + slab_f32[Int32(wi * num_tiles * nt * red_stride) + idx1]
+                m0 = t_i * Int32(_TILE_M) + gh * Int32(2) + Int32(8) * (e_i // Int32(2))
+                r = j_i * Int32(8) + Int32(2) * quad_i + e_i % Int32(2)
+                if m0 < nrows:
+                    if r < out_limit:
+                        pk[0] = s0.to(BFloat16)
+                        pk[1] = s1.to(BFloat16)
+                        out[r, n0 + m0] = pk[0]
+                        out[r, n0 + m0 + Int32(1)] = pk[1]
+                        # Every row of out is published, the zero rows past
+                        # the live ones included: the hop polls all rows of
+                        # its residual (out_rows of them).
+                        word = pk32[0]
+                        lo = word & Int32(0xFFFF)
+                        hi = (word >> Int32(16)) & Int32(0xFFFF)
+                        if lo == Int32(0x8000):
+                            lo = Int32(0)
+                        if hi == Int32(0x8000):
+                            hi = Int32(0)
+                        word = lo | (hi << Int32(16))
+                        if cutlass.const_expr(vec_publish):
+                            slab32[Int32(stage0) + r * Int32(stage_pitch) + m0 // Int32(2)] = word
+                        else:
+                            widx = buf_word0 + (
+                                (r * Int32(_NUM_RANKS) + Int32(publish_rank)) * Int32(_HIDDEN)
+                                + n0
+                                + m0
+                            ) // Int32(2)
+                            st_global_b32(mc_base + Int64(widx) * Int64(4), word)
+                p = p + Int32(_THREADS)
+            if cutlass.const_expr(vec_publish):
+                cute.arch.sync_threads()
+                if tidx < out_limit * Int32(chunks):
+                    pr = tidx // Int32(chunks)
+                    pc = tidx % Int32(chunks)
+                    sb = Int32(stage0) + pr * Int32(stage_pitch) + pc * Int32(4)
+                    vidx = buf_word0 + (
+                        (pr * Int32(_NUM_RANKS) + Int32(publish_rank)) * Int32(_HIDDEN)
+                        + n0
+                        + pc * Int32(8)
+                    ) // Int32(2)
+                    st_global_v4_b32(
+                        mc_base + Int64(vidx) * Int64(4),
+                        slab32[sb],
+                        slab32[sb + Int32(1)],
+                        slab32[sb + Int32(2)],
+                        slab32[sb + Int32(3)],
+                    )
+
+        @cute.jit
+        def __call__(
+            self,
+            v: cute.Tensor,
+            w: cute.Tensor,
+            out: cute.Tensor,
+            mc: cute.Tensor,
+            flags: cute.Tensor,
+            stream,
+        ):
+            self.kernel(v, w, out, mc, flags).launch(
+                grid=(grid, 1, 1),
+                block=(_THREADS, 1, 1),
+                stream=stream,
+                use_pdl=use_pdl,
+                min_blocks_per_mp=1,
+            )
+
+    rows_sym, out_sym = cute.sym_int(), cute.sym_int()
+    v_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (rows_sym, k_dim), stride_order=(1, 0), assumed_align=16
+    )
+    w_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (n_dim, k_dim), stride_order=(1, 0), assumed_align=16
+    )
+    out_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (out_sym, n_dim), stride_order=(1, 0), assumed_align=16
+    )
+    mc_fake = cute_runtime.make_fake_compact_tensor(
+        Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=16
+    )
+    flags_fake = cute_runtime.make_fake_compact_tensor(
+        Int32, (4,), stride_order=(0,), assumed_align=16
+    )
+    stream = cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        StreamOprojIntoHop(),
+        v_fake,
+        w_fake,
+        out_fake,
+        mc_fake,
+        flags_fake,
+        stream=stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+_COMPILED = {}
+
+
+DEFAULT_EARLY_ROWS = 16
+
+
+def stream_oproj_into_hop(
+    v: torch.Tensor,
+    w_o: torch.Tensor,
+    out: torch.Tensor,
+    use_pdl: bool,
+    mc_words: torch.Tensor,
+    flags: torch.Tensor,
+    rank: int,
+    grid: Optional[int] = None,
+    early_rows: int = DEFAULT_EARLY_ROWS,
+    trigger_dependents: bool = False,
+    vector_publish: bool = True,
+) -> None:
+    """Launch the publishing streaming o_proj on the current stream.
+
+    Arguments as ``cute_oproj.stream_oproj``, plus the post_attention_hop workspace's
+    multicast int32 view ``mc_words``, its protocol state ``flags`` and this
+    rank. Every row of ``out`` is published (rows past ``v``'s live rows as
+    zeros), so ``out`` must have exactly the row count of the hop that
+    consumes the transport, at most 8 (the transport's row capacity).
+    ``vector_publish`` allows 16-byte publish stores where the column split
+    permits them.
+    """
+    n_dim, k_dim = w_o.shape
+    if grid is None:
+        grid = _sm_count(w_o.device.index or 0)
+    if not 1 <= v.shape[0] <= out.shape[0] <= 8 or n_dim != _HIDDEN:
+        raise ValueError(
+            "the post_attention_hop transport holds 1..8 rows of 6144 and out must hold "
+            f"every live row; got v rows {v.shape[0]}, out rows "
+            f"{out.shape[0]}, N {n_dim}"
+        )
+    n_tiles_r = 1 if out.shape[0] <= 8 else 2
+    if not use_pdl:
+        early_rows = 1 << 30
+    key = (
+        n_dim,
+        k_dim,
+        grid,
+        n_tiles_r,
+        use_pdl,
+        early_rows,
+        trigger_dependents,
+        rank,
+        vector_publish,
+    )
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        compiled = _build(
+            n_dim,
+            k_dim,
+            grid,
+            n_tiles_r,
+            use_pdl,
+            early_rows,
+            trigger_dependents,
+            False,
+            rank,
+            vector_publish,
+        )
+        _COMPILED[key] = compiled
+    compiled(v, w_o, out, mc_words, flags)
+
+
+__all__ = ["check_shapes", "plan", "stream_oproj_into_hop"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/kernels.py
new file mode 100644
index 000000000000..6b12ce06c8cc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/kernels.py
@@ -0,0 +1,612 @@
+# 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.
+"""attention_output_tail (the attention_output_tail contract): kernels and launch entry point.
+
+The tail maps the latent attention output to the rank-local o_proj partial::
+
+    v[r, h, :] = bf16(latent[r, h, :] @ w_uv[h].T)  # W_UV, FP32 accumulate
+    partial[r] = bf16(v[r].flatten() @ w_o.T)  # local o_proj, FP32 accumulate
+
+Both rounding points match the native ``bmm_out`` + ``F.linear`` sequence.
+The default engine (``cute_stream``) runs two attention_output_tail-owned kernels:
+
+* ``_wuv_kernel`` (Triton) computes ``v`` into a small scratch tensor. With
+  PDL it is launched as a dependent of the latent's producer, exactly like
+  the native cuBLAS kernel it replaces: it loads its ``w_uv`` tile, then
+  executes ``griddepcontrol.wait`` before reading the latent or writing.
+  It triggers its own dependent at entry.
+* ``cute_oproj.stream_oproj`` (CuTe DSL) computes the partial. It is a PDL
+  dependent of ``_wuv_kernel``: one persistent CTA per SM streams its rows of
+  ``w_o`` into shared memory (a bounded number before
+  ``griddepcontrol.wait``, the rest after), runs ``mma.sync`` with FP32
+  accumulation, adds the per-warp partials in fixed order and stores BF16
+  once. It lets its own PDL consumer launch once all copies are issued; the
+  consumer still waits for the whole grid.
+
+``_oproj_kernel`` (Triton) and ``cute_fused`` (one CuTe DSL
+kernel for both GEMMs) remain selectable through ``TailConfig`` for
+experiments. Every engine is deterministic; rows ``[num_rows, out_rows)`` of
+the partial are written as zeros.
+"""
+
+import os
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import torch
+import triton
+import triton.language as tl
+
+ROW_TILE = 16
+"""Row tile of both kernels; it is also the largest supported row count."""
+
+_PREFETCH_SPLITS = 128
+
+
+@dataclass(frozen=True)
+class TailConfig:
+    """Launch configuration of the two tail kernels."""
+
+    wuv_block_v: int = 16
+    wuv_warps: int = 4
+    oproj_block_n: int = 16
+    oproj_block_k: int = 256
+    oproj_warps: int = 4
+    oproj_stages: int = 4
+    prefetch: bool = False
+    """Triton engine: prefetch the w_o slab into L2 before waiting (PDL)."""
+    pdl_entry: bool = True
+    """Launch the first attention_output_tail kernel as a PDL dependent of the latent's producer
+    (as the native cuBLAS kernels of this region are). Before its
+    ``griddepcontrol.wait`` it touches only immutable weights and its own
+    shared memory; every latent read and every global write follows the
+    wait. Effective when ``TRTLLM_ENABLE_PDL`` is on."""
+    early_rows: int = 16
+    """``cute_stream`` engine: ``w_o`` rows per CTA streamed before the wait
+    on the W_UV kernel."""
+    trigger_dependents: bool = True
+    """``cute_stream`` engine: let the PDL-launched consumer of the partial
+    (the all-reduce) start its prologue once all ``w_o`` copies are issued,
+    as the native cuBLAS kernels do; it still waits for completion."""
+    oproj_grid: int = 0
+    """``cute_stream`` engine: CTAs of the o_proj kernel (0: one per SM).
+    Fewer CTAs leave SMs free for a successor that needs whole SMs; each CTA
+    then holds more ``w_o`` rows in shared memory."""
+    trigger_late: bool = False
+    """``cute_stream`` engine, experiments: with ``trigger_dependents``,
+    trigger only after each CTA's MMAs."""
+    into_hop_vector_publish: bool = True
+    """``attention_output_into_hop`` publish: 16-byte multicast stores when every CTA owns a
+    multiple of 8 output columns (False: 32-bit stores only, to compare)."""
+    oproj_engine: str = "cute_stream"
+    """``"cute_stream"`` (Triton W_UV, then a CuTe DSL o_proj that streams
+    w_o to shared memory around its wait on v; default), ``"fused"`` (one
+    CuTe DSL kernel; kept for experiments), ``"triton"`` (two
+    Triton kernels) or ``"native"`` (``attention_output_tail_prefetch``: the native W_UV BMM and
+    o_proj GEMM run unchanged, on weights the ``attention_output_tail_prefetch`` prefetch staged
+    into L2; the tail runs the native projection so that the prefetch stream can
+    be joined after it)."""
+    prefetch_l2: bool = False
+    """``attention_output_tail_prefetch``: when the attention of an armed call starts, prefetch this
+    layer's ``w_uv`` and ``w_o`` into L2 from a private auxiliary stream
+    (see ``prefetch.py``)."""
+    successor_prefetch: bool = False
+    """``attention_output_tail_prefetch``: also stage the successor's weights (set with
+    ``AttnTail.set_successor_weights``, e.g. the MoE router weight that the
+    post-attention hop or the native router GEMM loads right after the
+    tail) into L2 with the same off-path prefetch. Works with or without
+    ``prefetch_l2``."""
+    prefetch_layers: str = "all"
+    """``attention_output_tail_prefetch``: layers that prefetch: ``"all"``, ``"shared"`` (layers without
+    an indexer of their own) or ``"full"`` (layers with an indexer)."""
+    native_oproj: str = "cublas"
+    """``native`` engine: the o_proj kernel run by ``project``: ``"cublas"``
+    (the native GEMM) or ``"ring"`` (``cute_oproj_ring``: ``ring_grid`` CTAs,
+    a ``ring_slots``-tile shared-memory ring, ``ring_early`` tiles requested
+    before its PDL wait)."""
+    ring_grid: int = 96
+    ring_slots: int = 3
+    ring_early: int = 3
+    prefetch_fork: str = "attn"
+    """``attention_output_tail_prefetch``: where the prefetch is forked: ``"attn"`` (right before the
+    attention backend call, i.e. after the absorption BMM and the RoPE /
+    FP8-q / KV-append kernel; the adapter's hook), ``"start"`` (start of
+    the generation attention; experiments) or ``"prev"`` (forked ahead by
+    the previous layer, right before its hop, via ``AttnTail.fork_ahead``;
+    experiments)."""
+    prefetch_join: str = "after_oproj"
+    """``attention_output_tail_prefetch``: where the caller's stream joins the prefetch stream:
+    ``"before_wuv"``, ``"before_oproj"`` or ``"after_oproj"``."""
+    prefetch_chunk: int = 4096
+    prefetch_ctas: int = 52
+    prefetch_threads: int = 32
+    prefetch_gap_ns: int = 0
+
+
+DEFAULT_CONFIG = TailConfig()
+TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE = TailConfig(oproj_engine="native", prefetch_l2=True)
+"""Experiment: the native W_UV and o_proj on tail weights staged into L2 by
+the off-path prefetch."""
+
+TAIL_PREFETCH_SCHEDULE = TailConfig(oproj_grid=128, successor_prefetch=True)
+"""``attention_output_tail_prefetch`` with the native neighbours (TRTLLM-Gen FMHA before the tail, the
+native AR + router GEMM after it): the attention_output_tail kernels with the o_proj on 128
+CTAs, plus the router weight staged into L2 for the router GEMM. See
+``tail.tail_prefetch_config`` for the other neighbour sets."""
+
+TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE = TailConfig(successor_prefetch=True)
+"""``attention_output_tail_prefetch`` behind ``sparse_mla_decode``'s sparse attention (which releases its PDL
+dependents early) with the native successor."""
+
+TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE = TailConfig(oproj_grid=128, early_rows=32)
+"""``attention_output_tail_prefetch`` behind ``sparse_mla_decode`` with the ``post_attention_hop``
+hop as successor (staging the router weight hurts next to the hop). ``sparse_mla_decode`` releases
+its dependents early enough for each CTA to stream 32 of its 48 ``w_o`` rows before its wait."""
+
+INTO_HOP_SCHEDULE = TailConfig(oproj_grid=128)
+"""``attention_output_into_hop``: the attention_output_tail kernels with the o_proj on 128 CTAs (48 output
+columns each, so every published row chunk is a whole number of 16-byte
+stores, and 24 SMs stay free for the hop's CTAs); when the decoder layer
+requests it, the o_proj also publishes its partial into the post_attention_hop's
+transport. Used behind the native FMHA (streaming more rows before the
+wait loses there)."""
+
+INTO_HOP_SPARSE_DECODE_SCHEDULE = TailConfig(oproj_grid=128, early_rows=32)
+"""``attention_output_into_hop`` behind ``sparse_mla_decode``: as ``INTO_HOP_SCHEDULE`` with 32 of each CTA's 48 ``w_o``
+rows streamed before the PDL wait (as ``TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE``)."""
+# Shape checks use the coarsest tiles any config may use.
+_MAX_BLOCK_C = 512
+_MAX_BLOCK_V = 64
+_MAX_BLOCK_N = 64
+_MAX_BLOCK_K = 512
+
+
+def pdl_enabled() -> bool:
+    """Return whether attention_output_tail may use programmatic dependent launch at all."""
+    return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1"
+
+
+def pdl_entry_enabled() -> bool:
+    """Return whether the first attention_output_tail kernel is a PDL dependent of its producer
+    (whenever PDL is enabled)."""
+    return pdl_enabled()
+
+
+@triton.jit
+def _gdc_wait_with_memory_clobber():
+    tl.inline_asm_elementwise(
+        "griddepcontrol.wait; // dummy $0",
+        "=r,~{memory}",
+        [],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _prefetch_l2(addr, nbytes):
+    tl.inline_asm_elementwise(
+        "cp.async.bulk.prefetch.L2.global [$1], $2; // dummy $0",
+        "=r,l,r",
+        [addr, nbytes],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _wuv_kernel(
+    latent_ptr,
+    w_uv_ptr,
+    v_ptr,
+    num_rows,
+    stride_lat_row,
+    stride_lat_head,
+    stride_wuv_head,
+    stride_wuv_v,
+    stride_v_row,
+    C: tl.constexpr,
+    V: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    WAIT_PRODUCER: tl.constexpr,
+):
+    head = tl.program_id(0)
+    v_block = tl.program_id(1)
+    if USE_PDL:
+        # Let the dependent o_proj kernel start its weight stream now; it
+        # still waits for this whole grid before reading v.
+        tl.extra.cuda.gdc_launch_dependents()
+
+    offs_v = v_block * BLOCK_V + tl.arange(0, BLOCK_V)
+    offs_r = tl.arange(0, BLOCK_R)
+    offs_c = tl.arange(0, C)
+    row_mask = offs_r < num_rows
+    # Weights are immutable: load them before waiting for the producer.
+    w = tl.load(
+        w_uv_ptr + head * stride_wuv_head + offs_v[:, None] * stride_wuv_v + offs_c[None, :]
+    )
+    if WAIT_PRODUCER:
+        # Launched as a PDL dependent of the latent's producer: every read of
+        # the latent and every global write follows this wait.
+        _gdc_wait_with_memory_clobber()
+    lat_t = tl.load(
+        latent_ptr + offs_r[None, :] * stride_lat_row + head * stride_lat_head + offs_c[:, None],
+        mask=row_mask[None, :],
+        other=0.0,
+    )
+    acc = tl.dot(w, lat_t)
+    v_out = acc.to(v_ptr.dtype.element_ty)
+    tl.store(
+        v_ptr + offs_r[None, :] * stride_v_row + head * V + offs_v[:, None],
+        v_out,
+        mask=row_mask[None, :],
+    )
+
+
+@triton.jit
+def _oproj_kernel(
+    v_ptr,
+    w_o_ptr,
+    w_o_addr,
+    out_ptr,
+    num_rows,
+    out_rows,
+    stride_v_row,
+    stride_out_row,
+    K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    PREFETCH_SPLITS: tl.constexpr,
+    PREFETCH: tl.constexpr,
+    USE_PDL: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    if PREFETCH:
+        # The [BLOCK_N, K] slab of w_o is contiguous: split it into equal
+        # pieces, one per lane, and pull them into L2 while W_UV runs.
+        slab_bytes: tl.constexpr = BLOCK_N * K * 2
+        piece: tl.constexpr = slab_bytes // PREFETCH_SPLITS
+        offs_p = tl.arange(0, PREFETCH_SPLITS).to(tl.int64)
+        addr = w_o_addr + pid.to(tl.int64) * slab_bytes + offs_p * piece
+        _prefetch_l2(addr, tl.full([PREFETCH_SPLITS], piece, tl.int32))
+    if USE_PDL:
+        _gdc_wait_with_memory_clobber()
+
+    offs_r = tl.arange(0, BLOCK_R)
+    acc = tl.zeros([BLOCK_N, BLOCK_R], dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        offs_k = k0 + tl.arange(0, BLOCK_K)
+        w = tl.load(w_o_ptr + offs_n[:, None] * K + offs_k[None, :])
+        v_t = tl.load(
+            v_ptr + offs_r[None, :] * stride_v_row + offs_k[:, None],
+            mask=offs_r[None, :] < num_rows,
+            other=0.0,
+        )
+        acc = tl.dot(w, v_t, acc)
+    tl.store(
+        out_ptr + offs_r[None, :] * stride_out_row + offs_n[:, None],
+        acc.to(out_ptr.dtype.element_ty),
+        mask=offs_r[None, :] < out_rows,
+    )
+
+
+def check_weights(w_uv: torch.Tensor, w_o: torch.Tensor) -> Optional[str]:
+    """Return why the kernels cannot take these weights, or ``None``.
+
+    Args:
+        w_uv: ``[H, V, C]`` BF16 with a unit inner stride.
+        w_o: ``[N, H * V]`` contiguous BF16.
+    """
+    if w_uv.dim() != 3 or w_o.dim() != 2:
+        return "unexpected weight rank"
+    heads, v_dim, c_dim = w_uv.shape
+    n_dim, k_dim = w_o.shape
+    if w_uv.dtype != torch.bfloat16 or w_o.dtype != torch.bfloat16:
+        return "weight dtype is not bfloat16"
+    if not (w_uv.is_cuda and w_o.is_cuda):
+        return "weight not on CUDA"
+    if k_dim != heads * v_dim:
+        return "w_o in_features != heads * v_head_dim"
+    if w_uv.stride(2) != 1 or not w_o.is_contiguous():
+        return "weights not packed"
+    if c_dim & (c_dim - 1) or c_dim > _MAX_BLOCK_C or v_dim % _MAX_BLOCK_V:
+        return "W_UV shape not supported by the tiles"
+    if k_dim % _MAX_BLOCK_K or n_dim % _MAX_BLOCK_N:
+        return "o_proj shape not a tile multiple"
+    if (16 * k_dim * 2) % (_PREFETCH_SPLITS * 16):
+        return "prefetch slab not a 16-byte multiple per lane"
+    if w_o.data_ptr() % 16:
+        return "w_o is not 16-byte aligned"
+    from .cute_oproj import _sm_count
+    from .cute_oproj import check_shapes as check_stream_shapes
+
+    grid = _sm_count(w_o.device.index or 0)
+    reason = check_stream_shapes(n_dim, k_dim, grid)
+    if reason is not None:
+        return reason
+    from .cute_fused import check_fused_shapes
+
+    reason = check_fused_shapes(heads, v_dim, c_dim, n_dim, grid)
+    if reason is not None:
+        return reason
+    if not w_uv.is_contiguous():
+        return "w_uv is not contiguous"
+    return None
+
+
+def check_tail_shapes(
+    latent: torch.Tensor, w_uv: torch.Tensor, w_o: torch.Tensor, num_rows: int, out_rows: int
+) -> Optional[str]:
+    """Return why the kernels cannot take these tensors, or ``None``.
+
+    Args:
+        latent: ``[>= num_rows, H, C]`` or ``[>= num_rows, H * C]`` BF16 with
+            a unit inner stride.
+        w_uv: ``[H, V, C]`` BF16 with a unit inner stride.
+        w_o: ``[N, H * V]`` contiguous BF16.
+        num_rows: live rows to compute, at most ``ROW_TILE``.
+        out_rows: rows of the partial; rows past ``num_rows`` are zeroed.
+    """
+    reason = check_weights(w_uv, w_o)
+    if reason is not None:
+        return reason
+    heads, _, c_dim = w_uv.shape
+    if latent.dtype != torch.bfloat16 or not latent.is_cuda:
+        return "latent is not a CUDA bfloat16 tensor"
+    if not 1 <= num_rows <= ROW_TILE or out_rows < num_rows:
+        return f"rows {num_rows}/{out_rows} outside [1, {ROW_TILE}]"
+    if latent.shape[0] < num_rows:
+        return "latent has fewer rows than num_rows"
+    if latent.dim() == 2:
+        if latent.shape[1] != heads * c_dim or latent.stride(1) != 1:
+            return "latent is not [R, H * C] with unit inner stride"
+    elif latent.dim() == 3:
+        if (
+            tuple(latent.shape[1:]) != (heads, c_dim)
+            or latent.stride(2) != 1
+            or latent.stride(1) != c_dim
+        ):
+            return "latent is not [R, H, C] with packed heads"
+    else:
+        return "unexpected latent rank"
+    if latent.stride(0) != heads * c_dim or latent.data_ptr() % 16:
+        return "latent rows are not packed and 16-byte aligned"
+    return None
+
+
+def launch_wuv(
+    latent: torch.Tensor,
+    w_uv: torch.Tensor,
+    v: torch.Tensor,
+    num_rows: int,
+    config: TailConfig = DEFAULT_CONFIG,
+    trigger_dependents: bool = False,
+    wait_producer: bool = False,
+) -> None:
+    """Launch the W_UV kernel: ``v[:num_rows] = bf16(latent @ w_uv^T)`` per head.
+
+    ``wait_producer`` launches it as a PDL dependent of the previous kernel
+    on the stream; it loads ``w_uv`` first and executes
+    ``griddepcontrol.wait`` before reading the latent or writing ``v``.
+    """
+    heads, v_dim, c_dim = w_uv.shape
+    if latent.dim() == 2:
+        stride_lat_row, stride_lat_head = latent.stride(0), c_dim
+    else:
+        stride_lat_row, stride_lat_head = latent.stride(0), latent.stride(1)
+    _wuv_kernel[(heads, v_dim // config.wuv_block_v)](
+        latent,
+        w_uv,
+        v,
+        num_rows,
+        stride_lat_row,
+        stride_lat_head,
+        w_uv.stride(0),
+        w_uv.stride(1),
+        v.stride(0),
+        C=c_dim,
+        V=v_dim,
+        BLOCK_V=config.wuv_block_v,
+        BLOCK_R=ROW_TILE,
+        USE_PDL=trigger_dependents,
+        WAIT_PRODUCER=wait_producer,
+        num_warps=config.wuv_warps,
+        num_stages=1,
+        **({"launch_pdl": True} if wait_producer else {}),
+    )
+
+
+def launch_oproj(
+    v: torch.Tensor,
+    w_o: torch.Tensor,
+    out: torch.Tensor,
+    num_rows: int,
+    out_rows: int,
+    config: TailConfig = DEFAULT_CONFIG,
+    use_pdl: bool = False,
+    prefetch: Optional[bool] = None,
+) -> None:
+    """Launch the o_proj kernel: ``out[:out_rows] = bf16(v @ w_o^T)``.
+
+    ``use_pdl`` launches it as a programmatic dependent of the previous
+    kernel on the stream; ``prefetch`` (default: ``config.prefetch`` when
+    ``use_pdl``) issues the L2 prefetch of its ``w_o`` slab first.
+    """
+    n_dim, k_dim = w_o.shape
+    if prefetch is None:
+        prefetch = config.prefetch and use_pdl
+    launch_kwargs = {"launch_pdl": True} if use_pdl else {}
+    _oproj_kernel[(n_dim // config.oproj_block_n,)](
+        v,
+        w_o,
+        w_o.data_ptr(),
+        out,
+        num_rows,
+        out_rows,
+        v.stride(0),
+        out.stride(0),
+        K=k_dim,
+        BLOCK_N=config.oproj_block_n,
+        BLOCK_K=config.oproj_block_k,
+        BLOCK_R=ROW_TILE,
+        PREFETCH_SPLITS=_PREFETCH_SPLITS,
+        PREFETCH=prefetch,
+        USE_PDL=use_pdl,
+        num_warps=config.oproj_warps,
+        num_stages=config.oproj_stages,
+        **launch_kwargs,
+    )
+
+
+_DEFAULT_STATES = {}
+
+
+def default_state(device: torch.device) -> torch.Tensor:
+    """Per-device fused-kernel counters that a benchmark passes explicitly.
+
+    Only for strictly sequential callers on one stream that opt in by passing
+    it as ``state``; ``attn_tail`` never falls back to it. The MLA adapter
+    owns one state per layer instead.
+    """
+    from .cute_fused import new_state
+
+    key = device.index or 0
+    if key not in _DEFAULT_STATES:
+        _DEFAULT_STATES[key] = new_state(device)
+    return _DEFAULT_STATES[key]
+
+
+def attn_tail(
+    latent: torch.Tensor,
+    w_uv: torch.Tensor,
+    w_o: torch.Tensor,
+    num_rows: int,
+    out_rows: Optional[int] = None,
+    use_pdl: Optional[bool] = None,
+    v_scratch: Optional[torch.Tensor] = None,
+    config: TailConfig = DEFAULT_CONFIG,
+    state: Optional[torch.Tensor] = None,
+    between: Optional[Callable[[], None]] = None,
+    hop_transport=None,
+) -> torch.Tensor:
+    """Compute the rank-local o_proj partial from the latent attention output.
+
+    The caller validates the arguments with ``check_tail_shapes`` first.
+
+    Args:
+        latent: ``[>= num_rows, H, C]`` or ``[>= num_rows, H * C]`` BF16.
+        w_uv: ``[H, V, C]`` BF16 (the module's ``v_b_proj``).
+        w_o: ``[N, H * V]`` BF16 (the local ``o_proj`` weight).
+        num_rows: live rows.
+        out_rows: rows of the returned partial (default ``num_rows``); rows
+            past ``num_rows`` are zero.
+        use_pdl: launch the o_proj kernel as a PDL dependent of the W_UV
+            kernel (default: ``TRTLLM_ENABLE_PDL``).
+        v_scratch: optional ``[num_rows, H * V]`` BF16 buffer for the
+            intermediate (tests read it back); allocated per call otherwise.
+        config: kernel launch configuration.
+        state: fused-engine counters (``cute_fused.new_state``) owned by the
+            caller. When omitted, the fused engine gets zeroed counters
+            allocated for this call only (an invocation-local packet, never
+            module-global storage).
+        between: optional host callback run after the W_UV launch and before
+            the o_proj launch (two-kernel engines only).
+        hop_transport: the ``IntoHopTransport`` of ``attention_output_into_hop``; the
+            ``cute_stream`` o_proj then also publishes all ``out_rows`` rows of the partial
+            into it, the zero rows past ``num_rows`` included (``cute_oproj_into_hop``; at
+            most 8).
+
+    Returns:
+        ``[out_rows, N]`` BF16 partial on ``latent``'s device.
+    """
+    if out_rows is None:
+        out_rows = num_rows
+    if use_pdl is None:
+        use_pdl = pdl_enabled()
+    n_dim, k_dim = w_o.shape
+    v = v_scratch
+    if v is None:
+        v = torch.empty((num_rows, k_dim), dtype=w_uv.dtype, device=latent.device)
+    out = torch.empty((out_rows, n_dim), dtype=w_o.dtype, device=latent.device)
+    if out_rows > ROW_TILE:
+        out[ROW_TILE:].zero_()
+    if config.oproj_engine == "fused":
+        from .cute_fused import fused_tail, new_state
+
+        if state is None:
+            state = new_state(latent.device)
+        fused_tail(
+            latent[:num_rows].reshape(num_rows, -1),
+            w_uv,
+            w_o,
+            out,
+            state,
+            v,
+            pdl_entry=config.pdl_entry and pdl_entry_enabled(),
+        )
+        return out
+    launch_wuv(
+        latent,
+        w_uv,
+        v,
+        num_rows,
+        config,
+        trigger_dependents=use_pdl,
+        wait_producer=config.pdl_entry and pdl_entry_enabled(),
+    )
+    if between is not None:
+        between()
+    if config.oproj_engine == "cute_stream" and hop_transport is not None:
+        from .into_hop import stream_oproj_into_hop
+
+        stream_oproj_into_hop(
+            v,
+            w_o,
+            out,
+            use_pdl=use_pdl,
+            mc_words=hop_transport.mc_words,
+            flags=hop_transport.flags,
+            rank=hop_transport.rank,
+            grid=config.oproj_grid or None,
+            early_rows=config.early_rows,
+            trigger_dependents=config.trigger_dependents and pdl_enabled(),
+            vector_publish=config.into_hop_vector_publish,
+        )
+    elif config.oproj_engine == "cute_stream":
+        from .cute_oproj import check_shapes as stream_shapes
+        from .cute_oproj import stream_oproj
+
+        grid = config.oproj_grid or None
+        if grid is not None and stream_shapes(n_dim, k_dim, grid):
+            raise ValueError(f"oproj_grid={grid}: {stream_shapes(n_dim, k_dim, grid)}")
+        stream_oproj(
+            v,
+            w_o,
+            out,
+            use_pdl=use_pdl,
+            grid=grid,
+            early_rows=config.early_rows,
+            trigger_dependents=config.trigger_dependents and pdl_enabled(),
+            trigger_late=config.trigger_late,
+        )
+    else:
+        launch_oproj(v, w_o, out, num_rows, out_rows, config, use_pdl=use_pdl)
+    return out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/prefetch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/prefetch.py
new file mode 100644
index 000000000000..a987faa22135
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/prefetch.py
@@ -0,0 +1,217 @@
+# 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.
+"""Off-path L2 prefetch of the attention_output_tail weights (``attention_output_tail_prefetch``).
+
+At M = 6 the tail (W_UV + local o_proj) is a 26 MiB weight stream (``w_o``
+24 MiB, ``w_uv`` 2 MiB per rank) that starts only after the attention kernel
+has produced the latent. The attention region before it (absorption BMM,
+RoPE / FP8 quantization / cache append, top-k bridge, sparse attention) is
+latency-bound and leaves HBM mostly idle, and it lasts longer than the
+stream. ``WeightPrefetch`` uses that window: when the attention of an armed
+layer starts, it forks a private auxiliary stream from the caller's stream,
+launches a small kernel there that issues ``cp.async.bulk.prefetch.L2`` over
+the tail's weights, and records an event that the caller's stream waits on
+later in the same layer (``join``). The tail kernels then read their weights
+from L2.
+
+A prefetch has no architectural effect: it neither reads nor writes any
+tensor the model observes, so it needs no producer/consumer protocol and no
+memory-order contract with the attention kernels, and the results are
+identical with and without it. The fork and the join only add graph edges;
+``join`` must run before the layer's forward returns so that a CUDA-graph
+capture always sees the auxiliary stream rejoin the capturing stream.
+
+The kernel is a fire-and-forget issuer: ``ctas`` CTAs of ``threads`` threads
+each issue one ``chunk``-byte prefetch per iteration over the concatenated
+weight ranges, optionally pausing ``gap_ns`` between iterations to pace the
+traffic.
+"""
+
+import functools
+from dataclasses import dataclass
+from typing import List, Sequence
+
+import torch
+import triton
+import triton.language as tl
+
+
+@dataclass(frozen=True)
+class PrefetchConfig:
+    """Launch parameters of the L2 prefetch kernel."""
+
+    chunk: int = 4096
+    """Bytes per ``cp.async.bulk.prefetch.L2`` instruction (multiple of 16)."""
+    ctas: int = 52
+    """CTAs of the prefetch kernel."""
+    threads: int = 32
+    """Threads per CTA (each issues one prefetch per iteration)."""
+    gap_ns: int = 0
+    """Pause between iterations of one thread (0: issue back to back)."""
+
+
+@triton.jit
+def _bulk_prefetch_l2(addr, nbytes):
+    tl.inline_asm_elementwise(
+        "cp.async.bulk.prefetch.L2.global [$1], $2; // dummy $0",
+        "=r,l,r",
+        [addr, nbytes],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _nanosleep(ns):
+    tl.inline_asm_elementwise(
+        "nanosleep.u32 $1; // dummy $0",
+        "=r,r",
+        [ns],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _l2_prefetch_kernel(
+    addr0,
+    chunks0,
+    addr1,
+    chunks1,
+    addr2,
+    chunks2,
+    CHUNK: tl.constexpr,
+    NPROG: tl.constexpr,
+    NT: tl.constexpr,
+    ITERS: tl.constexpr,
+    GAP_NS: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    lanes = tl.arange(0, NT).to(tl.int64)
+    total = chunks0 + chunks1 + chunks2
+    for it in tl.static_range(ITERS):
+        c = (it * NPROG + pid) * NT + lanes
+        # Lanes past the end re-issue chunk 0, which is harmless.
+        c = tl.where(c < total, c, 0)
+        c12 = c - chunks0
+        addr = tl.where(
+            c < chunks0,
+            addr0 + c * CHUNK,
+            tl.where(c12 < chunks1, addr1 + c12 * CHUNK, addr2 + (c12 - chunks1) * CHUNK),
+        )
+        _bulk_prefetch_l2(addr, tl.full([NT], CHUNK, tl.int32))
+        if GAP_NS > 0:
+            _nanosleep(tl.full([NT], GAP_NS, tl.int32))
+
+
+def check_prefetch_tensors(tensors: Sequence[torch.Tensor], chunk: int) -> str:
+    """Return why ``tensors`` cannot be prefetched in ``chunk`` pieces, or ''."""
+    if not 1 <= len(tensors) <= 3:
+        return "expects one to three weight tensors"
+    if chunk % 16 or chunk <= 0:
+        return "chunk is not a positive multiple of 16 bytes"
+    for t in tensors:
+        if not t.is_cuda or not t.is_contiguous():
+            return "weight is not a contiguous CUDA tensor"
+        nbytes = t.numel() * t.element_size()
+        if t.data_ptr() % 16 or nbytes % chunk:
+            return "weight size or alignment is not a chunk multiple"
+    return ""
+
+
+def launch_prefetch(
+    tensors: Sequence[torch.Tensor], cfg: PrefetchConfig = PrefetchConfig()
+) -> None:
+    """Issue L2 prefetches of one to three contiguous tensors on the current
+    stream."""
+    regions = [(t.data_ptr(), t.numel() * t.element_size() // cfg.chunk) for t in tensors]
+    regions += [(regions[0][0], 0)] * (3 - len(regions))
+    per_iter = cfg.ctas * cfg.threads
+    iters = -(-sum(n for _, n in regions) // per_iter)
+    _l2_prefetch_kernel[(cfg.ctas,)](
+        regions[0][0],
+        regions[0][1],
+        regions[1][0],
+        regions[1][1],
+        regions[2][0],
+        regions[2][1],
+        CHUNK=cfg.chunk,
+        NPROG=cfg.ctas,
+        NT=cfg.threads,
+        ITERS=iters,
+        GAP_NS=cfg.gap_ns,
+        num_warps=max(1, cfg.threads // 32),
+        num_stages=1,
+    )
+
+
+@functools.lru_cache(maxsize=None)
+def _aux_stream(device_index: int) -> torch.cuda.Stream:
+    # One private stream per device for the package; every fork is joined
+    # within the same layer, so layers never overlap on it.
+    return torch.cuda.Stream(device=device_index)
+
+
+class WeightPrefetch:
+    """Fork/join of the off-path weight prefetch of one tail (layer, rank).
+
+    Owns its own fork and completion events, so no state is shared between
+    layers beyond the device's auxiliary stream.
+    """
+
+    def __init__(self, cfg: PrefetchConfig = PrefetchConfig()):
+        self.cfg = cfg
+        self._fork_event = torch.cuda.Event()
+        self._done_event = torch.cuda.Event()
+        self._forked = False
+        self.forks = 0
+
+    @property
+    def forked(self) -> bool:
+        return self._forked
+
+    def fork(self, tensors: List[torch.Tensor]) -> None:
+        """Launch the prefetch of ``tensors`` on the auxiliary stream."""
+        if self._forked:
+            raise RuntimeError(
+                "[glm_moe_dsa] attention_output_tail_prefetch: prefetch forked twice without a join"
+            )
+        caller = torch.cuda.current_stream()
+        aux = _aux_stream(caller.device.index)
+        self._fork_event.record(caller)
+        aux.wait_event(self._fork_event)
+        with torch.cuda.stream(aux):
+            launch_prefetch(tensors, self.cfg)
+            self._done_event.record(aux)
+        self._forked = True
+        self.forks += 1
+
+    def join(self) -> None:
+        """Make the caller's stream wait for the prefetch kernel (no-op if
+        nothing was forked)."""
+        if not self._forked:
+            return
+        torch.cuda.current_stream().wait_event(self._done_event)
+        self._forked = False
+
+
+__all__: List[str] = [
+    "PrefetchConfig",
+    "WeightPrefetch",
+    "check_prefetch_tensors",
+    "launch_prefetch",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail.py
new file mode 100644
index 000000000000..cff8a4259b39
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail.py
@@ -0,0 +1,628 @@
+# 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.
+"""Per-layer attention_output_tail adapter between the MLA generation tail and ``o_proj``.
+
+The native GEN tail is split across two call sites of one ``MLA.forward``:
+``forward_absorption_generation`` applies W_UV inside the attention custom op
+and ``_project_output`` applies ``o_proj``. ``AttnTail`` joins them in three
+steps:
+
+1. ``arm`` runs in ``MLA.forward`` before the attention op. It is the only
+   place that decides whether this call may take the attention_output_tail, so the W_UV
+   step inside the op and the projection step after it always see the same
+   decision. It declines, without touching any state, while Dynamo traces the
+   module, under TensorRT-LLM's torch.compile policy (``is_torch_compiling``)
+   and inside a breakable CUDA graph: in those modes the attention op runs
+   separately from the traced or captured projection site, so only the
+   native path is consistent.
+2. ``defer_w_uv`` runs at the W_UV step inside the op. It acts only on an
+   armed call and checks the latent's shape before any GPU work. When it
+   accepts, it keeps the native latent tensor in a packet tagged with the
+   output buffer the native W_UV would have written, and the caller skips the
+   native BMM. When it declines, nothing changes.
+3. ``project`` runs before the native projection of an armed call. It
+   consumes the packet, launches the tail and returns ``o_proj``'s rank-local
+   partial before any communication (the attention_output_tail contract). The caller reduces it as
+   the native ROW-parallel ``Linear.forward`` reduces this ``o_proj``'s GEMM
+   output. With no packet it returns ``None`` and the caller runs the native
+   projection.
+
+With ``attention_output_tail_prefetch`` (``TailConfig.prefetch_l2``) there is one more step between
+1 and 2: ``before_attention`` runs in the op's generation attention
+(``forward_absorption_generation``, right before the attention backend call)
+of an armed call and forks the off-path L2 prefetch of this layer's tail
+weights (``prefetch.py``). The prefetch stream is joined in ``project`` (at
+the configured point), which an armed call always reaches, so a CUDA-graph
+capture always sees the fork rejoined within the layer. A prefetch changes
+no tensor; outputs are identical with and without it. With the ``native``
+engine (``attention_output_tail_prefetch`` experiments) the native W_UV runs in the op and
+``project`` runs the native o_proj GEMM itself, then joins.
+
+No path of this module communicates: the all-reduce of the partial, when
+there is one, belongs to the caller.
+
+The packet belongs to one module instance and one forward call. A packet that
+is left over, or that does not match the projection's input buffer, raises
+instead of producing a wrong result.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch.utils import is_torch_compiling
+
+from ... import dispatch
+from .kernels import (
+    DEFAULT_CONFIG,
+    INTO_HOP_SCHEDULE,
+    INTO_HOP_SPARSE_DECODE_SCHEDULE,
+    ROW_TILE,
+    TAIL_PREFETCH_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+    TailConfig,
+    attn_tail,
+    check_tail_shapes,
+    check_weights,
+)
+from .prefetch import PrefetchConfig, WeightPrefetch, check_prefetch_tensors
+
+FEATURE_ID = "attention_output_tail"
+TAIL_PREFETCH_PATH = "attention_output_tail_prefetch"
+INTO_HOP_PATH = "attention_output_into_hop"
+_JOIN_POINTS = ("before_wuv", "before_oproj", "after_oproj")
+_PREFETCH_LAYERS = ("all", "shared", "full")
+MAX_ENGAGED_ROWS = 8
+"""Largest row count attention_output_tail takes: the kernels support 16, but above 8 rows the
+tail needs a second n8 MMA tile and more ``v`` traffic, and the native tail runs."""
+SUPPORTED_TP_SIZE = 8
+"""o_proj tensor-parallel width attention_output_tail is verified for (the GEN role's TP8)."""
+ENGAGED_DOMAIN = (
+    f"GEN-only calls of 1-{MAX_ENGAGED_ROWS} token rows, "
+    f"TP{SUPPORTED_TP_SIZE}, no attention DP; any draft length "
+    "and request mix within that row count"
+)
+"""Per-call domain logged with an engaged capability decision. The tail is
+row-wise, so the draft length and the request mix of a GEN batch do not change
+its arithmetic; only the row count, GEN-only and the static checks gate it."""
+MAX_INTO_HOP_ROWS = 8
+"""Row capacity of the post_attention_hop transport (``attention_output_into_hop``): output rows, padding
+included."""
+
+
+@dataclass
+class _Packet:
+    latent: torch.Tensor
+    output_ptr: int
+    num_rows: int
+
+
+def _static_reason(mla) -> Optional[str]:
+    """Return why ``mla`` can never take the attention_output_tail, or ``None``."""
+    from tensorrt_llm._torch.attention.backends.sparse.hooks import MLASparseHooks
+    from tensorrt_llm._torch.attention.mla import MLA
+    from tensorrt_llm._torch.modules.linear import TensorParallelMode, UnquantizedLinearMethod
+    from tensorrt_llm._utils import get_sm_version
+
+    if get_sm_version() not in (100, 103):
+        return f"SM{get_sm_version()} is not SM100/SM103"
+    if mla.mapping.enable_attention_dp:
+        return "attention data parallelism"
+    o_tp_size = getattr(mla, "mapping_o", mla.mapping).tp_size
+    if o_tp_size != SUPPORTED_TP_SIZE:
+        return f"o_proj tensor-parallel size {o_tp_size} is not {SUPPORTED_TP_SIZE}"
+    if mla.mapping.has_cp_helix() or mla.num_heads_tp != mla.num_heads_tp_cp:
+        return "helix context parallelism"
+    if type(mla)._apply_output_gate is not MLA._apply_output_gate:
+        return "module overrides the output gate"
+    hooks = mla.sparse_attn_hooks
+    if hooks is not None and (type(hooks).project_output is not MLASparseHooks.project_output):
+        return "sparse hooks own the output projection"
+    o_proj = getattr(mla, "o_proj", None)
+    w_uv = getattr(mla, "v_b_proj", None)
+    if o_proj is None or w_uv is None:
+        return "module has no default o_proj or v_b_proj"
+    # The FP8 linear methods subclass UnquantizedLinearMethod.
+    if type(o_proj.quant_method) is not UnquantizedLinearMethod:
+        return "o_proj is quantized"
+    if o_proj.tp_mode != TensorParallelMode.ROW:
+        return "o_proj is not ROW parallel"
+    if o_proj.bias is not None or o_proj.lora is not None:
+        return "o_proj has bias or LoRA"
+    if o_proj.use_fused_gemm_allreduce:
+        return "o_proj uses the fused GEMM+all-reduce"
+    # The native GEMM writes its output into an NCCL symmetric-memory window
+    # only through the custom cuBLAS path (UnquantizedLinearMethod.apply); the
+    # default F.linear output is an ordinary buffer, which attention_output_tail matches.
+    if (
+        getattr(o_proj, "use_custom_cublas_mm", False)
+        and o_proj.all_reduce is not None
+        and o_proj.all_reduce.uses_nccl_symmetric_memory_window()
+    ):
+        return "o_proj writes into an NCCL window buffer"
+    return check_weights(w_uv, o_proj.weight)
+
+
+class AttnTail:
+    """attention_output_tail state of one MLA module (one layer, one rank).
+
+    Args:
+        layer_key: name used in capability logs, e.g. ``"layer3"``.
+        config: kernel launch configuration (benchmarks select ablations).
+    """
+
+    def __init__(self, layer_key: str, config: TailConfig = DEFAULT_CONFIG):
+        self.layer_key = layer_key
+        self.config = config
+        self.engaged_calls = 0
+        self.declined_calls = 0
+        self._static_ok: Optional[bool] = None
+        self.static_reason: Optional[str] = None
+        """Why this layer never takes attention_output_tail (``None`` if it can); set on the
+        first ``arm``."""
+        self._armed = False
+        self._pending: Optional[_Packet] = None
+        # attention_output_tail_prefetch: off-path weight prefetch, created on first use.
+        self._prefetch: Optional[WeightPrefetch] = None
+        self.prefetch_forks = 0
+        # True while a prefetch forked ahead of this layer's call (by the
+        # previous layer's hook, prefetch_fork="prev") awaits the call.
+        self._ahead = False
+        # Weights of the op that follows the tail (successor_prefetch).
+        self._successor_weights = []
+        # attention_output_into_hop: (transport, hop rows) requested by the decoder layer for
+        # the next call, the request this call may honour, and whether it
+        # published.
+        self._into_hop_request = None
+        self._into_hop = None
+        self._into_hop_published = False
+        self.into_hop_published_calls = 0
+        # Requests that arm accepted but project did not publish (output
+        # rows differ from the hop's rows or exceed the transport).
+        self.into_hop_declined_calls = 0
+        # Counters of the (non-default) fused engine, private to this layer.
+        # Allocated on the first eager call so a CUDA graph never captures
+        # their zeroing.
+        self._state: Optional[torch.Tensor] = None
+
+    def _check_static(self, mla) -> bool:
+        if self._static_ok is None:
+            reason = _static_reason(mla)
+            self.static_reason = reason
+            self._static_ok = dispatch.log_capability(
+                FEATURE_ID, self.layer_key, reason is None, reason or ENGAGED_DOMAIN
+            )
+        return self._static_ok
+
+    def arm(self, mla, attn_metadata) -> bool:
+        """Decide, before the attention op runs, whether this call may use attention_output_tail.
+
+        Called from ``MLA.forward``. Returns False, and touches no state, while
+        Dynamo traces the module or TensorRT-LLM's torch.compile policy is on;
+        both call sites then take the native path. Otherwise records the
+        decision for ``defer_w_uv`` and returns it.
+
+        Args:
+            mla: the calling ``MLA`` module.
+            attn_metadata: the call's attention metadata.
+        """
+        if torch.compiler.is_compiling() or is_torch_compiling():
+            return False
+        from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph
+
+        if self._pending is not None or (
+            self._prefetch is not None and self._prefetch.forked and not self._ahead
+        ):
+            raise RuntimeError(
+                f"[glm_moe_dsa] attention_output_tail/{self.layer_key}: previous tail packet or "
+                "prefetch was not consumed by the output projection"
+            )
+        num_rows = attn_metadata.num_tokens
+        armed = (
+            self._check_static(mla)
+            and not is_in_breakable_cuda_graph()
+            and attn_metadata.num_contexts == 0
+            and 1 <= num_rows <= min(ROW_TILE, MAX_ENGAGED_ROWS)
+        )
+        if armed and self.config.oproj_engine == "fused" and self._state is None:
+            if torch.cuda.is_current_stream_capturing():
+                armed = False
+            else:
+                from .cute_fused import new_state
+
+                self._state = new_state(torch.device("cuda"))
+        self._armed = armed
+        # attention_output_into_hop: a call that takes the attention_output_tail kernels can also
+        # publish its partial into the post_attention_hop transport; the request is consumed either
+        # way.
+        request, self._into_hop_request = self._into_hop_request, None
+        self._into_hop = (
+            request
+            if armed and request is not None and self.config.oproj_engine == "cute_stream"
+            else None
+        )
+        self._into_hop_published = False
+        if self._ahead and not armed:
+            # A prefetch forked ahead of a call that takes the native path:
+            # join it right away so the fork never outlives the layer.
+            self._join_prefetch()
+        if not armed:
+            self.declined_calls += 1
+        return armed
+
+    def request_into_hop(self, transport, rows: int) -> None:
+        """``attention_output_into_hop``: ask the next call to publish its partial into
+        ``transport`` (the ``IntoHopTransport`` of the layer's post_attention_hop).
+
+        Called by the decoder layer before the attention, from rank-uniform
+        facts. ``rows`` is the row count the hop will poll (the rows of its
+        residual). The call honours the request only if ``arm`` accepts the
+        call, the attention_output_tail kernels run and the projection's output has exactly
+        ``rows`` rows (live rows plus padding, at most ``MAX_INTO_HOP_ROWS``);
+        otherwise nothing is published and the layer runs the post_attention_hop.
+        ``take_into_hop_published`` tells the layer which one happened.
+        """
+        self._into_hop_request = (transport, int(rows))
+
+    def take_into_hop_published(self) -> bool:
+        """``attention_output_into_hop``: whether the last call published its partial into the
+        requested transport; if so the layer must run the into-hop entry next.
+
+        Also drops a request the call did not consume (``arm`` returns early
+        under torch.compile before it reads the request), so a request never
+        carries over to a later call.
+        """
+        published, self._into_hop_published = self._into_hop_published, False
+        self._into_hop_request = None
+        return published
+
+    def set_successor_weights(self, weights) -> None:
+        """``attention_output_tail_prefetch``: weights the op after the tail reads first (staged into L2
+        with ``TailConfig.successor_prefetch``). Called once by the model
+        adapter; the tensors are read at fork time, so in-place weight
+        loading after construction is fine."""
+        self._successor_weights = list(weights)
+
+    def _prefetch_tensors(self, mla):
+        """The tensors this call stages, or an empty list (static facts)."""
+        cfg = self.config
+        tensors = []
+        if self._prefetch_wanted(mla):
+            tensors += [mla.v_b_proj, mla.o_proj.weight]
+        if cfg.successor_prefetch and self._successor_weights:
+            succ = [w.data if hasattr(w, "data") else w for w in self._successor_weights]
+            if check_prefetch_tensors(succ, cfg.prefetch_chunk) == "":
+                tensors += succ
+        return tensors[:3]
+
+    def _prefetch_wanted(self, mla) -> bool:
+        cfg = self.config
+        if not cfg.prefetch_l2:
+            return False
+        if (
+            cfg.prefetch_layers not in _PREFETCH_LAYERS
+            or cfg.prefetch_join not in _JOIN_POINTS
+            or cfg.prefetch_fork not in ("start", "attn", "prev")
+        ):
+            raise ValueError(
+                f"[glm_moe_dsa] attention_output_tail_prefetch/{self.layer_key}: bad prefetch "
+                f"config {cfg}"
+            )
+        if cfg.prefetch_layers != "all":
+            has_indexer = getattr(getattr(mla, "mqa", None), "indexer", None) is not None
+            if has_indexer != (cfg.prefetch_layers == "full"):
+                return False
+        weights = [mla.v_b_proj, mla.o_proj.weight]
+        if check_prefetch_tensors(weights, cfg.prefetch_chunk):
+            return False
+        if cfg.oproj_engine == "native" and cfg.native_oproj == "ring":
+            from .cute_oproj_ring import check_shapes as ring_shapes
+
+            n_dim, k_dim = mla.o_proj.weight.shape
+            return ring_shapes(n_dim, k_dim, cfg.ring_grid, cfg.ring_slots) is None
+        return True
+
+    def before_attention(self, mla, point: str = "start") -> None:
+        """``attention_output_tail_prefetch``: fork the off-path L2 prefetch of this layer's tail weights.
+
+        Called at the start of the generation attention of the op
+        (``point="start"``). Acts only on a call that ``arm`` accepted, only
+        when the config selects the prefetch and only at the configured fork
+        point; the decision uses static facts and the armed flag, before any
+        device work. The fork is joined in ``project``.
+        """
+        if not self._armed or point != self.config.prefetch_fork:
+            return
+        tensors = self._prefetch_tensors(mla)
+        if not tensors:
+            return
+        cfg = self.config
+        if self._prefetch is None:
+            self._prefetch = WeightPrefetch(
+                PrefetchConfig(
+                    chunk=cfg.prefetch_chunk,
+                    ctas=cfg.prefetch_ctas,
+                    threads=cfg.prefetch_threads,
+                    gap_ns=cfg.prefetch_gap_ns,
+                )
+            )
+        self._prefetch.fork(tensors)
+        self.prefetch_forks += 1
+
+    def fork_ahead(self, mla) -> None:
+        """``attention_output_tail_prefetch`` with ``prefetch_fork="prev"``: fork this layer's weight
+        prefetch now, ahead of its call.
+
+        Called by the previous layer right before its post-attention or
+        post-MoE hop (an NVLink-latency-bound window), from static facts only
+        (config, capability, layer selection); no-op otherwise. The fork is
+        joined in this layer's ``project``, or in ``arm`` if the call takes
+        the native path.
+        """
+        if self.config.prefetch_fork != "prev" or self._ahead or not self._check_static(mla):
+            return
+        tensors = self._prefetch_tensors(mla)
+        if not tensors:
+            return
+        if self._prefetch is None:
+            cfg = self.config
+            self._prefetch = WeightPrefetch(
+                PrefetchConfig(
+                    chunk=cfg.prefetch_chunk,
+                    ctas=cfg.prefetch_ctas,
+                    threads=cfg.prefetch_threads,
+                    gap_ns=cfg.prefetch_gap_ns,
+                )
+            )
+        self._prefetch.fork(tensors)
+        self._ahead = True
+        self.prefetch_forks += 1
+
+    def _join_prefetch(self) -> None:
+        self._ahead = False
+        if self._prefetch is not None:
+            self._prefetch.join()
+
+    def defer_w_uv(self, mla, attn_metadata, latent: torch.Tensor, output: torch.Tensor) -> bool:
+        """Take over this call's W_UV + o_proj if ``arm`` accepted the call.
+
+        Args:
+            mla: the calling ``MLA`` module.
+            attn_metadata: the call's attention metadata.
+            latent: ``[num_gen_tokens, H, C]`` native latent attention output.
+            output: the generation rows of ``attn_output[0]``, which the native
+                W_UV BMM would write.
+
+        Returns:
+            True if the caller must skip the native W_UV; the packet is then
+            consumed by ``project``.
+        """
+        armed, self._armed = self._armed, False
+        if not armed:
+            return False
+        if self._pending is not None:
+            raise RuntimeError(
+                f"[glm_moe_dsa] attention_output_tail/{self.layer_key}: previous tail packet was "
+                "not consumed by the output projection"
+            )
+        if self.config.oproj_engine == "native":
+            # Experiments only: keep the native W_UV/o_proj (the prefetch
+            # fork, if any, is still joined in project).
+            return False
+        num_rows = latent.shape[0]
+        eligible = (
+            attn_metadata.num_contexts == 0
+            and 1 <= num_rows <= min(ROW_TILE, MAX_ENGAGED_ROWS)
+            and output.shape[0] == num_rows
+            and check_tail_shapes(latent, mla.v_b_proj, mla.o_proj.weight, num_rows, num_rows)
+            is None
+        )
+        if not eligible:
+            self.declined_calls += 1
+            return False
+        self._pending = _Packet(latent=latent, output_ptr=output.data_ptr(), num_rows=num_rows)
+        self.engaged_calls += 1
+        return True
+
+    def project(
+        self, mla, attn_output: torch.Tensor, all_reduce_params=None
+    ) -> Optional[torch.Tensor]:
+        """Return this call's rank-local o_proj output, or ``None`` for the
+        native path.
+
+        The result is the partial before any communication: the caller
+        reduces it exactly as the ROW-parallel ``Linear.forward`` reduces the
+        output of this ``o_proj`` (no bias, LoRA or fused GEMM+all-reduce;
+        see ``_static_reason``), with the same ``AllReduceParams``.
+
+        Args:
+            mla: the calling ``MLA`` module.
+            attn_output: ``attn_output[0]``, the native o_proj input buffer
+                ``[out_rows, H * V]`` (its rows were not written by W_UV when
+                a packet is pending).
+            all_reduce_params: unused; the tail does not communicate.
+        """
+        del all_reduce_params
+        self._armed = False
+        packet = self._pending
+        if packet is None:
+            if (
+                self.config.oproj_engine == "native"
+                and self._prefetch is not None
+                and self._prefetch.forked
+            ):
+                # attention_output_tail_prefetch with the native kernels: the native W_UV already
+                # ran in the op; run the native o_proj GEMM here so that the join can
+                # follow it (as in MLA._project_output without helix CP,
+                # which the static check guarantees).
+                if self.config.prefetch_join != "after_oproj":
+                    self._join_prefetch()
+                if self.config.native_oproj == "ring":
+                    out = self._ring_oproj(mla, attn_output)
+                else:
+                    # The GEMM of Linear.forward for an o_proj without bias
+                    # or LoRA, without its all-reduce.
+                    out = mla.o_proj.apply_linear(attn_output, None, layer_idx=mla.layer_idx)
+                self._join_prefetch()
+                self.engaged_calls += 1
+                return out
+            # Declined at the W_UV step: the native projection follows.
+            self._join_prefetch()
+            return None
+        self._pending = None
+        if attn_output.data_ptr() != packet.output_ptr or attn_output.shape[0] < packet.num_rows:
+            self._join_prefetch()
+            raise RuntimeError(
+                f"[glm_moe_dsa] attention_output_tail/{self.layer_key}: tail packet does not match "
+                "the output projection input"
+            )
+        join = self.config.prefetch_join
+        if join == "before_wuv":
+            self._join_prefetch()
+        hop_transport = None
+        request, self._into_hop = self._into_hop, None
+        if request is not None:
+            transport, hop_rows = request
+            # The hop polls hop_rows rows and later resets exactly those rows
+            # of the buffer, so the producer must publish exactly them.
+            if attn_output.shape[0] == hop_rows <= MAX_INTO_HOP_ROWS:
+                hop_transport = transport
+            else:
+                self.into_hop_declined_calls += 1
+        partial = attn_tail(
+            packet.latent,
+            mla.v_b_proj,
+            mla.o_proj.weight,
+            packet.num_rows,
+            out_rows=attn_output.shape[0],
+            config=self.config,
+            state=self._state,
+            between=(self._join_prefetch if join == "before_oproj" else None),
+            hop_transport=hop_transport,
+        )
+        if hop_transport is not None:
+            self._into_hop_published = True
+            self.into_hop_published_calls += 1
+        # Always joined before the layer's forward returns.
+        self._join_prefetch()
+        return partial
+
+    def _ring_oproj(self, mla, attn_output: torch.Tensor) -> torch.Tensor:
+        """o_proj GEMM of the ``native`` engine with ``native_oproj="ring"``."""
+        from .cute_oproj_ring import ring_oproj
+        from .kernels import pdl_enabled
+
+        cfg = self.config
+        w_o = mla.o_proj.weight
+        out = torch.empty((attn_output.shape[0], w_o.shape[0]), dtype=w_o.dtype, device=w_o.device)
+        ring_oproj(
+            attn_output,
+            w_o,
+            out,
+            use_pdl=pdl_enabled(),
+            grid=cfg.ring_grid,
+            slots=cfg.ring_slots,
+            early=cfg.ring_early,
+            trigger=pdl_enabled(),
+        )
+        return out
+
+
+def _selected(feature_id: str) -> bool:
+    """Whether ``feature_id`` is selected; False on a build whose selector
+    does not know it."""
+    if feature_id not in dispatch.SELECTOR_PATHS:
+        return False
+    return dispatch.is_enabled(feature_id)
+
+
+def tail_prefetch_selected() -> bool:
+    """Whether ``attention_output_tail_prefetch`` (attention_output_tail with off-path weight prefetch) is selected."""
+    return _selected(TAIL_PREFETCH_PATH)
+
+
+def tail_prefetch_config() -> Optional[TailConfig]:
+    """``attention_output_tail_prefetch``'s tail schedule for the neighbours selected in this process.
+
+    The schedule depends on the attention kernel before the tail (``sparse_mla_decode``
+    releases its PDL dependents well before it ends, the native FMHA only
+    shortly before) and on the op after it (the ``post_attention_hop`` needs whole
+    SMs and gains nothing from router-weight staging; the native router
+    GEMM does). The selector is identical on every rank, so is the choice.
+    Returns None where the native tail is kept
+    (``post_attention_hop`` without ``sparse_mla_decode``).
+    """
+    post_hop, sparse_decode = _selected("post_attention_hop"), _selected("sparse_mla_decode")
+    if post_hop and not sparse_decode:
+        return None
+    if post_hop:
+        return TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE
+    if sparse_decode:
+        return TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE
+    return TAIL_PREFETCH_SCHEDULE
+
+
+def make_attn_tail(mla) -> Optional[AttnTail]:
+    """Create the attention_output_tail for a target-model MLA module, or ``None``.
+
+    Draft (MTP) layers, whose index is at or beyond the target layer count,
+    keep the native tail. With ``attention_output_into_hop`` and ``post_attention_hop`` selected the
+    tail uses ``INTO_HOP_SPARSE_DECODE_SCHEDULE`` behind ``sparse_mla_decode`` and
+    ``INTO_HOP_SCHEDULE`` behind the native FMHA; ``attention_output_into_hop``
+    without ``post_attention_hop`` is decided as for
+    ``attention_output_tail_prefetch``/``attention_output_tail``. With
+    ``attention_output_tail_prefetch`` selected the tail uses ``tail_prefetch_config()``; with only
+    ``attention_output_tail`` it uses the attention_output_tail default.
+    """
+    layer_idx = getattr(mla, "layer_idx", None)
+    num_layers = getattr(mla, "num_hidden_layers", None)
+    if layer_idx is None or (num_layers is not None and layer_idx >= num_layers):
+        dispatch.log_capability(FEATURE_ID, f"layer{layer_idx}", False, "draft or unindexed layer")
+        return None
+    key = f"layer{layer_idx}"
+    if _selected(INTO_HOP_PATH):
+        if _selected("post_attention_hop"):
+            return AttnTail(
+                key,
+                INTO_HOP_SPARSE_DECODE_SCHEDULE
+                if _selected("sparse_mla_decode")
+                else INTO_HOP_SCHEDULE,
+            )
+        if not (tail_prefetch_selected() or _selected(FEATURE_ID)):
+            dispatch.log_capability(INTO_HOP_PATH, key, False, "needs the post_attention_hop hop")
+            return None
+    if tail_prefetch_selected():
+        config = tail_prefetch_config()
+        if config is None:
+            if _selected(FEATURE_ID):
+                # This selection runs the default tail schedule instead.
+                dispatch.log_capability(
+                    TAIL_PREFETCH_PATH,
+                    "schedule",
+                    False,
+                    "post_attention_hop without sparse_mla_decode",
+                )
+                return AttnTail(f"layer{layer_idx}")
+            dispatch.log_capability(
+                TAIL_PREFETCH_PATH,
+                f"layer{layer_idx}",
+                False,
+                "post_attention_hop without sparse_mla_decode: native tail is faster",
+            )
+            return None
+        return AttnTail(f"layer{layer_idx}", config)
+    return AttnTail(f"layer{layer_idx}")
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py
new file mode 100644
index 000000000000..c90328e8a5b3
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py
@@ -0,0 +1,244 @@
+# 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.
+"""Logical tensor contracts shared by the decode kernels of the glm_moe_dsa targets.
+
+Some kernel packages add PDL inside the package, each with a documented
+producer/consumer memory order. The contracts are written for GEN-only,
+non-attention-DP, TP8, fixed MTP depth K=5. Each contract lists its logical
+tensors; physical layouts follow the native TensorRT-LLM producer and
+consumer at that seam. ``R`` is the row capacity of one call (live rows are
+selected by a separate mask), ``H`` the local attention-head count.
+"""
+
+from dataclasses import dataclass
+from typing import Tuple
+
+
+@dataclass(frozen=True)
+class GlmShape:
+    """Model-shape constants checked against the GLM-5.2 NVFP4 config."""
+
+    num_layers: int = 78
+    num_dense_layers: int = 3
+    hidden_size: int = 6144
+    num_heads: int = 64
+    q_lora_rank: int = 2048
+    kv_lora_rank: int = 512
+    qk_nope_head_dim: int = 192
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 256
+    index_n_heads: int = 32
+    index_head_dim: int = 128
+    index_topk: int = 2048
+    num_routed_experts: int = 256
+    experts_per_token: int = 8
+    moe_intermediate_size: int = 2048
+    vocab_size: int = 154880
+    tp_size: int = 8
+    mtp_depth: int = 5
+
+    @property
+    def local_heads(self) -> int:
+        return self.num_heads // self.tp_size
+
+
+GLM52 = GlmShape()
+
+
+@dataclass(frozen=True)
+class Contract:
+    """One seam between native TensorRT-LLM code and a package.
+
+    ``rules`` are obligations of the seam beyond its tensors: communication,
+    ordering and memory order.
+    """
+
+    name: str
+    owner: str
+    inputs: Tuple[str, ...]
+    outputs: Tuple[str, ...]
+    rules: Tuple[str, ...] = ()
+
+
+# Rules of the attention_output_tail contract. They bind attention_output_tail and its tail
+# continuations in ``attention_output_tail/`` (attention_output_tail_prefetch, and
+# attention_output_into_hop apart from the into-hop publication named in the first rule).
+ATTENTION_TAIL_COMMUNICATION_RULE = (
+    "attention_output_tail owns W_UV + local o_proj and returns the unreduced rank-local "
+    "partial. router_gemm/native all-reduce consumes the identical partial. No "
+    "communication is hidden inside attention_output_tail. Only attention_output_into_hop (spec "
+    "attention_output_into_hop, a declared "
+    "cross-boundary fusion) adds a publication: after its griddepcontrol.wait "
+    "the o_proj kernel also writes the partial (-0.0 mapped to +0.0) into "
+    "the post_attention_hop hop's Lamport transport under the attention_output_into_hop protocol of "
+    "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h, and the "
+    "ordinary partial stays available."
+)
+ATTENTION_TAIL_PDL_RULE = (
+    "attention_output_tail may be a PDL dependent of the latent's producer. Before its "
+    "griddepcontrol.wait a attention_output_tail kernel reads only immutable weights (v_b_proj, "
+    "local o_proj weight) and writes only its own shared memory "
+    "(cp.async.bulk global->shared, completed through mbarriers before use; "
+    "no async-proxy global writes). Every latent read and every global write "
+    "(v scratch, partial) follows the wait. attention_output_tail's o_proj kernel executes "
+    "griddepcontrol.launch_dependents before it writes the partial: every "
+    "PDL-launched consumer of the partial must execute griddepcontrol.wait "
+    "before reading the partial and before writing any global memory that "
+    "may alias a buffer attention_output_tail still reads (v scratch and the latent are pool "
+    "memory and may be reused after attention_output_tail returns). Kill switches: "
+    "TRTLLM_ENABLE_PDL=0 (all, entry included), "
+    "TailConfig(trigger_dependents=False) (exit)."
+)
+
+# Physical form of the routed_moe contract's unfinalized tuple as routed_moe returns it.
+ROUTED_MOE_TOKEN_MAJOR_RULE = (
+    "A routed-MoE package may return the unfinalized tuple in token-major "
+    "form: fc2_output[R*8,6144] BF16 whose row t*8+k holds token t's k-th "
+    "expert output (row pitch 6144 elements), expert scales [R,8] BF16 with "
+    "the native dtype, values and top-k order, and "
+    "expanded_idx_to_permuted_idx[R,8] int32 equal to the identity t*8+k "
+    "(every row valid). The native runner's expert-grouped, padded layout is "
+    "the other legal form; every unfinalized consumer (native MoEAllReduce "
+    "finalize + all-reduce, the post_moe_hop hop) reads fc2_output only through the "
+    "permutation, so it accepts either."
+)
+
+# Routed input (routed_moe_into_hop): routed_moe's routed result written into the post_moe_hop.
+POST_MOE_HOP_ROUTED_INPUT_RULE = (
+    "With routed_moe_into_hop selected, routed_moe writes its routed result into a routed-input "
+    "buffer owned by the post_moe_hop hop instead of returning the routed_moe contract's output. "
+    "PostMoeHop.enable_routed_input allocates the buffers zeroed in device "
+    "memory and the mode is the routed_moe_into_hop numerical variant "
+    "(TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION, parsed only "
+    "by dispatch.feature_variants), agreed by every rank of the group; a "
+    "rank that disagrees on the mode, the buffers or the compile keeps every "
+    "rank on the native seam. Modes: row = finalized rows [R,6144] BF16; "
+    "identity = token-major raw rows [R*8,6144] BF16 and scales [R,8] "
+    "(ROUTED_MOE_TOKEN_MAJOR_RULE); fp32 = an FP32 accumulator [R,6144] that is "
+    "zero when routed_moe's kernel starts, to which routed_moe adds the eight weighted rows "
+    "of a token with atomic FP32 adds (order run-to-run nondeterministic), "
+    "and whose rows the hop clears after reading them. routed_moe never reuses a "
+    "buffer it wrote into; the hop owns their lifetime. PDL edge routed_moe -> post_moe_hop: "
+    "routed_moe's expert-cluster kernel executes griddepcontrol.launch_dependents "
+    "after each CTA's last FC2 tile, and the post_moe_hop (or deferred_norm) hop reads the "
+    "routed-input buffer only after its own griddepcontrol.wait."
+)
+
+# Deferred-norm packet (deferred_norm): the packet at the post_moe_hop -> query_front seam.
+POST_MOE_HOP_DEFERRED_NORM_RULE = (
+    "With deferred_norm selected, the post-MoE hop returns a DeferredNormHidden packet "
+    "(deferred_norm/packet.py) in place of the normalized hidden state: a_frag = "
+    "bf16(x * gamma) in the mma.m16n8k16 B-fragment layout of packet.py, "
+    "sumsq[R,P] FP32 per-warp square sums in post_moe_hop's reduction order, residual "
+    "[R,6144] BF16 bit-identical to post_moe_hop's, the next layer's gamma and eps. "
+    "The packet is not a tensor: a consumer that does not know it fails at "
+    "its first tensor operation, and every consumer other than deferred_norm's "
+    "projection calls materialize(), which returns post_moe_hop's normalized output "
+    "bit for bit. deferred_norm's projection applies r = rsqrt(sum(sumsq) / H + eps) "
+    "after its GEMM. PDL: with trigger_early (the default) the hop executes "
+    "griddepcontrol.launch_dependents right after its own "
+    "griddepcontrol.wait; the dependent projection reads only immutable "
+    "weights before its griddepcontrol.wait and reads the packet after it. "
+    "Scope (deferred_norm variant; this target runs shared): all = "
+    "every boundary into a target layer; shared = only boundaries into "
+    "shared-indexer layers, "
+    "decided per layer at hop construction and agreed by every rank, so a "
+    "full-indexer layer receives post_moe_hop's normalized tensor."
+)
+
+# Shared communication state of the post_moe_hop contract (deferred_norm on the post_moe_hop
+# workspace).
+POST_MOE_HOP_SHARED_PROTOCOL_RULE = (
+    "deferred_norm's hop kernel (deferred_norm/hop_kernel.py) runs on the post_moe_hop hop's workspace and "
+    "implements post_moe_hop's protocol a second time: epoch % 3 buffer rotation, "
+    "clearing the dirty buffer by the previous call's row count, the "
+    "Lamport sentinel poll bounded by post_moe_hop's timeout_ns, the sticky per-buffer "
+    "error words and the arrival count; it takes only the state-word indices "
+    "from post_moe_hop/kernels.py (deferred_norm/ptx.py). post_moe_hop and deferred_norm hops of one group advance "
+    "one epoch and one error chain, so a change to either kernel's protocol "
+    "must change both in the same commit and pass the interleaved post_moe_hop/deferred_norm "
+    "tests (test_deferred_norm_hop_mpi.py, test_deferred_norm_routed_mpi.py)."
+)
+
+CONTRACTS = {
+    "query_front": Contract(
+        "query/projection packet",
+        "absorbed_query",
+        ("hidden[R,6144] or DeferredNormHidden (deferred_norm)", "native q/kv-a/indexer weights"),
+        (
+            "q[R,H,256] or (q_nope_absorbed[R,H,512], q_rope[R,H,64])",
+            "native compressed-kv/rope/indexer projection outputs",
+        ),
+        (POST_MOE_HOP_DEFERRED_NORM_RULE,),
+    ),
+    "index_selection": Contract(
+        "index selection",
+        "indexer_topk",
+        ("native projected query/key/scale/weights", "immutable cache view", "live lengths"),
+        ("topk_ids[R,2048]", "validity"),
+    ),
+    "sparse_attention": Contract(
+        "sparse attention",
+        "sparse_mla_decode",
+        ("prepared query + scales", "FP8 paged KV", "selection/validity", "causal positions"),
+        ("latent_out[R,H,512]",),
+    ),
+    "attention_output_tail": Contract(
+        "attention tail",
+        "attention_output_tail",
+        ("latent_out[R,H,512]", "W_UV", "local W_o"),
+        ("partial[R,6144]",),
+        (ATTENTION_TAIL_COMMUNICATION_RULE, ATTENTION_TAIL_PDL_RULE),
+    ),
+    "router": Contract(
+        "pre-MoE hop",
+        "router_gemm",
+        ("rank-local attention partial", "residual", "norm weight", "router weight"),
+        ("normed_hidden", "residual", "router_logits[R,256] or None"),
+    ),
+    "routed_moe": Contract(
+        "routed MoE",
+        "routed_moe",
+        ("hidden or quantized hidden", "router logits", "native expert weights/scales"),
+        (
+            "finalized local routed output",
+            "or (fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx)",
+            "or the post_moe_hop hop's routed-input buffer (routed_moe_into_hop)",
+        ),
+        (ROUTED_MOE_TOKEN_MAJOR_RULE, POST_MOE_HOP_ROUTED_INPUT_RULE),
+    ),
+    "post_moe_hop": Contract(
+        "post-MoE hop",
+        "post_moe_hop",
+        (
+            "native raw MoE tuple or its routed-input buffer (routed_moe_into_hop)",
+            "shared output or None",
+            "residual",
+            "norm weight",
+        ),
+        ("post-reduction hidden or DeferredNormHidden (deferred_norm)", "residual"),
+        (
+            POST_MOE_HOP_ROUTED_INPUT_RULE,
+            POST_MOE_HOP_DEFERRED_NORM_RULE,
+            POST_MOE_HOP_SHARED_PROTOCOL_RULE,
+        ),
+    ),
+    "draft_front": Contract(
+        "draft front",
+        "draft_front",
+        ("draft embedding", "previous hidden", "enorm/hnorm weights"),
+        ("TP-local packed eh_proj input",),
+    ),
+}
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/__init__.py
new file mode 100644
index 000000000000..aba54c3ee9da
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/__init__.py
@@ -0,0 +1,39 @@
+# 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.
+"""cross_rank_topk_cluster_merge: cross_rank_topk's context-partitioned top-k with a cluster merge
+(selected with ``cross_rank_topk_cluster_merge``).
+
+cross_rank_topk_cluster_merge extends cross_rank_topk (``cross_rank_topk``): the region, its policy
+and its CP logits kernel are cross_rank_topk's (``indexer_topk/cross_rank_*.py``); the select / exchange /
+merge op is ``glm_moe_dsa_cross_rank_topk_cluster_merge`` (``crossRankTopkClusterMerge.h``), which
+merges each verify row on the row's eight CTAs as one thread-block cluster and keeps per-CTA epoch
+words. Its output rows are byte-identical to cross_rank_topk's (set and order) on every rank.
+
+Registered entry points (``dispatch.get_impl("cross_rank_topk_cluster_merge", name)``):
+
+- ``cp_indexer_decode``: :func:`cp_indexer_decode.try_cp_indexer_decode`, the
+  adapter hook of the GEN region (tried before cross_rank_topk's when
+  ``cross_rank_topk_cluster_merge`` is selected).
+- ``cp_decode_topk``: :func:`cp_indexer_decode.cp_decode_topk`, its
+  tensor-level entry.
+"""
+
+from ...dispatch import register
+from .cp_indexer_decode import cp_decode_topk, cp_topk_from_local_logits, try_cp_indexer_decode
+
+register("cross_rank_topk_cluster_merge", "cp_indexer_decode", try_cp_indexer_decode)
+register("cross_rank_topk_cluster_merge", "cp_decode_topk", cp_decode_topk)
+
+__all__ = ["cp_decode_topk", "cp_topk_from_local_logits", "try_cp_indexer_decode"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_indexer_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_indexer_decode.py
new file mode 100644
index 000000000000..64acb9ba352f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_indexer_decode.py
@@ -0,0 +1,264 @@
+# 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.
+"""cross_rank_topk_cluster_merge: the cross_rank_topk GEN indexer region with the cluster-merge top-k op.
+
+Same region, policy, CP logits kernel and output contract as cross_rank_topk
+(``indexer_topk/cross_rank_indexer_decode.py``): rank ``r`` scores the 64-token pages
+``p % world == r``, and ``glm_moe_dsa_cross_rank_topk_cluster_merge`` selects, exchanges and merges
+the candidates into the tie-aware exact top-``index_topk`` of every row, the
+same set and order on every rank and the same rows cross_rank_topk writes.
+
+Numerical status: as cross_rank_topk (rank-consistent exact top-k over each page owner's
+indexer K); cross_rank_topk_cluster_merge's rows are byte-identical to cross_rank_topk's for the
+same local logits.
+"""
+
+from typing import Dict
+
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.logger import logger
+
+from ..indexer_topk import cross_rank_policy
+from . import cp_workspace
+
+__all__ = [
+    "CP_DECISION_COUNTS",
+    "SUPPORTED_WORLD_SIZES",
+    "cp_decode_topk",
+    "cp_topk_from_local_logits",
+    "try_cp_indexer_decode",
+    "world_size_reason",
+]
+
+# As cross_rank_topk: MTP draft step 0 runs the same region on the MTP layer.
+ALLOW_DRAFT_STEP = True
+
+# Tensor-parallel group sizes cross_rank_topk_cluster_merge's GPU tests verify (one node of 4 GPUs,
+# and a two-host GEN group); at any other size cross_rank_topk serves the region.
+SUPPORTED_WORLD_SIZES = (4, 8)
+
+
+def world_size_reason(world: int) -> str:
+    """Why cross_rank_topk_cluster_merge declines a group of ``world`` ranks ("" when it may engage).
+
+    Identical on every rank of the group, so the decline needs no agreement.
+    """
+    if world not in SUPPORTED_WORLD_SIZES:
+        return (
+            f"tensor-parallel size {world} not in cross_rank_topk_cluster_merge's verified sizes "
+            f"{SUPPORTED_WORLD_SIZES}"
+        )
+    return ""
+
+
+CP_DECISION_COUNTS: Dict[str, int] = {}
+
+
+def _record_decision(decision) -> None:
+    key = "engaged" if decision.eligible else decision.reason
+    count = CP_DECISION_COUNTS.get(key, 0)
+    CP_DECISION_COUNTS[key] = count + 1
+    if count == 0:
+        state = "engaged" if decision.eligible else f"declined ({decision.reason})"
+        logger.info(
+            f"[glm_moe_dsa] cross_rank_topk_cluster_merge/cp_indexer_decode: first '{state}' decision on this rank"
+        )
+
+
+def cp_topk_from_local_logits(
+    local_logits: torch.Tensor,
+    kv_lens: torch.Tensor,
+    out: torch.Tensor,
+    workspace,
+    next_n: int,
+    debug_times=None,
+) -> torch.Tensor:
+    """Select, exchange and merge from this rank's local logits.
+
+    Args:
+        local_logits: ``[next_n, local_capacity]`` float32 logits of this
+            rank's local columns (see ``indexer_topk.cross_rank_logits_runner.local_to_global``).
+        kv_lens: ``[1]`` int32 global indexer KV length.
+        out: ``[next_n, index_topk]`` int32 output.
+        workspace: ``cp_workspace.CpWorkspace`` of the group.
+        next_n: Verify rows of the request.
+        debug_times: Optional ``[next_n, 8, 64]`` int64 phase stamps.
+
+    Returns:
+        ``out`` holding the same global top-k selection on every rank.
+    """
+    torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_cluster_merge(
+        local_logits,
+        kv_lens,
+        out,
+        workspace.state,
+        workspace.mc_ptr,
+        workspace.uc_ptr,
+        workspace.rank,
+        workspace.world,
+        next_n,
+        debug_times,
+    )
+    return out
+
+
+def cp_decode_topk(
+    *,
+    q: torch.Tensor,
+    weights: torch.Tensor,
+    k_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    kv_lens: torch.Tensor,
+    max_seq_len: int,
+    out: torch.Tensor,
+    workspace,
+) -> torch.Tensor:
+    """Run the cross_rank_topk_cluster_merge region for one request (arguments as cross_rank_topk's
+    ``cp_decode_topk``)."""
+    from ..indexer_topk.cross_rank_logits_runner import paged_mqa_logits_cp
+
+    local_logits = paged_mqa_logits_cp(
+        q, k_cache, weights, kv_lens, block_table, max_seq_len, workspace.rank, workspace.world
+    )
+    return cp_topk_from_local_logits(local_logits, kv_lens, out, workspace, q.shape[1])
+
+
+def try_cp_indexer_decode(
+    indexer,
+    metadata,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+) -> bool:
+    """Adapter entry: run the GEN region of ``Indexer.sparse_attn_indexer`` with cross_rank_topk_cluster_merge.
+
+    The decision is cross_rank_topk's (``indexer_topk.cross_rank_policy``) on the same host state,
+    plus cross_rank_topk_cluster_merge's own workspace creation (collective, first eligible eager
+    call), whose capability agreement also covers the cluster residency. A ``False`` return leaves
+    no side effect; every rank of the group reaches the same decision for the same forward, and the
+    caller then tries cross_rank_topk's region.
+
+    Args/Returns: as ``indexer_topk.indexer_decode.try_indexer_decode``.
+    """
+    from ...dispatch import log_capability
+    from ..indexer_topk.cross_rank_logits_runner import ensure_compiled, local_capacity
+
+    num_contexts = metadata.num_contexts
+    num_generations = metadata.num_generations
+    gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
+    uniform = num_generations > 0 and bool(
+        (gen_seq_lens == num_gen_tokens // num_generations).all()
+    )
+    top_k = indexer.top_k
+    k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(indexer.layer_idx)
+    block_table = metadata.indexer_k_cache_block_offsets
+    mapping = getattr(metadata, "mapping", None)
+    max_seq_len = metadata.get_indexer_max_seq_len()
+    world = mapping.tp_size if mapping is not None else 1
+    info = cp_workspace.op_info(world) if world > 0 else None
+    capacity = local_capacity(max_seq_len, world) if world > 0 else 0
+    # Facts of this rank's build and device, agreed by the workspace creation.
+    local_reason = cross_rank_policy.rank_local_reason(
+        op_available=info is not None,
+        local_capacity=capacity,
+        max_local_capacity=info[cp_workspace.INFO_MAX_LOCAL_CAPACITY] if info is not None else 0,
+        sm_version=get_sm_version(),
+        cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE,
+    ) or (cp_workspace.cluster_reason(info) if info is not None else "")
+    decision = cross_rank_policy.decide_cp_uniform(
+        world_size=world,
+        local_capacity=capacity,
+        allow_draft_step=ALLOW_DRAFT_STEP,
+        native_uses_cute_dsl_logits=indexer.use_cute_dsl_paged_mqa_logits,
+        use_fp4_indexer=indexer.use_fp4,
+        use_gvr_emission=indexer.use_gvr_emission,
+        needs_gvr_prior=top_k.needs_gvr_prior,
+        is_ragged_verify=metadata.is_ragged_verify,
+        expanded_mtp_buffers=metadata.use_expanded_buffers_for_mtp,
+        in_mtp_draft_loop=metadata.in_mtp_draft_loop,
+        attention_dp=bool(mapping is not None and mapping.enable_attention_dp),
+        topk_decode_impl=top_k.decode_implementation.value,
+        topk_self_sampling=top_k.gvr_self_sampling,
+        index_topk=top_k.top_k,
+        compress_ratio=top_k.compress_ratio,
+        num_heads=indexer.n_heads,
+        head_dim=indexer.head_dim,
+        query_dtype=str(q_fp8.dtype),
+        query_shape=tuple(q_fp8.shape),
+        query_contiguous=q_fp8.is_contiguous(),
+        weights_dtype=str(weights.dtype),
+        weights_shape=tuple(weights.shape),
+        weights_contiguous=weights.is_contiguous(),
+        k_cache_dtype=str(k_cache.dtype),
+        k_cache_shape=tuple(k_cache.shape),
+        block_table_cols=block_table.shape[1],
+        max_seq_len=max_seq_len,
+        num_contexts=num_contexts,
+        token_offset=token_offset,
+        num_gen_tokens=num_gen_tokens,
+        num_generations=num_generations,
+        uniform_verify_width=uniform,
+    )
+    if decision.eligible and world_size_reason(world):
+        decision = type(decision)(False, world_size_reason(world))
+    if decision.eligible:
+        # The logits kernel is compiled before the group agrees on the
+        # workspace, so no rank waits in an exchange while a peer compiles.
+        workspace = cp_workspace.get_workspace(
+            mapping,
+            prepare=lambda: ensure_compiled(
+                q_fp8.device,
+                k_cache.shape[1],
+                indexer.n_heads,
+                indexer.head_dim,
+                num_gen_tokens // num_generations,
+                world,
+            ),
+            local_reason=local_reason,
+        )
+        if workspace is None:
+            reason = (
+                cp_workspace.failure_reason(mapping)
+                or "workspace not created before CUDA-graph capture"
+            )
+            decision = type(decision)(False, f"workspace unavailable: {reason}")
+    log_capability(
+        "cross_rank_topk_cluster_merge", "cp_indexer_decode", decision.eligible, decision.reason
+    )
+    _record_decision(decision)
+    if not decision.eligible:
+        return False
+
+    next_n = num_gen_tokens // num_generations
+    gen_rows = slice(token_offset, token_offset + num_gen_tokens)
+    gen_reqs = slice(num_contexts, num_contexts + num_generations)
+    kv_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+    assert kv_lens is not None
+    cp_decode_topk(
+        q=q_fp8[gen_rows].view(num_generations, next_n, indexer.n_heads, indexer.head_dim),
+        weights=weights[gen_rows],
+        k_cache=k_cache,
+        block_table=block_table[gen_reqs],
+        kv_lens=kv_lens,
+        max_seq_len=max_seq_len,
+        out=out,
+        workspace=workspace,
+    )
+    return True
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_workspace.py
new file mode 100644
index 000000000000..a72fcbdfecdc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/cp_workspace.py
@@ -0,0 +1,247 @@
+# 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.
+"""Exchange workspace of the cross_rank_topk_cluster_merge context-partitioned top-k.
+
+cross_rank_topk_cluster_merge owns one multicast workspace per tensor-parallel group, allocated with
+the native ``McastGPUBuffer`` and used by nothing else (cross_rank_topk's workspace is a
+separate one). Its layout and Lamport protocol are defined by the
+``glm_moe_dsa_cross_rank_topk_cluster_merge`` kernel (``crossRankTopkClusterMerge.h``): every 32-bit word starts as
+``0xFFFFFFFF`` and the kernel rotates three buffers with per-CTA epoch words
+held in a separate per-rank state tensor.
+
+Creation is collective over the group and only happens outside CUDA-graph
+capture, exactly as cross_rank_topk's (``indexer_topk/cross_rank_workspace.py``): each rank runs the
+caller's host-side preparation, the ranks agree on the capability (compiled
+op, supported world size, MPI, successful preparation, no rank-local reason,
+and enough co-resident clusters for the largest row count), allocate, and
+agree on success; if any step fails on any rank, every rank records the
+failure, so the collective count and order stay identical.
+"""
+
+import os
+import threading
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.logger import logger
+
+# Indices of torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_cluster_merge_workspace_info().
+INFO_WORKSPACE_WORDS = 0
+INFO_TOPK = 1
+INFO_MAX_RANKS = 2
+INFO_MAX_ROWS = 3
+INFO_PAGE_TOKENS = 4
+INFO_MAX_LOCAL_CAPACITY = 7
+INFO_STATE_WORDS = 8
+INFO_TIMEOUT_MS = 10
+INFO_MAX_ACTIVE_CLUSTERS = 11
+# First epoch word of the state (state[EPOCH_BASE + row * 8 + cta]).
+EPOCH_BASE = 4
+
+
+@dataclass
+class CpWorkspace:
+    """One rank's view of the cross_rank_topk_cluster_merge exchange workspace.
+
+    Attributes:
+        handle: The native ``McastGPUBuffer`` (keeps the mapping alive).
+        mc: Multicast int32 view of the whole workspace.
+        uc: This rank's unicast int32 view of the whole workspace.
+        state: int32 device state: [2] second rounds taken, [3] error word
+            written before a timeout abort, then one epoch word (exchange
+            launches) per (row, CTA) from ``EPOCH_BASE``.
+        rank: Rank in the group.
+        world: Group size.
+        comm: The group communicator the handle was created with (kept alive).
+    """
+
+    handle: object
+    mc: torch.Tensor
+    uc: torch.Tensor
+    state: torch.Tensor
+    rank: int
+    world: int
+    comm: object = None
+
+    @property
+    def mc_ptr(self) -> int:
+        return self.mc.data_ptr()
+
+    @property
+    def uc_ptr(self) -> int:
+        return self.uc.data_ptr()
+
+
+_LOCK = threading.Lock()
+_WORKSPACES: Dict[Tuple, CpWorkspace] = {}
+_FAILED: Dict[Tuple, str] = {}
+
+
+_OP_INFO: Dict[Tuple[int, int], list] = {}
+
+
+def op_info(world_size: int = 8) -> Optional[list]:
+    """Constants of the compiled kernel for a group size, or None when the op is not built.
+
+    Queried once per group size and current device (the region calls it on
+    every eager forward); a missing op is not cached, so an op loaded later
+    is seen.
+    """
+    key = (world_size, torch.cuda.current_device() if torch.cuda.is_available() else -1)
+    info = _OP_INFO.get(key)
+    if info is None:
+        try:
+            info = list(
+                torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_cluster_merge_workspace_info(
+                    world_size
+                )
+            )
+        except (AttributeError, RuntimeError):
+            return None
+        _OP_INFO[key] = info
+    return list(info)
+
+
+def cluster_reason(info: Optional[list]) -> str:
+    """Why this device cannot keep every row's cluster resident ("" when it can)."""
+    if info is None or len(info) <= INFO_MAX_ACTIVE_CLUSTERS:
+        return "glm_moe_dsa_cross_rank_topk_cluster_merge op is not built"
+    clusters, rows = info[INFO_MAX_ACTIVE_CLUSTERS], info[INFO_MAX_ROWS]
+    if clusters < rows:
+        return f"{clusters} co-resident 8-CTA clusters < {rows} rows"
+    return ""
+
+
+def _key(mapping) -> Tuple:
+    return (mapping.world_size, mapping.rank, mapping.tp_size, mapping.pp_size, mapping.cp_size)
+
+
+def failure_reason(mapping) -> Optional[str]:
+    """Reason recorded when this group's workspace could not be created."""
+    return _FAILED.get(_key(mapping))
+
+
+def get_workspace(
+    mapping, prepare: Optional[Callable[[], None]] = None, local_reason: str = ""
+) -> Optional[CpWorkspace]:
+    """Return the group's workspace, creating it collectively on first use.
+
+    Must be reached by every rank of the tensor-parallel group in the same
+    order. Returns None when the workspace is unavailable on any rank, and
+    during CUDA-graph capture when it was not created before.
+
+    Args:
+        mapping: The rank's ``Mapping``.
+        prepare: Optional host-side preparation run on every rank before the
+            capability agreement of the creation; an exception counts as a
+            local failure that every rank agrees on.
+        local_reason: Why this rank's build or device cannot run cross_rank_topk_cluster_merge (""
+        when
+            it can); part of the creation's capability agreement.
+    """
+    key = _key(mapping)
+    ws = _WORKSPACES.get(key)
+    if ws is not None or key in _FAILED:
+        return ws
+    if torch.cuda.is_current_stream_capturing():
+        return None
+    with _LOCK:
+        if key in _WORKSPACES or key in _FAILED:
+            return _WORKSPACES.get(key)
+        ws, reason = _create(mapping, prepare, local_reason)
+        if ws is None:
+            _FAILED[key] = reason
+            logger.info(
+                f"[glm_moe_dsa] cross_rank_topk_cluster_merge/workspace: not created ({reason})"
+            )
+        else:
+            _WORKSPACES[key] = ws
+            info = op_info(ws.world) or []
+            timeout = info[INFO_TIMEOUT_MS] if len(info) > INFO_TIMEOUT_MS else "n/a"
+            clusters = (
+                info[INFO_MAX_ACTIVE_CLUSTERS] if len(info) > INFO_MAX_ACTIVE_CLUSTERS else "n/a"
+            )
+            logger.info(
+                f"[glm_moe_dsa] cross_rank_topk_cluster_merge/workspace: created ({ws.uc.numel() * 4} bytes, "
+                f"rank {ws.rank}/{ws.world}, exchange wait timeout {timeout} ms, "
+                f"co-resident clusters {clusters})"
+            )
+        return ws
+
+
+def _local_capability_reason(mapping, info, prepare) -> str:
+    """This rank's op, group size, cluster residency and preparation check ("" when capable)."""
+    if info is None:
+        return "glm_moe_dsa_cross_rank_topk_cluster_merge op is not built"
+    if mapping.tp_size > info[INFO_MAX_RANKS]:
+        return f"tp_size {mapping.tp_size} > {info[INFO_MAX_RANKS]}"
+    reason = cluster_reason(info)
+    if reason:
+        return reason
+    if prepare is not None:
+        try:
+            prepare()
+        except Exception as e:  # noqa: BLE001 - any local failure must be agreed before use
+            return f"preparation failed ({type(e).__name__}: {e})"
+    return ""
+
+
+def _create(mapping, prepare=None, local_reason="") -> Tuple[Optional[CpWorkspace], str]:
+    from tensorrt_llm._torch.distributed.ops import (
+        _get_mnnvl_workspace_comm,
+        _mnnvl_device_index,
+        _mnnvl_workspace_all_succeeded,
+    )
+    from tensorrt_llm._utils import mpi_disabled
+
+    if mpi_disabled():
+        # Without MPI every rank takes this branch before any collective.
+        return None, "MPI is disabled"
+    comm = _get_mnnvl_workspace_comm(mapping)
+    info = op_info(mapping.tp_size)
+    if not local_reason:
+        local_reason = _local_capability_reason(mapping, info, prepare)
+    if not _mnnvl_workspace_all_succeeded(comm, not local_reason):
+        return None, local_reason or "capability not agreed by every rank"
+
+    words = info[INFO_WORKSPACE_WORDS]
+    error: Optional[Exception] = None
+    ws: Optional[CpWorkspace] = None
+    try:
+        from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer
+
+        use_fabric = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node()
+        device = _mnnvl_device_index(mapping)
+        handle = McastGPUBuffer(
+            words * 4, mapping.tp_size, mapping.tp_rank, device, use_fabric, comm.py2f()
+        )
+        uc = handle.get_uc_buffer(mapping.tp_rank, [words], torch.int32, 0)
+        mc = handle.get_mc_buffer([words], torch.int32, 0)
+        with torch.inference_mode():
+            uc.fill_(-1)
+            state = torch.zeros(
+                info[INFO_STATE_WORDS], dtype=torch.int32, device=torch.device("cuda", device)
+            )
+        torch.cuda.synchronize()
+        ws = CpWorkspace(handle, mc, uc, state, mapping.tp_rank, mapping.tp_size, comm)
+    except Exception as e:  # noqa: BLE001 - any local failure must be agreed before use
+        error = e
+    # Also fences the initialization: no rank sends before every rank cleared
+    # its buffer.
+    if not _mnnvl_workspace_all_succeeded(comm, error is None):
+        return None, f"allocation failed on a rank ({type(error).__name__ if error else 'peer'})"
+    return ws, ""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/__init__.py
new file mode 100644
index 000000000000..ef6080339557
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/__init__.py
@@ -0,0 +1,60 @@
+# 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.
+"""cross_rank_topk_global_emit: rope_kv_append's top-k bridge folded into
+cross_rank_topk_cluster_merge's CP top-k emission (selected with ``cross_rank_topk_global_emit``).
+
+cross_rank_topk_global_emit extends cross_rank_topk_cluster_merge
+(``cross_rank_topk_cluster_merge``), cross_rank_topk (``cross_rank_topk``) and rope_kv_append
+(``rope_kv_append``): the GEN indexer region is cross_rank_topk_cluster_merge's (cross_rank_topk's
+policy and CP logits kernel), and its top-k op ``glm_moe_dsa_cross_rank_topk_global_emit``
+(``crossRankTopkGlobalEmit.h``) writes, besides the id rows, the outputs of rope_kv_append's
+``topkBridgeKernel`` for the call: the global KV-pool rows of the layer's remap group and the native
+top-k copy rows. The backend's remap then returns those rows instead of launching the bridge, and
+sparse_mla_decode runs directly behind the top-k. rope_kv_append's bridge still serves every call
+cross_rank_topk_global_emit does not (MTP draft steps 1-4, which reuse the shared rows, and every
+decline).
+
+Registered entry points (``dispatch.get_impl("cross_rank_topk_global_emit", name)``):
+
+- ``new_bridge_packet``: :func:`bridge.new_bridge_packet` (rope_kv_append's packet with a plan
+  slot).
+- ``cp_indexer_decode``: :func:`cp_indexer_decode.try_cp_indexer_decode`, the
+  adapter hook of the GEN region (tried before cross_rank_topk_cluster_merge's when
+  ``cross_rank_topk_global_emit`` is selected).
+- ``remap_topk_to_global``: :func:`bridge.remap_topk_to_global`, the backend's
+  remap (rope_kv_append's when the call carries no plan).
+"""
+
+from ...dispatch import register
+from .bridge import (
+    GlobalEmitBridgePacket,
+    GlobalEmitContractError,
+    new_bridge_packet,
+    remap_topk_to_global,
+)
+from .cp_indexer_decode import cp_topk_from_local_logits, try_cp_indexer_decode
+
+register("cross_rank_topk_global_emit", "new_bridge_packet", new_bridge_packet)
+register("cross_rank_topk_global_emit", "cp_indexer_decode", try_cp_indexer_decode)
+register("cross_rank_topk_global_emit", "remap_topk_to_global", remap_topk_to_global)
+
+__all__ = [
+    "GlobalEmitBridgePacket",
+    "GlobalEmitContractError",
+    "cp_topk_from_local_logits",
+    "new_bridge_packet",
+    "remap_topk_to_global",
+    "try_cp_indexer_decode",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/bridge.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/bridge.py
new file mode 100644
index 000000000000..f77a04768b26
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/bridge.py
@@ -0,0 +1,286 @@
+# 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.
+"""cross_rank_topk_global_emit's side of rope_kv_append's top-k bridge: the packet, the plan and the remap.
+
+rope_kv_append (``rope_kv_append/ops.py``) collects the work of its fused bridge in a per-call
+``BridgePacket``: the indexer and the backend attach the native top-k copy of
+the call (``defer_shared_copy``, ``defer_reuse_copy``, ``defer_mtp_row_stash``)
+after the top-k, and ``remap_topk_to_global`` launches the bridge. cross_rank_topk_global_emit
+writes the bridge's outputs at the CP top-k's emission instead, so it has to know the
+call's copy and remap before the top-k runs:
+
+- ``predict_copy`` and ``predict_remap`` evaluate, at the indexer hook, the
+  rules the call sites apply afterwards (rope_kv_append's own ``defer_mtp_row_stash`` and
+  remap grouping rule, the backend's shared-copy condition), from host state
+  only;
+- the hook stores what it executed as the packet's ``GlobalEmitPlan``;
+- ``remap_topk_to_global`` (the backend's remap when ``cross_rank_topk_global_emit`` is selected)
+  consumes the plan once: it compares the packet's copy, source and top-k
+  tensor and the layer's remap decision with the plan and raises
+  ``GlobalEmitContractError`` on any difference (the prediction is a function of the
+  step's host state, so a wrong prediction fails the first eager forward,
+  before CUDA-graph capture); without a plan it is rope_kv_append's remap unchanged (draft
+  steps 1-4, every call cross_rank_topk_global_emit declined).
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from ..rope_kv_append import ops as rope_ops
+
+__all__ = [
+    "RemapPlan",
+    "GlobalEmitBridgePacket",
+    "GlobalEmitContractError",
+    "GlobalEmitPlan",
+    "copy_reason",
+    "new_bridge_packet",
+    "predict_copy",
+    "predict_remap",
+    "remap_topk_to_global",
+]
+
+
+class GlobalEmitContractError(RuntimeError):
+    """The bridge work cross_rank_topk_global_emit performed at the top-k differs from what the call sites requested."""
+
+
+@dataclass
+class RemapPlan:
+    """rope_kv_append's remap decision for a layer that owns an indexer.
+
+    ``grouped``: the leader of an active full + shared group remaps the whole
+    group (``layer_ids``, ``page_index_scale``) and publishes it in
+    ``metadata._group_remap_batched[leader]``; the layer's own rows are
+    ``out[slot]``. Otherwise the layer remaps alone (``page_index_scale``,
+    ``layer_offset``).
+    """
+
+    local_layer: int
+    grouped: bool
+    page_index_scale: int
+    leader: int = -1
+    slot: int = 0
+    layer_ids: Optional[torch.Tensor] = None
+    layer_offset: int = 0
+
+
+@dataclass
+class GlobalEmitPlan:
+    """What the cross_rank_topk_global_emit top-k wrote for one call: the id rows, the pool rows and the copy."""
+
+    ids: torch.Tensor
+    remap: RemapPlan
+    copy: Optional[rope_ops.BridgeCopy]
+    pool: Optional[torch.Tensor] = None
+
+
+@dataclass
+class GlobalEmitBridgePacket(rope_ops.BridgePacket):
+    """rope_kv_append's packet with the plan of the cross_rank_topk_global_emit top-k that served
+    the call (None: not served)."""
+
+    plan: Optional[GlobalEmitPlan] = None
+
+
+def new_bridge_packet(metadata, is_generation: bool) -> Optional[GlobalEmitBridgePacket]:
+    """A packet when rope_kv_append's fused bridge serves this ``sparse_attn_predict`` call, else None."""
+    if not rope_ops.bridge_eligible(metadata, is_generation):
+        return None
+    return GlobalEmitBridgePacket()
+
+
+def predict_remap(metadata, local_layer: int) -> Optional[RemapPlan]:
+    """rope_kv_append's remap decision (``rope_kv_append.ops.remap_topk_to_global``) for a layer with an indexer.
+
+    Returns None when rope_kv_append would take the native fallback (a group of more than
+    rope_kv_append's eight layers).
+    """
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
+
+    if dsa_backend._GROUP_REMAP and not metadata.in_mtp_draft_loop:
+        struct = metadata._ensure_group_remap_struct()
+        leader_of = struct.get("leader_of")
+        if leader_of is not None and local_layer < len(leader_of):
+            leader = leader_of[local_layer]
+            if leader >= 0 and struct["group_active"].get(leader, False):
+                layer_ids = struct["group_layer_ids"][leader]
+                if layer_ids.numel() > rope_ops._MAX_GROUP_SIZE:
+                    return None
+                return RemapPlan(
+                    local_layer=local_layer,
+                    grouped=True,
+                    page_index_scale=int(struct["group_scale"][leader]),
+                    leader=leader,
+                    slot=struct["slot_of"][local_layer],
+                    layer_ids=layer_ids,
+                )
+    scale, offset = metadata.kv_cache_manager.get_primary_pool_page_index_params(local_layer)
+    return RemapPlan(
+        local_layer=local_layer,
+        grouped=False,
+        page_index_scale=int(scale),
+        layer_offset=int(offset),
+    )
+
+
+def predict_copy(indexer, metadata, num_gen_tokens: int) -> Optional[rope_ops.BridgeCopy]:
+    """The native top-k copy the call sites attach to the packet after a served top-k.
+
+    Generation-only call with no context (the CP region's domain):
+
+    - MTP draft step with ``mtp_index_share``: the indexer's last-accepted-row
+      stash, exactly as ``rope_kv_append.ops.defer_mtp_row_stash`` decides it (evaluated on
+      a scratch packet); the backend then preserves the shared rows.
+    - otherwise: the backend's ``defer_shared_copy`` of the generation rows
+      into ``shared_topk_indices`` when that buffer exists.
+    """
+    num_contexts = metadata.num_contexts
+    num_generations = metadata.num_generations
+    shared = metadata.shared_topk_indices
+    if indexer.mtp_index_share and metadata.in_mtp_draft_loop:
+        num_accepted = metadata.mtp_num_accepted
+        if num_accepted is not None:
+            num_accepted = num_accepted[num_contexts : num_contexts + num_generations]
+        scratch = rope_ops.BridgePacket()
+        if rope_ops.defer_mtp_row_stash(
+            scratch,
+            shared,
+            num_contexts,
+            num_generations,
+            num_gen_tokens // num_generations,
+            num_accepted,
+        ):
+            return scratch.copy
+        return None
+    if shared is None:
+        return None
+    return rope_ops.BridgeCopy(dst=shared, row_start=metadata.num_ctx_tokens, rows=num_gen_tokens)
+
+
+def copy_reason(copy: Optional[rope_ops.BridgeCopy], rows: int, top_k: int) -> str:
+    """Why the cross_rank_topk_global_emit op cannot perform ``copy`` for ``rows`` top-k rows ("" when it can).
+
+    These are the checks of ``trtllm::glm_moe_dsa_topk_index_bridge``; a copy that fails them
+    is left to rope_kv_append's path, which reports it.
+    """
+    if copy is None or copy.rows <= 0:
+        return ""
+    dst = copy.dst
+    if not (
+        dst.is_cuda
+        and dst.dtype == torch.int32
+        and dst.dim() == 2
+        and dst.stride(1) == 1
+        and dst.shape[1] >= top_k
+        and dst.shape[0] >= copy.row_start + copy.rows
+        and copy.row_start >= 0
+    ):
+        return "copy destination outside the bridge's domain"
+    if copy.mode == 0:
+        return "" if copy.rows <= rows else "copy rows exceed the top-k rows"
+    if copy.mode != 1 or copy.next_n < 1 or copy.rows * copy.next_n > rows:
+        return "copy mode-1 rows exceed the top-k rows"
+    acc = copy.num_accepted
+    if acc is not None and not (
+        acc.is_cuda
+        and acc.dtype in (torch.int32, torch.int64)
+        and acc.dim() == 1
+        and acc.is_contiguous()
+        and acc.shape[0] >= copy.rows
+    ):
+        return "accepted counts outside the bridge's domain"
+    return ""
+
+
+def _same_tensor(a: Optional[torch.Tensor], b: Optional[torch.Tensor]) -> bool:
+    if a is None or b is None:
+        return a is None and b is None
+    return (
+        a.data_ptr() == b.data_ptr()
+        and a.dtype == b.dtype
+        and tuple(a.shape) == tuple(b.shape)
+        and tuple(a.stride()) == tuple(b.stride())
+    )
+
+
+def _same_copy(a: Optional[rope_ops.BridgeCopy], b: Optional[rope_ops.BridgeCopy]) -> bool:
+    if a is None or b is None:
+        return a is None and b is None
+    return (
+        _same_tensor(a.dst, b.dst)
+        and (a.row_start, a.rows, a.mode, a.next_n) == (b.row_start, b.rows, b.mode, b.next_n)
+        and _same_tensor(a.num_accepted, b.num_accepted)
+    )
+
+
+def _same_remap(a: Optional[RemapPlan], b: Optional[RemapPlan]) -> bool:
+    if a is None or b is None:
+        return False
+    return (a.local_layer, a.grouped, a.page_index_scale, a.leader, a.slot, a.layer_offset) == (
+        b.local_layer,
+        b.grouped,
+        b.page_index_scale,
+        b.leader,
+        b.slot,
+        b.layer_offset,
+    ) and a.layer_ids is b.layer_ids
+
+
+def _check_plan(plan: GlobalEmitPlan, backend, topk: torch.Tensor, metadata, packet) -> None:
+    """Raise ``GlobalEmitContractError`` when the call sites asked for other bridge work than
+    cross_rank_topk_global_emit performed."""
+    problems = []
+    if packet.source is not None:
+        problems.append("the packet reads other source rows")
+    if not _same_tensor(topk, plan.ids):
+        problems.append(
+            "the remapped top-k rows are not the rows cross_rank_topk_global_emit wrote"
+        )
+    if backend.indexer is None:
+        problems.append("the layer owns no indexer")
+    local_layer = backend.get_local_layer_idx(metadata)
+    if local_layer != plan.remap.local_layer:
+        problems.append(f"local layer {local_layer} != planned {plan.remap.local_layer}")
+    elif not _same_remap(predict_remap(metadata, local_layer), plan.remap):
+        problems.append("the layer's remap decision changed")
+    if not _same_copy(packet.copy, plan.copy):
+        problems.append(f"copy {packet.copy} != planned {plan.copy}")
+    if problems:
+        raise GlobalEmitContractError(
+            "cross_rank_topk_global_emit bridge plan mismatch: " + "; ".join(problems)
+        )
+
+
+def remap_topk_to_global(backend, topk: torch.Tensor, metadata, packet) -> torch.Tensor:
+    """``rope_kv_append.ops.remap_topk_to_global`` with the pool rows the
+    cross_rank_topk_global_emit top-k already wrote.
+
+    Without a plan (the call was not served by cross_rank_topk_global_emit) this is rope_kv_append's
+    remap. With one, the plan is consumed, checked against the packet and the layer, and
+    the pool rows are returned (and published for the group's followers, as
+    rope_kv_append's grouped remap does).
+    """
+    plan = getattr(packet, "plan", None)
+    if plan is None:
+        return rope_ops.remap_topk_to_global(backend, topk, metadata, packet)
+    packet.plan = None
+    _check_plan(plan, backend, topk, metadata, packet)
+    if plan.remap.grouped:
+        metadata._group_remap_batched[plan.remap.leader] = plan.pool
+        return plan.pool[plan.remap.slot]
+    return plan.pool
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_indexer_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_indexer_decode.py
new file mode 100644
index 000000000000..2a16162b99f4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_indexer_decode.py
@@ -0,0 +1,332 @@
+# 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.
+"""cross_rank_topk_global_emit: the cross_rank_topk GEN indexer region with
+cross_rank_topk_cluster_merge's top-k that also writes rope_kv_append's bridge outputs.
+
+Same region, policy, CP logits kernel and id output as cross_rank_topk_cluster_merge
+(``cross_rank_topk_cluster_merge/``), which are cross_rank_topk's (``indexer_topk/cross_rank_*.py``): rank
+``r`` scores the 64-token pages ``p % world == r``, and ``glm_moe_dsa_cross_rank_topk_global_emit``
+selects, exchanges and merges the candidates into the tie-aware exact top-``index_topk`` of every
+row, the same rows cross_rank_topk_cluster_merge writes. At the emission it also writes the call's
+rope_kv_append bridge outputs (``bridge.py``): the global KV-pool rows of the layer's remap group
+and the native top-k copy rows, the values ``trtllm::glm_moe_dsa_topk_index_bridge`` writes
+from the same rows.
+
+With ``indexer_layer_front`` also selected, the CP logits come from indexer_layer_front's launcher
+(``indexer_layer_front/cp_logits_runner.py``: cross_rank_topk's arguments, arena and values, the
+kernel waiting at entry and launched as the generation rope's programmatic
+dependent); cross_rank_topk_global_emit's top-k then runs on them unchanged.
+
+Numerical status: as cross_rank_topk_cluster_merge / cross_rank_topk (rank-consistent exact top-k
+over each page owner's indexer K); pool and copy rows bit-identical to cross_rank_topk_cluster_merge
++ rope_kv_append's bridge.
+"""
+
+from typing import Dict
+
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.logger import logger
+
+from ..indexer_topk import cross_rank_policy
+from . import bridge as global_emit_bridge
+from . import cp_workspace
+
+__all__ = [
+    "CP_DECISION_COUNTS",
+    "SUPPORTED_WORLD_SIZES",
+    "bridge_reason",
+    "cp_topk_from_local_logits",
+    "try_cp_indexer_decode",
+    "world_size_reason",
+]
+
+# As cross_rank_topk_cluster_merge / cross_rank_topk: MTP draft step 0 runs the same region on the
+# MTP layer.
+ALLOW_DRAFT_STEP = True
+
+# Tensor-parallel group sizes cross_rank_topk_global_emit's GPU tests verify
+# (cross_rank_topk_cluster_merge's).
+SUPPORTED_WORLD_SIZES = (4, 8)
+
+# Block size of the KV pool the kernel's remap assumes (the CP page).
+PAGE_TOKENS = 64
+
+
+def world_size_reason(world: int) -> str:
+    """Why cross_rank_topk_global_emit declines a group of ``world`` ranks ("" when it may engage)."""
+    if world not in SUPPORTED_WORLD_SIZES:
+        return (
+            f"tensor-parallel size {world} not in cross_rank_topk_global_emit's verified sizes "
+            f"{SUPPORTED_WORLD_SIZES}"
+        )
+    return ""
+
+
+CP_DECISION_COUNTS: Dict[str, int] = {}
+
+
+def _record_decision(eligible: bool, reason: str) -> None:
+    key = "engaged" if eligible else reason
+    count = CP_DECISION_COUNTS.get(key, 0)
+    CP_DECISION_COUNTS[key] = count + 1
+    if count == 0:
+        state = "engaged" if eligible else f"declined ({reason})"
+        logger.info(
+            f"[glm_moe_dsa] cross_rank_topk_global_emit/cp_indexer_decode: first '{state}' decision on this rank"
+        )
+
+
+def bridge_reason(indexer, metadata, packet, num_gen_tokens: int):
+    """cross_rank_topk_global_emit's bridge preconditions: ``(reason, remap, copy)``, reason "" when
+    cross_rank_topk_global_emit may engage.
+
+    Host state only, identical on every rank of the group for the same call.
+    """
+    if not isinstance(packet, global_emit_bridge.GlobalEmitBridgePacket):
+        return "no cross_rank_topk_global_emit bridge packet for this call", None, None
+    if packet.copy is not None or packet.source is not None or packet.plan is not None:
+        return "the packet already holds bridge work", None, None
+    kvm = metadata.kv_cache_manager
+    if kvm.tokens_per_block != PAGE_TOKENS:
+        return f"KV pool block size {kvm.tokens_per_block} != {PAGE_TOKENS}", None, None
+    remap = global_emit_bridge.predict_remap(metadata, kvm.layer_offsets[indexer.layer_idx])
+    if remap is None:
+        return "remap group larger than the bridge's group limit", None, None
+    copy = global_emit_bridge.predict_copy(indexer, metadata, num_gen_tokens)
+    reason = global_emit_bridge.copy_reason(copy, num_gen_tokens, indexer.top_k.top_k)
+    if reason:
+        return reason, None, None
+    return "", remap, copy
+
+
+def cp_topk_from_local_logits(
+    local_logits: torch.Tensor,
+    kv_lens: torch.Tensor,
+    out: torch.Tensor,
+    workspace,
+    next_n: int,
+    *,
+    req_idx: torch.Tensor,
+    req_idx_bias: int,
+    block_table: torch.Tensor,
+    stride_factor: int,
+    layer_ids=None,
+    layer_id: int = 0,
+    copy=None,
+    debug_times=None,
+) -> torch.Tensor:
+    """Select, exchange, merge and write the bridge outputs from this rank's local logits.
+
+    Args:
+        local_logits: ``[next_n, local_capacity]`` float32 logits of this
+            rank's local columns.
+        kv_lens: ``[1]`` int32 global indexer KV length.
+        out: ``[next_n, index_topk]`` int32 id rows.
+        workspace: ``cp_workspace.CpWorkspace`` of the group.
+        next_n: Verify rows of the request.
+        req_idx, req_idx_bias, block_table, stride_factor, layer_ids,
+            layer_id, copy: the arguments of ``rope_kv_append.ops.topk_bridge`` for these
+            rows (``copy`` a ``rope_kv_append.ops.BridgeCopy`` or None).
+        debug_times: Optional ``[next_n, 8, 64]`` int64 phase stamps.
+
+    Returns:
+        The pool rows: ``[len(layer_ids), next_n, index_topk]`` with
+        ``layer_ids``, else ``[next_n, index_topk]``.
+    """
+    return torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_global_emit(
+        local_logits,
+        kv_lens,
+        out,
+        workspace.state,
+        workspace.mc_ptr,
+        workspace.uc_ptr,
+        workspace.rank,
+        workspace.world,
+        next_n,
+        req_idx,
+        req_idx_bias,
+        block_table,
+        PAGE_TOKENS,
+        stride_factor,
+        layer_ids,
+        layer_id,
+        None if copy is None else copy.dst,
+        0 if copy is None else copy.row_start,
+        0 if copy is None else copy.rows,
+        0 if copy is None else copy.mode,
+        1 if copy is None else copy.next_n,
+        None if copy is None else copy.num_accepted,
+        debug_times,
+    )
+
+
+def try_cp_indexer_decode(
+    indexer,
+    metadata,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+    bridge=None,
+) -> bool:
+    """Adapter entry: run the GEN region with cross_rank_topk_global_emit and record its bridge work on ``bridge``.
+
+    The decision is cross_rank_topk_cluster_merge's (cross_rank_topk's
+    ``indexer_topk.cross_rank_policy`` on the same host state, cross_rank_topk_cluster_merge's group
+    sizes, cross_rank_topk_global_emit's own workspace creation), plus cross_rank_topk_global_emit's
+    bridge preconditions (``bridge_reason``), all taken before any device work. A ``False`` return
+    leaves no side effect; every rank of the group reaches the same decision for the same forward,
+    and the caller then tries cross_rank_topk_cluster_merge's region. ``True``: ``out`` holds the id
+    rows and ``bridge.plan`` the pool rows and the copy cross_rank_topk_global_emit wrote, which
+    ``bridge.remap_topk_to_global`` consumes.
+    """
+    from ...dispatch import is_enabled, log_capability
+    from ..indexer_topk.cross_rank_logits_runner import local_capacity
+
+    # indexer_layer_front's logits launcher is imported only with indexer_layer_front selected.
+    indexer_front_logits = is_enabled("indexer_layer_front")
+    if indexer_front_logits:
+        from ..indexer_layer_front.cp_logits_runner import ensure_compiled, paged_mqa_logits_cp
+    else:
+        from ..indexer_topk.cross_rank_logits_runner import ensure_compiled, paged_mqa_logits_cp
+
+    num_contexts = metadata.num_contexts
+    num_generations = metadata.num_generations
+    gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
+    uniform = num_generations > 0 and bool(
+        (gen_seq_lens == num_gen_tokens // num_generations).all()
+    )
+    top_k = indexer.top_k
+    k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(indexer.layer_idx)
+    indexer_block_table = metadata.indexer_k_cache_block_offsets
+    mapping = getattr(metadata, "mapping", None)
+    max_seq_len = metadata.get_indexer_max_seq_len()
+    world = mapping.tp_size if mapping is not None else 1
+    info = cp_workspace.op_info(world) if world > 0 else None
+    capacity = local_capacity(max_seq_len, world) if world > 0 else 0
+    # Facts of this rank's build and device, agreed by the workspace creation.
+    local_reason = cross_rank_policy.rank_local_reason(
+        op_available=info is not None,
+        local_capacity=capacity,
+        max_local_capacity=info[cp_workspace.INFO_MAX_LOCAL_CAPACITY] if info is not None else 0,
+        sm_version=get_sm_version(),
+        cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE,
+    ) or (cp_workspace.cluster_reason(info) if info is not None else "")
+    decision = cross_rank_policy.decide_cp_uniform(
+        world_size=world,
+        local_capacity=capacity,
+        allow_draft_step=ALLOW_DRAFT_STEP,
+        native_uses_cute_dsl_logits=indexer.use_cute_dsl_paged_mqa_logits,
+        use_fp4_indexer=indexer.use_fp4,
+        use_gvr_emission=indexer.use_gvr_emission,
+        needs_gvr_prior=top_k.needs_gvr_prior,
+        is_ragged_verify=metadata.is_ragged_verify,
+        expanded_mtp_buffers=metadata.use_expanded_buffers_for_mtp,
+        in_mtp_draft_loop=metadata.in_mtp_draft_loop,
+        attention_dp=bool(mapping is not None and mapping.enable_attention_dp),
+        topk_decode_impl=top_k.decode_implementation.value,
+        topk_self_sampling=top_k.gvr_self_sampling,
+        index_topk=top_k.top_k,
+        compress_ratio=top_k.compress_ratio,
+        num_heads=indexer.n_heads,
+        head_dim=indexer.head_dim,
+        query_dtype=str(q_fp8.dtype),
+        query_shape=tuple(q_fp8.shape),
+        query_contiguous=q_fp8.is_contiguous(),
+        weights_dtype=str(weights.dtype),
+        weights_shape=tuple(weights.shape),
+        weights_contiguous=weights.is_contiguous(),
+        k_cache_dtype=str(k_cache.dtype),
+        k_cache_shape=tuple(k_cache.shape),
+        block_table_cols=indexer_block_table.shape[1],
+        max_seq_len=max_seq_len,
+        num_contexts=num_contexts,
+        token_offset=token_offset,
+        num_gen_tokens=num_gen_tokens,
+        num_generations=num_generations,
+        uniform_verify_width=uniform,
+    )
+    eligible, reason = decision.eligible, decision.reason
+    remap = copy = None
+    if eligible and world_size_reason(world):
+        eligible, reason = False, world_size_reason(world)
+    if eligible:
+        reason, remap, copy = bridge_reason(indexer, metadata, bridge, num_gen_tokens)
+        eligible = not reason
+    if eligible:
+        # The logits kernel is compiled before the group agrees on the
+        # workspace, so no rank waits in an exchange while a peer compiles.
+        workspace = cp_workspace.get_workspace(
+            mapping,
+            prepare=lambda: ensure_compiled(
+                q_fp8.device,
+                k_cache.shape[1],
+                indexer.n_heads,
+                indexer.head_dim,
+                num_gen_tokens // num_generations,
+                world,
+            ),
+            local_reason=local_reason,
+        )
+        if workspace is None:
+            why = (
+                cp_workspace.failure_reason(mapping)
+                or "workspace not created before CUDA-graph capture"
+            )
+            eligible, reason = False, f"workspace unavailable: {why}"
+    log_capability("cross_rank_topk_global_emit", "cp_indexer_decode", eligible, reason)
+    if indexer_front_logits:
+        log_capability("indexer_layer_front", "cp_indexer_decode_global_emit", eligible, reason)
+    _record_decision(eligible, reason)
+    if not eligible:
+        return False
+
+    next_n = num_gen_tokens // num_generations
+    gen_rows = slice(token_offset, token_offset + num_gen_tokens)
+    gen_reqs = slice(num_contexts, num_contexts + num_generations)
+    kv_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+    assert kv_lens is not None
+    local_logits = paged_mqa_logits_cp(
+        q_fp8[gen_rows].view(num_generations, next_n, indexer.n_heads, indexer.head_dim),
+        k_cache,
+        weights[gen_rows],
+        kv_lens,
+        indexer_block_table[gen_reqs],
+        max_seq_len,
+        workspace.rank,
+        workspace.world,
+    )
+    # The bridge inputs of rope_kv_append.ops.topk_bridge for these rows.
+    tokens_per_block = metadata.kv_cache_manager.tokens_per_block
+    pool = cp_topk_from_local_logits(
+        local_logits,
+        kv_lens,
+        out,
+        workspace,
+        next_n,
+        req_idx=metadata.req_idx_per_token[metadata.num_ctx_tokens : metadata.num_tokens],
+        req_idx_bias=num_contexts,
+        block_table=metadata.block_table[num_contexts : metadata.num_seqs],
+        stride_factor=remap.page_index_scale * tokens_per_block,
+        layer_ids=remap.layer_ids if remap.grouped else None,
+        layer_id=0 if remap.grouped else remap.layer_offset,
+        copy=copy,
+    )
+    bridge.plan = global_emit_bridge.GlobalEmitPlan(ids=out, remap=remap, copy=copy, pool=pool)
+    return True
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_workspace.py
new file mode 100644
index 000000000000..e32134459fee
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/cp_workspace.py
@@ -0,0 +1,246 @@
+# 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.
+"""Exchange workspace of the cross_rank_topk_global_emit context-partitioned top-k (with the bridge outputs).
+
+cross_rank_topk_global_emit owns one multicast workspace per tensor-parallel group, allocated with
+the native ``McastGPUBuffer`` and used by nothing else (cross_rank_topk's and
+cross_rank_topk_cluster_merge's workspaces are separate ones). Its layout and Lamport protocol are
+cross_rank_topk_cluster_merge's, defined by the ``glm_moe_dsa_cross_rank_topk_global_emit`` kernel
+(``crossRankTopkGlobalEmit.h``): every 32-bit word starts as ``0xFFFFFFFF`` and the kernel rotates
+three buffers with per-CTA epoch words held in a separate per-rank state tensor.
+
+Creation is collective over the group and only happens outside CUDA-graph
+capture, exactly as cross_rank_topk's (``indexer_topk/cross_rank_workspace.py``): each rank runs the
+caller's host-side preparation, the ranks agree on the capability (compiled
+op, supported world size, MPI, successful preparation, no rank-local reason,
+and enough co-resident clusters for the largest row count), allocate, and
+agree on success; if any step fails on any rank, every rank records the
+failure, so the collective count and order stay identical.
+"""
+
+import os
+import threading
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.logger import logger
+
+# Indices of torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_global_emit_workspace_info().
+INFO_WORKSPACE_WORDS = 0
+INFO_TOPK = 1
+INFO_MAX_RANKS = 2
+INFO_MAX_ROWS = 3
+INFO_PAGE_TOKENS = 4
+INFO_MAX_LOCAL_CAPACITY = 7
+INFO_STATE_WORDS = 8
+INFO_TIMEOUT_MS = 10
+INFO_MAX_ACTIVE_CLUSTERS = 11
+INFO_MAX_GROUP_SIZE = 12
+# First epoch word of the state (state[EPOCH_BASE + row * 8 + cta]).
+EPOCH_BASE = 4
+
+
+@dataclass
+class CpWorkspace:
+    """One rank's view of the cross_rank_topk_global_emit exchange workspace.
+
+    Attributes:
+        handle: The native ``McastGPUBuffer`` (keeps the mapping alive).
+        mc: Multicast int32 view of the whole workspace.
+        uc: This rank's unicast int32 view of the whole workspace.
+        state: int32 device state: [2] second rounds taken, [3] error word
+            written before a timeout abort, then one epoch word (exchange
+            launches) per (row, CTA) from ``EPOCH_BASE``.
+        rank: Rank in the group.
+        world: Group size.
+        comm: The group communicator the handle was created with (kept alive).
+    """
+
+    handle: object
+    mc: torch.Tensor
+    uc: torch.Tensor
+    state: torch.Tensor
+    rank: int
+    world: int
+    comm: object = None
+
+    @property
+    def mc_ptr(self) -> int:
+        return self.mc.data_ptr()
+
+    @property
+    def uc_ptr(self) -> int:
+        return self.uc.data_ptr()
+
+
+_LOCK = threading.Lock()
+_WORKSPACES: Dict[Tuple, CpWorkspace] = {}
+_FAILED: Dict[Tuple, str] = {}
+
+
+_OP_INFO: Dict[Tuple[int, int], list] = {}
+
+
+def op_info(world_size: int = 8) -> Optional[list]:
+    """Constants of the compiled kernel for a group size, or None when the op is not built.
+
+    Queried once per group size and current device (the region calls it on
+    every eager forward); a missing op is not cached, so an op loaded later
+    is seen.
+    """
+    key = (world_size, torch.cuda.current_device() if torch.cuda.is_available() else -1)
+    info = _OP_INFO.get(key)
+    if info is None:
+        try:
+            info = list(
+                torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_global_emit_workspace_info(world_size)
+            )
+        except (AttributeError, RuntimeError):
+            return None
+        _OP_INFO[key] = info
+    return list(info)
+
+
+def cluster_reason(info: Optional[list]) -> str:
+    """Why this device cannot keep every row's cluster resident ("" when it can)."""
+    if info is None or len(info) <= INFO_MAX_ACTIVE_CLUSTERS:
+        return "glm_moe_dsa_cross_rank_topk_global_emit op is not built"
+    clusters, rows = info[INFO_MAX_ACTIVE_CLUSTERS], info[INFO_MAX_ROWS]
+    if clusters < rows:
+        return f"{clusters} co-resident 8-CTA clusters < {rows} rows"
+    return ""
+
+
+def _key(mapping) -> Tuple:
+    return (mapping.world_size, mapping.rank, mapping.tp_size, mapping.pp_size, mapping.cp_size)
+
+
+def failure_reason(mapping) -> Optional[str]:
+    """Reason recorded when this group's workspace could not be created."""
+    return _FAILED.get(_key(mapping))
+
+
+def get_workspace(
+    mapping, prepare: Optional[Callable[[], None]] = None, local_reason: str = ""
+) -> Optional[CpWorkspace]:
+    """Return the group's workspace, creating it collectively on first use.
+
+    Must be reached by every rank of the tensor-parallel group in the same
+    order. Returns None when the workspace is unavailable on any rank, and
+    during CUDA-graph capture when it was not created before.
+
+    Args:
+        mapping: The rank's ``Mapping``.
+        prepare: Optional host-side preparation run on every rank before the
+            capability agreement of the creation; an exception counts as a
+            local failure that every rank agrees on.
+        local_reason: Why this rank's build or device cannot run cross_rank_topk_global_emit (""
+        when
+            it can); part of the creation's capability agreement.
+    """
+    key = _key(mapping)
+    ws = _WORKSPACES.get(key)
+    if ws is not None or key in _FAILED:
+        return ws
+    if torch.cuda.is_current_stream_capturing():
+        return None
+    with _LOCK:
+        if key in _WORKSPACES or key in _FAILED:
+            return _WORKSPACES.get(key)
+        ws, reason = _create(mapping, prepare, local_reason)
+        if ws is None:
+            _FAILED[key] = reason
+            logger.info(
+                f"[glm_moe_dsa] cross_rank_topk_global_emit/workspace: not created ({reason})"
+            )
+        else:
+            _WORKSPACES[key] = ws
+            info = op_info(ws.world) or []
+            timeout = info[INFO_TIMEOUT_MS] if len(info) > INFO_TIMEOUT_MS else "n/a"
+            clusters = (
+                info[INFO_MAX_ACTIVE_CLUSTERS] if len(info) > INFO_MAX_ACTIVE_CLUSTERS else "n/a"
+            )
+            logger.info(
+                f"[glm_moe_dsa] cross_rank_topk_global_emit/workspace: created ({ws.uc.numel() * 4} bytes, "
+                f"rank {ws.rank}/{ws.world}, exchange wait timeout {timeout} ms, "
+                f"co-resident clusters {clusters})"
+            )
+        return ws
+
+
+def _local_capability_reason(mapping, info, prepare) -> str:
+    """This rank's op, group size, cluster residency and preparation check ("" when capable)."""
+    if info is None:
+        return "glm_moe_dsa_cross_rank_topk_global_emit op is not built"
+    if mapping.tp_size > info[INFO_MAX_RANKS]:
+        return f"tp_size {mapping.tp_size} > {info[INFO_MAX_RANKS]}"
+    reason = cluster_reason(info)
+    if reason:
+        return reason
+    if prepare is not None:
+        try:
+            prepare()
+        except Exception as e:  # noqa: BLE001 - any local failure must be agreed before use
+            return f"preparation failed ({type(e).__name__}: {e})"
+    return ""
+
+
+def _create(mapping, prepare=None, local_reason="") -> Tuple[Optional[CpWorkspace], str]:
+    from tensorrt_llm._torch.distributed.ops import (
+        _get_mnnvl_workspace_comm,
+        _mnnvl_device_index,
+        _mnnvl_workspace_all_succeeded,
+    )
+    from tensorrt_llm._utils import mpi_disabled
+
+    if mpi_disabled():
+        # Without MPI every rank takes this branch before any collective.
+        return None, "MPI is disabled"
+    comm = _get_mnnvl_workspace_comm(mapping)
+    info = op_info(mapping.tp_size)
+    if not local_reason:
+        local_reason = _local_capability_reason(mapping, info, prepare)
+    if not _mnnvl_workspace_all_succeeded(comm, not local_reason):
+        return None, local_reason or "capability not agreed by every rank"
+
+    words = info[INFO_WORKSPACE_WORDS]
+    error: Optional[Exception] = None
+    ws: Optional[CpWorkspace] = None
+    try:
+        from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer
+
+        use_fabric = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node()
+        device = _mnnvl_device_index(mapping)
+        handle = McastGPUBuffer(
+            words * 4, mapping.tp_size, mapping.tp_rank, device, use_fabric, comm.py2f()
+        )
+        uc = handle.get_uc_buffer(mapping.tp_rank, [words], torch.int32, 0)
+        mc = handle.get_mc_buffer([words], torch.int32, 0)
+        with torch.inference_mode():
+            uc.fill_(-1)
+            state = torch.zeros(
+                info[INFO_STATE_WORDS], dtype=torch.int32, device=torch.device("cuda", device)
+            )
+        torch.cuda.synchronize()
+        ws = CpWorkspace(handle, mc, uc, state, mapping.tp_rank, mapping.tp_size, comm)
+    except Exception as e:  # noqa: BLE001 - any local failure must be agreed before use
+        error = e
+    # Also fences the initialization: no rank sends before every rank cleared
+    # its buffer.
+    if not _mnnvl_workspace_all_succeeded(comm, error is None):
+        return None, f"allocation failed on a rank ({type(error).__name__ if error else 'peer'})"
+    return ws, ""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/__init__.py
new file mode 100644
index 000000000000..4ed230c018dd
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/__init__.py
@@ -0,0 +1,35 @@
+# 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.
+"""decode_prologue: the target decode-metadata prologue and the target embedding in one launch.
+
+decode_prologue extends draft_loop (``draft_loop``): its launch reuses draft_loop's DeepGEMM
+schedule reproduction (``draft_loop/draft_step.py``) and draft_loop's embedding semantics
+(``draft_loop/embedding.py``) for the target step instead of the draft steps.
+
+Entry point: :class:`prologue.DecodePrologueHook`, which ``runtime.make_runtime_hooks``
+declares to the engine as the step loop's ``decode_prologue`` and hands to the decode
+target. Calling it runs :func:`prologue.target_prologue` at the start of
+``DecoderRunner._preprocess_inputs`` (True when it replaced the native prologue, False
+to keep the native code); its ``take_embedding`` (:func:`prologue.take_embedding`)
+returns the rows the launch embedded to the decode target's forward, or None.
+"""
+
+from ...dispatch import register
+from .prologue import DecodePrologueHook, prologue_ineligibility, take_embedding, target_prologue
+
+register("decode_prologue", "target_prologue", target_prologue)
+register("decode_prologue", "take_embedding", take_embedding)
+
+__all__ = ["DecodePrologueHook", "prologue_ineligibility", "take_embedding", "target_prologue"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/prologue.py
new file mode 100644
index 000000000000..35eb784bd622
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/prologue.py
@@ -0,0 +1,596 @@
+# 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.
+"""Single-launch target decode-metadata prologue and target embedding (decode_prologue).
+
+A decode step's captured forward begins with ``DecoderRunner._preprocess_inputs``:
+``on_update_kv_lens()``, the overlap scheduler's accepted-count correction of
+the position ids and the KV lengths, and ``on_update_kv_lens()`` again. On the
+fused DSA decode-metadata path each call writes the 2-D KV-length view, three
+DeepGEMM paged-MQA-logits schedules (each computed into a fresh tensor and
+copied into its persistent buffer), the expanded lengths of the DSL atom split,
+the fused decode metadata (request index, indexer K-cache slot mappings,
+generation indptrs) and, when the indexer is skipped, the dense top-k rows.
+The target model then embeds the input ids with the replicated ``Embedding``
+(a range mask, ``F.embedding`` and ``masked_fill_``). About thirty serial
+launches precede the first layer.
+
+``target_prologue`` produces the identical device state with one launch:
+program 0 applies the correction and derives every metadata buffer from the
+corrected lengths held in registers (the state the second native call leaves;
+the first call's writes are all overwritten by the second); programs 1.. copy
+the embedding rows (``weight[id]``, or +0.0 for an id outside the table). The
+host replays the host-side statements of both native calls. The schedule is
+draft_loop's reproduction of DeepGEMM's SM100 partition, confirmed against the native
+op by draft_loop's ``schedule_self_check``. Every decision is made on the host before
+any device work; outside the supported domain the native code runs unchanged.
+
+The embedding rows are handed to the model forward through
+``take_embedding``: the launch records ``(input_ids, weight, output)`` on the
+step's attention metadata, and the model forward consumes it once for exactly
+that ``input_ids`` tensor and embedding weight (otherwise it embeds natively).
+"""
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ..draft_loop import draft_step as _draft_step
+from ..draft_loop import embedding as _draft_embedding
+from ..draft_loop.draft_step import _write_schedule
+
+_MAX_BATCH = 64
+_NUM_WARPS = 4
+_BLOCK_H = 2048
+# Elements of one dense top-k tile (BLOCK_N rows x BLOCK_K columns).
+_DENSE_TILE = 4096
+_PACKET_ATTR = "_decode_prologue_embedding"
+
+
+@triton.jit
+def _target_prologue_kernel(
+    kv_lens_ptr,  # int32 [>= B] (in/out)
+    kv_off_ptr,  # int32 [>= B] previous_kv_lens_offsets (CORRECT)
+    seq_lens_ptr,  # int32 [>= B]
+    pos_ptr,  # int32 [>= N_POS] position ids of the step (in/out, CORRECT)
+    pos_off_ptr,  # int32 [>= N_POS] previous_pos_id_offsets (CORRECT)
+    kv2d_ptr,  # int32 [max_seqs, cap] (out)
+    kv2d_stride,
+    sched_ptr,  # int32 [(NUM_SMS + 1) * 2] (out)
+    sched_full_ptr,  # int32 [(NUM_SMS + 1) * 2] (out)
+    kv_exp_ptr,  # int32 [>= B * EXPAND] (out, EXPAND > 1)
+    sched_exp_ptr,  # int32 [(NUM_SMS + 1) * 2] (out, EXPAND > 1)
+    block_offsets_ptr,  # int32 [max_seqs, max_blocks]
+    block_offsets_stride_0,
+    req_idx_ptr,  # int32 [>= num_tokens] (out)
+    slot_fp8_ptr,  # int64 [>= num_tokens] (out)
+    slot_scale_ptr,  # int64 [>= num_tokens] (out)
+    gen_kv_indptr_ptr,  # int64 [>= B + 1] (out)
+    gen_cached_indptr_ptr,  # int64 [>= B + 1] (out)
+    topk_ptr,  # int32 [>= num_tokens, topk_stride] (out, DENSE)
+    topk_stride,
+    ids_ptr,  # int32/int64 [N_EMB] (EMBED)
+    weight_ptr,  # [vocab_size, HIDDEN] (EMBED)
+    emb_ptr,  # [N_EMB, HIDDEN] (out, EMBED)
+    vocab_size,
+    weight_row_stride,
+    emb_row_stride,
+    num_tokens,
+    N_POS: tl.constexpr,
+    B: tl.constexpr,
+    CORRECT: tl.constexpr,
+    WIDTH: tl.constexpr,
+    EXPAND: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SPLIT_KV: tl.constexpr,
+    max_blocks: tl.constexpr,
+    block_offsets_stride_1: tl.constexpr,
+    tokens_per_block: tl.constexpr,
+    data_bytes_per_token: tl.constexpr,
+    scale_size: tl.constexpr,
+    block_stride: tl.constexpr,
+    scale_base_offset: tl.constexpr,
+    DENSE: tl.constexpr,
+    TOPK: tl.constexpr,
+    EMBED: tl.constexpr,
+    HIDDEN: tl.constexpr,
+    NCB: tl.constexpr,
+    BLOCK_B: tl.constexpr,
+    BLOCK_W: tl.constexpr,
+    BLOCK_E: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_SM: tl.constexpr,
+    BLOCK_P: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_B)
+        mask_b = offs_b < B
+
+        # ----- accepted-count correction (in place) and the step's lengths ---
+        kv = tl.load(kv_lens_ptr + offs_b, mask=mask_b, other=0)
+        if CORRECT:
+            kv = kv + tl.load(kv_off_ptr + offs_b, mask=mask_b, other=0)
+            tl.store(kv_lens_ptr + offs_b, kv, mask=mask_b)
+            offs_p = tl.arange(0, BLOCK_P)
+            mask_p = offs_p < N_POS
+            pos = tl.load(pos_ptr + offs_p, mask=mask_p, other=0)
+            pos = pos + tl.load(pos_off_ptr + offs_p, mask=mask_p, other=0)
+            tl.store(pos_ptr + offs_p, pos, mask=mask_p)
+        kv = tl.where(mask_b, kv, 0)
+        seq = tl.load(seq_lens_ptr + offs_b, mask=mask_b, other=0)
+        seq = tl.where(mask_b, seq, 0)
+
+        # ----- 2-D KV lengths and the paged-MQA-logits schedules -------------
+        offs_w = tl.arange(0, BLOCK_W)
+        kv2d_mask = mask_b[:, None] & (offs_w[None, :] < WIDTH)
+        tl.store(
+            kv2d_ptr + offs_b[:, None] * kv2d_stride + offs_w[None, :],
+            kv[:, None] + tl.zeros((BLOCK_B, BLOCK_W), tl.int32),
+            mask=kv2d_mask,
+        )
+        _write_schedule(kv, mask_b, B, sched_ptr, 1, NUM_SMS, SPLIT_KV, BLOCK_SM)
+        _write_schedule(kv, mask_b, B, sched_full_ptr, WIDTH, NUM_SMS, SPLIT_KV, BLOCK_SM)
+        if EXPAND > 1:
+            offs_e = tl.arange(0, BLOCK_E)
+            mask_e = offs_e < B * EXPAND
+            owner = (offs_e[:, None] // EXPAND) == offs_b[None, :]
+            kv_e = tl.sum(tl.where(owner, kv[None, :], 0), 1)
+            tl.store(kv_exp_ptr + offs_e, kv_e, mask=mask_e)
+            _write_schedule(kv_e, mask_e, B * EXPAND, sched_exp_ptr, 1, NUM_SMS, SPLIT_KV, BLOCK_SM)
+
+        # ----- fused DSA decode metadata (the native kernel's arithmetic) ----
+        kv64 = kv.to(tl.int64)
+        seq64 = seq.to(tl.int64)
+        tl.store(gen_kv_indptr_ptr, tl.full((), 0, tl.int64))
+        tl.store(gen_cached_indptr_ptr, tl.full((), 0, tl.int64))
+        tl.store(gen_kv_indptr_ptr + 1 + offs_b, tl.cumsum(kv64, 0), mask=mask_b)
+        tl.store(gen_cached_indptr_ptr + 1 + offs_b, tl.cumsum(kv64 - seq64, 0), mask=mask_b)
+
+        flat_start = tl.cumsum(seq64, 0) - seq64
+        start_pos = kv64 - seq64
+        offs_t = tl.arange(0, BLOCK_T).to(tl.int64)
+        g = flat_start[:, None] + offs_t[None, :]
+        mask_t = mask_b[:, None] & (offs_t[None, :] < seq64[:, None]) & (g < num_tokens)
+        tl.store(
+            req_idx_ptr + g, offs_b[:, None] + tl.zeros((BLOCK_B, BLOCK_T), tl.int32), mask=mask_t
+        )
+        gpos = start_pos[:, None] + offs_t[None, :]
+        blk = gpos // tokens_per_block
+        pos_in_blk = gpos % tokens_per_block
+        pos_in_blk = tl.where(pos_in_blk < 0, pos_in_blk + tokens_per_block, pos_in_blk)
+        blk = tl.minimum(tl.maximum(blk, 0), max_blocks - 1)
+        bid = tl.load(
+            block_offsets_ptr
+            + offs_b[:, None].to(tl.int64) * block_offsets_stride_0
+            + blk * block_offsets_stride_1,
+            mask=mask_t,
+            other=0,
+        ).to(tl.int64)
+        tl.store(
+            slot_fp8_ptr + g, bid * block_stride + pos_in_blk * data_bytes_per_token, mask=mask_t
+        )
+        tl.store(
+            slot_scale_ptr + g,
+            bid * block_stride + scale_base_offset + pos_in_blk * scale_size,
+            mask=mask_t,
+        )
+
+        # ----- dense top-k rows of a skipped indexer ------------------------
+        if DENSE:
+            # position(g) = (kv - seq)[b] + (g - flat_start[b]) for the request b
+            # whose token range holds g (``_get_dense_topk_indices``).
+            offs_n = tl.arange(0, BLOCK_N).to(tl.int64)
+            seq_end = tl.cumsum(seq64, 0)
+            owned = (seq_end[None, :] <= offs_n[:, None]) & mask_b[None, :]
+            req = tl.sum(owned.to(tl.int32), 1)
+            onehot = req[:, None] == offs_b[None, :]
+            shift = tl.sum(tl.where(onehot, (start_pos - flat_start)[None, :], 0), 1)
+            tok_pos = offs_n + shift
+            mask_n = offs_n < num_tokens
+            for k0 in tl.static_range(0, TOPK, BLOCK_K):
+                offs_k = k0 + tl.arange(0, BLOCK_K)
+                rows = tl.where(
+                    offs_k[None, :].to(tl.int64) <= tok_pos[:, None], offs_k[None, :], -1
+                )
+                tl.store(
+                    topk_ptr + offs_n[:, None] * topk_stride + offs_k[None, :],
+                    rows.to(tl.int32),
+                    mask=mask_n[:, None] & (offs_k[None, :] < TOPK),
+                )
+    else:
+        if EMBED:
+            e = pid - 1
+            row = e // NCB
+            cb = e % NCB
+            token = tl.load(ids_ptr + row).to(tl.int64)
+            valid = (token >= 0) & (token < vocab_size)
+            safe_token = tl.where(valid, token, 0)
+            cols = cb * BLOCK_H + tl.arange(0, BLOCK_H)
+            col_mask = cols < HIDDEN
+            values = tl.load(
+                weight_ptr + safe_token * weight_row_stride + cols,
+                mask=col_mask & valid,
+                other=0.0,
+            )
+            tl.store(emb_ptr + row.to(tl.int64) * emb_row_stride + cols, values, mask=col_mask)
+
+
+def _metadata_ineligibility(md, B: int) -> Optional[str]:
+    """Why the fused-metadata part cannot replace the native calls, or None."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import is_dsa_cache_manager
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+        _fused_dsa_meta_enabled,
+    )
+
+    if not isinstance(md, DSAtrtllmAttentionMetadata):
+        return "metadata is not DSAtrtllmAttentionMetadata"
+    cls = type(md)
+    for name in ("on_update_kv_lens", "_compute_kv_lens_row_reorder", "get_indexer_kv_lens"):
+        if getattr(cls, name) is not getattr(DSAtrtllmAttentionMetadata, name):
+            return f"metadata overrides {name}"
+    if md.kv_cache_manager is None or not is_dsa_cache_manager(md.kv_cache_manager):
+        return "KV cache manager is not a DSA cache manager"
+    if not _fused_dsa_meta_enabled():
+        return "native fused DSA metadata disabled"
+    if md.num_contexts != 0 or md.num_generations <= 0 or md.num_tokens <= 0:
+        return "not a generation-only step"
+    if md.use_fp8_ds_mla or getattr(md, "fp4_mla_state", None) is not None:
+        return "FP8-DS/FP4 MLA metadata state"
+    if getattr(md, "enable_helix", False):
+        return "helix metadata"
+    if md.enable_ragged_verification or md.is_ragged_verify:
+        return "ragged verification"
+    if md._indexer_compress_ratio > 1:
+        return "compressed indexer lengths"
+    if md.use_expanded_buffers_for_mtp:
+        return "expanded MTP buffers"
+    if md.max_draft_tokens <= 0:
+        return "no draft tokens"
+    if B <= 0 or B > _MAX_BATCH or B != md.num_seqs or B != md.num_generations:
+        return f"batch size {B} outside the fused domain"
+    if md.needs_gvr_prior and B * md.gen_token_stride >= 2 * md.num_sms:
+        return "GVR row reorder active"
+    host_seq = md.seq_lens[:B]
+    if int(host_seq.sum()) != md.num_tokens or int(host_seq.max()) > 1 + md.max_draft_tokens:
+        return "query lengths outside the decode layout"
+    int32_attrs = (
+        "kv_lens_cuda",
+        "seq_lens_cuda",
+        "kv_lens_cuda_2d",
+        "indexer_k_cache_block_offsets",
+        "scheduler_metadata_buffer",
+        "scheduler_metadata_buffer_full_next_n",
+        "req_idx_per_token",
+    )
+    for name in int32_attrs:
+        if getattr(md, name).dtype != torch.int32:
+            return f"unexpected dtype of {name}"
+    for name in (
+        "slot_mapping_fp8",
+        "slot_mapping_scale",
+        "gen_kv_indptr",
+        "gen_cached_token_indptr",
+    ):
+        if getattr(md, name).dtype != torch.int64:
+            return f"unexpected dtype of {name}"
+    if md.expand_for_dsl and md.dsl_expand_factor > 1:
+        if md.kv_lens_expanded_cuda.dtype != torch.int32:
+            return "unexpected dtype of kv_lens_expanded_cuda"
+        if B * md.dsl_expand_factor > md.kv_lens_expanded_cuda.shape[0]:
+            return "expanded lengths exceed their buffer"
+    if md.skip_indexer_for_gen_reqs:
+        topk = md.topk_indices_buffer
+        if topk.dtype != torch.int32 or topk.stride(1) != 1 or topk.shape[1] != md.num_sparse_topk:
+            return "unexpected dense top-k buffer layout"
+    if _draft_step.schedule_self_check(md.num_sms, md.kv_lens_cuda.device) is None:
+        return "DeepGEMM schedule not reproduced"
+    return None
+
+
+def _correction_ineligibility(runner, inputs, md, B: int) -> Optional[str]:
+    """Why the accepted-count correction cannot be fused, or None."""
+    pos = inputs.get("position_ids")
+    ids = inputs.get("input_ids")
+    if pos is None or ids is None:
+        return "no position ids or input ids"
+    if pos.dim() != 2 or pos.shape[0] != 1 or pos.dtype != torch.int32 or pos.stride(1) != 1:
+        return "position ids are not a [1, N] int32 row"
+    num_pos = ids.shape[0] - md.num_ctx_tokens
+    if md.num_ctx_tokens != 0 or pos.shape[1] != num_pos:
+        return "position ids do not match the input ids"
+    # Program 0 rewrites the positions while the embedding programs read the ids.
+    if pos.untyped_storage().data_ptr() == ids.untyped_storage().data_ptr():
+        return "position ids share storage with the input ids"
+    pos_off = getattr(runner, "previous_pos_id_offsets_cuda", None)
+    kv_off = getattr(runner, "previous_kv_lens_offsets_cuda", None)
+    if pos_off is None or kv_off is None:
+        return "no overlap-scheduler offsets"
+    if pos_off.dtype != torch.int32 or kv_off.dtype != torch.int32:
+        return "unexpected offset dtype"
+    if not pos_off.is_contiguous() or not kv_off.is_contiguous():
+        return "offsets are not contiguous"
+    if pos_off.shape[0] < num_pos or kv_off.shape[0] < B:
+        return "offsets shorter than the step"
+    if getattr(md, "num_chunked_ctx_requests", 0) != 0:
+        return "chunked context rows"
+    return None
+
+
+def _embedding_target(runner, inputs):
+    """The target ``Embedding`` decode_prologue may run, or (None, reason)."""
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model
+
+    if inputs.get("inputs_embeds") is not None:
+        return None, "inputs_embeds given"
+    target = getattr(getattr(runner, "model", None), "model", None)
+    if not isinstance(target, DeepseekV3Model):
+        return None, "target model is not DeepseekV3Model"
+    embed = getattr(target, "embed_tokens", None)
+    reason = _draft_embedding.module_ineligibility(embed)
+    if reason is not None:
+        return None, reason
+    ids = inputs.get("input_ids")
+    reason = _draft_embedding.ids_ineligibility(ids, embed.weight)
+    if reason is not None:
+        return None, reason
+    return embed, None
+
+
+def _replay_host_on_update_kv_lens(md) -> None:
+    """Host statements of one fused-eligible ``on_update_kv_lens()`` (no device work)."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+    )
+    from tensorrt_llm.logger import logger
+
+    super(DSAtrtllmAttentionMetadata, md).on_update_kv_lens()
+    md._invalidate_pool_view_cache()
+    md.gen_indexer_kv_lens_cuda_runtime = md.get_indexer_kv_lens(
+        md.kv_lens_cuda[md.num_contexts : md.num_seqs]
+    )
+    if not getattr(md, "_fused_dsa_meta_armed", False):
+        logger.info(
+            "Fused DSA decode-metadata "
+            f"kernel armed (num_seqs={md.num_seqs}, "
+            f"num_tokens={md.num_tokens})."
+        )
+        md._fused_dsa_meta_armed = True
+    # In the eligible domain the row reorder is off, so this only resets the
+    # host attribute; the dense top-k rows are written by the launch.
+    md._compute_kv_lens_row_reorder()
+
+
+def prologue_ineligibility(runner, inputs, enable_spec_decode: bool) -> Optional[str]:
+    """Return why ``target_prologue`` must leave the native code in place, or None."""
+    md = inputs.get("attn_metadata")
+    if md is None:
+        return "no attention metadata"
+    mapping = getattr(runner, "mapping", None)
+    if mapping is not None and mapping.has_cp_helix():
+        return "helix context parallelism"
+    B = md.num_seqs
+    reason = _metadata_ineligibility(md, B)
+    if reason is not None:
+        return reason
+    config = getattr(runner, "_config", None)
+    if config is None:
+        return "no runner config"
+    if enable_spec_decode and not config.disable_overlap_scheduler:
+        return _correction_ineligibility(runner, inputs, md, B)
+    return None
+
+
+def target_prologue(runner, inputs, *, enable_spec_decode: bool, runtime_draft_len: int) -> bool:
+    """Run the decode-metadata prologue (and the target embedding) in one launch.
+
+    Replaces the device work and reproduces the host statements of
+    ``DecoderRunner._preprocess_inputs`` up to (not including) the guided
+    decoder event, which the caller keeps.
+
+    Args:
+        runner: the ``DecoderRunner`` (reads ``_config``, ``mapping``,
+            ``model`` and the overlap-scheduler offset buffers).
+        inputs: the step's input dict (``attn_metadata``, ``input_ids``,
+            ``position_ids``).
+        enable_spec_decode: as passed to ``_preprocess_inputs``.
+        runtime_draft_len: as passed to ``_preprocess_inputs`` (unused by the
+            fused path; the DSA branch of the native code does not read it).
+
+    Returns:
+        True when the launch replaced the native prologue; False when the
+        caller must run the native code (nothing has been modified then).
+    """
+    from ... import dispatch
+
+    del runtime_draft_len
+    reason = prologue_ineligibility(runner, inputs, enable_spec_decode)
+    # One log line per distinct outcome: context and warmup steps decline
+    # before the first decode step engages.
+    key = "target_prologue" if reason is None else f"target_prologue[{reason}]"
+    dispatch.log_capability("decode_prologue", key, reason is None, reason or "")
+    if reason is not None:
+        return False
+
+    md = inputs["attn_metadata"]
+    B = md.num_seqs
+    correct = enable_spec_decode and not runner._config.disable_overlap_scheduler
+    embed, embed_reason = _embedding_target(runner, inputs)
+    key = "target_embedding" if embed is not None else f"target_embedding[{embed_reason}]"
+    dispatch.log_capability("decode_prologue", key, embed is not None, embed_reason or "")
+
+    kv_mgr = md.kv_cache_manager
+    index_head_dim = kv_mgr.index_head_dim
+    use_fp4 = getattr(kv_mgr, "use_fp4", False)
+    data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim
+    scale_size = index_head_dim // kv_mgr.quant_block_size * 4
+    tokens_per_block = md._tokens_per_block
+    block_stride = tokens_per_block * (data_bytes_per_token + scale_size)
+    width = min(md.gen_token_stride, md.kv_lens_cuda_2d.shape[1])
+    expand = md.dsl_expand_factor if md.expand_for_dsl else 1
+    num_sms = md.num_sms
+    block_offsets = md.indexer_k_cache_block_offsets
+    dense = bool(md.skip_indexer_for_gen_reqs)
+    num_tokens = md.num_tokens
+    dummy = md.kv_lens_cuda
+
+    if correct:
+        pos = inputs["position_ids"]
+        num_pos = inputs["input_ids"].shape[0] - md.num_ctx_tokens
+        pos_off = runner.previous_pos_id_offsets_cuda
+        kv_off = runner.previous_kv_lens_offsets_cuda
+    else:
+        pos = pos_off = kv_off = dummy
+        num_pos = 1
+
+    if embed is not None:
+        ids = inputs["input_ids"]
+        weight = embed.weight
+        num_rows = ids.shape[0]
+        hidden = weight.shape[1]
+        emb_out = torch.empty((num_rows, hidden), dtype=weight.dtype, device=weight.device)
+        ncb = triton.cdiv(hidden, _BLOCK_H)
+        grid = (1 + num_rows * ncb,)
+    else:
+        ids = weight = emb_out = dummy
+        num_rows = 0
+        hidden = 1
+        ncb = 1
+        grid = (1,)
+
+    topk = md.topk_indices_buffer if dense else dummy
+    block_n = max(2, triton.next_power_of_2(num_tokens))
+    _target_prologue_kernel[grid](
+        md.kv_lens_cuda,
+        kv_off,
+        md.seq_lens_cuda,
+        pos,
+        pos_off,
+        md.kv_lens_cuda_2d,
+        md.kv_lens_cuda_2d.stride(0),
+        md.scheduler_metadata_buffer,
+        md.scheduler_metadata_buffer_full_next_n,
+        md.kv_lens_expanded_cuda if expand > 1 else dummy,
+        md.scheduler_metadata_buffer_expanded if expand > 1 else dummy,
+        block_offsets,
+        block_offsets.stride(0),
+        md.req_idx_per_token,
+        md.slot_mapping_fp8,
+        md.slot_mapping_scale,
+        md.gen_kv_indptr,
+        md.gen_cached_token_indptr,
+        topk,
+        topk.stride(0) if dense else 0,
+        ids,
+        weight,
+        emb_out,
+        weight.shape[0] if embed is not None else 0,
+        weight.stride(0) if embed is not None else 0,
+        emb_out.stride(0) if embed is not None else 0,
+        num_tokens,
+        N_POS=num_pos,
+        B=B,
+        CORRECT=correct,
+        WIDTH=width,
+        EXPAND=expand,
+        NUM_SMS=num_sms,
+        SPLIT_KV=_draft_step._SPLIT_KV_BY_SMS[num_sms],
+        max_blocks=block_offsets.shape[1],
+        block_offsets_stride_1=block_offsets.stride(1),
+        tokens_per_block=tokens_per_block,
+        data_bytes_per_token=data_bytes_per_token,
+        scale_size=scale_size,
+        block_stride=block_stride,
+        scale_base_offset=tokens_per_block * data_bytes_per_token,
+        DENSE=dense,
+        TOPK=md.num_sparse_topk if dense else 1,
+        EMBED=embed is not None,
+        HIDDEN=hidden,
+        NCB=ncb,
+        BLOCK_B=max(2, triton.next_power_of_2(B)),
+        BLOCK_W=max(2, triton.next_power_of_2(width)),
+        BLOCK_E=max(2, triton.next_power_of_2(B * expand)),
+        BLOCK_T=max(2, triton.next_power_of_2(1 + md.max_draft_tokens)),
+        BLOCK_SM=triton.next_power_of_2(num_sms + 1),
+        BLOCK_P=max(2, triton.next_power_of_2(num_pos)),
+        BLOCK_N=block_n,
+        BLOCK_K=max(16, _DENSE_TILE // block_n),
+        BLOCK_H=_BLOCK_H,
+        num_warps=_NUM_WARPS,
+    )
+
+    # Host statements of the native sequence: on_update_kv_lens() (call A),
+    # the correction (device only), on_update_kv_lens() (call B).
+    _replay_host_on_update_kv_lens(md)
+    if correct:
+        _replay_host_on_update_kv_lens(md)
+
+    if embed is not None:
+        setattr(md, _PACKET_ATTR, (inputs["input_ids"], weight, emb_out))
+    return True
+
+
+def take_embedding(attn_metadata, input_ids: torch.Tensor, embed_tokens) -> Optional[torch.Tensor]:
+    """Return the target embedding rows the prologue launch produced, or None.
+
+    Consumes the step's packet. The rows are returned only for exactly the
+    ``input_ids`` tensor and the embedding weight the launch used; otherwise
+    the caller embeds natively (the launch's rows are then unused).
+
+    Args:
+        attn_metadata: the step's attention metadata.
+        input_ids: ``[N]`` input ids the model forward embeds.
+        embed_tokens: the model's ``Embedding`` module.
+
+    Returns:
+        ``[N, hidden]`` embedding rows, or None.
+    """
+    packet = vars(attn_metadata).pop(_PACKET_ATTR, None)
+    if packet is None:
+        return None
+    ids, weight, rows = packet
+    if ids is not input_ids or getattr(embed_tokens, "weight", None) is not weight:
+        from ... import dispatch
+
+        dispatch.log_capability("decode_prologue", "take_embedding", False, "packet does not match")
+        return None
+    return rows
+
+
+class DecodePrologueHook:
+    """The decode-metadata prologue and the target embedding in one launch (``DecodePrologue``).
+
+    The engine calls it first in the step's input preprocessing, inside CUDA graph capture
+    as well (``target_prologue``); the decode target then takes the rows that launch
+    embedded (``take_embedding``).
+    """
+
+    def __call__(self, runner, inputs, *, enable_spec_decode: bool, runtime_draft_len: int) -> bool:
+        return target_prologue(
+            runner,
+            inputs,
+            enable_spec_decode=enable_spec_decode,
+            runtime_draft_len=runtime_draft_len,
+        )
+
+    def take_embedding(
+        self, attn_metadata, input_ids: torch.Tensor, embed_tokens
+    ) -> Optional[torch.Tensor]:
+        return take_embedding(attn_metadata, input_ids, embed_tokens)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/__init__.py
new file mode 100644
index 000000000000..3760e330f54e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/__init__.py
@@ -0,0 +1,30 @@
+# 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.
+"""deferred_input_prep: inter-iteration decode input preparation in one launch.
+
+The target declares one ``PrepStager`` to the engine as its step-input stager and one
+``SamplerStoreWriter`` as its sampler store writer (``runtime.make_runtime_hooks``);
+``compat`` describes the three parts. Entry points (``dispatch.get_impl``):
+
+* ``prep_stager``: the ``PrepStager`` class.
+"""
+
+from ...dispatch import register
+from .sampler_store import SamplerStoreWriter
+from .stager import PrepStager
+
+register("deferred_input_prep", "prep_stager", PrepStager)
+
+__all__ = ["PrepStager", "SamplerStoreWriter"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py
new file mode 100644
index 000000000000..d0ae536316f7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py
@@ -0,0 +1,84 @@
+# 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.
+"""deferred_input_prep's executor-loop domain of its store writer, and its debug level.
+
+deferred_input_prep has three parts, each a fast path of the plan
+(``dispatch.PART_PATHS``) that ``runtime.make_runtime_hooks`` declares to the engine:
+the staged input preparation (``prep``, the step-input stager), the sampler store
+writer (``store``) and the overlap executor loop on the execution stream (``stream``).
+``TRTLLM_STAIRCASE_GLM_LL_DISABLE`` turns each of them off.
+
+deferred_input_prep's sampler store writer is validated only in the overlap executor
+loop of an executor without attention data parallelism. The target's executor events
+declare the overlap loop's attention-DP setting when it starts
+(``declare_overlap_loop``); until a loop of this process has declared that domain, and
+once any loop has declared another one, the writer keeps the native path
+(``loop_domain``). The input preparation stager checks attention DP and the overlap
+scheduler itself, per preparation.
+
+``TRTLLM_MODELING_V2_VALIDATE=1`` (``dispatch.validate_mode``) builds the stager
+with ``debug_level=1``, the guard: between the start of a staged preparation and
+its launch, a torch dispatch mode and a Triton launch hook raise on any operation
+whose tensor arguments overlap an already staged destination, and every staged
+destination is compared at the launch with its bytes at staging time (raises if
+anything wrote it in between). Destination bytes are unchanged, so oracle digests
+remain comparable with the native run. ``debug_level=2`` (tests only) adds poison:
+every destination that the launch overwrites is filled with ``0xA5`` bytes when it
+is staged, so a reader of any kind (also kernels the hooks cannot see) consumes
+garbage and the oracle fails.
+"""
+
+from typing import Optional
+
+
+# Executor-loop domain of this process: declared by the overlap loop at its start.
+_LOOP_STATE = {"declared": False, "reason": "", "logged": False}
+_LOOP_UNDECLARED = (
+    "no overlap executor loop declared its domain; the store writer is validated only there"
+)
+_LOOP_NO_ARGUMENT = "the overlap executor loop did not pass its attention-DP setting"
+_LOOP_ATTENTION_DP = "attention DP: the store writer is not validated with it"
+
+
+def declare_overlap_loop(enable_attention_dp: Optional[bool]) -> None:
+    """Record the attention-DP setting of an overlap executor loop that starts.
+
+    ``None``: the caller did not pass it. A declaration outside the validated
+    domain is final for the process.
+    """
+    if enable_attention_dp is None:
+        reason = _LOOP_NO_ARGUMENT
+    elif enable_attention_dp:
+        reason = _LOOP_ATTENTION_DP
+    else:
+        reason = ""
+    _LOOP_STATE["declared"] = True
+    if reason and not _LOOP_STATE["reason"]:
+        _LOOP_STATE["reason"] = reason
+
+
+def loop_domain() -> bool:
+    """Whether the store writer may engage.
+
+    Logs the first decision of the process (``deferred_input_prep/loop_domain``).
+    """
+    state = _LOOP_STATE
+    reason = state["reason"] if state["declared"] else _LOOP_UNDECLARED
+    if not state["logged"]:
+        from ...dispatch import log_capability
+
+        state["logged"] = True
+        log_capability("deferred_input_prep", "loop_domain", not reason, reason)
+    return not reason
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/guard.py
new file mode 100644
index 000000000000..2ff600c37299
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/guard.py
@@ -0,0 +1,233 @@
+# 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.
+"""Debug guard for the window between staging a write and the deferred_input_prep launch.
+
+A staged write reaches its destination only at the deferred_input_prep launch. Between the
+staging call and the launch, another device operation that reads the destination
+would see older bytes, and one that writes it would be overwritten by the
+launch. ``DebugWindow`` (``TRTLLM_MODELING_V2_VALIDATE=1``,
+see ``compat``) makes both cases fail loudly:
+
+* a torch dispatch mode and a hook on Triton's ``JITFunction.run`` raise
+  ``InputPrepHazardError`` when an operation's CUDA tensor arguments overlap a
+  destination that was already staged;
+* every staged destination is snapshotted when staged (level 1) or filled
+  with ``POISON_BYTE`` when the launch overwrites it entirely (level 2) and
+  compared right before the launch; a difference means some operation (of
+  any launcher type) wrote it inside the window.
+
+Operations issued by the stager itself run inside ``internal()`` and are not
+checked.
+"""
+
+import contextlib
+import traceback
+from dataclasses import dataclass
+from typing import Iterable, List, Optional
+
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+from torch.utils._pytree import tree_flatten
+
+POISON_BYTE = 0xA5
+
+# Operators that read only tensor metadata (shape, dtype, device).
+_METADATA_ONLY_OPS = frozenset(
+    {
+        "aten::alias",
+        "aten::detach",
+        "aten::empty_like",
+        "aten::full_like",
+        "aten::lift_fresh",
+        "aten::new_empty",
+        "aten::new_empty_strided",
+        "aten::new_full",
+        "aten::new_ones",
+        "aten::new_zeros",
+        "aten::ones_like",
+        "aten::record_stream",
+        "aten::zeros_like",
+    }
+)
+
+
+class InputPrepHazardError(RuntimeError):
+    """An operation touched a deferred_input_prep-staged destination before the deferred_input_prep launch."""
+
+
+def tensor_span(t: torch.Tensor) -> Optional[tuple]:
+    """Byte range ``[lo, hi)`` covered by a CUDA tensor view, or None."""
+    if not t.is_cuda or t.numel() == 0:
+        return None
+    last = 0
+    for size, stride in zip(t.shape, t.stride()):
+        if size > 1:
+            last += (size - 1) * stride
+    lo = t.data_ptr()
+    return lo, lo + (last + 1) * t.element_size()
+
+
+def _reads_only_metadata(func) -> bool:
+    schema = func._schema
+    if schema.name in _METADATA_ONLY_OPS:
+        return True
+    # View operators return an alias of an input without touching its data.
+    return any(r.alias_info is not None and not r.alias_info.is_write for r in schema.returns)
+
+
+def _call_site() -> str:
+    for frame in reversed(traceback.extract_stack(limit=12)[:-1]):
+        if "kernels/deferred_input_prep/" not in frame.filename.replace("\\", "/"):
+            return f"{frame.filename.rsplit('/', 1)[-1]}:{frame.lineno} {frame.name}"
+    return "deferred_input_prep"
+
+
+@dataclass
+class _Staged:
+    lo: int
+    hi: int
+    label: str
+    view: torch.Tensor
+    snapshot: Optional[torch.Tensor]  # None: poisoned
+
+
+class _GuardMode(TorchDispatchMode):
+    def __init__(self, window: "DebugWindow"):
+        super().__init__()
+        self.window = window
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        window = self.window
+        if window.internal_depth == 0 and window.staged and not _reads_only_metadata(func):
+            window.check(str(func), tree_flatten((args, kwargs))[0])
+        return func(*args, **kwargs)
+
+
+class DebugWindow:
+    """Guard state of one staged preparation (see the module docstring)."""
+
+    def __init__(self, level: int):
+        self.level = level
+        self.staged: List[_Staged] = []
+        self.internal_depth = 0
+        self.num_checked_ops = 0
+        self.num_windows = 0
+        self._mode: Optional[_GuardMode] = None
+        self._triton_jit = None
+        self._triton_run = None
+
+    # ---- window lifecycle ----------------------------------------------
+
+    def open(self) -> None:
+        self.close()
+        self.staged = []
+        self.num_windows += 1
+        self._mode = _GuardMode(self)
+        self._mode.__enter__()
+        self._hook_triton()
+
+    def close(self) -> None:
+        if self._mode is not None:
+            mode, self._mode = self._mode, None
+            mode.__exit__(None, None, None)
+        self._unhook_triton()
+
+    @contextlib.contextmanager
+    def internal(self):
+        """Operations of the stager itself (not checked)."""
+        self.internal_depth += 1
+        try:
+            yield
+        finally:
+            self.internal_depth -= 1
+
+    # ---- staging and checks --------------------------------------------
+
+    def register(self, dst: torch.Tensor, overwrite: bool) -> None:
+        """Record a staged destination; call inside ``internal()``."""
+        span = tensor_span(dst)
+        if span is None:
+            return
+        if not dst.is_contiguous():
+            raise InputPrepHazardError(
+                f"deferred_input_prep debug: non-contiguous staged destination at {_call_site()}"
+            )
+        if overwrite and self.level >= 2:
+            dst.view(torch.uint8).fill_(POISON_BYTE)
+            snapshot = None
+        else:
+            snapshot = dst.detach().clone()
+        self.staged.append(_Staged(span[0], span[1], _call_site(), dst, snapshot))
+
+    def check(self, what: str, values: Iterable) -> None:
+        self.num_checked_ops += 1
+        for v in values:
+            if not isinstance(v, torch.Tensor):
+                continue
+            span = tensor_span(v)
+            if span is None:
+                continue
+            for s in self.staged:
+                if span[0] < s.hi and s.lo < span[1]:
+                    raise InputPrepHazardError(
+                        f"deferred_input_prep debug guard: {what} touches a destination staged at {s.label} "
+                        "before the deferred_input_prep launch (it would read stale bytes or be overwritten)"
+                    )
+
+    def verify(self) -> None:
+        """Raise if any staged destination changed between staging and now."""
+        if not self.staged:
+            return
+        with self.internal():
+            flags = []
+            for s in self.staged:
+                cur = s.view.view(torch.uint8)
+                if s.snapshot is None:
+                    flags.append((cur != POISON_BYTE).any())
+                else:
+                    flags.append((cur != s.snapshot.view(torch.uint8)).any())
+            changed = torch.stack(flags).cpu().tolist()
+        hits = [s.label for s, c in zip(self.staged, changed) if c]
+        self.staged = []
+        if hits:
+            raise InputPrepHazardError(
+                "deferred_input_prep debug guard: destinations staged at "
+                f"{sorted(set(hits))} were written by another operation before the deferred_input_prep launch"
+            )
+
+    # ---- Triton ---------------------------------------------------------
+
+    def _hook_triton(self) -> None:
+        try:
+            from triton.runtime.jit import JITFunction
+        except ImportError:
+            return
+        window = self
+        orig = JITFunction.run
+
+        def run(fn_self, *args, **kwargs):
+            if window.internal_depth == 0 and window.staged:
+                name = getattr(fn_self, "__name__", "kernel")
+                window.check(f"triton:{name}", list(args) + list(kwargs.values()))
+            return orig(fn_self, *args, **kwargs)
+
+        self._triton_jit, self._triton_run = JITFunction, orig
+        JITFunction.run = run
+
+    def _unhook_triton(self) -> None:
+        if self._triton_jit is not None:
+            self._triton_jit.run = self._triton_run
+            self._triton_jit = self._triton_run = None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py
new file mode 100644
index 000000000000..cd322a249581
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py
@@ -0,0 +1,398 @@
+# 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.
+"""One-launch application of many small device-buffer updates.
+
+A ``PrepPacket`` collects the device writes of one input preparation as
+descriptors (plus a byte payload for host values that are not already in
+pinned memory) in pinned host memory. ``commit`` launches a single Triton
+kernel, one CTA per descriptor chunk, that reads the packet and the pinned
+host sources directly (zero-copy, no H2D memcpy op) and performs every write.
+Descriptors of one packet must write disjoint bytes and must not read bytes
+another descriptor of the same packet writes; the builder enforces the first
+rule (``PacketOverlap``), the callers guarantee the second.
+
+Descriptor kinds (32-bit words; a descriptor covers at most ``CHUNK`` words,
+longer writes are split by the builder):
+
+* ``COPY``: ``dst[i] = payload[i]``.
+* ``GATHER``: ``dst[i] = src[idx[i]] + addend`` for ``i < n_idx`` and
+  ``dst[i] = 0`` for ``n_idx <= i < n``; ``idx`` lives in the payload and
+  ``src`` is a device buffer the packet does not write.
+* ``BLOCK_OFFSETS``: the KV-cache-manager page-table expansion
+  ``dst_k[i] = 0 if v == bad else scale * v`` and
+  ``dst_v[i] = 0 if v == bad else scale * v + kv_offset`` with ``v`` read
+  from a pinned host page-table row.
+* ``CLAMP_MIN``: ``dst[i] = max(dst[i], lo)`` in place.
+* ``COPY_CLAMP_MIN``: ``dst[i] = max(payload[i], lo)``.
+* ``COPY_PTR`` / ``COPY_PTR_CLAMP_MIN``: as ``COPY`` / ``COPY_CLAMP_MIN`` but
+  reading a pinned host source in place, at kernel execution time, exactly
+  like the native asynchronous H2D copy from that source. The packet keeps
+  the source tensor alive until its launch has completed.
+"""
+
+from typing import List, Optional, Sequence
+
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+
+KIND_COPY = 0
+KIND_GATHER = 1
+KIND_BLOCK_OFFSETS = 2
+KIND_CLAMP_MIN = 3
+KIND_COPY_CLAMP_MIN = 4
+KIND_COPY_PTR = 5
+KIND_COPY_PTR_CLAMP_MIN = 6
+
+# Words per descriptor: kind, dst, src, n, aux0..aux3.
+DESC_WORDS = 8
+CHUNK = 4096
+
+
+class PacketFull(ValueError):
+    """The packet has no room for this write; nothing of the write was added."""
+
+
+class PacketOverlap(ValueError):
+    """The write overlaps a destination of this packet; nothing was added.
+
+    Descriptors of one launch run concurrently, so a later write to the same
+    bytes must go into a later launch to keep last-write-wins order.
+    """
+
+
+@triton.jit
+def _prep_commit_kernel(desc_addr, payload_addr, CHUNK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    desc = desc_addr.to(tl.pointer_type(tl.int64)) + pid * 8
+    kind = tl.load(desc)
+    dst = tl.load(desc + 1).to(tl.pointer_type(tl.int32))
+    src_word = tl.load(desc + 2)
+    n = tl.load(desc + 3).to(tl.int32)
+    aux0 = tl.load(desc + 4)
+    aux1 = tl.load(desc + 5)
+    aux2 = tl.load(desc + 6)
+    aux3 = tl.load(desc + 7)
+    payload = payload_addr.to(tl.pointer_type(tl.int32))
+    offs = tl.arange(0, CHUNK_SIZE)
+    mask = offs < n
+    if kind == 0:
+        vals = tl.load(payload + src_word + offs, mask=mask)
+        tl.store(dst + offs, vals, mask=mask)
+    elif kind == 1:
+        src = src_word.to(tl.pointer_type(tl.int32))
+        n_idx = aux2.to(tl.int32)
+        gmask = offs < n_idx
+        idx = tl.load(payload + aux0 + offs, mask=gmask, other=0)
+        vals = tl.load(src + idx, mask=gmask, other=0) + aux1.to(tl.int32)
+        vals = tl.where(gmask, vals, 0)
+        tl.store(dst + offs, vals, mask=mask)
+    elif kind == 2:
+        src = src_word.to(tl.pointer_type(tl.int32))
+        dst_v = aux0.to(tl.pointer_type(tl.int32))
+        scale = aux1.to(tl.int32)
+        kv_off = aux2.to(tl.int32)
+        bad = aux3.to(tl.int32)
+        v = tl.load(src + offs, mask=mask, other=0)
+        is_bad = v == bad
+        k_val = tl.where(is_bad, 0, scale * v)
+        v_val = tl.where(is_bad, 0, scale * v + kv_off)
+        tl.store(dst + offs, k_val, mask=mask)
+        tl.store(dst_v + offs, v_val, mask=mask)
+    elif kind == 3:
+        lo = aux0.to(tl.int32)
+        vals = tl.load(dst + offs, mask=mask, other=0)
+        tl.store(dst + offs, tl.maximum(vals, lo), mask=mask)
+    elif kind == 4:
+        lo = aux0.to(tl.int32)
+        vals = tl.load(payload + src_word + offs, mask=mask, other=0)
+        tl.store(dst + offs, tl.maximum(vals, lo), mask=mask)
+    elif kind == 5:
+        src = src_word.to(tl.pointer_type(tl.int32))
+        vals = tl.load(src + offs, mask=mask)
+        tl.store(dst + offs, vals, mask=mask)
+    else:
+        src = src_word.to(tl.pointer_type(tl.int32))
+        lo = aux0.to(tl.int32)
+        vals = tl.load(src + offs, mask=mask, other=0)
+        tl.store(dst + offs, tl.maximum(vals, lo), mask=mask)
+
+
+class _NeverDone:
+    @staticmethod
+    def query() -> bool:
+        return False
+
+
+_NEVER_DONE = _NeverDone()
+
+
+class _PinnedSlot:
+    """One packet's pinned storage and the event that marks its consumption."""
+
+    def __init__(self, max_desc: int, payload_words: int):
+        self.desc = torch.zeros((max_desc, DESC_WORDS), dtype=torch.int64, pin_memory=True)
+        self.payload = torch.zeros((payload_words,), dtype=torch.int32, pin_memory=True)
+        self.desc_np = self.desc.numpy()
+        self.payload_np = self.payload.numpy()
+        self.event: Optional[torch.cuda.Event] = None
+        # Pinned sources read in place by this slot's launch.
+        self.keepalive: List[torch.Tensor] = []
+
+    def free(self) -> bool:
+        return self.event is None or self.event.query()
+
+
+class PrepPacket:
+    """Builder for one commit; reuses a ring of pinned slots without host syncs.
+
+    A slot is reused only after the CUDA event recorded behind its commit has
+    completed (polled with ``query``); otherwise a new slot is allocated, so
+    building a packet never blocks the host.
+    """
+
+    def __init__(self, max_desc: int = 4096, payload_words: int = 1 << 18):
+        self.max_desc = max_desc
+        self.payload_words = payload_words
+        self._slots: List[_PinnedSlot] = []
+        self._cur: Optional[_PinnedSlot] = None
+        self._descs: List[tuple] = []
+        self._n_payload = 0
+        self._ranges: List[tuple] = []
+
+    # ---- building -------------------------------------------------------
+
+    def begin(self) -> None:
+        slot = next((s for s in self._slots if s.free()), None)
+        if slot is None:
+            slot = _PinnedSlot(self.max_desc, self.payload_words)
+            self._slots.append(slot)
+        slot.keepalive = []
+        self._cur = slot
+        self._descs = []
+        self._n_payload = 0
+        self._ranges = []
+
+    def abort(self) -> None:
+        if self._cur is not None:
+            self._cur.keepalive = []
+        self._cur = None
+        self._descs = []
+
+    @property
+    def num_descriptors(self) -> int:
+        return len(self._descs)
+
+    def _reserve(self, n_desc: int, n_words: int) -> None:
+        if len(self._descs) + n_desc > self.max_desc:
+            raise PacketFull("packet descriptor table is full")
+        if self._n_payload + n_words > self.payload_words:
+            raise PacketFull("packet payload is full")
+
+    @staticmethod
+    def _dst_ptr(dst: torch.Tensor, nbytes: int) -> int:
+        if not dst.is_cuda:
+            raise ValueError("packet destinations must be CUDA tensors")
+        if not dst.is_contiguous():
+            raise ValueError("packet destinations must be contiguous")
+        ptr = dst.data_ptr()
+        if ptr % 4 or nbytes % 4:
+            raise ValueError("packet destinations must be 4-byte aligned words")
+        return ptr
+
+    def _check_dsts(self, *dsts) -> List[int]:
+        """Validate ``(tensor, nbytes)`` destinations and claim their bytes.
+
+        All-or-nothing: on ``PacketOverlap`` no range is claimed.
+        """
+        ptrs = [self._dst_ptr(dst, nbytes) for dst, nbytes in dsts]
+        new = [(p, p + nbytes) for p, (_, nbytes) in zip(ptrs, dsts)]
+        for lo, hi in new:
+            for a, b in self._ranges:
+                if lo < b and a < hi:
+                    raise PacketOverlap("overlapping destinations in one packet")
+        self._ranges.extend(new)
+        return ptrs
+
+    def _check_dst(self, dst: torch.Tensor, nbytes: int) -> int:
+        return self._check_dsts((dst, nbytes))[0]
+
+    def _alloc_payload(self, nwords: int) -> int:
+        off = self._n_payload
+        self._n_payload = off + nwords
+        return off
+
+    def _emit_split(self, kind, dst_ptr, src, n, aux, src_step, aux0_step=0):
+        """Append ``kind`` descriptors covering ``n`` words in ``CHUNK`` pieces.
+
+        ``src_step`` / ``aux0_step``: per-word advance of ``src`` / ``aux0``.
+        """
+        a0, a1, a2, a3 = aux
+        self._descs.extend(
+            (
+                kind,
+                dst_ptr + 4 * s,
+                src + src_step * s,
+                min(CHUNK, n - s),
+                a0 + aux0_step * s,
+                a1,
+                a2,
+                a3,
+            )
+            for s in range(0, n, CHUNK)
+        )
+
+    @staticmethod
+    def _num_chunks(n: int) -> int:
+        return (n + CHUNK - 1) // CHUNK
+
+    def add_copy(self, dst: torch.Tensor, src: torch.Tensor, clamp_min: Optional[int] = None):
+        """Stage ``dst.copy_(src)`` for a CPU ``src`` with ``dst``'s dtype and shape.
+
+        A contiguous pinned ``src`` is read in place by the kernel; any other
+        ``src`` is copied into the payload now.
+        """
+        if src.dtype != dst.dtype or src.numel() != dst.numel():
+            raise ValueError("add_copy needs matching dtype and element count")
+        nbytes = dst.numel() * dst.element_size()
+        nwords = nbytes // 4
+        in_place = src.is_contiguous() and src.is_pinned() and src.data_ptr() % 4 == 0
+        self._reserve(self._num_chunks(nwords), 0 if in_place else nwords)
+        ptr = self._check_dst(dst, nbytes)
+        lo = 0 if clamp_min is None else clamp_min
+        if in_place:
+            self._cur.keepalive.append(src)
+            kind = KIND_COPY_PTR if clamp_min is None else KIND_COPY_PTR_CLAMP_MIN
+            self._emit_split(kind, ptr, src.data_ptr(), nwords, (lo, 0, 0, 0), 4)
+            return
+        off = self._alloc_payload(nwords)
+        src_words = src.contiguous().reshape(-1).view(torch.int32).numpy()
+        self._cur.payload_np[off : off + nwords] = src_words
+        kind = KIND_COPY if clamp_min is None else KIND_COPY_CLAMP_MIN
+        self._emit_split(kind, ptr, off, nwords, (lo, 0, 0, 0), 1)
+
+    def add_values(self, dst: torch.Tensor, values: Sequence[int]):
+        """Stage ``dst.copy_(torch.tensor(values, dtype=dst.dtype))`` for int32 ``dst``."""
+        if dst.dtype != torch.int32 or dst.numel() != len(values):
+            raise ValueError("add_values needs an int32 destination of len(values)")
+        n = len(values)
+        self._reserve(self._num_chunks(n), n)
+        ptr = self._check_dst(dst, 4 * n)
+        off = self._alloc_payload(n)
+        self._cur.payload_np[off : off + n] = values
+        self._emit_split(KIND_COPY, ptr, off, n, (0, 0, 0, 0), 1)
+
+    def add_gather(
+        self, dst: torch.Tensor, src: torch.Tensor, indices: Sequence[int], addend: int = 0
+    ):
+        """Stage ``dst[:k] = src.flatten()[indices] + addend`` and ``dst[k:] = 0``.
+
+        ``src`` is an int32 device buffer that this packet does not write;
+        ``indices`` are flat element offsets into ``src``'s storage.
+        """
+        if dst.dtype != torch.int32 or src.dtype != torch.int32:
+            raise ValueError("add_gather supports int32 only")
+        n = dst.numel()
+        if len(indices) > n or n > CHUNK:
+            raise ValueError("add_gather: too many indices or destination too long")
+        k = max(1, len(indices))
+        self._reserve(1, k)
+        ptr = self._check_dst(dst, 4 * n)
+        off = self._alloc_payload(k)
+        if indices:
+            self._cur.payload_np[off : off + len(indices)] = indices
+        self._descs.append((KIND_GATHER, ptr, src.data_ptr(), n, off, addend, len(indices), 0))
+
+    def add_zero(self, dst: torch.Tensor):
+        """Stage ``dst.zero_()`` for a 4-byte-word buffer."""
+        nbytes = dst.numel() * dst.element_size()
+        self._reserve(self._num_chunks(nbytes // 4), 0)
+        ptr = self._check_dst(dst, nbytes)
+        # A GATHER without indices writes zeros.
+        self._emit_split(KIND_GATHER, ptr, 0, nbytes // 4, (0, 0, 0, 0), 0)
+
+    def add_block_offsets(
+        self,
+        dst_k: torch.Tensor,
+        dst_v: torch.Tensor,
+        src_host_ptr: int,
+        n: int,
+        scale: int,
+        kv_offset: int,
+        bad: int,
+    ):
+        """Stage one (pool, sequence) page-table row expansion."""
+        self._reserve(self._num_chunks(n), 0)
+        ptr_k, ptr_v = self._check_dsts((dst_k, 4 * n), (dst_v, 4 * n))
+        self._emit_split(
+            KIND_BLOCK_OFFSETS,
+            ptr_k,
+            src_host_ptr,
+            n,
+            (ptr_v, scale, kv_offset, bad),
+            4,
+            aux0_step=4,
+        )
+
+    def add_clamp_min(self, dst: torch.Tensor, lo: int):
+        """Stage ``dst.clamp_(min=lo)`` for an int32 buffer."""
+        if dst.dtype != torch.int32:
+            raise ValueError("add_clamp_min supports int32 only")
+        n = dst.numel()
+        self._reserve(self._num_chunks(n), 0)
+        ptr = self._check_dst(dst, 4 * n)
+        self._emit_split(KIND_CLAMP_MIN, ptr, 0, n, (lo, 0, 0, 0), 0)
+
+    # ---- launch ---------------------------------------------------------
+
+    def commit(self, stream: Optional[torch.cuda.Stream] = None, record_event: bool = True) -> int:
+        """Launch the packet on ``stream`` (default: current).
+
+        The slot is reused once an event recorded after the launch completes.
+        With ``record_event=False`` the caller records one of its own events
+        right after the launch and passes it to ``set_completion_event`` (an
+        event record costs GPU time between tiny kernels).
+        """
+        slot = self._cur
+        descs = self._descs
+        self._cur = None
+        self._descs = []
+        n = len(descs)
+        if n == 0:
+            if slot is not None:
+                slot.keepalive = []
+            return 0
+        slot.desc_np[:n] = np.array(descs, dtype=np.int64)
+        stream = stream or torch.cuda.current_stream()
+        with torch.cuda.stream(stream):
+            _prep_commit_kernel[(n,)](
+                slot.desc.data_ptr(), slot.payload.data_ptr(), CHUNK_SIZE=CHUNK, num_warps=4
+            )
+            if record_event:
+                slot.event = torch.cuda.Event()
+                slot.event.record(stream)
+            else:
+                # Not reusable until the caller provides the completion event.
+                slot.event = _NEVER_DONE
+                self._last_slot = slot
+        return n
+
+    def set_completion_event(self, event: torch.cuda.Event) -> None:
+        """Mark the last launch consumed once ``event`` (recorded after it) completes."""
+        slot = getattr(self, "_last_slot", None)
+        if slot is not None:
+            slot.event = event
+            self._last_slot = None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/paged_mqa_schedule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/paged_mqa_schedule.py
new file mode 100644
index 000000000000..bfb42ed7c1a4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/paged_mqa_schedule.py
@@ -0,0 +1,70 @@
+# 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.
+"""Host replica of DeepGEMM's SM100 paged-MQA-logits schedule metadata.
+
+``deep_gemm.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms)``
+launches ``sm100_paged_mqa_logits_metadata`` (DeepGEMM scheduler header
+``sm100_paged_mqa_logits.cuh``, non-varlen branch), whose output is a pure
+function of the per-request context lengths:
+
+* ``work[r] = ceil(len[r] / SPLIT_KV) * next_n`` and its inclusive prefix sum;
+* for every ``sm`` in ``[0, num_sms]``: ``w = sm * q + min(sm, rem)`` with
+  ``q, rem = divmod(total_work, num_sms)``; the first request ``r`` with
+  ``prefix[r] > w`` gives ``(r * next_n, (w - prefix[r - 1]) // next_n)``,
+  and ``(num_requests * next_n, 0)`` past the end.
+
+``len[r]`` is ``context_lens[r, -1]`` for a 2-D ``[requests, next_n]``
+input. deferred_input_prep stages the result as a packet copy because the lengths are known on
+the host when the decode inputs are prepared; ``test_input_prep_kernels.py`` checks it
+against the native kernel.
+"""
+
+import bisect
+from typing import Sequence
+
+import numpy as np
+
+# DeepGEMM's KV split width for TRT-LLM's schedule block size (64 KV tokens
+# per page, 4 pages per split); characterized against the native kernel by
+# ``test_paged_mqa_schedule_host_replica_matches_deepgemm``.
+SPLIT_KV_FOR_BLOCK_KV = {64: 256}
+
+
+def paged_mqa_logits_schedule(
+    context_lens: Sequence[int], next_n: int, split_kv: int, num_sms: int
+) -> np.ndarray:
+    """Return the ``[num_sms + 1, 2]`` int32 schedule for ``len(context_lens)`` requests."""
+    n = len(context_lens)
+    out = np.zeros((num_sms + 1, 2), dtype=np.int32)
+    total_tokens = n * next_n
+    if total_tokens == 0:
+        return out
+    prefix = []
+    acc = 0
+    for length in context_lens:
+        acc += -(-int(length) // split_kv) * next_n
+        prefix.append(acc)
+    q, rem = divmod(acc, num_sms)
+    for sm in range(num_sms + 1):
+        w = sm * q + min(sm, rem)
+        r = bisect.bisect_right(prefix, w)
+        if r < n:
+            before = prefix[r - 1] if r > 0 else 0
+            out[sm, 0] = r * next_n
+            out[sm, 1] = (w - before) // next_n
+        else:
+            out[sm, 0] = total_tokens
+            out[sm, 1] = 0
+    return out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/sampler_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/sampler_store.py
new file mode 100644
index 000000000000..d6a953d44284
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/sampler_store.py
@@ -0,0 +1,248 @@
+# 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.
+"""One-launch store update and host copy for the one-model speculative sampler.
+
+Native ``SpecSampler.sample_async`` issues, after each decode-graph replay, an
+H2D copy of the slot list, four ``index_copy_`` kernels that scatter the
+worker outputs into the slot-indexed store, and three D2H copies of the store
+into fresh pinned host tensors. ``store_and_copy_out`` does the same in one
+single-CTA launch: it reads the slot list from pinned memory, scatters the
+rows, and then writes the three complete store tensors straight into the
+pinned host tensors (zero-copy), so the host tensors hold exactly what the
+native D2H copies would have produced.
+
+The kernel reads the pinned slot list and writes the pinned host tensors at
+execution time, possibly long after the call returned. Both are therefore
+owned by a ring entry until the completion event recorded behind the launch
+has completed: the slot list is reused and the host tensors are released only
+then (a native D2H copy gets the same protection from the caching host
+allocator's stream events).
+
+The writer is used only in the executor-loop domain deferred_input_prep is validated in
+(``compat.loop_domain``: the overlap loop of an executor without attention
+DP); ``supported`` checks it before every launch, so any other loop keeps the
+native path.
+"""
+
+from typing import Sequence
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._utils import prefer_pinned
+
+from .compat import loop_domain
+from .packet import _NEVER_DONE
+
+
+@triton.jit
+def _sampler_store_kernel(
+    slots_addr,
+    n_req,
+    o_new_ptr,
+    o_new_s0,
+    o_new_s1,
+    w_new,
+    o_next_new_ptr,
+    o_next_new_s0,
+    o_next_new_s1,
+    w_next_new,
+    o_lens_ptr,
+    o_lens_s0,
+    o_draft_ptr,
+    o_draft_s0,
+    o_draft_s1,
+    w_draft,
+    st_new_ptr,
+    st_new_s0,
+    st_new_s1,
+    st_next_new_ptr,
+    st_next_new_s0,
+    st_next_new_s1,
+    st_lens_ptr,
+    st_draft_ptr,
+    st_draft_s0,
+    n_st_new,
+    n_st_lens,
+    n_st_draft,
+    h_new_addr,
+    h_lens_addr,
+    h_draft_addr,
+    BLOCK: tl.constexpr,
+):
+    slots = slots_addr.to(tl.pointer_type(tl.int64))
+    offs = tl.arange(0, BLOCK)
+    for j in range(n_req):
+        slot = tl.load(slots + j)
+        m = offs < w_new
+        v = tl.load(o_new_ptr + j * o_new_s0 + offs * o_new_s1, mask=m)
+        tl.store(st_new_ptr + offs * st_new_s0 + slot * st_new_s1, v, mask=m)
+        m = offs < w_next_new
+        v = tl.load(o_next_new_ptr + j * o_next_new_s0 + offs * o_next_new_s1, mask=m)
+        tl.store(st_next_new_ptr + offs * st_next_new_s0 + slot * st_next_new_s1, v, mask=m)
+        v = tl.load(o_lens_ptr + j * o_lens_s0)
+        tl.store(st_lens_ptr + slot, v)
+        m = offs < w_draft
+        v = tl.load(o_draft_ptr + j * o_draft_s0 + offs * o_draft_s1, mask=m)
+        tl.store(st_draft_ptr + slot * st_draft_s0 + offs, v, mask=m)
+    tl.debug_barrier()
+    h_new = h_new_addr.to(tl.pointer_type(tl.int32))
+    h_lens = h_lens_addr.to(tl.pointer_type(tl.int32))
+    h_draft = h_draft_addr.to(tl.pointer_type(tl.int32))
+    for start in range(0, n_st_new, BLOCK):
+        m = start + offs < n_st_new
+        tl.store(h_new + start + offs, tl.load(st_new_ptr + start + offs, mask=m), mask=m)
+    for start in range(0, n_st_lens, BLOCK):
+        m = start + offs < n_st_lens
+        tl.store(h_lens + start + offs, tl.load(st_lens_ptr + start + offs, mask=m), mask=m)
+    for start in range(0, n_st_draft, BLOCK):
+        m = start + offs < n_st_draft
+        tl.store(h_draft + start + offs, tl.load(st_draft_ptr + start + offs, mask=m), mask=m)
+
+
+class SamplerStoreWriter:
+    """Holds the pinned slot-list ring for ``store_and_copy_out``."""
+
+    def __init__(self, ring: int = 4, max_requests: int = 1024):
+        # Pageable when prefer_pinned() is False; the sampler then never engages the writer.
+        self._slots = [
+            torch.zeros(max_requests, dtype=torch.int64, pin_memory=prefer_pinned())
+            for _ in range(ring)
+        ]
+        self._events = [None] * ring
+        # Zero-copy host targets of each entry's launch, held until its event completes.
+        self._keep = [()] * ring
+        self._next = 0
+        self.num_calls = 0
+
+    def _slot_buffer(self) -> torch.Tensor:
+        for _ in range(len(self._slots)):
+            i = self._next
+            self._next = (self._next + 1) % len(self._slots)
+            ev = self._events[i]
+            if ev is None or ev.query():
+                self._keep[i] = ()
+                return i
+        self._slots.append(torch.zeros_like(self._slots[0], pin_memory=prefer_pinned()))
+        self._events.append(None)
+        self._keep.append(())
+        return len(self._slots) - 1
+
+    def supported(self, store, o_new, o_next_new, o_lens, o_draft) -> bool:
+        """Whether this sampling step may use the writer (checked before any side effect)."""
+        return loop_domain() and self.tensors_supported(store, o_new, o_next_new, o_lens, o_draft)
+
+    @staticmethod
+    def tensors_supported(store, o_new, o_next_new, o_lens, o_draft) -> bool:
+        """Whether the store and output tensors have the layout the kernel handles."""
+        tensors = (
+            store.new_tokens,
+            store.next_new_tokens,
+            store.new_tokens_lens,
+            store.next_draft_tokens,
+            o_new,
+            o_next_new,
+            o_lens,
+            o_draft,
+        )
+        if any(t.dtype != torch.int32 or not t.is_cuda for t in tensors):
+            return False
+        if not all(t.is_contiguous() for t in tensors[:4]):
+            return False
+        if store.new_tokens.shape[-1] != 1 or store.next_new_tokens.shape[-1] != 1:
+            return False
+        return (
+            o_new.shape[1] == store.new_tokens.shape[0]
+            and o_next_new.shape[1] == store.next_new_tokens.shape[0]
+            and o_draft.shape[1] == store.next_draft_tokens.shape[1]
+            and max(o_new.shape[1], o_next_new.shape[1], o_draft.shape[1]) <= 64
+        )
+
+    def store_and_copy_out(
+        self,
+        store,
+        slots: Sequence[int],
+        o_new,
+        o_next_new,
+        o_lens,
+        o_draft,
+        h_new,
+        h_lens,
+        h_draft,
+        record_event: bool = True,
+    ) -> None:
+        """Scatter outputs into ``store`` rows ``slots`` and copy the store to host.
+
+        ``record_event=False``: the caller records an event right after this
+        call (the sampler event) and passes it to ``set_completion_event``.
+        """
+        i = self._slot_buffer()
+        buf = self._slots[i]
+        n = len(slots)
+        if n > buf.numel():
+            raise ValueError("too many sampling requests for the slot buffer")
+        buf.numpy()[:n] = slots
+        self._keep[i] = (h_new, h_lens, h_draft)
+        _sampler_store_kernel[(1,)](
+            buf.data_ptr(),
+            n,
+            o_new,
+            o_new.stride(0),
+            o_new.stride(1),
+            o_new.shape[1],
+            o_next_new,
+            o_next_new.stride(0),
+            o_next_new.stride(1),
+            o_next_new.shape[1],
+            o_lens,
+            o_lens.stride(0),
+            o_draft,
+            o_draft.stride(0),
+            o_draft.stride(1),
+            o_draft.shape[1],
+            store.new_tokens,
+            store.new_tokens.stride(0),
+            store.new_tokens.stride(1),
+            store.next_new_tokens,
+            store.next_new_tokens.stride(0),
+            store.next_new_tokens.stride(1),
+            store.new_tokens_lens,
+            store.next_draft_tokens,
+            store.next_draft_tokens.stride(0),
+            store.new_tokens.numel(),
+            store.new_tokens_lens.numel(),
+            store.next_draft_tokens.numel(),
+            h_new.data_ptr(),
+            h_lens.data_ptr(),
+            h_draft.data_ptr(),
+            BLOCK=64,
+            num_warps=1,
+        )
+        if record_event:
+            ev = torch.cuda.Event()
+            ev.record()
+            self._events[i] = ev
+        else:
+            self._events[i] = _NEVER_DONE
+            self._last = i
+        self.num_calls += 1
+
+    def set_completion_event(self, event: torch.cuda.Event) -> None:
+        """Mark the last launch's slot list reusable once ``event`` completes."""
+        i = getattr(self, "_last", None)
+        if i is not None:
+            self._events[i] = event
+            self._last = None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/stager.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/stager.py
new file mode 100644
index 000000000000..cd31cce01caf
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/stager.py
@@ -0,0 +1,703 @@
+# 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.
+"""deferred_input_prep staging of one decode input preparation into a single device launch.
+
+Native input preparation between two decode-graph replays issues ~50 small
+stream-ordered operations (H2D copies, index/elementwise kernels, a
+page-table expansion kernel). Each costs launch/dependency latency on the
+critical stream although it moves only bytes. While a ``PrepStager`` is
+active, the instrumented call sites hand those writes to the stager instead;
+``commit`` then applies all of them with one kernel that reads the host
+values straight from pinned memory, and finally runs the few native compute
+operations that must observe the staged values (``defer``).
+
+Semantics preserved per staged write:
+
+* H2D copies: the host bytes are read when ``commit`` runs, i.e. after the
+  whole host preparation, which is also when the native asynchronous copy
+  would have read them (the host runs ahead of the GPU and must not modify a
+  pending copy's source, see ``KVCacheManager._stage_block_offsets_for_copy``).
+* Device gathers from the previous iteration's sampler store read the same
+  device values the native index kernels read: the store is written before
+  the commit in stream order and the packet never writes it.
+* Order: staged writes are added to the packet in staging (native) order.
+  Descriptors of one launch run concurrently, so a write that overlaps an
+  earlier one of the same packet, or that does not fit, first flushes the
+  packet (one more launch); a write too large for an empty packet is issued
+  natively at that point. The final bytes therefore always equal the native
+  sequence of writes.
+* Deferral: a staged write lands at the launch instead of at its call site.
+  This is equivalent only if no other device operation in between touches
+  the destination (the window contract below).
+
+The target declares one ``PrepStager`` to the engine as
+``RuntimeHooks.step_input_stager`` (``runtime.make_runtime_hooks``). It is
+attached to the objects that own the instrumented call sites (their
+``step_input_stager`` attribute) only for the duration of one preparation; the
+native path runs whenever that attribute is ``None``.
+
+Deferred-write window. ``DecoderRunner._forward_decoder`` opens it with
+``begin_if_eligible`` right before ``_prepare_inputs`` and closes it with
+``end`` right before the decode-graph replay. In between, the instrumented
+call sites (on the runner, the attention and spec metadata and the KV cache
+managers, each marked "Staged write") only record their writes, and the
+commit at ``end`` performs them. Contract for code that runs inside the
+window:
+
+* No other device operation may read or write a buffer the stager was
+  handed. A read sees the bytes from before the staged write; a write is
+  overwritten by the commit.
+* Route new writes to such a buffer through the staging API (``copy``,
+  ``values``, ``gather``, ``zero``, ``clamp_min``, ``block_offsets``), or call
+  ``flush`` first, which issues every write staged so far in staging order.
+* Exception: the runner's ``input_ids_cuda`` and ``position_ids_cuda``.
+  ``overlap_inputs`` and ``position_ids`` also stage the decode graph's static
+  copies of them (``shared_static_tensors``), and the replay then skips its
+  own copy from the runner buffers, so a later staged write to either buffer
+  does not reach the graph. To change them inside the window, call ``flush``
+  and then write them natively; ``flush`` returns both to the replay's copy.
+* Deferred native operations (``defer``) run after the commit at ``end``,
+  also when ``flush`` was called. Their outputs (the indexer's scheduler
+  metadata when ``scheduler_metadata`` declines) are not written before
+  ``end``, and they read the values staged up to ``end``, including writes
+  staged after the ``defer`` call.
+* Nothing checks the contract in a production run. The debug guard
+  (``PrepStager(debug_level=1)``, which
+  ``TRTLLM_MODELING_V2_VALIDATE=1`` selects; ``guard``
+  lists what it covers) raises on an operation that touches a staged buffer
+  inside the window. ``TRTLLM_STAIRCASE_GLM_LL_DISABLE=deferred_input_prep``
+  keeps every preparation native.
+"""
+
+import contextlib
+import functools
+from typing import Callable, List, Optional, Sequence
+
+import torch
+
+from ...dispatch import DISABLE_ENV
+from .guard import DebugWindow
+from .packet import PacketFull, PacketOverlap, PrepPacket
+
+BAD_PAGE_INDEX = -1
+
+_PREP_NOTICE = (
+    "deferred-write window from DecoderRunner._forward_decoder's begin_if_eligible to "
+    "stager.end(): the buffers written through the deferred_input_prep stager are written "
+    "later by the commit; do not read or write them in between with other device ops; route "
+    "writes through the stager API or call stager.flush() first; for input_ids_cuda and "
+    "position_ids_cuda (also staged into the graph's static inputs) call stager.flush() "
+    f"and then write natively; disable with {DISABLE_ENV}=deferred_input_prep"
+)
+
+
+def _host_memory_pinned() -> bool:
+    """Whether TensorRT-LLM keeps its host staging buffers pinned."""
+    from tensorrt_llm._utils import prefer_pinned
+
+    return prefer_pinned()
+
+
+def _log_prep_decision(reason: Optional[str]) -> None:
+    """Log the first preparation decision of each kind (once per process and rank).
+
+    ``deferred_input_prep/prep`` is the engaged decision, logged with the window contract;
+    every decline reason gets its own key, so the log shows each kind of
+    preparation deferred_input_prep left native.
+    """
+    from ...dispatch import log_capability
+
+    if reason is None:
+        log_capability("deferred_input_prep", "prep", True, _PREP_NOTICE)
+    else:
+        log_capability("deferred_input_prep", f"prep[{reason}]", False, reason)
+
+
+def _stager_internal(method):
+    """Run a staging method as the stager's own work (not checked by the debug guard)."""
+
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        with self._internal():
+            return method(self, *args, **kwargs)
+
+    return wrapper
+
+
+class PrepStager:
+    """Collects one preparation's device writes; see the module docstring.
+
+    ``debug_level``: 0 (default) runs no checks; 1 turns on the debug guard; 2 the guard
+    plus poisoning of every staged destination (``guard``).
+    """
+
+    def __init__(self, debug_level: int = 0):
+        if debug_level not in (0, 1, 2):
+            raise ValueError(f"debug_level must be 0, 1 or 2, got {debug_level!r}")
+        self.packet = PrepPacket()
+        self.active = False
+        # Staged writes in native order: (kind, args).
+        self._writes: List[tuple] = []
+        self._deferred: List[Callable[[], None]] = []
+        self.num_commits = 0
+        self.last_num_descriptors = 0
+        self.last_reason: Optional[str] = None
+        self.decisions: dict = {}
+        self.num_native_copies = 0
+        self.num_flushes = 0
+        self._attached: list = []
+        self._static_input_ids = 0
+        self._static_positions = 0
+        # Why every preparation of this stager stays native (None: may engage).
+        self._decline_reason: Optional[str] = None
+        if not _host_memory_pinned():
+            # block_offsets stages copies from the pinned host page table; it
+            # is pageable when prefer_pinned() is False (Confidential Compute).
+            self._decline_reason = "host page tables are not pinned (prefer_pinned)"
+        self.debug: Optional[DebugWindow] = DebugWindow(debug_level) if debug_level else None
+
+    # ---- lifecycle ------------------------------------------------------
+
+    def begin(self) -> None:
+        if self.active:
+            # A previous preparation raised before its launch; its writes are void.
+            self.discard()
+        self.packet.begin()
+        self._writes = []
+        self._deferred = []
+        self._static_input_ids = 0
+        self._static_positions = 0
+        self.active = True
+        if self.debug is not None:
+            self.debug.open()
+
+    def commit(self, record_event: bool = True) -> None:
+        """Apply every staged write (one launch), then the deferred native ops.
+
+        ``record_event=False``: the caller records an event right after this
+        returns and hands it to ``set_completion_event``.
+        """
+        if not self.active:
+            return
+        try:
+            if self.debug is not None:
+                self.debug.close()
+                self.debug.verify()
+            for kind, args in self._writes:
+                self._apply(kind, args)
+            self.last_num_descriptors = self.packet.commit(record_event=record_event)
+            deferred = self._deferred
+        finally:
+            self._writes = []
+            self._deferred = []
+            self.active = False
+        self.num_commits += 1
+        if self.debug is not None and self.num_commits % 64 == 1:
+            from tensorrt_llm.logger import logger
+
+            logger.info(
+                f"[glm_moe_dsa] deferred_input_prep debug level {self.debug.level}: commits={self.num_commits} "
+                f"decisions={self.decisions} flushes={self.num_flushes} "
+                f"native={self.num_native_copies} guard_checked_ops={self.debug.num_checked_ops}"
+            )
+        for fn in deferred:
+            fn()
+
+    def flush(self) -> None:
+        """Issue every write staged so far, in staging order, before the window closes.
+
+        For code inside the deferred-write window that must touch a staged
+        buffer with another device operation: once ``flush`` returns, the
+        staged values precede that operation in stream order. Later stagings
+        go into a new packet; deferred native operations still run at the
+        commit. The graph replay copies ``input_ids`` and ``position_ids``
+        from the runner buffers again (``static_inputs_tokens`` is 0 unless
+        later stagings write both static inputs again), so a native write to
+        ``input_ids_cuda`` or ``position_ids_cuda`` after ``flush`` reaches
+        the graph.
+        """
+        if not self.active:
+            return
+        with self._internal():
+            if self.debug is not None:
+                self.debug.verify()
+            writes, self._writes = self._writes, []
+            for kind, args in writes:
+                self._apply(kind, args)
+            if self.packet.num_descriptors:
+                self._flush()
+        self._static_input_ids = 0
+        self._static_positions = 0
+
+    def _flush(self) -> None:
+        self.packet.commit(record_event=True)
+        self.packet.begin()
+        self.num_flushes += 1
+
+    def _apply(self, kind: str, args: tuple) -> None:
+        try:
+            self._add(kind, args)
+            return
+        except (PacketFull, PacketOverlap):
+            pass
+        if self.packet.num_descriptors:
+            self._flush()
+            try:
+                self._add(kind, args)
+                return
+            except PacketFull:
+                pass
+        self.num_native_copies += 1
+        self._native(kind, args)
+
+    def _add(self, kind: str, args: tuple) -> None:
+        pk = self.packet
+        if kind == "copy":
+            dst, src, clamp = args
+            pk.add_copy(dst, src, clamp_min=clamp)
+        elif kind == "values":
+            pk.add_values(*args)
+        elif kind == "gather":
+            pk.add_gather(*args)
+        elif kind == "zero":
+            pk.add_zero(*args)
+        elif kind == "clamp":
+            pk.add_clamp_min(*args)
+        elif kind == "block_row":
+            dst_k, dst_v, host_row, scale, kv_offset = args
+            pk.add_block_offsets(
+                dst_k,
+                dst_v,
+                host_row.data_ptr(),
+                host_row.numel(),
+                scale,
+                kv_offset,
+                BAD_PAGE_INDEX,
+            )
+        else:
+            raise ValueError(f"unknown staged write kind {kind!r}")
+
+    @staticmethod
+    def _native(kind: str, args: tuple) -> None:
+        """Issue one staged write with native torch operations."""
+        if kind == "copy":
+            dst, src, clamp = args
+            dst.copy_(src, non_blocking=True)
+            if clamp is not None:
+                dst.clamp_(min=clamp)
+        elif kind == "values":
+            dst, values = args
+            dst.copy_(torch.tensor(values, dtype=dst.dtype), non_blocking=True)
+        elif kind == "gather":
+            dst, src, indices, addend = args
+            flat = torch.as_strided(
+                src, (src.untyped_storage().nbytes() // 4,), (1,), src.storage_offset()
+            )
+            dst.zero_()
+            if indices:
+                idx = torch.tensor(indices, dtype=torch.long).to(src.device, non_blocking=True)
+                dst[: len(indices)].copy_(flat[idx] + addend)
+        elif kind == "zero":
+            args[0].zero_()
+        elif kind == "clamp":
+            dst, lo = args
+            dst.clamp_(min=lo)
+        elif kind == "block_row":
+            dst_k, dst_v, host_row, scale, kv_offset = args
+            v = host_row.to(dst_k.device, non_blocking=True)
+            bad = v == BAD_PAGE_INDEX
+            dst_k.copy_(torch.where(bad, 0, v * scale))
+            dst_v.copy_(torch.where(bad, 0, v * scale + kv_offset))
+        else:
+            raise ValueError(f"unknown staged write kind {kind!r}")
+
+    def _stage(self, kind: str, args: tuple, *dsts, overwrite: bool = True) -> None:
+        self._writes.append((kind, args))
+        if self.debug is not None:
+            for dst in dsts:
+                self.debug.register(dst, overwrite)
+
+    def _internal(self):
+        """Context for the stager's own operations (not checked by the debug guard)."""
+        return self.debug.internal() if self.debug is not None else contextlib.nullcontext()
+
+    # ---- engagement (decided before any device work of the preparation) -
+
+    _ATTACH_ATTR = "step_input_stager"
+
+    def begin_if_eligible(
+        self,
+        runner,
+        requests,
+        attn_metadata,
+        spec_metadata,
+        new_tensors_device,
+        *,
+        graph_replay: bool,
+        is_dummy: bool,
+        enable_spec_decode: bool,
+        promoted: bool,
+    ) -> Optional["PrepStager"]:
+        """Begin staging this preparation if it is in the deferred_input_prep support domain.
+
+        Support domain: replay of an already captured decode graph, no dummy or
+        promoted-context rows, generation-only batch whose rows all
+        carry the previous iteration's device tensors (overlap scheduler),
+        linear one-model speculative decoding, no beam search / helix /
+        attention-DP / pipeline parallelism / MRoPE / LoRA / guided decoding,
+        DSA attention metadata on a KV-cache manager V2, pinned host staging
+        memory (``prefer_pinned``). Anything else returns ``None`` (native
+        path). The first decision of each kind is logged once per process
+        (``deferred_input_prep/prep`` when engaged, with the window contract;
+        ``deferred_input_prep/prep[<reason>]`` per decline reason). An engaged call opens the
+        deferred-write window (module docstring).
+        """
+        reason = (
+            self._decline_reason
+            if self._decline_reason is not None
+            else self._ineligible_reason(
+                runner,
+                requests,
+                attn_metadata,
+                spec_metadata,
+                new_tensors_device,
+                graph_replay,
+                is_dummy,
+                enable_spec_decode,
+                promoted,
+            )
+        )
+        self.last_reason = reason
+        key = reason or "engaged"
+        seen = self.decisions.get(key, 0)
+        self.decisions[key] = seen + 1
+        if not seen:
+            _log_prep_decision(reason)
+        if reason is not None:
+            return None
+        self.begin()
+        self._attached = [runner, attn_metadata, spec_metadata, attn_metadata.kv_cache_manager]
+        draft_manager = getattr(attn_metadata, "draft_kv_cache_manager", None)
+        if draft_manager is not None:
+            self._attached.append(draft_manager)
+        for obj in self._attached:
+            setattr(obj, self._ATTACH_ATTR, self)
+        return self
+
+    @staticmethod
+    def _ineligible_reason(
+        runner,
+        requests,
+        attn_metadata,
+        spec_metadata,
+        new_tensors_device,
+        graph_replay,
+        is_dummy,
+        enable_spec_decode,
+        promoted,
+    ) -> Optional[str]:
+        cfg = runner._config
+        if not graph_replay or is_dummy or promoted:
+            return "not a steady decode-graph replay"
+        if not enable_spec_decode or spec_metadata is None or cfg.spec_config is None:
+            return "speculative decoding off"
+        mode = cfg.spec_config.spec_dec_mode
+        if not (mode.use_one_engine() and cfg.spec_config.is_linear_tree):
+            return "not linear one-model speculation"
+        if cfg.disable_overlap_scheduler or cfg.enable_attention_dp or cfg.use_mrope:
+            return "overlap off, attention DP or MRoPE"
+        if runner.use_beam_search or runner.mapping.has_cp_helix():
+            return "beam search or helix"
+        if getattr(runner.mapping, "pp_size", 1) > 1:
+            # deferred_input_prep is not validated with the pipeline-parallel executor loop.
+            return "pipeline parallelism"
+        if runner.guided_decoder is not None or runner._lora.cuda_graph_manager is not None:
+            return "guided decoding or LoRA graphs"
+        if getattr(new_tensors_device, "next_draft_tokens", None) is None:
+            return "no previous speculative device tensors"
+        if requests.num_context_requests != 0:
+            return "context rows"
+        for r in requests.generation_requests:
+            if r.is_dummy or r.py_batch_idx is None:
+                return "dummy row or row without previous tensors"
+        if type(attn_metadata).__name__ != "DSAtrtllmAttentionMetadata":
+            return "attention metadata is not DSA"
+        managers = [
+            attn_metadata.kv_cache_manager,
+            getattr(attn_metadata, "draft_kv_cache_manager", None),
+        ]
+        for manager in managers:
+            if manager is None:
+                continue
+            if not hasattr(manager, "index_mapper") or getattr(
+                manager, "_use_per_layer_page_tables", True
+            ):
+                return "KV cache manager is not V2 with a shared page table"
+        return None
+
+    def end(self, record_event: bool = True) -> None:
+        """Close the deferred-write window: detach and launch the staged preparation."""
+        self._detach()
+        self.commit(record_event=record_event)
+
+    def set_completion_event(self, event: torch.cuda.Event) -> None:
+        """Event recorded by the caller right after ``end(record_event=False)``."""
+        self.packet.set_completion_event(event)
+
+    def discard(self) -> None:
+        """Detach after a failed preparation; the staged writes are dropped."""
+        self._detach()
+        if self.debug is not None:
+            self.debug.close()
+            self.debug.staged = []
+        self._writes = []
+        self._deferred = []
+        self.packet.abort()
+        self.active = False
+
+    def _detach(self) -> None:
+        for obj in getattr(self, "_attached", ()):
+            if getattr(obj, self._ATTACH_ATTR, None) is self:
+                setattr(obj, self._ATTACH_ATTR, None)
+        self._attached = []
+
+    # ---- staging API used by the call sites -----------------------------
+
+    def copy(self, dst: torch.Tensor, src: torch.Tensor, clamp_min: Optional[int] = None) -> None:
+        """Stage ``dst.copy_(src, non_blocking=True)`` for a host ``src``.
+
+        ``src`` is kept alive and read at commit time.
+        """
+        with self._internal():
+            if src.is_cuda or src.numel() != dst.numel():
+                raise ValueError("PrepStager.copy needs a host source of dst's size")
+            if src.dtype != dst.dtype:
+                # Same conversion as ``copy_``.
+                src = src.to(dst.dtype)
+            if dst.is_contiguous():
+                self._stage("copy", (dst, src, clamp_min), dst)
+                return
+            if dst.dim() != 2 or dst.stride(1) != 1:
+                raise ValueError("PrepStager.copy needs contiguous rows")
+            src2 = src.reshape(dst.shape)
+            for r in range(dst.shape[0]):
+                self._stage("copy", (dst[r], src2[r], clamp_min), dst[r])
+
+    def values(self, dst: torch.Tensor, values: Sequence[int]) -> None:
+        """Stage ``dst.copy_(torch.tensor(values))`` for an int32 ``dst``."""
+        with self._internal():
+            if dst.dtype != torch.int32 or dst.numel() != len(values):
+                raise ValueError("values needs an int32 destination of len(values)")
+            self._stage("values", (dst, list(values)), dst)
+
+    def gather(
+        self, dst: torch.Tensor, src: torch.Tensor, indices: Sequence[int], addend: int = 0
+    ) -> None:
+        """Stage ``dst[:k] = src_storage[indices] + addend``, ``dst[k:] = 0``."""
+        with self._internal():
+            self._stage("gather", (dst, src, list(indices), addend), dst)
+
+    def zero(self, dst: torch.Tensor) -> None:
+        """Stage ``dst.zero_()``."""
+        with self._internal():
+            self._stage("zero", (dst,), dst)
+
+    @_stager_internal
+    def overlap_inputs(
+        self,
+        runner,
+        previous_batch_indices: Sequence[int],
+        previous_pos_indices: Sequence[int],
+        num_tokens: int,
+        num_draft_tokens: int,
+        tokens_per_gen_step: int,
+        num_extend_requests: int,
+        new_tokens: torch.Tensor,
+        next_draft_tokens: torch.Tensor,
+        new_tokens_lens: torch.Tensor,
+    ) -> None:
+        """Stage the overlap-scheduler block of ``DecoderRunner._prepare_tp_inputs``.
+
+        Native: zero ``previous_pos_id_offsets_cuda`` and
+        ``previous_kv_lens_offsets_cuda``; upload the previous slots and
+        per-token slots; gather the previous iteration's accepted tokens into
+        ``input_ids_cuda``, its draft tokens into ``draft_tokens_cuda``, its
+        accepted counts into the position offsets and ``count - T`` into the
+        KV-length offsets, at the rows of the requests that carry previous
+        tensors.
+        """
+        T = tokens_per_gen_step
+        D = T - 1
+        prev = len(previous_batch_indices)
+        pos_off = runner.previous_pos_id_offsets_cuda
+        kv_off = runner.previous_kv_lens_offsets_cuda
+        if prev == 0:
+            self.zero(pos_off)
+            self.zero(kv_off)
+            return
+        nt_s0, nt_s1 = new_tokens.stride(0), new_tokens.stride(1)
+        nd_s0, nd_s1 = next_draft_tokens.stride(0), next_draft_tokens.stride(1)
+        nl_s0 = new_tokens_lens.stride(0)
+        self.values(runner.previous_batch_indices_cuda[:prev], previous_batch_indices)
+        tok_idx = [t * nt_s0 + slot * nt_s1 for slot in previous_batch_indices for t in range(T)]
+        self.gather(runner.input_ids_cuda[num_tokens : num_tokens + prev * T], new_tokens, tok_idx)
+        static = self._static_tensors(runner)
+        if static is not None and num_tokens == 0:
+            # Also fill the graph's static input buffer; CUDAGraphRunner.replay
+            # then skips its device-to-device copy of input_ids.
+            self.gather(static["input_ids"][: prev * T], new_tokens, tok_idx)
+            self._static_input_ids = prev * T
+        if D > 0:
+            draft_idx = [
+                slot * nd_s0 + d * nd_s1 for slot in previous_batch_indices for d in range(D)
+            ]
+            self.gather(
+                runner.draft_tokens_cuda[num_draft_tokens : num_draft_tokens + prev * D],
+                next_draft_tokens,
+                draft_idx,
+            )
+        self.values(runner.previous_pos_indices_cuda[: prev * T], previous_pos_indices)
+        lo, hi = (num_extend_requests - prev) * T, num_extend_requests * T
+        self._zero_except(pos_off, lo, hi)
+        self.gather(pos_off[lo:hi], new_tokens_lens, [i * nl_s0 for i in previous_pos_indices])
+        lo, hi = num_extend_requests - prev, num_extend_requests
+        self._zero_except(kv_off, lo, hi)
+        self.gather(
+            kv_off[lo:hi], new_tokens_lens, [s * nl_s0 for s in previous_batch_indices], addend=-T
+        )
+
+    @staticmethod
+    def _static_tensors(runner):
+        cgr = getattr(runner, "cuda_graph_runner", None)
+        static = getattr(cgr, "shared_static_tensors", None)
+        if not static or "input_ids" not in static or static["position_ids"].dim() != 2:
+            return None
+        return static
+
+    def position_ids(self, runner, dst: torch.Tensor, host_position_ids: torch.Tensor) -> None:
+        """Stage ``dst.copy_(host_position_ids)`` and the graph's static copy of it."""
+        with self._internal():
+            n = dst.numel()
+            self.copy(dst, host_position_ids)
+            static = self._static_tensors(runner)
+            if static is not None:
+                self.copy(static["position_ids"][0, :n], host_position_ids)
+                self._static_positions = n
+
+    def static_inputs_tokens(self) -> int:
+        """Tokens of the static graph inputs written by the last launch (0: none)."""
+        if self._static_input_ids and self._static_input_ids == self._static_positions:
+            return self._static_input_ids
+        return 0
+
+    def _zero_except(self, buf: torch.Tensor, lo: int, hi: int) -> None:
+        if lo > 0:
+            self.zero(buf[:lo])
+        if hi < buf.numel():
+            self.zero(buf[hi:])
+
+    def clamp_min(self, dst: torch.Tensor, lo: int) -> None:
+        """Stage ``dst.clamp_(min=lo)`` (reads the destination's current bytes)."""
+        with self._internal():
+            if dst.dtype != torch.int32:
+                raise ValueError("clamp_min supports int32 only")
+            self._stage("clamp", (dst, lo), dst, overwrite=False)
+
+    def block_offsets(
+        self,
+        host_table: torch.Tensor,
+        dst: torch.Tensor,
+        copy_index: Sequence[int],
+        index_scales: Sequence[int],
+        kv_offsets: Sequence[int],
+    ) -> None:
+        """Stage ``copy_batch_block_offsets_to_device`` (KV cache manager V2).
+
+        ``host_table``: pinned ``[pools, src_seqs, 2, blocks]`` page table;
+        ``dst``: ``[pools, dst_seqs, 2, blocks]``. For every pool ``p`` and
+        output row ``s``: ``dst[p, s, 0] = scale * host[p, copy_index[s], 0]``
+        and ``dst[p, s, 1] = dst[p, s, 0] + kv_offset`` with ``BAD_PAGE_INDEX``
+        mapped to 0 in both.
+        """
+        with self._internal():
+            if not host_table.is_pinned() or host_table.dtype != torch.int32:
+                raise ValueError("block_offsets needs the pinned int32 host page table")
+            pools, _, kv, blocks = host_table.shape
+            if kv != 2 or dst.shape[0] != pools or dst.shape[3] != blocks:
+                raise ValueError("block_offsets: unexpected page-table layout")
+            for p in range(pools):
+                for s, src_row in enumerate(copy_index):
+                    self._stage(
+                        "block_row",
+                        (
+                            dst[p, s, 0],
+                            dst[p, s, 1],
+                            host_table[p, src_row, 0],
+                            int(index_scales[p]),
+                            int(kv_offsets[p]),
+                        ),
+                        dst[p, s, 0],
+                        dst[p, s, 1],
+                    )
+
+    def scheduler_metadata(self, metadata) -> bool:
+        """Stage ``Indexer.prepare_scheduler_metadata`` from host KV lengths.
+
+        Native: ``kv_lens_cuda_2d[:G, :W] = gen_lens`` (broadcast), then
+        ``scheduler_metadata_buffer`` / ``scheduler_metadata_buffer_full_next_n``
+        = DeepGEMM paged-MQA schedules for next_n = 1 / W. ``gen_lens`` are the
+        values this packet uploads into ``kv_lens_cuda``, known on the host
+        (``kv_lens_runtime``), so the three results are staged as copies of the
+        host replica (``paged_mqa_schedule``). Returns False (caller defers the
+        native op) outside the replicated configuration.
+        """
+        with self._internal():
+            return self._scheduler_metadata(metadata)
+
+    def _scheduler_metadata(self, metadata) -> bool:
+        from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _DG_SCHEDULE_BLOCK_KV
+
+        from .paged_mqa_schedule import SPLIT_KV_FOR_BLOCK_KV, paged_mqa_logits_schedule
+
+        split_kv = SPLIT_KV_FOR_BLOCK_KV.get(_DG_SCHEDULE_BLOCK_KV)
+        if (
+            split_kv is None
+            or metadata.use_expanded_buffers_for_mtp
+            or metadata._indexer_compress_ratio > 1
+        ):
+            return False
+        num_ctx, num_gen = metadata.num_contexts, metadata.num_generations
+        lens = metadata.kv_lens_runtime[num_ctx : num_ctx + num_gen].tolist()
+        num_sms = metadata.num_sms
+        width = min(metadata.gen_token_stride, metadata.kv_lens_cuda_2d.shape[1])
+        sched = metadata.scheduler_metadata_buffer
+        if tuple(sched.shape) != (num_sms + 1, 2) or sched.dtype != torch.int32:
+            return False
+        metadata.gen_indexer_kv_lens_cuda_runtime = metadata.get_indexer_kv_lens(
+            metadata.kv_lens_cuda_runtime[num_ctx : num_ctx + num_gen]
+        )
+        lens_2d = torch.tensor([[x] * width for x in lens], dtype=torch.int32)
+        self.copy(metadata.kv_lens_cuda_2d[:num_gen, :width], lens_2d)
+        self.copy(sched, torch.from_numpy(paged_mqa_logits_schedule(lens, 1, split_kv, num_sms)))
+        if metadata.max_draft_tokens > 0:
+            self.copy(
+                metadata.scheduler_metadata_buffer_full_next_n,
+                torch.from_numpy(paged_mqa_logits_schedule(lens, width, split_kv, num_sms)),
+            )
+        return True
+
+    def defer(self, fn: Callable[[], None]) -> None:
+        """Run ``fn`` (native device work that reads staged values) after the launch."""
+        self._deferred.append(fn)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/__init__.py
new file mode 100644
index 000000000000..53d1f868c0a6
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/__init__.py
@@ -0,0 +1,72 @@
+# 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.
+"""deferred_norm: deferred RMSNorm across the post-MoE hop and the next layer's
+projection, with staged projection weights (the deferred-norm packet; a numerical variant).
+
+Requires ``post_moe_hop`` (the hop) and ``absorbed_query`` (the next layer's query path). Registered
+entry points (``dispatch.get_impl("deferred_norm", name)``):
+
+* ``deferred_hop``: :class:`DeferredNormHop` factory, built at model
+  construction from a layer's engaged ``PostMoeHop`` (collective); with
+  ``routed_moe_into_hop`` also selected, its ``enable_routed_input()`` (collective) lets
+  ``hop_routed`` serve the routed input routed_moe writes into the post_moe_hop;
+* ``dsa_proj(mla, position_ids, packet, attn_metadata)``: the query_front
+  packet of a target layer whose hidden state is a
+  :class:`DeferredNormHidden` (absorbed_query / shared_indexer_absorbed_query kernels on the
+  deferred_norm projection), or ``None`` to fall back;
+* ``is_deferred(hidden_states)``: whether a hidden state is a deferred_norm packet.
+
+Rounding changes versus native (see ``hop_kernel.py`` and
+``proj_kernel.py``): the next-layer projection rounds ``bf16(x * gamma)``
+instead of ``bf16((x * r) * gamma)`` and applies ``r`` after its fp32
+accumulation (also for the indexer ``wk | weights_proj`` rows of
+full-indexer layers), with its own K accumulation order. The residual
+stream, ``r`` itself and the fallback materialization are bit-identical to
+the post_moe_hop.
+
+Scope (``scope.py``): the served scope is ``shared``, which serves only
+boundaries into shared-indexer layers, so full-indexer layers keep the post-MoE
+hop and their native or absorbed-query projection (no deferred-norm rounding
+change on their indexer inputs); ``all`` would serve every boundary into a
+target layer. The choice is made per layer at hop construction and agreed over
+the TP group.
+
+The hop runs on the post_moe_hop workspace with post_moe_hop's protocol, error handling
+included: device errors land in the workspace's sticky error words, so the
+post_moe_hop health check (``dispatch.get_impl("post_moe_hop", "health_check")``, called by the
+post_moe_hop adapter after warmup and at executor cleanup) also covers deferred_norm hops, and
+its call counter counts them. Every inline-PTX helper of the deferred_norm kernels is
+deferred_norm's own (``ptx.py``).
+"""
+
+from ... import dispatch
+from .hop import DeferredNormHop, materialize
+from .packet import DeferredNormHidden, is_deferred
+from .projection import DeferredNormProjConfig, DeferredNormProjection
+from .query_entry import deferred_norm_dsa_proj
+
+dispatch.register("deferred_norm", "deferred_hop", DeferredNormHop)
+dispatch.register("deferred_norm", "dsa_proj", deferred_norm_dsa_proj)
+dispatch.register("deferred_norm", "is_deferred", is_deferred)
+
+__all__ = [
+    "DeferredNormHidden",
+    "DeferredNormHop",
+    "DeferredNormProjConfig",
+    "DeferredNormProjection",
+    "is_deferred",
+    "materialize",
+    "deferred_norm_dsa_proj",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop.py
new file mode 100644
index 000000000000..2f608a914a66
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop.py
@@ -0,0 +1,666 @@
+# 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.
+"""deferred-norm post-MoE hop (deferred-norm packet, numerical variant).
+
+``DeferredNormHop`` wraps an engaged :class:`PostMoeHop` of one target MoE
+layer whose next layer is a target attention layer. It runs the same
+finalize / multicast all-reduce / residual protocol on the post_moe_hop workspace, but
+returns a :class:`DeferredNormHidden` packet instead of the normalized hidden
+state; the next layer's projection (``DeferredNormProjection``) applies the RMSNorm
+scale after its GEMM. The residual output is bit-identical to post_moe_hop's.
+
+The hop shares the post_moe_hop workspace protocol, including post_moe_hop's error handling:
+the post_moe_hop's poll deadline (``timeout_ns``) bounds every deferred_norm poll; device
+errors (a peer that never publishes, a permutation entry beyond the FC2
+rows) set the sticky per-buffer error words that post_moe_hop's ``check_health``
+reads, and every later post_moe_hop or deferred_norm call of the group publishes NaN until a
+collective workspace reset. After the MoE runner the hop applies post_moe_hop's
+raw-tuple contract: a dtype or shape that differs from the declared contract
+raises ``PostMoeHopContractError`` (the runner already skipped its finalize, so no
+native path is left), and a layout the kernel cannot read with 16-byte
+vectors is copied into an aligned contiguous tensor on this rank only
+(counted in the workspace's ``layout_repairs``).
+
+With routed_moe_into_hop selected, routed_moe writes its routed result into the post_moe_hop's
+routed-input buffer and the layer calls ``PostMoeHop.hop_routed``. After
+:meth:`DeferredNormHop.enable_routed_input` (collective; only where this hop
+is engaged and the post_moe_hop's routed input is enabled on every rank),
+:meth:`DeferredNormHop.hop_routed` takes the same inputs and runs the same
+protocol on the same routed buffer, and returns the deferred packet instead
+(its materialization equals ``hop_routed``'s output bit for bit).
+
+Eligibility is decided collectively once and per call from rank-identical
+host state only, before the MoE runner runs, so every rank takes the same hop
+kernel. The collective decision covers every kernel a packet can reach: the
+hop, the materialization and both served projection kernels of the next
+layer (shared- and full-indexer layers). A rank that cannot compile one of
+them disengages the hop on every rank, so no rank produces a packet that
+another rank's consumer could not project; the consumer's own compile is
+then a cache hit and its remaining declines are rank-identical static checks.
+
+The scope (``scope.py``; the served scope is ``shared``) is part of the same local
+capability: with ``shared``, a hop whose next layer is a full-indexer layer (or
+of unknown type) is incapable on this rank before any compile, and the single
+agreement disengages it on every rank, so that layer keeps the post-MoE hop. The
+number of collectives does not depend on the scope.
+"""
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from ... import dispatch
+from ..post_moe_hop import PostMoeHopContractError
+from ..post_moe_hop.kernels import STATE_WORDS
+from ..post_moe_hop.workspace import PostMoeHopWorkspace, _all_ranks_agree
+from .packet import DeferredNormHidden, frag_elems
+from .scope import UNRESTRICTED_SCOPE, decline_reason, next_layer_kind, selected_scope
+
+FEATURE_ID = "deferred_norm"
+MAX_ROWS = 8
+DEFAULT_CTAS = 8
+_SCALE_DTYPES = (torch.bfloat16, torch.float32)
+# routed_moe_into_hop's routed-input modes (PostMoeHop.ROUTED_MODES), in agreement order.
+_ROUTED_INPUT_MODES = ("row", "identity", "fp32")
+_VECTOR_BYTES = 16
+
+
+def _vector_ready(t: torch.Tensor, row_major_2d: bool) -> bool:
+    """Whether the hop kernel may read ``t`` with 16-byte vectors at its strides."""
+    if t.data_ptr() % _VECTOR_BYTES or t.stride(-1) != 1:
+        return False
+    if row_major_2d:
+        return (t.stride(0) * t.element_size()) % _VECTOR_BYTES == 0
+    return t.is_contiguous()
+
+
+def _graph_mode_ok() -> Tuple[bool, str]:
+    from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph
+    from tensorrt_llm._torch.utils import is_torch_compiling
+
+    if is_torch_compiling():
+        return False, "torch.compile graph keeps the native custom-op split"
+    if is_in_breakable_cuda_graph():
+        return False, "breakable CUDA graph keeps the native custom-op split"
+    return True, ""
+
+
+class DeferredNormHop:
+    """Deferred-norm post-MoE hop on the post_moe_hop workspace.
+
+    Args:
+        post_moe_hop: the layer's engaged ``PostMoeHop`` (its workspace, mapping,
+            dtype, PDL policy, poll deadline and test-only fault injection
+            are reused).
+        trigger_early: release the PDL dependents right after the hop's own
+            dependency wait (so the next projection stages its weights while
+            the hop communicates) instead of after the reduction.
+        ctas: CTAs per token row (8 columns per thread).
+        carveout: preferred shared-memory carveout of the hop kernel in percent
+            (-1: driver default).
+        next_layer_full_indexer: whether the layer after this hop's MoE layer
+            owns an indexer; read only when
+            the scope is not ``all``.
+            ``None``: resolved from the constructing decoder layer
+            (``scope.next_layer_kind``).
+    """
+
+    def __init__(
+        self,
+        post_moe_hop,
+        trigger_early: bool = True,
+        ctas: int = DEFAULT_CTAS,
+        carveout: int = -1,
+        next_layer_full_indexer: Optional[bool] = None,
+    ):
+        from .hop_kernel import hop_kernel_supported
+
+        self.post_moe_hop = post_moe_hop
+        self.mapping = post_moe_hop.mapping
+        self.hidden_size = post_moe_hop.hidden_size
+        self.top_k = post_moe_hop.top_k
+        self.dtype = post_moe_hop.dtype
+        self.use_pdl = post_moe_hop.use_pdl
+        self.timeout_ns = post_moe_hop.timeout_ns
+        self.fault_delay_ns = post_moe_hop.fault_delay_ns
+        self.trigger_early = trigger_early
+        self._repairs_logged = set()
+        # routed_moe_into_hop routed input (enable_routed_input): engaged on every rank or none.
+        self.routed_engaged = False
+        self.routed_mode: Optional[str] = None
+        self.routed_reason = "not enabled"
+        self.carveout = carveout
+        self.ctas = ctas
+        self.threads = self.hidden_size // (8 * ctas)
+        self.partials = ctas * (self.threads // 32)
+        self.engaged = False
+        reason = ""
+        self.scope = selected_scope()
+        # The next layer's indexer type matters only for a restricted scope.
+        self.next_layer_full: Optional[bool] = None
+        self.next_layer_idx: Optional[int] = None
+        local_ok = bool(post_moe_hop.engaged)
+        if not local_ok:
+            reason = "post_moe_hop hop not engaged"
+        else:
+            why = hop_kernel_supported(self.hidden_size, self.top_k, self.mapping.tp_size, ctas)
+            if why is not None:
+                local_ok, reason = False, why
+        if self.scope != UNRESTRICTED_SCOPE:
+            full, self.next_layer_idx, how = next_layer_kind(post_moe_hop, next_layer_full_indexer)
+            self.next_layer_full = full
+            why = decline_reason(self.scope, full, self.next_layer_idx, how)
+            if local_ok and why is not None:
+                local_ok, reason = False, why
+        if local_ok:
+            try:
+                why = self._precompile()
+            except Exception as exc:  # noqa: BLE001 - converged across ranks below
+                why = f"kernel compile failed: {exc!r}"
+            if why is not None:
+                local_ok, reason = False, why
+        if post_moe_hop.engaged:
+            # Same communicator and agreement helper as the post_moe_hop workspace.
+            comm = PostMoeHopWorkspace.comm_for(self.mapping)
+            if _all_ranks_agree(comm, local_ok):
+                self.engaged = True
+            else:
+                reason = reason or "another rank is not capable"
+        self.reason = reason
+        dispatch.log_capability(FEATURE_ID, "deferred_hop", self.engaged, reason)
+        if self.scope != UNRESTRICTED_SCOPE:
+            into = "?" if self.next_layer_idx is None else self.next_layer_idx
+            dispatch.log_capability(
+                FEATURE_ID, f"scope[{self.scope}]_into_layer{into}", self.engaged, reason
+            )
+
+    def _precompile(self) -> Optional[str]:
+        """Compile every kernel a packet of this hop can reach.
+
+        Returns why a served projection shape is unsupported, or ``None``;
+        compile errors propagate (converged across ranks by the caller).
+        """
+        from . import projection
+        from .hop_kernel import compiled_hop, compiled_materialize
+
+        for has_shared in (True, False):
+            for scale_fp32 in (False, True):
+                compiled_hop(
+                    self.hidden_size,
+                    self.top_k,
+                    self.mapping.tp_size,
+                    self.post_moe_hop.max_rows,
+                    self.ctas,
+                    has_shared,
+                    scale_fp32,
+                    self.use_pdl,
+                    self.trigger_early,
+                    STATE_WORDS,
+                    self.carveout,
+                    fault_delay_ns=self.fault_delay_ns,
+                )
+        compiled_materialize(self.hidden_size, self.ctas, self.partials, self.use_pdl)
+        return projection.precompile_served(torch.cuda.current_device())
+
+    def can_run(self, num_rows: int, attn_metadata) -> bool:
+        """Rank-identical per-call decision, taken before the MoE runner runs."""
+        if not self.engaged or not 0 < num_rows <= MAX_ROWS:
+            return False
+        if not self.post_moe_hop.can_run(num_rows):
+            return False
+        if (
+            attn_metadata is None
+            or attn_metadata.num_contexts != 0
+            or (attn_metadata.num_generations == 0)
+        ):
+            return False
+        ok, _ = _graph_mode_ok()
+        return ok
+
+    # ------------------------------------------------------------------
+    # routed_moe_into_hop routed input (deferred_norm composed with routed_moe_into_hop): reachable
+    # only after enable_routed_input(), i.e. with both "routed_moe_into_hop" and "deferred_norm"
+    # selected. ------------------------------------------------------------------
+    def _post_moe_routed_mode(self) -> Optional[str]:
+        """The post_moe_hop's routed-input mode, or ``None`` when it is not enabled."""
+        if not getattr(self.post_moe_hop, "_routed_enabled", False):
+            return None
+        return getattr(self.post_moe_hop, "routed_mode", None)
+
+    def enable_routed_input(self) -> bool:
+        """Collective: serve the post_moe_hop's routed input (routed_moe_into_hop) with the deferred hop.
+
+        Requires this hop to be engaged and the post_moe_hop's routed input to be
+        enabled (both rank-identical, agreed over the TP group). Compiles the
+        routed variants of the post_moe_hop's current mode (shared output present
+        and absent); every rank must succeed, otherwise the routed input
+        stays off on all of them and the layer keeps ``PostMoeHop.hop_routed``.
+        Call it again (collectively, outside capture) after the post_moe_hop
+        switches modes.
+        """
+        from .hop_kernel import compiled_hop
+
+        mode = self._post_moe_routed_mode()
+        local_ok, reason = True, ""
+        if not self.engaged:
+            local_ok, reason = False, "deferred_norm hop not engaged"
+        elif mode is None:
+            local_ok, reason = False, "post_moe_hop hop's routed input is not enabled"
+        else:
+            try:
+                for has_shared in (True, False):
+                    compiled_hop(
+                        self.hidden_size,
+                        self.top_k,
+                        self.mapping.tp_size,
+                        self.post_moe_hop.max_rows,
+                        self.ctas,
+                        has_shared,
+                        False,
+                        self.use_pdl,
+                        self.trigger_early,
+                        STATE_WORDS,
+                        self.carveout,
+                        fault_delay_ns=self.fault_delay_ns,
+                        routed_input=mode,
+                    )
+            except Exception as exc:  # noqa: BLE001 - converged across ranks below
+                local_ok, reason = False, f"routed kernel compile failed: {exc!r}"
+        agreed = False
+        if self.engaged and mode is not None:
+            # Same communicator and agreement helper as the post_moe_hop workspace;
+            # both conditions above are rank-identical, so either every rank
+            # reaches these collectives or none does. The post_moe_hop's mode comes
+            # from each rank's environment: require the same mode everywhere
+            # (one agreement per mode, in a fixed order).
+            comm = PostMoeHopWorkspace.comm_for(self.mapping)
+            same_mode = [m for m in _ROUTED_INPUT_MODES if _all_ranks_agree(comm, mode == m)]
+            if not same_mode:
+                local_ok, reason = (
+                    False,
+                    "the post_moe_hop hops' routed-input modes differ across ranks",
+                )
+            agreed = _all_ranks_agree(comm, local_ok)
+            if not agreed:
+                reason = reason or "another rank is not capable"
+        self.routed_engaged = agreed
+        self.routed_mode = mode if agreed else None
+        self.routed_reason = reason
+        dispatch.log_capability(FEATURE_ID, f"deferred_hop_routed[{mode}]", agreed, reason)
+        return agreed
+
+    def can_run_routed(self, num_rows: int, attn_metadata) -> bool:
+        """Rank-identical per-call decision for the routed input."""
+        return (
+            self.routed_engaged
+            and self._post_moe_routed_mode() == self.routed_mode
+            and self.can_run(num_rows, attn_metadata)
+        )
+
+    def _routed_view(self, num_rows: int):
+        """``(mode, buffer view, scales)`` of the post_moe_hop's routed input for ``num_rows``."""
+        base = self.post_moe_hop
+        if self.routed_mode == "fp32":
+            return "fp32", base.routed_acc_buffer(num_rows), None
+        if self.routed_mode == "identity":
+            rows, scales = base.routed_raw_buffers(num_rows)
+            return "identity", rows, scales
+        return "row", base.routed_input_buffer(num_rows), None
+
+    def hop_routed(
+        self,
+        routed: torch.Tensor,
+        shared_output: Optional[torch.Tensor],
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+    ) -> Tuple[DeferredNormHidden, torch.Tensor]:
+        """Return ``(packet, residual_out)``; same inputs as ``PostMoeHop.hop_routed``.
+
+        ``routed`` is the view of the post_moe_hop's routed-input buffer routed_moe wrote
+        (``routed_output_into_hop``). Same arithmetic as ``hop_routed`` up to
+        the residual; the RMSNorm is deferred into the packet.
+
+        Raises:
+            ValueError: the routed input is not engaged or ``can_run`` is
+                false for this row count.
+            PostMoeHopContractError: ``routed`` is not the hop's buffer for the row
+                count, or a dtype / shape differs from the contract.
+        """
+        num_rows = residual.shape[0]
+        if not (
+            self.routed_engaged
+            and self._post_moe_routed_mode() == self.routed_mode
+            and 0 < num_rows <= MAX_ROWS
+            and self.post_moe_hop.can_run(num_rows)
+        ):
+            raise ValueError(
+                "deferred_norm routed hop called for an ineligible shape or before "
+                "enable_routed_input(); check can_run_routed() first"
+            )
+        mode, expect, scales = self._routed_view(num_rows)
+        rows_h = (num_rows, self.hidden_size)
+        problems = []
+        if (
+            routed.data_ptr() != expect.data_ptr()
+            or routed.dtype != expect.dtype
+            or tuple(routed.shape) != tuple(expect.shape)
+            or tuple(routed.stride()) != tuple(expect.stride())
+        ):
+            problems.append(
+                f"routed {routed.dtype} {tuple(routed.shape)} is not the post_moe_hop hop's "
+                f"{mode} buffer"
+            )
+        if residual.dtype != self.dtype or tuple(residual.shape) != rows_h:
+            problems.append(f"residual {residual.dtype} {tuple(residual.shape)}")
+        if shared_output is not None and (
+            shared_output.dtype != self.dtype or tuple(shared_output.shape) != rows_h
+        ):
+            problems.append(f"shared_output {shared_output.dtype} {tuple(shared_output.shape)}")
+        if norm_weight.dtype != self.dtype or tuple(norm_weight.shape) != (self.hidden_size,):
+            problems.append(f"norm_weight {norm_weight.dtype} {tuple(norm_weight.shape)}")
+        if problems:
+            raise PostMoeHopContractError(
+                "routed_moe_into_hop routed input violates the hop's contract (deferred_norm hop): "
+                f"{'; '.join(problems)}"
+            )
+        if shared_output is not None:
+            shared_output = self._vector_layout("shared_output", shared_output)
+        residual = self._vector_layout("residual", residual)
+        norm_weight = self._vector_layout("norm_weight", norm_weight)
+        return self.launch_routed(routed, shared_output, residual, norm_weight, eps, scales=scales)
+
+    def launch_routed(
+        self,
+        routed,
+        shared_output,
+        residual,
+        norm_weight,
+        eps,
+        scales=None,
+        ts: Optional[torch.Tensor] = None,
+        row_bytes: Optional[int] = None,
+    ) -> Tuple[DeferredNormHidden, torch.Tensor]:
+        """Launch the routed variant on checked, vector-ready inputs (``hop_routed``).
+
+        Tests only: ``ts`` as in :meth:`launch`; ``row_bytes`` overrides the
+        routed buffer's row pitch (memcheck positive control).
+        """
+        if row_bytes is None:
+            row_bytes = routed.stride(0) * routed.element_size()
+        return self._launch(
+            routed.data_ptr(),
+            row_bytes,
+            routed.shape[0],
+            scales.data_ptr() if scales is not None else 0,
+            0,
+            False,
+            self.routed_mode,
+            shared_output,
+            residual,
+            norm_weight,
+            eps,
+            ts,
+        )
+
+    def _contract_violations(
+        self,
+        fc2_output,
+        expert_scale_factor,
+        expanded_idx_to_permuted_idx,
+        shared_output,
+        residual,
+        norm_weight,
+    ) -> List[str]:
+        """Dtype/shape differences from the declared raw-tuple contract (rank-identical)."""
+        num_rows = residual.shape[0]
+        rows_h = (num_rows, self.hidden_size)
+        problems = []
+        if residual.dtype != self.dtype or tuple(residual.shape) != rows_h:
+            problems.append(f"residual {residual.dtype} {tuple(residual.shape)}")
+        if (
+            fc2_output.dtype != self.dtype
+            or fc2_output.dim() != 2
+            or fc2_output.shape[1] < self.hidden_size
+        ):
+            problems.append(f"fc2_output {fc2_output.dtype} {tuple(fc2_output.shape)}")
+        if expert_scale_factor.dtype not in _SCALE_DTYPES or tuple(expert_scale_factor.shape) != (
+            num_rows,
+            self.top_k,
+        ):
+            problems.append(
+                f"expert_scale_factor {expert_scale_factor.dtype} "
+                f"{tuple(expert_scale_factor.shape)}"
+            )
+        if expanded_idx_to_permuted_idx.dtype != torch.int32 or tuple(
+            expanded_idx_to_permuted_idx.shape
+        ) != (num_rows, self.top_k):
+            problems.append(
+                f"expanded_idx_to_permuted_idx {expanded_idx_to_permuted_idx.dtype} "
+                f"{tuple(expanded_idx_to_permuted_idx.shape)}"
+            )
+        if shared_output is not None and (
+            shared_output.dtype != self.dtype or tuple(shared_output.shape) != rows_h
+        ):
+            problems.append(f"shared_output {shared_output.dtype} {tuple(shared_output.shape)}")
+        if norm_weight.dtype != self.dtype or tuple(norm_weight.shape) != (self.hidden_size,):
+            problems.append(f"norm_weight {norm_weight.dtype} {tuple(norm_weight.shape)}")
+        return problems
+
+    def _vector_layout(self, name: str, t: torch.Tensor, row_major_2d: bool = False):
+        """Return ``t`` or a 16-byte-aligned contiguous copy of it (rank-local)."""
+        if _vector_ready(t, row_major_2d):
+            return t
+        misalignment, strides = t.data_ptr() % _VECTOR_BYTES, tuple(t.stride())
+        if row_major_2d:
+            t = t[:, : self.hidden_size]
+        fixed = torch.empty(t.shape, dtype=t.dtype, device=t.device)
+        fixed.copy_(t)
+        self.post_moe_hop.workspace.layout_repairs += 1
+        if name not in self._repairs_logged:
+            self._repairs_logged.add(name)
+            from tensorrt_llm.logger import logger
+
+            logger.warning(
+                f"[glm_moe_dsa] deferred_norm/hop: {name} is not ready for 16-byte vector "
+                f"access (data_ptr % 16 = {misalignment}, strides {strides}); "
+                "copying it before every such call"
+            )
+        return fixed
+
+    def __call__(
+        self,
+        fc2_output: torch.Tensor,
+        expert_scale_factor: torch.Tensor,
+        expanded_idx_to_permuted_idx: torch.Tensor,
+        shared_output: Optional[torch.Tensor],
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+    ) -> Tuple[DeferredNormHidden, torch.Tensor]:
+        """Return ``(packet, residual_out)``; same inputs as ``PostMoeHop``.
+
+        Raises:
+            ValueError: ``can_run`` is false for this row count.
+            PostMoeHopContractError: a dtype or shape differs from the raw-tuple contract.
+        """
+        if not self.engaged or not 0 < residual.shape[0] <= MAX_ROWS:
+            raise ValueError(
+                "deferred_norm hop called for an ineligible shape; check can_run() first"
+            )
+        problems = self._contract_violations(
+            fc2_output,
+            expert_scale_factor,
+            expanded_idx_to_permuted_idx,
+            shared_output,
+            residual,
+            norm_weight,
+        )
+        if problems:
+            raise PostMoeHopContractError(
+                "raw MoE tuple violates the declared native raw-tuple contract after the runner "
+                f"skipped its finalize (deferred_norm hop): {'; '.join(problems)}"
+            )
+        fc2_output = self._vector_layout("fc2_output", fc2_output, row_major_2d=True)
+        expert_scale_factor = self._vector_layout("expert_scale_factor", expert_scale_factor)
+        expanded_idx_to_permuted_idx = self._vector_layout(
+            "expanded_idx_to_permuted_idx", expanded_idx_to_permuted_idx
+        )
+        if shared_output is not None:
+            shared_output = self._vector_layout("shared_output", shared_output)
+        residual = self._vector_layout("residual", residual)
+        norm_weight = self._vector_layout("norm_weight", norm_weight)
+        return self.launch(
+            fc2_output,
+            expert_scale_factor,
+            expanded_idx_to_permuted_idx,
+            shared_output,
+            residual,
+            norm_weight,
+            eps,
+        )
+
+    def launch(
+        self,
+        fc2_output,
+        expert_scale_factor,
+        expanded_idx_to_permuted_idx,
+        shared_output,
+        residual,
+        norm_weight,
+        eps,
+        ts: Optional[torch.Tensor] = None,
+        fc2_rows: Optional[int] = None,
+    ) -> Tuple[DeferredNormHidden, torch.Tensor]:
+        """Launch on checked, vector-ready inputs (``__call__``).
+
+        Tests only: ``ts``, int64 ``[R * ctas, 8]`` per-CTA timer stamps;
+        ``fc2_rows`` overrides the permutation bound (memcheck positive
+        control).
+        """
+        return self._launch(
+            fc2_output.data_ptr(),
+            fc2_output.stride(0) * 2,
+            fc2_output.shape[0] if fc2_rows is None else fc2_rows,
+            expert_scale_factor.data_ptr(),
+            expanded_idx_to_permuted_idx.data_ptr(),
+            expert_scale_factor.dtype.itemsize == 4,
+            None,
+            shared_output,
+            residual,
+            norm_weight,
+            eps,
+            ts,
+        )
+
+    def _launch(
+        self,
+        fc2_ptr: int,
+        fc2_row_bytes: int,
+        fc2_rows: int,
+        scales_ptr: int,
+        perm_ptr: int,
+        scale_fp32: bool,
+        routed_input: Optional[str],
+        shared_output,
+        residual,
+        norm_weight,
+        eps,
+        ts,
+    ) -> Tuple[DeferredNormHidden, torch.Tensor]:
+        from .hop_kernel import compiled_hop
+
+        num_rows = residual.shape[0]
+        ws = self.post_moe_hop.workspace
+        a_frag = torch.empty(
+            (frag_elems(self.hidden_size),), dtype=residual.dtype, device=residual.device
+        )
+        residual_out = torch.empty_like(residual)
+        sumsq = torch.empty((num_rows, self.partials), dtype=torch.float32, device=residual.device)
+        has_shared = shared_output is not None
+        compiled = compiled_hop(
+            self.hidden_size,
+            self.top_k,
+            self.mapping.tp_size,
+            ws.max_rows,
+            self.ctas,
+            has_shared,
+            scale_fp32,
+            self.use_pdl,
+            self.trigger_early,
+            ws.state.numel(),
+            self.carveout,
+            ts is not None,
+            self.fault_delay_ns,
+            routed_input,
+        )
+        compiled(
+            fc2_ptr,
+            fc2_row_bytes,
+            fc2_rows,
+            scales_ptr,
+            perm_ptr,
+            shared_output.data_ptr() if has_shared else 0,
+            residual.data_ptr(),
+            norm_weight.data_ptr(),
+            a_frag.data_ptr(),
+            residual_out.data_ptr(),
+            sumsq.data_ptr(),
+            ws.mc.data_ptr(),
+            ws.uc.data_ptr(),
+            ws.partials.data_ptr(),
+            0 if ts is None else ts.data_ptr(),
+            ws.state,
+            self.mapping.tp_rank,
+            self.timeout_ns,
+            num_rows,
+        )
+        packet = DeferredNormHidden(
+            a_frag=a_frag,
+            sumsq=sumsq,
+            residual=residual_out,
+            gamma=norm_weight,
+            eps=float(eps),
+            materializer=self.materialize,
+        )
+        return packet, residual_out
+
+    def materialize(self, packet: DeferredNormHidden) -> torch.Tensor:
+        """post_moe_hop's normalized output ``bf16((x * r) * gamma)`` from a packet."""
+        return materialize(packet, self.ctas, self.use_pdl)
+
+
+def materialize(
+    packet: DeferredNormHidden, ctas: int = DEFAULT_CTAS, use_pdl: Optional[bool] = None
+) -> torch.Tensor:
+    from .hop_kernel import compiled_materialize
+
+    if use_pdl is None:
+        from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+        use_pdl = get_env_enable_pdl()
+    x = packet.residual
+    out = torch.empty_like(x)
+    compiled = compiled_materialize(x.shape[1], ctas, packet.sumsq.shape[1], use_pdl)
+    compiled(
+        x.data_ptr(),
+        packet.gamma.data_ptr(),
+        packet.sumsq.data_ptr(),
+        out.data_ptr(),
+        packet.sumsq.view(torch.int32).view(-1),
+        packet.eps,
+        x.shape[0],
+    )
+    return out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop_kernel.py
new file mode 100644
index 000000000000..592ee2e0940c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/hop_kernel.py
@@ -0,0 +1,758 @@
+# 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.
+"""CuTe DSL kernels of the deferred-norm post-MoE hop.
+
+``DeferredNormHopKernel`` is the post_moe_hop cluster kernel (``post_moe_hop/cute_kernel.py``) with
+the RMSNorm deferred to the next layer's projection. Steps 1-5 and 7 are post_moe_hop's
+(finalize, multicast publish, clear of the previous call, Lamport poll,
+rank-ordered fp32 sum, bf16, residual add, arrival), on the same post_moe_hop
+workspace and with the same workspace protocol, so deferred_norm and post_moe_hop hops may
+follow each other on one workspace:
+
+* state words, epoch / previous rows / arrivals / call count (post_moe_hop's word
+  indices, ``ptx.py``);
+* error handling: every Lamport poll is bounded by the hop's ``timeout_ns``
+  deadline; on expiry the thread records the Lamport-timeout bit and its
+  poll count and writes a NaN residual (so the row's square sum, ``r`` and
+  every value derived from the packet are NaN); a permutation entry at or
+  beyond ``fc2_rows`` is skipped and records the out-of-range bit. Error
+  bits go into the error word of the call's Lamport buffer. A call reads the
+  previous buffer's error word (complete at the kernel boundary) and, when
+  it is nonzero, publishes NaN and carries the bits into its own word, so
+  every later call of the rank -- post_moe_hop or deferred_norm -- publishes NaN until the
+  collective workspace reset, and ``post_moe_hop.check_health`` sees deferred_norm's errors.
+
+Step 6 differs:
+
+* every thread forms post_moe_hop's per-thread sum of bf16-rounded squares of its 8
+  columns and the warp butterfly (16, 8, 4, 2, 1); lane 0 of every warp
+  stores the warp sum to ``sumsq[row, cta * WARPS + warp]`` (fp32). The
+  consumer sums the ``CTAS * WARPS`` values in index order, which is the
+  order of post_moe_hop's distributed-shared-memory reduction, so its ``rsqrt``
+  operand is bit-identical to post_moe_hop's;
+* every thread stores the gain-folded activation ``a = bf16(x * gamma)`` of
+  its 8 columns instead of the normalized output, in the MMA B-fragment
+  layout of ``packet.py`` (``scaled`` is the ``a_frag`` buffer), so that the
+  projection reads it with one coalesced 8-byte load per lane and K step.
+
+There is no cross-CTA dependency left inside a row, so the kernel is not a
+cluster launch. With ``trigger_early`` the PDL dependents (the next layer's
+projection, which reads only weights before its own dependency wait) are
+released right after this kernel's dependency wait instead of after the
+reduction.
+
+With ``routed_input`` (deferred_norm composed with routed_moe_into_hop) step 1 reads the
+routed input routed_moe wrote into the post_moe_hop's routed-input buffer instead of the
+raw tuple, with the same instructions as post_moe_hop's routed_moe_into_hop variant of the cluster
+kernel for each mode:
+
+* ``"row"``: the finalized bf16 row ``[R, H]``;
+* ``"identity"``: raw FC2 rows ``t * 8 + k`` with bf16 scales ``[R, 8]``, the
+  fp32 FMA chain in k order without the permutation;
+* ``"fp32"``: the FP32 accumulator ``[R, H]``, rounded to bf16; the thread
+  clears the 8 values it read (the buffer is zero before every call).
+
+Everything after step 1 is the same code as the raw-tuple variant, so the
+packet's materialization equals ``PostMoeHop.hop_routed``'s output bit
+for bit on the same routed input.
+
+``DeferredNormMaterializeKernel`` turns a deferred packet back into post_moe_hop's normalized
+output ``bf16((x * r) * gamma)`` (same instructions as post_moe_hop's step 6); it is
+the fallback when a consumer cannot take the packet.
+
+Every inline-PTX helper comes from ``ptx.py`` (owned by deferred_norm).
+"""
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Int64
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from ..post_moe_hop.kernels import ERROR_PERM_RANGE, MAX_CHUNKS
+from .ptx import (
+    FP32_SENTINEL,
+    SENTINEL_WORD,
+    _asm,
+    _v,
+    add,
+    arrive,
+    bits_f32,
+    error_word_address,
+    fault_delay,
+    flag_error,
+    flag_timeout,
+    fma,
+    inv_rms,
+    ld_v4,
+    load_state,
+    mul,
+    pack_bf16x2,
+    perm_check,
+    poison,
+    poll8,
+    round_bf16,
+    sanitize,
+    select3,
+    select_f32,
+    shfl_xor,
+    st_v4,
+    unpack_bf16x2,
+)
+
+
+@dsl_user_op
+def _st_f32_if(addr: Int64, value: Float32, enable: Int32, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [_v(addr, loc, ip), _v(value, loc, ip), _v(enable, loc, ip)],
+        "{\n.reg .pred p;\nsetp.ne.s32 p, $2, 0;\n@p st.global.f32 [$0], $1;\n}\n",
+        "l,f,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _st_b32(addr: Int64, word: Int32, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [_v(addr, loc, ip), _v(word, loc, ip)],
+        "st.global.b32 [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _ld_f32(addr: Int64, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(T.f32(), [_v(addr, loc, ip)], "ld.global.f32 $0, [$1];", "=f,l", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def _globaltimer_st(addr: Int64, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [_v(addr, loc, ip)],
+        "{\n.reg .u64 t;\nmov.u64 t, %globaltimer;\nst.global.u64 [$0], t;\n}\n",
+        "l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@cute.jit
+def _hop_stamp(
+    enabled: cutlass.Constexpr[bool], ts: Int64, sid: Int32, tid: Int32, k: cutlass.Constexpr[int]
+):
+    if cutlass.const_expr(enabled):
+        if tid == Int32(0):
+            _globaltimer_st(ts + Int64((sid * Int32(8) + Int32(k)) * Int32(8)))
+
+
+class DeferredNormHopKernel:
+    """Compile-time configuration of the deferred hop for one shape.
+
+    ``fault_delay_ns`` (tests only; the served hop keeps 0) delays about one
+    thread in four by a hashed amount below this bound before its multicast
+    publish, before its share of the clear and before the arrival (the sites
+    of post_moe_hop's fault injection), so the individual stores of a call land in a
+    scrambled order.
+
+    ``routed_input`` (``None``: the raw tuple; ``"row"``, ``"identity"`` or
+    ``"fp32"``: routed_moe_into_hop's routed input, module docstring) selects step 1; then
+    ``fc2`` is the routed buffer, ``fc2_row_bytes`` its row pitch, ``scales``
+    the identity-mode scales and ``perm`` / ``fc2_rows`` are unused.
+    """
+
+    def __init__(
+        self,
+        hidden: int,
+        top_k: int,
+        world: int,
+        max_rows: int,
+        ctas: int,
+        has_shared: bool,
+        scale_fp32: bool,
+        use_pdl: bool,
+        trigger_early: bool,
+        carveout: int = -1,
+        debug_ts: bool = False,
+        fault_delay_ns: int = 0,
+        routed_input: Optional[str] = None,
+    ):
+        assert top_k == 8 and world == 8, "the deferred_norm hop is specialized to top-8 and TP8"
+        assert routed_input in ROUTED_MODES, routed_input
+        assert hidden % (8 * 32 * ctas) == 0
+        self.hidden = hidden
+        self.top_k = top_k
+        self.world = world
+        self.max_rows = max_rows
+        self.ctas = ctas
+        self.threads = hidden // (8 * ctas)
+        self.warps = self.threads // 32
+        self.has_shared = has_shared
+        self.scale_fp32 = scale_fp32
+        self.use_pdl = use_pdl
+        self.trigger_early = trigger_early
+        # Preferred shared-memory carveout (percent; -1: driver default).
+        self.carveout = carveout
+        # Tests only: %globaltimer stamps per CTA into ``ts`` [R * CTAS, 8]
+        # (0 entry, 1 after the dependency wait, 2 published, 3 poll done,
+        # 4 exit).
+        self.debug_ts = debug_ts
+        self.fault_delay_ns = fault_delay_ns
+        self.routed_input = routed_input
+
+    @property
+    def partial_count(self) -> int:
+        """Number of fp32 square-sum partials per row."""
+        return self.ctas * self.warps
+
+    @cute.jit
+    def __call__(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        scaled: Int64,
+        residual_out: Int64,
+        sumsq: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        ts: Int64,
+        state: cute.Tensor,
+        rank: Int32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+        stream: CUstream,
+    ):
+        # ``state`` is passed as a GPU tensor so the TVM-FFI launcher can bind
+        # the caller's current stream; every other buffer is a raw address.
+        if cutlass.const_expr(self.carveout >= 0):
+            self.kernel(
+                fc2,
+                fc2_row_bytes,
+                fc2_rows,
+                scales,
+                perm,
+                shared,
+                residual,
+                gamma,
+                scaled,
+                residual_out,
+                sumsq,
+                mc,
+                uc,
+                partials,
+                ts,
+                state,
+                rank,
+                timeout_ns,
+                num_rows,
+            ).launch(
+                grid=[self.ctas, num_rows, 1],
+                block=[self.threads, 1, 1],
+                stream=stream,
+                use_pdl=self.use_pdl,
+                preferred_smem_carveout=self.carveout,
+            )
+        else:
+            self.kernel(
+                fc2,
+                fc2_row_bytes,
+                fc2_rows,
+                scales,
+                perm,
+                shared,
+                residual,
+                gamma,
+                scaled,
+                residual_out,
+                sumsq,
+                mc,
+                uc,
+                partials,
+                ts,
+                state,
+                rank,
+                timeout_ns,
+                num_rows,
+            ).launch(
+                grid=[self.ctas, num_rows, 1],
+                block=[self.threads, 1, 1],
+                stream=stream,
+                use_pdl=self.use_pdl,
+            )
+
+    @cute.kernel
+    def kernel(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        scaled: Int64,
+        residual_out: Int64,
+        sumsq: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        ts: Int64,
+        state_tensor: cute.Tensor,
+        rank: Int32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+    ):
+        hidden = self.hidden
+        world = self.world
+        threads = self.threads
+        warps = self.warps
+        ctas = self.ctas
+        buffer_bytes = self.max_rows * world * hidden * 2
+
+        tidx, _, _ = cute.arch.thread_idx()
+        cta, row, _ = cute.arch.block_idx()
+        lane = cute.arch.lane_idx()
+        warp = cute.arch.warp_idx()
+
+        state = Int64(state_tensor.iterator.toint())
+        sid = row * Int32(ctas) + cta
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 0)
+
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+            if cutlass.const_expr(self.trigger_early):
+                cute.arch.griddepcontrol_launch_dependents()
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 1)
+
+        # State words of the previous call (complete at the kernel boundary).
+        epoch, prev_rows, err_words = load_state(state)
+        cur = epoch % Int32(3)
+        dirty = (epoch + Int32(2)) % Int32(3)
+        # Bits of the previous call's buffer: when set, this call publishes
+        # NaN and carries them into its own error word.
+        sticky = select3(err_words, dirty)
+        err_cur = error_word_address(state, cur)
+
+        col = (cta * Int32(threads) + tidx) * Int32(8)
+        row_off = Int64(row) * Int64(hidden * 2) + Int64(col) * Int64(2)
+        res_w = ld_v4(residual + row_off)
+        gamma_w = ld_v4(gamma + Int64(col) * Int64(2))
+        if cutlass.const_expr(self.has_shared):
+            sh_w = ld_v4(shared + row_off)
+
+        # 1. Finalize (fp32 FMA chain in k order), bf16, + shared, sanitize.
+        if cutlass.const_expr(self.routed_input == "row"):
+            # routed_moe_into_hop: routed_moe already finalized this rank's routed output.
+            r_w = ld_v4(fc2 + Int64(row) * fc2_row_bytes + Int64(col) * Int64(2))
+            local = [r_w[j] for j in range(4)]
+            out_of_range = Int32(0)
+        elif cutlass.const_expr(self.routed_input == "fp32"):
+            # routed_moe_into_hop: FP32 accumulator of the top-k contributions; round, then
+            # clear it for the routed kernel of the next call.
+            a_addr = fc2 + Int64(row) * fc2_row_bytes + Int64(col) * Int64(4)
+            a_w = ld_v4(a_addr) + ld_v4(a_addr + Int64(16))
+            a_f = [bits_f32(w) for w in a_w]
+            local = [pack_bf16x2(a_f[2 * j], a_f[2 * j + 1]) for j in range(4)]
+            zero = Int32(0)
+            st_v4(a_addr, zero, zero, zero, zero)
+            st_v4(a_addr + Int64(16), zero, zero, zero, zero)
+            out_of_range = Int32(0)
+        elif cutlass.const_expr(self.routed_input == "identity"):
+            # routed_moe_into_hop: raw rows t * 8 + k, every slot valid; the finalize without
+            # the permutation round trip.
+            s_w = ld_v4(scales + Int64(row) * Int64(16))
+            weights = []
+            for w in s_w:
+                lo, hi = unpack_bf16x2(w)
+                weights += [lo, hi]
+            acc = [Float32(0.0)] * 8
+            for k in cutlass.range_constexpr(8):
+                x_w = ld_v4(
+                    fc2 + (Int64(row) * Int64(8) + Int64(k)) * fc2_row_bytes + Int64(col) * Int64(2)
+                )
+                for j in cutlass.range_constexpr(4):
+                    x_lo, x_hi = unpack_bf16x2(x_w[j])
+                    acc[2 * j] = fma(weights[k], x_lo, acc[2 * j])
+                    acc[2 * j + 1] = fma(weights[k], x_hi, acc[2 * j + 1])
+            local = [pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)]
+            out_of_range = Int32(0)
+        else:
+            perm_lo = ld_v4(perm + Int64(row) * Int64(32))
+            perm_hi = ld_v4(perm + Int64(row) * Int64(32) + Int64(16))
+            perm_k = perm_lo + perm_hi
+            if cutlass.const_expr(self.scale_fp32):
+                s_lo = ld_v4(scales + Int64(row) * Int64(32))
+                s_hi = ld_v4(scales + Int64(row) * Int64(32) + Int64(16))
+                weights = [bits_f32(w) for w in s_lo + s_hi]
+            else:
+                s_w = ld_v4(scales + Int64(row) * Int64(16))
+                weights = []
+                for w in s_w:
+                    lo, hi = unpack_bf16x2(w)
+                    weights += [lo, hi]
+            acc = [Float32(0.0)] * 8
+            out_of_range = Int32(0)
+            for k in cutlass.range_constexpr(8):
+                valid, safe, oob = perm_check(perm_k[k], fc2_rows)
+                out_of_range = out_of_range + oob
+                x_w = ld_v4(fc2 + Int64(safe) * fc2_row_bytes + Int64(col) * Int64(2))
+                for j in cutlass.range_constexpr(4):
+                    x_lo, x_hi = unpack_bf16x2(x_w[j])
+                    acc[2 * j] = select_f32(valid, fma(weights[k], x_lo, acc[2 * j]), acc[2 * j])
+                    acc[2 * j + 1] = select_f32(
+                        valid, fma(weights[k], x_hi, acc[2 * j + 1]), acc[2 * j + 1]
+                    )
+            local = [pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)]
+        if cutlass.const_expr(self.has_shared):
+            summed = []
+            for j in cutlass.range_constexpr(4):
+                r_lo, r_hi = unpack_bf16x2(local[j])
+                s_lo_f, s_hi_f = unpack_bf16x2(sh_w[j])
+                summed.append(pack_bf16x2(add(s_lo_f, r_lo), add(s_hi_f, r_hi)))
+            local = summed
+        # After an earlier error on this rank, publish NaN until the reset.
+        local = [poison(sanitize(w), sticky) for w in local]
+
+        # 2. Publish through multicast into slot [row][rank].
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4),
+                self.fault_delay_ns,
+                0,
+            )
+        slot = (
+            Int64(cur) * Int64(buffer_bytes)
+            + Int64(row) * Int64(world * hidden * 2)
+            + Int64(rank) * Int64(hidden * 2)
+            + Int64(col) * Int64(2)
+        )
+        st_v4(mc + slot, local[0], local[1], local[2], local[3])
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 2)
+        flag_error(err_cur, out_of_range, Int32(ERROR_PERM_RANGE))
+        flag_error(err_cur, sticky, sticky)
+
+        # 3. Clear the previous call's Lamport rows and partial slice.
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4) + Int32(1),
+                self.fault_delay_ns,
+                1,
+            )
+        total_threads = num_rows * Int32(ctas * threads)
+        gtid = (row * Int32(ctas) + cta) * Int32(threads) + tidx
+        clear_vecs = prev_rows * Int32(world * hidden // 8)
+        sentinel = Int32(SENTINEL_WORD)
+        dirty_base = uc + Int64(dirty) * Int64(buffer_bytes)
+        for v in cutlass.range(gtid, clear_vecs, total_threads):
+            st_v4(dirty_base + Int64(v) * Int64(16), sentinel, sentinel, sentinel, sentinel)
+        part_vecs = prev_rows * Int32(MAX_CHUNKS // 4)
+        part_base = partials + Int64(dirty) * Int64(self.max_rows * MAX_CHUNKS * 4)
+        fp32_sentinel = Int32(FP32_SENTINEL)
+        for v in cutlass.range(gtid, part_vecs, total_threads):
+            st_v4(
+                part_base + Int64(v) * Int64(16),
+                fp32_sentinel,
+                fp32_sentinel,
+                fp32_sentinel,
+                fp32_sentinel,
+            )
+
+        # 4. Per-thread Lamport poll (timer deadline), rank-ordered fp32 sum.
+        poll_base = (
+            uc
+            + Int64(cur) * Int64(buffer_bytes)
+            + Int64(row) * Int64(world * hidden * 2)
+            + Int64(col) * Int64(2)
+        )
+        polled = poll8(poll_base, Int64(hidden * 2), timeout_ns)
+        timed_out = polled[32]
+        flag_timeout(err_cur, state, timed_out, polled[33])
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 3)
+        reduced = [Float32(0.0)] * 8
+        for r in cutlass.range_constexpr(8):
+            for j in cutlass.range_constexpr(4):
+                v_lo, v_hi = unpack_bf16x2(polled[4 * r + j])
+                reduced[2 * j] = add(reduced[2 * j], v_lo)
+                reduced[2 * j + 1] = add(reduced[2 * j + 1], v_hi)
+        if cutlass.const_expr(self.use_pdl and not self.trigger_early):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        # 5. Residual added once, after the reduction; NaN if this poll timed out.
+        new_res = []
+        x = []
+        for j in cutlass.range_constexpr(4):
+            rb = pack_bf16x2(reduced[2 * j], reduced[2 * j + 1])
+            a_lo, a_hi = unpack_bf16x2(rb)
+            r_lo, r_hi = unpack_bf16x2(res_w[j])
+            word = poison(pack_bf16x2(add(a_lo, r_lo), add(a_hi, r_hi)), timed_out)
+            new_res.append(word)
+            n_lo, n_hi = unpack_bf16x2(word)
+            x += [n_lo, n_hi]
+        st_v4(residual_out + row_off, new_res[0], new_res[1], new_res[2], new_res[3])
+
+        # 6. Deferred RMSNorm: post_moe_hop's per-thread square sum and warp butterfly,
+        #    one fp32 partial per warp; gain-folded activation.
+        thread_sum = Float32(0.0)
+        for i in cutlass.range_constexpr(8):
+            thread_sum = add(thread_sum, round_bf16(mul(x[i], x[i])))
+        for mask in (16, 8, 4, 2, 1):
+            thread_sum = add(thread_sum, shfl_xor(thread_sum, mask))
+        slot_idx = Int64(row) * Int64(ctas * warps) + Int64(cta * Int32(warps) + warp)
+        _st_f32_if(sumsq + slot_idx * Int64(4), thread_sum, Int32(lane == Int32(0)))
+        # a = bf16(x * gamma) in the MMA B-fragment layout of packet.py: this
+        # thread's columns [col, col + 8) are half (col / 8) % 2 of 16-column
+        # block col / 16, pair j of lane 4 * row + j.
+        blk = col // Int32(16)
+        half = (col // Int32(8)) % Int32(2)
+        for j in cutlass.range_constexpr(4):
+            g_lo, g_hi = unpack_bf16x2(gamma_w[j])
+            word = pack_bf16x2(mul(x[2 * j], g_lo), mul(x[2 * j + 1], g_hi))
+            lane_slot = blk * Int32(32) + row * Int32(4) + Int32(j)
+            _st_b32(scaled + Int64(lane_slot * Int32(2) + half) * Int64(4), word)
+
+        # 7. Arrival; the last CTA publishes the next epoch.
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4) + Int32(2),
+                self.fault_delay_ns,
+                2,
+            )
+        cute.arch.sync_threads()
+        arrive(
+            state,
+            Int32(tidx == Int32(0)),
+            num_rows * Int32(ctas),
+            num_rows,
+            (epoch + Int32(1)) % Int32(3),
+        )
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 4)
+
+
+class DeferredNormMaterializeKernel:
+    """``h = bf16((x * r) * gamma)`` from a deferred packet (post_moe_hop's step 6 arithmetic)."""
+
+    def __init__(self, hidden: int, ctas: int, partials_per_row: int, use_pdl: bool):
+        assert hidden % (8 * 32 * ctas) == 0
+        self.hidden = hidden
+        self.ctas = ctas
+        self.threads = hidden // (8 * ctas)
+        self.partials_per_row = partials_per_row
+        self.use_pdl = use_pdl
+
+    @cute.jit
+    def __call__(
+        self,
+        residual_out: Int64,
+        gamma: Int64,
+        sumsq: Int64,
+        out: Int64,
+        anchor: cute.Tensor,
+        eps: Float32,
+        num_rows: Int32,
+        stream: CUstream,
+    ):
+        self.kernel(residual_out, gamma, sumsq, out, anchor, eps).launch(
+            grid=[self.ctas, num_rows, 1],
+            block=[self.threads, 1, 1],
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        residual_out: Int64,
+        gamma: Int64,
+        sumsq: Int64,
+        out: Int64,
+        anchor: cute.Tensor,
+        eps: Float32,
+    ):
+        hidden = self.hidden
+        tidx, _, _ = cute.arch.thread_idx()
+        cta, row, _ = cute.arch.block_idx()
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+        col = (cta * Int32(self.threads) + tidx) * Int32(8)
+        row_off = Int64(row) * Int64(hidden * 2) + Int64(col) * Int64(2)
+        x_w = ld_v4(residual_out + row_off)
+        gamma_w = ld_v4(gamma + Int64(col) * Int64(2))
+        square_sum = Float32(0.0)
+        base = sumsq + Int64(row) * Int64(self.partials_per_row * 4)
+        for i in cutlass.range_constexpr(self.partials_per_row):
+            square_sum = add(square_sum, _ld_f32(base + Int64(4 * i)))
+        r = inv_rms(square_sum, Float32(float(hidden)), eps)
+        out_w = []
+        for j in cutlass.range_constexpr(4):
+            x_lo, x_hi = unpack_bf16x2(x_w[j])
+            g_lo, g_hi = unpack_bf16x2(gamma_w[j])
+            out_w.append(pack_bf16x2(mul(mul(x_lo, r), g_lo), mul(mul(x_hi, r), g_hi)))
+        st_v4(out + row_off, out_w[0], out_w[1], out_w[2], out_w[3])
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_launch_dependents()
+
+
+_COMPILED = {}
+# Step-1 variants: the raw tuple (None) and routed_moe_into_hop's routed-input modes.
+ROUTED_MODES = (None, "row", "identity", "fp32")
+
+
+def compiled_hop(
+    hidden: int,
+    top_k: int,
+    world: int,
+    max_rows: int,
+    ctas: int,
+    has_shared: bool,
+    scale_fp32: bool,
+    use_pdl: bool,
+    trigger_early: bool,
+    state_words: int,
+    carveout: int = -1,
+    debug_ts: bool = False,
+    fault_delay_ns: int = 0,
+    routed_input: Optional[str] = None,
+):
+    """Return the compiled deferred hop for one shape (compiles on first use, no GPU work).
+
+    ``routed_input`` selects a routed_moe_into_hop routed-input variant (``DeferredNormHopKernel``);
+    the raw-tuple variants keep their key.
+    """
+    key = (
+        "hop",
+        hidden,
+        top_k,
+        world,
+        max_rows,
+        ctas,
+        has_shared,
+        scale_fp32,
+        use_pdl,
+        trigger_early,
+        state_words,
+        carveout,
+        debug_ts,
+        fault_delay_ns,
+    ) + ((routed_input,) if routed_input else ())
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = DeferredNormHopKernel(
+            hidden,
+            top_k,
+            world,
+            max_rows,
+            ctas,
+            has_shared,
+            scale_fp32,
+            use_pdl,
+            trigger_early,
+            carveout,
+            debug_ts,
+            fault_delay_ns,
+            routed_input,
+        )
+        zero64, zero32 = Int64(0), Int32(0)
+        state_fake = cute_runtime.make_fake_compact_tensor(
+            Int32, (state_words,), stride_order=(0,), assumed_align=16
+        )
+        compiled = cute.compile(
+            kernel,
+            zero64,
+            zero64,
+            zero32,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            state_fake,
+            zero32,
+            zero64,
+            zero32,
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+        _COMPILED[key] = compiled
+    return compiled
+
+
+def compiled_materialize(hidden: int, ctas: int, partials_per_row: int, use_pdl: bool):
+    key = ("materialize", hidden, ctas, partials_per_row, use_pdl)
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = DeferredNormMaterializeKernel(hidden, ctas, partials_per_row, use_pdl)
+        zero64 = Int64(0)
+        anchor = cute_runtime.make_fake_compact_tensor(
+            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
+        )
+        compiled = cute.compile(
+            kernel,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            anchor,
+            Float32(0.0),
+            Int32(0),
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+        _COMPILED[key] = compiled
+    return compiled
+
+
+def hop_kernel_supported(hidden: int, top_k: int, world: int, ctas: int) -> Optional[str]:
+    """Return why the deferred hop cannot serve a shape, or ``None``."""
+    if top_k != 8 or world != 8:
+        return "deferred hop specialized to top-8 / TP8"
+    if ctas < 1 or ctas > 8 or hidden % (8 * 32 * ctas):
+        return f"hidden {hidden} does not split into {ctas} CTAs of whole warps"
+    if hidden // (8 * ctas) > 1024:
+        return "too many threads per CTA"
+    return None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/packet.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/packet.py
new file mode 100644
index 000000000000..2d9e594b9ef9
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/packet.py
@@ -0,0 +1,105 @@
+# 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.
+"""The hidden-state packet a deferred_norm deferred hop hands to the next layer.
+
+``DeferredNormHidden`` stands in for the normalized hidden state ``h =
+bf16((x * r) * gamma)`` that the post_moe_hop would return. It is deliberately not
+a tensor: a consumer that does not know the packet fails at its first tensor
+operation instead of silently computing with the un-normalized activation.
+deferred_norm-aware consumers take ``a_frag`` (``a = bf16(x * gamma)`` in the MMA
+operand layout below) and the square-sum partials, and apply ``r`` after
+their projection; every other consumer calls :meth:`materialize`, which
+returns post_moe_hop's normalized output bit for bit.
+
+``a_frag`` layout (``FRAG_ROWS = 8`` token rows, ``K`` columns): for every
+16-column block ``s`` and lane ``l = 4 g + t`` of an ``mma.m16n8k16`` B
+fragment (token ``g``, column pair ``t``), four BF16 values ``a[g, 16 s + 2 t
++ (0, 1)]`` and ``a[g, 16 s + 8 + 2 t + (0, 1)]``, i.e. flat index ``((s *
+32 + l) * 2 + half) * 2 + pair``. A warp then reads one 256-byte run per K
+step. Rows ``>= R`` are not written.
+"""
+
+from dataclasses import dataclass
+from typing import Callable
+
+import torch
+
+FRAG_ROWS = 8
+
+
+def frag_elems(k: int) -> int:
+    return (k // 16) * 32 * 4
+
+
+def frag_from_rows(a: torch.Tensor) -> torch.Tensor:
+    """``a_frag`` of a ``[R <= 8, K]`` tensor (rows ``>= R`` zero)."""
+    rows, k = a.shape
+    pad = a.new_zeros(FRAG_ROWS, k)
+    pad[:rows] = a
+    return pad.view(FRAG_ROWS, k // 16, 2, 4, 2).permute(1, 0, 3, 2, 4).reshape(-1).contiguous()
+
+
+def rows_from_frag(a_frag: torch.Tensor, rows: int, k: int) -> torch.Tensor:
+    """Inverse of :func:`frag_from_rows` for the first ``rows`` rows."""
+    t = a_frag.view(k // 16, FRAG_ROWS, 4, 2, 2).permute(1, 0, 3, 2, 4).reshape(FRAG_ROWS, k)
+    return t[:rows].contiguous()
+
+
+@dataclass(frozen=True)
+class DeferredNormHidden:
+    """Deferred RMSNorm output of a post-MoE hop.
+
+    Attributes:
+        a_frag: bf16 ``[frag_elems(H)]``, ``bf16(x * gamma)`` in the MMA operand
+            layout (module docstring).
+        sumsq: fp32 ``[R, P]``, per-warp sums of bf16-rounded squares of ``x``
+            in the order the RMSNorm sums them (``P`` = 24 for H = 6144).
+        residual: bf16 ``[R, H]``, ``x`` (the hop's residual output); also the
+            tensor view used for shapes and output allocation.
+        gamma: bf16 ``[H]``, the next layer's input RMSNorm weight.
+        eps: RMSNorm epsilon.
+        materializer: launches the kernel behind :meth:`materialize`.
+    """
+
+    a_frag: torch.Tensor
+    sumsq: torch.Tensor
+    residual: torch.Tensor
+    gamma: torch.Tensor
+    eps: float
+    materializer: Callable[["DeferredNormHidden"], torch.Tensor]
+
+    @property
+    def shape(self) -> torch.Size:
+        return self.residual.shape
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.residual.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.residual.device
+
+    @property
+    def num_rows(self) -> int:
+        return self.residual.shape[0]
+
+    def materialize(self) -> torch.Tensor:
+        """Return the normalized hidden state ``bf16((x * r) * gamma)``."""
+        return self.materializer(self)
+
+
+def is_deferred(hidden_states) -> bool:
+    return isinstance(hidden_states, DeferredNormHidden)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/proj_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/proj_kernel.py
new file mode 100644
index 000000000000..479ef8d89dc2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/proj_kernel.py
@@ -0,0 +1,889 @@
+# 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.
+"""CuTe DSL kernel of the deferred_norm next-layer projection with staged weights.
+
+Computes, for ``R <= 8`` token rows,
+
+    y[t, j] = bf16(r[t] * sum_k a[t, k] * W[j, k])
+
+over the concatenation of up to two row-major BF16 weight tensors with the
+same ``K`` (``W0 [N0, K]`` = the fused ``[q_a | kv_a | k_pe]`` projection,
+``W1 [N1, K]`` = the indexer ``wk | weights_proj`` rows on full-indexer
+layers), writing ``out0 [R, N0]`` and ``out1 [R, N1]``. In deferred mode
+``a`` is the gain-folded activation ``bf16(x * gamma)`` of the deferred_norm hop and
+``r[t] = rsqrt.approx(div.rn(sum_i sumsq[t, i], H) + eps)`` with the square
+sum taken in the hop's partial order (bit-identical to the post_moe_hop's
+``rsqrt`` operand); in native-input mode ``a`` is an ordinary normalized
+hidden state and ``r = 1``.
+
+Schedule (one CTA per SM, ``NT`` threads):
+
+* CTA ``c`` owns a contiguous range of 17-19 output rows. Before
+  ``griddepcontrol.wait`` (the weights do not depend on the preceding hop)
+  one thread issues ``cp.async.bulk`` copies of up to ``SMEM_ROWS`` of its
+  rows into shared memory (row stride padded by 16 bytes), one mbarrier per
+  K chunk, and prefetches every row it cannot stage into L2;
+* after the wait every thread loads its activation fragments (``B`` of
+  ``mma.m16n8k16``, the token rows) from global memory, and warp ``w`` runs
+  the K steps ``s = w + NW * m`` of two ``m16n8`` tiles (rows 0-15 via
+  ``ldmatrix.x4``, rows 16-23 via ``ldmatrix.x2``; rows the CTA does not own
+  read a zero chunk);
+* the ``NW`` per-warp partial sums are reduced in warp order through shared
+  memory, scaled by ``r`` and rounded once to BF16;
+* a CTA that owns more rows than it can stage processes the extra row after
+  the main loop from a bulk copy of the (L2-prefetched) row.
+"""
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.utils.smem_allocator import SmemAllocator
+
+from .ptx import (
+    add,
+    bf16_bits,
+    bits_f32,
+    cp_async16,
+    cp_async_commit,
+    cp_async_wait_group,
+    inv_rms,
+    ldsm_x4,
+    mma_bf16,
+    mul,
+    saddr,
+    st_global_b16,
+    st_shared_v4,
+)
+
+MAX_TOKENS = 8
+SMEM_LIMIT = 227 * 1024
+
+
+def _asm(res, args, text, cons, side_effects=True, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=side_effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _bulk_g2s(saddr: Int32, gaddr: Int64, nbytes: Int32, mbar: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [saddr.ir_value(), gaddr.ir_value(), nbytes.ir_value(), mbar.ir_value()],
+        "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];",
+        "r,l,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _prefetch_l2(gaddr: Int64, nbytes: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [gaddr.ir_value(), nbytes.ir_value()],
+        "cp.async.bulk.prefetch.L2.global [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _mbar_init(mbar: Int32, count: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [mbar.ir_value(), count.ir_value()],
+        "mbarrier.init.shared::cta.b64 [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _mbar_fence_init(*, loc=None, ip=None):
+    _asm(
+        None,
+        [],
+        "fence.mbarrier_init.release.cluster;\n\tfence.proxy.async.shared::cta;",
+        "",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _mbar_expect_tx(mbar: Int32, nbytes: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [mbar.ir_value(), nbytes.ir_value()],
+        "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _mbar_wait(mbar: Int32, phase: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [mbar.ir_value(), phase.ir_value()],
+        "{\n\t.reg .pred p;\n\tNORM_PROJ_MBAR_WAIT:\n\t"
+        "mbarrier.try_wait.parity.shared::cta.b64 p, [$0], $1;\n\t"
+        "@!p bra NORM_PROJ_MBAR_WAIT;\n\t}",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _ldsm_x2(saddr: Int32, *, loc=None, ip=None):
+    out = _asm(
+        llvm.StructType.get_literal([T.i32()] * 2),
+        [saddr.ir_value()],
+        "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {$0, $1}, [$2];",
+        "=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(2))
+
+
+@dsl_user_op
+def _ld_nc_b32(gaddr: Int64, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(T.i32(), [gaddr.ir_value()], "ld.global.nc.b32 $0, [$1];", "=r,l", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def _st_shared_f32(saddr: Int32, v: Float32, *, loc=None, ip=None):
+    _asm(None, [saddr.ir_value(), v.ir_value()], "st.shared.f32 [$0], $1;", "r,f", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _ld_shared_f32(saddr: Int32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(T.f32(), [saddr.ir_value()], "ld.shared.f32 $0, [$1];", "=f,r", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def _sel_i32(pred: Int32, a: Int32, b: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [pred.ir_value(), a.ir_value(), b.ir_value()],
+            "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $1, 0;\n\tselp.b32 $0, $2, $3, p;\n\t}",
+            "=r,r,r,r",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _sel_i64(pred: Int32, a: Int64, b: Int64, *, loc=None, ip=None) -> Int64:
+    return Int64(
+        _asm(
+            T.i64(),
+            [pred.ir_value(), a.ir_value(), b.ir_value()],
+            "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $1, 0;\n\tselp.b64 $0, $2, $3, p;\n\t}",
+            "=l,r,l,l",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _min_i32(a: Int32, b: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [a.ir_value(), b.ir_value()],
+            "min.s32 $0, $1, $2;",
+            "=r,r,r",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _fence_proxy_async(*, loc=None, ip=None):
+    _asm(None, [], "fence.proxy.async.shared::cta;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _ld_b32(gaddr: Int64, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(T.i32(), [gaddr.ir_value()], "ld.global.b32 $0, [$1];", "=r,l", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def _ld_v2(gaddr: Int64, *, loc=None, ip=None):
+    out = _asm(
+        llvm.StructType.get_literal([T.i32()] * 2),
+        [gaddr.ir_value()],
+        "ld.global.v2.b32 {$0, $1}, [$2];",
+        "=r,=r,l",
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(2))
+
+
+@dsl_user_op
+def _ld_v2_if(gaddr: Int64, pred: Int32, *, loc=None, ip=None):
+    """8-byte load if ``pred`` else zeros."""
+    out = _asm(
+        llvm.StructType.get_literal([T.i32()] * 2),
+        [gaddr.ir_value(), pred.ir_value()],
+        "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $3, 0;\n\tmov.b32 $0, 0;\n\t"
+        "mov.b32 $1, 0;\n\t@p ld.global.v2.b32 {$0, $1}, [$2];\n\t}",
+        "=r,=r,l,r",
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(2))
+
+
+@dsl_user_op
+def _shfl_idx_f32(v: Float32, src: Int32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [v.ir_value(), src.ir_value()],
+            "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;",
+            "=f,f,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _globaltimer(*, loc=None, ip=None) -> Int64:
+    return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def _st_global_b64(gaddr: Int64, v: Int64, *, loc=None, ip=None):
+    _asm(None, [gaddr.ir_value(), v.ir_value()], "st.global.b64 [$0], $1;", "l,l", loc=loc, ip=ip)
+
+
+@cute.jit
+def _row_src(
+    w0: Int64,
+    w1: Int64,
+    gr: Int32,
+    K: cutlass.Constexpr[int],
+    N0: cutlass.Constexpr[int],
+    N1: cutlass.Constexpr[int],
+) -> Int64:
+    """Global address of row ``gr`` of the concatenation ``[W0; W1]``."""
+    k2 = Int64(K * 2)
+    src = w0 + Int64(gr) * k2
+    if cutlass.const_expr(N1 > 0):
+        src = _sel_i64(Int32(gr >= Int32(N0)), w1 + Int64(gr - Int32(N0)) * k2, src)
+    return src
+
+
+@cute.jit
+def _store_out(
+    out0: Int64,
+    out1: Int64,
+    gr: Int32,
+    tok: Int32,
+    value: Float32,
+    N0: cutlass.Constexpr[int],
+    N1: cutlass.Constexpr[int],
+):
+    """BF16 store of ``value`` into ``out0[tok, gr]`` or ``out1[tok, gr - N0]``."""
+    if gr < Int32(N0):
+        st_global_b16(out0 + Int64(tok * Int32(N0) + gr) * Int64(2), bf16_bits(value))
+    else:
+        if cutlass.const_expr(N1 > 0):
+            st_global_b16(
+                out1 + Int64(tok * Int32(N1) + gr - Int32(N0)) * Int64(2), bf16_bits(value)
+            )
+
+
+@cute.jit
+def _stamp(
+    enabled: cutlass.Constexpr[bool], ts: Int64, bid: Int32, tid: Int32, k: cutlass.Constexpr[int]
+):
+    """Tests only: ``%globaltimer`` of CTA ``bid`` at point ``k``."""
+    if cutlass.const_expr(enabled):
+        if tid == Int32(0):
+            _st_global_b64(ts + Int64((bid * Int32(8) + Int32(k)) * Int32(8)), _globaltimer())
+
+
+class DeferredNormProjKernel:
+    """Compile-time configuration of the staged projection (see the module docstring)."""
+
+    def __init__(
+        self,
+        k: int,
+        n0: int,
+        n1: int,
+        ctas: int,
+        threads: int = 512,
+        smem_rows: int = 18,
+        nchunk: int = 4,
+        deferred: bool = True,
+        partials_per_row: int = 24,
+        use_pdl: bool = True,
+        trigger: str = "start",
+        stage_before_wait: bool = True,
+        l2_prefetch: bool = True,
+        dist: int = 0,
+        copy_threads: int = 1,
+        copy_method: str = "bulk",
+        wait_staged_first: bool = False,
+        b_layout: str = "frag",
+        prefetch_next: bool = False,
+        debug_ts: bool = False,
+    ):
+        assert k % 16 == 0
+        self.k = k
+        self.n0 = n0
+        self.n1 = n1
+        self.ctas = ctas
+        self.nt = threads
+        self.nw = threads // 32
+        self.steps = k // 16
+        assert self.steps % self.nw == 0, "K steps must split evenly over the warps"
+        self.rounds = self.steps // self.nw
+        n = n0 + n1
+        self.base_rows = n // ctas
+        self.rem_rows = n % ctas
+        self.max_rows = self.base_rows + (1 if self.rem_rows else 0)
+        assert 16 < self.base_rows and self.max_rows <= 24, "17..24 rows per CTA supported"
+        self.smem_rows = min(smem_rows, self.max_rows)
+        self.overflow = self.max_rows - self.smem_rows
+        assert self.overflow in (0, 1), "at most one row per CTA beyond the staged rows"
+        assert self.smem_rows >= 17
+        self.stride = k * 2 + 16
+        self.nchunk = nchunk
+        assert k % nchunk == 0 and (k // nchunk) % (16 * self.nw) == 0
+        self.chunk_k = k // nchunk
+        self.rounds_per_chunk = self.chunk_k // (16 * self.nw)
+        self.deferred = deferred
+        self.partials_per_row = partials_per_row
+        # Each lane of the r warp loads a quarter of its token's partials.
+        assert partials_per_row % 8 == 0
+        self.use_pdl = use_pdl
+        assert trigger in ("start", "end", "none")
+        self.trigger = trigger
+        self.stage_before_wait = stage_before_wait
+        self.l2_prefetch = l2_prefetch
+        self.debug_ts = debug_ts
+        self.dist = dist
+        self.copy_threads = copy_threads
+        assert copy_method in ("bulk", "cpasync")
+        self.copy_method = copy_method
+        # "frag": the activation is in the MMA B-fragment layout of
+        # packet.py (one 8-byte load per lane and K step); "plain": a
+        # row-major [R, K] tensor (two 4-byte loads per lane and K step).
+        assert b_layout in ("frag", "plain")
+        self.b_layout = b_layout
+        self.prefetch_next = prefetch_next
+        # cp.async staging completes per thread: every thread waits for it
+        # (then a CTA barrier) before the first MMA, so it is always "first".
+        self.wait_staged_first = wait_staged_first or copy_method == "cpasync"
+        self.red_bytes = self.nw * 24 * MAX_TOKENS * 4
+        assert self.red_bytes <= 2 * self.stride
+        self.smem_bytes = self.smem_rows * self.stride + 1024
+        assert self.smem_bytes <= SMEM_LIMIT, f"{self.smem_bytes} B of shared memory"
+
+    @cute.jit
+    def __call__(
+        self,
+        act: Int64,
+        sumsq: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        pf0: Int64,
+        pf0_bytes: Int64,
+        pf1: Int64,
+        pf1_bytes: Int64,
+        pf2: Int64,
+        pf2_bytes: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        eps: Float32,
+        stream: CUstream,
+    ):
+        self.kernel(
+            act,
+            sumsq,
+            w0,
+            w1,
+            out0,
+            out1,
+            ts,
+            pf0,
+            pf0_bytes,
+            pf1,
+            pf1_bytes,
+            pf2,
+            pf2_bytes,
+            anchor,
+            rows,
+            eps,
+        ).launch(grid=[self.ctas, 1, 1], block=[self.nt, 1, 1], stream=stream, use_pdl=self.use_pdl)
+
+    @cute.jit
+    def _issue_copies(
+        self, sW: Int32, sMbar: Int32, w0: Int64, w1: Int64, row0: Int32, nsm: Int32, tid: Int32
+    ):
+        """Bulk copies of the staged rows, one mbarrier per K chunk. With
+        ``copy_threads > 1`` thread 0 posts every chunk's byte count first and
+        thread ``i < nsm`` then copies row ``i`` (transaction counts may
+        complete in any order)."""
+        chunk_bytes = Int32(self.chunk_k * 2)
+        if cutlass.const_expr(self.copy_method == "cpasync"):
+            pieces = self.k * 2 // 16
+            for i in cutlass.range_constexpr((self.smem_rows * pieces + self.nt - 1) // self.nt):
+                c = tid + Int32(self.nt * i)
+                r = c // Int32(pieces)
+                off = c % Int32(pieces)
+                if r < nsm:
+                    cp_async16(
+                        sW + r * Int32(self.stride) + off * Int32(16),
+                        _row_src(w0, w1, row0 + r, self.k, self.n0, self.n1)
+                        + Int64(off * Int32(16)),
+                    )
+            cp_async_commit()
+        elif cutlass.const_expr(self.copy_threads == 1):
+            if tid == Int32(0):
+                for j in cutlass.range_constexpr(self.nchunk):
+                    _mbar_expect_tx(sMbar + Int32(8 * j), nsm * chunk_bytes)
+                    for i in cutlass.range(nsm):
+                        _bulk_g2s(
+                            sW + i * Int32(self.stride) + Int32(j * self.chunk_k * 2),
+                            _row_src(w0, w1, row0 + i, self.k, self.n0, self.n1)
+                            + Int64(j * self.chunk_k * 2),
+                            chunk_bytes,
+                            sMbar + Int32(8 * j),
+                        )
+        else:
+            if tid == Int32(0):
+                for j in cutlass.range_constexpr(self.nchunk):
+                    _mbar_expect_tx(sMbar + Int32(8 * j), nsm * chunk_bytes)
+            if tid < nsm:
+                src = _row_src(w0, w1, row0 + tid, self.k, self.n0, self.n1)
+                for j in cutlass.range_constexpr(self.nchunk):
+                    _bulk_g2s(
+                        sW + tid * Int32(self.stride) + Int32(j * self.chunk_k * 2),
+                        src + Int64(j * self.chunk_k * 2),
+                        chunk_bytes,
+                        sMbar + Int32(8 * j),
+                    )
+
+    @cute.jit
+    def _prefetch_slice(self, base: Int64, nbytes: Int64, bid: Int32):
+        """``cp.async.bulk.prefetch.L2`` of slice ``bid`` of ``[base, base +
+        nbytes)`` in 16 KB pieces (slices are 16-byte multiples)."""
+        slice_bytes = ((nbytes + Int64(self.ctas * 16 - 1)) // Int64(self.ctas * 16)) * Int64(16)
+        start = Int64(bid) * slice_bytes
+        end = start + slice_bytes
+        if end > nbytes:
+            end = nbytes
+        piece = Int64(16384)
+        off = start
+        for _ in cutlass.range_constexpr(4):
+            if off < end:
+                n = end - off
+                if n > piece:
+                    n = piece
+                _prefetch_l2(base + off, Int32(n))
+                off = off + n
+
+    @cute.jit
+    def _wait_staged(self, sMbar: Int32):
+        """Every thread waits until all staged rows are in shared memory."""
+        if cutlass.const_expr(self.copy_method == "cpasync"):
+            cp_async_wait_group(0)
+            cute.arch.sync_threads()
+        else:
+            for j in cutlass.range_constexpr(self.nchunk):
+                _mbar_wait(sMbar + Int32(8 * j), Int32(0))
+
+    @cute.kernel
+    def kernel(
+        self,
+        act: Int64,
+        sumsq: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        pf0: Int64,
+        pf0_bytes: Int64,
+        pf1: Int64,
+        pf1_bytes: Int64,
+        pf2: Int64,
+        pf2_bytes: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        eps: Float32,
+    ):
+        """``debug_ts`` stamps per CTA (int64 ``[ctas, 8]``): 0 entry, 7 copies
+        issued (before the wait when staging early), 6 every staged row
+        landed (only with ``wait_staged_first``: all threads wait for it
+        before the dependency wait), 1 after the dependency wait, 2
+        activation fragments loaded, 3 main loop done, 4 rows stored, 5
+        exit."""
+        k = self.k
+        nw = self.nw
+        stride = self.stride
+        tid, _, _ = cute.arch.thread_idx()
+        bid, _, _ = cute.arch.block_idx()
+        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        lane = cute.arch.lane_idx()
+        _stamp(self.debug_ts, ts, bid, tid, 0)
+
+        smem = SmemAllocator()
+        sW = saddr(smem.allocate(self.smem_rows * stride, byte_alignment=128))
+        sMisc = saddr(smem.allocate(512, byte_alignment=128))
+        # sMisc: [0, 32) zero chunk, [64, 64 + 8 * (nchunk + 1)) mbarriers,
+        # [256, 288) r per token.
+        sZero = sMisc
+        sMbar = sMisc + Int32(64)
+        sR = sMisc + Int32(256)
+
+        # Rows of this CTA.
+        base = Int32(self.base_rows)
+        rem = Int32(self.rem_rows)
+        extra = Int32(bid < rem)
+        nrows = base + extra
+        row0 = bid * base + _min_i32(bid, rem)
+        nsm = _min_i32(nrows, Int32(self.smem_rows))
+
+        if tid < Int32(2):
+            st_shared_v4(sZero + tid * Int32(16), Int32(0), Int32(0), Int32(0), Int32(0))
+        if tid == Int32(0):
+            for j in cutlass.range_constexpr(self.nchunk + 1):
+                _mbar_init(sMbar + Int32(8 * j), Int32(1))
+            _mbar_fence_init()
+        cute.arch.sync_threads()
+
+        if cutlass.const_expr(self.stage_before_wait):
+            self._issue_copies(sW, sMbar, w0, w1, row0, nsm, tid)
+        _stamp(self.debug_ts, ts, bid, tid, 7)
+        if cutlass.const_expr(self.prefetch_next):
+            # L2 prefetch of the next kernels' weights (a hint; e.g. q_b and
+            # W_UK of the same layer): CTA b covers slice b of each tensor.
+            if tid < Int32(3):
+                pf_addr = _sel_i64(
+                    Int32(tid == Int32(0)), pf0, _sel_i64(Int32(tid == Int32(1)), pf1, pf2)
+                )
+                pf_bytes = _sel_i64(
+                    Int32(tid == Int32(0)),
+                    pf0_bytes,
+                    _sel_i64(Int32(tid == Int32(1)), pf1_bytes, pf2_bytes),
+                )
+                self._prefetch_slice(pf_addr, pf_bytes, bid)
+        if cutlass.const_expr(self.l2_prefetch):
+            # Rows that are not staged into shared memory before the wait.
+            first = Int32(0)
+            if cutlass.const_expr(self.stage_before_wait):
+                first = nsm
+            if tid >= first:
+                if tid < nrows:
+                    _prefetch_l2(_row_src(w0, w1, row0 + tid, k, self.n0, self.n1), Int32(k * 2))
+        if cutlass.const_expr(self.use_pdl and self.trigger == "start"):
+            cute.arch.griddepcontrol_launch_dependents()
+        if cutlass.const_expr(self.stage_before_wait and self.wait_staged_first):
+            # Every thread waits until the staged rows have landed before its
+            # dependency wait (they do not depend on the preceding grid).
+            self._wait_staged(sMbar)
+            _stamp(self.debug_ts, ts, bid, tid, 6)
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+        _stamp(self.debug_ts, ts, bid, tid, 1)
+
+        if cutlass.const_expr(not self.stage_before_wait):
+            self._issue_copies(sW, sMbar, w0, w1, row0, nsm, tid)
+            if cutlass.const_expr(self.copy_method == "cpasync"):
+                self._wait_staged(sMbar)
+        # Activation fragments: token g = lane / 4, K pair 2 * (lane % 4).
+        g = lane // Int32(4)
+        t4 = lane % Int32(4)
+        tok_ok = Int32(g < rows)
+        safe_g = _min_i32(g, rows - Int32(1))
+        # Square-sum partials for r: lane 4g + q of the last warp loads
+        # quarter q of token g's partials before its activation loads (they
+        # are summed after the main loop, in the hop's partial order).
+        pq = []
+        if cutlass.const_expr(self.deferred):
+            quarter = self.partials_per_row // 4
+            p_ok = Int32(warp == Int32(nw - 1)) * tok_ok
+            pbase = sumsq + Int64(
+                (g * Int32(self.partials_per_row) + t4 * Int32(quarter)) * Int32(4)
+            )
+            for j in cutlass.range_constexpr(quarter // 2):
+                lo, hi = _ld_v2_if(pbase + Int64(8 * j), p_ok)
+                pq.append(bits_f32(lo))
+                pq.append(bits_f32(hi))
+        bfr = []
+        if cutlass.const_expr(self.b_layout == "frag"):
+            # Lanes of token rows >= R load nothing (zero fragments).
+            a_lane = act + Int64(lane * Int32(8))
+            for m in cutlass.range_constexpr(self.rounds):
+                bfr.append(_ld_v2_if(a_lane + Int64((warp + Int32(nw * m)) * Int32(256)), tok_ok))
+        else:
+            a_row = act + Int64(safe_g * Int32(k) * Int32(2) + t4 * Int32(4))
+            for m in cutlass.range_constexpr(self.rounds):
+                off = Int64((warp + Int32(nw * m)) * Int32(32))
+                b0 = _ld_b32(a_row + off)
+                b1 = _ld_b32(a_row + off + Int64(16))
+                bfr.append((_sel_i32(tok_ok, b0, Int32(0)), _sel_i32(tok_ok, b1, Int32(0))))
+        _stamp(self.debug_ts, ts, bid, tid, 2)
+
+        # ldmatrix row addresses (per lane), advanced by 32 B per K step.
+        r16 = lane % Int32(16)
+        a_lane0 = sW + r16 * Int32(stride) + (lane // Int32(16)) * Int32(16)
+        ok0 = Int32(r16 < nsm)
+        r8 = Int32(16) + lane % Int32(8)
+        a_lane1 = sW + r8 * Int32(stride) + ((lane // Int32(8)) % Int32(2)) * Int32(16)
+        ok1 = Int32(r8 < nsm)
+
+        acc0 = [Float32(0.0)] * 4
+        acc1 = [Float32(0.0)] * 4
+        zero = Int32(0)
+        rpc = self.rounds_per_chunk
+
+        def load_a(n):
+            if cutlass.const_expr(n % rpc == 0 and self.copy_method == "bulk"):
+                _mbar_wait(sMbar + Int32(8 * (n // rpc)), Int32(0))
+            kb = (warp + Int32(nw * n)) * Int32(32)
+            return (
+                ldsm_x4(_sel_i32(ok0, a_lane0 + kb, sZero)),
+                _ldsm_x2(_sel_i32(ok1, a_lane1 + kb, sZero)),
+            )
+
+        frags = [load_a(n) for n in range(min(self.dist, self.rounds))]
+        for m in cutlass.range_constexpr(self.rounds):
+            if cutlass.const_expr(m + self.dist < self.rounds):
+                frags.append(load_a(m + self.dist))
+            fa, fb = frags[m]
+            b0, b1 = bfr[m]
+            acc0[0], acc0[1], acc0[2], acc0[3] = mma_bf16(
+                fa[0], fa[1], fa[2], fa[3], b0, b1, acc0[0], acc0[1], acc0[2], acc0[3]
+            )
+            acc1[0], acc1[1], acc1[2], acc1[3] = mma_bf16(
+                fb[0], zero, fb[1], zero, b0, b1, acc1[0], acc1[1], acc1[2], acc1[3]
+            )
+        _stamp(self.debug_ts, ts, bid, tid, 3)
+
+        # r per token (published by the barrier below): the serial sum of
+        # the 24 partials, quarter by quarter along the 4 lanes of token g.
+        if cutlass.const_expr(self.deferred):
+            if warp == Int32(nw - 1):
+                s = Float32(0.0)
+                for q in cutlass.range_constexpr(4):
+                    mine = s
+                    for v in pq:
+                        mine = add(mine, v)
+                    s = _shfl_idx_f32(mine, g * Int32(4) + Int32(q))
+                if t4 == Int32(0):
+                    _st_shared_f32(sR + g * Int32(4), inv_rms(s, Float32(float(k)), eps))
+        else:
+            if tid < Int32(MAX_TOKENS):
+                _st_shared_f32(sR + tid * Int32(4), Float32(1.0))
+
+        # Cross-warp reduction in warp order (the weight rows are dead now).
+        cute.arch.sync_threads()
+        if cutlass.const_expr(self.overflow > 0):
+            if tid == Int32(0):
+                if nrows > nsm:
+                    _fence_proxy_async()
+                    _mbar_expect_tx(sMbar + Int32(8 * self.nchunk), Int32(k * 2))
+                    _bulk_g2s(
+                        sW + Int32(2 * stride),
+                        _row_src(w0, w1, row0 + nsm, k, self.n0, self.n1),
+                        Int32(k * 2),
+                        sMbar + Int32(8 * self.nchunk),
+                    )
+        sRed = sW
+        wbase = sRed + warp * Int32(24 * MAX_TOKENS * 4)
+        c0 = Int32(2) * t4
+        _st_shared_f32(wbase + (g * Int32(MAX_TOKENS) + c0) * Int32(4), acc0[0])
+        _st_shared_f32(wbase + (g * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc0[1])
+        _st_shared_f32(wbase + ((g + Int32(8)) * Int32(MAX_TOKENS) + c0) * Int32(4), acc0[2])
+        _st_shared_f32(
+            wbase + ((g + Int32(8)) * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc0[3]
+        )
+        _st_shared_f32(wbase + ((g + Int32(16)) * Int32(MAX_TOKENS) + c0) * Int32(4), acc1[0])
+        _st_shared_f32(
+            wbase + ((g + Int32(16)) * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc1[1]
+        )
+        cute.arch.sync_threads()
+        if tid < Int32(24 * MAX_TOKENS):
+            row = tid // Int32(MAX_TOKENS)
+            tok = tid % Int32(MAX_TOKENS)
+            if row < nsm:
+                if tok < rows:
+                    s = Float32(0.0)
+                    for w in cutlass.range_constexpr(nw):
+                        s = add(
+                            s,
+                            _ld_shared_f32(sRed + Int32(w * 24 * MAX_TOKENS * 4) + tid * Int32(4)),
+                        )
+                    _store_out(
+                        out0,
+                        out1,
+                        row0 + row,
+                        tok,
+                        mul(_ld_shared_f32(sR + tok * Int32(4)), s),
+                        self.n0,
+                        self.n1,
+                    )
+        _stamp(self.debug_ts, ts, bid, tid, 4)
+
+        if cutlass.const_expr(self.overflow > 0):
+            if nrows > nsm:
+                # The extra row: rows 0-7 of an m16 tile through ldmatrix.x2
+                # with only row 0 real, K steps as in the main loop.
+                _mbar_wait(sMbar + Int32(8 * self.nchunk), Int32(0))
+                l8 = lane % Int32(8)
+                okx = Int32(l8 == Int32(0))
+                a_lanex = sW + Int32(2 * stride) + ((lane // Int32(8)) % Int32(2)) * Int32(16)
+                accx = [Float32(0.0)] * 4
+                for m in cutlass.range_constexpr(self.rounds):
+                    kb = (warp + Int32(nw * m)) * Int32(32)
+                    fx = _ldsm_x2(_sel_i32(okx, a_lanex + kb, sZero))
+                    b0, b1 = bfr[m]
+                    accx[0], accx[1], accx[2], accx[3] = mma_bf16(
+                        fx[0], zero, fx[1], zero, b0, b1, accx[0], accx[1], accx[2], accx[3]
+                    )
+                sRedX = sW + Int32(3 * stride)
+                if g == Int32(0):
+                    _st_shared_f32(sRedX + (warp * Int32(MAX_TOKENS) + c0) * Int32(4), accx[0])
+                    _st_shared_f32(
+                        sRedX + (warp * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), accx[1]
+                    )
+                cute.arch.sync_threads()
+                if tid < rows:
+                    s = Float32(0.0)
+                    for w in cutlass.range_constexpr(nw):
+                        s = add(
+                            s, _ld_shared_f32(sRedX + Int32(w * MAX_TOKENS * 4) + tid * Int32(4))
+                        )
+                    _store_out(
+                        out0,
+                        out1,
+                        row0 + nsm,
+                        tid,
+                        mul(_ld_shared_f32(sR + tid * Int32(4)), s),
+                        self.n0,
+                        self.n1,
+                    )
+        if cutlass.const_expr(self.use_pdl and self.trigger == "end"):
+            cute.arch.griddepcontrol_launch_dependents()
+        _stamp(self.debug_ts, ts, bid, tid, 5)
+
+
+_COMPILED = {}
+
+
+def is_compiled(**cfg) -> bool:
+    """Whether ``compiled_projection(**cfg)`` would be a cache hit."""
+    return tuple(sorted(cfg.items())) in _COMPILED
+
+
+def compiled_projection(**cfg):
+    """Return the compiled projection for one configuration (compiles once, no GPU work)."""
+    key = tuple(sorted(cfg.items()))
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = DeferredNormProjKernel(**cfg)
+        z64 = Int64(0)
+        anchor = cute_runtime.make_fake_compact_tensor(
+            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
+        )
+        compiled = cute.compile(
+            kernel,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            anchor,
+            Int32(1),
+            Float32(1e-5),
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+        _COMPILED[key] = compiled
+    return compiled
+
+
+def projection_supported(
+    k: int, n: int, ctas: int, threads: int, smem_rows: int, nchunk: int
+) -> Optional[str]:
+    """Return why the projection kernel cannot serve a shape, or ``None``."""
+    nw = threads // 32
+    if k % 16 or (k // 16) % nw:
+        return f"K={k} does not split into {nw} warps of 16-wide steps"
+    if k % nchunk or (k // nchunk) % (16 * nw):
+        return f"K={k} does not split into {nchunk} chunks of whole warp rounds"
+    base, rem = divmod(n, ctas)
+    max_rows = base + (1 if rem else 0)
+    if base <= 16 or max_rows > 24:
+        return f"{n} rows over {ctas} CTAs gives {base}..{max_rows} rows per CTA (17..24 supported)"
+    if max_rows - min(smem_rows, max_rows) > 1:
+        return f"{max_rows} rows per CTA exceed {smem_rows} staged rows by more than one"
+    if min(smem_rows, max_rows) * (k * 2 + 16) + 1024 > SMEM_LIMIT:
+        return "staged rows exceed the shared-memory limit"
+    return None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/projection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/projection.py
new file mode 100644
index 000000000000..fe237611a1a7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/projection.py
@@ -0,0 +1,305 @@
+# 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.
+"""deferred_norm next-layer projection: the fused-A GEMM (and the indexer wk|weights_proj
+GEMM on full-indexer layers) on a deferred-norm packet, with staged weights.
+
+``DeferredNormProjection`` is bound to one native DSA ``MLA`` module. It reads the
+module's own weight tensors in place (``kv_a_proj_with_mqa.weight`` and, with
+``TRTLLM_DSA_INDEXER_BF16=1``, ``indexer._fused_wk_wp_weight``); no packed or
+gain-folded copy exists, so the extra weight memory is zero. ``project``
+returns ``(kva [R, 2624], kw [R, 160] or None)`` with the dtypes and layouts
+of the native ``dsv3_fused_a_gemm_op`` / ``F.linear`` outputs.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+
+from .packet import DeferredNormHidden
+
+_K = 6144
+_N_KVA = 2624
+_N_KW = 160
+MAX_ROWS = 8
+
+
+@dataclass(frozen=True)
+class DeferredNormProjConfig:
+    """Launch configuration of the projection kernel.
+
+    ``ctas = 0`` means one CTA per SM. ``stage_before_wait`` copies the
+    weights into shared memory before the PDL dependency wait (the deferred_norm
+    schedule); ``False`` copies them after the wait (attribution arm).
+    ``trigger`` is when the kernel releases its own PDL dependents.
+    ``dist`` is how many K rounds ahead each warp loads its weight fragments
+    from shared memory.
+    """
+
+    ctas: int = 0
+    threads: int = 512
+    smem_rows: int = 18
+    nchunk: int = 4
+    trigger: str = "start"
+    stage_before_wait: bool = True
+    l2_prefetch: bool = True
+    dist: int = 4
+    copy_threads: int = 1
+    copy_method: str = "bulk"
+    wait_staged_first: bool = False
+    prefetch_next: bool = False
+
+
+def _sm_count(device) -> int:
+    return torch.cuda.get_device_properties(device).multi_processor_count
+
+
+def _kernel_config(
+    config: DeferredNormProjConfig,
+    n1: int,
+    ctas: int,
+    use_pdl: bool,
+    deferred: bool,
+    debug_ts: bool,
+) -> dict:
+    return dict(
+        k=_K,
+        n0=_N_KVA,
+        n1=n1,
+        ctas=ctas,
+        threads=config.threads,
+        smem_rows=config.smem_rows,
+        nchunk=config.nchunk,
+        deferred=deferred,
+        partials_per_row=24,
+        use_pdl=use_pdl,
+        trigger=config.trigger,
+        stage_before_wait=config.stage_before_wait,
+        l2_prefetch=config.l2_prefetch,
+        dist=config.dist,
+        copy_threads=config.copy_threads,
+        copy_method=config.copy_method,
+        wait_staged_first=config.wait_staged_first,
+        prefetch_next=config.prefetch_next,
+        debug_ts=debug_ts,
+    )
+
+
+# Indexer rows appended to the fused-A rows: shared-indexer layers (none) and
+# full-indexer layers (the BF16 wk | weights_proj rows).
+SERVED_N1 = (0, _N_KW)
+
+
+def served_kernel_configs(device) -> list:
+    """Kernel configurations the deferred_norm consumer launches in serving.
+
+    One per layer kind (``SERVED_N1``): deferred mode, fragment-layout
+    activation, no debug stamps; the default launch configuration
+    (``DeferredNormProjConfig()``) and the PDL policy from ``TRTLLM_ENABLE_PDL``,
+    exactly as :meth:`DeferredNormProjection.served_config` builds them.
+    """
+    from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+    config = DeferredNormProjConfig()
+    ctas = config.ctas or _sm_count(device)
+    use_pdl = get_env_enable_pdl()
+    out = []
+    for n1 in SERVED_N1:
+        cfg = _kernel_config(config, n1, ctas, use_pdl, deferred=True, debug_ts=False)
+        cfg["b_layout"] = "frag"
+        out.append(cfg)
+    return out
+
+
+def precompile_served(device) -> Optional[str]:
+    """Compile every served projection kernel (no GPU work).
+
+    Called by ``DeferredNormHop`` before its cross-rank agreement, so a
+    deferred-norm packet exists only when every rank compiled both kernels
+    and the consumer's compile is a cache hit. Returns why a served shape is
+    not supported (rank-identical static check), or ``None``; a compile
+    error propagates to the caller.
+    """
+    from .proj_kernel import compiled_projection, projection_supported
+
+    for cfg in served_kernel_configs(device):
+        why = projection_supported(
+            cfg["k"],
+            cfg["n0"] + cfg["n1"],
+            cfg["ctas"],
+            cfg["threads"],
+            cfg["smem_rows"],
+            cfg["nchunk"],
+        )
+        if why is not None:
+            return f"projection n1={cfg['n1']}: {why}"
+    for cfg in served_kernel_configs(device):
+        compiled_projection(**cfg)
+    return None
+
+
+def projection_eligibility(mla) -> Tuple[bool, str, Optional[torch.Tensor]]:
+    """Static check of one MLA module; returns ``(ok, reason, kw_weight)``.
+
+    Under the served ``shared`` scope a full-indexer module is declined: its
+    producing hop is post_moe_hop's in that scope, and a packet that reached it
+    anyway is materialized (post_moe_hop's normalized output) for the native path.
+    """
+    from tensorrt_llm._utils import get_sm_version
+
+    from .scope import SCOPE_LABEL, selected_scope
+
+    if get_sm_version() not in (100, 103):
+        return False, f"sm{get_sm_version()} not validated", None
+    lin = getattr(mla, "kv_a_proj_with_mqa", None)
+    if lin is None or not getattr(mla, "fuse_qkv_a_proj", False):
+        return False, "needs the fused [q_a | kv_a] projection", None
+    w = lin.weight
+    if (
+        lin.has_any_quant
+        or lin.bias is not None
+        or w.dtype != torch.bfloat16
+        or tuple(w.shape) != (_N_KVA, _K)
+        or not w.is_contiguous()
+    ):
+        return False, "kv_a_proj_with_mqa must be an unquantized BF16 [2624, 6144] weight", None
+    kw = None
+    indexer = mla.mqa.indexer if getattr(mla, "mqa", None) is not None else None
+    if indexer is not None:
+        scope = selected_scope()
+        if scope == "shared":
+            return (
+                False,
+                f"{SCOPE_LABEL}={scope}: full-indexer layer keeps its native projection",
+                None,
+            )
+        kw = getattr(indexer, "_fused_wk_wp_weight", None)
+        if (
+            not getattr(indexer, "_indexer_bf16", False)
+            or kw is None
+            or kw.dtype != torch.bfloat16
+            or tuple(kw.shape) != (_N_KW, _K)
+            or not kw.is_contiguous()
+            or getattr(indexer, "use_fp4", False)
+        ):
+            return False, "full-indexer layer needs the BF16 fused wk|weights_proj GEMM", None
+    return True, "", kw
+
+
+class DeferredNormProjection:
+    """Staged projection bound to one MLA module (see the module docstring)."""
+
+    def __init__(
+        self,
+        mla,
+        config: Optional[DeferredNormProjConfig] = None,
+        kw_weight: Optional[torch.Tensor] = None,
+    ):
+        from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+        self.mla = mla
+        self.config = config or DeferredNormProjConfig()
+        self.w0 = mla.kv_a_proj_with_mqa.weight
+        self.w1 = kw_weight
+        self.n1 = 0 if kw_weight is None else kw_weight.shape[0]
+        self.use_pdl = get_env_enable_pdl()
+        self.ctas = self.config.ctas or _sm_count(self.w0.device)
+        self.debug_ts = None
+        # (address, bytes) of up to three tensors the next kernels read
+        # (L2-prefetched before the dependency wait with ``prefetch_next``).
+        self.prefetch = []
+
+    def set_prefetch(self, tensors) -> None:
+        """Tensors the following kernels read (weights only; L2 hint)."""
+        self.prefetch = [t for t in tensors if t is not None and t.is_contiguous()][:3]
+
+    def kernel_config(self, deferred: bool = True, debug_ts: bool = False) -> dict:
+        return _kernel_config(self.config, self.n1, self.ctas, self.use_pdl, deferred, debug_ts)
+
+    def served_config(self) -> dict:
+        """Configuration :meth:`project` launches (one of ``served_kernel_configs``
+        when the module's config comes from the environment)."""
+        cfg = self.kernel_config()
+        cfg["b_layout"] = "frag"
+        return cfg
+
+    def supported(self) -> Optional[str]:
+        from .proj_kernel import projection_supported
+
+        cfg = self.config
+        return projection_supported(
+            _K, _N_KVA + self.n1, self.ctas, cfg.threads, cfg.smem_rows, cfg.nchunk
+        )
+
+    def precompile(self) -> None:
+        """Compile the served kernel; a cache hit after ``precompile_served``.
+        A compile error propagates."""
+        from .proj_kernel import compiled_projection
+
+        compiled_projection(**self.served_config())
+
+    def run(
+        self,
+        act: torch.Tensor,
+        sumsq: Optional[torch.Tensor],
+        eps: float,
+        rows: Optional[int] = None,
+        ts: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Launch on the current stream.
+
+        ``act`` 1-D: the ``a_frag`` buffer of a packet (``rows`` token rows);
+        2-D ``[R, K]``: a row-major activation. ``sumsq`` given: deferred mode
+        (``act`` holds ``bf16(x * gamma)``); ``None``: native-input mode
+        (``act`` is a normalized hidden state, ``r = 1``).
+        """
+        from .proj_kernel import compiled_projection
+
+        frag = act.dim() == 1
+        if not frag:
+            rows = act.shape[0]
+        assert rows is not None and 1 <= rows <= MAX_ROWS
+        assert act.is_contiguous() and act.dtype == torch.bfloat16
+        kva = act.new_empty(rows, _N_KVA)
+        kw = act.new_empty(rows, self.n1) if self.n1 else None
+        deferred = sumsq is not None
+        cfg = self.kernel_config(deferred, ts is not None)
+        cfg["b_layout"] = "frag" if frag else "plain"
+        compiled = compiled_projection(**cfg)
+        pf = [(t.data_ptr(), t.numel() * t.element_size()) for t in self.prefetch[:3]]
+        pf += [(0, 0)] * (3 - len(pf))
+        compiled(
+            act.data_ptr(),
+            sumsq.data_ptr() if deferred else 0,
+            self.w0.data_ptr(),
+            0 if self.w1 is None else self.w1.data_ptr(),
+            kva.data_ptr(),
+            0 if kw is None else kw.data_ptr(),
+            0 if ts is None else ts.data_ptr(),
+            pf[0][0],
+            pf[0][1],
+            pf[1][0],
+            pf[1][1],
+            pf[2][0],
+            pf[2][1],
+            kva.view(torch.int32).view(-1),
+            rows,
+            float(eps),
+        )
+        return kva, kw
+
+    def project(self, packet: DeferredNormHidden) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        assert packet.sumsq.shape[1] == 24 and packet.sumsq.is_contiguous()
+        return self.run(packet.a_frag, packet.sumsq, packet.eps, rows=packet.num_rows)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/ptx.py
new file mode 100644
index 000000000000..ae1271ed0958
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/ptx.py
@@ -0,0 +1,619 @@
+# 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.
+"""Inline-PTX helpers of the deferred_norm kernels (CuTe DSL).
+
+deferred_norm owns every device helper its kernels use; it imports none from another
+package. Two groups follow a contract defined elsewhere and must keep its
+exact semantics:
+
+* the arithmetic of the post_moe_hop (finalize FMA chain, bf16 packing/rounding,
+  sanitize, square sums, ``rsqrt`` operand), which makes the deferred_norm residual,
+  gain-folded activation, square-sum partials and ``r`` bit-identical to the
+  post_moe_hop's;
+* the post_moe_hop workspace protocol (``post_moe_hop/kernels.py`` and
+  ``post_moe_hop/workspace.py``), which the deferred_norm hop shares with the post_moe_hop hops of
+  the same TP group: Lamport sentinel, timer-bounded poll, per-buffer error words with sticky
+  poison, timeout poll count, epoch/arrival/call counters. Every state-word byte
+  offset is derived here from post_moe_hop's exported word indices, so the two hops
+  address the same words.
+
+The tests check both: bitwise equality with the post_moe_hop on the same workspace
+(also interleaved with post_moe_hop calls, in CUDA graphs, under scrambled store
+order), and the error protocol across both hop types.
+"""
+
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from ..post_moe_hop.kernels import (
+    ERROR_LAMPORT_TIMEOUT,
+    STATE_ARRIVALS,
+    STATE_CALLS,
+    STATE_EPOCH,
+    STATE_ERROR_WORDS,
+    STATE_PREV_ROWS,
+    STATE_TIMEOUT_POLLS,
+    STATE_WORDS,
+)
+
+# The hop loads the state words with two 16-byte vectors.
+assert STATE_WORDS == 8, "deferred_norm loads the post_moe_hop state words as two 16-byte vectors"
+
+SENTINEL_WORD = -2147450880  # 0x80008000: two bf16 -0.0 (empty Lamport element)
+FP32_SENTINEL = -2147483648  # 0x80000000: fp32 -0.0 (empty row-sum partial)
+# Polls between two deadline checks; the clock starts at the first check.
+POLLS_PER_CLOCK_CHECK = 1024
+
+_OFF_EPOCH = 4 * STATE_EPOCH
+_OFF_PREV_ROWS = 4 * STATE_PREV_ROWS
+_OFF_ARRIVALS = 4 * STATE_ARRIVALS
+_OFF_CALLS = 4 * STATE_CALLS
+_OFF_TIMEOUT_POLLS = 4 * STATE_TIMEOUT_POLLS
+ERROR_WORD_OFFSETS = tuple(4 * w for w in STATE_ERROR_WORDS)
+
+
+def _asm(res, args, text, cons, side_effects=True, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=side_effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+def _v(x, loc, ip):
+    return x.ir_value(loc=loc, ip=ip)
+
+
+def _struct(n, ty):
+    return llvm.StructType.get_literal([ty] * n)
+
+
+def _unpack(res, n, cls, ty, loc, ip):
+    return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n))
+
+
+# --------------------------------------------------------------------------
+# Global memory and the post_moe_hop workspace protocol.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def ld_v4(addr: Int64, *, loc=None, ip=None):
+    res = _asm(
+        _struct(4, T.i32()),
+        [_v(addr, loc, ip)],
+        "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,l",
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(res, 4, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def st_v4(addr: Int64, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [_v(addr, loc, ip)] + [_v(x, loc, ip) for x in (a, b, c, d)],
+        "st.global.v4.b32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+def load_state(state: Int64):
+    """The 8 state words of the rank's post_moe_hop workspace (two 16-byte loads).
+
+    Returns ``(epoch, prev_rows, error_words)`` with the three per-buffer
+    error words in buffer order. Called after the kernel's dependency wait;
+    the words were written by the previous hop call of this workspace, which
+    is complete at that point (kernel boundary).
+    """
+    words = ld_v4(state) + ld_v4(state + Int64(16))
+    return (words[STATE_EPOCH], words[STATE_PREV_ROWS], [words[w] for w in STATE_ERROR_WORDS])
+
+
+@dsl_user_op
+def poll8(base: Int64, stride: Int64, timeout_ns: Int64, *, loc=None, ip=None):
+    """Spin until 8 rank slots (16 bytes each, ``stride`` apart) hold no sentinel.
+
+    Volatile 16-byte loads of all 8 slots per round; the ``%globaltimer``
+    deadline is checked every ``POLLS_PER_CLOCK_CHECK`` rounds and its clock
+    starts at the first check, so a poll that completes within that many
+    rounds never reads the timer. Returns the 32 words, a nonzero flag when
+    the deadline expired, and the round count. Same loop as the post_moe_hop cluster
+    kernel's poll.
+    """
+    loads = []
+    checks = []
+    for r in range(8):
+        o = 4 * r
+        loads.append(
+            f"ld.volatile.global.v4.b32 {{${o}, ${o + 1}, ${o + 2}, ${o + 3}}}, [a];\n"
+            "add.s64 a, a, $35;\n"
+        )
+        for w in range(4):
+            checks.append(
+                f"and.b32 t, ${o + w}, 0xFFFF0000;\n"
+                "setp.eq.u32 p, t, 0x80000000;\n"
+                "selp.b32 t, 1, 0, p;\n"
+                "or.b32 d, d, t;\n"
+            )
+    period = POLLS_PER_CLOCK_CHECK
+    asm = (
+        "{\n.reg .pred p;\n.reg .b32 t, d, n;\n.reg .b64 a, t0, now;\n"
+        "mov.b32 n, 0;\nmov.b64 t0, 0;\n"
+        "NORM_POLL:\nmov.b64 a, $34;\nmov.b32 d, 0;\n"
+        + "".join(loads)
+        + "".join(checks)
+        + "add.u32 n, n, 1;\nsetp.eq.u32 p, d, 0;\n@p bra NORM_POLL_DONE;\n"
+        f"and.b32 t, n, {period - 1};\nsetp.ne.u32 p, t, 0;\n@p bra NORM_POLL;\n"
+        "mov.u64 now, %globaltimer;\n"
+        f"setp.eq.u32 p, n, {period};\n@p mov.b64 t0, now;\n"
+        "sub.s64 now, now, t0;\nsetp.lt.s64 p, now, $36;\n@p bra NORM_POLL;\n"
+        "NORM_POLL_DONE:\nmov.b32 $32, d;\nmov.b32 $33, n;\n}\n"
+    )
+    constraints = ",".join(["=r"] * 34 + ["l", "l", "l"])
+    res = _asm(
+        _struct(34, T.i32()),
+        [_v(base, loc, ip), _v(stride, loc, ip), _v(timeout_ns, loc, ip)],
+        asm,
+        constraints,
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(res, 34, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def select3(values: list, index: Int32, *, loc=None, ip=None) -> Int32:
+    """``values[index]`` for ``index`` in 0..2 (e.g. the error word of a buffer)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(index, loc, ip)] + [_v(w, loc, ip) for w in values],
+            "{\n.reg .pred p;\nsetp.eq.u32 p, $1, 1;\nselp.b32 $0, $3, $2, p;\n"
+            "setp.eq.u32 p, $1, 2;\nselp.b32 $0, $4, $0, p;\n}\n",
+            "=r,r,r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+def error_word_address(state: Int64, buffer: Int32) -> Int64:
+    """Address of the error word of Lamport buffer ``buffer`` (0..2)."""
+    return state + Int64(select3([Int32(o) for o in ERROR_WORD_OFFSETS], buffer))
+
+
+@dsl_user_op
+def flag_error(word: Int64, flag: Int32, bits: Int32, *, loc=None, ip=None) -> None:
+    """OR ``bits`` into the error word at ``word`` when ``flag`` is nonzero."""
+    _asm(
+        None,
+        [_v(word, loc, ip), _v(flag, loc, ip), _v(bits, loc, ip)],
+        "{\n.reg .pred p;\nsetp.ne.s32 p, $1, 0;\n@p red.relaxed.gpu.global.or.b32 [$0], $2;\n}\n",
+        "l,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def flag_timeout(
+    word: Int64, state: Int64, flag: Int32, polls: Int32, *, loc=None, ip=None
+) -> None:
+    """The Lamport-timeout bit into the error word and the largest poll count
+    at expiry into the ``timeout_polls`` state word."""
+    _asm(
+        None,
+        [_v(word, loc, ip), _v(state, loc, ip), _v(flag, loc, ip), _v(polls, loc, ip)],
+        "{\n.reg .pred p;\nsetp.ne.s32 p, $2, 0;\n"
+        f"@p red.relaxed.gpu.global.or.b32 [$0], {ERROR_LAMPORT_TIMEOUT};\n"
+        f"@p red.relaxed.gpu.global.max.u32 [$1+{_OFF_TIMEOUT_POLLS}], $3;\n}}\n",
+        "l,l,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def arrive(
+    state: Int64,
+    is_leader: Int32,
+    num_ctas: Int32,
+    num_rows: Int32,
+    next_epoch: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    """Leader thread of each CTA counts in; the last CTA records the row count
+    to clear next, resets the arrival count, counts the call and publishes
+    the next epoch (post_moe_hop's arrival protocol)."""
+    _asm(
+        None,
+        [
+            _v(state, loc, ip),
+            _v(is_leader, loc, ip),
+            _v(num_ctas, loc, ip),
+            _v(num_rows, loc, ip),
+            _v(next_epoch, loc, ip),
+        ],
+        "{\n.reg .pred p, q;\n.reg .b32 old, c;\n"
+        "setp.ne.s32 p, $1, 0;\nmov.b32 old, 0;\n"
+        f"@p atom.acq_rel.gpu.global.add.u32 old, [$0+{_OFF_ARRIVALS}], 1;\n"
+        "add.u32 c, old, 1;\nsetp.eq.u32 q, c, $2;\nand.pred q, q, p;\n"
+        f"@q st.global.u32 [$0+{_OFF_PREV_ROWS}], $3;\n"
+        f"@q st.global.u32 [$0+{_OFF_ARRIVALS}], 0;\n"
+        f"@q red.relaxed.gpu.global.add.u32 [$0+{_OFF_CALLS}], 1;\n"
+        f"@q st.global.u32 [$0+{_OFF_EPOCH}], $4;\n}}\n",
+        "l,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def fault_delay(a: Int32, b: Int32, c: Int32, max_ns: int, site: int, *, loc=None, ip=None) -> None:
+    """Test-only fault injection: about one thread in four spins before ``site``.
+
+    The delay (below ``max_ns``) is a hash of ``a``, ``b`` and ``c`` (post_moe_hop's
+    hash), so the individual 16-byte stores of one call land in a scrambled
+    order.
+    """
+    asm = (
+        "{\n.reg .pred p;\n.reg .b32 h, t;\n.reg .b64 t0, now, d;\n"
+        "mul.lo.u32 h, $0, 0x9E3779B1;\nmul.lo.u32 t, $1, 0x85EBCA77;\nxor.b32 h, h, t;\n"
+        "mul.lo.u32 t, $2, 0xC2B2AE3D;\nxor.b32 h, h, t;\nshr.u32 t, h, 15;\n"
+        "xor.b32 h, h, t;\nmul.lo.u32 h, h, 0x2C1B3C6D;\nshr.u32 t, h, 13;\n"
+        "xor.b32 h, h, t;\nand.b32 t, h, 3;\nsetp.ne.u32 p, t, 0;\n"
+        f"@p bra NORM_FAULT_END_{site};\nshr.u32 t, h, 2;\nrem.u32 t, t, {max_ns};\n"
+        "cvt.u64.u32 d, t;\nmov.u64 t0, %globaltimer;\n"
+        f"NORM_FAULT_LOOP_{site}:\nmov.u64 now, %globaltimer;\nsub.s64 now, now, t0;\n"
+        f"setp.lt.s64 p, now, d;\n@p bra NORM_FAULT_LOOP_{site};\nNORM_FAULT_END_{site}:\n}}\n"
+    )
+    _asm(None, [_v(a, loc, ip), _v(b, loc, ip), _v(c, loc, ip)], asm, "r,r,r", loc=loc, ip=ip)
+
+
+# --------------------------------------------------------------------------
+# Arithmetic of the post_moe_hop.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def unpack_bf16x2(word: Int32, *, loc=None, ip=None):
+    res = _asm(
+        _struct(2, T.f32()),
+        [_v(word, loc, ip)],
+        "{\n.reg .b32 t;\nshl.b32 t, $2, 16;\nmov.b32 $0, t;\n"
+        "and.b32 t, $2, 0xFFFF0000;\nmov.b32 $1, t;\n}\n",
+        "=f,=f,r",
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(res, 2, Float32, T.f32(), loc, ip)
+
+
+@dsl_user_op
+def pack_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Int32:
+    """Round-to-nearest-even of two fp32 values into one bf16x2 word (``lo`` low)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(lo, loc, ip), _v(hi, loc, ip)],
+            "cvt.rn.bf16x2.f32 $0, $2, $1;",
+            "=r,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def round_bf16(x: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(x, loc, ip)],
+            "{\n.reg .b32 t;\ncvt.rn.bf16x2.f32 t, $1, $1;\n"
+            "and.b32 t, t, 0xFFFF0000;\nmov.b32 $0, t;\n}\n",
+            "=f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def sanitize(word: Int32, *, loc=None, ip=None) -> Int32:
+    """bf16 -0.0 (the Lamport sentinel) -> +0.0 in both halves."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(word, loc, ip)],
+            "{\n.reg .b32 t;\n.reg .pred p, q;\n"
+            "and.b32 t, $1, 0xFFFF;\nsetp.eq.u32 p, t, 0x8000;\n"
+            "and.b32 t, $1, 0xFFFF0000;\nsetp.eq.u32 q, t, 0x80000000;\n"
+            "mov.b32 $0, $1;\n@p and.b32 $0, $0, 0xFFFF0000;\n"
+            "@q and.b32 $0, $0, 0xFFFF;\n}\n",
+            "=r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def poison(word: Int32, flag: Int32, *, loc=None, ip=None) -> Int32:
+    """Two bf16 NaNs instead of ``word`` when ``flag`` is nonzero."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(word, loc, ip), _v(flag, loc, ip)],
+            "{\n.reg .pred p;\nsetp.ne.s32 p, $2, 0;\nselp.b32 $0, 0x7FC07FC0, $1, p;\n}\n",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def perm_check(index: Int32, rows: Int32, *, loc=None, ip=None):
+    """Return ``(valid, safe_index, out_of_range)`` of one permutation entry.
+
+    ``valid`` is ``0 <= index < rows``; ``safe_index`` is ``index`` when valid
+    and 0 otherwise; ``out_of_range`` flags ``index >= rows`` (negative
+    entries are the native "no local row" marker, not an error).
+    """
+    res = _asm(
+        _struct(3, T.i32()),
+        [_v(index, loc, ip), _v(rows, loc, ip)],
+        "{\n.reg .pred p, q;\nsetp.lt.u32 p, $3, $4;\nselp.b32 $0, 1, 0, p;\n"
+        "selp.b32 $1, $3, 0, p;\nsetp.ge.s32 q, $3, $4;\nselp.b32 $2, 1, 0, q;\n}\n",
+        "=r,=r,=r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(res, 3, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def fma(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(a, loc, ip), _v(b, loc, ip), _v(c, loc, ip)],
+            "fma.rn.f32 $0, $1, $2, $3;",
+            "=f,f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def add(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(a, loc, ip), _v(b, loc, ip)],
+            "add.rn.f32 $0, $1, $2;",
+            "=f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def mul(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(a, loc, ip), _v(b, loc, ip)],
+            "mul.rn.f32 $0, $1, $2;",
+            "=f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def select_f32(pred: Int32, a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(pred, loc, ip), _v(a, loc, ip), _v(b, loc, ip)],
+            "{\n.reg .pred p;\nsetp.ne.s32 p, $1, 0;\nselp.f32 $0, $2, $3, p;\n}\n",
+            "=f,r,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def bits_f32(word: Int32, *, loc=None, ip=None) -> Float32:
+    return Float32(_asm(T.f32(), [_v(word, loc, ip)], "mov.b32 $0, $1;", "=f,r", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def inv_rms(square_sum: Float32, hidden: Float32, eps: Float32, *, loc=None, ip=None) -> Float32:
+    """``rsqrt.approx(div.rn(square_sum, hidden) + eps)`` (post_moe_hop's ``r``)."""
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(square_sum, loc, ip), _v(hidden, loc, ip), _v(eps, loc, ip)],
+            "{\n.reg .f32 t;\ndiv.rn.f32 t, $1, $2;\nadd.rn.f32 t, t, $3;\n"
+            "rsqrt.approx.f32 $0, t;\n}\n",
+            "=f,f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def shfl_xor(x: Float32, mask: int, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(x, loc, ip)],
+            f"shfl.sync.bfly.b32 $0, $1, {mask}, 0x1f, 0xffffffff;",
+            "=f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+# --------------------------------------------------------------------------
+# Projection kernel: shared memory, cp.async and MMA.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def saddr(p: cute.Pointer, *, loc=None, ip=None) -> Int32:
+    """32-bit shared-window address of a shared-memory pointer."""
+    return Int32(p.toint(loc=loc, ip=ip))
+
+
+@dsl_user_op
+def cp_async16(smem: Int32, gaddr: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [smem.ir_value(), gaddr.ir_value()],
+        "cp.async.cg.shared.global [$0], [$1], 16;",
+        "r,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def cp_async_commit(*, loc=None, ip=None):
+    _asm(None, [], "cp.async.commit_group;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def cp_async_wait_group(n: int, *, loc=None, ip=None):
+    """Wait until at most ``n`` committed cp.async groups are pending."""
+    _asm(None, [], f"cp.async.wait_group {n};", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def st_global_b16(gaddr: Int64, v: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [gaddr.ir_value(), v.ir_value()],
+        "{\n\t.reg .b16 lo;\n\tcvt.u16.u32 lo, $1;\n\tst.global.b16 [$0], lo;\n\t}",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_shared_v4(smem: Int32, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [smem.ir_value(), a.ir_value(), b.ir_value(), c.ir_value(), d.ir_value()],
+        "st.shared.v4.b32 [$0], {$1, $2, $3, $4};",
+        "r,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def ldsm_x4(smem: Int32, *, loc=None, ip=None):
+    out = _asm(
+        _struct(4, T.i32()),
+        [smem.ir_value()],
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(out, 4, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def mma_bf16(
+    a0: Int32,
+    a1: Int32,
+    a2: Int32,
+    a3: Int32,
+    b0: Int32,
+    b1: Int32,
+    c0: Float32,
+    c1: Float32,
+    c2: Float32,
+    c3: Float32,
+    *,
+    loc=None,
+    ip=None,
+):
+    """``mma.sync.m16n8k16`` BF16 x BF16 + FP32 (row.col)."""
+    out = _asm(
+        _struct(4, T.f32()),
+        [
+            a0.ir_value(),
+            a1.ir_value(),
+            a2.ir_value(),
+            a3.ir_value(),
+            b0.ir_value(),
+            b1.ir_value(),
+            c0.ir_value(),
+            c1.ir_value(),
+            c2.ir_value(),
+            c3.ir_value(),
+        ],
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
+        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+        side_effects=False,
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(out, 4, Float32, T.f32(), loc, ip)
+
+
+@dsl_user_op
+def bf16_bits(x: Float32, *, loc=None, ip=None) -> Int32:
+    """Round-to-nearest-even of one fp32 value to bf16 (in the low half)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [x.ir_value()],
+            "{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tcvt.u32.u16 $0, h;\n\t}",
+            "=r,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/query_entry.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/query_entry.py
new file mode 100644
index 000000000000..35e42ca1440d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/query_entry.py
@@ -0,0 +1,224 @@
+# 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.
+"""deferred_norm entry into the absorbed_query / shared_indexer_absorbed_query path (reachable only
+with ``deferred_norm``).
+
+When the previous layer's hop deferred its RMSNorm (``deferred_norm``), the
+hidden state of a target layer is a ``DeferredNormHidden`` packet. This entry
+replaces exactly the two native GEMM bridges of absorbed_query / shared_indexer_absorbed_query that
+consume the normalized hidden state -- the fused-A ``kv_a_proj_with_mqa`` GEMM and, on
+full-indexer layers, the indexer ``wk | weights_proj`` GEMM -- by the deferred_norm
+staged projection, and then runs the unchanged absorbed_query / shared_indexer_absorbed_query
+kernels on their outputs, in the order the adapter uses: shared_indexer_absorbed_query's
+shared-layer kernel when ``shared_indexer_absorbed_query`` is selected and the layer qualifies, else
+absorbed_query's. On full-indexer layers absorbed_query's fused kernel takes the native
+``q_a_layernorm`` and indexer ``k_norm`` outputs (absorbed_query's bridges); they run here on the
+projected outputs. Like every absorbed_query-package entry point it validates the package selectors
+(``check_package_selectors``) before its first per-module decision.
+
+It returns ``None`` before any GPU work outside the support domain; the
+caller then materializes the packet (post_moe_hop's normalized output) and runs the
+next path. absorbed_query, draft_absorbed_query and shared_indexer_absorbed_query entry points never
+reach this module.
+"""
+
+from typing import Optional
+
+import torch
+
+from ...dispatch import is_enabled, log_capability
+from ..absorbed_query.query_proj import (
+    _KV_LORA,
+    AbsorbedQueryPacket,
+    AbsorbedQueryProjection,
+    absorb_for,
+    call_eligibility,
+    check_package_selectors,
+    indexer_tail_from_projected,
+    module_eligibility,
+)
+
+_SHARED_QUERY_MAX_ROWS = 8
+
+
+def _deferred_norm_runner(mla):
+    """Per-module deferred_norm projection (cached in ``mla._deferred_norm_projection``) or False.
+
+    Every decline here is a static check of replicated module state and the
+    environment, identical on all TP ranks. The kernel itself was compiled
+    by the producing ``DeferredNormHop`` before its cross-rank agreement, so
+    ``precompile`` is a cache hit; a compile error is not masked by a
+    rank-local fallback (one rank rounding differently from the others would
+    diverge the replicated caches and top-k sets) and propagates.
+    """
+    runner = getattr(mla, "_deferred_norm_projection", None)
+    if runner is not None:
+        return runner
+    from .projection import DeferredNormProjection, projection_eligibility
+
+    check_package_selectors()
+    ok, reason = module_eligibility(mla)
+    kw = None
+    if ok:
+        ok, reason, kw = projection_eligibility(mla)
+    runner = False
+    if ok:
+        proj = DeferredNormProjection(mla, kw_weight=kw)
+        why = proj.supported()
+        if why is not None:
+            ok, reason = False, why
+        else:
+            proj.precompile()
+            runner = proj
+    log_capability("deferred_norm", f"projection_layer{mla.layer_idx}", ok, reason)
+    mla._deferred_norm_projection = runner
+    return runner
+
+
+def _shared_query_runner(mla):
+    """shared_indexer_absorbed_query's cached shared-layer runner when
+    ``shared_indexer_absorbed_query`` would serve this module."""
+    if not is_enabled("shared_indexer_absorbed_query"):
+        return False
+    runner = getattr(mla, "_shared_indexer_query", None)
+    if runner is None:
+        from ..absorbed_query.shared_indexer import (
+            SharedIndexerProjection,
+            shared_module_eligibility,
+        )
+
+        ok, reason = shared_module_eligibility(mla)
+        log_capability("shared_indexer_absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+        runner = SharedIndexerProjection(mla) if ok else False
+        mla._shared_indexer_query = runner
+    return runner
+
+
+def _absorbed_query_runner(mla):
+    runner = getattr(mla, "_absorbed_query", None)
+    if runner is None:
+        ok, reason = module_eligibility(mla)
+        log_capability("absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+        runner = AbsorbedQueryProjection(mla) if ok else False
+        mla._absorbed_query = runner
+    return runner
+
+
+# Placement of the native k_norm in the deferred_norm full-layer path: "aux" (on the
+# indexer's aux stream next to the q_a norm when multi-stream is enabled and
+# absorbed_query's indexer_k_on_aux is set, joined before absorbed_query's kernel), "k_first" or
+# "q_first" (both on the current stream, in that order).
+K_NORM_MODES = ("aux", "k_first", "q_first")
+DEFAULT_K_NORM_MODE = "aux"
+
+
+def indexer_query_from_projected(
+    query_runner: AbsorbedQueryProjection,
+    kva: torch.Tensor,
+    kw: torch.Tensor,
+    position_ids: Optional[torch.Tensor],
+    absorb: bool,
+    k_norm_mode: Optional[str] = None,
+) -> AbsorbedQueryPacket:
+    """absorbed_query's fused full-layer kernel on projected ``kva`` and ``kw``.
+
+    The kernel consumes the native ``q_a_layernorm`` of the q_a columns of
+    ``kva`` and the native ``k_norm`` of the key columns of ``kw`` (absorbed_query's
+    bridges). Both are computed here from the deferred_norm projection's outputs with
+    the native calls absorbed_query's ``native_bridges`` makes; ``k_norm_mode`` (default
+    ``DEFAULT_K_NORM_MODE``) only places them in the stream(s), the
+    arithmetic is the same.
+    """
+    mla = query_runner.mla
+    indexer = mla.mqa.indexer
+    prenorm = query_runner.config.prenorm
+    mode = DEFAULT_K_NORM_MODE if k_norm_mode is None else k_norm_mode
+    assert mode in K_NORM_MODES, mode
+
+    def main():
+        return query_runner.native_q_norm(kva) if prenorm else None
+
+    def aux():
+        return indexer.k_norm(kw[:, : indexer.head_dim])
+
+    if mode == "aux" and query_runner.config.indexer_k_on_aux:
+        from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel
+
+        qn, k_ln = maybe_execute_in_parallel(
+            main, aux, query_runner.events[0], query_runner.events[1], indexer.aux_stream
+        )
+    elif mode == "k_first":
+        k_ln = aux()
+        qn = main()
+    else:
+        qn = main()
+        k_ln = aux()
+    return query_runner.launch_indexer(kva, kw, k_ln, None, 0, position_ids, absorb, qn)
+
+
+def deferred_norm_dsa_proj(
+    mla, position_ids: Optional[torch.Tensor], packet, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """The query_front packet of a target layer whose input is a deferred_norm packet.
+
+    Same outputs and side effects as absorbed_query's ``dsa_proj`` (or
+    shared_indexer_absorbed_query's ``shared_dsa_proj`` when ``shared_indexer_absorbed_query`` is
+    selected), with the fused-A and indexer wk|weights_proj outputs taken from the deferred_norm
+    projection.
+    """
+    if not (is_enabled("deferred_norm") and is_enabled("absorbed_query")):
+        return None
+    runner = _deferred_norm_runner(mla)
+    if runner is False:
+        return None
+    ok, _ = call_eligibility(mla, packet.residual, attn_metadata)
+    if not ok or packet.num_rows > runner_max_rows():
+        return None
+    indexer = mla.mqa.indexer
+    shared_query = _shared_query_runner(mla) if indexer is None else False
+    query_runner = _absorbed_query_runner(mla) if not shared_query else False
+    if not shared_query and not query_runner:
+        return None
+    if query_runner and query_runner.config.prefetch:
+        return None
+    if query_runner and indexer is not None and not query_runner.indexer_fusable():
+        fused_tail = False
+    else:
+        fused_tail = indexer is not None
+
+    if runner.config.prefetch_next and not runner.prefetch:
+        # Weights of the kernels that follow on this layer (L2 hint).
+        nxt = [mla.q_b_proj.weight, mla.k_b_proj_trans]
+        if indexer is not None:
+            nxt.insert(1, indexer.wq_b.weight)
+        runner.set_prefetch(nxt)
+    kva, kw = runner.project(packet)
+    if shared_query:
+        q, latent, fused_q = shared_query.launch(kva)
+        return AbsorbedQueryPacket([q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent], fused_q)
+    absorb = absorb_for(mla)
+    if fused_tail:
+        return indexer_query_from_projected(query_runner, kva, kw, position_ids, absorb)
+    q, qr, latent, fused_q = query_runner.launch(kva, need_qr=indexer is not None, absorb=absorb)
+    outs = [q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent]
+    if indexer is not None:
+        outs.extend(indexer_tail_from_projected(indexer, qr, kw, position_ids))
+    return AbsorbedQueryPacket(outs, fused_q)
+
+
+def runner_max_rows() -> int:
+    from .projection import MAX_ROWS
+
+    return MAX_ROWS
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/scope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/scope.py
new file mode 100644
index 000000000000..8d6995655cf8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/scope.py
@@ -0,0 +1,134 @@
+# 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.
+"""Which post-MoE boundaries the deferred-norm hop serves.
+
+The scopes (the deferred-norm numerical variant):
+
+* ``all``: every boundary ``L -> L+1`` whose next layer is a target attention
+  layer (the hop's unrestricted scope).
+* ``shared``: only boundaries whose next layer is a shared-indexer layer (it
+  reuses the previous full layer's top-k and has no indexer of its own). On a
+  boundary into a full-indexer layer the hop is not engaged, so the layer
+  keeps the post-MoE hop and the next layer receives its normalized tensor and
+  runs its native or absorbed-query projection path; the indexer
+  ``wk | weights_proj`` inputs of full layers then carry no deferred-norm
+  rounding change.
+
+The served scope is ``shared`` (``selected_scope``).
+
+A layer's indexer type is the one its attention module is built with:
+``sparse_attention_config.to_sparse_params(pretrained_config,
+layer_idx).is_full_indexer_layer`` (``MLA`` builds an ``Indexer`` only on
+full layers).
+"""
+
+import sys
+from typing import Optional, Tuple
+
+#: How the decline reasons name the scope (``<label>=<scope>: ...``).
+SCOPE_LABEL = "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE"
+# Scope names.
+SCOPES = ("all", "shared")
+# The scope that serves every boundary, without looking at the next layer.
+UNRESTRICTED_SCOPE = SCOPES[0]
+# The scope of this target.
+DEFAULT_SCOPE = "shared"
+# Calling frames searched for the decoder layer that constructs a hop.
+_MAX_FRAMES = 8
+
+
+def selected_scope() -> str:
+    """The served scope, ``DEFAULT_SCOPE``."""
+    return DEFAULT_SCOPE
+
+
+def is_full_indexer_layer(model_config, layer_idx: int) -> Optional[bool]:
+    """Whether target layer ``layer_idx`` owns an indexer (None: unknown).
+
+    Same call as ``MLA.__init__`` makes to decide whether the layer builds its
+    ``Indexer``.
+    """
+    sparse = getattr(model_config, "sparse_attention_config", None)
+    if sparse is None or not hasattr(sparse, "to_sparse_params"):
+        return None
+    params = sparse.to_sparse_params(
+        pretrained_config=getattr(model_config, "pretrained_config", None), layer_idx=layer_idx
+    )
+    full = getattr(params, "is_full_indexer_layer", None)
+    return None if full is None else bool(full)
+
+
+def constructing_layer(post_moe_hop):
+    """The decoder layer whose constructor is building a deferred_norm hop for ``post_moe_hop``.
+
+    Searches the calling frames for a ``self`` that holds ``post_moe_hop`` as its
+    ``post_moe_hop`` and has an integer ``layer_idx`` and a ``model_config``:
+    the layer assigns its post_moe_hop before it asks the deferred_norm factory for the
+    deferred hop. Returns ``None`` if no such frame is found within
+    ``_MAX_FRAMES`` callers.
+    """
+    frame = sys._getframe(1)
+    try:
+        for _ in range(_MAX_FRAMES):
+            if frame is None:
+                return None
+            owner = frame.f_locals.get("self")
+            if (
+                owner is not None
+                and getattr(owner, "post_moe_hop", None) is post_moe_hop
+                and isinstance(getattr(owner, "layer_idx", None), int)
+                and getattr(owner, "model_config", None) is not None
+            ):
+                return owner
+            frame = frame.f_back
+        return None
+    finally:
+        del frame
+
+
+def next_layer_kind(
+    post_moe_hop, next_layer_full_indexer: Optional[bool] = None
+) -> Tuple[Optional[bool], Optional[int], str]:
+    """``(full, next_layer_idx, how)`` of the layer after the hop's MoE layer.
+
+    ``next_layer_full_indexer`` given: that value (``how = "argument"``).
+    Otherwise it is resolved from the constructing decoder layer
+    (:func:`constructing_layer`). ``full`` is ``None`` when it cannot be
+    resolved; ``how`` then says why.
+    """
+    if next_layer_full_indexer is not None:
+        return bool(next_layer_full_indexer), None, "argument"
+    layer = constructing_layer(post_moe_hop)
+    if layer is None:
+        return None, None, "no constructing decoder layer found"
+    nxt = layer.layer_idx + 1
+    full = is_full_indexer_layer(layer.model_config, nxt)
+    if full is None:
+        return None, nxt, "the model config has no sparse-attention indexer layout"
+    return full, nxt, "model config"
+
+
+def decline_reason(
+    scope: str, full: Optional[bool], next_layer_idx: Optional[int], how: str
+) -> Optional[str]:
+    """Why ``scope`` keeps post_moe_hop's hop on this boundary (None: deferred_norm may serve it)."""
+    if scope == "all":
+        return None
+    where = f"layer {next_layer_idx}" if next_layer_idx is not None else "the next layer"
+    if full is None:
+        return f"{SCOPE_LABEL}={scope}: indexer type of {where} unknown ({how}); post_moe_hop hop"
+    if full:
+        return f"{SCOPE_LABEL}={scope}: {where} is a full-indexer layer; post_moe_hop hop"
+    return None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/__init__.py
new file mode 100644
index 000000000000..20a175b1daf7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/__init__.py
@@ -0,0 +1,62 @@
+# 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.
+"""deferred_norm_evict_hints: the deferred_norm hop + next-layer projection chain without the HBM collision.
+
+Requires ``absorbed_query``, ``post_moe_hop`` and ``deferred_norm``. deferred_norm_evict_hints keeps
+deferred_norm's deferred-norm contract (same packet, residual, ``materialize()`` and projection
+outputs, bit for bit) and changes only how the chain uses the memory system. By default:
+
+* the hop (``EvictHintsHopKernel``) is deferred_norm's hop launch unchanged;
+* the projection (``EvictHintsProjKernel``) stages its weights before its wait as
+  deferred_norm's does, with an ``L2::evict_first`` policy on the weight copies: the
+  next layer's 32 MB are read once per iteration, and the hint keeps them
+  from evicting the lines of the hop that runs under that stream.
+
+The other schedule knobs (``EvictHintsHopConfig``: L2 prefetch of the poll slice,
+evict_last clears, the shared-memory carveout that lets every projection
+CTA co-reside with the hop; ``EvictHintsProjConfig``: register caps, warp-parallel
+or producer-warp copy issue, split / deferred staging) are experiment arms.
+
+Registered entry points (``dispatch.get_impl("deferred_norm_evict_hints", name)``):
+
+* ``deferred_hop``: :class:`EvictHintsDeferredHop` factory (same arguments and
+  collective behaviour as deferred_norm's ``deferred_hop``);
+* ``dsa_proj(mla, position_ids, packet, attn_metadata)``: the query_front
+  packet of a target layer whose hidden state is a deferred_norm packet, through the
+  deferred_norm_evict_hints projection, or ``None`` to fall back;
+* ``is_deferred(hidden_states)``: deferred_norm's packet test.
+
+The hop and the projection launch with ``EvictHintsHopConfig()`` and
+``EvictHintsProjConfig()``.
+"""
+
+from ... import dispatch
+from ..deferred_norm.packet import is_deferred
+from .entry import evict_hints_dsa_proj
+from .hop import EvictHintsDeferredHop, EvictHintsHopConfig
+from .projection import EvictHintsProjConfig, EvictHintsProjection
+
+dispatch.register("deferred_norm_evict_hints", "deferred_hop", EvictHintsDeferredHop)
+dispatch.register("deferred_norm_evict_hints", "dsa_proj", evict_hints_dsa_proj)
+dispatch.register("deferred_norm_evict_hints", "is_deferred", is_deferred)
+
+__all__ = [
+    "EvictHintsDeferredHop",
+    "EvictHintsHopConfig",
+    "EvictHintsProjConfig",
+    "EvictHintsProjection",
+    "is_deferred",
+    "evict_hints_dsa_proj",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/entry.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/entry.py
new file mode 100644
index 000000000000..c703bb570a56
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/entry.py
@@ -0,0 +1,122 @@
+# 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.
+"""deferred_norm_evict_hints entry into the absorbed_query / shared_indexer_absorbed_query path for
+a deferred-norm packet.
+
+Same decisions, outputs and side effects as deferred_norm's ``deferred_norm_dsa_proj``
+(``deferred_norm/query_entry.py``), with the fused-A (and indexer) projection launched by
+``EvictHintsProjection`` (bit-identical to ``DeferredNormProjection``). It returns ``None``
+before any GPU work outside the support domain; the caller then
+materializes the packet.
+"""
+
+from typing import Optional
+
+import torch
+
+from ...dispatch import is_enabled, log_capability
+from ..absorbed_query.query_proj import (
+    _KV_LORA,
+    AbsorbedQueryPacket,
+    absorb_for,
+    call_eligibility,
+    check_package_selectors,
+    indexer_tail_from_projected,
+    module_eligibility,
+)
+from ..deferred_norm.projection import MAX_ROWS, projection_eligibility
+from ..deferred_norm.query_entry import (
+    _absorbed_query_runner,
+    _shared_query_runner,
+    indexer_query_from_projected,
+)
+
+
+def _evict_hints_runner(mla):
+    """Per-module deferred_norm_evict_hints projection (cached in ``mla._evict_hints_projection``) or False.
+
+    Static, rank-identical checks only (as deferred_norm's runner); the kernel was
+    compiled by the producing ``EvictHintsDeferredHop`` before its agreement.
+    """
+    runner = getattr(mla, "_evict_hints_projection", None)
+    if runner is not None:
+        return runner
+    from .projection import EvictHintsProjection
+
+    check_package_selectors()
+    ok, reason = module_eligibility(mla)
+    kw = None
+    if ok:
+        ok, reason, kw = projection_eligibility(mla)
+    runner = False
+    if ok:
+        proj = EvictHintsProjection(mla, kw_weight=kw)
+        why = proj.supported()
+        if why is not None:
+            ok, reason = False, why
+        else:
+            proj.precompile()
+            runner = proj
+    log_capability("deferred_norm_evict_hints", f"projection_layer{mla.layer_idx}", ok, reason)
+    mla._evict_hints_projection = runner
+    return runner
+
+
+def evict_hints_dsa_proj(
+    mla, position_ids: Optional[torch.Tensor], packet, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """The query_front packet of a target layer whose input is a deferred_norm packet
+    (deferred_norm_evict_hints projection)."""
+    if not (
+        is_enabled("deferred_norm_evict_hints")
+        and is_enabled("deferred_norm")
+        and is_enabled("absorbed_query")
+    ):
+        return None
+    runner = _evict_hints_runner(mla)
+    if runner is False:
+        return None
+    ok, _ = call_eligibility(mla, packet.residual, attn_metadata)
+    if not ok or packet.num_rows > MAX_ROWS:
+        return None
+    indexer = mla.mqa.indexer
+    shared_query = _shared_query_runner(mla) if indexer is None else False
+    query_runner = _absorbed_query_runner(mla) if not shared_query else False
+    if not shared_query and not query_runner:
+        return None
+    if query_runner and query_runner.config.prefetch:
+        return None
+    if query_runner and indexer is not None and not query_runner.indexer_fusable():
+        fused_tail = False
+    else:
+        fused_tail = indexer is not None
+
+    if runner.config.prefetch_next and not runner.prefetch:
+        nxt = [mla.q_b_proj.weight, mla.k_b_proj_trans]
+        if indexer is not None:
+            nxt.insert(1, indexer.wq_b.weight)
+        runner.set_prefetch(nxt)
+    kva, kw = runner.project(packet)
+    if shared_query:
+        q, latent, fused_q = shared_query.launch(kva)
+        return AbsorbedQueryPacket([q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent], fused_q)
+    absorb = absorb_for(mla)
+    if fused_tail:
+        return indexer_query_from_projected(query_runner, kva, kw, position_ids, absorb)
+    q, qr, latent, fused_q = query_runner.launch(kva, need_qr=indexer is not None, absorb=absorb)
+    outs = [q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent]
+    if indexer is not None:
+        outs.extend(indexer_tail_from_projected(indexer, qr, kw, position_ids))
+    return AbsorbedQueryPacket(outs, fused_q)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop.py
new file mode 100644
index 000000000000..5bb3c29d89e8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop.py
@@ -0,0 +1,239 @@
+# 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.
+"""deferred_norm_evict_hints deferred-norm post-MoE hop: deferred_norm's hop with
+deferred_norm_evict_hints's chain schedule.
+
+``EvictHintsDeferredHop`` is :class:`DeferredNormHop` (same eligibility, scope,
+contract checks, collective decisions, post_moe_hop workspace protocol, packet type and
+materialization) launching ``EvictHintsHopKernel`` with ``EvictHintsHopConfig``'s L2 hints
+and carveout. Its precompile covers the deferred_norm_evict_hints hop variants, deferred_norm's
+materialization and the served ``EvictHintsProjection`` kernels, so that a packet
+exists only where every rank can project it with deferred_norm_evict_hints's projection.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+
+from ... import dispatch
+from ..deferred_norm.hop import _ROUTED_INPUT_MODES, DeferredNormHop
+from ..deferred_norm.packet import DeferredNormHidden, frag_elems
+from ..post_moe_hop.kernels import STATE_WORDS
+from ..post_moe_hop.workspace import PostMoeHopWorkspace, _all_ranks_agree
+from .projection import EvictHintsProjConfig
+
+FEATURE_ID = "deferred_norm_evict_hints"
+# Poll-slice prefetch modes issued before the hop's dependency wait: they touch producer-owned
+# Lamport addresses before the wait, so only tests build them (``projection.parse_cfg``).
+PRE_WAIT_PREFETCH = ("pre3", "tpre3")
+
+
+@dataclass(frozen=True)
+class EvictHintsHopConfig:
+    """Launch configuration of the deferred_norm_evict_hints hop (``hop_kernel.py``).
+
+    ``ctas`` is the number of CTAs per token row (deferred_norm: 8); fewer CTAs leave
+    more SMs free for the projection that the hop releases early. The
+    square-sum partial order does not depend on it (``CTAS * WARPS = 24``
+    partials of 256 columns each, in column order).
+
+    The defaults launch deferred_norm's hop schedule unchanged. ``carveout=100`` with
+    ``ffi=False`` (the only launcher whose kernels the carveout reaches) lets
+    every projection CTA start next to a hop CTA, and its weight stream then
+    delays the exchange. The other values are for experiments only.
+    """
+
+    carveout: int = -1
+    poll_prefetch: str = "none"
+    clear_hint: str = "none"
+    ctas: int = 8
+    # TVM-FFI launcher (True) or the DSL executor (False).
+    ffi: bool = True
+    # Unused dynamic shared memory of the hop launch (``hop_kernel.MAX_DYN_SMEM`` at most).
+    dyn_smem: int = 0
+
+
+class EvictHintsDeferredHop(DeferredNormHop):
+    """deferred_norm's deferred hop launching the deferred_norm_evict_hints kernel (module docstring).
+
+    Args:
+        post_moe_hop, trigger_early, ctas, next_layer_full_indexer: as
+            :class:`DeferredNormHop`.
+        config: hop launch configuration (default: ``EvictHintsHopConfig()``).
+        proj_config: the deferred_norm_evict_hints projection configuration the consumer will
+            launch (default: ``EvictHintsProjConfig()``); compiled before
+            the cross-rank agreement.
+    """
+
+    def __init__(
+        self,
+        post_moe_hop,
+        trigger_early: bool = True,
+        ctas: Optional[int] = None,
+        next_layer_full_indexer: Optional[bool] = None,
+        config: Optional[EvictHintsHopConfig] = None,
+        proj_config: Optional[EvictHintsProjConfig] = None,
+    ):
+        self.hop_config = config or EvictHintsHopConfig()
+        self.evict_hints_proj = proj_config or EvictHintsProjConfig()
+        ctas = self.hop_config.ctas if ctas is None else ctas
+        super().__init__(
+            post_moe_hop,
+            trigger_early=trigger_early,
+            ctas=ctas,
+            carveout=self.hop_config.carveout,
+            next_layer_full_indexer=next_layer_full_indexer,
+        )
+        dispatch.log_capability(FEATURE_ID, "deferred_hop", self.engaged, self.reason)
+
+    def _compiled(
+        self,
+        has_shared: bool,
+        scale_fp32: bool,
+        routed_input: Optional[str],
+        debug_ts: bool = False,
+    ):
+        from .hop_kernel import compiled_evict_hints_hop
+
+        return compiled_evict_hints_hop(
+            self.hidden_size,
+            self.top_k,
+            self.mapping.tp_size,
+            self.post_moe_hop.max_rows,
+            self.ctas,
+            has_shared,
+            scale_fp32,
+            self.use_pdl,
+            self.trigger_early,
+            STATE_WORDS,
+            self.carveout,
+            debug_ts,
+            self.fault_delay_ns,
+            routed_input,
+            self.hop_config.poll_prefetch,
+            self.hop_config.clear_hint,
+            self.hop_config.ffi,
+            self.hop_config.dyn_smem,
+        )
+
+    def _precompile(self) -> Optional[str]:
+        from ..deferred_norm.hop_kernel import compiled_materialize
+        from .projection import precompile_served
+
+        for has_shared in (True, False):
+            for scale_fp32 in (False, True):
+                self._compiled(has_shared, scale_fp32, None)
+        compiled_materialize(self.hidden_size, self.ctas, self.partials, self.use_pdl)
+        return precompile_served(torch.cuda.current_device(), self.evict_hints_proj)
+
+    def enable_routed_input(self) -> bool:
+        """Collective; as :meth:`DeferredNormHop.enable_routed_input` with the deferred_norm_evict_hints kernels."""
+        mode = self._post_moe_routed_mode()
+        local_ok, reason = True, ""
+        if not self.engaged:
+            local_ok, reason = False, "deferred_norm_evict_hints hop not engaged"
+        elif mode is None:
+            local_ok, reason = False, "post_moe_hop hop's routed input is not enabled"
+        else:
+            try:
+                for has_shared in (True, False):
+                    self._compiled(has_shared, False, mode)
+            except Exception as exc:  # noqa: BLE001 - converged across ranks below
+                local_ok, reason = False, f"routed kernel compile failed: {exc!r}"
+        agreed = False
+        if self.engaged and mode is not None:
+            comm = PostMoeHopWorkspace.comm_for(self.mapping)
+            same_mode = [m for m in _ROUTED_INPUT_MODES if _all_ranks_agree(comm, mode == m)]
+            if not same_mode:
+                local_ok, reason = (
+                    False,
+                    "the post_moe_hop hops' routed-input modes differ across ranks",
+                )
+            agreed = _all_ranks_agree(comm, local_ok)
+            if not agreed:
+                reason = reason or "another rank is not capable"
+        self.routed_engaged = agreed
+        self.routed_mode = mode if agreed else None
+        self.routed_reason = reason
+        dispatch.log_capability(FEATURE_ID, f"deferred_hop_routed[{mode}]", agreed, reason)
+        return agreed
+
+    def _launch(
+        self,
+        fc2_ptr: int,
+        fc2_row_bytes: int,
+        fc2_rows: int,
+        scales_ptr: int,
+        perm_ptr: int,
+        scale_fp32: bool,
+        routed_input: Optional[str],
+        shared_output,
+        residual,
+        norm_weight,
+        eps,
+        ts,
+    ) -> Tuple[DeferredNormHidden, torch.Tensor]:
+        num_rows = residual.shape[0]
+        ws = self.post_moe_hop.workspace
+        a_frag = torch.empty(
+            (frag_elems(self.hidden_size),), dtype=residual.dtype, device=residual.device
+        )
+        residual_out = torch.empty_like(residual)
+        sumsq = torch.empty((num_rows, self.partials), dtype=torch.float32, device=residual.device)
+        has_shared = shared_output is not None
+        assert ws.state.numel() == STATE_WORDS
+        compiled = self._compiled(has_shared, scale_fp32, routed_input, ts is not None)
+        args = (
+            fc2_ptr,
+            fc2_row_bytes,
+            fc2_rows,
+            scales_ptr,
+            perm_ptr,
+            shared_output.data_ptr() if has_shared else 0,
+            residual.data_ptr(),
+            norm_weight.data_ptr(),
+            a_frag.data_ptr(),
+            residual_out.data_ptr(),
+            sumsq.data_ptr(),
+            ws.mc.data_ptr(),
+            ws.uc.data_ptr(),
+            ws.partials.data_ptr(),
+            0 if ts is None else ts.data_ptr(),
+        )
+        if self.hop_config.ffi:
+            compiled(*args, ws.state, self.mapping.tp_rank, self.timeout_ns, num_rows)
+        else:
+            from cuda.bindings.driver import CUstream
+            from cutlass import Int32, Int64
+
+            ints = (Int64, Int64, Int32) + (Int64,) * 12
+            compiled(
+                *[t(v) for t, v in zip(ints, args)],
+                Int64(ws.state.data_ptr()),
+                Int32(self.mapping.tp_rank),
+                Int64(self.timeout_ns),
+                Int32(num_rows),
+                CUstream(torch.cuda.current_stream().cuda_stream),
+            )
+        packet = DeferredNormHidden(
+            a_frag=a_frag,
+            sumsq=sumsq,
+            residual=residual_out,
+            gamma=norm_weight,
+            eps=float(eps),
+            materializer=self.materialize,
+        )
+        return packet, residual_out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop_kernel.py
new file mode 100644
index 000000000000..cddc5a489a09
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/hop_kernel.py
@@ -0,0 +1,814 @@
+# 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.
+"""CuTe DSL kernel of the deferred_norm_evict_hints deferred-norm post-MoE hop.
+
+``EvictHintsHopKernel`` is deferred_norm's ``DeferredNormHopKernel``
+(``deferred_norm/hop_kernel.py``) with the same instructions for steps 1-7 on the same post_moe_hop
+workspace (so its packet, residual and materialization are bit-identical to deferred_norm's), plus
+L2 cache hints that keep the hop's Lamport lines out of the next layer's weight
+stream:
+
+* ``poll_prefetch``: ``cp.async.bulk.prefetch.L2`` of the CTA's poll slice
+  (8 rank slots x ``threads * 16`` bytes of its row):
+
+  - ``"cur"``: the current buffer, after the state words are loaded;
+  - ``"wait3"``: the slice of all three Lamport buffers right after the
+    dependency wait (no state needed), before the PDL trigger;
+  - ``"pre3"``: the same before the dependency wait (exception arm: the
+    addresses are producer-owned; a prefetch returns no value);
+  - ``"tcur"`` / ``"t3"`` / ``"tpre3"``: the same three placements with a
+    per-thread ``prefetch.global.L2`` of the thread's own 8 poll addresses
+    (per buffer) instead of the bulk prefetch, whose uniform operands make
+    the compiler serialize one issue per thread in warp 0;
+
+* ``clear_hint="evict_last"``: the clear of the previous call's buffer
+  (which is the buffer of call n + 2) stores with an ``L2::evict_last``
+  policy;
+* ``carveout``: the preferred shared-memory carveout of the launch, so an
+  SM that holds a hop CTA can also take a projection CTA;
+* ``dyn_smem``: an unused dynamic shared-memory allocation of the launch
+  (at most ``MAX_DYN_SMEM`` bytes, so that a projection CTA still fits next
+  to the hop CTA in the 228 KB configuration), for drivers that apply the
+  carveout only to a kernel that uses shared memory.
+
+None of these changes a value or the order of the hop's memory operations.
+``debug_ts`` adds the CTA's ``%smid`` in stamp slot 5.
+
+The carveout is a CUDA kernel attribute. The DSL records
+``preferred_smem_carveout`` but its CUDA-library loader does not apply it,
+so a DSL kernel runs with the driver's default configuration (32 KB of
+shared memory on the hop's SMs; a projection CTA, which needs the 228 KB
+configuration, cannot start there until the hop CTA exits).
+``compiled_evict_hints_hop`` therefore sets ``CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_
+MEMORY_CARVEOUT`` itself on every kernel of the compiled function's CUDA
+library for the current device (``set_library_carveout``).
+``ffi=False`` compiles ``call_raw`` (the state words and the stream passed
+explicitly) for the DSL executor instead of the TVM-FFI launcher.
+"""
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Int64
+
+from ..deferred_norm.hop_kernel import (
+    ROUTED_MODES,
+    DeferredNormHopKernel,
+    _hop_stamp,
+    _st_b32,
+    _st_f32_if,
+)
+from ..deferred_norm.ptx import (
+    FP32_SENTINEL,
+    SENTINEL_WORD,
+    add,
+    arrive,
+    bits_f32,
+    error_word_address,
+    fault_delay,
+    flag_error,
+    flag_timeout,
+    fma,
+    ld_v4,
+    load_state,
+    mul,
+    pack_bf16x2,
+    perm_check,
+    poison,
+    poll8,
+    round_bf16,
+    sanitize,
+    select3,
+    select_f32,
+    shfl_xor,
+    st_v4,
+    unpack_bf16x2,
+)
+from ..post_moe_hop.kernels import ERROR_PERM_RANGE, MAX_CHUNKS
+from .ptx import policy_evict_last, prefetch_l2, prefetch_line_l2, smid, st_b64, st_v4_hint
+
+POLL_PREFETCH_MODES = ("none", "cur", "wait3", "pre3", "tcur", "t3", "tpre3")
+CLEAR_HINTS = ("none", "evict_last")
+# 228 KB configuration - (projection CTA 222,080 B + 1 KB reserved) - 1 KB reserved.
+MAX_DYN_SMEM = 233472 - (222080 + 1024) - 1024
+
+
+@cute.jit
+def _smid_stamp(enabled: cutlass.Constexpr[bool], ts: Int64, sid: Int32, tid: Int32):
+    if cutlass.const_expr(enabled):
+        if tid == Int32(0):
+            st_b64(ts + Int64((sid * Int32(8) + Int32(5)) * Int32(8)), Int64(smid()))
+
+
+class EvictHintsHopKernel(DeferredNormHopKernel):
+    """deferred_norm's deferred hop with deferred_norm_evict_hints's L2 hints (module docstring)."""
+
+    def __init__(
+        self,
+        hidden: int,
+        top_k: int,
+        world: int,
+        max_rows: int,
+        ctas: int,
+        has_shared: bool,
+        scale_fp32: bool,
+        use_pdl: bool,
+        trigger_early: bool,
+        carveout: int = -1,
+        debug_ts: bool = False,
+        fault_delay_ns: int = 0,
+        routed_input: Optional[str] = None,
+        poll_prefetch: str = "none",
+        clear_hint: str = "none",
+        dyn_smem: int = 0,
+    ):
+        super().__init__(
+            hidden,
+            top_k,
+            world,
+            max_rows,
+            ctas,
+            has_shared,
+            scale_fp32,
+            use_pdl,
+            trigger_early,
+            carveout,
+            debug_ts,
+            fault_delay_ns,
+            routed_input,
+        )
+        assert poll_prefetch in POLL_PREFETCH_MODES, poll_prefetch
+        assert clear_hint in CLEAR_HINTS, clear_hint
+        assert 0 <= dyn_smem <= MAX_DYN_SMEM and dyn_smem % 16 == 0, dyn_smem
+        self.dyn_smem = dyn_smem
+        self.poll_prefetch = poll_prefetch
+        self.clear_hint = clear_hint
+
+    @cute.jit
+    def _prefetch_poll(self, uc: Int64, buf: Int32, row: Int32, cta: Int32, slot: Int32):
+        """L2 prefetch of rank slot ``slot`` of this CTA's columns in buffer ``buf``."""
+        hidden = self.hidden
+        base = (
+            uc
+            + Int64(buf) * Int64(self.max_rows * self.world * hidden * 2)
+            + Int64(row) * Int64(self.world * hidden * 2)
+            + Int64(slot) * Int64(hidden * 2)
+            + Int64(cta) * Int64(self.threads * 16)
+        )
+        prefetch_l2(base, Int32(self.threads * 16))
+
+    @cute.jit
+    def _prefetch_own(self, uc: Int64, buf: Int32, row: Int32, cta: Int32, tidx: Int32):
+        """Per-thread L2 prefetch of this thread's 8 poll addresses in buffer ``buf``."""
+        hidden = self.hidden
+        base = (
+            uc
+            + Int64(buf) * Int64(self.max_rows * self.world * hidden * 2)
+            + Int64(row) * Int64(self.world * hidden * 2)
+            + Int64((cta * Int32(self.threads) + tidx) * Int32(16))
+        )
+        for r in cutlass.range_constexpr(8):
+            prefetch_line_l2(base + Int64(r * hidden * 2))
+
+    @cute.jit
+    def _prefetch_own3(self, uc: Int64, row: Int32, cta: Int32, tidx: Int32):
+        for b in cutlass.range_constexpr(3):
+            self._prefetch_own(uc, Int32(b), row, cta, tidx)
+
+    @cute.jit
+    def _prefetch_three(self, uc: Int64, row: Int32, cta: Int32, tidx: Int32):
+        """Threads 0..23: slot ``tidx % 8`` of buffer ``tidx // 8``."""
+        if tidx < Int32(24):
+            self._prefetch_poll(uc, tidx // Int32(8), row, cta, tidx % Int32(8))
+
+    @cute.jit
+    def __call__(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        scaled: Int64,
+        residual_out: Int64,
+        sumsq: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        ts: Int64,
+        state: cute.Tensor,
+        rank: Int32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+        stream: CUstream,
+    ):
+        """TVM-FFI entry (``state`` as a GPU tensor binds the caller's stream)."""
+        self._launch_kernel(
+            fc2,
+            fc2_row_bytes,
+            fc2_rows,
+            scales,
+            perm,
+            shared,
+            residual,
+            gamma,
+            scaled,
+            residual_out,
+            sumsq,
+            mc,
+            uc,
+            partials,
+            ts,
+            Int64(state.iterator.toint()),
+            rank,
+            timeout_ns,
+            num_rows,
+            stream,
+        )
+
+    @cute.jit
+    def call_raw(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        scaled: Int64,
+        residual_out: Int64,
+        sumsq: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        ts: Int64,
+        state: Int64,
+        rank: Int32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+        stream: CUstream,
+    ):
+        """Executor entry (state-word address and stream explicit)."""
+        self._launch_kernel(
+            fc2,
+            fc2_row_bytes,
+            fc2_rows,
+            scales,
+            perm,
+            shared,
+            residual,
+            gamma,
+            scaled,
+            residual_out,
+            sumsq,
+            mc,
+            uc,
+            partials,
+            ts,
+            state,
+            rank,
+            timeout_ns,
+            num_rows,
+            stream,
+        )
+
+    @cute.jit
+    def _launch_kernel(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        scaled: Int64,
+        residual_out: Int64,
+        sumsq: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        ts: Int64,
+        state: Int64,
+        rank: Int32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+        stream: CUstream,
+    ):
+        if cutlass.const_expr(self.carveout >= 0):
+            self.kernel(
+                fc2,
+                fc2_row_bytes,
+                fc2_rows,
+                scales,
+                perm,
+                shared,
+                residual,
+                gamma,
+                scaled,
+                residual_out,
+                sumsq,
+                mc,
+                uc,
+                partials,
+                ts,
+                state,
+                rank,
+                timeout_ns,
+                num_rows,
+            ).launch(
+                grid=[self.ctas, num_rows, 1],
+                block=[self.threads, 1, 1],
+                smem=self.dyn_smem if self.dyn_smem else None,
+                stream=stream,
+                use_pdl=self.use_pdl,
+                preferred_smem_carveout=self.carveout,
+            )
+        else:
+            self.kernel(
+                fc2,
+                fc2_row_bytes,
+                fc2_rows,
+                scales,
+                perm,
+                shared,
+                residual,
+                gamma,
+                scaled,
+                residual_out,
+                sumsq,
+                mc,
+                uc,
+                partials,
+                ts,
+                state,
+                rank,
+                timeout_ns,
+                num_rows,
+            ).launch(
+                grid=[self.ctas, num_rows, 1],
+                block=[self.threads, 1, 1],
+                smem=self.dyn_smem if self.dyn_smem else None,
+                stream=stream,
+                use_pdl=self.use_pdl,
+            )
+
+    @cute.kernel
+    def kernel(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        scaled: Int64,
+        residual_out: Int64,
+        sumsq: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        ts: Int64,
+        state: Int64,
+        rank: Int32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+    ):
+        hidden = self.hidden
+        world = self.world
+        threads = self.threads
+        warps = self.warps
+        ctas = self.ctas
+        buffer_bytes = self.max_rows * world * hidden * 2
+
+        tidx, _, _ = cute.arch.thread_idx()
+        cta, row, _ = cute.arch.block_idx()
+        lane = cute.arch.lane_idx()
+        warp = cute.arch.warp_idx()
+
+        sid = row * Int32(ctas) + cta
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 0)
+        _smid_stamp(self.debug_ts, ts, sid, tidx)
+        if cutlass.const_expr(self.poll_prefetch == "pre3"):
+            self._prefetch_three(uc, row, cta, tidx)
+        if cutlass.const_expr(self.poll_prefetch == "tpre3"):
+            self._prefetch_own3(uc, row, cta, tidx)
+
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+            if cutlass.const_expr(self.poll_prefetch == "wait3"):
+                self._prefetch_three(uc, row, cta, tidx)
+            if cutlass.const_expr(self.poll_prefetch == "t3"):
+                self._prefetch_own3(uc, row, cta, tidx)
+            if cutlass.const_expr(self.trigger_early):
+                cute.arch.griddepcontrol_launch_dependents()
+        else:
+            if cutlass.const_expr(self.poll_prefetch == "wait3"):
+                self._prefetch_three(uc, row, cta, tidx)
+            if cutlass.const_expr(self.poll_prefetch == "t3"):
+                self._prefetch_own3(uc, row, cta, tidx)
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 1)
+
+        # State words of the previous call (complete at the kernel boundary).
+        epoch, prev_rows, err_words = load_state(state)
+        cur = epoch % Int32(3)
+        dirty = (epoch + Int32(2)) % Int32(3)
+        if cutlass.const_expr(self.poll_prefetch == "cur"):
+            if tidx < Int32(8):
+                self._prefetch_poll(uc, cur, row, cta, tidx)
+        if cutlass.const_expr(self.poll_prefetch == "tcur"):
+            self._prefetch_own(uc, cur, row, cta, tidx)
+        # Bits of the previous call's buffer: when set, this call publishes
+        # NaN and carries them into its own error word.
+        sticky = select3(err_words, dirty)
+        err_cur = error_word_address(state, cur)
+
+        col = (cta * Int32(threads) + tidx) * Int32(8)
+        row_off = Int64(row) * Int64(hidden * 2) + Int64(col) * Int64(2)
+        res_w = ld_v4(residual + row_off)
+        gamma_w = ld_v4(gamma + Int64(col) * Int64(2))
+        if cutlass.const_expr(self.has_shared):
+            sh_w = ld_v4(shared + row_off)
+
+        # 1. Finalize (fp32 FMA chain in k order), bf16, + shared, sanitize.
+        if cutlass.const_expr(self.routed_input == "row"):
+            r_w = ld_v4(fc2 + Int64(row) * fc2_row_bytes + Int64(col) * Int64(2))
+            local = [r_w[j] for j in range(4)]
+            out_of_range = Int32(0)
+        elif cutlass.const_expr(self.routed_input == "fp32"):
+            a_addr = fc2 + Int64(row) * fc2_row_bytes + Int64(col) * Int64(4)
+            a_w = ld_v4(a_addr) + ld_v4(a_addr + Int64(16))
+            a_f = [bits_f32(w) for w in a_w]
+            local = [pack_bf16x2(a_f[2 * j], a_f[2 * j + 1]) for j in range(4)]
+            zero = Int32(0)
+            st_v4(a_addr, zero, zero, zero, zero)
+            st_v4(a_addr + Int64(16), zero, zero, zero, zero)
+            out_of_range = Int32(0)
+        elif cutlass.const_expr(self.routed_input == "identity"):
+            s_w = ld_v4(scales + Int64(row) * Int64(16))
+            weights = []
+            for w in s_w:
+                lo, hi = unpack_bf16x2(w)
+                weights += [lo, hi]
+            acc = [Float32(0.0)] * 8
+            for k in cutlass.range_constexpr(8):
+                x_w = ld_v4(
+                    fc2 + (Int64(row) * Int64(8) + Int64(k)) * fc2_row_bytes + Int64(col) * Int64(2)
+                )
+                for j in cutlass.range_constexpr(4):
+                    x_lo, x_hi = unpack_bf16x2(x_w[j])
+                    acc[2 * j] = fma(weights[k], x_lo, acc[2 * j])
+                    acc[2 * j + 1] = fma(weights[k], x_hi, acc[2 * j + 1])
+            local = [pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)]
+            out_of_range = Int32(0)
+        else:
+            perm_lo = ld_v4(perm + Int64(row) * Int64(32))
+            perm_hi = ld_v4(perm + Int64(row) * Int64(32) + Int64(16))
+            perm_k = perm_lo + perm_hi
+            if cutlass.const_expr(self.scale_fp32):
+                s_lo = ld_v4(scales + Int64(row) * Int64(32))
+                s_hi = ld_v4(scales + Int64(row) * Int64(32) + Int64(16))
+                weights = [bits_f32(w) for w in s_lo + s_hi]
+            else:
+                s_w = ld_v4(scales + Int64(row) * Int64(16))
+                weights = []
+                for w in s_w:
+                    lo, hi = unpack_bf16x2(w)
+                    weights += [lo, hi]
+            acc = [Float32(0.0)] * 8
+            out_of_range = Int32(0)
+            for k in cutlass.range_constexpr(8):
+                valid, safe, oob = perm_check(perm_k[k], fc2_rows)
+                out_of_range = out_of_range + oob
+                x_w = ld_v4(fc2 + Int64(safe) * fc2_row_bytes + Int64(col) * Int64(2))
+                for j in cutlass.range_constexpr(4):
+                    x_lo, x_hi = unpack_bf16x2(x_w[j])
+                    acc[2 * j] = select_f32(valid, fma(weights[k], x_lo, acc[2 * j]), acc[2 * j])
+                    acc[2 * j + 1] = select_f32(
+                        valid, fma(weights[k], x_hi, acc[2 * j + 1]), acc[2 * j + 1]
+                    )
+            local = [pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)]
+        if cutlass.const_expr(self.has_shared):
+            summed = []
+            for j in cutlass.range_constexpr(4):
+                r_lo, r_hi = unpack_bf16x2(local[j])
+                s_lo_f, s_hi_f = unpack_bf16x2(sh_w[j])
+                summed.append(pack_bf16x2(add(s_lo_f, r_lo), add(s_hi_f, r_hi)))
+            local = summed
+        local = [poison(sanitize(w), sticky) for w in local]
+
+        # 2. Publish through multicast into slot [row][rank].
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4),
+                self.fault_delay_ns,
+                0,
+            )
+        slot = (
+            Int64(cur) * Int64(buffer_bytes)
+            + Int64(row) * Int64(world * hidden * 2)
+            + Int64(rank) * Int64(hidden * 2)
+            + Int64(col) * Int64(2)
+        )
+        st_v4(mc + slot, local[0], local[1], local[2], local[3])
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 2)
+        flag_error(err_cur, out_of_range, Int32(ERROR_PERM_RANGE))
+        flag_error(err_cur, sticky, sticky)
+
+        # 3. Clear the previous call's Lamport rows and partial slice.
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4) + Int32(1),
+                self.fault_delay_ns,
+                1,
+            )
+        total_threads = num_rows * Int32(ctas * threads)
+        gtid = (row * Int32(ctas) + cta) * Int32(threads) + tidx
+        clear_vecs = prev_rows * Int32(world * hidden // 8)
+        sentinel = Int32(SENTINEL_WORD)
+        dirty_base = uc + Int64(dirty) * Int64(buffer_bytes)
+        if cutlass.const_expr(self.clear_hint == "evict_last"):
+            policy = policy_evict_last()
+            for v in cutlass.range(gtid, clear_vecs, total_threads):
+                st_v4_hint(
+                    dirty_base + Int64(v) * Int64(16),
+                    sentinel,
+                    sentinel,
+                    sentinel,
+                    sentinel,
+                    policy,
+                )
+        else:
+            for v in cutlass.range(gtid, clear_vecs, total_threads):
+                st_v4(dirty_base + Int64(v) * Int64(16), sentinel, sentinel, sentinel, sentinel)
+        part_vecs = prev_rows * Int32(MAX_CHUNKS // 4)
+        part_base = partials + Int64(dirty) * Int64(self.max_rows * MAX_CHUNKS * 4)
+        fp32_sentinel = Int32(FP32_SENTINEL)
+        for v in cutlass.range(gtid, part_vecs, total_threads):
+            st_v4(
+                part_base + Int64(v) * Int64(16),
+                fp32_sentinel,
+                fp32_sentinel,
+                fp32_sentinel,
+                fp32_sentinel,
+            )
+
+        # 4. Per-thread Lamport poll (timer deadline), rank-ordered fp32 sum.
+        poll_base = (
+            uc
+            + Int64(cur) * Int64(buffer_bytes)
+            + Int64(row) * Int64(world * hidden * 2)
+            + Int64(col) * Int64(2)
+        )
+        polled = poll8(poll_base, Int64(hidden * 2), timeout_ns)
+        timed_out = polled[32]
+        flag_timeout(err_cur, state, timed_out, polled[33])
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 3)
+        reduced = [Float32(0.0)] * 8
+        for r in cutlass.range_constexpr(8):
+            for j in cutlass.range_constexpr(4):
+                v_lo, v_hi = unpack_bf16x2(polled[4 * r + j])
+                reduced[2 * j] = add(reduced[2 * j], v_lo)
+                reduced[2 * j + 1] = add(reduced[2 * j + 1], v_hi)
+        if cutlass.const_expr(self.use_pdl and not self.trigger_early):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        # 5. Residual added once, after the reduction; NaN if this poll timed out.
+        new_res = []
+        x = []
+        for j in cutlass.range_constexpr(4):
+            rb = pack_bf16x2(reduced[2 * j], reduced[2 * j + 1])
+            a_lo, a_hi = unpack_bf16x2(rb)
+            r_lo, r_hi = unpack_bf16x2(res_w[j])
+            word = poison(pack_bf16x2(add(a_lo, r_lo), add(a_hi, r_hi)), timed_out)
+            new_res.append(word)
+            n_lo, n_hi = unpack_bf16x2(word)
+            x += [n_lo, n_hi]
+        st_v4(residual_out + row_off, new_res[0], new_res[1], new_res[2], new_res[3])
+
+        # 6. Deferred RMSNorm: one fp32 partial per warp; gain-folded activation.
+        thread_sum = Float32(0.0)
+        for i in cutlass.range_constexpr(8):
+            thread_sum = add(thread_sum, round_bf16(mul(x[i], x[i])))
+        for mask in (16, 8, 4, 2, 1):
+            thread_sum = add(thread_sum, shfl_xor(thread_sum, mask))
+        slot_idx = Int64(row) * Int64(ctas * warps) + Int64(cta * Int32(warps) + warp)
+        _st_f32_if(sumsq + slot_idx * Int64(4), thread_sum, Int32(lane == Int32(0)))
+        blk = col // Int32(16)
+        half = (col // Int32(8)) % Int32(2)
+        for j in cutlass.range_constexpr(4):
+            g_lo, g_hi = unpack_bf16x2(gamma_w[j])
+            word = pack_bf16x2(mul(x[2 * j], g_lo), mul(x[2 * j + 1], g_hi))
+            lane_slot = blk * Int32(32) + row * Int32(4) + Int32(j)
+            _st_b32(scaled + Int64(lane_slot * Int32(2) + half) * Int64(4), word)
+
+        # 7. Arrival; the last CTA publishes the next epoch.
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4) + Int32(2),
+                self.fault_delay_ns,
+                2,
+            )
+        cute.arch.sync_threads()
+        arrive(
+            state,
+            Int32(tidx == Int32(0)),
+            num_rows * Int32(ctas),
+            num_rows,
+            (epoch + Int32(1)) % Int32(3),
+        )
+        _hop_stamp(self.debug_ts, ts, sid, tidx, 4)
+
+
+_COMPILED = {}
+
+
+def compiled_evict_hints_hop(
+    hidden: int,
+    top_k: int,
+    world: int,
+    max_rows: int,
+    ctas: int,
+    has_shared: bool,
+    scale_fp32: bool,
+    use_pdl: bool,
+    trigger_early: bool,
+    state_words: int,
+    carveout: int = -1,
+    debug_ts: bool = False,
+    fault_delay_ns: int = 0,
+    routed_input: Optional[str] = None,
+    poll_prefetch: str = "none",
+    clear_hint: str = "none",
+    ffi: bool = True,
+    dyn_smem: int = 0,
+):
+    """Return the compiled deferred_norm_evict_hints hop for one shape (compiles on first use, no GPU work).
+
+    ``carveout >= 0`` is set on the compiled kernels (module docstring).
+    ``ffi=False`` compiles ``call_raw`` for the DSL's executor; call it with
+    ``(..., state_address, rank, timeout_ns, num_rows, CUstream)``.
+    """
+    assert routed_input in ROUTED_MODES, routed_input
+    key = (
+        "evict_hints_hop",
+        hidden,
+        top_k,
+        world,
+        max_rows,
+        ctas,
+        has_shared,
+        scale_fp32,
+        use_pdl,
+        trigger_early,
+        state_words,
+        carveout,
+        debug_ts,
+        fault_delay_ns,
+        routed_input,
+        poll_prefetch,
+        clear_hint,
+        ffi,
+        dyn_smem,
+    )
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = EvictHintsHopKernel(
+            hidden,
+            top_k,
+            world,
+            max_rows,
+            ctas,
+            has_shared,
+            scale_fp32,
+            use_pdl,
+            trigger_early,
+            carveout,
+            debug_ts,
+            fault_delay_ns,
+            routed_input,
+            poll_prefetch,
+            clear_hint,
+            dyn_smem,
+        )
+        zero64, zero32 = Int64(0), Int32(0)
+        if ffi:
+            state_fake = cute_runtime.make_fake_compact_tensor(
+                Int32, (state_words,), stride_order=(0,), assumed_align=16
+            )
+            compiled = cute.compile(
+                kernel,
+                zero64,
+                zero64,
+                zero32,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                state_fake,
+                zero32,
+                zero64,
+                zero32,
+                stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+                options="--enable-tvm-ffi",
+            )
+        else:
+            import torch
+
+            stream = CUstream(torch.cuda.current_stream().cuda_stream)
+            compiled = cute.compile(
+                kernel.call_raw,
+                zero64,
+                zero64,
+                zero32,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero64,
+                zero32,
+                zero64,
+                zero32,
+                stream,
+            )
+        if carveout >= 0:
+            LIBRARY_CARVEOUT[key] = set_library_carveout(compiled, carveout)
+        _COMPILED[key] = compiled
+    return compiled
+
+
+# compile key -> number of kernels whose carveout was set, or the reason it was not.
+LIBRARY_CARVEOUT = {}
+
+
+def set_library_carveout(compiled, carveout: int):
+    """Set the preferred shared-memory carveout of every kernel in the CUDA
+    library of a compiled DSL function, for the current device.
+
+    Returns the number of kernels set, or a string with the reason the
+    library could not be reached (the kernels then keep the driver default).
+    """
+    import torch
+    from cuda.bindings import driver as cu
+
+    try:
+        library = compiled.library
+    except (AttributeError, RuntimeError) as exc:
+        return f"no CUDA library: {exc!r}"[:200]
+    lib = cu.CUlibrary(int(library))
+    err, count = cu.cuLibraryGetKernelCount(lib)
+    assert err == cu.CUresult.CUDA_SUCCESS, err
+    err, kernels = cu.cuLibraryEnumerateKernels(count, lib)
+    assert err == cu.CUresult.CUDA_SUCCESS, err
+    device = cu.CUdevice(torch.cuda.current_device())
+    attr = cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT
+    for kernel in kernels:
+        (err,) = cu.cuKernelSetAttribute(attr, carveout, kernel, device)
+        assert err == cu.CUresult.CUDA_SUCCESS, err
+    return len(kernels)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/proj_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/proj_kernel.py
new file mode 100644
index 000000000000..a8b7b54ebbde
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/proj_kernel.py
@@ -0,0 +1,703 @@
+# 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.
+"""CuTe DSL kernel of the deferred_norm_evict_hints next-layer projection.
+
+``EvictHintsProjKernel`` is deferred_norm's staged projection (``deferred_norm/proj_kernel.py``)
+with the same per-warp MMA chains, warp-order reduction and ``r`` scaling (so its
+outputs are bit-identical to deferred_norm's), and three schedule knobs:
+
+* ``maxrreg`` / ``maxntid``: a register cap (``ptxas --maxrregcount``, or a
+  ``maxntid`` launch bound) so that a projection CTA fits in the register
+  file next to a hop CTA: the projection then stages its weights on every
+  SM while the hop exchanges, instead of starting on the hop's SMs only when
+  the hop exits;
+* ``copy_hint="evict_first"``: the weight bulk copies carry an
+  ``L2::evict_first`` policy (the weights are read once per iteration), so
+  they do not push the hop's Lamport lines out of L2;
+* ``split_chunk`` / ``split_delay_ns``: thread 0 issues the copies of K
+  chunks ``>= split_chunk`` only ``split_delay_ns`` after its first copy
+  (a ``%globaltimer`` spin; the copies still read only immutable weights);
+* ``copy_warps``: lane 0 of warp ``w`` issues the copies of rows ``w``,
+  ``w + NW``, ... (chunk-major) instead of thread 0 issuing all of them.
+  A bulk copy takes uniform operands, so one thread issues its copies one
+  at a time; 16 issuing warps stage a CTA's rows at the copy engine's rate.
+  Thread 0 still posts every chunk's byte count (the mbarrier transaction
+  count may go below zero transiently; the phase completes only after the
+  single arrival and all bytes);
+* ``prod_warp``: the CTA gets a 17th warp whose lane 0 issues every bulk
+  copy (with ``split_chunk`` / ``split_delay_ns``: the chunks below
+  ``split_chunk`` at entry, the rest ``split_delay_ns`` after its first
+  copy) and then exits. The 16 MMA warps issue nothing, so a deferred
+  chunk never holds an MMA warp before its dependency wait; they run deferred_norm's
+  body (same K steps, MMA chains, warp-order reduction and ``r``) with the
+  two post-loop CTA barriers over the 512 MMA threads only (``bar.sync 1,
+  512``). The producer reads only immutable weights and never waits on the
+  preceding grid; every chunk's mbarrier is waited on by every MMA warp, so
+  the CTA outlives every copy. ``defer_tail`` defers every chunk of the
+  last ``defer_tail * R`` CTAs, and ``defer_wait`` issues deferred copies
+  after the producer's own ``griddepcontrol.wait`` instead of the timer.
+  Debug stamps: slot 7 = first-phase copies issued, slot 6 = every copy
+  issued (producer lane 0).
+
+Every change is to the copy issue, the barrier membership or the compile;
+every value comes from deferred_norm's instruction sequence.
+"""
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Int64
+from cutlass.utils.smem_allocator import SmemAllocator
+
+from ..deferred_norm.proj_kernel import (
+    MAX_TOKENS,
+    DeferredNormProjKernel,
+    _bulk_g2s,
+    _ld_b32,
+    _ld_shared_f32,
+    _ld_v2_if,
+    _ldsm_x2,
+    _mbar_expect_tx,
+    _mbar_fence_init,
+    _mbar_init,
+    _mbar_wait,
+    _min_i32,
+    _prefetch_l2,
+    _row_src,
+    _sel_i32,
+    _sel_i64,
+    _shfl_idx_f32,
+    _st_shared_f32,
+    _stamp,
+    _store_out,
+)
+from ..deferred_norm.ptx import add, bits_f32, inv_rms, ldsm_x4, mma_bf16, mul, saddr, st_shared_v4
+from .ptx import bar_sync_named, bulk_g2s_hint, globaltimer, policy_evict_first, spin_until, st_b64
+
+COPY_HINTS = ("none", "evict_first")
+
+
+class EvictHintsProjKernel(DeferredNormProjKernel):
+    """deferred_norm's staged projection with deferred_norm_evict_hints's copy-issue knobs (module docstring)."""
+
+    def __init__(
+        self,
+        maxrreg: int = 0,
+        maxntid: int = 0,
+        copy_hint: str = "none",
+        split_chunk: int = 0,
+        split_delay_ns: int = 0,
+        copy_warps: bool = False,
+        prod_warp: bool = False,
+        defer_tail: int = 0,
+        defer_wait: bool = False,
+        **deferred_norm_cfg,
+    ):
+        super().__init__(**deferred_norm_cfg)
+        assert copy_hint in COPY_HINTS, copy_hint
+        if copy_hint != "none" or split_chunk or copy_warps or prod_warp:
+            assert self.copy_method == "bulk" and self.copy_threads == 1, (
+                "copy hints, split staging and warp-parallel copies need the bulk copy path"
+            )
+        assert not (copy_warps and split_chunk), "split staging issues from thread 0 only"
+        assert prod_warp or not (defer_tail or defer_wait), "deferral needs the producer warp"
+        if prod_warp:
+            assert not copy_warps, "the producer warp issues every copy"
+            assert self.stage_before_wait and not self.wait_staged_first
+            assert self.overflow == 0, "the producer warp needs every row staged"
+        self.copy_warps = copy_warps
+        self.prod_warp = prod_warp
+        self.defer_tail = defer_tail
+        self.defer_wait = defer_wait
+        # Threads per CTA: the MMA warps plus the producer warp.
+        self.nt_launch = self.nt + (32 if prod_warp else 0)
+        assert 0 <= split_chunk < self.nchunk
+        assert maxntid == 0 or maxntid >= self.nt_launch
+        self.maxrreg = maxrreg
+        self.maxntid = maxntid
+        self.copy_hint = copy_hint
+        self.split_chunk = split_chunk
+        self.split_delay_ns = split_delay_ns
+
+    @cute.jit
+    def __call__(
+        self,
+        act: Int64,
+        sumsq: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        pf0: Int64,
+        pf0_bytes: Int64,
+        pf1: Int64,
+        pf1_bytes: Int64,
+        pf2: Int64,
+        pf2_bytes: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        eps: Float32,
+        stream: CUstream,
+    ):
+        if cutlass.const_expr(self.prod_warp):
+            if cutlass.const_expr(self.maxntid > 0):
+                self.kernel_pw(
+                    act,
+                    sumsq,
+                    w0,
+                    w1,
+                    out0,
+                    out1,
+                    ts,
+                    pf0,
+                    pf0_bytes,
+                    pf1,
+                    pf1_bytes,
+                    pf2,
+                    pf2_bytes,
+                    anchor,
+                    rows,
+                    eps,
+                ).launch(
+                    grid=[self.ctas, 1, 1],
+                    block=[self.nt_launch, 1, 1],
+                    stream=stream,
+                    use_pdl=self.use_pdl,
+                    max_number_threads=[self.maxntid, 1, 1],
+                )
+            else:
+                self.kernel_pw(
+                    act,
+                    sumsq,
+                    w0,
+                    w1,
+                    out0,
+                    out1,
+                    ts,
+                    pf0,
+                    pf0_bytes,
+                    pf1,
+                    pf1_bytes,
+                    pf2,
+                    pf2_bytes,
+                    anchor,
+                    rows,
+                    eps,
+                ).launch(
+                    grid=[self.ctas, 1, 1],
+                    block=[self.nt_launch, 1, 1],
+                    stream=stream,
+                    use_pdl=self.use_pdl,
+                )
+        elif cutlass.const_expr(self.maxntid > 0):
+            self.kernel(
+                act,
+                sumsq,
+                w0,
+                w1,
+                out0,
+                out1,
+                ts,
+                pf0,
+                pf0_bytes,
+                pf1,
+                pf1_bytes,
+                pf2,
+                pf2_bytes,
+                anchor,
+                rows,
+                eps,
+            ).launch(
+                grid=[self.ctas, 1, 1],
+                block=[self.nt, 1, 1],
+                stream=stream,
+                use_pdl=self.use_pdl,
+                max_number_threads=[self.maxntid, 1, 1],
+            )
+        else:
+            self.kernel(
+                act,
+                sumsq,
+                w0,
+                w1,
+                out0,
+                out1,
+                ts,
+                pf0,
+                pf0_bytes,
+                pf1,
+                pf1_bytes,
+                pf2,
+                pf2_bytes,
+                anchor,
+                rows,
+                eps,
+            ).launch(
+                grid=[self.ctas, 1, 1], block=[self.nt, 1, 1], stream=stream, use_pdl=self.use_pdl
+            )
+
+    @cute.jit
+    def _copy_chunk(
+        self,
+        sW: Int32,
+        sMbar: Int32,
+        w0: Int64,
+        w1: Int64,
+        row0: Int32,
+        nsm: Int32,
+        j: cutlass.Constexpr[int],
+        policy: Int64,
+    ):
+        chunk_bytes = Int32(self.chunk_k * 2)
+        _mbar_expect_tx(sMbar + Int32(8 * j), nsm * chunk_bytes)
+        for i in cutlass.range(nsm):
+            dst = sW + i * Int32(self.stride) + Int32(j * self.chunk_k * 2)
+            src = _row_src(w0, w1, row0 + i, self.k, self.n0, self.n1) + Int64(j * self.chunk_k * 2)
+            if cutlass.const_expr(self.copy_hint == "evict_first"):
+                bulk_g2s_hint(dst, src, chunk_bytes, sMbar + Int32(8 * j), policy)
+            else:
+                _bulk_g2s(dst, src, chunk_bytes, sMbar + Int32(8 * j))
+
+    @cute.jit
+    def _copy_one(
+        self,
+        sW: Int32,
+        sMbar: Int32,
+        w0: Int64,
+        w1: Int64,
+        row0: Int32,
+        r: Int32,
+        j: cutlass.Constexpr[int],
+        policy: Int64,
+    ):
+        chunk_bytes = Int32(self.chunk_k * 2)
+        dst = sW + r * Int32(self.stride) + Int32(j * self.chunk_k * 2)
+        src = _row_src(w0, w1, row0 + r, self.k, self.n0, self.n1) + Int64(j * self.chunk_k * 2)
+        if cutlass.const_expr(self.copy_hint == "evict_first"):
+            bulk_g2s_hint(dst, src, chunk_bytes, sMbar + Int32(8 * j), policy)
+        else:
+            _bulk_g2s(dst, src, chunk_bytes, sMbar + Int32(8 * j))
+
+    @cute.jit
+    def _issue_copies(
+        self, sW: Int32, sMbar: Int32, w0: Int64, w1: Int64, row0: Int32, nsm: Int32, tid: Int32
+    ):
+        if cutlass.const_expr(self.copy_warps):
+            chunk_bytes = Int32(self.chunk_k * 2)
+            if tid == Int32(0):
+                for j in cutlass.range_constexpr(self.nchunk):
+                    _mbar_expect_tx(sMbar + Int32(8 * j), nsm * chunk_bytes)
+            if tid % Int32(32) == Int32(0):
+                warp = tid // Int32(32)
+                policy = Int64(0)
+                if cutlass.const_expr(self.copy_hint == "evict_first"):
+                    policy = policy_evict_first()
+                per_warp = (self.smem_rows + self.nw - 1) // self.nw
+                for j in cutlass.range_constexpr(self.nchunk):
+                    for q in cutlass.range_constexpr(per_warp):
+                        r = warp + Int32(q * self.nw)
+                        if r < nsm:
+                            self._copy_one(sW, sMbar, w0, w1, row0, r, j, policy)
+        elif cutlass.const_expr(self.copy_hint == "none" and self.split_chunk == 0):
+            DeferredNormProjKernel._issue_copies(self, sW, sMbar, w0, w1, row0, nsm, tid)
+        else:
+            if tid == Int32(0):
+                policy = Int64(0)
+                if cutlass.const_expr(self.copy_hint == "evict_first"):
+                    policy = policy_evict_first()
+                t0 = Int64(0)
+                if cutlass.const_expr(self.split_chunk > 0):
+                    t0 = globaltimer()
+                for j in cutlass.range_constexpr(self.nchunk):
+                    if cutlass.const_expr(self.split_chunk > 0 and j == self.split_chunk):
+                        spin_until(t0 + Int64(self.split_delay_ns))
+                    self._copy_chunk(sW, sMbar, w0, w1, row0, nsm, j, policy)
+
+    @cute.jit
+    def _pw_stamp(self, ts: Int64, bid: Int32, k: cutlass.Constexpr[int]):
+        """Tests only: ``%globaltimer`` of CTA ``bid`` at point ``k`` (producer lane 0)."""
+        if cutlass.const_expr(self.debug_ts):
+            st_b64(ts + Int64((bid * Int32(8) + Int32(k)) * Int32(8)), globaltimer())
+
+    @cute.jit
+    def _produce(
+        self,
+        sW: Int32,
+        sMbar: Int32,
+        w0: Int64,
+        w1: Int64,
+        row0: Int32,
+        nsm: Int32,
+        ts: Int64,
+        bid: Int32,
+        lane: Int32,
+        rows: Int32,
+    ):
+        """Producer warp: lane 0 posts each chunk's byte count and issues its
+        copies. Deferred copies: the chunks ``>= split_chunk`` of every CTA,
+        and every chunk of the last ``defer_tail * rows`` CTAs (the CTA count
+        of the preceding hop at ``defer_tail`` CTAs per token row). They are
+        issued after ``split_delay_ns`` (``%globaltimer``) or, with
+        ``defer_wait``, after the producer's own ``griddepcontrol.wait`` (the
+        preceding grid's completion; the producer still reads no data it
+        wrote). Every lane executes the PDL trigger no later than deferred_norm's
+        thread 0 would (after the copies issued at entry)."""
+        if lane == Int32(0):
+            policy = Int64(0)
+            if cutlass.const_expr(self.copy_hint == "evict_first"):
+                policy = policy_evict_first()
+            t0 = globaltimer()
+            late = Int32(0)
+            if cutlass.const_expr(self.defer_tail > 0):
+                late = Int32(bid >= Int32(self.ctas) - Int32(self.defer_tail) * rows)
+            first = self.split_chunk if self.split_chunk > 0 else self.nchunk
+            if late == Int32(0):
+                for j in cutlass.range_constexpr(first):
+                    self._copy_chunk(sW, sMbar, w0, w1, row0, nsm, j, policy)
+            self._pw_stamp(ts, bid, 7)
+            if cutlass.const_expr(self.use_pdl and self.trigger == "start"):
+                cute.arch.griddepcontrol_launch_dependents()
+            if cutlass.const_expr(self.split_chunk > 0 or self.defer_tail > 0):
+                if cutlass.const_expr(self.split_chunk > 0):
+                    self._defer_point(t0)
+                    for j in cutlass.range_constexpr(self.nchunk):
+                        if cutlass.const_expr(j < first):
+                            if late != Int32(0):
+                                self._copy_chunk(sW, sMbar, w0, w1, row0, nsm, j, policy)
+                        else:
+                            self._copy_chunk(sW, sMbar, w0, w1, row0, nsm, j, policy)
+                else:
+                    if late != Int32(0):
+                        self._defer_point(t0)
+                        for j in cutlass.range_constexpr(self.nchunk):
+                            self._copy_chunk(sW, sMbar, w0, w1, row0, nsm, j, policy)
+            self._pw_stamp(ts, bid, 6)
+        else:
+            if cutlass.const_expr(self.use_pdl and self.trigger == "start"):
+                cute.arch.griddepcontrol_launch_dependents()
+
+    @cute.jit
+    def _defer_point(self, t0: Int64):
+        """Where the producer issues its deferred copies (bounded either way:
+        the timer, or the preceding grid's completion)."""
+        if cutlass.const_expr(self.defer_wait and self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+        else:
+            spin_until(t0 + Int64(self.split_delay_ns))
+
+    @cute.jit
+    def _compute(
+        self,
+        act: Int64,
+        sumsq: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        pf0: Int64,
+        pf0_bytes: Int64,
+        pf1: Int64,
+        pf1_bytes: Int64,
+        pf2: Int64,
+        pf2_bytes: Int64,
+        rows: Int32,
+        eps: Float32,
+        tid: Int32,
+        bid: Int32,
+        warp: Int32,
+        lane: Int32,
+        sW: Int32,
+        sMisc: Int32,
+        row0: Int32,
+        nrows: Int32,
+        nsm: Int32,
+    ):
+        """The 16 MMA warps: deferred_norm's kernel body after its copy issue, with the
+        post-loop barriers over the MMA threads only (no overflow row)."""
+        k = self.k
+        nw = self.nw
+        stride = self.stride
+        sZero = sMisc
+        sMbar = sMisc + Int32(64)
+        sR = sMisc + Int32(256)
+        if cutlass.const_expr(self.prefetch_next):
+            if tid < Int32(3):
+                pf_addr = _sel_i64(
+                    Int32(tid == Int32(0)), pf0, _sel_i64(Int32(tid == Int32(1)), pf1, pf2)
+                )
+                pf_bytes = _sel_i64(
+                    Int32(tid == Int32(0)),
+                    pf0_bytes,
+                    _sel_i64(Int32(tid == Int32(1)), pf1_bytes, pf2_bytes),
+                )
+                self._prefetch_slice(pf_addr, pf_bytes, bid)
+        if cutlass.const_expr(self.l2_prefetch):
+            # Rows that are not staged (none: the producer variant stages every row).
+            if tid >= nsm:
+                if tid < nrows:
+                    _prefetch_l2(_row_src(w0, w1, row0 + tid, k, self.n0, self.n1), Int32(k * 2))
+        if cutlass.const_expr(self.use_pdl and self.trigger == "start"):
+            cute.arch.griddepcontrol_launch_dependents()
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+        _stamp(self.debug_ts, ts, bid, tid, 1)
+
+        g = lane // Int32(4)
+        t4 = lane % Int32(4)
+        tok_ok = Int32(g < rows)
+        safe_g = _min_i32(g, rows - Int32(1))
+        pq = []
+        if cutlass.const_expr(self.deferred):
+            quarter = self.partials_per_row // 4
+            p_ok = Int32(warp == Int32(nw - 1)) * tok_ok
+            pbase = sumsq + Int64(
+                (g * Int32(self.partials_per_row) + t4 * Int32(quarter)) * Int32(4)
+            )
+            for j in cutlass.range_constexpr(quarter // 2):
+                lo, hi = _ld_v2_if(pbase + Int64(8 * j), p_ok)
+                pq.append(bits_f32(lo))
+                pq.append(bits_f32(hi))
+        bfr = []
+        if cutlass.const_expr(self.b_layout == "frag"):
+            a_lane = act + Int64(lane * Int32(8))
+            for m in cutlass.range_constexpr(self.rounds):
+                bfr.append(_ld_v2_if(a_lane + Int64((warp + Int32(nw * m)) * Int32(256)), tok_ok))
+        else:
+            a_row = act + Int64(safe_g * Int32(k) * Int32(2) + t4 * Int32(4))
+            for m in cutlass.range_constexpr(self.rounds):
+                off = Int64((warp + Int32(nw * m)) * Int32(32))
+                b0 = _ld_b32(a_row + off)
+                b1 = _ld_b32(a_row + off + Int64(16))
+                bfr.append((_sel_i32(tok_ok, b0, Int32(0)), _sel_i32(tok_ok, b1, Int32(0))))
+        _stamp(self.debug_ts, ts, bid, tid, 2)
+
+        r16 = lane % Int32(16)
+        a_lane0 = sW + r16 * Int32(stride) + (lane // Int32(16)) * Int32(16)
+        ok0 = Int32(r16 < nsm)
+        r8 = Int32(16) + lane % Int32(8)
+        a_lane1 = sW + r8 * Int32(stride) + ((lane // Int32(8)) % Int32(2)) * Int32(16)
+        ok1 = Int32(r8 < nsm)
+
+        acc0 = [Float32(0.0)] * 4
+        acc1 = [Float32(0.0)] * 4
+        zero = Int32(0)
+        rpc = self.rounds_per_chunk
+
+        def load_a(n):
+            if cutlass.const_expr(n % rpc == 0):
+                _mbar_wait(sMbar + Int32(8 * (n // rpc)), Int32(0))
+            kb = (warp + Int32(nw * n)) * Int32(32)
+            return (
+                ldsm_x4(_sel_i32(ok0, a_lane0 + kb, sZero)),
+                _ldsm_x2(_sel_i32(ok1, a_lane1 + kb, sZero)),
+            )
+
+        frags = [load_a(n) for n in range(min(self.dist, self.rounds))]
+        for m in cutlass.range_constexpr(self.rounds):
+            if cutlass.const_expr(m + self.dist < self.rounds):
+                frags.append(load_a(m + self.dist))
+            fa, fb = frags[m]
+            b0, b1 = bfr[m]
+            acc0[0], acc0[1], acc0[2], acc0[3] = mma_bf16(
+                fa[0], fa[1], fa[2], fa[3], b0, b1, acc0[0], acc0[1], acc0[2], acc0[3]
+            )
+            acc1[0], acc1[1], acc1[2], acc1[3] = mma_bf16(
+                fb[0], zero, fb[1], zero, b0, b1, acc1[0], acc1[1], acc1[2], acc1[3]
+            )
+        _stamp(self.debug_ts, ts, bid, tid, 3)
+
+        if cutlass.const_expr(self.deferred):
+            if warp == Int32(nw - 1):
+                s = Float32(0.0)
+                for q in cutlass.range_constexpr(4):
+                    mine = s
+                    for v in pq:
+                        mine = add(mine, v)
+                    s = _shfl_idx_f32(mine, g * Int32(4) + Int32(q))
+                if t4 == Int32(0):
+                    _st_shared_f32(sR + g * Int32(4), inv_rms(s, Float32(float(k)), eps))
+        else:
+            if tid < Int32(MAX_TOKENS):
+                _st_shared_f32(sR + tid * Int32(4), Float32(1.0))
+
+        # Cross-warp reduction in warp order (every MMA warp has waited on
+        # every chunk, so every copy has landed before a row is overwritten).
+        bar_sync_named(1, self.nt)
+        sRed = sW
+        wbase = sRed + warp * Int32(24 * MAX_TOKENS * 4)
+        c0 = Int32(2) * t4
+        _st_shared_f32(wbase + (g * Int32(MAX_TOKENS) + c0) * Int32(4), acc0[0])
+        _st_shared_f32(wbase + (g * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc0[1])
+        _st_shared_f32(wbase + ((g + Int32(8)) * Int32(MAX_TOKENS) + c0) * Int32(4), acc0[2])
+        _st_shared_f32(
+            wbase + ((g + Int32(8)) * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc0[3]
+        )
+        _st_shared_f32(wbase + ((g + Int32(16)) * Int32(MAX_TOKENS) + c0) * Int32(4), acc1[0])
+        _st_shared_f32(
+            wbase + ((g + Int32(16)) * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc1[1]
+        )
+        bar_sync_named(1, self.nt)
+        if tid < Int32(24 * MAX_TOKENS):
+            row = tid // Int32(MAX_TOKENS)
+            tok = tid % Int32(MAX_TOKENS)
+            if row < nsm:
+                if tok < rows:
+                    s = Float32(0.0)
+                    for w in cutlass.range_constexpr(nw):
+                        s = add(
+                            s,
+                            _ld_shared_f32(sRed + Int32(w * 24 * MAX_TOKENS * 4) + tid * Int32(4)),
+                        )
+                    _store_out(
+                        out0,
+                        out1,
+                        row0 + row,
+                        tok,
+                        mul(_ld_shared_f32(sR + tok * Int32(4)), s),
+                        self.n0,
+                        self.n1,
+                    )
+        _stamp(self.debug_ts, ts, bid, tid, 4)
+        if cutlass.const_expr(self.use_pdl and self.trigger == "end"):
+            cute.arch.griddepcontrol_launch_dependents()
+
+    @cute.kernel
+    def kernel_pw(
+        self,
+        act: Int64,
+        sumsq: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        pf0: Int64,
+        pf0_bytes: Int64,
+        pf1: Int64,
+        pf1_bytes: Int64,
+        pf2: Int64,
+        pf2_bytes: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        eps: Float32,
+    ):
+        """deferred_norm's projection with a producer warp (``prod_warp``; module docstring)."""
+        stride = self.stride
+        tid, _, _ = cute.arch.thread_idx()
+        bid, _, _ = cute.arch.block_idx()
+        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        lane = cute.arch.lane_idx()
+        _stamp(self.debug_ts, ts, bid, tid, 0)
+
+        smem = SmemAllocator()
+        sW = saddr(smem.allocate(self.smem_rows * stride, byte_alignment=128))
+        sMisc = saddr(smem.allocate(512, byte_alignment=128))
+        sMbar = sMisc + Int32(64)
+
+        base = Int32(self.base_rows)
+        rem = Int32(self.rem_rows)
+        extra = Int32(bid < rem)
+        nrows = base + extra
+        row0 = bid * base + _min_i32(bid, rem)
+        nsm = _min_i32(nrows, Int32(self.smem_rows))
+
+        if tid < Int32(2):
+            st_shared_v4(sMisc + tid * Int32(16), Int32(0), Int32(0), Int32(0), Int32(0))
+        if tid == Int32(0):
+            for j in cutlass.range_constexpr(self.nchunk + 1):
+                _mbar_init(sMbar + Int32(8 * j), Int32(1))
+            _mbar_fence_init()
+        cute.arch.sync_threads()
+
+        if warp == Int32(self.nw):
+            self._produce(sW, sMbar, w0, w1, row0, nsm, ts, bid, lane, rows)
+        else:
+            self._compute(
+                act,
+                sumsq,
+                w0,
+                w1,
+                out0,
+                out1,
+                ts,
+                pf0,
+                pf0_bytes,
+                pf1,
+                pf1_bytes,
+                pf2,
+                pf2_bytes,
+                rows,
+                eps,
+                tid,
+                bid,
+                warp,
+                lane,
+                sW,
+                sMisc,
+                row0,
+                nrows,
+                nsm,
+            )
+        _stamp(self.debug_ts, ts, bid, tid, 5)
+
+
+_COMPILED = {}
+
+
+def is_compiled(**cfg) -> bool:
+    return tuple(sorted(cfg.items())) in _COMPILED
+
+
+def compiled_evict_hints_projection(**cfg):
+    """Return the compiled deferred_norm_evict_hints projection (deferred_norm keys plus
+    ``projection.EVICT_HINTS_CONFIG_KEYS``); compiles once."""
+    key = tuple(sorted(cfg.items()))
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = EvictHintsProjKernel(**cfg)
+        options = "--enable-tvm-ffi"
+        if kernel.maxrreg:
+            options += f" --ptxas-options=--maxrregcount={kernel.maxrreg}"
+        z64 = Int64(0)
+        anchor = cute_runtime.make_fake_compact_tensor(
+            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
+        )
+        compiled = cute.compile(
+            kernel,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            anchor,
+            Int32(1),
+            Float32(1e-5),
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options=options,
+        )
+        _COMPILED[key] = compiled
+    return compiled
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/projection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/projection.py
new file mode 100644
index 000000000000..e01a8cd8dc85
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/projection.py
@@ -0,0 +1,227 @@
+# 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.
+"""deferred_norm_evict_hints next-layer projection: deferred_norm's staged projection with
+deferred_norm_evict_hints's schedule knobs.
+
+``EvictHintsProjection`` is ``DeferredNormProjection`` (same module binding, weights read in
+place, same outputs bit for bit) launching ``EvictHintsProjKernel``. The launch
+configuration is deferred_norm's served one plus ``EvictHintsProjConfig``'s fields.
+"""
+
+from dataclasses import asdict, dataclass, fields
+from typing import Optional, Tuple
+
+import torch
+
+from ..deferred_norm.projection import (
+    _N_KVA,
+    MAX_ROWS,
+    SERVED_N1,
+    DeferredNormProjConfig,
+    DeferredNormProjection,
+    _kernel_config,
+    _sm_count,
+)
+
+# EvictHintsProjKernel's keys beyond deferred_norm's launch configuration.
+EVICT_HINTS_CONFIG_KEYS = (
+    "maxrreg",
+    "maxntid",
+    "copy_hint",
+    "split_chunk",
+    "split_delay_ns",
+    "copy_warps",
+    "prod_warp",
+    "defer_tail",
+    "defer_wait",
+)
+
+
+@dataclass(frozen=True)
+class EvictHintsProjConfig(DeferredNormProjConfig):
+    """deferred_norm's launch configuration plus deferred_norm_evict_hints's knobs (``proj_kernel.py``).
+
+    The defaults are deferred_norm's launch with ``copy_hint="evict_first"`` (the
+    weight copies carry an ``L2::evict_first`` policy). ``maxrreg`` caps the
+    registers per thread (0: none). ``prod_warp`` adds a copy-issuing 17th
+    warp, so that ``split_chunk`` / ``split_delay_ns`` / ``defer_tail`` /
+    ``defer_wait`` defer part of the staging without holding an MMA warp.
+    """
+
+    maxrreg: int = 0
+    maxntid: int = 0
+    copy_hint: str = "evict_first"
+    split_chunk: int = 0
+    split_delay_ns: int = 0
+    copy_warps: bool = False
+    prod_warp: bool = False
+    defer_tail: int = 0
+    defer_wait: bool = False
+
+
+def parse_cfg(text: str, cls=None):
+    """``"key=value,..."`` into a config (bool / int / str by field type)."""
+    cls = cls or EvictHintsProjConfig
+    kwargs = {}
+    types = {f.name: f.type for f in fields(cls)}
+    for item in filter(None, text.replace(";", ",").split(",")):
+        key, value = item.split("=")
+        typ = types[key]
+        if typ in (bool, "bool"):
+            kwargs[key] = value.lower() in ("1", "true")
+        elif typ in (int, "int"):
+            kwargs[key] = int(value)
+        else:
+            kwargs[key] = value
+    return cls(**kwargs)
+
+
+def config_for(config: EvictHintsProjConfig, n1: int) -> EvictHintsProjConfig:
+    """The configuration deferred_norm_evict_hints launches for a layer kind.
+
+    ``config`` serves the shared-indexer layers (``n1 = 0``), the only
+    boundaries deferred_norm serves under its ``shared`` scope. A full-indexer layer
+    (``n1 > 0``, scope ``all`` only) keeps deferred_norm's served launch configuration (no
+    register cap, no copy hint, one CTA per SM).
+    """
+    return config if n1 == 0 else EvictHintsProjConfig(copy_hint="none")
+
+
+def _evict_hints_kernel_config(
+    config: EvictHintsProjConfig, n1: int, ctas: int, use_pdl: bool, deferred: bool, debug_ts: bool
+) -> dict:
+    cfg = _kernel_config(config, n1, ctas, use_pdl, deferred, debug_ts)
+    z = asdict(config)
+    for key in EVICT_HINTS_CONFIG_KEYS:
+        cfg[key] = z[key]
+    return cfg
+
+
+def served_kernel_configs(device, config: Optional[EvictHintsProjConfig] = None) -> list:
+    """Kernel configurations the deferred_norm_evict_hints consumer launches (one per ``SERVED_N1``)."""
+    from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+    config = config or EvictHintsProjConfig()
+    use_pdl = get_env_enable_pdl()
+    out = []
+    for n1 in SERVED_N1:
+        kind_config = config_for(config, n1)
+        ctas = kind_config.ctas or _sm_count(device)
+        cfg = _evict_hints_kernel_config(
+            kind_config, n1, ctas, use_pdl, deferred=True, debug_ts=False
+        )
+        cfg["b_layout"] = "frag"
+        out.append(cfg)
+    return out
+
+
+def precompile_served(device, config: Optional[EvictHintsProjConfig] = None) -> Optional[str]:
+    """Compile every served deferred_norm_evict_hints projection kernel (no GPU work); see deferred_norm's."""
+    from ..deferred_norm.proj_kernel import projection_supported
+    from .proj_kernel import compiled_evict_hints_projection
+
+    cfgs = served_kernel_configs(device, config)
+    for cfg in cfgs:
+        why = projection_supported(
+            cfg["k"],
+            cfg["n0"] + cfg["n1"],
+            cfg["ctas"],
+            cfg["threads"],
+            cfg["smem_rows"],
+            cfg["nchunk"],
+        )
+        if why is not None:
+            return f"projection n1={cfg['n1']}: {why}"
+    for cfg in cfgs:
+        compiled_evict_hints_projection(**cfg)
+    return None
+
+
+class EvictHintsProjection(DeferredNormProjection):
+    """``DeferredNormProjection`` launching the deferred_norm_evict_hints kernel."""
+
+    def __init__(
+        self,
+        mla,
+        config: Optional[EvictHintsProjConfig] = None,
+        kw_weight: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(mla, config or EvictHintsProjConfig(), kw_weight)
+        self.config = config_for(self.config, self.n1)
+        self.ctas = self.config.ctas or _sm_count(self.w0.device)
+
+    def kernel_config(self, deferred: bool = True, debug_ts: bool = False) -> dict:
+        return _evict_hints_kernel_config(
+            self.config, self.n1, self.ctas, self.use_pdl, deferred, debug_ts
+        )
+
+    def precompile(self) -> None:
+        from .proj_kernel import compiled_evict_hints_projection
+
+        compiled_evict_hints_projection(**self.served_config())
+
+    def run(
+        self,
+        act: torch.Tensor,
+        sumsq: Optional[torch.Tensor],
+        eps: float,
+        rows: Optional[int] = None,
+        ts: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Launch on the current stream (arguments as ``DeferredNormProjection.run``)."""
+        from .proj_kernel import compiled_evict_hints_projection
+
+        frag = act.dim() == 1
+        if not frag:
+            rows = act.shape[0]
+        assert rows is not None and 1 <= rows <= MAX_ROWS
+        assert act.is_contiguous() and act.dtype == torch.bfloat16
+        kva = act.new_empty(rows, _N_KVA)
+        kw = act.new_empty(rows, self.n1) if self.n1 else None
+        deferred = sumsq is not None
+        cfg = self.kernel_config(deferred, ts is not None)
+        cfg["b_layout"] = "frag" if frag else "plain"
+        compiled = compiled_evict_hints_projection(**cfg)
+        pf = [(t.data_ptr(), t.numel() * t.element_size()) for t in self.prefetch[:3]]
+        pf += [(0, 0)] * (3 - len(pf))
+        compiled(
+            act.data_ptr(),
+            sumsq.data_ptr() if deferred else 0,
+            self.w0.data_ptr(),
+            0 if self.w1 is None else self.w1.data_ptr(),
+            kva.data_ptr(),
+            0 if kw is None else kw.data_ptr(),
+            0 if ts is None else ts.data_ptr(),
+            pf[0][0],
+            pf[0][1],
+            pf[1][0],
+            pf[1][1],
+            pf[2][0],
+            pf[2][1],
+            kva.view(torch.int32).view(-1),
+            rows,
+            float(eps),
+        )
+        return kva, kw
+
+
+__all__ = [
+    "EvictHintsProjConfig",
+    "EvictHintsProjection",
+    "config_for",
+    "parse_cfg",
+    "precompile_served",
+    "served_kernel_configs",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/ptx.py
new file mode 100644
index 000000000000..302de51a04a8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/ptx.py
@@ -0,0 +1,166 @@
+# 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.
+"""Inline-PTX helpers that only the deferred_norm_evict_hints kernels use (CuTe DSL).
+
+None of them computes a value that reaches an output: they are L2 cache
+hints (bulk prefetch, eviction-priority policies on stores and bulk copies),
+a named CTA barrier, a ``%globaltimer`` spin and test-only timer / SM-id
+stamps. The arithmetic and the post_moe_hop workspace
+protocol come unchanged from deferred_norm's helpers (``deferred_norm/ptx.py``).
+"""
+
+from cutlass import Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+
+def _asm(res, args, text, cons, side_effects=True, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=side_effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+def _v(x, loc, ip):
+    return x.ir_value(loc=loc, ip=ip)
+
+
+@dsl_user_op
+def prefetch_l2(gaddr: Int64, nbytes: Int32, *, loc=None, ip=None) -> None:
+    """``cp.async.bulk.prefetch.L2`` of ``nbytes`` (16-byte multiple) at ``gaddr``."""
+    _asm(
+        None,
+        [_v(gaddr, loc, ip), _v(nbytes, loc, ip)],
+        "cp.async.bulk.prefetch.L2.global [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def prefetch_line_l2(gaddr: Int64, *, loc=None, ip=None) -> None:
+    """Per-thread ``prefetch.global.L2`` of the line holding ``gaddr`` (no
+    uniform-operand serialization, unlike the bulk prefetch)."""
+    _asm(None, [_v(gaddr, loc, ip)], "prefetch.global.L2 [$0];", "l", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def policy_evict_last(*, loc=None, ip=None) -> Int64:
+    """L2 cache policy: every access marked ``evict_last``."""
+    return Int64(
+        _asm(
+            T.i64(), [], "createpolicy.fractional.L2::evict_last.b64 $0, 1.0;", "=l", loc=loc, ip=ip
+        )
+    )
+
+
+@dsl_user_op
+def policy_evict_first(*, loc=None, ip=None) -> Int64:
+    """L2 cache policy: every access marked ``evict_first``."""
+    return Int64(
+        _asm(
+            T.i64(),
+            [],
+            "createpolicy.fractional.L2::evict_first.b64 $0, 1.0;",
+            "=l",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def st_v4_hint(
+    addr: Int64, a: Int32, b: Int32, c: Int32, d: Int32, policy: Int64, *, loc=None, ip=None
+) -> None:
+    """``st.global.v4.b32`` with an L2 cache-policy hint (same store otherwise)."""
+    _asm(
+        None,
+        [_v(addr, loc, ip)] + [_v(x, loc, ip) for x in (a, b, c, d)] + [_v(policy, loc, ip)],
+        "st.global.L2::cache_hint.v4.b32 [$0], {$1, $2, $3, $4}, $5;",
+        "l,r,r,r,r,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def bulk_g2s_hint(
+    saddr: Int32, gaddr: Int64, nbytes: Int32, mbar: Int32, policy: Int64, *, loc=None, ip=None
+) -> None:
+    """``cp.async.bulk`` global -> shared with an L2 cache-policy hint, completing
+    on ``mbar`` with ``complete_tx`` (deferred_norm's copy otherwise)."""
+    _asm(
+        None,
+        [
+            _v(saddr, loc, ip),
+            _v(gaddr, loc, ip),
+            _v(nbytes, loc, ip),
+            _v(mbar, loc, ip),
+            _v(policy, loc, ip),
+        ],
+        "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint "
+        "[$0], [$1], $2, [$3], $4;",
+        "r,l,r,r,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def globaltimer(*, loc=None, ip=None) -> Int64:
+    return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def smid(*, loc=None, ip=None) -> Int32:
+    return Int32(_asm(T.i32(), [], "mov.u32 $0, %smid;", "=r", loc=loc, ip=ip))
+
+
+@dsl_user_op
+def st_b64(gaddr: Int64, v: Int64, *, loc=None, ip=None) -> None:
+    _asm(
+        None, [_v(gaddr, loc, ip), _v(v, loc, ip)], "st.global.b64 [$0], $1;", "l,l", loc=loc, ip=ip
+    )
+
+
+@dsl_user_op
+def spin_until(deadline: Int64, *, loc=None, ip=None) -> None:
+    """Spin until ``%globaltimer`` reaches ``deadline`` (bounded by the clock only)."""
+    _asm(
+        None,
+        [_v(deadline, loc, ip)],
+        "{\n.reg .pred p;\n.reg .b64 now;\nHINTS_SPIN:\nmov.u64 now, %globaltimer;\n"
+        "setp.lt.s64 p, now, $0;\n@p bra HINTS_SPIN;\n}\n",
+        "l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def bar_sync_named(bar_id: int, count: int, *, loc=None, ip=None) -> None:
+    """``bar.sync bar_id, count``: a CTA barrier over ``count`` threads only
+    (``count`` a multiple of 32; ``bar_id`` != 0, which ``sync_threads`` uses)."""
+    assert 0 < bar_id < 16 and count % 32 == 0
+    _asm(None, [], f"bar.sync {bar_id}, {count};", "", loc=loc, ip=ip)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/__init__.py
new file mode 100644
index 000000000000..5c17dcc1948c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/__init__.py
@@ -0,0 +1,26 @@
+# 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.
+"""dense_mlp: small-M dense MLP of the dense decoder layers.
+
+Registered entry point: ``dense_mlp_forward(mlp, x, final_all_reduce_params,
+lora_params) -> Optional[Tensor]`` (``None`` means "run the native module").
+"""
+
+from ... import dispatch
+from .dense_mlp import check_eligibility, dense_mlp_forward, op_available
+
+dispatch.register("dense_mlp", "dense_mlp_forward", dense_mlp_forward)
+
+__all__ = ["check_eligibility", "dense_mlp_forward", "op_available"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp.py
new file mode 100644
index 000000000000..e8b1c4f45b90
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp.py
@@ -0,0 +1,152 @@
+# 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.
+"""Small-M BF16 dense MLP for the dense decoder layers (``dense_mlp``).
+
+Replaces the rank-local compute of a BF16 ``GatedMLP`` with plain SwiGLU at
+decode row counts: ``gate_up_proj`` GEMM, SwiGLU and ``down_proj`` GEMM become
+two weight-streaming kernels (``trtllm::glm_moe_dsa_dense_mlp``). The result is
+the same rank-local partial that native ``down_proj`` returns when its
+all-reduce is disabled; the caller's all-reduce, residual and normalization
+stay native. Intermediate rounding follows the native path: gate and up are
+rounded to BF16 before SwiGLU and the SwiGLU output is rounded to BF16.
+
+Every eligibility check reads only module attributes and tensor metadata, so a
+call that is not eligible returns ``None`` before any device work and the
+caller runs the native module.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+# Largest activation row count of the kernels (MMA N side).
+MAX_ROWS = 8
+# Contraction sizes must be multiples of this for the default tilings
+# (warps x 32 contraction elements per step x steps per chunk = 1024 for both).
+CONTRACTION_MULTIPLE = 1024
+# Output-column granularity of the gate/up (8) and down (16) CTA tiles.
+GATE_UP_COL_MULTIPLE = 8
+DOWN_ROW_MULTIPLE = 16
+# Tiling variants of the gate/up and down kernels (see denseMlp.cu) for
+# the GLM-5.2 GEN shapes (M = 6, H = 6144, rank-local I = 3072): gate/up
+# 8 warps x 4 steps at three CTAs per SM, down 4 warps x 8 steps, both with a
+# one-chunk L2 prefetch.
+GATE_UP_VARIANT = 1
+DOWN_VARIANT = 0
+
+_OP_NAME = "glm_moe_dsa_dense_mlp"
+
+
+def op_available() -> bool:
+    """Return True when the ``trtllm::glm_moe_dsa_dense_mlp`` op is registered."""
+    return hasattr(torch.ops.trtllm, _OP_NAME)
+
+
+def _linear_is_plain_bf16(linear) -> bool:
+    from tensorrt_llm._torch.modules.linear import UnquantizedLinearMethod
+
+    weight = getattr(linear, "weight", None)
+    return (
+        isinstance(getattr(linear, "quant_method", None), UnquantizedLinearMethod)
+        and weight is not None
+        and weight.dtype == torch.bfloat16
+        and weight.dim() == 2
+        and weight.is_contiguous()
+        and getattr(linear, "bias", None) is None
+    )
+
+
+def _output_allreduce_requested(mlp, final_all_reduce_params) -> bool:
+    down = mlp.down_proj
+    if not getattr(down, "reduce_output", False):
+        return False
+    if down.mapping.tp_size == 1:
+        return False
+    return final_all_reduce_params is None or final_all_reduce_params.enable_allreduce
+
+
+def check_eligibility(
+    mlp, x, final_all_reduce_params=None, lora_params: Optional[dict] = None
+) -> Tuple[bool, str]:
+    """Decide on the host whether ``mlp(x)`` can run on the dense_mlp kernels.
+
+    Returns:
+        ``(eligible, reason)``; ``reason`` names the first failed condition.
+    """
+    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+
+    if not isinstance(mlp, GatedMLP):
+        return False, "not a GatedMLP"
+    if getattr(getattr(mlp, "mapping", None), "enable_attention_dp", False):
+        # With attention DP the dense MLP runs unsharded (4x the rank-local
+        # intermediate of the validated TP8 GEN shapes); outside the domain.
+        return False, "attention data parallelism"
+    if lora_params:
+        return False, "LoRA"
+    if not isinstance(x, torch.Tensor):
+        return False, "input is not a plain tensor"
+    if not x.is_cuda or x.dtype != torch.bfloat16 or x.dim() != 2:
+        return False, "input is not a 2-D BF16 CUDA tensor"
+    if not x.is_contiguous() or x.data_ptr() % 16 != 0:
+        return False, "input is not contiguous and 16-byte aligned"
+    num_rows, hidden = x.shape
+    if num_rows < 1 or num_rows > MAX_ROWS:
+        return False, f"row count outside [1, {MAX_ROWS}]"
+    if mlp.split_gate_up:
+        return False, "split gate/up projections"
+    if mlp.activation is not F.silu or not mlp._is_plain_swiglu() or mlp.swiglu_limit is not None:
+        return False, "activation is not plain SwiGLU"
+    if not (_linear_is_plain_bf16(mlp.gate_up_proj) and _linear_is_plain_bf16(mlp.down_proj)):
+        return False, "projections are not unquantized BF16 without bias"
+    if _output_allreduce_requested(mlp, final_all_reduce_params):
+        return False, "down_proj all-reduce requested"
+    gate_up_w = mlp.gate_up_proj.weight
+    down_w = mlp.down_proj.weight
+    inter = down_w.shape[1]
+    if gate_up_w.shape[0] != 2 * inter or gate_up_w.shape[1] != hidden or down_w.shape[0] != hidden:
+        return False, "weight shapes do not match the input"
+    if (
+        hidden % CONTRACTION_MULTIPLE != 0
+        or inter % CONTRACTION_MULTIPLE != 0
+        or inter % GATE_UP_COL_MULTIPLE != 0
+        or hidden % DOWN_ROW_MULTIPLE != 0
+    ):
+        return False, "hidden/intermediate sizes not supported by the tiling"
+    if gate_up_w.device != x.device or down_w.device != x.device:
+        return False, "weights on a different device"
+    if not op_available():
+        return False, "trtllm::glm_moe_dsa_dense_mlp op not built"
+    return True, ""
+
+
+def dense_mlp_forward(
+    mlp, x, final_all_reduce_params=None, lora_params: Optional[dict] = None
+) -> Optional[torch.Tensor]:
+    """Run the dense MLP on the dense_mlp kernels, or return ``None`` if ineligible.
+
+    The returned tensor is the rank-local ``down_proj`` output (no all-reduce),
+    exactly what ``mlp(x, final_all_reduce_params=...)`` returns when its
+    all-reduce is disabled.
+    """
+    from ... import dispatch
+
+    eligible, reason = check_eligibility(mlp, x, final_all_reduce_params, lora_params)
+    dispatch.log_capability("dense_mlp", f"dense_mlp:{reason or 'engaged'}", eligible, reason)
+    if not eligible:
+        return None
+    return torch.ops.trtllm.glm_moe_dsa_dense_mlp(
+        x, mlp.gate_up_proj.weight, mlp.down_proj.weight, GATE_UP_VARIANT, DOWN_VARIANT
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/__init__.py
new file mode 100644
index 000000000000..3f46f586276f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/__init__.py
@@ -0,0 +1,42 @@
+# 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.
+"""draft_front: MTP draft front-end normalization and TP-local packing (the draft_front contract).
+
+Registered entry points (``dispatch.get_impl("draft_front", name)``):
+
+* ``draft_front``: :func:`pack_draft_front`, returns the TP-local ``eh_proj``
+  input or ``None`` to request the native path.
+* ``unsupported_reason``: :func:`unsupported_reason`, the host-side
+  support-domain check used by ``draft_front``.
+* ``native_norm_reason``: :func:`native_norm_reason`, for the caller that owns
+  the ``enorm``/``hnorm`` modules: ``None`` when they run the FlashInfer CuTe DSL
+  RMSNorm that draft_front mirrors, else why the native draft front must stay.
+"""
+
+from ... import dispatch
+from .draft_front import Segment, pack_draft_front, plan_segments, unsupported_reason
+from .native_norm import native_norm_reason
+
+dispatch.register("draft_front", "draft_front", pack_draft_front)
+dispatch.register("draft_front", "unsupported_reason", unsupported_reason)
+dispatch.register("draft_front", "native_norm_reason", native_norm_reason)
+
+__all__ = [
+    "Segment",
+    "native_norm_reason",
+    "pack_draft_front",
+    "plan_segments",
+    "unsupported_reason",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/draft_front.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/draft_front.py
new file mode 100644
index 000000000000..f89207a166bc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/draft_front.py
@@ -0,0 +1,223 @@
+# 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.
+"""The draft_front contract: the MTP draft front feeding the ROW-parallel ``eh_proj``.
+
+Native semantics (``DeepseekV3MTP.forward``), with ``H`` the hidden size::
+
+    x = concat([enorm(embeds), hnorm(hidden)], dim=-1)  # [R, 2H]
+    x = chunk(x, tp_size, dim=-1)[tp_rank]  # [R, 2H / tp_size] view
+    out = eh_proj(x)  # ROW-parallel, all-reduced
+
+``pack_draft_front`` returns the same ``eh_proj`` input. Only the source(s) that
+the rank's chunk covers are normalized and only the chunk's columns are written.
+The result is a view into a fresh ``[R, 2H]`` buffer with the native chunk's
+shape, strides and offset, so the ``eh_proj`` GEMM sees an operand laid out
+exactly like the native one. The other columns of the buffer are never read.
+
+Embedding lookup, ``eh_proj`` and its reduction stay with the caller.
+
+The kernel mirrors the FlashInfer CuTe DSL RMSNorm, so draft_front serves a call only when
+the native ``enorm``/``hnorm`` would run that kernel. :func:`unsupported_reason`
+checks the per-process conditions on every call; the caller checks the norm
+modules once with :func:`~.native_norm.native_norm_reason`.
+"""
+
+from typing import NamedTuple
+
+import torch
+
+from ... import dispatch
+from .native_norm import flashinfer_path_reason
+
+_SUPPORTED_DTYPES = (torch.bfloat16,)
+_VEC_ELEMS = 8  # 16-byte vectors of BF16
+_MAX_HIDDEN_SIZE = 16384  # one CTA per row, as in the native norm kernel
+
+
+class Segment(NamedTuple):
+    """Columns ``[begin, end)`` of concat source ``source`` (0: embeds, 1: hidden)."""
+
+    source: int
+    begin: int
+    end: int
+
+
+def plan_segments(hidden_size: int, tp_size: int, tp_rank: int) -> tuple[Segment, ...]:
+    """Map rank ``tp_rank``'s chunk of the ``[2H]`` concatenation onto its sources.
+
+    ``tp_size == 1`` means the whole concatenation (the unsharded native input).
+
+    Raises:
+        ValueError: if ``2 * hidden_size`` is not divisible by ``tp_size`` or the
+            rank is out of range.
+    """
+    total = 2 * hidden_size
+    if tp_size < 1 or total % tp_size != 0:
+        raise ValueError(f"2*hidden_size={total} is not divisible by tp_size={tp_size}")
+    if not 0 <= tp_rank < tp_size:
+        raise ValueError(f"tp_rank={tp_rank} out of range for tp_size={tp_size}")
+    width = total // tp_size
+    lo, hi = tp_rank * width, (tp_rank + 1) * width
+    segments = []
+    for source in (0, 1):
+        src_lo, src_hi = source * hidden_size, (source + 1) * hidden_size
+        begin, end = max(lo, src_lo), min(hi, src_hi)
+        if begin < end:
+            segments.append(Segment(source, begin - src_lo, end - src_lo))
+    return tuple(segments)
+
+
+def _aligned(t: torch.Tensor) -> bool:
+    return t.data_ptr() % 16 == 0
+
+
+def _row_operand_reason(name: str, t: torch.Tensor, rows: int, hidden_size: int) -> str | None:
+    if t.dim() != 2 or t.shape[0] != rows or t.shape[1] != hidden_size:
+        return f"{name} shape {tuple(t.shape)} != ({rows}, {hidden_size})"
+    if t.stride(0) % _VEC_ELEMS != 0:
+        return f"{name} row stride {t.stride(0)} not a multiple of {_VEC_ELEMS}"
+    if t.stride(1) != 1:
+        return f"{name} column stride {t.stride(1)} != 1"
+    if not _aligned(t):
+        return f"{name} not 16-byte aligned"
+    return None
+
+
+def unsupported_reason(
+    embeds: torch.Tensor,
+    hidden_states: torch.Tensor,
+    enorm_weight: torch.Tensor,
+    hnorm_weight: torch.Tensor,
+    tp_size: int,
+    tp_rank: int,
+) -> str | None:
+    """Return why the candidate cannot serve these operands, or ``None`` if it can.
+
+    The checks are host-side only and launch nothing, so a caller can take the
+    native path with no side effect from the candidate. They include whether a
+    native BF16 ``RMSNorm`` on this device runs the FlashInfer CuTe DSL kernel
+    that draft_front mirrors; the per-module conditions (gemma, cuda.tile, NVFP4) are the
+    caller's, see :func:`~.native_norm.native_norm_reason`.
+    """
+    if embeds.dim() != 2:
+        return f"embeds must be 2-D, got {embeds.dim()}-D"
+    rows, hidden_size = embeds.shape
+    dtype = embeds.dtype
+    if dtype not in _SUPPORTED_DTYPES:
+        return f"dtype {dtype} not supported"
+    tensors = (embeds, hidden_states, enorm_weight, hnorm_weight)
+    if any(t.dtype != dtype for t in tensors):
+        return "operand dtypes differ"
+    if any(not t.is_cuda or t.device != embeds.device for t in tensors):
+        return "operands must share one CUDA device"
+    if hidden_size % _VEC_ELEMS != 0 or hidden_size > _MAX_HIDDEN_SIZE:
+        return f"hidden_size {hidden_size} outside the supported domain"
+    if tp_size < 1 or (2 * hidden_size) % tp_size != 0 or not 0 <= tp_rank < tp_size:
+        return f"tp_size={tp_size}, tp_rank={tp_rank} invalid for hidden_size {hidden_size}"
+    if (2 * hidden_size // tp_size) % _VEC_ELEMS != 0:
+        return f"chunk width {2 * hidden_size // tp_size} not a multiple of {_VEC_ELEMS}"
+    for name, t in (("embeds", embeds), ("hidden_states", hidden_states)):
+        reason = _row_operand_reason(name, t, rows, hidden_size)
+        if reason is not None:
+            return reason
+    for name, w in (("enorm_weight", enorm_weight), ("hnorm_weight", hnorm_weight)):
+        if w.dim() != 1 or w.shape[0] != hidden_size or not w.is_contiguous() or not _aligned(w):
+            return f"{name} must be a contiguous, 16-byte aligned [{hidden_size}] tensor"
+    return flashinfer_path_reason(embeds.device)
+
+
+def _enable_pdl() -> bool:
+    from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+    return get_env_enable_pdl()
+
+
+def pack_draft_front(
+    embeds: torch.Tensor,
+    hidden_states: torch.Tensor,
+    enorm_weight: torch.Tensor,
+    enorm_eps: float,
+    hnorm_weight: torch.Tensor,
+    hnorm_eps: float,
+    tp_rank: int,
+    tp_size: int,
+) -> torch.Tensor | None:
+    """Build the TP-local ``eh_proj`` input of the MTP draft layer.
+
+    Args:
+        embeds: ``[R, H]`` BF16 draft-token embeddings (native ``embed_tokens``
+            output).
+        hidden_states: ``[R, H]`` BF16 hidden states entering the draft layer.
+        enorm_weight: ``[H]`` BF16 weight of ``enorm``.
+        enorm_eps: epsilon of ``enorm``.
+        hnorm_weight: ``[H]`` BF16 weight of ``hnorm``.
+        hnorm_eps: epsilon of ``hnorm``.
+        tp_rank: this rank's index in the ``eh_proj`` TP group.
+        tp_size: size of that group; 1 returns the full concatenation.
+
+    Returns:
+        A ``[R, 2H / tp_size]`` view equal to native chunk ``tp_rank`` of
+        ``concat([enorm(embeds), hnorm(hidden_states)], -1)``, with the native
+        chunk's strides; or ``None`` when the operands are outside the supported
+        domain, when the native norm on this device would not run the FlashInfer
+        CuTe DSL kernel, or when the kernel was not compiled before CUDA-graph
+        capture (nothing was launched; run the native path).
+    """
+    reason = unsupported_reason(embeds, hidden_states, enorm_weight, hnorm_weight, tp_size, tp_rank)
+    if reason is not None:
+        dispatch.log_capability("draft_front", "draft_front_call_fallback", False, reason)
+        return None
+    rows, hidden_size = embeds.shape
+    kernel = None
+    if rows > 0:
+        from .kernels import get_compiled_kernel
+
+        segments = tuple(
+            (s.source, s.begin, s.end) for s in plan_segments(hidden_size, tp_size, tp_rank)
+        )
+        kernel = get_compiled_kernel(
+            embeds.dtype,
+            hidden_size,
+            segments,
+            _enable_pdl(),
+            allow_compile=not torch.cuda.is_current_stream_capturing(),
+        )
+        if kernel is None:
+            dispatch.log_capability(
+                "draft_front",
+                "draft_front_call_fallback",
+                False,
+                "kernel not compiled before CUDA graph capture",
+            )
+            return None
+    dispatch.log_capability("draft_front", "draft_front_call", True)
+    width = 2 * hidden_size // tp_size
+    concat = torch.empty((rows, 2 * hidden_size), dtype=embeds.dtype, device=embeds.device)
+    if kernel is not None:
+        kernel(
+            embeds,
+            hidden_states,
+            enorm_weight,
+            hnorm_weight,
+            concat[:, :hidden_size],
+            concat[:, hidden_size:],
+            rows,
+            float(enorm_eps),
+            float(hnorm_eps),
+        )
+    return concat[:, tp_rank * width : (tp_rank + 1) * width]
+
+
+__all__ = ["Segment", "pack_draft_front", "plan_segments", "unsupported_reason"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py
new file mode 100644
index 000000000000..85077c66c76e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py
@@ -0,0 +1,445 @@
+# 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.
+"""CuTe DSL kernel for the draft_front: RMSNorm of the needed source + TP-local packing.
+
+The MTP draft layer feeds ``eh_proj`` with column chunk ``tp_rank`` of
+``concat([enorm(embeds), hnorm(hidden)], dim=-1)``. This kernel computes only
+the columns of that chunk. Each CTA handles ``rows_per_block`` rows of one
+source segment: it reads the full source row (the RMS needs it), and stores the
+normalized values of the segment's columns into the caller's concat-layout
+buffer.
+
+The per-row arithmetic deliberately mirrors the FlashInfer CuTe DSL
+``RMSNormKernel`` that TensorRT-LLM's ``RMSNorm`` runs for BF16 inputs: the same
+thread-value layout and vector width, the per-thread ``TensorSSA`` reduction of
+``x * x``, a butterfly warp reduction with offsets 1, 2, 4, 8, 16, a
+shared-memory combine across the warps of a row, ``rsqrt(sum / H + eps)`` with
+the fast-math rsqrt, and ``y = x * rstd * (w + 0)`` rounded once to the output
+dtype. With the same reduction tree the normalized values are bit-identical to
+the native norm, which the unit tests check.
+"""
+
+import math
+from typing import Callable
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import Float32, Int64
+
+_COPY_BITS = 128
+
+# One segment: (source, first column, end column) in source coordinates. Source 0
+# is the embedding half of the concatenation, source 1 the hidden half.
+Segment = tuple[int, int, int]
+
+
+def threads_per_row(hidden_size: int) -> int:
+    """Threads that cooperate on one row; same table as the native norm kernel."""
+    if hidden_size <= 64:
+        return 8
+    if hidden_size <= 128:
+        return 16
+    if hidden_size <= 3072:
+        return 32
+    if hidden_size <= 6144:
+        return 64
+    if hidden_size <= 16384:
+        return 128
+    return 256
+
+
+def threads_per_block(hidden_size: int) -> int:
+    """CTA size; same table as the native norm kernel."""
+    return 128 if hidden_size <= 16384 else 256
+
+
+@cute.jit
+def _warp_reduce_sum(val: Float32, width: cutlass.Constexpr[int]) -> Float32:
+    for i in cutlass.range_constexpr(int(math.log2(width))):
+        val = val + cute.arch.shuffle_sync_bfly(val, offset=1 << i)
+    return val
+
+
+@cute.jit
+def _row_reduce_sum(
+    x_sq: cute.TensorSSA,
+    row_threads: cutlass.Constexpr[int],
+    reduction_buffer: cute.Tensor,
+) -> Float32:
+    """Sum ``x_sq`` over the threads of one row (the row's warps share an smem slot row)."""
+    local_val = x_sq.reduce(cute.ReductionOp.ADD, init_val=Float32(0.0), reduction_profile=0)
+    warp_val = _warp_reduce_sum(local_val, min(row_threads, 32))
+    result = warp_val
+    if cutlass.const_expr(row_threads > 32):
+        lane_idx = cute.arch.lane_idx()
+        warp_idx = cute.arch.warp_idx()
+        warps_per_row = cute.size(reduction_buffer.shape[1])
+        row_idx = warp_idx // warps_per_row
+        col_idx = warp_idx % warps_per_row
+        if lane_idx == 0:
+            reduction_buffer[row_idx, col_idx] = warp_val
+        cute.arch.barrier()
+        block_val = Float32(0.0)
+        if lane_idx < warps_per_row:
+            block_val = reduction_buffer[row_idx, lane_idx]
+        result = _warp_reduce_sum(block_val, 32)
+    return result
+
+
+@cute.jit
+def _column_predicate(
+    tXcX: cute.Tensor, begin: cutlass.Constexpr[int], end: cutlass.Constexpr[int]
+) -> cute.Tensor:
+    """Per-vector predicate ``begin <= column < end`` for a partitioned coordinate tensor."""
+    pred = cute.make_rmem_tensor(
+        cute.make_layout(
+            (
+                cute.size(tXcX, mode=[0, 1]),
+                cute.size(tXcX, mode=[1]),
+                cute.size(tXcX, mode=[2]),
+            ),
+            stride=(cute.size(tXcX, mode=[2]), 0, 1),
+        ),
+        cutlass.Boolean,
+    )
+    for rest_v in cutlass.range_constexpr(pred.shape[0]):
+        for rest_k in cutlass.range_constexpr(pred.shape[2]):
+            col = tXcX[(0, rest_v), 0, rest_k][1]
+            if cutlass.const_expr(begin == 0):
+                pred[rest_v, 0, rest_k] = cute.elem_less(col, end)
+            else:
+                pred[rest_v, 0, rest_k] = cute.elem_less(col, end) & cute.elem_less(begin - 1, col)
+    return pred
+
+
+class DraftFrontPackKernel:
+    """Normalize the needed concat source and write this rank's ``eh_proj`` input columns.
+
+    Args:
+        dtype: element type of the activations, weights and output (BF16 in the
+            supported domain).
+        hidden_size: ``H``; each source row has ``H`` elements.
+        segments: one or two ``(source, begin, end)`` triples; grid dimension y
+            enumerates them.
+        enable_pdl: wait on / trigger programmatic dependent launch like the
+            native norm kernel.
+    """
+
+    def __init__(
+        self,
+        dtype: cutlass.Numeric,
+        hidden_size: int,
+        segments: tuple[Segment, ...],
+        enable_pdl: bool,
+    ):
+        if len(segments) not in (1, 2):
+            raise ValueError(f"expected one or two segments, got {segments}")
+        self.dtype = dtype
+        self.H = hidden_size
+        self.segments = segments
+        self.enable_pdl = enable_pdl
+
+        elem_bits = dtype.width
+        max_vec_size = _COPY_BITS // elem_bits
+        h_align = hidden_size & (-hidden_size)
+        self.vec_size = min(h_align, max_vec_size)
+        self.copy_bits = self.vec_size * elem_bits
+        self.threads_per_row = threads_per_row(hidden_size)
+        self.num_threads = threads_per_block(hidden_size)
+        self.rows_per_block = self.num_threads // self.threads_per_row
+        self.warps_per_row = max(self.threads_per_row // 32, 1)
+        self.num_vec_blocks = max(
+            1, (hidden_size // self.vec_size + self.threads_per_row - 1) // self.threads_per_row
+        )
+        self.cols_per_tile = self.vec_size * self.num_vec_blocks * self.threads_per_row
+
+    def _tv_layout(self) -> cute.Layout:
+        shape = (
+            (self.threads_per_row, self.rows_per_block),
+            (self.vec_size, self.num_vec_blocks),
+        )
+        stride = (
+            (self.vec_size * self.rows_per_block, 1),
+            (self.rows_per_block, self.rows_per_block * self.vec_size * self.threads_per_row),
+        )
+        return cute.make_layout(shape, stride=stride)
+
+    def smem_bytes(self) -> int:
+        return self.rows_per_block * self.warps_per_row * 4
+
+    @cute.jit
+    def __call__(
+        self,
+        mE: cute.Tensor,
+        mHs: cute.Tensor,
+        mWe: cute.Tensor,
+        mWh: cute.Tensor,
+        mYe: cute.Tensor,
+        mYh: cute.Tensor,
+        M: Int64,
+        eps_e: Float32,
+        eps_h: Float32,
+        stream,
+    ):
+        tv_layout = self._tv_layout()
+        tiler_mn = (self.rows_per_block, self.cols_per_tile)
+        self.kernel(mE, mHs, mWe, mWh, mYe, mYh, M, eps_e, eps_h, tv_layout, tiler_mn).launch(
+            grid=[cute.ceil_div(M, self.rows_per_block), len(self.segments), 1],
+            block=[self.num_threads, 1, 1],
+            smem=self.smem_bytes(),
+            stream=stream,
+            use_pdl=self.enable_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mE: cute.Tensor,
+        mHs: cute.Tensor,
+        mWe: cute.Tensor,
+        mWh: cute.Tensor,
+        mYe: cute.Tensor,
+        mYh: cute.Tensor,
+        M: Int64,
+        eps_e: Float32,
+        eps_h: Float32,
+        tv_layout: cute.Layout,
+        tiler_mn: cute.Shape,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, seg_idx, _ = cute.arch.block_idx()
+
+        if cutlass.const_expr(self.enable_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        smem = cutlass.utils.SmemAllocator()
+        reduction_buffer = smem.allocate_tensor(
+            Float32,
+            cute.make_layout((self.rows_per_block, self.warps_per_row)),
+            byte_alignment=4,
+        )
+
+        # Source 0 is the embedding half of the concatenation, source 1 the hidden half.
+        operands = ((mE, mWe, mYe, eps_e), (mHs, mWh, mYh, eps_h))
+        seg0 = self.segments[0]
+        if cutlass.const_expr(len(self.segments) == 1):
+            self._normalize_segment(
+                *operands[seg0[0]],
+                M,
+                seg0[1],
+                seg0[2],
+                tidx,
+                bidx,
+                reduction_buffer,
+                tv_layout,
+                tiler_mn,
+            )
+        else:
+            seg1 = self.segments[1]
+            if seg_idx == 0:
+                self._normalize_segment(
+                    *operands[seg0[0]],
+                    M,
+                    seg0[1],
+                    seg0[2],
+                    tidx,
+                    bidx,
+                    reduction_buffer,
+                    tv_layout,
+                    tiler_mn,
+                )
+            else:
+                self._normalize_segment(
+                    *operands[seg1[0]],
+                    M,
+                    seg1[1],
+                    seg1[2],
+                    tidx,
+                    bidx,
+                    reduction_buffer,
+                    tv_layout,
+                    tiler_mn,
+                )
+
+        if cutlass.const_expr(self.enable_pdl):
+            cute.arch.griddepcontrol_launch_dependents()
+
+    @cute.jit
+    def _normalize_segment(
+        self,
+        mX: cute.Tensor,
+        mW: cute.Tensor,
+        mY: cute.Tensor,
+        eps: Float32,
+        M: Int64,
+        begin: cutlass.Constexpr[int],
+        end: cutlass.Constexpr[int],
+        tidx,
+        bidx,
+        reduction_buffer: cute.Tensor,
+        tv_layout: cute.Layout,
+        tiler_mn: cute.Shape,
+    ):
+        H = self.H
+
+        idX = cute.make_identity_tensor(mX.shape)
+        gX = cute.local_tile(mX, tiler_mn, (bidx, 0))
+        gY = cute.local_tile(mY, tiler_mn, (bidx, 0))
+        cX = cute.local_tile(idX, tiler_mn, (bidx, 0))
+
+        mW_expanded_layout = cute.prepend(mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+        mW_2d = cute.make_tensor(mW.iterator, mW_expanded_layout)
+        gW = cute.local_tile(mW_2d, tiler_mn, (0, 0))
+
+        copy_atom_load = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), mX.element_type, num_bits_per_copy=self.copy_bits
+        )
+        copy_atom_store = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), mY.element_type, num_bits_per_copy=self.copy_bits
+        )
+        tiled_copy_load = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn)
+        tiled_copy_w = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn)
+        tiled_copy_store = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn)
+
+        thr_copy_x = tiled_copy_load.get_slice(tidx)
+        thr_copy_w = tiled_copy_w.get_slice(tidx)
+        thr_copy_o = tiled_copy_store.get_slice(tidx)
+
+        tXgX = thr_copy_x.partition_S(gX)
+        tXcX = thr_copy_x.partition_S(cX)
+        tXrX = cute.make_fragment_like(tXgX)
+
+        tWgW = thr_copy_w.partition_S(gW)
+        tWrW = cute.make_fragment_like(tWgW)
+        tXrW = thr_copy_x.retile(tWrW)
+
+        tXgO = thr_copy_o.partition_D(gY)
+        tXrO = cute.make_fragment_like(tXgO)
+
+        # Full-row columns feed the reduction; only [begin, end) is weighted and stored.
+        tXpX = _column_predicate(tXcX, 0, H)
+        tWcW = thr_copy_w.partition_S(cX)
+        tWpW = _column_predicate(tWcW, begin, end)
+        tXpO = _column_predicate(tXcX, begin, end)
+        row_coord = tXcX[(0, 0), 0, 0]
+        row_in_bounds = row_coord[0] < M
+
+        tXrX.store(cute.zeros_like(tXrX, dtype=mX.element_type))
+        if row_in_bounds:
+            cute.copy(copy_atom_load, tXgX, tXrX, pred=tXpX)
+        tWrW.store(cute.zeros_like(tWrW, dtype=mW.element_type))
+        cute.copy(copy_atom_load, tWgW, tWrW, pred=tWpW)
+
+        x = tXrX.load().to(Float32)
+        x_sq = x * x
+        sum_sq = _row_reduce_sum(x_sq, self.threads_per_row, reduction_buffer)
+
+        mean_sq = sum_sq / Float32(H)
+        rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True)
+
+        w = tXrW.load().to(Float32)
+        y = x * rstd * (w + Float32(0.0))
+        tXrO.store(y.to(mY.element_type))
+
+        if row_in_bounds:
+            cute.copy(copy_atom_store, tXrO, tXgO, pred=tXpO)
+
+
+def _torch_to_cutlass_dtype(dtype: torch.dtype) -> cutlass.Numeric:
+    mapping: dict[torch.dtype, cutlass.Numeric] = {
+        torch.bfloat16: cutlass.BFloat16,
+        torch.float16: cutlass.Float16,
+    }
+    if dtype not in mapping:
+        raise ValueError(f"unsupported dtype {dtype}")
+    return mapping[dtype]
+
+
+_COMPILED: dict[tuple, Callable] = {}
+
+
+def get_compiled_kernel(
+    dtype: torch.dtype,
+    hidden_size: int,
+    segments: tuple[Segment, ...],
+    enable_pdl: bool,
+    allow_compile: bool = True,
+) -> Callable | None:
+    """Return the draft-front pack kernel, compiling it once per process and configuration.
+
+    The returned callable takes torch tensors through TVM-FFI and launches on the
+    current torch stream:
+    ``fn(embeds[M,H], hidden[M,H], enorm_w[H], hnorm_w[H], out_e[M,H], out_h[M,H], M, eps_e,
+    eps_h)``. Every ``[M, H]`` operand needs unit column stride and a row stride
+    divisible by the vector width.
+
+    Args:
+        allow_compile: when False and the kernel is not compiled yet, return
+            ``None`` instead of compiling (used during CUDA-graph capture).
+    """
+    key = (dtype, hidden_size, segments, enable_pdl)
+    kernel = _COMPILED.get(key)
+    if kernel is None and allow_compile:
+        kernel = _compile_kernel(*key)
+        _COMPILED[key] = kernel
+    return kernel
+
+
+def _compile_kernel(
+    dtype: torch.dtype, hidden_size: int, segments: tuple[Segment, ...], enable_pdl: bool
+) -> Callable:
+    cutlass_dtype = _torch_to_cutlass_dtype(dtype)
+    kernel_obj = DraftFrontPackKernel(cutlass_dtype, hidden_size, segments, enable_pdl)
+    vec = kernel_obj.vec_size
+    sym_m = cute.sym_int(64)
+
+    def row_major_fake() -> cute.Tensor:
+        return cute.runtime.make_fake_tensor(
+            cutlass_dtype,
+            (sym_m, hidden_size),
+            (cute.sym_int64(divisibility=vec), 1),
+            assumed_align=16,
+        )
+
+    def weight_fake() -> cute.Tensor:
+        return cute.runtime.make_fake_compact_tensor(
+            cutlass_dtype, (hidden_size,), assumed_align=16
+        )
+
+    stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        kernel_obj,
+        row_major_fake(),
+        row_major_fake(),
+        weight_fake(),
+        weight_fake(),
+        row_major_fake(),
+        row_major_fake(),
+        Int64(1),
+        Float32(1e-6),
+        Float32(1e-6),
+        stream_fake,
+        options="--enable-tvm-ffi",
+    )
+
+
+__all__ = [
+    "DraftFrontPackKernel",
+    "Segment",
+    "get_compiled_kernel",
+    "threads_per_block",
+    "threads_per_row",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/native_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/native_norm.py
new file mode 100644
index 000000000000..49f4b3f78f2d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/native_norm.py
@@ -0,0 +1,137 @@
+# 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.
+"""Which path the native ``enorm``/``hnorm`` take, and whether draft_front mirrors it.
+
+draft_front is bit-identical to the native draft front because its kernel reproduces the
+FlashInfer CuTe DSL ``RMSNormKernel`` that TensorRT-LLM's ``RMSNorm.forward``
+reaches for a 2-D BF16 input without a residual. ``RMSNorm.forward`` can take
+other paths: the fused NVFP4-quantize epilogue, ``cuda.tile``, the gemma form
+``x * rstd * (w + 1)``, or the torch fallback when FlashInfer is unavailable.
+FlashInfer itself can select its CUDA JIT norm instead of the CuTe DSL kernel.
+Each of those rounds differently from draft_front (gemma also changes the result), so
+draft_front must decline before launching anything whenever one of them would run.
+
+* :func:`norm_module_reason` covers the per-module attributes; only a caller
+  holding the ``RMSNorm`` modules can check them (the ``DeepseekV3MTP`` adapter,
+  once at construction).
+* :func:`flashinfer_path_reason` covers the per-process and per-device
+  conditions; :func:`~.draft_front.unsupported_reason` applies it on every call.
+* :func:`native_norm_reason` combines both.
+"""
+
+import functools
+
+import torch
+from torch import nn
+
+# FlashInfer releases whose CuTe DSL ``RMSNormKernel`` the draft_front kernel was written
+# against and verified bit for bit (thread-value layout, reduction tree, rsqrt).
+VERIFIED_FLASHINFER_VERSIONS = ("0.6.18",)
+
+
+def norm_module_reason(name: str, norm: nn.Module) -> str | None:
+    """Return why ``norm`` would not run the plain FlashInfer RMSNorm, or ``None``.
+
+    Args:
+        name: label used in the returned reason (``"enorm"`` or ``"hnorm"``).
+        norm: the module the native draft front calls as ``norm(x)``.
+    """
+    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+
+    if not isinstance(norm, RMSNorm) or type(norm).forward is not RMSNorm.forward:
+        return f"{name} is {type(norm).__name__}, not the TensorRT-LLM RMSNorm"
+    if norm.use_cuda_tile:
+        return f"{name} runs the cuda.tile RMSNorm"
+    if norm.use_gemma:
+        return f"{name} is a gemma RMSNorm (weight + 1)"
+    if norm.is_nvfp4:
+        return f"{name} can take the fused NVFP4-quantize RMSNorm"
+    if norm.return_hp_output:
+        return f"{name} requests a high-precision auxiliary output"
+    return None
+
+
+@functools.cache
+def _flashinfer_path_reason(device_index: int) -> str | None:
+    # Every input is fixed for the life of the process: the import-time
+    # FlashInfer availability, its version, FLASHINFER_USE_CUDA_NORM (read when
+    # flashinfer.norm is imported) and the device's compute capability.
+    from tensorrt_llm._torch.modules import rms_norm
+
+    if not rms_norm.IS_FLASHINFER_AVAILABLE:
+        return "FlashInfer unavailable: native RMSNorm takes the torch path"
+    import flashinfer
+    import flashinfer.norm as flashinfer_norm
+
+    version = getattr(flashinfer, "__version__", None)
+    if version not in VERIFIED_FLASHINFER_VERSIONS:
+        return (
+            f"FlashInfer {version} not verified; the draft_front kernel mirrors the RMSNorm of "
+            f"{', '.join(VERIFIED_FLASHINFER_VERSIONS)}"
+        )
+    use_cuda_norm = getattr(flashinfer_norm, "_use_cuda_norm", None)
+    if use_cuda_norm is None:
+        return "flashinfer.norm has no _use_cuda_norm dispatch"
+    if use_cuda_norm(torch.device("cuda", device_index)):
+        return (
+            "FlashInfer selects its CUDA JIT RMSNorm (FLASHINFER_USE_CUDA_NORM=1, CuTe DSL "
+            "unavailable, or device not supported by the CuTe DSL)"
+        )
+    return None
+
+
+def flashinfer_path_reason(device: torch.device) -> str | None:
+    """Return why a native BF16 ``RMSNorm`` on ``device`` would not run FlashInfer's
+    CuTe DSL ``RMSNormKernel``, or ``None`` if it would.
+
+    Host-side only; the result is cached per device for the life of the process.
+    """
+    if device.type != "cuda":
+        return f"device {device} is not a CUDA device"
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    return _flashinfer_path_reason(index)
+
+
+def native_norm_reason(
+    enorm: nn.Module, hnorm: nn.Module, device: torch.device | None = None
+) -> str | None:
+    """Return why the native ``enorm``/``hnorm`` would not take the path draft_front mirrors.
+
+    Args:
+        enorm: the draft layer's embedding norm module.
+        hnorm: the draft layer's hidden-state norm module.
+        device: the CUDA device the norms run on; defaults to the current device.
+
+    Returns:
+        A reason string when draft_front must leave the draft front to the native path,
+        else ``None``.
+    """
+    for name, norm in (("enorm", enorm), ("hnorm", hnorm)):
+        reason = norm_module_reason(name, norm)
+        if reason is not None:
+            return reason
+    if device is None:
+        if not torch.cuda.is_available():
+            return "CUDA unavailable"
+        device = torch.device("cuda", torch.cuda.current_device())
+    return flashinfer_path_reason(device)
+
+
+__all__ = [
+    "VERIFIED_FLASHINFER_VERSIONS",
+    "flashinfer_path_reason",
+    "native_norm_reason",
+    "norm_module_reason",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/__init__.py
new file mode 100644
index 000000000000..09c970ba2a3e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/__init__.py
@@ -0,0 +1,50 @@
+# 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.
+"""Package draft_live_rows: MTP draft step-0 live-row compaction at the post-attention hop.
+
+Entry points registered with ``dispatch``:
+
+* ``make_live_row_hop(post_attn_hop)`` returns the :class:`LiveRowHop` on the
+  workspace of the MTP layer's post_attention_hop, or None (the layer keeps the post_attention_hop).
+  Call it at the same point of model construction on every rank: the first
+  call for a post_attention_hop agrees over its TP communicator whether every rank
+  carries the operator, so the answer is the same on all ranks.
+* ``select_step0(worker, draft_model, spec_metadata, draft_loop_ops)`` returns the
+  draft step-0 arming (:class:`Step0LiveRows`) of one one-model MTP forward,
+  or None to keep every row; decided on the host before any device work.
+* ``engagement_counts()`` returns the host-side call counters (eager calls
+  and CUDA graph captures).
+
+The package imports no other package's code.
+"""
+
+from ...dispatch import register
+from .hop import LiveRowHop, engagement_counts, op_available, required_sms, reset_engagement_counts
+from .step0 import Step0LiveRows, select_step0, step0_ineligibility
+
+register("draft_live_rows", "make_live_row_hop", LiveRowHop.for_post_attn_hop)
+register("draft_live_rows", "select_step0", select_step0)
+register("draft_live_rows", "engagement_counts", engagement_counts)
+
+__all__ = [
+    "LiveRowHop",
+    "Step0LiveRows",
+    "engagement_counts",
+    "op_available",
+    "required_sms",
+    "reset_engagement_counts",
+    "select_step0",
+    "step0_ineligibility",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/hop.py
new file mode 100644
index 000000000000..600671f8e2b0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/hop.py
@@ -0,0 +1,277 @@
+# 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.
+"""Live-row post-attention hop of the MTP layer (draft_live_rows).
+
+At MTP draft step 0 the layer runs ``B * (K + 1)`` rows (every draft input
+token of every request), but after the layer only the last accepted row of
+each request is read: the draft LM head and the next step consume
+``hidden_states[gather_ids]``, with ``gather_ids`` written on the device by
+draft_loop's accept kernel. Attention needs every row (it appends their KV), while
+everything after the post-attention hop is per row. :class:`LiveRowHop`
+therefore replaces the post_attention_hop of that call by a fork that reads ``gather_ids``
+after its programmatic dependency wait and reduces, adds the residual,
+normalizes and computes router logits for the ``B`` live rows only; the MoE
+and the post-MoE all-reduce that follow then run on ``B`` rows.
+
+Every per-row quantity is computed with the post_attention_hop's arithmetic (fixed rank
+order, the native RMS tree, the router MMA with the 8 slices summed in a
+fixed order), so live row ``r`` equals row ``gather_ids[r]`` of the post_attention_hop
+over all rows, bit for bit. The fork runs on the post_attention_hop's workspace and
+device-side protocol state; post_attention_hop and draft_live_rows calls may follow each other in
+any order on the stream that issues the TP group's hops.
+
+Eligibility is decided from rank-uniform host facts only (the selection, the
+post_attention_hop's agreed SM budget, shapes and dtypes), before any device work. The
+one rank-local fact, whether this process carries the operator, is agreed
+over the post_attention_hop's TP communicator when the live-row hop is built: it exists on
+every rank of the group or on none.
+"""
+
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.logger import logger
+
+from ...dispatch import log_capability
+
+FEATURE = "draft_live_rows"
+HIDDEN_SIZE = 6144
+NUM_RANKS = 8
+NUM_EXPERTS = 256
+MAX_ROWS = 8
+# The post_attention_hop workspace: [3 Lamport buffers][8 rows][8 ranks][6144] bf16.
+WORKSPACE_BYTES = 3 * MAX_ROWS * NUM_RANKS * HIDDEN_SIZE * 2
+# CTAs per cluster and the router layout's expert / row groups (kernel).
+CLUSTER_CTAS = 8
+ROUTER_EXPERT_GROUPS = 4
+ROUTER_MAX_ROW_GROUPS = 3
+_OP_NAME = "glm_moe_dsa_draft_live_row_hop"
+
+# Host-side call counters. They count Python calls (eager runs and CUDA graph
+# captures), not graph replays.
+_counts: Dict[str, int] = {}
+
+
+def _count(key: str) -> None:
+    _counts[key] = _counts.get(key, 0) + 1
+
+
+def engagement_counts() -> Dict[str, int]:
+    """Copy of the host-side call counters.
+
+    Keys: ``engaged`` (live-row hop calls), ``declined`` and
+    ``declined:<reason>`` (armed calls that ran the post_attention_hop instead),
+    ``step0_armed`` and ``step0_compacted`` (draft step-0 calls armed by the
+    draft loop and those whose MTP layer returned the live rows).
+    """
+    return dict(_counts)
+
+
+def reset_engagement_counts() -> None:
+    """Clear the host-side call counters."""
+    _counts.clear()
+
+
+def required_sms(rows: int) -> int:
+    """SMs one live-row hop of ``rows`` live rows needs (every CTA co-resident)."""
+    return CLUSTER_CTAS * max(rows, ROUTER_EXPERT_GROUPS * min(rows, ROUTER_MAX_ROW_GROUPS))
+
+
+def op_available() -> bool:
+    """Whether the compiled ``trtllm::glm_moe_dsa_draft_live_row_hop`` op is loaded."""
+    return hasattr(torch.ops.trtllm, _OP_NAME)
+
+
+def all_ranks_true(comm: Any, local: bool) -> bool:
+    """Collective AND of ``local`` over ``comm``, the post_attention_hop's TP communicator.
+
+    Every rank of ``comm`` calls it at the same point, whatever its own
+    ``local`` is.
+    """
+    from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded
+
+    return _mnnvl_workspace_all_succeeded(comm, bool(local))
+
+
+def _register_fake() -> None:
+    if not op_available():
+        return
+    try:
+
+        @torch.library.register_fake(f"trtllm::{_OP_NAME}")
+        def _(
+            input,
+            residual,
+            row_ids,
+            gamma,
+            eps,
+            router_weight,
+            uc_buffer,
+            mc_ptr,
+            flags,
+            rank,
+            nranks,
+        ):
+            rows = row_ids.shape[0]
+            return [
+                input.new_empty((rows, HIDDEN_SIZE)),
+                residual.new_empty((rows, HIDDEN_SIZE)),
+                input.new_empty((rows, NUM_EXPERTS), dtype=torch.float32),
+            ]
+    except RuntimeError as error:
+        # Already registered by an earlier import in this process.
+        logger.debug(f"[glm_moe_dsa] draft_live_rows fake registration skipped: {error}")
+
+
+class LiveRowHop:
+    """The live-row hop on the workspace of a post_attention_hop (one per TP group)."""
+
+    # id(post_attention_hop) -> (post_attention_hop, live-row hop or None, reason): one decision per
+    # post_attention_hop.
+    _instances: Dict[int, Tuple[Any, Optional["LiveRowHop"], str]] = {}
+
+    @classmethod
+    def for_post_attn_hop(cls, post_attn_hop: Any) -> Optional["LiveRowHop"]:
+        """Return the live-row hop that shares ``post_attn_hop``'s workspace, or None.
+
+        Must be called at the same point of model construction on every rank
+        of the post_attention_hop's TP group. The post_attention_hop exists on every rank or on
+        none; whether the operator is loaded is agreed over the post_attention_hop's
+        communicator the first time a post_attention_hop is seen (one collective, also on
+        a rank without the operator), so either every rank gets the live-row
+        hop or every rank keeps the post_attention_hop. Later calls for the same post_attention_hop
+        return the recorded decision without communicating.
+        """
+        if post_attn_hop is None or not getattr(post_attn_hop, "enabled", False):
+            log_capability(
+                FEATURE, "live_row_hop:no post_attention_hop", False, "no post_attention_hop hop"
+            )
+            return None
+        key = id(post_attn_hop)
+        entry = cls._instances.get(key)
+        if entry is None or entry[0] is not post_attn_hop:
+            hop, reason = cls._agreed(post_attn_hop)
+            entry = (post_attn_hop, hop, reason)
+            cls._instances[key] = entry
+            if hop is None:
+                logger.warning(
+                    f"[glm_moe_dsa] draft_live_rows live-row hop disabled on every rank: {reason}"
+                )
+        _, hop, reason = entry
+        log_capability(FEATURE, "live_row_hop", hop is not None, reason)
+        return hop
+
+    @classmethod
+    def _agreed(cls, post_attn_hop: Any) -> Tuple[Optional["LiveRowHop"], str]:
+        """Build the live-row hop if every rank of the post_attention_hop's group carries the operator."""
+        comm = getattr(post_attn_hop, "comm", None)
+        if comm is None:
+            # An enabled post_attention_hop always has its communicator; without one there is no
+            # group to agree in.
+            return None, "post_attention_hop hop has no communicator"
+        local = op_available()
+        if not all_ranks_true(comm, local):
+            if not local:
+                return None, f"trtllm::{_OP_NAME} is not loaded"
+            return None, f"another rank lacks trtllm::{_OP_NAME}"
+        return cls(post_attn_hop), ""
+
+    def __init__(self, post_attn_hop: Any):
+        self.w2 = post_attn_hop
+        self.rank = post_attn_hop.rank
+        self.uc_buffer = post_attn_hop.uc_buffer
+        self.mc_ptr = post_attn_hop.mc_ptr
+        self.flags = post_attn_hop.flags
+        self.min_sm_count = post_attn_hop.min_sm_count
+
+    def call_ineligibility(
+        self,
+        hidden_states: torch.Tensor,
+        row_ids: Optional[torch.Tensor],
+        router_weight: Optional[torch.Tensor],
+    ) -> Optional[str]:
+        """Reason this armed call must run the post_attention_hop instead, or None (host only).
+
+        The caller has already established that the post_attention_hop serves the call
+        (same input, residual and norm weight), so only the live-row specific
+        facts are checked here.
+        """
+        if router_weight is None:
+            return "no router weight"
+        if row_ids is None:
+            return "no live rows"
+        if (
+            not isinstance(row_ids, torch.Tensor)
+            or row_ids.dtype != torch.int64
+            or row_ids.dim() != 1
+            or not row_ids.is_contiguous()
+            or row_ids.device != hidden_states.device
+        ):
+            return "live-row ids layout"
+        input_rows = hidden_states.numel() // HIDDEN_SIZE
+        rows = row_ids.shape[0]
+        if not 1 <= rows <= input_rows or input_rows > MAX_ROWS:
+            return f"live rows {rows} of {input_rows}"
+        if required_sms(rows) > self.min_sm_count:
+            return "SM count"
+        return None
+
+    def can_run(
+        self,
+        hidden_states: torch.Tensor,
+        row_ids: Optional[torch.Tensor],
+        router_weight: Optional[torch.Tensor],
+    ) -> bool:
+        """Whether this armed call can use the live-row hop (rank-uniform, no device work)."""
+        reason = self.call_ineligibility(hidden_states, row_ids, router_weight)
+        if reason is not None:
+            _count("declined")
+            _count(f"declined:{reason}")
+            log_capability(FEATURE, f"live_row_hop_call:{reason}", False, reason)
+            return False
+        log_capability(FEATURE, "live_row_hop_call", True)
+        return True
+
+    def __call__(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+        router_weight: torch.Tensor,
+        row_ids: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Run the hop; returns ``(normed, residual_out, router_logits)`` of the live rows.
+
+        Callers check :meth:`can_run` first.
+        """
+        _count("engaged")
+        return getattr(torch.ops.trtllm, _OP_NAME)(
+            hidden_states.contiguous(),
+            residual.contiguous(),
+            row_ids,
+            norm_weight.contiguous(),
+            float(eps),
+            router_weight,
+            self.uc_buffer,
+            self.mc_ptr,
+            self.flags,
+            self.rank,
+            NUM_RANKS,
+        )
+
+
+_register_fake()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/step0.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/step0.py
new file mode 100644
index 000000000000..517a7e15da82
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/step0.py
@@ -0,0 +1,103 @@
+# 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.
+"""Draft-loop side of package draft_live_rows: arm the MTP layer's live-row hop at step 0.
+
+The target's one-model MTP draft loop (``GlmMtpWorker._forward_linear_draft_loop``)
+asks :func:`select_step0` once per forward, before any device work, whether
+draft step 0 may compact the MTP layer to its live rows. The answer is
+rank-uniform: it depends on the selection, on draft_loop's draft-loop ops existing
+for this forward (generation-only batch inside draft_loop's domain), on nothing
+capturing the MTP layer's full output (``layers_to_capture`` empty) and on
+the MTP layer owning a live-row hop. If it returns a :class:`Step0LiveRows`,
+the loop runs step 0 inside :meth:`Step0LiveRows.armed`; when the MTP layer
+took the live-row hop (``compacted``), the layer's output already holds one
+row per request in ``gather_ids`` order, so draft_loop's step-0 row gather becomes the
+identity.
+"""
+
+import contextlib
+from typing import Any, Optional
+
+import torch
+
+from ...dispatch import log_capability
+from .hop import _count
+
+FEATURE = "draft_live_rows"
+
+
+def step0_ineligibility(
+    worker: Any, draft_model: Any, spec_metadata: Any, draft_loop_ops: Any
+) -> Optional[str]:
+    """Reason draft step 0 must keep every row after the MTP layer, or None (host only)."""
+    if not getattr(worker, "is_mtp_eagle", False):
+        return "not MTP Eagle one-model"
+    if draft_loop_ops is None:
+        return "no draft_loop draft-loop ops"
+    if len(getattr(spec_metadata, "layers_to_capture", None) or ()) != 0:
+        return "MTP layer output captured"
+    layers = getattr(draft_model, "mtp_layers", None)
+    if not layers:
+        return "no MTP layer"
+    if getattr(layers[0], "draft_live_row_hop", None) is None:
+        return "MTP layer has no live-row hop"
+    gather_ids = getattr(draft_loop_ops, "gather_ids", None)
+    if (
+        gather_ids is None
+        or gather_ids.dtype != torch.int64
+        or gather_ids.dim() != 1
+        or gather_ids.shape[0] != draft_loop_ops.batch_size
+    ):
+        return "draft_loop gather ids"
+    return None
+
+
+class Step0LiveRows:
+    """Arms the MTP layer with draft_loop's live-row ids for one draft step-0 call.
+
+    The layer takes the ids through ``arm_live_rows`` and reports the outcome through
+    ``disarm_live_rows``.
+    """
+
+    def __init__(self, mtp_layer: Any, gather_ids: torch.Tensor):
+        self.mtp_layer = mtp_layer
+        self.gather_ids = gather_ids
+        self.compacted = False
+
+    @contextlib.contextmanager
+    def armed(self):
+        """Run the step-0 MTP call inside this context; ``compacted`` tells the outcome."""
+        layer = self.mtp_layer
+        layer.arm_live_rows(self.gather_ids)
+        self.compacted = False
+        _count("step0_armed")
+        try:
+            yield self
+        finally:
+            compacted = layer.disarm_live_rows()
+        self.compacted = compacted
+        if compacted:
+            _count("step0_compacted")
+
+
+def select_step0(
+    worker: Any, draft_model: Any, spec_metadata: Any, draft_loop_ops: Any
+) -> Optional[Step0LiveRows]:
+    """The step-0 arming of this forward, or None to keep every row (native draft_loop gather)."""
+    reason = step0_ineligibility(worker, draft_model, spec_metadata, draft_loop_ops)
+    log_capability(FEATURE, f"step0:{reason or 'engaged'}", reason is None, reason or "")
+    if reason is not None:
+        return None
+    return Step0LiveRows(draft_model.mtp_layers[0], draft_loop_ops.gather_ids)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/__init__.py
new file mode 100644
index 000000000000..7854839706da
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/__init__.py
@@ -0,0 +1,46 @@
+# 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.
+"""draft_loop: MTP draft-loop overhead (embedding, bookkeeping, DSA metadata).
+
+Entry points registered with ``dispatch`` (looked up by the call sites only
+after ``dispatch.is_enabled("draft_loop")``):
+
+``draft_embedding(native_embedding) -> callable``
+    Single-launch stand-in for the native draft-token embedding.
+``draft_loop_ops(worker, attn_metadata, spec_metadata) -> DraftLoopOps | None``
+    Per-forward fused acceptance, draft-0 preparation, step-0 row gather,
+    per-step metadata update and post-loop finalize; None keeps the native
+    one-model MTP loop.
+``draft_step_update(...)``
+    The per-step metadata update alone (used by ``DraftLoopOps``).
+"""
+
+from ... import dispatch
+from .draft_loop import DraftLoopOps, make_draft_loop_ops
+from .draft_step import draft_step_update
+from .embedding import DraftEmbedding, draft_embedding, wrap_draft_embedding
+
+dispatch.register("draft_loop", "draft_embedding", wrap_draft_embedding)
+dispatch.register("draft_loop", "draft_loop_ops", make_draft_loop_ops)
+dispatch.register("draft_loop", "draft_step_update", draft_step_update)
+
+__all__ = [
+    "DraftEmbedding",
+    "DraftLoopOps",
+    "draft_embedding",
+    "draft_step_update",
+    "make_draft_loop_ops",
+    "wrap_draft_embedding",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_loop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_loop.py
new file mode 100644
index 000000000000..4bb76882082b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_loop.py
@@ -0,0 +1,450 @@
+# 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.
+"""Fused bookkeeping of the MTP one-model draft loop (draft_loop).
+
+``DraftLoopOps`` is created once per one-model MTP forward. Its host-side
+eligibility check runs before any device work; when it returns None the
+caller keeps the native code everywhere. Otherwise it replaces, with the same
+results:
+
+* acceptance (strict, optionally with the forced-acceptance override) and the
+  draft-0 input ids and gather rows: one launch after the native target
+  sampler, instead of about 26 small kernels;
+* the step-0 last-accepted-row gather of the MTP hidden states: one launch;
+* the per-step position/KV-length/DSA-metadata update (``draft_step``);
+* the post-loop draft-token stack and ``next_new_tokens``: one launch.
+"""
+
+from typing import List, Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from . import draft_step as _draft_step
+from . import verify as _verify
+
+_MAX_DRAFT_LEN = 8
+_MAX_BATCH = 64
+_NUM_WARPS = 4
+
+
+@triton.jit
+def _accept_kernel(
+    target_ptr,  # int32/int64 [B * (K + 1)] sampled target tokens
+    draft_ptr,  # int32 [B * K] previous draft tokens
+    batch_idx_ptr,  # int32 [>= B] spec_metadata.batch_indices_cuda
+    counter_ptr,  # int64 [1] forced-acceptance draw counter (FORCE_MODE 2)
+    pool_ptr,  # float32 [POOL] forced-acceptance uniform pool (FORCE_MODE 2)
+    accepted_ptr,  # int32 [B, K + 1] (out)
+    num_accepted_ptr,  # int32 [B] (out)
+    input_ids_ptr,  # int32 [B * (K + 1)] draft-0 input ids (out)
+    gather_ptr,  # int64 [B] step-0 last-accepted rows (out)
+    frac,  # float32 fractional forced acceptance
+    counter_stride,
+    slot_stride,
+    pool_mask,
+    B: tl.constexpr,
+    K: tl.constexpr,
+    FORCE_MODE: tl.constexpr,
+    BASE_TOTAL: tl.constexpr,
+    MAX_TOTAL: tl.constexpr,
+    BLOCK_B: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    offs_b = tl.arange(0, BLOCK_B)
+    offs_k = tl.arange(0, BLOCK_K)
+    mask_b = offs_b < B
+    row_mask = mask_b[:, None] & (offs_k[None, :] < K + 1)
+    flat = offs_b[:, None] * (K + 1) + offs_k[None, :]
+    target = tl.load(target_ptr + flat, mask=row_mask, other=0)
+    target32 = target.to(tl.int32)
+    tl.store(accepted_ptr + flat, target32, mask=row_mask)
+    tl.store(input_ids_ptr + flat, target32, mask=row_mask)
+
+    draft_mask = mask_b[:, None] & (offs_k[None, :] < K)
+    draft = tl.load(draft_ptr + offs_b[:, None] * K + offs_k[None, :], mask=draft_mask, other=0)
+    mismatch = draft_mask & (draft.to(tl.int64) != target.to(tl.int64))
+    # Leading matches == sum(cumprod(draft == target)).
+    leading = tl.min(tl.where(mismatch, offs_k[None, :], K), 1)
+    num_accepted = 1 + leading
+
+    if FORCE_MODE == 1:
+        num_accepted = tl.full((BLOCK_B,), BASE_TOTAL, tl.int32)
+    elif FORCE_MODE == 2:
+        counter = tl.load(counter_ptr) + 1
+        tl.store(counter_ptr, counter)
+        index = (counter * counter_stride + offs_b.to(tl.int64) * slot_stride) & pool_mask
+        draw = tl.load(pool_ptr + index, mask=mask_b, other=1.0)
+        extra = (draw < frac).to(tl.int32)
+        num_accepted = tl.minimum(BASE_TOTAL + extra, MAX_TOTAL)
+    tl.store(num_accepted_ptr + offs_b, num_accepted, mask=mask_b)
+
+    batch_idx = tl.load(batch_idx_ptr + offs_b, mask=mask_b, other=0)
+    rows = (batch_idx * (K + 1)).to(tl.int64) + num_accepted.to(tl.int64) - 1
+    tl.store(gather_ptr + offs_b, rows, mask=mask_b)
+
+
+@triton.jit
+def _gather_rows_kernel(
+    src_ptr,  # [num_src_rows, H]
+    rows_ptr,  # int64 [B]
+    out_ptr,  # [B, H] (out)
+    hidden,
+    src_row_stride,
+    BLOCK_H: tl.constexpr,
+):
+    b = tl.program_id(0)
+    cols = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = cols < hidden
+    row = tl.load(rows_ptr + b)
+    values = tl.load(src_ptr + row * src_row_stride + cols, mask=mask)
+    tl.store(out_ptr + b.to(tl.int64) * hidden + cols, values, mask=mask)
+
+
+@triton.jit
+def _finalize_kernel(
+    d0,
+    d1,
+    d2,
+    d3,
+    d4,
+    d5,
+    d6,
+    d7,  # int32 [B] draft tokens of steps 0..K-1
+    accepted_ptr,  # int32 [B, ACC_COLS]
+    num_accepted_ptr,  # int32 [B]
+    batch_idx_ptr,  # int32 [>= B]
+    next_draft_ptr,  # int32 [B, K] (out)
+    next_new_ptr,  # int32 [B, K + 1] (out)
+    B: tl.constexpr,
+    K: tl.constexpr,
+    ACC_COLS: tl.constexpr,
+    BLOCK_B: tl.constexpr,
+):
+    offs_b = tl.arange(0, BLOCK_B)
+    mask = offs_b < B
+    batch_idx = tl.load(batch_idx_ptr + offs_b, mask=mask, other=0).to(tl.int64)
+    last = tl.load(num_accepted_ptr + offs_b, mask=mask, other=1).to(tl.int64) - 1
+    bonus = tl.load(accepted_ptr + batch_idx * ACC_COLS + last, mask=mask, other=0)
+    tl.store(next_new_ptr + offs_b * (K + 1), bonus, mask=mask)
+    for k in tl.static_range(K):
+        if k == 0:
+            tok = tl.load(d0 + offs_b, mask=mask)
+        elif k == 1:
+            tok = tl.load(d1 + offs_b, mask=mask)
+        elif k == 2:
+            tok = tl.load(d2 + offs_b, mask=mask)
+        elif k == 3:
+            tok = tl.load(d3 + offs_b, mask=mask)
+        elif k == 4:
+            tok = tl.load(d4 + offs_b, mask=mask)
+        elif k == 5:
+            tok = tl.load(d5 + offs_b, mask=mask)
+        elif k == 6:
+            tok = tl.load(d6 + offs_b, mask=mask)
+        else:
+            tok = tl.load(d7 + offs_b, mask=mask)
+        tl.store(next_draft_ptr + offs_b * K + k, tok, mask=mask)
+        tl.store(next_new_ptr + offs_b * (K + 1) + 1 + k, tok, mask=mask)
+
+
+def gather_rows(src: torch.Tensor, rows: torch.Tensor) -> torch.Tensor:
+    """``src[rows]`` for a 2-D contiguous ``src`` and int64 ``rows`` (one launch)."""
+    num_rows = rows.shape[0]
+    hidden = src.shape[1]
+    out = torch.empty((num_rows, hidden), dtype=src.dtype, device=src.device)
+    block_h = 2048
+    _gather_rows_kernel[(num_rows, triton.cdiv(hidden, block_h))](
+        src, rows, out, hidden, src.stride(0), BLOCK_H=block_h, num_warps=_NUM_WARPS
+    )
+    return out
+
+
+class DraftLoopOps:
+    """draft_loop operations for one one-model MTP forward (see module docstring)."""
+
+    def __init__(self, worker, attn_metadata, spec_metadata):
+        self._worker = worker
+        self._attn_metadata = attn_metadata
+        self._spec_metadata = spec_metadata
+        self.batch_size = attn_metadata.num_seqs
+        self.draft_len = spec_metadata.runtime_draft_len
+        self.input_ids0: Optional[torch.Tensor] = None
+        self.gather_ids: Optional[torch.Tensor] = None
+        self.accepted_tokens: Optional[torch.Tensor] = None
+        self.num_accepted_tokens: Optional[torch.Tensor] = None
+        self.next_new_tokens: Optional[torch.Tensor] = None
+        self.step_engaged = True
+
+    # ------------------------------------------------------------------ setup
+    @staticmethod
+    def ineligibility(worker, attn_metadata, spec_metadata) -> Optional[str]:
+        """Return why this forward must stay native, or None (host only)."""
+        if not getattr(worker, "is_mtp_eagle", False):
+            return "not MTP Eagle one-model"
+        if getattr(worker, "_uses_external_shared_target_kv", False):
+            return "external shared-target KV draft loop"
+        if worker.guided_decoder is not None or worker.sa_enhancer is not None:
+            return "guided decoding or SA enhancer"
+        if getattr(worker, "_auxiliary_state_handlers", None):
+            return "auxiliary speculative state handlers"
+        if getattr(worker.spec_config, "use_relaxed_acceptance_for_thinking", False):
+            return "relaxed acceptance"
+        if getattr(worker, "use_dynamic_tree", False):
+            return "dynamic tree"
+        if getattr(spec_metadata, "enable_penalty", False):
+            return "occurrence penalties"
+        if worker._can_use_rejection_sampling(spec_metadata):
+            return "rejection sampling"
+        if spec_metadata.wants_advanced_draft_sampling:
+            return "advanced draft sampling"
+        if attn_metadata.num_contexts != 0:
+            return "context rows present"
+        batch_size = attn_metadata.num_seqs
+        draft_len = spec_metadata.runtime_draft_len
+        if not 0 < draft_len <= _MAX_DRAFT_LEN or not 0 < batch_size <= _MAX_BATCH:
+            return "batch or draft length outside the fused domain"
+        draft_tokens = spec_metadata.draft_tokens
+        if (
+            draft_tokens is None
+            or draft_tokens.dtype != torch.int32
+            or draft_tokens.numel() != batch_size * draft_len
+            or not draft_tokens.is_contiguous()
+        ):
+            return "draft-token buffer layout"
+        if spec_metadata.batch_indices_cuda.dtype != torch.int32:
+            return "batch-index dtype"
+        model_config = worker.model_config
+        if model_config is not None and model_config.mapping.enable_lm_head_tp_in_adp:
+            return "LM-head TP in attention DP"
+        return None
+
+    # ------------------------------------------------------------- acceptance
+    def accept(self, logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Native target sampling, then fused strict acceptance and draft-0 prep.
+
+        Returns ``(accepted_tokens [B, K+1] int32, num_accepted_tokens [B] int32)``
+        exactly as ``_sample_and_accept_draft_tokens_base`` (with the forced
+        acceptance override) returns them; also prepares the draft-0 input ids
+        and the step-0 gather rows.
+        """
+        worker = self._worker
+        spec_metadata = self._spec_metadata
+        batch_size, draft_len = self.batch_size, self.draft_len
+        if logits.dim() == 1:
+            logits = logits.unsqueeze(0)
+        device = logits.device
+        target_tokens = worker._sample_tokens_for_batch(logits, spec_metadata, 0, batch_size)
+
+        force_mode, base_total, max_total, frac = 0, 0, draft_len + 1, 0.0
+        counter, pool = spec_metadata.batch_indices_cuda, spec_metadata.batch_indices_cuda
+        force = worker.force_num_accepted_tokens
+        is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing()
+        if force != 0.0 and not is_warmup:
+            int_part = int(force)
+            frac = force - int_part
+            base_total = min(int_part + 1, max_total)
+            if frac > 0.0 and base_total < max_total:
+                worker._ensure_force_accept_rng_state(device)
+                force_mode = 2
+                counter = worker._force_accept_rng_counter
+                pool = worker._force_accept_rng_pool
+            else:
+                force_mode = 1
+
+        from tensorrt_llm._torch.speculative import interface as _iface
+
+        accepted = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=device)
+        num_accepted = torch.empty((batch_size,), dtype=torch.int32, device=device)
+        self.input_ids0 = torch.empty(
+            (batch_size * (draft_len + 1),), dtype=torch.int32, device=device
+        )
+        self.gather_ids = torch.empty((batch_size,), dtype=torch.int64, device=device)
+        _accept_kernel[(1,)](
+            target_tokens,
+            spec_metadata.draft_tokens,
+            spec_metadata.batch_indices_cuda,
+            counter,
+            pool,
+            accepted,
+            num_accepted,
+            self.input_ids0,
+            self.gather_ids,
+            float(frac),
+            _iface._FORCE_ACCEPT_RNG_COUNTER_STRIDE,
+            _iface._FORCE_ACCEPT_RNG_SLOT_STRIDE,
+            _iface._FORCE_ACCEPT_RNG_POOL_SIZE - 1,
+            B=batch_size,
+            K=draft_len,
+            FORCE_MODE=force_mode,
+            BASE_TOTAL=base_total,
+            MAX_TOTAL=max_total,
+            BLOCK_B=max(2, triton.next_power_of_2(batch_size)),
+            BLOCK_K=max(2, triton.next_power_of_2(draft_len + 1)),
+            # One warp: every lane reads the draw counter before it is stored.
+            num_warps=1,
+        )
+        self.accepted_tokens, self.num_accepted_tokens = accepted, num_accepted
+        if _verify.verify_enabled():
+            _verify.verify_accept(
+                worker,
+                logits,
+                spec_metadata,
+                batch_size,
+                draft_len,
+                accepted,
+                num_accepted,
+                self.gather_ids,
+                self.input_ids0,
+            )
+        return accepted, num_accepted
+
+    # ----------------------------------------------------------- draft loop
+    def first_drafter_inputs(self, position_ids, hidden_states) -> dict:
+        """Inputs of draft step 0 (``prepare_1st_drafter_inputs`` without context rows)."""
+        return {
+            "input_ids": self.input_ids0,
+            "position_ids": position_ids,
+            "hidden_states": hidden_states,
+            "attn_metadata": self._attn_metadata,
+            "spec_metadata": self._spec_metadata,
+        }
+
+    def step0_rows(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """``(gather_ids, hidden_states[gather_ids])`` for draft step 0."""
+        if hidden_states.dim() != 2 or not hidden_states.is_contiguous():
+            from ... import dispatch
+
+            desc = (tuple(hidden_states.shape), tuple(hidden_states.stride()))
+            dispatch.log_capability("draft_loop", f"rows:{desc}", False, f"hidden {desc}")
+            return self.gather_ids, hidden_states[self.gather_ids]
+        rows = gather_rows(hidden_states, self.gather_ids)
+        if _verify.verify_enabled():
+            _verify.verify_rows(hidden_states, self.gather_ids, rows)
+        return self.gather_ids, rows
+
+    def step_update(
+        self, attn_metadata, step: int, num_accepted_tokens, position_ids, gather_ids
+    ) -> Optional["_draft_step.DraftStepResult"]:
+        """Fused per-step metadata update, or None for the native update."""
+        if not self.step_engaged:
+            return None
+
+        def run_draft_loop():
+            return _draft_step.draft_step_update(
+                attn_metadata,
+                step,
+                self.batch_size,
+                0,
+                self.draft_len,
+                num_accepted_tokens,
+                position_ids,
+                gather_ids,
+            )
+
+        from ... import dispatch
+
+        reason = _draft_step.step_ineligibility(attn_metadata, self.batch_size, 0)
+        dispatch.log_capability(
+            "draft_loop", f"draft_step:{reason or 'engaged'}", reason is None, reason or ""
+        )
+        if _verify.verify_enabled() and reason is None:
+            result = _verify.verified_step_update(
+                run_draft_loop,
+                attn_metadata,
+                step,
+                self.batch_size,
+                self.draft_len,
+                num_accepted_tokens,
+                position_ids,
+                gather_ids,
+            )
+        else:
+            result = run_draft_loop()
+        if result is None:
+            # The metadata state is per forward: keep the remaining steps native.
+            self.step_engaged = False
+        return result
+
+    def finalize(self, draft_tokens: List[torch.Tensor]) -> Optional[torch.Tensor]:
+        """Stack the draft tokens and build ``next_new_tokens`` in one launch.
+
+        Returns ``next_draft_tokens [B, K]`` (``next_new_tokens [B, K+1]`` is kept
+        in ``self.next_new_tokens``), or None when the native code must run.
+        """
+        batch_size, draft_len = self.batch_size, self.draft_len
+        accepted_tokens, num_accepted_tokens = self.accepted_tokens, self.num_accepted_tokens
+        if accepted_tokens is None:
+            return None
+        if len(draft_tokens) != draft_len or any(
+            t.dtype != torch.int32 or t.shape != (batch_size,) or not t.is_contiguous()
+            for t in draft_tokens
+        ):
+            from ... import dispatch
+
+            desc = [(str(t.dtype), tuple(t.shape), t.is_contiguous()) for t in draft_tokens[:1]]
+            dispatch.log_capability("draft_loop", f"finalize:{desc}", False, f"draft tokens {desc}")
+            return None
+        if accepted_tokens.dtype != torch.int32 or not accepted_tokens.is_contiguous():
+            return None
+        device = accepted_tokens.device
+        next_draft = torch.empty((batch_size, draft_len), dtype=torch.int32, device=device)
+        next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=device)
+        ptrs = list(draft_tokens) + [draft_tokens[0]] * (_MAX_DRAFT_LEN - draft_len)
+        _finalize_kernel[(1,)](
+            *ptrs,
+            accepted_tokens,
+            num_accepted_tokens,
+            self._spec_metadata.batch_indices_cuda,
+            next_draft,
+            next_new,
+            B=batch_size,
+            K=draft_len,
+            ACC_COLS=accepted_tokens.shape[1],
+            BLOCK_B=max(2, triton.next_power_of_2(batch_size)),
+            num_warps=1,
+        )
+        self.next_new_tokens = next_new
+        if _verify.verify_enabled():
+            _verify.verify_finalize(
+                self._worker,
+                draft_tokens,
+                accepted_tokens,
+                num_accepted_tokens,
+                self._spec_metadata.batch_indices_cuda,
+                batch_size,
+                next_draft,
+                next_new,
+            )
+        return next_draft
+
+
+def make_draft_loop_ops(worker, attn_metadata, spec_metadata) -> Optional[DraftLoopOps]:
+    """Return the draft_loop ops for this forward, or None to keep the native path."""
+    from ... import dispatch
+
+    if _verify.verify_enabled():
+        _verify.COUNTERS.ensure_pool(
+            attn_metadata.kv_lens_cuda.device if hasattr(attn_metadata, "kv_lens_cuda") else "cuda"
+        )
+        _verify.COUNTERS.report()
+    reason = DraftLoopOps.ineligibility(worker, attn_metadata, spec_metadata)
+    dispatch.log_capability(
+        "draft_loop", f"draft_loop:{reason or 'engaged'}", reason is None, reason or ""
+    )
+    if reason is not None:
+        return None
+    return DraftLoopOps(worker, attn_metadata, spec_metadata)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_step.py
new file mode 100644
index 000000000000..3092fd822773
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/draft_step.py
@@ -0,0 +1,497 @@
+# 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.
+"""Single-launch per-step attention-metadata update for the MTP draft loop.
+
+After each draft step the native one-model loop (``Eagle3OneModelWorker.
+_forward_linear_draft_loop``) advances the next step's position ids and the
+KV lengths, then calls ``DSAtrtllmAttentionMetadata.update_for_spec_dec``,
+which rebuilds the DSA decode metadata: the 2-D KV-length view, three DeepGEMM
+paged-MQA-logits schedules (each computed into a fresh tensor and copied into
+its persistent buffer), the expanded KV lengths for the CuTe DSL logits path,
+and the fused decode metadata (request index, indexer K-cache slot mappings and
+generation indptrs). That is about twelve small launches per step.
+
+``draft_step_update`` performs the same update with one Triton launch plus the
+native host-side bookkeeping. Every device tensor it writes holds exactly the
+value the native chain writes. The schedule is the DeepGEMM SM100 metadata
+partition (``deep_gemm/scheduler/sm100_paged_mqa_logits.cuh``): request work
+``ceil(ctx / SPLIT_KV) * next_n``, SM ``i`` starts at work unit
+``i * q + min(i, rem)`` and records ``(request_q_token_start,
+work_in_request // next_n)``, with a ``(num_q_tokens, 0)`` sentinel past the
+last request. ``SPLIT_KV`` is confirmed against the native op once per process
+(``schedule_self_check``); without a match the native path runs.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+_SCHEDULE_BLOCK_KV = 64
+_SPLIT_KV_CANDIDATES = (64, 128, 256, 512)
+_MAX_BATCH = 64
+_NUM_WARPS = 4
+
+
+@triton.jit
+def _write_schedule(
+    ctx,  # int32 [BLOCK_R] context length per logical request
+    rmask,  # bool [BLOCK_R] prefix mask r < num_requests
+    num_requests,
+    sched_ptr,  # int32 [(NUM_SMS + 1) * 2] (out)
+    NEXT_N: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SPLIT_KV: tl.constexpr,
+    BLOCK_SM: tl.constexpr,
+):
+    splits = (ctx + SPLIT_KV - 1) // SPLIT_KV
+    work = tl.where(rmask, splits * NEXT_N, 0)
+    prefix = tl.cumsum(work, 0)
+    total = tl.sum(work, 0)
+    q = total // NUM_SMS
+    rem = total % NUM_SMS
+    sm = tl.arange(0, BLOCK_SM)
+    sm_mask = sm <= NUM_SMS
+    w = sm * q + tl.minimum(sm, rem)
+    # First request whose inclusive prefix exceeds w == count of prefixes <= w
+    # (prefix is non-decreasing over the live requests).
+    owned = (prefix[None, :] <= w[:, None]) & rmask[None, :]
+    request = tl.sum(owned.to(tl.int32), 1)
+    work_before = tl.max(tl.where(owned, prefix[None, :], 0), 1)
+    in_range = request < num_requests
+    q_token = tl.where(in_range, request * NEXT_N, num_requests * NEXT_N)
+    kv_split = tl.where(in_range, (w - work_before) // NEXT_N, 0)
+    tl.store(sched_ptr + 2 * sm, q_token, mask=sm_mask)
+    tl.store(sched_ptr + 2 * sm + 1, kv_split, mask=sm_mask)
+
+
+@triton.jit
+def _draft_step_kernel(
+    kv_lens_ptr,  # int32 [>= B] (in/out)
+    seq_lens_ptr,  # int32 [>= B] (out when STEP0, else in)
+    num_accepted_ptr,  # int32 [B] (STEP0)
+    pos_in_ptr,  # int32/int64 position ids of this step (HAS_POS)
+    pos_out_ptr,  # same dtype [B] (out, HAS_POS)
+    gather_ptr,  # int64 [B] rows of pos_in to advance (STEP0 and HAS_POS)
+    kv2d_ptr,  # int32 [max_seqs, cap] (out)
+    kv2d_stride,
+    sched_ptr,  # int32 [(NUM_SMS + 1) * 2] (out)
+    sched_full_ptr,  # int32 [(NUM_SMS + 1) * 2] (out, HAS_FULL)
+    kv_exp_ptr,  # int32 [>= B * EXPAND] (out, EXPAND > 1)
+    sched_exp_ptr,  # int32 [(NUM_SMS + 1) * 2] (out, EXPAND > 1)
+    block_offsets_ptr,  # int32 [max_seqs, max_blocks]
+    block_offsets_stride_0,
+    req_idx_ptr,  # int32 [>= num_tokens] (out)
+    slot_fp8_ptr,  # int64 [>= num_tokens] (out)
+    slot_scale_ptr,  # int64 [>= num_tokens] (out)
+    gen_kv_indptr_ptr,  # int64 [>= B + 1] (out)
+    gen_cached_indptr_ptr,  # int64 [>= B + 1] (out)
+    num_tokens,
+    B: tl.constexpr,
+    DRAFT_LEN: tl.constexpr,
+    STEP0: tl.constexpr,
+    HAS_POS: tl.constexpr,
+    WIDTH: tl.constexpr,
+    HAS_FULL: tl.constexpr,
+    EXPAND: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SPLIT_KV: tl.constexpr,
+    max_blocks: tl.constexpr,
+    block_offsets_stride_1: tl.constexpr,
+    tokens_per_block: tl.constexpr,
+    data_bytes_per_token: tl.constexpr,
+    scale_size: tl.constexpr,
+    block_stride: tl.constexpr,
+    scale_base_offset: tl.constexpr,
+    BLOCK_B: tl.constexpr,
+    BLOCK_W: tl.constexpr,
+    BLOCK_E: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_SM: tl.constexpr,
+):
+    offs_b = tl.arange(0, BLOCK_B)
+    mask_b = offs_b < B
+
+    # ----- KV lengths and query lengths (the loop's in-place updates) -------
+    kv_old = tl.load(kv_lens_ptr + offs_b, mask=mask_b, other=0)
+    if STEP0:
+        accepted = tl.load(num_accepted_ptr + offs_b, mask=mask_b, other=0)
+        kv = kv_old - (DRAFT_LEN - accepted)
+        seq = tl.full((BLOCK_B,), 1, tl.int32)
+        tl.store(seq_lens_ptr + offs_b, seq, mask=mask_b)
+    else:
+        kv = kv_old + 1
+        seq = tl.load(seq_lens_ptr + offs_b, mask=mask_b, other=0)
+    tl.store(kv_lens_ptr + offs_b, kv, mask=mask_b)
+    kv = tl.where(mask_b, kv, 0)
+    seq = tl.where(mask_b, seq, 0)
+
+    # ----- next step's position ids ------------------------------------------
+    if HAS_POS:
+        if STEP0:
+            rows = tl.load(gather_ptr + offs_b, mask=mask_b, other=0)
+        else:
+            rows = offs_b.to(tl.int64)
+        pos = tl.load(pos_in_ptr + rows, mask=mask_b, other=0)
+        tl.store(pos_out_ptr + offs_b, pos + 1, mask=mask_b)
+
+    # ----- 2-D KV lengths and the paged-MQA-logits schedules -----------------
+    offs_w = tl.arange(0, BLOCK_W)
+    kv2d_mask = mask_b[:, None] & (offs_w[None, :] < WIDTH)
+    tl.store(
+        kv2d_ptr + offs_b[:, None] * kv2d_stride + offs_w[None, :],
+        kv[:, None] + tl.zeros((BLOCK_B, BLOCK_W), tl.int32),
+        mask=kv2d_mask,
+    )
+    _write_schedule(kv, mask_b, B, sched_ptr, 1, NUM_SMS, SPLIT_KV, BLOCK_SM)
+    if HAS_FULL:
+        _write_schedule(kv, mask_b, B, sched_full_ptr, WIDTH, NUM_SMS, SPLIT_KV, BLOCK_SM)
+    if EXPAND > 1:
+        offs_e = tl.arange(0, BLOCK_E)
+        mask_e = offs_e < B * EXPAND
+        # repeat_interleave from registers (the updated lengths are not
+        # re-read from memory within this launch).
+        owner = (offs_e[:, None] // EXPAND) == offs_b[None, :]
+        kv_e = tl.sum(tl.where(owner, kv[None, :], 0), 1)
+        tl.store(kv_exp_ptr + offs_e, kv_e, mask=mask_e)
+        _write_schedule(kv_e, mask_e, B * EXPAND, sched_exp_ptr, 1, NUM_SMS, SPLIT_KV, BLOCK_SM)
+
+    # ----- fused DSA decode metadata (same arithmetic as the native kernel) --
+    kv64 = kv.to(tl.int64)
+    seq64 = seq.to(tl.int64)
+    tl.store(gen_kv_indptr_ptr, tl.full((), 0, tl.int64))
+    tl.store(gen_cached_indptr_ptr, tl.full((), 0, tl.int64))
+    tl.store(gen_kv_indptr_ptr + 1 + offs_b, tl.cumsum(kv64, 0), mask=mask_b)
+    tl.store(gen_cached_indptr_ptr + 1 + offs_b, tl.cumsum(kv64 - seq64, 0), mask=mask_b)
+
+    flat_start = tl.cumsum(seq64, 0) - seq64
+    start_pos = kv64 - seq64
+    offs_t = tl.arange(0, BLOCK_T).to(tl.int64)
+    g = flat_start[:, None] + offs_t[None, :]
+    mask_t = mask_b[:, None] & (offs_t[None, :] < seq64[:, None]) & (g < num_tokens)
+    tl.store(req_idx_ptr + g, offs_b[:, None] + tl.zeros((BLOCK_B, BLOCK_T), tl.int32), mask=mask_t)
+    gpos = start_pos[:, None] + offs_t[None, :]
+    blk = gpos // tokens_per_block
+    pos_in_blk = gpos % tokens_per_block
+    pos_in_blk = tl.where(pos_in_blk < 0, pos_in_blk + tokens_per_block, pos_in_blk)
+    blk = tl.minimum(tl.maximum(blk, 0), max_blocks - 1)
+    bid = tl.load(
+        block_offsets_ptr
+        + offs_b[:, None].to(tl.int64) * block_offsets_stride_0
+        + blk * block_offsets_stride_1,
+        mask=mask_t,
+        other=0,
+    ).to(tl.int64)
+    tl.store(slot_fp8_ptr + g, bid * block_stride + pos_in_blk * data_bytes_per_token, mask=mask_t)
+    tl.store(
+        slot_scale_ptr + g,
+        bid * block_stride + scale_base_offset + pos_in_blk * scale_size,
+        mask=mask_t,
+    )
+
+
+@triton.jit
+def _schedule_probe_kernel(
+    ctx_ptr,
+    sched_ptr,
+    num_requests,
+    NEXT_N: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SPLIT_KV: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    BLOCK_SM: tl.constexpr,
+):
+    offs = tl.arange(0, BLOCK_R)
+    rmask = offs < num_requests
+    ctx = tl.load(ctx_ptr + offs * NEXT_N + NEXT_N - 1, mask=rmask, other=0)
+    _write_schedule(ctx, rmask, num_requests, sched_ptr, NEXT_N, NUM_SMS, SPLIT_KV, BLOCK_SM)
+
+
+def paged_mqa_schedule(
+    context_lens: torch.Tensor, num_sms: int, split_kv: int, out: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """draft_loop reproduction of ``get_paged_mqa_logits_metadata`` (non-varlen).
+
+    Args:
+        context_lens: int32 ``[num_requests, next_n]`` contiguous CUDA tensor.
+        num_sms: SM count the schedule is balanced over.
+        split_kv: KV split length of the DeepGEMM scheduler.
+        out: optional int32 ``[num_sms + 1, 2]`` output.
+
+    Returns:
+        int32 ``[num_sms + 1, 2]`` per-SM ``(q_token_start, kv_split)`` starts.
+    """
+    num_requests, next_n = context_lens.shape
+    if out is None:
+        out = torch.empty((num_sms + 1, 2), dtype=torch.int32, device=context_lens.device)
+    _schedule_probe_kernel[(1,)](
+        context_lens,
+        out,
+        num_requests,
+        NEXT_N=next_n,
+        NUM_SMS=num_sms,
+        SPLIT_KV=split_kv,
+        BLOCK_R=max(2, triton.next_power_of_2(num_requests)),
+        BLOCK_SM=triton.next_power_of_2(num_sms + 1),
+        num_warps=_NUM_WARPS,
+    )
+    return out
+
+
+_SPLIT_KV_BY_SMS: dict = {}
+
+
+def schedule_self_check(num_sms: int, device: torch.device) -> Optional[int]:
+    """Return the DeepGEMM ``SPLIT_KV`` reproduced exactly for ``num_sms``, or None.
+
+    Compares ``paged_mqa_schedule`` with the native
+    ``get_paged_mqa_logits_metadata(lens, 64, num_sms)`` on fixed probes (1-D and
+    2-D lengths, empty, short, block-boundary and near-2^20 contexts). Runs
+    eagerly once per process and SM count; must not be called during capture.
+    """
+    if num_sms in _SPLIT_KV_BY_SMS:
+        return _SPLIT_KV_BY_SMS[num_sms]
+    if torch.cuda.is_current_stream_capturing():
+        # The probe synchronizes; without an eager result the native path runs.
+        return None
+    from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata
+
+    probes = []
+    lengths = [0, 1, 63, 64, 65, 2047, 2048, 2049, 32768, 131072, 391061, 491067, (1 << 20) + 14]
+    for length in lengths:
+        for next_n in (1, 6):
+            probes.append(torch.full((1, next_n), length, dtype=torch.int32))
+    gen = torch.Generator().manual_seed(8)
+    for num_requests in (2, 3, 8):
+        for next_n in (1, 2, 6):
+            lens = torch.randint(0, 600000, (num_requests, 1), generator=gen, dtype=torch.int32)
+            probes.append(lens.expand(-1, next_n).contiguous())
+    found = None
+    for split_kv in _SPLIT_KV_CANDIDATES:
+        ok = True
+        for lens in probes:
+            lens = lens.to(device)
+            native = get_paged_mqa_logits_metadata(lens, _SCHEDULE_BLOCK_KV, num_sms)
+            if not torch.equal(native, paged_mqa_schedule(lens, num_sms, split_kv)):
+                ok = False
+                break
+        if ok:
+            found = split_kv
+            break
+    _SPLIT_KV_BY_SMS[num_sms] = found
+    return found
+
+
+@dataclass
+class DraftStepResult:
+    """What the draft_loop step update produced for the draft loop."""
+
+    position_ids: Optional[torch.Tensor]
+
+
+def step_ineligibility(md, batch_size: int, num_contexts: int) -> Optional[str]:
+    """Return why the fused step update cannot replace the native chain, or None.
+
+    Checked on the host before any device work. The supported domain is the
+    pure-decode DSA metadata configuration of the GLM-5.2 GEN server: no
+    context rows, the fused native decode-metadata path, uncompressed indexer
+    lengths, no expanded MTP buffers, no ragged verification, no FP4/FP8-DS MLA
+    state and no dense top-k fallback.
+    """
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+        _fused_dsa_meta_enabled,
+    )
+
+    if not isinstance(md, DSAtrtllmAttentionMetadata):
+        return "metadata is not DSAtrtllmAttentionMetadata"
+    cls = type(md)
+    for name in ("update_for_spec_dec", "on_update_kv_lens", "on_update"):
+        if getattr(cls, name) is not getattr(DSAtrtllmAttentionMetadata, name):
+            return f"metadata overrides {name}"
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import is_dsa_cache_manager
+
+    # With a separate draft KV cache the loop has rebound the manager and the
+    # indexer buffers to the draft ones; the update reads them at call time.
+    if md.kv_cache_manager is None or not is_dsa_cache_manager(md.kv_cache_manager):
+        return "KV cache manager is not a DSA cache manager"
+    if num_contexts != 0 or md.num_contexts != 0:
+        return "context rows present"
+    if batch_size <= 0 or batch_size > _MAX_BATCH or batch_size != md.num_seqs:
+        return f"batch size {batch_size} outside the fused domain"
+    if not _fused_dsa_meta_enabled():
+        return "native fused DSA metadata disabled"
+    if md.use_fp8_ds_mla or getattr(md, "fp4_mla_state", None) is not None:
+        return "FP8-DS/FP4 MLA metadata state"
+    if md.enable_ragged_verification or md.is_ragged_verify:
+        return "ragged verification"
+    if md._indexer_compress_ratio > 1:
+        return "compressed indexer lengths"
+    if getattr(md, "use_expanded_buffers_for_mtp", True):
+        return "expanded MTP buffers"
+    if md.max_draft_tokens <= 0:
+        return "no draft tokens"
+    if getattr(md, "skip_indexer_for_gen_reqs", True) or getattr(
+        md, "skip_indexer_for_ctx_reqs", True
+    ):
+        return "dense top-k fallback active"
+    if md.needs_gvr_prior and batch_size * md.gen_token_stride >= 2 * md.num_sms:
+        return "GVR row reorder active"
+    if md.kv_lens_cuda.dtype != torch.int32 or md.seq_lens_cuda.dtype != torch.int32:
+        return "unexpected length dtype"
+    if md.indexer_k_cache_block_offsets.dtype != torch.int32:
+        return "unexpected block-offset dtype"
+    if schedule_self_check(md.num_sms, md.kv_lens_cuda.device) is None:
+        return "DeepGEMM schedule not reproduced"
+    return None
+
+
+def draft_step_update(
+    md,
+    step: int,
+    batch_size: int,
+    num_contexts: int,
+    runtime_draft_len: int,
+    num_accepted_tokens: Optional[torch.Tensor],
+    position_ids: Optional[torch.Tensor],
+    gather_ids: Optional[torch.Tensor],
+) -> Optional[DraftStepResult]:
+    """Fused replacement of one draft step's metadata update, or None (native).
+
+    Reproduces, for draft step ``step``, the native loop tail: the next step's
+    position ids (``position_ids[gather_ids] + 1`` at step 0,
+    ``position_ids[:batch_size] + 1`` later; skipped when ``position_ids`` is
+    None), the step-0 switch to one query token per request, the KV-length
+    update and ``update_for_spec_dec``.
+
+    Args:
+        md: the draft loop's attention metadata.
+        step: draft step index.
+        batch_size: number of requests.
+        num_contexts: context requests at the start of the loop.
+        runtime_draft_len: draft length K of this iteration.
+        num_accepted_tokens: int32 ``[batch_size]`` accepted counts (step 0).
+        position_ids: 1-D position ids of this step, or None on the last step.
+        gather_ids: int64 ``[batch_size]`` rows of ``position_ids`` (step 0).
+
+    Returns:
+        ``DraftStepResult`` when the fused update ran, None when the caller
+        must run the native update (no state has been modified then).
+    """
+    if step_ineligibility(md, batch_size, num_contexts) is not None:
+        return None
+    if step == 0:
+        if (
+            num_accepted_tokens is None
+            or num_accepted_tokens.dtype != torch.int32
+            or not num_accepted_tokens.is_contiguous()
+        ):
+            return None
+        if position_ids is not None and (
+            gather_ids is None or gather_ids.dtype != torch.int64 or not gather_ids.is_contiguous()
+        ):
+            return None
+    if position_ids is not None and (position_ids.dim() != 1 or not position_ids.is_contiguous()):
+        return None
+
+    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+
+    if step == 0:
+        # Host side of the native step-0 switch to one query token per request.
+        md._seq_lens[:batch_size].fill_(1)
+        md.on_update()
+        if hasattr(md, "host_request_types"):
+            md.host_request_types[: md.num_contexts].fill_(1)
+        md.num_contexts = 0
+
+    kv_mgr = md.kv_cache_manager
+    index_head_dim = kv_mgr.index_head_dim
+    use_fp4 = getattr(kv_mgr, "use_fp4", False)
+    data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim
+    scale_size = index_head_dim // kv_mgr.quant_block_size * 4
+    tokens_per_block = md._tokens_per_block
+    block_stride = tokens_per_block * (data_bytes_per_token + scale_size)
+    width = min(md.gen_token_stride, md.kv_lens_cuda_2d.shape[1])
+    expand = md.dsl_expand_factor if md.expand_for_dsl else 1
+    num_sms = md.num_sms
+    block_offsets = md.indexer_k_cache_block_offsets
+    has_pos = position_ids is not None
+    pos_out = torch.empty_like(position_ids[:batch_size]) if has_pos else None
+    dummy = md.kv_lens_cuda
+
+    _draft_step_kernel[(1,)](
+        md.kv_lens_cuda,
+        md.seq_lens_cuda,
+        num_accepted_tokens if step == 0 else dummy,
+        position_ids if has_pos else dummy,
+        pos_out if has_pos else dummy,
+        gather_ids if (step == 0 and has_pos) else dummy,
+        md.kv_lens_cuda_2d,
+        md.kv_lens_cuda_2d.stride(0),
+        md.scheduler_metadata_buffer,
+        md.scheduler_metadata_buffer_full_next_n,
+        md.kv_lens_expanded_cuda if expand > 1 else dummy,
+        md.scheduler_metadata_buffer_expanded if expand > 1 else dummy,
+        block_offsets,
+        block_offsets.stride(0),
+        md.req_idx_per_token,
+        md.slot_mapping_fp8,
+        md.slot_mapping_scale,
+        md.gen_kv_indptr,
+        md.gen_cached_token_indptr,
+        md.num_tokens,
+        B=batch_size,
+        DRAFT_LEN=runtime_draft_len,
+        STEP0=step == 0,
+        HAS_POS=has_pos,
+        WIDTH=width,
+        HAS_FULL=True,
+        EXPAND=expand,
+        NUM_SMS=num_sms,
+        SPLIT_KV=_SPLIT_KV_BY_SMS[num_sms],
+        max_blocks=block_offsets.shape[1],
+        block_offsets_stride_1=block_offsets.stride(1),
+        tokens_per_block=tokens_per_block,
+        data_bytes_per_token=data_bytes_per_token,
+        scale_size=scale_size,
+        block_stride=block_stride,
+        scale_base_offset=tokens_per_block * data_bytes_per_token,
+        BLOCK_B=max(2, triton.next_power_of_2(batch_size)),
+        BLOCK_W=max(2, triton.next_power_of_2(width)),
+        BLOCK_E=max(2, triton.next_power_of_2(batch_size * expand)),
+        BLOCK_T=max(2, triton.next_power_of_2(1 + md.max_draft_tokens)),
+        BLOCK_SM=triton.next_power_of_2(num_sms + 1),
+        num_warps=_NUM_WARPS,
+    )
+
+    # Host side of update_for_spec_dec -> on_update_kv_lens (no device work in
+    # the eligible domain, checked above).
+    TrtllmAttentionMetadata.update_for_spec_dec(md)
+    md.max_ctx_kv_len = 0
+    md.num_ctx_cached_tokens = 0
+    md.max_gen_seq_len = 1
+    TrtllmAttentionMetadata.on_update_kv_lens(md)
+    md._invalidate_pool_view_cache()
+    md._fused_dsa_meta_armed = True
+    md.gen_indexer_kv_lens_cuda_runtime = md.get_indexer_kv_lens(
+        md.kv_lens_cuda[md.num_contexts : md.num_seqs]
+    )
+    md._compute_kv_lens_row_reorder()
+    md.prepare_dense_topk_indices(md.kv_lens_cuda, device=True)
+    if step == 0:
+        md.use_spec_decoding = False
+    return DraftStepResult(position_ids=pos_out)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/embedding.py
new file mode 100644
index 000000000000..855eb327ce73
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/embedding.py
@@ -0,0 +1,175 @@
+# 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.
+"""Single-launch draft-token embedding for the MTP draft loop.
+
+The native ``Embedding`` of a model built without tensor parallelism runs
+``pre_comm_embedding_ops``: a seven-kernel range mask, ``F.embedding`` and a
+``masked_fill_``. This module produces the identical tensor in one launch:
+row ``r`` is ``weight[ids[r]]`` when ``0 <= ids[r] < weight.shape[0]`` and
+``+0.0`` otherwise. Values are copied bit for bit, so the result equals the
+native output exactly.
+"""
+
+import weakref
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+_BLOCK_H = 2048
+_NUM_WARPS = 4
+_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
+_SUPPORTED_ID_DTYPES = (torch.int32, torch.int64)
+
+
+@triton.jit
+def _draft_embedding_kernel(
+    ids_ptr,  # int32/int64 [num_rows]
+    weight_ptr,  # [vocab_size, hidden_size]
+    out_ptr,  # [num_rows, hidden_size]
+    vocab_size,
+    hidden_size,
+    weight_row_stride,
+    out_row_stride,
+    BLOCK_H: tl.constexpr,
+):
+    row = tl.program_id(0)
+    col_block = tl.program_id(1)
+    token = tl.load(ids_ptr + row).to(tl.int64)
+    valid = (token >= 0) & (token < vocab_size)
+    safe_token = tl.where(valid, token, 0)
+    cols = col_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    col_mask = cols < hidden_size
+    values = tl.load(
+        weight_ptr + safe_token * weight_row_stride + cols, mask=col_mask & valid, other=0.0
+    )
+    tl.store(out_ptr + row.to(tl.int64) * out_row_stride + cols, values, mask=col_mask)
+
+
+def draft_embedding(input_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """Embed ``input_ids`` with the native out-of-range-to-zero semantics.
+
+    Args:
+        input_ids: ``[num_rows]`` int32 or int64 CUDA tensor, contiguous.
+        weight: ``[vocab_size, hidden_size]`` contiguous CUDA tensor
+            (bfloat16, float16 or float32).
+
+    Returns:
+        ``[num_rows, hidden_size]`` tensor with ``weight``'s dtype.
+    """
+    num_rows = input_ids.shape[0]
+    vocab_size, hidden_size = weight.shape
+    out = torch.empty((num_rows, hidden_size), dtype=weight.dtype, device=weight.device)
+    if num_rows == 0:
+        return out
+    grid = (num_rows, triton.cdiv(hidden_size, _BLOCK_H))
+    _draft_embedding_kernel[grid](
+        input_ids,
+        weight,
+        out,
+        vocab_size,
+        hidden_size,
+        weight.stride(0),
+        out.stride(0),
+        BLOCK_H=_BLOCK_H,
+        num_warps=_NUM_WARPS,
+    )
+    return out
+
+
+def module_ineligibility(embedding: torch.nn.Module) -> Optional[str]:
+    """Return why ``embedding`` cannot use the fused kernel, or None.
+
+    The fused kernel reproduces only the communication-free native path:
+    ``tp_size == 1`` (the replicated full-vocabulary table the DeepSeek-V3 /
+    GLM models build), an unquantized 2-D contiguous CUDA weight.
+    """
+    if getattr(embedding, "tp_size", None) != 1:
+        return "embedding is tensor-parallel"
+    if getattr(embedding, "has_any_quant", False):
+        return "embedding weight is quantized"
+    weight = getattr(embedding, "weight", None)
+    if not isinstance(weight, torch.Tensor):
+        return "embedding has no weight tensor"
+    if weight.dim() != 2 or not weight.is_cuda or not weight.is_contiguous():
+        return "embedding weight is not a contiguous 2-D CUDA tensor"
+    if weight.dtype not in _SUPPORTED_DTYPES:
+        return f"unsupported embedding dtype {weight.dtype}"
+    return None
+
+
+def ids_ineligibility(input_ids: torch.Tensor, weight: torch.Tensor) -> Optional[str]:
+    """Return why ``input_ids`` cannot use the fused kernel, or None."""
+    if input_ids.dim() != 1:
+        return "draft ids are not 1-D"
+    if input_ids.dtype not in _SUPPORTED_ID_DTYPES:
+        return f"unsupported id dtype {input_ids.dtype}"
+    if input_ids.device != weight.device or not input_ids.is_contiguous():
+        return "draft ids are not contiguous on the weight's device"
+    return None
+
+
+class DraftEmbedding:
+    """Callable stand-in for a native ``Embedding`` inside the draft loop.
+
+    ``__call__`` runs the fused kernel when the module and the ids are in the
+    supported domain and the native module otherwise; the decision is made on
+    the host before any device work. Attribute reads fall through to the native
+    module, so consumers that inspect ``weight`` or other fields see the native
+    object.
+    """
+
+    def __init__(self, native: torch.nn.Module):
+        self._native = native
+        self.ineligible_reason = module_ineligibility(native)
+
+    @property
+    def native(self) -> torch.nn.Module:
+        return self._native
+
+    def __call__(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if self.ineligible_reason is None:
+            weight = self._native.weight
+            if ids_ineligibility(input_ids, weight) is None:
+                return draft_embedding(input_ids, weight)
+        return self._native(input_ids)
+
+    def __getattr__(self, name):
+        if name == "_native":
+            raise AttributeError(name)
+        return getattr(self._native, name)
+
+
+_WRAPPERS: "weakref.WeakKeyDictionary[torch.nn.Module, DraftEmbedding]" = (
+    weakref.WeakKeyDictionary()
+)
+
+
+def wrap_draft_embedding(native: torch.nn.Module) -> DraftEmbedding:
+    """Return the cached ``DraftEmbedding`` for ``native``."""
+    wrapper = _WRAPPERS.get(native)
+    if wrapper is None:
+        wrapper = DraftEmbedding(native)
+        _WRAPPERS[native] = wrapper
+        from ... import dispatch
+
+        dispatch.log_capability(
+            "draft_loop",
+            "draft_embedding",
+            wrapper.ineligible_reason is None,
+            wrapper.ineligible_reason or "",
+        )
+    return wrapper
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/verify.py
new file mode 100644
index 000000000000..3d7b0731b6b5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/verify.py
@@ -0,0 +1,245 @@
+# 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.
+"""In-pipeline native oracle for draft_loop (diagnostics, ``TRTLLM_MODELING_V2_VALIDATE=1``).
+
+Each verified draft_loop operation first runs the native code on the live state,
+records what it produced, restores the device state, runs the draft_loop operation and
+counts, on the device, every field whose value differs. The comparison ops are
+ordinary stream work, so they are captured into CUDA graphs and evaluated on
+every replay. Counters are logged by the next eager draft_loop forward and at exit.
+The draft_loop results are the ones kept.
+"""
+
+import atexit
+import functools
+from typing import Dict, List, Optional
+
+import torch
+
+from ... import dispatch
+
+_MAX_SNAPSHOT_NUMEL = 1 << 22
+
+
+@functools.lru_cache(maxsize=1)
+def verify_enabled() -> bool:
+    """Whether the installed plan runs the diagnostics; read on the first draft-loop step."""
+    return dispatch.validate_mode()
+
+
+class _Counters:
+    """Device-side mismatch counters keyed by a field name.
+
+    One eagerly allocated pool backs all counters, so a name first seen under
+    stream capture takes a free slot instead of allocating (a captured
+    allocation would be re-zeroed on every replay).
+    """
+
+    _POOL = 2048
+
+    def __init__(self):
+        self.names: List[str] = []
+        self.slots: Dict[str, int] = {}
+        self.pool: Optional[torch.Tensor] = None
+        self.last_report: Optional[dict] = None
+
+    def ensure_pool(self, device) -> None:
+        if self.pool is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError("draft_loop verify counters must be created before capture")
+            self.pool = torch.zeros((self._POOL,), dtype=torch.int64, device=device)
+
+    def get(self, name: str, device) -> torch.Tensor:
+        self.ensure_pool(device)
+        slot = self.slots.get(name)
+        if slot is None:
+            slot = len(self.names)
+            if slot >= self._POOL:
+                raise RuntimeError("draft_loop verify counter pool exhausted")
+            self.slots[name] = slot
+            self.names.append(name)
+        return self.pool[slot]
+
+    def add_mismatch(self, name: str, native: torch.Tensor, candidate: torch.Tensor) -> None:
+        counter = self.get(name, native.device)
+        if native.shape != candidate.shape or native.dtype != candidate.dtype:
+            counter.add_(1)
+        else:
+            # Bitwise comparison: uninitialized float scratch may hold NaN patterns.
+            counter.add_((_bits(native) != _bits(candidate)).any().to(torch.int64))
+
+    def add_one(self, name: str, device) -> None:
+        self.get(name, device).add_(1)
+
+    def report(self, force: bool = False) -> None:
+        if self.pool is None or torch.cuda.is_current_stream_capturing():
+            return
+        host = self.pool[: len(self.names)].tolist()
+        values = dict(zip(self.names, host))
+        if force or values != self.last_report:
+            self.last_report = values
+            calls = {
+                k: v for k, v in values.items() if k.endswith(":calls") or k.endswith(":declined")
+            }
+            bad = {k: v for k, v in values.items() if v and k not in calls}
+            print(
+                f"[glm_moe_dsa] draft_loop verify: calls {calls} mismatching fields {bad or 'none'}",
+                flush=True,
+            )
+
+
+def _bits(t: torch.Tensor) -> torch.Tensor:
+    if t.is_floating_point() and t.element_size() in (1, 2, 4, 8):
+        return t.view(
+            {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[t.element_size()]
+        )
+    return t
+
+
+COUNTERS = _Counters()
+
+
+@atexit.register
+def _report_at_exit() -> None:
+    try:
+        COUNTERS.report(force=True)
+    except Exception:  # noqa: BLE001 - best effort at interpreter exit
+        pass
+
+
+def _device_fields(md) -> List[str]:
+    names = []
+    for name, value in vars(md).items():
+        if (
+            isinstance(value, torch.Tensor)
+            and value.is_cuda
+            and 0 < value.numel() <= _MAX_SNAPSHOT_NUMEL
+        ):
+            names.append(name)
+    return sorted(names)
+
+
+def _native_step_tail(
+    md,
+    i,
+    batch_size,
+    num_contexts,
+    runtime_draft_len,
+    num_accepted_tokens,
+    position_ids,
+    gather_ids,
+):
+    """The native per-step tail of ``Eagle3OneModelWorker._forward_linear_draft_loop``."""
+    from tensorrt_llm._torch.speculative.mtp import _select_mtp_position_ids
+
+    new_position_ids = None
+    if position_ids is not None:
+        rows = gather_ids if i == 0 else slice(0, batch_size)
+        new_position_ids = _select_mtp_position_ids(position_ids, rows) + 1
+    if i == 0:
+        md._seq_lens[:batch_size].fill_(1)
+        md._seq_lens_cuda[:batch_size].fill_(1)
+        md.on_update()
+        if hasattr(md, "host_request_types"):
+            md.host_request_types[: md.num_contexts].fill_(1)
+        md.num_contexts = 0
+        md.kv_lens_cuda[num_contexts:batch_size] -= (
+            runtime_draft_len - num_accepted_tokens[num_contexts:]
+        )
+        md.kv_lens_cuda[:num_contexts] += 1
+        md.update_for_spec_dec()
+        md.use_spec_decoding = False
+    else:
+        md.kv_lens_cuda[:batch_size] += 1
+        md.update_for_spec_dec()
+    return new_position_ids
+
+
+def verified_step_update(
+    run_draft_loop, md, step, batch_size, draft_len, num_accepted_tokens, position_ids, gather_ids
+):
+    """Run the native tail, restore, run ``run_draft_loop()`` and count differences."""
+    names = _device_fields(md)
+    before = {name: getattr(md, name).clone() for name in names}
+    native_pos = _native_step_tail(
+        md, step, batch_size, 0, draft_len, num_accepted_tokens, position_ids, gather_ids
+    )
+    after_native = {name: getattr(md, name).clone() for name in names}
+    for name in names:
+        current = getattr(md, name)
+        if current.shape == before[name].shape:
+            current.copy_(before[name])
+    result = run_draft_loop()
+    if result is None:
+        # draft_loop declined after the native tail already ran: keep the native state.
+        for name in names:
+            current = getattr(md, name)
+            if current.shape == after_native[name].shape:
+                current.copy_(after_native[name])
+        COUNTERS.add_one("step:declined", md.kv_lens_cuda.device)
+        return None
+    device = md.kv_lens_cuda.device
+    COUNTERS.add_one("step:calls", device)
+    for name in names:
+        COUNTERS.add_mismatch(f"step:{name}", after_native[name], getattr(md, name))
+    if native_pos is not None:
+        COUNTERS.add_mismatch("step:position_ids", native_pos, result.position_ids)
+    return result
+
+
+def verify_accept(
+    worker,
+    logits,
+    spec_metadata,
+    batch_size,
+    draft_len,
+    accepted,
+    num_accepted,
+    gather_ids,
+    input_ids0,
+) -> None:
+    """Compare the fused acceptance with the native one (greedy batches only)."""
+    if not spec_metadata.is_all_greedy_sample or worker.force_num_accepted_tokens != 0.0:
+        return
+    device = accepted.device
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+    draft = spec_metadata.draft_tokens.reshape(batch_size, draft_len)
+    acc_n, num_n = worker._sample_and_accept_draft_tokens_base(
+        logits, draft, 0, batch_size, spec_metadata
+    )
+    COUNTERS.add_one("accept:calls", device)
+    COUNTERS.add_mismatch("accept:accepted_tokens", acc_n, accepted)
+    COUNTERS.add_mismatch("accept:num_accepted", num_n, num_accepted)
+    COUNTERS.add_mismatch("accept:input_ids0", acc_n.reshape(-1), input_ids0)
+    rows = (spec_metadata.batch_indices_cuda[:batch_size] * (draft_len + 1)).long() + num_n - 1
+    COUNTERS.add_mismatch("accept:gather_ids", rows, gather_ids)
+
+
+def verify_rows(hidden, gather_ids, rows) -> None:
+    COUNTERS.add_one("rows:calls", rows.device)
+    COUNTERS.add_mismatch("rows:hidden", hidden[gather_ids], rows)
+
+
+def verify_finalize(
+    worker, draft_tokens, accepted, num_accepted, batch_indices, batch_size, next_draft, next_new
+) -> None:
+    stacked = torch.stack(draft_tokens, dim=1)
+    expected_new = worker._prepare_next_new_tokens(
+        accepted, stacked, batch_indices, batch_size, num_accepted
+    )
+    COUNTERS.add_one("finalize:calls", next_draft.device)
+    COUNTERS.add_mismatch("finalize:next_draft", stacked, next_draft)
+    COUNTERS.add_mismatch("finalize:next_new", expected_new, next_new)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/__init__.py
new file mode 100644
index 000000000000..16f9aa490ba0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/__init__.py
@@ -0,0 +1,31 @@
+# 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.
+"""draft_tail: draft selection tail.
+
+Entry points (``dispatch.get_impl("draft_tail", name)``):
+
+* ``select_draft_tail(worker, draft_model, spec_metadata)`` -> callable or
+  None; the callable maps the draft LM-head input rows to int32 draft tokens
+  (or None for a step that must stay native).
+* ``get_runner(mapping, options=None)`` -> ``DraftTailRunner``.
+"""
+
+from ...dispatch import register
+from .draft_tail import DraftTailOptions, DraftTailRunner, get_runner, select_draft_tail
+
+register("draft_tail", "select_draft_tail", select_draft_tail)
+register("draft_tail", "get_runner", get_runner)
+
+__all__ = ["DraftTailRunner", "DraftTailOptions", "get_runner", "select_draft_tail"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail.py
new file mode 100644
index 000000000000..f3d02821af22
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail.py
@@ -0,0 +1,564 @@
+# 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.
+"""draft_tail: greedy draft-token selection over a vocab-sharded draft LM head.
+
+Replaces, for one MTP draft step, the native sequence
+
+    logits = lm_head(hidden)                     # [R, V/tp] BF16, vocab shard
+    tokens = greedy_sample_draft_with_tp_gather(logits)
+
+with ``trtllm::glm_moe_dsa_draft_tail`` (GEMV, rank-local maximum and a
+compact cross-rank exchange in one launch) or, in ``logits`` mode, the native
+GEMM followed by ``trtllm::glm_moe_dsa_draft_tail_from_logits``. Both return, on
+every rank, the token id of the largest BF16 logit over all vocab shards,
+lowest global id on ties, NaN highest: the selection rule of the native
+gather. The two modes differ in whose BF16 logits they select from:
+
+* ``logits`` mode selects from the native GEMM's own BF16 logits, so its draft
+  tokens are the native ones by construction.
+* ``fused`` mode computes the logits itself (FP32 accumulation, then BF16
+  rounding). Where the two largest logits are within one BF16 rounding step
+  of each other (a BF16 near-tie), the token depends on the exact FP32 sum,
+  i.e. on the accumulation order. The native-order fused variants (11-14,
+  ``NATIVE_ORDER_VARIANTS``; one of them is the default) use the order of the
+  native cuBLAS LM-head GEMM selected for this shape on SM103 (K walked in
+  16-value tensor-core steps, one FP32 accumulator), so their BF16 logits,
+  and draft tokens, equal native bit for bit;
+  ``test_draft_tail_draft_tail.py::test_native_order_logits_bitwise_equal_native``
+  checks this against ``F.linear`` and must pass again whenever cuBLAS or
+  the GPU changes. The other fused variants (0-10) accumulate in other
+  orders: decision-level variants that can pick a different draft token at
+  BF16 near-ties (draft acceptance can change; target verification and the
+  emitted tokens stay native).
+
+The exchange buffer is an MNNVL multicast allocation owned by this package
+(one per TP group); rank-local scratch holds the partial keys, the CTA arrival
+counter and the exchange epoch. Both are created collectively, outside CUDA
+graph capture, by ``DraftTailRunner.prepare``; all ranks agree on the outcome
+before any rank uses the candidate path.
+"""
+
+import os
+import threading
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional
+
+import torch
+
+from ...dispatch import log_capability
+
+FEATURE = "draft_tail"
+MAX_ROWS = 8
+MAX_RANKS = 8
+HIDDEN_SIZE = 6144
+NUM_SLOTS = 3
+SCRATCH_WORDS = MAX_ROWS + 2
+MAX_DIRECT_CTAS = 512
+# Last-CTA region (3 slots x 8 ranks x 8 rows), then the direct region
+# (3 slots x 8 rows x 8 ranks x 512 CTAs); see draftTail.h.
+EXCHANGE_WORDS = (
+    NUM_SLOTS * MAX_RANKS * MAX_ROWS + NUM_SLOTS * MAX_ROWS * MAX_RANKS * MAX_DIRECT_CTAS
+)
+EXCHANGE_NONE, EXCHANGE_MULTICAST, EXCHANGE_UNICAST = 0, 1, 2
+EXCHANGE_MULTICAST_DIRECT, EXCHANGE_UNICAST_DIRECT = 3, 4
+_EXCHANGE_IDS = {
+    "multicast": EXCHANGE_MULTICAST,
+    "unicast": EXCHANGE_UNICAST,
+    "multicast-direct": EXCHANGE_MULTICAST_DIRECT,
+    "unicast-direct": EXCHANGE_UNICAST_DIRECT,
+}
+NUM_VARIANTS = 15
+FIRST_MMA_VARIANT = 7
+# Fused variants that accumulate in the native LM-head GEMM's order (they
+# differ only in how the weights reach the lanes; see draftTail.cu).
+NATIVE_ORDER_VARIANTS = (11, 12, 13, 14)
+# Bounded wait for a peer's key, in GPU clock cycles (2**38 is about 2.3 min
+# at 1.9 GHz): far above any legitimate skew between ranks of one graph
+# replay, after which the kernel records the missing peer and traps.
+DEFAULT_TIMEOUT_CYCLES = 1 << 38
+
+
+@dataclass(frozen=True)
+class DraftTailOptions:
+    """Launch options of the draft tail; the family's targets run the defaults (``served``).
+
+    ``mode``: ``"fused"`` (default; draft_tail's own GEMV) or ``"logits"`` (native
+    GEMM + draft_tail). Both give native draft tokens with the default
+    ``variant`` 13 (native accumulation order, see the module docstring);
+    fused variants outside ``NATIVE_ORDER_VARIANTS`` are decision-level
+    variants at BF16 near-ties.
+
+    One fused kernel shape serves every row count 1..8: a row's packed key
+    (BF16-rounded logit and id) is then independent of the other rows of the
+    call, so CUDA-graph padding or batch composition cannot change a live
+    row's draft token. Kernel families round differently at BF16 near-ties,
+    which is why the row count never selects the family. ``variant`` 13 is
+    the tensor-core (``mma.sync``) kernel in native accumulation order with
+    32-byte weight loads, whose cost is nearly flat in R; 11, 12 and 14 load
+    the same values in narrower pieces (same results); 7-10 are
+    tensor-core kernels in a permuted K order and 0-6 the CUDA-core family
+    (cost grows with R), used only when set explicitly. ``mode="logits"`` (native GEMM + draft_tail)
+    likewise serves every row count.
+    """
+
+    mode: str = "fused"  # "fused" (own GEMV) or "logits" (native GEMM)
+    exchange: str = "multicast"  # see _EXCHANGE_IDS
+    variant: int = 13
+    ctas_per_sm: int = 3  # CUDA-core variants, when num_ctas is 0
+    num_ctas: int = 0  # overrides the default grid when non-zero
+    pdl: bool = True
+    timeout_cycles: int = DEFAULT_TIMEOUT_CYCLES
+    # Per-mille of each warp's rows loaded with an L2 evict_last hint (fused
+    # mode, L2-hinted variants); 0 keeps the whole stream evict_first.
+    keep_permille: int = 0
+
+    def shape(self, sm_count: int):
+        """(variant, num_ctas) of the fused kernel, for every row count."""
+        if self.num_ctas:
+            return self.variant, self.num_ctas
+        if self.variant >= FIRST_MMA_VARIANT:
+            return self.variant, 0
+        return self.variant, self.ctas_per_sm * sm_count
+
+    @staticmethod
+    def served() -> "DraftTailOptions":
+        """The defaults, with PDL as the native kernels have it (``TRTLLM_ENABLE_PDL``, on
+        unless set to 0)."""
+        return DraftTailOptions(pdl=os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0")
+
+    def validate(self) -> None:
+        if self.mode not in ("fused", "logits"):
+            raise ValueError(f"draft_tail: unknown mode {self.mode!r}")
+        if self.exchange not in _EXCHANGE_IDS:
+            raise ValueError(f"draft_tail: unknown exchange {self.exchange!r}")
+        if not 0 <= self.variant < NUM_VARIANTS:
+            raise ValueError(f"draft_tail: unknown variant {self.variant}")
+
+
+class DraftTailRunner:
+    """Scratch, exchange workspace and launch for one TP group on this rank."""
+
+    def __init__(
+        self, tp_rank: int, tp_size: int, mapping=None, options: Optional[DraftTailOptions] = None
+    ):
+        if not 1 <= tp_size <= MAX_RANKS:
+            raise ValueError(f"draft_tail: tp_size {tp_size} unsupported")
+        self.tp_rank = tp_rank
+        self.tp_size = tp_size
+        self.mapping = mapping
+        self.options = options or DraftTailOptions()
+        self.options.validate()
+        self.scratch: Optional[torch.Tensor] = None
+        # 16-byte-aligned rows for step inputs whose layout the kernels
+        # cannot read directly (see ``conform``).
+        self._staging: Optional[torch.Tensor] = None
+        self.staged_steps = 0
+        self._mcast = None
+        self._uc_buffer: Optional[torch.Tensor] = None
+        self._mc_ptr = 0
+        self._uc_ptr = 0
+        self._peer_ptrs: Optional[torch.Tensor] = None
+        self._state = "new"  # new -> ready | failed
+        self.calls = 0
+        self.sm_count = 0
+
+    @property
+    def ready(self) -> bool:
+        return self._state == "ready"
+
+    @property
+    def exchange_mode(self) -> int:
+        if self.tp_size == 1:
+            return EXCHANGE_NONE
+        return _EXCHANGE_IDS[self.options.exchange]
+
+    def prepare(self) -> bool:
+        """Allocate scratch and the exchange workspace once, on every rank.
+
+        Collective for tp_size > 1: every rank of the TP group must call it at
+        the same point. Never runs under CUDA graph capture; a runner that is
+        not ready there keeps the native path for that capture on all ranks.
+        """
+        if self._state != "new":
+            return self.ready
+        if torch.cuda.is_current_stream_capturing():
+            return False
+        local_error: Optional[Exception] = None
+        try:
+            self.scratch = torch.zeros(SCRATCH_WORDS, dtype=torch.int64, device="cuda")
+            self._staging = torch.zeros(MAX_ROWS, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
+            if not _kernel_layout_ok(self._staging):
+                raise RuntimeError("draft_tail: staging rows not 16-byte aligned")
+            self.sm_count = torch.cuda.get_device_properties(
+                torch.cuda.current_device()
+            ).multi_processor_count
+            if self.tp_size > 1:
+                self._allocate_exchange()
+            torch.cuda.synchronize()
+        except Exception as error:  # noqa: BLE001
+            # Converged across ranks below: a rank that raised here must not
+            # leave its peers waiting in a collective it never joins.
+            local_error = error
+        ok = self._agree(local_error is None)
+        self._state = "ready" if ok else "failed"
+        log_capability(
+            FEATURE,
+            f"draft_tail[tp{self.tp_size}]",
+            ok,
+            "" if ok else f"workspace setup failed: {local_error!r}",
+        )
+        return ok
+
+    def _allocate_exchange(self) -> None:
+        from tensorrt_llm._torch.distributed.ops import (
+            _get_mnnvl_workspace_comm,
+            _make_mnnvl_mcast_buffer,
+        )
+
+        if self.mapping is None:
+            raise ValueError("draft_tail: exchange needs the TP mapping")
+        self._comm = _get_mnnvl_workspace_comm(self.mapping)
+        use_fabric = (
+            self.mapping.is_multi_node() or os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1"
+        )
+        self._mcast = _make_mnnvl_mcast_buffer(
+            self._comm, EXCHANGE_WORDS * 8, self.mapping, use_fabric
+        )
+        self._uc_buffer = self._mcast.get_uc_buffer(self.tp_rank, [EXCHANGE_WORDS], torch.int64, 0)
+        self._uc_buffer.zero_()
+        self._uc_ptr = self._uc_buffer.data_ptr()
+        self._mc_ptr = self._mcast.get_mc_buffer([EXCHANGE_WORDS], torch.int64, 0).data_ptr()
+        peers = [
+            self._mcast.get_uc_buffer(q, [EXCHANGE_WORDS], torch.int64, 0).data_ptr()
+            for q in range(self.tp_size)
+        ]
+        self._peer_ptrs = torch.tensor(peers, dtype=torch.int64, device="cuda")
+
+    def _agree(self, local_ok: bool) -> bool:
+        if self.tp_size == 1:
+            return local_ok
+        comm = getattr(self, "_comm", None)
+        if comm is None:
+            from tensorrt_llm._torch.distributed.ops import _get_mnnvl_workspace_comm
+
+            comm = _get_mnnvl_workspace_comm(self.mapping)
+            self._comm = comm
+        from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded
+
+        # The collective also acts as the barrier that orders every rank's
+        # zeroed exchange buffer before any rank's first publish.
+        return _mnnvl_workspace_all_succeeded(comm, local_ok)
+
+    def _common_args(self, rows: int, vocab_offset: int, keys_out):
+        tokens = torch.empty(rows, dtype=torch.int32, device="cuda")
+        return tokens, (
+            self.scratch,
+            keys_out,
+            vocab_offset,
+            self.tp_rank,
+            self.tp_size,
+            self.exchange_mode,
+            self._mc_ptr,
+            self._uc_ptr,
+            self._peer_ptrs,
+        )
+
+    def run_fused(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        vocab_offset: int,
+        keys_out: Optional[torch.Tensor] = None,
+        variant: Optional[int] = None,
+        num_ctas: Optional[int] = None,
+        keep_permille: Optional[int] = None,
+        logits_out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """``logits_out`` (tests only, tensor-core variants): receives the
+        kernel's BF16 logits [R, V_local]."""
+        opts = self.options
+        tokens, common = self._common_args(hidden.shape[0], vocab_offset, keys_out)
+        scratch, keys, off, rank, size, mode, mc, uc, peers = common
+        torch.ops.trtllm.glm_moe_dsa_draft_tail(
+            hidden,
+            weight,
+            tokens,
+            scratch,
+            keys,
+            off,
+            rank,
+            size,
+            mode,
+            mc,
+            uc,
+            peers,
+            opts.variant if variant is None else variant,
+            opts.num_ctas if num_ctas is None else num_ctas,
+            opts.keep_permille if keep_permille is None else keep_permille,
+            opts.pdl,
+            opts.timeout_cycles,
+            logits_out,
+        )
+        self.calls += 1
+        return tokens
+
+    def run_from_logits(
+        self, logits: torch.Tensor, vocab_offset: int, keys_out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        tokens, common = self._common_args(logits.shape[0], vocab_offset, keys_out)
+        scratch, keys, off, rank, size, mode, mc, uc, peers = common
+        torch.ops.trtllm.glm_moe_dsa_draft_tail_from_logits(
+            logits,
+            tokens,
+            scratch,
+            keys,
+            off,
+            rank,
+            size,
+            mode,
+            mc,
+            uc,
+            peers,
+            self.options.pdl,
+            self.options.timeout_cycles,
+        )
+        self.calls += 1
+        return tokens
+
+    def select(
+        self, hidden: torch.Tensor, lm_head, keys_out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Draft tokens for ``hidden`` [R, H], 1 <= R <= 8, with the
+        vocab-sharded head. The same kernel serves every R (see
+        ``DraftTailOptions``); ``keys_out`` optionally receives the winning keys."""
+        rows = hidden.shape[0]
+        if not 1 <= rows <= MAX_ROWS:
+            raise ValueError(f"draft_tail: {rows} rows outside [1, {MAX_ROWS}]")
+        weight = lm_head.weight
+        vocab_offset = self.tp_rank * weight.shape[0]
+        if self.options.mode == "fused":
+            variant, num_ctas = self.options.shape(self.sm_count)
+            return self.run_fused(
+                hidden, weight, vocab_offset, keys_out=keys_out, variant=variant, num_ctas=num_ctas
+            )
+        logits = torch.nn.functional.linear(hidden, weight)
+        return self.run_from_logits(logits, vocab_offset, keys_out=keys_out)
+
+    def conform(self, hidden: torch.Tensor) -> torch.Tensor:
+        """``hidden`` itself when the kernels can read its layout, otherwise
+        a copy in this runner's 16-byte-aligned staging rows.
+
+        Strides and alignment come from this rank's allocator, so they must
+        never decide between draft_tail and the native path (the two use different
+        collectives); a non-conforming layout costs one device copy instead.
+        The staging rows are allocated by ``prepare`` (outside capture), so a
+        captured copy keeps a stable address across graph replays.
+        """
+        if _kernel_layout_ok(hidden):
+            return hidden
+        staged = self._staging[: hidden.shape[0]]
+        staged.copy_(hidden)
+        self.staged_steps += 1
+        return staged
+
+    def step(self, hidden: torch.Tensor, lm_head) -> Optional[torch.Tensor]:
+        """Draft tokens of one draft step, or None for the native path.
+
+        None is returned only for reasons that every TP rank shares (see
+        ``_shared_step_reason``), so all ranks run the same exchange.
+        """
+        reason = _shared_step_reason(hidden)
+        if reason is not None:
+            log_capability(
+                FEATURE, f"draft_tail[step: {reason}]", False, "identical on every TP rank"
+            )
+            return None
+        if hidden.device != lm_head.weight.device:
+            raise RuntimeError(
+                f"draft_tail: draft rows on {hidden.device}, LM head on {lm_head.weight.device}"
+            )
+        return self.select(self.conform(hidden), lm_head)
+
+
+def _kernel_layout_ok(hidden: torch.Tensor) -> bool:
+    """Unit last stride, 16-byte row stride and base: the kernels' loads."""
+    return hidden.stride(1) == 1 and hidden.stride(0) % 8 == 0 and hidden.data_ptr() % 16 == 0
+
+
+def _shared_step_reason(hidden: torch.Tensor) -> Optional[str]:
+    """Why a draft step cannot use draft_tail, judged only from properties that are
+    identical on every rank of the TP group, or None.
+
+    The draft LM-head input is ``hidden_states[gather_ids]``: the same
+    requests (row count), model dtype and hidden width on every rank. Layout
+    (strides, alignment) is rank-local and handled by ``conform`` instead.
+    """
+    if hidden.dim() != 2:
+        return f"rows of rank {hidden.dim()}"
+    if not 1 <= hidden.shape[0] <= MAX_ROWS:
+        return f"{hidden.shape[0]} rows outside [1, {MAX_ROWS}]"
+    if hidden.dtype != torch.bfloat16:
+        return f"dtype {hidden.dtype}"
+    if hidden.shape[1] != HIDDEN_SIZE:
+        return f"hidden width {hidden.shape[1]}"
+    return None
+
+
+_RUNNERS: Dict[tuple, DraftTailRunner] = {}
+_RUNNERS_LOCK = threading.Lock()
+
+
+def get_runner(mapping, options: Optional[DraftTailOptions] = None) -> DraftTailRunner:
+    """One runner per (TP group, options) in this process."""
+    options = options or DraftTailOptions.served()
+    tp_size = getattr(mapping, "tp_size", 1) if mapping is not None else 1
+    tp_rank = getattr(mapping, "tp_rank", 0) if mapping is not None else 0
+    key = (getattr(mapping, "pp_rank", 0), tp_rank, tp_size, options)
+    with _RUNNERS_LOCK:
+        runner = _RUNNERS.get(key)
+        if runner is None:
+            runner = DraftTailRunner(tp_rank, tp_size, mapping, options)
+            _RUNNERS[key] = runner
+    return runner
+
+
+def _head_is_plain_vocab_shard(lm_head, mapping) -> Optional[str]:
+    """Reason the head is not a plain BF16 vocab shard, or None."""
+    from tensorrt_llm._torch.modules.embedding import LMHead
+    from tensorrt_llm._torch.modules.linear import TensorParallelMode
+
+    if not isinstance(lm_head, LMHead):
+        return "draft head is not an LMHead"
+    if getattr(lm_head, "has_any_quant", False):
+        return "quantized LM head"
+    weight = getattr(lm_head, "weight", None)
+    if (
+        weight is None
+        or weight.dtype != torch.bfloat16
+        or weight.dim() != 2
+        or not weight.is_contiguous()
+        or weight.shape[1] != HIDDEN_SIZE
+    ):
+        return "LM head weight is not a contiguous BF16 [V, 6144] tensor"
+    tp_size = getattr(mapping, "tp_size", 1)
+    if tp_size > 1 and lm_head.tp_mode != TensorParallelMode.COLUMN:
+        return "LM head is not vocab (column) parallel"
+    if weight.shape[0] * tp_size >= 2**32:
+        return "vocabulary too large for 32-bit ids"
+    return None
+
+
+def _max_draft_rows(spec_metadata) -> Optional[int]:
+    """Most rows the draft LM head can see per step in this configuration.
+
+    The draft loop gathers one row per request, padded requests included, so
+    the engine's sequence-slot pool (``num_seq_slots``) bounds it. That pool
+    is fixed per engine and kept by the CUDA-graph metadata copies, whose
+    ``max_num_requests`` shrinks to the graph batch size; without a recorded
+    pool, ``max_num_requests`` of the eager metadata (the maximum batch size)
+    is used, and a graph copy gives no usable bound.
+    """
+    slots = getattr(spec_metadata, "num_seq_slots", 0) or 0
+    if slots > 0:
+        return slots
+    if getattr(spec_metadata, "is_cuda_graph", False):
+        return None
+    return getattr(spec_metadata, "max_num_requests", 0) or None
+
+
+def select_draft_tail(worker, draft_model, spec_metadata) -> Optional[Callable]:
+    """Decide, before any device work, whether draft_tail serves this draft loop.
+
+    Returns ``tail(request_hidden_states) -> Optional[int32 tokens]`` or None.
+    The decision depends only on configuration and on batch-level flags that
+    are identical on every TP rank, so all ranks take the same path. draft_tail
+    engages only when no step of the configuration can have more than 8 rows,
+    so the row count never moves a step between draft_tail and the native rounding.
+    ``tail`` (``DraftTailRunner.step``) returns None (native path for that
+    step) only for reasons every TP rank shares (row count, dtype, hidden
+    width); a rank-local layout (strides, alignment) is staged into aligned
+    rows instead, so no rank can leave draft_tail's exchange while others enter it.
+    """
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3MTPHead
+
+    mapping = getattr(worker, "mapping", None)
+    tp_size = getattr(mapping, "tp_size", 1) if mapping is not None else 1
+    reason = None
+    if not getattr(worker, "is_mtp_eagle", False):
+        reason = "not MTP-Eagle"
+    elif getattr(worker, "guided_decoder", None) is not None:
+        reason = "guided decoding reads draft logits"
+    elif getattr(worker, "_d2t", None) is not None:
+        reason = "draft-to-target vocab map"
+    elif spec_metadata.wants_advanced_draft_sampling:
+        reason = "advanced (rejection) draft sampling"
+    elif tp_size > MAX_RANKS:
+        reason = f"tp_size {tp_size} > {MAX_RANKS}"
+    elif (
+        mapping is not None
+        and tp_size > 1
+        and (
+            getattr(mapping, "enable_attention_dp", False)
+            or getattr(mapping, "enable_lm_head_tp_in_adp", False)
+        )
+    ):
+        reason = "attention DP"
+    elif type(draft_model.mtp_layers[0].shared_head) is not DeepseekV3MTPHead:
+        reason = "MTP shared head is not DeepseekV3MTPHead"
+    else:
+        reason = _head_is_plain_vocab_shard(draft_model.lm_head, mapping)
+    if reason is None and tp_size > 1:
+        draft_vocab = (
+            getattr(spec_metadata, "draft_vocab_size", 0)
+            or getattr(spec_metadata, "vocab_size", 0)
+            or 0
+        )
+        if not draft_vocab or draft_model.lm_head.weight.shape[0] >= draft_vocab:
+            reason = "draft logits are not vocab-sharded"
+    if reason is None:
+        max_rows = _max_draft_rows(spec_metadata)
+        if max_rows is None or max_rows > MAX_ROWS:
+            reason = f"draft steps can have {max_rows} rows (> {MAX_ROWS} or unknown)"
+    # Logged once per rank for eager and once for CUDA-graph forwards.
+    log_key = (
+        "draft_tail[graph]"
+        if getattr(spec_metadata, "is_cuda_graph", False)
+        else "draft_tail[eager]"
+    )
+    if reason is not None:
+        log_capability(FEATURE, log_key, False, reason)
+        return None
+    runner = get_runner(mapping)
+    if not runner.prepare():
+        return None
+    lm_head = draft_model.lm_head
+
+    def tail(hidden: torch.Tensor) -> Optional[torch.Tensor]:
+        return runner.step(hidden, lm_head)
+
+    decision = (
+        "native tokens"
+        if runner.options.mode == "logits" or runner.options.variant in NATIVE_ORDER_VARIANTS
+        else "near-tie variant"
+    )
+    log_capability(
+        FEATURE,
+        log_key,
+        True,
+        f"mode={runner.options.mode} ({decision}) variant="
+        f"{runner.options.variant} exchange="
+        f"{runner.options.exchange} tp={tp_size}",
+    )
+    return tail
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/__init__.py
new file mode 100644
index 000000000000..cc937c8efb5c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/__init__.py
@@ -0,0 +1,38 @@
+# 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.
+"""fused_kv_indexer_projection: one staged projection launch for ``kva`` and ``kw`` on the calls whose
+hidden state is the native normalized tensor.
+
+Requires ``absorbed_query`` (the consumer path). Registered entry point
+(``dispatch.get_impl("fused_kv_indexer_projection", name)``):
+
+* ``dsa_proj(mla, position_ids, hidden_states, attn_metadata)``: the
+  query_front packet of a full-indexer target layer, the MTP draft layer (with
+  ``draft_absorbed_query``) or a shared-indexer target layer, with the fused-A output ``kva``
+  and the indexer ``wk | weights_proj`` output ``kw`` taken from one
+  ``FusedKvProjection`` launch, or ``None`` to run the next path.
+
+``kva`` keeps the native fused-A accumulation order and is bit-identical to
+``dsv3_fused_a_gemm_op``; see ``proj_kernel.py`` for the order and
+``projection.py`` for the launch configuration.
+"""
+
+from ... import dispatch
+from .entry import fused_kv_dsa_proj
+from .projection import FusedKvProjConfig, FusedKvProjection
+
+dispatch.register("fused_kv_indexer_projection", "dsa_proj", fused_kv_dsa_proj)
+
+__all__ = ["FusedKvProjConfig", "FusedKvProjection", "fused_kv_dsa_proj"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/entry.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/entry.py
new file mode 100644
index 000000000000..6594199d7937
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/entry.py
@@ -0,0 +1,208 @@
+# 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.
+"""fused_kv_indexer_projection entry into the absorbed_query / draft_absorbed_query /
+shared_indexer_absorbed_query path (reachable only with ``fused_kv_indexer_projection``).
+
+On a call whose hidden state is the native normalized tensor (not a deferred_norm
+packet), fused_kv_indexer_projection replaces the two native GEMM bridges that consume it -- the
+fused-A ``kv_a_proj_with_mqa`` GEMM and, on full-indexer layers, the indexer
+``wk | weights_proj`` GEMM -- by one ``FusedKvProjection`` launch, and then runs
+the unchanged consumer of the selected host package on its outputs:
+
+* full-indexer target layer (absorbed_query) and the MTP draft layer (draft_absorbed_query): the
+  native ``q_a_layernorm`` and indexer ``k_norm`` bridges on the projected outputs
+  (``k_norm`` on the indexer's aux stream next to the q_a norm, joined
+  before the kernel), then absorbed_query's fused indexer kernel;
+* shared-indexer target layer: shared_indexer_absorbed_query's kernel when
+  ``shared_indexer_absorbed_query`` is selected and the layer qualifies, else absorbed_query's query
+  kernel.
+
+The layer kinds are ``full``, ``shared``, ``dense_shared`` and ``draft``; the
+served ones are ``DEFAULT_LAYERS`` (``served_layer_kinds``).
+``dense_shared`` is the shared-indexer target layer whose input comes from the
+post-MLP AllReduce of a dense layer (GLM-5.2: layer 3, after the dense layers
+0-2); ``shared`` selects every shared-indexer layer. Every decline happens
+before any GPU work and returns ``None``; the caller then runs the next path of
+the adapter (absorbed_query / shared_indexer_absorbed_query / draft_absorbed_query, or native).
+"""
+
+from typing import Optional
+
+from ...dispatch import is_enabled, log_capability
+from ..absorbed_query.query_proj import (
+    _KV_LORA,
+    AbsorbedQueryPacket,
+    AbsorbedQueryProjection,
+    absorb_for,
+    call_eligibility,
+    check_package_selectors,
+    module_eligibility,
+)
+
+#: How the decline reasons name the layer selection (``<kind> layers not selected (<label>)``).
+LAYERS_LABEL = "TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS"
+LAYER_KINDS = ("full", "shared", "dense_shared", "draft")
+#: The layer kinds this target serves.
+DEFAULT_LAYERS = ("dense_shared",)
+# GLM-5.2 runs dense MLPs in layers 0-2 (first_k_dense_replace = 3), so the
+# shared-indexer layers up to layer 3 take their input from a dense layer's
+# post-MLP AllReduce.
+DENSE_FED_LAST_LAYER = 3
+
+
+def served_layer_kinds() -> tuple:
+    """The served layer kinds, ``DEFAULT_LAYERS``."""
+    return DEFAULT_LAYERS
+
+
+def layer_kind(mla) -> str:
+    num_layers = getattr(mla, "num_hidden_layers", None)
+    if num_layers is not None and mla.layer_idx is not None and mla.layer_idx >= num_layers:
+        return "draft"
+    if mla.mqa.indexer is not None:
+        return "full"
+    if mla.layer_idx is not None and mla.layer_idx <= DENSE_FED_LAST_LAYER:
+        return "dense_shared"
+    return "shared"
+
+
+def kind_selected(kind: str, served: tuple) -> bool:
+    """``shared`` also selects the ``dense_shared`` layer."""
+    return kind in served or (kind == "dense_shared" and "shared" in served)
+
+
+def _host_runner(mla, kind: str):
+    """The absorbed_query / draft_absorbed_query / shared_indexer_absorbed_query runner that
+    consumes the projection on this module (the same cached object the package's own entry point
+    would create)."""
+    if kind == "draft":
+        if not is_enabled("draft_absorbed_query"):
+            return False
+        runner = getattr(mla, "_draft_absorbed_query", None)
+        if runner is None:
+            ok, reason = module_eligibility(mla, draft=True)
+            log_capability("draft_absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+            runner = AbsorbedQueryProjection(mla) if ok else False
+            mla._draft_absorbed_query = runner
+        return runner
+    if kind in ("shared", "dense_shared") and is_enabled("shared_indexer_absorbed_query"):
+        runner = getattr(mla, "_shared_indexer_query", None)
+        if runner is None:
+            from ..absorbed_query.shared_indexer import (
+                SharedIndexerProjection,
+                shared_module_eligibility,
+            )
+
+            ok, reason = shared_module_eligibility(mla)
+            log_capability("shared_indexer_absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+            runner = SharedIndexerProjection(mla) if ok else False
+            mla._shared_indexer_query = runner
+        if runner is not False:
+            return runner
+    runner = getattr(mla, "_absorbed_query", None)
+    if runner is None:
+        ok, reason = module_eligibility(mla)
+        log_capability("absorbed_query", f"layer{mla.layer_idx}", ok, reason)
+        runner = AbsorbedQueryProjection(mla) if ok else False
+        mla._absorbed_query = runner
+    return runner
+
+
+def _fused_kv_runner(mla):
+    """Per-module ``(projection, host runner, kind)`` (cached in
+    ``mla._fused_kv_projection``) or False.
+
+    Every decline is a static check of replicated module state and the
+    environment, identical on all TP ranks. The kernel is compiled here (no
+    GPU work); a compile error is not masked by a rank-local fallback and
+    propagates.
+    """
+    cached = getattr(mla, "_fused_kv_projection", None)
+    if cached is not None:
+        return cached
+    from .projection import FusedKvProjection, projection_eligibility
+
+    check_package_selectors()
+    kind = layer_kind(mla)
+    ok, reason = True, ""
+    if not kind_selected(kind, served_layer_kinds()):
+        ok, reason = False, f"{kind} layers not selected ({LAYERS_LABEL})"
+    kw = None
+    if ok:
+        ok, reason = module_eligibility(mla, draft=kind == "draft")
+    if ok:
+        ok, reason, kw = projection_eligibility(mla)
+    host = False
+    if ok:
+        host = _host_runner(mla, kind)
+        if host is False:
+            ok, reason = (
+                False,
+                "no eligible absorbed_query / draft_absorbed_query / shared_indexer_absorbed_query consumer",
+            )
+    shared = kind in ("shared", "dense_shared")
+    if ok and not shared and host.config.prefetch:
+        ok, reason = False, "absorbed_query prefetch kernel configured"
+    if ok and not shared and not host.indexer_fusable():
+        ok, reason = False, "absorbed_query fused indexer kernel not eligible"
+    runner = False
+    if ok:
+        proj = FusedKvProjection(mla, kw_weight=kw)
+        why = proj.supported()
+        if why is not None:
+            ok, reason = False, why
+        else:
+            proj.precompile()
+            runner = (proj, host, kind)
+    log_capability("fused_kv_indexer_projection", f"layer{mla.layer_idx}", ok, reason or kind)
+    mla._fused_kv_projection = runner
+    return runner
+
+
+def fused_kv_dsa_proj(
+    mla, position_ids, hidden_states, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """The query_front packet of a call whose hidden state is the native
+    normalized tensor, with ``kva`` / ``kw`` from the fused_kv_indexer_projection.
+
+    Same outputs and side effects as absorbed_query's ``dsa_proj``, draft_absorbed_query's
+    ``mtp_dsa_proj`` or shared_indexer_absorbed_query's ``shared_dsa_proj`` on the same call;
+    ``None`` (before any GPU work) outside the support domain.
+    """
+    if not (is_enabled("fused_kv_indexer_projection") and is_enabled("absorbed_query")):
+        return None
+    runner = _fused_kv_runner(mla)
+    if runner is False:
+        return None
+    proj, host, kind = runner
+    ok, _ = call_eligibility(mla, hidden_states, attn_metadata)
+    if not ok or hidden_states.shape[0] > _max_rows() or not hidden_states.is_contiguous():
+        return None
+    kva, kw = proj.run(hidden_states)
+    if kind in ("shared", "dense_shared"):
+        if isinstance(host, AbsorbedQueryProjection):
+            q, _, latent, fused_q = host.launch(kva, need_qr=False, absorb=absorb_for(mla))
+        else:
+            q, latent, fused_q = host.launch(kva)
+        return AbsorbedQueryPacket([q, latent[:, :_KV_LORA], latent[:, _KV_LORA:], latent], fused_q)
+    from ..deferred_norm.query_entry import indexer_query_from_projected
+
+    return indexer_query_from_projected(host, kva, kw, position_ids, absorb_for(mla))
+
+
+def _max_rows() -> int:
+    from .projection import MAX_ROWS
+
+    return MAX_ROWS
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py
new file mode 100644
index 000000000000..9fe9cf2ea82c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py
@@ -0,0 +1,906 @@
+# 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.
+"""CuTe DSL kernel of the fused_kv_indexer_projection merged projection with staged weights.
+
+Computes, for ``R <= 8`` token rows of a row-major BF16 activation
+``a [R, K]``,
+
+    y[t, j] = bf16(sum_k a[t, k] * W[j, k])
+
+over the concatenation of two row-major BF16 weight tensors with the same
+``K`` (``W0 [N0, K]`` = the fused ``[q_a | kv_a | k_pe]`` projection,
+``W1 [N1, K]`` = the indexer ``wk | weights_proj`` rows of a full-indexer
+layer, or none), writing ``out0 [R, N0]`` and ``out1 [R, N1]``. Both weight
+tensors are read in place.
+
+Accumulation order (the order of the native ``fused_a_gemm_kernel`` for
+``nsplit = 4``): ``K`` is cut into ``nsplit`` contiguous slices; each slice
+is one sequential chain of ``mma.sync.m16n8k16`` steps (weights as the A
+operand, the activation as the B operand, K ascending, accumulator starting
+at +0); the per-row slice partials are summed in slice order and the sum
+``0 + s`` is rounded once to BF16.
+
+Schedule (one CTA per SM, ``warps`` MMA warps):
+
+* CTA ``c`` owns a contiguous range of output rows. Before
+  ``griddepcontrol.wait`` (the weights are immutable) warp 0 copies up to
+  ``smem_rows`` of them into shared memory with ``cp.async.bulk`` (row
+  stride padded by 16 bytes) in sub-chunks ordered so that every warp's
+  slice receives bytes at the same rate (``(w0, s0), (w1, s0), ..., (w0,
+  s1), ...``), one single-phase mbarrier per sub-chunk. A CTA that owns one
+  row more than it stages prefetches that row into L2 and reads it from
+  global memory in the main loop (the "direct" row).
+* After the wait, warp ``w`` loads its activation fragments from global
+  memory through a rolling window of ``b_dist`` K steps and runs the chains
+  of its K range for both m16 tiles of the CTA (rows 0-15 by
+  ``ldmatrix.x4``, rows 16-23 by ``ldmatrix.x2``; rows the CTA does not
+  stage read a zero chunk, and the lanes that hold the direct row take its
+  fragment from global memory instead). Each finished slice partial is
+  stored to a shared-memory reduction buffer.
+* After a CTA barrier, thread ``(row, token)`` sums the slice partials in
+  slice order and stores the BF16 result.
+
+No global state is kept between launches (no counters, flags or scratch).
+"""
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Int64
+from cutlass.utils.smem_allocator import SmemAllocator
+
+from .ptx import (
+    add_f32,
+    bf16_bits,
+    bulk_g2s,
+    cp_async16,
+    cp_async_mbar_arrive,
+    fence_tensormap_acquire,
+    globaltimer,
+    ld_ca_b32,
+    ld_ca_b32_if,
+    ld_cg_b32,
+    ld_cg_b32_if,
+    ld_nc_b32,
+    ld_shared_f32,
+    ldsm_x2,
+    ldsm_x4,
+    mbar_arrive_expect_tx,
+    mbar_fence_init,
+    mbar_init,
+    mbar_wait,
+    min_i32,
+    mma_bf16,
+    prefetch_l2,
+    saddr,
+    sel_i32,
+    sel_i64,
+    st_global_b16,
+    st_global_b64,
+    st_shared_f32,
+    st_shared_v4,
+    sync_warp,
+    tensor_g2s_3d,
+    xor_i32,
+)
+
+MAX_TOKENS = 8
+SMEM_LIMIT = 227 * 1024
+# Rows held by one CTA's two m16 tiles: rows 0-15 (tile 0) and 16-23 (tile 1).
+TILE_ROWS = 24
+_MISC_BYTES = 1024
+# Tensor (TMA) staging: 16 rows x 64 K elements (128 B, 128-byte swizzle) per
+# K group; the rows of m16 tile 0 come from shared memory, rows 16.. of the
+# CTA are read from global memory.
+TMA_ROWS = 16
+TMA_GROUP_K = 64
+TMA_GROUP_BYTES = TMA_ROWS * TMA_GROUP_K * 2
+
+
+def split_ctas(n0: int, n1: int, ctas: int):
+    """CTA split of the tensor-staged schedule: ``(c0, c1)`` CTAs for the
+    ``n0`` rows of ``W0`` and the ``n1`` rows of ``W1`` (no CTA straddles the
+    two tensors), minimizing the largest row count; ``None`` if no split
+    gives every CTA ``TMA_ROWS .. TILE_ROWS`` rows."""
+    if n1 == 0:
+        parts = [(0, ctas)]
+    else:
+        parts = [(c, ctas - c) for c in range(1, ctas)]
+    best = None
+    for c1, c0 in parts:
+        groups = [(n0, c0)] + ([(n1, c1)] if n1 else [])
+        lo = min(nn // cc for nn, cc in groups)
+        hi = max(-(-nn // cc) for nn, cc in groups)
+        if lo >= TMA_ROWS and hi <= TILE_ROWS and (best is None or hi < best[0]):
+            best = (hi, c0, c1)
+    return None if best is None else (best[1], best[2])
+
+
+def _ordered_sum(vals, order: str, ftz: bool):
+    """Trace-time helper: fp32 sum of ``vals`` in slice order (``seq``:
+    ``((v0 + v1) + v2) + ...``; ``tree``: pairwise by index)."""
+    vals = list(vals)
+    if order == "seq":
+        s = vals[0]
+        for v in vals[1:]:
+            s = add_f32(s, v, ftz)
+        return s
+    while len(vals) > 1:
+        nxt = [add_f32(vals[2 * q], vals[2 * q + 1], ftz) for q in range(len(vals) // 2)]
+        if len(vals) % 2:
+            nxt.append(vals[-1])
+        vals = nxt
+    return vals[0]
+
+
+@cute.jit
+def _row_src(
+    w0: Int64,
+    w1: Int64,
+    gr: Int32,
+    K: cutlass.Constexpr[int],
+    N0: cutlass.Constexpr[int],
+    N1: cutlass.Constexpr[int],
+) -> Int64:
+    """Global address of row ``gr`` of the concatenation ``[W0; W1]``."""
+    k2 = Int64(K * 2)
+    src = w0 + Int64(gr) * k2
+    if cutlass.const_expr(N1 > 0):
+        src = sel_i64(Int32(gr >= Int32(N0)), w1 + Int64(gr - Int32(N0)) * k2, src)
+    return src
+
+
+@cute.jit
+def _store_out(
+    out0: Int64,
+    out1: Int64,
+    gr: Int32,
+    tok: Int32,
+    value: Float32,
+    N0: cutlass.Constexpr[int],
+    N1: cutlass.Constexpr[int],
+):
+    """BF16 store of ``value`` into ``out0[tok, gr]`` or ``out1[tok, gr - N0]``."""
+    if gr < Int32(N0):
+        st_global_b16(out0 + Int64(tok * Int32(N0) + gr) * Int64(2), bf16_bits(value))
+    else:
+        if cutlass.const_expr(N1 > 0):
+            st_global_b16(
+                out1 + Int64(tok * Int32(N1) + gr - Int32(N0)) * Int64(2), bf16_bits(value)
+            )
+
+
+@cute.jit
+def _stamp(
+    enabled: cutlass.Constexpr[bool],
+    ts: Int64,
+    bid: Int32,
+    tid: Int32,
+    k: cutlass.Constexpr[int],
+    lead: cutlass.Constexpr[int] = 0,
+):
+    """Tests only: ``%globaltimer`` of CTA ``bid`` at point ``k`` (thread ``lead``)."""
+    if cutlass.const_expr(enabled):
+        if tid == Int32(lead):
+            st_global_b64(ts + Int64((bid * Int32(8) + Int32(k)) * Int32(8)), globaltimer())
+
+
+class FusedKvProjKernel:
+    """Compile-time configuration of the merged projection (module docstring)."""
+
+    def __init__(
+        self,
+        k: int,
+        n0: int,
+        n1: int,
+        ctas: int,
+        warps: int = 4,
+        smem_rows: int = 18,
+        nsplit: int = 4,
+        nsub: int = 4,
+        use_pdl: bool = True,
+        trigger: str = "start",
+        stage_before_wait: bool = True,
+        l2_prefetch: bool = True,
+        b_dist: int = 16,
+        a_dist: int = 2,
+        reduce: str = "seq",
+        ftz: bool = False,
+        producer_warp: bool = True,
+        copy_method: str = "bulk",
+        producer_warps: int = 1,
+        copy_only: bool = False,
+        debug_ts: bool = False,
+        chunks: tuple = (),
+        b_load: str = "cg",
+        diag_no_act: bool = False,
+        tma_groups: int = 1,
+        prewait_chunks: int = -1,
+    ):
+        assert k % 16 == 0
+        self.k = k
+        self.n0 = n0
+        self.n1 = n1
+        self.ctas = ctas
+        self.warps = warps
+        # A dedicated producer warp issues the bulk copies (the issue blocks
+        # while the copy engine is busy); otherwise MMA warp 0 does.
+        # ``copy_method``: "bulk" (1-D cp.async.bulk, one lane per row of the
+        # copy warp), "cpasync" (16-byte cp.async spread over
+        # ``producer_warps`` dedicated warps, one mbarrier arrival per thread
+        # and sub-chunk) or "tma" (3-D tensor copies of 16 rows x
+        # ``tma_groups`` K groups each, 128-byte swizzle; see ``tma``).
+        assert copy_method in ("bulk", "cpasync", "tma")
+        self.copy_method = copy_method
+        if copy_method == "cpasync":
+            assert producer_warp and producer_warps >= 1
+        else:
+            producer_warps = 1 if producer_warp else 0
+        self.producer_warp = producer_warp
+        self.producer_warps = producer_warps
+        self.copy_warp = warps if producer_warp else 0
+        self.nt = 32 * (warps + producer_warps)
+        # Tests only: stage the weights and exit (streaming floor).
+        self.copy_only = copy_only
+        steps = k // 16
+        assert nsplit % warps == 0 and steps % nsplit == 0, (nsplit, warps)
+        self.nsplit = nsplit
+        self.split_steps = steps // nsplit
+        self.warp_steps = steps // warps
+        self.warp_k = k // warps
+        # Staging chunks of each warp slice, in K steps (``chunks`` overrides
+        # ``nsub`` equal chunks); one mbarrier per (chunk, warp).
+        if chunks:
+            chunks = tuple(int(c) for c in chunks)
+        else:
+            assert self.warp_steps % nsub == 0, (self.warp_steps, nsub)
+            chunks = (self.warp_steps // nsub,) * nsub
+        assert all(c > 0 for c in chunks) and sum(chunks) == self.warp_steps, chunks
+        self.tma = copy_method == "tma"
+        self.tma_groups = tma_groups
+        if self.tma:
+            assert self.warp_k % TMA_GROUP_K == 0
+            assert tma_groups in (1, 3), "odd group counts keep ldmatrix conflict-free"
+            assert all(c % (4 * tma_groups) == 0 for c in chunks), (chunks, tma_groups)
+        self.chunks = chunks
+        self.nsub = len(chunks)
+        self.chunk_start = [sum(chunks[:i]) for i in range(len(chunks))]
+        self.chunk_at = {st: i for i, st in enumerate(self.chunk_start)}
+        self.nmbar = warps * self.nsub
+        # Activation fragment loads: "cg" (L2 only), "ca" (L1 and L2), "nc"
+        # (read-only path; benchmarks only). ``diag_no_act``: benchmarks only,
+        # no activation loads at all (the output is meaningless).
+        assert b_load in ("cg", "ca", "nc")
+        self.b_load = b_load
+        self.diag_no_act = diag_no_act
+        assert 64 + 8 * self.nmbar <= _MISC_BYTES
+        n = n0 + n1
+        self.stride = k * 2 + 16
+        if self.tma:
+            # ``tma``: CTAs [0, c0) own W0 rows, [c0, ctas) own W1 rows; each
+            # CTA stages its first 16 rows (m16 tile 0) by tensor copies into
+            # one 2 KiB block per K group ([16 rows][128 B], 128-byte
+            # swizzle), and the lanes of tile 1 read rows 16.. directly.
+            split = split_ctas(n0, n1, ctas)
+            assert split is not None, (n0, n1, ctas)
+            self.c0, self.c1 = split
+            groups = [(n0, self.c0)] + ([(n1, self.c1)] if n1 else [])
+            self.max_rows = max(-(-nn // cc) for nn, cc in groups)
+            self.smem_rows = TMA_ROWS
+            self.direct = self.max_rows - TMA_ROWS
+            self.wgroups = self.warp_k // TMA_GROUP_K
+            self.w_bytes = warps * self.wgroups * TMA_GROUP_BYTES + 1024
+        else:
+            self.base_rows = n // ctas
+            self.rem_rows = n % ctas
+            self.max_rows = self.base_rows + (1 if self.rem_rows else 0)
+            assert 1 <= self.base_rows and self.max_rows <= TILE_ROWS
+            self.smem_rows = min(smem_rows, self.max_rows)
+            self.direct = self.max_rows - self.smem_rows
+            assert self.direct in (0, 1), "at most one directly-read row per CTA"
+            if self.direct:
+                # The direct row sits in m16 tile 1 (rows 16-23).
+                assert 16 <= self.smem_rows < TILE_ROWS
+            assert self.smem_rows <= 32, "the copy warp issues one row per lane"
+            self.w_bytes = self.smem_rows * self.stride
+        self.use_pdl = use_pdl
+        assert trigger in ("start", "end", "none")
+        self.trigger = trigger
+        self.stage_before_wait = stage_before_wait
+        # Chunks staged before the dependency wait (the rest after it):
+        # all (-1) or the first ``prewait_chunks``; none without
+        # ``stage_before_wait``.
+        if not stage_before_wait:
+            self.pre_chunks = 0
+        elif prewait_chunks < 0:
+            self.pre_chunks = self.nsub
+        else:
+            self.pre_chunks = min(prewait_chunks, self.nsub)
+        self.l2_prefetch = l2_prefetch
+        self.b_dist = max(1, b_dist)
+        self.a_dist = max(1, a_dist)
+        assert reduce in ("seq", "tree")
+        self.reduce = reduce
+        self.ftz = ftz
+        self.debug_ts = debug_ts
+        self.red_bytes = nsplit * TILE_ROWS * MAX_TOKENS * 4
+        self.smem_bytes = self.w_bytes + _MISC_BYTES + self.red_bytes
+        assert self.smem_bytes <= SMEM_LIMIT, f"{self.smem_bytes} B of shared memory"
+
+    @cute.jit
+    def __call__(
+        self,
+        act: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        tmaps: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        stream: CUstream,
+    ):
+        self.kernel(act, w0, w1, out0, out1, ts, tmaps, anchor, rows).launch(
+            grid=[self.ctas, 1, 1], block=[self.nt, 1, 1], stream=stream, use_pdl=self.use_pdl
+        )
+
+    @cute.jit
+    def _issue_copies(
+        self,
+        sW: Int32,
+        sMbar: Int32,
+        w0: Int64,
+        w1: Int64,
+        row0: Int32,
+        nsm: Int32,
+        warp: Int32,
+        lane: Int32,
+        tid: Int32,
+        tmap: Int64,
+        lrow0: Int32,
+        nrows: Int32,
+        part: cutlass.Constexpr[int],
+    ):
+        """Stage the CTA's rows sub-chunk by sub-chunk in arrival order.
+
+        ``bulk``: the copy warp's lane 0 posts every sub-chunk's byte count,
+        then lane ``r < nsm`` copies row ``r``. ``cpasync``: producer thread
+        ``p`` copies 16-byte pieces ``p, p + P, ...`` of each sub-chunk and
+        then arrives on the sub-chunk's mbarrier when its copies have landed.
+        ``tma``: the copy warp's lane 0 posts every chunk's byte count and
+        issues the chunk's tensor copies (box ``(64, tma_groups, 16)`` at K
+        group ``gq``, local row ``lrow0``) into the group blocks; lanes 1..
+        prefetch the directly-read rows into L2.
+        ``part`` 0 issues the chunks staged before the dependency wait, 1 the
+        others.
+        """
+        s_lo = self.pre_chunks * part
+        s_hi = self.pre_chunks + (self.nsub - self.pre_chunks) * part
+        first = (1 - part) + int(self.pre_chunks == 0) > 0
+        if cutlass.const_expr(self.tma):
+            if warp == Int32(self.copy_warp):
+                if lane == Int32(0):
+                    if cutlass.const_expr(first):
+                        fence_tensormap_acquire(tmap)
+                    for ci in cutlass.range_constexpr((s_hi - s_lo) * self.warps):
+                        c = s_lo * self.warps + ci
+                        mbar_arrive_expect_tx(
+                            sMbar + Int32(8 * c),
+                            Int32(self.chunks[c // self.warps] // 4 * TMA_GROUP_BYTES),
+                        )
+                    for si in cutlass.range_constexpr(s_hi - s_lo):
+                        s = s_lo + si
+                        for w in cutlass.range_constexpr(self.warps):
+                            g0 = w * self.wgroups + self.chunk_start[s] // 4
+                            for o in cutlass.range_constexpr(
+                                self.chunks[s] // (4 * self.tma_groups)
+                            ):
+                                gq = g0 + o * self.tma_groups
+                                tensor_g2s_3d(
+                                    sW + Int32(gq * TMA_GROUP_BYTES),
+                                    tmap,
+                                    Int32(0),
+                                    Int32(gq),
+                                    lrow0,
+                                    sMbar + Int32(8 * (s * self.warps + w)),
+                                )
+                if cutlass.const_expr(self.direct > 0 and self.l2_prefetch and first):
+                    if lane >= Int32(1):
+                        if lane <= nrows - Int32(TMA_ROWS):
+                            prefetch_l2(
+                                _row_src(
+                                    w0,
+                                    w1,
+                                    row0 + Int32(TMA_ROWS - 1) + lane,
+                                    self.k,
+                                    self.n0,
+                                    self.n1,
+                                ),
+                                Int32(self.k * 2),
+                            )
+        elif cutlass.const_expr(self.copy_method == "bulk"):
+            if warp == Int32(self.copy_warp):
+                if lane == Int32(0):
+                    for ci in cutlass.range_constexpr((s_hi - s_lo) * self.warps):
+                        c = s_lo * self.warps + ci
+                        mbar_arrive_expect_tx(
+                            sMbar + Int32(8 * c), nsm * Int32(32 * self.chunks[c // self.warps])
+                        )
+                sync_warp()
+                if lane < nsm:
+                    src = _row_src(w0, w1, row0 + lane, self.k, self.n0, self.n1)
+                    dst = sW + lane * Int32(self.stride)
+                    for si in cutlass.range_constexpr(s_hi - s_lo):
+                        s = s_lo + si
+                        for w in cutlass.range_constexpr(self.warps):
+                            off = (w * self.warp_k + self.chunk_start[s] * 16) * 2
+                            bulk_g2s(
+                                dst + Int32(off),
+                                src + Int64(off),
+                                Int32(32 * self.chunks[s]),
+                                sMbar + Int32(8 * (s * self.warps + w)),
+                            )
+        else:
+            nprod = 32 * self.producer_warps
+            if warp >= Int32(self.copy_warp):
+                p = tid - Int32(32 * self.copy_warp)
+                for si in cutlass.range_constexpr(s_hi - s_lo):
+                    s = s_lo + si
+                    per_row = self.chunks[s] * 2
+                    total = nsm * Int32(per_row)
+                    iters = (self.smem_rows * per_row + nprod - 1) // nprod
+                    for w in cutlass.range_constexpr(self.warps):
+                        base_off = (w * self.warp_k + self.chunk_start[s] * 16) * 2
+                        for it in cutlass.range_constexpr(iters):
+                            piece = p + Int32(nprod * it)
+                            if piece < total:
+                                r = piece // Int32(per_row)
+                                off = Int32(base_off) + (piece % Int32(per_row)) * Int32(16)
+                                cp_async16(
+                                    sW + r * Int32(self.stride) + off,
+                                    _row_src(w0, w1, row0 + r, self.k, self.n0, self.n1)
+                                    + Int64(off),
+                                )
+                        cp_async_mbar_arrive(sMbar + Int32(8 * (s * self.warps + w)))
+
+    @cute.jit
+    def _mma(
+        self,
+        act: Int64,
+        w0: Int64,
+        w1: Int64,
+        ts: Int64,
+        sW: Int32,
+        sMbar: Int32,
+        sZero: Int32,
+        sRed: Int32,
+        bid: Int32,
+        tid: Int32,
+        warp: Int32,
+        lane: Int32,
+        rows: Int32,
+        row0: Int32,
+        nrows: Int32,
+        nsm: Int32,
+    ):
+        """MMA warp ``warp``: the chains of its K range for both m16 tiles,
+        each finished slice partial stored to ``sRed``."""
+        k = self.k
+        stride = self.stride
+        # Fragment coordinates: token / row group g = lane / 4, pair t = lane % 4.
+        g = lane // Int32(4)
+        t4 = lane % Int32(4)
+        tok_ok = Int32(g < rows)
+        safe_g = min_i32(g, rows - Int32(1))
+        kw0 = warp * Int32(self.warp_k)
+        # B fragment of K step j: a[g, k0 + 2 t4 .. +1] and a[g, k0 + 8 + 2 t4 .. +1].
+        b_lane = act + Int64((safe_g * Int32(k) + kw0 + t4 * Int32(2)) * Int32(2))
+        # ldmatrix row addresses (per lane), advanced by 32 B per K step.
+        r16 = lane % Int32(16)
+        a_lane0 = sW + r16 * Int32(stride) + (lane // Int32(16)) * Int32(16) + kw0 * Int32(2)
+        ok0 = Int32(r16 < nsm)
+        r8 = Int32(16) + lane % Int32(8)
+        a_lane1 = (
+            sW + r8 * Int32(stride) + ((lane // Int32(8)) % Int32(2)) * Int32(16) + kw0 * Int32(2)
+        )
+        ok1 = Int32(r8 < nsm)
+        hl = lane // Int32(16)
+        t_off = []
+        t_x = []
+        t_base = sW
+        if cutlass.const_expr(self.tma):
+            # Tile 0 from the swizzled group blocks: in a box of G groups the
+            # 128-byte line of (row r, group g) is u = r * G + g, and its
+            # 16-byte chunk c sits at chunk c ^ (u % 8).
+            t_rows = [r16 * Int32(self.tma_groups) + Int32(q) for q in range(self.tma_groups)]
+            t_off = [u * Int32(128) for u in t_rows]
+            t_x = [u % Int32(8) for u in t_rows]
+            t_base = sW + warp * Int32(self.wgroups * TMA_GROUP_BYTES)
+        dir_lane = Int32(0)
+        d_lane = Int64(0)
+        if cutlass.const_expr(self.direct > 0):
+            if cutlass.const_expr(self.tma):
+                # Lanes of fragment row group g of tile 1 hold CTA row 16 + g.
+                dir_lane = Int32(g < nrows - Int32(TMA_ROWS))
+                d_lane = _row_src(w0, w1, row0 + Int32(TMA_ROWS) + g, k, self.n0, self.n1) + Int64(
+                    (kw0 + t4 * Int32(2)) * Int32(2)
+                )
+            else:
+                # Lanes of fragment row group smem_rows - 16 of tile 1 hold the direct row.
+                dir_lane = Int32(g == Int32(self.smem_rows - 16)) * Int32(nrows > nsm)
+                d_lane = _row_src(
+                    w0, w1, row0 + Int32(self.smem_rows), k, self.n0, self.n1
+                ) + Int64((kw0 + t4 * Int32(2)) * Int32(2))
+
+        zero = Int32(0)
+        nsteps = self.warp_steps
+
+        ld_b = {"cg": ld_cg_b32, "ca": ld_ca_b32, "nc": ld_nc_b32}[self.b_load]
+
+        def load_b(j):
+            off = Int64(32 * j)
+            if cutlass.const_expr(self.diag_no_act):
+                b0 = lane + Int32(j)
+                b1 = lane - Int32(j)
+            else:
+                b0 = ld_b(b_lane + off)
+                b1 = ld_b(b_lane + off + Int64(16))
+            return sel_i32(tok_ok, b0, zero), sel_i32(tok_ok, b1, zero)
+
+        ld_d = ld_ca_b32_if if self.b_load == "ca" else ld_cg_b32_if
+
+        def load_d(j):
+            res = (zero, zero)
+            if cutlass.const_expr(self.direct > 0):
+                off = Int64(32 * j)
+                res = (ld_d(d_lane + off, dir_lane), ld_d(d_lane + off + Int64(16), dir_lane))
+            return res
+
+        def load_a(j):
+            if cutlass.const_expr(j in self.chunk_at):
+                mbar_wait(
+                    sMbar + Int32(8 * self.chunk_at[j] * self.warps) + warp * Int32(8), Int32(0)
+                )
+                if cutlass.const_expr(self.chunk_at[j] == self.nsub - 1):
+                    _stamp(self.debug_ts, ts, bid, tid, 7, 32 * (self.warps - 1))
+            if cutlass.const_expr(self.tma):
+                gk = j // 4
+                q = gk % self.tma_groups
+                ch = Int32(2 * (j % 4)) + hl
+                addr = (
+                    t_base
+                    + Int32((gk - q) * TMA_GROUP_BYTES)
+                    + t_off[q]
+                    + xor_i32(ch, t_x[q]) * Int32(16)
+                )
+                return ldsm_x4(addr), (zero, zero)
+            kb = Int32(32 * j)
+            return (
+                ldsm_x4(sel_i32(ok0, a_lane0 + kb, sZero)),
+                ldsm_x2(sel_i32(ok1, a_lane1 + kb, sZero)),
+            )
+
+        c0 = Int32(2) * t4
+        splits_per_warp = self.nsplit // self.warps
+
+        def stash(acc0, acc1, i):
+            """Slice partials of this warp's local slice ``i`` -> sRed."""
+            sl = warp * Int32(splits_per_warp) + Int32(i)
+            base_s = sRed + sl * Int32(TILE_ROWS * MAX_TOKENS * 4)
+            st_shared_f32(base_s + (g * Int32(MAX_TOKENS) + c0) * Int32(4), acc0[0])
+            st_shared_f32(base_s + (g * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc0[1])
+            st_shared_f32(base_s + ((g + Int32(8)) * Int32(MAX_TOKENS) + c0) * Int32(4), acc0[2])
+            st_shared_f32(
+                base_s + ((g + Int32(8)) * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc0[3]
+            )
+            st_shared_f32(base_s + ((g + Int32(16)) * Int32(MAX_TOKENS) + c0) * Int32(4), acc1[0])
+            st_shared_f32(
+                base_s + ((g + Int32(16)) * Int32(MAX_TOKENS) + c0 + Int32(1)) * Int32(4), acc1[1]
+            )
+
+        bq = {}
+        dq = {}
+        aq = {}
+        for j in cutlass.range_constexpr(min(self.b_dist, nsteps)):
+            bq[j] = load_b(j)
+            dq[j] = load_d(j)
+        for j in cutlass.range_constexpr(min(self.a_dist, nsteps)):
+            aq[j] = load_a(j)
+        acc0 = [Float32(0.0)] * 4
+        acc1 = [Float32(0.0)] * 4
+        for j in cutlass.range_constexpr(nsteps):
+            if cutlass.const_expr(j + self.b_dist < nsteps):
+                bq[j + self.b_dist] = load_b(j + self.b_dist)
+                dq[j + self.b_dist] = load_d(j + self.b_dist)
+            if cutlass.const_expr(j + self.a_dist < nsteps):
+                aq[j + self.a_dist] = load_a(j + self.a_dist)
+            if cutlass.const_expr(j > 0 and j % self.split_steps == 0):
+                stash(acc0, acc1, j // self.split_steps - 1)
+                acc0 = [Float32(0.0)] * 4
+                acc1 = [Float32(0.0)] * 4
+            fa, fb = aq.pop(j)
+            b0, b1 = bq.pop(j)
+            fb0, fb1 = fb
+            d0, d1 = dq.pop(j)
+            if cutlass.const_expr(self.direct > 0):
+                fb0 = sel_i32(dir_lane, d0, fb0)
+                fb1 = sel_i32(dir_lane, d1, fb1)
+            acc0[0], acc0[1], acc0[2], acc0[3] = mma_bf16(
+                fa[0], fa[1], fa[2], fa[3], b0, b1, acc0[0], acc0[1], acc0[2], acc0[3]
+            )
+            acc1[0], acc1[1], acc1[2], acc1[3] = mma_bf16(
+                fb0, zero, fb1, zero, b0, b1, acc1[0], acc1[1], acc1[2], acc1[3]
+            )
+        stash(acc0, acc1, splits_per_warp - 1)
+        _stamp(self.debug_ts, ts, bid, tid, 6)
+        _stamp(self.debug_ts, ts, bid, tid, 3, 32 * (self.warps - 1))
+
+    @cute.kernel
+    def kernel(
+        self,
+        act: Int64,
+        w0: Int64,
+        w1: Int64,
+        out0: Int64,
+        out1: Int64,
+        ts: Int64,
+        tmaps: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+    ):
+        """``debug_ts`` stamps per CTA (int64 ``[ctas, 8]``): 0 entry, 1 copies
+        issued (thread 0 passed the issue: copy warp 0 only), 2 after the
+        dependency wait, 7 last chunk of the last MMA warp landed, 3 main loop
+        done (last MMA warp), 6 main loop done (MMA warp 0), 4 rows stored,
+        5 exit."""
+        k = self.k
+        stride = self.stride
+        tid, _, _ = cute.arch.thread_idx()
+        bid, _, _ = cute.arch.block_idx()
+        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        lane = cute.arch.lane_idx()
+        _stamp(self.debug_ts, ts, bid, tid, 0)
+
+        smem = SmemAllocator()
+        sW = saddr(smem.allocate(self.w_bytes, byte_alignment=128))
+        if cutlass.const_expr(self.tma):
+            # Swizzled tensor-copy destinations start on 1024-byte boundaries.
+            sW = (sW + Int32(1023)) // Int32(1024) * Int32(1024)
+        sMisc = saddr(smem.allocate(_MISC_BYTES, byte_alignment=128))
+        sRed = saddr(smem.allocate(self.red_bytes, byte_alignment=128))
+        # sMisc: [0, 32) zero chunk, [64, 64 + 8 * nmbar) mbarriers.
+        sZero = sMisc
+        sMbar = sMisc + Int32(64)
+
+        # Rows of this CTA: [row0, row0 + nrows) of the concatenation [W0; W1]
+        # (local row lrow0 of tensor map tmap for ``tma``).
+        tmap = tmaps
+        if cutlass.const_expr(self.tma):
+            in1 = Int32(0)
+            if cutlass.const_expr(self.n1 > 0):
+                in1 = Int32(bid >= Int32(self.c0))
+            b = bid - in1 * Int32(self.c0)
+            base = sel_i32(in1, Int32(self.n1 // max(self.c1, 1)), Int32(self.n0 // self.c0))
+            rem = sel_i32(in1, Int32(self.n1 % max(self.c1, 1)), Int32(self.n0 % self.c0))
+            nrows = base + Int32(b < rem)
+            lrow0 = b * base + min_i32(b, rem)
+            row0 = lrow0 + in1 * Int32(self.n0)
+            nsm = Int32(TMA_ROWS)
+            tmap = tmaps + Int64(in1 * Int32(128))
+        else:
+            base = Int32(self.base_rows)
+            rem = Int32(self.rem_rows)
+            nrows = base + Int32(bid < rem)
+            row0 = bid * base + min_i32(bid, rem)
+            lrow0 = row0
+            nsm = min_i32(nrows, Int32(self.smem_rows))
+
+        if tid < Int32(2):
+            st_shared_v4(sZero + tid * Int32(16), Int32(0), Int32(0), Int32(0), Int32(0))
+        if tid == Int32(0):
+            for c in cutlass.range_constexpr(self.nmbar):
+                mbar_init(
+                    sMbar + Int32(8 * c),
+                    Int32(32 * self.producer_warps if self.copy_method == "cpasync" else 1),
+                )
+            mbar_fence_init()
+        cute.arch.sync_threads()
+
+        if cutlass.const_expr(self.pre_chunks > 0):
+            self._issue_copies(sW, sMbar, w0, w1, row0, nsm, warp, lane, tid, tmap, lrow0, nrows, 0)
+        if cutlass.const_expr(self.direct > 0 and self.l2_prefetch and not self.tma):
+            if tid == Int32(0 if self.producer_warp else 32 % self.nt):
+                if nrows > nsm:
+                    prefetch_l2(_row_src(w0, w1, row0 + nsm, k, self.n0, self.n1), Int32(k * 2))
+        _stamp(self.debug_ts, ts, bid, tid, 1, 32 * self.copy_warp)
+        if cutlass.const_expr(self.use_pdl and self.trigger == "start"):
+            cute.arch.griddepcontrol_launch_dependents()
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+        _stamp(self.debug_ts, ts, bid, tid, 2)
+        if cutlass.const_expr(self.pre_chunks < self.nsub):
+            self._issue_copies(sW, sMbar, w0, w1, row0, nsm, warp, lane, tid, tmap, lrow0, nrows, 1)
+
+        if cutlass.const_expr(self.copy_only):
+            if warp < Int32(self.warps):
+                for sub in cutlass.range_constexpr(self.nsub):
+                    mbar_wait(sMbar + Int32(8 * sub * self.warps) + warp * Int32(8), Int32(0))
+        elif cutlass.const_expr(self.producer_warp):
+            if warp < Int32(self.warps):
+                self._mma(
+                    act,
+                    w0,
+                    w1,
+                    ts,
+                    sW,
+                    sMbar,
+                    sZero,
+                    sRed,
+                    bid,
+                    tid,
+                    warp,
+                    lane,
+                    rows,
+                    row0,
+                    nrows,
+                    nsm,
+                )
+        else:
+            self._mma(
+                act,
+                w0,
+                w1,
+                ts,
+                sW,
+                sMbar,
+                sZero,
+                sRed,
+                bid,
+                tid,
+                warp,
+                lane,
+                rows,
+                row0,
+                nrows,
+                nsm,
+            )
+
+        cute.arch.sync_threads()
+        if cutlass.const_expr(not self.copy_only):
+            pairs = TILE_ROWS * MAX_TOKENS
+            for i in cutlass.range_constexpr((pairs + self.nt - 1) // self.nt):
+                idx = tid + Int32(self.nt * i)
+                row = idx // Int32(MAX_TOKENS)
+                tok = idx % Int32(MAX_TOKENS)
+                if idx < Int32(pairs):
+                    if row < nrows:
+                        if tok < rows:
+                            vals = [
+                                ld_shared_f32(sRed + Int32(sl * pairs * 4) + idx * Int32(4))
+                                for sl in range(self.nsplit)
+                            ]
+                            s = _ordered_sum(vals, self.reduce, self.ftz)
+                            _store_out(
+                                out0,
+                                out1,
+                                row0 + row,
+                                tok,
+                                add_f32(Float32(0.0), s, self.ftz),
+                                self.n0,
+                                self.n1,
+                            )
+        _stamp(self.debug_ts, ts, bid, tid, 4)
+        if cutlass.const_expr(self.use_pdl and self.trigger == "end"):
+            cute.arch.griddepcontrol_launch_dependents()
+        _stamp(self.debug_ts, ts, bid, tid, 5)
+
+
+_COMPILED = {}
+
+
+def is_compiled(**cfg) -> bool:
+    """Whether ``compiled_projection(**cfg)`` would be a cache hit."""
+    return tuple(sorted(cfg.items())) in _COMPILED
+
+
+def compiled_projection(**cfg):
+    """Return the compiled projection for one configuration (compiles once, no GPU work)."""
+    key = tuple(sorted(cfg.items()))
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = FusedKvProjKernel(**cfg)
+        z64 = Int64(0)
+        anchor = cute_runtime.make_fake_compact_tensor(
+            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
+        )
+        compiled = cute.compile(
+            kernel,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            z64,
+            anchor,
+            Int32(1),
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+        _COMPILED[key] = compiled
+    return compiled
+
+
+def projection_supported(
+    k: int,
+    n: int,
+    ctas: int,
+    warps: int,
+    smem_rows: int,
+    nsplit: int,
+    nsub: int,
+    chunks: tuple = (),
+    n1: int = 0,
+    copy_method: str = "bulk",
+    tma_groups: int = 1,
+) -> Optional[str]:
+    """Return why the projection kernel cannot serve a shape, or ``None``."""
+    steps = k // 16
+    if k % 16 or nsplit % warps or steps % nsplit:
+        return f"K={k} does not split into {nsplit} slices over {warps} warps"
+    if chunks:
+        if any(c <= 0 for c in chunks) or sum(chunks) != steps // warps:
+            return f"chunks {chunks} do not cover the {steps // warps} K steps per warp"
+    elif (steps // warps) % nsub:
+        return f"{steps // warps} K steps per warp do not split into {nsub} sub-chunks"
+    else:
+        chunks = ((steps // warps) // nsub,) * nsub
+    nsub = len(chunks)
+    if 64 + 8 * warps * nsub > _MISC_BYTES:
+        return f"{warps * nsub} mbarriers do not fit"
+    if copy_method == "tma":
+        if (k // warps) % TMA_GROUP_K:
+            return f"K slice {k // warps} is not a multiple of {TMA_GROUP_K}"
+        if tma_groups not in (1, 3) or any(c % (4 * tma_groups) for c in chunks):
+            return f"chunks {chunks} are not whole boxes of {tma_groups} K groups"
+        if split_ctas(n - n1, n1, ctas) is None:
+            return f"no CTA split of {n - n1} + {n1} rows gives {TMA_ROWS}..{TILE_ROWS} rows"
+        smem = (
+            (k // TMA_GROUP_K) * TMA_GROUP_BYTES
+            + 1024
+            + _MISC_BYTES
+            + nsplit * TILE_ROWS * MAX_TOKENS * 4
+        )
+        if smem > SMEM_LIMIT:
+            return f"{smem} B of shared memory exceed the limit"
+        return None
+    base, rem = divmod(n, ctas)
+    max_rows = base + (1 if rem else 0)
+    if base < 1 or max_rows > TILE_ROWS:
+        return f"{n} rows over {ctas} CTAs gives {base}..{max_rows} rows per CTA (1..24 supported)"
+    staged = min(smem_rows, max_rows)
+    if max_rows - staged > 1:
+        return f"{max_rows} rows per CTA exceed {smem_rows} staged rows by more than one"
+    if max_rows > staged and not 16 <= staged < TILE_ROWS:
+        return "a directly-read row must fall into rows 16..23 of the CTA"
+    smem = staged * (k * 2 + 16) + _MISC_BYTES + nsplit * TILE_ROWS * MAX_TOKENS * 4
+    if smem > SMEM_LIMIT:
+        return f"{smem} B of shared memory exceed the limit"
+    return None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/projection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/projection.py
new file mode 100644
index 000000000000..f64a8dfbbdb6
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/projection.py
@@ -0,0 +1,316 @@
+# 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.
+"""fused_kv_indexer_projection merged projection bound to one native DSA ``MLA`` module.
+
+``FusedKvProjection.run(h)`` takes the normalized hidden state ``h [R, 6144]``
+that the native ``kv_a_proj_with_mqa`` GEMM and the native indexer
+``wk | weights_proj`` GEMM consume, and returns ``(kva [R, 2624], kw [R,
+160] or None)`` from one kernel launch, with the dtypes, shapes and strides
+of the native ``dsv3_fused_a_gemm_op`` / ``F.linear`` outputs. It reads the
+module's own weight tensors in place (``kv_a_proj_with_mqa.weight`` and,
+with ``TRTLLM_DSA_INDEXER_BF16=1``, ``indexer._fused_wk_wp_weight``); no
+packed copy exists, so the extra weight memory is zero.
+"""
+
+from dataclasses import dataclass, fields, replace
+from typing import Optional, Tuple
+
+import torch
+
+_K = 6144
+_N_KVA = 2624
+_N_KW = 160
+MAX_ROWS = 8
+# K slices of the native fused-A GEMM (one per MMA warp): the order fused_kv_indexer_projection
+# keeps.
+NATIVE_NSPLIT = 4
+
+
+@dataclass(frozen=True)
+class FusedKvProjConfig:
+    """Launch configuration of the projection kernel (no numerical effect).
+
+    ``ctas = 0`` means one CTA per SM. ``stage_before_wait`` copies the
+    weights into shared memory before the PDL dependency wait; ``False``
+    copies them after the wait (attribution arm). ``nsub``: sub-chunks per
+    warp slice of the staging order; ``b_dist`` / ``a_dist``: K steps of
+    activation / weight fragments loaded ahead; ``producer_warp``: dedicated
+    copy warp(s) (else MMA warp 0 issues the bulk copies); ``copy_method``:
+    ``bulk`` (1-D ``cp.async.bulk``) or ``cpasync`` (16-byte ``cp.async``
+    over ``producer_warps`` warps); ``chunks``: unequal staging chunks of a
+    warp slice in K steps, ``-``-separated (overrides ``nsub``; empty means
+    ``nsub`` equal chunks); ``b_load``: cache operator of the activation
+    loads (``cg``: L2 only, ``ca``: L1 and L2); ``copy_method=tma``: 3-D
+    tensor copies of 16 rows x ``tma_groups`` K groups of 64 (``tma_groups``
+    1 or 3); ``prewait_chunks``: chunks staged before the dependency wait
+    (``-1``: all; the rest is staged after the wait).
+    """
+
+    ctas: int = 0
+    smem_rows: int = 18
+    nsub: int = 4
+    trigger: str = "start"
+    stage_before_wait: bool = True
+    l2_prefetch: bool = True
+    b_dist: int = 32
+    a_dist: int = 2
+    producer_warp: bool = True
+    copy_method: str = "bulk"
+    producer_warps: int = 1
+    chunks: str = "48-24-12-12"
+    b_load: str = "ca"
+    tma_groups: int = 1
+    prewait_chunks: int = -1
+    # Tests only: stage the weights and exit without computing (streaming floor).
+    copy_only: bool = False
+    # Benchmarks only: no activation loads (meaningless output; chain cost).
+    diag_no_act: bool = False
+
+    def chunk_steps(self) -> tuple:
+        return tuple(int(c) for c in self.chunks.split("-")) if self.chunks else ()
+
+    @classmethod
+    def parse(cls, text: str, base: Optional["FusedKvProjConfig"] = None) -> "FusedKvProjConfig":
+        """``base`` (default: the served configuration) with ``key=value,...``
+        applied; an ``nsub`` without ``chunks`` selects ``nsub`` equal chunks."""
+        kwargs = {}
+        types = {f.name: f.type for f in fields(cls)}
+        for item in filter(None, text.split(",")):
+            key, value = item.split("=")
+            typ = types[key]
+            if typ in (bool, "bool"):
+                kwargs[key] = value.lower() in ("1", "true")
+            elif typ in (int, "int"):
+                kwargs[key] = int(value)
+            else:
+                kwargs[key] = value
+        if "nsub" in kwargs and "chunks" not in kwargs:
+            kwargs["chunks"] = ""
+        return replace(base or cls(), **kwargs)
+
+
+def _sm_count(device) -> int:
+    return torch.cuda.get_device_properties(device).multi_processor_count
+
+
+def kernel_config(
+    config: FusedKvProjConfig,
+    n0: int,
+    n1: int,
+    ctas: int,
+    use_pdl: bool,
+    nsplit: int = NATIVE_NSPLIT,
+    warps: int = NATIVE_NSPLIT,
+    reduce: str = "seq",
+    ftz: bool = False,
+    debug_ts: bool = False,
+) -> dict:
+    """Keyword arguments of ``FusedKvProjKernel`` / ``compiled_projection``."""
+    return dict(
+        k=_K,
+        n0=n0,
+        n1=n1,
+        ctas=ctas,
+        warps=warps,
+        smem_rows=config.smem_rows,
+        nsplit=nsplit,
+        nsub=config.nsub,
+        use_pdl=use_pdl,
+        trigger=config.trigger,
+        stage_before_wait=config.stage_before_wait,
+        l2_prefetch=config.l2_prefetch,
+        b_dist=config.b_dist,
+        a_dist=config.a_dist,
+        reduce=reduce,
+        ftz=ftz,
+        producer_warp=config.producer_warp,
+        copy_method=config.copy_method,
+        producer_warps=config.producer_warps,
+        copy_only=config.copy_only,
+        debug_ts=debug_ts,
+        chunks=config.chunk_steps(),
+        b_load=config.b_load,
+        diag_no_act=config.diag_no_act,
+        tma_groups=config.tma_groups,
+        prewait_chunks=config.prewait_chunks,
+    )
+
+
+def _tensor_map_words(w: torch.Tensor, groups: int) -> list:
+    """128-byte ``CUtensorMap`` (16 signed 64-bit words) of the row-major BF16
+    weight ``w [N, K]`` viewed as ``(64, K / 64, N)`` (elements, K groups,
+    rows), box ``(64, groups, 16)``, 128-byte swizzle."""
+    import ctypes
+
+    from cuda.bindings import driver as cu
+
+    from .proj_kernel import TMA_GROUP_K, TMA_ROWS
+
+    n, k = w.shape
+    dims = [TMA_GROUP_K, k // TMA_GROUP_K, n]
+    strides = [TMA_GROUP_K * 2, k * 2]
+    box = [TMA_GROUP_K, groups, TMA_ROWS]
+    err, tmap = cu.cuTensorMapEncodeTiled(
+        cu.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
+        3,
+        w.data_ptr(),
+        [cu.cuuint64_t(d) for d in dims],
+        [cu.cuuint64_t(d) for d in strides],
+        [cu.cuuint32_t(d) for d in box],
+        [cu.cuuint32_t(1)] * 3,
+        cu.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE,
+        cu.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,
+        cu.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
+        cu.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
+    )
+    if err != cu.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"cuTensorMapEncodeTiled failed: {err}")
+    if hasattr(tmap, "getPtr"):
+        raw = ctypes.string_at(int(tmap.getPtr()), 128)
+        return [int.from_bytes(raw[8 * i : 8 * i + 8], "little", signed=True) for i in range(16)]
+    words = [int(v) for v in tmap.opaque]
+    return [v - (1 << 64) if v >= (1 << 63) else v for v in words]
+
+
+def projection_eligibility(mla) -> Tuple[bool, str, Optional[torch.Tensor]]:
+    """Static check of one MLA module; returns ``(ok, reason, kw_weight)``.
+
+    ``kw_weight`` is the indexer's BF16 fused ``wk | weights_proj`` weight on
+    a full-indexer layer (``None`` on a shared-indexer layer).
+    """
+    from tensorrt_llm._utils import get_sm_version
+
+    if get_sm_version() not in (100, 103):
+        return False, f"sm{get_sm_version()} not validated", None
+    lin = getattr(mla, "kv_a_proj_with_mqa", None)
+    if lin is None or not getattr(mla, "fuse_qkv_a_proj", False):
+        return False, "needs the fused [q_a | kv_a] projection", None
+    w = lin.weight
+    if (
+        lin.has_any_quant
+        or lin.bias is not None
+        or w.dtype != torch.bfloat16
+        or tuple(w.shape) != (_N_KVA, _K)
+        or not w.is_contiguous()
+    ):
+        return False, "kv_a_proj_with_mqa must be an unquantized BF16 [2624, 6144] weight", None
+    if getattr(lin, "use_cute_dsl_bf16_gemm", False):
+        return False, "kv_a_proj_with_mqa runs the CuTe DSL BF16 GEMM, not fused-A", None
+    kw = None
+    indexer = mla.mqa.indexer if getattr(mla, "mqa", None) is not None else None
+    if indexer is not None:
+        kw = getattr(indexer, "_fused_wk_wp_weight", None)
+        if (
+            not getattr(indexer, "_indexer_bf16", False)
+            or kw is None
+            or kw.dtype != torch.bfloat16
+            or tuple(kw.shape) != (_N_KW, _K)
+            or not kw.is_contiguous()
+            or getattr(indexer, "use_fp4", False)
+        ):
+            return False, "full-indexer layer needs the BF16 fused wk|weights_proj GEMM", None
+    return True, "", kw
+
+
+class FusedKvProjection:
+    """Merged projection bound to one MLA module (see the module docstring)."""
+
+    def __init__(
+        self,
+        mla,
+        config: Optional[FusedKvProjConfig] = None,
+        kw_weight: Optional[torch.Tensor] = None,
+    ):
+        from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+        self.mla = mla
+        self.config = config or FusedKvProjConfig()
+        self.w0 = mla.kv_a_proj_with_mqa.weight
+        self.w1 = kw_weight
+        self.n1 = 0 if kw_weight is None else kw_weight.shape[0]
+        self.use_pdl = get_env_enable_pdl()
+        self.ctas = self.config.ctas or _sm_count(self.w0.device)
+        # Tensor maps of the weights (``copy_method=tma``), built by precompile().
+        self.tmaps = None
+
+    def kernel_config(self, debug_ts: bool = False) -> dict:
+        return kernel_config(
+            self.config, _N_KVA, self.n1, self.ctas, self.use_pdl, debug_ts=debug_ts
+        )
+
+    def supported(self) -> Optional[str]:
+        from .proj_kernel import projection_supported
+
+        cfg = self.kernel_config()
+        if self.config.b_load not in ("cg", "ca"):
+            return f"activation load {self.config.b_load!r} is a benchmark-only mode (not served)"
+        if self.config.copy_only or self.config.diag_no_act:
+            return "copy_only / diag_no_act is a benchmark-only mode (not served)"
+        return projection_supported(
+            cfg["k"],
+            cfg["n0"] + cfg["n1"],
+            cfg["ctas"],
+            cfg["warps"],
+            cfg["smem_rows"],
+            cfg["nsplit"],
+            cfg["nsub"],
+            cfg["chunks"],
+            cfg["n1"],
+            cfg["copy_method"],
+            cfg["tma_groups"],
+        )
+
+    def precompile(self) -> None:
+        """Compile the served kernel and, for ``copy_method=tma``, build the
+        weights' tensor maps (one 256-byte device buffer). Errors propagate."""
+        from .proj_kernel import compiled_projection
+
+        compiled_projection(**self.kernel_config())
+        if self.config.copy_method == "tma" and self.tmaps is None:
+            words = _tensor_map_words(self.w0, self.config.tma_groups)
+            if self.w1 is not None:
+                words += _tensor_map_words(self.w1, self.config.tma_groups)
+            self.tmaps = torch.tensor(words, dtype=torch.int64).to(self.w0.device)
+
+    def run(
+        self, h: torch.Tensor, ts: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Launch on the current stream; ``h`` is the native normalized hidden state."""
+        from .proj_kernel import compiled_projection
+
+        rows = h.shape[0]
+        assert h.dim() == 2 and h.shape[1] == _K and 1 <= rows <= MAX_ROWS
+        assert h.is_contiguous() and h.dtype == torch.bfloat16
+        kva = h.new_empty(rows, _N_KVA)
+        kw = h.new_empty(rows, self.n1) if self.n1 else None
+        compiled = compiled_projection(**self.kernel_config(debug_ts=ts is not None))
+        if self.config.copy_method == "tma" and self.tmaps is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError(
+                    "fused_kv_indexer_projection tensor maps must be built (precompile) before capture"
+                )
+            self.precompile()
+        compiled(
+            h.data_ptr(),
+            self.w0.data_ptr(),
+            0 if self.w1 is None else self.w1.data_ptr(),
+            kva.data_ptr(),
+            0 if kw is None else kw.data_ptr(),
+            0 if ts is None else ts.data_ptr(),
+            0 if self.tmaps is None else self.tmaps.data_ptr(),
+            kva.view(torch.int32).view(-1),
+            rows,
+        )
+        return kva, kw
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/ptx.py
new file mode 100644
index 000000000000..d62d4670bc9b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/ptx.py
@@ -0,0 +1,475 @@
+# 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.
+"""Inline-PTX helpers of the fused_kv_indexer_projection kernel (CuTe DSL).
+
+fused_kv_indexer_projection owns every device helper its kernel uses. The arithmetic helpers spell
+out the rounding of each operation (``add.rn``, ``cvt.rn``) so that the
+kernel's result is defined by the instruction sequence alone.
+"""
+
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+
+def _asm(res, args, text, cons, side_effects=True, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=side_effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+def _struct(n, ty):
+    return llvm.StructType.get_literal([ty] * n)
+
+
+def _unpack(res, n, cls, ty, loc, ip):
+    return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n))
+
+
+# --------------------------------------------------------------------------
+# Shared memory, mbarriers and bulk copies.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def saddr(p: cute.Pointer, *, loc=None, ip=None) -> Int32:
+    """32-bit shared-window address of a shared-memory pointer."""
+    return Int32(p.toint(loc=loc, ip=ip))
+
+
+@dsl_user_op
+def bulk_g2s(smem: Int32, gaddr: Int64, nbytes: Int32, mbar: Int32, *, loc=None, ip=None):
+    """1-D bulk copy global -> shared that completes ``nbytes`` on ``mbar``."""
+    _asm(
+        None,
+        [smem.ir_value(), gaddr.ir_value(), nbytes.ir_value(), mbar.ir_value()],
+        "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];",
+        "r,l,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def tensor_g2s_3d(
+    smem: Int32, tmap: Int64, c0: Int32, c1: Int32, c2: Int32, mbar: Int32, *, loc=None, ip=None
+):
+    """3-D tensor (TMA) copy global -> shared of the box at ``(c0, c1, c2)`` of
+    the tensor map at global address ``tmap``; completes the box bytes on ``mbar``."""
+    _asm(
+        None,
+        [
+            smem.ir_value(),
+            tmap.ir_value(),
+            c0.ir_value(),
+            c1.ir_value(),
+            c2.ir_value(),
+            mbar.ir_value(),
+        ],
+        "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes "
+        "[$0], [$1, {$2, $3, $4}], [$5];",
+        "r,l,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def fence_tensormap_acquire(tmap: Int64, *, loc=None, ip=None):
+    """Acquire a tensor map that the host wrote into global memory (the
+    tensormap proxy is not acquired at kernel start)."""
+    _asm(
+        None,
+        [tmap.ir_value()],
+        "fence.proxy.tensormap::generic.acquire.sys [$0], 128;",
+        "l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def prefetch_l2(gaddr: Int64, nbytes: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [gaddr.ir_value(), nbytes.ir_value()],
+        "cp.async.bulk.prefetch.L2.global [$0], $1;",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mbar_init(mbar: Int32, count: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [mbar.ir_value(), count.ir_value()],
+        "mbarrier.init.shared::cta.b64 [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mbar_fence_init(*, loc=None, ip=None):
+    """Make the mbarrier initializations visible to the async proxy."""
+    _asm(
+        None,
+        [],
+        "fence.mbarrier_init.release.cluster;\n\tfence.proxy.async.shared::cta;",
+        "",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mbar_arrive_expect_tx(mbar: Int32, nbytes: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [mbar.ir_value(), nbytes.ir_value()],
+        "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], $1;",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mbar_wait(mbar: Int32, phase: Int32, *, loc=None, ip=None):
+    """Spin until the phase with parity ``phase`` of ``mbar`` has completed."""
+    _asm(
+        None,
+        [mbar.ir_value(), phase.ir_value()],
+        "{\n\t.reg .pred p;\n\tFUSED_KV_MBAR_WAIT:\n\t"
+        "mbarrier.try_wait.parity.shared::cta.b64 p, [$0], $1;\n\t"
+        "@!p bra FUSED_KV_MBAR_WAIT;\n\t}",
+        "r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def cp_async16(smem: Int32, gaddr: Int64, *, loc=None, ip=None):
+    """16-byte ``cp.async`` global -> shared (L2 only)."""
+    _asm(
+        None,
+        [smem.ir_value(), gaddr.ir_value()],
+        "cp.async.cg.shared.global [$0], [$1], 16;",
+        "r,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def cp_async_mbar_arrive(mbar: Int32, *, loc=None, ip=None):
+    """Arrive on ``mbar`` once this thread's prior ``cp.async`` copies have
+    completed (``.noinc``: the arrival counts against the initial count)."""
+    _asm(
+        None,
+        [mbar.ir_value()],
+        "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [$0];",
+        "r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def sync_warp(*, loc=None, ip=None):
+    _asm(None, [], "bar.warp.sync 0xffffffff;", "", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def st_shared_v4(smem: Int32, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [smem.ir_value(), a.ir_value(), b.ir_value(), c.ir_value(), d.ir_value()],
+        "st.shared.v4.b32 [$0], {$1, $2, $3, $4};",
+        "r,r,r,r,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_shared_f32(smem: Int32, v: Float32, *, loc=None, ip=None):
+    _asm(None, [smem.ir_value(), v.ir_value()], "st.shared.f32 [$0], $1;", "r,f", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def ld_shared_f32(smem: Int32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(T.f32(), [smem.ir_value()], "ld.shared.f32 $0, [$1];", "=f,r", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def ldsm_x4(smem: Int32, *, loc=None, ip=None):
+    out = _asm(
+        _struct(4, T.i32()),
+        [smem.ir_value()],
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(out, 4, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def ldsm_x2(smem: Int32, *, loc=None, ip=None):
+    out = _asm(
+        _struct(2, T.i32()),
+        [smem.ir_value()],
+        "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {$0, $1}, [$2];",
+        "=r,=r,r",
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(out, 2, Int32, T.i32(), loc, ip)
+
+
+# --------------------------------------------------------------------------
+# Global memory.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def ld_cg_b32(gaddr: Int64, *, loc=None, ip=None) -> Int32:
+    """4-byte global load cached in L2 only (no L1 line can be stale)."""
+    return Int32(
+        _asm(T.i32(), [gaddr.ir_value()], "ld.global.cg.b32 $0, [$1];", "=r,l", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def ld_ca_b32(gaddr: Int64, *, loc=None, ip=None) -> Int32:
+    """4-byte global load cached in L1 and L2 (the default ``ld.global``)."""
+    return Int32(
+        _asm(T.i32(), [gaddr.ir_value()], "ld.global.ca.b32 $0, [$1];", "=r,l", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def ld_nc_b32(gaddr: Int64, *, loc=None, ip=None) -> Int32:
+    """4-byte global load through the non-coherent read-only path (benchmarks only)."""
+    return Int32(
+        _asm(T.i32(), [gaddr.ir_value()], "ld.global.nc.b32 $0, [$1];", "=r,l", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def ld_cg_b32_if(gaddr: Int64, pred: Int32, *, loc=None, ip=None) -> Int32:
+    """``ld_cg_b32`` if ``pred`` else 0 (no access when ``pred`` is 0)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [gaddr.ir_value(), pred.ir_value()],
+            "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\tmov.b32 $0, 0;\n\t"
+            "@p ld.global.cg.b32 $0, [$1];\n\t}",
+            "=r,l,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def ld_ca_b32_if(gaddr: Int64, pred: Int32, *, loc=None, ip=None) -> Int32:
+    """``ld_ca_b32`` if ``pred`` else 0 (no access when ``pred`` is 0)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [gaddr.ir_value(), pred.ir_value()],
+            "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\tmov.b32 $0, 0;\n\t"
+            "@p ld.global.ca.b32 $0, [$1];\n\t}",
+            "=r,l,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def st_global_b16(gaddr: Int64, v: Int32, *, loc=None, ip=None):
+    _asm(
+        None,
+        [gaddr.ir_value(), v.ir_value()],
+        "{\n\t.reg .b16 lo;\n\tcvt.u16.u32 lo, $1;\n\tst.global.b16 [$0], lo;\n\t}",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def st_global_b64(gaddr: Int64, v: Int64, *, loc=None, ip=None):
+    _asm(None, [gaddr.ir_value(), v.ir_value()], "st.global.b64 [$0], $1;", "l,l", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def globaltimer(*, loc=None, ip=None) -> Int64:
+    return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc=loc, ip=ip))
+
+
+# --------------------------------------------------------------------------
+# Integer selects.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def sel_i32(pred: Int32, a: Int32, b: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [pred.ir_value(), a.ir_value(), b.ir_value()],
+            "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $1, 0;\n\tselp.b32 $0, $2, $3, p;\n\t}",
+            "=r,r,r,r",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def sel_i64(pred: Int32, a: Int64, b: Int64, *, loc=None, ip=None) -> Int64:
+    return Int64(
+        _asm(
+            T.i64(),
+            [pred.ir_value(), a.ir_value(), b.ir_value()],
+            "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $1, 0;\n\tselp.b64 $0, $2, $3, p;\n\t}",
+            "=l,r,l,l",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def xor_i32(a: Int32, b: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(T.i32(), [a.ir_value(), b.ir_value()], "xor.b32 $0, $1, $2;", "=r,r,r", loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def min_i32(a: Int32, b: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [a.ir_value(), b.ir_value()],
+            "min.s32 $0, $1, $2;",
+            "=r,r,r",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+# --------------------------------------------------------------------------
+# Arithmetic.
+# --------------------------------------------------------------------------
+
+
+@dsl_user_op
+def mma_bf16(
+    a0: Int32,
+    a1: Int32,
+    a2: Int32,
+    a3: Int32,
+    b0: Int32,
+    b1: Int32,
+    c0: Float32,
+    c1: Float32,
+    c2: Float32,
+    c3: Float32,
+    *,
+    loc=None,
+    ip=None,
+):
+    """``mma.sync.m16n8k16`` BF16 x BF16 + FP32 (row.col), the instruction of
+    the native fused-A GEMM."""
+    out = _asm(
+        _struct(4, T.f32()),
+        [
+            a0.ir_value(),
+            a1.ir_value(),
+            a2.ir_value(),
+            a3.ir_value(),
+            b0.ir_value(),
+            b1.ir_value(),
+            c0.ir_value(),
+            c1.ir_value(),
+            c2.ir_value(),
+            c3.ir_value(),
+        ],
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
+        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+        side_effects=False,
+        loc=loc,
+        ip=ip,
+    )
+    return _unpack(out, 4, Float32, T.f32(), loc, ip)
+
+
+@dsl_user_op
+def add_f32(a: Float32, b: Float32, ftz: bool = False, *, loc=None, ip=None) -> Float32:
+    """IEEE round-to-nearest fp32 add (``.ftz`` flushes subnormal operands and results)."""
+    op = "add.rn.ftz.f32" if ftz else "add.rn.f32"
+    return Float32(
+        _asm(
+            T.f32(),
+            [a.ir_value(), b.ir_value()],
+            f"{op} $0, $1, $2;",
+            "=f,f,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def bf16_bits(x: Float32, *, loc=None, ip=None) -> Int32:
+    """Round-to-nearest-even of one fp32 value to bf16 (in the low half)."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [x.ir_value()],
+            "{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tcvt.u32.u16 $0, h;\n\t}",
+            "=r,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/__init__.py
new file mode 100644
index 000000000000..3d4e4f667afa
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/__init__.py
@@ -0,0 +1,31 @@
+# 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.
+"""host_page_retention: rank-skew reduction at the decode hops by removing its host-side cause.
+
+On unless ``TRTLLM_STAIRCASE_GLM_LL_DISABLE`` names it or a path it needs
+(``post_attention_hop``). Entry point (``dispatch.get_impl``):
+
+* ``host_policy``: called once per rank after warmup; stops the rank's CPU
+  allocator from returning freed pages to the OS while it serves, so host
+  work during a decode step does not invalidate the GPU's address
+  translations (``host_policy`` module docstring).
+"""
+
+from ...dispatch import register
+from .host_policy import install_host_policy
+
+register("host_page_retention", "host_policy", install_host_policy)
+
+__all__ = ["install_host_policy"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/host_policy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/host_policy.py
new file mode 100644
index 000000000000..7e8f1ce2c1c6
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/host_policy.py
@@ -0,0 +1,288 @@
+# 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.
+"""host_page_retention host policy: no immediate CPU-allocator page purges in a serving decode rank.
+
+On a GPU that translates through ATS (an sm_103 GPU attached to Grace), every page-table change
+of a process holding a CUDA context on that GPU (``munmap``,
+``madvise(MADV_DONTNEED)``, ``mprotect``) also invalidates the GPU's address
+translations, so TLB-miss-heavy kernels running at that moment (the sparse KV
+gathers of MLA decode and of the indexer logits) slow down. PyTorch allocates
+CPU tensors through a mimalloc that is linked statically into ``libc10.so``;
+with ``MIMALLOC_PURGE_DELAY=0`` every free that empties an allocator page
+returns it to the OS at once, so the decode rank's ordinary host work (input
+preparation, sampling bookkeeping) issues such invalidations while the GPU
+runs the step.
+
+``install_host_policy`` sets mimalloc's ``purge_delay`` option of this process
+to ``PURGE_DELAY_MS``: a freed allocator page goes back to the OS only after
+it has stayed unused that long. The pages that each decode step's host work
+frees are allocated again well within that time, so they are reused instead
+of purged during the step, while memory that stays idle is still returned. It
+changes no device operation, no kernel and no allocation size.
+mimalloc exports none of its symbols from ``libc10.so``, so the option
+functions are located through the library's ELF symbol table and the option
+index is resolved by name from mimalloc's option table; every step is checked
+and any mismatch leaves the allocator untouched (native fallback).
+"""
+
+import ctypes
+import os
+import struct
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+# mimalloc's purge_delay: -1 never purges, 0 purges at once, N > 0 after N ms.
+PURGE_DELAY_MS = 1000
+PURGE_DELAY_OPTION = "purge_delay"
+
+_ELF_MAGIC = b"\x7fELF"
+_SHT_SYMTAB = 2
+_SHT_DYNSYM = 11
+_PT_LOAD = 1
+_STT_OBJECT = 1
+_STT_FUNC = 2
+_SHN_UNDEF = 0
+# mi_option_desc_t: long value; mi_init_t init; mi_option_t option;
+# const char* name; const char* legacy_name.
+_DESC = struct.Struct("<qiiQQ")
+_INIT_STATES = (0, 1, 2)  # UNINIT, DEFAULTED, INITIALIZED
+_NEEDED = ("mi_option_get", "mi_option_set", "options")
+
+
+class HostPolicyUnavailable(RuntimeError):
+    """The mimalloc option table could not be located or failed a check."""
+
+
+@dataclass(frozen=True)
+class _Symbol:
+    value: int
+    size: int
+    kind: int
+
+
+def _read_symbols(path: str) -> Tuple[Dict[str, List[_Symbol]], Dict[str, _Symbol], int]:
+    """Return the symbols this module needs and the link address of file offset 0.
+
+    (``.symtab`` symbols named in ``_NEEDED``, one defined ``.dynsym`` function,
+    ``p_vaddr - p_offset`` of the first ``PT_LOAD``). The ``.dynsym`` function
+    anchors the load address of the library.
+    """
+    with open(path, "rb") as f:
+        ident = f.read(16)
+        if len(ident) < 16 or ident[:4] != _ELF_MAGIC or ident[4] != 2 or ident[5] != 1:
+            raise HostPolicyUnavailable(f"{path}: not a 64-bit little-endian ELF file")
+        f.seek(0)
+        hdr = struct.unpack("<16sHHIQQQIHHHHHH", f.read(64))
+        e_phoff, e_shoff, e_phentsize, e_phnum = hdr[5], hdr[6], hdr[9], hdr[10]
+        e_shentsize, e_shnum = hdr[11], hdr[12]
+        if e_phoff == 0 or e_phentsize != 56 or e_phnum == 0:
+            raise HostPolicyUnavailable(f"{path}: no program headers")
+        f.seek(e_phoff)
+        loads = [
+            ph
+            for ph in (struct.unpack("<IIQQQQQQ", f.read(56)) for _ in range(e_phnum))
+            if ph[0] == _PT_LOAD
+        ]
+        if not loads:
+            raise HostPolicyUnavailable(f"{path}: no PT_LOAD segment")
+        first_load = min(loads, key=lambda ph: ph[2])
+        link_offset0 = first_load[3] - first_load[2]
+        if e_shoff == 0 or e_shentsize != 64 or e_shnum == 0:
+            raise HostPolicyUnavailable(f"{path}: no section headers")
+        f.seek(e_shoff)
+        sections = [struct.unpack("<IIQQQQIIQQ", f.read(64)) for _ in range(e_shnum)]
+
+        def read_section(index):
+            sec = sections[index]
+            f.seek(sec[4])
+            return f.read(sec[5])
+
+        wanted: Dict[str, List[_Symbol]] = {}
+        anchor: Dict[str, _Symbol] = {}
+        for index, sec in enumerate(sections):
+            sh_type, sh_link, sh_entsize = sec[1], sec[6], sec[9]
+            if sh_type not in (_SHT_SYMTAB, _SHT_DYNSYM) or sh_entsize != 24:
+                continue
+            data = read_section(index)
+            strtab = read_section(sh_link)
+            for off in range(0, len(data) - 23, 24):
+                st_name, st_info, _, st_shndx, st_value, st_size = struct.unpack_from(
+                    "<IBBHQQ", data, off
+                )
+                if st_shndx == _SHN_UNDEF or st_name == 0:
+                    continue
+                end = strtab.find(b"\0", st_name)
+                name = strtab[st_name:end].decode("ascii", "replace")
+                kind = st_info & 0xF
+                if sh_type == _SHT_SYMTAB and name in _NEEDED:
+                    wanted.setdefault(name, []).append(_Symbol(st_value, st_size, kind))
+                elif (
+                    sh_type == _SHT_DYNSYM
+                    and not anchor
+                    and kind == _STT_FUNC
+                    and st_value
+                    and st_size
+                ):
+                    anchor[name] = _Symbol(st_value, st_size, kind)
+    return wanted, anchor, link_offset0
+
+
+def _unique(symbols: Dict[str, List[_Symbol]], name: str, kind: int) -> _Symbol:
+    found = {s for s in symbols.get(name, ()) if s.kind == kind and s.value}
+    if len(found) != 1:
+        raise HostPolicyUnavailable(f"symbol {name!r}: expected one definition, found {len(found)}")
+    return found.pop()
+
+
+def _mapped_base(path: str) -> Optional[int]:
+    """Start of the first mapping of ``path`` at file offset 0 in this process."""
+    real = os.path.realpath(path)
+    try:
+        with open("/proc/self/maps") as f:
+            for line in f:
+                parts = line.split()
+                if (
+                    len(parts) >= 6
+                    and parts[2] == "00000000"
+                    and os.path.realpath(parts[5]) == real
+                ):
+                    return int(parts[0].split("-")[0], 16)
+    except OSError:
+        return None
+    return None
+
+
+class MimallocOptions:
+    """mimalloc's runtime options in an already-loaded library, by option name."""
+
+    def __init__(self, path: str):
+        if not os.path.isfile(path):
+            raise HostPolicyUnavailable(f"{path}: not found")
+        symbols, anchor, link_offset0 = _read_symbols(path)
+        missing = [n for n in _NEEDED if n not in symbols]
+        if missing:
+            raise HostPolicyUnavailable(f"{path}: no mimalloc symbols {missing}")
+        if not anchor:
+            raise HostPolicyUnavailable(f"{path}: no exported function to anchor the load address")
+        get_sym = _unique(symbols, "mi_option_get", _STT_FUNC)
+        set_sym = _unique(symbols, "mi_option_set", _STT_FUNC)
+        table = _unique(symbols, "options", _STT_OBJECT)
+        if table.size == 0 or table.size % _DESC.size:
+            raise HostPolicyUnavailable(
+                f"option table of {table.size} bytes is not a whole number of "
+                f"{_DESC.size}-byte descriptors"
+            )
+        try:
+            lib = ctypes.CDLL(path, mode=os.RTLD_NOLOAD | os.RTLD_LAZY)
+        except OSError as exc:
+            raise HostPolicyUnavailable(f"{path}: not loaded in this process ({exc})") from exc
+        ((anchor_name, anchor_sym),) = anchor.items()
+        bias = ctypes.cast(lib[anchor_name], ctypes.c_void_p).value - anchor_sym.value
+        mapped = _mapped_base(path)
+        if mapped is not None and mapped != bias + link_offset0:
+            raise HostPolicyUnavailable(
+                f"load address mismatch: anchor bias {bias:#x} + {link_offset0:#x}, "
+                f"/proc/self/maps {mapped:#x}"
+            )
+        self.path = path
+        self._lib = lib  # keeps the handle (and the library) alive
+        self._table = bias + table.value
+        self._count = table.size // _DESC.size
+        self._get = ctypes.CFUNCTYPE(ctypes.c_long, ctypes.c_int)(bias + get_sym.value)
+        self._set = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_long)(bias + set_sym.value)
+        self.names = self._check_table()
+
+    def _descriptor(self, index: int) -> Tuple[int, int, int, int, int]:
+        raw = ctypes.string_at(self._table + index * _DESC.size, _DESC.size)
+        return _DESC.unpack(raw)
+
+    def _check_table(self) -> List[str]:
+        names = []
+        for index in range(self._count):
+            _, init, option, name_ptr, _ = self._descriptor(index)
+            if option != index or init not in _INIT_STATES or not name_ptr:
+                raise HostPolicyUnavailable(
+                    f"option descriptor {index}: option={option} init={init}; not a "
+                    f"mimalloc option table"
+                )
+            name = ctypes.c_char_p(name_ptr).value or b""
+            if not 0 < len(name) <= 64 or not all(
+                c == 0x5F or 0x61 <= c <= 0x7A or 0x30 <= c <= 0x39 for c in name
+            ):
+                raise HostPolicyUnavailable(f"option descriptor {index}: bad name {name!r}")
+            names.append(name.decode())
+        return names
+
+    def index(self, name: str) -> int:
+        if self.names.count(name) != 1:
+            raise HostPolicyUnavailable(f"option {name!r} not found exactly once")
+        return self.names.index(name)
+
+    def get(self, name: str) -> int:
+        index = self.index(name)
+        value = self._get(index)
+        stored = self._descriptor(index)[0]
+        if value != stored:
+            raise HostPolicyUnavailable(
+                f"mi_option_get({name}) = {value} but the table holds {stored}"
+            )
+        return value
+
+    def set(self, name: str, value: int) -> None:
+        self._set(self.index(name), value)
+        if self.get(name) != value:
+            raise HostPolicyUnavailable(f"mi_option_set({name}, {value}) did not take effect")
+
+
+def torch_libc10_path() -> str:
+    import torch
+
+    return os.path.join(os.path.dirname(torch.__file__), "lib", "libc10.so")
+
+
+def apply_purge_policy(
+    path: Optional[str] = None, purge_delay_ms: int = PURGE_DELAY_MS
+) -> Tuple[int, int]:
+    """Set mimalloc's purge delay in this process; return (previous, new) in ms.
+
+    Raises:
+        HostPolicyUnavailable: before changing anything, if any check fails.
+    """
+    options = MimallocOptions(path or torch_libc10_path())
+    previous = options.get(PURGE_DELAY_OPTION)
+    if previous != purge_delay_ms:
+        options.set(PURGE_DELAY_OPTION, purge_delay_ms)
+    return previous, purge_delay_ms
+
+
+def install_host_policy(where: str = "after_warmup") -> bool:
+    """Entry point of the ``host_page_retention`` package (``host_policy``); returns True when engaged.
+
+    Called once per rank by the executor after warmup. Logs the decision once
+    (``[glm_moe_dsa] host_page_retention/host_policy``); on any failed check the allocator is left
+    as it was.
+    """
+    from ...dispatch import log_capability
+
+    try:
+        previous, new = apply_purge_policy()
+    except HostPolicyUnavailable as exc:
+        return log_capability("host_page_retention", "host_policy", False, f"{where}: {exc}")
+    return log_capability(
+        "host_page_retention",
+        "host_policy",
+        True,
+        f"{where}: mimalloc {PURGE_DELAY_OPTION} {previous} -> {new} ms in libc10.so",
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/__init__.py
new file mode 100644
index 000000000000..3064d739edb4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/__init__.py
@@ -0,0 +1,55 @@
+# 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.
+"""indexer_layer_front: the full-layer query front.
+
+Requires ``absorbed_query``, ``indexer_topk``, ``draft_absorbed_query``, ``rope_kv_append`` and
+``cross_rank_topk`` (the selector closure). Three forks, each byte-identical to the package it
+forks:
+
+* ``query_proj.dsa_proj``: absorbed_query's / draft_absorbed_query's full-layer query kernel
+  (``kernels.front_qproj_indexer_kernel``) triggering its dependents after
+  each program's last global load and reading positions after its wait.
+  By default (``IndexerFrontConfig.norm="native"``) it takes the native q_a RMSNorm
+  output as absorbed_query does; the options ``inline`` / ``programs`` normalize q_a
+  inside the kernel instead (no native RMSNorm launch);
+* ``rope.mla_rope_generation``: rope_kv_append's generation rope in which each thread
+  triggers the dependents right after the global reads of its work item
+  (``trtllm::glm_moe_dsa_indexer_layer_rope_kv_append``);
+* ``cp_indexer_decode.try_cp_indexer_decode``: the cross_rank_topk / cross_rank_topk_cluster_merge
+  GEN indexer region with cross_rank_topk's CP logits kernel waiting for its programmatic primary at
+  entry and launched with PDL (``fp8_paged_mqa_logits_cp``).
+
+Registered entry points (``dispatch.get_impl("indexer_layer_front", name)``): ``dsa_proj``,
+``rope_available``, ``mla_rope_generation``, ``cp_indexer_decode``.
+"""
+
+from ...dispatch import register
+from .cp_indexer_decode import try_cp_indexer_decode
+from .query_proj import IndexerFrontConfig, IndexerFrontQueryProjection, dsa_proj
+from .rope import mla_rope_generation, rope_available
+
+register("indexer_layer_front", "dsa_proj", dsa_proj)
+register("indexer_layer_front", "rope_available", rope_available)
+register("indexer_layer_front", "mla_rope_generation", mla_rope_generation)
+register("indexer_layer_front", "cp_indexer_decode", try_cp_indexer_decode)
+
+__all__ = [
+    "IndexerFrontConfig",
+    "IndexerFrontQueryProjection",
+    "dsa_proj",
+    "mla_rope_generation",
+    "rope_available",
+    "try_cp_indexer_decode",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_indexer_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_indexer_decode.py
new file mode 100644
index 000000000000..3a0f4fecf25a
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_indexer_decode.py
@@ -0,0 +1,248 @@
+# 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.
+"""indexer_layer_front: the cross_rank_topk / cross_rank_topk_cluster_merge GEN indexer region with
+the indexer_layer_front logits kernel.
+
+Same region, policy, workspace, top-k op and output contract as cross_rank_topk_cluster_merge
+(``cross_rank_topk_cluster_merge/cp_indexer_decode.py``, with ``cross_rank_topk_cluster_merge``
+selected) or cross_rank_topk (``indexer_topk/cross_rank_indexer_decode.py``): every decision input
+and every decline is theirs, evaluated from the same host state before any device work. The only
+difference is the logits launch: ``indexer_layer_front/cp_logits_runner.py`` (wait at entry,
+programmatic dependent of the generation rope), whose logits are
+bit-identical to the context-partitioned logits kernel's. The top-k
+(``glm_moe_dsa_cross_rank_topk_cluster_merge`` or ``glm_moe_dsa_cross_rank_topk``)
+then runs on them unchanged.
+
+A ``False`` return leaves no side effect; the adapter then tries cross_rank_topk_cluster_merge's and
+cross_rank_topk's own entries, which reach the same decision.
+"""
+
+from typing import Dict
+
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.logger import logger
+
+from ...dispatch import is_enabled, log_capability
+from ..indexer_topk import cross_rank_policy
+
+__all__ = ["CP_DECISION_COUNTS", "cp_decode_topk", "try_cp_indexer_decode"]
+
+# As cross_rank_topk / cross_rank_topk_cluster_merge: MTP draft step 0 runs the same region on the
+# MTP layer.
+ALLOW_DRAFT_STEP = True
+
+CP_DECISION_COUNTS: Dict[str, int] = {}
+
+
+def _record_decision(flavor: str, decision) -> None:
+    key = f"{flavor}:" + ("engaged" if decision.eligible else decision.reason)
+    count = CP_DECISION_COUNTS.get(key, 0)
+    CP_DECISION_COUNTS[key] = count + 1
+    if count == 0:
+        state = "engaged" if decision.eligible else f"declined ({decision.reason})"
+        logger.info(
+            f"[glm_moe_dsa] indexer_layer_front/cp_indexer_decode[{flavor}]: first '{state}' decision on this rank"
+        )
+
+
+def _modules(flavor: str):
+    """``(workspace module, top-k entry)`` of the region owner."""
+    if flavor == "cross_rank_topk_cluster_merge":
+        from ..cross_rank_topk_cluster_merge import cp_indexer_decode as cluster_merge_decode
+        from ..cross_rank_topk_cluster_merge import cp_workspace as cluster_merge_ws
+
+        return cluster_merge_ws, cluster_merge_decode.cp_topk_from_local_logits
+    from ..indexer_topk import cross_rank_indexer_decode as cross_rank_decode
+    from ..indexer_topk import cross_rank_workspace as cross_rank_ws
+
+    return cross_rank_ws, cross_rank_decode.cp_topk_from_local_logits
+
+
+def cp_decode_topk(
+    *,
+    q: torch.Tensor,
+    weights: torch.Tensor,
+    k_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    kv_lens: torch.Tensor,
+    max_seq_len: int,
+    out: torch.Tensor,
+    workspace,
+    flavor: str,
+) -> torch.Tensor:
+    """The region for one request (arguments as cross_rank_topk's ``cp_decode_topk``)."""
+    from .cp_logits_runner import paged_mqa_logits_cp
+
+    local_logits = paged_mqa_logits_cp(
+        q, k_cache, weights, kv_lens, block_table, max_seq_len, workspace.rank, workspace.world
+    )
+    _, topk = _modules(flavor)
+    return topk(local_logits, kv_lens, out, workspace, q.shape[1])
+
+
+def _try(flavor: str, indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out) -> bool:
+    from ..indexer_topk.cross_rank_logits_runner import local_capacity
+    from .cp_logits_runner import ensure_compiled
+
+    ws_mod, _ = _modules(flavor)
+    num_contexts = metadata.num_contexts
+    num_generations = metadata.num_generations
+    gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
+    uniform = num_generations > 0 and bool(
+        (gen_seq_lens == num_gen_tokens // num_generations).all()
+    )
+    top_k = indexer.top_k
+    k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(indexer.layer_idx)
+    block_table = metadata.indexer_k_cache_block_offsets
+    mapping = getattr(metadata, "mapping", None)
+    max_seq_len = metadata.get_indexer_max_seq_len()
+    world = mapping.tp_size if mapping is not None else 1
+    info = ws_mod.op_info(world) if world > 0 else None
+    capacity = local_capacity(max_seq_len, world) if world > 0 else 0
+    max_cap = (
+        (
+            info[ws_mod.INFO_MAX_LOCAL_CAPACITY]
+            if flavor == "cross_rank_topk_cluster_merge"
+            else info[7]
+        )
+        if info
+        else 0
+    )
+    local_reason = cross_rank_policy.rank_local_reason(
+        op_available=info is not None,
+        local_capacity=capacity,
+        max_local_capacity=max_cap,
+        sm_version=get_sm_version(),
+        cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE,
+    )
+    if flavor == "cross_rank_topk_cluster_merge" and not local_reason and info is not None:
+        local_reason = ws_mod.cluster_reason(info)
+    decision = cross_rank_policy.decide_cp_uniform(
+        world_size=world,
+        local_capacity=capacity,
+        allow_draft_step=ALLOW_DRAFT_STEP,
+        native_uses_cute_dsl_logits=indexer.use_cute_dsl_paged_mqa_logits,
+        use_fp4_indexer=indexer.use_fp4,
+        use_gvr_emission=indexer.use_gvr_emission,
+        needs_gvr_prior=top_k.needs_gvr_prior,
+        is_ragged_verify=metadata.is_ragged_verify,
+        expanded_mtp_buffers=metadata.use_expanded_buffers_for_mtp,
+        in_mtp_draft_loop=metadata.in_mtp_draft_loop,
+        attention_dp=bool(mapping is not None and mapping.enable_attention_dp),
+        topk_decode_impl=top_k.decode_implementation.value,
+        topk_self_sampling=top_k.gvr_self_sampling,
+        index_topk=top_k.top_k,
+        compress_ratio=top_k.compress_ratio,
+        num_heads=indexer.n_heads,
+        head_dim=indexer.head_dim,
+        query_dtype=str(q_fp8.dtype),
+        query_shape=tuple(q_fp8.shape),
+        query_contiguous=q_fp8.is_contiguous(),
+        weights_dtype=str(weights.dtype),
+        weights_shape=tuple(weights.shape),
+        weights_contiguous=weights.is_contiguous(),
+        k_cache_dtype=str(k_cache.dtype),
+        k_cache_shape=tuple(k_cache.shape),
+        block_table_cols=block_table.shape[1],
+        max_seq_len=max_seq_len,
+        num_contexts=num_contexts,
+        token_offset=token_offset,
+        num_gen_tokens=num_gen_tokens,
+        num_generations=num_generations,
+        uniform_verify_width=uniform,
+    )
+    if flavor == "cross_rank_topk_cluster_merge" and decision.eligible:
+        from ..cross_rank_topk_cluster_merge.cp_indexer_decode import world_size_reason
+
+        why = world_size_reason(world)
+        if why:
+            decision = type(decision)(False, why)
+    if decision.eligible:
+        # indexer_layer_front's logits kernel compiles before the group agrees on the
+        # workspace.
+        workspace = ws_mod.get_workspace(
+            mapping,
+            prepare=lambda: ensure_compiled(
+                q_fp8.device,
+                k_cache.shape[1],
+                indexer.n_heads,
+                indexer.head_dim,
+                num_gen_tokens // num_generations,
+                world,
+            ),
+            local_reason=local_reason,
+        )
+        if workspace is None:
+            reason = (
+                ws_mod.failure_reason(mapping) or "workspace not created before CUDA-graph capture"
+            )
+            decision = type(decision)(False, f"workspace unavailable: {reason}")
+    log_capability(
+        "indexer_layer_front", f"cp_indexer_decode_{flavor}", decision.eligible, decision.reason
+    )
+    _record_decision(flavor, decision)
+    if not decision.eligible:
+        return False
+
+    next_n = num_gen_tokens // num_generations
+    gen_rows = slice(token_offset, token_offset + num_gen_tokens)
+    gen_reqs = slice(num_contexts, num_contexts + num_generations)
+    kv_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+    assert kv_lens is not None
+    cp_decode_topk(
+        q=q_fp8[gen_rows].view(num_generations, next_n, indexer.n_heads, indexer.head_dim),
+        weights=weights[gen_rows],
+        k_cache=k_cache,
+        block_table=block_table[gen_reqs],
+        kv_lens=kv_lens,
+        max_seq_len=max_seq_len,
+        out=out,
+        workspace=workspace,
+        flavor=flavor,
+    )
+    return True
+
+
+def try_cp_indexer_decode(
+    indexer,
+    metadata,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+) -> bool:
+    """Adapter entry: the GEN region of ``Indexer.sparse_attn_indexer`` with the indexer_layer_front logits kernel.
+
+    Tries cross_rank_topk_cluster_merge's region (when ``cross_rank_topk_cluster_merge`` is
+    selected), then cross_rank_topk's, each with its own decision; returns False (no side effect)
+    when both decline.
+
+    Args/Returns: as ``indexer_topk.indexer_decode.try_indexer_decode``.
+    """
+    if not is_enabled("cross_rank_topk"):
+        return False
+    flavors = (
+        ("cross_rank_topk_cluster_merge", "cross_rank_topk")
+        if is_enabled("cross_rank_topk_cluster_merge")
+        else ("cross_rank_topk",)
+    )
+    for flavor in flavors:
+        if _try(flavor, indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out):
+            return True
+    return False
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_logits_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_logits_runner.py
new file mode 100644
index 000000000000..574b00a7dcb1
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/cp_logits_runner.py
@@ -0,0 +1,217 @@
+# 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.
+"""Compile cache and launcher of the indexer_layer_front context-partitioned logits kernel.
+
+Same configuration, arguments, output arena and values as cross_rank_topk's
+``indexer_topk/cross_rank_logits_runner.py``; the kernel
+(``indexer_layer_front/fp8_paged_mqa_logits_cp.py``) waits for its programmatic primary before any
+global access and is launched with programmatic dependent launch when ``TRTLLM_ENABLE_PDL`` is on
+(it then starts during the generation rope, which triggers right after its last read).
+"""
+
+from typing import Dict, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+
+from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers
+
+from ..indexer_topk.cross_rank_logits_runner import (
+    COMPUTE_BLOCK_KV,
+    LOGITS_BUFFER_NAME,
+    NUM_W_IN_REG,
+    SPLIT_KV,
+    TRIGGER_DEPENDENTS,
+    local_capacity,
+)
+from .fp8_paged_mqa_logits_cp import FP8MQALogitsCPZ14Kernel
+
+_KERNEL_CACHE: Dict[Tuple, object] = {}
+_NUM_SMS: Dict[int, int] = {}
+_UNUSED_SCHEDULE: Dict[int, torch.Tensor] = {}
+
+
+def use_pdl() -> bool:
+    """Whether the kernel is launched as a programmatic dependent (TensorRT-LLM's PDL switch)."""
+    from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+    return bool(get_env_enable_pdl())
+
+
+def _num_sms(device: torch.device) -> int:
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    if index not in _NUM_SMS:
+        _NUM_SMS[index] = torch.cuda.get_device_properties(index).multi_processor_count
+    return _NUM_SMS[index]
+
+
+def _unused_schedule(device: torch.device) -> torch.Tensor:
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    t = _UNUSED_SCHEDULE.get(index)
+    if t is None:
+        t = torch.zeros(1, 2, dtype=torch.int32, device=device)
+        _UNUSED_SCHEDULE[index] = t
+    return t
+
+
+def _compile(
+    phys_block_kv: int,
+    num_heads: int,
+    head_dim: int,
+    next_n: int,
+    num_sms: int,
+    world: int,
+    wait_at_entry: bool,
+    pdl: bool,
+):
+    n_cols = next_n * num_heads
+    block_bytes = phys_block_kv * (head_dim + 4)
+    sym_num_phys_blocks = cute.sym_int()
+    sym_b = cute.sym_int()
+    max_ctx = cute.sym_int()
+    max_blocks_per_seq = cute.sym_int()
+    num_ctas = cute.sym_int()
+    kv_fake = cute.runtime.make_fake_tensor(
+        cutlass.Uint8, (sym_num_phys_blocks, block_bytes), stride=(cute.sym_int64(), 1)
+    )
+    q_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Uint8, (n_cols, head_dim, sym_b), stride_order=(1, 0, 2)
+    )
+    w_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float32, (n_cols, sym_b), stride_order=(0, 1)
+    )
+    logits_fake = cute.runtime.make_fake_tensor(
+        cutlass.Float32, (cute.sym_int(), max_ctx), stride=(cute.sym_int64(), 1)
+    )
+    bt_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Int32, (sym_b, max_blocks_per_seq), stride_order=(1, 0)
+    )
+    cl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,))
+    sm_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Int32, (num_ctas, 2), stride_order=(1, 0)
+    )
+    fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    kernel = FP8MQALogitsCPZ14Kernel(
+        block_kv=COMPUTE_BLOCK_KV,
+        phys_block_kv=phys_block_kv,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        next_n=next_n,
+        num_sms=num_sms,
+        num_w_in_reg=NUM_W_IN_REG,
+        world=world,
+        trigger_dependents=TRIGGER_DEPENDENTS,
+        wait_at_entry=wait_at_entry,
+        use_pdl=pdl,
+    )
+    return cute.compile(
+        kernel,
+        kv_fake,
+        q_fake,
+        w_fake,
+        logits_fake,
+        bt_fake,
+        cl_fake,
+        sm_fake,
+        cutlass.Int32(1),
+        cutlass.Int32(1),
+        cutlass.Int32(0),
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+def _key(
+    device: torch.device,
+    phys_block_kv: int,
+    num_heads: int,
+    head_dim: int,
+    next_n: int,
+    world: int,
+    wait_at_entry: bool,
+    pdl: bool,
+) -> Tuple:
+    return (phys_block_kv, num_heads, head_dim, next_n, _num_sms(device), world, wait_at_entry, pdl)
+
+
+def ensure_compiled(
+    device: torch.device,
+    phys_block_kv: int,
+    num_heads: int,
+    head_dim: int,
+    next_n: int,
+    world: int,
+    wait_at_entry: bool = True,
+    pdl: bool = None,
+) -> None:
+    """Compile the kernel for this shape now (host only, no device work)."""
+    pdl = use_pdl() if pdl is None else pdl
+    key = _key(device, phys_block_kv, num_heads, head_dim, next_n, world, wait_at_entry, pdl)
+    if key not in _KERNEL_CACHE:
+        _KERNEL_CACHE[key] = _compile(*key)
+
+
+def paged_mqa_logits_cp(
+    q: torch.Tensor,
+    kv_fused: torch.Tensor,
+    weights: torch.Tensor,
+    kv_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    max_seq_len: int,
+    rank: int,
+    world: int,
+    wait_at_entry: bool = True,
+    pdl: bool = None,
+) -> torch.Tensor:
+    """cross_rank_topk's ``paged_mqa_logits_cp`` with the indexer_layer_front kernel (same arguments, arena and values).
+
+    ``wait_at_entry`` / ``pdl`` default to the served indexer_layer_front configuration (wait
+    at entry, PDL launch when TensorRT-LLM's PDL switch is on); tests pass
+    other values.
+    """
+    batch, next_n, num_heads, head_dim = q.shape
+    assert batch == 1
+    pdl = use_pdl() if pdl is None else pdl
+    phys_block_kv = kv_fused.shape[1]
+    num_phys_blocks = kv_fused.shape[0]
+    width = local_capacity(max_seq_len, world)
+    aligned = (width + SPLIT_KV - 1) // SPLIT_KV * SPLIT_KV
+    logits = get_memory_buffers().get_buffer(
+        [next_n, aligned],
+        torch.float32,
+        buffer_name=LOGITS_BUFFER_NAME,
+        reserve_buffer=torch.cuda.is_current_stream_capturing(),
+    )[:, :width]
+    key = _key(q.device, phys_block_kv, num_heads, head_dim, next_n, world, wait_at_entry, pdl)
+    compiled = _KERNEL_CACHE.get(key)
+    if compiled is None:
+        compiled = _compile(*key)
+        _KERNEL_CACHE[key] = compiled
+    q_3d = q.reshape(batch, next_n * num_heads, head_dim).permute(1, 2, 0)
+    w_2d = weights.reshape(batch, next_n * num_heads).t()
+    compiled(
+        kv_fused.reshape(num_phys_blocks, -1),
+        q_3d.view(torch.uint8),
+        w_2d,
+        logits,
+        block_table,
+        kv_lens,
+        _unused_schedule(q.device),
+        num_phys_blocks,
+        batch,
+        rank,
+    )
+    return logits
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py
new file mode 100644
index 000000000000..bdf3d8013d38
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py
@@ -0,0 +1,1844 @@
+# 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.
+"""
+indexer_layer_front fork of cross_rank_topk's context-partitioned FP8 paged MQA logits kernel
+(Blackwell).
+
+Identical to ``indexer_topk/fp8_paged_mqa_logits_cross_rank.py`` (cross_rank_topk) in every value it
+writes; only the memory-order prologue differs:
+
+- The kernel is launched with programmatic dependent launch (``use_pdl``) as
+  the dependent of the generation rope, and with ``wait_at_entry`` every thread
+  executes ``griddepcontrol.wait`` before its first global-memory access.
+- Before the wait it only triggers its own dependents (as cross_rank_topk), derives the
+  warp roles, prefetches the TMA descriptors (kernel parameters), carves the
+  shared memory, initializes the pipeline mbarriers and allocates TMEM (math
+  warp 0). The entry load ``global_kv = mContextLens[0]`` and the per-CTA
+  schedule derived from it, the block-table reads, every TMA load (Q, weights,
+  KV pages, scales) and every logits store come after the wait.
+
+TMEM before the wait: the only grid this kernel can wait on while it holds
+TMEM is the rope (no TMEM); every earlier grid of the chain has completed
+because the rope triggers only after its own wait. Without ``wait_at_entry``
+the kernel computes the same values without a wait (test reference only).
+
+cross_rank_topk notes follow.
+
+cross_rank_topk context-partitioned variant of the indexer_topk FP8 paged MQA logits kernel
+(Blackwell).
+
+Identical to ``fp8_paged_mqa_logits_rows`` (the shipped indexer_topk all-rows kernel)
+except for which KV tokens one launch scores and how CTAs find their work:
+
+- The request's indexer context is partitioned over ``world`` ranks by whole
+  64-token pages (global page ``p`` belongs to rank ``p % world``). A launch
+  scores only the pages of its runtime ``rank``: local page ``lp`` reads the
+  block-table entry of global page ``lp * world + rank``, and logits column
+  ``c`` holds global token ``((c // 64) * world + rank) * 64 + c % 64``.
+- One request (batch size 1). The local length is derived in-kernel from the
+  global indexer KV length (``context_lens[0]``), and each CTA computes its
+  contiguous share of the local tiles itself; ``schedule_meta`` is not read.
+
+Per-element arithmetic is unchanged, so every local logit is bit-identical to
+the native logit of the same global token.
+
+Original indexer_topk notes follow.
+
+indexer_topk variant of the CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100).
+
+Derived from ``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
+for the GLM-5.x indexer shape, where all verify rows of a request fit one UMMA
+tile (``N = next_n * num_heads <= 256``). The only functional change is the
+size of the per-thread register cache of epilogue weights
+(``num_w_in_reg`` heads per verify row; the remaining heads are read from
+shared memory). Keeping that cache bounded avoids register spills when
+``N`` is large; weight values, accumulation order and output are unchanged.
+
+Architecture:
+  - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA)
+  - 1 TMA per KV block [128, 128], UMMA iterates 4x K=32
+  - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2)
+  - Q reloaded via TMA pipeline when q_idx (batch) changes
+  - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs
+  - Weights cached in registers: preloaded once per q_idx change (not per KV block)
+  - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes)
+
+Merged KV+Scale pipeline:
+  - KV data and scales share a single TMA barrier per group
+  - TMA loads both KV and Scale under one barrier (combined tx_count)
+  - UMMA waits on merged barrier (for KV GEMM), does NOT release
+  - Math waits on merged barrier (for scale read), Math releases
+
+Fused KV layout:
+  - KV data and scales stored contiguously per physical block:
+    [num_phys_blocks, block_kv * (head_dim + 4)] bytes
+  - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)]
+  - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic
+
+Scheduler:
+  - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups)
+  - schedule_meta[sm_idx+1] = end boundary for this CTA
+  - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx)
+  - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step
+
+Dynamic shape support:
+  - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen
+  - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq,
+    num_ctas) are marked dynamic via mark_compact_shape_dynamic
+  - Allows JIT cache reuse across different batch sizes / sequence lengths
+
+Epilogue dtype flows (--acc_dtype / --epi_dtype):
+
+  Flow 1: --acc_dtype fp32 --epi_dtype fp16
+    Q(FP8) x K(FP8) -> MMA acc(FP32) -> TMEM(FP32)
+      -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16)
+      -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16)
+      -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype)
+
+  Flow 2: --acc_dtype fp16 --epi_dtype fp16
+    Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b)
+      -> LDTM -> Reg(FP16) -> ReLU(FP16)
+      -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16)
+      -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype)
+
+  --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion.
+  Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32
+"""
+
+from typing import Tuple
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass import Float16, Int32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import griddepcontrol_launch_dependents
+
+# CuTe DSL CUDA 13 validates rounding modes as string literals. The string
+# form is also accepted by older wrappers, so keep it version-independent.
+_RND_RN = "rn"
+
+
+@dsl_user_op
+def pack_f16x2(
+    a: Float16,
+    b: Float16,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    f16_ty = Float16.mlir_type
+    i32_ty = Int32.mlir_type
+    vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc)
+    v = vector.from_elements(
+        vec2_f16,
+        (Float16(a).ir_value(loc=loc, ip=ip), Float16(b).ir_value(loc=loc, ip=ip)),
+        loc=loc,
+        ip=ip,
+    )
+    return Int32(llvm.bitcast(i32_ty, v, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def unpack_f16x2(
+    packed: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Tuple[Float16, Float16]:
+    f16_ty = Float16.mlir_type
+    vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc)
+    v = llvm.bitcast(vec2_f16, Int32(packed).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    r0 = Float16(vector.extract(v, dynamic_position=[], static_position=[0], loc=loc, ip=ip))
+    r1 = Float16(vector.extract(v, dynamic_position=[], static_position=[1], loc=loc, ip=ip))
+    return r0, r1
+
+
+@dsl_user_op
+def fma_f16x2(
+    a: Int32,
+    b: Int32,
+    c: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [
+                Int32(a).ir_value(loc=loc, ip=ip),
+                Int32(b).ir_value(loc=loc, ip=ip),
+                Int32(c).ir_value(loc=loc, ip=ip),
+            ],
+            "fma.rn.f16x2 $0, $1, $2, $3;",
+            "=r,r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def max_f16x2(
+    a: Int32,
+    b: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "max.f16x2 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def add_f16x2(
+    a: Int32,
+    b: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "add.f16x2 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+class FP8MQALogitsCPZ14Kernel:
+    """FP8 paged MQA logits kernel for Blackwell (SM100).
+
+    Each CTA processes a range of (q_idx, kv_split) pairs.
+    A split = 2 consecutive KV blocks within a sequence (one per warp group).
+    Q is shared between warp groups and reloaded when q_idx changes.
+    """
+
+    def __init__(
+        self,
+        block_kv: int = 128,
+        phys_block_kv: int = 128,
+        num_heads: int = 64,
+        head_dim: int = 128,
+        next_n: int = 1,
+        num_sms: int = 148,
+        remove_kv_wait_in_epilogue: bool = False,
+        early_tmem_copy: bool = False,
+        smem_subpartition_opt: bool = False,
+        max_kv_pipeline: bool = False,
+        max_umma_pipeline: bool = False,
+        num_epi_subtiles: int = 1,
+        epi_dtype=cutlass.Float32,
+        acc_dtype=cutlass.Float32,
+        output_dtype=cutlass.Float32,
+        num_w_in_reg=None,
+        world: int = 8,
+        trigger_dependents: bool = True,
+        wait_at_entry: bool = True,
+        use_pdl: bool = True,
+    ):
+        if block_kv % phys_block_kv != 0 or phys_block_kv != 64:
+            raise ValueError("the context-partitioned kernel needs 64-token pages")
+        self.world = world
+        # Let the dependent grid (cross_rank_topk's select/exchange op, launched with programmatic
+        # stream serialization) launch while this kernel runs: the op waits on griddepcontrol.wait
+        # before it reads anything, so only its launch latency overlaps this kernel.
+        self.trigger_dependents = trigger_dependents
+        # indexer_layer_front: wait for the programmatic primary (the rope) before any global
+        # access, and launch as its programmatic dependent.
+        self.wait_at_entry = wait_at_entry
+        self.use_pdl = use_pdl
+        self.block_kv = block_kv
+        self.phys_block_kv = phys_block_kv
+        self.num_blocks_per_mma = block_kv // phys_block_kv
+        assert block_kv % phys_block_kv == 0, (
+            f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}"
+        )
+        assert self.num_blocks_per_mma <= 4, (
+            f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4"
+        )
+        self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue
+        self.early_tmem_copy = early_tmem_copy
+        self.smem_subpartition_opt = smem_subpartition_opt
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.next_n = next_n
+        self.N = next_n * num_heads
+        self.num_sms = num_sms
+        self.num_epi_subtiles = num_epi_subtiles
+        self.epi_dtype = epi_dtype
+        self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4
+        # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
+        # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1
+        # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest).
+        w_stage_bytes = self.N * self.epi_bytes
+        self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes
+        self.output_dtype = output_dtype
+        if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0:
+            raise ValueError("num_heads must be divisible by num_epi_subtiles")
+        if (num_heads // num_epi_subtiles) % 4 != 0:
+            raise ValueError(
+                "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)"
+            )
+        self.num_groups = 2
+
+        # Register-cached epilogue weights per verify row. None keeps the
+        # upstream heuristic; otherwise a multiple of 4 (FMA unroll) <= heads.
+        if num_w_in_reg is None:
+            if epi_dtype == cutlass.Float16:
+                max_w_in_reg = 64 if next_n <= 3 else 48
+            else:
+                max_w_in_reg = 64 if next_n == 1 else 40 if next_n >= 4 else 52
+            num_w_in_reg = min(max_w_in_reg, num_heads)
+        if num_w_in_reg < 0 or num_w_in_reg > num_heads or num_w_in_reg % 4 != 0:
+            raise ValueError(
+                f"num_w_in_reg={num_w_in_reg} must be a multiple of 4 in [0, {num_heads}]"
+            )
+        self.num_w_in_reg = num_w_in_reg
+
+        self.num_math_threads = 256
+        self.num_specialized_threads = 128
+        self.threads_per_cta = 384
+        self.num_math_warps = 8
+        self.tma_warp_base = 8
+        self.umma_warp_base = 10
+
+        self.num_q_stages = 3  # 3 stages for Q pipelining across batch sequences
+
+        # TMEM: 512 columns total, each group needs N columns per UMMA stage
+        # max_umma_stages = 512 // (2 * N)
+        TMEM_COLS = 512
+        if max_umma_pipeline:
+            self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N))
+        else:
+            self.num_umma_stages = 1
+
+        if max_kv_pipeline:
+            smem_capacity = utils.get_smem_capacity_in_bytes()
+            # Reserve ~1 KB for barriers and misc
+            SMEM_BUDGET = smem_capacity - 1024
+            # KV+Scale per stage (×2 groups):
+            #   2 * (block_kv * head_dim * 1B + block_kv * 4B)
+            kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4)
+            # Q+W per stage: Q is N * head_dim * 1B, W uses padded stride
+            qw_per_stage = self.N * head_dim + self.w_stage_stride * self.epi_bytes
+            qw_total = qw_per_stage * self.num_q_stages
+            self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage
+        else:
+            self.num_kv_stages = 3
+
+        # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB),
+        # avoiding sub-bank conflicts with UMMA reading sKV.
+        # Layout: barriers(~256B) | sKV_0 | sKV_1 | sQ | [pad] | sW | sScales
+        if self.smem_subpartition_opt:
+            BOUNDARY = 128 * 1024
+            used = (
+                256
+                + 2 * (block_kv * head_dim * self.num_kv_stages)
+                + self.N * head_dim * self.num_q_stages
+            )
+            used = ((used + 127) // 128) * 128
+            if used < BOUNDARY:
+                self.smem_pad_bytes = ((BOUNDARY - used + 1023) // 1024) * 1024
+            else:
+                self.smem_pad_bytes = 0
+        else:
+            self.smem_pad_bytes = 0
+
+        self.acc_dtype = acc_dtype
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.cluster_shape_mn = (1, 1)
+        self.mma_tiler_mn = (block_kv, self.N)
+
+    def _setup_mma(self, a_dtype, b_dtype, a_major, b_major):
+        self.a_dtype = a_dtype
+        self.b_dtype = b_dtype
+        self.a_major_mode = a_major
+        self.b_major_mode = b_major
+
+        self.mma_tiler = (*self.mma_tiler_mn, 1)
+        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            a_dtype,
+            a_major,
+            b_major,
+            self.acc_dtype,
+            self.cta_group,
+            self.mma_tiler_mn,
+        )
+        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])  # 32
+
+        # Full-K: tile K = head_dim (128), 1 TMA per block
+        mma_inst_tile_k = self.head_dim // mma_inst_shape_k  # 4
+        full_k = mma_inst_shape_k * mma_inst_tile_k  # 128
+        self.mma_tiler = (
+            self.mma_tiler_mn[0],
+            self.mma_tiler_mn[1],
+            full_k,
+        )
+
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (tiled_mma.thr_id.shape,),
+        )
+        self.epi_tile = self.cta_tile_shape_mnk[:2]
+
+        # KV SMEM: 3 stages per group, each stage holds full [128, 128]
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            self.mma_tiler,
+            a_dtype,
+            self.num_kv_stages,
+        )
+        # Q SMEM: 1 stage, holds full [N, 128]
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            self.mma_tiler,
+            b_dtype,
+            self.num_q_stages,
+        )
+
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+        self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake)
+        self.num_tmem_alloc_cols_total = (
+            self.num_tmem_alloc_cols * self.num_groups * self.num_umma_stages
+        )
+
+        return tiled_mma
+
+    @cute.jit
+    def __call__(
+        self,
+        kv_fused: cute.Tensor,  # Fused KV: [num_phys_blocks, block_bytes] FP8
+        b: cute.Tensor,  # Q: [N, head_dim, batch_size]
+        weights: cute.Tensor,  # [N, batch_size] (transposed for TMA)
+        logits: cute.Tensor,  # [batch_size * next_n, max_context_len]
+        block_table: cute.Tensor,  # [batch_size, max_blocks_per_seq]
+        context_lens: cute.Tensor,  # [batch_size]
+        schedule_meta: cute.Tensor,  # [num_sms+1, 2] int32
+        num_phys_blocks: cutlass.Int32,
+        batch_size: cutlass.Int32,
+        rank: cutlass.Int32,
+        stream: cuda.CUstream,
+    ):
+        # Derive KV and Scale views from fused buffer using CuTE ops.
+        # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)]
+        phys_block_kv = self.phys_block_kv
+        scale_offset_elems = phys_block_kv * self.head_dim  # in FP8 elements
+
+        # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference)
+        kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN)
+
+        # Q (b) was passed as uint8 to work around DLPack's lack of float8 support;
+        # recast back to FP8 so MMA type inference and TMA descriptors are correct.
+        b = cute.recast_tensor(b, cutlass.Float8E4M3FN)
+
+        # Read the real per-block stride (bytes = FP8 elements) from the input.
+        # When KV is the indexer K-cache pool view, the pool is laid out as
+        # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride =
+        # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes).
+        # Using the input stride keeps both contiguous test path and the
+        # strided prod path correct. FP8 = 1 byte per element, so the byte
+        # stride is the same as the element stride after recast.
+        kv_block_stride = kv_fused.layout.stride[0]
+
+        # KV view: [phys_block_kv, head_dim, num_phys_blocks] FP8
+        # Each TMA loads one physical block; multiple TMAs fill a compute tile.
+        kv_layout = cute.make_layout(
+            (phys_block_kv, self.head_dim, num_phys_blocks),
+            stride=(self.head_dim, 1, kv_block_stride),
+        )
+        a = cute.make_tensor(kv_fp8.iterator, kv_layout)
+
+        # Scale view: offset pointer to scale region, recast FP8 → Float32
+        # [phys_block_kv, num_phys_blocks] float32 (after recast)
+        scale_fp8_layout = cute.make_layout(
+            (phys_block_kv * 4, num_phys_blocks),
+            stride=(1, kv_block_stride),
+        )
+        scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout)
+        scales = cute.recast_tensor(scale_fp8, cutlass.Float32)
+
+        a_dtype = a.element_type
+        b_dtype = b.element_type
+        a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode()
+        b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+
+        tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major)
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        # TMA for KV (A) — fmha_decode_paged pattern.
+        # Build a TMA SMEM layout via tiled_divide on the full compute-tile
+        # layout, then select to drop trivial K dim. Atom uses mode [0] as
+        # single-tile SMEM layout and (phys, head) as cta_tiler.
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+        self.a_tma_view_layout = sm100_utils.make_smem_layout(
+            tcgen05.OperandMajorMode.K,
+            (self.block_kv, self.head_dim),
+            a_dtype,
+            self.num_kv_stages,
+        )
+        self.a_tma_view_layout = cute.tiled_divide(
+            self.a_tma_view_layout, (self.phys_block_kv, self.head_dim)
+        )
+        # ((tile_M, tile_K), rest_M, rest_K, stages) → drop trivial rest_K
+        self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3])
+        # ((tile_M, tile_K), rest_M=num_sub_blocks, stages)
+        tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            a,
+            self.a_tma_view_layout[0],  # atom SMEM = single-tile (mode 0)
+            (self.phys_block_kv, self.head_dim),
+        )
+
+        # TMA for Q (B) — full K=128, L dim = batch_size (unchanged)
+        b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            b_op,
+            b,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # TMA for Weights — [N, batch_size], tile [N], L=batch_size
+        self.w_smem_layout_staged = cute.make_layout(
+            (self.N, self.num_q_stages),
+            stride=(1, self.w_stage_stride),
+        )
+        w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0])
+        tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            weights,
+            w_smem_per_stage,
+            self.w_smem_layout_staged.shape[:1],
+        )
+
+        # TMA for Scales — [phys_block_kv, num_phys_blocks], tile [phys_block_kv]
+        # SMEM holds compute_block_kv scales per stage; filled by
+        # num_blocks_per_mma sub-block TMAs at consecutive offsets.
+        self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages))
+        s_smem_per_subblock = cute.make_layout((phys_block_kv,))
+        tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            scales,
+            s_smem_per_subblock,
+            (phys_block_kv,),
+        )
+
+        b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout)
+        w_copy_size = self.N * self.epi_bytes
+        # Per sub-block: phys_block_kv * head_dim (KV) + phys_block_kv * 4 (scales)
+        kv_tma_bytes_per_subblock = phys_block_kv * self.head_dim
+        scale_tma_bytes_per_subblock = phys_block_kv * 4
+        # Total per compute tile = num_blocks_per_mma sub-blocks
+        self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * (
+            kv_tma_bytes_per_subblock + scale_tma_bytes_per_subblock
+        )
+        # Q + Weights share barrier (like DeepGEMM)
+        self.num_q_tma_bytes = b_copy_size * atom_thr_size + w_copy_size
+
+        num_ctas = self.num_sms
+
+        @cute.struct
+        class SharedStorage:
+            kv_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            kv_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2]
+            umma_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            umma_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            tmem_holding_buf: cutlass.Int32
+
+        self.kernel(
+            tiled_mma,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_w,
+            tma_tensor_w,
+            tma_atom_s,
+            tma_tensor_s,
+            logits,
+            block_table,
+            context_lens,
+            schedule_meta,
+            batch_size,
+            rank,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.w_smem_layout_staged,
+            self.s_smem_layout_staged,
+            self.a_tma_view_layout,
+            self.epi_tile,
+            SharedStorage,
+        ).launch(
+            grid=(1, 1, num_ctas),
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,  # KV pool
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,  # Q (L dim = batch_size)
+        tma_atom_w: cute.CopyAtom,
+        mW_tma: cute.Tensor,  # Weights TMA coord tensor [N, batch_size]
+        tma_atom_s: cute.CopyAtom,
+        mS_tma: cute.Tensor,  # Scales TMA coord tensor [block_kv, num_phys_blocks]
+        mLogits: cute.Tensor,  # [batch_size * next_n, max_context_len]
+        mBlockTable: cute.Tensor,  # [batch_size, max_blocks_per_seq]
+        mContextLens: cute.Tensor,  # [batch_size]
+        mScheduleMeta: cute.Tensor,  # unused (kept for the indexer_topk launcher signature)
+        batch_size: cutlass.Int32,
+        rank: cutlass.Int32,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        w_smem_layout_staged: cute.Layout,
+        s_smem_layout_staged: cute.Layout,
+        a_tma_view_layout: cute.ComposedLayout,
+        epi_tile: cute.Tile,
+        SharedStorage: cutlass.Constexpr,
+    ):
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
+
+        bidx, bidy, bidz = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+        is_leader_cta = mma_tile_coord_v == 0
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        tidx, _, _ = cute.arch.thread_idx()
+        if cutlass.const_expr(self.trigger_dependents):
+            griddepcontrol_launch_dependents()
+
+        # Warp roles (matches DeepGEMM SM100)
+        warpgroup_idx = warp_idx // 4
+        is_math_warp = warp_idx < 8
+        is_tma_warp_0 = warp_idx == 8
+        is_tma_warp_1 = warp_idx == 9
+        is_tma_warp = is_tma_warp_0 | is_tma_warp_1
+        is_umma_warp_0 = warp_idx == 10
+        is_umma_warp_1 = warp_idx == 11
+
+        NUM_MATH_WG = 2  # kNumMathWarpGroups
+        NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma
+        sm_idx = bidz
+        world = self.world
+        page_tokens = self.phys_block_kv
+
+        if is_tma_warp:
+            cpasync.prefetch_descriptor(tma_atom_a)
+            cpasync.prefetch_descriptor(tma_atom_b)
+            cpasync.prefetch_descriptor(tma_atom_w)
+            cpasync.prefetch_descriptor(tma_atom_s)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        block_kv_val = self.block_kv
+        num_heads = self.num_heads
+        next_n = self.next_n
+        num_epi_subtiles = self.num_epi_subtiles
+
+        # === Pipelines ===
+        prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+
+        # Q pipeline: TMA producer → Math consumer (8 math warps)
+        # PipelineTmaAsync: consumer_release uses is_signalling_thread
+        # (lane 0 per warp). 8 math warps × 1 lane-0 = 8 arrives.
+        q_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 8)
+        q_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.q_mbar.data_ptr(),
+            num_stages=self.num_q_stages,
+            producer_group=prod_group,
+            consumer_group=q_cons_group,
+            tx_count=self.num_q_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        q_prod_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_q_stages
+        )
+        # Both Math WGs share the same pipeline state (advance in lockstep)
+        q_cons_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        # UMMA warps observe Q pipeline (wait only, no release)
+        # to ensure Q is in SMEM before GEMM. Critical for UMMA warp 1
+        # since TMA warp 1 only loads KV1 (not Q).
+        q_cons_state_umma_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        q_cons_state_umma_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+
+        # Merged KV+Scale pipelines (per-group, 3 stages each)
+        # Like DeepGEMM: KV data and scales share one barrier.
+        # TMA loads both under one barrier. Math is consumer (releases).
+        # UMMA also waits on this barrier (for KV GEMM) but does NOT release.
+        math_warps_per_group = self.num_math_warps // 2  # 4 warps
+        kv_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, math_warps_per_group)
+        kv_pipeline_0 = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar_0.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=kv_cons_group,
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        kv_pipeline_1 = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar_1.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=kv_cons_group,
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+
+        kv_prod_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        kv_prod_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        # UMMA consumer states (wait only, no release)
+        kv_cons_state_umma_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_umma_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        # Math consumer states (wait + release)
+        kv_cons_state_math_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_math_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+
+        # UMMA pipelines (per-group)
+        math_threads_per_group = self.num_math_threads // 2
+        umma_pipeline_0 = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar_0.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        umma_pipeline_1 = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar_1.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+
+        umma_prod_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_prod_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_cons_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+        umma_cons_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+
+        # TMEM — only Math warps (8×32=256) + UMMA warps (2×32=64) = 320 threads
+        # TMA warps do NOT participate, so they can start TMA loads earlier.
+        # Math warp 0 is the allocator (like fp16_gemm_3's epilogue warp 0),
+        # because math warps are the last TMEM consumers (epilogue reads).
+        tmem_alloc_num_threads = 320  # 10 warps: warp 0-7 (math) + warp 10-11 (umma)
+        tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=tmem_alloc_num_threads)
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=0,  # math warp 0 does alloc+free (last TMEM consumer)
+            is_two_cta=False,
+        )
+
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+
+        # SMEM allocation: per-group KV + shared Q
+        sKV_0 = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sKV_1 = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sQ = smem.allocate_tensor(
+            element_type=self.b_dtype,
+            layout=b_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=b_smem_layout_staged.inner,
+        )
+        # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB)
+        # to avoid sub-bank conflicts with UMMA reading sKV from
+        # sub-partition 0.
+        if cutlass.const_expr(self.smem_pad_bytes > 0):
+            _ = smem.allocate(self.smem_pad_bytes)
+        # Weights SMEM: [N, num_q_stages], shared Q barrier
+        sW = smem.allocate_tensor(
+            element_type=self.epi_dtype,
+            layout=w_smem_layout_staged,
+            byte_alignment=128,
+        )
+        # Scales SMEM: [block_kv, num_kv_stages] float32, per group
+        sScales_0 = smem.allocate_tensor(
+            element_type=cutlass.Float32,
+            layout=s_smem_layout_staged,
+            byte_alignment=128,
+        )
+        sScales_1 = smem.allocate_tensor(
+            element_type=cutlass.Float32,
+            layout=s_smem_layout_staged,
+            byte_alignment=128,
+        )
+
+        a_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+        )
+        b_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+        )
+
+        # Partition KV (A): fmha_decode_paged pattern.
+        # SMEM view is ((tile), num_sub_blocks, stages) — built in __call__.
+        # Use .outer (plain layout); swizzle is captured by sKV_0's iterator.
+        # GMEM: local_tile by (phys, head), then group first 2 modes into tile.
+        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+        sKV_0_for_tma = cute.make_tensor(sKV_0.iterator, a_tma_view_layout.outer)
+        sKV_1_for_tma = cute.make_tensor(sKV_1.iterator, a_tma_view_layout.outer)
+        gA = cute.local_tile(
+            mA_mkl,
+            (self.phys_block_kv, self.head_dim),
+            coord=(None, None, None),
+        )
+        tAsA_0, tAgA_0 = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            sKV_0_for_tma,
+            cute.group_modes(gA, 0, 2),
+        )
+        tAsA_1, tAgA_1 = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            sKV_1_for_tma,
+            cute.group_modes(gA, 0, 2),
+        )
+
+        # Partition Q (B): shared SMEM, L dim = batch_size
+        gB_nkl = cute.local_tile(
+            mB_nkl,
+            cute.slice_(self.mma_tiler, (0, None, None)),
+            (None, None, None),
+        )
+        tCgB = thr_mma.partition_B(gB_nkl)
+        b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            block_in_cluster_coord_vmnk[1],
+            b_cta_layout,
+            cute.group_modes(sQ, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+        tBgB = tBgB[(None, 0, None, None)]  # [tma, K, L]
+
+        # Partition Weights: standalone TMA, [N, batch_size] → [N] per stage
+        w_cta_layout = cute.make_layout((1,))
+        tWsW, tWgW = cpasync.tma_partition(
+            tma_atom_w,
+            0,
+            w_cta_layout,
+            cute.group_modes(sW, 0, 1),
+            cute.group_modes(mW_tma, 0, 1),
+        )
+
+        # Partition Scales: explicit sub_blocks + stages dims.
+        # Layout (phys_block_kv, num_sub, stages) K-major with custom strides.
+        s_tma_view_layout = cute.make_layout(
+            (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages),
+            stride=(1, self.phys_block_kv, self.block_kv),
+        )
+        sScales_0_for_tma = cute.make_tensor(sScales_0.iterator, s_tma_view_layout)
+        sScales_1_for_tma = cute.make_tensor(sScales_1.iterator, s_tma_view_layout)
+        # GMEM: local_tile by phys to match atom's tile size
+        gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None))
+        tSsS_0, tSgS_0 = cpasync.tma_partition(
+            tma_atom_s,
+            0,
+            cute.make_layout(1),
+            sScales_0_for_tma,
+            gS,
+        )
+        tSsS_1, tSgS_1 = cpasync.tma_partition(
+            tma_atom_s,
+            0,
+            cute.make_layout(1),
+            sScales_1_for_tma,
+            gS,
+        )
+
+        # MMA fragments
+        tCrA_0 = tiled_mma.make_fragment_A(sKV_0)
+        tCrA_1 = tiled_mma.make_fragment_A(sKV_1)
+        tCrB = tiled_mma.make_fragment_B(sQ)  # shared
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+
+        # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim
+        # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank
+        us = self.num_umma_stages
+        cols = self.num_tmem_alloc_cols
+        acc_shape_staged = cute.append(acc_shape, us)
+        tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged)
+
+        # TMEM layout info (allocation deferred to UMMA/Math warp branches)
+        cols_per_group = cols * us * (32 // self.acc_dtype.width)
+        num_tmem_alloc_cols_total = self.num_tmem_alloc_cols_total
+
+        # Epilogue setup
+        c_layout = utils.LayoutEnum.ROW_MAJOR
+        epi_sub_mn = (epi_tile[0], num_heads // num_epi_subtiles)
+        copy_atom_t2r = sm100_utils.get_tmem_load_op(
+            self.cta_tile_shape_mnk,
+            c_layout,
+            self.acc_dtype,
+            self.acc_dtype,
+            epi_sub_mn,
+            use_2cta_instrs,
+        )
+
+        # TMEM is allocated before the grid-dependency wait (math warp 0 is the
+        # allocator; see the module docstring for why this cannot deadlock).
+        if cutlass.const_expr(self.wait_at_entry):
+            tmem.allocate(num_tmem_alloc_cols_total)
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+        if cutlass.const_expr(self.wait_at_entry):
+            # Every global-memory access of this CTA comes after this point.
+            cute.arch.griddepcontrol_wait()
+
+        # Context partition (one request): this rank's local tokens are its
+        # pages among the global pages covering the indexer KV length. Loaded
+        # after the wait: the length is per-step metadata.
+        global_kv = mContextLens[0]
+        full_pages = global_kv // page_tokens
+        tail_tokens = global_kv - full_pages * page_tokens
+        local_len = ((full_pages - rank + world - 1) // world) * page_tokens
+        if tail_tokens > 0:
+            if full_pages % world == rank:
+                local_len = local_len + tail_tokens
+        num_local_pages = (local_len + page_tokens - 1) // page_tokens
+        current_num_kv = (local_len + self.block_kv - 1) // self.block_kv
+        # Contiguous share of the tile pairs per CTA (schedule units of
+        # NUM_MATH_WG tiles, as in the DeepGEMM schedule this replaces). A
+        # share that reaches the end of the request ends at (q = 1, kv = 0).
+        num_units = (current_num_kv + NUM_MATH_WG - 1) // NUM_MATH_WG
+        num_ctas_total = self.num_sms
+        start_unit = (sm_idx * num_units) // num_ctas_total
+        end_unit = ((sm_idx + 1) * num_units) // num_ctas_total
+        start_q = cutlass.Int32(0)
+        start_kv_half = start_unit
+        end_q_idx = cutlass.Int32(0)
+        end_kv_half = end_unit
+        if end_unit == num_units:
+            end_q_idx = cutlass.Int32(1)
+            end_kv_half = cutlass.Int32(0)
+        if start_unit == end_unit:
+            start_q = cutlass.Int32(1)
+            start_kv_half = cutlass.Int32(0)
+            end_q_idx = cutlass.Int32(1)
+            end_kv_half = cutlass.Int32(0)
+
+        # ===== SCHEDULER: derive values from the schedule computed above =====
+        end_kv_idx = end_kv_half * NUM_MATH_WG
+
+        # Convert start to KV block units (matching DeepGEMM)
+        current_q_idx = start_q
+        current_kv_idx = start_kv_half * NUM_MATH_WG
+
+        # ===== COMMON SCHEDULER STATE (before warp branches) =====
+        # Each warp role independently maintains its own copy of these
+        # variables (like DeepGEMM where each role creates its own scheduler).
+        # Pre-fetch first task (current_num_kv loaded early above for latency hiding)
+        next_q_idx = current_q_idx
+        next_kv_idx = current_kv_idx
+        next_num_kv = current_num_kv
+        # Sentinel: no previous batch (matches DeepGEMM's q_idx = batch_size)
+        q_idx = batch_size
+        # While-loop termination flag (matches DeepGEMM's fetch_next_task pattern).
+        # True if this CTA has work assigned (start != end in schedule_meta).
+        has_work = (current_q_idx != end_q_idx) | (current_kv_idx != end_kv_idx)
+
+        # ===== WARP-SPECIALIZED EXECUTION =====
+
+        if is_tma_warp_0:
+            # TMA warp 0: loads Q (prefetch) + KV for group 0
+            # Matches DeepGEMM's TMA warp with kv_group_idx == 0
+            cute.arch.warpgroup_reg_dealloc(24)
+            lane_idx = tidx % 32
+
+            # Block table prefetch: 32 lanes cache block indices,
+            # distributed via shuffle. Each lane holds num_blocks_per_mma
+            # physical block indices per compute tile. (DeepGEMM L233-244)
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)  # force prefetch on first use
+
+            # Prefetch first Q before loop (like DeepGEMM line 203-204)
+            q_pipeline.producer_acquire(q_prod_state)
+            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, 0, next_q_idx)],
+                tBsB[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+                mcast_mask=b_mcast_mask,
+            )
+            cute.copy(
+                tma_atom_w,
+                tWgW[(None, next_q_idx)],
+                tWsW[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+            )
+            q_prod_state.advance()
+
+            while has_work:
+                # fetch_next_task: commit next → current
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                # Q prefetch: when batch changes, load Q for NEXT batch
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)  # force re-prefetch
+                    prefetch_next = q_idx + 1
+                    if prefetch_next < end_q_idx:
+                        q_pipeline.producer_acquire(q_prod_state)
+                        q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                        cute.copy(
+                            tma_atom_b,
+                            tBgB[(None, 0, prefetch_next)],
+                            tBsB[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                            mcast_mask=b_mcast_mask,
+                        )
+                        cute.copy(
+                            tma_atom_w,
+                            tWgW[(None, prefetch_next)],
+                            tWsW[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                        )
+                        q_prod_state.advance()
+                    elif prefetch_next == end_q_idx:
+                        if end_kv_idx > 0:
+                            q_pipeline.producer_acquire(q_prod_state)
+                            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                            cute.copy(
+                                tma_atom_b,
+                                tBgB[(None, 0, prefetch_next)],
+                                tBsB[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                                mcast_mask=b_mcast_mask,
+                            )
+                            cute.copy(
+                                tma_atom_w,
+                                tWgW[(None, prefetch_next)],
+                                tWsW[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                            )
+                            q_prod_state.advance()
+
+                # Block table prefetch for group 0 (like DeepGEMM L233-241).
+                # Each lane loads num_blocks_per_mma physical block indices
+                # for one compute tile (kv_idx counts compute tiles).
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            local_page = min(base_phys + i, num_local_pages - 1)
+                            cached_blks[i] = mBlockTable[(q_idx, local_page * world + rank)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+
+                # Get block indices via shuffle before barrier (like DeepGEMM L244)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile.
+                kv_pipeline_0.producer_acquire(kv_prod_state_0)
+                bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0)
+                stage = kv_prod_state_0.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_0[(None, 0, 0, phys_blks[i])],
+                        tAsA_0[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_s,
+                        tSgS_0[(None, 0, phys_blks[i])],
+                        tSsS_0[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                    )
+                kv_prod_state_0.advance()
+
+                # Advance: inline fetch_next_task
+                next_kv_idx = kv_idx + NUM_MATH_WG
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                # Update while-loop condition
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_tma_warp_1:
+            # TMA warp 1: loads KV + Scale for group 1 only
+            # Matches DeepGEMM's TMA warp with kv_group_idx == 1
+            cute.arch.warpgroup_reg_dealloc(24)
+            lane_idx = tidx % 32
+
+            # Block table prefetch for group 1
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)  # force prefetch on first use
+
+            while has_work:
+                # fetch_next_task: commit next → current
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                # New q_idx → force block table re-prefetch
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)
+
+                # Block table prefetch for group 1 (like DeepGEMM L233-241).
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            local_page = min(base_phys + i, num_local_pages - 1)
+                            cached_blks[i] = mBlockTable[(q_idx, local_page * world + rank)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+
+                # Get block indices via shuffle before barrier (like DeepGEMM L244)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile.
+                kv_pipeline_1.producer_acquire(kv_prod_state_1)
+                bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1)
+                stage = kv_prod_state_1.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_1[(None, 0, 0, phys_blks[i])],
+                        tAsA_1[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_s,
+                        tSgS_1[(None, 0, phys_blks[i])],
+                        tSsS_1[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                    )
+                kv_prod_state_1.advance()
+
+                # Advance: inline fetch_next_task
+                next_kv_idx = kv_idx + NUM_MATH_WG
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                # Update while-loop condition
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp_0:
+            # UMMA warp for group 0
+            # Must wait on Q pipeline: TMA operations with different
+            # barriers are NOT visibility-ordered even within the same
+            # warp. KV0 barrier arriving does not guarantee Q SMEM
+            # writes are visible.
+            cute.arch.warpgroup_reg_dealloc(24)
+
+            # TMEM: wait for math warp 0's allocation, retrieve pointer
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            if is_leader_cta:
+                num_k_blocks = cute.size(tCrA_0.shape[2])
+                q_stage_0 = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Wait for Q pipeline when batch changes
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma_0.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma_0)
+                        q_stage_0 = q_cons_state_umma_0.index
+
+                    # Process KV block for group 0 (kv_idx + 0)
+                    # Unconditional UMMA (like DeepGEMM): OOB iterations
+                    # compute on garbage data; results written to aligned
+                    # padding region in logits buffer.
+                    # Wait KV first, then TMEM empty (like DeepGEMM)
+                    kv_pipeline_0.consumer_wait(kv_cons_state_umma_0)
+                    umma_pipeline_0.producer_acquire(umma_prod_state_0)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage = kv_cons_state_umma_0.index
+                    tCtAcc_0 = tCtAcc_base_0[(None, None, None, umma_prod_state_0.index)]
+                    for k_block in cutlass.range_constexpr(num_k_blocks):
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc_0,
+                            tCrA_0[None, None, k_block, kv_stage],
+                            tCrB[None, None, k_block, q_stage_0],
+                            tCtAcc_0,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    # No consumer_release here — Math WG 0 releases
+                    kv_cons_state_umma_0.advance()
+
+                    umma_pipeline_0.producer_commit(umma_prod_state_0)
+                    umma_prod_state_0.advance()
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp_1:
+            # UMMA warp for group 1
+            # Explicitly waits on Q pipeline — critical because TMA warp 1
+            # only loads KV1, not Q. Without this wait, UMMA warp 1 can
+            # start GEMM before TMA warp 0 finishes loading Q into SMEM.
+            cute.arch.warpgroup_reg_dealloc(24)
+
+            # TMEM: wait for umma_warp_0's allocation, retrieve pointer
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            if is_leader_cta:
+                num_k_blocks_1 = cute.size(tCrA_1.shape[2])
+                q_stage_1 = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Wait for Q pipeline when batch changes
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma_1.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma_1)
+                        q_stage_1 = q_cons_state_umma_1.index
+
+                    # Process KV block for group 1 (kv_idx + 1)
+                    # Unconditional UMMA (like DeepGEMM)
+                    # Wait KV first, then TMEM empty (like DeepGEMM)
+                    kv_pipeline_1.consumer_wait(kv_cons_state_umma_1)
+                    umma_pipeline_1.producer_acquire(umma_prod_state_1)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage_1 = kv_cons_state_umma_1.index
+                    tCtAcc_1 = tCtAcc_base_1[(None, None, None, umma_prod_state_1.index)]
+                    for k_block in cutlass.range_constexpr(num_k_blocks_1):
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc_1,
+                            tCrA_1[None, None, k_block, kv_stage_1],
+                            tCrB[None, None, k_block, q_stage_1],
+                            tCtAcc_1,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    # No consumer_release here — Math WG 1 releases
+                    kv_cons_state_umma_1.advance()
+
+                    umma_pipeline_1.producer_commit(umma_prod_state_1)
+                    umma_prod_state_1.advance()
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_math_warp:
+            cute.arch.warpgroup_reg_alloc(240)
+
+            # TMEM: math warp 0 is the allocator; all math warps wait + retrieve
+            # (with wait_at_entry the allocation was issued before the wait).
+            if cutlass.const_expr(not self.wait_at_entry):
+                tmem.allocate(num_tmem_alloc_cols_total)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            local_tidx = tidx % 128
+            cC = cute.make_identity_tensor(epi_sub_mn)
+
+            if warpgroup_idx == 0:
+                # Math WG 0: process group 0
+                # Reference setup (stage 0) for m_coord
+                # flat_divide by sub-tile to get sub-tile partitions
+                tAcc_0_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)]
+                tAcc_0_ref_epi = cute.flat_divide(tAcc_0_ref, epi_sub_mn)
+                tiled_copy_ref_0 = tcgen05.make_tmem_copy(
+                    copy_atom_t2r, tAcc_0_ref_epi[(None, None, 0, 0)]
+                )
+                thr_copy_ref_0 = tiled_copy_ref_0.get_slice(local_tidx)
+                tTR_cC = thr_copy_ref_0.partition_D(cC)
+                m_coord = tTR_cC[0][0]
+
+                tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+                # Weight register cache: only first NUM_W_IN_REG
+                # per next_n slot (like DeepGEMM min(52, kNumHeads)).
+                # Remaining weights read from SMEM in epilogue.
+                # FP16 weights use half the regs, so we can fit
+                # all heads for next_n <= 3.
+                NUM_W_IN_REG = self.num_w_in_reg
+                w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype)
+                q_stage_local = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Q pipeline consumer: wait for Q+Weights SMEM
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_pipeline.consumer_release(q_cons_state)
+                            q_cons_state.advance()
+                        q_pipeline.consumer_wait(q_cons_state)
+                        q_stage_local = q_cons_state.index
+                        # Preload first NUM_W_IN_REG weights per slot
+                        for t_i in cutlass.range_constexpr(next_n):
+                            for w_j in cutlass.range_constexpr(NUM_W_IN_REG):
+                                w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
+                                    (t_i * num_heads + w_j, q_stage_local)
+                                ]
+
+                    # Process KV block for group 0 (kv_idx + 0)
+                    # Unconditional Math (like DeepGEMM): OOB results
+                    # written to aligned padding region in logits buffer.
+                    kv_pos = kv_idx * block_kv_val + m_coord
+
+                    if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                        # Skip KV wait, rely on UMMA barrier's
+                        # transitive visibility.
+                        umma_pipeline_0.consumer_wait(umma_cons_state_0)
+                    else:
+                        # Default: wait KV first to overlap lds with
+                        # UMMA computation.
+                        kv_pipeline_0.consumer_wait(kv_cons_state_math_0)
+                        sc_stage_0 = kv_cons_state_math_0.index
+                        scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        umma_pipeline_0.consumer_wait(umma_cons_state_0)
+
+                    # --- TMEM sub-tile setup ---
+                    # flat_divide accumulator by sub-tile shape;
+                    # partition once, then loop over sub-tiles.
+                    tCtAcc_c0 = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)]
+                    tAcc_c0 = tCtAcc_c0[((None, None), 0, 0)]
+                    tAcc_c0_epi = cute.flat_divide(tAcc_c0, epi_sub_mn)
+                    tc_0 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c0_epi[(None, None, 0, 0)])
+                    tr_0 = tc_0.get_slice(local_tidx)
+                    tTR_0 = tr_0.partition_S(tAcc_c0_epi)
+
+                    # --- First sub-tile LDTM + KV release ---
+                    if cutlass.const_expr(self.early_tmem_copy):
+                        # Issue first sub-tile LDTM early
+                        cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc)
+                        # Scale LDS + KV release fill latency
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_0 = kv_cons_state_math_0.index
+                            scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        kv_pipeline_0.consumer_release(kv_cons_state_math_0)
+                        kv_cons_state_math_0.advance()
+                        cute.arch.fence_view_async_tmem_load()
+                    else:
+                        # Default: scale LDS + KV release first
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_0 = kv_cons_state_math_0.index
+                            scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        kv_pipeline_0.consumer_release(kv_cons_state_math_0)
+                        kv_cons_state_math_0.advance()
+                        cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc)
+                        cute.arch.fence_view_async_tmem_load()
+
+                    # --- Sub-tile compute loop ---
+                    # Each sub-tile: LDTM.xN → fence → load →
+                    # ReLU+FMA. Breaks FMA chain (16→4 per chunk)
+                    # and interleaves LDTM with FP32 compute to
+                    # reduce ShadowPipeThrottle.
+                    # Sub-tiles are within each t-slot (num_heads
+                    # // num_epi_subtiles wide). flat_divide yields
+                    # next_n * num_epi_subtiles sub-tiles total;
+                    # global index = t * num_epi_subtiles + i.
+                    subtile_n = num_heads // num_epi_subtiles
+                    if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                        packed_zero = pack_f16x2(Float16(0.0), Float16(0.0))
+                    for t in cutlass.range_constexpr(next_n):
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps0 = packed_zero
+                            ps1 = packed_zero
+                        else:
+                            s0x = cutlass.Float32(0.0)
+                            s0y = cutlass.Float32(0.0)
+                            s1x = cutlass.Float32(0.0)
+                            s1y = cutlass.Float32(0.0)
+                        for i in cutlass.range_constexpr(num_epi_subtiles):
+                            # LDTM for sub-tiles 1..N-1
+                            # (sub-tile 0 handled above)
+                            if t > 0 or i > 0:
+                                cute.copy(
+                                    tc_0,
+                                    tTR_0[(None, None, None, 0, t * num_epi_subtiles + i)],
+                                    tTR_rAcc,
+                                )
+                                cute.arch.fence_view_async_tmem_load()
+                            # Release UMMA after last LDTM+fence
+                            if t == next_n - 1 and i == num_epi_subtiles - 1:
+                                umma_pipeline_0.consumer_release(umma_cons_state_0)
+                                umma_cons_state_0.advance()
+                            acc_vec = tTR_rAcc.load()
+                            # Reg-path: weights from registers
+                            reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n))
+                            for h in cutlass.range_constexpr(0, reg_h_end, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1])
+                                    pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3])
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    w0 = w_cache[r0]
+                                    w1 = w_cache[r0 + 1]
+                                    w2 = w_cache[r0 + 2]
+                                    w3 = w_cache[r0 + 3]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                            # SMEM-path: weights from shared mem
+                            smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n)
+                            for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    pw01 = pack_f16x2(
+                                        sW[(t * num_heads + h_g, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 1, q_stage_local)],
+                                    )
+                                    pw23 = pack_f16x2(
+                                        sW[(t * num_heads + h_g + 2, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 3, q_stage_local)],
+                                    )
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    w0 = sW[(t * num_heads + h_g, q_stage_local)]
+                                    w1 = sW[(t * num_heads + h_g + 1, q_stage_local)]
+                                    w2 = sW[(t * num_heads + h_g + 2, q_stage_local)]
+                                    w3 = sW[(t * num_heads + h_g + 3, q_stage_local)]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps_sum = add_f16x2(ps0, ps1)
+                            sum_lo, sum_hi = unpack_f16x2(ps_sum)
+                            result_t = sum_lo + sum_hi
+                        else:
+                            result_t = s0x + s0y + s1x + s1y
+                        out_row = q_idx * next_n + t
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(
+                                result_t * Float16(scale_val)
+                            )
+                        else:
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val)
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+                # Release last Q stage (WG 0)
+                if q_idx < batch_size:
+                    q_pipeline.consumer_release(q_cons_state)
+                    q_cons_state.advance()
+
+            else:
+                # Math WG 1: process group 1
+                tAcc_1_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)]
+                tAcc_1_ref_epi = cute.flat_divide(tAcc_1_ref, epi_sub_mn)
+                tiled_copy_ref_1 = tcgen05.make_tmem_copy(
+                    copy_atom_t2r, tAcc_1_ref_epi[(None, None, 0, 0)]
+                )
+                thr_copy_ref_1 = tiled_copy_ref_1.get_slice(local_tidx)
+                tTR_cC = thr_copy_ref_1.partition_D(cC)
+                m_coord = tTR_cC[0][0]
+
+                tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+                # Weight register cache: only first NUM_W_IN_REG
+                # per next_n slot (like DeepGEMM min(52, kNumHeads)).
+                # FP16 weights use half the regs, so we can fit
+                # all heads for next_n <= 3.
+                NUM_W_IN_REG = self.num_w_in_reg
+                w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype)
+                q_stage_local = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Q pipeline consumer: wait for Q+Weights SMEM
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_pipeline.consumer_release(q_cons_state)
+                            q_cons_state.advance()
+                        q_pipeline.consumer_wait(q_cons_state)
+                        q_stage_local = q_cons_state.index
+                        # Preload first NUM_W_IN_REG weights per slot
+                        for t_i in cutlass.range_constexpr(next_n):
+                            for w_j in cutlass.range_constexpr(NUM_W_IN_REG):
+                                w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
+                                    (t_i * num_heads + w_j, q_stage_local)
+                                ]
+
+                    # Process KV block for group 1 (kv_idx + 1)
+                    # Unconditional Math (like DeepGEMM)
+                    kv_idx_1 = kv_idx + 1
+
+                    kv_pos = kv_idx_1 * block_kv_val + m_coord
+
+                    if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                        umma_pipeline_1.consumer_wait(umma_cons_state_1)
+                    else:
+                        kv_pipeline_1.consumer_wait(kv_cons_state_math_1)
+                        sc_stage_1 = kv_cons_state_math_1.index
+                        scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        umma_pipeline_1.consumer_wait(umma_cons_state_1)
+
+                    # --- TMEM sub-tile setup (WG1) ---
+                    tCtAcc_c1 = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)]
+                    tAcc_c1 = tCtAcc_c1[((None, None), 0, 0)]
+                    tAcc_c1_epi = cute.flat_divide(tAcc_c1, epi_sub_mn)
+                    tc_1 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c1_epi[(None, None, 0, 0)])
+                    tr_1 = tc_1.get_slice(local_tidx)
+                    tTR_1 = tr_1.partition_S(tAcc_c1_epi)
+
+                    # --- First sub-tile LDTM + KV release (WG1) ---
+                    if cutlass.const_expr(self.early_tmem_copy):
+                        cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc)
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_1 = kv_cons_state_math_1.index
+                            scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        kv_pipeline_1.consumer_release(kv_cons_state_math_1)
+                        kv_cons_state_math_1.advance()
+                        cute.arch.fence_view_async_tmem_load()
+                    else:
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_1 = kv_cons_state_math_1.index
+                            scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        kv_pipeline_1.consumer_release(kv_cons_state_math_1)
+                        kv_cons_state_math_1.advance()
+                        cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc)
+                        cute.arch.fence_view_async_tmem_load()
+
+                    # --- Sub-tile compute loop (WG1) ---
+                    subtile_n = num_heads // num_epi_subtiles
+                    if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                        packed_zero = pack_f16x2(Float16(0.0), Float16(0.0))
+                    for t in cutlass.range_constexpr(next_n):
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps0 = packed_zero
+                            ps1 = packed_zero
+                        else:
+                            s0x = cutlass.Float32(0.0)
+                            s0y = cutlass.Float32(0.0)
+                            s1x = cutlass.Float32(0.0)
+                            s1y = cutlass.Float32(0.0)
+                        for i in cutlass.range_constexpr(num_epi_subtiles):
+                            if t > 0 or i > 0:
+                                cute.copy(
+                                    tc_1,
+                                    tTR_1[(None, None, None, 0, t * num_epi_subtiles + i)],
+                                    tTR_rAcc,
+                                )
+                                cute.arch.fence_view_async_tmem_load()
+                            if t == next_n - 1 and i == num_epi_subtiles - 1:
+                                umma_pipeline_1.consumer_release(umma_cons_state_1)
+                                umma_cons_state_1.advance()
+                            acc_vec = tTR_rAcc.load()
+                            # Reg-path
+                            reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n))
+                            for h in cutlass.range_constexpr(0, reg_h_end, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1])
+                                    pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3])
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    w0 = w_cache[r0]
+                                    w1 = w_cache[r0 + 1]
+                                    w2 = w_cache[r0 + 2]
+                                    w3 = w_cache[r0 + 3]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                            # SMEM-path
+                            smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n)
+                            for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    pw01 = pack_f16x2(
+                                        sW[(t * num_heads + h_g, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 1, q_stage_local)],
+                                    )
+                                    pw23 = pack_f16x2(
+                                        sW[(t * num_heads + h_g + 2, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 3, q_stage_local)],
+                                    )
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    w0 = sW[(t * num_heads + h_g, q_stage_local)]
+                                    w1 = sW[(t * num_heads + h_g + 1, q_stage_local)]
+                                    w2 = sW[(t * num_heads + h_g + 2, q_stage_local)]
+                                    w3 = sW[(t * num_heads + h_g + 3, q_stage_local)]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps_sum = add_f16x2(ps0, ps1)
+                            sum_lo, sum_hi = unpack_f16x2(ps_sum)
+                            result_t = sum_lo + sum_hi
+                        else:
+                            result_t = s0x + s0y + s1x + s1y
+                        out_row = q_idx * next_n + t
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(
+                                result_t * Float16(scale_val)
+                            )
+                        else:
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val)
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+                # Release last Q stage (WG 1)
+                if q_idx < batch_size:
+                    q_pipeline.consumer_release(q_cons_state)
+                    q_cons_state.advance()
+
+            # TMEM dealloc: math warps are allocator + last consumer
+            tmem.relinquish_alloc_permit()
+            tmem.free(tmem_ptr)
+
+        else:
+            cute.arch.warpgroup_reg_dealloc(24)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/kernels.py
new file mode 100644
index 000000000000..e95448eff1e2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/kernels.py
@@ -0,0 +1,732 @@
+# 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.
+"""indexer_layer_front fork of absorbed_query's fused full-layer query kernel (``absorbed_query/kernels.py``).
+
+``front_qproj_indexer_kernel`` writes the same bytes as absorbed_query's
+``qproj_indexer_kernel`` on its default path (``PRENORM``: native
+``q_a_layernorm`` output feeding the ``q_b`` and indexer ``wq_b`` MMA loops).
+``NORM_MODE`` selects where q_a is normalized:
+
+* 0 (``native``): ``qn`` is the native ``q_a_layernorm`` output (the same
+  input as absorbed_query's PRENORM path); the producers run shared_indexer_absorbed_query's
+  ``_shared_query_gemm`` on it. The indexer_layer_front changes of this mode are the trigger
+  placement and the position reads after the wait;
+* 1 (``inline``): q_a is normalized inside the kernel and the native RMSNorm
+  launch disappears. Every ``q_b`` and indexer ``wq_b`` producer computes
+  ``rstd`` of its rows from ``kva`` with the bit-exact reproduction of the
+  native FlashInfer RMSNorm (``_fi_rmsnorm_rows_exact``'s arithmetic:
+  lane-to-column map, per-lane fold order, xor butterfly, ``rsqrt``) and forms
+  each K tile of ``qn = bf16((x * rstd) * (w + 0))`` in registers right before
+  absorbed_query's MMA on it (``_front_norm_gemm``: the same tiles, products and ascending
+  K order as absorbed_query's ``_prenorm_gemm``). No program waits for another to obtain
+  ``qn``, and ``qn`` is not stored (tests: ``WRITE_QN``, program 0);
+* 2 (``nr<k>``): q_a is normalized inside the kernel by norm programs (the
+  lowest program ids ``[0, n_norm)``, ``NORM_ROWS`` rows each) that store
+  ``qn`` to a per-call scratch with the same arithmetic and release counter
+  ``NS`` (``atom.add.release.gpu`` after a CTA barrier); the producers spin on
+  it (acquire) and run ``_shared_query_gemm`` on ``qn``.
+
+``PREFETCH_W``: each producer asks L2 for its weight tile (``BN`` rows of
+``wq_b`` / indexer ``wq_b``) before the grid-dependency wait. Every other
+role is absorbed_query's.
+
+Program roles, in program-id order: norm (``NORM_MODE`` 2), ``q_b``
+producers (absorbing producers run absorption job ``p - PRODUCER_JOB_SHIFT``),
+``N_IQ`` indexer ``wq_b`` producers, extra absorption programs, ``N_QPOST``
+indexer query post-processing programs (one per indexer head),
+``compressed_kv`` norm, indexer key post-processing. With the
+``skip_reused_iq`` option (MTP draft steps that reuse the shared top-k) the
+host passes ``N_IQ = N_QPOST = 0``: ``q_fp8``, ``q_scale`` and ``weights``
+are left unwritten and every other output is unchanged.
+``sync_ptr`` = [2 * H absorption counters, IH indexer-head counters, NS
+(released norm programs), NS + 1 (consumers that passed the norm wait)].
+
+Memory order: before ``griddepcontrol.wait`` a
+program reads only immutable weights (the absorbing producers' ``W_UK``
+slice, the producers' L2 prefetch of their weight tile); positions, cos/sin
+rows, activations, counters and every output come after it (absorbed_query reads
+positions before its wait; this fork does not). The last
+consumer of ``NS`` (``NS + 1`` reaching ``n_cons - 1``) resets both words
+after every consumer has passed its wait; counters are touched only after
+the grid-dependency wait, so a launch never sees the counters of an
+incomplete earlier launch. Every spin-wait targets counters released by
+lower program ids, so the launch completes with one resident program slot
+under in-order dispatch (absorbed_query's progress argument, norm programs first).
+
+Trigger (``EARLY_TRIGGER``): every program executes
+``griddepcontrol.launch_dependents`` right after its last global load, which
+in every role comes after all of that program's spin-waits. The dependent
+(the generation rope) is released only when every program has done so, i.e.
+when all programs are resident and past their waits, so its residency
+cannot delay this grid. The dependent reads this kernel's outputs only after
+its own ``griddepcontrol.wait``.
+"""
+
+import triton
+import triton.language as tl
+
+from ..absorbed_query.kernels import (
+    _fi_butterfly32,
+    _fp8_ue8m0_scale,
+    _ld_wuk,
+    _rope_pairs,
+    _stress_delay,
+    _v8_leaves,
+)
+from ..absorbed_query.shared_indexer import _l2_prefetch_1k, _shared_query_gemm
+
+
+@triton.jit
+def _trigger(EARLY_TRIGGER: tl.constexpr):
+    if EARLY_TRIGGER:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+@triton.jit
+def _front_rmsnorm_rows(
+    x_ptr,
+    stride_x,
+    w_ptr,
+    y_ptr,
+    stride_y,
+    rows,
+    rmask,
+    eps,
+    H: tl.constexpr,
+    ROWS: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+):
+    """``_fi_rmsnorm_rows_exact`` (absorbed_query) with the trigger between its last load and its store."""
+    lanes = tl.arange(0, 32)
+    vv = tl.arange(0, 8)
+    offs = rows[:, None, None] * stride_x + 8 * lanes[None, :, None] + vv[None, None, :]
+    s = tl.zeros([ROWS, 32], dtype=tl.float32)
+    for b in tl.static_range(H // 256):
+        x = tl.load(x_ptr + offs + 256 * b, mask=rmask[:, None, None], other=0.0).to(tl.float32)
+        v0, v1, v2, v3, v4, v5, v6, v7 = _v8_leaves(x, ROWS)
+        s = s + v0 * v0
+        s = s + v1 * v1
+        s = s + v2 * v2
+        s = s + v3 * v3
+        s = s + v4 * v4
+        s = s + v5 * v5
+        s = s + v6 * v6
+        s = s + v7 * v7
+    rstd = tl.rsqrt(_fi_butterfly32(s, ROWS) / H + eps)
+    cols = tl.arange(0, H)
+    x = tl.load(
+        x_ptr + rows[:, None] * stride_x + cols[None, :], mask=rmask[:, None], other=0.0
+    ).to(tl.float32)
+    w = tl.load(w_ptr + cols).to(tl.float32) + 0.0
+    _trigger(EARLY_TRIGGER)
+    y = (x * rstd[:, None]) * w[None, :]
+    tl.store(
+        y_ptr + rows[:, None] * stride_y + cols[None, :], y.to(tl.bfloat16), mask=rmask[:, None]
+    )
+
+
+@triton.jit
+def _front_rstd(x_ptr, stride_x, rows, rmask, eps, H: tl.constexpr, ROWS: tl.constexpr):
+    """``rstd`` of ``_fi_rmsnorm_rows_exact`` (absorbed_query): same loads, fold order, butterfly and ``rsqrt``."""
+    lanes = tl.arange(0, 32)
+    vv = tl.arange(0, 8)
+    offs = rows[:, None, None] * stride_x + 8 * lanes[None, :, None] + vv[None, None, :]
+    s = tl.zeros([ROWS, 32], dtype=tl.float32)
+    # A rolled loop over the 256-column blocks (same fold order) bounds the live registers.
+    for b in tl.range(0, H // 256, num_stages=2):
+        x = tl.load(x_ptr + offs + 256 * b, mask=rmask[:, None, None], other=0.0).to(tl.float32)
+        v0, v1, v2, v3, v4, v5, v6, v7 = _v8_leaves(x, ROWS)
+        s = s + v0 * v0
+        s = s + v1 * v1
+        s = s + v2 * v2
+        s = s + v3 * v3
+        s = s + v4 * v4
+        s = s + v5 * v5
+        s = s + v6 * v6
+        s = s + v7 * v7
+    return tl.rsqrt(_fi_butterfly32(s, ROWS) / H + eps)
+
+
+@triton.jit
+def _front_norm_gemm(
+    t,
+    x_ptr,
+    stride_x,
+    g_ptr,
+    w_ptr,
+    out_ptr,
+    OUT_COLS: tl.constexpr,
+    rows,
+    rmask,
+    eps,
+    qn_ptr,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    WRITE_QN: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+):
+    """``out[:, t*BN:(t+1)*BN] = qn @ W[t*BN:(t+1)*BN]^T`` with ``qn`` formed tile by tile.
+
+    ``qn = bf16((x * rstd) * (w + 0))`` element for element as the native
+    FlashInfer RMSNorm (``_fi_rmsnorm_rows_exact``), then absorbed_query's MMA sequence
+    (``_prenorm_gemm``: same tiles, products and ascending K order). With
+    ``WRITE_QN`` (tests) the caller's program 0 also stores ``qn``.
+    """
+    rstd = _front_rstd(x_ptr, stride_x, rows, rmask, eps, K, BLOCK_M)
+    cols = tl.arange(0, BK)
+    n_offs = t * BN + tl.arange(0, BN)
+    acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32)
+    for k0 in tl.range(0, K, BK):
+        x = tl.load(
+            x_ptr + rows[:, None] * stride_x + (k0 + cols)[None, :], mask=rmask[:, None], other=0.0
+        ).to(tl.float32)
+        g = tl.load(g_ptr + k0 + cols).to(tl.float32) + 0.0
+        xn = ((x * rstd[:, None]) * g[None, :]).to(tl.bfloat16)
+        if WRITE_QN:
+            if t == 0:
+                tl.store(qn_ptr + rows[:, None] * K + (k0 + cols)[None, :], xn, mask=rmask[:, None])
+        w = tl.load(w_ptr + n_offs[None, :] * K + (k0 + cols)[:, None])
+        acc = tl.dot(xn, w, acc)
+    _trigger(EARLY_TRIGGER)
+    tl.store(
+        out_ptr + rows[:, None] * OUT_COLS + n_offs[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def _prefetch_w_tile(t, w_ptr, K: tl.constexpr, BN: tl.constexpr, PREFETCH_W: tl.constexpr):
+    """L2 prefetch of weight rows ``[t*BN, (t+1)*BN)`` (``BN * K`` BF16, contiguous)."""
+    if PREFETCH_W:
+        n_chunks: tl.constexpr = BN * K * 2 // 1024
+        _l2_prefetch_1k(w_ptr + t * BN * K, n_chunks, n_chunks)
+
+
+@triton.jit
+def _wait_norm(sync_ptr, NS: tl.constexpr, n_norm, n_cons):
+    """Acquire-spin until every norm program has released ``qn``, then count
+    this consumer; the last one resets both counters for the next launch."""
+    ready = tl.atomic_add(sync_ptr + NS, 0, sem="acquire", scope="gpu")
+    while ready < n_norm:
+        ready = tl.atomic_add(sync_ptr + NS, 0, sem="acquire", scope="gpu")
+    tl.debug_barrier()
+    done = tl.atomic_add(sync_ptr + NS + 1, 1, sem="acq_rel", scope="gpu")
+    if done == n_cons - 1:
+        tl.atomic_xchg(sync_ptr + NS, 0, sem="relaxed", scope="gpu")
+        tl.atomic_xchg(sync_ptr + NS + 1, 0, sem="relaxed", scope="gpu")
+
+
+@triton.jit
+def _front_absorb_job(
+    a,
+    w_lo,
+    w_hi,
+    q_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    rows,
+    rmask,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    BN: tl.constexpr,
+    BC: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """absorbed_query's ``_absorb_job`` with the trigger after the ``q`` tile loads."""
+    chunks_per_head: tl.constexpr = KV_LORA // BC
+    h = a // chunks_per_head
+    c0 = (a % chunks_per_head) * BC
+    cc = tl.arange(0, BC)
+    d_lo = tl.arange(0, 128)
+    d_hi = tl.arange(0, NOPE - 128)
+    ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu")
+    while ready < NOPE // BN:
+        ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu")
+    tl.debug_barrier()
+    q_base = q_ptr + rows[:, None] * (H * HEAD_DIM) + h * HEAD_DIM
+    q_lo = tl.load(q_base + d_lo[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg")
+    q_hi = tl.load(
+        q_base + 128 + d_hi[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    )
+    _trigger(EARLY_TRIGGER)
+    acc = tl.dot(q_lo, w_lo)
+    acc = tl.dot(q_hi, w_hi, acc)
+    tl.store(
+        fq_ptr + rows[:, None] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[None, :],
+        acc.to(tl.bfloat16),
+        mask=rmask[:, None],
+    )
+    _stress_delay(a, 3, DELAY_NS)
+    tl.debug_barrier()
+    done = tl.atomic_add(sync_ptr + H + h, 1, sem="acq_rel", scope="gpu")
+    if done == chunks_per_head - 1:
+        tl.atomic_xchg(sync_ptr + h, 0, sem="relaxed", scope="gpu")
+        tl.atomic_xchg(sync_ptr + H + h, 0, sem="relaxed", scope="gpu")
+
+
+@triton.jit
+def _front_qpost_program(
+    i,
+    iq_ptr,
+    isync_ptr,
+    kw_ptr,
+    stride_kw,
+    pos_ptr,
+    cos_sin_ptr,
+    qfp8_ptr,
+    qscale_ptr,
+    wout_ptr,
+    w_scale,
+    rows,
+    rmask,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    BN: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """absorbed_query's ``_qpost_program``; positions and cos/sin are read after the wait."""
+    half: tl.constexpr = ROPE_I // 2
+    nope: tl.constexpr = IHD - ROPE_I
+    j = tl.arange(0, half)
+    c = tl.arange(0, nope)
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+    pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64)
+    cs_row = cos_sin_ptr + pos[:, None] * ROPE_I
+    cos = tl.load(cs_row + j[None, :], mask=rmask[:, None], other=0.0)
+    sin = tl.load(cs_row + half + j[None, :], mask=rmask[:, None], other=0.0)
+    w = tl.load(kw_ptr + rows * stride_kw + IHD + i, mask=rmask, other=0.0).to(tl.float32)
+
+    ready = tl.atomic_add(isync_ptr + i, 0, sem="acquire", scope="gpu")
+    while ready < IHD // BN:
+        ready = tl.atomic_add(isync_ptr + i, 0, sem="acquire", scope="gpu")
+    tl.debug_barrier()
+
+    base = iq_ptr + rows[:, None] * (IH * IHD) + i * IHD
+    xe = tl.load(base + 2 * j[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg").to(
+        tl.float32
+    )
+    xo = tl.load(
+        base + 2 * j[None, :] + 1, mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    ).to(tl.float32)
+    xn = tl.load(
+        base + ROPE_I + c[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg"
+    ).to(tl.float32)
+    _trigger(EARLY_TRIGGER)
+    oe, oo = _rope_pairs(xe, xo, cos, sin)
+    amax = tl.maximum(
+        tl.maximum(tl.max(tl.abs(oe), axis=1), tl.max(tl.abs(oo), axis=1)),
+        tl.max(tl.abs(xn), axis=1),
+    )
+    scale, inv = _fp8_ue8m0_scale(amax)
+    out = qfp8_ptr + rows[:, None] * (IH * IHD) + i * IHD
+    tl.store(out + 2 * j[None, :], (oe * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + 2 * j[None, :] + 1, (oo * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + ROPE_I + c[None, :], (xn * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(qscale_ptr + rows * IH + i, scale, mask=rmask)
+    tl.store(wout_ptr + rows * IH + i, w * scale * w_scale, mask=rmask)
+    # Single consumer per head: reset for the next launch.
+    _stress_delay(i, 7, DELAY_NS)
+    tl.atomic_xchg(isync_ptr + i, 0, sem="relaxed", scope="gpu")
+
+
+@triton.jit
+def _front_kpost_program(
+    kln_ptr,
+    stride_kln,
+    pos_ptr,
+    cos_sin_ptr,
+    kfp8_ptr,
+    kscale_ptr,
+    rows,
+    rmask,
+    IHD: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+):
+    """absorbed_query's ``_kpost_program``; positions and cos/sin are read after the wait."""
+    half: tl.constexpr = ROPE_I // 2
+    nope: tl.constexpr = IHD - ROPE_I
+    j = tl.arange(0, half)
+    c = tl.arange(0, nope)
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+    pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64)
+    cs_row = cos_sin_ptr + pos[:, None] * ROPE_I
+    cos = tl.load(cs_row + j[None, :], mask=rmask[:, None], other=0.0)
+    sin = tl.load(cs_row + half + j[None, :], mask=rmask[:, None], other=0.0)
+    base = kln_ptr + rows[:, None] * stride_kln
+    ye = tl.load(base + 2 * j[None, :], mask=rmask[:, None], other=0.0).to(tl.float32)
+    yo = tl.load(base + 2 * j[None, :] + 1, mask=rmask[:, None], other=0.0).to(tl.float32)
+    yn = tl.load(base + ROPE_I + c[None, :], mask=rmask[:, None], other=0.0).to(tl.float32)
+    _trigger(EARLY_TRIGGER)
+    oe, oo = _rope_pairs(ye, yo, cos, sin)
+    amax = tl.maximum(
+        tl.maximum(tl.max(tl.abs(oe), axis=1), tl.max(tl.abs(oo), axis=1)),
+        tl.max(tl.abs(yn), axis=1),
+    )
+    scale, inv = _fp8_ue8m0_scale(amax)
+    out = kfp8_ptr + rows[:, None] * IHD
+    tl.store(out + 2 * j[None, :], (oe * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + 2 * j[None, :] + 1, (oo * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(out + ROPE_I + c[None, :], (yn * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None])
+    tl.store(kscale_ptr + rows, scale, mask=rmask)
+
+
+@triton.jit
+def _front_kv_program(
+    kva_ptr,
+    stride_kva,
+    gkv_ptr,
+    latent_ptr,
+    rows,
+    rmask,
+    eps,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+):
+    """absorbed_query's ``_kv_program`` (native ``kv_a_layernorm``, bit for bit, plus ``k_pe``)."""
+    rc = tl.arange(0, ROPE)
+    k_pe = tl.load(
+        kva_ptr + rows[:, None] * stride_kva + Q_LORA + KV_LORA + rc[None, :],
+        mask=rmask[:, None],
+        other=0.0,
+    )
+    _front_rmsnorm_rows(
+        kva_ptr + Q_LORA,
+        stride_kva,
+        gkv_ptr,
+        latent_ptr,
+        KV_LORA + ROPE,
+        rows,
+        rmask,
+        eps,
+        KV_LORA,
+        BLOCK_M,
+        EARLY_TRIGGER,
+    )
+    tl.store(
+        latent_ptr + rows[:, None] * (KV_LORA + ROPE) + KV_LORA + rc[None, :],
+        k_pe,
+        mask=rmask[:, None],
+    )
+
+
+@triton.jit
+def front_qproj_indexer_kernel(
+    kva_ptr,
+    stride_kva,
+    gq_ptr,
+    gkv_ptr,
+    wqb_ptr,
+    wuk_ptr,
+    q_ptr,
+    latent_ptr,
+    fq_ptr,
+    stride_fq_row,
+    sync_ptr,
+    qn_ptr,
+    n_norm,
+    R,
+    eps_q,
+    eps_kv,
+    wiq_ptr,
+    iq_ptr,
+    kw_ptr,
+    stride_kw,
+    kln_ptr,
+    stride_kln,
+    pos_ptr,
+    cos_sin_ptr,
+    qfp8_ptr,
+    qscale_ptr,
+    kfp8_ptr,
+    kscale_ptr,
+    wout_ptr,
+    w_scale,
+    Q_LORA: tl.constexpr,
+    KV_LORA: tl.constexpr,
+    ROPE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    H: tl.constexpr,
+    IHD: tl.constexpr,
+    IH: tl.constexpr,
+    ROPE_I: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    BC: tl.constexpr,
+    N_QB: tl.constexpr,
+    N_IQ: tl.constexpr,
+    N_ABS: tl.constexpr,
+    DO_ABSORB: tl.constexpr,
+    NORM_ROWS: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    EARLY_TRIGGER: tl.constexpr,
+    ABSORB_IN_PRODUCER: tl.constexpr = False,
+    PRODUCER_JOB_SHIFT: tl.constexpr = 0,
+    ABS_JOB0: tl.constexpr = 0,
+    NORM_MODE: tl.constexpr = 0,
+    N_QPOST: tl.constexpr = 0,
+    PREFETCH_W: tl.constexpr = False,
+    WRITE_QN: tl.constexpr = False,
+    DELAY_NS: tl.constexpr = 0,
+):
+    """See the module docstring. ``n_norm`` (runtime) = ceil(R / NORM_ROWS) norm programs
+    with ``NORM_MODE`` 2, else 0."""
+    NS: tl.constexpr = 2 * H + IH
+    pid = tl.program_id(0)
+    rows = tl.arange(0, BLOCK_M)
+    rmask = rows < R
+    t = pid - n_norm
+    if pid < n_norm:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        nrows = pid * NORM_ROWS + tl.arange(0, NORM_ROWS)
+        _front_rmsnorm_rows(
+            kva_ptr,
+            stride_kva,
+            gq_ptr,
+            qn_ptr,
+            Q_LORA,
+            nrows,
+            nrows < R,
+            eps_q,
+            Q_LORA,
+            NORM_ROWS,
+            EARLY_TRIGGER,
+        )
+        _stress_delay(pid, 11, DELAY_NS)
+        tl.debug_barrier()
+        tl.atomic_add(sync_ptr + NS, 1, sem="release", scope="gpu")
+    elif t < N_QB:
+        _prefetch_w_tile(t, wqb_ptr, Q_LORA, BN, PREFETCH_W)
+        if ABSORB_IN_PRODUCER:
+            job = tl.maximum(t - PRODUCER_JOB_SHIFT, 0)
+            w_lo, w_hi = _ld_wuk(job, t >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC)
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        if NORM_MODE == 1:
+            _front_norm_gemm(
+                t,
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                wqb_ptr,
+                q_ptr,
+                H * HEAD_DIM,
+                rows,
+                rmask,
+                eps_q,
+                qn_ptr,
+                Q_LORA,
+                BLOCK_M,
+                BN,
+                BK,
+                WRITE_QN,
+                False,
+            )
+        else:
+            if NORM_MODE == 2:
+                _wait_norm(sync_ptr, NS, n_norm, N_QB + N_IQ)
+            _shared_query_gemm(
+                t,
+                qn_ptr,
+                Q_LORA,
+                wqb_ptr,
+                q_ptr,
+                H * HEAD_DIM,
+                rows,
+                rmask,
+                Q_LORA,
+                BLOCK_M,
+                BN,
+                BK,
+            )
+        if DO_ABSORB:
+            if (t % (HEAD_DIM // BN)) < (NOPE // BN):
+                _stress_delay(t, 1, DELAY_NS)
+                # Publish this no-RoPE tile of head t // (HEAD_DIM // BN).
+                tl.debug_barrier()
+                tl.atomic_add(sync_ptr + t // (HEAD_DIM // BN), 1, sem="release", scope="gpu")
+        if ABSORB_IN_PRODUCER:
+            if t >= PRODUCER_JOB_SHIFT:
+                _front_absorb_job(
+                    job,
+                    w_lo,
+                    w_hi,
+                    q_ptr,
+                    fq_ptr,
+                    stride_fq_row,
+                    sync_ptr,
+                    rows,
+                    rmask,
+                    KV_LORA,
+                    ROPE,
+                    HEAD_DIM,
+                    NOPE,
+                    H,
+                    BN,
+                    BC,
+                    EARLY_TRIGGER,
+                    DELAY_NS,
+                )
+            else:
+                _trigger(EARLY_TRIGGER)
+        else:
+            _trigger(EARLY_TRIGGER)
+    elif t < N_QB + N_IQ:
+        it = t - N_QB
+        _prefetch_w_tile(it, wiq_ptr, Q_LORA, BN, PREFETCH_W)
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        if NORM_MODE == 1:
+            # Triggers between its last load and its store.
+            _front_norm_gemm(
+                it,
+                kva_ptr,
+                stride_kva,
+                gq_ptr,
+                wiq_ptr,
+                iq_ptr,
+                IH * IHD,
+                rows,
+                rmask,
+                eps_q,
+                qn_ptr,
+                Q_LORA,
+                BLOCK_M,
+                BN,
+                BK,
+                False,
+                EARLY_TRIGGER,
+            )
+        else:
+            if NORM_MODE == 2:
+                _wait_norm(sync_ptr, NS, n_norm, N_QB + N_IQ)
+            _shared_query_gemm(
+                it, qn_ptr, Q_LORA, wiq_ptr, iq_ptr, IH * IHD, rows, rmask, Q_LORA, BLOCK_M, BN, BK
+            )
+            _trigger(EARLY_TRIGGER)
+        _stress_delay(it, 5, DELAY_NS)
+        tl.debug_barrier()
+        tl.atomic_add(sync_ptr + 2 * H + it // (IHD // BN), 1, sem="release", scope="gpu")
+    elif t < N_QB + N_IQ + N_ABS:
+        a = ABS_JOB0 + t - N_QB - N_IQ
+        w_lo, w_hi = _ld_wuk(a, True, wuk_ptr, KV_LORA, NOPE, BC)
+        # Weights are never written by an earlier grid; only the counter and
+        # buffer accesses wait for the programmatic dependency.
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        _front_absorb_job(
+            a,
+            w_lo,
+            w_hi,
+            q_ptr,
+            fq_ptr,
+            stride_fq_row,
+            sync_ptr,
+            rows,
+            rmask,
+            KV_LORA,
+            ROPE,
+            HEAD_DIM,
+            NOPE,
+            H,
+            BN,
+            BC,
+            EARLY_TRIGGER,
+            DELAY_NS,
+        )
+    elif t < N_QB + N_IQ + N_ABS + N_QPOST:
+        _front_qpost_program(
+            t - N_QB - N_IQ - N_ABS,
+            iq_ptr,
+            sync_ptr + 2 * H,
+            kw_ptr,
+            stride_kw,
+            pos_ptr,
+            cos_sin_ptr,
+            qfp8_ptr,
+            qscale_ptr,
+            wout_ptr,
+            w_scale,
+            rows,
+            rmask,
+            IHD,
+            IH,
+            ROPE_I,
+            BN,
+            USE_PDL,
+            EARLY_TRIGGER,
+            DELAY_NS,
+        )
+    elif t == N_QB + N_IQ + N_ABS + N_QPOST:
+        if USE_PDL:
+            tl.extra.cuda.gdc_wait()
+        _front_kv_program(
+            kva_ptr,
+            stride_kva,
+            gkv_ptr,
+            latent_ptr,
+            rows,
+            rmask,
+            eps_kv,
+            Q_LORA,
+            KV_LORA,
+            ROPE,
+            BLOCK_M,
+            EARLY_TRIGGER,
+        )
+    else:
+        _front_kpost_program(
+            kln_ptr,
+            stride_kln,
+            pos_ptr,
+            cos_sin_ptr,
+            kfp8_ptr,
+            kscale_ptr,
+            rows,
+            rmask,
+            IHD,
+            ROPE_I,
+            USE_PDL,
+            EARLY_TRIGGER,
+        )
+
+
+def grid_size(n_norm: int, n_qb: int, n_iq: int, n_abs: int, ih: int) -> int:
+    """Programs of one launch (see the role list in the module docstring)."""
+    return n_norm + n_qb + n_iq + n_abs + ih + 2
+
+
+def num_norm_programs(rows: int, norm_rows: int) -> int:
+    return triton.cdiv(rows, norm_rows)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/query_proj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/query_proj.py
new file mode 100644
index 000000000000..dfcb0883b07d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/query_proj.py
@@ -0,0 +1,403 @@
+# 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.
+"""indexer_layer_front entry into the absorbed_query / draft_absorbed_query path of full-indexer layers.
+
+``dsa_proj`` serves the full-indexer target layers (absorbed_query's domain) and the
+one-model MTP draft layer (draft_absorbed_query's domain). It runs the native bridges absorbed_query
+runs there: the kv_a GEMM (``kv_a_proj_with_mqa``) and, with
+``IndexerFrontConfig.norm="native"`` (default), the native ``q_a_layernorm`` on the
+current stream and, with absorbed_query's ``indexer_k_on_aux`` and multi-stream enabled,
+the BF16 ``wk|weights_proj`` GEMM and the native ``k_norm`` on the indexer's
+aux stream, joined before the kernel. With ``norm="inline"`` or
+``"programs"`` the native ``q_a_layernorm`` launch is dropped and the indexer_layer_front
+kernel (``kernels.front_qproj_indexer_kernel``) normalizes q_a itself. Every
+mode writes the same query_front packet as absorbed_query's PRENORM path, byte for byte:
+``q``, ``latent`` (``compressed_kv`` / ``k_pe`` views), ``fused_q[..., :512]``
+(preabsorbed form) and the indexer outputs ``q_fp8``, ``k_fp8``, ``k_scale``,
+``weights``, ``q_scale``.
+
+Every decline happens before any GPU work and returns ``None``; the adapter
+then runs absorbed_query's / draft_absorbed_query's entry unchanged. The decision is a static check
+of replicated module state (identical on every rank) plus absorbed_query's per-call check.
+
+Option ``skip_reused_iq`` (``IndexerFrontQueryProjection.skip_reused_iq``, off
+in the family's targets; tests set it on a runner) is a declared ABI change: on an MTP
+draft step whose indexer reuses the shared top-k (the indexer's own ``reuse_topk`` condition:
+``mtp_index_share``, ``in_mtp_draft_loop``, ``indexer_skip_topk`` and shared
+rows present), the kernel skips the indexer ``wq_b`` producers and the indexer
+query post-processing, leaving ``q_fp8``, ``q_scale`` and ``weights`` of the
+packet unwritten (the reusing indexer reads none of them). Every other output is
+unchanged.
+"""
+
+from collections import Counter
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+
+from ...dispatch import is_enabled, log_capability
+from ..absorbed_query.query_proj import (
+    _HEAD_DIM,
+    _KV_LORA,
+    _LOCAL_HEADS,
+    _NOPE,
+    _Q_LORA,
+    _ROPE,
+    AbsorbedQueryPacket,
+    AbsorbedQueryProjection,
+    KernelConfig,
+    absorb_for,
+    absorb_layout,
+    call_eligibility,
+    check_package_selectors,
+    module_eligibility,
+)
+from . import kernels
+
+FEATURE_ID = "indexer_layer_front"
+# Host-side decisions of skip_reused_iq modules ("skip" / "full"), for tests and decode boots.
+SKIP_DECISIONS: Counter = Counter()
+
+
+# Where q_a is normalized (kernel constexpr NORM_MODE).
+NORM_MODES = {"native": 0, "inline": 1, "programs": 2}
+
+
+@dataclass(frozen=True)
+class IndexerFrontConfig:
+    # "native": the native q_a_layernorm launch feeds the kernel (absorbed_query's PRENORM input);
+    # "inline": every q_b / indexer wq_b producer normalizes its own K tiles;
+    # "programs": norm programs publish qn to the producers through a counter.
+    norm: str = "native"
+    # Rows per norm program (norm="programs"; FlashInfer uses one warp per row).
+    norm_rows: int = 2
+    # L2 prefetch of each producer's weight tile before the grid-dependency wait.
+    prefetch_weights: bool = False
+    # Trigger the dependents (the generation rope) right after each program's last global load.
+    early_trigger: bool = True
+    # Tests only: skew producers / consumers by up to this many nanoseconds.
+    debug_delay_ns: int = 0
+
+
+def config_for(mode: str, **kwargs) -> IndexerFrontConfig:
+    """Tests / benches: ``native``, ``inline`` or ``nr<k>`` (norm programs of ``k`` rows), with
+    an optional ``+pf`` suffix (weight L2 prefetch)."""
+    base, _, suffix = mode.partition("+")
+    if suffix not in ("", "pf"):
+        raise ValueError(f"unknown indexer_layer_front mode suffix {mode!r}")
+    kwargs.setdefault("prefetch_weights", suffix == "pf")
+    if base in ("native", "inline"):
+        return IndexerFrontConfig(norm=base, **kwargs)
+    if base.startswith("nr"):
+        return IndexerFrontConfig(norm="programs", norm_rows=int(base[2:]), **kwargs)
+    raise ValueError(f"unknown indexer_layer_front norm mode {mode!r}")
+
+
+def reuses_shared_topk(mla, attn_metadata) -> bool:
+    """The indexer's ``reuse_topk`` condition of this call (host state only)."""
+    indexer = mla.mqa.indexer
+    return bool(
+        getattr(indexer, "mtp_index_share", False)
+        and getattr(attn_metadata, "in_mtp_draft_loop", False)
+        and getattr(attn_metadata, "indexer_skip_topk", False)
+        and getattr(attn_metadata, "shared_topk_indices", None) is not None
+    )
+
+
+def _native_q_norm_reproduced(mla) -> Tuple[bool, str]:
+    from ..absorbed_query.shared_indexer import _native_rmsnorm_is_flashinfer_cute
+
+    ok, reason = _native_rmsnorm_is_flashinfer_cute(mla.q_a_layernorm, _Q_LORA)
+    return ok, (f"q_a_layernorm: {reason}" if not ok else "")
+
+
+def front_module_eligibility(mla) -> Tuple[bool, str, bool]:
+    """Static per-module check; returns ``(ok, reason, draft)``. No GPU work."""
+    num_layers = getattr(mla, "num_hidden_layers", None)
+    draft = num_layers is not None and mla.layer_idx is not None and mla.layer_idx >= num_layers
+    if draft and not is_enabled("draft_absorbed_query"):
+        return False, "draft layer without draft_absorbed_query", draft
+    if not draft and not is_enabled("absorbed_query"):
+        return False, "target layer without absorbed_query", draft
+    ok, reason = module_eligibility(mla, draft=draft)
+    if not ok:
+        return False, reason, draft
+    if mla.mqa.indexer is None:
+        return (
+            False,
+            "shared-indexer layer (no q_a RMSNorm launch to remove; shared_indexer_absorbed_query / "
+            "absorbed_query serve it)",
+            draft,
+        )
+    ok, reason = _native_q_norm_reproduced(mla)
+    if not ok:
+        return False, reason, draft
+    return True, "", draft
+
+
+class IndexerFrontQueryProjection(AbsorbedQueryProjection):
+    """absorbed_query's full-layer projection with the indexer_layer_front query kernel
+    (``IndexerFrontConfig.norm`` selects the q_a norm)."""
+
+    def __init__(
+        self,
+        mla,
+        config: IndexerFrontConfig = IndexerFrontConfig(),
+        kernel_config: KernelConfig = KernelConfig(),
+    ):
+        # kernel_config: absorbed_query's tiling and options (tests pass debug options).
+        super().__init__(mla, kernel_config)
+        self.front_config = config
+        # skip_reused_iq on this module (off in this target; tests set it on MTP draft layers).
+        self.skip_reused_iq = False
+        indexer = mla.mqa.indexer
+        # [2 * H absorption, IH indexer-head, norm released, norm consumers].
+        self.indexer_front_sync = torch.zeros(
+            2 * _LOCAL_HEADS + indexer.n_heads + 2,
+            dtype=torch.int32,
+            device=mla.q_b_proj.weight.device,
+        )
+
+    def supported(self) -> Optional[str]:
+        cfg = self.config
+        if cfg.prefetch or cfg.own_hidden_gemv or not cfg.prenorm:
+            return "absorbed_query kernel configuration is not the PRENORM pipelined default"
+        if not self.indexer_fusable():
+            return "absorbed_query fused indexer kernel not eligible"
+        if self.front_config.norm not in NORM_MODES:
+            return f"norm mode {self.front_config.norm!r} unknown"
+        if self.front_config.norm == "programs" and (
+            self.front_config.norm_rows not in (1, 2, 4, 8, 16)
+            or self.front_config.norm_rows > cfg.block_m
+        ):
+            return f"norm_rows {self.front_config.norm_rows} unsupported"
+        chunks = cfg.bn * _Q_LORA * 2 // 1024  # 1 KiB L2 prefetches per producer weight tile
+        if self.front_config.prefetch_weights and (
+            chunks * 1024 != cfg.bn * _Q_LORA * 2 or chunks & (chunks - 1)
+        ):
+            return f"weight tile of {cfg.bn} x {_Q_LORA} is not a power-of-two number of KiB"
+        return None
+
+    def bridges(self, hidden_states: torch.Tensor):
+        """absorbed_query's ``native_bridges``; the native ``q_a_layernorm`` only with ``norm="native"``.
+
+        Returns ``(kva, qn, kw, k_ln)``: the native kv_a GEMM output, the
+        native ``q_a_layernorm`` of its q_a columns (None unless ``native``),
+        the native BF16 fused ``wk|weights_proj`` output and the native
+        ``k_norm`` of its key columns (the latter two on the indexer's aux
+        stream when absorbed_query's ``indexer_k_on_aux`` is set and multi-stream is
+        enabled, joined before returning).
+        """
+        mla = self.mla
+        indexer = mla.mqa.indexer
+        native = self.front_config.norm == "native"
+
+        def main():
+            kva = mla.kv_a_proj_with_mqa(hidden_states)
+            return kva, (self.native_q_norm(kva) if native else None)
+
+        def aux():
+            kw = torch.nn.functional.linear(hidden_states, indexer._fused_wk_wp_weight)
+            return kw, indexer.k_norm(kw[:, : indexer.head_dim])
+
+        if self.config.indexer_k_on_aux:
+            from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel
+
+            (kva, qn), (kw, k_ln) = maybe_execute_in_parallel(
+                main, aux, self.events[0], self.events[1], indexer.aux_stream
+            )
+        else:
+            kva, qn = main()
+            kw, k_ln = aux()
+        return kva, qn, kw, k_ln
+
+    def launch_front(
+        self,
+        kva: torch.Tensor,
+        kw: torch.Tensor,
+        k_ln: torch.Tensor,
+        position_ids: torch.Tensor,
+        absorb: bool,
+        early_trigger: Optional[bool] = None,
+        qn: Optional[torch.Tensor] = None,
+        skip_iq: bool = False,
+    ) -> AbsorbedQueryPacket:
+        """The indexer_layer_front kernel on the bridge outputs (same packet as ``launch_indexer`` with PRENORM).
+
+        ``qn``: the native ``q_a_layernorm`` output (``norm="native"`` only).
+        ``skip_iq``: ``skip_reused_iq`` (``q_fp8``, ``q_scale``, ``weights`` left unwritten).
+        """
+        mla = self.mla
+        indexer = mla.mqa.indexer
+        cfg = self.config
+        rows = kva.shape[0]
+        assert (
+            kva.stride(1) == 1 and kw.stride(1) == 1 and k_ln.stride(1) == 1 and rows <= cfg.block_m
+        )
+        ih, ihd = indexer.n_heads, indexer.head_dim
+        q = kva.new_empty(rows, _LOCAL_HEADS * _HEAD_DIM)
+        mode = NORM_MODES[self.front_config.norm]
+        keep_qn = cfg.debug_keep_intermediates
+        if mode == 0:
+            assert qn is not None and qn.shape == (rows, _Q_LORA) and qn.is_contiguous()
+        elif mode == 1:
+            # Formed in registers; stored only for tests.
+            qn = kva.new_empty(rows, _Q_LORA) if keep_qn else kva
+        else:
+            qn = kva.new_empty(rows, _Q_LORA)
+        latent = kva.new_empty(rows, _KV_LORA + _ROPE)
+        fused_q = kva.new_empty(rows, _LOCAL_HEADS, _KV_LORA + _ROPE) if absorb else None
+        iq = kva.new_empty(rows, ih * ihd)
+        q_fp8 = torch.empty(rows, ih, ihd, dtype=torch.float8_e4m3fn, device=kva.device)
+        q_scale = torch.empty(rows, ih, 1, dtype=torch.float32, device=kva.device)
+        k_fp8 = torch.empty(rows, ihd, dtype=torch.float8_e4m3fn, device=kva.device)
+        k_scale = torch.empty(rows, 1, dtype=torch.float32, device=kva.device)
+        weights = torch.empty(rows, ih, dtype=torch.float32, device=kva.device)
+        pos = position_ids.reshape(-1)
+        assert pos.numel() >= rows
+        n_qb = (_LOCAL_HEADS * _HEAD_DIM) // cfg.bn
+        n_iq = 0 if skip_iq else (ih * ihd) // cfg.bn
+        n_qpost = 0 if skip_iq else ih
+        in_producer = absorb and cfg.absorb_in_producer
+        layout = absorb_layout(n_qb, absorb, in_producer, cfg.absorb_prev_head, cfg.bc)
+        n_norm = kernels.num_norm_programs(rows, self.front_config.norm_rows) if mode == 2 else 0
+        grid = (kernels.grid_size(n_norm, n_qb, n_iq, layout["N_ABS"], n_qpost),)
+        trigger = self.front_config.early_trigger if early_trigger is None else early_trigger
+        kernels.front_qproj_indexer_kernel[grid](
+            kva,
+            kva.stride(0),
+            mla.q_a_layernorm.weight,
+            mla.kv_a_layernorm.weight,
+            mla.q_b_proj.weight,
+            mla.k_b_proj_trans,
+            q,
+            latent,
+            fused_q if absorb else q,
+            fused_q.stride(0) if absorb else 0,
+            self.indexer_front_sync,
+            qn,
+            n_norm,
+            rows,
+            mla.q_a_layernorm.variance_epsilon,
+            mla.kv_a_layernorm.variance_epsilon,
+            indexer.wq_b.weight,
+            iq,
+            kw,
+            kw.stride(0),
+            k_ln,
+            k_ln.stride(0),
+            pos,
+            indexer.rotary_emb.rotary_cos_sin,
+            q_fp8,
+            q_scale,
+            k_fp8,
+            k_scale,
+            weights,
+            indexer.weight_scale_factor,
+            Q_LORA=_Q_LORA,
+            KV_LORA=_KV_LORA,
+            ROPE=_ROPE,
+            HEAD_DIM=_HEAD_DIM,
+            NOPE=_NOPE,
+            H=_LOCAL_HEADS,
+            IHD=ihd,
+            IH=ih,
+            ROPE_I=indexer.rope_dim,
+            BLOCK_M=cfg.block_m,
+            BN=cfg.bn,
+            BK=cfg.bk,
+            N_QB=n_qb,
+            N_IQ=n_iq,
+            N_QPOST=n_qpost,
+            DO_ABSORB=absorb,
+            NORM_ROWS=self.front_config.norm_rows,
+            USE_PDL=self.use_pdl,
+            EARLY_TRIGGER=trigger,
+            ABSORB_IN_PRODUCER=in_producer,
+            NORM_MODE=mode,
+            PREFETCH_W=self.front_config.prefetch_weights,
+            WRITE_QN=mode == 1 and keep_qn,
+            DELAY_NS=self.front_config.debug_delay_ns,
+            **layout,
+            num_warps=cfg.num_warps,
+            num_stages=cfg.num_stages,
+            launch_pdl=self.use_pdl,
+        )
+        if cfg.debug_keep_intermediates:
+            self.debug["iq"] = iq
+            self.debug["qn"] = qn
+        outs = [
+            q,
+            latent[:, :_KV_LORA],
+            latent[:, _KV_LORA:],
+            latent,
+            q_fp8,
+            k_fp8,
+            k_scale,
+            weights,
+            q_scale,
+        ]
+        return AbsorbedQueryPacket(outs, fused_q)
+
+    def project_front(
+        self,
+        position_ids: torch.Tensor,
+        hidden_states: torch.Tensor,
+        absorb: bool,
+        skip_iq: bool = False,
+    ) -> AbsorbedQueryPacket:
+        kva, qn, kw, k_ln = self.bridges(hidden_states)
+        return self.launch_front(kva, kw, k_ln, position_ids, absorb, qn=qn, skip_iq=skip_iq)
+
+
+def _runner(mla):
+    """Per-module ``IndexerFrontQueryProjection`` (cached in ``mla._indexer_front_projection``) or False."""
+    runner = getattr(mla, "_indexer_front_projection", None)
+    if runner is not None:
+        return runner
+    check_package_selectors()
+    ok, reason, draft = front_module_eligibility(mla)
+    runner = False
+    if ok:
+        proj = IndexerFrontQueryProjection(mla)
+        why = proj.supported()
+        if why is not None:
+            ok, reason = False, why
+        else:
+            runner = proj
+    log_capability(FEATURE_ID, f"{'draft' if draft else 'target'}_layer{mla.layer_idx}", ok, reason)
+    mla._indexer_front_projection = runner
+    return runner
+
+
+def dsa_proj(
+    mla, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata
+) -> Optional[AbsorbedQueryPacket]:
+    """indexer_layer_front: the query_front packet of a full-indexer target layer (absorbed_query)
+    or the MTP draft layer (draft_absorbed_query), or ``None`` (before any GPU work) to run the
+    next path."""
+    if not is_enabled(FEATURE_ID):
+        return None
+    runner = _runner(mla)
+    if runner is False:
+        return None
+    ok, _ = call_eligibility(mla, hidden_states, attn_metadata)
+    if not ok or position_ids is None:
+        return None
+    skip_iq = runner.skip_reused_iq and reuses_shared_topk(mla, attn_metadata)
+    if runner.skip_reused_iq:
+        SKIP_DECISIONS["skip" if skip_iq else "full"] += 1
+    return runner.project_front(position_ids, hidden_states, absorb_for(mla), skip_iq=skip_iq)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/rope.py
new file mode 100644
index 000000000000..3370e27ca923
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/rope.py
@@ -0,0 +1,110 @@
+# 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.
+"""indexer_layer_front generation rope: rope_kv_append's fused rope with the trigger after the last read.
+
+``trtllm::glm_moe_dsa_indexer_layer_rope_kv_append``
+(``cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu``) takes the
+arguments of ``trtllm::glm_moe_dsa_rope_kv_append`` and writes the same bytes (FP8 query,
+latent and indexer K caches, FMHA scheduler buffers). Its kernel executes
+``griddepcontrol.launch_dependents`` right after its last global read instead
+of at its end (a block counts as triggered at its first thread's
+invocation), so the next kernel starts while the rope's stores are in
+flight: the indexer_layer_front CP logits kernel on full-indexer target layers and draft
+step 0, sparse_mla_decode's sparse MLA decode on shared-indexer layers, rope_kv_append's top-k
+bridge on draft steps 1-4, or their native fallbacks. Every such dependent reads the
+rope's outputs only after its own ``griddepcontrol.wait`` and writes no
+global memory before it, as rope_kv_append's dependents already must.
+
+The caller (the adapter's rope_kv_append branch of ``_mla_gen_rope``) checks rope_kv_append's
+``rope_eligible`` first and uses this op only when it is available
+(:func:`rope_available`); otherwise rope_kv_append's op runs.
+"""
+
+from typing import List, Optional
+
+import torch
+
+from ..rope_kv_append.ops import IndexerKScatter, _consume
+
+OP_NAME = "glm_moe_dsa_indexer_layer_rope_kv_append"
+
+
+def rope_available() -> bool:
+    """Whether this process has the indexer_layer_front rope op (built in, or a development build loaded)."""
+    return hasattr(torch.ops.trtllm, OP_NAME)
+
+
+def mla_rope_generation(
+    mqa,
+    fused_q,
+    q_pe,
+    latent_cache,
+    metadata,
+    cu_q_seqlens,
+    cu_kv_seqlens,
+    fmha_scheduler_counter,
+    mla_bmm1_scale,
+    mla_bmm2_scale,
+    quant_q_buffer,
+    indexer_k: Optional[IndexerKScatter] = None,
+    early_trigger: bool = True,
+) -> None:
+    """rope_kv_append's ``mla_rope_generation`` (same arguments and bytes) with the indexer_layer_front kernel."""
+    kvm = metadata.kv_cache_manager
+    mqa._ensure_rope_table_size(metadata.max_seq_len)
+    idx_args: List[Optional[torch.Tensor]] = [None, None, None, None, None]
+    idx_tokens = 0
+    if indexer_k is not None:
+        _consume(indexer_k, "front_rope")
+        if indexer_k.num_tokens > 0:
+            idx_args = [
+                indexer_k.k_fp8,
+                indexer_k.k_scale,
+                indexer_k.k_cache,
+                indexer_k.slot_mapping_fp8,
+                indexer_k.slot_mapping_scale,
+            ]
+            idx_tokens = indexer_k.num_tokens
+    getattr(torch.ops.trtllm, OP_NAME)(
+        fused_q,
+        q_pe,
+        latent_cache,
+        mqa.rotary_cos_sin,
+        cu_q_seqlens,
+        cu_kv_seqlens,
+        fmha_scheduler_counter,
+        mla_bmm1_scale,
+        mla_bmm2_scale,
+        quant_q_buffer,
+        metadata.kv_lens_cuda_runtime,
+        metadata.prompt_lens_cpu_runtime,
+        metadata.num_contexts,
+        metadata.kv_cache_block_offsets,
+        kvm.kv_cache_pool_pointers,
+        kvm.kv_cache_pool_mapping,
+        getattr(mqa, "_nvfp4_mla_kv_scale_orig_quant", None),
+        mqa.get_local_layer_idx(metadata),
+        mqa.num_heads,
+        kvm.tokens_per_block,
+        metadata.max_seq_len,
+        mqa.quant_mode,
+        mqa.q_scaling,
+        mqa.qk_nope_head_dim,
+        mqa.qk_rope_head_dim,
+        mqa.kv_lora_rank,
+        *idx_args,
+        idx_tokens,
+        early_trigger,
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/__init__.py
new file mode 100644
index 000000000000..4052d6f2035b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/__init__.py
@@ -0,0 +1,45 @@
+# 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.
+"""indexer_topk: indexer logits and top-k selection (the index_selection contract), GEN target full layers.
+
+Registered entry points (``dispatch.get_impl("indexer_topk", name)``):
+
+- ``indexer_decode``: :func:`indexer_decode.try_indexer_decode`, the adapter
+  hook for the generation branch of ``Indexer.sparse_attn_indexer``.
+- ``decode_topk``: :func:`indexer_decode.decode_topk`, the tensor-level entry
+  used by tests and microbenchmarks.
+- ``cp_indexer_decode``: :func:`cross_rank_indexer_decode.try_cp_indexer_decode`, the
+  adapter hook of the cross_rank_topk context-partitioned variant (selected with
+  ``cross_rank_topk``).
+- ``cp_decode_topk``: :func:`cross_rank_indexer_decode.cp_decode_topk`, its
+  tensor-level entry.
+"""
+
+from ...dispatch import register
+from .cross_rank_indexer_decode import cp_decode_topk, try_cp_indexer_decode
+from .indexer_decode import decode_topk, paged_mqa_logits, try_indexer_decode
+
+register("indexer_topk", "indexer_decode", try_indexer_decode)
+register("indexer_topk", "decode_topk", decode_topk)
+register("indexer_topk", "cp_indexer_decode", try_cp_indexer_decode)
+register("indexer_topk", "cp_decode_topk", cp_decode_topk)
+
+__all__ = [
+    "cp_decode_topk",
+    "decode_topk",
+    "paged_mqa_logits",
+    "try_cp_indexer_decode",
+    "try_indexer_decode",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_indexer_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_indexer_decode.py
new file mode 100644
index 000000000000..d0094adbdcb9
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_indexer_decode.py
@@ -0,0 +1,266 @@
+# 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.
+"""cross_rank_topk: context-partitioned GEN indexer logits + top-k over the TP group.
+
+The indexer projections are replicated, so every GEN rank computes the same
+indexer queries and weights and keeps its own copy of the indexer K cache.
+The ranks split the scan instead of repeating it: rank ``r`` scores the
+64-token pages ``p`` with ``p % world == r`` for every causal verify row
+(``fp8_paged_mqa_logits_cross_rank``, logits bit-identical to the native kernel's
+for the same K), and the ``glm_moe_dsa_cross_rank_topk`` op selects each rank's local
+candidates, exchanges them through cross_rank_topk's own multicast Lamport workspace and
+merges them into the tie-aware exact top-``index_topk`` of every row,
+identically on every rank. Rows with at most ``index_topk`` valid tokens get
+the identity selection.
+
+Numerical status: the result is the exact top-``index_topk`` over the K of
+each page's owner rank, the same set and order on every rank. It equals each
+rank's native selection up to equal keys only when the ranks' indexer K
+caches are bitwise identical. That is the precondition of stock equivalence;
+it is not checked at run time (a check would read every page). Without it
+(one rank storing a different FP8 value for a token), native ranks already
+disagree with each other and cross_rank_topk differs from that rank's native selection in
+the affected rows, so cross_rank_topk is reported as a rank-consistent variant of the
+native selection for quality sign-off.
+
+Cache append, cross-layer reuse, the MTP last-accepted-row stash and the GVR
+state stay with the native caller, exactly as for indexer_topk.
+"""
+
+from typing import Dict
+
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.logger import logger
+
+from . import cross_rank_policy
+
+__all__ = [
+    "CP_DECISION_COUNTS",
+    "cp_decode_topk",
+    "cp_topk_from_local_logits",
+    "try_cp_indexer_decode",
+]
+
+# MTP draft step 0 runs the same 6-row region on the MTP layer (its own
+# indexer K cache, swapped into the metadata by the draft KV-cache context);
+# cross_rank_topk serves it as well. Draft steps 1-4 reuse the step-0 selection and never
+# reach the hook; a draft step with another verify width is declined by the
+# policy.
+ALLOW_DRAFT_STEP = True
+
+CP_DECISION_COUNTS: Dict[str, int] = {}
+
+
+def _record_decision(decision) -> None:
+    key = "engaged" if decision.eligible else decision.reason
+    count = CP_DECISION_COUNTS.get(key, 0)
+    CP_DECISION_COUNTS[key] = count + 1
+    if count == 0:
+        state = "engaged" if decision.eligible else f"native fallback ({decision.reason})"
+        logger.info(
+            f"[glm_moe_dsa] cross_rank_topk/cp_indexer_decode: first '{state}' decision on this rank"
+        )
+
+
+def cp_topk_from_local_logits(
+    local_logits: torch.Tensor,
+    kv_lens: torch.Tensor,
+    out: torch.Tensor,
+    workspace,
+    next_n: int,
+) -> torch.Tensor:
+    """Select, exchange and merge from this rank's local logits.
+
+    Args:
+        local_logits: ``[next_n, local_capacity]`` float32 logits of this
+            rank's local columns (see ``cross_rank_logits_runner.local_to_global``).
+        kv_lens: ``[1]`` int32 global indexer KV length.
+        out: ``[next_n, index_topk]`` int32 output.
+        workspace: ``cross_rank_workspace.CpWorkspace`` of the group.
+        next_n: Verify rows of the request.
+
+    Returns:
+        ``out`` holding the same global top-k selection on every rank.
+    """
+    torch.ops.trtllm.glm_moe_dsa_cross_rank_topk(
+        local_logits,
+        kv_lens,
+        out,
+        workspace.state,
+        workspace.mc_ptr,
+        workspace.uc_ptr,
+        workspace.rank,
+        workspace.world,
+        next_n,
+    )
+    return out
+
+
+def cp_decode_topk(
+    *,
+    q: torch.Tensor,
+    weights: torch.Tensor,
+    k_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    kv_lens: torch.Tensor,
+    max_seq_len: int,
+    out: torch.Tensor,
+    workspace,
+) -> torch.Tensor:
+    """Run the cross_rank_topk region for one request.
+
+    Args:
+        q: ``[1, next_n, H, D]`` float8_e4m3fn indexer queries.
+        weights: ``[next_n, H]`` float32 per-head weights.
+        k_cache: ``[num_blocks, 64, 1, D + 4]`` uint8 indexer K page view.
+        block_table: ``[1, max_blocks]`` int32 page table.
+        kv_lens: ``[1]`` int32 global indexer KV length.
+        max_seq_len: Capture-stable indexer max sequence length.
+        out: ``[next_n, index_topk]`` int32 output.
+        workspace: ``cross_rank_workspace.CpWorkspace`` of the group.
+
+    Returns:
+        ``out``.
+    """
+    from .cross_rank_logits_runner import paged_mqa_logits_cp
+
+    local_logits = paged_mqa_logits_cp(
+        q, k_cache, weights, kv_lens, block_table, max_seq_len, workspace.rank, workspace.world
+    )
+    return cp_topk_from_local_logits(local_logits, kv_lens, out, workspace, q.shape[1])
+
+
+def try_cp_indexer_decode(
+    indexer,
+    metadata,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+) -> bool:
+    """Adapter entry: run the GEN region of ``Indexer.sparse_attn_indexer``.
+
+    The decision reads only host state, tensor metadata and the group's
+    workspace (created collectively on the first eligible eager call), all
+    before any device work, so a ``False`` return leaves no side effect and
+    the caller runs the native path. Every rank of the group reaches the same
+    decision for the same forward.
+
+    Args/Returns: as ``indexer_decode.try_indexer_decode``.
+    """
+    from ...dispatch import log_capability
+    from .cross_rank_logits_runner import ensure_compiled, local_capacity
+    from .cross_rank_workspace import failure_reason, get_workspace, op_info
+
+    num_contexts = metadata.num_contexts
+    num_generations = metadata.num_generations
+    gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
+    uniform = num_generations > 0 and bool(
+        (gen_seq_lens == num_gen_tokens // num_generations).all()
+    )
+    top_k = indexer.top_k
+    k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(indexer.layer_idx)
+    block_table = metadata.indexer_k_cache_block_offsets
+    mapping = getattr(metadata, "mapping", None)
+    max_seq_len = metadata.get_indexer_max_seq_len()
+    world = mapping.tp_size if mapping is not None else 1
+    info = op_info(world) if world > 0 else None
+    capacity = local_capacity(max_seq_len, world) if world > 0 else 0
+    # Facts of this rank's build and device: agreed by the workspace creation
+    # below instead of deciding this call on this rank alone.
+    local_reason = cross_rank_policy.rank_local_reason(
+        op_available=info is not None,
+        local_capacity=capacity,
+        max_local_capacity=info[7] if info is not None else 0,
+        sm_version=get_sm_version(),
+        cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE,
+    )
+    decision = cross_rank_policy.decide_cp_uniform(
+        world_size=world,
+        local_capacity=capacity,
+        allow_draft_step=ALLOW_DRAFT_STEP,
+        native_uses_cute_dsl_logits=indexer.use_cute_dsl_paged_mqa_logits,
+        use_fp4_indexer=indexer.use_fp4,
+        use_gvr_emission=indexer.use_gvr_emission,
+        needs_gvr_prior=top_k.needs_gvr_prior,
+        is_ragged_verify=metadata.is_ragged_verify,
+        expanded_mtp_buffers=metadata.use_expanded_buffers_for_mtp,
+        in_mtp_draft_loop=metadata.in_mtp_draft_loop,
+        attention_dp=bool(mapping is not None and mapping.enable_attention_dp),
+        topk_decode_impl=top_k.decode_implementation.value,
+        topk_self_sampling=top_k.gvr_self_sampling,
+        index_topk=top_k.top_k,
+        compress_ratio=top_k.compress_ratio,
+        num_heads=indexer.n_heads,
+        head_dim=indexer.head_dim,
+        query_dtype=str(q_fp8.dtype),
+        query_shape=tuple(q_fp8.shape),
+        query_contiguous=q_fp8.is_contiguous(),
+        weights_dtype=str(weights.dtype),
+        weights_shape=tuple(weights.shape),
+        weights_contiguous=weights.is_contiguous(),
+        k_cache_dtype=str(k_cache.dtype),
+        k_cache_shape=tuple(k_cache.shape),
+        block_table_cols=block_table.shape[1],
+        max_seq_len=max_seq_len,
+        num_contexts=num_contexts,
+        token_offset=token_offset,
+        num_gen_tokens=num_gen_tokens,
+        num_generations=num_generations,
+        uniform_verify_width=uniform,
+    )
+    if decision.eligible:
+        # The logits kernel is compiled before the group agrees on the
+        # workspace, so no rank waits in an exchange while a peer compiles.
+        workspace = get_workspace(
+            mapping,
+            prepare=lambda: ensure_compiled(
+                q_fp8.device,
+                k_cache.shape[1],
+                indexer.n_heads,
+                indexer.head_dim,
+                num_gen_tokens // num_generations,
+                world,
+            ),
+            local_reason=local_reason,
+        )
+        if workspace is None:
+            reason = failure_reason(mapping) or "workspace not created before CUDA-graph capture"
+            decision = type(decision)(False, f"workspace unavailable: {reason}")
+    log_capability("cross_rank_topk", "cp_indexer_decode", decision.eligible, decision.reason)
+    _record_decision(decision)
+    if not decision.eligible:
+        return False
+
+    next_n = num_gen_tokens // num_generations
+    gen_rows = slice(token_offset, token_offset + num_gen_tokens)
+    gen_reqs = slice(num_contexts, num_contexts + num_generations)
+    kv_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+    assert kv_lens is not None
+    cp_decode_topk(
+        q=q_fp8[gen_rows].view(num_generations, next_n, indexer.n_heads, indexer.head_dim),
+        weights=weights[gen_rows],
+        k_cache=k_cache,
+        block_table=block_table[gen_reqs],
+        kv_lens=kv_lens,
+        max_seq_len=max_seq_len,
+        out=out,
+        workspace=workspace,
+    )
+    return True
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_logits_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_logits_runner.py
new file mode 100644
index 000000000000..aac939edad4d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_logits_runner.py
@@ -0,0 +1,219 @@
+# 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.
+"""Compile cache and launcher of the cross_rank_topk context-partitioned logits kernel.
+
+The kernel (``fp8_paged_mqa_logits_cross_rank``) scores the 64-token pages of one
+rank of the tensor-parallel group for every verify row of one request, with
+the indexer_topk ``rows`` configuration (8 register-cached weights per row, fp32
+accumulation, epilogue and output).
+"""
+
+from typing import Dict, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+
+from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers
+
+from .fp8_paged_mqa_logits_cross_rank import FP8MQALogitsCPKernel
+
+COMPUTE_BLOCK_KV = 128
+# Two math warpgroups per CTA each take one 128-token tile per step.
+SPLIT_KV = 2 * COMPUTE_BLOCK_KV
+PAGE_TOKENS = 64
+LOGITS_BUFFER_NAME = "cross_rank_topk_logits"
+# indexer_topk default ("rows" kernel, 8 register-cached weights per verify row).
+NUM_W_IN_REG = 8
+# The kernel lets the dependent select/exchange op launch early (programmatic dependent launch).
+TRIGGER_DEPENDENTS = True
+
+_KERNEL_CACHE: Dict[Tuple, object] = {}
+_NUM_SMS: Dict[int, int] = {}
+# The kernel derives its schedule in-kernel; this placeholder fills the
+# schedule argument of the shared launcher signature and is never read.
+_UNUSED_SCHEDULE: Dict[int, torch.Tensor] = {}
+
+
+def _num_sms(device: torch.device) -> int:
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    if index not in _NUM_SMS:
+        _NUM_SMS[index] = torch.cuda.get_device_properties(index).multi_processor_count
+    return _NUM_SMS[index]
+
+
+def _unused_schedule(device: torch.device) -> torch.Tensor:
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    t = _UNUSED_SCHEDULE.get(index)
+    if t is None:
+        t = torch.zeros(1, 2, dtype=torch.int32, device=device)
+        _UNUSED_SCHEDULE[index] = t
+    return t
+
+
+def local_length(n: int, rank: int, world: int) -> int:
+    """Tokens of ``rank`` among the global tokens ``[0, n)`` (page-interleaved)."""
+    if n <= 0:
+        return 0
+    full_pages, tail = divmod(n, PAGE_TOKENS)
+    length = (full_pages - rank + world - 1) // world * PAGE_TOKENS
+    if tail and full_pages % world == rank:
+        length += tail
+    return length
+
+
+def local_capacity(max_seq_len: int, world: int) -> int:
+    """Local logits columns that cover any rank's share of ``max_seq_len`` tokens."""
+    return max(local_length(max_seq_len, r, world) for r in range(world))
+
+
+def local_to_global(local_cols: torch.Tensor, rank: int, world: int) -> torch.Tensor:
+    """Global token index of local logits columns of ``rank``."""
+    return (local_cols // PAGE_TOKENS * world + rank) * PAGE_TOKENS + local_cols % PAGE_TOKENS
+
+
+def _compile(
+    phys_block_kv: int, num_heads: int, head_dim: int, next_n: int, num_sms: int, world: int
+):
+    n_cols = next_n * num_heads
+    block_bytes = phys_block_kv * (head_dim + 4)
+    sym_num_phys_blocks = cute.sym_int()
+    sym_b = cute.sym_int()
+    max_ctx = cute.sym_int()
+    max_blocks_per_seq = cute.sym_int()
+    num_ctas = cute.sym_int()
+    kv_fake = cute.runtime.make_fake_tensor(
+        cutlass.Uint8, (sym_num_phys_blocks, block_bytes), stride=(cute.sym_int64(), 1)
+    )
+    q_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Uint8, (n_cols, head_dim, sym_b), stride_order=(1, 0, 2)
+    )
+    w_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float32, (n_cols, sym_b), stride_order=(0, 1)
+    )
+    logits_fake = cute.runtime.make_fake_tensor(
+        cutlass.Float32, (cute.sym_int(), max_ctx), stride=(cute.sym_int64(), 1)
+    )
+    bt_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Int32, (sym_b, max_blocks_per_seq), stride_order=(1, 0)
+    )
+    cl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,))
+    sm_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Int32, (num_ctas, 2), stride_order=(1, 0)
+    )
+    fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    kernel = FP8MQALogitsCPKernel(
+        block_kv=COMPUTE_BLOCK_KV,
+        phys_block_kv=phys_block_kv,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        next_n=next_n,
+        num_sms=num_sms,
+        num_w_in_reg=NUM_W_IN_REG,
+        world=world,
+        trigger_dependents=TRIGGER_DEPENDENTS,
+    )
+    return cute.compile(
+        kernel,
+        kv_fake,
+        q_fake,
+        w_fake,
+        logits_fake,
+        bt_fake,
+        cl_fake,
+        sm_fake,
+        cutlass.Int32(1),
+        cutlass.Int32(1),
+        cutlass.Int32(0),
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+def ensure_compiled(
+    device: torch.device, phys_block_kv: int, num_heads: int, head_dim: int, next_n: int, world: int
+) -> None:
+    """Compile the kernel for this shape now (host only, no device work).
+
+    cross_rank_topk calls this before the group agrees to create its workspace, so the
+    per-rank compile time is absorbed by that collective instead of making
+    the first exchange wait for the slowest compiling rank.
+    """
+    key = (phys_block_kv, num_heads, head_dim, next_n, _num_sms(device), world)
+    if key not in _KERNEL_CACHE:
+        _KERNEL_CACHE[key] = _compile(*key)
+
+
+def paged_mqa_logits_cp(
+    q: torch.Tensor,
+    kv_fused: torch.Tensor,
+    weights: torch.Tensor,
+    kv_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    max_seq_len: int,
+    rank: int,
+    world: int,
+) -> torch.Tensor:
+    """Score this rank's pages of one request for every verify row.
+
+    Args:
+        q: ``[1, next_n, H, D]`` float8_e4m3fn queries.
+        kv_fused: ``[num_blocks, 64, 1, D + 4]`` uint8 indexer K page view.
+        weights: ``[next_n, H]`` float32 epilogue weights.
+        kv_lens: ``[1]`` int32 global indexer KV length.
+        block_table: ``[1, max_blocks]`` int32 page table of the request.
+        max_seq_len: Capture-stable indexer max sequence length.
+        rank: This rank's index in the context-parallel group.
+        world: Ranks in the group.
+
+    Returns:
+        ``[next_n, local_capacity(max_seq_len, world)]`` float32 view of a
+        graph-stable arena; local column ``c`` holds global token
+        ``local_to_global(c)``. Columns at or beyond a row's local length are
+        unspecified.
+    """
+    batch, next_n, num_heads, head_dim = q.shape
+    assert batch == 1
+    phys_block_kv = kv_fused.shape[1]
+    num_phys_blocks = kv_fused.shape[0]
+    num_sms = _num_sms(q.device)
+    width = local_capacity(max_seq_len, world)
+    aligned = (width + SPLIT_KV - 1) // SPLIT_KV * SPLIT_KV
+    logits = get_memory_buffers().get_buffer(
+        [next_n, aligned],
+        torch.float32,
+        buffer_name=LOGITS_BUFFER_NAME,
+        reserve_buffer=torch.cuda.is_current_stream_capturing(),
+    )[:, :width]
+    key = (phys_block_kv, num_heads, head_dim, next_n, num_sms, world)
+    compiled = _KERNEL_CACHE.get(key)
+    if compiled is None:
+        compiled = _compile(*key)
+        _KERNEL_CACHE[key] = compiled
+    q_3d = q.reshape(batch, next_n * num_heads, head_dim).permute(1, 2, 0)
+    w_2d = weights.reshape(batch, next_n * num_heads).t()
+    compiled(
+        kv_fused.reshape(num_phys_blocks, -1),
+        q_3d.view(torch.uint8),
+        w_2d,
+        logits,
+        block_table,
+        kv_lens,
+        _unused_schedule(q.device),
+        num_phys_blocks,
+        batch,
+        rank,
+    )
+    return logits
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_policy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_policy.py
new file mode 100644
index 000000000000..cc56efa4c6bc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_policy.py
@@ -0,0 +1,143 @@
+# 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.
+"""Eligibility policy of the cross_rank_topk context-partitioned indexer region.
+
+Stdlib only, like ``policy``. cross_rank_topk runs inside the indexer_topk envelope (same GPU,
+indexer shape, page format, verify width, selection engine and
+``max_seq_len`` bound; see ``policy.decide``) and additionally needs:
+
+- exactly one generation request (the in-kernel schedule and the exchange
+  slots are per request);
+- a tensor-parallel group of a GPU-verified size, without attention DP;
+- the compiled ``glm_moe_dsa_cross_rank_topk`` op and a local logits width it accepts.
+
+The decision of a forward (``decide_cp_uniform``) reads only inputs that are
+identical on all ranks of the group for the same forward, so the ranks make
+the same decision. The facts of a rank's build and device (SM version, CuTe
+DSL runtime, compiled op, the kernel's local width limit; ``rank_local_reason``)
+are not decided per rank: the workspace creation that follows agrees on them
+collectively (``cross_rank_workspace``), so no rank declines on its own while its
+peers enter the exchange.
+"""
+
+from typing import Tuple
+
+from . import policy
+
+# Tensor-parallel group sizes the GPU tests verify (one node with 4 GPUs, and
+# a GEN group of 8 GPUs on two hosts).
+SUPPORTED_WORLD_SIZES: Tuple[int, ...] = (4, 8)
+
+
+def decide_cp(
+    *,
+    world_size: int,
+    op_available: bool,
+    local_capacity: int,
+    max_local_capacity: int,
+    allow_draft_step: bool,
+    **indexer_topk_inputs,
+) -> policy.Decision:
+    """Decide whether cross_rank_topk may run the GEN indexer logits + top-k region.
+
+    Args:
+        world_size: Tensor-parallel group size of the GEN ranks.
+        op_available: Whether ``torch.ops.trtllm.glm_moe_dsa_cross_rank_topk`` exists.
+        local_capacity: Local logits columns the group needs at the indexer
+            ``max_seq_len`` (largest per-rank page share).
+        max_local_capacity: Largest local width the compiled kernel accepts.
+        allow_draft_step: Whether an MTP draft step may engage cross_rank_topk.
+        **indexer_topk_inputs: Every keyword input of ``policy.decide``.
+
+    Returns:
+        The decision with a fallback reason when ineligible.
+    """
+    in_draft = bool(indexer_topk_inputs.get("in_mtp_draft_loop", False))
+    if in_draft and not allow_draft_step:
+        return policy.Decision(False, "MTP draft step (target layers only)")
+    base = policy.decide(**dict(indexer_topk_inputs, in_mtp_draft_loop=False))
+    if not base.eligible:
+        return base
+    if indexer_topk_inputs["num_generations"] != 1:
+        return policy.Decision(
+            False, f"{indexer_topk_inputs['num_generations']} generation requests != 1"
+        )
+    if world_size not in SUPPORTED_WORLD_SIZES:
+        return policy.Decision(
+            False, f"tensor-parallel size {world_size} not in {SUPPORTED_WORLD_SIZES}"
+        )
+    if not op_available:
+        return policy.Decision(False, "glm_moe_dsa_cross_rank_topk op is not built")
+    if local_capacity <= 0 or local_capacity > max_local_capacity:
+        return policy.Decision(
+            False, f"local logits width {local_capacity} > kernel limit {max_local_capacity}"
+        )
+    return policy.Decision(True)
+
+
+def rank_local_reason(
+    *,
+    op_available: bool,
+    local_capacity: int,
+    max_local_capacity: int,
+    sm_version: int,
+    cute_dsl_available: bool,
+) -> str:
+    """Why this rank's build or device cannot run cross_rank_topk ("" when it can).
+
+    These are the inputs of ``decide_cp`` that are facts of one rank rather
+    than of the forward. The caller hands the reason to the collective
+    workspace creation, which keeps every rank native when any rank has one.
+
+    Args:
+        op_available: Whether ``torch.ops.trtllm.glm_moe_dsa_cross_rank_topk`` exists.
+        local_capacity: Local logits columns the group needs (rank-uniform).
+        max_local_capacity: Largest local width this rank's compiled kernel
+            accepts (derived from the device's shared memory).
+        sm_version: This rank's compute capability, e.g. 103.
+        cute_dsl_available: Whether the CUTLASS DSL runtime is importable.
+    """
+    if sm_version not in policy.SUPPORTED_SM_VERSIONS:
+        return f"sm{sm_version} unsupported"
+    if not cute_dsl_available:
+        return "CUTLASS DSL unavailable"
+    if not op_available:
+        return "glm_moe_dsa_cross_rank_topk op is not built"
+    if local_capacity > max_local_capacity:
+        return f"local logits width {local_capacity} > kernel limit {max_local_capacity}"
+    return ""
+
+
+def decide_cp_uniform(
+    *, world_size: int, local_capacity: int, allow_draft_step: bool, **indexer_topk_inputs
+) -> policy.Decision:
+    """``decide_cp`` on the rank-uniform inputs of a forward.
+
+    The rank-local facts (``rank_local_reason``) count as capable here; they
+    are agreed by the workspace creation. ``indexer_topk_inputs`` are the keyword inputs
+    of ``policy.decide`` other than ``sm_version`` and ``cute_dsl_available``.
+    """
+    return decide_cp(
+        world_size=world_size,
+        op_available=True,
+        local_capacity=local_capacity,
+        max_local_capacity=max(local_capacity, 0),
+        allow_draft_step=allow_draft_step,
+        **dict(
+            indexer_topk_inputs,
+            sm_version=policy.SUPPORTED_SM_VERSIONS[0],
+            cute_dsl_available=True,
+        ),
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_workspace.py
new file mode 100644
index 000000000000..8e86b3686fcc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/cross_rank_workspace.py
@@ -0,0 +1,212 @@
+# 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.
+"""Exchange workspace of the cross_rank_topk context-partitioned indexer top-k.
+
+cross_rank_topk owns one multicast workspace per tensor-parallel group, allocated with the
+native ``McastGPUBuffer`` (the class behind the MNNVL all-reduce workspaces)
+and used by nothing else. Its layout and Lamport protocol are defined by the
+``glm_moe_dsa_cross_rank_topk`` kernel (``crossRankTopk.h``): every 32-bit word starts as
+``0xFFFFFFFF`` and the kernel rotates three buffers with a device epoch held
+in a separate per-rank state tensor.
+
+Creation is collective over the group and only happens outside CUDA-graph
+capture. Each rank first runs the caller's host-side preparation (cross_rank_topk
+compiles its logits kernel there, so compile-time skew between ranks is
+absorbed by the first agreement rather than by the first exchange), then the
+ranks agree on the capability (every rank must have the compiled op, a
+supported world size, MPI, a successful preparation and none of the rank-local
+reasons of ``cross_rank_policy.rank_local_reason``: SM version, CuTe DSL runtime and
+the kernel's local width limit), allocate, and agree
+on success; if any step fails on any rank, every rank records the failure and
+keeps the native path, so the collective count and order stay identical.
+"""
+
+import os
+import threading
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.logger import logger
+
+# Indices of torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_workspace_info().
+_INFO_WORKSPACE_WORDS = 0
+_INFO_TOPK = 1
+_INFO_MAX_RANKS = 2
+_INFO_MAX_ROWS = 3
+_INFO_PAGE_TOKENS = 4
+_INFO_MAX_LOCAL_CAPACITY = 7
+_INFO_STATE_WORDS = 8
+_INFO_TIMEOUT_MS = 10
+
+
+@dataclass
+class CpWorkspace:
+    """One rank's view of the cross_rank_topk exchange workspace.
+
+    Attributes:
+        handle: The native ``McastGPUBuffer`` (keeps the mapping alive).
+        mc: Multicast int32 view of the whole workspace.
+        uc: This rank's unicast int32 view of the whole workspace.
+        state: ``[4]`` int32 device state (epoch = exchange launches, arrivals,
+            second rounds taken, error word written before a timeout abort).
+        rank: Rank in the group.
+        world: Group size.
+        comm: The group communicator the handle was created with (kept alive).
+    """
+
+    handle: object
+    mc: torch.Tensor
+    uc: torch.Tensor
+    state: torch.Tensor
+    rank: int
+    world: int
+    comm: object = None
+
+    @property
+    def mc_ptr(self) -> int:
+        return self.mc.data_ptr()
+
+    @property
+    def uc_ptr(self) -> int:
+        return self.uc.data_ptr()
+
+
+_LOCK = threading.Lock()
+_WORKSPACES: Dict[Tuple, CpWorkspace] = {}
+_FAILED: Dict[Tuple, str] = {}
+
+
+def op_info(world_size: int = 8) -> Optional[list]:
+    """Constants of the compiled kernel for a group size, or None when the op is not built."""
+    try:
+        return list(torch.ops.trtllm.glm_moe_dsa_cross_rank_topk_workspace_info(world_size))
+    except (AttributeError, RuntimeError):
+        return None
+
+
+def _key(mapping) -> Tuple:
+    return (mapping.world_size, mapping.rank, mapping.tp_size, mapping.pp_size, mapping.cp_size)
+
+
+def failure_reason(mapping) -> Optional[str]:
+    """Reason recorded when this group's workspace could not be created."""
+    return _FAILED.get(_key(mapping))
+
+
+def get_workspace(
+    mapping, prepare: Optional[Callable[[], None]] = None, local_reason: str = ""
+) -> Optional[CpWorkspace]:
+    """Return the group's workspace, creating it collectively on first use.
+
+    Must be reached by every rank of the tensor-parallel group in the same
+    order (it is called from the same layer of the same forward on every
+    rank). Returns None when the workspace is unavailable on any rank, and
+    during CUDA-graph capture when it was not created before.
+
+    Args:
+        mapping: The rank's ``Mapping``.
+        prepare: Optional host-side preparation run on every rank before the
+            capability agreement of the creation (not called when the
+            workspace already exists); an exception counts as a local
+            failure that every rank agrees on.
+        local_reason: Why this rank's build or device cannot run cross_rank_topk
+            (``cross_rank_policy.rank_local_reason``; "" when it can). Part of the
+            creation's capability agreement, so it keeps every rank native.
+    """
+    key = _key(mapping)
+    ws = _WORKSPACES.get(key)
+    if ws is not None or key in _FAILED:
+        return ws
+    if torch.cuda.is_current_stream_capturing():
+        return None
+    with _LOCK:
+        if key in _WORKSPACES or key in _FAILED:
+            return _WORKSPACES.get(key)
+        ws, reason = _create(mapping, prepare, local_reason)
+        if ws is None:
+            _FAILED[key] = reason
+            logger.info(f"[glm_moe_dsa] cross_rank_topk/workspace: native fallback ({reason})")
+        else:
+            _WORKSPACES[key] = ws
+            info = op_info(ws.world) or []
+            timeout = info[_INFO_TIMEOUT_MS] if len(info) > _INFO_TIMEOUT_MS else "n/a"
+            logger.info(
+                f"[glm_moe_dsa] cross_rank_topk/workspace: created ({ws.uc.numel() * 4} bytes, "
+                f"rank {ws.rank}/{ws.world}, exchange wait timeout {timeout} ms)"
+            )
+        return ws
+
+
+def _local_capability_reason(mapping, info, prepare) -> str:
+    """This rank's op, group size and preparation check ("" when capable)."""
+    if info is None:
+        return "glm_moe_dsa_cross_rank_topk op is not built"
+    if mapping.tp_size > info[_INFO_MAX_RANKS]:
+        return f"tp_size {mapping.tp_size} > {info[_INFO_MAX_RANKS]}"
+    if prepare is not None:
+        try:
+            prepare()
+        except Exception as e:  # noqa: BLE001 - any local failure must be agreed before use
+            return f"preparation failed ({type(e).__name__}: {e})"
+    return ""
+
+
+def _create(mapping, prepare=None, local_reason="") -> Tuple[Optional[CpWorkspace], str]:
+    from tensorrt_llm._torch.distributed.ops import (
+        _get_mnnvl_workspace_comm,
+        _mnnvl_device_index,
+        _mnnvl_workspace_all_succeeded,
+    )
+    from tensorrt_llm._utils import mpi_disabled
+
+    if mpi_disabled():
+        # Without MPI every rank takes this branch before any collective.
+        return None, "MPI is disabled"
+    comm = _get_mnnvl_workspace_comm(mapping)
+    info = op_info(mapping.tp_size)
+    if not local_reason:
+        local_reason = _local_capability_reason(mapping, info, prepare)
+    if not _mnnvl_workspace_all_succeeded(comm, not local_reason):
+        return None, local_reason or "capability not agreed by every rank"
+
+    words = info[_INFO_WORKSPACE_WORDS]
+    error: Optional[Exception] = None
+    ws: Optional[CpWorkspace] = None
+    try:
+        from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer
+
+        use_fabric = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node()
+        device = _mnnvl_device_index(mapping)
+        handle = McastGPUBuffer(
+            words * 4, mapping.tp_size, mapping.tp_rank, device, use_fabric, comm.py2f()
+        )
+        uc = handle.get_uc_buffer(mapping.tp_rank, [words], torch.int32, 0)
+        mc = handle.get_mc_buffer([words], torch.int32, 0)
+        with torch.inference_mode():
+            uc.fill_(-1)
+            state = torch.zeros(
+                info[_INFO_STATE_WORDS], dtype=torch.int32, device=torch.device("cuda", device)
+            )
+        torch.cuda.synchronize()
+        ws = CpWorkspace(handle, mc, uc, state, mapping.tp_rank, mapping.tp_size, comm)
+    except Exception as e:  # noqa: BLE001 - any local failure must be agreed before use
+        error = e
+    # Also fences the initialization: no rank sends before every rank cleared
+    # its buffer.
+    if not _mnnvl_workspace_all_succeeded(comm, error is None):
+        return None, f"allocation failed on a rank ({type(error).__name__ if error else 'peer'})"
+    return ws, ""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py
new file mode 100644
index 000000000000..f064a7aa5871
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py
@@ -0,0 +1,1806 @@
+# 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.
+"""
+cross_rank_topk context-partitioned variant of the indexer_topk FP8 paged MQA logits kernel
+(Blackwell).
+
+Identical to ``fp8_paged_mqa_logits_rows`` (the shipped indexer_topk all-rows kernel)
+except for which KV tokens one launch scores and how CTAs find their work:
+
+- The request's indexer context is partitioned over ``world`` ranks by whole
+  64-token pages (global page ``p`` belongs to rank ``p % world``). A launch
+  scores only the pages of its runtime ``rank``: local page ``lp`` reads the
+  block-table entry of global page ``lp * world + rank``, and logits column
+  ``c`` holds global token ``((c // 64) * world + rank) * 64 + c % 64``.
+- One request (batch size 1). The local length is derived in-kernel from the
+  global indexer KV length (``context_lens[0]``), and each CTA computes its
+  contiguous share of the local tiles itself; ``schedule_meta`` is not read.
+
+Per-element arithmetic is unchanged, so every local logit is bit-identical to
+the native logit of the same global token.
+
+The notes of the indexer_topk ``rows`` kernel follow.
+
+indexer_topk variant of the CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100).
+
+Derived from ``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
+for the GLM-5.x indexer shape, where all verify rows of a request fit one UMMA
+tile (``N = next_n * num_heads <= 256``). The only functional change is the
+size of the per-thread register cache of epilogue weights
+(``num_w_in_reg`` heads per verify row; the remaining heads are read from
+shared memory). Keeping that cache bounded avoids register spills when
+``N`` is large; weight values, accumulation order and output are unchanged.
+
+Architecture:
+  - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA)
+  - 1 TMA per KV block [128, 128], UMMA iterates 4x K=32
+  - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2)
+  - Q reloaded via TMA pipeline when q_idx (batch) changes
+  - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs
+  - Weights cached in registers: preloaded once per q_idx change (not per KV block)
+  - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes)
+
+Merged KV+Scale pipeline:
+  - KV data and scales share a single TMA barrier per group
+  - TMA loads both KV and Scale under one barrier (combined tx_count)
+  - UMMA waits on merged barrier (for KV GEMM), does NOT release
+  - Math waits on merged barrier (for scale read), Math releases
+
+Fused KV layout:
+  - KV data and scales stored contiguously per physical block:
+    [num_phys_blocks, block_kv * (head_dim + 4)] bytes
+  - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)]
+  - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic
+
+Scheduler:
+  - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups)
+  - schedule_meta[sm_idx+1] = end boundary for this CTA
+  - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx)
+  - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step
+
+Dynamic shape support:
+  - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen
+  - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq,
+    num_ctas) are marked dynamic via mark_compact_shape_dynamic
+  - Allows JIT cache reuse across different batch sizes / sequence lengths
+
+Epilogue dtype flows (--acc_dtype / --epi_dtype):
+
+  Flow 1: --acc_dtype fp32 --epi_dtype fp16
+    Q(FP8) x K(FP8) -> MMA acc(FP32) -> TMEM(FP32)
+      -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16)
+      -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16)
+      -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype)
+
+  Flow 2: --acc_dtype fp16 --epi_dtype fp16
+    Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b)
+      -> LDTM -> Reg(FP16) -> ReLU(FP16)
+      -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16)
+      -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype)
+
+  --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion.
+  Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32
+"""
+
+from typing import Tuple
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass import Float16, Int32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import griddepcontrol_launch_dependents
+
+# CuTe DSL CUDA 13 validates rounding modes as string literals. The string
+# form is also accepted by older wrappers, so keep it version-independent.
+_RND_RN = "rn"
+
+
+@dsl_user_op
+def pack_f16x2(
+    a: Float16,
+    b: Float16,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    f16_ty = Float16.mlir_type
+    i32_ty = Int32.mlir_type
+    vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc)
+    v = vector.from_elements(
+        vec2_f16,
+        (Float16(a).ir_value(loc=loc, ip=ip), Float16(b).ir_value(loc=loc, ip=ip)),
+        loc=loc,
+        ip=ip,
+    )
+    return Int32(llvm.bitcast(i32_ty, v, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def unpack_f16x2(
+    packed: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Tuple[Float16, Float16]:
+    f16_ty = Float16.mlir_type
+    vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc)
+    v = llvm.bitcast(vec2_f16, Int32(packed).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    r0 = Float16(vector.extract(v, dynamic_position=[], static_position=[0], loc=loc, ip=ip))
+    r1 = Float16(vector.extract(v, dynamic_position=[], static_position=[1], loc=loc, ip=ip))
+    return r0, r1
+
+
+@dsl_user_op
+def fma_f16x2(
+    a: Int32,
+    b: Int32,
+    c: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [
+                Int32(a).ir_value(loc=loc, ip=ip),
+                Int32(b).ir_value(loc=loc, ip=ip),
+                Int32(c).ir_value(loc=loc, ip=ip),
+            ],
+            "fma.rn.f16x2 $0, $1, $2, $3;",
+            "=r,r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def max_f16x2(
+    a: Int32,
+    b: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "max.f16x2 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def add_f16x2(
+    a: Int32,
+    b: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "add.f16x2 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+class FP8MQALogitsCPKernel:
+    """FP8 paged MQA logits kernel for Blackwell (SM100).
+
+    Each CTA processes a range of (q_idx, kv_split) pairs.
+    A split = 2 consecutive KV blocks within a sequence (one per warp group).
+    Q is shared between warp groups and reloaded when q_idx changes.
+    """
+
+    def __init__(
+        self,
+        block_kv: int = 128,
+        phys_block_kv: int = 128,
+        num_heads: int = 64,
+        head_dim: int = 128,
+        next_n: int = 1,
+        num_sms: int = 148,
+        remove_kv_wait_in_epilogue: bool = False,
+        early_tmem_copy: bool = False,
+        smem_subpartition_opt: bool = False,
+        max_kv_pipeline: bool = False,
+        max_umma_pipeline: bool = False,
+        num_epi_subtiles: int = 1,
+        epi_dtype=cutlass.Float32,
+        acc_dtype=cutlass.Float32,
+        output_dtype=cutlass.Float32,
+        num_w_in_reg=None,
+        world: int = 8,
+        trigger_dependents: bool = True,
+    ):
+        if block_kv % phys_block_kv != 0 or phys_block_kv != 64:
+            raise ValueError("the context-partitioned kernel needs 64-token pages")
+        self.world = world
+        # Let the dependent grid (cross_rank_topk's select/exchange op, launched with programmatic
+        # stream serialization) launch while this kernel runs: the op waits on griddepcontrol.wait
+        # before it reads anything, so only its launch latency overlaps this kernel.
+        self.trigger_dependents = trigger_dependents
+        self.block_kv = block_kv
+        self.phys_block_kv = phys_block_kv
+        self.num_blocks_per_mma = block_kv // phys_block_kv
+        assert block_kv % phys_block_kv == 0, (
+            f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}"
+        )
+        assert self.num_blocks_per_mma <= 4, (
+            f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4"
+        )
+        self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue
+        self.early_tmem_copy = early_tmem_copy
+        self.smem_subpartition_opt = smem_subpartition_opt
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.next_n = next_n
+        self.N = next_n * num_heads
+        self.num_sms = num_sms
+        self.num_epi_subtiles = num_epi_subtiles
+        self.epi_dtype = epi_dtype
+        self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4
+        # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
+        # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1
+        # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest).
+        w_stage_bytes = self.N * self.epi_bytes
+        self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes
+        self.output_dtype = output_dtype
+        if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0:
+            raise ValueError("num_heads must be divisible by num_epi_subtiles")
+        if (num_heads // num_epi_subtiles) % 4 != 0:
+            raise ValueError(
+                "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)"
+            )
+        self.num_groups = 2
+
+        # Register-cached epilogue weights per verify row. None keeps the
+        # upstream heuristic; otherwise a multiple of 4 (FMA unroll) <= heads.
+        if num_w_in_reg is None:
+            if epi_dtype == cutlass.Float16:
+                max_w_in_reg = 64 if next_n <= 3 else 48
+            else:
+                max_w_in_reg = 64 if next_n == 1 else 40 if next_n >= 4 else 52
+            num_w_in_reg = min(max_w_in_reg, num_heads)
+        if num_w_in_reg < 0 or num_w_in_reg > num_heads or num_w_in_reg % 4 != 0:
+            raise ValueError(
+                f"num_w_in_reg={num_w_in_reg} must be a multiple of 4 in [0, {num_heads}]"
+            )
+        self.num_w_in_reg = num_w_in_reg
+
+        self.num_math_threads = 256
+        self.num_specialized_threads = 128
+        self.threads_per_cta = 384
+        self.num_math_warps = 8
+        self.tma_warp_base = 8
+        self.umma_warp_base = 10
+
+        self.num_q_stages = 3  # 3 stages for Q pipelining across batch sequences
+
+        # TMEM: 512 columns total, each group needs N columns per UMMA stage
+        # max_umma_stages = 512 // (2 * N)
+        TMEM_COLS = 512
+        if max_umma_pipeline:
+            self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N))
+        else:
+            self.num_umma_stages = 1
+
+        if max_kv_pipeline:
+            smem_capacity = utils.get_smem_capacity_in_bytes()
+            # Reserve ~1 KB for barriers and misc
+            SMEM_BUDGET = smem_capacity - 1024
+            # KV+Scale per stage (×2 groups):
+            #   2 * (block_kv * head_dim * 1B + block_kv * 4B)
+            kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4)
+            # Q+W per stage: Q is N * head_dim * 1B, W uses padded stride
+            qw_per_stage = self.N * head_dim + self.w_stage_stride * self.epi_bytes
+            qw_total = qw_per_stage * self.num_q_stages
+            self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage
+        else:
+            self.num_kv_stages = 3
+
+        # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB),
+        # avoiding sub-bank conflicts with UMMA reading sKV.
+        # Layout: barriers(~256B) | sKV_0 | sKV_1 | sQ | [pad] | sW | sScales
+        if self.smem_subpartition_opt:
+            BOUNDARY = 128 * 1024
+            used = (
+                256
+                + 2 * (block_kv * head_dim * self.num_kv_stages)
+                + self.N * head_dim * self.num_q_stages
+            )
+            used = ((used + 127) // 128) * 128
+            if used < BOUNDARY:
+                self.smem_pad_bytes = ((BOUNDARY - used + 1023) // 1024) * 1024
+            else:
+                self.smem_pad_bytes = 0
+        else:
+            self.smem_pad_bytes = 0
+
+        self.acc_dtype = acc_dtype
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.cluster_shape_mn = (1, 1)
+        self.mma_tiler_mn = (block_kv, self.N)
+
+    def _setup_mma(self, a_dtype, b_dtype, a_major, b_major):
+        self.a_dtype = a_dtype
+        self.b_dtype = b_dtype
+        self.a_major_mode = a_major
+        self.b_major_mode = b_major
+
+        self.mma_tiler = (*self.mma_tiler_mn, 1)
+        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            a_dtype,
+            a_major,
+            b_major,
+            self.acc_dtype,
+            self.cta_group,
+            self.mma_tiler_mn,
+        )
+        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])  # 32
+
+        # Full-K: tile K = head_dim (128), 1 TMA per block
+        mma_inst_tile_k = self.head_dim // mma_inst_shape_k  # 4
+        full_k = mma_inst_shape_k * mma_inst_tile_k  # 128
+        self.mma_tiler = (
+            self.mma_tiler_mn[0],
+            self.mma_tiler_mn[1],
+            full_k,
+        )
+
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (tiled_mma.thr_id.shape,),
+        )
+        self.epi_tile = self.cta_tile_shape_mnk[:2]
+
+        # KV SMEM: 3 stages per group, each stage holds full [128, 128]
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            self.mma_tiler,
+            a_dtype,
+            self.num_kv_stages,
+        )
+        # Q SMEM: 1 stage, holds full [N, 128]
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            self.mma_tiler,
+            b_dtype,
+            self.num_q_stages,
+        )
+
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+        self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake)
+        self.num_tmem_alloc_cols_total = (
+            self.num_tmem_alloc_cols * self.num_groups * self.num_umma_stages
+        )
+
+        return tiled_mma
+
+    @cute.jit
+    def __call__(
+        self,
+        kv_fused: cute.Tensor,  # Fused KV: [num_phys_blocks, block_bytes] FP8
+        b: cute.Tensor,  # Q: [N, head_dim, batch_size]
+        weights: cute.Tensor,  # [N, batch_size] (transposed for TMA)
+        logits: cute.Tensor,  # [batch_size * next_n, max_context_len]
+        block_table: cute.Tensor,  # [batch_size, max_blocks_per_seq]
+        context_lens: cute.Tensor,  # [batch_size]
+        schedule_meta: cute.Tensor,  # [num_sms+1, 2] int32
+        num_phys_blocks: cutlass.Int32,
+        batch_size: cutlass.Int32,
+        rank: cutlass.Int32,
+        stream: cuda.CUstream,
+    ):
+        # Derive KV and Scale views from fused buffer using CuTE ops.
+        # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)]
+        phys_block_kv = self.phys_block_kv
+        scale_offset_elems = phys_block_kv * self.head_dim  # in FP8 elements
+
+        # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference)
+        kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN)
+
+        # Q (b) was passed as uint8 to work around DLPack's lack of float8 support;
+        # recast back to FP8 so MMA type inference and TMA descriptors are correct.
+        b = cute.recast_tensor(b, cutlass.Float8E4M3FN)
+
+        # Read the real per-block stride (bytes = FP8 elements) from the input.
+        # When KV is the indexer K-cache pool view, the pool is laid out as
+        # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride =
+        # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes).
+        # Using the input stride keeps both contiguous test path and the
+        # strided prod path correct. FP8 = 1 byte per element, so the byte
+        # stride is the same as the element stride after recast.
+        kv_block_stride = kv_fused.layout.stride[0]
+
+        # KV view: [phys_block_kv, head_dim, num_phys_blocks] FP8
+        # Each TMA loads one physical block; multiple TMAs fill a compute tile.
+        kv_layout = cute.make_layout(
+            (phys_block_kv, self.head_dim, num_phys_blocks),
+            stride=(self.head_dim, 1, kv_block_stride),
+        )
+        a = cute.make_tensor(kv_fp8.iterator, kv_layout)
+
+        # Scale view: offset pointer to scale region, recast FP8 → Float32
+        # [phys_block_kv, num_phys_blocks] float32 (after recast)
+        scale_fp8_layout = cute.make_layout(
+            (phys_block_kv * 4, num_phys_blocks),
+            stride=(1, kv_block_stride),
+        )
+        scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout)
+        scales = cute.recast_tensor(scale_fp8, cutlass.Float32)
+
+        a_dtype = a.element_type
+        b_dtype = b.element_type
+        a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode()
+        b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+
+        tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major)
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        # TMA for KV (A) — fmha_decode_paged pattern.
+        # Build a TMA SMEM layout via tiled_divide on the full compute-tile
+        # layout, then select to drop trivial K dim. Atom uses mode [0] as
+        # single-tile SMEM layout and (phys, head) as cta_tiler.
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+        self.a_tma_view_layout = sm100_utils.make_smem_layout(
+            tcgen05.OperandMajorMode.K,
+            (self.block_kv, self.head_dim),
+            a_dtype,
+            self.num_kv_stages,
+        )
+        self.a_tma_view_layout = cute.tiled_divide(
+            self.a_tma_view_layout, (self.phys_block_kv, self.head_dim)
+        )
+        # ((tile_M, tile_K), rest_M, rest_K, stages) → drop trivial rest_K
+        self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3])
+        # ((tile_M, tile_K), rest_M=num_sub_blocks, stages)
+        tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            a,
+            self.a_tma_view_layout[0],  # atom SMEM = single-tile (mode 0)
+            (self.phys_block_kv, self.head_dim),
+        )
+
+        # TMA for Q (B) — full K=128, L dim = batch_size (unchanged)
+        b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            b_op,
+            b,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # TMA for Weights — [N, batch_size], tile [N], L=batch_size
+        self.w_smem_layout_staged = cute.make_layout(
+            (self.N, self.num_q_stages),
+            stride=(1, self.w_stage_stride),
+        )
+        w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0])
+        tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            weights,
+            w_smem_per_stage,
+            self.w_smem_layout_staged.shape[:1],
+        )
+
+        # TMA for Scales — [phys_block_kv, num_phys_blocks], tile [phys_block_kv]
+        # SMEM holds compute_block_kv scales per stage; filled by
+        # num_blocks_per_mma sub-block TMAs at consecutive offsets.
+        self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages))
+        s_smem_per_subblock = cute.make_layout((phys_block_kv,))
+        tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            scales,
+            s_smem_per_subblock,
+            (phys_block_kv,),
+        )
+
+        b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout)
+        w_copy_size = self.N * self.epi_bytes
+        # Per sub-block: phys_block_kv * head_dim (KV) + phys_block_kv * 4 (scales)
+        kv_tma_bytes_per_subblock = phys_block_kv * self.head_dim
+        scale_tma_bytes_per_subblock = phys_block_kv * 4
+        # Total per compute tile = num_blocks_per_mma sub-blocks
+        self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * (
+            kv_tma_bytes_per_subblock + scale_tma_bytes_per_subblock
+        )
+        # Q + Weights share barrier (like DeepGEMM)
+        self.num_q_tma_bytes = b_copy_size * atom_thr_size + w_copy_size
+
+        num_ctas = self.num_sms
+
+        @cute.struct
+        class SharedStorage:
+            kv_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            kv_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2]
+            umma_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            umma_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            tmem_holding_buf: cutlass.Int32
+
+        self.kernel(
+            tiled_mma,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_w,
+            tma_tensor_w,
+            tma_atom_s,
+            tma_tensor_s,
+            logits,
+            block_table,
+            context_lens,
+            schedule_meta,
+            batch_size,
+            rank,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.w_smem_layout_staged,
+            self.s_smem_layout_staged,
+            self.a_tma_view_layout,
+            self.epi_tile,
+            SharedStorage,
+        ).launch(
+            grid=(1, 1, num_ctas),
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,  # KV pool
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,  # Q (L dim = batch_size)
+        tma_atom_w: cute.CopyAtom,
+        mW_tma: cute.Tensor,  # Weights TMA coord tensor [N, batch_size]
+        tma_atom_s: cute.CopyAtom,
+        mS_tma: cute.Tensor,  # Scales TMA coord tensor [block_kv, num_phys_blocks]
+        mLogits: cute.Tensor,  # [batch_size * next_n, max_context_len]
+        mBlockTable: cute.Tensor,  # [batch_size, max_blocks_per_seq]
+        mContextLens: cute.Tensor,  # [batch_size]
+        mScheduleMeta: cute.Tensor,  # unused (kept for the indexer_topk launcher signature)
+        batch_size: cutlass.Int32,
+        rank: cutlass.Int32,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        w_smem_layout_staged: cute.Layout,
+        s_smem_layout_staged: cute.Layout,
+        a_tma_view_layout: cute.ComposedLayout,
+        epi_tile: cute.Tile,
+        SharedStorage: cutlass.Constexpr,
+    ):
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
+
+        bidx, bidy, bidz = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+        is_leader_cta = mma_tile_coord_v == 0
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        tidx, _, _ = cute.arch.thread_idx()
+        if cutlass.const_expr(self.trigger_dependents):
+            griddepcontrol_launch_dependents()
+
+        # Warp roles (matches DeepGEMM SM100)
+        warpgroup_idx = warp_idx // 4
+        is_math_warp = warp_idx < 8
+        is_tma_warp_0 = warp_idx == 8
+        is_tma_warp_1 = warp_idx == 9
+        is_tma_warp = is_tma_warp_0 | is_tma_warp_1
+        is_umma_warp_0 = warp_idx == 10
+        is_umma_warp_1 = warp_idx == 11
+
+        # Early schedule metadata load: issue global loads ASAP so their
+        # ~200-cycle L2 latency overlaps with subsequent prologue setup
+        # (SMEM alloc, TMA partition, MMA fragment creation, etc.)
+        NUM_MATH_WG = 2  # kNumMathWarpGroups
+        NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma
+        sm_idx = bidz
+        # Context partition (one request): this rank's local tokens are its
+        # pages among the global pages covering the indexer KV length.
+        world = self.world
+        page_tokens = self.phys_block_kv
+        global_kv = mContextLens[0]
+        full_pages = global_kv // page_tokens
+        tail_tokens = global_kv - full_pages * page_tokens
+        local_len = ((full_pages - rank + world - 1) // world) * page_tokens
+        if tail_tokens > 0:
+            if full_pages % world == rank:
+                local_len = local_len + tail_tokens
+        num_local_pages = (local_len + page_tokens - 1) // page_tokens
+        current_num_kv = (local_len + self.block_kv - 1) // self.block_kv
+        # Contiguous share of the tile pairs per CTA (schedule units of
+        # NUM_MATH_WG tiles, as in the DeepGEMM schedule this replaces). A
+        # share that reaches the end of the request ends at (q = 1, kv = 0).
+        num_units = (current_num_kv + NUM_MATH_WG - 1) // NUM_MATH_WG
+        num_ctas_total = self.num_sms
+        start_unit = (sm_idx * num_units) // num_ctas_total
+        end_unit = ((sm_idx + 1) * num_units) // num_ctas_total
+        start_q = cutlass.Int32(0)
+        start_kv_half = start_unit
+        end_q_idx = cutlass.Int32(0)
+        end_kv_half = end_unit
+        if end_unit == num_units:
+            end_q_idx = cutlass.Int32(1)
+            end_kv_half = cutlass.Int32(0)
+        if start_unit == end_unit:
+            start_q = cutlass.Int32(1)
+            start_kv_half = cutlass.Int32(0)
+            end_q_idx = cutlass.Int32(1)
+            end_kv_half = cutlass.Int32(0)
+
+        if is_tma_warp:
+            cpasync.prefetch_descriptor(tma_atom_a)
+            cpasync.prefetch_descriptor(tma_atom_b)
+            cpasync.prefetch_descriptor(tma_atom_w)
+            cpasync.prefetch_descriptor(tma_atom_s)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        block_kv_val = self.block_kv
+        num_heads = self.num_heads
+        next_n = self.next_n
+        num_epi_subtiles = self.num_epi_subtiles
+
+        # === Pipelines ===
+        prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+
+        # Q pipeline: TMA producer → Math consumer (8 math warps)
+        # PipelineTmaAsync: consumer_release uses is_signalling_thread
+        # (lane 0 per warp). 8 math warps × 1 lane-0 = 8 arrives.
+        q_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 8)
+        q_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.q_mbar.data_ptr(),
+            num_stages=self.num_q_stages,
+            producer_group=prod_group,
+            consumer_group=q_cons_group,
+            tx_count=self.num_q_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        q_prod_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_q_stages
+        )
+        # Both Math WGs share the same pipeline state (advance in lockstep)
+        q_cons_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        # UMMA warps observe Q pipeline (wait only, no release)
+        # to ensure Q is in SMEM before GEMM. Critical for UMMA warp 1
+        # since TMA warp 1 only loads KV1 (not Q).
+        q_cons_state_umma_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        q_cons_state_umma_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+
+        # Merged KV+Scale pipelines (per-group, 3 stages each)
+        # Like DeepGEMM: KV data and scales share one barrier.
+        # TMA loads both under one barrier. Math is consumer (releases).
+        # UMMA also waits on this barrier (for KV GEMM) but does NOT release.
+        math_warps_per_group = self.num_math_warps // 2  # 4 warps
+        kv_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, math_warps_per_group)
+        kv_pipeline_0 = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar_0.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=kv_cons_group,
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        kv_pipeline_1 = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar_1.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=kv_cons_group,
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+
+        kv_prod_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        kv_prod_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        # UMMA consumer states (wait only, no release)
+        kv_cons_state_umma_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_umma_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        # Math consumer states (wait + release)
+        kv_cons_state_math_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_math_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+
+        # UMMA pipelines (per-group)
+        math_threads_per_group = self.num_math_threads // 2
+        umma_pipeline_0 = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar_0.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        umma_pipeline_1 = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar_1.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+
+        umma_prod_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_prod_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_cons_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+        umma_cons_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+
+        # TMEM — only Math warps (8×32=256) + UMMA warps (2×32=64) = 320 threads
+        # TMA warps do NOT participate, so they can start TMA loads earlier.
+        # Math warp 0 is the allocator (like fp16_gemm_3's epilogue warp 0),
+        # because math warps are the last TMEM consumers (epilogue reads).
+        tmem_alloc_num_threads = 320  # 10 warps: warp 0-7 (math) + warp 10-11 (umma)
+        tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=tmem_alloc_num_threads)
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=0,  # math warp 0 does alloc+free (last TMEM consumer)
+            is_two_cta=False,
+        )
+
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+
+        # SMEM allocation: per-group KV + shared Q
+        sKV_0 = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sKV_1 = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sQ = smem.allocate_tensor(
+            element_type=self.b_dtype,
+            layout=b_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=b_smem_layout_staged.inner,
+        )
+        # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB)
+        # to avoid sub-bank conflicts with UMMA reading sKV from
+        # sub-partition 0.
+        if cutlass.const_expr(self.smem_pad_bytes > 0):
+            _ = smem.allocate(self.smem_pad_bytes)
+        # Weights SMEM: [N, num_q_stages], shared Q barrier
+        sW = smem.allocate_tensor(
+            element_type=self.epi_dtype,
+            layout=w_smem_layout_staged,
+            byte_alignment=128,
+        )
+        # Scales SMEM: [block_kv, num_kv_stages] float32, per group
+        sScales_0 = smem.allocate_tensor(
+            element_type=cutlass.Float32,
+            layout=s_smem_layout_staged,
+            byte_alignment=128,
+        )
+        sScales_1 = smem.allocate_tensor(
+            element_type=cutlass.Float32,
+            layout=s_smem_layout_staged,
+            byte_alignment=128,
+        )
+
+        a_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+        )
+        b_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+        )
+
+        # Partition KV (A): fmha_decode_paged pattern.
+        # SMEM view is ((tile), num_sub_blocks, stages) — built in __call__.
+        # Use .outer (plain layout); swizzle is captured by sKV_0's iterator.
+        # GMEM: local_tile by (phys, head), then group first 2 modes into tile.
+        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+        sKV_0_for_tma = cute.make_tensor(sKV_0.iterator, a_tma_view_layout.outer)
+        sKV_1_for_tma = cute.make_tensor(sKV_1.iterator, a_tma_view_layout.outer)
+        gA = cute.local_tile(
+            mA_mkl,
+            (self.phys_block_kv, self.head_dim),
+            coord=(None, None, None),
+        )
+        tAsA_0, tAgA_0 = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            sKV_0_for_tma,
+            cute.group_modes(gA, 0, 2),
+        )
+        tAsA_1, tAgA_1 = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            sKV_1_for_tma,
+            cute.group_modes(gA, 0, 2),
+        )
+
+        # Partition Q (B): shared SMEM, L dim = batch_size
+        gB_nkl = cute.local_tile(
+            mB_nkl,
+            cute.slice_(self.mma_tiler, (0, None, None)),
+            (None, None, None),
+        )
+        tCgB = thr_mma.partition_B(gB_nkl)
+        b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            block_in_cluster_coord_vmnk[1],
+            b_cta_layout,
+            cute.group_modes(sQ, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+        tBgB = tBgB[(None, 0, None, None)]  # [tma, K, L]
+
+        # Partition Weights: standalone TMA, [N, batch_size] → [N] per stage
+        w_cta_layout = cute.make_layout((1,))
+        tWsW, tWgW = cpasync.tma_partition(
+            tma_atom_w,
+            0,
+            w_cta_layout,
+            cute.group_modes(sW, 0, 1),
+            cute.group_modes(mW_tma, 0, 1),
+        )
+
+        # Partition Scales: explicit sub_blocks + stages dims.
+        # Layout (phys_block_kv, num_sub, stages) K-major with custom strides.
+        s_tma_view_layout = cute.make_layout(
+            (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages),
+            stride=(1, self.phys_block_kv, self.block_kv),
+        )
+        sScales_0_for_tma = cute.make_tensor(sScales_0.iterator, s_tma_view_layout)
+        sScales_1_for_tma = cute.make_tensor(sScales_1.iterator, s_tma_view_layout)
+        # GMEM: local_tile by phys to match atom's tile size
+        gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None))
+        tSsS_0, tSgS_0 = cpasync.tma_partition(
+            tma_atom_s,
+            0,
+            cute.make_layout(1),
+            sScales_0_for_tma,
+            gS,
+        )
+        tSsS_1, tSgS_1 = cpasync.tma_partition(
+            tma_atom_s,
+            0,
+            cute.make_layout(1),
+            sScales_1_for_tma,
+            gS,
+        )
+
+        # MMA fragments
+        tCrA_0 = tiled_mma.make_fragment_A(sKV_0)
+        tCrA_1 = tiled_mma.make_fragment_A(sKV_1)
+        tCrB = tiled_mma.make_fragment_B(sQ)  # shared
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+
+        # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim
+        # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank
+        us = self.num_umma_stages
+        cols = self.num_tmem_alloc_cols
+        acc_shape_staged = cute.append(acc_shape, us)
+        tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged)
+
+        # TMEM layout info (allocation deferred to UMMA/Math warp branches)
+        cols_per_group = cols * us * (32 // self.acc_dtype.width)
+        num_tmem_alloc_cols_total = self.num_tmem_alloc_cols_total
+
+        # Epilogue setup
+        c_layout = utils.LayoutEnum.ROW_MAJOR
+        epi_sub_mn = (epi_tile[0], num_heads // num_epi_subtiles)
+        copy_atom_t2r = sm100_utils.get_tmem_load_op(
+            self.cta_tile_shape_mnk,
+            c_layout,
+            self.acc_dtype,
+            self.acc_dtype,
+            epi_sub_mn,
+            use_2cta_instrs,
+        )
+
+        # ===== SCHEDULER: derive values from early-loaded schedule metadata =====
+        end_kv_idx = end_kv_half * NUM_MATH_WG
+
+        # Convert start to KV block units (matching DeepGEMM)
+        current_q_idx = start_q
+        current_kv_idx = start_kv_half * NUM_MATH_WG
+
+        # ===== COMMON SCHEDULER STATE (before warp branches) =====
+        # Each warp role independently maintains its own copy of these
+        # variables (like DeepGEMM where each role creates its own scheduler).
+        # Pre-fetch first task (current_num_kv loaded early above for latency hiding)
+        next_q_idx = current_q_idx
+        next_kv_idx = current_kv_idx
+        next_num_kv = current_num_kv
+        # Sentinel: no previous batch (matches DeepGEMM's q_idx = batch_size)
+        q_idx = batch_size
+        # While-loop termination flag (matches DeepGEMM's fetch_next_task pattern).
+        # True if this CTA has work assigned (start != end in schedule_meta).
+        has_work = (current_q_idx != end_q_idx) | (current_kv_idx != end_kv_idx)
+
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+
+        # ===== WARP-SPECIALIZED EXECUTION =====
+
+        if is_tma_warp_0:
+            # TMA warp 0: loads Q (prefetch) + KV for group 0
+            # Matches DeepGEMM's TMA warp with kv_group_idx == 0
+            cute.arch.warpgroup_reg_dealloc(24)
+            lane_idx = tidx % 32
+
+            # Block table prefetch: 32 lanes cache block indices,
+            # distributed via shuffle. Each lane holds num_blocks_per_mma
+            # physical block indices per compute tile. (DeepGEMM L233-244)
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)  # force prefetch on first use
+
+            # Prefetch first Q before loop (like DeepGEMM line 203-204)
+            q_pipeline.producer_acquire(q_prod_state)
+            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, 0, next_q_idx)],
+                tBsB[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+                mcast_mask=b_mcast_mask,
+            )
+            cute.copy(
+                tma_atom_w,
+                tWgW[(None, next_q_idx)],
+                tWsW[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+            )
+            q_prod_state.advance()
+
+            while has_work:
+                # fetch_next_task: commit next → current
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                # Q prefetch: when batch changes, load Q for NEXT batch
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)  # force re-prefetch
+                    prefetch_next = q_idx + 1
+                    if prefetch_next < end_q_idx:
+                        q_pipeline.producer_acquire(q_prod_state)
+                        q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                        cute.copy(
+                            tma_atom_b,
+                            tBgB[(None, 0, prefetch_next)],
+                            tBsB[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                            mcast_mask=b_mcast_mask,
+                        )
+                        cute.copy(
+                            tma_atom_w,
+                            tWgW[(None, prefetch_next)],
+                            tWsW[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                        )
+                        q_prod_state.advance()
+                    elif prefetch_next == end_q_idx:
+                        if end_kv_idx > 0:
+                            q_pipeline.producer_acquire(q_prod_state)
+                            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                            cute.copy(
+                                tma_atom_b,
+                                tBgB[(None, 0, prefetch_next)],
+                                tBsB[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                                mcast_mask=b_mcast_mask,
+                            )
+                            cute.copy(
+                                tma_atom_w,
+                                tWgW[(None, prefetch_next)],
+                                tWsW[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                            )
+                            q_prod_state.advance()
+
+                # Block table prefetch for group 0 (like DeepGEMM L233-241).
+                # Each lane loads num_blocks_per_mma physical block indices
+                # for one compute tile (kv_idx counts compute tiles).
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            local_page = min(base_phys + i, num_local_pages - 1)
+                            cached_blks[i] = mBlockTable[(q_idx, local_page * world + rank)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+
+                # Get block indices via shuffle before barrier (like DeepGEMM L244)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile.
+                kv_pipeline_0.producer_acquire(kv_prod_state_0)
+                bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0)
+                stage = kv_prod_state_0.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_0[(None, 0, 0, phys_blks[i])],
+                        tAsA_0[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_s,
+                        tSgS_0[(None, 0, phys_blks[i])],
+                        tSsS_0[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                    )
+                kv_prod_state_0.advance()
+
+                # Advance: inline fetch_next_task
+                next_kv_idx = kv_idx + NUM_MATH_WG
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                # Update while-loop condition
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_tma_warp_1:
+            # TMA warp 1: loads KV + Scale for group 1 only
+            # Matches DeepGEMM's TMA warp with kv_group_idx == 1
+            cute.arch.warpgroup_reg_dealloc(24)
+            lane_idx = tidx % 32
+
+            # Block table prefetch for group 1
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)  # force prefetch on first use
+
+            while has_work:
+                # fetch_next_task: commit next → current
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                # New q_idx → force block table re-prefetch
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)
+
+                # Block table prefetch for group 1 (like DeepGEMM L233-241).
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            local_page = min(base_phys + i, num_local_pages - 1)
+                            cached_blks[i] = mBlockTable[(q_idx, local_page * world + rank)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+
+                # Get block indices via shuffle before barrier (like DeepGEMM L244)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile.
+                kv_pipeline_1.producer_acquire(kv_prod_state_1)
+                bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1)
+                stage = kv_prod_state_1.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_1[(None, 0, 0, phys_blks[i])],
+                        tAsA_1[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_s,
+                        tSgS_1[(None, 0, phys_blks[i])],
+                        tSsS_1[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                    )
+                kv_prod_state_1.advance()
+
+                # Advance: inline fetch_next_task
+                next_kv_idx = kv_idx + NUM_MATH_WG
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                # Update while-loop condition
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp_0:
+            # UMMA warp for group 0
+            # Must wait on Q pipeline: TMA operations with different
+            # barriers are NOT visibility-ordered even within the same
+            # warp. KV0 barrier arriving does not guarantee Q SMEM
+            # writes are visible.
+            cute.arch.warpgroup_reg_dealloc(24)
+
+            # TMEM: wait for math warp 0's allocation, retrieve pointer
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            if is_leader_cta:
+                num_k_blocks = cute.size(tCrA_0.shape[2])
+                q_stage_0 = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Wait for Q pipeline when batch changes
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma_0.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma_0)
+                        q_stage_0 = q_cons_state_umma_0.index
+
+                    # Process KV block for group 0 (kv_idx + 0)
+                    # Unconditional UMMA (like DeepGEMM): OOB iterations
+                    # compute on garbage data; results written to aligned
+                    # padding region in logits buffer.
+                    # Wait KV first, then TMEM empty (like DeepGEMM)
+                    kv_pipeline_0.consumer_wait(kv_cons_state_umma_0)
+                    umma_pipeline_0.producer_acquire(umma_prod_state_0)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage = kv_cons_state_umma_0.index
+                    tCtAcc_0 = tCtAcc_base_0[(None, None, None, umma_prod_state_0.index)]
+                    for k_block in cutlass.range_constexpr(num_k_blocks):
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc_0,
+                            tCrA_0[None, None, k_block, kv_stage],
+                            tCrB[None, None, k_block, q_stage_0],
+                            tCtAcc_0,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    # No consumer_release here — Math WG 0 releases
+                    kv_cons_state_umma_0.advance()
+
+                    umma_pipeline_0.producer_commit(umma_prod_state_0)
+                    umma_prod_state_0.advance()
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp_1:
+            # UMMA warp for group 1
+            # Explicitly waits on Q pipeline — critical because TMA warp 1
+            # only loads KV1, not Q. Without this wait, UMMA warp 1 can
+            # start GEMM before TMA warp 0 finishes loading Q into SMEM.
+            cute.arch.warpgroup_reg_dealloc(24)
+
+            # TMEM: wait for umma_warp_0's allocation, retrieve pointer
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            if is_leader_cta:
+                num_k_blocks_1 = cute.size(tCrA_1.shape[2])
+                q_stage_1 = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Wait for Q pipeline when batch changes
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma_1.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma_1)
+                        q_stage_1 = q_cons_state_umma_1.index
+
+                    # Process KV block for group 1 (kv_idx + 1)
+                    # Unconditional UMMA (like DeepGEMM)
+                    # Wait KV first, then TMEM empty (like DeepGEMM)
+                    kv_pipeline_1.consumer_wait(kv_cons_state_umma_1)
+                    umma_pipeline_1.producer_acquire(umma_prod_state_1)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage_1 = kv_cons_state_umma_1.index
+                    tCtAcc_1 = tCtAcc_base_1[(None, None, None, umma_prod_state_1.index)]
+                    for k_block in cutlass.range_constexpr(num_k_blocks_1):
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc_1,
+                            tCrA_1[None, None, k_block, kv_stage_1],
+                            tCrB[None, None, k_block, q_stage_1],
+                            tCtAcc_1,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    # No consumer_release here — Math WG 1 releases
+                    kv_cons_state_umma_1.advance()
+
+                    umma_pipeline_1.producer_commit(umma_prod_state_1)
+                    umma_prod_state_1.advance()
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_math_warp:
+            cute.arch.warpgroup_reg_alloc(240)
+
+            # TMEM: math warp 0 is the allocator; all math warps wait + retrieve
+            tmem.allocate(num_tmem_alloc_cols_total)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            local_tidx = tidx % 128
+            cC = cute.make_identity_tensor(epi_sub_mn)
+
+            if warpgroup_idx == 0:
+                # Math WG 0: process group 0
+                # Reference setup (stage 0) for m_coord
+                # flat_divide by sub-tile to get sub-tile partitions
+                tAcc_0_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)]
+                tAcc_0_ref_epi = cute.flat_divide(tAcc_0_ref, epi_sub_mn)
+                tiled_copy_ref_0 = tcgen05.make_tmem_copy(
+                    copy_atom_t2r, tAcc_0_ref_epi[(None, None, 0, 0)]
+                )
+                thr_copy_ref_0 = tiled_copy_ref_0.get_slice(local_tidx)
+                tTR_cC = thr_copy_ref_0.partition_D(cC)
+                m_coord = tTR_cC[0][0]
+
+                tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+                # Weight register cache: only first NUM_W_IN_REG
+                # per next_n slot (like DeepGEMM min(52, kNumHeads)).
+                # Remaining weights read from SMEM in epilogue.
+                # FP16 weights use half the regs, so we can fit
+                # all heads for next_n <= 3.
+                NUM_W_IN_REG = self.num_w_in_reg
+                w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype)
+                q_stage_local = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Q pipeline consumer: wait for Q+Weights SMEM
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_pipeline.consumer_release(q_cons_state)
+                            q_cons_state.advance()
+                        q_pipeline.consumer_wait(q_cons_state)
+                        q_stage_local = q_cons_state.index
+                        # Preload first NUM_W_IN_REG weights per slot
+                        for t_i in cutlass.range_constexpr(next_n):
+                            for w_j in cutlass.range_constexpr(NUM_W_IN_REG):
+                                w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
+                                    (t_i * num_heads + w_j, q_stage_local)
+                                ]
+
+                    # Process KV block for group 0 (kv_idx + 0)
+                    # Unconditional Math (like DeepGEMM): OOB results
+                    # written to aligned padding region in logits buffer.
+                    kv_pos = kv_idx * block_kv_val + m_coord
+
+                    if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                        # Skip KV wait, rely on UMMA barrier's
+                        # transitive visibility.
+                        umma_pipeline_0.consumer_wait(umma_cons_state_0)
+                    else:
+                        # Default: wait KV first to overlap lds with
+                        # UMMA computation.
+                        kv_pipeline_0.consumer_wait(kv_cons_state_math_0)
+                        sc_stage_0 = kv_cons_state_math_0.index
+                        scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        umma_pipeline_0.consumer_wait(umma_cons_state_0)
+
+                    # --- TMEM sub-tile setup ---
+                    # flat_divide accumulator by sub-tile shape;
+                    # partition once, then loop over sub-tiles.
+                    tCtAcc_c0 = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)]
+                    tAcc_c0 = tCtAcc_c0[((None, None), 0, 0)]
+                    tAcc_c0_epi = cute.flat_divide(tAcc_c0, epi_sub_mn)
+                    tc_0 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c0_epi[(None, None, 0, 0)])
+                    tr_0 = tc_0.get_slice(local_tidx)
+                    tTR_0 = tr_0.partition_S(tAcc_c0_epi)
+
+                    # --- First sub-tile LDTM + KV release ---
+                    if cutlass.const_expr(self.early_tmem_copy):
+                        # Issue first sub-tile LDTM early
+                        cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc)
+                        # Scale LDS + KV release fill latency
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_0 = kv_cons_state_math_0.index
+                            scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        kv_pipeline_0.consumer_release(kv_cons_state_math_0)
+                        kv_cons_state_math_0.advance()
+                        cute.arch.fence_view_async_tmem_load()
+                    else:
+                        # Default: scale LDS + KV release first
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_0 = kv_cons_state_math_0.index
+                            scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        kv_pipeline_0.consumer_release(kv_cons_state_math_0)
+                        kv_cons_state_math_0.advance()
+                        cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc)
+                        cute.arch.fence_view_async_tmem_load()
+
+                    # --- Sub-tile compute loop ---
+                    # Each sub-tile: LDTM.xN → fence → load →
+                    # ReLU+FMA. Breaks FMA chain (16→4 per chunk)
+                    # and interleaves LDTM with FP32 compute to
+                    # reduce ShadowPipeThrottle.
+                    # Sub-tiles are within each t-slot (num_heads
+                    # // num_epi_subtiles wide). flat_divide yields
+                    # next_n * num_epi_subtiles sub-tiles total;
+                    # global index = t * num_epi_subtiles + i.
+                    subtile_n = num_heads // num_epi_subtiles
+                    if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                        packed_zero = pack_f16x2(Float16(0.0), Float16(0.0))
+                    for t in cutlass.range_constexpr(next_n):
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps0 = packed_zero
+                            ps1 = packed_zero
+                        else:
+                            s0x = cutlass.Float32(0.0)
+                            s0y = cutlass.Float32(0.0)
+                            s1x = cutlass.Float32(0.0)
+                            s1y = cutlass.Float32(0.0)
+                        for i in cutlass.range_constexpr(num_epi_subtiles):
+                            # LDTM for sub-tiles 1..N-1
+                            # (sub-tile 0 handled above)
+                            if t > 0 or i > 0:
+                                cute.copy(
+                                    tc_0,
+                                    tTR_0[(None, None, None, 0, t * num_epi_subtiles + i)],
+                                    tTR_rAcc,
+                                )
+                                cute.arch.fence_view_async_tmem_load()
+                            # Release UMMA after last LDTM+fence
+                            if t == next_n - 1 and i == num_epi_subtiles - 1:
+                                umma_pipeline_0.consumer_release(umma_cons_state_0)
+                                umma_cons_state_0.advance()
+                            acc_vec = tTR_rAcc.load()
+                            # Reg-path: weights from registers
+                            reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n))
+                            for h in cutlass.range_constexpr(0, reg_h_end, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1])
+                                    pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3])
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    w0 = w_cache[r0]
+                                    w1 = w_cache[r0 + 1]
+                                    w2 = w_cache[r0 + 2]
+                                    w3 = w_cache[r0 + 3]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                            # SMEM-path: weights from shared mem
+                            smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n)
+                            for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    pw01 = pack_f16x2(
+                                        sW[(t * num_heads + h_g, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 1, q_stage_local)],
+                                    )
+                                    pw23 = pack_f16x2(
+                                        sW[(t * num_heads + h_g + 2, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 3, q_stage_local)],
+                                    )
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    w0 = sW[(t * num_heads + h_g, q_stage_local)]
+                                    w1 = sW[(t * num_heads + h_g + 1, q_stage_local)]
+                                    w2 = sW[(t * num_heads + h_g + 2, q_stage_local)]
+                                    w3 = sW[(t * num_heads + h_g + 3, q_stage_local)]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps_sum = add_f16x2(ps0, ps1)
+                            sum_lo, sum_hi = unpack_f16x2(ps_sum)
+                            result_t = sum_lo + sum_hi
+                        else:
+                            result_t = s0x + s0y + s1x + s1y
+                        out_row = q_idx * next_n + t
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(
+                                result_t * Float16(scale_val)
+                            )
+                        else:
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val)
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+                # Release last Q stage (WG 0)
+                if q_idx < batch_size:
+                    q_pipeline.consumer_release(q_cons_state)
+                    q_cons_state.advance()
+
+            else:
+                # Math WG 1: process group 1
+                tAcc_1_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)]
+                tAcc_1_ref_epi = cute.flat_divide(tAcc_1_ref, epi_sub_mn)
+                tiled_copy_ref_1 = tcgen05.make_tmem_copy(
+                    copy_atom_t2r, tAcc_1_ref_epi[(None, None, 0, 0)]
+                )
+                thr_copy_ref_1 = tiled_copy_ref_1.get_slice(local_tidx)
+                tTR_cC = thr_copy_ref_1.partition_D(cC)
+                m_coord = tTR_cC[0][0]
+
+                tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+                # Weight register cache: only first NUM_W_IN_REG
+                # per next_n slot (like DeepGEMM min(52, kNumHeads)).
+                # FP16 weights use half the regs, so we can fit
+                # all heads for next_n <= 3.
+                NUM_W_IN_REG = self.num_w_in_reg
+                w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype)
+                q_stage_local = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Q pipeline consumer: wait for Q+Weights SMEM
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_pipeline.consumer_release(q_cons_state)
+                            q_cons_state.advance()
+                        q_pipeline.consumer_wait(q_cons_state)
+                        q_stage_local = q_cons_state.index
+                        # Preload first NUM_W_IN_REG weights per slot
+                        for t_i in cutlass.range_constexpr(next_n):
+                            for w_j in cutlass.range_constexpr(NUM_W_IN_REG):
+                                w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
+                                    (t_i * num_heads + w_j, q_stage_local)
+                                ]
+
+                    # Process KV block for group 1 (kv_idx + 1)
+                    # Unconditional Math (like DeepGEMM)
+                    kv_idx_1 = kv_idx + 1
+
+                    kv_pos = kv_idx_1 * block_kv_val + m_coord
+
+                    if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                        umma_pipeline_1.consumer_wait(umma_cons_state_1)
+                    else:
+                        kv_pipeline_1.consumer_wait(kv_cons_state_math_1)
+                        sc_stage_1 = kv_cons_state_math_1.index
+                        scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        umma_pipeline_1.consumer_wait(umma_cons_state_1)
+
+                    # --- TMEM sub-tile setup (WG1) ---
+                    tCtAcc_c1 = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)]
+                    tAcc_c1 = tCtAcc_c1[((None, None), 0, 0)]
+                    tAcc_c1_epi = cute.flat_divide(tAcc_c1, epi_sub_mn)
+                    tc_1 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c1_epi[(None, None, 0, 0)])
+                    tr_1 = tc_1.get_slice(local_tidx)
+                    tTR_1 = tr_1.partition_S(tAcc_c1_epi)
+
+                    # --- First sub-tile LDTM + KV release (WG1) ---
+                    if cutlass.const_expr(self.early_tmem_copy):
+                        cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc)
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_1 = kv_cons_state_math_1.index
+                            scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        kv_pipeline_1.consumer_release(kv_cons_state_math_1)
+                        kv_cons_state_math_1.advance()
+                        cute.arch.fence_view_async_tmem_load()
+                    else:
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_1 = kv_cons_state_math_1.index
+                            scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        kv_pipeline_1.consumer_release(kv_cons_state_math_1)
+                        kv_cons_state_math_1.advance()
+                        cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc)
+                        cute.arch.fence_view_async_tmem_load()
+
+                    # --- Sub-tile compute loop (WG1) ---
+                    subtile_n = num_heads // num_epi_subtiles
+                    if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                        packed_zero = pack_f16x2(Float16(0.0), Float16(0.0))
+                    for t in cutlass.range_constexpr(next_n):
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps0 = packed_zero
+                            ps1 = packed_zero
+                        else:
+                            s0x = cutlass.Float32(0.0)
+                            s0y = cutlass.Float32(0.0)
+                            s1x = cutlass.Float32(0.0)
+                            s1y = cutlass.Float32(0.0)
+                        for i in cutlass.range_constexpr(num_epi_subtiles):
+                            if t > 0 or i > 0:
+                                cute.copy(
+                                    tc_1,
+                                    tTR_1[(None, None, None, 0, t * num_epi_subtiles + i)],
+                                    tTR_rAcc,
+                                )
+                                cute.arch.fence_view_async_tmem_load()
+                            if t == next_n - 1 and i == num_epi_subtiles - 1:
+                                umma_pipeline_1.consumer_release(umma_cons_state_1)
+                                umma_cons_state_1.advance()
+                            acc_vec = tTR_rAcc.load()
+                            # Reg-path
+                            reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n))
+                            for h in cutlass.range_constexpr(0, reg_h_end, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1])
+                                    pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3])
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    w0 = w_cache[r0]
+                                    w1 = w_cache[r0 + 1]
+                                    w2 = w_cache[r0 + 2]
+                                    w3 = w_cache[r0 + 3]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                            # SMEM-path
+                            smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n)
+                            for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    pw01 = pack_f16x2(
+                                        sW[(t * num_heads + h_g, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 1, q_stage_local)],
+                                    )
+                                    pw23 = pack_f16x2(
+                                        sW[(t * num_heads + h_g + 2, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 3, q_stage_local)],
+                                    )
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    w0 = sW[(t * num_heads + h_g, q_stage_local)]
+                                    w1 = sW[(t * num_heads + h_g + 1, q_stage_local)]
+                                    w2 = sW[(t * num_heads + h_g + 2, q_stage_local)]
+                                    w3 = sW[(t * num_heads + h_g + 3, q_stage_local)]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps_sum = add_f16x2(ps0, ps1)
+                            sum_lo, sum_hi = unpack_f16x2(ps_sum)
+                            result_t = sum_lo + sum_hi
+                        else:
+                            result_t = s0x + s0y + s1x + s1y
+                        out_row = q_idx * next_n + t
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(
+                                result_t * Float16(scale_val)
+                            )
+                        else:
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val)
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+                # Release last Q stage (WG 1)
+                if q_idx < batch_size:
+                    q_pipeline.consumer_release(q_cons_state)
+                    q_cons_state.advance()
+
+            # TMEM dealloc: math warps are allocator + last consumer
+            tmem.relinquish_alloc_permit()
+            tmem.free(tmem_ptr)
+
+        else:
+            cute.arch.warpgroup_reg_dealloc(24)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py
new file mode 100644
index 000000000000..32c115e2a7da
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py
@@ -0,0 +1,1747 @@
+# 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.
+"""
+indexer_topk variant of the CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100).
+
+Derived from ``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
+for the GLM-5.x indexer shape, where all verify rows of a request fit one UMMA
+tile (``N = next_n * num_heads <= 256``). The only functional change is the
+size of the per-thread register cache of epilogue weights
+(``num_w_in_reg`` heads per verify row; the remaining heads are read from
+shared memory). Keeping that cache bounded avoids register spills when
+``N`` is large; weight values, accumulation order and output are unchanged.
+
+Architecture:
+  - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA)
+  - 1 TMA per KV block [128, 128], UMMA iterates 4x K=32
+  - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2)
+  - Q reloaded via TMA pipeline when q_idx (batch) changes
+  - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs
+  - Weights cached in registers: preloaded once per q_idx change (not per KV block)
+  - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes)
+
+Merged KV+Scale pipeline:
+  - KV data and scales share a single TMA barrier per group
+  - TMA loads both KV and Scale under one barrier (combined tx_count)
+  - UMMA waits on merged barrier (for KV GEMM), does NOT release
+  - Math waits on merged barrier (for scale read), Math releases
+
+Fused KV layout:
+  - KV data and scales stored contiguously per physical block:
+    [num_phys_blocks, block_kv * (head_dim + 4)] bytes
+  - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)]
+  - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic
+
+Scheduler:
+  - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups)
+  - schedule_meta[sm_idx+1] = end boundary for this CTA
+  - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx)
+  - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step
+
+Dynamic shape support:
+  - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen
+  - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq,
+    num_ctas) are marked dynamic via mark_compact_shape_dynamic
+  - Allows JIT cache reuse across different batch sizes / sequence lengths
+
+Epilogue dtype flows (--acc_dtype / --epi_dtype):
+
+  Flow 1: --acc_dtype fp32 --epi_dtype fp16
+    Q(FP8) x K(FP8) -> MMA acc(FP32) -> TMEM(FP32)
+      -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16)
+      -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16)
+      -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype)
+
+  Flow 2: --acc_dtype fp16 --epi_dtype fp16
+    Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b)
+      -> LDTM -> Reg(FP16) -> ReLU(FP16)
+      -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16)
+      -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype)
+
+  --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion.
+  Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32
+"""
+
+from typing import Tuple
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass import Float16, Int32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+# CuTe DSL CUDA 13 validates rounding modes as string literals. The string
+# form is also accepted by older wrappers, so keep it version-independent.
+_RND_RN = "rn"
+
+
+@dsl_user_op
+def pack_f16x2(
+    a: Float16,
+    b: Float16,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    f16_ty = Float16.mlir_type
+    i32_ty = Int32.mlir_type
+    vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc)
+    v = vector.from_elements(
+        vec2_f16,
+        (Float16(a).ir_value(loc=loc, ip=ip), Float16(b).ir_value(loc=loc, ip=ip)),
+        loc=loc,
+        ip=ip,
+    )
+    return Int32(llvm.bitcast(i32_ty, v, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def unpack_f16x2(
+    packed: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Tuple[Float16, Float16]:
+    f16_ty = Float16.mlir_type
+    vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc)
+    v = llvm.bitcast(vec2_f16, Int32(packed).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    r0 = Float16(vector.extract(v, dynamic_position=[], static_position=[0], loc=loc, ip=ip))
+    r1 = Float16(vector.extract(v, dynamic_position=[], static_position=[1], loc=loc, ip=ip))
+    return r0, r1
+
+
+@dsl_user_op
+def fma_f16x2(
+    a: Int32,
+    b: Int32,
+    c: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [
+                Int32(a).ir_value(loc=loc, ip=ip),
+                Int32(b).ir_value(loc=loc, ip=ip),
+                Int32(c).ir_value(loc=loc, ip=ip),
+            ],
+            "fma.rn.f16x2 $0, $1, $2, $3;",
+            "=r,r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def max_f16x2(
+    a: Int32,
+    b: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "max.f16x2 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def add_f16x2(
+    a: Int32,
+    b: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    i32_ty = Int32.mlir_type
+    return Int32(
+        llvm.inline_asm(
+            i32_ty,
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "add.f16x2 $0, $1, $2;",
+            "=r,r,r",
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+class FP8MQALogitsKernel:
+    """FP8 paged MQA logits kernel for Blackwell (SM100).
+
+    Each CTA processes a range of (q_idx, kv_split) pairs.
+    A split = 2 consecutive KV blocks within a sequence (one per warp group).
+    Q is shared between warp groups and reloaded when q_idx changes.
+    """
+
+    def __init__(
+        self,
+        block_kv: int = 128,
+        phys_block_kv: int = 128,
+        num_heads: int = 64,
+        head_dim: int = 128,
+        next_n: int = 1,
+        num_sms: int = 148,
+        remove_kv_wait_in_epilogue: bool = False,
+        early_tmem_copy: bool = False,
+        smem_subpartition_opt: bool = False,
+        max_kv_pipeline: bool = False,
+        max_umma_pipeline: bool = False,
+        num_epi_subtiles: int = 1,
+        epi_dtype=cutlass.Float32,
+        acc_dtype=cutlass.Float32,
+        output_dtype=cutlass.Float32,
+        num_w_in_reg=None,
+    ):
+        self.block_kv = block_kv
+        self.phys_block_kv = phys_block_kv
+        self.num_blocks_per_mma = block_kv // phys_block_kv
+        assert block_kv % phys_block_kv == 0, (
+            f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}"
+        )
+        assert self.num_blocks_per_mma <= 4, (
+            f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4"
+        )
+        self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue
+        self.early_tmem_copy = early_tmem_copy
+        self.smem_subpartition_opt = smem_subpartition_opt
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.next_n = next_n
+        self.N = next_n * num_heads
+        self.num_sms = num_sms
+        self.num_epi_subtiles = num_epi_subtiles
+        self.epi_dtype = epi_dtype
+        self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4
+        # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
+        # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1
+        # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest).
+        w_stage_bytes = self.N * self.epi_bytes
+        self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes
+        self.output_dtype = output_dtype
+        if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0:
+            raise ValueError("num_heads must be divisible by num_epi_subtiles")
+        if (num_heads // num_epi_subtiles) % 4 != 0:
+            raise ValueError(
+                "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)"
+            )
+        self.num_groups = 2
+
+        # Register-cached epilogue weights per verify row. None keeps the
+        # upstream heuristic; otherwise a multiple of 4 (FMA unroll) <= heads.
+        if num_w_in_reg is None:
+            if epi_dtype == cutlass.Float16:
+                max_w_in_reg = 64 if next_n <= 3 else 48
+            else:
+                max_w_in_reg = 64 if next_n == 1 else 40 if next_n >= 4 else 52
+            num_w_in_reg = min(max_w_in_reg, num_heads)
+        if num_w_in_reg < 0 or num_w_in_reg > num_heads or num_w_in_reg % 4 != 0:
+            raise ValueError(
+                f"num_w_in_reg={num_w_in_reg} must be a multiple of 4 in [0, {num_heads}]"
+            )
+        self.num_w_in_reg = num_w_in_reg
+
+        self.num_math_threads = 256
+        self.num_specialized_threads = 128
+        self.threads_per_cta = 384
+        self.num_math_warps = 8
+        self.tma_warp_base = 8
+        self.umma_warp_base = 10
+
+        self.num_q_stages = 3  # 3 stages for Q pipelining across batch sequences
+
+        # TMEM: 512 columns total, each group needs N columns per UMMA stage
+        # max_umma_stages = 512 // (2 * N)
+        TMEM_COLS = 512
+        if max_umma_pipeline:
+            self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N))
+        else:
+            self.num_umma_stages = 1
+
+        if max_kv_pipeline:
+            smem_capacity = utils.get_smem_capacity_in_bytes()
+            # Reserve ~1 KB for barriers and misc
+            SMEM_BUDGET = smem_capacity - 1024
+            # KV+Scale per stage (×2 groups):
+            #   2 * (block_kv * head_dim * 1B + block_kv * 4B)
+            kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4)
+            # Q+W per stage: Q is N * head_dim * 1B, W uses padded stride
+            qw_per_stage = self.N * head_dim + self.w_stage_stride * self.epi_bytes
+            qw_total = qw_per_stage * self.num_q_stages
+            self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage
+        else:
+            self.num_kv_stages = 3
+
+        # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB),
+        # avoiding sub-bank conflicts with UMMA reading sKV.
+        # Layout: barriers(~256B) | sKV_0 | sKV_1 | sQ | [pad] | sW | sScales
+        if self.smem_subpartition_opt:
+            BOUNDARY = 128 * 1024
+            used = (
+                256
+                + 2 * (block_kv * head_dim * self.num_kv_stages)
+                + self.N * head_dim * self.num_q_stages
+            )
+            used = ((used + 127) // 128) * 128
+            if used < BOUNDARY:
+                self.smem_pad_bytes = ((BOUNDARY - used + 1023) // 1024) * 1024
+            else:
+                self.smem_pad_bytes = 0
+        else:
+            self.smem_pad_bytes = 0
+
+        self.acc_dtype = acc_dtype
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.cluster_shape_mn = (1, 1)
+        self.mma_tiler_mn = (block_kv, self.N)
+
+    def _setup_mma(self, a_dtype, b_dtype, a_major, b_major):
+        self.a_dtype = a_dtype
+        self.b_dtype = b_dtype
+        self.a_major_mode = a_major
+        self.b_major_mode = b_major
+
+        self.mma_tiler = (*self.mma_tiler_mn, 1)
+        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            a_dtype,
+            a_major,
+            b_major,
+            self.acc_dtype,
+            self.cta_group,
+            self.mma_tiler_mn,
+        )
+        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])  # 32
+
+        # Full-K: tile K = head_dim (128), 1 TMA per block
+        mma_inst_tile_k = self.head_dim // mma_inst_shape_k  # 4
+        full_k = mma_inst_shape_k * mma_inst_tile_k  # 128
+        self.mma_tiler = (
+            self.mma_tiler_mn[0],
+            self.mma_tiler_mn[1],
+            full_k,
+        )
+
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (tiled_mma.thr_id.shape,),
+        )
+        self.epi_tile = self.cta_tile_shape_mnk[:2]
+
+        # KV SMEM: 3 stages per group, each stage holds full [128, 128]
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            self.mma_tiler,
+            a_dtype,
+            self.num_kv_stages,
+        )
+        # Q SMEM: 1 stage, holds full [N, 128]
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            self.mma_tiler,
+            b_dtype,
+            self.num_q_stages,
+        )
+
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+        self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake)
+        self.num_tmem_alloc_cols_total = (
+            self.num_tmem_alloc_cols * self.num_groups * self.num_umma_stages
+        )
+
+        return tiled_mma
+
+    @cute.jit
+    def __call__(
+        self,
+        kv_fused: cute.Tensor,  # Fused KV: [num_phys_blocks, block_bytes] FP8
+        b: cute.Tensor,  # Q: [N, head_dim, batch_size]
+        weights: cute.Tensor,  # [N, batch_size] (transposed for TMA)
+        logits: cute.Tensor,  # [batch_size * next_n, max_context_len]
+        block_table: cute.Tensor,  # [batch_size, max_blocks_per_seq]
+        context_lens: cute.Tensor,  # [batch_size]
+        schedule_meta: cute.Tensor,  # [num_sms+1, 2] int32
+        num_phys_blocks: cutlass.Int32,
+        batch_size: cutlass.Int32,
+        stream: cuda.CUstream,
+    ):
+        # Derive KV and Scale views from fused buffer using CuTE ops.
+        # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)]
+        phys_block_kv = self.phys_block_kv
+        scale_offset_elems = phys_block_kv * self.head_dim  # in FP8 elements
+
+        # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference)
+        kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN)
+
+        # Q (b) was passed as uint8 to work around DLPack's lack of float8 support;
+        # recast back to FP8 so MMA type inference and TMA descriptors are correct.
+        b = cute.recast_tensor(b, cutlass.Float8E4M3FN)
+
+        # Read the real per-block stride (bytes = FP8 elements) from the input.
+        # When KV is the indexer K-cache pool view, the pool is laid out as
+        # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride =
+        # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes).
+        # Using the input stride keeps both contiguous test path and the
+        # strided prod path correct. FP8 = 1 byte per element, so the byte
+        # stride is the same as the element stride after recast.
+        kv_block_stride = kv_fused.layout.stride[0]
+
+        # KV view: [phys_block_kv, head_dim, num_phys_blocks] FP8
+        # Each TMA loads one physical block; multiple TMAs fill a compute tile.
+        kv_layout = cute.make_layout(
+            (phys_block_kv, self.head_dim, num_phys_blocks),
+            stride=(self.head_dim, 1, kv_block_stride),
+        )
+        a = cute.make_tensor(kv_fp8.iterator, kv_layout)
+
+        # Scale view: offset pointer to scale region, recast FP8 → Float32
+        # [phys_block_kv, num_phys_blocks] float32 (after recast)
+        scale_fp8_layout = cute.make_layout(
+            (phys_block_kv * 4, num_phys_blocks),
+            stride=(1, kv_block_stride),
+        )
+        scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout)
+        scales = cute.recast_tensor(scale_fp8, cutlass.Float32)
+
+        a_dtype = a.element_type
+        b_dtype = b.element_type
+        a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode()
+        b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+
+        tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major)
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        # TMA for KV (A) — fmha_decode_paged pattern.
+        # Build a TMA SMEM layout via tiled_divide on the full compute-tile
+        # layout, then select to drop trivial K dim. Atom uses mode [0] as
+        # single-tile SMEM layout and (phys, head) as cta_tiler.
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+        self.a_tma_view_layout = sm100_utils.make_smem_layout(
+            tcgen05.OperandMajorMode.K,
+            (self.block_kv, self.head_dim),
+            a_dtype,
+            self.num_kv_stages,
+        )
+        self.a_tma_view_layout = cute.tiled_divide(
+            self.a_tma_view_layout, (self.phys_block_kv, self.head_dim)
+        )
+        # ((tile_M, tile_K), rest_M, rest_K, stages) → drop trivial rest_K
+        self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3])
+        # ((tile_M, tile_K), rest_M=num_sub_blocks, stages)
+        tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            a,
+            self.a_tma_view_layout[0],  # atom SMEM = single-tile (mode 0)
+            (self.phys_block_kv, self.head_dim),
+        )
+
+        # TMA for Q (B) — full K=128, L dim = batch_size (unchanged)
+        b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            b_op,
+            b,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # TMA for Weights — [N, batch_size], tile [N], L=batch_size
+        self.w_smem_layout_staged = cute.make_layout(
+            (self.N, self.num_q_stages),
+            stride=(1, self.w_stage_stride),
+        )
+        w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0])
+        tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            weights,
+            w_smem_per_stage,
+            self.w_smem_layout_staged.shape[:1],
+        )
+
+        # TMA for Scales — [phys_block_kv, num_phys_blocks], tile [phys_block_kv]
+        # SMEM holds compute_block_kv scales per stage; filled by
+        # num_blocks_per_mma sub-block TMAs at consecutive offsets.
+        self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages))
+        s_smem_per_subblock = cute.make_layout((phys_block_kv,))
+        tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            scales,
+            s_smem_per_subblock,
+            (phys_block_kv,),
+        )
+
+        b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout)
+        w_copy_size = self.N * self.epi_bytes
+        # Per sub-block: phys_block_kv * head_dim (KV) + phys_block_kv * 4 (scales)
+        kv_tma_bytes_per_subblock = phys_block_kv * self.head_dim
+        scale_tma_bytes_per_subblock = phys_block_kv * 4
+        # Total per compute tile = num_blocks_per_mma sub-blocks
+        self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * (
+            kv_tma_bytes_per_subblock + scale_tma_bytes_per_subblock
+        )
+        # Q + Weights share barrier (like DeepGEMM)
+        self.num_q_tma_bytes = b_copy_size * atom_thr_size + w_copy_size
+
+        num_ctas = self.num_sms
+
+        @cute.struct
+        class SharedStorage:
+            kv_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            kv_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2]
+            umma_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            umma_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            tmem_holding_buf: cutlass.Int32
+
+        self.kernel(
+            tiled_mma,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_w,
+            tma_tensor_w,
+            tma_atom_s,
+            tma_tensor_s,
+            logits,
+            block_table,
+            context_lens,
+            schedule_meta,
+            batch_size,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.w_smem_layout_staged,
+            self.s_smem_layout_staged,
+            self.a_tma_view_layout,
+            self.epi_tile,
+            SharedStorage,
+        ).launch(
+            grid=(1, 1, num_ctas),
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,  # KV pool
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,  # Q (L dim = batch_size)
+        tma_atom_w: cute.CopyAtom,
+        mW_tma: cute.Tensor,  # Weights TMA coord tensor [N, batch_size]
+        tma_atom_s: cute.CopyAtom,
+        mS_tma: cute.Tensor,  # Scales TMA coord tensor [block_kv, num_phys_blocks]
+        mLogits: cute.Tensor,  # [batch_size * next_n, max_context_len]
+        mBlockTable: cute.Tensor,  # [batch_size, max_blocks_per_seq]
+        mContextLens: cute.Tensor,  # [batch_size]
+        mScheduleMeta: cute.Tensor,  # [num_sms+1, 2] int32
+        batch_size: cutlass.Int32,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        w_smem_layout_staged: cute.Layout,
+        s_smem_layout_staged: cute.Layout,
+        a_tma_view_layout: cute.ComposedLayout,
+        epi_tile: cute.Tile,
+        SharedStorage: cutlass.Constexpr,
+    ):
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
+
+        bidx, bidy, bidz = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+        is_leader_cta = mma_tile_coord_v == 0
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        tidx, _, _ = cute.arch.thread_idx()
+
+        # Warp roles (matches DeepGEMM SM100)
+        warpgroup_idx = warp_idx // 4
+        is_math_warp = warp_idx < 8
+        is_tma_warp_0 = warp_idx == 8
+        is_tma_warp_1 = warp_idx == 9
+        is_tma_warp = is_tma_warp_0 | is_tma_warp_1
+        is_umma_warp_0 = warp_idx == 10
+        is_umma_warp_1 = warp_idx == 11
+
+        # Early schedule metadata load: issue global loads ASAP so their
+        # ~200-cycle L2 latency overlaps with subsequent prologue setup
+        # (SMEM alloc, TMA partition, MMA fragment creation, etc.)
+        NUM_MATH_WG = 2  # kNumMathWarpGroups
+        NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma
+        sm_idx = bidz
+        start_q = mScheduleMeta[(sm_idx, 0)]
+        start_kv_half = mScheduleMeta[(sm_idx, 1)]
+        end_q_idx = mScheduleMeta[(sm_idx + 1, 0)]
+        end_kv_half = mScheduleMeta[(sm_idx + 1, 1)]
+        # Early mContextLens load: overlap ~200-cycle L2 latency with the
+        # entire prologue setup (pipelines, SMEM alloc, TMA partition, etc.)
+        # Clamp to avoid OOB when start_q == batch_size (zero-work CTA sentinel).
+        # Note: zero-work CTAs get a stale current_num_kv (from the last batch
+        # element), but it is never used because has_work will be False.
+        start_q_clamped = min(start_q, batch_size - 1)
+        current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv
+
+        if is_tma_warp:
+            cpasync.prefetch_descriptor(tma_atom_a)
+            cpasync.prefetch_descriptor(tma_atom_b)
+            cpasync.prefetch_descriptor(tma_atom_w)
+            cpasync.prefetch_descriptor(tma_atom_s)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        block_kv_val = self.block_kv
+        num_heads = self.num_heads
+        next_n = self.next_n
+        num_epi_subtiles = self.num_epi_subtiles
+
+        # === Pipelines ===
+        prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+
+        # Q pipeline: TMA producer → Math consumer (8 math warps)
+        # PipelineTmaAsync: consumer_release uses is_signalling_thread
+        # (lane 0 per warp). 8 math warps × 1 lane-0 = 8 arrives.
+        q_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 8)
+        q_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.q_mbar.data_ptr(),
+            num_stages=self.num_q_stages,
+            producer_group=prod_group,
+            consumer_group=q_cons_group,
+            tx_count=self.num_q_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        q_prod_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_q_stages
+        )
+        # Both Math WGs share the same pipeline state (advance in lockstep)
+        q_cons_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        # UMMA warps observe Q pipeline (wait only, no release)
+        # to ensure Q is in SMEM before GEMM. Critical for UMMA warp 1
+        # since TMA warp 1 only loads KV1 (not Q).
+        q_cons_state_umma_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        q_cons_state_umma_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+
+        # Merged KV+Scale pipelines (per-group, 3 stages each)
+        # Like DeepGEMM: KV data and scales share one barrier.
+        # TMA loads both under one barrier. Math is consumer (releases).
+        # UMMA also waits on this barrier (for KV GEMM) but does NOT release.
+        math_warps_per_group = self.num_math_warps // 2  # 4 warps
+        kv_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, math_warps_per_group)
+        kv_pipeline_0 = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar_0.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=kv_cons_group,
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        kv_pipeline_1 = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar_1.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=kv_cons_group,
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+
+        kv_prod_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        kv_prod_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        # UMMA consumer states (wait only, no release)
+        kv_cons_state_umma_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_umma_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        # Math consumer states (wait + release)
+        kv_cons_state_math_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_math_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+
+        # UMMA pipelines (per-group)
+        math_threads_per_group = self.num_math_threads // 2
+        umma_pipeline_0 = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar_0.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        umma_pipeline_1 = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar_1.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+
+        umma_prod_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_prod_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_cons_state_0 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+        umma_cons_state_1 = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+
+        # TMEM — only Math warps (8×32=256) + UMMA warps (2×32=64) = 320 threads
+        # TMA warps do NOT participate, so they can start TMA loads earlier.
+        # Math warp 0 is the allocator (like fp16_gemm_3's epilogue warp 0),
+        # because math warps are the last TMEM consumers (epilogue reads).
+        tmem_alloc_num_threads = 320  # 10 warps: warp 0-7 (math) + warp 10-11 (umma)
+        tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=tmem_alloc_num_threads)
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=0,  # math warp 0 does alloc+free (last TMEM consumer)
+            is_two_cta=False,
+        )
+
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+
+        # SMEM allocation: per-group KV + shared Q
+        sKV_0 = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sKV_1 = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sQ = smem.allocate_tensor(
+            element_type=self.b_dtype,
+            layout=b_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=b_smem_layout_staged.inner,
+        )
+        # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB)
+        # to avoid sub-bank conflicts with UMMA reading sKV from
+        # sub-partition 0.
+        if cutlass.const_expr(self.smem_pad_bytes > 0):
+            _ = smem.allocate(self.smem_pad_bytes)
+        # Weights SMEM: [N, num_q_stages], shared Q barrier
+        sW = smem.allocate_tensor(
+            element_type=self.epi_dtype,
+            layout=w_smem_layout_staged,
+            byte_alignment=128,
+        )
+        # Scales SMEM: [block_kv, num_kv_stages] float32, per group
+        sScales_0 = smem.allocate_tensor(
+            element_type=cutlass.Float32,
+            layout=s_smem_layout_staged,
+            byte_alignment=128,
+        )
+        sScales_1 = smem.allocate_tensor(
+            element_type=cutlass.Float32,
+            layout=s_smem_layout_staged,
+            byte_alignment=128,
+        )
+
+        a_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+        )
+        b_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+        )
+
+        # Partition KV (A): fmha_decode_paged pattern.
+        # SMEM view is ((tile), num_sub_blocks, stages) — built in __call__.
+        # Use .outer (plain layout); swizzle is captured by sKV_0's iterator.
+        # GMEM: local_tile by (phys, head), then group first 2 modes into tile.
+        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+        sKV_0_for_tma = cute.make_tensor(sKV_0.iterator, a_tma_view_layout.outer)
+        sKV_1_for_tma = cute.make_tensor(sKV_1.iterator, a_tma_view_layout.outer)
+        gA = cute.local_tile(
+            mA_mkl,
+            (self.phys_block_kv, self.head_dim),
+            coord=(None, None, None),
+        )
+        tAsA_0, tAgA_0 = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            sKV_0_for_tma,
+            cute.group_modes(gA, 0, 2),
+        )
+        tAsA_1, tAgA_1 = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            sKV_1_for_tma,
+            cute.group_modes(gA, 0, 2),
+        )
+
+        # Partition Q (B): shared SMEM, L dim = batch_size
+        gB_nkl = cute.local_tile(
+            mB_nkl,
+            cute.slice_(self.mma_tiler, (0, None, None)),
+            (None, None, None),
+        )
+        tCgB = thr_mma.partition_B(gB_nkl)
+        b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            block_in_cluster_coord_vmnk[1],
+            b_cta_layout,
+            cute.group_modes(sQ, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+        tBgB = tBgB[(None, 0, None, None)]  # [tma, K, L]
+
+        # Partition Weights: standalone TMA, [N, batch_size] → [N] per stage
+        w_cta_layout = cute.make_layout((1,))
+        tWsW, tWgW = cpasync.tma_partition(
+            tma_atom_w,
+            0,
+            w_cta_layout,
+            cute.group_modes(sW, 0, 1),
+            cute.group_modes(mW_tma, 0, 1),
+        )
+
+        # Partition Scales: explicit sub_blocks + stages dims.
+        # Layout (phys_block_kv, num_sub, stages) K-major with custom strides.
+        s_tma_view_layout = cute.make_layout(
+            (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages),
+            stride=(1, self.phys_block_kv, self.block_kv),
+        )
+        sScales_0_for_tma = cute.make_tensor(sScales_0.iterator, s_tma_view_layout)
+        sScales_1_for_tma = cute.make_tensor(sScales_1.iterator, s_tma_view_layout)
+        # GMEM: local_tile by phys to match atom's tile size
+        gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None))
+        tSsS_0, tSgS_0 = cpasync.tma_partition(
+            tma_atom_s,
+            0,
+            cute.make_layout(1),
+            sScales_0_for_tma,
+            gS,
+        )
+        tSsS_1, tSgS_1 = cpasync.tma_partition(
+            tma_atom_s,
+            0,
+            cute.make_layout(1),
+            sScales_1_for_tma,
+            gS,
+        )
+
+        # MMA fragments
+        tCrA_0 = tiled_mma.make_fragment_A(sKV_0)
+        tCrA_1 = tiled_mma.make_fragment_A(sKV_1)
+        tCrB = tiled_mma.make_fragment_B(sQ)  # shared
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+
+        # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim
+        # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank
+        us = self.num_umma_stages
+        cols = self.num_tmem_alloc_cols
+        acc_shape_staged = cute.append(acc_shape, us)
+        tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged)
+
+        # TMEM layout info (allocation deferred to UMMA/Math warp branches)
+        cols_per_group = cols * us * (32 // self.acc_dtype.width)
+        num_tmem_alloc_cols_total = self.num_tmem_alloc_cols_total
+
+        # Epilogue setup
+        c_layout = utils.LayoutEnum.ROW_MAJOR
+        epi_sub_mn = (epi_tile[0], num_heads // num_epi_subtiles)
+        copy_atom_t2r = sm100_utils.get_tmem_load_op(
+            self.cta_tile_shape_mnk,
+            c_layout,
+            self.acc_dtype,
+            self.acc_dtype,
+            epi_sub_mn,
+            use_2cta_instrs,
+        )
+
+        # ===== SCHEDULER: derive values from early-loaded schedule metadata =====
+        end_kv_idx = end_kv_half * NUM_MATH_WG
+
+        # Convert start to KV block units (matching DeepGEMM)
+        current_q_idx = start_q
+        current_kv_idx = start_kv_half * NUM_MATH_WG
+
+        # ===== COMMON SCHEDULER STATE (before warp branches) =====
+        # Each warp role independently maintains its own copy of these
+        # variables (like DeepGEMM where each role creates its own scheduler).
+        # Pre-fetch first task (current_num_kv loaded early above for latency hiding)
+        next_q_idx = current_q_idx
+        next_kv_idx = current_kv_idx
+        next_num_kv = current_num_kv
+        # Sentinel: no previous batch (matches DeepGEMM's q_idx = batch_size)
+        q_idx = batch_size
+        # While-loop termination flag (matches DeepGEMM's fetch_next_task pattern).
+        # True if this CTA has work assigned (start != end in schedule_meta).
+        has_work = (current_q_idx != end_q_idx) | (current_kv_idx != end_kv_idx)
+
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+
+        # ===== WARP-SPECIALIZED EXECUTION =====
+
+        if is_tma_warp_0:
+            # TMA warp 0: loads Q (prefetch) + KV for group 0
+            # Matches DeepGEMM's TMA warp with kv_group_idx == 0
+            cute.arch.warpgroup_reg_dealloc(24)
+            lane_idx = tidx % 32
+
+            # Block table prefetch: 32 lanes cache block indices,
+            # distributed via shuffle. Each lane holds num_blocks_per_mma
+            # physical block indices per compute tile. (DeepGEMM L233-244)
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)  # force prefetch on first use
+
+            # Prefetch first Q before loop (like DeepGEMM line 203-204)
+            q_pipeline.producer_acquire(q_prod_state)
+            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, 0, next_q_idx)],
+                tBsB[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+                mcast_mask=b_mcast_mask,
+            )
+            cute.copy(
+                tma_atom_w,
+                tWgW[(None, next_q_idx)],
+                tWsW[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+            )
+            q_prod_state.advance()
+
+            while has_work:
+                # fetch_next_task: commit next → current
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                # Q prefetch: when batch changes, load Q for NEXT batch
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)  # force re-prefetch
+                    prefetch_next = q_idx + 1
+                    if prefetch_next < end_q_idx:
+                        q_pipeline.producer_acquire(q_prod_state)
+                        q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                        cute.copy(
+                            tma_atom_b,
+                            tBgB[(None, 0, prefetch_next)],
+                            tBsB[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                            mcast_mask=b_mcast_mask,
+                        )
+                        cute.copy(
+                            tma_atom_w,
+                            tWgW[(None, prefetch_next)],
+                            tWsW[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                        )
+                        q_prod_state.advance()
+                    elif prefetch_next == end_q_idx:
+                        if end_kv_idx > 0:
+                            q_pipeline.producer_acquire(q_prod_state)
+                            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                            cute.copy(
+                                tma_atom_b,
+                                tBgB[(None, 0, prefetch_next)],
+                                tBsB[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                                mcast_mask=b_mcast_mask,
+                            )
+                            cute.copy(
+                                tma_atom_w,
+                                tWgW[(None, prefetch_next)],
+                                tWsW[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                            )
+                            q_prod_state.advance()
+
+                # Block table prefetch for group 0 (like DeepGEMM L233-241).
+                # Each lane loads num_blocks_per_mma physical block indices
+                # for one compute tile (kv_idx counts compute tiles).
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = mBlockTable[(q_idx, base_phys + i)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+
+                # Get block indices via shuffle before barrier (like DeepGEMM L244)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile.
+                kv_pipeline_0.producer_acquire(kv_prod_state_0)
+                bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0)
+                stage = kv_prod_state_0.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_0[(None, 0, 0, phys_blks[i])],
+                        tAsA_0[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_s,
+                        tSgS_0[(None, 0, phys_blks[i])],
+                        tSsS_0[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                    )
+                kv_prod_state_0.advance()
+
+                # Advance: inline fetch_next_task
+                next_kv_idx = kv_idx + NUM_MATH_WG
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                # Update while-loop condition
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_tma_warp_1:
+            # TMA warp 1: loads KV + Scale for group 1 only
+            # Matches DeepGEMM's TMA warp with kv_group_idx == 1
+            cute.arch.warpgroup_reg_dealloc(24)
+            lane_idx = tidx % 32
+
+            # Block table prefetch for group 1
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)  # force prefetch on first use
+
+            while has_work:
+                # fetch_next_task: commit next → current
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                # New q_idx → force block table re-prefetch
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)
+
+                # Block table prefetch for group 1 (like DeepGEMM L233-241).
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = mBlockTable[(q_idx, base_phys + i)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+
+                # Get block indices via shuffle before barrier (like DeepGEMM L244)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile.
+                kv_pipeline_1.producer_acquire(kv_prod_state_1)
+                bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1)
+                stage = kv_prod_state_1.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_1[(None, 0, 0, phys_blks[i])],
+                        tAsA_1[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    cute.copy(
+                        tma_atom_s,
+                        tSgS_1[(None, 0, phys_blks[i])],
+                        tSsS_1[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                    )
+                kv_prod_state_1.advance()
+
+                # Advance: inline fetch_next_task
+                next_kv_idx = kv_idx + NUM_MATH_WG
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                # Update while-loop condition
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp_0:
+            # UMMA warp for group 0
+            # Must wait on Q pipeline: TMA operations with different
+            # barriers are NOT visibility-ordered even within the same
+            # warp. KV0 barrier arriving does not guarantee Q SMEM
+            # writes are visible.
+            cute.arch.warpgroup_reg_dealloc(24)
+
+            # TMEM: wait for math warp 0's allocation, retrieve pointer
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            if is_leader_cta:
+                num_k_blocks = cute.size(tCrA_0.shape[2])
+                q_stage_0 = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Wait for Q pipeline when batch changes
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma_0.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma_0)
+                        q_stage_0 = q_cons_state_umma_0.index
+
+                    # Process KV block for group 0 (kv_idx + 0)
+                    # Unconditional UMMA (like DeepGEMM): OOB iterations
+                    # compute on garbage data; results written to aligned
+                    # padding region in logits buffer.
+                    # Wait KV first, then TMEM empty (like DeepGEMM)
+                    kv_pipeline_0.consumer_wait(kv_cons_state_umma_0)
+                    umma_pipeline_0.producer_acquire(umma_prod_state_0)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage = kv_cons_state_umma_0.index
+                    tCtAcc_0 = tCtAcc_base_0[(None, None, None, umma_prod_state_0.index)]
+                    for k_block in cutlass.range_constexpr(num_k_blocks):
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc_0,
+                            tCrA_0[None, None, k_block, kv_stage],
+                            tCrB[None, None, k_block, q_stage_0],
+                            tCtAcc_0,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    # No consumer_release here — Math WG 0 releases
+                    kv_cons_state_umma_0.advance()
+
+                    umma_pipeline_0.producer_commit(umma_prod_state_0)
+                    umma_prod_state_0.advance()
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp_1:
+            # UMMA warp for group 1
+            # Explicitly waits on Q pipeline — critical because TMA warp 1
+            # only loads KV1, not Q. Without this wait, UMMA warp 1 can
+            # start GEMM before TMA warp 0 finishes loading Q into SMEM.
+            cute.arch.warpgroup_reg_dealloc(24)
+
+            # TMEM: wait for umma_warp_0's allocation, retrieve pointer
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            if is_leader_cta:
+                num_k_blocks_1 = cute.size(tCrA_1.shape[2])
+                q_stage_1 = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Wait for Q pipeline when batch changes
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma_1.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma_1)
+                        q_stage_1 = q_cons_state_umma_1.index
+
+                    # Process KV block for group 1 (kv_idx + 1)
+                    # Unconditional UMMA (like DeepGEMM)
+                    # Wait KV first, then TMEM empty (like DeepGEMM)
+                    kv_pipeline_1.consumer_wait(kv_cons_state_umma_1)
+                    umma_pipeline_1.producer_acquire(umma_prod_state_1)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage_1 = kv_cons_state_umma_1.index
+                    tCtAcc_1 = tCtAcc_base_1[(None, None, None, umma_prod_state_1.index)]
+                    for k_block in cutlass.range_constexpr(num_k_blocks_1):
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc_1,
+                            tCrA_1[None, None, k_block, kv_stage_1],
+                            tCrB[None, None, k_block, q_stage_1],
+                            tCtAcc_1,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    # No consumer_release here — Math WG 1 releases
+                    kv_cons_state_umma_1.advance()
+
+                    umma_pipeline_1.producer_commit(umma_prod_state_1)
+                    umma_prod_state_1.advance()
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_math_warp:
+            cute.arch.warpgroup_reg_alloc(240)
+
+            # TMEM: math warp 0 is the allocator; all math warps wait + retrieve
+            tmem.allocate(num_tmem_alloc_cols_total)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout)
+
+            local_tidx = tidx % 128
+            cC = cute.make_identity_tensor(epi_sub_mn)
+
+            if warpgroup_idx == 0:
+                # Math WG 0: process group 0
+                # Reference setup (stage 0) for m_coord
+                # flat_divide by sub-tile to get sub-tile partitions
+                tAcc_0_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)]
+                tAcc_0_ref_epi = cute.flat_divide(tAcc_0_ref, epi_sub_mn)
+                tiled_copy_ref_0 = tcgen05.make_tmem_copy(
+                    copy_atom_t2r, tAcc_0_ref_epi[(None, None, 0, 0)]
+                )
+                thr_copy_ref_0 = tiled_copy_ref_0.get_slice(local_tidx)
+                tTR_cC = thr_copy_ref_0.partition_D(cC)
+                m_coord = tTR_cC[0][0]
+
+                tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+                # Weight register cache: only first NUM_W_IN_REG
+                # per next_n slot (like DeepGEMM min(52, kNumHeads)).
+                # Remaining weights read from SMEM in epilogue.
+                # FP16 weights use half the regs, so we can fit
+                # all heads for next_n <= 3.
+                NUM_W_IN_REG = self.num_w_in_reg
+                w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype)
+                q_stage_local = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Q pipeline consumer: wait for Q+Weights SMEM
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_pipeline.consumer_release(q_cons_state)
+                            q_cons_state.advance()
+                        q_pipeline.consumer_wait(q_cons_state)
+                        q_stage_local = q_cons_state.index
+                        # Preload first NUM_W_IN_REG weights per slot
+                        for t_i in cutlass.range_constexpr(next_n):
+                            for w_j in cutlass.range_constexpr(NUM_W_IN_REG):
+                                w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
+                                    (t_i * num_heads + w_j, q_stage_local)
+                                ]
+
+                    # Process KV block for group 0 (kv_idx + 0)
+                    # Unconditional Math (like DeepGEMM): OOB results
+                    # written to aligned padding region in logits buffer.
+                    kv_pos = kv_idx * block_kv_val + m_coord
+
+                    if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                        # Skip KV wait, rely on UMMA barrier's
+                        # transitive visibility.
+                        umma_pipeline_0.consumer_wait(umma_cons_state_0)
+                    else:
+                        # Default: wait KV first to overlap lds with
+                        # UMMA computation.
+                        kv_pipeline_0.consumer_wait(kv_cons_state_math_0)
+                        sc_stage_0 = kv_cons_state_math_0.index
+                        scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        umma_pipeline_0.consumer_wait(umma_cons_state_0)
+
+                    # --- TMEM sub-tile setup ---
+                    # flat_divide accumulator by sub-tile shape;
+                    # partition once, then loop over sub-tiles.
+                    tCtAcc_c0 = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)]
+                    tAcc_c0 = tCtAcc_c0[((None, None), 0, 0)]
+                    tAcc_c0_epi = cute.flat_divide(tAcc_c0, epi_sub_mn)
+                    tc_0 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c0_epi[(None, None, 0, 0)])
+                    tr_0 = tc_0.get_slice(local_tidx)
+                    tTR_0 = tr_0.partition_S(tAcc_c0_epi)
+
+                    # --- First sub-tile LDTM + KV release ---
+                    if cutlass.const_expr(self.early_tmem_copy):
+                        # Issue first sub-tile LDTM early
+                        cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc)
+                        # Scale LDS + KV release fill latency
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_0 = kv_cons_state_math_0.index
+                            scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        kv_pipeline_0.consumer_release(kv_cons_state_math_0)
+                        kv_cons_state_math_0.advance()
+                        cute.arch.fence_view_async_tmem_load()
+                    else:
+                        # Default: scale LDS + KV release first
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_0 = kv_cons_state_math_0.index
+                            scale_val = sScales_0[(m_coord, sc_stage_0)]
+                        kv_pipeline_0.consumer_release(kv_cons_state_math_0)
+                        kv_cons_state_math_0.advance()
+                        cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc)
+                        cute.arch.fence_view_async_tmem_load()
+
+                    # --- Sub-tile compute loop ---
+                    # Each sub-tile: LDTM.xN → fence → load →
+                    # ReLU+FMA. Breaks FMA chain (16→4 per chunk)
+                    # and interleaves LDTM with FP32 compute to
+                    # reduce ShadowPipeThrottle.
+                    # Sub-tiles are within each t-slot (num_heads
+                    # // num_epi_subtiles wide). flat_divide yields
+                    # next_n * num_epi_subtiles sub-tiles total;
+                    # global index = t * num_epi_subtiles + i.
+                    subtile_n = num_heads // num_epi_subtiles
+                    if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                        packed_zero = pack_f16x2(Float16(0.0), Float16(0.0))
+                    for t in cutlass.range_constexpr(next_n):
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps0 = packed_zero
+                            ps1 = packed_zero
+                        else:
+                            s0x = cutlass.Float32(0.0)
+                            s0y = cutlass.Float32(0.0)
+                            s1x = cutlass.Float32(0.0)
+                            s1y = cutlass.Float32(0.0)
+                        for i in cutlass.range_constexpr(num_epi_subtiles):
+                            # LDTM for sub-tiles 1..N-1
+                            # (sub-tile 0 handled above)
+                            if t > 0 or i > 0:
+                                cute.copy(
+                                    tc_0,
+                                    tTR_0[(None, None, None, 0, t * num_epi_subtiles + i)],
+                                    tTR_rAcc,
+                                )
+                                cute.arch.fence_view_async_tmem_load()
+                            # Release UMMA after last LDTM+fence
+                            if t == next_n - 1 and i == num_epi_subtiles - 1:
+                                umma_pipeline_0.consumer_release(umma_cons_state_0)
+                                umma_cons_state_0.advance()
+                            acc_vec = tTR_rAcc.load()
+                            # Reg-path: weights from registers
+                            reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n))
+                            for h in cutlass.range_constexpr(0, reg_h_end, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1])
+                                    pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3])
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    w0 = w_cache[r0]
+                                    w1 = w_cache[r0 + 1]
+                                    w2 = w_cache[r0 + 2]
+                                    w3 = w_cache[r0 + 3]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                            # SMEM-path: weights from shared mem
+                            smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n)
+                            for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    pw01 = pack_f16x2(
+                                        sW[(t * num_heads + h_g, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 1, q_stage_local)],
+                                    )
+                                    pw23 = pack_f16x2(
+                                        sW[(t * num_heads + h_g + 2, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 3, q_stage_local)],
+                                    )
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    w0 = sW[(t * num_heads + h_g, q_stage_local)]
+                                    w1 = sW[(t * num_heads + h_g + 1, q_stage_local)]
+                                    w2 = sW[(t * num_heads + h_g + 2, q_stage_local)]
+                                    w3 = sW[(t * num_heads + h_g + 3, q_stage_local)]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps_sum = add_f16x2(ps0, ps1)
+                            sum_lo, sum_hi = unpack_f16x2(ps_sum)
+                            result_t = sum_lo + sum_hi
+                        else:
+                            result_t = s0x + s0y + s1x + s1y
+                        out_row = q_idx * next_n + t
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(
+                                result_t * Float16(scale_val)
+                            )
+                        else:
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val)
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+                # Release last Q stage (WG 0)
+                if q_idx < batch_size:
+                    q_pipeline.consumer_release(q_cons_state)
+                    q_cons_state.advance()
+
+            else:
+                # Math WG 1: process group 1
+                tAcc_1_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)]
+                tAcc_1_ref_epi = cute.flat_divide(tAcc_1_ref, epi_sub_mn)
+                tiled_copy_ref_1 = tcgen05.make_tmem_copy(
+                    copy_atom_t2r, tAcc_1_ref_epi[(None, None, 0, 0)]
+                )
+                thr_copy_ref_1 = tiled_copy_ref_1.get_slice(local_tidx)
+                tTR_cC = thr_copy_ref_1.partition_D(cC)
+                m_coord = tTR_cC[0][0]
+
+                tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+                # Weight register cache: only first NUM_W_IN_REG
+                # per next_n slot (like DeepGEMM min(52, kNumHeads)).
+                # FP16 weights use half the regs, so we can fit
+                # all heads for next_n <= 3.
+                NUM_W_IN_REG = self.num_w_in_reg
+                w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype)
+                q_stage_local = cutlass.Int32(0)
+
+                while has_work:
+                    # fetch_next_task: commit next → current
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+
+                    # Q pipeline consumer: wait for Q+Weights SMEM
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_pipeline.consumer_release(q_cons_state)
+                            q_cons_state.advance()
+                        q_pipeline.consumer_wait(q_cons_state)
+                        q_stage_local = q_cons_state.index
+                        # Preload first NUM_W_IN_REG weights per slot
+                        for t_i in cutlass.range_constexpr(next_n):
+                            for w_j in cutlass.range_constexpr(NUM_W_IN_REG):
+                                w_cache[t_i * NUM_W_IN_REG + w_j] = sW[
+                                    (t_i * num_heads + w_j, q_stage_local)
+                                ]
+
+                    # Process KV block for group 1 (kv_idx + 1)
+                    # Unconditional Math (like DeepGEMM)
+                    kv_idx_1 = kv_idx + 1
+
+                    kv_pos = kv_idx_1 * block_kv_val + m_coord
+
+                    if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                        umma_pipeline_1.consumer_wait(umma_cons_state_1)
+                    else:
+                        kv_pipeline_1.consumer_wait(kv_cons_state_math_1)
+                        sc_stage_1 = kv_cons_state_math_1.index
+                        scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        umma_pipeline_1.consumer_wait(umma_cons_state_1)
+
+                    # --- TMEM sub-tile setup (WG1) ---
+                    tCtAcc_c1 = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)]
+                    tAcc_c1 = tCtAcc_c1[((None, None), 0, 0)]
+                    tAcc_c1_epi = cute.flat_divide(tAcc_c1, epi_sub_mn)
+                    tc_1 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c1_epi[(None, None, 0, 0)])
+                    tr_1 = tc_1.get_slice(local_tidx)
+                    tTR_1 = tr_1.partition_S(tAcc_c1_epi)
+
+                    # --- First sub-tile LDTM + KV release (WG1) ---
+                    if cutlass.const_expr(self.early_tmem_copy):
+                        cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc)
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_1 = kv_cons_state_math_1.index
+                            scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        kv_pipeline_1.consumer_release(kv_cons_state_math_1)
+                        kv_cons_state_math_1.advance()
+                        cute.arch.fence_view_async_tmem_load()
+                    else:
+                        if cutlass.const_expr(self.remove_kv_wait_in_epilogue):
+                            sc_stage_1 = kv_cons_state_math_1.index
+                            scale_val = sScales_1[(m_coord, sc_stage_1)]
+                        kv_pipeline_1.consumer_release(kv_cons_state_math_1)
+                        kv_cons_state_math_1.advance()
+                        cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc)
+                        cute.arch.fence_view_async_tmem_load()
+
+                    # --- Sub-tile compute loop (WG1) ---
+                    subtile_n = num_heads // num_epi_subtiles
+                    if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                        packed_zero = pack_f16x2(Float16(0.0), Float16(0.0))
+                    for t in cutlass.range_constexpr(next_n):
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps0 = packed_zero
+                            ps1 = packed_zero
+                        else:
+                            s0x = cutlass.Float32(0.0)
+                            s0y = cutlass.Float32(0.0)
+                            s1x = cutlass.Float32(0.0)
+                            s1y = cutlass.Float32(0.0)
+                        for i in cutlass.range_constexpr(num_epi_subtiles):
+                            if t > 0 or i > 0:
+                                cute.copy(
+                                    tc_1,
+                                    tTR_1[(None, None, None, 0, t * num_epi_subtiles + i)],
+                                    tTR_rAcc,
+                                )
+                                cute.arch.fence_view_async_tmem_load()
+                            if t == next_n - 1 and i == num_epi_subtiles - 1:
+                                umma_pipeline_1.consumer_release(umma_cons_state_1)
+                                umma_cons_state_1.advance()
+                            acc_vec = tTR_rAcc.load()
+                            # Reg-path
+                            reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n))
+                            for h in cutlass.range_constexpr(0, reg_h_end, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1])
+                                    pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3])
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    r0 = t * NUM_W_IN_REG + h_g
+                                    w0 = w_cache[r0]
+                                    w1 = w_cache[r0 + 1]
+                                    w2 = w_cache[r0 + 2]
+                                    w3 = w_cache[r0 + 3]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                            # SMEM-path
+                            smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n)
+                            for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4):
+                                n0 = h
+                                h_g = i * subtile_n + h
+                                if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                                    pa01 = pack_f16x2(
+                                        Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1])
+                                    )
+                                    pa23 = pack_f16x2(
+                                        Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3])
+                                    )
+                                    pa01 = max_f16x2(pa01, packed_zero)
+                                    pa23 = max_f16x2(pa23, packed_zero)
+                                    pw01 = pack_f16x2(
+                                        sW[(t * num_heads + h_g, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 1, q_stage_local)],
+                                    )
+                                    pw23 = pack_f16x2(
+                                        sW[(t * num_heads + h_g + 2, q_stage_local)],
+                                        sW[(t * num_heads + h_g + 3, q_stage_local)],
+                                    )
+                                    ps0 = fma_f16x2(pa01, pw01, ps0)
+                                    ps1 = fma_f16x2(pa23, pw23, ps1)
+                                else:
+                                    a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0))
+                                    a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0))
+                                    a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0))
+                                    a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0))
+                                    w0 = sW[(t * num_heads + h_g, q_stage_local)]
+                                    w1 = sW[(t * num_heads + h_g + 1, q_stage_local)]
+                                    w2 = sW[(t * num_heads + h_g + 2, q_stage_local)]
+                                    w3 = sW[(t * num_heads + h_g + 3, q_stage_local)]
+                                    s0x, s0y = cute.arch.fma_packed_f32x2(
+                                        (a0, a1), (w0, w1), (s0x, s0y), rnd=_RND_RN
+                                    )
+                                    s1x, s1y = cute.arch.fma_packed_f32x2(
+                                        (a2, a3), (w2, w3), (s1x, s1y), rnd=_RND_RN
+                                    )
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            ps_sum = add_f16x2(ps0, ps1)
+                            sum_lo, sum_hi = unpack_f16x2(ps_sum)
+                            result_t = sum_lo + sum_hi
+                        else:
+                            result_t = s0x + s0y + s1x + s1y
+                        out_row = q_idx * next_n + t
+                        if cutlass.const_expr(self.epi_dtype == cutlass.Float16):
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(
+                                result_t * Float16(scale_val)
+                            )
+                        else:
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val)
+
+                    # Advance: inline fetch_next_task
+                    next_kv_idx = kv_idx + NUM_MATH_WG
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    # Update while-loop condition
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+                # Release last Q stage (WG 1)
+                if q_idx < batch_size:
+                    q_pipeline.consumer_release(q_cons_state)
+                    q_cons_state.advance()
+
+            # TMEM dealloc: math warps are allocator + last consumer
+            tmem.relinquish_alloc_permit()
+            tmem.free(tmem_ptr)
+
+        else:
+            cute.arch.warpgroup_reg_dealloc(24)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py
new file mode 100644
index 000000000000..4cbd26c8b892
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py
@@ -0,0 +1,773 @@
+# 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.
+"""indexer_topk row-split FP8 paged MQA logits kernel for Blackwell (SM100 family).
+
+Scores all ``next_n`` verify rows of a request against one stream of the
+paged FP8 indexer K cache, derived from the native CuTe DSL kernel
+``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``.
+
+Differences from the native kernel (same TMA/UMMA building blocks, same
+epilogue arithmetic per element):
+
+- One KV stream instead of two KV groups: a single TMA warp loads every
+  128-token compute tile into a deeper KV ring, a single UMMA warp issues
+  ``M=128 x N=next_n*H x K=128`` per tile into a two-stage TMEM ring, so the
+  MMA of tile ``j+1`` overlaps the epilogue of tile ``j``.
+- All math warpgroups (2 or 3) consume the **same** tile: warpgroup ``w``
+  owns verify rows ``[w*r, (w+1)*r)`` with ``r = next_n / num_math_wg`` (the
+  TMEM columns of those rows). Each math thread still owns one token (TMEM
+  lane) and keeps all weights of its rows in registers. Three warpgroups
+  give the schedulers more math warps to hide latency with.
+
+Per (row t, token j) the epilogue computes, exactly as the native kernel,
+``(sum_h w[t,h] * relu(acc[t,h,j])) * k_scale[j]`` with four fp32 partial
+sums over heads in steps of 4, so logits are bit-identical to the native
+launch for any row grouping.
+
+Requirements: ``next_n`` divisible by ``num_math_wg``, ``N = next_n *
+num_heads`` a multiple of 16, ``2 * N <= 512`` TMEM columns,
+``num_heads % 4 == 0``, fp32 everywhere.
+
+Scheduler: the DeepGEMM paged-MQA schedule (``schedule_meta[sm] = (q_idx,
+kv_split)`` with 256-token splits) is consumed as in the native kernel, with
+tiles advanced one at a time instead of in pairs.
+"""
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+_RND_RN = "rn"
+_TMEM_COLS = 512
+_REGS_PER_SM = 65536
+_PRODUCER_REGS = 24
+# Schedule splits are 256 tokens = 2 compute tiles of 128 tokens.
+_TILES_PER_SPLIT = 2
+
+
+class FP8MQALogitsRowSplitKernel:
+    """Persistent paged MQA logits kernel with a row-split epilogue.
+
+    Args:
+        block_kv: Compute tile in tokens (MMA M), 128.
+        phys_block_kv: Cache page in tokens; must divide ``block_kv``.
+        num_heads: Indexer heads ``H``.
+        head_dim: Indexer head dimension ``D`` (MMA K), 128.
+        next_n: Verify rows per request (divisible by ``num_math_wg``).
+        num_sms: Persistent grid size; must match the schedule.
+        num_kv_stages: KV ring depth in compute tiles.
+        num_q_stages: Query/weight ring depth.
+        num_math_wg: Math warpgroups sharing each tile (2 or 3).
+        debug: Performance-diagnosis bit flags, never for production (results
+            are invalid when non-zero): 1 skips the epilogue reduction (stores
+            one accumulator per row), 2 skips the MMA (the accumulator stage
+            is committed unwritten), 4 skips the scale TMA loads (scale 1).
+    """
+
+    def __init__(
+        self,
+        block_kv: int = 128,
+        phys_block_kv: int = 64,
+        num_heads: int = 32,
+        head_dim: int = 128,
+        next_n: int = 6,
+        num_sms: int = 148,
+        num_kv_stages: int = 6,
+        num_q_stages: int = 2,
+        num_math_wg: int = 2,
+        debug: int = 0,
+    ):
+        if block_kv % phys_block_kv != 0 or block_kv // phys_block_kv > 4:
+            raise ValueError(f"phys_block_kv={phys_block_kv} must divide {block_kv} (<= 4 pages)")
+        if num_math_wg not in (2, 3) or next_n % num_math_wg != 0:
+            raise ValueError(f"next_n={next_n} must split over {num_math_wg} math warpgroups")
+        if num_heads % 4 != 0:
+            raise ValueError("num_heads must be a multiple of 4")
+        n_cols = next_n * num_heads
+        if n_cols % 16 != 0 or n_cols > 256 or 2 * n_cols > _TMEM_COLS:
+            raise ValueError(f"next_n * num_heads = {n_cols} does not fit two TMEM stages")
+        self.block_kv = block_kv
+        self.phys_block_kv = phys_block_kv
+        self.num_blocks_per_mma = block_kv // phys_block_kv
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.next_n = next_n
+        self.num_math_wg = num_math_wg
+        self.rows_per_wg = next_n // num_math_wg
+        self.N = n_cols
+        self.num_sms = num_sms
+        self.num_kv_stages = num_kv_stages
+        self.num_q_stages = num_q_stages
+        self.num_umma_stages = 2
+        self.debug = debug
+        self.acc_dtype = cutlass.Float32
+        self.w_stage_stride = (self.N * 4 + 127) // 128 * 128 // 4
+        self.num_math_warps = 4 * num_math_wg
+        self.num_math_threads = 128 * num_math_wg
+        self.threads_per_cta = self.num_math_threads + 128
+        self.tma_warp = self.num_math_warps
+        self.umma_warp = self.num_math_warps + 2
+        # Registers left for the math warpgroups once the producer warpgroup
+        # drops to its minimum (setmaxnreg takes multiples of 8, at most 240).
+        math_regs = (_REGS_PER_SM - 128 * _PRODUCER_REGS) // self.num_math_threads
+        self.math_regs = min(240, math_regs // 8 * 8)
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.cluster_shape_mn = (1, 1)
+        self.mma_tiler_mn = (block_kv, self.N)
+
+    def _setup_mma(self, a_dtype, b_dtype, a_major, b_major):
+        self.a_dtype = a_dtype
+        self.b_dtype = b_dtype
+        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            a_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler_mn
+        )
+        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+        full_k = mma_inst_shape_k * (self.head_dim // mma_inst_shape_k)
+        self.mma_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], full_k)
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
+        )
+        self.epi_tile = self.cta_tile_shape_mnk[:2]
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma, self.mma_tiler, a_dtype, self.num_kv_stages
+        )
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma, self.mma_tiler, b_dtype, self.num_q_stages
+        )
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+        self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake)
+        self.num_tmem_alloc_cols_total = self.num_tmem_alloc_cols * self.num_umma_stages
+        if self.num_tmem_alloc_cols_total > _TMEM_COLS:
+            raise ValueError("two TMEM stages exceed the 512 TMEM columns")
+        return tiled_mma
+
+    @cute.jit
+    def __call__(
+        self,
+        kv_fused: cute.Tensor,  # [num_phys_blocks, block_bytes] uint8
+        b: cute.Tensor,  # Q: [N, head_dim, batch_size] uint8 (FP8 bytes)
+        weights: cute.Tensor,  # [N, batch_size] fp32
+        logits: cute.Tensor,  # [batch_size * next_n, max_context_len] fp32
+        block_table: cute.Tensor,  # [batch_size, max_blocks_per_seq] int32
+        context_lens: cute.Tensor,  # [batch_size] int32
+        schedule_meta: cute.Tensor,  # [num_sms + 1, 2] int32
+        num_phys_blocks: cutlass.Int32,
+        batch_size: cutlass.Int32,
+        stream: cuda.CUstream,
+    ):
+        phys_block_kv = self.phys_block_kv
+        scale_offset_elems = phys_block_kv * self.head_dim
+        kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN)
+        b = cute.recast_tensor(b, cutlass.Float8E4M3FN)
+        kv_block_stride = kv_fused.layout.stride[0]
+        kv_layout = cute.make_layout(
+            (phys_block_kv, self.head_dim, num_phys_blocks),
+            stride=(self.head_dim, 1, kv_block_stride),
+        )
+        a = cute.make_tensor(kv_fp8.iterator, kv_layout)
+        scale_fp8_layout = cute.make_layout(
+            (phys_block_kv * 4, num_phys_blocks), stride=(1, kv_block_stride)
+        )
+        scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout)
+        scales = cute.recast_tensor(scale_fp8, cutlass.Float32)
+
+        a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode()
+        b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+        tiled_mma = self._setup_mma(a.element_type, b.element_type, a_major, b_major)
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+        self.a_tma_view_layout = sm100_utils.make_smem_layout(
+            tcgen05.OperandMajorMode.K,
+            (self.block_kv, self.head_dim),
+            a.element_type,
+            self.num_kv_stages,
+        )
+        self.a_tma_view_layout = cute.tiled_divide(
+            self.a_tma_view_layout, (self.phys_block_kv, self.head_dim)
+        )
+        self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3])
+        tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom(
+            tma_load_op, a, self.a_tma_view_layout[0], (self.phys_block_kv, self.head_dim)
+        )
+
+        b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            b_op,
+            b,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        self.w_smem_layout_staged = cute.make_layout(
+            (self.N, self.num_q_stages), stride=(1, self.w_stage_stride)
+        )
+        w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0])
+        tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom(
+            tma_load_op, weights, w_smem_per_stage, self.w_smem_layout_staged.shape[:1]
+        )
+
+        self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages))
+        s_smem_per_subblock = cute.make_layout((phys_block_kv,))
+        tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom(
+            tma_load_op, scales, s_smem_per_subblock, (phys_block_kv,)
+        )
+
+        b_copy_size = cute.size_in_bytes(b.element_type, b_smem_layout)
+        w_copy_size = self.N * 4
+        scale_bytes = 0 if self.debug & 4 else phys_block_kv * 4
+        self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * (
+            phys_block_kv * self.head_dim + scale_bytes
+        )
+        self.num_q_tma_bytes = b_copy_size * atom_thr_size + w_copy_size
+
+        @cute.struct
+        class SharedStorage:
+            kv_mbar: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2]
+            q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2]
+            umma_mbar: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2]
+            tmem_holding_buf: cutlass.Int32
+
+        self.kernel(
+            tiled_mma,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_w,
+            tma_tensor_w,
+            tma_atom_s,
+            tma_tensor_s,
+            logits,
+            block_table,
+            context_lens,
+            schedule_meta,
+            batch_size,
+            self.cluster_layout_vmnk,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.w_smem_layout_staged,
+            self.s_smem_layout_staged,
+            self.a_tma_view_layout,
+            self.epi_tile,
+            SharedStorage,
+        ).launch(
+            grid=(1, 1, self.num_sms),
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,
+        tma_atom_w: cute.CopyAtom,
+        mW_tma: cute.Tensor,
+        tma_atom_s: cute.CopyAtom,
+        mS_tma: cute.Tensor,
+        mLogits: cute.Tensor,
+        mBlockTable: cute.Tensor,
+        mContextLens: cute.Tensor,
+        mScheduleMeta: cute.Tensor,
+        batch_size: cutlass.Int32,
+        cluster_layout_vmnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        w_smem_layout_staged: cute.Layout,
+        s_smem_layout_staged: cute.Layout,
+        a_tma_view_layout: cute.ComposedLayout,
+        epi_tile: cute.Tile,
+        SharedStorage: cutlass.Constexpr,
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
+        bidx, bidy, bidz = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+        is_leader_cta = mma_tile_coord_v == 0
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        tidx, _, _ = cute.arch.thread_idx()
+
+        warpgroup_idx = warp_idx // 4
+        is_math_warp = warp_idx < self.num_math_warps
+        is_tma_warp = warp_idx == self.tma_warp
+        is_umma_warp = warp_idx == self.umma_warp
+
+        NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma
+        sm_idx = bidz
+        start_q = mScheduleMeta[(sm_idx, 0)]
+        start_kv_split = mScheduleMeta[(sm_idx, 1)]
+        end_q_idx = mScheduleMeta[(sm_idx + 1, 0)]
+        end_kv_split = mScheduleMeta[(sm_idx + 1, 1)]
+        start_q_clamped = min(start_q, batch_size - 1)
+        current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv
+
+        if is_tma_warp:
+            cpasync.prefetch_descriptor(tma_atom_a)
+            cpasync.prefetch_descriptor(tma_atom_b)
+            cpasync.prefetch_descriptor(tma_atom_w)
+            cpasync.prefetch_descriptor(tma_atom_s)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        block_kv_val = self.block_kv
+        num_heads = self.num_heads
+        next_n = self.next_n
+        rows_per_wg = self.rows_per_wg
+
+        # === Pipelines ===
+        prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        # Q + weights: TMA producer, every math warp releases (lane 0 each).
+        q_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.q_mbar.data_ptr(),
+            num_stages=self.num_q_stages,
+            producer_group=prod_group,
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_math_warps),
+            tx_count=self.num_q_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        q_prod_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_q_stages
+        )
+        q_cons_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        q_cons_state_umma = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_q_stages
+        )
+        # KV + scales: TMA producer; UMMA waits (no release); every math
+        # warp reads the tile's scales and releases.
+        kv_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.kv_mbar.data_ptr(),
+            num_stages=self.num_kv_stages,
+            producer_group=prod_group,
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_math_warps),
+            tx_count=self.num_kv_scale_tma_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            tidx=tidx,
+            defer_sync=True,
+        )
+        kv_prod_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_kv_stages
+        )
+        kv_cons_state_umma = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        kv_cons_state_math = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_kv_stages
+        )
+        # Accumulator ring: UMMA producer, all math threads consume.
+        umma_pipeline = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.umma_mbar.data_ptr(),
+            num_stages=self.num_umma_stages,
+            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_math_threads),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        umma_prod_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.num_umma_stages
+        )
+        umma_cons_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_umma_stages
+        )
+
+        # TMEM: math warps + UMMA warp; math warp 0 allocates.
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=1, num_threads=self.num_math_threads + 32
+        )
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=0,
+            is_two_cta=False,
+        )
+
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+
+        sKV = smem.allocate_tensor(
+            element_type=self.a_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        sQ = smem.allocate_tensor(
+            element_type=self.b_dtype,
+            layout=b_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=b_smem_layout_staged.inner,
+        )
+        sW = smem.allocate_tensor(
+            element_type=cutlass.Float32, layout=w_smem_layout_staged, byte_alignment=128
+        )
+        sScales = smem.allocate_tensor(
+            element_type=cutlass.Float32, layout=s_smem_layout_staged, byte_alignment=128
+        )
+
+        a_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+        )
+        b_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+        )
+
+        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+        sKV_for_tma = cute.make_tensor(sKV.iterator, a_tma_view_layout.outer)
+        gA = cute.local_tile(mA_mkl, (self.phys_block_kv, self.head_dim), coord=(None, None, None))
+        tAsA, tAgA = cpasync.tma_partition(
+            tma_atom_a, 0, cute.make_layout(1), sKV_for_tma, cute.group_modes(gA, 0, 2)
+        )
+        gB_nkl = cute.local_tile(
+            mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+        )
+        tCgB = thr_mma.partition_B(gB_nkl)
+        b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            block_in_cluster_coord_vmnk[1],
+            b_cta_layout,
+            cute.group_modes(sQ, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+        tBgB = tBgB[(None, 0, None, None)]
+        tWsW, tWgW = cpasync.tma_partition(
+            tma_atom_w,
+            0,
+            cute.make_layout((1,)),
+            cute.group_modes(sW, 0, 1),
+            cute.group_modes(mW_tma, 0, 1),
+        )
+        s_tma_view_layout = cute.make_layout(
+            (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages),
+            stride=(1, self.phys_block_kv, self.block_kv),
+        )
+        sScales_for_tma = cute.make_tensor(sScales.iterator, s_tma_view_layout)
+        gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None))
+        tSsS, tSgS = cpasync.tma_partition(tma_atom_s, 0, cute.make_layout(1), sScales_for_tma, gS)
+
+        tCrA = tiled_mma.make_fragment_A(sKV)
+        tCrB = tiled_mma.make_fragment_B(sQ)
+        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+        acc_shape_staged = cute.append(acc_shape, self.num_umma_stages)
+        tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged)
+
+        # Epilogue: one TMEM load per warpgroup and tile covering its
+        # rows_per_wg contiguous verify rows (rows_per_wg * num_heads columns).
+        epi_sub_mn = (epi_tile[0], rows_per_wg * num_heads)
+        copy_atom_t2r = sm100_utils.get_tmem_load_op(
+            self.cta_tile_shape_mnk,
+            utils.LayoutEnum.ROW_MAJOR,
+            self.acc_dtype,
+            self.acc_dtype,
+            epi_sub_mn,
+            use_2cta_instrs,
+        )
+
+        # Scheduler state (tiles of block_kv tokens, advanced one at a time).
+        end_kv_idx = end_kv_split * _TILES_PER_SPLIT
+        next_q_idx = start_q
+        next_kv_idx = start_kv_split * _TILES_PER_SPLIT
+        next_num_kv = current_num_kv
+        q_idx = batch_size
+        has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+
+        if is_tma_warp:
+            cute.arch.warpgroup_reg_dealloc(_PRODUCER_REGS)
+            lane_idx = tidx % 32
+            cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)]
+            kv_blk_ptr = cutlass.Int32(32)
+
+            q_pipeline.producer_acquire(q_prod_state)
+            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, 0, next_q_idx)],
+                tBsB[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+                mcast_mask=b_mcast_mask,
+            )
+            cute.copy(
+                tma_atom_w,
+                tWgW[(None, next_q_idx)],
+                tWsW[(None, q_prod_state.index)],
+                tma_bar_ptr=q_bar,
+            )
+            q_prod_state.advance()
+
+            while has_work:
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                if q_idx != q_idx_old:
+                    kv_blk_ptr = cutlass.Int32(32)
+                    prefetch_next = q_idx + 1
+                    if prefetch_next < end_q_idx:
+                        q_pipeline.producer_acquire(q_prod_state)
+                        q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                        cute.copy(
+                            tma_atom_b,
+                            tBgB[(None, 0, prefetch_next)],
+                            tBsB[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                            mcast_mask=b_mcast_mask,
+                        )
+                        cute.copy(
+                            tma_atom_w,
+                            tWgW[(None, prefetch_next)],
+                            tWsW[(None, q_prod_state.index)],
+                            tma_bar_ptr=q_bar,
+                        )
+                        q_prod_state.advance()
+                    elif prefetch_next == end_q_idx:
+                        if end_kv_idx > 0:
+                            q_pipeline.producer_acquire(q_prod_state)
+                            q_bar = q_pipeline.producer_get_barrier(q_prod_state)
+                            cute.copy(
+                                tma_atom_b,
+                                tBgB[(None, 0, prefetch_next)],
+                                tBsB[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                                mcast_mask=b_mcast_mask,
+                            )
+                            cute.copy(
+                                tma_atom_w,
+                                tWgW[(None, prefetch_next)],
+                                tWsW[(None, q_prod_state.index)],
+                                tma_bar_ptr=q_bar,
+                            )
+                            q_prod_state.advance()
+
+                # Page-table prefetch: lane l caches the pages of tile kv_idx + l.
+                if kv_blk_ptr == 32:
+                    kv_blk_ptr = cutlass.Int32(0)
+                    prefetch_kv = kv_idx + lane_idx
+                    if prefetch_kv < num_kv:
+                        base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = mBlockTable[(q_idx, base_phys + i)]
+                    else:
+                        for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                            cached_blks[i] = cutlass.Int32(0)
+                phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr)
+                kv_blk_ptr = kv_blk_ptr + 1
+
+                kv_pipeline.producer_acquire(kv_prod_state)
+                bar = kv_pipeline.producer_get_barrier(kv_prod_state)
+                stage = kv_prod_state.index
+                for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA):
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA[(None, 0, 0, phys_blks[i])],
+                        tAsA[(None, i, stage)],
+                        tma_bar_ptr=bar,
+                        mcast_mask=a_mcast_mask,
+                    )
+                    if cutlass.const_expr(not self.debug & 4):
+                        cute.copy(
+                            tma_atom_s,
+                            tSgS[(None, 0, phys_blks[i])],
+                            tSsS[(None, i, stage)],
+                            tma_bar_ptr=bar,
+                        )
+                kv_prod_state.advance()
+
+                next_kv_idx = kv_idx + 1
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_umma_warp:
+            cute.arch.warpgroup_reg_dealloc(_PRODUCER_REGS)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            if is_leader_cta:
+                num_k_blocks = cute.size(tCrA.shape[2])
+                q_stage = cutlass.Int32(0)
+                while has_work:
+                    q_idx_old = q_idx
+                    q_idx = next_q_idx
+                    kv_idx = next_kv_idx
+                    num_kv = next_num_kv
+                    if q_idx != q_idx_old:
+                        if q_idx_old < batch_size:
+                            q_cons_state_umma.advance()
+                        q_pipeline.consumer_wait(q_cons_state_umma)
+                        q_stage = q_cons_state_umma.index
+
+                    kv_pipeline.consumer_wait(kv_cons_state_umma)
+                    umma_pipeline.producer_acquire(umma_prod_state)
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    kv_stage = kv_cons_state_umma.index
+                    tCtAcc = tCtAcc_base[(None, None, None, umma_prod_state.index)]
+                    if cutlass.const_expr(not self.debug & 2):
+                        for k_block in cutlass.range_constexpr(num_k_blocks):
+                            cute.gemm(
+                                tiled_mma,
+                                tCtAcc,
+                                tCrA[None, None, k_block, kv_stage],
+                                tCrB[None, None, k_block, q_stage],
+                                tCtAcc,
+                            )
+                            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    kv_cons_state_umma.advance()
+                    umma_pipeline.producer_commit(umma_prod_state)
+                    umma_prod_state.advance()
+
+                    next_kv_idx = kv_idx + 1
+                    if next_kv_idx >= num_kv:
+                        next_q_idx = q_idx + 1
+                        next_kv_idx = 0
+                        if next_q_idx < batch_size:
+                            next_num_kv = (
+                                mContextLens[next_q_idx] + block_kv_val - 1
+                            ) // block_kv_val
+                    has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+        elif is_math_warp:
+            cute.arch.warpgroup_reg_alloc(self.math_regs)
+            tmem.allocate(self.num_tmem_alloc_cols_total)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout)
+            local_tidx = tidx % 128
+            cC = cute.make_identity_tensor(epi_sub_mn)
+
+            tAcc_ref = tCtAcc_base[(None, None, None, 0)][((None, None), 0, 0)]
+            tAcc_ref_epi = cute.flat_divide(tAcc_ref, epi_sub_mn)
+            tiled_copy_ref = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_ref_epi[(None, None, 0, 0)])
+            thr_copy_ref = tiled_copy_ref.get_slice(local_tidx)
+            tTR_cC = thr_copy_ref.partition_D(cC)
+            m_coord = tTR_cC[0][0]
+            tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype)
+
+            # This warpgroup's verify rows: [row0, row0 + rows_per_wg).
+            row0 = warpgroup_idx * rows_per_wg
+            w_cache = cute.make_rmem_tensor(rows_per_wg * num_heads, cutlass.Float32)
+            q_stage_local = cutlass.Int32(0)
+
+            while has_work:
+                q_idx_old = q_idx
+                q_idx = next_q_idx
+                kv_idx = next_kv_idx
+                num_kv = next_num_kv
+
+                if q_idx != q_idx_old:
+                    if q_idx_old < batch_size:
+                        q_pipeline.consumer_release(q_cons_state)
+                        q_cons_state.advance()
+                    q_pipeline.consumer_wait(q_cons_state)
+                    q_stage_local = q_cons_state.index
+                    for t_i in cutlass.range_constexpr(rows_per_wg):
+                        for w_j in cutlass.range_constexpr(num_heads):
+                            w_cache[t_i * num_heads + w_j] = sW[
+                                ((row0 + t_i) * num_heads + w_j, q_stage_local)
+                            ]
+
+                kv_pos = kv_idx * block_kv_val + m_coord
+                kv_pipeline.consumer_wait(kv_cons_state_math)
+                scale_val = sScales[(m_coord, kv_cons_state_math.index)]
+                umma_pipeline.consumer_wait(umma_cons_state)
+                kv_pipeline.consumer_release(kv_cons_state_math)
+                kv_cons_state_math.advance()
+
+                tAcc = tCtAcc_base[(None, None, None, umma_cons_state.index)]
+                tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0)], epi_sub_mn)
+                tc = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)])
+                tr = tc.get_slice(local_tidx)
+                tTR = tr.partition_S(tAcc_epi)
+
+                # One load of all this warpgroup's columns, one wait, then the
+                # accumulator stage goes back to the UMMA warp.
+                cute.copy(tc, tTR[(None, None, None, 0, warpgroup_idx)], tTR_rAcc)
+                cute.arch.fence_view_async_tmem_load()
+                umma_pipeline.consumer_release(umma_cons_state)
+                umma_cons_state.advance()
+                acc_vec = tTR_rAcc.load()
+
+                if cutlass.const_expr(self.debug & 1):
+                    for t in cutlass.range_constexpr(rows_per_wg):
+                        mLogits[(q_idx * next_n + row0 + t, kv_pos)] = (
+                            acc_vec[t * num_heads] * scale_val
+                        )
+                for t in cutlass.range_constexpr(0 if self.debug & 1 else rows_per_wg):
+                    s0x = cutlass.Float32(0.0)
+                    s0y = cutlass.Float32(0.0)
+                    s1x = cutlass.Float32(0.0)
+                    s1y = cutlass.Float32(0.0)
+                    for h in cutlass.range_constexpr(0, num_heads, 4):
+                        c0 = t * num_heads + h
+                        a0 = cutlass.max(acc_vec[c0], cutlass.Float32(0.0))
+                        a1 = cutlass.max(acc_vec[c0 + 1], cutlass.Float32(0.0))
+                        a2 = cutlass.max(acc_vec[c0 + 2], cutlass.Float32(0.0))
+                        a3 = cutlass.max(acc_vec[c0 + 3], cutlass.Float32(0.0))
+                        r0 = t * num_heads + h
+                        s0x, s0y = cute.arch.fma_packed_f32x2(
+                            (a0, a1), (w_cache[r0], w_cache[r0 + 1]), (s0x, s0y), rnd=_RND_RN
+                        )
+                        s1x, s1y = cute.arch.fma_packed_f32x2(
+                            (a2, a3), (w_cache[r0 + 2], w_cache[r0 + 3]), (s1x, s1y), rnd=_RND_RN
+                        )
+                    result_t = s0x + s0y + s1x + s1y
+                    mLogits[(q_idx * next_n + row0 + t, kv_pos)] = result_t * scale_val
+
+                next_kv_idx = kv_idx + 1
+                if next_kv_idx >= num_kv:
+                    next_q_idx = q_idx + 1
+                    next_kv_idx = 0
+                    if next_q_idx < batch_size:
+                        next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val
+                has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx)
+
+            if q_idx < batch_size:
+                q_pipeline.consumer_release(q_cons_state)
+                q_cons_state.advance()
+
+            tmem.relinquish_alloc_permit()
+            tmem.free(tmem_ptr)
+
+        else:
+            cute.arch.warpgroup_reg_dealloc(_PRODUCER_REGS)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/indexer_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/indexer_decode.py
new file mode 100644
index 000000000000..b06d1d3609b8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/indexer_decode.py
@@ -0,0 +1,296 @@
+# 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.
+"""indexer_topk: GEN indexer logits and top-k selection for target full-indexer layers.
+
+The index_selection contract: native projected indexer query/weights, the
+immutable indexer K cache view and live KV lengths in; native-compatible
+logical top-k IDs ``[R, index_topk]`` (``-1`` marks invalid slots) out. Cache
+append, cross-layer reuse, MTP last-accepted-row selection and GVR prior
+state stay with the native caller.
+
+Every verify row of a request is scored in one paged-MQA
+logits launch, so the indexer K cache is streamed once per request instead of
+once per row group (the native caller splits ``next_n`` rows into groups of at
+most four and re-reads the cache for each group). The kernel is an indexer_topk-owned
+variant of the native CuTe-DSL FP8 paged-MQA logits kernel with a bounded
+register cache of epilogue weights (see ``fp8_paged_mqa_logits_rows``); its
+per-element arithmetic does not depend on the row grouping, so the logits are
+bit-identical to the native split launch. Selection is the layer's native
+``TopK`` module (heuristic self-sampling GVR, as ``policy`` requires).
+"""
+
+from typing import Dict, Optional
+
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.logger import logger
+
+from . import policy
+
+__all__ = [
+    "DECISION_COUNTS",
+    "paged_mqa_logits",
+    "decode_topk",
+    "try_indexer_decode",
+]
+
+# Selection engine of decode_topk: "native" (the layer's TopK module) or
+# "radix" (topk_select.radix_topk_decode, exact, graph-safe).
+DEFAULT_SELECTOR = "native"
+
+# Hook decisions on this rank, keyed by "engaged" or the fallback reason. The
+# hook runs once per eager forward and once per CUDA-graph capture (never per
+# replay), so these counts are host-side and cost nothing in a replayed graph.
+DECISION_COUNTS: Dict[str, int] = {}
+
+
+def _record_decision(decision: policy.Decision) -> None:
+    key = "engaged" if decision.eligible else decision.reason
+    count = DECISION_COUNTS.get(key, 0)
+    DECISION_COUNTS[key] = count + 1
+    if count == 0:
+        state = "engaged" if decision.eligible else f"native fallback ({decision.reason})"
+        logger.info(
+            f"[glm_moe_dsa] indexer_topk/indexer_decode: first '{state}' decision on this rank"
+        )
+
+
+def paged_mqa_logits(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    weights: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    schedule_meta: torch.Tensor,
+    max_seq_len: int,
+    config=None,
+) -> torch.Tensor:
+    """Score all verify rows of every request in one launch.
+
+    Args:
+        q: ``[B, next_n, H, D]`` float8_e4m3fn indexer queries.
+        k_cache: ``[num_blocks, tokens_per_block, 1, D + 4]`` uint8 fused
+            indexer K page view of this layer (FP8 values then fp32 scales
+            per page).
+        weights: ``[B * next_n, H]`` float32 per-head weights (already scaled
+            by the query scale and softmax factor).
+        context_lens: ``[B]`` int32 indexer KV lengths.
+        block_table: ``[B, max_blocks]`` int32 page table of this layer.
+        schedule_meta: ``[num_sms + 1, 2]`` int32 DeepGEMM schedule built from
+            ``context_lens.view(-1, 1)`` with 64-token schedule blocks.
+        max_seq_len: Capture-stable logits width (indexer max sequence length).
+        config: Optional ``logits_runner.LogitsConfig``; None selects the
+            default indexer_topk configuration.
+
+    Returns:
+        ``[B * next_n, max_seq_len]`` float32 logits view of the native
+        256-column-aligned arena owned by indexer_topk. Columns at or beyond each row's
+        causal length are unspecified.
+    """
+    from .logits_runner import DEFAULT_CONFIG, paged_mqa_logits_rows
+
+    return paged_mqa_logits_rows(
+        q,
+        k_cache,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_seq_len,
+        config=DEFAULT_CONFIG if config is None else config,
+    )
+
+
+def decode_topk(
+    *,
+    q: torch.Tensor,
+    weights: torch.Tensor,
+    k_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    context_lens: torch.Tensor,
+    kv_lens: torch.Tensor,
+    schedule_meta: torch.Tensor,
+    max_seq_len: int,
+    top_k: torch.nn.Module,
+    out: torch.Tensor,
+    radix_aux_indices: Optional[torch.Tensor],
+    radix_aux_logits: Optional[torch.Tensor],
+    logits_config=None,
+    selector: Optional[str] = None,
+) -> torch.Tensor:
+    """Run indexer_topk logits + selection for uniform-width generation rows.
+
+    Args:
+        q: ``[B, next_n, H, D]`` float8_e4m3fn indexer queries.
+        weights: ``[B * next_n, H]`` float32 per-head weights.
+        k_cache: Fused indexer K page view, see :func:`paged_mqa_logits`.
+        block_table: ``[B, max_blocks]`` int32 page table.
+        context_lens: ``[B]`` int32 indexer (score-column) KV lengths.
+        kv_lens: ``[B]`` int32 logical KV lengths (equal to ``context_lens``
+            without index compression).
+        schedule_meta: ``[num_sms + 1, 2]`` int32 schedule, see
+            :func:`paged_mqa_logits`.
+        max_seq_len: Capture-stable indexer max sequence length.
+        top_k: The layer's native ``TopK`` module (selection policy).
+        out: ``[B * next_n, index_topk]`` int32 caller-owned output.
+        radix_aux_indices: Native radix scratch (used only by its fallback).
+        radix_aux_logits: Native radix scratch (used only by its fallback).
+        logits_config: Optional ``logits_runner.LogitsConfig``.
+        selector: ``"native"`` or ``"radix"``; None uses ``DEFAULT_SELECTOR``.
+            ``"radix"`` requires an uncompressed index (``compress_ratio`` 1).
+
+    Returns:
+        ``out`` with each row's logical top-k IDs and ``-1`` padding.
+    """
+    next_n = q.shape[1]
+    logits = paged_mqa_logits(
+        q,
+        k_cache,
+        weights,
+        context_lens,
+        block_table,
+        schedule_meta,
+        max_seq_len,
+        config=logits_config,
+    )
+    selector = DEFAULT_SELECTOR if selector is None else selector
+    if selector == "radix" and getattr(top_k, "compress_ratio", 1) == 1:
+        from .topk_select import default_chunks, radix_topk_decode
+
+        num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
+        return radix_topk_decode(
+            logits, kv_lens, out, next_n, chunks=default_chunks(out.shape[0], num_sms)
+        )
+    top_k(
+        logits,
+        out,
+        is_prefill=False,
+        sequence_lengths=kv_lens,
+        scan_lengths=context_lens,
+        row_kv_lens=None,
+        next_n=next_n,
+        max_seq_len=max_seq_len,
+        gvr_ext_kwargs=None,
+        radix_aux_indices=radix_aux_indices,
+        radix_aux_logits=radix_aux_logits,
+    )
+    return out
+
+
+def try_indexer_decode(
+    indexer,
+    metadata,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+) -> bool:
+    """Adapter entry: run the GEN region of ``Indexer.sparse_attn_indexer``.
+
+    The eligibility decision (``policy.decide``) reads only host state and
+    tensor metadata and runs before any device work, so a ``False`` return
+    leaves no side effect and the caller runs the native path. indexer_topk engages only
+    inside its GPU-verified envelope; see ``policy``. Each distinct decision is
+    logged once per rank and counted in ``DECISION_COUNTS``.
+
+    Args:
+        indexer: The native ``Indexer`` of a full-indexer layer.
+        metadata: The DSA attention metadata of this forward.
+        q_fp8: ``[num_tokens, H, D]`` float8_e4m3fn projected queries of the
+            phase passed to the indexer.
+        weights: ``[num_tokens, H]`` float32 projected per-head weights.
+        token_offset: First generation row in ``q_fp8`` / ``weights``.
+        num_gen_tokens: Generation rows in this call.
+        out: ``[num_gen_tokens, index_topk]`` int32 slice of the native
+            top-k buffer to fill.
+
+    Returns:
+        True when indexer_topk filled ``out``; False when the native path must run.
+    """
+    from ...dispatch import log_capability
+
+    num_contexts = metadata.num_contexts
+    num_generations = metadata.num_generations
+    gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
+    uniform = num_generations > 0 and bool(
+        (gen_seq_lens == num_gen_tokens // num_generations).all()
+    )
+    top_k = indexer.top_k
+    # Host-side view of the layer's page pool; no device work.
+    k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(indexer.layer_idx)
+    block_table = metadata.indexer_k_cache_block_offsets
+    mapping = getattr(metadata, "mapping", None)
+    max_seq_len = metadata.get_indexer_max_seq_len()
+    decision = policy.decide(
+        sm_version=get_sm_version(),
+        cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE,
+        native_uses_cute_dsl_logits=indexer.use_cute_dsl_paged_mqa_logits,
+        use_fp4_indexer=indexer.use_fp4,
+        use_gvr_emission=indexer.use_gvr_emission,
+        needs_gvr_prior=top_k.needs_gvr_prior,
+        is_ragged_verify=metadata.is_ragged_verify,
+        expanded_mtp_buffers=metadata.use_expanded_buffers_for_mtp,
+        in_mtp_draft_loop=metadata.in_mtp_draft_loop,
+        attention_dp=bool(mapping is not None and mapping.enable_attention_dp),
+        topk_decode_impl=top_k.decode_implementation.value,
+        topk_self_sampling=top_k.gvr_self_sampling,
+        index_topk=top_k.top_k,
+        compress_ratio=top_k.compress_ratio,
+        num_heads=indexer.n_heads,
+        head_dim=indexer.head_dim,
+        query_dtype=str(q_fp8.dtype),
+        query_shape=tuple(q_fp8.shape),
+        query_contiguous=q_fp8.is_contiguous(),
+        weights_dtype=str(weights.dtype),
+        weights_shape=tuple(weights.shape),
+        weights_contiguous=weights.is_contiguous(),
+        k_cache_dtype=str(k_cache.dtype),
+        k_cache_shape=tuple(k_cache.shape),
+        block_table_cols=block_table.shape[1],
+        max_seq_len=max_seq_len,
+        num_contexts=num_contexts,
+        token_offset=token_offset,
+        num_gen_tokens=num_gen_tokens,
+        num_generations=num_generations,
+        uniform_verify_width=uniform,
+    )
+    log_capability("indexer_topk", "indexer_decode", decision.eligible, decision.reason)
+    _record_decision(decision)
+    if not decision.eligible:
+        return False
+
+    next_n = num_gen_tokens // num_generations
+    gen_rows = slice(token_offset, token_offset + num_gen_tokens)
+    gen_reqs = slice(num_contexts, num_contexts + num_generations)
+    context_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+    assert context_lens is not None
+    decode_topk(
+        q=q_fp8[gen_rows].view(num_generations, next_n, indexer.n_heads, indexer.head_dim),
+        weights=weights[gen_rows],
+        k_cache=k_cache,
+        block_table=block_table[gen_reqs],
+        context_lens=context_lens,
+        kv_lens=metadata.kv_lens_cuda_runtime[gen_reqs],
+        schedule_meta=metadata.scheduler_metadata_buffer,
+        max_seq_len=max_seq_len,
+        top_k=top_k,
+        out=out,
+        radix_aux_indices=metadata.radix_aux_indices,
+        radix_aux_logits=metadata.radix_aux_logits,
+    )
+    return True
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/logits_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/logits_runner.py
new file mode 100644
index 000000000000..9e349dd2ee90
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/logits_runner.py
@@ -0,0 +1,237 @@
+# 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.
+"""Compile cache and launcher of the indexer_topk all-rows FP8 paged-MQA logits kernels.
+
+Mirrors the native ``CuteDSLPagedMQALogitsRunner`` (fake-tensor compile with
+TVM FFI, persistent graph-stable logits arena) for the two indexer_topk kernels, with
+fp32 accumulation, epilogue and output only (the native runner's
+configuration for this model):
+
+- ``rowsplit`` (:mod:`fp8_paged_mqa_logits_rowsplit`): one KV stream, both
+  math warpgroups split the verify rows of each tile, two-stage TMEM ring.
+  Needs an even ``next_n`` and ``2 * next_n * H <= 512``.
+- ``rows`` (:mod:`fp8_paged_mqa_logits_rows`): the native two-group kernel
+  at the full verify width with a bounded register weight cache.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+
+from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers
+
+from .fp8_paged_mqa_logits_rows import FP8MQALogitsKernel
+from .fp8_paged_mqa_logits_rowsplit import FP8MQALogitsRowSplitKernel
+
+COMPUTE_BLOCK_KV = 128
+# Two math warpgroups per CTA each take one 128-token tile per step.
+SPLIT_KV = 2 * COMPUTE_BLOCK_KV
+LOGITS_BUFFER_NAME = "indexer_topk_mqa_logits"
+
+
+@dataclass(frozen=True)
+class LogitsConfig:
+    """Compile-time knobs of the indexer_topk logits kernel.
+
+    Attributes:
+        variant: ``"rowsplit"`` (default; falls back to ``"rows"`` when the
+            shape does not allow it) or ``"rows"``.
+        num_kv_stages: KV ring depth of the ``rowsplit`` kernel.
+        num_q_stages: Query/weight ring depth of the ``rowsplit`` kernel.
+        num_math_wg: Math warpgroups of the ``rowsplit`` kernel (2 or 3).
+        debug: ``rowsplit`` diagnosis mode (non-zero values break results).
+        num_w_in_reg: ``rows`` kernel: register-cached epilogue weights per
+            verify row; None keeps the upstream heuristic.
+        num_epi_subtiles: TMEM-load sub-tiles per verify row (1, 2, 4 or 8
+            for 32 heads).
+        max_kv_pipeline: Fill shared memory with KV stages.
+        smem_subpartition_opt: Pad shared memory so weights/scales sit in the
+            second sub-partition.
+        early_tmem_copy: Issue the first TMEM load before the KV release.
+    """
+
+    variant: str = "rowsplit"
+    num_kv_stages: int = 6
+    num_q_stages: int = 2
+    num_math_wg: int = 2
+    debug: int = 0
+    num_w_in_reg: Optional[int] = 8
+    num_epi_subtiles: int = 1
+    max_kv_pipeline: bool = False
+    smem_subpartition_opt: bool = False
+    early_tmem_copy: bool = False
+
+
+# "rows" with 8 register-cached weights is the default variant; "rowsplit"
+# is the alternative for shorter indexer contexts.
+DEFAULT_CONFIG = LogitsConfig(variant="rows", num_w_in_reg=8)
+
+_KERNEL_CACHE: Dict[Tuple, object] = {}
+_NUM_SMS: Dict[int, int] = {}
+
+
+def _num_sms(device: torch.device) -> int:
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    if index not in _NUM_SMS:
+        _NUM_SMS[index] = torch.cuda.get_device_properties(index).multi_processor_count
+    return _NUM_SMS[index]
+
+
+def rowsplit_supported(next_n: int, num_heads: int, num_math_wg: int = 2) -> bool:
+    """Whether the ``rowsplit`` kernel can hold this verify width."""
+    n_cols = next_n * num_heads
+    return next_n % num_math_wg == 0 and n_cols % 16 == 0 and n_cols <= 256
+
+
+def _compile(
+    phys_block_kv: int,
+    num_heads: int,
+    head_dim: int,
+    next_n: int,
+    num_sms: int,
+    config: LogitsConfig,
+):
+    n_cols = next_n * num_heads
+    block_bytes = phys_block_kv * (head_dim + 4)
+    sym_num_phys_blocks = cute.sym_int()
+    sym_b = cute.sym_int()
+    max_ctx = cute.sym_int()
+    max_blocks_per_seq = cute.sym_int()
+    num_ctas = cute.sym_int()
+    kv_fake = cute.runtime.make_fake_tensor(
+        cutlass.Uint8, (sym_num_phys_blocks, block_bytes), stride=(cute.sym_int64(), 1)
+    )
+    q_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Uint8, (n_cols, head_dim, sym_b), stride_order=(1, 0, 2)
+    )
+    w_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float32, (n_cols, sym_b), stride_order=(0, 1)
+    )
+    logits_fake = cute.runtime.make_fake_tensor(
+        cutlass.Float32, (cute.sym_int(), max_ctx), stride=(cute.sym_int64(), 1)
+    )
+    bt_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Int32, (sym_b, max_blocks_per_seq), stride_order=(1, 0)
+    )
+    cl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,))
+    sm_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Int32, (num_ctas, 2), stride_order=(1, 0)
+    )
+    fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    if config.variant == "rowsplit" and rowsplit_supported(next_n, num_heads, config.num_math_wg):
+        kernel = FP8MQALogitsRowSplitKernel(
+            block_kv=COMPUTE_BLOCK_KV,
+            phys_block_kv=phys_block_kv,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            next_n=next_n,
+            num_sms=num_sms,
+            num_kv_stages=config.num_kv_stages,
+            num_q_stages=config.num_q_stages,
+            num_math_wg=config.num_math_wg,
+            debug=config.debug,
+        )
+    else:
+        kernel = FP8MQALogitsKernel(
+            block_kv=COMPUTE_BLOCK_KV,
+            phys_block_kv=phys_block_kv,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            next_n=next_n,
+            num_sms=num_sms,
+            early_tmem_copy=config.early_tmem_copy,
+            smem_subpartition_opt=config.smem_subpartition_opt,
+            max_kv_pipeline=config.max_kv_pipeline,
+            num_epi_subtiles=config.num_epi_subtiles,
+            num_w_in_reg=config.num_w_in_reg,
+        )
+    return cute.compile(
+        kernel,
+        kv_fake,
+        q_fake,
+        w_fake,
+        logits_fake,
+        bt_fake,
+        cl_fake,
+        sm_fake,
+        cutlass.Int32(1),
+        cutlass.Int32(1),
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+def paged_mqa_logits_rows(
+    q: torch.Tensor,
+    kv_fused: torch.Tensor,
+    weights: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    schedule_meta: torch.Tensor,
+    max_context_len: int,
+    config: LogitsConfig = DEFAULT_CONFIG,
+) -> torch.Tensor:
+    """Score every verify row of every request in one launch.
+
+    Args:
+        q: ``[B, next_n, H, D]`` float8_e4m3fn queries, ``next_n * H <= 256``.
+        kv_fused: ``[num_blocks, phys_block_kv, 1, D + 4]`` uint8 page view
+            (FP8 values then fp32 scales per page); dim-0 may be strided.
+        weights: ``[B * next_n, H]`` float32 epilogue weights.
+        context_lens: ``[B]`` int32 KV lengths.
+        block_table: ``[B, max_blocks]`` int32 page table.
+        schedule_meta: ``[num_sms + 1, 2]`` int32 DeepGEMM schedule built from
+            ``context_lens.view(-1, 1)`` with 64-token schedule blocks.
+        max_context_len: Capture-stable logits width.
+        config: Compile-time kernel knobs.
+
+    Returns:
+        ``[B * next_n, max_context_len]`` float32 view of a 256-column-aligned,
+        graph-stable arena; columns at or beyond a row's KV length are
+        unspecified.
+    """
+    batch, next_n, num_heads, head_dim = q.shape
+    phys_block_kv = kv_fused.shape[1]
+    num_phys_blocks = kv_fused.shape[0]
+    num_sms = _num_sms(q.device)
+    aligned = (max_context_len + SPLIT_KV - 1) // SPLIT_KV * SPLIT_KV
+    logits = get_memory_buffers().get_buffer(
+        [batch * next_n, aligned],
+        torch.float32,
+        buffer_name=LOGITS_BUFFER_NAME,
+        reserve_buffer=torch.cuda.is_current_stream_capturing(),
+    )[:, :max_context_len]
+    key = (phys_block_kv, num_heads, head_dim, next_n, num_sms, config)
+    compiled = _KERNEL_CACHE.get(key)
+    if compiled is None:
+        compiled = _compile(phys_block_kv, num_heads, head_dim, next_n, num_sms, config)
+        _KERNEL_CACHE[key] = compiled
+    q_3d = q.reshape(batch, next_n * num_heads, head_dim).permute(1, 2, 0)
+    w_2d = weights.reshape(batch, next_n * num_heads).t()
+    compiled(
+        kv_fused.reshape(num_phys_blocks, -1),
+        q_3d.view(torch.uint8),
+        w_2d,
+        logits,
+        block_table,
+        context_lens,
+        schedule_meta,
+        num_phys_blocks,
+        batch,
+    )
+    return logits
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/policy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/policy.py
new file mode 100644
index 000000000000..6db4932a1837
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/policy.py
@@ -0,0 +1,240 @@
+# 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.
+"""Eligibility policy of the indexer_topk indexer decode package (the index_selection contract).
+
+This module imports only the standard library so the decision table can be
+unit-tested without torch or the compiled extensions. Every input is a plain
+host value that the caller reads from the native ``Indexer``, its ``TopK``
+module, the DSA metadata and the tensors' dtypes/shapes before any indexer_topk kernel
+runs; an ineligible call returns a reason and the caller keeps the native
+path.
+
+indexer_topk engages only inside the envelope that its GPU tests verify (sm_103, the
+GLM-5.2 indexer shape, 64-token FP8 pages, ``next_n = 1 + K`` with K = 5,
+up to four generation requests, the native self-sampling GVR selection,
+indexer ``max_seq_len`` up to ``2**20 + 14``). Every precondition of the indexer_topk
+kernel launch that the caller can observe on the host is checked here, so an
+eligible call cannot fail after the decision.
+"""
+
+from dataclasses import dataclass
+from typing import Tuple
+
+# GPU-verified support envelope (the indexer_topk kernel tests on sm_103).
+SUPPORTED_SM_VERSIONS = (103,)
+SUPPORTED_NUM_HEADS = 32
+SUPPORTED_HEAD_DIM = 128
+SUPPORTED_TOKENS_PER_BLOCK = 64
+# 1 target row + K = 5 MTP draft rows per request.
+SUPPORTED_NEXT_N = 6
+MAX_GENERATIONS = 4
+# The indexer max_seq_len is the KV-cache manager's max_seq_len: the engine's
+# net max_seq_len plus (tokens_per_gen_step - 1) with the overlap scheduler,
+# plus (max_draft_len - 1) extra one-model KV tokens, plus
+# (tokens_per_gen_step - 1). For K = 5 that slack is 5 + 4 + 5 = 14, so a 2**20
+# net envelope serves an indexer max_seq_len of 2**20 + 14. Live KV lengths
+# stay at or below it; the logits arena and the native TopK envelope are sized
+# from it.
+MTP_K5_MAX_SEQ_LEN_SLACK = 14
+MAX_INDEXER_SEQ_LEN = (1 << 20) + MTP_K5_MAX_SEQ_LEN_SLACK
+# The logits kernel computes 128-token tiles and loads both page IDs of every
+# tile it computes, so a request whose last tile covers one page reads the
+# table entry after that page. Requiring the table width to be a multiple of
+# the pages per tile keeps that read inside the request's row for every KV
+# length the row can describe.
+KV_TILE_TOKENS = 128
+SUPPORTED_INDEX_TOPK = 2048
+SUPPORTED_COMPRESS_RATIO = 1
+SUPPORTED_TOPK_DECODE_IMPL = "cute_dsl_gvr"
+
+# Tensor formats of the indexer_topk launch: contiguous FP8 queries ``[tokens, H, D]``,
+# contiguous fp32 per-head weights ``[tokens, H]``, fused uint8 K pages
+# ``[pages, tokens_per_block, 1, D + 4]`` (D FP8 bytes plus one fp32 scale).
+QUERY_DTYPE = "torch.float8_e4m3fn"
+WEIGHTS_DTYPE = "torch.float32"
+K_CACHE_DTYPE = "torch.uint8"
+K_CACHE_SCALE_BYTES = 4
+
+
+@dataclass(frozen=True)
+class Decision:
+    """Outcome of an indexer_topk eligibility check.
+
+    Attributes:
+        eligible: True when indexer_topk may replace the native logits + select region.
+        reason: Short reason for a native fallback; empty when eligible.
+    """
+
+    eligible: bool
+    reason: str = ""
+
+
+def decide(
+    *,
+    sm_version: int,
+    cute_dsl_available: bool,
+    native_uses_cute_dsl_logits: bool,
+    use_fp4_indexer: bool,
+    use_gvr_emission: bool,
+    needs_gvr_prior: bool,
+    is_ragged_verify: bool,
+    expanded_mtp_buffers: bool,
+    in_mtp_draft_loop: bool,
+    attention_dp: bool,
+    topk_decode_impl: str,
+    topk_self_sampling: bool,
+    index_topk: int,
+    compress_ratio: int,
+    num_heads: int,
+    head_dim: int,
+    query_dtype: str,
+    query_shape: Tuple[int, ...],
+    query_contiguous: bool,
+    weights_dtype: str,
+    weights_shape: Tuple[int, ...],
+    weights_contiguous: bool,
+    k_cache_dtype: str,
+    k_cache_shape: Tuple[int, ...],
+    block_table_cols: int,
+    max_seq_len: int,
+    num_contexts: int,
+    token_offset: int,
+    num_gen_tokens: int,
+    num_generations: int,
+    uniform_verify_width: bool,
+) -> Decision:
+    """Decide whether indexer_topk may run the GEN indexer logits + top-k region.
+
+    Args:
+        sm_version: Compute capability as an integer, e.g. 103.
+        cute_dsl_available: Whether the CUTLASS DSL runtime is importable.
+        native_uses_cute_dsl_logits: The indexer's
+            ``use_cute_dsl_paged_mqa_logits`` flag.
+        use_fp4_indexer: Whether the indexer K cache is MXFP4.
+        use_gvr_emission: Whether the temporal GVR emission tier is enabled.
+        needs_gvr_prior: Whether the decode top-k consumes a previous-step
+            prior (temporal GVR engine).
+        is_ragged_verify: Whether generation rows have per-request widths.
+        expanded_mtp_buffers: The metadata ``use_expanded_buffers_for_mtp``
+            flag; when set, the per-request schedule indexer_topk consumes is not built.
+        in_mtp_draft_loop: Whether this forward is a draft step of the MTP loop.
+        attention_dp: Whether attention data parallelism is enabled.
+        topk_decode_impl: Value of the layer ``TopK.decode_implementation``.
+        topk_self_sampling: The layer ``TopK.gvr_self_sampling`` flag.
+        index_topk: Selected IDs per row (``TopK.top_k``).
+        compress_ratio: Index compression ratio (``TopK.compress_ratio``).
+        num_heads: Indexer heads ``H``.
+        head_dim: Indexer head dimension ``D``.
+        query_dtype: ``str(dtype)`` of the projected indexer queries.
+        query_shape: Shape of the projected queries ``[tokens, H, D]``.
+        query_contiguous: Whether the projected queries are contiguous.
+        weights_dtype: ``str(dtype)`` of the per-head weights.
+        weights_shape: Shape of the per-head weights ``[tokens, H]``.
+        weights_contiguous: Whether the per-head weights are contiguous.
+        k_cache_dtype: ``str(dtype)`` of the fused indexer K page view.
+        k_cache_shape: Shape of the fused K page view
+            ``[pages, tokens_per_block, 1, D + 4]``.
+        block_table_cols: Columns (pages per request) of the indexer
+            block table ``indexer_k_cache_block_offsets``.
+        max_seq_len: Capture-stable indexer max sequence length.
+        num_contexts: Context requests in this forward.
+        token_offset: First generation row in the query/weight tensors.
+        num_gen_tokens: Generation query rows in this call.
+        num_generations: Generation requests in this call.
+        uniform_verify_width: Whether every generation request has the same
+            number of query rows (from the host ``seq_lens``).
+
+    Returns:
+        The decision with a fallback reason when ineligible.
+    """
+    if sm_version not in SUPPORTED_SM_VERSIONS:
+        return Decision(False, f"sm{sm_version} unsupported")
+    if not cute_dsl_available:
+        return Decision(False, "CUTLASS DSL unavailable")
+    if not native_uses_cute_dsl_logits:
+        return Decision(False, "native path is not CuTe-DSL paged-MQA logits")
+    if use_fp4_indexer:
+        return Decision(False, "FP4 indexer keys")
+    if use_gvr_emission:
+        return Decision(False, "GVR emission tier enabled")
+    if needs_gvr_prior:
+        return Decision(False, "temporal GVR prior in use")
+    if is_ragged_verify:
+        return Decision(False, "ragged verification")
+    if expanded_mtp_buffers:
+        return Decision(False, "expanded MTP buffers (no per-request schedule)")
+    if in_mtp_draft_loop:
+        return Decision(False, "MTP draft step (target layers only)")
+    if attention_dp:
+        return Decision(False, "attention data parallelism")
+    if topk_decode_impl != SUPPORTED_TOPK_DECODE_IMPL or not topk_self_sampling:
+        return Decision(
+            False, f"top-k engine {topk_decode_impl} (self_sampling={topk_self_sampling})"
+        )
+    if index_topk != SUPPORTED_INDEX_TOPK:
+        return Decision(False, f"index_topk {index_topk} != {SUPPORTED_INDEX_TOPK}")
+    if compress_ratio != SUPPORTED_COMPRESS_RATIO:
+        return Decision(False, f"index compress_ratio {compress_ratio}")
+    if num_heads != SUPPORTED_NUM_HEADS:
+        return Decision(False, f"num_heads {num_heads} != {SUPPORTED_NUM_HEADS}")
+    if head_dim != SUPPORTED_HEAD_DIM:
+        return Decision(False, f"head_dim {head_dim} != {SUPPORTED_HEAD_DIM}")
+    k_cache_shape = tuple(k_cache_shape)
+    if len(k_cache_shape) != 4 or k_cache_shape[2] != 1:
+        return Decision(False, f"K cache view shape {k_cache_shape}")
+    if k_cache_shape[1] != SUPPORTED_TOKENS_PER_BLOCK:
+        return Decision(
+            False, f"tokens_per_block {k_cache_shape[1]} != {SUPPORTED_TOKENS_PER_BLOCK}"
+        )
+    if k_cache_dtype != K_CACHE_DTYPE or k_cache_shape[3] != head_dim + K_CACHE_SCALE_BYTES:
+        return Decision(
+            False, f"K cache layout {k_cache_dtype} x {k_cache_shape[3]} bytes per token"
+        )
+    if query_dtype != QUERY_DTYPE:
+        return Decision(False, f"query dtype {query_dtype}")
+    if weights_dtype != WEIGHTS_DTYPE:
+        return Decision(False, f"weights dtype {weights_dtype}")
+    query_shape, weights_shape = tuple(query_shape), tuple(weights_shape)
+    if (
+        len(query_shape) != 3
+        or query_shape[1:] != (num_heads, head_dim)
+        or len(weights_shape) != 2
+        or weights_shape[1] != num_heads
+        or not query_contiguous
+        or not weights_contiguous
+    ):
+        return Decision(False, f"query {query_shape} / weights {weights_shape} layout")
+    if max_seq_len <= 0 or max_seq_len > MAX_INDEXER_SEQ_LEN:
+        return Decision(False, f"indexer max_seq_len {max_seq_len} > {MAX_INDEXER_SEQ_LEN}")
+    pages_per_tile = KV_TILE_TOKENS // k_cache_shape[1]
+    if block_table_cols <= 0 or block_table_cols % pages_per_tile != 0:
+        return Decision(
+            False,
+            f"block table width {block_table_cols} is not a multiple of {pages_per_tile} pages",
+        )
+    if num_contexts != 0:
+        return Decision(False, "context requests in the same forward")
+    if num_generations <= 0 or num_gen_tokens <= 0:
+        return Decision(False, "no generation rows")
+    if num_generations > MAX_GENERATIONS:
+        return Decision(False, f"{num_generations} generation requests > {MAX_GENERATIONS}")
+    if token_offset < 0 or token_offset + num_gen_tokens > min(query_shape[0], weights_shape[0]):
+        return Decision(False, "generation rows exceed the query/weight tensors")
+    if num_gen_tokens % num_generations != 0 or not uniform_verify_width:
+        return Decision(False, "non-uniform verify width")
+    next_n = num_gen_tokens // num_generations
+    if next_n != SUPPORTED_NEXT_N:
+        return Decision(False, f"next_n {next_n} != {SUPPORTED_NEXT_N}")
+    return Decision(True)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/region.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/region.py
new file mode 100644
index 000000000000..88279f9f08a7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/region.py
@@ -0,0 +1,99 @@
+# 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.
+"""The generation region of ``Indexer.sparse_attn_indexer`` on a decode step.
+
+``GlmIndexer._fill_decode_topk`` asks these entries, in order, to fill the generation rows of
+the top-k buffer. Each decides from host state before any device work, so a ``False`` return
+leaves ``out`` untouched for the next entry and, last, for the built-in decode branch.
+"""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from ... import dispatch
+
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+    )
+
+
+def cross_rank_topk_decode(
+    indexer: "Indexer",
+    metadata: "DSAtrtllmAttentionMetadata",
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+    bridge=None,
+) -> bool:
+    """Run the context-partitioned (cross-rank) top-k of ``cross_rank_topk`` when selected.
+
+    Returns True when it filled ``out`` (the generation rows of the top-k buffer). The decision
+    reads host state that is identical on every rank of the tensor-parallel group. Among the
+    selected variants the first that serves the call wins:
+
+    1. ``cross_rank_topk_global_emit``: the cluster-merge top-k that also writes the call's
+       bridge outputs (the shared top-k copy and the global remap) and records them on
+       ``bridge``;
+    2. ``indexer_layer_front``: the same region with its logits kernel launched as the rope's
+       programmatic dependent; when it declines, the variants below decline too;
+    3. ``cross_rank_topk_cluster_merge``: the cluster-merge top-k (rows byte-identical to the
+       plain cross-rank top-k);
+    4. ``cross_rank_topk`` itself.
+    """
+    if not dispatch.is_enabled("cross_rank_topk"):
+        return False
+    if dispatch.is_enabled("cross_rank_topk_global_emit") and dispatch.get_impl(
+        "cross_rank_topk_global_emit", "cp_indexer_decode"
+    )(indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out, bridge):
+        return True
+    # The same region with the logits kernel waiting at entry and launched as the rope's
+    # programmatic dependent.
+    if dispatch.is_enabled("indexer_layer_front") and dispatch.get_impl(
+        "indexer_layer_front", "cp_indexer_decode"
+    )(indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out):
+        return True
+    if dispatch.is_enabled("cross_rank_topk_cluster_merge") and dispatch.get_impl(
+        "cross_rank_topk_cluster_merge", "cp_indexer_decode"
+    )(indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out):
+        return True
+    return dispatch.get_impl("indexer_topk", "cp_indexer_decode")(
+        indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out
+    )
+
+
+def indexer_topk_decode(
+    indexer: "Indexer",
+    metadata: "DSAtrtllmAttentionMetadata",
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    token_offset: int,
+    num_gen_tokens: int,
+    out: torch.Tensor,
+) -> bool:
+    """Run the ``indexer_topk`` decode top-k when it is selected.
+
+    Returns True when it filled ``out`` (the generation rows of the top-k buffer). Its
+    eligibility is checked from host state before any device work.
+    """
+    if not dispatch.is_enabled("indexer_topk"):
+        return False
+    return dispatch.get_impl("indexer_topk", "indexer_decode")(
+        indexer, metadata, q_fp8, weights, token_offset, num_gen_tokens, out
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/topk_select.py
new file mode 100644
index 000000000000..8d0707114438
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/topk_select.py
@@ -0,0 +1,475 @@
+# 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.
+"""indexer_topk exact decode top-k selection over the indexer logits arena (Triton).
+
+Experimental and opt-in (``decode_topk(selector="radix")``); the adapter
+path uses the native ``TopK`` module. For the 6-row decode shape it is
+exact, and the native self-sampling engine stays the default, so it is kept
+as a tested reference only.
+
+Same contract as the native decode ``TopK`` (self-sampling GVR) on the
+uniform-width GEN path: row ``r`` of ``B * next_n`` rows has
+``n_r = kv_len[r // next_n] - next_n + r % next_n + 1`` valid logits; the
+output row holds the indices of the ``k`` largest logits in any order
+(tie-interchangeable: the selected value multiset equals ``torch.topk``),
+rows with ``0 < n_r <= k`` hold ``0 .. n_r - 1`` then ``-1``, and rows with
+``n_r <= 0`` hold ``-1``. No host synchronization; graph-capturable.
+
+Algorithm (MSD radix select on order-preserving 32-bit keys):
+
+1. ``_hist_kernel``, grid ``(chunks, rows)``: per-row histogram of the top
+   12 key bits, accumulated per CTA and flushed with atomics; the last CTA of
+   a row (acq-rel ticket) finds the threshold bin and the count above it,
+   publishes both, and clears the histogram for the next call.
+2. ``_select_kernel``, same grid: every CTA writes the indices of its keys
+   above the threshold bin (slots from an atomic counter) and appends the
+   threshold-bin keys to a candidate buffer; the last CTA of a row resolves
+   the remaining 20 key bits over the candidates (or, if they overflowed the
+   buffer, over the whole row) and writes the rest, then resets the counters.
+"""
+
+from typing import Dict, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+TOP_BITS = 12
+NUM_TOP_BINS = 1 << TOP_BITS
+DEFAULT_CANDIDATE_CAPACITY = 16384
+_NUM_WARPS = 8
+
+
+@triton.jit
+def _keys(x):
+    """Order-preserving uint32 keys of fp32 values (larger value, larger key)."""
+    i = x.to(tl.int32, bitcast=True)
+    # Arithmetic shift: all ones for negative values, else only the sign bit.
+    flip = (i >> 31) | -2147483648
+    return (i ^ flip).to(tl.uint32, bitcast=True)
+
+
+@triton.jit
+def _row_len(kv_lens_ptr, row, next_n: tl.constexpr):
+    kv = tl.load(kv_lens_ptr + row // next_n)
+    return tl.maximum(kv - next_n + row % next_n + 1, 0)
+
+
+@triton.jit
+def _fence_all(NUM_THREADS: tl.constexpr):
+    # Every thread publishes its prior global writes at GPU scope (one
+    # element per thread, so each thread runs the fence once).
+    tl.inline_asm_elementwise(
+        "fence.acq_rel.gpu;\n\tmov.u32 $0, $1;",
+        "=r,r",
+        [tl.arange(0, NUM_THREADS)],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _hist_kernel(
+    logits_ptr,
+    stride_row,
+    kv_lens_ptr,
+    hist_ptr,
+    thr_ptr,
+    ctr_ptr,
+    next_n: tl.constexpr,
+    K: tl.constexpr,
+    CHUNKS: tl.constexpr,
+    BLOCK: tl.constexpr,
+    NB: tl.constexpr,
+    NUM_THREADS: tl.constexpr,
+):
+    chunk = tl.program_id(0)
+    row = tl.program_id(1)
+    n = _row_len(kv_lens_ptr, row, next_n)
+    if n > K:
+        span = tl.cdiv(tl.cdiv(n, CHUNKS), BLOCK) * BLOCK
+        start = chunk * span
+        end = tl.minimum(start + span, n)
+        offs = tl.arange(0, BLOCK)
+        local = tl.zeros([NB], dtype=tl.int32)
+        invalid = tl.full([], 0, tl.int32)
+        for base in range(start, end, BLOCK):
+            idx = base + offs
+            valid = idx < end
+            x = tl.load(logits_ptr + row * stride_row + idx, mask=valid, other=0.0)
+            b = tl.where(valid, (_keys(x) >> (32 - 12)).to(tl.int32), 0)
+            local += tl.histogram(b, NB)
+            invalid += tl.sum((~valid).to(tl.int32))
+        bins = tl.arange(0, NB)
+        local = local - tl.where(bins == 0, invalid, 0)
+        tl.atomic_add(hist_ptr + row * NB + bins, local, mask=local != 0, sem="relaxed")
+        _fence_all(NUM_THREADS)
+        tl.debug_barrier()
+        ticket = tl.atomic_add(ctr_ptr + row * 4, 1, sem="acq_rel", scope="gpu")
+        if ticket == CHUNKS - 1:
+            tl.debug_barrier()
+            h = tl.load(hist_ptr + row * NB + bins, cache_modifier=".cg")
+            total = tl.sum(h)
+            prefix = tl.cumsum(h, 0)
+            at_or_above = total - prefix + h
+            b1 = tl.max(tl.where(at_or_above >= K, bins, -1))
+            above = tl.sum(tl.where(bins > b1, h, 0))
+            tl.store(thr_ptr + row * 2, b1)
+            tl.store(thr_ptr + row * 2 + 1, above)
+            tl.store(hist_ptr + row * NB + bins, tl.zeros([NB], dtype=tl.int32))
+            tl.store(ctr_ptr + row * 4, 0)
+
+
+@triton.jit
+def _load_source(
+    logits_ptr,
+    stride_row,
+    cand_key_ptr,
+    cand_idx_ptr,
+    row,
+    i,
+    inside,
+    from_buffer,
+    CAP: tl.constexpr,
+):
+    """Keys and row positions of the candidate buffer or of the full row."""
+    if from_buffer:
+        key = tl.load(cand_key_ptr + row * CAP + i, mask=inside, other=0, cache_modifier=".cg")
+        idx = tl.load(cand_idx_ptr + row * CAP + i, mask=inside, other=0, cache_modifier=".cg")
+    else:
+        x = tl.load(logits_ptr + row * stride_row + i, mask=inside, other=0.0)
+        key = _keys(x)
+        idx = i
+    return key, idx
+
+
+@triton.jit
+def _digit_pass(
+    logits_ptr,
+    stride_row,
+    cand_key_ptr,
+    cand_idx_ptr,
+    row,
+    src_len,
+    from_buffer,
+    prefix_key,
+    prefix_mask,
+    want,
+    SHIFT: tl.constexpr,
+    NBITS: tl.constexpr,
+    CAP: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """One MSD digit: histogram of keys matching the prefix, then the digit
+    ``t`` holding the ``want``-th largest key. Returns the extended prefix,
+    mask and the number still wanted among keys equal to the new prefix."""
+    NUM: tl.constexpr = 1 << NBITS
+    offs = tl.arange(0, BLOCK)
+    bins = tl.arange(0, NUM)
+    h = tl.zeros([NUM], dtype=tl.int32)
+    for base in range(0, src_len, BLOCK):
+        i = base + offs
+        inside = i < src_len
+        key, _ = _load_source(
+            logits_ptr, stride_row, cand_key_ptr, cand_idx_ptr, row, i, inside, from_buffer, CAP
+        )
+        active = inside & ((key & prefix_mask) == prefix_key)
+        d = tl.where(active, ((key >> SHIFT) & (NUM - 1)).to(tl.int32), 0)
+        h += tl.histogram(d, NUM)
+        h -= tl.where(bins == 0, tl.sum((~active).to(tl.int32)), 0)
+    total = tl.sum(h)
+    prefix = tl.cumsum(h, 0)
+    at_or_above = total - prefix + h
+    t = tl.max(tl.where(at_or_above >= want, bins, -1))
+    above = tl.sum(tl.where(bins > t, h, 0))
+    new_key = prefix_key | (t.to(tl.uint32) << SHIFT)
+    new_mask = prefix_mask | ((NUM - 1) << SHIFT)
+    return new_key, new_mask, want - above
+
+
+@triton.jit
+def _final_select(
+    logits_ptr,
+    stride_row,
+    cand_key_ptr,
+    cand_idx_ptr,
+    out_ptr,
+    row,
+    n,
+    b1,
+    count_above,
+    num_cand,
+    K: tl.constexpr,
+    CAP: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Pick ``K - count_above`` of the threshold-bin keys (last CTA of a row)."""
+    from_buffer = num_cand <= CAP
+    src_len = tl.where(from_buffer, num_cand, n)
+    offs = tl.arange(0, BLOCK)
+    want = K - count_above
+    prefix_key = b1.to(tl.uint32) << 20
+    prefix_mask = tl.full([], 0xFFF00000, tl.uint32)
+    # Remaining 20 bits in digits of 8, 8 and 4 bits.
+    prefix_key, prefix_mask, want = _digit_pass(
+        logits_ptr,
+        stride_row,
+        cand_key_ptr,
+        cand_idx_ptr,
+        row,
+        src_len,
+        from_buffer,
+        prefix_key,
+        prefix_mask,
+        want,
+        12,
+        8,
+        CAP,
+        BLOCK,
+    )
+    prefix_key, prefix_mask, want = _digit_pass(
+        logits_ptr,
+        stride_row,
+        cand_key_ptr,
+        cand_idx_ptr,
+        row,
+        src_len,
+        from_buffer,
+        prefix_key,
+        prefix_mask,
+        want,
+        4,
+        8,
+        CAP,
+        BLOCK,
+    )
+    prefix_key, prefix_mask, want = _digit_pass(
+        logits_ptr,
+        stride_row,
+        cand_key_ptr,
+        cand_idx_ptr,
+        row,
+        src_len,
+        from_buffer,
+        prefix_key,
+        prefix_mask,
+        want,
+        0,
+        4,
+        CAP,
+        BLOCK,
+    )
+    # Emit keys above the exact threshold key, then `want` of the ties.
+    threshold = prefix_key
+    slot = count_above
+    ties_left = want
+    for base in range(0, src_len, BLOCK):
+        i = base + offs
+        inside = i < src_len
+        key, idx = _load_source(
+            logits_ptr, stride_row, cand_key_ptr, cand_idx_ptr, row, i, inside, from_buffer, CAP
+        )
+        in_bin = inside & ((key >> 20) == (threshold >> 20))
+        gt = in_bin & (key > threshold)
+        eq = in_bin & (key == threshold)
+        eq_rank = tl.cumsum(eq.to(tl.int32), 0)
+        take_eq = eq & (eq_rank <= ties_left)
+        sel = gt | take_eq
+        pos = slot + tl.cumsum(sel.to(tl.int32), 0) - 1
+        tl.store(out_ptr + row * K + pos, idx, mask=sel)
+        slot += tl.sum(sel.to(tl.int32))
+        ties_left -= tl.sum(take_eq.to(tl.int32))
+
+
+@triton.jit
+def _select_kernel(
+    logits_ptr,
+    stride_row,
+    kv_lens_ptr,
+    out_ptr,
+    thr_ptr,
+    ctr_ptr,
+    cand_key_ptr,
+    cand_idx_ptr,
+    next_n: tl.constexpr,
+    K: tl.constexpr,
+    CHUNKS: tl.constexpr,
+    BLOCK: tl.constexpr,
+    CAP: tl.constexpr,
+    NUM_THREADS: tl.constexpr,
+):
+    chunk = tl.program_id(0)
+    row = tl.program_id(1)
+    n = _row_len(kv_lens_ptr, row, next_n)
+    offs = tl.arange(0, BLOCK)
+    if n <= K:
+        if chunk == 0:
+            for base in range(0, K, BLOCK):
+                j = base + offs
+                tl.store(out_ptr + row * K + j, tl.where(j < n, j, -1), mask=j < K)
+    else:
+        b1 = tl.load(thr_ptr + row * 2, cache_modifier=".cg")
+        count_above = tl.load(thr_ptr + row * 2 + 1, cache_modifier=".cg")
+        span = tl.cdiv(tl.cdiv(n, CHUNKS), BLOCK) * BLOCK
+        start = chunk * span
+        end = tl.minimum(start + span, n)
+        for base in range(start, end, BLOCK):
+            idx = base + offs
+            valid = idx < end
+            x = tl.load(logits_ptr + row * stride_row + idx, mask=valid, other=0.0)
+            key = _keys(x)
+            b = (key >> 20).to(tl.int32)
+            gt = valid & (b > b1)
+            eq = valid & (b == b1)
+            n_gt = tl.sum(gt.to(tl.int32))
+            n_eq = tl.sum(eq.to(tl.int32))
+            if n_gt > 0:
+                slot = tl.atomic_add(ctr_ptr + row * 4 + 1, n_gt, sem="relaxed")
+                pos = slot + tl.cumsum(gt.to(tl.int32), 0) - 1
+                tl.store(out_ptr + row * K + pos, idx, mask=gt)
+            if n_eq > 0:
+                cslot = tl.atomic_add(ctr_ptr + row * 4 + 2, n_eq, sem="relaxed")
+                cpos = cslot + tl.cumsum(eq.to(tl.int32), 0) - 1
+                keep = eq & (cpos < CAP)
+                tl.store(cand_key_ptr + row * CAP + cpos, key, mask=keep)
+                tl.store(cand_idx_ptr + row * CAP + cpos, idx, mask=keep)
+        _fence_all(NUM_THREADS)
+        tl.debug_barrier()
+        ticket = tl.atomic_add(ctr_ptr + row * 4 + 3, 1, sem="acq_rel", scope="gpu")
+        if ticket == CHUNKS - 1:
+            tl.debug_barrier()
+            num_cand = tl.load(ctr_ptr + row * 4 + 2, cache_modifier=".cg")
+            _final_select(
+                logits_ptr,
+                stride_row,
+                cand_key_ptr,
+                cand_idx_ptr,
+                out_ptr,
+                row,
+                n,
+                b1,
+                count_above,
+                num_cand,
+                K,
+                CAP,
+                BLOCK,
+            )
+            tl.store(ctr_ptr + row * 4 + 1, 0)
+            tl.store(ctr_ptr + row * 4 + 2, 0)
+            tl.store(ctr_ptr + row * 4 + 3, 0)
+
+
+class RadixTopKWorkspace:
+    """Graph-stable scratch of the radix select for up to ``max_rows`` rows.
+
+    Counters and histograms are zero-initialized once and left zeroed by
+    every call, so the scratch can be reused across calls and CUDA-graph
+    replays on one stream. Concurrent calls on different streams need
+    separate workspaces.
+    """
+
+    def __init__(self, max_rows: int, device, capacity: int = DEFAULT_CANDIDATE_CAPACITY):
+        self.max_rows = max_rows
+        self.capacity = capacity
+        self.hist = torch.zeros(max_rows, NUM_TOP_BINS, dtype=torch.int32, device=device)
+        self.thr = torch.zeros(max_rows, 2, dtype=torch.int32, device=device)
+        self.ctr = torch.zeros(max_rows, 4, dtype=torch.int32, device=device)
+        self.cand_key = torch.empty(max_rows, capacity, dtype=torch.uint32, device=device)
+        self.cand_idx = torch.empty(max_rows, capacity, dtype=torch.int32, device=device)
+
+
+_WORKSPACES: Dict[Tuple[int, int], RadixTopKWorkspace] = {}
+
+
+def get_workspace(rows: int, device: torch.device) -> RadixTopKWorkspace:
+    """Per-device workspace sized for at least ``rows`` rows (grown on demand)."""
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    ws = _WORKSPACES.get(index)
+    if ws is None or ws.max_rows < rows:
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError("radix top-k workspace must be created before graph capture")
+        ws = RadixTopKWorkspace(max(rows, 8), device)
+        _WORKSPACES[index] = ws
+    return ws
+
+
+def default_chunks(rows: int, num_sms: int) -> int:
+    """CTAs per row: about two waves of the grid over the SMs."""
+    return max(1, min(64, (2 * num_sms) // max(rows, 1)))
+
+
+def radix_topk_decode(
+    logits: torch.Tensor,
+    kv_lens: torch.Tensor,
+    out: torch.Tensor,
+    next_n: int,
+    chunks: int,
+    block: int = 4096,
+    workspace: RadixTopKWorkspace = None,
+) -> torch.Tensor:
+    """Exact decode top-k over a row-major fp32 logits view.
+
+    Args:
+        logits: ``[rows, >= max n_r]`` float32 view with unit inner stride.
+        kv_lens: ``[rows // next_n]`` int32 logical KV lengths.
+        out: ``[rows, k]`` int32 contiguous output.
+        next_n: Verify rows per request.
+        chunks: CTAs per row (fixed per captured graph).
+        block: Elements per CTA iteration.
+        workspace: Scratch; defaults to the per-device workspace.
+
+    Returns:
+        ``out``.
+    """
+    rows, k = out.shape
+    assert logits.dtype == torch.float32 and logits.stride(1) == 1
+    assert out.dtype == torch.int32 and out.is_contiguous() and kv_lens.dtype == torch.int32
+    assert rows % next_n == 0 and kv_lens.numel() == rows // next_n
+    ws = workspace or get_workspace(rows, logits.device)
+    grid = (chunks, rows)
+    _hist_kernel[grid](
+        logits,
+        logits.stride(0),
+        kv_lens,
+        ws.hist,
+        ws.thr,
+        ws.ctr,
+        next_n=next_n,
+        K=k,
+        CHUNKS=chunks,
+        BLOCK=block,
+        NB=NUM_TOP_BINS,
+        NUM_THREADS=32 * _NUM_WARPS,
+        num_warps=_NUM_WARPS,
+    )
+    _select_kernel[grid](
+        logits,
+        logits.stride(0),
+        kv_lens,
+        out,
+        ws.thr,
+        ws.ctr,
+        ws.cand_key,
+        ws.cand_idx,
+        next_n=next_n,
+        K=k,
+        CHUNKS=chunks,
+        BLOCK=block,
+        CAP=ws.capacity,
+        NUM_THREADS=32 * _NUM_WARPS,
+        num_warps=_NUM_WARPS,
+    )
+    return out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/__init__.py
new file mode 100644
index 000000000000..3d2e34f685ae
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/__init__.py
@@ -0,0 +1,48 @@
+# 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.
+"""oproj_weight_staging: weight staging inside already-waiting kernels, plus cache-policy hints.
+
+Site ``oproj_weight_staging`` (attention window, hosted by attention_output_tail's o_proj): the
+served attention_output_tail/attention_output_into_hop tail streams part of the local ``W_o`` into
+shared memory before its PDL wait and the rest after it; oproj_weight_staging stages the after-wait
+rows before the wait, from a kernel that is already resident and waiting (the o_proj itself, or
+attention_output_tail's W_UV kernel), with L2 cache-policy hints (staging prefetch normal or
+``evict_last``, the consumer's copies optionally ``evict_first``). The
+arithmetic is unchanged, so the partial is bit-identical to attention_output_tail's. No launch
+and no stream is added.
+
+Registered entry points (``dispatch.get_impl("oproj_weight_staging", name)`` after this
+package is imported):
+
+* ``wrap_attn_tail(mla, tail) -> AttnTail | None``: wraps the tail that attention_output_tail's
+  ``make_attn_tail`` created for an MLA module (see ``tail.py``).
+* ``stage_tail(...)``: the functional tail (see ``kernels.py``).
+"""
+
+from ... import dispatch
+from .kernels import DEFAULT_STAGE, StageConfig, stage_tail
+from .tail import FEATURE_ID, StagedAttnTail, wrap_attn_tail
+
+dispatch.register(FEATURE_ID, "wrap_attn_tail", wrap_attn_tail)
+dispatch.register(FEATURE_ID, "stage_tail", stage_tail)
+
+__all__ = [
+    "DEFAULT_STAGE",
+    "FEATURE_ID",
+    "StageConfig",
+    "StagedAttnTail",
+    "stage_tail",
+    "wrap_attn_tail",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py
new file mode 100644
index 000000000000..1ee494a105f2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py
@@ -0,0 +1,841 @@
+# 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.
+"""oproj_weight_staging ``oproj_weight_staging``: the served
+attention_output_tail/attention_output_into_hop streaming o_proj with after-wait row staging.
+
+A fork of ``attention_output_tail/into_hop.py`` (the o_proj of ``INTO_HOP_SPARSE_DECODE_SCHEDULE``:
+one CTA per row block streams its ``w_o`` rows into a padded shared-memory slab
+around its PDL wait, 8 K-split warps run ``mma.sync``, fixed-order
+cross-warp reduction, one BF16 rounding; optionally the attention_output_into_hop publish into
+the post_attention_hop transport). The arithmetic, the row split (even boundaries), the
+epilogue and the publish are unchanged, so ``out`` and the published words
+equal ``cute_oproj_into_hop.stream_oproj_into_hop``'s bit for bit, and with
+``publish=False`` ``out`` equals ``cute_oproj.stream_oproj``'s.
+
+What changes is when the rows that the served schedule copies only after
+``griddepcontrol.wait`` (rows ``[early_rows, nrows)`` of each CTA) leave
+DRAM:
+
+* ``stage="l2"``: rows ``[early_rows, early_rows + stage_rows)`` are
+  prefetched into L2 (``cp.async.bulk.prefetch.L2``, policy
+  ``prefetch_policy``) before the wait. They are still copied into shared
+  memory after the wait, exactly as served, so correctness never depends on
+  the prefetch.
+* ``stage="smem"``: those rows are copied into the slab before the wait
+  (the same ``cp.async.bulk`` + tile mbarrier path as the early rows).
+
+``when`` places the staging: ``"entry"`` right after the early copies (by
+the copy warp), ``"landed"`` by a second warp once every tile that holds
+only early rows has landed, so the staging traffic never queues behind (or
+in front of) the early stream; ``"chain"`` (``stage="smem"`` only) by the
+second warp one tile at a time, each tile's staged rows once the previous
+tile has landed. ``copy_policy="evict_first"`` adds an L2
+``evict_first`` cache hint to every ``w_o`` row copy (the consumer read of a
+single-use stream). Rows below ``early_rows`` are never staged again.
+
+Memory order: before the wait only immutable ``w_o``
+rows are read (copies into the CTA's own shared memory, L2 prefetches); every
+read of ``v`` and of the into-hop protocol state and every global write follows
+the wait; the only new wait is a warp waiting for copies its own CTA issued.
+
+``trace=True`` (tests and diagnostics only) records ``%globaltimer`` per CTA
+at entry, after the wait, after warp 0's MMA loop and at exit.
+"""
+
+import functools
+import math
+from typing import Optional
+
+import torch
+
+_THREADS = 256
+_HIDDEN = 6144
+_NUM_RANKS = 8
+_LAMPORT_BUFFER_WORDS = 8 * _NUM_RANKS * _HIDDEN * 2 // 4
+_WARPS = _THREADS // 32
+_ROW_PAD = 8  # BF16 elements of padding per shared-memory row (16 bytes)
+_TILE_M = 16
+_STAGE_WARP = 1
+TRACE_POINTS = 4
+
+STAGES = ("none", "l2", "smem")
+WHENS = ("entry", "landed", "chain")
+PREFETCH_POLICIES = ("normal", "evict_last")
+COPY_POLICIES = ("none", "evict_first")
+
+
+@functools.lru_cache(maxsize=None)
+def _sm_count(device_index: int) -> int:
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def plan(n_dim: int, k_dim: int, grid: int) -> dict:
+    """Compile-time geometry of the kernel for one weight shape (as
+    ``cute_oproj_into_hop.plan``)."""
+    rows_max = 2 * math.ceil(n_dim / 2 / grid)
+    pitch = k_dim + _ROW_PAD
+    return {
+        "rows_max": rows_max,
+        "rows_min": 2 * ((n_dim // 2) // grid),
+        "num_tiles": math.ceil(rows_max / _TILE_M),
+        "pitch": pitch,
+        "smem_bytes": rows_max * pitch * 2,
+        "k_per_warp": k_dim // _WARPS,
+    }
+
+
+def check_shapes(n_dim: int, k_dim: int, grid: int) -> Optional[str]:
+    """Return why the kernel cannot take this ``[n_dim, k_dim]`` weight."""
+    if k_dim % (_WARPS * 16):
+        return "K is not a multiple of 128"
+    if n_dim % 2:
+        return "N is odd"
+    p = plan(n_dim, k_dim, grid)
+    red_bytes = _WARPS * p["num_tiles"] * 2 * 32 * 4 * 4
+    if p["smem_bytes"] < red_bytes:
+        return "slab smaller than the reduction buffer"
+    if p["smem_bytes"] > 226 * 1024:
+        return "weight slab does not fit in shared memory"
+    if (k_dim * 2) % 16:
+        return "row bytes not a 16-byte multiple"
+    return None
+
+
+def check_stage(stage: str, when: str, prefetch_policy: str, copy_policy: str) -> Optional[str]:
+    """Return why a staging configuration is invalid, or ``None``."""
+    if stage not in STAGES:
+        return f"stage {stage!r} not in {STAGES}"
+    if when not in WHENS:
+        return f"when {when!r} not in {WHENS}"
+    if prefetch_policy not in PREFETCH_POLICIES:
+        return f"prefetch_policy {prefetch_policy!r} not in {PREFETCH_POLICIES}"
+    if copy_policy not in COPY_POLICIES:
+        return f"copy_policy {copy_policy!r} not in {COPY_POLICIES}"
+    if when == "chain" and stage == "l2":
+        return "chain pacing waits for staged rows to land: smem staging only"
+    return None
+
+
+def _build(
+    n_dim: int,
+    k_dim: int,
+    grid: int,
+    n_tiles_r: int,
+    use_pdl: bool,
+    early_rows: int,
+    trigger_dependents: bool,
+    publish: bool,
+    publish_rank: int,
+    vector_publish: bool,
+    stage: str,
+    when: str,
+    stage_rows: int,
+    prefetch_policy: str,
+    copy_policy: str,
+    trace: bool,
+):
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import BFloat16, Float32, Int32, Int64
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cute import runtime as cute_runtime
+    from cutlass.cutlass_dsl import T, dsl_user_op
+    from cutlass.utils.smem_allocator import SmemAllocator
+
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+
+    geo = plan(n_dim, k_dim, grid)
+    rows_max = geo["rows_max"]
+    num_tiles = geo["num_tiles"]
+    pitch = geo["pitch"]
+    half_pitch = pitch // 2
+    k_per_warp = geo["k_per_warp"]
+    steps = k_per_warp // 16
+    row_bytes = k_dim * 2
+    nt = n_tiles_r
+    red_stride = 128  # 32 lanes x 4 accumulators
+    vec_publish = publish and vector_publish and n_dim % (8 * grid) == 0
+    cols = n_dim // grid
+    chunks = cols // 8
+    stage_pitch = cols // 2
+    stage0 = _WARPS * num_tiles * nt * red_stride
+    if vec_publish:
+        assert stage0 + 8 * stage_pitch <= rows_max * pitch // 2
+    # Rows [early_rows, stage_end) of every CTA are staged before the wait.
+    staging = use_pdl and stage != "none" and stage_rows > 0
+    stage_end = early_rows + stage_rows if staging else early_rows
+    # Tiles that hold only early rows in every CTA: their mbarriers complete
+    # without any copy issued after the wait.
+    early_tiles = min(early_rows, geo["rows_min"]) // _TILE_M
+    copy_hint = copy_policy == "evict_first"
+    pf_hint = prefetch_policy == "evict_last"
+    landed = staging and when == "landed"
+    chained = staging and when == "chain"
+    # chain: tiles holding staged rows, each issued once the previous landed.
+    first_tile = min(early_rows, rows_max) // _TILE_M
+    last_tile = min(stage_end - 1, rows_max - 1) // _TILE_M
+
+    @dsl_user_op
+    def bulk_g2s(dst_smem, src_gmem, num_bytes, mbar, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                dst_smem.toint(loc=loc, ip=ip).ir_value(),
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+                mbar.toint(loc=loc, ip=ip).ir_value(),
+            ],
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            " [$0], [$1], $2, [$3];",
+            "r,l,r,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def bulk_g2s_evict_first(dst_smem, src_gmem, num_bytes, mbar, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                dst_smem.toint(loc=loc, ip=ip).ir_value(),
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+                mbar.toint(loc=loc, ip=ip).ir_value(),
+            ],
+            "{\n\t.reg .b64 pol;\n\t"
+            "createpolicy.fractional.L2::evict_first.b64 pol, 1.0;\n\t"
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            ".L2::cache_hint [$0], [$1], $2, [$3], pol;\n\t}",
+            "r,l,r,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def prefetch_l2(src_gmem, num_bytes, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+            ],
+            "cp.async.bulk.prefetch.L2.global [$0], $1;",
+            "l,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def prefetch_l2_evict_last(src_gmem, num_bytes, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                src_gmem.toint(loc=loc, ip=ip).ir_value(),
+                Int32(num_bytes).ir_value(loc=loc, ip=ip),
+            ],
+            "{\n\t.reg .b64 pol;\n\t"
+            "createpolicy.fractional.L2::evict_last.b64 pol, 1.0;\n\t"
+            "cp.async.bulk.prefetch.L2.global.L2::cache_hint [$0], $1, pol;"
+            "\n\t}",
+            "l,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def grid_launch_dependents(*, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [],
+            "griddepcontrol.launch_dependents;",
+            "",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def ld_volatile_u32(addr, *, loc=None, ip=None):
+        return Int32(
+            llvm.inline_asm(
+                T.i32(),
+                [Int64(addr).ir_value(loc=loc, ip=ip)],
+                "ld.volatile.global.u32 $0, [$1];",
+                "=r,l",
+                has_side_effects=True,
+                asm_dialect=0,
+                loc=loc,
+                ip=ip,
+            )
+        )
+
+    @dsl_user_op
+    def st_global_b32(addr, val, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [Int64(addr).ir_value(loc=loc, ip=ip), Int32(val).ir_value(loc=loc, ip=ip)],
+            "st.global.b32 [$0], $1;",
+            "l,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def st_global_v4_b32(addr, w0, w1, w2, w3, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [
+                Int64(addr).ir_value(loc=loc, ip=ip),
+                Int32(w0).ir_value(loc=loc, ip=ip),
+                Int32(w1).ir_value(loc=loc, ip=ip),
+                Int32(w2).ir_value(loc=loc, ip=ip),
+                Int32(w3).ir_value(loc=loc, ip=ip),
+            ],
+            "st.global.v4.b32 [$0], {$1, $2, $3, $4};",
+            "l,r,r,r,r",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def st_global_b64(addr, val, *, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [Int64(addr).ir_value(loc=loc, ip=ip), Int64(val).ir_value(loc=loc, ip=ip)],
+            "st.global.b64 [$0], $1;",
+            "l,l",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def globaltimer(*, loc=None, ip=None):
+        return Int64(
+            llvm.inline_asm(
+                T.i64(),
+                [],
+                "mov.u64 $0, %globaltimer;",
+                "=l",
+                has_side_effects=True,
+                asm_dialect=0,
+                loc=loc,
+                ip=ip,
+            )
+        )
+
+    @dsl_user_op
+    def grid_dependency_wait(*, loc=None, ip=None):
+        llvm.inline_asm(
+            None,
+            [],
+            "griddepcontrol.wait;",
+            "",
+            has_side_effects=True,
+            asm_dialect=0,
+            loc=loc,
+            ip=ip,
+        )
+
+    def copy_row(slab, w, mbar, n0, row):
+        """``cp.async.bulk`` of ``w_o`` row ``n0 + row`` into the slab
+        (traced inline; the branch is resolved at compile time)."""
+        if copy_hint:
+            bulk_g2s_evict_first(
+                slab.iterator + row * Int32(pitch),
+                w.iterator + (n0 + row) * Int32(k_dim),
+                row_bytes,
+                mbar.iterator + row // Int32(_TILE_M),
+            )
+        else:
+            bulk_g2s(
+                slab.iterator + row * Int32(pitch),
+                w.iterator + (n0 + row) * Int32(k_dim),
+                row_bytes,
+                mbar.iterator + row // Int32(_TILE_M),
+            )
+
+    def stage_row(slab, w, mbar, n0, row):
+        """Stage after-wait row ``row`` before the wait (L2 or smem)."""
+        if stage == "smem":
+            copy_row(slab, w, mbar, n0, row)
+        elif pf_hint:
+            prefetch_l2_evict_last(w.iterator + (n0 + row) * Int32(k_dim), row_bytes)
+        else:
+            prefetch_l2(w.iterator + (n0 + row) * Int32(k_dim), row_bytes)
+
+    def stamp_addr(tr, bidx, k):
+        return tr.iterator.toint() + Int64(bidx * Int32(TRACE_POINTS) + Int32(k)) * Int64(8)
+
+    class StreamOprojStage:
+        @cute.kernel
+        def kernel(
+            self,
+            v: cute.Tensor,
+            w: cute.Tensor,
+            out: cute.Tensor,
+            mc: cute.Tensor,
+            flags: cute.Tensor,
+            tr: cute.Tensor,
+        ):
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            warp = tidx // Int32(32)
+            lane = tidx % Int32(32)
+            grp = lane // Int32(4)
+            quad = lane % Int32(4)
+            if cutlass.const_expr(trace):
+                if tidx == Int32(0):
+                    st_global_b64(stamp_addr(tr, bidx, 0), globaltimer())
+
+            num_rows = Int32(v.shape[0])
+            out_rows = Int32(out.shape[0])
+            n0 = Int32(2) * ((bidx * Int32(n_dim // 2)) // Int32(grid))
+            n1 = Int32(2) * (((bidx + Int32(1)) * Int32(n_dim // 2)) // Int32(grid))
+            nrows = n1 - n0
+
+            smem = SmemAllocator()
+            slab = smem.allocate_tensor(
+                BFloat16, cute.make_layout((rows_max * pitch,)), byte_alignment=128
+            )
+            mbar = smem.allocate_tensor(
+                cutlass.Int64, cute.make_layout((num_tiles,)), byte_alignment=8
+            )
+            slab32 = cute.recast_tensor(slab, Int32)
+            slab_f32 = cute.recast_tensor(slab, Float32)
+            v32 = cute.recast_tensor(v, Int32)
+
+            if tidx == Int32(0):
+                for t in cutlass.range_constexpr(num_tiles):
+                    cute.arch.mbarrier_init(mbar.iterator + t, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+
+            # Lane 0 arms every tile mbarrier with all of its rows' bytes,
+            # whichever warp and phase issues the copies.
+            if tidx == Int32(0):
+                for t in cutlass.range_constexpr(num_tiles):
+                    tile_rows = nrows - Int32(t * _TILE_M)
+                    if tile_rows > Int32(_TILE_M):
+                        tile_rows = Int32(_TILE_M)
+                    if tile_rows < Int32(0):
+                        tile_rows = Int32(0)
+                    cute.arch.mbarrier_arrive_and_expect_tx(
+                        mbar.iterator + t, tile_rows * Int32(row_bytes)
+                    )
+            cute.arch.sync_warp()
+            # Served: rows [0, early_rows) before the wait on v's producer.
+            if warp == Int32(0):
+                for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                    row = lane + Int32(32 * q)
+                    if row < nrows:
+                        if row < Int32(early_rows):
+                            copy_row(slab, w, mbar, n0, row)
+                if cutlass.const_expr(staging and when == "entry"):
+                    # oproj_weight_staging, entry: stage the after-wait rows right after the
+                    # early copies.
+                    for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                        row = lane + Int32(32 * q)
+                        if row < nrows:
+                            if row >= Int32(early_rows):
+                                if row < Int32(stage_end):
+                                    stage_row(slab, w, mbar, n0, row)
+            if cutlass.const_expr(landed):
+                # oproj_weight_staging, landed: a second warp stages the after-wait rows once
+                # the early-only tiles are in shared memory (copies this CTA
+                # issued above; no other CTA or kernel is waited on).
+                if warp == Int32(_STAGE_WARP):
+                    for t in cutlass.range_constexpr(early_tiles):
+                        cute.arch.mbarrier_wait(mbar.iterator + t, 0)
+                    for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                        row = lane + Int32(32 * q)
+                        if row < nrows:
+                            if row >= Int32(early_rows):
+                                if row < Int32(stage_end):
+                                    stage_row(slab, w, mbar, n0, row)
+            if cutlass.const_expr(chained):
+                # oproj_weight_staging, chain: the second warp issues the staged rows of tile t
+                # once tile t - 1 has landed (copies its own CTA issued).
+                if warp == Int32(_STAGE_WARP):
+                    for i in cutlass.range_constexpr(last_tile + 1 - first_tile):
+                        if cutlass.const_expr(first_tile + i >= 1):
+                            cute.arch.mbarrier_wait(mbar.iterator + (first_tile + i - 1), 0)
+                        row = Int32((first_tile + i) * _TILE_M) + lane
+                        if lane < Int32(_TILE_M):
+                            if row < nrows:
+                                if row >= Int32(early_rows):
+                                    if row < Int32(stage_end):
+                                        stage_row(slab, w, mbar, n0, row)
+
+            if cutlass.const_expr(use_pdl):
+                grid_dependency_wait()
+            if cutlass.const_expr(trace):
+                if tidx == Int32(0):
+                    st_global_b64(stamp_addr(tr, bidx, 1), globaltimer())
+
+            # Rows not yet copied into shared memory go out after the wait
+            # (with stage="l2" these copies read the prefetched lines).
+            if warp == Int32(0):
+                for q in cutlass.range_constexpr((rows_max + 31) // 32):
+                    row = lane + Int32(32 * q)
+                    if row < nrows:
+                        if cutlass.const_expr(staging and stage == "smem"):
+                            if row >= Int32(stage_end):
+                                copy_row(slab, w, mbar, n0, row)
+                        else:
+                            if row >= Int32(early_rows):
+                                copy_row(slab, w, mbar, n0, row)
+            if cutlass.const_expr(trigger_dependents):
+                # Dependents still wait for this whole grid (and its memory)
+                # before they read anything it produces.
+                grid_launch_dependents()
+
+            # v fragments of this warp's K range (B operand, [n][k] storage).
+            vb = cute.make_rmem_tensor((4, nt, steps), BFloat16)
+            vb32 = cute.recast_tensor(vb, Int32)
+            k_word0 = warp * Int32(k_per_warp // 2) + quad
+            for s in cutlass.range_constexpr(steps):
+                for j in cutlass.range_constexpr(nt):
+                    r = Int32(8 * j) + grp
+                    lo = Int32(0)
+                    hi = Int32(0)
+                    if r < num_rows:
+                        lo = v32[r, k_word0 + Int32(8 * s)]
+                        hi = v32[r, k_word0 + Int32(8 * s + 4)]
+                    vb32[0, j, s] = lo
+                    vb32[1, j, s] = hi
+
+            acc = cute.make_rmem_tensor((4, nt, num_tiles), Float32)
+            acc.fill(0.0)
+            a = cute.make_rmem_tensor((8,), BFloat16)
+            a32 = cute.recast_tensor(a, Int32)
+            last = nrows - Int32(1)
+            for t in cutlass.range_constexpr(num_tiles):
+                cute.arch.mbarrier_wait(mbar.iterator + t, 0)
+                row_lo = Int32(t * _TILE_M) + grp
+                row_hi = row_lo + Int32(8)
+                if row_lo > last:
+                    row_lo = last
+                if row_hi > last:
+                    row_hi = last
+                base_lo = row_lo * Int32(half_pitch) + k_word0
+                base_hi = row_hi * Int32(half_pitch) + k_word0
+                for s in cutlass.range_constexpr(steps):
+                    a32[0] = slab32[base_lo + Int32(8 * s)]
+                    a32[1] = slab32[base_hi + Int32(8 * s)]
+                    a32[2] = slab32[base_lo + Int32(8 * s + 4)]
+                    a32[3] = slab32[base_hi + Int32(8 * s + 4)]
+                    for j in cutlass.range_constexpr(nt):
+                        acc[None, j, t] = mma_sync(a, vb[None, j, s], acc[None, j, t])
+            if cutlass.const_expr(trace):
+                if tidx == Int32(0):
+                    st_global_b64(stamp_addr(tr, bidx, 2), globaltimer())
+
+            # Fixed-order cross-warp reduction through the (now idle) slab.
+            cute.arch.sync_threads()
+            for t in cutlass.range_constexpr(num_tiles):
+                for j in cutlass.range_constexpr(nt):
+                    base = ((warp * Int32(num_tiles) + Int32(t)) * Int32(nt) + Int32(j)) * Int32(
+                        red_stride
+                    ) + lane * Int32(4)
+                    for e in cutlass.range_constexpr(4):
+                        slab_f32[base + Int32(e)] = acc[e, j, t]
+            cute.arch.sync_threads()
+
+            out_limit = out_rows
+            if out_limit > Int32(8 * nt):
+                out_limit = Int32(8 * nt)
+            buf_word0 = Int32(0)
+            mc_base = Int64(0)
+            if cutlass.const_expr(publish):
+                # After the wait: the protocol state written by the previous
+                # hop is final.
+                buf_word0 = ld_volatile_u32(flags.iterator.toint()) * Int32(_LAMPORT_BUFFER_WORDS)
+                mc_base = mc.iterator.toint()
+            pk = cute.make_rmem_tensor((2,), BFloat16)
+            pk32 = cute.recast_tensor(pk, Int32)
+            total_pairs = num_tiles * nt * 64
+            p = tidx
+            while p < Int32(total_pairs):
+                t_i = p // Int32(nt * 64)
+                j_i = (p // Int32(64)) % Int32(nt)
+                q = p % Int32(64)
+                gh = q // Int32(16)
+                quad_i = (q // Int32(4)) % Int32(4)
+                e_i = q % Int32(4)
+                idx0 = (
+                    (t_i * Int32(nt) + j_i) * Int32(red_stride)
+                    + (gh * Int32(8) + quad_i) * Int32(4)
+                    + e_i
+                )
+                idx1 = idx0 + Int32(16)
+                s0 = Float32(0.0)
+                s1 = Float32(0.0)
+                for wi in cutlass.range_constexpr(_WARPS):
+                    s0 = s0 + slab_f32[Int32(wi * num_tiles * nt * red_stride) + idx0]
+                    s1 = s1 + slab_f32[Int32(wi * num_tiles * nt * red_stride) + idx1]
+                m0 = t_i * Int32(_TILE_M) + gh * Int32(2) + Int32(8) * (e_i // Int32(2))
+                r = j_i * Int32(8) + Int32(2) * quad_i + e_i % Int32(2)
+                if m0 < nrows:
+                    if r < out_limit:
+                        pk[0] = s0.to(BFloat16)
+                        pk[1] = s1.to(BFloat16)
+                        out[r, n0 + m0] = pk[0]
+                        out[r, n0 + m0 + Int32(1)] = pk[1]
+                        if cutlass.const_expr(publish):
+                            word = pk32[0]
+                            lo = word & Int32(0xFFFF)
+                            hi = (word >> Int32(16)) & Int32(0xFFFF)
+                            if lo == Int32(0x8000):
+                                lo = Int32(0)
+                            if hi == Int32(0x8000):
+                                hi = Int32(0)
+                            word = lo | (hi << Int32(16))
+                            if cutlass.const_expr(vec_publish):
+                                slab32[Int32(stage0) + r * Int32(stage_pitch) + m0 // Int32(2)] = (
+                                    word
+                                )
+                            else:
+                                widx = buf_word0 + (
+                                    (r * Int32(_NUM_RANKS) + Int32(publish_rank)) * Int32(_HIDDEN)
+                                    + n0
+                                    + m0
+                                ) // Int32(2)
+                                st_global_b32(mc_base + Int64(widx) * Int64(4), word)
+                p = p + Int32(_THREADS)
+            if cutlass.const_expr(vec_publish):
+                cute.arch.sync_threads()
+                if tidx < out_limit * Int32(chunks):
+                    pr = tidx // Int32(chunks)
+                    pc = tidx % Int32(chunks)
+                    sb = Int32(stage0) + pr * Int32(stage_pitch) + pc * Int32(4)
+                    vidx = buf_word0 + (
+                        (pr * Int32(_NUM_RANKS) + Int32(publish_rank)) * Int32(_HIDDEN)
+                        + n0
+                        + pc * Int32(8)
+                    ) // Int32(2)
+                    st_global_v4_b32(
+                        mc_base + Int64(vidx) * Int64(4),
+                        slab32[sb],
+                        slab32[sb + Int32(1)],
+                        slab32[sb + Int32(2)],
+                        slab32[sb + Int32(3)],
+                    )
+            if cutlass.const_expr(trace):
+                if tidx == Int32(0):
+                    st_global_b64(stamp_addr(tr, bidx, 3), globaltimer())
+
+        @cute.jit
+        def __call__(
+            self,
+            v: cute.Tensor,
+            w: cute.Tensor,
+            out: cute.Tensor,
+            mc: cute.Tensor,
+            flags: cute.Tensor,
+            tr: cute.Tensor,
+            stream,
+        ):
+            self.kernel(v, w, out, mc, flags, tr).launch(
+                grid=(grid, 1, 1),
+                block=(_THREADS, 1, 1),
+                stream=stream,
+                use_pdl=use_pdl,
+                min_blocks_per_mp=1,
+            )
+
+    rows_sym, out_sym = cute.sym_int(), cute.sym_int()
+    v_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (rows_sym, k_dim), stride_order=(1, 0), assumed_align=16
+    )
+    w_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (n_dim, k_dim), stride_order=(1, 0), assumed_align=16
+    )
+    out_fake = cute_runtime.make_fake_compact_tensor(
+        BFloat16, (out_sym, n_dim), stride_order=(1, 0), assumed_align=16
+    )
+    mc_fake = cute_runtime.make_fake_compact_tensor(
+        Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=16
+    )
+    flags_fake = cute_runtime.make_fake_compact_tensor(
+        Int32, (4,), stride_order=(0,), assumed_align=16
+    )
+    tr_fake = cute_runtime.make_fake_compact_tensor(
+        Int64, (cute.sym_int(),), stride_order=(0,), assumed_align=16
+    )
+    stream = cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        StreamOprojStage(),
+        v_fake,
+        w_fake,
+        out_fake,
+        mc_fake,
+        flags_fake,
+        tr_fake,
+        stream=stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+_COMPILED = {}
+_DUMMIES = {}
+
+
+def _dummies(device: torch.device):
+    """Placeholder publish / trace tensors for variants that do not use them
+    (never read or written by those variants)."""
+    key = device.index or 0
+    if key not in _DUMMIES:
+        _DUMMIES[key] = (
+            torch.zeros(4, dtype=torch.int32, device=device),
+            torch.zeros(4, dtype=torch.int32, device=device),
+            torch.zeros(TRACE_POINTS, dtype=torch.int64, device=device),
+        )
+    return _DUMMIES[key]
+
+
+def stream_oproj_stage(
+    v: torch.Tensor,
+    w_o: torch.Tensor,
+    out: torch.Tensor,
+    use_pdl: bool,
+    early_rows: int,
+    grid: Optional[int] = None,
+    trigger_dependents: bool = False,
+    stage: str = "l2",
+    when: str = "landed",
+    stage_rows: int = -1,
+    prefetch_policy: str = "normal",
+    copy_policy: str = "none",
+    hop_transport=None,
+    vector_publish: bool = True,
+    trace: Optional[torch.Tensor] = None,
+) -> None:
+    """Launch the staging o_proj on the current stream.
+
+    Args:
+        v: ``[num_rows, K]`` contiguous BF16 (``num_rows <= 16``).
+        w_o: ``[N, K]`` contiguous BF16, 16-byte aligned.
+        out: ``[out_rows, N]`` contiguous BF16; rows past ``num_rows`` are
+            written as zeros (at most 16 rows written).
+        use_pdl: launch as a PDL dependent of the previous kernel.
+        early_rows: rows per CTA copied before ``griddepcontrol.wait``
+            (served: ``TailConfig.early_rows``).
+        grid: CTA count (default: the device's SM count).
+        trigger_dependents: execute ``griddepcontrol.launch_dependents`` once
+            every row copy is issued.
+        stage, when, stage_rows, prefetch_policy, copy_policy: oproj_weight_staging
+            (module docstring); ``stage_rows=-1`` stages every after-wait
+            row.
+        hop_transport: the ``IntoHopTransport`` of ``attention_output_into_hop``; ``out``
+            (1..8 rows, the hop's rows) is then also published into it.
+        vector_publish: 16-byte publish stores where the split permits.
+        trace: optional int64 ``[grid * TRACE_POINTS]`` buffer for
+            ``%globaltimer`` stamps (diagnostics only).
+    """
+    n_dim, k_dim = w_o.shape
+    if grid is None:
+        grid = _sm_count(w_o.device.index or 0)
+    reason = check_stage(stage, when, prefetch_policy, copy_policy)
+    if reason is not None:
+        raise ValueError(reason)
+    publish = hop_transport is not None
+    if publish:
+        if not 1 <= v.shape[0] <= out.shape[0] <= 8 or n_dim != _HIDDEN:
+            raise ValueError(
+                "the post_attention_hop transport holds 1..8 rows of 6144 and out must hold "
+                f"every live row; got v rows {v.shape[0]}, out rows "
+                f"{out.shape[0]}, N {n_dim}"
+            )
+    n_tiles_r = 1 if out.shape[0] <= 8 else 2
+    if not use_pdl:
+        early_rows = 1 << 30
+    rows_max = plan(n_dim, k_dim, grid)["rows_max"]
+    late = max(0, rows_max - min(early_rows, rows_max))
+    if stage_rows < 0 or stage_rows > late:
+        stage_rows = late
+    if stage == "none":
+        stage_rows = 0
+    rank = hop_transport.rank if publish else 0
+    key = (
+        n_dim,
+        k_dim,
+        grid,
+        n_tiles_r,
+        use_pdl,
+        early_rows,
+        trigger_dependents,
+        publish,
+        rank,
+        vector_publish and publish,
+        stage,
+        when,
+        stage_rows,
+        prefetch_policy,
+        copy_policy,
+        trace is not None,
+    )
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        compiled = _build(
+            n_dim,
+            k_dim,
+            grid,
+            n_tiles_r,
+            use_pdl,
+            early_rows,
+            trigger_dependents,
+            publish,
+            rank,
+            vector_publish and publish,
+            stage,
+            when,
+            stage_rows,
+            prefetch_policy,
+            copy_policy,
+            trace is not None,
+        )
+        _COMPILED[key] = compiled
+    mc_dummy, flags_dummy, tr_dummy = _dummies(w_o.device)
+    mc = hop_transport.mc_words if publish else mc_dummy
+    flags = hop_transport.flags if publish else flags_dummy
+    compiled(v, w_o, out, mc, flags, trace if trace is not None else tr_dummy)
+
+
+__all__ = [
+    "COPY_POLICIES",
+    "PREFETCH_POLICIES",
+    "STAGES",
+    "TRACE_POINTS",
+    "WHENS",
+    "check_shapes",
+    "check_stage",
+    "plan",
+    "stream_oproj_stage",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/kernels.py
new file mode 100644
index 000000000000..8202cfd6ba0f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/kernels.py
@@ -0,0 +1,389 @@
+# 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.
+"""oproj_weight_staging ``oproj_weight_staging``: staging configuration and the functional attention tail.
+
+``stage_tail`` computes exactly what attention_output_tail's ``attn_tail`` computes with a
+``cute_stream`` ``TailConfig`` (served: ``INTO_HOP_SPARSE_DECODE_SCHEDULE``)::
+
+    v = bf16(latent @ w_uv ^ T)  # attention_output_tail's _wuv_kernel (or its fork)
+    partial = bf16(v @ w_o ^ T)  # cute_oproj_stage (fork of the
+    # served streaming o_proj)
+
+with the same kernels' arithmetic, so the partial (and the attention_output_into_hop words)
+are bit-identical to attention_output_tail's. ``StageConfig`` only decides how the o_proj rows
+that the served schedule streams after its PDL wait are staged before it:
+
+* host ``"oproj"``: the o_proj kernel itself, parked in
+  ``griddepcontrol.wait``, prefetches them into L2 (``stage="l2"``) or
+  copies them into its shared-memory slab (``stage="smem"``), right after
+  its early copies (``when="entry"``), once they landed
+  (``when="landed"``) or tile by tile behind them (``when="chain"``);
+* host ``"wuv"``: a fork of attention_output_tail's ``_wuv_kernel`` prefetches them into L2
+  before its own wait on the attention kernel (one o_proj CTA's rows per
+  W_UV CTA), i.e. earlier than the o_proj becomes resident.
+
+No launch and no stream is added in any configuration.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ..attention_output_tail.kernels import (
+    ROW_TILE,
+    TailConfig,
+    launch_wuv,
+    pdl_enabled,
+    pdl_entry_enabled,
+)
+from .cute_oproj_stage import check_shapes, check_stage, plan, stream_oproj_stage
+
+HOSTS = ("oproj", "wuv")
+_WUV_PREFETCH_PIECES = 128
+
+
+@dataclass(frozen=True)
+class StageConfig:
+    """How the after-wait ``w_o`` rows of the o_proj are staged (oproj_weight_staging)."""
+
+    stage: str = "l2"
+    """``"none"`` (served schedule, fork only), ``"l2"`` (L2 prefetch before
+    the wait; the rows are still copied after it) or ``"smem"`` (copied into
+    the slab before the wait)."""
+    host: str = "oproj"
+    """Kernel that stages: ``"oproj"`` or ``"wuv"`` (``stage="l2"`` only)."""
+    when: str = "landed"
+    """``host="oproj"``: ``"entry"`` (right after the early copies),
+    ``"landed"`` (once the early-only tiles are in shared memory) or
+    ``"chain"`` (``stage="smem"``: tile by tile, each once the previous tile
+    is in shared memory)."""
+    rows: int = -1
+    """After-wait rows per CTA to stage (-1: all of them)."""
+    prefetch_policy: str = "normal"
+    """L2 policy of the staging prefetch: ``"normal"`` or ``"evict_last"``
+    (never evict-first)."""
+    copy_policy: str = "none"
+    """L2 hint of the o_proj's ``w_o`` copies: ``"none"`` or
+    ``"evict_first"`` (consumer read of a single-use stream)."""
+    early_rows: Optional[int] = None
+    """Override of ``TailConfig.early_rows`` (None: the tail's)."""
+
+
+DEFAULT_STAGE = StageConfig(stage="smem", when="landed", early_rows=24, copy_policy="evict_first")
+"""Default schedule: 24 rows per CTA issued at entry, the other 24 copied
+into shared memory by the second warp once tile 0 has landed, all before
+the wait; every ``w_o`` copy (the single-use consumer read) with an L2
+``evict_first`` hint."""
+
+
+def check_stage_config(stage: StageConfig) -> Optional[str]:
+    """Return why ``stage`` is invalid, or ``None``."""
+    if stage.host not in HOSTS:
+        return f"host {stage.host!r} not in {HOSTS}"
+    if stage.host == "wuv" and stage.stage not in ("l2", "none"):
+        return "the W_UV host stages into L2 only"
+    return check_stage(stage.stage, stage.when, stage.prefetch_policy, stage.copy_policy)
+
+
+def check_tail_config(config: TailConfig, n_dim: int, k_dim: int, sm_count: int) -> Optional[str]:
+    """Return why oproj_weight_staging cannot replace the o_proj of this attention_output_tail config."""
+    if config.oproj_engine != "cute_stream":
+        return f"o_proj engine {config.oproj_engine!r} is not cute_stream"
+    if config.prefetch_l2 or config.successor_prefetch:
+        return "attention_output_tail_prefetch off-path prefetch configs keep the attention_output_tail tail"
+    if config.trigger_late:
+        return "trigger_late is an experiment of the attention_output_tail tail"
+    return check_shapes(n_dim, k_dim, config.oproj_grid or sm_count)
+
+
+@triton.jit
+def _gdc_wait_with_memory_clobber():
+    tl.inline_asm_elementwise(
+        "griddepcontrol.wait; // dummy $0",
+        "=r,~{memory}",
+        [],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _prefetch_l2(addr, nbytes):
+    tl.inline_asm_elementwise(
+        "cp.async.bulk.prefetch.L2.global [$1], $2; // dummy $0",
+        "=r,l,r",
+        [addr, nbytes],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _prefetch_l2_evict_last(addr, nbytes):
+    tl.inline_asm_elementwise(
+        "{ .reg .b64 pol; "
+        "createpolicy.fractional.L2::evict_last.b64 pol, 1.0; "
+        "cp.async.bulk.prefetch.L2.global.L2::cache_hint [$1], $2, pol; } "
+        "// dummy $0",
+        "=r,l,r",
+        [addr, nbytes],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _wuv_stage_kernel(
+    latent_ptr,
+    w_uv_ptr,
+    v_ptr,
+    w_o_addr,
+    num_rows,
+    stride_lat_row,
+    stride_lat_head,
+    stride_wuv_head,
+    stride_wuv_v,
+    stride_v_row,
+    C: tl.constexpr,
+    V: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    WAIT_PRODUCER: tl.constexpr,
+    OPROJ_GRID: tl.constexpr,
+    HALF_N: tl.constexpr,
+    ROW_BYTES: tl.constexpr,
+    STAGE_FIRST: tl.constexpr,
+    STAGE_BYTES: tl.constexpr,
+    PIECES: tl.constexpr,
+    EVICT_LAST: tl.constexpr,
+):
+    """attention_output_tail's ``_wuv_kernel`` plus, before its wait, the L2 prefetch of the
+    after-wait ``w_o`` rows of o_proj CTA ``pid`` (rows ``n0 + STAGE_FIRST``
+    onward, ``STAGE_BYTES`` bytes; ``n0`` as in the o_proj's even split)."""
+    head = tl.program_id(0)
+    v_block = tl.program_id(1)
+    if USE_PDL:
+        # Let the dependent o_proj kernel start its weight stream now; it
+        # still waits for this whole grid before reading v.
+        tl.extra.cuda.gdc_launch_dependents()
+
+    offs_v = v_block * BLOCK_V + tl.arange(0, BLOCK_V)
+    offs_r = tl.arange(0, BLOCK_R)
+    offs_c = tl.arange(0, C)
+    row_mask = offs_r < num_rows
+    # Weights are immutable: load them before waiting for the producer.
+    w = tl.load(
+        w_uv_ptr + head * stride_wuv_head + offs_v[:, None] * stride_wuv_v + offs_c[None, :]
+    )
+    pid = head * tl.num_programs(1) + v_block
+    if pid < OPROJ_GRID:
+        # oproj_weight_staging: stage o_proj CTA pid's after-wait rows into L2 (immutable
+        # weight; a prefetch has no architectural effect).
+        n0 = 2 * ((pid * HALF_N) // OPROJ_GRID)
+        piece: tl.constexpr = STAGE_BYTES // PIECES
+        offs_p = tl.arange(0, PIECES).to(tl.int64)
+        addr = w_o_addr + (n0 + STAGE_FIRST).to(tl.int64) * ROW_BYTES + offs_p * piece
+        if EVICT_LAST:
+            _prefetch_l2_evict_last(addr, tl.full([PIECES], piece, tl.int32))
+        else:
+            _prefetch_l2(addr, tl.full([PIECES], piece, tl.int32))
+    if WAIT_PRODUCER:
+        # Launched as a PDL dependent of the latent's producer: every read of
+        # the latent and every global write follows this wait.
+        _gdc_wait_with_memory_clobber()
+    lat_t = tl.load(
+        latent_ptr + offs_r[None, :] * stride_lat_row + head * stride_lat_head + offs_c[:, None],
+        mask=row_mask[None, :],
+        other=0.0,
+    )
+    acc = tl.dot(w, lat_t)
+    v_out = acc.to(v_ptr.dtype.element_ty)
+    tl.store(
+        v_ptr + offs_r[None, :] * stride_v_row + head * V + offs_v[:, None],
+        v_out,
+        mask=row_mask[None, :],
+    )
+
+
+def wuv_stage_geometry(
+    n_dim: int, k_dim: int, grid: int, early_rows: int, rows: int
+) -> Optional[tuple]:
+    """``(first_row, nbytes)`` of the per-CTA prefetch of the W_UV host, or
+    ``None`` when every CTA's after-wait rows cannot be covered evenly."""
+    geo = plan(n_dim, k_dim, grid)
+    if geo["rows_min"] != geo["rows_max"]:
+        return None
+    late = geo["rows_max"] - min(early_rows, geo["rows_max"])
+    if rows < 0 or rows > late:
+        rows = late
+    nbytes = rows * k_dim * 2
+    if nbytes == 0 or nbytes % (_WUV_PREFETCH_PIECES * 16):
+        return None
+    return early_rows, nbytes
+
+
+def launch_wuv_stage(
+    latent: torch.Tensor,
+    w_uv: torch.Tensor,
+    v: torch.Tensor,
+    num_rows: int,
+    w_o: torch.Tensor,
+    oproj_grid: int,
+    early_rows: int,
+    stage: StageConfig,
+    config: TailConfig,
+    trigger_dependents: bool,
+    wait_producer: bool,
+) -> None:
+    """Launch the W_UV fork that also stages the o_proj's after-wait rows."""
+    heads, v_dim, c_dim = w_uv.shape
+    n_dim, k_dim = w_o.shape
+    if latent.dim() == 2:
+        stride_lat_row, stride_lat_head = latent.stride(0), c_dim
+    else:
+        stride_lat_row, stride_lat_head = latent.stride(0), latent.stride(1)
+    geo = wuv_stage_geometry(n_dim, k_dim, oproj_grid, early_rows, stage.rows)
+    if geo is None:
+        raise ValueError("W_UV host cannot cover the o_proj's after-wait rows")
+    first, nbytes = geo
+    _wuv_stage_kernel[(heads, v_dim // config.wuv_block_v)](
+        latent,
+        w_uv,
+        v,
+        w_o.data_ptr(),
+        num_rows,
+        stride_lat_row,
+        stride_lat_head,
+        w_uv.stride(0),
+        w_uv.stride(1),
+        v.stride(0),
+        C=c_dim,
+        V=v_dim,
+        BLOCK_V=config.wuv_block_v,
+        BLOCK_R=ROW_TILE,
+        USE_PDL=trigger_dependents,
+        WAIT_PRODUCER=wait_producer,
+        OPROJ_GRID=oproj_grid,
+        HALF_N=n_dim // 2,
+        ROW_BYTES=k_dim * 2,
+        STAGE_FIRST=first,
+        STAGE_BYTES=nbytes,
+        PIECES=_WUV_PREFETCH_PIECES,
+        EVICT_LAST=stage.prefetch_policy == "evict_last",
+        num_warps=config.wuv_warps,
+        num_stages=1,
+        **({"launch_pdl": True} if wait_producer else {}),
+    )
+
+
+def stage_tail(
+    latent: torch.Tensor,
+    w_uv: torch.Tensor,
+    w_o: torch.Tensor,
+    num_rows: int,
+    out_rows: Optional[int] = None,
+    use_pdl: Optional[bool] = None,
+    v_scratch: Optional[torch.Tensor] = None,
+    config: TailConfig = TailConfig(oproj_grid=128, early_rows=32),
+    stage: StageConfig = DEFAULT_STAGE,
+    hop_transport=None,
+    trace: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """attention_output_tail's tail (``cute_stream`` engine) with oproj_weight_staging of the o_proj rows.
+
+    Arguments as ``attention_output_tail.kernels.attn_tail`` (the caller validates them with
+    ``check_tail_shapes`` and the config with ``check_tail_config``), plus
+    the staging configuration and an optional ``trace`` buffer for the
+    o_proj's ``%globaltimer`` stamps (diagnostics only).
+
+    Returns:
+        ``[out_rows, N]`` BF16 partial, bit-identical to ``attn_tail``'s.
+    """
+    if out_rows is None:
+        out_rows = num_rows
+    if use_pdl is None:
+        use_pdl = pdl_enabled()
+    n_dim, k_dim = w_o.shape
+    v = v_scratch
+    if v is None:
+        v = torch.empty((num_rows, k_dim), dtype=w_uv.dtype, device=latent.device)
+    out = torch.empty((out_rows, n_dim), dtype=w_o.dtype, device=latent.device)
+    if out_rows > ROW_TILE:
+        out[ROW_TILE:].zero_()
+    grid = config.oproj_grid or torch.cuda.get_device_properties(w_o.device).multi_processor_count
+    early = config.early_rows if stage.early_rows is None else stage.early_rows
+    wait_producer = config.pdl_entry and pdl_entry_enabled()
+    wuv_host = stage.host == "wuv" and stage.stage == "l2" and use_pdl
+    if wuv_host:
+        launch_wuv_stage(
+            latent,
+            w_uv,
+            v,
+            num_rows,
+            w_o,
+            grid,
+            early,
+            stage,
+            config,
+            trigger_dependents=use_pdl,
+            wait_producer=wait_producer,
+        )
+    else:
+        launch_wuv(
+            latent,
+            w_uv,
+            v,
+            num_rows,
+            config,
+            trigger_dependents=use_pdl,
+            wait_producer=wait_producer,
+        )
+    stream_oproj_stage(
+        v,
+        w_o,
+        out,
+        use_pdl=use_pdl,
+        early_rows=early,
+        grid=grid,
+        trigger_dependents=config.trigger_dependents and pdl_enabled(),
+        stage="none" if wuv_host else stage.stage,
+        when=stage.when,
+        stage_rows=stage.rows,
+        prefetch_policy=stage.prefetch_policy,
+        copy_policy=stage.copy_policy,
+        hop_transport=hop_transport,
+        vector_publish=config.into_hop_vector_publish,
+        trace=trace,
+    )
+    return out
+
+
+__all__ = [
+    "DEFAULT_STAGE",
+    "HOSTS",
+    "StageConfig",
+    "check_stage_config",
+    "check_tail_config",
+    "launch_wuv_stage",
+    "stage_tail",
+    "wuv_stage_geometry",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/tail.py
new file mode 100644
index 000000000000..cfca1f72ab7f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/tail.py
@@ -0,0 +1,208 @@
+# 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.
+"""oproj_weight_staging ``oproj_weight_staging`` adapter: the attention_output_tail attention tail
+with staged o_proj rows.
+
+``wrap_attn_tail(mla, tail)`` is called by the MLA adapter right after attention_output_tail's
+``make_attn_tail`` when ``oproj_weight_staging`` is selected. It returns a ``StagedAttnTail``
+that carries the same tail schedule (``TailConfig``) and takes every
+per-call decision exactly as attention_output_tail's ``AttnTail`` (``arm``, ``defer_w_uv``, the
+attention_output_into_hop request), so a declined call runs the native path with no
+oproj_weight_staging work issued. Only the launch in ``project`` differs: ``kernels.stage_tail``
+instead of ``attention_output_tail.kernels.attn_tail``, with the same arithmetic.
+
+The decisions are static (logged once per rank and layer): at construction
+oproj_weight_staging takes only one schedule, attention_output_tail's
+``INTO_HOP_SPARSE_DECODE_SCHEDULE`` (128 CTAs, 32 early rows: the served tail behind
+sparse_mla_decode with the attention_output_into_hop publish; otherwise attention_output_tail's
+tail is returned unchanged); on the first ``arm``, together with attention_output_tail's own static
+decision and before any device work, it checks the staging geometry on the layer's weights
+(otherwise the layer runs attention_output_tail's ``project``).
+"""
+
+from typing import Optional
+
+import torch
+
+from ... import dispatch
+from ..attention_output_tail.kernels import INTO_HOP_SPARSE_DECODE_SCHEDULE
+from ..attention_output_tail.tail import MAX_INTO_HOP_ROWS, AttnTail
+from .kernels import (
+    DEFAULT_STAGE,
+    StageConfig,
+    check_stage_config,
+    check_tail_config,
+    stage_tail,
+    wuv_stage_geometry,
+)
+
+FEATURE_ID = "oproj_weight_staging"
+
+
+class StagedAttnTail(AttnTail):
+    """attention_output_tail whose o_proj stages its after-wait ``w_o`` rows (oproj_weight_staging).
+
+    Args:
+        layer_key: name used in capability logs, e.g. ``"layer3"``.
+        config: the attention_output_tail schedule (served: ``INTO_HOP_SPARSE_DECODE_SCHEDULE``).
+        stage: the oproj_weight_staging configuration.
+    """
+
+    def __init__(self, layer_key: str, config, stage: StageConfig):
+        super().__init__(layer_key, config)
+        self.stage = stage
+        self.staging_calls = 0
+        self.staging_reason: Optional[str] = None
+        """Why this layer's o_proj is not staged (``None`` if it is); set
+        with attention_output_tail's static decision on the first ``arm``."""
+        self._staging_ok: Optional[bool] = None
+        self.trace: Optional[torch.Tensor] = None
+        """Diagnostics only (benchmarks): int64 ``[grid * TRACE_POINTS]``
+        buffer for the o_proj's ``%globaltimer`` stamps of every call."""
+
+    def _check_static(self, mla) -> bool:
+        ok = super()._check_static(mla)
+        if self._staging_ok is None:
+            reason = None if ok else "attention_output_tail tail not engaged on this layer"
+            if reason is None:
+                reason = geometry_reason(mla, self.config, self.stage)
+            self.staging_reason = reason
+            self._staging_ok = dispatch.log_capability(
+                FEATURE_ID,
+                self.layer_key,
+                reason is None,
+                reason
+                or (
+                    f"attention_output_tail tail o_proj (early_rows={self.config.early_rows}, "
+                    f"grid={self.config.oproj_grid or 'SMs'}) with "
+                    f"{self.stage}"
+                ),
+            )
+        return ok
+
+    def project(
+        self, mla, attn_output: torch.Tensor, all_reduce_params=None
+    ) -> Optional[torch.Tensor]:
+        """As ``AttnTail.project`` for a ``cute_stream`` schedule without the
+        attention_output_tail_prefetch, launching ``stage_tail``. Returns ``None`` (native
+        projection) when ``defer_w_uv`` declined the call; runs attention_output_tail's own
+        ``project`` when the staging geometry does not fit this layer."""
+        if not self._staging_ok:
+            return super().project(mla, attn_output, all_reduce_params)
+        del all_reduce_params
+        self._armed = False
+        packet = self._pending
+        if packet is None:
+            self._join_prefetch()
+            return None
+        self._pending = None
+        if attn_output.data_ptr() != packet.output_ptr or attn_output.shape[0] < packet.num_rows:
+            self._join_prefetch()
+            raise RuntimeError(
+                f"[glm_moe_dsa] oproj_weight_staging/{self.layer_key}: tail packet does not match "
+                "the output projection input"
+            )
+        hop_transport = None
+        request, self._into_hop = self._into_hop, None
+        if request is not None:
+            transport, hop_rows = request
+            # The hop polls hop_rows rows and later resets exactly those rows
+            # of the buffer, so the producer must publish exactly them.
+            if attn_output.shape[0] == hop_rows <= MAX_INTO_HOP_ROWS:
+                hop_transport = transport
+            else:
+                self.into_hop_declined_calls += 1
+        partial = stage_tail(
+            packet.latent,
+            mla.v_b_proj,
+            mla.o_proj.weight,
+            packet.num_rows,
+            out_rows=attn_output.shape[0],
+            config=self.config,
+            stage=self.stage,
+            hop_transport=hop_transport,
+            trace=self.trace,
+        )
+        if hop_transport is not None:
+            self._into_hop_published = True
+            self.into_hop_published_calls += 1
+        self.staging_calls += 1
+        self._join_prefetch()
+        return partial
+
+
+def config_reason(config, stage: StageConfig) -> Optional[str]:
+    """Static reason (tail schedule and staging configuration only) why oproj_weight_staging
+    cannot wrap a tail with ``config``, or ``None``."""
+    reason = check_stage_config(stage)
+    if reason is not None:
+        return reason
+    if config.oproj_engine != "cute_stream":
+        return f"o_proj engine {config.oproj_engine!r} is not cute_stream"
+    if config.prefetch_l2 or config.successor_prefetch:
+        return "attention_output_tail_prefetch off-path prefetch configs keep the attention_output_tail tail"
+    if config.trigger_late:
+        return "trigger_late is an experiment of the attention_output_tail tail"
+    if config != INTO_HOP_SPARSE_DECODE_SCHEDULE:
+        return (
+            "oproj_weight_staging is measured only on the served attention_output_into_hop/sparse_mla_decode schedule "
+            "(128 CTAs, 32 early rows)"
+        )
+    return None
+
+
+def geometry_reason(mla, config, stage: StageConfig) -> Optional[str]:
+    """Reason from the module's weights (first ``arm``) why the staging
+    o_proj cannot run this layer, or ``None``."""
+    w_o = mla.o_proj.weight
+    n_dim, k_dim = w_o.shape
+    sm_count = torch.cuda.get_device_properties(w_o.device).multi_processor_count
+    reason = check_tail_config(config, n_dim, k_dim, sm_count)
+    if reason is not None:
+        return reason
+    if stage.host == "wuv" and stage.stage == "l2":
+        grid = config.oproj_grid or sm_count
+        heads, v_dim, _ = mla.v_b_proj.shape
+        if heads * (v_dim // config.wuv_block_v) < grid:
+            return "W_UV grid smaller than the o_proj grid"
+        early = config.early_rows if stage.early_rows is None else stage.early_rows
+        if wuv_stage_geometry(n_dim, k_dim, grid, early, stage.rows) is None:
+            return "W_UV host cannot cover the o_proj's after-wait rows"
+    return None
+
+
+def wrap_attn_tail(mla, tail: Optional[AttnTail], stage: Optional[StageConfig] = None):
+    """Return a ``StagedAttnTail`` replacing attention_output_tail's ``tail``, or ``tail`` itself.
+
+    Called once per MLA module at construction, after attention_output_tail's
+    ``make_attn_tail``; ``None`` (no attention_output_tail on this layer) stays ``None``.
+    ``stage`` defaults to ``DEFAULT_STAGE``. The schedule-level
+    decision is made here (logged when it declines); the weight-geometry
+    decision is made with attention_output_tail's static decision on the first ``arm``, before any
+    device work.
+    """
+    del mla
+    if tail is None:
+        return None
+    if stage is None:
+        stage = DEFAULT_STAGE
+    reason = config_reason(tail.config, stage)
+    if reason is not None:
+        dispatch.log_capability(FEATURE_ID, tail.layer_key, False, reason)
+        return tail
+    return StagedAttnTail(tail.layer_key, tail.config, stage)
+
+
+__all__ = ["FEATURE_ID", "StagedAttnTail", "config_reason", "geometry_reason", "wrap_attn_tail"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/__init__.py
new file mode 100644
index 000000000000..f9babcaff681
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/__init__.py
@@ -0,0 +1,74 @@
+# 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.
+"""Package post_attention_hop: fused post-attention hop (cross-rank sum + residual + RMSNorm + router).
+
+Entry points registered with ``dispatch``:
+
+* ``make_post_attn_hop(mapping, hidden_size, dtype, native_allreduce=None)``
+  returns the TP group's :class:`PostAttnHop` (or None for the native path)
+  for a call site whose native ``AllReduce`` is ``native_allreduce``. Call it at
+  the same point of model construction on every rank.
+* ``engagement_counts()`` returns the host-side engaged/native-fallback call
+  counters (eager calls and CUDA graph captures).
+* ``make_into_hop_transport(hop)`` (``attention_output_into_hop`` only) returns the into-hop
+  transport of a hop (see ``from_partials.py``), or None.
+* ``into_hop_engagement_counts()`` (``attention_output_into_hop`` only) returns the
+  host-side into-hop call counters; an into-hop call replaces a post_attention_hop call.
+"""
+
+from ...dispatch import register
+from .hop import (
+    PostAttnHop,
+    engagement_counts,
+    native_allreduce_ineligibility,
+    op_available,
+    required_sms,
+    reset_engagement_counts,
+    static_ineligibility,
+)
+
+register("post_attention_hop", "make_post_attn_hop", PostAttnHop.get_or_create)
+register("post_attention_hop", "engagement_counts", engagement_counts)
+
+
+def _make_into_hop_transport(hop):
+    """``attention_output_into_hop`` only: the into-hop transport of ``hop`` (imports ``from_partials`` lazily, so
+    the post_attention_hop itself is unaffected)."""
+    from .from_partials import IntoHopTransport
+
+    return IntoHopTransport.create(hop)
+
+
+register("post_attention_hop", "make_into_hop_transport", _make_into_hop_transport)
+
+
+def _into_hop_engagement_counts():
+    """``attention_output_into_hop`` only: the into-hop call counters (imports ``from_partials`` lazily)."""
+    from .from_partials import into_hop_engagement_counts
+
+    return into_hop_engagement_counts()
+
+
+register("post_attention_hop", "into_hop_engagement_counts", _into_hop_engagement_counts)
+
+__all__ = [
+    "PostAttnHop",
+    "engagement_counts",
+    "native_allreduce_ineligibility",
+    "op_available",
+    "required_sms",
+    "reset_engagement_counts",
+    "static_ineligibility",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/from_partials.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/from_partials.py
new file mode 100644
index 000000000000..247a3ccc8225
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/from_partials.py
@@ -0,0 +1,207 @@
+# 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.
+"""into-hop entry of the post_attention_hop (``attention_output_into_hop``).
+
+Reachable only when ``attention_output_into_hop`` is selected: the attention tail's o_proj kernel
+writes its rank-local partial directly into the current Lamport buffer of the
+post_attention_hop workspace (multicast, sanitized like post_attention_hop's own publish), and the
+hop then runs ``trtllm::glm_moe_dsa_post_attention_hop_from_partials``, which takes every rank's
+partial from that transport. Its normalized hidden, residual and router logits are
+bit-identical to the post_attention_hop consuming the ordinary partial (see
+``cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h`` for the
+arithmetic and the memory-order contract with the producer).
+
+``IntoHopTransport`` exposes what the producer needs (the multicast view of the
+workspace, the protocol state and the rank) and runs the hop. It never
+changes the post_attention_hop object or its behaviour; both kinds of calls share the
+workspace and its buffer rotation. Its eligibility check is the post_attention_hop's
+per-call check without the hop's fallback counting (the layer evaluates the
+post_attention_hop itself when the into-hop path is not taken).
+
+Host-side counters (:func:`into_hop_engagement_counts`) record into-hop calls (eager
+runs and CUDA graph captures, not replays). An into-hop call replaces a post_attention_hop
+call, so the post-attention hop calls of a run are the post_attention_hop's ``engaged``
+count plus ``into_hop_engaged``.
+"""
+
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from ...dispatch import log_capability
+from .hop import HIDDEN_SIZE, LAMPORT_BUFFER_BYTES, NUM_RANKS, WORKSPACE_BYTES, PostAttnHop
+
+FEATURE = "attention_output_into_hop"
+_OP_NAME = "glm_moe_dsa_post_attention_hop_from_partials"
+LAMPORT_BUFFER_WORDS = LAMPORT_BUFFER_BYTES // 4
+
+# Host-side call counters; they count Python calls (eager runs and CUDA graph
+# captures), not graph replays.
+_counts: Dict[str, int] = {}
+
+
+def _count(key: str) -> None:
+    _counts[key] = _counts.get(key, 0) + 1
+
+
+def into_hop_engagement_counts() -> Dict[str, int]:
+    """Copy of the into-hop call counters.
+
+    Keys: ``into_hop_engaged`` (into-hop calls) and ``into_hop_engaged_router`` (into-hop entry
+    calls that also produced router logits).
+    """
+    return dict(_counts)
+
+
+def reset_into_hop_engagement_counts() -> None:
+    """Clear the into-hop call counters."""
+    _counts.clear()
+
+
+def op_available() -> bool:
+    """Whether ``trtllm::glm_moe_dsa_post_attention_hop_from_partials`` is loaded."""
+    return hasattr(torch.ops.trtllm, _OP_NAME)
+
+
+def _agreed_op_reason(hop: PostAttnHop) -> str:
+    """Why no rank of ``hop``'s group may use the into-hop entry ("" when every rank may).
+
+    The op's availability is a fact of each rank's build. It is agreed once
+    per hop over the hop's TP communicator, so a rank without the op cannot
+    leave its peers publishing into a transport whose hop it never runs.
+    """
+    cached = getattr(hop, "_into_hop_op_reason", None)
+    if cached is not None:
+        return cached
+    local = "" if op_available() else f"trtllm::{_OP_NAME} is not loaded"
+    comm = getattr(hop, "comm", None)
+    agreed = comm is None or comm.allreduce(0 if local else 1) == comm.Get_size()
+    reason = local or ("" if agreed else f"trtllm::{_OP_NAME} is not loaded on another rank")
+    hop._into_hop_op_reason = reason
+    return reason
+
+
+class IntoHopTransport:
+    """Producer view of one post_attention_hop workspace plus the transport-input hop."""
+
+    def __init__(self, hop: PostAttnHop):
+        self.hop = hop
+        self.rank = hop.rank
+        self.flags = hop.flags
+        # Multicast view of the whole workspace as 32-bit words: a store to
+        # word i reaches word i of every rank's unicast buffer.
+        self.mc_words = hop.handle.get_mc_buffer((WORKSPACE_BYTES // 4,), torch.int32, 0)
+        # Poll the transport only after griddepcontrol.wait: the hop's CTAs
+        # can be resident next to the producer's, and polling before then
+        # competes with the running o_proj for the SM's issue slots.
+        self.early_poll = False
+        # Where the hop loads its router-weight fragments (0: after the
+        # barrier that ends the rank sum; 1: right after the rank sum; 2:
+        # right before the MMAs). Same outputs. 1 issues the loads early
+        # enough that their latency overlaps the RMS tree and the cluster
+        # exchange, while a hop CTA still fits next to a CTA of the
+        # publishing o_proj.
+        self.router_load = 1
+
+    @classmethod
+    def create(cls, hop: Optional[PostAttnHop]) -> Optional["IntoHopTransport"]:
+        """Return the transport of ``hop``, or None (the native path, or the
+        post_attention_hop without the into-hop transport).
+
+        Rank-uniform: ``hop`` exists on all ranks of the group or on none
+        (its creation is agreed), and the availability of the into-hop op is agreed
+        over the hop's TP group on the first call for that hop, which every
+        rank makes at the same point of model construction.
+        """
+        reason = "post_attention_hop hop unavailable" if hop is None else _agreed_op_reason(hop)
+        log_capability(FEATURE, "into_hop_transport", not reason, reason)
+        return None if reason else cls(hop)
+
+    def can_run(
+        self,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        router_weight: Optional[torch.Tensor],
+    ) -> bool:
+        """Same eligibility as the post_attention_hop for this call (rank-uniform facts,
+        no device work). Not counted as a post_attention_hop fallback: when the into-hop path is
+        not taken, the layer runs the post_attention_hop's own check."""
+        return self.hop.call_ineligibility(residual, residual, norm_weight, router_weight) is None
+
+    def __call__(
+        self,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+        router_weight: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Run the hop on the partials the producer already published."""
+        _count("into_hop_engaged")
+        if router_weight is not None:
+            _count("into_hop_engaged_router")
+        normed, residual_out, logits = getattr(torch.ops.trtllm, _OP_NAME)(
+            residual.contiguous(),
+            norm_weight.contiguous(),
+            float(eps),
+            router_weight,
+            self.hop.uc_buffer,
+            self.flags,
+            self.rank,
+            NUM_RANKS,
+            self.early_poll,
+            self.router_load,
+        )
+        return normed, residual_out, (logits if router_weight is not None else None)
+
+
+def _register_fake() -> None:
+    if not op_available():
+        return
+    try:
+
+        @torch.library.register_fake(f"trtllm::{_OP_NAME}")
+        def _(
+            residual,
+            gamma,
+            eps,
+            router_weight,
+            uc_buffer,
+            flags,
+            rank,
+            nranks,
+            early_poll,
+            router_load=1,
+        ):
+            rows = residual.numel() // HIDDEN_SIZE
+            logits = (
+                residual.new_empty((rows, 256), dtype=torch.float32)
+                if router_weight is not None
+                else residual.new_empty((0,), dtype=torch.float32)
+            )
+            return [torch.empty_like(residual), torch.empty_like(residual), logits]
+    except RuntimeError:
+        # Already registered by an earlier import in this process.
+        pass
+
+
+_register_fake()
+
+__all__ = [
+    "LAMPORT_BUFFER_WORDS",
+    "IntoHopTransport",
+    "op_available",
+    "reset_into_hop_engagement_counts",
+    "into_hop_engagement_counts",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/hop.py
new file mode 100644
index 000000000000..7dff81a27112
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/hop.py
@@ -0,0 +1,430 @@
+# 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.
+"""post_attention_hop: fused post-attention hop (the router contract).
+
+One call replaces the native fused MNNVL one-shot ``AllReduce(RESIDUAL_RMS_NORM)``
+of the post-attention boundary (``PRE_MOE_FUSION`` of MoE and MTP layers) and,
+optionally, the router GEMM that consumes its output:
+
+* ``residual_out = bf16(sum_rank partial) + residual`` and
+  ``normed = RMSNorm(residual_out) * gamma`` are bit-identical to the native
+  ``AllReduce(RESIDUAL_RMS_NORM)`` result on the MNNVL path, which normalizes
+  the bf16-rounded ``residual_out``. They are not bit-identical to the unfused
+  chain of a plain ``AllReduce`` followed by a separate add+RMSNorm (the native
+  path of layers without ``PRE_MLP_FUSION``, e.g. the dense layers), which
+  normalizes the FP32 sum: there only ``residual_out`` matches and ``normed``
+  can differ by up to two bf16 ulps. Such call sites are not a supported
+  destination.
+* ``router_logits = normed @ router_weight.T`` in FP32 (``[R, 256]``), computed
+  from the reduced, normalized rows, identical on every rank. They are not
+  bit-identical to the native ``dsv3_router_gemm_op`` logits: both accumulate
+  the same bf16 operands in FP32, but in a different order (tensor-core tiles
+  summed in a fixed order), with an error against an FP64 reference of the
+  same size as native's. Routing that consumes them can
+  therefore select a different expert set than native on rows whose 8th and
+  9th expert scores lie within that rounding error.
+
+The exchange uses its own multicast workspace (three rotating Lamport buffers
+and a device-side protocol state), separate from every native or other-package
+workspace. One workspace serves all layers, CUDA graphs and eager calls of a
+TP group, which issue their hops in order on the model's main stream; every rank
+must issue the same sequence of hops. The native MNNVL all-reduce has the same
+contract: one workspace and one device-side buffer state per TP group
+(``MNNVLAllReduce.allreduce_mnnvl_workspaces``).
+
+Bit-identity holds against the native MNNVL one-shot kernel, which the native
+``AllReduce`` uses for these calls only when it selected MNNVL (strategy AUTO
+across nodes, or MNNVL) and the GPU has at least ``8 * rows`` SMs (cluster of
+8 CTAs per row). The hop therefore engages only for a call site whose native
+``AllReduce`` is on MNNVL and for row counts whose grid fits the smallest SM
+count of the group.
+
+Eligibility is decided from rank-uniform facts only (feature selection, the
+group's capability and SM count agreed at construction, the call site's native
+all-reduce configuration, dtype and shape of the call), so all ranks take the
+same path for every call and the collective sequence never diverges.
+Unsupported calls stay on the native path before any device work. Host-side
+counters (:func:`engagement_counts`) record engaged and native-fallback calls
+per Python call, i.e. eager runs and CUDA graph captures, not graph replays.
+"""
+
+import os
+import threading
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.logger import logger
+from tensorrt_llm.mapping import Mapping
+
+from ...dispatch import log_capability
+
+FEATURE = "post_attention_hop"
+HIDDEN_SIZE = 6144
+NUM_RANKS = 8
+NUM_EXPERTS = 256
+MAX_ROWS = 8
+NUM_LAMPORT_BUFFERS = 3
+LAMPORT_BUFFER_BYTES = MAX_ROWS * NUM_RANKS * HIDDEN_SIZE * 2
+WORKSPACE_BYTES = LAMPORT_BUFFER_BYTES * NUM_LAMPORT_BUFFERS
+# CTAs per cluster (one per 768-element slice of the hidden dimension) and the
+# router layout's expert groups (kernel: kSlices, launchRouter).
+CLUSTER_CTAS = 8
+ROUTER_EXPERT_GROUPS = 4
+ROUTER_MAX_ROW_GROUPS = 3
+_OP_NAME = "glm_moe_dsa_post_attention_hop"
+
+# Host-side call counters for dispatch-engagement reporting. They count Python
+# calls (eager runs and CUDA graph captures), not graph replays.
+_counts: Dict[str, int] = {}
+
+
+def _count(key: str) -> None:
+    _counts[key] = _counts.get(key, 0) + 1
+
+
+def engagement_counts() -> Dict[str, int]:
+    """Copy of the host-side call counters.
+
+    Keys: ``engaged`` (hop calls), ``engaged_router`` (hop calls that also
+    produced router logits), ``native_fallback`` (calls ``can_run`` sent to the
+    native path) and ``native_fallback:<reason>`` per reason.
+    """
+    return dict(_counts)
+
+
+def reset_engagement_counts() -> None:
+    """Clear the host-side call counters."""
+    _counts.clear()
+
+
+def required_sms(rows: int, router: bool) -> int:
+    """SMs one hop of ``rows`` rows needs.
+
+    The native one-shot kernel uses one 8-CTA cluster per row only while
+    ``8 * rows`` fits the SM count; the hop reproduces that layout. With router
+    logits every CTA of the grid (``4 * min(rows, 3)`` clusters of 8 CTAs, one
+    CTA per SM) must be co-resident, because CTAs spin on data that other CTAs
+    of every rank publish.
+    """
+    need = CLUSTER_CTAS * rows
+    if router:
+        need = max(need, CLUSTER_CTAS * ROUTER_EXPERT_GROUPS * min(rows, ROUTER_MAX_ROW_GROUPS))
+    return need
+
+
+def native_allreduce_ineligibility(
+    native_allreduce: Optional[Any], mapping: Mapping
+) -> Optional[str]:
+    """Reason the call site's native all-reduce is not the MNNVL path, or None.
+
+    Args:
+        native_allreduce: the call site's native ``AllReduce`` module (the one
+            the hop replaces), or None when the call site does not pass it.
+        mapping: the TP group's mapping.
+
+    With ``native_allreduce`` the check is exact: the native fused
+    ``RESIDUAL_RMS_NORM`` call runs the MNNVL one-shot kernel iff the module
+    created its ``MNNVLAllReduce`` (strategy AUTO across nodes, or MNNVL).
+    Without it, the hop requires that some native module of this TP group
+    created the MNNVL workspace, which no strategy other than AUTO/MNNVL does.
+    The answer depends only on configuration that is identical on all ranks.
+    """
+    from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce
+
+    if native_allreduce is not None:
+        if getattr(native_allreduce, "mnnvl_allreduce", None) is None:
+            strategy = getattr(native_allreduce, "strategy", None)
+            name = getattr(strategy, "name", str(strategy))
+            return f"native all-reduce ({name}) does not use MNNVL"
+        return None
+    if mapping not in MNNVLAllReduce.allreduce_mnnvl_workspaces:
+        return "no native MNNVL all-reduce on this TP group"
+    return None
+
+
+def op_available() -> bool:
+    """Whether the compiled ``trtllm::glm_moe_dsa_post_attention_hop`` op is loaded."""
+    return hasattr(torch.ops.trtllm, _OP_NAME)
+
+
+def _register_fake() -> None:
+    if not op_available():
+        return
+    try:
+
+        @torch.library.register_fake(f"trtllm::{_OP_NAME}")
+        def _(input, residual, gamma, eps, router_weight, uc_buffer, mc_ptr, flags, rank, nranks):
+            rows = input.numel() // input.shape[-1]
+            logits = (
+                input.new_empty((rows, NUM_EXPERTS), dtype=torch.float32)
+                if router_weight is not None
+                else input.new_empty((0,), dtype=torch.float32)
+            )
+            return [torch.empty_like(input), torch.empty_like(residual), logits]
+    except RuntimeError as error:
+        # Already registered by an earlier import in this process.
+        logger.debug(f"[glm_moe_dsa] post_attention_hop fake registration skipped: {error}")
+
+
+def static_ineligibility(mapping: Mapping, hidden_size: int, dtype: torch.dtype) -> Optional[str]:
+    """Reason the configuration can never use post_attention_hop, or None.
+
+    Every input is identical on all ranks of the TP group, so the answer is
+    rank-uniform without communication.
+    """
+    from tensorrt_llm._utils import mpi_disabled
+
+    if mapping.tp_size != NUM_RANKS:
+        return f"tp_size {mapping.tp_size} != {NUM_RANKS}"
+    if mapping.enable_attention_dp:
+        return "attention DP has no post-attention all-reduce"
+    if mapping.has_cp():
+        return "context parallelism is not supported"
+    if hidden_size != HIDDEN_SIZE:
+        return f"hidden size {hidden_size} != {HIDDEN_SIZE}"
+    if dtype != torch.bfloat16:
+        return f"dtype {dtype} is not bfloat16"
+    if mpi_disabled():
+        return "workspace setup needs MPI"
+    return None
+
+
+class PostAttnHop:
+    """Workspace and entry point of the fused post-attention hop for one TP group."""
+
+    _instances: Dict[Mapping, Optional["PostAttnHop"]] = {}
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_or_create(
+        cls,
+        mapping: Mapping,
+        hidden_size: int = HIDDEN_SIZE,
+        dtype: torch.dtype = torch.bfloat16,
+        native_allreduce: Optional[Any] = None,
+    ) -> Optional["PostAttnHop"]:
+        """Return the group's hop, creating its workspace collectively on first use.
+
+        Must be called at the same point of model construction on every rank of
+        the TP group.
+
+        Args:
+            mapping: the TP group's mapping.
+            hidden_size: model hidden size.
+            dtype: activation dtype.
+            native_allreduce: the call site's native ``AllReduce`` module whose
+                fused ``RESIDUAL_RMS_NORM`` call the hop replaces; the hop
+                serves this call site only if that module uses MNNVL.
+
+        Returns:
+            The hop, or None (native path) when the call site or the
+            configuration is ineligible or any rank failed to set up the
+            workspace.
+        """
+        with cls._lock:
+            site_reason = native_allreduce_ineligibility(native_allreduce, mapping)
+            if site_reason is not None:
+                log_capability(FEATURE, f"post_attn_hop:{site_reason}", False, site_reason)
+                return None
+            if mapping in cls._instances:
+                return cls._instances[mapping]
+            reason = static_ineligibility(mapping, hidden_size, dtype)
+            hop = None
+            if reason is None:
+                hop = cls(mapping)
+                if not hop.enabled:
+                    reason = hop.disabled_reason
+                    hop = None
+            log_capability(FEATURE, "post_attn_hop", hop is not None, reason or "")
+            cls._instances[mapping] = hop
+            return hop
+
+    def __init__(self, mapping: Mapping):
+        from tensorrt_llm._torch.distributed.mnnvl_memory import MnnvlMemory
+        from tensorrt_llm._torch.distributed.ops import (
+            _get_mnnvl_workspace_comm,
+            _make_mnnvl_mcast_buffer,
+            _mnnvl_device_index,
+        )
+        from tensorrt_llm._utils import mpi_comm
+
+        self.mapping = mapping
+        self.rank = mapping.tp_rank
+        self.enabled = False
+        self.disabled_reason = ""
+        self.handle = None
+        self.uc_buffer = None
+        self.mc_ptr = 0
+        self.flags = None
+        self.comm = None
+        self.min_sm_count = 0
+
+        # Local capability first, agreed over the world communicator so that
+        # either every rank creates the TP communicator and workspace or none.
+        local_reason = ""
+        if not op_available():
+            local_reason = "trtllm::glm_moe_dsa_post_attention_hop is not loaded"
+        elif not MnnvlMemory.supports_mnnvl():
+            local_reason = "MNNVL is not supported on this device"
+        world = mpi_comm()
+        if world.allreduce(0 if local_reason else 1) != world.Get_size():
+            self.disabled_reason = local_reason or "another rank is ineligible"
+            return
+        # The smallest SM count of all ranks bounds the row counts every rank
+        # accepts (see required_sms), so per-call decisions stay rank-uniform.
+        from mpi4py import MPI
+
+        self.min_sm_count = world.allreduce(
+            torch.cuda.get_device_properties(_mnnvl_device_index(mapping)).multi_processor_count,
+            op=MPI.MIN,
+        )
+
+        comm = _get_mnnvl_workspace_comm(mapping)
+        error: Optional[Exception] = None
+        try:
+            force_mn = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1"
+            self.handle = _make_mnnvl_mcast_buffer(
+                comm,
+                WORKSPACE_BYTES,
+                mapping,
+                use_fabric_handle=force_mn or mapping.is_multi_node(),
+            )
+            self.uc_buffer = self.handle.get_uc_buffer(
+                mapping.tp_rank, (WORKSPACE_BYTES // 4,), torch.float32, 0
+            )
+            self.mc_ptr = self.handle.get_mc_buffer((1,), torch.float32, 0).data_ptr()
+            # Built like the native MNNVL flags: torch.tensor is not redirected
+            # to the meta device under MetaInitMode during model construction.
+            self.flags = torch.tensor(
+                [0, 0, 0, 0],
+                dtype=torch.int32,
+                device=torch.device("cuda", _mnnvl_device_index(mapping)),
+            )
+            self.reset_protocol()
+        except (RuntimeError, ValueError) as err:
+            error = err
+        ok = comm.allreduce(0 if error is not None else 1) == comm.Get_size()
+        if not ok:
+            self.disabled_reason = (
+                f"workspace setup failed: {error}"
+                if error is not None
+                else "workspace setup failed on another rank"
+            )
+            self.handle = None
+            self.uc_buffer = None
+            self.flags = None
+            return
+        self.comm = comm
+        self.enabled = True
+
+    def reset_protocol(self) -> None:
+        """Fill every Lamport buffer with the sentinel and zero the protocol state.
+
+        Collective in effect: every rank must reset before any rank issues the
+        next hop, which the caller guarantees with a barrier.
+        """
+        with torch.inference_mode():
+            self.uc_buffer.fill_(-0.0)
+            self.flags.zero_()
+        torch.cuda.synchronize()
+
+    def call_ineligibility(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        norm_weight: torch.Tensor,
+        router_weight: Optional[torch.Tensor] = None,
+    ) -> Optional[str]:
+        """Reason this call must take the native path, or None (no device work)."""
+        if not self.enabled:
+            return "hop disabled"
+        if residual is None:
+            return "no residual"
+        if (
+            hidden_states.dtype != torch.bfloat16
+            or residual.dtype != torch.bfloat16
+            or norm_weight.dtype != torch.bfloat16
+        ):
+            return "dtype"
+        if (
+            hidden_states.dim() < 2
+            or hidden_states.shape[-1] != HIDDEN_SIZE
+            or residual.shape != hidden_states.shape
+            or norm_weight.shape != (HIDDEN_SIZE,)
+        ):
+            return "shape"
+        rows = hidden_states.numel() // HIDDEN_SIZE
+        if not 1 <= rows <= MAX_ROWS:
+            return f"rows outside [1, {MAX_ROWS}]"
+        if router_weight is not None and (
+            router_weight.dtype != torch.bfloat16
+            or router_weight.shape != (NUM_EXPERTS, HIDDEN_SIZE)
+            or not router_weight.is_contiguous()
+        ):
+            return "router weight"
+        if required_sms(rows, router_weight is not None) > self.min_sm_count:
+            return "SM count"
+        return None
+
+    def can_run(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        norm_weight: torch.Tensor,
+        router_weight: Optional[torch.Tensor] = None,
+    ) -> bool:
+        """Whether this call can use the fused hop (rank-uniform, no device work).
+
+        A False answer is counted as a native fallback (engagement_counts).
+        """
+        reason = self.call_ineligibility(hidden_states, residual, norm_weight, router_weight)
+        if reason is not None:
+            _count("native_fallback")
+            _count(f"native_fallback:{reason}")
+            return False
+        return True
+
+    def __call__(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+        router_weight: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Run the hop; returns ``(normed, residual_out, router_logits or None)``.
+
+        Callers check :meth:`can_run` first.
+        """
+        _count("engaged")
+        if router_weight is not None:
+            _count("engaged_router")
+        normed, residual_out, logits = getattr(torch.ops.trtllm, _OP_NAME)(
+            hidden_states.contiguous(),
+            residual.contiguous(),
+            norm_weight.contiguous(),
+            float(eps),
+            router_weight,
+            self.uc_buffer,
+            self.mc_ptr,
+            self.flags,
+            self.rank,
+            NUM_RANKS,
+        )
+        return normed, residual_out, (logits if router_weight is not None else None)
+
+
+_register_fake()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/__init__.py
new file mode 100644
index 000000000000..cd9c76712ab5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/__init__.py
@@ -0,0 +1,45 @@
+# 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.
+"""post_moe_hop: multi-node post-MoE finalize/reduction/norm (the post_moe_hop contract).
+
+Registered entry points (``dispatch.get_impl("post_moe_hop", name)``):
+
+* ``post_moe_hop``: :class:`PostMoeHop` factory, built once per target MoE
+  layer at model construction (collective on first use).
+* ``health_check``: :func:`check_health`, the host-side report of the device
+  error words and the engagement counter; call it outside timed regions.
+
+Two kernels implement the hop with the same workspace protocol and
+arithmetic: a CuTe DSL kernel with one thread block cluster per token row
+(``cute_kernel.py``, the default) and a Triton kernel (``kernels.py``) that
+serves every shape the cluster kernel does not.
+"""
+
+from ... import dispatch
+from .health import PostMoeHopError, PostMoeHopHealth, check_health
+from .hop import PostMoeHop, PostMoeHopContractError, native_finalize_reference, producer_capability
+
+dispatch.register("post_moe_hop", "post_moe_hop", PostMoeHop)
+dispatch.register("post_moe_hop", "health_check", check_health)
+
+__all__ = [
+    "PostMoeHopContractError",
+    "PostMoeHopHealth",
+    "PostMoeHopError",
+    "PostMoeHop",
+    "check_health",
+    "native_finalize_reference",
+    "producer_capability",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/cute_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/cute_kernel.py
new file mode 100644
index 000000000000..4d69c3a9bdfb
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/cute_kernel.py
@@ -0,0 +1,1078 @@
+# 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.
+"""CuTe DSL kernel of the post_moe_hop (the post_moe_hop contract), cluster variant.
+
+Same protocol and arithmetic as the Triton kernel in ``kernels.py`` and the
+same workspace, but one thread block cluster per token row: ``CLUSTER``
+CTAs of ``THREADS`` threads, every thread owning 8 consecutive hidden
+columns (one 16-byte access per tensor). Per thread:
+
+1. fp32 FMA finalize of the raw tuple in ``k`` order, bf16 rounding, bf16
+   add of the shared-expert partial, ``-0.0`` sanitize;
+2. 16-byte store through the multicast address into slot ``[row][rank]``;
+3. share of the clear of the previous call's Lamport rows (and partial
+   slice, which only the Triton kernel uses but both keep clean);
+4. per-thread Lamport poll of the 8 rank slots (one inline-PTX spin loop,
+   no block-wide reduction per iteration, ``%globaltimer`` deadline checked
+   every 1,024 polls), rank-ordered fp32 sum, bf16;
+5. bf16 residual add and store;
+6. RMSNorm: per-thread sum of bf16-rounded squares, warp butterfly, the
+   per-warp sums are written into every CTA of the cluster through
+   distributed shared memory, one cluster barrier, then every thread sums
+   the same ``CLUSTER * WARPS`` values in the same order.
+
+All global accesses go through inline PTX on raw addresses so the memory
+operations (volatile polls, multicast stores, acquire/release atomics) are
+explicit. Every tensor is therefore accessed with 16-byte vectors: the caller
+(``hop.py``) guarantees 16-byte aligned bases and row pitches.
+
+Error handling is the same as in the Triton kernel: an expired deadline sets
+error bit 1 and the largest poll count, and NaN replaces this thread's
+residual (hence the row's normalized output); a permutation entry at or beyond
+``fc2_rows`` is skipped and sets error bit 4. Error bits go into the error
+word of the call's Lamport buffer; a call reads the previous buffer's word
+(complete at the kernel boundary) and, when it is nonzero, publishes NaN and
+carries the bits into its own word, so every later call of this rank
+publishes NaN until the collective workspace reset (see ``kernels.py``).
+"""
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.utils.smem_allocator import SmemAllocator
+
+from .kernels import MAX_CHUNKS
+
+_SENTINEL_WORD = -2147450880  # 0x80008000: two bf16 -0.0
+_FP32_SENTINEL = -2147483648  # 0x80000000
+# Polls between two deadline checks; the clock starts at the first check.
+_POLLS_PER_CLOCK_CHECK = 1024
+
+
+def _asm(res, operands, asm, constraints, loc, ip):
+    return llvm.inline_asm(
+        res,
+        operands,
+        asm,
+        constraints,
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+def _v(x, loc, ip):
+    return x.ir_value(loc=loc, ip=ip)
+
+
+def _struct(n, ty):
+    return llvm.StructType.get_literal([ty] * n)
+
+
+def _unpack(res, n, cls, ty, loc, ip):
+    return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n))
+
+
+@dsl_user_op
+def _ld_v4(addr: Int64, *, loc=None, ip=None):
+    res = _asm(
+        _struct(4, T.i32()),
+        [_v(addr, loc, ip)],
+        "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+        "=r,=r,=r,=r,l",
+        loc,
+        ip,
+    )
+    return _unpack(res, 4, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def _st_v4(addr: Int64, a: Int32, b: Int32, c: Int32, d: Int32, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [_v(addr, loc, ip)] + [_v(x, loc, ip) for x in (a, b, c, d)],
+        "st.global.v4.b32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        loc,
+        ip,
+    )
+
+
+@dsl_user_op
+def _poll8(base: Int64, stride: Int64, timeout_ns: Int64, *, loc=None, ip=None):
+    """Spin until 8 rank slots (16 bytes each, ``stride`` apart) hold no sentinel.
+
+    The ``%globaltimer`` deadline is checked every ``_POLLS_PER_CLOCK_CHECK``
+    polls and its clock starts at the first check, so a poll that finds every
+    slot within that many iterations never reads the timer. Returns the 32
+    words, a nonzero flag when the deadline expired, and the poll count.
+    """
+    loads = []
+    checks = []
+    for r in range(8):
+        o = 4 * r
+        loads.append(
+            f"ld.volatile.global.v4.b32 {{${o}, ${o + 1}, ${o + 2}, ${o + 3}}}, [a];\n"
+            "add.s64 a, a, $35;\n"
+        )
+        for w in range(4):
+            checks.append(
+                f"and.b32 t, ${o + w}, 0xFFFF0000;\n"
+                "setp.eq.u32 p, t, 0x80000000;\n"
+                "selp.b32 t, 1, 0, p;\n"
+                "or.b32 d, d, t;\n"
+            )
+    period = _POLLS_PER_CLOCK_CHECK
+    asm = (
+        "{\n.reg .pred p;\n.reg .b32 t, d, n;\n.reg .b64 a, t0, now;\n"
+        "mov.b32 n, 0;\nmov.b64 t0, 0;\n"
+        "POST_MOE_POLL:\nmov.b64 a, $34;\nmov.b32 d, 0;\n"
+        + "".join(loads)
+        + "".join(checks)
+        + "add.u32 n, n, 1;\nsetp.eq.u32 p, d, 0;\n@p bra POST_MOE_POLL_DONE;\n"
+        f"and.b32 t, n, {period - 1};\nsetp.ne.u32 p, t, 0;\n@p bra POST_MOE_POLL;\n"
+        "mov.u64 now, %globaltimer;\n"
+        f"setp.eq.u32 p, n, {period};\n@p mov.b64 t0, now;\n"
+        "sub.s64 now, now, t0;\nsetp.lt.s64 p, now, $36;\n@p bra POST_MOE_POLL;\n"
+        "POST_MOE_POLL_DONE:\nmov.b32 $32, d;\nmov.b32 $33, n;\n}\n"
+    )
+    constraints = ",".join(["=r"] * 34 + ["l", "l", "l"])
+    res = _asm(
+        _struct(34, T.i32()),
+        [_v(base, loc, ip), _v(stride, loc, ip), _v(timeout_ns, loc, ip)],
+        asm,
+        constraints,
+        loc,
+        ip,
+    )
+    return _unpack(res, 34, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def _unpack_bf16x2(word: Int32, *, loc=None, ip=None):
+    res = _asm(
+        _struct(2, T.f32()),
+        [_v(word, loc, ip)],
+        "{\n.reg .b32 t;\nshl.b32 t, $2, 16;\nmov.b32 $0, t;\n"
+        "and.b32 t, $2, 0xFFFF0000;\nmov.b32 $1, t;\n}\n",
+        "=f,=f,r",
+        loc,
+        ip,
+    )
+    return _unpack(res, 2, Float32, T.f32(), loc, ip)
+
+
+@dsl_user_op
+def _pack_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(lo, loc, ip), _v(hi, loc, ip)],
+            "cvt.rn.bf16x2.f32 $0, $2, $1;",
+            "=r,f,f",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _round_bf16(x: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(x, loc, ip)],
+            "{\n.reg .b32 t;\ncvt.rn.bf16x2.f32 t, $1, $1;\n"
+            "and.b32 t, t, 0xFFFF0000;\nmov.b32 $0, t;\n}\n",
+            "=f,f",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _sanitize(word: Int32, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(word, loc, ip)],
+            "{\n.reg .b32 t;\n.reg .pred p, q;\n"
+            "and.b32 t, $1, 0xFFFF;\nsetp.eq.u32 p, t, 0x8000;\n"
+            "and.b32 t, $1, 0xFFFF0000;\nsetp.eq.u32 q, t, 0x80000000;\n"
+            "mov.b32 $0, $1;\n@p and.b32 $0, $0, 0xFFFF0000;\n"
+            "@q and.b32 $0, $0, 0xFFFF;\n}\n",
+            "=r,r",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _fma(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(a, loc, ip), _v(b, loc, ip), _v(c, loc, ip)],
+            "fma.rn.f32 $0, $1, $2, $3;",
+            "=f,f,f,f",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _add(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(T.f32(), [_v(a, loc, ip), _v(b, loc, ip)], "add.rn.f32 $0, $1, $2;", "=f,f,f", loc, ip)
+    )
+
+
+@dsl_user_op
+def _mul(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(T.f32(), [_v(a, loc, ip), _v(b, loc, ip)], "mul.rn.f32 $0, $1, $2;", "=f,f,f", loc, ip)
+    )
+
+
+@dsl_user_op
+def _select(pred: Int32, a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(pred, loc, ip), _v(a, loc, ip), _v(b, loc, ip)],
+            "{\n.reg .pred p;\nsetp.ne.s32 p, $1, 0;\nselp.f32 $0, $2, $3, p;\n}\n",
+            "=f,r,f,f",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _perm_check(index: Int32, rows: Int32, *, loc=None, ip=None):
+    """Return ``(valid, safe_index, out_of_range)`` of one permutation entry.
+
+    ``valid`` is ``0 <= index < rows``; ``safe_index`` is ``index`` when valid
+    and 0 otherwise; ``out_of_range`` flags ``index >= rows`` (negative
+    entries are the native "no local row" marker, not an error).
+    """
+    res = _asm(
+        _struct(3, T.i32()),
+        [_v(index, loc, ip), _v(rows, loc, ip)],
+        "{\n.reg .pred p, q;\nsetp.lt.u32 p, $3, $4;\nselp.b32 $0, 1, 0, p;\n"
+        "selp.b32 $1, $3, 0, p;\nsetp.ge.s32 q, $3, $4;\nselp.b32 $2, 1, 0, q;\n}\n",
+        "=r,=r,=r,r,r",
+        loc,
+        ip,
+    )
+    return _unpack(res, 3, Int32, T.i32(), loc, ip)
+
+
+@dsl_user_op
+def _poison(word: Int32, flag: Int32, *, loc=None, ip=None) -> Int32:
+    """Two bf16 NaNs instead of ``word`` when ``flag`` is nonzero."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(word, loc, ip), _v(flag, loc, ip)],
+            "{\n.reg .pred p;\nsetp.ne.s32 p, $2, 0;\nselp.b32 $0, 0x7FC07FC0, $1, p;\n}\n",
+            "=r,r,r",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _bits_f32(word: Int32, *, loc=None, ip=None) -> Float32:
+    return Float32(_asm(T.f32(), [_v(word, loc, ip)], "mov.b32 $0, $1;", "=f,r", loc, ip))
+
+
+@dsl_user_op
+def _inv_rms(square_sum: Float32, hidden: Float32, eps: Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(square_sum, loc, ip), _v(hidden, loc, ip), _v(eps, loc, ip)],
+            "{\n.reg .f32 t;\ndiv.rn.f32 t, $1, $2;\nadd.rn.f32 t, t, $3;\n"
+            "rsqrt.approx.f32 $0, t;\n}\n",
+            "=f,f,f,f",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _shfl_xor(x: Float32, mask: int, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [_v(x, loc, ip)],
+            f"shfl.sync.bfly.b32 $0, $1, {mask}, 0x1f, 0xffffffff;",
+            "=f,f",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _st_cluster_f32(
+    local_addr: Int32, target_cta: Int32, value: Float32, enable: Int32, *, loc=None, ip=None
+) -> None:
+    _asm(
+        None,
+        [_v(local_addr, loc, ip), _v(target_cta, loc, ip), _v(value, loc, ip), _v(enable, loc, ip)],
+        "{\n.reg .pred p;\n.reg .b32 r;\nsetp.ne.s32 p, $3, 0;\n"
+        "@p mapa.shared::cluster.u32 r, $0, $1;\n@p st.shared::cluster.f32 [r], $2;\n}\n",
+        "r,r,f,r",
+        loc,
+        ip,
+    )
+
+
+@dsl_user_op
+def _ld_shared_f32(addr: Int32, *, loc=None, ip=None) -> Float32:
+    return Float32(_asm(T.f32(), [_v(addr, loc, ip)], "ld.shared.f32 $0, [$1];", "=f,r", loc, ip))
+
+
+@dsl_user_op
+def _error_word(words: list, index: Int32, *, loc=None, ip=None) -> Int32:
+    """Select error word ``index`` (0..2) of the three loaded state words."""
+    return Int32(
+        _asm(
+            T.i32(),
+            [_v(index, loc, ip)] + [_v(w, loc, ip) for w in words],
+            "{\n.reg .pred p;\nsetp.eq.u32 p, $1, 1;\nselp.b32 $0, $3, $2, p;\n"
+            "setp.eq.u32 p, $1, 2;\nselp.b32 $0, $4, $0, p;\n}\n",
+            "=r,r,r,r,r",
+            loc,
+            ip,
+        )
+    )
+
+
+@dsl_user_op
+def _flag_error(word: Int64, flag: Int32, bits: Int32, *, loc=None, ip=None) -> None:
+    """OR ``bits`` into the error word at ``word`` when ``flag`` is nonzero."""
+    _asm(
+        None,
+        [_v(word, loc, ip), _v(flag, loc, ip), _v(bits, loc, ip)],
+        "{\n.reg .pred p;\nsetp.ne.s32 p, $1, 0;\n@p red.relaxed.gpu.global.or.b32 [$0], $2;\n}\n",
+        "l,r,r",
+        loc,
+        ip,
+    )
+
+
+@dsl_user_op
+def _flag_timeout(
+    word: Int64, state: Int64, flag: Int32, polls: Int32, *, loc=None, ip=None
+) -> None:
+    """Error bit 1 into the error word and the largest poll count at expiry (word 5)."""
+    _asm(
+        None,
+        [_v(word, loc, ip), _v(state, loc, ip), _v(flag, loc, ip), _v(polls, loc, ip)],
+        "{\n.reg .pred p;\nsetp.ne.s32 p, $2, 0;\n"
+        "@p red.relaxed.gpu.global.or.b32 [$0], 1;\n"
+        "@p red.relaxed.gpu.global.max.u32 [$1+20], $3;\n}\n",
+        "l,l,r,r",
+        loc,
+        ip,
+    )
+
+
+@dsl_user_op
+def _fault_delay(
+    a: Int32, b: Int32, c: Int32, max_ns: int, site: int, *, loc=None, ip=None
+) -> None:
+    """Test-only fault injection: about one thread in four spins before ``site``.
+
+    The delay (below ``max_ns``) is a hash of ``a``, ``b`` and ``c``, so the
+    individual 16-byte stores of one call land in a scrambled order.
+    """
+    asm = (
+        "{\n.reg .pred p;\n.reg .b32 h, t;\n.reg .b64 t0, now, d;\n"
+        "mul.lo.u32 h, $0, 0x9E3779B1;\nmul.lo.u32 t, $1, 0x85EBCA77;\nxor.b32 h, h, t;\n"
+        "mul.lo.u32 t, $2, 0xC2B2AE3D;\nxor.b32 h, h, t;\nshr.u32 t, h, 15;\n"
+        "xor.b32 h, h, t;\nmul.lo.u32 h, h, 0x2C1B3C6D;\nshr.u32 t, h, 13;\n"
+        "xor.b32 h, h, t;\nand.b32 t, h, 3;\nsetp.ne.u32 p, t, 0;\n"
+        f"@p bra POST_MOE_FAULT_END_{site};\nshr.u32 t, h, 2;\nrem.u32 t, t, {max_ns};\n"
+        "cvt.u64.u32 d, t;\nmov.u64 t0, %globaltimer;\n"
+        f"POST_MOE_FAULT_LOOP_{site}:\nmov.u64 now, %globaltimer;\nsub.s64 now, now, t0;\n"
+        f"setp.lt.s64 p, now, d;\n@p bra POST_MOE_FAULT_LOOP_{site};\n"
+        f"POST_MOE_FAULT_END_{site}:\n}}\n"
+    )
+    _asm(None, [_v(a, loc, ip), _v(b, loc, ip), _v(c, loc, ip)], asm, "r,r,r", loc, ip)
+
+
+@dsl_user_op
+def _arrive(
+    state: Int64,
+    is_leader: Int32,
+    num_ctas: Int32,
+    num_rows: Int32,
+    next_epoch: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    """Leader thread of each CTA counts in; the last CTA publishes the next epoch."""
+    _asm(
+        None,
+        [
+            _v(state, loc, ip),
+            _v(is_leader, loc, ip),
+            _v(num_ctas, loc, ip),
+            _v(num_rows, loc, ip),
+            _v(next_epoch, loc, ip),
+        ],
+        "{\n.reg .pred p, q;\n.reg .b32 old, c;\n"
+        "setp.ne.s32 p, $1, 0;\nmov.b32 old, 0;\n"
+        "@p atom.acq_rel.gpu.global.add.u32 old, [$0+8], 1;\n"
+        "add.u32 c, old, 1;\nsetp.eq.u32 q, c, $2;\nand.pred q, q, p;\n"
+        "@q st.global.u32 [$0+4], $3;\n@q st.global.u32 [$0+8], 0;\n"
+        "@q red.relaxed.gpu.global.add.u32 [$0+16], 1;\n"
+        "@q st.global.u32 [$0], $4;\n}\n",
+        "l,r,r,r,r",
+        loc,
+        ip,
+    )
+
+
+class PostMoeHopClusterKernel:
+    """Compile-time configuration of the cluster kernel for one hop shape.
+
+    Test-only knobs (the hop never sets them):
+
+    * ``race_probe``: fault injection for the compute-sanitizer positive
+      control; the reduction buffer is read between the cluster arrive and
+      wait, so the distributed shared memory stores of the other CTAs race
+      with the reads.
+    * ``fault_delay_ns``: delays about one thread in four by a hashed amount
+      below this bound before its multicast publish, before its share of the
+      clear and before the arrival, so the individual stores of a call land
+      in a scrambled order (protocol fault injection).
+
+    routed_moe_into_hop entry (``routed_input``, reachable only through
+    ``PostMoeHop.enable_routed_input``; the default ``False`` is the post_moe_hop
+    on the raw tuple): only how the rank-local routed partial is formed
+    changes, everything after it is the same code.
+
+    * ``"row"`` (or ``True``): ``fc2`` is the rank's finalized routed output
+      ``[R, H]`` bf16 (no scales, no permutation);
+    * ``"identity"``: raw FC2 rows in the identity layout (row ``t * 8 + k``)
+      with bf16 scales ``[R, 8]``; the native finalize without the
+      permutation round trip;
+    * ``"fp32"``: ``fc2`` is an FP32 ``[R, H]`` accumulator of the top-k
+      contributions, rounded to bf16 here and cleared after the read.
+    """
+
+    def __init__(
+        self,
+        hidden: int,
+        top_k: int,
+        world: int,
+        max_rows: int,
+        cluster: int,
+        has_shared: bool,
+        scale_fp32: bool,
+        use_pdl: bool,
+        race_probe: bool = False,
+        fault_delay_ns: int = 0,
+        routed_input=False,
+    ):
+        assert top_k == 8 and world == 8, "the cluster kernel is specialized to top-8 and TP8"
+        assert hidden % (8 * 32 * cluster) == 0
+        self.hidden = hidden
+        self.top_k = top_k
+        self.world = world
+        self.max_rows = max_rows
+        self.cluster = cluster
+        self.threads = hidden // (8 * cluster)
+        self.warps = self.threads // 32
+        self.has_shared = has_shared
+        self.scale_fp32 = scale_fp32
+        self.use_pdl = use_pdl
+        self.race_probe = race_probe
+        self.fault_delay_ns = fault_delay_ns
+        self.routed_input = "row" if routed_input is True else routed_input
+
+    @cute.jit
+    def __call__(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        out: Int64,
+        residual_out: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        state: cute.Tensor,
+        rank: Int32,
+        eps: Float32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+        stream: CUstream,
+    ):
+        # ``state`` is passed as a GPU tensor so the TVM-FFI launcher can bind
+        # the caller's current stream; every other buffer is a raw address.
+        self.kernel(
+            fc2,
+            fc2_row_bytes,
+            fc2_rows,
+            scales,
+            perm,
+            shared,
+            residual,
+            gamma,
+            out,
+            residual_out,
+            mc,
+            uc,
+            partials,
+            state,
+            rank,
+            eps,
+            timeout_ns,
+            num_rows,
+        ).launch(
+            grid=[self.cluster, num_rows, 1],
+            block=[self.threads, 1, 1],
+            cluster=[self.cluster, 1, 1],
+            smem=4 * self.cluster * self.warps + 64,
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        fc2: Int64,
+        fc2_row_bytes: Int64,
+        fc2_rows: Int32,
+        scales: Int64,
+        perm: Int64,
+        shared: Int64,
+        residual: Int64,
+        gamma: Int64,
+        out: Int64,
+        residual_out: Int64,
+        mc: Int64,
+        uc: Int64,
+        partials: Int64,
+        state_tensor: cute.Tensor,
+        rank: Int32,
+        eps: Float32,
+        timeout_ns: Int64,
+        num_rows: Int32,
+    ):
+        hidden = self.hidden
+        world = self.world
+        threads = self.threads
+        warps = self.warps
+        cluster = self.cluster
+        buffer_bytes = self.max_rows * world * hidden * 2
+
+        tidx, _, _ = cute.arch.thread_idx()
+        cta, row, _ = cute.arch.block_idx()
+        lane = cute.arch.lane_idx()
+        warp = cute.arch.warp_idx()
+
+        state = Int64(state_tensor.iterator.toint())
+
+        smem = SmemAllocator()
+        reduce_buf = smem.allocate_array(Float32, num_elems=cluster * warps)
+        reduce_addr = Int32(reduce_buf.toint())
+
+        cute.arch.cluster_arrive_relaxed()
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        # Words 0-7: epoch, prev_rows, arrivals, error[0], calls, timeout
+        # polls, error[1], error[2]; error[b] collects the bits of the calls
+        # that use Lamport buffer b.
+        epoch, prev_rows, _, err0 = _ld_v4(state)
+        _, _, err1, err2 = _ld_v4(state + Int64(16))
+        cur = epoch % Int32(3)
+        dirty = (epoch + Int32(2)) % Int32(3)
+        # Bits of the previous call (complete at the kernel boundary): when set,
+        # this call publishes NaN and carries them into its own error word.
+        sticky = _error_word([err0, err1, err2], dirty)
+        err_cur = state + Int64(_error_word([Int32(12), Int32(24), Int32(28)], cur))
+
+        col = (cta * Int32(threads) + tidx) * Int32(8)
+        row_off = Int64(row) * Int64(hidden * 2) + Int64(col) * Int64(2)
+        res_w = _ld_v4(residual + row_off)
+        gamma_w = _ld_v4(gamma + Int64(col) * Int64(2))
+        if cutlass.const_expr(self.has_shared):
+            sh_w = _ld_v4(shared + row_off)
+
+        # 1. Finalize (fp32 FMA chain in k order), bf16, + shared, sanitize.
+        if cutlass.const_expr(self.routed_input == "row"):
+            # routed_moe_into_hop: the routed kernel already finalized this rank's output
+            r_w = _ld_v4(fc2 + Int64(row) * fc2_row_bytes + Int64(col) * Int64(2))
+            local = [r_w[j] for j in range(4)]
+            out_of_range = Int32(0)
+        elif cutlass.const_expr(self.routed_input == "fp32"):
+            # routed_moe_into_hop: FP32 accumulator of the top-k contributions; round, then
+            # clear it for the routed kernel of the next call
+            a_addr = fc2 + Int64(row) * fc2_row_bytes + Int64(col) * Int64(4)
+            a_w = _ld_v4(a_addr) + _ld_v4(a_addr + Int64(16))
+            a_f = [_bits_f32(w) for w in a_w]
+            local = [_pack_bf16x2(a_f[2 * j], a_f[2 * j + 1]) for j in range(4)]
+            zero = Int32(0)
+            _st_v4(a_addr, zero, zero, zero, zero)
+            _st_v4(a_addr + Int64(16), zero, zero, zero, zero)
+            out_of_range = Int32(0)
+        elif cutlass.const_expr(self.routed_input == "identity"):
+            # routed_moe_into_hop: raw rows t * 8 + k, every slot valid; the native finalize
+            # without the permutation round trip
+            s_w = _ld_v4(scales + Int64(row) * Int64(16))
+            weights = []
+            for w in s_w:
+                lo, hi = _unpack_bf16x2(w)
+                weights += [lo, hi]
+            acc = [Float32(0.0)] * 8
+            for k in cutlass.range_constexpr(8):
+                x_w = _ld_v4(
+                    fc2 + (Int64(row) * Int64(8) + Int64(k)) * fc2_row_bytes + Int64(col) * Int64(2)
+                )
+                for j in cutlass.range_constexpr(4):
+                    x_lo, x_hi = _unpack_bf16x2(x_w[j])
+                    acc[2 * j] = _fma(weights[k], x_lo, acc[2 * j])
+                    acc[2 * j + 1] = _fma(weights[k], x_hi, acc[2 * j + 1])
+            local = [_pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)]
+            out_of_range = Int32(0)
+        else:
+            perm_lo = _ld_v4(perm + Int64(row) * Int64(32))
+            perm_hi = _ld_v4(perm + Int64(row) * Int64(32) + Int64(16))
+            perm_k = perm_lo + perm_hi
+            if cutlass.const_expr(self.scale_fp32):
+                s_lo = _ld_v4(scales + Int64(row) * Int64(32))
+                s_hi = _ld_v4(scales + Int64(row) * Int64(32) + Int64(16))
+                weights = [_bits_f32(w) for w in s_lo + s_hi]
+            else:
+                s_w = _ld_v4(scales + Int64(row) * Int64(16))
+                weights = []
+                for w in s_w:
+                    lo, hi = _unpack_bf16x2(w)
+                    weights += [lo, hi]
+            acc = [Float32(0.0)] * 8
+            out_of_range = Int32(0)
+            for k in cutlass.range_constexpr(8):
+                valid, safe, oob = _perm_check(perm_k[k], fc2_rows)
+                out_of_range = out_of_range + oob
+                x_w = _ld_v4(fc2 + Int64(safe) * fc2_row_bytes + Int64(col) * Int64(2))
+                for j in cutlass.range_constexpr(4):
+                    x_lo, x_hi = _unpack_bf16x2(x_w[j])
+                    acc[2 * j] = _select(valid, _fma(weights[k], x_lo, acc[2 * j]), acc[2 * j])
+                    acc[2 * j + 1] = _select(
+                        valid, _fma(weights[k], x_hi, acc[2 * j + 1]), acc[2 * j + 1]
+                    )
+            local = [_pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)]
+        if cutlass.const_expr(self.has_shared):
+            summed = []
+            for j in cutlass.range_constexpr(4):
+                r_lo, r_hi = _unpack_bf16x2(local[j])
+                s_lo_f, s_hi_f = _unpack_bf16x2(sh_w[j])
+                summed.append(_pack_bf16x2(_add(s_lo_f, r_lo), _add(s_hi_f, r_hi)))
+            local = summed
+        # After an earlier error on this rank, publish NaN until the reset.
+        local = [_poison(_sanitize(w), sticky) for w in local]
+
+        # 2. Publish through multicast into slot [row][rank].
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            _fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4),
+                self.fault_delay_ns,
+                0,
+            )
+        slot = (
+            Int64(cur) * Int64(buffer_bytes)
+            + Int64(row) * Int64(world * hidden * 2)
+            + Int64(rank) * Int64(hidden * 2)
+            + Int64(col) * Int64(2)
+        )
+        _st_v4(mc + slot, local[0], local[1], local[2], local[3])
+        _flag_error(err_cur, out_of_range, Int32(4))
+        _flag_error(err_cur, sticky, sticky)
+
+        # 3. Clear the previous call's Lamport rows and partial slice.
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            _fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4) + Int32(1),
+                self.fault_delay_ns,
+                1,
+            )
+        total_threads = num_rows * Int32(cluster * threads)
+        gtid = (row * Int32(cluster) + cta) * Int32(threads) + tidx
+        clear_vecs = prev_rows * Int32(world * hidden // 8)
+        sentinel = Int32(_SENTINEL_WORD)
+        dirty_base = uc + Int64(dirty) * Int64(buffer_bytes)
+        for v in cutlass.range(gtid, clear_vecs, total_threads):
+            _st_v4(dirty_base + Int64(v) * Int64(16), sentinel, sentinel, sentinel, sentinel)
+        part_vecs = prev_rows * Int32(MAX_CHUNKS // 4)
+        part_base = partials + Int64(dirty) * Int64(self.max_rows * MAX_CHUNKS * 4)
+        fp32_sentinel = Int32(_FP32_SENTINEL)
+        for v in cutlass.range(gtid, part_vecs, total_threads):
+            _st_v4(
+                part_base + Int64(v) * Int64(16),
+                fp32_sentinel,
+                fp32_sentinel,
+                fp32_sentinel,
+                fp32_sentinel,
+            )
+
+        # 4. Per-thread Lamport poll, rank-ordered fp32 sum.
+        poll_base = (
+            uc
+            + Int64(cur) * Int64(buffer_bytes)
+            + Int64(row) * Int64(world * hidden * 2)
+            + Int64(col) * Int64(2)
+        )
+        polled = _poll8(poll_base, Int64(hidden * 2), timeout_ns)
+        timed_out = polled[32]
+        _flag_timeout(err_cur, state, timed_out, polled[33])
+        reduced = [Float32(0.0)] * 8
+        for r in cutlass.range_constexpr(8):
+            for j in cutlass.range_constexpr(4):
+                v_lo, v_hi = _unpack_bf16x2(polled[4 * r + j])
+                reduced[2 * j] = _add(reduced[2 * j], v_lo)
+                reduced[2 * j + 1] = _add(reduced[2 * j + 1], v_hi)
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        # 5. Residual added once, after the reduction; NaN if this poll timed out.
+        new_res = []
+        x = []
+        for j in cutlass.range_constexpr(4):
+            rb = _pack_bf16x2(reduced[2 * j], reduced[2 * j + 1])
+            a_lo, a_hi = _unpack_bf16x2(rb)
+            r_lo, r_hi = _unpack_bf16x2(res_w[j])
+            word = _poison(_pack_bf16x2(_add(a_lo, r_lo), _add(a_hi, r_hi)), timed_out)
+            new_res.append(word)
+            n_lo, n_hi = _unpack_bf16x2(word)
+            x += [n_lo, n_hi]
+        _st_v4(residual_out + row_off, new_res[0], new_res[1], new_res[2], new_res[3])
+
+        # 6. RMSNorm with a cluster-wide square sum.
+        thread_sum = Float32(0.0)
+        for i in cutlass.range_constexpr(8):
+            thread_sum = _add(thread_sum, _round_bf16(_mul(x[i], x[i])))
+        for mask in (16, 8, 4, 2, 1):
+            thread_sum = _add(thread_sum, _shfl_xor(thread_sum, mask))
+        cute.arch.cluster_wait()
+        my_slot = reduce_addr + (cta * Int32(warps) + warp) * Int32(4)
+        _st_cluster_f32(my_slot, lane, thread_sum, Int32(lane < Int32(cluster)))
+        cute.arch.cluster_arrive()
+        if cutlass.const_expr(not self.race_probe):
+            cute.arch.cluster_wait()
+        square_sum = Float32(0.0)
+        for i in cutlass.range_constexpr(cluster * warps):
+            square_sum = _add(square_sum, _ld_shared_f32(reduce_addr + Int32(4 * i)))
+        if cutlass.const_expr(self.race_probe):
+            cute.arch.cluster_wait()
+        inv_rms = _inv_rms(square_sum, Float32(float(hidden)), eps)
+        out_w = []
+        for j in cutlass.range_constexpr(4):
+            g_lo, g_hi = _unpack_bf16x2(gamma_w[j])
+            out_w.append(
+                _pack_bf16x2(
+                    _mul(_mul(x[2 * j], inv_rms), g_lo), _mul(_mul(x[2 * j + 1], inv_rms), g_hi)
+                )
+            )
+        _st_v4(out + row_off, out_w[0], out_w[1], out_w[2], out_w[3])
+
+        # 7. Arrival; the last CTA publishes the next epoch.
+        if cutlass.const_expr(self.fault_delay_ns > 0):
+            _fault_delay(
+                tidx + cta * Int32(1024),
+                row + rank * Int32(64),
+                epoch * Int32(4) + Int32(2),
+                self.fault_delay_ns,
+                2,
+            )
+        cute.arch.sync_threads()
+        _arrive(
+            state,
+            Int32(tidx == Int32(0)),
+            num_rows * Int32(cluster),
+            num_rows,
+            (epoch + Int32(1)) % Int32(3),
+        )
+
+
+_COMPILED = {}
+
+
+def compiled_cluster_hop(
+    hidden: int,
+    top_k: int,
+    world: int,
+    max_rows: int,
+    cluster: int,
+    has_shared: bool,
+    scale_fp32: bool,
+    use_pdl: bool,
+    state_words: int,
+    race_probe: bool = False,
+    fault_delay_ns: int = 0,
+    routed_input=False,
+):
+    """Return the compiled kernel for one shape, compiling it on first use (no GPU work).
+
+    ``routed_input`` selects a routed_moe_into_hop variant (see ``PostMoeHopClusterKernel``); the
+    default variants keep their key.
+    """
+    mode = "row" if routed_input is True else routed_input
+    key = (
+        hidden,
+        top_k,
+        world,
+        max_rows,
+        cluster,
+        has_shared,
+        scale_fp32,
+        use_pdl,
+        state_words,
+        race_probe,
+        fault_delay_ns,
+    ) + ((mode,) if mode else ())
+    compiled = _COMPILED.get(key)
+    if compiled is None:
+        from cutlass.cute import runtime as cute_runtime
+
+        kernel = PostMoeHopClusterKernel(
+            hidden,
+            top_k,
+            world,
+            max_rows,
+            cluster,
+            has_shared,
+            scale_fp32,
+            use_pdl,
+            race_probe,
+            fault_delay_ns,
+            mode,
+        )
+        zero64, zero32 = Int64(0), Int32(0)
+        state_fake = cute_runtime.make_fake_compact_tensor(
+            Int32, (state_words,), stride_order=(0,), assumed_align=16
+        )
+        compiled = cute.compile(
+            kernel,
+            zero64,
+            zero64,
+            zero32,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            zero64,
+            state_fake,
+            zero32,
+            Float32(0.0),
+            zero64,
+            zero32,
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+        _COMPILED[key] = compiled
+    return compiled
+
+
+def launch_cluster_hop(
+    *,
+    fc2_output,
+    expert_scale_factor,
+    expanded_idx_to_permuted_idx,
+    shared_output,
+    residual,
+    norm_weight,
+    out,
+    residual_out,
+    mc_buffer,
+    uc_buffer,
+    partials,
+    state,
+    rank: int,
+    world: int,
+    max_rows: int,
+    eps: float,
+    cluster: int,
+    use_pdl: bool,
+    timeout_ns: int,
+    race_probe: bool = False,
+    fault_delay_ns: int = 0,
+    fc2_rows: Optional[int] = None,
+) -> None:
+    """Launch the cluster kernel on the current stream (compiled once per shape).
+
+    Every tensor must have a 16-byte aligned base, ``fc2_output`` a 16-byte
+    row pitch, and the other tensors must be contiguous (``hop.py`` enforces
+    this). ``race_probe`` and ``fault_delay_ns`` select test-only
+    fault-injected variants (see ``PostMoeHopClusterKernel``); ``fc2_rows`` overrides
+    the permutation bound for the memcheck positive control only. The hop
+    always launches the regular kernel.
+    """
+    num_rows, hidden = residual.shape
+    top_k = expanded_idx_to_permuted_idx.shape[1]
+    has_shared = shared_output is not None
+    scale_fp32 = expert_scale_factor.dtype.itemsize == 4
+    compiled = compiled_cluster_hop(
+        hidden,
+        top_k,
+        world,
+        max_rows,
+        cluster,
+        has_shared,
+        scale_fp32,
+        use_pdl,
+        state.numel(),
+        race_probe,
+        fault_delay_ns,
+    )
+    compiled(
+        fc2_output.data_ptr(),
+        fc2_output.stride(0) * 2,
+        fc2_output.shape[0] if fc2_rows is None else fc2_rows,
+        expert_scale_factor.data_ptr(),
+        expanded_idx_to_permuted_idx.data_ptr(),
+        shared_output.data_ptr() if has_shared else 0,
+        residual.data_ptr(),
+        norm_weight.data_ptr(),
+        out.data_ptr(),
+        residual_out.data_ptr(),
+        mc_buffer.data_ptr(),
+        uc_buffer.data_ptr(),
+        partials.data_ptr(),
+        state,
+        rank,
+        eps,
+        timeout_ns,
+        num_rows,
+    )
+
+
+def launch_cluster_hop_routed(
+    *,
+    routed,
+    shared_output,
+    residual,
+    norm_weight,
+    out,
+    residual_out,
+    mc_buffer,
+    uc_buffer,
+    partials,
+    state,
+    rank: int,
+    world: int,
+    max_rows: int,
+    eps: float,
+    cluster: int,
+    use_pdl: bool,
+    timeout_ns: int,
+    top_k: int = 8,
+    mode: str = "row",
+    scales=None,
+    race_probe: bool = False,
+    fault_delay_ns: int = 0,
+) -> None:
+    """routed_moe_into_hop entry: the cluster kernel on a routed input prepared by routed_moe.
+
+    Same workspace protocol and post-finalize arithmetic as the regular hop
+    (``bf16(routed + shared)``, rank-ordered FP32 sum, residual, RMSNorm);
+    only the finalize gather of the raw tuple changes with ``mode`` (see
+    ``PostMoeHopClusterKernel``): "row" loads the finalized bf16 row, "identity"
+    gathers the identity-layout raw rows with ``scales``, "fp32" rounds the
+    FP32 accumulator and clears it. The buffers are hop-owned (16-byte
+    aligned, contiguous); ``race_probe`` and ``fault_delay_ns`` select the
+    test-only fault-injected builds, as for ``launch_cluster_hop``.
+    """
+    num_rows, hidden = residual.shape
+    has_shared = shared_output is not None
+    compiled = compiled_cluster_hop(
+        hidden,
+        top_k,
+        world,
+        max_rows,
+        cluster,
+        has_shared,
+        False,
+        use_pdl,
+        state.numel(),
+        race_probe,
+        fault_delay_ns,
+        mode,
+    )
+    compiled(
+        routed.data_ptr(),
+        routed.stride(0) * routed.element_size(),
+        routed.shape[0],
+        scales.data_ptr() if scales is not None else 0,
+        0,
+        shared_output.data_ptr() if has_shared else 0,
+        residual.data_ptr(),
+        norm_weight.data_ptr(),
+        out.data_ptr(),
+        residual_out.data_ptr(),
+        mc_buffer.data_ptr(),
+        uc_buffer.data_ptr(),
+        partials.data_ptr(),
+        state,
+        rank,
+        eps,
+        timeout_ns,
+        num_rows,
+    )
+
+
+def cluster_kernel_supported(hidden: int, top_k: int, world: int, cluster: int) -> Optional[str]:
+    """Return why the cluster kernel cannot serve a shape, or ``None``."""
+    if top_k != 8 or world != 8:
+        return "cluster kernel specialized to top-8 / TP8"
+    if cluster < 1 or cluster > 8 or hidden % (8 * 32 * cluster):
+        return f"hidden {hidden} does not split into {cluster} CTAs of whole warps"
+    if hidden // (8 * cluster) > 1024:
+        return "too many threads per CTA"
+    return None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/health.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/health.py
new file mode 100644
index 000000000000..c4e03c8b810d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/health.py
@@ -0,0 +1,119 @@
+# 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.
+"""Host-side health report of the post_moe_hop post-MoE hops (bounded error reporting).
+
+The post_moe_hop kernels never hang and never report errors through the host on their
+own: a timed-out poll or an out-of-range permutation entry sets a sticky bit
+in the device error word of the rank's post_moe_hop workspace, and every later post_moe_hop call
+of that rank publishes NaN until a collective workspace reset, which makes
+the outputs of the whole TP group NaN instead of silently wrong.
+``check_health`` reads those words and is
+meant to be called outside timed regions, e.g. once after warmup and CUDA
+graph capture and once at the end of a serving window. It costs one device
+synchronization and one 32-byte copy per workspace (one workspace serves all
+layers of a TP group).
+"""
+
+from dataclasses import dataclass
+
+from .workspace import PostMoeHopWorkspace, _all_ranks_agree
+
+ERROR_BIT_NAMES = {1: "lamport_timeout", 2: "row_sum_timeout", 4: "permutation_out_of_range"}
+
+
+class PostMoeHopError(RuntimeError):
+    """A post_moe_hop reported a device error; the run that produced it is invalid."""
+
+
+@dataclass(frozen=True)
+class PostMoeHopHealth:
+    """Counters of every post_moe_hop workspace on this rank at one check point.
+
+    Attributes:
+        where: caller-supplied label of the check point.
+        workspaces: post_moe_hop workspaces in this process (0 when no hop engaged).
+        calls: completed hop launches since the last workspace reset, summed
+            over workspaces (the dispatch engagement counter; graph replays
+            included).
+        error_bits: OR of the sticky device error words on this rank.
+        timeout_polls: largest poll count at which a poll gave up.
+        layout_repairs: host-side count of calls that needed an aligned copy.
+        group_error: some rank of the TP group reported an error (collective
+            checks only; equals ``error_bits != 0`` otherwise).
+    """
+
+    where: str
+    workspaces: int
+    calls: int
+    error_bits: int
+    timeout_polls: int
+    layout_repairs: int
+    group_error: bool
+
+    @property
+    def errors(self) -> list:
+        return [name for bit, name in ERROR_BIT_NAMES.items() if self.error_bits & bit]
+
+
+def check_health(
+    where: str, *, collective: bool = False, raise_on_error: bool = True
+) -> PostMoeHopHealth:
+    """Report the post_moe_hop counters of this rank and fail on any device error.
+
+    Args:
+        where: label for the log line, e.g. ``"after_warmup"``.
+        collective: also agree on the verdict across each workspace's TP group
+            (every rank of the group must call this at the same point); then
+            every rank fails together when any rank has an error.
+        raise_on_error: raise ``PostMoeHopError`` on an error (after logging it).
+
+    Returns:
+        The counters; logged once per call at INFO (ERROR on failure).
+
+    Raises:
+        PostMoeHopError: a device error word is nonzero (on this rank, or on any
+            rank of the group when ``collective``) and ``raise_on_error``.
+    """
+    from tensorrt_llm.logger import logger
+
+    workspaces = PostMoeHopWorkspace.instances()
+    calls = bits = polls = repairs = 0
+    group_error = False
+    for workspace in workspaces:
+        counters = workspace.counters()
+        calls += counters["calls"]
+        bits |= counters["error_bits"]
+        polls = max(polls, counters["timeout_polls"])
+        repairs += counters["layout_repairs"]
+        if collective:
+            group_error |= not _all_ranks_agree(workspace.comm, counters["error_bits"] == 0)
+    group_error |= bits != 0
+    rank = workspaces[0].mapping.tp_rank if workspaces else -1
+    health = PostMoeHopHealth(where, len(workspaces), calls, bits, polls, repairs, group_error)
+    summary = (
+        f"[glm_moe_dsa] post_moe_hop health {where}: tp_rank={rank} workspaces={len(workspaces)} "
+        f"calls={calls} error_bits={bits:#x} {health.errors} timeout_polls={polls} "
+        f"layout_repairs={repairs} group_error={group_error}"
+    )
+    if not group_error:
+        logger.info(summary)
+        return health
+    logger.error(summary.replace("post_moe_hop health", "POST_MOE_HOP ERROR"))
+    if raise_on_error:
+        raise PostMoeHopError(
+            f"{summary}; the run is invalid (a post_moe_hop workspace needs a collective "
+            "reset before further use)"
+        )
+    return health
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/hop.py
new file mode 100644
index 000000000000..b12b1f5e87d8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/hop.py
@@ -0,0 +1,785 @@
+# 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.
+"""The post_moe_hop contract: native raw NVFP4 MoE tuple -> post-reduction hidden/residual.
+
+``PostMoeHop`` replaces, for one target MoE layer, the native multi-node
+chain "MoE runner finalize -> shared add -> MNNVL all-reduce with residual and
+RMSNorm" by one fused launch that consumes the runner's ``do_finalize=False``
+outputs. It returns the same ``(normed_hidden, new_residual)`` pair as
+``AllReduce(fusion_op=RESIDUAL_RMS_NORM)``.
+
+Eligibility is decided before the MoE runner is asked for raw output, from
+rank-identical facts only, so ranks never diverge in collective count or
+order:
+
+* when the hop is built (collective): the rank-local capability, the raw-tuple
+  contract the MoE producer declares (``supports_unfinalized_output()``,
+  hidden size, dtype, top-k), the post_moe_hop workspace allocation on every rank;
+* per call (``can_run``): the token count.
+
+After the runner the hop checks the raw tuple's metadata. A dtype or shape
+that differs from the declared contract is a producer error and raises
+``PostMoeHopContractError`` (the runner already skipped its finalize, so there is no
+native path left to take). A layout the kernels cannot read directly (inner
+stride, contiguity, a base or row pitch not aligned to 16 bytes) is copied
+into an aligned contiguous tensor on this rank only, which leaves the
+collective and the arithmetic unchanged; such copies are counted
+(``PostMoeHopWorkspace.layout_repairs``) and logged once per tensor.
+
+Device errors (a peer that never publishes within ``timeout_s``, a
+permutation entry beyond the FC2 rows) set the sticky per-rank error word;
+see ``health.check_health`` for the host-side report.
+"""
+
+import contextlib
+import os
+import platform
+from typing import Optional, Tuple
+
+import torch
+
+from tensorrt_llm.mapping import Mapping
+
+from ... import dispatch
+from .workspace import PostMoeHopWorkspace, _all_ranks_agree
+
+FEATURE_ID = "post_moe_hop"
+DEFAULT_MAX_ROWS = 16
+DEFAULT_BLOCK = 1024
+DEFAULT_NUM_WARPS = 4
+# Upper bound of any legitimate cross-rank skew at the hop, including the
+# first decode-shaped forward, in which neighbouring packages may compile
+# kernels per rank; below the executor's 300 s hang detector.
+DEFAULT_TIMEOUT_S = 120.0
+DEFAULT_CLUSTER = 8
+IMPLS = ("triton", "cluster")
+_SUPPORTED_WORLDS = (2, 4, 8)
+_SCALE_DTYPES = (torch.bfloat16, torch.float32)
+_VECTOR_BYTES = 16
+
+
+class PostMoeHopContractError(RuntimeError):
+    """The raw MoE tuple does not match the contract its producer declared."""
+
+
+def _local_capability(
+    mapping: Mapping,
+    hidden_size: int,
+    dtype: torch.dtype,
+    block: int,
+    num_warps: int,
+    require_multi_node: bool,
+) -> Tuple[bool, str]:
+    """Rank-local preconditions; the caller converges them across the TP group."""
+    if mapping.enable_attention_dp:
+        return False, "attention DP"
+    if mapping.has_cp():
+        return False, "context parallelism"
+    if mapping.tp_size not in _SUPPORTED_WORLDS:
+        return False, f"tp_size {mapping.tp_size} not in {_SUPPORTED_WORLDS}"
+    if require_multi_node and not mapping.is_multi_node():
+        return False, "single-node TP group (native fused MoE all-reduce applies)"
+    if dtype != torch.bfloat16:
+        return False, f"dtype {dtype}"
+    if block < 128 or block & (block - 1):
+        return False, f"block {block} is not a power of two >= 128"
+    if block < 256 * num_warps:
+        return False, f"block {block} < 256 * num_warps ({256 * num_warps})"
+    if hidden_size % 8 or -(-hidden_size // block) > 64:
+        return False, f"hidden {hidden_size} unsupported for block {block}"
+    if not torch.cuda.is_available():
+        return False, "no CUDA device"
+    major, _ = torch.cuda.get_device_capability()
+    if major < 9:
+        return False, f"SM {major}x < 90"
+    if "aarch64" not in platform.machine().lower() and os.environ.get("TLLM_TEST_MNNVL") != "1":
+        return False, "MNNVL multicast workspace requires an aarch64 NVLink domain"
+    try:
+        import triton  # noqa: F401
+
+        from tensorrt_llm._torch.distributed.mnnvl_memory import MnnvlMemory
+    except ImportError as exc:
+        return False, f"import failed: {exc}"
+    if not MnnvlMemory.supports_mnnvl() and os.environ.get("TLLM_TEST_MNNVL") != "1":
+        return False, "MNNVL not supported on this device"
+    return True, ""
+
+
+def producer_capability(
+    producer, hidden_size: int, top_k: int, dtype: torch.dtype
+) -> Tuple[bool, str]:
+    """Static raw-tuple contract of a MoE module, read from its own declaration.
+
+    ``producer`` is the layer's MoE module (``ConfigurableMoE`` or a backend
+    leaf). It qualifies when its backend declares that ``do_finalize=False``
+    returns the unfinalized ``(gemm2_output, expert_weights,
+    expanded_idx_to_permuted_idx)`` triple, leaves the reduction to the
+    caller, and matches the hop's hidden size, activation dtype and top-k.
+    """
+    backend = getattr(producer, "backend", producer)
+    declares = getattr(backend, "supports_unfinalized_output", None)
+    if not callable(declares) or not declares():
+        return False, (
+            f"MoE backend {type(backend).__name__} does not declare the unfinalized "
+            "raw-tuple output"
+        )
+    if getattr(producer, "reduce_results", False):
+        return False, "MoE producer reduces its own output"
+    if getattr(producer, "hidden_size", hidden_size) != hidden_size:
+        return False, f"MoE hidden size {producer.hidden_size} != {hidden_size}"
+    if getattr(producer, "dtype", None) not in (None, dtype):
+        return False, f"MoE dtype {producer.dtype} != {dtype}"
+    routing = getattr(producer, "routing_method", None)
+    experts_per_token = getattr(routing, "experts_per_token", top_k)
+    if experts_per_token != top_k:
+        return False, f"MoE top-k {experts_per_token} != {top_k}"
+    return True, ""
+
+
+def _vector_ready(t: torch.Tensor, row_major_2d: bool) -> bool:
+    """Whether the kernels may read ``t`` with 16-byte vectors at its strides."""
+    if t.data_ptr() % _VECTOR_BYTES or t.stride(-1) != 1:
+        return False
+    if row_major_2d:
+        return (t.stride(0) * t.element_size()) % _VECTOR_BYTES == 0
+    return t.is_contiguous()
+
+
+class PostMoeHop:
+    """Fused post-MoE finalize + cross-rank reduction + residual + RMSNorm.
+
+    Args:
+        mapping: TP mapping of the target model (non-ADP).
+        hidden_size: model hidden size ``H``.
+        top_k: routed experts per token.
+        dtype: activation dtype (bf16 only).
+        max_rows: largest token count the hop accepts per call.
+        block: hidden columns per program (power of two dividing ``H``).
+        num_warps: Triton warps per program.
+        timeout_s: deadline of every cross-rank poll; on expiry the call sets
+            the sticky error word, writes NaN outputs and completes.
+        use_pdl: launch with programmatic dependent launch; defaults to the
+            native ``TRTLLM_ENABLE_PDL`` policy (on for SM90+ unless set to 0).
+        impl: ``"triton"`` (chunked programs, square sums exchanged through a
+            local Lamport slot) or ``"cluster"`` (CuTe DSL, one thread block
+            cluster per row, square sums reduced through distributed shared
+            memory). Shapes the cluster kernel does not cover use Triton.
+        cluster: CTAs per row for the cluster kernel.
+        require_multi_node: restrict to TP groups that span hosts (the native
+            single-node fused path already exists); tests may disable it.
+        raw_producer: the layer's MoE module; its declared raw-tuple contract
+            (``producer_capability``) is part of the collective decision.
+            ``None`` skips that check (tests that feed synthetic tuples).
+        local_veto: test-only reason that makes this rank incapable.
+        fault_delay_ns: test-only protocol fault injection (scrambled store
+            delays inside the kernels); production keeps 0.
+    """
+
+    def __init__(
+        self,
+        mapping: Mapping,
+        hidden_size: int,
+        top_k: int,
+        dtype: torch.dtype = torch.bfloat16,
+        max_rows: int = DEFAULT_MAX_ROWS,
+        block: int = DEFAULT_BLOCK,
+        num_warps: int = DEFAULT_NUM_WARPS,
+        timeout_s: float = DEFAULT_TIMEOUT_S,
+        use_pdl: Optional[bool] = None,
+        impl: str = "cluster",
+        cluster: int = DEFAULT_CLUSTER,
+        require_multi_node: bool = True,
+        raw_producer: Optional[object] = None,
+        local_veto: Optional[str] = None,
+        fault_delay_ns: int = 0,
+    ):
+        self.mapping = mapping
+        self.hidden_size = hidden_size
+        self.top_k = top_k
+        self.dtype = dtype
+        self.max_rows = max_rows
+        self.block = block
+        self.num_warps = num_warps
+        if timeout_s <= 0:
+            raise ValueError(f"timeout_s must be positive, got {timeout_s}")
+        self.timeout_ns = int(timeout_s * 1e9)
+        self.fault_delay_ns = fault_delay_ns
+        self.workspace: Optional[PostMoeHopWorkspace] = None
+        if use_pdl is None:
+            use_pdl = (
+                os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0"
+                and torch.cuda.is_available()
+                and torch.cuda.get_device_capability()[0] >= 9
+            )
+        self.use_pdl = use_pdl
+        if impl not in IMPLS:
+            raise ValueError(f"impl must be one of {IMPLS}, got {impl!r}")
+        self.cluster = cluster
+        self.impl = impl
+        if impl == "cluster":
+            from .cute_kernel import cluster_kernel_supported
+
+            why = cluster_kernel_supported(hidden_size, top_k, mapping.tp_size, cluster)
+            if why is not None:
+                dispatch.log_capability(FEATURE_ID, "cluster_kernel", False, why)
+                self.impl = "triton"
+
+        local_ok, reason = _local_capability(
+            mapping, hidden_size, dtype, block, num_warps, require_multi_node
+        )
+        if local_ok and raw_producer is not None:
+            local_ok, reason = producer_capability(raw_producer, hidden_size, top_k, dtype)
+        if local_veto is not None:
+            local_ok, reason = False, local_veto
+        self.engaged = False
+        if mapping.tp_size > 1 and not mapping.enable_attention_dp:
+            comm = PostMoeHopWorkspace.comm_for(mapping)
+            if not _all_ranks_agree(comm, local_ok):
+                reason = reason or "another rank is not capable"
+            else:
+                self.workspace, reason = PostMoeHopWorkspace.get_or_create(
+                    mapping, hidden_size, max_rows
+                )
+                self.engaged = self.workspace is not None
+        if self.engaged and self.impl == "cluster":
+            self._precompile_cluster_kernel(comm)
+        self.reason = reason
+        self._repairs_logged = set()
+        dispatch.log_capability(FEATURE_ID, f"post_moe_hop[{self.impl}]", self.engaged, reason)
+
+    def _precompile_cluster_kernel(self, comm) -> None:
+        """Compile the cluster kernel now, outside any CUDA graph capture.
+
+        Covers both shared-output variants with bf16 routing weights (the
+        TRTLLM backend's runner emits them in the activation dtype); other
+        variants compile on first use. Ranks agree on the outcome, so a
+        compile failure moves every rank to the Triton kernel together.
+        """
+        from .cute_kernel import compiled_cluster_hop
+        from .kernels import STATE_WORDS
+
+        ok = True
+        try:
+            for has_shared in (True, False):
+                compiled_cluster_hop(
+                    self.hidden_size,
+                    self.top_k,
+                    self.mapping.tp_size,
+                    self.max_rows,
+                    self.cluster,
+                    has_shared,
+                    False,
+                    self.use_pdl,
+                    STATE_WORDS,
+                    False,
+                    self.fault_delay_ns,
+                )
+        except Exception as exc:  # noqa: BLE001 - converged across ranks below
+            ok = False
+            dispatch.log_capability(FEATURE_ID, "cluster_kernel_compile", False, repr(exc))
+        if not _all_ranks_agree(comm, ok):
+            self.impl = "triton"
+
+    def can_run(self, num_rows: int) -> bool:
+        """Whether a call with ``num_rows`` tokens takes the post_moe_hop path (rank-identical)."""
+        return self.engaged and 0 < num_rows <= self.max_rows
+
+    # ------------------------------------------------------------------
+    # routed_moe_into_hop entry: reachable only when the ``routed_moe_into_hop`` feature is
+    # selected; nothing else calls it, so the post_moe_hop is unchanged otherwise.
+    # ------------------------------------------------------------------
+    ROUTED_MODES = ("row", "identity", "fp32")
+    # fp32: no finalize kernel and no raw rows; identity
+    # is the bit-exact alternative. The selector's variant table of
+    # routed_moe_into_hop is the one source of the default and of the names
+    # its knob selects; row stays a mode of the hop that no knob selects.
+    DEFAULT_ROUTED_MODE = dispatch.NUMERICAL_VARIANTS["routed_moe_into_hop"][1][0]
+    assert set(dispatch.NUMERICAL_VARIANTS["routed_moe_into_hop"][1]) <= set(ROUTED_MODES)
+
+    @staticmethod
+    def selected_routed_mode() -> str:
+        """routed_moe_into_hop's routed-input mode selected by the environment.
+
+        The variable
+        (``TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION``) is
+        parsed by the selector (``dispatch.feature_variants``), the parser of
+        the feature record, so the record and the hop name the same variant;
+        unset or empty selects ``DEFAULT_ROUTED_MODE``.
+
+        Raises:
+            ValueError: the variable names no routed_moe_into_hop mode.
+        """
+        return dispatch.feature_variants(("routed_moe_into_hop",))["routed_moe_into_hop"]
+
+    def enable_routed_input(self, mode: Optional[str] = None) -> bool:
+        """Collective: prepare the hop for a routed input written by routed_moe.
+
+        ``mode`` (default: ``selected_routed_mode()``, i.e.
+        ``TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION`` as the
+        selector parses it) selects what routed_moe writes and the cluster-kernel
+        variant that reads it: "row" a finalized bf16 row, "identity" the
+        identity-layout raw rows and their scales, "fp32" an FP32 accumulator
+        of the top-k contributions (cleared by the hop after reading).
+        Allocates the graph-stable buffer of that mode in device memory (also
+        when called during the served model construction, see
+        ``_routed_device_buffer``), checks it, and precompiles its variants;
+        every TP rank must succeed with the same mode, otherwise the feature
+        stays off on all of them. May be called again (collectively, outside
+        capture) to switch modes.
+        """
+        mode = mode or self.selected_routed_mode()
+        if mode not in self.ROUTED_MODES:
+            raise ValueError(
+                f"routed_moe_into_hop routed-input mode {mode!r} not in {self.ROUTED_MODES}"
+            )
+        if getattr(self, "_routed_enabled", None) is not None and self.routed_mode == mode:
+            return self._routed_enabled
+        ok = self.engaged and self.impl == "cluster"
+        if ok:
+            try:
+                from .cute_kernel import compiled_cluster_hop
+                from .kernels import STATE_WORDS
+
+                for has_shared in (True, False):
+                    compiled_cluster_hop(
+                        self.hidden_size,
+                        self.top_k,
+                        self.mapping.tp_size,
+                        self.max_rows,
+                        self.cluster,
+                        has_shared,
+                        False,
+                        self.use_pdl,
+                        STATE_WORDS,
+                        False,
+                        self.fault_delay_ns,
+                        mode,
+                    )
+                if mode == "row" and self._needs_routed_buffer("_routed_buffer"):
+                    self._routed_buffer = self._routed_device_buffer(
+                        (self.max_rows, self.hidden_size), self.dtype
+                    )
+                if mode == "identity" and (
+                    self._needs_routed_buffer("_routed_rows")
+                    or self._needs_routed_buffer("_routed_scales")
+                ):
+                    self._routed_rows = self._routed_device_buffer(
+                        (self.max_rows * self.top_k, self.hidden_size), self.dtype
+                    )
+                    self._routed_scales = self._routed_device_buffer(
+                        (self.max_rows, self.top_k), self.dtype
+                    )
+                if mode == "fp32" and self._needs_routed_buffer("_routed_acc"):
+                    # the routed kernel accumulates into it; the hop clears
+                    # the rows it read, so it is zero before every call
+                    self._routed_acc = self._routed_device_buffer(
+                        (self.max_rows, self.hidden_size), torch.float32
+                    )
+                problem = self._routed_buffer_problem(mode)
+                if problem:
+                    dispatch.log_capability(
+                        "routed_moe_into_hop", "post_moe_routed_input_buffer", False, problem
+                    )
+                    ok = False
+            except Exception as exc:  # noqa: BLE001 - converged across ranks below
+                dispatch.log_capability(
+                    "routed_moe_into_hop", "post_moe_routed_input_compile", False, repr(exc)
+                )
+                ok = False
+        comm = PostMoeHopWorkspace.comm_for(self.mapping)
+        # The mode comes from each rank's environment: require the same mode
+        # on every rank (one agreement per mode, in a fixed order), so no rank
+        # writes or reads another variant than its peers.
+        if not [m for m in self.ROUTED_MODES if _all_ranks_agree(comm, mode == m)]:
+            dispatch.log_capability(
+                "routed_moe_into_hop",
+                "post_moe_routed_input_mode",
+                False,
+                f"routed-input modes differ across ranks (this rank: {mode})",
+            )
+            ok = False
+        self._routed_enabled = _all_ranks_agree(comm, ok)
+        self.routed_mode = mode
+        dispatch.log_capability(
+            "routed_moe_into_hop", f"post_moe_routed_input[{mode}]", self._routed_enabled
+        )
+        return self._routed_enabled
+
+    _ROUTED_BUFFERS = {
+        "row": ("_routed_buffer",),
+        "identity": ("_routed_rows", "_routed_scales"),
+        "fp32": ("_routed_acc",),
+    }
+
+    def _needs_routed_buffer(self, name: str) -> bool:
+        buf = getattr(self, name, None)
+        return buf is None or buf.is_meta
+
+    def _routed_device_buffer(self, shape: Tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
+        """A zeroed routed-input buffer in device memory on the workspace's device.
+
+        The served model loader constructs the model, and with it this hop,
+        inside ``MetaInitMode``: a dispatch mode that turns ``empty`` /
+        ``empty_like`` into meta tensors (data pointer 0) and afterwards
+        materializes only module parameters and buffers. The hop is not a
+        module, so its buffers are allocated with every dispatch mode
+        suspended.
+        """
+        try:
+            from torch.utils._python_dispatch import _disable_current_modes
+        except ImportError:  # the result is still checked by _routed_buffer_problem
+            _disable_current_modes = contextlib.nullcontext
+        with _disable_current_modes():
+            return torch.zeros(shape, dtype=dtype, device=self.workspace.state.device)
+
+    def _routed_buffer_problem(self, mode: str) -> str:
+        """Why ``mode``'s buffers cannot be handed to the kernels ("" when they can)."""
+        device = self.workspace.state.device
+        for name in self._ROUTED_BUFFERS[mode]:
+            buf = getattr(self, name, None)
+            if buf is None or buf.is_meta or buf.device != device:
+                where = None if buf is None else buf.device
+                return f"{name} is on {where}, not in device memory on {device}"
+            if not buf.is_contiguous() or buf.data_ptr() % _VECTOR_BYTES:
+                return f"{name} is not a contiguous 16-byte aligned buffer"
+        return ""
+
+    def routed_input_buffer(self, num_rows: int) -> torch.Tensor:
+        """Mode "row": the ``[num_rows, H]`` view routed_moe writes its finalized output into."""
+        return self._routed_buffer[:num_rows]
+
+    def routed_raw_buffers(self, num_rows: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Mode "identity": ``(rows [num_rows * top_k, H], scales [num_rows, top_k])``."""
+        return self._routed_rows[: num_rows * self.top_k], self._routed_scales[:num_rows]
+
+    def routed_acc_buffer(self, num_rows: int) -> torch.Tensor:
+        """Mode "fp32": the zeroed FP32 ``[num_rows, H]`` accumulator."""
+        return self._routed_acc[:num_rows]
+
+    def _routed_view(self, num_rows: int) -> Tuple[str, torch.Tensor, Optional[torch.Tensor]]:
+        """``(mode, buffer view, scales)`` the hop reads in its routed-input mode."""
+        if self.routed_mode == "fp32":
+            return "fp32", self.routed_acc_buffer(num_rows), None
+        if self.routed_mode == "identity":
+            rows, scales = self.routed_raw_buffers(num_rows)
+            return "identity", rows, scales
+        return "row", self.routed_input_buffer(num_rows), None
+
+    def hop_routed(
+        self,
+        routed: torch.Tensor,
+        shared_output: Optional[torch.Tensor],
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """The hop on a routed input written by routed_moe (routed_moe_into_hop).
+
+        ``routed`` is the buffer view ``routed_output_into_hop`` returned (a
+        view of this hop's own buffer in its routed-input mode). Same
+        arithmetic as ``__call__`` after its finalize: ``bf16(routed +
+        shared)``, rank-ordered FP32 sum over the TP group, residual added
+        once, RMSNorm.
+
+        Raises:
+            ValueError: ``can_run`` is false or the routed input is not enabled.
+            PostMoeHopContractError: ``routed`` is not this hop's buffer for the row
+                count, or a dtype / shape differs from the contract.
+        """
+        num_rows = residual.shape[0]
+        if not (self.can_run(num_rows) and getattr(self, "_routed_enabled", False)):
+            raise ValueError(
+                "routed_moe_into_hop hop called for an ineligible shape or before enable_routed_input()"
+            )
+        mode, expect, scales = self._routed_view(num_rows)
+        rows_h = (num_rows, self.hidden_size)
+        problems = []
+        if (
+            routed.data_ptr() != expect.data_ptr()
+            or routed.dtype != expect.dtype
+            or tuple(routed.shape) != tuple(expect.shape)
+            or tuple(routed.stride()) != tuple(expect.stride())
+        ):
+            problems.append(
+                f"routed {routed.dtype} {tuple(routed.shape)} is not the hop's {mode} buffer"
+            )
+        if residual.dtype != self.dtype or tuple(residual.shape) != rows_h:
+            problems.append(f"residual {residual.dtype} {tuple(residual.shape)}")
+        if shared_output is not None and (
+            shared_output.dtype != self.dtype or tuple(shared_output.shape) != rows_h
+        ):
+            problems.append(f"shared_output {shared_output.dtype} {tuple(shared_output.shape)}")
+        if norm_weight.dtype != self.dtype or tuple(norm_weight.shape) != (self.hidden_size,):
+            problems.append(f"norm_weight {norm_weight.dtype} {tuple(norm_weight.shape)}")
+        if problems:
+            raise PostMoeHopContractError(
+                "routed_moe_into_hop routed input violates the hop's contract: "
+                f"{'; '.join(problems)}"
+            )
+        if shared_output is not None:
+            shared_output = self._vector_layout("shared_output", shared_output)
+        residual = self._vector_layout("residual", residual)
+        norm_weight = self._vector_layout("norm_weight", norm_weight)
+        from .cute_kernel import launch_cluster_hop_routed
+
+        out = torch.empty_like(residual)
+        residual_out = torch.empty_like(residual)
+        workspace = self.workspace
+        launch_cluster_hop_routed(
+            routed=routed,
+            shared_output=shared_output,
+            residual=residual,
+            norm_weight=norm_weight,
+            out=out,
+            residual_out=residual_out,
+            mc_buffer=workspace.mc,
+            uc_buffer=workspace.uc,
+            partials=workspace.partials,
+            state=workspace.state,
+            rank=self.mapping.tp_rank,
+            world=self.mapping.tp_size,
+            max_rows=self.max_rows,
+            eps=eps,
+            cluster=self.cluster,
+            use_pdl=self.use_pdl,
+            timeout_ns=self.timeout_ns,
+            top_k=self.top_k,
+            mode=mode,
+            scales=scales,
+            fault_delay_ns=self.fault_delay_ns,
+        )
+        return out, residual_out
+
+    def _contract_violations(
+        self,
+        fc2_output,
+        expert_scale_factor,
+        expanded_idx_to_permuted_idx,
+        shared_output,
+        residual,
+        norm_weight,
+    ) -> list:
+        """Dtype/shape differences from the declared raw-tuple contract (rank-identical)."""
+        num_rows = residual.shape[0]
+        rows_h = (num_rows, self.hidden_size)
+        problems = []
+        if residual.dtype != self.dtype or tuple(residual.shape) != rows_h:
+            problems.append(f"residual {residual.dtype} {tuple(residual.shape)}")
+        if (
+            fc2_output.dtype != self.dtype
+            or fc2_output.dim() != 2
+            or fc2_output.shape[1] < self.hidden_size
+        ):
+            problems.append(f"fc2_output {fc2_output.dtype} {tuple(fc2_output.shape)}")
+        if expert_scale_factor.dtype not in _SCALE_DTYPES or tuple(expert_scale_factor.shape) != (
+            num_rows,
+            self.top_k,
+        ):
+            problems.append(
+                f"expert_scale_factor {expert_scale_factor.dtype} "
+                f"{tuple(expert_scale_factor.shape)}"
+            )
+        if expanded_idx_to_permuted_idx.dtype != torch.int32 or tuple(
+            expanded_idx_to_permuted_idx.shape
+        ) != (num_rows, self.top_k):
+            problems.append(
+                f"expanded_idx_to_permuted_idx {expanded_idx_to_permuted_idx.dtype} "
+                f"{tuple(expanded_idx_to_permuted_idx.shape)}"
+            )
+        if shared_output is not None and (
+            shared_output.dtype != self.dtype or tuple(shared_output.shape) != rows_h
+        ):
+            problems.append(f"shared_output {shared_output.dtype} {tuple(shared_output.shape)}")
+        if norm_weight.dtype != self.dtype or tuple(norm_weight.shape) != (self.hidden_size,):
+            problems.append(f"norm_weight {norm_weight.dtype} {tuple(norm_weight.shape)}")
+        return problems
+
+    def _vector_layout(self, name: str, t: torch.Tensor, row_major_2d: bool = False):
+        """Return ``t`` or a 16-byte-aligned contiguous copy of it (rank-local)."""
+        if _vector_ready(t, row_major_2d):
+            return t
+        misalignment, strides = t.data_ptr() % _VECTOR_BYTES, tuple(t.stride())
+        if row_major_2d:
+            t = t[:, : self.hidden_size]
+        fixed = torch.empty(t.shape, dtype=t.dtype, device=t.device)
+        fixed.copy_(t)
+        self.workspace.layout_repairs += 1
+        if name not in self._repairs_logged:
+            self._repairs_logged.add(name)
+            from tensorrt_llm.logger import logger
+
+            logger.warning(
+                f"[glm_moe_dsa] post_moe_hop/hop: {name} is not ready for 16-byte vector "
+                f"access (data_ptr % 16 = {misalignment}, strides {strides}); "
+                "copying it before every such call"
+            )
+        return fixed
+
+    def __call__(
+        self,
+        fc2_output: torch.Tensor,
+        expert_scale_factor: torch.Tensor,
+        expanded_idx_to_permuted_idx: torch.Tensor,
+        shared_output: Optional[torch.Tensor],
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        eps: float,
+        *,
+        trace: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return ``(rmsnorm(reduced + residual), reduced + residual)``.
+
+        Args:
+            fc2_output: bf16 ``[N_padded, >=H]`` raw FC2 rows of the native runner.
+            expert_scale_factor: ``[R, top_k]`` routing weights (bf16 or fp32).
+            expanded_idx_to_permuted_idx: int32 ``[R, top_k]``; negative = no local
+                row; entries ``>= N_padded`` are skipped and set error bit 4.
+            shared_output: bf16 ``[R, H]`` rank-local shared-expert partial, or ``None``.
+            residual: bf16 ``[R, H]`` stock residual (added once, after the reduction).
+            norm_weight: bf16 ``[H]`` next-layer RMSNorm weight.
+            eps: RMSNorm epsilon.
+            trace: diagnostics only; int64 ``[programs, 8]`` phase timestamps.
+
+        Raises:
+            ValueError: ``can_run`` is false for this row count.
+            PostMoeHopContractError: a dtype or shape differs from the raw-tuple contract.
+        """
+        if not self.can_run(residual.shape[0]):
+            raise ValueError(
+                "post_moe_hop hop called for an ineligible shape; check can_run() first"
+            )
+        problems = self._contract_violations(
+            fc2_output,
+            expert_scale_factor,
+            expanded_idx_to_permuted_idx,
+            shared_output,
+            residual,
+            norm_weight,
+        )
+        if problems:
+            raise PostMoeHopContractError(
+                "raw MoE tuple violates the declared native raw-tuple contract after the runner "
+                f"skipped its finalize: {'; '.join(problems)}"
+            )
+        fc2_output = self._vector_layout("fc2_output", fc2_output, row_major_2d=True)
+        expert_scale_factor = self._vector_layout("expert_scale_factor", expert_scale_factor)
+        expanded_idx_to_permuted_idx = self._vector_layout(
+            "expanded_idx_to_permuted_idx", expanded_idx_to_permuted_idx
+        )
+        if shared_output is not None:
+            shared_output = self._vector_layout("shared_output", shared_output)
+        residual = self._vector_layout("residual", residual)
+        norm_weight = self._vector_layout("norm_weight", norm_weight)
+
+        out = torch.empty_like(residual)
+        residual_out = torch.empty_like(residual)
+        workspace = self.workspace
+        if self.impl == "cluster" and trace is None:
+            from .cute_kernel import launch_cluster_hop
+
+            launch_cluster_hop(
+                fc2_output=fc2_output,
+                expert_scale_factor=expert_scale_factor,
+                expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+                shared_output=shared_output,
+                residual=residual,
+                norm_weight=norm_weight,
+                out=out,
+                residual_out=residual_out,
+                mc_buffer=workspace.mc,
+                uc_buffer=workspace.uc,
+                partials=workspace.partials,
+                state=workspace.state,
+                rank=self.mapping.tp_rank,
+                world=self.mapping.tp_size,
+                max_rows=self.max_rows,
+                eps=eps,
+                cluster=self.cluster,
+                use_pdl=self.use_pdl,
+                timeout_ns=self.timeout_ns,
+                fault_delay_ns=self.fault_delay_ns,
+            )
+            return out, residual_out
+        from .kernels import launch_post_moe_hop
+
+        launch_post_moe_hop(
+            fc2_output=fc2_output,
+            expert_scale_factor=expert_scale_factor,
+            expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+            shared_output=shared_output,
+            residual=residual,
+            norm_weight=norm_weight,
+            out=out,
+            residual_out=residual_out,
+            mc_buffer=workspace.mc,
+            uc_buffer=workspace.uc,
+            partials=workspace.partials,
+            state=workspace.state,
+            rank=self.mapping.tp_rank,
+            world=self.mapping.tp_size,
+            max_rows=self.max_rows,
+            eps=eps,
+            block=self.block,
+            num_warps=self.num_warps,
+            timeout_ns=self.timeout_ns,
+            use_pdl=self.use_pdl,
+            trace=trace,
+            fault_delay_ns=self.fault_delay_ns,
+        )
+        return out, residual_out
+
+    def error_bits(self) -> int:
+        """Sticky device error word of this rank (synchronizes); 0 means healthy."""
+        return 0 if self.workspace is None else self.workspace.error_bits()
+
+    def raise_if_error(self) -> None:
+        """Raise if any program of any earlier call failed on this rank."""
+        bits = self.error_bits()
+        if bits:
+            raise RuntimeError(
+                f"post_moe_hop post-MoE hop reported error bits {bits:#x} on TP rank "
+                f"{self.mapping.tp_rank}; the run is invalid and the workspace "
+                "needs a collective reset()"
+            )
+
+    def reset(self) -> bool:
+        """Collective workspace reset after an error; see ``PostMoeHopWorkspace.reset``."""
+        return self.workspace.reset() if self.workspace is not None else True
+
+
+def native_finalize_reference(
+    fc2_output: torch.Tensor,
+    expert_scale_factor: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    hidden_size: int,
+) -> torch.Tensor:
+    """Torch version of the runner's finalize: fp32 FMA chain over top-k, one bf16 rounding.
+
+    Test oracle. Each step ``acc = fma(w_k, x_k, acc)`` is evaluated in fp64
+    and rounded to fp32 (the bf16 x fp32 product is exact in fp64), which
+    reproduces the native kernel's fused multiply-add in ``k`` order. Entries
+    outside ``[0, N_padded)`` contribute nothing, like in the post_moe_hop kernels.
+    """
+    num_rows, top_k = expanded_idx_to_permuted_idx.shape
+    idx = expanded_idx_to_permuted_idx.to(torch.long)
+    valid = (idx >= 0) & (idx < fc2_output.shape[0])
+    safe = torch.where(valid, idx, torch.zeros_like(idx))
+    rows = fc2_output[safe.view(-1), :hidden_size].view(num_rows, top_k, hidden_size)
+    weights = expert_scale_factor.double()
+    acc = torch.zeros((num_rows, hidden_size), dtype=torch.float32, device=fc2_output.device)
+    for k in range(top_k):
+        step = (acc.double() + weights[:, k : k + 1] * rows[:, k].double()).float()
+        acc = torch.where(valid[:, k : k + 1], step, acc)
+    return acc.to(fc2_output.dtype)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/kernels.py
new file mode 100644
index 000000000000..dace8cd9254e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/kernels.py
@@ -0,0 +1,446 @@
+# 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.
+"""Triton kernel of the post_moe_hop (the post_moe_hop contract).
+
+One launch per target MoE layer and rank. Program ``(row, chunk)`` owns
+``BLOCK`` hidden columns of one token row and:
+
+1. finalizes the raw NVFP4 MoE tuple locally, with the arithmetic of the
+   native multi-node chain: fp32 FMA over the top-k expert rows in ``k``
+   order (``-1`` entries skipped), one bf16 rounding, then the bf16 add of the
+   shared-expert partial;
+2. publishes the bf16 local partial through the multicast address of the
+   post_moe_hop workspace into slot ``[row][rank]`` of every rank's Lamport buffer;
+3. clears the Lamport buffer used by the previous call (all programs share
+   the work);
+4. polls the local copy of slot ``[row][*]`` until no element holds the
+   bf16 ``-0.0`` sentinel, sums the ranks in rank order in fp32 and rounds
+   once to bf16, like the native MNNVL one-shot all-reduce;
+5. adds the residual once (bf16 add) and writes the new residual;
+6. exchanges per-chunk sums of bf16-rounded squares of the new residual
+   through a local Lamport buffer, so every program of the row sees the same
+   row sum, and writes ``bf16(x * rsqrt(sum / H + eps) * gamma)``.
+
+The last program to arrive advances the epoch (Lamport buffer index) and
+records the row count to clear on the next call.
+
+Errors. Every poll is bounded by a ``%globaltimer`` deadline (``timeout_ns``),
+read through a program-wide reduction so that all warps of a program leave
+the loop at the same iteration; on expiry the program sets a bit in the
+per-rank error word, records its poll count, writes NaN residual/normalized
+outputs and finishes, so a missing peer is reported instead of hanging the
+stream. A permutation entry at or beyond
+the FC2 row count is skipped and also sets an error bit. Each Lamport buffer
+has its own error word, which collects the bits of the calls that use the
+buffer. A call reads the error word of the previous call's buffer, which is
+complete at the kernel boundary; when it is nonzero, the call publishes NaN
+instead of its partial and carries the bits into its own error word. Errors
+are therefore sticky until a collective workspace reset, the group's outputs
+become NaN rather than silently wrong (the ranks' epochs are not guaranteed
+to agree after a timeout), and a call's own result does not depend on the
+order in which its programs record errors.
+
+Workspace layout (see ``workspace.py``):
+
+* ``lamport`` bf16 ``[3, max_rows, world, hidden]`` per rank (multicast
+  stores, unicast polls). ``-0.0`` (0x8000) marks an empty element.
+* ``partials`` fp32 ``[3, max_rows, MAX_CHUNKS]`` per rank, local only.
+  ``-0.0`` marks an empty element.
+* ``state`` int32: ``[epoch, prev_rows, arrivals, error[0], calls,
+  timeout_polls, error[1], error[2]]``; ``error[b]`` is the error word of
+  Lamport buffer ``b`` (the rank's error bits are their OR) and
+  ``timeout_polls`` the largest poll count at which a poll of this rank gave
+  up (diagnostics).
+"""
+
+import triton
+import triton.language as tl
+
+STATE_EPOCH = 0
+STATE_PREV_ROWS = 1
+STATE_ARRIVALS = 2
+STATE_CALLS = 4
+STATE_TIMEOUT_POLLS = 5
+# Error word of Lamport buffer 0, 1, 2.
+STATE_ERROR_WORDS = (3, 6, 7)
+STATE_WORDS = 8
+
+ERROR_LAMPORT_TIMEOUT = 1
+ERROR_PARTIAL_TIMEOUT = 2
+ERROR_PERM_RANGE = 4
+
+MAX_CHUNKS = 64
+
+TRACE_POINTS = 8
+
+_BF16_NEG_ZERO_BITS = tl.constexpr(-32768)  # 0x8000 as int16
+_FP32_NEG_ZERO_BITS = tl.constexpr(-2147483648)  # 0x80000000 as int32
+_BF16_NAN_BITS = tl.constexpr(0x7FC0)
+# Polls between two reads of the program clock (each read is a reduction).
+_CLOCK_EVERY = tl.constexpr(16)
+_FP32_NAN_BITS = tl.constexpr(0x7FC00000)
+
+
+@triton.jit
+def _globaltimer(like):
+    return tl.inline_asm_elementwise(
+        "mov.u64 $0, %globaltimer;", "=l,r", [like], dtype=tl.int64, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _program_clock(lanes):
+    """``%globaltimer`` as one value that every thread of the program agrees on.
+
+    Each thread reads the timer for its elements of ``lanes`` (the program's
+    column vector, which spans every warp when ``BLOCK >= 256 * num_warps``)
+    and a program-wide max reduction combines them. The poll loops contain
+    program-wide reductions, so every quantity in their exit condition must be
+    uniform across warps; a per-thread clock, or a pending count reduced only
+    within a warp, lets warps leave a loop at different iterations and pair up
+    different barriers.
+    """
+    return tl.max(_globaltimer(lanes), axis=0)
+
+
+@triton.jit
+def _trace(trace_ptr, pid, point: tl.constexpr, TRACE: tl.constexpr):
+    if TRACE:
+        tl.store(trace_ptr + pid * 8 + point, _globaltimer(pid))
+
+
+@triton.jit
+def _fault_delay(pid, epoch, rank, site: tl.constexpr, FAULT_DELAY_NS: tl.constexpr):
+    """Test-only fault injection: about one program in four spins before ``site``.
+
+    The delay (up to ``FAULT_DELAY_NS``) is a hash of program, epoch, rank and
+    site, so individual stores of one call land in a scrambled order.
+    """
+    if FAULT_DELAY_NS > 0:
+        h = (pid * 40503 + epoch * 2531 + rank * 977 + site * 7919) % 65521
+        if h % 4 == 0:
+            delay = (h * 6151) % FAULT_DELAY_NS
+            start = _globaltimer(pid)
+            now = start
+            while now - start < delay:
+                now = _globaltimer(pid)
+
+
+@triton.jit
+def _post_moe_hop_kernel(
+    fc2_ptr,
+    fc2_row_stride,
+    scale_ptr,
+    perm_ptr,
+    shared_ptr,
+    residual_ptr,
+    gamma_ptr,
+    out_ptr,
+    residual_out_ptr,
+    mc_ptr,
+    uc_ptr,
+    partials_ptr,
+    state_ptr,
+    trace_ptr,
+    fc2_rows,
+    rank,
+    eps,
+    timeout_ns,
+    HIDDEN: tl.constexpr,
+    TOP_K: tl.constexpr,
+    WORLD: tl.constexpr,
+    MAX_ROWS: tl.constexpr,
+    BLOCK: tl.constexpr,
+    NUM_CHUNKS: tl.constexpr,
+    CHUNKS_P2: tl.constexpr,
+    CLEAR_BLOCK: tl.constexpr,
+    PARTIAL_STRIDE: tl.constexpr,
+    HAS_SHARED: tl.constexpr,
+    USE_PDL: tl.constexpr,
+    TRACE: tl.constexpr,
+    FAULT_DELAY_NS: tl.constexpr,
+):
+    row = tl.program_id(0)
+    chunk = tl.program_id(1)
+    num_rows = tl.num_programs(0)
+    pid = row * NUM_CHUNKS + chunk
+    num_programs = num_rows * NUM_CHUNKS
+    _trace(trace_ptr, pid, 0, TRACE)
+    if USE_PDL:
+        # Every global access below may depend on the previous kernel,
+        # including the state words written by the previous post_moe_hop call.
+        tl.extra.cuda.gdc_wait()
+    _trace(trace_ptr, pid, 1, TRACE)
+
+    buffer_elems = MAX_ROWS * WORLD * HIDDEN
+    epoch = tl.load(state_ptr + 0)
+    prev_rows = tl.load(state_ptr + 1)
+    cur = epoch % 3
+    dirty = (epoch + 2) % 3
+    # Error words of this call's buffer and of the previous call's buffer.
+    err_ptr = state_ptr + tl.where(cur == 0, 3, 5 + cur)
+    sticky_bits = tl.load(state_ptr + tl.where(dirty == 0, 3, 5 + dirty))
+    sticky = sticky_bits != 0
+
+    cols = chunk * BLOCK + tl.arange(0, BLOCK)
+    col_mask = cols < HIDDEN
+
+    # Inputs that do not depend on the exchange are fetched up front so their
+    # latency overlaps the finalize and the cross-rank wait.
+    residual = tl.load(residual_ptr + row * HIDDEN + cols, mask=col_mask, other=0.0)
+    gamma = tl.load(gamma_ptr + cols, mask=col_mask, other=0.0).to(tl.float32)
+
+    # 1. Local finalize, native multi-node arithmetic.
+    acc = tl.zeros([BLOCK], dtype=tl.float32)
+    for k in tl.static_range(TOP_K):
+        permuted = tl.load(perm_ptr + row * TOP_K + k)
+        weight = tl.load(scale_ptr + row * TOP_K + k).to(tl.float32)
+        valid = (permuted >= 0) & (permuted < fc2_rows)
+        if permuted >= fc2_rows:
+            tl.atomic_or(err_ptr, 4)
+        safe_row = tl.where(valid, permuted, 0).to(tl.int64)
+        expert_row = tl.load(
+            fc2_ptr + safe_row * fc2_row_stride + cols, mask=col_mask, other=0.0
+        ).to(tl.float32)
+        acc = tl.where(valid, tl.fma(weight, expert_row, acc), acc)
+    local = acc.to(tl.bfloat16)
+    if HAS_SHARED:
+        shared = tl.load(shared_ptr + row * HIDDEN + cols, mask=col_mask, other=0.0)
+        local = (shared.to(tl.float32) + local.to(tl.float32)).to(tl.bfloat16)
+    local_bits = local.to(tl.int16, bitcast=True)
+    local_bits = tl.where(local_bits == _BF16_NEG_ZERO_BITS, tl.zeros_like(local_bits), local_bits)
+    local_bits = tl.where(sticky, tl.full([BLOCK], _BF16_NAN_BITS, tl.int16), local_bits)
+    local = local_bits.to(tl.bfloat16, bitcast=True)
+
+    # 2. Publish through multicast into slot [row][rank] of every rank.
+    row_base = row * WORLD * HIDDEN
+    _fault_delay(pid, epoch, rank, 0, FAULT_DELAY_NS)
+    tl.store(mc_ptr + cur * buffer_elems + row_base + rank * HIDDEN + cols, local, mask=col_mask)
+    if sticky:
+        # After an error recorded by an earlier call: NaN until the reset.
+        tl.atomic_or(err_ptr, sticky_bits)
+    _trace(trace_ptr, pid, 2, TRACE)
+
+    # 3. Clear the previous call's Lamport rows and partials.
+    _fault_delay(pid, epoch, rank, 1, FAULT_DELAY_NS)
+    sentinel_bf16 = tl.full([CLEAR_BLOCK], _BF16_NEG_ZERO_BITS, tl.int16).to(
+        tl.bfloat16, bitcast=True
+    )
+    clear_elems = prev_rows * WORLD * HIDDEN
+    clear_offsets = tl.arange(0, CLEAR_BLOCK)
+    for start in range(pid * CLEAR_BLOCK, clear_elems, num_programs * CLEAR_BLOCK):
+        offs = start + clear_offsets
+        tl.store(uc_ptr + dirty * buffer_elems + offs, sentinel_bf16, mask=offs < clear_elems)
+    # Unconditional: hops with different launch configs share the workspace,
+    # so a single-chunk call must still clear what a chunked call left behind.
+    if pid == 0:
+        sentinel_fp32 = tl.full([PARTIAL_STRIDE], _FP32_NEG_ZERO_BITS, tl.int32).to(
+            tl.float32, bitcast=True
+        )
+        part_offsets = tl.arange(0, PARTIAL_STRIDE)
+        for r in range(0, prev_rows):
+            tl.store(
+                partials_ptr + (dirty * MAX_ROWS + r) * PARTIAL_STRIDE + part_offsets, sentinel_fp32
+            )
+
+    _trace(trace_ptr, pid, 3, TRACE)
+
+    # 4. Lamport poll of every rank's slot, then a rank-ordered fp32 sum.
+    rank_ids = tl.arange(0, WORLD)
+    poll_ptrs = uc_ptr + cur * buffer_elems + row_base + rank_ids[:, None] * HIDDEN + cols[None, :]
+    poll_mask = col_mask[None, :] & (rank_ids[:, None] < WORLD)
+    values = tl.load(poll_ptrs, mask=poll_mask, other=0.0, volatile=True)
+    pending = tl.sum((values.to(tl.int16, bitcast=True) == _BF16_NEG_ZERO_BITS).to(tl.int32))
+    spins = 0
+    poll_start = _program_clock(cols)
+    elapsed = poll_start - poll_start
+    while (pending > 0) & (elapsed < timeout_ns):
+        values = tl.load(poll_ptrs, mask=poll_mask, other=0.0, volatile=True)
+        pending = tl.sum((values.to(tl.int16, bitcast=True) == _BF16_NEG_ZERO_BITS).to(tl.int32))
+        spins += 1
+        if spins % _CLOCK_EVERY == 0:
+            elapsed = _program_clock(cols) - poll_start
+    timed_out = pending > 0
+    if timed_out:
+        tl.atomic_or(err_ptr, 1)
+        tl.atomic_max(state_ptr + 5, spins)
+    _trace(trace_ptr, pid, 4, TRACE)
+    values_f32 = values.to(tl.float32)
+    reduced = tl.zeros([BLOCK], dtype=tl.float32)
+    for r in tl.static_range(WORLD):
+        reduced += tl.sum(tl.where(rank_ids[:, None] == r, values_f32, 0.0), axis=0)
+    reduced_bf16 = reduced.to(tl.bfloat16)
+    if USE_PDL:
+        # Dependents still wait for this grid to finish; only their launch
+        # latency overlaps the epilogue.
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # 5. Residual added once, after the reduction.
+    new_residual = (reduced_bf16.to(tl.float32) + residual.to(tl.float32)).to(tl.bfloat16)
+    new_residual = tl.where(
+        timed_out,
+        tl.full([BLOCK], _BF16_NAN_BITS, tl.int16),
+        new_residual.to(tl.int16, bitcast=True),
+    ).to(tl.bfloat16, bitcast=True)
+    tl.store(residual_out_ptr + row * HIDDEN + cols, new_residual, mask=col_mask)
+
+    # 6. RMSNorm over the full row.
+    x = new_residual.to(tl.float32)
+    square_sum = tl.sum((x * x).to(tl.bfloat16).to(tl.float32))
+    if NUM_CHUNKS > 1:
+        part_base = partials_ptr + (cur * MAX_ROWS + row) * PARTIAL_STRIDE
+        tl.store(part_base + chunk, square_sum)
+        # The poll reads the row's partials through a vector spread over every
+        # warp (lane i reads partial i % CHUNKS_P2), so the pending count is a
+        # program-wide sum and all warps leave the loop at the same iteration.
+        lane_chunk = tl.arange(0, BLOCK) % CHUNKS_P2
+        lane_mask = lane_chunk < NUM_CHUNKS
+        lane_parts = tl.load(part_base + lane_chunk, mask=lane_mask, other=0.0, volatile=True)
+        part_pending = tl.sum(
+            (lane_parts.to(tl.int32, bitcast=True) == _FP32_NEG_ZERO_BITS).to(tl.int32)
+        )
+        part_spins = 0
+        part_start = _program_clock(cols)
+        part_elapsed = part_start - part_start
+        while (part_pending > 0) & (part_elapsed < timeout_ns):
+            lane_parts = tl.load(part_base + lane_chunk, mask=lane_mask, other=0.0, volatile=True)
+            part_pending = tl.sum(
+                (lane_parts.to(tl.int32, bitcast=True) == _FP32_NEG_ZERO_BITS).to(tl.int32)
+            )
+            part_spins += 1
+            if part_spins % _CLOCK_EVERY == 0:
+                part_elapsed = _program_clock(cols) - part_start
+        if part_pending > 0:
+            tl.atomic_or(err_ptr, 2)
+            tl.atomic_max(state_ptr + 5, part_spins)
+        # Each partial is written once per call, so once none is pending every
+        # warp reads the same final values and sums them in the same order.
+        chunk_ids = tl.arange(0, CHUNKS_P2)
+        parts = tl.load(
+            part_base + chunk_ids, mask=chunk_ids < NUM_CHUNKS, other=0.0, volatile=True
+        )
+        fp32_nan = tl.full([], _FP32_NAN_BITS, tl.int32).to(tl.float32, bitcast=True)
+        square_sum = tl.where(part_pending > 0, fp32_nan, tl.sum(parts))
+    _trace(trace_ptr, pid, 5, TRACE)
+    inv_rms = tl.rsqrt(square_sum / HIDDEN + eps)
+    tl.store(out_ptr + row * HIDDEN + cols, ((x * inv_rms) * gamma).to(tl.bfloat16), mask=col_mask)
+
+    _trace(trace_ptr, pid, 6, TRACE)
+
+    # 7. The last program to arrive advances the epoch for the next call.
+    _fault_delay(pid, epoch, rank, 2, FAULT_DELAY_NS)
+    arrived = tl.atomic_add(state_ptr + 2, 1, sem="acq_rel", scope="gpu")
+    if arrived == num_programs - 1:
+        tl.store(state_ptr + 1, num_rows)
+        tl.store(state_ptr + 2, 0)
+        tl.atomic_add(state_ptr + 4, 1, sem="relaxed", scope="gpu")
+        tl.store(state_ptr + 0, (epoch + 1) % 3)
+
+
+def launch_post_moe_hop(
+    *,
+    fc2_output,
+    expert_scale_factor,
+    expanded_idx_to_permuted_idx,
+    shared_output,
+    residual,
+    norm_weight,
+    out,
+    residual_out,
+    mc_buffer,
+    uc_buffer,
+    partials,
+    state,
+    rank: int,
+    world: int,
+    max_rows: int,
+    eps: float,
+    block: int,
+    num_warps: int,
+    timeout_ns: int,
+    use_pdl: bool,
+    trace=None,
+    clear_block: int = 1024,
+    fault_delay_ns: int = 0,
+    fc2_rows=None,
+) -> None:
+    """Launch the fused hop on the current stream.
+
+    Args:
+        fc2_output: bf16 ``[N_padded, >=H]`` permuted expert rows, unit column
+            stride; permutation entries ``>= N_padded`` are skipped and flagged.
+        expert_scale_factor: bf16 or fp32 ``[R, top_k]`` contiguous.
+        expanded_idx_to_permuted_idx: int32 ``[R, top_k]`` contiguous; ``-1`` skips a slot.
+        shared_output: bf16 ``[R, H]`` contiguous, or ``None``.
+        residual: bf16 ``[R, H]`` contiguous.
+        norm_weight: bf16 ``[H]``.
+        out, residual_out: bf16 ``[R, H]`` outputs.
+        mc_buffer, uc_buffer: bf16 views of this rank's Lamport region
+            (``[3 * max_rows * world * H]``), multicast and local unicast.
+        partials: fp32 ``[3 * max_rows * MAX_CHUNKS]`` local buffer.
+        state: int32 ``[STATE_WORDS]`` local state words.
+        timeout_ns: deadline of every poll, in ``%globaltimer`` nanoseconds.
+        trace: optional int64 ``[R * chunks, TRACE_POINTS]`` buffer that receives
+            per-program ``%globaltimer`` stamps at the phase boundaries
+            (diagnostics only).
+        fault_delay_ns: test-only fault injection (see ``_fault_delay``); 0 in
+            production.
+        fc2_rows: test-only override of the permutation bound (memcheck
+            positive control); defaults to ``fc2_output.shape[0]``.
+    """
+    if block < 256 * num_warps:
+        raise ValueError(
+            f"block {block} < 256 * num_warps ({256 * num_warps}): the poll loops "
+            "need the column vector to span every warp of the program"
+        )
+    num_rows, hidden = residual.shape
+    top_k = expanded_idx_to_permuted_idx.shape[1]
+    num_chunks = triton.cdiv(hidden, block)
+    grid = (num_rows, num_chunks)
+    _post_moe_hop_kernel[grid](
+        fc2_output,
+        fc2_output.stride(0),
+        expert_scale_factor,
+        expanded_idx_to_permuted_idx,
+        shared_output if shared_output is not None else residual,
+        residual,
+        norm_weight,
+        out,
+        residual_out,
+        mc_buffer,
+        uc_buffer,
+        partials,
+        state,
+        trace if trace is not None else state,
+        fc2_output.shape[0] if fc2_rows is None else fc2_rows,
+        rank,
+        eps,
+        timeout_ns,
+        HIDDEN=hidden,
+        TOP_K=top_k,
+        WORLD=world,
+        MAX_ROWS=max_rows,
+        BLOCK=block,
+        NUM_CHUNKS=num_chunks,
+        CHUNKS_P2=triton.next_power_of_2(num_chunks),
+        CLEAR_BLOCK=clear_block,
+        PARTIAL_STRIDE=MAX_CHUNKS,
+        HAS_SHARED=shared_output is not None,
+        USE_PDL=use_pdl,
+        TRACE=trace is not None,
+        FAULT_DELAY_NS=fault_delay_ns,
+        num_warps=num_warps,
+        launch_pdl=use_pdl,
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/workspace.py
new file mode 100644
index 000000000000..22de592ba9c1
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/workspace.py
@@ -0,0 +1,200 @@
+# 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.
+"""post_moe_hop-owned MNNVL workspace and the rank-agreed capability decision.
+
+The workspace is separate from the native MNNVL all-reduce workspace: post_moe_hop keeps
+its own multicast buffer, its own Lamport epoch and its own error word, so its
+call sequence never interleaves with another all-reduce's protocol state. One
+workspace serves every post_moe_hop of a TP group; hops run in the same order on
+every rank, which keeps the per-rank epochs in lockstep.
+"""
+
+import threading
+from dataclasses import dataclass
+from typing import ClassVar, Dict, List, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.mapping import Mapping
+
+from .kernels import MAX_CHUNKS, STATE_CALLS, STATE_ERROR_WORDS, STATE_TIMEOUT_POLLS, STATE_WORDS
+
+_BF16_NEG_ZERO_BITS = -32768
+
+
+def _workspace_comm(mapping: Mapping):
+    from tensorrt_llm._torch.distributed.ops import _get_mnnvl_workspace_comm
+
+    return _get_mnnvl_workspace_comm(mapping)
+
+
+def _all_ranks_agree(comm, local_ok: bool) -> bool:
+    from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded
+
+    return _mnnvl_workspace_all_succeeded(comm, local_ok)
+
+
+@dataclass(frozen=True)
+class _WorkspaceKey:
+    mapping: Mapping
+    hidden_size: int
+    max_rows: int
+
+
+class PostMoeHopWorkspace:
+    """Multicast Lamport buffers plus local partials/state for one TP group.
+
+    Attributes:
+        mc: bf16 ``[3 * max_rows * world * hidden]`` multicast view.
+        uc: bf16 view of the same region through this rank's unicast address.
+        partials: fp32 ``[3 * max_rows * MAX_CHUNKS]`` local row-sum exchange.
+        state: int32 ``[STATE_WORDS]`` epoch / previous rows / arrivals / errors /
+            calls / poll count at the latest timeout.
+        layout_repairs: host-side count of post_moe_hop calls whose raw tuple needed an
+            aligned copy (see ``PostMoeHop``); a CUDA graph capture counts
+            once, its replays do not.
+    """
+
+    _instances: ClassVar[Dict[_WorkspaceKey, "PostMoeHopWorkspace"]] = {}
+    _comms: ClassVar[Dict[Mapping, object]] = {}
+    _lock: ClassVar[threading.Lock] = threading.Lock()
+
+    def __init__(self, mapping: Mapping, hidden_size: int, max_rows: int, comm):
+        from tensorrt_llm._torch.distributed.ops import (
+            _make_mnnvl_mcast_buffer,
+            _mnnvl_device_index,
+        )
+
+        self.mapping = mapping
+        self.hidden_size = hidden_size
+        self.max_rows = max_rows
+        self.world = mapping.tp_size
+        self.comm = comm
+        self.lamport_elems = 3 * max_rows * self.world * hidden_size
+        nbytes = self.lamport_elems * torch.bfloat16.itemsize
+        self.handle = _make_mnnvl_mcast_buffer(comm, nbytes, mapping, mapping.is_multi_node())
+        self.uc = self.handle.get_uc_buffer(
+            mapping.tp_rank, (self.lamport_elems,), torch.bfloat16, 0
+        )
+        self.mc = self.handle.get_mc_buffer((self.lamport_elems,), torch.bfloat16, 0)
+        device = torch.device("cuda", _mnnvl_device_index(mapping))
+        self.partials = torch.full(
+            (3 * max_rows * MAX_CHUNKS,), -0.0, dtype=torch.float32, device=device
+        )
+        self.state = torch.zeros((STATE_WORDS,), dtype=torch.int32, device=device)
+        # The kernels access every buffer with 16-byte vectors.
+        for name in ("uc", "mc", "partials", "state"):
+            if getattr(self, name).data_ptr() % 16:
+                raise RuntimeError(f"post_moe_hop workspace buffer {name} is not 16-byte aligned")
+        self.layout_repairs = 0
+        self._reset_local()
+
+    @classmethod
+    def comm_for(cls, mapping: Mapping):
+        """Return the cached TP-group communicator (collective on first use)."""
+        with cls._lock:
+            comm = cls._comms.get(mapping)
+            if comm is None:
+                comm = _workspace_comm(mapping)
+                cls._comms[mapping] = comm
+            return comm
+
+    @classmethod
+    def instances(cls) -> List["PostMoeHopWorkspace"]:
+        """Every workspace created in this process (one per TP group and shape)."""
+        with cls._lock:
+            return [ws for ws in cls._instances.values() if ws is not None]
+
+    @classmethod
+    def get_or_create(
+        cls, mapping: Mapping, hidden_size: int, max_rows: int
+    ) -> Tuple[Optional["PostMoeHopWorkspace"], str]:
+        """Collectively allocate (once) and return the workspace, or ``None``.
+
+        Every rank of the TP group must call this in the same order. The
+        allocation either succeeds on all ranks or is discarded on all ranks.
+        """
+        key = _WorkspaceKey(mapping, hidden_size, max_rows)
+        with cls._lock:
+            existing = cls._instances.get(key)
+        if existing is not None:
+            return existing, ""
+        comm = cls.comm_for(mapping)
+        workspace = None
+        error: Optional[Exception] = None
+        try:
+            workspace = cls(mapping, hidden_size, max_rows, comm)
+        except Exception as exc:  # noqa: BLE001 - converged across ranks below
+            error = exc
+        if not _all_ranks_agree(comm, error is None):
+            return None, f"workspace allocation failed on some rank ({error!r})"
+        with cls._lock:
+            cls._instances[key] = workspace
+        return workspace, ""
+
+    def _reset_local(self) -> None:
+        with torch.inference_mode():
+            self.uc.view(torch.int16).fill_(_BF16_NEG_ZERO_BITS)
+            self.partials.fill_(-0.0)
+            self.state.zero_()
+        torch.cuda.synchronize()
+
+    def reset(self) -> bool:
+        """Collectively return every rank to epoch 0 with empty buffers.
+
+        Needed only after an error bit was raised, because a timed-out call
+        leaves the ranks' epochs out of step. Returns True when every rank reset.
+        """
+        ok = True
+        try:
+            torch.cuda.synchronize()
+        except RuntimeError:
+            ok = False
+        # Every rank's in-flight multicast stores have landed once all ranks
+        # drained their streams; only then may a rank refill its buffers.
+        if not _all_ranks_agree(self.comm, ok):
+            return False
+        try:
+            self._reset_local()
+        except RuntimeError:
+            ok = False
+        return _all_ranks_agree(self.comm, ok)
+
+    def error_bits(self) -> int:
+        """Return this rank's sticky error bits (synchronizes the device)."""
+        return self.counters()["error_bits"]
+
+    def calls(self) -> int:
+        """Return how many hop launches completed on this rank (synchronizes)."""
+        return int(self.state[STATE_CALLS].item())
+
+    def counters(self) -> Dict[str, int]:
+        """Return the state words of this rank in one copy (synchronizes the device).
+
+        Keys: ``calls`` (completed hop launches since the last reset),
+        ``error_bits`` (sticky; bit 1 Lamport timeout, 2 row-sum timeout,
+        4 permutation index out of range), ``timeout_polls`` (largest poll
+        count at a timeout) and the host-side ``layout_repairs``.
+        """
+        torch.cuda.synchronize(self.state.device)
+        words = self.state.cpu().tolist()
+        return {
+            "calls": words[STATE_CALLS],
+            "error_bits": words[STATE_ERROR_WORDS[0]]
+            | words[STATE_ERROR_WORDS[1]]
+            | words[STATE_ERROR_WORDS[2]],
+            "timeout_polls": words[STATE_TIMEOUT_POLLS],
+            "layout_repairs": self.layout_repairs,
+        }
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/__init__.py
new file mode 100644
index 000000000000..db16aad81405
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/__init__.py
@@ -0,0 +1,70 @@
+# 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.
+"""rope_kv_append: fused per-layer RoPE / FP8 query / KV append, indexer K-cache scatter
+and top-k -> global KV index bridge for GLM-5.x DSA decode."""
+
+from ... import dispatch
+from .ops import (
+    BridgeCopy,
+    BridgePacket,
+    IndexerKScatter,
+    RopeKvAppendContractError,
+    bridge_eligible,
+    check_indexer_k_scatter_consumed,
+    defer_indexer_k_scatter,
+    defer_mtp_row_stash,
+    defer_reuse_copy,
+    defer_shared_copy,
+    mla_rope_generation,
+    new_bridge_packet,
+    remap_topk_to_global,
+    rope_eligible,
+    run_indexer_k_scatter_native,
+    topk_bridge,
+)
+
+_ENTRY_POINTS = {
+    "rope_eligible": rope_eligible,
+    "defer_indexer_k_scatter": defer_indexer_k_scatter,
+    "check_indexer_k_scatter_consumed": check_indexer_k_scatter_consumed,
+    "mla_rope_generation": mla_rope_generation,
+    "run_indexer_k_scatter_native": run_indexer_k_scatter_native,
+    "new_bridge_packet": new_bridge_packet,
+    "defer_shared_copy": defer_shared_copy,
+    "defer_reuse_copy": defer_reuse_copy,
+    "defer_mtp_row_stash": defer_mtp_row_stash,
+    "remap_topk_to_global": remap_topk_to_global,
+}
+for _name, _fn in _ENTRY_POINTS.items():
+    dispatch.register("rope_kv_append", _name, _fn)
+
+__all__ = [
+    "BridgeCopy",
+    "BridgePacket",
+    "IndexerKScatter",
+    "RopeKvAppendContractError",
+    "bridge_eligible",
+    "check_indexer_k_scatter_consumed",
+    "defer_indexer_k_scatter",
+    "defer_mtp_row_stash",
+    "defer_reuse_copy",
+    "defer_shared_copy",
+    "mla_rope_generation",
+    "new_bridge_packet",
+    "remap_topk_to_global",
+    "rope_eligible",
+    "run_indexer_k_scatter_native",
+    "topk_bridge",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/ops.py
new file mode 100644
index 000000000000..3a93b9e7e577
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/ops.py
@@ -0,0 +1,517 @@
+# 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.
+"""rope_kv_append: fused per-layer MLA RoPE / FP8 query / latent KV append, DSA indexer
+K-cache scatter, and top-k id -> global KV index bridge.
+
+Two CUDA ops back this package (``cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu``):
+
+* ``trtllm::glm_moe_dsa_rope_kv_append`` writes exactly the bytes of the
+  native ``trtllm::mla_rope_generation`` (FP8 MLA cache, absorption mode, no
+  helix) and, when given the indexer inputs, of
+  ``trtllm::indexer_k_cache_scatter_op`` in one launch.
+* ``trtllm::glm_moe_dsa_topk_index_bridge`` writes the values of
+  ``trtllm::convert_req_index_to_global`` (or its grouped variant) and the
+  native top-k reuse copy in one launch.
+
+Every helper here checks its support domain before any device work; a caller
+that gets ``None``/``False`` runs the native path unchanged.
+
+Memory order: both kernels may launch early as programmatic dependents (PDL).
+They read producer-written data and per-step metadata only after
+``griddepcontrol.wait``; only the static quantization scales are read before
+it.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+from tensorrt_llm._utils import get_sm_version
+
+from ... import dispatch
+
+FEATURE_ID = "rope_kv_append"
+
+_KV_LORA_RANK = 512
+_ROPE_DIM = 64
+
+
+@dataclass
+class IndexerKScatter:
+    """Deferred native indexer K-cache scatter of one layer and one forward.
+
+    Single use: exactly one of the fused generation rope
+    (:func:`mla_rope_generation`) or the native fallback
+    (:func:`run_indexer_k_scatter_native`) consumes it, before any reader of
+    the indexer K cache runs. The call site that deferred the scatter checks
+    with :func:`check_indexer_k_scatter_consumed` right after the generation
+    call, so a packet that did not reach its consumer fails the first forward
+    (the CUDA-graph warmup) instead of silently leaving the step's keys out of
+    the cache.
+    """
+
+    k_fp8: torch.Tensor
+    k_scale: torch.Tensor
+    k_cache: torch.Tensor
+    slot_mapping_fp8: torch.Tensor
+    slot_mapping_scale: torch.Tensor
+    num_tokens: int
+    layer_idx: int = -1
+    consumer: Optional[str] = None
+
+
+class RopeKvAppendContractError(RuntimeError):
+    """A rope_kv_append call-site contract was violated (lost or doubly consumed deferred work)."""
+
+
+def _consume(indexer_k: IndexerKScatter, consumer: str) -> None:
+    if indexer_k.consumer is not None:
+        raise RopeKvAppendContractError(
+            f"rope_kv_append deferred indexer K-cache scatter of layer {indexer_k.layer_idx} consumed twice "
+            f"({indexer_k.consumer!r}, then {consumer!r})"
+        )
+    indexer_k.consumer = consumer
+
+
+@dataclass
+class BridgeCopy:
+    """Native top-k reuse copy folded into the bridge.
+
+    mode 0: ``dst[row_start + i] = topk[i]`` for ``i < rows``.
+    mode 1: ``dst[row_start + i] = topk[i * next_n + clamp(num_accepted[i] - 1)]``
+    (the MTP last-accepted row of request ``i``).
+    """
+
+    dst: torch.Tensor
+    row_start: int
+    rows: int
+    mode: int = 0
+    next_n: int = 1
+    num_accepted: Optional[torch.Tensor] = None
+
+
+_MAX_GROUP_SIZE = 8
+
+# The rope_kv_append kernels read per-step metadata (cache lengths, KV block offsets,
+# indexer slot mappings, request indices, MTP accepted counts) and every
+# producer-written input only after ``griddepcontrol.wait``; only the static
+# quantization scales are read earlier.
+
+
+def _sm_ok() -> bool:
+    return get_sm_version() >= 100
+
+
+def _static_rope_domain(mqa, metadata) -> bool:
+    """Module/metadata conditions of the fused rope that do not depend on the call tensors."""
+    kvm = metadata.kv_cache_manager
+    return bool(
+        kvm is not None
+        and getattr(mqa, "is_mla_enable", False)
+        and not mqa.uses_fp4_mla_attention
+        and getattr(mqa, "has_fp8_kv_cache", False)
+        and not getattr(mqa, "has_fp4_kv_cache", False)
+        and mqa.kv_lora_rank == _KV_LORA_RANK
+        and mqa.qk_rope_head_dim == _ROPE_DIM
+        and bool(mqa.rope_append)
+        and mqa.num_kv_heads == 1
+        and getattr(kvm, "mla_kv_cache_residual_dim", 0) == 0
+        and metadata.beam_width == 1
+        and metadata.num_contexts == 0
+        and metadata.num_generations > 0
+        and getattr(metadata, "helix_position_offsets", None) is None
+        and getattr(metadata, "helix_is_inactive_rank", None) is None
+        and not getattr(metadata, "_helix_spec_tokens_valid", False)
+        and metadata.token_major_gen_view() is None
+        and not getattr(metadata, "is_ragged_verify", False)
+        and mqa.rotary_cos_sin is not None
+        and mqa.rotary_cos_sin.dtype == torch.float32
+        and mqa.rotary_cos_sin.is_contiguous()
+    )
+
+
+def rope_eligible(
+    mqa, metadata, fused_q, q_pe, latent_cache, quant_q_buffer, mla_bmm1_scale, mla_bmm2_scale
+) -> bool:
+    """Whether the fused rope can replace ``mqa.mla_rope_generation`` for this call."""
+    if not dispatch.is_enabled(FEATURE_ID) or not _sm_ok():
+        return False
+    ok = (
+        _static_rope_domain(mqa, metadata)
+        and fused_q is not None
+        and q_pe is not None
+        and quant_q_buffer is not None
+        and mla_bmm1_scale is not None
+        and mla_bmm2_scale is not None
+        and fused_q.dtype == torch.bfloat16
+        and q_pe.dtype == torch.bfloat16
+        and latent_cache.dtype == torch.bfloat16
+        and latent_cache.dim() == 2
+        and latent_cache.stride(-1) == 1
+        and latent_cache.stride(0) == _KV_LORA_RANK + _ROPE_DIM
+        and fused_q.is_contiguous()
+    )
+    return dispatch.log_capability(
+        FEATURE_ID, "rope", bool(ok), "" if ok else "outside the FP8 MLA decode domain"
+    )
+
+
+def defer_indexer_k_scatter(
+    mla, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata, *, forwards_packet: bool = False
+) -> Optional[IndexerKScatter]:
+    """Describe the native ``Indexer._update_k_cache`` call without running it.
+
+    The scatter is deferred into the generation rope of the same layer, which
+    still precedes every reader of the indexer K cache (the indexer logits).
+    The caller must forward the returned packet to
+    ``forward_absorption_generation(deferred_indexer_k=...)`` and must call
+    :func:`check_indexer_k_scatter_consumed` after that call; it declares both
+    with ``forwards_packet=True``. A caller that does not declare it gets None
+    and runs the native scatter.
+
+    Returns None when the native scatter must run now: rope_kv_append off, a caller that
+    does not forward the packet, a context or mixed batch, or a call outside
+    the fused rope domain.
+    """
+    if not dispatch.is_enabled(FEATURE_ID) or not _sm_ok():
+        return None
+    if not forwards_packet:
+        dispatch.log_capability(
+            FEATURE_ID, "scatter_defer", False, "call site does not forward the deferred scatter"
+        )
+        return None
+    mqa = mla.mqa
+    if (
+        metadata.num_contexts != 0
+        or metadata.num_generations == 0
+        or mla.apply_rotary_emb
+        or getattr(mla, "kv_cache_dtype", None) == "fp8_ds_mla"
+        or not _static_rope_domain(mqa, metadata)
+    ):
+        return None
+    indexer = mqa.indexer
+    if metadata.kv_cache_manager is None or metadata.slot_mapping_fp8 is None:
+        return None
+    k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(indexer.layer_idx)
+    if k_scale.element_size() == 1:
+        k_scale = k_scale.view(torch.int32)
+    if (
+        k_fp8.dim() != 2
+        or k_fp8.element_size() != 1
+        or not k_fp8.is_contiguous()
+        or k_scale.dim() != 2
+        or k_scale.element_size() * k_scale.shape[1] != 4
+        or not k_scale.is_contiguous()
+        or k_fp8.shape[1] not in (64, 128)
+    ):
+        return None
+    return IndexerKScatter(
+        k_fp8=k_fp8,
+        k_scale=k_scale,
+        k_cache=k_cache,
+        slot_mapping_fp8=metadata.slot_mapping_fp8,
+        slot_mapping_scale=metadata.slot_mapping_scale,
+        num_tokens=int(k_fp8.shape[0]),
+        layer_idx=int(indexer.layer_idx),
+    )
+
+
+def check_indexer_k_scatter_consumed(indexer_k: Optional[IndexerKScatter]) -> None:
+    """Raise unless a deferred scatter was executed by the generation call.
+
+    Called by the deferring call site right after the generation call. A
+    packet that never reached the rope means the native scatter was skipped
+    and the indexer read stale keys; there is no correct recovery afterwards.
+    """
+    if indexer_k is not None and indexer_k.consumer is None:
+        raise RopeKvAppendContractError(
+            f"rope_kv_append deferred indexer K-cache scatter of layer {indexer_k.layer_idx} was not consumed: "
+            "the call site skipped the native Indexer._update_k_cache but did not pass the packet "
+            "to the generation rope (forward_absorption_generation(deferred_indexer_k=...))"
+        )
+
+
+def mla_rope_generation(
+    mqa,
+    fused_q,
+    q_pe,
+    latent_cache,
+    metadata,
+    cu_q_seqlens,
+    cu_kv_seqlens,
+    fmha_scheduler_counter,
+    mla_bmm1_scale,
+    mla_bmm2_scale,
+    quant_q_buffer,
+    indexer_k: Optional[IndexerKScatter] = None,
+) -> None:
+    """Fused replacement of ``mqa.mla_rope_generation`` (+ the deferred indexer scatter).
+
+    Callers check :func:`rope_eligible` first.
+    """
+    kvm = metadata.kv_cache_manager
+    mqa._ensure_rope_table_size(metadata.max_seq_len)
+    idx_args: List[Optional[torch.Tensor]] = [None, None, None, None, None]
+    idx_tokens = 0
+    if indexer_k is not None:
+        _consume(indexer_k, "fused_rope")
+        if indexer_k.num_tokens > 0:
+            idx_args = [
+                indexer_k.k_fp8,
+                indexer_k.k_scale,
+                indexer_k.k_cache,
+                indexer_k.slot_mapping_fp8,
+                indexer_k.slot_mapping_scale,
+            ]
+            idx_tokens = indexer_k.num_tokens
+    torch.ops.trtllm.glm_moe_dsa_rope_kv_append(
+        fused_q,
+        q_pe,
+        latent_cache,
+        mqa.rotary_cos_sin,
+        cu_q_seqlens,
+        cu_kv_seqlens,
+        fmha_scheduler_counter,
+        mla_bmm1_scale,
+        mla_bmm2_scale,
+        quant_q_buffer,
+        metadata.kv_lens_cuda_runtime,
+        metadata.prompt_lens_cpu_runtime,
+        metadata.num_contexts,
+        metadata.kv_cache_block_offsets,
+        kvm.kv_cache_pool_pointers,
+        kvm.kv_cache_pool_mapping,
+        getattr(mqa, "_nvfp4_mla_kv_scale_orig_quant", None),
+        mqa.get_local_layer_idx(metadata),
+        mqa.num_heads,
+        kvm.tokens_per_block,
+        metadata.max_seq_len,
+        mqa.quant_mode,
+        mqa.q_scaling,
+        mqa.qk_nope_head_dim,
+        mqa.qk_rope_head_dim,
+        mqa.kv_lora_rank,
+        *idx_args,
+        idx_tokens,
+    )
+
+
+def run_indexer_k_scatter_native(indexer_k: IndexerKScatter) -> None:
+    """Native fallback for a deferred scatter (same op and arguments as ``_update_k_cache``)."""
+    _consume(indexer_k, "native")
+    torch.ops.trtllm.indexer_k_cache_scatter_op(
+        indexer_k.k_fp8,
+        indexer_k.k_scale,
+        indexer_k.k_cache,
+        indexer_k.slot_mapping_fp8,
+        indexer_k.slot_mapping_scale,
+        indexer_k.num_tokens,
+    )
+
+
+@dataclass
+class BridgePacket:
+    """Per-call handoff between the indexer and the fused bridge of one layer.
+
+    ``source`` replaces the rows the remap reads (draft reuse reads the shared
+    rows directly); ``copy`` is the one native top-k copy the bridge performs.
+    """
+
+    source: Optional[torch.Tensor] = None
+    copy: Optional[BridgeCopy] = None
+
+
+def new_bridge_packet(metadata, is_generation: bool) -> Optional[BridgePacket]:
+    """A packet when the fused bridge serves this ``sparse_attn_predict`` call, else None."""
+    if not bridge_eligible(metadata, is_generation):
+        return None
+    return BridgePacket()
+
+
+def bridge_eligible(metadata, is_generation: bool) -> bool:
+    """Whether the fused bridge can replace the native remap for this call."""
+    if not dispatch.is_enabled(FEATURE_ID) or not _sm_ok():
+        return False
+    import tensorrt_llm.bindings
+
+    kvm = metadata.kv_cache_manager
+    ok = (
+        is_generation
+        and kvm is not None
+        and getattr(kvm, "dtype", None) != tensorrt_llm.bindings.DataType.NVFP4
+        and metadata.num_contexts == 0
+        and metadata.num_generations > 0
+    )
+    return dispatch.log_capability(
+        FEATURE_ID, "bridge", bool(ok), "" if ok else "outside the generation remap domain"
+    )
+
+
+def topk_bridge(
+    topk: torch.Tensor,
+    metadata,
+    num_rows: int,
+    *,
+    layer_offset: Optional[int] = None,
+    layer_ids: Optional[torch.Tensor] = None,
+    page_index_scale: int,
+    copy: Optional[BridgeCopy] = None,
+) -> torch.Tensor:
+    """Remap local top-k rows to global KV-pool indices (+ the native reuse copy).
+
+    Exactly one of ``layer_offset`` (per-layer remap, returns ``[rows, k]``) or
+    ``layer_ids`` (grouped remap, returns ``[group, rows, k]``) is given.
+    """
+    # Same inputs as ``_ensure_pool_view_cached`` derives for the native remap;
+    # the ``req_idx - num_contexts`` subtraction happens inside the kernel.
+    # Unlike the native remap this does not populate the metadata's pool-view
+    # memo (``_pool_cache_valid``/``_cached_*``). Every reader of that memo calls
+    # ``_ensure_pool_view_cached()`` itself, so a later reader in the same step
+    # populates it (and launches the subtraction) on demand; the subtraction is
+    # saved only in steps where no other reader needs the memo.
+    tokens_per_block = metadata.kv_cache_manager.tokens_per_block
+    block_table = metadata.block_table[metadata.num_contexts : metadata.num_seqs]
+    req_idx = metadata.req_idx_per_token[metadata.num_ctx_tokens : metadata.num_tokens]
+    return torch.ops.trtllm.glm_moe_dsa_topk_index_bridge(
+        topk,
+        num_rows,
+        req_idx,
+        metadata.num_contexts,
+        block_table,
+        tokens_per_block,
+        page_index_scale * tokens_per_block,
+        layer_ids,
+        0 if layer_offset is None else layer_offset,
+        None if copy is None else copy.dst,
+        0 if copy is None else copy.row_start,
+        0 if copy is None else copy.rows,
+        0 if copy is None else copy.mode,
+        1 if copy is None else copy.next_n,
+        None if copy is None else copy.num_accepted,
+    )
+
+
+def defer_shared_copy(
+    packet: BridgePacket, shared: torch.Tensor, row_start: int, topk: torch.Tensor
+) -> None:
+    """``shared[row_start:row_start + rows, :k].copy_(topk)`` performed by the bridge."""
+    assert packet.copy is None
+    packet.copy = BridgeCopy(dst=shared, row_start=row_start, rows=int(topk.shape[0]))
+
+
+def defer_reuse_copy(
+    packet: Optional[BridgePacket],
+    topk_buffer: torch.Tensor,
+    row_start: int,
+    shared_rows: torch.Tensor,
+    gvr_prior_active: bool,
+) -> bool:
+    """Draft reuse ``topk_buffer[row_start:+rows] = shared_rows``; True when deferred."""
+    if (
+        packet is None
+        or gvr_prior_active
+        or packet.copy is not None
+        or shared_rows.shape[1] != topk_buffer.shape[1]
+    ):
+        return False
+    packet.source = shared_rows
+    packet.copy = BridgeCopy(dst=topk_buffer, row_start=row_start, rows=int(shared_rows.shape[0]))
+    return True
+
+
+def defer_mtp_row_stash(
+    packet: Optional[BridgePacket],
+    shared: Optional[torch.Tensor],
+    row_start: int,
+    num_generations: int,
+    next_n: int,
+    num_accepted: Optional[torch.Tensor],
+) -> bool:
+    """MTP draft step 0 last-accepted-row stash into ``shared``; True when deferred."""
+    if packet is None or shared is None or packet.copy is not None:
+        return False
+    if num_accepted is not None and num_accepted.dtype not in (torch.int32, torch.int64):
+        return False
+    packet.copy = BridgeCopy(
+        dst=shared,
+        row_start=row_start,
+        rows=num_generations,
+        mode=1,
+        next_n=next_n,
+        num_accepted=num_accepted,
+    )
+    return True
+
+
+def remap_topk_to_global(
+    backend, topk: torch.Tensor, metadata, packet: BridgePacket
+) -> torch.Tensor:
+    """Fused ``DSATrtllmAttention._remap_topk_to_global`` (+ the packet's copy).
+
+    Follows the native grouping rules: the leader of an active full+shared
+    group computes the whole group in one launch and publishes it in
+    ``metadata._group_remap_batched``; followers slice it; everything else
+    remaps per layer.
+    """
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
+
+    src = topk if packet.source is None else packet.source
+    num_rows = int(src.shape[0])
+    copy = packet.copy
+    local_layer_idx = backend.get_local_layer_idx(metadata)
+    if dsa_backend._GROUP_REMAP and not metadata.in_mtp_draft_loop:
+        struct = metadata._ensure_group_remap_struct()
+        leader_of = struct.get("leader_of")
+        if leader_of is not None and local_layer_idx < len(leader_of):
+            leader = leader_of[local_layer_idx]
+            if leader >= 0 and struct["group_active"].get(leader, False):
+                slot = struct["slot_of"][local_layer_idx]
+                if backend.indexer is not None:
+                    if struct["group_layer_ids"][leader].numel() > _MAX_GROUP_SIZE:
+                        if copy is not None:
+                            copy.dst[
+                                copy.row_start : copy.row_start + copy.rows, : src.shape[1]
+                            ].copy_(src)
+                        return backend._remap_topk_to_global(topk, metadata, True)
+                    batched = topk_bridge(
+                        src,
+                        metadata,
+                        num_rows,
+                        layer_ids=struct["group_layer_ids"][leader],
+                        page_index_scale=struct["group_scale"][leader],
+                        copy=copy,
+                    )
+                    metadata._group_remap_batched[leader] = batched
+                    return batched[slot]
+                batched = metadata._group_remap_batched.get(leader, None)
+                if (
+                    copy is None
+                    and batched is not None
+                    and slot < batched.shape[0]
+                    and batched.shape[1:] == src.shape
+                ):
+                    return batched[slot]
+    page_index_scale, layer_offset = metadata.kv_cache_manager.get_primary_pool_page_index_params(
+        local_layer_idx
+    )
+    return topk_bridge(
+        src,
+        metadata,
+        num_rows,
+        layer_offset=layer_offset,
+        page_index_scale=page_index_scale,
+        copy=copy,
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/__init__.py
new file mode 100644
index 000000000000..fa39194543a2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/__init__.py
@@ -0,0 +1,40 @@
+# 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.
+"""Package router_gemm: pre-MoE reduction/norm/router boundary (the router contract).
+
+Registered entry points (looked up with ``dispatch.get_impl("router_gemm", name)``):
+
+* ``router_logits(hidden, gate_weight) -> Tensor | None``
+* ``pre_moe_hop(allreduce, hidden, residual, norm_weight, eps, gate_weight)
+  -> (hidden, residual, Tensor | None)``
+* ``make_moe_router(gate, layer_idx, model_config) -> Callable | None``
+"""
+
+from ... import dispatch
+from .hop import engagement_counts, make_moe_router, pre_moe_hop, router_logits
+from .router_kernel import router_gemm, router_gemm_supported
+
+dispatch.register("router_gemm", "router_logits", router_logits)
+dispatch.register("router_gemm", "pre_moe_hop", pre_moe_hop)
+dispatch.register("router_gemm", "make_moe_router", make_moe_router)
+
+__all__ = [
+    "engagement_counts",
+    "make_moe_router",
+    "pre_moe_hop",
+    "router_gemm",
+    "router_gemm_supported",
+    "router_logits",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/hop.py
new file mode 100644
index 000000000000..793f335d5935
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/hop.py
@@ -0,0 +1,146 @@
+# 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.
+"""The router contract: the pre-MoE reduction/norm/router hop of package router_gemm.
+
+The hop keeps the native all-reduce + residual + RMSNorm (one fused native
+kernel, unchanged) and replaces only the native router GEMM that consumes its
+normalized output. The router logits are bitwise identical to the native
+``DeepseekV3Gate`` logits, so expert selection, which stays in the native MoE,
+is unchanged.
+
+Entry points (registered in ``router_gemm/__init__.py``):
+
+* ``router_logits(hidden, gate_weight)``: router_gemm logits or ``None`` when the call
+  is outside the support domain; on ``None`` the caller runs the native gate.
+* ``pre_moe_hop(...)``: native fused AR/residual/RMSNorm followed by
+  ``router_logits``; returns ``(hidden, residual, logits_or_None)``.
+* ``make_moe_router(gate, layer_idx, model_config)``: per-layer capability
+  decision made once at model construction; returns a callable for
+  ``Deepseekv3MoE.compute_routed_output`` or ``None`` (native gate).
+"""
+
+from typing import Callable, Dict, Optional, Tuple
+
+import torch
+
+from ... import dispatch
+from .router_kernel import pdl_enabled, router_gemm, router_gemm_supported
+
+FEATURE_ID = "router_gemm"
+
+# Host-side call counters for dispatch-engagement reporting. They count Python
+# calls (eager runs and CUDA graph captures), not graph replays.
+_counts: Dict[str, int] = {"engaged": 0, "native_fallback": 0}
+
+
+def engagement_counts() -> Dict[str, int]:
+    """Copy of the host-side engaged/fallback call counters."""
+    return dict(_counts)
+
+
+def router_logits(hidden_states: torch.Tensor, gate_weight: torch.Tensor) -> Optional[torch.Tensor]:
+    """FP32 router logits ``[R, E]`` or ``None`` (caller uses native gate).
+
+    Ineligible calls return before launching anything, so the native fallback
+    happens before any side effect.
+    """
+    eligible, _ = router_gemm_supported(hidden_states, gate_weight)
+    if not eligible:
+        _counts["native_fallback"] += 1
+        return None
+    _counts["engaged"] += 1
+    return router_gemm(hidden_states, gate_weight)
+
+
+def pre_moe_hop(
+    allreduce,
+    hidden_states: torch.Tensor,
+    residual: torch.Tensor,
+    norm_weight: torch.Tensor,
+    eps: float,
+    gate_weight: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Reduce -> residual -> RMSNorm (native) -> router (router_gemm).
+
+    Args:
+        allreduce: the layer's native ``AllReduce`` module.
+        hidden_states: rank-local attention output partial ``[R, H]``.
+        residual: residual stream ``[R, H]``.
+        norm_weight, eps: ``post_attention_layernorm`` parameters.
+        gate_weight: router weight ``[E, H]`` (BF16).
+
+    Returns:
+        ``(normed_hidden, residual, router_logits_or_None)``; the first two
+        are exactly the native fused all-reduce outputs.
+    """
+    from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams
+
+    hidden_states, residual = allreduce(
+        hidden_states,
+        all_reduce_params=AllReduceParams(
+            fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+            residual=residual,
+            norm_weight=norm_weight,
+            eps=eps,
+            trigger_completion_at_end=False,
+        ),
+    )
+    return hidden_states, residual, router_logits(hidden_states, gate_weight)
+
+
+def _layer_domain(gate, layer_idx: Optional[int], model_config) -> str:
+    """Return '' when the layer is in the support domain, else the reason."""
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate
+    from tensorrt_llm._utils import get_sm_version, is_sm_100f
+
+    if type(gate) is not DeepseekV3Gate:
+        return f"gate class {type(gate).__name__}"
+    mapping = model_config.mapping
+    if mapping.enable_attention_dp:
+        return "attention DP"
+    if mapping.tp_size <= 1:
+        return "no tensor parallelism"
+    num_target_layers = model_config.pretrained_config.num_hidden_layers
+    if layer_idx is None or layer_idx >= num_target_layers:
+        return "not a target layer"
+    if gate.use_cute_dsl_bf16_gemm and is_sm_100f():
+        return "native gate uses the CuTe DSL BF16 GEMM"
+    if not 100 <= get_sm_version() < 110:
+        return f"sm{get_sm_version()}"
+    return ""
+
+
+def make_moe_router(gate, layer_idx: Optional[int], model_config) -> Optional[Callable]:
+    """Per-layer router factory used by the ``Deepseekv3MoE`` adapter.
+
+    Returns ``None`` when the layer is outside the support domain (the adapter
+    then never calls router_gemm for it), otherwise ``fn(hidden) -> logits or None``.
+    The decision is the same on every rank: it depends only on the model and
+    mapping configuration.
+    """
+    reason = _layer_domain(gate, layer_idx, model_config)
+    dispatch.log_capability(
+        FEATURE_ID,
+        f"router:{reason or 'engaged'}",
+        not reason,
+        reason or f"first layer {layer_idx}, pdl={int(pdl_enabled())}",
+    )
+    if reason:
+        return None
+
+    def _route(hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
+        return router_logits(hidden_states, gate.weight)
+
+    return _route
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/router_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/router_kernel.py
new file mode 100644
index 000000000000..3c411e62fc9c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/router_kernel.py
@@ -0,0 +1,571 @@
+# 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.
+"""Small-M MoE router GEMM for the GLM pre-MoE hop (the router contract).
+
+``router_gemm(hidden, weight)`` returns ``hidden @ weight.T`` as FP32 logits
+``[R, E]`` for BF16 ``hidden [R, K]`` and BF16 ``weight [E, K]``.
+
+Every logit is computed with exactly the floating-point operations of the
+native ``trtllm::dsv3_router_gemm_op`` small-M kernel, in the same order, so
+the result is bitwise identical to it:
+
+* the 128 threads that own one expert split K the native way: thread ``t``
+  owns elements ``k = i * 1024 + t * 8 + j`` (``i < K / 1024``, ``j < 8``) and
+  accumulates their BF16 x BF16 products into an FP32 register in ``(i, j)``
+  order, one FMA each (the products are exact in FP32);
+* the per-thread partials are combined by the XOR butterfly 16, 8, 4, 2, 1
+  inside each warp;
+* the four warp sums are added in warp order onto 0.0.
+
+What differs is the data movement and the schedule:
+
+* Operands are staged in shared memory with ``cp.async``. The expert-weight
+  slice does not depend on the preceding kernel, so its copies are issued
+  before the programmatic grid dependency wait; under PDL the weight fetch
+  overlaps the tail of the all-reduce/RMSNorm producer. (Register loads issued
+  before ``griddepcontrol.wait`` are not kept there by ptxas, which schedules
+  the wait first; asynchronous copies have no register results to sink.)
+* With one expert per CTA (the default) every thread stages exactly the
+  activation vectors it consumes, one ``cp.async`` group per K step, and
+  computes step ``i`` while later steps are still landing. With several
+  experts per CTA the CTA stages the activation block once for all of them.
+* Dependents are released as soon as the activation copies are issued; they
+  still wait for this grid to complete before reading the logits.
+
+The CuTe DSL is imported lazily so that this module can be imported (and its
+eligibility checks used) where the DSL is absent.
+"""
+
+import os
+import threading
+from typing import Dict, Optional, Tuple
+
+import torch
+
+# Geometry of the native kernel whose arithmetic is reproduced.
+THREADS_PER_EXPERT = 128
+ELEMS_PER_VECTOR = 8  # one 16-byte vector of BF16
+K_PER_ITERATION = THREADS_PER_EXPERT * ELEMS_PER_VECTOR  # 1024
+WARPS_PER_EXPERT = THREADS_PER_EXPERT // 32
+
+NUM_EXPERTS = 256
+SUPPORTED_HIDDEN_SIZES = (4096, 6144, 7168)
+MAX_ROWS = 16
+
+# Dynamic shared memory available to one CTA on SM100-family GPUs.
+_MAX_SMEM_BYTES = 227 * 1024
+_SMEM_RESERVE_BYTES = 4 * 1024  # warp-sum scratch and alignment slack
+
+# Launch shape of this target: one expert per CTA, single-element FMAs.
+DEFAULT_EXPERTS_PER_CTA = 1
+DEFAULT_PAIRED_FMA = False
+
+
+def pdl_enabled() -> bool:
+    """Whether kernels launch with programmatic dependent launch."""
+    return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1"
+
+
+def router_gemm_supported(hidden: torch.Tensor, weight: torch.Tensor) -> Tuple[bool, str]:
+    """Return ``(eligible, reason)`` for ``router_gemm(hidden, weight)``.
+
+    Checks only tensor metadata, never device data, so it is safe inside CUDA
+    graph capture and launches nothing.
+    """
+    if not (hidden.is_cuda and weight.is_cuda):
+        return False, "tensors not on CUDA"
+    if hidden.device != weight.device:
+        return False, "device mismatch"
+    if hidden.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
+        return False, f"dtype {hidden.dtype}/{weight.dtype} is not bf16/bf16"
+    if hidden.dim() != 2 or weight.dim() != 2:
+        return False, "expects 2-D hidden and weight"
+    rows, hidden_size = hidden.shape
+    if not 1 <= rows <= MAX_ROWS:
+        return False, f"rows={rows} outside [1, {MAX_ROWS}]"
+    if tuple(weight.shape) != (NUM_EXPERTS, hidden_size):
+        return False, (f"weight shape {tuple(weight.shape)} is not [{NUM_EXPERTS}, {hidden_size}]")
+    if hidden_size not in SUPPORTED_HIDDEN_SIZES:
+        return False, f"hidden size {hidden_size} unsupported"
+    if not (hidden.is_contiguous() and weight.is_contiguous()):
+        return False, "non-contiguous operand"
+    if hidden.data_ptr() % 16 or weight.data_ptr() % 16:
+        return False, "operand not 16-byte aligned"
+    if _smem_bytes(rows, hidden_size, 1) > _MAX_SMEM_BYTES:
+        return False, f"[{rows}, {hidden_size}] activation exceeds shared memory"
+    major, minor = torch.cuda.get_device_capability(hidden.device)
+    if not 100 <= major * 10 + minor < 110:
+        return False, f"sm{major}{minor} outside the validated SM100 family"
+    return True, ""
+
+
+def _smem_bytes(num_rows: int, hidden_size: int, experts_per_cta: int) -> int:
+    """Shared memory of one CTA: activation block, weight slices, scratch."""
+    return (num_rows + experts_per_cta) * hidden_size * 2 + _SMEM_RESERVE_BYTES
+
+
+def fit_experts_per_cta(num_rows: int, hidden_size: int, requested: int) -> int:
+    """Largest power-of-two experts-per-CTA <= requested whose staging fits."""
+    epc = requested
+    while epc > 1 and _smem_bytes(num_rows, hidden_size, epc) > _MAX_SMEM_BYTES:
+        epc //= 2
+    return epc
+
+
+_compile_lock = threading.Lock()
+_compiled: Dict[tuple, object] = {}
+
+
+def _build_kernel_class():
+    """Define the CuTe DSL kernel; imported only on first compile."""
+    import cutlass
+    import cutlass.cute as cute
+    from cuda.bindings.driver import CUstream
+    from cutlass import BFloat16, Float32, Int32, Int64
+    from cutlass._mlir.dialects import llvm
+    from cutlass.cutlass_dsl import T, dsl_user_op
+    from cutlass.utils.smem_allocator import SmemAllocator
+
+    def _asm(res, operands, text, constraints, side_effects, loc, ip):
+        return llvm.inline_asm(
+            res,
+            operands,
+            text,
+            constraints,
+            has_side_effects=side_effects,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def griddep_wait(*, loc=None, ip=None) -> None:
+        _asm(None, [], "griddepcontrol.wait;", "~{memory}", True, loc, ip)
+
+    @dsl_user_op
+    def griddep_launch_dependents(*, loc=None, ip=None) -> None:
+        _asm(None, [], "griddepcontrol.launch_dependents;", "", True, loc, ip)
+
+    @dsl_user_op
+    def cp_async_16(smem_addr, gmem_addr, *, loc=None, ip=None) -> None:
+        # .cg: through L2 only; every vector is staged once per CTA.
+        _asm(
+            None,
+            [Int32(smem_addr).ir_value(loc=loc, ip=ip), Int64(gmem_addr).ir_value(loc=loc, ip=ip)],
+            "cp.async.cg.shared.global [$0], [$1], 16;",
+            "r,l,~{memory}",
+            True,
+            loc,
+            ip,
+        )
+
+    @dsl_user_op
+    def cp_async_commit(*, loc=None, ip=None) -> None:
+        _asm(None, [], "cp.async.commit_group;", "~{memory}", True, loc, ip)
+
+    @dsl_user_op
+    def cp_async_wait_all(*, loc=None, ip=None) -> None:
+        _asm(None, [], "cp.async.wait_group 0;", "~{memory}", True, loc, ip)
+
+    @dsl_user_op
+    def lds_v4(smem_addr, *, loc=None, ip=None):
+        ret = _asm(
+            llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
+            [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
+            "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+            "=r,=r,=r,=r,r,~{memory}",
+            True,
+            loc,
+            ip,
+        )
+        return tuple(Int32(llvm.extractvalue(T.i32(), ret, [i], loc=loc, ip=ip)) for i in range(4))
+
+    @dsl_user_op
+    def bf16x2_to_f32x2(word, *, loc=None, ip=None):
+        # Exact widening of the low and high BF16 halves of one 32-bit word.
+        ret = _asm(
+            llvm.StructType.get_literal([T.f32(), T.f32()]),
+            [Int32(word).ir_value(loc=loc, ip=ip)],
+            "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, -65536;",
+            "=f,=f,r",
+            False,
+            loc,
+            ip,
+        )
+        return (
+            Float32(llvm.extractvalue(T.f32(), ret, [0], loc=loc, ip=ip)),
+            Float32(llvm.extractvalue(T.f32(), ret, [1], loc=loc, ip=ip)),
+        )
+
+    @dsl_user_op
+    def fma_rn(a, b, c, *, loc=None, ip=None):
+        return Float32(
+            _asm(
+                T.f32(),
+                [
+                    Float32(a).ir_value(loc=loc, ip=ip),
+                    Float32(b).ir_value(loc=loc, ip=ip),
+                    Float32(c).ir_value(loc=loc, ip=ip),
+                ],
+                "fma.rn.f32 $0, $1, $2, $3;",
+                "=f,f,f,f",
+                False,
+                loc,
+                ip,
+            )
+        )
+
+    @dsl_user_op
+    def add_rn(a, b, *, loc=None, ip=None):
+        return Float32(
+            _asm(
+                T.f32(),
+                [
+                    Float32(a).ir_value(loc=loc, ip=ip),
+                    Float32(b).ir_value(loc=loc, ip=ip),
+                ],
+                "add.rn.f32 $0, $1, $2;",
+                "=f,f,f",
+                False,
+                loc,
+                ip,
+            )
+        )
+
+    @dsl_user_op
+    def fma_rn_x2(a0, a1, b, c0, c1, *, loc=None, ip=None):
+        # Two independent IEEE FMAs (one per lane): (a0*b+c0, a1*b+c1).
+        ret = _asm(
+            llvm.StructType.get_literal([T.f32(), T.f32()]),
+            [
+                Float32(a0).ir_value(loc=loc, ip=ip),
+                Float32(a1).ir_value(loc=loc, ip=ip),
+                Float32(b).ir_value(loc=loc, ip=ip),
+                Float32(c0).ir_value(loc=loc, ip=ip),
+                Float32(c1).ir_value(loc=loc, ip=ip),
+            ],
+            "{\n\t.reg .b64 fa, fb, fc;\n\tmov.b64 fa, {$2, $3};\n\t"
+            "mov.b64 fb, {$4, $4};\n\tmov.b64 fc, {$5, $6};\n\t"
+            "fma.rn.f32x2 fc, fa, fb, fc;\n\tmov.b64 {$0, $1}, fc;\n\t}",
+            "=f,=f,f,f,f,f,f",
+            False,
+            loc,
+            ip,
+        )
+        return (
+            Float32(llvm.extractvalue(T.f32(), ret, [0], loc=loc, ip=ip)),
+            Float32(llvm.extractvalue(T.f32(), ret, [1], loc=loc, ip=ip)),
+        )
+
+    @dsl_user_op
+    def cp_async_wait(pending: int, *, loc=None, ip=None) -> None:
+        _asm(None, [], f"cp.async.wait_group {pending};", "~{memory}", True, loc, ip)
+
+    # Plain Python helpers, inlined while the kernel is traced.
+    def widen(words):
+        vals = []
+        for w in range(4):
+            lo, hi = bf16x2_to_f32x2(words[w])
+            vals.append(lo)
+            vals.append(hi)
+        return vals
+
+    def accumulate(acc, a_words, b_words, paired):
+        """acc[m] += a[m, j] * b[j] for j = 0..7 in order, one FMA each."""
+        num_rows = len(acc)
+        b_vals = widen(b_words)
+        a_vals = [widen(words) for words in a_words]
+        if paired:
+            for m in range(0, num_rows - 1, 2):
+                for j in range(ELEMS_PER_VECTOR):
+                    acc[m], acc[m + 1] = fma_rn_x2(
+                        a_vals[m][j], a_vals[m + 1][j], b_vals[j], acc[m], acc[m + 1]
+                    )
+            if num_rows % 2:
+                for j in range(ELEMS_PER_VECTOR):
+                    acc[-1] = fma_rn(a_vals[-1][j], b_vals[j], acc[-1])
+        else:
+            for m in range(num_rows):
+                for j in range(ELEMS_PER_VECTOR):
+                    acc[m] = fma_rn(a_vals[m][j], b_vals[j], acc[m])
+        return acc
+
+    class RouterGemmKernel:
+        """``out[m, e] = sum_k a[m, k] * b[e, k]`` with native arithmetic.
+
+        ``experts_per_cta == 1``: every thread stages exactly the activation
+        and weight vectors it consumes, one ``cp.async`` group per K step, so
+        K step ``i`` is computed while later steps are still landing and no
+        CTA barrier is needed before the final reduction. ``> 1``: the CTA
+        stages the activation block once for all its experts.
+        """
+
+        def __init__(
+            self,
+            num_rows: int,
+            hidden_size: int,
+            num_experts: int,
+            experts_per_cta: int,
+            use_pdl: bool,
+            prefetch_weights: bool,
+            paired_fma: bool,
+            attribution: str,
+        ):
+            assert hidden_size % K_PER_ITERATION == 0
+            assert num_experts % experts_per_cta == 0
+            assert attribution in ("", "no_activation", "no_compute")
+            self.num_rows = num_rows
+            self.hidden_size = hidden_size
+            self.num_experts = num_experts
+            self.experts_per_cta = experts_per_cta
+            self.use_pdl = use_pdl
+            self.prefetch_weights = prefetch_weights
+            self.paired_fma = paired_fma
+            self.attribution = attribution
+            self.k_iterations = hidden_size // K_PER_ITERATION
+            self.threads = THREADS_PER_EXPERT * experts_per_cta
+            self.a_vectors = num_rows * hidden_size // ELEMS_PER_VECTOR
+
+        @cute.jit
+        def __call__(self, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, stream: CUstream):
+            self.kernel(mA, mB, mC).launch(
+                grid=(self.num_experts // self.experts_per_cta, 1, 1),
+                block=(self.threads, 1, 1),
+                stream=stream,
+                use_pdl=self.use_pdl,
+            )
+
+        @cute.kernel
+        def kernel(self, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
+            M = self.num_rows
+            K = self.hidden_size
+            KI = self.k_iterations
+            EPC = self.experts_per_cta
+            NT = self.threads
+            row_bytes = K * 2
+            iter_bytes = K_PER_ITERATION * 2
+            private = EPC == 1
+            stage_a = self.attribution != "no_activation"
+            compute = self.attribution != "no_compute"
+
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            slot = tidx // THREADS_PER_EXPERT
+            t = tidx % THREADS_PER_EXPERT
+            warp = t // 32
+            lane = tidx % 32
+            expert = bidx * EPC + slot
+
+            smem = SmemAllocator()
+            s_a = smem.allocate_tensor(BFloat16, cute.make_layout(M * K), byte_alignment=128)
+            s_b = smem.allocate_tensor(BFloat16, cute.make_layout(EPC * K), byte_alignment=128)
+            warp_sums = smem.allocate_tensor(Float32, cute.make_layout((WARPS_PER_EXPERT, M, EPC)))
+
+            s_a_base = Int32(s_a.iterator.toint())
+            s_b_base = Int32(s_b.iterator.toint())
+            g_a_base = Int64(mA.iterator.toint())
+            # This thread's weight vectors: expert `expert`, k = i*1024+t*8.
+            b_off = slot * row_bytes + t * 16
+            g_b_mine = (
+                Int64(mB.iterator.toint()) + Int64(bidx * EPC) * Int64(row_bytes) + Int64(b_off)
+            )
+            s_b_mine = s_b_base + b_off
+            g_a_mine = g_a_base + Int64(t * 16)
+            # Shared-memory address of activation vector (i, m) read by `t`.
+            # Private staging packs a thread's vectors K-step-major; shared
+            # staging keeps the row-major [M, K] image of the activation.
+            if cutlass.const_expr(private):
+                s_a_mine = s_a_base + t * 16
+                s_a_step = M * THREADS_PER_EXPERT * 16
+                s_a_row = THREADS_PER_EXPERT * 16
+            else:
+                s_a_mine = s_a_base + t * 16
+                s_a_step = iter_bytes
+                s_a_row = row_bytes
+
+            if cutlass.const_expr(self.prefetch_weights):
+                for i in cutlass.range_constexpr(KI):
+                    cp_async_16(s_b_mine + i * iter_bytes, g_b_mine + Int64(i * iter_bytes))
+                cp_async_commit()
+
+            if cutlass.const_expr(self.use_pdl):
+                griddep_wait()
+
+            if cutlass.const_expr(not self.prefetch_weights):
+                for i in cutlass.range_constexpr(KI):
+                    cp_async_16(s_b_mine + i * iter_bytes, g_b_mine + Int64(i * iter_bytes))
+                cp_async_commit()
+
+            if cutlass.const_expr(private):
+                for i in cutlass.range_constexpr(KI):
+                    if cutlass.const_expr(stage_a):
+                        for m in cutlass.range_constexpr(M):
+                            cp_async_16(
+                                s_a_mine + i * s_a_step + m * s_a_row,
+                                g_a_mine + Int64(m * row_bytes + i * iter_bytes),
+                            )
+                    cp_async_commit()
+            else:
+                # The activation block [M, K] is one contiguous range.
+                if cutlass.const_expr(stage_a):
+                    for v in cutlass.range_constexpr((self.a_vectors + NT - 1) // NT):
+                        idx = tidx + v * NT
+                        if cutlass.const_expr(self.a_vectors % NT == 0):
+                            cp_async_16(s_a_base + idx * 16, g_a_base + Int64(idx) * Int64(16))
+                        else:
+                            if idx < self.a_vectors:
+                                cp_async_16(s_a_base + idx * 16, g_a_base + Int64(idx) * Int64(16))
+                cp_async_commit()
+                cp_async_wait(0)
+                cute.arch.sync_threads()
+
+            if cutlass.const_expr(self.use_pdl):
+                griddep_launch_dependents()
+
+            acc = [Float32(0.0)] * M
+            for i in cutlass.range_constexpr(KI):
+                if cutlass.const_expr(private):
+                    # Groups complete in order: weights, then K steps 0..i.
+                    cp_async_wait(KI - 1 - i)
+                if cutlass.const_expr(compute):
+                    b_words = lds_v4(s_b_mine + i * iter_bytes)
+                    a_words = []
+                    for m in cutlass.range_constexpr(M):
+                        a_words.append(lds_v4(s_a_mine + i * s_a_step + m * s_a_row))
+                    acc = accumulate(acc, a_words, b_words, self.paired_fma)
+
+            for m in cutlass.range_constexpr(M):
+                s = acc[m]
+                for r in cutlass.range_constexpr(5):
+                    s = add_rn(s, cute.arch.shuffle_sync_bfly(s, offset=16 >> r))
+                if lane == 0:
+                    warp_sums[warp, m, slot] = s
+
+            cute.arch.sync_threads()
+
+            if t < M:
+                total = Float32(0.0)
+                for w in cutlass.range_constexpr(WARPS_PER_EXPERT):
+                    total = add_rn(total, warp_sums[w, t, slot])
+                mC[t, expert] = total
+
+    return cutlass, cute, RouterGemmKernel
+
+
+_kernel_class = None
+
+
+def _get_compiled(
+    num_rows: int,
+    hidden_size: int,
+    num_experts: int,
+    experts_per_cta: int,
+    use_pdl: bool,
+    prefetch_weights: bool,
+    paired_fma: bool,
+    attribution: str,
+):
+    key = (
+        num_rows,
+        hidden_size,
+        num_experts,
+        experts_per_cta,
+        use_pdl,
+        prefetch_weights,
+        paired_fma,
+        attribution,
+    )
+    fn = _compiled.get(key)
+    if fn is not None:
+        return fn
+    global _kernel_class
+    with _compile_lock:
+        fn = _compiled.get(key)
+        if fn is not None:
+            return fn
+        if _kernel_class is None:
+            _kernel_class = _build_kernel_class()
+        cutlass, cute, kernel_cls = _kernel_class
+        a_fake = cute.runtime.make_fake_compact_tensor(
+            cutlass.BFloat16, (num_rows, hidden_size), stride_order=(1, 0), assumed_align=16
+        )
+        b_fake = cute.runtime.make_fake_compact_tensor(
+            cutlass.BFloat16, (num_experts, hidden_size), stride_order=(1, 0), assumed_align=16
+        )
+        c_fake = cute.runtime.make_fake_compact_tensor(
+            cutlass.Float32, (num_rows, num_experts), stride_order=(1, 0), assumed_align=16
+        )
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        kernel = kernel_cls(
+            num_rows,
+            hidden_size,
+            num_experts,
+            experts_per_cta,
+            use_pdl,
+            prefetch_weights,
+            paired_fma,
+            attribution,
+        )
+        fn = cute.compile(kernel, a_fake, b_fake, c_fake, stream, options="--enable-tvm-ffi")
+        _compiled[key] = fn
+        return fn
+
+
+def router_gemm(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    experts_per_cta: Optional[int] = None,
+    use_pdl: Optional[bool] = None,
+    prefetch_weights: bool = True,
+    paired_fma: bool = DEFAULT_PAIRED_FMA,
+    _attribution: str = "",
+) -> torch.Tensor:
+    """FP32 router logits ``hidden @ weight.T``, bitwise equal to native.
+
+    Args:
+        hidden: ``[R, K]`` BF16, contiguous, ``1 <= R <= 16``.
+        weight: ``[256, K]`` BF16, contiguous.
+        out: optional ``[R, 256]`` FP32 contiguous output.
+
+    The caller must have checked ``router_gemm_supported``. The first call
+    for a geometry compiles the kernel, so it must happen outside CUDA graph
+    capture (the executor's eager warmup covers every captured shape).
+    ``_attribution`` ("no_activation" / "no_compute") drops one phase for
+    timing attribution only; its output is meaningless.
+    """
+    rows, hidden_size = hidden.shape
+    num_experts = weight.shape[0]
+    if experts_per_cta is None:
+        experts_per_cta = DEFAULT_EXPERTS_PER_CTA
+    experts_per_cta = fit_experts_per_cta(rows, hidden_size, experts_per_cta)
+    if use_pdl is None:
+        use_pdl = pdl_enabled()
+    if out is None:
+        out = torch.empty((rows, num_experts), dtype=torch.float32, device=hidden.device)
+    fn = _get_compiled(
+        rows,
+        hidden_size,
+        num_experts,
+        experts_per_cta,
+        use_pdl,
+        prefetch_weights,
+        paired_fma,
+        _attribution,
+    )
+    fn(hidden, weight, out)
+    return out
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/__init__.py
new file mode 100644
index 000000000000..b8515112d68e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/__init__.py
@@ -0,0 +1,31 @@
+# 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.
+"""sparse_mla_decode: sparse MLA generation attention (the sparse_attention contract).
+
+Entry points registered with ``dispatch``:
+
+* ``fmha_lib``: ``SparseMlaDecodeFmha``, an FMHA library class that the
+  attention FMHA registry places ahead of the native libraries when ``sparse_mla_decode`` is
+  selected.
+* ``sparse_mla_decode``: the functional kernel wrapper.
+"""
+
+from ... import dispatch
+from .sparse_mla import SparseMlaDecodeFmha, sparse_mla_decode
+
+dispatch.register("sparse_mla_decode", "fmha_lib", SparseMlaDecodeFmha)
+dispatch.register("sparse_mla_decode", "sparse_mla_decode", sparse_mla_decode)
+
+__all__ = ["SparseMlaDecodeFmha", "sparse_mla_decode"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla.py
new file mode 100644
index 000000000000..5ac2e95f88dc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla.py
@@ -0,0 +1,335 @@
+# 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.
+"""Sparse MLA generation attention for GLM DSA layers with an FP8 latent cache.
+
+The sparse_attention contract: the native ``mla_rope_generation`` has already
+rotated and FP8-quantized the query into ``quant_q_buffer`` and appended this
+step's latent rows to the paged pool; the native DSA predictor has converted
+the top-k selection into global pool rows. This module computes only the
+attention over those rows and writes the latent output ``[R, H * 512]`` that
+the native W_UV projection consumes.
+
+``SparseMlaDecodeFmha`` is an FMHA library for ``TrtllmAttention``'s library
+list. Every eligibility decision is made in ``is_available`` / ``is_supported``,
+so an ineligible layer or request falls through to the native library before
+any device work. A layer that is structurally outside the domain (not DSA, other
+head counts, BF16 cache, another SM family) takes the native path; a selected sparse_mla_decode
+whose kernel is missing from the build, or whose device cannot launch it, raises
+instead of falling back on one rank.
+
+Rows: uniform generation batches (``seq_len_q`` query rows per sequence) and the
+ragged-verification layout, which ``TrtllmAttention.forward`` presents to the
+FMHA library as one row per query token with that row's causal KV length
+(``presented_token_major``); both give each row the same causal extent.
+
+Index semantics: the native rules (positions at or past a row's causal length
+are ignored, repeats count) with one deliberate difference: a negative entry at a
+position below the causal length is skipped, while the native generation kernel
+reads the pool row that ``-1`` addresses. Native producers of the global indices
+(the indexer top-k and the local-to-global remap) never emit such an entry; the
+tests check that premise for the native and grouped remaps.
+
+Dispatch accounting: ``engagement_snapshot()`` returns the per-layer static
+decisions, the number of host-side calls that ran sparse_mla_decode per layer (eager calls and
+CUDA-graph captures; a replay runs the captured kernel without host code) and
+the declined selections per reason. Each layer's first static decision and first
+sparse_mla_decode call are logged once.
+"""
+
+from typing import TYPE_CHECKING, Dict, Optional
+
+import torch
+
+from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha, FmhaPhase
+from tensorrt_llm._torch.attention.backends.interface import (
+    AttentionForwardArgs,
+    AttentionInputType,
+)
+from tensorrt_llm._utils import get_sm_version
+
+from ... import dispatch
+
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.attention.backends.trtllm import (
+        TrtllmAttention,
+        TrtllmAttentionMetadata,
+    )
+
+FEATURE_ID = "sparse_mla_decode"
+NUM_HEADS = 8
+KV_LORA_RANK = 512
+QK_ROPE_HEAD_DIM = 64
+HEAD_DIM_QK = KV_LORA_RANK + QK_ROPE_HEAD_DIM
+TOP_K = 2048
+SUPPORTED_SM_VERSIONS = (100, 103)
+
+
+def op_available() -> bool:
+    """Whether the loaded TensorRT-LLM build registers the sparse_mla_decode kernel."""
+    return hasattr(torch.ops.trtllm, "glm_moe_dsa_sparse_mla_decode")
+
+
+_DEVICE_SUPPORT: Dict[int, bool] = {}
+
+# Dispatch accounting of this process (see the module docstring).
+_STATIC_DECISIONS: Dict[int, str] = {}
+_ENGAGED_CALLS: Dict[int, int] = {}
+_DECLINED_SELECTIONS: Dict[str, int] = {}
+
+
+def device_supported(device_index: int) -> bool:
+    """Whether the device can launch the 16-CTA cluster kernel (cached per device)."""
+    if device_index not in _DEVICE_SUPPORT:
+        _DEVICE_SUPPORT[device_index] = bool(
+            op_available()
+            and get_sm_version() in SUPPORTED_SM_VERSIONS
+            and torch.ops.trtllm.glm_moe_dsa_sparse_mla_decode_supported(device_index)
+        )
+    return _DEVICE_SUPPORT[device_index]
+
+
+def engagement_snapshot() -> dict:
+    """Copy of this process's dispatch accounting.
+
+    Returns:
+        ``static``: layer index -> "" when the layer may use sparse_mla_decode, else the reason;
+        ``engaged_calls``: layer index -> host-side calls that ran sparse_mla_decode;
+        ``declined_selections``: reason -> uncached library selections that declined.
+    """
+    return {
+        "static": dict(_STATIC_DECISIONS),
+        "engaged_calls": dict(_ENGAGED_CALLS),
+        "declined_selections": dict(_DECLINED_SELECTIONS),
+    }
+
+
+def sparse_mla_decode(
+    q_fp8: torch.Tensor,
+    kv_pool_ptr: int,
+    indices: torch.Tensor,
+    bmm1_scale: torch.Tensor,
+    bmm2_scale: torch.Tensor,
+    seq_lens_kv: torch.Tensor,
+    seq_len_q: int,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Run sparse MLA attention for ``R`` query rows; writes and returns ``out``.
+
+    The rounding and masking rules are those of the native generation FMHA for
+    this layout (see ``sparseMlaDecode.h``): index-list positions at or past a
+    row's causal KV length are ignored, negative entries are skipped, repeated
+    rows count per occurrence, probabilities enter P*V as ``e4m3(448 * p)``.
+
+    Args:
+        q_fp8: ``[R, 8, 576]`` e4m3 query bytes (``quant_q_buffer``).
+        kv_pool_ptr: Device address of the FP8 latent pool base.
+        indices: ``[R, 2048]`` int32 global pool rows.
+        bmm1_scale: ``[2]`` float32 device scales from ``mla_rope_generation``.
+        bmm2_scale: ``[1]`` float32 device scale from ``mla_rope_generation``.
+        seq_lens_kv: ``[R / seq_len_q]`` int32 KV lengths including this step's rows.
+        seq_len_q: Query rows per sequence.
+        out: ``[R, 8 * 512]`` BF16 output.
+    """
+    torch.ops.trtllm.glm_moe_dsa_sparse_mla_decode(
+        q_fp8, kv_pool_ptr, indices, bmm1_scale, bmm2_scale, seq_lens_kv, seq_len_q, out
+    )
+    return out
+
+
+def pool_base_pointer(metadata: "TrtllmAttentionMetadata", local_layer_idx: int) -> int:
+    """Pool base address the native sparse MLA path uses for this layer."""
+    pool_index = int(metadata.host_kv_cache_pool_mapping[local_layer_idx, 0])
+    pointers = metadata.host_kv_cache_pool_pointers
+    base = pointers[pool_index, 0, 0] if pointers.dim() == 3 else pointers[pool_index, 0]
+    return int(base)
+
+
+def static_ineligibility(attn: "TrtllmAttention") -> Optional[str]:
+    """Return why ``attn`` can never use sparse_mla_decode, or None if it may."""
+    if getattr(attn.sparse_params, "algorithm", None) != "dsa":
+        return "not a DSA layer"
+    if not attn.is_mla_enable:
+        return "not MLA"
+    if attn.num_heads != NUM_HEADS or attn.num_kv_heads != 1:
+        return f"local heads {attn.num_heads}/{attn.num_kv_heads}"
+    if attn.kv_lora_rank != KV_LORA_RANK or attn.qk_rope_head_dim != QK_ROPE_HEAD_DIM:
+        return "latent dims"
+    if not attn.rope_append:
+        return "rope_append disabled"
+    if not getattr(attn, "has_fp8_kv_cache", False) or getattr(attn, "has_fp4_kv_cache", False):
+        return "KV cache is not FP8"
+    if getattr(attn, "kv_cache_dtype", None) == "fp8_ds_mla":
+        return "fp8_ds_mla cache layout"
+    if attn.attention_chunk_size is not None:
+        return "chunked attention"
+    if get_sm_version() not in SUPPORTED_SM_VERSIONS:
+        return f"SM {get_sm_version()}"
+    # The layer is in the domain from here on: a missing kernel or a failed device
+    # query is an error of the selected candidate, not a per-rank fallback.
+    if not op_available():
+        raise RuntimeError(
+            "the plan enables sparse_mla_decode, but this "
+            "TensorRT-LLM build does not register "
+            "torch.ops.trtllm.glm_moe_dsa_sparse_mla_decode"
+        )
+    device = torch.cuda.current_device()
+    if not device_supported(device):
+        raise RuntimeError(
+            f"the plan enables sparse_mla_decode, but device "
+            f"{device} (SM {get_sm_version()}) cannot launch its 16-CTA cluster kernel: the "
+            "cluster occupancy query failed or returned no active cluster"
+        )
+    return None
+
+
+def request_ineligibility(
+    q: torch.Tensor,
+    k: Optional[torch.Tensor],
+    v: Optional[torch.Tensor],
+    metadata: "TrtllmAttentionMetadata",
+    forward_args: AttentionForwardArgs,
+) -> Optional[str]:
+    """Return why this request must take the native path, or None."""
+    if forward_args.attention_input_type != AttentionInputType.generation_only:
+        return "not generation-only"
+    if k is not None or v is not None:
+        return "separate K/V"
+    if metadata.num_generations <= 0:
+        return "no generation rows"
+    if metadata.effective_beam_width != 1:
+        return "beam search"
+    if metadata.helix_position_offsets is not None:
+        return "helix context parallelism"
+    if metadata.is_spec_dec_tree:
+        return "tree speculative decoding"
+    if metadata.num_sparse_topk != TOP_K:
+        return f"top-k {metadata.num_sparse_topk}"
+    if q.dim() != 2 or q.shape[1] != NUM_HEADS * HEAD_DIM_QK:
+        return "query shape"
+    num_rows = q.shape[0]
+    if (
+        forward_args.quant_q_buffer is None
+        or forward_args.mla_bmm1_scale is None
+        or forward_args.mla_bmm2_scale is None
+    ):
+        return "no FP8 query"
+    if forward_args.quant_scale_qkv is not None:
+        return "pre-quantized query path"
+    if forward_args.quant_q_buffer.numel() != num_rows * NUM_HEADS * HEAD_DIM_QK:
+        return "query buffer shape"
+    output = forward_args.output
+    if (
+        output is None
+        or output.dtype != torch.bfloat16
+        or forward_args.output_sf is not None
+        or not output.is_contiguous()
+        or output.numel() != num_rows * NUM_HEADS * KV_LORA_RANK
+    ):
+        return "output format"
+    if forward_args.softmax_stats_tensor is not None or forward_args.attention_sinks is not None:
+        return "softmax statistics or sinks requested"
+    if forward_args.enable_dsv4_epilogue_fusion:
+        return "epilogue fusion"
+    sparse = forward_args.sparse_runtime_params
+    indices = sparse.sparse_attn_indices
+    if indices is None or indices.dtype != torch.int32 or indices.dim() != 2:
+        return "no sparse indices"
+    if indices.shape[0] != num_rows or indices.shape[1] != TOP_K or indices.stride(1) != 1:
+        return "sparse index layout"
+    if (
+        sparse.sparse_attn_kv_lens is not None
+        or sparse.aux_kv_cache_pool_ptr is not None
+        or sparse.sparse_kv_indices is not None
+    ):
+        return "dynamic or auxiliary-pool sparse attention"
+    # KV lengths per generation sequence, this step's rows included. Under the
+    # ragged-verification presentation every query token is its own sequence
+    # (seq_len_q 1) carrying its causal length, so the same rule applies.
+    kv_lens = metadata.kv_lens_cuda_runtime
+    if kv_lens.dtype != torch.int32 or not kv_lens.is_contiguous():
+        return "KV length layout"
+    gen_seqs = kv_lens.shape[0] - metadata.num_contexts
+    if gen_seqs <= 0 or num_rows % gen_seqs != 0:
+        return "rows do not tile the generation sequences"
+    return None
+
+
+class SparseMlaDecodeFmha(Fmha):
+    """Generation-only sparse MLA library for GLM DSA layers with an FP8 cache."""
+
+    @classmethod
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
+        reason = static_ineligibility(attn)
+        _STATIC_DECISIONS[attn.layer_idx] = reason or ""
+        return dispatch.log_capability(
+            FEATURE_ID, f"layer{attn.layer_idx}", reason is None, reason or ""
+        )
+
+    def _is_supported(
+        self,
+        q: torch.Tensor,
+        k: Optional[torch.Tensor],
+        v: Optional[torch.Tensor],
+        metadata: "TrtllmAttentionMetadata",
+        forward_args: AttentionForwardArgs,
+        *,
+        phase: Optional[FmhaPhase] = None,
+    ) -> bool:
+        if phase is not None and phase != FmhaPhase.GENERATION:
+            return False
+        reason = request_ineligibility(q, k, v, metadata, forward_args)
+        if reason is not None:
+            _DECLINED_SELECTIONS[reason] = _DECLINED_SELECTIONS.get(reason, 0) + 1
+            dispatch.log_capability(
+                FEATURE_ID,
+                f"request:{reason}",
+                False,
+                f"{reason}; first seen at layer {self.attn.layer_idx} with {q.shape[0]} rows",
+            )
+            return False
+        return True
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: Optional[torch.Tensor],
+        v: Optional[torch.Tensor],
+        metadata: "TrtllmAttentionMetadata",
+        forward_args: AttentionForwardArgs,
+    ) -> None:
+        attn = self.attn
+        num_rows = q.shape[0]
+        gen_lens = metadata.kv_lens_cuda_runtime[metadata.num_contexts :]
+        seq_len_q = num_rows // gen_lens.shape[0]
+        calls = _ENGAGED_CALLS.get(attn.layer_idx, 0)
+        _ENGAGED_CALLS[attn.layer_idx] = calls + 1
+        if calls == 0:
+            ragged = metadata.token_major_gen_view() is not None
+            dispatch.log_capability(
+                FEATURE_ID,
+                f"call:layer{attn.layer_idx}",
+                True,
+                f"first call: {num_rows} rows, seq_len_q {seq_len_q}, ragged {ragged}",
+            )
+        sparse_mla_decode(
+            forward_args.quant_q_buffer.view(num_rows, NUM_HEADS, HEAD_DIM_QK),
+            pool_base_pointer(metadata, attn.local_layer_idx),
+            forward_args.sparse_runtime_params.sparse_attn_indices,
+            forward_args.mla_bmm1_scale,
+            forward_args.mla_bmm2_scale,
+            gen_lens,
+            seq_len_q,
+            forward_args.output,
+        )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/layers.py
new file mode 100644
index 000000000000..dd2d8ae7f0f2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/layers.py
@@ -0,0 +1,1189 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The decoder layers, the draft layer and the MoE modules of the glm_moe_dsa targets.
+
+``adopt`` re-classes a built-in module in place: a ``DeepseekV3DecoderLayer`` as
+``GlmDecoderLayer``, the ``DeepseekV3MTP`` draft layer as ``GlmMtpLayer`` and a
+``Deepseekv3MoE`` as ``GlmMoE``. It then attaches the fast-path objects the installed plan
+enables, in the order the built-in constructor builds the modules they wrap: the attention
+modules and the attention tail (``attention.GlmMLA``), the post-attention hop and its
+transports, the router GEMM, the post-MoE hops, the attention tail's successor weights and, on
+the draft layer, the draft front and the live-row hop. The
+classes add methods and class-level defaults only (no parameter, buffer or submodule), so
+the module tree, the weight names and the loaded weights stay the built-in model's.
+
+Each forward override runs the family's decode composition when the step is a decode step
+(``dispatch.Phase.DECODE``) and the built-in method otherwise, so a prefill or mixed step and
+a call outside a step run the built-in model.
+
+Every target of the family runs this composition; the targets differ in the plan they install
+(``dispatch.TargetPaths``). A fast path's block runs only when the plan enables that path
+(``dispatch.is_enabled``) or when an object only that path attaches is set, so the block of a
+path a target does not serve never runs for that target. The routed-MoE paths, for instance,
+take tensor-parallel experts: a target whose experts are expert-parallel does not serve them,
+and its MoE runs the built-in experts on the router logits the other paths produce.
+
+``adopt_model`` adopts a model's layers and its speculative worker; ``adopted_state`` lists
+what the adoption attached, for the rank-uniform check.
+"""
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from tensorrt_llm._torch.attention.attention import maybe_slice_for_helix_cp
+from tensorrt_llm._torch.attention.backends import AttentionMetadata
+from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams, MoEAllReduceParams
+from tensorrt_llm._torch.models.modeling_deepseekv3 import (
+    DeepseekV3DecoderLayer,
+    DeepseekV3Gate,
+    Deepseekv3MoE,
+    DeepseekV3MTP,
+    GlmMoeDsaForCausalLM,
+    moe_reduce_add_shared_output,
+)
+from tensorrt_llm._torch.modules.embedding import Embedding
+from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel
+from tensorrt_llm._torch.speculative import SpecMetadata
+from tensorrt_llm._torch.utils import EventType, Fp4QuantizedTensor
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.bindings.internal.thop import BufferKind
+
+from . import dispatch
+from .attention import GlmMLA, reclass
+from .spec_worker import GlmMtpWorker
+
+
+class GlmMoE(Deepseekv3MoE):
+    """The MoE of a decoder layer or of the draft layer: the router-logit producers and the
+    routed fast paths in front of the built-in experts."""
+
+    # The fast-path objects ``adopt`` and the layer's adoption attach; None keeps the built-in
+    # computation.
+    router_gemm = None
+    _routed_into_hop = None
+
+    @classmethod
+    def adopt(cls, moe: Deepseekv3MoE, model_config, layer_idx: int) -> "GlmMoE":
+        """Re-class a built-in MoE and attach its router GEMM.
+
+        ``model_config`` and ``layer_idx`` are the ones the built-in constructor received.
+
+        Raises:
+            TypeError: if ``moe`` is not exactly the built-in ``Deepseekv3MoE``.
+        """
+        reclass(moe, Deepseekv3MoE, cls)
+        # The router GEMM may replace this layer's gate GEMM; None keeps the native gate.
+        if dispatch.is_enabled("router_gemm"):
+            moe.router_gemm = dispatch.get_impl("router_gemm", "make_moe_router")(
+                gate=moe.gate, layer_idx=layer_idx, model_config=model_config
+            )
+        return moe
+
+    def compute_routed_output(
+        self,
+        hidden_states,
+        hidden_states_fp4,
+        all_rank_num_tokens,
+        do_finalize,
+        router_logits=None,
+        hop_routing_out=None,
+    ):
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().compute_routed_output(
+                hidden_states=hidden_states,
+                hidden_states_fp4=hidden_states_fp4,
+                all_rank_num_tokens=all_rank_num_tokens,
+                do_finalize=do_finalize,
+            )
+        # max-throughput
+        use_dp_padding = False
+        # Add DP padding on SM120 for context comm performance
+        # TODO: Move this model-agonostic part to MoE
+        if self.use_dp and self.mapping.tp_size > 1 and get_sm_version() == 120:
+            use_dp_padding = True
+            hidden_states = torch.nn.functional.pad(
+                hidden_states, (0, 0, 0, max(all_rank_num_tokens) - hidden_states.shape[0])
+            )
+
+        # Router logits come from at most one producer: the post-attention hop (the
+        # argument), the router GEMM, or the native gate.
+        if router_logits is None and self.router_gemm is not None and not use_dp_padding:
+            # Returns None, before launching anything, outside its domain.
+            router_logits = self.router_gemm(hidden_states)
+        if router_logits is None:
+            router_logits = self.gate(hidden_states)
+
+        # The draft layer's BF16 routed experts (the FlashInfer TRT-LLM-Gen BF16 leaf) through
+        # the draft routed-MoE kernels on the leaf's own buffers, bit-identical to the native
+        # call. Its eligibility (leaf type, routing, shapes, dtypes, device, finalized output)
+        # is decided before any device work; None keeps the path below (NVFP4 layers always
+        # decline).
+        if (
+            hidden_states_fp4 is None
+            and not self.use_dp
+            and not use_dp_padding
+            and dispatch.is_enabled("draft_routed_moe")
+        ):
+            routed_output = dispatch.get_impl("draft_routed_moe", "routed_output")(
+                self.experts,
+                self.gate.e_score_correction_bias,
+                hidden_states,
+                router_logits,
+                do_finalize,
+            )
+            if routed_output is not None:
+                return routed_output
+
+        # The routed MoE: routing + activation quant + FC1 + FC2 + finalize in one CuTe DSL
+        # launch on the native expert buffers. Its eligibility (shapes/dtypes/device/routing/
+        # bias) is decided before any device work; None means "run the native path below".
+        if (
+            hidden_states_fp4 is None
+            and not self.use_dp
+            and not use_dp_padding
+            and dispatch.is_enabled("routed_moe")
+        ):
+            # With the post-MoE hop as the consumer, the routed MoE writes its per-rank routed
+            # result straight into the hop's routed-input buffer (in the hop's mode) and
+            # returns [buffer]; None keeps the raw tuple below. The hop is set only when every
+            # TP rank enabled it.
+            routed_into_hop = getattr(self, "_routed_into_hop", None)
+            if (
+                routed_into_hop is not None
+                and not do_finalize
+                and dispatch.is_enabled("routed_moe_into_hop")
+            ):
+                # The hop-routing variant that also prefetches this layer's shared-expert
+                # weights into L2 before its dependency wait, so the shared expert on the
+                # auxiliary stream reads them from L2. None (decided before any device work)
+                # keeps the hop-routing variant.
+                if hop_routing_out is not None and dispatch.is_enabled("shared_expert_prefetch"):
+                    routed_output = dispatch.get_impl(
+                        "shared_expert_prefetch", "routed_output_into_hop"
+                    )(
+                        self.experts,
+                        self.gate.e_score_correction_bias,
+                        hidden_states,
+                        router_logits,
+                        hop_routing_out,
+                        routed_into_hop,
+                        self.shared_experts,
+                    )
+                    if routed_output is not None:
+                        return routed_output
+                # The routed MoE on the routing the pre-MoE hop computed (set only when the
+                # hop-routing hop ran); None keeps the routed MoE, which routes from the same
+                # logits to the same experts.
+                if hop_routing_out is not None:
+                    routed_output = dispatch.get_impl("hop_routing", "routed_output_into_hop")(
+                        self.experts,
+                        self.gate.e_score_correction_bias,
+                        hidden_states,
+                        router_logits,
+                        hop_routing_out,
+                        routed_into_hop,
+                    )
+                    if routed_output is not None:
+                        return routed_output
+                routed_output = dispatch.get_impl("routed_moe", "routed_output_into_hop")(
+                    self.experts,
+                    self.gate.e_score_correction_bias,
+                    hidden_states,
+                    router_logits,
+                    routed_into_hop,
+                )
+                if routed_output is not None:
+                    return routed_output
+            if hop_routing_out is not None:
+                routed_output = dispatch.get_impl("hop_routing", "routed_output")(
+                    self.experts,
+                    self.gate.e_score_correction_bias,
+                    hidden_states,
+                    router_logits,
+                    hop_routing_out,
+                    do_finalize,
+                )
+                if routed_output is not None:
+                    return routed_output
+            routed_output = dispatch.get_impl("routed_moe", "routed_output")(
+                self.experts,
+                self.gate.e_score_correction_bias,
+                hidden_states,
+                router_logits,
+                do_finalize,
+            )
+            if routed_output is not None:
+                return routed_output
+
+        routed_output = self.experts(
+            hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states,
+            router_logits,
+            do_finalize=do_finalize,
+            output_dtype=hidden_states.dtype,
+            all_rank_num_tokens=all_rank_num_tokens,
+            use_dp_padding=use_dp_padding,
+        )
+
+        return routed_output
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_fp4: Optional[Fp4QuantizedTensor] = None,
+        all_rank_num_tokens: Optional[list[int]] = None,
+        final_all_reduce_params: Optional[AllReduceParams] = None,
+        do_finalize: Optional[bool] = True,
+        router_logits: Optional[torch.Tensor] = None,
+        hop_routing_out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward(
+                hidden_states=hidden_states,
+                hidden_states_fp4=hidden_states_fp4,
+                all_rank_num_tokens=all_rank_num_tokens,
+                final_all_reduce_params=final_all_reduce_params,
+                do_finalize=do_finalize,
+            )
+        if not do_finalize:
+            assert not self.use_dp
+        # Precomputed logits (and the hop routing) come only from the TP post-attention hop.
+        assert router_logits is None or not self.use_dp
+        assert hop_routing_out is None or router_logits is not None
+
+        def _compute_shared_output():
+            shared_input = (
+                hidden_states_fp4
+                if (hidden_states_fp4 is not None and self.shared_experts_use_fp4)
+                else hidden_states
+            )
+            shared_output = self.shared_experts(shared_input)
+            if self.shared_output_scale is not None:
+                shared_output *= self.shared_output_scale
+            return shared_output
+
+        def _compute_routed_output():
+            routed_output = self.compute_routed_output(
+                hidden_states,
+                hidden_states_fp4,
+                all_rank_num_tokens,
+                do_finalize,
+                router_logits,
+                hop_routing_out,
+            )
+            return routed_output
+
+        # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames
+
+        if self.shared_experts is not None:
+            routed_output, shared_output = maybe_execute_in_parallel(
+                _compute_routed_output,
+                _compute_shared_output,
+                self.event_dict[EventType.Main],
+                self.event_dict[EventType.MoeShared],
+                self.aux_stream,
+                disable_on_compile=True,
+            )
+        else:
+            # Shared experts have been fused into the routed experts (see post_load_weights);
+            # routed_output already contains the shared expert contribution.
+            shared_output = None
+            routed_output = _compute_routed_output()
+
+        if not do_finalize:
+            return [shared_output, *routed_output]
+        else:
+            if shared_output is None:
+                final_hidden_states = routed_output
+            elif not isinstance(shared_output, torch.Tensor):
+                final_hidden_states = shared_output + routed_output
+                if not self.use_dp and self.mapping.tp_size > 1:
+                    final_hidden_states = self.allreduce(
+                        final_hidden_states, all_reduce_params=final_all_reduce_params
+                    )
+                return final_hidden_states
+            else:
+                output_tensor = None
+                if not self.use_dp and self.mapping.tp_size > 1:
+                    w, actual_kind = torch.ops.trtllm.allocate_output(
+                        shared_output, self.allreduce.output_buffer_kind, self.mapping.tp_group
+                    )
+                    if actual_kind == int(BufferKind.NCCL_WINDOW):
+                        output_tensor = w
+                if routed_output.dim() == 3:
+                    assert shared_output.numel() * self.top_k == routed_output.numel(), (
+                        "unmatched tensor shape"
+                    )
+                    final_hidden_states = moe_reduce_add_shared_output(
+                        routed_output, shared_output, out=output_tensor
+                    )
+                else:
+                    assert shared_output.size() == routed_output.size(), "unmatched tensor shape"
+                    if output_tensor is not None:
+                        final_hidden_states = torch.add(
+                            shared_output, routed_output, out=output_tensor
+                        )
+                    else:
+                        # In-place add to avoid allocating a temporary tensor, reducing peak memory
+                        final_hidden_states = shared_output.add_(routed_output)
+
+            if not self.use_dp and self.mapping.tp_size > 1:
+                final_hidden_states = self.allreduce(
+                    final_hidden_states, all_reduce_params=final_all_reduce_params
+                )
+            return final_hidden_states
+
+
+class _PostAttentionHopMixin:
+    """The post-attention hop of a decoder layer or of the draft layer: when it serves a call
+    and how it is called."""
+
+    def _post_attn_router_weight(self) -> Optional[torch.Tensor]:
+        gate = getattr(self.mlp, "gate", None)
+        return gate.weight if type(gate) is DeepseekV3Gate else None
+
+    def _can_use_post_attn_hop(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        router_weight: Optional[torch.Tensor],
+    ) -> bool:
+        """Whether the post-attention hop serves this call.
+
+        Decided from shapes/dtypes that are identical on all TP ranks, before
+        any device work.
+        """
+        return (
+            self.post_attn_hop is not None
+            and self.post_attention_layernorm.nvfp4_scale is None
+            and self.post_attn_hop.can_run(
+                hidden_states, residual, self.post_attention_layernorm.weight, router_weight
+            )
+        )
+
+    def _post_attn_hop(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        router_weight: Optional[torch.Tensor],
+    ):
+        return self.post_attn_hop(
+            hidden_states,
+            residual,
+            self.post_attention_layernorm.weight,
+            self.post_attention_layernorm.variance_epsilon,
+            router_weight,
+        )
+
+
+class _AdoptedLayerMixin:
+    """The fast-path state of a decoder layer or of the draft layer, and how it is attached.
+
+    The adoption steps run in the order the built-in constructor builds the modules they
+    wrap: the attention modules and the attention tail (the attention is built first), the
+    post-attention hop and its transports (built before the MLP), the router GEMM (the MoE's),
+    the post-MoE hops (after the MoE) and the attention tail's successor weights (last). They
+    are methods of the layer: the deferred-norm hop finds the layer it serves through its
+    caller's ``self`` (``kernels/deferred_norm/scope.py``).
+    """
+
+    # The fast-path objects the adoption attaches; None keeps the built-in computation.
+    post_attn_hop = None
+    _into_hop = None
+    _hop_routing = None
+    post_moe_hop = None
+    deferred_norm_hop = None
+    # Bound to the next MoE layer by the shell's ``setup_aliases``.
+    _expert_prefetch_link = None
+
+    def _adopt_fast_paths(self) -> None:
+        """Attach this layer's fast-path objects, in the built-in constructor's order."""
+        GlmMLA.adopt(self.self_attn)
+        self._adopt_pre_moe()
+        if isinstance(self.mlp, Deepseekv3MoE):
+            GlmMoE.adopt(self.mlp, self.model_config, self.layer_idx)
+        self._adopt_post_moe()
+        self._adopt_tail_prefetch()
+
+    def _adopt_pre_moe(self) -> None:
+        """The post-attention hop, the into-hop transport and the hop routing."""
+        config = self.model_config.pretrained_config
+        # The post-attention hop of the pre-MoE fused boundary. It is shared per TP group and
+        # created on every rank in layer order; it serves a layer only if that layer's
+        # all-reduce uses MNNVL.
+        if self.allreduce is not None and dispatch.is_enabled("post_attention_hop"):
+            self.post_attn_hop = dispatch.get_impl("post_attention_hop", "make_post_attn_hop")(
+                self.mapping,
+                config.hidden_size,
+                config.torch_dtype,
+                native_allreduce=self.allreduce,
+            )
+        # The attention tail publishes its partial straight into the hop's transport; the hop
+        # then reads it from there.
+        if (
+            self.post_attn_hop is not None
+            and dispatch.is_enabled("attention_output_into_hop")
+            and getattr(self.self_attn, "_attn_tail", None) is not None
+        ):
+            self._into_hop = dispatch.get_impl("post_attention_hop", "make_into_hop_transport")(
+                self.post_attn_hop
+            )
+        # The into-hop also computes the native top-8 routing, which the routed MoE consumes
+        # instead of routing again.
+        if self._into_hop is not None and dispatch.is_enabled("hop_routing"):
+            self._hop_routing = dispatch.get_impl("hop_routing", "make_transport")(self._into_hop)
+
+    def _adopt_post_moe(self) -> None:
+        """The post-MoE hop, its routed input and the deferred-norm hop (MoE layers only)."""
+        config = self.model_config.pretrained_config
+        layer_idx = self.layer_idx
+        if not (
+            config.n_routed_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % config.moe_layer_freq == 0
+        ):
+            return
+        # Fused finalize + MNNVL all-reduce + RMSNorm on the raw NVFP4 MoE output.
+        # Construction negotiates capability, including the raw-tuple contract the MoE module
+        # declares, and allocates its workspace collectively; it stays None (native path)
+        # unless every TP rank engages.
+        if (
+            dispatch.is_enabled("post_moe_hop")
+            and self.fusion_config.POST_MOE_FUSION
+            and self.is_nvfp4
+            and self.model_config.moe_backend == "TRTLLM"
+        ):
+            hop = dispatch.get_impl("post_moe_hop", "post_moe_hop")(
+                mapping=self.mapping,
+                hidden_size=config.hidden_size,
+                top_k=self.top_k,
+                dtype=config.torch_dtype,
+                raw_producer=self.mlp.experts,
+            )
+            self.post_moe_hop = hop if hop.engaged else None
+            # The post-MoE hop consumes the routed MoE's result directly (collective enable).
+            if (
+                self.post_moe_hop is not None
+                and dispatch.is_enabled("routed_moe_into_hop")
+                and self.post_moe_hop.enable_routed_input()
+            ):
+                self.mlp._routed_into_hop = self.post_moe_hop
+
+        # The hop defers the next layer's RMSNorm into that layer's projection. Only where the
+        # next layer is a target attention layer; construction is collective and the hop
+        # stays None unless every TP rank engages.
+        if (
+            self.post_moe_hop is not None
+            and dispatch.is_enabled("deferred_norm")
+            and dispatch.is_enabled("absorbed_query")
+            and layer_idx + 1 < config.num_hidden_layers
+        ):
+            # The evict-hints variant is the same deferred hop with another chain schedule
+            # (bit-identical packet).
+            deferred_norm_hop = dispatch.get_impl(
+                "deferred_norm_evict_hints"
+                if dispatch.is_enabled("deferred_norm_evict_hints")
+                else "deferred_norm",
+                "deferred_hop",
+            )(self.post_moe_hop)
+            self.deferred_norm_hop = deferred_norm_hop if deferred_norm_hop.engaged else None
+            # The deferred hop also serves the routed input the routed MoE writes into the
+            # post-MoE hop (collective enable; it stays off on every rank unless all of them
+            # engage).
+            if (
+                self.deferred_norm_hop is not None
+                and dispatch.is_enabled("routed_moe_into_hop")
+                and getattr(self.mlp, "_routed_into_hop", None) is not None
+            ):
+                self.deferred_norm_hop.enable_routed_input()
+
+    def _adopt_tail_prefetch(self) -> None:
+        """The router weight the attention tail stages into L2 (a cache hint)."""
+        # The attention tail also stages the router weight, which the op after the tail reads
+        # first, into L2 (no tensor changes).
+        if (
+            dispatch.is_enabled("attention_output_tail_prefetch")
+            and isinstance(self.mlp, Deepseekv3MoE)
+            and getattr(self.mlp, "gate", None) is not None
+            and getattr(self.self_attn, "_attn_tail", None) is not None
+        ):
+            self.self_attn._attn_tail.set_successor_weights([self.mlp.gate.weight])
+
+
+class GlmDecoderLayer(_AdoptedLayerMixin, _PostAttentionHopMixin, DeepseekV3DecoderLayer):
+    """A decoder layer: the built-in layer, with the family's decode composition."""
+
+    @classmethod
+    def adopt(cls, layer: DeepseekV3DecoderLayer) -> "GlmDecoderLayer":
+        """Re-class a built-in decoder layer and attach its fast-path objects.
+
+        Raises:
+            TypeError: if ``layer``, its attention modules or its MoE is not exactly the
+                built-in class.
+        """
+        reclass(layer, DeepseekV3DecoderLayer, cls)
+        layer._adopt_fast_paths()
+        return layer
+
+    def forward(
+        self,
+        position_ids: torch.IntTensor,
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        residual: torch.Tensor,
+        spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                residual=residual,
+                spec_metadata=spec_metadata,
+                **kwargs,
+            )
+        return self.forward_decode(
+            position_ids, hidden_states, attn_metadata, residual, spec_metadata, **kwargs
+        )
+
+    def forward_decode(
+        self,
+        position_ids: torch.IntTensor,
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        residual: torch.Tensor,
+        spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """The layer on a decode step."""
+        if residual is None:
+            residual = hidden_states
+            if self.input_layernorm.nvfp4_scale is not None:
+                # Layer-0 prologue: the residual-less input_layernorm folds the
+                # NVFP4 kv_a_proj input-quant via its attached .nvfp4_scale
+                # (RMSNorm.forward), returning an Fp4QuantizedTensor.
+                # return_norm_out stashes the BF16 view DSA's pre_indexer_proj
+                # needs on that tensor's unquantized_hidden_states.
+                hidden_states = self.input_layernorm(hidden_states, return_norm_out=True)[0]
+            else:
+                hidden_states = self.input_layernorm(hidden_states)
+        # Ask the attention tail to publish into the hop transport when this call will take
+        # the hop (rank-uniform facts, no device work).
+        into_hop_tail = None
+        if (
+            self._into_hop is not None
+            and isinstance(self.mlp, Deepseekv3MoE)
+            and self.fusion_config.PRE_MOE_FUSION
+            and self.post_attention_layernorm.nvfp4_scale is None
+            and self._into_hop.can_run(
+                residual, self.post_attention_layernorm.weight, self._post_attn_router_weight()
+            )
+        ):
+            into_hop_tail = self.self_attn._attn_tail
+            into_hop_tail.request_into_hop(self._into_hop, residual.shape[0])
+        # Self Attention
+        hidden_states = self.self_attn(
+            position_ids=position_ids,
+            hidden_states=hidden_states,
+            attn_metadata=attn_metadata,
+            all_reduce_params=AllReduceParams(enable_allreduce=not (self.disable_attn_allreduce)),
+            **kwargs,
+        )
+        into_hop_published = into_hop_tail is not None and into_hop_tail.take_into_hop_published()
+        residual = maybe_slice_for_helix_cp(
+            residual, attn_metadata, self.mapping_with_cp, self.layer_idx
+        )
+        if isinstance(self.mlp, Deepseekv3MoE):
+            if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+                self.fusion_config.POST_MOE_FUSION = False
+            return self.forward_MoE(
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                residual=residual,
+                spec_metadata=spec_metadata,
+                into_hop_published=into_hop_published,
+            )
+        else:
+            if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+                self.fusion_config.POST_MLP_FUSION = False
+            assert isinstance(self.mlp, GatedMLP)
+            return self.forward_mlp(
+                hidden_states=hidden_states,
+                residual=residual,
+                spec_metadata=spec_metadata,
+            )
+
+    def forward_MoE(
+        self,
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        residual: torch.Tensor,
+        spec_metadata: Optional[SpecMetadata] = None,
+        into_hop_published: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward_MoE(
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                residual=residual,
+                spec_metadata=spec_metadata,
+            )
+
+        def _run_MoE(hidden_states, hidden_states_fp4, do_finalize):
+            return self.mlp(
+                hidden_states,
+                hidden_states_fp4,
+                all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
+                final_all_reduce_params=AllReduceParams(
+                    enable_allreduce=not (
+                        self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1
+                    )
+                ),
+                do_finalize=do_finalize,
+                router_logits=router_logits,
+                hop_routing_out=hop_routing_out,
+            )
+
+        router_logits = None
+        hop_routing_out = None
+        router_weight = self._post_attn_router_weight()
+        routing_bias = self.mlp.gate.e_score_correction_bias if router_weight is not None else None
+        if (
+            into_hop_published
+            and self._hop_routing is not None
+            and self._hop_routing.can_run(
+                residual, self.post_attention_layernorm.weight, router_weight, routing_bias
+            )
+        ):
+            # The into-hop's outputs bit for bit, plus the native top-8 routing of every row
+            # (rank-uniform host-only decision).
+            (hidden_states, residual, router_logits, hop_routing_out) = self._hop_routing(
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.variance_epsilon,
+                router_weight,
+                routing_bias,
+            )
+        elif into_hop_published:
+            # The tail already published every rank's partial into the hop transport; same
+            # arithmetic and bits as the post-attention hop below.
+            hidden_states, residual, router_logits = self._into_hop(
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.variance_epsilon,
+                router_weight,
+            )
+        elif self.fusion_config.PRE_MOE_FUSION and self._can_use_post_attn_hop(
+            hidden_states, residual, router_weight
+        ):
+            # Same normalized hidden and residual bits as the fused all-reduce
+            # below. The router logits replace self.gate's: same FP32 GEMM,
+            # different summation order, so they are not bit-identical and a
+            # row whose 8th/9th expert scores lie within that rounding error
+            # can select a different expert set than native.
+            hidden_states, residual, router_logits = self._post_attn_hop(
+                hidden_states, residual, router_weight
+            )
+        elif self.fusion_config.PRE_MOE_FUSION:
+            # moe_backend can be either CUTLASS or TRTLLM here
+            # TODO: unify the two min-latency MoE backends by enabling quant fusion
+            hidden_states, residual = self.allreduce(
+                hidden_states,
+                all_reduce_params=AllReduceParams(
+                    fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                    residual=residual,
+                    norm_weight=self.post_attention_layernorm.weight,
+                    eps=self.post_attention_layernorm.variance_epsilon,
+                    trigger_completion_at_end=False,
+                ),
+            )
+        else:
+            # No fusion
+            hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+
+        # Cross-layer expert prefetch: predict the next MoE layer's routed experts from this
+        # layer's normalized hidden and prefetch their weights into L2 off the critical path
+        # (no tensor the model reads changes).
+        expert_prefetch = getattr(self, "_expert_prefetch_link", None)
+        if expert_prefetch is not None:
+            expert_prefetch.on_hop(hidden_states, router_logits)
+
+        # The post-MoE hop decision depends only on rank-identical state, so every TP rank
+        # takes the same path before the MoE runner produces any output.
+        use_post_moe_hop = (
+            self.post_moe_hop is not None
+            and self.fusion_config.POST_MOE_FUSION
+            and self.mlp.experts.has_nvfp4
+            and self.post_moe_hop.can_run(hidden_states.shape[0])
+        )
+
+        # Note: this fusion pattern is only supported for single-node TRTLLM-nvfp4 backend now
+        do_finalize = not use_post_moe_hop and (
+            self.mapping.is_multi_node()
+            or (
+                not (
+                    self.fusion_config.POST_MOE_FUSION
+                    and hidden_states.shape[0] <= self.moe_allreduce.max_token
+                    and self.model_config.moe_backend == "TRTLLM"
+                    and self.mlp.experts.has_nvfp4
+                    and self.is_p2p_supported
+                )
+            )
+        )
+
+        hidden_states = _run_MoE(hidden_states, hidden_states_fp4=None, do_finalize=do_finalize)
+        if expert_prefetch is not None:
+            expert_prefetch.after_moe()
+
+        if self.fusion_config.POST_MOE_FUSION:
+            if do_finalize:
+                hidden_states, residual = self.allreduce(
+                    hidden_states,
+                    all_reduce_params=AllReduceParams(
+                        fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                        residual=residual,
+                        norm_weight=self.next_layer_layernorm.weight,
+                        eps=self.next_layer_layernorm.variance_epsilon,
+                        trigger_completion_at_end=False,
+                    ),
+                )
+            elif use_post_moe_hop and len(hidden_states) == 2:
+                # [shared, routed], with routed written into the hop by the routed MoE.
+                shared_output, routed_output = hidden_states
+                # The deferred-norm hop reads the same routed input with the same collective
+                # protocol and returns a deferred-norm packet (rank-identical decision).
+                hop = (
+                    self.deferred_norm_hop
+                    if self.deferred_norm_hop is not None
+                    and self.deferred_norm_hop.can_run_routed(residual.shape[0], attn_metadata)
+                    else self.post_moe_hop
+                )
+                hidden_states, residual = hop.hop_routed(
+                    routed_output,
+                    shared_output,
+                    residual,
+                    self.next_layer_layernorm.weight,
+                    self.next_layer_layernorm.variance_epsilon,
+                )
+            elif use_post_moe_hop:
+                (shared_output, fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx) = (
+                    hidden_states
+                )
+                # The deferred-norm hop: same inputs and collective protocol; returns a
+                # deferred-norm packet for the next layer instead of the normalized hidden
+                # state (rank-identical decision).
+                hop = (
+                    self.deferred_norm_hop
+                    if self.deferred_norm_hop is not None
+                    and self.deferred_norm_hop.can_run(residual.shape[0], attn_metadata)
+                    else self.post_moe_hop
+                )
+                hidden_states, residual = hop(
+                    fc2_output,
+                    expert_scale_factor,
+                    expanded_idx_to_permuted_idx,
+                    shared_output,
+                    residual,
+                    self.next_layer_layernorm.weight,
+                    self.next_layer_layernorm.variance_epsilon,
+                )
+            else:
+                assert len(hidden_states) == 4, "hidden_states must have 4 elements"
+
+                shared_output = hidden_states[0]
+                fc2_output = hidden_states[1]
+                expert_scale_factor = hidden_states[2]
+                expanded_idx_to_permuted_idx = hidden_states[3]
+
+                moe_all_reduce_params = MoEAllReduceParams(
+                    expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+                    expert_scale_factor=expert_scale_factor,
+                    shared_expert_output=shared_output,
+                    residual=residual,
+                    norm_weight=self.next_layer_layernorm.weight,
+                    eps=self.next_layer_layernorm.variance_epsilon,
+                    is_cutlass_min_latency=False,
+                )
+                hidden_states, residual = self.moe_allreduce(
+                    fc2_output, all_reduce_params=moe_all_reduce_params
+                )
+        else:
+            if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+                spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual)
+            hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual)
+
+        if expert_prefetch is not None:
+            expert_prefetch.after_post_moe()
+        return hidden_states, residual
+
+    def forward_mlp(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        spec_metadata: Optional[SpecMetadata] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward_mlp(
+                hidden_states=hidden_states,
+                residual=residual,
+                spec_metadata=spec_metadata,
+            )
+
+        if self.fusion_config.PRE_MLP_FUSION:
+            if self.mlp.gate_up_proj.has_nvfp4:
+                act_fp4, act_sf, residual = self.allreduce(
+                    hidden_states,
+                    all_reduce_params=AllReduceParams(
+                        fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4,
+                        residual=residual,
+                        norm_weight=self.post_attention_layernorm.weight,
+                        scale=self.mlp.gate_up_proj.input_scale,
+                        eps=self.post_attention_layernorm.variance_epsilon,
+                    ),
+                )
+                hidden_states = Fp4QuantizedTensor(act_fp4, act_sf)
+            else:
+                hidden_states, residual = self.allreduce(
+                    hidden_states,
+                    all_reduce_params=AllReduceParams(
+                        fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                        residual=residual,
+                        norm_weight=self.post_attention_layernorm.weight,
+                        eps=self.post_attention_layernorm.variance_epsilon,
+                    ),
+                )
+        elif self.post_attention_layernorm.nvfp4_scale is not None:
+            # Attention-DP path (no allreduce): post_attention_layernorm folds
+            # the dense MLP's NVFP4 gate_up_proj input-quant into one kernel via
+            # its attached .nvfp4_scale (RMSNorm.forward), mirroring the
+            # PRE_MLP_FUSION quant but without the allreduce.
+            hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        else:
+            # No fusion
+            # We need to add twoshot allreduce here to avoid modifying MLA logic
+            hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+
+        final_all_reduce_params = AllReduceParams(
+            enable_allreduce=not (self.fusion_config.POST_MLP_FUSION or self.mlp_tp_size == 1)
+        )
+        mlp_output = None
+        if dispatch.is_enabled("dense_mlp"):
+            # Eligibility is decided on the host before any device work;
+            # None selects the native module below.
+            mlp_output = dispatch.get_impl("dense_mlp", "dense_mlp_forward")(
+                self.mlp, hidden_states, final_all_reduce_params
+            )
+        if mlp_output is None:
+            mlp_output = self.mlp(hidden_states, final_all_reduce_params=final_all_reduce_params)
+        hidden_states = mlp_output
+
+        if self.fusion_config.POST_MLP_FUSION:
+            hidden_states, residual = self.allreduce(
+                hidden_states,
+                all_reduce_params=AllReduceParams(
+                    fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                    residual=residual,
+                    norm_weight=self.next_layer_layernorm.weight,
+                    eps=self.next_layer_layernorm.variance_epsilon,
+                ),
+            )
+        else:
+            if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+                spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual)
+            hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual)
+
+        return hidden_states, residual
+
+
+class GlmMtpLayer(_AdoptedLayerMixin, _PostAttentionHopMixin, DeepseekV3MTP):
+    """The draft layer: the built-in MTP layer, with the family's decode composition."""
+
+    # The draft front packs this rank's chunk of the projection input; None keeps the
+    # built-in norms and concatenation.
+    draft_front = None
+    # The live-row hop of the first draft step and the rows the draft loop arms it with.
+    draft_live_row_hop = None
+    _draft_live_rows = None
+    _draft_live_rows_compacted = False
+
+    @classmethod
+    def adopt(cls, layer: DeepseekV3MTP) -> "GlmMtpLayer":
+        """Re-class the built-in draft layer and attach its fast-path objects.
+
+        Raises:
+            TypeError: if ``layer``, its attention modules or its MoE is not exactly the
+                built-in class.
+        """
+        reclass(layer, DeepseekV3MTP, cls)
+        layer._adopt_fast_paths()
+        layer._adopt_draft()
+        return layer
+
+    def _adopt_draft(self) -> None:
+        """The draft front and the live-row hop, after the draft layer's own modules."""
+        # Normalize and pack only this rank's chunk of the eh_proj input. Non-ADP TP only, and
+        # only while enorm/hnorm run the FlashInfer RMSNorm the draft front mirrors.
+        if dispatch.is_enabled("draft_front"):
+            mapping = self.model_config.mapping
+            eligible = not mapping.enable_attention_dp and mapping.tp_size > 1
+            reason = "" if eligible else "requires non-ADP tensor parallelism"
+            if eligible:
+                reason = (
+                    dispatch.get_impl("draft_front", "native_norm_reason")(self.enorm, self.hnorm)
+                    or ""
+                )
+                eligible = not reason
+            dispatch.log_capability("draft_front", "draft_front", eligible, reason)
+            if eligible:
+                self.draft_front = dispatch.get_impl("draft_front", "draft_front")
+
+        # On draft step 0 the draft loop arms the layer with the live rows it selected
+        # (``arm_live_rows``); the post-attention hop then keeps only those rows, and the MoE
+        # and the post-MoE all-reduce run on them.
+        if self.post_attn_hop is not None and dispatch.is_enabled("draft_live_rows"):
+            self.draft_live_row_hop = dispatch.get_impl("draft_live_rows", "make_live_row_hop")(
+                self.post_attn_hop
+            )
+
+    def arm_live_rows(self, rows: torch.Tensor) -> None:
+        """Arm the next forward with the live rows of draft step 0 (``rows``: int64 row ids)."""
+        self._draft_live_rows = rows
+        self._draft_live_rows_compacted = False
+
+    def disarm_live_rows(self) -> bool:
+        """Clear the arming; returns whether the armed forward kept only the live rows."""
+        compacted = bool(self._draft_live_rows_compacted)
+        self._draft_live_rows = None
+        self._draft_live_rows_compacted = False
+        return compacted
+
+    def forward(
+        self,
+        input_ids: torch.IntTensor,
+        position_ids: torch.IntTensor,
+        hidden_states: torch.Tensor,
+        embed_tokens: Embedding,
+        attn_metadata: AttentionMetadata,
+        all_rank_num_tokens: Optional[List[int]] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super().forward(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                embed_tokens=embed_tokens,
+                attn_metadata=attn_metadata,
+                all_rank_num_tokens=all_rank_num_tokens,
+                spec_metadata=spec_metadata,
+                **kwargs,
+            )
+        # Split hidden_states columnwise based on TP
+        tp_size = self.model_config.mapping.tp_size
+        tp_rank = self.model_config.mapping.tp_rank
+
+        embeds = None
+        packed = None
+        if self.draft_front is not None:
+            # Returns None, before launching anything, outside its support
+            # domain; the native path below then reuses the embeddings.
+            embeds = embed_tokens(input_ids)
+            packed = self.draft_front(
+                embeds,
+                hidden_states,
+                self.enorm.weight,
+                self.enorm.variance_epsilon,
+                self.hnorm.weight,
+                self.hnorm.variance_epsilon,
+                tp_rank,
+                tp_size,
+            )
+
+        if packed is not None:
+            hidden_states = packed
+        else:
+
+            def norm_embeds():
+                return self.enorm(
+                    embed_tokens(input_ids) if embeds is None else embeds
+                )  # emdedding
+
+            def norm_hidden():
+                return self.hnorm(hidden_states)
+
+            inputs_embeds, hidden_states = maybe_execute_in_parallel(
+                norm_embeds,
+                norm_hidden,
+                self.event_dict[EventType.Main],
+                self.event_dict[EventType.MoeShared],
+                self.aux_stream,
+                disable_on_compile=True,
+            )
+            hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
+
+            if tp_size > 1 and not (self.model_config.mapping.enable_attention_dp):
+                hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank]
+        hidden_states = self.eh_proj(hidden_states)
+
+        # Input layer norm
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states = self.self_attn(
+            position_ids=position_ids,
+            hidden_states=hidden_states,
+            attn_metadata=attn_metadata,
+            all_reduce_params=AllReduceParams(enable_allreduce=not (self.disable_attn_allreduce)),
+            **kwargs,
+        )
+
+        # MTP Layer Must have sparse MOE
+        router_logits = None
+        router_weight = self._post_attn_router_weight()
+        use_post_attn_hop = self.fusion_config.PRE_MOE_FUSION and self._can_use_post_attn_hop(
+            hidden_states, residual, router_weight
+        )
+        live_rows = self._draft_live_rows
+        if (
+            use_post_attn_hop
+            and live_rows is not None
+            and self.draft_live_row_hop is not None
+            and self.draft_live_row_hop.can_run(hidden_states, live_rows, router_weight)
+        ):
+            # The live rows only, each as the post-attention hop below.
+            hidden_states, residual, router_logits = self.draft_live_row_hop(
+                hidden_states,
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.variance_epsilon,
+                router_weight,
+                live_rows,
+            )
+            self._draft_live_rows_compacted = True
+            if all_rank_num_tokens is not None:
+                all_rank_num_tokens = [hidden_states.shape[0]] * len(all_rank_num_tokens)
+        elif use_post_attn_hop:
+            # Numerics as in DeepseekV3DecoderLayer.forward_MoE.
+            hidden_states, residual, router_logits = self._post_attn_hop(
+                hidden_states, residual, router_weight
+            )
+        elif self.fusion_config.PRE_MOE_FUSION:
+            hidden_states, residual = self.allreduce(
+                hidden_states,
+                all_reduce_params=AllReduceParams(
+                    fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                    residual=residual,
+                    norm_weight=self.post_attention_layernorm.weight,
+                    eps=self.post_attention_layernorm.variance_epsilon,
+                ),
+            )
+        else:
+            hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+
+        # MoE
+        hidden_states = self.mlp(
+            hidden_states,
+            all_rank_num_tokens=all_rank_num_tokens,
+            final_all_reduce_params=AllReduceParams(
+                enable_allreduce=not (
+                    self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1
+                )
+            ),
+            router_logits=router_logits,
+        )
+
+        if self.fusion_config.POST_MOE_FUSION:
+            hidden_states, residual = self.allreduce(
+                hidden_states,
+                all_reduce_params=AllReduceParams(
+                    fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                    residual=residual,
+                    norm_weight=self.shared_head.norm.weight,
+                    eps=self.shared_head.norm.variance_epsilon,
+                ),
+            )
+        else:
+            hidden_states, _ = self.shared_head.norm(hidden_states, residual)
+
+        # It's for 2-model path, capture the hidden states
+        if spec_metadata is not None:
+            spec_metadata.maybe_capture_hidden_states(0, hidden_states, None)
+
+        return hidden_states
+
+
+def adopt_model(model: "GlmMoeDsaForCausalLM", plan: dispatch.Plan) -> None:
+    """Re-class the built-in decoder layers, the draft layer, their attention modules, their
+    MoE modules and the speculative worker, and attach the fast-path objects ``plan`` enables.
+
+    Layer by layer in index order, the decoder layers first and the draft layer last, as the
+    built-in model builds them; within a layer in the order its constructor builds the modules
+    the objects wrap (this module). The speculative worker comes last, as the built-in model
+    builds it after the draft layer; its draft-loop fast paths are decided per forward. The
+    fast paths read the installed plan, which is ``plan``.
+
+    Raises:
+        TypeError: if a layer, its attention modules, its MoE or the speculative worker is not
+            exactly the built-in class.
+    """
+    core = model.model
+    num_layers = core.num_hidden_layers
+    for layer in core.layers[:num_layers]:
+        GlmDecoderLayer.adopt(layer)
+    for layer in core.layers[num_layers:]:
+        GlmMtpLayer.adopt(layer)
+    if model.spec_worker is not None:
+        GlmMtpWorker.adopt(model.spec_worker)
+
+
+# The fast-path objects a decoder layer (and the draft layer) carries: per submodule (empty for
+# the layer itself), the attributes that hold them. An object a target's plan never attaches
+# reads None and is left out of the state.
+_LAYER_STATE: tuple[tuple[str, tuple[str, ...]], ...] = (
+    (
+        "",
+        (
+            "post_attn_hop",
+            "_into_hop",
+            "_hop_routing",
+            "post_moe_hop",
+            "deferred_norm_hop",
+            "_expert_prefetch_link",
+            "draft_front",
+            "draft_live_row_hop",
+        ),
+    ),
+    ("self_attn", ("_attn_tail",)),
+    ("mlp", ("router_gemm", "_routed_into_hop")),
+)
+# Attributes of a fast-path object that every rank must agree on; none of them holds a rank,
+# a device or an address.
+_STATE_FLAGS = ("engaged", "impl", "routed_mode", "routed_engaged")
+
+
+def adopted_state(model: "GlmMoeDsaForCausalLM") -> dispatch.AdoptedState:
+    """The fast-path state attached to the decoder layers, layer by layer.
+
+    Per layer, each attached object as its attribute, its class name and its
+    ``_STATE_FLAGS``; ``dispatch.check_rank_uniform_adopted`` compares it across the ranks.
+    """
+    state = []
+    for index, layer in enumerate(model.model.layers):
+        attached = []
+        for submodule, names in _LAYER_STATE:
+            owner = getattr(layer, submodule, None) if submodule else layer
+            for name in names:
+                obj = getattr(owner, name, None)
+                if obj is None:
+                    continue
+                flags = tuple(
+                    f"{flag}={getattr(obj, flag)!r}" for flag in _STATE_FLAGS if hasattr(obj, flag)
+                )
+                where = f"{submodule}.{name}" if submodule else name
+                attached.append((where, type(obj).__name__) + flags)
+        state.append((f"model.layers.{index}", tuple(attached)))
+    return tuple(state)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/runtime.py
new file mode 100644
index 000000000000..2a64a12ea65a
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/runtime.py
@@ -0,0 +1,103 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The step-loop hooks a glm_moe_dsa target declares to the engine (``RuntimeHooks``).
+
+``make_runtime_hooks`` declares the engine-side fast paths that the plan runs: the
+decode step's input preparation recorded and issued as one launch
+(``deferred_input_prep``), the speculative sampler's store update and host copies
+in one launch (``sampler_store``), the overlap executor loop on the execution stream
+(``execution_stream_loop``), and the executor events of the post-MoE hop health
+check (``post_moe_hop``) and the host page retention policy (``host_page_retention``).
+Every hook it leaves undeclared stays off, and the engine runs that part of the step
+loop as it does for the built-in model.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
+
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+    from tensorrt_llm.mapping import Mapping
+
+    from .dispatch import Plan
+
+
+class GlmExecutorEvents:
+    """The executor events of a glm_moe_dsa target (``ExecutorEvents``).
+
+    A kernel package is imported only when its fast path runs.
+    """
+
+    def __init__(self, plan: "Plan", mapping: "Mapping") -> None:
+        self._plan = plan
+        self._mapping = mapping
+
+    def after_warmup(self) -> None:
+        if self._plan.enabled("post_moe_hop"):
+            from .kernels.post_moe_hop import check_health
+
+            # Collective over each post-MoE hop workspace's TP group: every rank
+            # fails startup together when any hop reported a device error during
+            # warmup or graph capture.
+            check_health("after_warmup", collective=True)
+        if self._plan.enabled("host_page_retention"):
+            from .kernels.host_page_retention import install_host_policy
+
+            # From here on this rank's CPU allocator keeps freed pages instead of
+            # returning them to the OS at once, so host work during a decode step
+            # does not invalidate the GPU's address translations.
+            install_host_policy("after_warmup")
+
+    def before_overlap_loop(self) -> None:
+        if self._plan.enabled("sampler_store") or self._plan.enabled("execution_stream_loop"):
+            from .kernels.deferred_input_prep import compat
+
+            # The store writer engages only after an overlap loop without attention
+            # DP declared its domain; logs the decision once per process.
+            compat.declare_overlap_loop(self._mapping.enable_attention_dp)
+            compat.loop_domain()
+
+    def at_shutdown(self) -> None:
+        if self._plan.enabled("post_moe_hop"):
+            from .kernels.post_moe_hop import check_health
+
+            # End-of-serving report of this rank; rank-local and non-raising inside
+            # cleanup (a device error is logged at ERROR).
+            check_health("end_of_serving", raise_on_error=False)
+
+
+def make_runtime_hooks(model: "GlmMoeDsaForCausalLM") -> RuntimeHooks:
+    """The model's hooks: every engine-side fast path its plan runs.
+
+    A kernel package is imported only when its fast path runs, as ``dispatch.get_impl``
+    imports the packages of the model's own fast paths.
+    """
+    plan = model.plan
+    step_input_stager = None
+    if plan.enabled("deferred_input_prep"):
+        from .kernels.deferred_input_prep import PrepStager
+
+        # The guard (debug level 1) runs under TRTLLM_MODELING_V2_VALIDATE.
+        step_input_stager = PrepStager(debug_level=1 if plan.validate else 0)
+    sampler_store_writer = None
+    if plan.enabled("sampler_store"):
+        from .kernels.deferred_input_prep import SamplerStoreWriter
+
+        sampler_store_writer = SamplerStoreWriter()
+    decode_prologue = None
+    if plan.enabled("decode_prologue"):
+        from .kernels.decode_prologue import DecodePrologueHook
+
+        # The decode target takes the embedding rows of the prologue's launch from the
+        # same object (the core's ``decode_prologue``).
+        decode_prologue = model.model.decode_prologue = DecodePrologueHook()
+    return RuntimeHooks(
+        loop_on_execution_stream=plan.enabled("execution_stream_loop"),
+        step_input_stager=step_input_stager,
+        sampler_store_writer=sampler_store_writer,
+        decode_prologue=decode_prologue,
+        executor_events=GlmExecutorEvents(plan, model.model_config.mapping),
+    )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/spec_worker.py
new file mode 100644
index 000000000000..4081133efb4f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/spec_worker.py
@@ -0,0 +1,621 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The speculative worker of the glm_moe_dsa targets: the one-model MTP-Eagle worker with the
+draft-loop fast paths.
+
+``GlmMtpWorker.adopt`` re-classes the built-in ``MTPEagleWorker`` the shell built in place, so
+the worker the engine and the shell's epilogue hold stays the same object. The class adds
+methods and a class-level default only; nothing is built at adoption, every fast path of the
+draft loop is decided per forward.
+
+Each override runs the family's draft loop when the step is a decode step
+(``dispatch.Phase.DECODE``) and the built-in method otherwise:
+
+* ``_forward_impl`` accepts the target's tokens and prepares the first draft input and the next
+  new tokens through the draft-loop operations (``draft_loop``);
+* ``_forward_linear_draft_loop`` runs the draft steps with the draft tail (``draft_tail``), the
+  step-0 live rows of the draft layer (``draft_live_rows``) and the draft-loop step updates;
+* ``_run_draft_forward`` calls the draft layer with the draft-loop embedding.
+"""
+
+import contextlib
+
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata
+from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager
+from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker, _reset_mtp_index_share
+from tensorrt_llm._torch.speculative.mtp import _select_mtp_position_ids
+
+from . import dispatch
+from .attention import reclass
+
+
+class GlmMtpWorker(MTPEagleWorker):
+    """The one-model MTP-Eagle worker, with the family's draft loop on decode steps."""
+
+    # The draft-loop operations of the running forward; None keeps the built-in acceptance,
+    # first draft inputs, step updates and draft-token stacking.
+    _draft_loop_ops = None
+
+    @classmethod
+    def adopt(cls, worker: MTPEagleWorker) -> "GlmMtpWorker":
+        """Re-class the built-in MTP-Eagle worker.
+
+        Raises:
+            TypeError: if ``worker`` is not exactly the built-in ``MTPEagleWorker``.
+        """
+        reclass(worker, MTPEagleWorker, cls)
+        return worker
+
+    def _forward_impl(
+        self,
+        input_ids,
+        position_ids,
+        hidden_states,
+        logits,
+        attn_metadata,
+        spec_metadata,
+        draft_model,
+        resource_manager=None,
+    ):
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super()._forward_impl(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                logits=logits,
+                attn_metadata=attn_metadata,
+                spec_metadata=spec_metadata,
+                draft_model=draft_model,
+                resource_manager=resource_manager,
+            )
+
+        runtime_draft_len = spec_metadata.runtime_draft_len
+        # skip the draft forward if the runtime draft length is 0
+        if runtime_draft_len == 0:
+            return self.skip_drafting(
+                input_ids,
+                position_ids,
+                hidden_states,
+                logits,
+                attn_metadata,
+                spec_metadata,
+                draft_model,
+            )
+
+        batch_size = attn_metadata.num_seqs
+        num_contexts = attn_metadata.num_contexts
+        num_gens = batch_size - num_contexts
+
+        raw_logits = logits
+
+        self._execute_guided_decoder_if_present(logits)
+
+        # The draft loop's fused bookkeeping (draft_loop), decided on the host before any
+        # device work; None keeps the native code below.
+        draft_loop_ops = None
+        if self.is_mtp_eagle and dispatch.is_enabled("draft_loop"):
+            draft_loop_ops = dispatch.get_impl("draft_loop", "draft_loop_ops")(
+                self, attn_metadata, spec_metadata
+            )
+        self._draft_loop_ops = draft_loop_ops
+
+        # Sample and accept tokens. ``input_ids`` is required by the relaxed-
+        # acceptance path (scans for thinking-phase tokens); ignored otherwise.
+        if draft_loop_ops is not None:
+            accepted_tokens, num_accepted_tokens = draft_loop_ops.accept(logits)
+        else:
+            accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens(
+                input_ids, logits, attn_metadata, spec_metadata
+            )
+
+        # Mamba hybrid models need state updates after token acceptance because
+        # the accepted token count affects which Mamba states are valid. The
+        # isinstance check below naturally no-ops on non-Mamba kv_cache_managers,
+        # so this is safe to run unconditionally regardless of spec mode (Eagle3
+        # over a Mamba-style draft is plausible, even if no such draft exists today).
+        if self._is_mamba_hybrid_cache is None:
+            self._is_mamba_hybrid_cache = isinstance(
+                attn_metadata.kv_cache_manager, MambaHybridCacheManager
+            )
+        if num_gens > 0 and self._is_mamba_hybrid_cache:
+            attn_metadata.kv_cache_manager.update_mamba_states(
+                attn_metadata=attn_metadata,
+                num_accepted_tokens=num_accepted_tokens,
+                state_indices=attn_metadata.mamba_metadata.state_indices,
+            )
+
+        sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None)
+        if self.sa_enhancer is not None and sa_manager is not None:
+            self.sa_enhancer.extend_and_prepare(
+                sa_manager=sa_manager,
+                request_ids=spec_metadata.request_ids,
+                accepted_tokens=accepted_tokens,
+                num_accepted_tokens=num_accepted_tokens,
+                num_gens=num_gens,
+                num_contexts=num_contexts,
+                max_draft_len=runtime_draft_len,
+            )
+
+        if self._uses_external_shared_target_kv:
+            next_draft_tokens = self._forward_external_shared_target_kv_draft_loop(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+                spec_metadata=spec_metadata,
+                draft_model=draft_model,
+                accepted_tokens=accepted_tokens,
+                num_accepted_tokens=num_accepted_tokens,
+                num_contexts=num_contexts,
+                batch_size=batch_size,
+            )
+        else:
+            # Save the old attn_metadata and spec_metadata
+            self._prepare_attn_metadata_for_spec_dec(attn_metadata)
+
+            # Prepare inputs for the 1st draft model forward
+            position_ids = position_ids.squeeze(0)
+            if draft_loop_ops is not None:
+                inputs = draft_loop_ops.first_drafter_inputs(position_ids, hidden_states)
+            else:
+                inputs = self.prepare_1st_drafter_inputs(
+                    input_ids=input_ids,
+                    position_ids=position_ids,
+                    hidden_states=hidden_states,
+                    accepted_tokens=accepted_tokens,
+                    attn_metadata=attn_metadata,
+                    spec_metadata=spec_metadata,
+                    draft_model=draft_model,
+                )
+
+            # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here
+            # so the post-loop restore (below) can put attn_metadata back into a
+            # state the target model expects.
+            original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens
+
+            # Get the draft KV cache manager if using separate layouts
+            draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager)
+
+            next_draft_tokens = self._forward_draft_loop(
+                inputs,
+                attn_metadata,
+                spec_metadata,
+                draft_model,
+                draft_kv_cache_manager,
+                num_contexts,
+                num_gens,
+                batch_size,
+                num_accepted_tokens,
+                original_all_rank_num_tokens,
+                resource_manager,
+            )
+            # restore attn_metadata to support cuda graph
+            self._restore_attn_metadata_from_spec_dec(attn_metadata)
+            # restore all_rank_num_tokens for attention DP
+            if original_all_rank_num_tokens is not None:
+                attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
+
+        # prepare next new tokens to support overlap scheduler
+        if draft_loop_ops is not None and draft_loop_ops.next_new_tokens is not None:
+            next_new_tokens = draft_loop_ops.next_new_tokens
+        else:
+            next_new_tokens = self._prepare_next_new_tokens(
+                accepted_tokens,
+                next_draft_tokens,
+                spec_metadata.batch_indices_cuda,
+                batch_size,
+                num_accepted_tokens,
+            )
+        self._draft_loop_ops = None
+
+        attn_metadata.use_spec_decoding = True
+
+        # Commit model-owned target state only after draft preparation has
+        # succeeded; SpecWorkerBase aborts the still-pending snapshots on any
+        # exception above this point.
+        if self._auxiliary_state_handlers:
+            self.commit_auxiliary_speculative_states(
+                num_accepted_tokens,
+                attn_metadata.mamba_metadata.state_indices[:batch_size],
+                num_contexts,
+            )
+
+        return {
+            "logits": raw_logits,
+            "new_tokens": accepted_tokens,
+            "new_tokens_lens": num_accepted_tokens,
+            "next_draft_tokens": next_draft_tokens,
+            "next_new_tokens": next_new_tokens,
+        }
+
+    def _forward_linear_draft_loop(
+        self,
+        inputs,
+        attn_metadata,
+        spec_metadata,
+        draft_model,
+        draft_kv_cache_manager,
+        num_contexts,
+        batch_size,
+        num_accepted_tokens,
+        original_all_rank_num_tokens,
+    ):
+        """Linear draft loop, unified for Eagle3 and MTP Eagle."""
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super()._forward_linear_draft_loop(
+                inputs=inputs,
+                attn_metadata=attn_metadata,
+                spec_metadata=spec_metadata,
+                draft_model=draft_model,
+                draft_kv_cache_manager=draft_kv_cache_manager,
+                num_contexts=num_contexts,
+                batch_size=batch_size,
+                num_accepted_tokens=num_accepted_tokens,
+                original_all_rank_num_tokens=original_all_rank_num_tokens,
+            )
+        from tensorrt_llm._torch.attention.backends.sparse.dsa import is_dsa_cache_manager
+
+        runtime_draft_len = spec_metadata.runtime_draft_len
+        num_gens = batch_size - num_contexts
+        next_draft_tokens = []
+        # The draft-loop operations run only without context rows, so the last-token scan is
+        # unused with them.
+        draft_loop_ops = getattr(self, "_draft_loop_ops", None)
+        if draft_loop_ops is None:
+            last_tokens_idx = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
+        position_ids = inputs["position_ids"]
+
+        # Sparse backends that can reuse one indexer selection across the draft
+        # loop expose this setter; each indexer still decides whether its own
+        # config turns the reuse on.
+        uses_mtp_index_share = self.is_mtp_eagle and isinstance(
+            attn_metadata, MTPIndexShareMetadata
+        )
+        # The draft tail (draft_tail): one launch for the draft LM head, the vocab-shard max
+        # and the cross-rank argmax. Decided once per forward, before any device work,
+        # identically on every TP rank.
+        draft_tail = None
+        if self.is_mtp_eagle and dispatch.is_enabled("draft_tail"):
+            draft_tail = dispatch.get_impl("draft_tail", "select_draft_tail")(
+                self, draft_model, spec_metadata
+            )
+        # The step-0 live rows (draft_live_rows): draft step 0 runs the MTP layer's MoE and
+        # post-MoE all-reduce on the live rows the draft-loop operations selected only, and
+        # the layer returns those rows. Decided once per forward, before any device work,
+        # identically on every TP rank.
+        live_rows_step0 = None
+        if (
+            self.is_mtp_eagle
+            and draft_loop_ops is not None
+            and dispatch.is_enabled("draft_live_rows")
+        ):
+            live_rows_step0 = dispatch.get_impl("draft_live_rows", "select_step0")(
+                self, draft_model, spec_metadata, draft_loop_ops
+            )
+        with contextlib.ExitStack() as draft_scope:
+            if uses_mtp_index_share:
+                attn_metadata.set_in_mtp_draft_loop(True)
+                # Accepted counts let the indexer stash each gen's last-accepted row.
+                attn_metadata.set_mtp_num_accepted(num_accepted_tokens)
+                # The metadata outlives this call, so clear the state even when a
+                # draft step raises: a stale draft-loop flag would make the next
+                # target forward reuse these selections. The loop writes
+                # set_skip_topk on the metadata yielded below, which is this same
+                # object: draft_kv_cache_context swaps buffers in place for every
+                # TrtllmAttentionMetadata, and sparse metadata all derive from it.
+                draft_scope.callback(_reset_mtp_index_share, attn_metadata)
+            draft_attn_metadata = draft_scope.enter_context(
+                self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager)
+            )
+            attn_metadata = draft_attn_metadata
+            inputs["attn_metadata"] = draft_attn_metadata
+            if self.is_mtp_eagle and is_dsa_cache_manager(draft_kv_cache_manager):
+                # Overlap scheduling corrects kv_lens_cuda from the runtime
+                # accepted-token counts inside the captured graph. The target
+                # forward refreshes target slot mappings during that correction;
+                # refresh once more after rebinding so the first draft forward
+                # writes the separate DSA indexer cache at the same positions.
+                attn_metadata.on_update_kv_lens()
+            for i in range(runtime_draft_len):
+                if uses_mtp_index_share:
+                    attn_metadata.set_skip_topk(i > 0)
+                # Run draft model (mode-specific via helper). The helper
+                # passes ``all_rank_num_tokens`` as a kwarg so the draft model
+                # handles save/restore internally (Eagle3DraftModel.forward
+                # uses try/finally); attn_metadata is left untouched here.
+                live_rows_compacted = False
+                if i == 0 and live_rows_step0 is not None:
+                    with live_rows_step0.armed() as live_rows:
+                        hidden_states, hidden_states_to_save = self._run_draft_forward(
+                            draft_model, inputs, spec_metadata, i
+                        )
+                    live_rows_compacted = live_rows.compacted
+                else:
+                    hidden_states, hidden_states_to_save = self._run_draft_forward(
+                        draft_model, inputs, spec_metadata, i
+                    )
+
+                # Compute gather_ids: on the first draft step each generation
+                # request may have accepted multiple tokens, so we index into
+                # the flattened token sequence to find the last accepted one.
+                # From step 1 onwards the loop feeds exactly one token per
+                # request, in request order, so every per-request row
+                # selection is the prefix slice [0, batch_size); each
+                # ``[gather_ids]`` site below takes a slice as readily as an
+                # index tensor.
+                draft_rows = None
+                if live_rows_compacted:
+                    # The MTP layer already returned the live rows.
+                    gather_ids, draft_rows = draft_loop_ops.gather_ids, hidden_states
+                elif i == 0 and draft_loop_ops is not None:
+                    gather_ids, draft_rows = draft_loop_ops.step0_rows(hidden_states)
+                elif i == 0:
+                    start_ids_gen = (
+                        spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1)
+                    ).long()
+                    gather_ids_gen = (
+                        start_ids_gen
+                        + num_accepted_tokens[num_contexts:]
+                        - 1
+                        + attn_metadata.num_ctx_tokens
+                    )
+                    gather_ids = torch.concat(
+                        [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0
+                    )
+                else:
+                    gather_ids = slice(0, batch_size)
+
+                if self.guided_decoder is not None:
+                    new_tokens = inputs["input_ids"][gather_ids]
+                    self.guided_decoder.add_draft_batch(
+                        new_tokens, num_accepted_tokens, draft_step=i
+                    )
+
+                # Compute logits.
+                # MTP Eagle: shared_head of the MTP layer, with optional
+                #   ADP+LM-head-TP padding to ``max_num_requests`` so every TP
+                #   rank produces logits of the same shape.
+                # Eagle3: logits_processor of the EAGLE draft model.
+                #
+                # The LM-head-TP fast path (group-stacked rows, vocab-sharded
+                # weight) is only usable when the consumer is an argmax: the
+                # distributed greedy sampler recovers the global argmax from
+                # vocab-shard maxima without materializing full distributions.
+                # Advanced (rejection) sampling needs each request's full-vocab
+                # distribution, so a batch headed for the advanced path
+                # bypasses LM-head-TP and computes full-vocab logits locally --
+                # under ADP the lm_head weight is replicated (the sliced-shard
+                # trick is a runtime optimization), exactly like the target
+                # head. With rejection off, non-greedy batches keep the
+                # LM-head-TP argmax path unconditionally, where the per-rank
+                # greedy flag never enters control flow. This branch is safe
+                # to take group-uniformly because is_all_greedy_sample is
+                # group-synchronized whenever rejection+ADP+LM-head-TP are
+                # combined -- see SpecMetadata.group_all_greedy_sample (anchor
+                # for the group-sync semantics).
+                advanced_draft_sampling = spec_metadata.wants_advanced_draft_sampling
+                # enable_lm_head_tp_in_adp implies enable_attention_dp
+                # (asserted in Mapping.__init__); no separate ADP check.
+                lm_head_tp_in_adp_configured = (
+                    self.is_mtp_eagle
+                    and self.model_config is not None
+                    and self.model_config.mapping.enable_lm_head_tp_in_adp
+                )
+                use_lm_head_tp_in_adp = lm_head_tp_in_adp_configured and not advanced_draft_sampling
+                draft_tokens = None
+                if self.is_mtp_eagle:
+                    # One row per request: the LM-head input on this step and
+                    # the draft input on the next.
+                    request_hidden_states = (
+                        draft_rows if draft_rows is not None else hidden_states[gather_ids]
+                    )
+                    if draft_tail is not None:
+                        draft_tokens = draft_tail(request_hidden_states)
+                    if draft_tokens is not None:
+                        logits = None
+                    elif use_lm_head_tp_in_adp:
+                        token_count = request_hidden_states.view(
+                            -1, request_hidden_states.shape[-1]
+                        ).shape[0]
+                        max_num_requests = spec_metadata.max_num_requests
+                        pad_len = max_num_requests - token_count
+                        if pad_len > 0:
+                            padded_hidden_states = F.pad(
+                                request_hidden_states.view(-1, request_hidden_states.shape[-1]),
+                                (0, 0, 0, pad_len),
+                                mode="constant",
+                                value=0,
+                            )
+                        elif pad_len == 0:
+                            padded_hidden_states = request_hidden_states.view(
+                                -1, request_hidden_states.shape[-1]
+                            )
+                        else:
+                            raise ValueError(
+                                "Eagle3OneModelWorker (MTP Eagle mode): "
+                                "token_count > max_num_requests, which is not supported"
+                            )
+                        logits = draft_model.mtp_layers[0].shared_head(
+                            padded_hidden_states, draft_model.lm_head, attn_metadata, True
+                        )
+                    elif lm_head_tp_in_adp_configured:
+                        # Advanced-sampling bypass: the model's shared_head
+                        # would re-apply the LM-head-TP stacked/sharded path
+                        # from config on its own. Under ADP the LMHead weight
+                        # is replicated, so project this rank's own rows to the
+                        # full vocabulary. Model-specific shared heads may need
+                        # preprocessing before that local projection.
+                        shared_head = draft_model.mtp_layers[0].shared_head
+                        local_full_vocab_forward = getattr(
+                            shared_head, "forward_local_full_vocab", None
+                        )
+                        if local_full_vocab_forward is None:
+                            logits = draft_model.lm_head(request_hidden_states)
+                        else:
+                            logits = local_full_vocab_forward(
+                                request_hidden_states,
+                                draft_model.lm_head,
+                                attn_metadata,
+                                True,
+                            )
+                    else:
+                        logits = draft_model.mtp_layers[0].shared_head(
+                            request_hidden_states, draft_model.lm_head, attn_metadata, True
+                        )
+                else:
+                    logits = draft_model.logits_processor(
+                        hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True
+                    )
+
+                if self.guided_decoder is not None:
+                    if self.is_mtp_eagle:
+                        self.guided_decoder.execute_draft_batch(logits, draft_step=i)
+                    else:
+                        self.guided_decoder.execute_draft_batch(logits, self._d2t, draft_step=i)
+
+                # ADP+LM-head-TP logits are the LM-head-TP group's row-stacked
+                # batch (each rank's rows padded to max_num_requests, then
+                # all-gathered along dim 0) with the vocab sharded across the
+                # group. Rows [:token_count] would be group rank 0's requests,
+                # not this rank's, and a per-rank argmax would return a
+                # shard-local index -- so keep the full stacked logits and let
+                # greedy_sample_draft_with_tp_gather combine the group's vocab
+                # shards and slice this rank's own row segment; only then trim
+                # the max_num_requests padding down to token_count.
+                mapping_lm_head_tp = None
+                if use_lm_head_tp_in_adp:
+                    # The MTP head built this per-forward mapping when producing
+                    # the vocab-sharded logits; the sampler needs it to gather.
+                    mapping_lm_head_tp = getattr(
+                        draft_model.mtp_layers[0].shared_head, "mapping_lm_head_tp", None
+                    )
+                if draft_tokens is not None:
+                    new_draft_token = draft_tokens
+                else:
+                    new_draft_token = self.sample_draft_tokens(
+                        logits,
+                        spec_metadata,
+                        batch_size,
+                        draft_step=i,
+                        mapping_lm_head_tp=mapping_lm_head_tp,
+                    )
+                if use_lm_head_tp_in_adp:
+                    new_draft_token = new_draft_token[:token_count]
+                next_draft_tokens.append(new_draft_token)
+
+                # Build the next iteration's inputs; the last iteration has
+                # none to build.
+                # MTP Eagle: the MTP layer returns a single tensor; the rows
+                #   selected above for the LM head carry over as they are.
+                # Eagle3: the EAGLE draft model returns a secondary
+                #   ``hidden_states_to_save`` specifically for this purpose.
+                # The draft-loop step update: one launch for the position advance
+                # and the metadata update below; None keeps the native update.
+                draft_step = None
+                if draft_loop_ops is not None:
+                    draft_step = draft_loop_ops.step_update(
+                        attn_metadata,
+                        i,
+                        num_accepted_tokens,
+                        inputs["position_ids"] if i + 1 < runtime_draft_len else None,
+                        gather_ids if i == 0 else None,
+                    )
+                if i + 1 < runtime_draft_len:
+                    if self.is_mtp_eagle:
+                        hidden_states = request_hidden_states
+                    else:
+                        hidden_states = hidden_states_to_save[gather_ids]
+                    if draft_step is not None:
+                        position_ids = draft_step.position_ids
+                    else:
+                        position_ids = (
+                            _select_mtp_position_ids(inputs["position_ids"], gather_ids) + 1
+                        )
+                    inputs = {
+                        "input_ids": new_draft_token,
+                        "position_ids": position_ids,
+                        "hidden_states": hidden_states,
+                        "attn_metadata": attn_metadata,
+                        "spec_metadata": spec_metadata,
+                    }
+
+                # Update attn_metadata for the next iteration.
+                if draft_step is not None:
+                    pass  # The draft-loop step update updated the metadata above.
+                elif i == 0:
+                    attn_metadata._seq_lens[:batch_size].fill_(1)
+                    attn_metadata._seq_lens_cuda[:batch_size].fill_(1)
+                    attn_metadata.on_update()
+                    has_kv_cache = inputs["attn_metadata"].kv_cache_manager is not None
+                    if has_kv_cache and hasattr(attn_metadata, "host_request_types"):
+                        attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1)
+                    if has_kv_cache:
+                        attn_metadata.num_contexts = 0
+                    if hasattr(attn_metadata, "kv_lens_cuda"):
+                        attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (
+                            runtime_draft_len - num_accepted_tokens[num_contexts:]
+                        )
+                        attn_metadata.kv_lens_cuda[:num_contexts] += 1
+
+                    if has_kv_cache:
+                        self._prepare_flash_mla_generation_layout(
+                            attn_metadata, num_contexts, batch_size
+                        )
+                    if hasattr(attn_metadata, "kv_lens_cuda"):
+                        attn_metadata.update_for_spec_dec()
+
+                    # Both Eagle3 and MTP Eagle drafters take ``draft_len + 1``
+                    # tokens in the first draft step (attention runs in spec-dec
+                    # mode), then 1 token per step in subsequent iterations.
+                    # Disable spec_decoding here so the masks/positions stay
+                    # correct on subsequent iters.
+                    attn_metadata.use_spec_decoding = False
+                else:
+                    if hasattr(attn_metadata, "kv_lens_cuda"):
+                        attn_metadata.kv_lens_cuda[:batch_size] += 1
+                        attn_metadata.update_for_spec_dec()
+        stacked = draft_loop_ops.finalize(next_draft_tokens) if draft_loop_ops is not None else None
+        next_draft_tokens = (
+            stacked if stacked is not None else torch.stack(next_draft_tokens, dim=1)
+        )
+
+        # Override with SA draft tokens after all draft layers have run,
+        # so that draft layers never see SA tokens in their inputs.
+        if self.sa_enhancer is not None:
+            gen_draft_tokens = next_draft_tokens[num_contexts:]
+            gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens(gen_draft_tokens)
+            next_draft_tokens[num_contexts:] = gen_draft_tokens
+
+        return next_draft_tokens
+
+    def _run_draft_forward(self, draft_model, inputs, spec_metadata, step_idx: int):
+        """Invoke the draft model for one iteration, branching on mode.
+
+        ``all_rank_num_tokens`` is passed as a kwarg in both modes. For MTP
+        Eagle it goes directly to ``mtp_layers[0]``; for Eagle3 it goes to
+        ``Eagle3DraftModel.forward`` which guards it with a try/finally so
+        attn_metadata sees the original value on return.
+        """
+        if dispatch.phase() is not dispatch.Phase.DECODE:
+            return super()._run_draft_forward(
+                draft_model=draft_model,
+                inputs=inputs,
+                spec_metadata=spec_metadata,
+                step_idx=step_idx,
+            )
+        all_rank_num_tokens = self._get_step_all_rank_num_tokens(spec_metadata, step_idx)
+
+        if self.is_mtp_eagle:
+            embed_tokens = draft_model.embed_tokens
+            if dispatch.is_enabled("draft_loop"):
+                embed_tokens = dispatch.get_impl("draft_loop", "draft_embedding")(embed_tokens)
+            hidden_states = draft_model.mtp_layers[0](
+                embed_tokens=embed_tokens, all_rank_num_tokens=all_rank_num_tokens, **inputs
+            )
+            return hidden_states, None
+
+        inputs["all_rank_num_tokens"] = all_rank_num_tokens
+        hidden_states, hidden_states_to_save = draft_model.model(**inputs)
+        return hidden_states, hidden_states_to_save

From 70a0cbbf0bc448cd59e723c80463907d2ac6d8ee Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 09:10:39 -0700
Subject: [PATCH 04/19] [None][test] modeling_v2: tests of the GLM-5.2 targets

Unit tests of the glm_moe_dsa family: routing and target claims,
isolation from other models, the per-target plans, the adopted layers and
the runtime hooks, plus kernel tests against the native implementations.
Tests that need the GLM-5.2 checkpoint find it under llm_models_root()
and skip when it is not available.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../test_modeling_v2_glm_dep_plan.py          |  411 ++++++
 .../test_modeling_v2_glm_dep_target.py        |  501 +++++++
 .../test_modeling_v2_glm_tep_layers.py        |  328 +++++
 .../test_modeling_v2_glm_tep_plan.py          |  355 +++++
 .../test_expert_prefetch.py                   |  267 ++++
 ...expert_prefetch_composition_hop_routing.py |  353 +++++
 ...osition_shared_prefetch_expert_prefetch.py |  352 +++++
 ...eferred_norm_routed_into_hop_wiring_mpi.py |  343 +++++
 .../draft_live_rows/test_live_rows_oracle.py  |  270 ++++
 .../draft_routed_moe/draft_moe_fixtures.py    |  315 +++++
 .../kernels/draft_routed_moe/draft_moe_jit.py |  101 ++
 .../test_draft_moe_adapter.py                 |  152 +++
 .../draft_routed_moe/test_draft_moe_oracle.py |  775 +++++++++++
 .../kernels/hop_routing/hop_routing_jit.py    |  161 +++
 .../kernels/hop_routing/hop_routing_oracle.py |  214 +++
 .../hop_routing/test_hop_routing_adapter.py   |  245 ++++
 .../test_hop_routing_hop_single_gpu.py        |  234 ++++
 .../test_hop_routing_routed_moe_fork.py       |  200 +++
 .../test_attn_hop_integration_contract.py     |  161 +++
 .../routed_moe/routed_into_hop_checks.py      |  492 +++++++
 .../routed_moe/routed_moe_chain_bench.py      |  764 +++++++++++
 .../kernels/routed_moe/routed_moe_fixtures.py | 1036 ++++++++++++++
 .../routed_moe/test_routed_into_hop_mpi.py    |  464 +++++++
 .../routed_moe/test_routed_moe_integration.py |  400 ++++++
 .../test_routed_moe_native_consumer_mpi.py    |  176 +++
 .../routed_moe/test_routed_moe_routed_moe.py  |  722 ++++++++++
 .../test_routed_moe_spec_prefetch.py          |  184 +++
 .../routed_moe/test_routed_moe_stress.py      |  812 +++++++++++
 ...test_router_gemm_integration_invariants.py |  135 ++
 .../shared_prefetch_fixtures.py               |  104 ++
 .../shared_prefetch_se.py                     |  251 ++++
 .../shared_prefetch_se_kernel.py              |  453 +++++++
 .../shared_prefetch_se_ptx.py                 |  120 ++
 .../test_shared_prefetch_adapter.py           |  251 ++++
 .../test_shared_prefetch_routed.py            |  454 +++++++
 .../test_shared_prefetch_shared.py            |  269 ++++
 .../test_modeling_v2_glm_adopt.py             |  152 +++
 .../test_modeling_v2_glm_layers.py            |  888 ++++++++++++
 .../test_modeling_v2_glm_phase.py             |  213 +++
 .../test_modeling_v2_glm_selection.py         |  573 ++++++++
 .../test_modeling_v2_glm_spec_worker.py       |  514 +++++++
 .../test_modeling_v2_glm_targets.py           |  597 +++++++++
 .../test_modeling_v2_glm_vocabulary.py        |  198 +++
 .../absorbed_query/absorbed_query_fixture.py  |  223 ++++
 .../shared/kernels/absorbed_query/conftest.py |   22 +
 ...st_absorbed_query_attention_integration.py |  870 ++++++++++++
 .../test_absorbed_query_proj.py               |  869 ++++++++++++
 ...test_draft_query_draft_loop_integration.py | 1103 +++++++++++++++
 .../test_draft_query_fused_dsa_metadata.py    |  745 +++++++++++
 .../test_draft_query_mtp_query_proj.py        |  439 ++++++
 ...test_shared_query_attention_integration.py |  447 +++++++
 .../test_shared_query_shared_absorb.py        |  599 +++++++++
 .../attention_output_tail/real_mla_harness.py |  485 +++++++
 .../attention_output_tail/tail_fixtures.py    |  191 +++
 .../attention_output_tail/tail_reference.py   |  101 ++
 .../test_attn_tail_all_layers.py              |  130 ++
 .../test_attn_tail_attn_tail.py               |  263 ++++
 .../test_attn_tail_bridge_mpi.py              |  244 ++++
 .../test_attn_tail_contract_o.py              |  290 ++++
 .../test_attn_tail_contract_o_forward.py      |  276 ++++
 .../test_attn_tail_gates.py                   |  132 ++
 .../test_attn_tail_mla_adapter.py             |  359 +++++
 .../test_attn_tail_mla_forward.py             |  269 ++++
 .../test_attn_tail_mla_forward_mpi.py         |  423 ++++++
 .../test_attn_tail_pdl_stress.py              |  506 +++++++
 .../test_attn_tail_sanitizer_probe.py         |  115 ++
 .../test_into_hop_producer.py                 |  358 +++++
 .../test_tail_prefetch.py                     |  531 ++++++++
 .../test_tail_prefetch_mla_forward.py         |  142 ++
 .../test_cluster_merge_adapter.py             |  171 +++
 .../test_global_emit_adapter.py               |  535 ++++++++
 ...t_composition_global_emit_indexer_front.py |  387 ++++++
 .../_decode_prologue_oracle.py                |  355 +++++
 .../decode_prologue/test_decode_prologue.py   |  421 ++++++
 .../deferred_input_prep/input_prep_testlib.py |  230 ++++
 .../test_input_prep_compat.py                 |  340 +++++
 .../test_input_prep_kernels.py                |  872 ++++++++++++
 .../shared/kernels/deferred_norm/conftest.py  |   24 +
 .../deferred_norm/deferred_norm_utils.py      |  299 +++++
 ...est_deferred_norm_attention_integration.py |  495 +++++++
 .../test_deferred_norm_controls.py            |  425 ++++++
 .../test_deferred_norm_hop_mpi.py             |  835 ++++++++++++
 .../test_deferred_norm_null_controls.py       |  559 ++++++++
 .../test_deferred_norm_projection.py          |  293 ++++
 .../test_deferred_norm_routed_mpi.py          |  946 +++++++++++++
 .../deferred_norm/test_deferred_norm_scope.py |  327 +++++
 .../test_deferred_norm_scope_mpi.py           |  342 +++++
 .../test_evict_hints_adapter.py               |  126 ++
 .../test_evict_hints_config.py                |   85 ++
 .../test_evict_hints_hop_mpi.py               |  424 ++++++
 .../test_evict_hints_projection.py            |  156 +++
 .../kernels/dense_mlp/dense_mlp_utils.py      |  190 +++
 .../shared/kernels/dense_mlp/pdl_delay.cu     |   73 +
 .../shared/kernels/dense_mlp/test_adapter.py  |  171 +++
 .../kernels/dense_mlp/test_dense_mlp.py       |  410 ++++++
 .../test_draft_front_draft_front.py           |  853 ++++++++++++
 .../draft_live_rows/test_live_rows_adapter.py |  472 +++++++
 .../draft_loop/_draft_loop_fixtures.py        |  311 +++++
 .../draft_loop/test_draft_loop_draft_loop.py  |  614 +++++++++
 .../draft_loop/test_draft_loop_draft_step.py  |  314 +++++
 .../draft_loop/test_draft_loop_embedding.py   |  177 +++
 .../kernels/draft_tail/draft_tail_nearties.py |  296 ++++
 .../kernels/draft_tail/draft_tail_oracle.py   |  148 ++
 .../draft_tail/test_draft_tail_draft_tail.py  |  736 ++++++++++
 .../draft_tail/test_draft_tail_eligibility.py |  239 ++++
 .../draft_tail/test_draft_tail_nearties.py    |   89 ++
 .../fused_kv_indexer_projection/conftest.py   |   24 +
 .../fused_kv_utils.py                         |  288 ++++
 .../test_fused_kv_attention_integration.py    |  275 ++++
 .../test_fused_kv_entry.py                    |  294 ++++
 .../test_fused_kv_projection.py               |  331 +++++
 .../test_host_pages_host_policy.py            |  176 +++
 .../test_host_pages_selector.py               |   82 ++
 .../test_indexer_front_adapter.py             |  246 ++++
 .../indexer_topk/test_cross_rank_cp_policy.py |  195 +++
 .../test_indexer_topk_indexer_decode.py       | 1081 +++++++++++++++
 .../test_indexer_topk_pipeline.py             | 1188 +++++++++++++++++
 .../indexer_topk/test_indexer_topk_policy.py  |  243 ++++
 .../oproj_weight_staging/test_oproj_stage.py  |  839 ++++++++++++
 .../shared/kernels/plan_testlib.py            |   92 ++
 .../post_attention_hop/into_hop_jit.py        |   81 ++
 .../test_attn_hop_static.py                   |  230 ++++
 .../test_into_hop_single_gpu.py               |  276 ++++
 .../post_moe_hop/post_moe_hop_mpi_utils.py    |  261 ++++
 .../post_moe_hop/post_moe_hop_real_fixture.py |  211 +++
 .../test_post_moe_hop_post_moe_mpi.py         |  796 +++++++++++
 .../test_post_moe_hop_protocol_model.py       |  287 ++++
 .../test_post_moe_hop_real_weights_mpi.py     |  218 +++
 .../shared/kernels/rope_kv_append/conftest.py |   22 +
 .../kernels/rope_kv_append/pdl_producer.py    |  194 +++
 .../kernels/rope_kv_append/rope_append_env.py |  624 +++++++++
 .../rope_kv_append/rope_append_stack.py       |  727 ++++++++++
 .../test_rope_append_composition.py           |  226 ++++
 .../test_rope_append_integration.py           |  501 +++++++
 .../rope_kv_append/test_rope_append_ops.py    |  597 +++++++++
 .../router_gemm/test_pre_moe_hop_mpi.py       |  241 ++++
 .../kernels/router_gemm/test_router_gemm.py   |  446 +++++++
 .../kernels/sparse_mla_decode/conftest.py     |  100 ++
 .../sparse_mla_decode/sparse_mla_case.py      |  397 ++++++
 .../sparse_mla_decode/sparse_mla_harness.py   | 1026 ++++++++++++++
 .../sparse_mla_pdl_helpers.py                 |  204 +++
 .../sparse_mla_real_fixture.py                |  310 +++++
 .../sparse_mla_stack_case.py                  |  480 +++++++
 .../test_sparse_mla_adapter.py                |  100 ++
 .../test_sparse_mla_composition.py            |  213 +++
 .../test_sparse_mla_dispatch_accounting.py    |  123 ++
 .../test_sparse_mla_native_semantics.py       |  171 +++
 .../test_sparse_mla_pdl_stress.py             |  479 +++++++
 .../test_sparse_mla_ragged_verify.py          |   99 ++
 .../test_sparse_mla_real_weights.py           |  404 ++++++
 .../test_sparse_mla_rope_append_bridge.py     |  195 +++
 .../test_sparse_mla_sanitizer_cases.py        |  150 +++
 .../test_sparse_mla_sparse_mla.py             |  285 ++++
 .../test_sparse_mla_stack_composition.py      |  350 +++++
 .../shared/kernels/test_kernels_contracts.py  |  157 +++
 .../shared/kernels/test_kernels_op_names.py   |  131 ++
 .../shared/test_modeling_v2_glm_attention.py  |  754 +++++++++++
 .../test_modeling_v2_glm_decline_kinds.py     |  534 ++++++++
 .../shared/test_modeling_v2_glm_dispatch.py   |  380 ++++++
 .../test_modeling_v2_glm_moe_dsa_isolation.py |  297 +++++
 .../test_modeling_v2_glm_moe_dsa_routing.py   |  700 ++++++++++
 .../modeling_v2/test_modeling_v2_claims.py    |    7 +
 162 files changed, 59170 insertions(+)
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py
 create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py

diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py
new file mode 100644
index 000000000000..564c5b927364
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py
@@ -0,0 +1,411 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The fast paths the GLM-5.2 dep target serves in each of its cells, how it picks the cell, and
+what its plan does with the declines its deployments log.
+
+Host-only: the plan is resolved and installed from the cell's ``TargetPaths`` exactly as the
+shell does at construction. The engagement of each deployment is the per-rank decision set the
+fast paths log there (the same on every rank): the paths that engage, and the (path, key,
+reason) triples of the paths the cell serves that decline. None of the declines may be a miss:
+under ``modeling_v2='require'`` a miss raises.
+"""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+pytestmark = pytest.mark.cpu_only
+
+D = dispatch.Decline
+REQUIRE, AUTO = ModelingV2Mode.REQUIRE, ModelingV2Mode.AUTO
+
+# The deployments, by (ranks, LM head tensor-parallel in attention DP, draft tokens per request),
+# with their KV-cache dtype and MoE backend.
+DEP16 = (16, False, 3)
+DEP8 = (8, False, 3)
+DEP8_LM_HEAD_TP = (8, True, 5)
+DEPLOYMENTS = {
+    DEP16: (QuantAlgo.FP8, "TRTLLM"),
+    DEP8: (QuantAlgo.FP8, "CUTEDSL"),
+    DEP8_LM_HEAD_TP: (QuantAlgo.NVFP4, "CUTEDSL"),
+}
+IDS = ["dep16", "dep8", "dep8-lm-head-tp"]
+
+# The paths that engage at each deployment (device work on every rank).
+ENGAGED = {
+    DEP16: {
+        "rope_kv_append",
+        "indexer_layer_front",
+        "draft_loop",
+        "decode_prologue",
+        "host_page_retention",
+    },
+    DEP8: {
+        "rope_kv_append",
+        "indexer_layer_front",
+        "draft_loop",
+        "decode_prologue",
+        "host_page_retention",
+    },
+    DEP8_LM_HEAD_TP: {"draft_loop", "decode_prologue", "host_page_retention"},
+}
+# The bases the engaged paths extend; they decline every call at these deployments, or (the
+# post-attention hop) are never built.
+BASES = {
+    DEP16: {
+        "absorbed_query",
+        "draft_absorbed_query",
+        "indexer_topk",
+        "cross_rank_topk",
+        "post_attention_hop",
+    },
+    DEP8: {
+        "absorbed_query",
+        "draft_absorbed_query",
+        "indexer_topk",
+        "cross_rank_topk",
+        "post_attention_hop",
+    },
+    DEP8_LM_HEAD_TP: {"post_attention_hop"},
+}
+SERVED = {cell: ENGAGED[cell] | BASES[cell] for cell in DEPLOYMENTS}
+# Paths that log an engaged decision at these deployments without device work there: the
+# step-input stager's feature check (its stager, store writer and execution-stream loop all
+# decline attention DP) and the cross-layer expert prefetch links (the prefetch is issued only
+# by the routed-MoE launch, which does not run with expert-parallel experts).
+LOG_ONLY = {"deferred_input_prep", "sampler_store", "execution_stream_loop"} | {
+    "cross_layer_expert_prefetch"
+}
+
+ADP = "attention DP / CP not supported"
+ADP_TOPK = "attention data parallelism"
+
+# The declines the served paths log with an FP8 KV cache, on every rank.
+DECLINES_FP8_KV = (
+    *(("absorbed_query", f"layer{i}", ADP) for i in range(78)),
+    ("absorbed_query", "layer78", "target layers only"),
+    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
+    ("draft_absorbed_query", "layer78", ADP),
+    *(("indexer_layer_front", f"target_layer{i}", ADP) for i in range(78)),
+    ("indexer_layer_front", "draft_layer78", ADP),
+    ("indexer_layer_front", "cp_indexer_decode_cross_rank_topk", ADP_TOPK),
+    ("indexer_topk", "indexer_decode", ADP_TOPK),
+    ("indexer_topk", "indexer_decode", "MTP draft step (target layers only)"),
+    ("cross_rank_topk", "cp_indexer_decode", ADP_TOPK),
+    ("draft_loop", "draft_loop:context rows present", "context rows present"),
+    ("draft_loop", "draft_step:dense top-k fallback active", "dense top-k fallback active"),
+    (
+        "decode_prologue",
+        "target_prologue[not a generation-only step]",
+        "not a generation-only step",
+    ),
+)
+# The declines the served paths log with an NVFP4 KV cache and the LM head tensor-parallel in
+# attention DP, on every rank.
+DECLINES_NVFP4_KV = (
+    ("draft_loop", "draft_loop:LM-head TP in attention DP", "LM-head TP in attention DP"),
+    ("draft_loop", "draft_loop:context rows present", "context rows present"),
+    (
+        "decode_prologue",
+        "target_prologue[not a generation-only step]",
+        "not a generation-only step",
+    ),
+)
+DECLINES = {DEP16: DECLINES_FP8_KV, DEP8: DECLINES_FP8_KV, DEP8_LM_HEAD_TP: DECLINES_NVFP4_KV}
+# The engaged decisions the served paths log, as (path, key).
+ENGAGED_KEYS = {
+    DEP16: {
+        ("rope_kv_append", "rope"),
+        ("rope_kv_append", "bridge"),
+        ("indexer_layer_front", "rope"),
+        ("draft_loop", "draft_embedding"),
+        ("draft_loop", "draft_loop:engaged"),
+        ("draft_loop", "draft_step:engaged"),
+        ("decode_prologue", "target_prologue"),
+        ("decode_prologue", "target_embedding"),
+        ("host_page_retention", "host_policy"),
+    },
+    DEP8_LM_HEAD_TP: {
+        ("draft_loop", "draft_embedding"),
+        ("decode_prologue", "target_prologue"),
+        ("decode_prologue", "target_embedding"),
+        ("host_page_retention", "host_policy"),
+    },
+}
+ENGAGED_KEYS[DEP8] = ENGAGED_KEYS[DEP16]
+
+
+def _model_config(cell, spec=True, **overrides):
+    """What ``target_paths`` reads off the model config of a deployment; ``overrides`` replace
+    its ``kv_cache`` and ``moe_backend``."""
+    ranks, lm_head_tp, draft = cell
+    kv_cache, moe_backend = DEPLOYMENTS.get(cell, (QuantAlgo.FP8, "TRTLLM"))
+    facts = {"kv_cache": kv_cache, "moe_backend": moe_backend, **overrides}
+    return types.SimpleNamespace(
+        mapping=types.SimpleNamespace(
+            world_size=ranks, enable_attention_dp=True, enable_lm_head_tp_in_adp=lm_head_tp
+        ),
+        spec_config=types.SimpleNamespace(max_draft_len=draft) if spec else None,
+        quant_config=types.SimpleNamespace(kv_cache_quant_algo=facts["kv_cache"]),
+        moe_backend=facts["moe_backend"],
+    )
+
+
+@pytest.fixture
+def installed(monkeypatch):
+    """Resolve and install the plan of a deployment's cell, as the shell does."""
+    for name in ("_INSTALLED", "_VALIDATE", "_require", "_TARGET"):
+        monkeypatch.setattr(dispatch, name, list(getattr(dispatch, name)))
+    monkeypatch.setattr(dispatch, "_logged", set())
+    monkeypatch.setattr(dispatch, "_missed", {})
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+    dispatch.classify_decline.cache_clear()
+
+    def install(cell, environ=None, mode=REQUIRE):
+        paths = modeling.target_paths(_model_config(cell), mode)
+        plan = dispatch.resolve(mode, environ or {}, paths)
+        dispatch.install(plan, environ or {})
+        return plan
+
+    yield install
+    dispatch.classify_decline.cache_clear()
+
+
+# --- the cells ------------------------------------------------------------------------------
+
+
+def test_the_cells_are_the_three_deployments():
+    assert set(modeling.CELLS) == set(DEPLOYMENTS)
+    for cell, (kv_cache, moe_backend) in DEPLOYMENTS.items():
+        assert (modeling.CELLS[cell].kv_cache, modeling.CELLS[cell].moe_backend) == (
+            kv_cache,
+            moe_backend,
+        )
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+@pytest.mark.parametrize("mode", [REQUIRE, AUTO], ids=["require", "auto"])
+def test_each_deployment_selects_its_cells_paths(cell, mode):
+    assert modeling.target_paths(_model_config(cell), mode) is modeling.CELLS[cell].paths
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_the_kv_cache_dtype_may_be_the_enum_or_its_value(cell):
+    kv_cache = DEPLOYMENTS[cell][0]
+    config = _model_config(cell, kv_cache=kv_cache.value)
+    assert not isinstance(config.quant_config.kv_cache_quant_algo, QuantAlgo)
+    assert modeling.target_paths(config, REQUIRE) is modeling.CELLS[cell].paths
+
+
+# A deployment ``route`` claims whose KV-cache dtype or MoE backend is not its cell's, and one
+# outside the cells: (cell, the facts that differ, what the error names).
+_OUTSIDE = {
+    "dep16-nvfp4-kv": (DEP16, {"kv_cache": QuantAlgo.NVFP4}, "KV cache NVFP4, MoE backend TRTLLM"),
+    "dep16-cutedsl": (DEP16, {"moe_backend": "CUTEDSL"}, "KV cache FP8, MoE backend CUTEDSL"),
+    "dep8-trtllm": (DEP8, {"moe_backend": "TRTLLM"}, "KV cache FP8, MoE backend TRTLLM"),
+    "dep8-nvfp4-kv": (DEP8, {"kv_cache": QuantAlgo.NVFP4}, "KV cache NVFP4, MoE backend CUTEDSL"),
+    "dep8-unquantized-kv": (DEP8, {"kv_cache": None}, "KV cache None, MoE backend CUTEDSL"),
+    "dep8-lm-head-tp-fp8-kv": (
+        DEP8_LM_HEAD_TP,
+        {"kv_cache": QuantAlgo.FP8},
+        "KV cache FP8, MoE backend CUTEDSL",
+    ),
+    "dep8-lm-head-tp-cutlass": (
+        DEP8_LM_HEAD_TP,
+        {"moe_backend": "CUTLASS"},
+        "KV cache NVFP4, MoE backend CUTLASS",
+    ),
+    "dep16-mtp5": (
+        (16, False, 5),
+        {},
+        "16 ranks, LM head tensor-parallel in attention DP False, 5",
+    ),
+}
+
+
+@pytest.mark.parametrize("cell, overrides, named", list(_OUTSIDE.values()), ids=list(_OUTSIDE))
+def test_a_deployment_outside_the_cells_raises_under_require(cell, overrides, named):
+    config = _model_config(cell, **overrides)
+    with pytest.raises(ValueError, match="is not a deployment this target serves") as excinfo:
+        modeling.target_paths(config, REQUIRE)
+    message = str(excinfo.value)
+    assert named in message, message
+    # The error names the three cells.
+    assert message.count("draft tokens: KV cache") == 3, message
+
+
+@pytest.mark.parametrize("cell, overrides, named", list(_OUTSIDE.values()), ids=list(_OUTSIDE))
+def test_a_deployment_outside_the_cells_serves_nothing_under_auto(
+    monkeypatch, cell, overrides, named
+):
+    warnings = []
+    monkeypatch.setattr(modeling.logger, "warning", warnings.append)
+    config = _model_config(cell, **overrides)
+    assert modeling.target_paths(config, AUTO) is modeling.NO_PATHS
+    assert modeling.NO_PATHS.served == frozenset()
+    (warning,) = warnings
+    assert named in warning and warning.endswith("no fast path runs"), warning
+
+
+def test_without_speculation_no_cell_matches():
+    config = _model_config(DEP8, spec=False)
+    with pytest.raises(ValueError, match="None draft tokens"):
+        modeling.target_paths(config, REQUIRE)
+
+
+def test_the_auto_plan_outside_the_cells_runs_no_fast_path(installed, monkeypatch):
+    monkeypatch.setattr(modeling.logger, "warning", lambda message: None)
+    paths = modeling.target_paths(_model_config(DEP8, moe_backend="TRTLLM"), AUTO)
+    plan = dispatch.resolve(AUTO, {}, paths)
+    dispatch.install(plan, {})
+    assert dispatch.installed() == frozenset()
+    assert plan.describe().startswith("modeling_v2=auto, 0/0 fast paths")
+    assert not dispatch.require_mode()
+
+
+# --- the plans ------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_the_cell_serves_the_engaged_paths_and_their_bases(installed, cell):
+    plan = installed(cell)
+    assert set(plan.served) == SERVED[cell]
+    assert dispatch.installed() == frozenset(SERVED[cell])
+    count = len(SERVED[cell])
+    assert plan.describe().startswith(f"modeling_v2=require, {count}/{count} fast paths")
+
+
+def test_the_cells_serve_ten_ten_and_four_paths():
+    assert [len(SERVED[cell]) for cell in DEPLOYMENTS] == [10, 10, 4]
+    assert modeling.CELLS[DEP16].paths is modeling.CELLS[DEP8].paths is modeling.FP8_KV_PATHS
+    assert modeling.CELLS[DEP8_LM_HEAD_TP].paths is modeling.NVFP4_KV_PATHS
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_every_base_of_a_served_path_is_served(cell):
+    served = SERVED[cell]
+    for path in served:
+        assert dispatch.REQUIRED_FEATURES.get(path, frozenset()) <= served, path
+    # The bases are exactly what the engaged paths extend.
+    required = set().union(*(dispatch.REQUIRED_FEATURES.get(p, frozenset()) for p in ENGAGED[cell]))
+    assert required - ENGAGED[cell] == BASES[cell]
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_no_log_only_path_is_served(cell):
+    """The step-input stager's parts and the expert prefetch do no device work under attention
+    DP; the execution-stream loop has no attention-DP gate of its own, so it must not be
+    served."""
+    assert not (SERVED[cell] & LOG_ONLY)
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_every_kernel_package_is_the_familys(cell):
+    assert dict(modeling.CELLS[cell].paths.kernel_modules) == {}
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_the_record_line_lists_the_served_paths(installed, monkeypatch, cell):
+    lines = []
+    monkeypatch.setattr(dispatch, "_log_info", lines.append)
+    installed(cell)
+    (record,) = [line for line in lines if "record" in line]
+    listed = record.split("paths=", 1)[1].split(" ", 1)[0].split(",")
+    assert {name.split("[", 1)[0] for name in listed} == SERVED[cell]
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_disable_may_name_only_a_served_path(installed, cell):
+    with pytest.raises(ValueError, match="not fast paths of this target"):
+        installed(cell, {dispatch.DISABLE_ENV: "execution_stream_loop"})
+    plan = installed(cell, {dispatch.DISABLE_ENV: "decode_prologue"})
+    assert not plan.enabled("decode_prologue") and plan.enabled("draft_loop")
+
+
+# Each path of the failure drop order with what it turns off (itself and its dependents).
+_DROP_ORDER = {
+    "decode_prologue": {"decode_prologue"},
+    "indexer_layer_front": {"indexer_layer_front"},
+    "draft_loop": {"draft_loop", "decode_prologue"},
+    "rope_kv_append": {"rope_kv_append", "indexer_layer_front"},
+}
+
+
+@pytest.mark.parametrize("path", list(_DROP_ORDER))
+def test_each_path_of_the_drop_order_can_be_disabled(installed, path):
+    plan = installed(DEP8, {dispatch.DISABLE_ENV: path})
+    off = {name for name in SERVED[DEP8] if not plan.enabled(name)}
+    assert off == _DROP_ORDER[path]
+    assert plan.describe().startswith(f"modeling_v2=require, {10 - len(off)}/10 fast paths")
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_the_routed_accumulation_knob_is_refused(installed, cell):
+    with pytest.raises(ValueError, match="does not serve routed_moe_into_hop"):
+        installed(cell, {dispatch.ACCUMULATION_ENV: "identity"})
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_the_cross_rank_wait_timeout_is_read(installed, cell):
+    plan = installed(cell, {dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV: "60000"})
+    assert plan.knobs == ((dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV, "60000"),)
+
+
+# --- the declines ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_no_logged_decline_is_a_miss(installed, cell):
+    """Under require, a decline decided once inside a path's cells would raise."""
+    installed(cell)
+    assert dispatch.require_mode()
+    for path, key, reason in DECLINES[cell]:
+        assert path in SERVED[cell], path
+        assert dispatch.classify_decline(path, key, reason) is D.UNCOVERED, (path, key, reason)
+        assert dispatch.log_capability(path, key, False, reason) is False
+    assert dispatch.missed_counts() == {}
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_the_engaged_decisions_are_of_served_paths(installed, cell):
+    installed(cell)
+    for path, key in ENGAGED_KEYS[cell]:
+        assert path in ENGAGED[cell], path
+        assert dispatch.log_capability(path, key, True) is True
+    assert {path for path, _ in ENGAGED_KEYS[cell]} == ENGAGED[cell]
+
+
+@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
+def test_every_base_declines_every_call(cell):
+    """A base the cell serves logs no engaged decision; every one but the never-built
+    post-attention hop declines."""
+    declining = {path for path, _, _ in DECLINES[cell]}
+    assert BASES[cell] - {"post_attention_hop"} <= declining
+    assert not ({path for path, _ in ENGAGED_KEYS[cell]} & BASES[cell])
+
+
+def test_the_family_cells_would_raise_on_the_query_fronts(installed):
+    """The family's cells hold every target layer for the query fronts; the dep table is what
+    keeps an attention-DP deployment from raising on their declines under require."""
+    plan = dispatch.resolve(REQUIRE, {}, dispatch.TargetPaths(served=SERVED[DEP8]))
+    dispatch.install(plan, {})
+    misses = {
+        (path, key)
+        for path, key, reason in DECLINES_FP8_KV
+        if dispatch.classify_decline(path, key, reason) is D.MISSED
+    }
+    assert ("absorbed_query", "layer0") in misses
+    assert ("draft_absorbed_query", "layer78") in misses
+    assert ("indexer_layer_front", "target_layer0") in misses
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py
new file mode 100644
index 000000000000..3925082357b9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py
@@ -0,0 +1,501 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 dep target's construction, adoption, decode loop and step-loop hooks.
+
+Host-only. Adoption runs on the stand-in built-in modules and kernel factories of the tp8
+target's layer test (``test_modeling_v2_glm_layers.py``, loaded by path), with the fast paths
+of a dep cell: the attention modules and the speculative worker are re-classed, and the
+decoder layers, the MoE modules and the draft layer stay the built-in classes. The shell's
+construction runs with the built-in constructor stubbed and a recorder in place of
+``adopt_model``, as the tp8 target's shell test does.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import inspect
+import os
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2 import _router_index
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    attention,
+    dispatch,
+    layers,
+    runtime,
+    spec_worker,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
+from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
+from tensorrt_llm._torch.models.modeling_deepseekv3 import (
+    DeepseekV3DecoderLayer,
+    DeepseekV3Model,
+    Deepseekv3MoE,
+    DeepseekV3MTP,
+    DeepseekV32Attention,
+    GlmMoeDsaForCausalLM,
+)
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of
+from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+pytestmark = pytest.mark.cpu_only
+
+PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
+REQUIRE, AUTO = _router_index.ModelingV2Mode.REQUIRE, _router_index.ModelingV2Mode.AUTO
+
+
+def _load_tp8_layer_test():
+    path = (
+        Path(__file__).resolve().parents[1]
+        / "glm_5_2_nvfp4__sm_103__tp8"
+        / "test_modeling_v2_glm_layers.py"
+    )
+    spec = importlib.util.spec_from_file_location("_glm_tp8_layer_test_for_dep", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+T = _load_tp8_layer_test()
+
+DEP16, DEP8, DEP8_LM_HEAD_TP = (16, False, 3), (8, False, 3), (8, True, 5)
+CELLS = [DEP16, DEP8, DEP8_LM_HEAD_TP]
+IDS = ["dep16", "dep8", "dep8-lm-head-tp"]
+
+
+def _served(cell):
+    return modeling.CELLS[cell].paths.served
+
+
+@pytest.fixture
+def fast_paths(monkeypatch):
+    def install(enabled):
+        paths = T._FastPaths(frozenset(enabled))
+        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+        monkeypatch.setattr(dispatch, "log_capability", paths.log_capability)
+        return paths
+
+    return install
+
+
+def _state(model):
+    return dict(modeling.adopted_state(types.SimpleNamespace(model=model.model)))
+
+
+# --- adoption -------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("cell", [DEP16, DEP8], ids=IDS[:2])
+def test_with_an_fp8_kv_cache_the_attention_and_the_worker_are_adopted(fast_paths, cell):
+    paths = fast_paths(_served(cell))
+    model = T._built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+
+    # No fast-path object is built: the attention tail, the hops, the router GEMM and the draft
+    # front are not served.
+    assert paths.events == []
+    dense, moe, last, draft = model.model.layers
+    assert [type(m) for m in (dense, moe, last)] == [DeepseekV3DecoderLayer] * 3
+    assert type(draft) is DeepseekV3MTP
+    assert [type(m.mlp) for m in (moe, last, draft)] == [Deepseekv3MoE] * 3
+    for layer in (dense, moe, last, draft):
+        mla = layer.self_attn
+        assert type(mla) is attention.GlmMLA
+        assert type(mla.sparse_attn_hooks) is attention.GlmDsaHooks
+        assert type(mla.mqa) is attention.GlmDsaBackend
+        assert mla._attn_tail is None and mla._tail_prefetch is False
+        # The FMHA libraries stay the registry's.
+        assert "prepended_fmha_lib_classes" not in vars(mla.mqa)
+    assert [type(m.self_attn.mqa.indexer) for m in (dense, last)] == [attention.GlmIndexer] * 2
+    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
+
+
+def test_with_an_nvfp4_kv_cache_only_the_worker_is_adopted(fast_paths):
+    paths = fast_paths(_served(DEP8_LM_HEAD_TP))
+    model = T._built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+
+    assert paths.events == []
+    for layer in model.model.layers:
+        assert type(layer.self_attn) is DeepseekV32Attention
+        assert type(layer.self_attn.sparse_attn_hooks) is DSAHooks
+        assert type(layer.self_attn.mqa) is DSATrtllmAttention
+    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
+
+
+@pytest.mark.parametrize(
+    "disabled, attention_adopted, worker_adopted",
+    [
+        ({"rope_kv_append", "indexer_layer_front"}, False, True),
+        ({"indexer_layer_front"}, True, True),
+        ({"draft_loop", "decode_prologue"}, True, False),
+        (set(), True, True),
+    ],
+    ids=["no rope", "no indexer front", "no draft loop", "all"],
+)
+def test_each_adoption_follows_its_fast_path(
+    fast_paths, disabled, attention_adopted, worker_adopted
+):
+    fast_paths(_served(DEP8) - disabled)
+    model = T._built_in_model([])
+    modeling.adopt_model(model, plan=None)
+    expected = attention.GlmMLA if attention_adopted else DeepseekV32Attention
+    assert [type(layer.self_attn) for layer in model.model.layers] == [expected] * 4
+    worker = spec_worker.GlmMtpWorker if worker_adopted else MTPEagleWorker
+    assert type(model.spec_worker) is worker
+
+
+def test_without_a_plan_nothing_is_adopted(fast_paths):
+    fast_paths(modeling.NO_PATHS.served)
+    model = T._built_in_model([])
+    modeling.adopt_model(model, plan=None)
+    assert [type(layer.self_attn) for layer in model.model.layers] == [DeepseekV32Attention] * 4
+    assert type(model.spec_worker) is MTPEagleWorker
+
+
+def test_without_a_worker_the_attention_is_adopted_alone(fast_paths):
+    fast_paths(_served(DEP8))
+    model = T._built_in_model([])
+    model.spec_worker = None
+    modeling.adopt_model(model, plan=None)
+    assert type(model.model.layers[0].self_attn) is attention.GlmMLA
+
+
+@pytest.mark.parametrize("cell", CELLS, ids=IDS)
+def test_adoption_attaches_no_layer_state(fast_paths, cell):
+    """Every rank checks the same empty per-layer state: no fast-path object is attached to the
+    layers, their attention or their MoE."""
+    fast_paths(_served(cell))
+    model = T._built_in_model([])
+    modeling.adopt_model(model, plan=None)
+    state = _state(model)
+    assert list(state) == [f"model.layers.{i}" for i in range(4)]
+    assert all(attached == () for attached in state.values())
+
+
+@pytest.mark.parametrize("cell", CELLS, ids=IDS)
+def test_adoption_adds_no_parameter_buffer_or_module(fast_paths, cell):
+    fast_paths(_served(cell))
+    model = T._built_in_model([])
+    before = [T._keys(layer) for layer in model.model.layers]
+    modeling.adopt_model(model, plan=None)
+    assert [T._keys(layer) for layer in model.model.layers] == before
+
+
+def test_adoption_refuses_an_attention_module_that_is_not_the_builtin_class(fast_paths):
+    fast_paths(_served(DEP8))
+    model = T._built_in_model([])
+    attention.GlmMLA.adopt(model.model.layers[1].self_attn)
+    with pytest.raises(TypeError, match="GlmMLA adopts the built-in DeepseekV32Attention"):
+        modeling.adopt_model(model, plan=None)
+
+
+def test_the_target_uses_no_layer_class_of_the_family():
+    """The decoder layers, the MoE modules and the draft layer stay the built-in classes: the
+    target names none of the family's layer classes or their adoption."""
+    source = inspect.getsource(modeling)
+    for name in ("GlmDecoderLayer", "GlmMtpLayer", "GlmMoE", "layers.adopt_model"):
+        assert name not in source, name
+    assert modeling.adopted_state is layers.adopted_state
+
+
+# --- the core and its two targets -----------------------------------------------------------
+
+
+def _builtin_text_model() -> DeepseekV3Model:
+    model = DeepseekV3Model.__new__(DeepseekV3Model)
+    nn.Module.__init__(model)
+    return model
+
+
+def test_adopt_reclasses_the_builtin_text_model_in_place():
+    model = _builtin_text_model()
+    core = modeling.GlmMoeDsaDepModelingV2.adopt(model)
+    assert core is model and type(core) is modeling.GlmMoeDsaDepModelingV2
+    assert isinstance(core._prefill, modeling.PrefillTarget) and core._prefill.core is core
+    assert isinstance(core._decode, modeling.DecodeTarget) and core._decode.core is core
+    assert list(core.named_modules()) == [("", core)]
+    with pytest.raises(TypeError, match="GlmMoeDsaDepModelingV2"):
+        modeling.GlmMoeDsaDepModelingV2.adopt(core)
+
+
+@pytest.mark.parametrize("num_contexts, target", [(0, "_decode"), (1, "_prefill"), (4, "_prefill")])
+def test_each_step_goes_to_the_target_of_its_phase(num_contexts, target):
+    core = modeling.GlmMoeDsaDepModelingV2.adopt(_builtin_text_model())
+    md = types.SimpleNamespace(num_contexts=num_contexts)
+    assert core._select_target(md) is getattr(core, target)
+
+
+@pytest.fixture
+def decode_core():
+    calls = []
+    core = DeepseekV3Model.__new__(DeepseekV3Model)
+    nn.Module.__init__(core)
+    core.embed_tokens = T._Embedding(calls)
+    core.layers = nn.ModuleList([T._RecordingLayer(i, calls) for i in range(T.NUM_LAYERS + 1)])
+    core.num_hidden_layers = T.NUM_LAYERS
+    core.mapping_with_cp = None
+    core = modeling.GlmMoeDsaDepModelingV2.adopt(core)
+    for layer in core.layers:
+        layer.register_forward_pre_hook(
+            lambda m, a, k: calls.append(("hook", m.index)), with_kwargs=True
+        )
+    return core, calls
+
+
+@pytest.mark.parametrize("prologue_rows", [None, "rows"], ids=["no packet", "packet"])
+def test_the_decode_loop_runs_the_decoder_layers_as_modules(decode_core, prologue_rows):
+    core, calls = decode_core
+    rows = torch.full((2, T.HIDDEN), 7.0) if prologue_rows else None
+    seen = []
+
+    def take_embedding(attn_metadata, input_ids, embed_tokens):
+        seen.append((attn_metadata, input_ids, embed_tokens))
+        return rows
+
+    core.decode_prologue = types.SimpleNamespace(take_embedding=take_embedding)
+    md = T._NoContexts()
+    with dispatch.step(DECODE):
+        out = core._decode.forward(md, input_ids="ids", position_ids="pos", spec_metadata="spec")
+
+    assert seen == [(md, "ids", core.embed_tokens)]
+    embedded = [] if prologue_rows else [("embed", "ids")]
+    start = rows if prologue_rows else torch.zeros(2, T.HIDDEN)
+    assert torch.equal(out, start + T.NUM_LAYERS)
+    layer_calls = []
+    for i in range(T.NUM_LAYERS):
+        kwargs = calls[len(embedded) + 2 * i + 1][1]
+        assert kwargs["residual"] == (None if i == 0 else f"residual {i - 1}")
+        assert kwargs["attn_metadata"] is md and kwargs["spec_metadata"] == "spec"
+        layer_calls += [("hook", i), (i, kwargs, DECODE)]
+    # The draft layer is not run, and nothing follows the last decoder layer.
+    assert calls == embedded + layer_calls
+
+
+def test_without_the_decode_prologue_the_decode_loop_is_the_builtin_loop(decode_core):
+    core, calls = decode_core
+    with dispatch.step(DECODE):
+        out = core._decode.forward(T._NoContexts(), input_ids="ids")
+    assert torch.equal(out, torch.zeros(2, T.HIDDEN) + T.NUM_LAYERS)
+    assert [c[0] for c in calls] == ["embed"] + ["hook", 0, "hook", 1, "hook", 2]
+
+
+def test_the_decode_loop_reads_no_fast_path_and_never_the_step_kind():
+    source = inspect.getsource(modeling.DecodeTarget.forward)
+    assert "dispatch." not in source
+    assert "num_contexts" not in source and "num_ctx" not in source
+
+
+# --- construction ---------------------------------------------------------------------------
+
+
+def _model_config(cell, **facts):
+    ranks, lm_head_tp, draft = cell
+    defaults = {
+        DEP16: (QuantAlgo.FP8, "TRTLLM"),
+        DEP8: (QuantAlgo.FP8, "CUTEDSL"),
+        DEP8_LM_HEAD_TP: (QuantAlgo.NVFP4, "CUTEDSL"),
+    }[cell]
+    return types.SimpleNamespace(
+        mapping=types.SimpleNamespace(
+            world_size=ranks,
+            enable_attention_dp=True,
+            enable_lm_head_tp_in_adp=lm_head_tp,
+        ),
+        spec_config=types.SimpleNamespace(max_draft_len=draft),
+        quant_config=types.SimpleNamespace(kv_cache_quant_algo=facts.get("kv_cache", defaults[0])),
+        moe_backend=facts.get("moe_backend", defaults[1]),
+    )
+
+
+RANK_UNIFORM_ENV = dispatch.SHARED_ENV + (_router_index.MODELING_V2_VALIDATE_ENV,)
+
+
+@pytest.fixture
+def shell_construction(monkeypatch):
+    """The shell class with the built-in constructor stubbed, and the events of a construction."""
+    events = []
+
+    def adopt_model(model, plan):
+        events.append(("adopt", type(model.model).__name__, plan))
+
+    def builtin_init(self, model_config):
+        events.append("built-in model")
+        events.append(("installed", dispatch.installed()))
+        nn.Module.__init__(self)
+        self.model_config = model_config
+        self.model = _builtin_text_model()
+        self.model.layers = nn.ModuleList([nn.Module(), nn.Module()])
+
+    def allgather(value, world_size):
+        events.append(("allgather", value, world_size))
+        return [value] * world_size
+
+    shell_cls = modeling.ModelingV2Glm52Nvfp4Sm103Dep
+    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__init__", builtin_init)
+    monkeypatch.setattr(modeling, "adopt_model", adopt_model)
+    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__post_init__", lambda self: None)
+    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__pp_init__", lambda self: None)
+    monkeypatch.setattr(dispatch, "_require", [False])
+    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
+    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
+    monkeypatch.setattr(dispatch, "_TARGET", [dispatch.DEFAULT_TARGET])
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: events.append(message))
+    monkeypatch.setattr(modeling.logger, "warning", lambda message: events.append(message))
+    for name in list(os.environ):
+        if name.startswith(dispatch.KNOB_PREFIX) or name in RANK_UNIFORM_ENV:
+            monkeypatch.delenv(name)
+
+    def construct(mode, model_config):
+        monkeypatch.setitem(_router_index._RESOLVED_MODES, shell_cls, mode)
+        return shell_cls(model_config)
+
+    return construct, events
+
+
+@pytest.mark.parametrize("cell", CELLS, ids=IDS)
+def test_the_shell_checks_the_cell_the_ranks_and_the_plan_then_adopts(shell_construction, cell):
+    construct, events = shell_construction
+    shell = construct(REQUIRE, _model_config(cell))
+
+    allgather, record, built_in, installed, adopt, adopted = events
+    assert allgather == ("allgather", ((dispatch.MODE_KEY, "require"),), cell[0])
+    served = modeling.CELLS[cell].paths.served
+    assert shell.plan.describe().startswith(
+        f"modeling_v2=require, {len(served)}/{len(served)} fast paths"
+    )
+    assert record.startswith("[glm_moe_dsa] dispatch: record paths=")
+    assert built_in == "built-in model"
+    assert installed == ("installed", served)
+    assert adopt == ("adopt", "GlmMoeDsaDepModelingV2", shell.plan)
+    assert adopted == ("allgather", modeling.adopted_state(shell), cell[0])
+    assert shell.plan.target is modeling.CELLS[cell].paths
+    assert type(shell.model) is modeling.GlmMoeDsaDepModelingV2
+    assert "_runtime_hooks" not in vars(shell)
+
+
+@pytest.mark.parametrize(
+    "cell, facts",
+    [
+        (DEP16, {"moe_backend": "CUTEDSL"}),
+        (DEP8, {"kv_cache": QuantAlgo.NVFP4}),
+        (DEP8_LM_HEAD_TP, {"kv_cache": QuantAlgo.FP8}),
+    ],
+    ids=IDS,
+)
+def test_outside_the_cells_require_raises_before_any_collective(shell_construction, cell, facts):
+    construct, events = shell_construction
+    with pytest.raises(ValueError, match="is not a deployment this target serves"):
+        construct(REQUIRE, _model_config(cell, **facts))
+    assert events == []
+
+
+def test_outside_the_cells_auto_builds_the_builtin_model_without_fast_paths(shell_construction):
+    construct, events = shell_construction
+    shell = construct(AUTO, _model_config(DEP8, moe_backend="TRTLLM"))
+
+    warning, allgather, record = events[:3]
+    assert "is not a deployment this target serves" in warning
+    assert allgather == ("allgather", ((dispatch.MODE_KEY, "auto"),), 8)
+    assert record.startswith("[glm_moe_dsa] dispatch: record paths=<none> ")
+    assert shell.plan.describe().startswith("modeling_v2=auto, 0/0 fast paths")
+    assert ("installed", frozenset()) in events
+    assert shell.plan.target is modeling.NO_PATHS
+    assert not dispatch.require_mode()
+
+
+def test_a_disable_knob_turns_its_path_and_dependents_off(monkeypatch, shell_construction):
+    construct, _ = shell_construction
+    monkeypatch.setenv(dispatch.DISABLE_ENV, "draft_loop")
+    shell = construct(REQUIRE, _model_config(DEP8))
+    off = {name for name in dispatch.FAST_PATH_NAMES if not dispatch.is_enabled(name)}
+    assert off >= {"draft_loop", "decode_prologue"}
+    assert dispatch.installed() == modeling.CELLS[DEP8].paths.served - {
+        "draft_loop",
+        "decode_prologue",
+    }
+    assert shell.plan.describe().startswith("modeling_v2=require, 8/10 fast paths")
+
+
+# --- the step-loop hooks --------------------------------------------------------------------
+
+
+def _model_of_cell(cell):
+    paths = modeling.CELLS[cell].paths
+    return types.SimpleNamespace(
+        plan=dispatch.resolve(REQUIRE, {}, paths),
+        model=types.SimpleNamespace(decode_prologue=None),
+        model_config=types.SimpleNamespace(mapping=types.SimpleNamespace(enable_attention_dp=True)),
+    )
+
+
+@pytest.fixture
+def quiet_dispatch(monkeypatch):
+    monkeypatch.setattr(dispatch, "_require", [False])
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+
+
+@pytest.mark.usefixtures("quiet_dispatch")
+@pytest.mark.parametrize("cell", CELLS, ids=IDS)
+def test_the_hooks_declare_the_decode_prologue_and_no_step_loop_part(cell):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E501
+        DecodePrologueHook,
+    )
+
+    model = _model_of_cell(cell)
+    hooks = runtime.make_runtime_hooks(model)
+    assert isinstance(hooks, RuntimeHooks)
+    assert not hooks.loop_on_execution_stream
+    assert hooks.step_input_stager is None and hooks.sampler_store_writer is None
+    assert isinstance(hooks.decode_prologue, DecodePrologueHook)
+    assert model.model.decode_prologue is hooks.decode_prologue
+    assert isinstance(hooks.executor_events, runtime.GlmExecutorEvents)
+
+
+@pytest.mark.usefixtures("quiet_dispatch")
+@pytest.mark.parametrize("cell", CELLS, ids=IDS)
+def test_the_executor_events_install_only_the_host_policy(monkeypatch, cell):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
+        host_page_retention,
+    )
+
+    calls = []
+    monkeypatch.setattr(host_page_retention, "install_host_policy", calls.append)
+    events = runtime.make_runtime_hooks(_model_of_cell(cell)).executor_events
+    events.after_warmup()
+    events.before_overlap_loop()
+    events.at_shutdown()
+    assert calls == ["after_warmup"]
+
+
+def test_the_shell_builds_its_hooks_on_the_first_read(monkeypatch, shell_construction):
+    construct, _ = shell_construction
+    shell = construct(REQUIRE, _model_config(DEP8))
+    built = []
+
+    def make_runtime_hooks(model):
+        built.append(model)
+        return RuntimeHooks()
+
+    monkeypatch.setattr(runtime, "make_runtime_hooks", make_runtime_hooks)
+    hooks = runtime_hooks_of(shell)
+    assert built == [shell] and runtime_hooks_of(shell) is hooks
+
+    def broken(model):
+        raise AttributeError("plan")
+
+    shell2 = construct(REQUIRE, _model_config(DEP8))
+    monkeypatch.setattr(runtime, "make_runtime_hooks", broken)
+    with pytest.raises(RuntimeError, match="step-loop hooks"):
+        runtime_hooks_of(shell2)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py
new file mode 100644
index 000000000000..f22584d3245e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py
@@ -0,0 +1,328 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 tep target's adoption at each tensor-parallel size, the family's layer
+composition under its plans, its decode loop, and how its ``modeling.py`` differs from the tp8
+target's.
+
+Host-only, with the stand-in built-in modules and kernel factories of the tp8 target's layer
+test (``test_modeling_v2_glm_layers.py``, loaded by path): the fast paths the stand-ins see are
+the ones the tep plan serves at that size.
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib.util
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tep import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
+    modeling as tp8_modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    attention,
+    dispatch,
+    layers,
+    spec_worker,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+
+def _load_tp8_layer_test():
+    path = (
+        Path(__file__).resolve().parents[1]
+        / "glm_5_2_nvfp4__sm_103__tp8"
+        / "test_modeling_v2_glm_layers.py"
+    )
+    spec = importlib.util.spec_from_file_location("_glm_tp8_layer_test", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+T = _load_tp8_layer_test()
+
+
+def _served(tp_size):
+    return modeling.TARGET_PATHS[tp_size].served
+
+
+@pytest.fixture
+def fast_paths(monkeypatch):
+    def install(enabled, post_attn_hop=True):
+        paths = T._FastPaths(frozenset(enabled))
+        if not post_attn_hop:
+            # As at four ranks: the call site's native all-reduce is not multi-node NVLink, so
+            # the hop factory returns None.
+            paths.entries[("post_attention_hop", "make_post_attn_hop")] = lambda *a, **k: None
+        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+        monkeypatch.setattr(dispatch, "log_capability", paths.log_capability)
+        return paths
+
+    return install
+
+
+def _state(model):
+    return dict(modeling.adopted_state(types.SimpleNamespace(model=model.model)))
+
+
+def test_at_eight_ranks_adoption_builds_the_tp8_objects_but_the_routed_moe_ones(fast_paths):
+    paths = fast_paths(_served(8))
+    model = T._built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+
+    tail = [("make", "attn_tail"), ("make", "staged_attn_tail")]
+    pre_moe = tail + [("make", "post_attn_hop"), ("make", "into_hop")]
+    post_moe = [("make", "router_gemm"), ("make", "post_moe_hop")]
+    expected = (
+        pre_moe
+        + pre_moe
+        + post_moe
+        + [("make", "deferred_norm_evict_hints_hop"), ("deferred hop serves", 2, "model config")]
+        + pre_moe
+        + post_moe
+        + [
+            ("make", "post_attn_hop"),
+            ("make", "router_gemm"),
+            ("log_capability", "draft_front", "draft_front", True, ""),
+            ("make", "draft_live_row_hop"),
+        ]
+    )
+    assert [e for e in paths.events if e[0] != "set_successor_weights"] == expected
+    dense, moe, last, draft = model.model.layers
+    assert [type(m) for m in (dense, moe, last)] == [layers.GlmDecoderLayer] * 3
+    assert type(draft) is layers.GlmMtpLayer
+    assert [type(m.mlp) for m in (moe, last, draft)] == [layers.GlmMoE] * 3
+    assert type(moe.self_attn) is attention.GlmMLA
+    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
+    state = _state(model)
+    assert [entry[0] for entry in state["model.layers.1"]] == [
+        "post_attn_hop",
+        "_into_hop",
+        "post_moe_hop",
+        "deferred_norm_hop",
+        "self_attn._attn_tail",
+        "mlp.router_gemm",
+    ]
+    assert [entry[0] for entry in state["model.layers.3"]] == [
+        "post_attn_hop",
+        "draft_front",
+        "draft_live_row_hop",
+        "mlp.router_gemm",
+    ]
+
+
+def test_at_four_ranks_adoption_builds_the_router_gemm_and_the_draft_front(fast_paths):
+    paths = fast_paths(_served(4), post_attn_hop=False)
+    model = T._built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+
+    assert [e for e in paths.events if e[0] == "make"] == [("make", "router_gemm")] * 3
+    state = _state(model)
+    assert [entry[0] for entry in state["model.layers.0"]] == []
+    assert [entry[0] for entry in state["model.layers.1"]] == ["mlp.router_gemm"]
+    assert [entry[0] for entry in state["model.layers.3"]] == ["draft_front", "mlp.router_gemm"]
+    for layer in model.model.layers:
+        assert layer.post_attn_hop is None and layer._into_hop is None
+        assert layer.post_moe_hop is None and layer.deferred_norm_hop is None
+        assert layer.self_attn._attn_tail is None
+
+
+# The fast paths only the tp8 target serves: their kernels take tensor-parallel experts. The
+# objects only they attach, as (owner submodule, attribute, path).
+_ROUTED_MOE_PATHS = frozenset(
+    {
+        "routed_moe",
+        "routed_moe_into_hop",
+        "hop_routing",
+        "shared_expert_prefetch",
+        "draft_routed_moe",
+        "cross_layer_expert_prefetch",
+    }
+)
+_ROUTED_MOE_STATE = (
+    ("", "_hop_routing", "hop_routing"),
+    ("mlp", "_routed_into_hop", "routed_moe_into_hop"),
+)
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_adoption_attaches_no_routed_moe_object(fast_paths, tp_size):
+    paths = fast_paths(_served(tp_size), post_attn_hop=tp_size == 8)
+    model = T._built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+    for layer in model.model.layers:
+        for submodule, name, _ in _ROUTED_MOE_STATE:
+            owner = getattr(layer, submodule) if submodule else layer
+            assert getattr(owner, name, None) is None, (submodule, name)
+        assert layer._expert_prefetch_link is None
+
+
+def _enclosing_tests(func, target):
+    """The tests of the ``if`` statements whose body contains ``target``."""
+    found = []
+
+    def visit(node, stack):
+        if node is target:
+            found.extend(stack)
+            return True
+        if isinstance(node, ast.If):
+            return any(visit(c, stack + [ast.unparse(node.test)]) for c in node.body) or any(
+                visit(c, stack) for c in node.orelse
+            )
+        return any(visit(c, stack) for c in ast.iter_child_nodes(node))
+
+    visit(func, [])
+    return found
+
+
+def _paths_named(tree):
+    """Every fast path a ``dispatch.is_enabled`` / ``dispatch.get_impl`` call names."""
+    return {
+        call.args[0].value
+        for call in ast.walk(tree)
+        if isinstance(call, ast.Call)
+        and ast.unparse(call.func) in ("dispatch.is_enabled", "dispatch.get_impl")
+        and call.args
+        and isinstance(call.args[0], ast.Constant)
+    }
+
+
+def test_every_path_the_composition_names_is_served_or_routed_moe():
+    """The family's layers and this target's ``modeling.py`` name, besides the paths the target
+    serves at some size, only routed-MoE paths, which it serves at no size: their blocks stay
+    off under both plans."""
+    served = _served(4) | _served(8)
+    assert not served & _ROUTED_MOE_PATHS
+    for module in (layers, modeling):
+        named = _paths_named(ast.parse(Path(module.__file__).read_text()))
+        assert named <= set(dispatch.FAST_PATH_NAMES), module.__name__
+        unserved = named - served
+        assert unserved <= _ROUTED_MOE_PATHS, (module.__name__, sorted(unserved))
+    named = _paths_named(ast.parse(Path(layers.__file__).read_text()))
+    assert named & _ROUTED_MOE_PATHS == _ROUTED_MOE_PATHS - {"cross_layer_expert_prefetch"}
+
+
+def test_the_routed_moe_objects_are_attached_only_under_their_paths():
+    """``_hop_routing`` and the MoE's ``_routed_into_hop`` are set only inside a block gated on
+    their path, so under a plan without it they keep their class default None."""
+    tree = ast.parse(Path(layers.__file__).read_text())
+    seen = set()
+    for func in ast.walk(tree):
+        if not isinstance(func, ast.FunctionDef):
+            continue
+        for node in ast.walk(func):
+            if not isinstance(node, ast.Assign):
+                continue
+            for target in node.targets:
+                for submodule, name, path in _ROUTED_MOE_STATE:
+                    owner = f"self.{submodule}" if submodule else "self"
+                    if ast.unparse(target) == f"{owner}.{name}":
+                        seen.add(name)
+                        assert any(
+                            f"dispatch.is_enabled({path!r})" in test
+                            for test in _enclosing_tests(func, node)
+                        ), (func.name, name)
+    assert seen == {name for _, name, _ in _ROUTED_MOE_STATE}
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_the_moe_runs_the_builtin_experts_on_the_produced_logits(fast_paths, monkeypatch, tp_size):
+    fast_paths(_served(tp_size))
+    calls = []
+    moe = T._moe(calls)
+    moe.__class__ = layers.GlmMoE
+    moe.router_gemm = lambda hidden: "router gemm logits"
+    monkeypatch.setattr(dispatch, "phase", lambda: dispatch.Phase.DECODE)
+
+    hidden = torch.zeros(2, T.HIDDEN)
+    assert moe.compute_routed_output(hidden, None, [2], True) == "routed"
+    assert calls[-1][0][1] == "router gemm logits"
+    assert moe.compute_routed_output(hidden, None, [2], False, "hop logits") == "routed"
+    assert calls[-1][0][1] == "hop logits" and calls[-1][1]["do_finalize"] is False
+    moe.router_gemm = None
+    moe.compute_routed_output(hidden, None, [2], True)
+    assert calls[-1][0][1] == "gate logits"
+
+
+def test_the_decode_loop_has_no_expert_prefetch_join():
+    source = inspect.getsource(modeling.DecodeTarget.forward)
+    assert "cross_layer_expert_prefetch" not in source
+    assert "setup_aliases" not in vars(modeling.ModelingV2Glm52Nvfp4Sm103Tep)
+
+
+def _function_bodies(module):
+    """Every function of ``module`` as (qualified name, AST dump without docstrings)."""
+    tree = ast.parse(Path(module.__file__).read_text())
+    out = {}
+
+    def strip(node):
+        for child in ast.walk(node):
+            body = getattr(child, "body", None)
+            if (
+                isinstance(body, list)
+                and body
+                and isinstance(body[0], ast.Expr)
+                and isinstance(getattr(body[0], "value", None), ast.Constant)
+                and isinstance(body[0].value.value, str)
+            ):
+                child.body = body[1:] or [ast.Pass()]
+        return node
+
+    for cls in tree.body:
+        if isinstance(cls, ast.ClassDef):
+            for fn in cls.body:
+                if isinstance(fn, ast.FunctionDef):
+                    out[f"{cls.name}.{fn.name}"] = ast.dump(strip(fn))
+        elif isinstance(cls, ast.FunctionDef):
+            out[cls.name] = ast.dump(strip(cls))
+    return out
+
+
+# The tp8 target's class each tep class corresponds to, and the functions the tep target
+# changes: its plan per tensor-parallel size, and no cross-layer expert prefetch (a routed-MoE
+# path) in its decode loop and shell.
+_RENAMES = {
+    "GlmMoeDsaTepModelingV2": "GlmMoeDsaTp8ModelingV2",
+    "ModelingV2Glm52Nvfp4Sm103Tep": "ModelingV2Glm52Nvfp4Sm103Tp8",
+}
+_CHANGED = {"DecodeTarget.forward", "ModelingV2Glm52Nvfp4Sm103Tp8.__init__"}
+
+
+def test_the_targets_modeling_differs_only_in_the_declared_functions():
+    """Every function of the tep ``modeling.py`` equals the tp8 one's (with the tep class
+    names), except exactly the declared ones; the tp8 module has only the declared extra
+    function. Both adopt their layers with the family's ``adopt_model``."""
+    ours = {}
+    for name, body in _function_bodies(modeling).items():
+        for new, old in _RENAMES.items():
+            name = name.replace(new, old)
+            body = body.replace(f"'{new}'", f"'{old}'")
+        ours[name] = body
+    theirs = _function_bodies(tp8_modeling)
+    differ = {name for name in ours if name in theirs and ours[name] != theirs[name]}
+    assert differ == _CHANGED, (sorted(differ - _CHANGED), sorted(_CHANGED - differ))
+    only_tp8 = set(theirs) - set(ours)
+    assert only_tp8 <= {"ModelingV2Glm52Nvfp4Sm103Tp8.setup_aliases"}, sorted(only_tp8)
+    assert set(ours) - set(theirs) == set()
+    assert modeling.adopt_model is tp8_modeling.adopt_model is layers.adopt_model
+    assert modeling.adopted_state is tp8_modeling.adopted_state is layers.adopted_state
+
+
+def test_the_classes_add_no_parameter_buffer_or_module(fast_paths):
+    paths = fast_paths(_served(8))
+    model = T._built_in_model(paths.events)
+    before = [T._keys(layer) for layer in model.model.layers]
+    modeling.adopt_model(model, plan=None)
+    assert [T._keys(layer) for layer in model.model.layers] == before
+    assert all(isinstance(layer, nn.Module) for layer in model.model.layers)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py
new file mode 100644
index 000000000000..9fea5839cfd3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py
@@ -0,0 +1,355 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The fast paths the GLM-5.2 tep target serves at each tensor-parallel size, and what its
+plan does with the declines its deployments log.
+
+Host-only: the plan is resolved and installed from the target's ``TARGET_PATHS`` exactly as
+the shell does at construction; the declines are the (path, key, reason) triples the kernels
+log at those deployments, decided at construction or in the warm-up, and none of them may be
+a miss: under ``modeling_v2='require'`` a miss raises.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tep import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
+    modeling as tp8_modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+D = dispatch.Decline
+REQUIRE = ModelingV2Mode.REQUIRE
+
+# The paths whose kernels take only tensor-parallel experts (every expert on every rank).
+ROUTED_MOE = {
+    "routed_moe",
+    "routed_moe_into_hop",
+    "hop_routing",
+    "shared_expert_prefetch",
+    "draft_routed_moe",
+    "cross_layer_expert_prefetch",
+}
+# At four ranks: the paths that engage there, and the bases they extend.
+TEP4_ENGAGED = {
+    "indexer_layer_front",
+    "rope_kv_append",
+    "indexer_topk",
+    "cross_rank_topk",
+    "cross_rank_topk_cluster_merge",
+    "cross_rank_topk_global_emit",
+    "router_gemm",
+    "dense_mlp",
+    "draft_front",
+    "draft_tail",
+    "draft_loop",
+    "decode_prologue",
+    "deferred_input_prep",
+    "sampler_store",
+    "execution_stream_loop",
+    "host_page_retention",
+}
+TEP4_BASES = {"absorbed_query", "draft_absorbed_query", "post_attention_hop"}
+SERVED = {
+    4: TEP4_ENGAGED | TEP4_BASES,
+    8: set(dispatch.FAST_PATH_NAMES) - ROUTED_MOE,
+}
+
+
+def _layers(spec: str) -> list:
+    out = []
+    for part in spec.split(","):
+        lo, _, hi = part.partition("-")
+        out += list(range(int(lo), int(hi or lo) + 1))
+    return out
+
+
+FULL = "0-2,6,10,14,18,22,26,30,34,38,42,46,50,54,58,62,66,70,74"
+SHARED = (
+    "3-5,7-9,11-13,15-17,19-21,23-25,27-29,31-33,35-37,39-41,43-45,47-49,51-53,55-57,59-61,"
+    "63-65,67-69,71-73,75-77"
+)
+TP8_DOMAIN = "shape outside the GLM-5.2 TP8 domain"
+
+# The declines the served paths log at four ranks (bs 2 and bs 5), on every rank.
+DECLINES_TP4 = (
+    *(("absorbed_query", f"layer{i}", TP8_DOMAIN) for i in range(78)),
+    ("absorbed_query", "layer78", "target layers only"),
+    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
+    ("draft_absorbed_query", "layer78", TP8_DOMAIN),
+    *(("indexer_layer_front", f"target_layer{i}", TP8_DOMAIN) for i in range(78)),
+    ("indexer_layer_front", "draft_layer78", TP8_DOMAIN),
+    (
+        "post_attention_hop",
+        "post_attn_hop:native all-reduce (AUTO) does not use MNNVL",
+        "native all-reduce (AUTO) does not use MNNVL",
+    ),
+    ("router_gemm", "router:not a target layer", "not a target layer"),
+    ("indexer_topk", "indexer_decode", "MTP draft step (target layers only)"),
+    ("indexer_topk", "indexer_decode", "5 generation requests > 4"),
+    ("cross_rank_topk", "cp_indexer_decode", "2 generation requests != 1"),
+    ("cross_rank_topk_cluster_merge", "cp_indexer_decode", "4 generation requests != 1"),
+    ("cross_rank_topk_global_emit", "cp_indexer_decode", "5 generation requests > 4"),
+    ("dense_mlp", "dense_mlp:row count outside [1, 8]", "row count outside [1, 8]"),
+    ("draft_loop", "draft_loop:context rows present", "context rows present"),
+    ("draft_loop", "draft_step:dense top-k fallback active", "dense top-k fallback active"),
+    (
+        "deferred_input_prep",
+        "prep[not a steady decode-graph replay]",
+        "not a steady decode-graph replay",
+    ),
+    (
+        "deferred_input_prep",
+        "prep[no previous speculative device tensors]",
+        "no previous speculative device tensors",
+    ),
+    (
+        "deferred_input_prep",
+        "prep[dummy row or row without previous tensors]",
+        "dummy row or row without previous tensors",
+    ),
+    (
+        "decode_prologue",
+        "target_prologue[not a generation-only step]",
+        "not a generation-only step",
+    ),
+)
+
+# The declines the served paths log at eight ranks (bs 5), on every rank.
+DECLINES_TP8 = (
+    ("absorbed_query", "layer78", "target layers only"),
+    ("attention_output_tail", "layer78", "draft or unindexed layer"),
+    *(
+        (
+            "shared_indexer_absorbed_query",
+            f"layer{i}",
+            "full indexer layer (shared_indexer_absorbed_query serves shared layers)",
+        )
+        for i in _layers(FULL)
+    ),
+    ("shared_indexer_absorbed_query", "layer78", "target layers only"),
+    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
+    *(
+        (
+            "deferred_norm",
+            f"scope[shared]_into_layer{i}",
+            "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE=shared: layer "
+            f"{i} is a full-indexer layer; post_moe_hop hop",
+        )
+        for i in _layers(FULL)[3:]
+    ),
+    *(
+        (
+            "indexer_layer_front",
+            f"target_layer{i}",
+            "shared-indexer layer (no q_a RMSNorm launch to remove; "
+            "shared_indexer_absorbed_query / absorbed_query serve it)",
+        )
+        for i in _layers(SHARED)
+    ),
+    *(
+        (
+            "fused_kv_indexer_projection",
+            f"layer{i}",
+            "full layers not selected (TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS)",
+        )
+        for i in _layers(FULL)
+    ),
+    (
+        "fused_kv_indexer_projection",
+        "layer78",
+        "draft layers not selected (TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS)",
+    ),
+    ("router_gemm", "router:not a target layer", "not a target layer"),
+    (
+        "sparse_mla_decode",
+        "request:not generation-only",
+        "not generation-only; first seen at layer 3 with 1 rows",
+    ),
+    ("indexer_topk", "indexer_decode", "MTP draft step (target layers only)"),
+    ("indexer_topk", "indexer_decode", "5 generation requests > 4"),
+    ("cross_rank_topk", "cp_indexer_decode", "2 generation requests != 1"),
+    ("cross_rank_topk_global_emit", "cp_indexer_decode", "5 generation requests > 4"),
+    ("dense_mlp", "dense_mlp:row count outside [1, 8]", "row count outside [1, 8]"),
+    ("draft_loop", "draft_loop:context rows present", "context rows present"),
+    (
+        "deferred_input_prep",
+        "prep[not a steady decode-graph replay]",
+        "not a steady decode-graph replay",
+    ),
+    (
+        "decode_prologue",
+        "target_prologue[not a generation-only step]",
+        "not a generation-only step",
+    ),
+)
+DECLINES = {4: DECLINES_TP4, 8: DECLINES_TP8}
+
+# The layer-keyed capabilities that engage at eight ranks: (path, key family, layers).
+ENGAGED_TP8 = (
+    ("absorbed_query", "layer", "0-77"),
+    ("attention_output_tail", "layer", "0-77"),
+    ("oproj_weight_staging", "layer", "0-77"),
+    ("shared_indexer_absorbed_query", "layer", SHARED),
+    ("draft_absorbed_query", "layer", "78"),
+    ("sparse_mla_decode", "layer", "0-78"),
+    ("sparse_mla_decode", "call:layer", "0-78"),
+    ("fused_kv_indexer_projection", "layer", "3"),
+    ("indexer_layer_front", "target_layer", FULL),
+    ("indexer_layer_front", "draft_layer", "78"),
+    ("deferred_norm", "scope[shared]_into_layer", SHARED.replace("3-5", "4-5")),
+    ("deferred_norm_evict_hints", "projection_layer", SHARED.replace("3-5", "4-5")),
+)
+
+
+@pytest.fixture
+def installed(monkeypatch):
+    """Resolve and install the target's plan at a tensor-parallel size, as the shell does."""
+    for name in ("_INSTALLED", "_VALIDATE", "_require", "_TARGET"):
+        monkeypatch.setattr(dispatch, name, list(getattr(dispatch, name)))
+    monkeypatch.setattr(dispatch, "_logged", set())
+    monkeypatch.setattr(dispatch, "_missed", {})
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+    dispatch.classify_decline.cache_clear()
+
+    def install(tp_size, environ=None):
+        plan = dispatch.resolve(REQUIRE, environ or {}, modeling.TARGET_PATHS[tp_size])
+        dispatch.install(plan, environ or {})
+        return plan
+
+    yield install
+    dispatch.classify_decline.cache_clear()
+
+
+def test_the_target_has_a_table_per_tensor_parallel_size():
+    assert sorted(modeling.TARGET_PATHS) == [4, 8]
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_the_target_serves_its_paths_at_each_size(installed, tp_size):
+    plan = installed(tp_size)
+    assert set(plan.served) == SERVED[tp_size]
+    assert dispatch.installed() == frozenset(SERVED[tp_size])
+    assert plan.describe().startswith(
+        f"modeling_v2=require, {len(SERVED[tp_size])}/{len(SERVED[tp_size])} fast paths"
+    )
+
+
+def test_the_sizes_serve_nineteen_and_thirty_paths():
+    assert len(SERVED[4]) == 19 and len(SERVED[8]) == 30
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_no_served_path_extends_one_the_target_does_not_serve(tp_size):
+    served = SERVED[tp_size]
+    for path in served:
+        assert dispatch.REQUIRED_FEATURES.get(path, frozenset()) <= served, path
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_the_target_never_serves_a_routed_moe_path(tp_size):
+    assert not (SERVED[tp_size] & ROUTED_MOE)
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_every_kernel_package_is_the_familys(tp_size):
+    """The target has no kernel package of its own: every served path imports from shared."""
+    assert dict(modeling.TARGET_PATHS[tp_size].kernel_modules) == {}
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_the_record_line_lists_the_served_paths(installed, monkeypatch, tp_size):
+    lines = []
+    monkeypatch.setattr(dispatch, "_log_info", lines.append)
+    installed(tp_size)
+    (record,) = [line for line in lines if "record" in line]
+    listed = record.split("paths=", 1)[1].split(" ", 1)[0].split(",")
+    assert {name.split("[", 1)[0] for name in listed} == SERVED[tp_size]
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_disable_may_name_only_a_served_path(installed, tp_size):
+    with pytest.raises(ValueError, match="not fast paths of this target"):
+        installed(tp_size, {dispatch.DISABLE_ENV: "routed_moe"})
+    plan = installed(tp_size, {dispatch.DISABLE_ENV: "draft_tail"})
+    assert not plan.enabled("draft_tail") and plan.enabled("draft_loop")
+
+
+def test_disable_turns_the_dependents_off_at_four_ranks(installed):
+    plan = installed(4, {dispatch.DISABLE_ENV: "absorbed_query"})
+    assert not plan.enabled("indexer_layer_front")
+    assert plan.describe().startswith("modeling_v2=require, 17/19 fast paths")
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_the_routed_accumulation_knob_is_refused(installed, tp_size):
+    with pytest.raises(ValueError, match="does not serve routed_moe_into_hop"):
+        installed(tp_size, {dispatch.ACCUMULATION_ENV: "identity"})
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_the_cross_rank_wait_timeout_is_read(installed, tp_size):
+    plan = installed(tp_size, {dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV: "60000"})
+    assert plan.knobs == ((dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV, "60000"),)
+
+
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_no_logged_decline_is_a_miss(installed, tp_size):
+    """Under require, a decline decided once inside a path's cells would raise."""
+    installed(tp_size)
+    assert dispatch.require_mode()
+    for path, key, reason in DECLINES[tp_size]:
+        assert path in SERVED[tp_size], path
+        assert dispatch.classify_decline(path, key, reason) is not D.MISSED, (path, key, reason)
+        assert dispatch.log_capability(path, key, False, reason) is False
+    assert all(count >= 0 for count in dispatch.missed_counts().values())
+
+
+def test_at_four_ranks_the_bases_decline_as_uncovered(installed):
+    installed(4)
+    for path, key, reason in DECLINES_TP4:
+        if path in TEP4_BASES or key.startswith(("target_layer", "draft_layer")):
+            assert dispatch.classify_decline(path, key, reason) is D.UNCOVERED, (path, key)
+
+
+def test_the_four_rank_tables_would_raise_at_the_tp8_cells(installed):
+    """The tp8 cells hold the query fronts and the hop that four ranks decline on: the tep
+    table is what keeps a four-rank deployment from raising under require."""
+    installed(8)
+    misses = [
+        (path, key)
+        for path, key, reason in DECLINES_TP4
+        if dispatch.classify_decline(path, key, reason) is D.MISSED
+    ]
+    assert ("absorbed_query", "layer0") in misses
+    assert ("post_attention_hop", "post_attn_hop:native all-reduce (AUTO) does not use MNNVL") in (
+        misses
+    )
+
+
+def test_at_eight_ranks_every_engaged_layer_is_inside_the_cells(installed):
+    target = modeling.TARGET_PATHS[8]
+    for path, family, layers in ENGAGED_TP8:
+        cells = target.layer_cells[path][family]
+        assert set(_layers(layers)) <= cells, (path, family)
+
+
+def test_the_tp8_target_serves_every_path_with_the_family_tables():
+    target = tp8_modeling.TARGET_PATHS
+    assert target.served == frozenset(dispatch.FAST_PATH_NAMES)
+    assert target.layer_cells is dispatch.LAYER_CELLS
+    assert dict(target.uncovered) == {}
+    assert {path for path in target.kernel_modules} == ROUTED_MOE
+    assert all(
+        module.startswith(tp8_modeling.__package__ + ".kernels.")
+        for module in target.kernel_modules.values()
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py
new file mode 100644
index 000000000000..724c6c33a153
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py
@@ -0,0 +1,267 @@
+# 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.
+"""Tests of package cross_layer_expert_prefetch (cross-layer speculative routed-expert L2 prefetch)."""
+
+import pytest
+import torch
+from torch import nn
+
+cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+NE = 256
+HIDDEN = 6144
+KMAX = 32
+STAT_KS = (8, 12, 18, 24, 32)
+
+
+def _ref_union(logits, bias, rows):
+    s = torch.sigmoid(logits[:rows].float()) + bias.float()[None, :]
+    thr = s.sort(dim=1, descending=True).values[:, 7:8]
+    sel = s >= thr
+    return sel.sum(0).to(torch.int32), s.max(0).values
+
+
+def _ref_rank(logits, bias, rows):
+    cnt, mx = _ref_union(logits, bias, rows)
+    order = sorted(range(NE), key=lambda e: (-int(cnt[e]), -float(mx[e]), e))
+    return order[:KMAX]
+
+
+@pytest.fixture()
+def spec():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
+        spec_experts,
+    )
+
+    return spec_experts
+
+
+def test_the_served_config(spec):
+    cfg = spec.config()
+    assert cfg is spec.SERVED_CONFIG
+    assert (
+        cfg.k,
+        cfg.start,
+        cfg.pred,
+        cfg.fused,
+        cfg.delay_ns,
+        cfg.chunk,
+        cfg.ctas,
+        cfg.issue,
+        cfg.xc_ctas,
+    ) == (8, "moe", "self", True, 0, 16384, 16, "xc", 32)
+    assert cfg.start in spec.STARTS and cfg.pred in spec.PREDS and cfg.issue in spec.ISSUES
+    assert spec.ExpertPrefetchConfig(8, "moe", "self", True, 0, 16384, 16) == cfg
+
+
+@cuda
+@pytest.mark.parametrize("rows,ties", [(1, False), (6, False), (8, False), (1, True), (6, True)])
+def test_rank_and_stats_match_reference(rows, ties):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
+        kernels,
+    )
+
+    g = torch.Generator(device="cuda").manual_seed(rows)
+    dev = "cuda"
+    nxt = torch.randn(8, NE, device=dev, generator=g) * 3
+    nbias = (torch.randn(NE, device=dev, generator=g) * 0.5).bfloat16()
+    cur = torch.randn(8, NE, device=dev, generator=g) * 3
+    cbias = (torch.randn(NE, device=dev, generator=g) * 0.5).bfloat16()
+    if ties:
+        # 13 integer logits and no bias: many equal scores within and across
+        # rows (resolved by ID), distinct values far apart
+        nxt = nxt.round().clamp(-6, 6)
+        cur = cur.round().clamp(-6, 6)
+        nbias.zero_()
+        cbias.zero_()
+    cur_ids = torch.randperm(NE, device=dev, generator=g)[:KMAX].to(torch.int32)
+    out_ids = torch.zeros(KMAX, dtype=torch.int32, device=dev)
+    counters = torch.zeros(kernels.NUM_COUNTERS, dtype=torch.int64, device=dev)
+    kernels.rank_and_stats(nxt, nbias, out_ids, cur, cbias, cur_ids, counters, rows, True, True)
+    torch.cuda.synchronize()
+    assert out_ids.tolist() == _ref_rank(nxt, nbias, rows)
+    cnt, _ = _ref_union(cur, cbias, rows)
+    actual = cnt > 0
+    pos = {int(e): j for j, e in enumerate(cur_ids.tolist())}
+    exp = [1, int(actual.sum())]
+    for k in STAT_KS:
+        exp.append(sum(1 for e in range(NE) if actual[e] and pos.get(e, KMAX) < k))
+    assert counters.tolist() == exp
+
+
+def _exact_weight(g):
+    """Gate weights in multiples of 1/256: with hidden values in {-1, 0, 1} every
+    logit is exact in FP32 whatever the summation order."""
+    return (torch.randint(-2, 3, (NE, HIDDEN), device="cuda", generator=g) / 256).bfloat16()
+
+
+def _exact_hidden(rows, g):
+    return torch.randint(-1, 2, (rows, HIDDEN), device="cuda", generator=g).bfloat16()
+
+
+@cuda
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_gemm_rank_matches_reference(rows):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
+        kernels,
+    )
+
+    g = torch.Generator(device="cuda").manual_seed(100 + rows)
+    h = _exact_hidden(rows, g)
+    w = _exact_weight(g)
+    bias = (torch.randn(NE, device="cuda", generator=g) * 0.1).bfloat16()
+    scores = torch.zeros(kernels.MAX_ROWS, NE, device="cuda")
+    ticket = torch.zeros(1, dtype=torch.int32, device="cuda")
+    out = torch.zeros(KMAX, dtype=torch.int32, device="cuda")
+    ref = _ref_rank(h.float() @ w.float().t(), bias, rows)
+    for _ in range(3):  # every launch resets the ticket for the next one
+        out.zero_()
+        kernels.gemm_rank(h, w, bias, scores, ticket, out, rows)
+        torch.cuda.synchronize()
+        assert out.tolist() == ref
+        assert ticket.item() == 0
+    graph = torch.cuda.CUDAGraph()
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        with torch.cuda.graph(graph, stream=s):
+            kernels.gemm_rank(h, w, bias, scores, ticket, out, rows)
+    torch.cuda.current_stream().wait_stream(s)
+    out.zero_()
+    for _ in range(3):
+        graph.replay()
+    torch.cuda.synchronize()
+    assert out.tolist() == ref and ticket.item() == 0
+
+
+class _Gate(nn.Module):
+    def __init__(self, seed):
+        super().__init__()
+        g = torch.Generator(device="cuda").manual_seed(seed)
+        self.weight = nn.Parameter(_exact_weight(g), requires_grad=False)
+        self.e_score_correction_bias = nn.Parameter(
+            (torch.randn(NE, device="cuda", generator=g) * 0.1).bfloat16(), requires_grad=False
+        )
+
+    def forward(self, h):
+        return h.float() @ self.weight.float().t()
+
+
+class _Experts(nn.Module):
+    def __init__(self):
+        super().__init__()
+        u8 = dict(dtype=torch.uint8, device="cuda")
+        self.w3_w1_weight = nn.Parameter(
+            torch.zeros(NE, 512, HIDDEN // 2, **u8), requires_grad=False
+        )
+        self.w3_w1_weight_scale = nn.Parameter(
+            torch.zeros(NE, 512, HIDDEN // 16, **u8), requires_grad=False
+        )
+        self.w2_weight = nn.Parameter(torch.zeros(NE, HIDDEN, 128, **u8), requires_grad=False)
+        self.w2_weight_scale = nn.Parameter(torch.zeros(NE, HIDDEN, 16, **u8), requires_grad=False)
+
+
+class _Mlp(nn.Module):
+    def __init__(self, seed, experts):
+        super().__init__()
+        self.gate = _Gate(seed)
+        self.experts = experts
+
+
+class _Layer(nn.Module):
+    def __init__(self, mlp):
+        super().__init__()
+        self.mlp = mlp
+
+
+@cuda
+def test_link_layers_binds_routed_kernel_hooks(spec, monkeypatch):
+    """At the fixed depth, predictor, issuing kernel and CTA count, every MoE layer's routed
+    kernel prefetches the experts its own prediction names."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe,
+    )
+
+    monkeypatch.setattr(routed_moe, "resolve_expert_owner", lambda m: (m, ""))
+    monkeypatch.setattr(spec.dispatch, "is_enabled", lambda f: True)
+    spec._STATES.clear()
+    gates = [_Gate(i) for i in range(3)]
+    leaves = [nn.Module() for _ in range(3)]
+    layers = [nn.Identity()]
+    for gate, leaf in zip(gates, leaves):
+        mlp = nn.Module()
+        mlp.gate, mlp.experts = gate, leaf
+        layers.append(_Layer(mlp))
+    assert spec.link_layers(layers, 4) == 3
+    bound = [hasattr(leaf, "_expert_prefetch_spec") for leaf in leaves]
+    # Every layer's prediction feeds its own routed kernel.
+    assert bound == [True, True, True]
+    assert leaves[2]._expert_prefetch_spec() is None  # before the device state exists
+    rows = 6
+    h = torch.zeros(rows, HIDDEN, device="cuda", dtype=torch.bfloat16)
+    logits = torch.zeros(rows, NE, device="cuda", dtype=torch.float32)
+    for i in (1, 2, 3):
+        link = layers[i]._expert_prefetch_link
+        link.on_hop(h, logits)
+        link.after_moe()
+        link.after_post_moe()
+    spec.join()
+    torch.cuda.synchronize()
+    st = spec._STATES[torch.cuda.current_device()]
+    ids, k, ctas, self_store = leaves[2]._expert_prefetch_spec()
+    assert ids.data_ptr() == st.pred_ids[3].data_ptr() and (k, ctas) == (8, 32)
+    assert self_store
+    # the routed kernel issues the prefetch: no prefetch kernel is launched
+    assert st.prefetch_calls == 0
+    assert not st.forked and not st.counters.any()
+    # without routed_moe nothing is bound
+    spec._STATES.clear()
+    for leaf in leaves:
+        leaf.__dict__.pop("_expert_prefetch_spec", None)
+    monkeypatch.setattr(spec.dispatch, "is_enabled", lambda f: f != "routed_moe")
+    spec.link_layers(layers, 4)
+    assert not any(hasattr(leaf, "_expert_prefetch_spec") for leaf in leaves)
+    spec._STATES.clear()
+
+
+@cuda
+def test_declines_outside_decode_shape(spec):
+    spec._STATES.clear()
+    experts = _Experts()
+    layers = [_Layer(_Mlp(1, experts)), _Layer(_Mlp(2, experts))]
+    spec.link_layers(layers, 2)
+    link = layers[0]._expert_prefetch_link
+    h = torch.zeros(16, HIDDEN, device="cuda", dtype=torch.bfloat16)
+    link.on_hop(h, None)
+    link.after_moe()
+    link.after_post_moe()
+    spec.join()
+    assert not spec._STATES
+
+
+@cuda
+def test_delay_kernel_waits():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
+        kernels,
+    )
+
+    kernels.delay(1000)  # compile
+    torch.cuda.synchronize()
+    t0, t1 = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    t0.record()
+    kernels.delay(50_000)
+    t1.record()
+    torch.cuda.synchronize()
+    assert t0.elapsed_time(t1) >= 0.045  # ms
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py
new file mode 100644
index 000000000000..4070e5ca1c39
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py
@@ -0,0 +1,353 @@
+# 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.
+"""Package cross_layer_expert_prefetch's speculative expert prefetch in hop_routing's routed_moe fork.
+
+* Every routed-MoE runner class of the package tree (routed_moe's runner and every
+  subclass, e.g. hop_routing's fork) compiles the prefetch when its key carries a
+  cross_layer_expert_prefetch spec (``key[6]``): ``spec_k`` / ``spec_ctas`` / ``spec_self`` of the
+  kernel equal the hook's values, ``spec_k == 0`` without a hook, and the
+  capability line names the runner's package. A runner that dropped the spec
+  would compile the prefetch out silently (the call still passes the spec
+  row's address). Kernel compilation is stubbed: no GPU needed.
+* On an SM100-family GPU with real GLM-5.2 weights (``GLM_MOE_DSA_REAL_MOE_INPUTS_DIR``), the
+  fork's outputs with the prefetch equal the fork's outputs without it bit
+  for bit (raw tuple, finalized output; the FP32 accumulator within the FP32
+  reorder bound of the exact sum), for valid and out-of-range expert IDs, R
+  = 1 / 6 / 8 and K = 8 / 12 with all or a subset of issuing CTAs; the
+  self-prediction store writes the call's distinct experts ranked by row
+  count; a CUDA graph with the prefetch replays bit-identically.
+"""
+
+import importlib
+import os
+import pkgutil
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_ROOT = os.path.dirname(_HERE)
+sys.path.insert(0, os.path.join(_ROOT, "routed_moe"))
+sys.path.insert(0, os.path.join(_ROOT, "hop_routing"))
+
+KMAX = 32
+REORDER_REL = 2.0**-20
+_PKG = "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels"
+
+sm100 = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+
+# ---------------------------------------------------------------------------
+# Structural guard: the spec reaches every runner's kernel (no compilation)
+# ---------------------------------------------------------------------------
+def _runner_classes():
+    """routed_moe's routed-MoE runner and every subclass in the package tree."""
+    pkg = importlib.import_module(_PKG)
+    for mod in pkgutil.iter_modules(pkg.__path__):
+        if mod.ispkg:
+            try:
+                importlib.import_module(f"{_PKG}.{mod.name}.routed_moe")
+            except ModuleNotFoundError:
+                pass
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    out, stack = {RoutedMoe}, [RoutedMoe]
+    while stack:
+        for sub in stack.pop().__subclasses__():
+            if sub not in out:
+                out.add(sub)
+                stack.append(sub)
+    return sorted(out, key=lambda c: f"{c.__module__}.{c.__qualname__}")
+
+
+class _Owner:
+    """Stands in for the NVFP4 expert leaf: only what the compile path reads."""
+
+    def __init__(self):
+        for name in ("w3_w1_weight", "w3_w1_weight_scale", "w2_weight", "w2_weight_scale"):
+            setattr(self, name, torch.zeros(64, dtype=torch.uint8))
+
+
+class _Stream:
+    cuda_stream = 0
+
+
+def _compile_kernels(cls, monkeypatch, hook):
+    """The kernel objects ``cls._compiled`` constructs for an FP32-accumulator call (stubbed compile)."""
+    import cutlass.cute as cute
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe as routed_moe_mod,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    built = []
+
+    def fake_compile(fn, *args, **kwargs):
+        built.append(getattr(fn, "__self__", fn))
+        return object()
+
+    logged = []
+    monkeypatch.setattr(cute, "compile", fake_compile)
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda *a, **k: _Stream())
+    monkeypatch.setattr(
+        dispatch, "log_capability", lambda f, k, e, r="": logged.append((f, k, e)) or e
+    )
+    # 24 co-resident four-CTA clusters: a 96-CTA expert-cluster grid, as on the served GPU.
+    monkeypatch.setattr(routed_moe_mod, "_max_active_clusters", lambda device, cluster: 24)
+    # Each class keeps its own compile cache: start from empty ones.
+    monkeypatch.setattr(routed_moe_mod, "_COMPILED", {})
+    mod = sys.modules[cls.__module__]
+    if hasattr(mod, "_COMPILED"):
+        monkeypatch.setattr(mod, "_COMPILED", {})
+    owner = _Owner()
+    if hook is not None:
+        owner._expert_prefetch_spec = hook
+    runner = cls.__new__(cls)
+    runner.module = owner
+    runner.bias = torch.zeros(256, dtype=torch.bfloat16)
+    runner.device = torch.device("cpu")
+    runner.num_ctas = 152
+    key = runner._key(False, 0, 6, True)
+    runner._compiled(False, 0, 6, True)
+    kernels = [k for k in built if hasattr(k, "spec_k")]
+    assert len(kernels) == 1, f"{cls.__qualname__}: expected one expert-cluster kernel, got {built}"
+    return key, kernels[0], logged
+
+
+@pytest.mark.parametrize("spec", [(12, 0, True), (8, 32, True), (8, 40, False), (0, 0, True)])
+def test_every_runner_compiles_the_expert_prefetch_spec(monkeypatch, spec):
+    ids = torch.zeros(KMAX, dtype=torch.int32)
+    classes = _runner_classes()
+    assert any(c.__module__.endswith(".hop_routing.routed_moe") for c in classes), classes
+    for cls in classes:
+        key, kern, logged = _compile_kernels(cls, monkeypatch, lambda: (ids,) + spec)
+        assert key[6] == spec, (cls.__qualname__, key)
+        got = (kern.spec_k, kern.spec_ctas, kern.spec_self)
+        assert got == spec, (
+            f"{cls.__qualname__} dropped the cross_layer_expert_prefetch spec: kernel {got}, key {spec}"
+        )
+        if spec[0]:
+            pkg = cls.__module__.split(".")[-2]
+            line = f"expert_prefetch[k={spec[0]},ctas={spec[1]},self={int(spec[2])}]"
+            assert (pkg, line, True) in logged, (cls.__qualname__, logged)
+
+
+def test_runner_without_expert_prefetch_has_no_prefetch(monkeypatch):
+    for cls in _runner_classes():
+        key, kern, logged = _compile_kernels(cls, monkeypatch, None)
+        assert key[6] is None
+        assert (kern.spec_k, kern.spec_ctas, kern.spec_self) == (0, 0, False), cls.__qualname__
+        assert not [entry for entry in logged if entry[1].startswith("expert_prefetch")], logged
+
+
+def test_expert_prefetch_requires_routed_moe():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    assert dispatch.REQUIRED_FEATURES["cross_layer_expert_prefetch"] == frozenset({"routed_moe"})
+    auto = dispatch.ModelingV2Mode.AUTO
+    assert dispatch.resolve(auto, {}).enabled("cross_layer_expert_prefetch")
+    plan = dispatch.resolve(auto, {dispatch.DISABLE_ENV: "routed_moe"})
+    assert not plan.enabled("cross_layer_expert_prefetch")
+
+
+# ---------------------------------------------------------------------------
+# Real weights: the fork with and without the prefetch
+# ---------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def layer():
+    import routed_moe_fixtures as fx
+
+    torch.manual_seed(0)
+    return fx.build_native_layer(
+        layer=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10")),
+        tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")),
+    )
+
+
+@pytest.fixture(scope="module")
+def fork(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing.routed_moe import (  # noqa: E501
+        HopRoutedMoe,
+    )
+
+    r = HopRoutedMoe(layer.module, layer.bias)
+    if r.schedule(8, False, True) != "xc":
+        pytest.skip("the speculative prefetch is part of the expert-cluster schedule")
+    return r
+
+
+class _Spec:
+    """Install a cross_layer_expert_prefetch hook on the fork's expert leaf."""
+
+    def __init__(self, runner, ids, k, ctas=0, self_store=False):
+        self.owner = runner.module
+        self.spec = (ids, k, ctas, self_store)
+
+    def __enter__(self):
+        self.owner._expert_prefetch_spec = lambda: self.spec
+        return self
+
+    def __exit__(self, *exc):
+        del self.owner._expert_prefetch_spec
+
+
+def _pack(logits, bias):
+    import hop_routing_oracle as orc
+
+    ids, w = orc.native_routing(logits, bias)
+    routing = torch.empty(2, ids.shape[0], 8, dtype=torch.int32, device="cuda")
+    routing[0].copy_(ids)
+    routing[1].copy_(w.float().view(torch.int32))
+    return routing
+
+
+def _windows(layer, rows, count, offset=0):
+    import routed_moe_fixtures as fx
+
+    win = [
+        (x, lg)
+        for x, lg in fx.real_windows(fx.load_real_capture(layer.layer), rows)
+        if x is not None
+    ]
+    if not win:
+        pytest.skip(
+            f"no captured hidden states of layer {layer.layer} (GLM_MOE_DSA_REAL_MOE_INPUTS_DIR)"
+        )
+    step = max(1, len(win) // count)
+    out = []
+    for i in range(count):
+        x, lg = win[(offset + i * step) % len(win)]
+        x = x.cuda().contiguous()
+        lg = lg.float().cuda().contiguous()
+        out.append((x, _pack(lg, layer.bias)))
+    return out
+
+
+def _ids(kind, seed):
+    g = torch.Generator().manual_seed(seed)
+    if kind == "valid":
+        ids = torch.randperm(256, generator=g)[:KMAX]
+    else:  # the kernel masks each ID into [0, 256)
+        ids = torch.randint(-(2**31), 2**31 - 1, (KMAX,), generator=g)
+    return ids.to(torch.int32).cuda()
+
+
+def _outputs(fork, x, routing):
+    fin = fork(x, routing)
+    fc2, scale, _ = fork(x, routing, do_finalize=False)
+    acc = torch.zeros(x.shape[0], 6144, dtype=torch.float32, device="cuda")
+    fork(x, routing, red_out=acc)
+    torch.cuda.synchronize()
+    return fin.clone(), fc2.clone(), scale.clone(), acc
+
+
+def _check_same(ref, got, what):
+    fin, fc2, scale, _ = ref
+    gfin, gfc2, gscale, gacc = got
+    assert torch.equal(gfin, fin), f"{what}: finalized output differs"
+    assert torch.equal(gfc2, fc2), f"{what}: raw FC2 rows differ"
+    assert torch.equal(gscale, scale), f"{what}: routing weights differ"
+    rows = gacc.shape[0]
+    terms = scale.double()[:, :, None] * fc2.view(rows, 8, -1).double()
+    exact = terms.sum(1)
+    bound = REORDER_REL * terms.abs().sum(1)
+    assert ((gacc.double() - exact).abs() <= bound).all(), f"{what}: accumulator"
+
+
+@sm100
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("k,ctas", [(12, 0), (8, 32), (8, 0), (12, 40)])
+@pytest.mark.parametrize("ids_kind", ["valid", "out_of_range"])
+def test_fork_prefetch_outputs_bit_identical(layer, fork, rows, k, ctas, ids_kind):
+    for i, (x, routing) in enumerate(_windows(layer, rows, 2, offset=rows)):
+        ref = _outputs(fork, x, routing)
+        ids = _ids(ids_kind, seed=rows * 100 + i)
+        with _Spec(fork, ids, k, ctas):
+            got = _outputs(fork, x, routing)
+        _check_same(ref, got, f"rows={rows} k={k} ctas={ctas} {ids_kind} window {i}")
+
+
+def _ranked_experts(routing):
+    ids = routing[0].flatten().long().cpu()
+    counts = torch.bincount(ids, minlength=256)
+    distinct = [e for e in range(256) if counts[e] > 0]
+    return sorted(distinct, key=lambda e: (-int(counts[e]), e))
+
+
+@sm100
+@pytest.mark.parametrize("rows", [1, 3, 6, 8])
+@pytest.mark.parametrize("k,ctas", [(12, 0), (8, 32)])
+def test_fork_self_store_ranks_by_row_count(layer, fork, rows, k, ctas):
+    for i, (x, routing) in enumerate(_windows(layer, rows, 3, offset=5 * rows)):
+        ref = _outputs(fork, x, routing)
+        ids = torch.full((KMAX,), 999, dtype=torch.int32, device="cuda")
+        with _Spec(fork, ids, k, ctas, self_store=True):
+            acc = torch.zeros(rows, 6144, dtype=torch.float32, device="cuda")
+            fork(x, routing, red_out=acc)
+            torch.cuda.synchronize()
+            got = ids.tolist()
+            # the store is idempotent for the same call (the prefetch reads the
+            # previous call's row first)
+            fc2, scale, _ = fork(x, routing, do_finalize=False)
+            gfin = fork(x, routing)
+            torch.cuda.synchronize()
+        assert torch.equal(fc2, ref[1]) and torch.equal(scale, ref[2]), (rows, i)
+        assert torch.equal(gfin, ref[0]), (rows, i)
+        want = _ranked_experts(routing)[:KMAX]
+        assert got[: len(want)] == want, (rows, i, got, want)
+        assert all(v == 999 for v in got[len(want) :]), (rows, i, got)
+
+
+@sm100
+def test_fork_prefetch_cuda_graph_replay(layer, fork):
+    rows = 6
+    wins = _windows(layer, rows, 6, offset=3)
+    # References: the fork without the prefetch, eager.
+    refs = []
+    for x, routing in wins:
+        fc2, scale, _ = fork(x, routing, do_finalize=False)
+        torch.cuda.synchronize()
+        refs.append((fc2.clone(), scale.clone()))
+    x_static = wins[0][0].clone()
+    r_static = wins[0][1].clone()
+    ids = torch.zeros(KMAX, dtype=torch.int32, device="cuda")
+    with _Spec(fork, ids, 8, 32, self_store=True):
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            fork(x_static, r_static, do_finalize=False)
+        torch.cuda.current_stream().wait_stream(stream)
+        # the eager call compiled the prefetch variant: the capture records it
+        assert fork.ready_for_capture(False, rows=rows)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            fc2_s, scale_s, _ = fork(x_static, r_static, do_finalize=False)
+    for it, (x, routing) in enumerate(wins):
+        x_static.copy_(x)
+        r_static.copy_(routing)
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(fc2_s, refs[it][0]) and torch.equal(scale_s, refs[it][1]), it
+        # the replay stored this call's ranked experts for the next one
+        want = _ranked_experts(routing)[:KMAX]
+        assert ids.tolist()[: len(want)] == want, it
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py
new file mode 100644
index 000000000000..73142cbf4ea7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py
@@ -0,0 +1,352 @@
+# 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.
+"""The cross_layer_expert_prefetch expert prefetch in the shared_expert_prefetch fork.
+
+``cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py`` checks that every
+routed-MoE runner class of the tree, shared_expert_prefetch's included, compiles the prefetch from a
+key with a cross_layer_expert_prefetch spec. This file checks what is specific to
+shared_expert_prefetch's port:
+
+* the merged pre-wait prologue of ``SharedPrefetchExpertClusterKernel``: the shared-expert chunks
+  are spread over cross_layer_expert_prefetch's issuing CTAs (``spec_ctas``) when
+  cross_layer_expert_prefetch restricts them, over the whole grid otherwise
+  (shared_expert_prefetch's mapping without cross_layer_expert_prefetch), and the chunks cover the
+  shared expert's bytes;
+* ``SharedPrefetchRoutedMoe.__call__`` passes cross_layer_expert_prefetch's row of predicted experts
+  as the 18th routed_moe address (routed_moe's dummy without a hook) before the two shared-expert
+  addresses, i.e. ``_NUM_XC_ADDRS + 2`` addresses in the kernel's order.
+  Both need no GPU (kernel compilation and launch are stubbed);
+* on an SM100-family GPU with real GLM-5.2 weights (``GLM_MOE_DSA_REAL_MOE_INPUTS_DIR``,
+  ``GLM_MOE_DSA_ROUTED_MOE_CKPT_DIR``): shared_expert_prefetch's FP32 accumulator with the prefetch (K = 8 / 12, all
+  or 32 issuing CTAs, valid and out-of-range IDs) is the exact sum of routed_moe's weighted rows within the
+  FP32 reorder bound (routed_moe_into_hop's fp32 class, as shared_expert_prefetch without
+  cross_layer_expert_prefetch); the self-store writes the call's distinct experts ranked by row
+  count; a CUDA graph with the prefetch replays within the bound and keeps storing.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_ROOT = os.path.dirname(_HERE)
+for _sub in ("routed_moe", "hop_routing", "shared_expert_prefetch"):
+    sys.path.insert(0, os.path.join(_ROOT, _sub))
+
+KMAX = 32
+REORDER_REL = 2.0**-20
+LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
+
+sm100 = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+
+# ---------------------------------------------------------------------------
+# The merged prologue's issuers (no compilation)
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "num_ctas,spec_k,spec_ctas,want",
+    [
+        (96, 0, 0, 96),  # shared_expert_prefetch without cross_layer_expert_prefetch: every CTA
+        (96, 0, 32, 96),  # spec_ctas, prefetch not compiled in: shared_expert_prefetch's mapping
+        (96, 8, 0, 96),  # cross_layer_expert_prefetch on every CTA
+        (96, 8, 32, 96 // 3),  # the prefetch on the early-resident CTAs: the same CTAs
+        (96, 12, 40, 40),
+        (96, 8, 200, 96),
+        (144, 8, 56, 56),
+    ],
+)
+def test_prologue_issuers_and_coverage(num_ctas, spec_k, spec_ctas, want):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch import (  # noqa: E501
+        kernel_xc as zk,
+    )
+
+    kern = zk.SharedPrefetchExpertClusterKernel(
+        num_ctas, spec_k=spec_k, spec_ctas=spec_ctas, spec_self=True, red_out=True
+    )
+    assert kern.se_ctas == want
+    assert (kern.spec_k, kern.spec_ctas) == (spec_k, spec_ctas)
+    for chunk, total in ((kern.se_chunk1, zk.SE_W1_BYTES), (kern.se_chunk2, zk.SE_W2_BYTES)):
+        assert chunk % 16 == 0
+        lanes = kern.se_ctas * zk.PF_LANES
+        assert lanes * chunk >= total, (chunk, total)
+        # every issuing CTA owns a chunk (no CTA past se_ctas is needed)
+        assert (lanes - zk.PF_LANES) * chunk < total
+    if spec_k == 0:
+        assert (kern.se_chunk1, kern.se_chunk2) == (
+            zk._chunk(zk.SE_W1_BYTES, num_ctas),
+            zk._chunk(zk.SE_W2_BYTES, num_ctas),
+        )
+
+
+# ---------------------------------------------------------------------------
+# SharedPrefetchRoutedMoe.__call__ passes cross_layer_expert_prefetch's row (no launch)
+# ---------------------------------------------------------------------------
+class _Stream:
+    cuda_stream = 0
+
+
+class _Owner:
+    """The NVFP4 expert leaf: only what the call path reads."""
+
+    def __init__(self):
+        for name in (
+            "w3_w1_weight",
+            "w3_w1_weight_scale",
+            "w2_weight",
+            "w2_weight_scale",
+            "fc31_alpha",
+            "fc31_scale_c",
+            "fc2_alpha",
+            "fc31_input_scale",
+        ):
+            setattr(self, name, torch.zeros(64, dtype=torch.uint8))
+
+
+class _Ws:
+    def __init__(self):
+        for name in ("rows", "scales", "dummy", "tile_count"):
+            setattr(self, name, torch.zeros(16, dtype=torch.uint8))
+        self.num_ctas = 152
+
+    def order_on(self, stream):
+        pass
+
+
+@pytest.mark.parametrize("hooked", [True, False])
+def test_call_passes_the_expert_prefetch_row(monkeypatch, hooked):
+    import cutlass
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe as routed_moe_mod,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
+        SharedPrefetchRoutedMoe,
+    )
+
+    calls = []
+    ws = _Ws()
+    monkeypatch.setattr(routed_moe_mod, "_bulk_copy_ready", lambda x: x)
+    monkeypatch.setattr(routed_moe_mod, "_workspace", lambda device, n: ws)
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda *a, **k: _Stream())
+    # 24 co-resident four-CTA clusters: a 96-CTA expert-cluster grid, as on the served GPU.
+    monkeypatch.setattr(routed_moe_mod, "_max_active_clusters", lambda device, cluster: 24)
+    owner = _Owner()
+    ids = torch.zeros(KMAX, dtype=torch.int32)
+    if hooked:
+        owner._expert_prefetch_spec = lambda: (ids, 8, 32, True)
+    runner = SharedPrefetchRoutedMoe.__new__(SharedPrefetchRoutedMoe)
+    runner.module = owner
+    runner.bias = torch.zeros(256, dtype=torch.bfloat16)
+    runner.device = torch.device("cpu")
+    runner.num_ctas = 152
+    monkeypatch.setattr(runner, "_compiled", lambda *a, **k: (lambda *args: calls.append(args)))
+    x = torch.zeros(6, 16, dtype=torch.bfloat16)
+    routing = torch.zeros(2, 6, 8, dtype=torch.int32)
+    acc = torch.zeros(6, 16, dtype=torch.float32)
+    se = (torch.zeros(8, dtype=torch.bfloat16), torch.zeros(8, dtype=torch.bfloat16))
+    assert runner(x, routing, acc, se) == [acc]
+    assert len(calls) == 1
+    ints = [a for a in calls[0] if isinstance(a, cutlass.Int64)]
+    addrs = [int(a.value) if hasattr(a, "value") else int(a) for a in ints]
+    assert len(addrs) == routed_moe_mod._NUM_XC_ADDRS + 2 == 20
+    assert (
+        addrs[0] == x.data_ptr() and addrs[1] == routing.data_ptr() and addrs[10] == acc.data_ptr()
+    )
+    assert addrs[17] == (ids.data_ptr() if hooked else ws.dummy.data_ptr())
+    assert addrs[18:] == [se[0].data_ptr(), se[1].data_ptr()]
+
+
+# ---------------------------------------------------------------------------
+# Real weights: shared_expert_prefetch's fork with and without the prefetch
+# ---------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def layer():
+    import routed_moe_fixtures as fx
+
+    torch.manual_seed(0)
+    return fx.build_native_layer(
+        layer=LAYER, tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
+    )
+
+
+@pytest.fixture(scope="module")
+def runners(layer):
+    import shared_prefetch_fixtures as zf
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
+        SharedPrefetchRoutedMoe,
+        shared_weights,
+    )
+
+    mlp = zf.build_shared_mlp(LAYER, int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")))
+    se, reason = shared_weights(mlp)
+    assert se is not None, reason
+    shared_expert_prefetch = SharedPrefetchRoutedMoe(layer.module, layer.bias)
+    if shared_expert_prefetch.schedule(8, False, True) != "xc":
+        pytest.skip("the prefetch is part of the expert-cluster schedule")
+    return RoutedMoe.from_native(layer.module, layer.bias), shared_expert_prefetch, se, mlp
+
+
+class _Spec:
+    """Install a cross_layer_expert_prefetch hook on the fork's expert leaf."""
+
+    def __init__(self, runner, ids, k, ctas=0, self_store=False):
+        self.owner = runner.module
+        self.spec = (ids, k, ctas, self_store)
+
+    def __enter__(self):
+        self.owner._expert_prefetch_spec = lambda: self.spec
+        return self
+
+    def __exit__(self, *exc):
+        del self.owner._expert_prefetch_spec
+
+
+def _pack(logits, bias):
+    import hop_routing_oracle as orc
+
+    ids, w = orc.native_routing(logits, bias)
+    routing = torch.empty(2, ids.shape[0], 8, dtype=torch.int32, device="cuda")
+    routing[0].copy_(ids)
+    routing[1].copy_(w.float().view(torch.int32))
+    return routing
+
+
+def _windows(layer, rows, count, offset=0):
+    import routed_moe_fixtures as fx
+
+    win = [
+        (x, lg)
+        for x, lg in fx.real_windows(fx.load_real_capture(layer.layer), rows)
+        if x is not None
+    ]
+    if not win:
+        pytest.skip(
+            f"no captured hidden states of layer {layer.layer} (GLM_MOE_DSA_REAL_MOE_INPUTS_DIR)"
+        )
+    step = max(1, len(win) // count)
+    out = []
+    for i in range(count):
+        x, lg = win[(offset + i * step) % len(win)]
+        x = x.cuda().contiguous()
+        lg = lg.float().cuda().contiguous()
+        out.append((x, lg, _pack(lg, layer.bias)))
+    return out
+
+
+def _ids(kind, seed):
+    g = torch.Generator().manual_seed(seed)
+    if kind == "valid":
+        ids = torch.randperm(256, generator=g)[:KMAX]
+    else:  # the kernel masks each ID into [0, 256)
+        ids = torch.randint(-(2**31), 2**31 - 1, (KMAX,), generator=g)
+    return ids.to(torch.int32).cuda()
+
+
+def _exact(routed_runner, x, lg):
+    fc2, scale, idx = routed_runner(x, lg, do_finalize=False)
+    torch.cuda.synchronize()
+    terms = scale.double()[:, :, None] * fc2[idx.long()].double()
+    return terms.sum(1), REORDER_REL * terms.abs().sum(1)
+
+
+def _shared_prefetch_acc(shared_expert_prefetch, se, x, routing):
+    acc = torch.zeros(x.shape[0], 6144, dtype=torch.float32, device="cuda")
+    shared_expert_prefetch(x, routing, acc, se)
+    torch.cuda.synchronize()
+    return acc
+
+
+def _ranked_experts(routing):
+    ids = routing[0].flatten().long().cpu()
+    counts = torch.bincount(ids, minlength=256)
+    distinct = [e for e in range(256) if counts[e] > 0]
+    return sorted(distinct, key=lambda e: (-int(counts[e]), e))
+
+
+@sm100
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("k,ctas", [(8, 32), (12, 0), (8, 0), (12, 40)])
+@pytest.mark.parametrize("ids_kind", ["valid", "out_of_range"])
+def test_shared_prefetch_accumulator_within_bound(layer, runners, rows, k, ctas, ids_kind):
+    routed_runner, shared_expert_prefetch, se, _ = runners
+    for i, (x, lg, routing) in enumerate(_windows(layer, rows, 2, offset=rows)):
+        exact, bound = _exact(routed_runner, x, lg)
+        ref = _shared_prefetch_acc(shared_expert_prefetch, se, x, routing)
+        assert ((ref.double() - exact).abs() <= bound).all(), ("no prefetch", rows, i)
+        ids = _ids(ids_kind, seed=rows * 100 + i)
+        with _Spec(shared_expert_prefetch, ids, k, ctas):
+            got = _shared_prefetch_acc(shared_expert_prefetch, se, x, routing)
+        assert ((got.double() - exact).abs() <= bound).all(), (rows, k, ctas, ids_kind, i)
+
+
+@sm100
+@pytest.mark.parametrize("rows", [1, 3, 6, 8])
+@pytest.mark.parametrize("k,ctas", [(8, 32), (12, 0)])
+def test_shared_prefetch_self_store_ranks_by_row_count(layer, runners, rows, k, ctas):
+    routed_runner, shared_expert_prefetch, se, _ = runners
+    for i, (x, lg, routing) in enumerate(_windows(layer, rows, 3, offset=5 * rows)):
+        exact, bound = _exact(routed_runner, x, lg)
+        ids = torch.full((KMAX,), 999, dtype=torch.int32, device="cuda")
+        with _Spec(shared_expert_prefetch, ids, k, ctas, self_store=True):
+            acc = _shared_prefetch_acc(shared_expert_prefetch, se, x, routing)
+            got = ids.tolist()
+        assert ((acc.double() - exact).abs() <= bound).all(), (rows, i)
+        want = _ranked_experts(routing)[:KMAX]
+        assert got[: len(want)] == want, (rows, i, got, want)
+        assert all(v == 999 for v in got[len(want) :]), (rows, i, got)
+
+
+@sm100
+def test_shared_prefetch_cuda_graph_replay(layer, runners):
+    routed_runner, shared_expert_prefetch, se, _ = runners
+    rows = 6
+    wins = _windows(layer, rows, 6, offset=3)
+    refs = [_exact(routed_runner, x, lg) for x, lg, _ in wins]
+    x_static = wins[0][0].clone()
+    r_static = wins[0][2].clone()
+    acc = torch.zeros(rows, 6144, dtype=torch.float32, device="cuda")
+    ids = torch.zeros(KMAX, dtype=torch.int32, device="cuda")
+    with _Spec(shared_expert_prefetch, ids, 8, 32, self_store=True):
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            shared_expert_prefetch(x_static, r_static, acc, se)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        # the eager call compiled the prefetch variant: the capture records it
+        assert shared_expert_prefetch.ready_for_capture(False, rows=rows, red=True)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            shared_expert_prefetch(x_static, r_static, acc, se)
+    for it, (x, _, routing) in enumerate(wins):
+        x_static.copy_(x)
+        r_static.copy_(routing)
+        acc.zero_()
+        graph.replay()
+        torch.cuda.synchronize()
+        exact, bound = refs[it]
+        assert ((acc.double() - exact).abs() <= bound).all(), it
+        want = _ranked_experts(routing)[:KMAX]
+        assert ids.tolist()[: len(want)] == want, it
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py
new file mode 100644
index 000000000000..d04624e97387
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py
@@ -0,0 +1,343 @@
+# 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.
+"""deferred_norm with routed_moe_into_hop through the target's decoder layer (8 ranks over 2 hosts).
+
+Runs the target's ``GlmDecoderLayer.forward_MoE`` as a decode step on a decoder-layer
+stand-in that carries a real MoE of layer 10 (``Deepseekv3MoE`` adopted by ``GlmMoE``;
+routed_moe's fixtures: real NVFP4 experts and router, captured real hidden-state windows),
+the layer's post_moe_hop with its routed_moe_into_hop routed input enabled, and the
+deferred_norm hop with its routed input enabled. Features
+``absorbed_query,post_moe_hop,routed_moe,routed_moe_into_hop,deferred_norm``: routed_moe writes the
+routed result into the post_moe_hop and ``forward_MoE`` takes the routed_moe_into_hop branch.
+
+Checks, per routed_moe_into_hop mode:
+
+* the routed_moe_into_hop branch returns a deferred_norm packet (the deferred_norm routed hop
+  engaged);
+* the packet's materialization and the residual equal
+  ``PostMoeHop.hop_routed`` on the routed content routed_moe wrote for that call
+  (captured before the hop consumed it) bit for bit;
+* without the deferred_norm hop (as with ``deferred_norm`` off) the same branch returns
+  post_moe_hop's normalized tensor, bit-identical to the materialized packet in the exact
+  routed_moe_into_hop modes (routed_moe's ``fp32`` mode accumulates with ``red.add``, so two
+  routed_moe calls can differ by FP32 reordering; there the two runs are compared
+  within routed_moe_into_hop's reorder class only through the captured-input check above);
+* a context-phase batch keeps post_moe_hop's ``hop_routed``;
+* a CUDA graph of the branch (multi-stream capture as served) replays with
+  fresh windows and matches the eager deferred_norm path.
+
+The decoder layer's adoption wiring is checked on the target's source: the routed enable is
+called only with ``routed_moe_into_hop`` selected and the post_moe_hop's routed input set,
+after the deferred_norm hop is built.
+"""
+
+import inspect
+import os
+import sys
+import types
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "..", "..", "shared", "kernels", "post_moe_hop"))
+sys.path.insert(0, os.path.join(_HERE, "..", "routed_moe"))
+sys.path.insert(0, os.path.join(_HERE, "..", "..", "..", "shared", "kernels"))
+
+import plan_testlib  # noqa: E402
+import post_moe_hop_mpi_utils as M  # noqa: E402
+
+LAYER = 10
+FEATURES = "absorbed_query,post_moe_hop,routed_moe,routed_moe_into_hop,deferred_norm"
+MODES = ("fp32", "identity", "row")
+
+
+def _set_features(value: str) -> None:
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    plan_testlib.select(dispatch, value)
+
+
+@pytest.fixture(autouse=True)
+def decode_step():
+    """Every test runs as a decode step, where the target's decoder layer takes its fast paths."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    with dispatch.step(dispatch.Phase.DECODE):
+        yield
+
+
+@pytest.fixture(scope="module")
+def env():
+    e = M.mpi_env()
+    if e.world != 8:
+        pytest.skip("needs the TP8 group")
+    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
+    return e
+
+
+@pytest.fixture(scope="module")
+def layer(env):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    if not hasattr(PostMoeHop, "enable_routed_input"):
+        pytest.skip("the installed post_moe_hop hop has no routed_moe_into_hop routed input")
+    import routed_moe_chain_bench as CB
+    import routed_moe_fixtures as fx
+
+    _set_features("")
+    moes, gammas = CB.build_layers(env, [LAYER])
+    moe = moes[LAYER]
+    hop = PostMoeHop(env.mapping, fx.HIDDEN, fx.TOP_K, torch.bfloat16, raw_producer=moe.experts)
+    assert M.all_ranks(hop.engaged), hop.reason
+    win = [x for x, _ in fx.real_windows(fx.load_real_capture(LAYER), 6)]
+    step = max(1, len(win) // 8)
+    windows = [win[i * step].cuda() for i in range(8)]
+    yield types.SimpleNamespace(moe=moe, gamma=gammas[LAYER], hop=hop, windows=windows)
+    _set_features("")
+
+
+def _stub(layer, deferred_norm, gamma):
+    """The decoder-layer attributes ``forward_MoE`` reads on the post_moe_hop path."""
+    return types.SimpleNamespace(
+        _post_attn_router_weight=lambda: None,
+        _into_hop=None,
+        fusion_config=types.SimpleNamespace(PRE_MOE_FUSION=False, POST_MOE_FUSION=True),
+        post_attention_layernorm=lambda h, r: (h, r),
+        allreduce=None,
+        moe_allreduce=None,
+        model_config=None,
+        is_p2p_supported=False,
+        mapping=M.mpi_env().mapping,
+        mlp=layer.moe,
+        post_moe_hop=layer.hop,
+        deferred_norm_hop=deferred_norm,
+        next_layer_layernorm=types.SimpleNamespace(weight=gamma, variance_epsilon=M.GLM_EPS),
+    )
+
+
+def _metadata(rows, context=False):
+    return types.SimpleNamespace(
+        all_rank_num_tokens=None,
+        num_contexts=1 if context else 0,
+        num_generations=0 if context else rows,
+    )
+
+
+def _forward(stub, x, residual, md):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmDecoderLayer,
+    )
+
+    return GlmDecoderLayer.forward_MoE(stub, x, md, residual)
+
+
+def _enable(layer, mode):
+    """Collective: routed_moe_into_hop routed input in ``mode`` and a deferred_norm hop serving it."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+
+    assert M.all_ranks(layer.hop.enable_routed_input(mode))
+    layer.moe._routed_into_hop = layer.hop
+    key = f"_deferred_norm_{id(layer.hop)}"
+    deferred_norm = getattr(layer, key, None)
+    if deferred_norm is None:
+        deferred_norm = DeferredNormHop(layer.hop)
+        setattr(layer, key, deferred_norm)
+    assert M.all_ranks(deferred_norm.engaged and deferred_norm.enable_routed_input()), (
+        deferred_norm.routed_reason
+    )
+    return deferred_norm
+
+
+class _Recorder:
+    """Wraps ``DeferredNormHop.hop_routed`` to keep a copy of the routed input."""
+
+    def __init__(self, deferred_norm):
+        self.deferred_norm = deferred_norm
+        self.inner = deferred_norm.hop_routed
+        self.calls = []
+
+    def __call__(self, routed, shared, residual, gamma, eps):
+        mode, _, scales = self.deferred_norm._routed_view(residual.shape[0])
+        self.calls.append(
+            (
+                routed.clone(),
+                None if scales is None else scales.clone(),
+                None if shared is None else shared.clone(),
+                residual.clone(),
+            )
+        )
+        return self.inner(routed, shared, residual, gamma, eps)
+
+
+def _replay_post_moe_hop(hop, mode, routed, scales, shared, residual, gamma):
+    rows = residual.shape[0]
+    if mode == "fp32":
+        buf = hop.routed_acc_buffer(rows)
+    elif mode == "identity":
+        buf, s = hop.routed_raw_buffers(rows)
+        s.copy_(scales)
+    else:
+        buf = hop.routed_input_buffer(rows)
+    buf.copy_(routed)
+    return hop.hop_routed(buf, shared, residual, gamma, M.GLM_EPS)
+
+
+def test_adoption_wiring_in_target_source(env, layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmDecoderLayer,
+    )
+
+    init = inspect.getsource(GlmDecoderLayer._adopt_post_moe)
+    body = inspect.getsource(GlmDecoderLayer.forward_MoE)
+    i_routed_into_hop = init.find("self.mlp._routed_into_hop = self.post_moe_hop")
+    i_deferred_norm = init.find('"deferred_hop",')
+    i_en = init.find("self.deferred_norm_hop.enable_routed_input()")
+    ok = 0 <= i_routed_into_hop < i_deferred_norm < i_en
+    ok &= 'dispatch.is_enabled("routed_moe_into_hop")' in init[i_deferred_norm:i_en]
+    ok &= "_routed_into_hop" in init[i_deferred_norm:i_en]
+    i_branch = body.find("elif use_post_moe_hop and len(hidden_states) == 2:")
+    i_raw = body.find("elif use_post_moe_hop:", i_branch + 1)
+    ok &= 0 <= i_branch < body.find("can_run_routed", i_branch) < i_raw
+    ok &= body.find("hop.hop_routed(", i_branch) < i_raw
+    assert ok
+
+
+@pytest.mark.parametrize("mode", MODES)
+def test_routed_into_hop_branch_returns_packet_equal_to_hop_routed(env, layer, mode):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        is_deferred,
+    )
+
+    deferred_norm = _enable(layer, mode)
+    rec = _Recorder(deferred_norm)
+    deferred_norm.hop_routed = rec
+    _set_features(FEATURES)
+    g = torch.Generator(device="cuda").manual_seed(300)
+    ok, n = True, 0
+    try:
+        with torch.inference_mode():
+            for i, x in enumerate(layer.windows):
+                rows = x.shape[0]
+                residual = (torch.randn(rows, 6144, device="cuda", generator=g) * 4.0).bfloat16()
+                out, res = _forward(
+                    _stub(layer, deferred_norm, layer.gamma), x, residual, _metadata(rows)
+                )
+                engaged = is_deferred(out) and len(rec.calls) == i + 1
+                ok &= engaged
+                if not engaged:
+                    print(
+                        f"GLM_MOE_DSA_DEFERRED_NORM_WIRING_NOT_ENGAGED mode={mode} call={i} rank={env.rank}",
+                        flush=True,
+                    )
+                    continue
+                h = out.materialize()
+                routed, scales, shared, res_in = rec.calls[-1]
+                h_ref, r_ref = _replay_post_moe_hop(
+                    layer.hop, mode, routed, scales, shared, res_in, layer.gamma
+                )
+                torch.cuda.synchronize()
+                same = torch.equal(h, h_ref) and torch.equal(res, r_ref)
+                ok &= same
+                # deferred_norm off: the same branch returns post_moe_hop's normalized tensor.
+                out_off, res_off = _forward(
+                    _stub(layer, None, layer.gamma), x, residual, _metadata(rows)
+                )
+                torch.cuda.synchronize()
+                ok &= isinstance(out_off, torch.Tensor) and not is_deferred(out_off)
+                if mode != "fp32":
+                    ok &= torch.equal(out_off, h) and torch.equal(res_off, res)
+                n += 1
+    finally:
+        del deferred_norm.hop_routed
+        _set_features("")
+    stats = env.comm.allgather((env.rank, n, ok))
+    if env.rank == 0:
+        print(
+            f"GLM_MOE_DSA_DEFERRED_NORM_WIRING mode={mode}: (rank, windows, ok) {stats}", flush=True
+        )
+    assert M.all_ranks(ok and n == len(layer.windows))
+
+
+def test_context_batch_keeps_hop_routed(env, layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        is_deferred,
+    )
+
+    deferred_norm = _enable(layer, "identity")
+    _set_features(FEATURES)
+    try:
+        with torch.inference_mode():
+            x = layer.windows[0]
+            residual = torch.zeros(x.shape[0], 6144, dtype=torch.bfloat16, device="cuda")
+            out, _ = _forward(
+                _stub(layer, deferred_norm, layer.gamma),
+                x,
+                residual,
+                _metadata(x.shape[0], context=True),
+            )
+            ok = isinstance(out, torch.Tensor) and not is_deferred(out)
+    finally:
+        _set_features("")
+    assert M.all_ranks(ok)
+
+
+def test_routed_into_hop_branch_cuda_graph_replay(env, layer):
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    mode = "identity"
+    deferred_norm = _enable(layer, mode)
+    _set_features(FEATURES)
+    try:
+        with torch.inference_mode():
+            rows = layer.windows[0].shape[0]
+            x = layer.windows[0].clone()
+            g = torch.Generator(device="cuda").manual_seed(7)
+            residual = (torch.randn(rows, 6144, device="cuda", generator=g) * 4.0).bfloat16()
+            md = _metadata(rows)
+            stub = _stub(layer, deferred_norm, layer.gamma)
+            stream = torch.cuda.Stream()
+            stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(stream):
+                for _ in range(2):
+                    _forward(stub, x, residual, md)
+            torch.cuda.current_stream().wait_stream(stream)
+            torch.cuda.synchronize()
+            env.comm.barrier()
+            graph = torch.cuda.CUDAGraph()
+            with with_multi_stream(True):
+                with torch.cuda.graph(graph, stream=stream):
+                    g_out, g_res = _forward(stub, x, residual, md)
+            torch.cuda.synchronize()
+            ok = True
+            for w in layer.windows[1:5]:
+                x.copy_(w)
+                env.comm.barrier()
+                graph.replay()
+                torch.cuda.synchronize()
+                h_graph = g_out.materialize()
+                e_out, e_res = _forward(stub, x, residual, md)
+                h_eager = e_out.materialize()
+                torch.cuda.synchronize()
+                ok &= torch.equal(h_graph, h_eager) and torch.equal(g_res, e_res)
+    finally:
+        _set_features("")
+    assert M.all_ranks(ok)
+    layer.hop.raise_if_error()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py
new file mode 100644
index 000000000000..2da3c9dcc750
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py
@@ -0,0 +1,270 @@
+# 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.
+"""Gate C1(b, c): per-row R invariance of the MTP layer's MoE on real weights (one GPU, one TP8 rank shard).
+
+draft_live_rows runs the MTP layer's routed MoE, shared expert and their add on the live
+rows only. That is bit-identical to the served step 0 only if each of them
+computes a row the same way whatever the other rows of the call are. Checked
+here, bitwise, on rank ``GLM_MOE_DSA_DRAFT_LIVE_ROWS_TP_RANK``'s shard of ``model.layers.78``:
+
+* routed MoE, native FlashInfer TRT-LLM-Gen BF16 leaf and draft_routed_moe (each with its
+  own launch configuration per row count): row g of the R = 6 call equals the
+  R = 1 call on row g, for every g, with the six rows chosen to cover
+  D = 8..48 distinct experts (Gaussian rows and real token-embedding rows);
+  rows (g1, g2) of an R = 8 call equal the R = 2 call on them (B = 2 at K = 3);
+* shared expert (the served BF16 ``GatedMLP``, TP8 shard, real weights): row g
+  at M = 6 equals M = 1, rows (g1, g2) at M = 8 equal M = 2;
+* the same under CUDA-graph replay (graphs captured at R = 6 / R = 1, replayed
+  with new rows copied into their static inputs);
+* negative control: row g of the R = 6 call against the R = 1 call on another
+  row differs.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "draft_routed_moe"))
+import draft_moe_fixtures as fx  # noqa: E402
+
+TP_RANK = int(os.environ.get("GLM_MOE_DSA_DRAFT_LIVE_ROWS_TP_RANK", "0"))
+D_TARGETS = (8, 12, 16, 24, 31, 40, 48)
+
+
+def _bits(t):
+    return t.contiguous().view(torch.int16)
+
+
+def _eq(a, b):
+    return torch.equal(_bits(a), _bits(b))
+
+
+@pytest.fixture(scope="module")
+def layer():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("needs an SM100-family GPU")
+    import tensorrt_llm  # noqa: F401  (loads th_common: draft_routed_moe and the native MoE ops)
+
+    if not hasattr(torch.ops.trtllm, "glm_moe_dsa_draft_routed_moe"):
+        pytest.skip("draft_routed_moe op not in this build")
+    return fx.build_native(TP_RANK)
+
+
+@pytest.fixture(scope="module")
+def pools(layer):
+    gauss = fx.RowPool.gaussian(layer, num=1024, seed=10_000)
+    g = torch.Generator().manual_seed(77)
+    ids = torch.randint(0, 100_000, (1024,), generator=g).tolist()
+    xs = fx.embedding_rows(layer, ids)
+    emb = fx.RowPool(xs, fx.top8_sets(layer.bias, layer.logits(xs)))
+    return {"gauss": gauss, "emb": emb}
+
+
+def _routed(kind, layer, x, logits):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    if kind == "draft_routed_moe":
+        return rm.run(layer.module, layer.bias, x, logits)
+    return layer.native(x, logits)
+
+
+@pytest.mark.parametrize("kind", ["native", "draft_routed_moe"])
+@pytest.mark.parametrize("pool", ["gauss", "emb"])
+def test_routed_row_invariance_r6(layer, pools, kind, pool):
+    bad, cases, negatives = [], 0, 0
+    for t, target in enumerate(D_TARGETS):
+        for rep in range(3):
+            x6, d = pools[pool].take(6, target, seed=1000 * t + rep)
+            l6 = layer.logits(x6)
+            out6 = _routed(kind, layer, x6, l6)
+            for g in range(6):
+                out1 = _routed(kind, layer, x6[g : g + 1].contiguous(), l6[g : g + 1].contiguous())
+                cases += 1
+                if not _eq(out6[g : g + 1], out1):
+                    bad.append((target, d, rep, g))
+            # Negative control: another (different) row's R = 1 output differs from row 0's.
+            others = [h for h in range(1, 6) if not _eq(x6[h], x6[0])]
+            if others:
+                h = others[0]
+                out1o = _routed(kind, layer, x6[h : h + 1].contiguous(), l6[h : h + 1].contiguous())
+                assert not _eq(out6[0:1], out1o), "negative control not caught"
+                negatives += 1
+    torch.cuda.synchronize()
+    assert negatives >= len(D_TARGETS), negatives
+    print(
+        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE routed {kind} {pool} R=6 vs R=1 cases={cases} mismatching={len(bad)} "
+        f"negative_controls={negatives} {bad[:4]}",
+        flush=True,
+    )
+    assert not bad
+
+
+@pytest.mark.parametrize("kind", ["native", "draft_routed_moe"])
+def test_routed_pairs_r8(layer, pools, kind):
+    bad, cases = [], 0
+    for t, target in enumerate((16, 31, 48, 64)):
+        x8, d = pools["gauss"].take(8, target, seed=5000 + t)
+        l8 = layer.logits(x8)
+        out8 = _routed(kind, layer, x8, l8)
+        for g1, g2 in ((0, 7), (7, 0), (3, 4), (5, 2), (1, 6)):
+            idx = torch.tensor([g1, g2], device="cuda")
+            out2 = _routed(kind, layer, x8[idx].contiguous(), l8[idx].contiguous())
+            cases += 1
+            if not _eq(out8[idx], out2):
+                bad.append((target, d, g1, g2))
+    torch.cuda.synchronize()
+    print(
+        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE routed {kind} R=8 vs R=2 cases={cases} mismatching={len(bad)} {bad[:4]}",
+        flush=True,
+    )
+    assert not bad
+
+
+def _shared_expert(layer):
+    """The served BF16 shared expert of the MTP layer, TP rank shard, real weights (Deepseekv3MoE's construction)."""
+    import copy
+    import math
+
+    sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
+    import routed_moe_fixtures as wfx
+
+    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+    from tensorrt_llm.quantization.mode import QuantAlgo
+
+    mc = wfx._served_model_config(fx.MTP_LAYER, TP_RANK)
+    mc.quant_config.exclude_modules = [f"model.layers.{fx.MTP_LAYER}*"]
+    shared_mc = copy.copy(mc)
+    shared_mc.quant_config = QuantConfig(quant_algo=None, kv_cache_quant_algo=QuantAlgo.FP8)
+    tp = math.gcd(fx.INTER // 1, mc.mapping.tp_size)
+    se = GatedMLP(
+        hidden_size=fx.HIDDEN,
+        intermediate_size=fx.INTER,
+        bias=False,
+        dtype=torch.bfloat16,
+        config=shared_mc,
+        overridden_tp_size=tp,
+        reduce_output=False,
+    )
+    rd = fx.Reader()
+    pre = f"model.layers.{fx.MTP_LAYER}.mlp.shared_experts"
+    full = {p: rd.full(f"{pre}.{p}.weight") for p in ("gate_proj", "up_proj", "down_proj")}
+    if hasattr(se, "gate_up_proj"):
+        se.gate_up_proj.load_weights([{"weight": full["gate_proj"]}, {"weight": full["up_proj"]}])
+    else:
+        se.gate_proj.load_weights([{"weight": full["gate_proj"]}])
+        se.up_proj.load_weights([{"weight": full["up_proj"]}])
+    se.down_proj.load_weights([{"weight": full["down_proj"]}])
+    wfx._walk_post_load(se)
+    return se.cuda(), tp
+
+
+@pytest.fixture(scope="module")
+def shared(layer):
+    se, tp = _shared_expert(layer)
+    print(
+        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE shared expert TP{tp} rank {TP_RANK} {type(se).__name__}",
+        flush=True,
+    )
+    return se
+
+
+def test_shared_row_invariance(layer, pools, shared):
+    bad, cases = [], 0
+    with torch.inference_mode():
+        for seed in range(12):
+            x6, _ = pools["emb" if seed % 2 else "gauss"].take(6, 8 + 4 * seed, seed=seed)
+            out6 = shared(x6)
+            for g in range(6):
+                cases += 1
+                if not _eq(out6[g : g + 1], shared(x6[g : g + 1].contiguous())):
+                    bad.append((seed, g))
+            x8, _ = pools["gauss"].take(8, 40, seed=100 + seed)
+            out8 = shared(x8)
+            for g1, g2 in ((0, 7), (6, 1), (3, 4)):
+                idx = torch.tensor([g1, g2], device="cuda")
+                cases += 1
+                if not _eq(out8[idx], shared(x8[idx].contiguous())):
+                    bad.append((seed, g1, g2))
+        other = shared(x6[1:2].contiguous())
+        assert not _eq(out6[0:1], other) or _eq(x6[0], x6[1]), "negative control not caught"
+    torch.cuda.synchronize()
+    print(
+        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE shared M=6/8 vs M=1/2 cases={cases} mismatching={len(bad)} {bad[:4]}",
+        flush=True,
+    )
+    assert not bad
+
+
+def _capture(fn, *static):
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        fn(*static)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        out = fn(*static)
+    torch.cuda.synchronize()
+    return graph, out
+
+
+@pytest.mark.parametrize("kind", ["native", "draft_routed_moe"])
+def test_graph_replay(layer, pools, shared, kind):
+    """R = 6 and R = 1 graphs of the routed MoE and of the shared expert, replayed with new rows."""
+    x6s = torch.empty(6, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
+    l6s = torch.empty(6, fx.NUM_EXPERTS, dtype=torch.float32, device="cuda")
+    x1s = torch.empty(1, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
+    l1s = torch.empty(1, fx.NUM_EXPERTS, dtype=torch.float32, device="cuda")
+    x0, _ = pools["gauss"].take(6, 31, seed=9)
+    x6s.copy_(x0)
+    l6s.copy_(layer.logits(x0))
+    x1s.copy_(x0[:1])
+    l1s.copy_(layer.logits(x0[:1]))
+    with torch.inference_mode():
+        g6, o6 = _capture(lambda x, lg: _routed(kind, layer, x, lg), x6s, l6s)
+        g1, o1 = _capture(lambda x, lg: _routed(kind, layer, x, lg), x1s, l1s)
+        s6, so6 = _capture(shared, x6s)
+        s1, so1 = _capture(shared, x1s)
+        bad, cases = [], 0
+        for trial, target in enumerate(D_TARGETS):
+            x, _ = pools["emb" if trial % 2 else "gauss"].take(6, target, seed=300 + trial)
+            lg = layer.logits(x)
+            x6s.copy_(x)
+            l6s.copy_(lg)
+            g6.replay()
+            s6.replay()
+            r6, sh6 = o6.clone(), so6.clone()
+            for g in range(6):
+                x1s.copy_(x[g : g + 1])
+                l1s.copy_(lg[g : g + 1])
+                g1.replay()
+                s1.replay()
+                cases += 1
+                if not (_eq(r6[g : g + 1], o1) and _eq(sh6[g : g + 1], so1)):
+                    bad.append((target, g))
+        torch.cuda.synchronize()
+    print(
+        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE graph {kind} routed+shared R=6 vs R=1 cases={cases} "
+        f"mismatching={len(bad)} {bad[:4]}",
+        flush=True,
+    )
+    assert not bad
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py
new file mode 100644
index 000000000000..dae482dcc4ac
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py
@@ -0,0 +1,315 @@
+# 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.
+"""Real-weight fixtures and the native oracle for the draft_routed_moe tests.
+
+The fixture is the MTP layer (``model.layers.78``, BF16 in the GLM-5.2 NVFP4
+checkpoint) as TP rank ``r`` of TP8/EP1 sees it: every routed expert's
+gate/up rows ``[256r, 256r+256)`` and down columns ``[256r, 256r+256)``, the
+router weight and the BF16 router bias. The native module is the stock
+FlashInfer TRT-LLM-Gen BF16 leaf (``trtllm_gen_leaf(None)``) loaded through its
+own ``load_weights``/``post_load_weights``, so its buffers carry exactly the
+shuffled BlockMajorK layouts of a served model. Slicing the checkpoint before
+loading a ``tp_size=1`` leaf is what the TP8 loader's ``load_weight_shard``
+does (COLUMN for w1/w3, ROW for w2).
+"""
+
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+import torch
+from utils.llm_data import llm_models_root
+
+CKPT_DIR = os.environ.get(
+    "GLM_MOE_DSA_DRAFT_ROUTED_MOE_CKPT_DIR",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+CACHE_DIR = os.environ.get(
+    "GLM_MOE_DSA_DRAFT_ROUTED_MOE_FIXTURE_CACHE",
+    os.path.join(tempfile.gettempdir(), "glm_moe_dsa_draft_routed_moe_fixtures"),
+)
+
+MTP_LAYER = 78
+HIDDEN = 6144
+INTER = 2048
+NUM_EXPERTS = 256
+TOP_K = 8
+TP_SIZE = 8
+ROUTED_SCALING = 2.5
+RMS_EPS = 1e-5
+
+
+class Reader:
+    """Opens each safetensors shard once and slices tensors lazily."""
+
+    def __init__(self):
+        from safetensors import safe_open
+
+        self._safe_open = safe_open
+        with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
+            self._map = json.load(f)["weight_map"]
+        self._open = {}
+
+    def _handle(self, name: str):
+        shard = self._map[name]
+        if shard not in self._open:
+            self._open[shard] = self._safe_open(
+                os.path.join(CKPT_DIR, shard), framework="pt", device="cpu"
+            )
+        return self._open[shard]
+
+    def full(self, name: str) -> torch.Tensor:
+        return self._handle(name).get_tensor(name)
+
+    def rows(self, name: str, start: int, stop: int) -> torch.Tensor:
+        return self._handle(name).get_slice(name)[start:stop].clone()
+
+    def cols(self, name: str, start: int, stop: int) -> torch.Tensor:
+        return self._handle(name).get_slice(name)[:, start:stop].clone()
+
+
+def load_rank(tp_rank: int, layer: int = MTP_LAYER) -> Dict[str, torch.Tensor]:
+    """Rank-sliced raw tensors of the MTP layer's MoE (CPU, cached under ``CACHE_DIR``)."""
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    path = os.path.join(CACHE_DIR, f"glm52_L{layer}_bf16_tp{TP_SIZE}r{tp_rank}.pt")
+    if os.path.exists(path):
+        return torch.load(path, map_location="cpu")
+    rd = Reader()
+    inter_rank = INTER // TP_SIZE
+    r0, r1 = tp_rank * inter_rank, (tp_rank + 1) * inter_rank
+    pre = f"model.layers.{layer}.mlp"
+    out: Dict[str, torch.Tensor] = {}
+    for e in range(NUM_EXPERTS):
+        out[f"{e}.w1.weight"] = rd.rows(f"{pre}.experts.{e}.gate_proj.weight", r0, r1)
+        out[f"{e}.w3.weight"] = rd.rows(f"{pre}.experts.{e}.up_proj.weight", r0, r1)
+        out[f"{e}.w2.weight"] = rd.cols(f"{pre}.experts.{e}.down_proj.weight", r0, r1)
+    out["gate.weight"] = rd.full(f"{pre}.gate.weight")
+    out["gate.e_score_correction_bias"] = rd.full(f"{pre}.gate.e_score_correction_bias")
+    out["post_attention_layernorm.weight"] = rd.full(
+        f"model.layers.{layer}.post_attention_layernorm.weight"
+    )
+    tmp = path + f".tmp{os.getpid()}"
+    torch.save(out, tmp)
+    os.replace(tmp, path)
+    return out
+
+
+@dataclass
+class NativeBf16Layer:
+    """The native FlashInfer TRT-LLM-Gen BF16 MoE leaf of one TP rank plus its router."""
+
+    module: torch.nn.Module
+    gate_weight: torch.Tensor  # [256, 6144] bf16
+    bias: torch.Tensor  # [256] bf16 (the TRTLLM gate dtype)
+    norm_weight: torch.Tensor  # [6144] bf16
+    tp_rank: int
+
+    def logits(self, x: torch.Tensor) -> torch.Tensor:
+        """FP32 router logits of ``x`` (FP32 GEMM; any FP32 logits are valid oracle inputs)."""
+        return (x.float() @ self.gate_weight.float().t()).contiguous()
+
+    def native(self, x: torch.Tensor, logits: torch.Tensor, do_finalize: bool = True):
+        """``trtllm_bf16_moe`` exactly as ``FlashinferTrtllmGenBf16Impl.run_moe`` calls it."""
+        from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.kernel_inputs import (
+            to_trtllm_gen_act_type,
+        )
+
+        m = self.module
+        return m.op_backend.run_bf16_moe(
+            router_logits=logits,
+            routing_bias=self.bias,
+            hidden_states=x,
+            gemm1_weights=m.w3_w1_weight,
+            gemm2_weights=m.w2_weight,
+            num_experts=m.num_slots,
+            top_k=TOP_K,
+            n_group=1,
+            topk_group=1,
+            intermediate_size=m.intermediate_size_per_partition,
+            local_expert_offset=m.slot_start,
+            local_num_experts=m.expert_size_per_partition,
+            routed_scaling_factor=ROUTED_SCALING,
+            routing_method_type=m.routing_method.routing_method_type,
+            topk_weights=None,
+            topk_ids=None,
+            gated_act_type=to_trtllm_gen_act_type(m.activation_type),
+            output=None,
+            use_shuffled_weight=m.quant_method.use_shuffled_weight,
+            weight_layout=m.quant_method.weight_layout,
+            do_finalize=do_finalize,
+        )
+
+
+def build_native(
+    tp_rank: int = 0, max_num_tokens: int = 128, identity_fc2: bool = False
+) -> NativeBf16Layer:
+    """The rank's native leaf.
+
+    ``identity_fc2`` replaces every expert's down projection by ``W[i, i % 256] = 1`` so that
+    ``gemm2_output[:, :256]`` is the native FC1 + SwiGLU intermediate (diagnostics only).
+    """
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
+    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
+    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import trtllm_gen_leaf
+    from tensorrt_llm.mapping import Mapping
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+
+    raw = load_rank(tp_rank)
+    bias = raw["gate.e_score_correction_bias"].to(torch.bfloat16).cuda()
+    routing = DeepSeekV3MoeRoutingMethod(
+        top_k=TOP_K,
+        n_group=1,
+        topk_group=1,
+        routed_scaling_factor=ROUTED_SCALING,
+        callable_e_score_correction_bias=lambda: bias,
+        is_fused=True,
+    )
+    model_config = ModelConfig(
+        quant_config=QuantConfig(),
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        moe_backend="TRTLLM",
+        max_num_tokens=max_num_tokens,
+    )
+    leaf_cls = trtllm_gen_leaf(None)
+    module = leaf_cls(
+        routing_method=routing,
+        num_experts=NUM_EXPERTS,
+        hidden_size=HIDDEN,
+        intermediate_size=INTER // TP_SIZE,
+        dtype=torch.bfloat16,
+        reduce_results=False,
+        model_config=model_config,
+        init_load_balancer=False,
+        weight_loading_mode=MoEWeightLoadingMode.VANILLA,
+    )
+    weights = {k: v for k, v in raw.items() if k[0].isdigit()}
+    if identity_fc2:
+        eye = torch.zeros(HIDDEN, INTER // TP_SIZE, dtype=torch.bfloat16)
+        eye[torch.arange(HIDDEN), torch.arange(HIDDEN) % (INTER // TP_SIZE)] = 1.0
+        weights = {k: (eye if k.endswith(".w2.weight") else v) for k, v in weights.items()}
+    module.load_weights([weights])
+    module.post_load_weights()
+    module = module.cuda()
+    return NativeBf16Layer(
+        module=module,
+        gate_weight=raw["gate.weight"].to(torch.bfloat16).cuda().contiguous(),
+        bias=bias,
+        norm_weight=raw["post_attention_layernorm.weight"].to(torch.bfloat16).cuda(),
+        tp_rank=tp_rank,
+    )
+
+
+def hidden_rows(layer: NativeBf16Layer, num: int, seed: int, scale: float = 1.0) -> torch.Tensor:
+    """``num`` RMS-normalized Gaussian rows times the layer's post-attention norm weight (BF16)."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    h = torch.randn(num, HIDDEN, device="cuda", generator=g) * scale
+    y = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + RMS_EPS)
+    return (y * layer.norm_weight.float()).to(torch.bfloat16).contiguous()
+
+
+def chain_like_rows(layer: NativeBf16Layer, num: int, seed: int) -> torch.Tensor:
+    """hidden_rows with the distinct-expert count of the draft chain at R = 6.
+
+    Independent rows select about 40 distinct experts, the chain's six draft rows about 31: rows 4 and 5 repeat
+    rows 0 and 1. Other row counts are hidden_rows unchanged.
+    """
+    x = hidden_rows(layer, num, seed)
+    if num == 6:
+        x[4:] = x[:2]
+    return x
+
+
+def distinct_experts(layer: NativeBf16Layer, logits: torch.Tensor) -> int:
+    """Distinct experts of the top-8 selection (sigmoid + correction bias), for reports."""
+    scores = torch.sigmoid(logits.float()) + layer.bias.float()
+    return int(scores.topk(8, dim=-1).indices.unique().numel())
+
+
+def top8_sets(bias: torch.Tensor, logits: torch.Tensor) -> List[frozenset]:
+    """Per row, the set of its top-8 experts (sigmoid + correction bias) as frozensets (for row selection)."""
+    scores = torch.sigmoid(logits.float()) + bias.float().to(logits.device)
+    return [frozenset(r) for r in scores.topk(TOP_K, dim=-1).indices.cpu().tolist()]
+
+
+def select_distinct(
+    sets: List[frozenset], rows: int, target: int, seed: int, steps: int = 600, sample: int = 256
+) -> Tuple[List[int], int]:
+    """Indices into ``sets`` (repeats allowed) of ``rows`` rows whose top-8 sets cover ``target`` distinct experts.
+
+    Seeded local search: start from random rows, then repeatedly replace one row by the candidate (a random sample
+    of the pool plus the selection's other rows, whose repeat lowers the count) that brings the union closest to the
+    target. Returns ``(indices, achieved count)``; the count differs from the target only when the search does not
+    reach it.
+    """
+    import random
+
+    rng = random.Random(seed)
+    n = len(sets)
+    pick = [rng.randrange(n) for _ in range(rows)]
+
+    def count(p):
+        u = set()
+        for i in p:
+            u |= sets[i]
+        return len(u)
+
+    best, best_d = list(pick), count(pick)
+    for _ in range(steps):
+        d = count(pick)
+        if abs(d - target) < abs(best_d - target):
+            best, best_d = list(pick), d
+        if d == target:
+            return pick, d
+        pos = rng.randrange(rows)
+        cands = rng.sample(range(n), min(n, sample)) + [pick[j] for j in range(rows) if j != pos]
+        scored = [
+            (abs(count(pick[:pos] + [c] + pick[pos + 1 :]) - target), rng.random(), c)
+            for c in cands
+        ]
+        pick[pos] = min(scored)[2]
+    d = count(pick)
+    if abs(d - target) < abs(best_d - target):
+        best, best_d = list(pick), d
+    return best, best_d
+
+
+class RowPool:
+    """Candidate rows and their top-8 sets, from which ``take`` draws rows with a chosen distinct-expert count."""
+
+    def __init__(self, xs: torch.Tensor, sets: List[frozenset]):
+        self.xs, self.sets = xs, sets
+
+    @classmethod
+    def gaussian(cls, layer: NativeBf16Layer, num: int = 1024, seed: int = 10_000) -> "RowPool":
+        xs = hidden_rows(layer, num, seed)
+        return cls(xs, top8_sets(layer.bias, layer.logits(xs)))
+
+    def take(self, rows: int, target: int, seed: int) -> Tuple[torch.Tensor, int]:
+        """``rows`` rows (repeats allowed) covering ``target`` distinct experts, and the count achieved."""
+        idx, d = select_distinct(self.sets, rows, target, seed)
+        return self.xs[torch.tensor(idx, device=self.xs.device)].contiguous(), d
+
+
+def embedding_rows(layer: NativeBf16Layer, token_ids: List[int]) -> torch.Tensor:
+    """RMSNorm(embedding) * post_attention_layernorm.weight for real token ids (BF16)."""
+    rd = Reader()
+    name = "model.embed_tokens.weight"
+    sl = rd._handle(name).get_slice(name)
+    emb = torch.stack([sl[t : t + 1][0] for t in token_ids]).cuda().float()
+    y = emb * torch.rsqrt(emb.pow(2).mean(-1, keepdim=True) + RMS_EPS)
+    return (y * layer.norm_weight.float()).to(torch.bfloat16).contiguous()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py
new file mode 100644
index 000000000000..9e57868210e0
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py
@@ -0,0 +1,101 @@
+# 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.
+"""Load ``trtllm::glm_moe_dsa_draft_routed_moe`` for tests.
+
+A TensorRT-LLM build of this branch registers the op in ``th_common``. On an
+older image (development overlay) it is compiled from this checkout's C++
+sources with ``torch.utils.cpp_extension``, unmodified and with the same op
+schema; results obtained that way are labelled DEV_JIT.
+"""
+
+import hashlib
+import os
+from pathlib import Path
+from typing import Callable, Sequence
+
+import torch
+
+_REPO = Path(__file__).resolve().parents[9]
+_SOURCES = (
+    _REPO / "cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp",
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu",
+)
+_HEADERS = (_REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h",)
+OP_NAME = "glm_moe_dsa_draft_routed_moe"
+_NAME = "draft_moe_jit"
+
+
+def _digest() -> str:
+    h = hashlib.sha256()
+    for path in (*_SOURCES, *_HEADERS):
+        h.update(path.read_bytes())
+    return h.hexdigest()[:16]
+
+
+def ensure_op(
+    build_root: str,
+    is_builder: bool = True,
+    barrier: Callable[[], None] = lambda: None,
+    defines: Sequence[str] = (),
+) -> str:
+    """Make the op available in this process; returns where it came from.
+
+    ``is_builder`` must be true on exactly one process per shared build root;
+    ``barrier`` synchronizes the processes after that one has built it.
+    ``defines`` (e.g. ``("GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1",)``) builds a probe variant of
+    the op; it needs a process in which the op is not registered yet.
+    """
+    import tensorrt_llm  # noqa: F401  (loads th_common and the trtllm namespace)
+
+    if hasattr(torch.ops.trtllm, OP_NAME):
+        if defines:
+            raise RuntimeError(
+                f"{OP_NAME} is already registered; a probe build needs a fresh process"
+            )
+        barrier()
+        return "builtin"
+    tag = _digest()
+    name = _NAME
+    if defines:
+        tag += "-" + hashlib.sha256(",".join(defines).encode()).hexdigest()[:8]
+        name += "_probe"
+    build_dir = Path(build_root) / f"draft_routed_moe-{tag}"
+    library = build_dir / f"{name}.so"
+    if is_builder and not library.exists():
+        from torch.utils.cpp_extension import load
+
+        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
+        build_dir.mkdir(parents=True, exist_ok=True)
+        extra = [f"-D{d}" for d in defines]
+        load(
+            name=name,
+            sources=[str(s) for s in _SOURCES],
+            extra_include_paths=[str(_REPO / "cpp/include"), str(_REPO / "cpp")],
+            extra_cflags=["-O3", "-std=c++17", *extra],
+            extra_cuda_cflags=[
+                "-O3",
+                "-std=c++17",
+                "--expt-relaxed-constexpr",
+                "-Xptxas=-v",
+                *extra,
+            ],
+            build_directory=str(build_dir),
+            is_python_module=False,
+            verbose=True,
+        )
+    barrier()
+    if not hasattr(torch.ops.trtllm, OP_NAME):
+        torch.ops.load_library(str(library))
+    return f"DEV_JIT {library}"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py
new file mode 100644
index 000000000000..115e9acfae19
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py
@@ -0,0 +1,152 @@
+# 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.
+"""The draft_routed_moe call site in the target's MoE.
+
+* a plan runs draft_routed_moe alone and next to F's GEN set, draft_routed_moe extends no
+  package and has a single numerical form; deferred_input_prep's gate
+  classifies draft_routed_moe's launch-configuration selectors as draft_routed_moe's;
+* ``GlmMoE.compute_routed_output`` dispatches to draft_routed_moe only when it is selected, before
+  routed_moe, and keeps the native path when draft_routed_moe declines;
+* an NVFP4 MoE layer (every target layer) declines draft_routed_moe before any device work,
+  so routed_moe / hop_routing see exactly the calls they see in F;
+* through ``GlmMoE.compute_routed_output`` on a decode step the MTP rank
+  shard returns draft_routed_moe's output, bit-identical to the native leaf.
+"""
+
+import inspect
+import os
+import re
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "..", "..", "shared", "kernels"))
+import plan_testlib  # noqa: E402
+
+F_GEN = (
+    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,router_gemm,"
+    "post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,rope_kv_append,"
+    "draft_loop,shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
+    "attention_output_tail_prefetch,attention_output_into_hop,routed_moe_into_hop,"
+    "deferred_norm,dense_mlp,deferred_input_prep,hop_routing,oproj_weight_staging,"
+    "fused_kv_indexer_projection,host_page_retention"
+)
+
+
+def _dispatch():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    return dispatch
+
+
+def test_selector_facts():
+    d = _dispatch()
+    assert "draft_routed_moe" in d.SELECTOR_PATHS
+    assert "draft_routed_moe" not in d.REQUIRED_FEATURES
+    assert "draft_routed_moe" not in d.NUMERICAL_VARIANTS
+    assert plan_testlib.plan_running(d, "draft_routed_moe").enabled("draft_routed_moe")
+    no_input_prep = frozenset(F_GEN.split(",")) - {"deferred_input_prep"}
+    assert plan_testlib.plan_running(d, no_input_prep | {"draft_routed_moe"}).enabled(
+        "draft_routed_moe"
+    )
+
+
+def test_call_site_order():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmMoE,
+    )
+
+    src = inspect.getsource(GlmMoE.compute_routed_output)
+    # The formatter wraps long statements: compare with the whitespace collapsed.
+    src = re.sub(r"(?<=[(\[])\s+|\s+(?=[)\]])", "", " ".join(src.split()))
+    draft_routed_moe = src.index('is_enabled("draft_routed_moe")')
+    assert draft_routed_moe < src.index('is_enabled("routed_moe")')
+    assert src.index('get_impl("draft_routed_moe", "routed_output")') > draft_routed_moe
+
+
+def test_nvfp4_layer_declines():
+    if not torch.cuda.is_available():
+        pytest.skip("needs a GPU")
+    import draft_moe_jit
+    import routed_moe_fixtures as wfx
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    # With the op loaded, the decline names the leaf type, not the missing op.
+    draft_moe_jit.ensure_op(
+        os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT", "/tmp/draft_routed_moe_jit")
+    )
+    layer = wfx.build_native_layer(10, 0)
+    x = torch.zeros(6, 6144, dtype=torch.bfloat16, device="cuda")
+    lg = torch.zeros(6, 256, dtype=torch.float32, device="cuda")
+    ok, reason = rm.eligibility(layer.module, layer.bias, x, lg, True)
+    print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_NVFP4_DECLINE {reason}", flush=True)
+    assert not ok and "BF16 leaf" in reason
+
+
+def test_patched_call_site_bit_exact(monkeypatch):
+    """``GlmMoE.compute_routed_output`` on a decode step, on a stand-in module that holds the MTP leaf."""
+    if not torch.cuda.is_available():
+        pytest.skip("needs a GPU")
+    import draft_moe_fixtures as fx
+    import draft_moe_jit
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmMoE,
+    )
+
+    draft_moe_jit.ensure_op(
+        os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT", "/tmp/draft_routed_moe_jit")
+    )
+    layer = fx.build_native(0)
+    plan_testlib.select(_dispatch(), "draft_routed_moe", monkeypatch)
+
+    class Gate:
+        e_score_correction_bias = layer.bias
+
+    class Stub:
+        use_dp = False
+        mapping = None
+        router_gemm = None
+        gate = Gate()
+        experts = layer.module
+
+    calls = {"n": 0}
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    inner = dispatch.get_impl("draft_routed_moe", "routed_output")
+
+    def counted(*a, **k):
+        out = inner(*a, **k)
+        calls["n"] += out is not None
+        return out
+
+    dispatch.register("draft_routed_moe", "routed_output", counted)
+    try:
+        for rows in (1, 6):
+            x = fx.hidden_rows(layer, rows, seed=40 + rows)
+            lg = layer.logits(x)
+            with dispatch.step(dispatch.Phase.DECODE):
+                out = GlmMoE.compute_routed_output(Stub(), x, None, None, True, lg)
+            assert torch.equal(out.view(torch.int16), layer.native(x, lg).view(torch.int16))
+    finally:
+        dispatch.register("draft_routed_moe", "routed_output", inner)
+    assert calls["n"] == 2
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py
new file mode 100644
index 000000000000..68b5934c391d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py
@@ -0,0 +1,775 @@
+# 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.
+"""draft_routed_moe against the native ``trtllm_bf16_moe`` on the real MTP-layer rank shard (single GPU).
+
+Every comparison is bit for bit on the BF16 outputs. On a mismatch the stage
+diagnostics compare the routing weights with the native ``expert_weights`` and
+draft_routed_moe's per-(token, slot) FC2 rows (``debug_fc2``) with the native unfinalized
+``gemm2_output`` through ``expanded_idx_to_permuted_idx``.
+
+Environment: ``GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT`` (JIT build directory when the image has no
+op), ``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TP_RANK`` (rank shard, default 0),
+``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TRIALS`` (calls per row count).
+"""
+
+import contextlib
+import json
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+import draft_moe_fixtures as fx  # noqa: E402
+import draft_moe_jit  # noqa: E402
+
+TRIALS = int(os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_TRIALS", "24"))
+
+# Keys of a launch configuration (``_launch``): FC1 and FC2 "ctas,warps" and the flag bits, each
+# "<value>" or "<one row>+<more rows>".
+FC1, FC2, FLAGS = "fc1", "fc2", "flags"
+
+# Kernel families: the default selectors (TMA-fed kernels: 64 FC1 CTAs of 4 MMA warps at R = 1, 152 CTAs with two
+# chains per warp at R > 1 and the triples path above 38 distinct experts with its early trigger, FC1 loads L2
+# evict-first at R > 1, the 3-warp FC2 on 128 CTAs, FC2 weight prefetch at R = 1), the
+# cp.async kernels, the TMA-fed kernels on other grids (flag 64) and the TMA-fed kernels with two chains per FC1 warp,
+# two experts per FC2 warp step, FC1 L2 prefetch and the FC2 weight prefetch (flags 8, 16, 128, 2; the FC1 ring of
+# 32 KB stages is 6 deep at R <= 2, 4 deep at R = 3..7 and 3 deep at R = 8), and the FC1 triples path on a 120-CTA
+# grid (flag 4096: R > 1 calls whose 128-row tiles exceed the grid but whose 64-row sub-tiles fit three per CTA; the
+# others run the regular path, e.g. R = 8 with about 50 distinct experts; flag 8192: the early trigger on that path
+# only), and the default R > 1 kernels with the activation streamed through the FC1 ring (flag 65536). All must be
+# bit-exact.
+FAMILIES = {
+    "default": {},
+    "cpasync": {
+        FC1: "128,2",
+        FC2: "96,4",
+        FLAGS: "2",
+    },
+    "tma": {
+        FC1: "128,2+152,4",
+        FC2: "96,4",
+        FLAGS: "96",
+    },
+    "tmapairs": {
+        FC1: "64,2+152,4",
+        FC2: "96,4",
+        FLAGS: "250",
+    },
+    "triples": {
+        FC1: "64,4+120,4",
+        FC2: "128,3",
+        FLAGS: "66+12360",
+    },
+    "streamx": {
+        FC1: "64,4+152,4",
+        FC2: "128,3",
+        FLAGS: "66+77896",
+    },
+}
+DELAYED_FLAGS = {
+    "default": ("66+45128", "70+76"),
+    "cpasync": ("2", "6"),
+    "tma": ("96", "100"),
+    "tmapairs": ("250", "254"),
+    "triples": ("66+12360", "70+4172"),
+    "streamx": ("66+77896", "70+77900"),
+}
+
+
+def _per_rows(raw, parse):
+    """``(one row, more rows)`` of a "<value>" or "<one row>+<more rows>" entry."""
+    parts = raw.split("+")
+    assert len(parts) <= 2, raw
+    return parse(parts[0]), parse(parts[-1])
+
+
+def _ctas_warps(raw):
+    ctas, warps = (int(v) for v in raw.split(","))
+    return ctas, warps
+
+
+@contextlib.contextmanager
+def _launch(config):
+    """``config`` set over draft_routed_moe's launch defaults for the duration of the block."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    saved = rm.DEFAULT_FC1, rm.DEFAULT_FC2, rm.DEFAULT_FLAGS
+    config = config or {}
+    try:
+        if FC1 in config:
+            rm.DEFAULT_FC1 = _per_rows(config[FC1], _ctas_warps)
+        if FC2 in config:
+            rm.DEFAULT_FC2 = _per_rows(config[FC2], _ctas_warps)
+        if FLAGS in config:
+            rm.DEFAULT_FLAGS = _per_rows(config[FLAGS], int)
+        yield
+    finally:
+        rm.DEFAULT_FC1, rm.DEFAULT_FC2, rm.DEFAULT_FLAGS = saved
+
+
+@pytest.fixture(scope="module")
+def layer():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("needs an SM100-family GPU")
+    origin = draft_moe_jit.ensure_op(
+        os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT", "/tmp/draft_routed_moe_jit")
+    )
+    print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_OP_ORIGIN {origin}", flush=True)
+    return fx.build_native(int(os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_TP_RANK", "0")))
+
+
+def _draft_moe(layer, x, logits, debug=False, config=None):
+    """draft_routed_moe's output, its debug FC2 rows and whether the TMA-fed kernels ran."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    with _launch(config):
+        dbg = (
+            torch.full(
+                (x.shape[0], fx.TOP_K, fx.HIDDEN), float("nan"), dtype=torch.bfloat16, device="cuda"
+            )
+            if debug
+            else None
+        )
+        tma = bool(rm.launch_config(x.shape[0])[2] & rm.FLAG_TMA)
+        out = rm.run(layer.module, layer.bias, x, logits, debug_fc2=dbg)
+        return out, dbg, tma
+
+
+def _bits(t):
+    return t.contiguous().view(torch.int16)
+
+
+_IDENTITY = {}
+
+
+def _identity_layer(layer):
+    """The same rank's native leaf with identity-like down projections (built on the first mismatch only)."""
+    if layer.tp_rank not in _IDENTITY:
+        _IDENTITY[layer.tp_rank] = fx.build_native(layer.tp_rank, identity_fc2=True)
+    return _IDENTITY[layer.tp_rank]
+
+
+def _h_image_order(rows: int = 8, inter: int = 256) -> torch.Tensor:
+    """Index map of an intermediate entry stored as its shared-memory image (TMA family): plain [tok, j] -> image."""
+    tok = torch.arange(rows)[:, None]
+    j = torch.arange(inter)[None, :]
+    u = j // 8
+    return tok * inter + (((u & ~7) | ((u & 7) ^ (tok & 7))) * 8) + (j & 7)
+
+
+def _diagnose(layer, x, logits, out_draft_moe, dbg, tma=False):
+    """Stage-level comparison of one mismatching call (returns a dict of counts).
+
+    Reads draft_routed_moe's routing record and intermediate scratch, which still hold this call's
+    values.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    scratch = rm._scratch(x.device, create=False)
+    route = scratch.route.cpu()
+    h_draft_moe = scratch.h.view(64, 8 * (fx.INTER // fx.TP_SIZE)).clone()
+    if tma:
+        h_draft_moe = h_draft_moe[:, _h_image_order().flatten().to(h_draft_moe.device)]
+    h_draft_moe = h_draft_moe.view(64, 8, fx.INTER // fx.TP_SIZE)
+    gemm2, weights, idx = layer.native(x, logits, do_finalize=False)[:3]
+    rows = x.shape[0]
+    idx = idx.reshape(-1).long()
+    w_draft_moe = route[129 : 129 + rows * fx.TOP_K].view(torch.float32).reshape(rows, fx.TOP_K)
+    slot_d = route[65 : 65 + rows * fx.TOP_K].reshape(rows, fx.TOP_K).long()
+    w_nat = weights.float().cpu().reshape(rows, fx.TOP_K)
+    fc2_nat = gemm2[idx].reshape(rows, fx.TOP_K, fx.HIDDEN)
+    fc2_diff = (_bits(fc2_nat) != _bits(dbg)).sum(-1)
+    # Native FC1 + SwiGLU of every (token, slot): gemm2 output of the identity leaf, columns 0..255.
+    ident = _identity_layer(layer)
+    g_id, _, idx_id = ident.native(x, logits, do_finalize=False)[:3]
+    h_nat = g_id[idx_id.reshape(-1).long()].reshape(rows, fx.TOP_K, fx.HIDDEN)[
+        ..., : fx.INTER // fx.TP_SIZE
+    ]
+    h_mine = torch.stack(
+        [torch.stack([h_draft_moe[slot_d[t, k], t] for k in range(fx.TOP_K)]) for t in range(rows)]
+    )
+    h_diff = (h_nat.float() != h_mine.float()).sum(-1)
+    return {
+        "weights_differ": int((w_draft_moe != w_nat).sum()),
+        "fc1_h_differ_per_slot": h_diff.tolist(),
+        "fc2_rows_differ_per_slot": fc2_diff.tolist(),
+        "out_elems_differ": int((_bits(out_draft_moe) != _bits(layer.native(x, logits))).sum()),
+    }
+
+
+def _check(layer, x, logits, label, config=None):
+    native = layer.native(x, logits)
+    out, dbg, tma = _draft_moe(layer, x, logits, debug=True, config=config)
+    torch.cuda.synchronize()
+    differ = int((_bits(out) != _bits(native)).sum())
+    if differ:
+        info = _diagnose(layer, x, logits, out, dbg, tma)
+        print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_MISMATCH {label} " + json.dumps(info), flush=True)
+    return differ
+
+
+@pytest.mark.parametrize("family", list(FAMILIES))
+@pytest.mark.parametrize("rows", [1, 2, 3, 4, 5, 6, 7, 8])
+def test_bit_exact_gaussian(layer, rows, family):
+    bad = 0
+    for trial in range(TRIALS):
+        x = fx.hidden_rows(layer, rows, seed=1000 * rows + trial)
+        bad += (
+            _check(
+                layer,
+                x,
+                layer.logits(x),
+                f"gauss {family} R={rows} trial={trial}",
+                config=FAMILIES[family],
+            )
+            > 0
+        )
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE gaussian {family} R={rows} calls={TRIALS} mismatching_calls={bad}",
+        flush=True,
+    )
+    assert bad == 0
+
+
+@pytest.mark.parametrize("family", list(FAMILIES))
+@pytest.mark.parametrize("rows", [1, 6])
+def test_bit_exact_embeddings(layer, rows, family):
+    """Consecutive real token embeddings (normalized) as activation rows."""
+    ids = list(range(1000, 1000 + 64 * rows))
+    bad = 0
+    for trial in range(8):
+        x = fx.embedding_rows(layer, ids[trial * rows : (trial + 1) * rows])
+        bad += (
+            _check(
+                layer,
+                x,
+                layer.logits(x),
+                f"emb {family} R={rows} trial={trial}",
+                config=FAMILIES[family],
+            )
+            > 0
+        )
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE embeddings {family} R={rows} calls=8 mismatching_calls={bad}",
+        flush=True,
+    )
+    assert bad == 0
+
+
+@pytest.mark.parametrize("family", list(FAMILIES))
+def test_routing_ties(layer, family):
+    """Equal selection keys.
+
+    Equal logits for every expert (the bias decides, ties to the lower index), and pairs of experts with equal
+    bias and equal logits.
+    """
+    bad = 0
+    x = fx.hidden_rows(layer, 6, seed=7)
+    flat = torch.zeros(6, fx.NUM_EXPERTS, device="cuda")
+    flat[1] += 3.0
+    flat[2] -= 3.0
+    bias = layer.bias.float()
+    # rows 3..5: logits that make sigmoid + bias equal for groups of experts with equal bias
+    for r in range(3, 6):
+        flat[r] = torch.where(bias == bias[r * 7], 1.0, -2.0)
+    bad += _check(layer, x, flat.contiguous(), "ties", config=FAMILIES[family]) > 0
+    big = torch.randn(6, fx.NUM_EXPERTS, device="cuda") * 40.0
+    bad += _check(layer, x, big.contiguous(), "saturated logits", config=FAMILIES[family]) > 0
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE ties {family} calls=2 mismatching_calls={bad}",
+        flush=True,
+    )
+    assert bad == 0
+
+
+@pytest.mark.parametrize(
+    "config",
+    [
+        {
+            FC1: "148,1",
+            FC2: "148,1",
+            FLAGS: "0",
+        },
+        {
+            FC1: "64,4",
+            FC2: "48,2",
+            FLAGS: "6",
+        },
+        {
+            FC1: "1,1",
+            FC2: "1,1",
+            FLAGS: "2",
+        },
+        {
+            FC1: "128,2",
+            FC2: "96,4",
+            FLAGS: "26",
+        },
+        {
+            FC1: "37,4",
+            FC2: "5,2",
+            FLAGS: "24",
+        },
+        {
+            FC1: "1,1+148,4",
+            FC2: "148,1+5,2",
+            FLAGS: "0+30",
+        },
+        {
+            FC1: "148,2+64,4",
+            FC2: "96,4",
+            FLAGS: "34+58",
+        },
+        {
+            FC1: "128,4",
+            FC2: "37,1",
+            FLAGS: "62",
+        },
+        # Delayed producer of the FC2 grid: one FC1 warp (milliseconds) with the early trigger, so the FC2 CTAs
+        # are resident long before the routing record and the intermediate exist.
+        {
+            FC1: "1,1",
+            FC2: "148,4",
+            FLAGS: "6",
+        },
+        # TMA-fed kernels: one CTA running every tile (persistent loop, ring wrap-around across tiles), deep rings
+        # with the early trigger, other grid shapes, and a one-CTA FC1 producer ahead of a full FC2 grid.
+        {
+            FC1: "1,1",
+            FC2: "1,1",
+            FLAGS: "64",
+        },
+        {
+            FC1: "37,4",
+            FC2: "5,2",
+            FLAGS: "100",
+        },
+        {
+            FC1: "152,2",
+            FC2: "152,4",
+            FLAGS: "96",
+        },
+        {
+            FC1: "64,4+256,1",
+            FC2: "192,2+96,4",
+            FLAGS: "64+96",
+        },
+        {
+            FC1: "1,1",
+            FC2: "152,4",
+            FLAGS: "68",
+        },
+        # TMA-fed with two chains per FC1 warp: one CTA of one warp, an odd grid with the early trigger, L2 prefetch
+        # and the early FC2 prefetch, grids with more CTAs than tiles, and one 4-warp CTA ahead of a full FC2 grid.
+        {
+            FC1: "1,1",
+            FC2: "1,1",
+            FLAGS: "72",
+        },
+        {
+            FC1: "37,4",
+            FC2: "5,2",
+            FLAGS: "460",
+        },
+        {
+            FC1: "152,2+100,4",
+            FC2: "152,4",
+            FLAGS: "226+98",
+        },
+        {
+            FC1: "1,4",
+            FC2: "152,4",
+            FLAGS: "236",
+        },
+        # TMA-fed FC2 with 3 MMA warps (48-row groups): the full grid of 128 groups and an odd persistent grid.
+        {
+            FC1: "128,2+152,4",
+            FC2: "128,3",
+            FLAGS: "66+68",
+        },
+        {
+            FC1: "64,4",
+            FC2: "37,3",
+            FLAGS: "64",
+        },
+        # 3-warp FC2 with expert pairs, after FC1 tile pairs at R > 1: the full grid and an odd persistent grid.
+        {
+            FC1: "64,4+152,4",
+            FC2: "128,3",
+            FLAGS: "82+90",
+        },
+        {
+            FC1: "32,4+136,4",
+            FC2: "37,3",
+            FLAGS: "74+88",
+        },
+        # FC1 triples (flag 4096): the default grid with tile pairs, without tile pairs (64-row tiles), with the FC2
+        # weight prefetch after the last copy and right after routing, and a grid too small for one round of
+        # triples (regular path).
+        {
+            FC1: "64,4+152,4",
+            FC2: "128,3",
+            FLAGS: "66+4172",
+        },
+        {
+            FC1: "152,4",
+            FC2: "128,3",
+            FLAGS: "4164",
+        },
+        {
+            FC1: "112,4",
+            FC2: "96,4",
+            FLAGS: "4170+4426",
+        },
+        {
+            FC1: "20,4",
+            FC2: "37,3",
+            FLAGS: "4168",
+        },
+    ],
+)
+def test_launch_configs(layer, config):
+    bad = 0
+    for rows in (1, 3, 6, 8):
+        x = fx.hidden_rows(layer, rows, seed=99 + rows)
+        bad += _check(layer, x, layer.logits(x), f"config {config} R={rows}", config=config) > 0
+    assert bad == 0
+
+
+def _distinct(layer, logits):
+    """Distinct experts of a call (float32 restatement of the selection; only used to pick test cases)."""
+    keys = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5 + layer.bias.float()[None, :]
+    return len(set(keys.topk(fx.TOP_K, dim=-1).indices.flatten().tolist()))
+
+
+@pytest.mark.parametrize("flags", ["24", "58", "16", "80", "90"])
+def test_pairs_odd_distinct(layer, flags):
+    """Expert pairs in FC2 / tile pairs in FC1 with an odd number of distinct experts.
+
+    An unpaired expert per item makes the pair parity alternate across a warp's items; every warp has several
+    items.
+    """
+    config = {
+        FC1: "37,2",
+        FC2: "48,1",
+        FLAGS: flags,
+    }
+    bad, odd, seed = 0, 0, 0
+    while odd < 8 and seed < 400:
+        rows = (2, 3, 5, 7)[seed % 4]
+        x = fx.hidden_rows(layer, rows, seed=7000 + seed)
+        lg = layer.logits(x)
+        seed += 1
+        if _distinct(layer, lg) % 2 == 0:
+            continue
+        odd += 1
+        bad += (
+            _check(layer, x, lg, f"odd D flags={flags} R={rows} seed={seed - 1}", config=config) > 0
+        )
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE pairs flags={flags} odd-D calls={odd} mismatching_calls={bad}",
+        flush=True,
+    )
+    assert odd == 8 and bad == 0
+
+
+# Grids and flags of the triples test: the default grid with tile pairs (one round of triples up to 57 distinct
+# experts; early trigger always or on the triples path only), 64-row tiles (8 per expert), and a 120-CTA grid
+# (triples from 31 to 45 distinct experts).
+TRIPLES_CONFIGS = (
+    ("152,4", "128,3", "4172"),
+    ("152,4", "128,3", "12360"),
+    ("152,4", "128,3", "4160"),
+    ("120,4", "96,4", "4168"),
+    ("152,4", "128,3", "45128"),
+    ("152,4", "128,3", "77896"),
+    ("152,4", "128,3", "110664"),
+)
+
+
+@pytest.mark.parametrize("config", TRIPLES_CONFIGS, ids=lambda c: "/".join(c))
+def test_triples_distinct(layer, config):
+    """R = 5..8 with chosen distinct-expert counts around the FC1 path boundaries.
+
+    The boundaries are 38/39 and 57/58 distinct experts on the 152-CTA grid with tile pairs and 30/31 and 45/46 on
+    120 CTAs; every call is bit-exact whichever FC1 path it takes.
+    """
+    fc1, fc2, flags = config
+    pool = fx.RowPool.gaussian(layer)
+    cases = [
+        (5, 39),
+        (5, 40),
+        (6, 30),
+        (6, 31),
+        (6, 38),
+        (6, 39),
+        (6, 44),
+        (6, 45),
+        (6, 46),
+        (6, 48),
+    ]
+    cases += [(7, 41), (7, 50), (7, 56), (8, 57), (8, 58), (8, 64)]
+    bad, reached = 0, 0
+    for i, (rows, target) in enumerate(cases):
+        x, d = pool.take(rows, target, seed=3000 + i)
+        reached += d == target
+        bad += (
+            _check(
+                layer,
+                x,
+                layer.logits(x),
+                f"triples {config} R={rows} D={d}",
+                config={
+                    FC1: fc1,
+                    FC2: fc2,
+                    FLAGS: flags,
+                },
+            )
+            > 0
+        )
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE triples {config} calls={len(cases)} at target D={reached} "
+        f"mismatching_calls={bad}",
+        flush=True,
+    )
+    assert bad == 0 and reached >= len(cases) - 2
+
+
+def _route_weight_bits(layer, x, logits):
+    """(draft_routed_moe routing record weights, native expert_weights) of one call as int32 bit
+    patterns of their FP32 values."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    rows = x.shape[0]
+    out = rm.run(layer.module, layer.bias, x, logits)
+    scratch = rm._scratch(x.device, create=False)
+    fast_w = scratch.route[129 : 129 + rows * fx.TOP_K].clone()
+    nat_w = layer.native(x, logits, do_finalize=False)[1].reshape(-1).float().view(torch.int32)
+    return fast_w, nat_w, out
+
+
+def test_route_record(layer):
+    """draft_routed_moe's routing weights (its routing record) against the native ``expert_weights``, bit for bit.
+
+    102,400 rows (R = 8; logits of several scales, every eighth call rounded so that keys tie), plus the outputs of
+    every 64th call.
+    """
+    calls, rows = 12800, 8
+    x = fx.hidden_rows(layer, rows, seed=4)
+    g = torch.Generator(device="cuda").manual_seed(12)
+    differ = torch.zeros((), dtype=torch.int64, device="cuda")
+    out_bad = 0
+    for i in range(calls):
+        scale = (0.5, 2.0, 6.0, 30.0)[i % 4]
+        lg = torch.randn(rows, fx.NUM_EXPERTS, device="cuda", generator=g) * scale
+        if i % 8 == 7:
+            lg = torch.round(lg * 4.0) / 4.0
+        fast_w, nat_w, out = _route_weight_bits(layer, x, lg.contiguous())
+        differ += (fast_w != nat_w).sum()
+        if i % 64 == 0:
+            out_bad += int((_bits(out) != _bits(layer.native(x, lg.contiguous()))).sum()) > 0
+    differ = int(differ)
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE route record rows={calls * rows} differing weights={differ} "
+        f"mismatching outputs={out_bad}/{calls // 64}",
+        flush=True,
+    )
+    assert differ == 0 and out_bad == 0
+
+
+def test_route_record_nonfinite(layer):
+    """Logits with NaN, +Inf and -Inf entries.
+
+    Routing weights and outputs against native, bit for bit (NaN payloads included).
+    """
+    rows = 8
+    x = fx.hidden_rows(layer, rows, seed=4)
+    g = torch.Generator(device="cuda").manual_seed(13)
+    w_bad, out_bad, calls = 0, 0, 64
+    for i in range(calls):
+        lg = torch.randn(rows, fx.NUM_EXPERTS, device="cuda", generator=g) * 2.0
+        cols = torch.randint(0, fx.NUM_EXPERTS, (3,), device="cuda", generator=g)
+        lg[i % rows, cols[0]] = float("inf")
+        lg[(i + 1) % rows, cols[1]] = float("-inf")
+        if i % 2:
+            lg[(i + 2) % rows, cols[2]] = float("nan")
+        fast_w, nat_w, out = _route_weight_bits(layer, x, lg.contiguous())
+        nat_out = layer.native(x, lg.contiguous())
+        w_differ = (fast_w != nat_w).view(rows, fx.TOP_K)
+        o_differ = (_bits(out) != _bits(nat_out)).sum(-1)
+        if (bool(w_differ.any()) or bool(o_differ.any())) and w_bad + out_bad < 4:
+            from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+                routed_moe as rm,
+            )
+
+            route = rm._scratch(x.device, create=False).route.cpu()
+            distinct = route[1:65]
+            ids = distinct[route[65 : 65 + rows * fx.TOP_K].long()].view(rows, fx.TOP_K)
+            zw, nw = fast_w.view(rows, -1).tolist(), nat_w.view(rows, -1).tolist()
+            zo, no = _bits(out), _bits(nat_out)
+            for r in range(rows):
+                if bool(w_differ[r].any()) or int(o_differ[r]) > 0:
+                    print(
+                        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_NONFINITE call={i} row={r} "
+                        f"nan_in_row={bool(lg[r].isnan().any())} "
+                        f"draft_moe_ids={ids[r].tolist()} fast_w={[hex(v & 0xFFFFFFFF) for v in zw[r]]} "
+                        f"native_w={[hex(v & 0xFFFFFFFF) for v in nw[r]]} out_differ={int(o_differ[r])} "
+                        f"draft_moe_out_nan={int(out[r].isnan().sum())} native_out_nan={int(nat_out[r].isnan().sum())} "
+                        f"draft_moe_out0={hex(int(zo[r, 0]) & 0xFFFF)} native_out0={hex(int(no[r, 0]) & 0xFFFF)}",
+                        flush=True,
+                    )
+        w_bad += bool(w_differ.any())
+        out_bad += bool(o_differ.any())
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE route record non-finite calls={calls} mismatching "
+        f"weights={w_bad} outputs={out_bad}",
+        flush=True,
+    )
+    assert w_bad == 0 and out_bad == 0
+
+
+@pytest.mark.parametrize("family", list(FAMILIES))
+def test_delayed_producer(layer, family):
+    """Inputs written late by their producer.
+
+    A copy writes the inputs after a long device-side sleep on the same stream, eagerly and inside a CUDA graph
+    (PDL edges as captured), with the default and the early-trigger configurations.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    bad = 0
+    flag_sets = DELAYED_FLAGS[family]
+    for flags in flag_sets:
+        with _launch({**FAMILIES[family], FLAGS: flags}):
+            for rows in (1, 6):
+                x0 = fx.hidden_rows(layer, rows, seed=11 + rows)
+                lg0 = layer.logits(x0)
+                ref = layer.native(x0, lg0)
+                xs = torch.zeros_like(x0)
+                ls = torch.zeros_like(lg0)
+                torch.cuda.synchronize()
+                torch.cuda._sleep(20_000_000)
+                xs.copy_(x0)
+                ls.copy_(lg0)
+                out = rm.run(layer.module, layer.bias, xs, ls)
+                torch.cuda.synchronize()
+                bad += int((_bits(out) != _bits(ref)).sum()) > 0
+                xs.zero_()
+                ls.zero_()
+                graph = torch.cuda.CUDAGraph()
+                stream = torch.cuda.Stream()
+                with torch.cuda.stream(stream):
+                    with torch.cuda.graph(graph, stream=stream):
+                        torch.cuda._sleep(20_000_000)
+                        xs.copy_(x0)
+                        ls.copy_(lg0)
+                        out_g = rm.run(layer.module, layer.bias, xs, ls)
+                for _ in range(3):
+                    xs.zero_()
+                    ls.zero_()
+                    torch.cuda.synchronize()
+                    graph.replay()
+                    torch.cuda.synchronize()
+                    bad += int((_bits(out_g) != _bits(ref)).sum()) > 0
+    print(
+        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE delayed producer {family} calls=16 mismatching={bad}",
+        flush=True,
+    )
+    assert bad == 0
+
+
+@pytest.mark.parametrize("family", list(FAMILIES))
+def test_graph_replay(layer, family):
+    """Captured once, replayed with new inputs copied into the static buffers."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    with _launch(FAMILIES[family]):
+        _graph_replay(layer, rm, family)
+
+
+def _graph_replay(layer, rm, family):
+    for rows in (1, 6):
+        xs = torch.empty(rows, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
+        ls = torch.empty(rows, fx.NUM_EXPERTS, dtype=torch.float32, device="cuda")
+        x0 = fx.hidden_rows(layer, rows, seed=5)
+        xs.copy_(x0)
+        ls.copy_(layer.logits(x0))
+        rm.run(layer.module, layer.bias, xs, ls)  # scratch exists before capture
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        stream = torch.cuda.Stream()
+        with torch.cuda.stream(stream):
+            with torch.cuda.graph(graph, stream=stream):
+                out = rm.run(layer.module, layer.bias, xs, ls)
+        bad = 0
+        for trial in range(6):
+            x = fx.hidden_rows(layer, rows, seed=500 + trial)
+            xs.copy_(x)
+            ls.copy_(layer.logits(x))
+            graph.replay()
+            torch.cuda.synchronize()
+            bad += int((_bits(out) != _bits(layer.native(xs, ls))).sum()) > 0
+        print(
+            f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE graph {family} R={rows} replays=6 mismatching={bad}",
+            flush=True,
+        )
+        assert bad == 0
+
+
+def test_eligibility(layer):
+    """Every unsupported call declines before any device work."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+
+    x = fx.hidden_rows(layer, 6, seed=3)
+    lg = layer.logits(x)
+    assert rm.eligibility(layer.module, layer.bias, x, lg, True) == (True, "")
+    cases = {
+        "unfinalized": (layer.module, layer.bias, x, lg, False),
+        "rows 9": (
+            layer.module,
+            layer.bias,
+            fx.hidden_rows(layer, 9, seed=1),
+            layer.logits(fx.hidden_rows(layer, 9, seed=1)),
+            True,
+        ),
+        "fp16 activation": (layer.module, layer.bias, x.half(), lg, True),
+        "bf16 logits": (layer.module, layer.bias, x, lg.bfloat16(), True),
+        "strided logits": (layer.module, layer.bias, x, torch.cat([lg, lg], 1)[:, ::2], True),
+        "other bias": (layer.module, layer.bias.clone(), x, lg, True),
+        "fp32 bias": (layer.module, layer.bias.float(), x, lg, True),
+        "not a module": (object(), layer.bias, x, lg, True),
+    }
+    for name, args in cases.items():
+        ok, reason = rm.eligibility(*args)
+        print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_DECLINE {name}: {reason}", flush=True)
+        assert not ok and reason, name
+    assert rm.routed_output(layer.module, layer.bias, x, lg, False) is None
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py
new file mode 100644
index 000000000000..62c71bb72228
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py
@@ -0,0 +1,161 @@
+# 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.
+"""Load ``trtllm::glm_moe_dsa_post_attention_hop_routing`` (and the test occupier) for tests.
+
+A TensorRT-LLM build of this branch registers the routing hop in
+``th_common``. On an older image (development overlay) it is compiled from
+this checkout's C++ sources with ``torch.utils.cpp_extension``, unmodified and
+with the same op schema. ``GLM_MOE_DSA_HOP_ROUTING_JIT_DEFINES`` adds compile definitions for
+diagnostic builds (``-DTLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS``); such a build is always
+compiled here and registers ``trtllm::glm_moe_dsa_post_attention_hop_routing_diag``.
+
+``ensure_occupier`` builds ``z2test::occupy(int blocks, int ns)``, a test-only
+kernel of ``blocks`` one-warp CTAs that each spin for ``ns`` nanoseconds; the
+progress tests run it on another stream to hold SMs while the hop waits.
+"""
+
+import hashlib
+import os
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+_REPO = Path(__file__).resolve().parents[9]
+_SOURCES = (
+    _REPO / "cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp",
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu",
+)
+_HEADERS = (
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h",
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h",
+)
+_EXTRA = [f for f in os.environ.get("GLM_MOE_DSA_HOP_ROUTING_JIT_DEFINES", "").split() if f]
+DIAG = "-DTLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS" in _EXTRA
+# The diagnostic build registers ``glm_moe_dsa_post_attention_hop_routing_diag``.
+OP_NAME = (
+    "glm_moe_dsa_post_attention_hop_routing_diag"
+    if DIAG
+    else "glm_moe_dsa_post_attention_hop_routing"
+)
+_NAME = "hop_routing_diag_jit" if DIAG else "hop_routing_jit"
+
+_OCCUPIER_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+__global__ void z2OccupyKernel(long long ns)
+{
+    unsigned long long start;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(start));
+    unsigned long long now = start;
+    while (now - start < static_cast<unsigned long long>(ns))
+    {
+        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now));
+    }
+}
+
+void occupy(int64_t blocks, int64_t ns)
+{
+    // One CTA per SM: all of an SM's shared memory, so no other CTA fits next to it.
+    int device = 0;
+    cudaGetDevice(&device);
+    int smem = 0;
+    cudaDeviceGetAttribute(&smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+    cudaFuncSetAttribute(z2OccupyKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+    z2OccupyKernel<<<static_cast<unsigned>(blocks), 32, smem, at::cuda::getCurrentCUDAStream()>>>(ns);
+}
+
+// No tensor arguments: registered as a catch-all kernel.
+TORCH_LIBRARY(z2test, m)
+{
+    m.def("occupy(int blocks, int ns) -> ()", &occupy);
+}
+"""
+
+
+def _digest() -> str:
+    h = hashlib.sha256(" ".join(_EXTRA).encode())
+    for path in (*_SOURCES, *_HEADERS):
+        h.update(path.read_bytes())
+    return h.hexdigest()[:16]
+
+
+def ensure_op(build_root: str, is_builder: bool, barrier: Callable[[], None]) -> str:
+    """Make the routing hop available on every rank; returns where it came from.
+
+    ``is_builder`` must be true on exactly one rank per shared build root;
+    ``barrier`` synchronizes the ranks after that rank has built it.
+    """
+    import tensorrt_llm  # noqa: F401  (loads th_common and the trtllm namespace)
+
+    if hasattr(torch.ops.trtllm, OP_NAME) and not _EXTRA:
+        barrier()
+        return "builtin"
+    build_dir = Path(build_root) / f"hop_routing-{_digest()}"
+    library = build_dir / f"{_NAME}.so"
+    if is_builder and not library.exists():
+        from torch.utils.cpp_extension import load
+
+        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
+        build_dir.mkdir(parents=True, exist_ok=True)
+        load(
+            name=_NAME,
+            sources=[str(s) for s in _SOURCES],
+            extra_include_paths=[str(_REPO / "cpp/include"), str(_REPO / "cpp")],
+            extra_cflags=["-O3", "-std=c++17", *_EXTRA],
+            extra_cuda_cflags=[
+                "-O3",
+                "-std=c++17",
+                "--expt-relaxed-constexpr",
+                "-Xptxas=-v",
+                *_EXTRA,
+            ],
+            build_directory=str(build_dir),
+            is_python_module=False,
+            verbose=True,
+        )
+    barrier()
+    if not hasattr(torch.ops.trtllm, OP_NAME):
+        torch.ops.load_library(str(library))
+    return str(library)
+
+
+def ensure_occupier(build_root: str, is_builder: bool, barrier: Callable[[], None]) -> None:
+    """Make ``torch.ops.z2test.occupy`` available on every rank."""
+    if hasattr(torch.ops, "z2test") and hasattr(torch.ops.z2test, "occupy"):
+        barrier()
+        return
+    digest = hashlib.sha256(_OCCUPIER_SRC.encode()).hexdigest()[:16]
+    build_dir = Path(build_root) / f"hop_routing_occupier-{digest}"
+    library = build_dir / "z2test_occupy.so"
+    if is_builder and not library.exists():
+        from torch.utils.cpp_extension import load
+
+        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
+        build_dir.mkdir(parents=True, exist_ok=True)
+        src = build_dir / "z2test_occupy.cu"
+        src.write_text(_OCCUPIER_SRC)
+        load(
+            name="z2test_occupy",
+            sources=[str(src)],
+            extra_cuda_cflags=["-O3", "-std=c++17"],
+            build_directory=str(build_dir),
+            is_python_module=False,
+            verbose=False,
+        )
+    barrier()
+    if not (hasattr(torch.ops, "z2test") and hasattr(torch.ops.z2test, "occupy")):
+        torch.ops.load_library(str(library))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py
new file mode 100644
index 000000000000..4a85a90d7f6e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py
@@ -0,0 +1,214 @@
+# 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.
+"""Native routing oracle and boundary logits for the hop_routing tests.
+
+``native_routing`` runs the served TRT-LLM-Gen routing runner (DeepSeekV3,
+n_group = topk_group = 1, top_k = 8, routed scaling 2.5, bf16 bias) through
+``trtllm::moe_topk_sort`` and reads the expert of every ``(t, k)`` back
+through the permutation: ``expanded_idx_to_permuted_idx[t, k]`` lies in
+tile ``p // TILE`` whose expert is ``tile_idx_to_expert_idx``. The weights
+are the runner's bf16 ``token_final_scales`` in top-k order.
+
+``boundary_rows`` builds FP32 logit rows whose native selection sits on an
+exact tie or a one-ulp key difference at the 8th/9th boundary (same warp,
+different warps, different clusters of the hop), plus NaN, +-Inf, -0.0,
+huge and subnormal logits; each row is labelled.
+"""
+
+import os
+import sys
+from typing import List, Optional, Tuple
+
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
+import routed_moe_fixtures as fx  # noqa: E402
+
+NUM_EXPERTS = 256
+TOP_K = 8
+TILE = 8
+ROUTED_SCALING = 2.5
+DEEPSEEK_V3 = 2  # RoutingMethodType.DeepSeekV3
+
+
+def native_routing(
+    logits: torch.Tensor, bias_bf16: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """``(ids int32 [R, 8], weights bf16 [R, 8])`` of the native runner."""
+    logits = logits.float().contiguous()
+    out = torch.ops.trtllm.moe_topk_sort(
+        logits,
+        bias_bf16.contiguous(),
+        NUM_EXPERTS,
+        TOP_K,
+        1,
+        1,
+        0,
+        NUM_EXPERTS,
+        ROUTED_SCALING,
+        TILE,
+        DEEPSEEK_V3,
+    )
+    tile_expert, tile_limit, exp_to_perm = out[0], out[1], out[2]
+    scales = out[6]
+    perm = exp_to_perm.long()
+    tile = perm // TILE
+    ids = tile_expert[tile]
+    assert bool((perm < tile_limit[tile]).all()), "permuted index outside its tile"
+    return ids.to(torch.int32), scales
+
+
+def reference_routing(logits: torch.Tensor, bias_bf16: torch.Tensor):
+    """The bit-faithful PyTorch model of routed_moe's routing (secondary check)."""
+    return fx.routing_reference(logits.float(), bias_bf16)
+
+
+def weight_bits(w_fp32: torch.Tensor) -> torch.Tensor:
+    return w_fp32.contiguous().view(torch.int32)
+
+
+def native_weight_bits(w_bf16: torch.Tensor) -> torch.Tensor:
+    return w_bf16.float().contiguous().view(torch.int32)
+
+
+def compare_routing(ids, weights, n_ids, n_w) -> Optional[str]:
+    """None when ids are equal and weights are bit-identical (NaN included)."""
+    if not torch.equal(ids, n_ids):
+        bad = (ids != n_ids).any(-1).nonzero().flatten().tolist()
+        return f"ids differ in rows {bad[:8]}: got {ids[bad[0]].tolist()} native {n_ids[bad[0]].tolist()}"
+    a, b = weight_bits(weights), native_weight_bits(n_w)
+    if not torch.equal(a, b):
+        bad = (a != b).any(-1).nonzero().flatten().tolist()
+        return f"weights differ in rows {bad[:8]}: got {weights[bad[0]].tolist()} native {n_w[bad[0]].float().tolist()}"
+    if not torch.equal(weight_bits(weights.to(torch.bfloat16).float()), a):
+        return "weights are not bf16-representable"
+    return None
+
+
+def _key(logit: float, bias_val: torch.Tensor) -> torch.Tensor:
+    lg = torch.tensor([logit], dtype=torch.float32, device=bias_val.device)
+    return fx._key_of(lg, bias_val)[0]
+
+
+def _logit_for_sigmoid(s: float) -> float:
+    return float(torch.logit(torch.tensor(s, dtype=torch.float64)))
+
+
+LOW_LOGIT = -12.0
+
+
+def boundary_rows(bias_bf16: torch.Tensor, seed: int = 0) -> List[Tuple[str, torch.Tensor]]:
+    """Labelled boundary rows (each ``[256]`` FP32 on the bias's device).
+
+    Every expert not placed explicitly gets a logit near ``LOW_LOGIT`` (key
+    ~ its bias). With ``k_low`` the largest such key, the placed experts aim
+    at keys ``k_low + 0.25`` (the 8th/9th pair) and ``k_low + 0.4 + ...``
+    (seven clear winners), so only experts whose bias lets their sigmoid
+    reach those keys are placed.
+    """
+    dev = bias_bf16.device
+    b = bias_bf16.float()
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    rows: List[Tuple[str, torch.Tensor]] = []
+    low_sig = torch.sigmoid(torch.tensor(LOW_LOGIT + 0.5)).item()
+    k_low = float((b + low_sig).max().item())
+    pair_key = k_low + 0.25
+    win_key = k_low + 0.40
+    reach = [
+        e
+        for e in range(NUM_EXPERTS)
+        if 1e-3 < pair_key - float(b[e]) and win_key + 0.08 - float(b[e]) < 1 - 1e-3
+    ]
+
+    def base():
+        return (torch.rand(NUM_EXPERTS, generator=g) * 0.5 + LOW_LOGIT).to(dev)
+
+    def place_winners(row, avoid):
+        winners = [e for e in reach if e not in avoid]
+        perm = torch.randperm(len(winners), generator=g).tolist()
+        for i, j in enumerate(perm[:7]):
+            e = winners[j]
+            row[e] = _logit_for_sigmoid(win_key + 0.01 * i - float(b[e]))
+
+    if len(reach) >= 9:
+        # 8th/9th boundary pairs (e1 < e2): same warp, different warps, different hop clusters.
+        cand = sorted(reach)
+        pairs = [
+            (cand[0], cand[1]),
+            (cand[0], cand[-1]),
+            (cand[len(cand) // 2], cand[len(cand) // 2 + 1]),
+            (cand[1], cand[-2]),
+        ]
+        same_warp = [(a, c) for a, c in zip(cand, cand[1:]) if a // 32 == c // 32]
+        if same_warp:
+            pairs.append(same_warp[0])
+        cross_cluster = [(a, c) for a in cand for c in cand if a < 64 and c >= 192]
+        if cross_cluster:
+            pairs.append(cross_cluster[0])
+        for e1, e2 in pairs:
+            for offset, label in ((0, "tie"), (1, "e2+1ulp"), (-1, "e2-1ulp")):
+                row = base()
+                place_winners(row, (e1, e2))
+                l1 = _logit_for_sigmoid(pair_key - float(b[e1]))
+                k1 = _key(l1, b[e1])
+                l2 = fx.logit_for_key(float(k1.item()), b[e2], offset)
+                if l2 is None:
+                    continue
+                row[e1] = l1
+                row[e2] = l2
+                rows.append((f"{label}({e1},{e2})", row))
+        # Exact ties inside the top 8 (order among equal keys) across warps and clusters.
+        picks = [cand[(i * len(cand)) // 8] for i in range(8)]
+        if len(set(picks)) == 8:
+            row = base()
+            k = _key(_logit_for_sigmoid(pair_key - float(b[picks[0]])), b[picks[0]])
+            ok = True
+            for e in picks:
+                lg = fx.logit_for_key(float(k.item()), b[e])
+                if lg is None:
+                    ok = False
+                    break
+                row[e] = lg
+            if ok:
+                rows.append(("tie-top8", row))
+
+    row = base()
+    row[77] = float("nan")
+    rows.append(("nan1", row))
+    row = base()
+    row[200] = float("nan")
+    row[13] = float("nan")
+    rows.append(("nan2", row))
+    rows.append(("nan-all", torch.full((NUM_EXPERTS,), float("nan"), device=dev)))
+    row = base()
+    row[[1, 9, 100, 250]] = float("inf")
+    row[[2, 3]] = float("-inf")
+    rows.append(("inf", row))
+    rows.append(("neginf-all", torch.full((NUM_EXPERTS,), float("-inf"), device=dev)))
+    rows.append(("posinf-all", torch.full((NUM_EXPERTS,), float("inf"), device=dev)))
+    rows.append(("zero-all", torch.zeros(NUM_EXPERTS, device=dev)))
+    row = torch.zeros(NUM_EXPERTS, device=dev)
+    row[::2] = -0.0
+    rows.append(("negzero-mix", row))
+    row = base()
+    row[[4, 44, 144]] = 3.0e38
+    row[[5, 55]] = -3.0e38
+    row[[6, 66]] = 1.0e-42
+    row[[7, 77]] = -1.0e-42
+    rows.append(("huge-subnormal", row))
+    for scale in (0.5, 2.0, 8.0):
+        rows.append((f"randn*{scale}", (torch.randn(NUM_EXPERTS, generator=g) * scale).to(dev)))
+    return rows
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py
new file mode 100644
index 000000000000..fa011031dd45
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py
@@ -0,0 +1,245 @@
+# 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.
+"""Control flow of the hop_routing call sites in the target's decoder layer and MoE.
+
+With fake collaborators (no device work), on a decode step, checks that
+
+* the plan runs hop_routing with its closure, and not without it;
+* ``GlmDecoderLayer.forward_MoE`` runs the hop_routing hop instead of the into-hop entry's
+  exactly when the tail published and the hop_routing transport can run, hands the hop's
+  routing and logits to the MoE, and otherwise runs the into-hop entry's hop with no routing;
+* ``GlmMoE.compute_routed_output`` tries routed_moe's hop_routing fork first only when routing
+  is present (routed_moe_into_hop entry with the hop, the plain entry without), and falls back
+  to routed_moe on the same logits when the fork returns None.
+"""
+
+import os
+import sys
+import types
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
+    dispatch as _dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+    GlmDecoderLayer,
+    GlmMoE,
+)
+
+sys.path.insert(
+    0,
+    os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "..", "shared", "kernels"
+    ),
+)
+
+import plan_testlib  # noqa: E402
+
+FEATURES = (
+    "attention_output_tail,post_moe_hop,post_attention_hop,attention_output_into_hop,"
+    "routed_moe,routed_moe_into_hop,hop_routing"
+)
+
+
+@pytest.fixture(autouse=True)
+def decode_step():
+    """The owned call sites run on decode steps."""
+    with _dispatch.step(_dispatch.Phase.DECODE):
+        yield
+
+
+@pytest.fixture
+def selector(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+    plan_testlib.select(dispatch, FEATURES, monkeypatch)
+    return dispatch
+
+
+def test_selector_accepts_hop_routing(selector):
+    dispatch = selector
+    assert dispatch.is_enabled("hop_routing")
+    for base in dispatch.REQUIRED_FEATURES["hop_routing"]:
+        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
+        assert not plan.enabled("hop_routing"), base
+
+
+class _Calls:
+    def __init__(self):
+        self.log = []
+
+    def fn(self, name, result):
+        def call(*args, **kwargs):
+            self.log.append((name, args, kwargs))
+            return result
+
+        return call
+
+
+def _decoder(hop_routing_can_run, calls):
+    dec = GlmDecoderLayer.__new__(GlmDecoderLayer)
+    torch.nn.Module.__init__(dec)
+    rows = 6
+    normed = torch.zeros(rows, 4)
+    res = torch.ones(rows, 4)
+    logits = torch.full((rows, 256), 2.0)
+    routing = torch.zeros(2, rows, 8, dtype=torch.int32)
+    bias = torch.zeros(256, dtype=torch.bfloat16)
+    weight = torch.zeros(256, 4)
+
+    class _HOP_ROUTING:
+        def can_run(self, *args, **kwargs):
+            calls.log.append(("hop_routing.can_run", args, kwargs))
+            return hop_routing_can_run
+
+        def __call__(self, *args, **kwargs):
+            calls.log.append(("hop_routing", args, kwargs))
+            return normed, res, logits, routing
+
+    class _INTO_HOP:
+        def __call__(self, *args, **kwargs):
+            calls.log.append(("into_hop", args, kwargs))
+            return normed, res, logits
+
+    class _MoE:
+        gate = types.SimpleNamespace(e_score_correction_bias=bias)
+        experts = types.SimpleNamespace(has_nvfp4=True)
+
+        def __call__(self, hidden_states, hidden_states_fp4, **kwargs):
+            calls.log.append(("mlp", (hidden_states,), kwargs))
+            return hidden_states
+
+    dec.__dict__.update(
+        _hop_routing=_HOP_ROUTING(),
+        _into_hop=_INTO_HOP(),
+        mlp=_MoE(),
+        _post_attn_router_weight=lambda: weight,
+        post_attention_layernorm=types.SimpleNamespace(
+            weight=torch.ones(4), variance_epsilon=1e-5, nvfp4_scale=None
+        ),
+        fusion_config=types.SimpleNamespace(PRE_MOE_FUSION=True, POST_MOE_FUSION=False),
+        mapping=types.SimpleNamespace(tp_size=8, is_multi_node=lambda: True),
+        post_moe_hop=None,
+        next_layer_layernorm=None,
+        layer_idx=3,
+        _apply_next_layer_layernorm=lambda h, r: (h, r),
+    )
+    return dec, (normed, res, logits, routing, bias, weight)
+
+
+@pytest.mark.parametrize("hop_routing_can_run", [True, False])
+def test_forward_moe_selects_the_hop(hop_routing_can_run):
+    calls = _Calls()
+    dec, (normed, res, logits, routing, bias, weight) = _decoder(hop_routing_can_run, calls)
+    meta = types.SimpleNamespace(all_rank_num_tokens=[6])
+    GlmDecoderLayer.forward_MoE(dec, torch.zeros(6, 4), meta, res, into_hop_published=True)
+    names = [c[0] for c in calls.log]
+    assert names[0] == "hop_routing.can_run"
+    can = calls.log[0]
+    assert can[1][2] is weight and can[1][3] is bias
+    mlp = [c for c in calls.log if c[0] == "mlp"][0]
+    assert mlp[2]["router_logits"] is logits
+    if hop_routing_can_run:
+        assert "hop_routing" in names and "into_hop" not in names
+        assert mlp[2]["hop_routing_out"] is routing
+    else:
+        assert "into_hop" in names and "hop_routing" not in names
+        assert mlp[2]["hop_routing_out"] is None
+
+
+def _moe(selector, calls, hop_routing_into, hop_routing_plain, hop):
+    dispatch = selector
+    for feature, name, result in (
+        ("hop_routing", "routed_output_into_hop", hop_routing_into),
+        ("hop_routing", "routed_output", hop_routing_plain),
+        ("routed_moe", "routed_output_into_hop", ["routed_moe_into"]),
+        ("routed_moe", "routed_output", ["routed_moe_plain"]),
+    ):
+        dispatch.register(feature, name, calls.fn(f"{feature}.{name}", result))
+    moe = GlmMoE.__new__(GlmMoE)
+    torch.nn.Module.__init__(moe)
+    moe.__dict__.update(
+        use_dp=False,
+        mapping=types.SimpleNamespace(tp_size=8),
+        router_gemm=None,
+        gate=types.SimpleNamespace(e_score_correction_bias=torch.zeros(256)),
+        experts=object(),
+        _routed_into_hop=hop,
+    )
+    return moe
+
+
+@pytest.fixture
+def restore_registry():
+    import importlib
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.modeling import (  # noqa: E501
+        TARGET_PATHS,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    # Import both packages (which registers their real entries) before the snapshot: a package registers only on
+    # its first import, so entries dropped by a restore would never come back.
+    for feature in ("hop_routing", "routed_moe"):
+        importlib.import_module(TARGET_PATHS.kernel_modules[feature])
+        dispatch.get_impl(feature, "routed_output")
+    saved = {k: dict(v) for k, v in dispatch._REGISTRY.items()}
+    yield
+    dispatch._REGISTRY.clear()
+    dispatch._REGISTRY.update(saved)
+
+
+@pytest.mark.parametrize(
+    "hop, routing_given, fork_result, want",
+    [
+        ("hop", True, ["fork_into"], ["hop_routing.routed_output_into_hop"]),
+        (
+            "hop",
+            True,
+            None,
+            ["hop_routing.routed_output_into_hop", "routed_moe.routed_output_into_hop"],
+        ),
+        ("hop", False, ["fork_into"], ["routed_moe.routed_output_into_hop"]),
+        (None, True, ["fork_plain"], ["hop_routing.routed_output"]),
+        (None, True, None, ["hop_routing.routed_output", "routed_moe.routed_output"]),
+        (None, False, ["fork_plain"], ["routed_moe.routed_output"]),
+    ],
+)
+def test_compute_routed_output_order(
+    selector, restore_registry, hop, routing_given, fork_result, want
+):
+    calls = _Calls()
+    moe = _moe(selector, calls, fork_result, fork_result, object() if hop else None)
+    x = torch.zeros(6, 4)
+    logits = torch.zeros(6, 256)
+    routing = torch.zeros(2, 6, 8, dtype=torch.int32) if routing_given else None
+    out = moe.compute_routed_output(
+        x, None, None, False, router_logits=logits, hop_routing_out=routing
+    )
+    assert [c[0] for c in calls.log] == want
+    for name, args, _ in calls.log:
+        assert args[3] is logits  # every entry receives the hop's logits
+        if name.startswith("hop_routing."):
+            assert args[4] is routing
+    results = {
+        "hop_routing.routed_output_into_hop": fork_result,
+        "hop_routing.routed_output": fork_result,
+        "routed_moe.routed_output_into_hop": ["routed_moe_into"],
+        "routed_moe.routed_output": ["routed_moe_plain"],
+    }
+    assert out == results[want[-1]]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py
new file mode 100644
index 000000000000..9b0329202532
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py
@@ -0,0 +1,234 @@
+# 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.
+"""The hop_routing hop on one GPU, over a simulated transport (hop_routing).
+
+Like the into-hop entry, the hop_routing hop reads every rank's partial from the local
+(unicast) view of the post_attention_hop workspace only, so a plain device buffer stands in
+for it: the test writes all eight ranks' sanitized partials into the current
+Lamport buffer, sets the protocol state, and runs the hop as rank 0 of 8.
+This needs no MNNVL and runs under every compute-sanitizer tool. Checks:
+
+* normed, residual and router logits are bit-identical to the into-hop op's on the
+  same transport (R = 1, 3, 6, 8, every ``router_load``, both poll modes);
+* ids and weights equal the native TRT-LLM-Gen routing runner on those
+  logits (ids equal, weight bits equal);
+* the gather protocol: the previous parity's buffer is reset to the sentinel
+  (it holds junk before the call), this call's buffer holds a score for
+  every live row and the sentinel elsewhere, the parity flips, the post_attention_hop state
+  advances as the into-hop entry's;
+* boundary rows through the ``logits_in`` instance;
+* CUDA-graph capture and replay with fresh partials and alternating parity.
+
+Partials, weights and bias are synthetic (seeded).
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(
+    0, os.path.join(os.path.dirname(_HERE), "..", "..", "shared", "kernels", "post_attention_hop")
+)
+import hop_routing_jit  # noqa: E402
+import hop_routing_oracle as orc  # noqa: E402
+import into_hop_jit  # noqa: E402
+
+H = 6144
+RANKS = 8
+MAX_ROWS = 8
+NUM_BUFFERS = 3
+BUFFER_BYTES = MAX_ROWS * RANKS * H * 2
+WORKSPACE_BYTES = BUFFER_BYTES * NUM_BUFFERS
+SENTINEL = -(2**31)  # 0x80000000 as int32
+GATHER_WORDS = 2 * MAX_ROWS * 256
+EPS = 1e-5
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="the hop_routing hop targets SM100/SM103",
+    ),
+    pytest.mark.threadleak(enabled=False),
+]
+
+_OPS = {}
+
+
+def _ops():
+    if not _OPS:
+        root = os.environ.get("GLM_MOE_DSA_HOP_ROUTING_JIT_ROOT", "/tmp/hop_routing_jit")
+        into_hop_jit.ensure_op(root, True, lambda: None)
+        hop_routing_jit.ensure_op(root, True, lambda: None)
+        _OPS["into_hop"] = torch.ops.trtllm.glm_moe_dsa_post_attention_hop_from_partials
+        _OPS["hop_routing"] = getattr(torch.ops.trtllm, hop_routing_jit.OP_NAME)
+    return _OPS
+
+
+def _sanitized(x):
+    bits = x.contiguous().view(torch.int16).clone()
+    bits[bits == -32768] = 0  # bf16 -0.0 -> +0.0
+    return bits.view(torch.bfloat16)
+
+
+class Transport:
+    """A device buffer laid out like the post_attention_hop workspace, its state and the hop_routing gather state."""
+
+    def __init__(self):
+        self.uc = torch.full((WORKSPACE_BYTES // 4,), SENTINEL, dtype=torch.int32, device="cuda")
+        self.flags = torch.zeros(4, dtype=torch.int32, device="cuda")
+        self.gather = torch.full((GATHER_WORDS,), SENTINEL, dtype=torch.int32, device="cuda")
+        self.zstate = torch.zeros(4, dtype=torch.int32, device="cuda")
+
+    def buffers(self):
+        return self.uc.view(torch.bfloat16).view(NUM_BUFFERS, MAX_ROWS, RANKS, H)
+
+    def publish(self, partials, cur, prev_rows):
+        prev = (cur + NUM_BUFFERS - 1) % NUM_BUFFERS
+        self.uc.fill_(SENTINEL)
+        bufs = self.buffers()
+        bufs[prev, :prev_rows].fill_(1.5)
+        for s, part in enumerate(partials):
+            bufs[cur, : part.shape[0], s] = _sanitized(part)
+        self.flags.copy_(torch.tensor([cur, prev_rows * RANKS * H * 2, 0, 0], dtype=torch.int32))
+
+
+def _inputs(rows, seed):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    partials = [
+        (torch.randn(rows, H, generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
+        for _ in range(RANKS)
+    ]
+    residual = torch.randn(rows, H, generator=gen, device="cuda").to(torch.bfloat16)
+    gamma = (1 + 0.1 * torch.randn(H, generator=gen, device="cuda")).to(torch.bfloat16)
+    weight = (
+        (torch.randn(256, H, generator=gen, device="cuda") * 0.02).to(torch.bfloat16).contiguous()
+    )
+    bias = (torch.randn(256, generator=gen, device="cuda") * 0.05).to(torch.bfloat16).contiguous()
+    return partials, residual, gamma, weight, bias
+
+
+def _into_hop(tr, residual, gamma, weight, early_poll, router_load):
+    return _ops()["into_hop"](
+        residual, gamma, EPS, weight, tr.uc, tr.flags, 0, RANKS, early_poll, router_load
+    )
+
+
+def _hop_routing(
+    tr, residual, gamma, weight, bias, early_poll=False, router_load=1, logits_in=None
+):
+    return _ops()["hop_routing"](
+        residual,
+        gamma,
+        EPS,
+        weight,
+        bias,
+        tr.uc,
+        tr.flags,
+        tr.gather,
+        tr.zstate,
+        0,
+        RANKS,
+        early_poll,
+        router_load,
+        2.5,
+        logits_in,
+        None,
+    )
+
+
+def _check_routing(routing, logits, bias, rows):
+    ids, w = routing[0][:rows], routing[1][:rows].view(torch.float32)
+    n_ids, n_w = orc.native_routing(logits[:rows], bias)
+    err = orc.compare_routing(ids, w, n_ids, n_w)
+    assert err is None, err
+
+
+@pytest.mark.parametrize("rows", [1, 3, 6, 8])
+def test_identity_routing_and_protocol(rows):
+    partials, residual, gamma, weight, bias = _inputs(rows, 100 + rows)
+    tr = Transport()
+    for early_poll in (False, True):
+        for router_load in (0, 1, 2):
+            cur, prev_rows = (early_poll + 2 * router_load) % 3, 8 - rows + 1
+            tr.publish(partials, cur, prev_rows)
+            ref = _into_hop(tr, residual, gamma, weight, early_poll, router_load)
+            torch.cuda.synchronize()
+            tr.publish(partials, cur, prev_rows)
+            parity = int(tr.zstate[0].item()) & 1
+            gview = tr.gather.view(2, MAX_ROWS, 256)
+            gview[parity ^ 1].fill_(12345)  # junk in the previous call's buffer
+            got = _hop_routing(tr, residual, gamma, weight, bias, early_poll, router_load)
+            torch.cuda.synchronize()
+            for a, b in zip(got[:3], ref):
+                assert torch.equal(a, b)
+            _check_routing(got[3], got[2], bias, rows)
+            assert bool((gview[parity ^ 1] == SENTINEL).all()), "previous gather buffer not reset"
+            assert bool((gview[parity, :rows] != SENTINEL).all()), "missing scores"
+            assert bool((gview[parity, rows:] == SENTINEL).all()), "scores outside the call's rows"
+            assert int(tr.zstate[0].item()) == parity ^ 1
+            want = [(cur + 1) % NUM_BUFFERS, rows * RANKS * H * 2, 0, 0]
+            assert tr.flags.tolist() == want, (tr.flags.tolist(), want)
+            # This call's buffer must be all-sentinel again before its parity comes back.
+            gview[parity].fill_(SENTINEL)
+
+
+def test_boundary_rows_logits_in():
+    partials, residual, gamma, weight, bias = _inputs(8, 7)
+    rows_all = orc.boundary_rows(bias, seed=3)
+    mat = torch.stack([r for _, r in rows_all]).float()
+    tr = Transport()
+    cur = 0
+    for start in range(0, mat.shape[0], 8):
+        chunk = mat[start : start + 8].contiguous()
+        n = chunk.shape[0]
+        tr.publish([p[:n] for p in partials], cur, n)
+        got = _hop_routing(tr, residual[:n].contiguous(), gamma, weight, bias, logits_in=chunk)
+        torch.cuda.synchronize()
+        tr.gather.view(2, MAX_ROWS, 256)[int(tr.zstate[0].item()) ^ 1].fill_(SENTINEL)
+        _check_routing(got[3], chunk, bias, n)
+        cur = (cur + 1) % NUM_BUFFERS
+
+
+def test_graph_replay_alternates_parity():
+    rows = 6
+    partials, residual, gamma, weight, bias = _inputs(rows, 11)
+    tr = Transport()
+    tr.publish(partials, 0, rows)
+    _hop_routing(tr, residual, gamma, weight, bias)
+    torch.cuda.synchronize()
+    tr.publish(partials, 1, rows)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = _hop_routing(tr, residual, gamma, weight, bias)
+    for rep in range(4):
+        fresh, res_new, _, _, _ = _inputs(rows, 300 + rep)
+        residual.copy_(res_new)
+        parity = int(tr.zstate[0].item()) & 1
+        tr.publish(fresh, rep % NUM_BUFFERS, rows)
+        graph.replay()
+        torch.cuda.synchronize()
+        assert int(tr.zstate[0].item()) == parity ^ 1
+        got = [t.clone() for t in captured]
+        tr.publish(fresh, rep % NUM_BUFFERS, rows)
+        ref = _into_hop(tr, residual, gamma, weight, False, 1)
+        torch.cuda.synchronize()
+        for a, b in zip(got[:3], ref):
+            assert torch.equal(a, b)
+        _check_routing(got[3], got[2], bias, rows)
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py
new file mode 100644
index 000000000000..a5df11349cb7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py
@@ -0,0 +1,200 @@
+# 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.
+"""The hop_routing routed_moe fork against routed_moe on one GPU (real GLM-5.2 NVFP4 weights).
+
+routed_moe computes the routing from the logits; the fork reads the same routing
+(the native runner's ids and weights for those logits, packed as the hop_routing hop
+writes them). Given that, every output must be routed_moe's bit for bit:
+
+* the routing tables the kernel builds (``DEBUG_ROUTING`` dump: top-k ids
+  and weights, slot experts, slot masks, slot count);
+* the raw tuple (FC2 rows and expert scales) and the finalized output;
+* the routed_moe_into_hop FP32 accumulator within the FP32 reorder bound (routed_moe's own output
+  class; the atomic adds are run-to-run nondeterministic), and identical to
+  routed_moe's from the same rows;
+* padded rows leave the live rows unchanged; CUDA-graph replay with fresh
+  routing; host-only declines (wrong routing shape or dtype) before any
+  device work.
+
+Inputs: real captured logits and post-norm hidden states of layer 10
+(``GLM_MOE_DSA_REAL_MOE_INPUTS_DIR``), plus random activations with the real router GEMM.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
+import hop_routing_oracle as orc  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+DEBUG_ROUTING = 1
+REORDER_REL = 2.0**-20
+LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
+
+
+@pytest.fixture(scope="module")
+def layer():
+    torch.manual_seed(0)
+    return fx.build_native_layer(
+        layer=LAYER, tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
+    )
+
+
+@pytest.fixture(scope="module")
+def runners(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing.routed_moe import (  # noqa: E501
+        HopRoutedMoe,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    return RoutedMoe.from_native(layer.module, layer.bias), HopRoutedMoe(layer.module, layer.bias)
+
+
+def _pack(logits, bias):
+    ids, w = orc.native_routing(logits, bias)
+    rows = ids.shape[0]
+    routing = torch.empty(2, rows, 8, dtype=torch.int32, device="cuda")
+    routing[0].copy_(ids)
+    routing[1].copy_(w.float().view(torch.int32))
+    return routing
+
+
+def _real(layer, rows, count, offset=0):
+    win = fx.real_windows(fx.load_real_capture(layer.layer), rows, stride=rows * 7)
+    out = []
+    for x, lg in win[offset : offset + count]:
+        out.append((x.cuda().contiguous(), lg.float().cuda().contiguous()))
+    return out
+
+
+def _random(layer, rows, seed):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+    return x, layer.router_logits(x).float().contiguous()
+
+
+def _cases(layer, rows):
+    return _real(layer, rows, 3, offset=rows) + [_random(layer, rows, 40 + rows)]
+
+
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+def test_routing_tables_raw_and_finalized_equal_routed_moe(layer, runners, rows):
+    routed_runner, fork = runners
+    for i, (x, logits) in enumerate(_cases(layer, rows)):
+        routing = _pack(logits, layer.bias)
+        dbg_a = torch.zeros(2 * 1024 * 1024, dtype=torch.int32, device="cuda")
+        dbg_b = torch.zeros_like(dbg_a)
+        routed_runner(x, logits, debug=DEBUG_ROUTING, debug_buffer=dbg_a)
+        fork(x, routing, debug=DEBUG_ROUTING, debug_buffer=dbg_b)
+        raw_a = routed_runner(x, logits, do_finalize=False)
+        raw_b = fork(x, routing, do_finalize=False)
+        fin_a = routed_runner(x, logits)
+        fin_b = fork(x, routing)
+        torch.cuda.synchronize()
+        assert torch.equal(dbg_a[:2048], dbg_b[:2048]), f"case {i}: routing tables differ"
+        for a, b in zip(raw_a, raw_b):
+            assert torch.equal(a, b), f"case {i}: raw tuple differs"
+        assert torch.equal(fin_a, fin_b), f"case {i}: finalized output differs"
+
+
+def test_fp32_accumulator_matches_routed_moe_class(layer, runners):
+    routed_runner, fork = runners
+    rows = 6
+    if fork.schedule(rows, False, True) != "xc":
+        pytest.skip("FP32 accumulator output needs the expert-cluster schedule")
+    for x, logits in _real(layer, rows, 2, offset=3):
+        routing = _pack(logits, layer.bias)
+        fc2, scale, idx = routed_runner(x, logits, do_finalize=False)
+        terms = scale.double()[:, :, None] * fc2[idx.long()].double()
+        exact = terms.sum(1)
+        bound = REORDER_REL * terms.abs().sum(1)
+        for _ in range(4):
+            acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
+            fork(x, routing, red_out=acc)
+            torch.cuda.synchronize()
+            assert bool(((acc.double() - exact).abs() <= bound).all())
+
+
+@pytest.mark.parametrize("live", [1, 5])
+def test_padded_rows_do_not_change_live_rows(layer, runners, live):
+    _, fork = runners
+    ((x, logits),) = _real(layer, 8, 1, offset=1)
+    full = fork(x, _pack(logits, layer.bias), do_finalize=False)
+    part = fork(
+        x[:live].contiguous(), _pack(logits[:live].contiguous(), layer.bias), do_finalize=False
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(full[0][: live * 8], part[0])
+    assert torch.equal(full[1][:live], part[1])
+
+
+def test_cuda_graph_replay(layer, runners):
+    _, fork = runners
+    rows = 6
+    cases = _real(layer, rows, 4, offset=2)
+    x = cases[0][0].clone()
+    routing = _pack(cases[0][1], layer.bias)
+    fork(x, routing, do_finalize=False)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = fork(x, routing, do_finalize=False)
+    for xn, ln in cases[1:]:
+        x.copy_(xn)
+        routing.copy_(_pack(ln, layer.bias))
+        graph.replay()
+        torch.cuda.synchronize()
+        got = [t.clone() for t in captured]
+        want = fork(x, routing, do_finalize=False)
+        torch.cuda.synchronize()
+        for a, b in zip(got, want):
+            assert torch.equal(a, b)
+    del graph
+
+
+def test_declines_before_device_work(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing import (  # noqa: E501
+        routed_moe as hop_routed_mod,
+    )
+
+    x, logits = _random(layer, 6, 3)
+    good = _pack(logits, layer.bias)
+    bad = [
+        None,
+        good[:, :5].contiguous(),
+        good.to(torch.int64),
+        good.transpose(1, 2).contiguous().transpose(1, 2),
+    ]
+    torch.cuda.synchronize()
+    for routing in bad:
+        before = torch.cuda.memory_allocated()
+        assert (
+            hop_routed_mod.routed_output(layer.module, layer.bias, x, logits, routing, False)
+            is None
+        )
+        assert torch.cuda.memory_allocated() == before
+    assert hop_routed_mod.routing_reason(good, x) == ""
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py
new file mode 100644
index 000000000000..438ff7dba012
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py
@@ -0,0 +1,161 @@
+# 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.
+"""Call-site invariants of the post_attention_hop in the target's layers.
+
+The post_attention_hop keeps one workspace and one device-side buffer state per TP group
+for every layer, CUDA graph and eager call, as the native MNNVL all-reduce
+does. That is correct only while every rank issues the same sequence of hop
+calls on one stream, deciding each call from rank-uniform facts. These checks
+pin the preconditions on the target's layer module:
+
+* the hop and the into-hop transport (attention_output_into_hop) are created at one place each, at
+  layer adoption, through the dispatch entry points;
+* every hop call (post_attention_hop, into-hop entry) is issued directly on the layer's stream:
+  never inside a ``with`` block (e.g. ``torch.cuda.stream``) and never inside
+  a nested function or lambda (e.g. a closure handed to
+  ``maybe_execute_in_parallel``, which runs one branch on an auxiliary
+  stream);
+* the per-call decision reads only shapes, dtypes, contiguity and the
+  rank-agreed SM count: no device-to-host transfer or tensor value;
+* dense layers keep the native post-attention path: ``forward_mlp`` never
+  reaches the hop (scope of the served post_attention_hop: MoE and MTP layers).
+"""
+
+import ast
+import inspect
+import textwrap
+
+import pytest
+
+modeling = pytest.importorskip(
+    "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers"
+)
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop import (  # noqa: E402, E501
+    hop,
+)
+
+_HOP_CALLS = ("_post_attn_hop", "_into_hop", "post_attn_hop")
+
+
+@pytest.fixture(scope="module")
+def tree():
+    return ast.parse(inspect.getsource(modeling))
+
+
+def _self_attr_calls(node, attrs):
+    return [
+        n
+        for n in ast.walk(node)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Attribute)
+        and n.func.attr in attrs
+        and isinstance(n.func.value, ast.Name)
+        and n.func.value.id == "self"
+    ]
+
+
+def _parents(tree):
+    parent = {}
+    for node in ast.walk(tree):
+        for child in ast.iter_child_nodes(node):
+            parent[child] = node
+    return parent
+
+
+def _methods(tree, cls_name):
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ClassDef) and node.name == cls_name:
+            return {n.name: n for n in node.body if isinstance(n, ast.FunctionDef)}
+    raise AssertionError(f"{cls_name} not found")
+
+
+def test_hop_created_once_through_dispatch(tree):
+    calls = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Call)
+        and isinstance(n.func.func, ast.Attribute)
+        and n.func.func.attr == "get_impl"
+        and len(n.func.args) == 2
+        and all(isinstance(a, ast.Constant) for a in n.func.args)
+    ]
+    entries = [(c.func.args[0].value, c.func.args[1].value) for c in calls]
+    assert entries.count(("post_attention_hop", "make_post_attn_hop")) == 1, entries
+    assert entries.count(("post_attention_hop", "make_into_hop_transport")) <= 1, entries
+    methods = _methods(tree, "_AdoptedLayerMixin")
+    init = methods["_adopt_pre_moe"]
+    in_init = [
+        c for c in ast.walk(init) if c in calls and c.func.args[0].value == "post_attention_hop"
+    ]
+    assert len(in_init) == len(
+        [e for e in entries if e[0] == "post_attention_hop" and e[1].startswith("make_")]
+    )
+
+
+def test_hop_calls_are_issued_on_the_layer_stream(tree):
+    parent = _parents(tree)
+    sites = []
+    for call in _self_attr_calls(tree, _HOP_CALLS):
+        node = call
+        while node in parent:
+            node = parent[node]
+            if isinstance(node, (ast.With, ast.AsyncWith)):
+                raise AssertionError(
+                    f"hop call at line {call.lineno} inside a with block: "
+                    f"{ast.unparse(node.items[0])}"
+                )
+            if isinstance(node, ast.Lambda):
+                raise AssertionError(f"hop call at line {call.lineno} inside a lambda")
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                outer = parent.get(node)
+                assert isinstance(outer, ast.ClassDef), (
+                    f"hop call at line {call.lineno} inside nested function {node.name}"
+                )
+                sites.append((outer.name, node.name))
+                break
+    # The served call sites: target MoE layers (forward_MoE) and the MTP
+    # layer; _post_attn_hop itself wraps the hop object.
+    assert ("GlmDecoderLayer", "forward_MoE") in sites, sites
+    assert ("GlmMtpLayer", "forward") in sites, sites
+
+
+def test_dense_layers_never_reach_the_hop(tree):
+    methods = _methods(tree, "GlmDecoderLayer")
+    forward_mlp = methods["forward_mlp"]
+    names = {n.attr for n in ast.walk(forward_mlp) if isinstance(n, ast.Attribute)}
+    assert not names & {"post_attn_hop", "_post_attn_hop", "_can_use_post_attn_hop", "_into_hop"}, (
+        names
+    )
+
+
+def test_call_decision_reads_no_tensor_values():
+    for fn in (hop.PostAttnHop.call_ineligibility, hop.PostAttnHop.can_run, hop.required_sms):
+        node = ast.parse(textwrap.dedent(inspect.getsource(fn)))
+        attrs = {n.attr for n in ast.walk(node) if isinstance(n, ast.Attribute)}
+        assert not attrs & {
+            "item",
+            "tolist",
+            "cpu",
+            "numpy",
+            "any",
+            "all",
+            "sum",
+            "max",
+            "min",
+            "nonzero",
+            "synchronize",
+        }, (fn.__name__, attrs)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py
new file mode 100644
index 000000000000..a7994e880c34
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py
@@ -0,0 +1,492 @@
+# 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.
+"""Collective checks of routed_moe_into_hop: routed_moe's finalize kernel -> post_moe_hop routed-input hop.
+
+Used by ``test_routed_into_hop_mpi.py``. Every function here is collective over the two-host TP8
+group (call it in the same order on every rank).
+
+A per-rank *producer* decides what a rank hands to the post_moe_hop for one call:
+
+* ``routed_moe_into_hop``: routed_moe writes into the hop's routed-input buffer ->
+  ``hop.hop_routed`` (the cluster kernel's routed-input variant of the hop's mode,
+  ``TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION``, default "fp32":
+  "row" = finalized by ``RoutedMoeFinalizeKernel``, "identity" = identity-layout raw
+  rows, "fp32" = FP32 accumulator added by the routed kernel);
+* ``routed_moe``: routed_moe raw tuple -> ``hop(...)`` (default variant);
+* ``native``: the native TRT-LLM-Gen raw tuple -> ``hop(...)``.
+
+The hop is constructed, and its routed input enabled, inside the served
+model loader's ``MetaInitMode`` (``served_construction``), as
+the target's decoder-layer adoption does in the server.
+
+Ranks may pick different producers for the same call (a rank whose routed_moe
+declines routed_moe_into_hop keeps the raw tuple); the collective order is one hop per call
+on every rank either way. The oracle of a call is the same per-rank choice
+with ``routed_moe_into_hop`` replaced by ``routed_moe``, i.e. every rank on the default (raw-tuple)
+variant. In modes "row" and "identity" routed_moe_into_hop's local partial must be
+bit-identical to the raw variant's, so the reduced outputs must be
+bit-identical too. In mode "fp32" the top-k sum order differs from the native
+finalize: the published partial must lie within the FP32 reorder bound of the
+raw variant's partial (``REORDER_REL * sum_k |w y| + ulp_bf16``), and the
+outputs must equal the native-arithmetic emulation of the published partials.
+"""
+
+import os
+import random
+import sys
+
+import torch
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(HERE), "routed_moe"))
+sys.path.insert(
+    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
+)
+import post_moe_hop_mpi_utils as U  # noqa: E402
+import routed_moe_chain_bench as CB  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+FEATURES = {
+    "routed_moe_into_hop": "post_moe_hop,routed_moe,routed_moe_into_hop",
+    "routed_moe": "routed_moe",
+    "native": "",
+}
+NORM_MAX_ULP = 2
+# eight exact FP32 products summed in another order: |delta| <= 7 * 2^-24 * sum|w y|
+REORDER_REL = 2.0**-20
+
+
+_WINDOWS = {}
+
+
+def windows(layer: int, rows: int, count: int, offset: int = 0):
+    """``count`` real hidden-state windows of the captured fixture (spread over it)."""
+    key = (layer, rows)
+    if key not in _WINDOWS:
+        _WINDOWS[key] = [x for x, _ in fx.real_windows(fx.load_real_capture(layer), rows)]
+    win = _WINDOWS[key]
+    step = max(1, len(win) // count)
+    return [win[(offset + i * step) % len(win)].cuda() for i in range(count)]
+
+
+def served_construction():
+    """The context in which the served model loader constructs the model.
+
+    ``PyTorchModelLoader`` builds the model inside ``MetaInitMode``: ``empty``
+    and ``empty_like`` produce meta tensors, and afterwards only module
+    parameters and buffers are materialized. The target's decoder-layer adoption
+    builds the post_moe_hop and enables its routed input inside that context, so
+    the harness does the same.
+    """
+    from tensorrt_llm._torch.models.modeling_utils import MetaInitMode
+
+    return MetaInitMode()
+
+
+_ROUTED_BUFFERS = {
+    "row": ("_routed_buffer",),
+    "identity": ("_routed_rows", "_routed_scales"),
+    "fp32": ("_routed_acc",),
+}
+
+
+def routed_buffers(hop):
+    """``{attribute: tensor or None}`` of the routed-input buffers of the hop's mode."""
+    return {name: getattr(hop, name, None) for name in _ROUTED_BUFFERS[hop.routed_mode]}
+
+
+def buffer_report(hop):
+    """Where each routed-input buffer of the hop's mode lives (for logs)."""
+    return {
+        name: None if t is None else f"{t.device} meta={t.is_meta} ptr={t.data_ptr():#x}"
+        for name, t in routed_buffers(hop).items()
+    }
+
+
+# MoE layers of GLM-5.2 that build a post_moe_hop (layers 3..77)
+MOE_LAYERS = 75
+
+
+def buffer_bytes(hop) -> int:
+    """Device bytes of the routed-input buffers of the hop's mode."""
+    return sum(t.numel() * t.element_size() for t in routed_buffers(hop).values() if t is not None)
+
+
+def buffer_problems(hop):
+    """Why the routed-input buffers are not usable device memory ([] when they are)."""
+    device = hop.workspace.state.device
+    problems = []
+    for name, t in routed_buffers(hop).items():
+        if t is None:
+            problems.append(f"{name} missing")
+        elif t.is_meta:
+            problems.append(f"{name} is a meta tensor (data_ptr {t.data_ptr():#x})")
+        elif t.device != device:
+            problems.append(f"{name} on {t.device}, the hop's workspace on {device}")
+        elif t.data_ptr() == 0 or t.data_ptr() % 16:
+            problems.append(f"{name} data_ptr {t.data_ptr():#x}")
+    return problems
+
+
+def build_hop(env, producer, mode=None):
+    """The post_moe_hop with routed input, constructed as the served model does.
+
+    Inside ``served_construction()``, with the layer's MoE module as the
+    declared raw-tuple producer, followed by the collective
+    ``enable_routed_input(mode)``. Returns ``(hop, enabled)``.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    with served_construction():
+        hop = PostMoeHop(env.mapping, fx.HIDDEN, fx.TOP_K, torch.bfloat16, raw_producer=producer)
+        enabled = hop.engaged and hop.enable_routed_input(mode)
+    return hop, enabled
+
+
+def build_stack(env, layers, check_buffers=True):
+    """Served MoE layers (real weights), the post_moe_hop with routed input on.
+
+    With ``check_buffers`` the routed-input buffers must be device memory on
+    every rank: the served construction must not leave them on the meta device.
+    """
+    CB.set_features("")
+    moes, gammas = CB.build_layers(env, layers)
+    hop, enabled = build_hop(env, next(iter(moes.values())).experts)
+    assert hop.engaged, hop.reason
+    assert enabled
+    if check_buffers:
+        problems = buffer_problems(hop)
+        assert U.all_ranks(not problems), f"rank {env.rank}: {problems}"
+    for moe in moes.values():
+        moe._routed_into_hop = hop
+    return moes, gammas, hop
+
+
+def produce(moe, x, producer: str):
+    """The MoE half of one call (a decode step); returns ``(kind, payload)`` without the hop."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    CB.set_features(FEATURES[producer])
+    with dispatch.step(dispatch.Phase.DECODE):
+        out = moe.forward(x, do_finalize=False)
+    if len(out) == 2:
+        assert producer == "routed_moe_into_hop", f"{producer}: unexpected [shared, routed] result"
+        return "routed", out
+    assert producer != "routed_moe_into_hop", (
+        "routed_moe_into_hop did not engage (raw tuple returned)"
+    )
+    return "raw", out
+
+
+def hop_call(hop, kind, payload, res, gamma, with_shared: bool):
+    """The post_moe_hop half of one call (collective)."""
+    if kind == "routed":
+        shared, routed = payload
+        return hop.hop_routed(routed, shared if with_shared else None, res, gamma, fx.RMS_EPS)
+    shared, fc2, scale, perm = payload
+    return hop(fc2, scale, perm, shared if with_shared else None, res, gamma, fx.RMS_EPS)
+
+
+def local_partial(kind, payload, with_shared: bool, hidden: int = fx.HIDDEN, hop=None):
+    """The rank partial the hop publishes: ``sanitize(bf16(routed + shared))``."""
+    if kind == "routed":
+        shared, routed = payload
+        if routed.dtype == torch.float32:
+            routed = routed.to(torch.bfloat16)
+        elif routed.shape[0] != shared.shape[0]:
+            # identity-layout raw rows; the routing weights are in the hop
+            rows = shared.shape[0]
+            ident = torch.arange(rows * fx.TOP_K, dtype=torch.int32, device=routed.device)
+            routed = U.reference_local_partial(
+                routed, hop._routed_scales[:rows], ident.view(rows, fx.TOP_K), None, hidden
+            )
+        local = routed if not with_shared else shared + routed
+    else:
+        shared, fc2, scale, perm = payload
+        local = U.reference_local_partial(fc2, scale, perm, shared if with_shared else None, hidden)
+    return U.sanitize_bf16(local)
+
+
+def call(moe, hop, x, res, gamma, producer: str, with_shared: bool):
+    """One call; returns ``(outputs, kind, payload)``.
+
+    In mode "fp32" the returned payload holds a copy of the accumulator taken
+    before the hop read and cleared the hop's own buffer.
+    """
+    kind, payload = produce(moe, x, producer)
+    seen = payload
+    if kind == "routed" and payload[1].dtype == torch.float32:
+        seen = [payload[0], payload[1].clone()]
+    out = hop_call(hop, kind, payload, res, gamma, with_shared)
+    return out, kind, seen
+
+
+def exact_mode(hop) -> bool:
+    """Whether routed_moe_into_hop's partial must equal the raw variant's bit for bit."""
+    return getattr(hop, "routed_mode", "row") != "fp32"
+
+
+def reorder_ok(local, ref_payload, with_shared: bool):
+    """FP32-reorder bound of a routed_moe_into_hop "fp32" partial against the raw variant's."""
+    shared, fc2, scale, perm = ref_payload
+    ref = local_partial("raw", ref_payload, with_shared)
+    routed_ref = local_partial("raw", ref_payload, False)
+    idx = perm.to(torch.long).clamp(min=0)
+    rows = fc2[idx.view(-1)].view(perm.shape[0], perm.shape[1], -1).float()
+    mag = (scale.float()[:, :, None] * rows).abs().sum(1)
+    # one BF16 rounding of the routed sum, one of the shared add (|x| 2^-7 >= ulp(x))
+    ulp = (ref.float().abs() + routed_ref.float().abs()).clamp(min=2.0**-126) * 2.0**-7
+    bound = REORDER_REL * mag + ulp
+    diff = (local.float() - ref.float()).abs()
+    return bool((diff <= bound).all()), int((diff > 0).sum())
+
+
+def oracle_producer(producer: str) -> str:
+    return "routed_moe" if producer == "routed_moe_into_hop" else producer
+
+
+def check_emulation(env, hop, kind, payload, with_shared, res, gamma, out, res_out):
+    """Published slots == gathered expected partials; outputs == native emulation.
+
+    Collective (gathers every rank's expected partial). Must directly follow
+    the checked call (reads the workspace's last buffer).
+    """
+    rows = res.shape[0]
+    mine = local_partial(kind, payload, with_shared, hop=hop)
+    expected = U.gather_bf16(mine).permute(1, 0, 2).cuda()
+    published = U.published_partials(hop, rows)
+    slots_exact = torch.equal(published.view(torch.int16), expected.view(torch.int16))
+    emu_norm, emu_res = U.emulate_post_reduction(expected, res, gamma, fx.RMS_EPS)
+    res_exact = torch.equal(res_out.view(torch.int16), emu_res.view(torch.int16))
+    norm_ulp = U.compare_stats(out, emu_norm)["max_ulp"]
+    return slots_exact and res_exact and norm_ulp <= NORM_MAX_ULP, {
+        "slots_exact": slots_exact,
+        "res_exact": res_exact,
+        "norm_max_ulp": norm_ulp,
+    }
+
+
+def same_or_close(hop, got, ref) -> bool:
+    """Bit-equal in the exact modes, close in mode "fp32".
+
+    In mode "fp32" the routed_moe_into_hop ranks' top-k sums are reordered, so the reduced
+    outputs differ by the BF16 rounding of the reordered partials (bounded
+    exactly per call by ``run_mixed``; here a small bound on residual and norm).
+    """
+    if exact_mode(hop):
+        return identical(got, ref)
+    ok = True
+    for a, b in zip(got, ref):
+        d = (a.float() - b.float()).abs()
+        ok &= bool((d <= 2.0**-6 * b.float().abs() + 2.0**-10).all())
+    return ok
+
+
+def identical(a, b) -> bool:
+    return torch.equal(a[0].view(torch.int16), b[0].view(torch.int16)) and torch.equal(
+        a[1].view(torch.int16), b[1].view(torch.int16)
+    )
+
+
+def warm(env, moes, rows_list, producers=("routed_moe_into_hop", "routed_moe", "native")):
+    """Compile every producer variant on every rank before any collective call.
+
+    A rank that compiles (seconds) while its peers already poll inside the post_moe_hop
+    kernel would make them hit the spin limit; this runs only the MoE half.
+    """
+    for moe in moes.values():
+        for rows in rows_list:
+            x = torch.zeros(rows, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
+            for producer in producers:
+                produce(moe, x, producer)
+    CB.set_features("")
+    torch.cuda.synchronize()
+    env.comm.barrier()
+
+
+def mixed_pattern(env, pattern: str):
+    """Per-rank producer of a mixed-variant call."""
+    if pattern == "odd-routed_moe":
+        return "routed_moe" if env.rank % 2 else "routed_moe_into_hop"
+    if pattern == "rank0-routed_moe":
+        return "routed_moe" if env.rank == 0 else "routed_moe_into_hop"
+    if pattern == "host1-routed_moe":
+        return "routed_moe" if env.rank >= env.world // 2 else "routed_moe_into_hop"
+    if pattern == "odd-native":
+        return "native" if env.rank % 2 else "routed_moe_into_hop"
+    if pattern == "all-routed_moe_into_hop":
+        return "routed_moe_into_hop"
+    raise ValueError(pattern)
+
+
+def run_mixed(env, moes, gammas, hop, pattern, with_shared, rows_list, per_rows=3, seed=0):
+    """Mixed-variant calls vs their all-raw-variant oracle; returns (ok, n, failures)."""
+    producer = mixed_pattern(env, pattern)
+    g = torch.Generator(device="cuda").manual_seed(900 + seed)
+    failures, n = [], 0
+    for layer in moes:
+        for rows in rows_list:
+            for i, x in enumerate(windows(layer, rows, per_rows, offset=seed)):
+                res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
+                got, kind, payload = call(
+                    moes[layer], hop, x, res, gammas[layer], producer, with_shared
+                )
+                torch.cuda.synchronize()
+                emu_ok, stats = check_emulation(
+                    env, hop, kind, payload, with_shared, res, gammas[layer], *got
+                )
+                mine = local_partial(kind, payload, with_shared, hop=hop)
+                ref, ref_kind, ref_payload = call(
+                    moes[layer], hop, x, res, gammas[layer], oracle_producer(producer), with_shared
+                )
+                torch.cuda.synchronize()
+                if kind == "routed" and not exact_mode(hop):
+                    same, n_diff = reorder_ok(mine, ref_payload, with_shared)
+                    stats["reorder_elems"] = n_diff
+                else:
+                    same = identical(got, ref)
+                if not (same and emu_ok):
+                    failures.append((layer, rows, i, stats, same))
+                n += 1
+    CB.set_features("")
+    return not failures, n, failures
+
+
+def run_stress(
+    env,
+    moes,
+    gammas,
+    hop,
+    batches: int,
+    batch: int,
+    max_delay_us: float,
+    rows_list=(1, 6, 8),
+    seed: int = 0,
+):
+    """Back-to-back batches with random per-rank producers and device delays.
+
+    No host synchronization separates the calls of a batch; every call is
+    checked bit for bit against its all-raw-variant oracle afterwards. The row
+    count, layer and window of a call are drawn from a generator seeded
+    identically on every rank (collective shapes must agree); the producer,
+    the shared-output choice and the delay are drawn per rank. Kernel-internal
+    delays come from ``routed_moe.DELAY_INJECTION`` if the caller set it (and compiled them
+    with ``warm``).
+    """
+    common = random.Random(31 + seed)
+    local = random.Random(1000 + 17 * seed + env.rank)
+    cycles_per_us = 2000.0  # approximate sm_103 SM clock; only the magnitude matters
+    layers = list(moes)
+    gen = torch.Generator(device="cuda").manual_seed(4000 + seed)
+    failures, calls, delayed, producers_seen = [], 0, 0, set()
+    for b in range(batches):
+        plan = []
+        for j in range(batch):
+            layer = layers[common.randrange(len(layers))]
+            rows = rows_list[common.randrange(len(rows_list))]
+            x = windows(layer, rows, 1, offset=common.randrange(1 << 16))[0]
+            res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=gen) * 4.0).bfloat16()
+            producer = ("routed_moe_into_hop", "routed_moe_into_hop", "routed_moe", "native")[
+                local.randrange(4)
+            ]
+            with_shared = local.random() < 0.75
+            delay = local.uniform(0.0, max_delay_us) if local.random() < 0.3 else 0.0
+            plan.append((layer, x, res, producer, with_shared, delay))
+            producers_seen.add(producer)
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        outs = []
+        for layer, x, res, producer, with_shared, delay in plan:
+            if delay:
+                torch.cuda._sleep(int(delay * cycles_per_us))
+                delayed += 1
+            got, _, _ = call(moes[layer], hop, x, res, gammas[layer], producer, with_shared)
+            outs.append(got)
+        torch.cuda.synchronize()
+        for j, ((layer, x, res, producer, with_shared, _), got) in enumerate(zip(plan, outs)):
+            ref, _, _ = call(
+                moes[layer], hop, x, res, gammas[layer], oracle_producer(producer), with_shared
+            )
+            torch.cuda.synchronize()
+            if not same_or_close(hop, got, ref):
+                failures.append((b, j, layer, x.shape[0], producer, with_shared))
+        calls += len(plan)
+    CB.set_features("")
+    bits = hop.error_bits()
+    return not failures and bits == 0, {
+        "calls": calls,
+        "delayed_on_rank": delayed,
+        "producers_on_rank": sorted(producers_seen),
+        "error_bits": bits,
+        "failures": failures[:6],
+    }
+
+
+def run_graph(env, moes, gammas, hop, pattern: str, replays: int = 4, rows: int = 6):
+    """A multi-stream CUDA graph (as served) of four mixed-producer calls.
+
+    Per-rank producers come from ``pattern``, the shared output alternates;
+    the graph is replayed on fresh real windows and every replay is checked
+    bit for bit against the eager all-raw-variant oracle. Returns
+    ``(ok, replays_checked)``.
+    """
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    layers = list(moes)
+    producer = mixed_pattern(env, pattern)
+    plan = [(layers[i % len(layers)], i % 2 == 0) for i in range(4)]
+    wins = {layer: windows(layer, rows, replays + 1, offset=11) for layer in layers}
+    xs = [wins[layer][0].clone() for layer, _ in plan]
+    g = torch.Generator(device="cuda").manual_seed(77)
+    ress = [
+        (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16() for _ in plan
+    ]
+
+    def body():
+        return [
+            call(moes[layer], hop, x, r, gammas[layer], producer, sh)[0]
+            for (layer, sh), x, r in zip(plan, xs, ress)
+        ]
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            body()
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    graph = torch.cuda.CUDAGraph()
+    with with_multi_stream(True):
+        with torch.cuda.graph(graph, stream=stream):
+            outs = body()
+    torch.cuda.synchronize()
+    ok = True
+    for rep in range(1, replays + 1):
+        for (layer, _), x in zip(plan, xs):
+            x.copy_(wins[layer][rep])
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        graph.replay()
+        torch.cuda.synchronize()
+        for (layer, sh), x, r, got in zip(plan, xs, ress, outs):
+            ref, _, _ = call(moes[layer], hop, x, r, gammas[layer], oracle_producer(producer), sh)
+            torch.cuda.synchronize()
+            ok &= same_or_close(hop, got, ref)
+    CB.set_features("")
+    del graph
+    return ok and hop.error_bits() == 0, replays
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py
new file mode 100644
index 000000000000..7b2677e9845e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py
@@ -0,0 +1,764 @@
+# 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.
+"""Real-neighbour benchmark of the routed MoE inside the served MoE block.
+
+Runs one process per GPU of the two-host TP8 group (``srun --mpi=pmix``, 8
+tasks). Every call is one target MoE layer of the served graph, built from the
+model's own modules with real GLM-5.2 NVFP4 weights:
+
+    router GEMM (main stream) || BF16 shared expert (aux stream)
+    -> routed MoE (native TRT-LLM-Gen runner, or routed_moe through the target's MoE)
+    -> stream join -> post-MoE hop
+
+The MoE is the built-in ``Deepseekv3MoE`` adopted by the target's ``GlmMoE``; its forward
+runs inside a decode step, with multi-stream on (as the CUDA-graph runner captures it) and
+``TRTLLM_ENABLE_PDL`` as served. The post-MoE hop is
+
+* ``fin`` (post_moe_hop off): ``shared + routed`` inside ``Deepseekv3MoE.forward``, then
+  the native MNNVL ``AllReduce(RESIDUAL_RMS_NORM)`` with the next layer's norm;
+* ``raw`` (post_moe_hop on): ``forward(do_finalize=False)`` and the post_moe_hop on the raw
+  tuple, exactly as the target's decoder layer calls it in ``forward_MoE``.
+
+Each call's hop is the predecessor of the next call's router GEMM. Calls
+rotate over several layers' weights (``--layers``), so a call never finds its
+experts in L2. One CUDA graph holds ``--calls`` calls; a trial replays it
+``--replays`` times after an MPI barrier; per-call time = trial time /
+(replays * calls); the summary is the median trial of the slowest rank.
+
+Inputs: ``--mode real`` uses the captured post-norm hidden states (real
+routing through the real router GEMM, consecutive-token windows of R rows,
+drawn from every position of the captured prompts);
+``--mode sweep`` replaces the router GEMM's result (the GEMM still runs) by
+constructed logits that route to exactly D distinct experts; ``--mode mix``
+draws every call's D from a served distribution (``--mix``, see ``MIXES``)
+and constructs its logits the same way, with fresh experts per call.
+
+Router-logit producer (``--logits``):
+
+* ``gate``: the MoE's own router GEMM runs on the main stream before the
+  routed MoE (native, router_gemm/W2 off). The shared expert starts during that GEMM.
+* ``supplied``: the router contract with logits present (post_attention_hop/router_gemm on): the
+  pre-MoE hop already produced the FP32 logits, the MoE's gate is skipped (it returns the
+  call's logits without launching anything), so the routed MoE and the
+  shared expert become ready at the same moment. Config labels get a
+  ``+sup`` suffix. The logits are the real router GEMM's output for the
+  call's hidden states (computed once before capture), or the sweep logits.
+
+Optional served neighbours before every call (identical in every arm):
+``--attn-mb N`` streams N MB of weights (a reduction over a rotating buffer,
+the stand-in for the attention block's weight reads, which also flushes L2
+as the served graph does), then ``--pre-hop ar`` runs the native two-host
+``AllReduce(RESIDUAL_RMS_NORM, trigger_completion_at_end=False)`` on a
+rank-local partial (the stand-in for the pre-MoE hop that precedes the MoE
+block; with ``--logits supplied`` it plays post_attention_hop's hop, which emits the logits).
+"""
+
+import argparse
+import json
+import os
+import statistics
+import sys
+import time
+
+import torch
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+sys.path.insert(
+    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
+)
+sys.path.insert(0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels"))
+import plan_testlib  # noqa: E402
+import post_moe_hop_mpi_utils as U  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+ACCUMULATION_ENV = "TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION"
+# routed_moe_into_hop routed-input modes (see routed_moe.routed_moe.routed_output_into_hop); the
+# routed kernels run at fixed launch parameters, so the routed-input mode is the only arm knob.
+ARM_ENV = {
+    "routed_moe_into_hop_I": {ACCUMULATION_ENV: "identity"},
+    "routed_moe_into_hop_F": {ACCUMULATION_ENV: "fp32"},
+    "routed_moe_into_hop_R": {ACCUMULATION_ENV: "row"},
+}
+ARM_KNOBS = (ACCUMULATION_ENV,)
+# Generic arm syntax ``routed_moe[K=V+K=V]`` / ``routed_moe_into_hop[...]`` with these short keys.
+ARM_KEYS = {"MODE": ACCUMULATION_ENV}
+# Served distinct-expert distributions at R = 6, from a server profile of
+# the routed_moe arms (fixture routed_moe-server-profile,
+# xc_end_histogram_rank0.csv: the share of two-round XC calls, i.e. D above
+# the cluster count, at 24 / 28 / 32 / 36 clusters). Bins (lo, hi, share);
+# inside a bin D is uniform, except the open top bin, which decays by 0.7
+# per expert from 37 (the real prefill windows reach p90 38, max 47).
+MIXES = {
+    # late in a decode segment: the distribution of D over its bins
+    "late": ((17, 24, 0.058), (25, 28, 0.077), (29, 32, 0.213), (33, 36, 0.195), (37, 48, 0.457)),
+    # early in a segment: fewer large D; the split below 33 follows the late
+    # window's proportions
+    "early": ((17, 24, 0.084), (25, 28, 0.111), (29, 32, 0.308), (33, 36, 0.380), (37, 48, 0.117)),
+}
+ARMS_REPORTED = tuple(["routed_moe", "routed_moe_into_hop"] + list(ARM_ENV))
+NUM_LAYERS = 78
+
+
+class _PassAllReduce(torch.nn.Module):
+    """Pass-through stand-in for the MoE module's own AllReduce.
+
+    The decoder layer disables that all-reduce (``enable_allreduce=False``
+    under POST_MOE_FUSION), so it only ever passes its input through.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    @property
+    def output_buffer_kind(self) -> int:
+        from tensorrt_llm.bindings.internal.thop import BufferKind
+
+        return int(BufferKind.DEFAULT)
+
+    def forward(self, input, *, all_reduce_params=None):
+        assert all_reduce_params is not None and not all_reduce_params.enable_allreduce
+        return input
+
+
+def arm_env(impl: str) -> dict:
+    """Environment of an arm: ``ARM_ENV[impl]`` or the generic ``base[K=V+...]``."""
+    if "[" not in impl:
+        return dict(ARM_ENV.get(impl, {}))
+    base, _, spec = impl.partition("[")
+    assert base in ("routed_moe", "routed_moe_into_hop") and spec.endswith("]"), impl
+    env = {}
+    for item in [v for v in spec[:-1].split("+") if v]:
+        key, _, val = item.partition("=")
+        env[ARM_KEYS[key]] = val
+    return env
+
+
+def mix_distinct(name: str, calls: int, rows: int, seed: int):
+    """Per-call distinct-expert counts drawn from ``MIXES[name]``."""
+    values, weights = [], []
+    for lo, hi, share in MIXES[name]:
+        if hi == fx.TOP_K * rows and lo > 32:
+            span = [0.7 ** (d - lo) for d in range(lo, hi + 1)]
+        else:
+            span = [1.0] * (hi - lo + 1)
+        total = sum(span)
+        for d, w in zip(range(lo, hi + 1), span):
+            values.append(d)
+            weights.append(share * w / total)
+    g = torch.Generator().manual_seed(seed)
+    idx = torch.multinomial(torch.tensor(weights), calls, replacement=True, generator=g)
+    return [values[i] for i in idx.tolist()]
+
+
+def mix_logits(moes, order, rows: int, distinct_per_call, seed: int):
+    """Constructed logits routing call ``i`` to ``distinct_per_call[i]`` experts,
+    drawn fresh per call (the attention stand-in evicts L2 between calls)."""
+    out = []
+    for i, (layer, d) in enumerate(zip(order, distinct_per_call)):
+        sets = fx.expert_sets_with_distinct(rows, d, seed=seed + i)
+        perm = torch.randperm(
+            fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(seed + 7919 * i)
+        ).tolist()
+        uniq = sorted(set(e for row in sets for e in row))
+        remap = {e: perm[j] for j, e in enumerate(uniq)}
+        sets = [[remap[e] for e in row] for row in sets]
+        bias = moes[layer].gate.e_score_correction_bias.to(torch.bfloat16)
+        lg = fx.logits_for_expert_sets(sets, bias, seed=seed + i)
+        ids, _ = fx.routing_reference(lg, bias)
+        assert len(set(ids.flatten().tolist())) == d
+        out.append(lg)
+    return out
+
+
+def distinct_stats(moes, order, logits_per_call):
+    """Distinct experts per call (native routing of the logits the MoE reads)."""
+    ds = []
+    for layer, lg in zip(order, logits_per_call):
+        bias = moes[layer].gate.e_score_correction_bias.to(torch.bfloat16)
+        ids, _ = fx.routing_reference(lg.float().cpu(), bias.cpu())
+        ds.append(len(set(ids.flatten().tolist())))
+    n = len(ds)
+    above = {c: round(100.0 * sum(d > c for d in ds) / n, 1) for c in (24, 26, 28, 32, 36)}
+    return {"mean": round(sum(ds) / n, 2), "median": sorted(ds)[n // 2], "pct_above": above}
+
+
+def set_features(ids: str) -> None:
+    """Install, for this process, the plan that runs exactly the fast paths ``ids`` (comma-separated)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    plan_testlib.select(dispatch, ids)
+
+
+def next_norm_weight(layer: int) -> torch.Tensor:
+    name = (
+        "model.norm.weight"
+        if layer + 1 >= NUM_LAYERS
+        else f"model.layers.{layer + 1}.input_layernorm.weight"
+    )
+    return fx._Reader().full(name).to(torch.bfloat16).cuda()
+
+
+class Engagement:
+    """Counts routed_moe adapter calls that engaged (returned a result)."""
+
+    def __init__(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+        dispatch.get_impl("routed_moe", "routed_output")  # import + register
+        self.counts = {}
+        for name in ("routed_output", "routed_output_into_hop"):
+            inner = dispatch._REGISTRY["routed_moe"].get(name)
+            if inner is None:
+                continue
+
+            def counted(*args, _inner=inner, _name=name, **kwargs):
+                out = _inner(*args, **kwargs)
+                key = (_name, out is not None)
+                self.counts[key] = self.counts.get(key, 0) + 1
+                return out
+
+            dispatch.register("routed_moe", name, counted)
+
+    def engaged(self, name: str) -> int:
+        return self.counts.get((name, True), 0)
+
+
+class Neighbours:
+    """Served kernels that precede each MoE block (identical in every arm).
+
+    ``attn_mb``: a weight-streaming stand-in for the attention block (one
+    reduction over ``attn_mb`` MB, two buffers used alternately so the reads
+    miss L2 as the served attention weights do). ``pre_hop``: the native
+    two-host fused all-reduce + residual + RMSNorm on a rank-local partial.
+    """
+
+    def __init__(self, env, rows: int, calls: int, attn_mb: int, pre_hop: str):
+        g = torch.Generator(device="cuda").manual_seed(500 + env.rank)
+        self.attn = []
+        if attn_mb > 0:
+            numel = attn_mb * (1 << 20) // 2
+            self.attn = [
+                torch.randn(numel, device="cuda", generator=g).bfloat16() for _ in range(2)
+            ]
+        self.pre = None
+        if pre_hop == "ar":
+            common = torch.Generator(device="cuda").manual_seed(600)
+            self.pre = [
+                (
+                    (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.05).bfloat16(),
+                    (torch.randn(rows, fx.HIDDEN, device="cuda", generator=common)).bfloat16(),
+                )
+                for _ in range(calls)
+            ]
+            self.gamma = (torch.rand(fx.HIDDEN, device="cuda", generator=common) + 0.5).bfloat16()
+        self.sink = []
+
+    def step(self, i: int):
+        if self.attn:
+            self.sink.append(self.attn[i % 2].amax())
+        if self.pre is not None:
+            partial, residual = self.pre[i]
+            self.sink.append(U.native_fused_allreduce(partial, residual, self.gamma, fx.RMS_EPS))
+
+
+class Chain:
+    """The per-rank call sequence of one arm, captured in one CUDA graph."""
+
+    def __init__(
+        self,
+        env,
+        moes,
+        gammas,
+        order,
+        xs,
+        residuals,
+        logits,
+        hop,
+        mode,
+        impl,
+        supplied=None,
+        neighbours=None,
+    ):
+        self.env, self.moes, self.gammas = env, moes, gammas
+        self.order, self.xs, self.residuals, self.logits = order, xs, residuals, logits
+        self.hop, self.mode, self.impl = hop, mode, impl
+        self.supplied, self.neighbours = supplied, neighbours
+        self.outs = []
+
+    def run(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+        with dispatch.step(dispatch.Phase.DECODE):
+            self._run()
+
+    def _run(self):
+        from tensorrt_llm._torch.distributed import AllReduceParams
+
+        self.outs = []
+        if self.neighbours is not None:
+            self.neighbours.sink = []
+        for i, layer in enumerate(self.order):
+            moe = self.moes[layer]
+            moe._routed_moe_bench_logits = None if self.logits is None else self.logits[i]
+            moe._routed_moe_bench_supplied = None if self.supplied is None else self.supplied[i]
+            if self.neighbours is not None:
+                self.neighbours.step(i)
+            x, res, gamma = self.xs[i], self.residuals[i], self.gammas[layer]
+            if self.mode == "fin":
+                h = moe.forward(
+                    x,
+                    do_finalize=True,
+                    final_all_reduce_params=AllReduceParams(enable_allreduce=False),
+                )
+                out = U.native_fused_allreduce(h, res, gamma, fx.RMS_EPS)
+            else:
+                res_moe = moe.forward(x, do_finalize=False)
+                if len(res_moe) == 2:
+                    # routed_moe_into_hop: [shared, routed] with routed finalized into the hop's
+                    # buffer
+                    shared, routed = res_moe
+                    out = self.hop.hop_routed(routed, shared, res, gamma, fx.RMS_EPS)
+                else:
+                    shared, fc2, scale, perm = res_moe
+                    out = self.hop(fc2, scale, perm, shared, res, gamma, fx.RMS_EPS)
+            self.outs.append(out)
+
+
+def build_layers(env, layers):
+    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe  # noqa: E501, F401  (registers routed_moe)
+
+    saved = fx._NoAllReduce
+    fx._NoAllReduce = _PassAllReduce
+    moes, gammas = {}, {}
+    try:
+        for layer in layers:
+            t0 = time.time()
+            full = fx.load_layer_full(layer)
+            moes[layer] = fx.build_deepseek_moe(layer, env.rank, full)
+            del full
+            gammas[layer] = next_norm_weight(layer)
+            if env.rank == 0:
+                print(f"[chain] layer {layer} built in {time.time() - t0:.1f}s", flush=True)
+    finally:
+        fx._NoAllReduce = saved
+    for moe in moes.values():
+        gate = moe.gate
+        orig = gate.forward
+
+        def forward(h, moe=moe, orig=orig):
+            if moe._routed_moe_bench_supplied is not None:
+                # the router contract with logits present: the gate is skipped
+                return moe._routed_moe_bench_supplied
+            out = orig(h)
+            if moe._routed_moe_bench_logits is not None:
+                # sweep: the router GEMM still runs as served; the routed MoE
+                # reads the constructed logits instead (no extra kernel)
+                out = moe._routed_moe_bench_logits
+            return out
+
+        gate.forward = forward
+        moe._routed_moe_bench_logits = None
+        moe._routed_moe_bench_supplied = None
+        moe._routed_moe_bench_router = orig
+    return moes, gammas
+
+
+def supplied_logits(moes, order, xs, logits):
+    """Per call, the FP32 logits a pre-MoE hop would hand over (static buffers)."""
+    out = []
+    with torch.no_grad():
+        for i, layer in enumerate(order):
+            if logits is not None:
+                lg = logits[i].float().contiguous().clone()
+            else:
+                lg = moes[layer]._routed_moe_bench_router(xs[i]).float().contiguous().clone()
+            out.append(lg)
+    torch.cuda.synchronize()
+    return out
+
+
+def real_inputs(layers, rows: int, calls: int, seed: int):
+    """Per call: (layer, x) from consecutive-token windows of the capture."""
+    per_layer = {}
+    for layer in layers:
+        win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer), rows)]
+        assert win and win[0] is not None, f"no hidden states captured for layer {layer}"
+        g = torch.Generator().manual_seed(seed + layer)
+        per_layer[layer] = [win[i] for i in torch.randperm(len(win), generator=g).tolist()]
+    order, xs = [], []
+    for i in range(calls):
+        layer = layers[i % len(layers)]
+        order.append(layer)
+        xs.append(per_layer[layer][(i // len(layers)) % len(per_layer[layer])].cuda())
+    return order, xs
+
+
+def sweep_logits(moes, order, rows: int, distinct: int, seed: int):
+    """Constructed logits routing each call to ``distinct`` experts.
+
+    Consecutive uses of one layer draw from disjoint expert pools.
+    """
+    pools = max(1, fx.NUM_EXPERTS // distinct)
+    perm = torch.randperm(fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(seed)).tolist()
+    out = []
+    for i, layer in enumerate(order):
+        sets = fx.expert_sets_with_distinct(rows, distinct, seed=seed + i)
+        p = (i // len(set(order))) % pools
+        pool = perm[p * distinct : (p + 1) * distinct]
+        uniq = sorted(set(e for row in sets for e in row))
+        remap = {e: pool[j] for j, e in enumerate(uniq)}
+        sets = [[remap[e] for e in row] for row in sets]
+        bias = moes[layer].gate.e_score_correction_bias.to(torch.bfloat16)
+        lg = fx.logits_for_expert_sets(sets, bias, seed=seed + i)
+        ids, _ = fx.routing_reference(lg, bias)
+        assert len(set(ids.flatten().tolist())) == distinct
+        out.append(lg)
+    return out
+
+
+def capture(env, chain):
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            chain.run()
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    graph = torch.cuda.CUDAGraph()
+    with with_multi_stream(True):
+        with torch.cuda.graph(graph, stream=stream):
+            chain.run()
+    torch.cuda.synchronize()
+    for _ in range(2):
+        env.comm.barrier()
+        graph.replay()
+        torch.cuda.synchronize()
+    return graph
+
+
+def time_graph(env, graph, replays: int):
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    start.record()
+    for _ in range(replays):
+        graph.replay()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0
+
+
+def compare(env, a_outs, b_outs, label):
+    """Max |a - b| of the final hidden and residual over every call (sanity)."""
+    worst = [0.0, 0.0]
+    for a, b in zip(a_outs, b_outs):
+        for j in range(2):
+            worst[j] = max(worst[j], (a[j].float() - b[j].float()).abs().max().item())
+    allw = env.comm.allgather(worst)
+    if env.rank == 0:
+        print(
+            f"[chain] check {label}: max|d norm_out| {max(w[0] for w in allw):.3e} "
+            f"max|d residual| {max(w[1] for w in allw):.3e}",
+            flush=True,
+        )
+
+
+def run_config(
+    env, args, moes, gammas, hop, eng, label, order, xs, logits, supplied=None, neighbours=None
+):
+    torch.manual_seed(0)
+    g = torch.Generator(device="cuda").manual_seed(11)
+    residuals = [
+        (torch.randn(args.rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
+        for _ in order
+    ]
+    arms = [a for a in args.arms.split(",") if a]
+    graphs, chains = {}, {}
+    # Every arm is captured ``--captures`` times, in forward and then reverse
+    # arm order: an occasional graph replays with a per-call inter-kernel gap
+    # that is a property of that one instantiation (identical arms captured
+    # elsewhere do not show it), so an arm reports its fastest capture.
+    passes = [arms if p % 2 == 0 else arms[::-1] for p in range(args.captures)]
+    for p, seq in enumerate(passes):
+        for arm in seq:
+            impl, mode = arm.split(":")
+            # "routed_moe" is the routed_moe schedule, "routed_moe_into_hop" the same schedule
+            # writing into the post_moe_hop (raw mode)
+            feats = (
+                ""
+                if impl == "native"
+                else (
+                    "post_moe_hop,routed_moe,routed_moe_into_hop"
+                    if impl.startswith("routed_moe_into_hop")
+                    else "routed_moe"
+                )
+            )
+            if supplied is not None:
+                # the composition's logit producer (post_attention_hop's hop) is selected
+                feats = ",".join(f for f in (feats, "post_attention_hop") if f)
+            set_features(feats)
+            for knob in ARM_KNOBS:
+                os.environ.pop(knob, None)
+            os.environ.update(arm_env(impl))
+            if impl.startswith("routed_moe_into_hop"):
+                # collective, outside capture: the hop's routed-input mode of this arm
+                assert hop.enable_routed_input(os.environ.get(ACCUMULATION_ENV) or None)
+            entry = (
+                "routed_output_into_hop"
+                if impl.startswith("routed_moe_into_hop")
+                else "routed_output"
+            )
+            e0 = eng.engaged(entry)
+            chain = Chain(
+                env,
+                moes,
+                gammas,
+                order,
+                xs,
+                residuals,
+                None if supplied is not None else logits,
+                hop,
+                mode,
+                impl,
+                supplied=supplied,
+                neighbours=neighbours,
+            )
+            graphs[(arm, p)] = capture(env, chain)
+            if p == 0:
+                chains[arm] = chain
+            engaged = eng.engaged(entry) - e0
+            want = 0 if impl == "native" else 3 * len(order)  # two warmups + capture
+            ok = U.all_ranks(engaged == want)
+            if env.rank == 0:
+                print(
+                    f"[chain] {label} arm {arm} capture {p}: routed_moe engaged {engaged}/{want} "
+                    f"all_ranks={ok}"
+                )
+            assert ok, f"{arm}: unexpected routed_moe engagement"
+    set_features("")
+    for knob in ARM_KNOBS:
+        os.environ.pop(knob, None)
+    per = {key: [] for key in graphs}
+    for _ in range(args.trials):
+        for key in graphs:
+            per[key].append(time_graph(env, graphs[key], args.replays))
+    if args.profile_config == label:
+        # one replay per arm inside a cudaProfilerApi capture range, each in
+        # an NVTX range named after the arm (nsys --capture-range=cudaProfilerApi)
+        torch.cuda.synchronize()
+        torch.cuda.profiler.start()
+        for arm, p in graphs:
+            env.comm.barrier()
+            torch.cuda.nvtx.range_push(f"{label}:{arm}" + (f"#{p}" if p else ""))
+            graphs[(arm, p)].replay()
+            torch.cuda.synchronize()
+            torch.cuda.nvtx.range_pop()
+        torch.cuda.profiler.stop()
+    rows = []
+    lg_seen = supplied if supplied is not None else logits
+    dstat = None
+    if lg_seen is not None:
+        dstat = distinct_stats(moes, order, lg_seen)
+    elif env.rank == 0:
+        dstat = distinct_stats(
+            moes,
+            order,
+            [moes[layer]._routed_moe_bench_router(x).float() for layer, x in zip(order, xs)],
+        )
+    if env.rank == 0:
+        print(f"CHAIN_DISTINCT {label} {json.dumps(dstat)}", flush=True)
+    for arm in arms:
+        caps = []
+        for p in range(args.captures):
+            per_call = [t / (args.replays * len(order)) for t in per[(arm, p)]]
+            meds = env.comm.allgather(statistics.median(per_call))
+            caps.append(
+                {
+                    "median": max(meds),
+                    "spread": [min(meds), max(meds)],
+                    "best": max(env.comm.allgather(min(per_call))),
+                    "mean": max(env.comm.allgather(statistics.mean(per_call))),
+                }
+            )
+        best = min(caps, key=lambda c: c["median"])
+        rows.append(
+            {
+                "config": label,
+                "arm": arm,
+                "per_call_us_median_slowest_rank": best["median"],
+                "per_call_us_median_rank_spread": best["spread"],
+                "per_call_us_best_trial_slowest_rank": best["best"],
+                "per_call_us_mean_slowest_rank": best["mean"],
+                "per_capture_median_slowest_rank": [c["median"] for c in caps],
+                "distinct": dstat,
+                "calls_per_graph": len(order),
+                "replays": args.replays,
+                "trials": args.trials,
+            }
+        )
+        if env.rank == 0:
+            print("CHAIN_BENCH " + json.dumps(rows[-1]), flush=True)
+            if len(caps) > 1:
+                m = [c["median"] for c in caps]
+                flag = "  <-- capture spread > 1 us" if max(m) - min(m) > 1.0 else ""
+                print(
+                    f"CHAIN_CAPTURES {label} {arm}: " + " ".join(f"{v:.2f}" for v in m) + flag,
+                    flush=True,
+                )
+    by = {r["arm"]: r["per_call_us_median_slowest_rank"] for r in rows}
+    by_mean = {r["arm"]: r["per_call_us_mean_slowest_rank"] for r in rows}
+    impls = [a.split(":")[0] for a in arms if not a.startswith("native:")]
+    impls = list(dict.fromkeys(impls))
+    if env.rank == 0:
+        for mode in ("fin", "raw"):
+            for impl in impls:
+                if f"native:{mode}" in by and f"{impl}:{mode}" in by:
+                    d = by[f"native:{mode}"] - by[f"{impl}:{mode}"]
+                    dm = by_mean[f"native:{mode}"] - by_mean[f"{impl}:{mode}"]
+                    print(
+                        f"CHAIN_DELTA {label} {mode}: native - {impl} = {d:+.2f} us/call "
+                        f"(trial mean {dm:+.2f})",
+                        flush=True,
+                    )
+        for impl in ("routed_moe", "routed_moe_into_hop"):
+            if "native:fin" in by and f"{impl}:raw" in by:
+                d = by["native:fin"] - by[f"{impl}:raw"]
+                print(f"CHAIN_DELTA {label} native:fin - {impl}:raw = {d:+.2f} us/call", flush=True)
+        if "routed_moe:raw" in by and "routed_moe_into_hop:raw" in by:
+            d = by["routed_moe:raw"] - by["routed_moe_into_hop:raw"]
+            print(
+                f"CHAIN_DELTA {label} routed_moe:raw - routed_moe_into_hop:raw = {d:+.2f} us/call (routed input)",
+                flush=True,
+            )
+    for mode in ("fin", "raw"):
+        for impl in impls:
+            if f"native:{mode}" in chains and f"{impl}:{mode}" in chains:
+                compare(
+                    env,
+                    chains[f"native:{mode}"].outs,
+                    chains[f"{impl}:{mode}"].outs,
+                    f"{label} {mode} {impl}",
+                )
+    del graphs
+    torch.cuda.synchronize()
+    return rows
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--layers", default="10,30,50,77")
+    ap.add_argument("--rows", type=int, default=6)
+    ap.add_argument("--calls", type=int, default=48)
+    ap.add_argument("--replays", type=int, default=10)
+    ap.add_argument("--trials", type=int, default=9)
+    ap.add_argument("--captures", type=int, default=2, help="graphs per arm (best one reported)")
+    ap.add_argument("--arms", default="native:fin,routed_moe:fin,native:raw,routed_moe:raw")
+    ap.add_argument("--mode", default="real", help="real | sweep | mix (comma list)")
+    ap.add_argument("--mix", default="early,late", help="served D distributions (MIXES)")
+    ap.add_argument("--distinct", default="8,14,21,25,31,38,48")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument(
+        "--logits", default="gate", help="gate | supplied | gate,supplied (router-logit producer)"
+    )
+    ap.add_argument("--attn-mb", type=int, default=0, help="attention weight-stream stand-in")
+    ap.add_argument("--pre-hop", default="none", choices=("none", "ar"))
+    ap.add_argument("--out", default="")
+    ap.add_argument(
+        "--profile-config", default="", help="config label replayed once per arm under nsys"
+    )
+    args = ap.parse_args(argv)
+
+    env = U.mpi_env()
+    assert env.world == 8 and env.multi_node, "needs the two-host TP8 group"
+    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
+    set_features("")
+    layers = [int(v) for v in args.layers.split(",")]
+    moes, gammas = build_layers(env, layers)
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    # as the target's decoder layer builds it: the layer's MoE module declares the
+    # raw-tuple contract the hop consumes
+    hop = PostMoeHop(
+        env.mapping,
+        fx.HIDDEN,
+        fx.TOP_K,
+        torch.bfloat16,
+        raw_producer=next(iter(moes.values())).experts,
+    )
+    assert hop.engaged, hop.reason
+    if any(a.startswith("routed_moe_into_hop") for a in args.arms.split(",")):
+        assert hop.enable_routed_input(), (
+            "routed_moe_into_hop routed input not enabled on every rank"
+        )
+        for moe in moes.values():
+            moe._routed_into_hop = hop
+    eng = Engagement()
+    if env.rank == 0:
+        print(
+            f"[chain] ranks={env.world} layers={layers} rows={args.rows} calls={args.calls} "
+            f"pdl={os.environ.get('TRTLLM_ENABLE_PDL')} post_moe_hop={hop.impl}",
+            flush=True,
+        )
+    results = []
+    modes = args.mode.split(",")
+    producers = [p for p in args.logits.split(",") if p]
+    assert set(producers) <= {"gate", "supplied"}, args.logits
+    order, xs = real_inputs(layers, args.rows, args.calls, args.seed)
+    neighbours = None
+    if args.attn_mb > 0 or args.pre_hop != "none":
+        neighbours = Neighbours(env, args.rows, len(order), args.attn_mb, args.pre_hop)
+    if env.rank == 0:
+        print(
+            f"[chain] logits={producers} attn_mb={args.attn_mb} pre_hop={args.pre_hop}",
+            flush=True,
+        )
+
+    def configs(label, lg):
+        out = []
+        for producer in producers:
+            sup = supplied_logits(moes, order, xs, lg) if producer == "supplied" else None
+            name = label + ("+sup" if producer == "supplied" else "")
+            out += run_config(
+                env, args, moes, gammas, hop, eng, name, order, xs, lg, sup, neighbours
+            )
+        return out
+
+    if "real" in modes:
+        results += configs("real", None)
+    if "mix" in modes:
+        for name in [v for v in args.mix.split(",") if v]:
+            ds = mix_distinct(name, len(order), args.rows, args.seed + 5000)
+            lg = mix_logits(moes, order, args.rows, ds, args.seed + 6000)
+            results += configs(f"mix{name}", lg)
+    if "sweep" in modes:
+        for d in [int(v) for v in args.distinct.split(",") if v]:
+            if d < fx.TOP_K or d > fx.TOP_K * args.rows:
+                continue
+            lg = sweep_logits(moes, order, args.rows, d, args.seed + 1000 * d)
+            results += configs(f"D{d}", lg)
+    if env.rank == 0 and args.out:
+        with open(args.out, "w") as f:
+            json.dump(results, f, indent=1)
+    hop.raise_if_error()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py
new file mode 100644
index 000000000000..f97e676e55e7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py
@@ -0,0 +1,1036 @@
+# 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.
+"""Real-weight fixtures and the native oracle for the routed-MoE tests.
+
+The fixture is one GLM-5.2 NVFP4 MoE layer as TP rank ``r`` of TP8/EP1 sees
+it: every routed expert's gate/up rows ``[256r, 256r+256)`` and down columns
+``[256r, 256r+256)``, with block and global scales, the router weight and the
+router bias. The native module is the stock TRT-LLM-Gen NVFP4 leaf loaded
+through its own ``load_weights``/``post_load_weights``, so its buffers carry
+exactly the shuffled/interleaved layouts and derived scalars of a served model.
+Slicing the checkpoint before loading a ``tp_size=1`` module is equivalent to
+the TP8 loader's ``load_weight_shard`` (COLUMN for w1/w3, ROW for w2).
+
+Hidden-state proxies are real token embeddings passed through the layer's
+``post_attention_layernorm`` (labelled semi-synthetic: the served MoE input is
+the normalized residual stream of that layer, which needs a full forward).
+"""
+
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+from utils.llm_data import llm_models_root
+
+CKPT_DIR = os.environ.get(
+    "GLM_MOE_DSA_ROUTED_MOE_CKPT_DIR",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+CACHE_DIR = os.environ.get(
+    "GLM_MOE_DSA_ROUTED_MOE_FIXTURE_CACHE",
+    os.path.join(tempfile.gettempdir(), "glm_moe_dsa_routed_moe_fixtures"),
+)
+
+HIDDEN = 6144
+INTER = 2048
+NUM_EXPERTS = 256
+TOP_K = 8
+TP_SIZE = 8
+ROUTED_SCALING = 2.5
+RMS_EPS = 1e-5
+
+# A fixed English/code text; consecutive tokens of it stand in for the
+# consecutive verify rows of one MTP request.
+PROMPT_TEXT = (
+    "def merge_sorted(a, b):\n    i = j = 0\n    out = []\n"
+    "    while i < len(a) and j < len(b):\n        if a[i] <= b[j]:\n"
+    "            out.append(a[i]); i += 1\n        else:\n"
+    "            out.append(b[j]); j += 1\n    return out + a[i:] + b[j:]\n"
+    "The quick brown fox jumps over the lazy dog while the compiler "
+    "schedules tensor-core instructions across the streaming multiprocessors."
+)
+
+
+def _index() -> Dict[str, str]:
+    with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
+        return json.load(f)["weight_map"]
+
+
+class _Reader:
+    """Opens each safetensors shard once and slices tensors lazily."""
+
+    def __init__(self):
+        from safetensors import safe_open
+
+        self._safe_open = safe_open
+        self._map = _index()
+        self._open = {}
+
+    def _handle(self, name: str):
+        shard = self._map[name]
+        if shard not in self._open:
+            self._open[shard] = self._safe_open(
+                os.path.join(CKPT_DIR, shard), framework="pt", device="cpu"
+            )
+        return self._open[shard]
+
+    def full(self, name: str) -> torch.Tensor:
+        return self._handle(name).get_tensor(name)
+
+    def rows(self, name: str, start: int, stop: int) -> torch.Tensor:
+        # clone: a row slice may be a view of the whole tensor's storage,
+        # which torch.save would otherwise write out in full
+        return self._handle(name).get_slice(name)[start:stop].clone()
+
+    def cols(self, name: str, start: int, stop: int) -> torch.Tensor:
+        return self._handle(name).get_slice(name)[:, start:stop].clone()
+
+
+def load_layer_rank(layer: int, tp_rank: int) -> Dict[str, torch.Tensor]:
+    """Rank-sliced raw checkpoint tensors of one MoE layer (CPU, cached)."""
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    path = os.path.join(CACHE_DIR, f"glm52_L{layer}_tp{TP_SIZE}r{tp_rank}.pt")
+    if os.path.exists(path):
+        return torch.load(path, map_location="cpu")
+    rd = _Reader()
+    inter_rank = INTER // TP_SIZE
+    r0, r1 = tp_rank * inter_rank, (tp_rank + 1) * inter_rank
+    pre = f"model.layers.{layer}.mlp"
+    out: Dict[str, torch.Tensor] = {}
+    for e in range(NUM_EXPERTS):
+        for proj, key in (("gate_proj", "w1"), ("up_proj", "w3")):
+            base = f"{pre}.experts.{e}.{proj}"
+            out[f"{e}.{key}.weight"] = rd.rows(f"{base}.weight", r0, r1)
+            out[f"{e}.{key}.weight_scale"] = rd.rows(f"{base}.weight_scale", r0, r1)
+            out[f"{e}.{key}.weight_scale_2"] = rd.full(f"{base}.weight_scale_2")
+            out[f"{e}.{key}.input_scale"] = rd.full(f"{base}.input_scale")
+        base = f"{pre}.experts.{e}.down_proj"
+        out[f"{e}.w2.weight"] = rd.cols(f"{base}.weight", r0 // 2, r1 // 2)
+        out[f"{e}.w2.weight_scale"] = rd.cols(f"{base}.weight_scale", r0 // 16, r1 // 16)
+        out[f"{e}.w2.weight_scale_2"] = rd.full(f"{base}.weight_scale_2")
+        out[f"{e}.w2.input_scale"] = rd.full(f"{base}.input_scale")
+    out["gate.weight"] = rd.full(f"{pre}.gate.weight")
+    out["gate.e_score_correction_bias"] = rd.full(f"{pre}.gate.e_score_correction_bias")
+    out["post_attention_layernorm.weight"] = rd.full(
+        f"model.layers.{layer}.post_attention_layernorm.weight"
+    )
+    tmp = path + f".tmp{os.getpid()}"
+    torch.save(out, tmp)
+    os.replace(tmp, path)
+    return out
+
+
+def load_layer_full(layer: int) -> Dict[str, torch.Tensor]:
+    """Unsliced checkpoint tensors of one MoE layer (CPU, not cached).
+
+    Keys follow the MoE loader's convention (``{e}.w1.weight`` ...), so the
+    TP8 module's own ``load_weights`` does the rank slicing.
+    """
+    rd = _Reader()
+    pre = f"model.layers.{layer}.mlp"
+    out: Dict[str, torch.Tensor] = {}
+    for e in range(NUM_EXPERTS):
+        for proj, key in (("gate_proj", "w1"), ("up_proj", "w3"), ("down_proj", "w2")):
+            base = f"{pre}.experts.{e}.{proj}"
+            for suffix in ("weight", "weight_scale", "weight_scale_2", "input_scale"):
+                out[f"{e}.{key}.{suffix}"] = rd.full(f"{base}.{suffix}")
+    out["gate.weight"] = rd.full(f"{pre}.gate.weight")
+    out["gate.e_score_correction_bias"] = rd.full(f"{pre}.gate.e_score_correction_bias")
+    out["post_attention_layernorm.weight"] = rd.full(
+        f"model.layers.{layer}.post_attention_layernorm.weight"
+    )
+    return out
+
+
+def _served_model_config(layer: int, tp_rank: int, max_num_tokens: int = 128):
+    """ModelConfig of one served GEN rank: TP8/EP1 over two 4-GPU hosts,
+    no attention DP, TRTLLM MoE backend, NVFP4 experts, BF16 shared expert.
+    """
+    from transformers import PretrainedConfig
+
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm.mapping import Mapping
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+    from tensorrt_llm.quantization.mode import QuantAlgo
+
+    quant_config = QuantConfig(
+        quant_algo=QuantAlgo.NVFP4,
+        kv_cache_quant_algo=QuantAlgo.FP8,
+        group_size=16,
+        exclude_modules=[f"model.layers.{layer}.mlp.shared_experts*"],
+    )
+    pretrained = PretrainedConfig(
+        hidden_size=HIDDEN,
+        n_routed_experts=NUM_EXPERTS,
+        num_experts_per_tok=TOP_K,
+        moe_intermediate_size=INTER,
+        n_shared_experts=1,
+        n_group=1,
+        topk_group=1,
+        routed_scaling_factor=ROUTED_SCALING,
+    )
+    mapping = Mapping(
+        world_size=TP_SIZE,
+        rank=tp_rank,
+        gpus_per_node=4,
+        tp_size=TP_SIZE,
+        moe_tp_size=TP_SIZE,
+        moe_ep_size=1,
+    )
+    return ModelConfig(
+        pretrained_config=pretrained,
+        mapping=mapping,
+        quant_config=quant_config,
+        moe_backend="TRTLLM",
+        max_num_tokens=max_num_tokens,
+    )
+
+
+def _walk_post_load(root: torch.nn.Module) -> None:
+    """The model loader's full post-load walk over one module tree."""
+    for mod in root.modules():
+        hook = getattr(mod, "post_load_weights", None)
+        if hook is not None and not getattr(mod, "_weights_removed", False):
+            hook()
+
+
+def _aux_streams():
+    from tensorrt_llm._torch.utils import AuxStreamType
+
+    return {t: torch.cuda.Stream() for t in AuxStreamType}
+
+
+def build_served_experts(layer: int, tp_rank: int, full: Dict[str, torch.Tensor]):
+    """The routed experts exactly as ``Deepseekv3MoE`` builds them: ``create_moe``
+    with the served TP8 ``ModelConfig`` (a ``ConfigurableMoE`` around the
+    resolved TRT-LLM-Gen leaf), loaded from unsliced checkpoint tensors.
+
+    Returns ``(wrapper, bias)`` with the BF16 router bias the gate routes with.
+    """
+    from tensorrt_llm._torch.moe.fused_moe import create_moe
+    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
+
+    bias = full["gate.e_score_correction_bias"].to(torch.bfloat16).cuda()
+    routing = DeepSeekV3MoeRoutingMethod(
+        top_k=TOP_K,
+        n_group=1,
+        topk_group=1,
+        routed_scaling_factor=ROUTED_SCALING,
+        callable_e_score_correction_bias=lambda: bias,
+        is_fused=True,
+    )
+    mc = _served_model_config(layer, tp_rank)
+    wrapper = create_moe(
+        num_experts=NUM_EXPERTS,
+        routing_method=routing,
+        hidden_size=HIDDEN,
+        intermediate_size=INTER,
+        dtype=torch.bfloat16,
+        reduce_results=False,
+        model_config=mc,
+        override_quant_config=mc.quant_config,
+        aux_stream_dict=_aux_streams(),
+        layer_idx=layer,
+    )
+    weights = {k: v for k, v in full.items() if k[0].isdigit()}
+    wrapper.load_weights([weights])
+    _walk_post_load(wrapper)
+    wrapper = wrapper.cuda()
+    return wrapper, bias
+
+
+class _NoAllReduce(torch.nn.Module):
+    """Stand-in for the cross-rank AllReduce of a single-GPU TP8-shard test;
+    the paths exercised here never reduce.
+    """
+
+    output_buffer_kind = 0
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        raise RuntimeError("cross-rank reduction is not available in a single-GPU test")
+
+
+def build_deepseek_moe(layer: int, tp_rank: int, full: Dict[str, torch.Tensor]):
+    """One served MoE (TP8 rank shard): the built-in ``Deepseekv3MoE`` through its own
+    constructor, adopted by the target's ``GlmMoE`` as the target's decoder layer adopts it.
+
+    The NVFP4 experts and the gate carry the real checkpoint tensors; the BF16
+    shared expert gets random weights (it is identical in both arms of a
+    comparison). The cross-rank AllReduce is replaced by a stub, so only the
+    paths that do not reduce (``forward(do_finalize=False)`` and
+    ``compute_routed_output``) are usable. The target's routed fast paths run
+    only inside a decode step (``dispatch.step(dispatch.Phase.DECODE)``).
+    """
+    import tensorrt_llm._torch.distributed as trtllm_dist
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmMoE,
+    )
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import Deepseekv3MoE
+
+    mc = _served_model_config(layer, tp_rank)
+    saved = trtllm_dist.AllReduce
+    trtllm_dist.AllReduce = _NoAllReduce
+    try:
+        moe = Deepseekv3MoE(
+            num_experts=NUM_EXPERTS,
+            top_k=TOP_K,
+            hidden_size=HIDDEN,
+            intermediate_size=INTER,
+            shared_expert_intermediate_size=INTER,
+            aux_stream_dict=_aux_streams(),
+            dtype=torch.bfloat16,
+            model_config=mc,
+            override_quant_config=mc.quant_config,
+            layer_idx=layer,
+        )
+    finally:
+        trtllm_dist.AllReduce = saved
+    GlmMoE.adopt(moe, mc, layer)
+    moe.gate.load_weights(
+        [
+            {
+                "weight": full["gate.weight"],
+                "e_score_correction_bias": full["gate.e_score_correction_bias"],
+            }
+        ]
+    )
+    weights = {k: v for k, v in full.items() if k[0].isdigit()}
+    moe.experts.load_weights([weights])
+    g = torch.Generator().manual_seed(1234)
+    for p in moe.shared_experts.parameters():
+        if p.dtype.is_floating_point:
+            p.data.copy_(torch.randn(p.shape, generator=g) * 0.02)
+    _walk_post_load(moe)
+    return moe.cuda()
+
+
+def load_router_bias(layer: int) -> torch.Tensor:
+    """BF16 ``e_score_correction_bias`` of one MoE layer.
+
+    BF16 is what the TRT-LLM-Gen backend routes with.
+    """
+    return (
+        _Reader()
+        .full(f"model.layers.{layer}.mlp.gate.e_score_correction_bias")
+        .to(torch.bfloat16)
+        .cuda()
+    )
+
+
+def load_embedding_rows(token_ids: List[int]) -> torch.Tensor:
+    """Rows of ``model.embed_tokens.weight`` (BF16, CPU)."""
+    rd = _Reader()
+    name = "model.embed_tokens.weight"
+    sl = rd._handle(name).get_slice(name)
+    return torch.stack([sl[t : t + 1][0] for t in token_ids]).contiguous()
+
+
+def prompt_token_ids(max_tokens: int = 256) -> List[int]:
+    cache = os.path.join(CACHE_DIR, "prompt_ids.json")
+    if os.path.exists(cache):
+        with open(cache) as f:
+            return json.load(f)[:max_tokens]
+    # The checkpoint config does not validate under the image's transformers
+    # AutoConfig, so read the tokenizer file directly.
+    from tokenizers import Tokenizer
+
+    tok = Tokenizer.from_file(os.path.join(CKPT_DIR, "tokenizer.json"))
+    ids = tok.encode(PROMPT_TEXT).ids
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    with open(cache, "w") as f:
+        json.dump(ids, f)
+    return ids[:max_tokens]
+
+
+@dataclass
+class NativeLayer:
+    """A native TRT-LLM-Gen NVFP4 MoE layer plus its router (one TP rank)."""
+
+    module: torch.nn.Module
+    gate_weight: torch.Tensor  # [256, 6144] bf16
+    bias: torch.Tensor  # [256] bf16, the native TRTLLM gate dtype
+    norm_weight: torch.Tensor  # [6144] bf16
+    layer: int
+    tp_rank: int
+
+    def router_logits(self, x: torch.Tensor) -> torch.Tensor:
+        """Native min-latency router GEMM (FP32 logits), as DeepseekV3Gate."""
+        return torch.ops.trtllm.dsv3_router_gemm_op(
+            x, self.gate_weight.t(), bias=None, out_dtype=torch.float32
+        )
+
+    def native_forward(
+        self,
+        x: torch.Tensor,
+        logits: torch.Tensor,
+        do_finalize: bool = True,
+        bias: Optional[torch.Tensor] = None,
+    ):
+        """The native region of the routed_moe contract: FP4 quant + routing/FC1/FC2/finalize.
+
+        Mirrors ``TRTLLMGenFp4BlockScaleBase.run_moe`` and
+        ``TRTLLMGenNvfp4Base.quantize_input`` argument for argument.
+        ``bias`` overrides the router bias (routing studies with another
+        layer's logits and bias).
+        """
+        m = self.module
+        x_fp4, x_sf = m.quantize_input(x)
+        inter_padded = m.w3_w1_weight.shape[-2] // 2
+        outputs = m.op_backend.run_fp4_block_scale_moe(
+            logits,
+            self.bias if bias is None else bias,
+            x_fp4,
+            x_sf,
+            m.w3_w1_weight,
+            m.w3_w1_weight_scale,
+            None,
+            m.act_alpha,
+            m.act_beta,
+            m.act_clamp,
+            m.w2_weight,
+            m.w2_weight_scale,
+            None,
+            m.fc31_scale_c.data,
+            m.fc31_alpha.data,
+            m.fc2_alpha.data,
+            m.num_slots,
+            TOP_K,
+            1,
+            1,
+            inter_padded,
+            m.slot_start,
+            m.expert_size_per_partition,
+            ROUTED_SCALING,
+            int(m.routing_method.routing_method_type),
+            do_finalize=do_finalize,
+            topk_weights=None,
+            topk_ids=None,
+            valid_hidden_size=m.hidden_size,
+            valid_intermediate_size=getattr(
+                m.quant_method, "intermediate_size_per_partition_lean", None
+            ),
+            gated_act_type=0,
+            output=None,
+            tune_max_num_tokens=m.max_num_tokens,
+            use_dp=False,
+        )
+        if do_finalize:
+            return outputs[0] if isinstance(outputs, (list, tuple)) else outputs
+        return outputs
+
+    def hidden_proxy(self, token_ids: List[int]) -> torch.Tensor:
+        """RMSNorm(embedding) * post_attention_layernorm.weight, BF16 on GPU."""
+        emb = load_embedding_rows(token_ids).cuda().float()
+        var = emb.pow(2).mean(-1, keepdim=True)
+        y = emb * torch.rsqrt(var + RMS_EPS)
+        return (y * self.norm_weight.float()).to(torch.bfloat16).contiguous()
+
+
+def build_native_layer(layer: int = 10, tp_rank: int = 0, max_num_tokens: int = 128) -> NativeLayer:
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
+    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
+    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import trtllm_gen_leaf
+    from tensorrt_llm.mapping import Mapping
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+    from tensorrt_llm.quantization.mode import QuantAlgo
+
+    raw = load_layer_rank(layer, tp_rank)
+    bias = raw["gate.e_score_correction_bias"].to(torch.bfloat16).cuda()
+    routing = DeepSeekV3MoeRoutingMethod(
+        top_k=TOP_K,
+        n_group=1,
+        topk_group=1,
+        routed_scaling_factor=ROUTED_SCALING,
+        callable_e_score_correction_bias=lambda: bias,
+        is_fused=True,
+    )
+    quant_config = QuantConfig(
+        quant_algo=QuantAlgo.NVFP4, kv_cache_quant_algo=QuantAlgo.FP8, group_size=16
+    )
+    model_config = ModelConfig(
+        quant_config=quant_config,
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        moe_backend="TRTLLM",
+        max_num_tokens=max_num_tokens,
+    )
+    leaf_cls = trtllm_gen_leaf(QuantAlgo.NVFP4)
+    module = leaf_cls(
+        routing_method=routing,
+        num_experts=NUM_EXPERTS,
+        hidden_size=HIDDEN,
+        intermediate_size=INTER // TP_SIZE,
+        dtype=torch.bfloat16,
+        reduce_results=False,
+        model_config=model_config,
+        init_load_balancer=False,
+        weight_loading_mode=MoEWeightLoadingMode.VANILLA,
+    )
+    weights = {k: v for k, v in raw.items() if k[0].isdigit()}
+    module.load_weights([weights])
+    module.post_load_weights()
+    module = module.cuda()
+    return NativeLayer(
+        module=module,
+        gate_weight=raw["gate.weight"].cuda().contiguous(),
+        bias=bias,
+        norm_weight=raw["post_attention_layernorm.weight"].cuda(),
+        layer=layer,
+        tp_rank=tp_rank,
+    )
+
+
+def _sum8(v: torch.Tensor, order: str) -> torch.Tensor:
+    """Sum of the 8 top-k values in a fixed FP32 association order."""
+    c = [v[:, i] for i in range(TOP_K)]
+    if order == "xor_desc":
+        # butterfly with offsets 16, 8, 4, 2, 1 over lanes 0..7 (zeros above)
+        return ((c[0] + c[4]) + (c[2] + c[6])) + ((c[1] + c[5]) + (c[3] + c[7]))
+    if order == "xor_asc":
+        return ((c[0] + c[1]) + (c[2] + c[3])) + ((c[4] + c[5]) + (c[6] + c[7]))
+    if order == "seq":
+        acc = c[0]
+        for x in c[1:]:
+            acc = acc + x
+        return acc
+    raise ValueError(order)
+
+
+ROUTING_SUM_ORDER = "xor_desc"
+
+
+def routing_reference(logits: torch.Tensor, bias_bf16: torch.Tensor, sum_order: str = None):
+    """Bit-faithful PyTorch model of the native SigmoidBias/ScaledSumNormalize.
+
+    Returns ``(ids[R,8] int32 in descending-key order, weights[R,8] bf16)``.
+    Selection key: FP32 ``0.5*tanh(0.5*x)+0.5 + float(bias)``; ties go to the
+    lower expert index (the native 64-bit key packs ``65535-idx`` below the
+    score bits).
+    """
+    s = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5
+    score = s + bias_bf16.float()[None, :]
+    # Stable descending sort keeps the lower index first among equal scores.
+    order = torch.sort(score, dim=-1, descending=True, stable=True).indices
+    ids = order[:, :TOP_K]
+    sel = torch.gather(score, 1, ids)
+    sig = sel - bias_bf16.float()[ids]
+    total = _sum8(sig, sum_order or ROUTING_SUM_ORDER)[:, None]
+    w = sig * ROUTED_SCALING / total
+    return ids.to(torch.int32), w.to(torch.bfloat16)
+
+
+def logits_for_expert_sets(
+    expert_sets: List[List[int]], bias_bf16: torch.Tensor, seed: int = 0
+) -> torch.Tensor:
+    """FP32 router logits whose native top-8 for row ``t`` is ``expert_sets[t]``.
+
+    Chosen experts get logits that make ``sigmoid + bias`` clear every
+    unchosen expert's maximum possible score by a margin; values within the
+    chosen set are distinct so the top-8 order is well defined.
+    """
+    g = torch.Generator().manual_seed(seed)
+    bias = bias_bf16.float().cpu()
+    num_rows = len(expert_sets)
+    logits = torch.full((num_rows, NUM_EXPERTS), -12.0)
+    logits += 0.5 * torch.rand((num_rows, NUM_EXPERTS), generator=g)
+    for t, chosen in enumerate(expert_sets):
+        assert len(chosen) == TOP_K and len(set(chosen)) == TOP_K
+        unchosen = [e for e in range(NUM_EXPERTS) if e not in set(chosen)]
+        ceiling = max(bias[e].item() for e in unchosen) + 0.05
+        for k, e in enumerate(chosen):
+            # need sigmoid(l) + bias[e] > ceiling + small spread
+            target = ceiling - bias[e].item() + 0.02 * (TOP_K - k)
+            target = min(max(target, 1e-3), 1.0 - 1e-6)
+            logit = torch.logit(torch.tensor(target, dtype=torch.float64))
+            logits[t, e] = float(logit) + 0.5
+    return logits.cuda()
+
+
+def expert_sets_with_distinct(
+    num_rows: int, num_distinct: int, seed: int = 0
+) -> Optional[List[List[int]]]:
+    """``num_rows`` top-8 sets whose union has exactly ``num_distinct`` experts."""
+    if num_distinct < TOP_K or num_distinct > TOP_K * num_rows:
+        return None
+    if num_distinct > NUM_EXPERTS:
+        return None
+    g = torch.Generator().manual_seed(seed)
+    pool = torch.randperm(NUM_EXPERTS, generator=g)[:num_distinct].tolist()
+    sets = []
+    # Cover the pool round-robin so every expert is used at least once.
+    cursor = 0
+    for _ in range(num_rows):
+        row = []
+        while len(row) < TOP_K:
+            e = pool[cursor % num_distinct]
+            cursor += 1
+            if e not in row:
+                row.append(e)
+        sets.append(row)
+    used = set(e for row in sets for e in row)
+    assert len(used) == num_distinct, (num_rows, num_distinct, len(used))
+    return sets
+
+
+# ---------------------------------------------------------------------------
+# Independent emulation of the native NVFP4 MoE arithmetic (reference only).
+# ---------------------------------------------------------------------------
+_E2M1 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+
+
+def e2m1_lut(device) -> torch.Tensor:
+    vals = _E2M1 + [-v for v in _E2M1]
+    return torch.tensor(vals, dtype=torch.float64, device=device)
+
+
+def unpack_fp4(packed: torch.Tensor) -> torch.Tensor:
+    """uint8 [..., K/2] -> float64 [..., K]; element 0 in the low nibble."""
+    lut = e2m1_lut(packed.device)
+    lo = lut[(packed & 0xF).long()]
+    hi = lut[(packed >> 4).long()]
+    return torch.stack([lo, hi], dim=-1).flatten(-2)
+
+
+def unswizzle_sf(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
+    """128x4-interleaved scale bytes -> linear [rows, cols] (uint8)."""
+    flat = sf.contiguous().view(torch.uint8).flatten()
+    r = torch.arange(rows, device=flat.device)[:, None]
+    c = torch.arange(cols, device=flat.device)[None, :]
+    idx = ((r // 128) * (cols // 4) + c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4
+    return flat[idx]
+
+
+def e4m3_bytes_to_float(b: torch.Tensor) -> torch.Tensor:
+    return b.contiguous().view(torch.float8_e4m3fn).to(torch.float64)
+
+
+def quant_e2m1_rn(v: torch.Tensor) -> torch.Tensor:
+    """FP32 -> E2M1 value with round-to-nearest-even and saturation."""
+    a = v.abs().double()
+    grid = torch.tensor(_E2M1, dtype=torch.float64, device=v.device)
+    # nearest grid point; ties go to the even code
+    d = (a[..., None] - grid).abs()
+    best = d.min(-1, keepdim=True).values
+    cand = d == best
+    codes = torch.arange(8, device=v.device)
+    even = cand & (codes % 2 == 0)
+    pick = torch.where(even.any(-1, keepdim=True), even, cand)
+    code = (pick.int() * (8 - codes)).argmax(-1)  # lowest qualifying code
+    q = grid[code]
+    q = torch.where(a >= 6.0, torch.full_like(q, 6.0), q)
+    return torch.copysign(q, v.double())
+
+
+def quant_nvfp4_unit(h: torch.Tensor):
+    """NVFP4 quantization with unit global scale, block 16 along the last dim.
+
+    Returns the dequantized values (q * sf) as float64.
+    """
+    shp = h.shape
+    hb = h.reshape(*shp[:-1], shp[-1] // 16, 16).float()
+    amax = hb.abs().amax(-1, keepdim=True)
+    sf = (amax * (1.0 / 6.0)).to(torch.float8_e4m3fn).to(torch.float32)
+    osc = torch.where(amax > 0, 1.0 / sf, torch.zeros_like(sf))
+    q = quant_e2m1_rn(hb * osc)
+    return (q * sf.double()).reshape(shp)
+
+
+def stored_row_maps():
+    """Stored-row -> (is_gate, j) for w3_w1 and stored-row -> hidden for w2."""
+    from tensorrt_llm.quantization.utils.fp4_utils import (
+        get_reorder_rows_for_gated_act_gemm_row_indices,
+        get_shuffle_matrix_a_row_indices,
+    )
+
+    dummy13 = torch.empty(2 * (INTER // TP_SIZE), 1)
+    p0 = get_reorder_rows_for_gated_act_gemm_row_indices(dummy13)
+    p1 = get_shuffle_matrix_a_row_indices(dummy13, epilogue_tile_m=128)
+    orig13 = p0[p1]  # stored row -> row of cat([w3, w1])
+    inter = INTER // TP_SIZE
+    is_gate = orig13 >= inter
+    j13 = torch.where(is_gate, orig13 - inter, orig13)
+    dummy2 = torch.empty(HIDDEN, 1)
+    orig2 = get_shuffle_matrix_a_row_indices(dummy2, epilogue_tile_m=128)
+    return is_gate, j13, orig2
+
+
+def dequant_expert(module, e: int):
+    """Stored-row dequantized FC1/FC2 weights of expert ``e`` (no global)."""
+    w13 = unpack_fp4(module.w3_w1_weight[e])
+    sf13 = e4m3_bytes_to_float(
+        unswizzle_sf(module.w3_w1_weight_scale[e], 2 * (INTER // TP_SIZE), HIDDEN // 16)
+    )
+    w13 = w13 * sf13.repeat_interleave(16, dim=1)
+    w2 = unpack_fp4(module.w2_weight[e])
+    sf2 = e4m3_bytes_to_float(
+        unswizzle_sf(module.w2_weight_scale[e], HIDDEN, (INTER // TP_SIZE) // 16)
+    )
+    w2 = w2 * sf2.repeat_interleave(16, dim=1)
+    return w13, w2
+
+
+def _device_row_maps(device):
+    is_gate, j13, orig2 = stored_row_maps()
+    return is_gate.to(device), j13.to(device), orig2.to(device)
+
+
+def emulate_expert_row(module, e: int, xd_row: torch.Tensor, maps, cache=None) -> torch.Tensor:
+    """FC2 output row (BF16 values as float32) of expert ``e`` for one
+    dequantized activation row, with FP64 accumulation and the native
+    SwiGLU / NVFP4 / scaling steps.
+    """
+    is_gate, j13, orig2 = maps
+    inter = INTER // TP_SIZE
+    if cache is not None and e in cache:
+        w13, w2 = cache[e]
+    else:
+        w13, w2 = dequant_expert(module, e)
+        if cache is not None:
+            cache[e] = (w13, w2)
+    a13 = w13 @ xd_row  # stored-row accumulators
+    up = torch.zeros(inter, dtype=torch.float64, device=xd_row.device)
+    gate = torch.zeros(inter, dtype=torch.float64, device=xd_row.device)
+    up[j13[~is_gate]] = a13[~is_gate]
+    gate[j13[is_gate]] = a13[is_gate]
+    gr = (gate * float(module.fc31_alpha[e])).float()
+    us = (up * float(module.fc31_scale_c[e])).float()
+    h = us * (gr * torch.sigmoid(gr))
+    hd = quant_nvfp4_unit(h[None, :])[0]
+    a2 = w2 @ hd  # stored rows
+    y = torch.zeros(HIDDEN, dtype=torch.float64, device=xd_row.device)
+    y[orig2] = a2
+    return (y * float(module.fc2_alpha[e])).to(torch.bfloat16).float()
+
+
+def dequant_activations(module, x: torch.Tensor) -> torch.Tensor:
+    """The stock ``fp4_quantize`` of ``x``, dequantized to float64."""
+    rows = x.shape[0]
+    xq, xsf = module.quantize_input(x)
+    return unpack_fp4(xq.view(torch.uint8)) * e4m3_bytes_to_float(
+        xsf.view(torch.uint8).reshape(rows, HIDDEN // 16)
+    ).repeat_interleave(16, dim=1)
+
+
+def emulate_moe(layer: "NativeLayer", x: torch.Tensor, logits: torch.Tensor):
+    """FP64-accumulated emulation of the native region (finalized output).
+
+    Activations are quantized with the stock ``fp4_quantize`` op; weights are
+    the stock module buffers; routing is ``routing_reference``.
+    """
+    m = layer.module
+    rows = x.shape[0]
+    xd = dequant_activations(m, x)
+    ids, wts = routing_reference(logits, layer.bias)
+    maps = _device_row_maps(x.device)
+    out = torch.zeros(rows, HIDDEN, dtype=torch.float32, device=x.device)
+    cache = {}
+    for t in range(rows):
+        acc = torch.zeros(HIDDEN, dtype=torch.float32, device=x.device)
+        for k in range(TOP_K):
+            y = emulate_expert_row(m, int(ids[t, k]), xd[t], maps, cache)
+            acc = acc + wts[t, k].float() * y
+        out[t] = acc
+    return out.to(torch.bfloat16)
+
+
+def error_stats(a: torch.Tensor, b: torch.Tensor) -> dict:
+    a = a.float()
+    b = b.float()
+    diff = a - b
+    ref_rms = b.pow(2).mean().sqrt().item()
+    return {
+        "nrms": (diff.pow(2).mean().sqrt().item() / max(ref_rms, 1e-30)),
+        "max_abs": diff.abs().max().item(),
+        "ref_rms": ref_rms,
+        "cos": torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item(),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Native capture of real routed-MoE inputs (read-only shared fixture).
+# ---------------------------------------------------------------------------
+REAL_DIR = os.environ.get("GLM_MOE_DSA_REAL_MOE_INPUTS_DIR", "")
+
+
+def load_real_capture(layer: int, real_dir: str = REAL_DIR):
+    """Per prompt, ``(hidden [T, 6144] bf16 or None, logits [T, 256] fp32)``
+    that the native ``compute_routed_output`` of ``layer`` received (CPU).
+    Hidden states were captured for layers 10, 30, 50 and 77 only.
+    """
+    import glob
+
+    out = []
+    for path in sorted(glob.glob(os.path.join(real_dir, f"rows*_layer{layer}.pt"))):
+        obj = torch.load(path, map_location="cpu")
+        out.append((obj.get("hidden"), obj["logits"]))
+    return out
+
+
+def real_windows(capture, rows: int, stride: Optional[int] = None):
+    """Consecutive-token windows ``(x, logits)`` (CPU); ``x`` is None when the
+    capture has no hidden states for this layer.
+    """
+    stride = stride or rows
+    win = []
+    for hidden, logits in capture:
+        for start in range(0, logits.shape[0] - rows + 1, stride):
+            x = None if hidden is None else hidden[start : start + rows]
+            win.append((x, logits[start : start + rows]))
+    return win
+
+
+def routing_keys(logits: torch.Tensor, bias_bf16: torch.Tensor) -> torch.Tensor:
+    """FP32 selection keys ``sigmoid + bias`` of the native routing."""
+    s = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5
+    return s + bias_bf16.float()[None, :].to(s.device)
+
+
+def boundary_gap(logits: torch.Tensor, bias_bf16: torch.Tensor) -> torch.Tensor:
+    """Per row, key(8th) - key(9th) in FP32 (0 means an exact boundary tie)."""
+    k = routing_keys(logits, bias_bf16)
+    top = torch.topk(k, TOP_K + 1, dim=-1).values
+    return top[:, TOP_K - 1] - top[:, TOP_K]
+
+
+def _key_of(logit: torch.Tensor, bias_val: torch.Tensor) -> torch.Tensor:
+    s = 0.5 * torch.tanh(0.5 * logit.float()) + 0.5
+    return s + bias_val.float()
+
+
+def logits_for_key(target_key: float, bias_val: torch.Tensor, offset_ulps: int = 0):
+    """All FP32 logits within +-2^16 logit ulps of the analytic solution whose
+    native key equals ``target_key`` shifted by ``offset_ulps`` key ulps (the
+    key arithmetic is FP32, evaluated on the device). May be empty.
+    """
+    dev = bias_val.device
+    want = torch.tensor(target_key, dtype=torch.float32, device=dev)
+    toward = torch.tensor(float("inf") if offset_ulps > 0 else -float("inf"), device=dev)
+    for _ in range(abs(offset_ulps)):
+        want = torch.nextafter(want, toward)
+    s_target = float(want.double() - bias_val.double())
+    s_target = min(max(s_target, 1e-4), 1 - 1e-4)
+    l0 = torch.logit(torch.tensor(s_target, dtype=torch.float64)).float().reshape(1).to(dev)
+    steps = torch.arange(-(1 << 16), 1 << 16, device=dev, dtype=torch.int32)
+    cand = (l0.view(torch.int32) + steps).view(torch.float32)
+    keys = _key_of(cand, bias_val)
+    return cand[keys == want]
+
+
+def logit_for_key(target_key: float, bias_val: torch.Tensor, offset_ulps: int = 0, avoid=None):
+    """One logit from ``logits_for_key`` (the one farthest from ``avoid`` when
+    given, else the middle hit) or ``None``.
+    """
+    hits = logits_for_key(target_key, bias_val, offset_ulps)
+    if hits.numel() == 0:
+        return None
+    if avoid is None:
+        return hits[hits.numel() // 2].item()
+    return hits[(hits - avoid).abs().argmax()].item()
+
+
+def raw_rows_by_topk(raw) -> torch.Tensor:
+    """``[R, 8, 6144]`` FC2 rows in top-k order from a raw tuple
+    ``(fc2_output, scales, expanded_idx_to_permuted_idx)``.
+    """
+    fc2, _, idx = raw
+    return fc2[idx.long()]
+
+
+def finalize_raw(raw) -> torch.Tensor:
+    """The stock finalize arithmetic: ``sum_k scale[t,k] * fc2[idx[t,k]]`` in
+    FP32, top-k order, rounded to BF16.
+    """
+    fc2, scale, idx = raw
+    rows = scale.shape[0]
+    fin = torch.zeros(rows, fc2.shape[1], device=fc2.device, dtype=torch.float32)
+    for k in range(TOP_K):
+        fin += scale[:, k : k + 1].float() * fc2[idx[:, k].long()].float()
+    return fin.to(torch.bfloat16)
+
+
+# ---------------------------------------------------------------------------
+# routed_moe-versus-native classification.
+#
+# Discrete invariants are exact: the expert scales of every (t, k) are equal
+# bit for bit and every (t, k) uses the same expert as native. Expert identity
+# is read from the FC2 rows: rows within 5 % relative L2 of native's row are
+# the same expert (rows of two different experts are far apart); any other
+# window is decoded exactly by ``ExpertDecoder`` (a probe activation shared by
+# all tokens turns every FC2 row into a per-expert fingerprint).
+# Floating-point results fall in one of three classes:
+#   exact    finalized outputs bitwise equal;
+#   reorder  every FC2 row bitwise equal to native's and the finalized output
+#            within the bound of the same eight exact products summed in a
+#            different FP32 order (routed_moe sums per expert slot and across split
+#            partials, the native finalize in top-k order):
+#            |a - b| <= 2^-20 * sum_k |w_k y_k| + ulp_bf16(max(|a|, |b|));
+#   budget   some FC2 rows differ (an FP4 rounding of the intermediate flipped
+#            because the split-K FC1 accumulation is associated differently);
+#            checked against the frozen error budget
+#            (nrms <= 1.5 * nrms(emul - native) + 2e-3, no row above 4x).
+# ---------------------------------------------------------------------------
+SAME_EXPERT_REL = 0.05
+
+
+def bf16_ulp(v: torch.Tensor) -> torch.Tensor:
+    _, e = torch.frexp(v.float().abs())
+    return torch.ldexp(torch.ones_like(v, dtype=torch.float32), e - 8)
+
+
+def raw_report(raw_a, raw_n) -> dict:
+    ra = raw_rows_by_topk(raw_a).float()
+    rn = raw_rows_by_topk(raw_n).float()
+    num = (ra - rn).pow(2).sum(-1).sqrt()
+    den = rn.pow(2).sum(-1).sqrt()
+    rel = torch.where(den > 0, num / den.clamp_min(1e-30), (num > 0).float() * 1e30)
+    rows_exact = (ra == rn).all(-1)
+    return {
+        "scale_exact": bool(torch.equal(raw_a[1], raw_n[1])),
+        "same_expert": bool((rel < SAME_EXPERT_REL).all()),
+        "rows_exact": int(rows_exact.sum()),
+        "rows": rows_exact.numel(),
+        "max_row_rel": float(rel.max()),
+    }
+
+
+def reorder_close(fin_a: torch.Tensor, fin_n: torch.Tensor, raw_n) -> bool:
+    fc2, scale, idx = raw_n
+    terms = (scale.float()[:, :, None] * fc2[idx.long()].float()).abs().sum(1)
+    a = fin_a.float()
+    b = fin_n.float()
+    bound = terms * 2.0**-20 + bf16_ulp(torch.maximum(a.abs(), b.abs()))
+    return bool(((a - b).abs() <= bound).all())
+
+
+def budget_check(out: torch.Tensor, native: torch.Tensor, emul: torch.Tensor):
+    """The frozen error budget; returns ``(ok, stats)``."""
+    e_native = error_stats(out, native)
+    e_emul = error_stats(emul, native)
+    budget = 1.5 * e_emul["nrms"] + 2e-3
+    row_ref = native.float().pow(2).mean(-1).sqrt()
+    row_err = (out.float() - native.float()).pow(2).mean(-1).sqrt()
+    ok = e_native["nrms"] <= budget and bool(torch.all(row_err <= 4 * budget * row_ref + 1e-6))
+    return ok, {"routed_moe_vs_native": e_native, "emul_vs_native": e_emul, "budget": budget}
+
+
+class ExpertDecoder:
+    """Exact expert IDs of a raw tuple, routed_moe's or native's.
+
+    With every token row equal to one probe activation, the FC2 row of
+    ``(t, k)`` depends only on the expert, so it is matched (nearest
+    neighbour, L2) against per-expert fingerprints: native FC2 rows of every
+    local expert for the probe, produced with constructed logits whose
+    expert of every ``(t, k)`` is known. Routing comes from the logits and
+    bias only, so a decoded call uses the window's logits with the probe as
+    activation.
+    """
+
+    def __init__(self, layer: "NativeLayer", probe: torch.Tensor):
+        self.layer = layer
+        self.probe = probe.reshape(1, HIDDEN).to(torch.bfloat16).cuda()
+        # Fingerprints: every expert's FC2 row for the probe, computed by the
+        # FP64 emulation of the native arithmetic from the native buffers
+        # (independent of any routing). A native or routed_moe row of the same
+        # expert differs from it at most by an FP4 rounding of the
+        # intermediate; rows of different experts are far apart.
+        m = layer.module
+        xd = dequant_activations(m, self.probe)[0]
+        maps = _device_row_maps(xd.device)
+        self.fps = torch.stack(
+            [emulate_expert_row(m, e, xd, maps) for e in range(NUM_EXPERTS)]
+        ).float()
+        norms = self.fps.norm(dim=-1)
+        d = torch.cdist(self.fps, self.fps) + torch.eye(NUM_EXPERTS, device="cuda") * 1e30
+        self.min_norm_ratio = float(norms.min() / norms.median())
+        self.min_separation = float(d.min() / norms.median())
+        assert self.min_norm_ratio > 0.05 and self.min_separation > 0.2, (
+            self.min_norm_ratio,
+            self.min_separation,
+        )
+
+    def _x(self, rows: int) -> torch.Tensor:
+        return self.probe.expand(rows, HIDDEN).contiguous()
+
+    def decode(self, raw):
+        rows = raw_rows_by_topk(raw).float()
+        flat = rows.reshape(-1, HIDDEN)
+        d = torch.cdist(flat, self.fps)
+        best = d.topk(2, dim=-1, largest=False)
+        margin = best.values[:, 1] / best.values[:, 0].clamp_min(1e-30)
+        return best.indices[:, 0].reshape(rows.shape[:2]), float(margin.min())
+
+    def compare(self, runner, logits, bias=None) -> dict:
+        """Decode native's and routed_moe's experts for ``logits`` (and ``bias``)."""
+        x = self._x(logits.shape[0])
+        n_raw = self.layer.native_forward(x, logits, do_finalize=False, bias=bias)
+        w_raw = runner(x, logits, do_finalize=False)
+        n_ids, n_margin = self.decode(n_raw)
+        w_ids, w_margin = self.decode(w_raw)
+        return {
+            "ids_equal": bool(torch.equal(n_ids, w_ids)),
+            "scale_exact": bool(torch.equal(n_raw[1], w_raw[1])),
+            "margin": min(n_margin, w_margin),
+            "native_ids": n_ids,
+            "routed_moe_ids": w_ids,
+        }
+
+
+def classify_vs_native(w_fin, w_raw, n_fin, n_raw, emul_fn=None, label="", id_check=None):
+    """Assert the discrete invariants and return ``(class, report)``.
+
+    ``emul_fn`` computes the FP64 emulation for the budget class; without it
+    a window of that class is returned as ``"budget-unchecked"``.
+    ``id_check`` (e.g. ``lambda: decoder.compare(runner, logits)``) decodes
+    the experts exactly when some FC2 row is not within 5 % of native's.
+    """
+    rep = raw_report(w_raw, n_raw)
+    assert rep["scale_exact"], f"{label}: expert scales differ from native"
+    if not rep["same_expert"]:
+        assert id_check is not None, f"{label}: FC2 rows far from native and no id check {rep}"
+        dec = id_check()
+        rep["decoded"] = {k: dec[k] for k in ("ids_equal", "scale_exact", "margin")}
+        assert dec["ids_equal"] and dec["scale_exact"] and dec["margin"] > 4, (
+            f"{label}: decoded experts differ from native {rep}"
+        )
+    if torch.equal(w_fin, n_fin):
+        return "exact", rep
+    if rep["rows_exact"] == rep["rows"]:
+        assert reorder_close(w_fin, n_fin, n_raw), f"{label}: beyond the FP32 reorder bound"
+        return "reorder", rep
+    if emul_fn is None:
+        return "budget-unchecked", rep
+    ok, stats = budget_check(w_fin, n_fin, emul_fn())
+    rep["budget"] = stats
+    assert ok, f"{label}: outside the frozen error budget {stats}"
+    return "budget", rep
+
+
+def check_call(layer, runner, x, logits, fin=None, raw=None, emul=True, label="", decoder=None):
+    """Run routed_moe (unless outputs are given) and native on one input; classify."""
+    if fin is None:
+        fin = runner(x, logits)
+    if raw is None:
+        raw = runner(x, logits, do_finalize=False)
+    n_fin = layer.native_forward(x, logits)
+    n_raw = layer.native_forward(x, logits, do_finalize=False)
+    emul_fn = (lambda: emulate_moe(layer, x, logits)) if emul else None
+    id_check = (lambda: decoder.compare(runner, logits)) if decoder is not None else None
+    return classify_vs_native(fin, raw, n_fin, n_raw, emul_fn, label, id_check)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py
new file mode 100644
index 000000000000..6f773bad1aa6
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py
@@ -0,0 +1,464 @@
+# 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.
+"""routed_moe_into_hop on the two-host TP8 group.
+
+routed_moe writes its finalized routed output into the post_moe_hop, which adds the shared
+output, reduces across the group and normalizes.
+
+Run every test on all 8 ranks (``srun --mpi=pmix``, two hosts). The MoE is the built-in
+``Deepseekv3MoE`` adopted by the target's ``GlmMoE``, and every test runs as a decode step.
+Real GLM-5.2 weights (layers 10 and 77) and captured real hidden-state windows.
+
+Oracles, per call (identical inputs on every arm):
+
+* ``raw``: ``Deepseekv3MoE.forward(do_finalize=False)`` with routed_moe (raw tuple)
+  -> post_moe_hop on the raw tuple, i.e. the composition without routed_moe_into_hop;
+* ``native``: native MoE finalized -> native MNNVL ``AllReduce`` with
+  RESIDUAL_RMS_NORM.
+
+routed_moe_into_hop must equal ``raw`` bit for bit (same routed values, same post_moe_hop arithmetic
+after the finalize) and match ``native`` like the composition does.
+
+The hop is built, and its routed input enabled, inside the served model
+loader's ``MetaInitMode``, as the target's decoder-layer adoption does in the
+server (``routed_into_hop_checks.served_construction``).
+
+Tests of the routed-input variant and the routed_moe finalize -> post_moe_hop
+hand-off (helpers in ``routed_into_hop_checks.py``):
+
+* ``test_routed_into_hop_served_construction_buffers``: per mode, a hop built inside the
+  served loader's ``MetaInitMode`` has its routed-input buffers in device
+  memory, and its routed_moe_into_hop calls equal the raw variant;
+* ``test_routed_into_hop_shared_absent_variant``: the ``has_shared=False`` routed variant;
+* ``test_routed_into_hop_mixed_rank_variants``: some ranks keep the raw tuple (routed_moe raw or
+  the native runner) while the others use routed_moe_into_hop, on the shared post_moe_hop workspace;
+  no hang, bit-equal to every rank on the raw variant, published partials
+  and outputs equal to the native-arithmetic emulation;
+* ``test_routed_into_hop_mixed_rank_graph``: the same mix inside a multi-stream CUDA
+  graph, replayed on fresh windows;
+* ``test_routed_into_hop_delayed_stress``: back-to-back batches with random per-rank
+  producers, shared output and device delays, plus delay injection inside
+  the routed_moe kernels (``routed_moe.DELAY_INJECTION``: CTA entry of the
+  routed kernel and the finalize kernel's stores).
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(HERE), "routed_moe"))
+sys.path.insert(
+    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
+)
+sys.path.insert(0, HERE)
+import post_moe_hop_mpi_utils as U  # noqa: E402
+import routed_into_hop_checks as G  # noqa: E402
+import routed_moe_chain_bench as CB  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+LAYERS = [
+    int(v) for v in os.environ.get("GLM_MOE_DSA_ROUTED_MOE_INTO_HOP_LAYERS", "10,77").split(",")
+]
+
+
+@pytest.fixture(autouse=True)
+def decode_step():
+    """Every test runs as a decode step, where the target's MoE takes its routed fast paths."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    with dispatch.step(dispatch.Phase.DECODE):
+        yield
+
+
+@pytest.fixture(scope="module")
+def env():
+    e = U.mpi_env()
+    if e.world != 8 or not e.multi_node:
+        pytest.skip("needs the two-host TP8 group")
+    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
+    return e
+
+
+@pytest.fixture(scope="module")
+def stack(env):
+    # as the served model builds it: inside the loader's MetaInitMode, with the
+    # layer's MoE module declaring the raw-tuple contract the hop consumes
+    moes, gammas, hop = G.build_stack(env, LAYERS)
+    eng = CB.Engagement()
+    return moes, gammas, hop, eng
+
+
+@pytest.mark.parametrize("mode", ["row", "identity", "fp32"])
+def test_routed_into_hop_served_construction_buffers(env, stack, mode):
+    """Every mode's routed-input buffers are device memory when the hop is
+    built inside the served loader's ``MetaInitMode``, and a routed_moe_into_hop call in that
+    mode equals the raw variant bit for bit (modes "row", "identity") or within
+    the FP32 reorder bound ("fp32")."""
+    moes, gammas, _, _ = stack
+    with G.served_construction():
+        probe = torch.empty(4, device="cuda")
+    assert probe.is_meta, "the harness does not reproduce the served loader's MetaInitMode"
+    layer = LAYERS[0]
+    hop, enabled = G.build_hop(env, moes[layer].experts, mode)
+    assert hop.engaged and enabled and hop.routed_mode == mode
+    problems = G.buffer_problems(hop)
+    assert U.all_ranks(not problems), f"rank {env.rank}: {problems}"
+    saved = {m: m._routed_into_hop for m in moes.values()}
+    try:
+        for m in moes.values():
+            m._routed_into_hop = hop
+        G.warm(env, {layer: moes[layer]}, [6], producers=("routed_moe_into_hop", "routed_moe"))
+        ok, n, fails = G.run_mixed(
+            env, {layer: moes[layer]}, gammas, hop, "all-routed_moe_into_hop", True, [6], 2
+        )
+        hop.raise_if_error()
+    finally:
+        for m, h in saved.items():
+            m._routed_into_hop = h
+        CB.set_features("")
+    if env.rank == 0:
+        per_hop = G.buffer_bytes(hop)
+        print(
+            f"[routed_moe_into_hop] served construction mode={mode}: buffers {G.buffer_report(hop)}, {n} calls; "
+            f"{per_hop} bytes per hop, {per_hop * G.MOE_LAYERS / 2**20:.1f} MiB per rank over "
+            f"{G.MOE_LAYERS} MoE layers"
+        )
+    assert U.all_ranks(ok), f"rank {env.rank}: {fails}"
+
+
+@pytest.mark.parametrize("bad_rank", [0, 5])
+def test_routed_into_hop_buffer_problem_turns_routed_into_hop_off_on_every_rank(
+    env, stack, monkeypatch, bad_rank
+):
+    """The buffer check's failure path: if one rank's routed-input buffer is not
+    device memory (as an allocation under ``MetaInitMode`` leaves it),
+    ``enable_routed_input`` returns False on every rank and the MoE keeps the
+    raw tuple, before any device work touches the buffer."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        hop as hop_mod,
+    )
+
+    moes, gammas, _, _ = stack
+    layer = LAYERS[0]
+    if env.rank == bad_rank:
+        monkeypatch.setattr(
+            hop_mod.PostMoeHop,
+            "_routed_device_buffer",
+            lambda self, shape, dtype: torch.empty(shape, dtype=dtype, device="meta"),
+        )
+    hop, enabled = G.build_hop(env, moes[layer].experts, "identity")
+    monkeypatch.undo()
+    assert hop.engaged
+    assert U.all_ranks(not enabled and not hop._routed_enabled), (
+        "routed_moe_into_hop stayed on somewhere"
+    )
+    moe = moes[layer]
+    saved = moe._routed_into_hop
+    try:
+        moe._routed_into_hop = hop
+        CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
+        out = moe.forward(_windows(layer, 6, 1)[0], do_finalize=False)
+        torch.cuda.synchronize()
+        assert len(out) == 4, (
+            "routed_moe_into_hop must keep the raw tuple when its hop declined routed input"
+        )
+    finally:
+        moe._routed_into_hop = saved
+        CB.set_features("")
+    if env.rank == 0:
+        print(
+            f"[routed_moe_into_hop] buffer problem on rank {bad_rank}: routed_moe_into_hop off on all {env.world} ranks"
+        )
+
+
+@pytest.mark.parametrize("odd_rank", [0, 6])
+def test_routed_into_hop_modes_differing_across_ranks_turn_routed_into_hop_off_everywhere(
+    env, stack, monkeypatch, odd_rank
+):
+    """``TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION`` is read per rank.
+
+    It comes from each rank's environment. If one rank selects another mode
+    than its peers, ``enable_routed_input`` agrees on that and turns routed_moe_into_hop off on
+    every rank (the MoE keeps the raw tuple), instead of enabling a different
+    routed-input variant per rank.
+    """
+    moes, _, _, _ = stack
+    layer = LAYERS[0]
+    monkeypatch.setenv(
+        "TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION",
+        "identity" if env.rank == odd_rank else "fp32",
+    )
+    hop, enabled = G.build_hop(env, moes[layer].experts)
+    monkeypatch.delenv("TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION")
+    assert hop.engaged
+    assert U.all_ranks(not enabled and not hop._routed_enabled), (
+        "routed_moe_into_hop stayed on somewhere"
+    )
+    # The same environment on every rank, spelled as the selector accepts it,
+    # enables routed_moe_into_hop everywhere with the mode the feature record names.
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    monkeypatch.setenv("TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION", " Identity ")
+    hop2, enabled2 = G.build_hop(env, moes[layer].experts)
+    record = dispatch.feature_record({"post_moe_hop", "routed_moe", "routed_moe_into_hop"})
+    monkeypatch.delenv("TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION")
+    assert U.all_ranks(enabled2 and hop2.routed_mode == "identity")
+    assert "routed_moe_into_hop[identity]" in record
+    if env.rank == 0:
+        print(
+            f"[routed_moe_into_hop] mode differs on rank {odd_rank}: routed_moe_into_hop off on all {env.world} ranks"
+        )
+
+
+def _windows(layer: int, rows: int, count: int):
+    win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer), rows)]
+    step = max(1, len(win) // count)
+    return [win[i * step].cuda() for i in range(count)]
+
+
+def _arm(moe, hop, x, res, gamma, features: str):
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    CB.set_features(features)
+    if features == "":
+        h = moe.forward(
+            x, do_finalize=True, final_all_reduce_params=AllReduceParams(enable_allreduce=False)
+        )
+        return U.native_fused_allreduce(h, res, gamma, fx.RMS_EPS), None
+    out = moe.forward(x, do_finalize=False)
+    if len(out) == 2:
+        shared, routed = out
+        return hop.hop_routed(routed, shared, res, gamma, fx.RMS_EPS), routed
+    shared, fc2, scale, perm = out
+    return hop(fc2, scale, perm, shared, res, gamma, fx.RMS_EPS), None
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_routed_into_hop_equals_composition_and_matches_native(env, stack, rows):
+    moes, gammas, hop, eng = stack
+    mode = hop.routed_mode
+    buffers = {
+        "row": lambda: hop.routed_input_buffer(rows),
+        "identity": lambda: hop.routed_raw_buffers(rows)[0],
+        "fp32": lambda: hop.routed_acc_buffer(rows),
+    }
+    g = torch.Generator(device="cuda").manual_seed(100 + rows)
+    worst_native = [0, 0.0]
+    exact_native = 0
+    total = 0
+    reordered = 0
+    for layer in LAYERS:
+        for x in _windows(layer, rows, 6):
+            res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
+            e0 = eng.engaged("routed_output_into_hop")
+            CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
+            out = moes[layer].forward(x, do_finalize=False)
+            assert eng.engaged("routed_output_into_hop") == e0 + 1, (
+                "routed_moe_into_hop did not engage"
+            )
+            assert len(out) == 2 and out[1].data_ptr() == buffers[mode]().data_ptr()
+            # the hop clears the FP32 accumulator after reading: keep a copy
+            seen = [out[0], out[1].clone() if mode == "fp32" else out[1]]
+            w_out, w_res = G.hop_call(hop, "routed", out, res, gammas[layer], True)
+            mine = G.local_partial("routed", seen, True, hop=hop)
+            (r_out, r_res), _, r_payload = G.call(
+                moes[layer], hop, x, res, gammas[layer], "routed_moe", True
+            )
+            (n_out, n_res), _ = _arm(moes[layer], hop, x, res, gammas[layer], "")
+            torch.cuda.synchronize()
+            if G.exact_mode(hop):
+                assert torch.equal(w_res, r_res), (
+                    f"L{layer} R={rows}: residual differs from routed_moe raw -> post_moe_hop"
+                )
+                assert torch.equal(w_out, r_out), (
+                    f"L{layer} R={rows}: output differs from routed_moe raw -> post_moe_hop"
+                )
+            else:
+                ok, n_diff = G.reorder_ok(mine, r_payload, True)
+                assert ok, f"L{layer} R={rows}: partial outside the FP32 reorder bound"
+                reordered += n_diff
+            st = U.compare_stats(w_res, n_res)
+            exact_native += int(st["max_ulp"] == 0)
+            worst_native[0] = max(worst_native[0], st["max_ulp"])
+            worst_native[1] = max(worst_native[1], st["max_abs"])
+            torch.testing.assert_close(w_out, n_out, rtol=0.05, atol=0.15)
+            total += 1
+    CB.set_features("")
+    allw = env.comm.allgather((exact_native, total, worst_native, reordered))
+    if env.rank == 0:
+        rel = "bit for bit" if G.exact_mode(hop) else "within the FP32 reorder bound"
+        print(
+            f"[routed_moe_into_hop] mode={mode} R={rows}: routed_moe_into_hop == composition (routed_moe "
+            f"raw -> post_moe_hop) {rel} on all ranks; "
+            f"residual vs native exact in {[a[0] for a in allw]} of {total} windows per rank, "
+            f"worst ulp/abs {[a[2] for a in allw]}, reordered partial elements {[a[3] for a in allw]}"
+        )
+
+
+def test_routed_into_hop_cuda_graph_replay(env, stack):
+    moes, gammas, hop, eng = stack
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    rows = 6
+    layer = LAYERS[0]
+    wins = _windows(layer, rows, 5)
+    x = wins[0].clone()
+    g = torch.Generator(device="cuda").manual_seed(7)
+    res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
+    CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            _arm(
+                moes[layer],
+                hop,
+                x,
+                res,
+                gammas[layer],
+                "post_moe_hop,routed_moe,routed_moe_into_hop",
+            )
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    graph = torch.cuda.CUDAGraph()
+    with with_multi_stream(True):
+        with torch.cuda.graph(graph, stream=stream):
+            (g_out, g_res), _ = _arm(
+                moes[layer],
+                hop,
+                x,
+                res,
+                gammas[layer],
+                "post_moe_hop,routed_moe,routed_moe_into_hop",
+            )
+    torch.cuda.synchronize()
+    for w in wins[1:]:
+        x.copy_(w)
+        env.comm.barrier()
+        graph.replay()
+        torch.cuda.synchronize()
+        (e_out, e_res), _ = _arm(
+            moes[layer], hop, x, res, gammas[layer], "post_moe_hop,routed_moe,routed_moe_into_hop"
+        )
+        torch.cuda.synchronize()
+        assert G.same_or_close(hop, (g_out, g_res), (e_out, e_res)), "graph replay differs"
+    CB.set_features("")
+    hop.raise_if_error()
+
+
+def test_routed_into_hop_falls_back_to_raw_without_routed_input(env, stack):
+    moes, gammas, hop, eng = stack
+    layer = LAYERS[0]
+    x = _windows(layer, 6, 1)[0]
+    moe = moes[layer]
+    saved = moe._routed_into_hop
+    try:
+        moe._routed_into_hop = None
+        CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
+        out = moe.forward(x, do_finalize=False)
+        assert len(out) == 4, (
+            "without the hop's routed input routed_moe_into_hop must keep the raw tuple"
+        )
+    finally:
+        moe._routed_into_hop = saved
+        CB.set_features("")
+
+
+_WARM = set()
+
+
+def _warm(env, moes, rows_list):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe as routed_moe_mod,
+    )
+
+    key = (tuple(rows_list), routed_moe_mod.DELAY_INJECTION)
+    if key not in _WARM:
+        G.warm(env, moes, rows_list)
+        _WARM.add(key)
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_routed_into_hop_shared_absent_variant(env, stack, rows):
+    moes, gammas, hop, eng = stack
+    _warm(env, moes, [1, 6, 8])
+    ok, n, fails = G.run_mixed(
+        env, moes, gammas, hop, "all-routed_moe_into_hop", False, [rows], 3, seed=rows
+    )
+    hop.raise_if_error()
+    if env.rank == 0:
+        print(
+            f"[routed_moe_into_hop] shared absent R={rows}: {n} calls, routed variant == raw variant"
+        )
+    assert U.all_ranks(ok), f"rank {env.rank}: {fails}"
+
+
+@pytest.mark.parametrize("with_shared", [True, False], ids=["shared", "no_shared"])
+@pytest.mark.parametrize(
+    "pattern", ["odd-routed_moe", "rank0-routed_moe", "host1-routed_moe", "odd-native"]
+)
+def test_routed_into_hop_mixed_rank_variants(env, stack, pattern, with_shared):
+    moes, gammas, hop, eng = stack
+    _warm(env, moes, [1, 6, 8])
+    ok, n, fails = G.run_mixed(env, moes, gammas, hop, pattern, with_shared, [1, 6, 8], 2)
+    hop.raise_if_error()
+    producers = env.comm.allgather(G.mixed_pattern(env, pattern))
+    if env.rank == 0:
+        print(
+            f"[routed_moe_into_hop] mixed {pattern} shared={with_shared}: producers {producers}, {n} calls "
+            f"bit-equal to the all-raw variant"
+        )
+    assert U.all_ranks(ok), f"rank {env.rank}: {fails}"
+
+
+@pytest.mark.parametrize("pattern", ["odd-routed_moe", "odd-native"])
+def test_routed_into_hop_mixed_rank_graph(env, stack, pattern):
+    moes, gammas, hop, eng = stack
+    _warm(env, moes, [1, 6, 8])
+    ok, replays = G.run_graph(env, moes, gammas, hop, pattern)
+    hop.raise_if_error()
+    if env.rank == 0:
+        print(f"[routed_moe_into_hop] graph mixed {pattern}: {replays} replays x 4 calls bit-equal")
+    assert U.all_ranks(ok)
+
+
+@pytest.mark.parametrize("kernel_delay", [False, True], ids=["no_kernel_delay", "kernel_delay"])
+def test_routed_into_hop_delayed_stress(env, stack, monkeypatch, kernel_delay):
+    moes, gammas, hop, eng = stack
+    if kernel_delay:
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+            routed_moe as routed_moe_mod,
+        )
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel import (  # noqa: E501
+            DELAY_ENTRY,
+        )
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel_xc import (  # noqa: E501
+            DELAY_FINALIZE,
+        )
+
+        monkeypatch.setattr(routed_moe_mod, "DELAY_INJECTION", (DELAY_ENTRY | DELAY_FINALIZE, 3000))
+    _warm(env, moes, [1, 6, 8])
+    ok, stats = G.run_stress(
+        env, moes, gammas, hop, batches=6, batch=16, max_delay_us=1500.0, seed=int(kernel_delay)
+    )
+    allstats = env.comm.allgather(stats)
+    if env.rank == 0:
+        print(f"[routed_moe_into_hop] delayed stress kernel_delay={kernel_delay}: {allstats}")
+    hop.raise_if_error()
+    assert U.all_ranks(ok), f"rank {env.rank}: {stats}"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py
new file mode 100644
index 000000000000..24abaaf27e14
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py
@@ -0,0 +1,400 @@
+# 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.
+"""routed_moe on the served module types (sm_103, real GLM-5.2 NVFP4 weights).
+
+The routed experts are built the way ``Deepseekv3MoE`` builds them:
+``create_moe`` with the served TP8/EP1 ``ModelConfig`` (rank ``GLM_MOE_DSA_ROUTED_MOE_RANK``),
+which returns a ``ConfigurableMoE`` around the resolved TRT-LLM-Gen leaf, and
+loaded from unsliced checkpoint tensors so the native loader does the TP
+slicing. routed_moe is called through ``dispatch.get_impl("routed_moe", "routed_output")``
+with that wrapper and compared bit for bit with ``wrapper.forward`` (the
+served native path) and with the hand-mirrored ``native_forward`` oracle of
+the other tests.
+
+``test_adopted_deepseek_moe_*`` run the served MoE: the built-in
+``Deepseekv3MoE`` adopted by the target's ``GlmMoE``, called inside a decode
+step.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(
+    0,
+    os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "..", "shared", "kernels"
+    ),
+)
+import plan_testlib  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
+RANK = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def full():
+    return fx.load_layer_full(LAYER)
+
+
+@pytest.fixture(scope="module")
+def served(full):
+    wrapper, bias = fx.build_served_experts(LAYER, RANK, full)
+    return wrapper, bias
+
+
+@pytest.fixture(scope="module")
+def leaf_layer():
+    return fx.build_native_layer(layer=LAYER, tp_rank=RANK)
+
+
+@pytest.fixture(scope="module")
+def leaf_runner(leaf_layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    return RoutedMoe.from_native(leaf_layer.module, leaf_layer.bias)
+
+
+@pytest.fixture(scope="module")
+def decoder(leaf_layer):
+    return fx.ExpertDecoder(leaf_layer, leaf_layer.hidden_proxy(fx.prompt_token_ids()[:1])[0])
+
+
+def _impl():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    return dispatch.get_impl("routed_moe", "routed_output")
+
+
+def _inputs(leaf_layer, rows, kind, seed):
+    ids = fx.prompt_token_ids()
+    if kind == "proxy":
+        start = (seed * 7) % (len(ids) - rows)
+        x = leaf_layer.hidden_proxy(ids[start : start + rows])
+        return x, leaf_layer.router_logits(x)
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+    if kind == "random":
+        return x, leaf_layer.router_logits(x)
+    d = int(kind[len("distinct") :])
+    sets = fx.expert_sets_with_distinct(rows, d, seed=seed)
+    return x, fx.logits_for_expert_sets(sets, leaf_layer.bias, seed=seed)
+
+
+def test_served_wrapper_type_and_resolution(served):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        eligibility,
+        resolve_expert_owner,
+    )
+    from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE
+
+    wrapper, _ = served
+    leaf = wrapper.backend
+    print(
+        f"[routed_moe] served experts: {type(wrapper).__name__} -> {type(leaf).__name__} "
+        f"provider={getattr(type(leaf), 'provider', None)} comm={wrapper.comm} "
+        f"tp={leaf.tp_size}/{leaf.tp_rank} ep={leaf.ep_size} "
+        f"w3_w1={tuple(leaf.w3_w1_weight.shape)}"
+    )
+    assert isinstance(wrapper, ConfigurableMoE)
+    # the wrapper only delegates: the tensors live on the leaf
+    assert "w3_w1_weight" not in wrapper._parameters
+    owner, reason = resolve_expert_owner(wrapper)
+    assert owner is leaf, reason
+    x = torch.zeros(6, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
+    lg = torch.zeros(6, fx.NUM_EXPERTS, device="cuda", dtype=torch.float32)
+    ok, reason = eligibility(wrapper, x, lg)
+    assert ok, reason
+
+
+def test_served_weights_equal_fixture_slices(served, leaf_layer):
+    """The native TP8 loader's rank shard equals the fixture's hand-sliced
+    shard byte for byte, so the leaf-level tests run on served buffers.
+    """
+    leaf = served[0].backend
+    ref = leaf_layer.module
+    for name in (
+        "w3_w1_weight",
+        "w3_w1_weight_scale",
+        "w2_weight",
+        "w2_weight_scale",
+        "fc31_alpha",
+        "fc31_scale_c",
+        "fc2_alpha",
+        "fc31_input_scale",
+    ):
+        a, b = getattr(leaf, name), getattr(ref, name)
+        assert a.shape == b.shape and a.dtype == b.dtype, name
+        assert torch.equal(a.reshape(-1).view(torch.uint8), b.reshape(-1).view(torch.uint8)), name
+    assert torch.equal(served[1], leaf_layer.bias)
+
+
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+@pytest.mark.parametrize("kind", ["proxy", "random", "distinct40"])
+def test_served_wrapper_matches_served_forward(
+    served, leaf_layer, leaf_runner, decoder, rows, kind
+):
+    wrapper, bias = served
+    if kind == "distinct40" and rows * fx.TOP_K < 40:
+        kind = "distinct8"
+    x, logits = _inputs(leaf_layer, rows, kind, seed=50 + rows)
+    out = _impl()(wrapper, bias, x, logits, True)
+    assert out is not None, "routed_moe did not engage on the served wrapper"
+    served_native = wrapper.forward(x, logits, do_finalize=True, output_dtype=torch.bfloat16)
+    mirrored = leaf_layer.native_forward(x, logits)
+    torch.cuda.synchronize()
+    # the hand-mirrored oracle is the served path, bit for bit
+    assert torch.equal(mirrored, served_native)
+    raw = _impl()(wrapper, bias, x, logits, False)
+    raw_native = wrapper.forward(x, logits, do_finalize=False, output_dtype=torch.bfloat16)
+    mirrored_raw = leaf_layer.native_forward(x, logits, do_finalize=False)
+    assert torch.equal(fx.raw_rows_by_topk(mirrored_raw), fx.raw_rows_by_topk(raw_native))
+    cls, rep = fx.classify_vs_native(
+        out,
+        raw,
+        served_native,
+        raw_native,
+        lambda: fx.emulate_moe(leaf_layer, x, logits),
+        f"rows={rows} {kind}",
+        lambda: decoder.compare(leaf_runner, logits),
+    )
+    print(f"[routed_moe] served wrapper rows={rows} {kind}: {cls} {rep}")
+
+
+def test_served_wrapper_fallbacks(served, leaf_layer):
+    """Wrapper-level behaviour routed_moe does not implement selects the native path
+    before any device work.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        _WORKSPACES,
+    )
+
+    wrapper, bias = served
+    x, logits = _inputs(leaf_layer, 6, "proxy", seed=3)
+    snapshot = {k: v.ws.clone() for k, v in _WORKSPACES.items()}
+    for attr, value in (
+        ("unfinalized_combine_fn", lambda *a, **k: None),
+        ("reduce_results", True),
+        ("apply_router_weight_on_input", True),
+        ("_enable_perfect_router", True),
+    ):
+        saved = getattr(wrapper, attr, None)
+        setattr(wrapper, attr, value)
+        try:
+            assert _impl()(wrapper, bias, x, logits, True) is None, attr
+        finally:
+            setattr(wrapper, attr, saved)
+    # FP32 activations and more than eight rows stay native
+    assert _impl()(wrapper, bias, x.float(), logits, True) is None
+    x9 = torch.cat([x, x[:3]])
+    assert _impl()(wrapper, bias, x9, torch.cat([logits, logits[:3]]), True) is None
+    torch.cuda.synchronize()
+    for k, v in _WORKSPACES.items():
+        assert torch.equal(v.ws, snapshot[k])
+    assert _impl()(wrapper, bias, x, logits, True) is not None
+
+
+# ---------------------------------------------------------------------------
+# The served MoE: Deepseekv3MoE adopted by the target's GlmMoE.
+# ---------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def deepseek_moe(full):
+    return fx.build_deepseek_moe(LAYER, RANK, full)
+
+
+class _Features:
+    """Run exactly the fast paths ``value`` in a block run as a decode step, and count
+    routed_moe engagements; the selector state before the block is restored after it.
+    """
+
+    def __init__(self, value: str):
+        self.value = value
+        self.engaged = 0
+        self.declined = 0
+
+    def __enter__(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+        self.dispatch = dispatch
+        self.step = dispatch.step(dispatch.Phase.DECODE)
+        self.step.__enter__()
+        self.saved_state = (dispatch._INSTALLED[0], dispatch._VALIDATE[0], dispatch._require[0])
+        plan_testlib.select(dispatch, self.value)
+        dispatch.get_impl("routed_moe", "routed_output")
+        self.saved_fn = dispatch._REGISTRY["routed_moe"]["routed_output"]
+
+        def counting(*args, **kwargs):
+            out = self.saved_fn(*args, **kwargs)
+            if out is None:
+                self.declined += 1
+            else:
+                self.engaged += 1
+            return out
+
+        dispatch._REGISTRY["routed_moe"]["routed_output"] = counting
+        return self
+
+    def __exit__(self, *exc):
+        self.dispatch._REGISTRY["routed_moe"]["routed_output"] = self.saved_fn
+        (
+            self.dispatch._INSTALLED[0],
+            self.dispatch._VALIDATE[0],
+            self.dispatch._require[0],
+        ) = self.saved_state
+        self.step.__exit__(*exc)
+        return False
+
+
+def _hidden(rows, seed):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    return (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+
+
+def _assert_forward_equivalent(on, off, emul_fn, label, id_check=None):
+    """forward(do_finalize=False) results [shared, fc2, scales, idx]: the
+    shared expert output is identical; the routed raw tuples agree on every
+    expert and scale, and their stock finalization is exact or in budget.
+    """
+    assert torch.equal(on[0], off[0]), f"{label}: shared expert output differs"
+    raw_on, raw_off = on[1:], off[1:]
+    cls, rep = fx.classify_vs_native(
+        fx.finalize_raw(raw_on),
+        raw_on,
+        fx.finalize_raw(raw_off),
+        raw_off,
+        emul_fn,
+        label,
+        id_check,
+    )
+    return cls
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_adopted_deepseek_moe_forward(deepseek_moe, leaf_layer, leaf_runner, decoder, rows):
+    """The served MoE forward (shared expert on the aux stream in
+    parallel, gate, routed experts) with ``routed_moe`` on matches it with ``routed_moe``
+    off (same experts and scales, shared output identical, routed output exact
+    or in budget), and routed_moe actually engaged.
+    """
+    moe = deepseek_moe
+    classes = []
+    for it in range(4):
+        h = _hidden(rows, 1000 * rows + it)
+        with _Features("routed_moe") as f:
+            on = moe.forward(h, do_finalize=False)
+            fin_on = moe.compute_routed_output(h, None, None, True)
+        with _Features("") as f_off:
+            off = moe.forward(h, do_finalize=False)
+            fin_off = moe.compute_routed_output(h, None, None, True)
+        torch.cuda.synchronize()
+        assert f.engaged == 2 and f.declined == 0, (f.engaged, f.declined)
+        assert f_off.engaged == 0 and f_off.declined == 0
+        logits = moe.gate(h)
+
+        def emul(h=h, logits=logits):
+            return fx.emulate_moe(leaf_layer, h, logits)
+
+        def ids(logits=logits):
+            return decoder.compare(leaf_runner, logits)
+
+        label = f"rows={rows} it={it}"
+        classes.append(_assert_forward_equivalent(on, off, emul, label, ids))
+        cls, _ = fx.classify_vs_native(fin_on, on[1:], fin_off, off[1:], emul, label, ids)
+        classes.append(cls)
+    print(f"[routed_moe] adopted Deepseekv3MoE rows={rows}: classes {classes}")
+
+
+def test_adopted_deepseek_moe_graph(deepseek_moe, leaf_layer, leaf_runner, decoder):
+    """CUDA graph of the served forward with ``routed_moe`` on (eager warmup first,
+    as the model's graph runner does), replayed on new inputs and compared
+    with eager native forwards; overlapping shared-expert GEMMs on the aux
+    stream run in every replay.
+    """
+    moe = deepseek_moe
+    rows = 6
+    h_static = _hidden(rows, 1)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with _Features("routed_moe") as f:
+        with torch.cuda.stream(stream):
+            for _ in range(2):
+                moe.forward(h_static, do_finalize=False)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            out_static = moe.forward(h_static, do_finalize=False)
+        torch.cuda.synchronize()
+    assert f.engaged == 3 and f.declined == 0, (f.engaged, f.declined)
+    for it in range(16):
+        h = _hidden(rows, 500 + it)
+        h_static.copy_(h)
+        graph.replay()
+        torch.cuda.synchronize()
+        with _Features("") as f_off:
+            ref = moe.forward(h, do_finalize=False)
+        torch.cuda.synchronize()
+        assert f_off.engaged == 0
+        logits = moe.gate(h)
+        _assert_forward_equivalent(
+            [t.clone() for t in out_static],
+            ref,
+            lambda h=h, logits=logits: fx.emulate_moe(leaf_layer, h, logits),
+            f"graph replay {it}",
+            lambda logits=logits: decoder.compare(leaf_runner, logits),
+        )
+
+
+def test_adopted_deepseek_moe_capture_without_warmup_falls_back(full):
+    """A routed_moe call whose first use is inside a CUDA graph capture takes the
+    native path (no compilation or workspace initialization is recorded).
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe,
+    )
+
+    moe = fx.build_deepseek_moe(LAYER, RANK, full)
+    h = _hidden(6, 7)
+    # warm the native path, which the capture below records
+    with _Features(""):
+        moe.compute_routed_output(h, None, None, True)
+    torch.cuda.synchronize()
+    saved_ws = dict(routed_moe._WORKSPACES)
+    saved_fn = dict(routed_moe._COMPILED)
+    routed_moe._WORKSPACES.clear()
+    routed_moe._COMPILED.clear()
+    try:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with _Features("routed_moe") as f:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                moe.compute_routed_output(h, None, None, True)
+        assert f.declined == 1 and f.engaged == 0
+        assert not routed_moe._WORKSPACES and not routed_moe._COMPILED
+    finally:
+        routed_moe._WORKSPACES.update(saved_ws)
+        routed_moe._COMPILED.update(saved_fn)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py
new file mode 100644
index 000000000000..4840fd9ba137
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py
@@ -0,0 +1,176 @@
+# 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.
+"""routed_moe's raw tuple through the native unfinalized consumer (one 4-GPU node).
+
+On a single-node TP group the served DeepSeek-V3 layer hands the
+``do_finalize=False`` MoE output to ``MoEAllReduce``, i.e. the native
+``moe_finalize_allreduce`` op: a local finalize that follows the
+permutation, the shared-expert add, the all-reduce, the residual add and the
+RMSNorm. Run every test on the 4 GPUs of one node::
+
+    srun -N1 --ntasks-per-node=4 --mpi=pmix python3 -m pytest ...
+    test_routed_moe_native_consumer_mpi.py
+
+Rank ``r`` holds TP8 shard ``r`` of a real GLM-5.2 NVFP4 layer (``GLM_MOE_DSA_ROUTED_MOE_LAYER``,
+default 10) and runs captured real hidden-state windows of that layer. Per call
+every rank feeds the same consumer three raw tuples:
+
+* ``routed_moe``: routed_moe's token-major rows (row ``t * 8 + k``) with the identity
+  permutation;
+* ``routed_moe_relocated``: the same rows written at the native runner's permuted
+  positions, with the native permutation and the native expert scales;
+* ``native``: the native runner's raw tuple.
+
+``routed_moe`` and ``routed_moe_relocated`` must give identical outputs bit for bit: the
+consumer reads the token-major layout exactly like the native expert-grouped
+layout. Where every rank's routed_moe rows equal the native rows bit for bit, ``routed_moe``
+and ``native`` must give identical outputs as well; the other calls (an FP4
+rounding flip inside routed_moe's frozen error budget) are counted and must stay
+close.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+sys.path.insert(
+    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
+)
+import post_moe_hop_mpi_utils as U  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
+WINDOWS = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_CONSUMER_WINDOWS", "8"))
+
+
+@pytest.fixture(scope="module")
+def env():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("needs an SM100-family GPU")
+    e = U.mpi_env()
+    if e.world < 2 or e.multi_node:
+        pytest.skip("needs a single-node TP group of at least 2 ranks")
+    return e
+
+
+@pytest.fixture(scope="module")
+def setup(env):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+    from tensorrt_llm._torch.distributed import MoEAllReduce
+
+    layer = fx.build_native_layer(layer=LAYER, tp_rank=env.rank)
+    runner = RoutedMoe.from_native(layer.module, layer.bias)
+    consumer = MoEAllReduce(env.mapping)
+    # compile routed_moe's variants on every rank before the first collective call
+    for rows in (1, 6, 8):
+        x = torch.zeros(rows, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
+        logits = layer.router_logits(x)
+        runner(x, logits, do_finalize=False)
+        layer.native_forward(x, logits, do_finalize=False)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    return layer, runner, consumer
+
+
+def _windows(rows: int, count: int):
+    win = [x for x, _ in fx.real_windows(fx.load_real_capture(LAYER), rows)]
+    assert win and win[0] is not None, f"no captured hidden states for layer {LAYER}"
+    step = max(1, len(win) // count)
+    return [win[i * step].cuda().contiguous() for i in range(count)]
+
+
+def _consume(consumer, raw, shared, residual, gamma):
+    from tensorrt_llm._torch.distributed import MoEAllReduceParams
+
+    fc2, scale, idx = raw
+    params = MoEAllReduceParams(
+        expanded_idx_to_permuted_idx=idx,
+        expert_scale_factor=scale,
+        shared_expert_output=shared,
+        residual=residual,
+        norm_weight=gamma,
+        eps=fx.RMS_EPS,
+        is_cutlass_min_latency=False,
+    )
+    return consumer(fc2, all_reduce_params=params)
+
+
+def _relocate(routed_moe_raw, native_raw):
+    """routed_moe's rows at the native runner's permuted positions."""
+    fc2_w, _, idx_w = routed_moe_raw
+    fc2_n, scale_n, idx_n = native_raw
+    assert fc2_n.shape[1] == fc2_w.shape[1], (fc2_n.shape, fc2_w.shape)
+    out = torch.zeros_like(fc2_n)
+    out[idx_n.reshape(-1).long()] = fc2_w[idx_w.reshape(-1).long()]
+    return out, scale_n, idx_n
+
+
+def _same(a, b) -> bool:
+    return all(torch.equal(x.view(torch.int16), y.view(torch.int16)) for x, y in zip(a, b))
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_raw_tuple_through_moe_finalize_allreduce(env, setup, rows):
+    layer, runner, consumer = setup
+    g_rank = torch.Generator(device="cuda").manual_seed(300 + 17 * env.rank + rows)
+    g_common = torch.Generator(device="cuda").manual_seed(500 + rows)
+    gamma = layer.norm_weight
+    exact_calls, calls, worst = 0, 0, [0, 0.0]
+    for x in _windows(rows, WINDOWS):
+        logits = layer.router_logits(x)
+        shared = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g_rank) * 0.1).bfloat16()
+        residual = torch.randn(rows, fx.HIDDEN, device="cuda", generator=g_common) * 4.0
+        residual = residual.bfloat16()
+        native_raw = layer.native_forward(x, logits, do_finalize=False)
+        routed_moe_raw = runner(x, logits, do_finalize=False)
+        torch.cuda.synchronize()
+        assert routed_moe_raw[0].shape == (rows * fx.TOP_K, fx.HIDDEN)
+        ident = torch.arange(rows * fx.TOP_K, dtype=torch.int32)
+        assert torch.equal(routed_moe_raw[2].view(-1).cpu(), ident), (
+            "routed_moe permutation is not t * 8 + k"
+        )
+        assert (native_raw[2] >= 0).all(), "TP8/EP1: every (t, k) has a local native row"
+        assert torch.equal(routed_moe_raw[1], native_raw[1]), "expert scales differ from native"
+        rep = fx.raw_report(routed_moe_raw, native_raw)
+        assert rep["same_expert"], rep
+        rows_exact = U.all_ranks(rep["rows_exact"] == rep["rows"])
+
+        out_w = _consume(consumer, routed_moe_raw, shared, residual, gamma)
+        out_r = _consume(consumer, _relocate(routed_moe_raw, native_raw), shared, residual, gamma)
+        out_n = _consume(consumer, native_raw, shared, residual, gamma)
+        torch.cuda.synchronize()
+        assert _same(out_w, out_r), "token-major layout read differently from the native layout"
+        if rows_exact:
+            assert _same(out_w, out_n), "identical rows, different consumer output"
+            exact_calls += 1
+        st = U.compare_stats(out_w[1], out_n[1])
+        worst[0] = max(worst[0], st["max_ulp"])
+        worst[1] = max(worst[1], st["max_abs"])
+        torch.testing.assert_close(out_w[0], out_n[0], rtol=0.05, atol=0.15)
+        calls += 1
+    allw = env.comm.allgather((exact_calls, calls, worst))
+    if env.rank == 0:
+        print(
+            f"[routed_moe] native consumer R={rows}: routed_moe raw == routed_moe relocated to the native layout "
+            f"bit for bit on {env.world} ranks; routed_moe == native in {[a[0] for a in allw]} of "
+            f"{calls} calls (all-rank-exact rows); residual vs native worst ulp/abs "
+            f"{[a[2] for a in allw]}"
+        )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py
new file mode 100644
index 000000000000..09abab9d414a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py
@@ -0,0 +1,722 @@
+# 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.
+"""sm_103 correctness tests of the routed_moe routed NVFP4 MoE against the native region.
+
+Real GLM-5.2 NVFP4 weights (layer ``GLM_MOE_DSA_ROUTED_MOE_LAYER``, default 10, of TP8 rank
+``GLM_MOE_DSA_ROUTED_MOE_RANK``, default 0) through the stock TRT-LLM-Gen leaf are the oracle.
+Error budget (frozen before any performance tuning):
+expert IDs and routing weights bit-exact; finalized output
+``nrms(routed_moe - native) <= 1.5 * nrms(emul - native) + 2e-3`` where ``emul`` is
+the FP64-accumulated emulation, and no row may exceed 4x the per-tensor
+budget.
+
+Native expert IDs and their top-k order are compared through the native raw
+tuple: for every ``(t, k)`` the expert scale equals native's bit for bit and
+the routed_moe FC2 row is the row of the same expert as native's row at
+``expanded_idx_to_permuted_idx[t, k]`` (see ``routed_moe_fixtures.classify_vs_native``
+for the exact / reorder / budget classes of the floating-point results).
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import routed_moe_fixtures as fx  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+DEBUG_ROUTING = 1
+DEBUG_QUANT = 2
+DEBUG_FC1 = 4
+
+
+@pytest.fixture(scope="module")
+def layer():
+    torch.manual_seed(0)
+    return fx.build_native_layer(
+        layer=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10")),
+        tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")),
+    )
+
+
+@pytest.fixture(scope="module")
+def runner(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    return RoutedMoe.from_native(layer.module, layer.bias)
+
+
+@pytest.fixture(scope="module")
+def decoder(layer):
+    return fx.ExpertDecoder(layer, layer.hidden_proxy(fx.prompt_token_ids()[:1])[0])
+
+
+def _inputs(layer, rows: int, kind: str, seed: int = 0):
+    ids = fx.prompt_token_ids()
+    if kind == "proxy":
+        start = (seed * 7) % (len(ids) - rows)
+        x = layer.hidden_proxy(ids[start : start + rows])
+        return x, layer.router_logits(x)
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+    if kind == "random":
+        return x, layer.router_logits(x)
+    if kind.startswith("distinct"):
+        d = int(kind[len("distinct") :])
+        sets = fx.expert_sets_with_distinct(rows, d, seed=seed)
+        return x, fx.logits_for_expert_sets(sets, layer.bias, seed=seed)
+    raise ValueError(kind)
+
+
+def _dbg_buffer():
+    return torch.zeros(2 * 1024 * 1024, dtype=torch.int32, device="cuda")
+
+
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+@pytest.mark.parametrize("kind", ["proxy", "random"])
+def test_routing_matches_native_bitwise(layer, runner, decoder, rows, kind):
+    x, logits = _inputs(layer, rows, kind, seed=rows)
+    dbg = _dbg_buffer()
+    runner(x, logits, debug=DEBUG_ROUTING, debug_buffer=dbg)
+    torch.cuda.synchronize()
+    ids_ref, w_ref = fx.routing_reference(logits, layer.bias)
+    ids = dbg[:64].view(8, 8)[:rows].cpu()
+    w = dbg[64:128].view(torch.float32).view(8, 8)[:rows].cpu()
+    assert torch.equal(ids, ids_ref.cpu()), (ids, ids_ref)
+    assert torch.equal(w, w_ref.float().cpu())
+    num_slots = int(dbg[256].item())
+    expect = sorted(set(ids_ref.flatten().tolist()))
+    assert num_slots == len(expect)
+    assert dbg[128 : 128 + num_slots].cpu().tolist() == expect
+    # native unfinalized expert weights, same top-k order
+    outs = layer.native_forward(x, logits, do_finalize=False)
+    assert torch.equal(outs[1].float().cpu(), w_ref.float().cpu())
+    # native expert IDs and order, through the raw rows
+    _assert_raw_matches_native(
+        runner(x, logits, do_finalize=False), outs, lambda: decoder.compare(runner, logits)
+    )
+    # and decoded exactly through the probe fingerprints
+    dec = decoder.compare(runner, logits)
+    assert dec["ids_equal"] and dec["scale_exact"], dec
+    assert torch.equal(dec["native_ids"].cpu().int(), ids_ref.cpu()), (dec["native_ids"], ids_ref)
+
+
+def _assert_raw_matches_native(raw_routed_moe, raw_native, id_check=None):
+    """Every (t, k): same expert scale bits and same expert as native."""
+    rep = fx.raw_report(raw_routed_moe, raw_native)
+    assert rep["scale_exact"], (raw_routed_moe[1], raw_native[1])
+    if not rep["same_expert"]:
+        assert id_check is not None, rep
+        dec = id_check()
+        assert dec["ids_equal"] and dec["scale_exact"] and dec["margin"] > 4, (rep, dec)
+    return rep
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_activation_quant_matches_native_bitwise(layer, runner, rows):
+    x, logits = _inputs(layer, rows, "proxy", seed=3)
+    dbg = _dbg_buffer()
+    runner(x, logits, debug=DEBUG_QUANT, debug_buffer=dbg)
+    torch.cuda.synchronize()
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel import (  # noqa: E501
+        SFB_BYTES,
+        SFB_CHUNK_STRIDE,
+        SFB_STAGE_STRIDE,
+    )
+
+    raw = dbg.view(torch.uint8)[4096 : 4096 + 24 * 1024 + SFB_BYTES]
+    sb = raw[: 24 * 1024].view(24, 1024)
+    sfb_flat = raw[24 * 1024 :]
+    xq, xsf = layer.module.quantize_input(x)
+    xq = xq.view(torch.uint8).view(rows, fx.HIDDEN // 2)
+    xsf = xsf.view(torch.uint8).view(rows, fx.HIDDEN // 16)
+    kb = torch.arange(24, device="cuda")[:, None, None]
+    n = torch.arange(rows, device="cuda")[None, :, None]
+    b = torch.arange(128, device="cuda")[None, None, :]
+    phys = n * 128 + (((b // 16) ^ n) * 16) + (b % 16)
+    got = torch.gather(sb[:, None, :].expand(24, rows, 1024), 2, phys.expand(24, rows, 128))
+    want = xq.view(rows, 24, 128).permute(1, 0, 2)
+    assert torch.equal(got, want)
+    f = torch.arange(16, device="cuda")[None, None, :]
+    sf_phys = kb * SFB_STAGE_STRIDE + (f // 4) * SFB_CHUNK_STRIDE + n * 16 + f % 4
+    got_sf = sfb_flat[sf_phys.expand(24, rows, 16)]
+    want_sf = xsf.view(rows, 24, 16).permute(1, 0, 2)
+    assert torch.equal(got_sf, want_sf)
+
+
+@pytest.mark.parametrize("rows,kind", [(1, "proxy"), (6, "proxy"), (8, "distinct40")])
+def test_fc1_accumulators_match_exact(layer, runner, rows, kind):
+    x, logits = _inputs(layer, rows, kind, seed=5)
+    dbg = _dbg_buffer()
+    runner(x, logits, debug=DEBUG_ROUTING | DEBUG_FC1, debug_buffer=dbg)
+    torch.cuda.synchronize()
+    num_slots = int(dbg[256].item())
+    slots = dbg[128 : 128 + num_slots].tolist()
+    acc = dbg.view(torch.float32)[65536 : 65536 + num_slots * 4 * 128 * 8].view(
+        num_slots, 4 * 128, 8
+    )
+    m = layer.module
+    xq, xsf = m.quantize_input(x)
+    xd = fx.unpack_fp4(xq.view(torch.uint8)) * fx.e4m3_bytes_to_float(
+        xsf.view(torch.uint8).reshape(rows, fx.HIDDEN // 16)
+    ).repeat_interleave(16, dim=1)
+    worst = 0.0
+    for s, e in enumerate(slots):
+        w13, _ = fx.dequant_expert(m, e)
+        ref = w13 @ xd.T  # [512, rows]
+        got = acc[s, :, :rows].double()
+        scale = ref.abs().max().item() + 1e-30
+        worst = max(worst, ((got - ref).abs().max().item()) / scale)
+    assert worst < 1e-5, worst
+
+
+def _check_vs_native(layer, out, x, logits, label):
+    native = layer.native_forward(x, logits)
+    emul = fx.emulate_moe(layer, x, logits)
+    e_native = fx.error_stats(out, native)
+    e_emul = fx.error_stats(emul, native)
+    e_self = fx.error_stats(out, emul)
+    print(
+        f"[routed_moe] {label}: routed_moe-vs-native {e_native} emul-vs-native {e_emul} routed_moe-vs-emul {e_self}"
+    )
+    budget = 1.5 * e_emul["nrms"] + 2e-3
+    assert e_native["nrms"] <= budget, (e_native, e_emul)
+    row_ref = native.float().pow(2).mean(-1).sqrt()
+    row_err = (out.float() - native.float()).pow(2).mean(-1).sqrt()
+    assert torch.all(row_err <= 4 * budget * row_ref + 1e-6), (row_err, row_ref)
+    return e_native
+
+
+@pytest.mark.parametrize("rows", [1, 2, 4, 6, 8])
+@pytest.mark.parametrize("kind", ["proxy", "random", "distinct8"])
+def test_finalized_matches_native(layer, runner, rows, kind):
+    x, logits = _inputs(layer, rows, kind, seed=11 + rows)
+    out = runner(x, logits)
+    torch.cuda.synchronize()
+    _check_vs_native(layer, out, x, logits, f"rows={rows} {kind}")
+
+
+def test_raw_tuple_consumable_by_stock_finalizer(layer, runner):
+    rows = 6
+    x, logits = _inputs(layer, rows, "proxy", seed=2)
+    fc2, scale, idx = runner(x, logits, do_finalize=False)
+    out_final = runner(x, logits)
+    torch.cuda.synchronize()
+    assert fc2.shape == (rows * 8, fx.HIDDEN) and fc2.dtype == torch.bfloat16
+    assert scale.shape == (rows, 8) and scale.dtype == torch.bfloat16
+    assert idx.shape == (rows, 8) and idx.dtype == torch.int32
+    _, w_ref = fx.routing_reference(logits, layer.bias)
+    assert torch.equal(scale, w_ref)
+    # stock finalize semantics: sum_k scale[t,k] * fc2[idx[t,k]] in FP32
+    fin = torch.zeros(rows, fx.HIDDEN, device="cuda", dtype=torch.float32)
+    for k in range(8):
+        fin += scale[:, k : k + 1].float() * fc2[idx[:, k].long()].float()
+    fin = fin.to(torch.bfloat16)
+    _check_vs_native(layer, fin, x, logits, "raw tuple")
+    e = fx.error_stats(fin, out_final)
+    assert e["nrms"] < 2e-3, e
+
+
+@pytest.mark.parametrize("offset", [1, 4, 7])
+def test_misaligned_activation_is_copied_not_faulted(layer, runner, offset):
+    """A contiguous activation view whose base is not 16-byte aligned (the bulk
+    copy's requirement) gives the same raw and finalized outputs as the aligned
+    tensor: the runner copies it instead of faulting or changing path."""
+    rows = 6
+    x, logits = _inputs(layer, rows, "proxy", seed=3)
+    store = torch.empty(rows * fx.HIDDEN + offset, device="cuda", dtype=torch.bfloat16)
+    x_mis = store[offset:].view(rows, fx.HIDDEN)
+    x_mis.copy_(x)
+    assert x_mis.is_contiguous() and x_mis.data_ptr() % 16 != 0
+    if runner.schedule(rows, True) != "xc" or runner.schedule(rows, False) != "xc":
+        pytest.skip("bit-exact repeat calls are a property of the expert-cluster schedule")
+    # Outputs may live in the runner's workspace: copy each before the next call.
+    ref = runner(x, logits).clone()
+    got = runner(x_mis, logits).clone()
+    raw_ref = [t.clone() for t in runner(x, logits, do_finalize=False)]
+    raw_got = [t.clone() for t in runner(x_mis, logits, do_finalize=False)]
+    torch.cuda.synchronize()
+    assert torch.equal(got, ref)
+    for a, b in zip(raw_got, raw_ref):
+        assert torch.equal(a, b)
+
+
+def test_cuda_graph_replay(layer, runner):
+    rows = 6
+    x0, l0 = _inputs(layer, rows, "proxy", seed=1)
+    x_static = x0.clone()
+    l_static = l0.clone()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        runner(x_static, l_static)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        out_static = runner(x_static, l_static)
+    for it in range(12):
+        kind = ["proxy", "random", "distinct8", "distinct48"][it % 4]
+        x, logits = _inputs(layer, rows, kind, seed=100 + it)
+        x_static.copy_(x)
+        l_static.copy_(logits)
+        graph.replay()
+        torch.cuda.synchronize()
+        eager = runner(x, logits)
+        torch.cuda.synchronize()
+        e = fx.error_stats(out_static, eager)
+        assert e["nrms"] < 1e-6, (it, kind, e)
+        if it < 4:
+            _check_vs_native(layer, out_static.clone(), x, logits, f"graph replay {it} {kind}")
+
+
+def _real_inputs(layer, rows: int, count: int, offset: int = 0):
+    """``count`` windows of the captured hidden states of this layer and their
+    router logits (GPU)."""
+    win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer.layer), rows)]
+    assert win and win[0] is not None, f"no captured hidden states for layer {layer.layer}"
+    step = max(1, len(win) // count)
+    xs = [win[(offset + i * step) % len(win)].cuda().contiguous() for i in range(count)]
+    return [(x, layer.router_logits(x)) for x in xs]
+
+
+def _expert_ids(logits, bias):
+    ids, _ = fx.routing_reference(logits, bias)
+    return set(ids.flatten().tolist())
+
+
+def _paddings(layer, x, logits, pads: int, other, seed: int):
+    """Padded-row contents ``{name: (x_pad, logits_pad)}`` for ``pads`` rows."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    rand = (torch.randn(pads, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+    out = {
+        "zeros": torch.zeros(pads, fx.HIDDEN, dtype=torch.bfloat16, device="cuda"),
+        "scaled_copy": (x[:1].float() * 3.0).to(torch.bfloat16).repeat(pads, 1),
+        "real": other[:pads].contiguous(),
+        "random": rand,
+    }
+    out = {k: (v, layer.router_logits(v)) for k, v in out.items()}
+    # padded rows whose top-8 sets avoid every live expert: the call's
+    # distinct-expert count grows by 8 per padded row (two rounds when it
+    # exceeds the cluster count)
+    used = _expert_ids(logits, layer.bias)
+    free = [
+        e
+        for e in torch.randperm(
+            fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(seed)
+        ).tolist()
+        if e not in used
+    ]
+    sets = [free[8 * i : 8 * i + 8] for i in range(pads)]
+    out["new_experts"] = (rand, fx.logits_for_expert_sets(sets, layer.bias, seed=seed))
+    return out
+
+
+@pytest.mark.parametrize("live", [1, 2, 4, 6])
+def test_padded_rows_do_not_change_live_rows(layer, runner, live):
+    """Padding invariance of the default (expert-cluster) schedule, bit for bit.
+
+    Real hidden-state windows of ``live`` rows are padded to 8 rows with
+    zeros, a scaled copy of a live row, rows of another real window, random
+    rows, and rows routed to experts no live row uses (more distinct experts,
+    up to a two-round call). The live rows' raw FC2 rows, expert scales and
+    finalized output must equal the unpadded call's bit for bit.
+    """
+    if runner.schedule(8, True) != "xc" or runner.schedule(8, False) != "xc":
+        pytest.skip("padding invariance is a property of the expert-cluster schedule")
+    pads = 8 - live
+    clusters = runner._xc_ctas(8, True) // 4
+    others = _real_inputs(layer, 8, 3, offset=5)
+    seen = {}
+    for i, (x, logits) in enumerate(_real_inputs(layer, live, 3, offset=live)):
+        fin = runner(x, logits).clone()
+        raw = [t.clone() for t in runner(x, logits, do_finalize=False)]
+        torch.cuda.synchronize()
+        for name, (xp, lp) in _paddings(layer, x, logits, pads, others[i][0], 70 + i).items():
+            xa = torch.cat([x, xp]).contiguous()
+            la = torch.cat([logits, lp.float()]).contiguous()
+            fin_p = runner(xa, la)
+            raw_p = runner(xa, la, do_finalize=False)
+            torch.cuda.synchronize()
+            d = len(_expert_ids(la, layer.bias))
+            seen.setdefault(name, []).append(d)
+            tag = f"live={live} window {i} padding {name} D={d}"
+            assert torch.equal(fin_p[:live].view(torch.int16), fin.view(torch.int16)), tag
+            rows_p = raw_p[0][: live * 8].view(torch.int16)
+            assert torch.equal(rows_p, raw[0].view(torch.int16)), tag
+            assert torch.equal(raw_p[1][:live], raw[1]), tag
+            assert torch.isfinite(fin_p).all(), tag
+    two_round = [d for ds in seen.values() for d in ds if d > clusters]
+    print(
+        f"[routed_moe] padding invariance live={live}: bit-exact for paddings {seen} "
+        f"(distinct experts per call; {len(two_round)} two-round calls, {clusters} clusters)"
+    )
+    assert two_round, "no padded call reached the two-round schedule"
+
+
+REORDER_REL = 2.0**-20
+
+
+def test_fp32_accumulator_output_reorder_class(layer, runner):
+    """routed_moe_into_hop "fp32" output: the kernel adds the eight weighted rows of a token
+    with atomic FP32 adds. Every run, with or without padded rows, lies within
+    the FP32 reorder bound of the exact sum of the same products (the raw
+    rows' ``scale * row``, each exact in FP32); the run-to-run spread is
+    counted, not assumed."""
+    rows = 6
+    if runner.schedule(8, False, True) != "xc":
+        pytest.skip("FP32 accumulator output needs the expert-cluster schedule")
+    others = _real_inputs(layer, 8, 3, offset=9)
+    distinct_runs = []
+    for i, (x, logits) in enumerate(_real_inputs(layer, rows, 3, offset=2)):
+        fc2, scale, idx = runner(x, logits, do_finalize=False)
+        terms = scale.double()[:, :, None] * fc2[idx.long()].double()
+        exact = terms.sum(1)
+        bound = REORDER_REL * terms.abs().sum(1)
+        runs = []
+        for _ in range(8):
+            acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
+            runner(x, logits, red_out=acc)
+            runs.append(acc)
+        xp = torch.cat([x, others[i][0][:2]]).contiguous()
+        lp = torch.cat([logits, layer.router_logits(others[i][0][:2]).float()]).contiguous()
+        acc8 = torch.zeros(8, fx.HIDDEN, dtype=torch.float32, device="cuda")
+        runner(xp, lp, red_out=acc8)
+        torch.cuda.synchronize()
+        for j, acc in enumerate(runs + [acc8[:rows]]):
+            assert ((acc.double() - exact).abs() <= bound).all(), f"window {i} run {j}"
+        distinct_runs.append(len({r.cpu().numpy().tobytes() for r in runs}))
+    print(
+        f"[routed_moe] fp32 accumulator: within the FP32 reorder bound in every run; distinct "
+        f"results over 8 runs per window {distinct_runs}"
+    )
+
+
+def test_workspace_created_once_in_device_memory(layer, runner):
+    """The per-device workspace is device memory even when first requested
+    inside the served loader's ``MetaInitMode``, is never replaced, and a
+    request for more CTA slots than it holds raises."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe as rm,
+    )
+    from tensorrt_llm._torch.models.modeling_utils import MetaInitMode
+
+    dev = runner.device
+    key = dev.index if dev.index is not None else torch.cuda.current_device()
+    saved = rm._WORKSPACES.pop(key, None)
+    try:
+        with MetaInitMode():
+            assert torch.empty(1, device="cuda").is_meta, "MetaInitMode is not active"
+            fresh = rm.RoutedMoe.from_native(layer.module, layer.bias)
+            ws = rm._workspace(fresh.device, fresh.num_ctas)
+        for name, t in ws.tensors().items():
+            assert not t.is_meta and t.is_cuda and t.device.index == key, name
+            assert t.data_ptr() != 0, name
+        assert ws.num_ctas >= rm._sm_count(dev)
+        assert rm._workspace(dev, fresh.num_ctas) is ws
+        with pytest.raises(RuntimeError):
+            rm._workspace(dev, ws.num_ctas + 4)
+        assert rm._WORKSPACES[key] is ws
+        print(f"[routed_moe] workspace cuda:{key}: {ws.num_ctas} CTA slots, {ws.nbytes()} bytes")
+    finally:
+        if saved is not None:
+            rm._WORKSPACES[key] = saved
+
+
+def test_eager_calls_from_two_streams_are_ordered(layer, runner):
+    """A call issued on another stream than the previous call waits for the
+    previous stream's submitted work before it uses the shared workspace."""
+    rows = 6
+    x, logits = _inputs(layer, rows, "proxy", seed=4)
+    x2, l2 = _inputs(layer, rows, "distinct40", seed=6)
+    ref = runner(x, logits).clone()
+    ref2 = runner(x2, l2).clone()
+    torch.cuda.synchronize()
+    a, b = torch.cuda.Stream(), torch.cuda.Stream()
+    a.wait_stream(torch.cuda.current_stream())
+    b.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(a):
+        torch.cuda._sleep(100_000_000)  # tens of milliseconds of device time
+        out_a = runner(x, logits)
+    with torch.cuda.stream(b):
+        out_b = runner(x2, l2)
+        done_b = torch.cuda.Event()
+        done_b.record(b)
+    done_b.synchronize()
+    assert a.query(), "the call on the second stream did not wait for the first stream"
+    torch.cuda.synchronize()
+    assert torch.equal(out_a.view(torch.int16), ref.view(torch.int16))
+    assert torch.equal(out_b.view(torch.int16), ref2.view(torch.int16))
+
+
+def test_fallback_is_decided_before_device_work(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        _WORKSPACES,
+        routed_output,
+    )
+
+    x9 = torch.randn(9, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
+    l9 = torch.randn(9, 256, device="cuda")
+    snapshot = {k: v.ws.clone() for k, v in _WORKSPACES.items()}
+    assert routed_output(layer.module, layer.bias, x9, l9, True) is None
+    x1 = torch.randn(1, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
+    assert (
+        routed_output(
+            layer.module,
+            layer.bias,
+            x1,
+            torch.randn(1, 256, device="cuda").to(torch.bfloat16),
+            True,
+        )
+        is None
+    )
+    assert (
+        routed_output(
+            layer.module, layer.bias, x1.float(), torch.randn(1, 256, device="cuda"), True
+        )
+        is None
+    )
+    # FP32 router bias (non-TRTLLM gate), a bias tensor that is not the routing
+    # method's, and modules without registered expert parameters stay native
+    l1 = torch.randn(1, 256, device="cuda")
+    assert routed_output(layer.module, layer.bias.float(), x1, l1, True) is None
+    assert routed_output(layer.module, layer.bias.clone(), x1, l1, True) is None
+    assert routed_output(torch.nn.Linear(4, 4), layer.bias, x1, l1, True) is None
+
+    class _DelegatingOnly(torch.nn.Module):
+        """Exposes the expert tensors as properties but owns none of them."""
+
+        w3_w1_weight = property(lambda self: layer.module.w3_w1_weight)
+        w2_weight = property(lambda self: layer.module.w2_weight)
+
+    assert routed_output(_DelegatingOnly(), layer.bias, x1, l1, True) is None
+    torch.cuda.synchronize()
+    for k, v in _WORKSPACES.items():
+        assert torch.equal(v.ws, snapshot[k])
+
+
+def test_dispatch_registration():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    fn = dispatch.get_impl("routed_moe", "routed_output")
+    assert callable(fn)
+    assert callable(dispatch.get_impl("routed_moe", "eligibility"))
+
+
+# ---------------------------------------------------------------------------
+# Near-tie routing: FP32 logits built so that the 8th/9th selection keys tie
+# exactly, or differ by one key ulp, including experts with equal BF16 bias.
+# ---------------------------------------------------------------------------
+TIE_PATTERNS = (
+    "same_logit",  # equal bias, identical logits: exact key tie
+    "same_key",  # equal bias, different logits, exact key tie
+    "plus_ulp",  # equal bias, key(b) = key(a) + 1 ulp
+    "minus_ulp",  # equal bias, key(b) = key(a) - 1 ulp
+    "cross_bias_lo",  # bias(a) < bias(b), exact key tie
+    "cross_bias_hi",  # bias(a) > bias(b), exact key tie
+    "three_way",  # three equal keys for the last two slots
+    "inner_tie",  # two winners with equal keys (top-k order only)
+)
+_WIN_LOGIT = 6.0
+_CONT_LOGIT = 0.3
+_BOOST = 0.02
+
+
+def _tie_row(bias, pattern: str, gen: torch.Generator):
+    """One row of logits for ``pattern`` plus its resolution variants.
+
+    Returns ``(row, variants)``: ``variants`` are logit rows in which each
+    possible resolution of the tie wins clearly (by ``_BOOST`` in logit); the
+    native outcome of ``row`` must equal exactly one of them.
+    """
+    b = bias.float().cpu()
+    vals = sorted(set(b.tolist()))
+    groups = [[e for e in range(fx.NUM_EXPERTS) if b[e].item() == v] for v in vals]
+    big = [g for g in groups if len(g) >= 3]
+    perm = lambda n: torch.randperm(n, generator=gen).tolist()  # noqa: E731
+    row = -12.0 + 0.5 * torch.rand(fx.NUM_EXPERTS, generator=gen)
+    if pattern in ("cross_bias_lo", "cross_bias_hi"):
+        # cont[0] is the lower expert index; its bias is the lower one for
+        # "lo" and the higher one for "hi"
+        lo_g = groups[0] + (groups[1] if len(groups) > 1 else [])
+        hi_g = groups[-1] + (groups[-2] if len(groups) > 2 else [])
+        want_lo_first = pattern == "cross_bias_lo"
+        cands = [(x, y) for x in lo_g for y in hi_g if (x < y) == want_lo_first]
+        assert cands, pattern
+        x, y = cands[perm(len(cands))[0]]
+        cont = [x, y] if want_lo_first else [y, x]
+    else:
+        g = big[perm(len(big))[0]]
+        pick = [g[i] for i in perm(len(g))[:3]]
+        cont = sorted(pick[:3] if pattern == "three_way" else pick[:2])
+    n_win = {"three_way": 6, "inner_tie": 6}.get(pattern, 7)
+    others = [e for e in perm(fx.NUM_EXPERTS) if e not in cont]
+    winners = others[:n_win]
+    for i, e in enumerate(winners):
+        row[e] = _WIN_LOGIT + 0.01 * i
+    bdev = bias.float().cuda()
+    a = cont[0]
+    if pattern == "inner_tie":
+        # two contenders become winners with identical keys; no boundary tie
+        for e in cont[:2]:
+            row[e] = _WIN_LOGIT + 0.5
+    else:
+        for e in cont:
+            row[e] = _CONT_LOGIT
+    key_a = fx._key_of(torch.tensor(row[a].item(), device="cuda"), bdev[a]).item()
+    c = cont[1]
+    if pattern == "same_key" or pattern.startswith("cross_bias"):
+        hit = fx.logit_for_key(key_a, bdev[c], 0, avoid=row[a].item())
+        assert hit is not None, pattern
+        row[c] = hit
+    elif pattern in ("plus_ulp", "minus_ulp"):
+        hit = fx.logit_for_key(key_a, bdev[c], 1 if pattern == "plus_ulp" else -1)
+        assert hit is not None, pattern
+        row[c] = hit
+    variants = []
+    if pattern == "three_way":
+        for pair in ((0, 1), (0, 2), (1, 2)):
+            v = row.clone()
+            for i in pair:
+                v[cont[i]] += _BOOST
+            variants.append(v)
+    else:
+        for e in cont[:2]:
+            v = row.clone()
+            v[e] += _BOOST
+            variants.append(v)
+    return row, variants
+
+
+def _tie_batch(bias, patterns, seed):
+    gen = torch.Generator().manual_seed(seed)
+    rows, variants = [], []
+    for p in patterns:
+        r, v = _tie_row(bias, p, gen)
+        rows.append(r)
+        variants.append(v)
+    logits = torch.stack(rows).cuda().contiguous()
+    nvar = max(len(v) for v in variants)
+    var_batches = [
+        torch.stack([v[min(i, len(v) - 1)] for v in variants]).cuda().contiguous()
+        for i in range(nvar)
+    ]
+    return logits, var_batches, [len(v) for v in variants]
+
+
+def _check_ties(layer, runner, decoder, x, logits, var_batches, nvars, label):
+    native_raw = layer.native_forward(x, logits, do_finalize=False)
+    native_fin = layer.native_forward(x, logits)
+    routed_moe_raw = runner(x, logits, do_finalize=False)
+    routed_moe_fin = runner(x, logits)
+    torch.cuda.synchronize()
+    tie_rows = fx.raw_rows_by_topk(native_raw)
+    var_rows = [
+        fx.raw_rows_by_topk(layer.native_forward(x, v, do_finalize=False)) for v in var_batches
+    ]
+    resolution = []
+    for t in range(x.shape[0]):
+        match = [i for i in range(nvars[t]) if torch.equal(tie_rows[t], var_rows[i][t])]
+        distinct = len({tuple(var_rows[i][t].flatten()[::97].tolist()) for i in range(nvars[t])})
+        assert distinct == nvars[t], f"{label} row {t}: variants not distinguishable"
+        assert len(match) == 1, f"{label} row {t}: native outcome matches variants {match}"
+        resolution.append(match[0])
+    cls, rep = fx.classify_vs_native(
+        routed_moe_fin,
+        routed_moe_raw,
+        native_fin,
+        native_raw,
+        lambda: fx.emulate_moe(layer, x, logits),
+        label,
+        lambda: decoder.compare(runner, logits),
+    )
+    # the decoded native experts of every tie row, independent of the rows
+    dec = decoder.compare(runner, logits)
+    assert dec["ids_equal"] and dec["scale_exact"], (label, dec)
+    gaps = fx.boundary_gap(logits, layer.bias)
+    print(
+        f"[routed_moe] ties {label}: resolution={resolution} class={cls} "
+        f"rows_exact={rep['rows_exact']}/{rep['rows']} decode_margin={dec['margin']:.1f} "
+        f"boundary_gap={gaps.tolist()}"
+    )
+    return resolution
+
+
+@pytest.mark.parametrize("seed", [0, 1, 2, 3])
+def test_boundary_ties_match_native(layer, runner, decoder, seed):
+    """All eight tie patterns in one R=8 call; routed_moe equals native bit for bit
+    and the native outcome is identified as one clear resolution.
+    """
+    logits, var_batches, nvars = _tie_batch(layer.bias, TIE_PATTERNS, seed)
+    x, _ = _inputs(layer, len(TIE_PATTERNS), "proxy", seed=20 + seed)
+    res = _check_ties(layer, runner, decoder, x, logits, var_batches, nvars, f"R=8 seed={seed}")
+    # identical logits and bias give identical keys in any arithmetic; they
+    # resolve to the lower expert index (variant 0), as the native 64-bit key
+    # (score bits, 65535 - index) orders them
+    for p, r in zip(TIE_PATTERNS, res):
+        if p in ("same_logit", "three_way", "inner_tie"):
+            assert r == 0, (p, r)
+
+
+@pytest.mark.parametrize("pattern", TIE_PATTERNS)
+def test_boundary_tie_single_row(layer, runner, decoder, pattern):
+    logits, var_batches, nvars = _tie_batch(layer.bias, [pattern], 7)
+    x, _ = _inputs(layer, 1, "proxy", seed=31)
+    _check_ties(layer, runner, decoder, x, logits, var_batches, nvars, f"R=1 {pattern}")
+
+
+def test_near_tie_sweep_matches_native(layer, runner, decoder):
+    """Randomized near-boundary rows: the 8th/9th keys at -2..+2 key ulps
+    (random contender biases and logits); routed_moe equals native bit for bit.
+    """
+    gen = torch.Generator().manual_seed(77)
+    bdev = layer.bias.float().cuda()
+    batches = 48
+    exact_ties = 0
+    classes = {}
+    for i in range(batches):
+        rows = []
+        for t in range(8):
+            row = -12.0 + 0.5 * torch.rand(fx.NUM_EXPERTS, generator=gen)
+            order = torch.randperm(fx.NUM_EXPERTS, generator=gen).tolist()
+            for k, e in enumerate(order[:7]):
+                row[e] = _WIN_LOGIT + 0.01 * k
+            a, c = order[7], order[8]
+            row[a] = 0.1 + 0.6 * torch.rand(1, generator=gen).item()
+            key_a = fx._key_of(torch.tensor(row[a].item(), device="cuda"), bdev[a]).item()
+            off = int(torch.randint(-2, 3, (1,), generator=gen).item())
+            hit = fx.logit_for_key(key_a, bdev[c], off)
+            row[c] = hit if hit is not None else row[a]
+            rows.append(row)
+        logits = torch.stack(rows).cuda().contiguous()
+        exact_ties += int((fx.boundary_gap(logits, layer.bias) == 0).sum().item())
+        x, _ = _inputs(layer, 8, "random", seed=200 + i)
+        cls, _ = fx.check_call(
+            layer, runner, x, logits, label=f"near-tie batch {i}", decoder=decoder
+        )
+        dec = decoder.compare(runner, logits)
+        assert dec["ids_equal"] and dec["scale_exact"], (i, dec)
+        classes[cls] = classes.get(cls, 0) + 1
+    print(
+        f"[routed_moe] near-tie sweep: {batches * 8} rows, {exact_ties} exact boundary key ties, "
+        f"batch classes {classes}"
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py
new file mode 100644
index 000000000000..c3a19dd2c013
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py
@@ -0,0 +1,184 @@
+# 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.
+"""routed_moe expert-cluster kernel with package cross_layer_expert_prefetch's speculative expert prefetch.
+
+The prefetch (``spec_k`` > 0) only issues ``cp.async.bulk.prefetch.L2``, so
+every output must equal the prefetch-off kernel bit for bit, for any content
+of the expert-ID row (including out-of-range IDs). The self-prediction store
+(``spec_self``) must write the call's distinct experts ranked by row count.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import routed_moe_fixtures as fx  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+KMAX = 32
+REORDER_REL = 2.0**-20
+
+
+@pytest.fixture(scope="module")
+def layer():
+    torch.manual_seed(0)
+    return fx.build_native_layer(
+        layer=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10")),
+        tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")),
+    )
+
+
+@pytest.fixture(scope="module")
+def runner(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    r = RoutedMoe.from_native(layer.module, layer.bias)
+    if r.schedule(8, False, True) != "xc":
+        pytest.skip("the speculative prefetch is part of the expert-cluster schedule")
+    return r
+
+
+class _Spec:
+    """Install a cross_layer_expert_prefetch hook on the runner's expert leaf."""
+
+    def __init__(self, runner, ids, k, ctas=0, self_store=False):
+        self.owner = runner.module
+        self.spec = (ids, k, ctas, self_store)
+
+    def __enter__(self):
+        self.owner._expert_prefetch_spec = lambda: self.spec
+        return self
+
+    def __exit__(self, *exc):
+        del self.owner._expert_prefetch_spec
+
+
+def _windows(layer, rows, count, offset=0):
+    win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer.layer), rows)]
+    if not win:
+        pytest.skip(
+            f"no captured hidden states of layer {layer.layer} (GLM_MOE_DSA_REAL_MOE_INPUTS_DIR)"
+        )
+    step = max(1, len(win) // count)
+    xs = [win[(offset + i * step) % len(win)].cuda().contiguous() for i in range(count)]
+    return [(x, layer.router_logits(x)) for x in xs]
+
+
+def _ids(kind, seed):
+    g = torch.Generator().manual_seed(seed)
+    if kind == "valid":
+        ids = torch.randperm(fx.NUM_EXPERTS, generator=g)[:KMAX]
+    else:  # out of range: the kernel masks each ID into [0, 256)
+        ids = torch.randint(-(2**31), 2**31 - 1, (KMAX,), generator=g)
+    return ids.to(torch.int32).cuda()
+
+
+def _outputs(runner, x, logits):
+    fin = runner(x, logits)
+    fc2, scale, _ = runner(x, logits, do_finalize=False)
+    acc = torch.zeros(x.shape[0], fx.HIDDEN, dtype=torch.float32, device="cuda")
+    runner(x, logits, red_out=acc)
+    torch.cuda.synchronize()
+    return fin.clone(), fc2.clone(), scale.clone(), acc
+
+
+def _check_same(ref, got, what):
+    fin, fc2, scale, _ = ref
+    gfin, gfc2, gscale, gacc = got
+    assert torch.equal(gfin, fin), f"{what}: finalized output differs"
+    assert torch.equal(gfc2, fc2), f"{what}: raw FC2 rows differ"
+    assert torch.equal(gscale, scale), f"{what}: routing weights differ"
+    # FP32 accumulator: atomic adds, compare within the FP32 reorder bound of
+    # the exact sum of the same raw products
+    rows = gacc.shape[0]
+    terms = scale.double()[:, :, None] * fc2.view(rows, 8, -1).double()
+    exact = terms.sum(1)
+    bound = REORDER_REL * terms.abs().sum(1)
+    assert ((gacc.double() - exact).abs() <= bound).all(), f"{what}: accumulator"
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("k,ctas", [(12, 0), (32, 48), (1, 4)])
+@pytest.mark.parametrize("ids_kind", ["valid", "out_of_range"])
+def test_prefetch_outputs_bit_identical(layer, runner, rows, k, ctas, ids_kind):
+    for i, (x, logits) in enumerate(_windows(layer, rows, 2, offset=rows)):
+        ref = _outputs(runner, x, logits)
+        ids = _ids(ids_kind, seed=rows * 100 + i)
+        with _Spec(runner, ids, k, ctas):
+            got = _outputs(runner, x, logits)
+        _check_same(ref, got, f"rows={rows} k={k} ctas={ctas} {ids_kind} window {i}")
+
+
+def test_prefetch_cuda_graph_replay(layer, runner):
+    rows = 6
+    wins = _windows(layer, rows, 6, offset=3)
+    x_static, l_static = (t.clone() for t in wins[0])
+    ids = _ids("valid", seed=7)
+    with _Spec(runner, ids, 12):
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            runner(x_static, l_static, do_finalize=False)
+        torch.cuda.current_stream().wait_stream(stream)
+        # the eager call compiled the prefetch variant: the capture records it
+        assert runner.ready_for_capture(False, rows=rows)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            fc2_s, scale_s, _ = runner(x_static, l_static, do_finalize=False)
+    for it, (x, logits) in enumerate(wins):
+        x_static.copy_(x)
+        l_static.copy_(logits)
+        ids.copy_(_ids("valid", seed=100 + it))
+        graph.replay()
+        torch.cuda.synchronize()
+        fc2, scale, _ = runner(x, logits, do_finalize=False)
+        torch.cuda.synchronize()
+        assert torch.equal(fc2_s, fc2) and torch.equal(scale_s, scale), it
+
+
+def _ranked_experts(logits, bias):
+    ids, _ = fx.routing_reference(logits, bias)
+    counts = torch.bincount(ids.flatten().long().cpu(), minlength=fx.NUM_EXPERTS)
+    distinct = [e for e in range(fx.NUM_EXPERTS) if counts[e] > 0]
+    return sorted(distinct, key=lambda e: (-int(counts[e]), e))
+
+
+@pytest.mark.parametrize("rows", [1, 3, 6, 8])
+def test_self_store_ranks_by_row_count(layer, runner, rows):
+    for i, (x, logits) in enumerate(_windows(layer, rows, 3, offset=5 * rows)):
+        ref = _outputs(runner, x, logits)
+        ids = torch.full((KMAX,), 999, dtype=torch.int32, device="cuda")
+        with _Spec(runner, ids, 12, self_store=True):
+            fc2, scale, _ = runner(x, logits, do_finalize=False)
+            torch.cuda.synchronize()
+            got = ids.tolist()
+            # the store is idempotent for the same call (the prefetch reads
+            # the previous call's row first)
+            gfin = runner(x, logits)
+            torch.cuda.synchronize()
+        assert torch.equal(fc2, ref[1]) and torch.equal(scale, ref[2]), (rows, i)
+        assert torch.equal(gfin, ref[0]), (rows, i)
+        want = _ranked_experts(logits, layer.bias)[:KMAX]
+        assert got[: len(want)] == want, (rows, i, got, want)
+        assert all(v == 999 for v in got[len(want) :]), (rows, i, got)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py
new file mode 100644
index 000000000000..afe1bf6302c9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py
@@ -0,0 +1,812 @@
+# 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.
+"""Stress, scheduling and real-input tests of the routed_moe cross-CTA protocol (sm_103).
+
+* ``test_sanitizer_cells``: routed_moe-vs-native cells in both output modes plus a
+  short graph, meant to be run under ``compute-sanitizer`` (memcheck /
+  synccheck / racecheck, kernel filter ``kns=RoutedMoeKernel``) in both
+  FC1->FC2 hand-off modes.
+* ``test_delayed_*``: the kernel's delay injection (``routed_moe.DELAY_INJECTION``) holds a
+  pseudo-random third of the CTAs (re-drawn every launch) at CTA entry, before
+  publishing FC1 tiles and split partials, before split owners gather and
+  before FC2 operand fetches; outputs must stay bit-identical to undelayed
+  launches.
+* ``test_corun_*``: other kernels on a second stream hold SMs while routed_moe runs
+  (a shared-memory-heavy spinning kernel, cuBLAS GEMMs), so not all routed_moe CTAs
+  are co-resident; routed_moe must complete with bit-identical outputs.
+* ``test_graph_many_launches_real_layers``: one CUDA graph with dozens of
+  back-to-back routed_moe launches over several real layers (one shared workspace),
+  every output equal to eager routed_moe and classified against native.
+* ``test_real_*``: a native capture of real routed-MoE inputs.
+
+Comparisons with native use ``routed_moe_fixtures.classify_vs_native`` (exact
+expert IDs and scales; finalized output exact, FP32-reorder-close or within
+the frozen budget).
+
+Layers: ``GLM_MOE_DSA_ROUTED_MOE_REAL_LAYERS`` (default ``10,30,50,77``, the layers with captured
+hidden states) of TP8 rank ``GLM_MOE_DSA_ROUTED_MOE_RANK``.
+"""
+
+import faulthandler
+import os
+import sys
+import time
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import pytest
+import torch
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.nvgpu import tcgen05
+from cutlass.cutlass_dsl import Int32, Int64, T, Uint32, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+    ptx,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import routed_moe_fixtures as fx  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+RANK = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
+LAYERS = [
+    int(v) for v in os.environ.get("GLM_MOE_DSA_ROUTED_MOE_REAL_LAYERS", "10,30,50,77").split(",")
+]
+WATCHDOG_S = 180
+
+
+@pytest.fixture(scope="module")
+def layers():
+    return {lay: fx.build_native_layer(layer=lay, tp_rank=RANK) for lay in LAYERS}
+
+
+@pytest.fixture(scope="module")
+def runners(layers):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    return {lay: RoutedMoe.from_native(m.module, m.bias) for lay, m in layers.items()}
+
+
+@pytest.fixture(scope="module")
+def decoders(layers):
+    probe = layers[LAYERS[0]].hidden_proxy(fx.prompt_token_ids()[:1])[0]
+    return {lay: fx.ExpertDecoder(m, probe) for lay, m in layers.items()}
+
+
+@pytest.fixture(scope="module")
+def captures():
+    return {
+        lay: [(None if h is None else h.cuda(), lg.cuda()) for h, lg in fx.load_real_capture(lay)]
+        for lay in LAYERS
+    }
+
+
+def _synthetic(layer, rows, distinct, seed):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+    sets = fx.expert_sets_with_distinct(rows, distinct, seed=seed)
+    return x, fx.logits_for_expert_sets(sets, layer.bias, seed=seed)
+
+
+def _real(captures, lay, rows, index):
+    win = fx.real_windows(captures[lay], rows)
+    x, lg = win[(index * 131) % len(win)]
+    return x.cuda().contiguous(), lg.cuda().contiguous()
+
+
+CELLS = [(1, 8), (2, 14), (4, 25), (6, 21), (6, 38), (8, 48)]
+
+
+def _cells(layer, captures, lay, seed=0):
+    out = [(f"R{r}D{d}", *_synthetic(layer, r, d, seed + r * 100 + d)) for r, d in CELLS]
+    for i in range(3):
+        out.append((f"real{lay}R6#{i}", *_real(captures, lay, 6, seed + i)))
+    out.append((f"real{lay}R8", *_real(captures, lay, 8, seed + 7)))
+    return out
+
+
+def _identical(a, b, label):
+    """routed_moe against routed_moe: bit-identical finalized outputs or raw tuples."""
+    if isinstance(a, torch.Tensor):
+        assert torch.equal(a, b), label
+    else:
+        for i in range(3):
+            assert torch.equal(a[i], b[i]), f"{label} (raw part {i})"
+
+
+def _routed_moe_both(runner, x, lg):
+    return runner(x, lg), runner(x, lg, do_finalize=False)
+
+
+def _classify(layer, x, lg, fin, raw, label, emul=True):
+    n_fin = layer.native_forward(x, lg)
+    n_raw = layer.native_forward(x, lg, do_finalize=False)
+    emul_fn = (lambda: fx.emulate_moe(layer, x, lg)) if emul else None
+    return fx.classify_vs_native(fin, raw, n_fin, n_raw, emul_fn, label)
+
+
+# ---------------------------------------------------------------------------
+# (a) compute-sanitizer target
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("mode", ["fin", "raw"])
+def test_sanitizer_cells(layers, runners, captures, mode):
+    """One output mode per run, so a sanitizer report names one compiled
+    kernel. Outputs are compared with native loosely here (the exact checks
+    are in the other tests). The graph segment needs the caching allocator
+    (``PYTORCH_NO_CUDA_MEMORY_CACHING=1`` would cudaMalloc during capture).
+    """
+    lay = LAYERS[0]
+    layer, runner = layers[lay], runners[lay]
+    fin = mode == "fin"
+    for label, x, lg in _cells(layer, captures, lay):
+        out = runner(x, lg, do_finalize=fin)
+        ref = layer.native_forward(x, lg, do_finalize=fin)
+        torch.cuda.synchronize()
+        if fin:
+            assert fx.error_stats(out, ref)["nrms"] < 2e-3, label
+        else:
+            rep = fx.raw_report(out, ref)
+            assert rep["scale_exact"] and rep["same_expert"], (label, rep)
+    if os.environ.get("PYTORCH_NO_CUDA_MEMORY_CACHING") == "1":
+        return
+    x, lg = _real(captures, lay, 6, 11)
+    eager = runner(x, lg, do_finalize=fin)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        outs = [runner(x, lg, do_finalize=fin) for _ in range(4)]
+    graph.replay()
+    torch.cuda.synchronize()
+    for i, o in enumerate(outs):
+        _identical(o, eager, f"graph launch {i}")
+
+
+class _BulkCopyControl:
+    """Racecheck control: the bulk-copy -> mbarrier (transaction count) ->
+    shared-load pattern routed_moe uses for its activation rows and FC2 operand
+    images, in isolation and correctly synchronized. Two phases on one
+    buffer: copy, wait, every thread reads, CTA barrier, copy again (the
+    write-after-read reuse), wait, read.
+    """
+
+    NBYTES = 2048
+
+    @cute.jit
+    def run(self, src: Int64, dst: Int64, stream):
+        @cute.struct
+        class Storage:
+            mbar: cute.struct.MemRange[cutlass.Int64, 2]
+            buf: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, 2048], 128]
+
+        self.storage_t = Storage
+        self.kernel(src, dst).launch(
+            grid=[1, 1, 1], block=[128, 1, 1], smem=Storage.size_in_bytes(), stream=stream
+        )
+
+    @cute.kernel
+    def kernel(self, src: Int64, dst: Int64):
+        tidx, _, _ = cute.arch.thread_idx()
+        smem = utils.SmemAllocator()
+        st = smem.allocate(self.storage_t)
+        bar = st.mbar.data_ptr()
+        buf = Int32(st.buf.data_ptr().toint())
+        if tidx == 0:
+            cute.arch.mbarrier_init(bar, 1)
+            cute.arch.mbarrier_init(bar + 1, 1)
+        cute.arch.mbarrier_init_fence()
+        cute.arch.sync_threads()
+        for phase in cutlass.range_constexpr(2):
+            if tidx == 0:
+                b = Int32((bar + phase).toint())
+                ptx.mbar_arrive_expect_tx(b, Int32(self.NBYTES))
+                ptx.bulk_g2s(buf, src + Int64(phase * self.NBYTES), Int32(self.NBYTES), b)
+            cute.arch.mbarrier_wait(bar + phase, 0)
+            for j in cutlass.range_constexpr(4):
+                v = ptx.lds_u32(buf + (j * 128 + tidx) * 4)
+                ptx.st_u32(dst + Int64(phase * self.NBYTES) + Int64((j * 128 + tidx) * 4), v)
+            # generic reads before the next async-proxy write of the buffer
+            ptx.fence_proxy_async_shared()
+            cute.arch.sync_threads()
+
+
+class _ChainControl:
+    """Racecheck control for routed_moe's shared-buffer reuse (FC1 B operand stages
+    overwritten by FC2 operands): warp 0 writes a buffer with generic stores,
+    proxy fence, mbarrier arrive; warp 1 waits and signals warp 2 through a
+    second mbarrier, with a plain arrive or with ``tcgen05.commit`` (routed_moe's
+    FC1-done signal); warp 2 waits, overwrites the buffer with a bulk copy,
+    waits for it and reads it back. Correctly ordered in both variants.
+    """
+
+    NBYTES = 1024
+
+    def __init__(self, use_commit: bool):
+        self.use_commit = use_commit
+
+    @cute.jit
+    def run(self, src: Int64, dst: Int64, stream):
+        @cute.struct
+        class Storage:
+            mbar: cute.struct.MemRange[cutlass.Int64, 3]
+            buf: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, 1024], 128]
+
+        self.storage_t = Storage
+        self.kernel(src, dst).launch(
+            grid=[1, 1, 1], block=[96, 1, 1], smem=Storage.size_in_bytes(), stream=stream
+        )
+
+    @cute.kernel
+    def kernel(self, src: Int64, dst: Int64):
+        tidx, _, _ = cute.arch.thread_idx()
+        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        lane = tidx % 32
+        smem = utils.SmemAllocator()
+        st = smem.allocate(self.storage_t)
+        bar = st.mbar.data_ptr()
+        buf = Int32(st.buf.data_ptr().toint())
+        if tidx == 0:
+            cute.arch.mbarrier_init(bar, 1)
+            cute.arch.mbarrier_init(bar + 1, 1)
+            cute.arch.mbarrier_init(bar + 2, 1)
+        cute.arch.mbarrier_init_fence()
+        cute.arch.sync_threads()
+        if warp == 0:
+            for j in cutlass.range_constexpr(8):
+                ptx.sts_u32(buf + (j * 32 + lane) * 4, Uint32(lane + 32 * j))
+            ptx.fence_proxy_async_shared()
+            cute.arch.sync_warp()
+            if lane == 0:
+                cute.arch.mbarrier_arrive(bar)
+        if warp == 1:
+            cute.arch.mbarrier_wait(bar, 0)
+            if cutlass.const_expr(self.use_commit):
+                with cute.arch.elect_one():
+                    tcgen05.commit(bar + 1)
+            else:
+                if lane == 0:
+                    cute.arch.mbarrier_arrive(bar + 1)
+        if warp == 2:
+            cute.arch.mbarrier_wait(bar + 1, 0)
+            if lane == 0:
+                b2 = Int32((bar + 2).toint())
+                ptx.mbar_arrive_expect_tx(b2, Int32(self.NBYTES))
+                ptx.bulk_g2s(buf, src, Int32(self.NBYTES), b2)
+            cute.arch.mbarrier_wait(bar + 2, 0)
+            for j in cutlass.range_constexpr(8):
+                v = ptx.lds_u32(buf + (j * 32 + lane) * 4)
+                ptx.st_u32(dst + Int64((j * 32 + lane) * 4), v)
+
+
+@pytest.mark.parametrize("use_commit", [False, True])
+def test_racecheck_control_chain(use_commit):
+    from cuda.bindings import driver as cuda
+
+    src = torch.randint(0, 255, (_ChainControl.NBYTES,), dtype=torch.uint8, device="cuda")
+    dst = torch.zeros_like(src)
+    cur = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    fn = cute.compile(_ChainControl(use_commit).run, cutlass.Int64(0), cutlass.Int64(0), cur)
+    fn(cutlass.Int64(src.data_ptr()), cutlass.Int64(dst.data_ptr()), cur)
+    torch.cuda.synchronize()
+    assert torch.equal(src, dst)
+
+
+def test_racecheck_control_bulk_copy():
+    from cuda.bindings import driver as cuda
+
+    src = torch.randint(0, 255, (2 * _BulkCopyControl.NBYTES,), dtype=torch.uint8, device="cuda")
+    dst = torch.zeros_like(src)
+    cur = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    ctl = _BulkCopyControl()
+    fn = cute.compile(ctl.run, cutlass.Int64(0), cutlass.Int64(0), cur)
+    fn(cutlass.Int64(src.data_ptr()), cutlass.Int64(dst.data_ptr()), cur)
+    torch.cuda.synchronize()
+    assert torch.equal(src, dst)
+
+
+# ---------------------------------------------------------------------------
+# (b) delayed producers / owners / consumers
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("delay_ns", [2000, 9000])
+@pytest.mark.parametrize("finalize", [True, False])
+def test_delayed_ctas_exact(layers, runners, captures, monkeypatch, delay_ns, finalize):
+    """Delayed launches are bit-identical to undelayed ones (eager and in a
+    replayed graph); the undelayed outputs are classified against native.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe as routed_moe_mod,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel import (  # noqa: E501
+        DELAY_CONSUMER,
+        DELAY_ENTRY,
+        DELAY_OWNER,
+        DELAY_PRODUCER,
+    )
+
+    lay = LAYERS[0]
+    layer, runner = layers[lay], runners[lay]
+    cells = _cells(layer, captures, lay, seed=delay_ns)
+    assert routed_moe_mod.DELAY_INJECTION[0] == 0
+    undelayed = []
+    classes = {}
+    for label, x, lg in cells:
+        fin, raw = _routed_moe_both(runner, x, lg)
+        cls, _ = _classify(layer, x, lg, fin, raw, label)
+        classes[cls] = classes.get(cls, 0) + 1
+        undelayed.append(fin if finalize else raw)
+    torch.cuda.synchronize()
+    sites = DELAY_ENTRY | DELAY_PRODUCER | DELAY_OWNER | DELAY_CONSUMER
+    monkeypatch.setattr(routed_moe_mod, "DELAY_INJECTION", (sites, delay_ns))
+    t0 = time.perf_counter()
+    for rep in range(3):  # every launch draws a new delayed set
+        for (label, x, lg), base in zip(cells, undelayed):
+            out = runner(x, lg, do_finalize=finalize)
+            torch.cuda.synchronize()
+            _identical(out, base, f"{label} rep{rep} delay={delay_ns}")
+    eager_s = time.perf_counter() - t0
+    # graph: all cells back to back, replayed (epochs advance per replay)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    static_inputs = [(x.clone(), lg.clone()) for _, x, lg in cells]
+    with torch.cuda.stream(stream):
+        runner(static_inputs[0][0], static_inputs[0][1], do_finalize=finalize)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        outs = [runner(x, lg, do_finalize=finalize) for x, lg in static_inputs]
+    for rep in range(4):
+        graph.replay()
+        torch.cuda.synchronize()
+        for (label, _, _), o, base in zip(cells, outs, undelayed):
+            _identical(o, base, f"graph {label} rep{rep}")
+    print(
+        f"[routed_moe] delayed CTAs bit-identical: delay_ns={delay_ns} finalize={finalize} "
+        f"eager {eager_s:.2f}s; undelayed vs native {classes}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# (c) co-running kernels that hold SMs
+# ---------------------------------------------------------------------------
+@dsl_user_op
+def _sreg_u32(name: str, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [],
+            f"mov.u32 $0, %{name};",
+            "=r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+class _SmHog:
+    """Test-only kernel: ``ctas`` CTAs that each reserve ``smem_bytes`` of
+    shared memory (more than half an SM, so one CTA per SM and no routed_moe CTA
+    beside it) and spin for ``spin_ns`` on the global timer. Each CTA records
+    ``started``, its SM id and its dynamic shared-memory size.
+    """
+
+    def __init__(self, ctas: int, smem_bytes: int, spin_ns: int):
+        self.ctas = ctas
+        self.smem_bytes = smem_bytes
+        self.spin_ns = spin_ns
+
+    @cute.jit
+    def run(self, info_addr: Int64, stream):
+        self.kernel(info_addr).launch(
+            grid=[self.ctas, 1, 1], block=[32, 1, 1], smem=self.smem_bytes, stream=stream
+        )
+
+    @cute.kernel
+    def kernel(self, info_addr: Int64):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        rec = info_addr + Int64(bidx) * Int64(16)
+        if tidx == 0:
+            ptx.st_u32(rec + Int64(4), _sreg_u32("smid"))
+            ptx.st_u32(rec + Int64(8), _sreg_u32("dynamic_smem_size"))
+            ptx.st_relaxed_gpu_u32(rec, Uint32(1))
+        t0 = ptx.globaltimer_u64()
+        t = t0
+        while t - t0 < Int64(self.spin_ns):
+            t = ptx.globaltimer_u64()
+        if tidx == 0:
+            ptx.st_u32(rec + Int64(12), Uint32(1))
+
+
+class _WaitStarted:
+    """Test-only kernel: one thread spins until ``n`` hog CTAs have started,
+    so work enqueued after it on the same stream runs while they are all
+    resident.
+    """
+
+    def __init__(self, n: int):
+        self.n = n
+
+    @cute.jit
+    def run(self, info_addr: Int64, stream):
+        self.kernel(info_addr).launch(grid=[1, 1, 1], block=[32, 1, 1], stream=stream)
+
+    @cute.kernel
+    def kernel(self, info_addr: Int64):
+        tidx, _, _ = cute.arch.thread_idx()
+        if tidx == 0:
+            i = Int32(0)
+            while i < self.n:
+                if ptx.ld_relaxed_gpu_u32(info_addr + Int64(i) * Int64(16)) != Uint32(0):
+                    i = i + 1
+
+
+def _make_hog(ctas: int, smem_bytes: int, spin_ns: int):
+    """Returns ``(launch_hog(stream), wait_started(stream), info)``; ``info``
+    is int32 ``[ctas, 4]``: started, SM id, dynamic smem bytes, done.
+    """
+    from cuda.bindings import driver as cuda
+
+    cur = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    hog_fn = cute.compile(_SmHog(ctas, smem_bytes, spin_ns).run, cutlass.Int64(0), cur)
+    wait_fn = cute.compile(_WaitStarted(ctas).run, cutlass.Int64(0), cur)
+    info = torch.zeros(ctas, 4, dtype=torch.int32, device="cuda")
+
+    def launch(stream: torch.cuda.Stream):
+        hog_fn(cutlass.Int64(info.data_ptr()), cuda.CUstream(stream.cuda_stream))
+
+    def wait_started(stream: torch.cuda.Stream):
+        wait_fn(cutlass.Int64(info.data_ptr()), cuda.CUstream(stream.cuda_stream))
+
+    return launch, wait_started, info
+
+
+class _Watchdog:
+    def __enter__(self):
+        faulthandler.dump_traceback_later(WATCHDOG_S, exit=True)
+        return self
+
+    def __exit__(self, *exc):
+        faulthandler.cancel_dump_traceback_later()
+        return False
+
+
+def _refs(layer, runner, cells):
+    """Eager routed_moe outputs without co-runners (the reference the co-run
+    outputs must equal bit for bit), classified against native.
+    """
+    refs, classes = [], {}
+    for label, x, lg in cells:
+        fin, raw = _routed_moe_both(runner, x, lg)
+        cls, _ = _classify(layer, x, lg, fin, raw, label)
+        classes[cls] = classes.get(cls, 0) + 1
+        refs.append(fin)
+    torch.cuda.synchronize()
+    return refs, classes
+
+
+def _ev():
+    return torch.cuda.Event(enable_timing=True)
+
+
+def test_corun_sm_hog(layers, runners, captures):
+    """routed_moe while a spinning kernel holds SMs (one CTA per SM, more shared
+    memory than a routed_moe CTA leaves free). A and C enqueue routed_moe only after every
+    hog CTA is resident, so routed_moe starts with part (C) or none (A) of its grid
+    resident and must finish once the hog exits; B starts the hog at
+    staggered points of routed_moe's execution; D replays a routed_moe graph against hogs.
+    Outputs must equal the undisturbed routed_moe outputs bit for bit.
+    """
+    lay = LAYERS[0]
+    layer, runner = layers[lay], runners[lay]
+    sms = torch.cuda.get_device_properties(0).multi_processor_count
+    smem = 120 * 1024
+    hog_all, wait_all, info_all = _make_hog(sms, smem, 300_000)  # every SM, 300 us
+    hog_half, wait_half, info_half = _make_hog(sms // 2, smem, 2_000_000)  # half, 2 ms
+    cells = _cells(layer, captures, lay, seed=5)
+    refs, classes = _refs(layer, runner, cells)
+    side = torch.cuda.Stream()
+    main = torch.cuda.current_stream()
+    timing = {}
+
+    def check(outs, label):
+        for i, (o, r) in enumerate(zip(outs, refs)):
+            _identical(o, r, f"{label} #{i}")
+
+    def held(hog, wait, info, label):
+        """Hog on ``side``; routed_moe batch on ``main`` once every hog CTA runs."""
+        info.zero_()
+        e0, e_hog, e_first, e_all = _ev(), _ev(), _ev(), _ev()
+        side.wait_stream(main)
+        e0.record(side)
+        hog(side)
+        e_hog.record(side)
+        main.wait_event(e0)
+        wait(main)
+        outs = []
+        for i, (_, x, lg) in enumerate(cells):
+            outs.append(runner(x, lg))
+            if i == 0:
+                e_first.record(main)
+        e_all.record(main)
+        torch.cuda.synchronize()
+        check(outs, label)
+        assert int(info[:, 0].sum()) == info.shape[0] and int(info[:, 3].sum()) == info.shape[0]
+        timing[f"{label}_hog_end_ms"] = round(e0.elapsed_time(e_hog), 3)
+        timing[f"{label}_first_routed_moe_end_ms"] = round(e0.elapsed_time(e_first), 3)
+        timing[f"{label}_all_routed_moe_end_ms"] = round(e0.elapsed_time(e_all), 3)
+        timing[f"{label}_hog_sms"] = len(set(info[:, 1].tolist()))
+        timing[f"{label}_hog_dyn_smem"] = sorted(set(info[:, 2].tolist()))
+
+    with _Watchdog():
+        e0, e1 = _ev(), _ev()
+        e0.record(main)
+        outs = [runner(x, lg) for _, x, lg in cells]
+        e1.record(main)
+        torch.cuda.synchronize()
+        check(outs, "alone")
+        timing["alone_batch_ms"] = round(e0.elapsed_time(e1), 3)
+        held(hog_all, wait_all, info_all, "A")
+        held(hog_half, wait_half, info_half, "C")
+        # B: routed_moe first, the hog launched at staggered offsets
+        for delay_cycles in (0, 2000, 10000, 30000, 60000, 120000):
+            for rep in range(4):
+                info_all.zero_()
+                side.wait_stream(main)
+                outs = [runner(x, lg) for _, x, lg in cells[:4]]
+                with torch.cuda.stream(side):
+                    torch.cuda._sleep(delay_cycles)
+                hog_all(side)
+                outs += [runner(x, lg) for _, x, lg in cells[4:]]
+                torch.cuda.synchronize()
+                check(outs, f"hog after {delay_cycles} cycles rep{rep}")
+                assert int(info_all[:, 3].sum()) == sms
+        # D: a captured graph of routed_moe launches replayed against repeated hogs
+        static_inputs = [(x.clone(), lg.clone()) for _, x, lg in cells]
+        gstream = torch.cuda.Stream()
+        gstream.wait_stream(main)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=gstream):
+            gouts = [runner(x, lg) for x, lg in static_inputs]
+        for rep in range(8):
+            side.wait_stream(main)
+            if rep % 2 == 0:
+                hog_all(side)
+                wait_all(main)
+            graph.replay()
+            if rep % 2 == 1:
+                hog_half(side)
+            torch.cuda.synchronize()
+            check(gouts, f"graph vs hog rep{rep}")
+    print(f"[routed_moe] co-run SM hog ok: timing {timing}; refs vs native {classes}")
+    # the hogs really held one SM per CTA with their full shared memory
+    assert timing["C_hog_sms"] == sms // 2 and timing["A_hog_sms"] == sms, timing
+    assert timing["A_hog_dyn_smem"] == [smem] and timing["C_hog_dyn_smem"] == [smem], timing
+
+
+def test_corun_gemm_side_stream(layers, runners, captures):
+    """CuBLAS GEMMs (shared-memory-heavy, many CTAs) on a second stream while
+    a graph of routed_moe launches replays, as the shared expert's GEMMs overlap the
+    routed experts in the served model.
+    """
+    lay = LAYERS[0]
+    layer, runner = layers[lay], runners[lay]
+    cells = _cells(layer, captures, lay, seed=9)
+    refs, classes = _refs(layer, runner, cells)
+    static_inputs = [(x.clone(), lg.clone()) for _, x, lg in cells]
+    main = torch.cuda.current_stream()
+    gstream = torch.cuda.Stream()
+    gstream.wait_stream(main)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=gstream):
+        gouts = [runner(x, lg) for x, lg in static_inputs]
+    side = torch.cuda.Stream()
+    a = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
+    small_w = torch.randn(512, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
+    with _Watchdog():
+        for rep in range(10):
+            side.wait_stream(main)
+            with torch.cuda.stream(side):
+                for _ in range(3):
+                    torch.matmul(a, b)
+                    # shared-expert-sized GEMMs as in the served overlap
+                    for x, _ in static_inputs[:4]:
+                        torch.nn.functional.linear(x, small_w)
+            graph.replay()
+            torch.cuda.synchronize()
+            for i, (o, r) in enumerate(zip(gouts, refs)):
+                _identical(o, r, f"gemm co-run rep{rep} #{i}")
+    print(f"[routed_moe] co-run GEMM ok; refs vs native {classes}")
+
+
+# ---------------------------------------------------------------------------
+# (d) many back-to-back launches in one graph, several real layers
+# ---------------------------------------------------------------------------
+def test_graph_many_launches_real_layers(layers, runners, captures):
+    """One graph, 12 launches per layer interleaved over the layers (one
+    shared workspace, every fourth launch in raw mode, an unrelated consumer
+    after each); every replayed output equals the eager routed_moe output for its
+    input bit for bit and is classified against native.
+    """
+    per_layer = 12
+    plan = []  # (layer, finalize)
+    for i in range(per_layer):
+        for lay in LAYERS:
+            plan.append((lay, (i % 4) != 3))
+    wins = {lay: fx.real_windows(captures[lay], 6) for lay in LAYERS}
+    static_inputs = [
+        (
+            torch.empty(6, fx.HIDDEN, device="cuda", dtype=torch.bfloat16),
+            torch.empty(6, fx.NUM_EXPERTS, device="cuda", dtype=torch.float32),
+        )
+        for _ in plan
+    ]
+
+    def fill(rep):
+        used = []
+        for j, ((lay, _), (xs, ls)) in enumerate(zip(plan, static_inputs)):
+            x, lg = wins[lay][(rep * 977 + j * 31) % len(wins[lay])]
+            xs.copy_(x)
+            ls.copy_(lg)
+            used.append((x, lg))
+        return used
+
+    used = fill(0)
+    main = torch.cuda.current_stream()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(main)
+    with torch.cuda.stream(stream):
+        for (lay, fin), (xs, ls) in zip(plan, static_inputs):
+            runners[lay](xs, ls, do_finalize=fin)
+    main.wait_stream(stream)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    sink = torch.zeros(6, fx.HIDDEN, device="cuda", dtype=torch.float32)
+    with torch.cuda.graph(graph, stream=stream):
+        outs = []
+        for (lay, fin), (xs, ls) in zip(plan, static_inputs):
+            o = runners[lay](xs, ls, do_finalize=fin)
+            outs.append(o)
+            # an unrelated consumer between launches, as the model's layers
+            sink.add_(o.float() if fin else o[0][:6].float())
+    classes = {}
+    emul_left = 8
+    for rep in range(6):
+        if rep:
+            used = fill(rep)
+        graph.replay()
+        torch.cuda.synchronize()
+        for j, ((lay, fin), o, (x, lg)) in enumerate(zip(plan, outs, used)):
+            tag = f"rep{rep} launch{j} layer{lay}"
+            e_fin, e_raw = _routed_moe_both(runners[lay], x, lg)
+            _identical(o, e_fin if fin else e_raw, tag)
+            cls, _ = _classify(layers[lay], x, lg, e_fin, e_raw, tag, emul=emul_left > 0)
+            if cls == "budget":
+                emul_left -= 1
+            classes[cls] = classes.get(cls, 0) + 1
+    print(
+        f"[routed_moe] graph of {len(plan)} launches over layers {LAYERS}: 6 replays bit-identical "
+        f"to eager routed_moe; vs native {classes}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Real routed-MoE inputs (native capture): correctness, not only timing
+# ---------------------------------------------------------------------------
+EMUL_PER_CASE = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_EMUL_PER_CASE", "24"))
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_real_windows_match_native(layers, runners, decoders, captures, rows):
+    """Every real window of every captured layer: exact expert IDs and
+    scales; finalized output exact, FP32-reorder-close, or (sampled, FP64
+    emulation) within the frozen budget.
+    """
+    stride = {1: 7, 6: 6, 8: 8}[rows]
+    summary = {}
+    for lay in LAYERS:
+        layer, runner = layers[lay], runners[lay]
+        wins = fx.real_windows(captures[lay], rows, stride)
+        assert wins and wins[0][0] is not None, f"layer {lay}: no captured hidden states"
+        classes = {}
+        nrms = []
+        emul_left = EMUL_PER_CASE
+        for i, (x, lg) in enumerate(wins):
+            fin, raw = _routed_moe_both(runner, x, lg)
+            n_fin = layer.native_forward(x, lg)
+            n_raw = layer.native_forward(x, lg, do_finalize=False)
+            emul_fn = None
+            if emul_left > 0:
+
+                def emul_fn(x=x, lg=lg):
+                    return fx.emulate_moe(layer, x, lg)
+
+            cls, rep = fx.classify_vs_native(
+                fin,
+                raw,
+                n_fin,
+                n_raw,
+                emul_fn,
+                f"layer {lay} window {i}",
+                lambda lg=lg: decoders[lay].compare(runner, lg),
+            )
+            if cls == "budget":
+                emul_left -= 1
+            classes[cls] = classes.get(cls, 0) + 1
+            if cls != "exact":
+                nrms.append(fx.error_stats(fin, n_fin)["nrms"])
+        gaps = torch.cat([fx.boundary_gap(lg, layer.bias) for _, lg in wins])
+        summary[lay] = {
+            "windows": len(wins),
+            "classes": classes,
+            "max_nrms_non_exact": max(nrms) if nrms else 0.0,
+            "rows_exact_boundary_tie": int((gaps == 0).sum()),
+            "rows_gap_le_1ulp": int((gaps <= 4.8e-7).sum()),
+            "min_gap": float(gaps.min()),
+        }
+    print(f"[routed_moe] real windows R={rows} rank={RANK}: {summary}")
+
+
+def test_real_routing_all_layers(layers, decoders, captures):
+    """Routing on the real logits of every captured MoE layer with that
+    layer's own router bias: the experts of every (t, k), decoded exactly
+    from the FC2 rows of a probe activation (``ExpertDecoder``), and the
+    expert scale bits equal native's. The expert weights are layer
+    ``LAYERS[0]``'s; routing depends only on logits and bias.
+    """
+    import glob
+    import re
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+
+    base = LAYERS[0]
+    layer = layers[base]
+    decoder = decoders[base]
+    all_layers = sorted(
+        {
+            int(re.search(r"_layer(\d+)\.pt$", p).group(1))
+            for p in glob.glob(os.path.join(fx.REAL_DIR, "rows*_layer*.pt"))
+        }
+    )
+    totals = {"layers": len(all_layers), "windows": 0, "min_margin": float("inf")}
+    totals.update({"exact_ties": 0, "gap_le_1ulp": 0, "tie_windows_checked": 0})
+    for lay in all_layers:
+        cap = [(None, lg.cuda()) for _, lg in fx.load_real_capture(lay)]
+        bias = fx.load_router_bias(lay)
+        runner = RoutedMoe.from_native(layer.module, bias)
+        for i, (_, lg) in enumerate(fx.real_windows(cap, 6)):
+            dec = decoder.compare(runner, lg, bias=bias)
+            assert dec["ids_equal"] and dec["scale_exact"], (lay, i, dec)
+            totals["windows"] += 1
+            totals["min_margin"] = min(totals["min_margin"], dec["margin"])
+            if bool((fx.boundary_gap(lg, bias) <= 4.8e-7).any()):
+                totals["tie_windows_checked"] += 1
+        gaps = torch.cat([fx.boundary_gap(lg, bias) for _, lg in cap])
+        totals["exact_ties"] += int((gaps == 0).sum())
+        totals["gap_le_1ulp"] += int((gaps <= 4.8e-7).sum())
+    print(f"[routed_moe] real routing over all layers (decoded experts): {totals}")
+    assert totals["min_margin"] > 4, totals
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py
new file mode 100644
index 000000000000..29835a49f106
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py
@@ -0,0 +1,135 @@
+# 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.
+"""Router-producer invariants of the target's MoE and layers (router_gemm, post_attention_hop).
+
+* one ``is_enabled("router_gemm")`` gate, at ``GlmMoE`` adoption, and the router GEMM is
+  None (the native gate) unless adoption attaches it;
+* ``GlmMoE.compute_routed_output`` has exactly one router_gemm router call and one native
+  gate call; router_gemm runs only when no logits were passed in, the gate only when neither
+  the caller (post_attention_hop) nor router_gemm produced logits, in that order;
+* the post_attention_hop's logits reach ``compute_routed_output`` only through
+  ``router_logits`` and only on the tensor-parallel path.
+"""
+
+import ast
+import inspect
+import textwrap
+
+import pytest
+
+modeling = pytest.importorskip(
+    "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers"
+)
+
+
+def _function(cls, name):
+    source = textwrap.dedent(inspect.getsource(getattr(cls, name)))
+    return ast.parse(source).body[0]
+
+
+def _self_calls(node, attr):
+    return [
+        n
+        for n in ast.walk(node)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Attribute)
+        and n.func.attr == attr
+        and isinstance(n.func.value, ast.Name)
+        and n.func.value.id == "self"
+    ]
+
+
+def _enclosing_ifs(func, target):
+    """Tests of the ``if`` statements that contain ``target`` (outermost first)."""
+    tests = []
+
+    def visit(node, stack):
+        if node is target:
+            tests.extend(stack)
+            return True
+        children = ast.iter_child_nodes(node)
+        if isinstance(node, ast.If):
+            for child in node.body:
+                if visit(child, stack + [node.test]):
+                    return True
+            for child in node.orelse:
+                if visit(child, stack):
+                    return True
+            return False
+        return any(visit(child, stack) for child in children)
+
+    visit(func, [])
+    return [ast.unparse(t) for t in tests]
+
+
+@pytest.fixture(scope="module")
+def integrated():
+    return inspect.getsource(modeling)
+
+
+def test_single_router_gemm_gate(integrated):
+    assert integrated.count('is_enabled("router_gemm")') == 1
+    adopt = inspect.getsource(modeling.GlmMoE.adopt)
+    assert 'is_enabled("router_gemm")' in adopt
+    assert modeling.GlmMoE.router_gemm is None
+
+
+def test_router_producer_order(integrated):
+    func = _function(modeling.GlmMoE, "compute_routed_output")
+    args = [a.arg for a in func.args.args]
+    assert "router_logits" in args
+    defaults = dict(zip(args[-len(func.args.defaults) :], func.args.defaults))
+    assert ast.unparse(defaults["router_logits"]) == "None"
+
+    router_gemm_calls = _self_calls(func, "router_gemm")
+    gate_calls = _self_calls(func, "gate")
+    assert len(router_gemm_calls) == 1 and len(gate_calls) == 1
+    assert router_gemm_calls[0].lineno < gate_calls[0].lineno
+
+    router_gemm_guard = " and ".join(_enclosing_ifs(func, router_gemm_calls[0]))
+    assert "router_logits is None" in router_gemm_guard
+    assert "self.router_gemm is not None" in router_gemm_guard
+    gate_guard = _enclosing_ifs(func, gate_calls[0])
+    assert gate_guard == ["router_logits is None"], gate_guard
+
+    # router_logits is assigned only by the two producers.
+    assigned = [
+        n.value
+        for n in ast.walk(func)
+        if isinstance(n, ast.Assign)
+        and any(isinstance(t, ast.Name) and t.id == "router_logits" for t in n.targets)
+    ]
+    assert sorted(ast.unparse(v) for v in assigned) == sorted(
+        [
+            ast.unparse(router_gemm_calls[0]),
+            ast.unparse(gate_calls[0]),
+        ]
+    )
+
+
+def test_attn_hop_logits_only_through_the_argument(integrated):
+    forward = inspect.getsource(modeling.GlmMoE.forward)
+    assert "assert router_logits is None or not self.use_dp" in forward
+    for cls, name in (
+        (modeling.GlmDecoderLayer, "forward_MoE"),
+        (modeling.GlmMtpLayer, "forward"),
+    ):
+        func = _function(cls, name)
+        hop_calls = _self_calls(func, "_post_attn_hop")
+        assert len(hop_calls) == 1, (cls.__name__, name)
+        guard = " and ".join(_enclosing_ifs(func, hop_calls[0]))
+        assert "PRE_MOE_FUSION" in guard and "_can_use_post_attn_hop" in guard
+        source = ast.unparse(func)
+        assert "router_logits=router_logits" in source
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py
new file mode 100644
index 000000000000..0f9203385777
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py
@@ -0,0 +1,104 @@
+# 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.
+"""Real GLM-5.2 shared-expert fixtures for the shared_expert_prefetch tests.
+
+* ``build_shared_mlp(layer, tp_rank)``: the served BF16 shared expert of one
+  MoE layer (one TP8 rank shard), built through ``GatedMLP``'s constructor
+  with the served model configuration and loaded from the checkpoint's
+  ``shared_experts.{gate,up,down}_proj`` tensors through its own loaders (the
+  rank slicing is the module's).
+* ``load_real_shared(mlp, layer)``: the same loading into an existing module
+  (the chain bench's MoE modules carry random shared-expert weights).
+* ``reference(mlp, x)``: the shared expert with the native rounding points and
+  exact (FP64) GEMM sums: ``h = bf16(x attention_output_into_hop^T)``, ``a`` = the native Triton
+  SiLU-and-mul of ``h``, ``y = bf16(a W2^T)``.
+"""
+
+import copy
+import os
+import sys
+
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
+import routed_moe_fixtures as fx  # noqa: E402
+
+TARGET_LAYERS = tuple(range(3, 78))
+
+
+def _shared_tensors(layer: int):
+    rd = fx._Reader()
+    pre = f"model.layers.{layer}.mlp.shared_experts"
+    return {
+        p: rd.full(f"{pre}.{p}.weight").to(torch.bfloat16)
+        for p in ("gate_proj", "up_proj", "down_proj")
+    }
+
+
+def load_real_shared(mlp, layer: int) -> None:
+    w = _shared_tensors(layer)
+    dev = mlp.gate_up_proj.weight.device
+    mlp.gate_up_proj.load_weights([{"weight": w["gate_proj"]}, {"weight": w["up_proj"]}])
+    mlp.down_proj.load_weights([{"weight": w["down_proj"]}])
+    fx._walk_post_load(mlp)
+    assert mlp.gate_up_proj.weight.device == dev
+
+
+def build_shared_mlp(layer: int, tp_rank: int):
+    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+
+    mc = fx._served_model_config(layer, tp_rank)
+    smc = copy.copy(mc)
+    smc.quant_config = QuantConfig()
+    mlp = GatedMLP(
+        hidden_size=fx.HIDDEN,
+        intermediate_size=fx.INTER,
+        bias=False,
+        dtype=torch.bfloat16,
+        config=smc,
+        overridden_tp_size=fx.TP_SIZE,
+        reduce_output=False,
+    ).cuda()
+    load_real_shared(mlp, layer)
+    return mlp
+
+
+def reference(mlp, x: torch.Tensor) -> torch.Tensor:
+    """FP64 GEMM sums with the native BF16 rounding points (see the module docstring)."""
+    w1 = mlp.gate_up_proj.weight.double()
+    w2 = mlp.down_proj.weight.double()
+    h = (x.double() @ w1.t()).to(torch.bfloat16)
+    a = torch.ops.trtllm.silu_and_mul(h)
+    return (a.double() @ w2.t()).to(torch.bfloat16)
+
+
+def native_in_batch(mlp, x: torch.Tensor, batch: int = 64) -> torch.Tensor:
+    """The native module on ``x`` placed in a ``batch``-row input (another cuBLAS
+    algorithm than at R rows): native algorithm-to-algorithm spread."""
+    big = torch.zeros(batch, x.shape[1], dtype=x.dtype, device=x.device)
+    big[: x.shape[0]] = x
+    return mlp(big)[: x.shape[0]].contiguous()
+
+
+def ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """|ordinal(a) - ordinal(b)| of BF16 values (sign-magnitude to ordinal)."""
+
+    def ordinal(t):
+        i = t.contiguous().view(torch.int16).to(torch.int32)
+        return torch.where(i < 0, -(i & 0x7FFF), i)
+
+    return (ordinal(a) - ordinal(b)).abs()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py
new file mode 100644
index 000000000000..0299068cdba8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py
@@ -0,0 +1,251 @@
+# 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.
+"""Host side of shared_expert_prefetch's two-kernel BF16 shared expert
+(``shared_prefetch_se_kernel.py``), an alternative.
+
+It is not part of package shared_expert_prefetch: in the two-host chain it ends later than the
+native shared expert beside the routed kernel,
+so the served path keeps the native ``GatedMLP``. The chain bench's ``se7``
+arms and ``test_shared_prefetch_shared.py`` run it from here.
+
+``shared_output(mlp, x)`` returns the shared expert's BF16 ``[R, 6144]``
+output for a served ``GatedMLP`` (the BF16 shared expert of a GLM-5.2 MoE
+layer, one TP8 rank shard), or ``None`` when the call is not eligible; the
+decision depends only on host-side state (module configuration, shapes,
+dtypes, device) and is taken before any device work, so ``None`` leaves the
+caller on the native ``GatedMLP`` path with nothing launched.
+
+The two kernels run on the current stream (the MoE's auxiliary stream in the
+served graph) with programmatic dependent launch. Numerical variant
+``se-reorder`` (see ``kernel_shared``). Side selectors (tests and
+measurements): ``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_UP_CTAS`` (up-kernel CTAs, 32 or 16),
+``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS`` (down-kernel CTAs, 48 or 24),
+``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_L2_PREFETCH`` (L2 prefetch bits, default 3),
+``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_SILU_VARIANT`` (SiLU instruction variant, tests only),
+``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS`` (1: weights staged in shared memory by bulk
+copies before the dependency wait, bit-identical outputs; default 0),
+``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_TRIGGER`` (``end`` or ``early``: where the up kernel
+releases the down kernel; default ``end``).
+"""
+
+import os
+import threading
+from typing import Dict, Optional, Tuple
+
+import torch
+
+FEATURE = "shared_expert_prefetch"
+HIDDEN = 6144
+INTER = 256
+MAX_ROWS = 8
+
+_LOCK = threading.Lock()
+_COMPILED: Dict[tuple, tuple] = {}
+
+
+def _knobs() -> tuple:
+    import shared_prefetch_se_ptx as ptx
+
+    up = int(
+        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_UP_CTAS") or 32
+    )
+    down = int(
+        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS") or 48
+    )
+    pf = int(
+        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_L2_PREFETCH") or 3
+    )
+    silu = int(
+        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_SILU_VARIANT")
+        or ptx.SILU_DEFAULT
+    )
+    stage = (
+        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS")
+        or "0"
+    ) == "1"
+    trig = (
+        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_TRIGGER") or "end"
+    )
+    return up, down, pf, silu, stage, trig
+
+
+def _compiled(key: tuple):
+    with _LOCK:
+        fns = _COMPILED.get(key)
+        if fns is not None:
+            return fns
+        import cutlass
+        import cutlass.cute as cute
+        from cuda.bindings import driver as cuda
+        from shared_prefetch_se_kernel import (
+            PF_W2_OWN,
+            SharedExpertTwoKernelDown,
+            SharedExpertTwoKernelUp,
+        )
+
+        up, down, pf, silu, stage, trig = key
+        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+        z = cutlass.Int64(0)
+        up_fn = cute.compile(
+            SharedExpertTwoKernelUp(up, prefetch=pf & 3, silu=silu, stage=stage, trigger=trig).run,
+            z,
+            z,
+            z,
+            z,
+            cutlass.Int32(1),
+            stream,
+        )
+        down_fn = cute.compile(
+            SharedExpertTwoKernelDown(down, prefetch=pf & PF_W2_OWN, stage=stage).run,
+            z,
+            z,
+            z,
+            cutlass.Int32(1),
+            stream,
+        )
+        _COMPILED[key] = (up_fn, down_fn)
+        return up_fn, down_fn
+
+
+def eligibility(mlp, x) -> Tuple[bool, str]:
+    """Host-only check that the shared_expert_prefetch kernels compute ``mlp(x)`` (rank-identical inputs)."""
+    import torch.nn.functional as F
+
+    if not isinstance(x, torch.Tensor):
+        return False, "input is not a plain tensor"
+    if x.dtype != torch.bfloat16 or x.dim() != 2 or x.shape[1] != HIDDEN:
+        return False, f"input {x.dtype} {tuple(x.shape)}"
+    rows = x.shape[0]
+    if not 1 <= rows <= MAX_ROWS:
+        return False, f"rows {rows} outside 1..{MAX_ROWS}"
+    if not x.is_cuda or not x.is_contiguous() or x.data_ptr() % 16:
+        return False, "input is not a contiguous 16-byte aligned device tensor"
+    if torch.cuda.get_device_capability(x.device)[0] != 10:
+        return False, "not an SM100-family device"
+    if type(mlp).__name__ != "GatedMLP" or getattr(mlp, "split_gate_up", False):
+        return False, "not a fused gate/up GatedMLP"
+    if mlp.activation is not F.silu:
+        return False, "activation is not SiLU"
+    if (
+        mlp.swiglu_limit is not None
+        or mlp.swiglu_alpha not in (None, 1.0)
+        or mlp.swiglu_beta not in (None, 0.0)
+    ):
+        return False, "SwiGLU limit/alpha/beta"
+    gu, dn = mlp.gate_up_proj, mlp.down_proj
+    for name, lin, shape in (("gate_up", gu, (2 * INTER, HIDDEN)), ("down", dn, (HIDDEN, INTER))):
+        w = getattr(lin, "weight", None)
+        if (
+            w is None
+            or w.dtype != torch.bfloat16
+            or tuple(w.shape) != shape
+            or not w.is_contiguous()
+        ):
+            return False, f"{name} weight is not a contiguous BF16 {shape}"
+        if w.device != x.device or w.data_ptr() % 16:
+            return False, f"{name} weight is not a 16-byte aligned tensor on the input's device"
+        if getattr(lin, "bias", None) is not None:
+            return False, f"{name} has a bias"
+        if getattr(lin, "has_any_quant", False):
+            return False, f"{name} is quantized"
+    if getattr(dn, "reduce_output", False) and getattr(dn, "tp_size", 1) > 1:
+        return False, "down projection reduces across ranks"
+    for probe in (
+        "_can_fuse_gate_up_swiglu_fp4out",
+        "_can_fuse_gate_up_swiglu",
+        "_can_fuse_swiglu_fp8_quant",
+    ):
+        fn = getattr(mlp, probe, None)
+        if fn is not None and fn():
+            return False, f"native path is {probe[5:]}"
+    return True, ""
+
+
+class SharedExpertTwoKernel:
+    """The shared_expert_prefetch shared expert of one ``GatedMLP`` (one per module: its scratch is private)."""
+
+    def __init__(self, mlp):
+        self.w1 = mlp.gate_up_proj.weight
+        self.w2 = mlp.down_proj.weight
+        self.a = torch.zeros(MAX_ROWS, INTER, dtype=torch.bfloat16, device=self.w1.device)
+
+    @staticmethod
+    def ready_for_capture() -> bool:
+        return _knobs() in _COMPILED
+
+    def __call__(self, x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+        import cutlass
+        from cuda.bindings import driver as cuda
+
+        up_fn, down_fn = _compiled(_knobs())
+        rows = x.shape[0]
+        if out is None:
+            out = torch.empty(rows, HIDDEN, dtype=torch.bfloat16, device=x.device)
+        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+        n = cutlass.Int32(rows)
+        up_fn(
+            cutlass.Int64(x.data_ptr()),
+            cutlass.Int64(self.w1.data_ptr()),
+            cutlass.Int64(self.w2.data_ptr()),
+            cutlass.Int64(self.a.data_ptr()),
+            n,
+            stream,
+        )
+        down_fn(
+            cutlass.Int64(self.a.data_ptr()),
+            cutlass.Int64(self.w2.data_ptr()),
+            cutlass.Int64(out.data_ptr()),
+            n,
+            stream,
+        )
+        return out
+
+
+def shared_eligible(mlp, x) -> bool:
+    """Whether ``shared_output(mlp, x)`` engages (host-only, same decision)."""
+    ok, _ = eligibility(mlp, x)
+    if ok and torch.cuda.is_current_stream_capturing():
+        runner = getattr(mlp, "_shared_expert_two_kernel", None)
+        ok = runner is not None and runner.ready_for_capture()
+    return ok
+
+
+def shared_output(mlp, x) -> Optional[torch.Tensor]:
+    """The shared_expert_prefetch shared expert's output for ``mlp(x)``, or ``None`` (native path; nothing launched)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    ok, reason = eligibility(mlp, x)
+    runner = getattr(mlp, "_shared_expert_two_kernel", None) if ok else None
+    if (
+        ok
+        and runner is not None
+        and (runner.w1 is not mlp.gate_up_proj.weight or runner.w2 is not mlp.down_proj.weight)
+    ):
+        runner = None
+    if (
+        ok
+        and torch.cuda.is_current_stream_capturing()
+        and (runner is None or not runner.ready_for_capture())
+    ):
+        ok, reason = False, "not warmed up before CUDA graph capture"
+    dispatch.log_capability(
+        FEATURE, f"shared_expert rows={x.shape[0] if ok else '-'} {reason or 'engaged'}", ok, reason
+    )
+    if not ok:
+        return None
+    if runner is None:
+        runner = SharedExpertTwoKernel(mlp)
+        mlp._shared_expert_two_kernel = runner
+    return runner(x)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py
new file mode 100644
index 000000000000..8aac4f631df7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py
@@ -0,0 +1,453 @@
+# 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.
+"""BF16 shared expert of a GLM-5.2 MoE layer in two small kernels (CuTe DSL, SM100 family); an
+alternative to the native shared expert, used by ``shared_prefetch_se.py``.
+
+The per-rank shared expert is ``y = down(silu(gate(x)) * up(x))`` with
+``gate_up`` ``[512, 6144]`` (rows 0..255 gate, 256..511 up) and ``down``
+``[6144, 256]``, all BF16, for ``R <= 8`` rows. Rounding points are the
+native ones: gate and up rounded to BF16 after an FP32 accumulation, the
+SiLU-and-mul in FP32 from those BF16 values (the native Triton instruction
+sequence, ``ptx.silu_mul_triton``) rounded to BF16, the down projection
+accumulated in FP32 and rounded to BF16. Only the FP32 summation order inside
+each GEMM differs from cuBLAS (numerical variant ``se-reorder``).
+
+``SharedExpertTwoKernelUp`` (``ctas`` CTAs of 8 warps): CTA ``c`` owns the
+intermediate columns ``[c * cols, (c + 1) * cols)``; its gate and up rows form
+``cols / 8`` tiles of 16 rows (8 gate rows over the 8 up rows of the same
+columns), so the ``mma.sync`` m16n8k16 accumulator of a thread holds the gate
+and the up value of one column for two tokens. The 8 warps split K = 6144
+(768 each) and their FP32 partial tiles are summed in shared memory in warp
+order 0..7. It writes ``a`` (BF16 ``[R, 256]``, a per-module scratch buffer).
+
+``SharedExpertTwoKernelDown`` (``ctas`` CTAs of 8 warps): CTA ``c`` owns output
+features ``[c * F, (c + 1) * F)``, ``F = 6144 / ctas``; each warp computes
+16-feature tiles over the full K = 256 from ``a`` and writes ``y`` (BF16
+``[R, 6144]``).
+
+Both kernels read every weight with 16-byte non-coherent loads straight into
+``mma.sync`` fragments: thread ``(g, t)`` loads 8 consecutive K elements of
+its rows per 32-wide K chunk and the activation row ``g`` at the same 8 K
+positions, so A and B use the same permutation of K inside the chunk.
+Activation columns of tokens ``>= R`` read row ``R - 1`` and their outputs are
+never stored (MMA columns are independent).
+
+With ``stage`` the kernels read the same weight bytes from shared memory
+instead: before ``griddepcontrol.wait`` each warp bulk-copies
+(``cp.async.bulk``, completion on its own mbarrier) the weight rows it reads
+into a padded shared-memory image (row pitch + 16 bytes, so that the 16-byte
+fragment loads of a warp are bank-conflict free), waits for its mbarrier after
+the dependency wait, and runs the same ``mma.sync`` sequence on ``ld.shared``
+fragments. The arithmetic and its order are unchanged, so the outputs are
+bit-identical to the unstaged kernels. The up kernel stages its 16 rows
+(192 KB, ``ctas = 32`` only); the down kernel its ``F`` rows (64 KB at 48
+CTAs, 128 KB at 24). With ``trigger = "early"`` the up kernel releases its
+dependent (the down kernel) right after its dependency wait, so that the down
+kernel can stage its rows while the up kernel computes.
+
+Memory order and progress (programmatic dependent launch as served):
+
+* Before ``griddepcontrol.wait`` the up kernel reads nothing and writes
+  nothing; with ``prefetch`` it issues ``cp.async.bulk.prefetch.L2`` of its
+  own gate/up rows (bit 1) and of one 1/ctas slice of ``down`` (bit 2), all
+  immutable weights. The down kernel likewise prefetches its own ``down``
+  rows (bit 4).
+* After the wait, the up kernel reads ``x`` (written by the pre-MoE hop,
+  plain coherent loads) and writes ``a``; the down kernel reads ``a`` (written
+  by the up kernel, complete when its wait returns) and writes ``y``. ``a``'s
+  previous reader is the down kernel of the module's previous call, complete
+  before this call's predecessors ran; ``y`` is a fresh output tensor.
+* With ``stage``, the pre-wait bulk copies read only immutable weights and
+  write only the CTA's own shared memory; ``griddepcontrol.wait`` still
+  precedes every read of ``x`` or ``a``. The early trigger changes only when
+  the down kernel is scheduled: its own wait returns after the up kernel has
+  completed and its writes of ``a`` are visible.
+* No CTA waits for another CTA: the only cross-CTA dependency is the kernel
+  boundary between the two launches, so progress does not depend on
+  co-residency.
+"""
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import shared_prefetch_se_ptx as ptx
+from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+    ptx as wptx,
+)
+
+HIDDEN = 6144
+INTER = 256  # per-rank shared intermediate size (2048 / TP8)
+NTOK = 8
+THREADS = 256
+WARPS = THREADS // 32
+W1_ROW_BYTES = HIDDEN * 2
+W2_ROW_BYTES = INTER * 2
+PF_W1 = 1
+PF_W2_SLICE = 2
+PF_W2_OWN = 4
+PAD = 16  # shared-memory row padding of the staged weight images
+ROWP1 = W1_ROW_BYTES + PAD
+ROWP2 = W2_ROW_BYTES + PAD
+SLICE1 = HIDDEN // WARPS * 2  # bytes of one gate/up row in one warp's K range
+
+
+class SharedExpertTwoKernelUp:
+    """gate/up GEMM + SiLU-and-mul (see the module docstring)."""
+
+    def __init__(
+        self,
+        ctas: int = 32,
+        prefetch: int = PF_W1 | PF_W2_SLICE,
+        unroll: int = 0,
+        silu: int = ptx.SILU_DEFAULT,
+        carveout: int = 100,
+        stage: bool = False,
+        trigger: str = "end",
+    ):
+        assert ctas in (16, 32), ctas
+        assert not stage or ctas == 32, "the staged up kernel holds 16 rows per CTA"
+        assert trigger in ("end", "early"), trigger
+        self.stage = stage
+        self.trigger = trigger
+        self.ctas = ctas
+        self.cols = INTER // ctas
+        self.tiles = self.cols // 8
+        self.chunks = HIDDEN // (WARPS * 32)
+        unroll = unroll or 8 // self.tiles
+        assert self.chunks % unroll == 0
+        self.unroll = unroll
+        self.prefetch = prefetch
+        self.silu = silu
+        self.carveout = carveout
+
+    @cute.jit
+    def run(
+        self,
+        x_addr: Int64,
+        w1_addr: Int64,
+        w2_addr: Int64,
+        a_addr: Int64,
+        num_tokens: Int32,
+        stream,
+    ):
+        @cute.struct
+        class SharedStorage:
+            part: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Float32, WARPS * self.tiles * 128], 16
+            ]
+
+        @cute.struct
+        class StagedStorage:
+            part: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Float32, WARPS * self.tiles * 128], 16
+            ]
+            mbar: cute.struct.MemRange[cutlass.Int64, WARPS]
+            wsm: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Int64, 2 * self.cols * ROWP1 // 8], 128
+            ]
+
+        if cutlass.const_expr(self.stage):
+            self.shared_storage = StagedStorage
+        else:
+            self.shared_storage = SharedStorage
+        self.kernel(x_addr, w1_addr, w2_addr, a_addr, num_tokens).launch(
+            grid=[self.ctas, 1, 1],
+            block=[THREADS, 1, 1],
+            smem=self.shared_storage.size_in_bytes(),
+            stream=stream,
+            use_pdl=True,
+            preferred_smem_carveout=self.carveout,
+        )
+
+    @cute.kernel
+    def kernel(
+        self, x_addr: Int64, w1_addr: Int64, w2_addr: Int64, a_addr: Int64, num_tokens: Int32
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        warp = Int32(tidx) // 32
+        lane = Int32(tidx) % 32
+        g = lane // 4
+        t = lane % 4
+        c = Int32(bidx)
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        part = Int32(storage.part.data_ptr().toint())
+        if cutlass.const_expr(self.stage):
+            # Pre-wait: each warp bulk-copies its K range of the CTA's 2 * cols weight rows (immutable) into the
+            # padded image; row r < cols is gate row c * cols + r, row cols + r is up row INTER + c * cols + r.
+            mbar = storage.mbar.data_ptr()
+            wsm = Int32(storage.wsm.data_ptr().toint())
+            if tidx == 0:
+                for w in cutlass.range_constexpr(WARPS):
+                    cute.arch.mbarrier_init(mbar + w, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+            if lane == 0:
+                bar = Int32((mbar + warp).toint())
+                wptx.mbar_arrive_expect_tx(bar, Int32(2 * self.cols * SLICE1))
+                for r in cutlass.range_constexpr(2 * self.cols):
+                    src_row = c * self.cols + r
+                    if cutlass.const_expr(r >= self.cols):
+                        src_row = INTER + c * self.cols + (r - self.cols)
+                    wptx.bulk_g2s(
+                        wsm + r * ROWP1 + warp * SLICE1,
+                        w1_addr + Int64(src_row) * W1_ROW_BYTES + Int64(warp * SLICE1),
+                        Int32(SLICE1),
+                        bar,
+                    )
+
+        if cutlass.const_expr(self.prefetch & (PF_W1 | PF_W2_SLICE)):
+            if tidx == 0:
+                if cutlass.const_expr(self.prefetch & PF_W1):
+                    wptx.prefetch_l2_bulk(
+                        w1_addr + Int64(c * self.cols) * W1_ROW_BYTES,
+                        Int32(self.cols * W1_ROW_BYTES),
+                    )
+                    wptx.prefetch_l2_bulk(
+                        w1_addr + Int64(INTER + c * self.cols) * W1_ROW_BYTES,
+                        Int32(self.cols * W1_ROW_BYTES),
+                    )
+                if cutlass.const_expr(self.prefetch & PF_W2_SLICE):
+                    rows2 = HIDDEN // self.ctas
+                    wptx.prefetch_l2_bulk(
+                        w2_addr + Int64(c * rows2) * W2_ROW_BYTES, Int32(rows2 * W2_ROW_BYTES)
+                    )
+        cute.arch.griddepcontrol_wait()
+        if cutlass.const_expr(self.trigger == "early"):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        xrow = g
+        if xrow >= num_tokens:
+            xrow = num_tokens - 1
+        xbase = x_addr + Int64(xrow * HIDDEN + warp * (HIDDEN // WARPS) + 8 * t) * 2
+        wbase = w1_addr + Int64((c * self.cols + g) * HIDDEN + warp * (HIDDEN // WARPS) + 8 * t) * 2
+        sbase = Int32(0)
+        if cutlass.const_expr(self.stage):
+            sbase = wsm + g * ROWP1 + warp * SLICE1 + 16 * t
+            cute.arch.mbarrier_wait(mbar + warp, 0)
+        acc = cute.make_rmem_tensor((self.tiles * 4,), Float32)
+        for i in cutlass.range_constexpr(self.tiles * 4):
+            acc[i] = Float32(0.0)
+        xv = cute.make_rmem_tensor((self.unroll * 4,), Uint32)
+        wv = cute.make_rmem_tensor((self.unroll * self.tiles * 8,), Uint32)
+        q0 = Int32(0)
+        while q0 < self.chunks:
+            for u in cutlass.range_constexpr(self.unroll):
+                off = Int64((q0 + u) * 64)
+                v = wptx.ld_v4_u32(xbase + off)
+                for i in cutlass.range_constexpr(4):
+                    xv[u * 4 + i] = v[i]
+                for j in cutlass.range_constexpr(self.tiles):
+                    if cutlass.const_expr(self.stage):
+                        soff = Int32((q0 + u) * 64)
+                        rg = wptx.lds_v4_u32(sbase + j * 8 * ROWP1 + soff)
+                        ru = wptx.lds_v4_u32(sbase + (self.cols + j * 8) * ROWP1 + soff)
+                    else:
+                        rg = wptx.ld_nc_v4_u32(wbase + Int64(j * 8 * W1_ROW_BYTES) + off)
+                        ru = wptx.ld_nc_v4_u32(wbase + Int64((INTER + j * 8) * W1_ROW_BYTES) + off)
+                    for i in cutlass.range_constexpr(4):
+                        wv[(u * self.tiles + j) * 8 + i] = rg[i]
+                        wv[(u * self.tiles + j) * 8 + 4 + i] = ru[i]
+            for u in cutlass.range_constexpr(self.unroll):
+                for j in cutlass.range_constexpr(self.tiles):
+                    b = (u * self.tiles + j) * 8
+                    for h in cutlass.range_constexpr(2):
+                        d = ptx.mma_bf16_16816(
+                            wv[b + 2 * h],
+                            wv[b + 4 + 2 * h],
+                            wv[b + 2 * h + 1],
+                            wv[b + 4 + 2 * h + 1],
+                            xv[u * 4 + 2 * h],
+                            xv[u * 4 + 2 * h + 1],
+                            acc[j * 4],
+                            acc[j * 4 + 1],
+                            acc[j * 4 + 2],
+                            acc[j * 4 + 3],
+                        )
+                        for i in cutlass.range_constexpr(4):
+                            acc[j * 4 + i] = d[i]
+            q0 = q0 + self.unroll
+
+        for j in cutlass.range_constexpr(self.tiles):
+            for i in cutlass.range_constexpr(4):
+                wptx.sts_u32(
+                    part + (((warp * self.tiles + j) * 32 + lane) * 4 + i) * 4,
+                    wptx.f32_bits(acc[j * 4 + i]),
+                )
+        cute.arch.sync_threads()
+        if warp < self.tiles:
+            j = warp
+            # warp order 0..7, starting from warp 0's partial (no +0.0 term)
+            s = cute.make_rmem_tensor((4,), Float32)
+            for i in cutlass.range_constexpr(4):
+                s[i] = wptx.bits_f32(wptx.lds_u32(part + ((j * 32 + lane) * 4 + i) * 4))
+            for w in cutlass.range_constexpr(WARPS - 1):
+                for i in cutlass.range_constexpr(4):
+                    s[i] = s[i] + wptx.bits_f32(
+                        wptx.lds_u32(part + ((((w + 1) * self.tiles + j) * 32 + lane) * 4 + i) * 4)
+                    )
+            col = c * self.cols + j * 8 + g
+            for h in cutlass.range_constexpr(2):
+                tok = 2 * t + h
+                gate = wptx.bf16_round(s[h])
+                up = wptx.bf16_round(s[2 + h])
+                av = ptx.silu_mul_triton(gate, up, self.silu)
+                if tok < num_tokens:
+                    ptx.st_u16(a_addr + Int64(tok * INTER + col) * 2, wptx.bf16_bits(av))
+        if cutlass.const_expr(self.trigger == "end"):
+            cute.arch.griddepcontrol_launch_dependents()
+
+
+class SharedExpertTwoKernelDown:
+    """down GEMM (see the module docstring)."""
+
+    def __init__(self, ctas: int = 48, prefetch: int = 0, carveout: int = 100, stage: bool = False):
+        assert ctas in (24, 48), ctas
+        self.ctas = ctas
+        self.feats = HIDDEN // ctas
+        self.tpw = self.feats // (WARPS * 16)
+        self.prefetch = prefetch
+        self.carveout = carveout
+        self.stage = stage
+
+    @cute.jit
+    def run(self, a_addr: Int64, w2_addr: Int64, y_addr: Int64, num_tokens: Int32, stream):
+        if cutlass.const_expr(self.stage):
+
+            @cute.struct
+            class StagedStorage:
+                mbar: cute.struct.MemRange[cutlass.Int64, WARPS]
+                wsm: cute.struct.Align[
+                    cute.struct.MemRange[cutlass.Int64, self.feats * ROWP2 // 8], 128
+                ]
+
+            self.shared_storage = StagedStorage
+            self.kernel(a_addr, w2_addr, y_addr, num_tokens).launch(
+                grid=[self.ctas, 1, 1],
+                block=[THREADS, 1, 1],
+                smem=self.shared_storage.size_in_bytes(),
+                stream=stream,
+                use_pdl=True,
+                preferred_smem_carveout=self.carveout,
+            )
+        else:
+            self.kernel(a_addr, w2_addr, y_addr, num_tokens).launch(
+                grid=[self.ctas, 1, 1],
+                block=[THREADS, 1, 1],
+                stream=stream,
+                use_pdl=True,
+                preferred_smem_carveout=self.carveout,
+            )
+
+    @cute.kernel
+    def kernel(self, a_addr: Int64, w2_addr: Int64, y_addr: Int64, num_tokens: Int32):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        warp = Int32(tidx) // 32
+        lane = Int32(tidx) % 32
+        g = lane // 4
+        t = lane % 4
+        c = Int32(bidx)
+        rows_w = 16 * self.tpw
+        if cutlass.const_expr(self.stage):
+            # Pre-wait: each warp bulk-copies its 16 * tpw down rows (immutable) into the padded image.
+            smem = utils.SmemAllocator()
+            storage = smem.allocate(self.shared_storage)
+            mbar = storage.mbar.data_ptr()
+            wsm = Int32(storage.wsm.data_ptr().toint())
+            if tidx == 0:
+                for w in cutlass.range_constexpr(WARPS):
+                    cute.arch.mbarrier_init(mbar + w, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+            if lane == 0:
+                bar = Int32((mbar + warp).toint())
+                wptx.mbar_arrive_expect_tx(bar, Int32(rows_w * W2_ROW_BYTES))
+                for r in cutlass.range_constexpr(rows_w):
+                    wptx.bulk_g2s(
+                        wsm + (warp * rows_w + r) * ROWP2,
+                        w2_addr + Int64(c * self.feats + warp * rows_w + r) * W2_ROW_BYTES,
+                        Int32(W2_ROW_BYTES),
+                        bar,
+                    )
+        if cutlass.const_expr(self.prefetch & PF_W2_OWN):
+            if tidx == 0:
+                wptx.prefetch_l2_bulk(
+                    w2_addr + Int64(c * self.feats) * W2_ROW_BYTES, Int32(self.feats * W2_ROW_BYTES)
+                )
+        cute.arch.griddepcontrol_wait()
+
+        arow = g
+        if arow >= num_tokens:
+            arow = num_tokens - 1
+        abase = a_addr + Int64(arow * INTER + 8 * t) * 2
+        f0 = c * self.feats + warp * (16 * self.tpw)
+        wbase = w2_addr + Int64((f0 + g) * INTER + 8 * t) * 2
+        sbase = Int32(0)
+        if cutlass.const_expr(self.stage):
+            sbase = wsm + (warp * rows_w + g) * ROWP2 + 16 * t
+            cute.arch.mbarrier_wait(mbar + warp, 0)
+        acc = cute.make_rmem_tensor((self.tpw * 4,), Float32)
+        for i in cutlass.range_constexpr(self.tpw * 4):
+            acc[i] = Float32(0.0)
+        half = 8 // self.tpw
+        av = cute.make_rmem_tensor((half * 4,), Uint32)
+        wv = cute.make_rmem_tensor((half * self.tpw * 8,), Uint32)
+        for p in cutlass.range_constexpr(self.tpw):
+            for u in cutlass.range_constexpr(half):
+                off = Int64((p * half + u) * 64)
+                v = wptx.ld_v4_u32(abase + off)
+                for i in cutlass.range_constexpr(4):
+                    av[u * 4 + i] = v[i]
+                for j in cutlass.range_constexpr(self.tpw):
+                    if cutlass.const_expr(self.stage):
+                        soff = Int32((p * half + u) * 64)
+                        r0 = wptx.lds_v4_u32(sbase + j * 16 * ROWP2 + soff)
+                        r1 = wptx.lds_v4_u32(sbase + (j * 16 + 8) * ROWP2 + soff)
+                    else:
+                        r0 = wptx.ld_nc_v4_u32(wbase + Int64(j * 16 * W2_ROW_BYTES) + off)
+                        r1 = wptx.ld_nc_v4_u32(wbase + Int64((j * 16 + 8) * W2_ROW_BYTES) + off)
+                    for i in cutlass.range_constexpr(4):
+                        wv[(u * self.tpw + j) * 8 + i] = r0[i]
+                        wv[(u * self.tpw + j) * 8 + 4 + i] = r1[i]
+            for u in cutlass.range_constexpr(half):
+                for j in cutlass.range_constexpr(self.tpw):
+                    b = (u * self.tpw + j) * 8
+                    for h in cutlass.range_constexpr(2):
+                        d = ptx.mma_bf16_16816(
+                            wv[b + 2 * h],
+                            wv[b + 4 + 2 * h],
+                            wv[b + 2 * h + 1],
+                            wv[b + 4 + 2 * h + 1],
+                            av[u * 4 + 2 * h],
+                            av[u * 4 + 2 * h + 1],
+                            acc[j * 4],
+                            acc[j * 4 + 1],
+                            acc[j * 4 + 2],
+                            acc[j * 4 + 3],
+                        )
+                        for i in cutlass.range_constexpr(4):
+                            acc[j * 4 + i] = d[i]
+        for j in cutlass.range_constexpr(self.tpw):
+            for h in cutlass.range_constexpr(2):
+                tok = 2 * t + h
+                if tok < num_tokens:
+                    row = y_addr + Int64(tok * HIDDEN + f0 + j * 16 + g) * 2
+                    ptx.st_u16(row, wptx.bf16_bits(acc[j * 4 + h]))
+                    ptx.st_u16(row + 16, wptx.bf16_bits(acc[j * 4 + 2 + h]))
+        cute.arch.griddepcontrol_launch_dependents()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py
new file mode 100644
index 000000000000..78200b103e7e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py
@@ -0,0 +1,120 @@
+# 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.
+"""Inline-PTX primitives of the shared_expert_prefetch shared-expert kernels
+(``shared_prefetch_se_kernel.py``), in addition to routed_moe's."""
+
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import Float32, Int64, T, Uint32, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.ptx import (  # noqa: E501
+    _asm,
+)
+
+
+@dsl_user_op
+def mma_bf16_16816(
+    a0: Uint32,
+    a1: Uint32,
+    a2: Uint32,
+    a3: Uint32,
+    b0: Uint32,
+    b1: Uint32,
+    c0: Float32,
+    c1: Float32,
+    c2: Float32,
+    c3: Float32,
+    *,
+    loc=None,
+    ip=None,
+):
+    """``mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32``: D = A x B + C.
+
+    A fragment (row-major 16 x 16): ``a0`` = (row g, k 2t..2t+1), ``a1`` =
+    (row g+8, k 2t..), ``a2`` = (row g, k 2t+8..), ``a3`` = (row g+8, k
+    2t+8..); B fragment (16 x 8, column n = g): ``b0`` = (k 2t..2t+1),
+    ``b1`` = (k 2t+8..); C/D: (row g, n 2t / 2t+1), (row g+8, n 2t / 2t+1),
+    with g = lane / 4 and t = lane % 4.
+    """
+    ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in (a0, a1, a2, a3, b0, b1)]
+    ops += [Float32(v).ir_value(loc=loc, ip=ip) for v in (c0, c1, c2, c3)]
+    res = _asm(
+        llvm.StructType.get_literal([T.f32()] * 4),
+        ops,
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
+        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+        loc=loc,
+        ip=ip,
+    )
+    return [Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def st_u16(addr: Int64, value: Uint32, *, loc=None, ip=None):
+    """16-bit global store of the low half of ``value``."""
+    _asm(
+        None,
+        [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
+        "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.global.b16 [$0], h;\n}",
+        "l,r",
+        loc=loc,
+        ip=ip,
+    )
+
+
+# Candidate instruction sequences of the native Triton SiLU-and-mul for an FP32
+# gate ``$1`` and up ``$2`` (``a * sigmoid(a) * (b + 0.0)`` with ``sigmoid(a) =
+# 1 / (1 + exp(-a))``): the exponential as ``ex2.approx`` of ``-a * log2(e)``
+# (with or without flush-to-zero) and the reciprocal as an IEEE or an
+# approximate division. ``SILU_DEFAULT`` is Triton's lowering of FP32 ``exp``
+# (``ex2.approx.f32``) and FP32 ``/`` (``div.full.f32``); the shared_expert_prefetch tests check
+# every candidate against the Triton kernel over every BF16 gate value.
+_SILU_EX2 = ("ex2.approx.f32", "ex2.approx.ftz.f32")
+_SILU_DIV = ("div.rn.f32", "div.full.f32", "div.approx.f32")
+SILU_VARIANTS = tuple((e, d) for d in _SILU_DIV for e in _SILU_EX2)
+SILU_DEFAULT = 2
+
+
+def _silu_text(variant: int) -> str:
+    ex2, div = SILU_VARIANTS[variant]
+    return (
+        "{\n"
+        ".reg .f32 n, e, d, s, gs, ub;\n"
+        "mul.rn.f32 n, $1, 0fBFB8AA3B;\n"
+        f"{ex2} e, n;\n"
+        "add.rn.f32 d, e, 0f3F800000;\n"
+        f"{div} s, 0f3F800000, d;\n"
+        "mul.rn.f32 gs, $1, s;\n"
+        "add.rn.f32 ub, $2, 0f00000000;\n"
+        "mul.rn.f32 $0, gs, ub;\n"
+        "}"
+    )
+
+
+@dsl_user_op
+def silu_mul_triton(
+    g: Float32, u: Float32, variant: int = SILU_DEFAULT, *, loc=None, ip=None
+) -> Float32:
+    """``g * sigmoid(g) * (u + 0.0)`` in FP32 with one of ``SILU_VARIANTS``."""
+    return Float32(
+        _asm(
+            T.f32(),
+            [Float32(g).ir_value(loc=loc, ip=ip), Float32(u).ir_value(loc=loc, ip=ip)],
+            _silu_text(variant),
+            "=f,f,f",
+            loc=loc,
+            ip=ip,
+        )
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py
new file mode 100644
index 000000000000..35121c45a059
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py
@@ -0,0 +1,251 @@
+# 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.
+"""Control flow of the shared_expert_prefetch call site in the target's MoE.
+
+With fake collaborators (no device work), on a decode step, checks that
+
+* the plan runs shared_expert_prefetch with its closure (hop_routing's fork into
+  routed_moe_into_hop's FP32 accumulator), not without it, and records no numerical variant
+  for it;
+* deferred_input_prep's tables cover the shared_expert_prefetch directory (with hop_routing's and
+  routed_moe's) and its side selectors;
+* ``GlmMoE.compute_routed_output`` tries ``shared_expert_prefetch.routed_output_into_hop`` (with the
+  module's shared expert) before hop_routing's fork, only with the hop's routing and
+  only when shared_expert_prefetch is selected; ``None`` falls through to hop_routing's fork; the
+  shared branch stays the native shared expert.
+"""
+
+import os
+import sys
+import types
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
+    dispatch as _dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
+
+sys.path.insert(
+    0,
+    os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "..", "shared", "kernels"
+    ),
+)
+
+import plan_testlib  # noqa: E402
+
+F_FEATURES = (
+    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,"
+    "router_gemm,post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,"
+    "rope_kv_append,draft_loop,shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
+    "attention_output_tail_prefetch,attention_output_into_hop,routed_moe_into_hop,"
+    "deferred_norm,dense_mlp,deferred_input_prep,hop_routing,oproj_weight_staging,"
+    "fused_kv_indexer_projection,host_page_retention"
+)
+
+
+@pytest.fixture(autouse=True)
+def decode_step():
+    """The owned call site runs on decode steps."""
+    with _dispatch.step(_dispatch.Phase.DECODE):
+        yield
+
+
+def _selector(monkeypatch, features):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+    plan_testlib.select(dispatch, features, monkeypatch)
+    return dispatch
+
+
+@pytest.fixture
+def shared_prefetch_on(monkeypatch):
+    return _selector(monkeypatch, F_FEATURES + ",shared_expert_prefetch")
+
+
+@pytest.fixture
+def shared_prefetch_off(monkeypatch):
+    return _selector(monkeypatch, F_FEATURES)
+
+
+def test_selector(shared_prefetch_on, monkeypatch):
+    dispatch = shared_prefetch_on
+    assert dispatch.is_enabled("shared_expert_prefetch")
+    assert "shared_expert_prefetch" not in dispatch.NUMERICAL_VARIANTS
+    paths = dispatch.feature_record(dispatch.enabled_features()).split()[0]
+    assert paths.startswith("paths=") and "shared_expert_prefetch" in paths[6:].split(",")
+    closure = (
+        "attention_output_tail",
+        "post_moe_hop",
+        "routed_moe",
+        "post_attention_hop",
+        "attention_output_into_hop",
+        "routed_moe_into_hop",
+        "hop_routing",
+    )
+    assert set(dispatch.REQUIRED_FEATURES["shared_expert_prefetch"]) == set(closure)
+    _selector(monkeypatch, closure + ("shared_expert_prefetch",))
+    assert dispatch.installed() == frozenset(closure + ("shared_expert_prefetch",))
+    for missing in closure:
+        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: missing})
+        assert not plan.enabled("shared_expert_prefetch"), missing
+
+
+@pytest.fixture
+def restore_registry():
+    import importlib
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.modeling import (  # noqa: E501
+        TARGET_PATHS,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    for feature in ("hop_routing", "routed_moe", "shared_expert_prefetch"):
+        importlib.import_module(TARGET_PATHS.kernel_modules[feature])
+        dispatch.get_impl(feature, "routed_output_into_hop")
+    saved = {k: dict(v) for k, v in dispatch._REGISTRY.items()}
+    yield
+    dispatch._REGISTRY.clear()
+    dispatch._REGISTRY.update(saved)
+
+
+class _Log:
+    def __init__(self):
+        self.calls = []
+
+    def fn(self, name, result):
+        def call(*args, **kwargs):
+            self.calls.append((name, args, kwargs))
+            return result
+
+        return call
+
+    def names(self):
+        return [c[0] for c in self.calls]
+
+
+def _moe(dispatch, log, r7_result, monkeypatch):
+    from tensorrt_llm._torch.utils import EventType
+
+    for feature, result in (
+        ("shared_expert_prefetch", r7_result),
+        ("hop_routing", ["hop_routing_into"]),
+        ("routed_moe", ["routed_moe_into"]),
+    ):
+        dispatch.register(
+            feature, "routed_output_into_hop", log.fn(f"{feature}.routed_output_into_hop", result)
+        )
+
+    def sequential(fn0, fn1, *args, **kwargs):
+        log.calls.append(("parallel", (), {}))
+        return fn0(), fn1()
+
+    monkeypatch.setattr(layers, "maybe_execute_in_parallel", sequential)
+    moe = layers.GlmMoE.__new__(layers.GlmMoE)
+    torch.nn.Module.__init__(moe)
+
+    class _Shared:
+        def __call__(self, x):
+            log.calls.append(("native_shared", (x,), {}))
+            return torch.full((1,), 7.0)
+
+    moe.__dict__.update(
+        use_dp=False,
+        mapping=types.SimpleNamespace(tp_size=8),
+        router_gemm=None,
+        gate=types.SimpleNamespace(e_score_correction_bias=torch.zeros(256)),
+        experts=object(),
+        _routed_into_hop=object(),
+        shared_experts=_Shared(),
+        shared_output_scale=None,
+        shared_experts_use_fp4=False,
+        event_dict={k: None for k in EventType},
+        aux_stream=None,
+    )
+    return moe
+
+
+@pytest.mark.parametrize(
+    "routing_given, r7_result, want_routed, want_out",
+    [
+        (
+            True,
+            ["shared_prefetch_into"],
+            ["shared_expert_prefetch.routed_output_into_hop"],
+            ["shared_prefetch_into"],
+        ),
+        (
+            True,
+            None,
+            ["shared_expert_prefetch.routed_output_into_hop", "hop_routing.routed_output_into_hop"],
+            ["hop_routing_into"],
+        ),
+        (
+            False,
+            ["shared_prefetch_into"],
+            ["routed_moe.routed_output_into_hop"],
+            ["routed_moe_into"],
+        ),
+    ],
+)
+def test_routed_order(
+    shared_prefetch_on,
+    restore_registry,
+    monkeypatch,
+    routing_given,
+    r7_result,
+    want_routed,
+    want_out,
+):
+    log = _Log()
+    moe = _moe(shared_prefetch_on, log, r7_result, monkeypatch)
+    x = torch.zeros(6, 4)
+    logits = torch.zeros(6, 256)
+    routing = torch.zeros(2, 6, 8, dtype=torch.int32) if routing_given else None
+    out = moe.forward(x, do_finalize=False, router_logits=logits, hop_routing_out=routing)
+    names = log.names()
+    assert [n for n in names if n.endswith("routed_output_into_hop")] == want_routed
+    assert names.count("native_shared") == 1
+    assert float(out[0][0]) == 7.0 and out[1:] == want_out
+    for name, args, kwargs in log.calls:
+        if name == "shared_expert_prefetch.routed_output_into_hop":
+            want = (
+                moe.experts,
+                moe.gate.e_score_correction_bias,
+                x,
+                logits,
+                routing,
+                moe._routed_into_hop,
+                moe.shared_experts,
+            )
+            assert not kwargs and len(args) == len(want) and all(a is b for a, b in zip(args, want))
+
+
+def test_not_selected_never_calls_shared_prefetch(
+    shared_prefetch_off, restore_registry, monkeypatch
+):
+    log = _Log()
+    moe = _moe(shared_prefetch_off, log, ["shared_prefetch_into"], monkeypatch)
+    out = moe.forward(
+        torch.zeros(6, 4),
+        do_finalize=False,
+        router_logits=torch.zeros(6, 256),
+        hop_routing_out=torch.zeros(2, 6, 8, dtype=torch.int32),
+    )
+    assert not [n for n in log.names() if n.startswith("shared_expert_prefetch.")]
+    assert float(out[0][0]) == 7.0 and out[1:] == ["hop_routing_into"]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py
new file mode 100644
index 000000000000..2f499ef2e4b5
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py
@@ -0,0 +1,454 @@
+# 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.
+"""shared_expert_prefetch's routed fork (``SharedPrefetchExpertClusterKernel``: hop_routing's fork
+plus the pre-wait L2 prefetch of the shared expert's weights, on the shared_expert_prefetch grid) on
+one GPU, real GLM-5.2 NVFP4 weights.
+
+Given the same routing, neither the grid nor the prefetch may change any expert row:
+
+* raw FC2 rows and scales of hop_routing's fork at the shared_expert_prefetch grid (the fork's
+  grid set to 4 x shared_expert_prefetch's clusters) equal those at F's grid (24 clusters,
+  raised to ``ceil(R * 8 / 2)`` for R > 6) bit for bit, for every D bucket 8..48 at R = 6 and
+  for R = 1..8 (real windows);
+* shared_expert_prefetch's FP32 accumulator equals the exact sum of routed_moe's weighted rows
+  within the FP32 reorder bound (routed_moe_into_hop's fp32 class), for the same cases and in every
+  test below. The accumulator's atomic adds come in any order, so two runs
+  need not be bit-identical; every check is against that bound;
+* CUDA-graph replay with fresh routing and interleaved row counts (one graph
+  per R, replayed in alternating order);
+* a delayed producer: a kernel that releases its dependents first and writes
+  the routing and the activation only after a delay, so a read before the
+  dependency wait would see stale data;
+* the native shared expert co-running on a second stream (as served) changes
+  no routed result, and its own output stays bit-identical while the routed
+  kernel prefetches its weights;
+* host-only declines.
+
+Routings: constructed expert sets with exactly D distinct experts (fresh
+experts per case), routed through the native routing of logits built for
+them, and real captured windows of layer ``GLM_MOE_DSA_ROUTED_MOE_LAYER``.
+"""
+
+import contextlib
+import os
+import sys
+from unittest import mock
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_ROOT = os.path.dirname(_HERE)
+sys.path.insert(0, _HERE)
+for sub in ("routed_moe", "hop_routing"):
+    sys.path.insert(0, os.path.join(_ROOT, sub))
+import hop_routing_oracle as orc  # noqa: E402
+import routed_moe_fixtures as fx  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM100-family GPU",
+)
+
+REORDER_REL = 2.0**-20
+WIDE_CLUSTERS = 36
+LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
+F_CLUSTERS = 24
+D_CASES = (8, 12, 16, 20, 23, 24, 25, 26, 28, 30, 31, 32, 33, 35, 36, 37, 40, 44, 48)
+
+
+def _grid(runner, ctas):
+    """``runner``'s expert-cluster grid set to ``ctas`` CTAs inside the block."""
+    return mock.patch.object(runner, "_xc_ctas", lambda *args, **kwargs: ctas)
+
+
+def _clusters(n):
+    """shared_expert_prefetch's cluster count set to ``n`` inside the block (None: unchanged)."""
+    if n is None:
+        return contextlib.nullcontext()
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch import (  # noqa: E501
+        routed_moe,
+    )
+
+    return mock.patch.object(routed_moe, "DEFAULT_CLUSTERS", n)
+
+
+@pytest.fixture(scope="module")
+def layer():
+    torch.manual_seed(0)
+    return fx.build_native_layer(
+        layer=LAYER, tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
+    )
+
+
+class _SHARED_PREFETCH:
+    """shared_expert_prefetch's runner with the layer's real shared expert ``mlp`` (its weights are
+    the prefetch target)."""
+
+    def __init__(self, runner, mlp, se):
+        self.runner, self.mlp, self.se = runner, mlp, se
+
+    def __call__(self, x, routing, red_out):
+        return self.runner(x, routing, red_out, self.se)
+
+    def __getattr__(self, name):
+        return getattr(self.runner, name)
+
+
+@pytest.fixture(scope="module")
+def runners(layer):
+    import shared_prefetch_fixtures as zf
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing.routed_moe import (  # noqa: E501
+        HopRoutedMoe,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
+        RoutedMoe,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
+        SharedPrefetchRoutedMoe,
+        shared_weights,
+    )
+
+    mlp = zf.build_shared_mlp(LAYER, int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")))
+    se, reason = shared_weights(mlp)
+    assert se is not None, reason
+    return (
+        RoutedMoe.from_native(layer.module, layer.bias),
+        HopRoutedMoe(layer.module, layer.bias),
+        _SHARED_PREFETCH(SharedPrefetchRoutedMoe(layer.module, layer.bias), mlp, se),
+    )
+
+
+def _pack(logits, bias):
+    ids, w = orc.native_routing(logits, bias)
+    routing = torch.empty(2, ids.shape[0], 8, dtype=torch.int32, device="cuda")
+    routing[0].copy_(ids)
+    routing[1].copy_(w.float().view(torch.int32))
+    return routing
+
+
+def _constructed(layer, rows, d, seed):
+    sets = fx.expert_sets_with_distinct(rows, d, seed=seed)
+    perm = torch.randperm(
+        fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(1000 + seed)
+    ).tolist()
+    uniq = sorted(set(e for row in sets for e in row))
+    remap = {e: perm[j] for j, e in enumerate(uniq)}
+    sets = [[remap[e] for e in row] for row in sets]
+    lg = fx.logits_for_expert_sets(sets, layer.bias.cpu(), seed=seed).float().cuda().contiguous()
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
+    routing = _pack(lg, layer.bias)
+    assert len(set(routing[0].flatten().tolist())) == d
+    return x, lg, routing
+
+
+def _real(layer, rows, count, offset=0):
+    win = fx.real_windows(fx.load_real_capture(layer.layer), rows, stride=rows * 7)
+    out = []
+    for x, lg in win[offset : offset + count]:
+        x, lg = x.cuda().contiguous(), lg.float().cuda().contiguous()
+        out.append((x, lg, _pack(lg, layer.bias)))
+    return out
+
+
+def _cases(layer):
+    out = [("D%d" % d, 6) + _constructed(layer, 6, d, seed=d) for d in D_CASES]
+    for rows in range(1, 9):
+        out += [
+            ("real R%d #%d" % (rows, i), rows) + c
+            for i, c in enumerate(_real(layer, rows, 2, offset=rows))
+        ]
+    return out
+
+
+def _shared_prefetch_ctas(runner, rows):
+    return runner._xc_ctas(rows, False, True)
+
+
+def _f_ctas(runner, rows):
+    """F's FP32-accumulator grid for ``rows`` rows: 24 clusters, raised to cover rows x 8 experts in two rounds."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        routed_moe as routed_moe_mod,
+    )
+
+    co = routed_moe_mod._max_active_clusters(runner.device, 4)
+    return 4 * min(co, max(F_CLUSTERS, -(-rows * fx.TOP_K // 2)))
+
+
+def _acc_within_bound(routed_runner, x, lg, acc):
+    """Whether ``acc`` is the exact FP32-class sum of routed_moe's weighted raw rows for ``(x,
+    lg)``: the atomics add the top-k contributions in any order, so the check is the reorder bound
+    against the exact sum, not run-to-run equality."""
+    fc2, scale, idx = routed_runner(x, lg, do_finalize=False)
+    terms = scale.double()[:, :, None] * fc2[idx.long()].double()
+    exact = terms.sum(1)
+    bound = REORDER_REL * terms.abs().sum(1)
+    return bool(((acc.double() - exact).abs() <= bound).all())
+
+
+def test_grid(layer, runners):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
+        DEFAULT_CLUSTERS,
+    )
+
+    _, _, shared_expert_prefetch = runners
+    co = __import__(
+        "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe",
+        fromlist=["x"],
+    )
+    clusters = co._max_active_clusters(shared_expert_prefetch.device, 4)
+    print(
+        f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_GRID co-resident 4-CTA clusters {clusters}; "
+        f"shared_expert_prefetch grid R6 {_shared_prefetch_ctas(shared_expert_prefetch, 6)} CTAs",
+        flush=True,
+    )
+    assert _shared_prefetch_ctas(shared_expert_prefetch, 6) == 4 * min(clusters, DEFAULT_CLUSTERS)
+    for want in (28, 32, 36):
+        with _clusters(want):
+            assert _shared_prefetch_ctas(shared_expert_prefetch, 6) == 4 * min(clusters, want)
+            assert _shared_prefetch_ctas(shared_expert_prefetch, 8) == 4 * min(
+                clusters, max(want, 32)
+            )  # two rounds cover 8 x 8 experts
+    assert _shared_prefetch_ctas(shared_expert_prefetch, 8) == 4 * min(
+        clusters, max(DEFAULT_CLUSTERS, 32)
+    )
+    assert shared_expert_prefetch.schedule(6, False, True) == "xc"
+
+
+def test_raw_rows_grid_independent_and_fp32_class(layer, runners):
+    routed_runner, hop_routing, shared_expert_prefetch = runners
+    bad = []
+    for name, rows, x, lg, routing in _cases(layer):
+        with _grid(hop_routing, _f_ctas(shared_expert_prefetch, rows)):
+            raw_f = hop_routing(x, routing, do_finalize=False)
+        with _clusters(WIDE_CLUSTERS):
+            wide = _shared_prefetch_ctas(shared_expert_prefetch, rows)
+        with _grid(hop_routing, wide):
+            raw_7 = hop_routing(x, routing, do_finalize=False)
+        fc2, scale, idx = routed_runner(x, lg, do_finalize=False)
+        torch.cuda.synchronize()
+        same = all(torch.equal(a, b) for a, b in zip(raw_f, raw_7))
+        same_routed_moe = all(torch.equal(a, b) for a, b in zip(raw_f, (fc2, scale, idx)))
+        terms = scale.double()[:, :, None] * fc2[idx.long()].double()
+        exact = terms.sum(1)
+        bound = REORDER_REL * terms.abs().sum(1)
+        ok_acc = True
+        for clusters in (None, WIDE_CLUSTERS):
+            with _clusters(clusters):
+                for _ in range(2):
+                    acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
+                    shared_expert_prefetch(x, routing, red_out=acc)
+                    torch.cuda.synchronize()
+                    ok_acc = ok_acc and bool(((acc.double() - exact).abs() <= bound).all())
+        print(
+            f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_ROUTED {name}: raw F-grid == {wide}-CTA grid {same}, == "
+            f"routed_moe {same_routed_moe}, "
+            f"fp32 within bound (shared_expert_prefetch default / wide) {ok_acc}",
+            flush=True,
+        )
+        if not (same and same_routed_moe and ok_acc):
+            bad.append(name)
+    assert not bad, bad
+
+
+def test_graph_replay_interleaved_rows(layer, runners):
+    routed_runner, _, shared_expert_prefetch = runners
+    graphs, state = {}, {}
+    for rows in (1, 3, 6, 8):
+        cases = _real(layer, rows, 3, offset=1)
+        x = cases[0][0].clone()
+        routing = cases[0][2].clone()
+        acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
+        shared_expert_prefetch(x, routing, red_out=acc)
+        torch.cuda.synchronize()
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            shared_expert_prefetch(x, routing, red_out=acc)
+        graphs[rows] = g
+        state[rows] = (x, routing, acc, cases)
+    for step in range(2):
+        for rows in (1, 3, 6, 8) if step == 0 else (8, 6, 3, 1):
+            x, routing, acc, cases = state[rows]
+            xn, ln, rn = cases[1 + step]
+            x.copy_(xn)
+            routing.copy_(rn)
+            acc.zero_()
+            graphs[rows].replay()
+            torch.cuda.synchronize()
+            assert _acc_within_bound(routed_runner, xn, ln, acc), (rows, step)
+    del graphs
+
+
+def _delay_kernel():
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass.cutlass_dsl import Int32, Int64
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        ptx as wptx,
+    )
+
+    class DelayedCopy:
+        """Releases its dependents at entry, spins ``ns``, then copies ``words`` 16-byte words."""
+
+        @cute.jit
+        def run(self, src: Int64, dst: Int64, words: Int32, ns: Int64, stream):
+            self.kernel(src, dst, words, ns).launch(
+                grid=[1, 1, 1], block=[256, 1, 1], stream=stream, use_pdl=True
+            )
+
+        @cute.kernel
+        def kernel(self, src: Int64, dst: Int64, words: Int32, ns: Int64):
+            tidx, _, _ = cute.arch.thread_idx()
+            cute.arch.griddepcontrol_launch_dependents()
+            t0 = wptx.globaltimer_u64()
+            tn = t0
+            while tn - t0 < ns:
+                tn = wptx.globaltimer_u64()
+            i = Int32(tidx)
+            while i < words:
+                v = wptx.ld_v4_u32(src + Int64(i) * 16)
+                wptx.st_v4_u32(dst + Int64(i) * 16, v[0], v[1], v[2], v[3])
+                i = i + 256
+
+    from cuda.bindings import driver as cuda
+
+    z = cutlass.Int64(0)
+    fn = cute.compile(
+        DelayedCopy().run,
+        z,
+        z,
+        cutlass.Int32(1),
+        z,
+        cuda.CUstream(torch.cuda.current_stream().cuda_stream),
+    )
+
+    def launch(src, dst, ns):
+        assert src.nbytes % 16 == 0 and src.nbytes == dst.nbytes
+        fn(
+            cutlass.Int64(src.data_ptr()),
+            cutlass.Int64(dst.data_ptr()),
+            cutlass.Int32(src.nbytes // 16),
+            cutlass.Int64(ns),
+            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
+        )
+
+    return launch
+
+
+def test_delayed_producer(layer, runners):
+    """Routing and activation written 20 us after the producer releases the routed kernel."""
+    routed_runner, _, shared_expert_prefetch = runners
+    launch = _delay_kernel()
+    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
+    for name, rows, x, lg, routing in _cases(layer)[::3]:
+        # packed source: [routing | x] so that one delayed copy writes both
+        src = torch.cat([routing.flatten().view(torch.uint8), x.flatten().view(torch.uint8)])
+        pad = (-src.numel()) % 16
+        src = torch.cat([src, torch.zeros(pad, dtype=torch.uint8, device="cuda")])
+        # stale bytes (huge activations, invalid expert ids), so that a read before the dependency wait changes
+        # the result
+        dst = torch.full_like(src, 0x7F)
+        r_view = dst[: routing.nbytes].view(torch.int32).view(2, rows, 8)
+        x_view = (
+            dst[routing.nbytes : routing.nbytes + x.nbytes]
+            .view(torch.bfloat16)
+            .view(rows, fx.HIDDEN)
+        )
+        acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
+        launch(src, dst, 20000)
+        shared_expert_prefetch(x_view, r_view, red_out=acc)
+        torch.cuda.synchronize()
+        assert _acc_within_bound(routed_runner, x, lg, acc), name
+
+
+def test_corun_shared_expert(layer, runners):
+    """shared_expert_prefetch's routed fork and the native shared expert (whose weights it
+    prefetches) on two streams, as served: same results as alone."""
+    routed_runner, _, shared_expert_prefetch = runners
+    mlp = shared_expert_prefetch.mlp
+    side = torch.cuda.Stream()
+    for name, rows, x, lg, routing in _cases(layer)[::2]:
+        want_acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
+        shared_expert_prefetch(x, routing, red_out=want_acc)
+        want_se = mlp(x)
+        torch.cuda.synchronize()
+        acc = torch.zeros_like(want_acc)
+        side.wait_stream(torch.cuda.current_stream())
+        shared_expert_prefetch(x, routing, red_out=acc)
+        with torch.cuda.stream(side):
+            se = mlp(x)
+        torch.cuda.current_stream().wait_stream(side)
+        torch.cuda.synchronize()
+        assert _acc_within_bound(routed_runner, x, lg, acc), name
+        assert torch.equal(se.view(torch.int16), want_se.view(torch.int16)), name
+
+
+def test_declines_before_device_work(layer):
+    import shared_prefetch_fixtures as zf
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch import (  # noqa: E501
+        routed_moe as shared_prefetch_mod,
+    )
+
+    class _Hop:
+        _routed_enabled = True
+
+        def __init__(self, mode):
+            self.routed_mode = mode
+
+        def can_run(self, rows):
+            return rows <= 8
+
+    mlp = zf.build_shared_mlp(LAYER, 0)
+    x, lg, routing = _constructed(layer, 6, 30, seed=5)
+    torch.cuda.synchronize()
+    cases = (
+        (_Hop("identity"), routing, mlp),
+        (_Hop("row"), routing, mlp),
+        (_Hop("fp32"), None, mlp),
+        (_Hop("fp32"), routing[:, :5].contiguous(), mlp),
+        (None, routing, mlp),
+        (_Hop("fp32"), routing, None),
+        (_Hop("fp32"), routing, torch.nn.Linear(4, 4)),
+    )
+    for hop, r, shared in cases:
+        before = torch.cuda.memory_allocated()
+        assert (
+            shared_prefetch_mod.routed_output_into_hop(
+                layer.module, layer.bias, x, lg, r, hop, shared
+            )
+            is None
+        )
+        assert torch.cuda.memory_allocated() == before
+
+
+@pytest.mark.parametrize("live", [1, 5, 6])
+def test_padded_rows_do_not_change_live_rows(layer, runners, live):
+    """An 8-row call whose rows past ``live`` are padding: the live rows' FP32 accumulator equals the live-row call's
+    within the reorder bound."""
+    routed_runner, _, shared_expert_prefetch = runners
+    ((x, lg, routing),) = _real(layer, 8, 1, offset=4)
+    full = torch.zeros(8, fx.HIDDEN, dtype=torch.float32, device="cuda")
+    shared_expert_prefetch(x, routing, red_out=full)
+    xl = x[:live].contiguous()
+    rl = _pack(lg[:live].contiguous(), layer.bias)
+    part = torch.zeros(live, fx.HIDDEN, dtype=torch.float32, device="cuda")
+    shared_expert_prefetch(xl, rl, red_out=part)
+    torch.cuda.synchronize()
+    assert _acc_within_bound(routed_runner, xl, lg[:live].contiguous(), part)
+    assert _acc_within_bound(routed_runner, xl, lg[:live].contiguous(), full[:live])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py
new file mode 100644
index 000000000000..315c577296e8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py
@@ -0,0 +1,269 @@
+# 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.
+"""Single-GPU tests of shared_expert_prefetch's two-kernel BF16 shared expert
+(``shared_prefetch_se.py``, an alternative to the native shared expert) on real GLM-5.2 weights.
+
+* ``test_silu_variants``: every candidate SiLU-and-mul instruction sequence
+  (``shared_prefetch_se_ptx.SILU_VARIANTS``) against the native Triton ``silu_and_mul`` over
+  all 65,536 BF16 gate values times 24 up values; the default must be
+  bit-identical.
+* ``test_shared_expert``: the shared_expert_prefetch kernels against the native ``GatedMLP`` and
+  the FP64 reference with native rounding points (``shared_prefetch_fixtures.reference``)
+  on real hidden states (captured windows of layers 10, 30, 50, 77, all
+  rows R = 1..8) for the layers in ``GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_LAYERS`` (default
+  3,10,30,50,77; ``all`` = every target layer) and the TP ranks in
+  ``GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_RANKS`` (default 0). Prints the error statistics per layer
+  (``GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SE_STATS``) and checks the frozen error budget of
+  ``se-reorder``.
+* ``test_eligibility``: declines (host-only, nothing launched).
+"""
+
+import json
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
+import routed_moe_fixtures as fx  # noqa: E402
+import shared_prefetch_fixtures as zf  # noqa: E402
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+
+def _silu_probe(variant: int):
+    import cutlass
+    import cutlass.cute as cute
+    import shared_prefetch_se_ptx as ptx
+    from cutlass.cutlass_dsl import Int32, Int64
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
+        ptx as wptx,
+    )
+
+    class SiluProbe:
+        @cute.jit
+        def run(self, g_addr: Int64, u_addr: Int64, o_addr: Int64, n: Int32, stream):
+            self.kernel(g_addr, u_addr, o_addr, n).launch(
+                grid=[(n + 255) // 256, 1, 1], block=[256, 1, 1], stream=stream
+            )
+
+        @cute.kernel
+        def kernel(self, g_addr: Int64, u_addr: Int64, o_addr: Int64, n: Int32):
+            tidx, _, _ = cute.arch.thread_idx()
+            bidx, _, _ = cute.arch.block_idx()
+            i = Int32(bidx) * 256 + Int32(tidx)
+            if i < n:
+                gv = wptx.bf16x2_to_f32(wptx.ld_nc_u16(g_addr + Int64(i) * 2), False)
+                uv = wptx.bf16x2_to_f32(wptx.ld_nc_u16(u_addr + Int64(i) * 2), False)
+                r = ptx.silu_mul_triton(gv, uv, variant)
+                ptx.st_u16(o_addr + Int64(i) * 2, wptx.bf16_bits(r))
+
+    from cuda.bindings import driver as cuda
+
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    z = cutlass.Int64(0)
+    fn = cute.compile(SiluProbe().run, z, z, z, cutlass.Int32(1), stream)
+
+    def call(g, u):
+        out = torch.empty_like(g)
+        fn(
+            cutlass.Int64(g.data_ptr()),
+            cutlass.Int64(u.data_ptr()),
+            cutlass.Int64(out.data_ptr()),
+            cutlass.Int32(g.numel()),
+            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
+        )
+        return out
+
+    return call
+
+
+def test_silu_variants():
+    import shared_prefetch_se_ptx as ptx
+
+    import tensorrt_llm  # noqa: F401  (registers trtllm ops)
+
+    gates = torch.arange(-32768, 32768, dtype=torch.int32).to(torch.int16).view(torch.bfloat16)
+    gen = torch.Generator().manual_seed(7)
+    ups = [1.0, -1.0, 0.5, 2.0, 3.140625, -0.0078125, 1e-3, 1e-20, 300.0, -6e4] + list(
+        (torch.randn(14, generator=gen) * 4).tolist()
+    )
+    ups = torch.tensor(ups, dtype=torch.float32).to(torch.bfloat16)
+    rows = ups.numel()
+    d = gates.numel()
+    x = torch.empty(rows, 2 * d, dtype=torch.bfloat16)
+    x[:, :d] = gates[None, :]
+    x[:, d:] = ups[:, None]
+    x = x.cuda()
+    native = torch.ops.trtllm.silu_and_mul(x)
+    g = x[:, :d].contiguous().flatten()
+    u = x[:, d:].contiguous().flatten()
+    nat = native.flatten()
+    nat_nan = torch.isnan(nat)
+    results = {}
+    for v in range(len(ptx.SILU_VARIANTS)):
+        out = _silu_probe(v)(g, u)
+        same = (out.view(torch.int16) == nat.view(torch.int16)) | (torch.isnan(out) & nat_nan)
+        results[v] = int((~same).sum().item())
+        print(
+            f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SILU variant {v} {ptx.SILU_VARIANTS[v]}: {results[v]} of "
+            f"{nat.numel()} differ",
+            flush=True,
+        )
+    assert results[ptx.SILU_DEFAULT] == 0, results
+
+
+def _windows():
+    out = []
+    for layer in (10, 30, 50, 77):
+        cap = fx.load_real_capture(layer)
+        for rows in range(1, 9):
+            for x, _ in fx.real_windows(cap, rows, stride=97)[:12]:
+                out.append(x)
+    return out
+
+
+def _layers():
+    raw = os.environ.get("GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_LAYERS", "3,10,30,50,77")
+    return list(zf.TARGET_LAYERS) if raw == "all" else [int(v) for v in raw.split(",")]
+
+
+def _ranks():
+    return [
+        int(v) for v in os.environ.get("GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_RANKS", "0").split(",")
+    ]
+
+
+def _stats(y, ref):
+    d = y.float() - ref.float()
+    nrms = (d.pow(2).mean().sqrt() / ref.float().pow(2).mean().sqrt()).item()
+    ulp = zf.ulp_distance(y, ref)
+    return {
+        "nrms": nrms,
+        "max_ulp": int(ulp.max().item()),
+        "frac_diff": (ulp > 0).float().mean().item(),
+        "max_abs": d.abs().max().item(),
+    }
+
+
+def test_shared_expert():
+    from shared_prefetch_se import shared_output
+
+    wins = [w.cuda() for w in _windows()]
+    failures = []
+    for rank in _ranks():
+        for layer in _layers():
+            mlp = zf.build_shared_mlp(layer, rank)
+            acc = {k: ([], []) for k in ("shared_expert_prefetch", "native", "native64")}
+            repeat_ok = True
+            for x in wins:
+                ref = zf.reference(mlp, x)
+                y7 = shared_output(mlp, x)
+                assert y7 is not None
+                y7b = shared_output(mlp, x)
+                repeat_ok = repeat_ok and torch.equal(y7.view(torch.int16), y7b.view(torch.int16))
+                for k, y in (
+                    ("shared_expert_prefetch", y7),
+                    ("native", mlp(x)),
+                    ("native64", zf.native_in_batch(mlp, x)),
+                ):
+                    acc[k][0].append(y.flatten())
+                    acc[k][1].append(ref.flatten())
+            st = {k: _stats(torch.cat(v[0]), torch.cat(v[1])) for k, v in acc.items()}
+            shared_prefetch_vs_native = _stats(
+                torch.cat(acc["shared_expert_prefetch"][0]), torch.cat(acc["native"][0])
+            )
+            nat_spread = _stats(torch.cat(acc["native64"][0]), torch.cat(acc["native"][0]))
+            row = {
+                "layer": layer,
+                "rank": rank,
+                "windows": len(wins),
+                "deterministic": repeat_ok,
+                "vs_ref": st,
+                "shared_prefetch_vs_native": shared_prefetch_vs_native,
+                "native64_vs_native": nat_spread,
+            }
+            print("GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SE_STATS " + json.dumps(row), flush=True)
+            # Frozen se-reorder budget: against the FP64 reference with native rounding points,
+            # shared_expert_prefetch's normalized RMS error is at most 1.1x the larger of the two
+            # native algorithms' and its largest ulp distance at most theirs + 1; the output is
+            # run-to-run deterministic.
+            nat_nrms = max(st["native"]["nrms"], st["native64"]["nrms"])
+            nat_ulp = max(st["native"]["max_ulp"], st["native64"]["max_ulp"])
+            if not (
+                repeat_ok
+                and st["shared_expert_prefetch"]["nrms"] <= 1.1 * nat_nrms
+                and st["shared_expert_prefetch"]["max_ulp"] <= nat_ulp + 1
+            ):
+                failures.append(row)
+            del mlp
+    assert not failures, failures
+
+
+@pytest.mark.parametrize("trig", ["end", "early"])
+def test_staged_bit_identical(trig, monkeypatch):
+    """Weights staged in shared memory (bulk copies before the dependency wait) give the unstaged outputs bit for
+    bit: same fragments, same ``mma.sync`` order."""
+    from shared_prefetch_se import shared_output
+
+    wins = [w.cuda() for w in _windows()]
+    bad = []
+    for layer in _layers()[:2]:
+        mlp = zf.build_shared_mlp(layer, _ranks()[0])
+        for i, x in enumerate(wins):
+            monkeypatch.delenv(
+                "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS",
+                raising=False,
+            )
+            want = shared_output(mlp, x)
+            for down in ("48", "24"):
+                monkeypatch.setenv(
+                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS", "1"
+                )
+                monkeypatch.setenv(
+                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_TRIGGER", trig
+                )
+                monkeypatch.setenv(
+                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS", down
+                )
+                got = shared_output(mlp, x)
+                monkeypatch.delenv(
+                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS"
+                )
+                torch.cuda.synchronize()
+                if not torch.equal(got.view(torch.int16), want.view(torch.int16)):
+                    bad.append((layer, i, x.shape[0], down))
+        del mlp
+    print(f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SE_STAGED trig={trig} mismatches {bad}", flush=True)
+    assert not bad
+
+
+def test_eligibility():
+    from shared_prefetch_se import eligibility
+
+    mlp = zf.build_shared_mlp(10, 0)
+    x = torch.zeros(6, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
+    assert eligibility(mlp, x) == (True, "")
+    assert not eligibility(mlp, torch.zeros(9, fx.HIDDEN, dtype=torch.bfloat16, device="cuda"))[0]
+    assert not eligibility(mlp, x.float())[0]
+    assert not eligibility(mlp, x[:, :-8])[0]
+    assert not eligibility(
+        mlp, torch.zeros(6, 2 * fx.HIDDEN, dtype=torch.bfloat16, device="cuda")[:, ::2]
+    )[0]
+    assert not eligibility(mlp, x.cpu())[0]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py
new file mode 100644
index 000000000000..7e67cf291de1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py
@@ -0,0 +1,152 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 tp8 target's class swaps side by side: the eight built-in classes the target
+adopts, what each adoption refuses, and the state the target classes may declare.
+
+Host-only: the objects are created without their constructors. What each adoption attaches is
+covered by the layer, attention and speculative-worker tests of this target.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    attention,
+    layers,
+    spec_worker,
+)
+from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.indexer import DeepseekV4Indexer
+from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
+from tensorrt_llm._torch.models.modeling_deepseekv3 import (
+    DeepseekV3DecoderLayer,
+    Deepseekv3MoE,
+    DeepseekV3MTP,
+    DeepseekV32Attention,
+)
+from tensorrt_llm._torch.models.modeling_exaone_moe import ExaoneMoeSparseMoEBlock
+from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+from tensorrt_llm._torch.speculative.mtp_dynamic_tree import MTPEagleDynamicTreeWorker
+
+pytestmark = pytest.mark.cpu_only
+
+# The target and the family code it adopts with (``shared``).
+TARGET_PACKAGE = layers.__name__.rpartition(".")[0].rpartition(".")[0]
+
+
+def _adopt_moe(moe):
+    return layers.GlmMoE.adopt(moe, None, 0)
+
+
+# (target class, the built-in class it adopts, the adoption)
+SWAPS = [
+    (attention.GlmMLA, DeepseekV32Attention, attention.GlmMLA.adopt),
+    (attention.GlmDsaHooks, DSAHooks, attention.GlmDsaHooks.adopt),
+    (attention.GlmDsaBackend, DSATrtllmAttention, attention.GlmDsaBackend.adopt),
+    (attention.GlmIndexer, Indexer, attention.GlmIndexer.adopt),
+    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, layers.GlmDecoderLayer.adopt),
+    (layers.GlmMtpLayer, DeepseekV3MTP, layers.GlmMtpLayer.adopt),
+    (layers.GlmMoE, Deepseekv3MoE, _adopt_moe),
+    (spec_worker.GlmMtpWorker, MTPEagleWorker, spec_worker.GlmMtpWorker.adopt),
+]
+SWAP_IDS = [cls.__name__ for cls, _, _ in SWAPS]
+
+# Every attribute the target's code reads that adoption may leave unset, with its default.
+CLASS_DEFAULTS = {
+    layers.GlmDecoderLayer: {
+        "post_attn_hop": None,
+        "_into_hop": None,
+        "_hop_routing": None,
+        "post_moe_hop": None,
+        "deferred_norm_hop": None,
+        "_expert_prefetch_link": None,
+    },
+    layers.GlmMtpLayer: {
+        "post_attn_hop": None,
+        "_into_hop": None,
+        "_hop_routing": None,
+        "post_moe_hop": None,
+        "deferred_norm_hop": None,
+        "_expert_prefetch_link": None,
+        "draft_front": None,
+        "draft_live_row_hop": None,
+        "_draft_live_rows": None,
+        "_draft_live_rows_compacted": False,
+    },
+    layers.GlmMoE: {"router_gemm": None, "_routed_into_hop": None},
+    attention.GlmMLA: {"_attn_tail": None, "_tail_prefetch": False},
+    attention.GlmIndexer: {"_topk_bridge": None},
+    spec_worker.GlmMtpWorker: {"_draft_loop_ops": None},
+}
+
+
+def _bare(cls):
+    obj = cls.__new__(cls)
+    if isinstance(obj, nn.Module):
+        nn.Module.__init__(obj)
+    return obj
+
+
+def _target_classes(cls):
+    return [klass for klass in cls.__mro__ if klass.__module__.startswith(TARGET_PACKAGE)]
+
+
+@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
+def test_each_target_class_extends_only_its_builtin_class(cls, builtin, adopt):
+    assert issubclass(cls, builtin)
+    added = [klass for klass in cls.__mro__ if klass not in builtin.__mro__]
+    assert added == _target_classes(cls), added
+
+
+@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
+def test_adoption_refuses_a_subclass_of_the_builtin_class(cls, builtin, adopt):
+    subclass = type(f"_Derived{builtin.__name__}", (builtin,), {})
+    obj = _bare(subclass)
+    with pytest.raises(TypeError, match=subclass.__name__):
+        adopt(obj)
+    assert type(obj) is subclass
+
+
+@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
+def test_adoption_refuses_an_adopted_object(cls, builtin, adopt):
+    obj = _bare(cls)
+    with pytest.raises(TypeError, match=cls.__name__):
+        adopt(obj)
+    assert type(obj) is cls
+
+
+@pytest.mark.parametrize(
+    "upstream, adopt",
+    [
+        (ExaoneMoeSparseMoEBlock, _adopt_moe),
+        (DeepseekV4Indexer, attention.GlmIndexer.adopt),
+        (MTPEagleDynamicTreeWorker, spec_worker.GlmMtpWorker.adopt),
+    ],
+    ids=lambda value: getattr(value, "__name__", None),
+)
+def test_adoption_refuses_the_upstream_subclasses(upstream, adopt):
+    obj = _bare(upstream)
+    with pytest.raises(TypeError, match=upstream.__name__):
+        adopt(obj)
+    assert type(obj) is upstream
+
+
+@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
+def test_the_target_classes_keep_the_builtin_layout(cls, builtin, adopt):
+    for klass in _target_classes(cls):
+        own = vars(klass)
+        assert "__init__" not in own and "__new__" not in own and "__slots__" not in own, klass
+        for name, value in own.items():
+            assert not isinstance(value, (torch.Tensor, nn.Module)), (klass, name)
+
+
+@pytest.mark.parametrize("cls", list(CLASS_DEFAULTS), ids=lambda cls: cls.__name__)
+def test_the_class_level_defaults(cls):
+    builtin = next(builtin for target, builtin, _ in SWAPS if target is cls)
+    for name, default in CLASS_DEFAULTS[cls].items():
+        assert getattr(cls, name) is default, name
+        assert not hasattr(builtin, name), name
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py
new file mode 100644
index 000000000000..a3c3df4e39bb
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py
@@ -0,0 +1,888 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 tp8 target's decoder layers, draft layer and MoE modules: their adoption, the
+phase gate of every forward override, the decode loop of ``DecodeTarget`` and the shell's
+aliases.
+
+Host-only: the built-in modules are created without their constructors and carry only the
+attributes that adoption and the phase checks read. The kernel factories are stand-ins that
+record their calls, and the built-in methods an override falls back to are recorders, so each
+test sees which call a step makes and in which order adoption builds the fast-path objects.
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib.util
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    attention,
+    dispatch,
+    layers,
+    spec_worker,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
+from tensorrt_llm._torch.models.modeling_deepseekv3 import (
+    DeepseekV3DecoderLayer,
+    DeepseekV3Gate,
+    DeepseekV3Model,
+    Deepseekv3MoE,
+    DeepseekV3MTP,
+    DeepseekV32Attention,
+    GlmMoeDsaForCausalLM,
+)
+from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+
+pytestmark = pytest.mark.cpu_only
+
+PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
+NUM_LAYERS = 3  # layer 0 dense, layers 1 and 2 MoE, layer 3 the draft layer
+HIDDEN = 16
+
+
+def _load_scope():
+    """``kernels/deferred_norm/scope.py`` without importing the deferred-norm kernels."""
+    path = Path(dispatch.__file__).parent / "kernels" / "deferred_norm" / "scope.py"
+    spec = importlib.util.spec_from_file_location(
+        f"{dispatch.__package__}.kernels.deferred_norm.scope", path
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+SCOPE = _load_scope()
+
+
+# --- built-in modules without their constructors ------------------------------------------
+
+
+def _bare(cls, **attributes):
+    module = cls.__new__(cls)
+    nn.Module.__init__(module)
+    for name, value in attributes.items():
+        setattr(module, name, value)
+    return module
+
+
+class _Gate(nn.Module):
+    """Stands in for the router gate: a weight, a bias and logits of its own."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(4, HIDDEN))
+        self.e_score_correction_bias = nn.Parameter(torch.zeros(4))
+        self.register_buffer("routing_table", torch.zeros(4))
+
+    def forward(self, hidden_states):
+        return "gate logits"
+
+
+class _Experts(nn.Module):
+    def __init__(self, calls=None) -> None:
+        super().__init__()
+        self.calls = calls if calls is not None else []
+        self.weight = nn.Parameter(torch.zeros(2, 2))
+
+    def forward(self, *args, **kwargs):
+        self.calls.append((args, kwargs))
+        return "routed"
+
+
+class _Tail:
+    """Stands in for an attention tail."""
+
+    def __init__(self, events) -> None:
+        self.events = events
+
+    def set_successor_weights(self, weights) -> None:
+        self.events.append(("set_successor_weights", tuple(weights)))
+
+
+class _SparseMlaDecodeFmha:
+    """Stands in for the sparse-MLA decode FMHA library."""
+
+
+def _attention(layer_idx):
+    """A layer's built-in attention modules: the MLA, its DSA hooks, its DSA backend and, on a
+    full-indexer layer (even index here), its indexer."""
+    backend = DSATrtllmAttention.__new__(DSATrtllmAttention)
+    backend.indexer = _bare(Indexer) if layer_idx % 2 == 0 else None
+    return _bare(
+        DeepseekV32Attention,
+        layer_idx=layer_idx,
+        sparse_attn_hooks=DSAHooks.__new__(DSAHooks),
+        mqa=backend,
+        mha=None,
+        indexer=backend.indexer,
+    )
+
+
+class _SparseParams:
+    def to_sparse_params(self, pretrained_config, layer_idx):
+        return types.SimpleNamespace(is_full_indexer_layer=layer_idx % 2 == 0)
+
+
+def _config():
+    return types.SimpleNamespace(
+        n_routed_experts=4,
+        first_k_dense_replace=1,
+        moe_layer_freq=1,
+        num_hidden_layers=NUM_LAYERS,
+        hidden_size=HIDDEN,
+        torch_dtype=torch.bfloat16,
+    )
+
+
+def _model_config(config):
+    mapping = types.SimpleNamespace(enable_attention_dp=False, tp_size=8, tp_rank=0)
+    return types.SimpleNamespace(
+        pretrained_config=config,
+        mapping=mapping,
+        moe_backend="TRTLLM",
+        sparse_attention_config=_SparseParams(),
+    )
+
+
+def _moe(experts_calls=None):
+    return _bare(
+        Deepseekv3MoE,
+        gate=_Gate(),
+        experts=_Experts(experts_calls),
+        use_dp=False,
+        mapping=types.SimpleNamespace(tp_size=8),
+        shared_experts=None,
+    )
+
+
+def _layer(layer_idx, events, model_config, cls=DeepseekV3DecoderLayer, nvfp4=True):
+    config = model_config.pretrained_config
+    moe = layer_idx >= config.first_k_dense_replace
+    return _bare(
+        cls,
+        model_config=model_config,
+        layer_idx=layer_idx,
+        mapping=model_config.mapping,
+        allreduce=object(),
+        self_attn=_attention(layer_idx),
+        mlp=_moe() if moe else nn.Linear(2, 2),
+        fusion_config=types.SimpleNamespace(PRE_MOE_FUSION=moe, POST_MOE_FUSION=moe),
+        is_nvfp4=nvfp4,
+        top_k=2,
+    )
+
+
+def _mtp(events, model_config):
+    layer = _layer(NUM_LAYERS, events, model_config, cls=DeepseekV3MTP, nvfp4=False)
+    layer.enorm = nn.LayerNorm(2)
+    layer.hnorm = nn.LayerNorm(2)
+    return layer
+
+
+def _keys(module):
+    return (
+        [name for name, _ in module.named_parameters()],
+        [name for name, _ in module.named_buffers()],
+        [name for name, _ in module.named_modules()],
+    )
+
+
+# --- stand-in kernel factories ------------------------------------------------------------
+
+
+class _Made:
+    """A fast-path object a factory made, with the call that made it."""
+
+    def __init__(self, events, name, *args, engaged=True, **kwargs):
+        self.events, self.name, self.args, self.kwargs = events, name, args, kwargs
+        self.engaged = engaged
+        events.append(("make", name))
+
+    def enable_routed_input(self) -> bool:
+        self.events.append(("enable_routed_input", self.name))
+        return True
+
+
+class _FastPaths:
+    """``dispatch.is_enabled`` / ``get_impl`` / ``log_capability`` for the adoption tests."""
+
+    def __init__(self, enabled, post_moe_engaged=True):
+        self.events = []
+        self.enabled = enabled
+        self.post_moe_engaged = post_moe_engaged
+        events = self.events
+
+        class DeferredNormHop(_Made):
+            def __init__(self, post_moe_hop, name="deferred_norm_hop"):
+                super().__init__(events, name, post_moe_hop)
+                # As the real hop: it names its post-MoE hop too, so the search for the
+                # constructing layer must skip it.
+                self.post_moe_hop = post_moe_hop
+                _, next_layer, how = SCOPE.next_layer_kind(post_moe_hop)
+                events.append(("deferred hop serves", next_layer, how))
+
+        class EvictHintsDeferredHop(DeferredNormHop):
+            def __init__(self, post_moe_hop):
+                super().__init__(post_moe_hop, name="deferred_norm_evict_hints_hop")
+
+        def draft_front(*args):
+            return None
+
+        def make_attn_tail(mla):
+            # As the real factory: the draft layer keeps the native tail.
+            if mla.layer_idx >= NUM_LAYERS:
+                return None
+            events.append(("make", "attn_tail"))
+            return _Tail(events)
+
+        def wrap_attn_tail(mla, tail):
+            if tail is not None:
+                events.append(("make", "staged_attn_tail"))
+            return tail
+
+        self.draft_front = draft_front
+        self.entries = {
+            ("attention_output_tail", "make_attn_tail"): make_attn_tail,
+            ("oproj_weight_staging", "wrap_attn_tail"): wrap_attn_tail,
+            ("sparse_mla_decode", "fmha_lib"): _SparseMlaDecodeFmha,
+            ("post_attention_hop", "make_post_attn_hop"): lambda *a, **k: _Made(
+                events, "post_attn_hop", *a, **k
+            ),
+            ("post_attention_hop", "make_into_hop_transport"): lambda hop: _Made(
+                events, "into_hop", hop
+            ),
+            ("hop_routing", "make_transport"): lambda hop: _Made(events, "hop_routing", hop),
+            ("router_gemm", "make_moe_router"): lambda **k: _Made(events, "router_gemm", **k),
+            ("post_moe_hop", "post_moe_hop"): lambda **k: _Made(
+                events, "post_moe_hop", engaged=self.post_moe_engaged, **k
+            ),
+            ("deferred_norm", "deferred_hop"): DeferredNormHop,
+            ("deferred_norm_evict_hints", "deferred_hop"): EvictHintsDeferredHop,
+            ("draft_front", "native_norm_reason"): lambda enorm, hnorm: "",
+            ("draft_front", "draft_front"): draft_front,
+            ("draft_live_rows", "make_live_row_hop"): lambda hop: _Made(
+                events, "draft_live_row_hop", hop
+            ),
+        }
+
+    def is_enabled(self, path):
+        assert path in dispatch.FAST_PATH_NAMES, path
+        return path in self.enabled
+
+    def get_impl(self, path, name):
+        assert path in self.enabled, path
+        return self.entries[(path, name)]
+
+    def log_capability(self, path, key, eligible, reason=""):
+        self.events.append(("log_capability", path, key, eligible, reason))
+        return eligible
+
+
+@pytest.fixture
+def fast_paths(monkeypatch):
+    def install(enabled=dispatch.FAST_PATH_NAMES, post_moe_engaged=True):
+        paths = _FastPaths(frozenset(enabled), post_moe_engaged)
+        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+        monkeypatch.setattr(dispatch, "log_capability", paths.log_capability)
+        return paths
+
+    return install
+
+
+def _built_in_model(events):
+    """The built-in model as ``adopt_model`` sees it: decoder layers, then the draft layer, and
+    the speculative worker."""
+    model_config = _model_config(_config())
+    core = DeepseekV3Model.__new__(DeepseekV3Model)
+    nn.Module.__init__(core)
+    core.num_hidden_layers = NUM_LAYERS
+    core.layers = nn.ModuleList(
+        [_layer(i, events, model_config) for i in range(NUM_LAYERS)] + [_mtp(events, model_config)]
+    )
+    return types.SimpleNamespace(model=core, spec_worker=_bare(MTPEagleWorker))
+
+
+# --- adoption -----------------------------------------------------------------------------
+
+
+def test_adoption_builds_each_object_in_the_builtin_constructor_order(fast_paths):
+    """Layer by layer: the attention tail, the post-attention hop and its transports, the router
+    GEMM, the post-MoE hops, the tail's successor weights; on the draft layer then its front and
+    live-row hop."""
+    paths = fast_paths()
+    model = _built_in_model(paths.events)
+
+    modeling.adopt_model(model, plan=None)
+
+    tail = [("make", "attn_tail"), ("make", "staged_attn_tail")]
+    pre_moe = tail + [("make", "post_attn_hop"), ("make", "into_hop"), ("make", "hop_routing")]
+    post_moe = [
+        ("make", "router_gemm"),
+        ("make", "post_moe_hop"),
+        ("enable_routed_input", "post_moe_hop"),
+    ]
+    layers_0 = pre_moe
+    layers_1 = (
+        pre_moe
+        + post_moe
+        + [
+            ("make", "deferred_norm_evict_hints_hop"),
+            ("deferred hop serves", 2, "model config"),
+            ("enable_routed_input", "deferred_norm_evict_hints_hop"),
+        ]
+    )
+    layers_2 = pre_moe + post_moe  # the last decoder layer: no next layer to defer into
+    tails = [e for e in paths.events if e[0] == "set_successor_weights"]
+    draft = [
+        ("make", "post_attn_hop"),
+        ("make", "router_gemm"),
+        ("log_capability", "draft_front", "draft_front", True, ""),
+        ("make", "draft_live_row_hop"),
+    ]
+    events = [e for e in paths.events if e[0] != "set_successor_weights"]
+    assert events == layers_0 + layers_1 + layers_2 + draft
+    # The tail stages its MoE layer's router weight, after that layer's post-MoE hops.
+    core = model.model
+    assert tails == [("set_successor_weights", (core.layers[i].mlp.gate.weight,)) for i in (1, 2)]
+    assert paths.events.index(tails[0]) == len(layers_0 + layers_1)
+
+
+def test_adoption_attaches_what_the_constructors_attached(fast_paths):
+    paths = fast_paths()
+    model = _built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+    dense, moe, last, draft = model.model.layers
+
+    assert [type(m) for m in (dense, moe, last)] == [layers.GlmDecoderLayer] * 3
+    assert type(draft) is layers.GlmMtpLayer
+    assert [type(m.mlp) for m in (moe, last, draft)] == [layers.GlmMoE] * 3
+    assert type(dense.mlp) is nn.Linear
+    for layer in (dense, moe, last, draft):
+        mla = layer.self_attn
+        assert type(mla) is attention.GlmMLA
+        assert type(mla.sparse_attn_hooks) is attention.GlmDsaHooks
+        assert type(mla.mqa) is attention.GlmDsaBackend
+        assert mla.mqa.prepended_fmha_lib_classes == (_SparseMlaDecodeFmha,)
+        assert mla._tail_prefetch is True
+    # Layers 0 and 2 run a full indexer, layers 1 and 3 reuse the shared top-k.
+    assert [type(m.self_attn.mqa.indexer) for m in (dense, last)] == [attention.GlmIndexer] * 2
+    assert moe.self_attn.mqa.indexer is None and draft.self_attn.mqa.indexer is None
+    assert isinstance(moe.self_attn._attn_tail, _Tail) and draft.self_attn._attn_tail is None
+
+    hop = moe.post_attn_hop
+    assert hop.args == (moe.mapping, HIDDEN, torch.bfloat16) and hop.kwargs == {
+        "native_allreduce": moe.allreduce
+    }
+    assert moe._into_hop.args == (hop,) and moe._hop_routing.args == (moe._into_hop,)
+    assert moe.mlp.router_gemm.kwargs == {
+        "gate": moe.mlp.gate,
+        "layer_idx": 1,
+        "model_config": moe.model_config,
+    }
+    assert moe.post_moe_hop.kwargs == {
+        "mapping": moe.mapping,
+        "hidden_size": HIDDEN,
+        "top_k": 2,
+        "dtype": torch.bfloat16,
+        "raw_producer": moe.mlp.experts,
+    }
+    assert moe.mlp._routed_into_hop is moe.post_moe_hop
+    assert moe.deferred_norm_hop.post_moe_hop is moe.post_moe_hop
+    assert last.deferred_norm_hop is None and last.post_moe_hop is not None
+    assert dense.post_moe_hop is None and dense.post_attn_hop is not None
+    assert draft._into_hop is None and draft.post_moe_hop is None
+    assert draft.draft_front is paths.draft_front
+    assert draft.draft_live_row_hop.args == (draft.post_attn_hop,)
+    assert draft._draft_live_rows is None and draft._draft_live_rows_compacted is False
+    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
+    # The adopted state names every object, layer by layer.
+    state = dict(modeling.adopted_state(types.SimpleNamespace(model=model.model)))
+    assert [entry[0] for entry in state["model.layers.1"]] == [
+        "post_attn_hop",
+        "_into_hop",
+        "_hop_routing",
+        "post_moe_hop",
+        "deferred_norm_hop",
+        "self_attn._attn_tail",
+        "mlp.router_gemm",
+        "mlp._routed_into_hop",
+    ]
+    assert [entry[0] for entry in state["model.layers.3"]] == [
+        "post_attn_hop",
+        "draft_front",
+        "draft_live_row_hop",
+        "mlp.router_gemm",
+    ]
+
+
+def test_the_deferred_norm_variant_without_evict_hints(fast_paths):
+    paths = fast_paths(set(dispatch.FAST_PATH_NAMES) - {"deferred_norm_evict_hints"})
+    model = _built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+    assert ("make", "deferred_norm_hop") in paths.events
+    assert ("deferred hop serves", 2, "model config") in paths.events
+
+
+def test_without_the_fast_paths_adoption_attaches_nothing(fast_paths):
+    paths = fast_paths(enabled=())
+    model = _built_in_model(paths.events)
+    before = [_keys(layer) for layer in model.model.layers]
+
+    modeling.adopt_model(model, plan=None)
+
+    assert paths.events == []
+    for layer in model.model.layers:
+        assert type(layer.self_attn) is attention.GlmMLA and layer.self_attn._attn_tail is None
+        assert layer.self_attn.mqa.prepended_fmha_lib_classes == ()
+        assert "prepended_fmha_lib_classes" not in vars(layer.self_attn.mqa)
+        for name in ("post_attn_hop", "_into_hop", "_hop_routing", "post_moe_hop"):
+            assert getattr(layer, name) is None and name not in vars(layer)
+        assert layer.deferred_norm_hop is None and layer._expert_prefetch_link is None
+    draft = model.model.layers[NUM_LAYERS]
+    assert draft.draft_front is None and draft.draft_live_row_hop is None
+    assert model.model.layers[1].mlp.router_gemm is None
+    assert model.model.layers[1].mlp._routed_into_hop is None
+    assert [_keys(layer) for layer in model.model.layers] == before
+
+
+def test_adoption_adds_no_parameter_buffer_or_module(fast_paths):
+    paths = fast_paths()
+    model = _built_in_model(paths.events)
+    before = [_keys(layer) for layer in model.model.layers]
+    modeling.adopt_model(model, plan=None)
+    assert [_keys(layer) for layer in model.model.layers] == before
+    assert "mlp.gate.weight" in before[1][0] and before[1][1] == ["mlp.gate.routing_table"]
+
+
+def test_a_post_moe_hop_that_does_not_engage_is_not_attached(fast_paths):
+    paths = fast_paths(post_moe_engaged=False)
+    model = _built_in_model(paths.events)
+    modeling.adopt_model(model, plan=None)
+    moe = model.model.layers[1]
+    assert moe.post_moe_hop is None and moe.deferred_norm_hop is None
+    assert moe.mlp._routed_into_hop is None
+    assert ("enable_routed_input", "post_moe_hop") not in paths.events
+
+
+def test_the_draft_front_needs_non_adp_tensor_parallelism(fast_paths):
+    paths = fast_paths()
+    events = paths.events
+    model_config = _model_config(_config())
+    model_config.mapping.enable_attention_dp = True
+    draft = layers.GlmMtpLayer.adopt(_mtp(events, model_config))
+    assert draft.draft_front is None
+    assert (
+        "log_capability",
+        "draft_front",
+        "draft_front",
+        False,
+        "requires non-ADP tensor parallelism",
+    ) in events
+
+
+class _OtherDecoderLayer(DeepseekV3DecoderLayer):
+    pass
+
+
+class _OtherMoE(Deepseekv3MoE):
+    pass
+
+
+class _OtherAttention(DeepseekV32Attention):
+    pass
+
+
+def _refusal(case):
+    """``(adopt, module, the class the refusal names)``."""
+    events = []
+    model_config = _model_config(_config())
+    if case == "the draft layer as a decoder layer":
+        return layers.GlmDecoderLayer.adopt, _mtp(events, model_config), "DeepseekV3MTP"
+    if case == "a decoder layer subclass":
+        module = _layer(1, events, model_config, cls=_OtherDecoderLayer)
+        return layers.GlmDecoderLayer.adopt, module, "_OtherDecoderLayer"
+    if case == "a module as a decoder layer":
+        return layers.GlmDecoderLayer.adopt, nn.Module(), "Module"
+    if case == "a decoder layer as the draft layer":
+        return layers.GlmMtpLayer.adopt, _layer(0, events, model_config), "DeepseekV3DecoderLayer"
+    if case == "a decoder layer with an attention subclass":
+        module = _layer(1, events, model_config)
+        module.self_attn.__class__ = _OtherAttention
+        return layers.GlmDecoderLayer.adopt, module, "_OtherAttention"
+    module = _layer(1, events, model_config)
+    module.mlp = _bare(_OtherMoE, gate=_Gate())
+    return layers.GlmDecoderLayer.adopt, module, "_OtherMoE"
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        "the draft layer as a decoder layer",
+        "a decoder layer subclass",
+        "a module as a decoder layer",
+        "a decoder layer as the draft layer",
+        "a decoder layer with an attention subclass",
+        "a decoder layer with a MoE subclass",
+    ],
+)
+def test_adoption_refuses_anything_but_the_exact_builtin_class(fast_paths, case):
+    fast_paths(enabled=())
+    adopt, module, got = _refusal(case)
+    with pytest.raises(TypeError, match=f"adopts the built-in .*, got {got}$"):
+        adopt(module)
+
+
+def test_the_moe_is_adopted_once(fast_paths):
+    fast_paths(enabled=())
+    model_config = _model_config(_config())
+    moe = layers.GlmMoE.adopt(_moe(), model_config, 1)
+    with pytest.raises(TypeError, match="GlmMoE adopts the built-in Deepseekv3MoE, got GlmMoE"):
+        layers.GlmMoE.adopt(moe, model_config, 1)
+
+
+def test_the_classes_declare_no_state_of_their_own():
+    for cls in (layers.GlmDecoderLayer, layers.GlmMtpLayer, layers.GlmMoE):
+        for base in cls.__mro__[: cls.__mro__.index(nn.Module)]:
+            if base.__module__ == layers.__name__:
+                assert "__slots__" not in vars(base) and "__init__" not in vars(base), base
+
+
+# --- the phase gate -----------------------------------------------------------------------
+
+# Each forward override and the built-in method it falls back to.
+OVERRIDES = [
+    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, "forward"),
+    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, "forward_MoE"),
+    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, "forward_mlp"),
+    (layers.GlmMtpLayer, DeepseekV3MTP, "forward"),
+    (layers.GlmMoE, Deepseekv3MoE, "forward"),
+    (layers.GlmMoE, Deepseekv3MoE, "compute_routed_output"),
+]
+
+
+def _gate_of(cls, name):
+    tree = ast.parse(Path(layers.__file__).read_text())
+    node = next(
+        n
+        for c in tree.body
+        if isinstance(c, ast.ClassDef) and c.name == cls.__name__
+        for n in c.body
+        if isinstance(n, ast.FunctionDef) and n.name == name
+    )
+    return node, node.body[0]
+
+
+@pytest.mark.parametrize("cls, builtin, name", OVERRIDES, ids=lambda x: getattr(x, "__name__", x))
+def test_each_override_starts_with_the_phase_gate(cls, builtin, name):
+    """``if dispatch.phase() is not DECODE: return super().<name>(<every built-in parameter by
+    name>)``; the parameters only the decode composition takes are dropped."""
+    node, gate = _gate_of(cls, name)
+    assert ast.unparse(gate.test) == "dispatch.phase() is not dispatch.Phase.DECODE"
+    assert len(gate.body) == 1 and not gate.orelse
+    call = gate.body[0].value
+    assert ast.unparse(call.func) == f"super().{name}" and not call.args
+    passed = [k.arg for k in call.keywords]
+    parameters = inspect.signature(getattr(builtin, name)).parameters
+    expected = [
+        p.name if p.kind is not p.VAR_KEYWORD else None
+        for p in parameters.values()
+        if p.name != "self"
+    ]
+    assert passed == expected
+    # The override keeps the built-in parameters, in order, and may add keywords after them.
+    own = list(inspect.signature(getattr(cls, name)).parameters)
+    assert own[: len(parameters)] == list(parameters)
+
+
+def _decode_layer():
+    model_config = _model_config(_config())
+    layer = _layer(1, [], model_config)
+    layer.__class__ = layers.GlmDecoderLayer
+    layer.mlp.__class__ = layers.GlmMoE
+    return layer
+
+
+def _calls_through(monkeypatch, owner, name, calls):
+    def recorder(self, *args, **kwargs):
+        calls.append((name, self, args, kwargs))
+        return f"{owner.__name__}.{name}"
+
+    monkeypatch.setattr(owner, name, recorder)
+
+
+LAYER_ARGS = dict(position_ids="pos", hidden_states="h", attn_metadata="md", residual="r")
+
+
+@pytest.mark.parametrize("phase", [PREFILL, None], ids=["prefill", "outside a step"])
+def test_off_decode_every_override_runs_the_builtin_method(monkeypatch, phase):
+    calls = []
+    for owner, name in [
+        (DeepseekV3DecoderLayer, "forward"),
+        (DeepseekV3DecoderLayer, "forward_MoE"),
+        (DeepseekV3DecoderLayer, "forward_mlp"),
+        (DeepseekV3MTP, "forward"),
+        (Deepseekv3MoE, "forward"),
+        (Deepseekv3MoE, "compute_routed_output"),
+    ]:
+        _calls_through(monkeypatch, owner, name, calls)
+    monkeypatch.setattr(layers.GlmDecoderLayer, "forward_decode", lambda *a, **k: pytest.fail())
+    layer = _decode_layer()
+    mtp = _bare(layers.GlmMtpLayer)
+
+    def run():
+        out = [
+            layer.forward(**LAYER_ARGS, spec_metadata="spec", lora_params="lora"),
+            layer.forward_MoE("h", "md", "r", "spec", into_hop_published=True),
+            layer.forward_mlp("h", "r"),
+            mtp.forward("ids", "pos", "h", "embed", "md", [2], "spec", lora_params="lora"),
+            layer.mlp.forward("h", None, [2], "params", False, "logits", "routing"),
+            layer.mlp.compute_routed_output("h", None, [2], True, "logits", "routing"),
+        ]
+        assert out == [
+            "DeepseekV3DecoderLayer.forward",
+            "DeepseekV3DecoderLayer.forward_MoE",
+            "DeepseekV3DecoderLayer.forward_mlp",
+            "DeepseekV3MTP.forward",
+            "Deepseekv3MoE.forward",
+            "Deepseekv3MoE.compute_routed_output",
+        ]
+
+    if phase is None:
+        run()
+    else:
+        with dispatch.step(phase):
+            run()
+    assert [(c[0], c[2], c[3]) for c in calls] == [
+        ("forward", (), {**LAYER_ARGS, "spec_metadata": "spec", "lora_params": "lora"}),
+        (
+            "forward_MoE",
+            (),
+            dict(hidden_states="h", attn_metadata="md", residual="r", spec_metadata="spec"),
+        ),
+        ("forward_mlp", (), dict(hidden_states="h", residual="r", spec_metadata=None)),
+        (
+            "forward",
+            (),
+            dict(
+                input_ids="ids",
+                position_ids="pos",
+                hidden_states="h",
+                embed_tokens="embed",
+                attn_metadata="md",
+                all_rank_num_tokens=[2],
+                spec_metadata="spec",
+                lora_params="lora",
+            ),
+        ),
+        (
+            "forward",
+            (),
+            dict(
+                hidden_states="h",
+                hidden_states_fp4=None,
+                all_rank_num_tokens=[2],
+                final_all_reduce_params="params",
+                do_finalize=False,
+            ),
+        ),
+        (
+            "compute_routed_output",
+            (),
+            dict(
+                hidden_states="h", hidden_states_fp4=None, all_rank_num_tokens=[2], do_finalize=True
+            ),
+        ),
+    ]
+
+
+def test_on_a_decode_step_the_layer_runs_its_decode_composition(monkeypatch):
+    calls = []
+    _calls_through(monkeypatch, DeepseekV3DecoderLayer, "forward", calls)
+    _calls_through(monkeypatch, layers.GlmDecoderLayer, "forward_decode", calls)
+    layer = _decode_layer()
+    with dispatch.step(DECODE):
+        assert layer.forward(**LAYER_ARGS, lora_params="lora") == "GlmDecoderLayer.forward_decode"
+    assert calls == [
+        ("forward_decode", layer, ("pos", "h", "md", "r", None), {"lora_params": "lora"})
+    ]
+
+
+@pytest.mark.parametrize("phase", [DECODE, PREFILL], ids=["decode", "prefill"])
+def test_the_router_gemm_replaces_the_gate_on_decode_steps_only(fast_paths, phase):
+    paths = fast_paths(enabled={"router_gemm"})
+    paths.entries[("router_gemm", "make_moe_router")] = lambda **k: (lambda h: "router logits")
+    experts = []
+    moe = layers.GlmMoE.adopt(_moe(experts), _model_config(_config()), 1)
+    hidden = torch.zeros(2, HIDDEN)
+
+    with dispatch.step(phase):
+        assert moe.compute_routed_output(hidden, None, [2], True) == "routed"
+
+    logits = "router logits" if phase is DECODE else "gate logits"
+    assert experts == [
+        (
+            (hidden, logits),
+            dict(
+                do_finalize=True,
+                output_dtype=hidden.dtype,
+                all_rank_num_tokens=[2],
+                use_dp_padding=False,
+            ),
+        )
+    ]
+
+
+def test_the_post_attention_hop_reads_the_builtin_gates_weight_only():
+    layer = _decode_layer()
+    assert layer._post_attn_router_weight() is None  # a stand-in gate
+    layer.mlp.gate = _bare(DeepseekV3Gate, weight=nn.Parameter(torch.zeros(4, HIDDEN)))
+    assert layer._post_attn_router_weight() is layer.mlp.gate.weight
+
+
+# --- DecodeTarget -------------------------------------------------------------------------
+
+
+class _NoContexts:
+    """Attention metadata a decode step passes; reading the context count fails the test."""
+
+    def __getattr__(self, name):
+        raise AssertionError(f"DecodeTarget read attn_metadata.{name}")
+
+
+class _RecordingLayer(nn.Module):
+    def __init__(self, index, calls) -> None:
+        super().__init__()
+        self.index, self.calls = index, calls
+
+    def forward(self, **kwargs):
+        self.calls.append((self.index, kwargs, dispatch.phase()))
+        return kwargs["hidden_states"] + 1, f"residual {self.index}"
+
+
+class _Embedding(nn.Module):
+    def __init__(self, calls) -> None:
+        super().__init__()
+        self.calls = calls
+
+    def forward(self, input_ids):
+        self.calls.append(("embed", input_ids))
+        return torch.zeros(2, HIDDEN)
+
+
+@pytest.fixture
+def decode_core(monkeypatch):
+    calls = []
+    core = DeepseekV3Model.__new__(DeepseekV3Model)
+    nn.Module.__init__(core)
+    core.embed_tokens = _Embedding(calls)
+    core.layers = nn.ModuleList([_RecordingLayer(i, calls) for i in range(NUM_LAYERS + 1)])
+    core.num_hidden_layers = NUM_LAYERS
+    core.mapping_with_cp = None
+    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(core)
+    for layer in core.layers:
+        layer.register_forward_pre_hook(
+            lambda m, a, k: calls.append(("hook", m.index)), with_kwargs=True
+        )
+    return core, calls
+
+
+@pytest.mark.parametrize("prologue_rows", [None, "rows"], ids=["no packet", "packet"])
+def test_the_decode_loop_runs_the_decoder_layers_as_modules(fast_paths, decode_core, prologue_rows):
+    core, calls = decode_core
+    paths = fast_paths(enabled={"decode_prologue", "cross_layer_expert_prefetch"})
+    rows = torch.full((2, HIDDEN), 7.0) if prologue_rows else None
+    seen = []
+
+    def take_embedding(attn_metadata, input_ids, embed_tokens):
+        seen.append((attn_metadata, input_ids, embed_tokens))
+        return rows
+
+    core.decode_prologue = types.SimpleNamespace(take_embedding=take_embedding)
+    paths.entries[("cross_layer_expert_prefetch", "join")] = lambda: calls.append(("join",))
+    md = _NoContexts()
+
+    with dispatch.step(DECODE):
+        out = core._decode.forward(
+            md, input_ids="ids", position_ids="pos", spec_metadata="spec", lora_params="x"
+        )
+
+    assert seen == [(md, "ids", core.embed_tokens)]
+    embedded = [] if prologue_rows else [("embed", "ids")]
+    expected_start = (rows if prologue_rows else torch.zeros(2, HIDDEN)) + NUM_LAYERS
+    assert torch.equal(out, expected_start)
+    layer_calls = []
+    for i in range(NUM_LAYERS):
+        kwargs = calls[len(embedded) + 2 * i + 1][1]
+        assert set(kwargs) == {
+            "position_ids",
+            "hidden_states",
+            "attn_metadata",
+            "residual",
+            "spec_metadata",
+        }
+        assert kwargs["residual"] == (None if i == 0 else f"residual {i - 1}")
+        assert kwargs["attn_metadata"] is md and kwargs["spec_metadata"] == "spec"
+        layer_calls += [("hook", i), (i, kwargs, DECODE)]
+    assert calls == embedded + layer_calls + [("join",)]
+
+
+def test_the_decode_loop_without_its_fast_paths_is_the_builtin_loop(fast_paths, decode_core):
+    core, calls = decode_core
+    fast_paths(enabled=())
+    with dispatch.step(DECODE):
+        out = core._decode.forward(_NoContexts(), input_ids="ids")
+    assert torch.equal(out, torch.zeros(2, HIDDEN) + NUM_LAYERS)
+    assert [c[0] for c in calls] == ["embed"] + ["hook", 0, "hook", 1, "hook", 2]
+
+
+def test_the_decode_loop_takes_ids_or_embeddings(fast_paths, decode_core):
+    core, _ = decode_core
+    fast_paths(enabled=())
+    with pytest.raises(ValueError, match="cannot specify both input_ids and inputs_embeds"):
+        core._decode.forward(_NoContexts(), input_ids="ids", inputs_embeds=torch.zeros(1))
+    with pytest.raises(ValueError, match="must specify either one"):
+        core._decode.forward(_NoContexts())
+
+
+def test_the_decode_loop_never_reads_the_step_kind():
+    source = inspect.getsource(modeling.DecodeTarget.forward)
+    assert "num_contexts" not in source and "num_ctx" not in source
+
+
+# --- the shell's aliases ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("linked", [True, False], ids=["prefetch", "no prefetch"])
+def test_the_shell_links_the_expert_prefetch_after_the_builtin_aliases(
+    monkeypatch, fast_paths, linked
+):
+    paths = fast_paths(enabled={"cross_layer_expert_prefetch"} if linked else ())
+    events = paths.events
+    paths.entries[("cross_layer_expert_prefetch", "link_layers")] = (
+        lambda layers_, n: events.append(("link_layers", layers_, n))
+    )
+    monkeypatch.setattr(
+        GlmMoeDsaForCausalLM, "setup_aliases", lambda self: events.append("built-in aliases")
+    )
+    shell = _bare(modeling.ModelingV2Glm52Nvfp4Sm103Tp8)
+    shell.model = types.SimpleNamespace(layers=["layer"])
+    shell.model_config = types.SimpleNamespace(
+        pretrained_config=types.SimpleNamespace(num_hidden_layers=78)
+    )
+
+    shell.setup_aliases()
+
+    assert events == ["built-in aliases"] + ([("link_layers", ["layer"], 78)] if linked else [])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py
new file mode 100644
index 000000000000..4bc0a870e4ff
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py
@@ -0,0 +1,213 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The phase of a model step of the GLM-5.2 tp8 target: the scope, the targets' phases and the
+shell's forward, which runs every step under the phase of the target that takes it.
+
+Host-only: the shell is created without its constructor around an adopted built-in text model
+without layers, and the built-in causal LM's forward is replaced by a recorder that also calls
+a stand-in speculative worker, so each test sees the phase every part of a step runs under.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import pytest
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model
+from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM
+
+pytestmark = pytest.mark.cpu_only
+
+PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
+
+
+class _Metadata:
+    def __init__(self, num_contexts: int) -> None:
+        self.num_contexts = num_contexts
+
+
+@pytest.fixture(autouse=True)
+def _no_phase():
+    assert dispatch.phase() is None
+    yield
+    assert dispatch.phase() is None
+
+
+# --- the scope ----------------------------------------------------------------------------
+
+
+def test_there_is_no_phase_outside_a_step():
+    assert dispatch.phase() is None
+
+
+@pytest.mark.parametrize("phase", [PREFILL, DECODE])
+def test_a_step_runs_under_its_phase(phase):
+    with dispatch.step(phase):
+        assert dispatch.phase() is phase
+    assert dispatch.phase() is None
+
+
+def test_steps_nest_and_restore_the_enclosing_phase():
+    with dispatch.step(PREFILL):
+        with dispatch.step(DECODE):
+            assert dispatch.phase() is DECODE
+            with dispatch.step(PREFILL):
+                assert dispatch.phase() is PREFILL
+            assert dispatch.phase() is DECODE
+        assert dispatch.phase() is PREFILL
+
+
+def test_a_step_that_raises_restores_the_enclosing_phase():
+    with dispatch.step(PREFILL):
+        with pytest.raises(ValueError, match="inside"):
+            with dispatch.step(DECODE):
+                raise ValueError("inside")
+        assert dispatch.phase() is PREFILL
+
+
+def test_the_phases():
+    assert [p.value for p in dispatch.Phase] == ["prefill", "decode"]
+
+
+# --- the targets --------------------------------------------------------------------------
+
+
+def test_each_target_names_its_phase():
+    assert modeling.PrefillTarget.phase is PREFILL
+    assert modeling.DecodeTarget.phase is DECODE
+
+
+@pytest.mark.parametrize(
+    "num_contexts, phase",
+    [(0, DECODE), (1, PREFILL), (3, PREFILL)],
+    ids=["decode", "prefill", "mixed"],
+)
+def test_the_selected_target_carries_the_phase_of_the_step(num_contexts, phase):
+    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
+    assert core._select_target(_Metadata(num_contexts)).phase is phase
+
+
+# --- the shell's forward ------------------------------------------------------------------
+
+
+def _builtin_text_model() -> DeepseekV3Model:
+    model = DeepseekV3Model.__new__(DeepseekV3Model)
+    nn.Module.__init__(model)
+    return model
+
+
+@pytest.fixture
+def shell_forward(monkeypatch):
+    """A shell around an adopted core, and the phases its step's parts ran under."""
+    seen = []
+
+    def builtin_forward(
+        self,
+        attn_metadata,
+        input_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        return_context_logits=False,
+        spec_metadata=None,
+        resource_manager=None,
+        **kwargs,
+    ):
+        seen.append(("model", dispatch.phase()))
+        seen.append(
+            (
+                "arguments",
+                attn_metadata,
+                input_ids,
+                position_ids,
+                inputs_embeds,
+                return_context_logits,
+                spec_metadata,
+                resource_manager,
+                kwargs,
+            )
+        )
+        # The speculative worker and its draft loop run inside the built-in forward.
+        seen.append(("spec worker", dispatch.phase()))
+        if kwargs.get("fail"):
+            raise RuntimeError("forward failed")
+        return "logits"
+
+    monkeypatch.setattr(SpecDecOneEngineForCausalLM, "forward", builtin_forward)
+    shell = modeling.ModelingV2Glm52Nvfp4Sm103Tp8.__new__(modeling.ModelingV2Glm52Nvfp4Sm103Tp8)
+    nn.Module.__init__(shell)
+    shell.model = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
+    return shell, seen
+
+
+@pytest.mark.parametrize(
+    "num_contexts, phase",
+    [(0, DECODE), (1, PREFILL), (2, PREFILL)],
+    ids=["decode", "prefill", "mixed"],
+)
+def test_the_shell_runs_the_whole_step_under_the_phase_of_its_target(
+    shell_forward, num_contexts, phase
+):
+    shell, seen = shell_forward
+    md = _Metadata(num_contexts)
+
+    out = shell.forward(
+        attn_metadata=md,
+        input_ids="ids",
+        position_ids="positions",
+        inputs_embeds=None,
+        return_context_logits=True,
+        spec_metadata="spec",
+        resource_manager="resources",
+        lora_params="lora",
+    )
+
+    assert out == "logits"
+    assert seen == [
+        ("model", phase),
+        (
+            "arguments",
+            md,
+            "ids",
+            "positions",
+            None,
+            True,
+            "spec",
+            "resources",
+            {"lora_params": "lora"},
+        ),
+        ("spec worker", phase),
+    ]
+    assert dispatch.phase() is None
+
+
+def test_the_shell_takes_the_builtin_arguments_positionally(shell_forward):
+    shell, seen = shell_forward
+    md = _Metadata(0)
+
+    assert shell.forward(md, "ids", "positions", None, False, "spec", "resources") == "logits"
+    assert seen[1] == ("arguments", md, "ids", "positions", None, False, "spec", "resources", {})
+
+
+def test_a_step_that_fails_leaves_no_phase_behind(shell_forward):
+    shell, seen = shell_forward
+    with pytest.raises(RuntimeError, match="forward failed"):
+        shell.forward(attn_metadata=_Metadata(0), fail=True)
+    assert seen[-1] == ("spec worker", DECODE)
+    assert dispatch.phase() is None
+
+
+def test_the_shell_forward_takes_the_builtin_parameters():
+    builtin = inspect.signature(SpecDecOneEngineForCausalLM.forward)
+    shell = inspect.signature(modeling.ModelingV2Glm52Nvfp4Sm103Tp8.forward)
+    assert list(shell.parameters) == list(builtin.parameters)
+    for name, parameter in builtin.parameters.items():
+        assert shell.parameters[name].default == parameter.default, name
+        assert shell.parameters[name].kind == parameter.kind, name
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py
new file mode 100644
index 000000000000..27afa09d82d3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py
@@ -0,0 +1,573 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Which fast paths of the GLM-5.2 tp8 target run: the installed plan, the knobs the target
+refuses, the kernel registry, the capability log and the knobs folded into constants.
+
+Host-only: the knobs are set per test, the log goes into a list and the kernel imports are
+recorded, so nothing here needs a GPU or a kernel package.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import types
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
+    MODELING_V2_VALIDATE_ENV,
+    ModelingV2Mode,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+AUTO = ModelingV2Mode.AUTO
+ACCUMULATION = dispatch.KNOB_PREFIX + "ROUTED_MOE_INTO_HOP_ACCUMULATION"
+SCOPE = dispatch.KNOB_PREFIX + "DEFERRED_NORM_SCOPE"
+PARTS = ("sampler_store", "execution_stream_loop")
+
+# The record of every fast path but the parts that only deferred_input_prep brings along.
+SERVED_PATHS = (
+    "paths=absorbed_query,attention_output_into_hop,attention_output_tail,"
+    "attention_output_tail_prefetch,cross_layer_expert_prefetch,cross_rank_topk[rank-consistent],"
+    "cross_rank_topk_cluster_merge[rank-consistent],cross_rank_topk_global_emit[rank-consistent],"
+    "decode_prologue,deferred_input_prep,deferred_norm[shared],deferred_norm_evict_hints,dense_mlp,"
+    "draft_absorbed_query,draft_front,draft_live_rows,draft_loop,draft_routed_moe,draft_tail,"
+    "fused_kv_indexer_projection[kw-fused-a-order],hop_routing,host_page_retention,"
+    "indexer_layer_front,indexer_topk,oproj_weight_staging,post_attention_hop,post_moe_hop,"
+    "rope_kv_append,routed_moe,routed_moe_into_hop[fp32],router_gemm,shared_expert_prefetch,"
+    "shared_indexer_absorbed_query,sparse_mla_decode"
+)
+
+# The paths each path has no effect without: the closure of ``depends_on``.
+REQUIRED = {
+    "shared_indexer_absorbed_query": {"absorbed_query"},
+    "attention_output_tail_prefetch": {"attention_output_tail"},
+    "attention_output_into_hop": {"attention_output_tail", "post_attention_hop"},
+    "routed_moe_into_hop": {"post_moe_hop", "routed_moe"},
+    "deferred_norm": {"absorbed_query", "post_moe_hop"},
+    "cross_layer_expert_prefetch": {"routed_moe"},
+    "hop_routing": {
+        "attention_output_tail",
+        "post_attention_hop",
+        "attention_output_into_hop",
+        "routed_moe",
+    },
+    "oproj_weight_staging": {
+        "attention_output_tail",
+        "post_attention_hop",
+        "sparse_mla_decode",
+        "attention_output_into_hop",
+    },
+    "fused_kv_indexer_projection": {"absorbed_query"},
+    "host_page_retention": {"post_attention_hop"},
+    "shared_expert_prefetch": {
+        "attention_output_tail",
+        "post_moe_hop",
+        "routed_moe",
+        "post_attention_hop",
+        "attention_output_into_hop",
+        "routed_moe_into_hop",
+        "hop_routing",
+    },
+    "cross_rank_topk_cluster_merge": {"cross_rank_topk"},
+    "deferred_norm_evict_hints": {"absorbed_query", "post_moe_hop", "deferred_norm"},
+    "cross_rank_topk_global_emit": {
+        "rope_kv_append",
+        "cross_rank_topk",
+        "cross_rank_topk_cluster_merge",
+    },
+    "decode_prologue": {"draft_loop"},
+    "draft_live_rows": {"post_attention_hop", "draft_loop"},
+    "indexer_layer_front": {
+        "absorbed_query",
+        "indexer_topk",
+        "draft_absorbed_query",
+        "rope_kv_append",
+        "cross_rank_topk",
+    },
+}
+
+RECORD_PREFIX = "[glm_moe_dsa] dispatch: record "
+_RECORD = re.compile(r"^paths=(\S+) side_selectors=(\{.*\}) shared_env=(\{.*\})$")
+
+
+@pytest.fixture
+def lines(monkeypatch):
+    """A fresh selector state that logs into the returned list; no knob of the target is set."""
+    for name in list(os.environ):
+        if name.startswith(dispatch.SIDE_SELECTOR_PREFIXES) or name in dispatch.SHARED_ENV:
+            monkeypatch.delenv(name)
+    logged = []
+    monkeypatch.setattr(dispatch, "_log_info", logged.append)
+    monkeypatch.delenv(MODELING_V2_VALIDATE_ENV, raising=False)
+    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
+    monkeypatch.setattr(dispatch, "_VALIDATE", [False])
+    monkeypatch.setattr(dispatch, "_require", [False])
+    monkeypatch.setattr(dispatch, "_REGISTRY", {})
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    return logged
+
+
+def _parse(record):
+    m = _RECORD.match(record)
+    assert m, record
+    return m.group(1), json.loads(m.group(2)), json.loads(m.group(3))
+
+
+def _install(environ):
+    """Resolve and install the plan of ``environ``; return the paths it runs."""
+    dispatch.install(dispatch.resolve(AUTO, environ), environ)
+    return dispatch.installed()
+
+
+# --- the vocabulary -----------------------------------------------------------------------
+
+
+def test_the_selector_takes_every_fast_path_but_the_parts():
+    assert len(dispatch.FAST_PATH_NAMES) == 36
+    assert dispatch.SELECTOR_PATHS == tuple(n for n in dispatch.FAST_PATH_NAMES if n not in PARTS)
+    assert len(dispatch.SELECTOR_PATHS) == 34
+    assert dict(dispatch.PART_PATHS) == {
+        "prep": "deferred_input_prep",
+        "store": "sampler_store",
+        "stream": "execution_stream_loop",
+    }
+    assert dispatch.PARTS_ENV == "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_INPUT_PREP_PARTS"
+
+
+def test_the_prefetch_schedule_extends_the_attention_tail():
+    names = list(dispatch.FAST_PATH_NAMES)
+    assert names.index("attention_output_tail_prefetch") == names.index("attention_output_tail") + 1
+    prefetch = dispatch.FAST_PATHS[names.index("attention_output_tail_prefetch")]
+    assert prefetch.depends_on == ("attention_output_tail",)
+
+
+def test_every_fast_path_has_its_kernel_directory():
+    directories = set(dispatch.FAST_PATH_NAMES) - set(dispatch.KERNEL_DIRECTORY)
+    assert len(directories) == 28
+    assert len(dispatch.KERNEL_DIRECTORY) == 8
+    assert set(dispatch.KERNEL_DIRECTORY) <= set(dispatch.FAST_PATH_NAMES)
+    assert set(dispatch.KERNEL_DIRECTORY.values()) <= directories
+    assert dispatch.KERNELS_PACKAGE == dispatch.__name__.rpartition(".")[0] + ".kernels"
+
+
+def test_required_features_are_the_closure_of_the_table():
+    assert {k: set(v) for k, v in dispatch.REQUIRED_FEATURES.items()} == REQUIRED
+
+
+def test_the_variant_table_names_fast_paths_and_knobs():
+    for path, (name, choices) in dispatch.NUMERICAL_VARIANTS.items():
+        assert path in dispatch.FAST_PATH_NAMES
+        assert choices and len(set(choices)) == len(choices)
+        assert all(re.fullmatch(r"[a-z0-9][a-z0-9-]*", c) for c in choices)
+        assert name is None or name.startswith(dispatch.KNOB_PREFIX)
+    assert dispatch.NUMERICAL_VARIANTS["routed_moe_into_hop"] == (
+        ACCUMULATION,
+        ("fp32", "identity"),
+    )
+    # The deferred-norm scope is fixed in this target: no knob selects a variant.
+    assert dispatch.NUMERICAL_VARIANTS["deferred_norm"] == (None, ("shared",))
+
+
+# --- a path and its bases -----------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path", dispatch.FAST_PATH_NAMES)
+def test_every_path_runs_on_its_minimal_base(lines, path):
+    """Turning off every path outside a path and its bases leaves exactly those running."""
+    kept = {path} | set(dispatch.REQUIRED_FEATURES.get(path, ()))
+    off = ",".join(sorted(set(dispatch.FAST_PATH_NAMES) - kept))
+    assert _install({dispatch.DISABLE_ENV: off} if off else {}) == kept
+
+
+# --- variants and the record --------------------------------------------------------------
+
+
+def test_the_served_record_names_every_variant(lines):
+    paths = frozenset(dispatch.SELECTOR_PATHS)
+    assert (
+        dispatch.feature_record(paths, {}) == f"{SERVED_PATHS} side_selectors={{}} shared_env={{}}"
+    )
+    assert dispatch.feature_variants(paths, {}) == {
+        "cross_rank_topk": "rank-consistent",
+        "cross_rank_topk_cluster_merge": "rank-consistent",
+        "cross_rank_topk_global_emit": "rank-consistent",
+        "deferred_norm": "shared",
+        "fused_kv_indexer_projection": "kw-fused-a-order",
+        "routed_moe_into_hop": "fp32",
+    }
+
+
+@pytest.mark.parametrize("mode", ["fp32", "identity", " Identity "])
+def test_the_accumulation_variant_follows_its_knob(mode):
+    paths = {"post_moe_hop", "routed_moe", "routed_moe_into_hop"}
+    env = {ACCUMULATION: mode}
+    listed, side, _ = _parse(dispatch.feature_record(paths, env))
+    assert listed == f"post_moe_hop,routed_moe,routed_moe_into_hop[{mode.strip().lower()}]"
+    assert side == env
+
+
+@pytest.mark.parametrize("scope", ["all", "shared", " Shared ", ""])
+def test_the_deferred_norm_variant_is_the_fixed_scope(scope):
+    """The record names the scope the hop runs; ``resolve`` refuses the scope variable."""
+    paths = {"absorbed_query", "post_moe_hop", "deferred_norm"}
+    env = {SCOPE: scope}
+    listed, side, _ = _parse(dispatch.feature_record(paths, env))
+    assert listed == "absorbed_query,deferred_norm[shared],post_moe_hop"
+    assert side == env
+
+
+@pytest.mark.parametrize(
+    "raw, env",
+    [
+        ("post_moe_hop,routed_moe,routed_moe_into_hop", {ACCUMULATION: "fp16"}),
+        # The finalized-row input stays a mode of the post-MoE hop, but no knob selects it.
+        ("post_moe_hop,routed_moe,routed_moe_into_hop", {ACCUMULATION: "row"}),
+    ],
+)
+def test_an_unknown_variant_of_a_selected_path_fails(raw, env):
+    with pytest.raises(ValueError, match="is not a .* variant"):
+        dispatch.feature_record(frozenset(raw.split(",")), env)
+
+
+def test_the_variant_knob_of_an_unselected_path_is_only_recorded():
+    env = {ACCUMULATION: "fp16", SCOPE: "x"}
+    listed, side, _ = _parse(dispatch.feature_record({"absorbed_query"}, env))
+    assert listed == "absorbed_query"
+    assert side == env
+
+
+def test_side_selectors_and_shared_env():
+    env = {
+        dispatch.DISABLE_ENV: "draft_tail",
+        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS": "100",
+        ACCUMULATION: "identity",
+        "TRTLLM_ENABLE_PDL": "1",
+        "TRTLLM_FORCE_MNNVL_AR": "0",
+        "PATH": "/usr/bin",
+    }
+    assert dispatch.side_selectors(env) == {
+        dispatch.DISABLE_ENV: "draft_tail",
+        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS": "100",
+        ACCUMULATION: "identity",
+    }
+    assert dispatch.shared_env(env) == {"TRTLLM_ENABLE_PDL": "1", "TRTLLM_FORCE_MNNVL_AR": "0"}
+    assert _parse(dispatch.feature_record({"absorbed_query"}, env)) == (
+        "absorbed_query",
+        dispatch.side_selectors(env),
+        dispatch.shared_env(env),
+    )
+
+
+# --- without a plan ---------------------------------------------------------------------
+
+
+def test_without_a_plan_no_fast_path_runs_and_nothing_is_logged(lines, monkeypatch):
+    """A process that does not build the target runs no fast path, whatever its environment."""
+    monkeypatch.setenv(dispatch.FEATURE_ENV, ",".join(dispatch.SELECTOR_PATHS))
+    monkeypatch.setenv(ACCUMULATION, "identity")
+    monkeypatch.setenv("TRTLLM_ENABLE_PDL", "1")
+    assert dispatch.installed() is None
+    assert dispatch.enabled_features() == frozenset()
+    assert not any(dispatch.is_enabled(path) for path in dispatch.FAST_PATH_NAMES)
+    assert not dispatch.validate_mode()
+    assert lines == []
+
+
+def test_asking_for_something_that_is_not_a_fast_path_fails(lines):
+    with pytest.raises(ValueError, match="'draft_tails' is not a fast path"):
+        dispatch.is_enabled("draft_tails")
+    assert lines == []
+
+
+# --- the installed plan -------------------------------------------------------------------
+
+
+def test_an_installed_plan_answers_every_query(lines):
+    plan = dispatch.resolve(AUTO, {})
+    lines.clear()
+    assert dispatch.installed() is None
+    dispatch.install(plan, {})
+    assert dispatch.installed() == frozenset(dispatch.FAST_PATH_NAMES)
+    assert all(dispatch.is_enabled(path) for path in dispatch.FAST_PATH_NAMES)
+    assert len(lines) == 1
+    listed, side, shared = _parse(lines[0].removeprefix(RECORD_PREFIX))
+    served = SERVED_PATHS.removeprefix("paths=").split(",")
+    assert listed.split(",") == sorted(served + list(PARTS))
+    assert (side, shared) == ({}, {})
+
+
+@pytest.mark.parametrize(
+    "part, other_parts",
+    [
+        ("deferred_input_prep", {"sampler_store", "execution_stream_loop"}),
+        ("sampler_store", {"deferred_input_prep", "execution_stream_loop"}),
+        ("execution_stream_loop", {"deferred_input_prep", "sampler_store"}),
+    ],
+)
+def test_each_part_of_the_input_preparation_is_a_path_of_its_own(lines, part, other_parts):
+    paths = _install({dispatch.DISABLE_ENV: part})
+    assert paths == frozenset(dispatch.FAST_PATH_NAMES) - {part}
+    assert other_parts <= paths
+
+
+@pytest.mark.parametrize("raw, validate", [("", False), ("0", False), ("1", True), (" On ", True)])
+def test_validate_follows_the_modeling_v2_switch(lines, monkeypatch, raw, validate):
+    monkeypatch.setenv(MODELING_V2_VALIDATE_ENV, raw)
+    plan = dispatch.resolve(AUTO, {})
+    assert plan.validate is validate
+    assert not dispatch.validate_mode()
+    dispatch.install(plan, {})
+    assert dispatch.validate_mode() is validate
+    dispatch.reset_fallback_state()
+    assert not dispatch.validate_mode()
+
+
+def test_disable_turns_off_a_path_and_its_dependents(lines):
+    dispatch.install(dispatch.resolve(AUTO, {dispatch.DISABLE_ENV: "post_attention_hop"}), {})
+    off = {path for path in dispatch.FAST_PATH_NAMES if not dispatch.is_enabled(path)}
+    assert off == {
+        "post_attention_hop",
+        "attention_output_into_hop",
+        "oproj_weight_staging",
+        "hop_routing",
+        "shared_expert_prefetch",
+        "draft_live_rows",
+        "host_page_retention",
+    }
+    listed, _, _ = _parse(lines[-1].removeprefix(RECORD_PREFIX))
+    assert set(listed.split(",")).isdisjoint(off)
+
+
+def test_a_bad_variant_of_an_enabled_path_fails_at_install(lines):
+    plan = dispatch.resolve(AUTO, {})
+    with pytest.raises(ValueError, match="is not a routed_moe_into_hop variant"):
+        dispatch.install(plan, {ACCUMULATION: "fp16"})
+    assert dispatch.installed() is None
+
+
+def test_the_variant_knob_of_a_disabled_path_is_only_recorded(lines):
+    env = {dispatch.DISABLE_ENV: "routed_moe_into_hop", ACCUMULATION: "fp16"}
+    dispatch.install(dispatch.resolve(AUTO, env), env)
+    assert not dispatch.is_enabled("routed_moe_into_hop")
+
+
+@pytest.mark.parametrize(
+    "knob", sorted(dispatch.REMOVED_KNOBS), ids=lambda k: k.removeprefix(dispatch.KNOB_PREFIX)
+)
+@pytest.mark.parametrize("value", ["deferred_input_prep", ""], ids=["set", "empty"])
+def test_a_removed_knob_is_refused_with_what_to_use_instead(lines, knob, value):
+    with pytest.raises(ValueError) as excinfo:
+        dispatch.resolve(AUTO, {knob: value})
+    assert str(excinfo.value) == (
+        f"this target does not read {knob} ({dispatch.REMOVED_KNOBS[knob]}); unset them"
+    )
+    assert dispatch.installed() is None
+    assert lines == []
+
+
+def test_every_removed_knob_set_is_named_in_one_refusal(lines):
+    env = {knob: "1" for knob in dispatch.REMOVED_KNOBS}
+    env[dispatch.DISABLE_ENV] = "draft_tail"
+    with pytest.raises(ValueError) as excinfo:
+        dispatch.resolve(AUTO, env)
+    message = str(excinfo.value)
+    assert all(f"{knob} (" in message for knob in dispatch.REMOVED_KNOBS)
+    assert f"{dispatch.DISABLE_ENV} (" not in message
+
+
+def test_the_removed_knobs_point_at_what_replaces_them():
+    hints = dispatch.REMOVED_KNOBS
+    assert set(hints) == {
+        dispatch.FEATURE_ENV,
+        dispatch.PARTS_ENV,
+        dispatch.KNOB_PREFIX + "DEFERRED_INPUT_PREP_DEBUG",
+        dispatch.KNOB_PREFIX + "DRAFT_LOOP_VERIFY",
+        dispatch.KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_STATS",
+        dispatch.KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_REPORT_EVERY",
+        dispatch.KNOB_PREFIX + "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE",
+    }
+    assert "DISABLE" in hints[dispatch.KNOB_PREFIX + "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE"]
+    assert dispatch.DISABLE_ENV in hints[dispatch.FEATURE_ENV]
+    assert dispatch.DISABLE_ENV in hints[dispatch.PARTS_ENV]
+    for part, path in dispatch.PART_PATHS.items():
+        assert f"{part} = {path}" in hints[dispatch.PARTS_ENV]
+    for knob in ("DEFERRED_INPUT_PREP_DEBUG", "DRAFT_LOOP_VERIFY"):
+        assert MODELING_V2_VALIDATE_ENV in hints[dispatch.KNOB_PREFIX + knob]
+
+
+def test_resetting_the_fallback_state_uninstalls_the_plan(lines):
+    dispatch.install(dispatch.resolve(AUTO, {dispatch.DISABLE_ENV: "draft_tail"}), {})
+    assert not dispatch.is_enabled("draft_tail")
+    dispatch.reset_fallback_state()
+    assert dispatch.installed() is None
+
+
+# --- the kernel registry ------------------------------------------------------------------
+
+
+def test_register_and_look_up(lines, monkeypatch):
+    monkeypatch.setattr(
+        dispatch, "importlib", types.SimpleNamespace(import_module=lambda name: pytest.fail(name))
+    )
+
+    def impl():
+        return 42
+
+    assert dispatch.register("draft_front", "draft_front", impl) is impl
+    assert dispatch.get_impl("draft_front", "draft_front")() == 42
+    with pytest.raises(KeyError, match="'draft_front' does not register 'missing'"):
+        dispatch.get_impl("draft_front", "missing")
+    with pytest.raises(ValueError, match="'not_a_path' is not a fast path"):
+        dispatch.register("not_a_path", "draft_front", impl)
+
+
+@pytest.mark.parametrize("path", dispatch.FAST_PATH_NAMES)
+def test_get_impl_imports_the_kernel_directory_of_the_path(lines, monkeypatch, path):
+    imported = []
+    monkeypatch.setattr(dispatch, "importlib", types.SimpleNamespace(import_module=imported.append))
+    with pytest.raises(KeyError, match=f"{path!r} does not register 'entry'"):
+        dispatch.get_impl(path, "entry")
+    directory = dispatch.KERNEL_DIRECTORY.get(path, path)
+    assert imported == [f"{dispatch.KERNELS_PACKAGE}.{directory}"]
+
+
+def test_capability_lines_are_logged_once_per_key(lines):
+    assert dispatch.log_capability("draft_front", "draft_front", True)
+    assert not dispatch.log_capability("attention_output_tail", "layer3", False, "dense")
+    # One line per (path, key) and process; the call still returns its own decision.
+    assert not dispatch.log_capability("draft_front", "draft_front", False)
+    assert lines == [
+        "[glm_moe_dsa] draft_front/draft_front: engaged",
+        "[glm_moe_dsa] attention_output_tail/layer3: native fallback (dense)",
+    ]
+
+
+# --- knobs folded into constants --------------------------------------------------------
+
+# Every variable with the knob prefix that the kernels of this target once read; they run at
+# fixed launch parameters now, so ``resolve`` refuses each of them when it is set.
+FOLDED_KNOBS = tuple(
+    dispatch.KNOB_PREFIX + name
+    for name in (
+        "DEFERRED_NORM_HOP_CARVEOUT",
+        "DEFERRED_NORM_PROJECTION_CONFIG",
+        "DEFERRED_NORM_EVICT_HINTS_HOP_CONFIG",
+        "DEFERRED_NORM_EVICT_HINTS_PROJECTION_CONFIG",
+        "ABSORBED_QUERY_FORM",
+        "FUSED_KV_INDEXER_PROJECTION_LAYERS",
+        "FUSED_KV_INDEXER_PROJECTION_CONFIG",
+        "INDEXER_LAYER_FRONT_SKIP_REUSED_INDEXER_Q",
+        "CROSS_RANK_TOPK_LOGITS_PDL",
+        "ATTENTION_OUTPUT_TAIL_PDL_ENTRY",
+        "OPROJ_WEIGHT_STAGING_CONFIG",
+        "ROUTER_GEMM_EXPERTS_PER_CTA",
+        "ROUTER_GEMM_PAIRED_FMA",
+        "ROUTED_MOE_SCHEDULE",
+        "ROUTED_MOE_EXPERT_CLUSTER_CTAS",
+        "ROUTED_MOE_FINALIZE_MODE",
+        "ROUTED_MOE_PDL_TRIGGER",
+        "ROUTED_MOE_MIN_CLUSTERS",
+        "ROUTED_MOE_STAGES",
+        "ROUTED_MOE_GRID_WIDE_L2_PREFETCH",
+        "ROUTED_MOE_MIN_FC1_KBLOCKS",
+        "ROUTED_MOE_ACCUMULATOR_STAGES",
+        "ROUTED_MOE_HANDOFF",
+        "ROUTED_MOE_DELAY_SITES",
+        "ROUTED_MOE_DELAY_NS",
+        "ROUTED_MOE_EXPERT_CLUSTER_L2_PREFETCH",
+        "ROUTED_MOE_TWO_ROUND_CLUSTERS",
+        "SHARED_EXPERT_PREFETCH_CLUSTERS",
+        "SHARED_EXPERT_PREFETCH_WEIGHTS",
+        "DRAFT_ROUTED_MOE_FC1_CONFIG",
+        "DRAFT_ROUTED_MOE_FC2_CONFIG",
+        "DRAFT_ROUTED_MOE_FLAGS",
+        "DRAFT_TAIL_MODE",
+        "DRAFT_TAIL_EXCHANGE",
+        "DRAFT_TAIL_VARIANT",
+        "DRAFT_TAIL_CTAS_PER_SM",
+        "DRAFT_TAIL_NUM_CTAS",
+        "DRAFT_TAIL_PDL",
+        "DRAFT_TAIL_TIMEOUT_CYCLES",
+        "DRAFT_TAIL_KEEP_PERMILLE",
+        "CROSS_LAYER_EXPERT_PREFETCH_START_POINT",
+        "CROSS_LAYER_EXPERT_PREFETCH_CHUNK",
+        "CROSS_LAYER_EXPERT_PREFETCH_CTAS",
+        "CROSS_LAYER_EXPERT_PREFETCH_DELAY_NS",
+        "CROSS_LAYER_EXPERT_PREFETCH_FUSED",
+        "DRAFT_LOOP_PARTS",
+        "DEFERRED_NORM_SCOPE",
+        "CROSS_LAYER_EXPERT_PREFETCH_DEPTH",
+        "CROSS_LAYER_EXPERT_PREFETCH_PREDICTOR",
+        "CROSS_LAYER_EXPERT_PREFETCH_ISSUE_KERNEL",
+        "CROSS_LAYER_EXPERT_PREFETCH_ROUTED_KERNEL_CTAS",
+        "DRAFT_ABSORBED_QUERY_SKIP_REUSED_INDEXER_Q",
+        "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE",
+        "ROPE_KV_APPEND_METADATA_HINTS",
+        "ROPE_KV_APPEND_METADATA_PREFETCH",
+    )
+)
+
+
+def test_the_knobs_of_the_target_are_the_plan_knobs():
+    assert dispatch.KNOBS == {
+        dispatch.DISABLE_ENV,
+        ACCUMULATION,
+        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS",
+    }
+    assert not set(FOLDED_KNOBS) & dispatch.KNOBS
+
+
+@pytest.mark.parametrize("knob", FOLDED_KNOBS, ids=lambda k: k.removeprefix(dispatch.KNOB_PREFIX))
+def test_a_folded_knob_is_refused_even_at_its_value(lines, knob):
+    with pytest.raises(ValueError) as excinfo:
+        dispatch.resolve(AUTO, {knob: "0"})
+    hint = dispatch.REMOVED_KNOBS.get(
+        knob, "the fast paths of this target run at fixed launch parameters"
+    )
+    assert str(excinfo.value) == f"this target does not read {knob} ({hint}); unset them"
+    assert dispatch.installed() is None
+    assert lines == []
+
+
+def test_the_knobs_of_the_target_resolve(lines):
+    env = {
+        dispatch.DISABLE_ENV: "draft_tail",
+        ACCUMULATION: "identity",
+        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS": "100",
+    }
+    plan = dispatch.resolve(AUTO, env)
+    assert plan.disabled == {"draft_tail"}
+    assert plan.knobs == tuple(sorted(env.items()))
+
+
+# Shared TensorRT-LLM switches a kernel of this target may read next to the plan.
+_SHARED_READS = {"TRTLLM_ENABLE_PDL", "TRTLLM_FORCE_MNNVL_AR", "TLLM_TEST_MNNVL"}
+
+
+def test_only_dispatch_reads_the_knobs():
+    """No module of the family but ``dispatch`` reads the environment, except the shared switches."""
+    root = os.path.dirname(os.path.dirname(dispatch.__file__))
+    reads = []
+    for directory, _, files in os.walk(root):
+        for name in sorted(files):
+            if not name.endswith(".py") or os.path.join(directory, name) == dispatch.__file__:
+                continue
+            path = os.path.join(directory, name)
+            with open(path) as f:
+                text = f.read()
+            for match in re.finditer(r"os\.(?:environ|getenv)\b", text):
+                read = re.match(
+                    r"os\.(?:environ(?:\.get)?|getenv)\s*[\[(]\s*[\"']([A-Z_]+)[\"']",
+                    text[match.start() :],
+                )
+                if read is None or read.group(1) not in _SHARED_READS:
+                    line = text.count("\n", 0, match.start()) + 1
+                    reads.append(f"{os.path.relpath(path, root)}:{line}")
+    assert reads == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py
new file mode 100644
index 000000000000..951159049d65
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py
@@ -0,0 +1,514 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 tp8 target's speculative worker: its adoption, the phase gate of every override,
+the draft loop it runs on a decode step, and the draft layer's live-row arming.
+
+Host-only: the built-in worker is created without its constructor and carries only the
+attributes the overrides read. The kernel entry points are stand-ins, and the built-in methods
+an override falls back to or calls are recorders.
+"""
+
+from __future__ import annotations
+
+import ast
+import contextlib
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+    layers,
+    spec_worker,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows import (  # noqa: E501
+    step0,
+)
+from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker, MTPEagleWorker
+from tensorrt_llm._torch.speculative.mtp_dynamic_tree import MTPEagleDynamicTreeWorker
+
+pytestmark = pytest.mark.cpu_only
+
+PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
+GlmMtpWorker = spec_worker.GlmMtpWorker
+OVERRIDES = ["_forward_impl", "_forward_linear_draft_loop", "_run_draft_forward"]
+
+
+def _bare(cls, **attributes):
+    obj = cls.__new__(cls)
+    nn.Module.__init__(obj)
+    for name, value in attributes.items():
+        setattr(obj, name, value)
+    return obj
+
+
+def _keys(module):
+    return (
+        [name for name, _ in module.named_parameters()],
+        [name for name, _ in module.named_buffers()],
+        [name for name, _ in module.named_modules()],
+    )
+
+
+# --- adoption -----------------------------------------------------------------------------
+
+
+def test_adoption_reclasses_the_builtin_worker_in_place():
+    worker = _bare(MTPEagleWorker, is_mtp_eagle=True)
+    worker.register_buffer("probe", torch.zeros(1))
+    before = _keys(worker)
+
+    assert GlmMtpWorker.adopt(worker) is worker
+
+    assert type(worker) is GlmMtpWorker
+    assert worker._draft_loop_ops is None and "_draft_loop_ops" not in vars(worker)
+    assert _keys(worker) == before
+
+
+def _other_worker():
+    return _bare(type("_OtherMTPEagleWorker", (MTPEagleWorker,), {}))
+
+
+@pytest.mark.parametrize(
+    "make, got",
+    [
+        (lambda: _bare(Eagle3OneModelWorker), "Eagle3OneModelWorker"),
+        (lambda: _bare(MTPEagleDynamicTreeWorker), "MTPEagleDynamicTreeWorker"),
+        (_other_worker, "_OtherMTPEagleWorker"),
+        (lambda: GlmMtpWorker.adopt(_bare(MTPEagleWorker)), "GlmMtpWorker"),
+    ],
+    ids=["the Eagle3 worker", "the dynamic-tree worker", "a subclass", "an adopted worker"],
+)
+def test_adoption_refuses_anything_but_the_exact_builtin_worker(make, got):
+    with pytest.raises(TypeError, match=f"adopts the built-in MTPEagleWorker, got {got}$"):
+        GlmMtpWorker.adopt(make())
+
+
+def test_the_class_declares_no_state_of_its_own():
+    assert not {"__slots__", "__init__", "forward"} & set(vars(GlmMtpWorker))
+    assert GlmMtpWorker._draft_loop_ops is None
+
+
+@pytest.mark.parametrize("has_worker", [True, False], ids=["a worker", "no worker"])
+def test_the_worker_is_adopted_after_the_draft_layer(monkeypatch, has_worker):
+    order = []
+    for cls, label in (
+        (layers.GlmDecoderLayer, "decoder layer"),
+        (layers.GlmMtpLayer, "draft layer"),
+        (GlmMtpWorker, "worker"),
+    ):
+        monkeypatch.setattr(
+            cls, "adopt", classmethod(lambda c, module, label=label: order.append((label, module)))
+        )
+    core = types.SimpleNamespace(num_hidden_layers=2, layers=["layer 0", "layer 1", "draft"])
+    worker = object() if has_worker else None
+
+    modeling.adopt_model(types.SimpleNamespace(model=core, spec_worker=worker), plan=None)
+
+    expected = [
+        ("decoder layer", "layer 0"),
+        ("decoder layer", "layer 1"),
+        ("draft layer", "draft"),
+    ]
+    assert order == expected + ([("worker", worker)] if has_worker else [])
+
+
+# --- the phase gate -----------------------------------------------------------------------
+
+
+def _method(name):
+    tree = ast.parse(Path(spec_worker.__file__).read_text())
+    cls = next(c for c in tree.body if isinstance(c, ast.ClassDef) and c.name == "GlmMtpWorker")
+    return next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == name)
+
+
+@pytest.mark.parametrize("name", OVERRIDES)
+def test_each_override_starts_with_the_phase_gate(name):
+    """``if dispatch.phase() is not DECODE: return super().<name>(<every built-in parameter by
+    name>)``, after the built-in docstring if the method has one."""
+    node = _method(name)
+    docstring = isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant)
+    gate = node.body[1] if docstring else node.body[0]
+    assert ast.unparse(gate.test) == "dispatch.phase() is not dispatch.Phase.DECODE"
+    assert len(gate.body) == 1 and not gate.orelse
+    call = gate.body[0].value
+    assert ast.unparse(call.func) == f"super().{name}" and not call.args
+    parameters = inspect.signature(getattr(Eagle3OneModelWorker, name)).parameters
+    assert [k.arg for k in call.keywords] == [p for p in parameters if p != "self"]
+    assert all(k.arg == ast.unparse(k.value) for k in call.keywords)
+    assert list(inspect.signature(getattr(GlmMtpWorker, name)).parameters) == list(parameters)
+
+
+def _record(monkeypatch, owner, name, calls):
+    def recorder(*args, **kwargs):
+        calls.append((name, args, kwargs))
+        return f"{owner.__name__}.{name}"
+
+    monkeypatch.setattr(owner, name, recorder)
+
+
+@pytest.mark.parametrize("phase", [PREFILL, None], ids=["prefill", "outside a step"])
+def test_off_decode_every_override_runs_the_builtin_method(monkeypatch, phase):
+    worker = GlmMtpWorker.adopt(_bare(MTPEagleWorker, is_mtp_eagle=True))
+    calls = []
+    for name in OVERRIDES:
+        _record(monkeypatch, Eagle3OneModelWorker, name, calls)
+    monkeypatch.setattr(dispatch, "is_enabled", lambda path: pytest.fail(path))
+    monkeypatch.setattr(dispatch, "get_impl", lambda *a: pytest.fail(str(a)))
+
+    with dispatch.step(phase) if phase is not None else contextlib.nullcontext():
+        results = [
+            worker._forward_impl("ids", "pos", "h", "logits", "md", "spec", "draft", "rm"),
+            worker._forward_linear_draft_loop(
+                "inputs", "md", "spec", "draft", "kv", 0, 2, "accepted", "ranks"
+            ),
+            worker._run_draft_forward("draft", "inputs", "spec", 1),
+        ]
+
+    assert results == [f"Eagle3OneModelWorker.{name}" for name in OVERRIDES]
+    assert all(args == (worker,) for _, args, _ in calls)
+    assert [(name, kwargs) for name, _, kwargs in calls] == [
+        (
+            "_forward_impl",
+            dict(
+                input_ids="ids",
+                position_ids="pos",
+                hidden_states="h",
+                logits="logits",
+                attn_metadata="md",
+                spec_metadata="spec",
+                draft_model="draft",
+                resource_manager="rm",
+            ),
+        ),
+        (
+            "_forward_linear_draft_loop",
+            dict(
+                inputs="inputs",
+                attn_metadata="md",
+                spec_metadata="spec",
+                draft_model="draft",
+                draft_kv_cache_manager="kv",
+                num_contexts=0,
+                batch_size=2,
+                num_accepted_tokens="accepted",
+                original_all_rank_num_tokens="ranks",
+            ),
+        ),
+        (
+            "_run_draft_forward",
+            dict(draft_model="draft", inputs="inputs", spec_metadata="spec", step_idx=1),
+        ),
+    ]
+    assert "_draft_loop_ops" not in vars(worker)
+
+
+# --- the decode draft loop ----------------------------------------------------------------
+
+BATCH = 2
+HIDDEN = 4
+
+
+class _DraftLoopOps:
+    """Stands in for the draft-loop operations of one forward; records each call."""
+
+    def __init__(self, events, next_new_tokens="ops next new tokens"):
+        self.events = events
+        self.batch_size = BATCH
+        self.gather_ids = torch.arange(BATCH, dtype=torch.int64)
+        self.next_new_tokens = next_new_tokens
+
+    def accept(self, logits):
+        self.events.append(("accept", logits))
+        return "ops accepted", "ops num accepted"
+
+    def first_drafter_inputs(self, position_ids, hidden_states):
+        self.events.append(("first_drafter_inputs", tuple(position_ids.shape), hidden_states))
+        return "ops first inputs"
+
+    def step0_rows(self, hidden_states):
+        self.events.append(("step0_rows", hidden_states.shape[0]))
+        return self.gather_ids, hidden_states[:BATCH]
+
+    def step_update(self, attn_metadata, step, num_accepted, position_ids, gather_ids):
+        self.events.append(("step_update", step, position_ids is not None, gather_ids is not None))
+        return types.SimpleNamespace(position_ids=f"positions after step {step}")
+
+    def finalize(self, tokens):
+        self.events.append(("finalize", list(tokens)))
+        return "ops stacked"
+
+
+def _worker(events, **methods):
+    worker = GlmMtpWorker.adopt(
+        _bare(
+            MTPEagleWorker,
+            is_mtp_eagle=True,
+            sa_enhancer=None,
+            guided_decoder=None,
+            model_config=None,
+            _uses_external_shared_target_kv=False,
+            _is_mamba_hybrid_cache=False,
+            _auxiliary_state_handlers=[],
+        )
+    )
+    for name, result in methods.items():
+
+        def method(*args, _name=name, _result=result, **kwargs):
+            events.append((_name,))
+            return _result(worker) if callable(_result) else _result
+
+        setattr(worker, name, method)
+    return worker
+
+
+class _FastPaths:
+    """``dispatch.is_enabled`` / ``get_impl`` with stand-in draft-loop entry points."""
+
+    def __init__(self, events, enabled, entries):
+        self.events, self.enabled, self.entries = events, frozenset(enabled), entries
+
+    def is_enabled(self, path):
+        assert path in dispatch.FAST_PATH_NAMES, path
+        return path in self.enabled
+
+    def get_impl(self, path, name):
+        assert path in self.enabled, path
+        entry = self.entries[(path, name)]
+
+        def call(*args):
+            self.events.append((name,) + args)
+            return entry(*args)
+
+        return call
+
+
+@pytest.mark.parametrize("mode", ["engaged", "ops without next new tokens", "declined", "off"])
+def test_the_decode_forward_runs_the_draft_loop_operations(monkeypatch, mode):
+    events = []
+    ops = _DraftLoopOps(events, None if mode == "ops without next new tokens" else "ops tokens")
+    worker = _worker(
+        events,
+        _execute_guided_decoder_if_present=None,
+        sample_and_accept_draft_tokens=("accepted", "num accepted"),
+        _prepare_attn_metadata_for_spec_dec=None,
+        prepare_1st_drafter_inputs="first inputs",
+        get_draft_kv_cache_manager="draft kv",
+        _forward_draft_loop=lambda w: ("draft tokens", w._draft_loop_ops),
+        _restore_attn_metadata_from_spec_dec=None,
+        _prepare_next_new_tokens="next new tokens",
+    )
+    paths = _FastPaths(
+        events,
+        () if mode == "off" else {"draft_loop"},
+        {("draft_loop", "draft_loop_ops"): lambda *a: None if mode == "declined" else ops},
+    )
+    monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+    monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+    md = types.SimpleNamespace(num_seqs=BATCH, num_contexts=0, all_rank_num_tokens=None)
+    spec = types.SimpleNamespace(
+        runtime_draft_len=2, spec_resource_manager=None, batch_indices_cuda="batch indices"
+    )
+
+    with dispatch.step(DECODE):
+        out = worker._forward_impl(
+            "ids", torch.zeros(1, 6), "hidden", "logits", md, spec, "draft model", "rm"
+        )
+
+    engaged = mode in ("engaged", "ops without next new tokens")
+    start = [("_execute_guided_decoder_if_present",)]
+    if mode != "off":
+        start.append(("draft_loop_ops", worker, md, spec))
+    if engaged:
+        assert events == start + [
+            ("accept", "logits"),
+            ("_prepare_attn_metadata_for_spec_dec",),
+            ("first_drafter_inputs", (6,), "hidden"),
+            ("get_draft_kv_cache_manager",),
+            ("_forward_draft_loop",),
+            ("_restore_attn_metadata_from_spec_dec",),
+        ] + ([("_prepare_next_new_tokens",)] if ops.next_new_tokens is None else [])
+        assert out["new_tokens"] == "ops accepted"
+        assert out["new_tokens_lens"] == "ops num accepted"
+        assert out["next_draft_tokens"] == ("draft tokens", ops)
+        assert out["next_new_tokens"] == (ops.next_new_tokens or "next new tokens")
+    else:
+        assert events == start + [
+            ("sample_and_accept_draft_tokens",),
+            ("_prepare_attn_metadata_for_spec_dec",),
+            ("prepare_1st_drafter_inputs",),
+            ("get_draft_kv_cache_manager",),
+            ("_forward_draft_loop",),
+            ("_restore_attn_metadata_from_spec_dec",),
+            ("_prepare_next_new_tokens",),
+        ]
+        assert out["new_tokens"] == "accepted"
+        assert out["next_draft_tokens"] == ("draft tokens", None)
+        assert out["next_new_tokens"] == "next new tokens"
+    # The operations live for one forward only.
+    assert worker._draft_loop_ops is None and md.use_spec_decoding is True
+    assert out["logits"] == "logits"
+
+
+class _DraftLayerForward:
+    """The draft layer's decode forward as the draft loop sees it: it records the rows it was
+    armed with and, when armed, returns only the live rows, as the live-row hop does."""
+
+    def __init__(self, layer, compacts):
+        self.layer, self.compacts, self.calls = layer, compacts, []
+
+    def __call__(self, **kwargs):
+        rows = self.layer._draft_live_rows
+        self.calls.append((rows, kwargs["embed_tokens"], kwargs["all_rank_num_tokens"]))
+        if rows is not None and self.compacts:
+            self.layer._draft_live_rows_compacted = True
+            return torch.ones(rows.shape[0], HIDDEN)
+        return torch.zeros(3 * BATCH if len(self.calls) == 1 else BATCH, HIDDEN)
+
+
+@pytest.mark.parametrize("compacts", [True, False], ids=["compacted", "kept every row"])
+def test_the_decode_draft_loop_runs_the_draft_fast_paths(monkeypatch, compacts):
+    """Each draft-loop fast path is decided once before the steps; step 0 runs armed with the
+    draft-loop operations' live rows; the draft tail makes each step's tokens; the step updates
+    replace the native position advance and metadata update."""
+    events = []
+    ops = _DraftLoopOps(events)
+    draft_layer = _bare(layers.GlmMtpLayer)
+    forward = _DraftLayerForward(draft_layer, compacts)
+    draft_layer.forward = forward
+    draft_model = types.SimpleNamespace(embed_tokens="embedding table", mtp_layers=[draft_layer])
+    worker = _worker(events, draft_kv_cache_context=lambda w: contextlib.nullcontext(md))
+    worker._draft_loop_ops = ops
+    tail_tokens = iter(["tokens 0", "tokens 1"])
+    paths = _FastPaths(
+        events,
+        {"draft_loop", "draft_tail", "draft_live_rows"},
+        {
+            ("draft_tail", "select_draft_tail"): lambda *a: (
+                lambda rows: (events.append(("draft tail", rows.shape[0])), next(tail_tokens))[1]
+            ),
+            ("draft_live_rows", "select_step0"): lambda w, dm, s, o: step0.Step0LiveRows(
+                dm.mtp_layers[0], o.gather_ids
+            ),
+            ("draft_loop", "draft_embedding"): lambda table: ("draft embedding", table),
+        },
+    )
+    monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+    monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+    md = types.SimpleNamespace()
+    spec = types.SimpleNamespace(
+        runtime_draft_len=2,
+        all_rank_num_tokens=[3 * BATCH],
+        subseq_all_rank_num_tokens=[BATCH],
+        wants_advanced_draft_sampling=False,
+    )
+    inputs = {"input_ids": "ids", "position_ids": "positions", "hidden_states": "h"}
+
+    with dispatch.step(DECODE):
+        out = worker._forward_linear_draft_loop(
+            inputs, md, spec, draft_model, None, 0, BATCH, "num accepted", None
+        )
+
+    assert out == "ops stacked"
+    assert events[0] == ("select_draft_tail", worker, draft_model, spec)
+    assert events[1] == ("select_step0", worker, draft_model, spec, ops)
+    assert events[2] == ("draft_kv_cache_context",)
+    step_events = [e for e in events if e[0] in ("step0_rows", "draft tail", "step_update")]
+    assert step_events == ([] if compacts else [("step0_rows", 3 * BATCH)]) + [
+        ("draft tail", BATCH),
+        ("step_update", 0, True, True),
+        ("draft tail", BATCH),
+        ("step_update", 1, False, False),
+    ]
+    assert events[-1] == ("finalize", ["tokens 0", "tokens 1"])
+    # Step 0 ran armed with the live rows; step 1 did not; the layer is disarmed afterwards.
+    armed = [rows for rows, _, _ in forward.calls]
+    assert armed[0] is ops.gather_ids and armed[1] is None
+    assert [c[1:] for c in forward.calls] == [
+        (("draft embedding", "embedding table"), [3 * BATCH]),
+        (("draft embedding", "embedding table"), [BATCH]),
+    ]
+    assert draft_layer._draft_live_rows is None and draft_layer._draft_live_rows_compacted is False
+
+
+@pytest.mark.parametrize("draft_loop", [True, False], ids=["draft loop", "no draft loop"])
+def test_the_decode_draft_forward_takes_the_draft_loop_embedding(monkeypatch, draft_loop):
+    events = []
+    calls = []
+    layer = lambda **kwargs: calls.append(kwargs) or "draft layer output"  # noqa: E731
+    draft_model = types.SimpleNamespace(embed_tokens="embedding table", mtp_layers=[layer])
+    worker = _worker(events)
+    paths = _FastPaths(
+        events,
+        {"draft_loop"} if draft_loop else (),
+        {("draft_loop", "draft_embedding"): lambda table: ("draft embedding", table)},
+    )
+    monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+    monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+    spec = types.SimpleNamespace(all_rank_num_tokens=[6], subseq_all_rank_num_tokens=[2])
+
+    with dispatch.step(DECODE):
+        out = worker._run_draft_forward(draft_model, {"input_ids": "ids"}, spec, 0)
+
+    assert out == ("draft layer output", None)
+    table = ("draft embedding", "embedding table") if draft_loop else "embedding table"
+    assert calls == [dict(embed_tokens=table, all_rank_num_tokens=[6], input_ids="ids")]
+
+
+# --- the draft layer's live-row arming ----------------------------------------------------
+
+
+def _recording_draft_layer(writes):
+    class _Recording(layers.GlmMtpLayer):
+        def __setattr__(self, name, value):
+            if name.startswith("_draft_live_rows"):
+                writes.append((name, value))
+            super().__setattr__(name, value)
+
+    return _bare(_Recording)
+
+
+def test_the_draft_layer_is_armed_and_disarmed_through_its_api():
+    layer = _bare(layers.GlmMtpLayer)
+    rows = torch.arange(BATCH)
+    layer.arm_live_rows(rows)
+    assert layer._draft_live_rows is rows and layer._draft_live_rows_compacted is False
+    layer._draft_live_rows_compacted = True
+    assert layer.disarm_live_rows() is True
+    assert layer._draft_live_rows is None and layer._draft_live_rows_compacted is False
+    assert layer.disarm_live_rows() is False
+
+
+@pytest.mark.parametrize("outcome", ["compacted", "kept every row", "raised"])
+def test_step0_arming_writes_the_layer_state_in_the_builtin_order(outcome):
+    """Arm (rows, then not compacted), the draft layer's own write, then clear (rows, then
+    compacted), also when the step raises."""
+    writes = []
+    layer = _recording_draft_layer(writes)
+    rows = torch.arange(BATCH)
+    live = step0.Step0LiveRows(layer, rows)
+
+    with pytest.raises(RuntimeError) if outcome == "raised" else contextlib.nullcontext():
+        with live.armed() as armed:
+            assert armed is live and layer._draft_live_rows is rows
+            if outcome != "kept every row":
+                layer._draft_live_rows_compacted = True
+            if outcome == "raised":
+                raise RuntimeError("draft step failed")
+
+    assert live.compacted is (outcome == "compacted")
+    own = [("_draft_live_rows_compacted", True)] if outcome != "kept every row" else []
+    assert writes == [
+        ("_draft_live_rows", rows),
+        ("_draft_live_rows_compacted", False),
+        *own,
+        ("_draft_live_rows", None),
+        ("_draft_live_rows_compacted", False),
+    ]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py
new file mode 100644
index 000000000000..97e626cf595b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py
@@ -0,0 +1,597 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 tp8 core and its two targets, and what the shell keeps from the built-in model.
+
+Host-only: the core is adopted on a built-in text model created without its constructor
+(which builds CUDA streams and the decoder layers), and the built-in forward is replaced by
+a recorder, so each test sees exactly which call a step makes. The shell's construction runs
+with the built-in constructor stubbed, stand-in decoder layers, a recorder in place of
+``adopt_model`` (``test_modeling_v2_glm_layers.py`` covers it) and a stand-in for the
+sparse-MLA decode FMHA library.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+import sys
+import types
+
+import pytest
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2 import _router_index
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
+    modeling,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+    runtime,
+)
+from tensorrt_llm._torch.attention.backends.fmha import manager as fmha_manager
+from tensorrt_llm._torch.attention.backends.fmha import registry as fmha_registry
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of
+
+pytestmark = pytest.mark.cpu_only
+
+
+def _builtin_text_model() -> DeepseekV3Model:
+    """A built-in text model as far as class and module state go, without its layers."""
+    model = DeepseekV3Model.__new__(DeepseekV3Model)
+    nn.Module.__init__(model)
+    return model
+
+
+class _Hop:
+    """Stands in for a post-MoE hop object a decoder layer carries."""
+
+    def __init__(self, engaged: bool = True) -> None:
+        self.engaged = engaged
+        self.impl = "cluster"
+
+
+class _Layer(nn.Module):
+    """Stands in for a decoder layer with the fast-path attributes the constructors set."""
+
+    def __init__(self, post_moe_hop=None, router_gemm=None) -> None:
+        super().__init__()
+        self.self_attn = nn.Module()
+        self.self_attn._attn_tail = None
+        self.mlp = nn.Module()
+        self.mlp.router_gemm = router_gemm
+        self.post_attn_hop = None
+        self.post_moe_hop = post_moe_hop
+
+
+class _Metadata:
+    def __init__(self, num_contexts: int) -> None:
+        self.num_contexts = num_contexts
+
+
+def test_adopt_reclasses_the_builtin_text_model_in_place():
+    model = _builtin_text_model()
+    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(model)
+    assert core is model
+    assert type(core) is modeling.GlmMoeDsaTp8ModelingV2
+    assert isinstance(core, DeepseekV3Model)
+    assert isinstance(core._prefill, modeling.PrefillTarget) and core._prefill.core is core
+    assert isinstance(core._decode, modeling.DecodeTarget) and core._decode.core is core
+    # Plain attributes only: no module or parameter joins the built-in tree.
+    assert list(core.named_modules()) == [("", core)]
+    assert list(core.named_parameters()) == []
+
+
+def test_adopt_refuses_anything_but_the_builtin_text_model():
+    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
+    with pytest.raises(TypeError, match="GlmMoeDsaTp8ModelingV2"):
+        modeling.GlmMoeDsaTp8ModelingV2.adopt(core)
+    with pytest.raises(TypeError, match="Module"):
+        modeling.GlmMoeDsaTp8ModelingV2.adopt(nn.Module())
+
+
+def _recorder(calls):
+    def forward(
+        self,
+        attn_metadata,
+        input_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        spec_metadata=None,
+        **kwargs,
+    ):
+        calls.append((self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata))
+        calls.append(kwargs)
+        return "hidden states"
+
+    return forward
+
+
+def _run_step_both_ways(core, md, calls, receiver):
+    # Called the way the speculative causal LM calls its text model: keywords throughout.
+    out = core(
+        input_ids="ids",
+        attn_metadata=md,
+        position_ids="positions",
+        inputs_embeds=None,
+        spec_metadata="spec",
+        lora_params="lora",
+    )
+    assert out == "hidden states"
+    assert calls == [(receiver, md, "ids", "positions", None, "spec"), {"lora_params": "lora"}]
+
+    # And positionally, in the built-in forward's parameter order.
+    calls.clear()
+    assert core(md, "ids", "positions", None, "spec") == "hidden states"
+    assert calls == [(receiver, md, "ids", "positions", None, "spec"), {}]
+
+
+@pytest.mark.parametrize("num_contexts", [1, 4], ids=["prefill", "mixed"])
+def test_each_step_runs_the_builtin_forward_through_its_target(monkeypatch, num_contexts):
+    calls = []
+    monkeypatch.setattr(DeepseekV3Model, "forward", _recorder(calls))
+    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
+    md = _Metadata(num_contexts)
+
+    assert type(core._select_target(md)) is modeling.PrefillTarget
+    _run_step_both_ways(core, md, calls, core)
+
+
+def test_a_decode_step_runs_the_decode_targets_loop(monkeypatch):
+    calls, builtin = [], []
+    monkeypatch.setattr(DeepseekV3Model, "forward", _recorder(builtin))
+    monkeypatch.setattr(modeling.DecodeTarget, "forward", _recorder(calls))
+    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
+    md = _Metadata(0)
+
+    assert type(core._select_target(md)) is modeling.DecodeTarget
+    _run_step_both_ways(core, md, calls, core._decode)
+    assert builtin == []
+
+
+def test_the_forwards_take_the_builtin_parameters():
+    builtin = inspect.signature(DeepseekV3Model.forward)
+    for forward in (
+        modeling.GlmMoeDsaTp8ModelingV2.forward,
+        modeling.PrefillTarget.forward,
+        modeling.DecodeTarget.forward,
+    ):
+        assert list(inspect.signature(forward).parameters) == list(builtin.parameters), forward
+
+
+def test_the_shell_is_the_builtin_causal_lm():
+    shell = modeling.ModelingV2Glm52Nvfp4Sm103Tp8
+    assert issubclass(shell, GlmMoeDsaForCausalLM)
+    # The model loader inspects load_weights' parameters to decide what it passes.
+    assert (
+        inspect.getfullargspec(shell.load_weights).args
+        == inspect.getfullargspec(GlmMoeDsaForCausalLM.load_weights).args
+    )
+
+
+def _hooks_of_plan(disable: str = "") -> RuntimeHooks:
+    """``make_runtime_hooks`` of a model whose plan disables ``disable``."""
+    env = {dispatch.DISABLE_ENV: disable} if disable else {}
+    model = types.SimpleNamespace(
+        plan=dispatch.resolve(_router_index.ModelingV2Mode.AUTO, env),
+        model=types.SimpleNamespace(decode_prologue=None),
+        model_config=types.SimpleNamespace(
+            mapping=types.SimpleNamespace(enable_attention_dp=False)
+        ),
+    )
+    return runtime.make_runtime_hooks(model)
+
+
+class _StoreWriter:
+    """Stands in for ``SamplerStoreWriter``, whose constructor allocates pinned host memory."""
+
+
+@pytest.fixture
+def host_only_store_writer(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
+        deferred_input_prep,
+    )
+
+    monkeypatch.setattr(deferred_input_prep, "SamplerStoreWriter", _StoreWriter)
+
+
+@pytest.mark.usefixtures("host_only_store_writer")
+def test_the_target_declares_the_step_input_stager_of_its_plan(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        PrepStager,
+    )
+
+    validate_env = _router_index.MODELING_V2_VALIDATE_ENV
+    monkeypatch.delenv(validate_env, raising=False)
+    hooks = _hooks_of_plan()
+    assert isinstance(hooks.step_input_stager, PrepStager)
+    assert hooks.step_input_stager.debug is None
+
+    monkeypatch.setenv(validate_env, "1")
+    assert _hooks_of_plan().step_input_stager.debug.level == 1
+    monkeypatch.delenv(validate_env)
+
+    assert _hooks_of_plan("deferred_input_prep").step_input_stager is None
+
+
+@pytest.mark.usefixtures("host_only_store_writer")
+def test_the_target_declares_the_decode_prologue_of_its_plan():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E501
+        DecodePrologueHook,
+    )
+
+    hooks = _hooks_of_plan()
+    assert isinstance(hooks.decode_prologue, DecodePrologueHook)
+    assert _hooks_of_plan("decode_prologue").decode_prologue is None
+
+
+@pytest.mark.usefixtures("host_only_store_writer")
+def test_the_target_declares_the_sampler_store_writer_of_its_plan():
+    assert isinstance(_hooks_of_plan().sampler_store_writer, _StoreWriter)
+    hooks = _hooks_of_plan("sampler_store")
+    assert hooks.sampler_store_writer is None
+    assert hooks.step_input_stager is not None
+
+
+@pytest.mark.usefixtures("host_only_store_writer")
+@pytest.mark.parametrize(
+    "disabled, off",
+    [
+        ("", ()),
+        ("deferred_input_prep", ("step_input_stager",)),
+        ("sampler_store", ("sampler_store_writer",)),
+        ("execution_stream_loop", ("loop_on_execution_stream",)),
+        (
+            "deferred_input_prep,sampler_store,execution_stream_loop",
+            ("step_input_stager", "sampler_store_writer", "loop_on_execution_stream"),
+        ),
+    ],
+)
+def test_each_part_of_deferred_input_prep_is_a_hook_of_its_own(disabled, off):
+    """The three parts are separate fast paths: disabling one leaves the other two declared."""
+    hooks = _hooks_of_plan(disabled)
+    for field in ("step_input_stager", "sampler_store_writer", "loop_on_execution_stream"):
+        assert bool(getattr(hooks, field)) is (field not in off), field
+    assert isinstance(hooks.executor_events, runtime.GlmExecutorEvents)
+
+
+@pytest.fixture
+def event_calls(monkeypatch):
+    """The kernel entry points the executor events call, recorded in order."""
+    calls = []
+    kernels = runtime.__name__.rsplit(".", 1)[0] + ".kernels"
+    monkeypatch.setitem(
+        sys.modules,
+        kernels + ".post_moe_hop",
+        types.SimpleNamespace(
+            check_health=lambda where, **kwargs: calls.append(("check_health", where, kwargs))
+        ),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        kernels + ".host_page_retention",
+        types.SimpleNamespace(
+            install_host_policy=lambda where: calls.append(("install_host_policy", where))
+        ),
+    )
+    compat = types.SimpleNamespace(
+        declare_overlap_loop=lambda adp: calls.append(("declare_overlap_loop", adp)),
+        loop_domain=lambda: calls.append(("loop_domain",)),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        kernels + ".deferred_input_prep",
+        types.SimpleNamespace(compat=compat),
+    )
+    return calls
+
+
+def _events(paths, enable_attention_dp=False):
+    plan = types.SimpleNamespace(enabled=lambda name: name in paths)
+    mapping = types.SimpleNamespace(enable_attention_dp=enable_attention_dp)
+    return runtime.GlmExecutorEvents(plan, mapping)
+
+
+def test_the_executor_events_act_for_the_fast_paths_of_the_plan(event_calls):
+    events = _events({"post_moe_hop", "host_page_retention", "sampler_store"})
+
+    events.after_warmup()
+    events.before_overlap_loop()
+    events.at_shutdown()
+    assert event_calls == [
+        ("check_health", "after_warmup", {"collective": True}),
+        ("install_host_policy", "after_warmup"),
+        ("declare_overlap_loop", False),
+        ("loop_domain",),
+        ("check_health", "end_of_serving", {"raise_on_error": False}),
+    ]
+
+
+@pytest.mark.parametrize("path", ["sampler_store", "execution_stream_loop"])
+def test_the_overlap_loop_declares_its_domain_for_either_part(event_calls, path):
+    _events({path}, enable_attention_dp=True).before_overlap_loop()
+    assert event_calls == [("declare_overlap_loop", True), ("loop_domain",)]
+
+
+def test_the_executor_events_do_nothing_without_their_fast_paths(event_calls):
+    events = _events(set())
+
+    events.after_warmup()
+    events.before_overlap_loop()
+    events.at_shutdown()
+    assert event_calls == []
+
+
+def test_the_stager_attaches_the_attribute_every_engine_call_site_reads():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        PrepStager,
+    )
+    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+    from tensorrt_llm._torch.speculative.interface import SpecMetadata
+
+    for cls in (DecoderRunner, TrtllmAttentionMetadata, SpecMetadata, KVCacheManagerV2):
+        assert PrepStager._ATTACH_ATTR in vars(cls), cls
+        assert getattr(cls, PrepStager._ATTACH_ATTR) is None, cls
+
+
+class _ModelConfig:
+    """What the shell and its hooks read off the model config."""
+
+    class mapping:
+        world_size = 8
+        enable_attention_dp = False
+
+
+class _SparseMlaDecodeFmha:
+    """Stands in for the FMHA library the ``sparse_mla_decode`` kernel package registers."""
+
+
+SPARSE_MLA_LIB = "glm_moe_dsa_sparse_mla_decode"
+RANK_UNIFORM_ENV = dispatch.SHARED_ENV + (_router_index.MODELING_V2_VALIDATE_ENV,)
+
+
+@pytest.fixture
+def shell_construction(monkeypatch):
+    """The shell class with the built-in constructor stubbed, and the events of a construction.
+
+    The built-in constructor records the installed fast paths and the FMHA library list it
+    would build its attention layers from; ``adopt_model`` records the plan it attaches. The
+    selector state and the FMHA registry module are restored afterwards. The store writer the
+    hooks build is a host-only stand-in.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
+        deferred_input_prep,
+    )
+
+    monkeypatch.setattr(deferred_input_prep, "SamplerStoreWriter", _StoreWriter)
+    events = []
+
+    def adopt_model(model, plan):
+        events.append(("adopt", type(model.model).__name__, plan))
+
+    def builtin_init(self, model_config):
+        events.append("built-in model")
+        events.append(("installed", dispatch.installed()))
+        events.append(("fmha", list(fmha_registry.FMHA_LIBS)))
+        nn.Module.__init__(self)
+        self.model_config = model_config
+        self.model = _builtin_text_model()
+        self.model.layers = nn.ModuleList([_Layer(), _Layer(post_moe_hop=_Hop())])
+
+    def allgather(value, world_size):
+        events.append(("allgather", value, world_size))
+        return [value] * world_size
+
+    shell_cls = modeling.ModelingV2Glm52Nvfp4Sm103Tp8
+    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__init__", builtin_init)
+    monkeypatch.setattr(modeling, "adopt_model", adopt_model)
+    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__post_init__", lambda self: None)
+    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__pp_init__", lambda self: None)
+    monkeypatch.setitem(
+        _router_index._RESOLVED_MODES, shell_cls, _router_index.ModelingV2Mode.REQUIRE
+    )
+    monkeypatch.setattr(dispatch, "_require", [False])
+    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
+    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: events.append(message))
+    monkeypatch.setattr(
+        dispatch, "_REGISTRY", {"sparse_mla_decode": {"fmha_lib": _SparseMlaDecodeFmha}}
+    )
+    for name in list(os.environ):
+        if name.startswith(dispatch.KNOB_PREFIX) or name in RANK_UNIFORM_ENV:
+            monkeypatch.delenv(name)
+    registry_namespace = dict(vars(fmha_registry))
+    yield shell_cls, events
+    vars(fmha_registry).update(registry_namespace)
+
+
+def test_the_shell_checks_the_ranks_resolves_the_plan_and_adopts_the_core(
+    monkeypatch, shell_construction
+):
+    """The knobs are checked across the ranks before any other construction step, and the
+    attached fast-path state after the model is built."""
+    shell_cls, events = shell_construction
+    monkeypatch.setenv(dispatch.DISABLE_ENV, "draft_tail")
+
+    shell = shell_cls(_ModelConfig())
+
+    allgather, record, built_in, installed, fmha, adopt, adopted = events
+    assert allgather == (
+        "allgather",
+        ((dispatch.DISABLE_ENV, "draft_tail"), (dispatch.MODE_KEY, "require")),
+        8,
+    )
+    # The layers are adopted after the core and before their state is checked.
+    assert adopt == ("adopt", "GlmMoeDsaTp8ModelingV2", shell.plan)
+    assert adopted == ("allgather", modeling.adopted_state(shell), 8)
+    assert record.startswith("[glm_moe_dsa] dispatch: record paths=")
+    assert built_in == "built-in model"
+    assert installed == ("installed", frozenset(dispatch.FAST_PATH_NAMES) - {"draft_tail"})
+    assert fmha[0] == "fmha"
+    assert type(shell.model) is modeling.GlmMoeDsaTp8ModelingV2
+    assert shell.plan.require and dispatch.require_mode()
+    assert not shell.plan.enabled("draft_tail") and shell.plan.enabled("routed_moe")
+    # The hooks are built on the engine's first read, not during construction.
+    assert "_runtime_hooks" not in vars(shell)
+    hooks = runtime_hooks_of(shell)
+    assert hooks.step_input_stager is not None
+    assert isinstance(hooks.sampler_store_writer, _StoreWriter)
+    assert hooks.loop_on_execution_stream
+    assert isinstance(hooks.executor_events, runtime.GlmExecutorEvents)
+    assert runtime_hooks_of(shell) is hooks
+
+
+def test_an_attribute_error_while_building_the_hooks_is_not_read_as_no_hooks(
+    monkeypatch, shell_construction
+):
+    shell_cls, _ = shell_construction
+    shell = shell_cls(_ModelConfig())
+
+    def broken(model):
+        raise AttributeError("plan")
+
+    monkeypatch.setattr(runtime, "make_runtime_hooks", broken)
+    with pytest.raises(RuntimeError, match="step-loop hooks"):
+        runtime_hooks_of(shell)
+
+
+def test_the_fast_paths_read_the_installed_plan(monkeypatch, shell_construction):
+    """The plan is installed before the built-in model is built; DISABLE turns paths off."""
+    shell_cls, events = shell_construction
+    monkeypatch.setenv(dispatch.DISABLE_ENV, "post_attention_hop")
+
+    shell = shell_cls(_ModelConfig())
+
+    off = {name for name in dispatch.FAST_PATH_NAMES if not dispatch.is_enabled(name)}
+    assert off == shell.plan.disabled | shell.plan.disabled_dependents
+    assert "post_attention_hop" in off and "hop_routing" in off
+    assert ("installed", frozenset(dispatch.FAST_PATH_NAMES) - off) in events
+    # A knob the target does not read, set after construction, changes nothing.
+    monkeypatch.setenv(dispatch.FEATURE_ENV, "post_attention_hop")
+    assert not dispatch.is_enabled("post_attention_hop")
+    record = next(e for e in events if isinstance(e, str) and "dispatch: record" in e)
+    assert "post_attention_hop" not in record.split(" side_selectors=")[0]
+
+
+def test_the_fmha_registry_stays_the_builtin_one(monkeypatch, shell_construction):
+    """sparse_mla_decode joins no process-wide FMHA library list: adoption declares it on the
+    GLM attention backends (``test_modeling_v2_glm_attention.py``)."""
+    shell_cls, events = shell_construction
+    builtin_libs = list(fmha_registry.FMHA_LIBS)
+    builtin_enabled = fmha_manager.get_enabled_fmha_lib_classes()
+
+    shell_cls(_ModelConfig())
+    monkeypatch.setenv(dispatch.DISABLE_ENV, "sparse_mla_decode")
+    shell_cls(_ModelConfig())
+
+    assert SPARSE_MLA_LIB not in builtin_libs
+    assert [e for e in events if isinstance(e, tuple) and e[0] == "fmha"] == [
+        ("fmha", builtin_libs)
+    ] * 2
+    assert list(fmha_registry.FMHA_LIBS) == builtin_libs
+    assert fmha_manager.get_enabled_fmha_lib_classes() == builtin_enabled
+
+
+@pytest.mark.parametrize(
+    "knob",
+    sorted(dispatch.REMOVED_KNOBS),
+    ids=lambda knob: knob.removeprefix(dispatch.KNOB_PREFIX),
+)
+def test_a_removed_knob_fails_before_the_builtin_model_is_built(
+    monkeypatch, shell_construction, knob
+):
+    shell_cls, events = shell_construction
+    monkeypatch.setenv(knob, "deferred_input_prep")
+
+    with pytest.raises(ValueError, match=f"this target does not read {knob} "):
+        shell_cls(_ModelConfig())
+    assert [event[0] for event in events] == ["allgather"]
+    assert dispatch.installed() is None
+
+
+def test_a_knob_value_no_rank_can_parse_fails_before_the_builtin_model_is_built(
+    monkeypatch, shell_construction
+):
+    shell_cls, events = shell_construction
+    monkeypatch.setenv(dispatch.DISABLE_ENV, "not_a_fast_path")
+
+    with pytest.raises(ValueError, match="not_a_fast_path"):
+        shell_cls(_ModelConfig())
+    assert [event[0] for event in events] == ["allgather"]
+
+
+def test_the_ranks_check_the_shared_switches_and_the_mode_too(monkeypatch, shell_construction):
+    shell_cls, events = shell_construction
+    monkeypatch.setenv("TRTLLM_ENABLE_PDL", "1")
+    monkeypatch.setenv(_router_index.MODELING_V2_VALIDATE_ENV, "1")
+
+    shell_cls(_ModelConfig())
+
+    assert events[0] == (
+        "allgather",
+        (
+            ("TRTLLM_ENABLE_PDL", "1"),
+            (_router_index.MODELING_V2_VALIDATE_ENV, "1"),
+            (dispatch.MODE_KEY, "require"),
+        ),
+        8,
+    )
+
+
+def test_ranks_that_disagree_fail_before_the_builtin_model_is_built(
+    monkeypatch, shell_construction
+):
+    shell_cls, events = shell_construction
+
+    def allgather(value, world_size):
+        events.append(("allgather", value, world_size))
+        return [value, value + (("TRTLLM_FORCE_MNNVL_AR", "1"),)]
+
+    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
+    with pytest.raises(RuntimeError, match="TRTLLM_FORCE_MNNVL_AR='1'"):
+        shell_cls(_ModelConfig())
+    assert [event[0] for event in events] == ["allgather"]
+    assert dispatch.installed() is None
+
+
+def test_the_adopted_state_names_each_attached_fast_path_object():
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[_Layer(), _Layer()]))
+    model.model.layers[1].post_moe_hop = _Hop()
+    model.model.layers[1].mlp.router_gemm = _Hop(engaged=False)
+    model.model.layers[1].self_attn._attn_tail = object()
+
+    assert modeling.adopted_state(model) == (
+        ("model.layers.0", ()),
+        (
+            "model.layers.1",
+            (
+                ("post_moe_hop", "_Hop", "engaged=True", "impl='cluster'"),
+                ("self_attn._attn_tail", "object"),
+                ("mlp.router_gemm", "_Hop", "engaged=False", "impl='cluster'"),
+            ),
+        ),
+    )
+
+
+def test_ranks_that_attached_different_state_fail_after_the_model_is_built(
+    monkeypatch, shell_construction
+):
+    shell_cls, events = shell_construction
+    hooks = []
+    monkeypatch.setattr(runtime, "make_runtime_hooks", lambda model: hooks.append(model))
+
+    def allgather(value, world_size):
+        events.append(("allgather", value, world_size))
+        if value and value[0][0] == "model.layers.0":
+            # A peer whose post-MoE hop on layer 1 did not engage.
+            return [value, (value[0], ("model.layers.1", ()))]
+        return [value] * world_size
+
+    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
+    with pytest.raises(RuntimeError, match=r"rank 1: \('model.layers.1', \(\)\)") as excinfo:
+        shell_cls(_ModelConfig())
+    assert "rank 0: ('model.layers.1', (('post_moe_hop', '_Hop'" in str(excinfo.value)
+    assert "built-in model" in events and hooks == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py
new file mode 100644
index 000000000000..5642cf41d4bb
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py
@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""One name per fast path of the GLM-5.2 tp8 target, used the same way everywhere.
+
+The fast-path names of ``dispatch.FAST_PATHS`` are the ``DISABLE`` tokens, the kernel package
+directories (or a directory declared in ``KERNEL_DIRECTORY``), the log tags, the knob infixes,
+the kernel registry keys and the kernel test directories. Read off disk: no kernel module is
+imported.
+"""
+
+from __future__ import annotations
+
+import ast
+import re
+from pathlib import Path
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+# The family's code: its shared package and every target of it, each with its kernel packages.
+_TARGET = Path(dispatch.__file__).resolve().parents[1]
+_KERNEL_ROOTS = sorted(_TARGET.glob("*/kernels"))
+_TESTS = Path(__file__).resolve().parents[1]
+_PATHS = frozenset(dispatch.FAST_PATH_NAMES)
+
+_LOG_TAG = "[glm_moe_dsa] "
+# After the tag: a fast-path name and a separator, or the dispatcher's own lines.
+_TAGGED = re.compile(r"(?P<name>[a-z0-9_]+)(?P<sep>[/ :])")
+_KNOB = re.compile(re.escape(dispatch.KNOB_PREFIX) + r"(?P<infix>[A-Z0-9_]*)")
+# Knobs that select among the fast paths rather than tune one.
+_SELECTOR_KNOBS = ("FEATURES", "DISABLE")
+
+
+def _sources(root: Path):
+    return sorted(p for p in root.rglob("*.py") if "__pycache__" not in p.parts)
+
+
+def _kernel_directories() -> set:
+    """The kernel packages of the family, the shared ones and every target's own."""
+    return {
+        d.name
+        for kernels in _KERNEL_ROOTS
+        for d in kernels.iterdir()
+        if d.is_dir() and (d / "__init__.py").is_file()
+    }
+
+
+def _string_constants(tree: ast.AST):
+    """Every ``str`` constant, the constant parts of f-strings included."""
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Constant) and isinstance(node.value, str):
+            yield node
+
+
+def test_each_fast_path_has_one_name():
+    assert len(dispatch.FAST_PATH_NAMES) == 36
+    assert len(_PATHS) == 36
+
+
+def test_the_kernel_directories_are_the_fast_paths_without_a_declared_directory():
+    directories = _kernel_directories()
+    assert len(directories) == 28
+    assert directories == _PATHS - set(dispatch.KERNEL_DIRECTORY)
+    assert set(dispatch.KERNEL_DIRECTORY) <= _PATHS
+    assert set(dispatch.KERNEL_DIRECTORY.values()) <= directories
+
+
+def test_every_log_tag_names_a_fast_path_or_the_dispatcher():
+    offences = []
+    for path in _sources(_TARGET):
+        for node in _string_constants(ast.parse(path.read_text(encoding="utf-8"))):
+            text = node.value
+            start = text.find(_LOG_TAG)
+            while start >= 0:
+                match = _TAGGED.match(text, start + len(_LOG_TAG))
+                name = match.group("name") if match else None
+                if not (
+                    match and (name in _PATHS or (name == "dispatch" and match.group("sep") == ":"))
+                ):
+                    offences.append(f"{path.relative_to(_TARGET)}:{node.lineno}: {text[start:]!r}")
+                start = text.find(_LOG_TAG, start + 1)
+    assert not offences, "log lines must name a fast path after the tag:\n" + "\n".join(offences)
+
+
+def _knob_owner(infix: str):
+    """The fast path whose name starts the knob's infix (the longest such name), if any."""
+    owners = [p for p in _PATHS if infix.startswith(p.upper() + "_")]
+    return max(owners, key=len) if owners else None
+
+
+def test_every_knob_name_starts_with_its_fast_path():
+    offences = []
+    for path in _sources(_TARGET):
+        text = path.read_text(encoding="utf-8")
+        for match in _KNOB.finditer(text):
+            infix = match.group("infix")
+            if infix and infix not in _SELECTOR_KNOBS and _knob_owner(infix) is None:
+                line = text.count("\n", 0, match.start()) + 1
+                offences.append(f"{path.relative_to(_TARGET)}:{line}: {match.group(0)}")
+    assert not offences, "knob names must start with a fast-path name:\n" + "\n".join(offences)
+
+
+def _imports_register_from_dispatch(tree: ast.AST) -> bool:
+    return any(
+        isinstance(node, ast.ImportFrom)
+        and (node.module or "").split(".")[-1] == "dispatch"
+        and any(alias.name == "register" and alias.asname is None for alias in node.names)
+        for node in ast.walk(tree)
+    )
+
+
+def _module_strings(path: Path, tree: ast.Module) -> dict:
+    """Module-level ``NAME = "<str>"`` bindings, and those imported from a sibling module."""
+    out = {}
+    for stmt in tree.body:
+        if (
+            isinstance(stmt, ast.Assign)
+            and isinstance(stmt.value, ast.Constant)
+            and isinstance(stmt.value.value, str)
+        ):
+            out.update({t.id: stmt.value.value for t in stmt.targets if isinstance(t, ast.Name)})
+        elif isinstance(stmt, ast.ImportFrom) and stmt.level == 1 and stmt.module:
+            sibling = path.parent / f"{stmt.module}.py"
+            if sibling.is_file():
+                strings = _module_strings(sibling, ast.parse(sibling.read_text(encoding="utf-8")))
+                out.update(
+                    {
+                        alias.asname or alias.name: strings[alias.name]
+                        for alias in stmt.names
+                        if alias.name in strings
+                    }
+                )
+    return out
+
+
+def _registered_paths(path: Path):
+    """``(line, name)`` of every ``dispatch.register(<name>, …)`` call of one kernel module."""
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    bare = _imports_register_from_dispatch(tree)
+    names = _module_strings(path, tree)
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call) or not node.args:
+            continue
+        func = node.func
+        if not (
+            (isinstance(func, ast.Name) and func.id == "register" and bare)
+            or (
+                isinstance(func, ast.Attribute)
+                and func.attr == "register"
+                and isinstance(func.value, ast.Name)
+                and func.value.id == "dispatch"
+            )
+        ):
+            continue
+        first = node.args[0]
+        if isinstance(first, ast.Constant) and isinstance(first.value, str):
+            yield node.lineno, first.value
+        elif isinstance(first, ast.Name) and first.id in names:
+            yield node.lineno, names[first.id]
+        else:
+            yield node.lineno, ast.unparse(first)
+
+
+def test_every_registry_entry_is_keyed_by_the_fast_path_of_its_directory():
+    offences, registering = [], set()
+    for kernels in _KERNEL_ROOTS:
+        for path in _sources(kernels):
+            relative = path.relative_to(kernels)
+            if len(relative.parts) < 2:
+                continue
+            directory = relative.parts[0]
+            for line, name in _registered_paths(path):
+                if name not in _PATHS or dispatch.KERNEL_DIRECTORY.get(name, name) != directory:
+                    offences.append(
+                        f"{kernels.parent.name}/kernels/{relative}:{line}: register({name!r}, ...)"
+                    )
+                else:
+                    registering.add(directory)
+    assert not offences, (
+        "get_impl imports kernels.<directory of the path>, so each package registers its "
+        "entry points under its own path:\n" + "\n".join(offences)
+    )
+    assert registering == _kernel_directories()
+
+
+def test_every_kernel_test_directory_is_a_kernel_directory():
+    """A target's tests of how it composes a shared package live under the package's name."""
+    for kernels in _KERNEL_ROOTS:
+        test_root = _TESTS / kernels.parent.name / "kernels"
+        if not test_root.is_dir():
+            continue
+        tests = {d.name for d in test_root.iterdir() if d.is_dir() and d.name != "__pycache__"}
+        assert tests <= _kernel_directories(), sorted(tests - _kernel_directories())
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py
new file mode 100644
index 000000000000..c63a5c287a14
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py
@@ -0,0 +1,223 @@
+# 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.
+"""Real-weight native fixtures for the absorbed_query/projection package.
+
+``build_native_attention`` constructs the built-in ``DeepseekV32Attention``
+module that the GLM-5.2 decode engine uses for one target layer at TP8 rank
+``tp_rank``, adopts it as the target's ``GlmMLA`` right after construction (as the
+target's model does; on a decode step it runs the selected fast paths, and with none
+selected or on any other step it computes the built-in attention) and places the
+checkpoint tensors exactly as ``DeepseekV3WeightLoader`` does for the BF16 attention
+weights:
+
+* ``kv_a_proj_with_mqa.weight = cat([q_a_proj, kv_a_proj_with_mqa])``
+* ``q_b_proj`` through ``Linear.load_weights`` (column-parallel slicing)
+* ``k_b_proj_trans`` / ``v_b_proj`` / ``kv_b_proj`` from the TP slice of
+  ``kv_b_proj`` (heads ``[r*8, r*8+8)``)
+* indexer ``wq_b``/``wk``/``weights_proj``/``k_norm`` and
+  ``Indexer.post_load_weights`` (fused ``wk``+``weights_proj``).
+
+The attention linears of this checkpoint are excluded from NVFP4, so the
+module is built with a KV-cache-only ``QuantConfig`` (FP8 KV), which is the
+configuration ``apply_quant_config_exclude_modules`` leaves on them.
+
+Fixture labels: weights are **real** (checkpoint tensors); activations are
+**synthetic** unless a caller supplies captured ones.
+"""
+
+import json
+import os
+from typing import Dict, List, Optional
+
+import torch
+from utils.llm_data import llm_models_root
+
+CKPT_DIR = os.environ.get(
+    "GLM52_CKPT_DIR",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+TP_SIZE = 8
+
+# Full (own indexer) and shared (reused top-k) target layers of GLM-5.2.
+FULL_LAYER = 6
+SHARED_LAYER = 7
+
+
+# Feature names of the neighbour packages every arm of a comparison runs on
+# top of, the native reference included (e.g. "rope_kv_append" on an
+# integration image that contains rope_kv_append).
+BASE_FEATURES_ENV = "GLM_MOE_DSA_TEST_BASE_FEATURES"
+
+
+def with_base_features(features: str) -> str:
+    """Return the fast paths ``features`` as a comma-separated list.
+
+    The fast paths of ``GLM_MOE_DSA_TEST_BASE_FEATURES`` follow them.
+    """
+    names = [f.strip() for f in features.split(",") if f.strip()]
+    base = [f.strip() for f in os.environ.get(BASE_FEATURES_ENV, "").split(",") if f.strip()]
+    return ",".join(names + [b for b in base if b not in names])
+
+
+def checkpoint_available() -> bool:
+    return os.path.isfile(os.path.join(CKPT_DIR, "model.safetensors.index.json"))
+
+
+def load_layer_tensors(
+    layer_idx: int, names: Optional[List[str]] = None
+) -> Dict[str, torch.Tensor]:
+    """Load ``model.layers.<layer_idx>.self_attn.*`` tensors (CPU)."""
+    from safetensors import safe_open
+
+    with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
+        weight_map = json.load(f)["weight_map"]
+    prefix = f"model.layers.{layer_idx}.self_attn."
+    keys = [k for k in weight_map if k.startswith(prefix)]
+    if names is not None:
+        keys = [prefix + n for n in names if prefix + n in weight_map]
+    by_shard: Dict[str, List[str]] = {}
+    for k in keys:
+        by_shard.setdefault(weight_map[k], []).append(k)
+    out = {}
+    for shard, shard_keys in by_shard.items():
+        with safe_open(os.path.join(CKPT_DIR, shard), framework="pt") as f:
+            for k in shard_keys:
+                out[k[len(prefix) :]] = f.get_tensor(k)
+    return out
+
+
+def build_model_config(tp_rank: int = 0):
+    """Native ``ModelConfig`` for GLM-5.2 at TP8 with the served DSA flags."""
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, MTPDecodingConfig
+    from tensorrt_llm.mapping import Mapping
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+    from tensorrt_llm.quantization.mode import QuantAlgo
+
+    mapping = Mapping(world_size=TP_SIZE, tp_size=TP_SIZE, rank=tp_rank, gpus_per_node=4)
+    sparse = DeepSeekSparseAttentionConfig(
+        enable_heuristic_topk=True, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=True
+    )
+    model_config = ModelConfig.from_pretrained(
+        CKPT_DIR,
+        trust_remote_code=True,
+        mapping=mapping,
+        sparse_attention_config=sparse,
+        spec_config=MTPDecodingConfig(max_draft_len=5),
+        attn_backend="TRTLLM",
+    )
+    model_config._frozen = False
+    # Attention linears are excluded from NVFP4 in this checkpoint; the
+    # KV cache stays FP8 as in the served configuration.
+    model_config.quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8)
+    if model_config.pretrained_config.model_type == "glm_moe_dsa":
+        model_config.pretrained_config.model_type = "deepseek_v32"
+    model_config._frozen = True
+    return model_config
+
+
+def build_native_attention(
+    layer_idx: int,
+    tp_rank: int = 0,
+    model_config=None,
+    tensors: Optional[Dict[str, torch.Tensor]] = None,
+    aux_stream: Optional[torch.cuda.Stream] = None,
+):
+    """Build the built-in ``DeepseekV32Attention``, adopted as ``GlmMLA``, with real weights on GPU."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+        GlmMLA,
+    )
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+    from tensorrt_llm._torch.utils import AuxStreamType
+
+    if model_config is None:
+        model_config = build_model_config(tp_rank)
+    if tensors is None:
+        tensors = load_layer_tensors(layer_idx)
+    if aux_stream is None:
+        aux_stream = torch.cuda.Stream()
+    attn = DeepseekV32Attention(
+        model_config, layer_idx=layer_idx, aux_stream_dict={AuxStreamType.Attention: aux_stream}
+    )
+    GlmMLA.adopt(attn)
+    # The model builds the FMHA libraries in ``create_weights`` after adoption; a module built
+    # alone rebuilds them for the library adoption declared.
+    for backend in (attn.mha, attn.mqa):
+        if backend is not None and backend.prepended_fmha_lib_classes:
+            backend.update_quant_config(backend.quant_config)
+    attn = attn.cuda()
+
+    cfg = model_config.pretrained_config
+    num_heads = cfg.num_attention_heads
+    nope, vdim = cfg.qk_nope_head_dim, cfg.v_head_dim
+    local_heads = num_heads // TP_SIZE
+
+    with torch.no_grad():
+        fused_a = torch.cat(
+            [tensors["q_a_proj.weight"], tensors["kv_a_proj_with_mqa.weight"]], dim=0
+        )
+        attn.kv_a_proj_with_mqa.weight.data.copy_(fused_a)
+        attn.q_a_layernorm.weight.data.copy_(tensors["q_a_layernorm.weight"])
+        attn.kv_a_layernorm.weight.data.copy_(tensors["kv_a_layernorm.weight"])
+        attn.q_b_proj.load_weights(weights=[{"weight": tensors["q_b_proj.weight"]}])
+
+        kv_b = tensors["kv_b_proj.weight"].unflatten(0, [num_heads, nope + vdim])
+        kv_b = torch.chunk(kv_b, TP_SIZE, dim=0)[tp_rank].contiguous()
+        k_nope, v_w = kv_b.split([nope, vdim], dim=1)
+        attn.k_b_proj_trans.data.copy_(k_nope.transpose(2, 1).contiguous())
+        attn.kv_b_proj.weight.data.copy_(
+            torch.cat(
+                [k_nope.reshape(local_heads * nope, -1), v_w.reshape(local_heads * vdim, -1)], dim=0
+            )
+        )
+        attn.v_b_proj = torch.nn.Parameter(v_w.contiguous().cuda(), requires_grad=False)
+
+        indexer = attn.mqa.indexer
+        if indexer is not None:
+            indexer.wq_b.load_weights(weights=[{"weight": tensors["indexer.wq_b.weight"]}])
+            indexer.wk.load_weights(weights=[{"weight": tensors["indexer.wk.weight"]}])
+            indexer.weights_proj.load_weights(
+                weights=[{"weight": tensors["indexer.weights_proj.weight"]}]
+            )
+            indexer.k_norm.weight.data.copy_(tensors["indexer.k_norm.weight"])
+            indexer.k_norm.bias.data.copy_(tensors["indexer.k_norm.bias"])
+            indexer.post_load_weights()
+    attn.post_load_weights()
+    return attn
+
+
+def make_hidden(
+    num_rows: int, hidden_size: int = 6144, seed: int = 0, gamma: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """Synthetic post-RMSNorm hidden state ``[R, hidden]`` (BF16).
+
+    Unit-RMS Gaussian rows scaled by a real ``input_layernorm`` gain when one
+    is given, so activation magnitudes match what the layer sees.
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(num_rows, hidden_size, generator=g, device="cuda", dtype=torch.float32)
+    if gamma is not None:
+        x = x * gamma.float().cuda()
+    return x.to(torch.bfloat16)
+
+
+def load_input_layernorm(layer_idx: int) -> torch.Tensor:
+    from safetensors import safe_open
+
+    with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
+        weight_map = json.load(f)["weight_map"]
+    key = f"model.layers.{layer_idx}.input_layernorm.weight"
+    with safe_open(os.path.join(CKPT_DIR, weight_map[key]), framework="pt") as f:
+        return f.get_tensor(key)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py
new file mode 100644
index 000000000000..24857f178128
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py
@@ -0,0 +1,22 @@
+# 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.
+"""Make the absorbed_query fixture helpers importable under ``--import-mode=importlib``."""
+
+import sys
+from pathlib import Path
+
+_HERE = str(Path(__file__).resolve().parent)
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py
new file mode 100644
index 000000000000..b65f038f957a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py
@@ -0,0 +1,870 @@
+# 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.
+"""Stock producer -> absorbed_query -> native DSA attention, through the target's attention.
+
+The attention modules are the target's adopted ones (``GlmMLA``, built by
+``absorbed_query_fixture.build_native_attention``); every generation step runs as a decode
+step (``dispatch.Phase.DECODE``) and the prefill as a prefill step. Real
+GLM-5.2 weights of the full-indexer layer 6 and the shared layer 7 at TP8
+rank 0; synthetic activations; FP8 KV cache (V1 ``DSACacheManager``).
+
+Six requests with a 2,600-token context are natively prefilled through both
+layers (the indexer then selects 2,048 of 2,600 keys). A generation step
+then presents six query rows in one of two layouts: six requests x one token,
+or one request x six tokens (the served MTP verify layout, K = 5).
+
+Both call paths are exercised: ``forward_impl`` (``forward_sparse_attn``)
+and ``custom_op`` (``MLA.forward``'s ``_forward_custom_op`` with the
+registered metadata, which the model uses). Every step, eager or captured,
+first calls ``on_update_kv_lens()`` on the metadata, as the decoder
+runner's ``_forward_step`` does through ``_preprocess_inputs``.
+
+Checks against the same step with absorbed_query off (all asserted unless noted):
+
+* only the new-token slots of every cache (latent caches of both layers,
+  indexer-key cache of layer 6) change, and every byte equals native;
+* the MLA query after native RoPE/quantization is bit-identical: the
+  absorbed half ``fused_q[..., :512]`` and the FP8 ``quant_q_buffer``
+  handed to the attention kernel, and the BMM scales written next to them
+  (recorded by wrapping the native ``mla_rope_generation``; in the FP8-KV
+  path the RoPE half of ``fused_q`` is not written, the rotated q-rope goes
+  only to ``quant_q_buffer``);
+* the indexer inputs (``q_fp8``, ``weights``) are bit-identical to native,
+  and the indexer's logical top-k (layer 6) and the converted sparse indices
+  of both layers select the same key sets (rows compared after sorting;
+  exact order and, on a failure, the per-row overlap reported);
+* the attention outputs: the native sparse attention is not bitwise
+  run-to-run deterministic, so the outputs are bounded by its measured
+  repeat difference over ``NATIVE_RUNS`` native runs of the same step:
+  ``max|absorbed_query - native| <= 2 * max(native repeat max diff, one BF16 ulp of
+  max|native|)`` and ``relRMS(absorbed_query - native) <= max(2 * native repeat relRMS,
+  2**-10)``. This is a secondary check: the negative control shows that a
+  zeroed absorbed query moves the shared layer's outputs by only a few
+  times that noise, while the bit-identity checks of the attention inputs
+  above flag it directly;
+* absorbed_query itself is deterministic: a repeated absorbed_query step, 20 more absorbed_query
+  steps and every CUDA-graph replay reproduce the caches and recorded attention inputs of
+  the eager absorbed_query step bit for bit.
+
+The graph test captures the adopted layer forward (absorbed_query, ``_forward_dsa_attn``
+and the native cache append) for CUDA-graph generation metadata, replays it
+once and 20 times with fresh hidden states, and compares each replay with an
+eager native step and with a captured native graph.
+"""
+
+import inspect
+import math
+import os
+import sys
+
+import absorbed_query_fixture
+import pytest
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+
+import plan_testlib  # noqa: E402
+
+CTX = 2600
+REQS = 6
+LAYERS = (absorbed_query_fixture.FULL_LAYER, absorbed_query_fixture.SHARED_LAYER)
+# layout name -> (requests, query tokens per request)
+LAYOUTS = {"6req_x1tok": (6, 1), "1req_x6tok": (1, 6)}
+REL_RMS_FLOOR = 2.0**-10
+
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+
+def _set_absorbed_query(enabled: bool, form: str = "preabsorbed"):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    query_proj.SERVED_FORM = form
+    plan_testlib.select(
+        dispatch, absorbed_query_fixture.with_base_features("absorbed_query" if enabled else "")
+    )
+
+
+class _Recorder:
+    """Records, per layer, what the native attention consumes.
+
+    Wraps the native ``mla_rope_generation`` (post-RoPE ``fused_q``, FP8
+    ``quant_q_buffer``, BMM scales), the indexer's ``sparse_attn_indexer``
+    (logical top-k) and ``sparse_attn_predict`` (converted indices). Inside
+    a CUDA-graph capture the clones are graph-pool tensors that every replay
+    refreshes. When the installed package registers rope_kv_append's fused replacement
+    of ``mla_rope_generation`` (integration images), that entry point is
+    wrapped as well and records the same tensors.
+    """
+
+    def __init__(self, attn_by_layer):
+        self.attn = attn_by_layer
+        self.data = {layer: {} for layer in attn_by_layer}
+        self._installed = []
+        self._rope_append = None
+
+    @staticmethod
+    def _rope_record(rec, kv_lora, a):
+        # The FP8-KV path writes the rotated q-rope only into
+        # quant_q_buffer; fused_q[..., kv_lora:] stays unwritten.
+        rec["absorbed_q"] = a["fused_q"][..., :kv_lora].clone()
+        for key, arg in (
+            ("quant_q", "quant_q_buffer"),
+            ("bmm1_scale", "mla_bmm1_scale"),
+            ("bmm2_scale", "mla_bmm2_scale"),
+        ):
+            if a.get(arg) is not None:
+                rec[key] = a[arg].clone()
+
+    def _wrap_rope_append_rope(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+        try:
+            impl = dispatch.get_impl("rope_kv_append", "mla_rope_generation")
+        except (ImportError, KeyError, ValueError):
+            return
+        sig = inspect.signature(impl)
+        by_mqa = {
+            id(attn.mqa): (self.data[layer], attn.kv_lora_rank) for layer, attn in self.attn.items()
+        }
+
+        def hook(*args, **kwargs):
+            out = impl(*args, **kwargs)
+            a = sig.bind(*args, **kwargs).arguments
+            hit = by_mqa.get(id(a["mqa"]))
+            if hit is not None:
+                self._rope_record(hit[0], hit[1], a)
+                # Marks the call as served by rope_kv_append (compared like any record).
+                hit[0]["rope_append_rope"] = torch.ones(
+                    1, dtype=torch.int32, device=a["fused_q"].device
+                )
+            return out
+
+        dispatch._REGISTRY["rope_kv_append"]["mla_rope_generation"] = hook
+        self._rope_append = (dispatch, impl)
+
+    def _wrap(self, obj, name, post):
+        orig = getattr(obj, name)
+        sig = inspect.signature(orig)
+
+        def hook(*args, **kwargs):
+            out = orig(*args, **kwargs)
+            post(sig.bind(*args, **kwargs).arguments, out)
+            return out
+
+        setattr(obj, name, hook)
+        self._installed.append((obj, name))
+
+    def __enter__(self):
+        for layer, attn in self.attn.items():
+            rec = self.data[layer]
+
+            def rope_post(a, _out, rec=rec, kv_lora=attn.kv_lora_rank):
+                self._rope_record(rec, kv_lora, a)
+
+            def predict_post(_a, out, rec=rec):
+                if isinstance(out, tuple) and out[0] is not None:
+                    rec["sparse_indices"] = out[0].clone()
+
+            self._wrap(attn.mqa, "mla_rope_generation", rope_post)
+            self._wrap(attn.mqa, "sparse_attn_predict", predict_post)
+            if attn.mqa.indexer is not None:
+
+                def topk_post(a, out, rec=rec):
+                    rec["topk"] = out.clone()
+                    rec["indexer_q_fp8"] = a["q_fp8"].clone()
+                    rec["indexer_weights"] = a["weights"].clone()
+
+                self._wrap(attn.mqa.indexer, "sparse_attn_indexer", topk_post)
+        self._wrap_rope_append_rope()
+        return self
+
+    def __exit__(self, *exc):
+        for obj, name in self._installed:
+            delattr(obj, name)
+        self._installed.clear()
+        if self._rope_append is not None:
+            dispatch, impl = self._rope_append
+            dispatch._REGISTRY["rope_kv_append"]["mla_rope_generation"] = impl
+            self._rope_append = None
+        return False
+
+    def snapshot(self):
+        return {layer: {k: v.clone() for k, v in rec.items()} for layer, rec in self.data.items()}
+
+
+class _Harness:
+    def __init__(self):
+        import tensorrt_llm.bindings
+        from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
+        from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+        from tensorrt_llm._torch.metadata import KVCacheParams
+        from tensorrt_llm._torch.model_config import ModelConfig
+        from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
+        from tensorrt_llm.llmapi.llm_args import KvCacheConfig
+        from tensorrt_llm.mapping import Mapping
+        from tensorrt_llm.models.modeling_utils import QuantConfig
+        from tensorrt_llm.quantization.mode import QuantAlgo
+
+        self.KVCacheParams = KVCacheParams
+        self.model_config = absorbed_query_fixture.build_model_config(0)
+        pretrained = self.model_config.pretrained_config
+        sparse = self.model_config.sparse_attention_config
+        self.attn = {
+            layer: absorbed_query_fixture.build_native_attention(layer, 0, self.model_config)
+            for layer in LAYERS
+        }
+        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        kv_model_config = ModelConfig(
+            mapping=self.mapping,
+            sparse_attention_config=sparse,
+            pretrained_config=pretrained,
+            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8),
+        )
+        self.kvm = DSACacheManager(
+            KvCacheConfig(
+                dtype="fp8",
+                max_tokens=REQS * (CTX + 128),
+                enable_block_reuse=False,
+                event_buffer_max_size=0,
+            ),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+            num_layers=max(LAYERS) + 1,
+            num_kv_heads=1,
+            head_dim=576,
+            tokens_per_block=64,
+            max_seq_len=CTX + 128,
+            max_batch_size=REQS,
+            mapping=self.mapping,
+            dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
+            sparse_attention_config=sparse,
+            model_config=kv_model_config,
+            pretrained_config=pretrained,
+            vocab_size=pretrained.vocab_size,
+            max_num_tokens=REQS * CTX,
+        )
+        self.req_ids = list(range(REQS))
+        for rid in self.req_ids:
+            # CTX + 6 new tokens fit the blocks of CTX + 1 tokens (64/block).
+            self.kvm.add_dummy_requests(
+                request_ids=[rid], token_nums=[CTX + 1], is_gen=False, prepare_resource=True
+            )
+        sparse_params = sparse.to_sparse_params(pretrained_config=pretrained)
+        self.Metadata = get_attention_backend("TRTLLM", sparse_params=sparse_params).Metadata
+        self.meta_params = sparse.to_sparse_metadata_params(pretrained_config=pretrained)
+        self.gamma = absorbed_query_fixture.load_input_layernorm(LAYERS[0])
+
+        # Native prefill through both layers fills the latent and indexer
+        # caches the way the stock producers do.
+        _set_absorbed_query(False)
+        hidden = absorbed_query_fixture.make_hidden(REQS * CTX, seed=11, gamma=self.gamma)
+        positions = torch.arange(CTX, dtype=torch.int32, device="cuda").repeat(REQS).unsqueeze(0)
+        ctx_meta = self.Metadata(
+            seq_lens=torch.tensor([CTX] * REQS, dtype=torch.int),
+            request_ids=self.req_ids,
+            max_num_requests=REQS,
+            num_contexts=REQS,
+            prompt_lens=[CTX] * REQS,
+            max_num_tokens=REQS * CTX,
+            kv_cache_manager=self.kvm,
+            kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0] * REQS),
+            mapping=self.mapping,
+            sparse_metadata_params=self.meta_params,
+        )
+        ctx_meta.prepare()
+        self.step(ctx_meta, positions, hidden, "forward_impl")
+        torch.cuda.synchronize()
+
+    # -- generation inputs -------------------------------------------------
+
+    def gen_inputs(self, layout, seed):
+        nreq, ntok = LAYOUTS[layout]
+        hidden = absorbed_query_fixture.make_hidden(nreq * ntok, seed=seed, gamma=self.gamma)
+        pos = (CTX + torch.arange(ntok, dtype=torch.int32, device="cuda")).repeat(nreq).unsqueeze(0)
+        return hidden, pos
+
+    def gen_metadata(self, layout, cuda_graph=False):
+        """Generation metadata as the runner builds it: graph metadata from
+        ``create_cuda_graph_metadata`` with a buffer registry (reserved,
+        graph-stable buffers, as the graph runner passes its own), MTP verify
+        width from ``update_spec_dec_param``, then ``prepare()``."""
+        nreq, ntok = LAYOUTS[layout]
+        req_ids = self.req_ids[:nreq]
+        kv_params = self.KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[CTX] * nreq)
+        meta = self.Metadata(
+            seq_lens=torch.tensor([ntok] * nreq, dtype=torch.int),
+            request_ids=req_ids,
+            max_num_requests=nreq,
+            num_contexts=0,
+            prompt_lens=[CTX] * nreq,
+            max_num_tokens=nreq * ntok,
+            kv_cache_manager=self.kvm,
+            kv_cache_params=kv_params,
+            mapping=self.mapping,
+            sparse_metadata_params=self.meta_params,
+        )
+        if cuda_graph:
+            from tensorrt_llm._torch.memory_buffer_utils import Buffers
+
+            meta = meta.create_cuda_graph_metadata(nreq, False, ntok - 1, Buffers())
+            assert meta.is_cuda_graph
+            meta.request_ids = req_ids
+            meta.prompt_lens = [CTX] * nreq
+            meta.num_contexts = 0
+            meta.kv_cache_params = kv_params
+        if ntok > 1:
+            meta.update_spec_dec_param(
+                batch_size=nreq,
+                is_spec_decoding_enabled=True,
+                is_spec_dec_tree=False,
+                is_spec_dec_dynamic_tree=False,
+                max_draft_len=ntok - 1,
+                max_total_draft_tokens=ntok - 1,
+            )
+        meta.prepare()
+        return meta
+
+    # -- forwards ----------------------------------------------------------
+
+    def forward(self, layer, positions, hidden, meta, path):
+        attn = self.attn[layer]
+        if path == "forward_impl":
+            out = attn._create_outputs(hidden, meta)
+            attn.forward_impl(positions, hidden, meta, attn_output=out)
+            return out[0]
+        # The model's path: MLA.forward with register_to_config goes through
+        # _forward_custom_op and the registered metadata.
+        import weakref
+
+        from tensorrt_llm._torch.utils import model_extra_attrs
+
+        extra = self.model_config.extra_attrs
+        extra["attention_metadata"] = weakref.ref(meta)
+        with model_extra_attrs(extra):
+            out = [torch.ops.trtllm.create_mla_outputs(hidden, attn.layer_idx_str)]
+            attn._forward_custom_op(hidden, positions, out, None)
+        return out[0]
+
+    def step(self, meta, positions, hidden, path):
+        """One forward step of both layers, as the decoder runner's
+        ``_forward_step`` runs it (eager or captured): ``_preprocess_inputs``
+        first refreshes the metadata derived from the KV lengths
+        (``on_update_kv_lens``), then the model forward, under the step's phase (a
+        generation-only step is a decode step)."""
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+            dispatch,
+        )
+
+        phase = dispatch.Phase.DECODE if meta.num_contexts == 0 else dispatch.Phase.PREFILL
+        meta.on_update_kv_lens()
+        with dispatch.step(phase):
+            return {layer: self.forward(layer, positions, hidden, meta, path) for layer in LAYERS}
+
+    def gen_step(self, layout, path, hidden, pos, meta=None):
+        meta = meta if meta is not None else self.gen_metadata(layout)
+        outs = {layer: out.clone() for layer, out in self.step(meta, pos, hidden, path).items()}
+        torch.cuda.synchronize()
+        return outs
+
+    def recorded_step(self, layout, path, hidden, pos):
+        with _Recorder(self.attn) as rec:
+            out = self.gen_step(layout, path, hidden, pos)
+        return out, rec.snapshot()
+
+    # -- caches ------------------------------------------------------------
+
+    def caches(self):
+        bufs = {f"latent{layer}": self.kvm.get_buffers(layer) for layer in LAYERS}
+        bufs[f"index_k{LAYERS[0]}"] = self.kvm.get_indexer_k_cache_buffers(LAYERS[0])
+        return bufs
+
+    def snapshot(self):
+        return {k: v.clone() for k, v in self.caches().items()}
+
+    def restore(self, snap):
+        for k, v in self.caches().items():
+            v.copy_(snap[k])
+        torch.cuda.synchronize()
+
+    def reset_absorbed_query_runners(self):
+        for layer in LAYERS:
+            self.attn[layer]._absorbed_query = None
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = _Harness()
+    yield h
+    _set_absorbed_query(False)
+
+
+# -- comparison helpers -----------------------------------------------------
+
+
+def _bytes(t):
+    return t.contiguous().view(torch.uint8)
+
+
+def _changed_bytes(before, after):
+    return _bytes(before) != _bytes(after)
+
+
+def _changed_latent_slots(before, after):
+    """(block, token) pairs whose bytes changed in a latent cache buffer
+    ``[blocks, kv_factor, tokens_per_block, heads, head_dim]``."""
+    diff = before.view(torch.uint8) != after.view(torch.uint8)
+    diff = diff.any(dim=4).any(dim=3).any(dim=1)
+    return {tuple(x) for x in diff.nonzero().tolist()}
+
+
+def _index_k_tokens(buf):
+    """Indexer-key cache blocks ``[64 x 128 FP8 | 64 x FP32 scale]`` as
+    (codes [blocks, 64, 128] uint8, scales [blocks, 64] float32)."""
+    raw = buf.contiguous().view(torch.uint8).reshape(buf.shape[0], -1)
+    tokens = buf.shape[1]
+    data = raw[:, : tokens * 128].reshape(-1, tokens, 128)
+    scale = raw[:, tokens * 128 : tokens * 132].contiguous().view(torch.float32).reshape(-1, tokens)
+    return data, scale
+
+
+def _changed_index_k_tokens(before, after):
+    db, sb = _index_k_tokens(before)
+    da, sa = _index_k_tokens(after)
+    changed = (db != da).any(dim=-1) | (_bytes(sb) != _bytes(sa)).view(
+        sb.shape[0], sb.shape[1], 4
+    ).any(dim=-1)
+    return {tuple(x) for x in changed.nonzero().tolist()}
+
+
+def _compare_caches(snap, ref_cache, cand_cache, n_new_tokens, report, tag):
+    """Only the new-token slots may differ from the pre-step snapshot, and
+    the candidate's bytes must equal the reference's."""
+    fails = []
+    for key in cand_cache:
+        ref, cand = ref_cache[key], cand_cache[key]
+        n_diff = int(_changed_bytes(ref, cand).sum())
+        report[f"{tag}:{key}_bytes_differing"] = n_diff
+        if key.startswith("latent"):
+            slots = _changed_latent_slots(snap[key], cand)
+            ref_slots = _changed_latent_slots(snap[key], ref)
+        else:
+            slots = _changed_index_k_tokens(snap[key], cand)
+            ref_slots = _changed_index_k_tokens(snap[key], ref)
+        if len(slots) != n_new_tokens or slots != ref_slots:
+            fails.append(
+                f"{tag}: {key} changed slots {sorted(slots)} vs reference {sorted(ref_slots)}"
+            )
+        if n_diff:
+            fails.append(f"{tag}: {key} differs in {n_diff} bytes")
+    return fails
+
+
+def _compare_records(ref, cand, report, tag):
+    """Bit equality of the recorded MLA attention inputs and indexer inputs;
+    top-k / sparse indices must select equal key sets (row-wise sorted; the
+    per-row overlap is reported when they do not)."""
+    fails = []
+    for layer in LAYERS:
+        r, c = ref[layer], cand[layer]
+        if set(r) != set(c) or "absorbed_q" not in r:
+            fails.append(f"{tag}: layer {layer} recorded {sorted(c)} vs {sorted(r)}")
+            return fails
+        for key in sorted(r):
+            if key in ("topk", "sparse_indices"):
+                rs = r[key].sort(dim=-1).values
+                cs = c[key].sort(dim=-1).values
+                report[f"{tag}:L{layer}.{key}_same_order"] = torch.equal(r[key], c[key])
+                if torch.equal(rs, cs):
+                    continue
+                overlap = min(
+                    len(set(x) & set(y)) / max(len(set(x)), 1)
+                    for x, y in zip(rs.tolist(), cs.tolist())
+                )
+                n_rows = int((rs != cs).any(dim=-1).sum())
+                report[f"{tag}:L{layer}.{key}_rows_differing"] = n_rows
+                report[f"{tag}:L{layer}.{key}_min_row_overlap"] = overlap
+                fails.append(f"{tag}: layer {layer} {key} selects different keys in {n_rows} rows")
+            else:
+                n_diff = int((_bytes(r[key]) != _bytes(c[key])).sum())
+                report[f"{tag}:L{layer}.{key}_bytes_differing"] = n_diff
+                if n_diff:
+                    fails.append(f"{tag}: layer {layer} {key} differs in {n_diff} bytes")
+    return fails
+
+
+def _rel_rms(x, ref):
+    return (
+        (x.double() - ref.double()).pow(2).mean().sqrt()
+        / ref.double().pow(2).mean().sqrt().clamp_min(1e-30)
+    ).item()
+
+
+def _bf16_ulp(value):
+    return 2.0 ** (math.floor(math.log2(max(value, 2.0**-126))) - 7)
+
+
+def _output_bounds(native_runs):
+    """Per-layer bounds from repeated native runs of the same step."""
+    ref = native_runs[0]
+    bounds = {}
+    for layer in LAYERS:
+        noise = max(
+            [
+                (run[layer].float() - ref[layer].float()).abs().max().item()
+                for run in native_runs[1:]
+            ]
+            + [0.0]
+        )
+        rel = max([_rel_rms(run[layer], ref[layer]) for run in native_runs[1:]] + [0.0])
+        ulp = _bf16_ulp(ref[layer].float().abs().max().item())
+        bounds[layer] = (2.0 * max(noise, ulp), max(2.0 * rel, REL_RMS_FLOOR), noise, rel)
+    return bounds
+
+
+def _compare_outputs(ref, cand, bounds, report, tag):
+    fails = []
+    for layer in LAYERS:
+        max_bound, rel_bound, noise, _ = bounds[layer]
+        diff = (cand[layer].float() - ref[layer].float()).abs().max().item()
+        rel = _rel_rms(cand[layer], ref[layer])
+        report[f"{tag}:L{layer}.out_max_diff"] = diff
+        report[f"{tag}:L{layer}.out_rel_rms"] = rel
+        report[f"L{layer}.native_repeat_max_diff"] = noise
+        if diff > max_bound or rel > rel_bound:
+            fails.append(
+                f"{tag}: layer {layer} output max diff {diff:.3e} "
+                f"(bound {max_bound:.3e}), relRMS {rel:.3e} (bound "
+                f"{rel_bound:.3e})"
+            )
+    return fails
+
+
+# Native runs per step used to estimate the attention's run-to-run noise
+# (one reference and four repeats).
+NATIVE_RUNS = 5
+
+
+def _native_runs(h, layout, path, hidden, pos, snap, repeats=NATIVE_RUNS):
+    """Repeated native steps from the same cache snapshot: outputs, the
+    recorded attention inputs and the caches after the first run. Native
+    caches and recorded inputs must themselves be repeatable."""
+    _set_absorbed_query(False)
+    outs, recs, cache, fails = [], [], None, []
+    for i in range(repeats):
+        h.restore(snap)
+        out, rec = h.recorded_step(layout, path, hidden, pos)
+        outs.append(out)
+        recs.append(rec)
+        if i == 0:
+            cache = h.snapshot()
+        else:
+            report = {}
+            fails += _compare_caches(
+                snap, cache, h.snapshot(), math.prod(LAYOUTS[layout]), report, f"native_repeat{i}"
+            )
+            fails += _compare_records(recs[0], rec, report, f"native_repeat{i}")
+    h.restore(snap)
+    assert not fails, f"native step is not repeatable: {fails}"
+    return outs, recs, cache
+
+
+# -- tests ------------------------------------------------------------------
+
+
+def _compare_absorbed_query_to_native(snap, native, candidate, n_new, bounds, report, tag):
+    """absorbed_query step vs native step: caches, MLA and indexer inputs bitwise, equal
+    top-k key sets, noise-bounded outputs."""
+    native_out, native_rec, native_cache = native
+    absorbed_query_out, absorbed_query_rec, absorbed_query_cache = candidate
+    fails = _compare_caches(snap, native_cache, absorbed_query_cache, n_new, report, tag)
+    fails += _compare_records(native_rec, absorbed_query_rec, report, tag)
+    fails += _compare_outputs(native_out, absorbed_query_out, bounds, report, tag)
+    return fails
+
+
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("form", ["auto", "preabsorbed", "native_q"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_absorbed_query_into_native_attention(harness, layout, form, path):
+    h = harness
+    n_new = math.prod(LAYOUTS[layout])
+    hidden, pos = h.gen_inputs(layout, seed=12)
+    snap = h.snapshot()
+    try:
+        native_outs, native_recs, native_cache = _native_runs(h, layout, path, hidden, pos, snap)
+        bounds = _output_bounds(native_outs)
+
+        _set_absorbed_query(True, form)
+        h.reset_absorbed_query_runners()
+        absorbed_query_out, absorbed_query_rec = h.recorded_step(layout, path, hidden, pos)
+        for layer in LAYERS:
+            runner = h.attn[layer]._absorbed_query
+            assert runner not in (None, False), f"absorbed_query did not engage on {layer}"
+        absorbed_query_cache = h.snapshot()
+
+        report, fails = {}, []
+        fails += _compare_absorbed_query_to_native(
+            snap,
+            (native_outs[0], native_recs[0], native_cache),
+            (absorbed_query_out, absorbed_query_rec, absorbed_query_cache),
+            n_new,
+            bounds,
+            report,
+            "absorbed_query",
+        )
+
+        # absorbed_query is deterministic: the same step again from the snapshot ...
+        h.restore(snap)
+        absorbed_query_again, absorbed_query_rec_again = h.recorded_step(layout, path, hidden, pos)
+        fails += _compare_caches(
+            snap, absorbed_query_cache, h.snapshot(), n_new, report, "absorbed_query_again"
+        )
+        fails += _compare_records(
+            absorbed_query_rec, absorbed_query_rec_again, report, "absorbed_query_again"
+        )
+        # ... and 20 more steps rewrite the same slots with the same bytes.
+        for _ in range(20):
+            again = h.gen_step(layout, path, hidden, pos)
+        fails += _compare_caches(
+            snap, absorbed_query_cache, h.snapshot(), n_new, report, "absorbed_query_repeat20"
+        )
+        fails += _compare_outputs(native_outs[0], again, bounds, report, "absorbed_query_repeat20")
+        print(f"layout={layout} form={form} path={path} {report}")
+        assert not fails, fails
+    finally:
+        h.restore(snap)
+        _set_absorbed_query(False)
+
+
+# Graphs captured by the CUDA-graph test: absorbed_query alone, native alone, or both in
+# one process sharing one memory pool (as the graph runner keeps several
+# graphs).
+GRAPH_SETS = {
+    "p1only": ("absorbed_query",),
+    "nativeonly": ("native",),
+    "both": ("absorbed_query", "native"),
+}
+
+
+@pytest.mark.parametrize("graph_set", list(GRAPH_SETS))
+@pytest.mark.parametrize("form", ["auto", "preabsorbed", "native_q"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_absorbed_query_cuda_graph_replay_into_native_attention(harness, layout, form, graph_set):
+    """Captured absorbed_query + native attention + native cache append, replayed once
+    and 20 times with fresh inputs. Each absorbed_query replay must reproduce the eager
+    absorbed_query step bit for bit (caches, recorded attention inputs) and compare with
+    the eager native step as the eager test does; a captured native graph
+    must reproduce the eager native step bit for bit.
+
+    Capture follows the graph runner: every configuration is warmed up on
+    the current stream first, then all are captured on ``torch.cuda.graph``'s
+    own capture stream into one shared memory pool."""
+    h = harness
+    n_new = math.prod(LAYOUTS[layout])
+    path = "custom_op"
+    snap = h.snapshot()
+    graphs = {}
+    try:
+        _graph_replay_checks(h, layout, form, n_new, path, snap, graphs, GRAPH_SETS[graph_set])
+    except BaseException as exc:
+        # Report before the graphs are destroyed: a CUDA error raised while
+        # tearing them down would otherwise hide the original failure.
+        print(f"graph test failure: {exc!r}", flush=True)
+        raise
+    finally:
+        graphs.clear()
+        h.reset_absorbed_query_runners()
+        h.restore(snap)
+        _set_absorbed_query(False)
+
+
+def _graph_replay_checks(h, layout, form, n_new, path, snap, graphs, wanted):
+    """Body of the CUDA-graph test; ``graphs`` is owned by the caller."""
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    configs = [
+        (name, name == "absorbed_query") for name in ("absorbed_query", "native") if name in wanted
+    ]
+    # Diagnostics only: GLM_MOE_DSA_ABSORBED_QUERY_IT_MULTI_STREAM=0 warms up and captures without
+    # the aux-stream overlap the graph runner enables.
+    multi_stream = os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_IT_MULTI_STREAM", "1") != "0"
+
+    def step(meta, static_pos, static_hidden):
+        # The runner's step: on_update_kv_lens() is captured too, so the
+        # graph recomputes the graph-temporary metadata (indexer scan
+        # lengths, slot mappings) instead of addressing tensors that a later
+        # prepare() rebinds.
+        return h.step(meta, static_pos, static_hidden, path)
+
+    # Pass 1: warm up every configuration (graph metadata, runners, JIT).
+    warm = {}
+    for name, enabled in configs:
+        _set_absorbed_query(enabled, form)
+        h.reset_absorbed_query_runners()
+        meta = h.gen_metadata(layout, cuda_graph=True)
+        static_hidden, static_pos = h.gen_inputs(layout, seed=200)
+        with with_multi_stream(multi_stream):
+            for _ in range(2):
+                step(meta, static_pos, static_hidden)
+        torch.cuda.synchronize()
+        print(f"[graph test] warmed up {name}", flush=True)
+        # The captured graph addresses its runners' counters: keep them.
+        runners = [h.attn[layer]._absorbed_query for layer in LAYERS]
+        if enabled:
+            for layer, runner in zip(LAYERS, runners):
+                assert runner not in (None, False), f"absorbed_query did not engage on {layer}"
+        warm[name] = (meta, static_hidden, static_pos, runners)
+    h.restore(snap)
+    # Pass 2: capture into one shared pool.
+    pool = None
+    for name, enabled in configs:
+        meta, static_hidden, static_pos, runners = warm[name]
+        _set_absorbed_query(enabled, form)
+        for layer, runner in zip(LAYERS, runners):
+            h.attn[layer]._absorbed_query = runner
+        graph = torch.cuda.CUDAGraph()
+        with with_multi_stream(multi_stream), _Recorder(h.attn) as rec:
+            with torch.cuda.graph(graph, pool=pool):
+                outs = step(meta, static_pos, static_hidden)
+        pool = graph.pool()
+        graphs[name] = (graph, meta, static_hidden, outs, rec, runners)
+        torch.cuda.synchronize()
+        print(f"[graph test] captured {name}", flush=True)
+    # Eager absorbed_query steps below use their own runners.
+    h.reset_absorbed_query_runners()
+
+    report, fails = {}, []
+    for trial, (seed, replays) in enumerate([(301, 1), (302, 1), (303, 20)]):
+        hidden, pos = h.gen_inputs(layout, seed=seed)
+        native_outs, native_recs, native_cache = _native_runs(h, layout, path, hidden, pos, snap)
+        bounds = _output_bounds(native_outs)
+        _set_absorbed_query(True, form)
+        h.restore(snap)
+        absorbed_query_eager = h.recorded_step(layout, path, hidden, pos)
+        absorbed_query_eager_cache = h.snapshot()
+        _set_absorbed_query(False)
+        print(f"[graph test] trial {trial}: eager steps done", flush=True)
+        replayed = {}
+        for name, (graph, meta, static_hidden, outs, rec, runners) in graphs.items():
+            h.restore(snap)
+            meta.prepare()  # host-side per-step preparation (no realloc)
+            static_hidden.copy_(hidden)
+            for _ in range(replays):
+                graph.replay()
+            torch.cuda.synchronize()
+            print(f"[graph test] trial {trial}: {name} x{replays} replayed", flush=True)
+            replayed[name] = (
+                {layer: outs[layer].clone() for layer in LAYERS},
+                rec.snapshot(),
+                h.snapshot(),
+            )
+        tag = f"t{trial}x{replays}"
+        if "absorbed_query" in replayed:
+            absorbed_query_out, absorbed_query_rec, absorbed_query_cache = replayed[
+                "absorbed_query"
+            ]
+            # Replays reproduce the eager absorbed_query step bit for bit ...
+            fails += _compare_caches(
+                snap,
+                absorbed_query_eager_cache,
+                absorbed_query_cache,
+                n_new,
+                report,
+                f"{tag}:absorbed_query_graph_vs_absorbed_query_eager",
+            )
+            fails += _compare_records(
+                absorbed_query_eager[1],
+                absorbed_query_rec,
+                report,
+                f"{tag}:absorbed_query_graph_vs_absorbed_query_eager",
+            )
+            # ... and compare with native as in the eager test.
+            fails += _compare_absorbed_query_to_native(
+                snap,
+                (native_outs[0], native_recs[0], native_cache),
+                (absorbed_query_out, absorbed_query_rec, absorbed_query_cache),
+                n_new,
+                bounds,
+                report,
+                f"{tag}:absorbed_query_graph_vs_native",
+            )
+        if "native" in replayed:
+            g_out, g_rec, g_cache = replayed["native"]
+            fails += _compare_caches(
+                snap, native_cache, g_cache, n_new, report, f"{tag}:native_graph_vs_eager"
+            )
+            fails += _compare_records(native_recs[0], g_rec, report, f"{tag}:native_graph_vs_eager")
+            fails += _compare_outputs(
+                native_outs[0], g_out, bounds, report, f"{tag}:native_graph_vs_eager"
+            )
+        absorbed_query_runners = graphs["absorbed_query"][5] if "absorbed_query" in graphs else []
+        for runner in absorbed_query_runners + [h.attn[layer]._absorbed_query for layer in LAYERS]:
+            if runner not in (None, False):
+                assert int(runner.sync.abs().sum()) == 0
+    print(f"graph layout={layout} form={form} graphs={wanted} {report}")
+    assert not fails, fails
+
+
+@torch.inference_mode()
+def test_checks_detect_a_wrong_preabsorbed_query(harness, monkeypatch):
+    """Negative control: an absorbed_query packet whose absorbed query is left unwritten
+    (zeroed) must fail both the absorbed-query / quant_q checks and the
+    output bound."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    h = harness
+    layout, path = "6req_x1tok", "forward_impl"
+    hidden, pos = h.gen_inputs(layout, seed=12)
+    snap = h.snapshot()
+    orig_project = query_proj.AbsorbedQueryProjection.project
+
+    def broken_project(self, *args, **kwargs):
+        packet = orig_project(self, *args, **kwargs)
+        if packet.fused_q is not None:
+            packet.fused_q[..., : self.mla.kv_lora_rank].zero_()
+        return packet
+
+    try:
+        native_outs, native_recs, _ = _native_runs(h, layout, path, hidden, pos, snap)
+        bounds = _output_bounds(native_outs)
+        monkeypatch.setattr(query_proj.AbsorbedQueryProjection, "project", broken_project)
+        _set_absorbed_query(True, "preabsorbed")
+        h.reset_absorbed_query_runners()
+        bad_out, bad_rec = h.recorded_step(layout, path, hidden, pos)
+        report = {}
+        rec_fails = _compare_records(native_recs[0], bad_rec, report, "bad")
+        out_fails = _compare_outputs(native_outs[0], bad_out, bounds, report, "bad")
+        print(f"negative control {report}")
+        assert any("absorbed_q" in f for f in rec_fails), rec_fails
+        assert any("quant_q" in f for f in rec_fails), rec_fails
+        assert len(out_fails) == len(LAYERS), out_fails
+    finally:
+        monkeypatch.undo()
+        h.restore(snap)
+        _set_absorbed_query(False)
+        h.reset_absorbed_query_runners()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py
new file mode 100644
index 000000000000..58ec3fd16b63
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py
@@ -0,0 +1,869 @@
+# 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.
+"""absorbed_query: query_front packet correctness against the native TRT-LLM region.
+
+Fixtures: real GLM-5.2 checkpoint weights of one full-indexer layer (6) and
+one shared layer (7) at TP8 rank 0 (see ``absorbed_query_fixture``); synthetic
+activations (unit-RMS rows times the layer's real ``input_layernorm`` gain).
+
+Asserted properties:
+
+* The frozen error budget (fixed before tuning) for every floating output:
+  ``allclose(candidate, native, rtol=0.1, atol=1e-2)`` -- the tolerance of
+  the native fused-A GEMM test -- and normalized RMS error vs an FP64 oracle
+  (with the native BF16 intermediate roundings) ``<= 1.25 * native_error +
+  2**-12``.
+* Native agreement (``_native_agreement``), checked on every launch of the
+  default kernels that is compared with native: every output of the 4/9
+  list and the absorbed query ``fused_q[..., :512]`` is bitwise identical
+  to native -- ``q``, ``compressed_kv`` / ``latent[:, :512]`` (FlashInfer
+  RMSNorm fold order), ``k_pe`` / ``latent[:, 512:]``, the indexer
+  ``q_fp8`` / ``k_fp8`` codes, their UE8M0 scales and the head ``weights``.
+  The indexer key LayerNorm is the native ``k_norm`` call (its inductor
+  kernel is shared with the native region), the RoPE is FlashInfer's
+  ``fma(x, cos, rn(+-x_pair * sin))`` and the FP8 tail is the native
+  ``fused_cat_fp8`` arithmetic. The experimental prefetch kernel (own q_a
+  norm) is held to the frozen budget only.
+
+  ``test_native_agreement_statistics`` reports how often each output
+  differs over many inputs (with the first differing elements), and
+  whether the indexer's BF16 ``wq_b`` output already differs from the
+  native GEMM.
+* absorbed_query determinism: padded rows never change live rows, CUDA-graph replays
+  and every stress launch reproduce an undisturbed eager absorbed_query launch bit for
+  bit, and the per-layer sync counters are zero after every launch.
+* Stress (new inter-program synchronization): delayed producers/consumers,
+  many concurrent launches on several streams, and progress while a
+  concurrent kernel holds all but ``k`` SMs.
+
+``GLM_MOE_DSA_ABSORBED_QUERY_TEST_ROWS`` (comma list) and
+``GLM_MOE_DSA_ABSORBED_QUERY_TEST_PREFETCH`` (0/1) restrict the parametrization, e.g. for
+compute-sanitizer runs.
+"""
+
+import os
+import time
+
+import absorbed_query_fixture
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+_EPS_FLOOR = 2.0**-12
+
+
+def _env_list(name, default):
+    value = os.environ.get(name)
+    if not value:
+        return default
+    return [type(default[0])(int(x)) for x in value.split(",")]
+
+
+ROWS = _env_list("GLM_MOE_DSA_ABSORBED_QUERY_TEST_ROWS", [1, 2, 3, 4, 5, 6, 8, 16])
+PREFETCH = _env_list("GLM_MOE_DSA_ABSORBED_QUERY_TEST_PREFETCH", [True, False])
+STRESS_ROWS = _env_list("GLM_MOE_DSA_ABSORBED_QUERY_TEST_ROWS", [1, 6, 16])
+STRESS_LAUNCHES = int(os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_STRESS_LAUNCHES", "8"))
+LAYERS = [absorbed_query_fixture.FULL_LAYER, absorbed_query_fixture.SHARED_LAYER]
+
+# Kernel configurations whose inter-program waits are exercised by the
+# stress tests: the default (absorption of head h by the producers of head
+# h + 1), separate spinning absorption programs, the prefetching kernel, and
+# the legacy own-head-group layout (reference only).
+STRESS_CONFIGS = {
+    "default": {},
+    "separate_consumers": dict(absorb_in_producer=False),
+    "prefetch": dict(prefetch=True),
+    "legacy_own_head": dict(absorb_prev_head=False),
+}
+
+
+@pytest.fixture(scope="module")
+def model_config():
+    return absorbed_query_fixture.build_model_config(0)
+
+
+_ATTN_CACHE = {}
+
+
+def _attn(model_config, layer):
+    if layer not in _ATTN_CACHE:
+        _ATTN_CACHE[layer] = absorbed_query_fixture.build_native_attention(layer, 0, model_config)
+    return _ATTN_CACHE[layer]
+
+
+def _gamma(layer):
+    return absorbed_query_fixture.load_input_layernorm(layer)
+
+
+def _native(attn, hidden, position_ids):
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_dsa_proj
+
+    outs = forward_dsa_proj(attn, position_ids, hidden, None)
+    rows = hidden.shape[0]
+    fused_q = torch.zeros(
+        rows,
+        attn.num_heads_tp,
+        attn.kv_lora_rank + attn.qk_rope_head_dim,
+        dtype=hidden.dtype,
+        device=hidden.device,
+    )
+    q_nope = outs[0].view(rows, attn.num_heads_tp, attn.qk_head_dim)[..., : attn.qk_nope_head_dim]
+    attn._bmm_bf16_out(
+        q_nope.transpose(0, 1),
+        attn.k_b_proj_trans,
+        attn.k_b_proj_trans.transpose(1, 2),
+        fused_q[..., : attn.kv_lora_rank].transpose(0, 1),
+    )
+    return outs, fused_q
+
+
+def _oracle(attn, hidden):
+    """FP64 reference with the native BF16 intermediate roundings."""
+    kva = (hidden.double() @ attn.kv_a_proj_with_mqa.weight.double().t()).to(torch.bfloat16)
+    rows = hidden.shape[0]
+    q_a = kva[:, : attn.q_lora_rank].double()
+    eps = attn.q_a_layernorm.variance_epsilon
+    qn = q_a * torch.rsqrt(q_a.pow(2).mean(-1, keepdim=True) + eps)
+    qn = (qn * attn.q_a_layernorm.weight.double()).to(torch.bfloat16)
+    q = qn.double() @ attn.q_b_proj.weight.double().t()
+    q_b16 = q.to(torch.bfloat16).view(rows, attn.num_heads_tp, attn.qk_head_dim)
+    absorbed = torch.einsum(
+        "rhd,hcd->rhc", q_b16[..., : attn.qk_nope_head_dim].double(), attn.k_b_proj_trans.double()
+    )
+    ckv = kva[:, attn.q_lora_rank : attn.q_lora_rank + attn.kv_lora_rank].double()
+    ckv = ckv * torch.rsqrt(
+        ckv.pow(2).mean(-1, keepdim=True) + attn.kv_a_layernorm.variance_epsilon
+    )
+    ckv = ckv * attn.kv_a_layernorm.weight.double()
+    k_pe = kva[:, attn.q_lora_rank + attn.kv_lora_rank :]
+    return {"qr": qn.double(), "q": q, "absorbed": absorbed, "ckv": ckv, "k_pe": k_pe}
+
+
+def _rel_rms(x, ref):
+    x = x.double()
+    return ((x - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-30)).item()
+
+
+def _bytes(t):
+    return t.contiguous().view(torch.uint8)
+
+
+def _assert_bitwise(name, cand, native):
+    assert cand.shape == native.shape and cand.dtype == native.dtype, name
+    diff = _bytes(cand) != _bytes(native)
+    if bool(diff.any()):
+        max_abs = (cand.float() - native.float()).abs().max().item()
+        raise AssertionError(
+            f"{name}: {int(diff.sum())} of {diff.numel()} bytes differ "
+            f"(max abs difference {max_abs:.3e})"
+        )
+
+
+def _ordered(t):
+    """Order-preserving integer view: adjacent representable values differ
+    by 1 (+0 and -0 both map to 0)."""
+    t = t.contiguous()
+    if t.dtype == torch.float8_e4m3fn:
+        u = t.view(torch.uint8).to(torch.int32)
+        mag = u & 0x7F
+        return torch.where(u >= 0x80, -mag, mag)
+    if t.dtype == torch.bfloat16:
+        i = t.view(torch.int16).to(torch.int32)
+        return torch.where(i < 0, -(i & 0x7FFF), i)
+    if t.dtype == torch.float32:
+        i = t.view(torch.int32).to(torch.int64)
+        return torch.where(i < 0, -(i & 0x7FFFFFFF), i)
+    raise TypeError(t.dtype)
+
+
+def _steps(cand, native):
+    """Per-element distance in representable values (ulps / codes)."""
+    assert cand.shape == native.shape and cand.dtype == native.dtype
+    return (_ordered(cand) - _ordered(native)).abs()
+
+
+def _first_diffs(cand, ref, n=4):
+    """Up to ``n`` differing elements as ``(index, candidate, native)``."""
+    idx = (_ordered(cand) != _ordered(ref)).nonzero()[:n].tolist()
+    return [(tuple(i), cand[tuple(i)].float().item(), ref[tuple(i)].float().item()) for i in idx]
+
+
+def _native_agreement(outs, fused_q, native_outs, native_fq, kv_lora, rows=slice(None), check=True):
+    """Per-output (differing elements, elements, max steps) statistics of
+    absorbed_query vs native; with ``check`` asserts that every output is bitwise
+    identical to native (the module docstring)."""
+    out = dict(zip(_NAMES, outs))
+    nat = dict(zip(_NAMES, native_outs))
+    pairs = {name: (out[name][rows], nat[name][rows]) for name in ("q", "k_pe", "compressed_kv")}
+    latent_c, latent_n = out["latent_cache"][rows], nat["latent_cache"][rows]
+    pairs["latent_nope"] = (latent_c[:, :kv_lora], latent_n[:, :kv_lora])
+    pairs["latent_rope"] = (latent_c[:, kv_lora:], latent_n[:, kv_lora:])
+    if fused_q is not None:
+        pairs["absorbed"] = (fused_q[rows, :, :kv_lora], native_fq[rows, :, :kv_lora])
+    if "q_fp8" in out:
+        rope = 64
+        for side in ("q", "k"):
+            codes_c, codes_n = out[f"{side}_fp8"][rows], nat[f"{side}_fp8"][rows]
+            pairs[f"{side}_fp8_rope"] = (codes_c[..., :rope], codes_n[..., :rope])
+            pairs[f"{side}_fp8_nope"] = (codes_c[..., rope:], codes_n[..., rope:])
+            pairs[f"{side}_scale"] = (out[f"{side}_scale"][rows], nat[f"{side}_scale"][rows])
+        pairs["weights"] = (out["weights"][rows], nat["weights"][rows])
+    stats, fails = {}, []
+    for name, (cand, ref) in pairs.items():
+        steps = _steps(cand, ref)
+        n_diff = int((steps != 0).sum())
+        stats[name] = (n_diff, steps.numel(), int(steps.max()) if steps.numel() else 0)
+        if n_diff:
+            fails.append(
+                f"{name} differs from native: {stats[name]}, first "
+                f"(index, absorbed_query, native) {_first_diffs(cand, ref)}"
+            )
+    if check:
+        assert not fails, f"{fails} (stats {stats})"
+    return stats, fails
+
+
+def _assert_same_packet(a, b, name):
+    """Two absorbed_query launches on the same inputs must agree bit for bit."""
+    assert len(a.proj_outputs) == len(b.proj_outputs)
+    for out_name, x, y in zip(_NAMES, a.proj_outputs, b.proj_outputs):
+        _assert_bitwise(f"{name} {out_name}", x, y)
+    assert (a.fused_q is None) == (b.fused_q is None)
+    if a.fused_q is not None:
+        kv_lora = a.proj_outputs[1].shape[1]
+        _assert_bitwise(f"{name} absorbed", a.fused_q[..., :kv_lora], b.fused_q[..., :kv_lora])
+
+
+def _check_budget(name, cand, native, ref, report):
+    torch.testing.assert_close(
+        cand.float(), native.float(), rtol=0.1, atol=1e-2, msg=f"{name}: candidate vs native"
+    )
+    e_c, e_n = _rel_rms(cand, ref), _rel_rms(native, ref)
+    report[name] = (e_c, e_n)
+    assert e_c <= 1.25 * e_n + _EPS_FLOOR, (
+        f"{name}: candidate rel-RMS {e_c:.3e} vs native {e_n:.3e}"
+    )
+
+
+# Index of each output in the native 4/9-tensor list.
+_NAMES = (
+    "q",
+    "compressed_kv",
+    "k_pe",
+    "latent_cache",
+    "q_fp8",
+    "k_fp8",
+    "k_scale",
+    "weights",
+    "q_scale",
+)
+
+
+def _inputs(layer, rows, seed=0, padded_from=None):
+    hidden = absorbed_query_fixture.make_hidden(rows, seed=seed + layer, gamma=_gamma(layer))
+    if padded_from is not None:
+        hidden[padded_from:] = 0
+    position_ids = torch.arange(491061, 491061 + rows, dtype=torch.int32, device="cuda").unsqueeze(
+        0
+    )
+    return hidden, position_ids
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("rows", ROWS)
+@pytest.mark.parametrize("absorb", [True, False])
+@pytest.mark.parametrize("fuse_indexer", [True, False])
+@pytest.mark.parametrize("prefetch", PREFETCH)
+@torch.inference_mode()
+def test_contract_q_matches_native(model_config, layer, rows, absorb, fuse_indexer, prefetch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    if fuse_indexer and layer != absorbed_query_fixture.FULL_LAYER:
+        pytest.skip("indexer tail exists on full indexer layers only")
+    attn = _attn(model_config, layer)
+    hidden, position_ids = _inputs(layer, rows)
+    native_outs, native_fq = _native(attn, hidden, position_ids)
+    runner = AbsorbedQueryProjection(
+        attn, KernelConfig(fuse_indexer=fuse_indexer, prefetch=prefetch)
+    )
+    packet = runner.project(position_ids, hidden, absorb=absorb)
+    torch.cuda.synchronize()
+    ref = _oracle(attn, hidden)
+    outs = packet.proj_outputs
+    report = {}
+
+    assert len(outs) == len(native_outs) == (9 if attn.mqa.indexer else 4)
+    for cand, nat in zip(outs, native_outs):
+        assert cand.shape == nat.shape and cand.dtype == nat.dtype
+    assert (packet.fused_q is not None) == absorb
+    # Views of the latent row, as in the native list.
+    assert torch.equal(outs[3][:, : attn.kv_lora_rank], outs[1])
+    assert torch.equal(outs[3][:, attn.kv_lora_rank :], outs[2])
+    # Frozen budget on every floating output.
+    _check_budget("q", outs[0], native_outs[0], ref["q"], report)
+    _check_budget("compressed_kv", outs[1], native_outs[1], ref["ckv"], report)
+    if absorb:
+        _check_budget(
+            "absorbed",
+            packet.fused_q[..., : attn.kv_lora_rank],
+            native_fq[..., : attn.kv_lora_rank],
+            ref["absorbed"],
+            report,
+        )
+    if attn.mqa.indexer is not None:
+        q_fp8, k_fp8, k_scale, weights, q_scale = outs[4:]
+        nq_fp8, nk_fp8, nk_scale, nweights, nq_scale = native_outs[4:]
+        torch.testing.assert_close(
+            k_fp8.float() * k_scale, nk_fp8.float() * nk_scale, rtol=0.1, atol=1e-2
+        )
+        torch.testing.assert_close(
+            q_fp8.float() * q_scale, nq_fp8.float() * nq_scale, rtol=0.1, atol=1e-2
+        )
+        torch.testing.assert_close(weights, nweights, rtol=0.1, atol=1e-2)
+    # Bitwise native agreement; the experimental prefetch kernel normalizes
+    # q_a itself and is held to the frozen budget only.
+    stats, _ = _native_agreement(
+        outs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank, check=not prefetch
+    )
+    print(
+        f"layer={layer} rows={rows} absorb={absorb} "
+        f"fuse_indexer={fuse_indexer} prefetch={prefetch} "
+        f"budget(cand_relrms, native_relrms)={report} "
+        f"agreement(n_diff, n, max_steps)={stats}"
+    )
+    assert int(runner.sync.abs().sum()) == 0, "counters must reset"
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("rows", [6, 16])
+@torch.inference_mode()
+def test_native_agreement_statistics(model_config, layer, rows):
+    """Default kernel over many inputs: how often each output differs from
+    native (and by how many steps), and whether the fused indexer's BF16
+    wq_b output already differs from the native GEMM. Every input must
+    agree with native bit for bit."""
+    import json
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    attn = _attn(model_config, layer)
+    indexer = attn.mqa.indexer
+    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_keep_intermediates=True))
+    n_inputs = int(os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_AGREEMENT_INPUTS", "48"))
+    totals, inputs_with_diff, fails = {}, {}, []
+    for i in range(n_inputs):
+        hidden, position_ids = _inputs(layer, rows, seed=1000 + 17 * i)
+        native_outs, native_fq = _native(attn, hidden, position_ids)
+        packet = runner.project(position_ids, hidden, absorb=True)
+        stats, input_fails = _native_agreement(
+            packet.proj_outputs,
+            packet.fused_q,
+            native_outs,
+            native_fq,
+            attn.kv_lora_rank,
+            check=False,
+        )
+        if indexer is not None:
+            kva = attn.kv_a_proj_with_mqa(hidden)
+            qr = attn._q_a_layernorm_maybe_fused(kva[:, : attn.q_lora_rank])
+            steps = _steps(runner.debug["iq"], indexer.wq_b(qr))
+            stats["indexer_wq_b_bf16"] = (int((steps != 0).sum()), steps.numel(), int(steps.max()))
+        fails += [f"input {i}: {f}" for f in input_fails]
+        for name, (n_diff, n, max_steps) in stats.items():
+            tot = totals.setdefault(name, [0, 0, 0])
+            tot[0] += n_diff
+            tot[1] += n
+            tot[2] = max(tot[2], max_steps)
+            inputs_with_diff[name] = inputs_with_diff.get(name, 0) + (n_diff > 0)
+    summary = {
+        name: dict(
+            inputs_with_diff=inputs_with_diff[name],
+            inputs=n_inputs,
+            diff_elements=tot[0],
+            elements=tot[1],
+            max_steps=tot[2],
+        )
+        for name, tot in totals.items()
+    }
+    print(f"AGREEMENT layer={layer} rows={rows} {json.dumps(summary)}")
+    assert not fails, fails
+    assert int(runner.sync.abs().sum()) == 0
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("absorb", [True, False])
+@torch.inference_mode()
+def test_padded_rows_do_not_affect_live_rows(model_config, layer, absorb):
+    """Rows are independent: padded rows (zero or 1e3 garbage) never change
+    any live-row output (all 4/9 outputs and the absorbed query, bit for
+    bit), live rows agree with native bit for bit, and zero padded rows
+    stay finite."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+    )
+
+    attn = _attn(model_config, layer)
+    runner = AbsorbedQueryProjection(attn)
+    live = slice(0, 3)
+    hidden_zero, position_ids = _inputs(layer, 6, padded_from=3)
+    hidden_junk = hidden_zero.clone()
+    hidden_junk[3:] = 1e3 * torch.randn_like(hidden_junk[3:])
+    a = runner.project(position_ids, hidden_zero, absorb=absorb)
+    b = runner.project(position_ids, hidden_junk, absorb=absorb)
+    native_outs, native_fq = _native(attn, hidden_junk, position_ids)
+    torch.cuda.synchronize()
+    assert len(a.proj_outputs) == (9 if attn.mqa.indexer else 4)
+    for name, x, y in zip(_NAMES, a.proj_outputs, b.proj_outputs):
+        _assert_bitwise(f"padded {name}", x[live], y[live])
+    if absorb:
+        _assert_bitwise(
+            "padded absorbed",
+            a.fused_q[live, :, : attn.kv_lora_rank],
+            b.fused_q[live, :, : attn.kv_lora_rank],
+        )
+    _native_agreement(b.proj_outputs, b.fused_q, native_outs, native_fq, attn.kv_lora_rank, live)
+    for x in a.proj_outputs:
+        assert torch.isfinite(x[3:].float()).all()
+    assert int(runner.sync.abs().sum()) == 0
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("absorb", [True, False])
+@pytest.mark.parametrize("delay_ns", [0, 20000])
+@torch.inference_mode()
+def test_cuda_graph_replay(model_config, layer, absorb, delay_ns):
+    """One and many replays of a captured absorbed_query call (also with delayed
+    producers/consumers) reproduce an eager absorbed_query launch on the same inputs
+    bit for bit (all 4/9 outputs and the absorbed query), agree with
+    native bit for bit, and leave the per-layer counters at zero."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    attn = _attn(model_config, layer)
+    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_delay_ns=delay_ns))
+    eager = AbsorbedQueryProjection(attn)
+    rows = 6
+    static_hidden, position_ids = _inputs(layer, rows)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            runner.project(position_ids, static_hidden, absorb=absorb)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        packet = runner.project(position_ids, static_hidden, absorb=absorb)
+    assert len(packet.proj_outputs) == (9 if attn.mqa.indexer else 4)
+    for step, replays in enumerate([1, 1, 50]):
+        hidden, _ = _inputs(layer, rows, seed=100 + step)
+        static_hidden.copy_(hidden)
+        for _ in range(replays):
+            graph.replay()
+        torch.cuda.synchronize()
+        reference = eager.project(position_ids, hidden, absorb=absorb)
+        native_outs, native_fq = _native(attn, hidden, position_ids)
+        torch.cuda.synchronize()
+        _assert_same_packet(packet, reference, f"replay x{replays}")
+        _native_agreement(
+            packet.proj_outputs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
+        )
+        assert int(runner.sync.abs().sum()) == 0
+        assert int(eager.sync.abs().sum()) == 0
+
+
+def _stress_runner(attn, config_name, delay_ns=0):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    return AbsorbedQueryProjection(
+        attn, KernelConfig(debug_delay_ns=delay_ns, **STRESS_CONFIGS[config_name])
+    )
+
+
+def _check_stressed(packet, reference, native, kv_lora, config_name, tag):
+    """A stressed launch must equal the undisturbed launch of the same
+    configuration bit for bit, and (except the experimental prefetch
+    kernel) agree with native bit for bit."""
+    _assert_same_packet(packet, reference, tag)
+    native_outs, native_fq = native
+    _native_agreement(
+        packet.proj_outputs,
+        packet.fused_q,
+        native_outs,
+        native_fq,
+        kv_lora,
+        check=config_name != "prefetch",
+    )
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("rows", STRESS_ROWS)
+@pytest.mark.parametrize("absorb", [True, False])
+@pytest.mark.parametrize("config_name", list(STRESS_CONFIGS))
+@torch.inference_mode()
+def test_delayed_producers_and_consumers(model_config, layer, rows, absorb, config_name):
+    """Producers delayed before they publish and consumers delayed before
+    they release/reset (up to 50 us, deterministic per program): many
+    back-to-back launches with fresh inputs reproduce the undelayed launch
+    bit for bit and leave the counters at zero."""
+    attn = _attn(model_config, layer)
+    runner = _stress_runner(attn, config_name, delay_ns=50000)
+    plain = _stress_runner(attn, config_name)
+    launches = []
+    for i in range(STRESS_LAUNCHES):
+        hidden, position_ids = _inputs(layer, rows, seed=500 + i)
+        launches.append(
+            (
+                runner.project(position_ids, hidden, absorb=absorb),
+                plain.project(position_ids, hidden, absorb=absorb),
+                _native(attn, hidden, position_ids),
+            )
+        )
+    torch.cuda.synchronize()
+    for i, (packet, reference, native) in enumerate(launches):
+        _check_stressed(
+            packet, reference, native, attn.kv_lora_rank, config_name, f"delayed launch {i}"
+        )
+    assert int(runner.sync.abs().sum()) == 0
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("config_name", list(STRESS_CONFIGS))
+@torch.inference_mode()
+def test_concurrent_launches_on_many_streams(model_config, layer, config_name):
+    """Six independent absorbed_query instances (own counters, like six layers) launched
+    concurrently on six streams, 30 rounds, with and without delays: every
+    stream's last result equals its undisturbed launch bit for bit, the
+    counters are zero, no hang."""
+    attn = _attn(model_config, layer)
+    n_streams, rounds, rows = 6, 30, 6
+    streams = [torch.cuda.Stream() for _ in range(n_streams)]
+    runners = [
+        _stress_runner(attn, config_name, delay_ns=(5000 if s % 2 else 0)) for s in range(n_streams)
+    ]
+    plain = _stress_runner(attn, config_name)
+    inputs = [_inputs(layer, rows, seed=700 + s) for s in range(n_streams)]
+    references = [plain.project(p, h, absorb=True) for h, p in inputs]
+    natives = [_native(attn, h, p) for h, p in inputs]
+    torch.cuda.synchronize()
+    last = [None] * n_streams
+    for _ in range(rounds):
+        for s in range(n_streams):
+            with torch.cuda.stream(streams[s]):
+                hidden, position_ids = inputs[s]
+                last[s] = runners[s].project(position_ids, hidden, absorb=True)
+    torch.cuda.synchronize()
+    for s in range(n_streams):
+        _check_stressed(
+            last[s], references[s], natives[s], attn.kv_lora_rank, config_name, f"stream {s}"
+        )
+        assert int(runners[s].sync.abs().sum()) == 0
+
+
+def _memcpy_copy(t):
+    """A new device tensor initialized by a host-to-device memcpy only."""
+    return t.cpu().to(t.device)
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("rows", STRESS_ROWS)
+@pytest.mark.parametrize("absorb", [True, False])
+@torch.inference_mode()
+def test_absorbed_query_kernel_with_memcpy_initialized_inputs(model_config, layer, rows, absorb):
+    """The absorbed_query kernel launch alone, with every tensor it reads placed by a
+    host-to-device copy: the kv_a output, the q_a_layernorm output, the
+    indexer wk/weights_proj output, the native k_norm output, positions,
+    RoPE-table rows, sync counters and the weights (swapped in for the
+    launch). The result must
+    equal the normal launch bit for bit.
+
+    For ``compute-sanitizer --tool initcheck --kernel-name kns=qproj
+    --check-api-memory-access no`` the suite sets ``GLM_MOE_DSA_ABSORBED_QUERY_INITCHECK_ONLY=1``:
+    then only the memcpy-fed launch runs, so the only instrumented kernel
+    reads only memory placed by memcpys or written by itself, and any report
+    is a read of memory nothing initialized."""
+    initcheck_only = os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_INITCHECK_ONLY") == "1"
+    attn = _attn(model_config, layer)
+    runner = _stress_runner(attn, "default")
+    hidden, _ = _inputs(layer, rows, seed=1200)
+    # Small positions so that a copy of the first RoPE-table rows suffices.
+    position_ids = torch.arange(rows, dtype=torch.int32).unsqueeze(0).cuda()
+    kva = attn.kv_a_proj_with_mqa(hidden)
+    fused = runner.indexer_fusable()
+    kw = torch.nn.functional.linear(hidden, attn.mqa.indexer._fused_wk_wp_weight) if fused else None
+    k_ln = attn.mqa.indexer.k_norm(kw[:, : attn.mqa.indexer.head_dim]) if fused else None
+    qn = runner.native_q_norm(kva)
+    reference = None
+    if not initcheck_only:
+        reference = (
+            runner.launch_indexer(kva, kw, k_ln, None, 0, position_ids, absorb)
+            if fused
+            else runner.launch(kva, need_qr=False, absorb=absorb)
+        )
+    torch.cuda.synchronize()
+
+    probe = _stress_runner(attn, "default")
+    probe.sync = _memcpy_copy(torch.zeros_like(probe.sync))
+    inputs = dict(kva=_memcpy_copy(kva), qn=_memcpy_copy(qn))
+    probe.native_q_norm = lambda _kva: inputs["qn"]
+    params = [
+        attn.q_b_proj.weight,
+        attn.k_b_proj_trans,
+        attn.kv_a_layernorm.weight,
+        attn.q_a_layernorm.weight,
+    ]
+    rope = attn.mqa.indexer.rotary_emb if fused else None
+    saved_table = rope.rotary_cos_sin if fused else None
+    if fused:
+        inputs.update(
+            kw=_memcpy_copy(kw),
+            k_ln=_memcpy_copy(k_ln),
+            pos=_memcpy_copy(position_ids),
+            rope=_memcpy_copy(saved_table[:rows].contiguous()),
+        )
+        indexer = attn.mqa.indexer
+        params += [indexer.wq_b.weight]
+    saved_params = [p.data for p in params]
+    try:
+        for p in params:
+            p.data = _memcpy_copy(p.data)
+        if fused:
+            rope.rotary_cos_sin = inputs["rope"]
+        torch.cuda.synchronize()
+        if fused:
+            result = probe.launch_indexer(
+                inputs["kva"], inputs["kw"], inputs["k_ln"], None, 0, inputs["pos"], absorb
+            )
+        else:
+            result = probe.launch(inputs["kva"], need_qr=False, absorb=absorb)
+        torch.cuda.synchronize()
+    finally:
+        for p, data in zip(params, saved_params):
+            p.data = data
+        if fused:
+            rope.rotary_cos_sin = saved_table
+    assert int(probe.sync.abs().sum()) == 0
+    if initcheck_only:
+        return
+    if fused:
+        _assert_same_packet(result, reference, "memcpy inputs")
+    else:
+        for name, x, y in zip(("q", "qr", "latent", "fused_q"), result, reference):
+            if x is None:
+                assert y is None
+                continue
+            if name == "fused_q":
+                x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
+            _assert_bitwise(f"memcpy inputs {name}", x, y)
+
+
+@triton.jit
+def _occupier_kernel(
+    a_ptr,
+    b_ptr,
+    out_ptr,
+    smid_ptr,
+    hold_ns,
+    ITERS: tl.constexpr,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+    NS: tl.constexpr,
+):
+    """Holds one SM per program for ``hold_ns``: its pipelined dot loop
+    needs ~192 KB of shared memory, so no absorbed_query program fits next to it."""
+    pid = tl.program_id(0)
+    tl.store(smid_ptr + pid, tl.extra.cuda.smid())
+    offs_m = tl.arange(0, BM)
+    offs_n = tl.arange(0, BN)
+    offs_k = tl.arange(0, BK)
+    acc = tl.zeros([BM, BN], dtype=tl.float32)
+    for k in tl.range(0, ITERS, num_stages=NS):
+        a = tl.load(a_ptr + offs_m[:, None] * (ITERS * BK) + (k * BK + offs_k)[None, :])
+        b = tl.load(b_ptr + (k * BK + offs_k)[:, None] * BN + offs_n[None, :])
+        acc = tl.dot(a, b, acc)
+    t0 = tl.extra.cuda.globaltimer()
+    while tl.extra.cuda.globaltimer() - t0 < hold_ns:
+        pass
+    tl.store(out_ptr + pid * BM * BN + offs_m[:, None] * BN + offs_n[None, :], acc)
+
+
+def _compiled_shared_bytes(jit_fn):
+    """Shared memory of the variants of a Triton kernel compiled so far."""
+    sizes = []
+    for caches in jit_fn.device_caches.values():
+        for kernel in caches[0].values():
+            sizes.append(kernel.metadata.shared)
+    return sizes
+
+
+def _launch_absorbed_query_kernel_only(runner, kva, kw, position_ids):
+    """The absorbed_query launch of ``project`` after its native bridges."""
+    if kw is not None:
+        k_ln = runner.mla.mqa.indexer.k_norm(kw[:, : runner.mla.mqa.indexer.head_dim])
+        return runner.launch_indexer(kva, kw, k_ln, None, 0, position_ids, True)
+    q, _, latent, fused_q = runner.launch(kva, need_qr=False, absorb=True)
+    return q, latent, fused_q
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("free_sms", [1, 2, 4, 8])
+@pytest.mark.parametrize("config_name", ["default", "separate_consumers", "legacy_own_head"])
+@torch.inference_mode()
+def test_progress_with_restricted_sms(model_config, layer, free_sms, config_name):
+    """A concurrent kernel holds all but ``free_sms`` SMs for 0.5 s (its
+    programs leave no room for an absorbed_query program on those SMs). The absorbed_query kernel
+    (preabsorbed form; on the full layer with the fused indexer tail) must
+    complete while the SMs are still held, with outputs equal to its
+    unconstrained launch and native bit for bit.
+    The legacy own-head-group layout is reported, not asserted: it needs
+    NOPE // BN co-resident programs of one head group."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        kernels,
+    )
+
+    attn = _attn(model_config, layer)
+    runner = _stress_runner(attn, config_name)
+    rows = 6
+    hidden, position_ids = _inputs(layer, rows, seed=900)
+    native_outs, native_fq = _native(attn, hidden, position_ids)
+    kva = attn.kv_a_proj_with_mqa(hidden)
+    kw = (
+        torch.nn.functional.linear(hidden, attn.mqa.indexer._fused_wk_wp_weight)
+        if runner.indexer_fusable()
+        else None
+    )
+    # Warm up (compile/load) before the timed launch; also the reference.
+    reference = _launch_absorbed_query_kernel_only(runner, kva, kw, position_ids)
+    torch.cuda.synchronize()
+
+    n_sm = torch.cuda.get_device_properties(0).multi_processor_count
+    n_occ = n_sm - free_sms
+    bm = bn = 128
+    bk, iters = 64, 32
+    a = torch.randn(bm, iters * bk, device="cuda").to(torch.bfloat16)
+    b = torch.randn(iters * bk, bn, device="cuda").to(torch.bfloat16)
+    out = torch.empty(n_occ, bm, bn, device="cuda")
+    smids = torch.full((n_occ,), -1, dtype=torch.int32, device="cuda")
+    occ_stream, absorbed_query_stream = torch.cuda.Stream(), torch.cuda.Stream()
+    ev = {k: torch.cuda.Event(enable_timing=True) for k in ("s", "o", "p")}
+    hold_ns = 500_000_000
+    torch.cuda.synchronize()
+    with torch.cuda.stream(occ_stream):
+        ev["s"].record()
+        occ = _occupier_kernel[(n_occ,)](
+            a, b, out, smids, hold_ns, iters, bm, bn, bk, 6, num_warps=4, num_stages=6
+        )
+        ev["o"].record()
+    time.sleep(0.05)  # the occupier is resident long before absorbed_query is launched
+    with torch.cuda.stream(absorbed_query_stream):
+        result = _launch_absorbed_query_kernel_only(runner, kva, kw, position_ids)
+        ev["p"].record()
+    torch.cuda.synchronize()
+    t_absorbed_query = ev["s"].elapsed_time(ev["p"])
+    t_occ = ev["s"].elapsed_time(ev["o"])
+    occ_smem = occ.metadata.shared
+    jit_fn = kernels.qproj_indexer_kernel if kw is not None else (kernels.qproj_kernel)
+    absorbed_query_smem = min(_compiled_shared_bytes(jit_fn) or [0])
+    held = int(torch.unique(smids).numel())
+    report = dict(
+        free_sms=free_sms,
+        held_sms=held,
+        occupier_smem=occ_smem,
+        absorbed_query_min_smem=absorbed_query_smem,
+        absorbed_query_done_ms=round(t_absorbed_query, 3),
+        occupier_done_ms=round(t_occ, 3),
+    )
+    print(f"layer={layer} config={config_name} {report}")
+    assert held == n_occ, report  # one occupier program per SM
+    assert occ_smem + absorbed_query_smem > 228 * 1024, report
+
+    if kw is not None:
+        _assert_same_packet(result, reference, "restricted")
+        _native_agreement(
+            result.proj_outputs, result.fused_q, native_outs, native_fq, attn.kv_lora_rank
+        )
+    else:
+        for name, x, y in zip(("q", "latent", "fused_q"), result, reference):
+            if name == "fused_q":
+                x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
+            _assert_bitwise(f"restricted {name}", x, y)
+        q, latent, fused_q = result
+        _assert_bitwise("q vs native", q, native_outs[0])
+        _assert_bitwise(
+            "absorbed vs native",
+            fused_q[..., : attn.kv_lora_rank],
+            native_fq[..., : attn.kv_lora_rank],
+        )
+        _assert_bitwise("latent vs native", latent, native_outs[3])
+    assert int(runner.sync.abs().sum()) == 0
+    if config_name != "legacy_own_head":
+        assert t_absorbed_query < t_occ - 100.0, (
+            f"absorbed_query did not complete while {n_occ} SMs were held: {report}"
+        )
+
+
+@torch.inference_mode()
+def test_eligibility_and_fallback(model_config):
+    """Unsupported calls return None before any GPU work."""
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    attn = _attn(model_config, absorbed_query_fixture.SHARED_LAYER)
+    ok, reason = query_proj.module_eligibility(attn)
+    assert ok, reason
+    gen = SimpleNamespace(num_contexts=0, num_generations=1)
+    ctx = SimpleNamespace(num_contexts=1, num_generations=0)
+    mixed = SimpleNamespace(num_contexts=1, num_generations=1)
+    hidden, position_ids = _inputs(absorbed_query_fixture.SHARED_LAYER, 6)
+    assert query_proj.call_eligibility(attn, hidden, gen)[0]
+    for meta in (None, ctx, mixed):
+        assert not query_proj.call_eligibility(attn, hidden, meta)[0]
+        attn._absorbed_query = None
+        assert query_proj.dsa_proj(attn, position_ids, hidden, meta) is None
+    big = torch.zeros(17, hidden.shape[1], dtype=hidden.dtype, device="cuda")
+    assert not query_proj.call_eligibility(attn, big, gen)[0]
+    assert not query_proj.call_eligibility(attn, hidden.float(), gen)[0]
+    # Draft (MTP) layers are outside the target-only domain.
+    saved = attn.layer_idx
+    try:
+        attn.layer_idx = attn.num_hidden_layers
+        assert not query_proj.module_eligibility(attn)[0]
+    finally:
+        attn.layer_idx = saved
+    attn._absorbed_query = None
+    packet = query_proj.dsa_proj(attn, position_ids, hidden, gen)
+    assert packet is not None and len(packet.proj_outputs) == 4
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py
new file mode 100644
index 000000000000..37044e1331ac
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py
@@ -0,0 +1,1103 @@
+# 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.
+"""draft_absorbed_query inside the MTP Eagle draft loop, through the target's attention and worker.
+
+Modules: the built-in ``DeepseekV32Attention`` (adopted as the target's ``GlmMLA``) of the
+full-indexer target
+layer 6 (stand-in for the target verify forward) and of the MTP layer 78,
+both with real GLM-5.2 weights at TP8 rank 0; FP8 KV cache (V1
+``DSACacheManager`` holding layers 6 and 78); synthetic activations. Two
+requests with a 2,600-token context are prefilled natively through both
+layers.
+
+One decode iteration runs as the GEN CUDA graph runs it, as a decode step:
+
+1. ``on_update_kv_lens()`` (the runner's ``_preprocess_inputs``);
+2. target verify: layer 6 on ``B x 6`` rows (MTP verify layout);
+3. ``_prepare_attn_metadata_for_spec_dec``, then the draft loop of the target's
+   adopted worker (``GlmMtpWorker._forward_linear_draft_loop`` on an ``MTPEagleWorker``),
+   then ``_restore_attn_metadata_from_spec_dec``. The loop is the real
+   one: it marks the draft loop and the accepted counts for the indexer,
+   calls ``set_skip_topk(i > 0)``, gathers each request's last accepted row,
+   advances positions, sets ``seq_lens`` to 1 and updates ``kv_lens_cuda``.
+   Its ``draft_model`` is a stub whose ``mtp_layers[0]`` runs the native
+   layer-78 attention (``MLA.forward``'s custom-op path, or ``forward_impl``)
+   on the rows the loop hands it and returns an attention-independent next
+   hidden state (a fixed column rotation), and whose ``shared_head`` returns
+   ``hidden[:, :1024]`` as logits. So draft step 0 runs layer 78 on
+   ``B x 6`` rows and steps 1-4 on ``B`` rows, exactly the draft-loop shapes.
+
+Asserted:
+
+* gating through the target's ``_absorbed_query_packet``: ``draft_absorbed_query`` engages only layer
+  78 (all five draft steps, rows ``[6B, B, B, B, B]``), ``absorbed_query`` only layer 6,
+  ``absorbed_query,draft_absorbed_query`` both, none of them nothing; layers that are not engaged
+  produce caches and attention inputs byte-identical to native;
+* indexer reuse semantics: in every run (native and draft_absorbed_query) draft step 0 scores
+  the top-k (``in_mtp_draft_loop`` set, ``indexer_skip_topk`` clear) and
+  steps 1-4 reuse it (both set), and the top-k of steps 1-4 equals, row for
+  row and in order, the step-0 top-k row of each request's last accepted
+  token;
+* cache writes: per draft step the indexer keys handed to
+  ``_update_k_cache`` (``k_fp8``, ``k_scale``) and, after the iteration,
+  every cache (latent and indexer-key caches of layers 6 and 78) are
+  bit-identical to native, and exactly the same slots changed;
+* MLA inputs per call (absorbed query, FP8 ``quant_q_buffer``, BMM scales)
+  and indexer inputs per call (FP8 query, head weights) bit-identical to
+  native; the top-k and sparse-index key sets of every call equal native's;
+* attention outputs per call within the native run-to-run noise bound (as
+  in the absorbed_query integration test);
+* determinism: a repeated draft_absorbed_query iteration reproduces caches and recorded inputs
+  bit for bit;
+* CUDA graph: the whole iteration (target layer + five draft steps) captured
+  in one graph, replayed 1, 1 and 20 times with fresh hidden states and
+  fresh accepted counts, reproduces the eager draft_absorbed_query iteration bit for bit and
+  compares with native as above; a captured native graph reproduces eager
+  native;
+* a negative control (zeroed absorbed query on layer 78) is detected.
+
+Every draft step launches the full fused-indexer kernel and writes every
+packet output, including the indexer query inputs that the reusing steps 1-4
+do not read.
+"""
+
+import inspect
+import os
+import sys
+from types import SimpleNamespace
+
+import absorbed_query_fixture
+import pytest
+import test_absorbed_query_attention_integration as aq_integration
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+
+import plan_testlib  # noqa: E402
+
+TARGET = absorbed_query_fixture.FULL_LAYER
+MTP = 78
+LAYERS = (TARGET, MTP)
+CTX = 2600
+MAX_REQS = 2
+DRAFT_LEN = 5
+VERIFY_ROWS = DRAFT_LEN + 1
+STUB_VOCAB = 1024
+# Native iterations per comparison. The output bound is twice the largest
+# native repeat difference, so it needs enough repeats to sample the native
+# sparse attention's run-to-run spread (graph replays included).
+NATIVE_RUNS = int(os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_IT_NATIVE_RUNS", "9"))
+# Accepted counts per request (including the target-sampled token, 1..6).
+ACCEPTED = {1: [3], 2: [6, 1]}
+
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+
+def _set_features(features: str, form: str = "auto"):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    query_proj.SERVED_FORM = form
+    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
+
+
+class _LaunchLog:
+    """Records ``(layer, rows)`` of every fused-indexer absorbed_query launch."""
+
+    def __enter__(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+            query_proj,
+        )
+
+        self.cls = query_proj.AbsorbedQueryProjection
+        self.orig = self.cls.launch_indexer
+        self.calls = []
+        log, orig = self.calls, self.orig
+
+        def launch_indexer(runner, kva, *args, **kwargs):
+            log.append((runner.mla.layer_idx, kva.shape[0]))
+            return orig(runner, kva, *args, **kwargs)
+
+        self.cls.launch_indexer = launch_indexer
+        return self
+
+    def __exit__(self, *exc):
+        self.cls.launch_indexer = self.orig
+        return False
+
+
+class _Recorder:
+    """Per-call records of what the native attention consumes, per layer:
+    lists indexed by call (one call for layer 6, five for layer 78).
+
+    Wraps ``mla_rope_generation`` (absorbed query, ``quant_q_buffer``, BMM
+    scales), ``sparse_attn_predict`` (converted indices), the indexer's
+    ``_update_k_cache`` (keys appended to the indexer cache) and
+    ``sparse_attn_indexer`` (logical top-k, indexer query inputs and the
+    draft-loop flags it sees). Inside a capture the clones are graph-pool
+    tensors that every replay refreshes; the flags are host values seen at
+    capture."""
+
+    def __init__(self, attn_by_layer):
+        self.attn = attn_by_layer
+        self.data = {layer: {} for layer in attn_by_layer}
+        self._installed = []
+
+    def _wrap(self, obj, name, post):
+        orig = getattr(obj, name)
+        sig = inspect.signature(orig)
+
+        def hook(*args, **kwargs):
+            out = orig(*args, **kwargs)
+            post(sig.bind(*args, **kwargs).arguments, out)
+            return out
+
+        setattr(obj, name, hook)
+        self._installed.append((obj, name))
+
+    def __enter__(self):
+        for layer, attn in self.attn.items():
+            rec = self.data[layer]
+
+            def add(key, value, rec=rec):
+                rec.setdefault(key, []).append(value)
+
+            def rope_post(a, _out, add=add, kv_lora=attn.kv_lora_rank):
+                add("absorbed_q", a["fused_q"][..., :kv_lora].clone())
+                for key, arg in (
+                    ("quant_q", "quant_q_buffer"),
+                    ("bmm1_scale", "mla_bmm1_scale"),
+                    ("bmm2_scale", "mla_bmm2_scale"),
+                ):
+                    if a.get(arg) is not None:
+                        add(key, a[arg].clone())
+
+            def predict_post(_a, out, add=add):
+                if isinstance(out, tuple) and out[0] is not None:
+                    add("sparse_indices", out[0].clone())
+
+            def kcache_post(a, _out, add=add):
+                add("index_k_fp8", a["k_fp8"].clone())
+                add("index_k_scale", a["k_scale"].clone())
+
+            def topk_post(a, out, add=add):
+                meta = a["metadata"]
+                add("topk", out.clone())
+                add("indexer_q_fp8", a["q_fp8"].clone())
+                add("indexer_weights", a["weights"].clone())
+                add("flags", (bool(meta.in_mtp_draft_loop), bool(meta.indexer_skip_topk)))
+
+            self._wrap(attn.mqa, "mla_rope_generation", rope_post)
+            self._wrap(attn.mqa, "sparse_attn_predict", predict_post)
+            self._wrap(attn.mqa.indexer, "_update_k_cache", kcache_post)
+            self._wrap(attn.mqa.indexer, "sparse_attn_indexer", topk_post)
+        return self
+
+    def __exit__(self, *exc):
+        for obj, name in self._installed:
+            delattr(obj, name)
+        self._installed.clear()
+        return False
+
+    def snapshot(self):
+        return {
+            layer: {
+                key: [v.clone() if torch.is_tensor(v) else v for v in vals]
+                for key, vals in rec.items()
+            }
+            for layer, rec in self.data.items()
+        }
+
+
+class _Engagement:
+    """Counts the absorbed_query package entry points the absorbed-query packet calls:
+    ``(entry, layer, rows, engaged)`` per call."""
+
+    def __enter__(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501, F401
+            absorbed_query,
+        )
+
+        self.registry = dispatch._REGISTRY["absorbed_query"]
+        self.saved = dict(self.registry)
+        self.calls = []
+        for name in ("dsa_proj", "mtp_dsa_proj"):
+            fn = self.saved[name]
+
+            def wrapped(mla, position_ids, hidden_states, attn_metadata, fn=fn, name=name):
+                packet = fn(mla, position_ids, hidden_states, attn_metadata)
+                self.calls.append((name, mla.layer_idx, hidden_states.shape[0], packet is not None))
+                return packet
+
+            self.registry[name] = wrapped
+        return self
+
+    def __exit__(self, *exc):
+        self.registry.update(self.saved)
+        return False
+
+
+class _StubShared:
+    def __call__(self, hidden_states, lm_head, attn_metadata, return_context_logits=False):
+        return hidden_states[:, :STUB_VOCAB].float()
+
+
+class _StubMTPLayer:
+    """``mtp_layers[0]``: the native layer-78 attention on the loop's rows;
+    the next hidden state does not depend on the attention output, so native
+    and draft_absorbed_query runs see identical inputs at every draft step."""
+
+    def __init__(self, harness, path, outputs):
+        self.h, self.path, self.outputs = harness, path, outputs
+        self.shared_head = _StubShared()
+
+    def __call__(
+        self,
+        input_ids,
+        position_ids,
+        hidden_states,
+        embed_tokens,
+        attn_metadata,
+        all_rank_num_tokens=None,
+        spec_metadata=None,
+        **kwargs,
+    ):
+        out = self.h.forward(MTP, position_ids, hidden_states, attn_metadata, self.path)
+        self.outputs.append(out.clone())
+        return torch.roll(hidden_states, shifts=97, dims=-1)
+
+
+class _StubDraftModel:
+    def __init__(self, mtp_layer):
+        self.mtp_layers = [mtp_layer]
+        self.embed_tokens = None
+        self.lm_head = None
+
+
+class _Harness:
+    def __init__(self):
+        import tensorrt_llm.bindings
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
+            GlmMtpWorker,
+        )
+        from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
+        from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+        from tensorrt_llm._torch.metadata import KVCacheParams
+        from tensorrt_llm._torch.model_config import ModelConfig
+        from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+        from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
+        from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig
+        from tensorrt_llm.mapping import Mapping
+        from tensorrt_llm.models.modeling_utils import QuantConfig
+        from tensorrt_llm.quantization.mode import QuantAlgo
+
+        self.KVCacheParams = KVCacheParams
+        self.model_config = absorbed_query_fixture.build_model_config(0)
+        pretrained = self.model_config.pretrained_config
+        sparse = self.model_config.sparse_attention_config
+        self.attn = {
+            layer: absorbed_query_fixture.build_native_attention(layer, 0, self.model_config)
+            for layer in LAYERS
+        }
+        assert self.attn[MTP].mqa.indexer.mtp_index_share
+        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        kv_model_config = ModelConfig(
+            mapping=self.mapping,
+            sparse_attention_config=sparse,
+            pretrained_config=pretrained,
+            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8),
+        )
+        self.kvm = DSACacheManager(
+            KvCacheConfig(
+                dtype="fp8",
+                max_tokens=MAX_REQS * (CTX + 128),
+                enable_block_reuse=False,
+                event_buffer_max_size=0,
+            ),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+            num_layers=len(LAYERS),
+            layer_mask=[i in LAYERS for i in range(MTP + 1)],
+            num_kv_heads=1,
+            head_dim=576,
+            tokens_per_block=64,
+            max_seq_len=CTX + 128,
+            max_batch_size=MAX_REQS,
+            mapping=self.mapping,
+            dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
+            sparse_attention_config=sparse,
+            model_config=kv_model_config,
+            pretrained_config=pretrained,
+            vocab_size=pretrained.vocab_size,
+            max_num_tokens=MAX_REQS * CTX,
+        )
+        self.req_ids = list(range(MAX_REQS))
+        for rid in self.req_ids:
+            # CTX + 10 tokens (6 verify slots, up to 4 more draft slots) fit
+            # the blocks of CTX + 1 tokens (64 per block).
+            self.kvm.add_dummy_requests(
+                request_ids=[rid], token_nums=[CTX + 1], is_gen=False, prepare_resource=True
+            )
+        sparse_params = sparse.to_sparse_params(pretrained_config=pretrained)
+        self.Metadata = get_attention_backend("TRTLLM", sparse_params=sparse_params).Metadata
+        self.meta_params = sparse.to_sparse_metadata_params(pretrained_config=pretrained)
+        assert self.meta_params.mtp_index_share
+        self.gamma = {layer: absorbed_query_fixture.load_input_layernorm(layer) for layer in LAYERS}
+        self.worker = GlmMtpWorker.adopt(
+            MTPEagleWorker(MTPDecodingConfig(max_draft_len=5), mapping=self.mapping)
+        )
+        assert self.worker.is_mtp_eagle
+        self.batch_indices = torch.arange(MAX_REQS, dtype=torch.int32, device="cuda")
+
+        # Native prefill of both requests through both layers.
+        _set_features("")
+        hidden = absorbed_query_fixture.make_hidden(
+            MAX_REQS * CTX, seed=11, gamma=self.gamma[TARGET]
+        )
+        positions = (
+            torch.arange(CTX, dtype=torch.int32, device="cuda").repeat(MAX_REQS).unsqueeze(0)
+        )
+        ctx_meta = self.Metadata(
+            seq_lens=torch.tensor([CTX] * MAX_REQS, dtype=torch.int),
+            request_ids=self.req_ids,
+            max_num_requests=MAX_REQS,
+            num_contexts=MAX_REQS,
+            prompt_lens=[CTX] * MAX_REQS,
+            max_num_tokens=MAX_REQS * CTX,
+            kv_cache_manager=self.kvm,
+            kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0] * MAX_REQS),
+            mapping=self.mapping,
+            sparse_metadata_params=self.meta_params,
+        )
+        ctx_meta.prepare()
+        ctx_meta.on_update_kv_lens()
+        for layer in LAYERS:
+            self.forward(layer, positions, hidden, ctx_meta, "forward_impl")
+        torch.cuda.synchronize()
+
+    # -- inputs and metadata -----------------------------------------------
+
+    def iteration_inputs(self, nreq, seed, accepted=None):
+        rows = nreq * VERIFY_ROWS
+        accepted = ACCEPTED[nreq] if accepted is None else accepted
+        return SimpleNamespace(
+            hidden_t=absorbed_query_fixture.make_hidden(rows, seed=seed, gamma=self.gamma[TARGET]),
+            pos_t=(CTX + torch.arange(VERIFY_ROWS, dtype=torch.int32, device="cuda"))
+            .repeat(nreq)
+            .unsqueeze(0),
+            hidden_d=absorbed_query_fixture.make_hidden(rows, seed=seed + 1, gamma=self.gamma[MTP]),
+            num_accepted=torch.tensor(accepted, dtype=torch.int32, device="cuda"),
+            input_ids=torch.arange(rows, dtype=torch.int32, device="cuda"),
+        )
+
+    def gen_metadata(self, nreq, cuda_graph=False):
+        """Verify-step metadata as the runner builds it (see the absorbed_query
+        integration test): ``nreq`` requests x 6 tokens, MTP verify width
+        from ``update_spec_dec_param``, then ``prepare()``."""
+        req_ids = self.req_ids[:nreq]
+        kv_params = self.KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[CTX] * nreq)
+        meta = self.Metadata(
+            seq_lens=torch.tensor([VERIFY_ROWS] * nreq, dtype=torch.int),
+            request_ids=req_ids,
+            max_num_requests=nreq,
+            num_contexts=0,
+            prompt_lens=[CTX] * nreq,
+            max_num_tokens=nreq * VERIFY_ROWS,
+            kv_cache_manager=self.kvm,
+            kv_cache_params=kv_params,
+            mapping=self.mapping,
+            sparse_metadata_params=self.meta_params,
+        )
+        if cuda_graph:
+            from tensorrt_llm._torch.memory_buffer_utils import Buffers
+
+            meta = meta.create_cuda_graph_metadata(nreq, False, DRAFT_LEN, Buffers())
+            assert meta.is_cuda_graph
+            meta.request_ids = req_ids
+            meta.prompt_lens = [CTX] * nreq
+            meta.num_contexts = 0
+            meta.kv_cache_params = kv_params
+        meta.update_spec_dec_param(
+            batch_size=nreq,
+            is_spec_decoding_enabled=True,
+            is_spec_dec_tree=False,
+            is_spec_dec_dynamic_tree=False,
+            max_draft_len=DRAFT_LEN,
+            max_total_draft_tokens=DRAFT_LEN,
+        )
+        meta.prepare()
+        return meta
+
+    # -- forwards ----------------------------------------------------------
+
+    def forward(self, layer, positions, hidden, meta, path):
+        attn = self.attn[layer]
+        if path == "forward_impl":
+            out = attn._create_outputs(hidden, meta)
+            attn.forward_impl(positions, hidden, meta, attn_output=out)
+            return out[0]
+        import weakref
+
+        from tensorrt_llm._torch.utils import model_extra_attrs
+
+        extra = self.model_config.extra_attrs
+        extra["attention_metadata"] = weakref.ref(meta)
+        with model_extra_attrs(extra):
+            out = [torch.ops.trtllm.create_mla_outputs(hidden, attn.layer_idx_str)]
+            attn._forward_custom_op(hidden, positions, out, None)
+        return out[0]
+
+    def iteration(self, meta, it, path):
+        """Target verify (layer 6) and the draft loop (layer 78), as the GEN graph
+        runs one decode iteration (a decode step). Returns the attention
+        outputs per layer and call."""
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+            dispatch,
+        )
+
+        with dispatch.step(dispatch.Phase.DECODE):
+            return self._iteration(meta, it, path)
+
+    def _iteration(self, meta, it, path):
+        nreq = it.num_accepted.shape[0]
+        meta.on_update_kv_lens()
+        outs = {TARGET: [], MTP: []}
+        outs[TARGET].append(self.forward(TARGET, it.pos_t, it.hidden_t, meta, path).clone())
+        draft_model = _StubDraftModel(_StubMTPLayer(self, path, outs[MTP]))
+        spec_metadata = SimpleNamespace(
+            runtime_draft_len=DRAFT_LEN,
+            batch_indices_cuda=self.batch_indices[:nreq],
+            all_rank_num_tokens=None,
+            subseq_all_rank_num_tokens=None,
+            wants_advanced_draft_sampling=False,
+            max_num_requests=nreq,
+        )
+        self.worker._prepare_attn_metadata_for_spec_dec(meta)
+        try:
+            inputs = {
+                "input_ids": it.input_ids,
+                "position_ids": it.pos_t.squeeze(0),
+                "hidden_states": it.hidden_d,
+                "attn_metadata": meta,
+                "spec_metadata": spec_metadata,
+            }
+            outs["tokens"] = self.worker._forward_linear_draft_loop(
+                inputs, meta, spec_metadata, draft_model, None, 0, nreq, it.num_accepted, None
+            )
+        finally:
+            self.worker._restore_attn_metadata_from_spec_dec(meta)
+            meta.use_spec_decoding = True
+        return outs
+
+    def run(self, nreq, it, path, meta=None):
+        """One eager iteration with the current features: outputs, records,
+        entry-point engagement."""
+        meta = meta if meta is not None else self.gen_metadata(nreq)
+        with _Recorder(self.attn) as rec, _Engagement() as eng:
+            outs = self.iteration(meta, it, path)
+        torch.cuda.synchronize()
+        return outs, rec.snapshot(), eng.calls
+
+    # -- caches ------------------------------------------------------------
+
+    def caches(self):
+        bufs = {}
+        for layer in LAYERS:
+            bufs[f"latent{layer}"] = self.kvm.get_buffers(layer)
+            bufs[f"index_k{layer}"] = self.kvm.get_indexer_k_cache_buffers(layer)
+        return bufs
+
+    def snapshot(self):
+        return {k: v.clone() for k, v in self.caches().items()}
+
+    def restore(self, snap):
+        for k, v in self.caches().items():
+            v.copy_(snap[k])
+        torch.cuda.synchronize()
+
+    def reset_runners(self):
+        for attn in self.attn.values():
+            attn._absorbed_query = None
+            attn._draft_absorbed_query = None
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = _Harness()
+    yield h
+    _set_features("")
+
+
+# -- comparison helpers -----------------------------------------------------
+
+
+def _expected_slots(nreq, accepted):
+    """Changed token slots per layer after one iteration: six verify slots
+    per request on both layers; on the MTP layer steps 1-4 append at
+    positions ``accepted .. accepted + 3`` after the context."""
+    mtp = sum(len(set(range(VERIFY_ROWS)) | set(range(a, a + DRAFT_LEN - 1))) for a in accepted)
+    return {TARGET: nreq * VERIFY_ROWS, MTP: mtp}
+
+
+def _compare_caches(snap, ref_cache, cand_cache, expected, report, tag):
+    """Changed slots equal the reference's (and the expected count) and
+    every byte equals the reference's."""
+    fails = []
+    for key in cand_cache:
+        ref, cand = ref_cache[key], cand_cache[key]
+        layer = int(key.replace("latent", "").replace("index_k", ""))
+        n_diff = int(aq_integration._changed_bytes(ref, cand).sum())
+        report[f"{tag}:{key}_bytes_differing"] = n_diff
+        if key.startswith("latent"):
+            slots = aq_integration._changed_latent_slots(snap[key], cand)
+            ref_slots = aq_integration._changed_latent_slots(snap[key], ref)
+        else:
+            slots = aq_integration._changed_index_k_tokens(snap[key], cand)
+            ref_slots = aq_integration._changed_index_k_tokens(snap[key], ref)
+        if slots != ref_slots or len(ref_slots) != expected[layer]:
+            fails.append(
+                f"{tag}: {key} changed {len(slots)} slots vs "
+                f"reference {len(ref_slots)} (expected "
+                f"{expected[layer]})"
+            )
+        if n_diff:
+            fails.append(f"{tag}: {key} differs in {n_diff} bytes")
+    return fails
+
+
+def _check_reuse(rec, accepted, report, tag):
+    """Indexer reuse semantics of one run: layer 6 outside the draft loop;
+    layer 78 scores at step 0 and reuses at steps 1-4, the reused rows being
+    the step-0 rows of each request's last accepted token, in order."""
+    fails = []
+    if rec[TARGET]["flags"] != [(False, False)]:
+        fails.append(f"{tag}: target flags {rec[TARGET]['flags']}")
+    flags = rec[MTP]["flags"]
+    want = [(True, False)] + [(True, True)] * (DRAFT_LEN - 1)
+    if flags != want:
+        fails.append(f"{tag}: draft-loop flags {flags}, expected {want}")
+    topk = rec[MTP]["topk"]
+    if len(topk) != DRAFT_LEN:
+        return fails + [f"{tag}: {len(topk)} layer-78 top-k calls"]
+    rows = torch.tensor(
+        [r * VERIFY_ROWS + min(max(a - 1, 0), DRAFT_LEN) for r, a in enumerate(accepted)],
+        device=topk[0].device,
+    )
+    stash = topk[0][rows]
+    for step in range(1, DRAFT_LEN):
+        if not torch.equal(topk[step], stash):
+            fails.append(f"{tag}: step {step} top-k is not the step-0 last-accepted row")
+    report[f"{tag}:reuse_checked_steps"] = DRAFT_LEN - 1
+    return fails
+
+
+def _same_sets(a, b):
+    return torch.equal(a.sort(dim=-1).values, b.sort(dim=-1).values)
+
+
+def _compare_records(ref, cand, report, tag, layers=LAYERS):
+    """Per call: MLA inputs (absorbed query, FP8 query, BMM scales), indexer
+    inputs (FP8 query, head weights) and appended indexer keys bitwise; top-k
+    and sparse-index key sets equal."""
+    fails = []
+    for layer in layers:
+        r, c = ref[layer], cand[layer]
+        if set(r) != set(c) or "absorbed_q" not in r:
+            return [f"{tag}: layer {layer} recorded {sorted(c)} vs {sorted(r)}"]
+        for key in sorted(r):
+            if len(r[key]) != len(c[key]):
+                fails.append(f"{tag}: layer {layer} {key}: {len(c[key])} calls vs {len(r[key])}")
+        if fails:
+            return fails
+        for i in range(len(r["absorbed_q"])):
+            name = f"{tag}:L{layer}.call{i}"
+            for key in (
+                "absorbed_q",
+                "quant_q",
+                "bmm1_scale",
+                "bmm2_scale",
+                "indexer_q_fp8",
+                "indexer_weights",
+                "index_k_fp8",
+                "index_k_scale",
+            ):
+                if key not in r:
+                    continue
+                n_diff = int(
+                    (aq_integration._bytes(r[key][i]) != aq_integration._bytes(c[key][i])).sum()
+                )
+                report[f"{name}.{key}_bytes_differing"] = n_diff
+                if n_diff:
+                    fails.append(f"{name}: {key} differs in {n_diff} bytes")
+            for key in ("topk", "sparse_indices"):
+                same = _same_sets(r[key][i], c[key][i])
+                report[f"{name}.{key}_same_sets"] = same
+                if not same:
+                    fails.append(f"{name}: {key} selects different keys")
+    return fails
+
+
+def _output_bounds(native_runs):
+    """Per (layer, call) bounds from repeated native iterations."""
+    ref = native_runs[0]
+    bounds = {}
+    for layer in LAYERS:
+        for i in range(len(ref[layer])):
+            base = ref[layer][i].float()
+            noise = max(
+                [(run[layer][i].float() - base).abs().max().item() for run in native_runs[1:]]
+                + [0.0]
+            )
+            rel = max(
+                [aq_integration._rel_rms(run[layer][i], ref[layer][i]) for run in native_runs[1:]]
+                + [0.0]
+            )
+            ulp = aq_integration._bf16_ulp(base.abs().max().item())
+            bounds[(layer, i)] = (
+                2.0 * max(noise, ulp),
+                max(2.0 * rel, aq_integration.REL_RMS_FLOOR),
+                noise,
+            )
+    return bounds
+
+
+def _compare_outputs(ref, cand, bounds, report, tag):
+    fails = []
+    for (layer, i), (max_bound, rel_bound, noise) in bounds.items():
+        diff = (cand[layer][i].float() - ref[layer][i].float()).abs().max().item()
+        rel = aq_integration._rel_rms(cand[layer][i], ref[layer][i])
+        report[f"{tag}:L{layer}.call{i}.out_max_diff"] = diff
+        report[f"{tag}:L{layer}.call{i}.out_rel_rms"] = rel
+        report[f"L{layer}.call{i}.native_repeat_max_diff"] = noise
+        if diff > max_bound or rel > rel_bound:
+            fails.append(
+                f"{tag}: L{layer} call {i} output max diff "
+                f"{diff:.3e} (bound {max_bound:.3e}), relRMS "
+                f"{rel:.3e} (bound {rel_bound:.3e})"
+            )
+    return fails
+
+
+def _native_runs(h, nreq, it, path, snap, repeats=NATIVE_RUNS):
+    """Repeated native iterations from the same snapshot; the native
+    iteration must itself be repeatable in caches and recorded inputs."""
+    _set_features("")
+    h.reset_runners()
+    outs, recs, cache, fails = [], [], None, []
+    accepted = it.num_accepted.tolist()
+    expected = _expected_slots(nreq, accepted)
+    for i in range(repeats):
+        h.restore(snap)
+        out, rec, calls = h.run(nreq, it, "custom_op" if path is None else path)
+        assert calls == [], f"native run called absorbed_query entry points: {calls}"
+        outs.append(out)
+        recs.append(rec)
+        report = {}
+        fails += _check_reuse(rec, accepted, report, f"native{i}")
+        if i == 0:
+            cache = h.snapshot()
+        else:
+            fails += _compare_caches(
+                snap, cache, h.snapshot(), expected, report, f"native_repeat{i}"
+            )
+            fails += _compare_records(recs[0], rec, report, f"native_repeat{i}")
+    h.restore(snap)
+    assert not fails, f"native iteration is not repeatable: {fails}"
+    return outs, recs, cache
+
+
+def _expected_calls(features, nreq):
+    """Entry-point calls the absorbed-query packet makes in one iteration."""
+    selected = set(features.split(","))
+    target_query, draft_query = "absorbed_query" in selected, "draft_absorbed_query" in selected
+    target_rows, steps = nreq * VERIFY_ROWS, [nreq * VERIFY_ROWS] + [nreq] * (DRAFT_LEN - 1)
+    calls = []
+    if target_query:
+        calls.append(("dsa_proj", TARGET, target_rows, True))
+    elif draft_query:
+        calls.append(("mtp_dsa_proj", TARGET, target_rows, False))
+    for rows in steps:
+        if target_query:
+            calls.append(("dsa_proj", MTP, rows, False))
+        if draft_query:
+            calls.append(("mtp_dsa_proj", MTP, rows, True))
+    return calls
+
+
+def _compare_to_native(snap, native, cand, accepted, nreq, bounds, report, tag):
+    """Caches and recorded inputs bit-identical to native, the indexer reuse
+    semantics, and outputs within the native run-to-run bound."""
+    native_out, native_rec, native_cache = native
+    out, rec, cache = cand
+    fails = _compare_caches(snap, native_cache, cache, _expected_slots(nreq, accepted), report, tag)
+    fails += _compare_records(native_rec, rec, report, tag)
+    fails += _check_reuse(rec, accepted, report, tag)
+    fails += _compare_outputs(native_out, out, bounds, report, tag)
+    return fails
+
+
+def _expected_launches(nreq):
+    """Fused-indexer launches of layer 78 in one iteration (draft_absorbed_query only)."""
+    return [(MTP, nreq * VERIFY_ROWS)] + [(MTP, nreq)] * (DRAFT_LEN - 1)
+
+
+# -- tests ------------------------------------------------------------------
+
+FEATURE_SETS = {
+    "none": "",
+    "absorbed_query": "absorbed_query",
+    "draft_absorbed_query": "draft_absorbed_query",
+    "absorbed_query_draft_query": "absorbed_query,draft_absorbed_query",
+}
+
+
+@pytest.mark.parametrize("features", list(FEATURE_SETS))
+@pytest.mark.parametrize("nreq", [1, 2])
+@torch.inference_mode()
+def test_feature_gating_through_adapter(harness, features, nreq):
+    """Which layers each selection engages, through the absorbed-query packet,
+    and native agreement of the whole iteration."""
+    h = harness
+    sel = FEATURE_SETS[features]
+    it = h.iteration_inputs(nreq, seed=12)
+    accepted = it.num_accepted.tolist()
+    snap = h.snapshot()
+    try:
+        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, "custom_op", snap)
+        bounds = _output_bounds(native_outs)
+        _set_features(sel)
+        h.reset_runners()
+        h.restore(snap)
+        out, rec, calls = h.run(nreq, it, "custom_op")
+        cache = h.snapshot()
+        expected = _expected_calls(sel, nreq)
+        assert calls == expected, f"engagement {calls} != {expected}"
+        engaged = {layer for _, layer, _, ok in calls if ok}
+        report = {}
+        fails = _compare_to_native(
+            snap,
+            (native_outs[0], native_recs[0], native_cache),
+            (out, rec, cache),
+            accepted,
+            nreq,
+            bounds,
+            report,
+            f"{features}",
+        )
+        print(f"features={sel!r} nreq={nreq} engaged={sorted(engaged)} {report}")
+        assert not fails, fails
+    finally:
+        h.restore(snap)
+        _set_features("")
+        h.reset_runners()
+
+
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("form", ["auto", "preabsorbed", "native_q"])
+@pytest.mark.parametrize("nreq", [1, 2])
+@torch.inference_mode()
+def test_draft_query_draft_loop_matches_native(harness, nreq, form, path):
+    """draft_absorbed_query alone (layer 78 on all five draft steps) against native: caches,
+    MLA inputs, indexer keys and reuse, outputs; then determinism (the same
+    iteration again reproduces caches and records bit for bit)."""
+    h = harness
+    it = h.iteration_inputs(nreq, seed=20 + nreq)
+    accepted = it.num_accepted.tolist()
+    snap = h.snapshot()
+    try:
+        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, path, snap)
+        bounds = _output_bounds(native_outs)
+        _set_features("draft_absorbed_query", form)
+        h.reset_runners()
+        h.restore(snap)
+        with _LaunchLog() as launches:
+            out, rec, calls = h.run(nreq, it, path)
+        cache = h.snapshot()
+        assert calls == _expected_calls("draft_absorbed_query", nreq), calls
+        assert launches.calls == _expected_launches(nreq), launches.calls
+        report = {}
+        fails = _compare_to_native(
+            snap,
+            (native_outs[0], native_recs[0], native_cache),
+            (out, rec, cache),
+            accepted,
+            nreq,
+            bounds,
+            report,
+            "draft_absorbed_query",
+        )
+        h.restore(snap)
+        _, rec_again, _ = h.run(nreq, it, path)
+        fails += _compare_caches(
+            snap, cache, h.snapshot(), _expected_slots(nreq, accepted), report, "draft_query_again"
+        )
+        fails += _compare_records(rec, rec_again, report, "draft_query_again")
+        print(f"nreq={nreq} form={form} path={path} {report}")
+        assert not fails, fails
+    finally:
+        h.restore(snap)
+        _set_features("")
+        h.reset_runners()
+
+
+GRAPH_SETS = {
+    "w5only": ("draft_absorbed_query",),
+    "nativeonly": ("native",),
+    "both": ("draft_absorbed_query", "native"),
+}
+# Diagnostics: restrict the graph cases (comma lists).
+GRAPH_NREQ = [
+    int(x)
+    for x in os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_IT_GRAPH_NREQ", "1,2").split(",")
+]
+GRAPH_SET_NAMES = os.environ.get(
+    "GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_IT_GRAPH_SETS", ",".join(GRAPH_SETS)
+).split(",")
+
+
+def _tensor_ptrs(meta):
+    """data_ptr of every CUDA tensor attribute of the metadata (what a
+    captured graph may have baked in), as the graph runner's strict buffer
+    check records them."""
+    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
+
+    return CUDAGraphRunner._snapshot_graph_tensor_ptrs(meta)
+
+
+def _moved(before, after):
+    return sorted(k for k, v in before.items() if after.get(k) != v)
+
+
+@pytest.mark.parametrize("graph_set", GRAPH_SET_NAMES)
+@pytest.mark.parametrize("form", ["auto", "native_q"])
+@pytest.mark.parametrize("nreq", GRAPH_NREQ)
+@torch.inference_mode()
+def test_draft_query_draft_loop_cuda_graph(harness, nreq, form, graph_set):
+    """The whole iteration (target layer + native draft loop over layer 78)
+    captured in one CUDA graph, as the GEN graph runner does: warm-up of
+    every configuration first, then capture on ``torch.cuda.graph``'s stream
+    into one shared pool. Replays 1, 1 and 20 times with fresh hidden states
+    and fresh accepted counts must reproduce the eager iteration of the same
+    selection bit for bit (caches, records) and compare with eager native as
+    the eager test does."""
+    h = harness
+    graphs = {}
+    snap = h.snapshot()
+    try:
+        _graph_checks(h, nreq, form, snap, graphs, GRAPH_SETS[graph_set])
+    except BaseException as exc:
+        print(f"graph test failure: {exc!r}", flush=True)
+        raise
+    finally:
+        graphs.clear()
+        h.reset_runners()
+        h.restore(snap)
+        _set_features("")
+
+
+def _graph_checks(h, nreq, form, snap, graphs, wanted):
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    path = "custom_op"
+    features = {"draft_absorbed_query": "draft_absorbed_query", "native": ""}
+    configs = [name for name in ("draft_absorbed_query", "native") if name in wanted]
+    warm = {}
+    for name in configs:
+        _set_features(features[name], form)
+        h.reset_runners()
+        meta = h.gen_metadata(nreq, cuda_graph=True)
+        static = h.iteration_inputs(nreq, seed=200)
+        with with_multi_stream(True):
+            for _ in range(2):
+                # Host-side per-step preparation, as before every replay
+                # (the draft loop moved kv_lens_cuda on the device).
+                meta.prepare()
+                h.iteration(meta, static, path)
+        torch.cuda.synchronize()
+        runners = {
+            layer: (h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query)
+            for layer in LAYERS
+        }
+        if name == "draft_absorbed_query":
+            assert runners[MTP][1] not in (None, False), "draft_absorbed_query did not engage"
+        warm[name] = (meta, static, runners)
+    h.restore(snap)
+    pool = None
+    for name in configs:
+        meta, static, runners = warm[name]
+        _set_features(features[name], form)
+        for layer, (r_absorbed_query, r_draft_query) in runners.items():
+            h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query = (
+                r_absorbed_query,
+                r_draft_query,
+            )
+        graph = torch.cuda.CUDAGraph()
+        ptrs_before = _tensor_ptrs(meta)
+        with (
+            with_multi_stream(True),
+            _Recorder(h.attn) as rec,
+            _Engagement() as eng,
+            _LaunchLog() as launches,
+        ):
+            with torch.cuda.graph(graph, pool=pool):
+                outs = h.iteration(meta, static, path)
+        pool = graph.pool()
+        ptrs_captured = _tensor_ptrs(meta)
+        print(
+            f"[graph test] {name}: metadata tensors rebound during "
+            f"capture: {_moved(ptrs_before, ptrs_captured)}; new: "
+            f"{sorted(set(ptrs_captured) - set(ptrs_before))}",
+            flush=True,
+        )
+        assert eng.calls == _expected_calls(features[name], nreq), eng.calls
+        if name == "draft_absorbed_query":
+            assert launches.calls == _expected_launches(nreq), launches.calls
+        graphs[name] = (graph, meta, static, outs, rec, runners, ptrs_captured)
+        torch.cuda.synchronize()
+        print(f"[graph test] captured {name}", flush=True)
+    h.reset_runners()
+
+    report, fails = {}, []
+    trials = [(301, 1, None), (302, 1, [1] * nreq), (303, 20, [VERIFY_ROWS] * nreq)]
+    for trial, (seed, replays, accepted) in enumerate(trials):
+        it = h.iteration_inputs(nreq, seed=seed, accepted=accepted)
+        accepted = it.num_accepted.tolist()
+        expected = _expected_slots(nreq, accepted)
+        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, path, snap)
+        bounds = _output_bounds(native_outs)
+        _set_features("draft_absorbed_query", form)
+        h.reset_runners()
+        h.restore(snap)
+        draft_query_eager = h.run(nreq, it, path)
+        draft_query_eager_cache = h.snapshot()
+        _set_features("")
+        h.reset_runners()
+        tag = f"t{trial}x{replays}"
+        for name, (graph, meta, static, outs, rec, runners, ptrs_captured) in graphs.items():
+            h.restore(snap)
+            meta.prepare()
+            moved = _moved(ptrs_captured, _tensor_ptrs(meta))
+            print(
+                f"[graph test] trial {trial} {name}: metadata tensors "
+                f"rebound by prepare() since capture: {moved}",
+                flush=True,
+            )
+            for key in ("hidden_t", "pos_t", "hidden_d", "num_accepted", "input_ids"):
+                getattr(static, key).copy_(getattr(it, key))
+            for i in range(replays):
+                if i:
+                    # The captured draft loop advances kv_lens_cuda on the
+                    # device; like the runner before every step, rebuild the
+                    # step's metadata from the host first (the replays
+                    # rewrite the same slots with the same bytes).
+                    meta.prepare()
+                graph.replay()
+            torch.cuda.synchronize()
+            print(f"[graph test] trial {trial} {name}: replayed x{replays}", flush=True)
+            g_out = {layer: [o.clone() for o in outs[layer]] for layer in LAYERS}
+            g_rec, g_cache = rec.snapshot(), h.snapshot()
+            fails += _check_reuse(g_rec, accepted, report, f"{tag}:{name}")
+            if name == "draft_absorbed_query":
+                fails += _compare_caches(
+                    snap,
+                    draft_query_eager_cache,
+                    g_cache,
+                    expected,
+                    report,
+                    f"{tag}:draft_query_graph_vs_draft_query_eager",
+                )
+                fails += _compare_records(
+                    draft_query_eager[1],
+                    g_rec,
+                    report,
+                    f"{tag}:draft_query_graph_vs_draft_query_eager",
+                )
+                fails += _compare_to_native(
+                    snap,
+                    (native_outs[0], native_recs[0], native_cache),
+                    (g_out, g_rec, g_cache),
+                    accepted,
+                    nreq,
+                    bounds,
+                    report,
+                    f"{tag}:draft_query_graph_vs_native",
+                )
+            else:
+                fails += _compare_caches(
+                    snap, native_cache, g_cache, expected, report, f"{tag}:native_graph_vs_eager"
+                )
+                fails += _compare_records(
+                    native_recs[0], g_rec, report, f"{tag}:native_graph_vs_eager"
+                )
+                fails += _compare_outputs(
+                    native_outs[0], g_out, bounds, report, f"{tag}:native_graph_vs_eager"
+                )
+            for r_absorbed_query, r_draft_query in runners.values():
+                for runner in (r_absorbed_query, r_draft_query):
+                    if runner not in (None, False):
+                        assert int(runner.sync.abs().sum()) == 0
+    print(f"graph nreq={nreq} form={form} graphs={wanted} {report}")
+    assert not fails, fails
+
+
+@torch.inference_mode()
+def test_checks_detect_a_wrong_mtp_query(harness, monkeypatch):
+    """Negative control: draft_absorbed_query packets whose absorbed query is zeroed must fail
+    the absorbed-query / quant_q checks on layer 78."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    h = harness
+    nreq, path = 1, "custom_op"
+    it = h.iteration_inputs(nreq, seed=12)
+    snap = h.snapshot()
+    orig_project = query_proj.AbsorbedQueryProjection.project
+
+    def broken_project(self, *args, **kwargs):
+        packet = orig_project(self, *args, **kwargs)
+        if packet.fused_q is not None:
+            packet.fused_q[..., : self.mla.kv_lora_rank].zero_()
+        return packet
+
+    try:
+        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, path, snap)
+        monkeypatch.setattr(query_proj.AbsorbedQueryProjection, "project", broken_project)
+        _set_features("draft_absorbed_query", "preabsorbed")
+        h.reset_runners()
+        h.restore(snap)
+        _, bad_rec, calls = h.run(nreq, it, path)
+        assert calls == _expected_calls("draft_absorbed_query", nreq)
+        report = {}
+        fails = _compare_records(native_recs[0], bad_rec, report, "bad")
+        print(f"negative control {fails}")
+        for step in range(DRAFT_LEN):
+            assert any(f"L{MTP}.call{step}: absorbed_q" in f for f in fails), fails
+            assert any(f"L{MTP}.call{step}: quant_q" in f for f in fails), fails
+        assert not any(f"L{TARGET}." in f for f in fails), fails
+    finally:
+        monkeypatch.undo()
+        h.restore(snap)
+        _set_features("")
+        h.reset_runners()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py
new file mode 100644
index 000000000000..7c7084e0eeb1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py
@@ -0,0 +1,745 @@
+# 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.
+"""draft_absorbed_query in the MTP draft loop on the DSA decode-metadata path of the served env.
+
+The served GEN environment exports ``TRTLLM_FUSED_DSA_METADATA=1``. TensorRT-LLM
+does not read that name. The fused DSA decode-metadata kernel
+(``dsa/kernels.py::fused_dsa_decode_metadata``) is on by default and is
+turned off only by ``TRTLLM_DISABLE_FUSED_DSA_METADATA``
+(``dsa/metadata.py::_fused_dsa_meta_enabled``, cached per process). On a
+pure-decode step, ``DSAtrtllmAttentionMetadata.on_update_kv_lens`` then fills
+the request index per token, the indexer K-cache slot mappings and the two
+generation indptrs with one launch instead of the eager chain.
+
+One decode iteration (the harness of ``test_draft_query_draft_loop_integration.py``:
+target layer 6, then the native MTP Eagle draft loop over layer 78, real
+weights, FP8 KV cache, two requests with a 2,600-token context) updates the
+DSA decode metadata six times:
+
+* ``u0``: ``on_update_kv_lens()`` at the start of the step (B x 6 rows);
+* ``u1``..``u5``: after each of the five draft steps (B rows), through
+  ``update_for_spec_dec()`` -> ``on_update_kv_lens()`` (native step update),
+  or through draft_loop's one-launch ``draft_step_update`` when the step update of
+  the served draft_loop package is attached to the worker.
+
+``u0`` feeds layer 6 and draft step 0, and ``u<i>`` feeds draft step ``i``.
+Every update is recorded: which path ran (native or draft_loop), how many fused
+metadata kernels it launched, and the values it left in the request index,
+the slot mappings, the indptrs, the KV lengths (1-D and 2-D views) and the
+paged-MQA schedules.
+
+Arms per case (custom-op or ``forward_impl`` attention path; native or draft_loop
+draft-step update):
+
+* native, eager metadata chain (``TRTLLM_DISABLE_FUSED_DSA_METADATA=1``);
+* native, served chain (fused kernel), repeated for the output-noise bound;
+* draft_absorbed_query (layer 78 at every draft step), served chain, then repeated;
+* draft_absorbed_query, eager chain.
+
+Asserted:
+
+* served env: no TensorRT-LLM source or library reads
+  ``TRTLLM_FUSED_DSA_METADATA``; the fused path is on with the served env
+  and off only with the opt-out;
+* the update sequence: served chain -> one fused launch at each native update
+  (six per iteration, or one at ``u0`` with the draft_loop step update, whose five
+  steps then run draft_loop's kernel); eager chain -> no fused launch;
+* every recorded metadata value of every update is identical between the
+  eager and the fused chain, with and without draft_absorbed_query, and with the draft_loop step
+  update;
+* draft_absorbed_query under the served chain against native under the served chain: caches,
+  MLA inputs, indexer inputs and appended indexer keys bit-identical; top-k
+  sets equal; draft-loop top-k reuse at steps 1-4; outputs within the
+  native run-to-run bound; draft_absorbed_query repeat bit-identical; draft_absorbed_query under the
+  eager chain bit-identical to draft_absorbed_query under the served chain;
+* CUDA graph: the whole iteration captured with the served chain (draft_absorbed_query graph
+  and native graph in one pool), replayed 1, 1 and 20 times with fresh
+  inputs and accepted counts: the captured fused launches, the replayed
+  metadata, caches and records equal the eager runs as above;
+* a negative control: a slot mapping corrupted after the fused launch of one
+  draft step is caught in that update's metadata and in the layer-78 indexer
+  K cache.
+"""
+
+import contextlib
+import mmap
+import os
+import re
+from types import SimpleNamespace
+
+import pytest
+import test_draft_query_draft_loop_integration as dq_loop_integration
+import torch
+
+TARGET, MTP, LAYERS = (
+    dq_loop_integration.TARGET,
+    dq_loop_integration.MTP,
+    dq_loop_integration.LAYERS,
+)
+DRAFT_LEN, VERIFY_ROWS = dq_loop_integration.DRAFT_LEN, dq_loop_integration.VERIFY_ROWS
+NATIVE_RUNS = dq_loop_integration.NATIVE_RUNS
+SERVED_ENV = "TRTLLM_FUSED_DSA_METADATA"
+DISABLE_ENV = "TRTLLM_DISABLE_FUSED_DSA_METADATA"
+# Metadata outputs of one update: per token, per sequence.
+TOKEN_FIELDS = ("req_idx_per_token", "slot_mapping_fp8", "slot_mapping_scale")
+SEQ_FIELDS = ("gen_kv_indptr", "gen_cached_token_indptr")
+SCHEDULE_FIELDS = ("scheduler_metadata_buffer", "scheduler_metadata_buffer_full_next_n")
+# One indexer key in the FP8 slot mapping (bytes per token).
+INDEX_KEY_BYTES = 128
+
+pytestmark = dq_loop_integration.pytestmark
+harness = dq_loop_integration.harness
+
+
+def _select_chain(chain: str):
+    """Served env, with the fused metadata chain ("fused") or the opt-out
+    ("eager"). The gate is cached per process; it is reset here and never
+    flipped while a graph is being captured."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import _fused_dsa_meta_enabled
+
+    os.environ[SERVED_ENV] = "1"
+    if chain == "fused":
+        os.environ.pop(DISABLE_ENV, None)
+    else:
+        os.environ[DISABLE_ENV] = "1"
+    _fused_dsa_meta_enabled.cache_clear()
+    assert _fused_dsa_meta_enabled() == (chain == "fused")
+
+
+def _features(step: str):
+    """Feature selections of the native and the draft_absorbed_query arm. The draft_loop arm attaches
+    draft_loop's draft-loop operations to the worker (``_draft_loop_step_update``) instead
+    of selecting ``draft_loop``: the selection would also wrap the draft embedding,
+    which the harness's stub draft model does not have and which plays no
+    part in the decode metadata."""
+    del step
+    return {"draft_absorbed_query": "draft_absorbed_query", "native": ""}
+
+
+def _inputs(h, nreq, seed, accepted=None):
+    it = h.iteration_inputs(nreq, seed=seed, accepted=accepted)
+    # Step-0 rows of the last accepted tokens (the draft_loop step-0 gather).
+    it.gather_ids = (
+        torch.arange(nreq, device="cuda", dtype=torch.int64) * VERIFY_ROWS
+        + it.num_accepted.to(torch.int64)
+        - 1
+    )
+    return it
+
+
+class _MetaLog:
+    """Records every DSA decode-metadata update of the metadata object
+    ``meta``: the native ``on_update_kv_lens`` and draft_loop's step update. Inside a
+    capture the clones are graph-pool tensors that every replay refreshes;
+    the update kinds and launch counts are host values seen at capture.
+
+    ``corrupt_update``: index of the update after whose fused launch the FP8
+    slot mapping is moved by one indexer key (negative control)."""
+
+    def __init__(self, meta, corrupt_update=None):
+        self.meta = meta
+        self.corrupt_update = corrupt_update
+        self.updates = []
+        self.fused_rows = []
+
+    def __enter__(self):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E501
+            draft_loop as draft_loop_mod,
+        )
+        from tensorrt_llm._torch.attention.backends.sparse.dsa import metadata as dsa_metadata
+
+        log = self
+        self._mod = dsa_metadata
+        self._fused = dsa_metadata.fused_dsa_decode_metadata
+        self._ops_cls = draft_loop_mod.DraftLoopOps
+        self._step_update = draft_loop_mod.DraftLoopOps.step_update
+        native_update = self.meta.on_update_kv_lens
+
+        def fused(*args, **kwargs):
+            out = log._fused(*args, **kwargs)
+            log.fused_rows.append(kwargs["num_tokens"])
+            if log.corrupt_update == len(log.updates):
+                args[4].add_(INDEX_KEY_BYTES)  # slot_mapping_fp8[:num_tokens]
+            return out
+
+        def on_update_kv_lens():
+            before = len(log.fused_rows)
+            native_update()
+            log._record("native", len(log.fused_rows) - before)
+
+        def step_update(ops, attn_metadata, *args, **kwargs):
+            before = len(log.fused_rows)
+            result = log._step_update(ops, attn_metadata, *args, **kwargs)
+            if result is not None and attn_metadata is log.meta:
+                log._record("draft_loop", len(log.fused_rows) - before)
+            return result
+
+        dsa_metadata.fused_dsa_decode_metadata = fused
+        draft_loop_mod.DraftLoopOps.step_update = step_update
+        self.meta.on_update_kv_lens = on_update_kv_lens
+        return self
+
+    def __exit__(self, *exc):
+        self._mod.fused_dsa_decode_metadata = self._fused
+        self._ops_cls.step_update = self._step_update
+        del self.meta.on_update_kv_lens
+        return False
+
+    def _record(self, kind, fused):
+        m = self.meta
+        n, s, g = m.num_tokens, m.num_seqs, m.num_generations
+        fields = {k: getattr(m, k)[:n].clone() for k in TOKEN_FIELDS}
+        fields.update({k: getattr(m, k)[: s + 1].clone() for k in SEQ_FIELDS})
+        fields["kv_lens_cuda"] = m.kv_lens_cuda[:s].clone()
+        width = min(m.gen_token_stride, m.kv_lens_cuda_2d.shape[1])
+        fields["kv_lens_cuda_2d"] = m.kv_lens_cuda_2d[:g, :width].clone()
+        for k in SCHEDULE_FIELDS:
+            t = getattr(m, k, None)
+            if torch.is_tensor(t):
+                fields[k] = t.clone()
+        self.updates.append(SimpleNamespace(kind=kind, rows=n, seqs=s, fused=fused, fields=fields))
+
+    def snapshot(self):
+        return [
+            SimpleNamespace(
+                kind=u.kind,
+                rows=u.rows,
+                seqs=u.seqs,
+                fused=u.fused,
+                fields={k: v.clone() for k, v in u.fields.items()},
+            )
+            for u in self.updates
+        ]
+
+
+@contextlib.contextmanager
+def _draft_loop_step_update(h, meta, it, enabled):
+    """Attach draft_loop's draft-loop operations to the worker, as the served
+    ``_forward_impl`` does for an eligible forward with ``draft_loop`` selected. The
+    harness has no target sampling and no draft embedding, so only draft_loop's
+    step-0 row gather and its per-step metadata update run; acceptance, the
+    draft-token stack and the embedding stay native."""
+    if not enabled:
+        yield None
+        return
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E501
+        draft_loop as draft_loop_mod,
+    )
+
+    ops = draft_loop_mod.DraftLoopOps(h.worker, meta, SimpleNamespace(runtime_draft_len=DRAFT_LEN))
+    ops.gather_ids = it.gather_ids
+    h.worker._draft_loop_ops = ops
+    try:
+        yield ops
+    finally:
+        h.worker._draft_loop_ops = None
+
+
+def _run(h, nreq, it, features, step, chain, path="custom_op", corrupt_update=None):
+    """One eager iteration from the current caches."""
+    _select_chain(chain)
+    dq_loop_integration._set_features(features)
+    h.reset_runners()
+    meta = h.gen_metadata(nreq)
+    with (
+        _MetaLog(meta, corrupt_update) as ml,
+        _draft_loop_step_update(h, meta, it, step == "draft_loop") as ops,
+        dq_loop_integration._LaunchLog() as launches,
+    ):
+        out, rec, calls = h.run(nreq, it, path, meta=meta)
+    return SimpleNamespace(
+        out=out,
+        rec=rec,
+        calls=calls,
+        meta=ml.updates,
+        launches=launches.calls,
+        cache=h.snapshot(),
+        draft_loop_engaged=ops is not None and ops.step_engaged,
+    )
+
+
+def _expected_updates(nreq, step, fused):
+    """``(kind, rows, fused launches)`` of the six metadata updates."""
+    draft_loop_pkg = step == "draft_loop" and fused
+    first = ("native", nreq * VERIFY_ROWS, int(fused))
+    later = ("draft_loop", nreq, 0) if draft_loop_pkg else ("native", nreq, int(fused))
+    return [first] + [later] * DRAFT_LEN
+
+
+def _check_updates(updates, nreq, step, fused, tag):
+    got = [(u.kind, u.rows, u.fused) for u in updates]
+    want = _expected_updates(nreq, step, fused)
+    if got != want or any(u.seqs != nreq for u in updates):
+        return [f"{tag}: metadata updates {got}, expected {want}"]
+    return []
+
+
+def _compare_meta(ref, cand, report, tag):
+    """Every recorded value of every update bit-identical."""
+    if [(u.rows, u.seqs) for u in ref] != [(u.rows, u.seqs) for u in cand]:
+        return [
+            f"{tag}: metadata updates {[(u.rows, u.seqs) for u in cand]}"
+            f" vs {[(u.rows, u.seqs) for u in ref]}"
+        ]
+    fails = []
+    for i, (r, c) in enumerate(zip(ref, cand)):
+        if set(r.fields) != set(c.fields):
+            fails.append(f"{tag}: u{i} fields {sorted(c.fields)} vs {sorted(r.fields)}")
+            continue
+        for key, a in r.fields.items():
+            b = c.fields[key]
+            same = a.dtype == b.dtype and a.shape == b.shape and torch.equal(a, b)
+            report[f"{tag}:u{i}.{key}_same"] = same
+            if not same:
+                fails.append(f"{tag}: u{i} {key} differs")
+    return fails
+
+
+def _summary(report):
+    """Counters of a report: byte-difference counters (all must be 0),
+    key-set and metadata comparisons (all must hold)."""
+    nbytes = [v for k, v in report.items() if k.endswith("_bytes_differing")]
+    sets = [v for k, v in report.items() if k.endswith("_same_sets")]
+    meta = [v for k, v in report.items() if k.endswith("_same")]
+    return (
+        f"byte counters {len(nbytes)} (nonzero "
+        f"{sum(1 for v in nbytes if v)}), key-set comparisons "
+        f"{len(sets)} (unequal {sum(1 for v in sets if not v)}), "
+        f"metadata comparisons {len(meta)} (unequal "
+        f"{sum(1 for v in meta if not v)})"
+    )
+
+
+def _repeat_native(h, nreq, it, features, step, snap, path="custom_op"):
+    """Native iterations under the served chain from one snapshot; they must
+    repeat bit for bit in caches, recorded inputs and metadata."""
+    runs, fails = [], []
+    accepted = it.num_accepted.tolist()
+    expected = dq_loop_integration._expected_slots(nreq, accepted)
+    for i in range(NATIVE_RUNS):
+        h.restore(snap)
+        run = _run(h, nreq, it, features, step, "fused", path)
+        assert run.calls == [], f"native run called absorbed_query entry points: {run.calls}"
+        report = {}
+        fails += _check_updates(run.meta, nreq, step, True, f"native{i}")
+        fails += dq_loop_integration._check_reuse(run.rec, accepted, report, f"native{i}")
+        if i:
+            fails += dq_loop_integration._compare_caches(
+                snap, runs[0].cache, run.cache, expected, report, f"native_repeat{i}"
+            )
+            fails += dq_loop_integration._compare_records(
+                runs[0].rec, run.rec, report, f"native_repeat{i}"
+            )
+            fails += _compare_meta(runs[0].meta, run.meta, report, f"native_repeat{i}")
+        runs.append(run)
+    h.restore(snap)
+    assert not fails, f"native iteration is not repeatable: {fails}"
+    return runs
+
+
+def _output_bounds(native_outs):
+    """Per (layer, call) output bounds from native iterations of one input:
+    twice the largest difference over all pairs of native runs (at least
+    one BF16 ulp, relative RMS at least ``REL_RMS_FLOOR``).
+
+    The attention inputs of these runs are bit-identical (checked
+    separately), so their outputs differ only by the native layer-78
+    attention's run-to-run nondeterminism, which leaves some runs
+    identical and others about one BF16 ulp apart. Every pair samples it;
+    differences from one reference run alone underestimate it."""
+    aq_integration = dq_loop_integration.aq_integration
+    bounds = {}
+    for layer in LAYERS:
+        for i in range(len(native_outs[0][layer])):
+            outs = [run[layer][i] for run in native_outs]
+            noise, rel = 0.0, 0.0
+            for a in range(len(outs)):
+                for b in range(a + 1, len(outs)):
+                    noise = max(noise, (outs[a].float() - outs[b].float()).abs().max().item())
+                    rel = max(rel, aq_integration._rel_rms(outs[a], outs[b]))
+            ulp = aq_integration._bf16_ulp(outs[0].float().abs().max().item())
+            bounds[(layer, i)] = (
+                2.0 * max(noise, ulp),
+                max(2.0 * rel, aq_integration.REL_RMS_FLOOR),
+                noise,
+            )
+    return bounds
+
+
+def _restore_defaults(h, snap):
+    h.worker._draft_loop_ops = None
+    h.restore(snap)
+    dq_loop_integration._set_features("")
+    h.reset_runners()
+    _select_chain("fused")
+
+
+# -- tests ------------------------------------------------------------------
+
+
+def _names_served_env(data) -> bool:
+    """Whether ``data`` contains ``TRTLLM_FUSED_DSA_METADATA`` as a whole
+    name (not as the tail of ``TRTLLM_DISABLE_FUSED_DSA_METADATA``)."""
+    name = SERVED_ENV.encode()
+    start = 0
+    while True:
+        i = data.find(name, start)
+        if i < 0:
+            return False
+        if i == 0 or not re.match(rb"[A-Za-z0-9_]", data[i - 1 : i]):
+            return True
+        start = i + 1
+
+
+def test_served_env_selects_fused_metadata():
+    """The served ``TRTLLM_FUSED_DSA_METADATA=1`` is read nowhere in the
+    installed TensorRT-LLM (Python sources and shared libraries); the fused
+    chain is on with or without it and off only with the opt-out."""
+    import tensorrt_llm
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import _fused_dsa_meta_enabled
+
+    root = os.path.dirname(tensorrt_llm.__file__)
+    scanned, readers = {"py": 0, "so": 0}, []
+    for dirpath, _, files in os.walk(root):
+        for name in files:
+            kind = "py" if name.endswith(".py") else "so" if ".so" in name else None
+            if kind is None:
+                continue
+            path = os.path.join(dirpath, name)
+            with open(path, "rb") as fh:
+                if os.fstat(fh.fileno()).st_size == 0:
+                    continue
+                with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as data:
+                    scanned[kind] += 1
+                    if _names_served_env(data):
+                        readers.append(os.path.relpath(path, root))
+    print(f"scanned {scanned} under {root}; readers of {SERVED_ENV}: {readers}")
+    assert scanned["py"] > 1000 and scanned["so"] > 0, scanned
+    assert readers == [], readers
+
+    saved = {k: os.environ.get(k) for k in (SERVED_ENV, DISABLE_ENV)}
+    try:
+        os.environ.pop(DISABLE_ENV, None)
+        for served in ("1", "0", None):
+            if served is None:
+                os.environ.pop(SERVED_ENV, None)
+            else:
+                os.environ[SERVED_ENV] = served
+            _fused_dsa_meta_enabled.cache_clear()
+            assert _fused_dsa_meta_enabled(), served
+        os.environ[SERVED_ENV] = "1"
+        os.environ[DISABLE_ENV] = "1"
+        _fused_dsa_meta_enabled.cache_clear()
+        assert not _fused_dsa_meta_enabled()
+    finally:
+        for k, v in saved.items():
+            if v is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = v
+        _fused_dsa_meta_enabled.cache_clear()
+
+
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("step", ["native", "draft_loop"])
+@pytest.mark.parametrize("nreq", [1, 2])
+@torch.inference_mode()
+def test_draft_query_draft_loop_on_fused_metadata(harness, nreq, step, path):
+    """draft_absorbed_query on layer 78 (all five draft steps) with the served metadata chain,
+    against native on the served chain and against the eager chain."""
+    h = harness
+    feats = _features(step)
+    it = _inputs(h, nreq, seed=40 + nreq)
+    accepted = it.num_accepted.tolist()
+    expected = dq_loop_integration._expected_slots(nreq, accepted)
+    snap = h.snapshot()
+    try:
+        natives = _repeat_native(h, nreq, it, feats["native"], step, snap, path)
+        native = natives[0]
+        bounds = _output_bounds([r.out for r in natives])
+        h.restore(snap)
+        chain = _run(h, nreq, it, feats["native"], step, "eager", path)
+        h.restore(snap)
+        fused_run = _run(h, nreq, it, feats["draft_absorbed_query"], step, "fused", path)
+        h.restore(snap)
+        draft_query_again = _run(h, nreq, it, feats["draft_absorbed_query"], step, "fused", path)
+        h.restore(snap)
+        draft_query_chain = _run(h, nreq, it, feats["draft_absorbed_query"], step, "eager", path)
+
+        assert fused_run.calls == dq_loop_integration._expected_calls(
+            "draft_absorbed_query", nreq
+        ), fused_run.calls
+        assert fused_run.launches == dq_loop_integration._expected_launches(nreq), (
+            fused_run.launches
+        )
+        assert draft_query_chain.calls == fused_run.calls and chain.calls == []
+        if step == "draft_loop":
+            assert native.draft_loop_engaged and fused_run.draft_loop_engaged
+
+        report, fails = {}, []
+        for tag, run, fused in (
+            ("native_fused", native, True),
+            ("native_eager", chain, False),
+            ("draft_query_fused", fused_run, True),
+            ("draft_query_again", draft_query_again, True),
+            ("draft_query_eager", draft_query_chain, False),
+        ):
+            fails += _check_updates(run.meta, nreq, step, fused, tag)
+        # Eager vs fused chain, native: metadata, caches and inputs.
+        fails += _compare_meta(chain.meta, native.meta, report, "native_fused_vs_eager")
+        fails += dq_loop_integration._compare_caches(
+            snap, chain.cache, native.cache, expected, report, "native_fused_vs_eager"
+        )
+        fails += dq_loop_integration._compare_records(
+            chain.rec, native.rec, report, "native_fused_vs_eager"
+        )
+        fails += dq_loop_integration._compare_outputs(
+            native.out, chain.out, bounds, report, "native_eager_vs_fused"
+        )
+        # draft_absorbed_query vs native, both on the served chain.
+        fails += _compare_meta(native.meta, fused_run.meta, report, "draft_query_vs_native")
+        fails += dq_loop_integration._compare_to_native(
+            snap,
+            (native.out, native.rec, native.cache),
+            (fused_run.out, fused_run.rec, fused_run.cache),
+            accepted,
+            nreq,
+            bounds,
+            report,
+            "draft_query_vs_native",
+        )
+        # draft_absorbed_query determinism and draft_absorbed_query eager chain vs served chain.
+        for tag, other in (
+            ("draft_query_again", draft_query_again),
+            ("draft_query_eager", draft_query_chain),
+        ):
+            fails += _compare_meta(fused_run.meta, other.meta, report, tag)
+            fails += dq_loop_integration._compare_caches(
+                snap, fused_run.cache, other.cache, expected, report, tag
+            )
+            fails += dq_loop_integration._compare_records(fused_run.rec, other.rec, report, tag)
+            fails += dq_loop_integration._check_reuse(other.rec, accepted, report, tag)
+        print(
+            f"nreq={nreq} step={step} path={path} accepted={accepted} "
+            f"updates={[(u.kind, u.rows, u.fused) for u in fused_run.meta]} "
+            f"summary: {_summary(report)}"
+        )
+        print(f"report {report}")
+        assert not fails, fails
+    finally:
+        _restore_defaults(h, snap)
+
+
+@pytest.mark.parametrize("step", ["native", "draft_loop"])
+@pytest.mark.parametrize("nreq", [1, 2])
+@torch.inference_mode()
+def test_draft_query_fused_metadata_cuda_graph(harness, nreq, step):
+    """The whole iteration with the served chain captured in one CUDA graph
+    per selection (draft_absorbed_query and native, one pool), replayed 1, 1 and 20 times with
+    fresh inputs and accepted counts (``prepare()`` before every replay)."""
+    h = harness
+    snap = h.snapshot()
+    graphs = {}
+    try:
+        _graph_checks(h, nreq, step, snap, graphs)
+    except BaseException as exc:
+        print(f"graph test failure: {exc!r}", flush=True)
+        raise
+    finally:
+        graphs.clear()
+        _restore_defaults(h, snap)
+
+
+def _graph_checks(h, nreq, step, snap, graphs):
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    path = "custom_op"
+    feats = _features(step)
+    _select_chain("fused")
+    warm = {}
+    for name, sel in feats.items():
+        dq_loop_integration._set_features(sel)
+        h.reset_runners()
+        meta = h.gen_metadata(nreq, cuda_graph=True)
+        static = _inputs(h, nreq, seed=200)
+        with with_multi_stream(True):
+            for _ in range(2):
+                meta.prepare()
+                with _draft_loop_step_update(h, meta, static, step == "draft_loop"):
+                    h.iteration(meta, static, path)
+        torch.cuda.synchronize()
+        runners = {
+            layer: (h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query)
+            for layer in LAYERS
+        }
+        if name == "draft_absorbed_query":
+            assert runners[MTP][1] not in (None, False), "draft_absorbed_query did not engage"
+        warm[name] = (meta, static, runners)
+    h.restore(snap)
+    pool = None
+    for name, sel in feats.items():
+        meta, static, runners = warm[name]
+        dq_loop_integration._set_features(sel)
+        for layer, (r_absorbed_query, r_draft_query) in runners.items():
+            h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query = (
+                r_absorbed_query,
+                r_draft_query,
+            )
+        meta.prepare()
+        graph = torch.cuda.CUDAGraph()
+        with (
+            with_multi_stream(True),
+            dq_loop_integration._Recorder(h.attn) as rec,
+            dq_loop_integration._Engagement() as eng,
+            dq_loop_integration._LaunchLog() as launches,
+            _MetaLog(meta) as ml,
+            _draft_loop_step_update(h, meta, static, step == "draft_loop") as ops,
+        ):
+            with torch.cuda.graph(graph, pool=pool):
+                outs = h.iteration(meta, static, path)
+        pool = graph.pool()
+        assert eng.calls == dq_loop_integration._expected_calls(sel, nreq), eng.calls
+        if name == "draft_absorbed_query":
+            assert launches.calls == dq_loop_integration._expected_launches(nreq), launches.calls
+        if step == "draft_loop":
+            assert ops.step_engaged, "draft_loop step update declined at capture"
+        captured = _check_updates(ml.updates, nreq, step, True, f"capture:{name}")
+        assert not captured, captured
+        print(
+            f"[fused-meta graph] captured {name}: updates "
+            f"{[(u.kind, u.rows, u.fused) for u in ml.updates]}",
+            flush=True,
+        )
+        graphs[name] = (graph, meta, static, outs, rec, ml, runners)
+        torch.cuda.synchronize()
+    h.reset_runners()
+
+    report, fails = {}, []
+    trials = [(311, 1, None), (312, 1, [1] * nreq), (313, 20, [VERIFY_ROWS] * nreq)]
+    for trial, (seed, replays, accepted) in enumerate(trials):
+        it = _inputs(h, nreq, seed=seed, accepted=accepted)
+        accepted = it.num_accepted.tolist()
+        expected = dq_loop_integration._expected_slots(nreq, accepted)
+        natives = _repeat_native(h, nreq, it, feats["native"], step, snap, path)
+        native = natives[0]
+        bounds = _output_bounds([r.out for r in natives])
+        h.restore(snap)
+        chain = _run(h, nreq, it, feats["native"], step, "eager", path)
+        h.restore(snap)
+        draft_query_eager = _run(h, nreq, it, feats["draft_absorbed_query"], step, "fused", path)
+        _select_chain("fused")
+        dq_loop_integration._set_features("")
+        h.reset_runners()
+        tag = f"t{trial}x{replays}"
+        for name, (graph, meta, static, outs, rec, ml, runners) in graphs.items():
+            h.restore(snap)
+            meta.prepare()
+            for key in ("hidden_t", "pos_t", "hidden_d", "num_accepted", "input_ids", "gather_ids"):
+                getattr(static, key).copy_(getattr(it, key))
+            for i in range(replays):
+                if i:
+                    meta.prepare()
+                graph.replay()
+            torch.cuda.synchronize()
+            g_out = {layer: [o.clone() for o in outs[layer]] for layer in LAYERS}
+            g_rec, g_meta, g_cache = rec.snapshot(), ml.snapshot(), h.snapshot()
+            fails += dq_loop_integration._check_reuse(g_rec, accepted, report, f"{tag}:{name}")
+            fails += _compare_meta(
+                native.meta, g_meta, report, f"{tag}:{name}_graph_vs_native_fused"
+            )
+            fails += _compare_meta(
+                chain.meta, g_meta, report, f"{tag}:{name}_graph_vs_native_eager"
+            )
+            if name == "draft_absorbed_query":
+                fails += dq_loop_integration._compare_caches(
+                    snap,
+                    draft_query_eager.cache,
+                    g_cache,
+                    expected,
+                    report,
+                    f"{tag}:draft_query_graph_vs_draft_query_eager",
+                )
+                fails += dq_loop_integration._compare_records(
+                    draft_query_eager.rec,
+                    g_rec,
+                    report,
+                    f"{tag}:draft_query_graph_vs_draft_query_eager",
+                )
+                fails += dq_loop_integration._compare_to_native(
+                    snap,
+                    (native.out, native.rec, native.cache),
+                    (g_out, g_rec, g_cache),
+                    accepted,
+                    nreq,
+                    bounds,
+                    report,
+                    f"{tag}:draft_query_graph_vs_native",
+                )
+            else:
+                fails += dq_loop_integration._compare_caches(
+                    snap, native.cache, g_cache, expected, report, f"{tag}:native_graph_vs_eager"
+                )
+                fails += dq_loop_integration._compare_records(
+                    native.rec, g_rec, report, f"{tag}:native_graph_vs_eager"
+                )
+                fails += dq_loop_integration._compare_outputs(
+                    native.out, g_out, bounds, report, f"{tag}:native_graph_vs_eager"
+                )
+            for r_absorbed_query, r_draft_query in runners.values():
+                for runner in (r_absorbed_query, r_draft_query):
+                    if runner not in (None, False):
+                        assert int(runner.sync.abs().sum()) == 0
+            print(
+                f"[fused-meta graph] trial {trial} {name}: replayed "
+                f"x{replays}, accepted={accepted}",
+                flush=True,
+            )
+    print(f"graph nreq={nreq} step={step} summary: {_summary(report)}")
+    print(f"report {report}")
+    assert not fails, fails
+
+
+@torch.inference_mode()
+def test_checks_detect_corrupted_draft_step_metadata(harness):
+    """Negative control: the FP8 slot mapping of the update that feeds draft
+    step 2 (``u2``), moved by one indexer key after the fused launch, must be
+    flagged in that update only and in the layer-78 indexer K cache only."""
+    h = harness
+    nreq = 1
+    it = _inputs(h, nreq, seed=12)
+    accepted = it.num_accepted.tolist()
+    snap = h.snapshot()
+    try:
+        h.restore(snap)
+        native = _run(h, nreq, it, "", "native", "fused")
+        h.restore(snap)
+        bad = _run(h, nreq, it, "draft_absorbed_query", "native", "fused", corrupt_update=2)
+        assert bad.calls == dq_loop_integration._expected_calls("draft_absorbed_query", nreq)
+        report = {}
+        meta_fails = _compare_meta(native.meta, bad.meta, report, "bad")
+        cache_fails = dq_loop_integration._compare_caches(
+            snap,
+            native.cache,
+            bad.cache,
+            dq_loop_integration._expected_slots(nreq, accepted),
+            report,
+            "bad",
+        )
+        print(f"negative control: metadata {meta_fails}; caches {cache_fails}")
+        assert meta_fails == ["bad: u2 slot_mapping_fp8 differs"], meta_fails
+        assert any(f"index_k{MTP}" in f for f in cache_fails), cache_fails
+        assert not any(f"index_k{TARGET}" in f or "latent" in f for f in cache_fails), cache_fails
+    finally:
+        _restore_defaults(h, snap)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py
new file mode 100644
index 000000000000..a6bae5b8fc1d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py
@@ -0,0 +1,439 @@
+# 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.
+"""draft_absorbed_query: the absorbed_query query_front kernel on the MTP draft layer, against native.
+
+Fixtures: the real GLM-5.2 checkpoint weights of the one-model MTP layer
+(``model.layers.78.self_attn.*``, BF16, target-layer shapes, full indexer)
+at TP8 rank 0, placed exactly as for target layers (``absorbed_query_fixture``), and the
+layer's real ``input_layernorm`` gain; synthetic activations. The module is
+the native ``DeepseekV32Attention`` with ``layer_idx = 78``, as
+``DeepseekV3MTP`` builds it in the one-model engine.
+
+Draft-step shapes of the MTP Eagle loop at batch 1: step 0 presents the six
+verify rows (``R = 6``, positions ``p .. p + 5``), steps 1-4 one row each
+(``R = 1``, position ``p + accepted + i - 1``).
+
+Asserted, with the absorbed_query checks (``test_absorbed_query_proj``):
+
+* eligibility: absorbed_query's ``dsa_proj`` still declines layer 78 ("target layers
+  only"); ``mtp_dsa_proj`` engages on layer 78 only, and only for
+  generation-only calls inside the absorbed_query domain;
+* the frozen absorbed_query error budget and bit identity with native for every
+  output (``aq_proj._native_agreement``) for R in {1, 6} and the other absorbed_query row
+  counts, both forms, fused and native indexer tail;
+* padded rows (whole padded requests, as CUDA-graph padding adds them) never
+  change live rows;
+* the draft loop's five calls of one layer captured back to back in one CUDA
+  graph, sharing the layer's sync counters (also with delayed producers /
+  consumers), reproduce eager absorbed_query calls bit for bit on 1, 1 and 20 replays
+  and leave the counters at zero.
+"""
+
+import os
+
+import absorbed_query_fixture
+import pytest
+import test_absorbed_query_proj as aq_proj
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+MTP_LAYER = 78
+# Context length at the start of the decode iteration (the controlled
+# decode workload) and draft-loop shape at batch 1: R_d = [6, 1, 1, 1, 1].
+CTX = 491061
+DRAFT_LEN = 5
+ROWS = aq_proj._env_list("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_TEST_ROWS", [1, 6, 2, 3, 4, 5, 8, 16])
+
+_ATTN = {}
+
+
+@pytest.fixture(scope="module")
+def model_config():
+    return absorbed_query_fixture.build_model_config(0)
+
+
+def _attn(model_config, layer=MTP_LAYER):
+    if layer not in _ATTN:
+        _ATTN[layer] = absorbed_query_fixture.build_native_attention(layer, 0, model_config)
+    return _ATTN[layer]
+
+
+def _gamma():
+    return absorbed_query_fixture.load_input_layernorm(MTP_LAYER)
+
+
+def _hidden(rows, seed):
+    return absorbed_query_fixture.make_hidden(rows, seed=seed, gamma=_gamma())
+
+
+def _draft_step_inputs(step, accepted=3, seed=0, batch=1):
+    """Hidden rows and positions of draft step ``step`` (batch requests)."""
+    if step == 0:
+        rows = batch * (DRAFT_LEN + 1)
+        pos = (CTX + torch.arange(DRAFT_LEN + 1, dtype=torch.int32)).repeat(batch)
+    else:
+        rows = batch
+        pos = torch.full((batch,), CTX + accepted + step - 1, dtype=torch.int32)
+    # Step >= 1 position ids are 1-D (one per request), as the draft loop
+    # builds them; step 0 reuses the target's [1, num_tokens] position ids.
+    pos = pos.cuda()
+    if step == 0:
+        pos = pos.unsqueeze(0)
+    return _hidden(rows, seed=seed + 31 * step), pos
+
+
+@torch.inference_mode()
+def test_mtp_layer_fixture(model_config):
+    """Layer 78 is the MTP layer: full indexer with draft-loop top-k reuse,
+    BF16 attention weights of the target-layer shapes."""
+    attn = _attn(model_config)
+    target = _attn(model_config, absorbed_query_fixture.FULL_LAYER)
+    assert attn.layer_idx == MTP_LAYER == attn.num_hidden_layers
+    indexer = attn.mqa.indexer
+    assert indexer is not None and indexer.mtp_index_share
+    for name in ("kv_a_proj_with_mqa", "q_b_proj"):
+        a, t = getattr(attn, name).weight, getattr(target, name).weight
+        assert a.dtype == t.dtype == torch.bfloat16 and a.shape == t.shape
+        assert not torch.equal(a, t), f"{name}: layer 78 weights not loaded"
+    assert attn.k_b_proj_trans.shape == target.k_b_proj_trans.shape
+    assert indexer.wq_b.weight.shape == target.mqa.indexer.wq_b.weight.shape
+
+
+@torch.inference_mode()
+def test_eligibility(model_config):
+    """absorbed_query declines layer 78 as before; draft_absorbed_query engages on layer 78 only, for
+    generation-only calls inside the absorbed_query domain, before any GPU work."""
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    mtp = _attn(model_config)
+    target = _attn(model_config, absorbed_query_fixture.FULL_LAYER)
+    assert query_proj.module_eligibility(mtp) == (False, "target layers only")
+    assert query_proj.module_eligibility(mtp, draft=True) == (True, "")
+    assert query_proj.module_eligibility(target) == (True, "")
+    assert query_proj.module_eligibility(target, draft=True) == (False, "draft (MTP) layers only")
+
+    gen = SimpleNamespace(num_contexts=0, num_generations=1)
+    ctx = SimpleNamespace(num_contexts=1, num_generations=0)
+    mixed = SimpleNamespace(num_contexts=1, num_generations=1)
+    for step in range(DRAFT_LEN):
+        hidden, pos = _draft_step_inputs(step)
+        for attr in ("_absorbed_query", "_draft_absorbed_query"):
+            for mod in (mtp, target):
+                setattr(mod, attr, None)
+        assert query_proj.dsa_proj(mtp, pos, hidden, gen) is None
+        assert mtp._absorbed_query is False
+        assert query_proj.mtp_dsa_proj(target, pos, hidden, gen) is None
+        assert target._draft_absorbed_query is False
+        packet = query_proj.mtp_dsa_proj(mtp, pos, hidden, gen)
+        assert packet is not None and len(packet.proj_outputs) == 9
+        assert packet.fused_q is not None  # auto form: preabsorbed (full)
+        for meta in (None, ctx, mixed):
+            assert query_proj.mtp_dsa_proj(mtp, pos, hidden, meta) is None
+    big = torch.zeros(17, 6144, dtype=torch.bfloat16, device="cuda")
+    assert query_proj.mtp_dsa_proj(mtp, None, big, gen) is None
+    assert query_proj.mtp_dsa_proj(mtp, None, big[:6].float(), gen) is None
+    for mod in (mtp, target):
+        mod._absorbed_query = mod._draft_absorbed_query = None
+
+
+@pytest.mark.parametrize("rows", ROWS)
+@pytest.mark.parametrize("absorb", [True, False])
+@pytest.mark.parametrize("fuse_indexer", [True, False])
+@torch.inference_mode()
+def test_mtp_contract_q_matches_native(model_config, rows, absorb, fuse_indexer):
+    """The query_front contract on layer 78 against the native region, absorbed_query budget and
+    bitwise agreement; R = 6 and R = 1 are the draft-loop shapes."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    attn = _attn(model_config)
+    hidden = _hidden(rows, seed=rows)
+    pos = (CTX + torch.arange(rows, dtype=torch.int32, device="cuda")).unsqueeze(0)
+    native_outs, native_fq = aq_proj._native(attn, hidden, pos)
+    runner = AbsorbedQueryProjection(attn, KernelConfig(fuse_indexer=fuse_indexer))
+    assert runner.indexer_fusable() == fuse_indexer
+    packet = runner.project(pos, hidden, absorb=absorb)
+    torch.cuda.synchronize()
+    ref = aq_proj._oracle(attn, hidden)
+    outs = packet.proj_outputs
+    report = {}
+    assert len(outs) == len(native_outs) == 9
+    for cand, nat in zip(outs, native_outs):
+        assert cand.shape == nat.shape and cand.dtype == nat.dtype
+    assert (packet.fused_q is not None) == absorb
+    assert torch.equal(outs[3][:, : attn.kv_lora_rank], outs[1])
+    assert torch.equal(outs[3][:, attn.kv_lora_rank :], outs[2])
+    aq_proj._check_budget("q", outs[0], native_outs[0], ref["q"], report)
+    aq_proj._check_budget("compressed_kv", outs[1], native_outs[1], ref["ckv"], report)
+    if absorb:
+        aq_proj._check_budget(
+            "absorbed",
+            packet.fused_q[..., : attn.kv_lora_rank],
+            native_fq[..., : attn.kv_lora_rank],
+            ref["absorbed"],
+            report,
+        )
+    q_fp8, k_fp8, k_scale, weights, q_scale = outs[4:]
+    nq_fp8, nk_fp8, nk_scale, nweights, nq_scale = native_outs[4:]
+    torch.testing.assert_close(
+        k_fp8.float() * k_scale, nk_fp8.float() * nk_scale, rtol=0.1, atol=1e-2
+    )
+    torch.testing.assert_close(
+        q_fp8.float() * q_scale, nq_fp8.float() * nq_scale, rtol=0.1, atol=1e-2
+    )
+    torch.testing.assert_close(weights, nweights, rtol=0.1, atol=1e-2)
+    stats, _ = aq_proj._native_agreement(
+        outs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
+    )
+    print(
+        f"layer={MTP_LAYER} rows={rows} absorb={absorb} "
+        f"fuse_indexer={fuse_indexer} budget={report} agreement={stats}"
+    )
+    assert int(runner.sync.abs().sum()) == 0
+
+
+@pytest.mark.parametrize("step", list(range(DRAFT_LEN)))
+@torch.inference_mode()
+def test_mtp_entry_point_per_draft_step(model_config, step):
+    """``mtp_dsa_proj`` at each draft step's shape and positions equals the
+    default absorbed_query runner and native bit for bit."""
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    attn = _attn(model_config)
+    attn._draft_absorbed_query = None
+    gen = SimpleNamespace(num_contexts=0, num_generations=1)
+    hidden, pos = _draft_step_inputs(step, seed=40)
+    packet = query_proj.mtp_dsa_proj(attn, pos, hidden, gen)
+    reference = query_proj.AbsorbedQueryProjection(attn).project(
+        pos, hidden, absorb=query_proj.absorb_for(attn)
+    )
+    native_outs, native_fq = aq_proj._native(attn, hidden, pos)
+    torch.cuda.synchronize()
+    assert packet.proj_outputs[0].shape[0] == (6 if step == 0 else 1)
+    aq_proj._assert_same_packet(packet, reference, f"step {step}")
+    aq_proj._native_agreement(
+        packet.proj_outputs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
+    )
+    assert int(attn._draft_absorbed_query.sync.abs().sum()) == 0
+    attn._draft_absorbed_query = None
+
+
+@pytest.mark.parametrize("rows", [6, 1])
+@torch.inference_mode()
+def test_mtp_native_agreement_statistics(model_config, rows):
+    """Default kernel on layer 78 over many inputs: per-output difference
+    statistics against native; every input must agree bit for bit."""
+    import json
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    attn = _attn(model_config)
+    indexer = attn.mqa.indexer
+    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_keep_intermediates=True))
+    n_inputs = int(os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_AGREEMENT_INPUTS", "48"))
+    totals, inputs_with_diff, fails = {}, {}, []
+    for i in range(n_inputs):
+        hidden = _hidden(rows, seed=2000 + 17 * i)
+        pos = (CTX + i + torch.arange(rows, dtype=torch.int32, device="cuda")).unsqueeze(0)
+        native_outs, native_fq = aq_proj._native(attn, hidden, pos)
+        packet = runner.project(pos, hidden, absorb=True)
+        stats, input_fails = aq_proj._native_agreement(
+            packet.proj_outputs,
+            packet.fused_q,
+            native_outs,
+            native_fq,
+            attn.kv_lora_rank,
+            check=False,
+        )
+        kva = attn.kv_a_proj_with_mqa(hidden)
+        qr = attn._q_a_layernorm_maybe_fused(kva[:, : attn.q_lora_rank])
+        steps = aq_proj._steps(runner.debug["iq"], indexer.wq_b(qr))
+        stats["indexer_wq_b_bf16"] = (int((steps != 0).sum()), steps.numel(), int(steps.max()))
+        fails += [f"input {i}: {f}" for f in input_fails]
+        for name, (n_diff, n, max_steps) in stats.items():
+            tot = totals.setdefault(name, [0, 0, 0])
+            tot[0] += n_diff
+            tot[1] += n
+            tot[2] = max(tot[2], max_steps)
+            inputs_with_diff[name] = inputs_with_diff.get(name, 0) + (n_diff > 0)
+    summary = {
+        name: dict(
+            inputs_with_diff=inputs_with_diff[name],
+            inputs=n_inputs,
+            diff_elements=tot[0],
+            elements=tot[1],
+            max_steps=tot[2],
+        )
+        for name, tot in totals.items()
+    }
+    print(f"AGREEMENT layer={MTP_LAYER} rows={rows} {json.dumps(summary)}")
+    assert not fails, fails
+    assert int(runner.sync.abs().sum()) == 0
+
+
+@pytest.mark.parametrize("rows,live", [(12, 6), (6, 3), (2, 1)])
+@pytest.mark.parametrize("absorb", [True, False])
+@torch.inference_mode()
+def test_mtp_padded_rows_do_not_affect_live_rows(model_config, rows, live, absorb):
+    """Graph padding adds whole dummy requests: 12 = 2 x 6 step-0 rows with
+    one live request, 2 step >= 1 rows with one live request. Padded rows
+    (zero or 1e3 garbage) never change any live-row output bit for bit, and
+    live rows agree with native bit for bit."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+    )
+
+    attn = _attn(model_config)
+    runner = AbsorbedQueryProjection(attn)
+    live_rows = slice(0, live)
+    hidden_zero = _hidden(rows, seed=77)
+    hidden_zero[live:] = 0
+    hidden_junk = hidden_zero.clone()
+    hidden_junk[live:] = 1e3 * torch.randn_like(hidden_junk[live:])
+    pos = (CTX + torch.arange(rows, dtype=torch.int32, device="cuda") % (DRAFT_LEN + 1)).unsqueeze(
+        0
+    )
+    a = runner.project(pos, hidden_zero, absorb=absorb)
+    b = runner.project(pos, hidden_junk, absorb=absorb)
+    native_outs, native_fq = aq_proj._native(attn, hidden_junk, pos)
+    torch.cuda.synchronize()
+    for name, x, y in zip(aq_proj._NAMES, a.proj_outputs, b.proj_outputs):
+        aq_proj._assert_bitwise(f"padded {name}", x[live_rows], y[live_rows])
+    if absorb:
+        aq_proj._assert_bitwise(
+            "padded absorbed",
+            a.fused_q[live_rows, :, : attn.kv_lora_rank],
+            b.fused_q[live_rows, :, : attn.kv_lora_rank],
+        )
+    aq_proj._native_agreement(
+        b.proj_outputs, b.fused_q, native_outs, native_fq, attn.kv_lora_rank, live_rows
+    )
+    for x in a.proj_outputs:
+        assert torch.isfinite(x[live:].float()).all()
+    assert int(runner.sync.abs().sum()) == 0
+
+
+def _draft_sequence(runner, static_inputs, absorb):
+    """The five draft-step calls of the MTP layer, back to back."""
+    return [runner.project(pos, hidden, absorb=absorb) for hidden, pos in static_inputs]
+
+
+def _check_step_vs_native(attn, packet, hidden, pos):
+    """Every output bit-identical to the native region."""
+    native_outs, native_fq = aq_proj._native(attn, hidden, pos)
+    torch.cuda.synchronize()
+    aq_proj._native_agreement(
+        packet.proj_outputs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
+    )
+
+
+@pytest.mark.parametrize("absorb", [True, False])
+@pytest.mark.parametrize("delay_ns", [0, 20000])
+@torch.inference_mode()
+def test_mtp_draft_loop_graph_replay(model_config, absorb, delay_ns):
+    """One CUDA graph holding the draft loop's five calls of layer 78
+    (R = 6, then four R = 1 calls), all on the same runner and sync
+    counters, replayed 1, 1 and 20 times with fresh inputs: every call equals
+    an eager absorbed_query call and native bit for bit, and
+    the counters are zero afterwards (also with delayed producers and
+    consumers, which stresses counter reuse between adjacent launches)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    attn = _attn(model_config)
+    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_delay_ns=delay_ns))
+    eager = AbsorbedQueryProjection(attn)
+    static_inputs = [_draft_step_inputs(step, seed=300) for step in range(DRAFT_LEN)]
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            _draft_sequence(runner, static_inputs, absorb)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        packets = _draft_sequence(runner, static_inputs, absorb)
+    for trial, replays in enumerate([1, 1, 20]):
+        fresh = [
+            _draft_step_inputs(step, accepted=1 + trial, seed=400 + 10 * trial)
+            for step in range(DRAFT_LEN)
+        ]
+        for (s_hidden, s_pos), (hidden, pos) in zip(static_inputs, fresh):
+            s_hidden.copy_(hidden)
+            s_pos.copy_(pos)
+        for _ in range(replays):
+            graph.replay()
+        torch.cuda.synchronize()
+        for step, (packet, (hidden, pos)) in enumerate(zip(packets, fresh)):
+            reference = eager.project(pos, hidden, absorb=absorb)
+            torch.cuda.synchronize()
+            tag = f"trial {trial} x{replays} step {step}"
+            aq_proj._assert_same_packet(packet, reference, tag)
+            _check_step_vs_native(attn, packet, hidden, pos)
+        assert int(runner.sync.abs().sum()) == 0
+        assert int(eager.sync.abs().sum()) == 0
+    del graph
+
+
+@pytest.mark.parametrize("absorb", [True, False])
+@torch.inference_mode()
+def test_mtp_draft_loop_back_to_back_stress(model_config, absorb):
+    """Many eager draft-loop sequences (five back-to-back launches on one
+    runner) with delayed producers/consumers and fresh inputs reproduce the
+    undelayed launches bit for bit; counters stay zero."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        KernelConfig,
+    )
+
+    attn = _attn(model_config)
+    delayed = AbsorbedQueryProjection(attn, KernelConfig(debug_delay_ns=50000))
+    plain = AbsorbedQueryProjection(attn)
+    results = []
+    for it in range(int(os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_STRESS_ITERS", "6"))):
+        inputs = [
+            _draft_step_inputs(step, accepted=1 + it % 6, seed=600 + it)
+            for step in range(DRAFT_LEN)
+        ]
+        results.append(
+            (_draft_sequence(delayed, inputs, absorb), _draft_sequence(plain, inputs, absorb))
+        )
+    torch.cuda.synchronize()
+    for it, (got, want) in enumerate(results):
+        for step, (a, b) in enumerate(zip(got, want)):
+            aq_proj._assert_same_packet(a, b, f"iteration {it} step {step}")
+    assert int(delayed.sync.abs().sum()) == 0
+    assert int(plain.sync.abs().sum()) == 0
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py
new file mode 100644
index 000000000000..2c3539d2c41d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py
@@ -0,0 +1,447 @@
+# 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.
+"""Stock producer -> shared_indexer_absorbed_query -> native DSA attention, through the target's
+attention (the absorbed-query packet of absorbed_query, draft_absorbed_query and
+shared_indexer_absorbed_query).
+
+Uses the absorbed_query integration harness
+(``absorbed_query/test_absorbed_query_attention_integration .py``): real weights of the full-indexer
+layer 6 and the shared layer 7 at TP8 rank 0, six requests natively prefilled with 2,600 tokens, FP8
+KV cache, generation steps in the 6 x 1 and the 1 x 6 (MTP verify) layouts, both call
+paths, every step preceded by ``on_update_kv_lens()`` as in the runner.
+
+With the fast paths ``absorbed_query,shared_indexer_absorbed_query`` running, layer 7 takes
+shared_indexer_absorbed_query (native kv_a GEMM + the shared_indexer_absorbed_query
+kernel, preabsorbed packet, no native BMM) and layer 6 takes absorbed_query.
+Checks against the same step with every fast path off:
+
+* engagement: shared_indexer_absorbed_query serves layer 7 only; absorbed_query serves layer 6 only;
+* layer 7 (shared_indexer_absorbed_query): its latent cache is **byte-identical** to native, and the
+  attention inputs recorded at the native ``mla_rope_generation`` (absorbed
+  query, FP8 ``quant_q_buffer``, BMM scales) and the converted sparse
+  indices are bit-identical / select the same keys;
+* layer 6 (absorbed_query): bit-identical caches, attention and indexer inputs and
+  equal top-k key sets (the absorbed_query harness's comparison);
+* outputs within the native-noise bound; a repeated shared_indexer_absorbed_query step and 20 more
+  steps reproduce the shared_indexer_absorbed_query caches and recorded inputs bit for bit;
+* CUDA graphs (shared_indexer_absorbed_query alone, native alone, both in one pool): replays 1, 1
+  and 20 times with fresh hidden states reproduce the eager shared_indexer_absorbed_query step bit
+  for bit and compare with native as above; counters are zero after each trial;
+* feature gating through the target's ``_absorbed_query_packet``: ``""``, ``absorbed_query``,
+  ``shared_indexer_absorbed_query``, ``absorbed_query,shared_indexer_absorbed_query``,
+  ``absorbed_query,draft_absorbed_query,shared_indexer_absorbed_query``;
+* negative control: a shared_indexer_absorbed_query packet whose absorbed query is zeroed is
+  flagged.
+"""
+
+import math
+import os
+import sys
+
+import absorbed_query_fixture
+import pytest
+import test_absorbed_query_attention_integration as aq_integration
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+
+import plan_testlib  # noqa: E402
+
+LAYERS = aq_integration.LAYERS
+SHARED = absorbed_query_fixture.SHARED_LAYER
+FULL = absorbed_query_fixture.FULL_LAYER
+LAYOUTS = aq_integration.LAYOUTS
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+
+def _set_features(features: str, form: str = "auto"):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    query_proj.SERVED_FORM = form
+    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
+
+
+def _reset_runners(h):
+    for layer in LAYERS:
+        h.attn[layer]._absorbed_query = None
+        h.attn[layer]._shared_indexer_query = None
+
+
+def _runners(h):
+    return {
+        layer: (
+            getattr(h.attn[layer], "_absorbed_query", None),
+            getattr(h.attn[layer], "_shared_indexer_query", None),
+        )
+        for layer in LAYERS
+    }
+
+
+def _assert_engaged(h):
+    r = _runners(h)
+    assert r[SHARED][1] not in (None, False), "shared_indexer_absorbed_query did not engage on L7"
+    assert r[SHARED][0] is None, (
+        "absorbed_query must not run on the shared_indexer_absorbed_query layer"
+    )
+    assert r[FULL][1] is False, "shared_indexer_absorbed_query must decline the full layer"
+    assert r[FULL][0] not in (None, False), "absorbed_query did not engage on L6"
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = aq_integration._Harness()
+    yield h
+    _set_features("")
+
+
+def _compare_shared_query_to_native(snap, native, cand, n_new, bounds, report, tag):
+    """absorbed_query harness comparison (bitwise on both layers) plus an explicit check
+    of the shared_indexer_absorbed_query layer's latent cache."""
+    fails = aq_integration._compare_absorbed_query_to_native(
+        snap, native, cand, n_new, bounds, report, tag
+    )
+    n_diff = report.get(f"{tag}:latent{SHARED}_bytes_differing")
+    if n_diff != 0:
+        fails.append(
+            f"{tag}: shared_indexer_absorbed_query layer latent cache differs from native in "
+            f"{n_diff} bytes"
+        )
+    return fails
+
+
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_shared_query_into_native_attention(harness, layout, path):
+    h = harness
+    n_new = math.prod(LAYOUTS[layout])
+    hidden, pos = h.gen_inputs(layout, seed=12)
+    snap = h.snapshot()
+    try:
+        native_outs, native_recs, native_cache = aq_integration._native_runs(
+            h, layout, path, hidden, pos, snap
+        )
+        bounds = aq_integration._output_bounds(native_outs)
+
+        _set_features("absorbed_query,shared_indexer_absorbed_query")
+        _reset_runners(h)
+        shared_query_out, shared_query_rec = h.recorded_step(layout, path, hidden, pos)
+        _assert_engaged(h)
+        shared_query_cache = h.snapshot()
+        rope_served = {layer: "rope_append_rope" in shared_query_rec[layer] for layer in LAYERS}
+        print(f"rope_kv_append rope served: {rope_served}", flush=True)
+        if path == "custom_op" and "rope_kv_append" in absorbed_query_fixture.with_base_features(
+            ""
+        ).split(","):
+            assert rope_served[SHARED], (
+                "rope_kv_append did not serve the shared_indexer_absorbed_query layer's rope"
+            )
+
+        report, fails = {}, []
+        fails += _compare_shared_query_to_native(
+            snap,
+            (native_outs[0], native_recs[0], native_cache),
+            (shared_query_out, shared_query_rec, shared_query_cache),
+            n_new,
+            bounds,
+            report,
+            "shared_indexer_absorbed_query",
+        )
+
+        h.restore(snap)
+        _, rec_again = h.recorded_step(layout, path, hidden, pos)
+        fails += aq_integration._compare_caches(
+            snap, shared_query_cache, h.snapshot(), n_new, report, "shared_query_again"
+        )
+        fails += aq_integration._compare_records(
+            shared_query_rec, rec_again, report, "shared_query_again"
+        )
+        for _ in range(20):
+            again = h.gen_step(layout, path, hidden, pos)
+        fails += aq_integration._compare_caches(
+            snap, shared_query_cache, h.snapshot(), n_new, report, "shared_query_repeat20"
+        )
+        fails += aq_integration._compare_outputs(
+            native_outs[0], again, bounds, report, "shared_query_repeat20"
+        )
+        assert h.attn[SHARED]._shared_indexer_query.counters_quiescent()
+        print(f"layout={layout} path={path} {report}")
+        assert not fails, fails
+    finally:
+        h.restore(snap)
+        _set_features("")
+        _reset_runners(h)
+
+
+@pytest.mark.parametrize(
+    "features,expect",
+    [
+        ("", {FULL: (None, None), SHARED: (None, None)}),
+        ("absorbed_query", {FULL: ("absorbed_query", None), SHARED: ("absorbed_query", None)}),
+        ("shared_indexer_absorbed_query", {FULL: (None, None), SHARED: (None, None)}),
+        (
+            "absorbed_query,shared_indexer_absorbed_query",
+            {FULL: ("absorbed_query", False), SHARED: (None, "shared_indexer_absorbed_query")},
+        ),
+        (
+            "absorbed_query,draft_absorbed_query,shared_indexer_absorbed_query",
+            {FULL: ("absorbed_query", False), SHARED: (None, "shared_indexer_absorbed_query")},
+        ),
+    ],
+)
+@torch.inference_mode()
+def test_feature_gating_through_adapter(harness, features, expect):
+    """Which runner serves which layer for each feature selection (the draft_absorbed_query
+    entry point declines target layers, so ``draft_absorbed_query`` changes nothing here).
+    ``shared_indexer_absorbed_query`` without ``absorbed_query`` is refused by the selection itself
+    (``_set_features``), before any runner is attached, when the installed selector has the
+    dependency rule (``REQUIRED_FEATURES``); without that rule it runs native."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+        shared_indexer,
+    )
+
+    h = harness
+    layout, path = "1req_x6tok", "custom_op"
+    hidden, pos = h.gen_inputs(layout, seed=21)
+    snap = h.snapshot()
+    ids = features.split(",")
+    rejected = (
+        "shared_indexer_absorbed_query" in ids
+        and "absorbed_query" not in ids
+        and "shared_indexer_absorbed_query" in getattr(dispatch, "REQUIRED_FEATURES", {})
+    )
+    try:
+        if rejected:
+            _reset_runners(h)
+            with pytest.raises(ValueError, match="which they extend"):
+                _set_features(features)
+            assert _runners(h) == {layer: (None, None) for layer in LAYERS}, _runners(h)
+            return
+        _set_features(features)
+        _reset_runners(h)
+        h.gen_step(layout, path, hidden, pos)
+        got = _runners(h)
+        for layer, (want_absorbed_query, want_shared_query) in expect.items():
+            absorbed_query_r, shared_query_r = got[layer]
+            if want_absorbed_query == "absorbed_query":
+                assert isinstance(absorbed_query_r, query_proj.AbsorbedQueryProjection), (
+                    features,
+                    layer,
+                    got,
+                )
+            else:
+                assert absorbed_query_r is want_absorbed_query, (features, layer, got)
+            if want_shared_query == "shared_indexer_absorbed_query":
+                assert isinstance(shared_query_r, shared_indexer.SharedIndexerProjection), (
+                    features,
+                    layer,
+                    got,
+                )
+            else:
+                assert shared_query_r is want_shared_query, (features, layer, got)
+    finally:
+        h.restore(snap)
+        _set_features("")
+        _reset_runners(h)
+
+
+GRAPH_SETS = {
+    "w10only": ("shared_indexer_absorbed_query",),
+    "nativeonly": ("native",),
+    "both": ("shared_indexer_absorbed_query", "native"),
+}
+
+
+@pytest.mark.parametrize("graph_set", list(GRAPH_SETS))
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_shared_query_cuda_graph_replay_into_native_attention(harness, layout, graph_set):
+    """Captured shared_indexer_absorbed_query (+ absorbed_query on layer 6) + native attention + native cache
+    append, replayed once and 20 times with fresh inputs (the absorbed_query harness's
+    graph protocol: warm-up pass, capture into one shared pool, the runner's
+    on_update_kv_lens() inside every step)."""
+    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+
+    h = harness
+    n_new = math.prod(LAYOUTS[layout])
+    path = "custom_op"
+    snap = h.snapshot()
+    graphs = {}
+    wanted = GRAPH_SETS[graph_set]
+    configs = [
+        (
+            name,
+            "absorbed_query,shared_indexer_absorbed_query"
+            if name == "shared_indexer_absorbed_query"
+            else "",
+        )
+        for name in ("shared_indexer_absorbed_query", "native")
+        if name in wanted
+    ]
+    try:
+        warm = {}
+        for name, features in configs:
+            _set_features(features)
+            _reset_runners(h)
+            meta = h.gen_metadata(layout, cuda_graph=True)
+            static_hidden, static_pos = h.gen_inputs(layout, seed=200)
+            with with_multi_stream(True):
+                for _ in range(2):
+                    h.step(meta, static_pos, static_hidden, path)
+            torch.cuda.synchronize()
+            if name == "shared_indexer_absorbed_query":
+                _assert_engaged(h)
+            warm[name] = (meta, static_hidden, static_pos, _runners(h))
+        h.restore(snap)
+        pool = None
+        for name, features in configs:
+            meta, static_hidden, static_pos, runners = warm[name]
+            _set_features(features)
+            for layer, (absorbed_query_r, shared_query_r) in runners.items():
+                h.attn[layer]._absorbed_query = absorbed_query_r
+                h.attn[layer]._shared_indexer_query = shared_query_r
+            graph = torch.cuda.CUDAGraph()
+            with with_multi_stream(True), aq_integration._Recorder(h.attn) as rec:
+                with torch.cuda.graph(graph, pool=pool):
+                    outs = h.step(meta, static_pos, static_hidden, path)
+            pool = graph.pool()
+            graphs[name] = (graph, meta, static_hidden, outs, rec, runners)
+            torch.cuda.synchronize()
+        _reset_runners(h)
+
+        report, fails = {}, []
+        for trial, (seed, replays) in enumerate([(301, 1), (302, 1), (303, 20)]):
+            hidden, pos = h.gen_inputs(layout, seed=seed)
+            native_outs, native_recs, native_cache = aq_integration._native_runs(
+                h, layout, path, hidden, pos, snap
+            )
+            bounds = aq_integration._output_bounds(native_outs)
+            _set_features("absorbed_query,shared_indexer_absorbed_query")
+            h.restore(snap)
+            shared_query_eager = h.recorded_step(layout, path, hidden, pos)
+            shared_query_eager_cache = h.snapshot()
+            _set_features("")
+            replayed = {}
+            for name, (graph, meta, static_hidden, outs, rec, runners) in graphs.items():
+                h.restore(snap)
+                meta.prepare()
+                static_hidden.copy_(hidden)
+                for _ in range(replays):
+                    graph.replay()
+                torch.cuda.synchronize()
+                replayed[name] = (
+                    {layer: outs[layer].clone() for layer in LAYERS},
+                    rec.snapshot(),
+                    h.snapshot(),
+                )
+            tag = f"t{trial}x{replays}"
+            if "shared_indexer_absorbed_query" in replayed:
+                g_out, g_rec, g_cache = replayed["shared_indexer_absorbed_query"]
+                fails += aq_integration._compare_caches(
+                    snap,
+                    shared_query_eager_cache,
+                    g_cache,
+                    n_new,
+                    report,
+                    f"{tag}:shared_query_graph_vs_eager",
+                )
+                fails += aq_integration._compare_records(
+                    shared_query_eager[1], g_rec, report, f"{tag}:shared_query_graph_vs_eager"
+                )
+                fails += _compare_shared_query_to_native(
+                    snap,
+                    (native_outs[0], native_recs[0], native_cache),
+                    (g_out, g_rec, g_cache),
+                    n_new,
+                    bounds,
+                    report,
+                    f"{tag}:shared_query_graph_vs_native",
+                )
+            if "native" in replayed:
+                g_out, g_rec, g_cache = replayed["native"]
+                fails += aq_integration._compare_caches(
+                    snap, native_cache, g_cache, n_new, report, f"{tag}:native_graph_vs_eager"
+                )
+                fails += aq_integration._compare_records(
+                    native_recs[0], g_rec, report, f"{tag}:native_graph_vs_eager"
+                )
+            for _, (graph, meta, sh, outs, rec, runners) in graphs.items():
+                for absorbed_query_r, shared_query_r in runners.values():
+                    if absorbed_query_r not in (None, False):
+                        assert int(absorbed_query_r.sync.abs().sum()) == 0
+                    if shared_query_r not in (None, False):
+                        assert shared_query_r.counters_quiescent()
+        print(f"graph layout={layout} graphs={wanted} {report}")
+        assert not fails, fails
+    except BaseException as exc:
+        print(f"graph test failure: {exc!r}", flush=True)
+        raise
+    finally:
+        graphs.clear()
+        _reset_runners(h)
+        h.restore(snap)
+        _set_features("")
+
+
+@torch.inference_mode()
+def test_checks_detect_a_wrong_shared_query_absorbed_query(harness, monkeypatch):
+    """Negative control: zeroing the shared_indexer_absorbed_query absorbed query must be flagged by
+    the absorbed-query and quant_q checks of layer 7."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        shared_indexer,
+    )
+
+    h = harness
+    layout, path = "1req_x6tok", "custom_op"
+    hidden, pos = h.gen_inputs(layout, seed=12)
+    snap = h.snapshot()
+    orig = shared_indexer.SharedIndexerProjection.project
+
+    def broken(self, hidden_states):
+        packet = orig(self, hidden_states)
+        packet.fused_q[..., : self.mla.kv_lora_rank].zero_()
+        return packet
+
+    try:
+        native_outs, native_recs, _ = aq_integration._native_runs(
+            h, layout, path, hidden, pos, snap
+        )
+        monkeypatch.setattr(shared_indexer.SharedIndexerProjection, "project", broken)
+        _set_features("absorbed_query,shared_indexer_absorbed_query")
+        _reset_runners(h)
+        _, bad_rec = h.recorded_step(layout, path, hidden, pos)
+        report = {}
+        fails = aq_integration._compare_records(native_recs[0], bad_rec, report, "bad")
+        print(f"negative control {report}")
+        assert any(f"layer {SHARED} absorbed_q" in f for f in fails), fails
+        assert any(f"layer {SHARED} quant_q" in f for f in fails), fails
+        assert not any(f"layer {FULL} absorbed_q" in f for f in fails), fails
+    finally:
+        monkeypatch.undo()
+        h.restore(snap)
+        _set_features("")
+        _reset_runners(h)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py
new file mode 100644
index 000000000000..7aaa78cc645c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py
@@ -0,0 +1,599 @@
+# 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.
+"""Contract tests of shared_indexer_absorbed_query (absorbed_query shared-indexer-layer path, preabsorbed form).
+
+Real GLM-5.2 weights of the shared layer 7 (and the full layer 6 for the
+fallback checks) at TP8 rank 0; synthetic activations (unit-RMS rows times
+the layer's real ``input_layernorm`` gain) unless stated.
+
+Acceptance (frozen before tuning): every shared_indexer_absorbed_query output is **bit-identical**
+to the native region it replaces -- ``forward_dsa_proj`` (``q``,
+``compressed_kv``, ``k_pe``, ``latent_cache``) followed by the native
+absorption BMM (``fused_q[..., :512]``) -- because the kernel reproduces the
+native RMSNorm arithmetic and uses absorbed_query's MMA sequences (absorbed_query is held to the
+same bitwise acceptance). The normalized q_a rows the kernel shares between its programs must equal the
+native ``q_a_layernorm`` output bit for bit as well.
+
+Also covered: every kernel configuration (the CuTe-DSL cluster kernel with
+both q exchanges, early PDL trigger on (default) / off and weights before /
+after the dependency wait; the Triton kernel's poll protocols, L2 prefetch,
+W_UK preload, early trigger, rows per norm program, tile shapes and
+orientations, counter protocols), R in {1..8, 16}, padded rows,
+CUDA-graph capture and replay (also with delayed programs), back-to-back and
+multi-stream stress, progress with only a few free SMs, memcpy-fed inputs
+(for compute-sanitizer initcheck), eligibility / fallback (returns
+``None`` before any GPU work), and the package selectors: ``shared_indexer_absorbed_query`` is off
+without ``absorbed_query``, and the absorbed_query form is logged.
+"""
+
+import os
+import sys
+import time
+from types import SimpleNamespace
+
+import absorbed_query_fixture
+import pytest
+import test_absorbed_query_proj as aq_proj
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+
+import plan_testlib  # noqa: E402
+
+SHARED = absorbed_query_fixture.SHARED_LAYER
+FULL = absorbed_query_fixture.FULL_LAYER
+ROWS = aq_proj._env_list(
+    "GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_TEST_ROWS", [1, 2, 3, 4, 5, 6, 7, 8, 16]
+)
+STRESS_ROWS = aq_proj._env_list("GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_TEST_ROWS", [1, 6, 16])
+STRESS_LAUNCHES = int(
+    os.environ.get("GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_STRESS_LAUNCHES", "8")
+)
+
+
+def _triton(**kwargs):
+    return dict(cute=False, **kwargs)
+
+
+# Kernel configurations that must all give the same (native) bits.
+# "default" is the production configuration (CuTe-DSL cluster kernel for
+# R <= 8, Triton kernel for R = 16); the Triton configurations set
+# cute=False.
+CONFIGS = {
+    "default": {},
+    "triton": _triton(),
+    "poll_atomic": _triton(poll_mode=0),
+    "poll_sleep": _triton(poll_mode=2),
+    "no_prefetch": _triton(prefetch_l2=False),
+    "no_preload": _triton(preload_wuk=False),
+    "no_trigger": _triton(early_trigger=False),
+    "norm_rows_2": _triton(norm_rows=2),
+    "norm_rows_8": _triton(norm_rows=8),
+    "norm_scalar": _triton(norm_vectorized=False),
+    "transposed_bn32": _triton(gemm_transposed=True, absorb_transposed=True),
+    "transposed_bn64": _triton(
+        bn=64, gemm_transposed=True, absorb_transposed=True, preload_wuk=False
+    ),
+    "bn64": _triton(bn=64, preload_wuk=False),
+    "epoch_bn16": _triton(bn=16, num_stages=8),
+    "absorbed_query_protocol": _triton(epoch_sync=False),
+    # CuTe-DSL cluster kernel (R <= 8; R = 16 takes the Triton kernel).
+    "cute": dict(cute=True),
+    "cute_no_trigger": dict(cute=True, early_trigger=False),
+    "cute_weights_after_wait": dict(cute=True, cute_weights_before_wait=False),
+    "cute_cluster_barrier": dict(cute=True, cute_push_async=False),
+}
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+
+@pytest.fixture(scope="module")
+def model_config():
+    return absorbed_query_fixture.build_model_config(0)
+
+
+def _runner(attn, **kwargs):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        shared_indexer,
+    )
+
+    return shared_indexer.SharedIndexerProjection(
+        attn, shared_indexer.SharedQueryKernelConfig(**kwargs)
+    )
+
+
+def _assert_packet_native(packet, native, kv_lora, tag, rows=slice(None)):
+    """shared_indexer_absorbed_query packet vs the native region: bit for bit."""
+    native_outs, native_fq = native
+    names = ("q", "compressed_kv", "k_pe", "latent_cache")
+    assert len(packet.proj_outputs) == 4, tag
+    for name, x, y in zip(names, packet.proj_outputs, native_outs):
+        assert x.shape == y.shape and x.dtype == y.dtype, (tag, name)
+        aq_proj._assert_bitwise(f"{tag} {name}", x[rows], y[rows])
+    assert packet.fused_q is not None
+    aq_proj._assert_bitwise(
+        f"{tag} absorbed", packet.fused_q[rows, :, :kv_lora], native_fq[rows, :, :kv_lora]
+    )
+
+
+def _assert_same(a, b, kv_lora, tag):
+    for i, (x, y) in enumerate(zip(a.proj_outputs, b.proj_outputs)):
+        aq_proj._assert_bitwise(f"{tag} out{i}", x, y)
+    aq_proj._assert_bitwise(f"{tag} absorbed", a.fused_q[..., :kv_lora], b.fused_q[..., :kv_lora])
+
+
+def _inputs(rows, seed=0, scale=1.0):
+    hidden, position_ids = aq_proj._inputs(SHARED, rows, seed=seed)
+    if scale != 1.0:
+        hidden = (hidden.float() * scale).to(torch.bfloat16)
+    return hidden, position_ids
+
+
+@pytest.mark.parametrize("kind", ["real_q_a", "real_ckv", "wide"])
+@pytest.mark.parametrize("rows", [1, 6, 64, 4096])
+@torch.inference_mode()
+def test_norm_reproduction_is_bitwise(model_config, kind, rows):
+    """The Triton reproduction of the native RMSNorm (sequential per-lane
+    fold) equals the native module call bit for bit, on strided kv_a views
+    and on wide-dynamic-range rows."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        shared_indexer,
+    )
+
+    attn = aq_proj._attn(model_config, SHARED)
+    gen = torch.Generator(device="cuda").manual_seed(rows)
+    if kind == "wide":
+        mag = torch.exp(torch.empty(rows, 2048, device="cuda").uniform_(-6, 6, generator=gen))
+        x = (torch.randn(rows, 2048, device="cuda", generator=gen) * mag).to(torch.bfloat16)
+        norm = attn.q_a_layernorm
+    else:
+        hidden = absorbed_query_fixture.make_hidden(rows, seed=rows, gamma=aq_proj._gamma(SHARED))
+        kva = attn.kv_a_proj_with_mqa(hidden)
+        if kind == "real_q_a":
+            x, norm = kva[:, :2048], attn.q_a_layernorm
+        else:
+            x, norm = kva[:, 2048:2560], attn.kv_a_layernorm
+    native = norm(x)
+    for rpp in (1, 4, 8):
+        for vec in (False, True):
+            rep = shared_indexer.fi_rmsnorm_reproduction(
+                x, norm.weight, norm.variance_epsilon, order=0, rows_per_program=rpp, vectorized=vec
+            )
+            aq_proj._assert_bitwise(f"{kind} rows={rows} rpp={rpp} vec={vec}", rep, native)
+
+
+@pytest.mark.parametrize("rows", ROWS)
+@pytest.mark.parametrize("config_name", list(CONFIGS))
+@torch.inference_mode()
+def test_contract_matches_native_bitwise(model_config, rows, config_name):
+    """Every output of the shared_indexer_absorbed_query packet equals the native region bit for bit
+    for every configuration; the shared normalized rows equal the native
+    q_a_layernorm output; ``q`` equals absorbed_query's; the counters return to zero."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+    )
+
+    if config_name not in ("default", "triton") and rows not in (1, 6, 16):
+        pytest.skip("non-default configurations run at R in {1, 6, 16}")
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn, debug_keep_qn=True, **CONFIGS[config_name])
+    hidden, position_ids = _inputs(rows, seed=rows)
+    packet = runner.project(hidden)
+    native = aq_proj._native(attn, hidden, position_ids)
+    kva = attn.kv_a_proj_with_mqa(hidden)
+    native_qn = attn.q_a_layernorm(kva[:, :2048])
+    aq_packet = AbsorbedQueryProjection(attn).project(position_ids, hidden, absorb=False)
+    torch.cuda.synchronize()
+    _assert_packet_native(packet, native, attn.kv_lora_rank, f"R={rows} {config_name}")
+    aq_proj._assert_bitwise("qn", runner.debug["qn"], native_qn)
+    aq_proj._assert_bitwise(
+        "q vs absorbed_query", packet.proj_outputs[0], aq_packet.proj_outputs[0]
+    )
+    # Packet layout: compressed_kv / k_pe are views of latent_cache.
+    q, ckv, k_pe, latent = packet.proj_outputs
+    assert ckv.data_ptr() == latent.data_ptr()
+    assert k_pe.data_ptr() == latent.data_ptr() + 512 * latent.element_size()
+    assert packet.fused_q.shape == (
+        rows,
+        attn.num_heads_tp,
+        attn.kv_lora_rank + attn.qk_rope_head_dim,
+    )
+    assert runner.counters_quiescent(), runner.epochs
+
+
+@pytest.mark.parametrize("rows", [1, 2, 3, 4, 5, 6, 7, 8])
+@torch.inference_mode()
+def test_cute_kernel_all_rows_back_to_back(model_config, rows):
+    """The CuTe-DSL cluster kernel at every R it serves: 8 back-to-back
+    launches with fresh inputs (also scaled), each bit-identical to the
+    native region, the shared normalized rows equal to the native
+    q_a_layernorm output."""
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn, cute=True, debug_keep_qn=True)
+    launches = []
+    for i in range(8):
+        hidden, position_ids = _inputs(rows, seed=4000 + i, scale=(1.0, 1e-3, 30.0)[i % 3])
+        packet = runner.project(hidden)
+        qn = runner.debug["qn"].clone()
+        kva = attn.kv_a_proj_with_mqa(hidden)
+        launches.append(
+            (
+                packet,
+                qn,
+                attn.q_a_layernorm(kva[:, :2048]),
+                aq_proj._native(attn, hidden, position_ids),
+            )
+        )
+    torch.cuda.synchronize()
+    for i, (packet, qn, native_qn, native) in enumerate(launches):
+        aq_proj._assert_bitwise(f"cute qn {i}", qn, native_qn)
+        _assert_packet_native(packet, native, attn.kv_lora_rank, f"cute launch {i}")
+
+
+@pytest.mark.parametrize("rows", [6, 16])
+@pytest.mark.parametrize("scale", [1.0, 1e-3, 30.0])
+@torch.inference_mode()
+def test_native_agreement_statistics(model_config, rows, scale):
+    """48 fresh inputs per case (also scaled hidden states, which move the
+    RMSNorm sums across many rounding boundaries): zero differing elements
+    in every output; absorbed_query's own compressed_kv norm is reported for
+    comparison."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+    )
+
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn)
+    aq_projection = AbsorbedQueryProjection(attn)
+    absorbed_query_ckv_diff = 0
+    for i in range(48):
+        hidden, position_ids = _inputs(rows, seed=3000 + i, scale=scale)
+        packet = runner.project(hidden)
+        native = aq_proj._native(attn, hidden, position_ids)
+        aq_packet = aq_projection.project(position_ids, hidden, absorb=False)
+        torch.cuda.synchronize()
+        _assert_packet_native(packet, native, attn.kv_lora_rank, f"input {i}")
+        absorbed_query_ckv_diff += int(
+            (aq_proj._steps(aq_packet.proj_outputs[1], native[0][1]) != 0).sum()
+        )
+    print(
+        f"rows={rows} scale={scale}: shared_indexer_absorbed_query 0 differing elements in 48 "
+        f"inputs; absorbed_query compressed_kv differing elements {absorbed_query_ckv_diff}"
+    )
+
+
+@pytest.mark.parametrize("cute", [False, True], ids=["triton", "cute"])
+@torch.inference_mode()
+def test_padded_rows_do_not_affect_live_rows(model_config, cute):
+    """Rows are independent: padded rows (zero or 1e3 garbage) never change a
+    live row, live rows equal native bit for bit, zero rows stay finite."""
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn, cute=cute)
+    live = slice(0, 3)
+    hidden_zero, position_ids = aq_proj._inputs(SHARED, 6, padded_from=3)
+    hidden_junk = hidden_zero.clone()
+    hidden_junk[3:] = 1e3 * torch.randn_like(hidden_junk[3:])
+    a = runner.project(hidden_zero)
+    b = runner.project(hidden_junk)
+    native = aq_proj._native(attn, hidden_junk, position_ids)
+    torch.cuda.synchronize()
+    for i, (x, y) in enumerate(zip(a.proj_outputs, b.proj_outputs)):
+        aq_proj._assert_bitwise(f"padded out{i}", x[live], y[live])
+    aq_proj._assert_bitwise(
+        "padded absorbed",
+        a.fused_q[live, :, : attn.kv_lora_rank],
+        b.fused_q[live, :, : attn.kv_lora_rank],
+    )
+    _assert_packet_native(b, native, attn.kv_lora_rank, "padded", live)
+    for x in a.proj_outputs:
+        assert torch.isfinite(x[3:].float()).all()
+    assert runner.counters_quiescent(), runner.epochs
+
+
+@pytest.mark.parametrize(
+    "delay_ns,cute",
+    [(0, False), (20000, False), (0, True)],
+    ids=["triton", "triton_delayed", "cute"],
+)
+@pytest.mark.parametrize("rows", [1, 6])
+@torch.inference_mode()
+def test_cuda_graph_replay(model_config, delay_ns, cute, rows):
+    """One and many replays of a captured shared_indexer_absorbed_query call (native kv_a GEMM +
+    shared_indexer_absorbed_query kernel, PDL-chained) with fresh inputs equal the native region bit
+    for bit and leave the counters at zero."""
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn, debug_delay_ns=delay_ns, cute=cute)
+    static_hidden, _ = _inputs(rows)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            runner.project(static_hidden)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        packet = runner.project(static_hidden)
+    for step, replays in enumerate([1, 1, 50]):
+        hidden, position_ids = _inputs(rows, seed=100 + step)
+        static_hidden.copy_(hidden)
+        for _ in range(replays):
+            graph.replay()
+        torch.cuda.synchronize()
+        native = aq_proj._native(attn, hidden, position_ids)
+        torch.cuda.synchronize()
+        _assert_packet_native(packet, native, attn.kv_lora_rank, f"replay x{replays}")
+        assert runner.counters_quiescent(), runner.epochs
+
+
+@pytest.mark.parametrize("rows", STRESS_ROWS)
+@pytest.mark.parametrize(
+    "config_name", ["triton", "poll_atomic", "no_preload", "norm_rows_2", "absorbed_query_protocol"]
+)
+@torch.inference_mode()
+def test_delayed_programs_back_to_back(model_config, rows, config_name):
+    """Norm programs delayed before they release, producers before they
+    publish, absorption jobs before they count out (up to 50 us,
+    deterministic per program): back-to-back launches with fresh inputs
+    equal the native region bit for bit; counters end at zero."""
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn, debug_delay_ns=50000, **CONFIGS[config_name])
+    launches = []
+    for i in range(STRESS_LAUNCHES):
+        hidden, position_ids = _inputs(rows, seed=500 + i)
+        launches.append((runner.project(hidden), aq_proj._native(attn, hidden, position_ids)))
+    torch.cuda.synchronize()
+    for i, (packet, native) in enumerate(launches):
+        _assert_packet_native(packet, native, attn.kv_lora_rank, f"delayed launch {i}")
+    assert runner.counters_quiescent(), runner.epochs
+
+
+@pytest.mark.parametrize(
+    "config_name",
+    ["default", "triton", "poll_atomic", "absorbed_query_protocol", "cute_cluster_barrier"],
+)
+@torch.inference_mode()
+def test_concurrent_launches_on_many_streams(model_config, config_name):
+    """Six shared_indexer_absorbed_query instances (own counters, like six layers) concurrently on six
+    streams, 30 rounds, with and without delays: each stream's last result
+    equals the native region bit for bit; no hang; counters at zero."""
+    attn = aq_proj._attn(model_config, SHARED)
+    n_streams, rounds, rows = 6, 30, 6
+    streams = [torch.cuda.Stream() for _ in range(n_streams)]
+    cute = CONFIGS[config_name].get("cute", False)
+    runners = [
+        _runner(attn, debug_delay_ns=(5000 if s % 2 and not cute else 0), **CONFIGS[config_name])
+        for s in range(n_streams)
+    ]
+    inputs = [_inputs(rows, seed=700 + s) for s in range(n_streams)]
+    natives = [aq_proj._native(attn, h, p) for h, p in inputs]
+    torch.cuda.synchronize()
+    last = [None] * n_streams
+    for _ in range(rounds):
+        for s in range(n_streams):
+            with torch.cuda.stream(streams[s]):
+                last[s] = runners[s].project(inputs[s][0])
+    torch.cuda.synchronize()
+    for s in range(n_streams):
+        _assert_packet_native(last[s], natives[s], attn.kv_lora_rank, f"stream {s}")
+        assert runners[s].counters_quiescent()
+
+
+@pytest.mark.parametrize("free_sms", [1, 2, 4, 8])
+@torch.inference_mode()
+def test_progress_with_restricted_sms(model_config, free_sms):
+    """A concurrent kernel holds all but ``free_sms`` SMs for 0.5 s (no shared_indexer_absorbed_query
+    program fits beside one of its programs). The shared_indexer_absorbed_query kernel must complete
+    while those SMs are still held, with native bits: every spin-wait
+    targets lower program ids (norm programs, then producers, then the
+    extra absorption programs). Triton kernel only: the CuTe kernel's
+    clusters are co-scheduled by the hardware and do not wait on each
+    other."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        shared_indexer,
+    )
+
+    attn = aq_proj._attn(model_config, SHARED)
+    runner = _runner(attn, cute=False)
+    rows = 6
+    hidden, position_ids = _inputs(rows, seed=900)
+    native = aq_proj._native(attn, hidden, position_ids)
+    kva = attn.kv_a_proj_with_mqa(hidden)
+    reference = runner.launch(kva)
+    torch.cuda.synchronize()
+
+    n_sm = torch.cuda.get_device_properties(0).multi_processor_count
+    n_occ = n_sm - free_sms
+    bm = bn = 128
+    bk, iters = 64, 32
+    a = torch.randn(bm, iters * bk, device="cuda").to(torch.bfloat16)
+    b = torch.randn(iters * bk, bn, device="cuda").to(torch.bfloat16)
+    out = torch.empty(n_occ, bm, bn, device="cuda")
+    smids = torch.full((n_occ,), -1, dtype=torch.int32, device="cuda")
+    occ_stream, shared_query_stream = torch.cuda.Stream(), torch.cuda.Stream()
+    ev = {k: torch.cuda.Event(enable_timing=True) for k in ("s", "o", "p")}
+    torch.cuda.synchronize()
+    with torch.cuda.stream(occ_stream):
+        ev["s"].record()
+        occ = aq_proj._occupier_kernel[(n_occ,)](
+            a, b, out, smids, 500_000_000, iters, bm, bn, bk, 6, num_warps=4, num_stages=6
+        )
+        ev["o"].record()
+    time.sleep(0.05)
+    with torch.cuda.stream(shared_query_stream):
+        result = runner.launch(kva)
+        ev["p"].record()
+    torch.cuda.synchronize()
+    t_shared_query, t_occ = ev["s"].elapsed_time(ev["p"]), ev["s"].elapsed_time(ev["o"])
+    shared_query_smem = min(
+        aq_proj._compiled_shared_bytes(shared_indexer.shared_query_kernel) or [0]
+    )
+    held = int(torch.unique(smids).numel())
+    report = dict(
+        free_sms=free_sms,
+        held_sms=held,
+        occupier_smem=occ.metadata.shared,
+        shared_query_min_smem=shared_query_smem,
+        shared_query_done_ms=round(t_shared_query, 3),
+        occupier_done_ms=round(t_occ, 3),
+    )
+    print(f"restricted SMs {report}")
+    assert held == n_occ, report
+    assert occ.metadata.shared + shared_query_smem > 228 * 1024, report
+    for name, x, y in zip(("q", "latent", "fused_q"), result, reference):
+        if name == "fused_q":
+            x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
+        aq_proj._assert_bitwise(f"restricted {name}", x, y)
+    aq_proj._assert_bitwise("restricted q vs native", result[0], native[0][0])
+    assert runner.counters_quiescent(), runner.epochs
+    assert t_shared_query < t_occ - 100.0, (
+        f"shared_indexer_absorbed_query did not complete while {n_occ} SMs were held: {report}"
+    )
+
+
+@pytest.mark.parametrize("cute", [False, True], ids=["triton", "cute"])
+@pytest.mark.parametrize("rows", STRESS_ROWS)
+@torch.inference_mode()
+def test_kernel_with_memcpy_initialized_inputs(model_config, rows, cute):
+    """The shared_indexer_absorbed_query kernel alone with every tensor it reads placed by a
+    host-to-device copy (kv_a output, weights, norm gains, counters) equals
+    the normal launch bit for bit. ``GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_INITCHECK_ONLY=1``
+    (initcheck run) skips the reference launch so that only memcpy-fed memory is read by
+    instrumented kernels."""
+    initcheck_only = (
+        os.environ.get("GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_INITCHECK_ONLY") == "1"
+    )
+    attn = aq_proj._attn(model_config, SHARED)
+    hidden, _ = _inputs(rows, seed=1200)
+    kva = attn.kv_a_proj_with_mqa(hidden)
+    reference = None
+    if not initcheck_only:
+        reference = _runner(attn, cute=cute).launch(kva)
+    torch.cuda.synchronize()
+    probe = _runner(attn, cute=cute)
+    probe.sync = aq_proj._memcpy_copy(torch.zeros_like(probe.sync))
+    probe.epochs = aq_proj._memcpy_copy(torch.zeros_like(probe.epochs))
+    kva_copy = aq_proj._memcpy_copy(kva)
+    params = [
+        attn.q_b_proj.weight,
+        attn.k_b_proj_trans,
+        attn.kv_a_layernorm.weight,
+        attn.q_a_layernorm.weight,
+    ]
+    saved = [p.data for p in params]
+    try:
+        for p in params:
+            p.data = aq_proj._memcpy_copy(p.data)
+        torch.cuda.synchronize()
+        result = probe.launch(kva_copy)
+        torch.cuda.synchronize()
+    finally:
+        for p, data in zip(params, saved):
+            p.data = data
+    assert probe.counters_quiescent()
+    if initcheck_only:
+        return
+    for name, x, y in zip(("q", "latent", "fused_q"), result, reference):
+        if name == "fused_q":
+            x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
+        aq_proj._assert_bitwise(f"memcpy inputs {name}", x, y)
+
+
+@torch.inference_mode()
+def test_eligibility_and_fallback(model_config, monkeypatch):
+    """``shared_dsa_proj`` returns None before any GPU work outside its
+    domain (full-indexer and draft layers, context / mixed / missing
+    metadata, R > 16, non-BF16 input, force-native switch) and the packet
+    of a generation-only call on a shared layer otherwise."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        shared_indexer,
+    )
+
+    shared = aq_proj._attn(model_config, SHARED)
+    full = aq_proj._attn(model_config, FULL)
+    gen = SimpleNamespace(num_contexts=0, num_generations=1)
+    ctx = SimpleNamespace(num_contexts=1, num_generations=0)
+    mixed = SimpleNamespace(num_contexts=1, num_generations=1)
+    hidden, position_ids = _inputs(6)
+
+    assert shared_indexer.shared_module_eligibility(shared)[0]
+    ok, reason = shared_indexer.shared_module_eligibility(full)
+    assert not ok and "full indexer" in reason
+    full._shared_indexer_query = None
+    assert shared_indexer.shared_dsa_proj(full, position_ids, hidden, gen) is None
+    assert full._shared_indexer_query is False
+    for meta in (None, ctx, mixed):
+        shared._shared_indexer_query = None
+        assert shared_indexer.shared_dsa_proj(shared, position_ids, hidden, meta) is None
+    big = torch.zeros(17, hidden.shape[1], dtype=hidden.dtype, device="cuda")
+    assert shared_indexer.shared_dsa_proj(shared, None, big, gen) is None
+    assert shared_indexer.shared_dsa_proj(shared, None, hidden.float(), gen) is None
+    saved = shared.layer_idx
+    try:
+        shared.layer_idx = shared.num_hidden_layers
+        assert not shared_indexer.shared_module_eligibility(shared)[0]
+    finally:
+        shared.layer_idx = saved
+    shared._shared_indexer_query = None
+    packet = shared_indexer.shared_dsa_proj(shared, position_ids, hidden, gen)
+    assert packet is not None and len(packet.proj_outputs) == 4
+    assert packet.fused_q is not None
+    _assert_packet_native(
+        packet, aq_proj._native(shared, hidden, position_ids), shared.kv_lora_rank, "entry point"
+    )
+    shared._shared_indexer_query = None
+
+
+def test_package_selectors():
+    """``shared_indexer_absorbed_query`` needs ``absorbed_query`` (the plan turns it off with
+    ``absorbed_query``, so it never runs silently native); the absorbed_query form is
+    logged once next to the feature set."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    assert dispatch.REQUIRED_FEATURES["shared_indexer_absorbed_query"] == {"absorbed_query"}
+    plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: "absorbed_query"})
+    assert not plan.enabled("shared_indexer_absorbed_query")
+    assert plan_testlib.plan_running(
+        dispatch, "absorbed_query,shared_indexer_absorbed_query"
+    ).enabled("shared_indexer_absorbed_query")
+    try:
+        dispatch._CAPABILITY_LOGGED.pop(("absorbed_query", "form"), None)
+        assert query_proj.check_package_selectors() == query_proj.SERVED_FORM == "auto"
+        assert ("absorbed_query", "form") in dispatch._CAPABILITY_LOGGED
+    finally:
+        dispatch._CAPABILITY_LOGGED.pop(("absorbed_query", "form"), None)
+
+
+def test_registered_entry_point():
+    """``shared_dsa_proj`` is reachable through the dispatch registry under
+    the absorbed_query package, and the absorbed_query / draft_absorbed_query entries are
+    unchanged."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+        shared_indexer,
+    )
+
+    assert dispatch.get_impl("absorbed_query", "shared_dsa_proj") is shared_indexer.shared_dsa_proj
+    assert dispatch.get_impl("absorbed_query", "dsa_proj") is query_proj.dsa_proj
+    assert dispatch.get_impl("absorbed_query", "mtp_dsa_proj") is query_proj.mtp_dsa_proj
+    assert "shared_indexer_absorbed_query" in dispatch.SELECTOR_PATHS
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py
new file mode 100644
index 000000000000..7f384ebceb21
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py
@@ -0,0 +1,485 @@
+# 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.
+"""Real GLM-5.2 DSA attention layers for attention_output_tail end-to-end tests and benchmarks.
+
+``AttentionStack`` builds, for one TP8 rank, the attention module the model
+builds for ``glm_moe_dsa`` checkpoints (``DeepseekV32Attention``) for a run of
+consecutive layers. It loads every attention weight of the rank's shard from
+the checkpoint the way ``DeepseekV3WeightLoader`` does, fills an FP8
+``DSACacheManager`` with a real prefill forward, and prepares GEN-only
+attention metadata for one request with ``rows`` query tokens (MTP target
+verification with ``rows - 1`` draft tokens) the way the decoder runner does:
+CUDA-graph metadata, ``update_spec_dec_param``, then ``prepare``.
+
+Each module is adopted as the target's ``GlmMLA`` right after construction, as the target's
+model adopts it, and the GEN forward runs as a decode step (``scope()``), so ``GlmMLA.forward``
+then runs exactly as in the model: ``register_to_config`` is set, so the GEN attention runs
+inside the target's DSA attention op and ``o_proj`` runs in ``GlmMLA.forward``. With a plan
+running ``attention_output_tail`` installed at construction (``fast_paths``, or ``select``),
+every target layer owns an attention_output_tail ``AttnTail``; ``set_attn_tail`` swaps it in and out to compare
+against the native tail on the same module, cache and inputs. The prefill runs as a prefill
+step (the built-in attention).
+
+Fixture labels: the weights (all attention projections, norms and indexer
+weights of the rank's shard) are real. The activations are synthetic: the
+prefill and GEN inputs are the layer's real ``input_layernorm`` applied to
+real token embeddings, not the residual stream of a full forward.
+
+The cache manager and the attention metadata use a one-rank mapping: an MLA
+cache holds one latent head, so it is not TP-sharded, and a one-rank mapping
+keeps their constructors free of collectives in a process that stands in for
+one of eight TP ranks.
+"""
+
+import contextlib
+import os
+import sys
+import weakref
+from typing import Dict, List, Optional
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+import tail_fixtures  # noqa: E402
+
+TP8 = 8
+REQUEST_ID = 7
+TOKEN_OFFSET = 1000
+
+
+def load_pretrained(ckpt: str):
+    """Load the checkpoint's config as the model does for ``glm_moe_dsa``."""
+    from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config
+
+    cfg = load_pretrained_config(ckpt, trust_remote_code=True)
+    # DeepseekV3ForCausalLM renames glm_moe_dsa to deepseek_v32, which makes
+    # the decoder layers build DeepseekV32Attention.
+    cfg.model_type = "deepseek_v32"
+    dtype = getattr(cfg, "torch_dtype", None)
+    if not isinstance(dtype, torch.dtype):
+        resolved = getattr(torch, str(dtype or "bfloat16").replace("torch.", ""))
+        try:
+            cfg.torch_dtype = resolved
+        except AttributeError:
+            cfg.dtype = resolved
+    return cfg
+
+
+def make_model_config(pretrained, tp_rank: int, max_draft_len: int):
+    """``ModelConfig`` of the served GEN role for TP8 rank ``tp_rank``.
+
+    Attention weights of this checkpoint are excluded from NVFP4, so the
+    module is built unquantized with an FP8 KV cache, as the model's
+    exclusion handling leaves it.
+    """
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, MTPDecodingConfig
+    from tensorrt_llm.mapping import Mapping
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+    from tensorrt_llm.quantization.mode import QuantAlgo
+
+    quant = QuantConfig()
+    quant.kv_cache_quant_algo = QuantAlgo.FP8
+    sparse = DeepSeekSparseAttentionConfig(
+        enable_heuristic_topk=True, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=True
+    )
+    spec = MTPDecodingConfig(max_draft_len=max_draft_len) if max_draft_len > 0 else None
+    model_config = ModelConfig(
+        pretrained_config=pretrained,
+        mapping=Mapping(world_size=TP8, tp_size=TP8, rank=tp_rank),
+        quant_config=quant,
+        sparse_attention_config=sparse,
+        spec_config=spec,
+    )
+    model_config.extra_attrs["kv_cache_dtype"] = "fp8"
+    return model_config
+
+
+def build_attention(model_config, layer_idx: int, aux_stream, reduce_output: bool):
+    """Build ``DeepseekV32Attention`` for one layer on the current device, adopted as ``GlmMLA``.
+
+    The module has TP8 shapes, but this process may run without the other
+    ranks, so the all-reduce modules its ``Linear``s would create are
+    replaced by ``None`` during construction; callers that reduce attach a
+    real ``AllReduce`` afterwards.
+    """
+    import tensorrt_llm._torch.distributed as distributed
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+        GlmMLA,
+    )
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+    from tensorrt_llm._torch.utils import AuxStreamType
+
+    saved = distributed.AllReduce
+    distributed.AllReduce = lambda *args, **kwargs: None
+    try:
+        with torch.device("cuda"):
+            attn = DeepseekV32Attention(
+                model_config,
+                layer_idx=layer_idx,
+                aux_stream_dict={AuxStreamType.Attention: aux_stream},
+                reduce_output=reduce_output,
+            )
+            GlmMLA.adopt(attn)
+            # The model builds the FMHA libraries in ``create_weights`` after adoption; a module
+            # built alone rebuilds them for the library adoption declared.
+            for backend in (attn.mha, attn.mqa):
+                if backend is not None and backend.prepended_fmha_lib_classes:
+                    backend.update_quant_config(backend.quant_config)
+    finally:
+        distributed.AllReduce = saved
+    return attn.cuda()
+
+
+def load_attention_weights(attn, ckpt: str, layer: int, tp_rank: int, tp_size: int = TP8) -> None:
+    """Load the attention weights of ``layer`` like ``DeepseekV3WeightLoader``."""
+    from tensorrt_llm._torch.models.modeling_utils import filter_weights
+
+    prefix = f"model.layers.{layer}.self_attn"
+    weights = tail_fixtures.read_prefix(ckpt, prefix + ".")
+    heads = attn.num_heads
+    nope, v_dim, kv_lora = (attn.qk_nope_head_dim, attn.v_head_dim, attn.kv_lora_rank)
+    with torch.no_grad():
+        for name, module in attn.named_modules():
+            if not name or len(module._parameters) == 0:
+                continue
+            full = f"{prefix}.{name}"
+            leaf = name.split(".")[-1]
+            if leaf == "kv_b_proj":
+                kv_b = weights[f"{full}.weight"].unflatten(0, [heads, nope + v_dim])
+                kv_b = torch.chunk(kv_b, tp_size, dim=0)[tp_rank]
+                k_nope, v_w = kv_b.split([nope, v_dim], dim=1)
+                local = kv_b.shape[0]
+                module.weight.data.copy_(
+                    torch.cat(
+                        [k_nope.reshape(local * nope, kv_lora), v_w.reshape(local * v_dim, kv_lora)]
+                    )
+                )
+                attn.v_b_proj = torch.nn.Parameter(
+                    module.weight.data[local * nope :].view(local, v_dim, kv_lora),
+                    requires_grad=False,
+                )
+                attn.k_b_proj_trans.data.copy_(
+                    k_nope.transpose(2, 1).contiguous().reshape(attn.k_b_proj_trans.shape)
+                )
+            elif leaf == "kv_a_proj_with_mqa":
+                fused = torch.cat([weights[f"{prefix}.q_a_proj.weight"], weights[f"{full}.weight"]])
+                module.weight.data[: fused.shape[0]].copy_(fused)
+            else:
+                module_weights = filter_weights(full, weights)
+                if hasattr(module, "load_weights"):
+                    module.load_weights(weights=[module_weights])
+                else:
+                    for n, p in module.named_parameters():
+                        p.data.copy_(module_weights[n][:])
+        for m in attn.modules():
+            post = getattr(m, "post_load_weights", None)
+            if post is not None and not getattr(m, "_weights_removed", False):
+                post()
+
+
+def layer_inputs(ckpt: str, layer: int, first_token: int, num_tokens: int) -> torch.Tensor:
+    """Real ``input_layernorm`` of ``layer`` over real token embeddings."""
+    return tail_fixtures.normed_embeddings(ckpt, layer, first_token, num_tokens)
+
+
+def select(fast_paths: str) -> None:
+    """Install, for this process, the plan that runs exactly ``fast_paths`` (comma-separated)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+        dispatch,
+    )
+
+    plan_testlib.select(dispatch, fast_paths)
+
+
+class AttentionStack:
+    """Consecutive real DSA attention layers of one TP8 rank with a filled cache.
+
+    Args:
+        ckpt: checkpoint directory.
+        layers: consecutive layer indices; the first must run a full indexer.
+        tp_rank: TP8 rank whose weight shard is loaded.
+        ctx: cached context length per layer (filled by a real prefill).
+        rows: GEN query rows of the single request (``1 + max_draft_len``).
+        reduce_output: build ``o_proj`` with ``reduce_output`` (dense layers).
+        fast_paths: when given, the plan that runs exactly these fast paths is installed
+            (``select``) before the layers are built; otherwise the installed plan is used.
+    """
+
+    def __init__(
+        self,
+        ckpt: str,
+        layers: List[int],
+        tp_rank: int,
+        ctx: int = 4096,
+        rows: int = 6,
+        reduce_output: bool = False,
+        fast_paths: Optional[str] = None,
+    ):
+        from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+
+        if fast_paths is not None:
+            select(fast_paths)
+
+        self.ckpt = ckpt
+        self.layers = list(layers)
+        self.tp_rank = tp_rank
+        self.ctx = ctx
+        self.rows = rows
+        self.max_draft_len = rows - 1
+        self.pretrained = load_pretrained(ckpt)
+        self.model_config = make_model_config(self.pretrained, tp_rank, self.max_draft_len)
+        sparse = self.model_config.sparse_attention_config
+        self.sparse_params = sparse.to_sparse_params(pretrained_config=self.pretrained)
+        self.sparse_metadata_params = sparse.to_sparse_metadata_params(
+            pretrained_config=self.pretrained
+        )
+        self.backend = get_attention_backend("TRTLLM", sparse_params=self.sparse_params)
+        self.aux_stream = torch.cuda.Stream()
+        self.attn: Dict[int, torch.nn.Module] = {}
+        for layer in self.layers:
+            attn = build_attention(self.model_config, layer, self.aux_stream, reduce_output)
+            load_attention_weights(attn, ckpt, layer, tp_rank)
+            self.attn[layer] = attn
+        self.attn_tail_tails = {
+            layer: getattr(attn, "_attn_tail", None) for layer, attn in self.attn.items()
+        }
+        self.kv_cache_manager = self._make_cache()
+        self.kv_cache_manager.add_dummy_requests(
+            request_ids=[REQUEST_ID], token_nums=[ctx + 64], is_gen=False, prepare_resource=True
+        )
+        self._prefill()
+        self.gen_metadata = self._gen_metadata()
+
+    # -- construction --------------------------------------------------------
+    def _one_rank_mapping(self):
+        from tensorrt_llm.mapping import Mapping
+
+        return Mapping(world_size=1, tp_size=1, rank=0)
+
+    def _make_cache(self):
+        import tensorrt_llm.bindings
+        from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
+        from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
+        from tensorrt_llm.llmapi.llm_args import KvCacheConfig
+
+        num_layers = max(self.layers) + 1
+        layer_mask = [i in self.layers for i in range(num_layers)]
+        max_tokens = self.ctx + 256
+        return DSACacheManager(
+            KvCacheConfig(
+                dtype="fp8",
+                max_tokens=max_tokens,
+                enable_block_reuse=False,
+                event_buffer_max_size=0,
+            ),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+            num_layers=len(self.layers),
+            layer_mask=layer_mask,
+            num_kv_heads=1,
+            head_dim=self.pretrained.kv_lora_rank + self.pretrained.qk_rope_head_dim,
+            tokens_per_block=64,
+            max_seq_len=max_tokens,
+            max_batch_size=1,
+            mapping=self._one_rank_mapping(),
+            dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
+            max_num_tokens=max(8192, self.ctx),
+            sparse_attention_config=self.model_config.sparse_attention_config,
+            model_config=self.model_config,
+            pretrained_config=self.pretrained,
+            vocab_size=self.pretrained.vocab_size,
+        )
+
+    def _prefill(self) -> None:
+        from tensorrt_llm._torch.metadata import KVCacheParams
+
+        meta = self.backend.Metadata(
+            seq_lens=torch.tensor([self.ctx], dtype=torch.int),
+            request_ids=[REQUEST_ID],
+            max_num_requests=1,
+            num_contexts=1,
+            prompt_lens=[self.ctx],
+            max_num_tokens=self.ctx,
+            kv_cache_manager=self.kv_cache_manager,
+            kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]),
+            mapping=self._one_rank_mapping(),
+            sparse_metadata_params=self.sparse_metadata_params,
+        )
+        meta.prepare()
+        position_ids = torch.arange(self.ctx, dtype=torch.int32, device="cuda")[None]
+        from tensorrt_llm._torch.distributed import AllReduceParams
+
+        no_ar = AllReduceParams(enable_allreduce=False)
+        with self.scope(meta):
+            for layer in self.layers:
+                hidden = layer_inputs(self.ckpt, layer, TOKEN_OFFSET, self.ctx)
+                self.set_attn_tail(False)
+                self.attn[layer](
+                    position_ids=position_ids,
+                    hidden_states=hidden,
+                    attn_metadata=meta,
+                    all_reduce_params=no_ar,
+                )
+        torch.cuda.synchronize()
+
+    def _gen_metadata(self):
+        from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures
+        from tensorrt_llm._torch.metadata import KVCacheParams
+
+        base = self.backend.Metadata(
+            max_num_requests=1,
+            max_num_tokens=128,
+            max_num_sequences=1,
+            kv_cache_manager=self.kv_cache_manager,
+            mapping=self._one_rank_mapping(),
+            runtime_features=AttentionRuntimeFeatures(
+                has_speculative_draft_tokens=self.max_draft_len > 0
+            ),
+            enable_flash_mla=False,
+            enable_context_mla_with_cached_kv=False,
+            num_heads_per_kv=1,
+            sparse_metadata_params=self.sparse_metadata_params,
+        )
+        meta = base.create_cuda_graph_metadata(1, False, self.max_draft_len)
+        meta.beam_width = 1
+        meta.request_ids = [REQUEST_ID]
+        meta.prompt_lens = [self.ctx]
+        meta.num_contexts = 0
+        meta.num_chunked_ctx_requests = 0
+        meta.kv_cache_params = KVCacheParams(
+            use_cache=True,
+            num_cached_tokens_per_seq=[self.ctx],
+            num_extra_kv_tokens=max(self.max_draft_len - 1, 0),
+        )
+        meta.kv_cache_manager = self.kv_cache_manager
+        if self.max_draft_len > 0:
+            meta.update_spec_dec_param(
+                batch_size=1,
+                is_spec_decoding_enabled=True,
+                is_spec_dec_tree=False,
+                is_spec_dec_dynamic_tree=False,
+                max_draft_len=self.max_draft_len,
+                max_total_draft_tokens=self.max_draft_len,
+                num_contexts=0,
+            )
+        meta.prepare()
+        return meta
+
+    # -- running -------------------------------------------------------------
+    @contextlib.contextmanager
+    def scope(self, meta=None):
+        """Model extra attrs, multi-stream mode and the step's phase, as the engine forward
+        sets them: the GEN metadata (``meta`` None) is a decode step, ``meta`` a prefill."""
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+            dispatch,
+        )
+        from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
+        from tensorrt_llm._torch.utils import model_extra_attrs
+
+        phase = dispatch.Phase.DECODE if meta is None else dispatch.Phase.PREFILL
+        attrs = dict(self.model_config.extra_attrs)
+        attrs["attention_metadata"] = weakref.ref(meta or self.gen_metadata)
+        with model_extra_attrs(attrs), with_multi_stream(True), dispatch.step(phase):
+            yield
+
+    def set_attn_tail(self, enabled: bool, config=None) -> None:
+        """Attach (or detach) each layer's attention_output_tail; optionally set its config."""
+        for layer, attn in self.attn.items():
+            tail = self.attn_tail_tails[layer]
+            if enabled and tail is None:
+                raise RuntimeError(
+                    f"layer {layer} has no attention_output_tail tail "
+                    "(the installed plan does not run attention_output_tail)"
+                )
+            if tail is not None and config is not None:
+                tail.config = config
+            attn._attn_tail = tail if enabled else None
+
+    def gen_inputs(self, layer: int, step: int = 0) -> torch.Tensor:
+        """GEN hidden states ``[rows, hidden]`` for ``layer`` at ``step``."""
+        first = TOKEN_OFFSET + self.ctx + step * self.rows
+        return layer_inputs(self.ckpt, layer, first, self.rows)
+
+    def position_ids(self) -> torch.Tensor:
+        return torch.arange(self.ctx, self.ctx + self.rows, dtype=torch.int32, device="cuda")[None]
+
+    def forward(
+        self, layer: int, hidden: torch.Tensor, position_ids, all_reduce_params
+    ) -> torch.Tensor:
+        """One GEN ``GlmMLA.forward`` of ``layer``; call inside ``scope()``."""
+        return self.attn[layer](
+            position_ids=position_ids,
+            hidden_states=hidden,
+            attn_metadata=self.gen_metadata,
+            all_reduce_params=all_reduce_params,
+        )
+
+    def shutdown(self) -> None:
+        self.kv_cache_manager.shutdown()
+
+
+class LatentRecorder:
+    """Record the latent each layer's GEN tail receives (eager runs only).
+
+    Wraps the native W_UV entry (``MLA._bmm_bf16_out``) and the attention_output_tail entry
+    (``AttnTail.defer_w_uv``) on the given modules, so the test can build the
+    FP64 oracle from the latent the real attention produced.
+    """
+
+    def __init__(self, stack: AttentionStack):
+        self.stack = stack
+        self.latents: Dict[int, torch.Tensor] = {}
+        self._saved: List = []
+
+    def __enter__(self):
+        for layer, attn in self.stack.attn.items():
+            native = attn._bmm_bf16_out
+
+            def bmm(a, b, b_t, out, *rest, _layer=layer, _native=native):
+                # a: latent transposed to [H, R, C].
+                self.latents[_layer] = a.transpose(0, 1).contiguous().clone()
+                return _native(a, b, b_t, out, *rest)
+
+            attn._bmm_bf16_out = bmm
+            self._saved.append((attn, "_bmm_bf16_out", None))
+            tail = self.stack.attn_tail_tails[layer]
+            if tail is not None:
+                defer = tail.defer_w_uv
+
+                def wrapped(mla, meta, latent, output, _layer=layer, _defer=defer):
+                    self.latents[_layer] = latent.contiguous().clone()
+                    return _defer(mla, meta, latent, output)
+
+                tail.defer_w_uv = wrapped
+                self._saved.append((tail, "defer_w_uv", None))
+        return self
+
+    def __exit__(self, *exc):
+        for obj, name, _ in self._saved:
+            delattr(obj, name)
+        self._saved.clear()
+        return False
+
+
+def native_reference(attn, latent: torch.Tensor) -> Optional[torch.Tensor]:
+    """FP64 oracle of the tail on a recorded latent (native BF16 ``v``)."""
+    import tail_reference
+
+    ref, _ = tail_reference.oracle_attn_tail(
+        latent, attn.v_b_proj, attn.o_proj.weight, latent.shape[0]
+    )
+    return ref
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py
new file mode 100644
index 000000000000..65eec6af6e4e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py
@@ -0,0 +1,191 @@
+# 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.
+"""Real-weight fixtures for the attention_output_tail, read from a GLM-5.x checkpoint.
+
+The TP slicing mirrors ``DeepseekV3WeightLoader``: ``kv_b_proj`` is unflattened
+to ``[heads, qk_nope + v, kv_lora]``, chunked along heads per rank, and the
+``v`` rows form ``v_b_proj [heads / tp, v, kv_lora]``; the ROW-parallel
+``o_proj`` takes the rank's chunk of input features (columns).
+
+``compressed_kv_rows`` and ``mix_latent`` build a latent attention output from
+real layer weights (token embeddings → input RMSNorm → ``kv_a_proj_with_mqa``
+→ ``kv_a_layernorm``), mixed by random softmax weights. It is a real-weight, synthetic-activation
+fixture: the rows are convex combinations of real compressed-KV vectors, but
+the attention probabilities are random, not model activations.
+"""
+
+import json
+import os
+from typing import Dict, Iterable, Tuple
+
+import torch
+from utils.llm_data import llm_models_root
+
+GLM52_CKPT_ENV = "GLM52_NVFP4_CKPT"
+
+
+def checkpoint_dir() -> str:
+    """The GLM-5.2 NVFP4 checkpoint: ``$GLM52_NVFP4_CKPT``, else under ``llm_models_root()``."""
+    path = os.environ.get(GLM52_CKPT_ENV)
+    if path:
+        return path
+    root = llm_models_root()
+    return os.path.join(str(root), "GLM-5.2-NVFP4") if root is not None else ""
+
+
+def _index(ckpt_dir: str) -> Dict[str, str]:
+    with open(os.path.join(ckpt_dir, "model.safetensors.index.json")) as f:
+        return json.load(f)["weight_map"]
+
+
+def _read(
+    ckpt_dir: str, weight_map: Dict[str, str], name: str, rows: slice = slice(None)
+) -> torch.Tensor:
+    from safetensors import safe_open
+
+    path = os.path.join(ckpt_dir, weight_map[name])
+    with safe_open(path, framework="pt", device="cpu") as f:
+        return f.get_slice(name)[rows].contiguous()
+
+
+def read_tensor(ckpt_dir: str, name: str) -> torch.Tensor:
+    """Read one checkpoint tensor by its safetensors name (CPU)."""
+    return _read(ckpt_dir, _index(ckpt_dir), name)
+
+
+def read_prefix(ckpt_dir: str, prefix: str) -> Dict[str, torch.Tensor]:
+    """Read every checkpoint tensor whose name starts with ``prefix`` (CPU)."""
+    from safetensors import safe_open
+
+    weight_map = _index(ckpt_dir)
+    by_file: Dict[str, list] = {}
+    for name, fname in weight_map.items():
+        if name.startswith(prefix):
+            by_file.setdefault(fname, []).append(name)
+    out = {}
+    for fname, names in by_file.items():
+        with safe_open(os.path.join(ckpt_dir, fname), framework="pt", device="cpu") as f:
+            for name in names:
+                out[name] = f.get_tensor(name)
+    return out
+
+
+def normed_embeddings(
+    ckpt_dir: str, layer: int, first_token: int, num_tokens: int, device: str = "cuda"
+) -> torch.Tensor:
+    """``input_layernorm`` of ``layer`` over the embeddings of consecutive
+    token IDs ``first_token ... first_token + num_tokens - 1`` (BF16)."""
+    with open(os.path.join(ckpt_dir, "config.json")) as f:
+        eps = json.load(f).get("rms_norm_eps", 1e-6)
+    weight_map = _index(ckpt_dir)
+    emb = _read(
+        ckpt_dir,
+        weight_map,
+        "model.embed_tokens.weight",
+        slice(first_token, first_token + num_tokens),
+    ).to(device)
+    ln = _read(ckpt_dir, weight_map, f"model.layers.{layer}.input_layernorm.weight").to(device)
+    return _rms_norm(emb, ln, eps).contiguous()
+
+
+def load_tail_weights(
+    ckpt_dir: str, layer: int, ranks: Iterable[int], tp_size: int = 8
+) -> Dict[int, Tuple[torch.Tensor, torch.Tensor]]:
+    """Load ``(v_b_proj [H/tp, V, C], o_proj [N, H*V/tp])`` per TP rank (CPU)."""
+    with open(os.path.join(ckpt_dir, "config.json")) as f:
+        cfg = json.load(f)
+    heads = cfg["num_attention_heads"]
+    nope = cfg["qk_nope_head_dim"]
+    v_dim = cfg["v_head_dim"]
+    weight_map = _index(ckpt_dir)
+    prefix = f"model.layers.{layer}.self_attn"
+    kv_b = _read(ckpt_dir, weight_map, f"{prefix}.kv_b_proj.weight")
+    kv_b = kv_b.unflatten(0, [heads, nope + v_dim])
+    o_proj = _read(ckpt_dir, weight_map, f"{prefix}.o_proj.weight")
+    out = {}
+    for rank in ranks:
+        kv_b_rank = torch.chunk(kv_b, tp_size, dim=0)[rank]
+        _, v_rank = kv_b_rank.split([nope, v_dim], dim=1)
+        o_rank = torch.chunk(o_proj, tp_size, dim=1)[rank]
+        out[rank] = (v_rank.contiguous(), o_rank.contiguous())
+    return out
+
+
+def load_full_o_proj(ckpt_dir: str, layer: int) -> torch.Tensor:
+    """Return the unsharded ``o_proj`` weight ``[N, H*V]`` (CPU)."""
+    return _read(ckpt_dir, _index(ckpt_dir), f"model.layers.{layer}.self_attn.o_proj.weight")
+
+
+def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    xf = x.float()
+    y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
+    return (y * weight.float()).to(x.dtype)
+
+
+def compressed_kv_rows(
+    ckpt_dir: str, layer: int, context: int = 2048, token_offset: int = 1000, device: str = "cuda"
+) -> torch.Tensor:
+    """Return ``[context, kv_lora]`` normalized compressed-KV rows of ``layer``.
+
+    Computed from the embeddings of token IDs ``token_offset`` to
+    ``token_offset + context - 1``: input RMSNorm, the compressed part of
+    ``kv_a_proj_with_mqa``, then ``kv_a_layernorm``.
+    """
+    with open(os.path.join(ckpt_dir, "config.json")) as f:
+        cfg = json.load(f)
+    eps = cfg.get("rms_norm_eps", 1e-6)
+    kv_lora = cfg["kv_lora_rank"]
+    weight_map = _index(ckpt_dir)
+    prefix = f"model.layers.{layer}"
+    emb = _read(
+        ckpt_dir,
+        weight_map,
+        "model.embed_tokens.weight",
+        slice(token_offset, token_offset + context),
+    ).to(device)
+    ln = _read(ckpt_dir, weight_map, f"{prefix}.input_layernorm.weight").to(device)
+    kv_a = _read(ckpt_dir, weight_map, f"{prefix}.self_attn.kv_a_proj_with_mqa.weight").to(device)
+    kv_ln = _read(ckpt_dir, weight_map, f"{prefix}.self_attn.kv_a_layernorm.weight").to(device)
+    hidden = _rms_norm(emb, ln, eps)
+    compressed = torch.nn.functional.linear(hidden, kv_a)[:, :kv_lora]
+    return _rms_norm(compressed, kv_ln, eps)
+
+
+def mix_latent(
+    compressed: torch.Tensor,
+    num_rows: int,
+    num_local_heads: int = 8,
+    score_scale: float = 3.0,
+    seed: int = 0,
+) -> torch.Tensor:
+    """Return a ``[num_rows, num_local_heads, kv_lora]`` BF16 latent.
+
+    Each (row, head) is a softmax-weighted average of the ``compressed`` rows
+    with random scores, like a sparse-attention output over a real cache.
+    """
+    device = compressed.device
+    gen = torch.Generator(device=device).manual_seed(seed)
+    scores = (
+        torch.randn(
+            (num_rows, num_local_heads, compressed.shape[0]),
+            generator=gen,
+            device=device,
+            dtype=torch.float32,
+        )
+        * score_scale
+    )
+    probs = torch.softmax(scores, dim=-1)
+    latent = torch.einsum("rht,tc->rhc", probs, compressed.float())
+    return latent.to(torch.bfloat16).contiguous()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py
new file mode 100644
index 000000000000..c13f46af2f37
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py
@@ -0,0 +1,101 @@
+# 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.
+"""Oracles for the attention_output_tail.
+
+``native_attn_tail`` runs the exact native TensorRT-LLM operations of the GEN
+tail: ``MLA._bmm_bf16_out`` for W_UV (the BF16 branch of
+``MLA.forward_absorption_generation``) followed by the ROW-parallel ``o_proj``
+GEMM (``UnquantizedLinearMethod.apply``). ``oracle_attn_tail`` is an
+independent FP64 reference that keeps the native BF16 rounding of the
+intermediate ``v``.
+"""
+
+from types import SimpleNamespace
+from typing import Optional, Tuple
+
+import torch
+
+
+def native_attn_tail(
+    latent: torch.Tensor,
+    w_uv: torch.Tensor,
+    w_o: torch.Tensor,
+    num_rows: int,
+    out_rows: Optional[int] = None,
+    o_proj: Optional[torch.nn.Module] = None,
+    return_v: bool = False,
+):
+    """Run the native W_UV BMM and o_proj GEMM.
+
+    Args:
+        latent: ``[>= num_rows, H * C]`` or ``[>= num_rows, H, C]`` BF16.
+        w_uv: ``[H, V, C]`` BF16 ``v_b_proj``.
+        w_o: ``[N, H * V]`` BF16 local o_proj weight (ignored when ``o_proj``
+            is given).
+        num_rows: live rows written by W_UV.
+        out_rows: rows of the o_proj input buffer (native ``create_output``
+            rows); rows past ``num_rows`` are zero-filled here so the oracle is
+            defined, whereas the native buffer leaves them uninitialized.
+        o_proj: optional native ``Linear`` (ROW parallel, no reduction) to run
+            instead of ``F.linear`` on ``w_o``.
+        return_v: also return the ``[out_rows, H * V]`` intermediate.
+    """
+    from tensorrt_llm._torch.attention.mla import MLA
+
+    heads, v_dim, c_dim = w_uv.shape
+    if out_rows is None:
+        out_rows = num_rows
+    attn_out = torch.zeros((out_rows, heads * v_dim), dtype=w_uv.dtype, device=latent.device)
+    lat = latent[:num_rows].reshape(num_rows, heads, c_dim)
+    v_view = attn_out[:num_rows].view(num_rows, heads, v_dim)
+    native_cfg = SimpleNamespace(use_cute_dsl_bf16_bmm=False)
+    MLA._bmm_bf16_out(
+        native_cfg, lat.transpose(0, 1), w_uv, w_uv.transpose(1, 2), v_view.transpose(0, 1)
+    )
+    if o_proj is not None:
+        partial = o_proj(attn_out)
+    else:
+        partial = torch.nn.functional.linear(attn_out, w_o)
+    if return_v:
+        return partial, attn_out
+    return partial
+
+
+def oracle_attn_tail(
+    latent: torch.Tensor, w_uv: torch.Tensor, w_o: torch.Tensor, num_rows: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """FP64 reference with the native BF16 rounding of the intermediate.
+
+    Returns:
+        ``(partial [num_rows, N] float64, v [num_rows, H * V] bfloat16)``.
+    """
+    heads, v_dim, c_dim = w_uv.shape
+    lat = latent[:num_rows].reshape(num_rows, heads, c_dim).double()
+    v = torch.einsum("rhc,hvc->rhv", lat, w_uv.double())
+    v_bf16 = v.reshape(num_rows, heads * v_dim).to(torch.bfloat16)
+    partial = v_bf16.double() @ w_o.double().t()
+    return partial, v_bf16
+
+
+def error_stats(x: torch.Tensor, ref: torch.Tensor) -> dict:
+    """Normalized RMS / max errors of ``x`` against ``ref`` (both cast to FP64)."""
+    diff = x.double() - ref.double()
+    ref_rms = ref.double().pow(2).mean().sqrt().item()
+    return {
+        "rel_rms": diff.pow(2).mean().sqrt().item() / max(ref_rms, 1e-300),
+        "max_abs": diff.abs().max().item(),
+        "ref_rms": ref_rms,
+        "ref_max": ref.double().abs().max().item(),
+    }
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py
new file mode 100644
index 000000000000..95bac645109e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py
@@ -0,0 +1,130 @@
+# 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.
+"""Default attention_output_tail engine on the real tail weights of all 78 target layers.
+
+For every target layer (3 dense destinations, 75 MoE destinations) the test
+loads the real TP8 ``v_b_proj`` / ``o_proj`` shards of one rank (rank
+``layer % 8``, so every rank is covered across the sweep) and of all eight
+ranks for layers 1, 41 and 76, and runs the default ``cute_stream`` engine
+with PDL at R ∈ {1, 6, 8} on real-weight, synthetic-activation latents (real
+compressed-KV rows of the same layer, random softmax mixtures).
+
+Checks per (layer, rank, rows):
+
+* the frozen partial-output budget of ``test_attn_tail_attn_tail.py`` against the
+  FP64 oracle with the native BF16 rounding of ``v``;
+* the strict intermediate check for the default engine: attention_output_tail's BF16 ``v``
+  equals native ``v`` (``MLA._bmm_bf16_out``) bit for bit, or differs by at
+  most one BF16 ulp per element; the count of non-identical elements is
+  printed. (The relaxed FP64-bound check of ``test_attn_tail_attn_tail.py`` applies
+  to the non-default engines.)
+
+Weights are re-read per layer from the checkpoint, so the device sees 78
+distinct weight sets (no warm reuse across layers).
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
+    DEFAULT_CONFIG,
+    attn_tail,
+    check_tail_shapes,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+TP_SIZE = 8
+LOCAL_HEADS, V_DIM, C_DIM = 8, 256, 512
+NUM_TARGET_LAYERS = 78
+ALL_RANK_LAYERS = (1, 41, 76)
+ROWS = (1, 6, 8)
+_CKPT = tail_fixtures.checkpoint_dir()
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="attention_output_tail kernels target SM100/SM103",
+    ),
+    pytest.mark.skipif(
+        not (_CKPT and os.path.isdir(_CKPT)), reason=f"{tail_fixtures.GLM52_CKPT_ENV} not available"
+    ),
+]
+
+
+def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Per-element distance in BF16 ulps (+0 and -0 are equal)."""
+
+    def ordered(x):
+        bits = x.contiguous().view(torch.int16).to(torch.int32)
+        mag = bits & 0x7FFF
+        return torch.where(bits < 0, -mag, mag)
+
+    return (ordered(a) - ordered(b)).abs()
+
+
+@pytest.mark.parametrize("layer", range(NUM_TARGET_LAYERS))
+def test_default_engine_layer(layer):
+    ranks = list(range(TP_SIZE)) if layer in ALL_RANK_LAYERS else [layer % TP_SIZE]
+    weights = tail_fixtures.load_tail_weights(_CKPT, layer, ranks, TP_SIZE)
+    compressed = tail_fixtures.compressed_kv_rows(_CKPT, layer)
+    destination = "dense" if layer < 3 else "moe"
+    for rank in ranks:
+        w_uv, w_o = (t.cuda() for t in weights[rank])
+        for rows in ROWS:
+            latent = tail_fixtures.mix_latent(
+                compressed, rows, num_local_heads=LOCAL_HEADS, seed=layer * 100 + rank * 10 + rows
+            )
+            assert check_tail_shapes(latent, w_uv, w_o, rows, rows) is None
+            v_attn_tail = torch.empty(
+                (rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda"
+            )
+            attention_output_tail = attn_tail(
+                latent, w_uv, w_o, rows, v_scratch=v_attn_tail, config=DEFAULT_CONFIG
+            )
+            native, v_nat = tail_reference.native_attn_tail(latent, w_uv, w_o, rows, return_v=True)
+            ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, w_o, rows)
+            e_attn_tail = tail_reference.error_stats(attention_output_tail, ref)
+            e_nat = tail_reference.error_stats(native, ref)
+            label = (layer, destination, rank, rows)
+            assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (
+                label,
+                e_attn_tail,
+                e_nat,
+            )
+            assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_nat["ref_max"]), (
+                label,
+                e_attn_tail,
+                e_nat,
+            )
+            torch.testing.assert_close(
+                attention_output_tail, native, rtol=2**-6, atol=2**-6 * e_nat["ref_rms"]
+            )
+            ulps = bf16_ulp_distance(v_attn_tail, v_nat)
+            mismatches = int((ulps > 0).sum())
+            assert int(ulps.max()) <= 1, (label, int(ulps.max()), mismatches)
+            print(
+                f"[attention_output_tail-pdl-stress-layer] layer={layer} dest={destination} "
+                f"rank={rank} rows={rows} rel_rms attention_output_tail={e_attn_tail['rel_rms']:.3e} "
+                f"native={e_nat['rel_rms']:.3e} "
+                f"v_mismatch={mismatches}/{v_attn_tail.numel()} "
+                f"v_max_ulp={int(ulps.max())}"
+            )
+        del w_uv, w_o
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py
new file mode 100644
index 000000000000..b6fc4bc1bb36
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py
@@ -0,0 +1,263 @@
+# 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.
+"""Single-GPU correctness of the attention_output_tail (the attention_output_tail contract).
+
+Oracles: the native TensorRT-LLM ops (``MLA._bmm_bf16_out`` + ROW-parallel
+``Linear``) and an FP64 reference that keeps the native BF16 rounding of the
+W_UV intermediate.
+
+Error budget, frozen before kernel tuning (normalized against the FP64
+reference ``ref``; ``err(x) = ||x - ref||_2 / ||ref||_2``):
+
+* ``err(attention_output_tail) <= 1.25 * err(native) + 1e-6`` on every (layer, rank, rows) case;
+* ``max|attention_output_tail - ref| <= 2 * max|native - ref| + 2**-8 * max|ref|``;
+* ``torch.testing.assert_close(attention_output_tail, native, rtol=2**-6, atol=2**-6 * rms(ref))``.
+
+Fixtures: real GLM-5.2 ``kv_b_proj``/``o_proj`` weights sliced per TP8 rank
+(``GLM52_NVFP4_CKPT``) with two latent kinds: ``derived`` (real-weight,
+synthetic-activation: softmax mixtures of real compressed-KV rows) and
+``randn`` (synthetic). Tests that need the checkpoint skip without it.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
+    ROW_TILE,
+    TailConfig,
+    attn_tail,
+    check_tail_shapes,
+    check_weights,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+TP_SIZE = 8
+HEADS, V_DIM, C_DIM, HIDDEN = 64, 256, 512, 6144
+LOCAL_HEADS = HEADS // TP_SIZE
+# Dense layers 0 and 2, MoE layers 3, 40 and 77 (all TP8 ranks of each).
+LAYERS = (0, 2, 3, 40, 77)
+ROWS = (1, 2, 3, 4, 5, 6, 7, 8, 12, 16)
+
+_CKPT = tail_fixtures.checkpoint_dir()
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="attention_output_tail kernels target SM100/SM103",
+)
+
+needs_ckpt = pytest.mark.skipif(
+    not (_CKPT and os.path.isdir(_CKPT)),
+    reason=f"{tail_fixtures.GLM52_CKPT_ENV} checkpoint not available",
+)
+
+ENGINES = {
+    "fused": TailConfig(oproj_engine="fused"),
+    "cute_stream": TailConfig(oproj_engine="cute_stream"),
+    "triton": TailConfig(oproj_engine="triton"),
+}
+
+_WEIGHTS = {}
+_COMPRESSED = {}
+
+
+def _layer_weights(layer):
+    if layer not in _WEIGHTS:
+        _WEIGHTS[layer] = tail_fixtures.load_tail_weights(_CKPT, layer, range(TP_SIZE), TP_SIZE)
+    return _WEIGHTS[layer]
+
+
+def _latent(kind, layer, rows, seed):
+    if kind == "derived":
+        if layer not in _COMPRESSED:
+            _COMPRESSED[layer] = tail_fixtures.compressed_kv_rows(_CKPT, layer)
+        return tail_fixtures.mix_latent(
+            _COMPRESSED[layer], rows, num_local_heads=LOCAL_HEADS, seed=seed
+        )
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    return torch.randn(
+        (rows, LOCAL_HEADS, C_DIM), generator=gen, device="cuda", dtype=torch.float32
+    ).to(torch.bfloat16)
+
+
+def _check_budget(attention_output_tail, native, ref, label):
+    e_attn_tail = tail_reference.error_stats(attention_output_tail, ref)
+    e_nat = tail_reference.error_stats(native, ref)
+    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (label, e_attn_tail, e_nat)
+    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_nat["ref_max"]), (
+        label,
+        e_attn_tail,
+        e_nat,
+    )
+    torch.testing.assert_close(
+        attention_output_tail, native, rtol=2**-6, atol=2**-6 * e_nat["ref_rms"]
+    )
+    return e_attn_tail, e_nat
+
+
+@needs_ckpt
+@pytest.mark.parametrize("layer", LAYERS)
+def test_o_proj_slice_matches_native_row_linear(layer):
+    """The fixture's o_proj slice equals what a native ROW Linear loads."""
+    from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
+    from tensorrt_llm.mapping import Mapping
+
+    full = tail_fixtures.load_full_o_proj(_CKPT, layer)
+    weights = _layer_weights(layer)
+    for rank in (0, 3, 7):
+        linear = Linear(
+            HEADS * V_DIM,
+            HIDDEN,
+            bias=False,
+            dtype=torch.bfloat16,
+            mapping=Mapping(world_size=TP_SIZE, tp_size=TP_SIZE, rank=rank),
+            tensor_parallel_mode=TensorParallelMode.ROW,
+            reduce_output=False,
+        ).cuda()
+        linear.load_weights([{"weight": full}])
+        assert torch.equal(linear.weight.cpu(), weights[rank][1])
+
+
+@needs_ckpt
+@pytest.mark.parametrize("layer", LAYERS)
+@pytest.mark.parametrize("kind", ["derived", "randn"])
+@pytest.mark.parametrize("use_pdl", [True, False])
+@pytest.mark.parametrize("engine", list(ENGINES))
+def test_tail_matches_native_real_weights(layer, kind, use_pdl, engine):
+    """All TP8 ranks and row counts: attention_output_tail vs native vs FP64 within the budget."""
+    worst = {"attention_output_tail": 0.0, "native": 0.0}
+    for rank, (w_uv_cpu, w_o_cpu) in _layer_weights(layer).items():
+        w_uv, w_o = w_uv_cpu.cuda(), w_o_cpu.cuda()
+        for rows in ROWS:
+            latent = _latent(kind, layer, rows, seed=rank * 100 + rows)
+            assert check_tail_shapes(latent, w_uv, w_o, rows, rows) is None
+            attention_output_tail = attn_tail(
+                latent, w_uv, w_o, rows, use_pdl=use_pdl, config=ENGINES[engine]
+            )
+            native = tail_reference.native_attn_tail(latent, w_uv, w_o, rows)
+            ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, w_o, rows)
+            e_attn_tail, e_nat = _check_budget(
+                attention_output_tail, native, ref, (layer, rank, rows, kind)
+            )
+            worst["attention_output_tail"] = max(
+                worst["attention_output_tail"], e_attn_tail["rel_rms"]
+            )
+            worst["native"] = max(worst["native"], e_nat["rel_rms"])
+    print(
+        f"[attention_output_tail-err] layer={layer} kind={kind} pdl={use_pdl} engine={engine} "
+        f"worst_rel_rms attention_output_tail={worst['attention_output_tail']:.3e} native={worst['native']:.3e}"
+    )
+
+
+@needs_ckpt
+@pytest.mark.parametrize("layer", (2, 40))
+@pytest.mark.parametrize("engine", list(ENGINES))
+def test_intermediate_rounding_matches_native(layer, engine):
+    """The W_UV intermediate is a correctly rounded BF16 of the exact product.
+
+    Both attention_output_tail and native ``v`` must lie within one BF16 rounding step of the
+    FP64 product plus an FP32 accumulation-order term (``2**-20`` of the sum
+    of absolute products, which matters only under cancellation); elements
+    that differ between attention_output_tail and native must stay rare.
+    """
+    for rank in (0, 5):
+        w_uv_cpu, w_o_cpu = _layer_weights(layer)[rank]
+        w_uv, w_o = w_uv_cpu.cuda(), w_o_cpu.cuda()
+        rows = 6
+        latent = _latent("derived", layer, rows, seed=rank)
+        v_attn_tail = torch.empty((rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
+        attn_tail(latent, w_uv, w_o, rows, v_scratch=v_attn_tail, config=ENGINES[engine])
+        _, v_nat = tail_reference.native_attn_tail(latent, w_uv, w_o, rows, return_v=True)
+        lat64 = latent.reshape(rows, LOCAL_HEADS, C_DIM).double()
+        v64 = torch.einsum("rhc,hvc->rhv", lat64, w_uv.double()).reshape(rows, -1)
+        mag = torch.einsum("rhc,hvc->rhv", lat64.abs(), w_uv.double().abs()).reshape(rows, -1)
+        bound = torch.finfo(torch.bfloat16).eps * v64.abs() + 2**-20 * mag
+        for name, v in (("attention_output_tail", v_attn_tail), ("native", v_nat)):
+            err = (v.double() - v64).abs()
+            assert bool((err <= bound).all()), (name, (err - bound).max())
+        mismatch = (v_attn_tail != v_nat).float().mean().item()
+        assert mismatch < 0.01, mismatch
+        print(
+            f"[attention_output_tail-v] layer={layer} rank={rank} engine={engine} "
+            f"v_mismatch_fraction={mismatch:.2e}"
+        )
+
+
+@pytest.mark.parametrize(
+    "live,out_rows", [(1, 4), (6, 8), (6, 9), (6, 16), (6, 20), (9, 12), (16, 40)]
+)
+@pytest.mark.parametrize("engine", list(ENGINES))
+def test_padded_rows_are_zero(live, out_rows, engine):
+    torch.manual_seed(live)
+    w_uv = (torch.randn(LOCAL_HEADS, V_DIM, C_DIM, device="cuda") * 0.02).bfloat16()
+    w_o = (torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda") * 0.02).bfloat16()
+    latent = torch.randn(live, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
+    out = attn_tail(latent, w_uv, w_o, live, out_rows=out_rows, config=ENGINES[engine])
+    assert out.shape == (out_rows, HIDDEN)
+    assert bool((out[live:] == 0).all())
+    native = tail_reference.native_attn_tail(latent, w_uv, w_o, live)
+    ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, w_o, live)
+    _check_budget(out[:live], native, ref, (live, out_rows))
+
+
+@pytest.mark.parametrize("use_pdl", [True, False])
+@pytest.mark.parametrize("engine", list(ENGINES))
+def test_cuda_graph_replay(use_pdl, engine):
+    """Replay a captured tail with fresh latent contents and padded rows."""
+    torch.manual_seed(0)
+    rows, out_rows = 6, 8
+    w_uv = (torch.randn(LOCAL_HEADS, V_DIM, C_DIM, device="cuda") * 0.02).bfloat16()
+    w_o = (torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda") * 0.02).bfloat16()
+    # Latent rows live inside a wider buffer, like the backend's output.
+    latent_buf = torch.zeros(rows, LOCAL_HEADS * C_DIM, device="cuda").bfloat16()
+    cfg = ENGINES[engine]
+    attn_tail(latent_buf, w_uv, w_o, rows, out_rows=out_rows, use_pdl=use_pdl, config=cfg)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = attn_tail(latent_buf, w_uv, w_o, rows, out_rows=out_rows, use_pdl=use_pdl, config=cfg)
+    for step in range(64):
+        latent_buf.copy_(torch.randn(rows, LOCAL_HEADS * C_DIM, device="cuda").bfloat16())
+        graph.replay()
+        native = tail_reference.native_attn_tail(latent_buf, w_uv, w_o, rows)
+        ref, _ = tail_reference.oracle_attn_tail(latent_buf, w_uv, w_o, rows)
+        _check_budget(out[:rows].clone(), native, ref, ("graph", step))
+        assert bool((out[rows:] == 0).all())
+
+
+def test_rejects_unsupported_inputs():
+    w_uv = torch.zeros(LOCAL_HEADS, V_DIM, C_DIM, device="cuda").bfloat16()
+    w_o = torch.zeros(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda").bfloat16()
+    latent = torch.zeros(6, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
+    assert check_tail_shapes(latent, w_uv, w_o, 6, 6) is None
+    assert check_tail_shapes(latent, w_uv, w_o, ROW_TILE + 1, ROW_TILE + 1) is not None
+    assert check_tail_shapes(latent, w_uv, w_o, 0, 0) is not None
+    assert check_tail_shapes(latent.float(), w_uv, w_o, 6, 6) is not None
+    assert check_weights(w_uv.to(torch.float8_e4m3fn), w_o) is not None
+    assert check_weights(w_uv, w_o.t()) is not None
+    assert check_weights(w_uv, w_o[:, :1024]) is not None
+    assert check_tail_shapes(latent.transpose(0, 1), w_uv, w_o, 6, 6) is not None
+
+
+def test_registered_entry_points():
+    assert dispatch.get_impl("attention_output_tail", "attn_tail") is attn_tail
+    make_tail = dispatch.get_impl("attention_output_tail", "make_attn_tail")
+    assert callable(make_tail)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py
new file mode 100644
index 000000000000..3ba4d02e55dc
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py
@@ -0,0 +1,244 @@
+# 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.
+"""attention_output_tail contract bridge with a synthetic latent, into the native reduction/norm.
+
+The latent is ``tail_fixtures.mix_latent`` (softmax mixes of real compressed-KV
+rows with random scores; no attention kernel runs), so this file checks only
+the consumer side of the bridge. The native attention producer → attention_output_tail → native
+reduction/norm chain runs in ``test_attn_tail_mla_forward_mpi.py``.
+
+Launch one process per GPU under MPI (``srun --mpi=pmix -n <tp>``); every
+rank runs the same tests in the same order. Each rank holds its own TP8 shard
+of real layer weights (rank ``r`` uses TP8 shard ``r``; with fewer than eight
+ranks the all-reduce sums a subset of shards, which still exercises the
+bridge). Both destinations are covered:
+
+* MoE layer (3): ``o_proj`` with ``enable_allreduce=False`` → decoder-level
+  ``AllReduce(RESIDUAL_RMS_NORM)`` with the real ``post_attention_layernorm``.
+* Dense layer (0): ``o_proj`` reduces inside ``Linear.forward``, then the real
+  ``post_attention_layernorm(hidden, residual)``.
+
+The native chain uses the native W_UV BMM and the native ROW ``Linear``; the
+attention_output_tail chain goes through ``AttnTail.defer_w_uv`` / ``AttnTail.project``, which
+returns the rank-local partial without communicating, and then reduces it as
+``MLA.forward`` does (``o_proj.all_reduce`` with the caller's params, only
+where ``o_proj`` reduces), with the same ``Linear`` and ``AllReduce``
+modules. ``project`` must not call the all-reduce. Outputs must match within
+the frozen attention_output_tail tolerance and be bit-identical across ranks.
+"""
+
+import hashlib
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
+    AttnTail,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+TP8 = 8
+HEADS, V_DIM, C_DIM, HIDDEN = 64, 256, 512, 6144
+LOCAL_HEADS = HEADS // TP8
+EPS = 1e-5
+
+_CKPT = tail_fixtures.checkpoint_dir()
+
+
+def _comm():
+    from mpi4py import MPI
+
+    return MPI.COMM_WORLD
+
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"),
+    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
+    pytest.mark.skipif(_comm().Get_size() < 2, reason="needs >= 2 MPI ranks"),
+]
+
+
+_SETUP = {}
+
+
+def _setup(layer):
+    if layer not in _SETUP:
+        _SETUP[layer] = _build(layer)
+    return _SETUP[layer]
+
+
+def _build(layer):
+    from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
+    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+    from tensorrt_llm.mapping import Mapping
+
+    comm = _comm()
+    rank, world = comm.Get_rank(), comm.Get_size()
+    torch.cuda.set_device(rank % torch.cuda.device_count())
+    mapping = Mapping(
+        world_size=world, tp_size=world, rank=rank, gpus_per_node=torch.cuda.device_count()
+    )
+    w_uv, w_o = tail_fixtures.load_tail_weights(_CKPT, layer, [rank], TP8)[rank]
+    w_uv = w_uv.cuda()
+    # A ROW Linear over `world` ranks whose local slice is this rank's TP8
+    # shard: load the [N, world * K_local] matrix built from TP8 shards.
+    o_proj = Linear(
+        world * LOCAL_HEADS * V_DIM,
+        HIDDEN,
+        bias=False,
+        dtype=torch.bfloat16,
+        mapping=mapping,
+        tensor_parallel_mode=TensorParallelMode.ROW,
+        reduce_output=True,
+    ).cuda()
+    o_proj.weight.data.copy_(w_o.cuda())
+    norm_w = tail_fixtures.read_tensor(
+        _CKPT, f"model.layers.{layer}.post_attention_layernorm.weight"
+    ).cuda()
+    norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda()
+    norm.weight.data.copy_(norm_w)
+    return rank, world, mapping, w_uv, o_proj, norm
+
+
+def _digest(t):
+    return hashlib.sha256(t.contiguous().view(torch.uint8).cpu().numpy().tobytes()).hexdigest()
+
+
+def _tail_ns(w_uv, o_proj):
+    tail = AttnTail("bridge")
+    tail._static_ok = True  # The weights are real; MLA-type checks need a module.
+    mla = SimpleNamespace(v_b_proj=w_uv, o_proj=o_proj)
+    return tail, mla
+
+
+def _run_tail(tail, mla, latent, out_rows, all_reduce_params):
+    """The attention_output_tail, then the reduction ``MLA.forward`` applies to it."""
+    rows = latent.shape[0]
+    attn_out = torch.empty((out_rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
+    meta = SimpleNamespace(num_contexts=0, num_tokens=rows)
+    o_proj = mla.o_proj
+    # Count calls of o_proj's AllReduce module while the tail runs.
+    reduce = o_proj.all_reduce
+    native_forward = reduce.forward
+    calls = []
+
+    def counting(*args, **kwargs):
+        calls.append(args)
+        return native_forward(*args, **kwargs)
+
+    reduce.forward = counting
+    try:
+        assert tail.arm(mla, meta)
+        assert tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
+        partial = tail.project(mla, attn_out)
+    finally:
+        del reduce.forward
+    assert not calls, "the attention_output_tail tail called the all-reduce"
+    if o_proj.reduce_output:
+        partial = o_proj.all_reduce(partial, all_reduce_params=all_reduce_params)
+    return partial
+
+
+@pytest.mark.parametrize("rows", [1, 6])
+def test_moe_layer_bridge(rows):
+    from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams
+
+    rank, world, mapping, w_uv, o_proj, norm = _setup(3)
+    allreduce = AllReduce(mapping=mapping, dtype=torch.bfloat16)
+    comp = tail_fixtures.compressed_kv_rows(_CKPT, 3)
+    latent = tail_fixtures.mix_latent(comp, rows, LOCAL_HEADS, seed=rank)
+    gen = torch.Generator(device="cuda").manual_seed(1234)
+    residual = torch.randn(
+        (rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
+    ).bfloat16()
+
+    def fused(partial):
+        return allreduce(
+            partial,
+            all_reduce_params=AllReduceParams(
+                fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                residual=residual,
+                norm_weight=norm.weight,
+                eps=EPS,
+                trigger_completion_at_end=False,
+            ),
+        )
+
+    no_ar = AllReduceParams(enable_allreduce=False)
+    native_partial = tail_reference.native_attn_tail(
+        latent, w_uv, None, rows, o_proj=lambda x: o_proj(x, all_reduce_params=no_ar)
+    )
+    hs_nat, res_nat = fused(native_partial)
+    tail, mla = _tail_ns(w_uv, o_proj)
+    attn_tail_partial = _run_tail(tail, mla, latent, rows, no_ar)
+    hs_attn_tail, res_attn_tail = fused(attn_tail_partial)
+    torch.cuda.synchronize()
+
+    ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, o_proj.weight, rows)
+    e_attn_tail = tail_reference.error_stats(attn_tail_partial, ref)
+    e_nat = tail_reference.error_stats(native_partial, ref)
+    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (e_attn_tail, e_nat)
+    torch.testing.assert_close(res_attn_tail, res_nat, rtol=2**-6, atol=2**-6)
+    torch.testing.assert_close(hs_attn_tail, hs_nat, rtol=2**-6, atol=2**-5)
+    digests = _comm().allgather((_digest(hs_attn_tail), _digest(res_attn_tail)))
+    assert len(set(digests)) == 1, "post-norm outputs differ across ranks"
+    print(
+        f"[attention_output_tail-bridge] moe rank={rank}/{world} rows={rows} "
+        f"partial_err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e} "
+        f"hs_maxdiff={(hs_attn_tail.float() - hs_nat.float()).abs().max().item():.3e}"
+    )
+
+
+@pytest.mark.parametrize("rows", [1, 6])
+def test_dense_layer_bridge(rows):
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    rank, world, mapping, w_uv, o_proj, norm = _setup(0)
+    comp = tail_fixtures.compressed_kv_rows(_CKPT, 0)
+    latent = tail_fixtures.mix_latent(comp, rows, LOCAL_HEADS, seed=rank)
+    gen = torch.Generator(device="cuda").manual_seed(4321)
+    residual = torch.randn(
+        (rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
+    ).bfloat16()
+    with_ar = AllReduceParams(enable_allreduce=True)
+
+    native_reduced = tail_reference.native_attn_tail(
+        latent, w_uv, None, rows, o_proj=lambda x: o_proj(x, all_reduce_params=with_ar)
+    )
+    hs_nat, res_nat = norm(native_reduced, residual.clone())
+    tail, mla = _tail_ns(w_uv, o_proj)
+    attn_tail_reduced = _run_tail(tail, mla, latent, rows, with_ar)
+    hs_attn_tail, res_attn_tail = norm(attn_tail_reduced, residual.clone())
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        attn_tail_reduced,
+        native_reduced,
+        rtol=2**-6,
+        atol=2**-6 * native_reduced.float().pow(2).mean().sqrt().item(),
+    )
+    torch.testing.assert_close(hs_attn_tail, hs_nat, rtol=2**-6, atol=2**-5)
+    digests = _comm().allgather((_digest(attn_tail_reduced), _digest(hs_attn_tail)))
+    assert len(set(digests)) == 1, "reduced outputs differ across ranks"
+    print(
+        f"[attention_output_tail-bridge] dense rank={rank}/{world} rows={rows} "
+        f"hs_maxdiff={(hs_attn_tail.float() - hs_nat.float()).abs().max().item():.3e}"
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py
new file mode 100644
index 000000000000..5cba06b6e302
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py
@@ -0,0 +1,290 @@
+# 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.
+"""The attention_output_tail contract on one GPU: no attention_output_tail entry point communicates.
+
+The attention_output_tail contract: "attention_output_tail owns W_UV + local o_proj.
+router_gemm/native all-reduce consumes the identical partial. No communication is hidden inside
+attention_output_tail."
+
+The real DSA ``MLA`` of TP8 rank 0 with real layer-3 weights gets an
+``o_proj`` that reduces its output, as in the model at TP8
+(``reduce_output=True``). A recording stand-in (``SpyAllReduce``) replaces
+the all-reduce module. It reports an NCCL symmetric-memory window, as the
+model's AUTO all-reduce does by default. The ``AttnTail`` entry points
+(``arm``, ``before_attention``, ``defer_w_uv``, ``project``) are driven with
+synthetic latents for every engine and named schedule the package ships:
+attention_output_tail's default, the ``attention_output_tail_prefetch`` and
+``attention_output_into_hop`` schedules (``INTO_HOP_SPARSE_DECODE_SCHEDULE`` is the served
+one), the into-hop publication, the ``fused`` and ``triton`` engines, and
+``attention_output_tail_prefetch``'s ``native`` engine with the cuBLAS and ring o_proj. The decoder
+layer's two ``AllReduceParams`` are used: enabled (dense layers, where ``Linear.forward`` reduces)
+and disabled (MoE layers, where the reduction is fused downstream).
+
+Checks:
+
+* no all-reduce call happens while any tail entry point runs;
+* ``project`` returns the rank-local partial: for the attention_output_tail kernels it equals
+  the functional ``attn_tail`` bit for bit; the ``native`` engine returns
+  the native ``o_proj`` GEMM output (bit-identical for cuBLAS, within the
+  frozen budget for the ring kernel) before any reduction;
+* captured in a CUDA graph, the tail still enqueues no all-reduce.
+
+``MLA.forward``'s reduction of this partial is tested at the real call site
+in ``test_attn_tail_contract_o_forward.py`` and ``test_attn_tail_mla_forward_mpi.py``.
+"""
+
+import dataclasses
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+import tail_reference  # noqa: E402
+import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
+    DEFAULT_CONFIG,
+    INTO_HOP_SCHEDULE,
+    INTO_HOP_SPARSE_DECODE_SCHEDULE,
+    TAIL_PREFETCH_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+    TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE,
+    TailConfig,
+    attn_tail,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
+    AttnTail,
+)
+
+pytestmark = tail_adapter.pytestmark
+
+ROWS = 6
+_R = dataclasses.replace
+
+ATTN_TAIL_KERNEL_CONFIGS = {
+    "attention_output_tail": DEFAULT_CONFIG,
+    "attention_output_tail_prefetch": TAIL_PREFETCH_SCHEDULE,
+    "tail_prefetch_sparse_mla": TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+    "tail_prefetch_sparse_mla_attn_hop": TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+    "attention_output_into_hop": INTO_HOP_SCHEDULE,
+    "into_hop_sparse_mla": INTO_HOP_SPARSE_DECODE_SCHEDULE,
+    "fused": TailConfig(oproj_engine="fused"),
+    "triton": TailConfig(oproj_engine="triton"),
+}
+
+
+TAIL_ENTRY_POINTS = ("arm", "before_attention", "defer_w_uv", "project")
+
+
+class SpyAllReduce:
+    """Stand-in for ``o_proj.all_reduce`` that records every call.
+
+    ``flag.in_tail`` is set while an attention_output_tail entry point runs. An enabled call
+    returns twice its input (two identical ranks), so a reduction that
+    happens is visible in the values; a disabled call returns its input, as
+    ``AllReduce.forward`` does.
+    """
+
+    def __init__(self, flag=None):
+        self.calls = []
+        self.flag = flag if flag is not None else SimpleNamespace(in_tail=False)
+
+    def uses_nccl_symmetric_memory_window(self) -> bool:
+        return True
+
+    def __call__(self, input, all_reduce_params=None):
+        self.calls.append(
+            SimpleNamespace(
+                in_tail=self.flag.in_tail, params=all_reduce_params, data_ptr=input.data_ptr()
+            )
+        )
+        if all_reduce_params is not None and not all_reduce_params.enable_allreduce:
+            return input
+        return input * 2
+
+
+def guard_tail(tail: AttnTail, spy: SpyAllReduce) -> None:
+    """Flag every all-reduce call made while a tail entry point runs."""
+    for name in TAIL_ENTRY_POINTS:
+        entry = getattr(tail, name)
+
+        def guarded(*args, _entry=entry, **kwargs):
+            spy.flag.in_tail = True
+            try:
+                return _entry(*args, **kwargs)
+            finally:
+                spy.flag.in_tail = False
+
+        setattr(tail, name, guarded)
+
+
+_MLA = {}
+
+
+def _reducing_mla():
+    """The real DSA MLA with an ``o_proj`` that reduces into a spy."""
+    if "m" not in _MLA:
+        mla = tail_adapter._build_mla()
+        mla.o_proj.reduce_output = True
+        _MLA["m"] = mla
+    mla = _MLA["m"]
+    spy = SpyAllReduce()
+    mla.o_proj.all_reduce = spy
+    return mla, spy
+
+
+def _params(enable: bool):
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    return AllReduceParams(enable_allreduce=enable)
+
+
+def _transport(rank=3):
+    """A post_attention_hop transport stand-in on one GPU (see attention_output_tail/test_into_hop_producer.py)."""
+    words = 8 * 8 * 6144 * 2 // 4
+    mc = torch.full((3 * words,), -(2**31), dtype=torch.int32, device="cuda")
+    flags = torch.tensor([1, 0, 0, 0], dtype=torch.int32, device="cuda")
+    return SimpleNamespace(mc_words=mc, flags=flags, rank=rank, cur=1)
+
+
+def _tail(cfg, mla, spy):
+    tail = AttnTail("layer3", cfg)
+    if cfg.successor_prefetch:
+        # The router weight an attention_output_tail_prefetch tail stages for the op after it.
+        tail.set_successor_weights([torch.zeros((256, 6144), dtype=torch.bfloat16, device="cuda")])
+    guard_tail(tail, spy)
+    return tail
+
+
+def _attn_tail_call(tail, mla, latent, params, hop_transport=None):
+    rows = latent.shape[0]
+    attn_out = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    if hop_transport is not None:
+        tail.request_into_hop(hop_transport, rows)
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point=tail.config.prefetch_fork)
+    assert tail.defer_w_uv(mla, meta, latent, attn_out)
+    return tail.project(mla, attn_out, params)
+
+
+@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
+@pytest.mark.parametrize("name", list(ATTN_TAIL_KERNEL_CONFIGS))
+def test_attn_tail_kernels_return_partial_without_communication(name, enable):
+    cfg = ATTN_TAIL_KERNEL_CONFIGS[name]
+    mla, spy = _reducing_mla()
+    latent = tail_adapter._latent(ROWS, seed=11)
+    want = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, ROWS, config=cfg)
+    tail = _tail(cfg, mla, spy)
+    got = _attn_tail_call(tail, mla, latent, _params(enable))
+    torch.cuda.synchronize()
+    assert spy.calls == [], "the attention_output_tail tail called the all-reduce"
+    # Engaged with a window all-reduce attached, as at TP8 in the model.
+    assert tail.static_reason is None
+    assert tail.engaged_calls == 1 and tail._pending is None
+    assert torch.equal(got, want), name
+    if cfg.successor_prefetch:
+        assert tail.prefetch_forks == 1 and not tail._prefetch.forked
+
+
+@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
+@pytest.mark.parametrize("name", ["attention_output_into_hop", "into_hop_sparse_mla"])
+def test_into_hop_publication_has_no_communication(name, enable):
+    cfg = ATTN_TAIL_KERNEL_CONFIGS[name]
+    mla, spy = _reducing_mla()
+    latent = tail_adapter._latent(ROWS, seed=12)
+    want = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, ROWS, config=cfg)
+    tail = _tail(cfg, mla, spy)
+    got = _attn_tail_call(tail, mla, latent, _params(enable), hop_transport=_transport())
+    torch.cuda.synchronize()
+    assert tail.take_into_hop_published()
+    assert spy.calls == []
+    assert torch.equal(got, want)
+
+
+def _native_engine_call(tail, mla, latent, params):
+    """What the op and MLA.forward do with attention_output_tail_prefetch's ``native`` engine."""
+    from tensorrt_llm._torch.attention.mla import MLA
+
+    rows = latent.shape[0]
+    attn_out = torch.zeros((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point="attn")
+    assert not tail.defer_w_uv(mla, meta, latent, attn_out)
+    v_view = attn_out.view(rows, mla.num_heads_tp, mla.v_head_dim)
+    MLA._bmm_bf16_out(
+        mla,
+        latent.transpose(0, 1),
+        mla.v_b_proj,
+        mla.v_b_proj.transpose(1, 2),
+        v_view.transpose(0, 1),
+    )
+    return tail.project(mla, attn_out, params), attn_out
+
+
+@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
+@pytest.mark.parametrize("join", ["before_oproj", "after_oproj"])
+@pytest.mark.parametrize("oproj", ["cublas", "ring"])
+def test_native_engine_returns_gemm_output_without_communication(oproj, join, enable):
+    mla, spy = _reducing_mla()
+    cfg = _R(TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE, native_oproj=oproj, prefetch_join=join)
+    tail = _tail(cfg, mla, spy)
+    latent = tail_adapter._latent(ROWS, seed=13)
+    got, attn_out = _native_engine_call(tail, mla, latent, _params(enable))
+    gemm = mla.o_proj.apply_linear(attn_out, None)
+    torch.cuda.synchronize()
+    assert spy.calls == []
+    assert tail.engaged_calls == 1 and tail.prefetch_forks == 1
+    assert not tail._prefetch.forked
+    if oproj == "cublas":
+        assert torch.equal(got, gemm)
+    else:
+        ref = attn_out.double() @ mla.o_proj.weight.double().t()
+        e_ring = tail_reference.error_stats(got, ref)
+        e_nat = tail_reference.error_stats(gemm, ref)
+        assert e_ring["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6
+        assert e_ring["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_ring["ref_max"])
+
+
+@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
+@pytest.mark.parametrize(
+    "name", ["attention_output_tail", "into_hop_sparse_mla", "attention_output_tail_prefetch"]
+)
+def test_graph_capture_enqueues_no_communication(name, enable):
+    cfg = ATTN_TAIL_KERNEL_CONFIGS[name]
+    mla, spy = _reducing_mla()
+    tail = _tail(cfg, mla, spy)
+    lat = tail_adapter._latent(ROWS, seed=14).contiguous()
+    params = _params(enable)
+    _attn_tail_call(tail, mla, lat, params)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = _attn_tail_call(tail, mla, lat, params)
+    assert spy.calls == []
+    for seed in range(4):
+        lat.copy_(tail_adapter._latent(ROWS, seed=140 + seed))
+        graph.replay()
+        torch.cuda.synchronize()
+        want = attn_tail(lat, mla.v_b_proj, mla.o_proj.weight, ROWS, config=cfg)
+        assert torch.equal(captured, want)
+    assert spy.calls == []
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py
new file mode 100644
index 000000000000..4817785f65e8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py
@@ -0,0 +1,276 @@
+# 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.
+"""The attention_output_tail contract at the real call site: ``GlmMLA.forward`` reduces attention_output_tail's partial.
+
+Runs with the plan that runs only ``attention_output_tail`` (opt-in:
+``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED=1``).
+
+Builds TP8 rank 0 of real GLM-5.2 layers 2 and 3 (``real_mla_harness``: all
+attention weights real, FP8 KV cache filled by a real prefill) and runs the
+GEN step (R=6) as a decode step through ``GlmMLA.forward``, so the latent comes from the native
+DSA attention kernel. ``o_proj`` reduces its output, as in the model at TP8;
+on one GPU a recording stand-in (``test_attn_tail_contract_o.SpyAllReduce``)
+replaces the all-reduce. Both ``AllReduceParams`` of the decoder layer are
+used: enabled (dense destination, where ``Linear.forward`` reduces) and
+disabled (MoE destination, where the reduction is fused downstream). Tails:
+attention_output_tail's default schedule and ``INTO_HOP_SPARSE_DECODE_SCHEDULE``, the served
+one.
+
+Per layer call:
+
+* native: exactly one all-reduce call (from ``Linear.forward``) with the
+  call's ``AllReduceParams`` object;
+* attention_output_tail: the tail engages; no all-reduce call while a tail entry point runs;
+  exactly one call, from ``GlmMLA.forward``, with the same params object, on
+  the tensor ``project`` returned; the module output is that call's result,
+  and each arm is within the frozen attention_output_tail budget against the FP64 oracle of
+  the latent it received;
+* CUDA graph: the attention_output_tail step captured with both params records the same calls
+  and every replay with new inputs is within the frozen budget of the
+  native eager step on the same inputs.
+"""
+
+import os
+import sys
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
+    DEFAULT_CONFIG,
+    INTO_HOP_SPARSE_DECODE_SCHEDULE,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
+    AttnTail,
+)
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+import real_mla_harness as harness  # noqa: E402
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+from test_attn_tail_contract_o import TAIL_ENTRY_POINTS, SpyAllReduce  # noqa: E402
+
+_CKPT = tail_fixtures.checkpoint_dir()
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="attention_output_tail targets SM100/SM103",
+    ),
+    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
+    pytest.mark.skipif(
+        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
+        reason="needs attention_output_tail selected (opt-in)",
+    ),
+    pytest.mark.threadleak(enabled=False),
+]
+
+LAYERS = [2, 3]
+CONFIGS = {
+    "attention_output_tail": DEFAULT_CONFIG,
+    "into_hop_sparse_mla": INTO_HOP_SPARSE_DECODE_SCHEDULE,
+}
+_STACK = {}
+
+
+@pytest.fixture
+def probe(monkeypatch):
+    """Instrument every ``AttnTail`` entry point (class level, so the
+    harness's per-instance latent recorder composes with it): set
+    ``in_tail`` while one runs and keep what ``project`` returned."""
+    state = SimpleNamespace(in_tail=False, returned={})
+    for name in TAIL_ENTRY_POINTS:
+        entry = getattr(AttnTail, name)
+
+        def run(self, *args, _entry=entry, _name=name, **kwargs):
+            state.in_tail = True
+            try:
+                out = _entry(self, *args, **kwargs)
+            finally:
+                state.in_tail = False
+            if _name == "project" and out is not None:
+                state.returned[self.layer_key] = out
+            return out
+
+        monkeypatch.setattr(AttnTail, name, run)
+    return state
+
+
+@pytest.fixture
+def setup(probe):
+    """The stack with a fresh spy all-reduce on every layer's ``o_proj``."""
+    if "s" not in _STACK:
+        stack = harness.AttentionStack(
+            _CKPT, LAYERS, tp_rank=0, ctx=4096, rows=6, fast_paths="attention_output_tail"
+        )
+        for layer in LAYERS:
+            assert stack.attn_tail_tails[layer] is not None, (
+                "no attention_output_tail tail: call sites or attention_output_tail missing"
+            )
+            stack.attn[layer].o_proj.reduce_output = True
+        _STACK["s"] = stack
+    stack = _STACK["s"]
+    spies = {}
+    for layer in LAYERS:
+        spies[layer] = SpyAllReduce(flag=probe)
+        stack.attn[layer].o_proj.all_reduce = spies[layer]
+    yield stack, spies, probe
+    stack.set_attn_tail(True, config=DEFAULT_CONFIG)
+
+
+def _params(enable):
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    return AllReduceParams(enable_allreduce=enable)
+
+
+def _step(stack, attention_output_tail, inputs, params, record=False, config=None):
+    """One GEN step over the stack; outputs are cloned as produced."""
+    stack.set_attn_tail(attention_output_tail, config=config)
+    pos = stack.position_ids()
+    outs = {}
+    recorder = harness.LatentRecorder(stack) if record else nullcontext()
+    with stack.scope(), recorder:
+        for layer in LAYERS:
+            outs[layer] = stack.forward(layer, inputs[layer], pos, params[layer]).clone()
+    torch.cuda.synchronize()
+    return outs, (dict(recorder.latents) if record else {})
+
+
+def _budget(stack, layer, attn_tail_out, lat_attn_tail, native_out, lat_native, scale):
+    """Each arm against the FP64 oracle of the latent it received."""
+    attn = stack.attn[layer]
+    e_attn_tail = tail_reference.error_stats(
+        attn_tail_out.float() / scale, harness.native_reference(attn, lat_attn_tail)
+    )
+    e_nat = tail_reference.error_stats(
+        native_out.float() / scale, harness.native_reference(attn, lat_native)
+    )
+    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (layer, e_attn_tail, e_nat)
+    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]), (
+        layer,
+        e_attn_tail,
+        e_nat,
+    )
+    return e_attn_tail, e_nat
+
+
+def _one_call_outside_tail(spy, params, label):
+    assert len(spy.calls) == 1, (label, spy.calls)
+    call = spy.calls[0]
+    assert not call.in_tail, f"{label}: all-reduce called inside attention_output_tail"
+    assert call.params is params, f"{label}: params object differs"
+    spy.calls.clear()
+    return call
+
+
+@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_call_site_reduces_like_native(setup, name, enable):
+    stack, spies, probe = setup
+    inputs = {layer: stack.gen_inputs(layer, step=0) for layer in LAYERS}
+    params = {layer: _params(enable) for layer in LAYERS}
+    scale = 2.0 if enable else 1.0
+
+    native, lat_native = _step(stack, False, inputs, params, record=True)
+    for layer in LAYERS:
+        _one_call_outside_tail(spies[layer], params[layer], f"native layer{layer}")
+    assert probe.returned == {}
+
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    attention_output_tail, lat_attn_tail = _step(
+        stack, True, inputs, params, record=True, config=CONFIGS[name]
+    )
+    for layer in LAYERS:
+        tail = stack.attn_tail_tails[layer]
+        assert tail.static_reason is None
+        assert tail.engaged_calls == engaged[layer] + 1, layer
+        call = _one_call_outside_tail(
+            spies[layer], params[layer], f"attention_output_tail layer{layer}"
+        )
+        partial = probe.returned[tail.layer_key]
+        assert call.data_ptr == partial.data_ptr(), (
+            "the all-reduce input is not the tensor attention_output_tail returned"
+        )
+        assert torch.equal(attention_output_tail[layer], partial * 2 if enable else partial)
+        e_attn_tail, e_nat = _budget(
+            stack,
+            layer,
+            attention_output_tail[layer],
+            lat_attn_tail[layer],
+            native[layer],
+            lat_native[layer],
+            scale,
+        )
+        print(
+            f"[attention_output_tail-contract-o] tail={name} allreduce={enable} layer={layer} "
+            f"err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e}"
+        )
+
+
+@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_call_site_graph_replay(setup, name, enable):
+    stack, spies, probe = setup
+    params = {layer: _params(enable) for layer in LAYERS}
+    scale = 2.0 if enable else 1.0
+    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
+    stack.set_attn_tail(True, config=CONFIGS[name])
+    pos = stack.position_ids()
+    with stack.scope():
+        for _ in range(2):
+            for layer in LAYERS:
+                stack.forward(layer, static_in[layer], pos, params[layer])
+        torch.cuda.synchronize()
+        for layer in LAYERS:
+            spies[layer].calls.clear()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph), harness.LatentRecorder(stack) as rec:
+            captured = {
+                layer: stack.forward(layer, static_in[layer], pos, params[layer])
+                for layer in LAYERS
+            }
+        captured_lat = dict(rec.latents)
+    for layer in LAYERS:
+        call = _one_call_outside_tail(spies[layer], params[layer], f"capture layer{layer}")
+        partial = probe.returned[stack.attn_tail_tails[layer].layer_key]
+        assert call.data_ptr == partial.data_ptr()
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    for step in range(1, 4):
+        for layer in LAYERS:
+            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
+        graph.replay()
+        torch.cuda.synchronize()
+        got = {layer: captured[layer].clone() for layer in LAYERS}
+        lat = {layer: captured_lat[layer].clone() for layer in LAYERS}
+        native, lat_native = _step(
+            stack, False, {layer: static_in[layer].clone() for layer in LAYERS}, params, record=True
+        )
+        for layer in LAYERS:
+            _one_call_outside_tail(spies[layer], params[layer], f"native step{step} layer{layer}")
+            e_attn_tail, e_nat = _budget(
+                stack, layer, got[layer], lat[layer], native[layer], lat_native[layer], scale
+            )
+            print(
+                f"[attention_output_tail-contract-o-graph] tail={name} allreduce={enable} "
+                f"step={step} layer={layer} err attention_output_tail={e_attn_tail['rel_rms']:.3e} "
+                f"native={e_nat['rel_rms']:.3e}"
+            )
+    # Replays run no Python: the tails' counters do not move.
+    assert engaged == {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py
new file mode 100644
index 000000000000..03b999ef8aa1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py
@@ -0,0 +1,132 @@
+# 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.
+"""attention_output_tail eligibility gates and registered-API state.
+
+* The static capability check declines attention-DP modules and o_proj
+  tensor-parallel widths other than 8 explicitly, before any other check or
+  GPU work, and the engaged decision logs its per-call domain.
+* The registered functional entry point never falls back to module-global
+  counters for the non-default fused engine: without a caller state it uses
+  zeroed counters private to the call, eagerly and under CUDA-graph replay.
+"""
+
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail import (  # noqa: E501
+    kernels,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
+    TailConfig,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
+    ENGAGED_DOMAIN,
+    SUPPORTED_TP_SIZE,
+    AttnTail,
+    _static_reason,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import tail_reference  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="attention_output_tail targets SM100/SM103",
+)
+
+LOCAL_HEADS, V_DIM, C_DIM, HIDDEN = 8, 256, 512, 6144
+
+
+def _module(tp_size, attention_dp):
+    from tensorrt_llm.mapping import Mapping
+
+    mapping = Mapping(world_size=8, tp_size=8, rank=0, enable_attention_dp=attention_dp)
+    mapping_o = Mapping(world_size=tp_size, tp_size=tp_size, rank=0)
+    # Anything past the gates would fail on this stub, so a returned reason
+    # proves the gate fired first.
+    return SimpleNamespace(mapping=mapping, mapping_o=mapping_o)
+
+
+def test_attention_dp_declines():
+    assert _static_reason(_module(8, True)) == "attention data parallelism"
+
+
+@pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
+def test_other_tp_sizes_decline(tp_size):
+    reason = _static_reason(_module(tp_size, False))
+    assert reason == (f"o_proj tensor-parallel size {tp_size} is not {SUPPORTED_TP_SIZE}")
+
+
+def test_declined_layer_takes_native_path_without_side_effects():
+    tail = AttnTail("adp")
+    meta = SimpleNamespace(num_contexts=0, num_tokens=6)
+    assert not tail.arm(_module(8, True), meta)
+    assert tail.static_reason == "attention data parallelism"
+    assert tail._pending is None and not tail._armed
+    assert tail._state is None
+    assert tail.engaged_calls == 0 and tail.declined_calls == 1
+    # The op-side hook then declines and touches nothing.
+    latent = torch.zeros(6, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
+    out = torch.zeros(6, LOCAL_HEADS * V_DIM, device="cuda").bfloat16()
+    assert not tail.defer_w_uv(None, meta, latent, out)
+    assert tail.project(None, out, None) is None
+
+
+def test_engaged_domain_text():
+    assert "GEN-only" in ENGAGED_DOMAIN and "TP8" in ENGAGED_DOMAIN
+    assert "1-8 token rows" in ENGAGED_DOMAIN
+
+
+def test_registered_fused_entry_has_no_module_state():
+    torch.manual_seed(5)
+    w_uv = (torch.randn(LOCAL_HEADS, V_DIM, C_DIM, device="cuda") * 0.02).bfloat16()
+    w_o = (torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda") * 0.02).bfloat16()
+    latent = torch.randn(6, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
+    fused = TailConfig(oproj_engine="fused")
+    entry = dispatch.get_impl("attention_output_tail", "attn_tail")
+    assert entry is kernels.attn_tail
+    states_before = dict(kernels._DEFAULT_STATES)
+    first = entry(latent, w_uv, w_o, 6, config=fused)
+    second = entry(latent, w_uv, w_o, 6, config=fused)
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_fused import (  # noqa: E501
+        new_state,
+    )
+
+    own = new_state(latent.device)
+    explicit = entry(latent, w_uv, w_o, 6, config=fused, state=own)
+    torch.cuda.synchronize()
+    assert kernels._DEFAULT_STATES == states_before
+    assert torch.equal(first, second) and torch.equal(first, explicit)
+    native = tail_reference.native_attn_tail(latent, w_uv, w_o, 6)
+    torch.testing.assert_close(first, native, rtol=2**-6, atol=2**-6)
+
+    # Graph replay: the per-call counters are re-zeroed by every replay.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = entry(latent, w_uv, w_o, 6, config=fused)
+    for step in range(16):
+        latent.copy_(torch.randn_like(latent, dtype=torch.float32).bfloat16())
+        graph.replay()
+        torch.cuda.synchronize()
+        expected = entry(latent, w_uv, w_o, 6, config=fused, state=own)
+        torch.cuda.synchronize()
+        assert torch.equal(captured, expected), step
+    assert kernels._DEFAULT_STATES == states_before
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py
new file mode 100644
index 000000000000..beb358aa9516
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py
@@ -0,0 +1,359 @@
+# 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.
+"""``AttnTail`` protocol on a real DSA ``MLA`` module with synthetic latents.
+
+Builds the GLM-5.2 attention module of TP8 rank 0 (DSA sparse attention, FP8
+KV cache, BF16 projections) with real ``v_b_proj``/``o_proj`` weights and
+drives the three ``AttnTail`` entry points directly (``arm``, ``defer_w_uv``,
+``project``) with synthetic latents (``tail_fixtures.mix_latent``) and
+``SimpleNamespace`` metadata. No attention kernel runs here; the real
+``MLA.forward`` path with the native attention producer is covered by
+``test_attn_tail_mla_forward.py`` and ``test_attn_tail_mla_forward_mpi.py``.
+
+* the static capability check accepts the module;
+* ``arm`` + ``defer_w_uv`` + ``project`` equal the native GEN tail
+  (``MLA._bmm_bf16_out`` into ``attn_output[0]`` followed by the native
+  ``MLA._project_output``) within the frozen attention_output_tail budget, for every row count;
+* ``defer_w_uv`` acts only on an armed call; ``arm`` declines (and touches no
+  state) under torch.compile tracing, TensorRT-LLM's compile policy and
+  breakable CUDA graphs; unsupported calls (context present, more than
+  ``MAX_ENGAGED_ROWS`` rows) are declined before any GPU work; a mismatched or
+  leftover packet raises; draft (MTP) layers get no tail;
+* the tail can be captured in a CUDA graph and replayed.
+"""
+
+import math
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
+    AttnTail,
+    make_attn_tail,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+_CKPT = tail_fixtures.checkpoint_dir()
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="attention_output_tail targets SM100/SM103",
+    ),
+    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
+]
+
+LAYER = 3
+
+
+def _rope_config():
+    return SimpleNamespace(
+        hidden_size=6144,
+        num_attention_heads=64,
+        rope_scaling={
+            "beta_fast": 32,
+            "beta_slow": 1,
+            "factor": 40,
+            "mscale": 1.0,
+            "mscale_all_dim": 1.0,
+            "original_max_position_embeddings": 4096,
+            "type": "yarn",
+        },
+        max_position_embeddings=4096,
+        rope_theta=10000.0,
+        qk_rope_head_dim=64,
+        model_type="deepseek_v2",
+    )
+
+
+def _build_mla(layer_idx=LAYER, num_hidden_layers=None):
+    from tensorrt_llm._torch.attention.mla import MLA
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
+    from tensorrt_llm.mapping import Mapping
+    from tensorrt_llm.models.modeling_utils import QuantConfig
+    from tensorrt_llm.quantization.mode import QuantAlgo
+
+    quant = QuantConfig()
+    quant.kv_cache_quant_algo = QuantAlgo.FP8
+    pretrained = SimpleNamespace(rms_norm_eps=1e-5)
+    if num_hidden_layers is not None:
+        pretrained.num_hidden_layers = num_hidden_layers
+    model_config = ModelConfig(
+        mapping=Mapping(world_size=8, tp_size=8, rank=0),
+        sparse_attention_config=DeepSeekSparseAttentionConfig(
+            index_n_heads=32, index_head_dim=128, index_topk=2048
+        ),
+        pretrained_config=pretrained,
+        quant_config=quant,
+    )
+    model_config.extra_attrs["kv_cache_dtype"] = "fp8"
+    # One process stands in for TP8 rank 0: the COLUMN Linears would build
+    # all-reduce workspaces over 8 local GPUs, which this test never uses.
+    import tensorrt_llm._torch.distributed as distributed
+
+    saved_all_reduce = distributed.AllReduce
+    distributed.AllReduce = lambda *args, **kwargs: None
+    try:
+        mla = _construct(MLA, model_config, layer_idx)
+    finally:
+        distributed.AllReduce = saved_all_reduce
+    w_uv, w_o = tail_fixtures.load_tail_weights(_CKPT, LAYER, [0])[0]
+    with torch.no_grad():
+        mla.v_b_proj.data = w_uv.cuda().contiguous()
+        mla.o_proj.weight.data.copy_(w_o.cuda())
+    return mla
+
+
+def _build_adopted_mla(layer_idx=LAYER, num_hidden_layers=None):
+    """``_build_mla``, adopted as the target's ``GlmMLA`` (which attaches the attention tail)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+        GlmMLA,
+    )
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+
+    mla = _build_mla(layer_idx=layer_idx, num_hidden_layers=num_hidden_layers)
+    # The model's DeepseekV32Attention adds only its constructor to MLA.
+    mla.__class__ = DeepseekV32Attention
+    return GlmMLA.adopt(mla)
+
+
+def _construct(MLA, model_config, layer_idx):
+    from tensorrt_llm._torch.attention.backends.interface import (
+        PositionalEmbeddingParams,
+        RopeParams,
+    )
+    from tensorrt_llm.functional import PositionEmbeddingType
+
+    return MLA(
+        hidden_size=6144,
+        num_attention_heads=64,
+        num_key_value_heads=64,
+        qk_nope_head_dim=192,
+        qk_rope_head_dim=64,
+        v_head_dim=256,
+        q_lora_rank=2048,
+        kv_lora_rank=512,
+        predicted_tokens_per_seq=6,
+        max_position_embeddings=4096,
+        bias=False,
+        pos_embd_params=PositionalEmbeddingParams(
+            type=PositionEmbeddingType.yarn,
+            rope=RopeParams.from_config(_rope_config()),
+            is_neox=False,
+        ),
+        layer_idx=layer_idx,
+        dtype=torch.bfloat16,
+        config=model_config,
+        reduce_output=False,
+    ).cuda()
+
+
+_MLA = {}
+
+
+def _mla():
+    if "m" not in _MLA:
+        _MLA["m"] = _build_mla()
+    return _MLA["m"]
+
+
+def _native(mla, latent, out_rows):
+    from tensorrt_llm._torch.attention.mla import MLA
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    rows = latent.shape[0]
+    attn_out = torch.zeros(
+        (out_rows, mla.num_heads_tp * mla.v_head_dim), dtype=torch.bfloat16, device="cuda"
+    )
+    v_view = attn_out[:rows].view(rows, mla.num_heads_tp, mla.v_head_dim)
+    MLA._bmm_bf16_out(
+        mla,
+        latent.transpose(0, 1),
+        mla.v_b_proj,
+        mla.v_b_proj.transpose(1, 2),
+        v_view.transpose(0, 1),
+    )
+    meta = SimpleNamespace(num_contexts=0)
+    return mla._project_output([attn_out], None, meta, AllReduceParams(enable_allreduce=False))
+
+
+def _meta(rows, num_contexts=0):
+    return SimpleNamespace(num_contexts=num_contexts, num_tokens=rows)
+
+
+def _latent(rows, seed):
+    if "comp" not in _MLA:
+        _MLA["comp"] = tail_fixtures.compressed_kv_rows(_CKPT, LAYER)
+    return tail_fixtures.mix_latent(_MLA["comp"], rows, 8, seed=seed)
+
+
+def test_static_capability_accepts_dsa_mla():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
+        _static_reason,
+    )
+
+    mla = _mla()
+    assert mla.sparse_attn_hooks is not None
+    assert _static_reason(mla) is None
+
+
+@pytest.mark.parametrize("rows,out_rows", [(1, 1), (2, 2), (4, 4), (6, 6), (6, 8), (6, 20), (8, 8)])
+def test_defer_and_project_match_native(rows, out_rows):
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    mla = _mla()
+    latent = _latent(rows, seed=rows)
+    native = _native(mla, latent, out_rows)
+    tail = make_attn_tail(mla)
+    attn_out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = _meta(rows)
+    assert tail.arm(mla, meta)
+    assert tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
+    attention_output_tail = tail.project(mla, attn_out, AllReduceParams(enable_allreduce=False))
+    assert tail.project(mla, attn_out, None) is None  # packet consumed
+    ref, _ = tail_reference.oracle_attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows)
+    e_attn_tail = tail_reference.error_stats(attention_output_tail[:rows], ref)
+    e_nat = tail_reference.error_stats(native[:rows], ref)
+    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6
+    torch.testing.assert_close(
+        attention_output_tail[:rows], native[:rows], rtol=2**-6, atol=2**-6 * e_nat["ref_rms"]
+    )
+    assert bool((attention_output_tail[rows:] == 0).all())
+    assert tail.engaged_calls == 1 and tail.declined_calls == 0
+
+
+def test_declines_before_any_work():
+    mla = _mla()
+    tail = make_attn_tail(mla)
+    lat = _latent(6, seed=0)
+    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
+    # Not armed: the op-side entry never engages on its own.
+    assert not tail.defer_w_uv(mla, _meta(6), lat, out)
+    assert not tail.arm(mla, _meta(6, num_contexts=1))
+    assert not tail.defer_w_uv(mla, _meta(6, num_contexts=1), lat, out)
+    assert not tail.arm(mla, _meta(9))
+    # Armed, but the latent the op hands over does not fit the call.
+    assert tail.arm(mla, _meta(6))
+    big = torch.zeros((9, 8, 512), dtype=torch.bfloat16, device="cuda")
+    big_out = torch.empty((9, 2048), dtype=torch.bfloat16, device="cuda")
+    assert not tail.defer_w_uv(mla, _meta(6), big, big_out)
+    assert tail.declined_calls == 3 and tail.engaged_calls == 0
+    assert not tail._armed
+    assert tail.project(mla, out, None) is None
+
+
+def test_arm_declines_under_compile_and_breakable_graphs():
+    from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import enable_breakable_cuda_graph
+    from tensorrt_llm._torch.utils import torch_compiling
+
+    mla = _mla()
+    tail = make_attn_tail(mla)
+    with torch_compiling(True):
+        assert not tail.arm(mla, _meta(6))
+    with enable_breakable_cuda_graph():
+        assert not tail.arm(mla, _meta(6))
+    assert tail.declined_calls == 1  # the compile check touches no state
+
+    meta = _meta(6)
+
+    def traced(x):
+        return x + 1 if tail.arm(mla, meta) else x - 1
+
+    x = torch.zeros(4, device="cuda")
+    out = torch.compile(traced, backend="eager", fullgraph=True)(x)
+    assert torch.equal(out, x - 1)
+    assert not tail._armed and tail.declined_calls == 1
+
+
+def test_packet_errors_are_loud():
+    mla = _mla()
+    tail = make_attn_tail(mla)
+    meta = _meta(6)
+    lat = _latent(6, seed=1)
+    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
+    other = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
+    assert tail.arm(mla, meta)
+    assert tail.defer_w_uv(mla, meta, lat, out)
+    with pytest.raises(RuntimeError, match="not consumed"):
+        tail.arm(mla, meta)
+    with pytest.raises(RuntimeError, match="does not match"):
+        tail.project(mla, other, None)
+
+
+def test_draft_layers_keep_native_tail():
+    mla = _build_mla(layer_idx=78, num_hidden_layers=78)
+    assert make_attn_tail(mla) is None
+
+
+def test_graph_capture_and_replay():
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    mla = _mla()
+    tail = make_attn_tail(mla)
+    rows = 6
+    meta = _meta(rows)
+    lat = _latent(rows, seed=2).contiguous()
+    out = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    params = AllReduceParams(enable_allreduce=False)
+    # Eager warm-up (JIT, per-layer state) before capture, like the engine.
+    assert tail.arm(mla, meta)
+    assert tail.defer_w_uv(mla, meta, lat, out)
+    tail.project(mla, out, params)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        assert tail.arm(mla, meta)
+        assert tail.defer_w_uv(mla, meta, lat, out)
+        captured = tail.project(mla, out, params)
+    for seed in range(8):
+        lat.copy_(_latent(rows, seed=100 + seed))
+        graph.replay()
+        native = _native(mla, lat, rows)
+        torch.testing.assert_close(
+            captured.clone(),
+            native,
+            rtol=2**-6,
+            atol=2**-6 * native.float().pow(2).mean().sqrt().item(),
+        )
+    assert not math.isnan(captured.float().sum().item())
+
+
+@pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
+    reason="opt-in",
+)
+def test_adoption_attaches_tail_to_target_layers(monkeypatch):
+    """With attention_output_tail running, the adopted attention (``GlmMLA``) owns an attention_output_tail.
+
+    The call sites themselves run in ``test_attn_tail_mla_forward.py``.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+        dispatch,
+    )
+
+    plan_testlib.select(dispatch, "attention_output_tail", monkeypatch)
+    mla = _build_adopted_mla(layer_idx=LAYER, num_hidden_layers=78)
+    assert isinstance(mla._attn_tail, AttnTail)
+    draft = _build_adopted_mla(layer_idx=78, num_hidden_layers=78)
+    assert draft._attn_tail is None
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py
new file mode 100644
index 000000000000..45a3230bbfb4
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py
@@ -0,0 +1,269 @@
+# 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.
+"""attention_output_tail inside the target's ``GlmMLA.forward``: native DSA attention → attention_output_tail → output.
+
+Runs with the plan that runs only ``attention_output_tail`` (opt-in:
+``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED=1``).
+
+Builds TP8 rank 0 of real GLM-5.2 layers 2 (dense destination, full indexer)
+and 3 (MoE destination, shared indexer) as ``DeepseekV32Attention`` with all
+attention weights real, adopted as the target's ``GlmMLA``, fills an FP8 KV cache with a real
+4096-token prefill, and runs GEN-only ``GlmMLA.forward`` as a decode step through the
+``register_to_config`` path (the target's DSA attention op), as the model does. Every check compares
+attention_output_tail with the native tail on the same module, cache and inputs:
+
+* eager, R=6 (one request, MTP draft 5) and R=1: the latent the tail receives
+  comes from the native attention kernel and is identical in both arms; the
+  attention_output_tail output is within the frozen budget against the FP64 oracle of that
+  latent and matches native; attention_output_tail engages exactly once per layer and call;
+* CUDA-graph capture of the GEN step, replayed with new inputs;
+* under torch.compile (Dynamo traces ``MLA.forward``, the attention op stays
+  opaque) and under TensorRT-LLM's torch.compile policy, attention_output_tail declines in
+  ``MLA.forward`` before the op runs: the output equals native and no packet
+  or armed decision is left behind; a breakable CUDA graph also declines.
+
+Fixture labels: weights real, activations synthetic (real input_layernorm of
+real token embeddings).
+"""
+
+import os
+import sys
+from contextlib import nullcontext
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import real_mla_harness as harness  # noqa: E402
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+_CKPT = tail_fixtures.checkpoint_dir()
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="attention_output_tail targets SM100/SM103",
+    ),
+    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
+    pytest.mark.skipif(
+        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
+        reason="needs attention_output_tail selected (opt-in)",
+    ),
+    pytest.mark.threadleak(enabled=False),
+]
+
+LAYERS = [2, 3]
+CTX = 4096
+TAIL_PATHS = "attention_output_tail"
+_STACKS = {}
+
+
+def _stack(rows, fast_paths=TAIL_PATHS):
+    """The stack of ``rows`` GEN rows built (once) under the plan that runs ``fast_paths``."""
+    key = (rows, fast_paths)
+    if key not in _STACKS:
+        _STACKS[key] = harness.AttentionStack(
+            _CKPT, LAYERS, tp_rank=0, ctx=CTX, rows=rows, fast_paths=fast_paths
+        )
+    else:
+        harness.select(fast_paths)
+    return _STACKS[key]
+
+
+def _no_ar():
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    return AllReduceParams(enable_allreduce=False)
+
+
+def _run(stack, attention_output_tail, inputs, record=False, position_ids=None):
+    """One GEN step over the stack; returns per-layer outputs and latents."""
+    stack.set_attn_tail(attention_output_tail)
+    pos = stack.position_ids() if position_ids is None else position_ids
+    outs = {}
+    recorder = harness.LatentRecorder(stack) if record else nullcontext()
+    with stack.scope(), recorder:
+        for layer in stack.layers:
+            outs[layer] = stack.forward(layer, inputs[layer], pos, _no_ar()).clone()
+    torch.cuda.synchronize()
+    return outs, (dict(recorder.latents) if record else {})
+
+
+def _check_budget(stack, layer, attn_tail_out, native_out, latent):
+    ref = harness.native_reference(stack.attn[layer], latent)
+    e_attn_tail = tail_reference.error_stats(attn_tail_out, ref)
+    e_nat = tail_reference.error_stats(native_out, ref)
+    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (layer, e_attn_tail, e_nat)
+    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]), (
+        layer,
+        e_attn_tail,
+        e_nat,
+    )
+    torch.testing.assert_close(attn_tail_out, native_out, rtol=2**-6, atol=2**-6 * e_nat["ref_rms"])
+    return e_attn_tail, e_nat
+
+
+@pytest.mark.parametrize("rows", [6, 1])
+def test_eager_forward_matches_native(rows):
+    stack = _stack(rows)
+    inputs = {layer: stack.gen_inputs(layer, step=0) for layer in LAYERS}
+    native, lat_native = _run(stack, False, inputs, record=True)
+    # A second native run characterizes the run-to-run variation of the
+    # native attention latent itself.
+    _, lat_native2 = _run(stack, False, inputs, record=True)
+    before = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    attention_output_tail, lat_attn_tail = _run(stack, True, inputs, record=True)
+    for layer in LAYERS:
+        tail = stack.attn_tail_tails[layer]
+        assert tail.static_reason is None
+        assert tail.engaged_calls == before[layer] + 1, layer
+        assert tail._pending is None and not tail._armed
+        assert lat_native[layer].shape == (rows, 8, 512)
+        same = torch.equal(lat_native[layer], lat_attn_tail[layer])
+        native_rerun_same = torch.equal(lat_native[layer], lat_native2[layer])
+        native_rerun_diff = (
+            (lat_native[layer].float() - lat_native2[layer].float()).abs().max().item()
+        )
+        print(
+            f"[attention_output_tail-mla-latent] rows={rows} layer={layer} "
+            f"native_rerun_equal={native_rerun_same} "
+            f"native_rerun_maxdiff={native_rerun_diff:.3e} "
+            f"attn_tail_vs_native_equal={same}"
+        )
+        torch.testing.assert_close(lat_attn_tail[layer], lat_native[layer])
+        e_attn_tail, e_nat = _check_budget(
+            stack, layer, attention_output_tail[layer], native[layer], lat_attn_tail[layer]
+        )
+        print(
+            f"[attention_output_tail-mla-forward] rows={rows} layer={layer} latent_equal={same} "
+            f"err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e} "
+            f"maxdiff={(attention_output_tail[layer].float() - native[layer].float()).abs().max().item():.3e}"
+        )
+
+
+def test_cuda_graph_capture_and_replay():
+    stack = _stack(6)
+    pos = stack.position_ids()
+    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
+    stack.set_attn_tail(True)
+    with stack.scope():
+        for _ in range(2):
+            for layer in LAYERS:
+                stack.forward(layer, static_in[layer], pos, _no_ar())
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = {
+                layer: stack.forward(layer, static_in[layer], pos, _no_ar()) for layer in LAYERS
+            }
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    for step in range(1, 5):
+        for layer in LAYERS:
+            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
+        graph.replay()
+        torch.cuda.synchronize()
+        got = {layer: captured[layer].clone() for layer in LAYERS}
+        native, lat = _run(
+            stack,
+            False,
+            {layer: static_in[layer].clone() for layer in LAYERS},
+            record=True,
+            position_ids=pos,
+        )
+        for layer in LAYERS:
+            e_attn_tail, e_nat = _check_budget(stack, layer, got[layer], native[layer], lat[layer])
+            print(
+                f"[attention_output_tail-mla-graph] step={step} layer={layer} "
+                f"err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e}"
+            )
+    # Replays run no Python: the engagement counters do not move.
+    assert engaged == {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    del graph
+
+
+def _assert_declined_and_native(stack, layer, out, native, engaged_before):
+    tail = stack.attn_tail_tails[layer]
+    assert tail.engaged_calls == engaged_before
+    assert tail._pending is None and not tail._armed
+    print(
+        f"[attention_output_tail-mla-declined] layer={layer} "
+        f"bitwise_equal_to_native={torch.equal(out, native)}"
+    )
+    torch.testing.assert_close(
+        out, native, rtol=2**-6, atol=2**-6 * native.float().pow(2).mean().sqrt().item()
+    )
+
+
+def test_torch_compile_takes_native_path():
+    stack = _stack(6)
+    inputs = {layer: stack.gen_inputs(layer, step=7) for layer in LAYERS}
+    native, _ = _run(stack, False, inputs)
+    pos = stack.position_ids()
+    stack.set_attn_tail(True)
+    compiled = {
+        layer: torch.compile(stack.attn[layer], backend="eager", dynamic=False) for layer in LAYERS
+    }
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    with stack.scope():
+        for _ in range(2):
+            outs = {
+                layer: compiled[layer](
+                    position_ids=pos,
+                    hidden_states=inputs[layer],
+                    attn_metadata=stack.gen_metadata,
+                    all_reduce_params=_no_ar(),
+                ).clone()
+                for layer in LAYERS
+            }
+            torch.cuda.synchronize()
+            for layer in LAYERS:
+                _assert_declined_and_native(
+                    stack, layer, outs[layer], native[layer], engaged[layer]
+                )
+    # The same module still takes attention_output_tail eagerly afterwards.
+    attention_output_tail, _ = _run(stack, True, inputs)
+    for layer in LAYERS:
+        assert stack.attn_tail_tails[layer].engaged_calls == engaged[layer] + 1
+        torch.testing.assert_close(
+            attention_output_tail[layer],
+            native[layer],
+            rtol=2**-6,
+            atol=2**-6 * native[layer].float().pow(2).mean().sqrt().item(),
+        )
+
+
+def test_trtllm_compile_policy_takes_native_path():
+    from tensorrt_llm._torch.utils import torch_compiling
+
+    stack = _stack(6)
+    inputs = {layer: stack.gen_inputs(layer, step=8) for layer in LAYERS}
+    native, _ = _run(stack, False, inputs)
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    with torch_compiling(True):
+        outs, _ = _run(stack, True, inputs)
+    for layer in LAYERS:
+        _assert_declined_and_native(stack, layer, outs[layer], native[layer], engaged[layer])
+
+
+def test_breakable_cuda_graph_declines():
+    from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import enable_breakable_cuda_graph
+
+    stack = _stack(6)
+    layer = LAYERS[0]
+    tail = stack.attn_tail_tails[layer]
+    with enable_breakable_cuda_graph():
+        assert not tail.arm(stack.attn[layer], stack.gen_metadata)
+    assert not tail._armed
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py
new file mode 100644
index 000000000000..853abf9552b7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py
@@ -0,0 +1,423 @@
+# 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.
+"""Native attention → attention_output_tail → native reduction/norm, multi-rank, the target's GlmMLA.forward.
+
+Runs with the plan that runs only ``attention_output_tail`` (opt-in:
+``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED=1``),
+one process per GPU under
+``srun --mpi=pmix``. Rank ``r`` builds TP8 rank ``r`` of real layers 2 and 3
+(``real_mla_harness.AttentionStack``: all attention weights real, FP8 KV
+cache filled by a real prefill) and runs the GEN step (R=6, one request, MTP
+draft 5) as a decode step through ``GlmMLA.forward``, so the latent comes from the native DSA
+attention kernel inside the target's DSA attention op. Each rank's output
+feeds the native reduction exactly as the decoder layer does:
+
+* layer 2 (dense destination): ``o_proj`` reduces inside ``Linear.forward``
+  (``AllReduceParams(enable_allreduce=True)``), then the real
+  ``post_attention_layernorm(hidden, residual)``;
+* layer 3 (MoE destination): ``o_proj`` returns the rank-local partial
+  (``enable_allreduce=False``), then ``AllReduce(RESIDUAL_RMS_NORM)`` with the
+  real ``post_attention_layernorm`` weight.
+
+With eight ranks on two nodes (``-N2``) this is the GEN role's TP8 group;
+with fewer ranks the reduction sums the shards of ranks ``0 .. world-1``
+only, which still exercises the bridge.
+
+The native DSA attention output varies slightly between identical runs, so
+each arm is judged against its own oracle: every rank records the latent its
+tail received (``LatentRecorder``), the FP64 tail oracle of that latent (with
+the native BF16 rounding of ``v``) is summed over ranks in FP64, and the
+residual add and RMSNorm are applied in FP64. attention_output_tail's normalized RMS and max
+errors against its oracle must stay within the frozen attention_output_tail budget relative to
+native's errors against native's oracle. A second native run (A/A) is
+reported for scale. Results are gathered from all ranks before any assertion,
+so a failure on one rank cannot leave the others waiting in a collective.
+
+The attention_output_tail contract ("no communication is hidden inside attention_output_tail"): each
+``o_proj`` gets its own view of the shared all-reduce that counts calls and flags any call
+made while an ``AttnTail`` entry point runs. In both arms every layer call
+must reach the all-reduce exactly once, outside the tail, with the params
+object the caller passed: from ``Linear.forward`` natively, from
+``MLA.forward`` with attention_output_tail. A second test captures the attention_output_tail step
+(both destinations, real all-reduces) in a CUDA graph and replays it with new
+inputs against the same per-arm FP64 oracles.
+"""
+
+import hashlib
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import real_mla_harness as harness  # noqa: E402
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+_CKPT = tail_fixtures.checkpoint_dir()
+HIDDEN = 6144
+LAYERS = [2, 3]
+
+
+def _comm():
+    from mpi4py import MPI
+
+    return MPI.COMM_WORLD
+
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"),
+    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
+    pytest.mark.skipif(
+        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
+        reason="needs attention_output_tail selected (opt-in)",
+    ),
+    pytest.mark.skipif(_comm().Get_size() < 2, reason="needs >= 2 MPI ranks"),
+    pytest.mark.threadleak(enabled=False),
+]
+
+_SETUP = {}
+TAIL_ENTRY_POINTS = ("arm", "before_attention", "defer_w_uv", "project")
+
+
+class _InTail:
+    """Set while an ``AttnTail`` entry point runs (class-level hooks, so
+    they compose with the harness's per-instance latent recorder)."""
+
+    active = False
+
+
+def _hook_tail_entry_points():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
+        AttnTail,
+    )
+
+    if getattr(AttnTail, "_contract_o_hooked", False):
+        return
+    for name in TAIL_ENTRY_POINTS:
+        entry = getattr(AttnTail, name)
+
+        def run(self, *args, _entry=entry, **kwargs):
+            _InTail.active = True
+            try:
+                return _entry(self, *args, **kwargs)
+            finally:
+                _InTail.active = False
+
+        setattr(AttnTail, name, run)
+    AttnTail._contract_o_hooked = True
+
+
+class _CountingAllReduce:
+    """One ``o_proj``'s view of the shared ``AllReduce``: forwards every
+    call and records whether it was made inside the attention_output_tail and with which
+    params object."""
+
+    def __init__(self, inner):
+        self.inner = inner
+        self.calls = []
+
+    def __call__(self, input, all_reduce_params=None):
+        self.calls.append((_InTail.active, all_reduce_params))
+        return self.inner(input, all_reduce_params=all_reduce_params)
+
+    def __getattr__(self, name):
+        return getattr(self.inner, name)
+
+
+def _setup():
+    if "s" in _SETUP:
+        return _SETUP["s"]
+    from tensorrt_llm._torch.distributed import AllReduce
+    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+    from tensorrt_llm.mapping import Mapping
+
+    comm = _comm()
+    rank, world = comm.Get_rank(), comm.Get_size()
+    torch.cuda.set_device(rank % torch.cuda.device_count())
+    mapping = Mapping(
+        world_size=world, tp_size=world, rank=rank, gpus_per_node=torch.cuda.device_count()
+    )
+    stack = harness.AttentionStack(
+        _CKPT, LAYERS, tp_rank=rank, ctx=4096, rows=6, fast_paths="attention_output_tail"
+    )
+    # As in the model: o_proj owns an all-reduce; the decoder layer enables
+    # it for dense layers and fuses it into the post-attention norm for MoE.
+    allreduce = AllReduce(mapping=mapping, dtype=torch.bfloat16)
+    for attn in stack.attn.values():
+        attn.o_proj.reduce_output = True
+        attn.o_proj.all_reduce = _CountingAllReduce(allreduce)
+    _hook_tail_entry_points()
+    eps = stack.pretrained.rms_norm_eps
+    norms = {}
+    for layer in LAYERS:
+        norm = RMSNorm(hidden_size=HIDDEN, eps=eps, dtype=torch.bfloat16).cuda()
+        norm.weight.data.copy_(
+            tail_fixtures.read_tensor(
+                _CKPT, f"model.layers.{layer}.post_attention_layernorm.weight"
+            ).cuda()
+        )
+        norms[layer] = norm
+    _SETUP["s"] = (rank, world, stack, norms, allreduce, eps)
+    return _SETUP["s"]
+
+
+def _digest(t):
+    return hashlib.sha256(t.contiguous().view(torch.uint8).cpu().numpy().tobytes()).hexdigest()
+
+
+def _chain(
+    stack, norms, allreduce, eps, attention_output_tail, inputs, residual, params=None, sync=True
+):
+    """GEN step of layers 2 (dense) and 3 (MoE) with their native reductions.
+
+    Returns the outputs and the latent each layer's tail received; ``params``
+    (if given) receives the ``AllReduceParams`` object passed to each layer.
+    """
+    from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams
+
+    stack.set_attn_tail(attention_output_tail)
+    pos = stack.position_ids()
+    params = {} if params is None else params
+    params[2] = AllReduceParams(enable_allreduce=True)
+    params[3] = AllReduceParams(enable_allreduce=False)
+    # Outputs are cloned right after they are produced: an all-reduce result
+    # may live in a workspace that the next all-reduce reuses.
+    out = {}
+    with stack.scope(), harness.LatentRecorder(stack) as recorder:
+        reduced = stack.forward(2, inputs[2], pos, params[2])
+        out["dense_reduced"] = reduced.clone()
+        hs2, res2 = norms[2](reduced, residual.clone())
+        out["dense_hidden"], out["dense_residual"] = hs2.clone(), res2.clone()
+        partial = stack.forward(3, inputs[3], pos, params[3])
+        out["moe_partial"] = partial.clone()
+        hs3, res3 = allreduce(
+            partial,
+            all_reduce_params=AllReduceParams(
+                fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                residual=residual.clone(),
+                norm_weight=norms[3].weight,
+                eps=eps,
+                trigger_completion_at_end=False,
+            ),
+        )
+        out["moe_hidden"], out["moe_residual"] = hs3.clone(), res3.clone()
+        latents = dict(recorder.latents)
+    if sync:
+        torch.cuda.synchronize()
+    return out, latents
+
+
+def _take_calls(stack, params, label):
+    """The attention_output_tail contract: one all-reduce call per layer, outside the tail, with
+    the caller's params object. Returns failure strings and clears."""
+    failures = []
+    for layer in LAYERS:
+        calls = stack.attn[layer].o_proj.all_reduce.calls
+        ok = len(calls) == 1 and not calls[0][0] and calls[0][1] is params[layer]
+        if not ok:
+            failures.append(
+                f"{label} layer{layer}: o_proj all-reduce calls "
+                f"(in_tail, same_params) = "
+                f"{[(c[0], c[1] is params[layer]) for c in calls]}"
+            )
+        calls.clear()
+    return failures
+
+
+def _oracle(stack, norms, eps, latents, residual):
+    """FP64 references of ``_chain``'s outputs from the recorded latents."""
+    import numpy as np
+    from mpi4py import MPI
+
+    def summed(local):
+        send = local.cpu().numpy()
+        recv = np.empty_like(send)
+        _comm().Allreduce(send, recv, op=MPI.SUM)
+        return torch.from_numpy(recv).cuda()
+
+    def norm(x, weight):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight.double()
+
+    res_in = residual.double()
+    partial = {}
+    for layer in LAYERS:
+        attn = stack.attn[layer]
+        partial[layer], _ = tail_reference.oracle_attn_tail(
+            latents[layer], attn.v_b_proj, attn.o_proj.weight, stack.rows
+        )
+    dense = summed(partial[2])
+    moe = summed(partial[3])
+    return {
+        "dense_reduced": dense,
+        "dense_hidden": norm(dense + res_in, norms[2].weight),
+        "dense_residual": dense + res_in,
+        "moe_partial": partial[3],
+        "moe_hidden": norm(moe + res_in, norms[3].weight),
+        "moe_residual": moe + res_in,
+    }
+
+
+@pytest.mark.parametrize("step", [0, 1])
+def test_native_producer_attn_tail_native_reduction(step):
+    rank, world, stack, norms, allreduce, eps = _setup()
+    inputs = {layer: stack.gen_inputs(layer, step=step) for layer in LAYERS}
+    gen = torch.Generator(device="cuda").manual_seed(4321 + step)
+    residual = torch.randn(
+        (stack.rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
+    ).bfloat16()
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    for layer in LAYERS:
+        stack.attn[layer].o_proj.all_reduce.calls.clear()
+    failures, report = [], []
+    params = {}
+    native, lat_native = _chain(stack, norms, allreduce, eps, False, inputs, residual, params)
+    failures += _take_calls(stack, params, "native")
+    native_b, _ = _chain(stack, norms, allreduce, eps, False, inputs, residual, params)
+    failures += _take_calls(stack, params, "native_b")
+    attention_output_tail, lat_attn_tail = _chain(
+        stack, norms, allreduce, eps, True, inputs, residual, params
+    )
+    failures += _take_calls(stack, params, "attention_output_tail")
+    report.append("contract_o_calls=" + ("ok" if not failures else "BAD"))
+    ref_native = _oracle(stack, norms, eps, lat_native, residual)
+    ref_attn_tail = _oracle(stack, norms, eps, lat_attn_tail, residual)
+
+    for layer in LAYERS:
+        tail = stack.attn_tail_tails[layer]
+        # As in the model (default TLLM_NCCL_SYMMETRIC_ZERO_COPY, AUTO
+        # strategy, TP > 1), the o_proj all-reduce may use an NCCL window,
+        # but the native o_proj GEMM output is an ordinary F.linear buffer,
+        # so attention_output_tail still takes the call.
+        report.append(
+            f"layer{layer}: nccl_window_allreduce="
+            f"{stack.attn[layer].o_proj.all_reduce.uses_nccl_symmetric_memory_window()} "
+            f"static_reason={tail.static_reason} "
+            f"engaged={tail.engaged_calls - engaged[layer]}"
+        )
+        if tail.static_reason is not None or (tail.engaged_calls != engaged[layer] + 1):
+            failures.append(report[-1])
+    for key in native:
+        e_attn_tail = tail_reference.error_stats(attention_output_tail[key], ref_attn_tail[key])
+        e_nat = tail_reference.error_stats(native[key], ref_native[key])
+        aa = tail_reference.error_stats(native_b[key], native[key].double())
+        report.append(
+            f"{key}: rel_rms attention_output_tail={e_attn_tail['rel_rms']:.2e} "
+            f"native={e_nat['rel_rms']:.2e} "
+            f"max attention_output_tail={e_attn_tail['max_abs']:.2e} "
+            f"native={e_nat['max_abs']:.2e} "
+            f"(A/A rel_rms {aa['rel_rms']:.2e})"
+        )
+        if (
+            e_attn_tail["rel_rms"] > 1.25 * e_nat["rel_rms"] + 1e-6
+            or e_attn_tail["max_abs"] > 2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]
+        ):
+            failures.append(report[-1])
+    digests = tuple(
+        _digest(attention_output_tail[k])
+        for k in ("dense_reduced", "dense_hidden", "moe_hidden", "moe_residual")
+    )
+    gathered = _comm().allgather((rank, failures, report, digests))
+    if rank == 0:
+        for r, _, rep, _ in gathered:
+            print(
+                f"[attention_output_tail-mla-forward-mpi] rank={r}/{world} step={step} | "
+                + " | ".join(rep)
+            )
+    bad = [(r, f) for r, f, _, _ in gathered if f]
+    assert not bad, bad
+    assert len({d for _, _, _, d in gathered}) == 1, "reduced outputs differ across ranks"
+
+
+def test_graph_replay_native_reduction():
+    """The attention_output_tail step of both destinations, real all-reduces included, captured
+    in one CUDA graph and replayed with new inputs and residuals; each replay
+    is judged like the eager step, against native eager on the same inputs."""
+    rank, world, stack, norms, allreduce, eps = _setup()
+    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
+    static_res = torch.zeros((stack.rows, HIDDEN), dtype=torch.bfloat16, device="cuda")
+    for layer in LAYERS:
+        stack.attn[layer].o_proj.all_reduce.calls.clear()
+    failures, report = [], []
+    params = {}
+    for _ in range(2):  # eager warm-up (JIT, per-layer state), as the engine
+        _chain(stack, norms, allreduce, eps, True, static_in, static_res, params)
+        failures += _take_calls(stack, params, "warm-up")
+    _comm().Barrier()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured, captured_lat = _chain(
+            stack, norms, allreduce, eps, True, static_in, static_res, params, sync=False
+        )
+    failures += _take_calls(stack, params, "capture")
+    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
+    digests = []
+    for step in range(1, 4):
+        for layer in LAYERS:
+            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
+        gen = torch.Generator(device="cuda").manual_seed(4321 + step)
+        static_res.copy_(
+            torch.randn(
+                (stack.rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
+            ).bfloat16()
+        )
+        _comm().Barrier()
+        graph.replay()
+        torch.cuda.synchronize()
+        got = {k: v.clone() for k, v in captured.items()}
+        lat = {k: v.clone() for k, v in captured_lat.items()}
+        native, lat_native = _chain(
+            stack,
+            norms,
+            allreduce,
+            eps,
+            False,
+            {layer: static_in[layer].clone() for layer in LAYERS},
+            static_res.clone(),
+            params,
+        )
+        failures += _take_calls(stack, params, f"native step{step}")
+        ref_attn_tail = _oracle(stack, norms, eps, lat, static_res)
+        ref_native = _oracle(stack, norms, eps, lat_native, static_res)
+        for key in got:
+            e_attn_tail = tail_reference.error_stats(got[key], ref_attn_tail[key])
+            e_nat = tail_reference.error_stats(native[key], ref_native[key])
+            report.append(
+                f"step{step} {key}: rel_rms "
+                f"graph={e_attn_tail['rel_rms']:.2e} "
+                f"native={e_nat['rel_rms']:.2e}"
+            )
+            if (
+                e_attn_tail["rel_rms"] > 1.25 * e_nat["rel_rms"] + 1e-6
+                or e_attn_tail["max_abs"] > 2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]
+            ):
+                failures.append(report[-1])
+        digests.append(
+            tuple(_digest(got[k]) for k in ("dense_reduced", "dense_hidden", "moe_hidden"))
+        )
+    if engaged != {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}:
+        failures.append("replays ran tail Python")
+    gathered = _comm().allgather((rank, failures, report, tuple(digests)))
+    if rank == 0:
+        for r, _, rep, _ in gathered:
+            print(
+                f"[attention_output_tail-mla-forward-mpi-graph] rank={r}/{world} | "
+                + " | ".join(rep)
+            )
+    bad = [(r, f) for r, f, _, _ in gathered if f]
+    assert not bad, bad
+    assert len({d for _, _, _, d in gathered}) == 1, "replayed reduced outputs differ across ranks"
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py
new file mode 100644
index 000000000000..a2313306d43e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py
@@ -0,0 +1,506 @@
+# 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.
+"""PDL boundary stress of the default attention_output_tail engine (``cute_stream``).
+
+The default attention_output_tail is two kernels with programmatic dependent launch (PDL) at
+both package boundaries:
+
+* ``_wuv_kernel`` is a PDL dependent of the latent's producer. It triggers its
+  own dependent at entry and loads ``w_uv`` before ``griddepcontrol.wait``;
+  it reads the latent and writes ``v`` only after the wait.
+* ``stream_oproj`` is a PDL dependent of ``_wuv_kernel``. It copies ``w_o``
+  rows into its own shared memory before its wait, reads ``v`` and writes the
+  partial after it, and triggers its own PDL consumer before it has written
+  the partial.
+
+The boundary contract is therefore: before its wait an attention_output_tail kernel touches only
+immutable weights and its own shared memory, and every consumer of the
+partial waits before reading it. These tests exercise that contract with
+synthetic neighbour kernels that stretch the overlap windows:
+
+* **Delayed producer.** A predecessor triggers its dependents at entry, spins
+  for ``delay`` µs, then writes the latent. The latent and ``v`` are poisoned
+  with NaN before the producer runs. An attention_output_tail read of the latent or of ``v`` that
+  is not ordered after the producer (or after ``_wuv_kernel``) shows up as
+  NaN or as another iteration's values.
+* **Delayed consumer with rapid buffer reuse.** Every kernel of the chain
+  producer → attention_output_tail → consumer is PDL-launched and triggers its dependents at
+  entry, so the next iteration's producer and attention_output_tail kernels may become resident
+  while the current consumer still runs. The consumer waits, spins, then
+  copies the partial. attention_output_tail's ``v`` and partial come from the CUDA-graph pool and
+  are reused across iterations. An attention_output_tail global write issued before its wait would
+  overwrite the partial the sleeping consumer has not read yet.
+
+Every captured iteration has its own latent, and the partial the consumer
+copies must equal, bit for bit, an eager attention_output_tail run on that latent
+(attention_output_tail is deterministic). The graphs are replayed many times with new latents. The
+consumer records ``%globaltimer`` at entry and after its wait, so the tests
+can report (and ``test_pdl_overlap_witness`` asserts) that the consumer really
+was resident before attention_output_tail finished.
+
+Weights: real GLM-5.2 ``v_b_proj`` / ``o_proj`` TP8 shards when
+``GLM52_NVFP4_CKPT`` is set, otherwise synthetic. Latents: real-weight,
+synthetic-activation mixtures of real compressed-KV rows when the checkpoint
+is available, otherwise random.
+
+``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_QUICK=1`` shrinks iteration and replay counts
+(compute-sanitizer runs); ``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SANITIZER=1`` skips the timing
+witness, because the sanitizer serializes kernels;
+``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SYNTHETIC=1`` uses synthetic weights and latents even when
+the checkpoint is available.
+"""
+
+import os
+import statistics
+import sys
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
+    DEFAULT_CONFIG,
+    TailConfig,
+    attn_tail,
+    check_tail_shapes,
+    pdl_enabled,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import tail_fixtures  # noqa: E402
+import tail_reference  # noqa: E402
+
+TP_SIZE = 8
+LOCAL_HEADS, V_DIM, C_DIM, HIDDEN = 8, 256, 512, 6144
+QUICK = os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_QUICK") == "1"
+UNDER_SANITIZER = os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SANITIZER") == "1"
+N_ITERS = 3 if QUICK else 16
+N_REPLAYS = 2 if QUICK else 24
+_CKPT = tail_fixtures.checkpoint_dir()
+# GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SYNTHETIC=1 forces synthetic weights and latents built by
+# elementwise kernels only (compute-sanitizer initcheck: the real-weight
+# fixture runs a cuBLAS GEMM whose output initcheck cannot see as written).
+_HAVE_CKPT = (
+    bool(_CKPT and os.path.isdir(_CKPT))
+    and os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SYNTHETIC") != "1"
+)
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="attention_output_tail kernels target SM100/SM103",
+    ),
+    pytest.mark.skipif(not pdl_enabled(), reason="the stress needs TRTLLM_ENABLE_PDL=1"),
+]
+
+# A fully serialized attention_output_tail: no PDL at entry, inside, or at exit.
+SERIAL_CONFIG = TailConfig(pdl_entry=False, trigger_dependents=False)
+
+# ---------------------------------------------------------------------------
+# Synthetic neighbour kernels.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _clock_fenced():
+    # The memory clobber keeps loads and stores from moving across the read.
+    return tl.inline_asm_elementwise(
+        "mov.u64 $0, %globaltimer;", "=l,~{memory}", [], dtype=tl.int64, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _gdc_wait():
+    tl.inline_asm_elementwise(
+        "griddepcontrol.wait; // dummy $0",
+        "=r,~{memory}",
+        [],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _gdc_launch_dependents():
+    tl.inline_asm_elementwise(
+        "griddepcontrol.launch_dependents; // dummy $0",
+        "=r,~{memory}",
+        [],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _spin(delay_ns):
+    t0 = _clock_fenced()
+    t = t0
+    while t - t0 < delay_ns:
+        t = _clock_fenced()
+    return t
+
+
+@triton.jit
+def _delayed_producer(
+    src_ptr,
+    dst_ptr,
+    stamp_ptr,
+    numel,
+    delay_ns,
+    BLOCK: tl.constexpr,
+    TRIGGER: tl.constexpr,
+    WAIT: tl.constexpr,
+):
+    """Copy ``src`` to ``dst`` after spinning ``delay_ns``.
+
+    ``TRIGGER`` lets the PDL dependent (attention_output_tail's W_UV kernel) launch at entry;
+    ``WAIT`` (when itself PDL-launched) waits for its predecessor before the
+    spin. ``stamp[0]`` = time the copy starts.
+    """
+    pid = tl.program_id(0)
+    if TRIGGER:
+        _gdc_launch_dependents()
+    if WAIT:
+        _gdc_wait()
+    t = _spin(delay_ns)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    x = tl.load(src_ptr + offs, mask=mask)
+    tl.store(dst_ptr + offs, x, mask=mask)
+    if pid == 0:
+        tl.store(stamp_ptr, t)
+
+
+@triton.jit
+def _delayed_consumer(
+    src_ptr, dst_ptr, stamp_ptr, numel, delay_ns, BLOCK: tl.constexpr, TRIGGER: tl.constexpr
+):
+    """PDL consumer of the attention_output_tail partial: wait, spin ``delay_ns``, copy.
+
+    ``stamp[0]`` = entry time (before the wait), ``stamp[1]`` = time the
+    wait returned, ``stamp[2]`` = time the copy starts.
+    """
+    pid = tl.program_id(0)
+    t_entry = _clock_fenced()
+    if TRIGGER:
+        _gdc_launch_dependents()
+    _gdc_wait()
+    t_wait = _clock_fenced()
+    t_copy = _spin(delay_ns)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    x = tl.load(src_ptr + offs, mask=mask)
+    tl.store(dst_ptr + offs, x, mask=mask)
+    if pid == 0:
+        tl.store(stamp_ptr, t_entry)
+        tl.store(stamp_ptr + 1, t_wait)
+        tl.store(stamp_ptr + 2, t_copy)
+
+
+_BLOCK = 1024
+
+
+def launch_producer(src, dst, stamp, delay_us, trigger, wait, pdl):
+    numel = dst.numel()
+    _delayed_producer[(triton.cdiv(numel, _BLOCK),)](
+        src,
+        dst,
+        stamp,
+        numel,
+        int(delay_us * 1000),
+        BLOCK=_BLOCK,
+        TRIGGER=trigger,
+        WAIT=wait,
+        num_warps=4,
+        **({"launch_pdl": True} if pdl else {}),
+    )
+
+
+def launch_consumer(src, dst, stamp, delay_us, trigger):
+    numel = src.numel()
+    _delayed_consumer[(triton.cdiv(numel, _BLOCK),)](
+        src,
+        dst,
+        stamp,
+        numel,
+        int(delay_us * 1000),
+        BLOCK=_BLOCK,
+        TRIGGER=trigger,
+        num_warps=4,
+        launch_pdl=True,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Fixtures.
+# ---------------------------------------------------------------------------
+
+_CACHE = {}
+
+
+def tail_weights(layer=3, rank=0):
+    key = ("w", layer, rank)
+    if key not in _CACHE:
+        if _HAVE_CKPT:
+            w_uv, w_o = tail_fixtures.load_tail_weights(_CKPT, layer, [rank], TP_SIZE)[rank]
+            _CACHE[key] = (w_uv.cuda(), w_o.cuda(), "real")
+        else:
+            gen = torch.Generator(device="cuda").manual_seed(layer * 8 + rank)
+            w_uv = (
+                torch.randn(LOCAL_HEADS, V_DIM, C_DIM, generator=gen, device="cuda") * 0.02
+            ).bfloat16()
+            w_o = (
+                torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, generator=gen, device="cuda") * 0.02
+            ).bfloat16()
+            _CACHE[key] = (w_uv, w_o, "synthetic")
+    return _CACHE[key]
+
+
+def make_latents(count, rows, seed, layer=3):
+    """``[count, rows, 8, 512]`` BF16 latents, one per captured iteration."""
+    if _HAVE_CKPT:
+        key = ("kv", layer)
+        if key not in _CACHE:
+            _CACHE[key] = tail_fixtures.compressed_kv_rows(_CKPT, layer)
+        return torch.stack(
+            [
+                tail_fixtures.mix_latent(
+                    _CACHE[key], rows, num_local_heads=LOCAL_HEADS, seed=seed + i
+                )
+                for i in range(count)
+            ]
+        )
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    return torch.randn((count, rows, LOCAL_HEADS, C_DIM), generator=gen, device="cuda").bfloat16()
+
+
+def eager_expected(latents, w_uv, w_o, rows, out_rows, config):
+    """attention_output_tail on each latent, one isolated call at a time."""
+    outs = []
+    for i in range(latents.shape[0]):
+        lat = latents[i].clone()
+        torch.cuda.synchronize()
+        outs.append(attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows, config=config).clone())
+        torch.cuda.synchronize()
+    return torch.stack(outs)
+
+
+def check_against_native(latents, expected, w_uv, w_o, rows, label):
+    """Frozen attention_output_tail budget (test_attn_tail_attn_tail) on the first latent."""
+    native = tail_reference.native_attn_tail(latents[0], w_uv, w_o, rows)
+    ref, _ = tail_reference.oracle_attn_tail(latents[0], w_uv, w_o, rows)
+    e_attn_tail = tail_reference.error_stats(expected[0][:rows], ref)
+    e_nat = tail_reference.error_stats(native, ref)
+    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (label, e_attn_tail, e_nat)
+    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_nat["ref_max"]), (
+        label,
+        e_attn_tail,
+        e_nat,
+    )
+
+
+def assert_bitwise(caps, expected, label):
+    assert not bool(torch.isnan(caps).any()), (label, "NaN in a partial")
+    bad = [i for i in range(caps.shape[0]) if not torch.equal(caps[i], expected[i])]
+    if bad:
+        i = bad[0]
+        diff = (caps[i].float() - expected[i].float()).abs()
+        pytest.fail(
+            f"{label}: {len(bad)}/{caps.shape[0]} iterations differ "
+            f"from the eager attention_output_tail result (first {i}: "
+            f"{int((diff > 0).sum())} elements, max |diff| "
+            f"{diff.max().item():.3e})"
+        )
+
+
+def _median_us(values):
+    return statistics.median(values) / 1000.0 if values else float("nan")
+
+
+# ---------------------------------------------------------------------------
+# Tests.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1), (8, 8)])
+def test_eager_exact_buffers(rows, out_rows):
+    """Exactly sized latent / v / partial (memcheck with the caching
+    allocator off sees every out-of-bounds access), default config vs the
+    fully serialized config, bit for bit, and the frozen budget vs native."""
+    w_uv, w_o, kind = tail_weights()
+    lat = make_latents(1, rows, seed=11)[0].contiguous()
+    assert check_tail_shapes(lat, w_uv, w_o, rows, out_rows) is None
+    v = torch.empty((rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
+    out = attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows, v_scratch=v)
+    serial = attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows, use_pdl=False, config=SERIAL_CONFIG)
+    torch.cuda.synchronize()
+    assert torch.equal(out, serial), "PDL and serialized attention_output_tail differ"
+    assert bool((out[rows:] == 0).all())
+    check_against_native(lat[None], out[None], w_uv, w_o, rows, ("eager", kind, rows, out_rows))
+    _, v_nat = tail_reference.native_attn_tail(lat, w_uv, w_o, rows, return_v=True)
+    print(
+        f"[attention_output_tail-pdl-stress-eager] weights={kind} rows={rows} out_rows={out_rows} "
+        f"pdl_equals_serial=True v_equals_native={torch.equal(v, v_nat)}"
+    )
+
+
+def _run_delayed_producer(rows, out_rows, delay_us, config):
+    w_uv, w_o, kind = tail_weights()
+    n = LOCAL_HEADS * V_DIM
+    latent = torch.empty((rows, LOCAL_HEADS, C_DIM), dtype=torch.bfloat16, device="cuda")
+    v = torch.empty((rows, n), dtype=torch.bfloat16, device="cuda")
+    srcs = make_latents(N_ITERS, rows, seed=1000)
+    caps = torch.empty((N_ITERS, out_rows, HIDDEN), dtype=torch.bfloat16, device="cuda")
+    stamps = torch.zeros((N_ITERS, 4), dtype=torch.int64, device="cuda")
+
+    def body():
+        for i in range(N_ITERS):
+            # Poison (ordinary, fully serialized kernels).
+            latent.fill_(float("nan"))
+            v.fill_(float("nan"))
+            # Not PDL-launched itself, but lets attention_output_tail's W_UV launch at its entry.
+            launch_producer(srcs[i], latent, stamps[i], delay_us, True, False, False)
+            out = attn_tail(latent, w_uv, w_o, rows, out_rows=out_rows, v_scratch=v, config=config)
+            launch_consumer(out, caps[i], stamps[i, 1:], 0, False)
+
+    body()  # eager pass: JIT compile + a first stressed run
+    torch.cuda.synchronize()
+    expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, config)
+    assert_bitwise(caps, expected, "delayed-producer eager")
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        body()
+    gaps = []
+    for rep in range(N_REPLAYS):
+        srcs.copy_(make_latents(N_ITERS, rows, seed=2000 + 100 * rep))
+        caps.fill_(float("nan"))
+        torch.cuda.synchronize()
+        graph.replay()
+        torch.cuda.synchronize()
+        expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, config)
+        assert_bitwise(caps, expected, f"delayed-producer replay {rep}")
+        st = stamps.cpu()
+        # Producer copy start -> consumer wait return: attention_output_tail's critical path
+        # after the latent becomes available (plus the copy itself).
+        gaps += [int(st[i, 2] - st[i, 0]) for i in range(N_ITERS)]
+    check_against_native(srcs, expected, w_uv, w_o, rows, ("delayed-producer", kind))
+    del graph
+    return kind, gaps
+
+
+@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1), (8, 8)])
+@pytest.mark.parametrize("delay_us", [0, 3, 20, 50])
+@pytest.mark.parametrize("entry", ["pdl_entry", "no_entry"])
+def test_delayed_producer(rows, out_rows, delay_us, entry):
+    config = DEFAULT_CONFIG if entry == "pdl_entry" else TailConfig(pdl_entry=False)
+    kind, gaps = _run_delayed_producer(rows, out_rows, delay_us, config)
+    print(
+        f"[attention_output_tail-pdl-stress-producer] weights={kind} rows={rows} "
+        f"out_rows={out_rows} delay_us={delay_us} entry={entry} "
+        f"iterations={len(gaps)} bitwise_ok=True "
+        f"producer_copy_to_consumer_wait_us_median={_median_us(gaps):.2f}"
+    )
+
+
+def _run_delayed_consumer(rows, out_rows, producer_us, consumer_us):
+    w_uv, w_o, kind = tail_weights()
+    latent = torch.empty((rows, LOCAL_HEADS, C_DIM), dtype=torch.bfloat16, device="cuda")
+    srcs = make_latents(N_ITERS, rows, seed=3000)
+    caps = torch.empty((N_ITERS, out_rows, HIDDEN), dtype=torch.bfloat16, device="cuda")
+    stamps = torch.zeros((N_ITERS, 4), dtype=torch.int64, device="cuda")
+
+    def body():
+        out_ptrs, v_ptrs = [], []
+        for i in range(N_ITERS):
+            # PDL-launched and triggering at entry: the chain may run ahead
+            # into this iteration while the previous consumer still spins.
+            launch_producer(srcs[i], latent, stamps[i], producer_us, True, True, True)
+            v = torch.empty((rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
+            out = attn_tail(latent, w_uv, w_o, rows, out_rows=out_rows, v_scratch=v)
+            out_ptrs.append(out.data_ptr())
+            v_ptrs.append(v.data_ptr())
+            launch_consumer(out, caps[i], stamps[i, 1:], consumer_us, True)
+            del out, v  # back to the pool: the next iteration reuses them
+        return out_ptrs, v_ptrs
+
+    body()
+    torch.cuda.synchronize()
+    expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, DEFAULT_CONFIG)
+    assert_bitwise(caps, expected, "delayed-consumer eager")
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out_ptrs, v_ptrs = body()
+    early = []
+    for rep in range(N_REPLAYS):
+        srcs.copy_(make_latents(N_ITERS, rows, seed=4000 + 100 * rep))
+        caps.fill_(float("nan"))
+        torch.cuda.synchronize()
+        graph.replay()
+        torch.cuda.synchronize()
+        expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, DEFAULT_CONFIG)
+        assert_bitwise(caps, expected, f"delayed-consumer replay {rep}")
+        st = stamps.cpu()
+        # Consumer resident before attention_output_tail finished: time spent in its wait.
+        early += [int(st[i, 2] - st[i, 1]) for i in range(N_ITERS)]
+    check_against_native(srcs, expected, w_uv, w_o, rows, ("delayed-consumer", kind))
+    del graph
+    reuse = {
+        "distinct_out_ptrs": len(set(out_ptrs)),
+        "distinct_v_ptrs": len(set(v_ptrs)),
+        "v_overlaps_a_partial": any(
+            any(o <= p < o + out_rows * HIDDEN * 2 for o in out_ptrs) for p in v_ptrs
+        ),
+    }
+    return kind, early, reuse
+
+
+@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1)])
+@pytest.mark.parametrize("producer_us,consumer_us", [(0, 0), (0, 20), (20, 20), (5, 50)])
+def test_delayed_consumer_buffer_reuse(rows, out_rows, producer_us, consumer_us):
+    kind, early, reuse = _run_delayed_consumer(rows, out_rows, producer_us, consumer_us)
+    # The pool hands the same partial buffer to consecutive iterations, so
+    # the stress covers a reused output buffer.
+    assert reuse["distinct_out_ptrs"] < N_ITERS, reuse
+    frac = sum(e > 1000 for e in early) / max(len(early), 1)
+    print(
+        f"[attention_output_tail-pdl-stress-consumer] weights={kind} rows={rows} "
+        f"out_rows={out_rows} producer_us={producer_us} "
+        f"consumer_us={consumer_us} iterations={len(early)} bitwise_ok=True "
+        f"reuse={reuse} consumer_wait_us_median={_median_us(early):.2f} "
+        f"fraction_waited_over_1us={frac:.2f}"
+    )
+
+
+@pytest.mark.skipif(UNDER_SANITIZER, reason="the sanitizer serializes kernel launches")
+def test_pdl_overlap_witness():
+    """The stress is not vacuous: the PDL consumer of attention_output_tail is resident and
+    waiting before attention_output_tail has finished (attention_output_tail triggers its dependents
+    early), and attention_output_tail's PDL entry shortens the producer → consumer path (W_UV
+    launched under the producer)."""
+    _, early, _ = _run_delayed_consumer(6, 6, 0, 20)
+    frac = sum(e > 1000 for e in early) / len(early)
+    _, gaps_entry = _run_delayed_producer(6, 6, 20, DEFAULT_CONFIG)
+    _, gaps_noentry = _run_delayed_producer(6, 6, 20, TailConfig(pdl_entry=False))
+    print(
+        f"[attention_output_tail-pdl-stress-witness] consumer_wait_us_median={_median_us(early):.2f} "
+        f"fraction_waited_over_1us={frac:.2f} "
+        f"producer_to_consumer_us_median pdl_entry="
+        f"{_median_us(gaps_entry):.2f} no_entry={_median_us(gaps_noentry):.2f}"
+    )
+    assert frac >= 0.5, frac
+    assert _median_us(gaps_entry) < _median_us(gaps_noentry)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py
new file mode 100644
index 000000000000..20f230b6e21b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py
@@ -0,0 +1,115 @@
+# 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.
+"""Attribution probes for compute-sanitizer initcheck reports.
+
+Run under ``compute-sanitizer --tool initcheck`` with
+``PYTORCH_NO_CUDA_MEMORY_CACHING=1``.
+Each probe produces one fully written tensor and then reads it once with an
+ordinary PyTorch conversion kernel, so pytest's per-test output attributes
+every report to one producer: native producers (H2D copy, elementwise,
+cuBLAS GEMMs, the native oracle tail) and the default attention_output_tail. A report under
+a native probe concerns how that producer's writes are tracked; a report
+under the attention_output_tail probe would mean attention_output_tail left part of its partial
+unwritten.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+def _read(y):
+    return float(y.float().abs().sum())
+
+
+def test_probe_h2d_copy_then_read():
+    x = torch.randn(2048, 512, dtype=torch.bfloat16).cuda()
+    torch.cuda.synchronize()
+    assert _read(x) > 0
+
+
+def test_probe_elementwise_then_read():
+    x = torch.randn(2048, 512, device="cuda").bfloat16()
+    torch.cuda.synchronize()
+    assert _read(x) > 0
+
+
+def test_probe_cublas_linear_then_read():
+    x = torch.randn(2048, 6144, device="cuda").bfloat16()
+    w = torch.randn(576, 6144, device="cuda").bfloat16()
+    y = torch.nn.functional.linear(x, w)
+    torch.cuda.synchronize()
+    assert _read(y) > 0
+
+
+def test_probe_cublas_linear_small_m_then_read():
+    x = torch.randn(6, 2048, device="cuda").bfloat16()
+    w = torch.randn(6144, 2048, device="cuda").bfloat16()
+    y = torch.nn.functional.linear(x, w)
+    torch.cuda.synchronize()
+    assert _read(y) > 0
+
+
+def test_probe_cublas_bmm_then_read():
+    a = torch.randn(8, 6, 512, device="cuda").bfloat16()
+    b = torch.randn(8, 512, 256, device="cuda").bfloat16()
+    y = torch.bmm(a, b)
+    torch.cuda.synchronize()
+    assert _read(y) > 0
+
+
+def _tail_inputs(rows):
+    gen = torch.Generator(device="cuda").manual_seed(rows)
+    w_uv = (torch.randn(8, 256, 512, generator=gen, device="cuda") * 0.02).bfloat16()
+    w_o = (torch.randn(6144, 2048, generator=gen, device="cuda") * 0.02).bfloat16()
+    lat = torch.randn(rows, 8, 512, generator=gen, device="cuda").bfloat16()
+    return lat, w_uv, w_o
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_probe_native_tail_then_read(rows):
+    """The native oracle tail (bmm into the o_proj input buffer, then
+    F.linear), read once by a conversion kernel."""
+    import tail_reference
+
+    lat, w_uv, w_o = _tail_inputs(rows)
+    native = tail_reference.native_attn_tail(lat, w_uv, w_o, rows)
+    torch.cuda.synchronize()
+    assert _read(native) > 0
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="attention_output_tail kernels target SM100/SM103",
+)
+@pytest.mark.parametrize("rows,out_rows", [(1, 1), (6, 6), (6, 8), (8, 8)])
+def test_probe_attn_tail_then_read(rows, out_rows):
+    """The default attention_output_tail (PDL on), its partial read once by a conversion
+    kernel: every element of the partial, padded rows included, must be
+    written by attention_output_tail."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
+        attn_tail,
+    )
+
+    lat, w_uv, w_o = _tail_inputs(rows)
+    out = attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows)
+    torch.cuda.synchronize()
+    assert _read(out) > 0
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py
new file mode 100644
index 000000000000..86f3015b779a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py
@@ -0,0 +1,358 @@
+# 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.
+"""attention_output_into_hop producer (single GPU): the publishing attention_output_tail o_proj and the tail protocol.
+
+A plain device buffer stands in for the post_attention_hop workspace (multicast view and
+local copy are the same memory on one GPU), filled with the Lamport sentinel,
+with a protocol state selecting the current buffer. On the real DSA ``MLA``
+of TP8 rank 0 with real layer-3 weights and synthetic latents:
+
+* ``request_into_hop`` → ``arm`` → ``defer_w_uv`` → ``project`` publishes exactly
+  slot ``[row][rank]`` of the current buffer for every output row (live rows
+  and the zero rows of padding the post_attention_hop also sums), as 32-bit words equal
+  to the BF16 partial with ``-0.0`` mapped to ``+0.0``; every other word
+  stays the sentinel; ``take_into_hop_published`` reports it once;
+* the ordinary partial is bit-identical to the attention_output_tail without publishing;
+* a request is consumed by the next ``arm`` even when the call declines;
+  non-attention_output_tail engines never publish; a request whose hop row count differs from
+  the output rows publishes nothing (the layer then runs the post_attention_hop);
+* CUDA-graph capture and replay publish into the same slots, also with
+  padded rows;
+* both publish widths (32-bit words on 152 CTAs or on request, 16-byte
+  stores on 128 CTAs) write the same words;
+* ``make_attn_tail`` engages the attention_output_into_hop schedules only next to the
+  post_attention_hop (``INTO_HOP_SPARSE_DECODE_SCHEDULE`` behind ``sparse_mla_decode``,
+  ``INTO_HOP_SCHEDULE`` behind the native FMHA); the plan turns attention_output_into_hop
+  off with either path it extends.
+"""
+
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
+sys.path.insert(0, os.path.dirname(_HERE))
+import plan_testlib  # noqa: E402
+import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
+    dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
+    DEFAULT_CONFIG,
+    INTO_HOP_SCHEDULE,
+    INTO_HOP_SPARSE_DECODE_SCHEDULE,
+    TAIL_PREFETCH_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+    TailConfig,
+    attn_tail,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
+    AttnTail,
+    make_attn_tail,
+)
+
+pytestmark = tail_adapter.pytestmark
+
+H = 6144
+RANKS = 8
+MAX_ROWS = 8
+BUF_WORDS = MAX_ROWS * RANKS * H * 2 // 4
+SENTINEL = -(2**31)  # 0x80000000 as int32
+
+
+def _transport(cur=1, rank=3):
+    mc = torch.full((3 * BUF_WORDS,), SENTINEL, dtype=torch.int32, device="cuda")
+    flags = torch.tensor([cur, 0, 0, 0], dtype=torch.int32, device="cuda")
+    return SimpleNamespace(mc_words=mc, flags=flags, rank=rank, cur=cur)
+
+
+def _sanitized_words(partial):
+    bits = partial.contiguous().view(torch.int16).clone()
+    bits[bits == -32768] = 0  # bf16 -0.0 (0x8000) -> +0.0
+    return bits.view(torch.int32)
+
+
+def _check_slots(tr, partial, rows, live=None):
+    """Slots ``[:rows, rank]`` hold ``partial``'s words; rows past ``live``
+    are zero words; every other word of every buffer is the sentinel."""
+    buf = tr.mc_words[tr.cur * BUF_WORDS : (tr.cur + 1) * BUF_WORDS].view(MAX_ROWS, RANKS, H // 2)
+    want = _sanitized_words(partial[:rows])
+    assert torch.equal(buf[:rows, tr.rank], want)
+    assert not bool((buf[:rows, tr.rank] == SENTINEL).any())
+    if live is not None:
+        assert bool((buf[live:rows, tr.rank] == 0).all())
+    others = torch.ones(MAX_ROWS, RANKS, dtype=torch.bool, device="cuda")
+    others[:rows, tr.rank] = False
+    assert bool((buf[others] == SENTINEL).all())
+    for b in range(3):
+        if b != tr.cur:
+            assert bool((tr.mc_words[b * BUF_WORDS : (b + 1) * BUF_WORDS] == SENTINEL).all())
+
+
+def _step(tail, mla, latent, tr, out_rows=None, hop_rows=None):
+    """One armed call; the attention output has ``out_rows`` rows (live rows
+    first, as in ``MLA.forward``) and the request names ``hop_rows``."""
+    rows = latent.shape[0]
+    out_rows = rows if out_rows is None else out_rows
+    hop_rows = out_rows if hop_rows is None else hop_rows
+    out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    tail.request_into_hop(tr, hop_rows)
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point="attn")
+    assert tail.defer_w_uv(mla, meta, latent, out[:rows])
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    return tail.project(mla, out, AllReduceParams(enable_allreduce=False))
+
+
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+@pytest.mark.parametrize("grid,vector", [(0, True), (128, True), (128, False)])
+def test_publish_slots_and_partial(rows, grid, vector):
+    mla = tail_adapter._mla()
+    cfg = TailConfig(oproj_grid=grid, into_hop_vector_publish=vector)
+    tail = AttnTail("layer3", cfg)
+    tr = _transport()
+    latent = tail_adapter._latent(rows, seed=rows)
+    partial = _step(tail, mla, latent, tr)
+    torch.cuda.synchronize()
+    assert tail.take_into_hop_published() and not tail.take_into_hop_published()
+    ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows, config=cfg)
+    assert torch.equal(partial, ref)
+    _check_slots(tr, partial, rows)
+
+
+@pytest.mark.parametrize("live,out_rows", [(1, 4), (2, 3), (6, 8), (7, 8), (1, 8)])
+@pytest.mark.parametrize("grid,vector", [(0, True), (128, True), (128, False)])
+def test_publish_padded_rows(live, out_rows, grid, vector):
+    """Padded output rows (``hidden_states`` rows > ``num_tokens``): every
+    output row is published, the padding as zero words, so the hop, which
+    polls the residual's rows, sees all of them (and sums zeros, as the post_attention_hop
+    hop does from the zero rows of the ordinary partial)."""
+    mla = tail_adapter._mla()
+    cfg = TailConfig(oproj_grid=grid, into_hop_vector_publish=vector)
+    tail = AttnTail("layer3", cfg)
+    tr = _transport(cur=0, rank=6)
+    latent = tail_adapter._latent(live, seed=30 + live)
+    partial = _step(tail, mla, latent, tr, out_rows=out_rows)
+    torch.cuda.synchronize()
+    assert tail.take_into_hop_published()
+    assert tail.into_hop_published_calls == 1 and tail.into_hop_declined_calls == 0
+    ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, live, out_rows=out_rows, config=cfg)
+    assert partial.shape[0] == out_rows
+    assert torch.equal(partial, ref)
+    assert bool((partial[live:] == 0).all())
+    _check_slots(tr, partial, out_rows, live=live)
+
+
+@pytest.mark.parametrize("live,out_rows,hop_rows", [(6, 8, 6), (6, 6, 8), (4, 4, 3)])
+def test_hop_row_mismatch_publishes_nothing(live, out_rows, hop_rows):
+    """A request whose hop row count differs from the output rows is not
+    honoured: nothing is published (the hop resets exactly the rows it
+    polls, so extra or missing rows would corrupt the transport), the
+    ordinary partial is unchanged and the layer runs the post_attention_hop."""
+    mla = tail_adapter._mla()
+    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
+    tr = _transport()
+    latent = tail_adapter._latent(live, seed=70 + live)
+    partial = _step(tail, mla, latent, tr, out_rows=out_rows, hop_rows=hop_rows)
+    torch.cuda.synchronize()
+    assert not tail.take_into_hop_published()
+    assert tail.into_hop_published_calls == 0 and tail.into_hop_declined_calls == 1
+    assert tail.engaged_calls == 1
+    ref = attn_tail(
+        latent, mla.v_b_proj, mla.o_proj.weight, live, out_rows=out_rows, config=INTO_HOP_SCHEDULE
+    )
+    assert torch.equal(partial, ref)
+    assert bool((tr.mc_words == SENTINEL).all())
+
+
+@pytest.mark.parametrize("live,out_rows", [(6, 6), (6, 8), (1, 4), (8, 8)])
+@pytest.mark.parametrize("grid,vector", [(0, True), (128, True), (128, False)])
+def test_publisher_synthetic_inputs(live, out_rows, grid, vector):
+    """The publishing o_proj alone on fully initialized synthetic inputs
+    (no MLA module, so a compute-sanitizer run sees only this kernel and
+    its reference): ordinary partial bit-identical to ``stream_oproj``,
+    every output row published."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_oproj import (  # noqa: E501
+        stream_oproj,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.into_hop import (  # noqa: E501
+        stream_oproj_into_hop,
+    )
+
+    gen = torch.Generator(device="cuda").manual_seed(17 * live + out_rows)
+    v = (torch.randn(live, 2048, generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
+    w_o = (torch.randn(H, 2048, generator=gen, device="cuda") * 0.02).to(torch.bfloat16)
+    tr = _transport(cur=2, rank=1)
+    out = torch.empty((out_rows, H), dtype=torch.bfloat16, device="cuda")
+    stream_oproj_into_hop(
+        v,
+        w_o,
+        out,
+        True,
+        tr.mc_words,
+        tr.flags,
+        tr.rank,
+        grid=grid or None,
+        trigger_dependents=True,
+        vector_publish=vector,
+    )
+    ref = torch.empty((out_rows, H), dtype=torch.bfloat16, device="cuda")
+    stream_oproj(v, w_o, ref, use_pdl=True, grid=grid or None)
+    torch.cuda.synchronize()
+    assert torch.equal(out, ref)
+    _check_slots(tr, out, out_rows, live=live)
+
+
+def test_publisher_rejects_rows_outside_transport():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.into_hop import (  # noqa: E501
+        stream_oproj_into_hop,
+    )
+
+    mla = tail_adapter._mla()
+    tr = _transport()
+    v = torch.zeros((6, 2048), dtype=torch.bfloat16, device="cuda")
+    for out_rows in (9, 5):
+        out = torch.empty((out_rows, H), dtype=torch.bfloat16, device="cuda")
+        with pytest.raises(ValueError):
+            stream_oproj_into_hop(v, mla.o_proj.weight, out, True, tr.mc_words, tr.flags, tr.rank)
+
+
+def test_request_consumed_and_engines():
+    mla = tail_adapter._mla()
+    tr = _transport()
+    # Declined call: the request is consumed, nothing is published.
+    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
+    tail.request_into_hop(tr, 6)
+    assert not tail.arm(mla, tail_adapter._meta(6, num_contexts=1))
+    assert tail._into_hop_request is None and tail._into_hop is None
+    assert not tail.take_into_hop_published()
+    # Without a request the attention_output_tail does not publish.
+    latent = tail_adapter._latent(6, seed=9)
+    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(6)
+    assert tail.arm(mla, meta)
+    assert tail.defer_w_uv(mla, meta, latent, out)
+    tail.project(mla, out, None)
+    torch.cuda.synchronize()
+    assert not tail.take_into_hop_published()
+    assert bool((tr.mc_words == SENTINEL).all())
+    # A non-attention_output_tail engine never publishes.
+    native = AttnTail("layer3", TailConfig(oproj_engine="triton"))
+    native.request_into_hop(tr, 6)
+    assert native.arm(mla, meta)
+    assert native._into_hop is None
+
+
+def test_request_dropped_after_a_compile_declined_call():
+    """Under torch.compile ``arm`` returns before it reads the request; the
+    layer's ``take_into_hop_published`` drops it, so a later call without a request
+    never honours it (no publication into the transport)."""
+    from tensorrt_llm._torch.utils import torch_compiling
+
+    mla = tail_adapter._mla()
+    tr = _transport()
+    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
+    tail.request_into_hop(tr, 6)
+    with torch_compiling(True):
+        assert not tail.arm(mla, tail_adapter._meta(6))
+    assert tail._into_hop_request is not None  # arm touches no state under compile
+    assert not tail.take_into_hop_published()
+    assert tail._into_hop_request is None
+    latent = tail_adapter._latent(6, seed=9)
+    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(6)
+    assert tail.arm(mla, meta)
+    assert tail._into_hop is None
+    assert tail.defer_w_uv(mla, meta, latent, out)
+    tail.project(mla, out, None)
+    torch.cuda.synchronize()
+    assert not tail.take_into_hop_published()
+    assert bool((tr.mc_words == SENTINEL).all())
+
+
+@pytest.mark.parametrize("out_rows", [6, 8])
+def test_graph_capture_and_replay_publish(out_rows):
+    mla = tail_adapter._mla()
+    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
+    rows = 6
+    tr = _transport(cur=2, rank=5)
+    lat = tail_adapter._latent(rows, seed=21).contiguous()
+    _step(tail, mla, lat, tr, out_rows=out_rows)
+    tail.take_into_hop_published()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = _step(tail, mla, lat, tr, out_rows=out_rows)
+    assert tail.take_into_hop_published()
+    for seed in range(4):
+        tr.mc_words.fill_(SENTINEL)
+        lat.copy_(tail_adapter._latent(rows, seed=800 + seed))
+        graph.replay()
+        torch.cuda.synchronize()
+        ref = attn_tail(
+            lat, mla.v_b_proj, mla.o_proj.weight, rows, out_rows=out_rows, config=INTO_HOP_SCHEDULE
+        )
+        assert torch.equal(captured, ref)
+        _check_slots(tr, captured, out_rows, live=rows)
+
+
+def test_make_attn_tail_selects_into_hop(monkeypatch):
+    mla = tail_adapter._build_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
+    cases = [
+        (
+            "attention_output_tail,attention_output_into_hop,post_attention_hop,sparse_mla_decode",
+            INTO_HOP_SPARSE_DECODE_SCHEDULE,
+        ),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,attention_output_into_hop,post_attention_hop,sparse_mla_decode",
+            INTO_HOP_SPARSE_DECODE_SCHEDULE,
+        ),
+        ("attention_output_tail,attention_output_into_hop,post_attention_hop", INTO_HOP_SCHEDULE),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,attention_output_into_hop,post_attention_hop",
+            INTO_HOP_SCHEDULE,
+        ),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,post_attention_hop,sparse_mla_decode",
+            TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+        ),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,sparse_mla_decode",
+            TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+        ),
+        # attention_output_tail_prefetch has no schedule next to the post_attention_hop
+        # without sparse_mla_decode; the tail keeps its default schedule.
+        ("attention_output_tail,attention_output_tail_prefetch,post_attention_hop", DEFAULT_CONFIG),
+        ("attention_output_tail,attention_output_tail_prefetch", TAIL_PREFETCH_SCHEDULE),
+        ("attention_output_tail", DEFAULT_CONFIG),
+    ]
+    for features, want in cases:
+        plan_testlib.select(dispatch, features, monkeypatch)
+        tail = make_attn_tail(mla)
+        assert tail is not None and tail.config == want, features
+
+
+def test_attention_output_into_hop_is_off_without_the_paths_it_extends():
+    for base in ("attention_output_tail", "post_attention_hop"):
+        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
+        assert not plan.enabled("attention_output_into_hop"), base
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py
new file mode 100644
index 000000000000..c07963d4f391
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py
@@ -0,0 +1,531 @@
+# 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.
+"""attention_output_tail_prefetch: off-path L2 prefetch of the attention_output_tail weights (protocol and kernel).
+
+Uses the real DSA ``MLA`` module of TP8 rank 0 with real layer-3
+``v_b_proj``/``o_proj`` weights and synthetic latents (the attention_output_tail adapter test's
+builders); the real ``MLA.forward`` path is in ``test_tail_prefetch_mla_forward.py``.
+
+* the prefetch kernel runs on every config the benchmarks use and changes no
+  tensor;
+* attention_output_tail_prefetch's default tail (native W_UV + native o_proj on prefetched weights)
+  is bit-identical to the native tail, and with the attention_output_tail kernels on prefetched
+  weights it is bit-identical to the attention_output_tail: the prefetch has no effect on
+  results;
+* fork/join protocol: ``before_attention`` forks only on an armed call and
+  only for layers the config selects; ``project`` always joins (engaged or
+  declined), at every configured join point; a leftover fork raises in the
+  next ``arm``;
+* CUDA-graph capture of arm → before_attention → defer → project replays
+  correctly with the auxiliary stream rejoined;
+* ``make_attn_tail`` picks the attention_output_tail_prefetch schedule of the running
+  neighbours exactly when ``attention_output_tail_prefetch`` runs;
+* the adopted attention (``GlmMLA``) carries the attention_output_tail_prefetch flag and its tail
+  uses ``TAIL_PREFETCH_SCHEDULE``; the layer adoption hands the MoE gate weight to the tail.
+"""
+
+import dataclasses
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
+sys.path.insert(0, os.path.dirname(_HERE))
+import plan_testlib  # noqa: E402
+import tail_reference  # noqa: E402
+import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
+    dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail import (  # noqa: E402, E501
+    prefetch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
+    DEFAULT_CONFIG,
+    TAIL_PREFETCH_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+    TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE,
+    attn_tail,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
+    TailConfig as TailConfig_,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
+    AttnTail,
+    make_attn_tail,
+)
+
+pytestmark = tail_adapter.pytestmark
+
+_R = dataclasses.replace
+STAGED = TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE
+# attention_output_tail_prefetch's prefetch protocol with the attention_output_tail kernels (the
+# default uses the native ones).
+TAIL_PREFETCH_KERNEL_CONFIGS = _R(STAGED, oproj_engine="cute_stream")
+
+
+def _no_ar():
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    return AllReduceParams(enable_allreduce=False)
+
+
+@pytest.mark.parametrize(
+    "cfg",
+    [
+        prefetch.PrefetchConfig(),
+        prefetch.PrefetchConfig(ctas=152),
+        prefetch.PrefetchConfig(chunk=16384, ctas=26),
+        prefetch.PrefetchConfig(gap_ns=500),
+    ],
+)
+def test_prefetch_kernel_changes_nothing(cfg):
+    mla = tail_adapter._mla()
+    weights = [mla.v_b_proj, mla.o_proj.weight]
+    assert prefetch.check_prefetch_tensors(weights, cfg.chunk) == ""
+    before = [w.clone() for w in weights]
+    probe = torch.randn(4096, device="cuda")
+    probe_before = probe.clone()
+    prefetch.launch_prefetch(weights, cfg)
+    torch.cuda.synchronize()
+    for w, b in zip(weights, before):
+        assert torch.equal(w, b)
+    assert torch.equal(probe, probe_before)
+
+
+def test_prefetch_rejects_bad_inputs():
+    a = torch.zeros(4096, dtype=torch.bfloat16, device="cuda")
+    assert prefetch.check_prefetch_tensors([a, a, a, a], 4096) != ""
+    assert prefetch.check_prefetch_tensors([a, a], 4000) != ""
+    odd = torch.zeros(1000, dtype=torch.bfloat16, device="cuda")
+    assert prefetch.check_prefetch_tensors([a, odd], 4096) != ""
+
+
+def _drive(tail, mla, latent, out_rows=None, decline=False):
+    rows = latent.shape[0]
+    out_rows = out_rows or rows
+    attn_out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point=tail.config.prefetch_fork)
+    if decline:
+        big = torch.zeros((9, 8, 512), dtype=torch.bfloat16, device="cuda")
+        assert not tail.defer_w_uv(mla, meta, big, attn_out[:rows])
+    else:
+        assert tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
+    return tail.project(mla, attn_out, _no_ar())
+
+
+@pytest.mark.parametrize("join", ["before_wuv", "before_oproj", "after_oproj"])
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_tail_prefetch_tail_bitwise_equals_attn_tail(join, rows):
+    mla = tail_adapter._mla()
+    latent = tail_adapter._latent(rows, seed=rows)
+    ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows, config=DEFAULT_CONFIG)
+    tail = AttnTail("layer3", _R(TAIL_PREFETCH_KERNEL_CONFIGS, prefetch_join=join))
+    got = _drive(tail, mla, latent)
+    torch.cuda.synchronize()
+    assert torch.equal(got, ref)
+    assert tail.prefetch_forks == 1
+    assert not tail._prefetch.forked and not tail._armed
+    assert tail._pending is None
+
+
+def test_fork_only_on_armed_calls_and_join_on_decline():
+    mla = tail_adapter._mla()
+    tail = AttnTail("layer3", TAIL_PREFETCH_KERNEL_CONFIGS)
+    # Not armed: no fork.
+    tail.before_attention(mla, point="attn")
+    assert tail.prefetch_forks == 0
+    assert not tail.arm(mla, tail_adapter._meta(6, num_contexts=1))
+    tail.before_attention(mla, point="attn")
+    assert tail.prefetch_forks == 0
+    # Armed, but called at the other fork point: no fork.
+    assert tail.arm(mla, tail_adapter._meta(6))
+    tail.before_attention(mla, point="start")
+    assert tail.prefetch_forks == 0
+    tail._armed = False
+    # Armed, forked, then declined at the W_UV step: project joins.
+    assert _drive(tail, mla, tail_adapter._latent(6, seed=3), decline=True) is None
+    assert tail.prefetch_forks == 1 and not tail._prefetch.forked
+    # An attention_output_tail-only tail never forks.
+    served_tail = AttnTail("layer3", DEFAULT_CONFIG)
+    out = _drive(served_tail, mla, tail_adapter._latent(6, seed=4))
+    assert out is not None and served_tail.prefetch_forks == 0
+    # A leftover fork is an error at the next arm.
+    assert tail.arm(mla, tail_adapter._meta(6))
+    tail.before_attention(mla, point="attn")
+    tail._armed = False
+    with pytest.raises(RuntimeError, match="not consumed"):
+        tail.arm(mla, tail_adapter._meta(6))
+    tail._prefetch.join()
+
+
+def test_layer_selection():
+    mla = tail_adapter._mla()
+    has_indexer = getattr(mla.mqa, "indexer", None) is not None
+    for layers in ("all", "shared", "full"):
+        tail = AttnTail("layer3", _R(TAIL_PREFETCH_KERNEL_CONFIGS, prefetch_layers=layers))
+        _drive(tail, mla, tail_adapter._latent(6, seed=5))
+        want = layers == "all" or (layers == "full") == has_indexer
+        assert tail.prefetch_forks == int(want), (layers, has_indexer)
+
+
+def _native_step(tail, mla, latent, out_rows=None):
+    """What the patched op and MLA.forward do with attention_output_tail_prefetch's native engine."""
+    from tensorrt_llm._torch.attention.mla import MLA
+
+    rows = latent.shape[0]
+    out_rows = out_rows or rows
+    attn_out = torch.zeros((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point="attn")
+    # The op's native W_UV (the tail declines the W_UV step).
+    assert not tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
+    v_view = attn_out[:rows].view(rows, mla.num_heads_tp, mla.v_head_dim)
+    MLA._bmm_bf16_out(
+        mla,
+        latent.transpose(0, 1),
+        mla.v_b_proj,
+        mla.v_b_proj.transpose(1, 2),
+        v_view.transpose(0, 1),
+    )
+    return tail.project(mla, attn_out, _no_ar())
+
+
+@pytest.mark.parametrize("join", ["before_oproj", "after_oproj"])
+@pytest.mark.parametrize("rows,out_rows", [(1, 1), (6, 6), (6, 8), (8, 8)])
+def test_native_engine_bitwise_equals_native(join, rows, out_rows):
+    mla = tail_adapter._mla()
+    latent = tail_adapter._latent(rows, seed=10 + rows)
+    native = tail_adapter._native(mla, latent, out_rows)
+    tail = AttnTail("layer3", _R(STAGED, prefetch_join=join))
+    got = _native_step(tail, mla, latent, out_rows)
+    torch.cuda.synchronize()
+    assert torch.equal(got[:rows], native[:rows])
+    assert tail.prefetch_forks == 1 and tail.engaged_calls == 1
+    assert not tail._prefetch.forked and not tail._armed
+
+
+def test_native_engine_without_fork_takes_native_path():
+    mla = tail_adapter._mla()
+    # Layer selection excludes this layer: no fork, so the caller's native
+    # projection runs (project returns None).
+    has_indexer = getattr(mla.mqa, "indexer", None) is not None
+    other = "shared" if has_indexer else "full"
+    tail = AttnTail("layer3", _R(STAGED, prefetch_layers=other))
+    rows = 6
+    attn_out = torch.zeros((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point="attn")
+    assert tail.prefetch_forks == 0
+    assert not tail.defer_w_uv(mla, meta, tail_adapter._latent(rows, seed=6), attn_out)
+    assert tail.project(mla, attn_out, _no_ar()) is None
+
+
+def test_native_engine_graph_capture_and_replay():
+    mla = tail_adapter._mla()
+    tail = AttnTail("layer3", STAGED)
+    rows = 6
+    lat = tail_adapter._latent(rows, seed=17).contiguous()
+    _native_step(tail, mla, lat)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = _native_step(tail, mla, lat)
+    assert not tail._prefetch.forked
+    for seed in range(6):
+        lat.copy_(tail_adapter._latent(rows, seed=300 + seed))
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(captured, tail_adapter._native(mla, lat, rows))
+
+
+def test_graph_capture_and_replay():
+    mla = tail_adapter._mla()
+    tail = AttnTail("layer3", TAIL_PREFETCH_KERNEL_CONFIGS)
+    rows = 6
+    meta = tail_adapter._meta(rows)
+    lat = tail_adapter._latent(rows, seed=7).contiguous()
+    out = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
+
+    def step():
+        assert tail.arm(mla, meta)
+        tail.before_attention(mla, point="attn")
+        assert tail.defer_w_uv(mla, meta, lat, out)
+        return tail.project(mla, out, _no_ar())
+
+    step()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = step()
+    assert not tail._prefetch.forked
+    for seed in range(6):
+        lat.copy_(tail_adapter._latent(rows, seed=200 + seed))
+        graph.replay()
+        torch.cuda.synchronize()
+        ref = attn_tail(lat, mla.v_b_proj, mla.o_proj.weight, rows, config=DEFAULT_CONFIG)
+        assert torch.equal(captured, ref)
+
+
+def test_make_attn_tail_selects_config(monkeypatch):
+    mla = tail_adapter._build_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
+    cases = [
+        ("attention_output_tail", DEFAULT_CONFIG),
+        ("attention_output_tail,attention_output_tail_prefetch", TAIL_PREFETCH_SCHEDULE),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,sparse_mla_decode",
+            TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
+        ),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,sparse_mla_decode,post_attention_hop",
+            TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+        ),
+        (
+            "attention_output_tail,attention_output_tail_prefetch,post_attention_hop",
+            DEFAULT_CONFIG,
+        ),
+    ]
+    for features, want in cases:
+        plan_testlib.select(dispatch, features, monkeypatch)
+        tail = make_attn_tail(mla)
+        assert tail is not None and tail.config == want, features
+
+
+@pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED") != "1",
+    reason="opt-in",
+)
+def test_adoption_selects_tail_prefetch(monkeypatch):
+    import inspect
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+        attention_forward,
+    )
+
+    assert "before_attention" in inspect.getsource(attention_forward.forward_absorption_generation)
+    plan_testlib.select(
+        dispatch, "attention_output_tail,attention_output_tail_prefetch", monkeypatch
+    )
+    mla = tail_adapter._build_adopted_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
+    assert mla._tail_prefetch
+    assert mla._attn_tail.config == TAIL_PREFETCH_SCHEDULE
+    plan_testlib.select(dispatch, "attention_output_tail", monkeypatch)
+    assert not tail_adapter._build_adopted_mla(
+        layer_idx=tail_adapter.LAYER, num_hidden_layers=78
+    )._tail_prefetch
+    draft = tail_adapter._build_adopted_mla(layer_idx=78, num_hidden_layers=78)
+    assert draft._attn_tail is None
+
+
+@pytest.mark.parametrize("grid,slots,early", [(96, 3, 3), (96, 3, 0), (128, 3, 3), (64, 3, 2)])
+@pytest.mark.parametrize("rows,out_rows", [(1, 1), (6, 6), (6, 8), (8, 8), (8, 16)])
+def test_ring_oproj_within_budget(grid, slots, early, rows, out_rows):
+    """The small-footprint o_proj against the FP64 oracle of the same ``v``
+    (frozen attention_output_tail budget relative to the native GEMM), padded rows zero."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_oproj_ring import (  # noqa: E501
+        ring_oproj,
+    )
+
+    mla = tail_adapter._mla()
+    latent = tail_adapter._latent(rows, seed=40 + rows)
+    native, v = tail_reference.native_attn_tail(
+        latent, mla.v_b_proj, mla.o_proj.weight, rows, out_rows=out_rows, return_v=True
+    )
+    out = torch.full(
+        (out_rows, mla.o_proj.weight.shape[0]), 7.0, dtype=torch.bfloat16, device="cuda"
+    )
+    ring_oproj(v, mla.o_proj.weight, out, use_pdl=False, grid=grid, slots=slots, early=early)
+    torch.cuda.synchronize()
+    ref = v[:rows].double() @ mla.o_proj.weight.double().t()
+    e_ring = tail_reference.error_stats(out[:rows], ref)
+    e_nat = tail_reference.error_stats(native[:rows], ref)
+    assert e_ring["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (e_ring, e_nat)
+    assert e_ring["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_ring["ref_max"])
+    assert bool((out[rows:] == 0).all())
+
+
+def test_ring_engine_graph_replay_matches_native():
+    mla = tail_adapter._mla()
+    cfg = _R(STAGED, native_oproj="ring")
+    tail = AttnTail("layer3", cfg)
+    rows = 6
+    lat = tail_adapter._latent(rows, seed=51).contiguous()
+    _native_step(tail, mla, lat)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = _native_step(tail, mla, lat)
+    for seed in range(6):
+        lat.copy_(tail_adapter._latent(rows, seed=500 + seed))
+        graph.replay()
+        torch.cuda.synchronize()
+        native = tail_adapter._native(mla, lat, rows)
+        torch.testing.assert_close(
+            captured.clone(),
+            native,
+            rtol=2**-6,
+            atol=2**-6 * native.float().pow(2).mean().sqrt().item(),
+        )
+    assert tail.prefetch_forks == 2 and not tail._prefetch.forked
+
+
+def test_fork_ahead_protocol():
+    mla = tail_adapter._mla()
+    prev = _R(STAGED, prefetch_fork="prev")
+    # Other fork points ignore fork_ahead.
+    tail = AttnTail("layer3", STAGED)
+    tail.fork_ahead(mla)
+    assert tail.prefetch_forks == 0
+    # Forked ahead, then an armed call: before_attention does not fork
+    # again, project runs the native o_proj and joins; bit-identical.
+    tail = AttnTail("layer3", prev)
+    rows = 6
+    latent = tail_adapter._latent(rows, seed=60)
+    tail.fork_ahead(mla)
+    assert tail.prefetch_forks == 1 and tail._prefetch.forked
+    meta = tail_adapter._meta(rows)
+    attn_out = torch.zeros((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    assert tail.arm(mla, meta)
+    tail.before_attention(mla, point="attn")
+    tail.before_attention(mla, point="start")
+    assert tail.prefetch_forks == 1
+    from tensorrt_llm._torch.attention.mla import MLA
+
+    assert not tail.defer_w_uv(mla, meta, latent, attn_out)
+    v_view = attn_out.view(rows, mla.num_heads_tp, mla.v_head_dim)
+    MLA._bmm_bf16_out(
+        mla,
+        latent.transpose(0, 1),
+        mla.v_b_proj,
+        mla.v_b_proj.transpose(1, 2),
+        v_view.transpose(0, 1),
+    )
+    got = tail.project(mla, attn_out, _no_ar())
+    torch.cuda.synchronize()
+    assert torch.equal(got, tail_adapter._native(mla, latent, rows))
+    assert not tail._prefetch.forked and not tail._ahead
+    # Forked ahead, then a call that declines: arm joins at once.
+    tail.fork_ahead(mla)
+    assert not tail.arm(mla, tail_adapter._meta(rows, num_contexts=1))
+    assert not tail._prefetch.forked and not tail._ahead
+    # Twice ahead without a call in between is a no-op the second time.
+    tail.fork_ahead(mla)
+    tail.fork_ahead(mla)
+    assert tail.prefetch_forks == 3
+    assert not tail.arm(mla, tail_adapter._meta(9))
+    assert not tail._prefetch.forked
+
+
+def test_fork_ahead_graph_capture():
+    mla = tail_adapter._mla()
+    tail = AttnTail("layer3", _R(STAGED, prefetch_fork="prev"))
+    rows = 6
+    lat = tail_adapter._latent(rows, seed=61).contiguous()
+
+    def step():
+        tail.fork_ahead(mla)
+        return _native_step(tail, mla, lat)
+
+    step()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = step()
+    assert not tail._prefetch.forked
+    for seed in range(4):
+        lat.copy_(tail_adapter._latent(rows, seed=700 + seed))
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(captured, tail_adapter._native(mla, lat, rows))
+
+
+def test_successor_prefetch_protocol_and_bitwise():
+    mla = tail_adapter._mla()
+    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
+    before = router.clone()
+    for staged in (False, True):
+        cfg = TailConfig_(successor_prefetch=True, prefetch_l2=staged, oproj_engine="cute_stream")
+        tail = AttnTail("layer3", cfg)
+        rows = 6
+        latent = tail_adapter._latent(rows, seed=80)
+        # Without successor weights and without tail staging: no fork.
+        if not staged:
+            assert _drive(tail, mla, latent) is not None
+            assert tail.prefetch_forks == 0
+        tail.set_successor_weights([router])
+        got = _drive(tail, mla, latent)
+        torch.cuda.synchronize()
+        ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows, config=DEFAULT_CONFIG)
+        assert torch.equal(got, ref)
+        assert tail.prefetch_forks == 1 and not tail._prefetch.forked
+    assert torch.equal(router, before)
+
+
+def test_prefetch_kernel_three_regions():
+    mla = tail_adapter._mla()
+    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
+    tensors = [mla.v_b_proj, mla.o_proj.weight, router]
+    assert prefetch.check_prefetch_tensors(tensors, 4096) == ""
+    copies = [t.clone() for t in tensors]
+    prefetch.launch_prefetch(tensors, prefetch.PrefetchConfig())
+    prefetch.launch_prefetch([router], prefetch.PrefetchConfig(ctas=8))
+    torch.cuda.synchronize()
+    for t, c in zip(tensors, copies):
+        assert torch.equal(t, c)
+
+
+@pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED") != "1",
+    reason="needs attention_output_tail and attention_output_tail_prefetch selected",
+)
+def test_adoption_registers_router_weight():
+    """The layer adoption hands the MoE gate weight to the tail."""
+    import inspect
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmDecoderLayer,
+    )
+
+    src = inspect.getsource(GlmDecoderLayer._adopt_tail_prefetch)
+    assert "set_successor_weights" in src and 'is_enabled("attention_output_tail_prefetch")' in src
+
+
+def test_ring_oproj_refuses_publishing():
+    """The ring o_proj has no live-rows-only publish path (the into-hop entry needs
+    every hop row); a caller that asks for it gets an error, not a partial
+    publication.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_oproj_ring import (  # noqa: E501
+        ring_oproj,
+    )
+
+    mla = tail_adapter._mla()
+    v = torch.zeros((6, mla.o_proj.weight.shape[1]), dtype=torch.bfloat16, device="cuda")
+    out = torch.empty((8, mla.o_proj.weight.shape[0]), dtype=torch.bfloat16, device="cuda")
+    words = torch.zeros(16, dtype=torch.int32, device="cuda")
+    flags = torch.zeros(4, dtype=torch.int32, device="cuda")
+    with pytest.raises(ValueError, match="withdrawn"):
+        ring_oproj(
+            v, mla.o_proj.weight, out, use_pdl=False, mc_words=words, flags=flags, publish_rank=0
+        )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py
new file mode 100644
index 000000000000..70d710a337a3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py
@@ -0,0 +1,142 @@
+# 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.
+"""attention_output_tail_prefetch inside the target's ``GlmMLA.forward`` (native DSA attention →
+attention_output_tail_prefetch tail).
+
+Runs with the plan that runs only ``attention_output_tail`` and ``attention_output_tail_prefetch``
+(opt-in: ``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED=1``). Builds TP8 rank 0 of
+real GLM-5.2 layers 2 (full indexer) and 3 (shared indexer) with the attention_output_tail test
+harness (real weights, FP8 KV cache filled by a real prefill, GEN metadata as the decoder runner
+prepares it) and runs GEN-only ``GlmMLA.forward`` as a decode step through the target's DSA
+attention op:
+
+* eager R = 6 and R = 1: with a successor weight registered (a router-shaped
+  tensor, as the layer adoption registers the MoE gate weight), the
+  target's ``forward_absorption_generation`` forks the staging exactly once
+  per layer call (right before the attention backend call) and ``project``
+  joins it; the attention_output_tail_prefetch output is bit-identical to the attention_output_tail
+  kernels of ``TAIL_PREFETCH_SCHEDULE`` applied to the latent the attention produced (staging
+  changes no result) and within the frozen attention_output_tail budget against the FP64
+  oracle;
+* CUDA-graph capture of the GEN step with the fork/join inside, replayed
+  with new inputs, within budget of a native eager rerun.
+
+Fixture labels: weights real, activations synthetic.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
+import test_attn_tail_mla_forward as tail_forward  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
+    TAIL_PREFETCH_SCHEDULE,
+    attn_tail,
+)
+
+pytestmark = [
+    *tail_forward.pytestmark[:2],
+    pytest.mark.skipif(
+        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED") != "1",
+        reason="needs attention_output_tail and attention_output_tail_prefetch selected (opt-in)",
+    ),
+    pytest.mark.threadleak(enabled=False),
+]
+
+LAYERS = tail_forward.LAYERS
+PREFETCH_PATHS = "attention_output_tail,attention_output_tail_prefetch"
+
+
+@pytest.mark.parametrize("rows", [6, 1])
+def test_eager_forward_forks_joins_and_matches_attn_tail(rows):
+    stack = tail_forward._stack(rows, PREFETCH_PATHS)
+    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
+    for layer in LAYERS:
+        assert stack.attn[layer]._tail_prefetch
+        assert stack.attn_tail_tails[layer].config == TAIL_PREFETCH_SCHEDULE
+        stack.attn_tail_tails[layer].set_successor_weights([router])
+    inputs = {layer: stack.gen_inputs(layer, step=0) for layer in LAYERS}
+    native, lat_native = tail_forward._run(stack, False, inputs, record=True)
+    forks = {layer: stack.attn_tail_tails[layer].prefetch_forks for layer in LAYERS}
+    attention_output_tail, lat = tail_forward._run(stack, True, inputs, record=True)
+    for layer in LAYERS:
+        tail = stack.attn_tail_tails[layer]
+        attn = stack.attn[layer]
+        assert tail.prefetch_forks == forks[layer] + 1, layer
+        assert not tail._prefetch.forked and not tail._armed
+        ref = attn_tail(
+            lat[layer], attn.v_b_proj, attn.o_proj.weight, rows, config=TAIL_PREFETCH_SCHEDULE
+        )
+        assert torch.equal(attention_output_tail[layer], ref), layer
+        e_tail_prefetch, e_nat = tail_forward._check_budget(
+            stack, layer, attention_output_tail[layer], native[layer], lat[layer]
+        )
+        print(
+            f"[attention_output_tail_prefetch-mla-forward] rows={rows} layer={layer} "
+            f"err attention_output_tail_prefetch={e_tail_prefetch['rel_rms']:.3e} native={e_nat['rel_rms']:.3e}"
+        )
+
+
+def test_cuda_graph_capture_and_replay():
+    stack = tail_forward._stack(6, PREFETCH_PATHS)
+    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
+    for layer in LAYERS:
+        stack.attn_tail_tails[layer].set_successor_weights([router])
+    pos = stack.position_ids()
+    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
+    stack.set_attn_tail(True)
+    with stack.scope():
+        for _ in range(2):
+            for layer in LAYERS:
+                stack.forward(layer, static_in[layer], pos, tail_forward._no_ar())
+        torch.cuda.synchronize()
+        forks = {layer: stack.attn_tail_tails[layer].prefetch_forks for layer in LAYERS}
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = {
+                layer: stack.forward(layer, static_in[layer], pos, tail_forward._no_ar())
+                for layer in LAYERS
+            }
+    for layer in LAYERS:
+        assert stack.attn_tail_tails[layer].prefetch_forks == forks[layer] + 1
+        assert not stack.attn_tail_tails[layer]._prefetch.forked
+    for step in range(1, 5):
+        for layer in LAYERS:
+            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
+        graph.replay()
+        torch.cuda.synchronize()
+        got = {layer: captured[layer].clone() for layer in LAYERS}
+        native, lat = tail_forward._run(
+            stack,
+            False,
+            {layer: static_in[layer].clone() for layer in LAYERS},
+            record=True,
+            position_ids=pos,
+        )
+        for layer in LAYERS:
+            e_tail_prefetch, e_nat = tail_forward._check_budget(
+                stack, layer, got[layer], native[layer], lat[layer]
+            )
+            print(
+                f"[attention_output_tail_prefetch-mla-graph] step={step} layer={layer} "
+                f"err attention_output_tail_prefetch={e_tail_prefetch['rel_rms']:.3e} "
+                f"native={e_nat['rel_rms']:.3e}"
+            )
+    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py
new file mode 100644
index 000000000000..c1ee2d65b758
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py
@@ -0,0 +1,171 @@
+# 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.
+"""Selector and call site of cross_rank_topk_cluster_merge.
+
+- ``cross_rank_topk_cluster_merge`` runs only with ``cross_rank_topk`` (disabling ``cross_rank_topk``
+  turns it off; disabling it keeps ``cross_rank_topk``) and is recorded as a
+  cross_rank_topk-identical ``rank-consistent`` variant.
+- The cross_rank_topk region of the indexer's decode top-k (``region.cross_rank_topk_decode``,
+  which ``GlmIndexer._fill_decode_topk`` asks first) tries cross_rank_topk_cluster_merge
+  first when ``cross_rank_topk_cluster_merge`` is selected and falls back to cross_rank_topk when
+  cross_rank_topk_cluster_merge declines; without ``cross_rank_topk_cluster_merge`` it calls
+  cross_rank_topk only (composition F unchanged), and without ``cross_rank_topk`` nothing.
+- The cross_rank_topk_cluster_merge package registers its entry points and declines a device that
+  cannot keep every row's cluster resident, and a group size its GPU tests
+  do not cover.
+"""
+
+import os
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+
+def _plan(disabled=()):
+    environ = {dispatch.DISABLE_ENV: ",".join(disabled)} if disabled else {}
+    return dispatch.resolve(dispatch.ModelingV2Mode.AUTO, environ)
+
+
+def test_selector_closure():
+    assert dispatch.REQUIRED_FEATURES["cross_rank_topk_cluster_merge"] == {"cross_rank_topk"}
+    assert _plan().enabled("cross_rank_topk_cluster_merge")
+    assert not _plan(["cross_rank_topk"]).enabled("cross_rank_topk_cluster_merge")
+    record = dispatch.feature_record(
+        ["cross_rank_topk", "cross_rank_topk_cluster_merge"], environ={}
+    )
+    assert (
+        "cross_rank_topk_cluster_merge[rank-consistent]" in record
+        and "cross_rank_topk[rank-consistent]" in record
+    )
+
+
+def test_disabling_cluster_merge_keeps_cross_rank_topk():
+    plan = _plan(["cross_rank_topk_cluster_merge"])
+    assert plan.enabled("cross_rank_topk")
+    assert not plan.enabled("cross_rank_topk_cluster_merge")
+
+
+class _Calls:
+    def __init__(self, cluster_merge_result):
+        self.cluster_merge_result = cluster_merge_result
+        self.log = []
+
+    def cross_rank_topk_cluster_merge(self, *args):
+        self.log.append("cross_rank_topk_cluster_merge")
+        return self.cluster_merge_result
+
+    def indexer_topk(self, *args):
+        self.log.append("indexer_topk")
+        return True
+
+
+@pytest.mark.parametrize(
+    "features, cluster_merge_result, expected, calls",
+    [
+        (
+            frozenset({"cross_rank_topk", "cross_rank_topk_cluster_merge"}),
+            True,
+            True,
+            ["cross_rank_topk_cluster_merge"],
+        ),
+        (
+            frozenset({"cross_rank_topk", "cross_rank_topk_cluster_merge"}),
+            False,
+            True,
+            ["cross_rank_topk_cluster_merge", "indexer_topk"],
+        ),
+        (frozenset({"cross_rank_topk"}), True, True, ["indexer_topk"]),
+        (frozenset(), True, False, []),
+    ],
+)
+def test_indexer_hook(monkeypatch, features, cluster_merge_result, expected, calls):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        region,
+    )
+
+    rec = _Calls(cluster_merge_result)
+    monkeypatch.setattr(dispatch, "enabled_features", lambda: features)
+    impls = {
+        ("cross_rank_topk_cluster_merge", "cp_indexer_decode"): rec.cross_rank_topk_cluster_merge,
+        ("indexer_topk", "cp_indexer_decode"): rec.indexer_topk,
+    }
+    monkeypatch.setattr(dispatch, "get_impl", lambda fid, name: impls[(fid, name)])
+    got = region.cross_rank_topk_decode(None, None, None, None, 0, 6, None)
+    assert got is expected
+    assert rec.log == calls
+
+
+def test_package_registration_and_cluster_reason():
+    impl = dispatch.get_impl("cross_rank_topk_cluster_merge", "cp_indexer_decode")
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_cluster_merge import (  # noqa: E501
+        cp_workspace,
+        try_cp_indexer_decode,
+    )
+
+    assert impl is try_cp_indexer_decode
+    assert dispatch.get_impl("cross_rank_topk_cluster_merge", "cp_decode_topk") is not None
+    info = [0] * 12
+    info[cp_workspace.INFO_MAX_ROWS] = 8
+    info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] = 16
+    assert cp_workspace.cluster_reason(info) == ""
+    info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] = 7
+    assert "7 co-resident" in cp_workspace.cluster_reason(info)
+    assert "not built" in cp_workspace.cluster_reason(None)
+
+
+def test_world_size_domain():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_cluster_merge import (  # noqa: E501
+        cp_indexer_decode,
+    )
+
+    assert cp_indexer_decode.SUPPORTED_WORLD_SIZES == (4, 8)
+    for world in (4, 8):
+        assert cp_indexer_decode.world_size_reason(world) == ""
+    for world in (1, 2, 3, 5, 6, 7, 16):
+        assert (
+            "not in cross_rank_topk_cluster_merge's verified sizes"
+            in cp_indexer_decode.world_size_reason(world)
+        )
+
+
+def test_op_info_when_built():
+    import torch
+
+    import tensorrt_llm  # noqa: F401  (the trtllm op namespace)
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_cluster_merge import (  # noqa: E501
+        cp_workspace,
+    )
+
+    if not torch.cuda.is_available():
+        pytest.skip("no GPU")
+    library = os.environ.get("GLM_MOE_DSA_CROSS_RANK_TOPK_CLUSTER_MERGE_JIT_LIB", "")
+    if not hasattr(
+        torch.ops.trtllm, "glm_moe_dsa_cross_rank_topk_cluster_merge_workspace_info"
+    ) and os.path.exists(library):
+        # A JIT build of the op on an image without it.
+        torch.ops.load_library(library)
+    info = cp_workspace.op_info(8)
+    if info is None:
+        pytest.skip(
+            "glm_moe_dsa_cross_rank_topk_cluster_merge is not built into this image (no development build given)"
+        )
+    assert info[cp_workspace.INFO_STATE_WORDS] == cp_workspace.EPOCH_BASE + 64
+    assert info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= info[cp_workspace.INFO_MAX_ROWS]
+    # Cached per group size and device; callers get their own copy.
+    again = cp_workspace.op_info(8)
+    assert again == info and again is not info
+    info4 = cp_workspace.op_info(4)
+    assert info4[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= info4[cp_workspace.INFO_MAX_ROWS]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py
new file mode 100644
index 000000000000..aa01442ce1b6
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py
@@ -0,0 +1,535 @@
+# 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.
+"""Selector, call sites and bridge plan of cross_rank_topk_global_emit.
+
+- ``cross_rank_topk_global_emit`` runs only with ``rope_kv_append``, ``cross_rank_topk``
+  and ``cross_rank_topk_cluster_merge`` and is recorded as a ``rank-consistent`` variant; disabling it
+  keeps the paths it extends.
+- The cross_rank_topk region of the indexer's decode top-k (``region.cross_rank_topk_decode``)
+  tries cross_rank_topk_global_emit first (with the call's rope_kv_append packet, which
+  ``GlmIndexer._fill_decode_topk`` passes) when ``cross_rank_topk_global_emit`` is selected,
+  then cross_rank_topk_cluster_merge, then cross_rank_topk; without ``cross_rank_topk_global_emit``
+  it is B5's hook.
+- The backend's packet and remap (``GlmDsaBackend.sparse_attn_predict``) come from
+  cross_rank_topk_global_emit when it is selected.
+- ``bridge.predict_copy`` / ``predict_remap`` reproduce the call sites' rules
+  (rope_kv_append's stash, the backend's shared copy, rope_kv_append's grouping), and
+  ``remap_topk_to_global`` returns the planned pool rows, falls back to rope_kv_append's
+  remap without a plan, and raises on any difference from the plan.
+- The package registers its entry points; group sizes and op info.
+- deferred_input_prep's tables list cross_rank_topk_global_emit's package directories and its side
+  selector.
+"""
+
+import inspect
+import os
+import re
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+
+def _plan(disabled=()):
+    environ = {dispatch.DISABLE_ENV: ",".join(disabled)} if disabled else {}
+    return dispatch.resolve(dispatch.ModelingV2Mode.AUTO, environ)
+
+
+def test_selector_closure():
+    assert dispatch.REQUIRED_FEATURES["cross_rank_topk_global_emit"] == {
+        "rope_kv_append",
+        "cross_rank_topk",
+        "cross_rank_topk_cluster_merge",
+    }
+    assert _plan().enabled("cross_rank_topk_global_emit")
+    for base in ("rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"):
+        assert not _plan([base]).enabled("cross_rank_topk_global_emit"), base
+    record = dispatch.feature_record(
+        [
+            "rope_kv_append",
+            "cross_rank_topk",
+            "cross_rank_topk_cluster_merge",
+            "cross_rank_topk_global_emit",
+        ],
+        environ={},
+    )
+    assert (
+        "cross_rank_topk_global_emit[rank-consistent]" in record
+        and "cross_rank_topk_cluster_merge[rank-consistent]" in record
+    )
+
+
+def test_disabling_global_emit_keeps_cluster_merge():
+    plan = _plan(["cross_rank_topk_global_emit"])
+    assert not plan.enabled("cross_rank_topk_global_emit")
+    for base in ("rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"):
+        assert plan.enabled(base), base
+
+
+class _Calls:
+    def __init__(self, global_emit_result, cluster_merge_result):
+        self.global_emit_result = global_emit_result
+        self.cluster_merge_result = cluster_merge_result
+        self.log = []
+        self.packets = []
+
+    def cross_rank_topk_global_emit(self, *args):
+        self.log.append("cross_rank_topk_global_emit")
+        self.packets.append(args[-1])
+        assert len(args) == 8
+        return self.global_emit_result
+
+    def cross_rank_topk_cluster_merge(self, *args):
+        self.log.append("cross_rank_topk_cluster_merge")
+        assert len(args) == 7
+        return self.cluster_merge_result
+
+    def indexer_topk(self, *args):
+        self.log.append("indexer_topk")
+        return True
+
+
+@pytest.mark.parametrize(
+    "features, global_emit_result, cluster_merge_result, calls",
+    [
+        (
+            frozenset(
+                {
+                    "rope_kv_append",
+                    "cross_rank_topk",
+                    "cross_rank_topk_cluster_merge",
+                    "cross_rank_topk_global_emit",
+                }
+            ),
+            True,
+            True,
+            ["cross_rank_topk_global_emit"],
+        ),
+        (
+            frozenset(
+                {
+                    "rope_kv_append",
+                    "cross_rank_topk",
+                    "cross_rank_topk_cluster_merge",
+                    "cross_rank_topk_global_emit",
+                }
+            ),
+            False,
+            True,
+            ["cross_rank_topk_global_emit", "cross_rank_topk_cluster_merge"],
+        ),
+        (
+            frozenset(
+                {
+                    "rope_kv_append",
+                    "cross_rank_topk",
+                    "cross_rank_topk_cluster_merge",
+                    "cross_rank_topk_global_emit",
+                }
+            ),
+            False,
+            False,
+            ["cross_rank_topk_global_emit", "cross_rank_topk_cluster_merge", "indexer_topk"],
+        ),
+        (
+            frozenset({"rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"}),
+            True,
+            True,
+            ["cross_rank_topk_cluster_merge"],
+        ),
+        (frozenset({"rope_kv_append", "cross_rank_topk"}), True, True, ["indexer_topk"]),
+        (frozenset(), True, True, []),
+    ],
+)
+def test_indexer_hook(monkeypatch, features, global_emit_result, cluster_merge_result, calls):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        region,
+    )
+
+    rec = _Calls(global_emit_result, cluster_merge_result)
+    monkeypatch.setattr(dispatch, "enabled_features", lambda: features)
+    impls = {
+        ("cross_rank_topk_global_emit", "cp_indexer_decode"): rec.cross_rank_topk_global_emit,
+        ("cross_rank_topk_cluster_merge", "cp_indexer_decode"): rec.cross_rank_topk_cluster_merge,
+        ("indexer_topk", "cp_indexer_decode"): rec.indexer_topk,
+    }
+    monkeypatch.setattr(dispatch, "get_impl", lambda fid, name: impls[(fid, name)])
+    packet = object()
+    got = region.cross_rank_topk_decode(None, None, None, None, 0, 6, None, packet)
+    # Each tried region but the last declined; the last one served the call.
+    assert got is bool(calls)
+    assert rec.log == calls
+    assert all(p is packet for p in rec.packets)
+
+
+def test_indexer_call_site_passes_the_packet():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+        GlmIndexer,
+    )
+
+    src = inspect.getsource(GlmIndexer._fill_decode_topk)
+    call = src[src.index("region.cross_rank_topk_decode(") :]
+    call = call[: call.index(")\n")]
+    assert "self._topk_bridge" in call
+
+
+def test_backend_uses_global_emit_packet_and_remap_only_when_selected():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+        GlmDsaBackend,
+    )
+
+    src = inspect.getsource(GlmDsaBackend.sparse_attn_predict)
+    # The formatter wraps long statements: compare with the whitespace collapsed.
+    src = re.sub(r"(?<=[(\[])\s+|\s+(?=[)\]])", "", " ".join(src.split()))
+    assert (
+        'bridge_path = ("cross_rank_topk_global_emit" if '
+        'dispatch.is_enabled("cross_rank_topk_global_emit") else "rope_kv_append")' in src
+    )
+    assert 'get_impl(bridge_path, "new_bridge_packet")' in src
+    assert 'bridge_path, "remap_topk_to_global"' in src
+    # Every other rope_kv_append call site of the backend is rope_kv_append's own.
+    assert src.count('get_impl("rope_kv_append", "defer_shared_copy")') == 1
+
+
+def test_package_registration():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
+        cross_rank_topk_global_emit,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+        cp_indexer_decode,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        ops as rope_ops,
+    )
+
+    assert (
+        dispatch.get_impl("cross_rank_topk_global_emit", "cp_indexer_decode")
+        is cp_indexer_decode.try_cp_indexer_decode
+    )
+    assert (
+        dispatch.get_impl("cross_rank_topk_global_emit", "new_bridge_packet")
+        is bridge.new_bridge_packet
+    )
+    assert (
+        dispatch.get_impl("cross_rank_topk_global_emit", "remap_topk_to_global")
+        is bridge.remap_topk_to_global
+    )
+    assert issubclass(cross_rank_topk_global_emit.GlobalEmitBridgePacket, rope_ops.BridgePacket)
+    assert cross_rank_topk_global_emit.GlobalEmitBridgePacket().plan is None
+
+
+def test_world_size_domain():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        cp_indexer_decode,
+    )
+
+    assert cp_indexer_decode.SUPPORTED_WORLD_SIZES == (4, 8)
+    for world in (4, 8):
+        assert cp_indexer_decode.world_size_reason(world) == ""
+    for world in (1, 2, 3, 5, 6, 7, 16):
+        assert (
+            "not in cross_rank_topk_global_emit's verified sizes"
+            in cp_indexer_decode.world_size_reason(world)
+        )
+
+
+# ---- Bridge plan ------------------------------------------------------------------------------------------------
+
+
+class _Kvm:
+    def __init__(self, tokens_per_block=64):
+        self.tokens_per_block = tokens_per_block
+        self.layer_offsets = {layer: layer for layer in range(100)}
+
+    def get_primary_pool_page_index_params(self, local_layer):
+        return 79, local_layer
+
+
+def _struct(group_ids):
+    """Group struct of rope_kv_append's grouping: layers 0, 1 singletons, then groups of len(group_ids)."""
+    n = 2 + 2 * len(group_ids)
+    leader_of = [0, 1] + [2] * len(group_ids) + [2 + len(group_ids)] * len(group_ids)
+    slot_of = [0, 0] + list(range(len(group_ids))) * 2
+    return {
+        "leader_of": leader_of[:n],
+        "slot_of": slot_of[:n],
+        "group_active": {0: False, 1: False, 2: True, 2 + len(group_ids): True},
+        "group_scale": {2: 79, 2 + len(group_ids): 79},
+        "group_layer_ids": {
+            2: torch.tensor(group_ids, dtype=torch.int32),
+            2 + len(group_ids): torch.tensor(group_ids, dtype=torch.int32),
+        },
+    }
+
+
+def _metadata(*, draft=False, shared=True, accepted=torch.int32, group=4, rows=6, struct=None):
+    st = _struct(list(range(group))) if struct is None else struct
+    return SimpleNamespace(
+        num_contexts=0,
+        num_generations=1,
+        num_ctx_tokens=0,
+        num_tokens=rows,
+        in_mtp_draft_loop=draft,
+        shared_topk_indices=torch.zeros(8, 2048, dtype=torch.int32) if shared else None,
+        mtp_num_accepted=None if accepted is None else torch.tensor([3, 1], dtype=accepted),
+        kv_cache_manager=_Kvm(),
+        _ensure_group_remap_struct=lambda: st,
+        _group_remap_batched={},
+    )
+
+
+def _indexer(share=True, layer=2):
+    return SimpleNamespace(
+        mtp_index_share=share, layer_idx=layer, top_k=SimpleNamespace(top_k=2048)
+    )
+
+
+def test_predict_copy_rules():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+    )
+
+    # Target layer: the backend's shared copy of the generation rows.
+    md = _metadata()
+    c = bridge.predict_copy(_indexer(), md, 6)
+    assert (c.dst is md.shared_topk_indices, c.row_start, c.rows, c.mode) == (True, 0, 6, 0)
+    assert bridge.predict_copy(_indexer(), _metadata(shared=False), 6) is None
+    # MTP draft step 0 with index sharing: rope_kv_append's stash (mode 1, the request's accepted
+    # count).
+    md = _metadata(draft=True)
+    c = bridge.predict_copy(_indexer(), md, 6)
+    assert (c.dst is md.shared_topk_indices, c.row_start, c.rows, c.mode, c.next_n) == (
+        True,
+        0,
+        1,
+        1,
+        6,
+    )
+    assert c.num_accepted.dtype == torch.int32 and c.num_accepted.tolist() == [3]
+    c = bridge.predict_copy(_indexer(), _metadata(draft=True, accepted=None), 6)
+    assert c.mode == 1 and c.num_accepted is None
+    c = bridge.predict_copy(_indexer(), _metadata(draft=True, accepted=torch.int64), 6)
+    assert c.num_accepted.dtype == torch.int64
+    # The stash is not deferred (rope_kv_append declines): the native stash runs and the backend
+    # keeps the rows.
+    assert bridge.predict_copy(_indexer(), _metadata(draft=True, shared=False), 6) is None
+    assert bridge.predict_copy(_indexer(), _metadata(draft=True, accepted=torch.float32), 6) is None
+    # A draft step without index sharing takes the backend's shared copy.
+    c = bridge.predict_copy(_indexer(share=False), _metadata(draft=True), 6)
+    assert c.mode == 0 and c.rows == 6
+
+
+def test_predict_remap_rules(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+    )
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
+
+    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
+    md = _metadata(group=4)
+    r = bridge.predict_remap(md, 2)
+    assert (r.grouped, r.leader, r.slot, r.page_index_scale) == (True, 2, 0, 79)
+    assert r.layer_ids is md._ensure_group_remap_struct()["group_layer_ids"][2]
+    r = bridge.predict_remap(md, 1)
+    assert (r.grouped, r.page_index_scale, r.layer_offset) == (False, 79, 1)
+    assert bridge.predict_remap(_metadata(group=9), 2) is None
+    r = bridge.predict_remap(_metadata(draft=True), 2)
+    assert not r.grouped and r.layer_offset == 2
+    r = bridge.predict_remap(_metadata(struct={}), 2)
+    assert not r.grouped
+    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", False)
+    assert not bridge.predict_remap(_metadata(), 2).grouped
+
+
+def _backend(layer=2, indexer=True):
+    return SimpleNamespace(
+        indexer=object() if indexer else None, get_local_layer_idx=lambda metadata: layer
+    )
+
+
+def _planned(md, layer=2, grouped=True):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+    )
+
+    ids = torch.zeros(6, 2048, dtype=torch.int32)
+    remap = bridge.predict_remap(md, layer)
+    copy = bridge.predict_copy(_indexer(layer=layer), md, 6)
+    pool = (
+        torch.arange(4 * 6 * 2048, dtype=torch.int32).view(4, 6, 2048)
+        if remap.grouped
+        else torch.ones(6, 2048, dtype=torch.int32)
+    )
+    packet = bridge.GlobalEmitBridgePacket()
+    packet.plan = bridge.GlobalEmitPlan(ids=ids, remap=remap, copy=copy, pool=pool)
+    return packet, ids, pool, copy
+
+
+def test_remap_uses_the_plan_once(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        ops as rope_ops,
+    )
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
+
+    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
+    rope_append_calls = []
+    monkeypatch.setattr(
+        rope_ops, "remap_topk_to_global", lambda *a: rope_append_calls.append(a) or "rope_kv_append"
+    )
+    md = _metadata()
+    packet, ids, pool, copy = _planned(md, layer=3)
+    rope_ops.defer_shared_copy(packet, md.shared_topk_indices, 0, ids)
+    out = bridge.remap_topk_to_global(_backend(layer=3), ids, md, packet)
+    assert out.data_ptr() == pool[1].data_ptr() and md._group_remap_batched[2] is pool
+    assert packet.plan is None and not rope_append_calls
+    # A packet without a plan (draft reuse, declines) is rope_kv_append's remap.
+    assert (
+        bridge.remap_topk_to_global(_backend(layer=3), ids, md, packet) == "rope_kv_append"
+        and len(rope_append_calls) == 1
+    )
+    # Per-layer plan.
+    packet, ids, pool, copy = _planned(md, layer=1)
+    rope_ops.defer_shared_copy(packet, md.shared_topk_indices, 0, ids)
+    assert bridge.remap_topk_to_global(_backend(layer=1), ids, md, packet) is pool
+
+
+def test_remap_rejects_differences(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        ops as rope_ops,
+    )
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
+
+    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
+
+    def check(mutate, match):
+        md = _metadata()
+        packet, ids, pool, copy = _planned(md, layer=2)
+        rope_ops.defer_shared_copy(packet, md.shared_topk_indices, 0, ids)
+        args = mutate(md, packet, ids)
+        with pytest.raises(bridge.GlobalEmitContractError, match=match):
+            bridge.remap_topk_to_global(*args)
+
+    check(lambda md, p, ids: (_backend(), ids, md, setattr(p, "copy", None) or p), "copy")
+    check(lambda md, p, ids: (_backend(), ids, md, setattr(p.copy, "rows", 5) or p), "copy")
+    check(lambda md, p, ids: (_backend(), ids, md, setattr(p, "source", ids) or p), "source")
+    check(
+        lambda md, p, ids: (_backend(), ids.clone(), md, p),
+        "rows cross_rank_topk_global_emit wrote",
+    )
+    check(lambda md, p, ids: (_backend(layer=3), ids, md, p), "local layer")
+    check(lambda md, p, ids: (_backend(indexer=False), ids, md, p), "no indexer")
+
+    def regroup(md, p, ids):
+        md._ensure_group_remap_struct = lambda: {}
+        return _backend(), ids, md, p
+
+    check(regroup, "remap decision")
+
+
+def test_copy_reason_matches_the_bridge_domain():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        ops as rope_ops,
+    )
+
+    if not torch.cuda.is_available():
+        pytest.skip("no GPU")
+    dst = torch.zeros(8, 2048, dtype=torch.int32, device="cuda")
+    assert bridge.copy_reason(None, 6, 2048) == ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 6), 6, 2048) == ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 3, 6), 6, 2048) != ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 7), 6, 2048) != ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst.cpu(), 0, 6), 6, 2048) != ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst[:, :100], 0, 6), 6, 2048) != ""
+    acc = torch.tensor([2], dtype=torch.int32, device="cuda")
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 1, 1, 6, acc), 6, 2048) == ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 2, 1, 6, None), 6, 2048) != ""
+    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 1, 1, 6, acc.float()), 6, 2048) != ""
+
+
+def test_bridge_reason_declines_before_device_work(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge,
+        cp_indexer_decode,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        ops as rope_ops,
+    )
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
+
+    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
+    reason, _, _ = cp_indexer_decode.bridge_reason(
+        _indexer(), _metadata(), rope_ops.BridgePacket(), 6
+    )
+    assert "no cross_rank_topk_global_emit bridge packet" in reason
+    reason, _, _ = cp_indexer_decode.bridge_reason(_indexer(), _metadata(), None, 6)
+    assert "no cross_rank_topk_global_emit bridge packet" in reason
+    md = _metadata()
+    md.kv_cache_manager.tokens_per_block = 32
+    reason, _, _ = cp_indexer_decode.bridge_reason(
+        _indexer(), md, bridge.GlobalEmitBridgePacket(), 6
+    )
+    assert "block size" in reason
+    reason, _, _ = cp_indexer_decode.bridge_reason(
+        _indexer(), _metadata(group=9), bridge.GlobalEmitBridgePacket(), 6
+    )
+    assert "group" in reason
+    used = bridge.GlobalEmitBridgePacket()
+    used.copy = rope_ops.BridgeCopy(torch.zeros(1, 1), 0, 1)
+    reason, _, _ = cp_indexer_decode.bridge_reason(_indexer(), _metadata(), used, 6)
+    assert "already holds" in reason
+
+
+def test_op_info_when_built():
+    import tensorrt_llm  # noqa: F401  (the trtllm op namespace)
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        cp_workspace,
+    )
+
+    if not torch.cuda.is_available():
+        pytest.skip("no GPU")
+    library = os.environ.get("GLM_MOE_DSA_CROSS_RANK_TOPK_GLOBAL_EMIT_JIT_LIB", "")
+    if not hasattr(
+        torch.ops.trtllm, "glm_moe_dsa_cross_rank_topk_global_emit_workspace_info"
+    ) and os.path.exists(library):
+        # A JIT build of the op on an image without it.
+        torch.ops.load_library(library)
+    info = cp_workspace.op_info(8)
+    if info is None:
+        pytest.skip(
+            "glm_moe_dsa_cross_rank_topk_global_emit is not built into this image (no development build given)"
+        )
+    assert info[cp_workspace.INFO_STATE_WORDS] == cp_workspace.EPOCH_BASE + 64
+    assert info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= info[cp_workspace.INFO_MAX_ROWS]
+    assert info[cp_workspace.INFO_MAX_GROUP_SIZE] == 8
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk.cross_rank_logits_runner import (  # noqa: E501
+        local_capacity,
+    )
+
+    for world in (4, 8):
+        winfo = cp_workspace.op_info(world)
+        assert winfo[cp_workspace.INFO_MAX_LOCAL_CAPACITY] >= local_capacity((1 << 20) + 14, world)
+        assert winfo[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= winfo[cp_workspace.INFO_MAX_ROWS]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py
new file mode 100644
index 000000000000..14fc28efef99
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py
@@ -0,0 +1,387 @@
+# 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.
+"""cross_rank_topk_global_emit and indexer_layer_front composed in the cross_rank_topk CP indexer
+region (host-side checks).
+
+cross_rank_topk_global_emit forks the region's top-k (its emission also writes rope_kv_append's
+bridge outputs) and indexer_layer_front forks the region's CP logits launch (entry wait, PDL
+launch). With both selected, cross_rank_topk_global_emit serves the region and takes
+indexer_layer_front's logits launcher; the indexer hook tries cross_rank_topk_global_emit, then
+indexer_layer_front (cross_rank_topk_cluster_merge's or cross_rank_topk's top-k), then
+cross_rank_topk_cluster_merge, then cross_rank_topk. These tests mock every device step: they check
+which entry serves a call, which logits launcher cross_rank_topk_global_emit uses, and the
+capability lines; the numerics are covered by the GPU tests of each package.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+
+class _Calls:
+    def __init__(self, results):
+        self.results = results
+        self.log = []
+
+    def make(self, name, nargs):
+        def entry(*args):
+            self.log.append(name)
+            assert len(args) == nargs, (name, len(args))
+            return self.results[name]
+
+        return entry
+
+
+_ALL = {
+    "cross_rank_topk_global_emit": True,
+    "indexer_layer_front": True,
+    "cross_rank_topk_cluster_merge": True,
+    "indexer_topk": True,
+}
+
+
+@pytest.mark.parametrize(
+    "features, results, calls",
+    [
+        # F4: cross_rank_topk_global_emit serves (with indexer_layer_front's logits).
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "cross_rank_topk_global_emit",
+                "indexer_layer_front",
+            },
+            _ALL,
+            ["cross_rank_topk_global_emit"],
+        ),
+        # cross_rank_topk_global_emit declines: indexer_layer_front's region
+        # (cross_rank_topk_cluster_merge's top-k) serves.
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "cross_rank_topk_global_emit",
+                "indexer_layer_front",
+            },
+            {**_ALL, "cross_rank_topk_global_emit": False},
+            ["cross_rank_topk_global_emit", "indexer_layer_front"],
+        ),
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "cross_rank_topk_global_emit",
+                "indexer_layer_front",
+            },
+            {**_ALL, "cross_rank_topk_global_emit": False, "indexer_layer_front": False},
+            ["cross_rank_topk_global_emit", "indexer_layer_front", "cross_rank_topk_cluster_merge"],
+        ),
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "cross_rank_topk_global_emit",
+                "indexer_layer_front",
+            },
+            {
+                "cross_rank_topk_global_emit": False,
+                "indexer_layer_front": False,
+                "cross_rank_topk_cluster_merge": False,
+                "indexer_topk": True,
+            },
+            [
+                "cross_rank_topk_global_emit",
+                "indexer_layer_front",
+                "cross_rank_topk_cluster_merge",
+                "indexer_topk",
+            ],
+        ),
+        # F4 - cross_rank_topk_global_emit: indexer_layer_front first.
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "indexer_layer_front",
+            },
+            _ALL,
+            ["indexer_layer_front"],
+        ),
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "indexer_layer_front",
+            },
+            {**_ALL, "indexer_layer_front": False},
+            ["indexer_layer_front", "cross_rank_topk_cluster_merge"],
+        ),
+        # Without indexer_layer_front: cross_rank_topk_global_emit first, then
+        # cross_rank_topk_cluster_merge.
+        (
+            {
+                "rope_kv_append",
+                "cross_rank_topk",
+                "cross_rank_topk_cluster_merge",
+                "cross_rank_topk_global_emit",
+            },
+            {**_ALL, "cross_rank_topk_global_emit": False},
+            ["cross_rank_topk_global_emit", "cross_rank_topk_cluster_merge"],
+        ),
+        # cross_rank_topk_cluster_merge only.
+        (
+            {"rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"},
+            _ALL,
+            ["cross_rank_topk_cluster_merge"],
+        ),
+        (set(), _ALL, []),
+    ],
+)
+def test_indexer_hook_order(monkeypatch, features, results, calls):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        region,
+    )
+
+    rec = _Calls(results)
+    impls = {
+        ("cross_rank_topk_global_emit", "cp_indexer_decode"): rec.make(
+            "cross_rank_topk_global_emit", 8
+        ),
+        ("indexer_layer_front", "cp_indexer_decode"): rec.make("indexer_layer_front", 7),
+        ("cross_rank_topk_cluster_merge", "cp_indexer_decode"): rec.make(
+            "cross_rank_topk_cluster_merge", 7
+        ),
+        ("indexer_topk", "cp_indexer_decode"): rec.make("indexer_topk", 7),
+    }
+    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(features))
+    monkeypatch.setattr(dispatch, "get_impl", lambda fid, name: impls[(fid, name)])
+    got = region.cross_rank_topk_decode(None, None, None, None, 0, 6, None, object())
+    assert rec.log == calls
+    assert got is bool(calls) and (not calls or results[calls[-1]])
+
+
+def _region_inputs(num_heads=64, head_dim=128, next_n=6):
+    top_k = SimpleNamespace(
+        needs_gvr_prior=False,
+        decode_implementation=SimpleNamespace(value="default"),
+        gvr_self_sampling=False,
+        top_k=2048,
+        compress_ratio=1,
+    )
+    indexer = SimpleNamespace(
+        top_k=top_k,
+        use_cute_dsl_paged_mqa_logits=False,
+        use_fp4=False,
+        use_gvr_emission=False,
+        n_heads=num_heads,
+        head_dim=head_dim,
+        layer_idx=3,
+    )
+    kvm = SimpleNamespace(
+        get_indexer_k_cache_buffers=lambda layer: torch.zeros(
+            4, 64, 1, head_dim + 4, dtype=torch.uint8
+        ),
+        tokens_per_block=64,
+        layer_offsets={3: 3},
+    )
+    metadata = SimpleNamespace(
+        num_contexts=0,
+        num_generations=1,
+        seq_lens=torch.tensor([next_n], dtype=torch.int32),
+        kv_cache_manager=kvm,
+        indexer_k_cache_block_offsets=torch.zeros(1, 16, dtype=torch.int32),
+        mapping=SimpleNamespace(tp_size=8, enable_attention_dp=False),
+        get_indexer_max_seq_len=lambda: 32768,
+        is_ragged_verify=False,
+        use_expanded_buffers_for_mtp=False,
+        in_mtp_draft_loop=False,
+        gen_indexer_kv_lens_cuda_runtime=torch.tensor([32768], dtype=torch.int32),
+        req_idx_per_token=torch.zeros(next_n, dtype=torch.int32),
+        num_ctx_tokens=0,
+        num_tokens=next_n,
+        block_table=torch.zeros(1, 512, dtype=torch.int32),
+        num_seqs=1,
+    )
+    q_fp8 = torch.zeros(next_n, num_heads, head_dim)
+    weights = torch.zeros(next_n, num_heads)
+    out = torch.zeros(next_n, 2048, dtype=torch.int32)
+    return indexer, metadata, q_fp8, weights, out
+
+
+@pytest.mark.parametrize("with_indexer_front", [True, False])
+def test_global_emit_takes_the_logits_launcher_of_indexer_front(monkeypatch, with_indexer_front):
+    """cross_rank_topk_global_emit's region launches indexer_layer_front's logits kernel exactly
+    when indexer_layer_front is selected."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        bridge as global_emit_bridge,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        cp_indexer_decode as global_emit_decode,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        cp_workspace as global_emit_ws,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
+        cp_logits_runner as indexer_front_runner,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        cross_rank_logits_runner as indexer_topk_runner,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        cross_rank_policy,
+    )
+
+    features = {
+        "absorbed_query",
+        "indexer_topk",
+        "draft_absorbed_query",
+        "rope_kv_append",
+        "cross_rank_topk",
+        "cross_rank_topk_cluster_merge",
+        "cross_rank_topk_global_emit",
+    } | ({"indexer_layer_front"} if with_indexer_front else set())
+    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(features))
+    launched = []
+
+    def runner(owner):
+        def ensure_compiled(*args, **kwargs):
+            launched.append((owner, "compile"))
+
+        def paged_mqa_logits_cp(
+            q, k_cache, weights, kv_lens, block_table, max_seq_len, rank, world
+        ):
+            launched.append((owner, "logits"))
+            assert (rank, world) == (2, 8)
+            return torch.zeros(q.shape[1], 16)
+
+        return ensure_compiled, paged_mqa_logits_cp
+
+    for owner, mod in (
+        ("indexer_layer_front", indexer_front_runner),
+        ("indexer_topk", indexer_topk_runner),
+    ):
+        ensure, logits = runner(owner)
+        monkeypatch.setattr(mod, "ensure_compiled", ensure)
+        monkeypatch.setattr(mod, "paged_mqa_logits_cp", logits)
+    monkeypatch.setattr(
+        cross_rank_policy,
+        "decide_cp_uniform",
+        lambda **kw: SimpleNamespace(eligible=True, reason=""),
+    )
+    monkeypatch.setattr(cross_rank_policy, "rank_local_reason", lambda **kw: "")
+    monkeypatch.setattr(global_emit_ws, "op_info", lambda world: (0,) * 16)
+    monkeypatch.setattr(global_emit_ws, "cluster_reason", lambda info: "")
+    workspace = SimpleNamespace(rank=2, world=8, state=None, mc_ptr=0, uc_ptr=0)
+
+    def get_workspace(mapping, prepare=None, local_reason=""):
+        prepare()
+        return workspace
+
+    monkeypatch.setattr(global_emit_ws, "get_workspace", get_workspace)
+    remap = SimpleNamespace(page_index_scale=1, grouped=False, layer_ids=None, layer_offset=3)
+    monkeypatch.setattr(global_emit_decode, "bridge_reason", lambda *a: ("", remap, None))
+    topk_calls = []
+
+    def topk(local_logits, kv_lens, out, ws, next_n, **kw):
+        topk_calls.append((tuple(local_logits.shape), next_n, kw["layer_id"]))
+        return torch.zeros(next_n, 2048, dtype=torch.int32)
+
+    monkeypatch.setattr(global_emit_decode, "cp_topk_from_local_logits", topk)
+    logged = {}
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", logged)
+
+    indexer, metadata, q_fp8, weights, out = _region_inputs()
+    packet = SimpleNamespace(plan=None)
+    assert global_emit_decode.try_cp_indexer_decode(
+        indexer, metadata, q_fp8, weights, 0, 6, out, packet
+    )
+    owner = "indexer_layer_front" if with_indexer_front else "indexer_topk"
+    assert launched == [(owner, "compile"), (owner, "logits")]
+    assert topk_calls == [((6, 16), 6, 3)]
+    assert isinstance(packet.plan, global_emit_bridge.GlobalEmitPlan) and packet.plan.ids is out
+    assert logged.get(("cross_rank_topk_global_emit", "cp_indexer_decode")) is True
+    assert (
+        ("indexer_layer_front", "cp_indexer_decode_global_emit") in logged
+    ) is with_indexer_front
+    if with_indexer_front:
+        assert logged[("indexer_layer_front", "cp_indexer_decode_global_emit")] is True
+
+
+def test_global_emit_decline_is_logged_for_both(monkeypatch):
+    """A cross_rank_topk_global_emit decline with indexer_layer_front selected is visible on both capability keys."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        cp_indexer_decode as global_emit_decode,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
+        cp_workspace as global_emit_ws,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        cross_rank_policy,
+    )
+
+    features = {
+        "absorbed_query",
+        "indexer_topk",
+        "draft_absorbed_query",
+        "rope_kv_append",
+        "cross_rank_topk",
+        "cross_rank_topk_cluster_merge",
+        "cross_rank_topk_global_emit",
+        "indexer_layer_front",
+    }
+    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(features))
+    monkeypatch.setattr(
+        cross_rank_policy,
+        "decide_cp_uniform",
+        lambda **kw: SimpleNamespace(eligible=False, reason="why"),
+    )
+    monkeypatch.setattr(cross_rank_policy, "rank_local_reason", lambda **kw: "")
+    monkeypatch.setattr(global_emit_ws, "op_info", lambda world: (0,) * 16)
+    monkeypatch.setattr(global_emit_ws, "cluster_reason", lambda info: "")
+    logged = {}
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", logged)
+    indexer, metadata, q_fp8, weights, out = _region_inputs()
+    packet = SimpleNamespace(plan=None)
+    assert not global_emit_decode.try_cp_indexer_decode(
+        indexer, metadata, q_fp8, weights, 0, 6, out, packet
+    )
+    assert packet.plan is None
+    assert logged == {
+        ("cross_rank_topk_global_emit", "cp_indexer_decode"): False,
+        ("indexer_layer_front", "cp_indexer_decode_global_emit"): False,
+    }
+
+
+def test_mla_rope_line():
+    """The rope_kv_append rope branch logs indexer_layer_front/rope and picks indexer_layer_front's
+    rope only when its op is loaded."""
+    import inspect
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+        attention_forward,
+    )
+
+    src = inspect.getsource(attention_forward.forward_absorption_generation)
+    assert '"indexer_layer_front", "rope_available"' in src
+    assert '"indexer_layer_front",\n' in src and '"rope",' in src
+    assert 'rope_owner = "rope_kv_append"' in src
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py
new file mode 100644
index 000000000000..45eb2446a16a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py
@@ -0,0 +1,355 @@
+# 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.
+"""Native oracle of the decode_prologue target prologue (shared by the tests and the benchmark).
+
+The native reference is the literal ``DecoderRunner._preprocess_inputs``
+(called unbound on a runner stub that holds the overlap-scheduler offset
+buffers, the configuration and the mapping) followed by the target model's
+native ``Embedding``. It runs on a real CUDA-graph ``DSAtrtllmAttentionMetadata``
+prepared from a real ``DSACacheManager`` at the served GEN DSA
+configuration (draft_loop's fixtures), exactly as the decoder runner prepares a decode
+step: ``create_cuda_graph_metadata``, ``update_spec_dec_param``,
+``KVCacheParams``, ``prepare()``.
+"""
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Dict, List, Optional, Sequence
+
+import torch
+
+_DRAFT_LOOP_TESTS = Path(__file__).resolve().parent.parent / "draft_loop"
+sys.path.insert(0, str(_DRAFT_LOOP_TESTS))
+
+from _draft_loop_fixtures import (  # noqa: E402
+    MTP_K,
+    add_gen_requests,
+    load_checkpoint_tensor,
+    make_dsa_cache_manager,
+)
+
+VOCAB = 154880
+HIDDEN = 6144
+# Served GEN: max_num_tokens 128, max_batch_size 1 (2 for the padded case).
+MAX_NUM_TOKENS = 128
+POISON = 0xA5
+PACKET_ATTR = "_decode_prologue_embedding"
+
+# Every metadata buffer the native prologue writes; poisoned identically in
+# both arms before each run and compared in full afterwards.
+DEST_FIELDS = (
+    "kv_lens_cuda_2d",
+    "scheduler_metadata_buffer",
+    "scheduler_metadata_buffer_full_next_n",
+    "kv_lens_expanded_cuda",
+    "scheduler_metadata_buffer_expanded",
+    "req_idx_per_token",
+    "slot_mapping_fp8",
+    "slot_mapping_scale",
+    "gen_kv_indptr",
+    "gen_cached_token_indptr",
+    "topk_indices_buffer",
+)
+
+
+def make_graph_metadata(
+    cache_manager, mapping, sparse_config, request_ids, past_kv_lens, draft_len: int = MTP_K
+):
+    """A prepared CUDA-graph decode metadata (the runner's per-step preparation)."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata
+    from tensorrt_llm._torch.metadata import KVCacheParams
+
+    batch_size = len(request_ids)
+    base = DSAtrtllmAttentionMetadata(
+        seq_lens=torch.full((batch_size,), draft_len + 1, dtype=torch.int),
+        request_ids=list(request_ids),
+        max_num_requests=batch_size,
+        num_contexts=0,
+        prompt_lens=list(past_kv_lens),
+        max_num_tokens=MAX_NUM_TOKENS,
+        kv_cache_manager=cache_manager,
+        mapping=mapping,
+        sparse_attention_config=sparse_config,
+    )
+    md = base.create_cuda_graph_metadata(batch_size, False, draft_len)
+    md.request_ids = list(request_ids)
+    md.prompt_lens = list(past_kv_lens)
+    md.update_spec_dec_param(
+        batch_size=batch_size,
+        is_spec_decoding_enabled=True,
+        is_spec_dec_tree=False,
+        is_spec_dec_dynamic_tree=False,
+        max_draft_len=draft_len,
+        max_total_draft_tokens=draft_len,
+        spec_metadata=None,
+        spec_tree_manager=None,
+        num_contexts=0,
+    )
+    md.kv_cache_params = KVCacheParams(
+        use_cache=True,
+        num_cached_tokens_per_seq=list(past_kv_lens),
+        num_extra_kv_tokens=draft_len - 1,
+    )
+    md.prepare()
+    # Set by the runner's input preparation (no chunked-context rows in a decode step).
+    md.num_chunked_ctx_requests = 0
+    return md, base
+
+
+def make_runner_stub(model, num_tokens: int, batch_size: int):
+    """The attributes ``_preprocess_inputs`` and ``target_prologue`` read from a runner."""
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS
+    from tensorrt_llm.mapping import Mapping
+
+    return SimpleNamespace(
+        _config=SimpleNamespace(disable_overlap_scheduler=False),
+        previous_pos_id_offsets_cuda=torch.zeros(
+            (MAX_NUM_TOKENS,), dtype=torch.int32, device="cuda"
+        ),
+        previous_kv_lens_offsets_cuda=torch.zeros(
+            (max(batch_size, 1),), dtype=torch.int32, device="cuda"
+        ),
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        guided_decoder=None,
+        model=model,
+        runtime_hooks=DEFAULT_RUNTIME_HOOKS,
+    )
+
+
+def make_target_model(weight: torch.Tensor):
+    """A ``DeepseekV3Model`` shell holding only the native replicated ``Embedding``."""
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model
+    from tensorrt_llm._torch.modules.embedding import Embedding
+
+    embed = Embedding(weight.shape[0], weight.shape[1], dtype=weight.dtype).cuda()
+    embed.weight.data = weight
+    target = DeepseekV3Model.__new__(DeepseekV3Model)
+    torch.nn.Module.__init__(target)
+    target.embed_tokens = embed
+    return SimpleNamespace(model=target), embed
+
+
+def embedding_weight(synthetic: bool = False) -> torch.Tensor:
+    """The GLM-5.2 embedding table (synthetic only when asked or unavailable)."""
+    if not synthetic:
+        weight = load_checkpoint_tensor("model.embed_tokens.weight")
+        if weight is not None:
+            assert weight.shape == (VOCAB, HIDDEN) and weight.dtype == torch.bfloat16
+            return weight
+    gen = torch.Generator(device="cuda").manual_seed(12)
+    weight = torch.randn((VOCAB, HIDDEN), dtype=torch.bfloat16, device="cuda", generator=gen)
+    weight[3, :4] = torch.tensor([-0.0, float("inf"), float("-inf"), float("nan")])
+    return weight
+
+
+class Arm:
+    """One arm (native or decode_prologue): metadata, runner stub, static inputs."""
+
+    def __init__(self, name, past_kv_lens, weight, cache_manager, mapping, sparse_config):
+        self.name = name
+        self.batch_size = len(past_kv_lens)
+        ids = list(range(self.batch_size))
+        self.md, self._base = make_graph_metadata(
+            cache_manager, mapping, sparse_config, ids, past_kv_lens
+        )
+        self.num_tokens = self.md.num_tokens
+        self.model, self.embed = make_target_model(weight)
+        self.runner = make_runner_stub(self.model, self.num_tokens, self.batch_size)
+        self.input_ids = torch.zeros((MAX_NUM_TOKENS,), dtype=torch.int32, device="cuda")
+        self.position_ids = torch.zeros((1, MAX_NUM_TOKENS), dtype=torch.int32, device="cuda")
+        self.out = None
+
+    def inputs(self) -> dict:
+        n = self.num_tokens
+        return {
+            "attn_metadata": self.md,
+            "input_ids": self.input_ids[:n],
+            "position_ids": self.position_ids[:, :n],
+        }
+
+    def set_step(self, kv_lens, positions, ids, kv_off, pos_off) -> None:
+        """Install one step's inputs (eager copies outside any graph)."""
+        b, n = self.batch_size, self.num_tokens
+        self.md.kv_lens_cuda[:b].copy_(kv_lens)
+        self.position_ids[0, :n].copy_(positions)
+        self.input_ids[:n].copy_(ids)
+        self.runner.previous_kv_lens_offsets_cuda[:b].copy_(kv_off)
+        self.runner.previous_pos_id_offsets_cuda[:n].copy_(pos_off)
+
+    def poison(self) -> None:
+        for name in DEST_FIELDS:
+            t = getattr(self.md, name, None)
+            if isinstance(t, torch.Tensor):
+                t.view(torch.uint8).fill_(POISON) if t.is_contiguous() else t.fill_(-1)
+
+
+def run_native(
+    arm: Arm, inputs: Optional[dict] = None, enable_spec_decode: bool = True
+) -> torch.Tensor:
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+
+    inputs = arm.inputs() if inputs is None else inputs
+    DecoderRunner._preprocess_inputs(
+        arm.runner, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=MTP_K
+    )
+    return arm.embed(inputs["input_ids"])
+
+
+def run_decode_prologue(
+    arm: Arm, inputs: Optional[dict] = None, enable_spec_decode: bool = True
+) -> torch.Tensor:
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E501
+        prologue,
+    )
+
+    inputs = arm.inputs() if inputs is None else inputs
+    done = prologue.target_prologue(
+        arm.runner, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=MTP_K
+    )
+    assert done, prologue.prologue_ineligibility(arm.runner, inputs, enable_spec_decode)
+    rows = prologue.take_embedding(inputs["attn_metadata"], inputs["input_ids"], arm.embed)
+    assert rows is not None
+    return rows
+
+
+def state(arm: Arm) -> Dict[str, object]:
+    """Every tensor (cloned) and scalar attribute of the metadata, plus the step tensors."""
+    out = {}
+    for name, value in vars(arm.md).items():
+        if isinstance(value, torch.Tensor):
+            out[name] = value.detach().clone()
+        elif isinstance(value, (bool, int, float, str, type(None))):
+            out[name] = value
+    out["@position_ids"] = arm.position_ids.clone()
+    out["@input_ids"] = arm.input_ids.clone()
+    out["@kv_off"] = arm.runner.previous_kv_lens_offsets_cuda.clone()
+    out["@pos_off"] = arm.runner.previous_pos_id_offsets_cuda.clone()
+    return out
+
+
+def _bits(t: torch.Tensor) -> torch.Tensor:
+    if t.dtype in (torch.bfloat16, torch.float16):
+        return t.view(torch.int16)
+    if t.dtype == torch.float32:
+        return t.view(torch.int32)
+    if t.dtype == torch.float64:
+        return t.view(torch.int64)
+    return t
+
+
+def _same(a, b) -> bool:
+    if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
+        if a.shape != b.shape or a.dtype != b.dtype:
+            return False
+        return bool(torch.equal(_bits(a).cpu(), _bits(b).cpu()))
+    return type(a) is type(b) and a == b
+
+
+def changed(before: dict, after: dict) -> set:
+    return {
+        k
+        for k in set(before) | set(after)
+        if not _same(before.get(k, "<missing>"), after.get(k, "<missing>"))
+    }
+
+
+def compare(native_before, native_after, z_before, z_after) -> List[str]:
+    """Names whose bytes (or host values) differ between the arms after the run.
+
+    Compared: every key either arm changed, every destination buffer in full,
+    and the step tensors.
+    """
+    keys = changed(native_before, native_after) | changed(z_before, z_after)
+    keys |= {k for k in DEST_FIELDS if k in native_after}
+    keys |= {"@position_ids", "@input_ids", "@kv_off", "@pos_off", "kv_lens_cuda"}
+    bad = []
+    for k in sorted(keys):
+        if not _same(native_after.get(k, "<missing>"), z_after.get(k, "<missing>")):
+            bad.append(k)
+    return bad
+
+
+def step_inputs(
+    gen: torch.Generator,
+    kv_base: Sequence[int],
+    accepted: Sequence[int],
+    num_tokens: int,
+    draft_len: int = MTP_K,
+    id_edges: bool = False,
+):
+    """One step's (kv_lens, positions, input ids, kv offsets, position offsets).
+
+    ``kv_base[r]`` is request r's KV length after the correction; the eager
+    preparation wrote ``kv_base[r] - offset`` (the overlap scheduler assumed
+    every draft accepted), and the correction adds the offset
+    ``accepted[r] - draft_len`` back. Positions follow the same rule per token.
+    """
+    b = len(kv_base)
+    kv_off = torch.tensor([a - draft_len for a in accepted], dtype=torch.int32)
+    kv_lens = torch.tensor(kv_base, dtype=torch.int32) - kv_off
+    per_req = num_tokens // b
+    pos = []
+    pos_off = []
+    for r in range(b):
+        start = kv_base[r] - per_req
+        for t in range(per_req):
+            pos.append(start + t - (accepted[r] + 1))
+            pos_off.append(accepted[r] + 1)
+    ids = torch.randint(0, VOCAB, (num_tokens,), generator=gen, dtype=torch.int32)
+    if id_edges:
+        edge = torch.tensor([-1, VOCAB, VOCAB + 7, 0, VOCAB - 1, 3], dtype=torch.int32)
+        ids[: min(num_tokens, edge.numel())] = edge[: min(num_tokens, edge.numel())]
+    return (
+        kv_lens.cuda(),
+        torch.tensor(pos, dtype=torch.int32).cuda(),
+        ids.cuda(),
+        kv_off.cuda(),
+        torch.tensor(pos_off, dtype=torch.int32).cuda(),
+    )
+
+
+def make_pair(
+    past_kv_lens, weight, max_kv: Optional[int] = None, sparse_update: Optional[dict] = None
+):
+    """A cache manager and two identically prepared arms on it.
+
+    ``sparse_update`` overrides fields of the arms' DSA configuration (the
+    metadata's parameters), e.g. the paged-MQA logits kernel choice.
+    """
+    batch_size = len(past_kv_lens)
+    cache_manager, mapping, sparse = make_dsa_cache_manager(
+        batch_size, max_kv or max(past_kv_lens) + 4096
+    )
+    if sparse_update:
+        sparse = sparse.model_copy(update=sparse_update)
+    add_gen_requests(cache_manager, list(range(batch_size)), past_kv_lens)
+    native = Arm("native", past_kv_lens, weight, cache_manager, mapping, sparse)
+    decode_prologue = Arm("decode_prologue", past_kv_lens, weight, cache_manager, mapping, sparse)
+    return cache_manager, native, decode_prologue
+
+
+def capture(arm: Arm, fn, debug_dot: Optional[str] = None):
+    """Warm up ``fn(arm)`` eagerly once, then capture it; returns (graph, output)."""
+    fn(arm)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    if debug_dot:
+        graph.enable_debug_mode()
+    with torch.cuda.graph(graph):
+        out = fn(arm)
+    if debug_dot:
+        graph.debug_dump(debug_dot)
+    return graph, out
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py
new file mode 100644
index 000000000000..b4218b33aa75
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py
@@ -0,0 +1,421 @@
+# 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.
+"""decode_prologue target prologue versus the native prologue, byte for byte (TP1 oracle).
+
+Native: the literal ``DecoderRunner._preprocess_inputs`` (two
+``on_update_kv_lens()`` calls around the overlap scheduler's accepted-count
+correction) and the target model's native ``Embedding``. Candidate: decode_prologue's one
+launch. Both run on identically prepared CUDA-graph DSA metadata of a real
+``DSACacheManager`` (served GEN DSA configuration) with the GLM-5.2
+embedding table. Before every run every destination buffer is poisoned
+identically (0xA5) in both arms; afterwards every metadata attribute either
+arm changed, every destination buffer in full, the position ids, the input ids,
+the offsets and the embedding rows must be bitwise equal.
+
+Cases: accepted counts 0-5 (KV offsets -5..0), contexts 2K-1M with lengths
+at k*256-1, k*256 and k*256+1 (DeepGEMM's SPLIT_KV), B in {1, 2}, a padded
+B = 2 graph whose second row is a short dummy, every DSL atom-split factor the
+picker yields (6, 3, 2), the dense top-k (skipped indexer) graph, out-of-range
+input ids, the single-call prologue without the correction (overlap scheduler
+off, or spec decode off for the step), the DeepGEMM paged-MQA path without the
+DSL atom split (factor 1); eager and CUDA graph with changing inputs at every
+replay.
+"""
+
+import os
+import re
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+import _decode_prologue_oracle as orc  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E402, E501
+    prologue,
+)
+
+SYNTHETIC = os.environ.get("GLM_MOE_DSA_DECODE_PROLOGUE_SYNTHETIC_EMBEDDING") == "1"
+
+
+@pytest.fixture(scope="module")
+def weight():
+    return orc.embedding_weight(synthetic=SYNTHETIC)
+
+
+def _boundary(k: int, d: int) -> int:
+    return k * 256 + d
+
+
+# (past KV lengths at prepare, list of per-replay post-correction KV lengths)
+CONTEXT_CASES = {
+    # B = 1: atom split factor 6 (2K), 3 (~12.8K), 2 (>= 32K); SPLIT_KV edges.
+    "b1_2k": ([2100], [[_boundary(9, d)] for d in (-1, 0, 1)]),
+    "b1_12k": ([12700], [[_boundary(50, d)] for d in (-1, 0, 1)]),
+    "b1_32k": ([32768], [[_boundary(128, d)] for d in (-1, 0, 1, 7)]),
+    "b1_131k": ([131072], [[_boundary(512, d)] for d in (-1, 0, 1)]),
+    "b1_491k": (
+        [491061],
+        [[491067], [_boundary(1918, 1)], [_boundary(1919, 0)], [_boundary(1919, -1)]],
+    ),
+    "b1_1m": ([1048000], [[_boundary(4094, d)] for d in (-1, 0, 1)]),
+    # B = 2 live rows (different contexts).
+    "b2_mixed": ([491061, 32768], [[491067, 32774], [_boundary(1919, 1), _boundary(128, -1)]]),
+    # Padded: a B = 2 graph whose second row is a short dummy request.
+    "b2_padded": ([131072, 1], [[131078, 7], [_boundary(512, 0), 7]]),
+}
+
+ACCEPTED_SWEEP = [0, 1, 2, 3, 4, 5]
+
+
+def _run_eager_pair(native, decode_prologue, step):
+    for arm in (native, decode_prologue):
+        arm.set_step(*step)
+        arm.poison()
+    torch.cuda.synchronize()
+    nb, zb = orc.state(native), orc.state(decode_prologue)
+    out_n = orc.run_native(native)
+    out_z = orc.run_decode_prologue(decode_prologue)
+    torch.cuda.synchronize()
+    na, za = orc.state(native), orc.state(decode_prologue)
+    bad = orc.compare(nb, na, zb, za)
+    assert out_n.shape == out_z.shape and out_n.dtype == out_z.dtype
+    assert torch.equal(orc._bits(out_n), orc._bits(out_z)), "embedding rows differ"
+    return bad
+
+
+@pytest.mark.parametrize("case", sorted(CONTEXT_CASES))
+def test_eager_matches_native(case, weight):
+    past, kv_steps = CONTEXT_CASES[case]
+    cm, native, decode_prologue = orc.make_pair(past, weight)
+    try:
+        md = decode_prologue.md
+        reason = prologue.prologue_ineligibility(
+            decode_prologue.runner, decode_prologue.inputs(), True
+        )
+        assert reason is None, reason
+        print(
+            f"[decode_prologue] {case}: B={md.num_seqs} N={md.num_tokens} expand_for_dsl={md.expand_for_dsl} "
+            f"factor={md.dsl_expand_factor} width={md.gen_token_stride} "
+            f"dense={md.skip_indexer_for_gen_reqs} num_sms={md.num_sms}"
+        )
+        gen = torch.Generator().manual_seed(len(case))
+        trials = 0
+        for kv_base in kv_steps:
+            for a in ACCEPTED_SWEEP:
+                accepted = [a] + [(a + 3) % 6 for _ in kv_base[1:]]
+                if case == "b2_padded":
+                    accepted[1] = 5  # a padding row has offset 0
+                step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=a == 0)
+                bad = _run_eager_pair(native, decode_prologue, step)
+                assert not bad, (case, kv_base, accepted, bad)
+                trials += 1
+        print(f"[decode_prologue] {case}: {trials} eager trials identical")
+    finally:
+        cm.shutdown()
+
+
+@pytest.mark.parametrize("case", ["b1_491k", "b1_32k", "b1_2k", "b2_mixed", "b2_padded"])
+def test_cuda_graph_replays_match_native(case, weight):
+    past, kv_steps = CONTEXT_CASES[case]
+    cm, native, decode_prologue = orc.make_pair(past, weight)
+    try:
+        g_n, out_n = orc.capture(native, orc.run_native)
+        g_z, out_z = orc.capture(decode_prologue, orc.run_decode_prologue)
+        gen = torch.Generator().manual_seed(7)
+        replays = 0
+        for rep in range(4):
+            for kv_base in kv_steps:
+                accepted = [int(x) for x in torch.randint(0, 6, (len(kv_base),), generator=gen)]
+                if case == "b2_padded":
+                    accepted[1] = 5
+                step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=rep == 1)
+                for arm in (native, decode_prologue):
+                    arm.set_step(*step)
+                    arm.poison()
+                torch.cuda.synchronize()
+                nb, zb = orc.state(native), orc.state(decode_prologue)
+                g_n.replay()
+                g_z.replay()
+                torch.cuda.synchronize()
+                bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
+                assert not bad, (case, rep, kv_base, accepted, bad)
+                assert torch.equal(orc._bits(out_n), orc._bits(out_z)), (case, rep, "embedding")
+                replays += 1
+        assert replays >= 8
+        print(f"[decode_prologue] {case}: {replays} graph replays identical")
+    finally:
+        cm.shutdown()
+
+
+def test_dense_topk_graph_matches_native(weight):
+    """Short context: the indexer is skipped and the dense top-k rows are written."""
+    cm, native, decode_prologue = orc.make_pair([1000], weight, max_kv=4096)
+    try:
+        assert decode_prologue.md.skip_indexer_for_gen_reqs
+        reason = prologue.prologue_ineligibility(
+            decode_prologue.runner, decode_prologue.inputs(), True
+        )
+        assert reason is None, reason
+        g_n, out_n = orc.capture(native, orc.run_native)
+        g_z, out_z = orc.capture(decode_prologue, orc.run_decode_prologue)
+        gen = torch.Generator().manual_seed(3)
+        for kv_base in ([1006], [1001], [2040], [7]):
+            for a in ACCEPTED_SWEEP:
+                step = orc.step_inputs(gen, kv_base, [a], native.num_tokens)
+                for arm in (native, decode_prologue):
+                    arm.set_step(*step)
+                    arm.poison()
+                torch.cuda.synchronize()
+                nb, zb = orc.state(native), orc.state(decode_prologue)
+                g_n.replay()
+                g_z.replay()
+                torch.cuda.synchronize()
+                bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
+                assert not bad, (kv_base, a, bad)
+                assert torch.equal(orc._bits(out_n), orc._bits(out_z))
+    finally:
+        cm.shutdown()
+
+
+def _eager_and_graph(
+    native, decode_prologue, kv_steps, run_n, run_z, label, sweep=(0, 2, 5), reps=2
+):
+    """Eager trials over ``sweep`` accepted counts, then CUDA-graph replays; returns the counts."""
+    gen = torch.Generator().manual_seed(11)
+    trials = 0
+    for kv_base in kv_steps:
+        for a in sweep:
+            accepted = [a] + [(a + 3) % 6 for _ in kv_base[1:]]
+            step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=a == 0)
+            for arm in (native, decode_prologue):
+                arm.set_step(*step)
+                arm.poison()
+            torch.cuda.synchronize()
+            nb, zb = orc.state(native), orc.state(decode_prologue)
+            out_n, out_z = run_n(native), run_z(decode_prologue)
+            torch.cuda.synchronize()
+            bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
+            assert not bad, (label, kv_base, accepted, bad)
+            assert torch.equal(orc._bits(out_n), orc._bits(out_z)), (label, "embedding")
+            trials += 1
+
+    g_n, out_n = orc.capture(native, run_n)
+    g_z, out_z = orc.capture(decode_prologue, run_z)
+    replays = 0
+    for rep in range(reps):
+        for kv_base in kv_steps:
+            accepted = [int(x) for x in torch.randint(0, 6, (len(kv_base),), generator=gen)]
+            step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=rep == 1)
+            for arm in (native, decode_prologue):
+                arm.set_step(*step)
+                arm.poison()
+            torch.cuda.synchronize()
+            nb, zb = orc.state(native), orc.state(decode_prologue)
+            g_n.replay()
+            g_z.replay()
+            torch.cuda.synchronize()
+            bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
+            assert not bad, (label, rep, kv_base, accepted, bad)
+            assert torch.equal(orc._bits(out_n), orc._bits(out_z)), (label, rep, "embedding")
+            replays += 1
+    return trials, replays
+
+
+@pytest.mark.parametrize("mode", ["no_overlap", "no_spec_decode"])
+@pytest.mark.parametrize("case", ["b1_491k", "b1_2k", "b2_mixed"])
+def test_without_correction_matches_native(case, mode, weight):
+    """Without the overlap correction the native prologue is one ``on_update_kv_lens()``.
+
+    ``no_overlap``: spec decode with the overlap scheduler disabled;
+    ``no_spec_decode``: the step runs with ``enable_spec_decode=False``.
+    """
+    past, kv_steps = CONTEXT_CASES[case]
+    cm, native, decode_prologue = orc.make_pair(past, weight)
+    try:
+        spec = mode != "no_spec_decode"
+        if mode == "no_overlap":
+            for arm in (native, decode_prologue):
+                arm.runner._config.disable_overlap_scheduler = True
+        reason = prologue.prologue_ineligibility(
+            decode_prologue.runner, decode_prologue.inputs(), spec
+        )
+        assert reason is None, reason
+
+        def run_n(arm):
+            return orc.run_native(arm, enable_spec_decode=spec)
+
+        def run_z(arm):
+            return orc.run_decode_prologue(arm, enable_spec_decode=spec)
+
+        trials, replays = _eager_and_graph(
+            native, decode_prologue, kv_steps, run_n, run_z, (case, mode)
+        )
+        print(
+            f"[decode_prologue] {case} {mode}: {trials} eager trials and {replays} graph replays identical"
+        )
+    finally:
+        cm.shutdown()
+
+
+@pytest.mark.parametrize("case", ["b1_491k", "b1_2k", "b2_mixed"])
+def test_without_dsl_expansion_matches_native(case, weight):
+    """DeepGEMM paged-MQA logits without the DSL atom split: no expanded schedule (factor 1)."""
+    past, kv_steps = CONTEXT_CASES[case]
+    cm, native, decode_prologue = orc.make_pair(
+        past, weight, sparse_update={"use_cute_dsl_paged_mqa_logits": False}
+    )
+    try:
+        md = decode_prologue.md
+        assert not md.expand_for_dsl and not md.use_expanded_buffers_for_mtp
+        reason = prologue.prologue_ineligibility(
+            decode_prologue.runner, decode_prologue.inputs(), True
+        )
+        assert reason is None, reason
+        trials, replays = _eager_and_graph(
+            native,
+            decode_prologue,
+            kv_steps,
+            orc.run_native,
+            orc.run_decode_prologue,
+            (case, "no_dsl"),
+            sweep=tuple(ACCEPTED_SWEEP),
+        )
+        print(
+            f"[decode_prologue] {case} no_dsl: {trials} eager trials and {replays} graph replays identical"
+        )
+    finally:
+        cm.shutdown()
+
+
+class _HelixMapping:
+    @staticmethod
+    def has_cp_helix():
+        return True
+
+
+def test_declines_leave_native_untouched(weight):
+    """Out-of-domain steps: target_prologue returns False and writes nothing."""
+    cm, native, decode_prologue = orc.make_pair([32768], weight)
+    try:
+        md = decode_prologue.md
+        inputs = decode_prologue.inputs()
+        mrope = dict(inputs)
+        mrope["position_ids"] = torch.zeros((3, 1, md.num_tokens), dtype=torch.int32, device="cuda")
+        # (expected reason, attribute owner, attribute, value, inputs)
+        cases = [
+            ("expanded MTP buffers", md, "use_expanded_buffers_for_mtp", True, inputs),
+            ("ragged verification", md, "enable_ragged_verification", True, inputs),
+            (
+                "helix context parallelism",
+                decode_prologue.runner,
+                "mapping",
+                _HelixMapping(),
+                inputs,
+            ),
+            ("not a generation-only step", md, "_num_contexts", 1, inputs),
+            ("position ids are not a [1, N] int32 row", None, None, None, mrope),
+        ]
+        for expect, owner, name, value, case_inputs in cases:
+            saved = getattr(owner, name) if owner is not None else None
+            if owner is not None:
+                setattr(owner, name, value)
+            try:
+                reason = prologue.prologue_ineligibility(decode_prologue.runner, case_inputs, True)
+                before = orc.state(decode_prologue)
+                done = prologue.target_prologue(
+                    decode_prologue.runner,
+                    case_inputs,
+                    enable_spec_decode=True,
+                    runtime_draft_len=orc.MTP_K,
+                )
+                torch.cuda.synchronize()
+                after = orc.state(decode_prologue)
+            finally:
+                if owner is not None:
+                    setattr(owner, name, saved)
+            assert reason == expect, (expect, reason)
+            assert not done, expect
+            assert not orc.changed(before, after), (expect, orc.changed(before, after))
+            print(f"[decode_prologue] decline '{expect}' ok")
+    finally:
+        cm.shutdown()
+
+
+def test_take_embedding_requires_the_same_ids(weight):
+    cm, native, decode_prologue = orc.make_pair([32768], weight)
+    try:
+        inputs = decode_prologue.inputs()
+        assert prologue.target_prologue(
+            decode_prologue.runner, inputs, enable_spec_decode=True, runtime_draft_len=5
+        )
+        other_ids = inputs["input_ids"].clone()
+        assert prologue.take_embedding(decode_prologue.md, other_ids, decode_prologue.embed) is None
+        # The packet is consumed either way.
+        assert (
+            prologue.take_embedding(decode_prologue.md, inputs["input_ids"], decode_prologue.embed)
+            is None
+        )
+    finally:
+        cm.shutdown()
+
+
+_KINDS = ("KERNEL", "MEMCPY", "MEMSET", "EMPTY", "HOST", "EVENT", "MEM_ALLOC", "MEM_FREE")
+
+
+def _dot_nodes(path):
+    """(kind, label text) per node of a ``cudaGraphDebugDotPrint`` file, in file order."""
+    text = Path(path).read_text(errors="replace")
+    nodes = []
+    for label in re.findall(r'"graph_\d+_node_\d+"\[[^\]]*?label="(.*?)"\];', text, flags=re.S):
+        flat = re.sub(r"\\[<>]|[{}|]|\s+", " ", label).strip()
+        kind = next((k for k in _KINDS if flat.startswith(k)), "OTHER")
+        width = re.search(r"Width\s+(\d+)", flat)
+        name = re.search(r"topoId:\s*\d+\)\s+(\S+)", flat)
+        if kind == "MEMCPY" and width:
+            flat = f"MEMCPY {width.group(1)} B"
+        elif kind == "KERNEL" and name:
+            flat = name.group(1)
+        nodes.append((kind, flat[:160]))
+    return nodes
+
+
+def test_native_chain_census(weight, tmp_path):
+    """Print the captured native prologue's nodes (memcpy sizes included) and decode_prologue's."""
+    cm, native, decode_prologue = orc.make_pair([491061], weight)
+    try:
+        out_dir = Path(os.environ.get("GLM_MOE_DSA_DECODE_PROLOGUE_DOT_DIR", str(tmp_path)))
+        out_dir.mkdir(parents=True, exist_ok=True)
+        dot_n = str(out_dir / "native_prologue.dot")
+        dot_z = str(out_dir / "decode_prologue.dot")
+        orc.capture(native, orc.run_native, debug_dot=dot_n)
+        orc.capture(decode_prologue, orc.run_decode_prologue, debug_dot=dot_z)
+        for tag, path in (("native", dot_n), ("decode_prologue", dot_z)):
+            nodes = _dot_nodes(path)
+            kinds = [k for k, _ in nodes]
+            print(
+                f"[decode_prologue] census {tag}: {len(nodes)} nodes: "
+                f"{kinds.count('KERNEL')} kernels, {kinds.count('MEMCPY')} memcpy, "
+                f"{kinds.count('MEMSET')} memset; dot {path}"
+            )
+            for i, (kind, detail) in enumerate(nodes):
+                print(f"[decode_prologue]   {tag} {i:2d} {kind} {detail}")
+    finally:
+        cm.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py
new file mode 100644
index 000000000000..90cf8c0922cc
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py
@@ -0,0 +1,230 @@
+# 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.
+"""Helpers for the CPU tests of deferred_input_prep's host side (stager and loop domain).
+
+``load()`` returns deferred_input_prep's ``dispatch``, ``compat``, ``packet`` and ``stager``
+modules. Without a built ``tensorrt_llm`` it imports them from the source tree
+under a stand-in package, so ``PrepStager`` runs with CPU tensors. Logging and
+``prefer_pinned`` then need ``tensorrt_llm``: tests replace
+``dispatch._log_info`` and ``stager._host_memory_pinned``.
+
+``FakePacket`` replaces ``PrepPacket`` for CPU tensors. It performs the
+staged writes at its ``commit``, through NumPy views of the destinations.
+"""
+
+import importlib
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+REPO = Path(__file__).resolve().parents[9]
+MODELING_V2 = REPO / "tensorrt_llm" / "_torch" / "_experimental" / "modeling_v2"
+TARGET = MODELING_V2 / "models" / "glm_moe_dsa" / "shared"
+PACKAGES = TARGET / "kernels"
+_STANDALONE = "_input_prep_testlib"
+
+
+def _module_from(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def load_dispatch_standalone(name):
+    """The target's ``dispatch.py`` as module ``name``, without importing ``tensorrt_llm``.
+
+    Its absolute imports, ``tensorrt_llm.logger`` and the modeling_v2 router index, use
+    only the standard library: when they are not imported yet, they are loaded from their
+    source files and registered under their own names only while ``dispatch.py`` runs.
+    """
+    deps = (
+        ("tensorrt_llm.logger", REPO / "tensorrt_llm" / "logger.py"),
+        (
+            "tensorrt_llm._torch._experimental.modeling_v2._router_index",
+            MODELING_V2 / "_router_index.py",
+        ),
+    )
+    added = []
+    try:
+        for dep, path in deps:
+            if dep not in sys.modules:
+                added.append(dep)
+                _module_from(dep, path)
+        return _module_from(name, TARGET / "dispatch.py")
+    finally:
+        for dep in added:
+            sys.modules.pop(dep, None)
+
+
+def _standalone():
+    if _STANDALONE not in sys.modules:
+        for name, path in (
+            (_STANDALONE, TARGET),
+            (f"{_STANDALONE}.kernels", PACKAGES),
+            (f"{_STANDALONE}.kernels.deferred_input_prep", PACKAGES / "deferred_input_prep"),
+        ):
+            package = types.ModuleType(name)
+            package.__path__ = [str(path)]
+            sys.modules[name] = package
+        sys.modules[_STANDALONE].dispatch = load_dispatch_standalone(f"{_STANDALONE}.dispatch")
+    return SimpleNamespace(
+        dispatch=importlib.import_module(f"{_STANDALONE}.dispatch"),
+        compat=importlib.import_module(f"{_STANDALONE}.kernels.deferred_input_prep.compat"),
+        packet=importlib.import_module(f"{_STANDALONE}.kernels.deferred_input_prep.packet"),
+        stager=importlib.import_module(f"{_STANDALONE}.kernels.deferred_input_prep.stager"),
+    )
+
+
+def load():
+    """deferred_input_prep's host modules: the installed package, else the source tree standalone."""
+    try:
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+            compat,
+            packet,
+            stager,
+        )
+    except (ImportError, OSError):
+        return _standalone()
+    return SimpleNamespace(dispatch=dispatch, compat=compat, packet=packet, stager=stager)
+
+
+class FakePacket:
+    """``PrepPacket`` stand-in for CPU tensors (see the module docstring).
+
+    ``max_desc`` descriptors fit in one packet (one per write); a write whose
+    bytes overlap an earlier write of the packet raises ``PacketOverlap``.
+    """
+
+    def __init__(self, packet_module, max_desc=1 << 20):
+        self._full = packet_module.PacketFull
+        self._overlap = packet_module.PacketOverlap
+        self.max_desc = max_desc
+        self._ops = []
+        self._ranges = []
+        self.num_commits = 0
+        self.num_writes = 0
+
+    def begin(self):
+        self._ops, self._ranges = [], []
+
+    def abort(self):
+        self._ops, self._ranges = [], []
+
+    @property
+    def num_descriptors(self):
+        return len(self._ops)
+
+    def _claim(self, *dsts):
+        if len(self._ops) + 1 > self.max_desc:
+            raise self._full("fake packet is full")
+        new = [(t.data_ptr(), t.data_ptr() + t.numel() * t.element_size()) for t in dsts]
+        for lo, hi in new:
+            for a, b in self._ranges:
+                if lo < b and a < hi:
+                    raise self._overlap("overlapping destinations in one packet")
+        self._ranges.extend(new)
+
+    def _add(self, fn, *dsts):
+        self._claim(*dsts)
+        self._ops.append(fn)
+
+    def add_copy(self, dst, src, clamp_min=None):
+        def op():
+            values = src.reshape(dst.shape).numpy()
+            dst.numpy()[...] = values if clamp_min is None else values.clip(min=clamp_min)
+
+        self._add(op, dst)
+
+    def add_values(self, dst, values):
+        def op():
+            dst.numpy()[...] = list(values)
+
+        self._add(op, dst)
+
+    def add_gather(self, dst, src, indices, addend=0):
+        def op():
+            flat = src.reshape(-1).numpy()
+            out = dst.numpy()
+            out[...] = 0
+            out[: len(indices)] = flat[list(indices)] + addend
+
+        self._add(op, dst)
+
+    def add_zero(self, dst):
+        def op():
+            dst.numpy()[...] = 0
+
+        self._add(op, dst)
+
+    def add_clamp_min(self, dst, lo):
+        def op():
+            out = dst.numpy()
+            out[...] = out.clip(min=lo)
+
+        self._add(op, dst)
+
+    def commit(self, stream=None, record_event=True):
+        ops, self._ops, self._ranges = self._ops, [], []
+        for op in ops:
+            op()
+        self.num_commits += 1
+        self.num_writes += len(ops)
+        return len(ops)
+
+    def set_completion_event(self, event):
+        pass
+
+
+def eligibility_case(**overrides):
+    """Arguments of ``PrepStager.begin_if_eligible`` for a steady served decode step."""
+    mode = SimpleNamespace(use_one_engine=lambda: True)
+    cfg = SimpleNamespace(
+        spec_config=SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True),
+        disable_overlap_scheduler=overrides.pop("disable_overlap_scheduler", False),
+        enable_attention_dp=overrides.pop("enable_attention_dp", False),
+        use_mrope=False,
+    )
+    runner = SimpleNamespace(
+        _config=cfg,
+        use_beam_search=False,
+        mapping=SimpleNamespace(has_cp_helix=lambda: False, pp_size=overrides.pop("pp_size", 1)),
+        guided_decoder=None,
+        _lora=SimpleNamespace(cuda_graph_manager=None),
+    )
+    rows = overrides.pop("rows", [SimpleNamespace(is_dummy=False, py_batch_idx=0)])
+    requests = SimpleNamespace(num_context_requests=0, generation_requests=rows)
+    DSAtrtllmAttentionMetadata = type("DSAtrtllmAttentionMetadata", (), {})
+    manager = SimpleNamespace(index_mapper=object(), _use_per_layer_page_tables=False)
+    attn = DSAtrtllmAttentionMetadata()
+    attn.kv_cache_manager = manager
+    attn.draft_kv_cache_manager = None
+    args = dict(
+        runner=runner,
+        requests=requests,
+        attn_metadata=attn,
+        spec_metadata=SimpleNamespace(),
+        new_tensors_device=SimpleNamespace(next_draft_tokens=object()),
+        graph_replay=True,
+        is_dummy=False,
+        enable_spec_decode=True,
+        promoted=False,
+    )
+    args.update(overrides)
+    return args
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py
new file mode 100644
index 000000000000..a212ca16f519
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py
@@ -0,0 +1,340 @@
+# 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.
+"""deferred_input_prep's executor-loop domain, its engagement rule, its debug level and its window (CPU only).
+
+The step-input stager engages when the target declares it (the plan runs
+``deferred_input_prep``) and the step is in its support domain; the store writer engages
+once a non-attention-DP overlap loop declared its domain. Nothing depends on which other
+packages are selected or on their sources. An engaged preparation opens the deferred-write window, which
+``PrepStager.flush`` can close early for the writes staged so far, also handing ``input_ids`` /
+``position_ids`` back to the graph replay's own copy.
+"""
+
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import input_prep_testlib  # noqa: E402
+
+HOST = torch.arange(10, 14, dtype=torch.int32)
+
+
+@pytest.fixture
+def deferred_input_prep(monkeypatch):
+    """deferred_input_prep's host modules with logging captured in ``deferred_input_prep.lines``."""
+    modules = input_prep_testlib.load()
+    lines = []
+    monkeypatch.setattr(modules.dispatch, "_log_info", lines.append)
+    monkeypatch.setattr(modules.dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(modules.dispatch, "_INSTALLED", [None])
+    monkeypatch.setattr(
+        modules.compat, "_LOOP_STATE", {"declared": False, "reason": "", "logged": False}
+    )
+    modules.lines = lines
+    return modules
+
+
+@pytest.fixture
+def compat(deferred_input_prep):
+    return deferred_input_prep.compat
+
+
+def test_loop_domain_state(compat):
+    """Store writer: only after a non-ADP overlap loop declared."""
+    state = compat._LOOP_STATE
+    assert compat.loop_domain() is False
+    compat.declare_overlap_loop(False)
+    assert compat.loop_domain() is True
+    compat.declare_overlap_loop(True)
+    assert compat.loop_domain() is False and "attention DP" in state["reason"]
+    compat.declare_overlap_loop(False)
+    assert compat.loop_domain() is False
+    state.update(declared=False, reason="")
+    compat.declare_overlap_loop(None)
+    assert compat.loop_domain() is False and "did not pass" in state["reason"]
+
+
+def test_every_part_engages_without_any_record(deferred_input_prep, monkeypatch):
+    """A shared switch set to a value no oracle ran: deferred_input_prep engages.
+
+    Engagement reads no package source and no table of package versions, so
+    editing another package's sources cannot turn deferred_input_prep off.
+    """
+    monkeypatch.setenv("TRTLLM_ENABLE_PDL", "0")
+    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
+    st = deferred_input_prep.stager.PrepStager()
+    st.packet = input_prep_testlib.FakePacket(deferred_input_prep.packet)
+    case = input_prep_testlib.eligibility_case()
+    attached = (
+        case["runner"],
+        case["attn_metadata"],
+        case["spec_metadata"],
+        case["attn_metadata"].kv_cache_manager,
+    )
+    assert st.begin_if_eligible(**case) is st
+    # The attribute every engine call site reads (RuntimeHooks.step_input_stager).
+    assert all(obj.step_input_stager is st for obj in attached)
+    st.end()
+    assert all(obj.step_input_stager is None for obj in attached)
+    assert st.decisions == {"engaged": 1} and st.num_commits == 1
+    assert deferred_input_prep.lines[0].startswith(
+        "[glm_moe_dsa] deferred_input_prep/prep: engaged"
+    )
+    deferred_input_prep.compat.declare_overlap_loop(False)
+    assert deferred_input_prep.compat.loop_domain()
+
+
+def test_prep_declines_only_for_pageable_host_memory(deferred_input_prep, monkeypatch):
+    """Whether a stager exists is the plan's decision; the stager declines only pageable memory."""
+    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
+    st = deferred_input_prep.stager.PrepStager()
+    st.packet = input_prep_testlib.FakePacket(deferred_input_prep.packet)
+    assert st.begin_if_eligible(**input_prep_testlib.eligibility_case()) is st
+    st.end()
+    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: False)
+    st = deferred_input_prep.stager.PrepStager()
+    assert st.begin_if_eligible(**input_prep_testlib.eligibility_case()) is None
+    assert st.last_reason == "host page tables are not pinned (prefer_pinned)"
+    assert not st.active
+
+
+def test_debug_level_is_the_constructor_argument(deferred_input_prep, monkeypatch):
+    """The stager reads no environment: its debug level is the argument (0, 1 or 2)."""
+    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
+    monkeypatch.setenv("TRTLLM_MODELING_V2_VALIDATE", "1")
+    PrepStager = deferred_input_prep.stager.PrepStager
+    assert PrepStager().debug is None
+    assert PrepStager(debug_level=1).debug.level == 1
+    assert PrepStager(debug_level=2).debug.level == 2
+    with pytest.raises(ValueError, match="0, 1 or 2"):
+        PrepStager(debug_level=3)
+
+
+def _engaged(deferred_input_prep, monkeypatch, packet=None, debug_level=0):
+    """An engaged stager on a CPU case whose runner owns a 16-element ``input_ids_cuda``."""
+    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
+    case = input_prep_testlib.eligibility_case()
+    case["runner"].input_ids_cuda = torch.zeros(16, dtype=torch.int32)
+    st = deferred_input_prep.stager.PrepStager(debug_level=debug_level)
+    st.packet = (
+        packet if packet is not None else input_prep_testlib.FakePacket(deferred_input_prep.packet)
+    )
+    assert st.begin_if_eligible(**case) is st
+    return st, case
+
+
+def test_engaged_line_states_the_window_contract(deferred_input_prep, monkeypatch):
+    st, _ = _engaged(deferred_input_prep, monkeypatch)
+    st.end()
+    assert deferred_input_prep.lines == [
+        "[glm_moe_dsa] deferred_input_prep/prep: engaged (deferred-write window from "
+        "DecoderRunner._forward_decoder's begin_if_eligible to stager.end(): the buffers written "
+        "through the deferred_input_prep stager are written later by the commit; do not read or "
+        "write them in between with other device ops; route writes through the stager API or "
+        "call stager.flush() first; for input_ids_cuda and position_ids_cuda (also staged into "
+        "the graph's static inputs) call stager.flush() and then write natively; "
+        "disable with TRTLLM_STAIRCASE_GLM_LL_DISABLE=deferred_input_prep)"
+    ]
+
+
+def test_flush_then_another_write_keeps_native_order(deferred_input_prep, monkeypatch):
+    """After flush, a write to a staged buffer lands after the staged write, as natively."""
+    st, case = _engaged(deferred_input_prep, monkeypatch)
+    buf = case["runner"].input_ids_cuda
+    st.copy(buf[:4], HOST)
+    st.flush()
+    assert buf[:4].tolist() == HOST.tolist() and st.num_flushes == 1
+    buf[:4].add_(100)
+    st.zero(buf[4:])
+    st.end()
+    assert st.packet.num_commits == 2 and st.num_commits == 1
+    assert buf.tolist() == (HOST + 100).tolist() + [0] * 12
+
+
+@pytest.mark.parametrize("max_desc", [0, 2, 1 << 20])
+def test_flush_issues_several_writes_in_staging_order(deferred_input_prep, monkeypatch, max_desc):
+    """Overlapping writes staged before one flush land in staging order.
+
+    ``max_desc`` 0: every write is native; 2: the packet fills up; otherwise
+    only the overlaps split the packet.
+    """
+    packet = input_prep_testlib.FakePacket(deferred_input_prep.packet, max_desc=max_desc)
+    st, case = _engaged(deferred_input_prep, monkeypatch, packet=packet)
+    buf = case["runner"].input_ids_cuda
+    buf.fill_(5)
+    st.copy(buf[:4], torch.tensor([-3, 1, 7, -9], dtype=torch.int32))
+    st.zero(buf[4:])
+    st.clamp_min(buf[:4], 0)
+    st.values(buf[:2], [-1, -2])
+    st.flush()
+    expected = [-1, -2, 7, 0] + [0] * 12
+    assert buf.tolist() == expected
+    st.end()
+    assert buf.tolist() == expected
+
+
+def test_flush_and_commit_fall_back_to_native_writes(deferred_input_prep, monkeypatch):
+    """With a packet that holds nothing, flush and the commit issue each write natively."""
+    full = input_prep_testlib.FakePacket(deferred_input_prep.packet, max_desc=0)
+    st, case = _engaged(deferred_input_prep, monkeypatch, packet=full)
+    buf = case["runner"].input_ids_cuda
+    st.copy(buf[:4], HOST)
+    st.flush()
+    assert buf[:4].tolist() == HOST.tolist() and st.num_native_copies == 1
+    st.values(buf[4:8], [1, 2, 3, 4])
+    st.zero(buf[8:])
+    st.end()
+    assert st.num_native_copies == 3 and full.num_writes == 0
+    assert buf.tolist() == HOST.tolist() + [1, 2, 3, 4] + [0] * 8
+    st.flush()  # outside a window: nothing to do
+    assert st.num_native_copies == 3
+
+
+def test_flush_ends_the_guarded_span_of_the_flushed_writes(deferred_input_prep, monkeypatch):
+    """Debug guard: a flushed buffer may be touched; a buffer staged later stays guarded."""
+    guard = sys.modules[deferred_input_prep.stager.DebugWindow.__module__]
+
+    def cpu_span(t):
+        lo = t.data_ptr()
+        return (lo, lo + t.numel() * t.element_size()) if t.numel() else None
+
+    monkeypatch.setattr(guard, "tensor_span", cpu_span)
+    st, case = _engaged(deferred_input_prep, monkeypatch, debug_level=1)
+    buf = case["runner"].input_ids_cuda
+    st.copy(buf[:4], HOST)
+    st.flush()
+    buf[:4].add_(100)
+    st.zero(buf[4:])
+    try:
+        with pytest.raises(guard.InputPrepHazardError, match="touches a destination staged"):
+            buf[4:].add_(1)
+    finally:
+        st.discard()
+    assert buf.tolist() == (HOST + 100).tolist() + [0] * 12
+
+
+STEP_TOKENS = 4
+STAGED_INPUT_IDS = [11, 21, 12, 22]
+STAGED_POSITIONS = [7, 8, 9, 10]
+
+
+def _stage_input_ids(st, runner):
+    """``overlap_inputs`` of a steady step: two rows with previous tensors, 2 tokens per row."""
+    new_tokens = torch.tensor([[11, 12, 13, 14], [21, 22, 23, 24]], dtype=torch.int32)
+    next_draft_tokens = torch.tensor([[31], [32], [33], [34]], dtype=torch.int32)
+    new_tokens_lens = torch.tensor([2, 1, 2, 2], dtype=torch.int32)
+    st.overlap_inputs(
+        runner, [0, 1], [0, 1, 2, 3], 0, 0, 2, 2, new_tokens, next_draft_tokens, new_tokens_lens
+    )
+
+
+def _stage_position_ids(st, runner):
+    st.position_ids(
+        runner,
+        runner.position_ids_cuda[:STEP_TOKENS],
+        torch.tensor(STAGED_POSITIONS, dtype=torch.int32),
+    )
+
+
+def _graph_step(deferred_input_prep, monkeypatch):
+    """An engaged step whose runner has the decode graph's static input buffers."""
+    st, case = _engaged(deferred_input_prep, monkeypatch)
+    runner = case["runner"]
+    for name in (
+        "position_ids_cuda",
+        "draft_tokens_cuda",
+        "previous_batch_indices_cuda",
+        "previous_pos_indices_cuda",
+        "previous_pos_id_offsets_cuda",
+        "previous_kv_lens_offsets_cuda",
+    ):
+        setattr(runner, name, torch.zeros(16, dtype=torch.int32))
+    static = {
+        "input_ids": torch.zeros(16, dtype=torch.int32),
+        "position_ids": torch.zeros((1, 16), dtype=torch.int32),
+    }
+    runner.cuda_graph_runner = SimpleNamespace(shared_static_tensors=static)
+    return st, runner, static
+
+
+def _graph_inputs(st, runner, static):
+    """``input_ids`` / ``position_ids`` as ``CUDAGraphRunner.replay`` hands them to the graph."""
+    n = STEP_TOKENS
+    if st.static_inputs_tokens() != n:
+        static["input_ids"][:n].copy_(runner.input_ids_cuda[:n])
+        static["position_ids"][0, :n].copy_(runner.position_ids_cuda[:n])
+    return static["input_ids"][:n].tolist(), static["position_ids"][0, :n].tolist()
+
+
+def test_staged_step_writes_the_graph_inputs(deferred_input_prep, monkeypatch):
+    """The commit writes the static inputs too, so the replay skips its own copy."""
+    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
+    _stage_input_ids(st, runner)
+    _stage_position_ids(st, runner)
+    st.end()
+    assert st.static_inputs_tokens() == STEP_TOKENS
+    assert runner.input_ids_cuda[:STEP_TOKENS].tolist() == STAGED_INPUT_IDS
+    assert _graph_inputs(st, runner, static) == (STAGED_INPUT_IDS, STAGED_POSITIONS)
+
+
+def test_native_writes_after_flush_reach_the_graph_inputs(deferred_input_prep, monkeypatch):
+    """flush, then native writes to input_ids_cuda / position_ids_cuda: the graph sees them."""
+    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
+    _stage_input_ids(st, runner)
+    _stage_position_ids(st, runner)
+    st.flush()
+    runner.input_ids_cuda[:STEP_TOKENS].copy_(torch.tensor([101, 102, 103, 104]))
+    runner.position_ids_cuda[:STEP_TOKENS].add_(100)
+    st.end()
+    assert st.static_inputs_tokens() == 0
+    assert _graph_inputs(st, runner, static) == (
+        [101, 102, 103, 104],
+        [p + 100 for p in STAGED_POSITIONS],
+    )
+
+
+def test_flush_between_input_ids_and_position_ids(deferred_input_prep, monkeypatch):
+    """Only position_ids is staged again after the flush: the replay copies both inputs."""
+    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
+    _stage_input_ids(st, runner)
+    st.flush()
+    runner.input_ids_cuda[:STEP_TOKENS].add_(100)
+    _stage_position_ids(st, runner)
+    st.end()
+    assert st.static_inputs_tokens() == 0
+    assert _graph_inputs(st, runner, static) == (
+        [t + 100 for t in STAGED_INPUT_IDS],
+        STAGED_POSITIONS,
+    )
+
+
+def test_staged_write_to_input_ids_does_not_reach_the_graph(deferred_input_prep, monkeypatch):
+    """Why the window contract sends new input_ids_cuda writes through flush.
+
+    The static inputs staged by ``overlap_inputs`` are what the graph replays;
+    a later staged write to ``input_ids_cuda`` changes only the runner buffer.
+    """
+    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
+    _stage_input_ids(st, runner)
+    _stage_position_ids(st, runner)
+    st.copy(runner.input_ids_cuda[:STEP_TOKENS], torch.tensor([101, 102, 103, 104]))
+    st.end()
+    assert runner.input_ids_cuda[:STEP_TOKENS].tolist() == [101, 102, 103, 104]
+    assert _graph_inputs(st, runner, static)[0] == STAGED_INPUT_IDS
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py
new file mode 100644
index 000000000000..543ccdd38174
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py
@@ -0,0 +1,872 @@
+# 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.
+"""Byte-exact checks of the deferred_input_prep kernels against the native operations they replace.
+
+Also: ordering (last write wins), native fallbacks, delayed-producer stress of
+the pinned rings (GPU held back by a sleep kernel while many launches are
+staged) and the debug guard. Set ``GLM_MOE_DSA_DEFERRED_INPUT_PREP_UNDER_SANITIZER=1`` when running
+under compute-sanitizer: it slows the host enough that the "GPU still busy"
+precondition of the stress tests may not hold, so only that precondition is
+not asserted.
+"""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+UNDER_SANITIZER = os.environ.get("GLM_MOE_DSA_DEFERRED_INPUT_PREP_UNDER_SANITIZER") == "1"
+# About one second of GPU spin on sm_103: far longer than staging the launches.
+HOLD_CYCLES = 2_000_000_000
+
+
+def _stager(debug_level=0):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.stager import (  # noqa: E501
+        PrepStager,
+    )
+
+    return PrepStager(debug_level=debug_level)
+
+
+@triton.jit
+def _copy_tensor_kernel(src, dst, n, BLOCK: tl.constexpr):
+    offs = tl.arange(0, BLOCK)
+    m = offs < n
+    tl.store(dst + offs, tl.load(src + offs, mask=m), mask=m)
+
+
+@triton.jit
+def _copy_addr_kernel(src_addr, dst_addr, n, BLOCK: tl.constexpr):
+    src = src_addr.to(tl.pointer_type(tl.int32))
+    dst = dst_addr.to(tl.pointer_type(tl.int32))
+    offs = tl.arange(0, BLOCK)
+    m = offs < n
+    tl.store(dst + offs, tl.load(src + offs, mask=m), mask=m)
+
+
+def _rand_i32(*shape, lo=-1000, hi=1 << 20):
+    return torch.randint(lo, hi, shape, dtype=torch.int32)
+
+
+def test_copy_values_gather_and_zero_tail():
+    torch.manual_seed(0)
+    dev = "cuda"
+    a = _rand_i32(37).to(dev)
+    b = torch.zeros(5000, dtype=torch.int32, device=dev)  # multi-chunk copy
+    c = _rand_i32(128).to(dev)  # gather target with zero tail (previous_pos_id_offsets)
+    d = torch.zeros(4, dtype=torch.int64, device=dev)  # int64 copy
+    f = torch.zeros(3, dtype=torch.float32, device=dev)
+    src = _rand_i32(7, 1, 3).to(dev)  # sampler store [W, S, B]
+    a_ref, b_ref, c_ref, d_ref, f_ref = a.clone(), b.clone(), c.clone(), d.clone(), f.clone()
+
+    host_b = _rand_i32(5000).pin_memory()
+    host_d = torch.tensor([1, -2, 1 << 40, 7], dtype=torch.int64).pin_memory()
+    host_f = torch.tensor([0.5, -1.25, 3.0], dtype=torch.float32)
+    values = list(range(10, 47))
+    idx = [5 * 3 + 2, 2 * 3 + 2, 0]
+
+    # Native reference.
+    a_ref.copy_(torch.tensor(values, dtype=torch.int32), non_blocking=True)
+    b_ref.copy_(host_b, non_blocking=True)
+    c_ref *= 0
+    c_ref[10:13] = src.flatten()[torch.tensor(idx, device=dev)] - 6
+    d_ref.copy_(host_d, non_blocking=True)
+    f_ref.copy_(host_f, non_blocking=True)
+
+    st = _stager()
+    st.begin()
+    st.values(a, values)
+    st.copy(b, host_b)
+    st.gather(c[:10], src, [], 0)
+    st.gather(c[10:13], src, idx, -6)
+    st.gather(c[13:], src, [], 0)
+    st.copy(d, host_d)
+    st.copy(f, host_f)
+    st.commit()
+    torch.cuda.synchronize()
+    for x, y in ((a, a_ref), (b, b_ref), (c, c_ref), (d, d_ref), (f, f_ref)):
+        assert torch.equal(x.view(torch.uint8), y.view(torch.uint8))
+
+
+def test_host_bytes_read_at_commit_time():
+    dev = "cuda"
+    dst = torch.zeros(8, dtype=torch.int32, device=dev)
+    host = torch.zeros(8, dtype=torch.int32).pin_memory()
+    st = _stager()
+    st.begin()
+    st.copy(dst, host)
+    host.fill_(42)  # a later in-place host update before commit is what the copy sees
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(dst.cpu(), torch.full((8,), 42, dtype=torch.int32))
+
+
+def test_block_offsets_matches_native_kernel():
+    from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import (
+        copy_batch_block_offsets_to_device,
+    )
+
+    torch.manual_seed(1)
+    pools, src_seqs, dst_seqs, blocks = 2, 3, 2, 16388
+    host = _rand_i32(pools, src_seqs, 2, blocks, lo=-1, hi=50000).pin_memory()
+    host[:, :, :, 7000:] = -1
+    scales = torch.tensor([3, 5], dtype=torch.int32).pin_memory()
+    kv_off = torch.tensor([11, 0], dtype=torch.int32).pin_memory()
+    copy_idx = torch.tensor([2, 0], dtype=torch.int32).pin_memory()
+    ref = _rand_i32(pools, dst_seqs + 1, 2, blocks).cuda()
+    out = ref.clone()
+    copy_batch_block_offsets_to_device(
+        host, ref, copy_idx, scales, kv_off, torch.cuda.current_stream().cuda_stream
+    )
+    st = _stager()
+    st.begin()
+    st.block_offsets(host, out, copy_idx.tolist(), scales.tolist(), kv_off.tolist())
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(out, ref)
+
+
+def test_copy_clamp_and_clamp_rest():
+    dev = "cuda"
+    buf = _rand_i32(4, 300).to(dev)
+    ref = buf.clone()
+    host = _rand_i32(1, 300, lo=-5, hi=5).pin_memory()
+    ref[:1].copy_(host, non_blocking=True)
+    ref.clamp_(min=0)
+    st = _stager()
+    st.begin()
+    st.copy(buf[:1], host, clamp_min=0)
+    st.clamp_min(buf[1:], 0)
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(buf, ref)
+
+
+def test_overlapping_writes_keep_native_order():
+    """Overlapping staged writes land in staging order (one more launch each)."""
+    torch.manual_seed(6)
+    dst = _rand_i32(64).cuda()
+    ref = dst.clone()
+    host = _rand_i32(64, lo=-20, hi=20).pin_memory()
+    ref.copy_(host, non_blocking=True)
+    ref[4:12] = torch.arange(8, dtype=torch.int32, device="cuda")
+    ref[:8].clamp_(min=5)
+    ref[60:].zero_()
+    st = _stager()
+    st.begin()
+    st.copy(dst, host)
+    st.values(dst[4:12], list(range(8)))
+    st.clamp_min(dst[:8], 5)
+    st.zero(dst[60:])
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(dst, ref)
+    assert st.num_flushes == 2 and st.num_native_copies == 0
+
+
+def test_flush_then_native_op_keeps_native_order():
+    """``flush`` in the window: a native op after it sees the staged writes."""
+    torch.manual_seed(8)
+    dst = _rand_i32(64).cuda()
+    ref = dst.clone()
+    host = _rand_i32(64, lo=-20, hi=20).pin_memory()
+    ref.copy_(host, non_blocking=True)
+    ref[:8].clamp_(min=5)
+    ref[4:12] += 7
+    ref[60:].zero_()
+    st = _stager()
+    st.begin()
+    st.copy(dst, host)
+    st.clamp_min(dst[:8], 5)
+    st.flush()
+    dst[4:12] += 7
+    st.zero(dst[60:])
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(dst, ref)
+    assert st.num_flushes == 2 and st.num_native_copies == 0 and st.num_commits == 1
+
+
+def test_every_kind_falls_back_to_native_in_order():
+    """Writes too large for an empty packet are issued natively, in order."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.packet import (  # noqa: E501
+        PrepPacket,
+    )
+
+    torch.manual_seed(7)
+    dev = "cuda"
+    a = _rand_i32(8).to(dev)
+    b = _rand_i32(5000).to(dev)
+    c = _rand_i32(32).to(dev)
+    e = _rand_i32(5000).to(dev)
+    f = _rand_i32(5000, lo=-9, hi=9).to(dev)
+    tail = _rand_i32(2).to(dev)
+    src = _rand_i32(6, 1, 1).to(dev)
+    host_b = _rand_i32(5000).pin_memory()
+    pools, blocks = 2, 9000
+    host_t = _rand_i32(pools, 3, 2, blocks, lo=-1, hi=50000).pin_memory()
+    bo = _rand_i32(pools, 2, 2, blocks).to(dev)
+    refs = [t.clone() for t in (a, b, c, e, f, tail, bo)]
+    ra, rb, rc, re_, rf, rtail, rbo = refs
+    ra.copy_(torch.arange(8, dtype=torch.int32), non_blocking=True)
+    rb.copy_(host_b, non_blocking=True)
+    rc.zero_()
+    rc[:6] = src.flatten() + 2
+    re_.zero_()
+    rf.clamp_(min=0)
+    for p in range(pools):
+        for s, row in enumerate([2, 0]):
+            v = host_t[p, row, 0].cuda()
+            rbo[p, s, 0] = torch.where(v == -1, 0, v * (p + 2))
+            rbo[p, s, 1] = torch.where(v == -1, 0, v * (p + 2) + 3)
+    rtail.copy_(torch.tensor([5, 6], dtype=torch.int32), non_blocking=True)
+
+    st = _stager()
+    st.packet = PrepPacket(max_desc=1, payload_words=4)
+    st.begin()
+    st.values(a, list(range(8)))  # payload 8 > 4 words
+    st.copy(b, host_b)  # 2 chunks > 1 descriptor
+    st.gather(c, src, list(range(6)), addend=2)  # payload 6 > 4
+    st.zero(e)  # 2 chunks
+    st.clamp_min(f, 0)  # 2 chunks
+    st.block_offsets(host_t, bo, [2, 0], [2, 3], [3, 3])  # 3 chunks per row
+    st.values(tail, [5, 6])  # fits
+    st.commit()
+    torch.cuda.synchronize()
+    for got, want in zip((a, b, c, e, f, tail, bo), refs):
+        assert torch.equal(got, want)
+    assert st.num_native_copies == 5 + 2 * pools  # a, b, c, e, f + every block row
+
+
+@pytest.mark.parametrize("caller_event", [False, True])
+def test_packet_ring_under_gpu_lag(caller_event):
+    """Delayed producer: many packets staged while the GPU is held back.
+
+    Every launch reads its descriptors, its payload and the caller's pinned
+    sources long after ``commit`` returned. The caller drops its pinned
+    temporary right after staging and the pinned allocator is churned with
+    junk of the same size, so a slot or source reused before its launch ran
+    would show up as wrong destination bytes.
+    """
+    torch.manual_seed(8)
+    dev = "cuda"
+    n_iter, words = 24, 3 * 4096 + 7
+    big = [torch.zeros(words, dtype=torch.int32, device=dev) for _ in range(n_iter)]
+    small = [torch.zeros(9, dtype=torch.int32, device=dev) for _ in range(n_iter)]
+    vals = [torch.zeros(5, dtype=torch.int32, device=dev) for _ in range(n_iter)]
+    st = _stager()
+    st.begin()
+    st.values(vals[0], [0] * 5)  # compile the kernel outside the lag window
+    st.commit()
+    torch.cuda.synchronize()
+
+    expected, junk = [], []
+    torch.cuda._sleep(HOLD_CYCLES)
+    held = torch.cuda.Event()
+    held.record()
+    for i in range(n_iter):
+        host_big = _rand_i32(words).pin_memory()  # read in place by the kernel
+        host_small = _rand_i32(9)  # pageable: goes through the payload
+        values = [100 * i + k for k in range(5)]
+        expected.append((host_big.clone(), host_small.clone(), values))
+        st.begin()
+        st.copy(big[i], host_big)
+        st.copy(small[i], host_small)
+        st.values(vals[i], values)
+        if caller_event:
+            st.commit(record_event=False)
+            ev = torch.cuda.Event()
+            ev.record()
+            st.set_completion_event(ev)
+        else:
+            st.commit()
+        del host_big, host_small
+        junk.append(torch.full((words,), -7, dtype=torch.int32).pin_memory())
+    lagging = not held.query()
+    slots_while_lagging = len(st.packet._slots)
+    torch.cuda.synchronize()
+    for i, (hb, hs, values) in enumerate(expected):
+        assert torch.equal(big[i].cpu(), hb), i
+        assert torch.equal(small[i].cpu(), hs), i
+        assert vals[i].tolist() == values, i
+    assert all(bool((j == -7).all()) for j in junk)
+    if not UNDER_SANITIZER:
+        assert lagging, "the GPU finished its hold before staging ended; inconclusive"
+        assert slots_while_lagging >= n_iter, slots_while_lagging
+
+
+def test_sampler_store_ring_under_gpu_lag():
+    """Delayed producer for the sampler store writer (slot ring + zero-copy host targets)."""
+    import random
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
+        SamplerStoreWriter,
+    )
+
+    torch.manual_seed(9)
+    rng = random.Random(9)
+    dev = "cuda"
+    num_slots, w, d, n_iter = 8, 6, 5, 24
+
+    def zeros(*shape):
+        return torch.zeros(shape, dtype=torch.int32, device=dev)
+
+    store = SimpleNamespace(
+        new_tokens=zeros(w, num_slots, 1),
+        next_new_tokens=zeros(w, num_slots, 1),
+        new_tokens_lens=zeros(num_slots),
+        next_draft_tokens=zeros(num_slots, d),
+    )
+    ref = SimpleNamespace(**{k: v.cpu() for k, v in vars(store).items()})
+    calls = []
+    for i in range(n_iter):
+        slots = rng.sample(range(num_slots), rng.randint(1, 3))
+        n = len(slots)
+        calls.append(
+            (
+                slots,
+                _rand_i32(n, w).to(dev),
+                _rand_i32(n, w).to(dev),
+                _rand_i32(n, lo=1, hi=7).to(dev),
+                _rand_i32(n, d).to(dev),
+            )
+        )
+    calls_cpu = [tuple(t.cpu() for t in c[1:]) for c in calls]
+    writer = SamplerStoreWriter()
+    warm = [
+        torch.empty_like(t, device="cpu", pin_memory=True)
+        for t in (store.new_tokens, store.new_tokens_lens, store.next_draft_tokens)
+    ]
+    scratch = SimpleNamespace(**{k: v.clone() for k, v in vars(store).items()})
+    # Compile every argument specialization (request count, alignments) of the
+    # kernel outside the lag window: a compile inside it outlasts the hold.
+    for call in calls:
+        writer.store_and_copy_out(scratch, *call, *warm)
+    torch.cuda.synchronize()
+
+    kept, junk = [], []
+    torch.cuda._sleep(HOLD_CYCLES)
+    held = torch.cuda.Event()
+    held.record()
+    for i, ((slots, o_new, o_next, o_lens, o_draft), cpu) in enumerate(zip(calls, calls_cpu)):
+        # Host reference from host copies made before the hold (a D2H copy here
+        # would wait for the GPU).
+        s = torch.tensor(slots)
+        ref.new_tokens[:, s, 0] = cpu[0].T
+        ref.next_new_tokens[:, s, 0] = cpu[1].T
+        ref.new_tokens_lens[s] = cpu[2]
+        ref.next_draft_tokens[s] = cpu[3]
+        want = (ref.new_tokens.clone(), ref.new_tokens_lens.clone(), ref.next_draft_tokens.clone())
+        host = [
+            torch.empty_like(t, device="cpu", pin_memory=True)
+            for t in (store.new_tokens, store.new_tokens_lens, store.next_draft_tokens)
+        ]
+        writer.store_and_copy_out(
+            store, slots, o_new, o_next, o_lens, o_draft, *host, record_event=False
+        )
+        ev = torch.cuda.Event()
+        ev.record()
+        writer.set_completion_event(ev)
+        if i % 2:
+            kept.append((host, want))
+        else:
+            # Dropped targets must stay owned by the writer until its launch ran:
+            # junk allocations of the same size must not receive the store bytes.
+            del host
+            junk.extend(torch.full_like(t, -9).pin_memory() for t in want)
+    lagging = not held.query()
+    ring_while_lagging = len(writer._slots)
+    torch.cuda.synchronize()
+    for host, want in kept:
+        for h, x in zip(host, want):
+            assert torch.equal(h, x)
+    assert all(bool((j == -9).all()) for j in junk)
+    for k in vars(store):
+        assert torch.equal(getattr(store, k).cpu(), getattr(ref, k)), k
+    if not UNDER_SANITIZER:
+        assert lagging, "the GPU finished its hold before staging ended; inconclusive"
+        assert ring_while_lagging > 4, ring_while_lagging
+
+
+@pytest.mark.parametrize("num_slots,slots", [(1, [0]), (4, [2, 0]), (8, [7, 3, 5])])
+def test_sampler_store_matches_index_copy_and_d2h(num_slots, slots):
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
+        SamplerStoreWriter,
+    )
+
+    torch.manual_seed(2)
+    dev = "cuda"
+    w_new, w_next, d = 6, 6, 5
+
+    def make_store():
+        return SimpleNamespace(
+            new_tokens=_rand_i32(w_new, num_slots, 1).to(dev),
+            next_new_tokens=_rand_i32(w_next, num_slots, 1).to(dev),
+            new_tokens_lens=_rand_i32(num_slots).to(dev),
+            next_draft_tokens=_rand_i32(num_slots, d).to(dev),
+        )
+
+    ref = make_store()
+    out = SimpleNamespace(**{k: v.clone() for k, v in vars(ref).items()})
+    n = len(slots)
+    big_new = _rand_i32(n + 2, w_new).to(dev)
+    o_new = big_new[1 : 1 + n]
+    o_next = _rand_i32(n, w_next).to(dev)
+    o_lens = _rand_i32(n, lo=1, hi=7).to(dev)
+    o_draft = _rand_i32(n, d).to(dev)
+
+    s = torch.tensor(slots, dtype=torch.long).to(dev, non_blocking=True)
+    ref.new_tokens.squeeze(-1).T.index_copy_(0, s, o_new)
+    ref.next_new_tokens.squeeze(-1).T.index_copy_(0, s, o_next)
+    ref.new_tokens_lens.index_copy_(0, s, o_lens)
+    ref.next_draft_tokens.index_copy_(0, s, o_draft)
+    h_ref = [
+        torch.empty_like(t, device="cpu", pin_memory=True)
+        for t in (ref.new_tokens, ref.new_tokens_lens, ref.next_draft_tokens)
+    ]
+    for h, t in zip(h_ref, (ref.new_tokens, ref.new_tokens_lens, ref.next_draft_tokens)):
+        h.copy_(t, non_blocking=True)
+
+    writer = SamplerStoreWriter()
+    assert writer.tensors_supported(out, o_new, o_next, o_lens, o_draft)
+    h_out = [
+        torch.empty_like(t, device="cpu", pin_memory=True)
+        for t in (out.new_tokens, out.new_tokens_lens, out.next_draft_tokens)
+    ]
+    writer.store_and_copy_out(out, slots, o_new, o_next, o_lens, o_draft, *h_out)
+    torch.cuda.synchronize()
+    for k in vars(ref):
+        assert torch.equal(getattr(out, k), getattr(ref, k)), k
+    for a, b in zip(h_out, h_ref):
+        assert torch.equal(a, b)
+
+
+def test_sampler_store_without_sampling_requests():
+    """No sampling request (a prefill chunk): empty outputs, the store is copied out unchanged.
+
+    The empty outputs have the null address; the launcher must accept it.
+    """
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
+        SamplerStoreWriter,
+    )
+
+    torch.manual_seed(3)
+    dev = "cuda"
+    num_slots, w, d = 4, 6, 5
+    store = SimpleNamespace(
+        new_tokens=_rand_i32(w, num_slots, 1).to(dev),
+        next_new_tokens=_rand_i32(w, num_slots, 1).to(dev),
+        new_tokens_lens=_rand_i32(num_slots).to(dev),
+        next_draft_tokens=_rand_i32(num_slots, d).to(dev),
+    )
+    ref = {k: v.clone() for k, v in vars(store).items()}
+    o_new = torch.empty((0, w), dtype=torch.int32, device=dev)
+    o_next = torch.empty((0, w), dtype=torch.int32, device=dev)
+    o_lens = torch.empty((0,), dtype=torch.int32, device=dev)
+    o_draft = torch.empty((0, d), dtype=torch.int32, device=dev)
+    assert o_new.data_ptr() == 0
+    writer = SamplerStoreWriter()
+    h_out = [
+        torch.full_like(t, -1, device="cpu").pin_memory()
+        for t in (store.new_tokens, store.new_tokens_lens, store.next_draft_tokens)
+    ]
+    writer.store_and_copy_out(store, [], o_new, o_next, o_lens, o_draft, *h_out)
+    torch.cuda.synchronize()
+    for k, v in ref.items():
+        assert torch.equal(getattr(store, k), v), k
+    for h, k in zip(h_out, ("new_tokens", "new_tokens_lens", "next_draft_tokens")):
+        assert torch.equal(h, ref[k].cpu()), k
+
+
+def test_pinned_in_place_and_pageable_payload_sources():
+    dev = "cuda"
+    big = torch.zeros(3 * 4096 + 17, dtype=torch.int32, device=dev)  # multi-chunk, in place
+    small = torch.zeros(9, dtype=torch.int32, device=dev)  # pageable -> payload
+    host_big = _rand_i32(big.numel()).pin_memory()
+    host_small = _rand_i32(9)
+    st = _stager()
+    st.begin()
+    st.copy(big, host_big)
+    st.copy(small, host_small)
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(big.cpu(), host_big) and torch.equal(small.cpu(), host_small)
+
+
+def test_packet_full_falls_back_to_native_writes():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.packet import (  # noqa: E501
+        PrepPacket,
+    )
+
+    torch.manual_seed(3)
+    dev = "cuda"
+    st = _stager()
+    st.packet = PrepPacket(max_desc=3, payload_words=64)
+    a = torch.zeros(5000, dtype=torch.int32, device=dev)
+    b = torch.zeros(100, dtype=torch.int32, device=dev)
+    c = _rand_i32(64).to(dev)
+    src = _rand_i32(6, 1, 1).to(dev)
+    host_a = _rand_i32(5000).pin_memory()
+    host_b = _rand_i32(100)
+    c_ref = torch.zeros_like(c)
+    c_ref[:6] = src.flatten() + 3
+    st.begin()
+    st.copy(a, host_a)  # 2 descriptors
+    st.copy(b, host_b)  # payload too small -> native
+    st.gather(c, src, list(range(6)), addend=3)  # 1 descriptor
+    st.zero(c[:0]) if c[:0].numel() else None
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(a.cpu(), host_a) and torch.equal(b.cpu(), host_b)
+    assert torch.equal(c, c_ref)
+    assert st.num_native_copies >= 1
+
+
+@pytest.mark.parametrize("next_n", [1, 6])
+def test_paged_mqa_schedule_host_replica_matches_deepgemm(next_n):
+    """The host replica equals ``get_paged_mqa_logits_metadata`` bit for bit."""
+    import random
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.paged_mqa_schedule import (  # noqa: E501
+        SPLIT_KV_FOR_BLOCK_KV,
+        paged_mqa_logits_schedule,
+    )
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _DG_SCHEDULE_BLOCK_KV
+    from tensorrt_llm.deep_gemm import get_num_sms, get_paged_mqa_logits_metadata
+
+    num_sms = get_num_sms()
+    split_kv = SPLIT_KV_FOR_BLOCK_KV[_DG_SCHEDULE_BLOCK_KV]
+    max_len = (1 << 20) + 14
+    rng = random.Random(next_n)
+    cases = [
+        [k * split_kv + d]
+        for k in range(0, max_len // split_kv + 1)
+        for d in (-1, 0, 1)
+        if 1 <= k * split_kv + d <= max_len
+    ]
+    cases += [[rng.randint(1, max_len) for _ in range(rng.randint(2, 8))] for _ in range(300)]
+    cases += [[rng.randint(1, 4096) for _ in range(rng.randint(2, 8))] for _ in range(300)]
+    mismatches = 0
+    for lens in cases:
+        t = torch.tensor(lens, dtype=torch.int32, device="cuda").view(-1, 1)
+        if next_n > 1:
+            t = t.expand(-1, next_n).contiguous()
+        native = get_paged_mqa_logits_metadata(t, _DG_SCHEDULE_BLOCK_KV, num_sms).cpu().numpy()
+        host = paged_mqa_logits_schedule(lens, next_n, split_kv, num_sms)
+        if native.shape != host.shape or not (native == host).all():
+            mismatches += 1
+            if mismatches <= 3:
+                print("mismatch", lens[:4], native.reshape(-1)[:12], host.reshape(-1)[:12])
+    assert mismatches == 0, f"{mismatches}/{len(cases)} schedules differ"
+
+
+# ---- debug guard ------------------------------------------------------------
+
+
+def test_debug_guard_flags_torch_op_on_staged_destination():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
+        InputPrepHazardError,
+    )
+
+    a = torch.zeros(8, dtype=torch.int32, device="cuda")
+    b = torch.zeros(8, dtype=torch.int32, device="cuda")
+    st = _stager(1)
+    st.begin()
+    a.add_(1)  # before staging: native order is preserved
+    st.values(a, list(range(8)))
+    b.add_(1)  # unrelated buffer
+    _ = a[2:4]  # a view reads no data
+    with pytest.raises(InputPrepHazardError, match="touches a destination staged"):
+        _ = a[2:4] * 2
+    st.discard()
+    # The guard is gone after discard.
+    _ = a * 2
+    torch.cuda.synchronize()
+
+
+def test_debug_guard_flags_triton_kernel_on_staged_destination():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
+        InputPrepHazardError,
+    )
+
+    a = torch.zeros(8, dtype=torch.int32, device="cuda")
+    out = torch.zeros(8, dtype=torch.int32, device="cuda")
+    st = _stager(1)
+    st.begin()
+    st.values(a, list(range(8)))
+    with pytest.raises(InputPrepHazardError, match="triton:_copy_tensor_kernel"):
+        _copy_tensor_kernel[(1,)](a, out, 8, BLOCK=8)
+    st.discard()
+
+
+def test_debug_guard_flags_hidden_writer_at_launch():
+    """A writer the hooks cannot see (direct CUDA launch) is caught before the launch."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
+        InputPrepHazardError,
+    )
+
+    a = torch.zeros(8, dtype=torch.int32, device="cuda")
+    junk = torch.full((8,), 3, dtype=torch.int32, device="cuda")
+    st = _stager(1)
+    st.begin()
+    st.zero(a)
+    _copy_addr_kernel[(1,)](junk.data_ptr(), a.data_ptr(), 8, BLOCK=8)
+    with pytest.raises(InputPrepHazardError, match="written by another operation"):
+        st.commit()
+    assert not st.active
+
+
+@pytest.mark.parametrize("level", [1, 2])
+def test_debug_levels_keep_final_bytes(level):
+    """The guard and the poison do not change what the launch writes."""
+    torch.manual_seed(10)
+    dst = _rand_i32(300).cuda()
+    ref = dst.clone()
+    host = _rand_i32(100).pin_memory()
+    ref[:100].copy_(host, non_blocking=True)
+    ref[100:200].clamp_(min=0)
+    ref[200:].zero_()
+    st = _stager(level)
+    st.begin()
+    st.copy(dst[:100], host)
+    st.clamp_min(dst[100:200], 0)
+    st.zero(dst[200:])
+    st.commit()
+    torch.cuda.synchronize()
+    assert torch.equal(dst, ref)
+
+
+def test_debug_poison_reaches_hidden_reader():
+    """Level 2: a reader the hooks cannot see consumes poison, not stale bytes."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
+        POISON_BYTE,
+    )
+
+    a = torch.full((8,), 7, dtype=torch.int32, device="cuda")
+    out = torch.zeros(8, dtype=torch.int32, device="cuda")
+    st = _stager(2)
+    st.begin()
+    st.values(a, list(range(8)))
+    _copy_addr_kernel[(1,)](a.data_ptr(), out.data_ptr(), 8, BLOCK=8)
+    st.commit()
+    torch.cuda.synchronize()
+    assert a.tolist() == list(range(8))
+    assert bool((out.view(torch.uint8) == POISON_BYTE).all())
+
+
+# ---- support domain: executor loop and preparation eligibility ---------------
+
+
+@pytest.fixture
+def loop_state(monkeypatch):
+    """A fresh executor-loop domain state for the process."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        compat,
+    )
+
+    state = {"declared": False, "reason": "", "logged": False}
+    monkeypatch.setattr(compat, "_LOOP_STATE", state)
+    return state
+
+
+def _store_args(dev="cuda", num_slots=4, w=6, d=5):
+    store = SimpleNamespace(
+        new_tokens=torch.zeros(w, num_slots, 1, dtype=torch.int32, device=dev),
+        next_new_tokens=torch.zeros(w, num_slots, 1, dtype=torch.int32, device=dev),
+        new_tokens_lens=torch.zeros(num_slots, dtype=torch.int32, device=dev),
+        next_draft_tokens=torch.zeros(num_slots, d, dtype=torch.int32, device=dev),
+    )
+    outs = (
+        torch.zeros(1, w, dtype=torch.int32, device=dev),
+        torch.zeros(1, w, dtype=torch.int32, device=dev),
+        torch.zeros(1, dtype=torch.int32, device=dev),
+        torch.zeros(1, d, dtype=torch.int32, device=dev),
+    )
+    return store, outs
+
+
+def test_store_writer_needs_the_validated_loop_domain(loop_state):
+    """The writer keeps the native path until a non-ADP overlap loop declared its domain."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        compat,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
+        SamplerStoreWriter,
+    )
+
+    writer = SamplerStoreWriter()
+    store, outs = _store_args()
+    assert writer.tensors_supported(store, *outs)
+    # No overlap loop yet (also: the non-overlap and pipeline-parallel loops never declare).
+    assert not writer.supported(store, *outs)
+    compat.declare_overlap_loop(False)
+    assert writer.supported(store, *outs)
+    # A later loop with attention DP: final for the process.
+    compat.declare_overlap_loop(True)
+    assert not writer.supported(store, *outs)
+    compat.declare_overlap_loop(False)
+    assert not writer.supported(store, *outs)
+    assert writer.num_calls == 0
+
+
+@pytest.mark.parametrize(
+    "adp, expected",
+    [(False, True), (True, False), (None, False)],
+)
+def test_the_loop_domain_declines_under_attention_dp(loop_state, adp, expected):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        compat,
+    )
+
+    compat.declare_overlap_loop(adp)  # None: a caller without the setting
+    assert loop_state["declared"]
+    assert compat.loop_domain() is expected
+
+
+def _eligibility_case(**overrides):
+    """Arguments of ``PrepStager._ineligible_reason`` for a steady served decode step."""
+    mode = SimpleNamespace(use_one_engine=lambda: True)
+    cfg = SimpleNamespace(
+        spec_config=SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True),
+        disable_overlap_scheduler=overrides.pop("disable_overlap_scheduler", False),
+        enable_attention_dp=overrides.pop("enable_attention_dp", False),
+        use_mrope=False,
+    )
+    runner = SimpleNamespace(
+        _config=cfg,
+        use_beam_search=False,
+        mapping=SimpleNamespace(has_cp_helix=lambda: False, pp_size=overrides.pop("pp_size", 1)),
+        guided_decoder=None,
+        _lora=SimpleNamespace(cuda_graph_manager=None),
+    )
+    rows = overrides.pop("rows", [SimpleNamespace(is_dummy=False, py_batch_idx=0)])
+    requests = SimpleNamespace(num_context_requests=0, generation_requests=rows)
+    DSAtrtllmAttentionMetadata = type("DSAtrtllmAttentionMetadata", (), {})
+    manager = SimpleNamespace(index_mapper=object(), _use_per_layer_page_tables=False)
+    attn = DSAtrtllmAttentionMetadata()
+    attn.kv_cache_manager = manager
+    attn.draft_kv_cache_manager = None
+    args = dict(
+        runner=runner,
+        requests=requests,
+        attn_metadata=attn,
+        spec_metadata=object(),
+        new_tensors_device=SimpleNamespace(next_draft_tokens=object()),
+        graph_replay=True,
+        is_dummy=False,
+        enable_spec_decode=True,
+        promoted=False,
+    )
+    args.update(overrides)
+    return args
+
+
+def test_preparation_eligibility_declines_attention_dp_and_new_arrivals():
+    """Prep stays native under attention DP and for rows without previous device tensors.
+
+    A request that just arrived (also a disaggregated generation request whose
+    KV cache came from a context server) has no previous iteration's device
+    tensors (``py_batch_idx`` None), so its first preparation is native.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.stager import (  # noqa: E501
+        PrepStager,
+    )
+
+    reason = PrepStager._ineligible_reason
+    assert reason(**_eligibility_case()) is None
+    assert (
+        reason(**_eligibility_case(enable_attention_dp=True))
+        == "overlap off, attention DP or MRoPE"
+    )
+    assert reason(**_eligibility_case(disable_overlap_scheduler=True)) == (
+        "overlap off, attention DP or MRoPE"
+    )
+    arrival = [
+        SimpleNamespace(is_dummy=False, py_batch_idx=0),
+        SimpleNamespace(is_dummy=False, py_batch_idx=None),
+    ]
+    assert reason(**_eligibility_case(rows=arrival)) == "dummy row or row without previous tensors"
+    dummy = [SimpleNamespace(is_dummy=True, py_batch_idx=0)]
+    assert reason(**_eligibility_case(rows=dummy)) == "dummy row or row without previous tensors"
+    assert reason(**_eligibility_case(new_tensors_device=None)) == (
+        "no previous speculative device tensors"
+    )
+
+
+def test_preparation_eligibility_declines_pipeline_parallelism():
+    """deferred_input_prep is not validated with the pipeline-parallel executor loop: prep stays native."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.stager import (  # noqa: E501
+        PrepStager,
+    )
+
+    reason = PrepStager._ineligible_reason
+    assert reason(**_eligibility_case(pp_size=1)) is None
+    assert reason(**_eligibility_case(pp_size=2)) == "pipeline parallelism"
+
+
+def test_prep_declines_without_pinned_host_memory(monkeypatch):
+    """Confidential Compute (``prefer_pinned()`` False) keeps every preparation native.
+
+    The staged page-table copy reads the pinned host table; with pageable host
+    tables it would raise in the middle of a preparation instead of declining.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        stager as stager_mod,
+    )
+
+    monkeypatch.setattr(stager_mod, "_host_memory_pinned", lambda: False)
+    st = stager_mod.PrepStager()
+    got = st.begin_if_eligible(
+        **_eligibility_case(),
+    )
+    assert got is None
+    assert st.last_reason == "host page tables are not pinned (prefer_pinned)"
+    assert not st.active
+
+
+def test_prep_decisions_are_logged_once_per_kind(monkeypatch):
+    """Every kind of preparation decision is logged once, later ones only counted."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
+        stager as stager_mod,
+    )
+
+    lines = []
+    monkeypatch.setattr(dispatch, "_log_info", lines.append)
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(stager_mod, "_host_memory_pinned", lambda: True)
+    st = stager_mod.PrepStager()
+    for _ in range(3):
+        assert st.begin_if_eligible(**_eligibility_case(pp_size=2)) is None
+    for _ in range(2):
+        assert st.begin_if_eligible(**_eligibility_case(new_tensors_device=None)) is None
+    assert st.decisions == {
+        "pipeline parallelism": 3,
+        "no previous speculative device tensors": 2,
+    }
+    assert lines == [
+        "[glm_moe_dsa] deferred_input_prep/prep[pipeline parallelism]: native fallback (pipeline parallelism)",
+        "[glm_moe_dsa] deferred_input_prep/prep[no previous speculative device tensors]: native fallback "
+        "(no previous speculative device tensors)",
+    ]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py
new file mode 100644
index 000000000000..c48abebdbaa9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py
@@ -0,0 +1,24 @@
+# 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.
+"""Make the deferred_norm helpers and the absorbed_query / post_moe_hop fixtures importable under
+``--import-mode=importlib``."""
+
+import sys
+from pathlib import Path
+
+_HERE = Path(__file__).resolve().parent
+for _path in (str(_HERE), str(_HERE.parent / "absorbed_query"), str(_HERE.parent / "post_moe_hop")):
+    if _path not in sys.path:
+        sys.path.insert(0, _path)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py
new file mode 100644
index 000000000000..3a103b099ceb
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py
@@ -0,0 +1,299 @@
+# 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.
+"""Shared helpers of the deferred_norm tests and benchmarks.
+
+* ``hop_sumsq_partials``: torch emulation of the deferred_norm / post_moe_hop's square-sum
+  partials (per thread: 8 columns, sequential fp32 sum of bf16-rounded
+  squares; warp butterfly 16, 8, 4, 2, 1; one value per warp), bit-exact by
+  construction because every step is one IEEE fp32 operation.
+* ``fp64_oracle``: exact projection of the un-rounded RMSNorm output.
+* ``ProxyProducer``: a CuTe DSL stand-in for a hop kernel. It is launched
+  with programmatic dependent launch, releases its dependents at entry (or
+  at the end), spins for a given time and only then writes the packet
+  tensors; a dependent that reads them before its dependency wait would see
+  stale data.
+* error statistics (normalized RMS error, bf16 ulp distance).
+"""
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import dsl_user_op
+
+H = 6144
+EPS = 1e-5
+HOP_CTAS = 8
+HOP_THREADS = H // (8 * HOP_CTAS)  # 96
+HOP_WARPS = HOP_THREADS // 32  # 3
+PARTIALS = HOP_CTAS * HOP_WARPS  # 24
+
+
+def hop_sumsq_partials(x: torch.Tensor) -> torch.Tensor:
+    """fp32 ``[R, 24]`` partials exactly as the hop computes them from bf16 ``x``."""
+    rows = x.shape[0]
+    xf = x.float().view(rows, HOP_CTAS, HOP_WARPS, 32, 8)
+    sq = (xf * xf).bfloat16().float()
+    s = torch.zeros(rows, HOP_CTAS, HOP_WARPS, 32, dtype=torch.float32, device=x.device)
+    for i in range(8):
+        s = s + sq[..., i]
+    lanes = torch.arange(32, device=x.device)
+    for mask in (16, 8, 4, 2, 1):
+        s = s + s[..., lanes ^ mask]
+    return s[..., 0].reshape(rows, PARTIALS).contiguous()
+
+
+def square_sum_in_order(partials: torch.Tensor) -> torch.Tensor:
+    """fp32 sum of the partials in index order (the consumer's order)."""
+    s = torch.zeros(partials.shape[0], dtype=torch.float32, device=partials.device)
+    for i in range(partials.shape[1]):
+        s = s + partials[:, i]
+    return s
+
+
+def make_frag_packet(x: torch.Tensor, gamma: torch.Tensor, eps: float = EPS):
+    """Packet of the hop computed in torch (bit-exact ``a`` and partials).
+
+    ``sumsq`` is a view of an 8-row buffer (``sumsq_full(packet)`` is the
+    full buffer) so a proxy producer can copy whole 8-row units.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHidden,
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        frag_from_rows,
+    )
+
+    rows = x.shape[0]
+    a = (x.float() * gamma.float()).bfloat16()
+    ss = torch.zeros(8, PARTIALS, dtype=torch.float32, device=x.device)
+    ss[:rows] = hop_sumsq_partials(x)
+    return DeferredNormHidden(frag_from_rows(a), ss[:rows], x, gamma, eps, materialize)
+
+
+def sumsq_full(p):
+    """The 8-row partial buffer behind a test packet's ``sumsq`` view (test
+    packets allocate 8 rows; inference-mode views do not keep ``_base``)."""
+    return torch.as_strided(p.sumsq, (8, PARTIALS), (PARTIALS, 1))
+
+
+def empty_like_packet(p):
+    """Same-shaped packet with fresh (uninitialized) a_frag / partials buffers."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHidden,
+    )
+
+    ss = torch.empty(8, PARTIALS, dtype=torch.float32, device=p.sumsq.device)
+    return DeferredNormHidden(
+        torch.empty_like(p.a_frag), ss[: p.num_rows], p.residual, p.gamma, p.eps, p.materializer
+    )
+
+
+def fp64_oracle(
+    x: torch.Tensor, gamma: torch.Tensor, weight: torch.Tensor, eps: float = EPS
+) -> torch.Tensor:
+    """``y* = r* · Σ_k x_k γ_k W_jk`` in fp64 (``r*`` from exact squares)."""
+    xd = x.double()
+    r = torch.rsqrt(xd.pow(2).mean(-1, keepdim=True) + eps)
+    return (xd * gamma.double() * r) @ weight.double().t()
+
+
+def norm_rms_err(y: torch.Tensor, ref: torch.Tensor) -> float:
+    ref = ref.double()
+    return float((y.double() - ref).norm() / ref.norm().clamp_min(1e-300))
+
+
+def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    def ordered(t):
+        v = t.contiguous().view(torch.int16).to(torch.int32)
+        return torch.where(v < 0, -(v + 32768), v)
+
+    return (ordered(a) - ordered(b)).abs()
+
+
+def compare(cand: torch.Tensor, native: torch.Tensor, oracle: torch.Tensor) -> dict:
+    ulp = bf16_ulp_distance(cand, native)
+    return {
+        "allclose_native": bool(torch.allclose(cand.float(), native.float(), rtol=0.1, atol=1e-2)),
+        "e_cand": norm_rms_err(cand, oracle),
+        "e_native": norm_rms_err(native, oracle),
+        "max_ulp_vs_native": int(ulp.max()),
+        "frac_diff_vs_native": float((ulp > 0).float().mean()),
+        "max_abs_vs_native": float((cand.float() - native.float()).abs().max()),
+    }
+
+
+def within_budget(stats: dict) -> bool:
+    return stats["allclose_native"] and stats["e_cand"] <= 1.25 * stats["e_native"] + 2.0**-12
+
+
+def _asm(res, args, text, cons, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _spin_ns(ns: Int64, sleep_ns: int = 0, *, loc=None, ip=None):
+    nap = f"\t@p nanosleep.u32 {int(sleep_ns)};\n" if sleep_ns else ""
+    _asm(
+        None,
+        [ns.ir_value()],
+        "{\n\t.reg .u64 t0, t;\n\t.reg .pred p;\n\tmov.u64 t0, %globaltimer;\n\t"
+        "GLM_MOE_DSA_DEFERRED_NORM_SPIN:\n\tmov.u64 t, %globaltimer;\n\tsub.u64 t, t, t0;\n\t"
+        "setp.lt.u64 p, t, $0;\n" + nap + "\t@p bra GLM_MOE_DSA_DEFERRED_NORM_SPIN;\n\t}",
+        "l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _copy16(dst: Int64, src: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [dst.ir_value(), src.ir_value()],
+        "{\n\t.reg .b32 a, b, c, d;\n\tld.global.v4.b32 {a, b, c, d}, [$1];\n\t"
+        "st.global.v4.b32 [$0], {a, b, c, d};\n\t}",
+        "l,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _copy4(dst: Int64, src: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [dst.ir_value(), src.ir_value()],
+        "{\n\t.reg .b32 a;\n\tld.global.b32 a, [$1];\n\tst.global.b32 [$0], a;\n\t}",
+        "l,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+class _ProxyKernel:
+    def __init__(self, trigger_at_start: bool, sleep_ns: int = 0, carveout: int = -1):
+        self.trigger_at_start = trigger_at_start
+        self.sleep_ns = sleep_ns
+        self.carveout = carveout
+
+    @cute.jit
+    def __call__(
+        self,
+        act_src: Int64,
+        act_dst: Int64,
+        ss_src: Int64,
+        ss_dst: Int64,
+        delay_ns: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        stream: CUstream,
+    ):
+        if cutlass.const_expr(self.carveout >= 0):
+            self.kernel(act_src, act_dst, ss_src, ss_dst, delay_ns, anchor).launch(
+                grid=[HOP_CTAS, rows, 1],
+                block=[HOP_THREADS, 1, 1],
+                stream=stream,
+                use_pdl=True,
+                preferred_smem_carveout=self.carveout,
+            )
+        else:
+            self.kernel(act_src, act_dst, ss_src, ss_dst, delay_ns, anchor).launch(
+                grid=[HOP_CTAS, rows, 1], block=[HOP_THREADS, 1, 1], stream=stream, use_pdl=True
+            )
+
+    @cute.kernel
+    def kernel(
+        self,
+        act_src: Int64,
+        act_dst: Int64,
+        ss_src: Int64,
+        ss_dst: Int64,
+        delay_ns: Int64,
+        anchor: cute.Tensor,
+    ):
+        tid, _, _ = cute.arch.thread_idx()
+        cta, row, _ = cute.arch.block_idx()
+        cute.arch.griddepcontrol_wait()
+        if cutlass.const_expr(self.trigger_at_start):
+            cute.arch.griddepcontrol_launch_dependents()
+        _spin_ns(delay_ns, self.sleep_ns)
+        off = Int64(row) * Int64(H * 2) + Int64((cta * Int32(HOP_THREADS) + tid) * Int32(16))
+        _copy16(act_dst + off, act_src + off)
+        if tid < Int32(HOP_WARPS):
+            o4 = Int64(row) * Int64(PARTIALS * 4) + Int64((cta * Int32(HOP_WARPS) + tid) * 4)
+            _copy4(ss_dst + o4, ss_src + o4)
+        if cutlass.const_expr(not self.trigger_at_start):
+            cute.arch.griddepcontrol_launch_dependents()
+
+
+_PROXY = {}
+
+
+def proxy_producer(
+    act_src, act_dst, ss_src, ss_dst, delay_ns: int, trigger_at_start: bool = True, rows: int = None
+):
+    """Launch the proxy producer on the current stream (PDL).
+
+    Copies ``rows`` x 12,288 bytes of ``act_src`` and ``rows`` x 24 fp32 of
+    ``ss_src`` (``rows`` defaults to ``act_src.shape[0]``; an ``a_frag``
+    buffer is 8 rows' worth).
+    """
+    import os
+
+    sleep_ns = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_PROXY_SLEEP_NS", "0"))
+    carveout = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_PROXY_CARVEOUT", "-1"))
+    key = (trigger_at_start, sleep_ns, carveout)
+    if key not in _PROXY:
+        from cutlass.cute import runtime as cute_runtime
+
+        z = Int64(0)
+        anchor = cute_runtime.make_fake_compact_tensor(
+            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
+        )
+        _PROXY[key] = cute.compile(
+            _ProxyKernel(trigger_at_start, sleep_ns, carveout),
+            z,
+            z,
+            z,
+            z,
+            z,
+            anchor,
+            Int32(1),
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+    _PROXY[key](
+        act_src.data_ptr(),
+        act_dst.data_ptr(),
+        ss_src.data_ptr(),
+        ss_dst.data_ptr(),
+        int(delay_ns),
+        ss_dst.view(torch.int32).view(-1),
+        rows if rows is not None else act_src.shape[0],
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py
new file mode 100644
index 000000000000..10382e53ddbb
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py
@@ -0,0 +1,495 @@
+# 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.
+"""deferred_norm packet -> layer L+1 DSA attention through the target's adopted attention
+(absorbed_query, draft_absorbed_query, shared_indexer_absorbed_query, post_moe_hop and
+deferred_norm; single GPU; generation steps run as decode steps).
+
+Uses the absorbed_query integration harness (real weights of the full-indexer layer 6
+and the shared layer 7 at TP8 rank 0, six requests natively prefilled with
+2,600 tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6 (MTP verify), both call
+paths, ``on_update_kv_lens()`` before every step). Each layer gets its own
+input: residual-like rows ``x`` and that layer's real ``input_layernorm``
+gain; the deferred_norm packet carries ``a = bf16(x * gamma)`` and the hop-order
+square-sum partials, the native input is ``h = materialize(packet)`` (the post_moe_hop
+hop's output bit for bit). The target's deferred-norm scope is ``shared``: the hop defers
+the norm only into shared-indexer layers, so a deferred_norm step gives layer 7 the packet
+and layer 6 the normalized ``h`` (``SERVED_LAYERS``, ``_served``).
+
+* deviation (report + sanity bound): relative RMS vs the native
+  step of the attention inputs recorded at ``mla_rope_generation`` and the
+  indexer, top-k set agreement per row (identical rows, mean Jaccard), and
+  the attention output. Reported next to two native-only references: a
+  rerun of the native step (run-to-run noise) and the *rounding control*,
+  the native step on ``bf16(bf16(x * gamma) * r)`` (the normalized hidden
+  state with the gain applied first, a one-rounding perturbation of the
+  native input in native code only). Bound: output relRMS <= max(2**-7,
+  2 * control) and mean Jaccard >= 0.98 on every layer (2**-7 alone is
+  below the control: the attention output amplifies rounding-level input
+  differences through the FP8 query and the sparse softmax);
+* fallback exactness: with the deferred_norm projection declined (``_deferred_norm_projection =
+  False``) the packet is materialized and every cache, recorded input and
+  output equals the native step bit for bit;
+* engagement: the deferred_norm projection serves layer 7 and is never asked on layer 6,
+  shared_indexer_absorbed_query serves layer 7 under
+  ``absorbed_query,shared_indexer_absorbed_query,deferred_norm`` and absorbed_query serves layer 6;
+* CUDA graph: a captured deferred_norm step replays (1, 1, 20 times with fresh
+  inputs) bit-identically to the eager deferred_norm step;
+* a projection compile error on layer 7 propagates; on layer 6 the projection declines
+  before any compile;
+* negative control: a packet whose ``a_frag`` lacks the gain (``bf16(x)``
+  instead of ``bf16(x * gamma)``) violates the bound.
+
+Fallback and graph comparisons of attention outputs whose inputs and caches
+are bitwise equal allow 4 bf16 ulps of the output scale: the native sparse
+attention is not bitwise run-to-run stable.
+
+Every arm (native reference included) also runs the fast paths in
+``GLM_MOE_DSA_TEST_BASE_FEATURES`` (e.g. ``rope_kv_append``), as the absorbed_query harness
+does.
+"""
+
+import os
+import sys
+
+import absorbed_query_fixture
+import deferred_norm_utils as U
+import pytest
+import test_absorbed_query_attention_integration as aq_integration
+import torch
+
+LAYERS = aq_integration.LAYERS
+SHARED = absorbed_query_fixture.SHARED_LAYER
+FULL = absorbed_query_fixture.FULL_LAYER
+LAYOUTS = aq_integration.LAYOUTS
+# The layers the deferred_norm projection serves: the shared-indexer ones.
+SERVED_LAYERS = (SHARED,)
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+
+import plan_testlib  # noqa: E402
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+
+def _set_features(features: str):
+    """Run exactly the fast paths ``features`` plus ``GLM_MOE_DSA_TEST_BASE_FEATURES``.
+
+    The latter are the neighbour paths every arm runs on, e.g.
+    ``rope_kv_append`` on an integration image (absorbed_query harness rule).
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
+        query_proj,
+    )
+
+    query_proj.SERVED_FORM = "auto"
+    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
+
+
+def _served(packets):
+    """The inputs of a deferred_norm step: the packet on the layers the projection serves, the
+    normalized hidden state (the post-MoE hop's output) on the others."""
+    return {layer: p if layer in SERVED_LAYERS else p.materialize() for layer, p in packets.items()}
+
+
+def _reset_runners(h):
+    for layer in LAYERS:
+        for name in ("_absorbed_query", "_shared_indexer_query", "_deferred_norm_projection"):
+            setattr(h.attn[layer], name, None)
+
+
+class _DeferredNormHarness(aq_integration._Harness):
+    """absorbed_query harness whose layers take per-layer inputs (tensor or deferred_norm packet)."""
+
+    def __init__(self):
+        super().__init__()
+        self.gammas = {
+            layer: absorbed_query_fixture.load_input_layernorm(layer).cuda() for layer in LAYERS
+        }
+
+    def forward(self, layer, positions, hidden, meta, path):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+            is_deferred,
+        )
+
+        view = hidden.residual if is_deferred(hidden) else hidden
+        attn = self.attn[layer]
+        if path == "forward_impl":
+            out = attn._create_outputs(view, meta)
+            attn.forward_impl(positions, hidden, meta, attn_output=out)
+            return out[0]
+        import weakref
+
+        from tensorrt_llm._torch.utils import model_extra_attrs
+
+        extra = self.model_config.extra_attrs
+        extra["attention_metadata"] = weakref.ref(meta)
+        with model_extra_attrs(extra):
+            out = [torch.ops.trtllm.create_mla_outputs(view, attn.layer_idx_str)]
+            attn._forward_custom_op(hidden, positions, out, None)
+        return out[0]
+
+    def step(self, meta, positions, hidden, path):
+        """Per-layer inputs when ``hidden`` is a dict (layer -> tensor or
+        packet); otherwise the same tensor for every layer (native prefill). A
+        generation-only step runs as a decode step."""
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+            dispatch,
+        )
+
+        phase = dispatch.Phase.DECODE if meta.num_contexts == 0 else dispatch.Phase.PREFILL
+        meta.on_update_kv_lens()
+        with dispatch.step(phase):
+            return {
+                layer: self.forward(
+                    layer,
+                    positions,
+                    hidden[layer] if isinstance(hidden, dict) else hidden,
+                    meta,
+                    path,
+                )
+                for layer in LAYERS
+            }
+
+    def packets(self, layout, seed, scale=1.0):
+        nreq, ntok = LAYOUTS[layout]
+        out = {}
+        for layer in LAYERS:
+            g = torch.Generator(device="cuda").manual_seed(seed * 10 + layer)
+            x = torch.randn(nreq * ntok, U.H, generator=g, device="cuda") * 2.0
+            x = (x * scale).bfloat16()
+            out[layer] = U.make_frag_packet(x, self.gammas[layer])
+        pos = (
+            (aq_integration.CTX + torch.arange(ntok, dtype=torch.int32, device="cuda"))
+            .repeat(nreq)
+            .unsqueeze(0)
+        )
+        return out, pos
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = _DeferredNormHarness()
+    yield h
+    _set_features("")
+
+
+def _run(h, features, layout, path, inputs, pos, snap):
+    _set_features(features)
+    h.restore(snap)
+    out, rec = h.recorded_step(layout, path, inputs, pos)
+    cache = h.snapshot()
+    return out, rec, cache
+
+
+def _jaccard_rows(a, b):
+    rows = []
+    for ra, rb in zip(a.tolist(), b.tolist()):
+        sa = {v for v in ra if v >= 0}
+        sb = {v for v in rb if v >= 0}
+        union = len(sa | sb)
+        rows.append(1.0 if union == 0 else len(sa & sb) / union)
+    return rows
+
+
+def _output_stats(native, cand, layout):
+    """Deviation statistics per layer."""
+    n_out, n_rec, _ = native
+    c_out, c_rec, _ = cand
+    stats = {}
+    for layer in LAYERS:
+        s = {"out_rel_rms": aq_integration._rel_rms(c_out[layer], n_out[layer])}
+        nr, cr = n_rec[layer], c_rec[layer]
+        if "absorbed_q" in nr:
+            s["absorbed_q_rel_rms"] = aq_integration._rel_rms(cr["absorbed_q"], nr["absorbed_q"])
+        if "quant_q" in nr:
+            s["quant_q_bytes_diff_frac"] = float(
+                (cr["quant_q"].view(torch.uint8) != nr["quant_q"].view(torch.uint8)).float().mean()
+            )
+        if "indexer_weights" in nr:
+            s["indexer_weights_rel_rms"] = aq_integration._rel_rms(
+                cr["indexer_weights"], nr["indexer_weights"]
+            )
+            s["indexer_q_fp8_bytes_diff_frac"] = float(
+                (cr["indexer_q_fp8"].view(torch.uint8) != nr["indexer_q_fp8"].view(torch.uint8))
+                .float()
+                .mean()
+            )
+        if "topk" in nr:
+            jac = _jaccard_rows(cr["topk"], nr["topk"])
+            s["topk_rows_identical"] = sum(1 for j in jac if j == 1.0) / len(jac)
+            s["topk_mean_jaccard"] = sum(jac) / len(jac)
+        stats[layer] = s
+    return stats
+
+
+def _control_inputs(packets):
+    """Rounding control: ``bf16(bf16(x * gamma) * r)`` per layer (native code
+    consumes it; ``r`` in torch from the hop-order square sum)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    out = {}
+    for layer, p in packets.items():
+        rows, hidden = p.residual.shape
+        a = rows_from_frag(p.a_frag, rows, hidden)
+        r = torch.rsqrt(U.square_sum_in_order(p.sumsq) / hidden + p.eps)
+        out[layer] = (a.float() * r.unsqueeze(-1)).bfloat16()
+    return out
+
+
+def _check_output_bound(stats, control):
+    fails = []
+    for layer, s in stats.items():
+        bound = max(2.0**-7, 2.0 * control[layer]["out_rel_rms"])
+        if s["out_rel_rms"] > bound:
+            fails.append(f"L{layer} output relRMS {s['out_rel_rms']:.3e} > {bound:.3e}")
+        if s.get("topk_mean_jaccard", 1.0) < 0.98:
+            fails.append(f"L{layer} top-k mean Jaccard {s['topk_mean_jaccard']:.4f} < 0.98")
+    return fails
+
+
+@pytest.mark.parametrize(
+    "features",
+    [
+        "absorbed_query,post_moe_hop,deferred_norm",
+        "absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm",
+    ],
+)
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_deferred_norm_packet_into_attention(harness, layout, path, features):
+    h = harness
+    packets, pos = h.packets(layout, seed=3)
+    native_in = {layer: p.materialize() for layer, p in packets.items()}
+    snap = h.snapshot()
+    try:
+        base = ",".join(f for f in features.split(",") if f != "deferred_norm")
+        _reset_runners(h)
+        native = _run(h, base, layout, path, native_in, pos, snap)
+        _reset_runners(h)
+        cand = _run(h, features, layout, path, _served(packets), pos, snap)
+        for layer in LAYERS:
+            engaged = h.attn[layer]._deferred_norm_projection not in (None, False)
+            assert engaged == (layer in SERVED_LAYERS), f"deferred_norm on L{layer}: {engaged}"
+        if "shared_indexer_absorbed_query" in features:
+            assert h.attn[SHARED]._shared_indexer_query not in (None, False)
+        assert h.attn[FULL]._absorbed_query not in (None, False)
+        _reset_runners(h)
+        rerun = _run(h, base, layout, path, native_in, pos, snap)
+        _reset_runners(h)
+        control = _run(h, base, layout, path, _control_inputs(packets), pos, snap)
+        stats = _output_stats(native, cand, layout)
+        noise = _output_stats(native, rerun, layout)
+        ctrl = _output_stats(native, control, layout)
+        tag = f"layout={layout} path={path} features={features}"
+        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_BOUND {tag} {stats}", flush=True)
+        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_NOISE {tag} {noise}", flush=True)
+        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_CONTROL {tag} {ctrl}", flush=True)
+        fails = _check_output_bound(stats, ctrl)
+        assert not fails, fails
+    finally:
+        h.restore(snap)
+        _set_features("")
+
+
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_fallback_materialize_is_exact(harness, layout, path):
+    """deferred_norm projection declined -> materialized packet -> the absorbed_query path, bit for
+    bit equal to the native-input step."""
+    h = harness
+    packets, pos = h.packets(layout, seed=5)
+    native_in = {layer: p.materialize() for layer, p in packets.items()}
+    snap = h.snapshot()
+    try:
+        _reset_runners(h)
+        n_out, n_rec, n_cache = _run(h, "absorbed_query", layout, path, native_in, pos, snap)
+        _reset_runners(h)
+        for layer in LAYERS:
+            h.attn[layer]._deferred_norm_projection = False
+        _set_features("absorbed_query,post_moe_hop,deferred_norm")
+        h.restore(snap)
+        c_out, c_rec = h.recorded_step(layout, path, packets, pos)
+        c_cache = h.snapshot()
+        for key in n_cache:
+            assert torch.equal(c_cache[key], n_cache[key]), key
+        for layer in LAYERS:
+            for key, val in n_rec[layer].items():
+                if key in ("topk", "sparse_indices"):
+                    continue  # native top-k order is not run-to-run stable
+                assert torch.equal(c_rec[layer][key], val), (layer, key)
+            if "topk" in n_rec[layer]:
+                assert min(_jaccard_rows(c_rec[layer]["topk"], n_rec[layer]["topk"])) == 1.0
+        noise_runs = [n_out]
+        _reset_runners(h)
+        noise_runs.append(_run(h, "absorbed_query", layout, path, native_in, pos, snap)[0])
+        for layer in LAYERS:
+            diff = (c_out[layer].float() - n_out[layer].float()).abs().max().item()
+            noise = (noise_runs[1][layer].float() - n_out[layer].float()).abs().max().item()
+            ulp = aq_integration._bf16_ulp(n_out[layer].float().abs().max().item())
+            assert diff <= max(2 * noise, 4 * ulp), (layer, diff, noise, ulp)
+    finally:
+        h.restore(snap)
+        _set_features("")
+
+
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_cuda_graph_replay(harness, layout):
+    h = harness
+    snap = h.snapshot()
+    try:
+        _set_features("absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm")
+        _reset_runners(h)
+        packets, pos = h.packets(layout, seed=7)
+        static = {}
+        for layer, p in _served(packets).items():
+            if layer in SERVED_LAYERS:
+                st = U.empty_like_packet(p)
+                st.a_frag.copy_(p.a_frag)
+                st.sumsq.copy_(p.sumsq)
+            else:
+                st = p.clone()
+            static[layer] = st
+        static_pos = pos.clone()
+        meta = h.gen_metadata(layout, cuda_graph=True)
+        h.restore(snap)
+        h.step(meta, static_pos, static, "custom_op")
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        h.restore(snap)
+        with torch.cuda.graph(graph):
+            g_out = h.step(meta, static_pos, static, "custom_op")
+        for trial, replays in enumerate((1, 1, 20)):
+            fresh, _ = h.packets(layout, seed=100 + trial)
+            fresh = _served(fresh)
+            for layer in LAYERS:
+                if layer in SERVED_LAYERS:
+                    static[layer].a_frag.copy_(fresh[layer].a_frag)
+                    static[layer].sumsq.copy_(fresh[layer].sumsq)
+                else:
+                    static[layer].copy_(fresh[layer])
+            for _ in range(replays):
+                h.restore(snap)
+                graph.replay()
+            torch.cuda.synchronize()
+            got = {layer: g_out[layer].clone() for layer in LAYERS}
+            got_cache = h.snapshot()
+            h.restore(snap)
+            eager = h.gen_step(
+                layout, "custom_op", fresh, pos, meta=h.gen_metadata(layout, cuda_graph=True)
+            )
+            eager_cache = h.snapshot()
+            for key in eager_cache:
+                assert torch.equal(got_cache[key], eager_cache[key]), (trial, key)
+            for layer in LAYERS:
+                diff = (got[layer].float() - eager[layer].float()).abs().max().item()
+                ref = eager[layer].float().abs().max().item()
+                # Native sparse attention is not bitwise run-to-run stable:
+                # bound by 4 bf16 ulps of the output scale.
+                assert diff <= 4 * aq_integration._bf16_ulp(ref), (trial, layer, diff)
+    finally:
+        h.restore(snap)
+        _set_features("")
+
+
+@torch.inference_mode()
+def test_projection_compile_error_is_not_masked(harness, monkeypatch):
+    """A projection compile error in the consumer propagates (it invalidates
+    the run) instead of a rank-local decline to the materialized path."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        projection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.query_entry import (  # noqa: E501
+        deferred_norm_dsa_proj,
+    )
+
+    h = harness
+    packets, pos = h.packets("1req_x6tok", seed=11)
+
+    def _fail(self):
+        raise RuntimeError("injected projection compile failure")
+
+    monkeypatch.setattr(projection.DeferredNormProjection, "precompile", _fail)
+    try:
+        _set_features("absorbed_query,post_moe_hop,deferred_norm")
+        _reset_runners(h)
+        for layer in SERVED_LAYERS:
+            with pytest.raises(RuntimeError, match="injected projection compile failure"):
+                deferred_norm_dsa_proj(h.attn[layer], pos, packets[layer], None)
+            assert h.attn[layer]._deferred_norm_projection is None, (
+                "a failed runner must not be cached"
+            )
+        # A full-indexer layer keeps its native projection: declined before any compile.
+        assert deferred_norm_dsa_proj(h.attn[FULL], pos, packets[FULL], None) is None
+        assert h.attn[FULL]._deferred_norm_projection is False
+    finally:
+        _reset_runners(h)
+        _set_features("")
+
+
+@torch.inference_mode()
+def test_negative_control_missing_gain(harness):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHidden,
+        materialize,
+    )
+
+    h = harness
+    layout = "1req_x6tok"
+    packets, pos = h.packets(layout, seed=9)
+    native_in = {layer: p.materialize() for layer, p in packets.items()}
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        frag_from_rows,
+    )
+
+    bad = {
+        layer: DeferredNormHidden(
+            frag_from_rows(p.residual), p.sumsq, p.residual, p.gamma, U.EPS, materialize
+        )
+        for layer, p in packets.items()
+    }
+    snap = h.snapshot()
+    try:
+        _reset_runners(h)
+        native = _run(h, "absorbed_query", layout, "custom_op", native_in, pos, snap)
+        _reset_runners(h)
+        cand = _run(
+            h, "absorbed_query,post_moe_hop,deferred_norm", layout, "custom_op", bad, pos, snap
+        )
+        _reset_runners(h)
+        control = _run(
+            h, "absorbed_query", layout, "custom_op", _control_inputs(packets), pos, snap
+        )
+        stats = _output_stats(native, cand, layout)
+        ctrl = _output_stats(native, control, layout)
+        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_NEGATIVE {stats}", flush=True)
+        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_NEGATIVE_CONTROL {ctrl}", flush=True)
+        assert _check_output_bound(stats, ctrl), (
+            "the output bound must flag a packet without the gain"
+        )
+    finally:
+        h.restore(snap)
+        _set_features("")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py
new file mode 100644
index 000000000000..c22bc6a7eb39
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py
@@ -0,0 +1,425 @@
+# 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.
+"""Diagnosis of deferred_norm's attention-output deviation next to native-only controls.
+
+Same single-GPU harness and metrics as ``test_deferred_norm_attention_integration.py``
+(real layer-6 / layer-7 weights at TP8 rank 0, six requests natively
+prefilled with 2,600 tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6, both
+call paths, ``absorbed_query`` and ``absorbed_query,shared_indexer_absorbed_query``). Every arm runs
+the same generation step from the same cache snapshot and is compared with the native reference step
+on ``h = materialize(packet)`` (the post_moe_hop's normalized output, bit for bit).
+
+Arms:
+
+* ``deferred_norm``: the deferred_norm packet through the deferred_norm projection (the candidate)
+  on the layers the projection serves (``dn_integration.SERVED_LAYERS``: the target's deferred-norm
+  scope is ``shared``), the normalized ``h`` on the full-indexer layer, as a served step
+  gives them.
+* ``rerun``: the native step again (run-to-run noise of native attention).
+* ``n1_rmsnorm``: native code on ``h'`` = TRT-LLM's ``RMSNorm`` module
+  (FlashInfer RMSNorm) applied to the same ``x`` and gain: a second in-tree
+  implementation of the same normalization.
+* ``n2_cublas``: native code with the ``kv_a_proj_with_mqa`` GEMM through
+  ``Linear.apply_linear`` (``cublas_mm``), the path ``DeepseekV3Linear``
+  takes for more than 16 tokens instead of ``dsv3_fused_a_gemm_op``; the
+  input ``h`` is unchanged.
+* ``n3_round``: native code on ``bf16(bf16(x * gamma) * r)`` (a single
+  extra rounding of the native input).
+* ``n4_dither1`` / ``n4_dither50``: native code on ``h`` with a random
+  +-1 bf16 ulp change in 1 % / 50 % of its elements.
+* ``ref``: native code whose projection outputs (``kva`` and, on the
+  full-indexer layer, the fused ``wk|weights_proj`` output ``kw``) are the
+  correctly rounded FP64 oracle ``bf16(r* * sum_k x_k gamma_k W_jk)``; every
+  arm is also reported as its distance to this step.
+* ``neg``: a packet without the gain through deferred_norm on the served layers (negative
+  control: a real bug).
+
+Reported per case, per arm and layer: the deviation metrics vs the native step
+(attention-output relRMS, absorbed-query relRMS, FP8 query bytes differing,
+indexer weights relRMS, FP8 indexer-query bytes differing, top-k rows
+identical, mean top-k Jaccard), the attention-output relRMS vs ``ref``, and
+whether the predeclared deviation bound (output relRMS <=
+2**-7 and mean Jaccard >= 0.98) holds. Per layout, per arm and served layer, the
+projection-level error ``e = ||y - y*|| / ||y*||`` of ``kva``.
+
+The numbers are evidence for the deviation bound, not a pass
+criterion: the assertions only check that every arm ran
+the intended path (deferred_norm engaged on the served layers only, hooks fired once, finite
+outputs) and that the experiment can see a real bug (on every served layer the negative
+control's output deviation exceeds every native-only control's).
+
+Fixture label: weights real (checkpoint), activations synthetic (Gaussian
+residual rows; the ``hot`` variant multiplies 12 channels by 30, as
+residual streams with outlier channels have).
+"""
+
+import contextlib
+import json
+import os
+
+import absorbed_query_fixture
+import deferred_norm_utils as U
+import pytest
+import test_absorbed_query_attention_integration as aq_integration
+import test_deferred_norm_attention_integration as dn_integration
+import torch
+
+LAYERS = aq_integration.LAYERS
+FULL = absorbed_query_fixture.FULL_LAYER
+LAYOUTS = aq_integration.LAYOUTS
+OUTPUT_BOUND = 2.0**-7
+NATIVE_CONTROLS = ("n1_rmsnorm", "n2_cublas", "n3_round", "n4_dither1", "n4_dither50")
+
+pytestmark = dn_integration.pytestmark
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = dn_integration._DeferredNormHarness()
+    yield h
+    dn_integration._set_features("")
+
+
+def _packets(h, layout, seed, hot):
+    """Per-layer packets of residual-like rows ``x`` (Gaussian x 2, optionally
+    12 channels x 30) with the layer's real gain."""
+    nreq, ntok = LAYOUTS[layout]
+    out = {}
+    for layer in LAYERS:
+        g = torch.Generator(device="cuda").manual_seed(seed * 10 + layer)
+        x = torch.randn(nreq * ntok, U.H, generator=g, device="cuda") * 2.0
+        if hot:
+            idx = torch.randint(0, U.H, (12,), generator=g, device="cuda")
+            x[:, idx] *= 30.0
+        out[layer] = U.make_frag_packet(x.bfloat16(), h.gammas[layer])
+    pos = (
+        (aq_integration.CTX + torch.arange(ntok, dtype=torch.int32, device="cuda"))
+        .repeat(nreq)
+        .unsqueeze(0)
+    )
+    return out, pos
+
+
+def _rmsnorm_inputs(packets):
+    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+
+    out = {}
+    for layer, p in packets.items():
+        norm = RMSNorm(hidden_size=U.H, eps=p.eps, dtype=torch.bfloat16, device="cuda")
+        norm.weight.data.copy_(p.gamma)
+        out[layer] = norm(p.residual)
+    return out
+
+
+def _dither(h, frac, seed):
+    """``h`` with a random +-1 ulp step in a fraction ``frac`` of its nonzero
+    elements (the bf16 bit pattern is sign-magnitude, so +-1 on the integer
+    view moves the magnitude by one ulp)."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    bits = h.contiguous().view(torch.int16)
+    pick = torch.rand(h.shape, generator=g, device="cuda") < frac
+    step = torch.where(torch.rand(h.shape, generator=g, device="cuda") < 0.5, -1, 1)
+    nonzero = (bits.to(torch.int32) & 0x7FFF) != 0
+    out = torch.where(pick & nonzero, bits.to(torch.int32) + step, bits.to(torch.int32))
+    return out.to(torch.int16).view(torch.bfloat16)
+
+
+@contextlib.contextmanager
+def _cublas_kv_a():
+    """``DeepseekV3Linear`` without its <= 16-token ``dsv3_fused_a_gemm_op``
+    branch: ``Linear.apply_linear`` (``cublas_mm``), as for larger batches."""
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Linear
+    from tensorrt_llm._torch.modules.linear import Linear
+
+    calls = [0]
+    own = DeepseekV3Linear.__dict__["apply_linear"]
+
+    def apply_linear(self, input, bias, lora_params=None, layer_idx=None):
+        calls[0] += 1
+        return Linear.apply_linear(self, input, bias, lora_params, layer_idx)
+
+    DeepseekV3Linear.apply_linear = apply_linear
+    try:
+        yield calls
+    finally:
+        DeepseekV3Linear.apply_linear = own
+
+
+def _oracle(p, weight):
+    return U.fp64_oracle(p.residual, p.gamma, weight, p.eps).to(torch.bfloat16)
+
+
+@contextlib.contextmanager
+def _oracle_projection(h, packets):
+    """The native step with ``kva`` / ``kw`` replaced by the correctly rounded
+    FP64 oracle of each layer's packet."""
+    f = torch.nn.functional
+    orig_linear = f.linear
+    kva = {
+        layer: _oracle(packets[layer], h.attn[layer].kv_a_proj_with_mqa.weight) for layer in LAYERS
+    }
+    kw = {}
+    for layer in LAYERS:
+        indexer = h.attn[layer].mqa.indexer
+        if indexer is not None:
+            w = indexer._fused_wk_wp_weight
+            kw[id(w)] = (layer, _oracle(packets[layer], w))
+    torch.cuda.synchronize()
+    calls = {f"kva{layer}": 0 for layer in LAYERS}
+    calls.update({f"kw{layer}": 0 for layer, _ in kw.values()})
+
+    def make_forward(layer):
+        def forward(input, *args, **kwargs):
+            ref = kva[layer]
+            assert input.shape[0] == ref.shape[0], (input.shape, ref.shape)
+            calls[f"kva{layer}"] += 1
+            return ref.clone()
+
+        return forward
+
+    def linear(input, weight, bias=None):
+        hit = kw.get(id(weight))
+        if hit is not None and bias is None:
+            layer, ref = hit
+            assert input.shape[0] == ref.shape[0], (input.shape, ref.shape)
+            calls[f"kw{layer}"] += 1
+            return ref.clone()
+        return orig_linear(input, weight, bias)
+
+    for layer in LAYERS:
+        h.attn[layer].kv_a_proj_with_mqa.forward = make_forward(layer)
+    f.linear = linear
+    try:
+        yield calls
+    finally:
+        f.linear = orig_linear
+        for layer in LAYERS:
+            del h.attn[layer].kv_a_proj_with_mqa.forward
+
+
+def _projection_errors(h, layer, p, inputs):
+    """``e`` of ``kva`` / ``kw`` per arm on one layer (projection level, vs
+    the FP64 oracle). ``inputs``: arm -> that arm's normalized hidden ``h``."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+    from tensorrt_llm._torch.modules.linear import Linear
+
+    attn = h.attn[layer]
+    lin = attn.kv_a_proj_with_mqa
+    indexer = attn.mqa.indexer
+    w_kw = indexer._fused_wk_wp_weight if indexer is not None else None
+    oracle = {"kva": U.fp64_oracle(p.residual, p.gamma, lin.weight, p.eps)}
+    if w_kw is not None:
+        oracle["kw"] = U.fp64_oracle(p.residual, p.gamma, w_kw, p.eps)
+    ok, reason, kw_weight = projection_eligibility(attn)
+    assert ok, reason
+    proj = DeferredNormProjection(attn, kw_weight=kw_weight)
+    assert proj.supported() is None, proj.supported()
+    kva_w, kw_w = proj.project(p)
+    arms = {"deferred_norm": {"kva": kva_w, "kw": kw_w}}
+    for arm, hid in inputs.items():
+        arms[arm] = {"kva": lin(hid)}
+        if w_kw is not None:
+            arms[arm]["kw"] = torch.nn.functional.linear(hid, w_kw)
+    arms["n2_cublas"] = {"kva": Linear.apply_linear(lin, inputs["native"], None)}
+    if w_kw is not None:
+        arms["n2_cublas"]["kw"] = torch.nn.functional.linear(inputs["native"], w_kw)
+    arms["ref"] = {k: v.to(torch.bfloat16) for k, v in oracle.items()}
+    torch.cuda.synchronize()
+    return {
+        arm: {
+            name: U.norm_rms_err(y, oracle[name])
+            for name, y in ys.items()
+            if name in oracle and y is not None
+        }
+        for arm, ys in arms.items()
+    }
+
+
+def _input_stats(inputs):
+    """Fraction of elements of each arm's ``h`` that differ from native and
+    their max bf16 ulp distance."""
+    out = {}
+    for arm, per_layer in inputs.items():
+        if arm == "native":
+            continue
+        out[arm] = {}
+        for layer, hid in per_layer.items():
+            ulp = U.bf16_ulp_distance(hid, inputs["native"][layer])
+            out[arm][layer] = {
+                "frac_diff": float((ulp > 0).float().mean()),
+                "max_ulp": int(ulp.max()),
+            }
+    return out
+
+
+def _vs_ref(arm_out, ref_out):
+    return {layer: aq_integration._rel_rms(arm_out[layer], ref_out[layer]) for layer in LAYERS}
+
+
+@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
+@pytest.mark.parametrize(
+    "features",
+    [
+        "absorbed_query,post_moe_hop,deferred_norm",
+        "absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm",
+    ],
+)
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_output_bound_with_native_controls(harness, layout, path, features, hot):
+    h = harness
+    base = ",".join(f for f in features.split(",") if f != "deferred_norm")
+    seed = 3 if not hot else 13
+    packets, pos = _packets(h, layout, seed, hot)
+    native_in = {layer: p.materialize() for layer, p in packets.items()}
+    inputs = {
+        "native": native_in,
+        "n1_rmsnorm": _rmsnorm_inputs(packets),
+        "n3_round": dn_integration._control_inputs(packets),
+        "n4_dither1": {layer: _dither(v, 0.01, 1000 + layer) for layer, v in native_in.items()},
+        "n4_dither50": {layer: _dither(v, 0.5, 2000 + layer) for layer, v in native_in.items()},
+    }
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHidden,
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        frag_from_rows,
+    )
+
+    neg = {
+        layer: DeferredNormHidden(
+            frag_from_rows(p.residual), p.sumsq, p.residual, p.gamma, p.eps, materialize
+        )
+        if layer in dn_integration.SERVED_LAYERS
+        else native_in[layer]
+        for layer, p in packets.items()
+    }
+    snap = h.snapshot()
+    runs = {}
+    try:
+
+        def native_run(hidden):
+            dn_integration._reset_runners(h)
+            return dn_integration._run(h, base, layout, path, hidden, pos, snap)
+
+        runs["native"] = native_run(native_in)
+        dn_integration._reset_runners(h)
+        runs["deferred_norm"] = dn_integration._run(
+            h, features, layout, path, dn_integration._served(packets), pos, snap
+        )
+        for layer in LAYERS:
+            engaged = h.attn[layer]._deferred_norm_projection not in (None, False)
+            assert engaged == (layer in dn_integration.SERVED_LAYERS), f"deferred_norm on L{layer}"
+        runs["rerun"] = native_run(native_in)
+        for arm in ("n1_rmsnorm", "n3_round", "n4_dither1", "n4_dither50"):
+            runs[arm] = native_run(inputs[arm])
+        with _cublas_kv_a() as calls:
+            runs["n2_cublas"] = native_run(native_in)
+        assert calls[0] == len(LAYERS), f"cuBLAS kv_a path ran {calls[0]} times"
+        with _oracle_projection(h, packets) as calls:
+            runs["ref"] = native_run(native_in)
+        assert all(v == 1 for v in calls.values()), calls
+        dn_integration._reset_runners(h)
+        runs["neg"] = dn_integration._run(h, features, layout, path, neg, pos, snap)
+
+        report = {
+            "layout": layout,
+            "path": path,
+            "features": features,
+            "data": "hot" if hot else "gauss",
+            "arms": {},
+        }
+        ref_out = runs["ref"][0]
+        for arm, run in runs.items():
+            for layer in LAYERS:
+                assert torch.isfinite(run[0][layer].float()).all(), (arm, layer)
+            if arm == "native":
+                continue
+            stats = dn_integration._output_stats(runs["native"], run, layout)
+            for layer, d in _vs_ref(run[0], ref_out).items():
+                stats[layer]["out_rel_rms_vs_ref"] = d
+                stats[layer]["original_output_bound_ok"] = bool(
+                    stats[layer]["out_rel_rms"] <= OUTPUT_BOUND
+                    and stats[layer].get("topk_mean_jaccard", 1.0) >= 0.98
+                )
+            report["arms"][arm] = stats
+        report["arms"]["native"] = {
+            layer: {"out_rel_rms_vs_ref": d}
+            for layer, d in _vs_ref(runs["native"][0], ref_out).items()
+        }
+        report["inputs"] = _input_stats(inputs)
+        print(
+            "GLM_MOE_DSA_DEFERRED_NORM_CONTROLS " + json.dumps(report, sort_keys=True), flush=True
+        )
+        out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_CONTROLS_OUT")
+        if out_dir:
+            os.makedirs(out_dir, exist_ok=True)
+            name = f"controls-{layout}-{path}-{features.replace(',', '_')}-{report['data']}.json"
+            with open(os.path.join(out_dir, name), "w") as fh:
+                json.dump(report, fh, indent=1, sort_keys=True)
+
+        for layer in dn_integration.SERVED_LAYERS:
+            native_max = max(report["arms"][a][layer]["out_rel_rms"] for a in NATIVE_CONTROLS)
+            assert report["arms"]["neg"][layer]["out_rel_rms"] > native_max, (
+                f"L{layer}: the negative control is not separated from the native controls"
+            )
+    finally:
+        h.restore(snap)
+        dn_integration._set_features("")
+
+
+@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_projection_error_of_every_arm(harness, layout, hot):
+    """The projection error (``e`` vs the FP64 oracle) of deferred_norm and of every
+    native-only input / GEMM variant of the diagnosis."""
+    h = harness
+    packets, _ = _packets(h, layout, 3 if not hot else 13, hot)
+    native_in = {layer: p.materialize() for layer, p in packets.items()}
+    per_arm_inputs = {
+        "native": native_in,
+        "n1_rmsnorm": _rmsnorm_inputs(packets),
+        "n3_round": dn_integration._control_inputs(packets),
+        "n4_dither1": {layer: _dither(v, 0.01, 1000 + layer) for layer, v in native_in.items()},
+        "n4_dither50": {layer: _dither(v, 0.5, 2000 + layer) for layer, v in native_in.items()},
+    }
+    errs = {
+        layer: _projection_errors(
+            h, layer, packets[layer], {arm: v[layer] for arm, v in per_arm_inputs.items()}
+        )
+        for layer in dn_integration.SERVED_LAYERS
+    }
+    report = {
+        "layout": layout,
+        "data": "hot" if hot else "gauss",
+        "e": errs,
+        "inputs": _input_stats(per_arm_inputs),
+    }
+    print(
+        "GLM_MOE_DSA_DEFERRED_NORM_CONTROLS_PROJ " + json.dumps(report, sort_keys=True), flush=True
+    )
+    for layer, arms in errs.items():
+        for arm, d in arms.items():
+            for name, e in d.items():
+                assert e == e and e < 0.05, (layer, arm, name, e)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py
new file mode 100644
index 000000000000..22a4391d31da
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py
@@ -0,0 +1,835 @@
+# 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.
+"""Collective tests of the deferred_norm deferred hop (8 ranks over 2 hosts, TP8).
+
+Run every test file process under ``srun --mpi=pmix`` (one process per GPU);
+every test is collective and runs in the same order on every rank.
+
+Exact invariants (the deferred_norm hop against the post_moe_hop on the same raw tuple and
+the same workspace):
+
+* ``residual_out`` bit-identical to post_moe_hop's;
+* the fragment-layout ``a_frag`` holds ``bf16(x * gamma)`` bitwise (torch);
+* ``sumsq`` bit-identical to the torch emulation of post_moe_hop's per-thread /
+  butterfly order;
+* ``materialize(packet)`` bit-identical to post_moe_hop's normalized output (this
+  checks that the deferred ``r`` equals post_moe_hop's ``r`` bit for bit).
+
+Workspace protocol shared with post_moe_hop (``post_moe_hop/kernels.py``; post_moe_hop's
+error handling), with deferred_norm and post_moe_hop hops on one workspace:
+
+* a rank that never joins makes every other rank's deferred_norm call time out after
+  the post_moe_hop's deadline, record the Lamport-timeout bit and poll count, and
+  return NaN (residual, activation, square sums, materialized output); post_moe_hop's
+  health check fails, and after a collective reset the next call is exact;
+* an error word left by an earlier call poisons every later deferred_norm and post_moe_hop call
+  of the group (NaN) and is carried from buffer to buffer until the reset;
+* a permutation entry at or beyond the FC2 row count is skipped (output as
+  post_moe_hop's) and reported (bit 4);
+* scrambled store order (fault-injected deferred_norm and post_moe_hop builds, chosen per rank
+  and per call, mixed with plain builds and random device delays) changes
+  no output bit;
+* contract violations raise ``PostMoeHopContractError`` before any device work;
+  layouts the kernel cannot read with 16-byte vectors are repaired on the
+  rank that has them (counted in the workspace), also in a CUDA graph;
+* the health check counts deferred_norm calls as hop calls.
+"""
+
+import functools
+import os
+import random
+import sys
+import time
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
+
+import deferred_norm_utils as U  # noqa: E402
+import post_moe_hop_mpi_utils as M  # noqa: E402
+
+
+def _env():
+    env = M.mpi_env()
+    if env.world != 8:
+        pytest.skip("the deferred_norm hop is specialized to TP8")
+    return env
+
+
+_HOPS = {}
+
+
+def _hops(trigger_early=True):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    if "post_moe_hop" not in _HOPS:
+        _HOPS["post_moe_hop"] = PostMoeHop(
+            env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False
+        )
+    post_moe_hop = _HOPS["post_moe_hop"]
+    key = ("deferred_norm", trigger_early)
+    if key not in _HOPS:
+        _HOPS[key] = DeferredNormHop(post_moe_hop, trigger_early=trigger_early)
+    return post_moe_hop, _HOPS[key]
+
+
+def _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
+    """This rank's packet invariants against post_moe_hop's outputs (not collective)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    torch.cuda.synchronize()
+    ok_res = torch.equal(res_deferred_norm, res_post_moe_hop)
+    scaled_ref = (res_deferred_norm.float() * gamma.float()).bfloat16()
+    ok_scaled = torch.equal(
+        rows_from_frag(packet.a_frag, res_deferred_norm.shape[0], res_deferred_norm.shape[1]),
+        scaled_ref,
+    )
+    ok_ss = torch.equal(packet.sumsq, U.hop_sumsq_partials(res_deferred_norm))
+    h = materialize(packet)
+    torch.cuda.synchronize()
+    ok_h = torch.equal(h, h_post_moe_hop)
+    ok = ok_res and ok_scaled and ok_ss and ok_h
+    if not ok:
+        print(
+            f"GLM_MOE_DSA_DEFERRED_NORM_HOP_FAIL {tag} rank={M.mpi_env().rank} res={ok_res} scaled={ok_scaled} "
+            f"sumsq={ok_ss} materialize={ok_h} "
+            f"max_ulp_h={int(U.bf16_ulp_distance(h, h_post_moe_hop).max())}",
+            flush=True,
+        )
+    return ok
+
+
+def _check(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
+    """Collective: ``_verdict`` on every rank."""
+    return M.all_ranks(
+        _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag)
+    )
+
+
+def _all_nan_packet(packet, res):
+    """Every value a consumer can derive from the packet is NaN."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    torch.cuda.synchronize()
+    a = rows_from_frag(packet.a_frag, res.shape[0], res.shape[1])
+    h = materialize(packet)
+    torch.cuda.synchronize()
+    return all(bool(torch.isnan(t.float()).all()) for t in (res, a, packet.sumsq, h))
+
+
+def _reset(post_moe_hop):
+    assert M.all_ranks(post_moe_hop.reset())
+    assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+def test_capability_engaged():
+    post_moe_hop, deferred_norm = _hops()
+    assert M.all_ranks(post_moe_hop.engaged and deferred_norm.engaged), deferred_norm.reason
+
+
+def test_hop_precompiles_served_projections():
+    """An engaged hop has compiled both served projection kernels on every
+    rank (the consumer's compile is a cache hit)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.proj_kernel import (  # noqa: E501
+        is_compiled,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        SERVED_N1,
+        served_kernel_configs,
+    )
+
+    _, deferred_norm = _hops()
+    cfgs = served_kernel_configs(torch.cuda.current_device())
+    ok = (
+        deferred_norm.engaged
+        and sorted(c["n1"] for c in cfgs) == sorted(SERVED_N1)
+        and all(is_compiled(**c) for c in cfgs)
+    )
+    assert M.all_ranks(ok)
+
+
+@pytest.mark.parametrize("failing_rank", [3, 7])
+def test_projection_compile_failure_disengages_every_rank(monkeypatch, failing_rank):
+    """A projection compile error on one rank disengages the hop on all
+    ranks (collective decision), so no rank produces a packet."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+        projection,
+    )
+
+    env = _env()
+    post_moe_hop, _ = _hops()
+    if env.rank == failing_rank:
+
+        def _fail(device):
+            raise RuntimeError("injected projection compile failure")
+
+        monkeypatch.setattr(projection, "precompile_served", _fail)
+    hop = DeferredNormHop(post_moe_hop)
+    if env.rank == failing_rank:
+        local = "injected projection compile failure" in hop.reason
+    else:
+        local = hop.reason == "another rank is not capable"
+    assert M.all_ranks(not hop.engaged and local), hop.reason
+
+
+def test_unsupported_projection_disengages_every_rank(monkeypatch):
+    """A served projection shape the kernel does not support (static,
+    rank-identical) disengages the hop on every rank."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501  # noqa: E501
+        DeferredNormHop,
+        projection,
+    )
+
+    post_moe_hop, _ = _hops()
+    monkeypatch.setattr(
+        projection,
+        "DeferredNormProjConfig",
+        functools.partial(projection.DeferredNormProjConfig, smem_rows=8),
+    )
+    hop = DeferredNormHop(post_moe_hop)
+    assert M.all_ranks(not hop.engaged and "staged rows" in hop.reason), hop.reason
+
+
+@pytest.mark.parametrize("trigger_early", [True, False])
+@pytest.mark.parametrize("scale_dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("with_shared", [True, False])
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+def test_deferred_matches_post_moe_hop(rows, with_shared, scale_dtype, trigger_early):
+    post_moe_hop, deferred_norm = _hops(trigger_early)
+    with torch.inference_mode():
+        fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(
+            rows, seed=40 + rows, scale_dtype=scale_dtype, with_shared=with_shared
+        )
+        h_post_moe_hop, res_post_moe_hop = post_moe_hop(
+            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
+        )
+        packet, res_deferred_norm = deferred_norm(
+            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
+        )
+        assert _check(
+            packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, f"rows={rows}"
+        )
+
+
+def test_empty_rows_and_large_residual():
+    post_moe_hop, deferred_norm = _hops()
+    with torch.inference_mode():
+        fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(6, seed=77, empty_rows=(1, 4))
+        residual = (residual.float() * 50.0).bfloat16()
+        h_post_moe_hop, res_post_moe_hop = post_moe_hop(
+            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
+        )
+        packet, res_deferred_norm = deferred_norm(
+            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
+        )
+        assert _check(
+            packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, "empty-rows"
+        )
+
+
+def test_back_to_back_mixed_with_post_moe_hop():
+    """deferred_norm and post_moe_hop hops interleaved on one workspace without host syncs."""
+    post_moe_hop, deferred_norm = _hops()
+    rows_seq = [6, 1, 8, 3, 6, 2, 6, 8, 5, 6, 1, 6]
+    with torch.inference_mode():
+        inputs = [M.make_raw_tuple(r, seed=900 + i) for i, r in enumerate(rows_seq)]
+        outs = []
+        for i, t in enumerate(inputs):
+            if i % 2 == 0:
+                outs.append(("deferred_norm", deferred_norm(*t, M.GLM_EPS)))
+            else:
+                outs.append(("post_moe_hop", post_moe_hop(*t, M.GLM_EPS)))
+        torch.cuda.synchronize()
+        # Reference: each call again, alone, through post_moe_hop.
+        failures = []
+        for i, (t, (kind, out)) in enumerate(zip(inputs, outs)):
+            h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+            if kind == "post_moe_hop":
+                ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
+            else:
+                ok = _check(out[0], out[1], h_ref, res_ref, t[5], f"b2b{i}")
+            if not ok:
+                failures.append(i)
+        assert M.all_ranks(not failures), failures
+        assert post_moe_hop.error_bits() == 0
+
+
+def test_cuda_graph_replay():
+    post_moe_hop, deferred_norm = _hops()
+    with torch.inference_mode():
+        static = [M.make_raw_tuple(6, seed=1200 + i) for i in range(3)]
+        bufs = [[x.clone() if x is not None else None for x in t] for t in static]
+
+        def body():
+            a = deferred_norm(*bufs[0], M.GLM_EPS)
+            b = post_moe_hop(*bufs[1], M.GLM_EPS)
+            c = deferred_norm(*bufs[2], M.GLM_EPS)
+            return a, b, c
+
+        body()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            ga, gb, gc = body()
+        for trial, replays in enumerate((1, 1, 20)):
+            fresh = [M.make_raw_tuple(6, seed=1300 + 10 * trial + i) for i in range(3)]
+            for dst, src in zip(bufs, fresh):
+                for d, s in zip(dst, src):
+                    if d is not None:
+                        d.copy_(s)
+            for _ in range(replays):
+                M.mpi_env().comm.barrier()
+                graph.replay()
+            torch.cuda.synchronize()
+            ok = True
+            for (kind, got), t in zip(
+                (("deferred_norm", ga), ("post_moe_hop", gb), ("deferred_norm", gc)), fresh
+            ):
+                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+                if kind == "post_moe_hop":
+                    ok &= torch.equal(got[0], h_ref) and torch.equal(got[1], res_ref)
+                else:
+                    ok &= _check(got[0], got[1], h_ref, res_ref, t[5], f"graph{trial}")
+            assert M.all_ranks(ok), f"trial {trial}"
+
+
+def test_delayed_rank_stress():
+    """Random per-rank device delays inside back-to-back batches (no host sync)."""
+    post_moe_hop, deferred_norm = _hops()
+    env = M.mpi_env()
+    gen = torch.Generator().manual_seed(4242)
+    with torch.inference_mode():
+        for batch in range(4):
+            inputs = []
+            outs = []
+            for i in range(12):
+                rows = int(torch.randint(1, 9, (1,), generator=gen))
+                t = M.make_raw_tuple(rows, seed=5000 + 100 * batch + i)
+                local = torch.Generator().manual_seed(77 * batch + 13 * i + env.rank)
+                if float(torch.rand(1, generator=local)) < 0.3:
+                    torch.cuda._sleep(int(float(torch.rand(1, generator=local)) * 2_000_000))
+                inputs.append(t)
+                outs.append(deferred_norm(*t, M.GLM_EPS))
+            torch.cuda.synchronize()
+            ok = True
+            for t, (packet, res) in zip(inputs, outs):
+                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+                ok &= _check(packet, res, h_ref, res_ref, t[5], f"stress{batch}")
+            assert M.all_ranks(ok), f"batch {batch}"
+        assert post_moe_hop.error_bits() == 0
+
+
+@pytest.mark.parametrize("layer", [10, 76])
+def test_real_weights(layer):
+    """Real routed experts / router / shared expert / next-layer gain (synthetic activations)."""
+    from post_moe_hop_real_fixture import build_real_moe_layer
+
+    post_moe_hop, deferred_norm = _hops()
+    env = M.mpi_env()
+    with torch.inference_mode():
+        real = build_real_moe_layer(env.mapping, layer, max_num_tokens=64)
+        ok = True
+        for rows in (1, 6, 8):
+            x = real.hidden_states(rows, seed=31 + rows)
+            logits = real.gate(x)
+            shared = real.shared_output(x)
+            residual = real.residual(rows, seed=41 + rows)
+            fc2, scales, perm = real.run_moe(x, logits, False)
+            h_post_moe_hop, res_post_moe_hop = post_moe_hop(
+                fc2, scales, perm, shared, residual, real.next_gamma, M.GLM_EPS
+            )
+            packet, res_deferred_norm = deferred_norm(
+                fc2, scales, perm, shared, residual, real.next_gamma, M.GLM_EPS
+            )
+            ok &= _check(
+                packet,
+                res_deferred_norm,
+                h_post_moe_hop,
+                res_post_moe_hop,
+                real.next_gamma,
+                f"L{layer}R{rows}",
+            )
+        real.moe.destroy()
+        assert M.all_ranks(ok)
+
+
+# ---------------------------------------------------------------------------
+# Workspace protocol shared with post_moe_hop: bounded errors, sticky poison, permutation
+# range, scrambled store order, raw-tuple contract and layout repairs.
+# ---------------------------------------------------------------------------
+
+_MISSING_TIMEOUT_S = 0.2
+
+
+def _impatient():
+    """A post_moe_hop with a short poll deadline on the same workspace, and the deferred_norm
+    hop built on it (collective; the deferred_norm kernels are compile-cache hits)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    if "impatient" not in _HOPS:
+        post_moe_hop = PostMoeHop(
+            env.mapping,
+            M.GLM_HIDDEN,
+            M.GLM_TOP_K,
+            require_multi_node=False,
+            timeout_s=_MISSING_TIMEOUT_S,
+        )
+        _HOPS["impatient"] = (post_moe_hop, DeferredNormHop(post_moe_hop))
+    return _HOPS["impatient"]
+
+
+@pytest.mark.parametrize("missing", [0, 5, 7])
+def test_missing_rank_times_out_with_nan_and_reset(missing):
+    """A rank that never joins makes every other rank's deferred_norm call time out,
+    report and return NaN instead of hanging; a collective reset recovers."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+
+    post_moe_hop, deferred_norm = _hops()
+    imp_post_moe_hop, imp_deferred_norm = _impatient()
+    env = M.mpi_env()
+    assert M.all_ranks(
+        imp_deferred_norm.engaged and imp_post_moe_hop.workspace is post_moe_hop.workspace
+    )
+    assert M.all_ranks(imp_deferred_norm.timeout_ns == int(_MISSING_TIMEOUT_S * 1e9))
+    rows = 6
+    with torch.inference_mode():
+        tensors = M.make_raw_tuple(rows, seed=600 + missing)
+        before = M.lamport_snapshot(post_moe_hop)
+        assert M.all_ranks(before["ok"]), before
+        env.comm.barrier()
+        start = time.monotonic()
+        packet = res = None
+        if env.rank != missing:
+            packet, res = imp_deferred_norm(*tensors, M.GLM_EPS)
+        torch.cuda.synchronize()
+        elapsed = time.monotonic() - start
+        counters = post_moe_hop.workspace.counters()
+        bits, polls = counters["error_bits"], counters["timeout_polls"]
+        if env.rank == missing:
+            reported = bits == 0 and polls == 0
+        else:
+            reported = (
+                bool(bits & 1)
+                and polls > 0
+                and _MISSING_TIMEOUT_S <= elapsed < 30.0
+                and _all_nan_packet(packet, res)
+            )
+        stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
+        if env.rank == 0:
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_MISSING rank {missing} skipped: (rank, elapsed_s, error_bits, "
+                f"timeout_polls, ok) {stats}",
+                flush=True,
+            )
+        assert M.all_ranks(reported), stats
+        try:
+            check_health(f"deferred_norm-missing-rank{missing}-local")
+            local_raised = False
+        except PostMoeHopError:
+            local_raised = True
+        try:
+            check_health(f"deferred_norm-missing-rank{missing}-collective", collective=True)
+            group_raised = False
+        except PostMoeHopError:
+            group_raised = True
+        assert M.all_ranks(local_raised == (env.rank != missing) and group_raised)
+        _reset(post_moe_hop)
+        t = M.make_raw_tuple(rows, seed=650 + missing)
+        packet, res = deferred_norm(*t, M.GLM_EPS)
+        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+        assert _check(packet, res, h_ref, res_ref, t[5], f"after-reset{missing}")
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize(
+    "sequence",
+    [
+        "deferred_norm,deferred_norm,deferred_norm",
+        "deferred_norm,post_moe_hop,deferred_norm",
+        "post_moe_hop,deferred_norm,post_moe_hop",
+    ],
+)
+def test_error_poisons_every_rank_until_reset(sequence):
+    """An error word left by an earlier call on one rank makes the next calls
+    of every rank NaN, deferred_norm and post_moe_hop alike, and is carried through all three
+    Lamport buffers until the collective reset."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
+        STATE_EPOCH,
+        STATE_ERROR_WORDS,
+    )
+
+    post_moe_hop, deferred_norm = _hops()
+    env = M.mpi_env()
+    victim = env.world - 2
+    with torch.inference_mode():
+        torch.cuda.synchronize()
+        if env.rank == victim:
+            # As if this rank's previous call had timed out.
+            epoch = int(post_moe_hop.workspace.state[STATE_EPOCH].item())
+            post_moe_hop.workspace.state[STATE_ERROR_WORDS[(epoch + 2) % 3]] = 1
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        poisoned = True
+        for call, kind in enumerate(sequence.split(",")):
+            t = M.make_raw_tuple(6, seed=860 + call)
+            if kind == "deferred_norm":
+                packet, res = deferred_norm(*t, M.GLM_EPS)
+                poisoned &= _all_nan_packet(packet, res)
+            else:
+                out, res = post_moe_hop(*t, M.GLM_EPS)
+                torch.cuda.synchronize()
+                poisoned &= bool(torch.isnan(out.float()).all() and torch.isnan(res.float()).all())
+        words = post_moe_hop.workspace.state.cpu().tolist()
+        carried = all(words[w] == 1 for w in STATE_ERROR_WORDS)
+        clean = not any(words[w] for w in STATE_ERROR_WORDS)
+        if env.rank == 0:
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_POISON {sequence}: per-rank error words "
+                f"{env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])}",
+                flush=True,
+            )
+        else:
+            env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])
+        assert M.all_ranks(poisoned and (carried if env.rank == victim else clean))
+        with pytest.raises(PostMoeHopError):
+            check_health("deferred_norm-poisoned", collective=True)
+        _reset(post_moe_hop)
+        t = M.make_raw_tuple(6, seed=870)
+        packet, res = deferred_norm(*t, M.GLM_EPS)
+        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+        assert _check(packet, res, h_ref, res_ref, t[5], "after-poison-reset")
+
+
+def test_permutation_out_of_range_is_skipped_and_reported():
+    """Entries at or beyond the FC2 row count are skipped (as by post_moe_hop and the
+    torch oracle) and set error bit 4 on the rank that has them. FC2 holds
+    exactly the referenced rows, so under memcheck without the caching
+    allocator a read of a skipped entry would be an invalid access."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
+        ERROR_PERM_RANGE,
+    )
+
+    post_moe_hop, deferred_norm = _hops()
+    env = M.mpi_env()
+    culprit = 2 % env.world
+    rows = 6
+    with torch.inference_mode():
+        t = list(M.make_raw_tuple(rows, seed=880, padded_rows=rows * M.GLM_TOP_K))
+        if env.rank == culprit:
+            t[2] = t[2].clone()
+            t[2][1, 3] = t[0].shape[0]
+            t[2][4, 0] = t[0].shape[0] + 1000
+        packet, res = deferred_norm(*t, M.GLM_EPS)
+        torch.cuda.synchronize()
+        bits = post_moe_hop.error_bits()
+        reported = bits == (ERROR_PERM_RANGE if env.rank == culprit else 0)
+        with pytest.raises(PostMoeHopError):
+            check_health("deferred_norm-perm-range", collective=True)
+        _reset(post_moe_hop)
+        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+        torch.cuda.synchronize()
+        ok = _verdict(packet, res, h_ref, res_ref, t[5], "perm-range") and reported
+        # The oracle skips entries outside [0, N_padded) like the kernels.
+        ref_local = M.sanitize_bf16(M.reference_local_partial(t[0], t[1], t[2], t[3]))
+        slots = M.gather_bf16(ref_local).permute(1, 0, 2).cuda()
+        _, emu_res = M.emulate_post_reduction(slots, t[4], t[5], M.GLM_EPS)
+        ok &= torch.equal(res.view(torch.int16), emu_res.view(torch.int16))
+        assert M.all_ranks(ok), bits
+        _reset(post_moe_hop)
+
+
+_FAULT_DELAYS_NS = (3000, 60000)
+
+
+def _fault_hops():
+    """post_moe_hop and deferred_norm hops with fault-injected kernel builds on the same workspace."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    post_moe_hop, deferred_norm = _hops()
+    if "fault" not in _HOPS:
+        hops = {"post_moe_hop": post_moe_hop, "deferred_norm": deferred_norm}
+        for d in _FAULT_DELAYS_NS:
+            fp6 = PostMoeHop(
+                env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False, fault_delay_ns=d
+            )
+            hops[f"post_moe_hop_fault{d}"] = fp6
+            hops[f"deferred_norm_fault{d}"] = DeferredNormHop(fp6)
+        _HOPS["fault"] = hops
+    return _HOPS["fault"]
+
+
+def test_fault_injected_store_order_mixed_with_post_moe_hop():
+    """Scrambled store order inside the kernels (fault-injected builds), a
+    per-rank choice of deferred_norm or post_moe_hop kernel for the same collective call, and
+    random per-rank device delays: every output bit as in an undisturbed post_moe_hop
+    call, the call counter advances once per call, no error bit."""
+    hops = _fault_hops()
+    post_moe_hop = hops["post_moe_hop"]
+    env = M.mpi_env()
+    names = sorted(hops)
+    assert M.all_ranks(all(h.engaged for h in hops.values()))
+    common = random.Random(4711)
+    local = random.Random(1000 + env.rank)
+    calls_before = post_moe_hop.workspace.calls()
+    failed, used, total = [], {}, 0
+    with torch.inference_mode():
+        for batch in range(6):
+            plan = [
+                (
+                    common.randint(1, 8),
+                    common.random() < 0.75,
+                    torch.float32 if common.random() < 0.2 else torch.bfloat16,
+                )
+                for _ in range(12)
+            ]
+            inputs = [
+                M.make_raw_tuple(r, seed=9000 + 100 * batch + j, with_shared=sh, scale_dtype=sd)
+                for j, (r, sh, sd) in enumerate(plan)
+            ]
+            picks = [local.choice(names) for _ in plan]
+            torch.cuda.synchronize()
+            env.comm.barrier()
+            outs = []
+            for name, t in zip(picks, inputs):
+                if local.random() < 0.3:
+                    torch.cuda._sleep(int(local.uniform(0.0, 50.0) * 2000))
+                used[name] = used.get(name, 0) + 1
+                outs.append(hops[name](*t, M.GLM_EPS))
+            torch.cuda.synchronize()
+            for j, (name, t, out) in enumerate(zip(picks, inputs, outs)):
+                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+                if name.startswith("deferred_norm"):
+                    ok = _verdict(out[0], out[1], h_ref, res_ref, t[5], f"fault{batch}.{j}")
+                else:
+                    torch.cuda.synchronize()
+                    ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
+                if not ok:
+                    failed.append((batch, j, name))
+            total += len(plan)
+        counted = post_moe_hop.workspace.calls() - calls_before
+    bits = post_moe_hop.error_bits()
+    if env.rank == 0:
+        print(
+            f"GLM_MOE_DSA_DEFERRED_NORM_FAULT calls={total} (+{total} reference calls) used_on_rank0={used}",
+            flush=True,
+        )
+    assert M.all_ranks(not failed and counted == 2 * total and bits == 0), (
+        failed[:4],
+        counted,
+        bits,
+    )
+
+
+@pytest.mark.parametrize(
+    "case",
+    ("fc2_dtype", "scale_shape", "perm_dtype", "residual_shape", "shared_dtype", "gamma_shape"),
+)
+def test_contract_violation_raises_before_any_device_work(case):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopContractError,
+    )
+
+    post_moe_hop, deferred_norm = _hops()
+    rows = 6
+    with torch.inference_mode():
+        fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(rows, seed=920)
+        if case == "fc2_dtype":
+            fc2 = fc2.half()
+        elif case == "scale_shape":
+            scales = scales[:, : M.GLM_TOP_K - 1]
+        elif case == "perm_dtype":
+            perm = perm.long()
+        elif case == "residual_shape":
+            residual = residual[:, : M.GLM_HIDDEN // 2]
+        elif case == "shared_dtype":
+            shared = shared.float()
+        else:
+            gamma = gamma[:-8]
+        calls_before = post_moe_hop.workspace.calls()
+        with pytest.raises(PostMoeHopContractError):
+            deferred_norm(fc2, scales, perm, shared, residual, gamma, M.GLM_EPS)
+        assert M.all_ranks(post_moe_hop.workspace.calls() == calls_before)
+        t = M.make_raw_tuple(rows, seed=921)
+        packet, res = deferred_norm(*t, M.GLM_EPS)
+        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+        assert _check(packet, res, h_ref, res_ref, t[5], f"after-contract-{case}")
+
+
+def _misaligned(t, offset_elems):
+    storage = torch.empty(t.numel() + 16, dtype=t.dtype, device=t.device)
+    view = storage[offset_elems : offset_elems + t.numel()].view(t.shape)
+    view.copy_(t)
+    return view
+
+
+def _transposed(t):
+    view = torch.empty(tuple(reversed(t.shape)), dtype=t.dtype, device=t.device).t()
+    view.copy_(t)
+    return view
+
+
+def _padded_rows(t, extra_cols):
+    storage = torch.zeros((t.shape[0], t.shape[1] + extra_cols), dtype=t.dtype, device=t.device)
+    view = storage[:, : t.shape[1]]
+    view.copy_(t)
+    return view
+
+
+def _needs_copy(t, row_major_2d=False):
+    if t.data_ptr() % 16 or t.stride(-1) != 1:
+        return True
+    if row_major_2d:
+        return (t.stride(0) * t.element_size()) % 16 != 0
+    return not t.is_contiguous()
+
+
+def _relayout(rank, tensors):
+    """Rank-dependent layouts; returns (tensors, expected repairs on this rank)."""
+    fc2, scales, perm, shared, residual, gamma = tensors
+    kind = rank % 4
+    if kind == 0:
+        fc2 = _misaligned(fc2, 1)
+        perm = _misaligned(perm, 1)
+        gamma = _misaligned(gamma, 2)
+    elif kind == 1:
+        fc2 = _padded_rows(fc2, 4)
+        scales = _transposed(scales)
+        residual = _misaligned(residual, 4)
+    elif kind == 2:
+        fc2 = _padded_rows(fc2, 8)  # aligned row pitch: read in place
+        if shared is not None:
+            shared = _transposed(shared)
+    repairs = int(_needs_copy(fc2, row_major_2d=True)) + sum(
+        int(_needs_copy(t)) for t in (scales, perm, shared, residual, gamma) if t is not None
+    )
+    return (fc2, scales, perm, shared, residual, gamma), repairs
+
+
+def test_unaligned_layouts_are_repaired_locally():
+    """Per-rank layout repairs keep the collective and every output bit."""
+    post_moe_hop, deferred_norm = _hops()
+    env = M.mpi_env()
+    ok = True
+    with torch.inference_mode():
+        for i, rows in enumerate((6, 1, 8)):
+            reference = M.make_raw_tuple(rows, seed=880 + i)
+            tensors, expected = _relayout(env.rank, reference)
+            before = post_moe_hop.workspace.layout_repairs
+            packet, res = deferred_norm(*tensors, M.GLM_EPS)
+            repaired = post_moe_hop.workspace.layout_repairs - before
+            h_ref, res_ref = post_moe_hop(*reference, M.GLM_EPS)
+            ok &= _verdict(packet, res, h_ref, res_ref, reference[5], f"relayout{rows}")
+            ok &= repaired == expected
+        # The repaired call inside a CUDA graph, replayed with fresh values.
+        rows = 6
+        static_ref = M.make_raw_tuple(rows, seed=890)
+        static, _ = _relayout(env.rank, static_ref)
+        deferred_norm(*static, M.GLM_EPS)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            g_packet, g_res = deferred_norm(*static, M.GLM_EPS)
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        for replay in range(4):
+            fresh = M.make_raw_tuple(rows, seed=900 + replay, padded_rows=static_ref[0].shape[0])
+            for dst, src in zip(static, fresh):
+                if dst is not None:
+                    dst.copy_(src)
+            graph.replay()
+            torch.cuda.synchronize()
+            h_ref, res_ref = post_moe_hop(*fresh, M.GLM_EPS)
+            ok &= _verdict(g_packet, g_res, h_ref, res_ref, fresh[5], f"relayout-graph{replay}")
+    assert M.all_ranks(ok and post_moe_hop.error_bits() == 0)
+
+
+def test_health_check_counts_deferred_norm_calls():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        check_health,
+    )
+
+    post_moe_hop, deferred_norm = _hops()
+    with torch.inference_mode():
+        before = check_health("deferred_norm-engagement-before", collective=True)
+        for i in range(5):
+            deferred_norm(*M.make_raw_tuple(6, seed=940 + i), M.GLM_EPS)
+        after = check_health("deferred_norm-engagement-after", collective=True)
+    assert M.all_ranks(
+        after.calls - before.calls == 5
+        and after.error_bits == 0
+        and not after.group_error
+        and after.calls == post_moe_hop.workspace.calls()
+    )
+
+
+def test_memcheck_positive_control():
+    """Tests only, run under compute-sanitizer memcheck by the sanitizer
+    suite (``GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL=1``): the kernel is told FC2 has one more
+    row than it does and the permutation references that row, so memcheck
+    must report an invalid global read in the deferred_norm hop."""
+    if os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL") != "1":
+        pytest.skip("memcheck positive control (sanitizer suite only)")
+    post_moe_hop, deferred_norm = _hops()
+    rows = 1
+    with torch.inference_mode():
+        t = list(M.make_raw_tuple(rows, seed=990, padded_rows=rows * M.GLM_TOP_K))
+        t[2] = t[2].clone()
+        t[2][0, 0] = t[0].shape[0]
+        deferred_norm.launch(*t, M.GLM_EPS, fc2_rows=t[0].shape[0] + 1)
+        torch.cuda.synchronize()
+    _reset(post_moe_hop)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py
new file mode 100644
index 000000000000..c7471662e07c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py
@@ -0,0 +1,559 @@
+# 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.
+"""Null distribution of deferred_norm's attention-output distances among correct
+rounding variants of the same projection.
+
+Same harness, cases and inputs as ``test_deferred_norm_controls.py`` (real layer-6 /
+layer-7 weights at TP8 rank 0, six requests natively prefilled with 2,600
+tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6, both call paths, ``absorbed_query`` and
+``absorbed_query,shared_indexer_absorbed_query``, Gaussian and hot-channel residual rows).
+
+The null family ``S_k`` (``k = 1..K``) models every correct implementation
+of "RMSNorm, then the projection" with one bf16 rounding of the normalized
+input and one of the output: native code whose projection outputs (``kva``
+and, on the full-indexer layer, ``kw``) are ``bf16(W . h_k)`` with ``h_k =
+h* + u_k * ulp(h*)``, where ``h* = x * gamma * r*`` is the exact normalized
+hidden state (FP64), ``ulp`` the bf16 ulp at ``h*`` and ``u_k`` i.i.d.
+uniform on [-1/2, 1/2]. Native (``h = bf16((x * r) * gamma)``, then the
+GEMM) and deferred_norm (``a = bf16(x * gamma)``, the GEMM, then ``r``) are both of
+this form; their rounding errors are deterministic draws, the family's are
+random draws.
+
+Per case, all on the same cache snapshot:
+
+* ``native``, ``rerun``, ``deferred_norm``: as in ``test_deferred_norm_controls.py`` (the
+  deferred_norm projection on the layers it serves, ``dn_integration.SERVED_LAYERS``: the target's
+  deferred-norm scope is ``shared``).
+* ``ref``: native code on the correctly rounded FP64 oracle projection.
+* ``swap``: native code whose projection outputs on the served layers are deferred_norm's
+  projection kernel outputs (``DeferredNormProjection.project`` on the same packet). It
+  differs from the ``deferred_norm`` arm only in the code path around the projection.
+* ``S_1 .. S_K``.
+
+Reported per case, arm and layer: attention-output relRMS vs ``ref`` and vs
+``native`` (overall and per token row), FP8 query bytes and mean top-k
+Jaccard vs ``ref``; the null family's order statistics and the counts of
+draws at or above ``deferred_norm`` / ``native``.
+
+``test_projection_null`` reports, per data set, served layer and projection output, the
+projection-level error ``e = ||y - y*|| / ||y*||`` of native, deferred_norm and
+``K_PROJ`` draws of the family; per-segment errors; the bias and the
+largest error in units of the bf16 ulp at ``y*``; and each kernel against
+its own arithmetic specification evaluated in FP64 (native: ``bf16(W .
+h)``; deferred_norm: ``bf16(r . (W . a))`` with the hop-order ``r``).
+
+Assertions are harness facts only (hooks fired once per projection and run,
+deferred_norm engaged on the served layers only, finite outputs). The decision rule that uses these numbers was
+written before the data.
+"""
+
+import contextlib
+import json
+import os
+
+import deferred_norm_utils as U
+import pytest
+import test_absorbed_query_attention_integration as aq_integration
+import test_deferred_norm_attention_integration as dn_integration
+import test_deferred_norm_controls as controls
+import torch
+
+LAYERS = aq_integration.LAYERS
+LAYOUTS = aq_integration.LAYOUTS
+K_ATTN = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_K", "48"))
+K_PROJ = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_K_PROJ", "1000"))
+SEED_BASE = 7000
+
+pytestmark = dn_integration.pytestmark
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = dn_integration._DeferredNormHarness()
+    yield h
+    dn_integration._set_features("")
+
+
+def _kw_weight(h, layer):
+    indexer = h.attn[layer].mqa.indexer
+    return None if indexer is None else indexer._fused_wk_wp_weight
+
+
+def _weights64(h):
+    """Per layer: name -> FP64 copy of the projection weight."""
+    out = {}
+    for layer in LAYERS:
+        out[layer] = {"kva": h.attn[layer].kv_a_proj_with_mqa.weight.double()}
+        w = _kw_weight(h, layer)
+        if w is not None:
+            out[layer]["kw"] = w.double()
+    return out
+
+
+def _h_star(p):
+    """Exact normalized hidden state ``x * gamma * r*`` (FP64)."""
+    xd = p.residual.double()
+    r = torch.rsqrt(xd.pow(2).mean(-1, keepdim=True) + p.eps)
+    return xd * p.gamma.double() * r
+
+
+def _bf16_ulp64(t):
+    """bf16 ulp at ``|t|`` (FP64; 0 where ``t == 0``)."""
+    _, e = torch.frexp(t)
+    ulp = torch.exp2((e - 8).to(torch.float64))
+    return torch.where(t == 0, torch.zeros_like(t), ulp)
+
+
+def _draw(hstar, ulp, w64, seed):
+    """One member of the null family: name -> ``bf16(W . (h* + u ulp))``."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    u = torch.rand(hstar.shape, generator=g, device="cuda", dtype=torch.float64) - 0.5
+    hk = hstar + u * ulp
+    return {name: (hk @ w.t()).to(torch.bfloat16) for name, w in w64.items()}
+
+
+def _seed(k, layer):
+    return SEED_BASE + 97 * k + layer
+
+
+@contextlib.contextmanager
+def _projection_override(h, kva, kw):
+    """Native step with ``kv_a_proj_with_mqa`` returning ``kva[layer]`` and the
+    indexer's fused ``wk|weights_proj`` GEMM returning ``kw[layer]``."""
+    f = torch.nn.functional
+    orig_linear = f.linear
+    by_weight = {id(_kw_weight(h, layer)): (layer, t) for layer, t in kw.items()}
+    calls = {f"kva{layer}": 0 for layer in kva}
+    calls.update({f"kw{layer}": 0 for layer in kw})
+
+    def make_forward(layer):
+        def forward(input, *args, **kwargs):
+            t = kva[layer]
+            assert input.shape[0] == t.shape[0], (input.shape, t.shape)
+            calls[f"kva{layer}"] += 1
+            return t.clone()
+
+        return forward
+
+    def linear(input, weight, bias=None):
+        hit = by_weight.get(id(weight))
+        if hit is not None and bias is None:
+            layer, t = hit
+            assert input.shape[0] == t.shape[0], (input.shape, t.shape)
+            calls[f"kw{layer}"] += 1
+            return t.clone()
+        return orig_linear(input, weight, bias)
+
+    for layer in kva:
+        h.attn[layer].kv_a_proj_with_mqa.forward = make_forward(layer)
+    f.linear = linear
+    try:
+        yield calls
+    finally:
+        f.linear = orig_linear
+        for layer in kva:
+            del h.attn[layer].kv_a_proj_with_mqa.forward
+
+
+def _split(per_layer):
+    kva = {layer: d["kva"] for layer, d in per_layer.items()}
+    kw = {layer: d["kw"] for layer, d in per_layer.items() if d.get("kw") is not None}
+    return kva, kw
+
+
+def _deferred_norm_projection_outputs(h, packets):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+
+    out = {}
+    for layer in dn_integration.SERVED_LAYERS:
+        attn = h.attn[layer]
+        ok, reason, kw_weight = projection_eligibility(attn)
+        assert ok, reason
+        proj = DeferredNormProjection(attn, kw_weight=kw_weight)
+        assert proj.supported() is None, proj.supported()
+        kva, kw = proj.project(packets[layer])
+        out[layer] = {"kva": kva, "kw": kw}
+    torch.cuda.synchronize()
+    return out
+
+
+def _row_rel(x, ref):
+    xd = x.double().reshape(x.shape[0], -1)
+    rd = ref.double().reshape(ref.shape[0], -1)
+    return ((xd - rd).pow(2).mean(1).sqrt() / rd.pow(2).mean(1).sqrt().clamp_min(1e-30)).tolist()
+
+
+def _order_stats(v):
+    s = sorted(v)
+    n = len(s)
+    return {
+        "min": s[0],
+        "p10": s[int(0.1 * (n - 1))],
+        "median": s[(n - 1) // 2],
+        "p90": s[int(0.9 * (n - 1))],
+        "max": s[-1],
+        "mean": sum(s) / n,
+    }
+
+
+@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
+@pytest.mark.parametrize(
+    "features",
+    [
+        "absorbed_query,post_moe_hop,deferred_norm",
+        "absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm",
+    ],
+)
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_attention_null(harness, layout, path, features, hot):
+    h = harness
+    base = ",".join(f for f in features.split(",") if f != "deferred_norm")
+    packets, pos = controls._packets(h, layout, 3 if not hot else 13, hot)
+    native_in = {layer: p.materialize() for layer, p in packets.items()}
+    w64 = _weights64(h)
+    hstar = {layer: _h_star(p) for layer, p in packets.items()}
+    ulp = {layer: _bf16_ulp64(v) for layer, v in hstar.items()}
+    oracle = {
+        layer: {
+            name: U.fp64_oracle(p.residual, p.gamma, w, p.eps).to(torch.bfloat16)
+            for name, w in w64[layer].items()
+        }
+        for layer, p in packets.items()
+    }
+    swap = _deferred_norm_projection_outputs(h, packets)
+    snap = h.snapshot()
+    runs = {}
+    try:
+
+        def native_run(hidden=native_in):
+            dn_integration._reset_runners(h)
+            return dn_integration._run(h, base, layout, path, hidden, pos, snap)
+
+        def override_run(per_layer):
+            kva, kw = _split(per_layer)
+            with _projection_override(h, kva, kw) as calls:
+                run = native_run()
+            assert all(v == 1 for v in calls.values()), calls
+            return run
+
+        runs["native"] = native_run()
+        dn_integration._reset_runners(h)
+        runs["deferred_norm"] = dn_integration._run(
+            h, features, layout, path, dn_integration._served(packets), pos, snap
+        )
+        for layer in LAYERS:
+            engaged = h.attn[layer]._deferred_norm_projection not in (None, False)
+            assert engaged == (layer in dn_integration.SERVED_LAYERS), f"deferred_norm on L{layer}"
+        runs["rerun"] = native_run()
+        runs["ref"] = override_run(oracle)
+        runs["swap"] = override_run(swap)
+        null = []
+        for k in range(1, K_ATTN + 1):
+            draw = {
+                layer: _draw(hstar[layer], ulp[layer], w64[layer], _seed(k, layer))
+                for layer in LAYERS
+            }
+            null.append(override_run(draw))
+
+        ref_out, nat_out = runs["ref"][0], runs["native"][0]
+        report = {
+            "layout": layout,
+            "path": path,
+            "features": features,
+            "data": "hot" if hot else "gauss",
+            "k": K_ATTN,
+            "layers": {},
+        }
+        for layer in LAYERS:
+            for run in list(runs.values()) + null:
+                assert torch.isfinite(run[0][layer].float()).all(), layer
+            arms = {}
+            for arm, run in runs.items():
+                vs_ref = dn_integration._output_stats(runs["ref"], run, layout)[layer]
+                arms[arm] = {
+                    "d_ref": aq_integration._rel_rms(run[0][layer], ref_out[layer]),
+                    "d_nat": aq_integration._rel_rms(run[0][layer], nat_out[layer]),
+                    "rows_d_ref": _row_rel(run[0][layer], ref_out[layer]),
+                    "quant_q_bytes_diff_vs_ref": vs_ref.get("quant_q_bytes_diff_frac"),
+                    "topk_jaccard_vs_ref": vs_ref.get("topk_mean_jaccard"),
+                }
+            d_ref = [aq_integration._rel_rms(r[0][layer], ref_out[layer]) for r in null]
+            d_nat = [aq_integration._rel_rms(r[0][layer], nat_out[layer]) for r in null]
+            rows = [_row_rel(r[0][layer], ref_out[layer]) for r in null]
+            vs_ref = [dn_integration._output_stats(runs["ref"], r, layout)[layer] for r in null]
+            q_bytes = [s.get("quant_q_bytes_diff_frac") for s in vs_ref]
+            jac = [s.get("topk_mean_jaccard") for s in vs_ref if "topk_mean_jaccard" in s]
+            nrows = len(rows[0])
+            null_stats = {
+                "d_ref": _order_stats(d_ref),
+                "d_nat": _order_stats(d_nat),
+                "d_ref_values": d_ref,
+                "d_nat_values": d_nat,
+                "rows_d_ref": [_order_stats([r[i] for r in rows]) for i in range(nrows)],
+                "quant_q_bytes_diff_vs_ref": _order_stats(q_bytes),
+            }
+            if jac:
+                null_stats["topk_jaccard_vs_ref"] = _order_stats(jac)
+            for arm in ("deferred_norm", "native", "swap"):
+                arms[arm]["null_ge_d_ref"] = sum(1 for v in d_ref if v >= arms[arm]["d_ref"])
+                arms[arm]["null_ge_d_nat"] = sum(1 for v in d_nat if v >= arms[arm]["d_nat"])
+            arms["deferred_norm"]["d_swap"] = aq_integration._rel_rms(
+                runs["deferred_norm"][0][layer], runs["swap"][0][layer]
+            )
+            report["layers"][str(layer)] = {"arms": arms, "null": null_stats}
+        print(
+            "GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS " + json.dumps(report, sort_keys=True),
+            flush=True,
+        )
+        out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_OUT")
+        if out_dir:
+            os.makedirs(out_dir, exist_ok=True)
+            name = (
+                f"null-controls-{layout}-{path}-{features.replace(',', '_')}-{report['data']}.json"
+            )
+            with open(os.path.join(out_dir, name), "w") as fh:
+                json.dump(report, fh, indent=1, sort_keys=True)
+    finally:
+        h.restore(snap)
+        dn_integration._set_features("")
+
+
+def _segments(h, layer, name, n):
+    attn = h.attn[layer]
+    if name == "kva":
+        kv = attn.kv_lora_rank
+        rope = getattr(attn, "qk_rope_head_dim", 64)
+        q = n - kv - rope
+        return {"q_a": (0, q), "kv_a": (q, q + kv), "k_rope": (q + kv, n)}
+    wk = attn.mqa.indexer.head_dim
+    return {"wk": (0, wk), "weights_proj": (wk, n)}
+
+
+def _audit(y, ystar, segments):
+    """Error statistics of ``y`` vs the FP64 oracle ``y*``."""
+    yd = y.double()
+    ulp = _bf16_ulp64(ystar)
+    err_ulp = torch.where(ulp > 0, (yd - ystar) / ulp.clamp_min(1e-300), torch.zeros_like(yd))
+    nearest = ystar.to(torch.bfloat16)
+    out = {
+        "e": U.norm_rms_err(y, ystar),
+        "frac_not_nearest": float((y != nearest).float().mean()),
+        "max_err_ulp": float(err_ulp.abs().max()),
+        "mean_err_ulp": float(err_ulp.mean()),
+        "frac_err_ge_1ulp": float((err_ulp.abs() >= 1.0).float().mean()),
+        "seg_e": {s: U.norm_rms_err(y[:, a:b], ystar[:, a:b]) for s, (a, b) in segments.items()},
+        "row_e": [U.norm_rms_err(y[i], ystar[i]) for i in range(y.shape[0])],
+    }
+    return out
+
+
+def _vs_spec(y, spec):
+    ulp = U.bf16_ulp_distance(y, spec)
+    return {"max_ulp": int(ulp.max()), "frac_diff": float((ulp > 0).float().mean())}
+
+
+@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
+@torch.inference_mode()
+def test_projection_null(harness, hot):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    h = harness
+    packets, _ = controls._packets(h, "6req_x1tok", 3 if not hot else 13, hot)
+    w64 = _weights64(h)
+    swap = _deferred_norm_projection_outputs(h, packets)
+    report = {"data": "hot" if hot else "gauss", "k": K_PROJ, "layers": {}}
+    for layer in dn_integration.SERVED_LAYERS:
+        p = packets[layer]
+        rows, hidden = p.residual.shape
+        hstar = _h_star(p)
+        ulp = _bf16_ulp64(hstar)
+        h_nat = p.materialize()
+        a = rows_from_frag(p.a_frag, rows, hidden)
+        r_w = torch.rsqrt(U.square_sum_in_order(p.sumsq) / hidden + p.eps)
+        layer_rep = {}
+        for name, w in w64[layer].items():
+            ystar = U.fp64_oracle(p.residual, p.gamma, w, p.eps)
+            segs = _segments(h, layer, name, w.shape[0])
+            if name == "kva":
+                y_nat = h.attn[layer].kv_a_proj_with_mqa(h_nat)
+            else:
+                y_nat = torch.nn.functional.linear(h_nat, _kw_weight(h, layer))
+            y_deferred_norm = swap[layer][name]
+            spec_nat = (h_nat.double() @ w.t()).to(torch.bfloat16)
+            spec_deferred_norm = (r_w.double().unsqueeze(-1) * (a.double() @ w.t())).to(
+                torch.bfloat16
+            )
+            e_null = []
+            for k in range(1, K_PROJ + 1):
+                g = torch.Generator(device="cuda").manual_seed(_seed(k, layer))
+                u = torch.rand(hstar.shape, generator=g, device="cuda", dtype=torch.float64) - 0.5
+                y = ((hstar + u * ulp) @ w.t()).to(torch.bfloat16)
+                e_null.append(U.norm_rms_err(y, ystar))
+            draw1 = _draw(hstar, ulp, {name: w}, _seed(1, layer))[name]
+            e_deferred_norm = U.norm_rms_err(y_deferred_norm, ystar)
+            e_nat = U.norm_rms_err(y_nat, ystar)
+            layer_rep[name] = {
+                "native": _audit(y_nat, ystar, segs),
+                "deferred_norm": _audit(y_deferred_norm, ystar, segs),
+                "ref": _audit(ystar.to(torch.bfloat16), ystar, segs),
+                "null_draw1": _audit(draw1, ystar, segs),
+                "null_e": _order_stats(e_null),
+                "null_ge_e_deferred_norm": sum(1 for v in e_null if v >= e_deferred_norm),
+                "null_ge_e_native": sum(1 for v in e_null if v >= e_nat),
+                "native_vs_spec": _vs_spec(y_nat, spec_nat),
+                "deferred_norm_vs_spec": _vs_spec(y_deferred_norm, spec_deferred_norm),
+                "deferred_norm_vs_native": _vs_spec(y_deferred_norm, y_nat),
+            }
+        report["layers"][str(layer)] = layer_rep
+    torch.cuda.synchronize()
+    print(
+        "GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_PROJ " + json.dumps(report, sort_keys=True),
+        flush=True,
+    )
+    out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_OUT")
+    if out_dir:
+        os.makedirs(out_dir, exist_ok=True)
+        with open(os.path.join(out_dir, f"null-controls-proj-{report['data']}.json"), "w") as fh:
+            json.dump(report, fh, indent=1, sort_keys=True)
+    for layer, d in report["layers"].items():
+        for name, s in d.items():
+            assert s["native"]["e"] == s["native"]["e"] and s["null_e"]["max"] < 0.05, (layer, name)
+
+
+def _fp32_reference(op, w64, scale=None):
+    """The same specification through an FP32 cuBLAS GEMM (TF32 off): a
+    third accumulation order of the same products."""
+    prev = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    try:
+        y = op.float() @ w64.float().t()
+        if scale is not None:
+            y = y * scale.unsqueeze(-1)
+        return y.to(torch.bfloat16)
+    finally:
+        torch.backends.cuda.matmul.allow_tf32 = prev
+
+
+def _accumulation_audit(y, s, t_abs, ystar, max_list=16):
+    """``y`` (bf16) against its exact pre-rounding specification ``s`` (FP64).
+
+    ``rho`` is the error beyond the output rounding (half a bf16 ulp at
+    ``max(|s|, |y|)``) and beyond a ``2**-21 |s|`` allowance for the scale
+    ``r`` (``rsqrt.approx``), in units of ``2**-24 * sum_k |t_jk|``: the
+    scale of FP32 accumulation error of the dot product's terms ``t_jk``.
+    """
+    yd = y.double()
+    err = (yd - s).abs()
+    half = 0.5 * _bf16_ulp64(torch.maximum(s.abs(), yd.abs()))
+    excess = (err - half - 2.0**-21 * s.abs()).clamp_min(0.0)
+    rho = excess / (2.0**-24 * t_abs).clamp_min(1e-300)
+    ulp = U.bf16_ulp_distance(y, s.to(torch.bfloat16))
+    big = (ulp > 1).nonzero().tolist()
+    flat = rho.flatten()
+    k999 = max(1, int(0.999 * flat.numel()))
+    out = {
+        "frac_diff": float((ulp > 0).float().mean()),
+        "max_ulp": int(ulp.max()),
+        "n_gt1ulp": len(big),
+        "rho_max": float(rho.max()),
+        "rho_p999": float(flat.kthvalue(k999).values),
+        "gt1ulp": [
+            {
+                "row": i,
+                "col": j,
+                "ulp": int(ulp[i, j]),
+                "cancel": float(s[i, j].abs() / t_abs[i, j]),
+                "rho": float(rho[i, j]),
+            }
+            for i, j in big[:max_list]
+        ],
+        "bias_rel": float((yd - ystar).mean() / ystar.pow(2).mean().sqrt()),
+        "frac_above_oracle": float(
+            ((yd - ystar) > 0).double().sum() / ((yd - ystar) != 0).double().sum().clamp_min(1)
+        ),
+    }
+    return out
+
+
+@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
+@torch.inference_mode()
+def test_projection_spec_accumulation(harness, hot):
+    """Each projection kernel (deferred_norm's and native's) against its own exact
+    arithmetic specification, with the error measured in units of the FP32
+    accumulation scale of its terms; next to an FP32 cuBLAS reference of the
+    same specification."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    h = harness
+    packets, _ = controls._packets(h, "6req_x1tok", 3 if not hot else 13, hot)
+    w64 = _weights64(h)
+    deferred_norm_out = _deferred_norm_projection_outputs(h, packets)
+    report = {"data": "hot" if hot else "gauss", "layers": {}}
+    for layer in dn_integration.SERVED_LAYERS:
+        p = packets[layer]
+        rows, hidden = p.residual.shape
+        h_nat = p.materialize()
+        a = rows_from_frag(p.a_frag, rows, hidden)
+        r_w = torch.rsqrt(U.square_sum_in_order(p.sumsq) / hidden + p.eps)
+        layer_rep = {}
+        for name, w in w64[layer].items():
+            ystar = U.fp64_oracle(p.residual, p.gamma, w, p.eps)
+            if name == "kva":
+                y_nat = h.attn[layer].kv_a_proj_with_mqa(h_nat)
+            else:
+                y_nat = torch.nn.functional.linear(h_nat, _kw_weight(h, layer))
+            s_deferred_norm = r_w.double().unsqueeze(-1) * (a.double() @ w.t())
+            t_deferred_norm = r_w.double().unsqueeze(-1) * (a.double().abs() @ w.abs().t())
+            s_nat = h_nat.double() @ w.t()
+            t_nat = h_nat.double().abs() @ w.abs().t()
+            layer_rep[name] = {
+                "deferred_norm_kernel": _accumulation_audit(
+                    deferred_norm_out[layer][name], s_deferred_norm, t_deferred_norm, ystar
+                ),
+                "deferred_norm_fp32_ref": _accumulation_audit(
+                    _fp32_reference(a, w, r_w), s_deferred_norm, t_deferred_norm, ystar
+                ),
+                "native_kernel": _accumulation_audit(y_nat, s_nat, t_nat, ystar),
+                "native_fp32_ref": _accumulation_audit(
+                    _fp32_reference(h_nat, w), s_nat, t_nat, ystar
+                ),
+            }
+        report["layers"][str(layer)] = layer_rep
+    torch.cuda.synchronize()
+    print(
+        "GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_SPEC " + json.dumps(report, sort_keys=True),
+        flush=True,
+    )
+    out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_OUT")
+    if out_dir:
+        os.makedirs(out_dir, exist_ok=True)
+        with open(os.path.join(out_dir, f"null-controls-spec-{report['data']}.json"), "w") as fh:
+            json.dump(report, fh, indent=1, sort_keys=True)
+    for layer, d in report["layers"].items():
+        for name, s in d.items():
+            for impl, v in s.items():
+                assert v["rho_max"] == v["rho_max"], (layer, name, impl)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py
new file mode 100644
index 000000000000..da9db76e5168
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py
@@ -0,0 +1,293 @@
+# 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.
+"""Single-GPU tests of the deferred_norm staged projection (real GLM-5.2 weights).
+
+Native chain: ``h = bf16((x * r) * gamma)`` (the post_moe_hop's output, produced
+here by the deferred_norm materialization kernel, which runs post_moe_hop's arithmetic) followed
+by the native ``kv_a_proj_with_mqa`` (``dsv3_fused_a_gemm_op``).
+deferred_norm: the staged projection on ``a = bf16(x * gamma)`` and the hop's
+square-sum partials (emulated bit-exactly in torch).
+
+The target's deferred-norm scope is ``shared``: the projection serves the shared-indexer
+layers (``LAYERS``); a full-indexer layer (``FULL_INDEXER_LAYERS``) keeps its native
+projection, and the projection's eligibility check declines it.
+
+Frozen error budget: per output tensor
+``allclose(y_deferred_norm, y_native, rtol=0.1, atol=1e-2)`` and normalized RMS error
+vs the FP64 oracle ``e_deferred_norm <= 1.25 * e_native + 2**-12``.
+
+Fixture label: weights real (checkpoint), activations synthetic.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
+
+import absorbed_query_fixture  # noqa: E402
+import deferred_norm_utils as U  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
+    reason="needs a GPU and the GLM-5.2 checkpoint",
+)
+
+LAYERS = [int(x) for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_TEST_LAYERS", "7").split(",")]
+FULL_INDEXER_LAYERS = [
+    int(x) for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_TEST_FULL_LAYERS", "6").split(",")
+]
+ROWS = [int(x) for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_TEST_ROWS", "1,6,8").split(",")]
+_MODS = {}
+
+
+def _module(layer):
+    if layer not in _MODS:
+        with torch.inference_mode():
+            cfg = absorbed_query_fixture.build_model_config(0)
+            _MODS[layer] = (
+                absorbed_query_fixture.build_native_attention(layer, 0, cfg),
+                absorbed_query_fixture.load_input_layernorm(layer).cuda(),
+            )
+    return _MODS[layer]
+
+
+def make_residual(rows, seed, scale=1.0):
+    """Residual-stream-like rows: unit Gaussian plus a few large channels."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(rows, U.H, generator=g, device="cuda") * 2.0
+    hot = torch.randint(0, U.H, (12,), generator=g, device="cuda")
+    x[:, hot] *= 30.0
+    return (x * scale).bfloat16()
+
+
+def make_packet(x, gamma):
+    return U.make_frag_packet(x, gamma)
+
+
+def _projection(attn):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+
+    ok, reason, kw = projection_eligibility(attn)
+    assert ok, reason
+    proj = DeferredNormProjection(attn, kw_weight=kw)
+    assert proj.supported() is None, proj.supported()
+    return proj
+
+
+def _native(attn, h):
+    kva = attn.kv_a_proj_with_mqa(h)
+    kw = None
+    if attn.mqa.indexer is not None:
+        kw = torch.nn.functional.linear(h, attn.mqa.indexer._fused_wk_wp_weight)
+    return kva, kw
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+def test_served_configs_cover_the_consumer(layer):
+    """``precompile_served`` (run by the producing hop before its cross-rank
+    agreement) compiles exactly the kernel the consumer launches on this
+    layer kind, so the consumer's ``precompile`` is a cache hit."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.proj_kernel import (  # noqa: E501
+        is_compiled,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        precompile_served,
+        served_kernel_configs,
+    )
+
+    device = torch.cuda.current_device()
+    assert precompile_served(device) is None
+    proj = _projection(_module(layer)[0])
+    assert proj.served_config() in served_kernel_configs(device)
+    assert is_compiled(**proj.served_config())
+
+
+@pytest.mark.parametrize("scale", [1.0, 1e-3, 30.0])
+@pytest.mark.parametrize("rows", ROWS)
+@pytest.mark.parametrize("layer", LAYERS)
+def test_projection_within_frozen_budget(layer, rows, scale):
+    attn, gamma = _module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        x = make_residual(rows, seed=100 + rows + layer, scale=scale)
+        packet = make_packet(x, gamma)
+        h = packet.materialize()
+        kva_n, kw_n = _native(attn, h)
+        kva_w, kw_w = proj.project(packet)
+        torch.cuda.synchronize()
+        outs = [("kva", kva_w, kva_n, attn.kv_a_proj_with_mqa.weight)]
+        if kw_n is not None:
+            outs.append(("kw", kw_w, kw_n, attn.mqa.indexer._fused_wk_wp_weight))
+        for name, cand, nat, w in outs:
+            oracle = U.fp64_oracle(x, gamma, w)
+            stats = U.compare(cand, nat, oracle)
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_BUDGET layer={layer} rows={rows} scale={scale} out={name} {stats}"
+            )
+            assert U.within_budget(stats), (name, stats)
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+def test_native_input_mode(layer):
+    """Same kernel on the normalized hidden with r = 1 (attribution arm)."""
+    attn, gamma = _module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        x = make_residual(6, seed=7)
+        h = make_packet(x, gamma).materialize()
+        kva_n, kw_n = _native(attn, h)
+        kva_w, kw_w = proj.run(h, None, U.EPS)
+        torch.cuda.synchronize()
+        for name, cand, nat in (("kva", kva_w, kva_n), ("kw", kw_w, kw_n)):
+            if nat is None:
+                continue
+            ulp = U.bf16_ulp_distance(cand, nat)
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_NATIVE_INPUT layer={layer} out={name} max_ulp={int(ulp.max())} "
+                f"frac_diff={float((ulp > 0).float().mean()):.3e}"
+            )
+            assert torch.allclose(cand.float(), nat.float(), rtol=0.1, atol=1e-2)
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+def test_rows_match_single_row_runs(layer):
+    """Every token row equals a launch of that row alone (no cross-row mixing)."""
+    attn, gamma = _module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        x = make_residual(8, seed=11)
+        kva, kw = proj.project(make_packet(x, gamma))
+        for t in range(8):
+            kva1, kw1 = proj.project(make_packet(x[t : t + 1].contiguous(), gamma))
+            assert torch.equal(kva1[0], kva[t])
+            if kw is not None:
+                assert torch.equal(kw1[0], kw[t])
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+def test_graph_replay_matches_eager(layer):
+    attn, gamma = _module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        x0 = make_residual(6, seed=21)
+        static = U.empty_like_packet(make_packet(x0, gamma))
+        ref0 = make_packet(x0, gamma)
+        static.a_frag.copy_(ref0.a_frag)
+        static.sumsq.copy_(ref0.sumsq)
+        proj.project(static)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            g_kva, g_kw = proj.project(static)
+        for trial, replays in enumerate((1, 1, 50)):
+            x = make_residual(6, seed=300 + trial)
+            ref = make_packet(x, gamma)
+            static.a_frag.copy_(ref.a_frag)
+            static.sumsq.copy_(ref.sumsq)
+            for _ in range(replays):
+                graph.replay()
+            e_kva, e_kw = proj.project(ref)
+            torch.cuda.synchronize()
+            assert torch.equal(g_kva, e_kva)
+            if e_kw is not None:
+                assert torch.equal(g_kw, e_kw)
+
+
+@pytest.mark.parametrize("stage_before_wait", [True, False])
+@pytest.mark.parametrize("delay_ns", [0, 5000, 50000])
+@pytest.mark.parametrize("layer", LAYERS)
+def test_pdl_delayed_producer(layer, delay_ns, stage_before_wait):
+    """The producer releases its dependents at entry and writes the packet only
+    after ``delay_ns``; the projection (launched with PDL, staging its weights
+    before its dependency wait) must still read the final packet."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormProjConfig,
+        DeferredNormProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+
+    attn, gamma = _module(layer)
+    _, _, kwt = projection_eligibility(attn)
+    proj = DeferredNormProjection(
+        attn, DeferredNormProjConfig(stage_before_wait=stage_before_wait), kw_weight=kwt
+    )
+    if not proj.use_pdl:
+        pytest.skip("PDL disabled (TRTLLM_ENABLE_PDL=0)")
+    with torch.inference_mode():
+        for trial in range(4):
+            x = make_residual(6, seed=500 + trial)
+            src = make_packet(x, gamma)
+            dst = U.empty_like_packet(src)
+            dst.a_frag.fill_(float("nan"))
+            U.sumsq_full(dst).fill_(float("nan"))
+            torch.cuda.synchronize()
+            U.proxy_producer(
+                src.a_frag, dst.a_frag, U.sumsq_full(src), U.sumsq_full(dst), delay_ns, rows=8
+            )
+            kva, kw = proj.project(dst)
+            ref_kva, ref_kw = proj.project(src)
+            torch.cuda.synchronize()
+            assert torch.equal(kva, ref_kva), f"trial {trial}"
+            if kw is not None:
+                assert torch.equal(kw, ref_kw)
+
+
+def test_materialize_close_to_torch():
+    """The fallback's normalized output vs a torch RMSNorm of the same x
+    (bit-exactness against the post_moe_hop is asserted by the MPI tests)."""
+    _, gamma = _module(LAYERS[0])
+    with torch.inference_mode():
+        x = make_residual(6, seed=3)
+        h = make_packet(x, gamma).materialize()
+        s = U.square_sum_in_order(U.hop_sumsq_partials(x))
+        r = torch.rsqrt(s / U.H + U.EPS).unsqueeze(-1)
+        ref = ((x.float() * r) * gamma.float()).bfloat16()
+        assert int(U.bf16_ulp_distance(h, ref).max()) <= 1
+
+
+@pytest.mark.parametrize("layer", FULL_INDEXER_LAYERS)
+def test_full_indexer_layer_keeps_its_native_projection(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+
+    attn, _ = _module(layer)
+    assert attn.mqa.indexer is not None
+    ok, reason, kw = projection_eligibility(attn)
+    assert not ok and kw is None
+    assert reason.endswith("=shared: full-indexer layer keeps its native projection"), reason
+
+
+def test_eligibility_rejects_out_of_domain():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.proj_kernel import (  # noqa: E501
+        projection_supported,
+    )
+
+    assert projection_supported(6144, 2624, 152, 512, 18, 4) is None
+    assert projection_supported(6144, 2784, 152, 512, 18, 4) is None
+    assert projection_supported(6144, 2624, 64, 512, 18, 4) is not None  # 41 rows per CTA
+    assert projection_supported(6144, 2784, 152, 512, 17, 4) is not None  # 2 rows over staged
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py
new file mode 100644
index 000000000000..05718956b17c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py
@@ -0,0 +1,946 @@
+# 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.
+"""Collective tests of the deferred_norm deferred hop on routed_moe_into_hop's routed input (8 ranks, 2 hosts).
+
+With ``routed_moe_into_hop`` selected, routed_moe writes its per-rank routed result into the
+post_moe_hop hop's routed-input buffer (``PostMoeHop.enable_routed_input``; modes
+``fp32`` (default), ``identity``, ``row``) and the layer calls
+``PostMoeHop.hop_routed``. ``DeferredNormHop.hop_routed`` takes the same
+inputs and returns the deferred packet. The routed content is written here
+as routed_moe would write it (per rank), from a synthetic raw tuple.
+
+Exact invariants, per mode, against ``PostMoeHop.hop_routed`` on the same
+routed content and the same workspace:
+
+* ``residual_out`` bit-identical;
+* ``a_frag`` holds ``bf16(x * gamma)`` bitwise; ``sumsq`` equals the torch
+  emulation of the hop's order; ``materialize(packet)`` equals
+  ``hop_routed``'s normalized output bit for bit;
+* ``fp32`` mode: the hop clears the accumulator rows it read (the buffer is
+  zero after every call, as after post_moe_hop's).
+
+post_moe_hop's error protocol on the routed variant, mixed with post_moe_hop routed, post_moe_hop
+raw and deferred_norm raw calls on one workspace: missing rank (deadline, NaN, health check,
+reset), sticky poison carried through every Lamport buffer, fault-injected
+store order with per-rank kernel choice, contract violations before any
+device work, per-rank layout repairs (also in a CUDA graph), graph replay,
+the health check's call count, and the collective enable (post_moe_hop routed input
+off, compile failure on one rank, modes that differ across ranks, mode
+switch).
+"""
+
+import os
+import random
+import sys
+import time
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
+
+import deferred_norm_utils as U  # noqa: E402
+import post_moe_hop_mpi_utils as M  # noqa: E402
+
+MODES = ("fp32", "identity", "row")
+
+
+def _env():
+    env = M.mpi_env()
+    if env.world != 8:
+        pytest.skip("the deferred_norm hop is specialized to TP8")
+    return env
+
+
+def _routed_api():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    if not hasattr(PostMoeHop, "enable_routed_input"):
+        pytest.skip("the installed post_moe_hop hop has no routed_moe_into_hop routed input")
+
+
+_HOPS = {}
+
+
+def _hops(mode, trigger_early=True, timeout_s=None, fault_delay_ns=0):
+    """A post_moe_hop with its routed input in ``mode`` and a deferred_norm hop on it (collective)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    _routed_api()
+    key = (mode, timeout_s, fault_delay_ns)
+    if key not in _HOPS:
+        kw = {}
+        if timeout_s is not None:
+            kw["timeout_s"] = timeout_s
+        if fault_delay_ns:
+            kw["fault_delay_ns"] = fault_delay_ns
+        post_moe_hop = PostMoeHop(
+            env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False, **kw
+        )
+        assert M.all_ranks(post_moe_hop.engaged and post_moe_hop.enable_routed_input(mode)), (
+            post_moe_hop.reason
+        )
+        _HOPS[key] = {"post_moe_hop": post_moe_hop}
+    entry = _HOPS[key]
+    if trigger_early not in entry:
+        deferred_norm = DeferredNormHop(entry["post_moe_hop"], trigger_early=trigger_early)
+        assert M.all_ranks(deferred_norm.engaged and deferred_norm.enable_routed_input()), (
+            deferred_norm.reason,
+            deferred_norm.routed_reason,
+        )
+        entry[trigger_early] = deferred_norm
+    return entry["post_moe_hop"], entry[trigger_early]
+
+
+def _source(mode, rows, seed, with_shared=True):
+    """Per-rank routed content as routed_moe writes it in ``mode``, plus shared /
+    residual / gamma (the residual and gamma are replicated)."""
+    fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(
+        rows, seed=seed, with_shared=with_shared
+    )
+    if mode == "row":
+        src = M.reference_local_partial(fc2, scales, perm, None)
+    elif mode == "identity":
+        valid = perm >= 0
+        safe = torch.where(valid, perm, torch.zeros_like(perm)).long()
+        src = (
+            fc2[safe.view(-1)].contiguous(),
+            torch.where(valid, scales, torch.zeros_like(scales)).bfloat16().contiguous(),
+        )
+    else:
+        valid = (perm >= 0).unsqueeze(-1)
+        safe = torch.where(perm >= 0, perm, torch.zeros_like(perm)).long()
+        prod = fc2[safe.view(-1)].view(rows, M.GLM_TOP_K, -1).float() * scales.float().unsqueeze(-1)
+        src = torch.where(valid, prod, torch.zeros_like(prod)).sum(1).contiguous()
+    return src, shared, residual, gamma
+
+
+def _fill(post_moe_hop, mode, src, rows):
+    """Write ``src`` into the post_moe_hop's routed buffer (stream-ordered); return the view."""
+    if mode == "row":
+        buf = post_moe_hop.routed_input_buffer(rows)
+        buf.copy_(src)
+        return buf
+    if mode == "identity":
+        buf, scales = post_moe_hop.routed_raw_buffers(rows)
+        buf.copy_(src[0])
+        scales.copy_(src[1])
+        return buf
+    buf = post_moe_hop.routed_acc_buffer(rows)
+    buf.copy_(src)
+    return buf
+
+
+def _post_moe_hop_reference(post_moe_hop, mode, src, rows, shared, residual, gamma):
+    routed = _fill(post_moe_hop, mode, src, rows)
+    return post_moe_hop.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
+
+
+def _deferred_norm_call(post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma):
+    routed = _fill(post_moe_hop, mode, src, rows)
+    return deferred_norm.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
+
+
+def _acc_clear(post_moe_hop, mode):
+    if mode != "fp32":
+        return True
+    torch.cuda.synchronize()
+    return bool((post_moe_hop.routed_acc_buffer(post_moe_hop.max_rows) == 0).all())
+
+
+def _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
+    """This rank's packet invariants against post_moe_hop's outputs (not collective)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    torch.cuda.synchronize()
+    ok_res = torch.equal(res_deferred_norm, res_post_moe_hop)
+    scaled_ref = (res_deferred_norm.float() * gamma.float()).bfloat16()
+    ok_scaled = torch.equal(
+        rows_from_frag(packet.a_frag, res_deferred_norm.shape[0], res_deferred_norm.shape[1]),
+        scaled_ref,
+    )
+    ok_ss = torch.equal(packet.sumsq, U.hop_sumsq_partials(res_deferred_norm))
+    h = materialize(packet)
+    torch.cuda.synchronize()
+    ok_h = torch.equal(h, h_post_moe_hop)
+    ok = ok_res and ok_scaled and ok_ss and ok_h
+    if not ok:
+        print(
+            f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_FAIL {tag} rank={M.mpi_env().rank} res={ok_res} scaled={ok_scaled} "
+            f"sumsq={ok_ss} materialize={ok_h} "
+            f"max_ulp_h={int(U.bf16_ulp_distance(h, h_post_moe_hop).max())}",
+            flush=True,
+        )
+    return ok
+
+
+def _check(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
+    return M.all_ranks(
+        _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag)
+    )
+
+
+def _all_nan_packet(packet, res):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    torch.cuda.synchronize()
+    a = rows_from_frag(packet.a_frag, res.shape[0], res.shape[1])
+    h = materialize(packet)
+    torch.cuda.synchronize()
+    return all(bool(torch.isnan(t.float()).all()) for t in (res, a, packet.sumsq, h))
+
+
+def _reset(post_moe_hop):
+    assert M.all_ranks(post_moe_hop.reset())
+    assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+# ---------------------------------------------------------------------------
+# Collective enable
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("mode", MODES)
+def test_routed_engaged(mode):
+    post_moe_hop, deferred_norm = _hops(mode)
+    assert M.all_ranks(deferred_norm.routed_engaged and deferred_norm.routed_mode == mode)
+    md = type("Md", (), {"num_contexts": 0, "num_generations": 6})()
+    ctx = type("Md", (), {"num_contexts": 1, "num_generations": 0})()
+    assert M.all_ranks(
+        deferred_norm.can_run_routed(6, md)
+        and not deferred_norm.can_run_routed(6, ctx)
+        and not deferred_norm.can_run_routed(9, md)
+        and not deferred_norm.can_run_routed(0, md)
+    )
+
+
+def test_routed_requires_post_moe_hop_routed_input():
+    """Without the post_moe_hop's routed input the routed variant stays off on every rank."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    _routed_api()
+    post_moe_hop = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
+    deferred_norm = DeferredNormHop(post_moe_hop)
+    ok = (
+        deferred_norm.engaged
+        and not deferred_norm.enable_routed_input()
+        and not deferred_norm.routed_engaged
+    )
+    ok &= "routed input is not enabled" in deferred_norm.routed_reason
+    md = type("Md", (), {"num_contexts": 0, "num_generations": 6})()
+    ok &= not deferred_norm.can_run_routed(6, md)
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source("fp32", 6, seed=31)
+        with pytest.raises(ValueError):
+            deferred_norm.hop_routed(src, shared, residual, gamma, M.GLM_EPS)
+    assert M.all_ranks(ok), deferred_norm.routed_reason
+
+
+@pytest.mark.parametrize("failing_rank", [2, 7])
+def test_routed_compile_failure_disengages_every_rank(monkeypatch, failing_rank):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+        hop_kernel,
+    )
+
+    env = _env()
+    post_moe_hop, _ = _hops("fp32")
+    deferred_norm = DeferredNormHop(post_moe_hop)
+    if env.rank == failing_rank:
+        real = hop_kernel.compiled_hop
+
+        def _fail(*args, **kwargs):
+            if kwargs.get("routed_input"):
+                raise RuntimeError("injected routed compile failure")
+            return real(*args, **kwargs)
+
+        monkeypatch.setattr(hop_kernel, "compiled_hop", _fail)
+    engaged = deferred_norm.enable_routed_input()
+    if env.rank == failing_rank:
+        local = "injected routed compile failure" in deferred_norm.routed_reason
+    else:
+        local = deferred_norm.routed_reason == "another rank is not capable"
+    assert M.all_ranks(
+        deferred_norm.engaged and not engaged and not deferred_norm.routed_engaged and local
+    ), deferred_norm.routed_reason
+
+
+def test_routed_modes_must_agree_across_ranks():
+    """post_moe_hop hops whose routed-input modes differ across ranks (environment)
+    leave the deferred_norm routed variant off on every rank."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    _routed_api()
+    post_moe_hop = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
+    assert M.all_ranks(post_moe_hop.enable_routed_input("fp32" if env.rank % 2 else "identity"))
+    deferred_norm = DeferredNormHop(post_moe_hop)
+    engaged = deferred_norm.enable_routed_input()
+    assert M.all_ranks(not engaged and "modes differ" in deferred_norm.routed_reason), (
+        deferred_norm.routed_reason
+    )
+
+
+def test_routed_mode_switch_needs_reenable():
+    """After the post_moe_hop switches modes (collective) the routed variant
+    declines until it is enabled again, then serves the new mode exactly."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    _routed_api()
+    post_moe_hop = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
+    assert M.all_ranks(post_moe_hop.enable_routed_input("fp32"))
+    deferred_norm = DeferredNormHop(post_moe_hop)
+    assert M.all_ranks(deferred_norm.enable_routed_input() and deferred_norm.routed_mode == "fp32")
+    md = type("Md", (), {"num_contexts": 0, "num_generations": 6})()
+    assert M.all_ranks(post_moe_hop.enable_routed_input("identity"))
+    ok = not deferred_norm.can_run_routed(6, md)
+    ok &= deferred_norm.enable_routed_input() and deferred_norm.routed_mode == "identity"
+    ok &= deferred_norm.can_run_routed(6, md)
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source("identity", 6, seed=33)
+        h_post_moe_hop, res_post_moe_hop = _post_moe_hop_reference(
+            post_moe_hop, "identity", src, 6, shared, residual, gamma
+        )
+        packet, res = _deferred_norm_call(
+            post_moe_hop, deferred_norm, "identity", src, 6, shared, residual, gamma
+        )
+        ok &= _verdict(packet, res, h_post_moe_hop, res_post_moe_hop, gamma, "mode-switch")
+    assert M.all_ranks(ok)
+
+
+# ---------------------------------------------------------------------------
+# Exactness against post_moe_hop's hop_routed
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("trigger_early", [True, False])
+@pytest.mark.parametrize("with_shared", [True, False])
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+@pytest.mark.parametrize("mode", MODES)
+def test_routed_packet_matches_hop_routed(mode, rows, with_shared, trigger_early):
+    post_moe_hop, deferred_norm = _hops(mode, trigger_early)
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, rows, seed=50 + rows, with_shared=with_shared)
+        h_post_moe_hop, res_post_moe_hop = _post_moe_hop_reference(
+            post_moe_hop, mode, src, rows, shared, residual, gamma
+        )
+        cleared_post_moe_hop = _acc_clear(post_moe_hop, mode)
+        packet, res_deferred_norm = _deferred_norm_call(
+            post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma
+        )
+        cleared_deferred_norm = _acc_clear(post_moe_hop, mode)
+        ok = _verdict(
+            packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, f"{mode} R={rows}"
+        )
+        assert M.all_ranks(ok and cleared_post_moe_hop and cleared_deferred_norm), (
+            cleared_post_moe_hop,
+            cleared_deferred_norm,
+        )
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize("mode", MODES)
+def test_routed_large_residual_and_zero_routed(mode):
+    post_moe_hop, deferred_norm = _hops(mode)
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, 6, seed=77)
+        residual = (residual.float() * 50.0).bfloat16()
+        if mode == "identity":
+            src = (src[0], torch.zeros_like(src[1]))
+        else:
+            src = torch.zeros_like(src)
+        h_post_moe_hop, res_post_moe_hop = _post_moe_hop_reference(
+            post_moe_hop, mode, src, 6, shared, residual, gamma
+        )
+        packet, res = _deferred_norm_call(
+            post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma
+        )
+        assert _check(packet, res, h_post_moe_hop, res_post_moe_hop, gamma, f"{mode}-large-res")
+
+
+@pytest.mark.parametrize("mode", MODES)
+def test_back_to_back_routed_and_raw_mixed_with_post_moe_hop(mode):
+    """deferred_norm routed, post_moe_hop routed, deferred_norm raw and post_moe_hop raw calls interleaved on one
+    workspace without host syncs (each routed call refills the buffer on
+    the stream, as routed_moe does)."""
+    post_moe_hop, deferred_norm = _hops(mode)
+    kinds = [
+        "routed_deferred_norm",
+        "routed_post_moe_hop",
+        "deferred_norm",
+        "post_moe_hop",
+        "routed_deferred_norm",
+        "routed_deferred_norm",
+        "post_moe_hop",
+        "deferred_norm",
+        "routed_post_moe_hop",
+        "routed_deferred_norm",
+        "deferred_norm",
+        "routed_post_moe_hop",
+    ]
+    rows_seq = [6, 1, 8, 3, 6, 2, 6, 8, 5, 6, 1, 6]
+    with torch.inference_mode():
+        inputs = []
+        for i, (kind, r) in enumerate(zip(kinds, rows_seq)):
+            if kind.startswith("routed_"):
+                inputs.append(_source(mode, r, seed=900 + i))
+            else:
+                inputs.append(M.make_raw_tuple(r, seed=900 + i))
+        outs = []
+        for kind, r, t in zip(kinds, rows_seq, inputs):
+            if kind == "routed_deferred_norm":
+                outs.append(_deferred_norm_call(post_moe_hop, deferred_norm, mode, t[0], r, *t[1:]))
+            elif kind == "routed_post_moe_hop":
+                outs.append(_post_moe_hop_reference(post_moe_hop, mode, t[0], r, *t[1:]))
+            elif kind == "deferred_norm":
+                outs.append(deferred_norm(*t, M.GLM_EPS))
+            else:
+                outs.append(post_moe_hop(*t, M.GLM_EPS))
+        torch.cuda.synchronize()
+        failures = []
+        for i, (kind, r, t, out) in enumerate(zip(kinds, rows_seq, inputs, outs)):
+            if kind.startswith("routed_"):
+                h_ref, res_ref = _post_moe_hop_reference(post_moe_hop, mode, t[0], r, *t[1:])
+                gamma = t[3]
+            else:
+                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
+                gamma = t[5]
+            torch.cuda.synchronize()
+            if kind.endswith("post_moe_hop"):
+                ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
+            else:
+                ok = _verdict(out[0], out[1], h_ref, res_ref, gamma, f"b2b{i}-{kind}")
+            if not ok:
+                failures.append((i, kind))
+        assert M.all_ranks(not failures and _acc_clear(post_moe_hop, mode)), failures
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize("mode", MODES)
+def test_routed_cuda_graph_replay(mode):
+    """Fill + deferred_norm routed, fill + post_moe_hop routed and a deferred_norm raw call in one CUDA
+    graph, replayed with fresh routed content."""
+    post_moe_hop, deferred_norm = _hops(mode)
+    rows = 6
+    with torch.inference_mode():
+        static = [_source(mode, rows, seed=1200 + i) for i in range(2)]
+        raw = [x.clone() if x is not None else None for x in M.make_raw_tuple(rows, seed=1250)]
+
+        def clone_src(s):
+            return tuple(x.clone() for x in s) if isinstance(s, tuple) else s.clone()
+
+        bufs = [(clone_src(s[0]),) + tuple(s[1:]) for s in static]
+
+        def body():
+            a = _deferred_norm_call(
+                post_moe_hop, deferred_norm, mode, bufs[0][0], rows, *bufs[0][1:]
+            )
+            b = _post_moe_hop_reference(post_moe_hop, mode, bufs[1][0], rows, *bufs[1][1:])
+            c = deferred_norm(*raw, M.GLM_EPS)
+            return a, b, c
+
+        body()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            ga, gb, gc = body()
+        for trial, replays in enumerate((1, 1, 20)):
+            fresh = [_source(mode, rows, seed=1300 + 10 * trial + i) for i in range(2)]
+            fresh_raw = M.make_raw_tuple(rows, seed=1350 + trial)
+            for dst, src in zip(bufs, fresh):
+                if isinstance(dst[0], tuple):
+                    for d, s in zip(dst[0], src[0]):
+                        d.copy_(s)
+                else:
+                    dst[0].copy_(src[0])
+                for d, s in zip(dst[1:], src[1:]):
+                    if d is not None:
+                        d.copy_(s)
+            for d, s in zip(raw, fresh_raw):
+                if d is not None:
+                    d.copy_(s)
+            for _ in range(replays):
+                M.mpi_env().comm.barrier()
+                graph.replay()
+            torch.cuda.synchronize()
+            ok = True
+            h0, r0 = _post_moe_hop_reference(post_moe_hop, mode, fresh[0][0], rows, *fresh[0][1:])
+            ok &= _verdict(ga[0], ga[1], h0, r0, fresh[0][3], f"graph{trial}-routed_deferred_norm")
+            h1, r1 = _post_moe_hop_reference(post_moe_hop, mode, fresh[1][0], rows, *fresh[1][1:])
+            torch.cuda.synchronize()
+            ok &= torch.equal(gb[0], h1) and torch.equal(gb[1], r1)
+            h2, r2 = post_moe_hop(*fresh_raw, M.GLM_EPS)
+            ok &= _verdict(gc[0], gc[1], h2, r2, fresh_raw[5], f"graph{trial}-deferred_norm")
+            assert M.all_ranks(ok and _acc_clear(post_moe_hop, mode)), f"trial {trial}"
+
+
+def test_routed_delayed_rank_stress():
+    """Random per-rank device delays inside back-to-back routed batches."""
+    mode = "fp32"
+    post_moe_hop, deferred_norm = _hops(mode)
+    env = M.mpi_env()
+    gen = torch.Generator().manual_seed(4243)
+    with torch.inference_mode():
+        for batch in range(3):
+            inputs, outs = [], []
+            for i in range(12):
+                rows = int(torch.randint(1, 9, (1,), generator=gen))
+                t = _source(mode, rows, seed=5000 + 100 * batch + i)
+                local = torch.Generator().manual_seed(77 * batch + 13 * i + env.rank)
+                if float(torch.rand(1, generator=local)) < 0.3:
+                    torch.cuda._sleep(int(float(torch.rand(1, generator=local)) * 2_000_000))
+                inputs.append((rows, t))
+                outs.append(
+                    _deferred_norm_call(post_moe_hop, deferred_norm, mode, t[0], rows, *t[1:])
+                )
+            torch.cuda.synchronize()
+            ok = True
+            for (rows, t), (packet, res) in zip(inputs, outs):
+                h_ref, res_ref = _post_moe_hop_reference(post_moe_hop, mode, t[0], rows, *t[1:])
+                ok &= _verdict(packet, res, h_ref, res_ref, t[3], f"stress{batch}")
+            assert M.all_ranks(ok), f"batch {batch}"
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+# ---------------------------------------------------------------------------
+# post_moe_hop's error protocol on the routed variant
+# ---------------------------------------------------------------------------
+
+_MISSING_TIMEOUT_S = 0.2
+
+
+@pytest.mark.parametrize("missing", [0, 5])
+@pytest.mark.parametrize("mode", ["fp32", "identity"])
+def test_routed_missing_rank_times_out_with_nan_and_reset(mode, missing):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+
+    post_moe_hop, deferred_norm = _hops(mode)
+    imp_post_moe_hop, imp_deferred_norm = _hops(mode, timeout_s=_MISSING_TIMEOUT_S)
+    env = M.mpi_env()
+    assert M.all_ranks(
+        imp_deferred_norm.routed_engaged and imp_post_moe_hop.workspace is post_moe_hop.workspace
+    )
+    rows = 6
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, rows, seed=600 + missing)
+        before = M.lamport_snapshot(post_moe_hop)
+        assert M.all_ranks(before["ok"]), before
+        env.comm.barrier()
+        start = time.monotonic()
+        packet = res = None
+        if env.rank != missing:
+            packet, res = _deferred_norm_call(
+                imp_post_moe_hop, imp_deferred_norm, mode, src, rows, shared, residual, gamma
+            )
+        torch.cuda.synchronize()
+        elapsed = time.monotonic() - start
+        counters = post_moe_hop.workspace.counters()
+        bits, polls = counters["error_bits"], counters["timeout_polls"]
+        if env.rank == missing:
+            reported = bits == 0 and polls == 0
+        else:
+            reported = (
+                bool(bits & 1)
+                and polls > 0
+                and _MISSING_TIMEOUT_S <= elapsed < 30.0
+                and _all_nan_packet(packet, res)
+                and _acc_clear(imp_post_moe_hop, mode)
+            )
+        stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
+        if env.rank == 0:
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_MISSING mode={mode} rank {missing} skipped: (rank, elapsed_s, "
+                f"error_bits, timeout_polls, ok) {stats}",
+                flush=True,
+            )
+        assert M.all_ranks(reported), stats
+        try:
+            check_health(f"routed_deferred_norm-missing-rank{missing}-local")
+            local_raised = False
+        except PostMoeHopError:
+            local_raised = True
+        try:
+            check_health(f"routed_deferred_norm-missing-rank{missing}-collective", collective=True)
+            group_raised = False
+        except PostMoeHopError:
+            group_raised = True
+        assert M.all_ranks(local_raised == (env.rank != missing) and group_raised)
+        _reset(post_moe_hop)
+        src, shared, residual, gamma = _source(mode, rows, seed=650 + missing)
+        packet, res = _deferred_norm_call(
+            post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma
+        )
+        h_ref, res_ref = _post_moe_hop_reference(
+            post_moe_hop, mode, src, rows, shared, residual, gamma
+        )
+        assert _check(packet, res, h_ref, res_ref, gamma, f"after-reset{missing}")
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize(
+    "sequence",
+    [
+        "routed_deferred_norm,routed_deferred_norm,routed_deferred_norm",
+        "routed_deferred_norm,routed_post_moe_hop,deferred_norm",
+        "routed_post_moe_hop,routed_deferred_norm,post_moe_hop",
+        "deferred_norm,routed_deferred_norm,routed_post_moe_hop",
+    ],
+)
+def test_routed_error_poisons_every_rank_until_reset(sequence):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
+        STATE_EPOCH,
+        STATE_ERROR_WORDS,
+    )
+
+    mode = "fp32"
+    post_moe_hop, deferred_norm = _hops(mode)
+    env = M.mpi_env()
+    victim = env.world - 3
+    with torch.inference_mode():
+        torch.cuda.synchronize()
+        if env.rank == victim:
+            epoch = int(post_moe_hop.workspace.state[STATE_EPOCH].item())
+            post_moe_hop.workspace.state[STATE_ERROR_WORDS[(epoch + 2) % 3]] = 1
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        poisoned = True
+        for call, kind in enumerate(sequence.split(",")):
+            if kind.startswith("routed_"):
+                src, shared, residual, gamma = _source(mode, 6, seed=860 + call)
+                if kind == "routed_deferred_norm":
+                    packet, res = _deferred_norm_call(
+                        post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma
+                    )
+                    poisoned &= _all_nan_packet(packet, res)
+                else:
+                    out, res = _post_moe_hop_reference(
+                        post_moe_hop, mode, src, 6, shared, residual, gamma
+                    )
+                    torch.cuda.synchronize()
+                    poisoned &= bool(
+                        torch.isnan(out.float()).all() and torch.isnan(res.float()).all()
+                    )
+            else:
+                t = M.make_raw_tuple(6, seed=860 + call)
+                if kind == "deferred_norm":
+                    packet, res = deferred_norm(*t, M.GLM_EPS)
+                    poisoned &= _all_nan_packet(packet, res)
+                else:
+                    out, res = post_moe_hop(*t, M.GLM_EPS)
+                    torch.cuda.synchronize()
+                    poisoned &= bool(
+                        torch.isnan(out.float()).all() and torch.isnan(res.float()).all()
+                    )
+        words = post_moe_hop.workspace.state.cpu().tolist()
+        carried = all(words[w] == 1 for w in STATE_ERROR_WORDS)
+        clean = not any(words[w] for w in STATE_ERROR_WORDS)
+        per_rank = env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])
+        if env.rank == 0:
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_POISON {sequence}: per-rank error words {per_rank}",
+                flush=True,
+            )
+        assert M.all_ranks(
+            poisoned
+            and (carried if env.rank == victim else clean)
+            and _acc_clear(post_moe_hop, mode)
+        )
+        with pytest.raises(PostMoeHopError):
+            check_health("routed_deferred_norm-poisoned", collective=True)
+        _reset(post_moe_hop)
+        src, shared, residual, gamma = _source(mode, 6, seed=870)
+        packet, res = _deferred_norm_call(
+            post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma
+        )
+        h_ref, res_ref = _post_moe_hop_reference(
+            post_moe_hop, mode, src, 6, shared, residual, gamma
+        )
+        assert _check(packet, res, h_ref, res_ref, gamma, "after-poison-reset")
+
+
+_FAULT_DELAYS_NS = (3000, 60000)
+
+
+@pytest.mark.parametrize("mode", ["fp32", "identity"])
+def test_routed_fault_injected_store_order_mixed_with_post_moe_hop(mode):
+    """Fault-injected deferred_norm / post_moe_hop routed builds chosen per rank and per call for
+    the same collective call, plus random device delays: every output bit as
+    in an undisturbed post_moe_hop routed call, one call count per call, no error."""
+    post_moe_hop, deferred_norm = _hops(mode)
+    hops = {
+        "post_moe_hop": (post_moe_hop, post_moe_hop),
+        "deferred_norm": (post_moe_hop, deferred_norm),
+    }
+    for d in _FAULT_DELAYS_NS:
+        faulty_post_moe, faulty_deferred_norm = _hops(mode, fault_delay_ns=d)
+        hops[f"post_moe_hop_fault{d}"] = (faulty_post_moe, faulty_post_moe)
+        hops[f"deferred_norm_fault{d}"] = (faulty_post_moe, faulty_deferred_norm)
+    env = M.mpi_env()
+    names = sorted(hops)
+    common = random.Random(4712)
+    local = random.Random(2000 + env.rank)
+    calls_before = post_moe_hop.workspace.calls()
+    failed, used, total = [], {}, 0
+    with torch.inference_mode():
+        for batch in range(5):
+            plan = [(common.randint(1, 8), common.random() < 0.75) for _ in range(12)]
+            inputs = [
+                _source(mode, r, seed=9100 + 100 * batch + j, with_shared=sh)
+                for j, (r, sh) in enumerate(plan)
+            ]
+            picks = [local.choice(names) for _ in plan]
+            torch.cuda.synchronize()
+            env.comm.barrier()
+            outs = []
+            for name, (r, _), t in zip(picks, plan, inputs):
+                if local.random() < 0.3:
+                    torch.cuda._sleep(int(local.uniform(0.0, 50.0) * 2000))
+                used[name] = used.get(name, 0) + 1
+                owner, hop = hops[name]
+                routed = _fill(owner, mode, t[0], r)
+                outs.append(hop.hop_routed(routed, t[1], t[2], t[3], M.GLM_EPS))
+            torch.cuda.synchronize()
+            for j, (name, (r, _), t, out) in enumerate(zip(picks, plan, inputs, outs)):
+                h_ref, res_ref = _post_moe_hop_reference(post_moe_hop, mode, t[0], r, *t[1:])
+                if name.startswith("deferred_norm"):
+                    ok = _verdict(out[0], out[1], h_ref, res_ref, t[3], f"fault{batch}.{j}")
+                else:
+                    torch.cuda.synchronize()
+                    ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
+                if not ok:
+                    failed.append((batch, j, name))
+            total += len(plan)
+        counted = post_moe_hop.workspace.calls() - calls_before
+    bits = post_moe_hop.error_bits()
+    if env.rank == 0:
+        print(
+            f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_FAULT mode={mode} calls={total} (+{total} reference calls) "
+            f"used_on_rank0={used}",
+            flush=True,
+        )
+    assert M.all_ranks(not failed and counted == 2 * total and bits == 0), (
+        failed[:4],
+        counted,
+        bits,
+    )
+
+
+@pytest.mark.parametrize(
+    "case", ("foreign_buffer", "row_count", "residual_dtype", "shared_shape", "gamma_shape")
+)
+def test_routed_contract_violation_raises_before_any_device_work(case):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopContractError,
+    )
+
+    mode = "fp32"
+    post_moe_hop, deferred_norm = _hops(mode)
+    rows = 6
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, rows, seed=920)
+        routed = _fill(post_moe_hop, mode, src, rows)
+        if case == "foreign_buffer":
+            routed = routed.clone()
+        elif case == "row_count":
+            routed = post_moe_hop.routed_acc_buffer(rows + 1)
+        elif case == "residual_dtype":
+            residual = residual.half()
+        elif case == "shared_shape":
+            shared = shared[:, : M.GLM_HIDDEN // 2]
+        else:
+            gamma = gamma[:-8]
+        calls_before = post_moe_hop.workspace.calls()
+        with pytest.raises(PostMoeHopContractError):
+            deferred_norm.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
+        assert M.all_ranks(post_moe_hop.workspace.calls() == calls_before)
+        post_moe_hop.routed_acc_buffer(post_moe_hop.max_rows).zero_()
+        src, shared, residual, gamma = _source(mode, rows, seed=921)
+        packet, res = _deferred_norm_call(
+            post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma
+        )
+        h_ref, res_ref = _post_moe_hop_reference(
+            post_moe_hop, mode, src, rows, shared, residual, gamma
+        )
+        assert _check(packet, res, h_ref, res_ref, gamma, f"after-contract-{case}")
+
+
+def _misaligned(t, offset_elems):
+    storage = torch.empty(t.numel() + 16, dtype=t.dtype, device=t.device)
+    view = storage[offset_elems : offset_elems + t.numel()].view(t.shape)
+    view.copy_(t)
+    return view
+
+
+def _transposed(t):
+    view = torch.empty(tuple(reversed(t.shape)), dtype=t.dtype, device=t.device).t()
+    view.copy_(t)
+    return view
+
+
+def _needs_copy(t):
+    """The hop's 16-byte vector rule for a 1-D / row-major operand."""
+    return bool(t.data_ptr() % 16 or t.stride(-1) != 1 or not t.is_contiguous())
+
+
+def _relayout(rank, shared, residual, gamma):
+    """Rank-dependent layouts; returns them and the repairs expected on this
+    rank (a one-row "transposed" tensor is still contiguous)."""
+    kind = rank % 4
+    if kind == 0:
+        gamma = _misaligned(gamma, 2)
+    elif kind == 1:
+        residual = _misaligned(residual, 4)
+    elif kind == 2 and shared is not None:
+        shared = _transposed(shared)
+    repairs = sum(int(_needs_copy(t)) for t in (shared, residual, gamma) if t is not None)
+    return shared, residual, gamma, repairs
+
+
+def test_routed_unaligned_layouts_are_repaired_locally():
+    mode = "fp32"
+    post_moe_hop, deferred_norm = _hops(mode)
+    env = M.mpi_env()
+    ok = True
+    with torch.inference_mode():
+        for i, rows in enumerate((6, 1, 8)):
+            src, shared, residual, gamma = _source(mode, rows, seed=880 + i)
+            r_shared, r_residual, r_gamma, expected = _relayout(env.rank, shared, residual, gamma)
+            before = post_moe_hop.workspace.layout_repairs
+            packet, res = _deferred_norm_call(
+                post_moe_hop, deferred_norm, mode, src, rows, r_shared, r_residual, r_gamma
+            )
+            repaired = post_moe_hop.workspace.layout_repairs - before
+            h_ref, res_ref = _post_moe_hop_reference(
+                post_moe_hop, mode, src, rows, shared, residual, gamma
+            )
+            ok &= _verdict(packet, res, h_ref, res_ref, gamma, f"relayout{rows}")
+            ok &= repaired == expected
+        rows = 6
+        src, shared, residual, gamma = _source(mode, rows, seed=890)
+        s_src = src.clone()
+        s_shared, s_residual, s_gamma, _ = _relayout(env.rank, shared, residual, gamma)
+        _deferred_norm_call(
+            post_moe_hop, deferred_norm, mode, s_src, rows, s_shared, s_residual, s_gamma
+        )
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            g_packet, g_res = _deferred_norm_call(
+                post_moe_hop, deferred_norm, mode, s_src, rows, s_shared, s_residual, s_gamma
+            )
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        for replay in range(4):
+            f_src, f_shared, f_residual, _ = _source(mode, rows, seed=900 + replay)
+            s_src.copy_(f_src)
+            s_shared.copy_(f_shared)
+            s_residual.copy_(f_residual)
+            graph.replay()
+            torch.cuda.synchronize()
+            h_ref, res_ref = _post_moe_hop_reference(
+                post_moe_hop, mode, f_src, rows, f_shared, f_residual, gamma
+            )
+            ok &= _verdict(g_packet, g_res, h_ref, res_ref, gamma, f"relayout-graph{replay}")
+    assert M.all_ranks(ok and post_moe_hop.error_bits() == 0)
+
+
+def test_health_check_counts_deferred_norm_routed_calls():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        check_health,
+    )
+
+    mode = "fp32"
+    post_moe_hop, deferred_norm = _hops(mode)
+    with torch.inference_mode():
+        before = check_health("routed_deferred_norm-engagement-before", collective=True)
+        for i in range(5):
+            src, shared, residual, gamma = _source(mode, 6, seed=940 + i)
+            _deferred_norm_call(post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma)
+        after = check_health("routed_deferred_norm-engagement-after", collective=True)
+    assert M.all_ranks(
+        after.calls - before.calls == 5
+        and after.error_bits == 0
+        and not after.group_error
+        and after.calls == post_moe_hop.workspace.calls()
+    )
+
+
+def test_routed_memcheck_positive_control():
+    """Tests only, run under compute-sanitizer memcheck by the sanitizer
+    suite (``GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL=1``): the routed variant reads (and, in
+    fp32 mode, clears) an exactly sized accumulator with twice its row
+    pitch, so row 1 is out of bounds and memcheck must report it."""
+    if os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL") != "1":
+        pytest.skip("memcheck positive control (sanitizer suite only)")
+    mode = "fp32"
+    post_moe_hop, deferred_norm = _hops(mode)
+    rows = 2
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, rows, seed=990)
+        acc = src.clone()
+        deferred_norm.launch_routed(
+            acc,
+            shared,
+            residual,
+            gamma,
+            M.GLM_EPS,
+            row_bytes=2 * acc.stride(0) * acc.element_size(),
+        )
+        torch.cuda.synchronize()
+    _reset(post_moe_hop)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py
new file mode 100644
index 000000000000..a1b2b6506bfd
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py
@@ -0,0 +1,327 @@
+# 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.
+"""Single-process tests of the deferred_norm scope (``deferred_norm/scope.py``).
+
+The scope logic (this target's ``shared`` scope, the next layer's indexer type, the constructing
+decoder layer found from the calling frames) needs only the standard library and the target's
+``dispatch`` and runs without a built ``tensorrt_llm``. The remaining tests need the installed
+package and the GLM-5.2 checkpoint:
+
+* the GLM-5.2 indexer layout seen by the scope (the model config's
+  ``to_sparse_params``) has 18 full-indexer layers among the 74 next layers of
+  the deferred_norm boundaries (6, 10, ..., 74) and agrees with the native attention
+  module (an ``Indexer`` exactly on full layers);
+* the consumer guard: under ``shared``, this target's scope, the deferred_norm projection
+  declines a full-indexer layer (so a packet there is materialized for the
+  native path) and accepts shared layers.
+"""
+
+import importlib
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+REPO = Path(__file__).resolve().parents[9]
+TARGET = (
+    REPO
+    / "tensorrt_llm"
+    / "_torch"
+    / "_experimental"
+    / "modeling_v2"
+    / "models"
+    / "glm_moe_dsa"
+    / "shared"
+)
+_STANDALONE = "_deferred_norm_scope_standalone"
+# GLM-5.2: 78 target layers, full-indexer layers where max(i - 2, 0) % 4 == 0.
+GLM52 = dict(
+    num_hidden_layers=78, index_topk_freq=4, index_skip_topk_offset=3, index_topk_pattern=None
+)
+FULL_NEXT = list(range(6, 78, 4))
+BOUNDARIES = list(range(3, 77))  # L -> L+1 with a deferred_norm hop (first_k_dense_replace = 3)
+
+
+def _load_dispatch_standalone(name):
+    """The target's ``dispatch.py`` as module ``name``, without importing ``tensorrt_llm``."""
+    spec = importlib.util.spec_from_file_location(
+        "deferred_norm_scope_input_prep_testlib",
+        Path(__file__).resolve().parents[1] / "deferred_input_prep" / "input_prep_testlib.py",
+    )
+    testlib = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(testlib)
+    return testlib.load_dispatch_standalone(name)
+
+
+def _load_scope():
+    """``scope.py``: the installed package, else the source tree under a stand-in package."""
+    try:
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+            scope,
+        )
+
+        return scope
+    except (ImportError, OSError):
+        pass
+    if _STANDALONE not in sys.modules:
+        for name, path in (
+            (_STANDALONE, TARGET),
+            (f"{_STANDALONE}.kernels", TARGET / "kernels"),
+            (f"{_STANDALONE}.kernels.deferred_norm", TARGET / "kernels" / "deferred_norm"),
+        ):
+            package = types.ModuleType(name)
+            package.__path__ = [str(path)]
+            sys.modules[name] = package
+        sys.modules[f"{_STANDALONE}.dispatch"] = _load_dispatch_standalone(
+            f"{_STANDALONE}.dispatch"
+        )
+    return importlib.import_module(f"{_STANDALONE}.kernels.deferred_norm.scope")
+
+
+@pytest.fixture
+def scope():
+    return _load_scope()
+
+
+class _Sparse:
+    """``to_sparse_params`` with GLM-5.2's frequency rule; records its calls."""
+
+    def __init__(self):
+        self.calls = []
+
+    def to_sparse_params(self, pretrained_config=None, layer_idx=None):
+        self.calls.append((pretrained_config, layer_idx))
+        offset = pretrained_config.index_skip_topk_offset
+        full = max(layer_idx - offset + 1, 0) % pretrained_config.index_topk_freq == 0
+        return types.SimpleNamespace(is_full_indexer_layer=full)
+
+
+def _model_config(sparse=None):
+    return types.SimpleNamespace(
+        sparse_attention_config=sparse or _Sparse(),
+        pretrained_config=types.SimpleNamespace(**GLM52),
+    )
+
+
+class _Hop:
+    """Stands in for ``DeferredNormHop.__init__``: resolves the next layer like it."""
+
+    def __init__(self, scope, post_moe_hop, arg=None):
+        self.kind = scope.next_layer_kind(post_moe_hop, arg)
+
+
+class _DecoderLayer:
+    """The order of ``GlmDecoderLayer._adopt_post_moe``: post_moe_hop first, then deferred_norm."""
+
+    def __init__(self, scope, layer_idx, model_config, post_moe_hop):
+        self.model_config = model_config
+        self.layer_idx = layer_idx
+        self.post_moe_hop = post_moe_hop
+        self.deferred_norm = _Hop(scope, self.post_moe_hop)
+
+
+def test_the_scope_of_this_target(scope):
+    assert scope.SCOPES == ("all", "shared") and scope.DEFAULT_SCOPE == "shared"
+    assert scope.UNRESTRICTED_SCOPE == "all"
+    assert scope.SCOPE_LABEL == "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE"
+    assert scope.selected_scope() == "shared"
+
+
+def test_the_record_names_the_scope_of_this_target(scope):
+    """The record's deferred_norm variant is the scope the hop runs, whatever the environment
+    holds (``resolve`` refuses the variable)."""
+    try:
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    except (ImportError, OSError):
+        dispatch = _load_dispatch_standalone("dispatch_standalone")
+    assert dispatch.NUMERICAL_VARIANTS["deferred_norm"] == (None, (scope.DEFAULT_SCOPE,))
+    for raw in (None, "", " ", "SHARED", " shared ", "all", "bogus"):
+        env = {} if raw is None else {scope.SCOPE_LABEL: raw}
+        assert dispatch.feature_variants(["deferred_norm"], env) == {
+            "deferred_norm": scope.DEFAULT_SCOPE
+        }
+
+
+def test_glm52_frequency_rule_boundaries(scope):
+    cfg = _model_config()
+    full = {i for i in range(78) if scope.is_full_indexer_layer(cfg, i)}
+    assert full == {0, 1, 2} | set(FULL_NEXT)
+    into_full = [b for b in BOUNDARIES if scope.is_full_indexer_layer(cfg, b + 1)]
+    assert into_full == [n - 1 for n in FULL_NEXT] and len(into_full) == 18
+    assert len(BOUNDARIES) - len(into_full) == 56
+
+
+def test_unknown_layout(scope):
+    assert (
+        scope.is_full_indexer_layer(types.SimpleNamespace(sparse_attention_config=None), 6) is None
+    )
+    params = types.SimpleNamespace(to_sparse_params=lambda **kw: types.SimpleNamespace())
+    assert (
+        scope.is_full_indexer_layer(
+            types.SimpleNamespace(sparse_attention_config=params, pretrained_config=None), 6
+        )
+        is None
+    )
+
+
+@pytest.mark.parametrize("layer_idx", [3, 5, 6, 9, 73, 75, 76])
+def test_next_layer_from_the_constructing_layer(scope, layer_idx):
+    hop_stub = object()
+    cfg = _model_config()
+    layer = _DecoderLayer(scope, layer_idx, cfg, hop_stub)
+    full, nxt, how = layer.deferred_norm.kind
+    assert nxt == layer_idx + 1 and how == "model config"
+    assert full is ((layer_idx + 1) in FULL_NEXT)
+    # The same call as MLA.__init__ of the next layer.
+    assert cfg.sparse_attention_config.calls == [(cfg.pretrained_config, layer_idx + 1)]
+
+
+def test_constructing_layer_must_hold_this_post_moe_hop(scope):
+    hop_stub, other = object(), object()
+    unresolved = (None, None, "no constructing decoder layer found")
+    assert _Hop(scope, hop_stub).kind == unresolved
+
+    class _OtherHop:
+        def __init__(self):
+            self.model_config = _model_config()
+            self.layer_idx = 5
+            self.post_moe_hop = other
+            self.deferred_norm = _Hop(scope, hop_stub)
+
+    assert _OtherHop().deferred_norm.kind == unresolved
+
+    class _NoIndex:
+        def __init__(self):
+            self.post_moe_hop = hop_stub
+            self.layer_idx = None
+            self.model_config = _model_config()
+            self.deferred_norm = _Hop(scope, hop_stub)
+
+    assert _NoIndex().deferred_norm.kind == unresolved
+
+
+def test_explicit_argument_wins(scope):
+    hop_stub = object()
+
+    class _Layer:
+        def __init__(self, arg):
+            self.model_config = _model_config()
+            self.layer_idx = 5  # next layer 6 is full
+            self.post_moe_hop = hop_stub
+            self.deferred_norm = _Hop(scope, hop_stub, arg)
+
+    assert _Layer(False).deferred_norm.kind == (False, None, "argument")
+    assert _Layer(True).deferred_norm.kind == (True, None, "argument")
+    assert _Layer(None).deferred_norm.kind == (True, 6, "model config")
+
+
+def test_unresolved_config(scope):
+    hop_stub = object()
+    cfg = types.SimpleNamespace(sparse_attention_config=None, pretrained_config=None)
+    layer = _DecoderLayer(scope, 5, cfg, hop_stub)
+    assert layer.deferred_norm.kind[0] is None and layer.deferred_norm.kind[1] == 6
+
+
+def test_decline_reason(scope):
+    assert scope.decline_reason("all", True, 6, "model config") is None
+    assert scope.decline_reason("all", None, None, "x") is None
+    assert scope.decline_reason("shared", False, 7, "model config") is None
+    why = scope.decline_reason("shared", True, 6, "model config")
+    assert "layer 6 is a full-indexer layer" in why and "post_moe_hop hop" in why
+    why = scope.decline_reason("shared", None, None, "no constructing decoder layer found")
+    assert "unknown" in why and "no constructing decoder layer found" in why
+
+
+# ---- installed package and checkpoint ------------------------------------------------
+
+
+def _installed():
+    try:
+        import torch  # noqa: F401
+
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501, F401
+            projection,
+        )
+    except (ImportError, OSError) as exc:
+        pytest.skip(f"needs the installed tensorrt_llm ({exc!r})")
+    import absorbed_query_fixture
+
+    if not absorbed_query_fixture.checkpoint_available():
+        pytest.skip("needs the GLM-5.2 checkpoint")
+    return absorbed_query_fixture
+
+
+_ATTN = {}
+
+
+def _attention(absorbed_query_fixture, layer):
+    import torch
+
+    if not torch.cuda.is_available():
+        pytest.skip("needs a GPU")
+    if "cfg" not in _ATTN:
+        _ATTN["cfg"] = absorbed_query_fixture.build_model_config(0)
+    if layer not in _ATTN:
+        with torch.inference_mode():
+            _ATTN[layer] = absorbed_query_fixture.build_native_attention(layer, 0, _ATTN["cfg"])
+    return _ATTN["cfg"], _ATTN[layer]
+
+
+def test_glm52_model_config_layout(scope):
+    """The scope's view of the real GLM-5.2 model config (no GPU work)."""
+    absorbed_query_fixture = _installed()
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        scope as installed,
+    )
+
+    cfg = absorbed_query_fixture.build_model_config(0)
+    full = [b + 1 for b in BOUNDARIES if installed.is_full_indexer_layer(cfg, b + 1)]
+    assert full == FULL_NEXT, full
+    assert len(BOUNDARIES) - len(full) == 56
+
+
+@pytest.mark.parametrize("layer", [6, 7])
+def test_scope_layout_matches_the_attention_module(layer):
+    """``is_full_indexer_layer`` is true exactly where the native module has an indexer."""
+    absorbed_query_fixture = _installed()
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        scope as installed,
+    )
+
+    cfg, attn = _attention(absorbed_query_fixture, layer)
+    has_indexer = attn.mqa.indexer is not None
+    assert installed.is_full_indexer_layer(cfg, layer) is has_indexer
+    assert has_indexer is (layer in FULL_NEXT)
+
+
+@pytest.mark.parametrize("layer", [6, 7])
+def test_consumer_guard(layer):
+    """Under ``shared`` the deferred_norm projection declines full-indexer layers only."""
+    absorbed_query_fixture = _installed()
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.scope import (  # noqa: E501
+        SCOPE_LABEL,
+    )
+
+    _, attn = _attention(absorbed_query_fixture, layer)
+    ok, reason, kw = projection_eligibility(attn)
+    if layer in FULL_NEXT:
+        assert not ok and kw is None and f"{SCOPE_LABEL}=shared" in reason, reason
+    else:
+        assert ok and kw is None, reason
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py
new file mode 100644
index 000000000000..e265c090fda3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py
@@ -0,0 +1,342 @@
+# 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.
+"""Collective tests of the deferred_norm scope (8 ranks over 2 hosts, TP8).
+
+The ``shared`` scope, this target's, keeps the deferred-norm hop only on boundaries whose next
+layer is a shared-indexer layer. The tests that compare against the hop's unrestricted scope
+``all`` build the hop with that scope directly. Checks:
+
+* unrestricted scope: the hop is the default deferred_norm hop (engaged, no next-layer
+  lookup, no per-boundary log key);
+* ``shared``, next layer full: the hop is disengaged on every rank and the
+  layer's choice (``can_run`` / ``can_run_routed`` false; the decoder layer
+  keeps ``None``) is post_moe_hop's hop, whose output is post_moe_hop's by construction; the
+  routed input is never enabled on such a hop;
+* ``shared``, next layer shared: the hop is engaged and its packets (raw
+  tuple and routed_moe_into_hop routed input in ``fp32`` and ``identity``) are bit-identical
+  to the unrestricted hop's packets on the same inputs and workspace, and
+  their materialization equals post_moe_hop's normalized output (``a_frag``, ``sumsq``,
+  residual, ``materialize``);
+* rank-differing scopes: one rank with another scope or an unknown
+  next-layer type can only disengage a boundary on every rank (one agreement,
+  the same number of collectives in every scope);
+* the decoder-layer path: a stand-in with the target decoder layer's adoption order
+  (``GlmDecoderLayer._adopt_post_moe``: post_moe_hop assigned, then the deferred_norm factory
+  called with only the post_moe_hop) and the real GLM-5.2 model config resolves layers 5, 9, 73 ->
+  full (post_moe_hop) and 3, 6, 76 -> shared (deferred_norm hop) on every rank.
+"""
+
+import importlib
+import os
+import sys
+import types
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
+sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
+
+import post_moe_hop_mpi_utils as M  # noqa: E402
+
+FULL_NEXT = set(range(6, 78, 4))
+
+
+def _env():
+    env = M.mpi_env()
+    if env.world != 8:
+        pytest.skip("the deferred_norm hop is specialized to TP8")
+    return env
+
+
+_POST_MOE_HOP = {}
+
+
+def _post_moe_hop(mode=None):
+    """One post_moe_hop per routed-input mode (``None``: raw tuple only), collective."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    if mode not in _POST_MOE_HOP:
+        hop_stub = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
+        ok = hop_stub.engaged
+        if mode is not None:
+            if not hasattr(hop_stub, "enable_routed_input"):
+                pytest.skip(
+                    "the installed post_moe_hop hop has no routed_moe_into_hop routed input"
+                )
+            ok = ok and hop_stub.enable_routed_input(mode)
+        assert M.all_ranks(ok), hop_stub.reason
+        _POST_MOE_HOP[mode] = hop_stub
+    return _POST_MOE_HOP[mode]
+
+
+UNRESTRICTED = "all"
+_HOP_MODULE = (
+    "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa."
+    "shared.kernels.deferred_norm.hop"
+)
+
+
+def _hop(monkeypatch, hop_stub, scope_value, **kwargs):
+    """A deferred_norm hop built under this target's scope (``scope_value`` None or ``shared``).
+
+    ``UNRESTRICTED``: the hop's unrestricted scope, which this target does not run; the hop
+    module's scope reader returns it for this construction.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+
+    with monkeypatch.context() as patch:
+        if scope_value == UNRESTRICTED:
+            hop_module = importlib.import_module(_HOP_MODULE)
+            patch.setattr(hop_module, "selected_scope", lambda: UNRESTRICTED)
+        return DeferredNormHop(hop_stub, **kwargs)
+
+
+def _metadata(rows):
+    return types.SimpleNamespace(all_rank_num_tokens=None, num_contexts=0, num_generations=rows)
+
+
+def _same_packet(pa, ra, pb, rb):
+    """Same packet content (the fragment buffer holds 8 rows; rows >= R are not written)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    torch.cuda.synchronize()
+    rows, k = ra.shape
+    return (
+        torch.equal(rows_from_frag(pa.a_frag, rows, k), rows_from_frag(pb.a_frag, rows, k))
+        and torch.equal(pa.sumsq, pb.sumsq)
+        and torch.equal(ra, rb)
+        and pa.eps == pb.eps
+    )
+
+
+def _materialized_is_post_moe_hop(packet, h_post_moe_hop):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        materialize,
+    )
+
+    h = materialize(packet)
+    torch.cuda.synchronize()
+    return torch.equal(h, h_post_moe_hop)
+
+
+def test_default_scope_is_the_b11_hop(monkeypatch):
+    hop_stub = _post_moe_hop()
+    # The next-layer argument is ignored by the unrestricted scope.
+    hop = _hop(monkeypatch, hop_stub, UNRESTRICTED, next_layer_full_indexer=True)
+    ok = (
+        hop.engaged
+        and hop.scope == "all"
+        and hop.next_layer_full is None
+        and hop.next_layer_idx is None
+        and hop.reason == ""
+    )
+    assert M.all_ranks(ok), (hop.scope, hop.reason)
+
+
+def test_shared_scope_full_next_layer_keeps_post_moe_hop(monkeypatch):
+    hop_stub = _post_moe_hop()
+    hop = _hop(monkeypatch, hop_stub, None, next_layer_full_indexer=True)
+    ok = (
+        not hop.engaged
+        and hop.scope == "shared"
+        and hop.next_layer_full is True
+        and "full-indexer layer" in hop.reason
+        and not hop.can_run(6, _metadata(6))
+        and not hop.can_run_routed(6, _metadata(6))
+    )
+    assert M.all_ranks(ok), hop.reason
+    # The decoder layer enables the routed input only on a kept (engaged) hop;
+    # called anyway, the enable declines on every rank without a collective.
+    assert M.all_ranks(not hop.enable_routed_input() and not hop.routed_engaged)
+
+
+@pytest.mark.parametrize("with_shared", [True, False])
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+def test_shared_scope_shared_next_layer_equals_default(monkeypatch, rows, with_shared):
+    hop_stub = _post_moe_hop()
+    default = _hop(monkeypatch, hop_stub, UNRESTRICTED)
+    shared = _hop(monkeypatch, hop_stub, None, next_layer_full_indexer=False)
+    assert M.all_ranks(
+        default.engaged
+        and shared.engaged
+        and shared.next_layer_full is False
+        and shared.can_run(rows, _metadata(rows))
+    ), shared.reason
+    with torch.inference_mode():
+        fc2, scales, perm, sh, residual, gamma = M.make_raw_tuple(
+            rows, seed=900 + rows, with_shared=with_shared
+        )
+        h_post_moe_hop, res_post_moe_hop = hop_stub(
+            fc2, scales, perm, sh, residual, gamma, M.GLM_EPS
+        )
+        pa, ra = default(fc2, scales, perm, sh, residual, gamma, M.GLM_EPS)
+        pb, rb = shared(fc2, scales, perm, sh, residual, gamma, M.GLM_EPS)
+        ok = (
+            _same_packet(pa, ra, pb, rb)
+            and torch.equal(rb, res_post_moe_hop)
+            and _materialized_is_post_moe_hop(pb, h_post_moe_hop)
+        )
+    assert M.all_ranks(ok)
+
+
+def _routed_source(mode, rows, seed):
+    fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(rows, seed=seed)
+    if mode == "identity":
+        valid = perm >= 0
+        safe = torch.where(valid, perm, torch.zeros_like(perm)).long()
+        src = (
+            fc2[safe.view(-1)].contiguous(),
+            torch.where(valid, scales, torch.zeros_like(scales)).bfloat16().contiguous(),
+        )
+    else:
+        valid = (perm >= 0).unsqueeze(-1)
+        safe = torch.where(perm >= 0, perm, torch.zeros_like(perm)).long()
+        prod = fc2[safe.view(-1)].view(rows, M.GLM_TOP_K, -1).float() * scales.float().unsqueeze(-1)
+        src = torch.where(valid, prod, torch.zeros_like(prod)).sum(1).contiguous()
+    return src, shared, residual, gamma
+
+
+def _fill(hop_stub, mode, src, rows):
+    if mode == "identity":
+        buf, scales = hop_stub.routed_raw_buffers(rows)
+        buf.copy_(src[0])
+        scales.copy_(src[1])
+        return buf
+    buf = hop_stub.routed_acc_buffer(rows)
+    buf.copy_(src)
+    return buf
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("mode", ["fp32", "identity"])
+def test_shared_scope_routed_equals_default(monkeypatch, mode, rows):
+    hop_stub = _post_moe_hop(mode)
+    default = _hop(monkeypatch, hop_stub, UNRESTRICTED)
+    shared = _hop(monkeypatch, hop_stub, None, next_layer_full_indexer=False)
+    ok = default.enable_routed_input() and shared.enable_routed_input()
+    assert M.all_ranks(ok and shared.can_run_routed(rows, _metadata(rows))), shared.routed_reason
+    with torch.inference_mode():
+        src, sh, residual, gamma = _routed_source(mode, rows, seed=950 + rows)
+        h_post_moe_hop, res_post_moe_hop = hop_stub.hop_routed(
+            _fill(hop_stub, mode, src, rows), sh, residual, gamma, M.GLM_EPS
+        )
+        pa, ra = default.hop_routed(
+            _fill(hop_stub, mode, src, rows), sh, residual, gamma, M.GLM_EPS
+        )
+        pb, rb = shared.hop_routed(_fill(hop_stub, mode, src, rows), sh, residual, gamma, M.GLM_EPS)
+        ok = (
+            _same_packet(pa, ra, pb, rb)
+            and torch.equal(rb, res_post_moe_hop)
+            and _materialized_is_post_moe_hop(pb, h_post_moe_hop)
+        )
+        if mode == "fp32":
+            torch.cuda.synchronize()
+            ok = ok and bool((hop_stub.routed_acc_buffer(hop_stub.max_rows) == 0).all())
+    assert M.all_ranks(ok)
+
+
+@pytest.mark.parametrize("next_full", [True, False])
+@pytest.mark.parametrize("odd_rank", [0, 3, 7])
+def test_rank_differing_scope_agrees(monkeypatch, odd_rank, next_full):
+    """One rank with the unrestricted scope: a full boundary is post_moe_hop everywhere, a
+    shared boundary is deferred_norm everywhere."""
+    env = _env()
+    hop_stub = _post_moe_hop()
+    value = UNRESTRICTED if env.rank == odd_rank else "shared"
+    hop = _hop(monkeypatch, hop_stub, value, next_layer_full_indexer=next_full)
+    if next_full:
+        local = (
+            "another rank is not capable" == hop.reason
+            if env.rank == odd_rank
+            else "full-indexer layer" in hop.reason
+        )
+        assert M.all_ranks(not hop.engaged and local), hop.reason
+    else:
+        assert M.all_ranks(hop.engaged), hop.reason
+
+
+@pytest.mark.parametrize("odd_rank", [2, 5])
+def test_unresolved_next_layer_on_one_rank(monkeypatch, odd_rank):
+    """A rank that cannot resolve the next layer keeps post_moe_hop on every rank."""
+    env = _env()
+    hop_stub = _post_moe_hop()
+    arg = None if env.rank == odd_rank else False
+    hop = _hop(monkeypatch, hop_stub, "shared", next_layer_full_indexer=arg)
+    if env.rank == odd_rank:
+        local = "unknown" in hop.reason and "no constructing decoder layer found" in hop.reason
+    else:
+        local = hop.reason == "another rank is not capable"
+    assert M.all_ranks(not hop.engaged and local), hop.reason
+
+
+def _glm52_model_config():
+    import absorbed_query_fixture
+
+    if not absorbed_query_fixture.checkpoint_available():
+        pytest.skip("needs the GLM-5.2 checkpoint config")
+    return absorbed_query_fixture.build_model_config(M.mpi_env().rank)
+
+
+class _DecoderLayerStandIn:
+    """``GlmDecoderLayer._adopt_post_moe``'s deferred_norm wiring, in its order."""
+
+    def __init__(self, model_config, layer_idx, post_moe_hop):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+        self.model_config = model_config
+        self.layer_idx = layer_idx
+        self.post_moe_hop = post_moe_hop
+        hop = dispatch.get_impl("deferred_norm", "deferred_hop")(self.post_moe_hop)
+        self.deferred_norm_hop = hop if hop.engaged else None
+        self.deferred_norm_reason = hop.reason
+
+
+def test_decoder_layer_resolves_the_next_layer():
+    """Under ``shared``, this target's scope, full next layers keep the post-MoE hop."""
+    hop_stub = _post_moe_hop()
+    cfg = _glm52_model_config()
+    ok = True
+    for layer_idx in (3, 5, 6, 9, 73, 76):
+        layer = _DecoderLayerStandIn(cfg, layer_idx, hop_stub)
+        hop = layer.deferred_norm_hop
+        if (layer_idx + 1) not in FULL_NEXT:
+            good = (
+                hop is not None
+                and hop.next_layer_idx == layer_idx + 1
+                and hop.next_layer_full is False
+            )
+        else:
+            good = (
+                hop is None
+                and f"layer {layer_idx + 1} is a full-indexer layer" in layer.deferred_norm_reason
+            )
+        if not good:
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_SCOPE_LAYER_FAIL rank={M.mpi_env().rank} layer={layer_idx} "
+                f"reason={layer.deferred_norm_reason}",
+                flush=True,
+            )
+        ok = ok and good
+    assert M.all_ranks(ok)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py
new file mode 100644
index 000000000000..f0b0b59ecdf9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py
@@ -0,0 +1,126 @@
+# 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.
+"""Control flow of the deferred_norm_evict_hints call sites in the target.
+
+Without device work, checks that
+
+* the plan runs deferred_norm_evict_hints with its closure ``(absorbed_query, post_moe_hop,
+  deferred_norm)`` and turns it off with any of them;
+* the DSA projection entry (``attention_forward._absorbed_query_packet``) hands a deferred_norm
+  packet to deferred_norm_evict_hints's ``dsa_proj``
+  exactly when deferred_norm_evict_hints is selected (deferred_norm's otherwise), and leaves other
+  hidden states to the next entry;
+* the layer adoption (``_adopt_post_moe``) asks deferred_norm_evict_hints's ``deferred_hop`` factory when
+  deferred_norm_evict_hints is selected (deferred_norm's otherwise), with the decision taken on the
+  host before the factory runs;
+* deferred_input_prep knows deferred_norm_evict_hints's package directories and side selectors.
+"""
+
+import inspect
+import os
+import re
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import plan_testlib  # noqa: E402
+
+BASE = "absorbed_query,post_moe_hop,routed_moe_into_hop,deferred_norm,routed_moe"
+
+
+@pytest.fixture
+def select(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+
+    def _select(features):
+        plan_testlib.select(dispatch, features, monkeypatch)
+        return dispatch
+
+    return _select
+
+
+def test_selector_closure(select):
+    dispatch = select(BASE + ",deferred_norm_evict_hints")
+    assert dispatch.is_enabled("deferred_norm_evict_hints") and dispatch.is_enabled("deferred_norm")
+    assert dispatch.REQUIRED_FEATURES["deferred_norm_evict_hints"] == frozenset(
+        {"absorbed_query", "post_moe_hop", "deferred_norm"}
+    )
+    for base in ("absorbed_query", "post_moe_hop", "deferred_norm"):
+        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
+        assert not plan.enabled("deferred_norm_evict_hints"), base
+
+
+@pytest.mark.parametrize("with_evict_hints", [True, False])
+def test_dsa_entry_routes_deferred_norm_packets(select, monkeypatch, with_evict_hints):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+        attention_forward,
+    )
+
+    dispatch = select(BASE + (",deferred_norm_evict_hints" if with_evict_hints else ""))
+    calls = []
+    packet = object()
+    real_get_impl = dispatch.get_impl
+
+    def fake_get_impl(feature, name):
+        if name == "is_deferred":
+            return lambda h: h is packet
+        if name == "dsa_proj" and feature in ("deferred_norm", "deferred_norm_evict_hints"):
+            return lambda *args: calls.append(feature) or f"{feature}-packet"
+        return real_get_impl(feature, name)
+
+    monkeypatch.setattr(dispatch, "get_impl", fake_get_impl)
+    out = attention_forward._absorbed_query_packet(object(), None, packet, None)
+    expect = "deferred_norm_evict_hints" if with_evict_hints else "deferred_norm"
+    assert out == f"{expect}-packet" and calls == [expect]
+
+
+def test_decoder_layer_selects_the_hop_factory():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmDecoderLayer,
+    )
+
+    src = inspect.getsource(GlmDecoderLayer._adopt_post_moe)
+    # The formatter wraps long statements: compare with the whitespace collapsed.
+    src = re.sub(r"(?<=[(\[])\s+|\s+(?=[)\]])", "", " ".join(src.split()))
+    assert (
+        '"deferred_norm_evict_hints" if dispatch.is_enabled("deferred_norm_evict_hints") else "deferred_norm"'
+        in src
+    )
+    # The choice is made before the factory is called (host-side selector).
+    head, tail = src.split(
+        '"deferred_norm_evict_hints" if dispatch.is_enabled("deferred_norm_evict_hints") else "deferred_norm"',
+        1,
+    )
+    assert head.rstrip().endswith("get_impl(")
+    assert re.match(r', "deferred_hop",?\)\(self\.post_moe_hop\)', tail), tail[:60]
+
+
+def test_evict_hints_registers_its_entry_points(select):
+    dispatch = select(BASE + ",deferred_norm_evict_hints")
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        is_deferred,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
+        EvictHintsDeferredHop,
+        evict_hints_dsa_proj,
+    )
+
+    assert dispatch.get_impl("deferred_norm_evict_hints", "deferred_hop") is EvictHintsDeferredHop
+    assert dispatch.get_impl("deferred_norm_evict_hints", "dsa_proj") is evict_hints_dsa_proj
+    assert dispatch.get_impl("deferred_norm_evict_hints", "is_deferred") is is_deferred
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py
new file mode 100644
index 000000000000..79ad8ea35702
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py
@@ -0,0 +1,85 @@
+# 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.
+"""Host-only tests of the deferred_norm_evict_hints configuration parsing (no GPU work)."""
+
+import dataclasses
+
+
+def test_pre_wait_prefetch_is_a_test_arm():
+    """The served hop configuration prefetches nothing before its dependency wait; tests build
+    the pre-wait arms explicitly."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.hop import (  # noqa: E501
+        PRE_WAIT_PREFETCH,
+        EvictHintsHopConfig,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
+        parse_cfg,
+    )
+
+    assert EvictHintsHopConfig().poll_prefetch == "none"
+    assert EvictHintsHopConfig().poll_prefetch not in PRE_WAIT_PREFETCH
+    for mode in PRE_WAIT_PREFETCH + ("none", "wait3", "t3"):
+        assert parse_cfg(f"poll_prefetch={mode}", EvictHintsHopConfig).poll_prefetch == mode
+
+
+def test_defaults_are_deferred_norm_launch():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormProjConfig,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
+        EvictHintsHopConfig,
+        EvictHintsProjConfig,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
+        EVICT_HINTS_CONFIG_KEYS,
+    )
+
+    z = EvictHintsProjConfig()
+    w = DeferredNormProjConfig()
+    for f in dataclasses.fields(DeferredNormProjConfig):
+        assert getattr(z, f.name) == getattr(w, f.name), f.name
+    extra = [
+        f.name
+        for f in dataclasses.fields(EvictHintsProjConfig)
+        if f.name not in {g.name for g in dataclasses.fields(DeferredNormProjConfig)}
+    ]
+    assert tuple(extra) == EVICT_HINTS_CONFIG_KEYS
+    # Served default: deferred_norm's launch plus the evict_first policy on the weight copies.
+    assert not any(getattr(z, k) for k in EVICT_HINTS_CONFIG_KEYS if k != "copy_hint")
+    assert z.copy_hint == "evict_first"
+    h = EvictHintsHopConfig()
+    assert (h.carveout, h.poll_prefetch, h.clear_hint, h.ctas, h.ffi, h.dyn_smem) == (
+        -1,
+        "none",
+        "none",
+        8,
+        True,
+        0,
+    )
+
+
+def test_full_indexer_layers_keep_deferred_norm_launch():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
+        EvictHintsProjConfig,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
+        config_for,
+        parse_cfg,
+    )
+
+    cfg = parse_cfg("prod_warp=1,split_chunk=1,split_delay_ns=3000", EvictHintsProjConfig)
+    assert config_for(cfg, 0) is cfg
+    assert config_for(cfg, 160) == EvictHintsProjConfig(copy_hint="none")
+    assert config_for(EvictHintsProjConfig(), 160).copy_hint == "none"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py
new file mode 100644
index 000000000000..26bc7fda6289
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py
@@ -0,0 +1,424 @@
+# 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.
+"""Collective tests of the deferred_norm_evict_hints hop against F's deferred_norm hop (8 ranks, 2 hosts).
+
+For every deferred_norm_evict_hints hop configuration (``CONFIGS``: four experiment arms, the
+defaults ``carveout=-1,poll_prefetch=none`` through TVM-FFI, and the
+co-resident executor hop ``carveout=100,poll_prefetch=none,ffi=0``), on the
+same post_moe_hop workspace and the same inputs as deferred_norm's hop (routed_moe_into_hop routed
+input in each mode, and the raw tuple), every test below runs per configuration:
+
+* ``residual_out``, ``a_frag``, ``sumsq`` and ``materialize()`` are bit-identical
+  to deferred_norm's (and so to post_moe_hop's ``hop_routed`` / ``__call__``, which deferred_norm's
+  own suite checks); the fp32 accumulator is zero after the call;
+* deferred_norm_evict_hints, deferred_norm and post_moe_hop calls interleave on one workspace
+  without host syncs;
+* a CUDA graph of deferred_norm_evict_hints / deferred_norm / post_moe_hop calls replays with fresh
+  inputs;
+* random per-rank device delays and fault-injected store order;
+* a missing rank times out with NaN, error bits and the post_moe_hop health check,
+  and a collective reset restores exact results.
+"""
+
+import os
+import sys
+import time
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
+sys.path.insert(0, os.path.join(_HERE, "..", "deferred_norm"))
+
+import post_moe_hop_mpi_utils as M  # noqa: E402
+from test_deferred_norm_routed_mpi import _fill, _source  # noqa: E402
+
+CONFIGS = os.environ.get(
+    "GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_HOP_CFGS",
+    "carveout=100,poll_prefetch=wait3,ffi=0;"
+    "carveout=100,poll_prefetch=cur,ffi=0;carveout=100,poll_prefetch=pre3,ffi=0;"
+    "carveout=100,poll_prefetch=wait3,clear_hint=evict_last,ffi=0;"
+    "carveout=-1,poll_prefetch=none;carveout=100,poll_prefetch=none,ffi=0",
+).split(";")
+MODES = ("fp32", "identity", "row")
+_HOPS = {}
+
+
+def _env():
+    env = M.mpi_env()
+    if env.world != 8:
+        pytest.skip("the hop is specialized to TP8")
+    return env
+
+
+def _hops(mode, cfg, timeout_s=None, fault_delay_ns=0):
+    """post_moe_hop (routed input ``mode``), deferred_norm hop and deferred_norm_evict_hints hop
+    ``cfg`` on one workspace."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormHop,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
+        EvictHintsDeferredHop,
+        EvictHintsHopConfig,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
+        parse_cfg,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    env = _env()
+    key = (mode, timeout_s, fault_delay_ns)
+    if key not in _HOPS:
+        kw = {}
+        if timeout_s is not None:
+            kw["timeout_s"] = timeout_s
+        if fault_delay_ns:
+            kw["fault_delay_ns"] = fault_delay_ns
+        post_moe_hop = PostMoeHop(
+            env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False, **kw
+        )
+        assert M.all_ranks(post_moe_hop.engaged and post_moe_hop.enable_routed_input(mode)), (
+            post_moe_hop.reason
+        )
+        deferred_norm = DeferredNormHop(post_moe_hop, trigger_early=True)
+        assert M.all_ranks(deferred_norm.engaged and deferred_norm.enable_routed_input())
+        _HOPS[key] = {"post_moe_hop": post_moe_hop, "deferred_norm": deferred_norm}
+    entry = _HOPS[key]
+    if cfg not in entry:
+        hop_config = EvictHintsDeferredHop(
+            entry["post_moe_hop"], trigger_early=True, config=parse_cfg(cfg, EvictHintsHopConfig)
+        )
+        assert M.all_ranks(hop_config.engaged and hop_config.enable_routed_input()), (
+            hop_config.reason,
+            hop_config.routed_reason,
+        )
+        entry[cfg] = hop_config
+    return entry["post_moe_hop"], entry["deferred_norm"], entry[cfg]
+
+
+def _bits(t):
+    return t.view(torch.int16) if t.element_size() == 2 else t.view(torch.int32)
+
+
+def _outs(packet, res):
+    """Residual, the R written rows of ``a_frag``, ``sumsq`` and the materialization."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        materialize,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
+        rows_from_frag,
+    )
+
+    h = materialize(packet)
+    torch.cuda.synchronize()
+    a = rows_from_frag(packet.a_frag, res.shape[0], res.shape[1])
+    return [res.clone(), a.clone(), packet.sumsq.clone(), h.clone()]
+
+
+def _same(a, b, tag):
+    same = [torch.equal(_bits(x), _bits(y)) for x, y in zip(a, b)]
+    if not all(same):
+        print(
+            f"GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_HOP_FAIL {tag} rank={M.mpi_env().rank} "
+            f"residual/a_frag/sumsq/h={same}",
+            flush=True,
+        )
+    return all(same)
+
+
+def _acc_clear(post_moe_hop, mode):
+    if mode != "fp32":
+        return True
+    torch.cuda.synchronize()
+    return bool((post_moe_hop.routed_acc_buffer(post_moe_hop.max_rows) == 0).all())
+
+
+def _routed(post_moe_hop, hop, mode, src, rows, shared, residual, gamma):
+    routed = _fill(post_moe_hop, mode, src, rows)
+    return hop.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
+
+
+@pytest.mark.parametrize("with_shared", [True, False])
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+@pytest.mark.parametrize("mode", MODES)
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_routed_matches_deferred_norm(cfg, mode, rows, with_shared):
+    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, rows, seed=50 + rows, with_shared=with_shared)
+        ref = _outs(*_routed(post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma))
+        got = _outs(*_routed(post_moe_hop, hop_config, mode, src, rows, shared, residual, gamma))
+        ok = _same(got, ref, f"{cfg} {mode} R={rows}") and _acc_clear(post_moe_hop, mode)
+        assert M.all_ranks(ok)
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_raw_tuple_matches_deferred_norm(cfg, rows):
+    post_moe_hop, deferred_norm, hop_config = _hops("fp32", cfg)
+    with torch.inference_mode():
+        t = M.make_raw_tuple(rows, seed=300 + rows)
+        ref = _outs(*deferred_norm(*t, M.GLM_EPS))
+        got = _outs(*hop_config(*t, M.GLM_EPS))
+        assert M.all_ranks(_same(got, ref, f"{cfg} raw R={rows}"))
+
+
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_interleaved_with_deferred_norm_and_post_moe_hop(cfg):
+    """deferred_norm_evict_hints / deferred_norm / post_moe_hop calls (routed and raw) back to back on one workspace."""
+    mode = "fp32"
+    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
+    kinds = [
+        "routed_deferred_norm_evict_hints",
+        "routed_deferred_norm",
+        "routed_post_moe_hop",
+        "deferred_norm_evict_hints",
+        "post_moe_hop",
+        "routed_deferred_norm_evict_hints",
+        "routed_deferred_norm_evict_hints",
+        "deferred_norm",
+        "routed_post_moe_hop",
+        "routed_deferred_norm_evict_hints",
+        "deferred_norm_evict_hints",
+        "routed_deferred_norm",
+    ]
+    rows_seq = [6, 1, 8, 3, 6, 2, 6, 8, 5, 6, 1, 6]
+    with torch.inference_mode():
+        inputs = [
+            (
+                _source(mode, r, seed=900 + i)
+                if k.startswith("routed_")
+                else M.make_raw_tuple(r, seed=900 + i)
+            )
+            for i, (k, r) in enumerate(zip(kinds, rows_seq))
+        ]
+        outs = []
+        for kind, r, t in zip(kinds, rows_seq, inputs):
+            if kind.startswith("routed_"):
+                hop = {
+                    "routed_deferred_norm_evict_hints": hop_config,
+                    "routed_deferred_norm": deferred_norm,
+                }.get(kind)
+                if hop is None:
+                    outs.append(
+                        post_moe_hop.hop_routed(
+                            _fill(post_moe_hop, mode, t[0], r), *t[1:], M.GLM_EPS
+                        )
+                    )
+                else:
+                    outs.append(_routed(post_moe_hop, hop, mode, t[0], r, *t[1:]))
+            else:
+                hop = {
+                    "deferred_norm_evict_hints": hop_config,
+                    "deferred_norm": deferred_norm,
+                    "post_moe_hop": post_moe_hop,
+                }[kind]
+                outs.append(hop(*t, M.GLM_EPS))
+        torch.cuda.synchronize()
+        snap = [
+            [x.clone() for x in o] if kind.endswith("post_moe_hop") else _outs(*o)
+            for kind, o in zip(kinds, outs)
+        ]
+        failures = []
+        for i, (kind, r, t, got) in enumerate(zip(kinds, rows_seq, inputs, snap)):
+            if kind.endswith("post_moe_hop"):
+                if kind.startswith("routed_"):
+                    ref = post_moe_hop.hop_routed(
+                        _fill(post_moe_hop, mode, t[0], r), *t[1:], M.GLM_EPS
+                    )
+                else:
+                    ref = post_moe_hop(*t, M.GLM_EPS)
+                torch.cuda.synchronize()
+                ok = torch.equal(got[0], ref[0]) and torch.equal(got[1], ref[1])
+            else:
+                if kind.startswith("routed_"):
+                    ref = _outs(*_routed(post_moe_hop, deferred_norm, mode, t[0], r, *t[1:]))
+                else:
+                    ref = _outs(*deferred_norm(*t, M.GLM_EPS))
+                ok = _same(got, ref, f"interleaved{i}-{kind}")
+            if not ok:
+                failures.append((i, kind))
+        assert M.all_ranks(not failures and _acc_clear(post_moe_hop, mode)), failures
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_cuda_graph_replay(cfg):
+    """Fill + deferred_norm_evict_hints routed, fill + deferred_norm routed, a
+    deferred_norm_evict_hints raw call in one graph, replayed."""
+    mode = "fp32"
+    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
+    rows = 6
+    with torch.inference_mode():
+        bufs = [list(_source(mode, rows, seed=1200 + i)) for i in range(2)]
+        raw = [x.clone() if x is not None else None for x in M.make_raw_tuple(rows, seed=1250)]
+
+        def body():
+            a = _routed(post_moe_hop, hop_config, mode, bufs[0][0], rows, *bufs[0][1:])
+            b = _routed(post_moe_hop, deferred_norm, mode, bufs[1][0], rows, *bufs[1][1:])
+            c = hop_config(*raw, M.GLM_EPS)
+            return a, b, c
+
+        body()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            ga, gb, gc = body()
+        for trial, replays in enumerate((1, 1, 20)):
+            fresh = [_source(mode, rows, seed=1300 + 10 * trial + i) for i in range(2)]
+            fresh_raw = M.make_raw_tuple(rows, seed=1350 + trial)
+            for dst, src in zip(bufs, fresh):
+                for d, s in zip(dst, src):
+                    if d is not None:
+                        d.copy_(s)
+            for d, s in zip(raw, fresh_raw):
+                if d is not None:
+                    d.copy_(s)
+            for _ in range(replays):
+                M.mpi_env().comm.barrier()
+                graph.replay()
+            torch.cuda.synchronize()
+            got = [_outs(*ga), _outs(*gb), _outs(*gc)]
+            refs = [
+                _outs(
+                    *_routed(post_moe_hop, deferred_norm, mode, fresh[0][0], rows, *fresh[0][1:])
+                ),
+                _outs(
+                    *_routed(post_moe_hop, deferred_norm, mode, fresh[1][0], rows, *fresh[1][1:])
+                ),
+                _outs(*deferred_norm(*fresh_raw, M.GLM_EPS)),
+            ]
+            ok = all(_same(g, r, f"graph{trial}-{j}") for j, (g, r) in enumerate(zip(got, refs)))
+            assert M.all_ranks(ok and _acc_clear(post_moe_hop, mode)), f"trial {trial}"
+
+
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_delayed_rank_stress(cfg):
+    """Random per-rank device delays inside back-to-back deferred_norm_evict_hints routed batches."""
+    mode = "fp32"
+    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
+    env = M.mpi_env()
+    gen = torch.Generator().manual_seed(4243)
+    with torch.inference_mode():
+        for batch in range(3):
+            inputs, outs = [], []
+            for i in range(12):
+                rows = int(torch.randint(1, 9, (1,), generator=gen))
+                t = _source(mode, rows, seed=5000 + 100 * batch + i)
+                local = torch.Generator().manual_seed(77 * batch + 13 * i + env.rank)
+                if float(torch.rand(1, generator=local)) < 0.3:
+                    torch.cuda._sleep(int(float(torch.rand(1, generator=local)) * 2_000_000))
+                inputs.append((rows, t))
+                outs.append(_routed(post_moe_hop, hop_config, mode, t[0], rows, *t[1:]))
+            torch.cuda.synchronize()
+            snap = [_outs(*o) for o in outs]
+            ok = True
+            for (rows, t), got in zip(inputs, snap):
+                ok &= _same(
+                    got,
+                    _outs(*_routed(post_moe_hop, deferred_norm, mode, t[0], rows, *t[1:])),
+                    f"stress{batch}",
+                )
+            assert M.all_ranks(ok), f"batch {batch}"
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_fault_injected_store_order(cfg):
+    """Scrambled store order (post_moe_hop's fault injection) in the deferred_norm_evict_hints kernel."""
+    mode = "fp32"
+    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg, fault_delay_ns=4000)
+    with torch.inference_mode():
+        inputs = [(r, _source(mode, r, seed=7000 + i)) for i, r in enumerate([6, 1, 8, 6, 3, 6])]
+        outs = [_routed(post_moe_hop, hop_config, mode, t[0], r, *t[1:]) for r, t in inputs]
+        torch.cuda.synchronize()
+        snap = [_outs(*o) for o in outs]
+        ok = True
+        for (r, t), got in zip(inputs, snap):
+            ok &= _same(
+                got, _outs(*_routed(post_moe_hop, deferred_norm, mode, t[0], r, *t[1:])), "fault"
+            )
+        assert M.all_ranks(ok)
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+
+
+_MISSING_TIMEOUT_S = 0.2
+
+
+@pytest.mark.parametrize("missing", [0, 5])
+@pytest.mark.parametrize("cfg", CONFIGS)
+def test_missing_rank_times_out_with_nan_and_reset(cfg, missing):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+
+    mode = "fp32"
+    post_moe_hop, deferred_norm, _ = _hops(mode, cfg)
+    imp_post_moe_hop, _, imp_evict_hints = _hops(mode, cfg, timeout_s=_MISSING_TIMEOUT_S)
+    env = M.mpi_env()
+    assert M.all_ranks(imp_post_moe_hop.workspace is post_moe_hop.workspace)
+    rows = 6
+    with torch.inference_mode():
+        src, shared, residual, gamma = _source(mode, rows, seed=600 + missing)
+        env.comm.barrier()
+        start = time.monotonic()
+        snap = None
+        if env.rank != missing:
+            snap = _outs(
+                *_routed(
+                    imp_post_moe_hop, imp_evict_hints, mode, src, rows, shared, residual, gamma
+                )
+            )
+        torch.cuda.synchronize()
+        elapsed = time.monotonic() - start
+        counters = post_moe_hop.workspace.counters()
+        bits, polls = counters["error_bits"], counters["timeout_polls"]
+        if env.rank == missing:
+            reported = bits == 0 and polls == 0
+        else:
+            nan = all(bool(torch.isnan(t.float()).all()) for t in (snap[0], snap[2], snap[3]))
+            reported = (
+                bool(bits & 1)
+                and polls > 0
+                and _MISSING_TIMEOUT_S <= elapsed < 30.0
+                and nan
+                and _acc_clear(imp_post_moe_hop, mode)
+            )
+        stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
+        if env.rank == 0:
+            print(
+                f"GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_MISSING cfg={cfg} rank {missing} skipped: {stats}",
+                flush=True,
+            )
+        assert M.all_ranks(reported), stats
+        try:
+            check_health(f"deferred_norm_evict_hints-missing-rank{missing}", collective=True)
+            raised = False
+        except PostMoeHopError:
+            raised = True
+        assert M.all_ranks(raised)
+        assert M.all_ranks(post_moe_hop.reset())
+        assert M.all_ranks(post_moe_hop.error_bits() == 0)
+        src, shared, residual, gamma = _source(mode, rows, seed=650 + missing)
+        _, _, hop_config = _hops(mode, cfg)
+        got = _outs(*_routed(post_moe_hop, hop_config, mode, src, rows, shared, residual, gamma))
+        ref = _outs(*_routed(post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma))
+        assert M.all_ranks(_same(got, ref, f"after-reset{missing}"))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py
new file mode 100644
index 000000000000..7ee030e1624f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py
@@ -0,0 +1,156 @@
+# 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.
+"""Single-GPU tests of the deferred_norm_evict_hints projection against F's deferred_norm projection (real weights).
+
+For every deferred_norm_evict_hints projection configuration (``CONFIGS``) and every layer
+the projection serves (``GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_LAYERS``: a shared-indexer
+target layer; the target's deferred-norm scope is ``shared``, so full-indexer layers keep their
+native projection), the fused-A output
+is bit-identical to ``DeferredNormProjection``'s on the same deferred packet, for R in
+1..8 and residual scales from 1e-3 to 30, in eager mode and in a CUDA graph
+with PDL. ``precompile_served`` covers the consumer's launch configuration.
+
+Fixture label: weights real (checkpoint), activations synthetic.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
+sys.path.insert(0, os.path.join(_HERE, "..", "deferred_norm"))
+
+import absorbed_query_fixture  # noqa: E402
+from test_deferred_norm_projection import _module, make_packet, make_residual  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
+    reason="needs a GPU and the GLM-5.2 checkpoint",
+)
+
+LAYERS = [
+    int(x)
+    for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_LAYERS", "7").split(",")
+]
+CONFIGS = os.environ.get(
+    "GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_PROJ_CFGS",
+    "maxrreg=112;maxrreg=0,maxntid=576;maxrreg=112,copy_hint=evict_first;"
+    "maxrreg=112,split_chunk=2,split_delay_ns=2500;maxrreg=0;copy_warps=1;prod_warp=1;"
+    "prod_warp=1,split_chunk=1,split_delay_ns=3000;"
+    "prod_warp=1,split_chunk=2,split_delay_ns=2500,copy_hint=evict_first;"
+    "prod_warp=1,defer_tail=8,defer_wait=1;prod_warp=1,split_chunk=1,defer_tail=3,defer_wait=1;"
+    "prod_warp=1,defer_tail=5,split_delay_ns=2000;",
+).split(";")
+
+
+def _pair(layer, cfg_text):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
+        DeferredNormProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
+        EvictHintsProjConfig,
+        EvictHintsProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
+        parse_cfg,
+    )
+
+    attn, gamma = _module(layer)
+    ok, reason, kw = projection_eligibility(attn)
+    assert ok, reason
+    deferred_norm = DeferredNormProjection(attn, kw_weight=kw)
+    hop_config = EvictHintsProjection(attn, parse_cfg(cfg_text, EvictHintsProjConfig), kw_weight=kw)
+    assert hop_config.supported() is None
+    return attn, gamma, deferred_norm, hop_config
+
+
+def _bits(t):
+    return None if t is None else t.view(torch.int16)
+
+
+def _equal(a, b):
+    return all((x is None and y is None) or torch.equal(_bits(x), _bits(y)) for x, y in zip(a, b))
+
+
+@pytest.mark.parametrize("cfg", CONFIGS)
+@pytest.mark.parametrize("layer", LAYERS)
+def test_served_configs_cover_the_consumer(layer, cfg):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.proj_kernel import (  # noqa: E501
+        is_compiled,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
+        EvictHintsProjConfig,
+        parse_cfg,
+        precompile_served,
+        served_kernel_configs,
+    )
+
+    config = parse_cfg(cfg, EvictHintsProjConfig)
+    device = torch.cuda.current_device()
+    assert precompile_served(device, config) is None
+    _, _, _, hop_config = _pair(layer, cfg)
+    assert hop_config.served_config() in served_kernel_configs(device, config)
+    assert is_compiled(**hop_config.served_config())
+
+
+@pytest.mark.parametrize("scale", [1.0, 1e-3, 30.0])
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+@pytest.mark.parametrize("cfg", CONFIGS)
+@pytest.mark.parametrize("layer", LAYERS)
+def test_bit_identical_to_deferred_norm(layer, cfg, rows, scale):
+    _, gamma, deferred_norm, hop_config = _pair(layer, cfg)
+    with torch.inference_mode():
+        x = make_residual(rows, seed=100 + rows + layer, scale=scale)
+        packet = make_packet(x, gamma)
+        ref = deferred_norm.project(packet)
+        got = hop_config.project(packet)
+        torch.cuda.synchronize()
+        assert _equal(got, ref)
+
+
+@pytest.mark.parametrize("cfg", CONFIGS)
+@pytest.mark.parametrize("layer", LAYERS)
+def test_graph_with_pdl_bit_identical(layer, cfg):
+    """deferred_norm_evict_hints then deferred_norm projection of 4 packets in one CUDA graph, replayed with
+    fresh packets (PDL chain between the projections)."""
+    _, gamma, deferred_norm, hop_config = _pair(layer, cfg)
+    with torch.inference_mode():
+        packets = [make_packet(make_residual(6, seed=500 + i), gamma) for i in range(4)]
+
+        def body():
+            return [(hop_config.project(p), deferred_norm.project(p)) for p in packets]
+
+        body()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            outs = body()
+        for trial in range(3):
+            for i, p in enumerate(packets):
+                fresh = make_packet(make_residual(6, seed=600 + 10 * trial + i), gamma)
+                p.a_frag.copy_(fresh.a_frag)
+                p.sumsq.copy_(fresh.sumsq)
+            graph.replay()
+            torch.cuda.synchronize()
+            for (got, ref), p in zip(outs, packets):
+                assert _equal(got, ref)
+                assert _equal(got, deferred_norm.project(p))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py
new file mode 100644
index 000000000000..b8c481449425
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py
@@ -0,0 +1,190 @@
+# 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.
+"""Shared helpers for the dense_mlp tests and microbenchmarks.
+
+* ``ensure_op``: make ``trtllm::glm_moe_dsa_dense_mlp*`` available. An image built from
+  a commit that contains the dense_mlp thop sources already registers the ops; for
+  in-job development on an older image the same sources are JIT-compiled.
+* Native oracle: the exact native BF16 ``GatedMLP`` compute (``F.linear`` ->
+  TRT-LLM ``swiglu`` -> ``F.linear``) and an FP64 reference with the native
+  BF16 intermediate rounding.
+* Real data: GLM-5.2 dense-layer weight shards as the GEN loader lays them out
+  (TP4 sub-group, fused [gate; up]) and real post-RMSNorm hidden rows.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from utils.llm_data import llm_models_root
+
+REPO_ROOT = Path(__file__).resolve().parents[9]
+DEFAULT_CKPT = os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4")
+DEFAULT_REAL_HIDDEN = os.environ.get("GLM_MOE_DSA_REAL_MOE_INPUTS_DIR", "")
+GLM52_HIDDEN = 6144
+GLM52_INTER = 12288
+GLM52_DENSE_TP = 4
+
+
+def ensure_op(build_dir: Optional[str] = None, verbose: bool = False) -> str:
+    """Return "builtin" or "jit" after making the dense_mlp ops available."""
+    import tensorrt_llm  # noqa: F401  (loads th_common and its op fragments)
+
+    if hasattr(torch.ops.trtllm, "glm_moe_dsa_dense_mlp"):
+        return "builtin"
+    from torch.utils.cpp_extension import load
+
+    libs_dir = Path(tensorrt_llm.__file__).resolve().parent / "libs"
+    build_dir = build_dir or os.environ.get("GLM_MOE_DSA_DENSE_MLP_JIT_BUILD_DIR")
+    if build_dir:
+        Path(build_dir).mkdir(parents=True, exist_ok=True)
+    load(
+        name="dense_mlp_jit",
+        sources=[
+            str(REPO_ROOT / "cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu"),
+            str(REPO_ROOT / "cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp"),
+        ],
+        extra_include_paths=[
+            str(REPO_ROOT / "cpp/include"),
+            str(REPO_ROOT / "cpp"),
+        ],
+        extra_cflags=["-O3", "-DENABLE_BF16", "-DENABLE_FP8"],
+        extra_cuda_cflags=[
+            "-O3",
+            "-DENABLE_BF16",
+            "-DENABLE_FP8",
+            "-gencode=arch=compute_103a,code=sm_103a",
+            "-lineinfo",
+        ],
+        extra_ldflags=[f"-L{libs_dir}", "-ltensorrt_llm", f"-Wl,-rpath,{libs_dir}"],
+        build_directory=build_dir,
+        is_python_module=False,
+        verbose=verbose,
+    )
+    assert hasattr(torch.ops.trtllm, "glm_moe_dsa_dense_mlp")
+    return "jit"
+
+
+def native_dense_mlp(
+    x: torch.Tensor, gate_up_w: torch.Tensor, down_w: torch.Tensor
+) -> torch.Tensor:
+    """Native BF16 GatedMLP compute without the output all-reduce."""
+    from tensorrt_llm._torch.modules.swiglu import swiglu
+
+    h1 = F.linear(x, gate_up_w)
+    h2 = swiglu(h1)
+    return F.linear(h2, down_w)
+
+
+def native_intermediate(x: torch.Tensor, gate_up_w: torch.Tensor) -> torch.Tensor:
+    from tensorrt_llm._torch.modules.swiglu import swiglu
+
+    return swiglu(F.linear(x, gate_up_w))
+
+
+def oracle_dense_mlp(
+    x: torch.Tensor, gate_up_w: torch.Tensor, down_w: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """FP64 sums with the native BF16 rounding points.
+
+    Returns ``(h, out)`` as FP64 tensors whose values are BF16-representable
+    where the native path rounds (gate, up, h); ``out`` is unrounded.
+    """
+    inter = down_w.shape[1]
+    gu = x.double() @ gate_up_w.double().t()
+    gate = gu[:, :inter].to(torch.bfloat16).double()
+    up = gu[:, inter:].to(torch.bfloat16).double()
+    h = (gate * torch.sigmoid(gate) * up).to(torch.bfloat16).double()
+    out = h @ down_w.double().t()
+    return h, out
+
+
+def error_stats(value: torch.Tensor, ref: torch.Tensor) -> Dict[str, float]:
+    diff = value.double() - ref.double()
+    scale = ref.double().pow(2).mean().sqrt().item()
+    return {
+        "rms": diff.pow(2).mean().sqrt().item(),
+        "max": diff.abs().max().item(),
+        "nrms": diff.pow(2).mean().sqrt().item() / max(scale, 1e-30),
+        "ref_rms": scale,
+    }
+
+
+def _safetensors_index(ckpt: str) -> Dict[str, str]:
+    with open(os.path.join(ckpt, "model.safetensors.index.json")) as f:
+        return json.load(f)["weight_map"]
+
+
+def load_dense_shard(
+    layer: int, shard: int, ckpt: str = DEFAULT_CKPT, tp: int = GLM52_DENSE_TP, device: str = "cuda"
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Return GEN-layout ``(gate_up [2I/tp, H], down [H, I/tp])`` of one rank.
+
+    Matches the native loader: gate rows ``[I/tp * shard, ...)`` then the same
+    up rows; down input columns of that shard.
+    """
+    from safetensors import safe_open
+
+    index = _safetensors_index(ckpt)
+    prefix = f"model.layers.{layer}.mlp."
+    tensors = {}
+    for name in ("gate_proj", "up_proj", "down_proj"):
+        key = prefix + name + ".weight"
+        with safe_open(os.path.join(ckpt, index[key]), framework="pt") as f:
+            sl = f.get_slice(key)
+            rows, cols = sl.get_shape()
+            if name == "down_proj":
+                width = cols // tp
+                tensors[name] = sl[:, width * shard : width * (shard + 1)]
+            else:
+                width = rows // tp
+                tensors[name] = sl[width * shard : width * (shard + 1), :]
+    gate_up = torch.cat([tensors["gate_proj"], tensors["up_proj"]], dim=0)
+    down = tensors["down_proj"].contiguous()
+    assert gate_up.dtype == torch.bfloat16 and down.dtype == torch.bfloat16
+    return gate_up.to(device).contiguous(), down.to(device).contiguous()
+
+
+def load_real_hidden(
+    fixture_dir: str = DEFAULT_REAL_HIDDEN, layer: int = 10, tokens: int = 2311
+) -> Optional[torch.Tensor]:
+    """Real post-RMSNorm hidden rows [tokens, 6144] (BF16) or None if absent.
+
+    The fixture captured the routed-MoE input (``post_attention_layernorm``
+    output) of the native model; it serves as a realistic activation for the
+    dense MLP, which consumes the same kind of tensor in layers 0-2.
+    """
+    path = os.path.join(fixture_dir, f"rows{tokens}_layer{layer}.pt")
+    if not os.path.exists(path):
+        return None
+    obj = torch.load(path, map_location="cpu", weights_only=False)
+    if isinstance(obj, dict):
+        for key in ("hidden_states", "hidden", "x", "input"):
+            if key in obj and isinstance(obj[key], torch.Tensor):
+                return obj[key].to(torch.bfloat16)
+        for value in obj.values():
+            if (
+                isinstance(value, torch.Tensor)
+                and value.dim() == 2
+                and value.shape[-1] == GLM52_HIDDEN
+            ):
+                return value.to(torch.bfloat16)
+        return None
+    if isinstance(obj, torch.Tensor):
+        return obj.to(torch.bfloat16)
+    return None
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
new file mode 100644
index 000000000000..19530f200648
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
@@ -0,0 +1,73 @@
+/*
+ * 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.
+ */
+
+// Test-only delayed producer for programmatic dependent launch: releases its dependents immediately,
+// spins for a given number of nanoseconds, then copies src into dst. A PDL consumer that reads dst
+// before griddepcontrol.wait would observe stale data.
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+
+namespace
+{
+
+__global__ void delayedCopyKernel(uint4 const* __restrict__ src, uint4* __restrict__ dst, int64_t numVec, int64_t delayNs)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+#endif
+    uint64_t start;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(start));
+    uint64_t now = start;
+    while (now - start < static_cast<uint64_t>(delayNs))
+    {
+        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now));
+    }
+    for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < numVec;
+         i += static_cast<int64_t>(gridDim.x) * blockDim.x)
+    {
+        dst[i] = src[i];
+    }
+}
+
+} // namespace
+
+void delayedCopy(torch::Tensor src, torch::Tensor dst, int64_t delayNs)
+{
+    TORCH_CHECK(src.is_contiguous() && dst.is_contiguous() && src.nbytes() == dst.nbytes() && src.nbytes() % 16 == 0);
+    auto const stream = at::cuda::getCurrentCUDAStream(src.get_device());
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(4);
+    config.blockDim = dim3(128);
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = 1;
+    config.attrs = attrs;
+    config.numAttrs = 1;
+    TORCH_CHECK(cudaLaunchKernelEx(&config, delayedCopyKernel, reinterpret_cast<uint4 const*>(src.data_ptr()),
+                    reinterpret_cast<uint4*>(dst.data_ptr()), static_cast<int64_t>(src.nbytes() / 16), delayNs)
+        == cudaSuccess);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("delayed_copy", &delayedCopy);
+}
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py
new file mode 100644
index 000000000000..60af1ceb683d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py
@@ -0,0 +1,171 @@
+# 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.
+"""Tests of the dense_mlp call site in the target's decoder layer.
+
+``GlmDecoderLayer.forward_mlp`` is exercised on a decode step on a minimal layer object
+(dense GatedMLP, single rank, no fusion), with and without ``dense_mlp`` selected.
+"""
+
+import os
+import sys
+import types
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import plan_testlib  # noqa: E402
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+
+HIDDEN = 2048
+INTER = 4096
+
+
+@pytest.fixture(scope="module", autouse=True)
+def dense_mlp_op():
+    import dense_mlp_utils as wu
+
+    wu.ensure_op()
+
+
+@pytest.fixture(autouse=True)
+def decode_step():
+    """The owned call site runs on decode steps."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    with dispatch.step(dispatch.Phase.DECODE):
+        yield
+
+
+def _select(monkeypatch, value):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    plan_testlib.select(dispatch, value, monkeypatch)
+
+
+@pytest.fixture
+def spy(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    dispatch.get_impl("dense_mlp", "dense_mlp_forward")
+    real = dispatch._REGISTRY["dense_mlp"]["dense_mlp_forward"]
+
+    class Calls(list):
+        pass
+
+    calls = Calls()
+
+    def wrapped(*args, **kwargs):
+        result = real(*args, **kwargs)
+        calls.append(result is not None)
+        if result is not None:
+            wrapped.last = (args[0], args[1].clone(), args[2], result.clone())
+        return result
+
+    wrapped.last = None
+
+    monkeypatch.setitem(dispatch._REGISTRY["dense_mlp"], "dense_mlp_forward", wrapped)
+    calls.wrapped = wrapped
+    return calls
+
+
+def _layer():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmDecoderLayer,
+    )
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.models.modeling_utils import EagerFusionConfig
+    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+
+    torch.manual_seed(0)
+    mlp = GatedMLP(
+        hidden_size=HIDDEN,
+        intermediate_size=INTER,
+        bias=False,
+        dtype=torch.bfloat16,
+        config=ModelConfig(),
+        reduce_output=False,
+    ).cuda()
+    with torch.no_grad():
+        mlp.gate_up_proj.weight.normal_(0, 0.02)
+        mlp.down_proj.weight.normal_(0, 0.02)
+    layer = types.SimpleNamespace(
+        mlp=mlp,
+        fusion_config=EagerFusionConfig(),
+        mlp_tp_size=1,
+        allreduce=None,
+        layer_idx=0,
+        post_attention_layernorm=RMSNorm(hidden_size=HIDDEN, eps=1e-5, dtype=torch.bfloat16).cuda(),
+        next_layer_layernorm=RMSNorm(hidden_size=HIDDEN, eps=1e-5, dtype=torch.bfloat16).cuda(),
+    )
+    layer._apply_next_layer_layernorm = types.MethodType(
+        GlmDecoderLayer._apply_next_layer_layernorm, layer
+    )
+    return layer, GlmDecoderLayer.forward_mlp
+
+
+def _run(layer, forward_mlp, rows, seed=1):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    hidden = torch.randn(rows, HIDDEN, device="cuda", generator=g).to(torch.bfloat16)
+    residual = torch.randn(rows, HIDDEN, device="cuda", generator=g).to(torch.bfloat16)
+    return forward_mlp(layer, hidden, residual)
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_dense_mlp_engages_and_matches_native(monkeypatch, spy, rows):
+    layer, forward_mlp = _layer()
+    _select(monkeypatch, "")
+    _run(layer, forward_mlp, rows)
+    assert spy == []
+    _select(monkeypatch, "dense_mlp")
+    hidden, residual = _run(layer, forward_mlp, rows)
+    assert spy == [True]
+    # The MLP output handed back to forward_mlp matches the native module on
+    # the same input within the BF16 GEMM tolerance.
+    mlp, x_in, params, dense_mlp_out = spy.wrapped.last
+    torch.testing.assert_close(
+        dense_mlp_out, mlp(x_in, final_all_reduce_params=params), rtol=1e-2, atol=5e-3
+    )
+    # Everything after the MLP stays native: the native forward_mlp fed with
+    # the same MLP output reproduces the dense_mlp layer outputs bit for bit.
+    _select(monkeypatch, "")
+    native_mlp = layer.mlp
+    layer.mlp = lambda x, final_all_reduce_params=None: dense_mlp_out.clone()
+    ref_hidden, ref_residual = _run(layer, forward_mlp, rows)
+    layer.mlp = native_mlp
+    assert torch.equal(hidden, ref_hidden)
+    assert torch.equal(residual, ref_residual)
+
+
+def test_ineligible_rows_take_native_path(monkeypatch, spy):
+    layer, forward_mlp = _layer()
+    _select(monkeypatch, "")
+    nat_hidden, nat_residual = _run(layer, forward_mlp, 12)
+    _select(monkeypatch, "dense_mlp")
+    hidden, residual = _run(layer, forward_mlp, 12)
+    assert spy == [False]
+    assert torch.equal(hidden, nat_hidden)
+    assert torch.equal(residual, nat_residual)
+
+
+def test_other_features_do_not_engage_dense_mlp(monkeypatch, spy):
+    layer, forward_mlp = _layer()
+    _select(monkeypatch, "absorbed_query")
+    _run(layer, forward_mlp, 6)
+    assert spy == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py
new file mode 100644
index 000000000000..6e266ee0e1c1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py
@@ -0,0 +1,410 @@
+# 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.
+"""GPU tests of package dense_mlp (small-M BF16 dense MLP).
+
+Oracle: the native BF16 GatedMLP compute (cuBLAS ``F.linear`` -> TRT-LLM
+``swiglu`` -> ``F.linear``) and an FP64 reference with native BF16 rounding.
+Tolerance against native: the existing TRT-LLM small-M BF16 GEMM tolerance
+(``tests/unittest/_torch/modules/test_low_m_gemm.py``: rtol 1e-2, atol 5e-3)
+for each GEMM stage. Error budget against the FP64 oracle, frozen before
+tuning: candidate normalized RMS error <= 1.25 x native + 1e-4.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from utils.llm_data import llm_models_root
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+
+NRMS_FACTOR = 1.25
+NRMS_SLACK = 1e-4
+
+
+@pytest.fixture(scope="module", autouse=True)
+def dense_mlp_op():
+    import dense_mlp_utils as wu
+
+    return wu.ensure_op()
+
+
+@pytest.fixture(autouse=True)
+def _seed_inputs():
+    """Seed the global generator so each case draws the same inputs in any order."""
+    torch.manual_seed(0)
+
+
+def _random_weights(hidden, inter, scale=0.02, seed=0):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    gate_up = (scale * torch.randn(2 * inter, hidden, device="cuda", generator=g)).to(
+        torch.bfloat16
+    )
+    down = (scale * torch.randn(hidden, inter, device="cuda", generator=g)).to(torch.bfloat16)
+    return gate_up, down
+
+
+def assert_within_bf16_ulps(actual, expected, ulps, abs_floor_rel=1e-5):
+    """|actual - expected| <= ulps BF16 ulps of the larger magnitude + floor.
+
+    The absolute floor (``abs_floor_rel`` x RMS of ``expected``) covers values
+    near zero, where a different FP32 summation order of a cancelling sum
+    moves the result by far more than its own (tiny) ulp.
+    """
+    a = actual.float()
+    e = expected.float()
+    mag = torch.maximum(a.abs(), e.abs()).clamp_min(2.0**-126)
+    ulp = torch.exp2(torch.floor(torch.log2(mag)) - 7)
+    floor = abs_floor_rel * e.pow(2).mean().sqrt().item()
+    excess = ((a - e).abs() - ulps * ulp - floor).max().item()
+    assert excess <= 0, f"difference exceeds {ulps} BF16 ulps by {excess}"
+
+
+# SwiGLU of BF16-rounded gate and up, rounded to BF16: when another FP32
+# summation order flips the rounding of gate and of up by one ulp each, the
+# product moves by up to two ulps, plus one ulp from rounding h.
+SWIGLU_ULPS = 3
+
+
+def _check_against_native_and_oracle(x, gate_up, down, variants=(0, 0)):
+    """Per-stage native comparison plus the end-to-end FP64-oracle budget.
+
+    The down stage (one BF16 rounding) uses the existing small-M BF16 GEMM
+    tolerance against the native down GEMM on the candidate's own SwiGLU
+    output. The gate/up + SwiGLU stage has three rounding points (gate, up,
+    h) and is bounded by SWIGLU_ULPS against native. End to end, the frozen
+    FP64-oracle error budget applies.
+    """
+    import dense_mlp_utils as wu
+
+    h = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up, variants[0])
+    assert_within_bf16_ulps(h, wu.native_intermediate(x, gate_up), SWIGLU_ULPS)
+    out = torch.ops.trtllm.glm_moe_dsa_dense_mlp_down(h, down, variants[1])
+    torch.testing.assert_close(out, F.linear(h, down), rtol=1e-2, atol=5e-3)
+    cand = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
+    assert torch.equal(cand, out)
+    nat = wu.native_dense_mlp(x, gate_up, down)
+    _, ref = wu.oracle_dense_mlp(x, gate_up, down)
+    cand_err = wu.error_stats(cand, ref)["nrms"]
+    nat_err = wu.error_stats(nat, ref)["nrms"]
+    assert cand_err <= NRMS_FACTOR * nat_err + NRMS_SLACK, (cand_err, nat_err)
+    return cand
+
+
+# Contraction multiple (warps x 32 x steps) of each variant in
+# denseMlp.cu.
+VARIANT_K_MULTIPLE = [4 * 32 * 8, 8 * 32 * 4, 8 * 32 * 4, 4 * 32 * 8, 8 * 32 * 8]
+
+# (gate/up variant, down variant) pairs exercised by the shape-generic tests.
+VARIANT_PAIRS = [(0, 0), (1, 0), (2, 1), (3, 3), (4, 3)]
+
+
+def _supported(variants, hidden, inter):
+    return (
+        hidden % VARIANT_K_MULTIPLE[variants[0]] == 0
+        and inter % VARIANT_K_MULTIPLE[variants[1]] == 0
+    )
+
+
+@pytest.mark.parametrize("num_rows", [1, 2, 3, 4, 5, 6, 7, 8])
+@pytest.mark.parametrize("hidden,inter", [(6144, 3072), (1024, 1024), (2048, 4096)])
+@pytest.mark.parametrize("variants", VARIANT_PAIRS)
+def test_matches_native_random(num_rows, hidden, inter, variants):
+    if not _supported(variants, hidden, inter):
+        pytest.skip("tiling does not divide the shape")
+    gate_up, down = _random_weights(hidden, inter)
+    x = torch.randn(num_rows, hidden, device="cuda", dtype=torch.bfloat16)
+    _check_against_native_and_oracle(x, gate_up, down, variants)
+
+
+@pytest.mark.parametrize("num_rows", [1, 6, 8])
+def test_intermediate_rounding_matches_native(num_rows):
+    """gate/up/h rounding follows native; h differs only by summation order."""
+    import dense_mlp_utils as wu
+
+    gate_up, down = _random_weights(6144, 3072, seed=3)
+    x = torch.randn(num_rows, 6144, device="cuda", dtype=torch.bfloat16)
+    h = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up)
+    nat_h = wu.native_intermediate(x, gate_up)
+    ref_h, _ = wu.oracle_dense_mlp(x, gate_up, down)
+    assert_within_bf16_ulps(h, nat_h, SWIGLU_ULPS)
+    # Values that native reproduces exactly from the FP64 oracle must mostly
+    # match exactly here too (only BF16 rounding-boundary cases may differ).
+    assert (h == nat_h).float().mean().item() > 0.97
+    assert wu.error_stats(h, ref_h)["nrms"] <= (
+        NRMS_FACTOR * wu.error_stats(nat_h, ref_h)["nrms"] + NRMS_SLACK
+    )
+
+
+@pytest.mark.parametrize("variant", range(len(VARIANT_K_MULTIPLE)))
+def test_all_variants(variant):
+    gate_up, down = _random_weights(6144, 3072, seed=5)
+    x = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
+    ref = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, 0, 0)
+    h = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up, variant)
+    h0 = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up, 0)
+    assert_within_bf16_ulps(h, h0, SWIGLU_ULPS)
+    if 3072 % VARIANT_K_MULTIPLE[variant] != 0:
+        with pytest.raises(RuntimeError):
+            torch.ops.trtllm.glm_moe_dsa_dense_mlp_down(h0, down, variant)
+        return
+    out = torch.ops.trtllm.glm_moe_dsa_dense_mlp_down(h0, down, variant)
+    torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-3)
+
+
+@pytest.mark.parametrize("variants", VARIANT_PAIRS)
+def test_deterministic(variants):
+    gate_up, down = _random_weights(6144, 3072, seed=7)
+    x = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
+    first = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
+    for _ in range(5):
+        assert torch.equal(
+            torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants), first
+        )
+
+
+@pytest.mark.parametrize("variants", VARIANT_PAIRS)
+def test_cuda_graph_replay_tracks_inputs(variants):
+    gate_up, down = _random_weights(6144, 3072, seed=11)
+    x_static = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_static, gate_up, down, *variants)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out_static = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_static, gate_up, down, *variants)
+    for seed in range(4):
+        torch.manual_seed(100 + seed)
+        new_x = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
+        x_static.copy_(new_x)
+        graph.replay()
+        eager = torch.ops.trtllm.glm_moe_dsa_dense_mlp(new_x, gate_up, down, *variants)
+        torch.cuda.synchronize()
+        assert torch.equal(out_static, eager)
+
+
+def test_rejects_unsupported_inputs():
+    gate_up, down = _random_weights(6144, 3072)
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.glm_moe_dsa_dense_mlp(
+            torch.randn(9, 6144, device="cuda", dtype=torch.bfloat16), gate_up, down
+        )
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.glm_moe_dsa_dense_mlp(
+            torch.randn(6, 6144, device="cuda", dtype=torch.float16), gate_up, down
+        )
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.glm_moe_dsa_dense_mlp(
+            torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16), gate_up, down, 99, 0
+        )
+
+
+def _make_gated_mlp(hidden, inter, tp_size=1, overridden_tp_size=None):
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+    from tensorrt_llm.mapping import Mapping
+
+    mapping = Mapping(world_size=tp_size, rank=0, tp_size=tp_size, gpus_per_node=max(tp_size, 1))
+    mlp = GatedMLP(
+        hidden_size=hidden,
+        intermediate_size=inter,
+        bias=False,
+        dtype=torch.bfloat16,
+        config=ModelConfig(mapping=mapping),
+        overridden_tp_size=overridden_tp_size,
+        reduce_output=tp_size > 1,
+    ).cuda()
+    gate_up, down = _random_weights(hidden, inter // (overridden_tp_size or tp_size), seed=13)
+    with torch.no_grad():
+        mlp.gate_up_proj.weight.copy_(gate_up)
+        mlp.down_proj.weight.copy_(down)
+    return mlp
+
+
+@pytest.mark.parametrize("num_rows", [1, 6, 8])
+def test_forward_matches_gated_mlp_module(num_rows):
+    import dense_mlp_utils as wu
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
+        dense_mlp,
+        dense_mlp_forward,
+    )
+
+    mlp = _make_gated_mlp(2048, 4096)
+    x = torch.randn(num_rows, 2048, device="cuda", dtype=torch.bfloat16)
+    cand = dense_mlp_forward(mlp, x)
+    assert cand is not None
+    # The entry point runs exactly the op with the module's weights and the
+    # package's default tilings; the module's native forward is the
+    # reference the op is checked against.
+    gate_up, down = mlp.gate_up_proj.weight, mlp.down_proj.weight
+    variants = (dense_mlp.GATE_UP_VARIANT, dense_mlp.DOWN_VARIANT)
+    assert torch.equal(cand, _check_against_native_and_oracle(x, gate_up, down, variants))
+    _, ref = wu.oracle_dense_mlp(x, gate_up, down)
+    assert wu.error_stats(cand, ref)["nrms"] <= (
+        NRMS_FACTOR * wu.error_stats(mlp(x), ref)["nrms"] + NRMS_SLACK
+    )
+
+
+def test_eligibility_falls_back_before_device_work():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
+        check_eligibility,
+        dense_mlp_forward,
+    )
+    from tensorrt_llm._torch.distributed import AllReduceParams
+    from tensorrt_llm._torch.utils import Fp4QuantizedTensor
+
+    mlp = _make_gated_mlp(2048, 4096)
+    x = torch.randn(6, 2048, device="cuda", dtype=torch.bfloat16)
+    assert check_eligibility(mlp, x)[0]
+    cases = {
+        "rows": torch.randn(9, 2048, device="cuda", dtype=torch.bfloat16),
+        "dtype": x.float(),
+        "noncontig": torch.randn(2048, 6, device="cuda", dtype=torch.bfloat16).t(),
+        "3d": x.view(1, 6, 2048),
+        "fp4": Fp4QuantizedTensor(x, x),
+    }
+    for name, bad in cases.items():
+        eligible, reason = check_eligibility(mlp, bad)
+        assert not eligible, name
+        assert dense_mlp_forward(mlp, bad) is None, name
+    assert not check_eligibility(mlp, x, lora_params={"x": 1})[0]
+    mlp.swiglu_limit = 7.0
+    assert not check_eligibility(mlp, x)[0]
+    mlp.swiglu_limit = None
+    # Output all-reduce requested on a TP>1 down_proj -> native path. The
+    # mapping is swapped in place so no communicator is needed.
+    from tensorrt_llm.mapping import Mapping
+
+    mlp.down_proj.reduce_output = True
+    mlp.down_proj.mapping = Mapping(world_size=2, rank=0, tp_size=2, gpus_per_node=2)
+    assert not check_eligibility(mlp, x)[0]
+    assert not check_eligibility(mlp, x, AllReduceParams(enable_allreduce=True))[0]
+    assert check_eligibility(mlp, x, AllReduceParams(enable_allreduce=False))[0]
+    mlp.activation = F.gelu
+    assert not check_eligibility(mlp, x, AllReduceParams(enable_allreduce=False))[0]
+
+
+def test_attention_dp_takes_native_path():
+    """Attention DP runs the dense MLP unsharded: outside dense_mlp's TP domain."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
+        check_eligibility,
+        dense_mlp_forward,
+    )
+    from tensorrt_llm.mapping import Mapping
+
+    mlp = _make_gated_mlp(2048, 4096)
+    x = torch.randn(6, 2048, device="cuda", dtype=torch.bfloat16)
+    assert check_eligibility(mlp, x)[0]
+    mlp.mapping = Mapping(
+        world_size=1, rank=0, tp_size=1, gpus_per_node=1, enable_attention_dp=True
+    )
+    assert check_eligibility(mlp, x) == (False, "attention data parallelism")
+    assert dense_mlp_forward(mlp, x) is None
+
+
+def test_registered_with_dispatch():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
+        dense_mlp_forward,
+    )
+
+    assert dispatch.get_impl("dense_mlp", "dense_mlp_forward") is dense_mlp_forward
+
+
+@pytest.mark.skipif(
+    not os.path.isdir(os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4")),
+    reason="GLM-5.2 checkpoint not mounted",
+)
+@pytest.mark.parametrize("layer", [0, 1, 2])
+def test_real_weights_real_hidden(layer):
+    import dense_mlp_utils as wu
+
+    gate_up, down = wu.load_dense_shard(layer, shard=layer % 4)
+    real = wu.load_real_hidden()
+    rows = [real[i : i + 6] for i in (0, 777, 1500)] if real is not None else []
+    rows.append(torch.randn(6, 6144, dtype=torch.bfloat16))
+    for x in rows:
+        for variants in VARIANT_PAIRS:
+            _check_against_native_and_oracle(x.cuda().contiguous(), gate_up, down, variants)
+
+
+def _pdl_delay_module():
+    from torch.utils.cpp_extension import load
+
+    build_dir = os.environ.get("GLM_MOE_DSA_DENSE_MLP_JIT_BUILD_DIR")
+    if build_dir:
+        os.makedirs(build_dir, exist_ok=True)
+    return load(
+        name="dense_mlp_pdl_delay",
+        sources=[os.path.join(os.path.dirname(os.path.abspath(__file__)), "pdl_delay.cu")],
+        extra_cuda_cflags=["-O3", "-gencode=arch=compute_103a,code=sm_103a"]
+        if torch.cuda.get_device_capability() == (10, 3)
+        else ["-O3"],
+        build_directory=build_dir,
+        verbose=False,
+    )
+
+
+@pytest.mark.skipif(
+    os.environ.get("TRTLLM_ENABLE_PDL") != "1",
+    reason="the dense_mlp kernels launch with PDL only when TRTLLM_ENABLE_PDL=1",
+)
+@pytest.mark.parametrize("variants", [(1, 0), (0, 0)])
+@pytest.mark.parametrize("delay_ns", [0, 20_000, 200_000])
+def test_pdl_delayed_producer(variants, delay_ns):
+    """Activations written late by a PDL predecessor are the ones consumed.
+
+    The predecessor releases its dependents immediately and writes x only
+    after a delay; the dense_mlp kernels issue weight loads before
+    griddepcontrol.wait but must read x after it, in eager mode and in a
+    replayed CUDA graph.
+    """
+    mod = _pdl_delay_module()
+    gate_up, down = _random_weights(6144, 3072, seed=17)
+    x_old = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
+    x_new = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
+    expected = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_new, gate_up, down, *variants)
+    stale = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_old, gate_up, down, *variants)
+    assert not torch.equal(expected, stale)
+    x = x_old.clone()
+    mod.delayed_copy(x_new, x, delay_ns)
+    out = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
+    torch.cuda.synchronize()
+    assert torch.equal(out, expected)
+
+    x.copy_(x_old)
+    graph = torch.cuda.CUDAGraph()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        mod.delayed_copy(x_new, x, delay_ns)
+        torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    with torch.cuda.graph(graph):
+        mod.delayed_copy(x_new, x, delay_ns)
+        out_g = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
+    for _ in range(3):
+        x.copy_(x_old)
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(out_g, expected)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py
new file mode 100644
index 000000000000..4fe078ecd6c1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py
@@ -0,0 +1,853 @@
+# 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.
+"""draft_front (the draft_front contract) against the native ``DeepseekV3MTP`` front end.
+
+The oracle runs the native ops of ``DeepseekV3MTP.forward``: TensorRT-LLM
+``RMSNorm`` modules for ``enorm``/``hnorm`` (run in parallel through
+``maybe_execute_in_parallel``), ``torch.concat`` and ``torch.chunk``. The
+candidate must reproduce chunk ``tp_rank`` bit for bit, with the native view's
+shape and strides.
+
+Fixtures:
+* synthetic: random activations and weights (labelled ``synthetic``).
+* real: ``enorm``/``hnorm``/``eh_proj``/``model.norm``/``shared_head.norm``
+  weights and embedding rows read from the GLM-5.2 NVFP4 checkpoint
+  (``GLM52_NVFP4_CKPT``); the hidden states entering the draft layer are
+  synthetic pre-norm activations passed through the real ``model.norm`` (first
+  draft pass) or ``shared_head.norm`` (later passes) weights.
+
+The native-norm-path tests check that draft_front serves a call only when the native
+``enorm``/``hnorm`` run the FlashInfer CuTe DSL RMSNorm that it mirrors, and
+leaves every other native path (gemma, cuda.tile, NVFP4, torch fallback,
+FlashInfer CUDA JIT norm, unverified FlashInfer) to the native front.
+"""
+
+import functools
+import json
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+from utils.llm_data import llm_models_root
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_front import (  # noqa: E501
+    native_norm,
+    native_norm_reason,
+    pack_draft_front,
+    plan_segments,
+    unsupported_reason,
+)
+from tensorrt_llm._torch.modules import rms_norm as rms_norm_module
+from tensorrt_llm._torch.modules.multi_stream_utils import (
+    maybe_execute_in_parallel,
+    with_multi_stream,
+)
+from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+
+GLM_H = 6144
+GLM_EPS = 1e-5
+GLM_TP = 8
+CKPT = os.environ.get(
+    "GLM52_NVFP4_CKPT",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+# Draft passes at K=5 for one request: first pass verifies 6 rows, then 1 row each.
+DRAFT_ROWS = (6, 1)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import plan_testlib  # noqa: E402
+
+
+class NativeDraftFront:
+    """Native ``DeepseekV3MTP`` front end: parallel enorm/hnorm, concat, chunk."""
+
+    def __init__(
+        self,
+        enorm_weight: torch.Tensor,
+        hnorm_weight: torch.Tensor,
+        enorm_eps: float = GLM_EPS,
+        hnorm_eps: float = GLM_EPS,
+    ):
+        hidden_size = enorm_weight.shape[0]
+        self.enorm = RMSNorm(hidden_size=hidden_size, eps=enorm_eps, dtype=torch.bfloat16).cuda()
+        self.hnorm = RMSNorm(hidden_size=hidden_size, eps=hnorm_eps, dtype=torch.bfloat16).cuda()
+        with torch.no_grad():
+            self.enorm.weight.copy_(enorm_weight)
+            self.hnorm.weight.copy_(hnorm_weight)
+        self.aux_stream = torch.cuda.Stream()
+        self.events = (torch.cuda.Event(), torch.cuda.Event())
+
+    def __call__(
+        self, embeds: torch.Tensor, hidden: torch.Tensor, tp_size: int, tp_rank: int
+    ) -> torch.Tensor:
+        e, h = maybe_execute_in_parallel(
+            lambda: self.enorm(embeds),
+            lambda: self.hnorm(hidden),
+            self.events[0],
+            self.events[1],
+            self.aux_stream,
+            disable_on_compile=True,
+        )
+        x = torch.concat([e, h], dim=-1)
+        if tp_size > 1:
+            x = torch.chunk(x, tp_size, dim=-1)[tp_rank]
+        return x
+
+    def candidate(
+        self, embeds: torch.Tensor, hidden: torch.Tensor, tp_size: int, tp_rank: int
+    ) -> torch.Tensor | None:
+        return pack_draft_front(
+            embeds,
+            hidden,
+            self.enorm.weight,
+            self.enorm.variance_epsilon,
+            self.hnorm.weight,
+            self.hnorm.variance_epsilon,
+            tp_rank,
+            tp_size,
+        )
+
+
+def assert_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor, what: str = "") -> None:
+    assert actual is not None, f"{what}: candidate fell back to native"
+    assert actual.shape == expected.shape, (what, actual.shape, expected.shape)
+    assert actual.stride() == expected.stride(), (what, actual.stride(), expected.stride())
+    assert actual.dtype == expected.dtype
+    a = actual.contiguous().view(torch.int16)
+    b = expected.contiguous().view(torch.int16)
+    mismatch = a != b
+    if bool(mismatch.any()):
+        diff = (actual.float() - expected.float()).abs()
+        pytest.fail(
+            f"{what}: {int(mismatch.sum())}/{mismatch.numel()} elements differ, "
+            f"max |diff| {float(diff.max()):.3e}"
+        )
+
+
+def make_inputs(
+    rows: int, hidden_size: int = GLM_H, seed: int = 0, scale_spread: bool = True
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Synthetic BF16 embeddings / hidden states with per-row scale spread and outliers."""
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    e = torch.randn(rows, hidden_size, device="cuda", generator=gen)
+    h = torch.randn(rows, hidden_size, device="cuda", generator=gen)
+    if scale_spread:
+        scales = torch.logspace(-2, 2, steps=max(rows, 1), device="cuda")[:rows, None]
+        e = e * scales.flip(0) * 0.05
+        h = h * scales
+        h[:, 7] *= 50.0  # an outlier channel, as in residual streams
+    return e.bfloat16(), h.bfloat16()
+
+
+def make_weights(hidden_size: int = GLM_H, seed: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    we = (1.0 + 0.5 * torch.randn(hidden_size, device="cuda", generator=gen)).bfloat16()
+    wh = (0.5 + 0.25 * torch.randn(hidden_size, device="cuda", generator=gen)).bfloat16()
+    return we, wh
+
+
+def fp64_reference(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
+    x64 = x.double()
+    return x64 * torch.rsqrt(x64.pow(2).mean(-1, keepdim=True) + eps) * w.double()
+
+
+# --------------------------------------------------------------------------- planning (CPU)
+
+
+@pytest.mark.parametrize("hidden_size", [GLM_H, 7168, 2048])
+@pytest.mark.parametrize("tp_size", [1, 2, 3, 4, 8])
+def test_plan_segments_cover_native_chunks(hidden_size, tp_size):
+    if (2 * hidden_size) % tp_size:
+        pytest.skip("not a native chunking")
+    width = 2 * hidden_size // tp_size
+    for rank in range(tp_size):
+        cols = []
+        for seg in plan_segments(hidden_size, tp_size, rank):
+            assert 0 <= seg.begin < seg.end <= hidden_size
+            cols.extend(
+                range(seg.source * hidden_size + seg.begin, seg.source * hidden_size + seg.end)
+            )
+        assert cols == list(range(rank * width, (rank + 1) * width))
+
+
+def test_plan_segments_glm_tp8_one_source_per_rank():
+    for rank in range(GLM_TP):
+        (seg,) = plan_segments(GLM_H, GLM_TP, rank)
+        assert seg.source == (0 if rank < 4 else 1)
+        assert seg.end - seg.begin == 1536
+
+
+def test_plan_segments_rejects_bad_tp():
+    with pytest.raises(ValueError):
+        plan_segments(GLM_H, 5, 0)
+    with pytest.raises(ValueError):
+        plan_segments(GLM_H, 8, 8)
+
+
+# --------------------------------------------------------------------------- equivalence
+
+
+@requires_cuda
+@pytest.mark.parametrize("rows", [1, 2, 3, 6, 7, 64, 128, 1000])
+@pytest.mark.parametrize("tp_size", [1, 2, 3, 4, 8])
+def test_matches_native_all_ranks_synthetic(rows, tp_size):
+    e, h = make_inputs(rows, seed=rows)
+    native = NativeDraftFront(*make_weights())
+    for rank in range(tp_size):
+        assert_bitwise_equal(
+            native.candidate(e, h, tp_size, rank),
+            native(e, h, tp_size, rank),
+            f"synthetic rows={rows} tp={tp_size} rank={rank}",
+        )
+
+
+@requires_cuda
+@pytest.mark.parametrize("hidden_size", [2048, 4096, 7168])
+@pytest.mark.parametrize("rows", [1, 6])
+def test_matches_native_other_hidden_sizes(hidden_size, rows):
+    e, h = make_inputs(rows, hidden_size, seed=3)
+    native = NativeDraftFront(*make_weights(hidden_size))
+    for rank in range(GLM_TP):
+        assert_bitwise_equal(
+            native.candidate(e, h, GLM_TP, rank),
+            native(e, h, GLM_TP, rank),
+            f"H={hidden_size} rows={rows} rank={rank}",
+        )
+
+
+@requires_cuda
+@pytest.mark.parametrize("enorm_eps, hnorm_eps", [(1e-5, 1e-5), (1e-6, 1e-3), (1e-2, 1e-6)])
+def test_epsilon_per_source(enorm_eps, hnorm_eps):
+    # Small magnitudes so that eps dominates mean(x^2) for some rows.
+    e, h = make_inputs(6, seed=11)
+    e = (e.float() * 1e-3).bfloat16()
+    h = (h.float() * 1e-3).bfloat16()
+    native = NativeDraftFront(*make_weights(), enorm_eps=enorm_eps, hnorm_eps=hnorm_eps)
+    for tp_size in (1, 3, 8):
+        for rank in range(tp_size):
+            assert_bitwise_equal(
+                native.candidate(e, h, tp_size, rank),
+                native(e, h, tp_size, rank),
+                f"eps=({enorm_eps},{hnorm_eps}) tp={tp_size} rank={rank}",
+            )
+
+
+@requires_cuda
+def test_extreme_rows():
+    rows = 8
+    e, h = make_inputs(rows, seed=5, scale_spread=False)
+    e[0].zero_()
+    h[1].zero_()
+    e[2] = 3.0e4
+    h[3] = -2.0e4
+    e[4] = 1.0e-30
+    h[5].copy_(torch.linspace(-1.0e3, 1.0e3, GLM_H, device="cuda").bfloat16())
+    h[6, ::2] = 0.0
+    e[7, 100] = 6.0e4
+    native = NativeDraftFront(*make_weights())
+    for rank in range(GLM_TP):
+        assert_bitwise_equal(
+            native.candidate(e, h, GLM_TP, rank), native(e, h, GLM_TP, rank), f"extreme rank={rank}"
+        )
+
+
+@requires_cuda
+def test_strided_and_gathered_inputs():
+    """Draft-loop operands: row slices of a larger buffer and gathered rows."""
+    native = NativeDraftFront(*make_weights())
+    big_e, big_h = make_inputs(16, seed=9)
+    wide = torch.empty(16, 2 * GLM_H, dtype=torch.bfloat16, device="cuda")
+    wide[:, GLM_H:] = big_h
+    cases = {
+        "row_slice": (big_e[3:9], big_h[3:9]),
+        "wide_row_stride": (big_e[:6], wide[:6, GLM_H:]),
+        "gathered": (
+            big_e[torch.tensor([5, 0, 9], device="cuda")],
+            big_h[torch.tensor([2, 2, 15], device="cuda")],
+        ),
+    }
+    for name, (e, h) in cases.items():
+        for rank in range(GLM_TP):
+            assert_bitwise_equal(
+                native.candidate(e, h, GLM_TP, rank),
+                native(e, h, GLM_TP, rank),
+                f"{name} rank={rank}",
+            )
+
+
+@requires_cuda
+def test_padded_rows_do_not_leak():
+    """Rows past the live count hold stale data; live rows must not depend on them."""
+    native = NativeDraftFront(*make_weights())
+    live = 2
+    e, h = make_inputs(6, seed=21)
+    e2, h2 = e.clone(), h.clone()
+    e2[live:] = 1.0e4
+    h2[live:] = -7.0
+    for rank in range(GLM_TP):
+        out_a = native.candidate(e, h, GLM_TP, rank)[:live].clone()
+        out_b = native.candidate(e2, h2, GLM_TP, rank)
+        assert torch.equal(out_a, out_b[:live])
+        assert_bitwise_equal(out_b, native(e2, h2, GLM_TP, rank), f"padded rank={rank}")
+
+
+@requires_cuda
+@pytest.mark.parametrize("pdl", ["0", "1"])
+def test_cuda_graph_replay(monkeypatch, pdl):
+    """Capture once per rank and draft-row count; replay with fresh inputs."""
+    monkeypatch.setenv("TRTLLM_ENABLE_PDL", pdl)
+    native = NativeDraftFront(*make_weights())
+    for rows in DRAFT_ROWS:
+        for rank in range(GLM_TP):
+            static_e, static_h = make_inputs(rows, seed=100)
+            native.candidate(static_e, static_h, GLM_TP, rank)  # compile outside capture
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                static_out = native.candidate(static_e, static_h, GLM_TP, rank)
+            for step in range(8):
+                e, h = make_inputs(rows, seed=200 + step)
+                static_e.copy_(e)
+                static_h.copy_(h)
+                graph.replay()
+                torch.cuda.synchronize()
+                assert_bitwise_equal(
+                    static_out,
+                    native(static_e, static_h, GLM_TP, rank),
+                    f"graph pdl={pdl} rows={rows} rank={rank} step={step}",
+                )
+
+
+@requires_cuda
+def test_graph_with_native_neighbors():
+    """Embedding lookup -> candidate -> eh_proj GEMM captured in one multi-stream graph."""
+    vocab = 4096
+    gen = torch.Generator(device="cuda").manual_seed(7)
+    table = torch.randn(vocab, GLM_H, device="cuda", generator=gen).bfloat16()
+    w_eh = (torch.randn(GLM_H, 2 * GLM_H, device="cuda", generator=gen) * 0.02).bfloat16()
+    native = NativeDraftFront(*make_weights())
+    width = 2 * GLM_H // GLM_TP
+    for rows in DRAFT_ROWS:
+        ids = torch.randint(0, vocab, (rows,), device="cuda", generator=gen)
+        _, h = make_inputs(rows, seed=31)
+        for rank in range(GLM_TP):
+            w_r = w_eh[:, rank * width : (rank + 1) * width].contiguous()
+
+            def native_step():
+                return F.linear(native(F.embedding(ids, table), h, GLM_TP, rank), w_r)
+
+            def candidate_step():
+                return F.linear(native.candidate(F.embedding(ids, table), h, GLM_TP, rank), w_r)
+
+            expected = native_step()
+            candidate_step()
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with with_multi_stream(True), torch.cuda.graph(graph):
+                out_native = native_step()
+                out_cand = candidate_step()
+            for step in range(4):
+                ids.copy_(torch.randint(0, vocab, (rows,), device="cuda", generator=gen))
+                graph.replay()
+                torch.cuda.synchronize()
+                expected = native_step()
+                assert_bitwise_equal(out_native, expected, f"native-in-graph rank={rank}")
+                assert_bitwise_equal(out_cand, expected, f"neighbors rows={rows} rank={rank}")
+
+
+@requires_cuda
+def test_capture_before_compile_falls_back():
+    """A configuration first seen inside CUDA-graph capture takes the native path."""
+    hidden_size, rank = 1024, 5
+    e, h = make_inputs(6, hidden_size, seed=55)
+    native = NativeDraftFront(*make_weights(hidden_size))
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        assert native.candidate(e, h, GLM_TP, rank) is None
+    assert_bitwise_equal(
+        native.candidate(e, h, GLM_TP, rank),
+        native(e, h, GLM_TP, rank),
+        "eager after capture fallback",
+    )
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        static_out = native.candidate(e, h, GLM_TP, rank)
+    assert static_out is not None
+    graph.replay()
+    torch.cuda.synchronize()
+    assert_bitwise_equal(static_out, native(e, h, GLM_TP, rank), "capture after compile")
+
+
+# --------------------------------------------------------------------------- served neighbours
+
+
+def make_eh_shard(tp_size: int = GLM_TP, seed: int = 9) -> torch.Tensor:
+    """A ``[H, 2H / tp]`` BF16 stand-in for the local ROW-parallel ``eh_proj`` weight shard."""
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    return (
+        torch.randn(GLM_H, 2 * GLM_H // tp_size, device="cuda", generator=gen) * 0.02
+    ).bfloat16()
+
+
+@requires_cuda
+@pytest.mark.parametrize("pdl", ["0", "1"])
+def test_graph_with_served_neighbors(monkeypatch, pdl):
+    """LM-head-sized L2 eviction -> embedding -> front -> eh_proj GEMM -> input RMSNorm."""
+    monkeypatch.setenv("TRTLLM_ENABLE_PDL", pdl)
+    vocab = 4096
+    gen = torch.Generator(device="cuda").manual_seed(17)
+    table = torch.randn(vocab, GLM_H, device="cuda", generator=gen).bfloat16()
+    flush_w = torch.randn(19360, GLM_H, device="cuda", generator=gen).bfloat16()
+    flush_x = torch.randn(1, GLM_H, device="cuda", generator=gen).bfloat16()
+    native = NativeDraftFront(*make_weights())
+    w_eh = make_eh_shard()
+    input_norm = RMSNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16).cuda()
+    for rows in DRAFT_ROWS:
+        ids = torch.randint(0, vocab, (rows,), device="cuda", generator=gen)
+        _, h = make_inputs(rows, seed=33)
+        for rank in (0, 5):
+
+            def native_step():
+                F.linear(flush_x, flush_w)
+                return input_norm(F.linear(native(F.embedding(ids, table), h, GLM_TP, rank), w_eh))
+
+            def candidate_step():
+                F.linear(flush_x, flush_w)
+                front = native.candidate(F.embedding(ids, table), h, GLM_TP, rank)
+                return input_norm(F.linear(front, w_eh))
+
+            native_step()
+            candidate_step()
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with with_multi_stream(True), torch.cuda.graph(graph):
+                out_native = native_step()
+                out_cand = candidate_step()
+            for step in range(4):
+                ids.copy_(torch.randint(0, vocab, (rows,), device="cuda", generator=gen))
+                graph.replay()
+                torch.cuda.synchronize()
+                assert_bitwise_equal(
+                    out_cand, out_native, f"served pdl={pdl} rows={rows} rank={rank}"
+                )
+
+
+# --------------------------------------------------------------------------- fallback
+
+
+@requires_cuda
+def test_unsupported_operands_take_native_path():
+    e, h = make_inputs(4)
+    we, wh = make_weights()
+    flat = torch.empty(4 * GLM_H + 8, dtype=torch.bfloat16, device="cuda")
+    misaligned = flat[1 : 1 + 4 * GLM_H].view(4, GLM_H)
+    wide = torch.empty(4, 2 * GLM_H, dtype=torch.bfloat16, device="cuda")
+    cases = {
+        "fp32": (e.float(), h.float(), we.float(), wh.float(), 8, 0),
+        "fp16_mixed": (e.half(), h, we, wh, 8, 0),
+        "misaligned": (misaligned, h, we, wh, 8, 0),
+        "column_stride": (e, wide[:, ::2], we, wh, 8, 0),
+        "rows_mismatch": (e, h[:3], we, wh, 8, 0),
+        "tp_not_dividing": (e, h, we, wh, 5, 0),
+        "rank_out_of_range": (e, h, we, wh, 8, 8),
+        "weight_shape": (e, h, we[:-8], wh, 8, 0),
+        "cpu": (e.cpu(), h.cpu(), we.cpu(), wh.cpu(), 8, 0),
+    }
+    for name, (ce, ch, cwe, cwh, tp_size, rank) in cases.items():
+        assert unsupported_reason(ce, ch, cwe, cwh, tp_size, rank) is not None, name
+        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
+            out = pack_draft_front(ce, ch, cwe, GLM_EPS, cwh, GLM_EPS, rank, tp_size)
+            torch.cuda.synchronize()
+        assert out is None, name
+        kernels = [
+            ev.name for ev in prof.events() if ev.device_type == torch.autograd.DeviceType.CUDA
+        ]
+        assert not kernels, (name, kernels)
+
+
+@requires_cuda
+def test_empty_rows():
+    e, h = make_inputs(1)
+    we, wh = make_weights()
+    out = pack_draft_front(e[:0], h[:0], we, GLM_EPS, wh, GLM_EPS, 3, GLM_TP)
+    assert out is not None and out.shape == (0, 2 * GLM_H // GLM_TP)
+
+
+# --------------------------------------------------------------------------- native norm path
+
+
+@pytest.fixture
+def fresh_norm_path():
+    """Recompute the per-device FlashInfer path decision inside the test and after it."""
+    native_norm._flashinfer_path_reason.cache_clear()
+    yield
+    native_norm._flashinfer_path_reason.cache_clear()
+
+
+def glm_norm(**kwargs) -> RMSNorm:
+    """``enorm``/``hnorm`` as ``DeepseekV3MTP`` builds them, on the CPU."""
+    return RMSNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16, **kwargs)
+
+
+def _with_attr(name: str, value) -> RMSNorm:
+    norm = glm_norm()
+    setattr(norm, name, value)
+    return norm
+
+
+class _OverridingNorm(RMSNorm):
+    def forward(self, hidden_states, *args, **kwargs):
+        return super().forward(hidden_states, *args, **kwargs)
+
+
+class _InheritingNorm(RMSNorm):
+    pass
+
+
+def _module_cases() -> dict:
+    """Norm modules whose native forward is not the plain FlashInfer RMSNorm."""
+    return {
+        "gemma": glm_norm(use_gemma=True),
+        "cuda_tile": _with_attr("use_cuda_tile", True),
+        "nvfp4": _with_attr("is_nvfp4", True),
+        "hp_output": _with_attr("return_hp_output", True),
+        "forward_override": _OverridingNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16),
+        "foreign_module": torch.nn.RMSNorm(GLM_H, eps=GLM_EPS),
+    }
+
+
+def test_native_norm_module_paths_decline():
+    plain = glm_norm()
+    assert native_norm.norm_module_reason("enorm", plain) is None
+    inheriting = _InheritingNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16)
+    assert native_norm.norm_module_reason("hnorm", inheriting) is None
+    for case, norm in _module_cases().items():
+        assert native_norm.norm_module_reason("enorm", norm) is not None, case
+        reason = native_norm_reason(norm, plain)
+        assert reason is not None and reason.startswith("enorm"), (case, reason)
+        reason = native_norm_reason(plain, norm)
+        assert reason is not None and reason.startswith("hnorm"), (case, reason)
+
+
+@requires_cuda
+def test_native_norm_glm_config_takes_mirrored_path(monkeypatch, fresh_norm_path):
+    """GLM's enorm/hnorm reach FlashInfer's CuTe DSL RMSNorm, so draft_front serves and is bitwise."""
+    import flashinfer.norm as flashinfer_norm
+
+    native = NativeDraftFront(*make_weights())
+    assert native_norm_reason(native.enorm, native.hnorm) is None
+    calls = {"cute": [], "cuda_jit": 0}
+    cute = flashinfer_norm.rmsnorm_cute
+    jit = flashinfer_norm.get_norm_module
+
+    def cute_spy(input, *args, **kwargs):
+        calls["cute"].append(input.dim())
+        return cute(input, *args, **kwargs)
+
+    def jit_spy():
+        calls["cuda_jit"] += 1
+        return jit()
+
+    monkeypatch.setattr(flashinfer_norm, "rmsnorm_cute", cute_spy)
+    monkeypatch.setattr(flashinfer_norm, "get_norm_module", jit_spy)
+    for rows in DRAFT_ROWS:
+        e, h = make_inputs(rows, seed=rows + 70)
+        assert unsupported_reason(e, h, native.enorm.weight, native.hnorm.weight, GLM_TP, 0) is None
+        for rank in (0, 4):
+            calls["cute"].clear()
+            expected = native(e, h, GLM_TP, rank)
+            assert calls["cute"] == [2, 2], (rows, rank, calls)
+            assert_bitwise_equal(
+                native.candidate(e, h, GLM_TP, rank),
+                expected,
+                f"mirrored path rows={rows} rank={rank}",
+            )
+    assert calls["cuda_jit"] == 0
+
+
+@requires_cuda
+def test_gemma_native_differs_from_candidate():
+    """A gemma enorm computes x * rstd * (w + 1): draft_front would change the result."""
+    we, wh = make_weights()
+    native = NativeDraftFront(we, wh)
+    gemma = glm_norm(use_gemma=True).cuda()
+    with torch.no_grad():
+        gemma.weight.copy_(we)
+    e, h = make_inputs(6, seed=81)
+    rank = 0  # an embedding-source rank
+    width = 2 * GLM_H // GLM_TP
+    native_gemma = gemma(e)[:, :width]
+    candidate = native.candidate(e, h, GLM_TP, rank)
+    assert not torch.equal(candidate, native_gemma)
+    assert native_norm_reason(gemma, native.hnorm) is not None
+
+
+def _process_case(monkeypatch, case: str) -> None:
+    import flashinfer
+    import flashinfer.norm as flashinfer_norm
+
+    if case == "flashinfer_unavailable":
+        monkeypatch.setattr(rms_norm_module, "IS_FLASHINFER_AVAILABLE", False)
+    elif case == "cuda_jit_norm":
+        monkeypatch.setattr(flashinfer_norm, "_use_cuda_norm", lambda device: True)
+    elif case == "unverified_flashinfer":
+        monkeypatch.setattr(flashinfer, "__version__", "0.0.0")
+    elif case == "no_norm_dispatch":
+        monkeypatch.delattr(flashinfer_norm, "_use_cuda_norm")
+    else:
+        raise AssertionError(case)
+
+
+@requires_cuda
+@pytest.mark.parametrize(
+    "case", ["flashinfer_unavailable", "cuda_jit_norm", "unverified_flashinfer", "no_norm_dispatch"]
+)
+def test_native_norm_process_paths_fall_back(monkeypatch, fresh_norm_path, case):
+    """Per-process conditions under which native RMSNorm leaves FlashInfer's CuTe kernel."""
+    native = NativeDraftFront(*make_weights())
+    e, h = make_inputs(6, seed=91)
+    rank = 4  # a hidden-source rank
+    candidate = native.candidate(e, h, GLM_TP, rank)
+    assert candidate is not None
+    _process_case(monkeypatch, case)
+    native_norm._flashinfer_path_reason.cache_clear()
+    assert native_norm.flashinfer_path_reason(e.device) is not None, case
+    assert native_norm_reason(native.enorm, native.hnorm) is not None, case
+    weights = (native.enorm.weight, native.hnorm.weight)
+    assert unsupported_reason(e, h, *weights, GLM_TP, rank) is not None, case
+    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
+        out = pack_draft_front(e, h, weights[0], GLM_EPS, weights[1], GLM_EPS, rank, GLM_TP)
+        torch.cuda.synchronize()
+    assert out is None, case
+    kernels = [ev.name for ev in prof.events() if ev.device_type == torch.autograd.DeviceType.CUDA]
+    assert not kernels, (case, kernels)
+    if case == "flashinfer_unavailable":
+        # The torch path rounds x * rstd to BF16 before the weight multiply.
+        torch_path = native(e, h, GLM_TP, rank)
+        assert not torch.equal(torch_path, candidate)
+
+
+@requires_cuda
+def test_adapter_native_norm_gate(monkeypatch, fresh_norm_path):
+    """The draft layer's adoption binds draft_front only when enorm/hnorm take the mirrored path.
+
+    Builds the built-in ``DeepseekV3MTP`` with the decoder-layer base, ``eh_proj`` and the
+    shared head stubbed out, then runs the draft step of its adoption as ``GlmMtpLayer``.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+        reclass,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmMtpLayer,
+    )
+    from tensorrt_llm._torch.models import modeling_deepseekv3 as model_module
+
+    plan_testlib.select(dispatch, "draft_front", monkeypatch)
+    monkeypatch.setattr(
+        model_module.DeepseekV3DecoderLayer,
+        "__init__",
+        lambda self, *args, **kwargs: torch.nn.Module.__init__(self),
+    )
+    monkeypatch.setattr(model_module, "Linear", lambda *args, **kwargs: torch.nn.Identity())
+    monkeypatch.setattr(
+        model_module, "DeepseekV3MTPHead", lambda *args, **kwargs: torch.nn.Identity()
+    )
+    aux_streams = {model_module.AuxStreamType.MoeShared: torch.cuda.Stream()}
+
+    def build(norm_factory=RMSNorm, attention_dp: bool = False, tp_size: int = GLM_TP):
+        monkeypatch.setattr(model_module, "RMSNorm", norm_factory)
+        config = SimpleNamespace(
+            pretrained_config=SimpleNamespace(
+                hidden_size=GLM_H,
+                moe_intermediate_size=2048,
+                n_routed_experts=256,
+                n_shared_experts=1,
+                num_experts_per_tok=8,
+                rms_norm_eps=GLM_EPS,
+                torch_dtype=torch.bfloat16,
+            ),
+            mapping=SimpleNamespace(enable_attention_dp=attention_dp, tp_size=tp_size, tp_rank=0),
+            skip_create_weights_in_init=True,
+            use_cute_dsl_blockscaling_mm=False,
+        )
+        layer = model_module.DeepseekV3MTP(config, 78, aux_streams)
+        # The stubbed decoder-layer constructor does not keep the model config.
+        layer.model_config = config
+        reclass(layer, model_module.DeepseekV3MTP, GlmMtpLayer)
+        layer._adopt_draft()
+        return layer.draft_front
+
+    def factory(name, value):
+        def make(**kwargs):
+            norm = RMSNorm(**kwargs)
+            setattr(norm, name, value)
+            return norm
+
+        return make
+
+    assert build() is pack_draft_front
+    assert build(functools.partial(RMSNorm, use_gemma=True)) is None
+    assert build(factory("use_cuda_tile", True)) is None
+    assert build(factory("is_nvfp4", True)) is None
+    assert build(attention_dp=True) is None
+    assert build(tp_size=1) is None
+    monkeypatch.setattr(rms_norm_module, "IS_FLASHINFER_AVAILABLE", False)
+    native_norm._flashinfer_path_reason.cache_clear()
+    assert build() is None
+
+
+def test_dispatch_registration(monkeypatch):
+    plan_testlib.select(dispatch, "draft_front", monkeypatch)
+    assert dispatch.is_enabled("draft_front")
+    assert dispatch.get_impl("draft_front", "draft_front") is pack_draft_front
+    assert dispatch.get_impl("draft_front", "unsupported_reason") is unsupported_reason
+    assert dispatch.get_impl("draft_front", "native_norm_reason") is native_norm_reason
+    # draft_front is the draft_front contract only: nothing that changes the eh_proj reduction
+    # or the draft layer's numerics is registered under this path.
+    with pytest.raises(KeyError):
+        dispatch.get_impl("draft_front", "input_norm_fold")
+
+
+def test_dispatch_default_off(monkeypatch):
+    """Without an installed plan (a process that does not build the target) draft_front is off."""
+    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
+    assert not dispatch.is_enabled("draft_front")
+
+
+# --------------------------------------------------------------------------- real checkpoint
+
+
+@functools.cache
+def _real_fixture() -> dict | None:
+    index_path = os.path.join(CKPT, "model.safetensors.index.json")
+    if not os.path.isfile(index_path):
+        return None
+    from safetensors import safe_open
+
+    with open(index_path) as f:
+        weight_map = json.load(f)["weight_map"]
+
+    def load(name: str, rows: list[int] | None = None) -> torch.Tensor:
+        with safe_open(os.path.join(CKPT, weight_map[name]), framework="pt", device="cpu") as f:
+            if rows is None:
+                return f.get_tensor(name)
+            sl = f.get_slice(name)
+            return torch.cat([sl[r : r + 1] for r in rows])
+
+    token_ids = [
+        0,
+        1,
+        2,
+        3,
+        13,
+        198,
+        271,
+        1000,
+        9999,
+        12345,
+        50000,
+        77777,
+        100000,
+        151329,
+        154000,
+        154879,
+    ]
+    return {
+        "enorm": load("model.layers.78.enorm.weight").cuda(),
+        "hnorm": load("model.layers.78.hnorm.weight").cuda(),
+        "eh_proj": load("model.layers.78.eh_proj.weight").cuda(),
+        "model_norm": load("model.norm.weight").cuda(),
+        "shared_head_norm": load("model.layers.78.shared_head.norm.weight").cuda(),
+        "embeds": load("model.embed_tokens.weight", token_ids).cuda(),
+        "token_ids": token_ids,
+    }
+
+
+def _real_hidden(rows: int, norm_weight: torch.Tensor, seed: int) -> torch.Tensor:
+    """Synthetic pre-norm residual activations through a real output RMSNorm weight."""
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    pre = torch.randn(rows, GLM_H, device="cuda", generator=gen) * 3.0
+    pre[:, 11] *= 40.0
+    pre = pre.bfloat16()
+    norm = RMSNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16).cuda()
+    with torch.no_grad():
+        norm.weight.copy_(norm_weight)
+    return norm(pre)
+
+
+@requires_cuda
+@pytest.mark.parametrize("draft_pass", [0, 1])
+def test_real_checkpoint_all_ranks(draft_pass):
+    fx = _real_fixture()
+    if fx is None:
+        pytest.skip(f"checkpoint not found at {CKPT}")
+    rows = DRAFT_ROWS[0] if draft_pass == 0 else DRAFT_ROWS[1]
+    e = fx["embeds"][draft_pass : draft_pass + rows]
+    h = _real_hidden(
+        rows, fx["model_norm"] if draft_pass == 0 else fx["shared_head_norm"], seed=40 + draft_pass
+    )
+    native = NativeDraftFront(fx["enorm"], fx["hnorm"])
+    width = 2 * GLM_H // GLM_TP
+    full_ref = torch.cat(
+        [fp64_reference(e, fx["enorm"], GLM_EPS), fp64_reference(h, fx["hnorm"], GLM_EPS)], dim=-1
+    )
+    for rank in range(GLM_TP):
+        out_native = native(e, h, GLM_TP, rank)
+        out_cand = native.candidate(e, h, GLM_TP, rank)
+        assert_bitwise_equal(out_cand, out_native, f"real pass={draft_pass} rank={rank}")
+        ref = full_ref[:, rank * width : (rank + 1) * width]
+        for label, out in (("native", out_native), ("candidate", out_cand)):
+            err = out.double() - ref
+            rel_rms = float(err.pow(2).mean().sqrt() / ref.pow(2).mean().sqrt())
+            print(
+                f"GLM_MOE_DSA_DRAFT_FRONT_STAT real pass={draft_pass} rank={rank} {label} "
+                f"max_abs={float(err.abs().max()):.3e} rel_rms={rel_rms:.3e}"
+            )
+        # eh_proj consumer: the local ROW-parallel GEMM sees an identical operand.
+        w_r = fx["eh_proj"][:, rank * width : (rank + 1) * width].contiguous()
+        assert_bitwise_equal(
+            F.linear(out_cand, w_r),
+            F.linear(out_native, w_r),
+            f"real eh_proj partial pass={draft_pass} rank={rank}",
+        )
+
+
+@requires_cuda
+def test_real_checkpoint_draft_loop_rows():
+    """All 16 real embedding rows through one first pass and later single-row passes."""
+    fx = _real_fixture()
+    if fx is None:
+        pytest.skip(f"checkpoint not found at {CKPT}")
+    native = NativeDraftFront(fx["enorm"], fx["hnorm"])
+    e_all = fx["embeds"]
+    h_all = _real_hidden(e_all.shape[0], fx["shared_head_norm"], seed=77)
+    for start in range(0, e_all.shape[0]):
+        for rows in (1, 6):
+            if start + rows > e_all.shape[0]:
+                continue
+            e, h = e_all[start : start + rows], h_all[start : start + rows]
+            for rank in range(GLM_TP):
+                assert_bitwise_equal(
+                    native.candidate(e, h, GLM_TP, rank),
+                    native(e, h, GLM_TP, rank),
+                    f"real loop start={start} rows={rows} rank={rank}",
+                )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py
new file mode 100644
index 000000000000..530dd3a0785e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py
@@ -0,0 +1,472 @@
+# 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.
+"""The draft_live_rows call sites in the target and their host-side decisions.
+
+* the plan runs draft_live_rows with its closure (post_attention_hop, draft_loop), next to
+  the served GEN set without deferred_input_prep, and turns it off with either base; draft_live_rows has
+  one numerical form; the package reads no environment variable and
+  imports no other package (no deferred_input_prep side-selector or package-directory entry);
+* the call sites: ``GlmMtpLayer`` adoption builds the live-row hop only when draft_live_rows
+  is selected and its forward tries it only on an armed call that the post_attention_hop
+  serves, before the post_attention_hop; the draft loop (``GlmMtpWorker``) decides once per
+  forward, arms step 0 only, and makes draft_loop's step-0 gather the identity only when the
+  layer compacted;
+* the host-side eligibility of the live-row hop and of the step-0 arming,
+  case by case, with no device work; the arming is cleared even when the
+  step raises;
+* the live-row hop is built on every rank of the post_attention_hop's TP group or on none:
+  whether the operator is loaded is agreed over the group's communicator
+  (mocked here; one collective per post_attention_hop, also on a rank without the
+  operator), with a negative control showing that the rank-local rule
+  disagrees whenever one rank lacks the operator.
+"""
+
+import ast
+import inspect
+import os
+import re
+import sys
+import threading
+from pathlib import Path
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import plan_testlib  # noqa: E402
+
+GEN_FAST_PATHS_WITHOUT_INPUT_PREP = (
+    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,"
+    "router_gemm,post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,"
+    "rope_kv_append,draft_loop,shared_indexer_absorbed_query,routed_moe,"
+    "cross_rank_topk,attention_output_tail_prefetch,attention_output_into_hop,"
+    "routed_moe_into_hop,deferred_norm,dense_mlp,hop_routing,oproj_weight_staging,"
+    "fused_kv_indexer_projection,host_page_retention,shared_expert_prefetch,"
+    "draft_routed_moe,cross_layer_expert_prefetch"
+)
+
+
+def _dispatch():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    return dispatch
+
+
+def test_selector_facts(monkeypatch):
+    d = _dispatch()
+    monkeypatch.setattr(d, "_log_info", lambda message: None)
+    assert "draft_live_rows" in d.SELECTOR_PATHS
+    assert set(d.REQUIRED_FEATURES["draft_live_rows"]) == {"post_attention_hop", "draft_loop"}
+    assert "draft_live_rows" not in d.NUMERICAL_VARIANTS
+    assert plan_testlib.select(
+        d, "post_attention_hop,draft_loop,draft_live_rows", monkeypatch
+    ) == frozenset({"post_attention_hop", "draft_loop", "draft_live_rows"})
+    no_input_prep = frozenset(GEN_FAST_PATHS_WITHOUT_INPUT_PREP.split(","))
+    assert plan_testlib.select(
+        d, GEN_FAST_PATHS_WITHOUT_INPUT_PREP + ",draft_live_rows", monkeypatch
+    ) == no_input_prep | {"draft_live_rows"}
+    assert plan_testlib.select(
+        d,
+        GEN_FAST_PATHS_WITHOUT_INPUT_PREP + ",cross_rank_topk_cluster_merge,draft_live_rows",
+        monkeypatch,
+    ) >= {
+        "cross_rank_topk_cluster_merge",
+        "draft_live_rows",
+    }
+    for base in ("post_attention_hop", "draft_loop"):
+        plan = d.resolve(d.ModelingV2Mode.AUTO, {d.DISABLE_ENV: base})
+        assert not plan.enabled("draft_live_rows"), base
+
+
+def _package_dir() -> Path:
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
+        draft_live_rows,
+    )
+
+    return Path(draft_live_rows.__file__).parent
+
+
+def test_package_is_self_contained():
+    """No side selector and no import of another package's code."""
+    for path in sorted(_package_dir().glob("*.py")):
+        src = path.read_text()
+        assert "environ" not in src and "getenv" not in src, path.name
+        for node in ast.walk(ast.parse(src)):
+            if isinstance(node, ast.ImportFrom) and node.level >= 2:
+                assert node.module == "dispatch", (path.name, node.module)
+
+
+def test_mtp_call_site():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmMtpLayer,
+    )
+
+    init = inspect.getsource(GlmMtpLayer._adopt_draft)
+    assert 'is_enabled("draft_live_rows")' in init and re.search(
+        r'get_impl\(\s*"draft_live_rows",\s*"make_live_row_hop"\s*\)', init
+    )
+    fwd = inspect.getsource(GlmMtpLayer.forward)
+    draft_live_rows = fwd.index("self.draft_live_row_hop.can_run(")
+    # The post_attention_hop decision is made once, before either hop; the live-row hop is tried
+    # first and only when armed.
+    assert (
+        fwd.index("use_post_attn_hop = ") < draft_live_rows < fwd.index("elif use_post_attn_hop:")
+    )
+    assert fwd.index("live_rows is not None") < draft_live_rows
+    assert fwd.count("self._can_use_post_attn_hop(") == 1
+    assert "self._draft_live_rows_compacted = True" in fwd
+
+
+def test_draft_loop_call_site():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
+        GlmMtpWorker,
+    )
+
+    src = inspect.getsource(GlmMtpWorker._forward_linear_draft_loop)
+    match = re.search(r'get_impl\(\s*"draft_live_rows",\s*"select_step0"\s*\)', src)
+    assert match is not None
+    sel = match.start()
+    assert sel < src.index("for i in range(runtime_draft_len):")
+    assert "if i == 0 and live_rows_step0 is not None:" in src
+    ident = src.index("if live_rows_compacted:")
+    assert ident < src.index("draft_loop_ops.step0_rows(hidden_states)")
+    assert "gather_ids, draft_rows = draft_loop_ops.gather_ids, hidden_states" in src
+
+
+class _AttnHopStub:
+    enabled = True
+    rank = 2
+    uc_buffer = None
+    mc_ptr = 1
+    flags = None
+    min_sm_count = 148
+
+
+def test_live_row_hop_eligibility():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows.hop import (  # noqa: E501
+        LiveRowHop,
+        required_sms,
+    )
+
+    hop = LiveRowHop(_AttnHopStub())
+    h6 = torch.zeros(6, 6144, dtype=torch.bfloat16)
+    w = torch.zeros(256, 6144, dtype=torch.bfloat16)
+    ids = torch.zeros(1, dtype=torch.int64)
+    assert hop.call_ineligibility(h6, ids, w) is None
+    assert hop.call_ineligibility(h6, torch.zeros(2, dtype=torch.int64), w) is None
+    assert hop.call_ineligibility(h6, ids, None) == "no router weight"
+    assert hop.call_ineligibility(h6, None, w) == "no live rows"
+    assert hop.call_ineligibility(h6, ids.int(), w) == "live-row ids layout"
+    assert hop.call_ineligibility(h6, ids.view(1, 1), w) == "live-row ids layout"
+    assert (
+        hop.call_ineligibility(h6, torch.zeros(4, dtype=torch.int64)[::2], w)
+        == "live-row ids layout"
+    )
+    assert hop.call_ineligibility(h6, torch.zeros(7, dtype=torch.int64), w).startswith(
+        "live rows 7 of 6"
+    )
+    assert hop.call_ineligibility(h6, torch.zeros(0, dtype=torch.int64), w).startswith(
+        "live rows 0"
+    )
+    h12 = torch.zeros(12, 6144, dtype=torch.bfloat16)
+    assert hop.call_ineligibility(h12, ids, w).startswith("live rows 1 of 12")
+    small = _AttnHopStub()
+    small.min_sm_count = 24
+    assert LiveRowHop(small).call_ineligibility(h6, ids, w) == "SM count"
+    assert [required_sms(r) for r in (1, 2, 3, 8)] == [32, 64, 96, 96]
+
+
+class _Comm:
+    """Mocked TP communicator (MPI style): this rank's value plus fixed peer values."""
+
+    def __init__(self, peers):
+        self.peers = [int(v) for v in peers]
+        self.calls = []
+
+    def allreduce(self, value, op=None):
+        self.calls.append(int(value))
+        return int(value) + sum(self.peers)
+
+    def Get_size(self):
+        return len(self.peers) + 1
+
+
+class _FakeAttnHop(_AttnHopStub):
+    def __init__(self, comm, rank=2):
+        self.comm = comm
+        self.rank = rank
+
+
+@pytest.fixture
+def live_row_hop(monkeypatch):
+    """The hop module with an empty decision table and the MPI agreement path."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows import (  # noqa: E501
+        hop,
+    )
+
+    monkeypatch.setattr(hop.LiveRowHop, "_instances", {})
+    monkeypatch.delenv("TLLM_DISABLE_MPI", raising=False)
+    return hop
+
+
+@pytest.mark.parametrize(
+    "local, peers, built, reason",
+    [
+        (True, [1] * 7, True, ""),
+        (True, [1] * 6 + [0], False, "another rank lacks trtllm::glm_moe_dsa_draft_live_row_hop"),
+        (
+            True,
+            [1] * 3 + [0] * 4,
+            False,
+            "another rank lacks trtllm::glm_moe_dsa_draft_live_row_hop",
+        ),
+        (False, [1] * 7, False, "trtllm::glm_moe_dsa_draft_live_row_hop is not loaded"),
+        (False, [0] * 7, False, "trtllm::glm_moe_dsa_draft_live_row_hop is not loaded"),
+    ],
+)
+def test_live_row_hop_agreed_over_tp_group(live_row_hop, monkeypatch, local, peers, built, reason):
+    monkeypatch.setattr(live_row_hop, "op_available", lambda: local)
+    comm = _Comm(peers)
+    w2 = _FakeAttnHop(comm)
+    hop = live_row_hop.LiveRowHop.for_post_attn_hop(w2)
+    assert (hop is not None) == built
+    # This rank takes part in the agreement whatever its own answer: a rank that
+    # skipped it would leave its peers waiting in the collective.
+    assert comm.calls == [int(local)]
+    assert live_row_hop.LiveRowHop._instances[id(w2)] == (w2, hop, reason)
+    if built:
+        assert hop.w2 is w2 and hop.rank == w2.rank
+    # Later calls for the same post_attention_hop return the recorded decision without
+    # communicating.
+    assert live_row_hop.LiveRowHop.for_post_attn_hop(w2) is hop
+    assert comm.calls == [int(local)]
+
+
+def test_live_row_hop_without_attn_hop_group(live_row_hop, monkeypatch):
+    """No enabled post_attention_hop (identical on all ranks) or no communicator: no hop, no collective."""
+    monkeypatch.setattr(live_row_hop, "op_available", lambda: True)
+    assert live_row_hop.LiveRowHop.for_post_attn_hop(None) is None
+    comm = _Comm([1] * 7)
+    off = _FakeAttnHop(comm)
+    off.enabled = False
+    assert live_row_hop.LiveRowHop.for_post_attn_hop(off) is None
+    lone = _FakeAttnHop(None)
+    assert live_row_hop.LiveRowHop.for_post_attn_hop(lone) is None
+    assert (
+        live_row_hop.LiveRowHop._instances[id(lone)][2]
+        == "post_attention_hop hop has no communicator"
+    )
+    assert comm.calls == []
+
+
+class _ThreadComm:
+    """One in-process TP group of ``size`` threads; a single integer-sum allreduce."""
+
+    def __init__(self, size):
+        self.size = size
+        self.total = 0
+        self.lock = threading.Lock()
+        self.barrier = threading.Barrier(size, timeout=60)
+
+    def allreduce(self, value, op=None):
+        with self.lock:
+            self.total += int(value)
+        self.barrier.wait()
+        return self.total
+
+    def Get_size(self):
+        return self.size
+
+
+@pytest.mark.parametrize("lacking", [(), (0,), (3,), (5,), (7,), (4, 5, 6, 7), tuple(range(8))])
+def test_live_row_hop_rank_agreement_eight_ranks(live_row_hop, monkeypatch, lacking):
+    """Eight ranks (threads) build the hop concurrently: every rank gets it, or none does."""
+    size = 8
+    tls = threading.local()
+    monkeypatch.setattr(live_row_hop, "op_available", lambda: tls.has_op)
+    comm = _ThreadComm(size)
+    built, local_rule, errors = [None] * size, [None] * size, []
+
+    def rank_main(r):
+        tls.has_op = r not in lacking
+        try:
+            built[r] = (
+                live_row_hop.LiveRowHop.for_post_attn_hop(_FakeAttnHop(comm, rank=r)) is not None
+            )
+            # The construction rule before the agreement: this rank's own operator only.
+            local_rule[r] = live_row_hop.op_available()
+        except Exception as error:  # noqa: BLE001  (reported below with the rank)
+            errors.append((r, repr(error)))
+
+    threads = [threading.Thread(target=rank_main, args=(r,)) for r in range(size)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=120)
+    assert not any(t.is_alive() for t in threads) and not errors, errors
+    assert built == [not lacking] * size
+    # Negative control: the rank-local rule splits the group whenever some but not all ranks lack the op.
+    split = 0 < len(lacking) < size
+    assert (len(set(local_rule)) == 2) == split
+
+
+class _Layer:
+    """The draft layer's live-row state and its arming API (``GlmMtpLayer``)."""
+
+    def __init__(self, hop):
+        self.draft_live_row_hop = hop
+        self._draft_live_rows = None
+        self._draft_live_rows_compacted = False
+
+    def arm_live_rows(self, rows):
+        self._draft_live_rows = rows
+        self._draft_live_rows_compacted = False
+
+    def disarm_live_rows(self):
+        compacted = bool(self._draft_live_rows_compacted)
+        self._draft_live_rows = None
+        self._draft_live_rows_compacted = False
+        return compacted
+
+
+class _Draft:
+    def __init__(self, hop):
+        self.mtp_layers = [_Layer(hop)]
+
+
+class _Spec:
+    layers_to_capture = ()
+
+
+class _DRAFT_LOOP:
+    def __init__(self, b):
+        self.batch_size = b
+        self.gather_ids = torch.zeros(b, dtype=torch.int64)
+
+
+class _Worker:
+    is_mtp_eagle = True
+
+
+def test_step0_eligibility():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows.step0 import (  # noqa: E501
+        select_step0,
+        step0_ineligibility,
+    )
+
+    hop = object()
+    draft_loop = _DRAFT_LOOP(1)
+    assert step0_ineligibility(_Worker(), _Draft(hop), _Spec(), draft_loop) is None
+    other = _Worker()
+    other.is_mtp_eagle = False
+    assert step0_ineligibility(other, _Draft(hop), _Spec(), draft_loop) == "not MTP Eagle one-model"
+    assert (
+        step0_ineligibility(_Worker(), _Draft(hop), _Spec(), None) == "no draft_loop draft-loop ops"
+    )
+    cap = _Spec()
+    cap.layers_to_capture = (0,)
+    assert (
+        step0_ineligibility(_Worker(), _Draft(hop), cap, draft_loop) == "MTP layer output captured"
+    )
+    none = _Spec()
+    none.layers_to_capture = None
+    assert step0_ineligibility(_Worker(), _Draft(hop), none, draft_loop) is None
+    assert (
+        step0_ineligibility(_Worker(), _Draft(None), _Spec(), draft_loop)
+        == "MTP layer has no live-row hop"
+    )
+    bad = _DRAFT_LOOP(2)
+    bad.gather_ids = torch.zeros(3, dtype=torch.int64)
+    assert step0_ineligibility(_Worker(), _Draft(hop), _Spec(), bad) == "draft_loop gather ids"
+    bad.gather_ids = None
+    assert step0_ineligibility(_Worker(), _Draft(hop), _Spec(), bad) == "draft_loop gather ids"
+    assert select_step0(_Worker(), _Draft(None), _Spec(), draft_loop) is None
+    assert select_step0(_Worker(), _Draft(hop), _Spec(), draft_loop) is not None
+
+
+def test_arming_is_cleared():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows.step0 import (  # noqa: E501
+        select_step0,
+    )
+
+    draft = _Draft(object())
+    draft_loop = _DRAFT_LOOP(1)
+    layer = draft.mtp_layers[0]
+    s0 = select_step0(_Worker(), draft, _Spec(), draft_loop)
+    with s0.armed() as live:
+        assert layer._draft_live_rows is draft_loop.gather_ids
+        layer._draft_live_rows_compacted = True
+    assert (
+        live.compacted and layer._draft_live_rows is None and not layer._draft_live_rows_compacted
+    )
+    with s0.armed() as live:
+        pass
+    assert not live.compacted
+    with pytest.raises(RuntimeError):
+        with s0.armed():
+            layer._draft_live_rows_compacted = True
+            raise RuntimeError("step failed")
+    assert layer._draft_live_rows is None and not layer._draft_live_rows_compacted
+
+
+def test_mtp_layer_default_off(monkeypatch):
+    """Without draft_live_rows the MTP layer's draft_live_rows attributes stay inert (no hop, never armed)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
+        GlmMtpLayer,
+    )
+
+    d = _dispatch()
+    monkeypatch.setattr(d, "_log_info", lambda message: None)
+    plan_testlib.select(d, GEN_FAST_PATHS_WITHOUT_INPUT_PREP, monkeypatch)
+    assert not d.is_enabled("draft_live_rows")
+    src = inspect.getsource(GlmMtpLayer._adopt_draft)
+    # The hop is created only under is_enabled("draft_live_rows"); the class defaults keep
+    # the attributes inert either way.
+    assert 'if self.post_attn_hop is not None and dispatch.is_enabled("draft_live_rows"):' in src
+    assert GlmMtpLayer.draft_live_row_hop is None and GlmMtpLayer._draft_live_rows is None
+
+
+@pytest.mark.skipif(
+    not os.environ.get("GLM_MOE_DSA_DRAFT_LIVE_ROWS_OP_LIB") and not torch.cuda.is_available(),
+    reason="needs the draft_live_rows op",
+)
+def test_op_registered_with_fake():
+    if not torch.cuda.is_available():
+        pytest.skip("needs a GPU")
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows import (  # noqa: E501
+        hop as live_row_hop,
+    )
+
+    if not live_row_hop.op_available():
+        pytest.skip("draft_live_rows op not loaded in this process")
+    live_row_hop._register_fake()  # idempotent
+    with torch._subclasses.fake_tensor.FakeTensorMode():
+        x = torch.empty(6, 6144, dtype=torch.bfloat16, device="cuda")
+        ids = torch.empty(2, dtype=torch.int64, device="cuda")
+        out = torch.ops.trtllm.glm_moe_dsa_draft_live_row_hop(
+            x,
+            x,
+            ids,
+            torch.empty(6144, dtype=torch.bfloat16, device="cuda"),
+            1e-5,
+            torch.empty(256, 6144, dtype=torch.bfloat16, device="cuda"),
+            torch.empty(10, device="cuda"),
+            1,
+            torch.empty(4, dtype=torch.int32, device="cuda"),
+            0,
+            8,
+        )
+        assert [tuple(t.shape) for t in out] == [(2, 6144), (2, 6144), (2, 256)]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py
new file mode 100644
index 000000000000..6d5aac009ff9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py
@@ -0,0 +1,311 @@
+# 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.
+"""Shared helpers for the draft_loop tests and microbenchmarks."""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional
+
+import torch
+from utils.llm_data import llm_models_root
+
+GLM52_CKPT_ENV = "GLM52_CKPT"
+_DEFAULT_CKPT = os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4")
+
+
+def checkpoint_dir() -> Optional[Path]:
+    """Return the GLM-5.2 NVFP4 checkpoint directory if it is readable."""
+    path = Path(os.environ.get(GLM52_CKPT_ENV, _DEFAULT_CKPT))
+    if (path / "model.safetensors.index.json").is_file():
+        return path
+    return None
+
+
+def load_checkpoint_tensor(name: str, device: str = "cuda") -> Optional[torch.Tensor]:
+    """Load one tensor of the GLM-5.2 checkpoint, or None if unavailable."""
+    ckpt = checkpoint_dir()
+    if ckpt is None:
+        return None
+    from safetensors import safe_open
+
+    weight_map = json.loads((ckpt / "model.safetensors.index.json").read_text())["weight_map"]
+    with safe_open(str(ckpt / weight_map[name]), framework="pt", device="cpu") as f:
+        return f.get_tensor(name).to(device)
+
+
+def real_token_ids(num_tokens: int) -> Optional[torch.Tensor]:
+    """Tokenize a fixed tool-use style text with the checkpoint tokenizer."""
+    ckpt = checkpoint_dir()
+    if ckpt is None:
+        return None
+    # The raw tokenizer avoids loading the model config through transformers.
+    from tokenizers import Tokenizer
+
+    tokenizer = Tokenizer.from_file(str(ckpt / "tokenizer.json"))
+    text = (
+        "<tool_call>search_files(path='/repo/src', pattern='def forward')</tool_call>\n"
+        "The draft loop embeds each proposed token, normalizes it, and projects the "
+        "concatenation with the previous hidden state; 你好，世界。 0x7f3a 3.14159 "
+    ) * 64
+    ids = tokenizer.encode(text, add_special_tokens=True).ids
+    assert len(ids) >= num_tokens
+    return torch.tensor(ids[:num_tokens], dtype=torch.int32)
+
+
+def bits(t: torch.Tensor) -> torch.Tensor:
+    """Bit view of a floating tensor (for bitwise equality incl. -0.0/NaN)."""
+    if t.dtype in (torch.bfloat16, torch.float16):
+        return t.view(torch.int16)
+    if t.dtype == torch.float32:
+        return t.view(torch.int32)
+    return t
+
+
+def graph_time_us(fn, reps: int = 50, replays: int = 20, warmup: int = 3) -> float:
+    """Mean device time of ``fn`` per call, measured as a CUDA graph of ``reps`` calls.
+
+    The graph holds ``reps`` back-to-back invocations so the result includes
+    the inter-kernel launch gaps a captured decode graph pays.
+    """
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(warmup):
+            fn()
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        for _ in range(reps):
+            fn()
+    for _ in range(3):
+        graph.replay()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(replays):
+        graph.replay()
+    end.record()
+    torch.cuda.synchronize()
+    elapsed_ms = start.elapsed_time(end)
+    del graph
+    return elapsed_ms * 1000.0 / (replays * reps)
+
+
+# ----------------------------------------------------------------------------
+# Real DSA generation metadata at the served GEN configuration
+# ----------------------------------------------------------------------------
+
+MTP_K = 5
+TOKENS_PER_BLOCK = 64
+
+
+def served_tp8_sparse_config():
+    """DSA config of the served GEN role (GLM-5.2 indexer geometry)."""
+    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
+
+    return DeepSeekSparseAttentionConfig(
+        index_n_heads=32,
+        index_head_dim=128,
+        index_topk=2048,
+        enable_heuristic_topk=True,
+        use_cute_dsl_paged_mqa_logits=True,
+        use_cute_dsl_topk=True,
+        indexer_k_dtype="fp8",
+    )
+
+
+def make_dsa_cache_manager(batch_size: int, max_kv_len: int, draft_len: int = MTP_K):
+    """A one-layer FP8 DSA cache manager sized for ``batch_size`` long requests."""
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm.bindings import DataType
+    from tensorrt_llm.bindings.executor import KvCacheConfig
+    from tensorrt_llm.bindings.internal.batch_manager import CacheType
+    from tensorrt_llm.mapping import Mapping
+
+    mapping = Mapping(world_size=1, tp_size=1, rank=0)
+    sparse_config = served_tp8_sparse_config()
+    model_config = ModelConfig(
+        mapping=mapping,
+        sparse_attention_config=sparse_config,
+        pretrained_config=SimpleNamespace(rms_norm_eps=1e-5),
+    )
+    max_seq_len = max_kv_len + 2 * (draft_len + 1) + TOKENS_PER_BLOCK
+    blocks = (max_seq_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK + 2
+    return (
+        DSACacheManager(
+            KvCacheConfig(
+                max_tokens=batch_size * blocks * TOKENS_PER_BLOCK, enable_block_reuse=False
+            ),
+            CacheType.SELFKONLY,
+            num_layers=1,
+            num_kv_heads=1,
+            head_dim=576,
+            tokens_per_block=TOKENS_PER_BLOCK,
+            max_seq_len=max_seq_len,
+            max_batch_size=batch_size,
+            mapping=mapping,
+            dtype=DataType.FP8,
+            sparse_attn_config=sparse_config,
+            model_config=model_config,
+        ),
+        mapping,
+        sparse_config,
+    )
+
+
+def make_gen_dsa_metadata(
+    cache_manager,
+    mapping,
+    sparse_config,
+    request_ids,
+    past_kv_lens,
+    draft_len: int = MTP_K,
+    draft_cache_manager=None,
+):
+    """Prepared generation metadata for one MTP verify step (K+1 query tokens each).
+
+    Follows the decoder runner's order: ``update_spec_dec_param`` (one-model MTP,
+    linear tree), then ``KVCacheParams`` with the one-model extra KV tokens and
+    ``prepare()``.
+    """
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata
+    from tensorrt_llm._torch.metadata import KVCacheParams
+
+    batch_size = len(request_ids)
+    md = DSAtrtllmAttentionMetadata(
+        seq_lens=torch.full((batch_size,), draft_len + 1, dtype=torch.int),
+        request_ids=list(request_ids),
+        max_num_requests=batch_size,
+        num_contexts=0,
+        prompt_lens=list(past_kv_lens),
+        max_num_tokens=batch_size * (draft_len + 1),
+        kv_cache_manager=cache_manager,
+        draft_kv_cache_manager=draft_cache_manager,
+        mapping=mapping,
+        sparse_attention_config=sparse_config,
+    )
+    md.update_spec_dec_param(
+        batch_size=batch_size,
+        is_spec_decoding_enabled=True,
+        is_spec_dec_tree=False,
+        is_spec_dec_dynamic_tree=False,
+        max_draft_len=draft_len,
+        max_total_draft_tokens=draft_len,
+        spec_metadata=None,
+        spec_tree_manager=None,
+        num_contexts=0,
+    )
+    md.kv_cache_params = KVCacheParams(
+        use_cache=True,
+        num_cached_tokens_per_seq=list(past_kv_lens),
+        num_extra_kv_tokens=draft_len - 1,
+    )
+    md.prepare()
+    return md
+
+
+def add_gen_requests(cache_manager, request_ids, past_kv_lens, draft_len: int = MTP_K):
+    """Allocate KV blocks for generation requests with the given past lengths."""
+    cache_manager.add_dummy_requests(
+        list(request_ids),
+        token_nums=[n + 1 for n in past_kv_lens],
+        is_gen=True,
+        prepare_resource=True,
+        max_num_draft_tokens=draft_len,
+    )
+
+
+def tensor_state(obj) -> dict:
+    """Clone every tensor and scalar attribute of ``obj`` (for state comparison)."""
+    state = {}
+    for name, value in vars(obj).items():
+        if isinstance(value, torch.Tensor):
+            state[name] = value.detach().clone()
+        elif isinstance(value, (bool, int, float, str, type(None))):
+            state[name] = value
+    return state
+
+
+def changed_keys(before: dict, after: dict) -> set:
+    """Attributes whose value differs between two ``tensor_state`` snapshots."""
+    keys = set()
+    for name in set(before) | set(after):
+        a, b = before.get(name, "<missing>"), after.get(name, "<missing>")
+        if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
+            if a.shape != b.shape or a.dtype != b.dtype or not torch.equal(a.cpu(), b.cpu()):
+                keys.add(name)
+        elif type(a) is not type(b) or a != b:
+            keys.add(name)
+    return keys
+
+
+def assert_same_state(native: dict, candidate: dict, keys, what: str) -> None:
+    """Assert that ``keys`` hold identical values in both snapshots."""
+    for name in sorted(keys):
+        a, b = native.get(name, "<missing>"), candidate.get(name, "<missing>")
+        if isinstance(a, torch.Tensor):
+            assert isinstance(b, torch.Tensor), f"{what}: {name} type differs"
+            assert a.shape == b.shape and a.dtype == b.dtype, f"{what}: {name} shape/dtype"
+            assert torch.equal(a.cpu(), b.cpu()), f"{what}: {name} differs"
+        else:
+            assert a == b, f"{what}: {name} differs ({a!r} vs {b!r})"
+
+
+def native_draft_step_tail(
+    md,
+    i,
+    batch_size,
+    num_contexts,
+    runtime_draft_len,
+    num_accepted_tokens,
+    position_ids,
+    gather_ids,
+):
+    """Literal copy of the native loop tail for draft step ``i`` (MTP Eagle, DSA)."""
+    from tensorrt_llm._torch.speculative.mtp import _select_mtp_position_ids
+
+    new_position_ids = None
+    if i + 1 < runtime_draft_len:
+        new_position_ids = _select_mtp_position_ids(position_ids, gather_ids) + 1
+    attn_metadata = md
+    if i == 0:
+        attn_metadata._seq_lens[:batch_size].fill_(1)
+        attn_metadata._seq_lens_cuda[:batch_size].fill_(1)
+        attn_metadata.on_update()
+        has_kv_cache = attn_metadata.kv_cache_manager is not None
+        if has_kv_cache and hasattr(attn_metadata, "host_request_types"):
+            attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1)
+        if has_kv_cache:
+            attn_metadata.num_contexts = 0
+        if hasattr(attn_metadata, "kv_lens_cuda"):
+            attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (
+                runtime_draft_len - num_accepted_tokens[num_contexts:]
+            )
+            attn_metadata.kv_lens_cuda[:num_contexts] += 1
+        # _prepare_flash_mla_generation_layout is a no-op without context rows.
+        if hasattr(attn_metadata, "kv_lens_cuda"):
+            attn_metadata.update_for_spec_dec()
+        attn_metadata.use_spec_decoding = False
+    else:
+        if hasattr(attn_metadata, "kv_lens_cuda"):
+            attn_metadata.kv_lens_cuda[:batch_size] += 1
+            attn_metadata.update_for_spec_dec()
+    return new_position_ids
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py
new file mode 100644
index 000000000000..77905e327352
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py
@@ -0,0 +1,614 @@
+# 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.
+"""draft-loop bookkeeping versus the native MTP one-model worker.
+
+Unit level: fused acceptance (strict + forced-acceptance override, including
+the draw counter) against ``_sample_and_accept_draft_tokens_base``; draft-0
+inputs against ``prepare_1st_drafter_inputs``; the row gather and the
+post-loop finalize against the native indexing.
+
+Worker level: the real ``MTPEagleWorker.forward`` of the target's adopted worker
+(``GlmMtpWorker``) on decode steps -- the native draft loop -- on real DSA metadata,
+with draft_loop off and on. A stub MTP layer (consuming the real embedding callable)
+records what every draft step sees; outputs, per-step inputs, per-step
+metadata and the final metadata state must be identical, eagerly and under
+CUDA-graph replay.
+"""
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from _draft_loop_fixtures import (  # noqa: E402
+    MTP_K,
+    add_gen_requests,
+    assert_same_state,
+    changed_keys,
+    make_dsa_cache_manager,
+    make_gen_dsa_metadata,
+    tensor_state,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
+    dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E402, E501
+    draft_loop as draft_loop_mod,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E402, E501
+    draft_step as draft_step_mod,
+)
+
+VOCAB = 1024
+HIDDEN = 6144
+
+
+def _make_worker(force: float):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
+        GlmMtpWorker,
+    )
+    from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+    from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
+
+    worker = GlmMtpWorker.adopt(MTPEagleWorker(MTPDecodingConfig(max_draft_len=MTP_K)))
+    worker.force_num_accepted_tokens = force
+    return worker
+
+
+def _spec_metadata(batch_size, draft_tokens, is_cuda_graph=False):
+    return SimpleNamespace(
+        runtime_draft_len=MTP_K,
+        draft_tokens=draft_tokens,
+        batch_indices_cuda=torch.arange(max(batch_size, 8), dtype=torch.int32, device="cuda"),
+        is_cuda_graph=is_cuda_graph,
+        is_all_greedy_sample=True,
+        use_rejection_sampling=False,
+        enable_penalty=False,
+        wants_advanced_draft_sampling=False,
+        all_rank_num_tokens=None,
+        subseq_all_rank_num_tokens=None,
+        spec_resource_manager=None,
+        max_num_requests=batch_size,
+        gather_ids=None,
+        context_prompt_lookahead_tokens=None,
+        is_ragged_verify=False,
+    )
+
+
+def _target_logits(targets: torch.Tensor) -> torch.Tensor:
+    """Float logits whose argmax per row is ``targets``."""
+    logits = torch.randn((targets.numel(), VOCAB), device="cuda")
+    logits[torch.arange(targets.numel(), device="cuda"), targets.long()] = 100.0
+    return logits
+
+
+def _random_case(gen, batch_size):
+    """Draft tokens and target tokens with a random leading-match length per row."""
+    draft = torch.randint(0, VOCAB, (batch_size, MTP_K), generator=gen, dtype=torch.int32)
+    target = torch.randint(0, VOCAB, (batch_size, MTP_K + 1), generator=gen, dtype=torch.int32)
+    for b in range(batch_size):
+        n = int(torch.randint(0, MTP_K + 1, (1,), generator=gen))
+        target[b, :n] = draft[b, :n]
+        if n < MTP_K and target[b, n] == draft[b, n]:
+            target[b, n] = (draft[b, n] + 1) % VOCAB
+    return draft.cuda().reshape(-1), target.cuda()
+
+
+@pytest.mark.parametrize("force", [0.0, 2.61, 0.5, 3.0, 5.0, 7.0])
+@pytest.mark.parametrize("batch_size", [1, 3, 8])
+def test_accept_matches_native(force, batch_size):
+    native, fused = _make_worker(force), _make_worker(force)
+    gen = torch.Generator().manual_seed(int(force * 100) + batch_size)
+    md_stub = SimpleNamespace(num_seqs=batch_size, num_contexts=0, num_ctx_tokens=0)
+    for call in range(40):
+        draft, target = _random_case(gen, batch_size)
+        logits = _target_logits(target.reshape(-1))
+        spec = _spec_metadata(batch_size, draft)
+        acc_n, num_n = native._sample_and_accept_draft_tokens_base(
+            logits, draft.reshape(batch_size, MTP_K), 0, batch_size, spec
+        )
+        ops = draft_loop_mod.DraftLoopOps(fused, md_stub, spec)
+        assert draft_loop_mod.DraftLoopOps.ineligibility(fused, md_stub, spec) is None
+        acc_w, num_w = ops.accept(logits)
+        assert acc_w.dtype == acc_n.dtype and torch.equal(acc_w, acc_n), call
+        assert num_w.dtype == num_n.dtype and torch.equal(num_w, num_n), call
+        if native._force_accept_rng_counter is not None:
+            assert torch.equal(native._force_accept_rng_counter, fused._force_accept_rng_counter)
+        # Draft-0 inputs and step-0 rows.
+        inputs = native.prepare_1st_drafter_inputs(
+            input_ids=torch.zeros((batch_size * (MTP_K + 1),), dtype=torch.int32, device="cuda"),
+            position_ids=None,
+            hidden_states=None,
+            accepted_tokens=acc_n,
+            attn_metadata=md_stub,
+            spec_metadata=spec,
+            draft_model=None,
+        )
+        assert torch.equal(ops.input_ids0, inputs["input_ids"])
+        assert ops.input_ids0.data_ptr() != acc_w.data_ptr()
+        rows = (spec.batch_indices_cuda[:batch_size] * (MTP_K + 1)).long() + num_n - 1 + 0
+        assert torch.equal(ops.gather_ids, rows)
+
+
+def test_gather_rows_and_finalize_match_native():
+    worker = _make_worker(0.0)
+    gen = torch.Generator(device="cuda").manual_seed(3)
+    for batch_size in (1, 2, 6):
+        hidden = torch.randn((batch_size * (MTP_K + 1), HIDDEN), generator=gen, device="cuda")
+        hidden = hidden.to(torch.bfloat16)
+        rows = torch.randint(0, hidden.shape[0], (batch_size,), generator=gen, device="cuda")
+        assert torch.equal(draft_loop_mod.gather_rows(hidden, rows), hidden[rows])
+
+        draft = [
+            torch.randint(0, VOCAB, (batch_size,), generator=gen, device="cuda", dtype=torch.int32)
+            for _ in range(MTP_K)
+        ]
+        accepted = torch.randint(
+            0, VOCAB, (batch_size, MTP_K + 1), generator=gen, device="cuda", dtype=torch.int32
+        )
+        num_acc = torch.randint(
+            1, MTP_K + 2, (batch_size,), generator=gen, device="cuda", dtype=torch.int32
+        )
+        spec = _spec_metadata(batch_size, None)
+        ops = draft_loop_mod.DraftLoopOps(
+            worker, SimpleNamespace(num_seqs=batch_size, num_contexts=0), spec
+        )
+        ops.accepted_tokens, ops.num_accepted_tokens = accepted, num_acc
+        stacked = torch.stack(draft, dim=1)
+        expected_new = worker._prepare_next_new_tokens(
+            accepted, stacked, spec.batch_indices_cuda, batch_size, num_acc
+        )
+        got = ops.finalize(draft)
+        assert torch.equal(got, stacked)
+        assert torch.equal(ops.next_new_tokens, expected_new)
+
+
+# ----------------------------------------------------------------------------
+# Worker level: the adopted worker's draft loop on decode steps
+# ----------------------------------------------------------------------------
+
+_STEP_FIELDS = (
+    "kv_lens_cuda",
+    "kv_lens_cuda_2d",
+    "scheduler_metadata_buffer",
+    "scheduler_metadata_buffer_full_next_n",
+    "kv_lens_expanded_cuda",
+    "scheduler_metadata_buffer_expanded",
+    "req_idx_per_token",
+    "slot_mapping_fp8",
+    "slot_mapping_scale",
+    "gen_kv_indptr",
+    "gen_cached_token_indptr",
+)
+
+
+class _Recorder:
+    """Per-step copies of what the MTP layer sees (static buffers under capture)."""
+
+    def __init__(self):
+        self.steps = []
+        self.static = None
+        self.index = 0
+
+    def record(self, md, input_ids, position_ids, hidden, emb):
+        items = {
+            "input_ids": input_ids,
+            "position_ids": position_ids,
+            "hidden": hidden,
+            "emb": emb,
+            "seq_lens_cuda": md.seq_lens_cuda,
+        }
+        items.update({f: getattr(md, f) for f in _STEP_FIELDS})
+        if torch.cuda.is_current_stream_capturing():
+            if self.static is None:
+                self.static = []
+            if self.index == len(self.static):
+                self.static.append({k: torch.empty_like(v) for k, v in items.items()})
+            for k, v in items.items():
+                self.static[self.index][k].copy_(v)
+            self.index += 1
+        else:
+            self.steps.append({k: v.clone() for k, v in items.items()})
+            self.steps[-1]["num_tokens"] = md.num_tokens
+
+
+class _StubHead:
+    def __call__(self, hidden, lm_head, attn_metadata, return_context_logits):
+        return hidden.float() @ lm_head.weight.float().t()
+
+
+class _StubMTPLayer(torch.nn.Module):
+    """Deterministic draft layer: embeds the ids with the given callable."""
+
+    def __init__(self, recorder):
+        super().__init__()
+        self.recorder = recorder
+        self.shared_head = _StubHead()
+
+    def forward(
+        self,
+        input_ids,
+        position_ids,
+        hidden_states,
+        embed_tokens,
+        attn_metadata,
+        all_rank_num_tokens=None,
+        spec_metadata=None,
+        **kwargs,
+    ):
+        emb = embed_tokens(input_ids)
+        self.recorder.record(attn_metadata, input_ids, position_ids, hidden_states, emb)
+        mix = emb.float() + 0.5 * hidden_states.float()
+        mix = mix + (position_ids.float() % 97).unsqueeze(-1) * 1e-2
+        return mix.to(hidden_states.dtype)
+
+
+def _draft_model(recorder, seed=11):
+    from tensorrt_llm._torch.modules.embedding import Embedding
+
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    embed = Embedding(VOCAB, HIDDEN, dtype=torch.bfloat16).cuda()
+    embed.weight.data = torch.randn((VOCAB, HIDDEN), generator=gen, device="cuda").to(
+        torch.bfloat16
+    )
+    lm_head = SimpleNamespace(
+        weight=torch.randn((VOCAB, HIDDEN), generator=gen, device="cuda").to(torch.bfloat16)
+    )
+    return SimpleNamespace(
+        mtp_layers=[_StubMTPLayer(recorder)], embed_tokens=embed, lm_head=lm_head
+    )
+
+
+def _forward_inputs(batch_size, past, gen):
+    draft, target = _random_case(gen, batch_size)
+    n = batch_size * (MTP_K + 1)
+    return {
+        "input_ids": torch.randint(0, VOCAB, (n,), generator=gen, dtype=torch.int32).cuda(),
+        "position_ids": torch.cat(
+            [torch.arange(p, p + MTP_K + 1, dtype=torch.int32) for p in past]
+        ).cuda()[None],
+        "hidden_states": torch.randn((n, HIDDEN), generator=gen).to(torch.bfloat16).cuda(),
+        "logits": _target_logits(target.reshape(-1)),
+        "draft_tokens": draft,
+    }
+
+
+class _ResourceManagerStub:
+    """Resource manager exposing only the separate draft KV cache manager."""
+
+    def __init__(self, draft_manager):
+        self.draft_manager = draft_manager
+
+    def get_resource_manager(self, kind):
+        from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
+
+        return self.draft_manager if kind == ResourceManagerType.DRAFT_KV_CACHE_MANAGER else None
+
+
+def _run_forward(worker, md, spec, draft_model, inputs):
+    resource_manager = None
+    if md.draft_kv_cache_manager is not None:
+        # One-model MTP with a separate draft KV cache, as the GEN runner sets up.
+        worker.use_separate_draft_kv_cache = True
+        resource_manager = _ResourceManagerStub(md.draft_kv_cache_manager)
+    with dispatch.step(dispatch.Phase.DECODE):
+        return worker.forward(
+            input_ids=inputs["input_ids"],
+            position_ids=inputs["position_ids"],
+            hidden_states=inputs["hidden_states"],
+            logits=inputs["logits"],
+            attn_metadata=md,
+            spec_metadata=spec,
+            draft_model=draft_model,
+            resource_manager=resource_manager,
+        )
+
+
+_OUTPUT_KEYS = ("new_tokens", "new_tokens_lens", "next_draft_tokens", "next_new_tokens")
+
+
+@pytest.mark.parametrize("force", [0.0, 2.61])
+@pytest.mark.parametrize("past", [[391061], [391061, 120000, 4100]])
+@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
+def test_worker_forward_matches_native(monkeypatch, force, past, separate_draft):
+    batch_size = len(past)
+    cache_manager, mapping, sparse = make_dsa_cache_manager(batch_size, max(past))
+    draft_manager = None
+    calls = {"accept": 0, "step": 0}
+    real_accept, real_step = draft_loop_mod.DraftLoopOps.accept, draft_step_mod.draft_step_update
+
+    def spy_accept(self, logits):
+        calls["accept"] += 1
+        return real_accept(self, logits)
+
+    def spy_step(*args, **kwargs):
+        calls["step"] += 1
+        return real_step(*args, **kwargs)
+
+    monkeypatch.setattr(draft_loop_mod.DraftLoopOps, "accept", spy_accept)
+    monkeypatch.setattr(draft_step_mod, "draft_step_update", spy_step)
+    try:
+        ids = list(range(batch_size))
+        add_gen_requests(cache_manager, ids, past)
+        if separate_draft:
+            draft_manager, _, _ = make_dsa_cache_manager(batch_size, max(past))
+            add_gen_requests(draft_manager, ids, past)
+        gen = torch.Generator().manual_seed(len(past) + int(force * 10))
+        for it in range(3):
+            inputs = _forward_inputs(batch_size, past, gen)
+            results = {}
+            for name, features in (
+                ("native", frozenset()),
+                ("draft_loop", frozenset({"draft_loop"})),
+            ):
+                monkeypatch.setattr(dispatch, "enabled_features", lambda f=features: f)
+                recorder = _Recorder()
+                worker = _make_worker(force)
+                md = make_gen_dsa_metadata(
+                    cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
+                )
+                base = tensor_state(md)
+                spec = _spec_metadata(batch_size, inputs["draft_tokens"].clone())
+                out = _run_forward(worker, md, spec, _draft_model(recorder), inputs)
+                torch.cuda.synchronize()
+                results[name] = (out, recorder.steps, base, tensor_state(md), worker)
+            assert calls["accept"] == it + 1 and calls["step"] == (it + 1) * MTP_K
+            (o_n, s_n, b_n, f_n, w_n), (o_w, s_w, b_w, f_w, w_w) = (
+                results["native"],
+                results["draft_loop"],
+            )
+            for k in _OUTPUT_KEYS:
+                assert o_n[k].dtype == o_w[k].dtype and torch.equal(o_n[k], o_w[k]), (it, k)
+            assert len(s_n) == len(s_w) == MTP_K
+            for i, (a, b) in enumerate(zip(s_n, s_w)):
+                for k in a:
+                    if isinstance(a[k], torch.Tensor):
+                        assert a[k].dtype == b[k].dtype and torch.equal(a[k], b[k]), (it, i, k)
+                    else:
+                        assert a[k] == b[k], (it, i, k)
+            keys = changed_keys(b_n, f_n) | changed_keys(b_w, f_w)
+            assert_same_state(f_n, f_w, keys, f"iteration {it} final metadata")
+            if w_n._force_accept_rng_counter is not None:
+                assert torch.equal(w_n._force_accept_rng_counter, w_w._force_accept_rng_counter)
+    finally:
+        cache_manager.shutdown()
+        if draft_manager is not None:
+            draft_manager.shutdown()
+
+
+def _capture_and_replay(
+    monkeypatch,
+    features,
+    cache_manager,
+    mapping,
+    sparse,
+    ids,
+    past,
+    replay_inputs,
+    between=None,
+    keep_graph=False,
+    draft_manager=None,
+):
+    """Capture one worker forward for ``features``, replay it per input set.
+
+    Returns per-replay clones of the outputs, the per-step records and the
+    draw counter. ``between`` runs after capture and before the replays.
+    """
+    batch_size = len(ids)
+    monkeypatch.setattr(dispatch, "enabled_features", lambda f=features: f)
+    recorder = _Recorder()
+    worker = _make_worker(2.61)
+    md = make_gen_dsa_metadata(
+        cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
+    )
+    kv_init = md.kv_lens_cuda[:batch_size].clone()
+    static = {k: v.clone() for k, v in replay_inputs[0].items()}
+    spec = _spec_metadata(batch_size, static["draft_tokens"], is_cuda_graph=True)
+    draft_model = _draft_model(recorder)
+    host_seq = md._seq_lens.clone()
+    # The forced-acceptance pool is built eagerly, never under capture.
+    worker._ensure_force_accept_rng_state(torch.device("cuda"))
+    for _ in range(2):  # eager warmup (forced acceptance is skipped here)
+        md.kv_lens_cuda[:batch_size].copy_(kv_init)
+        _run_forward(worker, md, spec, draft_model, static)
+        md._seq_lens.copy_(host_seq)
+        md.num_contexts = 0
+        md.on_update()
+    md.kv_lens_cuda[:batch_size].copy_(kv_init)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = _run_forward(worker, md, spec, draft_model, static)
+    md._seq_lens.copy_(host_seq)
+    if between is not None:
+        between()
+    results = []
+    for replay, inputs in enumerate(replay_inputs):
+        for k, v in inputs.items():
+            static[k].copy_(v)
+        md.kv_lens_cuda[:batch_size].copy_(kv_init)
+        graph.replay()
+        try:
+            torch.cuda.synchronize()
+        except (RuntimeError, torch.AcceleratorError) as err:
+            raise AssertionError(f"{sorted(features)} graph replay {replay} faulted") from err
+        results.append(
+            (
+                {k: out[k].clone() for k in _OUTPUT_KEYS},
+                [{k: v.clone() for k, v in step.items()} for step in recorder.static],
+                worker._force_accept_rng_counter.clone(),
+            )
+        )
+    if keep_graph:
+        return results, graph, md, static
+    del graph
+    torch.cuda.synchronize()
+    return results
+
+
+def _graph_setup():
+    past = [391061, 200000]
+    cache_manager, mapping, sparse = make_dsa_cache_manager(len(past), max(past) + 64)
+    ids = list(range(len(past)))
+    add_gen_requests(cache_manager, ids, past)
+    gen = torch.Generator().manual_seed(21)
+    replay_inputs = [_forward_inputs(len(past), past, gen) for _ in range(10)]
+    return cache_manager, mapping, sparse, ids, past, replay_inputs
+
+
+@pytest.mark.parametrize(
+    "scenario",
+    [
+        "native",
+        "native_then_self_check",
+        "draft_loop",
+        "native_then_new_metadata",
+        "native_then_native_eager",
+        "native_then_draft_loop_eager",
+        "native_then_native_capture",
+        "native_then_draft_loop_capture",
+    ],
+)
+def test_graph_isolation_diagnostics(monkeypatch, scenario):
+    """A captured graph replays cleanly after other work (fault localization)."""
+    cache_manager, mapping, sparse, ids, past, replay_inputs = _graph_setup()
+    try:
+        between = None
+        if scenario.startswith("native_then_") and scenario != "native_then_self_check":
+            eager_features = {
+                "native_then_native_eager": frozenset(),
+                "native_then_draft_loop_eager": frozenset({"draft_loop"}),
+                "native_then_native_capture": frozenset(),
+                "native_then_draft_loop_capture": frozenset({"draft_loop"}),
+            }.get(scenario)
+            second_graphs = []
+
+            def between():
+                if scenario.endswith("_capture"):
+                    # A second captured graph that stays alive during the replays.
+                    second_graphs.append(
+                        _capture_and_replay(
+                            monkeypatch,
+                            eager_features,
+                            cache_manager,
+                            mapping,
+                            sparse,
+                            ids,
+                            past,
+                            replay_inputs[:1],
+                            keep_graph=True,
+                        )
+                    )
+                    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset())
+                    return
+                md2 = make_gen_dsa_metadata(cache_manager, mapping, sparse, ids, past)
+                if eager_features is None:
+                    return
+                monkeypatch.setattr(dispatch, "enabled_features", lambda: eager_features)
+                worker2 = _make_worker(2.61)
+                spec2 = _spec_metadata(len(ids), replay_inputs[0]["draft_tokens"].clone())
+                for _ in range(2):
+                    _run_forward(worker2, md2, spec2, _draft_model(_Recorder()), replay_inputs[0])
+                torch.cuda.synchronize()
+                monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset())
+
+        if scenario == "native_then_self_check":
+
+            def between():
+                draft_step_mod._SPLIT_KV_BY_SMS.clear()
+                assert (
+                    draft_step_mod.schedule_self_check(
+                        torch.cuda.get_device_properties(0).multi_processor_count,
+                        torch.device("cuda"),
+                    )
+                    is not None
+                )
+
+        features = frozenset({"draft_loop"}) if scenario == "draft_loop" else frozenset()
+        results = _capture_and_replay(
+            monkeypatch,
+            features,
+            cache_manager,
+            mapping,
+            sparse,
+            ids,
+            past,
+            replay_inputs[:3],
+            between,
+        )
+        assert len(results) == 3
+    finally:
+        cache_manager.shutdown()
+
+
+@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
+def test_worker_forward_cuda_graph(monkeypatch, separate_draft):
+    """Captured native and draft_loop forwards replay identically (forced acceptance 2.61).
+
+    Outputs, per-step records and draw counters are compared over 10 replays
+    with fresh inputs.
+    """
+    # The draft_loop schedule self-check runs eagerly before any capture, as in a
+    # serving warmup.
+    assert draft_step_mod.schedule_self_check(
+        torch.cuda.get_device_properties(0).multi_processor_count, torch.device("cuda")
+    )
+    cache_manager, mapping, sparse, ids, past, replay_inputs = _graph_setup()
+    draft_manager = None
+    try:
+        if separate_draft:
+            draft_manager, _, _ = make_dsa_cache_manager(len(past), max(past) + 64)
+            add_gen_requests(draft_manager, ids, past)
+        native = _capture_and_replay(
+            monkeypatch,
+            frozenset(),
+            cache_manager,
+            mapping,
+            sparse,
+            ids,
+            past,
+            replay_inputs,
+            draft_manager=draft_manager,
+        )
+        fused = _capture_and_replay(
+            monkeypatch,
+            frozenset({"draft_loop"}),
+            cache_manager,
+            mapping,
+            sparse,
+            ids,
+            past,
+            replay_inputs,
+            draft_manager=draft_manager,
+        )
+        for replay, ((o_n, r_n, c_n), (o_w, r_w, c_w)) in enumerate(zip(native, fused)):
+            for k in _OUTPUT_KEYS:
+                assert torch.equal(o_n[k], o_w[k]), (replay, k)
+            assert len(r_n) == len(r_w) == MTP_K
+            for i in range(MTP_K):
+                for k in r_n[i]:
+                    assert torch.equal(r_n[i][k], r_w[i][k]), (replay, i, k)
+            assert torch.equal(c_n, c_w), replay
+    finally:
+        cache_manager.shutdown()
+        if draft_manager is not None:
+            draft_manager.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py
new file mode 100644
index 000000000000..a1a6147b6ec2
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py
@@ -0,0 +1,314 @@
+# 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.
+"""draft_loop fused draft-step metadata update versus the native draft-loop tail.
+
+The native reference is the literal per-step tail of
+``Eagle3OneModelWorker._forward_linear_draft_loop`` (position advance, step-0
+switch to one query token, KV-length update, ``update_for_spec_dec``) running
+on a real ``DSAtrtllmAttentionMetadata`` prepared from a real ``DSACacheManager``
+at the served GEN configuration. After every step every tensor and scalar
+attribute either path changed must be identical.
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from _draft_loop_fixtures import (  # noqa: E402
+    MTP_K,
+    add_gen_requests,
+    assert_same_state,
+    changed_keys,
+    make_dsa_cache_manager,
+    make_gen_dsa_metadata,
+    native_draft_step_tail,
+    tensor_state,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E402, E501
+    draft_step as draft_step_mod,
+)
+
+
+def _step(md, use_draft_loop, i, batch_size, num_accepted, position_ids, gather_ids):
+    pos_arg = position_ids if i + 1 < MTP_K else None
+    gather = gather_ids if i == 0 else slice(0, batch_size)
+    if use_draft_loop:
+        result = draft_step_mod.draft_step_update(
+            md, i, batch_size, 0, MTP_K, num_accepted, pos_arg, gather_ids if i == 0 else None
+        )
+        assert result is not None, draft_step_mod.step_ineligibility(md, batch_size, 0)
+        return result.position_ids
+    return native_draft_step_tail(md, i, batch_size, 0, MTP_K, num_accepted, position_ids, gather)
+
+
+def _enter_draft(md):
+    """The native loop's draft-KV entry: rebind to the draft buffers, refresh."""
+    if md.draft_kv_cache_manager is None:
+        return None
+    from tensorrt_llm._torch.speculative.interface import prepare_attn_metadata_for_draft_replay
+
+    saved = prepare_attn_metadata_for_draft_replay(md, md.draft_kv_cache_manager)
+    md.on_update_kv_lens()
+    return saved
+
+
+def _exit_draft(md, saved):
+    if saved is not None:
+        from tensorrt_llm._torch.speculative.interface import (
+            restore_attn_metadata_after_draft_replay,
+        )
+
+        restore_attn_metadata_after_draft_replay(md, saved)
+
+
+def _run_loop(md, use_draft_loop, num_accepted, position_ids, gather_ids):
+    batch_size = md.num_seqs
+    md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
+    saved = _enter_draft(md)
+    states = []
+    pos = position_ids
+    for i in range(MTP_K):
+        pos = _step(md, use_draft_loop, i, batch_size, num_accepted, pos, gather_ids)
+        torch.cuda.synchronize()
+        states.append((tensor_state(md), None if pos is None else pos.clone()))
+    _exit_draft(md, saved)
+    md.restore_from_spec_dec()
+    md.on_update()
+    return states, tensor_state(md)
+
+
+def _inputs(batch_size, past_kv_lens, accepted):
+    num_accepted = torch.tensor(accepted, dtype=torch.int32, device="cuda")
+    # Target verify positions: past_kv_len .. past_kv_len + K for each request.
+    pos = torch.cat(
+        [torch.arange(n, n + MTP_K + 1, dtype=torch.int32) for n in past_kv_lens]
+    ).cuda()
+    gather_ids = (
+        torch.arange(batch_size, dtype=torch.int64, device="cuda") * (MTP_K + 1)
+        + num_accepted.to(torch.int64)
+        - 1
+    )
+    return num_accepted, pos, gather_ids
+
+
+CASES = [
+    # (past KV lengths, accepted counts)
+    ([391061], [4]),
+    ([391061], [1]),
+    ([391061], [6]),
+    ([2100], [3]),
+    ([65536 - 3], [6]),
+    ([(1 << 20) - 64], [2]),
+    ([491061, 32768, 131072], [6, 1, 3]),
+    ([4096, 4097, 4098, 4099, 4100, 4101, 4102, 300000], [1, 2, 3, 4, 5, 6, 4, 3]),
+]
+
+
+@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
+@pytest.mark.parametrize("past_kv_lens, accepted", CASES)
+def test_step_update_matches_native(past_kv_lens, accepted, separate_draft):
+    batch_size = len(past_kv_lens)
+    cache_manager, mapping, sparse = make_dsa_cache_manager(batch_size, max(past_kv_lens))
+    draft_manager = None
+    try:
+        ids = list(range(batch_size))
+        add_gen_requests(cache_manager, ids, past_kv_lens)
+        if separate_draft:
+            draft_manager, _, _ = make_dsa_cache_manager(batch_size, max(past_kv_lens))
+            add_gen_requests(draft_manager, ids, past_kv_lens)
+        md_native = make_gen_dsa_metadata(
+            cache_manager, mapping, sparse, ids, past_kv_lens, draft_cache_manager=draft_manager
+        )
+        md_draft_loop = make_gen_dsa_metadata(
+            cache_manager, mapping, sparse, ids, past_kv_lens, draft_cache_manager=draft_manager
+        )
+        reason = draft_step_mod.step_ineligibility(md_draft_loop, batch_size, 0)
+        assert reason is None, reason
+        print(
+            f"expand_for_dsl={md_draft_loop.expand_for_dsl} factor={md_draft_loop.dsl_expand_factor} "
+            f"num_sms={md_draft_loop.num_sms} split_kv={draft_step_mod._SPLIT_KV_BY_SMS[md_draft_loop.num_sms]} "
+            f"width={md_draft_loop.gen_token_stride}"
+        )
+        base_native, base_draft_loop = tensor_state(md_native), tensor_state(md_draft_loop)
+        inputs = _inputs(batch_size, past_kv_lens, accepted)
+        steps_native, final_native = _run_loop(md_native, False, *inputs)
+        steps_draft_loop, final_draft_loop = _run_loop(md_draft_loop, True, *inputs)
+        for i, ((s_n, p_n), (s_w, p_w)) in enumerate(zip(steps_native, steps_draft_loop)):
+            keys = changed_keys(base_native, s_n) | changed_keys(base_draft_loop, s_w)
+            # The query lengths always change (K + 1 -> 1); the step buffers are
+            # compared whether or not they changed in this step.
+            assert "_seq_lens_cuda" in keys
+            assert_same_state(s_n, s_w, keys | set(_SNAPSHOT_FIELDS), f"step {i}")
+            if p_n is None:
+                assert p_w is None
+            else:
+                assert p_w is not None and p_w.dtype == p_n.dtype and torch.equal(p_w, p_n)
+        keys = changed_keys(base_native, final_native) | changed_keys(
+            base_draft_loop, final_draft_loop
+        )
+        assert_same_state(final_native, final_draft_loop, keys, "after restore")
+    finally:
+        cache_manager.shutdown()
+        if draft_manager is not None:
+            draft_manager.shutdown()
+
+
+def test_ineligible_short_context_is_untouched():
+    """Dense top-k (short context) keeps the native chain; draft_loop changes nothing."""
+    cache_manager, mapping, sparse = make_dsa_cache_manager(1, 1000)
+    try:
+        add_gen_requests(cache_manager, [0], [1000])
+        md = make_gen_dsa_metadata(cache_manager, mapping, sparse, [0], [1000])
+        assert draft_step_mod.step_ineligibility(md, 1, 0) == "dense top-k fallback active"
+        md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
+        before = tensor_state(md)
+        num_accepted, pos, gather_ids = _inputs(1, [1000], [3])
+        assert (
+            draft_step_mod.draft_step_update(md, 0, 1, 0, MTP_K, num_accepted, pos, gather_ids)
+            is None
+        )
+        torch.cuda.synchronize()
+        assert not changed_keys(before, tensor_state(md))
+        md.restore_from_spec_dec()
+    finally:
+        cache_manager.shutdown()
+
+
+def test_schedule_matches_deepgemm():
+    """The draft_loop schedule equals DeepGEMM's on random 1-D/2-D context lengths."""
+    from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata
+
+    num_sms = torch.cuda.get_device_properties(0).multi_processor_count
+    split_kv = draft_step_mod.schedule_self_check(num_sms, torch.device("cuda"))
+    assert split_kv is not None
+    gen = torch.Generator().manual_seed(80)
+    for trial in range(300):
+        num_requests = int(torch.randint(1, 17, (1,), generator=gen))
+        next_n = int(torch.randint(1, 9, (1,), generator=gen))
+        hi = [64, 4096, 1 << 20][trial % 3]
+        lens = torch.randint(0, hi + 15, (num_requests, 1), generator=gen, dtype=torch.int32)
+        lens = lens.expand(-1, next_n).contiguous().cuda()
+        native = get_paged_mqa_logits_metadata(lens, 64, num_sms)
+        ours = draft_step_mod.paged_mqa_schedule(lens, num_sms, split_kv)
+        assert torch.equal(native, ours), (lens.tolist(), next_n)
+
+
+_SNAPSHOT_FIELDS = (
+    "kv_lens_cuda",
+    "kv_lens_cuda_2d",
+    "scheduler_metadata_buffer",
+    "scheduler_metadata_buffer_full_next_n",
+    "kv_lens_expanded_cuda",
+    "scheduler_metadata_buffer_expanded",
+    "req_idx_per_token",
+    "slot_mapping_fp8",
+    "slot_mapping_scale",
+    "gen_kv_indptr",
+    "gen_cached_token_indptr",
+)
+
+
+def _capture_loop(md, use_draft_loop, kv_init, num_accepted, pos0, gather_ids):
+    """Capture prepare_for_spec_dec + 5 steps + restore, with in-graph snapshots."""
+    batch_size = md.num_seqs
+    snaps = [{f: torch.empty_like(getattr(md, f)) for f in _SNAPSHOT_FIELDS} for _ in range(MTP_K)]
+    seq_snaps = [torch.empty((batch_size,), dtype=torch.int32, device="cuda") for _ in range(MTP_K)]
+    pos_snaps = [
+        torch.empty((batch_size,), dtype=pos0.dtype, device="cuda") for _ in range(MTP_K - 1)
+    ]
+
+    def body():
+        md.kv_lens_cuda[:batch_size].copy_(kv_init)
+        md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
+        saved = _enter_draft(md)
+        pos = pos0
+        for i in range(MTP_K):
+            pos = _step(md, use_draft_loop, i, batch_size, num_accepted, pos, gather_ids)
+            for f in _SNAPSHOT_FIELDS:
+                snaps[i][f].copy_(getattr(md, f))
+            seq_snaps[i].copy_(md.seq_lens_cuda[:batch_size])
+            if pos is not None:
+                pos_snaps[i].copy_(pos)
+        _exit_draft(md, saved)
+        md.restore_from_spec_dec()
+        md.on_update()
+
+    # Warm up eagerly (JIT and the schedule self-check), then reset host state.
+    host_seq = md._seq_lens.clone()
+    num_contexts = md.num_contexts
+    body()
+    md._seq_lens.copy_(host_seq)
+    md.num_contexts = num_contexts
+    md.use_spec_decoding = True
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        body()
+    md._seq_lens.copy_(host_seq)
+    md.num_contexts = num_contexts
+    return graph, snaps, seq_snaps, pos_snaps
+
+
+@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
+def test_cuda_graph_replays_match_native(separate_draft):
+    past = [391061, 120000]
+    batch_size = len(past)
+    cache_manager, mapping, sparse = make_dsa_cache_manager(batch_size, max(past) + 64)
+    draft_manager = None
+    try:
+        ids = list(range(batch_size))
+        add_gen_requests(cache_manager, ids, past)
+        if separate_draft:
+            draft_manager, _, _ = make_dsa_cache_manager(batch_size, max(past) + 64)
+            add_gen_requests(draft_manager, ids, past)
+        md_native = make_gen_dsa_metadata(
+            cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
+        )
+        md_draft_loop = make_gen_dsa_metadata(
+            cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
+        )
+        kv_init = md_native.kv_lens_cuda[:batch_size].clone()
+        num_accepted, pos0, gather_ids = _inputs(batch_size, past, [3, 5])
+        g_native = _capture_loop(md_native, False, kv_init, num_accepted, pos0, gather_ids)
+        g_draft_loop = _capture_loop(md_draft_loop, True, kv_init, num_accepted, pos0, gather_ids)
+        gen = torch.Generator().manual_seed(5)
+        for replay in range(12):
+            shift = int(torch.randint(-40, 40, (1,), generator=gen))
+            kv_init.copy_(torch.tensor([n + MTP_K + 1 + shift for n in past], dtype=torch.int32))
+            num_accepted.copy_(
+                torch.randint(1, MTP_K + 2, (batch_size,), generator=gen, dtype=torch.int32)
+            )
+            gather_ids.copy_(torch.arange(batch_size) * (MTP_K + 1) + num_accepted.cpu() - 1)
+            pos0.copy_(torch.cat([torch.arange(n + shift, n + shift + MTP_K + 1) for n in past]))
+            g_native[0].replay()
+            g_draft_loop[0].replay()
+            torch.cuda.synchronize()
+            for i in range(MTP_K):
+                for f in _SNAPSHOT_FIELDS:
+                    assert torch.equal(g_native[1][i][f], g_draft_loop[1][i][f]), (replay, i, f)
+                assert torch.equal(g_native[2][i], g_draft_loop[2][i]), (replay, i, "seq_lens")
+                if i < MTP_K - 1:
+                    assert torch.equal(g_native[3][i], g_draft_loop[3][i]), (replay, i, "pos")
+    finally:
+        cache_manager.shutdown()
+        if draft_manager is not None:
+            draft_manager.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py
new file mode 100644
index 000000000000..913e918cb7e1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py
@@ -0,0 +1,177 @@
+# 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.
+"""draft_loop draft embedding versus the native ``Embedding`` (bitwise)."""
+
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from _draft_loop_fixtures import bits, load_checkpoint_tensor, real_token_ids  # noqa: E402
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop.embedding import (  # noqa: E402, E501
+    DraftEmbedding,
+    draft_embedding,
+    wrap_draft_embedding,
+)
+from tensorrt_llm._torch.modules.embedding import Embedding  # noqa: E402
+
+VOCAB = 154880
+HIDDEN = 6144
+
+
+def _native_embedding(weight: torch.Tensor) -> Embedding:
+    module = Embedding(weight.shape[0], weight.shape[1], dtype=weight.dtype).cuda()
+    module.weight.data = weight
+    return module
+
+
+def _edge_ids(vocab: int, dtype: torch.dtype) -> torch.Tensor:
+    info = torch.iinfo(dtype)
+    values = [0, 1, vocab - 1, -1, vocab, vocab + 7, info.min, info.max, 12345]
+    return torch.tensor(values, dtype=dtype, device="cuda")
+
+
+@pytest.fixture(scope="module")
+def synthetic_weight():
+    gen = torch.Generator(device="cuda").manual_seed(8)
+    weight = torch.randn((VOCAB, HIDDEN), dtype=torch.bfloat16, device="cuda", generator=gen)
+    # Keep special bit patterns in the table so a bit-copy is required.
+    weight[3, :4] = torch.tensor([-0.0, float("inf"), float("-inf"), float("nan")])
+    return weight
+
+
+@pytest.fixture(scope="module")
+def real_weight():
+    weight = load_checkpoint_tensor("model.embed_tokens.weight")
+    if weight is None:
+        pytest.skip("GLM-5.2 checkpoint not available")
+    assert weight.shape == (VOCAB, HIDDEN) and weight.dtype == torch.bfloat16
+    return weight
+
+
+@pytest.mark.parametrize("num_rows", [1, 2, 5, 6, 7, 64])
+@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
+def test_matches_native_synthetic(synthetic_weight, num_rows, id_dtype):
+    native = _native_embedding(synthetic_weight)
+    gen = torch.Generator(device="cuda").manual_seed(num_rows)
+    ids = torch.randint(0, VOCAB, (num_rows,), generator=gen, device="cuda", dtype=id_dtype)
+    ids[0] = 3
+    expected = native(ids)
+    got = draft_embedding(ids, synthetic_weight)
+    assert got.shape == expected.shape and got.dtype == expected.dtype
+    assert torch.equal(bits(got), bits(expected))
+
+
+@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
+def test_out_of_range_ids_give_zero_rows(synthetic_weight, id_dtype):
+    native = _native_embedding(synthetic_weight)
+    ids = _edge_ids(VOCAB, id_dtype)
+    expected = native(ids)
+    got = draft_embedding(ids, synthetic_weight)
+    assert torch.equal(bits(got), bits(expected))
+    # +0.0, not -0.0, for every invalid row.
+    invalid = (ids < 0) | (ids >= VOCAB)
+    assert torch.all(bits(got[invalid]) == 0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+def test_other_dtypes(dtype):
+    weight = torch.randn((1000, 384), dtype=dtype, device="cuda")
+    native = _native_embedding(weight)
+    ids = torch.tensor([0, 999, 1000, -5, 17], dtype=torch.int32, device="cuda")
+    assert torch.equal(bits(draft_embedding(ids, weight)), bits(native(ids)))
+
+
+def test_real_weights_real_ids(real_weight):
+    native = _native_embedding(real_weight)
+    ids_host = real_token_ids(512)
+    assert ids_host is not None
+    for chunk in (1, 6):
+        for start in range(0, 512, 37):
+            ids = ids_host[start : start + chunk].cuda()
+            assert torch.equal(bits(draft_embedding(ids, real_weight)), bits(native(ids)))
+    # Every vocabulary row once.
+    all_ids = torch.arange(VOCAB, dtype=torch.int32, device="cuda")
+    assert torch.equal(bits(draft_embedding(all_ids, real_weight)), bits(native(all_ids)))
+
+
+def test_cuda_graph_replay(synthetic_weight):
+    """Replays read the static id buffer; each replay matches the native result."""
+    native = _native_embedding(synthetic_weight)
+    static_ids = torch.zeros((6,), dtype=torch.int32, device="cuda")
+    draft_embedding(static_ids, synthetic_weight)  # compile outside capture
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        static_out = draft_embedding(static_ids, synthetic_weight)
+    gen = torch.Generator(device="cuda").manual_seed(1)
+    for replay in range(20):
+        new_ids = torch.randint(
+            -3, VOCAB + 3, (6,), generator=gen, device="cuda", dtype=torch.int32
+        )
+        static_ids.copy_(new_ids)
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(bits(static_out), bits(native(new_ids))), f"replay {replay}"
+
+
+def test_wrapper_uses_fused_path_and_proxies_attributes(synthetic_weight, monkeypatch):
+    native = _native_embedding(synthetic_weight)
+    wrapper = wrap_draft_embedding(native)
+    assert wrap_draft_embedding(native) is wrapper
+    assert wrapper.ineligible_reason is None
+    assert wrapper.weight is native.weight
+    assert wrapper.num_embeddings == VOCAB
+
+    calls = []
+    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop.embedding as emb  # noqa: E501
+
+    real = emb.draft_embedding
+    monkeypatch.setattr(emb, "draft_embedding", lambda *a: calls.append(1) or real(*a))
+    ids = torch.tensor([5, 6], dtype=torch.int32, device="cuda")
+    assert torch.equal(bits(wrapper(ids)), bits(native(ids)))
+    assert calls == [1]
+
+
+def test_wrapper_native_fallbacks(synthetic_weight, monkeypatch):
+    native = _native_embedding(synthetic_weight)
+    wrapper = DraftEmbedding(native)
+    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop.embedding as emb  # noqa: E501
+
+    def _fail(*_):
+        raise AssertionError("fused path taken for an ineligible input")
+
+    monkeypatch.setattr(emb, "draft_embedding", _fail)
+    # 2-D ids and non-contiguous ids take the native module.
+    ids_2d = torch.tensor([[1, 2], [3, 4]], dtype=torch.int32, device="cuda")
+    assert torch.equal(bits(wrapper(ids_2d)), bits(native(ids_2d)))
+    ids_strided = torch.arange(10, dtype=torch.int32, device="cuda")[::2]
+    assert torch.equal(bits(wrapper(ids_strided)), bits(native(ids_strided)))
+    # A tensor-parallel module is ineligible as a whole.
+    native.tp_size = 2
+    assert DraftEmbedding(native).ineligible_reason == "embedding is tensor-parallel"
+    native.tp_size = 1
+
+
+def test_empty_ids(synthetic_weight):
+    ids = torch.empty((0,), dtype=torch.int32, device="cuda")
+    out = draft_embedding(ids, synthetic_weight)
+    assert out.shape == (0, HIDDEN)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py
new file mode 100644
index 000000000000..fe2499054729
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py
@@ -0,0 +1,296 @@
+# 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.
+"""BF16 near-tie tooling for the draft_tail decision-exactness study.
+
+* Native decisions: the 8-shard native oracle (``draft_tail_oracle``) on BF16 logits
+  computed by ``F.linear`` (nvjet) at a given GEMM row count M.
+* Exact-rounding decisions: native semantics (largest BF16 logit, lowest id
+  on ties) applied to logits computed in FP64 and rounded once to BF16.
+* Adversarial near-ties: real hidden rows changed in a few BF16 elements so
+  the FP64 logits of their two best tokens straddle one BF16 rounding
+  midpoint at a chosen distance (in FP32 ulps at the midpoint). There the
+  BF16 decision depends on the GEMM's FP32 accumulation order.
+"""
+
+import math
+import os
+import sys
+from typing import Dict, List, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import draft_tail_oracle as oracle  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# BF16 grid helpers (finite values only)
+# ---------------------------------------------------------------------------
+
+
+def bf16_step(x: torch.Tensor, k) -> torch.Tensor:
+    """BF16 values ``k`` grid steps above ``x`` (BF16), along the real line;
+    -0 and +0 are one point. ``k`` is an int or a broadcastable int tensor."""
+    bits = x.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF
+    mag = bits & 0x7FFF
+    ordv = torch.where((bits & 0x8000) != 0, -mag, mag) + k
+    out = torch.where(ordv >= 0, ordv, 0x8000 | (-ordv))
+    out = torch.where(out >= 0x8000, out - 0x10000, out)
+    return out.to(torch.int16).view(torch.bfloat16)
+
+
+def bf16_round_exact(x: torch.Tensor) -> torch.Tensor:
+    """Round FP64 ``x`` to BF16 once (nearest, ties to even), avoiding the
+    double rounding of FP64 -> FP32 -> BF16."""
+    x = x.double()
+    trunc = (
+        (x.float().contiguous().view(torch.int32) & -65536).view(torch.float32).to(torch.bfloat16)
+    )  # toward zero, from FP32
+    # FP32 rounding of x may have crossed a BF16 grid point: step back.
+    lo = torch.where(
+        trunc.double().abs() > x.abs(), bf16_step(trunc, torch.where(x > 0, -1, 1)), trunc
+    )
+    hi = bf16_step(lo, torch.where(x >= 0, 1, -1))
+    lo_d, hi_d = lo.double(), hi.double()
+    mid = (lo_d + hi_d) / 2
+    away = x.abs() > mid.abs()
+    tie = x.abs() == mid.abs()
+    lo_even = (lo.view(torch.int16) & 1) == 0
+    pick_hi = away | (tie & ~lo_even)
+    return torch.where(pick_hi, hi, lo)
+
+
+def bf16_ulp(x: torch.Tensor) -> torch.Tensor:
+    """BF16 ulp at |x| (FP64), for normal values."""
+    e = torch.floor(torch.log2(x.abs().double().clamp_min(1e-30)))
+    return torch.pow(2.0, e - 7)
+
+
+def fp32_ulp(x: torch.Tensor) -> torch.Tensor:
+    e = torch.floor(torch.log2(x.abs().double().clamp_min(1e-30)))
+    return torch.pow(2.0, e - 23)
+
+
+# ---------------------------------------------------------------------------
+# Decisions
+# ---------------------------------------------------------------------------
+
+
+def native_bf16_logits(
+    hidden: torch.Tensor, shards: Sequence[torch.Tensor], m: int
+) -> List[torch.Tensor]:
+    """Per-shard native BF16 logits of ``hidden`` [N, H] computed with GEMMs
+    of ``m`` rows (the last chunk is padded with copies of its first row so
+    every GEMM has exactly ``m`` rows)."""
+    out = [[] for _ in shards]
+    n = hidden.shape[0]
+    for b in range(0, n, m):
+        chunk = hidden[b : b + m]
+        live = chunk.shape[0]
+        if live < m:
+            chunk = torch.cat([chunk, chunk[:1].expand(m - live, -1)])
+        for r, w in enumerate(shards):
+            out[r].append(F.linear(chunk, w)[:live])
+    return [torch.cat(o) for o in out]
+
+
+def native_tokens(shard_logits: Sequence[torch.Tensor], chunk: int = 512) -> torch.Tensor:
+    """8-shard native oracle (``_get_local_max_and_combined`` + gather +
+    ``_get_draft_tokens_from_gathered``) over many rows."""
+    n = shard_logits[0].shape[0]
+    toks = []
+    for b in range(0, n, chunk):
+        toks.append(oracle.native_tokens_from_logits([lg[b : b + chunk] for lg in shard_logits]))
+    return torch.cat(toks)
+
+
+def bf16_key(values_bf16: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
+    """draft_tail's packed key (ordered BF16 << 32 | ~id) as int64 (no NaN here)."""
+    bits = values_bf16.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF
+    ordered = torch.where((bits & 0x8000) != 0, 0x7FFF - (bits & 0x7FFF), bits | 0x8000)
+    # Same mapping as orderedBf16 for non-NaN values, up to a constant that
+    # does not change the order (-0 == +0 below).
+    ordered = torch.where((bits & 0x7FFF) == 0, torch.full_like(bits, 0x7FFF), ordered)
+    return (ordered << 32) | (0xFFFFFFFF - ids.to(torch.int64))
+
+
+def topk_fp64(
+    hidden: torch.Tensor, head: torch.Tensor, shard_logits: Sequence[torch.Tensor], k: int = 16
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Ids of the ``k`` largest native BF16 logits per row and their FP64
+    logits (sorted descending by FP64 value)."""
+    full = torch.cat(shard_logits, dim=-1).float()
+    ids = full.topk(k, dim=-1).indices
+    w = head[ids].double()  # [N, k, H]
+    vals = torch.einsum("nh,nkh->nk", hidden.double(), w)
+    order = vals.argsort(dim=-1, descending=True)
+    return ids.gather(1, order), vals.gather(1, order)
+
+
+def exact_tokens(ids: torch.Tensor, vals: torch.Tensor) -> torch.Tensor:
+    """Native semantics on once-rounded FP64 logits of the candidates."""
+    rounded = bf16_round_exact(vals)
+    keys = bf16_key(rounded, ids)
+    return ids.gather(1, keys.argmax(dim=1, keepdim=True)).squeeze(1).to(torch.int32)
+
+
+def near_tie_stats(ids: torch.Tensor, vals: torch.Tensor) -> Dict:
+    """Per-row near-tie measures from the FP64 candidates.
+
+    gap_bf16_ulps: (top1 - top2) / BF16 ulp at top1.
+    edge_fp32_ulps: for the candidates within one BF16 ulp of the top, the
+    smallest distance of any of them to a BF16 rounding midpoint, in FP32
+    ulps; small values mean a different FP32 accumulation order can change
+    that candidate's BF16 value.
+    """
+    top = vals[:, :1]
+    ulp_b = bf16_ulp(top)
+    gap = (vals[:, 0] - vals[:, 1]) / ulp_b[:, 0]
+    rounded = bf16_round_exact(vals).double()
+    # nearest midpoint: between rounded and its neighbour on x's side
+    side = torch.where(vals >= rounded, 1, -1)
+    neigh = bf16_step(bf16_round_exact(vals), side).double()
+    mid = (rounded + neigh) / 2
+    dist = (vals - mid).abs() / fp32_ulp(vals)
+    close = (top - vals) <= ulp_b
+    edge = torch.where(close, dist, torch.full_like(dist, float("inf")))
+    return {
+        "gap_bf16_ulps": gap,
+        "edge_fp32_ulps": edge.min(dim=1).values,
+    }
+
+
+def draft_tail_full_tokens(
+    runner, hidden: torch.Tensor, head: torch.Tensor, variant: int, chunk: int = 8
+) -> torch.Tensor:
+    """draft_tail fused-kernel decision over the whole vocabulary (one tp_size 1
+    shard; equals the 8-shard global decision, see test_draft_tail_draft_tail)."""
+    out = []
+    for b in range(0, hidden.shape[0], chunk):
+        h = hidden[b : b + chunk].contiguous()
+        keys = torch.zeros(h.shape[0], dtype=torch.int64, device="cuda")
+        runner.run_fused(h, head, 0, keys_out=keys, variant=variant)
+        out.append(oracle.keys_to_tokens(keys))
+    return torch.cat(out)
+
+
+def draft_tail_logits_tokens(
+    runner, shard_logits: Sequence[torch.Tensor], chunk: int = 8
+) -> torch.Tensor:
+    """draft_tail logits mode (native BF16 logits + draft_tail), 8 shards emulated."""
+    n = shard_logits[0].shape[0]
+    out = []
+    for b in range(0, n, chunk):
+        keys = []
+        for r, lg in enumerate(shard_logits):
+            k = torch.zeros(min(chunk, n - b), dtype=torch.int64, device="cuda")
+            runner.run_from_logits(lg[b : b + chunk].contiguous(), r * lg.shape[1], keys_out=k)
+            keys.append(k)
+        out.append(oracle.keys_to_tokens(oracle.global_key(keys)))
+    return torch.cat(out)
+
+
+# ---------------------------------------------------------------------------
+# Adversarial near-ties
+# ---------------------------------------------------------------------------
+
+_STEPS = (1, 2, 4, 8, 16, 32)
+
+
+def construct_straddle(
+    hidden: torch.Tensor,
+    head: torch.Tensor,
+    tok_hi: torch.Tensor,
+    tok_lo: torch.Tensor,
+    tau_fp32_ulps: torch.Tensor,
+    iters: int = 3000,
+) -> Dict[str, torch.Tensor]:
+    """Move rows ``hidden`` [B, H] (BF16) by BF16 grid steps of single
+    elements until the FP64 logit of ``tok_hi`` sits ``tau`` FP32 ulps above
+    a BF16 rounding midpoint m and that of ``tok_lo`` ``tau`` below it (tau
+    per row, times a random factor in [0.5, 1]); m is the midpoint nearest
+    to the mean of the two starting logits. Greedy: each iteration applies,
+    per row, the single (element, step) change that most reduces the
+    distance to the targets.
+    """
+    b, hdim = hidden.shape
+    h = hidden.clone()
+    w_hi = head[tok_hi].double()
+    w_lo = head[tok_lo].double()
+
+    def logits(hh):
+        hd = hh.double()
+        return (hd * w_hi).sum(-1), (hd * w_lo).sum(-1)
+
+    l_hi, l_lo = logits(h)
+    mean = (l_hi + l_lo) / 2
+    lo_grid = bf16_round_exact(mean)
+    lo_grid = torch.where(lo_grid.double() > mean, bf16_step(lo_grid, -1), lo_grid)
+    mid = (lo_grid.double() + bf16_step(lo_grid, 1).double()) / 2
+    gen = torch.Generator(device=hidden.device).manual_seed(7)
+    jitter = 0.5 + 0.5 * torch.rand(2, b, generator=gen, device=hidden.device, dtype=torch.float64)
+    ulp = fp32_ulp(mid)
+    t_hi = mid + tau_fp32_ulps.double() * ulp * jitter[0]
+    t_lo = mid - tau_fp32_ulps.double() * ulp * jitter[1]
+    rows = torch.arange(b, device=hidden.device)
+    for _ in range(iters):
+        r_hi, r_lo = t_hi - l_hi, t_lo - l_lo
+        cur = torch.maximum(r_hi.abs(), r_lo.abs())
+        done = cur <= 0.05 * tau_fp32_ulps.double() * ulp
+        if bool(done.all()):
+            break
+        best_cost = cur.clone()
+        best_idx = torch.full((b,), -1, dtype=torch.int64, device=hidden.device)
+        best_step = torch.zeros(b, dtype=torch.int64, device=hidden.device)
+        hd = h.double()
+        for s in _STEPS:
+            for sign in (1, -1):
+                moved = bf16_step(h, sign * s).double()
+                delta = moved - hd
+                c = torch.maximum(
+                    (r_hi[:, None] - delta * w_hi).abs(), (r_lo[:, None] - delta * w_lo).abs()
+                )
+                c = torch.where(torch.isfinite(moved), c, torch.full_like(c, float("inf")))
+                val, idx = c.min(dim=1)
+                better = val < best_cost
+                best_cost = torch.where(better, val, best_cost)
+                best_idx = torch.where(better, idx, best_idx)
+                best_step = torch.where(better, torch.full_like(best_step, sign * s), best_step)
+        move = (best_idx >= 0) & ~done
+        if not bool(move.any()):
+            break
+        sel = rows[move]
+        col = best_idx[move]
+        h[sel, col] = bf16_step(h[sel, col], best_step[move])
+        l_hi, l_lo = logits(h)
+    return {
+        "hidden": h,
+        "mid": mid,
+        "l_hi": l_hi,
+        "l_lo": l_lo,
+        "dist_hi_fp32_ulps": (l_hi - mid) / ulp,
+        "dist_lo_fp32_ulps": (l_lo - mid) / ulp,
+        "changed_elements": (h != hidden).sum(-1),
+    }
+
+
+def wilson_upper(k: int, n: int, z: float = 1.96) -> float:
+    if n == 0:
+        return float("nan")
+    p = k / n
+    den = 1 + z * z / n
+    c = p + z * z / (2 * n)
+    r = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
+    return (c + r) / den
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py
new file mode 100644
index 000000000000..a844bdfbd158
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py
@@ -0,0 +1,148 @@
+# 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.
+"""Native oracle and fixtures shared by the draft_tail tests and benchmarks.
+
+The oracle runs the native TensorRT-LLM code itself: ``F.linear`` (what
+``UnquantizedLinearMethod`` runs for the BF16 LM head) per vocab shard, then
+``SpecWorkerBase._get_local_max_and_combined`` per rank, the all-gather along
+the last dimension (a concatenation in rank order, as ``allgather(dim=-1)``
+returns) and the compiled ``SpecWorkerBase._get_draft_tokens_from_gathered``.
+On one GPU the eight ranks are emulated shard by shard.
+"""
+
+import json
+import os
+from types import SimpleNamespace
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from utils.llm_data import llm_models_root
+
+CKPT = os.environ.get(
+    "GLM_MOE_DSA_DRAFT_TAIL_CKPT",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+VOCAB = 154880
+HIDDEN = 6144
+TP = 8
+SHARD = VOCAB // TP
+
+
+def checkpoint_available() -> bool:
+    return os.path.exists(os.path.join(CKPT, "model.safetensors.index.json"))
+
+
+def _tensor_file(name: str) -> str:
+    index = json.load(open(os.path.join(CKPT, "model.safetensors.index.json")))
+    return os.path.join(CKPT, index["weight_map"][name])
+
+
+def load_lm_head_rows(begin: int, end: int, device="cuda") -> torch.Tensor:
+    from safetensors import safe_open
+
+    with safe_open(_tensor_file("lm_head.weight"), framework="pt", device="cpu") as f:
+        return f.get_slice("lm_head.weight")[begin:end].to(device).contiguous()
+
+
+def load_lm_head_shards(tp: int = TP, device="cuda") -> List[torch.Tensor]:
+    rows = VOCAB // tp
+    full = load_lm_head_rows(0, VOCAB, device=device)
+    return [full[r * rows : (r + 1) * rows].contiguous() for r in range(tp)]
+
+
+def load_tensor(name: str, device="cuda") -> torch.Tensor:
+    from safetensors import safe_open
+
+    with safe_open(_tensor_file(name), framework="pt", device="cpu") as f:
+        return f.get_tensor(name).to(device)
+
+
+def mtp_norm_weight(device="cuda") -> torch.Tensor:
+    return load_tensor("model.layers.78.shared_head.norm.weight", device)
+
+
+def rms_normalize(x: torch.Tensor, gamma: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
+    xf = x.float()
+    y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
+    return (y * gamma.float()).to(torch.bfloat16)
+
+
+def synthetic_hidden(
+    rows: int,
+    gamma: torch.Tensor,
+    generator: torch.Generator,
+    peaked_from: Optional[torch.Tensor] = None,
+    peak_scale: float = 0.0,
+) -> torch.Tensor:
+    """Normed hidden rows as the MTP shared-head RMSNorm emits them.
+
+    ``peaked_from`` (rows of the LM head) adds a component along chosen token
+    directions so the logits have a clear winner, like real decoding states;
+    without it the logits are flat and near-ties are frequent.
+    """
+    x = torch.randn(rows, HIDDEN, generator=generator, device="cuda", dtype=torch.float32)
+    if peaked_from is not None:
+        d = peaked_from.float()
+        d = d / d.norm(dim=-1, keepdim=True)
+        x = x + peak_scale * d * x.norm(dim=-1, keepdim=True)
+    return rms_normalize(x, gamma)
+
+
+class _Rank(SimpleNamespace):
+    pass
+
+
+def native_combined(logits: torch.Tensor, tp_rank: int) -> torch.Tensor:
+    from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
+
+    return SpecWorkerBase._get_local_max_and_combined(
+        None, logits, mapping_lm_tp=_Rank(tp_rank=tp_rank)
+    )
+
+
+def native_from_gathered(gathered: torch.Tensor) -> torch.Tensor:
+    from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
+
+    return SpecWorkerBase._get_draft_tokens_from_gathered(None, gathered)
+
+
+def native_tokens_from_logits(shard_logits: List[torch.Tensor]) -> torch.Tensor:
+    """Native greedy TP draft tokens from per-rank BF16 vocab-shard logits."""
+    combined = [native_combined(lg, r) for r, lg in enumerate(shard_logits)]
+    return native_from_gathered(torch.cat(combined, dim=-1))
+
+
+def native_tokens(hidden: torch.Tensor, shards: List[torch.Tensor]) -> torch.Tensor:
+    return native_tokens_from_logits([F.linear(hidden, w) for w in shards])
+
+
+def keys_to_tokens(keys: torch.Tensor) -> torch.Tensor:
+    low = keys & 0xFFFFFFFF
+    return (0xFFFFFFFF - low).to(torch.int32)
+
+
+def global_key(keys_per_rank: List[torch.Tensor]) -> torch.Tensor:
+    """Unsigned max of packed keys (all keys are < 2**63 here)."""
+    stacked = torch.stack(keys_per_rank, dim=0)
+    assert bool((stacked >= 0).all()), "keys must fit in int64 as non-negative"
+    return stacked.max(dim=0).values
+
+
+def fp64_logits(
+    hidden: torch.Tensor, rows: torch.Tensor, weight_full_rows: torch.Tensor
+) -> torch.Tensor:
+    """Exact-ish logits of selected vocab rows (FP64 accumulation)."""
+    return hidden.double() @ weight_full_rows.double().t()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py
new file mode 100644
index 000000000000..da0a918a2417
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py
@@ -0,0 +1,736 @@
+# 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.
+"""Single-GPU correctness of the draft_tail against the native oracle.
+
+The eight TP ranks of GLM-5.2 are emulated shard by shard on one GPU: each
+shard runs the draft_tail kernel without exchange (``tp_size=1`` runner, explicit
+vocab offset) and returns its packed key; the global decision is the unsigned
+max of the eight keys, exactly what the exchange computes on every rank. The
+cross-rank exchange itself needs 8 GPUs and is not tested here.
+"""
+
+import json
+import os
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import draft_tail_nearties as nt  # noqa: E402
+import draft_tail_oracle as oracle  # noqa: E402
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+NUM_VARIANTS = 15
+FIRST_MMA_VARIANT = 7
+NATIVE_ORDER_VARIANTS = (11, 12, 13, 14)
+
+
+def _family(variant):
+    """Kernels of one family accumulate a row in the same order: CUDA-core
+    (0-6), tensor-core permuted K (7-10), tensor-core native K order
+    (11-14)."""
+    if variant < FIRST_MMA_VARIANT:
+        return 0
+    return 11 if variant in NATIVE_ORDER_VARIANTS else FIRST_MMA_VARIANT
+
+
+def _draft_tail():
+    import tensorrt_llm  # noqa: F401  (loads the trtllm torch ops)
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail import (  # noqa: E501
+        draft_tail,
+    )
+
+    return draft_tail
+
+
+@pytest.fixture(scope="module")
+def runner():
+    dt = _draft_tail()
+    r = dt.DraftTailRunner(0, 1)
+    assert r.prepare()
+    return r
+
+
+@pytest.fixture(scope="module")
+def shards():
+    if not oracle.checkpoint_available():
+        pytest.skip("GLM-5.2 checkpoint not mounted")
+    return oracle.load_lm_head_shards()
+
+
+@pytest.fixture(scope="module")
+def gamma():
+    if not oracle.checkpoint_available():
+        pytest.skip("GLM-5.2 checkpoint not mounted")
+    return oracle.mtp_norm_weight()
+
+
+class _Head:
+    """Minimal stand-in exposing ``weight`` like LMHead."""
+
+    def __init__(self, weight):
+        self.weight = weight
+
+
+@pytest.fixture(scope="module")
+def full_head(shards):
+    """The whole vocabulary as one shard: with a tp_size 1 runner the draft_tail
+    decision over it is the global decision of the 8 shards (the packed-key
+    max is associative and every row's dot product is computed alike)."""
+    return _Head(torch.cat(shards, dim=0))
+
+
+def _full_keys(runner, hidden, head, variant):
+    keys = torch.zeros(hidden.shape[0], dtype=torch.int64, device="cuda")
+    runner.run_fused(hidden, head.weight, 0, keys_out=keys, variant=variant)
+    return keys
+
+
+def _e2e_near_tie_rows():
+    """Draft LM-head input rows of the real BF16 near-ties found end to end.
+
+    ``GLM_MOE_DSA_DRAFT_TAIL_NEAR_TIE_DIR`` is a directory with the native run's captured draft LM-head
+    inputs (with each draft-loop call's first row) and ``tie_analysis.json``.
+    """
+    d = os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_NEAR_TIE_DIR")
+    if not d:
+        return []
+    ties = [t for t in json.load(open(os.path.join(d, "tie_analysis.json"))) if t.get("near_tie")]
+    starts = torch.load(os.path.join(d, "native_eager_draft.pt"))["start_rows"]
+    rows = torch.load(os.path.join(d, "native_eager_hidden.pt"))["rows"]
+    return [
+        (f"e2e call {t['call']} step {t['step']}", rows[starts[t["call"]] + t["step"]].to("cuda"))
+        for t in ties
+    ]
+
+
+def _family_split_rows(runner, head, gamma, want=4):
+    """Constructed BF16 near-ties: seeded flat MTP-normed rows whose winner
+    the CUDA-core family (variant 4) and the tensor-core family (variant 10)
+    round differently. ``GLM_MOE_DSA_DRAFT_TAIL_FAMILY_SPLIT_BATCHES`` bounds the search."""
+    batches = int(os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_FAMILY_SPLIT_BATCHES", "512"))
+    gen = torch.Generator(device="cuda").manual_seed(2026)
+    found = []
+    for b in range(batches):
+        h = oracle.synthetic_hidden(8, gamma, gen)
+        a = oracle.keys_to_tokens(_full_keys(runner, h, head, 4))
+        c = oracle.keys_to_tokens(_full_keys(runner, h, head, 10))
+        for i in torch.nonzero(a != c).flatten().tolist():
+            found.append((f"family split, search batch {b} row {i}", h[i]))
+        if len(found) >= want:
+            break
+    return found[:want]
+
+
+@pytest.fixture(scope="module")
+def live_rows(runner, full_head, gamma):
+    """(label, hidden row) pairs used as live rows: real and constructed
+    BF16 near-ties on which the two kernel families disagree, plus random
+    flat and peaked rows."""
+    gen = torch.Generator(device="cuda").manual_seed(31)
+    rows = _e2e_near_tie_rows()
+    rows += _family_split_rows(runner, full_head, gamma)
+    rows.append(("flat", oracle.synthetic_hidden(1, gamma, gen)[0]))
+    peak = full_head.weight[123456:123457]
+    rows.append(("peaked", oracle.synthetic_hidden(1, gamma, gen, peak, 0.3)[0]))
+    return rows
+
+
+def _draft_tail_keys_fused(runner, hidden, shards, **kw):
+    keys = []
+    for r, w in enumerate(shards):
+        k = torch.zeros(hidden.shape[0], dtype=torch.int64, device="cuda")
+        runner.run_fused(hidden, w, r * w.shape[0], keys_out=k, **kw)
+        keys.append(k)
+    return keys
+
+
+def _draft_tail_keys_logits(runner, shard_logits):
+    keys = []
+    for r, lg in enumerate(shard_logits):
+        k = torch.zeros(lg.shape[0], dtype=torch.int64, device="cuda")
+        runner.run_from_logits(lg, r * lg.shape[1], keys_out=k)
+        keys.append(k)
+    return keys
+
+
+# ---------------------------------------------------------------------------
+# Selection semantics on given BF16 logits (logits kernel): exact equality.
+# ---------------------------------------------------------------------------
+
+
+def _random_logits(rows, gen, tp=oracle.TP, shard=oracle.SHARD):
+    return [
+        (torch.randn(rows, shard, generator=gen, device="cuda") * 3).to(torch.bfloat16)
+        for _ in range(tp)
+    ]
+
+
+@pytest.mark.parametrize("rows", [1, 2, 6, 8])
+def test_logits_selection_matches_native_random(runner, rows):
+    gen = torch.Generator(device="cuda").manual_seed(rows)
+    for _ in range(20):
+        lg = _random_logits(rows, gen)
+        ref = oracle.native_tokens_from_logits(lg)
+        got = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, lg)))
+        torch.testing.assert_close(got, ref, rtol=0, atol=0)
+
+
+def test_logits_selection_ties_and_specials(runner):
+    """Cross-shard and in-shard ties, -0/+0, infinities, NaN."""
+    gen = torch.Generator(device="cuda").manual_seed(7)
+    rows = 8
+    cases = []
+    lg = _random_logits(rows, gen)
+    # row 0: same max value in shards 5 and 2 -> lower global id (shard 2)
+    for r, c in ((5, 100), (2, 19000)):
+        lg[r][0, c] = 1000.0
+    # row 1: tie inside one shard at two columns -> lower column
+    lg[3][1, 777] = 1000.0
+    lg[3][1, 12] = 1000.0
+    # row 2: every logit equal -> global id 0
+    for r in range(oracle.TP):
+        lg[r][2].fill_(-1.5)
+    # row 3: max is +0 vs -0 at different ids -> treated equal -> lower id
+    for r in range(oracle.TP):
+        lg[r][3].fill_(-5.0)
+    lg[6][3, 3] = 0.0
+    lg[1][3, 9] = -0.0
+    # row 4: +inf in shard 4, large finite elsewhere
+    lg[4][4, 55] = float("inf")
+    lg[0][4, 1] = 3e38
+    # row 5: all -inf except one -inf-adjacent value
+    for r in range(oracle.TP):
+        lg[r][5].fill_(float("-inf"))
+    lg[7][5, 19359] = -3e38
+    # row 6: tie across the last column of shard 0 and first of shard 1
+    lg[0][6, oracle.SHARD - 1] = 500.0
+    lg[1][6, 0] = 500.0
+    # row 7: NaN in shard 3 (native: NaN ranks above every number)
+    lg[3][7, 4242] = float("nan")
+    cases.append(lg)
+    for lg in cases:
+        ref = oracle.native_tokens_from_logits(lg)
+        got = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, lg)))
+        torch.testing.assert_close(got, ref, rtol=0, atol=0)
+    assert got[0].item() == 2 * oracle.SHARD + 19000
+    assert got[1].item() == 3 * oracle.SHARD + 12
+    assert got[2].item() == 0
+
+
+def test_logits_selection_padded_vocab(runner):
+    """Vocab not divisible by TP: the last shard carries zero padding rows.
+
+    The native path includes the padding columns of the last shard in its
+    local max (they produce logit 0); draft_tail must make the same decision.
+    """
+    gen = torch.Generator(device="cuda").manual_seed(11)
+    tp, shard = 8, 1000
+    vocab = tp * shard - 37
+    for rows in (1, 3):
+        lg = [
+            (torch.randn(rows, shard, generator=gen, device="cuda") - 4).to(torch.bfloat16)
+            for _ in range(tp)
+        ]
+        lg[-1][:, vocab - (tp - 1) * shard :] = 0.0  # zero weight rows
+        ref = oracle.native_tokens_from_logits(lg)
+        got = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, lg)))
+        torch.testing.assert_close(got, ref, rtol=0, atol=0)
+
+
+# ---------------------------------------------------------------------------
+# Fused GEMV path on the real GLM-5.2 LM head.
+# ---------------------------------------------------------------------------
+
+
+def _assert_near_ties(reports):
+    """A token mismatch is acceptable only when the two winners' exact logits
+    are within one BF16 ulp of the maximum: both kernels then round FP32 sums
+    of different accumulation order across the same BF16 boundary."""
+    for rep in reports:
+        top = max(abs(v) for v in rep["fp64"])
+        ulp = 2.0 ** (int(torch.tensor(top).log2().floor().item()) - 7)
+        assert abs(rep["fp64_gap"]) <= ulp, rep
+
+
+def _decision_report(hidden, shards, got, ref):
+    """For each mismatch, the FP64 logit gap between the two winners."""
+    full = torch.cat(shards, dim=0)
+    out = []
+    for i in torch.nonzero(got != ref).flatten().tolist():
+        ids = torch.tensor([ref[i].item(), got[i].item()], device="cuda")
+        exact = hidden[i : i + 1].double() @ full[ids].double().t()
+        native_bf16 = F.linear(hidden[i : i + 1], full[ids])
+        out.append(
+            {
+                "row": i,
+                "native": ref[i].item(),
+                "draft_tail": got[i].item(),
+                "fp64_gap": (exact[0, 0] - exact[0, 1]).item(),
+                "fp64": exact[0].tolist(),
+                "native_bf16": native_bf16[0].float().tolist(),
+            }
+        )
+    return out
+
+
+@pytest.mark.parametrize("peak", [0.0, 0.3, 1.0])
+def test_fused_matches_native_real_lm_head(runner, shards, gamma, peak):
+    """Exact token equality vs native on the real head; any mismatch must be
+    a BF16 rounding near-tie (FP64 gap below one BF16 ulp of the max)."""
+    gen = torch.Generator(device="cuda").manual_seed(int(peak * 10) + 1)
+    full_rows = torch.cat(shards, dim=0)
+    rows_total, mismatches, reports = 0, 0, []
+    for it in range(16):
+        rows = 8
+        target = torch.randint(0, oracle.VOCAB, (rows,), generator=gen, device="cuda")
+        hidden = oracle.synthetic_hidden(
+            rows, gamma, gen, peaked_from=full_rows[target] if peak > 0 else None, peak_scale=peak
+        )
+        ref = oracle.native_tokens(hidden, shards)
+        got = oracle.keys_to_tokens(
+            oracle.global_key(_draft_tail_keys_fused(runner, hidden, shards))
+        )
+        rows_total += rows
+        bad = got != ref
+        if bool(bad.any()):
+            mismatches += int(bad.sum())
+            reports += _decision_report(hidden, shards, got, ref)
+    print(f"[draft_tail] peak={peak}: {mismatches}/{rows_total} token mismatches; {reports}")
+    _assert_near_ties(reports)
+
+
+@pytest.mark.parametrize("variant", list(range(NUM_VARIANTS)))
+def test_fused_variants_bitwise_identical(runner, shards, gamma, variant):
+    """Within a kernel family (CUDA-core 0-6, tensor-core 7-10, tensor-core
+    native order 11-14) every shape, CTA count and L2 keep hint accumulates
+    a row in the same order, so the keys are identical."""
+    base_variant = _family(variant)
+    gen = torch.Generator(device="cuda").manual_seed(3)
+    for rows in (1, 2, 5, 8):
+        hidden = oracle.synthetic_hidden(rows, gamma, gen)
+        base = _draft_tail_keys_fused(runner, hidden, shards[:2], variant=base_variant)
+        for num_ctas, keep in ((0, 0), (37, 0), (300, 0), (0, 500), (304, 1000)):
+            got = _draft_tail_keys_fused(
+                runner, hidden, shards[:2], variant=variant, num_ctas=num_ctas, keep_permille=keep
+            )
+            for a, b in zip(base, got):
+                torch.testing.assert_close(a, b, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("variant", [0, FIRST_MMA_VARIANT, 11, 13])
+@pytest.mark.parametrize("rows", [1, 3, 6, 8])
+def test_fused_families_match_native(runner, shards, gamma, variant, rows):
+    """Every kernel family makes the native decision on the real head for
+    every row count (mismatches only at BF16 near-ties; none at all for the
+    native-order family)."""
+    gen = torch.Generator(device="cuda").manual_seed(100 + rows)
+    reports = []
+    for it in range(8):
+        peak = (0.0, 0.3, 1.0)[it % 3]
+        hidden = oracle.synthetic_hidden(
+            rows, gamma, gen, torch.cat(shards[:1])[it * 8 : it * 8 + rows] if peak else None, peak
+        )
+        ref = oracle.native_tokens(hidden, shards)
+        got = oracle.keys_to_tokens(
+            oracle.global_key(_draft_tail_keys_fused(runner, hidden, shards, variant=variant))
+        )
+        reports += _decision_report(hidden, shards, got, ref)
+    print(f"[draft_tail] variant={variant} rows={rows}: mismatches {reports}")
+    _assert_near_ties(reports)
+    if variant in NATIVE_ORDER_VARIANTS:
+        assert not reports, reports
+
+
+@pytest.mark.parametrize("variant", [4, 10, 11, 13])
+def test_fused_family_rows_independent(runner, shards, gamma, variant):
+    """Within one kernel, row r's key does not depend on the other rows of
+    the call or on the row count."""
+    gen = torch.Generator(device="cuda").manual_seed(5)
+    live = oracle.synthetic_hidden(3, gamma, gen)
+    w = shards[4]
+    single = [
+        _draft_tail_keys_fused(runner, live[i : i + 1], [w], variant=variant)[0] for i in range(3)
+    ]
+    for rows in range(3, 9):
+        padded = torch.zeros(rows, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda")
+        padded[:3] = live
+        if rows > 3:
+            padded[3:] = oracle.synthetic_hidden(rows - 3, gamma, gen) * 7
+        keys = _draft_tail_keys_fused(runner, padded, [w], variant=variant)[0]
+        for i in range(3):
+            assert keys[i].item() == single[i].item()
+
+
+def test_select_live_row_key_independent_of_row_count(runner, full_head, gamma, live_rows):
+    """The shipped dispatch (``runner.select``, as the model adapter calls
+    it) gives a live row the same packed key, hence the same draft token,
+    for every row count R = 1..8 and position, so CUDA-graph padding and
+    batch composition cannot change a live row's decision. The live rows
+    include real and constructed BF16 near-ties on which the two kernel
+    families round differently: a dispatch that picked the family by row
+    count would fail here."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
+        DraftTailOptions,
+    )
+
+    assert runner.options == DraftTailOptions()  # the shipped defaults
+    gen = torch.Generator(device="cuda").manual_seed(17)
+    pad = torch.cat(
+        [
+            oracle.synthetic_hidden(4, gamma, gen) * 7,
+            torch.zeros(4, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda"),
+        ]
+    )
+    split = 0
+    for label, row in live_rows:
+        v4 = int(_full_keys(runner, row[None], full_head, 4)[0])
+        v10 = int(_full_keys(runner, row[None], full_head, 10)[0])
+        vdef = int(_full_keys(runner, row[None], full_head, DraftTailOptions().variant)[0])
+        logits = F.linear(row[None], full_head.weight)
+        native_key = int(
+            nt.bf16_key(logits, torch.arange(logits.shape[1], device="cuda")[None]).max()
+        )
+        split += int(v4 != v10)
+        seen = set()
+        for rows in range(1, 9):
+            for pos in sorted({0, rows - 1}):
+                batch = pad[:rows].clone()
+                batch[pos] = row
+                keys = torch.zeros(rows, dtype=torch.int64, device="cuda")
+                tokens = runner.select(batch, full_head, keys_out=keys)
+                assert int(tokens[pos]) == int(oracle.keys_to_tokens(keys[pos : pos + 1])[0])
+                seen.add(int(keys[pos]))
+        print(
+            f"[draft_tail-rowcount] {label}: select keys over R=1..8 {seen}; "
+            f"v4 token {int(oracle.keys_to_tokens(torch.tensor([v4])))} "
+            f"v10 token {int(oracle.keys_to_tokens(torch.tensor([v10])))}"
+        )
+        assert len(seen) == 1, (label, seen)
+        # The shipped kernel for every R is the default variant, and its key
+        # (winning BF16 value and id) is native's.
+        assert seen == {vdef}, (label, seen, vdef)
+        assert vdef == native_key, (label, vdef, native_key)
+    print(
+        f"[draft_tail-rowcount] live rows {len(live_rows)}, rows where the two "
+        f"families disagree: {split}"
+    )
+
+
+def test_row_count_characterization(runner, full_head, shards, gamma, live_rows):
+    """For the record (no assertion on native): the native decision on each
+    live row when it shares the LM-head GEMM with M - 1 other rows (8-shard
+    native oracle; cuBLAS picks its kernel by M), next to both draft_tail kernel
+    families and the FP64 logits of every candidate."""
+    gen = torch.Generator(device="cuda").manual_seed(23)
+    pad = oracle.synthetic_hidden(64, gamma, gen)
+    full = full_head.weight
+    for label, row in live_rows:
+        native = {}
+        bf16 = {}
+        for m in (1, 2, 3, 4, 5, 6, 7, 8, 16, 32, 64):
+            batch = pad[:m].clone()
+            batch[0] = row
+            logits = [F.linear(batch, w) for w in shards]
+            tok = int(oracle.native_tokens_from_logits(logits)[0])
+            native[m] = tok
+            bf16[m] = logits[tok // oracle.SHARD][0, tok % oracle.SHARD].item()
+        fam = {
+            v: int(oracle.keys_to_tokens(_full_keys(runner, row[None], full_head, v))[0])
+            for v in (4, 10)
+        }
+        cands = sorted(set(native.values()) | set(fam.values()))
+        exact = (row[None].double() @ full[cands].double().t())[0].tolist()
+        print(
+            "[draft_tail-rowcount-native] "
+            + json.dumps(
+                {
+                    "row": label,
+                    "native_token_by_M": native,
+                    "native_bf16_winner_logit_by_M": bf16,
+                    "draft_tail_v4": fam[4],
+                    "draft_tail_v10": fam[10],
+                    "fp64_logits": dict(zip(cands, exact)),
+                }
+            )
+        )
+
+
+def test_fused_cuda_graph_replay(runner, shards, gamma):
+    """Scratch keys and the arrival counter are reset by every launch:
+    repeated replays with changing inputs keep matching native."""
+    gen = torch.Generator(device="cuda").manual_seed(9)
+    rows = 2
+    static_h = torch.zeros(rows, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda")
+    static_keys = [torch.zeros(rows, dtype=torch.int64, device="cuda") for _ in shards]
+
+    def step():
+        for r, w in enumerate(shards):
+            runner.run_fused(static_h, w, r * w.shape[0], keys_out=static_keys[r])
+
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        step()
+    torch.cuda.current_stream().wait_stream(s)
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        step()
+    for it in range(50):
+        static_h.copy_(oracle.synthetic_hidden(rows, gamma, gen))
+        g.replay()
+        torch.cuda.synchronize()
+        got = oracle.keys_to_tokens(oracle.global_key(static_keys))
+        eager = oracle.keys_to_tokens(
+            oracle.global_key(_draft_tail_keys_fused(runner, static_h, shards))
+        )
+        torch.testing.assert_close(got, eager, rtol=0, atol=0)
+        ref = oracle.native_tokens(static_h, shards)
+        _assert_near_ties(_decision_report(static_h, shards, got, ref))
+    assert runner.scratch[:8].abs().sum().item() == 0
+    assert runner.scratch[8].item() & 0xFFFFFFFF == 0
+
+
+def test_runner_tp1_tokens_match_native_full_vocab(runner, shards, gamma):
+    """tp_size == 1 (no exchange): one shard is the whole local vocab; the
+    native path is a plain argmax over the logits."""
+    gen = torch.Generator(device="cuda").manual_seed(13)
+    w = shards[1]
+    for rows in (1, 4, 8):
+        h = oracle.synthetic_hidden(rows, gamma, gen, w[:rows], 1.0)
+        tokens = runner.run_fused(h, w, 0)
+        ref = torch.argmax(F.linear(h, w), dim=-1).to(torch.int32)
+        torch.testing.assert_close(tokens, ref, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    not os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE"),
+    reason="GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE (captured draft LM-head inputs) not set",
+)
+@pytest.mark.parametrize("variant", [0, FIRST_MMA_VARIANT, 11, 13])
+def test_real_hidden_states_exact_tokens(runner, shards, variant):
+    """Real MTP draft LM-head inputs captured from a native GLM-5.2 run:
+    token-exact vs native for every row, through
+    the fused kernel (both families) and the logits kernel; any mismatch must
+    be a BF16 near-tie."""
+    rows = torch.load(os.environ["GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE"])["rows"].to("cuda")
+    assert rows.dtype == torch.bfloat16 and rows.shape[1] == oracle.HIDDEN
+    total = fused_bad = logits_bad = 0
+    reports = []
+    for i in range(0, rows.shape[0], 8):
+        h = rows[i : i + 8].contiguous()
+        logits = [F.linear(h, w) for w in shards]
+        ref = oracle.native_tokens_from_logits(logits)
+        got_f = oracle.keys_to_tokens(
+            oracle.global_key(_draft_tail_keys_fused(runner, h, shards, variant=variant))
+        )
+        got_l = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, logits)))
+        # The logits kernel consumes the native BF16 logits: always exact.
+        torch.testing.assert_close(got_l, ref, rtol=0, atol=0)
+        total += h.shape[0]
+        fused_bad += int((got_f != ref).sum())
+        logits_bad += int((got_l != ref).sum())
+        reports += _decision_report(h, shards, got_f, ref)
+    print(
+        f"[draft_tail] real hidden rows={total} variant={variant}: fused "
+        f"mismatches={fused_bad}, logits mismatches={logits_bad}; "
+        f"{reports}"
+    )
+    _assert_near_ties(reports)
+    if variant in NATIVE_ORDER_VARIANTS:
+        assert fused_bad == 0
+
+
+def _exactness_rows(full_head, gamma):
+    """Rows for the bit-level comparison: real draft LM-head inputs
+    (``GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE``, first ``GLM_MOE_DSA_DRAFT_TAIL_EXACT_ROWS``),
+    knife-edge rows built from the real rows with the smallest top-2 gaps (their two best logits
+    straddle a BF16 rounding midpoint within a fraction of an FP32 ulp, so
+    any difference in their FP32 sums changes the BF16 values), and
+    synthetic flat / peaked MTP-normed rows."""
+    gen = torch.Generator(device="cuda").manual_seed(41)
+    rows = [oracle.synthetic_hidden(16, gamma, gen)]
+    rows.append(oracle.synthetic_hidden(16, gamma, gen, full_head.weight[1000:1016], 0.5))
+    path = os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE")
+    if path:
+        n = int(os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_EXACT_ROWS", "256"))
+        real = torch.load(path)["rows"][:n].to("cuda")
+        rows.append(real)
+        shards = [
+            full_head.weight[r * oracle.SHARD : (r + 1) * oracle.SHARD] for r in range(oracle.TP)
+        ]
+        logits = nt.native_bf16_logits(real, shards, 8)
+        ids, vals = nt.topk_fp64(real, full_head.weight, logits)
+        gaps = (vals[:, 0] - vals[:, 1]) / nt.bf16_ulp(vals[:, 0])
+        base = gaps.argsort()[:16]
+        lo = torch.minimum(ids[base, 0], ids[base, 1])
+        hi = torch.maximum(ids[base, 0], ids[base, 1])
+        for up, down in ((lo, hi), (hi, lo)):
+            res = nt.construct_straddle(
+                real[base], full_head.weight, up, down, torch.full((16,), 0.5, device="cuda")
+            )
+            rows.append(res["hidden"])
+    return torch.cat(rows)
+
+
+@pytest.mark.parametrize("variant", list(NATIVE_ORDER_VARIANTS))
+def test_native_order_logits_bitwise_equal_native(runner, shards, full_head, gamma, variant):
+    """The native-order fused variants compute every BF16 logit bit for bit
+    like the native LM-head GEMM (``F.linear``), for every row count, so
+    their draft tokens are native ones. Variant 10 (permuted K order) is
+    measured on the same rows for comparison."""
+    rows = _exactness_rows(full_head, gamma)
+    diffs = {variant: 0, 10: 0}
+    total = 0
+    for i in range(0, rows.shape[0], 8):
+        for r in (1, 3, 8):
+            h = rows[i : i + r].contiguous()
+            for s, w in enumerate(shards):
+                ref = F.linear(h, w)
+                for v in diffs:
+                    out = torch.empty_like(ref)
+                    runner.run_fused(h, w, s * w.shape[0], variant=v, logits_out=out)
+                    diffs[v] += int((out.view(torch.int16) != ref.view(torch.int16)).sum())
+                total += ref.numel()
+    print(
+        f"[draft_tail-exact] rows={rows.shape[0]} logits compared={total} "
+        f"BF16 logits != native: {diffs}"
+    )
+    assert diffs[variant] == 0
+
+
+def _loaded_cublas_versions():
+    """cuBLAS / cuBLASLt libraries mapped into this process and their
+    versions (the native LM-head GEMM's kernel is chosen by these)."""
+    import ctypes
+
+    paths = set()
+    with open("/proc/self/maps") as f:
+        for line in f:
+            p = line.split()[-1]
+            if "libcublas" in os.path.basename(p) and ".so" in p:
+                paths.add(p)
+    out = {}
+    for p in sorted(paths):
+        lib = ctypes.CDLL(p)
+        fn = "cublasLtGetProperty" if "cublasLt" in p else "cublasGetProperty"
+        try:
+            ver = []
+            for prop in range(3):  # MAJOR_VERSION, MINOR_VERSION, PATCH_LEVEL
+                v = ctypes.c_int(-1)
+                getattr(lib, fn)(prop, ctypes.byref(v))
+                ver.append(v.value)
+            out[p] = ".".join(str(v) for v in ver)
+        except AttributeError:
+            out[p] = None
+    return out
+
+
+def test_native_lm_head_kernel_record(shards):
+    """For the record (no assertion): the GEMM kernel this image's cuBLAS
+    picks for the native draft LM head (``F.linear``, one vocab shard) at
+    M = 1, 3, 8, and the library versions. The native-order bitwise claim
+    is tied to this kernel; a different name on another image is the cue to
+    re-run ``test_native_order_logits_bitwise_equal_native``."""
+    from torch.profiler import ProfilerActivity, profile
+
+    h = torch.randn(8, oracle.HIDDEN, device="cuda").to(torch.bfloat16)
+    w = shards[0]
+    kernels = {}
+    for m in (1, 3, 8):
+        F.linear(h[:m], w)
+        torch.cuda.synchronize()
+        with profile(activities=[ProfilerActivity.CUDA]) as prof:
+            F.linear(h[:m], w)
+            torch.cuda.synchronize()
+        kernels[m] = sorted(
+            {e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA}
+        )
+    try:
+        blas = str(torch.backends.cuda.preferred_blas_library())
+    except (AttributeError, RuntimeError):
+        blas = None
+    print(
+        "[draft_tail-nvjet-record] "
+        + json.dumps(
+            {
+                "torch": torch.__version__,
+                "cuda": torch.version.cuda,
+                "device": torch.cuda.get_device_name(0),
+                "capability": list(torch.cuda.get_device_capability(0)),
+                "preferred_blas": blas,
+                "cublas_libs": _loaded_cublas_versions(),
+                "env": {
+                    k: v
+                    for k, v in os.environ.items()
+                    if k.startswith(("CUBLAS", "TORCH_BLAS", "NVIDIA_TF32"))
+                },
+                "kernels_by_M": kernels,
+            }
+        )
+    )
+
+
+@pytest.mark.parametrize("m", [1, 8])
+def test_native_order_logits_bitwise_equal_native_in_graph(runner, shards, full_head, gamma, m):
+    """The served draft head runs inside a CUDA graph. Capture the native
+    LM-head GEMM and the shipped default draft_tail kernel (``DraftTailOptions().variant``)
+    per shard in graphs with M rows, replay them over the exactness rows
+    (real, knife-edge, synthetic) and require the draft_tail BF16 logits to equal
+    the graph-captured native logits bit for bit; the graph-captured native
+    logits are also compared with eager ``F.linear`` (printed)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
+        DraftTailOptions,
+    )
+
+    variant = DraftTailOptions().variant
+    rows = _exactness_rows(full_head, gamma)
+    rows = rows[: rows.shape[0] // m * m]
+    static_h = torch.zeros(m, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda")
+    nat = [None] * len(shards)
+    w3 = [torch.empty(m, w.shape[0], dtype=torch.bfloat16, device="cuda") for w in shards]
+
+    def step():
+        for s, w in enumerate(shards):
+            nat[s] = F.linear(static_h, w)
+            runner.run_fused(static_h, w, s * w.shape[0], variant=variant, logits_out=w3[s])
+
+    st = torch.cuda.Stream()
+    st.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(st):
+        step()
+    torch.cuda.current_stream().wait_stream(st)
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        step()
+    draft_tail_vs_graph = graph_vs_eager = total = 0
+    for i in range(0, rows.shape[0], m):
+        static_h.copy_(rows[i : i + m])
+        g.replay()
+        torch.cuda.synchronize()
+        for s, w in enumerate(shards):
+            eager = F.linear(rows[i : i + m].contiguous(), w)
+            draft_tail_vs_graph += int((w3[s].view(torch.int16) != nat[s].view(torch.int16)).sum())
+            graph_vs_eager += int((nat[s].view(torch.int16) != eager.view(torch.int16)).sum())
+            total += eager.numel()
+    print(
+        f"[draft_tail-exact-graph] M={m} variant={variant} rows={rows.shape[0]} "
+        f"logits compared={total} draft_tail != graph native: {draft_tail_vs_graph}; "
+        f"graph native != eager native: {graph_vs_eager}"
+    )
+    assert draft_tail_vs_graph == 0
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py
new file mode 100644
index 000000000000..876cca23c360
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py
@@ -0,0 +1,239 @@
+# 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.
+"""draft_tail dispatch: native fallback is chosen before any device work.
+
+Uses the real ``LMHead`` and ``DeepseekV3MTPHead`` classes with stand-in
+worker / metadata objects; tp_size 1, so no exchange workspace is needed.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+VOCAB = 4096
+
+
+def _objects(**over):
+    import tensorrt_llm  # noqa: F401
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3MTPHead
+    from tensorrt_llm._torch.modules.embedding import LMHead
+    from tensorrt_llm._torch.modules.linear import TensorParallelMode
+    from tensorrt_llm.mapping import Mapping
+
+    mapping = over.pop("mapping", Mapping(world_size=1, rank=0, tp_size=1))
+    head = LMHead(
+        VOCAB,
+        6144,
+        dtype=over.pop("head_dtype", torch.bfloat16),
+        mapping=mapping,
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
+    )
+    head.weight.data = (torch.randn(head.weight.shape, device="cuda") * 0.02).to(head.weight.dtype)
+    shared_head = over.pop("shared_head", object.__new__(DeepseekV3MTPHead))
+    worker = SimpleNamespace(
+        is_mtp_eagle=over.pop("is_mtp_eagle", True),
+        guided_decoder=over.pop("guided_decoder", None),
+        _d2t=over.pop("d2t", None),
+        mapping=mapping,
+    )
+    spec = SimpleNamespace(
+        wants_advanced_draft_sampling=over.pop("advanced", False),
+        draft_vocab_size=0,
+        vocab_size=VOCAB,
+        num_seq_slots=over.pop("num_seq_slots", 8),
+        max_num_requests=over.pop("max_num_requests", 8),
+        is_cuda_graph=over.pop("is_cuda_graph", False),
+    )
+    draft = SimpleNamespace(mtp_layers=[SimpleNamespace(shared_head=shared_head)], lm_head=head)
+    assert not over, over
+    return worker, draft, spec
+
+
+def _select(worker, draft, spec):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    return dispatch.get_impl("draft_tail", "select_draft_tail")(worker, draft, spec)
+
+
+def test_eligible_path_matches_native_argmax():
+    worker, draft, spec = _objects()
+    tail = _select(worker, draft, spec)
+    assert tail is not None
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    for rows in range(1, 9):
+        h = torch.randn(rows, 6144, generator=gen, device="cuda").to(torch.bfloat16)
+        tokens = tail(h)
+        ref = torch.argmax(F.linear(h, draft.lm_head.weight), dim=-1).to(torch.int32)
+        assert tokens is not None and tokens.dtype == torch.int32
+        torch.testing.assert_close(tokens, ref, rtol=0, atol=0)
+
+
+def test_per_step_fallback_only_for_rank_shared_properties():
+    """A step leaves draft_tail only for properties every TP rank shares: the row
+    count (same requests on every rank), dtype and hidden width (same
+    model)."""
+    worker, draft, spec = _objects()
+    tail = _select(worker, draft, spec)
+    h = torch.randn(9, 6144, device="cuda").to(torch.bfloat16)
+    assert tail(h) is None  # more than 8 rows
+    assert tail(h[:0]) is None  # no rows
+    assert tail(h[:2].half()) is None  # dtype
+    assert tail(h[:2, :4096].contiguous()) is None  # hidden size
+    assert tail(h[0]) is None  # not a row block
+
+
+def _misaligned_layouts(rows, gen):
+    """Row blocks whose layout the kernels cannot read directly: non-unit
+    last stride, 8-byte-aligned base, row stride not a multiple of 8."""
+    src = torch.randn(rows, 6144, generator=gen, device="cuda").to(torch.bfloat16)
+    strided = src.t().contiguous().t()
+    off = torch.empty(rows * 6144 + 4, dtype=torch.bfloat16, device="cuda")
+    shifted = off[4:].view(rows, 6144)
+    shifted.copy_(src)
+    wide = torch.empty(rows, 6148, dtype=torch.bfloat16, device="cuda")
+    pitched = wide[:, :6144]
+    pitched.copy_(src)
+    layouts = [shifted, pitched]
+    if rows > 1:  # one row has no meaningful last-dim transpose
+        layouts.append(strided)
+    return src, layouts
+
+
+def test_per_step_layout_is_staged_never_native():
+    """Strides and alignment are rank-local (allocator placement): they must
+    not move one rank to the native collective while its peers run draft_tail's
+    exchange. Such rows are copied into draft_tail's aligned staging rows."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
+        _kernel_layout_ok,
+    )
+
+    worker, draft, spec = _objects()
+    tail = _select(worker, draft, spec)
+    runner = dispatch.get_impl("draft_tail", "get_runner")(worker.mapping)
+    gen = torch.Generator(device="cuda").manual_seed(3)
+    for rows in (1, 2, 5, 8):
+        src, layouts = _misaligned_layouts(rows, gen)
+        ref = torch.argmax(F.linear(src, draft.lm_head.weight), dim=-1).to(torch.int32)
+        for x in layouts:
+            assert not _kernel_layout_ok(x)
+            before = runner.staged_steps
+            tokens = tail(x)
+            assert tokens is not None
+            assert runner.staged_steps == before + 1
+            torch.testing.assert_close(tokens, ref, rtol=0, atol=0)
+        before = runner.staged_steps
+        torch.testing.assert_close(tail(src), ref, rtol=0, atol=0)
+        assert runner.staged_steps == before  # aligned rows: no copy
+
+
+def test_per_step_staging_under_cuda_graph():
+    """The staging copy is captured with the step and replays with new
+    contents of the same (misaligned) source buffer."""
+    worker, draft, spec = _objects()
+    tail = _select(worker, draft, spec)
+    gen = torch.Generator(device="cuda").manual_seed(4)
+    off = torch.zeros(3 * 6144 + 4, dtype=torch.bfloat16, device="cuda")
+    x = off[4:].view(3, 6144)
+    x.copy_(torch.randn(3, 6144, generator=gen, device="cuda"))
+    tail(x)  # warm-up outside capture
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = tail(x)
+    for _ in range(5):
+        x.copy_(torch.randn(3, 6144, generator=gen, device="cuda"))
+        graph.replay()
+        torch.cuda.synchronize()
+        ref = torch.argmax(F.linear(x.contiguous(), draft.lm_head.weight), dim=-1).to(torch.int32)
+        torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_per_step_device_mismatch_raises():
+    worker, draft, spec = _objects()
+    tail = _select(worker, draft, spec)
+    with pytest.raises(RuntimeError):
+        tail(torch.zeros(1, 6144, dtype=torch.bfloat16))
+
+
+@pytest.mark.parametrize(
+    "change",
+    [
+        dict(is_mtp_eagle=False),
+        dict(guided_decoder=object()),
+        dict(d2t=torch.zeros(4, dtype=torch.int64)),
+        dict(advanced=True),
+        dict(shared_head=object()),
+        dict(head_dtype=torch.float16),
+    ],
+)
+def test_configuration_fallbacks(change):
+    worker, draft, spec = _objects(**change)
+    assert _select(worker, draft, spec) is None
+
+
+@pytest.mark.parametrize(
+    "change,eligible",
+    [
+        (dict(num_seq_slots=9), False),
+        (dict(num_seq_slots=16, max_num_requests=4, is_cuda_graph=True), False),
+        (dict(num_seq_slots=0, max_num_requests=16), False),
+        (dict(num_seq_slots=0, max_num_requests=4, is_cuda_graph=True), False),
+        (dict(num_seq_slots=0, max_num_requests=4), True),
+        (dict(num_seq_slots=2, max_num_requests=1, is_cuda_graph=True), True),
+    ],
+)
+def test_row_bound_decides_per_configuration(change, eligible):
+    """draft_tail engages only if no draft step of the configuration can exceed 8
+    rows, judged from the engine's sequence-slot pool, which CUDA-graph
+    metadata keeps: a step can then never switch between draft_tail and native
+    rounding because of its (padded) row count."""
+    worker, draft, spec = _objects(**change)
+    assert (_select(worker, draft, spec) is not None) == eligible
+
+
+def test_one_kernel_for_every_row_count():
+    """The shipped options use one kernel shape for every row count, in the
+    native accumulation order (native draft tokens, near-ties included)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
+        FIRST_MMA_VARIANT,
+        NATIVE_ORDER_VARIANTS,
+        DraftTailOptions,
+    )
+
+    opts = DraftTailOptions()
+    assert opts.mode == "fused" and opts.variant >= FIRST_MMA_VARIANT
+    assert opts.variant in NATIVE_ORDER_VARIANTS
+    assert opts.shape(152) == (opts.variant, 0)
+
+
+def test_attention_dp_falls_back():
+    from tensorrt_llm.mapping import Mapping
+
+    mapping = Mapping(world_size=2, rank=0, tp_size=2, enable_attention_dp=True)
+    worker, draft, spec = _objects(mapping=mapping)
+    assert _select(worker, draft, spec) is None
+
+
+def test_feature_off_by_default(monkeypatch):
+    """Without an installed plan (a process that does not build the target) draft_tail is off."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
+    assert not dispatch.is_enabled("draft_tail")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py
new file mode 100644
index 000000000000..b2a773ca0212
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py
@@ -0,0 +1,89 @@
+# 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.
+"""Self-checks of the near-tie tooling used by the flip-rate study."""
+
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import draft_tail_nearties as nt  # noqa: E402
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+
+def test_bf16_step_walks_the_grid():
+    x = torch.tensor([1.0, -1.0, 0.0, 3.0, 17.875], dtype=torch.bfloat16)
+    up = nt.bf16_step(x, 1).double()
+    down = nt.bf16_step(x, -1).double()
+    assert bool((up > x.double()).all()) and bool((down < x.double()).all())
+    # No BF16 value lies strictly between x and its neighbours.
+    mid_up = ((x.double() + up) / 2).float().to(torch.bfloat16).double()
+    assert bool(((mid_up == x.double()) | (mid_up == up)).all())
+    assert nt.bf16_step(x, 1)[4].item() == 18.0
+    assert nt.bf16_step(nt.bf16_step(x, 5), -5).tolist() == x.tolist()
+
+
+def test_bf16_round_exact_is_nearest_without_double_rounding():
+    gen = torch.Generator().manual_seed(0)
+    v = torch.randn(200000, generator=gen, dtype=torch.float64) * 20
+    r = nt.bf16_round_exact(v)
+    lo = nt.bf16_step(r, -1).double()
+    hi = nt.bf16_step(r, 1).double()
+    rd = r.double()
+    assert bool(((v - rd).abs() <= (v - lo).abs()).all())
+    assert bool(((v - rd).abs() <= (v - hi).abs()).all())
+    # Just above / below a midpoint, closer than FP32 can resolve.
+    b = torch.tensor([17.875], dtype=torch.bfloat16)
+    n = nt.bf16_step(b, 1)
+    m = (b.double() + n.double()) / 2
+    assert nt.bf16_round_exact(m + 1e-9).item() == n.item()
+    assert nt.bf16_round_exact(m - 1e-9).item() == b.item()
+    # Exactly at the midpoint: ties to even (17.875 has an odd mantissa).
+    assert nt.bf16_round_exact(m).item() == n.item()
+
+
+def test_bf16_key_orders_value_then_lower_id():
+    vals = torch.tensor([[1.0, 2.0, 2.0, -0.0, 0.0, -3.0]], dtype=torch.bfloat16)
+    ids = torch.arange(6)[None]
+    assert int(nt.bf16_key(vals, ids).argmax()) == 1
+    z = torch.tensor([[-0.0, 0.0]], dtype=torch.bfloat16)
+    k = nt.bf16_key(z, torch.arange(2)[None])
+    assert int(k.argmax()) == 0  # -0 == +0: lower id
+
+
+def test_construct_straddle_reaches_targets():
+    gen = torch.Generator(device="cuda").manual_seed(1)
+    head = (torch.randn(64, 6144, generator=gen, device="cuda") * 0.02).to(torch.bfloat16)
+    hidden = torch.randn(16, 6144, generator=gen, device="cuda").to(torch.bfloat16)
+    up = torch.full((16,), 3, device="cuda")
+    down = torch.full((16,), 7, device="cuda")
+    for tau in (1.0, 64.0):
+        res = nt.construct_straddle(hidden, head, up, down, torch.full((16,), tau, device="cuda"))
+        ok = (
+            (res["dist_hi_fp32_ulps"] > 0)
+            & (res["dist_lo_fp32_ulps"] < 0)
+            & (res["dist_hi_fp32_ulps"] <= 1.1 * tau)
+            & (res["dist_lo_fp32_ulps"] >= -1.1 * tau)
+        )
+        # Greedy single-element steps from large toy gaps (about 2): most
+        # rows reach the band (real near-tie rows start close); callers
+        # keep only rows that did.
+        assert int(ok.sum()) >= 8, (tau, res["dist_hi_fp32_ulps"], res["dist_lo_fp32_ulps"])
+        # The returned logits are those of the returned rows.
+        h = res["hidden"].double()
+        torch.testing.assert_close((h * head[3].double()).sum(-1), res["l_hi"], rtol=0, atol=1e-9)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py
new file mode 100644
index 000000000000..9cc26219a55b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py
@@ -0,0 +1,24 @@
+# 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.
+"""Make the fused_kv_indexer_projection helpers and the absorbed_query / post_moe_hop fixtures importable under
+``--import-mode=importlib``."""
+
+import sys
+from pathlib import Path
+
+_HERE = Path(__file__).resolve().parent
+for _path in (str(_HERE), str(_HERE.parent / "absorbed_query"), str(_HERE.parent / "post_moe_hop")):
+    if _path not in sys.path:
+        sys.path.insert(0, _path)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py
new file mode 100644
index 000000000000..a4be67eb41aa
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py
@@ -0,0 +1,288 @@
+# 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.
+"""Shared helpers of the fused_kv_indexer_projection tests and benchmarks.
+
+* ``module(layer)``: the native ``DeepseekV32Attention`` of a GLM-5.2 layer
+  at TP8 rank 0 with real checkpoint weights (absorbed_query's fixture; layer 78 is the
+  MTP layer), with its real ``input_layernorm`` gain.
+* ``native_hidden``: the tensor the native GEMMs consume: the native
+  ``RMSNorm`` module (FlashInfer kernel) applied with the real gain to
+  synthetic residual-stream rows; ``real_hidden``: rows of the real
+  normalized MoE-input fixture (real-text prefill, GLM-5.2 NVFP4).
+* ``native(attn, h)``: the native ``kv_a_proj_with_mqa`` call
+  (``dsv3_fused_a_gemm_op``) and ``F.linear`` on the indexer's fused
+  ``wk | weights_proj`` weight.
+* error statistics against an FP64 oracle and the frozen budget (the deferred_norm
+  budget for the same GEMM).
+* ``proxy_producer``: a CuTe DSL stand-in for the preceding kernel. It is
+  launched with programmatic dependent launch, releases its dependents at
+  entry, spins for a given time and only then writes the activation; a
+  dependent that read it before its dependency wait would see stale data.
+"""
+
+import glob
+import os
+import sys
+
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import Int32, Int64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import dsl_user_op
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+for _p in (_HERE, os.path.join(_HERE, "..", "absorbed_query")):
+    if _p not in sys.path:
+        sys.path.insert(0, _p)
+
+import absorbed_query_fixture  # noqa: E402
+
+H = 6144
+EPS = 1e-5
+MTP_LAYER = 78
+REAL_HIDDEN_DIR = os.environ.get("GLM_MOE_DSA_REAL_MOE_INPUTS_DIR", "")
+
+_MODS = {}
+_MODEL_CONFIG = {}
+
+
+def model_config(tp_rank: int = 0):
+    if tp_rank not in _MODEL_CONFIG:
+        _MODEL_CONFIG[tp_rank] = absorbed_query_fixture.build_model_config(tp_rank)
+    return _MODEL_CONFIG[tp_rank]
+
+
+def module(layer: int):
+    """``(attn, input_layernorm gain)`` of ``layer`` (cached)."""
+    if layer not in _MODS:
+        with torch.inference_mode():
+            _MODS[layer] = (
+                absorbed_query_fixture.build_native_attention(layer, 0, model_config(0)),
+                absorbed_query_fixture.load_input_layernorm(layer).cuda(),
+            )
+    return _MODS[layer]
+
+
+def make_residual(rows: int, seed: int, scale: float = 1.0) -> torch.Tensor:
+    """Residual-stream-like rows: Gaussian plus a few large channels."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(rows, H, generator=g, device="cuda") * 2.0
+    hot = torch.randint(0, H, (12,), generator=g, device="cuda")
+    x[:, hot] *= 30.0
+    return (x * scale).bfloat16()
+
+
+_NORMS = {}
+
+
+def native_rmsnorm(gamma: torch.Tensor):
+    """The native ``RMSNorm`` module (as the decoder layers build it)."""
+    key = gamma.data_ptr()
+    if key not in _NORMS:
+        from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+
+        norm = RMSNorm(hidden_size=H, eps=EPS, dtype=torch.bfloat16).cuda()
+        norm.weight.data.copy_(gamma)
+        _NORMS[key] = norm
+    return _NORMS[key]
+
+
+def native_hidden(rows: int, seed: int, gamma: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+    """Native RMSNorm output of synthetic residual rows, times ``scale``."""
+    h = native_rmsnorm(gamma)(make_residual(rows, seed))
+    if scale != 1.0:
+        h = (h.float() * scale).bfloat16()
+    return h.contiguous()
+
+
+_REAL = []
+
+
+def real_hidden(rows: int, offset: int = 0):
+    """``rows`` rows of a real normalized hidden state, or ``None`` if the
+    fixture is not mounted (first 2-D BF16 ``[N, 6144]`` tensor of the first
+    file)."""
+    if not _REAL:
+        files = sorted(glob.glob(os.path.join(REAL_HIDDEN_DIR, "*.pt")))
+        if not files:
+            return None
+        obj = torch.load(files[0], map_location="cpu", weights_only=False)
+        cands = (
+            obj.values()
+            if isinstance(obj, dict)
+            else (obj if isinstance(obj, (list, tuple)) else [obj])
+        )
+        found = None
+        for t in cands:
+            if (
+                isinstance(t, torch.Tensor)
+                and t.dim() == 2
+                and t.shape[1] == H
+                and t.dtype == torch.bfloat16
+            ):
+                found = t
+                break
+        if found is None:
+            return None
+        _REAL.append(found)
+    t = _REAL[0]
+    return t[offset : offset + rows].contiguous().cuda()
+
+
+def native(attn, h: torch.Tensor):
+    kva = attn.kv_a_proj_with_mqa(h)
+    kw = None
+    if attn.mqa.indexer is not None:
+        kw = torch.nn.functional.linear(h, attn.mqa.indexer._fused_wk_wp_weight)
+    return kva, kw
+
+
+def fp64_oracle(h: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return h.double() @ weight.double().t()
+
+
+def norm_rms_err(y: torch.Tensor, ref: torch.Tensor) -> float:
+    ref = ref.double()
+    return float((y.double() - ref).norm() / ref.norm().clamp_min(1e-300))
+
+
+def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    def ordered(t):
+        v = t.contiguous().view(torch.int16).to(torch.int32)
+        return torch.where(v < 0, -(v + 32768), v)
+
+    return (ordered(a) - ordered(b)).abs()
+
+
+def compare(cand: torch.Tensor, nat: torch.Tensor, oracle: torch.Tensor) -> dict:
+    ulp = bf16_ulp_distance(cand, nat)
+    return {
+        "bit_identical": bool(torch.equal(cand.view(torch.int16), nat.view(torch.int16))),
+        "allclose_native": bool(torch.allclose(cand.float(), nat.float(), rtol=0.1, atol=1e-2)),
+        "e_cand": norm_rms_err(cand, oracle),
+        "e_native": norm_rms_err(nat, oracle),
+        "max_ulp_vs_native": int(ulp.max()),
+        "frac_diff_vs_native": float((ulp > 0).float().mean()),
+        "max_abs_vs_native": float((cand.float() - nat.float()).abs().max()),
+    }
+
+
+def within_budget(stats: dict) -> bool:
+    """Frozen budget of a non-bit-identical output."""
+    return stats["allclose_native"] and stats["e_cand"] <= 1.25 * stats["e_native"] + 2.0**-12
+
+
+# --------------------------------------------------------------------------
+# Proxy producer (PDL): releases dependents at entry, writes after a delay.
+# --------------------------------------------------------------------------
+
+
+def _asm(res, args, text, cons, loc=None, ip=None):
+    return llvm.inline_asm(
+        res,
+        args,
+        text,
+        cons,
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _spin_ns(ns: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [ns.ir_value()],
+        "{\n\t.reg .u64 t0, t;\n\t.reg .pred p;\n\tmov.u64 t0, %globaltimer;\n\t"
+        "GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_SPIN:\n\tmov.u64 t, %globaltimer;\n\tsub.u64 t, t, t0;\n\t"
+        "setp.lt.u64 p, t, $0;\n\t@p bra GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_SPIN;\n\t}",
+        "l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def _copy16(dst: Int64, src: Int64, *, loc=None, ip=None):
+    _asm(
+        None,
+        [dst.ir_value(), src.ir_value()],
+        "{\n\t.reg .b32 a, b, c, d;\n\tld.global.v4.b32 {a, b, c, d}, [$1];\n\t"
+        "st.global.v4.b32 [$0], {a, b, c, d};\n\t}",
+        "l,l",
+        loc=loc,
+        ip=ip,
+    )
+
+
+class _ProxyKernel:
+    """Grid ``[8, rows]`` x 96 threads; CTA ``(c, r)`` copies 1536 B of row ``r``."""
+
+    @cute.jit
+    def __call__(
+        self,
+        src: Int64,
+        dst: Int64,
+        delay_ns: Int64,
+        anchor: cute.Tensor,
+        rows: Int32,
+        stream: CUstream,
+    ):
+        self.kernel(src, dst, delay_ns, anchor).launch(
+            grid=[8, rows, 1], block=[96, 1, 1], stream=stream, use_pdl=True
+        )
+
+    @cute.kernel
+    def kernel(self, src: Int64, dst: Int64, delay_ns: Int64, anchor: cute.Tensor):
+        tid, _, _ = cute.arch.thread_idx()
+        cta, row, _ = cute.arch.block_idx()
+        cute.arch.griddepcontrol_wait()
+        cute.arch.griddepcontrol_launch_dependents()
+        _spin_ns(delay_ns)
+        off = Int64(row) * Int64(H * 2) + Int64((cta * Int32(96) + tid) * Int32(16))
+        _copy16(dst + off, src + off)
+
+
+_PROXY = {}
+
+
+def proxy_producer(src: torch.Tensor, dst: torch.Tensor, delay_ns: int) -> None:
+    """Launch the proxy producer on the current stream (PDL): ``dst <- src``
+    (``[R, 6144]`` BF16) after spinning ``delay_ns``."""
+    if "k" not in _PROXY:
+        from cutlass.cute import runtime as cute_runtime
+
+        z = Int64(0)
+        anchor = cute_runtime.make_fake_compact_tensor(
+            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
+        )
+        _PROXY["k"] = cute.compile(
+            _ProxyKernel(),
+            z,
+            z,
+            z,
+            anchor,
+            Int32(1),
+            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+    assert src.shape == dst.shape and src.shape[1] == H and src.is_contiguous()
+    _PROXY["k"](
+        src.data_ptr(), dst.data_ptr(), int(delay_ns), dst.view(torch.int32).view(-1), src.shape[0]
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py
new file mode 100644
index 000000000000..76fe1fe14450
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py
@@ -0,0 +1,275 @@
+# 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.
+"""fused_kv_indexer_projection through the target's absorbed-query packet into the native DSA
+attention (single GPU; the target's adopted attention, generation steps as decode steps).
+
+Uses the absorbed_query integration harness (real weights of the full-indexer layer 6
+and the shared layer 7 at TP8 rank 0, six requests natively prefilled with
+2,600 tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6 (MTP verify), both call
+paths, ``on_update_kv_lens()`` before every step). The same native
+normalized hidden state goes through ``absorbed_query,shared_indexer_absorbed_query``
+(absorbed_query / shared_indexer_absorbed_query with their native GEMM bridges) and
+``absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection``
+(fused_kv_indexer_projection serves both layers: ``entry.DEFAULT_LAYERS`` set to
+``full,shared``).
+
+Checks (the fused_kv_indexer_projection contract and the attention-level sanity bound
+declared with it):
+
+* engagement: fused_kv_indexer_projection serves both layers;
+* every latent-cache byte and every recorded tensor derived only from
+  ``kva`` (absorbed / FP8 query, BMM scales, indexer FP8 query) is
+  bit-identical; the indexer-K cache, indexer weights and top-k sets are
+  bit-identical when ``kw`` is, else reported and bounded (top-k mean
+  Jaccard >= 0.98, output relRMS <= 2**-7);
+* attention outputs: within ``max(2 * native rerun noise, 4 bf16 ulps)``
+  (the native sparse attention is not bitwise run-to-run stable);
+* CUDA graph: a captured fused_kv_indexer_projection step replays (1, 1, 20 times, fresh inputs)
+  equal to the eager fused_kv_indexer_projection step;
+* with no layer kind selected, ``fused_kv_indexer_projection`` selected runs the absorbed_query /
+  shared_indexer_absorbed_query path bit for bit.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
+sys.path.insert(0, os.path.join(_HERE, ".."))
+
+import absorbed_query_fixture  # noqa: E402
+import plan_testlib  # noqa: E402
+import test_absorbed_query_attention_integration as aq_integration  # noqa: E402
+
+LAYERS = aq_integration.LAYERS
+LAYOUTS = aq_integration.LAYOUTS
+KW_EXACT = os.environ.get("GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW_EXACT", "0") == "1"
+KW_RECORDS = ("indexer_weights",)
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
+    pytest.mark.skipif(
+        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
+    ),
+]
+
+
+def _set_features(features: str, fused_kv_layers: str = "full,shared"):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        entry,
+    )
+
+    entry.DEFAULT_LAYERS = tuple(fused_kv_layers.split(","))
+    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
+
+
+def _reset_runners(h):
+    for layer in LAYERS:
+        for name in ("_absorbed_query", "_shared_indexer_query", "_fused_kv_projection"):
+            setattr(h.attn[layer], name, None)
+
+
+@pytest.fixture(scope="module")
+def harness():
+    with torch.inference_mode():
+        h = aq_integration._Harness()
+    yield h
+    _set_features("")
+
+
+def _run(h, features, layout, path, hidden, pos, snap, fused_kv_layers="full,shared"):
+    _set_features(features, fused_kv_layers)
+    _reset_runners(h)
+    h.restore(snap)
+    out, rec = h.recorded_step(layout, path, hidden, pos)
+    return out, rec, h.snapshot()
+
+
+def _jaccard_rows(a, b):
+    rows = []
+    for ra, rb in zip(a.tolist(), b.tolist()):
+        sa = {v for v in ra if v >= 0}
+        sb = {v for v in rb if v >= 0}
+        union = len(sa | sb)
+        rows.append(1.0 if union == 0 else len(sa & sb) / union)
+    return rows
+
+
+def _compare(native, rerun, cand, report):
+    n_out, n_rec, n_cache = native
+    c_out, c_rec, c_cache = cand
+    fails = []
+    for key in n_cache:
+        same = torch.equal(c_cache[key], n_cache[key])
+        report[f"cache_{key}_identical"] = same
+        if not same and (key.startswith("latent") or KW_EXACT):
+            fails.append(f"cache {key} differs")
+    for layer in LAYERS:
+        for key, val in n_rec[layer].items():
+            if key in ("topk", "sparse_indices"):
+                continue
+            same = torch.equal(c_rec[layer][key], val)
+            report[f"L{layer}_{key}_identical"] = same
+            if not same and (key not in KW_RECORDS or KW_EXACT):
+                fails.append(f"L{layer} record {key} differs")
+        if "topk" in n_rec[layer]:
+            jac = _jaccard_rows(c_rec[layer]["topk"], n_rec[layer]["topk"])
+            report[f"L{layer}_topk_mean_jaccard"] = sum(jac) / len(jac)
+            if min(jac) < 1.0 and KW_EXACT:
+                fails.append(f"L{layer} top-k sets differ")
+            if sum(jac) / len(jac) < 0.98:
+                fails.append(f"L{layer} top-k mean Jaccard {sum(jac) / len(jac):.4f} < 0.98")
+        diff = (c_out[layer].float() - n_out[layer].float()).abs().max().item()
+        noise = (rerun[0][layer].float() - n_out[layer].float()).abs().max().item()
+        ulp = aq_integration._bf16_ulp(n_out[layer].float().abs().max().item())
+        rel = aq_integration._rel_rms(c_out[layer], n_out[layer])
+        report[f"L{layer}_out_maxdiff"] = diff
+        report[f"L{layer}_out_noise"] = noise
+        report[f"L{layer}_out_rel_rms"] = rel
+        if KW_EXACT or LAYERS.index(layer) == 1:
+            # kva-only layer (shared) or exact kw: native run-to-run bound.
+            if diff > max(2 * noise, 4 * ulp):
+                fails.append(f"L{layer} output diff {diff:.3e} > max(2*{noise:.3e}, 4*{ulp:.3e})")
+        elif rel > 2.0**-7:
+            fails.append(f"L{layer} output relRMS {rel:.3e} > 2**-7")
+    return fails
+
+
+@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_fused_kv_into_attention(harness, layout, path):
+    h = harness
+    hidden, pos = h.gen_inputs(layout, seed=3)
+    snap = h.snapshot()
+    try:
+        native = _run(
+            h, "absorbed_query,shared_indexer_absorbed_query", layout, path, hidden, pos, snap
+        )
+        cand = _run(
+            h,
+            "absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection",
+            layout,
+            path,
+            hidden,
+            pos,
+            snap,
+        )
+        for layer in LAYERS:
+            assert h.attn[layer]._fused_kv_projection not in (None, False), (
+                f"fused_kv_indexer_projection not engaged on L{layer}"
+            )
+        rerun = _run(
+            h, "absorbed_query,shared_indexer_absorbed_query", layout, path, hidden, pos, snap
+        )
+        report = {}
+        fails = _compare(native, rerun, cand, report)
+        print(
+            f"GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_ATTN layout={layout} path={path} {report}",
+            flush=True,
+        )
+        assert not fails, fails
+    finally:
+        h.restore(snap)
+        _set_features("")
+
+
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_no_layer_kind_selected_is_native_path(harness, layout):
+    h = harness
+    hidden, pos = h.gen_inputs(layout, seed=5)
+    snap = h.snapshot()
+    try:
+        n_out, n_rec, n_cache = _run(
+            h,
+            "absorbed_query,shared_indexer_absorbed_query",
+            layout,
+            "custom_op",
+            hidden,
+            pos,
+            snap,
+        )
+        c_out, c_rec, c_cache = _run(
+            h,
+            "absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection",
+            layout,
+            "custom_op",
+            hidden,
+            pos,
+            snap,
+            fused_kv_layers="",
+        )
+        for layer in LAYERS:
+            assert h.attn[layer]._fused_kv_projection is False
+        for key in n_cache:
+            assert torch.equal(c_cache[key], n_cache[key]), key
+        for layer in LAYERS:
+            for key, val in n_rec[layer].items():
+                if key in ("topk", "sparse_indices"):
+                    continue
+                assert torch.equal(c_rec[layer][key], val), (layer, key)
+    finally:
+        h.restore(snap)
+        _set_features("")
+
+
+@pytest.mark.parametrize("layout", list(LAYOUTS))
+@torch.inference_mode()
+def test_fused_kv_cuda_graph_replay(harness, layout):
+    h = harness
+    snap = h.snapshot()
+    try:
+        _set_features("absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection")
+        _reset_runners(h)
+        hidden, pos = h.gen_inputs(layout, seed=7)
+        static = hidden.clone()
+        static_pos = pos.clone()
+        meta = h.gen_metadata(layout, cuda_graph=True)
+        h.restore(snap)
+        h.step(meta, static_pos, static, "custom_op")
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        h.restore(snap)
+        with torch.cuda.graph(graph):
+            g_out = h.step(meta, static_pos, static, "custom_op")
+        for trial, replays in enumerate((1, 1, 20)):
+            fresh, _ = h.gen_inputs(layout, seed=100 + trial)
+            static.copy_(fresh)
+            for _ in range(replays):
+                h.restore(snap)
+                graph.replay()
+            torch.cuda.synchronize()
+            got = {layer: g_out[layer].clone() for layer in LAYERS}
+            got_cache = h.snapshot()
+            h.restore(snap)
+            eager = h.gen_step(
+                layout, "custom_op", fresh, pos, meta=h.gen_metadata(layout, cuda_graph=True)
+            )
+            eager_cache = h.snapshot()
+            for key in eager_cache:
+                assert torch.equal(got_cache[key], eager_cache[key]), (trial, key)
+            for layer in LAYERS:
+                diff = (got[layer].float() - eager[layer].float()).abs().max().item()
+                ref = eager[layer].float().abs().max().item()
+                assert diff <= 4 * aq_integration._bf16_ulp(ref), (trial, layer, diff)
+    finally:
+        h.restore(snap)
+        _set_features("")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py
new file mode 100644
index 000000000000..fffdd1d29299
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py
@@ -0,0 +1,294 @@
+# 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.
+"""The query_front packet through the fused_kv_indexer_projection entry (``fused_kv_dsa_proj``)
+against the host package's own path on the same hidden state (real weights, single GPU).
+
+* full-indexer target layer 0 and the MTP layer 78 (draft_absorbed_query): absorbed_query's
+  ``project`` with its native bridges (``fused_a`` + q_a norm next to the
+  ``wk | weights_proj`` GEMM + k_norm on the aux stream) vs fused_kv_indexer_projection +
+  the same norms + the same absorbed_query kernel. Every packet tensor that depends only
+  on ``kva`` must be bit-identical; the tensors derived from ``kw`` (indexer
+  key FP8 + scale, indexer weights) are bit-identical when ``kw`` is, and
+  otherwise held to the frozen budget.
+* shared-indexer target layer 3: shared_indexer_absorbed_query's ``project`` (native ``fused_a`` +
+  shared_indexer_absorbed_query kernel) and absorbed_query's shared-layer path vs
+  fused_kv_indexer_projection + the same kernels: bit-identical.
+* static declines (layer kinds not selected, ineligible module) return
+  ``None`` before any GPU work; a compile error propagates and is not
+  cached.
+
+The feature selection is set by replacing ``dispatch.enabled_features``
+(test-only), so these tests do not need the target's model.
+"""
+
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
+
+import absorbed_query_fixture  # noqa: E402
+import fused_kv_utils as U  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
+    reason="needs a GPU and the GLM-5.2 checkpoint",
+)
+
+KW_EXACT = os.environ.get("GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW_EXACT", "0") == "1"
+CTX = 491061
+GEN_META = SimpleNamespace(num_contexts=0, num_generations=1)
+# Packet slots (AbsorbedQueryPacket.proj_outputs) of a full-indexer layer.
+KVA_SLOTS = (0, 1, 2, 3, 4, 8)  # q, compressed_kv, k_pe, latent, indexer q_fp8, q_scale
+KW_SLOTS = (5, 6, 7)  # indexer k_fp8, k_scale, weights
+
+
+@pytest.fixture
+def features(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        entry,
+    )
+
+    def select(ids):
+        monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(ids))
+
+    select(
+        {
+            "absorbed_query",
+            "draft_absorbed_query",
+            "shared_indexer_absorbed_query",
+            "fused_kv_indexer_projection",
+        }
+    )
+    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("full", "shared", "draft"))
+    return select
+
+
+def _reset(attn):
+    for name in (
+        "_absorbed_query",
+        "_draft_absorbed_query",
+        "_shared_indexer_query",
+        "_fused_kv_projection",
+    ):
+        setattr(attn, name, None)
+
+
+def _pos(rows):
+    return (CTX + torch.arange(rows, dtype=torch.int32, device="cuda")).unsqueeze(0)
+
+
+def _assert_same(a, b, tag):
+    assert a.shape == b.shape and a.dtype == b.dtype, tag
+    assert torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)), tag
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("layer", [0, U.MTP_LAYER])
+@torch.inference_mode()
+def test_full_layer_packet(features, layer, rows):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        absorb_for,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
+        fused_kv_dsa_proj,
+    )
+
+    attn, gamma = U.module(layer)
+    _reset(attn)
+    h = U.native_hidden(rows, seed=31 + rows, gamma=gamma)
+    pos = _pos(rows)
+    ref = AbsorbedQueryProjection(attn).project(pos, h, absorb=absorb_for(attn))
+    got = fused_kv_dsa_proj(attn, pos, h, GEN_META)
+    torch.cuda.synchronize()
+    assert got is not None and attn._fused_kv_projection not in (None, False)
+    assert len(got.proj_outputs) == len(ref.proj_outputs) == 9
+    for i in KVA_SLOTS:
+        _assert_same(got.proj_outputs[i], ref.proj_outputs[i], f"slot {i}")
+    assert (got.fused_q is None) == (ref.fused_q is None)
+    if ref.fused_q is not None:
+        _assert_same(got.fused_q[..., :512], ref.fused_q[..., :512], "fused_q")
+    k_fp8, k_scale, weights = (got.proj_outputs[i] for i in KW_SLOTS)
+    nk_fp8, nk_scale, nweights = (ref.proj_outputs[i] for i in KW_SLOTS)
+    if KW_EXACT:
+        for i in KW_SLOTS:
+            _assert_same(got.proj_outputs[i], ref.proj_outputs[i], f"slot {i}")
+    else:
+        torch.testing.assert_close(
+            k_fp8.float() * k_scale, nk_fp8.float() * nk_scale, rtol=0.1, atol=1e-2
+        )
+        torch.testing.assert_close(weights, nweights, rtol=0.1, atol=1e-2)
+    diff = {
+        f"slot{i}": float(
+            (
+                got.proj_outputs[i].contiguous().view(torch.uint8)
+                != ref.proj_outputs[i].contiguous().view(torch.uint8)
+            )
+            .float()
+            .mean()
+        )
+        for i in KW_SLOTS
+    }
+    print(
+        f"GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_ENTRY layer={layer} rows={rows} kw_derived_byte_diff={diff}"
+    )
+
+
+@pytest.mark.parametrize("host", ["shared_indexer_absorbed_query", "absorbed_query"])
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@torch.inference_mode()
+def test_shared_layer_packet(features, host, rows):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
+        AbsorbedQueryProjection,
+        absorb_for,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.shared_indexer import (  # noqa: E501
+        SharedIndexerProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
+        fused_kv_dsa_proj,
+    )
+
+    if host == "absorbed_query":
+        features({"absorbed_query", "fused_kv_indexer_projection"})
+    attn, gamma = U.module(3)
+    _reset(attn)
+    h = U.native_hidden(rows, seed=71 + rows, gamma=gamma)
+    pos = _pos(rows)
+    if host == "shared_indexer_absorbed_query":
+        ref = SharedIndexerProjection(attn).project(h)
+    else:
+        ref = AbsorbedQueryProjection(attn).project(pos, h, absorb=absorb_for(attn))
+    got = fused_kv_dsa_proj(attn, pos, h, GEN_META)
+    torch.cuda.synchronize()
+    assert got is not None
+    assert len(got.proj_outputs) == len(ref.proj_outputs) == 4
+    for i in range(4):
+        _assert_same(got.proj_outputs[i], ref.proj_outputs[i], f"slot {i}")
+    assert (got.fused_q is None) == (ref.fused_q is None)
+    if ref.fused_q is not None:
+        _assert_same(got.fused_q[..., :512], ref.fused_q[..., :512], "fused_q")
+
+
+@torch.inference_mode()
+def test_static_declines_before_gpu_work(features, monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        entry,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
+        fused_kv_dsa_proj,
+    )
+
+    attn, gamma = U.module(0)
+    h = U.native_hidden(6, seed=3, gamma=gamma)
+    # Layer kind not selected.
+    _reset(attn)
+    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("shared",))
+    assert fused_kv_dsa_proj(attn, _pos(6), h, GEN_META) is None
+    assert attn._fused_kv_projection is False
+    # fused_kv_indexer_projection not selected.
+    _reset(attn)
+    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("full",))
+    features({"absorbed_query"})
+    assert fused_kv_dsa_proj(attn, _pos(6), h, GEN_META) is None
+    # Per call: context batch, too many rows, non-contiguous input.
+    features({"absorbed_query", "fused_kv_indexer_projection"})
+    _reset(attn)
+    assert (
+        fused_kv_dsa_proj(attn, _pos(6), h, SimpleNamespace(num_contexts=1, num_generations=0))
+        is None
+    )
+    h9 = U.native_hidden(9, seed=4, gamma=gamma)
+    assert fused_kv_dsa_proj(attn, _pos(9), h9, GEN_META) is None
+    assert fused_kv_dsa_proj(attn, _pos(3), h9[::3], GEN_META) is None
+    _reset(attn)
+
+
+def test_layer_kinds_and_default_selection(monkeypatch):
+    """Layer 3 is the dense-fed shared layer; the default serves only it."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        entry,
+    )
+
+    def mla(idx, full):
+        return SimpleNamespace(
+            layer_idx=idx,
+            num_hidden_layers=78,
+            mqa=SimpleNamespace(indexer=object() if full else None),
+        )
+
+    assert entry.layer_kind(mla(0, True)) == "full"
+    assert entry.layer_kind(mla(3, False)) == "dense_shared"
+    assert entry.layer_kind(mla(7, False)) == "shared"
+    assert entry.layer_kind(mla(78, True)) == "draft"
+    served = entry.served_layer_kinds()
+    assert served == ("dense_shared",)
+    assert entry.kind_selected("dense_shared", served)
+    for kind in ("full", "shared", "draft"):
+        assert not entry.kind_selected(kind, served)
+    assert entry.kind_selected("dense_shared", ("shared",))
+    assert not entry.kind_selected("shared", ("dense_shared",))
+
+
+@torch.inference_mode()
+def test_default_selection_serves_layer3_only(features, monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        entry,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
+        fused_kv_dsa_proj,
+    )
+
+    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("dense_shared",))
+    full, gamma0 = U.module(0)
+    _reset(full)
+    assert (
+        fused_kv_dsa_proj(full, _pos(6), U.native_hidden(6, seed=8, gamma=gamma0), GEN_META) is None
+    )
+    assert full._fused_kv_projection is False
+    shared, gamma3 = U.module(3)
+    _reset(shared)
+    got = fused_kv_dsa_proj(shared, _pos(6), U.native_hidden(6, seed=9, gamma=gamma3), GEN_META)
+    assert got is not None
+    _reset(full)
+    _reset(shared)
+
+
+@torch.inference_mode()
+def test_compile_error_is_not_masked(features, monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        projection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
+        fused_kv_dsa_proj,
+    )
+
+    def _fail(self):
+        raise RuntimeError("injected projection compile failure")
+
+    attn, gamma = U.module(0)
+    _reset(attn)
+    monkeypatch.setattr(projection.FusedKvProjection, "precompile", _fail)
+    with pytest.raises(RuntimeError, match="injected projection compile failure"):
+        fused_kv_dsa_proj(attn, _pos(6), U.native_hidden(6, seed=5, gamma=gamma), GEN_META)
+    assert attn._fused_kv_projection is None, "a failed runner must not be cached"
+    _reset(attn)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py
new file mode 100644
index 000000000000..1dd8e71e8a00
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py
@@ -0,0 +1,331 @@
+# 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.
+"""Single-GPU tests of the fused_kv_indexer_projection merged projection (real GLM-5.2 weights).
+
+Oracle: the native ``kv_a_proj_with_mqa`` call (``dsv3_fused_a_gemm_op``)
+and ``F.linear`` on the indexer's fused ``wk | weights_proj`` weight, on the
+same normalized hidden state.
+
+Contract (frozen before measurement): ``kva`` is
+bit-identical to native for every R = 1..8; ``kw`` is bit-identical or, as
+the named variant, within the frozen budget ``allclose(rtol=0.1,
+atol=1e-2)`` and ``e_fused_kv <= 1.25 e_native + 2^-12`` against the FP64 oracle.
+
+Fixture labels: weights real (checkpoint); activations: native RMSNorm of
+synthetic residual rows with the real gain, scaled copies, plain Gaussian
+rows, and (when mounted) real normalized hidden rows of the captured fixture.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
+
+import absorbed_query_fixture  # noqa: E402
+import fused_kv_utils as U  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
+    reason="needs a GPU and the GLM-5.2 checkpoint",
+)
+
+LAYERS = [
+    int(x)
+    for x in os.environ.get(
+        "GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_TEST_LAYERS", "0,1,2,3,78"
+    ).split(",")
+]
+ROWS = [
+    int(x)
+    for x in os.environ.get(
+        "GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_TEST_ROWS", "1,2,3,4,5,6,7,8"
+    ).split(",")
+]
+# Whether kw must be bit-identical (set once the order probe found the native order).
+KW_EXACT = os.environ.get("GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW_EXACT", "0") == "1"
+INPUTS = ("norm", "norm_x1e-3", "norm_x30", "gauss", "real")
+
+
+def _projection(attn, config=None):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        FusedKvProjection,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.projection import (  # noqa: E501
+        projection_eligibility,
+    )
+
+    ok, reason, kw = projection_eligibility(attn)
+    assert ok, reason
+    proj = FusedKvProjection(attn, config, kw_weight=kw)
+    assert proj.supported() is None, proj.supported()
+    return proj
+
+
+def _parse_cfg(text: str, **extra):
+    from dataclasses import replace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
+        FusedKvProjConfig,
+    )
+
+    return replace(FusedKvProjConfig.parse(text), **extra)
+
+
+def _input(kind: str, rows: int, seed: int, gamma: torch.Tensor):
+    if kind == "norm":
+        return U.native_hidden(rows, seed, gamma)
+    if kind == "norm_x1e-3":
+        return U.native_hidden(rows, seed, gamma, scale=1e-3)
+    if kind == "norm_x30":
+        return U.native_hidden(rows, seed, gamma, scale=30.0)
+    if kind == "gauss":
+        g = torch.Generator(device="cuda").manual_seed(seed)
+        return torch.randn(rows, U.H, generator=g, device="cuda").bfloat16()
+    return U.real_hidden(rows, offset=97 * seed % 1024)
+
+
+def _check_outputs(attn, h, kva, kw, tag):
+    kva_n, kw_n = U.native(attn, h)
+    torch.cuda.synchronize()
+    assert kva.shape == kva_n.shape and kva.dtype == kva_n.dtype and kva.stride() == kva_n.stride()
+    ok = torch.equal(kva.view(torch.int16), kva_n.view(torch.int16))
+    if not ok:
+        stats = U.compare(kva, kva_n, U.fp64_oracle(h, attn.kv_a_proj_with_mqa.weight))
+        pytest.fail(f"kva not bit-identical {tag}: {stats}")
+    if kw_n is None:
+        assert kw is None
+        return None
+    assert kw.shape == kw_n.shape and kw.dtype == kw_n.dtype and kw.stride() == kw_n.stride()
+    stats = U.compare(kw, kw_n, U.fp64_oracle(h, attn.mqa.indexer._fused_wk_wp_weight))
+    print(f"GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW {tag} {stats}")
+    if KW_EXACT:
+        assert stats["bit_identical"], (tag, stats)
+    else:
+        assert U.within_budget(stats), (tag, stats)
+    return stats
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+def test_served_kernel_supported_and_compiled(layer):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.proj_kernel import (  # noqa: E501
+        is_compiled,
+    )
+
+    attn, _ = U.module(layer)
+    proj = _projection(attn)
+    proj.precompile()
+    assert is_compiled(**proj.kernel_config())
+    assert proj.n1 == (0 if attn.mqa.indexer is None else 160)
+
+
+@pytest.mark.parametrize("kind", INPUTS)
+@pytest.mark.parametrize("rows", ROWS)
+@pytest.mark.parametrize("layer", LAYERS)
+def test_kva_bit_identical_kw_contract(layer, rows, kind):
+    attn, gamma = U.module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        h = _input(kind, rows, seed=1000 + 10 * layer + rows, gamma=gamma)
+        if h is None:
+            pytest.skip("real hidden-state fixture not mounted")
+        kva, kw = proj.run(h)
+        _check_outputs(attn, h, kva, kw, f"layer={layer} rows={rows} input={kind}")
+
+
+@pytest.mark.parametrize("layer", LAYERS)
+def test_rows_match_single_row_runs(layer):
+    """Every token row equals a launch of that row alone (padded or other
+    rows of the batch never change a live row)."""
+    attn, gamma = U.module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        h = U.native_hidden(8, seed=11, gamma=gamma)
+        kva, kw = proj.run(h)
+        for t in range(8):
+            kva1, kw1 = proj.run(h[t : t + 1].contiguous())
+            assert torch.equal(kva1[0], kva[t])
+            if kw is not None:
+                assert torch.equal(kw1[0], kw[t])
+
+
+@pytest.mark.parametrize("layer", LAYERS[:2] + LAYERS[-1:])
+def test_graph_replay_changing_rows(layer):
+    """Graphs captured at R = 6 and R = 1 (the draft-loop shapes) replayed
+    interleaved with new inputs equal eager launches."""
+    attn, gamma = U.module(layer)
+    proj = _projection(attn)
+    with torch.inference_mode():
+        static_inputs, graphs, outs = {}, {}, {}
+        for rows in (6, 1):
+            static_inputs[rows] = U.native_hidden(rows, seed=rows, gamma=gamma)
+            proj.run(static_inputs[rows])
+            torch.cuda.synchronize()
+            graphs[rows] = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graphs[rows]):
+                outs[rows] = proj.run(static_inputs[rows])
+        for trial in range(6):
+            rows = (6, 1)[trial % 2]
+            static_inputs[rows].copy_(U.native_hidden(rows, seed=300 + trial, gamma=gamma))
+            for _ in range(1 + 20 * (trial == 4)):
+                graphs[rows].replay()
+            e_kva, e_kw = proj.run(static_inputs[rows])
+            torch.cuda.synchronize()
+            g_kva, g_kw = outs[rows]
+            assert torch.equal(g_kva, e_kva)
+            if e_kw is not None:
+                assert torch.equal(g_kw, e_kw)
+            _check_outputs(
+                attn, static_inputs[rows], g_kva, g_kw, f"graph layer={layer} rows={rows}"
+            )
+
+
+@pytest.mark.parametrize(
+    "cfg", ["b_load=cg", "b_load=ca", "copy_method=tma,b_load=ca", "prewait_chunks=1"]
+)
+@pytest.mark.parametrize("stage_before_wait", [True, False])
+@pytest.mark.parametrize("delay_ns", [0, 5000, 50000])
+@pytest.mark.parametrize("layer", LAYERS[:1] + LAYERS[-1:])
+def test_pdl_delayed_producer(layer, delay_ns, stage_before_wait, cfg):
+    """The producer releases its dependents at entry and writes the
+    activation only after ``delay_ns``; the projection (PDL, staging weights
+    before its dependency wait) must read the final activation. The same
+    activation buffer is rewritten every trial, so an activation line kept in
+    L1 from the previous trial (``b_load=ca``) would show up as a mismatch."""
+    attn, gamma = U.module(layer)
+    proj = _projection(attn, _parse_cfg(cfg, stage_before_wait=stage_before_wait))
+    if not proj.use_pdl:
+        pytest.skip("PDL disabled (TRTLLM_ENABLE_PDL=0)")
+    with torch.inference_mode():
+        dst = torch.empty(6, U.H, dtype=torch.bfloat16, device="cuda")
+        for trial in range(6):
+            src = U.native_hidden(6, seed=500 + trial, gamma=gamma)
+            dst.fill_(float("nan") if trial % 2 else 0.0)
+            torch.cuda.synchronize()
+            U.proxy_producer(src, dst, delay_ns)
+            kva, kw = proj.run(dst)
+            ref_kva, ref_kw = U.native(attn, src)
+            torch.cuda.synchronize()
+            assert torch.equal(kva, ref_kva), f"trial {trial}"
+            if kw is not None:
+                r_kva, r_kw = proj.run(src)
+                torch.cuda.synchronize()
+                assert torch.equal(kw, r_kw), f"trial {trial}"
+
+
+@pytest.mark.parametrize(
+    "cfg",
+    [
+        "stage_before_wait=0",
+        "trigger=end",
+        "trigger=none",
+        "nsub=8",
+        "nsub=2",
+        "b_dist=4",
+        "b_dist=24",
+        "b_dist=48",
+        "a_dist=4",
+        "ctas=148",
+        "ctas=150",
+        "l2_prefetch=0",
+        "producer_warp=0",
+        "copy_method=cpasync,producer_warps=4",
+        "copy_method=cpasync,producer_warps=2,nsub=8",
+        "b_load=ca",
+        "chunks=48-24-12-12",
+        "chunks=40-32-16-8,b_load=ca",
+        "chunks=64-32,copy_method=cpasync,producer_warps=4",
+        "copy_method=tma",
+        "copy_method=tma,b_load=ca,nsub=8",
+        "copy_method=tma,nsub=24",
+        "copy_method=tma,tma_groups=3,chunks=48-24-12-12",
+        "copy_method=tma,producer_warp=0",
+        "copy_method=tma,stage_before_wait=0,b_load=ca",
+        "copy_method=tma,ctas=148",
+        "b_load=cg",
+        "nsub=4,b_load=cg",
+        "prewait_chunks=1",
+        "prewait_chunks=2,copy_method=tma",
+        "prewait_chunks=0",
+        "prewait_chunks=3,copy_method=cpasync,producer_warps=4",
+    ],
+)
+@pytest.mark.parametrize("layer", LAYERS[:1] + LAYERS[3:4])
+def test_launch_configs_are_numerically_identical(cfg, layer):
+    """Launch-configuration knobs change the schedule only, never a bit."""
+    attn, gamma = U.module(layer)
+    proj = _projection(attn, _parse_cfg(cfg))
+    ref = _projection(attn)
+    with torch.inference_mode():
+        for rows in (1, 6, 8):
+            h = U.native_hidden(rows, seed=900 + rows, gamma=gamma)
+            a = proj.run(h)
+            b = ref.run(h)
+            torch.cuda.synchronize()
+            assert torch.equal(a[0], b[0])
+            if b[1] is not None:
+                assert torch.equal(a[1], b[1])
+
+
+def test_corunning_kernels_progress():
+    """The projection on one stream while another stream keeps native GEMMs
+    of the same weights in flight: every launch completes and is correct (no
+    inter-CTA dependency, no co-residency assumption)."""
+    attn, gamma = U.module(LAYERS[0])
+    proj = _projection(attn)
+    side = torch.cuda.Stream()
+    with torch.inference_mode():
+        hs = [U.native_hidden(6, seed=700 + i, gamma=gamma) for i in range(8)]
+        refs = [U.native(attn, h) for h in hs]
+        torch.cuda.synchronize()
+        side.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(side):
+            for _ in range(40):
+                for h in hs:
+                    U.native(attn, h)
+        outs = [proj.run(h) for h in hs for _ in range(5)]
+        torch.cuda.synchronize()
+        for i, (kva, kw) in enumerate(outs):
+            r_kva, r_kw = refs[i // 5]
+            assert torch.equal(kva, r_kva)
+            if r_kw is not None:
+                assert U.within_budget(
+                    U.compare(
+                        kw, r_kw, U.fp64_oracle(hs[i // 5], attn.mqa.indexer._fused_wk_wp_weight)
+                    )
+                )
+
+
+def test_supported_rejects_out_of_domain():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.proj_kernel import (  # noqa: E501
+        projection_supported,
+    )
+
+    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4) is None
+    assert projection_supported(6144, 2624, 152, 4, 18, 4, 4) is None
+    assert projection_supported(6144, 2784, 152, 4, 17, 4, 4) is not None  # 2 direct rows
+    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4, (48, 24, 12, 12)) is None
+    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4, (48, 24, 12)) is not None
+    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4, (), 160, "tma") is None
+    assert projection_supported(6144, 2624, 152, 4, 18, 4, 4, (), 0, "tma") is None
+    assert projection_supported(6144, 2784, 152, 4, 18, 4, 24, (), 160, "tma", 3) is not None
+    assert projection_supported(6144, 2784, 152, 4, 18, 4, 8, (), 160, "tma", 3) is None
+    assert projection_supported(6144, 2784, 100, 4, 18, 4, 4, (), 160, "tma") is not None
+    assert projection_supported(6144, 2784, 100, 4, 18, 4, 4) is not None  # 28 rows per CTA
+    assert projection_supported(6144, 2784, 152, 4, 18, 3, 4) is not None  # slices vs warps
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py
new file mode 100644
index 000000000000..45bfe5fbdbed
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py
@@ -0,0 +1,176 @@
+# 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.
+"""host_page_retention host policy: locating and setting mimalloc's purge delay in libc10.so.
+
+CPU only. Tests that need PyTorch's statically linked mimalloc skip on builds
+without it.
+"""
+
+import ctypes
+import json
+import os
+import subprocess
+import sys
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.host_page_retention import (  # noqa: E501
+    host_policy as hp,
+)
+
+
+def _options_or_skip():
+    try:
+        return hp.MimallocOptions(hp.torch_libc10_path())
+    except hp.HostPolicyUnavailable as exc:
+        pytest.skip(f"no mimalloc option table in this PyTorch build: {exc}")
+
+
+def _run(code, env_extra):
+    env = {k: v for k, v in os.environ.items() if not k.startswith("MIMALLOC_")}
+    env.update(env_extra)
+    out = subprocess.run(
+        [sys.executable, "-c", code],
+        env=env,
+        check=True,
+        capture_output=True,
+        text=True,
+        cwd="/tmp",
+        timeout=300,
+    )
+    return json.loads(out.stdout.strip().splitlines()[-1])
+
+
+_READ = """
+import json, torch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.host_page_retention import host_policy as hp
+opts = hp.MimallocOptions(hp.torch_libc10_path())
+before = opts.get("purge_delay")
+result = None
+if {apply}:
+    result = hp.apply_purge_policy()
+print(json.dumps({{"before": before, "after": opts.get("purge_delay"), "result": result}}))
+"""  # noqa: E501
+
+
+def test_option_table_is_found_and_consistent():
+    opts = _options_or_skip()
+    assert opts.names.count(hp.PURGE_DELAY_OPTION) == 1
+    assert len(opts.names) == len(set(opts.names))
+    for name in opts.names:
+        opts.get(name)  # mi_option_get agrees with the table for every option
+
+
+def test_set_takes_effect_and_restores():
+    opts = _options_or_skip()
+    old = opts.get(hp.PURGE_DELAY_OPTION)
+    try:
+        opts.set(hp.PURGE_DELAY_OPTION, 1234)
+        assert opts.get(hp.PURGE_DELAY_OPTION) == 1234
+    finally:
+        opts.set(hp.PURGE_DELAY_OPTION, old)
+    assert opts.get(hp.PURGE_DELAY_OPTION) == old
+
+
+def test_runtime_policy_equals_environment_setting():
+    _options_or_skip()
+    env_policy = _run(_READ.format(apply=False), {"MIMALLOC_PURGE_DELAY": str(hp.PURGE_DELAY_MS)})
+    served = _run(_READ.format(apply=False), {"MIMALLOC_PURGE_DELAY": "0"})
+    applied = _run(_READ.format(apply=True), {"MIMALLOC_PURGE_DELAY": "0"})
+    assert env_policy["before"] == hp.PURGE_DELAY_MS
+    assert served["before"] == 0 and served["after"] == 0
+    assert applied["result"] == [0, hp.PURGE_DELAY_MS]
+    assert applied["after"] == env_policy["before"]
+
+
+def test_already_set_is_left_alone():
+    _options_or_skip()
+    out = _run(_READ.format(apply=True), {"MIMALLOC_PURGE_DELAY": str(hp.PURGE_DELAY_MS)})
+    assert out["result"] == [hp.PURGE_DELAY_MS, hp.PURGE_DELAY_MS]
+
+
+def test_bad_descriptor_table_is_rejected():
+    opts = _options_or_skip()
+    good = ctypes.string_at(opts._table, opts._count * hp._DESC.size)
+    buf = ctypes.create_string_buffer(good, len(good))
+    value, init, option, name, legacy = hp._DESC.unpack_from(buf.raw, hp._DESC.size)
+    hp._DESC.pack_into(buf, hp._DESC.size, value, init, option + 7, name, legacy)
+    fake = object.__new__(hp.MimallocOptions)
+    fake._table = ctypes.addressof(buf)
+    fake._count = opts._count
+    with pytest.raises(hp.HostPolicyUnavailable, match="not a mimalloc option table"):
+        fake._check_table()
+
+
+@pytest.mark.parametrize("content", [b"", b"not an elf file at all" * 4])
+def test_non_elf_library_is_rejected(tmp_path, content):
+    path = tmp_path / "libc10.so"
+    path.write_bytes(content)
+    with pytest.raises(hp.HostPolicyUnavailable):
+        hp.MimallocOptions(str(path))
+
+
+def test_missing_library_is_rejected():
+    with pytest.raises(hp.HostPolicyUnavailable, match="not found"):
+        hp.MimallocOptions("/nonexistent/libc10.so")
+
+
+def test_library_not_loaded_is_rejected(tmp_path):
+    # A copy of libc10.so is an ELF file with the symbols but is not loaded
+    # under that path: the policy must not dlopen a second copy.
+    _options_or_skip()
+    copy = tmp_path / "libc10.so"
+    copy.write_bytes(open(hp.torch_libc10_path(), "rb").read())
+    with pytest.raises(hp.HostPolicyUnavailable, match="not loaded"):
+        hp.MimallocOptions(str(copy))
+
+
+def test_install_declines_without_side_effect(monkeypatch):
+    opts = _options_or_skip()
+    old = opts.get(hp.PURGE_DELAY_OPTION)
+    logged = []
+    monkeypatch.setattr(hp, "torch_libc10_path", lambda: "/nonexistent/libc10.so")
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(dispatch, "_log_info", logged.append)
+    assert hp.install_host_policy("test") is False
+    assert opts.get(hp.PURGE_DELAY_OPTION) == old
+    assert logged and "host_page_retention/host_policy: native fallback" in logged[0]
+
+
+def test_install_engages_and_logs(monkeypatch):
+    opts = _options_or_skip()
+    old = opts.get(hp.PURGE_DELAY_OPTION)
+    logged = []
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    monkeypatch.setattr(dispatch, "_log_info", logged.append)
+    try:
+        assert hp.install_host_policy("test") is True
+        assert opts.get(hp.PURGE_DELAY_OPTION) == hp.PURGE_DELAY_MS
+        assert "host_page_retention/host_policy: engaged" in logged[0]
+        assert f"{old} -> {hp.PURGE_DELAY_MS} ms" in logged[0]
+    finally:
+        opts.set(hp.PURGE_DELAY_OPTION, old)
+
+
+def test_package_registers_entry_point():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
+        host_page_retention,
+    )
+
+    assert (
+        dispatch.get_impl("host_page_retention", "host_policy")
+        is host_page_retention.install_host_policy
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py
new file mode 100644
index 000000000000..a4e3d588ca78
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py
@@ -0,0 +1,82 @@
+# 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.
+"""host_page_retention selection and call site.
+
+A plan runs ``host_page_retention`` with its closure, ``post_attention_hop``, and turning that
+off turns it off; the target's executor events install the host policy once after warmup,
+after the post-MoE hop health check.
+"""
+
+import os
+import sys
+import types
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
+    dispatch,
+    runtime,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import plan_testlib  # noqa: E402
+
+B1 = (
+    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,router_gemm,"
+    "post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,rope_kv_append,"
+    "draft_loop,shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
+    "attention_output_tail_prefetch,attention_output_into_hop,routed_moe_into_hop,deferred_norm,"
+    "dense_mlp,deferred_input_prep"
+)
+
+
+def test_closure_is_attn_hop():
+    assert dispatch.REQUIRED_FEATURES["host_page_retention"] == frozenset({"post_attention_hop"})
+    assert "host_page_retention" not in dispatch.NUMERICAL_VARIANTS  # allocator policy only
+
+
+def test_selected_with_adapter(monkeypatch):
+    paths = frozenset(B1.split(",")) | {"host_page_retention"}
+    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
+    assert plan_testlib.select(dispatch, paths, monkeypatch) == paths
+    assert dispatch.is_enabled("host_page_retention")
+    plan = dispatch.resolve(
+        dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: "post_attention_hop"}
+    )
+    assert not plan.enabled("host_page_retention")
+
+
+def test_the_host_policy_is_installed_after_warmup_after_the_hop_health_check(monkeypatch):
+    calls = []
+    kernels = runtime.__name__.rsplit(".", 1)[0] + ".kernels"
+    monkeypatch.setitem(
+        sys.modules,
+        kernels + ".post_moe_hop",
+        types.SimpleNamespace(check_health=lambda *args, **kwargs: calls.append(("health", args))),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        kernels + ".host_page_retention",
+        types.SimpleNamespace(install_host_policy=lambda where: calls.append(("policy", where))),
+    )
+    paths = {"post_moe_hop", "host_page_retention"}
+    plan = types.SimpleNamespace(enabled=lambda name: name in paths)
+    events = runtime.GlmExecutorEvents(plan, types.SimpleNamespace(enable_attention_dp=False))
+
+    events.after_warmup()
+    assert calls == [("health", ("after_warmup",)), ("policy", "after_warmup")]
+    paths.discard("post_moe_hop")
+    calls.clear()
+    events.after_warmup()
+    assert calls == [("policy", "after_warmup")]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py
new file mode 100644
index 000000000000..adcb534b1114
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py
@@ -0,0 +1,246 @@
+# 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.
+"""indexer_layer_front selector and call sites.
+
+* the plan: indexer_layer_front runs only with its closure (absorbed_query, indexer_topk,
+  draft_absorbed_query, rope_kv_append, cross_rank_topk) and has no numerical variant;
+* the registered entry points and their call sites in the target (the absorbed-query packet,
+  the cross-rank top-k region and the owned ``forward_absorption_generation``);
+* the ``skip_reused_iq`` selector (0 / 1, anything else fails before any GPU
+  work) and the reuse condition it acts on, checked against the indexer's
+  own source;
+* deferred_input_prep's source map lists indexer_layer_front's directories and side selector.
+"""
+
+import inspect
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import plan_testlib  # noqa: E402
+
+CLOSURE = (
+    "absorbed_query",
+    "indexer_topk",
+    "draft_absorbed_query",
+    "rope_kv_append",
+    "cross_rank_topk",
+)
+
+
+@pytest.fixture
+def dispatch(monkeypatch):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
+        dispatch as d,
+    )
+
+    return d
+
+
+def _select(monkeypatch, d, features):
+    monkeypatch.setattr(d, "_log_info", lambda message: None)
+    return plan_testlib.select(d, features, monkeypatch)
+
+
+def test_selector_tables(dispatch):
+    assert "indexer_layer_front" in dispatch.SELECTOR_PATHS
+    assert dispatch.REQUIRED_FEATURES["indexer_layer_front"] == frozenset(CLOSURE)
+    assert "indexer_layer_front" not in dispatch.NUMERICAL_VARIANTS
+
+
+def test_closure_enforced(dispatch, monkeypatch):
+    """Turning off any path of the closure turns indexer_layer_front off with it."""
+    monkeypatch.setattr(dispatch, "_require", [False])
+    for base in CLOSURE:
+        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
+        assert not plan.enabled("indexer_layer_front"), base
+    assert "indexer_layer_front" in _select(
+        monkeypatch, dispatch, CLOSURE + ("indexer_layer_front",)
+    )
+    assert "indexer_layer_front" in _select(
+        monkeypatch, dispatch, CLOSURE + ("cross_rank_topk_cluster_merge", "indexer_layer_front")
+    )
+
+
+def test_entry_points_registered(dispatch):
+    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front  # noqa: E501, F401
+
+    for name in ("dsa_proj", "rope_available", "mla_rope_generation", "cp_indexer_decode"):
+        assert callable(dispatch.get_impl("indexer_layer_front", name)), name
+
+
+def test_call_sites_patched():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+        attention_forward,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        region,
+    )
+
+    assert 'get_impl("indexer_layer_front", "dsa_proj")' in inspect.getsource(
+        attention_forward._absorbed_query_packet
+    )
+    src = inspect.getsource(region.cross_rank_topk_decode)
+    assert src.index('"indexer_layer_front", "cp_indexer_decode"') < src.index(
+        '"cross_rank_topk_cluster_merge", "cp_indexer_decode"'
+    )
+    assert '"indexer_layer_front", "rope_available"' in inspect.getsource(
+        attention_forward.forward_absorption_generation
+    )
+
+
+def test_skip_reused_iq_is_off_in_this_target():
+    """Every runner starts with the option off and the entry never turns it on."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
+        query_proj,
+    )
+
+    init = inspect.getsource(query_proj.IndexerFrontQueryProjection.__init__)
+    assert "self.skip_reused_iq = False" in init
+    assert "skip_reused_iq" not in inspect.getsource(query_proj._runner)
+
+
+def test_reuses_shared_topk_condition():
+    """The option fires only under the indexer's own ``reuse_topk`` condition."""
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
+        query_proj,
+    )
+
+    def mla(share):
+        return SimpleNamespace(mqa=SimpleNamespace(indexer=SimpleNamespace(mtp_index_share=share)))
+
+    def md(loop, skip, rows):
+        return SimpleNamespace(
+            in_mtp_draft_loop=loop, indexer_skip_topk=skip, shared_topk_indices=rows
+        )
+
+    assert query_proj.reuses_shared_topk(mla(True), md(True, True, object()))
+    assert not query_proj.reuses_shared_topk(mla(False), md(True, True, object()))
+    assert not query_proj.reuses_shared_topk(mla(True), md(False, True, object()))
+    assert not query_proj.reuses_shared_topk(mla(True), md(True, False, object()))
+    assert not query_proj.reuses_shared_topk(mla(True), md(True, True, None))
+
+
+def _conjuncts(node):
+    import ast
+
+    if isinstance(node, ast.BoolOp) and isinstance(node.op, ast.And):
+        return [c for v in node.values for c in _conjuncts(v)]
+    return [node]
+
+
+def _normalized(node):
+    """``getattr(obj, "x", default)`` -> ``obj.x``; the indexer and metadata names unified."""
+    import ast
+
+    roles = {"self": "indexer", "metadata": "metadata", "attn_metadata": "metadata"}
+
+    class Norm(ast.NodeTransformer):
+        def visit_Call(self, call):
+            self.generic_visit(call)
+            if isinstance(call.func, ast.Name) and call.func.id == "getattr":
+                return ast.Attribute(value=call.args[0], attr=call.args[1].value, ctx=ast.Load())
+            return call
+
+        def visit_Name(self, name):
+            return ast.Name(id=roles.get(name.id, name.id), ctx=ast.Load())
+
+    return ast.unparse(Norm().visit(node))
+
+
+def test_reuses_shared_topk_matches_indexer_source():
+    """``reuses_shared_topk`` is term for term the indexer's ``reuse_topk``.
+
+    The branch the indexer takes under it reads none of ``q_fp8`` / ``q_scale`` / ``weights``
+    (the outputs the ``skip_reused_iq`` option leaves unwritten).
+    """
+    import ast
+    import textwrap
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
+        query_proj,
+    )
+    from tensorrt_llm._torch.attention.backends.sparse.dsa import indexer as indexer_mod
+
+    tree = ast.parse(textwrap.dedent(inspect.getsource(indexer_mod.Indexer.sparse_attn_indexer)))
+    assigns = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Assign)
+        and any(isinstance(t, ast.Name) and t.id == "reuse_topk" for t in n.targets)
+    ]
+    assert len(assigns) == 1, "the indexer assigns reuse_topk once"
+    indexer_terms = sorted(_normalized(c) for c in _conjuncts(assigns[0].value))
+
+    fn = ast.parse(textwrap.dedent(inspect.getsource(query_proj.reuses_shared_topk)))
+    aliases = {
+        ast.unparse(s.targets[0]): ast.unparse(s.value)
+        for s in fn.body[0].body
+        if isinstance(s, ast.Assign)
+    }
+    assert aliases == {"indexer": "mla.mqa.indexer"}, aliases
+    ret = next(n for n in ast.walk(fn) if isinstance(n, ast.Return))
+    value = ret.value.args[0] if isinstance(ret.value, ast.Call) else ret.value  # bool(...)
+    indexer_front_terms = sorted(_normalized(c) for c in _conjuncts(value))
+    assert indexer_front_terms == indexer_terms, (indexer_front_terms, indexer_terms)
+
+    branches = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.If)
+        and isinstance(n.test, ast.BoolOp)
+        and ast.unparse(n.test.values[-1]) == "reuse_topk"
+    ]
+    assert len(branches) == 1
+    names = {n.id for s in branches[0].body for n in ast.walk(s) if isinstance(n, ast.Name)}
+    assert not names & {"q_fp8", "q_scale", "weights"}, names
+
+
+def test_entry_declines_without_feature(dispatch, monkeypatch):
+    """With indexer_layer_front off the entry returns None before touching the module."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
+        query_proj,
+    )
+
+    _select(monkeypatch, dispatch, [])
+
+    class Boom:
+        def __getattr__(self, name):
+            raise AssertionError(f"module touched: {name}")
+
+    assert not dispatch.is_enabled("indexer_layer_front")
+    assert query_proj.dsa_proj(Boom(), None, None, None) is None
+
+
+@pytest.mark.skipif(
+    not os.environ.get("GLM_MOE_DSA_INDEXER_LAYER_FRONT_JIT_LIB"),
+    reason="needs the development build of the indexer_layer_front op",
+)
+def test_rope_op_available():
+    import torch
+
+    import tensorrt_llm  # noqa: F401
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
+        rope,
+    )
+
+    if not rope.rope_available():
+        torch.ops.load_library(os.environ["GLM_MOE_DSA_INDEXER_LAYER_FRONT_JIT_LIB"])
+    assert rope.rope_available()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py
new file mode 100644
index 000000000000..57926839977b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py
@@ -0,0 +1,195 @@
+# 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.
+"""Eligibility of the cross_rank_topk context-partitioned indexer region (no GPU needed).
+
+``cross_rank_policy.py`` and ``policy.py`` are loaded from the indexer_topk package directory
+without executing the package ``__init__`` (which imports torch).
+"""
+
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+_INDEXER_TOPK_DIR = (
+    Path(__file__).resolve().parents[9]
+    / "tensorrt_llm"
+    / "_torch"
+    / "_experimental"
+    / "modeling_v2"
+    / "models"
+    / "glm_moe_dsa"
+    / "shared"
+    / "kernels"
+    / "indexer_topk"
+)
+
+
+def _load_cp_policy():
+    name = "_cross_rank_policy_pkg"
+    if name not in sys.modules:
+        pkg = types.ModuleType(name)
+        pkg.__path__ = [str(_INDEXER_TOPK_DIR)]
+        sys.modules[name] = pkg
+    return importlib.import_module(f"{name}.cross_rank_policy")
+
+
+cp_policy = _load_cp_policy()
+
+SERVED_MAX_SEQ_LEN = (1 << 20) + 14
+# Largest per-rank share of the served envelope at world 8, and the kernel limit on sm_103.
+LOCAL_CAPACITY_DRAFT_LOOP = 131136
+MAX_LOCAL_CAPACITY = 179712
+
+_SERVED_TP8 = dict(
+    world_size=8,
+    op_available=True,
+    local_capacity=LOCAL_CAPACITY_DRAFT_LOOP,
+    max_local_capacity=MAX_LOCAL_CAPACITY,
+    allow_draft_step=False,
+    sm_version=103,
+    cute_dsl_available=True,
+    native_uses_cute_dsl_logits=True,
+    use_fp4_indexer=False,
+    use_gvr_emission=False,
+    needs_gvr_prior=False,
+    is_ragged_verify=False,
+    expanded_mtp_buffers=False,
+    in_mtp_draft_loop=False,
+    attention_dp=False,
+    topk_decode_impl="cute_dsl_gvr",
+    topk_self_sampling=True,
+    index_topk=2048,
+    compress_ratio=1,
+    num_heads=32,
+    head_dim=128,
+    query_dtype="torch.float8_e4m3fn",
+    query_shape=(6, 32, 128),
+    query_contiguous=True,
+    weights_dtype="torch.float32",
+    weights_shape=(6, 32),
+    weights_contiguous=True,
+    k_cache_dtype="torch.uint8",
+    k_cache_shape=(16384, 64, 1, 132),
+    block_table_cols=16388,
+    max_seq_len=SERVED_MAX_SEQ_LEN,
+    num_contexts=0,
+    token_offset=0,
+    num_gen_tokens=6,
+    num_generations=1,
+    uniform_verify_width=True,
+)
+
+
+def test_tp8_decode_is_eligible():
+    d = cp_policy.decide_cp(**_SERVED_TP8)
+    assert d.eligible, d.reason
+
+
+@pytest.mark.parametrize("world", [4, 8])
+def test_supported_world_sizes(world):
+    assert cp_policy.decide_cp(**dict(_SERVED_TP8, world_size=world)).eligible
+
+
+FALLBACKS = {
+    "two_requests": dict(
+        num_generations=2, num_gen_tokens=12, query_shape=(12, 32, 128), weights_shape=(12, 32)
+    ),
+    "world_2": dict(world_size=2),
+    "world_16": dict(world_size=16),
+    "op_missing": dict(op_available=False),
+    "local_width_above_kernel_limit": dict(local_capacity=MAX_LOCAL_CAPACITY + 64),
+    "mtp_draft_step": dict(in_mtp_draft_loop=True),
+    "attention_dp": dict(attention_dp=True),
+    "indexer_topk_envelope_page32": dict(k_cache_shape=(32768, 32, 1, 132)),
+    "indexer_topk_envelope_next_n4": dict(
+        num_gen_tokens=4, query_shape=(4, 32, 128), weights_shape=(4, 32)
+    ),
+    "indexer_topk_envelope_max_seq_len": dict(max_seq_len=SERVED_MAX_SEQ_LEN + 1),
+    "indexer_topk_envelope_sm100": dict(sm_version=100),
+    "indexer_topk_envelope_fp4": dict(use_fp4_indexer=True),
+    "indexer_topk_envelope_ragged": dict(is_ragged_verify=True),
+    "indexer_topk_envelope_mixed_batch": dict(num_contexts=1),
+}
+
+
+@pytest.mark.parametrize("name", sorted(FALLBACKS))
+def test_native_fallback_cases(name):
+    d = cp_policy.decide_cp(**dict(_SERVED_TP8, **FALLBACKS[name]))
+    assert not d.eligible
+    assert d.reason
+
+
+def test_draft_step_only_when_allowed():
+    inputs = dict(_SERVED_TP8, in_mtp_draft_loop=True)
+    assert not cp_policy.decide_cp(**inputs).eligible
+    assert cp_policy.decide_cp(**dict(inputs, allow_draft_step=True)).eligible
+
+
+# The rank-local facts of a decision: a rank's build and device.
+_LOCAL_KEYS = ("op_available", "max_local_capacity", "sm_version", "cute_dsl_available")
+
+
+def _uniform(inputs):
+    return {k: v for k, v in inputs.items() if k not in _LOCAL_KEYS}
+
+
+def _local(inputs):
+    return dict(
+        op_available=inputs["op_available"],
+        local_capacity=inputs["local_capacity"],
+        max_local_capacity=inputs["max_local_capacity"],
+        sm_version=inputs["sm_version"],
+        cute_dsl_available=inputs["cute_dsl_available"],
+    )
+
+
+@pytest.mark.parametrize("name", sorted(FALLBACKS))
+def test_uniform_and_local_parts_reproduce_decide_cp(name):
+    """``decide_cp`` = rank-uniform decision AND no rank-local reason.
+
+    The adapter decides a forward with ``decide_cp_uniform`` (identical on
+    every rank) and hands ``rank_local_reason`` to the collective workspace
+    creation, so a fact of one rank's build or device keeps every rank native
+    instead of moving that rank alone out of the exchange.
+    """
+    inputs = dict(_SERVED_TP8, **FALLBACKS[name])
+    full = cp_policy.decide_cp(**inputs)
+    uniform = cp_policy.decide_cp_uniform(**_uniform(inputs))
+    local = cp_policy.rank_local_reason(**_local(inputs))
+    assert full.eligible == (uniform.eligible and not local), (full, uniform, local)
+
+
+@pytest.mark.parametrize(
+    "fact, value",
+    [
+        ("op_available", False),
+        ("sm_version", 100),
+        ("cute_dsl_available", False),
+        ("max_local_capacity", LOCAL_CAPACITY_DRAFT_LOOP - 64),
+    ],
+)
+def test_rank_local_facts_never_decide_a_forward(fact, value):
+    """A rank-local fact never changes the per-forward decision.
+
+    A rank whose own build or device differs makes the same decision as its
+    peers; only its local reason differs.
+    """
+    inputs = dict(_SERVED_TP8, **{fact: value})
+    assert cp_policy.decide_cp_uniform(**_uniform(inputs)).eligible
+    assert cp_policy.rank_local_reason(**_local(inputs))
+    assert not cp_policy.rank_local_reason(**_local(_SERVED_TP8))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py
new file mode 100644
index 000000000000..c1472647f66c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py
@@ -0,0 +1,1081 @@
+# 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.
+"""GPU correctness of the indexer_topk indexer decode package against native TRT-LLM.
+
+Native oracle: the CuTe-DSL FP8 paged-MQA logits launch exactly as
+``Indexer.sparse_attn_indexer`` issues it for the GEN DSL path (wave-aware
+row-group split from ``_pick_dsl_expand`` with expanded lengths, page tables
+and schedule), followed by the same ``TopK`` module. Most fixtures are
+synthetic (random FP8 keys with UE8M0 scales, random queries and weights).
+``test_real_weight_fixture`` uses a real GLM-5.2 layer-0 fixture when
+``GLM_MOE_DSA_INDEXER_TOPK_REAL_FIXTURE`` points at it.
+"""
+
+import os
+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+
+if not torch.cuda.is_available() or not IS_CUTLASS_DSL_AVAILABLE:
+    pytest.skip("CUDA and the CUTLASS DSL are required", allow_module_level=True)
+if get_sm_version() not in (100, 103):
+    pytest.skip("indexer_topk targets SM100/SM103", allow_module_level=True)
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+    GlmIndexer,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+    indexer_decode as indexer_topk,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk.logits_runner import (  # noqa: E501
+    DEFAULT_CONFIG,
+    LogitsConfig,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa import indexer as native_indexer_mod
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _pick_dsl_expand
+from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation
+from tensorrt_llm.deep_gemm import get_num_sms, get_paged_mqa_logits_metadata
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+
+HEADS = 32
+HEAD_DIM = 128
+BLOCK_KV = 64
+TOP_K = 2048
+NEXT_N = 6
+SCHEDULE_BLOCK_KV = 64
+FP8_ATOMS = (1, 2, 3, 4)
+_DEV = torch.device("cuda")
+
+
+def _ue8m0(x):
+    return torch.pow(2.0, torch.ceil(torch.log2(x.clamp_min(1e-12))))
+
+
+def make_pool(num_blocks, gen, k_values=None):
+    """Fused FP8 indexer K pages ``[num_blocks, 64, 1, 132]`` plus their parts."""
+    if k_values is None:
+        k_values = torch.randn(num_blocks, BLOCK_KV, HEAD_DIM, device=_DEV, generator=gen)
+    scale = _ue8m0(k_values.abs().amax(dim=-1) / 448.0)
+    k_fp8 = (k_values / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
+    data = k_fp8.view(torch.uint8).reshape(num_blocks, BLOCK_KV * HEAD_DIM)
+    scales = scale.float().contiguous().view(torch.uint8).reshape(num_blocks, BLOCK_KV * 4)
+    fused = (
+        torch.cat([data, scales], dim=1).contiguous().view(num_blocks, BLOCK_KV, 1, HEAD_DIM + 4)
+    )
+    return fused, k_fp8, scale.float()
+
+
+def make_top_k():
+    return TopK(
+        TOP_K, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=True
+    )
+
+
+class Case:
+    """B requests x NEXT_N verify rows over one layer's paged FP8 K cache."""
+
+    def __init__(
+        self,
+        kv_lens,
+        seed=0,
+        max_seq_len=1 << 20,
+        shuffle=True,
+        q_scale=1.0,
+        weight_sign=None,
+        distinct_pages=None,
+        block_table_cols=None,
+        boost_from_token=None,
+    ):
+        """Build the case.
+
+        Args:
+            block_table_cols: Block-table width; defaults to
+                ``ceil(max_seq_len / 64)``.
+            boost_from_token: When set, the pages holding logical tokens at or
+                beyond this position get 64x larger keys (so 64x larger
+                per-token scales and identical FP8 values), which puts those
+                tokens into the top-k when the weights are positive.
+        """
+        gen = torch.Generator(device=_DEV).manual_seed(seed)
+        self.B = len(kv_lens)
+        self.kv_lens_list = list(kv_lens)
+        self.max_seq_len = max_seq_len
+        blocks = [max(1, -(-n // BLOCK_KV)) for n in kv_lens]
+        max_blocks = block_table_cols or -(-max_seq_len // BLOCK_KV)
+        assert max(blocks) <= max_blocks
+        num_blocks = (distinct_pages or sum(blocks)) + 4
+        self.pool, self.k_fp8, self.k_scale = make_pool(num_blocks, gen)
+        bt = torch.zeros(self.B, max_blocks, dtype=torch.int32, device=_DEV)
+        base = 0
+        for b, nb in enumerate(blocks):
+            if distinct_pages:
+                pages = torch.randint(0, distinct_pages, (nb,), device=_DEV, generator=gen)
+            else:
+                pages = torch.arange(base, base + nb, device=_DEV)
+                if shuffle:
+                    pages = pages[torch.randperm(nb, device=_DEV, generator=gen)]
+                base += nb
+            bt[b, :nb] = pages.to(torch.int32)
+        self.block_table = bt
+        if boost_from_token is not None:
+            k_values = torch.randn(num_blocks, BLOCK_KV, HEAD_DIM, device=_DEV, generator=gen)
+            first_page = boost_from_token // BLOCK_KV
+            for b, nb in enumerate(blocks):
+                if nb > first_page:
+                    k_values[bt[b, first_page:nb].long()] *= 64.0
+            self.pool, self.k_fp8, self.k_scale = make_pool(num_blocks, gen, k_values)
+        self.kv_lens = torch.tensor(kv_lens, dtype=torch.int32, device=_DEV)
+        q = torch.randn(self.B, NEXT_N, HEADS, HEAD_DIM, device=_DEV, generator=gen) * q_scale
+        self.q = q.to(torch.float8_e4m3fn)
+        w = torch.randn(self.B * NEXT_N, HEADS, device=_DEV, generator=gen)
+        if weight_sign is not None:
+            w = w.abs() * weight_sign
+        self.weights = w.float()
+        self.num_sms = get_num_sms()
+        self.schedule = get_paged_mqa_logits_metadata(
+            self.kv_lens.view(-1, 1), SCHEDULE_BLOCK_KV, self.num_sms
+        )
+        self.factor, self.atom = _pick_dsl_expand(
+            NEXT_N,
+            batch_size=self.B,
+            max_ctx=max(kv_lens),
+            num_sms=self.num_sms,
+            kernel_atoms=FP8_ATOMS,
+        )
+        self.kv_lens_exp = self.kv_lens.repeat_interleave(self.factor)
+        self.block_table_exp = self.block_table.repeat_interleave(self.factor, dim=0)
+        self.schedule_exp = get_paged_mqa_logits_metadata(
+            self.kv_lens_exp.view(-1, 1), SCHEDULE_BLOCK_KV, self.num_sms
+        )
+        rows = self.B * NEXT_N
+        self.radix_idx = torch.empty(rows, 10, TOP_K, dtype=torch.int32, device=_DEV)
+        self.radix_val = torch.empty(rows, 10, TOP_K, dtype=torch.float32, device=_DEV)
+
+    @classmethod
+    def from_fixture(cls, fixture, kv_len, shuffle, seed=0):
+        """Real-weight case: the fixture's pages, queries and weights."""
+        case = cls([kv_len], seed=seed)
+        gen = torch.Generator(device=_DEV).manual_seed(seed)
+        nb = -(-kv_len // BLOCK_KV)
+        pages = fixture["pool"][:nb].to(_DEV)
+        perm = torch.arange(nb, device=_DEV)
+        if shuffle:
+            perm = perm[torch.randperm(nb, device=_DEV, generator=gen)]
+        pool = torch.zeros(nb + 4, BLOCK_KV, 1, HEAD_DIM + 4, dtype=torch.uint8, device=_DEV)
+        pool[perm] = pages
+        case.pool = pool
+        flat = pool.view(nb + 4, -1)
+        case.k_fp8 = flat[:, : BLOCK_KV * HEAD_DIM].contiguous().view(torch.float8_e4m3fn)
+        case.k_fp8 = case.k_fp8.view(nb + 4, BLOCK_KV, HEAD_DIM)
+        case.k_scale = flat[:, BLOCK_KV * HEAD_DIM :].contiguous().view(torch.float32)
+        case.block_table.zero_()
+        case.block_table[0, :nb] = perm.to(torch.int32)
+        case.block_table_exp = case.block_table.repeat_interleave(case.factor, dim=0)
+        case.q = fixture["q"][kv_len].to(_DEV).view(1, NEXT_N, HEADS, HEAD_DIM)
+        case.weights = fixture["weights"][kv_len].to(_DEV).float()
+        return case
+
+    @property
+    def row_lens(self):
+        return [
+            max(0, self.kv_lens_list[r // NEXT_N] - NEXT_N + r % NEXT_N + 1)
+            for r in range(self.B * NEXT_N)
+        ]
+
+    def native_logits(self):
+        """Native GEN DSL launch, including the wave-aware row-group split."""
+        if self.factor > 1:
+            q = self.q.reshape(self.B * self.factor, self.atom, HEADS, HEAD_DIM)
+            return torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits(
+                q,
+                self.pool,
+                self.weights,
+                self.kv_lens_exp,
+                self.block_table_exp,
+                self.schedule_exp,
+                self.max_seq_len,
+            )
+        return torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits(
+            self.q,
+            self.pool,
+            self.weights,
+            self.kv_lens,
+            self.block_table,
+            self.schedule,
+            self.max_seq_len,
+        )
+
+    def indexer_topk_logits(self, config=None):
+        return indexer_topk.paged_mqa_logits(
+            self.q,
+            self.pool,
+            self.weights,
+            self.kv_lens,
+            self.block_table,
+            self.schedule,
+            self.max_seq_len,
+            config=config,
+        )
+
+    def native_region(self, top_k, out):
+        top_k(
+            self.native_logits(),
+            out,
+            is_prefill=False,
+            sequence_lengths=self.kv_lens,
+            scan_lengths=self.kv_lens,
+            row_kv_lens=None,
+            next_n=NEXT_N,
+            max_seq_len=self.max_seq_len,
+            gvr_ext_kwargs=None,
+            radix_aux_indices=self.radix_idx,
+            radix_aux_logits=self.radix_val,
+        )
+        return out
+
+    def indexer_topk_region(self, top_k, out, selector=None, config=None):
+        return indexer_topk.decode_topk(
+            q=self.q,
+            weights=self.weights,
+            k_cache=self.pool,
+            block_table=self.block_table,
+            context_lens=self.kv_lens,
+            kv_lens=self.kv_lens,
+            schedule_meta=self.schedule,
+            max_seq_len=self.max_seq_len,
+            top_k=top_k,
+            out=out,
+            radix_aux_indices=self.radix_idx,
+            radix_aux_logits=self.radix_val,
+            logits_config=config,
+            selector=selector,
+        )
+
+    def new_out(self, fill=-7):
+        return torch.full((self.B * NEXT_N, TOP_K), fill, dtype=torch.int32, device=_DEV)
+
+
+def assert_logits_equal(case, native, cand):
+    for r, n in enumerate(case.row_lens):
+        if n > 0:
+            assert torch.equal(native[r, :n], cand[r, :n]), f"row {r}: logits differ"
+
+
+def assert_selection_contract(case, logits, sel):
+    """Native the index_selection contract semantics, tie-aware (value multiset vs torch.topk)."""
+    for r, n in enumerate(case.row_lens):
+        idx = sel[r].long()
+        if n <= 0:
+            assert bool((idx == -1).all()), f"row {r}: empty row must be all -1"
+        elif n <= TOP_K:
+            assert torch.equal(torch.sort(idx[:n]).values, torch.arange(n, device=_DEV)), r
+            assert bool((idx[n:] == -1).all()), f"row {r}: short row tail must be -1"
+        else:
+            assert int(idx.min()) >= 0 and int(idx.max()) < n, f"row {r}: index out of range"
+            assert int(torch.unique(idx).numel()) == TOP_K, f"row {r}: duplicate index"
+            ref = torch.topk(logits[r, :n], TOP_K).values
+            got = torch.sort(logits[r].gather(0, idx) + 0.0, descending=True).values
+            assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), (
+                f"row {r}: selected value multiset is not the exact top-k"
+            )
+
+
+def reference_logits(case):
+    """Independent fp32 oracle: scale_j * sum_h w_h * relu(q_h . k_j)."""
+    k_all = case.k_fp8.float()
+    out = torch.full(
+        (case.B * NEXT_N, max(case.kv_lens_list)), float("nan"), device=_DEV, dtype=torch.float32
+    )
+    for b, n in enumerate(case.kv_lens_list):
+        nb = -(-n // BLOCK_KV)
+        pages = case.block_table[b, :nb].long()
+        k = k_all[pages].reshape(-1, HEAD_DIM)[:n]
+        s = case.k_scale[pages].reshape(-1)[:n]
+        qk = torch.matmul(case.q[b].float().reshape(NEXT_N * HEADS, HEAD_DIM), k.t())
+        qk = torch.relu(qk).view(NEXT_N, HEADS, n)
+        w = case.weights[b * NEXT_N : (b + 1) * NEXT_N]
+        out[b * NEXT_N : (b + 1) * NEXT_N, :n] = (w.unsqueeze(-1) * qk).sum(dim=1) * s
+    return out
+
+
+# --- logits ---------------------------------------------------------------
+
+LENGTHS = [3, 6, 7, 2047, 2048, 2049, 2054, 32768, 131072, 491061, (1 << 20) - 1]
+
+
+# None runs DEFAULT_CONFIG (rows, 8 register-cached weights), the kernel the
+# adapter hook ships; the rowsplit variants are opt-in alternatives.
+VARIANTS = {
+    "default_rows_draft_loop": None,
+    "rowsplit": LogitsConfig(variant="rowsplit"),
+    "rowsplit_wg3": LogitsConfig(variant="rowsplit", num_math_wg=3),
+}
+
+
+def test_default_config_is_rows_draft_loop():
+    assert DEFAULT_CONFIG == LogitsConfig(variant="rows", num_w_in_reg=8)
+
+
+@pytest.mark.parametrize("variant", sorted(VARIANTS))
+@pytest.mark.parametrize("kv_len", LENGTHS)
+def test_logits_bit_exact_vs_native(kv_len, variant):
+    case = Case([kv_len], seed=kv_len)
+    native = case.native_logits()[:, :kv_len].clone()
+    cand = case.indexer_topk_logits(VARIANTS[variant])[:, :kv_len]
+    assert_logits_equal(case, native, cand)
+
+
+@pytest.mark.parametrize("variant", sorted(VARIANTS))
+@pytest.mark.parametrize(
+    "kv_lens",
+    [[4224, 64], [491061, 300], [7, 131072, 2049], [1, 2, 3, 4]],
+    ids=lambda v: "-".join(map(str, v)),
+)
+def test_logits_bit_exact_multi_request(kv_lens, variant):
+    # Odd and even tile counts per request and schedule splits that end at
+    # request boundaries.
+    case = Case([max(n, NEXT_N) for n in kv_lens], seed=sum(kv_lens))
+    width = max(case.kv_lens_list)
+    native = case.native_logits()[:, :width].clone()
+    cand = case.indexer_topk_logits(VARIANTS[variant])[:, :width]
+    assert_logits_equal(case, native, cand)
+
+
+@pytest.mark.parametrize("kv_len", [2049, 32768, 131072])
+def test_logits_vs_reference(kv_len):
+    case = Case([kv_len], seed=7 * kv_len)
+    cand = case.indexer_topk_logits()[:, :kv_len]
+    ref = reference_logits(case)
+    for r, n in enumerate(case.row_lens):
+        # Same tolerances as the native CuTe-DSL FP8 paged-MQA logits test.
+        torch.testing.assert_close(cand[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
+
+
+# --- selection --------------------------------------------------------------
+
+
+SELECTORS = ["native", "radix"]
+
+
+@pytest.mark.parametrize("selector", SELECTORS)
+@pytest.mark.parametrize("kv_len", LENGTHS)
+def test_selection_matches_native(kv_len, selector):
+    case = Case([kv_len], seed=11 + kv_len)
+    top_k = make_top_k()
+    native = case.native_region(top_k, case.new_out()).clone()
+    logits = case.native_logits()[:, :kv_len].clone()
+    cand = case.indexer_topk_region(top_k, case.new_out(), selector=selector)
+    assert_selection_contract(case, logits, cand)
+    # Continuous random logits: the top-k set is unique, so any exact
+    # selection engine must return the native set.
+    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
+
+
+@pytest.mark.parametrize("selector", SELECTORS)
+@pytest.mark.parametrize(
+    "kv_lens",
+    [[30000, 64], [491061, 7, 2048, 2055], [131072, 131073]],
+    ids=lambda v: "-".join(map(str, v)),
+)
+def test_selection_multi_request(kv_lens, selector):
+    case = Case(kv_lens, seed=sum(kv_lens) + 1)
+    top_k = make_top_k()
+    width = max(kv_lens)
+    logits = case.native_logits()[:, :width].clone()
+    native = case.native_region(top_k, case.new_out()).clone()
+    cand = case.indexer_topk_region(top_k, case.new_out(), selector=selector)
+    assert_selection_contract(case, logits, cand)
+    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
+
+
+@pytest.mark.parametrize("selector", SELECTORS)
+@pytest.mark.parametrize("kind", ["duplicate_pages", "all_zero", "negative_weights"])
+def test_ties(kind, selector):
+    kv_len = 40000
+    if kind == "duplicate_pages":
+        # 8 distinct pages repeated across the context: every logit value
+        # occurs ~78 times, so the top-k boundary falls inside a tie class.
+        case = Case([kv_len], seed=3, distinct_pages=8)
+    elif kind == "all_zero":
+        case = Case([kv_len], seed=4, q_scale=0.0)
+    else:
+        case = Case([kv_len], seed=5, weight_sign=-1.0)
+    top_k = make_top_k()
+    native = case.native_region(top_k, case.new_out()).clone()
+    logits = case.native_logits()[:, :kv_len].clone()
+    cand = case.indexer_topk_region(top_k, case.new_out(), selector=selector)
+    assert_logits_equal(case, logits, case.indexer_topk_logits()[:, :kv_len])
+    assert_selection_contract(case, logits, cand)
+    # Characterize tie handling (all-zero logits overflow the radix
+    # candidate buffer and exercise its full-row fallback).
+    same_ids = torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
+    print(f"[ties:{kind}:{selector}] identical index sets vs native: {same_ids}")
+
+
+def test_padded_batch_rows():
+    # Request 1 stands in for a CUDA-graph padding request with a short cache.
+    case = Case([30000, 64], seed=21)
+    top_k = make_top_k()
+    native = case.native_region(top_k, case.new_out()).clone()
+    logits = case.native_logits()[:, :30000].clone()
+    cand = case.indexer_topk_region(top_k, case.new_out())
+    assert_logits_equal(case, logits, case.indexer_topk_logits()[:, :30000])
+    assert_selection_contract(case, logits, cand)
+    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
+
+
+_FIXTURE_PATH = os.environ.get("GLM_MOE_DSA_INDEXER_TOPK_REAL_FIXTURE", "")
+
+
+@pytest.mark.skipif(
+    not os.path.isfile(_FIXTURE_PATH), reason="GLM_MOE_DSA_INDEXER_TOPK_REAL_FIXTURE not provided"
+)
+@pytest.mark.parametrize("shuffle", [False, True], ids=["contiguous", "shuffled"])
+def test_real_weight_fixture(shuffle):
+    fixture = torch.load(_FIXTURE_PATH, map_location="cpu", weights_only=False)
+    print(f"[real] {fixture['label']} sha={fixture['token_ids_sha256'][:16]}")
+    top_k = make_top_k()
+    for kv_len in fixture["kv_lens"]:
+        case = Case.from_fixture(fixture, kv_len, shuffle)
+        logits = case.native_logits()[:, :kv_len].clone()
+        assert_logits_equal(case, logits, case.indexer_topk_logits()[:, :kv_len])
+        native = case.native_region(top_k, case.new_out()).clone()
+        cand = case.indexer_topk_region(top_k, case.new_out())
+        assert_selection_contract(case, logits, cand)
+        assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
+        split = LogitsConfig(variant="rowsplit")
+        assert_logits_equal(case, logits, case.indexer_topk_logits(split)[:, :kv_len])
+        radix = case.indexer_topk_region(top_k, case.new_out(), selector="radix", config=split)
+        assert_selection_contract(case, logits, radix)
+        assert torch.equal(torch.sort(native, dim=1).values, torch.sort(radix, dim=1).values)
+        if kv_len <= 131072:
+            ref = reference_logits(case)
+            for r, n in enumerate(case.row_lens):
+                torch.testing.assert_close(logits[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
+        # Boundary sensitivity of the real distribution: gap between the k-th
+        # and (k+1)-th largest logits, and overlap of adjacent verify rows.
+        for r, n in enumerate(case.row_lens):
+            vals = torch.topk(logits[r, :n], TOP_K + 1).values
+            gap = (vals[TOP_K - 1] - vals[TOP_K]).item()
+            overlap = (
+                torch.isin(cand[r], cand[r - 1]).float().mean().item() if r > 0 else float("nan")
+            )
+            print(
+                f"[real] kv={kv_len} row={r} kth={vals[TOP_K - 1].item():.6g} "
+                f"gap={gap:.3g} overlap_prev_row={overlap:.4f}"
+            )
+
+
+# --- CUDA graph replay --------------------------------------------------------
+
+
+GRAPH_CONFIGS = {"default_rows_draft_loop": None, "rowsplit": LogitsConfig(variant="rowsplit")}
+GRAPH_STEPS = [20000, 20006, 20012, 35000, 9000, 65530]
+
+
+def _graph_step_lens(kv_len, batch):
+    # Request 0 follows the step length; later requests are shorter, down to
+    # a 6-token request (verify rows of 1..6 keys, padding-like).
+    return [max(NEXT_N, kv_len - 4099 * b) for b in range(batch)]
+
+
+def _native_oracle(lens, max_seq_len, pool, k_fp8, k_scale, block_table, q, weights, kv_lens):
+    """Eager native launch state for the given tensors (no new pool)."""
+    batch = len(lens)
+    ref = Case([NEXT_N] * batch, seed=0, max_seq_len=max_seq_len)
+    ref.kv_lens_list = list(lens)
+    ref.pool, ref.k_fp8, ref.k_scale = pool, k_fp8, k_scale
+    ref.block_table, ref.q, ref.weights, ref.kv_lens = block_table, q, weights, kv_lens
+    ref.schedule = get_paged_mqa_logits_metadata(kv_lens.view(-1, 1), 64, ref.num_sms)
+    ref.factor, ref.atom = _pick_dsl_expand(
+        NEXT_N, batch_size=batch, max_ctx=max(lens), num_sms=ref.num_sms, kernel_atoms=FP8_ATOMS
+    )
+    ref.kv_lens_exp = kv_lens.repeat_interleave(ref.factor)
+    ref.block_table_exp = block_table.repeat_interleave(ref.factor, dim=0)
+    ref.schedule_exp = get_paged_mqa_logits_metadata(ref.kv_lens_exp.view(-1, 1), 64, ref.num_sms)
+    return ref
+
+
+def _run_graph_replay(batch, logits_config, selector, max_seq_len, steps, block_table_cols=None):
+    """Capture indexer_topk once, then replay with new lengths, pages, tables and queries.
+
+    After every replay the indexer_topk logits (read from the graph-stable arena the
+    captured launch wrote) must be bit-exact to an eager native launch on the
+    same state, and the selection must match the native region.
+    """
+    case = Case(
+        _graph_step_lens(steps[0], batch),
+        seed=31,
+        max_seq_len=max_seq_len,
+        block_table_cols=block_table_cols,
+    )
+    # Persistent pool sized for the largest step; pages are rewritten between
+    # replays to model page reuse by later requests.
+    total_blocks = batch * -(-max_seq_len // BLOCK_KV) + 4
+    gen = torch.Generator(device=_DEV).manual_seed(99 + batch)
+    case.pool, case.k_fp8, case.k_scale = make_pool(total_blocks, gen)
+    top_k = make_top_k()
+    out = case.new_out()
+    q, weights = case.q.clone(), case.weights.clone()
+    kv_lens, block_table = case.kv_lens.clone(), case.block_table.clone()
+    schedule, pool = case.schedule.clone(), case.pool
+    arena = {}
+    indexer_topk_logits_fn = indexer_topk.paged_mqa_logits
+
+    def recording_logits(*args, **kwargs):
+        arena["logits"] = indexer_topk_logits_fn(*args, **kwargs)
+        return arena["logits"]
+
+    def indexer_topk_call():
+        return indexer_topk.decode_topk(
+            q=q,
+            weights=weights,
+            k_cache=pool,
+            block_table=block_table,
+            context_lens=kv_lens,
+            kv_lens=kv_lens,
+            schedule_meta=schedule,
+            max_seq_len=max_seq_len,
+            top_k=top_k,
+            out=out,
+            radix_aux_indices=case.radix_idx,
+            radix_aux_logits=case.radix_val,
+            logits_config=logits_config,
+            selector=selector,
+        )
+
+    with patch.object(indexer_topk, "paged_mqa_logits", side_effect=recording_logits):
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            indexer_topk_call()
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            indexer_topk_call()
+    captured_logits = arena["logits"]
+    assert captured_logits.shape == (batch * NEXT_N, max_seq_len)
+
+    for step, kv_len in enumerate(steps):
+        lens = _graph_step_lens(kv_len, batch)
+        new_pool, new_k, new_s = make_pool(total_blocks, gen)
+        pool.copy_(new_pool)
+        del new_pool
+        perm = torch.randperm(total_blocks, device=_DEV, generator=gen).to(torch.int32)
+        block_table.zero_()
+        base = 0
+        for b, n in enumerate(lens):
+            nb = -(-n // BLOCK_KV)
+            block_table[b, :nb] = perm[base : base + nb]
+            base += nb
+        kv_lens.copy_(torch.tensor(lens, dtype=torch.int32))
+        schedule.copy_(get_paged_mqa_logits_metadata(kv_lens.view(-1, 1), 64, case.num_sms))
+        q.copy_((torch.randn(q.shape, device=_DEV, generator=gen)).to(torch.float8_e4m3fn))
+        weights.copy_(torch.randn(weights.shape, device=_DEV, generator=gen))
+        out.fill_(-7)
+        graph.replay()
+        torch.cuda.synchronize()
+        ref = _native_oracle(
+            lens, max_seq_len, pool, new_k, new_s, block_table, q, weights, kv_lens
+        )
+        width = max(lens)
+        logits = ref.native_logits()[:, :width].clone()
+        assert_logits_equal(ref, logits, captured_logits[:, :width])
+        native = ref.native_region(make_top_k(), ref.new_out())
+        assert_selection_contract(ref, logits, out)
+        assert torch.equal(torch.sort(native, dim=1).values, torch.sort(out, dim=1).values), step
+
+
+@pytest.mark.parametrize("selector", SELECTORS)
+@pytest.mark.parametrize("config_name", sorted(GRAPH_CONFIGS))
+@pytest.mark.parametrize("batch", [1, 4])
+def test_cuda_graph_replay_growing_context_and_page_reuse(batch, config_name, selector):
+    _run_graph_replay(batch, GRAPH_CONFIGS[config_name], selector, 65536, GRAPH_STEPS)
+
+
+# --- served indexer max_seq_len (2**20 + 14) --------------------------------------
+
+# Served GEN: net max_seq_len 2**20, MTP K = 5, overlap scheduler on. The
+# indexer max_seq_len is 2**20 + 5 + 4 + 5 and KV cache manager V2 sizes the
+# block table at round_up(ceil((2**20 + 14 + 10) / 64), 4) columns; manager V1
+# would use ceil((2**20 + 14) / 64) columns.
+SERVED_MAX_SEQ_LEN = (1 << 20) + 14
+V2_BLOCK_TABLE_COLS = 16388
+V1_BLOCK_TABLE_COLS = 16385
+# Live indexer lengths around 2**20 up to the served bound (target layers
+# reach 2**20 + 10 with the overlap scheduler; the bound itself is included).
+BOUND_LENGTHS = [
+    (1 << 20) - 6,
+    (1 << 20) - 1,
+    1 << 20,
+    (1 << 20) + 1,
+    (1 << 20) + 5,
+    (1 << 20) + 10,
+    (1 << 20) + 13,
+    SERVED_MAX_SEQ_LEN,
+]
+
+
+def _bound_case(kv_lens, seed, **kwargs):
+    return Case(
+        kv_lens,
+        seed=seed,
+        max_seq_len=SERVED_MAX_SEQ_LEN,
+        block_table_cols=V2_BLOCK_TABLE_COLS,
+        **kwargs,
+    )
+
+
+def test_policy_envelope_is_the_served_bound():
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+        policy,
+    )
+
+    assert policy.MAX_INDEXER_SEQ_LEN == SERVED_MAX_SEQ_LEN
+    assert V2_BLOCK_TABLE_COLS % 2 == 0 and V1_BLOCK_TABLE_COLS % 2 == 1
+
+
+@pytest.mark.parametrize("kv_len", BOUND_LENGTHS)
+def test_served_bound_logits_and_selection(kv_len):
+    case = _bound_case([kv_len], seed=kv_len % 9973)
+    native = case.native_logits()[:, :kv_len].clone()
+    cand = case.indexer_topk_logits()
+    assert cand.shape == (NEXT_N, SERVED_MAX_SEQ_LEN)
+    assert cand.stride(0) % 256 == 0 and cand.stride(0) >= SERVED_MAX_SEQ_LEN
+    assert_logits_equal(case, native, cand[:, :kv_len])
+    top_k = make_top_k()
+    native_sel = case.native_region(top_k, case.new_out()).clone()
+    sel = case.indexer_topk_region(top_k, case.new_out())
+    assert_selection_contract(case, native, sel)
+    assert torch.equal(torch.sort(native_sel, dim=1).values, torch.sort(sel, dim=1).values)
+
+
+@pytest.mark.parametrize("kv_len", [(1 << 20) + 1, (1 << 20) + 10, SERVED_MAX_SEQ_LEN])
+def test_served_bound_selects_ids_beyond_2pow20(kv_len):
+    """Keys at the tail dominate, so IDs >= 2**20 must be selected exactly."""
+    boost_from = (1 << 20) - 64
+    case = _bound_case([kv_len], seed=kv_len % 7919, weight_sign=1.0, boost_from_token=boost_from)
+    native = case.native_logits()[:, :kv_len].clone()
+    assert_logits_equal(case, native, case.indexer_topk_logits()[:, :kv_len])
+    top_k = make_top_k()
+    native_sel = case.native_region(top_k, case.new_out()).clone()
+    sel = case.indexer_topk_region(top_k, case.new_out())
+    assert_selection_contract(case, native, sel)
+    assert torch.equal(torch.sort(native_sel, dim=1).values, torch.sort(sel, dim=1).values)
+    beyond = 0
+    for r, n in enumerate(case.row_lens):
+        tail = torch.arange(boost_from, n, device=_DEV, dtype=torch.int32)
+        assert bool(torch.isin(tail, sel[r]).all()), f"row {r}: boosted tail not selected"
+        beyond += int((sel[r] >= (1 << 20)).sum())
+    assert beyond > 0, "no ID >= 2**20 was selected"
+
+
+@pytest.mark.parametrize(
+    "kv_lens",
+    [
+        [SERVED_MAX_SEQ_LEN, (1 << 20) + 1],
+        [SERVED_MAX_SEQ_LEN, (1 << 20) + 10, 1 << 20, (1 << 20) - 1],
+        [(1 << 20) + 7, 491061, NEXT_N, 2049],
+    ],
+    ids=lambda v: "-".join(map(str, v)),
+)
+def test_served_bound_multi_request(kv_lens):
+    case = _bound_case(kv_lens, seed=sum(kv_lens) % 10007)
+    width = max(kv_lens)
+    native = case.native_logits()[:, :width].clone()
+    assert_logits_equal(case, native, case.indexer_topk_logits()[:, :width])
+    top_k = make_top_k()
+    native_sel = case.native_region(top_k, case.new_out()).clone()
+    sel = case.indexer_topk_region(top_k, case.new_out())
+    assert_selection_contract(case, native, sel)
+    assert torch.equal(torch.sort(native_sel, dim=1).values, torch.sort(sel, dim=1).values)
+
+
+def test_served_bound_vs_reference():
+    case = _bound_case([SERVED_MAX_SEQ_LEN], seed=17)
+    cand = case.indexer_topk_logits()[:, :SERVED_MAX_SEQ_LEN]
+    ref = reference_logits(case)
+    for r, n in enumerate(case.row_lens):
+        torch.testing.assert_close(cand[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("launch", ["indexer_topk", "native"])
+@pytest.mark.parametrize(
+    "cols", [V2_BLOCK_TABLE_COLS, V1_BLOCK_TABLE_COLS], ids=["v2_cols", "v1_cols"]
+)
+def test_served_bound_single_launch(launch, cols):
+    """One logits launch at the bound, checked against the fp32 oracle.
+
+    Each parameter runs exactly one kind of kernel launch so the test can run
+    under compute-sanitizer memcheck (with PYTORCH_NO_CUDA_MEMORY_CACHING=1)
+    one parameter at a time. With the odd V1 width the last 128-token tile's
+    second page ID is one entry past the row; the policy declines that width.
+    """
+    case = Case(
+        [SERVED_MAX_SEQ_LEN], seed=23, max_seq_len=SERVED_MAX_SEQ_LEN, block_table_cols=cols
+    )
+    logits = case.indexer_topk_logits() if launch == "indexer_topk" else case.native_logits()
+    torch.cuda.synchronize()
+    ref = reference_logits(case)
+    for r, n in enumerate(case.row_lens):
+        torch.testing.assert_close(logits[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
+
+
+SERVED_GRAPH_STEPS = [
+    (1 << 20) - 600,
+    1 << 20,
+    (1 << 20) + 1,
+    SERVED_MAX_SEQ_LEN,
+    491061,
+    (1 << 20) + 10,
+]
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+def test_cuda_graph_replay_across_2pow20(batch):
+    """Shipped kernel + native top-k, captured at the served max_seq_len."""
+    _run_graph_replay(
+        batch,
+        None,
+        "native",
+        SERVED_MAX_SEQ_LEN,
+        SERVED_GRAPH_STEPS,
+        block_table_cols=V2_BLOCK_TABLE_COLS,
+    )
+
+
+# --- adapter hook ---------------------------------------------------------------
+
+
+def _hook_objects(case, **overrides):
+    top_k = make_top_k()
+    indexer = SimpleNamespace(
+        use_cute_dsl_paged_mqa_logits=True,
+        use_fp4=False,
+        use_gvr_emission=False,
+        top_k=top_k,
+        n_heads=HEADS,
+        head_dim=HEAD_DIM,
+        layer_idx=0,
+    )
+    metadata = SimpleNamespace(
+        num_contexts=0,
+        num_generations=case.B,
+        seq_lens=torch.full((case.B,), NEXT_N, dtype=torch.int32),
+        is_ragged_verify=False,
+        in_mtp_draft_loop=False,
+        gen_indexer_kv_lens_cuda_runtime=case.kv_lens,
+        kv_lens_cuda_runtime=case.kv_lens,
+        use_expanded_buffers_for_mtp=False,
+        mapping=SimpleNamespace(enable_attention_dp=False),
+        kv_cache_manager=SimpleNamespace(get_indexer_k_cache_buffers=lambda idx: case.pool),
+        indexer_k_cache_block_offsets=case.block_table,
+        scheduler_metadata_buffer=case.schedule,
+        get_indexer_max_seq_len=lambda: case.max_seq_len,
+        radix_aux_indices=case.radix_idx,
+        radix_aux_logits=case.radix_val,
+    )
+    for key, value in overrides.items():
+        target = indexer if hasattr(indexer, key) else metadata
+        setattr(target, key, value)
+    return indexer, metadata
+
+
+HOOK_LENS = [50000, 4224, 64, 131072]
+
+
+@pytest.mark.parametrize("batch", [1, 2, 3, 4])
+def test_adapter_hook_engages_and_matches_native(batch):
+    case = Case(HOOK_LENS[:batch], seed=41 + batch)
+    indexer, metadata = _hook_objects(case)
+    rows = batch * NEXT_N
+    q_tok = case.q.reshape(rows, HEADS, HEAD_DIM)
+    out = case.new_out()
+    assert indexer_topk.try_indexer_decode(indexer, metadata, q_tok, case.weights, 0, rows, out)
+    native = case.native_region(make_top_k(), case.new_out())
+    assert_selection_contract(case, case.native_logits()[:, : max(HOOK_LENS[:batch])], out)
+    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(out, dim=1).values)
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+def test_adapter_hook_engages_at_served_bound(batch):
+    """Served GEN values: indexer max_seq_len 2**20 + 14, V2 table width."""
+    lens = [SERVED_MAX_SEQ_LEN, (1 << 20) + 1, 1 << 20, (1 << 20) - 1][:batch]
+    case = _bound_case(lens, seed=61 + batch)
+    indexer, metadata = _hook_objects(case)
+    rows = batch * NEXT_N
+    q_tok = case.q.reshape(rows, HEADS, HEAD_DIM)
+    out = case.new_out()
+    engaged_before = indexer_topk.DECISION_COUNTS.get("engaged", 0)
+    assert indexer_topk.try_indexer_decode(indexer, metadata, q_tok, case.weights, 0, rows, out)
+    assert indexer_topk.DECISION_COUNTS["engaged"] == engaged_before + 1
+    logits = case.native_logits()[:, : max(lens)].clone()
+    native = case.native_region(make_top_k(), case.new_out())
+    assert_selection_contract(case, logits, out)
+    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(out, dim=1).values)
+
+
+# Each case: indexer/metadata overrides plus the call shape. Every one must be
+# declined by the policy before any device work.
+_I32 = torch.int32
+FALLBACK_CASES = {
+    "fp4_keys": dict(overrides={"use_fp4": True}),
+    "non_dsl_native_logits": dict(overrides={"use_cute_dsl_paged_mqa_logits": False}),
+    "gvr_emission": dict(overrides={"use_gvr_emission": True}),
+    "ragged_verify": dict(overrides={"is_ragged_verify": True}),
+    "expanded_mtp_buffers": dict(overrides={"use_expanded_buffers_for_mtp": True}),
+    "mtp_draft_step": dict(overrides={"in_mtp_draft_loop": True}),
+    "attention_dp": dict(overrides={"mapping": SimpleNamespace(enable_attention_dp=True)}),
+    "radix_topk_engine": dict(radix_topk=True),
+    "sm100": dict(sm=100),
+    "page16": dict(page=16),
+    "page32": dict(page=32),
+    "page128": dict(page=128),
+    "non_uniform_seq_lens": dict(overrides={"seq_lens": torch.tensor([4], dtype=_I32)}),
+    "next_n4": dict(overrides={"seq_lens": torch.tensor([4], dtype=_I32)}, next_n=4),
+    # Constructor-invalid for the rows kernel (8 cached weights > 4 heads).
+    "heads4_next_n4": dict(
+        overrides={"n_heads": 4, "seq_lens": torch.tensor([4], dtype=_I32)}, heads=4, next_n=4
+    ),
+    "batch5": dict(
+        overrides={"num_generations": 5, "seq_lens": torch.full((5,), NEXT_N, dtype=_I32)},
+        batch=5,
+    ),
+    "mixed_context_batch": dict(
+        overrides={"num_contexts": 1, "seq_lens": torch.tensor([100, NEXT_N], dtype=_I32)}
+    ),
+    "max_seq_len_2m": dict(overrides={"get_indexer_max_seq_len": lambda: 1 << 21}),
+    "max_seq_len_above_served_bound": dict(
+        overrides={"get_indexer_max_seq_len": lambda: SERVED_MAX_SEQ_LEN + 1}
+    ),
+    # KV cache manager V1 at the served max_seq_len: odd block-table width.
+    "block_table_v1_odd_width": dict(
+        overrides={"get_indexer_max_seq_len": lambda: SERVED_MAX_SEQ_LEN},
+        block_table_cols=V1_BLOCK_TABLE_COLS,
+    ),
+    "bf16_weights": dict(weights_dtype=torch.bfloat16),
+    "non_contiguous_query": dict(non_contiguous_q=True),
+}
+
+
+@pytest.mark.parametrize("name", sorted(FALLBACK_CASES))
+def test_adapter_hook_falls_back_before_side_effects(name):
+    spec = FALLBACK_CASES[name]
+    case = Case([50000], seed=43)
+    overrides = dict(spec.get("overrides", {}))
+    if "page" in spec:
+        view = case.pool.view(-1, spec["page"], 1, HEAD_DIM + 4)
+        overrides["kv_cache_manager"] = SimpleNamespace(
+            get_indexer_k_cache_buffers=lambda idx: view
+        )
+    if spec.get("radix_topk"):
+        overrides["top_k"] = TopK(TOP_K, decode_implementation=TopKImplementation.CUDA_RADIX)
+    if "block_table_cols" in spec:
+        overrides["indexer_k_cache_block_offsets"] = torch.zeros(
+            1, spec["block_table_cols"], dtype=torch.int32, device=_DEV
+        )
+    indexer, metadata = _hook_objects(case, **overrides)
+    heads = spec.get("heads", HEADS)
+    rows = spec.get("next_n", NEXT_N) * spec.get("batch", 1)
+    q_wide = torch.zeros(rows, heads, 2 * HEAD_DIM, device=_DEV).to(torch.float8_e4m3fn)
+    q_tok = q_wide[..., :HEAD_DIM]
+    if not spec.get("non_contiguous_q"):
+        q_tok = q_tok.contiguous()
+    weights = torch.zeros(rows, heads, device=_DEV, dtype=spec.get("weights_dtype", torch.float32))
+    out = torch.full((rows, TOP_K), -7, dtype=torch.int32, device=_DEV)
+    torch.cuda.synchronize()
+    decisions_before = sum(indexer_topk.DECISION_COUNTS.values())
+    engaged_before = indexer_topk.DECISION_COUNTS.get("engaged", 0)
+    with patch.object(
+        indexer_topk, "get_sm_version", return_value=spec.get("sm", get_sm_version())
+    ):
+        with patch.object(indexer_topk, "decode_topk") as launched:
+            assert not indexer_topk.try_indexer_decode(
+                indexer, metadata, q_tok, weights, 0, rows, out
+            )
+    torch.cuda.synchronize()
+    assert not launched.called, "indexer_topk must decline before building or launching a kernel"
+    assert bool((out == -7).all()), "fallback must leave the output untouched"
+    assert sum(indexer_topk.DECISION_COUNTS.values()) == decisions_before + 1
+    assert indexer_topk.DECISION_COUNTS.get("engaged", 0) == engaged_before
+
+
+def test_adapter_hook_temporal_gvr_prior_falls_back():
+    case = Case([50000], seed=45)
+    indexer, metadata = _hook_objects(case)
+    indexer.top_k = TopK(
+        TOP_K, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=False
+    )
+    out = case.new_out(fill=-7)
+    q_tok = case.q.reshape(NEXT_N, HEADS, HEAD_DIM)
+    assert not indexer_topk.try_indexer_decode(
+        indexer, metadata, q_tok, case.weights, 0, NEXT_N, out
+    )
+    assert bool((out == -7).all())
+
+
+# --- the target's indexer against the built-in Indexer ------------------------------
+
+
+def _make_indexer(module):
+    from tensorrt_llm._torch.attention.backends.interface import (
+        PositionalEmbeddingParams,
+        RopeParams,
+    )
+    from tensorrt_llm.functional import PositionEmbeddingType
+    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
+
+    sparse_params = DeepSeekSparseAttentionConfig(
+        index_head_dim=HEAD_DIM,
+        index_n_heads=HEADS,
+        index_topk=TOP_K,
+        use_cute_dsl_topk=True,
+        use_cute_dsl_paged_mqa_logits=True,
+        enable_heuristic_topk=True,
+        index_share_for_mtp_iteration=True,
+    ).to_sparse_params()
+    pos = PositionalEmbeddingParams(
+        type=PositionEmbeddingType.rope_gpt_neox,
+        rope=RopeParams(dim=64, theta=10000.0, max_positions=4096),
+        is_neox=True,
+    )
+    mla = SimpleNamespace(hidden_size=6144, q_lora_rank=2048, qk_rope_head_dim=64)
+    with patch.object(module, "RotaryEmbedding", return_value=Mock()):
+        return module.Indexer(
+            quant_config=None,
+            pos_embd_params=pos,
+            mla_params=mla,
+            skip_create_weights_in_init=True,
+            sparse_params=sparse_params,
+            dtype=torch.bfloat16,
+            layer_idx=0,
+        )
+
+
+def _decode_metadata(case, in_mtp_draft_loop, num_accepted=None):
+    rows = case.B * NEXT_N
+    return SimpleNamespace(
+        kv_cache_manager=SimpleNamespace(
+            quant_block_size=128,
+            get_indexer_k_cache_buffers=lambda idx: case.pool,
+            layer_offsets={0: 0},
+        ),
+        num_contexts=0,
+        num_generations=case.B,
+        num_ctx_tokens=0,
+        num_tokens=rows,
+        get_empty=lambda buffers, shape, cache_name, dtype, capture_graph: torch.empty(
+            shape, dtype=dtype, device=_DEV
+        ),
+        cuda_graph_buffers=None,
+        is_cuda_graph=False,
+        skip_indexer_for_ctx_reqs=False,
+        skip_indexer_for_gen_reqs=False,
+        gvr_prior_indices=None,
+        in_mtp_draft_loop=in_mtp_draft_loop,
+        indexer_skip_topk=False,
+        shared_topk_indices=None,
+        mtp_num_accepted=num_accepted,
+        seq_lens=torch.full((case.B,), NEXT_N, dtype=torch.int32),
+        is_ragged_verify=False,
+        use_expanded_buffers_for_mtp=False,
+        kv_lens_cuda_2d=case.kv_lens.view(-1, 1).expand(-1, NEXT_N).contiguous(),
+        indexer_k_cache_block_offsets=case.block_table,
+        scheduler_metadata_buffer=case.schedule,
+        scheduler_metadata_buffer_full_next_n=case.schedule,
+        gen_indexer_kv_lens_cuda_runtime=case.kv_lens,
+        get_indexer_max_seq_len=lambda: case.max_seq_len,
+        dsl_expand_factor=case.factor,
+        dsl_atom=case.atom,
+        kv_lens_expanded_cuda=case.kv_lens_exp,
+        block_table_expanded=case.block_table_exp,
+        scheduler_metadata_buffer_expanded=case.schedule_exp,
+        kv_lens_row_reorder=None,
+        kv_lens_cuda_runtime=case.kv_lens,
+        ragged_row_kv_lens=lambda n: None,
+        radix_aux_indices=case.radix_idx,
+        radix_aux_logits=case.radix_val,
+        _tokens_per_block=BLOCK_KV,
+    )
+
+
+@pytest.mark.parametrize("in_mtp_draft_loop", [False, True], ids=["target", "draft_loop"])
+def test_adopted_indexer_preserves_native_indexer_semantics(monkeypatch, in_mtp_draft_loop):
+    """Stock projections -> the target's ``GlmIndexer.sparse_attn_indexer`` in a decode step
+    (indexer_topk on) == the built-in indexer.
+
+    Covers the target forward (indexer_topk engaged) and an MTP draft step (indexer_topk declines;
+    the native path computes the selection and stashes each request's last
+    accepted row for the reuse steps).
+    """
+    case = Case([60000, 45000], seed=51)
+    rows = case.B * NEXT_N
+    q_tok = case.q.reshape(rows, HEADS, HEAD_DIM)
+    k_dummy = torch.zeros(rows, HEAD_DIM, device=_DEV).to(torch.float8_e4m3fn)
+    k_scale_dummy = torch.ones(rows, 1, device=_DEV)
+    hidden = torch.zeros(rows, 6144, device=_DEV, dtype=torch.bfloat16)
+    accepted = torch.tensor([3, 6], dtype=torch.int32, device=_DEV)
+
+    native_indexer = _make_indexer(native_indexer_mod)
+    meta_native = _decode_metadata(case, in_mtp_draft_loop, accepted)
+    native = native_indexer.sparse_attn_indexer(
+        meta_native, hidden, q_tok, k_dummy, k_scale_dummy, case.weights, is_generation=True
+    ).clone()
+
+    with monkeypatch.context() as mp:
+        plan_testlib.select(dispatch, "indexer_topk", mp)
+        indexer_topk_indexer = GlmIndexer.adopt(_make_indexer(native_indexer_mod))
+        meta_indexer_topk = _decode_metadata(case, in_mtp_draft_loop, accepted)
+        with patch.object(
+            indexer_topk, "decode_topk", wraps=indexer_topk.decode_topk
+        ) as indexer_topk_launch:
+            with dispatch.step(dispatch.Phase.DECODE):
+                got = indexer_topk_indexer.sparse_attn_indexer(
+                    meta_indexer_topk,
+                    hidden,
+                    q_tok,
+                    k_dummy,
+                    k_scale_dummy,
+                    case.weights,
+                    is_generation=True,
+                )
+
+    # indexer_topk engages for the target forward and declines in the MTP draft step.
+    assert indexer_topk_launch.call_count == (0 if in_mtp_draft_loop else 1)
+    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(got, dim=1).values)
+    if in_mtp_draft_loop:
+        assert meta_native.shared_topk_indices is not None
+        # The native engine's output order within a row varies run to run.
+        assert torch.equal(
+            torch.sort(meta_native.shared_topk_indices, dim=1).values,
+            torch.sort(meta_indexer_topk.shared_topk_indices, dim=1).values,
+        )
+    else:
+        assert meta_indexer_topk.shared_topk_indices is None
+
+
+def test_without_a_plan_indexer_topk_is_off(monkeypatch):
+    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
+    assert not dispatch.is_enabled("indexer_topk")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py
new file mode 100644
index 000000000000..c5300f6c1a1f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py
@@ -0,0 +1,1188 @@
+# 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.
+"""indexer_topk inside the real DSA decode pipeline.
+
+Every object on the path to and from indexer_topk is the native one:
+
+- ``DSACacheManagerV2`` holding three GLM-5.2-shaped MLA layers: layer 0 owns
+  a full indexer, layer 1 is a shared-indexer layer, layer 2 is the MTP layer
+  (full indexer). FP8 MLA cache, FP8 indexer K, 64-token pages and the served
+  ``max_seq_len`` (2**20 + 14).
+- ``MLA`` modules with the per-rank GLM-5.2 TP8 attention shapes (8 heads,
+  q-nope 192, q-rope 64, v 256, kv-LoRA 512, q-LoRA 2048; indexer 32 x 128,
+  top-k 2048), adopted as the target's ``GlmMLA`` (with its DSA backend and
+  indexer), running the target's ``_forward_dsa_attn`` as a decode step:
+  indexer K append, the indexer
+  (native or indexer_topk), the backend's shared-layer copy, the local-to-global remap,
+  RoPE/FP8-query/KV append and sparse MLA (native, or sparse_mla_decode/rope_kv_append when
+  selected).
+- ``DSAtrtllmAttentionMetadata`` with ``update_spec_dec_param(max_draft_len=5)``
+  and ``prepare()`` per decode step, followed by the decode runner's
+  ``_preprocess_inputs`` sequence: ``on_update_kv_lens()``, the overlap
+  scheduler's accepted-token correction of ``kv_lens_cuda``, and
+  ``on_update_kv_lens()`` again.
+- The one-model MTP-Eagle draft loop of the target's worker,
+  ``GlmMtpWorker._forward_linear_draft_loop``, itself (index share on, accepted-row
+  stash, per-step metadata updates), with a stub draft model whose MTP layer runs the
+  real layer-2 attention. With
+  ``draft_loop`` selected the loop takes draft_loop's per-step metadata update.
+
+The context is written by the native producers (indexer K quantize + scatter,
+MLA RoPE-append of random latents). A decode sequence is four target verify
+steps (next_n = 6) separated by three draft loops with accepted counts
+``ACCEPTED``. Every arm replays the same sequence on the same caches; arms
+differ only in whether indexer_topk runs. ``GLM_MOE_DSA_TEST_FAST_PATHS``
+(``plan_testlib.FAST_PATHS_ENV``) is the base set of fast paths (for example
+``sparse_mla_decode,rope_kv_append`` or ``draft_loop``), whose plan is installed before any layer
+is built; the tests add ``indexer_topk`` on top of it for the indexer_topk arms.
+
+``test_rank_agreement`` runs the same sequence on every rank of an MPI job
+(8 ranks over 2 hosts) and compares digests across ranks.
+``test_indexer_topk_region_timing`` (``GLM_MOE_DSA_INDEXER_TOPK_PIPELINE_TIMING=1``) times the attention region of a
+full + shared layer pair in CUDA graphs with indexer_topk off and on.
+"""
+
+import hashlib
+import json
+import os
+import socket
+import sys
+import types
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Dict, List
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+from tensorrt_llm._utils import get_sm_version
+
+if not torch.cuda.is_available() or not IS_CUTLASS_DSL_AVAILABLE:
+    pytest.skip("CUDA and the CUTLASS DSL are required", allow_module_level=True)
+if get_sm_version() != 103:
+    pytest.skip("indexer_topk is verified on GB300 (sm103) only", allow_module_level=True)
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+    GlmMLA,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention_forward import (  # noqa: E501
+    _forward_dsa_attn,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+    indexer_decode as indexer_topk,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
+    GlmMtpWorker,
+)
+from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams
+from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManagerV2
+from tensorrt_llm._torch.attention.backends.sparse.dsa import indexer as indexer_mod
+from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+from tensorrt_llm._torch.attention.mla import MLA
+from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
+from tensorrt_llm.bindings import DataType
+from tensorrt_llm.bindings.internal.batch_manager import CacheType
+from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata
+from tensorrt_llm.functional import PositionEmbeddingType
+from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+# DSACacheManagerV2 owns background threads.
+pytestmark = pytest.mark.threadleak(enabled=False)
+
+LOCAL_HEADS = 8
+QK_NOPE = 192
+QK_ROPE = 64
+V_HEAD = 256
+KV_LORA = 512
+Q_LORA = 2048
+HIDDEN = 6144
+INDEX_HEADS = 32
+INDEX_DIM = 128
+TOP_K = 2048
+DRAFT_LEN = 5
+NEXT_N = DRAFT_LEN + 1
+TOKENS_PER_BLOCK = 64
+SCHEDULE_BLOCK_KV = indexer_mod._DG_SCHEDULE_BLOCK_KV
+# GLM-5.2: max_position_embeddings 2**20, rope theta 8e6, no rope scaling.
+MAX_POSITIONS = 1 << 20
+ROPE_THETA = 8.0e6
+# Served GEN: max_seq_len 2**20 plus the K = 5 one-model MTP slack.
+SERVED_MAX_SEQ_LEN = (1 << 20) + 14
+NUM_EXTRA_KV_TOKENS = DRAFT_LEN - 1
+LAYER_FULL, LAYER_SHARED, LAYER_MTP = 0, 1, 2
+DRAFT_VOCAB = 64
+# Accepted tokens (bonus included, 1..K+1) of the three draft loops.
+ACCEPTED = (3, 1, 6)
+_DEV = torch.device("cuda")
+BASE_FEATURES = os.environ.get(plan_testlib.FAST_PATHS_ENV, "")
+
+
+def _base_feature_set():
+    return set(plan_testlib.names(BASE_FEATURES))
+
+
+@pytest.fixture(scope="module", autouse=True)
+def base_plan():
+    """The plan of the base fast paths, installed before any layer is adopted."""
+    saved = (dispatch._INSTALLED[0], dispatch._VALIDATE[0], dispatch._require[0])
+    plan_testlib.select(dispatch, BASE_FEATURES)
+    yield
+    dispatch._INSTALLED[0], dispatch._VALIDATE[0], dispatch._require[0] = saved
+
+
+@pytest.fixture(autouse=True)
+def decode_step():
+    """Every test runs the target's decode steps."""
+    with dispatch.step(dispatch.Phase.DECODE):
+        yield
+
+
+@dataclass
+class _RopeConfig:
+    hidden_size: int
+    num_attention_heads: int
+    rope_scaling: object
+    max_position_embeddings: int
+    rope_theta: float
+    qk_rope_head_dim: int
+    model_type: str
+
+
+def _pos_embd_params():
+    rope_config = _RopeConfig(
+        hidden_size=HIDDEN,
+        num_attention_heads=LOCAL_HEADS,
+        rope_scaling=None,
+        max_position_embeddings=MAX_POSITIONS,
+        rope_theta=ROPE_THETA,
+        qk_rope_head_dim=QK_ROPE,
+        model_type="glm_moe_dsa",
+    )
+    # Same construction as the DeepSeek-V3-family attention layers.
+    return PositionalEmbeddingParams(
+        type=PositionEmbeddingType.yarn, rope=RopeParams.from_config(rope_config), is_neox=False
+    )
+
+
+def _expected_target_kv_lens(context_len: int) -> List[int]:
+    lens = [context_len]
+    for accepted in ACCEPTED:
+        lens.append(lens[-1] + accepted)
+    return lens
+
+
+@dataclass
+class StepInputs:
+    """Token-wise MLA/indexer inputs of one layer call."""
+
+    q: torch.Tensor
+    compressed_kv: torch.Tensor
+    k_pe: torch.Tensor
+    hidden: torch.Tensor
+    qr: torch.Tensor
+    position_ids: torch.Tensor
+
+
+class Recorder:
+    """Ordered record of one arm: indexer calls, shared rows and layer outputs."""
+
+    def __init__(self, name: str, canonical: bool):
+        self.name = name
+        self.canonical = canonical
+        self.events: List[dict] = []
+        self.capability: List[tuple] = []
+        self.sparse_mla_calls: List[tuple] = []
+        self.tag = None
+
+    def add(self, kind: str, **fields):
+        fields.update(kind=kind, tag=self.tag)
+        self.events.append(fields)
+
+
+def _canonicalize_(rows: torch.Tensor):
+    """Sort each top-k row in place (ascending local token ids)."""
+    assert bool((rows >= 0).all()), "tested lengths fill every top-k slot"
+    rows.copy_(torch.sort(rows, dim=1).values)
+
+
+class IndexerTap:
+    """Wraps ``Indexer.sparse_attn_indexer`` to record each call of an arm.
+
+    With ``canonical`` set, the rows every consumer reads (the returned rows,
+    and on a draft-loop reuse step the stashed row the reuse copies) are
+    sorted in place, so the attention input order does not depend on the
+    native top-k engine's in-row order, which varies run to run.
+    """
+
+    def __init__(self, recorder: Recorder):
+        self.rec = recorder
+
+    def wrap(self, original):
+        tap = self
+
+        def sparse_attn_indexer(indexer, metadata, *args, **kwargs):
+            in_draft = bool(metadata.in_mtp_draft_loop)
+            reuse = in_draft and bool(metadata.indexer_skip_topk)
+            pre = {}
+            if not in_draft:
+                # The per-request schedule and lengths indexer_topk consumes, as they are
+                # when the indexer runs.
+                pre["sched"] = metadata.scheduler_metadata_buffer.clone()
+                pre["kv"] = metadata.gen_indexer_kv_lens_cuda_runtime.clone()
+                pre["kv_lens_cuda"] = metadata.kv_lens_cuda[: metadata.num_seqs].clone()
+                pre["num_sms"] = metadata.num_sms
+                pre["block_table_cols"] = metadata.indexer_k_cache_block_offsets.shape[1]
+            if reuse:
+                stash = metadata.shared_topk_indices[: metadata.num_generations]
+                pre["stash"] = stash.clone()
+                if tap.rec.canonical:
+                    _canonicalize_(stash)
+            engaged = indexer_topk.DECISION_COUNTS.get("engaged", 0)
+            decided = sum(indexer_topk.DECISION_COUNTS.values())
+            out = original(indexer, metadata, *args, **kwargs)
+            rows = out[: metadata.num_tokens - metadata.num_ctx_tokens]
+            if tap.rec.canonical and not reuse:
+                _canonicalize_(rows)
+            tap.rec.add(
+                "indexer",
+                layer=indexer.layer_idx,
+                in_draft=in_draft,
+                reuse=reuse,
+                # On a reuse step rope_kv_append may defer the copy into its bridge, so the
+                # returned buffer is only meaningful when the indexer selected.
+                rows=None if reuse else torch.sort(rows, dim=1).values.clone(),
+                indexer_topk_engaged=indexer_topk.DECISION_COUNTS.get("engaged", 0) - engaged,
+                indexer_topk_decisions=sum(indexer_topk.DECISION_COUNTS.values()) - decided,
+                **pre,
+            )
+            return out
+
+        return sparse_attn_indexer
+
+
+class _StubMtpLayer:
+    """MTP layer of the stub draft model: runs the real layer-2 attention."""
+
+    def __init__(self, pipe: "GlmDsaPipeline", rec: Recorder, loop_idx: int):
+        self.pipe = pipe
+        self.rec = rec
+        self.loop_idx = loop_idx
+        self.step = 0
+
+    def __call__(
+        self,
+        *,
+        embed_tokens,
+        all_rank_num_tokens,
+        input_ids,
+        position_ids,
+        hidden_states,
+        attn_metadata,
+        spec_metadata,
+        **kwargs,
+    ):
+        step = self.step
+        self.step += 1
+        self.rec.tag = ("draft", self.loop_idx, step)
+        n = attn_metadata.num_tokens
+        positions = position_ids.reshape(-1)[:n].to(torch.int32).contiguous()
+        seed = 1000 + 37 * self.loop_idx + step
+        inputs = self.pipe.make_inputs(seed=seed, n=n, positions=positions)
+        out = self.pipe.layer_forward(LAYER_MTP, attn_metadata, inputs)
+        self.rec.add("output", layer=LAYER_MTP, out=out.clone())
+        # Hidden states that depend on the attention output, so the drafted
+        # tokens do as well.
+        hidden = torch.zeros(n, HIDDEN, dtype=torch.bfloat16, device=_DEV)
+        hidden[:, : out.shape[1]] = out
+        return hidden
+
+    @staticmethod
+    def shared_head(hidden_states, lm_head, attn_metadata, return_context_logits):
+        return hidden_states[:, :DRAFT_VOCAB].float()
+
+
+def _sample_draft_tokens(logits, spec_metadata, batch_size, draft_step=0, mapping_lm_head_tp=None):
+    return torch.argmax(logits, dim=-1).to(torch.int32)
+
+
+def _make_worker():
+    """The attributes of an MTP-Eagle one-model worker the draft loop reads."""
+    worker = SimpleNamespace(
+        is_mtp_eagle=True,
+        guided_decoder=None,
+        sa_enhancer=None,
+        model_config=None,
+        spec_config=SimpleNamespace(use_relaxed_acceptance_for_thinking=False),
+        _auxiliary_state_handlers=None,
+        _uses_external_shared_target_kv=False,
+        use_dynamic_tree=False,
+        _draft_loop_ops=None,
+    )
+    worker.draft_kv_cache_context = types.MethodType(SpecWorkerBase.draft_kv_cache_context, worker)
+    for name in (
+        "_run_draft_forward",
+        "_get_step_all_rank_num_tokens",
+        "_prepare_flash_mla_generation_layout",
+    ):
+        setattr(worker, name, types.MethodType(getattr(GlmMtpWorker, name), worker))
+    worker.sample_draft_tokens = _sample_draft_tokens
+    worker._can_use_rejection_sampling = lambda spec_metadata: False
+    return worker
+
+
+def _make_spec_metadata():
+    return SimpleNamespace(
+        runtime_draft_len=DRAFT_LEN,
+        batch_indices_cuda=torch.zeros(1, dtype=torch.int32, device=_DEV),
+        all_rank_num_tokens=None,
+        subseq_all_rank_num_tokens=None,
+        wants_advanced_draft_sampling=False,
+        max_num_requests=1,
+        draft_tokens=torch.zeros(DRAFT_LEN, dtype=torch.int32, device=_DEV),
+        enable_penalty=False,
+    )
+
+
+class GlmDsaPipeline:
+    """Three GLM-5.2-shaped DSA MLA layers over one ``DSACacheManagerV2``."""
+
+    def __init__(self, context_len: int, seed: int = 0):
+        # ``context_len`` is the KV length of the first target verify step
+        # (its six verify tokens included).
+        self.context_len = context_len
+        self.prompt_len = context_len - NEXT_N
+        self.capacity = context_len + sum(ACCEPTED) + 2 * NEXT_N + NUM_EXTRA_KV_TOKENS
+        assert self.capacity + NUM_EXTRA_KV_TOKENS <= SERVED_MAX_SEQ_LEN
+        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        self.pretrained = SimpleNamespace(
+            rms_norm_eps=1e-5,
+            num_hidden_layers=2,
+            index_topk_pattern=["F", "S"],
+            qk_rope_head_dim=QK_ROPE,
+        )
+        self.sparse_config = DeepSeekSparseAttentionConfig(
+            index_n_heads=INDEX_HEADS,
+            index_head_dim=INDEX_DIM,
+            index_topk=TOP_K,
+            use_cute_dsl_topk=True,
+            use_cute_dsl_paged_mqa_logits=True,
+            enable_heuristic_topk=True,
+            index_share_for_mtp_iteration=True,
+        )
+        quant_config = QuantConfig()
+        quant_config.kv_cache_quant_algo = QuantAlgo.FP8
+        self.model_config = ModelConfig(
+            mapping=self.mapping,
+            sparse_attention_config=self.sparse_config,
+            pretrained_config=self.pretrained,
+            quant_config=quant_config,
+        )
+        self.model_config.extra_attrs["kv_cache_dtype"] = "fp8"
+        pos = _pos_embd_params()
+        self.layers = []
+        for layer_idx in range(3):
+            mla = MLA(
+                hidden_size=HIDDEN,
+                num_attention_heads=LOCAL_HEADS,
+                num_key_value_heads=1,
+                qk_nope_head_dim=QK_NOPE,
+                qk_rope_head_dim=QK_ROPE,
+                v_head_dim=V_HEAD,
+                q_lora_rank=Q_LORA,
+                kv_lora_rank=KV_LORA,
+                predicted_tokens_per_seq=NEXT_N,
+                max_position_embeddings=MAX_POSITIONS,
+                bias=False,
+                pos_embd_params=pos,
+                layer_idx=layer_idx,
+                dtype=torch.bfloat16,
+                config=self.model_config,
+            )
+            # The model's DeepseekV32Attention adds only its constructor to MLA; the pipeline
+            # builds the MLA at GLM shapes directly and adopts it as the model's attention. The
+            # model builds the FMHA libraries in ``create_weights`` after adoption; a module built
+            # alone rebuilds them for the library adoption declared.
+            mla.__class__ = DeepseekV32Attention
+            GlmMLA.adopt(mla)
+            for backend in (mla.mha, mla.mqa):
+                if backend is not None and backend.prepended_fmha_lib_classes:
+                    backend.update_quant_config(backend.quant_config)
+            mla = mla.to(_DEV)
+            if mla.mqa.indexer is not None:
+                mla.indexer = mla.mqa.indexer.to(_DEV)
+            self.layers.append(mla)
+        assert self.layers[LAYER_FULL].mqa.indexer is not None
+        assert self.layers[LAYER_SHARED].mqa.indexer is None, "layer 1 must be shared"
+        assert self.layers[LAYER_MTP].mqa.indexer is not None
+        self._init_weights(seed)
+
+        self.kv_cache_manager = DSACacheManagerV2(
+            KvCacheConfig(
+                dtype="fp8",
+                # At least max_seq_len, or the manager clamps max_seq_len.
+                max_tokens=SERVED_MAX_SEQ_LEN + 64 * TOKENS_PER_BLOCK,
+                enable_block_reuse=False,
+                host_cache_size=0,
+            ),
+            CacheType.SELFKONLY,
+            num_layers=3,
+            num_kv_heads=1,
+            head_dim=KV_LORA + QK_ROPE,
+            tokens_per_block=TOKENS_PER_BLOCK,
+            max_seq_len=SERVED_MAX_SEQ_LEN,
+            max_batch_size=1,
+            mapping=self.mapping,
+            dtype=DataType.FP8,
+            sparse_attention_config=self.sparse_config,
+            pretrained_config=self.pretrained,
+        )
+        assert self.kv_cache_manager.max_seq_len == SERVED_MAX_SEQ_LEN
+        assert self.kv_cache_manager.indexer_k_cache_local_layer_mask == [True, False, True]
+        added = self.kv_cache_manager.add_dummy_requests(
+            request_ids=[0], token_nums=[self.capacity], is_gen=False, prepare_resource=True
+        )
+        assert added is not None, "KV cache allocation failed"
+        sparse_params = self.sparse_config.to_sparse_params(pretrained_config=self.pretrained)
+        self.attn_cls = get_attention_backend("TRTLLM", sparse_params=sparse_params)
+        self.metadata_params = self.sparse_config.to_sparse_metadata_params(
+            pretrained_config=self.pretrained
+        )
+        assert self.metadata_params.has_shared_indexer_layers
+        assert self.metadata_params.mtp_index_share
+        self._populate_context(seed)
+        self.md = self._new_metadata(
+            num_contexts=0,
+            seq_len=NEXT_N,
+            cached=self.prompt_len,
+            max_num_tokens=64,
+            extra_kv=NUM_EXTRA_KV_TOKENS,
+        )
+        with torch.no_grad():
+            self.embed_tokens = torch.nn.Embedding(DRAFT_VOCAB, HIDDEN).to(_DEV, torch.bfloat16)
+
+    # ------------------------------------------------------------------ setup
+    def _init_weights(self, seed: int):
+        gen = torch.Generator(device=_DEV).manual_seed(seed)
+
+        def randn_like(t):
+            return (torch.randn(t.shape, device=_DEV, generator=gen) * 0.02).to(t.dtype)
+
+        with torch.no_grad():
+            for mla in self.layers:
+                mla.kv_b_proj.weight.copy_(randn_like(mla.kv_b_proj.weight))
+                kv_b = mla.kv_b_proj.weight.data.view(LOCAL_HEADS, QK_NOPE + V_HEAD, KV_LORA)
+                mla.v_b_proj.data = kv_b[:, QK_NOPE:, :].contiguous()
+                mla.k_b_proj_trans.data = kv_b[:, :QK_NOPE, :].transpose(1, 2).contiguous()
+                indexer = mla.mqa.indexer
+                if indexer is not None:
+                    for linear in (indexer.wq_b, indexer.wk, indexer.weights_proj):
+                        linear.weight.copy_(randn_like(linear.weight))
+                    indexer.post_load_weights()
+
+    def _new_metadata(self, *, num_contexts, seq_len, cached, max_num_tokens, extra_kv):
+        return self.attn_cls.Metadata(
+            seq_lens=torch.tensor([seq_len], dtype=torch.int),
+            request_ids=[0],
+            max_num_requests=1,
+            num_contexts=num_contexts,
+            prompt_lens=[self.prompt_len],
+            max_num_tokens=max_num_tokens,
+            kv_cache_manager=self.kv_cache_manager,
+            kv_cache_params=KVCacheParams(
+                use_cache=True, num_cached_tokens_per_seq=[cached], num_extra_kv_tokens=extra_kv
+            ),
+            mapping=self.mapping,
+            sparse_metadata_params=self.metadata_params,
+        )
+
+    @torch.inference_mode()
+    def _populate_context(self, seed: int):
+        """Write the prompt with the native producers (one prefill-shaped call)."""
+        n = self.prompt_len
+        md = self._new_metadata(num_contexts=1, seq_len=n, cached=0, max_num_tokens=n, extra_kv=0)
+        md.prepare()
+        gen = torch.Generator(device=_DEV).manual_seed(seed + 1)
+        for layer_idx in (LAYER_FULL, LAYER_MTP):
+            indexer = self.layers[layer_idx].mqa.indexer
+            k = torch.randn(n, INDEX_DIM, device=_DEV, generator=gen).to(torch.bfloat16)
+            rope = indexer.rope_dim
+            # The indexer's own K quantizer and K-cache scatter.
+            k_fp8, k_scale = indexer._prep_q_or_k(
+                k[:, :rope].contiguous(), k[:, rope:].contiguous()
+            )
+            indexer._update_k_cache(k_fp8, k_scale, md)
+            del k, k_fp8, k_scale
+        for mla in self.layers:
+            latent = torch.randn(n, KV_LORA + QK_ROPE, device=_DEV, generator=gen)
+            latent = latent.to(torch.bfloat16)
+            dummy_q = torch.zeros(
+                n, LOCAL_HEADS * (QK_NOPE + QK_ROPE), device=_DEV, dtype=torch.bfloat16
+            )
+            mla.mqa.mla_rope_append_paged_kv_assign_q(dummy_q, latent, md)
+            del latent, dummy_q
+        torch.cuda.synchronize()
+        del md
+        torch.cuda.empty_cache()
+
+    def shutdown(self):
+        self.kv_cache_manager.shutdown()
+
+    # ------------------------------------------------------------- execution
+    def make_inputs(self, seed: int, n: int, positions: torch.Tensor) -> StepInputs:
+        gen = torch.Generator(device=_DEV).manual_seed(seed)
+
+        def rnd(*shape):
+            return torch.randn(*shape, device=_DEV, generator=gen).to(torch.bfloat16)
+
+        return StepInputs(
+            q=rnd(n, LOCAL_HEADS * (QK_NOPE + QK_ROPE)),
+            compressed_kv=rnd(n, KV_LORA),
+            k_pe=rnd(n, QK_ROPE),
+            hidden=rnd(n, HIDDEN),
+            qr=rnd(n, Q_LORA),
+            position_ids=positions,
+        )
+
+    def project(self, layer_idx: int, inputs: StepInputs):
+        """Graph-capturable indexer projections of a full-indexer layer (Op 1)."""
+        indexer = self.layers[layer_idx].mqa.indexer
+        if indexer is None:
+            return []
+        return list(indexer.pre_indexer_proj(inputs.qr, inputs.hidden, inputs.position_ids))
+
+    def attend(self, layer_idx: int, md, inputs: StepInputs, intermediates) -> torch.Tensor:
+        """``_forward_dsa_attn`` of one layer (Op 2) on copies of the inputs."""
+        mla = self.layers[layer_idx]
+        n = md.num_tokens
+        latent = torch.cat([inputs.compressed_kv, inputs.k_pe], dim=-1)
+        output = torch.empty(n, LOCAL_HEADS * V_HEAD, dtype=torch.bfloat16, device=_DEV)
+        _forward_dsa_attn(
+            mla,
+            inputs.q.clone(),
+            inputs.compressed_kv.clone(),
+            inputs.k_pe.clone(),
+            latent,
+            intermediates,
+            inputs.position_ids,
+            md,
+            output,
+        )
+        return output
+
+    def layer_forward(self, layer_idx: int, md, inputs: StepInputs) -> torch.Tensor:
+        return self.attend(layer_idx, md, inputs, self.project(layer_idx, inputs))
+
+    def prepare_target_step(self, host_cached: int, kv_offset: int):
+        """Runner order: host metadata + ``prepare()``, then ``_preprocess_inputs``."""
+        md = self.md
+        md.seq_lens = torch.tensor([NEXT_N], dtype=torch.int)
+        md.num_contexts = 0
+        md.request_ids = [0]
+        md.prompt_lens = [self.prompt_len]
+        md.kv_cache_params = KVCacheParams(
+            use_cache=True,
+            num_cached_tokens_per_seq=[host_cached],
+            num_extra_kv_tokens=NUM_EXTRA_KV_TOKENS,
+        )
+        md.update_spec_dec_param(
+            batch_size=1,
+            is_spec_decoding_enabled=True,
+            is_spec_dec_tree=False,
+            is_spec_dec_dynamic_tree=False,
+            max_draft_len=DRAFT_LEN,
+            max_total_draft_tokens=DRAFT_LEN,
+            spec_metadata=None,
+            spec_tree_manager=None,
+            num_contexts=0,
+        )
+        md.prepare()
+        # DecoderRunner._preprocess_inputs (overlap scheduler + spec decode).
+        md.on_update_kv_lens()
+        md.kv_lens_cuda[0:1] += kv_offset
+        md.on_update_kv_lens()
+        return md
+
+    @torch.inference_mode()
+    def run_sequence(self, rec: Recorder, draft_loop: bool):
+        """Four target steps separated by three draft loops."""
+        md = self.md
+        kv = self.context_len
+        cached = kv - NEXT_N
+        host_cached = cached
+        kv_offset = 0
+        for step in range(len(ACCEPTED) + 1):
+            rec.tag = ("target", step)
+            self.prepare_target_step(host_cached, kv_offset)
+            positions = torch.arange(cached, cached + NEXT_N, dtype=torch.int32, device=_DEV)
+            rec.add("target_kv", kv=kv)
+            for layer_idx in (LAYER_FULL, LAYER_SHARED):
+                if layer_idx == LAYER_SHARED:
+                    rec.add(
+                        "shared_rows",
+                        rows=torch.sort(md.shared_topk_indices[:NEXT_N], dim=1).values.clone(),
+                    )
+                inputs = self.make_inputs(seed=10 * step + layer_idx, n=NEXT_N, positions=positions)
+                out = self.layer_forward(layer_idx, md, inputs)
+                rec.add("output", layer=layer_idx, out=out.clone())
+            if step == len(ACCEPTED):
+                break
+            accepted = ACCEPTED[step]
+            self.run_draft_loop(rec, step, accepted, positions, draft_loop)
+            rec.tag = ("after_loop", step)
+            rec.add("stale_sched", sched=md.scheduler_metadata_buffer.clone())
+            # Overlap scheduler: the host assumes every draft was accepted and
+            # the next forward corrects kv_lens_cuda by (accepted - next_n).
+            host_cached = cached + NEXT_N
+            kv_offset = accepted - NEXT_N
+            cached += accepted
+            kv += accepted
+        torch.cuda.synchronize()
+        return rec
+
+    def run_draft_loop(
+        self, rec: Recorder, loop_idx: int, accepted: int, positions, draft_loop: bool
+    ):
+        md = self.md
+        worker = _make_worker()
+        spec_md = _make_spec_metadata()
+        num_accepted = torch.tensor([accepted], dtype=torch.int32, device=_DEV)
+        draft_model = SimpleNamespace(
+            embed_tokens=self.embed_tokens,
+            lm_head=None,
+            mtp_layers=[_StubMtpLayer(self, rec, loop_idx)],
+        )
+        if draft_loop:
+            ops = dispatch.get_impl("draft_loop", "draft_loop_ops")(worker, md, spec_md)
+            assert ops is not None, "draft_loop declined the draft loop (see the capability log)"
+            # Acceptance is not run here (there are no target logits); give draft_loop
+            # the step-0 rows the native loop computes for one generation request.
+            ops.gather_ids = (num_accepted.to(torch.int64) - 1).contiguous()
+            worker._draft_loop_ops = ops
+        gen = torch.Generator(device=_DEV).manual_seed(7000 + loop_idx)
+        inputs = {
+            "input_ids": torch.randint(0, DRAFT_VOCAB, (NEXT_N,), device=_DEV, generator=gen),
+            "position_ids": positions.clone(),
+            "hidden_states": torch.randn(NEXT_N, HIDDEN, device=_DEV, generator=gen).to(
+                torch.bfloat16
+            ),
+            "attn_metadata": md,
+            "spec_metadata": spec_md,
+        }
+        # The worker's _forward_impl around the loop.
+        md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
+        try:
+            tokens = GlmMtpWorker._forward_linear_draft_loop(
+                worker,
+                inputs,
+                md,
+                spec_md,
+                draft_model,
+                None,
+                0,
+                1,
+                num_accepted,
+                None,
+            )
+        finally:
+            md.restore_from_spec_dec()
+            md.on_update()
+            md.use_spec_decoding = True
+        rec.tag = ("draft_tokens", loop_idx)
+        rec.add("draft_tokens", tokens=tokens.clone())
+
+
+# --------------------------------------------------------------------- arms
+def _install_recorders(mp, rec: Recorder):
+    original_log = dispatch.log_capability
+
+    def log_capability(feature_id, key, eligible, reason=""):
+        rec.capability.append((feature_id, key, bool(eligible), reason, rec.tag))
+        return original_log(feature_id, key, eligible, reason)
+
+    mp.setattr(dispatch, "log_capability", log_capability)
+    tap = IndexerTap(rec)
+    original_indexer = indexer_mod.Indexer.sparse_attn_indexer
+    mp.setattr(indexer_mod.Indexer, "sparse_attn_indexer", tap.wrap(original_indexer))
+    if "sparse_mla_decode" in _base_feature_set():
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+            sparse_mla as sparse_mla_sparse_mla,
+        )
+
+        sparse_mla_cls = sparse_mla_sparse_mla.SparseMlaDecodeFmha
+        original_forward = sparse_mla_cls.forward
+
+        def forward(self, *args, **kwargs):
+            rec.sparse_mla_calls.append(rec.tag)
+            return original_forward(self, *args, **kwargs)
+
+        mp.setattr(sparse_mla_cls, "forward", forward)
+
+
+def _select(mp, with_indexer_topk: bool):
+    features = _base_feature_set()
+    if with_indexer_topk:
+        features.add("indexer_topk")
+    else:
+        features.discard("indexer_topk")
+    plan_testlib.select(dispatch, features, mp)
+    assert dispatch.is_enabled("indexer_topk") == with_indexer_topk
+
+
+def run_arm(pipe: GlmDsaPipeline, monkeypatch, name: str, with_indexer_topk: bool, canonical: bool):
+    rec = Recorder(name, canonical)
+    with monkeypatch.context() as mp:
+        _select(mp, with_indexer_topk)
+        _install_recorders(mp, rec)
+        pipe.run_sequence(rec, draft_loop="draft_loop" in _base_feature_set())
+    return rec
+
+
+# -------------------------------------------------------------- comparisons
+def _fresh_schedule(kv: torch.Tensor, num_sms: int) -> torch.Tensor:
+    return get_paged_mqa_logits_metadata(kv.view(-1, 1), SCHEDULE_BLOCK_KV, num_sms)
+
+
+def _indexer_calls(rec: Recorder, phase: str):
+    return [e for e in rec.events if e["kind"] == "indexer" and e["tag"][0] == phase]
+
+
+def check_arm_invariants(rec: Recorder, pipe: GlmDsaPipeline, with_indexer_topk: bool):
+    """Per-arm checks that do not need another arm."""
+    expected_kv = _expected_target_kv_lens(pipe.context_len)
+    target_calls = _indexer_calls(rec, "target")
+    assert len(target_calls) == len(expected_kv), "one full-layer indexer call per target step"
+    for call, kv in zip(target_calls, expected_kv):
+        assert call["layer"] == LAYER_FULL and not call["in_draft"]
+        # The lengths and the per-request schedule indexer_topk reads are this step's,
+        # not a draft step's leftovers.
+        assert call["kv"].tolist() == [kv], (rec.name, call["tag"], call["kv"].tolist(), kv)
+        assert call["kv_lens_cuda"].tolist() == [kv]
+        fresh = _fresh_schedule(call["kv"], call["num_sms"])
+        assert torch.equal(call["sched"], fresh), (rec.name, call["tag"], "stale schedule")
+        assert call["block_table_cols"] % 2 == 0
+        assert call["indexer_topk_decisions"] == (1 if with_indexer_topk else 0)
+        assert call["indexer_topk_engaged"] == (1 if with_indexer_topk else 0), (
+            rec.name,
+            call["tag"],
+        )
+    draft_calls = _indexer_calls(rec, "draft")
+    assert len(draft_calls) == DRAFT_LEN * len(ACCEPTED)
+    for call in draft_calls:
+        assert call["layer"] == LAYER_MTP and call["in_draft"]
+        assert call["reuse"] == (call["tag"][2] > 0)
+        # indexer_topk is target-only: it declines every draft selection before device
+        # work, and the reuse steps never reach it.
+        assert call["indexer_topk_engaged"] == 0
+        assert call["indexer_topk_decisions"] == (
+            1 if with_indexer_topk and not call["reuse"] else 0
+        )
+    # The shared layer reads exactly the full layer's rows of the same step.
+    for step in range(len(expected_kv)):
+        full = next(e for e in target_calls if e["tag"] == ("target", step))
+        shared = next(
+            e for e in rec.events if e["kind"] == "shared_rows" and e["tag"] == ("target", step)
+        )
+        assert torch.equal(full["rows"], shared["rows"]), (rec.name, step)
+    features = _base_feature_set()
+    caps = rec.capability
+    num_attention_calls = 2 * len(expected_kv) + DRAFT_LEN * len(ACCEPTED)
+    # The consumers of indexer_topk's rows are the target layers; draft-step decisions
+    # of the neighbour packages are reported, not asserted.
+    if "sparse_mla_decode" in features:
+        declined = [c for c in caps if c[0] == "sparse_mla_decode" and not c[2]]
+        assert not [c for c in declined if c[4] and c[4][0] == "target"], declined
+        target_sparse_mla = [t for t in rec.sparse_mla_calls if t[0] == "target"]
+        # Both target layers (full and shared) of every step attend through sparse_mla_decode.
+        assert len(target_sparse_mla) == 2 * len(expected_kv), rec.sparse_mla_calls
+        if len(rec.sparse_mla_calls) != num_attention_calls or declined:
+            print(
+                f"[indexer_topk] {rec.name}: sparse_mla_decode calls {len(rec.sparse_mla_calls)}, declines {declined}"
+            )
+    if "rope_kv_append" in features:
+        bridge = [c for c in caps if c[0] == "rope_kv_append" and c[1] == "bridge"]
+        target_bridge = [c for c in bridge if c[4][0] == "target"]
+        assert len(target_bridge) == 2 * len(expected_kv), bridge
+        assert all(c[2] for c in target_bridge), bridge
+        if len(bridge) != num_attention_calls or not all(c[2] for c in bridge):
+            print(
+                f"[indexer_topk] {rec.name}: rope_kv_append bridge decisions {[c[2] for c in bridge]}"
+            )
+    if "draft_loop" in features:
+        steps = [c for c in caps if c[0] == "draft_loop" and c[1].startswith("draft_step:")]
+        assert len(steps) == DRAFT_LEN * len(ACCEPTED), steps
+        assert all(c[1] == "draft_step:engaged" for c in steps), steps
+
+
+def _pairs(rec_a: Recorder, rec_b: Recorder):
+    assert [(e["kind"], e["tag"]) for e in rec_a.events] == [
+        (e["kind"], e["tag"]) for e in rec_b.events
+    ]
+    return zip(rec_a.events, rec_b.events)
+
+
+def compare_selection(rec_a: Recorder, rec_b: Recorder, draft_tokens: bool = True):
+    """Discrete invariants: identical selections, shared rows and stashes."""
+    for a, b in _pairs(rec_a, rec_b):
+        if a["kind"] == "indexer":
+            if a["rows"] is not None:
+                assert torch.equal(a["rows"], b["rows"]), (rec_a.name, rec_b.name, a["tag"])
+            if "stash" in a:
+                sa = torch.sort(a["stash"], dim=1).values
+                sb = torch.sort(b["stash"], dim=1).values
+                assert torch.equal(sa, sb), (rec_a.name, rec_b.name, a["tag"], "stash")
+        elif a["kind"] == "shared_rows":
+            assert torch.equal(a["rows"], b["rows"]), (rec_a.name, rec_b.name, a["tag"])
+        elif a["kind"] == "draft_tokens" and draft_tokens:
+            assert torch.equal(a["tokens"], b["tokens"]), (rec_a.name, rec_b.name, a["tag"])
+
+
+def output_diffs(rec_a: Recorder, rec_b: Recorder) -> Dict[tuple, float]:
+    diffs = {}
+    for a, b in _pairs(rec_a, rec_b):
+        if a["kind"] == "output":
+            fa, fb = a["out"].float(), b["out"].float()
+            assert torch.isfinite(fa).all() and torch.isfinite(fb).all()
+            diffs[(a["tag"], a["layer"])] = (fa - fb).abs().max().item()
+    return diffs
+
+
+def _report(title: str, diffs: Dict[tuple, float]):
+    worst = max(diffs.values()) if diffs else 0.0
+    nonzero = sum(1 for v in diffs.values() if v != 0.0)
+    print(
+        f"[indexer_topk] {title}: {len(diffs)} layer outputs, {nonzero} differ, max |d| = {worst:.3e}"
+    )
+
+
+# --------------------------------------------------------------------- tests
+# DeepGEMM's decode schedule balances ceil(kv / SPLIT_KV) work units per
+# request over the SMs (SPLIT_KV is one of 64..512 tokens), so it changes only
+# where kv crosses a multiple of SPLIT_KV. 66,046 and 1,048,062 put the KV
+# length the first draft loop leaves behind (L + a1 - 1) on a multiple of 512
+# and the next target step's (L + a1) one past it, so a leftover draft-step
+# schedule differs from the target schedule for every SPLIT_KV candidate.
+SCHEDULE_SPLIT_LCM = 512
+CONTEXT_LENS = [66046, 491061, 1048062]
+
+
+def _is_schedule_boundary(context_len: int) -> bool:
+    return (context_len + ACCEPTED[0] - 1) % SCHEDULE_SPLIT_LCM == 0
+
+
+@pytest.fixture(scope="module", params=CONTEXT_LENS, ids=lambda n: f"kv{n}")
+def pipeline(request):
+    pipe = GlmDsaPipeline(request.param)
+    yield pipe
+    pipe.shutdown()
+
+
+def _served_order_report(title: str, rec_a: Recorder, rec_b: Recorder) -> float:
+    diffs = output_diffs(rec_a, rec_b)
+    outputs_a = {(e["tag"], e["layer"]): e["out"] for e in rec_a.events if e["kind"] == "output"}
+    outputs_b = {(e["tag"], e["layer"]): e["out"] for e in rec_b.events if e["kind"] == "output"}
+    worst_cos = 1.0
+    for key in diffs:
+        cos = torch.nn.functional.cosine_similarity(
+            outputs_a[key].float().flatten(), outputs_b[key].float().flatten(), dim=0
+        ).item()
+        worst_cos = min(worst_cos, cos)
+    _report(title, diffs)
+    print(f"[indexer_topk] {title}: min cosine {worst_cos:.7f}")
+    return worst_cos
+
+
+def test_indexer_topk_in_real_decode_pipeline(pipeline, monkeypatch):
+    """indexer_topk on vs off in the real metadata/draft-loop pipeline (sets and outputs)."""
+    print(f"[indexer_topk] base features '{BASE_FEATURES}', kv {pipeline.context_len}")
+    native_a = run_arm(pipeline, monkeypatch, "native_a", with_indexer_topk=False, canonical=True)
+    native_b = run_arm(pipeline, monkeypatch, "native_b", with_indexer_topk=False, canonical=True)
+    with_indexer_topk = run_arm(
+        pipeline, monkeypatch, "indexer_topk", with_indexer_topk=True, canonical=True
+    )
+    for rec, on in ((native_a, False), (native_b, False), (with_indexer_topk, True)):
+        check_arm_invariants(rec, pipeline, on)
+    compare_selection(native_a, native_b)
+    compare_selection(native_a, with_indexer_topk)
+    native_env = output_diffs(native_a, native_b)
+    indexer_topk_diff = output_diffs(native_a, with_indexer_topk)
+    _report("native vs native (canonical order)", native_env)
+    _report("indexer_topk vs native (canonical order)", indexer_topk_diff)
+    for key, value in indexer_topk_diff.items():
+        # Identical index sets in the same order: the attention is the same
+        # computation, so allow only what native shows run to run.
+        assert value <= native_env[key], (key, value, native_env[key])
+
+    # Served order (no canonicalization): the sets still match. The outputs
+    # then depend on the native top-k's run-to-run in-row order (the FP8
+    # attention tiles the selected tokens in that order), for native and indexer_topk
+    # alike; this is characterized, not bounded by a tolerance.
+    native_raw = run_arm(
+        pipeline, monkeypatch, "native_raw", with_indexer_topk=False, canonical=False
+    )
+    native_raw_b = run_arm(
+        pipeline, monkeypatch, "native_raw_b", with_indexer_topk=False, canonical=False
+    )
+    indexer_topk_raw = run_arm(
+        pipeline, monkeypatch, "indexer_topk_raw", with_indexer_topk=True, canonical=False
+    )
+    check_arm_invariants(native_raw, pipeline, False)
+    check_arm_invariants(native_raw_b, pipeline, False)
+    check_arm_invariants(indexer_topk_raw, pipeline, True)
+    compare_selection(native_raw, native_raw_b, draft_tokens=False)
+    compare_selection(native_raw, indexer_topk_raw, draft_tokens=False)
+    compare_selection(native_a, indexer_topk_raw, draft_tokens=False)
+    native_cos = _served_order_report("native vs native (served order)", native_raw, native_raw_b)
+    indexer_topk_cos = _served_order_report(
+        "indexer_topk vs native (served order)", native_raw, indexer_topk_raw
+    )
+    assert native_cos > 0.999 and indexer_topk_cos > 0.999, (native_cos, indexer_topk_cos)
+
+    # The freshness check is not vacuous at the boundary lengths: the schedule
+    # the first draft loop leaves behind is that of its last step and differs
+    # from the next target step's.
+    stale = [e for e in native_a.events if e["kind"] == "stale_sched"]
+    targets = _indexer_calls(native_a, "target")
+    differs = [not torch.equal(s["sched"], t["sched"]) for s, t in zip(stale, targets[1:])]
+    print(
+        f"[indexer_topk] leftover draft schedule differs from the next target schedule: {differs}"
+    )
+    if _is_schedule_boundary(pipeline.context_len):
+        leftover_kv = torch.tensor(
+            [pipeline.context_len + ACCEPTED[0] - 1], dtype=torch.int32, device=_DEV
+        )
+        leftover = _fresh_schedule(leftover_kv, targets[1]["num_sms"])
+        assert torch.equal(stale[0]["sched"], leftover), "leftover is the last draft step's"
+        assert differs[0], "boundary case must exercise a schedule change"
+    caps = sorted({c[:3] for c in with_indexer_topk.capability}, key=str)
+    print(f"[indexer_topk] capability decisions in the indexer_topk arm: {caps}")
+    print(
+        f"[indexer_topk] sparse_mla_decode sparse-MLA calls per arm: {len(with_indexer_topk.sparse_mla_calls)}"
+    )
+    print(f"[indexer_topk] indexer block-table width: {targets[0]['block_table_cols']}")
+    print(f"[indexer_topk] indexer_topk decisions: {dict(indexer_topk.DECISION_COUNTS)}")
+
+
+def _mpi_comm():
+    try:
+        from mpi4py import MPI
+    except ImportError:
+        return None
+    comm = MPI.COMM_WORLD
+    return comm if comm.Get_size() > 1 else None
+
+
+def _digest(t: torch.Tensor) -> str:
+    t = t.detach().contiguous().cpu()
+    if t.dtype in (torch.bfloat16, torch.float16):
+        t = t.view(torch.int16)
+    elif t.element_size() == 1:
+        t = t.view(torch.uint8)
+    return hashlib.sha256(t.numpy().tobytes()).hexdigest()[:16]
+
+
+@pytest.mark.parametrize("context_len", [491061])
+def test_rank_agreement(context_len, monkeypatch):
+    """Every rank of the MPI job selects identical sets with indexer_topk on and off."""
+    comm = _mpi_comm()
+    if comm is None:
+        pytest.skip("needs an MPI job with more than one rank")
+    rank, size = comm.Get_rank(), comm.Get_size()
+    pipe = GlmDsaPipeline(context_len)
+    try:
+        native = run_arm(pipe, monkeypatch, "native", with_indexer_topk=False, canonical=True)
+        with_indexer_topk = run_arm(
+            pipe, monkeypatch, "indexer_topk", with_indexer_topk=True, canonical=True
+        )
+        check_arm_invariants(native, pipe, False)
+        check_arm_invariants(with_indexer_topk, pipe, True)
+        compare_selection(native, with_indexer_topk)
+        kvm = pipe.kv_cache_manager
+        first_page = (
+            int(kvm.get_pool_block_indices(1, request_ids=[0])[0, 0])
+            * kvm.indexer_k_cache_page_scale
+        )
+        props = torch.cuda.get_device_properties(torch.cuda.current_device())
+
+        def selections(rec):
+            return [
+                (str(e["tag"]), e["layer"], _digest(e["rows"]))
+                for e in rec.events
+                if e["kind"] == "indexer" and e["rows"] is not None
+            ]
+
+        digest = {
+            "host": socket.gethostname(),
+            "gpu": str(getattr(props, "uuid", props.name)),
+            "indexer_k_first_page": _digest(
+                kvm.get_indexer_k_cache_buffers(LAYER_FULL)[first_page]
+            ),
+            "selections": selections(with_indexer_topk),
+            "native_selections": selections(native),
+            "outputs": [
+                (str(e["tag"]), e["layer"], _digest(e["out"]))
+                for e in with_indexer_topk.events
+                if e["kind"] == "output"
+            ],
+            "indexer_topk_decisions": dict(indexer_topk.DECISION_COUNTS),
+        }
+    finally:
+        pipe.shutdown()
+    everyone = comm.allgather(digest)
+    if rank == 0:
+        hosts = sorted({d["host"] for d in everyone})
+        print(f"[indexer_topk] {size} ranks on {len(hosts)} hosts: {hosts}")
+        for r, d in enumerate(everyone):
+            print(
+                f"[indexer_topk] rank {r} {d['host']} {d['gpu']} indexer_topk decisions {d['indexer_topk_decisions']}"
+            )
+        print(f"[indexer_topk] rank 0 selection digests: {everyone[0]['selections']}")
+    ref = everyone[0]
+    for r, d in enumerate(everyone):
+        assert d["indexer_k_first_page"] == ref["indexer_k_first_page"], r
+        assert d["selections"] == ref["selections"], r
+        assert d["native_selections"] == ref["native_selections"], r
+        assert d["selections"] == d["native_selections"], r
+        assert d["indexer_topk_decisions"] == ref["indexer_topk_decisions"], r
+        # The attention is replicated per rank on identical inputs as well.
+        assert d["outputs"] == ref["outputs"], r
+
+
+# ------------------------------------------------------------------- timing
+TIMING_LENS = [4224, 16384, 32768, 65536, 131072, 314721, 391000, 491061]
+_TIMING_REPS = 16
+_TIMING_ROUNDS = 15
+_TIMING_REPLAYS = 5
+
+
+def _replay_us(graph: torch.cuda.CUDAGraph, replays: int) -> float:
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(replays):
+        graph.replay()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / replays
+
+
+def _median(values):
+    values = sorted(values)
+    return values[len(values) // 2]
+
+
+def _capture_region(pipe, monkeypatch, kv: int, with_indexer_topk: bool, flush: torch.Tensor):
+    """CUDA graph of ``_TIMING_REPS`` x (L2 flush, full layer + shared layer attention)."""
+    with monkeypatch.context() as mp, torch.inference_mode():
+        _select(mp, with_indexer_topk)
+        md = pipe.prepare_target_step(host_cached=kv - NEXT_N, kv_offset=0)
+        positions = torch.arange(kv - NEXT_N, kv, dtype=torch.int32, device=_DEV)
+        work = []
+        for layer_idx in (LAYER_FULL, LAYER_SHARED):
+            inputs = pipe.make_inputs(seed=500 + layer_idx, n=NEXT_N, positions=positions)
+            work.append((layer_idx, inputs, pipe.project(layer_idx, inputs)))
+
+        def region():
+            for layer_idx, inputs, intermediates in work:
+                pipe.attend(layer_idx, md, inputs, intermediates)
+
+        engaged = indexer_topk.DECISION_COUNTS.get("engaged", 0)
+        region()  # compiles the kernels of this shape outside the capture
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, capture_error_mode="thread_local"):
+            for _ in range(_TIMING_REPS):
+                flush.zero_()
+                region()
+        torch.cuda.synchronize()
+        engaged = indexer_topk.DECISION_COUNTS.get("engaged", 0) - engaged
+    assert engaged == ((_TIMING_REPS + 1) if with_indexer_topk else 0), engaged
+    return graph, work
+
+
+def _time_arm(pipe, monkeypatch, kv: int, with_indexer_topk: bool, flush, flush_graph):
+    """Per-region microseconds (median, spread) of one arm at one KV length.
+
+    The graph is replayed before the next ``prepare()``: the non-graph metadata
+    re-derives per-step tensors (such as the cached pool view) there, which a
+    graph captured earlier would still point at.
+    """
+    graph, work = _capture_region(pipe, monkeypatch, kv, with_indexer_topk, flush)
+    flush_samples, samples = [], []
+    for _ in range(_TIMING_ROUNDS):
+        flush_samples.append(_replay_us(flush_graph, _TIMING_REPLAYS))
+        samples.append(_replay_us(graph, _TIMING_REPLAYS))
+    region_us = (_median(samples) - _median(flush_samples)) / _TIMING_REPS
+    spread_us = (max(samples) - min(samples)) / _TIMING_REPS
+    del graph, work
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    return region_us, spread_us
+
+
+@pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_INDEXER_TOPK_PIPELINE_TIMING") != "1",
+    reason="set GLM_MOE_DSA_INDEXER_TOPK_PIPELINE_TIMING=1 to time",
+)
+def test_indexer_topk_region_timing(monkeypatch):
+    """Attention region of a full + shared layer pair, indexer_topk off vs on, in CUDA graphs.
+
+    The region is ``_forward_dsa_attn`` of both layers with the served
+    neighbours of the indexer (indexer K scatter, shared-layer copy, remap,
+    RoPE/FP8-query/KV append, sparse MLA). Each capture flushes L2 before
+    every region (the served K caches of 21 full layers are not L2 resident);
+    the flush alone is timed in its own graph and subtracted. Native is timed
+    before and after indexer_topk at each length; their difference is the drift.
+    """
+    pipe = GlmDsaPipeline(max(TIMING_LENS) + TOKENS_PER_BLOCK)
+    flush = torch.empty(1 << 29, dtype=torch.uint8, device=_DEV)
+    results = []
+    try:
+        flush_graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(flush_graph, capture_error_mode="thread_local"):
+            for _ in range(_TIMING_REPS):
+                flush.zero_()
+        for kv in TIMING_LENS:
+            native_1, spread_1 = _time_arm(pipe, monkeypatch, kv, False, flush, flush_graph)
+            indexer_topk_us, indexer_topk_spread = _time_arm(
+                pipe, monkeypatch, kv, True, flush, flush_graph
+            )
+            native_2, spread_2 = _time_arm(pipe, monkeypatch, kv, False, flush, flush_graph)
+            native_us = (native_1 + native_2) / 2
+            row = {
+                "kv": kv,
+                "base_features": BASE_FEATURES,
+                "native_region_us": round(native_us, 3),
+                "indexer_topk_region_us": round(indexer_topk_us, 3),
+                "delta_us": round(indexer_topk_us - native_us, 3),
+                "native_drift_us": round(native_2 - native_1, 3),
+                "native_spread_us": round(max(spread_1, spread_2), 3),
+                "indexer_topk_spread_us": round(indexer_topk_spread, 3),
+            }
+            results.append(row)
+            print(f"[indexer_topk-timing] {json.dumps(row)}")
+    finally:
+        pipe.shutdown()
+    assert len(results) == len(TIMING_LENS)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py
new file mode 100644
index 000000000000..3e737bd68d21
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py
@@ -0,0 +1,243 @@
+# 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.
+"""Eligibility table of the indexer_topk indexer decode package (no GPU needed).
+
+``policy.py`` is loaded by file path so this test needs neither torch nor the
+compiled extensions.
+"""
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+_POLICY = (
+    Path(__file__).resolve().parents[9]
+    / "tensorrt_llm"
+    / "_torch"
+    / "_experimental"
+    / "modeling_v2"
+    / "models"
+    / "glm_moe_dsa"
+    / "shared"
+    / "kernels"
+    / "indexer_topk"
+    / "policy.py"
+)
+
+
+def _load_policy():
+    spec = importlib.util.spec_from_file_location("indexer_topk_policy", _POLICY)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+# Served GEN: max_seq_len 2**20 with MTP K = 5 and the overlap
+# scheduler gives an indexer max_seq_len of 2**20 + 14, and KV cache manager V2
+# sizes the block table at round_up(ceil((2**20 + 14 + 10) / 64), 4) columns.
+SERVED_MAX_SEQ_LEN = (1 << 20) + 14
+V2_BLOCK_TABLE_COLS = 16388
+# KV cache manager V1 at the same max_seq_len: ceil((2**20 + 14) / 64) columns.
+V1_BLOCK_TABLE_COLS = 16385
+
+# Served GEN: sm103, CuTe-DSL logits, FP8 indexer keys, self-sampling GVR,
+# one request with 1 + 5 verify rows, 32 heads of dim 128, 64-token pages,
+# TP8 without attention DP, the served indexer max_seq_len.
+_SERVED_TP8 = dict(
+    sm_version=103,
+    cute_dsl_available=True,
+    native_uses_cute_dsl_logits=True,
+    use_fp4_indexer=False,
+    use_gvr_emission=False,
+    needs_gvr_prior=False,
+    is_ragged_verify=False,
+    expanded_mtp_buffers=False,
+    in_mtp_draft_loop=False,
+    attention_dp=False,
+    topk_decode_impl="cute_dsl_gvr",
+    topk_self_sampling=True,
+    index_topk=2048,
+    compress_ratio=1,
+    num_heads=32,
+    head_dim=128,
+    query_dtype="torch.float8_e4m3fn",
+    query_shape=(6, 32, 128),
+    query_contiguous=True,
+    weights_dtype="torch.float32",
+    weights_shape=(6, 32),
+    weights_contiguous=True,
+    k_cache_dtype="torch.uint8",
+    k_cache_shape=(16384, 64, 1, 132),
+    block_table_cols=V2_BLOCK_TABLE_COLS,
+    max_seq_len=SERVED_MAX_SEQ_LEN,
+    num_contexts=0,
+    token_offset=0,
+    num_gen_tokens=6,
+    num_generations=1,
+    uniform_verify_width=True,
+)
+
+
+def _with_shape(heads=32, head_dim=128, page=64, next_n=6, batch=1, **extra):
+    """Served decision inputs with a different indexer/verify shape."""
+    rows = next_n * batch
+    return {
+        **_SERVED_TP8,
+        "num_heads": heads,
+        "head_dim": head_dim,
+        "query_shape": (rows, heads, head_dim),
+        "weights_shape": (rows, heads),
+        "k_cache_shape": (16384, page, 1, head_dim + 4),
+        "num_gen_tokens": rows,
+        "num_generations": batch,
+        **extra,
+    }
+
+
+def test_the_served_tp8_shape_is_eligible():
+    decision = _load_policy().decide(**_SERVED_TP8)
+    assert decision.eligible
+    assert decision.reason == ""
+
+
+@pytest.mark.parametrize(
+    "override, reason_part",
+    [
+        ({"sm_version": 90}, "sm90"),
+        # SM100 has the same kernels but has not been validated on a GPU.
+        ({"sm_version": 100}, "sm100"),
+        ({"sm_version": 120}, "sm120"),
+        ({"cute_dsl_available": False}, "DSL"),
+        ({"native_uses_cute_dsl_logits": False}, "CuTe-DSL"),
+        ({"use_fp4_indexer": True}, "FP4"),
+        ({"use_gvr_emission": True}, "emission"),
+        ({"needs_gvr_prior": True}, "prior"),
+        ({"is_ragged_verify": True}, "ragged"),
+        ({"expanded_mtp_buffers": True}, "expanded"),
+        ({"in_mtp_draft_loop": True}, "draft"),
+        ({"attention_dp": True}, "data parallelism"),
+        ({"topk_decode_impl": "cuda_radix"}, "top-k engine"),
+        ({"topk_decode_impl": "cute_dsl_radix"}, "top-k engine"),
+        ({"topk_self_sampling": False}, "top-k engine"),
+        ({"index_topk": 512}, "index_topk"),
+        ({"compress_ratio": 4}, "compress_ratio"),
+        ({"query_dtype": "torch.bfloat16"}, "query dtype"),
+        ({"query_contiguous": False}, "layout"),
+        ({"query_shape": (6, 32, 64)}, "layout"),
+        ({"weights_dtype": "torch.bfloat16"}, "weights dtype"),
+        ({"weights_contiguous": False}, "layout"),
+        ({"weights_shape": (6, 16)}, "layout"),
+        ({"k_cache_dtype": "torch.float8_e4m3fn"}, "K cache layout"),
+        ({"k_cache_shape": (16384, 64, 1, 68)}, "K cache layout"),
+        ({"k_cache_shape": (16384, 64, 132)}, "K cache view"),
+        ({"max_seq_len": SERVED_MAX_SEQ_LEN + 1}, "max_seq_len"),
+        ({"max_seq_len": 1 << 21}, "max_seq_len"),
+        ({"max_seq_len": 0}, "max_seq_len"),
+        # V1 table at the served max_seq_len: the last 128-token tile of a
+        # 2**20 + 1 .. 2**20 + 14 token request would read one entry past the row.
+        ({"block_table_cols": V1_BLOCK_TABLE_COLS}, "block table width"),
+        ({"block_table_cols": 0}, "block table width"),
+        ({"num_contexts": 1}, "context requests"),
+        ({"token_offset": 1}, "exceed"),
+        ({"num_generations": 0}, "no generation"),
+        (
+            {
+                "num_gen_tokens": 7,
+                "num_generations": 2,
+                "query_shape": (7, 32, 128),
+                "weights_shape": (7, 32),
+            },
+            "non-uniform",
+        ),
+        ({"uniform_verify_width": False}, "non-uniform"),
+    ],
+)
+def test_native_fallback_cases(override, reason_part):
+    decision = _load_policy().decide(**{**_SERVED_TP8, **override})
+    assert not decision.eligible
+    assert reason_part in decision.reason
+
+
+@pytest.mark.parametrize(
+    "shape, reason_part",
+    [
+        # Constructor-invalid for the rows kernel (8 cached weights > 4 heads):
+        # must fall back before any kernel is built.
+        (dict(heads=4, next_n=4), "num_heads"),
+        (dict(heads=16), "num_heads"),
+        (dict(heads=64), "num_heads"),
+        (dict(head_dim=64), "head_dim"),
+        (dict(page=16), "tokens_per_block"),
+        (dict(page=32), "tokens_per_block"),
+        (dict(page=128), "tokens_per_block"),
+        (dict(next_n=1), "next_n"),
+        (dict(next_n=4), "next_n"),
+        (dict(next_n=8), "next_n"),
+        (dict(batch=5), "generation requests"),
+    ],
+    ids=lambda v: "-".join(f"{k}{w}" for k, w in v.items()) if isinstance(v, dict) else v,
+)
+def test_unverified_shapes_fall_back(shape, reason_part):
+    decision = _load_policy().decide(**_with_shape(**shape))
+    assert not decision.eligible
+    assert reason_part in decision.reason
+
+
+def test_envelope_is_the_served_bound():
+    policy = _load_policy()
+    assert policy.MAX_INDEXER_SEQ_LEN == SERVED_MAX_SEQ_LEN
+    # net max_seq_len + (K) overlap lead + (K - 1) extra KV tokens + (K) tokens
+    # per generation step beyond the golden token, for K = 5.
+    assert policy.MAX_INDEXER_SEQ_LEN == (1 << 20) + 5 + 4 + 5
+
+
+@pytest.mark.parametrize(
+    "max_seq_len, block_table_cols",
+    [
+        (SERVED_MAX_SEQ_LEN, V2_BLOCK_TABLE_COLS),
+        # V1 and V2 at the 2**20 net envelope without MTP slack.
+        (1 << 20, 16384),
+        (1 << 20, 16388),
+        (131072, 2048),
+    ],
+)
+def test_served_lengths_and_table_widths_are_eligible(max_seq_len, block_table_cols):
+    decision = _load_policy().decide(
+        **{**_SERVED_TP8, "max_seq_len": max_seq_len, "block_table_cols": block_table_cols}
+    )
+    assert decision.eligible, decision.reason
+
+
+@pytest.mark.parametrize("batch", [1, 2, 3, 4])
+def test_verified_batch_sizes(batch):
+    assert _load_policy().decide(**_with_shape(batch=batch)).eligible
+
+
+def test_eligible_set_is_the_verified_envelope():
+    """No shape outside the GPU-verified one is eligible."""
+    policy = _load_policy()
+    eligible = set()
+    for sm in (90, 100, 103, 120):
+        for heads in (4, 8, 16, 32, 64):
+            for head_dim in (64, 128, 256):
+                for page in (16, 32, 64, 128):
+                    for next_n in range(1, 10):
+                        for batch in range(1, 9):
+                            inputs = _with_shape(heads, head_dim, page, next_n, batch)
+                            inputs["sm_version"] = sm
+                            if policy.decide(**inputs).eligible:
+                                eligible.add((sm, heads, head_dim, page, next_n, batch))
+    assert eligible == {(103, 32, 128, 64, 6, b) for b in range(1, 5)}
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py
new file mode 100644
index 000000000000..81e6f86dfff8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py
@@ -0,0 +1,839 @@
+# 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.
+"""oproj_weight_staging ``oproj_weight_staging`` (single GPU): the staging o_proj against the served
+attention_output_tail.
+
+Oracle: the served attention_output_tail/attention_output_into_hop kernels
+(``attention_output_tail.kernels.attn_tail``, its ``stream_oproj`` and ``stream_oproj_into_hop``) on
+the same inputs. oproj_weight_staging changes no arithmetic, so every check is bit-for-bit
+(``torch.equal``), on real GLM-5.2 TP8 weight slices (``tail_fixtures.load_tail_weights``) and
+real-weight latents (``tail_fixtures.mix_latent``):
+
+* every staging configuration (L2 prefetch or shared-memory staging, at
+  entry or once the early rows landed, from the o_proj or from the W_UV
+  kernel, cache-policy hints, partial staging), on the served schedule
+  (``INTO_HOP_SPARSE_DECODE_SCHEDULE``: 128 CTAs, 32 early rows) and on other ``early_rows``
+  / grids; the BF16 ``v`` intermediate as well;
+* the attention_output_into_hop publish (a device buffer stands in for the post_attention_hop
+  transport, as in ``attention_output_tail/test_into_hop_producer.py``): same words, same untouched
+  sentinels;
+* padded rows (live < out rows, also out rows 9..16);
+* CUDA-graph capture and replay with fresh contents, alternating graphs of
+  different row counts;
+* a delayed producer of the latent (a kernel that releases its PDL
+  dependents at entry and writes the latent only after a ``%globaltimer``
+  spin; the tail's W_UV kernel waits for it) and a delayed producer of the
+  o_proj's own input ``v`` (the o_proj is its direct PDL dependent, so every
+  o_proj CTA is resident and done staging long before its wait returns), and
+  a co-running kernel on another stream;
+* ``%globaltimer`` stamps of the diagnostic build are ordered;
+* ``StagedAttnTail`` on the real DSA ``MLA`` module: the attention_output_tail protocol (arm,
+  defer_w_uv, project, into-hop request) is unchanged and the partial equals
+  attention_output_tail's; schedules oproj_weight_staging does not take keep attention_output_tail's
+  tail; geometry declines run attention_output_tail.
+
+Opt-in (``GLM_MOE_DSA_OPROJ_WEIGHT_STAGING_ADAPTER_PATCH_APPLIED=1``): the adopted attention
+(``GlmMLA.adopt``) wraps the tail;
+every set of fast paths containing ``oproj_weight_staging`` that a plan can run engages it (the
+closure is the engaged domain, so no runnable set leaves it inert); deferred_input_prep's gate
+maps ``oproj_weight_staging`` to its directories and side selector.
+"""
+
+import dataclasses
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
+import tail_fixtures  # noqa: E402
+import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
+
+sys.path.insert(0, os.path.dirname(_HERE))
+import plan_testlib  # noqa: E402
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
+    dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
+    INTO_HOP_SCHEDULE,
+    INTO_HOP_SPARSE_DECODE_SCHEDULE,
+    TAIL_PREFETCH_SCHEDULE,
+    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
+    TailConfig,
+    attn_tail,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
+    AttnTail,
+    make_attn_tail,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging import (  # noqa: E402, E501
+    DEFAULT_STAGE,
+    StageConfig,
+    StagedAttnTail,
+    stage_tail,
+    wrap_attn_tail,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging.cute_oproj_stage import (  # noqa: E402, E501
+    TRACE_POINTS,
+    stream_oproj_stage,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging.kernels import (  # noqa: E402, E501
+    check_stage_config,
+    wuv_stage_geometry,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging.tail import (  # noqa: E402, E501
+    config_reason,
+    geometry_reason,
+)
+
+_CKPT = tail_fixtures.checkpoint_dir()
+_PATCHED = os.environ.get("GLM_MOE_DSA_OPROJ_WEIGHT_STAGING_ADAPTER_PATCH_APPLIED") == "1"
+needs_patch = pytest.mark.skipif(not _PATCHED, reason="opt-in test of the adopted attention")
+# Served GEN selector (without deferred_input_prep, which the eligibility test adds).
+B1_GEN = (
+    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,"
+    "draft_front,router_gemm,post_attention_hop,draft_tail,"
+    "sparse_mla_decode,draft_absorbed_query,rope_kv_append,draft_loop,"
+    "shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
+    "attention_output_tail_prefetch,attention_output_into_hop,"
+    "routed_moe_into_hop,deferred_norm,dense_mlp"
+)
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="oproj_weight_staging targets SM100/SM103",
+    ),
+    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
+]
+
+TP_SIZE = 8
+LOCAL_HEADS = 8
+H = 6144
+RANKS = 8
+MAX_ROWS = 8
+BUF_WORDS = MAX_ROWS * RANKS * H * 2 // 4
+SENTINEL = -(2**31)
+
+S = StageConfig
+ARMS = {
+    "none": S(stage="none"),
+    "l2_entry": S(stage="l2", when="entry"),
+    "l2_landed": S(stage="l2", when="landed"),
+    "l2_landed_el": S(stage="l2", when="landed", prefetch_policy="evict_last"),
+    "l2_half": S(stage="l2", when="landed", rows=8),
+    "smem_entry": S(stage="smem", when="entry"),
+    "smem_landed": S(stage="smem", when="landed"),
+    "smem_half": S(stage="smem", when="landed", rows=8),
+    "wuv": S(stage="l2", host="wuv"),
+    "wuv_el": S(stage="l2", host="wuv", prefetch_policy="evict_last"),
+    "ef": S(stage="none", copy_policy="evict_first"),
+    "l2_landed_ef": S(stage="l2", when="landed", copy_policy="evict_first"),
+    "smem_landed_ef": S(stage="smem", when="landed", copy_policy="evict_first"),
+    "smem_chain": S(stage="smem", when="chain"),
+    "smem_chain_e16": S(stage="smem", when="chain", early_rows=16),
+    "smem_chain_e8_ef": S(stage="smem", when="chain", early_rows=8, copy_policy="evict_first"),
+    "default": DEFAULT_STAGE,
+}
+PUBLISH_ARMS = (
+    "none",
+    "l2_landed",
+    "smem_landed",
+    "wuv",
+    "l2_landed_ef",
+    "smem_chain_e16",
+    "default",
+)
+
+_W = {}
+_COMP = {}
+
+
+def _weights(layer, rank):
+    if layer not in _W:
+        _W[layer] = {
+            r: (w_uv.cuda().contiguous(), w_o.cuda().contiguous())
+            for r, (w_uv, w_o) in tail_fixtures.load_tail_weights(
+                _CKPT, layer, range(TP_SIZE), TP_SIZE
+            ).items()
+        }
+    return _W[layer][rank]
+
+
+def _latent(layer, rows, seed):
+    if layer not in _COMP:
+        _COMP[layer] = tail_fixtures.compressed_kv_rows(_CKPT, layer)
+    return tail_fixtures.mix_latent(_COMP[layer], rows, LOCAL_HEADS, seed=seed)
+
+
+def _transport(cur=1, rank=3):
+    mc = torch.full((3 * BUF_WORDS,), SENTINEL, dtype=torch.int32, device="cuda")
+    flags = torch.tensor([cur, 0, 0, 0], dtype=torch.int32, device="cuda")
+    return SimpleNamespace(mc_words=mc, flags=flags, rank=rank, cur=cur)
+
+
+def _pair(layer, rank, rows, out_rows, cfg, stage, seed, into_hop_pair=False):
+    """Run attention_output_tail's tail and oproj_weight_staging's tail on the same inputs; return both outputs
+    (and the v intermediates, and the two transports when ``into_hop_pair``)."""
+    w_uv, w_o = _weights(layer, rank)
+    lat = _latent(layer, rows, seed)
+    v_attn_tail = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    v_oproj_stage = torch.empty_like(v_attn_tail)
+    tr_attn_tail = _transport() if into_hop_pair else None
+    tr_oproj_stage = _transport() if into_hop_pair else None
+    ref = attn_tail(
+        lat,
+        w_uv,
+        w_o,
+        rows,
+        out_rows=out_rows,
+        v_scratch=v_attn_tail,
+        config=cfg,
+        hop_transport=tr_attn_tail,
+    )
+    got = stage_tail(
+        lat,
+        w_uv,
+        w_o,
+        rows,
+        out_rows=out_rows,
+        v_scratch=v_oproj_stage,
+        config=cfg,
+        stage=stage,
+        hop_transport=tr_oproj_stage,
+    )
+    torch.cuda.synchronize()
+    return ref, got, v_attn_tail, v_oproj_stage, tr_attn_tail, tr_oproj_stage
+
+
+@pytest.mark.parametrize("arm", list(ARMS))
+@pytest.mark.parametrize("layer,rank", [(3, 0), (77, 7)])
+@pytest.mark.parametrize("rows", [1, 6, 8])
+def test_served_schedule_bit_identical(arm, layer, rank, rows):
+    ref, got, v_attn_tail, v_oproj_stage, _, _ = _pair(
+        layer, rank, rows, rows, INTO_HOP_SPARSE_DECODE_SCHEDULE, ARMS[arm], seed=10 * rows + rank
+    )
+    assert torch.equal(v_attn_tail, v_oproj_stage)
+    assert torch.equal(ref, got), (arm, layer, rank, rows)
+
+
+@pytest.mark.parametrize("arm", PUBLISH_ARMS)
+@pytest.mark.parametrize("live,out_rows", [(1, 1), (6, 6), (8, 8), (6, 8), (1, 4)])
+def test_into_hop_publish_bit_identical(arm, live, out_rows):
+    ref, got, _, _, tr_attn_tail, tr_oproj_stage = _pair(
+        40,
+        3,
+        live,
+        out_rows,
+        INTO_HOP_SPARSE_DECODE_SCHEDULE,
+        ARMS[arm],
+        seed=live + out_rows,
+        into_hop_pair=True,
+    )
+    assert torch.equal(ref, got)
+    assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words), arm
+    buf = tr_oproj_stage.mc_words[
+        tr_oproj_stage.cur * BUF_WORDS : (tr_oproj_stage.cur + 1) * BUF_WORDS
+    ].view(MAX_ROWS, RANKS, H // 2)
+    assert not bool((buf[:out_rows, tr_oproj_stage.rank] == SENTINEL).any())
+
+
+@pytest.mark.parametrize(
+    "arm", ["none", "l2_entry", "l2_landed", "smem_landed", "smem_chain_e16", "wuv", "default"]
+)
+@pytest.mark.parametrize("live,out_rows,publish", [(6, 6, True), (1, 4, True), (6, 12, False)])
+def test_synthetic_inputs_bit_identical(arm, live, out_rows, publish):
+    """Every input written by plain stores (``torch.randn``/``full``), for
+    compute-sanitizer initcheck: the real-weight fixtures run cuBLAS GEMMs
+    whose outputs initcheck reports as uninitialized when torch reads them."""
+    gen = torch.Generator(device="cuda").manual_seed(live + out_rows)
+    w_uv = (torch.randn((8, 256, 512), generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
+    w_o = (torch.randn((H, 2048), generator=gen, device="cuda") * 0.02).to(torch.bfloat16)
+    lat = torch.randn((live, 8, 512), generator=gen, device="cuda").to(torch.bfloat16)
+    tr_attn_tail = _transport() if publish else None
+    tr_oproj_stage = _transport() if publish else None
+    ref = attn_tail(
+        lat,
+        w_uv,
+        w_o,
+        live,
+        out_rows=out_rows,
+        config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
+        hop_transport=tr_attn_tail,
+    )
+    got = stage_tail(
+        lat,
+        w_uv,
+        w_o,
+        live,
+        out_rows=out_rows,
+        config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
+        stage=ARMS[arm],
+        hop_transport=tr_oproj_stage,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(ref, got)
+    if publish:
+        assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words)
+
+
+@pytest.mark.parametrize(
+    "cfg_name,cfg",
+    [
+        ("attn_tail_default", TailConfig()),
+        ("g128_e24", TailConfig(oproj_grid=128, early_rows=24)),
+        ("g128_e40", TailConfig(oproj_grid=128, early_rows=40)),
+        ("g128_e48", TailConfig(oproj_grid=128, early_rows=48)),
+        ("attention_output_into_hop", INTO_HOP_SCHEDULE),
+    ],
+)
+@pytest.mark.parametrize(
+    "arm", ["l2_entry", "l2_landed", "smem_landed", "smem_chain", "default", "ef"]
+)
+def test_other_schedules_bit_identical(cfg_name, cfg, arm):
+    ref, got, v_attn_tail, v_oproj_stage, _, _ = _pair(3, 5, 6, 6, cfg, ARMS[arm], seed=7)
+    assert torch.equal(v_attn_tail, v_oproj_stage)
+    assert torch.equal(ref, got), (cfg_name, arm)
+
+
+@pytest.mark.parametrize("early", [16, 24, 40])
+def test_early_rows_override(early):
+    stage = S(stage="l2", when="landed", early_rows=early)
+    cfg = dataclasses.replace(INTO_HOP_SPARSE_DECODE_SCHEDULE, early_rows=early)
+    w_uv, w_o = _weights(3, 2)
+    lat = _latent(3, 6, 99)
+    ref = attn_tail(lat, w_uv, w_o, 6, config=cfg)
+    got = stage_tail(lat, w_uv, w_o, 6, config=INTO_HOP_SPARSE_DECODE_SCHEDULE, stage=stage)
+    torch.cuda.synchronize()
+    assert torch.equal(ref, got)
+
+
+@pytest.mark.parametrize("live,out_rows", [(6, 12), (8, 16), (1, 9)])
+@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
+def test_wide_padded_rows(live, out_rows, arm):
+    ref, got, _, _, _, _ = _pair(
+        2, 1, live, out_rows, INTO_HOP_SPARSE_DECODE_SCHEDULE, ARMS[arm], seed=live
+    )
+    assert torch.equal(ref, got)
+    assert bool((got[live:] == 0).all())
+
+
+@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
+def test_graph_replay_changing_rows(arm):
+    """Graphs for 1 live row (4 out rows) and 6 live rows (8 out rows),
+    replayed alternately with fresh latents; each replay equals attention_output_tail."""
+    w_uv, w_o = _weights(3, 0)
+    stage = ARMS[arm]
+    graphs = {}
+    for rows, out_rows in ((1, 4), (6, 8)):
+        lat = _latent(3, rows, 0).clone()
+        stage_tail(
+            lat,
+            w_uv,
+            w_o,
+            rows,
+            out_rows=out_rows,
+            config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
+            stage=stage,
+        )
+        torch.cuda.synchronize()
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            out = stage_tail(
+                lat,
+                w_uv,
+                w_o,
+                rows,
+                out_rows=out_rows,
+                config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
+                stage=stage,
+            )
+        graphs[rows] = (g, lat, out, out_rows)
+    for step in range(12):
+        rows = (1, 6)[step % 2]
+        g, lat, out, out_rows = graphs[rows]
+        lat.copy_(_latent(3, rows, 100 + step))
+        g.replay()
+        ref = attn_tail(
+            lat, w_uv, w_o, rows, out_rows=out_rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE
+        )
+        torch.cuda.synchronize()
+        assert torch.equal(out, ref), (arm, step)
+
+
+@triton.jit
+def _globaltimer(like):
+    return tl.inline_asm_elementwise(
+        "mov.u64 $0, %globaltimer;", "=l,r", [like], dtype=tl.int64, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _slow_copy_kernel(src, dst, n, delay_ns, stamps, BLOCK: tl.constexpr):
+    """One CTA, launched with PDL (Triton emits the ``griddepcontrol``
+    instructions only for such launches): waits for its predecessor,
+    releases its own dependents, spins ``delay_ns`` on ``%globaltimer``, then
+    copies ``src`` into ``dst`` (the latent). One CTA leaves SMs free, so
+    o_proj CTAs (197,400 B of shared memory each) can become resident next
+    to the parked W_UV grid."""
+    tl.extra.cuda.gdc_wait()
+    tl.extra.cuda.gdc_launch_dependents()
+    pid = tl.program_id(0)
+    start = _globaltimer(pid)
+    now = start
+    while now - start < delay_ns:
+        now = _globaltimer(pid)
+    for base in range(0, n, BLOCK):
+        offs = base + tl.arange(0, BLOCK)
+        mask = offs < n
+        tl.store(dst + offs, tl.load(src + offs, mask=mask), mask=mask)
+    # Start of the spin and end of the latent's stores (``%globaltimer``).
+    tl.store(stamps, start)
+    tl.store(stamps + 1, _globaltimer(pid))
+
+
+@triton.jit
+def _hold_kernel(delay_ns, BLOCK: tl.constexpr):
+    pid = tl.program_id(0)
+    start = _globaltimer(pid)
+    now = start
+    while now - start < delay_ns:
+        now = _globaltimer(pid)
+
+
+def _hold_stream(delay_ns=2_000_000):
+    """Occupy the stream with one spinning CTA (no PDL) while the host
+    enqueues the launches that follow. Without it the host's launch latency
+    (tens of us per Triton / CuTe launch) exceeds the producer's delay, the
+    producer has completed before its PDL dependents are even enqueued, and
+    they never park in ``griddepcontrol.wait``."""
+    _hold_kernel[(1,)](delay_ns, BLOCK=1)
+
+
+def _slow_produce(src, dst, delay_ns, stamps=None):
+    if stamps is None:
+        stamps = torch.zeros(2, dtype=torch.int64, device=src.device)
+    _slow_copy_kernel[(1,)](src, dst, src.numel(), delay_ns, stamps, BLOCK=1024, launch_pdl=True)
+    return stamps
+
+
+@pytest.mark.parametrize(
+    "arm",
+    ["l2_entry", "l2_landed", "smem_landed", "wuv", "l2_landed_ef", "smem_chain_e16", "default"],
+)
+@pytest.mark.parametrize("delay_ns", [0, 3000, 20000])
+def test_delayed_producer(arm, delay_ns):
+    """The latent's producer is late (it releases its PDL dependents, then
+    spins ``delay_ns`` before writing the latent; a held stream lets the
+    host enqueue the tail before the producer starts). Outputs equal attention_output_tail's
+    under the same producer, and no o_proj CTA leaves its wait before the
+    producer's last store (``%globaltimer`` stamps of both kernels)."""
+    w_uv, w_o = _weights(40, 6)
+    rows = 6
+    src = _latent(40, rows, 5).contiguous()
+    lat_attn_tail = torch.zeros_like(src)
+    lat_oproj_stage = torch.zeros_like(src)
+    tr = torch.zeros(128 * TRACE_POINTS, dtype=torch.int64, device="cuda")
+    for _ in range(2):
+        _hold_stream()
+        _slow_produce(src, lat_attn_tail, delay_ns)
+        ref = attn_tail(lat_attn_tail, w_uv, w_o, rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE)
+        _hold_stream()
+        ps = _slow_produce(src, lat_oproj_stage, delay_ns)
+        got = stage_tail(
+            lat_oproj_stage,
+            w_uv,
+            w_o,
+            rows,
+            config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
+            stage=ARMS[arm],
+            trace=tr,
+        )
+        torch.cuda.synchronize()
+        assert torch.equal(ref, got), (arm, delay_ns)
+        st = tr.view(128, TRACE_POINTS).cpu()
+        p_start, p_end = (int(x) for x in ps.cpu())
+        assert p_end - p_start >= delay_ns
+        # Every o_proj CTA passed its wait only after the producer's stores
+        # (through W_UV's completion), however early it became resident.
+        assert int(st[:, 1].min()) >= p_end, (arm, delay_ns)
+    parked = (st[:, 1] - st[:, 0]).float()
+    print(
+        f"delayed producer {delay_ns} ns, {arm}: o_proj resident "
+        f"{float(parked.mean()):.0f} ns (mean) / {float(parked.max()):.0f} "
+        f"ns (max) before its wait returned"
+    )
+
+
+@pytest.mark.parametrize(
+    "arm",
+    ["none", "l2_entry", "l2_landed", "smem_landed", "l2_landed_ef", "smem_chain_e16", "default"],
+)
+@pytest.mark.parametrize("publish", [False, True])
+def test_parked_oproj_bit_identical(arm, publish):
+    """The o_proj is the direct PDL dependent of a one-CTA producer that
+    releases its dependents at entry and writes ``v`` only after a 20 us
+    spin (behind a held stream, so the host has enqueued the o_proj before
+    the producer starts). All 128 o_proj CTAs fit next to it, so each one
+    issues its pre-wait copies and staging and then stays parked in
+    ``griddepcontrol.wait`` for most of the spin (the case where the staging
+    has long completed when the wait returns). Outputs (and the into-hop words) equal the served
+    tail's on the same latent, for the plain and the stamped build; no CTA leaves its wait
+    before the producer's last store; the median CTA was parked at least
+    half the delay."""
+    delay_ns = 20000
+    w_uv, w_o = _weights(40, 6)
+    rows = 6
+    stage = ARMS[arm]
+    early = (
+        INTO_HOP_SPARSE_DECODE_SCHEDULE.early_rows if stage.early_rows is None else stage.early_rows
+    )
+    lat = _latent(40, rows, 21)
+    v_src = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
+    tr_attn_tail = _transport() if publish else None
+    ref = attn_tail(
+        lat,
+        w_uv,
+        w_o,
+        rows,
+        v_scratch=v_src,
+        config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
+        hop_transport=tr_attn_tail,
+    )
+    torch.cuda.synchronize()
+    for trace in (False, True, True):
+        # The first stamped run compiles that build (host time beyond the
+        # held stream); the parked time is asserted on the second.
+        tr = torch.zeros(128 * TRACE_POINTS, dtype=torch.int64, device="cuda") if trace else None
+        v = torch.zeros_like(v_src)
+        got = torch.empty_like(ref)
+        tr_oproj_stage = _transport() if publish else None
+        torch.cuda.synchronize()
+        _hold_stream()
+        ps = _slow_produce(v_src, v, delay_ns)
+        stream_oproj_stage(
+            v,
+            w_o,
+            got,
+            use_pdl=True,
+            early_rows=early,
+            grid=128,
+            trigger_dependents=True,
+            stage=stage.stage,
+            when=stage.when,
+            stage_rows=stage.rows,
+            prefetch_policy=stage.prefetch_policy,
+            copy_policy=stage.copy_policy,
+            hop_transport=tr_oproj_stage,
+            trace=tr,
+        )
+        torch.cuda.synchronize()
+        assert torch.equal(ref, got), (arm, publish, trace)
+        if publish:
+            assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words), (arm, trace)
+        if not trace:
+            continue
+        st = tr.view(128, TRACE_POINTS).cpu()
+        p_start, p_end = (int(x) for x in ps.cpu())
+        assert p_end - p_start >= delay_ns
+        assert int(st[:, 1].min()) >= p_end, arm
+        parked = (st[:, 1] - st[:, 0]).float()
+        print(
+            f"parked o_proj {arm} publish={publish}: resident "
+            f"{float(parked.median()):.0f} ns (median) / "
+            f"{float(parked.min()):.0f} ns (min) before its wait returned"
+        )
+    assert float(parked.median()) >= delay_ns / 2, arm
+
+
+@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
+def test_co_running_kernel_progress(arm):
+    """A long GEMM on another stream occupies SMs while the tail runs; the
+    tail completes with attention_output_tail's result (no wait in oproj_weight_staging depends on
+    co-residency)."""
+    w_uv, w_o = _weights(3, 4)
+    rows = 6
+    lat = _latent(3, rows, 17)
+    ref = attn_tail(lat, w_uv, w_o, rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE)
+    a = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
+    side = torch.cuda.Stream()
+    torch.cuda.synchronize()
+    outs = []
+    for _ in range(3):
+        with torch.cuda.stream(side):
+            c = a @ a
+        outs.append(
+            stage_tail(
+                lat, w_uv, w_o, rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE, stage=ARMS[arm]
+            )
+        )
+    torch.cuda.synchronize()
+    del c
+    for got in outs:
+        assert torch.equal(ref, got)
+
+
+def test_trace_stamps_ordered():
+    w_uv, w_o = _weights(3, 0)
+    lat = _latent(3, 6, 3)
+    tr = torch.zeros(128 * TRACE_POINTS, dtype=torch.int64, device="cuda")
+    out = stage_tail(
+        lat, w_uv, w_o, 6, config=INTO_HOP_SPARSE_DECODE_SCHEDULE, stage=ARMS["l2_landed"], trace=tr
+    )
+    torch.cuda.synchronize()
+    ref = attn_tail(lat, w_uv, w_o, 6, config=INTO_HOP_SPARSE_DECODE_SCHEDULE)
+    assert torch.equal(out, ref)
+    st = tr.view(128, TRACE_POINTS).cpu()
+    assert bool((st > 0).all())
+    for k in range(TRACE_POINTS - 1):
+        assert bool((st[:, k + 1] >= st[:, k]).all()), k
+
+
+def test_stage_rejects_bad_config():
+    w_uv, w_o = _weights(3, 0)
+    v = torch.zeros(6, 2048, dtype=torch.bfloat16, device="cuda")
+    out = torch.empty(6, H, dtype=torch.bfloat16, device="cuda")
+    for bad in (
+        dict(stage="tma"),
+        dict(when="late"),
+        dict(prefetch_policy="evict_first"),
+        dict(copy_policy="evict_last"),
+        dict(when="chain"),
+    ):
+        kwargs = dict(stage="l2", when="landed", prefetch_policy="normal", copy_policy="none")
+        kwargs.update(bad)
+        with pytest.raises(ValueError):
+            stream_oproj_stage(v, w_o, out, use_pdl=True, early_rows=32, grid=128, **kwargs)
+
+
+def test_wuv_geometry():
+    assert wuv_stage_geometry(H, 2048, 128, 32, -1) == (32, 16 * 4096)
+    assert wuv_stage_geometry(H, 2048, 128, 32, 8) == (32, 8 * 4096)
+    # 152 CTAs: 40 or 42 rows per CTA, not covered evenly.
+    assert wuv_stage_geometry(H, 2048, 152, 16, -1) is None
+    assert wuv_stage_geometry(H, 2048, 128, 48, -1) is None
+
+
+def test_the_served_stage_is_valid():
+    assert check_stage_config(DEFAULT_STAGE) is None
+    assert (DEFAULT_STAGE.stage, DEFAULT_STAGE.when, DEFAULT_STAGE.early_rows) == (
+        "smem",
+        "landed",
+        24,
+    )
+    assert check_stage_config(dataclasses.replace(DEFAULT_STAGE, stage="l2", when="chain"))
+
+
+# ---- StagedAttnTail on the real DSA MLA module (attention_output_tail protocol unchanged) -------
+
+
+def _step(tail, mla, latent, tr=None, out_rows=None):
+    from tensorrt_llm._torch.distributed import AllReduceParams
+
+    rows = latent.shape[0]
+    out_rows = rows if out_rows is None else out_rows
+    out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
+    meta = tail_adapter._meta(rows)
+    if tr is not None:
+        tail.request_into_hop(tr, out_rows)
+    assert tail.arm(mla, meta)
+    assert tail.defer_w_uv(mla, meta, latent, out[:rows])
+    return tail.project(mla, out, AllReduceParams(enable_allreduce=False))
+
+
+@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
+@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1)])
+def test_oproj_stage_tail_matches_attn_tail(arm, rows, out_rows):
+    mla = tail_adapter._mla()
+    served_tail = AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE)
+    oproj_stage_tail = wrap_attn_tail(
+        mla, AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE), ARMS[arm]
+    )
+    assert isinstance(oproj_stage_tail, StagedAttnTail)
+    lat = tail_adapter._latent(rows, seed=rows + out_rows)
+    tr_attn_tail, tr_oproj_stage = _transport(cur=0, rank=2), _transport(cur=0, rank=2)
+    ref = _step(served_tail, mla, lat, tr_attn_tail, out_rows)
+    got = _step(oproj_stage_tail, mla, lat, tr_oproj_stage, out_rows)
+    torch.cuda.synchronize()
+    assert torch.equal(ref, got)
+    assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words)
+    assert served_tail.take_into_hop_published() and oproj_stage_tail.take_into_hop_published()
+    assert oproj_stage_tail.staging_calls == 1 and oproj_stage_tail.engaged_calls == 1
+    assert oproj_stage_tail.staging_reason is None
+    # No packet left: the native projection would follow.
+    assert oproj_stage_tail.project(mla, torch.empty(rows, 2048, device="cuda"), None) is None
+
+
+def test_oproj_stage_tail_declines_like_attn_tail():
+    mla = tail_adapter._mla()
+    tail = wrap_attn_tail(
+        mla, AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE), ARMS["l2_landed"]
+    )
+    lat = tail_adapter._latent(6, seed=0)
+    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
+    assert not tail.defer_w_uv(mla, tail_adapter._meta(6), lat, out)
+    assert not tail.arm(mla, tail_adapter._meta(6, num_contexts=1))
+    assert not tail.arm(mla, tail_adapter._meta(9))
+    assert tail.declined_calls == 2 and tail.staging_calls == 0
+    # torch.compile policy: arm declines without touching state.
+    from tensorrt_llm._torch.utils import torch_compiling
+
+    with torch_compiling(True):
+        assert not tail.arm(mla, tail_adapter._meta(6))
+    assert tail.declined_calls == 2
+
+
+def test_wrap_keeps_attn_tail_for_other_schedules():
+    mla = tail_adapter._mla()
+    for cfg in (
+        TAIL_PREFETCH_SCHEDULE,
+        TailConfig(oproj_engine="native"),
+        TailConfig(oproj_engine="fused"),
+        TailConfig(trigger_late=True),
+        TailConfig(),
+        INTO_HOP_SCHEDULE,
+    ):
+        tail = AttnTail("layer3", cfg)
+        assert config_reason(cfg, StageConfig()) is not None
+        assert wrap_attn_tail(mla, tail, StageConfig()) is tail
+    assert wrap_attn_tail(mla, None, StageConfig()) is None
+    served = AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE)
+    assert wrap_attn_tail(mla, served, StageConfig(host="sparse_mla_decode")) is served
+    wrapped = wrap_attn_tail(mla, served, None)
+    assert isinstance(wrapped, StagedAttnTail)
+    assert wrapped.stage == DEFAULT_STAGE
+
+
+def test_geometry_decline_runs_attn_tail():
+    """W_UV host on the 152-CTA attention_output_tail default schedule cannot cover the rows:
+    the layer logs an oproj_weight_staging decline and runs attention_output_tail's own project."""
+    mla = tail_adapter._mla()
+    cfg = TailConfig()
+    assert geometry_reason(mla, cfg, ARMS["wuv"]) is not None
+    assert geometry_reason(mla, INTO_HOP_SPARSE_DECODE_SCHEDULE, ARMS["wuv"]) is None
+    tail = StagedAttnTail("layer3", cfg, ARMS["wuv"])
+    lat = tail_adapter._latent(6, seed=4)
+    got = _step(tail, mla, lat)
+    ref = _step(AttnTail("layer3", cfg), mla, lat)
+    torch.cuda.synchronize()
+    assert torch.equal(ref, got)
+    assert tail.staging_calls == 0 and tail.staging_reason is not None
+
+
+@needs_patch
+def test_adoption_attaches_oproj_stage_tail(monkeypatch):
+    """With ``oproj_weight_staging`` running on its closure (alone and with
+    attention_output_tail_prefetch, as served), the attention adoption wraps the
+    attention_output_tail of a target layer; the MTP layer keeps no tail; without
+    ``oproj_weight_staging`` the plain attention_output_tail ``AttnTail``.
+    """
+    for feats in (
+        "attention_output_tail,post_attention_hop,"
+        "sparse_mla_decode,attention_output_into_hop,"
+        "oproj_weight_staging",
+        "attention_output_tail,post_attention_hop,"
+        "sparse_mla_decode,attention_output_tail_prefetch,"
+        "attention_output_into_hop,oproj_weight_staging",
+    ):
+        plan_testlib.select(dispatch, feats, monkeypatch)
+        mla = tail_adapter._build_adopted_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
+        assert isinstance(mla._attn_tail, StagedAttnTail), feats
+        assert mla._attn_tail.config == INTO_HOP_SPARSE_DECODE_SCHEDULE
+        assert mla._attn_tail.stage == DEFAULT_STAGE
+        mtp = tail_adapter._build_adopted_mla(layer_idx=78, num_hidden_layers=78)
+        assert mtp._attn_tail is None
+    plan_testlib.select(
+        dispatch,
+        "attention_output_tail,post_attention_hop,sparse_mla_decode,"
+        "attention_output_tail_prefetch,attention_output_into_hop",
+        monkeypatch,
+    )
+    mla = tail_adapter._build_adopted_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
+    assert type(mla._attn_tail) is AttnTail
+    assert mla._attn_tail.config == INTO_HOP_SPARSE_DECODE_SCHEDULE
+
+
+HOST_SCHEDULE_PATHS = (
+    "attention_output_tail",
+    "post_attention_hop",
+    "sparse_mla_decode",
+    "attention_output_into_hop",
+)
+
+
+@needs_patch
+def test_oproj_stage_runs_only_with_its_host_schedule():
+    """``oproj_weight_staging`` stages only in ``INTO_HOP_SPARSE_DECODE_SCHEDULE``
+    (attention_output_tail's tail with attention_output_into_hop and post_attention_hop behind
+    sparse_mla_decode): the plan turns it off with any of attention_output_tail,
+    post_attention_hop, sparse_mla_decode, attention_output_into_hop instead of running
+    attention_output_tail under oproj_weight_staging's name. This includes disabling only
+    ``attention_output_into_hop``: the attention_output_tail_prefetch schedule that then runs has
+    the same value as ``INTO_HOP_SPARSE_DECODE_SCHEDULE`` but no into-hop publication."""
+    assert TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE == INTO_HOP_SPARSE_DECODE_SCHEDULE
+    assert dispatch.REQUIRED_FEATURES["oproj_weight_staging"] == set(HOST_SCHEDULE_PATHS)
+    for base in HOST_SCHEDULE_PATHS:
+        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
+        assert not plan.enabled("oproj_weight_staging"), base
+    assert plan_testlib.plan_running(
+        dispatch, HOST_SCHEDULE_PATHS + ("oproj_weight_staging",)
+    ).enabled("oproj_weight_staging")
+    b1 = frozenset(B1_GEN.split(","))
+    for paths in (b1, b1 | {"deferred_input_prep"}):
+        assert plan_testlib.plan_running(dispatch, paths | {"oproj_weight_staging"}).enabled(
+            "oproj_weight_staging"
+        )
+
+
+@needs_patch
+def test_every_runnable_selection_engages_oproj_stage(monkeypatch):
+    """Over every subset of the paths ``make_attn_tail`` reads (attention_output_tail, post_attention_hop,
+    sparse_mla_decode, attention_output_tail_prefetch, attention_output_into_hop) plus
+    ``oproj_weight_staging``: each set a plan can run gives a target layer an
+    attention_output_tail with ``INTO_HOP_SPARSE_DECODE_SCHEDULE`` that ``wrap_attn_tail`` wraps, so
+    ``oproj_weight_staging`` never runs inert."""
+    import itertools
+
+    mla = SimpleNamespace(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
+    others = (
+        "attention_output_tail",
+        "post_attention_hop",
+        "sparse_mla_decode",
+        "attention_output_tail_prefetch",
+        "attention_output_into_hop",
+    )
+    runnable = []
+    for n in range(len(others) + 1):
+        for combo in itertools.combinations(others, n):
+            ids = combo + ("oproj_weight_staging",)
+            if not dispatch.REQUIRED_FEATURES["oproj_weight_staging"] <= set(ids):
+                continue
+            plan_testlib.select(dispatch, ids, monkeypatch)
+            tail = wrap_attn_tail(mla, make_attn_tail(mla))
+            assert isinstance(tail, StagedAttnTail), ids
+            assert tail.config == INTO_HOP_SPARSE_DECODE_SCHEDULE, ids
+            runnable.append(",".join(ids))
+    assert sorted(runnable) == sorted(
+        [
+            "attention_output_tail,post_attention_hop,sparse_mla_decode,attention_output_tail_prefetch,attention_output_into_hop,oproj_weight_staging",
+            "attention_output_tail,post_attention_hop,sparse_mla_decode,attention_output_into_hop,oproj_weight_staging",
+        ]
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py
new file mode 100644
index 000000000000..bcb6ffa7a21c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py
@@ -0,0 +1,92 @@
+# 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.
+"""Run a chosen set of the target's fast paths in a test.
+
+``select(dispatch, paths)`` resolves and installs the plan that runs exactly ``paths``: every
+other fast path is listed in ``TRTLLM_STAIRCASE_GLM_LL_DISABLE``, so the plan is built by the
+same ``resolve`` and ``install`` as the target's at construction. ``paths`` must hold every
+path that one of them extends (``depends_on``), since ``DISABLE`` would turn such a path off;
+``select`` raises otherwise. ``dispatch`` is the selector module the code under test reads
+(some harnesses load the target's modules under a stand-in package).
+
+With ``monkeypatch`` the selector state of the process (the installed plan, its validation
+switch and the ``require`` mode) is restored when the test ends; without it the selection
+stays installed, as a process-level harness wants.
+
+A test that runs its case in a child interpreter passes the paths in ``FAST_PATHS_ENV``; the
+child installs them with ``select_from_env``.
+"""
+
+import os
+from typing import FrozenSet, Iterable, Mapping, Optional, Union
+
+Paths = Union[str, Iterable[str]]
+
+FAST_PATHS_ENV = "GLM_MOE_DSA_TEST_FAST_PATHS"
+
+
+def names(paths: Paths) -> FrozenSet[str]:
+    """The fast-path names of a comma-separated string or an iterable of names."""
+    if isinstance(paths, str):
+        return frozenset(name.strip() for name in paths.split(",") if name.strip())
+    return frozenset(paths)
+
+
+def plan_running(dispatch, paths: Paths, mode=None):
+    """The resolved plan that runs exactly ``paths``.
+
+    Raises:
+        ValueError: if a name is not a fast path, or if ``paths`` lacks a path one of them
+            extends.
+    """
+    wanted = names(paths)
+    unknown = sorted(wanted - set(dispatch.FAST_PATH_NAMES))
+    if unknown:
+        raise ValueError(f"{unknown} are not fast paths of the target")
+    bases = {base for path in wanted for base in dispatch.REQUIRED_FEATURES.get(path, ())}
+    missing = sorted(bases - wanted)
+    if missing:
+        raise ValueError(f"{sorted(wanted)} lacks {missing}, which they extend")
+    off = sorted(set(dispatch.FAST_PATH_NAMES) - wanted)
+    environ = {dispatch.DISABLE_ENV: ",".join(off)} if off else {}
+    plan = dispatch.resolve(mode if mode is not None else dispatch.ModelingV2Mode.AUTO, environ)
+    running = frozenset(name for name in dispatch.FAST_PATH_NAMES if plan.enabled(name))
+    assert running == wanted, (sorted(running), sorted(wanted))
+    return plan
+
+
+def select(
+    dispatch,
+    paths: Paths,
+    monkeypatch=None,
+    environ: Optional[Mapping[str, str]] = None,
+    mode=None,
+) -> FrozenSet[str]:
+    """Install the plan that runs exactly ``paths`` and return the installed paths.
+
+    ``environ`` is what the record line reads (the process environment by default), so a
+    variant knob set there selects the numerical variant of a selected path.
+    """
+    if monkeypatch is not None:
+        for name in ("_INSTALLED", "_VALIDATE", "_require"):
+            monkeypatch.setattr(dispatch, name, list(getattr(dispatch, name)))
+    dispatch.install(plan_running(dispatch, paths, mode), environ)
+    return dispatch.installed()
+
+
+def select_from_env(dispatch, environ: Optional[Mapping[str, str]] = None) -> FrozenSet[str]:
+    """Install, for this process, the plan that runs exactly the paths in ``FAST_PATHS_ENV``."""
+    environ = os.environ if environ is None else environ
+    return select(dispatch, environ.get(FAST_PATHS_ENV, ""), environ=environ)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py
new file mode 100644
index 000000000000..62710ef7ae15
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py
@@ -0,0 +1,81 @@
+# 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.
+"""Load ``trtllm::glm_moe_dsa_post_attention_hop_from_partials`` for tests.
+
+A TensorRT-LLM build of this branch registers the op in ``th_common``. On an
+older image (development overlay) the op is compiled from this checkout's
+C++ sources with ``torch.utils.cpp_extension`` instead, unmodified and with
+the same op schema (``-DTLLM_GLM_MOE_DSA_ATTENTION_OUTPUT_INTO_HOP_JIT`` only names the library).
+"""
+
+import hashlib
+import os
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+_REPO = Path(__file__).resolve().parents[9]
+_SOURCES = (
+    _REPO / "cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp",
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu",
+)
+_HEADERS = (
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h",
+    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h",
+)
+_NAME = "into_hop_jit"
+_OP = "glm_moe_dsa_post_attention_hop_from_partials"
+
+
+def _digest() -> str:
+    h = hashlib.sha256()
+    for path in (*_SOURCES, *_HEADERS):
+        h.update(path.read_bytes())
+    return h.hexdigest()[:16]
+
+
+def ensure_op(build_root: str, is_builder: bool, barrier: Callable[[], None]) -> str:
+    """Make the op available on every rank; returns where it came from.
+
+    ``is_builder`` must be true on exactly one rank per shared build root;
+    ``barrier`` synchronizes the ranks after that rank has built it.
+    """
+    import tensorrt_llm  # noqa: F401  (loads th_common and the trtllm namespace)
+
+    if hasattr(torch.ops.trtllm, _OP):
+        barrier()
+        return "builtin"
+    build_dir = Path(build_root) / f"attention_output_into_hop-{_digest()}"
+    library = build_dir / f"{_NAME}.so"
+    if is_builder and not library.exists():
+        from torch.utils.cpp_extension import load
+
+        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
+        build_dir.mkdir(parents=True, exist_ok=True)
+        load(
+            name=_NAME,
+            sources=[str(s) for s in _SOURCES],
+            extra_include_paths=[str(_REPO / "cpp/include"), str(_REPO / "cpp")],
+            extra_cflags=["-O3", "-std=c++17"],
+            extra_cuda_cflags=["-O3", "-std=c++17", "--expt-relaxed-constexpr", "-Xptxas=-v"],
+            build_directory=str(build_dir),
+            is_python_module=False,
+            verbose=True,
+        )
+    barrier()
+    if not hasattr(torch.ops.trtllm, _OP):
+        torch.ops.load_library(str(library))
+    return str(library)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py
new file mode 100644
index 000000000000..dae079e8d365
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py
@@ -0,0 +1,230 @@
+# 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.
+"""Single-process checks of post_attention_hop (no collective, no GPU kernel).
+
+The eight-rank correctness checks need two 4-GPU nodes of one NVLink domain
+and run as collective tests.
+"""
+
+import re
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop import (  # noqa: E501
+    hop,
+)
+from tensorrt_llm.mapping import Mapping
+
+_HEADER = (
+    Path(__file__).resolve().parents[9] / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h"
+)
+
+
+def _header_constant(name: str) -> int:
+    match = re.search(rf"constexpr int {name} = (\d+);", _HEADER.read_text())
+    assert match, name
+    return int(match.group(1))
+
+
+def test_constants_match_kernel_header():
+    assert hop.HIDDEN_SIZE == _header_constant("kHiddenDim")
+    assert hop.NUM_RANKS == _header_constant("kNumRanks")
+    assert hop.NUM_EXPERTS == _header_constant("kNumExperts")
+    assert hop.MAX_ROWS == _header_constant("kMaxRows")
+    assert hop.NUM_LAMPORT_BUFFERS == _header_constant("kNumLamportBuffers")
+
+
+def test_entry_point_registered():
+    assert (
+        dispatch.get_impl("post_attention_hop", "make_post_attn_hop")
+        == hop.PostAttnHop.get_or_create
+    )
+    assert dispatch.get_impl("post_attention_hop", "engagement_counts") == hop.engagement_counts
+
+
+def test_router_layout_constants_match_kernel():
+    """required_sms mirrors the kernel's launch configuration."""
+    source = _HEADER.with_suffix(".cu").read_text()
+    assert f"constexpr int kSlices = {hop.CLUSTER_CTAS};" in source
+    for rows in range(1, hop.MAX_ROWS + 1):
+        groups = min(rows, hop.ROUTER_MAX_ROW_GROUPS)
+        assert (
+            f"case {rows}: return launchTyped<{rows}, "
+            f"{hop.ROUTER_EXPERT_GROUPS}, {groups}>(params);"
+        ) in source
+
+
+@pytest.mark.parametrize(
+    "rows, router, expected",
+    [
+        (1, False, 8),
+        (6, False, 48),
+        (8, False, 64),
+        (1, True, 32),
+        (2, True, 64),
+        (6, True, 96),
+        (8, True, 96),
+    ],
+)
+def test_required_sms(rows, router, expected):
+    assert hop.required_sms(rows, router) == expected
+
+
+def _mapping():
+    return Mapping(world_size=8, tp_size=8, gpus_per_node=4, rank=0)
+
+
+def test_native_allreduce_must_use_mnnvl():
+    from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce
+    from tensorrt_llm.functional import AllReduceStrategy
+
+    mapping = _mapping()
+    nccl = SimpleNamespace(strategy=AllReduceStrategy.NCCL, mnnvl_allreduce=None)
+    reason = hop.native_allreduce_ineligibility(nccl, mapping)
+    assert reason is not None and "NCCL" in reason
+    auto_single_node = SimpleNamespace(strategy=AllReduceStrategy.AUTO, mnnvl_allreduce=None)
+    assert hop.native_allreduce_ineligibility(auto_single_node, mapping) is not None
+    mnnvl = SimpleNamespace(strategy=AllReduceStrategy.AUTO, mnnvl_allreduce=object())
+    assert hop.native_allreduce_ineligibility(mnnvl, mapping) is None
+    # Call sites that do not pass their all-reduce: the TP group must have a
+    # native MNNVL workspace.
+    workspaces = MNNVLAllReduce.allreduce_mnnvl_workspaces
+    assert mapping not in workspaces
+    assert hop.native_allreduce_ineligibility(None, mapping) is not None
+    workspaces[mapping] = {}
+    try:
+        assert hop.native_allreduce_ineligibility(None, mapping) is None
+    finally:
+        del workspaces[mapping]
+
+
+def test_factory_declines_non_mnnvl_call_site_before_any_collective():
+    """A non-MNNVL call site returns None without creating the workspace."""
+    from tensorrt_llm.functional import AllReduceStrategy
+
+    mapping = _mapping()
+    nccl = SimpleNamespace(strategy=AllReduceStrategy.NCCL, mnnvl_allreduce=None)
+    assert (
+        hop.PostAttnHop.get_or_create(
+            mapping, hop.HIDDEN_SIZE, torch.bfloat16, native_allreduce=nccl
+        )
+        is None
+    )
+    assert mapping not in hop.PostAttnHop._instances
+
+
+def _enabled_hop(min_sm_count):
+    h = hop.PostAttnHop.__new__(hop.PostAttnHop)
+    h.enabled = True
+    h.min_sm_count = min_sm_count
+    return h
+
+
+def test_call_eligibility_and_fallback_counters():
+    h = _enabled_hop(152)
+    x6 = torch.zeros(6, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
+    x9 = torch.zeros(9, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
+    gamma = torch.zeros(hop.HIDDEN_SIZE, dtype=torch.bfloat16)
+    router = torch.zeros(hop.NUM_EXPERTS, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
+    hop.reset_engagement_counts()
+    assert h.can_run(x6, x6, gamma, router)
+    assert h.can_run(x6[:1], x6[:1], gamma)
+    assert not h.can_run(x9, x9, gamma, router)
+    assert not h.can_run(x6.half(), x6.half(), gamma)
+    assert not h.can_run(x6, None, gamma)
+    assert not h.can_run(x6, x6, gamma, router.float())
+    counts = hop.engagement_counts()
+    assert counts["native_fallback"] == 4
+    assert counts[f"native_fallback:rows outside [1, {hop.MAX_ROWS}]"] == 1
+    assert counts["native_fallback:dtype"] == 1
+    assert counts["native_fallback:no residual"] == 1
+    assert counts["native_fallback:router weight"] == 1
+    assert "engaged" not in counts
+    # A group whose smallest GPU has 64 SMs: the router layout needs 96 SMs
+    # from 3 rows on; the router-less layout fits up to 8 rows.
+    small = _enabled_hop(64)
+    assert small.call_ineligibility(x6[:2], x6[:2], gamma, router) is None
+    assert small.call_ineligibility(x6[:3], x6[:3], gamma, router) == "SM count"
+    x8 = torch.zeros(8, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
+    assert small.call_ineligibility(x8, x8, gamma) is None
+    hop.reset_engagement_counts()
+    assert hop.engagement_counts() == {}
+
+
+@pytest.mark.parametrize(
+    "kwargs, hidden, dtype, expected",
+    [
+        (dict(world_size=8, tp_size=8, gpus_per_node=4), 6144, torch.bfloat16, None),
+        (dict(world_size=4, tp_size=4, gpus_per_node=4), 6144, torch.bfloat16, "tp_size"),
+        (
+            dict(world_size=8, tp_size=8, gpus_per_node=4, enable_attention_dp=True),
+            6144,
+            torch.bfloat16,
+            "attention DP",
+        ),
+        (dict(world_size=8, tp_size=8, gpus_per_node=4), 7168, torch.bfloat16, "hidden size"),
+        (dict(world_size=8, tp_size=8, gpus_per_node=4), 6144, torch.float16, "dtype"),
+    ],
+)
+def test_static_ineligibility(kwargs, hidden, dtype, expected):
+    reason = hop.static_ineligibility(Mapping(rank=0, **kwargs), hidden, dtype)
+    if expected is None:
+        assert reason is None
+    else:
+        assert reason is not None and expected in reason
+
+
+def test_disabled_hop_never_runs():
+    disabled = hop.PostAttnHop.__new__(hop.PostAttnHop)
+    disabled.enabled = False
+    x = torch.zeros(6, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
+    assert not disabled.can_run(x, x, torch.zeros(hop.HIDDEN_SIZE, dtype=torch.bfloat16))
+
+
+def test_hop_call_sites_keep_dense_layers_native():
+    """Only PRE_MOE_FUSION sites (MoE and MTP layers) call the hop.
+
+    The dense layers' native chain (plain all-reduce in o_proj, then a separate
+    add+RMSNorm) normalizes the FP32 sum, which the hop does not reproduce, so
+    the target must not route ``forward_mlp`` through the hop, and the hop's adoption
+    must not touch o_proj's all-reduce.
+    """
+    import inspect
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
+
+    source = Path(layers.__file__).read_text()
+    assert source.count("self._post_attn_hop(\n") == 2
+    for method in (layers.GlmDecoderLayer.forward_MoE, layers.GlmMtpLayer.forward):
+        assert "self._post_attn_hop(\n" in inspect.getsource(method)
+    assert "post_attn_hop" not in inspect.getsource(layers.GlmDecoderLayer.forward_mlp)
+    adoption = inspect.getsource(layers.GlmDecoderLayer._adopt_pre_moe)
+    assert "enable_allreduce" not in adoption
+    assert "o_proj" not in adoption
+
+
+def test_adoption_passes_the_native_allreduce():
+    """The hop serves a layer only if that layer's native AllReduce is MNNVL."""
+    import inspect
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
+
+    adoption = inspect.getsource(layers.GlmDecoderLayer._adopt_pre_moe)
+    assert '"make_post_attn_hop")' in adoption
+    assert "native_allreduce=self.allreduce" in adoption
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py
new file mode 100644
index 000000000000..5a6d75c3428a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py
@@ -0,0 +1,276 @@
+# 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.
+"""The into-hop entry on one GPU, over a simulated transport (attention_output_into_hop).
+
+The hop from partials (``trtllm::glm_moe_dsa_post_attention_hop_from_partials``) reads every rank's
+partial from the local (unicast) view of the post_attention_hop workspace only, so on one GPU
+a plain device buffer can stand in for it: the test writes all eight ranks'
+sanitized partials into slot ``[row][rank]`` of the current Lamport buffer
+(what the eight producers' multicast stores leave there), sets the protocol
+state, and runs the hop as rank 0 of 8. This needs no MNNVL and runs under
+every compute-sanitizer tool, the transport included. Checks:
+
+* ``residual_out`` equals the rank-order FP32 sum rounded to BF16 plus the
+  residual, exactly; ``normed`` and the router logits agree with references
+  of the same math within one BF16 ulp / FP32 summation-order tolerance;
+* every ``early_poll`` mode and ``router_load`` point gives bit-identical
+  outputs;
+* the protocol: the previous call's buffer (``clearBytes`` of it) is reset
+  to the sentinel, the current buffer is left as published, the state
+  advances to the next buffer with this call's byte count;
+* rows 1, 6, 8, router on/off; CUDA-graph capture and replay with fresh
+  partials.
+
+The partial values are synthetic (seeded); the real transport needs eight
+ranks.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _HERE)
+import into_hop_jit  # noqa: E402
+
+H = 6144
+RANKS = 8
+MAX_ROWS = 8
+NUM_BUFFERS = 3
+BUFFER_BYTES = MAX_ROWS * RANKS * H * 2
+WORKSPACE_BYTES = BUFFER_BYTES * NUM_BUFFERS
+SENTINEL = -(2**31)  # 0x80000000 as int32
+EPS = 1e-5
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+        reason="the attention_output_into_hop hop targets SM100/SM103",
+    ),
+    pytest.mark.threadleak(enabled=False),
+]
+
+_OP = {}
+
+
+def _op():
+    if "op" not in _OP:
+        into_hop_jit.ensure_op(
+            os.environ.get(
+                "GLM_MOE_DSA_ATTENTION_OUTPUT_INTO_HOP_JIT_ROOT",
+                "/tmp/attention_output_into_hop_jit",
+            ),
+            True,
+            lambda: None,
+        )
+        _OP["op"] = torch.ops.trtllm.glm_moe_dsa_post_attention_hop_from_partials
+    return _OP["op"]
+
+
+def _sanitized(x):
+    bits = x.contiguous().view(torch.int16).clone()
+    bits[bits == -32768] = 0  # bf16 -0.0 -> +0.0
+    return bits.view(torch.bfloat16)
+
+
+class Transport:
+    """A device buffer laid out like the post_attention_hop workspace, plus its state."""
+
+    def __init__(self):
+        self.uc = torch.full((WORKSPACE_BYTES // 4,), SENTINEL, dtype=torch.int32, device="cuda")
+        self.flags = torch.zeros(4, dtype=torch.int32, device="cuda")
+
+    def buffers(self):
+        return self.uc.view(torch.bfloat16).view(NUM_BUFFERS, MAX_ROWS, RANKS, H)
+
+    def publish(self, partials, cur, prev_rows):
+        """State as after the previous hop (prev_rows rows) and the eight
+        producers of this call; the previous buffer holds stale data."""
+        prev = (cur + NUM_BUFFERS - 1) % NUM_BUFFERS
+        self.uc.fill_(SENTINEL)
+        bufs = self.buffers()
+        bufs[prev, :prev_rows].fill_(1.5)  # stale words to be reset
+        for s, part in enumerate(partials):
+            bufs[cur, : part.shape[0], s] = _sanitized(part)
+        self.flags.copy_(torch.tensor([cur, prev_rows * RANKS * H * 2, 0, 0], dtype=torch.int32))
+
+
+def _inputs(rows, seed, router):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    partials = [
+        (torch.randn(rows, H, generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
+        for _ in range(RANKS)
+    ]
+    residual = torch.randn(rows, H, generator=gen, device="cuda").to(torch.bfloat16)
+    gamma = (1 + 0.1 * torch.randn(H, generator=gen, device="cuda")).to(torch.bfloat16)
+    weight = (
+        (torch.randn(256, H, generator=gen, device="cuda") * 0.02).to(torch.bfloat16).contiguous()
+        if router
+        else None
+    )
+    return partials, residual, gamma, weight
+
+
+def _run(tr, residual, gamma, weight, early_poll=False, router_load=0):
+    normed, res_out, logits = _op()(
+        residual, gamma, EPS, weight, tr.uc, tr.flags, 0, RANKS, early_poll, router_load
+    )
+    return normed, res_out, (logits if weight is not None else None)
+
+
+def _reference(partials, residual, gamma, weight, normed_k):
+    acc = torch.zeros_like(partials[0], dtype=torch.float32)
+    for part in partials:  # rank order, FP32
+        acc = acc + _sanitized(part).float()
+    res_out = acc.to(torch.bfloat16) + residual  # bf16 add, as __hadd
+    y = res_out.float()
+    sq = (res_out * res_out).float().sum(-1, keepdim=True)
+    normed = (y * torch.rsqrt(sq / H + EPS) * gamma.float()).to(torch.bfloat16)
+    logits = normed_k.float() @ weight.float().t() if weight is not None else None
+    return normed, res_out, logits
+
+
+@pytest.mark.parametrize("rows", [1, 6, 8])
+@pytest.mark.parametrize("router", [True, False])
+def test_hop_math_variants_and_protocol(rows, router):
+    partials, residual, gamma, weight = _inputs(rows, 100 + rows, router)
+    tr = Transport()
+    outs = []
+    for early_poll in (False, True):
+        for router_load in (0, 1, 2) if router else (0,):
+            cur, prev_rows = (early_poll + 2 * router_load) % 3, 8 - rows + 1
+            tr.publish(partials, cur, prev_rows)
+            out = _run(tr, residual, gamma, weight, early_poll, router_load)
+            torch.cuda.synchronize()
+            bufs = tr.buffers()
+            prev = (cur + NUM_BUFFERS - 1) % NUM_BUFFERS
+            # The previous call's rows are reset to the sentinel ...
+            assert bool((tr.uc.view(NUM_BUFFERS, -1)[prev] == SENTINEL).all())
+            # ... the published buffer is left as written ...
+            for s, part in enumerate(partials):
+                assert torch.equal(bufs[cur, :rows, s], _sanitized(part))
+            # ... and the state advances with this call's byte count.
+            want = [(cur + 1) % NUM_BUFFERS, rows * RANKS * H * 2, 0, 0]
+            assert tr.flags.tolist() == want, (tr.flags.tolist(), want)
+            outs.append(out)
+    for out in outs[1:]:
+        for a, b in zip(out, outs[0]):
+            assert (a is None and b is None) or torch.equal(a, b)
+    normed, res_out, logits = outs[0]
+    ref_n, ref_r, ref_l = _reference(partials, residual, gamma, weight, normed)
+    assert torch.equal(res_out, ref_r)
+    torch.testing.assert_close(normed.float(), ref_n.float(), rtol=2**-7, atol=2**-7)
+    if router:
+        torch.testing.assert_close(logits, ref_l, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("router_load", [0, 1, 2])
+def test_hop_graph_replay(router_load):
+    rows = 6
+    partials, residual, gamma, weight = _inputs(rows, 7, True)
+    tr = Transport()
+    tr.publish(partials, 0, rows)
+    _run(tr, residual, gamma, weight, router_load=router_load)
+    torch.cuda.synchronize()
+    tr.publish(partials, 1, rows)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = _run(tr, residual, gamma, weight, router_load=router_load)
+    for rep in range(4):
+        fresh, res_new, _, _ = _inputs(rows, 300 + rep, True)
+        residual.copy_(res_new)
+        tr.publish(fresh, rep % NUM_BUFFERS, rows)
+        graph.replay()
+        torch.cuda.synchronize()
+        got = [t.clone() for t in captured]
+        tr.publish(fresh, rep % NUM_BUFFERS, rows)
+        want = _run(tr, residual, gamma, weight, router_load=router_load)
+        torch.cuda.synchronize()
+        for a, b in zip(got, want):
+            assert torch.equal(a, b)
+    del graph
+
+
+class _FakeHop:
+    """The parts of the post_attention_hop ``PostAttnHop`` that ``IntoHopTransport`` reads.
+
+    ``can_run`` (the post_attention_hop's own, fallback-counting check) must not be
+    called by the into-hop transport; ``call_ineligibility`` is the rank-uniform
+    per-call answer.
+    """
+
+    def __init__(self, uc, reason=None):
+        self.uc_buffer = uc
+        self.reason = reason
+        self.checked = []
+
+    def call_ineligibility(self, hidden_states, residual, norm_weight, router_weight=None):
+        self.checked.append(router_weight is not None)
+        return self.reason
+
+    def can_run(self, *args, **kwargs):
+        raise AssertionError("IntoHopTransport must not count post_attention_hop fallbacks")
+
+
+def _transport(tr, reason=None):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop.from_partials import (  # noqa: E501
+        IntoHopTransport,
+    )
+
+    into_hop = object.__new__(IntoHopTransport)
+    into_hop.hop = _FakeHop(tr.uc, reason)
+    into_hop.rank = 0
+    into_hop.flags = tr.flags
+    into_hop.early_poll = False
+    into_hop.router_load = 1
+    return into_hop
+
+
+def test_transport_eligibility_does_not_count_attn_hop_fallbacks():
+    tr = Transport()
+    residual = torch.zeros(6, H, dtype=torch.bfloat16, device="cuda")
+    gamma = torch.ones(H, dtype=torch.bfloat16, device="cuda")
+    assert _transport(tr).can_run(residual, gamma, None)
+    into_hop = _transport(tr, reason="SM count")
+    assert not into_hop.can_run(residual, gamma, gamma)
+    assert into_hop.hop.checked == [True]
+
+
+@pytest.mark.parametrize("router", [True, False])
+def test_transport_call_counts_and_matches_op(router):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop import (  # noqa: E501
+        from_partials as from_partials_mod,
+    )
+
+    rows = 6
+    partials, residual, gamma, weight = _inputs(rows, 500 + router, router)
+    tr = Transport()
+    tr.publish(partials, 2, rows)
+    want = _run(tr, residual, gamma, weight, router_load=1)
+    torch.cuda.synchronize()
+    from_partials_mod.reset_into_hop_engagement_counts()
+    tr.publish(partials, 2, rows)
+    got = _transport(tr)(residual, gamma, EPS, weight)
+    torch.cuda.synchronize()
+    for a, b in zip(got, want):
+        assert (a is None and b is None) or torch.equal(a, b)
+    counts = dispatch.get_impl("post_attention_hop", "into_hop_engagement_counts")()
+    assert counts == (
+        {"into_hop_engaged": 1, "into_hop_engaged_router": 1} if router else {"into_hop_engaged": 1}
+    )
+    from_partials_mod.reset_into_hop_engagement_counts()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py
new file mode 100644
index 000000000000..b88e35074ae6
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py
@@ -0,0 +1,261 @@
+# 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.
+"""Shared helpers for the post_moe_hop collective tests and microbenchmarks.
+
+Every process of the TP group runs the same code under ``srun --mpi=pmix``;
+helpers here are collective where noted and must be called in the same order
+on every rank.
+"""
+
+import os
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from mpi4py import MPI
+
+GLM_HIDDEN = 6144
+GLM_TOP_K = 8
+GLM_EPS = 1e-5
+
+
+@dataclass
+class MpiEnv:
+    comm: object
+    rank: int
+    world: int
+    gpus_per_node: int
+    mapping: object
+
+    @property
+    def multi_node(self) -> bool:
+        return self.mapping.is_multi_node()
+
+
+_ENV: Optional[MpiEnv] = None
+
+
+def mpi_env() -> MpiEnv:
+    """Per-process TP mapping over every MPI rank (collective-free)."""
+    global _ENV
+    if _ENV is None:
+        from tensorrt_llm.mapping import Mapping
+
+        comm = MPI.COMM_WORLD
+        rank, world = comm.Get_rank(), comm.Get_size()
+        gpus_per_node = int(
+            os.environ.get("GLM_MOE_DSA_POST_MOE_HOP_GPUS_PER_NODE", torch.cuda.device_count())
+        )
+        torch.cuda.set_device(rank % gpus_per_node)
+        mapping = Mapping(world_size=world, rank=rank, gpus_per_node=gpus_per_node, tp_size=world)
+        _ENV = MpiEnv(comm, rank, world, gpus_per_node, mapping)
+    return _ENV
+
+
+def all_ranks(ok: bool) -> bool:
+    """Collective AND of a per-rank verdict."""
+    return bool(mpi_env().comm.allreduce(int(bool(ok)), op=MPI.MIN))
+
+
+def gather_bf16(tensor: torch.Tensor) -> torch.Tensor:
+    """Collective all-gather of a bf16 tensor into a CPU ``[world, ...]`` tensor."""
+    host = tensor.detach().contiguous().view(torch.int16).cpu().numpy()
+    parts = mpi_env().comm.allgather(host)
+    return torch.stack([torch.from_numpy(p).view(torch.bfloat16) for p in parts])
+
+
+def make_raw_tuple(
+    rows: int,
+    hidden: int = GLM_HIDDEN,
+    top_k: int = GLM_TOP_K,
+    padded_rows: Optional[int] = None,
+    seed: int = 0,
+    scale_dtype: torch.dtype = torch.bfloat16,
+    invalid_fraction: float = 0.1,
+    with_shared: bool = True,
+    empty_rows: Tuple[int, ...] = (),
+):
+    """Synthetic rank-local raw MoE tuple plus replicated residual/gamma.
+
+    The FC2 rows, routing weights and permutation are rank-specific, the
+    residual and norm weight are identical on every rank (TP replication).
+    Rows listed in ``empty_rows`` have no valid expert slot.
+    """
+    env = mpi_env()
+    device = torch.device("cuda")
+    padded_rows = padded_rows or max(128, rows * top_k * 2)
+    local = torch.Generator(device=device).manual_seed(10_000 * seed + env.rank)
+    common = torch.Generator(device=device).manual_seed(20_000 * seed + 17)
+    fc2 = (torch.randn(padded_rows, hidden, device=device, generator=local) * 0.05).bfloat16()
+    perm = torch.randperm(padded_rows, device=device, generator=local)[: rows * top_k]
+    perm = perm.view(rows, top_k).to(torch.int32)
+    invalid = torch.rand(rows, top_k, device=device, generator=local) < invalid_fraction
+    perm = torch.where(invalid, torch.full_like(perm, -1), perm)
+    for row in empty_rows:
+        perm[row] = -1
+    scales = torch.rand(rows, top_k, device=device, generator=local) * 0.6
+    scales = scales.to(scale_dtype)
+    shared = None
+    if with_shared:
+        shared = (torch.randn(rows, hidden, device=device, generator=local) * 0.1).bfloat16()
+    residual = torch.randn(rows, hidden, device=device, generator=common).bfloat16()
+    gamma = (torch.rand(hidden, device=device, generator=common) + 0.5).bfloat16()
+    return fc2, scales, perm, shared, residual, gamma
+
+
+def reference_local_partial(fc2, scales, perm, shared, hidden: int = GLM_HIDDEN) -> torch.Tensor:
+    """Native multi-node local arithmetic: finalize (fp32, k order) -> bf16 -> + shared."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        native_finalize_reference,
+    )
+
+    routed = native_finalize_reference(fc2, scales, perm, hidden)
+    return routed if shared is None else shared + routed
+
+
+def published_partials(hop, rows: int) -> torch.Tensor:
+    """Read back ``[rows, world, H]`` partials of the hop's most recent call.
+
+    The buffer of the last completed call stays intact until the next call
+    clears it, so this reads exactly what every rank published.
+    """
+    ws = hop.workspace
+    torch.cuda.synchronize()
+    epoch = int(ws.state[0].item())
+    last = (epoch + 2) % 3
+    view = ws.uc.view(3, ws.max_rows, ws.world, ws.hidden_size)
+    return view[last, :rows].clone()
+
+
+def lamport_snapshot(hop) -> dict:
+    """Fill state of this rank's local Lamport copy (synchronizes the device).
+
+    At a quiescent point, when every rank's last call completed, the protocol
+    keeps exactly the rows of the last call's buffer filled in every slot and
+    every other element at the sentinel; ``ok`` reports that invariant.
+
+    Returns:
+        ``epoch``, ``prev_rows``, ``filled`` (non-sentinel element count per
+        ``[buffer][row][slot]``, int64 CPU), ``ok``, and ``partials_filled``
+        (non-sentinel row-sum partials per buffer).
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
+        MAX_CHUNKS,
+    )
+
+    ws = hop.workspace
+    torch.cuda.synchronize()
+    epoch, prev_rows = int(ws.state[0].item()), int(ws.state[1].item())
+    view = ws.uc.view(torch.int16).view(3, ws.max_rows, ws.world, ws.hidden_size)
+    filled = (view != -0x8000).sum(dim=-1).cpu()
+    expected = torch.zeros_like(filled)
+    expected[(epoch + 2) % 3, :prev_rows] = ws.hidden_size
+    parts = ws.partials.view(torch.int32).view(3, ws.max_rows * MAX_CHUNKS)
+    partials_filled = (parts != -0x80000000).sum(dim=-1).cpu().tolist()
+    return dict(
+        epoch=epoch,
+        prev_rows=prev_rows,
+        filled=filled,
+        ok=torch.equal(filled, expected),
+        partials_filled=partials_filled,
+    )
+
+
+def emulate_post_reduction(
+    partials: torch.Tensor, residual: torch.Tensor, gamma: torch.Tensor, eps: float
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Exact emulation of the native MNNVL one-shot arithmetic on given partials.
+
+    Args:
+        partials: bf16 ``[R, world, H]`` sanitized rank partials.
+
+    Returns:
+        ``(norm_out, residual_out)`` bf16 ``[R, H]``.
+    """
+    acc = torch.zeros(
+        partials.shape[0], partials.shape[2], dtype=torch.float32, device=partials.device
+    )
+    for r in range(partials.shape[1]):
+        acc = acc + partials[:, r].float()
+    reduced = acc.bfloat16()
+    new_residual = (reduced.float() + residual.float()).bfloat16()
+    x = new_residual.float()
+    square_sum = (x * x).bfloat16().float().sum(-1, keepdim=True)
+    inv_rms = torch.rsqrt(square_sum / x.shape[-1] + eps)
+    norm = ((x * inv_rms) * gamma.float()).bfloat16()
+    return norm, new_residual
+
+
+def sanitize_bf16(x: torch.Tensor) -> torch.Tensor:
+    bits = x.view(torch.int16)
+    return torch.where(bits == -32768, torch.zeros_like(bits), bits).view(torch.bfloat16)
+
+
+def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Elementwise distance in bf16 units in the last place."""
+
+    def ordered(x):
+        v = x.contiguous().view(torch.int16).to(torch.int32)
+        return torch.where(v < 0, -(v + 32768), v)
+
+    return (ordered(a) - ordered(b)).abs()
+
+
+def compare_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict:
+    ulps = bf16_ulp_distance(candidate, reference)
+    diff = (candidate.float() - reference.float()).abs()
+    rel = diff / reference.float().abs().clamp_min(1e-6)
+    return {
+        "max_abs": float(diff.max()),
+        "max_rel": float(rel.max()),
+        "max_ulp": int(ulps.max()),
+        "frac_ulp_gt0": float((ulps > 0).float().mean()),
+        "frac_ulp_gt1": float((ulps > 1).float().mean()),
+        "rms_err": float(diff.pow(2).mean().sqrt()),
+    }
+
+
+_NATIVE_AR = None
+
+
+def native_allreduce():
+    """Cached native ``AllReduce(AUTO, bf16)`` of the TP group (MNNVL across hosts)."""
+    global _NATIVE_AR
+    if _NATIVE_AR is None:
+        from tensorrt_llm._torch.distributed import AllReduce
+        from tensorrt_llm.functional import AllReduceStrategy
+
+        _NATIVE_AR = AllReduce(
+            mapping=mpi_env().mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16
+        )
+    return _NATIVE_AR
+
+
+def native_fused_allreduce(
+    local: torch.Tensor, residual: torch.Tensor, gamma: torch.Tensor, eps: float
+):
+    """The native post-MoE hop: ``AllReduce`` with RESIDUAL_RMS_NORM."""
+    from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams
+
+    return native_allreduce()(
+        local,
+        all_reduce_params=AllReduceParams(
+            fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+            residual=residual,
+            norm_weight=gamma,
+            eps=eps,
+            trigger_completion_at_end=False,
+        ),
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py
new file mode 100644
index 000000000000..a7c518c09f86
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py
@@ -0,0 +1,211 @@
+# 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.
+"""Real-weight native MoE fixture for one GLM-5.2 NVFP4 target layer.
+
+Builds, on every TP rank, the native TRT-LLM pieces that feed the post-MoE
+hop of the served GEN role (TP8/EP1, ``moe_backend="TRTLLM"``):
+
+* ``DeepseekV3Gate`` with the layer's real router weight and correction bias;
+* the native NVFP4 MoE (``create_moe`` -> TRTLLM-gen NVFP4 leaf) with the
+  layer's real 256 routed experts (TP-sliced by the module's own loader);
+* the rank's TP slice of the real BF16 shared expert (evaluated with torch);
+* the real RMSNorm weight of the next layer (``input_layernorm`` of
+  ``layer + 1``, or ``model.norm`` after the last layer).
+
+Weights are real; activations are synthetic (a unit-RMS random hidden state
+scaled by the layer's real ``post_attention_layernorm`` weight, and a random
+residual), identical on every rank as TP requires.
+"""
+
+import json
+import os
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+from utils.llm_data import llm_models_root
+
+CKPT = os.environ.get(
+    "GLM_MOE_DSA_POST_MOE_HOP_GLM52_CKPT",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+NUM_EXPERTS = 256
+TOP_K = 8
+HIDDEN = 6144
+MOE_INTER = 2048
+NUM_LAYERS = 78
+ROUTED_SCALING = 2.5
+EPS = 1e-5
+
+
+def _load(names: List[str]) -> Dict[str, torch.Tensor]:
+    from safetensors import safe_open
+
+    with open(os.path.join(CKPT, "model.safetensors.index.json")) as f:
+        weight_map = json.load(f)["weight_map"]
+    by_shard: Dict[str, List[str]] = {}
+    for name in names:
+        by_shard.setdefault(weight_map[name], []).append(name)
+    out = {}
+    for shard, shard_names in sorted(by_shard.items()):
+        with safe_open(os.path.join(CKPT, shard), framework="pt", device="cpu") as f:
+            for name in shard_names:
+                out[name] = f.get_tensor(name)
+    return out
+
+
+@dataclass
+class RealMoELayer:
+    layer: int
+    gate: torch.nn.Module
+    moe: torch.nn.Module
+    shared_gate: torch.Tensor
+    shared_up: torch.Tensor
+    shared_down: torch.Tensor
+    post_attn_gamma: torch.Tensor
+    next_gamma: torch.Tensor
+    expert_input_scale: float
+
+    def shared_output(self, x: torch.Tensor) -> torch.Tensor:
+        """Rank-local TP partial of the BF16 shared expert (no reduction)."""
+        gate = torch.nn.functional.linear(x, self.shared_gate)
+        up = torch.nn.functional.linear(x, self.shared_up)
+        return torch.nn.functional.linear(torch.nn.functional.silu(gate) * up, self.shared_down)
+
+    def hidden_states(self, rows: int, seed: int) -> torch.Tensor:
+        gen = torch.Generator(device="cuda").manual_seed(seed)
+        x = torch.randn(rows, HIDDEN, device="cuda", generator=gen)
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS)
+        return (x * self.post_attn_gamma.float()).bfloat16()
+
+    @staticmethod
+    def residual(rows: int, seed: int) -> torch.Tensor:
+        gen = torch.Generator(device="cuda").manual_seed(seed + 1)
+        return (torch.randn(rows, HIDDEN, device="cuda", generator=gen) * 4.0).bfloat16()
+
+    def run_moe(self, x: torch.Tensor, router_logits: torch.Tensor, do_finalize: bool):
+        return self.moe.forward(
+            x, router_logits, do_finalize=do_finalize, output_dtype=torch.bfloat16
+        )
+
+
+def build_real_moe_layer(mapping, layer: int, max_num_tokens: int = 64) -> RealMoELayer:
+    """Load layer ``layer`` of the checkpoint into the native modules on this rank."""
+    from transformers.configuration_utils import PretrainedConfig
+
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate
+    from tensorrt_llm._torch.moe.fused_moe import create_moe
+    from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
+
+    prefix = f"model.layers.{layer}"
+    next_norm = (
+        "model.norm.weight"
+        if layer + 1 >= NUM_LAYERS
+        else f"model.layers.{layer + 1}.input_layernorm.weight"
+    )
+    names = [
+        f"{prefix}.mlp.gate.weight",
+        f"{prefix}.mlp.gate.e_score_correction_bias",
+        f"{prefix}.post_attention_layernorm.weight",
+        next_norm,
+    ] + [f"{prefix}.mlp.shared_experts.{p}.weight" for p in ("gate_proj", "up_proj", "down_proj")]
+    suffixes = ("weight", "weight_scale", "weight_scale_2", "input_scale")
+    hf_to_native = {"gate_proj": "w1", "up_proj": "w3", "down_proj": "w2"}
+    for expert in range(NUM_EXPERTS):
+        for proj in hf_to_native:
+            names += [f"{prefix}.mlp.experts.{expert}.{proj}.{s}" for s in suffixes]
+    tensors = _load(names)
+
+    pretrained = PretrainedConfig()
+    pretrained.num_experts = NUM_EXPERTS
+    pretrained.n_routed_experts = NUM_EXPERTS
+    pretrained.hidden_size = HIDDEN
+    pretrained.intermediate_size = MOE_INTER
+    pretrained.moe_intermediate_size = MOE_INTER
+    pretrained.torch_dtype = torch.bfloat16
+    model_config = ModelConfig(
+        pretrained_config=pretrained,
+        mapping=mapping,
+        quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16),
+        moe_backend="TRTLLM",
+        max_num_tokens=max_num_tokens,
+    )
+
+    with torch.device("cuda"):
+        gate = DeepseekV3Gate(
+            HIDDEN,
+            NUM_EXPERTS,
+            top_k=TOP_K,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=ROUTED_SCALING,
+            dtype=torch.bfloat16,
+            fuse_routing_kernel=True,
+            apply_routing=False,
+            moe_backend="TRTLLM",
+        )
+        gate.load_weights(
+            [
+                {
+                    "weight": tensors[f"{prefix}.mlp.gate.weight"],
+                    "e_score_correction_bias": tensors[
+                        f"{prefix}.mlp.gate.e_score_correction_bias"
+                    ],
+                }
+            ]
+        )
+        moe = create_moe(
+            routing_method=gate.routing_method,
+            num_experts=NUM_EXPERTS,
+            hidden_size=HIDDEN,
+            intermediate_size=MOE_INTER,
+            dtype=torch.bfloat16,
+            reduce_results=False,
+            model_config=model_config,
+            layer_idx=layer,
+            allow_backend_degradation=False,
+        )
+        weights = {}
+        for expert in range(NUM_EXPERTS):
+            for proj, native in hf_to_native.items():
+                for s in suffixes:
+                    weights[f"{expert}.{native}.{s}"] = tensors[
+                        f"{prefix}.mlp.experts.{expert}.{proj}.{s}"
+                    ]
+        moe.load_weights([weights])
+        moe.post_load_weights()
+    moe.cuda()
+
+    tp, rank = mapping.tp_size, mapping.tp_rank
+    width = MOE_INTER // tp
+    rows = slice(rank * width, (rank + 1) * width)
+    shared_gate = tensors[f"{prefix}.mlp.shared_experts.gate_proj.weight"][rows].cuda()
+    shared_up = tensors[f"{prefix}.mlp.shared_experts.up_proj.weight"][rows].cuda()
+    shared_down = (
+        tensors[f"{prefix}.mlp.shared_experts.down_proj.weight"][:, rows].contiguous().cuda()
+    )
+    input_scale = float(tensors[f"{prefix}.mlp.experts.0.gate_proj.input_scale"])
+    return RealMoELayer(
+        layer,
+        gate,
+        moe,
+        shared_gate,
+        shared_up,
+        shared_down,
+        tensors[f"{prefix}.post_attention_layernorm.weight"].cuda(),
+        tensors[next_norm].cuda(),
+        input_scale,
+    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
new file mode 100644
index 000000000000..69641ad86883
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
@@ -0,0 +1,796 @@
+# 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.
+"""Collective correctness tests of the post_moe_hop (the post_moe_hop contract).
+
+Run one pytest process per GPU of the TP group, e.g. two 4-GPU hosts:
+
+    srun --mpi=pmix -N2 --ntasks-per-node=4 ... \
+        python3 -m pytest --import-mode=importlib -p no:cacheprovider -x -q \
+        tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
+
+Every test is collective: all ranks run the same tests in the same order and
+each verdict is AND-reduced across ranks before it is asserted, so a failure
+on one rank fails the test everywhere instead of deadlocking the next test.
+Inputs are synthetic raw tuples (rank-specific FC2 rows/weights/permutation,
+replicated residual and norm weight); the real-weight native-runner check is
+in ``test_post_moe_hop_real_weights_mpi.py``.
+
+Oracles:
+* the native finalize arithmetic in torch (``native_finalize_reference``) for
+  the rank-local partial, which must equal what the kernel published;
+* an exact emulation of the native MNNVL one-shot arithmetic applied to the
+  partials every rank published (residual output must be bit-exact);
+* the native TRT-LLM ``AllReduce`` (MNNVL across hosts) with RESIDUAL_RMS_NORM
+  on the same partial (residual output bit-exact, normalized output compared
+  in bf16 ulps and with the existing MNNVL test tolerance rtol=0.05/atol=0.15).
+"""
+
+import os
+import sys
+import time
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import post_moe_hop_mpi_utils as U  # noqa: E402
+
+ROW_CASES = (1, 2, 6, 16)
+NORM_MAX_ULP = 2
+
+
+def _log(env, message):
+    if env.rank == 0:
+        print(f"[post_moe_hop-test] {message}", flush=True)
+
+
+@pytest.fixture(scope="module")
+def env():
+    e = U.mpi_env()
+    if e.world < 2:
+        pytest.skip("needs a TP group of at least two ranks under srun --mpi=pmix")
+    _log(e, f"world={e.world} gpus_per_node={e.gpus_per_node} multi_node={e.multi_node}")
+    return e
+
+
+@pytest.fixture(scope="module", params=("triton", "cluster"))
+def hop(env, request):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    U.native_allreduce()
+    hop = PostMoeHop(
+        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, impl=request.param
+    )
+    assert U.all_ranks(hop.impl == request.param), f"{request.param} kernel unavailable"
+    _log(env, f"--- hop impl={hop.impl}")
+    return hop
+
+
+def _check_call(env, hop, fc2, scales, perm, shared, residual, gamma, out, res_out, published=None):
+    """Return (ok, stats) for one completed call against all three oracles."""
+    ref_local = U.sanitize_bf16(U.reference_local_partial(fc2, scales, perm, shared))
+    expected_slots = U.gather_bf16(ref_local).permute(1, 0, 2).cuda()  # [R, world, H]
+    stats = {}
+    ok = True
+    if published is not None:
+        slots_equal = torch.equal(published.view(torch.int16), expected_slots.view(torch.int16))
+        stats["slots_exact"] = slots_equal
+        ok &= slots_equal
+    emu_norm, emu_res = U.emulate_post_reduction(expected_slots, residual, gamma, U.GLM_EPS)
+    res_exact = torch.equal(res_out.view(torch.int16), emu_res.view(torch.int16))
+    norm_stats = U.compare_stats(out, emu_norm)
+    stats.update(res_exact=res_exact, norm_vs_emulation=norm_stats)
+    ok &= res_exact and norm_stats["max_ulp"] <= NORM_MAX_ULP
+    return ok, stats
+
+
+class _Backend:
+    def __init__(self, declares: bool):
+        self._declares = declares
+
+    def supports_unfinalized_output(self) -> bool:
+        return self._declares
+
+
+class _Routing:
+    def __init__(self, top_k: int):
+        self.experts_per_token = top_k
+
+
+class _Producer:
+    """Stand-in for a ``ConfigurableMoE``: the attributes ``producer_capability`` reads."""
+
+    def __init__(
+        self,
+        declares=True,
+        hidden=U.GLM_HIDDEN,
+        top_k=U.GLM_TOP_K,
+        dtype=torch.bfloat16,
+        reduce_results=False,
+    ):
+        self.backend = _Backend(declares)
+        self.routing_method = _Routing(top_k)
+        self.hidden_size = hidden
+        self.dtype = dtype
+        self.reduce_results = reduce_results
+
+
+def _reset_all(env, hop):
+    assert U.all_ranks(hop.reset())
+    assert U.all_ranks(hop.error_bits() == 0)
+
+
+def test_capability_and_dispatch(env, hop):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+        check_health,
+    )
+    from tensorrt_llm.mapping import Mapping
+
+    assert dispatch.get_impl("post_moe_hop", "post_moe_hop") is PostMoeHop
+    assert dispatch.get_impl("post_moe_hop", "health_check") is check_health
+    assert U.all_ranks(hop.engaged), f"hop not engaged: {hop.reason}"
+    assert U.all_ranks(
+        not hop.can_run(0)
+        and hop.can_run(6)
+        and hop.can_run(hop.max_rows)
+        and not hop.can_run(hop.max_rows + 1)
+    )
+
+    # One vetoing rank must disengage every rank, without a hang.
+    veto = "test veto" if env.rank == env.world - 1 else None
+    vetoed = PostMoeHop(
+        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, local_veto=veto
+    )
+    assert U.all_ranks(not vetoed.engaged and not vetoed.can_run(6))
+
+    # A TP1 mapping never negotiates (no collective) and stays native.
+    single = PostMoeHop(
+        Mapping(world_size=1, rank=0, tp_size=1),
+        U.GLM_HIDDEN,
+        U.GLM_TOP_K,
+        require_multi_node=False,
+    )
+    assert not single.engaged
+
+    # The default domain is multi-node only; the default kernel is the cluster kernel.
+    default = PostMoeHop(env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K)
+    assert U.all_ranks(default.engaged == env.multi_node)
+    assert U.all_ranks(not default.engaged or default.impl == "cluster")
+
+    # A cluster width that does not split H into whole warps falls back to Triton.
+    fallback = PostMoeHop(
+        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, impl="cluster", cluster=5
+    )
+    assert U.all_ranks(fallback.engaged and fallback.impl == "triton")
+    _log(env, f"capability ok; multi_node={env.multi_node}")
+
+
+@pytest.mark.parametrize(
+    "case", ("declares", "no_declaration", "reduces", "top_k_one_rank", "hidden", "dtype")
+)
+def test_producer_contract_decided_before_runner(env, hop, case):
+    """The producer's declared raw-tuple contract is part of the collective decision."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    producer = {
+        "declares": _Producer(),
+        "no_declaration": _Producer(declares=False),
+        "reduces": _Producer(reduce_results=True),
+        # A mismatch on one rank only must disengage every rank.
+        "top_k_one_rank": _Producer(top_k=6 if env.rank == 1 else U.GLM_TOP_K),
+        "hidden": _Producer(hidden=U.GLM_HIDDEN // 2),
+        "dtype": _Producer(dtype=torch.float16),
+    }[case]
+    candidate = PostMoeHop(
+        env.mapping,
+        U.GLM_HIDDEN,
+        U.GLM_TOP_K,
+        require_multi_node=False,
+        impl=hop.impl,
+        raw_producer=producer,
+    )
+    expected = case == "declares"
+    _log(
+        env,
+        f"[{hop.impl}] producer {case}: engaged={candidate.engaged} reason={candidate.reason!r}",
+    )
+    assert U.all_ranks(candidate.engaged == expected and candidate.can_run(6) == expected)
+
+
+@pytest.mark.parametrize("scale_dtype", (torch.bfloat16, torch.float32), ids=("w_bf16", "w_fp32"))
+@pytest.mark.parametrize("with_shared", (True, False), ids=("shared", "no_shared"))
+@pytest.mark.parametrize("rows", ROW_CASES)
+def test_matches_oracles(env, hop, rows, with_shared, scale_dtype):
+    seed = 100 + rows * 4 + 2 * int(with_shared) + int(scale_dtype == torch.float32)
+    fc2, scales, perm, shared, residual, gamma = U.make_raw_tuple(
+        rows, seed=seed, scale_dtype=scale_dtype, with_shared=with_shared
+    )
+    out, res_out = hop(fc2, scales, perm, shared, residual, gamma, U.GLM_EPS)
+    published = U.published_partials(hop, rows)
+    ok, stats = _check_call(
+        env, hop, fc2, scales, perm, shared, residual, gamma, out, res_out, published
+    )
+
+    # Native TRT-LLM hop on the identical rank partial.
+    native_norm, native_res = U.native_fused_allreduce(
+        published[:, env.rank].contiguous(), residual, gamma, U.GLM_EPS
+    )
+    torch.cuda.synchronize()
+    native_res_exact = torch.equal(native_res.view(torch.int16), res_out.view(torch.int16))
+    native_stats = U.compare_stats(out, native_norm)
+    try:
+        torch.testing.assert_close(out, native_norm, rtol=0.05, atol=0.15)
+        torch.testing.assert_close(res_out, native_res, rtol=0.05, atol=0.15)
+        within_tolerance = True
+    except AssertionError:
+        within_tolerance = False
+    ok &= native_res_exact and within_tolerance and native_stats["max_ulp"] <= NORM_MAX_ULP
+    ok &= hop.error_bits() == 0
+    _log(
+        env,
+        f"[{hop.impl}] rows={rows} shared={with_shared} w={scale_dtype}: {stats} "
+        f"native_res_exact={native_res_exact} norm_vs_native={native_stats}",
+    )
+    assert U.all_ranks(ok), f"rank {env.rank}: {stats} native={native_stats}"
+
+
+@pytest.mark.parametrize("with_shared", (True, False), ids=("shared", "no_shared"))
+def test_empty_and_padded_rows(env, hop, with_shared):
+    rows = 6
+    fc2, scales, perm, shared, residual, gamma = U.make_raw_tuple(
+        rows, seed=300 + int(with_shared), with_shared=with_shared, empty_rows=(1, 4)
+    )
+    residual[2].zero_()  # a padded token row
+    if shared is not None:
+        shared[2].zero_()
+    out, res_out = hop(fc2, scales, perm, shared, residual, gamma, U.GLM_EPS)
+    published = U.published_partials(hop, rows)
+    ok, stats = _check_call(
+        env, hop, fc2, scales, perm, shared, residual, gamma, out, res_out, published
+    )
+    if not with_shared:
+        ok &= bool((published[1].view(torch.int16) == 0).all())
+        ok &= bool((published[4].view(torch.int16) == 0).all())
+    ok &= bool(torch.isfinite(out.float()).all())
+    _log(env, f"empty/padded rows shared={with_shared}: {stats}")
+    assert U.all_ranks(ok), stats
+
+
+@pytest.mark.parametrize(
+    "block,warps,pdl", ((512, 2, False), (2048, 8, True), (8192, 16, True), (8192, 32, False))
+)
+def test_launch_configs(env, hop, block, warps, pdl):
+    """Chunked and single-CTA-per-row grids, with and without PDL, share one workspace."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    variant = PostMoeHop(
+        env.mapping,
+        U.GLM_HIDDEN,
+        U.GLM_TOP_K,
+        require_multi_node=False,
+        block=block,
+        num_warps=warps,
+        use_pdl=pdl,
+        impl="triton",
+    )
+    ok = variant.engaged and variant.workspace is hop.workspace
+    for rows in (1, 6):
+        tensors = U.make_raw_tuple(rows, seed=350 + rows + block)
+        out, res_out = variant(*tensors, U.GLM_EPS)
+        call_ok, stats = _check_call(
+            env, variant, *tensors, out, res_out, U.published_partials(variant, rows)
+        )
+        ok &= call_ok
+    ok &= variant.error_bits() == 0
+    _log(env, f"launch config block={block} warps={warps} pdl={pdl}: ok={ok}")
+    assert U.all_ranks(ok), stats
+
+
+def test_narrow_launch_config_disengages(env, hop):
+    """A Triton program whose column vector does not span every warp is refused."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    variant = PostMoeHop(
+        env.mapping,
+        U.GLM_HIDDEN,
+        U.GLM_TOP_K,
+        require_multi_node=False,
+        block=1024,
+        num_warps=8,
+        impl="triton",
+    )
+    _log(env, f"block=1024 num_warps=8: engaged={variant.engaged} reason={variant.reason!r}")
+    assert U.all_ranks(not variant.engaged and "256 * num_warps" in variant.reason)
+
+
+def _run_back_to_back(env, hops, sequence, seed):
+    """Launch every call without a host sync, then check each one."""
+    calls_before = hops[0].workspace.calls()
+    inputs, outputs = [], []
+    for i, rows in enumerate(sequence):
+        tensors = U.make_raw_tuple(
+            rows,
+            seed=seed + i,
+            with_shared=(i % 4 != 3),
+            scale_dtype=torch.float32 if i % 5 == 0 else torch.bfloat16,
+        )
+        inputs.append(tensors)
+        outputs.append(hops[i % len(hops)](*tensors, U.GLM_EPS))
+    torch.cuda.synchronize()
+    ok = hops[0].workspace.calls() - calls_before == len(sequence)
+    failed = []
+    for i, (tensors, (out, res_out)) in enumerate(zip(inputs, outputs)):
+        call_ok, stats = _check_call(env, hops[0], *tensors, out, res_out)
+        ok &= call_ok
+        if not call_ok:
+            failed.append(
+                (i, sequence[i], stats["res_exact"], stats["norm_vs_emulation"]["max_ulp"])
+            )
+    ok &= hops[0].error_bits() == 0
+    return ok, failed
+
+
+def test_rapid_reuse_varying_rows(env, hop):
+    sequence = [6, 1, 16, 3, 6, 6, 2, 16, 5, 6] * 3
+    ok, failed = _run_back_to_back(env, [hop], sequence, seed=400)
+    _log(env, f"[{hop.impl}] rapid reuse over {len(sequence)} calls ok={ok} failed={failed[:6]}")
+    assert U.all_ranks(ok), failed
+
+
+def test_mixed_launch_configs_back_to_back(env, hop):
+    """The fixture's kernel interleaved with chunked and single-program Triton hops."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    variants = [hop] + [
+        PostMoeHop(
+            env.mapping,
+            U.GLM_HIDDEN,
+            U.GLM_TOP_K,
+            require_multi_node=False,
+            block=block,
+            num_warps=warps,
+            use_pdl=pdl,
+            impl="triton",
+        )
+        for block, warps, pdl in ((8192, 16, True), (512, 2, False))
+    ]
+    sequence = [6, 16, 1, 6, 3, 16, 6, 2, 6, 5, 16, 6] * 2
+    ok, failed = _run_back_to_back(env, variants, sequence, seed=450)
+    _log(env, f"mixed configs over {len(sequence)} calls ok={ok} failed={failed[:6]}")
+    assert U.all_ranks(ok), failed
+
+
+def test_cuda_graph_replay_with_native_interleave(env, hop):
+    num_layers, rows = 4, 6
+    static = [U.make_raw_tuple(rows, seed=500 + layer) for layer in range(num_layers)]
+    native_in = torch.randn(rows, U.GLM_HIDDEN, device="cuda").bfloat16()
+    native_residual, native_gamma = static[0][4].clone(), static[0][5]
+
+    def body():
+        results = []
+        native_out = None
+        for layer, tensors in enumerate(static):
+            results.append(hop(*tensors, U.GLM_EPS))
+            if layer == 1:
+                native_out = U.native_fused_allreduce(
+                    native_in, native_residual, native_gamma, U.GLM_EPS
+                )
+        return results, native_out
+
+    body()  # eager warmup: Triton compile and native workspaces
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        graph_results, graph_native = body()
+    torch.cuda.synchronize()
+    env.comm.barrier()
+
+    ok = True
+    calls_before = hop.workspace.calls()
+    for replay in range(20):
+        for layer, tensors in enumerate(static):
+            fresh = U.make_raw_tuple(
+                rows, seed=600 + 10 * replay + layer, padded_rows=tensors[0].shape[0]
+            )
+            for dst, src in zip(tensors, fresh):
+                if dst is not None:
+                    dst.copy_(src)
+        native_in.copy_(torch.randn_like(native_in.float()).bfloat16())
+        graph.replay()
+        torch.cuda.synchronize()
+        for tensors, (out, res_out) in zip(static, graph_results):
+            call_ok, stats = _check_call(env, hop, *tensors, out, res_out)
+            ok &= call_ok
+        gathered = U.gather_bf16(native_in).permute(1, 0, 2).cuda()
+        _, native_res_ref = U.emulate_post_reduction(
+            gathered, native_residual, native_gamma, U.GLM_EPS
+        )
+        ok &= torch.equal(graph_native[1].view(torch.int16), native_res_ref.view(torch.int16))
+    ok &= hop.workspace.calls() - calls_before == 20 * num_layers
+    ok &= hop.error_bits() == 0
+    _log(env, f"[{hop.impl}] graph replay x20 with native interleave ok={ok}")
+    assert U.all_ranks(ok)
+
+
+def test_delayed_rank_stress(env, hop):
+    rows = 6
+    ok = True
+    for it in range(12):
+        victim = it % env.world
+        tensors = U.make_raw_tuple(rows, seed=700 + it)
+        env.comm.barrier()
+        if env.rank == victim:
+            if it % 3 == 0:
+                time.sleep(0.02)  # host-side launch delay
+            else:
+                torch.cuda._sleep(int(2e6 * (1 + it % 4)))  # ~1-4 ms of device delay
+        out, res_out = hop(*tensors, U.GLM_EPS)
+        torch.cuda.synchronize()
+        call_ok, stats = _check_call(env, hop, *tensors, out, res_out)
+        ok &= call_ok
+    ok &= hop.error_bits() == 0
+    _log(env, f"[{hop.impl}] delayed-rank stress ok={ok}")
+    assert U.all_ranks(ok)
+
+
+def _missing_rank_trial(env, hop, impatient, missing, seed):
+    """One call that rank ``missing`` skips; returns this rank's verdict."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+
+    rows, timeout_s = 6, impatient.timeout_ns / 1e9
+    tensors = U.make_raw_tuple(rows, seed=seed)
+    before = U.lamport_snapshot(hop)
+    _log(
+        env,
+        f"[{hop.impl}] Lamport invariant before rank {missing} skips, on every rank "
+        f"(rank, epoch, prev_rows, ok): "
+        f"{env.comm.allgather((env.rank, before['epoch'], before['prev_rows'], before['ok']))}",
+    )
+    assert U.all_ranks(before["ok"]), before
+    cur = before["epoch"] % 3
+    env.comm.barrier()
+    start = time.monotonic()
+    out = res_out = None
+    if env.rank != missing:
+        out, res_out = impatient(*tensors, U.GLM_EPS)
+    torch.cuda.synchronize()
+    elapsed = time.monotonic() - start
+    counters = hop.workspace.counters()
+    bits, polls = counters["error_bits"], counters["timeout_polls"]
+    nan_chunks = slot_filled = None
+    if env.rank == missing:
+        reported = bits == 0 and polls == 0
+    else:
+        # Every program of the call must time out (the skipped rank's slot
+        # stays empty) and therefore write NaN for all of its columns.
+        slot_filled = U.lamport_snapshot(hop)["filled"][cur, :rows, missing].tolist()
+        nan_chunks = [
+            torch.isnan(t.float()).view(rows, U.GLM_HIDDEN // 1024, 1024).all(-1).sum(-1).tolist()
+            for t in (out, res_out)
+        ]
+        reported = (
+            bool(bits & 1)
+            and polls > 0
+            and timeout_s <= elapsed < 30.0
+            and bool(torch.isnan(out.float()).all())
+            and bool(torch.isnan(res_out.float()).all())
+        )
+    stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
+    _log(
+        env,
+        f"[{hop.impl}] rank {missing} skipped; per rank (rank, elapsed_s, error_bits, "
+        f"timeout_polls, ok): {stats}",
+    )
+    if not U.all_ranks(reported):
+        detail = env.comm.allgather((env.rank, nan_chunks, slot_filled))
+        _log(
+            env,
+            f"[{hop.impl}] per rank (rank, all-NaN 1024-column chunks per row of "
+            f"[out, residual_out], filled elements per row of the skipped rank's slot): {detail}",
+        )
+        return False
+
+    try:
+        hop.raise_if_error()
+        raised = False
+    except RuntimeError:
+        raised = True
+    try:
+        check_health(f"missing-rank{missing}-local")
+        local_raised = False
+    except PostMoeHopError:
+        local_raised = True
+    try:
+        check_health(f"missing-rank{missing}-collective", collective=True)
+        group_raised = False
+    except PostMoeHopError:
+        group_raised = True
+    expected = env.rank != missing
+    return raised == expected and local_raised == expected and group_raised
+
+
+def test_bounded_error_and_collective_reset(env, hop):
+    """A rank that never joins makes its peers time out, poison and report, not hang."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+        check_health,
+    )
+
+    impatient = PostMoeHop(
+        env.mapping,
+        U.GLM_HIDDEN,
+        U.GLM_TOP_K,
+        require_multi_node=False,
+        timeout_s=0.2,
+        impl=hop.impl,
+    )
+    assert U.all_ranks(impatient.workspace is hop.workspace)
+    for trial, missing in enumerate(sorted({0, env.world // 2 + 1, env.world - 1})):
+        ok = _missing_rank_trial(env, hop, impatient, missing, seed=800 + trial)
+        assert U.all_ranks(ok), f"trial {trial}: rank {missing} skipped"
+        _reset_all(env, hop)
+
+    rows = 6
+    tensors = U.make_raw_tuple(rows, seed=810)
+    out, res_out = hop(*tensors, U.GLM_EPS)
+    torch.cuda.synchronize()
+    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, U.published_partials(hop, rows))
+    assert U.all_ranks(call_ok and hop.error_bits() == 0), stats
+    health = check_health("after-reset", collective=True)
+    assert U.all_ranks(not health.group_error and health.workspaces == 1)
+
+
+def test_error_poisons_every_rank_until_reset(env, hop):
+    """After an error recorded on one rank, every rank's next outputs are NaN."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
+        STATE_EPOCH,
+        STATE_ERROR_WORDS,
+    )
+
+    victim = env.world - 2
+    torch.cuda.synchronize()
+    if env.rank == victim:
+        # As if the previous call of this rank had timed out.
+        epoch = int(hop.workspace.state[STATE_EPOCH].item())
+        hop.workspace.state[STATE_ERROR_WORDS[(epoch + 2) % 3]] = 1
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    poisoned = True
+    # Three calls cover every Lamport buffer, so the bits must be carried from
+    # each buffer's error word into the next one.
+    for call in range(3):
+        tensors = U.make_raw_tuple(6, seed=850 + call)
+        out, res_out = hop(*tensors, U.GLM_EPS)
+        torch.cuda.synchronize()
+        poisoned &= bool(torch.isnan(out.float()).all()) and bool(
+            torch.isnan(res_out.float()).all()
+        )
+    words = hop.workspace.state.cpu().tolist()
+    carried = all(words[w] == 1 for w in STATE_ERROR_WORDS)
+    _log(
+        env,
+        f"[{hop.impl}] victim rank {victim} poisoned; three calls all NaN on this rank: "
+        f"{poisoned}; per-rank error words: "
+        f"{env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])}",
+    )
+    assert U.all_ranks(
+        poisoned
+        and (carried if env.rank == victim else not any(words[w] for w in STATE_ERROR_WORDS))
+    )
+    with pytest.raises(PostMoeHopError):
+        check_health("poisoned", collective=True)
+    _reset_all(env, hop)
+    tensors = U.make_raw_tuple(6, seed=860)
+    out, res_out = hop(*tensors, U.GLM_EPS)
+    torch.cuda.synchronize()
+    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, U.published_partials(hop, 6))
+    assert U.all_ranks(call_ok), stats
+
+
+def test_permutation_out_of_range_is_skipped_and_reported(env, hop):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopError,
+        check_health,
+    )
+
+    culprit = 2 % env.world
+    rows = 6
+    tensors = U.make_raw_tuple(rows, seed=870)
+    if env.rank == culprit:
+        tensors[2][1, 3] = tensors[0].shape[0]
+        tensors[2][4, 0] = tensors[0].shape[0] + 1000
+    out, res_out = hop(*tensors, U.GLM_EPS)
+    published = U.published_partials(hop, rows)
+    # The oracle skips entries outside [0, N_padded) like the kernels.
+    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, published)
+    bits = hop.error_bits()
+    call_ok &= bits == (4 if env.rank == culprit else 0)
+    _log(env, f"[{hop.impl}] out-of-range permutation on rank {culprit}: {stats} bits={bits}")
+    assert U.all_ranks(call_ok), (stats, bits)
+    with pytest.raises(PostMoeHopError):
+        check_health("perm-range", collective=True)
+    _reset_all(env, hop)
+
+
+def _misaligned(t: torch.Tensor, offset_elems: int) -> torch.Tensor:
+    """Same values in a view whose base is ``offset_elems`` past a fresh allocation."""
+    storage = torch.empty(t.numel() + 16, dtype=t.dtype, device=t.device)
+    view = storage[offset_elems : offset_elems + t.numel()].view(t.shape)
+    view.copy_(t)
+    return view
+
+
+def _transposed(t: torch.Tensor) -> torch.Tensor:
+    """Same values, non-contiguous (column-major) layout."""
+    view = torch.empty(tuple(reversed(t.shape)), dtype=t.dtype, device=t.device).t()
+    view.copy_(t)
+    return view
+
+
+def _padded_rows(t: torch.Tensor, extra_cols: int) -> torch.Tensor:
+    """Same values in a ``[:, :H]`` view of a wider row-major allocation."""
+    storage = torch.zeros((t.shape[0], t.shape[1] + extra_cols), dtype=t.dtype, device=t.device)
+    view = storage[:, : t.shape[1]]
+    view.copy_(t)
+    return view
+
+
+def _needs_copy(t: torch.Tensor, row_major_2d: bool = False) -> bool:
+    """The documented 16-byte vector-access contract of the post_moe_hop kernels."""
+    if t.data_ptr() % 16 or t.stride(-1) != 1:
+        return True
+    if row_major_2d:
+        return (t.stride(0) * t.element_size()) % 16 != 0
+    return not t.is_contiguous()
+
+
+def _relayout(env, tensors):
+    """Rank-dependent layouts; returns (tensors, expected repairs on this rank)."""
+    fc2, scales, perm, shared, residual, gamma = tensors
+    kind = env.rank % 4
+    if kind == 0:
+        fc2 = _misaligned(fc2, 1)  # base % 16 == 2
+        perm = _misaligned(perm, 1)  # base % 16 == 4
+        gamma = _misaligned(gamma, 2)  # base % 16 == 4
+    elif kind == 1:
+        fc2 = _padded_rows(fc2, 4)  # row pitch % 16 == 8
+        scales = _transposed(scales)  # not contiguous (for R > 1)
+        residual = _misaligned(residual, 4)  # base % 16 == 8
+    elif kind == 2:
+        fc2 = _padded_rows(fc2, 8)  # aligned row pitch: read in place
+        if shared is not None:
+            shared = _transposed(shared)  # inner stride != 1 (for R > 1)
+    repairs = int(_needs_copy(fc2, row_major_2d=True)) + sum(
+        int(_needs_copy(t)) for t in (scales, perm, shared, residual, gamma) if t is not None
+    )
+    return (fc2, scales, perm, shared, residual, gamma), repairs
+
+
+def test_unaligned_layouts_are_repaired_locally(env, hop):
+    """Per-rank layout repairs keep the collective and the arithmetic unchanged."""
+    ok = True
+    for i, rows in enumerate((6, 1, 16)):
+        reference = U.make_raw_tuple(rows, seed=880 + i)
+        tensors, expected = _relayout(env, reference)
+        before = hop.workspace.layout_repairs
+        out, res_out = hop(*tensors, U.GLM_EPS)
+        call_ok, stats = _check_call(
+            env, hop, *reference, out, res_out, U.published_partials(hop, rows)
+        )
+        repaired = hop.workspace.layout_repairs - before
+        call_ok &= repaired == expected and (expected > 0) == (
+            env.rank % 4 in (0, 1) or (rows > 1 and env.rank % 4 == 2)
+        )
+        ok &= call_ok
+        _log(
+            env,
+            f"[{hop.impl}] relayout rows={rows}: {stats} repairs={repaired} expected={expected}",
+        )
+
+    # The repaired call inside a CUDA graph, replayed with fresh values.
+    rows = 6
+    static_ref = U.make_raw_tuple(rows, seed=890)
+    static, expected = _relayout(env, static_ref)
+    hop(*static, U.GLM_EPS)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        graph_out, graph_res = hop(*static, U.GLM_EPS)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    for replay in range(4):
+        fresh = U.make_raw_tuple(rows, seed=900 + replay, padded_rows=static_ref[0].shape[0])
+        for dst, src in zip(static, fresh):
+            if dst is not None:
+                dst.copy_(src)
+        graph.replay()
+        torch.cuda.synchronize()
+        call_ok, stats = _check_call(env, hop, *fresh, graph_out, graph_res)
+        ok &= call_ok
+    ok &= hop.error_bits() == 0
+    _log(env, f"[{hop.impl}] relayout graph replay x4 ok={ok}")
+    assert U.all_ranks(ok)
+
+
+@pytest.mark.parametrize(
+    "case",
+    ("fc2_dtype", "scale_shape", "perm_dtype", "residual_shape", "shared_dtype", "gamma_shape"),
+)
+def test_contract_violation_raises_before_any_device_work(env, hop, case):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHopContractError,
+    )
+
+    rows = 6
+    fc2, scales, perm, shared, residual, gamma = U.make_raw_tuple(rows, seed=920)
+    if case == "fc2_dtype":
+        fc2 = fc2.half()
+    elif case == "scale_shape":
+        scales = scales[:, : U.GLM_TOP_K - 1]
+    elif case == "perm_dtype":
+        perm = perm.long()
+    elif case == "residual_shape":
+        residual = residual[:, : U.GLM_HIDDEN // 2]
+    elif case == "shared_dtype":
+        shared = shared.float()
+    else:
+        gamma = gamma[:-8]
+    calls_before = hop.workspace.calls()
+    with pytest.raises(PostMoeHopContractError):
+        hop(fc2, scales, perm, shared, residual, gamma, U.GLM_EPS)
+    assert U.all_ranks(hop.workspace.calls() == calls_before)
+    tensors = U.make_raw_tuple(rows, seed=921)
+    out, res_out = hop(*tensors, U.GLM_EPS)
+    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, U.published_partials(hop, rows))
+    assert U.all_ranks(call_ok and hop.error_bits() == 0), stats
+
+
+def test_health_check_reports_engagement(env, hop):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        check_health,
+    )
+
+    before = check_health("engagement-before", collective=True)
+    for i in range(5):
+        hop(*U.make_raw_tuple(6, seed=940 + i), U.GLM_EPS)
+    after = check_health("engagement-after", collective=True)
+    ok = (
+        after.calls - before.calls == 5
+        and after.error_bits == 0
+        and not after.group_error
+        and after.workspaces == 1
+        and after.calls == hop.workspace.calls()
+    )
+    _log(env, f"[{hop.impl}] health {before} -> {after}")
+    assert U.all_ranks(ok)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py
new file mode 100644
index 000000000000..f2d3e482bf3a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py
@@ -0,0 +1,287 @@
+# 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.
+"""Explicit-state model check of the post_moe_hop Lamport epoch/clear protocol (CPU only).
+
+compute-sanitizer has no race detector for global memory, so the safety of
+the post_moe_hop exchange protocol is checked here on an abstract model, exhaustively
+over every interleaving of its individual memory events for small bounds.
+
+Model of the cross-rank exchange (``kernels.py`` / ``cute_kernel.py``): ``W``
+ranks run the same sequence of calls; call ``n`` has ``rows[n]`` token rows
+and uses Lamport buffer ``n % 3`` (the epoch is the call index; the last
+arriving program advances it, and the next call reads it after the kernel
+boundary, so it is not modelled further). Every rank owns a local copy of the
+three buffers, each ``[row][slot]``; a slot is either the sentinel or the
+value ``(call, publisher)``. Within call ``n`` rank ``r`` performs, in any
+order (the kernel's threads are concurrent):
+
+* ``PUB(r, n, row, d)`` for every row and destination rank ``d``: one
+  multicast store into ``d``'s copy of slot ``[row][r]``. Each store is
+  issued when the call starts and may land at any later time, also after the
+  call ended (no completion guarantee is assumed for remote stores);
+* ``CLR(r, n, row, s)`` for every row of the previous call and every slot:
+  a local store of the sentinel into buffer ``(n - 1) % 3`` of ``r``'s copy;
+* ``POLL(r, n, row)``: enabled once no slot of the row in ``r``'s copy of
+  buffer ``n % 3`` holds the sentinel; it reads every slot.
+
+A call ends when its clears and polls are done; the next call of that rank
+starts only after that (stream order, kernel boundary). Safety: every
+completed poll reads exactly ``(n, s)`` in slot ``s``, every run reaches the
+end without a stuck poll, and every remote store lands. The Triton kernel's
+intra-GPU row-sum exchange is the same scheme with programs as "ranks", one
+shared copy, one program clearing, and a kernel boundary between calls
+(``shared_memory=True``), which is checked separately.
+
+Mutated protocols (two buffers, clearing the current buffer, no clear,
+clearing the current call's row count, polling only one slot) and a timed-out
+poll must be rejected, which shows the checker can find the failures it is
+meant to exclude.
+"""
+
+from collections import deque
+from dataclasses import dataclass, replace
+from typing import Optional, Tuple
+
+import pytest
+
+SENT = None
+
+
+@dataclass(frozen=True)
+class Protocol:
+    buffers: int = 3
+    clear: str = "previous"  # "previous", "current" or "none"
+    clear_rows: str = "previous_call"  # or "current_call"
+    poll_slots: str = "all"  # or "first"
+    shared_memory: bool = False  # Triton row-sum exchange: one copy, one clearing program
+    allow_timeout: bool = False  # a poll may give up once (error path)
+
+
+def _events(proto: Protocol, world: int, rows: Tuple[int, ...], rank: int, n: int):
+    """Local (clear, poll) events of call ``n`` of ``rank``; remote stores are separate."""
+    clears = []
+    if proto.clear != "none" and (not proto.shared_memory or rank == 0):
+        if proto.clear_rows == "previous_call":
+            clear_count = rows[n - 1] if n > 0 else 0
+        else:
+            clear_count = rows[n]
+        clears = [("CLR", r, s) for r in range(clear_count) for s in range(world)]
+    polls = [("POLL", r) for r in range(rows[n])]
+    return frozenset(clears + polls)
+
+
+class _Model:
+    def __init__(self, proto: Protocol, world: int, rows: Tuple[int, ...]):
+        self.proto = proto
+        self.world = world
+        self.rows = rows
+        self.calls = len(rows)
+        self.max_rows = max(rows)
+        self.copies = 1 if proto.shared_memory else world
+
+    def initial(self):
+        mem = tuple(
+            SENT for _ in range(self.copies * self.proto.buffers * self.max_rows * self.world)
+        )
+        # per rank: (call index, remaining local events, call started)
+        ranks = tuple((0, None) for _ in range(self.world))
+        return (ranks, frozenset(), mem, 0)  # (ranks, pending stores, memory, timeouts used)
+
+    def _addr(self, copy: int, buf: int, row: int, slot: int) -> int:
+        return ((copy * self.proto.buffers + buf) * self.max_rows + row) * self.world + slot
+
+    def _cur(self, n: int) -> int:
+        return n % self.proto.buffers
+
+    def _dirty(self, n: int) -> int:
+        if self.proto.clear == "current":
+            return self._cur(n)
+        return (n - 1) % self.proto.buffers
+
+    def successors(self, state):
+        """Yield ``(label, next_state, violation)`` for every enabled event."""
+        ranks, pending, mem, timeouts = state
+        for rank, (n, remaining) in enumerate(ranks):
+            if n >= self.calls:
+                continue
+            if remaining is None:
+                # Shared memory models one kernel's programs: call n starts
+                # only after every program ended call n - 1 (kernel boundary).
+                if self.proto.shared_memory and any(m < n for m, _ in ranks):
+                    continue
+                new_ranks = list(ranks)
+                new_ranks[rank] = (n, _events(self.proto, self.world, self.rows, rank, n))
+                stores = pending | frozenset(
+                    (rank, n, row, dest)
+                    for row in range(self.rows[n])
+                    for dest in ((0,) if self.proto.shared_memory else range(self.world))
+                )
+                yield f"START r{rank} n{n}", (tuple(new_ranks), stores, mem, timeouts), None
+                continue
+            for event in remaining:
+                kind = event[0]
+                if kind == "CLR":
+                    _, row, slot = event
+                    copy = 0 if self.proto.shared_memory else rank
+                    new_mem = list(mem)
+                    new_mem[self._addr(copy, self._dirty(n), row, slot)] = SENT
+                    nxt = self._advance(ranks, rank, n, remaining - {event})
+                    yield (
+                        f"CLR r{rank} n{n} row{row} s{slot}",
+                        (nxt, pending, tuple(new_mem), timeouts),
+                        None,
+                    )
+                else:
+                    _, row = event
+                    copy = 0 if self.proto.shared_memory else rank
+                    slots = range(self.world) if self.proto.poll_slots == "all" else range(1)
+                    values = [mem[self._addr(copy, self._cur(n), row, s)] for s in slots]
+                    timed_out = False
+                    if any(v is SENT for v in values):
+                        if not self.proto.allow_timeout or timeouts:
+                            continue
+                        timed_out = True
+                    violation = None
+                    if not timed_out:
+                        read = [
+                            mem[self._addr(copy, self._cur(n), row, s)] for s in range(self.world)
+                        ]
+                        expected = [(n, s) for s in range(self.world)]
+                        if read != expected:
+                            violation = (
+                                f"r{rank} call {n} row {row} read {read}, expected {expected}"
+                            )
+                    nxt = self._advance(ranks, rank, n, remaining - {event})
+                    label = f"{'TIMEOUT' if timed_out else 'POLL'} r{rank} n{n} row{row}"
+                    yield label, (nxt, pending, mem, timeouts + int(timed_out)), violation
+        for store in pending:
+            publisher, n, row, dest = store
+            new_mem = list(mem)
+            new_mem[self._addr(dest, self._cur(n), row, publisher)] = (n, publisher)
+            yield (
+                f"LAND r{publisher} n{n} row{row} -> r{dest}",
+                (ranks, pending - {store}, tuple(new_mem), timeouts),
+                None,
+            )
+
+    @staticmethod
+    def _advance(ranks, rank, n, remaining):
+        new_ranks = list(ranks)
+        new_ranks[rank] = (n + 1, None) if not remaining else (n, remaining)
+        return tuple(new_ranks)
+
+    def done(self, state) -> bool:
+        ranks, pending, _, _ = state
+        return not pending and all(n >= self.calls for n, _ in ranks)
+
+
+@dataclass
+class CheckResult:
+    states: int
+    violation: Optional[str]
+    trace: Tuple[str, ...]
+
+
+def explore(proto: Protocol, world: int, rows: Tuple[int, ...], max_states: int = 3_000_000):
+    """Breadth-first exploration; returns the first violation with its event trace."""
+    model = _Model(proto, world, rows)
+    start = model.initial()
+    parent = {start: None}
+    queue = deque([start])
+
+    def trace_to(state, last_label):
+        labels = [last_label]
+        while parent[state] is not None:
+            state, label = parent[state]
+            labels.append(label)
+        return tuple(reversed(labels))
+
+    while queue:
+        state = queue.popleft()
+        moved = False
+        for label, nxt, violation in model.successors(state):
+            moved = True
+            if violation is not None:
+                return CheckResult(len(parent), violation, trace_to(state, label))
+            if nxt not in parent:
+                parent[nxt] = (state, label)
+                if len(parent) > max_states:
+                    raise RuntimeError(f"state bound {max_states} exceeded")
+                queue.append(nxt)
+        if not moved and not model.done(state):
+            return CheckResult(
+                len(parent), "stuck: no event enabled before the end", trace_to(state, "<stuck>")
+            )
+    return CheckResult(len(parent), None, ())
+
+
+GOOD = Protocol()
+ROW_PATTERNS = ((1, 1, 1, 1, 1, 1, 1), (2, 1, 2, 2, 1, 1, 2))
+
+
+@pytest.mark.parametrize("rows", ROW_PATTERNS, ids=("rows1", "rows_varying"))
+def test_cross_rank_protocol_is_safe_two_ranks(rows):
+    result = explore(GOOD, world=2, rows=rows)
+    print(
+        f"[post_moe_hop-model] cross-rank W=2 rows={rows}: {result.states} states, "
+        f"violation={result.violation}"
+    )
+    assert result.violation is None, result
+
+
+def test_cross_rank_protocol_is_safe_three_ranks():
+    # Four calls are the fewest that reuse a buffer (call 3 reuses call 0's),
+    # after the clear of call 1 and with calls 1 and 2 in flight.
+    rows = (1, 1, 1, 1)
+    result = explore(GOOD, world=3, rows=rows, max_states=8_000_000)
+    print(
+        f"[post_moe_hop-model] cross-rank W=3 rows={rows}: {result.states} states, "
+        f"violation={result.violation}"
+    )
+    assert result.violation is None, result
+
+
+@pytest.mark.parametrize(
+    "rows", ((1, 1, 1, 1, 1, 1), (2, 1, 2, 1, 2, 2)), ids=("rows1", "rows_varying")
+)
+def test_intra_gpu_row_sum_exchange_is_safe(rows):
+    proto = replace(GOOD, shared_memory=True)
+    result = explore(proto, world=3, rows=rows)
+    print(
+        f"[post_moe_hop-model] row-sum exchange 3 programs rows={rows}: {result.states} states, "
+        f"violation={result.violation}"
+    )
+    assert result.violation is None, result
+
+
+MUTANTS = {
+    "two_buffers": (replace(GOOD, buffers=2), 2, (1, 1, 1, 1, 1)),
+    "clear_current_buffer": (replace(GOOD, clear="current"), 2, (1, 1, 1, 1)),
+    "no_clear": (replace(GOOD, clear="none"), 2, (1, 1, 1, 1, 1)),
+    "clear_current_row_count": (replace(GOOD, clear_rows="current_call"), 2, (2, 1, 1, 2, 2)),
+    "poll_first_slot_only": (replace(GOOD, poll_slots="first"), 2, (1, 1, 1)),
+    "timeout_without_reset": (replace(GOOD, allow_timeout=True), 2, (1, 1, 1, 1, 1)),
+}
+
+
+@pytest.mark.parametrize("name", sorted(MUTANTS))
+def test_mutated_protocol_is_rejected(name):
+    proto, world, rows = MUTANTS[name]
+    result = explore(proto, world=world, rows=rows)
+    print(
+        f"[post_moe_hop-model] mutant {name}: {result.states} states, violation={result.violation}; "
+        f"trace tail {result.trace[-6:]}"
+    )
+    assert result.violation is not None, f"mutant {name} was not rejected"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py
new file mode 100644
index 000000000000..ccb4b23d64a1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py
@@ -0,0 +1,218 @@
+# 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.
+"""post_moe_hop fed by the native NVFP4 runner with real GLM-5.2 weights (collective).
+
+Fixture: real weights of target layers (router, 256 routed NVFP4 experts,
+BF16 shared expert, next-layer RMSNorm weight) loaded on every TP rank through
+the native modules; synthetic activations (see ``post_moe_hop_real_fixture``).
+
+For each case the same input goes through
+
+* native chain: runner ``do_finalize=True`` -> ``shared + routed`` -> native
+  ``AllReduce`` RESIDUAL_RMS_NORM (MNNVL one-shot across hosts), i.e. the
+  served GEN multi-node path;
+* candidate chain: runner ``do_finalize=False`` -> ``PostMoeHop``;
+
+and must agree: the runner's own finalize equals the torch finalize oracle on
+its raw tuple, the residual outputs are bit-exact, the normalized outputs are
+within the error budget frozen in ``test_post_moe_hop_post_moe_mpi.py``. A CUDA-graph
+case replays the captured candidate chain with fresh inputs.
+
+Set ``GLM_MOE_DSA_POST_MOE_HOP_REAL_LAYERS`` (comma-separated layer ids, default ``10,77``).
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import post_moe_hop_mpi_utils as U  # noqa: E402
+
+LAYERS = tuple(
+    int(x) for x in os.environ.get("GLM_MOE_DSA_POST_MOE_HOP_REAL_LAYERS", "10,77").split(",")
+)
+NORM_MAX_ULP = 2
+
+
+def _log(env, message):
+    if env.rank == 0:
+        print(f"[post_moe_hop-real] {message}", flush=True)
+
+
+@pytest.fixture(scope="module")
+def env():
+    e = U.mpi_env()
+    if e.world < 2:
+        pytest.skip("needs a TP group of at least two ranks under srun --mpi=pmix")
+    return e
+
+
+@pytest.fixture(scope="module", params=("triton", "cluster"))
+def hop(env, request):
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+    )
+
+    U.native_allreduce()
+    hop = PostMoeHop(
+        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, impl=request.param
+    )
+    assert U.all_ranks(hop.impl == request.param), f"{request.param} kernel unavailable"
+    return hop
+
+
+@pytest.fixture(scope="module", params=LAYERS, ids=lambda layer: f"layer{layer}")
+def real_layer(env, request):
+    from post_moe_hop_real_fixture import build_real_moe_layer
+
+    layer = build_real_moe_layer(env.mapping, request.param)
+    torch.cuda.synchronize()
+    env.comm.barrier()
+    yield layer
+    layer.moe.destroy()
+
+
+def test_native_producer_declares_raw_contract(env, hop, real_layer):
+    """The native NVFP4 MoE module passes the pre-runner producer check."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
+        PostMoeHop,
+        producer_capability,
+    )
+
+    declared, reason = producer_capability(
+        real_layer.moe, U.GLM_HIDDEN, U.GLM_TOP_K, torch.bfloat16
+    )
+    bound = PostMoeHop(
+        env.mapping,
+        U.GLM_HIDDEN,
+        U.GLM_TOP_K,
+        require_multi_node=False,
+        impl=hop.impl,
+        raw_producer=real_layer.moe,
+    )
+    backend = type(getattr(real_layer.moe, "backend", real_layer.moe)).__name__
+    _log(
+        env,
+        f"[{hop.impl}] layer {real_layer.layer} producer {type(real_layer.moe).__name__}/"
+        f"{backend}: declared={declared} {reason!r} engaged={bound.engaged}",
+    )
+    assert U.all_ranks(declared and bound.engaged and bound.workspace is hop.workspace)
+
+
+def _native_chain(layer, x, logits, shared, residual):
+    routed = layer.run_moe(x, logits, do_finalize=True)
+    local = shared + routed
+    norm, new_residual = U.native_fused_allreduce(local, residual, layer.next_gamma, U.GLM_EPS)
+    return routed, norm, new_residual
+
+
+def _candidate_chain(layer, hop, x, logits, shared, residual):
+    fc2, scales, perm = layer.run_moe(x, logits, do_finalize=False)
+    out, new_residual = hop(fc2, scales, perm, shared, residual, layer.next_gamma, U.GLM_EPS)
+    return (fc2, scales, perm), out, new_residual
+
+
+@pytest.mark.parametrize("rows", (1, 6, 16))
+def test_real_runner_native_vs_post_moe_hop(env, hop, real_layer, rows):
+    with torch.inference_mode():
+        x = real_layer.hidden_states(rows, seed=1000 + rows)
+        residual = real_layer.residual(rows, seed=2000 + rows)
+        logits = real_layer.gate(x)
+        shared = real_layer.shared_output(x)
+        routed, native_norm, native_res = _native_chain(real_layer, x, logits, shared, residual)
+        raw, out, new_residual = _candidate_chain(real_layer, hop, x, logits, shared, residual)
+        torch.cuda.synchronize()
+
+    fc2, scales, perm = raw
+    layout = {
+        "fc2_base_mod16": fc2.data_ptr() % 16,
+        "fc2_strides": tuple(fc2.stride()),
+        "scales_contig_base_mod16": (scales.is_contiguous(), scales.data_ptr() % 16),
+        "perm_contig_base_mod16": (perm.is_contiguous(), perm.data_ptr() % 16),
+    }
+    finalize_oracle = U.reference_local_partial(fc2, scales, perm, None)
+    finalize_ulps = U.bf16_ulp_distance(finalize_oracle, routed)
+    res_exact = torch.equal(new_residual.view(torch.int16), native_res.view(torch.int16))
+    norm_stats = U.compare_stats(out, native_norm)
+    res_stats = U.compare_stats(new_residual, native_res)
+    try:
+        torch.testing.assert_close(out, native_norm, rtol=0.05, atol=0.15)
+        torch.testing.assert_close(new_residual, native_res, rtol=0.05, atol=0.15)
+        within_tolerance = True
+    except AssertionError:
+        within_tolerance = False
+    distinct = int(perm[perm >= 0].numel())
+    fixture = {
+        "layer": real_layer.layer,
+        "rows": rows,
+        "scale_dtype": str(scales.dtype),
+        "fc2_shape": list(fc2.shape),
+        "valid_slots": distinct,
+        "finalize_max_ulp": int(finalize_ulps.max()),
+        "finalize_frac_ulp_gt0": float((finalize_ulps > 0).float().mean()),
+    }
+    _log(
+        env,
+        f"[{hop.impl}] {fixture} raw_layout={layout} res_exact={res_exact} "
+        f"res_vs_native={res_stats} norm_vs_native={norm_stats} tol={within_tolerance}",
+    )
+    # The native producer needs no layout repair.
+    ok = (
+        int(finalize_ulps.max()) == 0
+        and res_exact
+        and within_tolerance
+        and norm_stats["max_ulp"] <= NORM_MAX_ULP
+        and hop.error_bits() == 0
+        and hop.workspace.layout_repairs == 0
+    )
+    assert U.all_ranks(ok), f"rank {env.rank}: {fixture} {res_stats} {norm_stats}"
+
+
+def test_real_runner_graph_replay(env, hop, real_layer):
+    rows = 6
+    with torch.inference_mode():
+        x = real_layer.hidden_states(rows, seed=3000)
+        residual = real_layer.residual(rows, seed=3001)
+        logits = real_layer.gate(x)
+        shared = real_layer.shared_output(x)
+        _candidate_chain(real_layer, hop, x, logits, shared, residual)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            _, graph_out, graph_res = _candidate_chain(real_layer, hop, x, logits, shared, residual)
+        torch.cuda.synchronize()
+        env.comm.barrier()
+        ok = True
+        worst = {}
+        for replay in range(8):
+            x.copy_(real_layer.hidden_states(rows, seed=3100 + replay))
+            residual.copy_(real_layer.residual(rows, seed=3200 + replay))
+            logits.copy_(real_layer.gate(x))
+            shared.copy_(real_layer.shared_output(x))
+            graph.replay()
+            torch.cuda.synchronize()
+            _, native_norm, native_res = _native_chain(real_layer, x, logits, shared, residual)
+            torch.cuda.synchronize()
+            res_exact = torch.equal(graph_res.view(torch.int16), native_res.view(torch.int16))
+            stats = U.compare_stats(graph_out, native_norm)
+            call_ok = res_exact and stats["max_ulp"] <= NORM_MAX_ULP
+            ok &= call_ok
+            if not call_ok and not worst:
+                worst = {"replay": replay, "res_exact": res_exact, **stats}
+        ok &= hop.error_bits() == 0
+    _log(env, f"[{hop.impl}] layer {real_layer.layer} graph replay x8 ok={ok} {worst}")
+    assert U.all_ranks(ok), worst
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py
new file mode 100644
index 000000000000..12e7dfa1b62d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py
@@ -0,0 +1,22 @@
+# 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.
+"""Make the shared rope_kv_append test helpers importable under ``--import-mode=importlib``."""
+
+import sys
+from pathlib import Path
+
+_HERE = str(Path(__file__).resolve().parent)
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py
new file mode 100644
index 000000000000..4463d6a6953b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py
@@ -0,0 +1,194 @@
+# 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.
+"""Test-only delayed metadata producer launched with programmatic dependent launch.
+
+``delayed_write(srcs, dsts, delay_cycles)`` enqueues one kernel on the current
+stream (capturable into a CUDA graph) that
+
+1. executes ``griddepcontrol.launch_dependents`` first, so a following kernel
+   launched with programmatic stream serialization may start right away;
+2. waits for its own predecessor (``griddepcontrol.wait``);
+3. spins for ``delay_cycles`` SM cycles;
+4. copies every ``srcs[i]`` into ``dsts[i]``.
+
+It stands for an in-graph metadata writer (for example a fused DSA metadata
+kernel) that is PDL-chained to the next kernel: a consumer is correct only if
+it reads ``dsts`` after its own ``griddepcontrol.wait``.
+
+``early_read(srcs, dsts)`` is the sensitivity control of that stress: a
+PDL-launched kernel that copies ``srcs`` into ``dsts`` *before* its own
+``griddepcontrol.wait`` (then waits and releases its dependents), i.e. a
+consumer that violates the rule. Placed right after ``delayed_write`` it must
+observe the stale values, which shows the stress opens the window.
+"""
+
+import functools
+from typing import List
+
+import torch
+
+_CPP_DECL = (
+    "void delayed_write(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts, "
+    "int64_t delay_cycles);\n"
+    "void early_read(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts);"
+)
+
+_CUDA_SRC = r"""
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace
+{
+constexpr int kMaxCopies = 8;
+
+struct Copies
+{
+    uint32_t const* src[kMaxCopies];
+    uint32_t* dst[kMaxCopies];
+    int64_t words[kMaxCopies];
+    int32_t n;
+};
+
+__global__ void delayedWriterKernel(Copies c, int64_t delayCycles)
+{
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+#endif
+    if (threadIdx.x == 0)
+    {
+        long long const start = clock64();
+        while (clock64() - start < delayCycles)
+        {
+        }
+    }
+    __syncthreads();
+    for (int32_t i = 0; i < c.n; ++i)
+    {
+        for (int64_t w = threadIdx.x; w < c.words[i]; w += blockDim.x)
+        {
+            c.dst[i][w] = c.src[i][w];
+        }
+    }
+}
+
+__global__ void earlyReaderKernel(Copies c)
+{
+    for (int32_t i = 0; i < c.n; ++i)
+    {
+        for (int64_t w = threadIdx.x; w < c.words[i]; w += blockDim.x)
+        {
+            c.dst[i][w] = c.src[i][w];
+        }
+    }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+#endif
+}
+
+Copies makeCopies(std::vector<torch::Tensor> const& srcs, std::vector<torch::Tensor> const& dsts)
+{
+    TORCH_CHECK(srcs.size() == dsts.size() && srcs.size() <= kMaxCopies, "up to 8 (src, dst) pairs");
+    Copies c{};
+    c.n = static_cast<int32_t>(srcs.size());
+    for (size_t i = 0; i < srcs.size(); ++i)
+    {
+        TORCH_CHECK(srcs[i].is_cuda() && dsts[i].is_cuda() && srcs[i].is_contiguous() && dsts[i].is_contiguous(),
+            "contiguous CUDA tensors expected");
+        TORCH_CHECK(srcs[i].nbytes() == dsts[i].nbytes() && srcs[i].nbytes() % 4 == 0, "equal sizes, 4-byte words");
+        c.src[i] = static_cast<uint32_t const*>(srcs[i].data_ptr());
+        c.dst[i] = static_cast<uint32_t*>(dsts[i].data_ptr());
+        c.words[i] = static_cast<int64_t>(srcs[i].nbytes() / 4);
+    }
+    return c;
+}
+
+template <typename Kernel, typename... Args>
+void launchPdl(Kernel kernel, Args... args)
+{
+    cudaLaunchConfig_t cfg{};
+    cfg.gridDim = dim3(1);
+    cfg.blockDim = dim3(256);
+    cfg.dynamicSmemBytes = 0;
+    cfg.stream = at::cuda::getCurrentCUDAStream();
+    cudaLaunchAttribute attr[1];
+    attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attr[0].val.programmaticStreamSerializationAllowed = 1;
+    cfg.attrs = attr;
+    cfg.numAttrs = 1;
+    C10_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kernel, args...));
+}
+} // namespace
+
+void delayed_write(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts, int64_t delay_cycles)
+{
+    launchPdl(delayedWriterKernel, makeCopies(srcs, dsts), delay_cycles);
+}
+
+void early_read(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts)
+{
+    launchPdl(earlyReaderKernel, makeCopies(srcs, dsts));
+}
+"""
+
+
+@functools.lru_cache(maxsize=1)
+def _module():
+    import os
+
+    from torch.utils.cpp_extension import load_inline
+
+    # Build for the visible GPU only (the image's TORCH_CUDA_ARCH_LIST includes pre-sm_90 targets).
+    major, minor = torch.cuda.get_device_capability()
+    saved = os.environ.get("TORCH_CUDA_ARCH_LIST")
+    os.environ["TORCH_CUDA_ARCH_LIST"] = f"{major}.{minor}"
+    try:
+        return load_inline(
+            name="rope_append_pdl_delayed_producer",
+            cpp_sources=[_CPP_DECL],
+            cuda_sources=[_CUDA_SRC],
+            functions=["delayed_write", "early_read"],
+            extra_cuda_cflags=["-O2"],
+            verbose=False,
+        )
+    finally:
+        if saved is None:
+            os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
+        else:
+            os.environ["TORCH_CUDA_ARCH_LIST"] = saved
+
+
+def available() -> bool:
+    """Whether the producer can be built (needs nvcc and a compiler at test time)."""
+    try:
+        _module()
+    except (OSError, RuntimeError, ImportError) as err:  # no toolchain in this environment
+        print(f"pdl_producer unavailable: {err}")
+        return False
+    return True
+
+
+def delayed_write(srcs: List[torch.Tensor], dsts: List[torch.Tensor], delay_cycles: int) -> None:
+    _module().delayed_write(list(srcs), list(dsts), int(delay_cycles))
+
+
+def early_read(srcs: List[torch.Tensor], dsts: List[torch.Tensor]) -> None:
+    _module().early_read(list(srcs), list(dsts))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py
new file mode 100644
index 000000000000..ea47ff28bd87
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py
@@ -0,0 +1,624 @@
+# 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.
+"""Native-state environment shared by the rope_kv_append tests and microbenchmark.
+
+It allocates a real ``DSACacheManager`` (FP8 MLA latent pool + FP8 indexer
+K-cache pool, tokens_per_block 64) at GLM-5.2 TP8 decode geometry and builds
+the per-step tensors the native ops read: cache lengths, KV block offsets,
+indexer slot mappings, the DSA block table and request indices. Both the
+native ops (``mla_rope_generation``, ``indexer_k_cache_scatter_op``,
+``convert_req_index_to_global[_grouped]``, the top-k reuse ``copy_``) and the
+rope_kv_append ops are called on the same state so that results can be compared byte for
+byte.
+"""
+
+import os
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import List, Optional
+
+import torch
+
+import tensorrt_llm
+import tensorrt_llm.bindings
+from tensorrt_llm._torch.attention.backends.interface import RopeParams
+from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _compute_slot_mappings
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm.bindings.executor import KvCacheConfig
+from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
+from tensorrt_llm.mapping import Mapping
+
+NUM_HEADS = 8  # GLM-5.2 64 heads at TP8
+Q_LORA_RANK = 2048
+KV_LORA_RANK = 512
+QK_NOPE = 192
+QK_ROPE = 64
+V_HEAD = 256
+LATENT = KV_LORA_RANK + QK_ROPE
+INDEX_HEAD_DIM = 128
+TOPK = 2048
+TOKENS_PER_BLOCK = 64
+QUANT_MODE_FP8_KV = 128
+ROPE_THETA = 8_000_000.0
+MAX_POSITIONS = 1 << 20
+
+
+def load_dev_library() -> None:
+    """Load a development build of the rope_kv_append ops when ``GLM_MOE_DSA_ROPE_KV_APPEND_DEV_LIB`` points to one."""
+    path = os.environ.get("GLM_MOE_DSA_ROPE_KV_APPEND_DEV_LIB")
+    if path and not hasattr(torch.ops.trtllm, "glm_moe_dsa_topk_index_bridge"):
+        torch.ops.load_library(path)
+
+
+CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1  # tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+PAD = None  # batch entry standing for the CUDA-graph padding dummy request
+PAD_DRAFT_LEN = 5  # served MTP depth; the padding dummy is created per draft length
+
+
+@dataclass
+class Step:
+    """One generation step of a batch of requests (``rows`` query tokens each).
+
+    Request-major like the native metadata: request ``i`` owns token rows
+    ``[i * rows, (i + 1) * rows)`` and row ``i`` of the block offsets, the
+    block table and the indexer block offsets.
+    """
+
+    request_ids: List[int]
+    ctx_lens: List[int]  # cached tokens before this step, per batch entry
+    rows: int
+    is_pad: List[bool]  # entries that are the CUDA-graph padding dummy
+    kv_lens_cuda: torch.Tensor  # [N] int32, cache length after the append
+    kv_lens_cpu: torch.Tensor
+    prompt_lens_cpu: torch.Tensor
+    block_offsets: torch.Tensor  # [num_pools, N, 2, max_blocks]
+    block_table: torch.Tensor  # [N, max_blocks] pool slots, -1 padding
+    req_idx: torch.Tensor  # [N * rows] int32 request index of each token
+    slot_fp8: torch.Tensor  # [N * rows] int64
+    slot_scale: torch.Tensor  # [N * rows] int64
+    idx_block_offsets: torch.Tensor  # [N, max_blocks] indexer K-cache pages
+
+    @property
+    def num_seqs(self) -> int:
+        return len(self.request_ids)
+
+    @property
+    def num_tokens(self) -> int:
+        return self.num_seqs * self.rows
+
+    @property
+    def request_id(self) -> int:
+        return self.request_ids[0]
+
+    @property
+    def ctx_len(self) -> int:
+        return self.ctx_lens[0]
+
+    def kv_lens(self) -> List[int]:
+        return [c + self.rows for c in self.ctx_lens]
+
+    def slot_mappings(self, kv_lens):
+        """Native ``_compute_slot_mappings`` for each request's ``rows`` tokens ending at ``kv_lens``."""
+        if isinstance(kv_lens, int):
+            kv_lens = [kv_lens] * self.num_seqs
+        dev = self.idx_block_offsets.device
+        positions = torch.cat(
+            [torch.arange(k - self.rows, k, dtype=torch.int64, device=dev) for k in kv_lens]
+        )
+        req_ind = torch.arange(self.num_seqs, dtype=torch.int64, device=dev).repeat_interleave(
+            self.rows
+        )
+        return _compute_slot_mappings(
+            positions,
+            self.idx_block_offsets,
+            req_ind,
+            INDEX_HEAD_DIM,
+            TOKENS_PER_BLOCK,
+            128,
+            data_bytes_per_token=INDEX_HEAD_DIM,
+        )
+
+    def set_kv_lens(self, kv_lens: List[int]) -> None:
+        """Advance the step in place (static addresses, as across CUDA-graph replays)."""
+        self.kv_lens_cuda.copy_(torch.tensor(kv_lens, dtype=torch.int32))
+        self.kv_lens_cpu.copy_(torch.tensor(kv_lens, dtype=torch.int32))
+        fp8, scl = self.slot_mappings(kv_lens)
+        self.slot_fp8.copy_(fp8)
+        self.slot_scale.copy_(scl)
+
+
+class RopeAppendEnv:
+    """Real DSA cache manager plus helpers to call native and rope_kv_append ops."""
+
+    def __init__(
+        self,
+        num_layers: int = 4,
+        max_tokens: int = (1 << 20) + 8192,
+        max_seq_len: int = (1 << 20) + 64,
+        max_batch_size: int = 40,
+    ):
+        self.device = torch.device("cuda")
+        self.num_layers = num_layers
+        self.max_seq_len = max_seq_len
+        mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        self.sparse_config = DeepSeekSparseAttentionConfig(
+            index_n_heads=32, index_head_dim=INDEX_HEAD_DIM, index_topk=TOPK
+        )
+        model_config = ModelConfig(
+            mapping=mapping,
+            sparse_attention_config=self.sparse_config,
+            pretrained_config=SimpleNamespace(rms_norm_eps=1e-5),
+        )
+        self.kvm = DSACacheManager(
+            KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+            num_layers=num_layers,
+            num_kv_heads=1,
+            head_dim=LATENT,
+            tokens_per_block=TOKENS_PER_BLOCK,
+            max_seq_len=max_seq_len,
+            max_batch_size=max_batch_size,
+            mapping=mapping,
+            dtype=tensorrt_llm.bindings.DataType.FP8,
+            sparse_attention_config=self.sparse_config,
+            model_config=model_config,
+        )
+        rope = RopeParams(
+            dim=QK_ROPE, theta=ROPE_THETA, max_positions=max_seq_len, duplicate_data=True
+        )
+        _, self.cos_sin = rope.create_rope_const_params()
+        self.kv_scale_orig_quant = torch.ones(1, dtype=torch.float32, device=self.device)
+        self._next_request = 1000
+        self._live: List = []
+        self._pad = {}
+        self._layer_ids = {}
+
+    # ----------------------------------------------------------------- state
+    def pools(self) -> List[torch.Tensor]:
+        """Every cache byte the ops may write: the MLA pool and the indexer pools."""
+        pools = [self.kvm.get_unique_primary_pool().view(torch.uint8)]
+        for pool in self.kvm.indexer_k_cache_pool_per_layer:
+            if pool is not None:
+                pools.append(pool.view(torch.uint8))
+        return pools
+
+    def snapshot(self) -> List[torch.Tensor]:
+        return [p.clone() for p in self.pools()]
+
+    def restore(self, snap: List[torch.Tensor]) -> None:
+        for p, s in zip(self.pools(), snap):
+            p.copy_(s)
+
+    def _padding_dummy(self, rows: int):
+        """The CUDA-graph padding request, created as ``CUDAGraphRunner._get_or_create_padding_dummy``."""
+        rid = CUDA_GRAPH_DUMMY_REQUEST_ID - PAD_DRAFT_LEN
+        if rid not in self._pad:
+            reqs = self.kvm.add_dummy_requests(
+                [rid], token_nums=None, is_gen=True, max_num_draft_tokens=PAD_DRAFT_LEN
+            )
+            assert reqs, "KV pool exhausted"
+            self._live.extend(reqs)
+            self._pad[rid] = reqs[0]
+        # Engine view of a generation dummy: past = prompt_len = token_num - 1 = draft_len.
+        assert rows <= 1 + PAD_DRAFT_LEN
+        return rid, PAD_DRAFT_LEN
+
+    def new_step(self, ctx_len: int, rows: int) -> Step:
+        """One generation request whose cache holds ``ctx_len`` tokens plus this step."""
+        return self.new_batch([ctx_len], rows)
+
+    def new_batch(self, ctx_lens: List[Optional[int]], rows: int) -> Step:
+        """A generation batch; ``PAD`` entries are the CUDA-graph padding dummy (repeated as padding is)."""
+        rids, ctxs, is_pad = [], [], []
+        for ctx in ctx_lens:
+            if ctx is PAD:
+                rid, past = self._padding_dummy(rows)
+                rids.append(rid)
+                ctxs.append(past)
+                is_pad.append(True)
+                continue
+            rid = self._next_request
+            self._next_request += 1
+            reqs = self.kvm.add_dummy_requests(
+                [rid], token_nums=[ctx + 1], is_gen=True, max_num_draft_tokens=rows - 1
+            )
+            assert reqs, "KV pool exhausted"
+            self._live.extend(reqs)
+            rids.append(rid)
+            ctxs.append(ctx)
+            is_pad.append(False)
+        n = len(rids)
+        kv_lens = [c + rows for c in ctxs]
+        dev = self.device
+        offsets = torch.zeros(
+            (self.kvm.num_pools, n, 2, self.kvm.max_blocks_per_seq), dtype=torch.int32, device=dev
+        )
+        self.kvm.copy_batch_block_offsets(offsets, rids, 1, 0, n)
+        pool_idx = self.kvm.get_pool_block_indices(n)  # [n, max_blocks], host
+        bt = pool_idx.clone()
+        for i, kv in enumerate(kv_lens):
+            bt[i, (kv + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK :] = -1
+        idx_offsets = (pool_idx * self.kvm.indexer_k_cache_page_scale).to(dev)
+        step = Step(
+            request_ids=rids,
+            ctx_lens=ctxs,
+            rows=rows,
+            is_pad=is_pad,
+            kv_lens_cuda=torch.tensor(kv_lens, dtype=torch.int32, device=dev),
+            kv_lens_cpu=torch.tensor(kv_lens, dtype=torch.int32),
+            prompt_lens_cpu=torch.tensor(ctxs, dtype=torch.int32),
+            block_offsets=offsets,
+            block_table=bt.to(dev),
+            req_idx=torch.arange(n, dtype=torch.int32, device=dev).repeat_interleave(rows),
+            slot_fp8=torch.empty(n * rows, dtype=torch.int64, device=dev),
+            slot_scale=torch.empty(n * rows, dtype=torch.int64, device=dev),
+            idx_block_offsets=idx_offsets,
+        )
+        fp8, scl = step.slot_mappings(kv_lens)
+        step.slot_fp8.copy_(fp8)
+        step.slot_scale.copy_(scl)
+        return step
+
+    def release(self) -> None:
+        for req in self._live:
+            self.kvm.free_resources(req)
+        self._live = []
+        self._pad = {}
+
+    def shutdown(self) -> None:
+        self.release()
+        self.kvm.shutdown()
+
+    # ------------------------------------------------------------ native ops
+    def native_rope(
+        self, step: Step, layer: int, inp: "RopeInputs", out: "RopeOutputs", q_scaling: float = 1.0
+    ) -> None:
+        """``TrtllmAttention._mla_rope_generation_impl`` with the DSA/MLA decode arguments."""
+        torch.ops.trtllm.mla_rope_generation(
+            inp.fused_q,
+            inp.q_pe,
+            inp.latent,
+            self.cos_sin,
+            out.cu_q,
+            out.cu_kv,
+            out.counter,
+            out.bmm1,
+            out.bmm2,
+            out.quant_q,
+            step.kv_lens_cuda,
+            step.kv_lens_cpu,
+            step.prompt_lens_cpu,
+            0,
+            step.block_offsets,
+            self.kvm.kv_cache_pool_pointers,
+            self.kvm.kv_cache_pool_mapping,
+            None,
+            None,
+            self.kv_scale_orig_quant,
+            None,
+            None,
+            [None, None],
+            1,
+            layer,
+            NUM_HEADS,
+            1,
+            LATENT,
+            0,
+            TOKENS_PER_BLOCK,
+            self.max_seq_len,
+            1,
+            QUANT_MODE_FP8_KV,
+            q_scaling,
+            Q_LORA_RANK,
+            KV_LORA_RANK,
+            QK_NOPE,
+            QK_ROPE,
+            V_HEAD,
+            True,
+            None,
+            1e-6,
+            False,
+            False,
+            False,
+            False,
+            None,
+        )
+
+    def native_scatter(self, step: Step, layer: int, inp: "RopeInputs") -> None:
+        """``Indexer._update_k_cache``."""
+        k_cache = self.kvm.get_indexer_k_cache_buffers(layer)
+        torch.ops.trtllm.indexer_k_cache_scatter_op(
+            inp.k_fp8, inp.k_scale, k_cache, step.slot_fp8, step.slot_scale, inp.k_fp8.shape[0]
+        )
+
+    def layer_ids(self, layers: List[int]) -> torch.Tensor:
+        """Device layer-offset tensor of a group, built once (as the native group struct does)."""
+        key = tuple(layers)
+        if key not in self._layer_ids:
+            self._layer_ids[key] = torch.tensor(
+                [self.kvm.get_primary_pool_page_index_params(m)[1] for m in layers],
+                dtype=torch.int32,
+                device=self.device,
+            )
+        return self._layer_ids[key]
+
+    def native_bridge(
+        self,
+        step: Step,
+        topk: torch.Tensor,
+        layers: List[int],
+        shared: Optional[torch.Tensor],
+        grouped: bool,
+    ) -> torch.Tensor:
+        """Top-k reuse copy + ``transform_local_topk_and_prepare_pool_view[_grouped]``."""
+        if shared is not None:
+            shared[0 : topk.shape[0], : topk.shape[1]].copy_(topk)
+        scale, _ = self.kvm.get_primary_pool_page_index_params(layers[0])
+        stride = scale * TOKENS_PER_BLOCK
+        if grouped:
+            layer_ids = self.layer_ids(layers)
+            return torch.ops.trtllm.convert_req_index_to_global_grouped(
+                step.req_idx,
+                step.block_table,
+                topk,
+                TOKENS_PER_BLOCK,
+                topk.shape[1],
+                stride,
+                layer_ids,
+            )
+        _, off = self.kvm.get_primary_pool_page_index_params(layers[0])
+        return torch.ops.trtllm.convert_req_index_to_global(
+            step.req_idx, step.block_table, topk, TOKENS_PER_BLOCK, topk.shape[1], stride, off
+        )
+
+    # --------------------------------------------------------------- rope_kv_append ops
+    def rope_append_rope(
+        self,
+        step: Step,
+        layer: int,
+        inp: "RopeInputs",
+        out: "RopeOutputs",
+        with_scatter: bool,
+        q_scaling: float = 1.0,
+    ) -> None:
+        idx = [None, None, None, None, None]
+        n_idx = 0
+        if with_scatter:
+            idx = [
+                inp.k_fp8,
+                inp.k_scale,
+                self.kvm.get_indexer_k_cache_buffers(layer),
+                step.slot_fp8,
+                step.slot_scale,
+            ]
+            n_idx = inp.k_fp8.shape[0]
+        torch.ops.trtllm.glm_moe_dsa_rope_kv_append(
+            inp.fused_q,
+            inp.q_pe,
+            inp.latent,
+            self.cos_sin,
+            out.cu_q,
+            out.cu_kv,
+            out.counter,
+            out.bmm1,
+            out.bmm2,
+            out.quant_q,
+            step.kv_lens_cuda,
+            step.prompt_lens_cpu,
+            0,
+            step.block_offsets,
+            self.kvm.kv_cache_pool_pointers,
+            self.kvm.kv_cache_pool_mapping,
+            self.kv_scale_orig_quant,
+            layer,
+            NUM_HEADS,
+            TOKENS_PER_BLOCK,
+            self.max_seq_len,
+            QUANT_MODE_FP8_KV,
+            q_scaling,
+            QK_NOPE,
+            QK_ROPE,
+            KV_LORA_RANK,
+            *idx,
+            n_idx,
+        )
+
+    def rope_append_bridge(
+        self,
+        step: Step,
+        topk: torch.Tensor,
+        layers: List[int],
+        shared: Optional[torch.Tensor],
+        grouped: bool,
+        copy_mode: int = 0,
+        next_n: int = 1,
+        num_accepted=None,
+        copy_rows: Optional[int] = None,
+        num_rows: Optional[int] = None,
+    ) -> torch.Tensor:
+        scale, off = self.kvm.get_primary_pool_page_index_params(layers[0])
+        layer_ids = None
+        if grouped:
+            layer_ids = self.layer_ids(layers)
+        rows = topk.shape[0] if num_rows is None else num_rows
+        crows = 0 if shared is None else (rows if copy_rows is None else copy_rows)
+        return torch.ops.trtllm.glm_moe_dsa_topk_index_bridge(
+            topk,
+            rows,
+            step.req_idx,
+            0,
+            step.block_table,
+            TOKENS_PER_BLOCK,
+            scale * TOKENS_PER_BLOCK,
+            layer_ids,
+            off,
+            shared,
+            0,
+            crows,
+            copy_mode,
+            next_n,
+            num_accepted,
+        )
+
+
+@dataclass
+class RopeInputs:
+    fused_q: torch.Tensor  # [R, H, 576] bf16
+    q: torch.Tensor  # [R, H * 256] bf16, q_pe is a strided view of it
+    q_pe: torch.Tensor
+    latent: torch.Tensor  # [R, 576] bf16
+    k_fp8: torch.Tensor  # [R, 128] fp8
+    k_scale: torch.Tensor  # [R, 1] fp32
+
+
+@dataclass
+class RopeOutputs:
+    quant_q: torch.Tensor
+    cu_q: torch.Tensor
+    cu_kv: torch.Tensor
+    counter: torch.Tensor
+    bmm1: torch.Tensor
+    bmm2: torch.Tensor
+
+    @staticmethod
+    def make(rows: int, device, fill: int = 0xA5, num_seqs: int = 1) -> "RopeOutputs":
+        """Outputs pre-filled with sentinels; the cu_seqlens buffers have 3 slack entries."""
+
+        def i32(n: int) -> torch.Tensor:
+            return torch.full((n,), -7, dtype=torch.int32, device=device)
+
+        return RopeOutputs(
+            quant_q=torch.full((rows, NUM_HEADS, LATENT), fill, dtype=torch.uint8, device=device),
+            cu_q=i32(num_seqs + 4),
+            cu_kv=i32(num_seqs + 4),
+            counter=torch.full((1,), 99, dtype=torch.int32, device=device).view(torch.uint32),
+            bmm1=torch.full((2,), -3.0, dtype=torch.float32, device=device),
+            bmm2=torch.full((1,), -3.0, dtype=torch.float32, device=device),
+        )
+
+    def tensors(self):
+        return [
+            self.quant_q,
+            self.cu_q,
+            self.cu_kv,
+            self.counter.view(torch.int32),
+            self.bmm1,
+            self.bmm2,
+        ]
+
+
+def make_inputs(
+    rows: int, device, scale: float = 1.0, seed: int = 0, specials: bool = False
+) -> RopeInputs:
+    g = torch.Generator(device=device).manual_seed(seed)
+    fused_q = (torch.randn((rows, NUM_HEADS, LATENT), generator=g, device=device) * scale).to(
+        torch.bfloat16
+    )
+    q = (
+        torch.randn((rows, NUM_HEADS * (QK_NOPE + QK_ROPE)), generator=g, device=device) * scale
+    ).to(torch.bfloat16)
+    latent = (torch.randn((rows, LATENT), generator=g, device=device) * scale).to(torch.bfloat16)
+    if specials:
+        vals = torch.tensor(
+            [
+                0.0,
+                -0.0,
+                1e-40,
+                -1e-38,
+                448.0,
+                449.0,
+                1e4,
+                -1e4,
+                float("inf"),
+                float("-inf"),
+                float("nan"),
+                2.0**-10,
+                3.0e-3,
+            ],
+            device=device,
+        ).to(torch.bfloat16)
+        for t in (fused_q.view(-1), q.view(-1), latent.view(-1)):
+            pos = torch.randint(0, t.numel(), (vals.numel() * 8,), generator=g, device=device)
+            t[pos] = vals.repeat(8)
+    q_pe = q.view(rows, NUM_HEADS, QK_NOPE + QK_ROPE)[..., QK_NOPE:]
+    k_fp8 = (
+        torch.randint(0, 256, (rows, INDEX_HEAD_DIM), generator=g, device=device, dtype=torch.int32)
+        .to(torch.uint8)
+        .view(torch.float8_e4m3fn)
+    )
+    k_scale = torch.rand((rows, 1), generator=g, device=device, dtype=torch.float32) + 0.01
+    return RopeInputs(fused_q=fused_q, q=q, q_pe=q_pe, latent=latent, k_fp8=k_fp8, k_scale=k_scale)
+
+
+def make_topk(
+    rows: int, kv_len: int, max_blocks: int, device, seed: int = 0, invalid_frac: float = 0.05
+) -> torch.Tensor:
+    """Top-k ids in [0, kv_len) plus -1 padding and out-of-table ids."""
+    g = torch.Generator(device=device).manual_seed(seed)
+    topk = torch.randint(
+        0, max(kv_len, 1), (rows, TOPK), generator=g, device=device, dtype=torch.int32
+    )
+    mask = torch.rand((rows, TOPK), generator=g, device=device) < invalid_frac
+    topk[mask] = -1
+    beyond = torch.rand((rows, TOPK), generator=g, device=device) < invalid_frac / 2
+    topk[beyond] = torch.randint(
+        kv_len,
+        kv_len + 4 * TOKENS_PER_BLOCK,
+        (int(beyond.sum()),),
+        generator=g,
+        device=device,
+        dtype=torch.int32,
+    )
+    far = torch.rand((rows, TOPK), generator=g, device=device) < invalid_frac / 4
+    topk[far] = max_blocks * TOKENS_PER_BLOCK + 5
+    return topk.contiguous()
+
+
+def make_batch_inputs(
+    step: Step, device, scale: float = 1.0, seed: int = 0, specials: bool = False
+) -> RopeInputs:
+    """Inputs for every token of ``step``.
+
+    Repeated padding entries are the same request writing the same cache slots,
+    so they carry identical rows (as padded graph replays feed identical dummy
+    tokens); otherwise the final cache bytes would depend on write order for
+    native and rope_kv_append alike.
+    """
+    inp = make_inputs(step.num_tokens, device, scale=scale, seed=seed, specials=specials)
+    pads = [i for i, p in enumerate(step.is_pad) if p]
+    r = step.rows
+    for i in pads[1:]:
+        src, dst = slice(pads[0] * r, pads[0] * r + r), slice(i * r, i * r + r)
+        for t in (inp.fused_q, inp.q, inp.latent, inp.k_fp8.view(torch.uint8), inp.k_scale):
+            t[dst] = t[src]
+    return inp
+
+
+def make_batch_topk(
+    step: Step,
+    device,
+    seed: int = 0,
+    rows: Optional[int] = None,
+    kv_lens: Optional[List[int]] = None,
+) -> torch.Tensor:
+    """Per-request top-k rows (``rows`` per request, default ``step.rows``) inside each cache length."""
+    rows = step.rows if rows is None else rows
+    kv_lens = step.kv_lens() if kv_lens is None else kv_lens
+    parts = [
+        make_topk(rows, kv, step.block_table.shape[1], device, seed=seed * 131 + i)
+        for i, kv in enumerate(kv_lens)
+    ]
+    return torch.cat(parts).contiguous()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py
new file mode 100644
index 000000000000..fa41b7101dd3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py
@@ -0,0 +1,727 @@
+# 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.
+"""Real DSA layer stack for the rope_kv_append call-site, composition and region tests.
+
+``Stack`` builds real ``MLA`` modules at GLM-5.2 TP8 attention geometry (8
+local heads, q-LoRA 2048, kv-LoRA 512, nope 192, rope 64, v 256, 32 x 128
+indexer, top-k 2048, FP8 KV cache, block 64, heuristic GVR top-k, CuTe DSL
+logits) in a full/shared indexer pattern plus the MTP layer, adopted as the target's
+``GlmMLA``, a real ``DSACacheManager`` filled with FP8 context, and real generation metadata.
+On decode steps, ``run_seg`` drives the target's ``_forward_dsa_attn`` with given projections;
+``run_proj`` drives the target's ``forward_sparse_attn`` from hidden states, so the
+projection producer (native, or absorbed_query/shared_indexer_absorbed_query/draft_absorbed_query
+when they run) is the real one. ``Engagement`` counts native, rope_kv_append and neighbour-package
+entry points and can record each layer's FP8 query/scheduler/scale buffers and global indices.
+"""
+
+import os
+import sys
+
+import torch
+from rope_append_env import load_dev_library  # noqa: F401  (re-exported for the tests)
+from transformers import PretrainedConfig
+
+import tensorrt_llm
+import tensorrt_llm.bindings
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+    GlmMLA,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention_forward import (  # noqa: E501
+    _forward_dsa_attn,
+    forward_sparse_attn,
+)
+from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams
+from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
+from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+from tensorrt_llm._torch.attention.mla import MLA
+from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+from tensorrt_llm.bindings.executor import KvCacheConfig
+from tensorrt_llm.functional import PositionEmbeddingType
+from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+
+GLM_HIDDEN = 6144
+
+HIDDEN = 1024  # run_seg does not exercise the projections
+HEADS = 8
+NOPE, ROPE, V, Q_LORA, KV_LORA = 192, 64, 256, 2048, 512
+PATTERN = ["F", "S", "S", "S"]
+MTP_LAYER = len(PATTERN)
+TOKENS_PER_BLOCK = 64
+CTX_REQ_ID = 9  # a request in its context phase (fallback cases)
+CTX_REQ_TOKENS = 256
+
+
+def set_features(ids: str) -> None:
+    """Install, for this process, the plan that runs exactly the fast paths ``ids`` (comma-separated)."""
+    plan_testlib.select(dispatch, ids)
+
+
+class Engagement:
+    """Counts native and rope_kv_append entry points so a silent native fallback cannot pass a comparison."""
+
+    def __init__(self, capture: bool = False):
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
+            rope_kv_append as rope_append_pkg,
+        )
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+            ops as rope_ops,
+        )
+        from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
+        from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+        from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention
+
+        self.counts = {}
+        self._undo = []
+        self.capture = capture
+        # Per rope call: (quant_q, cu_q, cu_kv, scheduler counter, bmm1, bmm2); per remap: indices.
+        self.rope_outputs = []
+        self.global_indices = []
+        self._wrap(Indexer, "_update_k_cache", "native_scatter")
+        self._wrap(TrtllmAttention, "mla_rope_generation", "native_rope")
+        self._wrap(DSATrtllmAttention, "_remap_topk_to_global", "native_remap")
+        self._wrap_neighbours()
+        self.bridge_checks = 0
+        orig_bridge = rope_ops.topk_bridge
+        counted_bridge = self._counted(orig_bridge, "rope_append_bridge_kernel")
+
+        def checked_bridge(topk, metadata, num_rows, **kwargs):
+            out = counted_bridge(topk, metadata, num_rows, **kwargs)
+            self._check_bridge(topk, metadata, num_rows, out, **kwargs)
+            return out
+
+        rope_ops.topk_bridge = checked_bridge
+        self._undo.append(lambda: setattr(rope_ops, "topk_bridge", orig_bridge))
+        del rope_append_pkg
+        registry = dispatch._REGISTRY["rope_kv_append"]
+        for name in (
+            "mla_rope_generation",
+            "remap_topk_to_global",
+            "run_indexer_k_scatter_native",
+            "check_indexer_k_scatter_consumed",
+        ):
+            orig = registry[name]
+            registry[name] = self._counted(orig, "rope_append_" + name)
+            self._undo.append(lambda n=name, o=orig: registry.__setitem__(n, o))
+        orig_defer = registry["defer_indexer_k_scatter"]
+
+        def counted_defer(*args, **kwargs):
+            pkt = orig_defer(*args, **kwargs)
+            key = (
+                "rope_append_deferred_scatter"
+                if pkt is not None
+                else "rope_append_scatter_not_deferred"
+            )
+            self.counts[key] = self.counts.get(key, 0) + 1
+            return pkt
+
+        registry["defer_indexer_k_scatter"] = counted_defer
+        self._undo.append(lambda: registry.__setitem__("defer_indexer_k_scatter", orig_defer))
+
+    def _counted(self, fn, key):
+        def wrapper(*args, **kwargs):
+            self.counts[key] = self.counts.get(key, 0) + 1
+            if key == "rope_append_mla_rope_generation" and kwargs.get("indexer_k") is not None:
+                self.counts["rope_append_fused_scatter"] = (
+                    self.counts.get("rope_append_fused_scatter", 0) + 1
+                )
+            out = fn(*args, **kwargs)
+            if self.capture and key in ("native_rope", "rope_append_mla_rope_generation"):
+                # (mqa, fused_q, q_pe, latent, metadata, cu_q, cu_kv, counter, bmm1, bmm2, quant_q)
+                md = args[4]
+                n = md.num_seqs + 1
+                cu_q, cu_kv, counter, bmm1, bmm2, quant_q = args[5:11]
+                self.rope_outputs.append(
+                    tuple(t.clone() for t in (quant_q, cu_q[:n], cu_kv[:n], counter, bmm1, bmm2))
+                )
+            if self.capture and key in ("native_remap", "rope_append_remap_topk_to_global"):
+                self.global_indices.append(out.clone())
+            return out
+
+        return wrapper
+
+    def _wrap_neighbours(self):
+        """Count engagement of the neighbour packages present in this tree."""
+
+        def wrap_registry(pkg, name, key, engaged):
+            try:
+                orig = dispatch.get_impl(pkg, name)
+            except (ImportError, KeyError):
+                return
+            registry = dispatch._REGISTRY[pkg]
+
+            def wrapper(*args, **kwargs):
+                out = orig(*args, **kwargs)
+                if engaged(out):
+                    self.counts[key] = self.counts.get(key, 0) + 1
+                return out
+
+            registry[name] = wrapper
+            self._undo.append(lambda: registry.__setitem__(name, orig))
+
+        wrap_registry("absorbed_query", "dsa_proj", "aq_packet", lambda o: o is not None)
+        wrap_registry(
+            "absorbed_query", "shared_dsa_proj", "shared_query_packet", lambda o: o is not None
+        )
+        wrap_registry(
+            "absorbed_query", "mtp_dsa_proj", "draft_query_packet", lambda o: o is not None
+        )
+        wrap_registry("indexer_topk", "indexer_decode", "indexer_topk_decode", lambda o: bool(o))
+        try:
+            fmha = dispatch.get_impl("sparse_mla_decode", "fmha_lib")
+        except (ImportError, KeyError):
+            fmha = None
+        if fmha is not None:
+            self._wrap(fmha, "forward", "sparse_mla_fmha")
+
+    def _check_bridge(
+        self,
+        topk,
+        metadata,
+        num_rows,
+        out,
+        *,
+        layer_offset=None,
+        layer_ids=None,
+        page_index_scale,
+        copy=None,
+    ):
+        """Per-call oracle at the real call site: native remap and native copy of the same rows.
+
+        Independent of which rows the (run-to-run nondeterministic) native
+        top-k selected in this run.
+        """
+        torch.cuda.synchronize()
+        tpb = metadata.kv_cache_manager.tokens_per_block
+        nc = metadata.num_contexts
+        block_table = metadata.block_table[nc : metadata.num_seqs]
+        req = (metadata.req_idx_per_token[metadata.num_ctx_tokens : metadata.num_tokens] - nc)[
+            :num_rows
+        ]
+        src = topk[:num_rows]
+        stride = page_index_scale * tpb
+        if layer_ids is not None:
+            want = torch.ops.trtllm.convert_req_index_to_global_grouped(
+                req.to(torch.int32).contiguous(),
+                block_table,
+                src,
+                tpb,
+                src.shape[1],
+                stride,
+                layer_ids,
+            )
+        else:
+            want = torch.ops.trtllm.convert_req_index_to_global(
+                req.to(torch.int32).contiguous(),
+                block_table,
+                src,
+                tpb,
+                src.shape[1],
+                stride,
+                layer_offset,
+            )
+        assert torch.equal(out, want), (
+            "rope_kv_append remap differs from the native remap of the same rows"
+        )
+        if copy is not None:
+            k = topk.shape[1]
+            got = copy.dst[copy.row_start : copy.row_start + copy.rows, :k]
+            if copy.mode == 0:
+                exp = topk[: copy.rows]
+            else:
+                base = torch.arange(copy.rows, device=topk.device, dtype=torch.long) * copy.next_n
+                if copy.num_accepted is None:
+                    off = torch.full_like(base, copy.next_n - 1)
+                else:
+                    off = (copy.num_accepted[: copy.rows].long() - 1).clamp(0, copy.next_n - 1)
+                exp = topk[base + off]
+            assert torch.equal(got, exp), (
+                f"rope_kv_append top-k copy (mode {copy.mode}) differs from native"
+            )
+        self.bridge_checks += 1
+
+    def _wrap(self, owner, attr, key):
+        orig = getattr(owner, attr)
+        setattr(owner, attr, self._counted(orig, key))
+        self._undo.append(lambda: setattr(owner, attr, orig))
+
+    def close(self):
+        for undo in reversed(self._undo):
+            undo()
+
+
+class Stack:
+    def __init__(self, ctx_len: int, pattern=None, real_proj: bool = False):
+        """``real_proj``: GLM hidden size 6144 and random projection weights for ``run_proj``."""
+        pattern = list(PATTERN if pattern is None else pattern)
+        self.pattern = pattern
+        self.mtp_layer = len(pattern)
+        self.hidden = GLM_HIDDEN if real_proj else HIDDEN
+        dev = torch.device("cuda")
+        self.dev = dev
+        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        self.pretrained = PretrainedConfig(
+            hidden_size=self.hidden,
+            num_attention_heads=HEADS,
+            num_key_value_heads=HEADS,
+            max_position_embeddings=1 << 20,
+            rope_parameters={"rope_theta": 8_000_000, "rope_type": "default"},
+            qk_rope_head_dim=ROPE,
+            qk_nope_head_dim=NOPE,
+            kv_lora_rank=KV_LORA,
+            q_lora_rank=Q_LORA,
+            v_head_dim=V,
+            model_type="glm_moe_dsa",
+            rms_norm_eps=1e-5,
+            num_hidden_layers=len(pattern),
+            index_topk_pattern=pattern,
+            index_share_for_mtp_iteration=True,
+            index_n_heads=32,
+            index_head_dim=128,
+            index_topk=2048,
+        )
+        self.sparse = DeepSeekSparseAttentionConfig(
+            index_n_heads=32,
+            index_head_dim=128,
+            index_topk=2048,
+            enable_heuristic_topk=True,
+            use_cute_dsl_paged_mqa_logits=True,
+            use_cute_dsl_topk=True,
+            skip_indexer_for_short_seqs=False,
+        )
+        self.model_config = ModelConfig(
+            mapping=self.mapping,
+            sparse_attention_config=self.sparse,
+            pretrained_config=self.pretrained,
+            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8),
+        )
+        g = torch.Generator(device=dev).manual_seed(1234)
+        self.layers = []
+        for idx in range(len(pattern) + 1):
+            mla = MLA(
+                hidden_size=self.hidden,
+                num_attention_heads=HEADS,
+                num_key_value_heads=1,
+                qk_nope_head_dim=NOPE,
+                qk_rope_head_dim=ROPE,
+                v_head_dim=V,
+                q_lora_rank=Q_LORA,
+                kv_lora_rank=KV_LORA,
+                predicted_tokens_per_seq=6,
+                max_position_embeddings=1 << 20,
+                bias=False,
+                pos_embd_params=PositionalEmbeddingParams(
+                    type=PositionEmbeddingType.yarn,
+                    rope=RopeParams.from_config(self.pretrained),
+                    is_neox=False,
+                ),
+                layer_idx=idx,
+                dtype=torch.bfloat16,
+                config=self.model_config,
+            )
+            # The model's DeepseekV32Attention adds only its constructor to MLA; the stack builds
+            # the MLA at GLM geometry directly and adopts it as the model's attention.
+            mla.__class__ = DeepseekV32Attention
+            GlmMLA.adopt(mla)
+            for backend in (mla.mha, mla.mqa):
+                if backend is not None and backend.prepended_fmha_lib_classes:
+                    backend.update_quant_config(backend.quant_config)
+            mla = mla.to(dev)
+            with torch.no_grad():
+                if real_proj:
+                    for name, prm in mla.named_parameters():
+                        if not prm.dtype.is_floating_point:
+                            continue
+                        if name.endswith("bias"):
+                            prm.zero_()
+                        elif "norm" in name:
+                            prm.fill_(1.0)
+                        else:
+                            prm.normal_(0.0, 0.02, generator=g)
+                k_nope = torch.randn(HEADS, NOPE, KV_LORA, generator=g, device=dev) * 0.05
+                v = torch.randn(HEADS, V, KV_LORA, generator=g, device=dev) * 0.05
+                mla.k_b_proj_trans.data = k_nope.transpose(1, 2).contiguous().to(torch.bfloat16)
+                mla.v_b_proj.data = v.contiguous().to(torch.bfloat16)
+                if mla.mqa.indexer is not None:
+                    mla.mqa.indexer.to(dev)
+                    for w in (
+                        mla.mqa.indexer.wq_b.weight,
+                        mla.mqa.indexer.wk.weight,
+                        mla.mqa.indexer.weights_proj.weight,
+                    ):
+                        w.normal_(0.0, 0.02, generator=g)
+                    mla.mqa.indexer.post_load_weights()
+            self.layers.append(mla)
+        assert [m.mqa.indexer is not None for m in self.layers] == [p == "F" for p in pattern] + [
+            True
+        ]
+
+        # A second request with a different cache length for the two-request cases.
+        self.ctx_b = ctx_len // 3 + 4099
+        # Block-table width = max_seq_len / 64 pages; keep it even so indexer_topk (whose logits
+        # tiles span two pages) stays inside its envelope at every context length.
+        pages_per_tile = 2
+        max_seq_len = -(-(ctx_len + 1024) // (pages_per_tile * TOKENS_PER_BLOCK)) * (
+            pages_per_tile * TOKENS_PER_BLOCK
+        )
+        self.kvm = DSACacheManager(
+            KvCacheConfig(
+                max_tokens=ctx_len + self.ctx_b + 8192 + CTX_REQ_TOKENS + 256,
+                enable_block_reuse=False,
+            ),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+            num_layers=len(pattern) + 1,
+            num_kv_heads=1,
+            head_dim=KV_LORA + ROPE,
+            tokens_per_block=TOKENS_PER_BLOCK,
+            max_seq_len=max_seq_len,
+            max_batch_size=3,
+            mapping=self.mapping,
+            dtype=tensorrt_llm.bindings.DataType.FP8,
+            sparse_attention_config=self.sparse,
+            model_config=self.model_config,
+            pretrained_config=self.pretrained,
+        )
+        self.ctx_len = ctx_len
+        self.rid = 7
+        self.rids = [7, 8]
+        self.ctxs = [ctx_len, self.ctx_b]
+        self.reqs = self.kvm.add_dummy_requests(
+            self.rids, token_nums=[c + 1 for c in self.ctxs], is_gen=True, max_num_draft_tokens=5
+        )
+        # FP8 context: valid E4M3 latents and indexer keys, positive indexer scales.
+        pool = self.kvm.get_unique_primary_pool()
+        flat = pool.view(torch.uint8).view(-1)
+        chunk = flat.numel() if flat.numel() <= (1 << 31) else (1 << 30)
+        for start in range(0, flat.numel(), chunk):
+            n = min(chunk, flat.numel() - start)
+            flat[start : start + n].copy_(
+                (torch.randn(n, generator=g, device=dev) * 0.5)
+                .to(torch.float8_e4m3fn)
+                .view(torch.uint8)
+            )
+        for ipool in self.kvm.indexer_k_cache_pool_per_layer:
+            if ipool is None:
+                continue
+            blocks = ipool.view(torch.uint8).view(self.kvm.num_blocks, -1)
+            data = TOKENS_PER_BLOCK * 128
+            blocks[:, :data].copy_(
+                (torch.randn(blocks.shape[0], data, generator=g, device=dev))
+                .to(torch.float8_e4m3fn)
+                .view(torch.uint8)
+            )
+            blocks[:, data:].view(torch.float32).copy_(
+                torch.rand(blocks.shape[0], TOKENS_PER_BLOCK, generator=g, device=dev) + 0.05
+            )
+        self.attn_cls = get_attention_backend(
+            "TRTLLM", sparse_params=self.sparse.to_sparse_params(pretrained_config=self.pretrained)
+        )
+
+    def cached(self, cached: int, nreq: int):
+        """Cached lengths of the first ``nreq`` requests; ``cached`` is request 0's."""
+        return [c + (cached - self.ctx_len) for c in self.ctxs[:nreq]]
+
+    def metadata(
+        self, rows: int, cached: int, nreq: int = 1, ctx_tokens: int = 0, ragged_lens=None
+    ):
+        """Generation metadata for ``nreq`` requests (``rows`` query tokens each).
+
+        ``ctx_tokens`` > 0 prepends the context request ``CTX_REQ_ID`` with that
+        many new tokens (mixed batch; ``nreq`` = 0 gives a pure context batch).
+        ``ragged_lens`` installs ragged verification with these per-request
+        query lengths.
+        """
+        nctx = 1 if ctx_tokens else 0
+        gen_lens = list(ragged_lens) if ragged_lens is not None else [rows] * nreq
+        kwargs = {}
+        if ragged_lens is not None:
+            kwargs.update(enable_ragged_verification=True)
+        md = self.attn_cls.Metadata(
+            seq_lens=torch.tensor([ctx_tokens] * nctx + gen_lens, dtype=torch.int),
+            request_ids=[CTX_REQ_ID] * nctx + self.rids[:nreq],
+            max_num_requests=nctx + nreq,
+            num_contexts=nctx,
+            prompt_lens=[ctx_tokens] * nctx + self.ctxs[:nreq],
+            max_num_tokens=128 + ctx_tokens,
+            kv_cache_manager=self.kvm,
+            kv_cache_params=KVCacheParams(
+                use_cache=True, num_cached_tokens_per_seq=[0] * nctx + self.cached(cached, nreq)
+            ),
+            mapping=self.mapping,
+            sparse_metadata_params=self.sparse.to_sparse_metadata_params(
+                pretrained_config=self.pretrained
+            ),
+            **kwargs,
+        )
+        if ragged_lens is not None:
+            md.ragged_verify_lens = list(ragged_lens)
+        md.update_spec_dec_param(
+            batch_size=nctx + nreq,
+            is_spec_decoding_enabled=True,
+            is_spec_dec_tree=False,
+            is_spec_dec_dynamic_tree=False,
+            max_draft_len=5,
+            max_total_draft_tokens=5,
+        )
+        md.prepare()
+        return md
+
+    def add_context_request(self):
+        """Allocate the context-phase request used by the mixed / context fallback cases."""
+        if not hasattr(self, "ctx_req"):
+            self.ctx_req = self.kvm.add_dummy_requests(
+                [CTX_REQ_ID], token_nums=[CTX_REQ_TOKENS], is_gen=False
+            )
+            self.reqs = list(self.reqs) + list(self.ctx_req)
+
+    def inputs(self, rows: int, cached: int, seed: int, nreq: int = 1):
+        return self.inputs_seg([(rows, c) for c in self.cached(cached, nreq)], seed)
+
+    def inputs_seg(self, segs, seed: int):
+        """Random layer inputs for request-major segments ``(rows, first_position)``."""
+        g = torch.Generator(device=self.dev).manual_seed(seed)
+        n = sum(r for r, _ in segs)
+        q = (torch.randn(n, HEADS * (NOPE + ROPE), generator=g, device=self.dev)).to(torch.bfloat16)
+        ckv = torch.randn(n, KV_LORA, generator=g, device=self.dev).to(torch.bfloat16)
+        k_pe = torch.randn(n, ROPE, generator=g, device=self.dev).to(torch.bfloat16)
+        qr = torch.randn(n, Q_LORA, generator=g, device=self.dev).to(torch.bfloat16)
+        hidden = torch.randn(n, self.hidden, generator=g, device=self.dev).to(torch.bfloat16)
+        pos = torch.cat(
+            [torch.arange(p, p + r, dtype=torch.int32, device=self.dev) for r, p in segs]
+        ).unsqueeze(0)
+        return q, ckv, k_pe, qr, hidden, pos
+
+    def state(self, md):
+        out = [p.view(torch.uint8).clone() for p in [self.kvm.get_unique_primary_pool()]]
+        out += [
+            p.view(torch.uint8).clone()
+            for p in self.kvm.indexer_k_cache_pool_per_layer
+            if p is not None
+        ]
+        if md.shared_topk_indices is not None:
+            out.append(md.shared_topk_indices.clone())
+        return out
+
+    def restore(self, md, state):
+        pools = [self.kvm.get_unique_primary_pool()] + [
+            p for p in self.kvm.indexer_k_cache_pool_per_layer if p is not None
+        ]
+        for p, s in zip(pools, state):
+            p.view(torch.uint8).copy_(s)
+        if md.shared_topk_indices is not None:
+            md.shared_topk_indices.copy_(state[-1])
+        md._group_remap_batched.clear()
+
+    @torch.inference_mode()
+    def run_layers(self, md, layer_ids, rows, cached, seed, nreq: int = 1):
+        return self.run_seg(md, layer_ids, [(rows, c) for c in self.cached(cached, nreq)], seed)
+
+    @torch.inference_mode()
+    def run_seg(self, md, layer_ids, segs, seed, latent_row_stride: int = 0):
+        """Run ``_forward_dsa_attn`` of ``layer_ids`` on request-major segments.
+
+        ``latent_row_stride`` > 576 passes ``latent_cache`` as a row-strided
+        view (outside the fused rope domain).
+        """
+        q, ckv, k_pe, qr, hidden, pos = self.inputs_seg(segs, seed)
+        n = q.shape[0]
+        outs = []
+        for idx in layer_ids:
+            mla = self.layers[idx]
+            inter = []
+            if mla.mqa.indexer is not None:
+                inter = list(mla.mqa.indexer.pre_indexer_proj(qr, hidden, pos))
+            latent = torch.cat([ckv, k_pe], dim=-1)
+            if latent_row_stride:
+                buf = torch.zeros(n, latent_row_stride, dtype=latent.dtype, device=self.dev)
+                buf[:, : KV_LORA + ROPE].copy_(latent)
+                latent = buf[:, : KV_LORA + ROPE]
+            out = torch.empty(n, HEADS * V, dtype=torch.bfloat16, device=self.dev)
+            with dispatch.step(dispatch.Phase.DECODE):
+                _forward_dsa_attn(
+                    mla, q.clone(), ckv.clone(), k_pe.clone(), latent, inter, pos, md, out
+                )
+            outs.append(out)
+        torch.cuda.synchronize()
+        return outs
+
+    @torch.inference_mode()
+    def run_proj(self, md, layer_ids, segs, seed):
+        """Run ``forward_sparse_attn`` (projections + attention) of ``layer_ids`` from hidden states."""
+        _, _, _, _, hidden, pos = self.inputs_seg(segs, seed)
+        outs = []
+        for idx in layer_ids:
+            out = torch.empty(hidden.shape[0], HEADS * V, dtype=torch.bfloat16, device=self.dev)
+            with dispatch.step(dispatch.Phase.DECODE):
+                forward_sparse_attn(self.layers[idx], pos, hidden, md, [out])
+            outs.append(out)
+        torch.cuda.synchronize()
+        return outs
+
+    @torch.inference_mode()
+    def indexer_k(self, mla, rows, cached, seed, nreq=1):
+        """Indexer K inputs of one layer call, as ``_forward_dsa_attn`` slices them."""
+        _, _, _, qr, hidden, pos = self.inputs(rows, cached, seed, nreq)
+        _, k_fp8, k_scale, _, _ = mla.mqa.indexer.pre_indexer_proj(qr, hidden, pos)
+        return k_fp8[: rows * nreq], k_scale[: rows * nreq]
+
+    @staticmethod
+    def _topk_set_diff(a, b):
+        """Per-row number of entries in one selection but not the other (order ignored)."""
+        diffs = []
+        for ra, rb in zip(a, b):
+            sa, sb = set(ra.tolist()), set(rb.tolist())
+            diffs.append(len(sa ^ sb))
+        return diffs
+
+    @staticmethod
+    def _state_equal(i, a, b, is_topk):
+        # The native GVR top-k emits its selection in a run-to-run varying order,
+        # so the top-k reuse buffer is compared as a per-row set.
+        if is_topk:
+            return torch.equal(torch.sort(a, dim=-1).values, torch.sort(b, dim=-1).values)
+        return torch.equal(a, b)
+
+    def _compare_captures(self):
+        """Per-layer FP8 query / scheduler / scale buffers and global indices vs the rope_kv_append-off runs.
+
+        A buffer whose two rope_kv_append-off runs agree must be byte-identical with rope_kv_append;
+        buffers that differ between the rope_kv_append-off runs (selection-dependent
+        global indices under a nondeterministic top-k) are counted and
+        reported instead.
+        """
+        (rope_a, idx_a), (rope_b, idx_b), (rope_w, idx_w) = self.captures
+        assert len(rope_a) == len(rope_b) == len(rope_w) > 0, (len(rope_a), len(rope_w))
+        assert len(idx_a) == len(idx_b) == len(idx_w) > 0, (len(idx_a), len(idx_w))
+        names = ("quant_q", "cu_q", "cu_kv", "counter", "bmm1", "bmm2")
+        report = {"rope_exact": 0, "rope_native_nondet": 0, "idx_exact": 0, "idx_native_nondet": 0}
+        for i, (a, b, w) in enumerate(zip(rope_a, rope_b, rope_w)):
+            for name, x, y, z in zip(names, a, b, w):
+                xb, yb, zb = (t.view(torch.uint8) for t in (x, y, z))
+                if torch.equal(xb, yb):
+                    assert torch.equal(xb, zb), (
+                        f"rope call {i}: {name} differs from the rope_kv_append-off run"
+                    )
+                    report["rope_exact"] += 1
+                else:
+                    report["rope_native_nondet"] += 1
+        # The native top-k emits each row's selection in a run-to-run varying order, so global
+        # indices are compared per row as sets when the two rope_kv_append-off runs differ only in
+        # order.
+        report.update(idx_set_exact=0)
+        for i, (a, b, w) in enumerate(zip(idx_a, idx_b, idx_w)):
+            if torch.equal(a, b):
+                assert torch.equal(a, w), (
+                    f"remap call {i}: global indices differ from the rope_kv_append-off run"
+                )
+                report["idx_exact"] += 1
+                continue
+            sa, sb, sw = (torch.sort(t, dim=-1).values for t in (a, b, w))
+            if torch.equal(sa, sb):
+                assert torch.equal(sa, sw), (
+                    f"remap call {i}: global index sets differ from rope_kv_append-off"
+                )
+                report["idx_set_exact"] += 1
+            else:
+                report["idx_native_nondet"] += 1
+        print("capture comparison:", report)
+        self.capture_report = report
+
+    def compare(self, md, step, base: str = "", capture: bool = False):
+        """Run ``step`` twice with features ``base`` and once with ``base`` + rope_kv_append from the same state.
+
+        Always: cache pools byte-identical to native, and every rope_kv_append bridge call
+        equal to the native remap/copy of the rows it was given (per-call
+        oracle). When the two native runs select the same top-k sets, rope_kv_append must
+        select the same sets too and its sparse-MLA output must be exact (or
+        within the native run-to-run spread). When the native selection itself
+        differs between runs (characterized and printed), the selection-
+        dependent outputs cannot be compared run to run; the per-call oracle
+        and the pools still are.
+        """
+        state0 = self.state(md)
+        runs = []
+        self.engagement = []
+        self.bridge_checks = []
+        self.captures = []
+        rope_append_features = ",".join([f for f in base.split(",") if f] + ["rope_kv_append"])
+        for features in (base, base, rope_append_features):
+            self.restore(md, state0)
+            set_features(features)
+            eng = Engagement(capture=capture)
+            try:
+                outs = step()
+            finally:
+                eng.close()
+            self.engagement.append(eng.counts)
+            self.bridge_checks.append(eng.bridge_checks)
+            self.captures.append((eng.rope_outputs, eng.global_indices))
+            runs.append((outs, self.state(md)))
+        set_features("")
+        if capture:
+            self._compare_captures()
+        print("engagement native/native/rope_kv_append:", self.engagement)
+        assert self.bridge_checks[2] == self.engagement[2].get("rope_append_bridge_kernel", 0)
+        (out_n, st_n), (out_n2, st_n2), (out_w, st_w) = runs
+        topk_idx = len(st_n) - 1 if md.shared_topk_indices is not None else -1
+        for i, (a, b, c) in enumerate(zip(st_n, st_n2, st_w)):
+            if i == topk_idx:
+                continue
+            assert torch.equal(a, b), f"native cache state {i} is not deterministic"
+            if not torch.equal(a, c):
+                raise AssertionError(f"cache state {i}: {(a != c).sum().item()} bytes differ")
+        assert any(not torch.equal(a, b) for a, b in zip(state0, st_n)), "native step wrote nothing"
+        self.native_topk_diff = None
+        if topk_idx >= 0:
+            self.native_topk_diff = self._topk_set_diff(st_n[topk_idx], st_n2[topk_idx])
+            if any(self.native_topk_diff):
+                nz = [(r, d) for r, d in enumerate(self.native_topk_diff) if d]
+                print(
+                    f"native top-k selection differs between native runs (row, |symmetric diff|): {nz}"
+                )
+                return None
+            rope_append_diff = self._topk_set_diff(st_n[topk_idx], st_w[topk_idx])
+            assert not any(rope_append_diff), (
+                f"rope_kv_append run selected different top-k sets: {rope_append_diff}"
+            )
+        report = []
+        for i, (a, a2, b) in enumerate(zip(out_n, out_n2, out_w)):
+            native_spread = (a.float() - a2.float()).abs().max().item()
+            rope_append_diff = (a.float() - b.float()).abs().max().item()
+            rope_append_diff2 = (a2.float() - b.float()).abs().max().item()
+            report.append((native_spread, min(rope_append_diff, rope_append_diff2)))
+            if native_spread == 0.0:
+                assert torch.equal(a.view(torch.int16), b.view(torch.int16)), f"output {i} differs"
+            else:
+                assert min(rope_append_diff, rope_append_diff2) <= 2 * native_spread, (
+                    i,
+                    native_spread,
+                    rope_append_diff,
+                )
+        print("output native-spread / rope_kv_append-diff per layer:", report)
+        return report
+
+    def shutdown(self):
+        for r in self.reqs:
+            self.kvm.free_resources(r)
+        self.kvm.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py
new file mode 100644
index 000000000000..6108700c5caf
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py
@@ -0,0 +1,226 @@
+# 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.
+"""rope_kv_append composed with its neighbour packages (needs a tree that contains them).
+
+Producers of rope_kv_append's inputs: the native projection + absorption BMM, absorbed_query (target
+layers), absorbed_query + shared_indexer_absorbed_query (shared-indexer target layers) and
+draft_absorbed_query (the MTP layer); consumers of its outputs: the native trtllm-gen sparse MLA or
+sparse_mla_decode, with the native or indexer_topk indexer logits in between. Every case runs
+``forward_sparse_attn`` (projections + attention) from hidden states with the neighbour set alone,
+twice, and with rope_kv_append added, from identical state. Compared byte for byte
+against the rope_kv_append-off runs: every KV and indexer-K pool byte, each layer's FP8
+query buffer, ``cu_q``/``cu_kv`` scheduler offsets, scheduler counter, BMM1/BMM2
+scales and global KV indices (a buffer that differs between the two rope_kv_append-off
+runs is reported, not compared). rope_kv_append and every selected neighbour must engage,
+and rope_kv_append must not change the neighbours' engagement.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+from rope_append_stack import MTP_LAYER, PATTERN, Stack, load_dev_library
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+
+
+def _have(pkg: str, name: str) -> bool:
+    try:
+        dispatch.get_impl(pkg, name)
+    except (ImportError, KeyError):
+        return False
+    return True
+
+
+NEIGHBOURS_PRESENT = all(
+    _have(p, n)
+    for p, n in (
+        ("absorbed_query", "dsa_proj"),
+        ("absorbed_query", "shared_dsa_proj"),
+        ("absorbed_query", "mtp_dsa_proj"),
+        ("sparse_mla_decode", "fmha_lib"),
+        ("indexer_topk", "indexer_decode"),
+    )
+)
+
+
+pytestmark = [
+    pytest.mark.threadleak(enabled=False),
+    pytest.mark.skipif(
+        not NEIGHBOURS_PRESENT,
+        reason="absorbed_query/draft_absorbed_query/shared_indexer_absorbed_query/"
+        "sparse_mla_decode/indexer_topk are not in this tree",
+    ),
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+        reason="rope_kv_append targets Blackwell decode",
+    ),
+]
+
+CTX = 32768
+TARGET_CASES = {
+    "native": "",
+    "absorbed_query": "absorbed_query",
+    "absorbed_query_shared_query": "absorbed_query,shared_indexer_absorbed_query",
+    "indexer_topk": "indexer_topk",
+    "sparse_mla_decode": "sparse_mla_decode",
+    "all": "absorbed_query,shared_indexer_absorbed_query,draft_absorbed_query,indexer_topk,sparse_mla_decode",
+}
+DRAFT_CASES = {
+    "native": "",
+    "draft_absorbed_query": "draft_absorbed_query",
+    "sparse_mla_decode": "sparse_mla_decode",
+    "all": "absorbed_query,shared_indexer_absorbed_query,draft_absorbed_query,indexer_topk,sparse_mla_decode",
+}
+NEIGHBOUR_KEYS = (
+    "aq_packet",
+    "shared_query_packet",
+    "draft_query_packet",
+    "indexer_topk_decode",
+    "sparse_mla_fmha",
+)
+
+
+@pytest.fixture(scope="module")
+def stack():
+    load_dev_library()
+    # The adopted layers declare the sparse-MLA decode library only under a plan that runs
+    # sparse_mla_decode: the stack is built under the plan of ``plan_testlib.FAST_PATHS_ENV``.
+    plan_testlib.select_from_env(dispatch)
+    s = Stack(CTX, real_proj=True)
+    yield s
+    s.shutdown()
+
+
+def _check_neighbours(features: str, base_counts, rope_append_counts, layers: int, draft: bool):
+    feats = {f for f in features.split(",") if f}
+    for key in NEIGHBOUR_KEYS:
+        assert base_counts.get(key, 0) == rope_append_counts.get(key, 0), (
+            key,
+            base_counts,
+            rope_append_counts,
+        )
+    if "absorbed_query" in feats and not draft:
+        shared = sum(1 for p in PATTERN[:layers] if p == "S")
+        assert (
+            rope_append_counts.get("aq_packet", 0)
+            + rope_append_counts.get("shared_query_packet", 0)
+            == layers
+        )
+        if "shared_indexer_absorbed_query" in feats:
+            assert rope_append_counts.get("shared_query_packet", 0) == shared
+    if "draft_absorbed_query" in feats and draft:
+        assert rope_append_counts.get("draft_query_packet", 0) == 1
+    if "sparse_mla_decode" in feats:
+        assert rope_append_counts.get("sparse_mla_fmha", 0) == layers
+
+
+def _sparse_mla_declared(stack) -> bool:
+    """Whether the stack's layers put the sparse-MLA decode library ahead of the registry's.
+
+    Once declared, sparse_mla_decode serves every eligible layer, so sparse_mla_decode cases run
+    in a process whose stack was built with ``sparse_mla_decode`` in
+    ``plan_testlib.FAST_PATHS_ENV``, the other cases in one without.
+    """
+    return any(getattr(mla.mqa, "prepended_fmha_lib_classes", ()) for mla in stack.layers)
+
+
+def _skip_unless_sparse_mla_matches(stack, features: str):
+    wants_sparse_mla = "sparse_mla_decode" in features.split(",")
+    if wants_sparse_mla != _sparse_mla_declared(stack):
+        pytest.skip(
+            "sparse_mla_decode cases need the stack built with sparse_mla_decode running; others one without"
+            if wants_sparse_mla
+            else "the stack declares sparse_mla_decode, which would serve this case"
+        )
+
+
+@pytest.mark.parametrize("case", list(TARGET_CASES))
+def test_target_layers_composed(stack, case):
+    """Verify step (M = 6) through full + 3 shared target layers."""
+    features = TARGET_CASES[case]
+    _skip_unless_sparse_mla_matches(stack, features)
+    md = stack.metadata(6, stack.ctx_len)
+    stack.compare(
+        md,
+        lambda: stack.run_proj(md, range(len(PATTERN)), [(6, stack.ctx_len)], seed=51),
+        base=features,
+        capture=True,
+    )
+    base_counts, _, rope_append_counts = stack.engagement
+    print(f"{case}: base {base_counts}\n{case}: rope_kv_append {rope_append_counts}")
+    assert (
+        rope_append_counts.get("rope_append_mla_rope_generation") == 4
+        and rope_append_counts.get("rope_append_fused_scatter") == 1
+    )
+    assert (
+        rope_append_counts.get("rope_append_remap_topk_to_global") == 4
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert rope_append_counts.get("rope_append_check_indexer_k_scatter_consumed") == 1
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+    _check_neighbours(features, base_counts, rope_append_counts, len(PATTERN), draft=False)
+    report = stack.capture_report
+    # quant_q may differ between runs without rope_kv_append; the rest may not.
+    assert report["rope_exact"] >= 4 * 5
+    assert report["idx_exact"] + report["idx_set_exact"] + report["idx_native_nondet"] == 4
+
+
+@pytest.mark.parametrize("case", list(DRAFT_CASES))
+def test_mtp_draft_steps_composed(stack, case):
+    """MTP layer: draft step 0 (6 rows, last-accepted-row stash) and a reuse step (1 row)."""
+    features = DRAFT_CASES[case]
+    _skip_unless_sparse_mla_matches(stack, features)
+    md0 = stack.metadata(6, stack.ctx_len)
+    md0.set_in_mtp_draft_loop(True)
+    md0.set_mtp_num_accepted(torch.tensor([4], dtype=torch.int32, device=stack.dev))
+    md0.set_skip_topk(False)
+    stack.compare(
+        md0,
+        lambda: stack.run_proj(md0, [MTP_LAYER], [(6, stack.ctx_len)], seed=52),
+        base=features,
+        capture=True,
+    )
+    base_counts, _, rope_append_counts = stack.engagement
+    print(f"{case} step0: base {base_counts}\n{case} step0: rope_kv_append {rope_append_counts}")
+    assert (
+        rope_append_counts.get("rope_append_fused_scatter") == 1
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+    _check_neighbours(features, base_counts, rope_append_counts, 1, draft=True)
+
+    md1 = stack.metadata(1, stack.ctx_len + 5)
+    md1.set_in_mtp_draft_loop(True)
+    md1.set_skip_topk(True)
+    md1.shared_topk_indices.copy_(md0.shared_topk_indices)
+    stack.compare(
+        md1,
+        lambda: stack.run_proj(md1, [MTP_LAYER], [(1, stack.ctx_len + 5)], seed=53),
+        base=features,
+        capture=True,
+    )
+    base_counts, _, rope_append_counts = stack.engagement
+    print(f"{case} step1: base {base_counts}\n{case} step1: rope_kv_append {rope_append_counts}")
+    assert (
+        rope_append_counts.get("rope_append_fused_scatter") == 1
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+    _check_neighbours(features, base_counts, rope_append_counts, 1, draft=True)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py
new file mode 100644
index 000000000000..6d21dc14199d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py
@@ -0,0 +1,501 @@
+# 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.
+"""rope_kv_append through the target's DSA call sites (the adopted attention, on decode steps).
+
+Real MLA/DSA modules at GLM-5.2 TP8 attention geometry (8 local heads,
+q-LoRA 2048, kv-LoRA 512, nope 192, rope 64, v 256, 32 x 128 indexer, top-k
+2048, FP8 KV cache, block 64, heuristic GVR top-k, CuTe DSL logits), four
+target layers in the GLM pattern full/shared/shared/shared plus the MTP layer,
+a real ``DSACacheManager`` filled with FP8 context and real generation
+metadata (verify M = 6, MTP draft rows 6 then 1), for one request and for
+two requests with different cache lengths. Each step runs the production
+``_forward_dsa_attn`` natively (no fast path running) and with rope_kv_append
+from identical state. Every KV and indexer-K pool byte must be identical, the
+shared top-k buffer must hold the same per-row selection and the attention
+outputs must be within the native run-to-run spread.
+"""
+
+import pytest
+import torch
+from rope_append_stack import (
+    CTX_REQ_TOKENS,
+    HEADS,
+    KV_LORA,
+    MTP_LAYER,
+    PATTERN,
+    ROPE,
+    Stack,
+    load_dev_library,
+    set_features,
+)
+
+import tensorrt_llm
+import tensorrt_llm.bindings
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+pytestmark = [
+    pytest.mark.threadleak(enabled=False),
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+        reason="rope_kv_append targets Blackwell decode",
+    ),
+]
+
+
+@pytest.fixture(scope="module", params=[32768, 491_061], ids=["ctx32k", "ctx491k"])
+def stack(request):
+    load_dev_library()
+    s = Stack(request.param)
+    yield s
+    s.shutdown()
+
+
+@pytest.fixture(scope="module")
+def stack32k():
+    """32K context stack for the call-site fallback cases."""
+    load_dev_library()
+    s = Stack(32768)
+    s.add_context_request()
+    yield s
+    s.shutdown()
+
+
+def test_native_state_characterization(stack):
+    """Native run-to-run behaviour: cache bytes are deterministic; the top-k selection is reported.
+
+    The heuristic GVR top-k emits its rows in a varying order and, on some
+    inputs, selects different sets in two runs from identical state; the
+    per-row symmetric difference is printed so it is characterized, not hidden.
+    """
+    md = stack.metadata(6, stack.ctx_len)
+    state0 = stack.state(md)
+    set_features("")
+    stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=1)
+    st_a = stack.state(md)
+    stack.restore(md, state0)
+    stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=1)
+    st_b = stack.state(md)
+    topk_idx = len(st_a) - 1 if md.shared_topk_indices is not None else -1
+    for i, (x, y) in enumerate(zip(st_a, st_b)):
+        if i != topk_idx:
+            assert torch.equal(x, y), f"native cache state {i} differs"
+    if topk_idx >= 0:
+        diff = Stack._topk_set_diff(st_a[topk_idx][:6], st_b[topk_idx][:6])
+        print(
+            f"ctx {stack.ctx_len}: native top-k per-row symmetric difference between runs: {diff}"
+        )
+
+
+def test_target_verify_layers(stack):
+    """Full layer (deferred scatter, grouped bridge with the reuse copy) + 3 shared layers."""
+    md = stack.metadata(6, stack.ctx_len)
+    stack.compare(md, lambda: stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=2))
+    native, _, rope_append_counts = stack.engagement
+    assert native.get("native_scatter") == 1 and native.get("native_rope") == 4
+    assert native.get("native_remap") == 4 and "rope_append_mla_rope_generation" not in native
+    # rope_kv_append: one fused rope per layer (the full layer carries the scatter), one grouped
+    # bridge launch for the whole group, no native rope/scatter/remap.
+    assert (
+        rope_append_counts.get("rope_append_mla_rope_generation") == 4
+        and rope_append_counts.get("rope_append_fused_scatter") == 1
+    )
+    assert (
+        rope_append_counts.get("rope_append_remap_topk_to_global") == 4
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+
+
+@pytest.mark.parametrize("accepted", [1, 4, 6])
+def test_mtp_draft_steps(stack, accepted):
+    """Draft step 0 (row stash of the last accepted row) and a reuse step (rows = 1)."""
+    md0 = stack.metadata(6, stack.ctx_len)
+    md0.set_in_mtp_draft_loop(True)
+    md0.set_mtp_num_accepted(torch.tensor([accepted], dtype=torch.int32, device=stack.dev))
+    md0.set_skip_topk(False)
+    stack.compare(md0, lambda: stack.run_layers(md0, [MTP_LAYER], 6, stack.ctx_len, seed=3))
+    rope_append_counts = stack.engagement[2]
+    assert (
+        rope_append_counts.get("rope_append_fused_scatter") == 1
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+
+    md1 = stack.metadata(1, stack.ctx_len + 5)
+    md1.set_in_mtp_draft_loop(True)
+    md1.set_skip_topk(True)
+    md1.shared_topk_indices.copy_(md0.shared_topk_indices)
+    stack.compare(md1, lambda: stack.run_layers(md1, [MTP_LAYER], 1, stack.ctx_len + 5, seed=4))
+    rope_append_counts = stack.engagement[2]
+    assert (
+        rope_append_counts.get("rope_append_fused_scatter") == 1
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+
+
+def test_target_verify_layers_two_requests(stack):
+    """Two generation requests with different cache lengths (request-major rows 0..5, 6..11)."""
+    md = stack.metadata(6, stack.ctx_len, nreq=2)
+    assert md.num_generations == 2 and md.num_tokens == 12
+    stack.compare(
+        md, lambda: stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=5, nreq=2)
+    )
+    rope_append_counts = stack.engagement[2]
+    assert (
+        rope_append_counts.get("rope_append_mla_rope_generation") == 4
+        and rope_append_counts.get("rope_append_fused_scatter") == 1
+    )
+    assert (
+        rope_append_counts.get("rope_append_remap_topk_to_global") == 4
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+
+
+@pytest.mark.parametrize("accepted", [(1, 6), (4, 2)])
+def test_mtp_draft_steps_two_requests(stack, accepted):
+    """Draft step 0 stashes each request's own last accepted row; reuse step with 2 x 1 rows."""
+    md0 = stack.metadata(6, stack.ctx_len, nreq=2)
+    md0.set_in_mtp_draft_loop(True)
+    md0.set_mtp_num_accepted(torch.tensor(list(accepted), dtype=torch.int32, device=stack.dev))
+    md0.set_skip_topk(False)
+    stack.compare(md0, lambda: stack.run_layers(md0, [MTP_LAYER], 6, stack.ctx_len, seed=6, nreq=2))
+    rope_append_counts = stack.engagement[2]
+    assert (
+        rope_append_counts.get("rope_append_fused_scatter") == 1
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+
+    md1 = stack.metadata(1, stack.ctx_len + 5, nreq=2)
+    md1.set_in_mtp_draft_loop(True)
+    md1.set_skip_topk(True)
+    md1.shared_topk_indices.copy_(md0.shared_topk_indices)
+    stack.compare(
+        md1, lambda: stack.run_layers(md1, [MTP_LAYER], 1, stack.ctx_len + 5, seed=7, nreq=2)
+    )
+    rope_append_counts = stack.engagement[2]
+    assert (
+        rope_append_counts.get("rope_append_fused_scatter") == 1
+        and rope_append_counts.get("rope_append_bridge_kernel") == 1
+    )
+    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
+
+
+def test_lost_deferred_scatter_raises(stack, monkeypatch):
+    """A call site that defers the indexer scatter but drops the packet fails at once.
+
+    Reproduces a call-site placement defect (the packet never reached the
+    generation rope, so nothing wrote the step's indexer keys): the deferring
+    call site's consumed check must raise during that forward.
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
+        attention_forward,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        RopeKvAppendContractError,
+    )
+
+    orig = attention_forward.forward_generation_sparse_attn
+
+    def dropping(*args, deferred_indexer_k=None, **kwargs):
+        return orig(*args, **kwargs)
+
+    monkeypatch.setattr(attention_forward, "forward_generation_sparse_attn", dropping)
+    md = stack.metadata(6, stack.ctx_len)
+    state0 = stack.state(md)
+    set_features("rope_kv_append")
+    try:
+        with pytest.raises(RopeKvAppendContractError, match="not consumed"):
+            stack.run_layers(md, [0], 6, stack.ctx_len, seed=8)
+    finally:
+        set_features("")
+        # The raise skipped the layer's join of the forked prev-top-k copy.
+        stack.layers[0].mqa.indexer.maybe_join_prev_topk_copy()
+        torch.cuda.synchronize()
+        stack.restore(md, state0)
+
+
+def test_deferred_scatter_contract(stack):
+    """The packet is single use, the check sees a missing consumer, undeclared callers get None."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
+        rope_kv_append,
+    )
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
+        RopeKvAppendContractError,
+    )
+
+    md = stack.metadata(6, stack.ctx_len)
+    state0 = stack.state(md)
+    mla = stack.layers[0]
+    k_fp8, k_scale = stack.indexer_k(mla, 6, stack.ctx_len, seed=9)
+    set_features("rope_kv_append")
+    try:
+        assert rope_kv_append.defer_indexer_k_scatter(mla, k_fp8, k_scale, md) is None
+        pkt = rope_kv_append.defer_indexer_k_scatter(mla, k_fp8, k_scale, md, forwards_packet=True)
+        assert pkt is not None and pkt.consumer is None and pkt.layer_idx == 0
+        with pytest.raises(RopeKvAppendContractError, match="not consumed"):
+            rope_kv_append.check_indexer_k_scatter_consumed(pkt)
+        rope_kv_append.run_indexer_k_scatter_native(pkt)
+        rope_kv_append.check_indexer_k_scatter_consumed(pkt)
+        with pytest.raises(RopeKvAppendContractError, match="consumed twice"):
+            rope_kv_append.run_indexer_k_scatter_native(pkt)
+        rope_kv_append.check_indexer_k_scatter_consumed(None)
+        torch.cuda.synchronize()
+        # The native fallback wrote exactly what the native scatter writes.
+        after_rope_append = stack.state(md)
+        stack.restore(md, state0)
+        set_features("")
+        mla.mqa.indexer._update_k_cache(k_fp8, k_scale, md)
+        torch.cuda.synchronize()
+        native = stack.state(md)
+        for i, (a, b) in enumerate(zip(after_rope_append, native)):
+            assert torch.equal(a, b), f"state {i} differs"
+    finally:
+        set_features("")
+        stack.restore(md, state0)
+
+
+def test_pool_view_memo_is_lazy(stack):
+    """The fused bridge leaves the pool-view memo unpopulated; a later reader populates it.
+
+    Accepted deviation from the native remap's side effect: every reader of
+    ``_cached_*`` calls ``_ensure_pool_view_cached()`` first, so values are
+    unchanged and only the memo's ``req_idx - num_contexts`` kernel moves to
+    the first such reader of the step (none on the rope_kv_append/sparse_mla_decode/native-FMHA
+    decode path).
+    """
+    md = stack.metadata(6, stack.ctx_len)
+    state0 = stack.state(md)
+    set_features("rope_kv_append")
+    try:
+        stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=10)
+        assert not md._pool_cache_valid
+        md._ensure_pool_view_cached()
+        nc = md.num_contexts
+        want_req = md.req_idx_per_token[md.num_ctx_tokens : md.num_tokens] - nc
+        assert torch.equal(md._cached_req_idx_gen, want_req)
+        assert torch.equal(md._cached_block_table_gen, md.block_table[nc : md.num_seqs])
+    finally:
+        set_features("")
+        stack.restore(md, state0)
+
+
+def test_default_off_dispatch(stack):
+    """With the selector empty, no rope_kv_append helper is consulted (native call graph)."""
+    set_features("")
+    md = stack.metadata(6, stack.ctx_len)
+    assert dispatch.enabled_features() == frozenset()
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
+        rope_kv_append,
+    )
+
+    assert rope_kv_append.defer_indexer_k_scatter(stack.layers[0], None, None, md) is None
+    assert rope_kv_append.new_bridge_packet(md, True) is None
+
+
+# ------------------------------------------------------------------------------------------
+# Call-site fallbacks: every rejection selects the native path before any rope_kv_append side
+# effect, the pools stay byte-identical to the rope_kv_append-off run and the engagement counters
+# show which path ran.
+
+ROPE_APPEND_ENTRY_KEYS = (
+    "rope_append_mla_rope_generation",
+    "rope_append_remap_topk_to_global",
+    "rope_append_bridge_kernel",
+    "rope_append_deferred_scatter",
+    "rope_append_fused_scatter",
+    "rope_append_run_indexer_k_scatter_native",
+)
+
+
+def _no_rope_append_engaged(counts):
+    return not any(counts.get(k) for k in ROPE_APPEND_ENTRY_KEYS)
+
+
+def test_fallback_rope_ineligible_after_deferral(stack32k):
+    """Rope outside the fused domain (row-strided latent): the deferred scatter runs natively."""
+    st = stack32k
+    md = st.metadata(6, st.ctx_len)
+    segs = [(6, st.ctx_len)]
+    st.compare(
+        md, lambda: st.run_seg(md, range(len(PATTERN)), segs, seed=31, latent_row_stride=640)
+    )
+    native, _, scatter_counts = st.engagement
+    assert native.get("native_scatter") == 1 and native.get("native_rope") == 4
+    assert (
+        scatter_counts.get("rope_append_deferred_scatter") == 1
+        and scatter_counts.get("rope_append_run_indexer_k_scatter_native") == 1
+    )
+    assert scatter_counts.get("rope_append_check_indexer_k_scatter_consumed") == 1
+    assert not scatter_counts.get("rope_append_mla_rope_generation") and not scatter_counts.get(
+        "native_scatter"
+    )
+    assert scatter_counts.get("native_rope") == 4
+    # The bridge does not depend on the latent layout and stays engaged.
+    assert (
+        scatter_counts.get("rope_append_remap_topk_to_global") == 4
+        and scatter_counts.get("rope_append_bridge_kernel") == 1
+    )
+
+
+@pytest.fixture(scope="module")
+def stack_g9():
+    """One full layer leading 8 shared layers: a group larger than the fused bridge supports."""
+    load_dev_library()
+    s = Stack(32768, pattern=["F"] + ["S"] * 8)
+    yield s
+    s.shutdown()
+
+
+def test_fallback_group_larger_than_8(stack_g9):
+    """Group of 9: the leader copies and remaps natively, followers slice the native batch."""
+    st = stack_g9
+    md = st.metadata(6, st.ctx_len)
+    st.compare(md, lambda: st.run_layers(md, range(len(st.pattern)), 6, st.ctx_len, seed=32))
+    _, _, scatter_counts = st.engagement
+    assert scatter_counts.get("native_remap") == 1 and not scatter_counts.get(
+        "rope_append_bridge_kernel"
+    )
+    assert scatter_counts.get("rope_append_remap_topk_to_global") == 9
+    assert (
+        scatter_counts.get("rope_append_mla_rope_generation") == 9
+        and scatter_counts.get("rope_append_fused_scatter") == 1
+    )
+
+
+def test_fallback_context_batch(stack32k):
+    """Pure context batch: no rope_kv_append entry engages."""
+    st = stack32k
+    md = st.metadata(0, st.ctx_len, nreq=0, ctx_tokens=CTX_REQ_TOKENS)
+    assert md.num_contexts == 1 and md.num_generations == 0
+    st.compare(md, lambda: st.run_seg(md, range(len(PATTERN)), [(CTX_REQ_TOKENS, 0)], seed=33))
+    native, _, scatter_counts = st.engagement
+    assert _no_rope_append_engaged(scatter_counts), scatter_counts
+    assert native.get("native_scatter") == 1 and scatter_counts.get("native_scatter") == 1
+
+
+def test_fallback_mixed_batch(stack32k):
+    """Context + generation in one batch: no rope_kv_append entry engages."""
+    st = stack32k
+    md = st.metadata(6, st.ctx_len, nreq=1, ctx_tokens=CTX_REQ_TOKENS)
+    assert md.num_contexts == 1 and md.num_generations == 1
+    segs = [(CTX_REQ_TOKENS, 0), (6, st.ctx_len)]
+    st.compare(md, lambda: st.run_seg(md, range(len(PATTERN)), segs, seed=34))
+    native, _, scatter_counts = st.engagement
+    assert _no_rope_append_engaged(scatter_counts), scatter_counts
+    assert scatter_counts.get("native_scatter") == native.get("native_scatter") == 1
+    assert scatter_counts.get("native_rope") == native.get("native_rope")
+
+
+def test_fallback_ragged_verify(stack32k):
+    """Ragged verification (token-major rows): rope and scatter native; bridge per-call exact."""
+    st = stack32k
+    md = st.metadata(6, st.ctx_len, nreq=2, ragged_lens=[6, 3])
+    assert md.is_ragged_verify and md.token_major_gen_view() is not None
+    c0, c1 = st.cached(st.ctx_len, 2)
+    st.compare(md, lambda: st.run_seg(md, range(len(PATTERN)), [(6, c0), (3, c1)], seed=35))
+    native, _, scatter_counts = st.engagement
+    assert not scatter_counts.get("rope_append_deferred_scatter") and not scatter_counts.get(
+        "rope_append_mla_rope_generation"
+    )
+    assert scatter_counts.get("native_scatter") == native.get("native_scatter") == 1
+    assert scatter_counts.get("native_rope") == native.get("native_rope")
+
+
+class _Over:
+    """Attribute overlay over a real object (unset attributes forward to it)."""
+
+    def __init__(self, base, **overrides):
+        self._base = base
+        self.__dict__.update(overrides)
+
+    def __getattr__(self, name):
+        return getattr(self._base, name)
+
+
+def test_eligibility_rejections(stack32k):
+    """Host-side predicates reject every unsupported mode before device work.
+
+    Starts from a real eligible GLM decode call and changes one condition at
+    a time (NVFP4 / FP4 cache or attention, beam > 1, helix, ragged, strided
+    latent, non-packed fused_q, residual-dim cache, missing KV manager).
+    """
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
+        rope_kv_append,
+    )
+
+    st = stack32k
+    md = st.metadata(6, st.ctx_len)
+    mla = st.layers[0]
+    mqa = mla.mqa
+    n = 6
+    fused_q = torch.zeros(n, HEADS, KV_LORA + ROPE, dtype=torch.bfloat16, device=st.dev)
+    q_pe = torch.zeros(n, HEADS, ROPE, dtype=torch.bfloat16, device=st.dev)
+    latent = torch.zeros(n, KV_LORA + ROPE, dtype=torch.bfloat16, device=st.dev)
+    quant_q = torch.zeros(n, HEADS, KV_LORA + ROPE, dtype=torch.uint8, device=st.dev)
+    s1 = torch.zeros(2, dtype=torch.float32, device=st.dev)
+    s2 = torch.zeros(1, dtype=torch.float32, device=st.dev)
+    k_fp8, k_scale = st.indexer_k(mla, 6, st.ctx_len, seed=36)
+
+    def rope(mqa_=mqa, md_=md, latent_=latent, fused_q_=fused_q):
+        return rope_kv_append.rope_eligible(mqa_, md_, fused_q_, q_pe, latent_, quant_q, s1, s2)
+
+    def defer(mla_=mla, md_=md):
+        return rope_kv_append.defer_indexer_k_scatter(
+            mla_, k_fp8, k_scale, md_, forwards_packet=True
+        )
+
+    set_features("rope_kv_append")
+    try:
+        assert rope() and rope_kv_append.bridge_eligible(md, True) and defer() is not None
+        nvfp4 = tensorrt_llm.bindings.DataType.NVFP4
+        cases = {
+            "fp4_kv_cache": dict(mqa_=_Over(mqa, has_fp4_kv_cache=True)),
+            "no_fp8_kv_cache": dict(mqa_=_Over(mqa, has_fp8_kv_cache=False)),
+            "fp4_mla_attention": dict(mqa_=_Over(mqa, uses_fp4_mla_attention=True)),
+            "beam2": dict(md_=_Over(md, beam_width=2)),
+            "helix": dict(md_=_Over(md, helix_position_offsets=torch.zeros(1))),
+            "ragged": dict(md_=_Over(md, is_ragged_verify=True)),
+            "token_major": dict(md_=_Over(md, token_major_gen_view=lambda: object())),
+            "context": dict(md_=_Over(md, num_contexts=1)),
+            "no_kv_manager": dict(md_=_Over(md, kv_cache_manager=None)),
+            "residual_dim_cache": dict(
+                md_=_Over(
+                    md, kv_cache_manager=_Over(md.kv_cache_manager, mla_kv_cache_residual_dim=64)
+                )
+            ),
+        }
+        for name, kw in cases.items():
+            assert not rope(**kw), name
+        strided = torch.zeros(n, 640, dtype=torch.bfloat16, device=st.dev)[:, : KV_LORA + ROPE]
+        assert not rope(latent_=strided)
+        assert not rope(fused_q_=fused_q.transpose(0, 1).contiguous().transpose(0, 1))
+        for name in ("fp4_kv_cache", "beam2", "helix", "ragged", "token_major", "context"):
+            kw = cases[name]
+            assert (
+                defer(
+                    mla_=_Over(mla, mqa=kw["mqa_"]) if "mqa_" in kw else mla, md_=kw.get("md_", md)
+                )
+                is None
+            ), name
+        assert not rope_kv_append.bridge_eligible(
+            _Over(md, kv_cache_manager=_Over(md.kv_cache_manager, dtype=nvfp4)), True
+        )
+        assert not rope_kv_append.bridge_eligible(_Over(md, num_contexts=1), True)
+        assert not rope_kv_append.bridge_eligible(md, False)
+    finally:
+        set_features("")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py
new file mode 100644
index 000000000000..818ff85ba7b1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py
@@ -0,0 +1,597 @@
+# 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.
+"""rope_kv_append fused ops versus the native TensorRT-LLM ops, byte for byte.
+
+The oracle is the native op sequence on a real ``DSACacheManager``:
+``indexer_k_cache_scatter_op`` + ``mla_rope_generation`` for the fused rope,
+and the top-k reuse ``copy_`` + ``convert_req_index_to_global[_grouped]`` for
+the bridge. Every cache pool byte, the FP8 query buffer, the FMHA scheduler
+buffers, the bmm scales, the global index tensor and the reuse buffer must be
+identical. Batches cover one request, several live requests with distinct
+cache lengths and the CUDA-graph padding dummy request; a captured-graph
+stress test puts a delayed PDL-launched metadata producer right before each
+rope_kv_append kernel.
+"""
+
+import os
+
+import pdl_producer
+import pytest
+import torch
+from rope_append_env import (
+    NUM_HEADS,
+    PAD,
+    TOKENS_PER_BLOCK,
+    TOPK,
+    RopeAppendEnv,
+    RopeOutputs,
+    load_dev_library,
+    make_batch_inputs,
+    make_batch_topk,
+    make_inputs,
+    make_topk,
+)
+
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+pytestmark = [
+    pytest.mark.threadleak(enabled=False),
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+        reason="rope_kv_append targets Blackwell decode",
+    ),
+]
+
+LONG_CTX = 491_061
+NEAR_MAX_CTX = (1 << 20) + 8
+
+
+@pytest.fixture(scope="module")
+def env():
+    load_dev_library()
+    e = RopeAppendEnv()
+    yield e
+    e.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def _release(env):
+    yield
+    env.release()
+
+
+def _assert_pools_equal(env, expected):
+    for i, (got, want) in enumerate(zip(env.pools(), expected)):
+        if not torch.equal(got, want):
+            diff = (got != want).nonzero()
+            raise AssertionError(
+                f"pool {i}: {diff.shape[0]} bytes differ, first at {diff[:4].tolist()}"
+            )
+
+
+def _assert_outputs_equal(out_w, out_n):
+    names = ["quant_q", "cu_q", "cu_kv", "counter", "bmm1", "bmm2"]
+    for name, got, want in zip(names, out_w.tensors(), out_n.tensors()):
+        g = got.view(torch.uint8) if got.dtype != torch.uint8 else got
+        w = want.view(torch.uint8) if want.dtype != torch.uint8 else want
+        assert torch.equal(g, w), f"{name} differs: {(g != w).sum().item()} bytes"
+
+
+def _run_rope_pair(env, step, layer, inp, with_scatter, q_scaling=1.0):
+    snap0 = env.snapshot()
+    out_n = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
+    if with_scatter:
+        env.native_scatter(step, layer, inp)
+    env.native_rope(step, layer, inp, out_n, q_scaling=q_scaling)
+    torch.cuda.synchronize()
+    native_pools = env.snapshot()
+    # The native ops must have written the cache (guards against a vacuous comparison).
+    assert any(not torch.equal(a, b) for a, b in zip(native_pools, snap0))
+
+    env.restore(snap0)
+    out_w = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
+    env.rope_append_rope(step, layer, inp, out_w, with_scatter, q_scaling=q_scaling)
+    torch.cuda.synchronize()
+    _assert_pools_equal(env, native_pools)
+    _assert_outputs_equal(out_w, out_n)
+    del snap0, native_pools
+    return out_n
+
+
+@pytest.mark.parametrize("with_scatter", [True, False], ids=["full_layer", "shared_layer"])
+@pytest.mark.parametrize("ctx", [1, 2047, 2048, 2049, 32768, 131072, LONG_CTX, NEAR_MAX_CTX])
+@pytest.mark.parametrize("rows", [6, 1])
+def test_rope_append_scatter_matches_native(env, rows, ctx, with_scatter):
+    step = env.new_step(ctx, rows)
+    # Distinct values per case: requests reuse the same pool blocks across tests.
+    inp = make_inputs(rows, env.device, scale=2.0, seed=ctx * 7 + rows * 3 + int(with_scatter) * 2)
+    layer = 3 if ctx % 2 else 0
+    _run_rope_pair(env, step, layer, inp, with_scatter)
+
+
+@pytest.mark.parametrize(
+    "scale,specials",
+    [(1e-3, False), (40.0, False), (300.0, True)],
+    ids=["tiny", "wide", "saturating+specials"],
+)
+@pytest.mark.parametrize("q_scaling", [1.0, 0.53366])
+def test_rope_value_ranges_match_native(env, scale, specials, q_scaling):
+    step = env.new_step(LONG_CTX, 6)
+    inp = make_inputs(
+        6, env.device, scale=scale, seed=int(scale * 1000) + int(q_scaling * 100), specials=specials
+    )
+    _run_rope_pair(env, step, 2, inp, True, q_scaling=q_scaling)
+
+
+def test_rope_noncontiguous_pages(env):
+    """Requests interleaved in the pool so the written block is not the first slot."""
+    for ctx in (64 * 37 + 5, 1000, 64 * 3 - 1):
+        env.new_step(ctx, 6)
+    step = env.new_step(LONG_CTX - 7, 6)
+    inp = make_inputs(6, env.device, seed=11)
+    _run_rope_pair(env, step, 1, inp, True)
+
+
+@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "per_layer"])
+@pytest.mark.parametrize("ctx", [1, 2047, 2049, 131072, LONG_CTX, NEAR_MAX_CTX])
+@pytest.mark.parametrize("rows", [6, 1])
+def test_bridge_matches_native(env, rows, ctx, grouped):
+    step = env.new_step(ctx, rows)
+    topk = make_topk(rows, ctx + rows, step.block_table.shape[1], env.device, seed=ctx)
+    layers = [0, 1, 2, 3] if grouped else [2]
+    shared_n = torch.full((16, TOPK), -5, dtype=torch.int32, device=env.device)
+    shared_w = shared_n.clone()
+    want = env.native_bridge(step, topk, layers, shared_n, grouped)
+    got = env.rope_append_bridge(step, topk, layers, shared_w, grouped)
+    torch.cuda.synchronize()
+    assert got.shape == want.shape
+    assert torch.equal(got, want)
+    assert torch.equal(shared_w, shared_n)
+
+
+def test_bridge_permuted_page_map(env):
+    """Non-contiguous page map and padding entries in the block table."""
+    step = env.new_step(LONG_CTX, 6)
+    nblk = (LONG_CTX + 6 + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
+    g = torch.Generator(device=env.device).manual_seed(3)
+    perm = torch.randperm(env.kvm.num_blocks, generator=g, device=env.device)[:nblk].to(torch.int32)
+    bt = torch.full_like(step.block_table, -1)
+    bt[0, :nblk] = perm
+    bt[0, 17] = -1  # a hole inside the used range
+    step.block_table = bt
+    topk = make_topk(6, LONG_CTX + 6, bt.shape[1], env.device, seed=5)
+    topk[0, :8] = torch.arange(17 * TOKENS_PER_BLOCK, 17 * TOKENS_PER_BLOCK + 8, device=env.device)
+    for grouped, layers in ((True, [0, 1, 2, 3]), (False, [1])):
+        shared_n = torch.full((8, TOPK), -5, dtype=torch.int32, device=env.device)
+        shared_w = shared_n.clone()
+        want = env.native_bridge(step, topk, layers, shared_n, grouped)
+        got = env.rope_append_bridge(step, topk, layers, shared_w, grouped)
+        assert torch.equal(got, want)
+        assert torch.equal(shared_w, shared_n)
+        assert (want == -1).any()
+
+
+@pytest.mark.parametrize("acc_dtype", [torch.int32, torch.int64, None])
+def test_bridge_mtp_last_accepted_copy(env, acc_dtype):
+    """Draft step 0: MTP last-accepted row stash (native ``_mtp_last_accepted_rows`` + copy)."""
+    next_n, num_gen = 6, 1
+    step = env.new_step(LONG_CTX, next_n)
+    topk = make_topk(next_n, LONG_CTX + next_n, step.block_table.shape[1], env.device, seed=9)
+    for accepted in (1, 3, 6, 0, 9):
+        acc = (
+            None
+            if acc_dtype is None
+            else torch.tensor([accepted], dtype=acc_dtype, device=env.device)
+        )
+        md = type("Md", (), {"mtp_num_accepted": acc})()
+        rows = Indexer._mtp_last_accepted_rows(None, topk, md, 0, num_gen, next_n)
+        shared_n = torch.full((4, TOPK), -5, dtype=torch.int32, device=env.device)
+        shared_n[0 : rows.shape[0], : rows.shape[1]].copy_(rows)
+        want = env.native_bridge(step, topk, [3], None, False)
+
+        shared_w = torch.full((4, TOPK), -5, dtype=torch.int32, device=env.device)
+        got = env.rope_append_bridge(
+            step,
+            topk,
+            [3],
+            shared_w,
+            False,
+            copy_mode=1,
+            next_n=next_n,
+            num_accepted=acc,
+            copy_rows=num_gen,
+        )
+        assert torch.equal(got, want)
+        assert torch.equal(shared_w, shared_n), f"accepted={accepted}"
+
+
+def test_bridge_draft_reuse_copy(env):
+    """Draft steps 1..4: reuse copy shared -> top-k buffer, remap from the shared rows."""
+    step = env.new_step(LONG_CTX + 3, 1)
+    shared = make_topk(4, LONG_CTX + 4, step.block_table.shape[1], env.device, seed=13)
+    buf_n = torch.full((8, TOPK), -5, dtype=torch.int32, device=env.device)
+    buf_n[0:1, :] = shared[:1, :]
+    want = env.native_bridge(step, buf_n[0:1], [3], None, False)
+    buf_w = torch.full((8, TOPK), -5, dtype=torch.int32, device=env.device)
+    got = env.rope_append_bridge(step, shared[:1], [3], buf_w, False, num_rows=1, copy_rows=1)
+    assert torch.equal(got, want)
+    assert torch.equal(buf_w, buf_n)
+
+
+# Multi-request batches. ``PAD`` is the CUDA-graph padding dummy request, created and repeated
+# in the batch exactly as ``CUDAGraphRunner`` pads a generation batch. b33 crosses the 32-request
+# chunk of the cu_kv_seqlens scan.
+BATCHES = {
+    "b2": [LONG_CTX, 2049],
+    "b3pad": [131072 + 13, 7, PAD],
+    "b33": [LONG_CTX] + [37 * i * i + 5 * i + 1 for i in range(1, 31)] + [PAD, PAD],
+}
+assert len(BATCHES["b33"]) == 33
+
+
+def _expected_cu(step):
+    kv = step.kv_lens()
+    cu_q = [NUM_HEADS * step.rows * i for i in range(step.num_seqs + 1)]
+    cu_kv = [sum(kv[:i]) for i in range(step.num_seqs + 1)]
+    return cu_q, cu_kv
+
+
+@pytest.mark.parametrize("with_scatter", [True, False], ids=["full_layer", "shared_layer"])
+@pytest.mark.parametrize("rows", [6, 1])
+@pytest.mark.parametrize("batch", list(BATCHES))
+def test_multi_request_rope_matches_native(env, batch, rows, with_scatter):
+    step = env.new_batch(BATCHES[batch], rows)
+    assert step.num_seqs == len(BATCHES[batch]) and any(step.is_pad) == (PAD in BATCHES[batch])
+    inp = make_batch_inputs(
+        step,
+        env.device,
+        scale=2.0,
+        seed=len(BATCHES[batch]) * 17 + rows * 3 + int(with_scatter) * 2,
+    )
+    out_n = _run_rope_pair(env, step, 3 if rows == 6 else 0, inp, with_scatter)
+    # Oracle sanity: per-request scheduler offsets, the scan carries across 32-request chunks.
+    cu_q, cu_kv = _expected_cu(step)
+    n = step.num_seqs + 1
+    assert out_n.cu_q[:n].tolist() == cu_q and out_n.cu_kv[:n].tolist() == cu_kv
+    assert (out_n.cu_q[n:] == -7).all() and (out_n.cu_kv[n:] == -7).all()
+
+
+@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "per_layer"])
+@pytest.mark.parametrize("rows", [6, 1])
+@pytest.mark.parametrize("batch", list(BATCHES))
+def test_multi_request_bridge_matches_native(env, batch, rows, grouped):
+    step = env.new_batch(BATCHES[batch], rows)
+    topk = make_batch_topk(step, env.device, seed=rows * 7 + int(grouped))
+    layers = [0, 1, 2, 3] if grouped else [2]
+    shared_n = torch.full((step.num_tokens + 8, TOPK), -5, dtype=torch.int32, device=env.device)
+    shared_w = shared_n.clone()
+    want = env.native_bridge(step, topk, layers, shared_n, grouped)
+    got = env.rope_append_bridge(step, topk, layers, shared_w, grouped)
+    torch.cuda.synchronize()
+    assert got.shape == want.shape
+    assert torch.equal(got, want)
+    assert torch.equal(shared_w, shared_n)
+    # Non-vacuous: the result depends on each row's request (different block-table rows).
+    req0 = torch.zeros_like(step.req_idx)
+    scale, off = env.kvm.get_primary_pool_page_index_params(layers[0])
+    as_req0 = torch.ops.trtllm.convert_req_index_to_global(
+        req0, step.block_table, topk, TOKENS_PER_BLOCK, TOPK, scale * TOKENS_PER_BLOCK, off
+    )
+    first = want[0] if grouped else want
+    assert not torch.equal(first[step.rows :], as_req0[step.rows :])
+
+
+@pytest.mark.parametrize("acc_dtype", [torch.int32, torch.int64, None])
+@pytest.mark.parametrize("batch", list(BATCHES))
+def test_multi_request_mtp_last_accepted_copy(env, batch, acc_dtype):
+    """Draft step 0 with several generation requests: stash row i * next_n + clamp(acc_i - 1)."""
+    next_n = 6
+    step = env.new_batch(BATCHES[batch], next_n)
+    num_gen = step.num_seqs
+    topk = make_batch_topk(step, env.device, seed=41)
+    pattern = [1, 3, 6, 0, 9, 2, 5, 4]
+    accepted = [pattern[(i * 5 + 1) % len(pattern)] for i in range(num_gen)]
+    acc = None if acc_dtype is None else torch.tensor(accepted, dtype=acc_dtype, device=env.device)
+    md = type("Md", (), {"mtp_num_accepted": acc})()
+    rows = Indexer._mtp_last_accepted_rows(None, topk, md, 0, num_gen, next_n)
+    assert rows.shape[0] == num_gen
+    shared_n = torch.full((num_gen + 4, TOPK), -5, dtype=torch.int32, device=env.device)
+    shared_n[0:num_gen, :].copy_(rows)
+    want = env.native_bridge(step, topk, [3], None, False)
+    shared_w = torch.full((num_gen + 4, TOPK), -5, dtype=torch.int32, device=env.device)
+    got = env.rope_append_bridge(
+        step,
+        topk,
+        [3],
+        shared_w,
+        False,
+        copy_mode=1,
+        next_n=next_n,
+        num_accepted=acc,
+        copy_rows=num_gen,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(got, want)
+    assert torch.equal(shared_w, shared_n), f"accepted={accepted}"
+
+
+@pytest.mark.parametrize("batch", list(BATCHES))
+def test_multi_request_draft_reuse_copy(env, batch):
+    """Draft steps 1..4 with several requests: one shared row per request."""
+    step = env.new_batch(BATCHES[batch], 1)
+    n = step.num_seqs
+    shared = make_batch_topk(step, env.device, seed=43, rows=1)
+    buf_n = torch.full((n + 8, TOPK), -5, dtype=torch.int32, device=env.device)
+    buf_n[0:n, :] = shared[:n, :]
+    want = env.native_bridge(step, buf_n[0:n], [3], None, False)
+    buf_w = torch.full((n + 8, TOPK), -5, dtype=torch.int32, device=env.device)
+    got = env.rope_append_bridge(step, shared[:n], [3], buf_w, False, num_rows=n, copy_rows=n)
+    torch.cuda.synchronize()
+    assert torch.equal(got, want)
+    assert torch.equal(buf_w, buf_n)
+
+
+GRAPH_BATCHES = {"single": [LONG_CTX], "b3pad": BATCHES["b3pad"], "b33": BATCHES["b33"]}
+
+
+@pytest.mark.parametrize("batch", list(GRAPH_BATCHES))
+def test_cuda_graph_replay(env, batch):
+    """Capture both rope_kv_append ops once and replay with new values and advancing cache lengths.
+
+    Live requests advance by ``rows`` per replay as in successive verify steps;
+    padding entries keep the dummy's fixed length, as padded replays do.
+    """
+    rows, spec = 6, GRAPH_BATCHES[batch]
+    alloc = [c if c is PAD else c + 4 * rows for c in spec]
+    step = env.new_batch(alloc, rows)
+    base = [c if pad else c - 4 * rows for c, pad in zip(step.ctx_lens, step.is_pad)]
+
+    def kv_at(i):
+        return [b + rows if pad else b + (i + 1) * rows for b, pad in zip(base, step.is_pad)]
+
+    inp = make_batch_inputs(step, env.device, seed=21)
+    out_w = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
+    topk = make_batch_topk(step, env.device, seed=22, kv_lens=kv_at(0))
+    shared_w = torch.full((step.num_tokens + 8, TOPK), -5, dtype=torch.int32, device=env.device)
+
+    def set_step(i):
+        step.set_kv_lens(kv_at(i))
+        new = make_batch_inputs(step, env.device, seed=100 + i)
+        for dst, src in (
+            (inp.fused_q, new.fused_q),
+            (inp.q, new.q),
+            (inp.latent, new.latent),
+            (inp.k_fp8, new.k_fp8),
+            (inp.k_scale, new.k_scale),
+        ):
+            dst.copy_(src)
+        topk.copy_(make_batch_topk(step, env.device, seed=200 + i, kv_lens=kv_at(i)))
+
+    set_step(0)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):  # warm up outside capture
+        env.rope_append_rope(step, 2, inp, out_w, True)
+        env.rope_append_bridge(step, topk, [0, 1, 2, 3], shared_w, True)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        env.rope_append_rope(step, 2, inp, out_w, True)
+        out_static = env.rope_append_bridge(step, topk, [0, 1, 2, 3], shared_w, True)
+
+    for i in range(1, 4):
+        set_step(i)
+        torch.cuda.synchronize()
+        snap0 = env.snapshot()
+        out_n = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
+        env.native_scatter(step, 2, inp)
+        env.native_rope(step, 2, inp, out_n)
+        shared_n = shared_w.clone()
+        want = env.native_bridge(step, topk, [0, 1, 2, 3], shared_n, True)
+        torch.cuda.synchronize()
+        native_pools = env.snapshot()
+        env.restore(snap0)
+        graph.replay()
+        torch.cuda.synchronize()
+        _assert_pools_equal(env, native_pools)
+        _assert_outputs_equal(out_w, out_n)
+        assert torch.equal(out_static, want)
+        assert torch.equal(shared_w, shared_n)
+    assert NUM_HEADS == 8
+
+
+# --------------------------------------------------------------------------------------------
+# Delayed-producer stress: in the captured graph a PDL-launched producer that releases its
+# dependents at entry and writes the per-step metadata ~40 us later runs immediately before each
+# rope_kv_append kernel. Correct rope_kv_append kernels read that metadata only after
+# griddepcontrol.wait.
+
+DELAY_CYCLES = 80_000  # ~40 us at the sm_103 SM clock
+STRESS_REPLAYS = 5
+
+
+def _pdl_enabled() -> bool:
+    return os.environ.get("TRTLLM_ENABLE_PDL", "0") == "1"
+
+
+def _delayed_producer_run(env):
+    """Return, per replay, the names of rope_kv_append outputs that differ from the native oracle.
+
+    Also returns, per replay, whether a control consumer that reads the cache
+    lengths *before* its own wait (``pdl_producer.early_read``) right after the
+    first producer saw the stale values, i.e. whether the stress opened the
+    window a pre-wait reader would fall into.
+    """
+    rows = 6
+    live = [LONG_CTX, 2049 + 64 * 5]
+    step = env.new_batch([c + rows for c in live] + [PAD], rows)
+    kv_new = step.kv_lens()  # this verify step (the allocated capacity) ...
+    kv_old = [k if pad else k - rows for k, pad in zip(kv_new, step.is_pad)]  # ... and the previous
+    dev = env.device
+
+    # "Old" (stale) and "new" values of every per-step metadata tensor rope_kv_append reads.
+    step.set_kv_lens(kv_old)
+    old = {
+        "kv": step.kv_lens_cuda.clone(),
+        "fp8": step.slot_fp8.clone(),
+        "scale": step.slot_scale.clone(),
+        "offsets": step.block_offsets.clone(),
+    }
+    step.set_kv_lens(kv_new)
+    new = {
+        "kv": step.kv_lens_cuda.clone(),
+        "fp8": step.slot_fp8.clone(),
+        "scale": step.slot_scale.clone(),
+        "offsets": step.block_offsets.clone(),
+    }
+    old["offsets"][:, [0, 1]] = new["offsets"][
+        :, [1, 0]
+    ]  # stale page map: the two live rows swapped
+    new["req"] = step.req_idx.clone()
+    old["req"] = new["req"].clone()
+    old["req"][:rows], old["req"][rows : 2 * rows] = 1, 0  # stale request index of the live rows
+    new["acc"] = torch.tensor([2, 5, 1], dtype=torch.int32, device=dev)
+    old["acc"] = torch.tensor([6, 1, 4], dtype=torch.int32, device=dev)
+    num_acc = old["acc"].clone()
+
+    live_rope = [step.kv_lens_cuda, step.slot_fp8, step.slot_scale, step.block_offsets]
+    new_rope = [new["kv"], new["fp8"], new["scale"], new["offsets"]]
+    old_rope = [old["kv"], old["fp8"], old["scale"], old["offsets"]]
+
+    def set_meta(src_rope, src_req, src_acc):
+        for dst, src in zip(live_rope, src_rope):
+            dst.copy_(src)
+        step.req_idx.copy_(src_req)
+        num_acc.copy_(src_acc)
+
+    inp = make_batch_inputs(step, dev, seed=71)
+    topk = make_batch_topk(step, dev, seed=72, kv_lens=kv_new)
+    layers = [0, 1, 2, 3]
+
+    # Native oracle on the new metadata.
+    set_meta(new_rope, new["req"], new["acc"])
+    snap0 = env.snapshot()
+    out_n = RopeOutputs.make(step.num_tokens, dev, num_seqs=step.num_seqs)
+    env.native_scatter(step, 2, inp)
+    env.native_rope(step, 2, inp, out_n)
+    shared_n = torch.full((step.num_tokens + 8, TOPK), -5, dtype=torch.int32, device=dev)
+    want_grouped = env.native_bridge(step, topk, layers, shared_n, True)
+    md = type("Md", (), {"mtp_num_accepted": num_acc})()
+    stash_rows = Indexer._mtp_last_accepted_rows(None, topk, md, 0, step.num_seqs, rows)
+    stash_n = torch.full((8, TOPK), -5, dtype=torch.int32, device=dev)
+    stash_n[0 : step.num_seqs].copy_(stash_rows)
+    want_stash = env.native_bridge(step, topk, [3], None, False)
+    torch.cuda.synchronize()
+    native_pools = env.snapshot()
+    env.restore(snap0)
+
+    out_w = RopeOutputs.make(step.num_tokens, dev, num_seqs=step.num_seqs)
+    shared_w = torch.full_like(shared_n, -5)
+    stash_w = torch.full_like(stash_n, -5)
+    seen_early = torch.zeros_like(step.kv_lens_cuda)
+    kv_copy = torch.empty_like(step.kv_lens_cuda)
+
+    def body():
+        # Control: a pre-wait reader in the same position as a rope_kv_append kernel.
+        pdl_producer.delayed_write([new["kv"]], [kv_copy], DELAY_CYCLES)
+        pdl_producer.early_read([kv_copy], [seen_early])
+        pdl_producer.delayed_write(new_rope, live_rope, DELAY_CYCLES)
+        env.rope_append_rope(step, 2, inp, out_w, True)
+        pdl_producer.delayed_write([new["req"]], [step.req_idx], DELAY_CYCLES)
+        grouped = env.rope_append_bridge(step, topk, layers, shared_w, True)
+        pdl_producer.delayed_write([new["acc"]], [num_acc], DELAY_CYCLES)
+        stash = env.rope_append_bridge(
+            step,
+            topk,
+            [3],
+            stash_w,
+            False,
+            copy_mode=1,
+            next_n=rows,
+            num_accepted=num_acc,
+            copy_rows=step.num_seqs,
+        )
+        return grouped, stash
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):  # warm up (and build the producer) outside capture
+        body()
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        got_grouped, got_stash = body()
+
+    report = []
+    stale_seen = []
+    for _ in range(STRESS_REPLAYS):
+        env.restore(snap0)
+        set_meta(old_rope, old["req"], old["acc"])
+        kv_copy.copy_(old["kv"])
+        seen_early.fill_(-1)
+        shared_w.fill_(-5)
+        stash_w.fill_(-5)
+        torch.cuda.synchronize()
+        graph.replay()
+        torch.cuda.synchronize()
+        bad = [
+            f"pool{i}"
+            for i, (a, b) in enumerate(zip(env.pools(), native_pools))
+            if not torch.equal(a, b)
+        ]
+        names = ["quant_q", "cu_q", "cu_kv", "counter", "bmm1", "bmm2"]
+        bad += [
+            n
+            for n, a, b in zip(names, out_w.tensors(), out_n.tensors())
+            if not torch.equal(a.view(torch.uint8), b.view(torch.uint8))
+        ]
+        bad += [
+            n
+            for n, a, b in (
+                ("grouped_remap", got_grouped, want_grouped),
+                ("shared_copy", shared_w, shared_n),
+                ("stash_remap", got_stash, want_stash),
+                ("stash_copy", stash_w, stash_n),
+            )
+            if not torch.equal(a, b)
+        ]
+        report.append(bad)
+        stale_seen.append(torch.equal(seen_early, old["kv"]))
+    print(
+        f"delayed producer: rope_kv_append outputs differing from native per replay {report}; "
+        f"control pre-wait reader saw stale cache lengths per replay {stale_seen}"
+    )
+    return report, stale_seen
+
+
+@pytest.fixture(scope="module")
+def pdl_producer_ok():
+    if not pdl_producer.available():
+        pytest.skip("cannot build the test PDL producer")
+    return True
+
+
+def test_pdl_delayed_producer(env, pdl_producer_ok):
+    """rope_kv_append is correct with a PDL-chained metadata writer right before each of its kernels.
+
+    The control consumer that reads before its wait must see the stale values
+    in at least one replay, so the stress does open the window (with PDL on).
+    """
+    report, stale_seen = _delayed_producer_run(env)
+    assert all(not bad for bad in report), report
+    if _pdl_enabled():
+        assert any(stale_seen), "the control pre-wait reader never saw stale metadata"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py
new file mode 100644
index 000000000000..13f7293c7e36
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py
@@ -0,0 +1,241 @@
+# 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.
+"""Multi-rank tests of the router_gemm pre-MoE hop with the native MNNVL all-reduce.
+
+Run one pytest process per GPU, every rank executing the same tests:
+
+    srun --mpi=pmix -N2 -n8 python3 -m pytest --import-mode=importlib <this>
+
+With 4 ranks on one node the MNNVL strategy is requested explicitly; with 8
+ranks over two nodes ``AUTO`` selects it as in the served model. The chain
+under test is: stock row-parallel o_proj -> native fused AR + residual +
+RMSNorm -> router_gemm router -> native routing (``noaux_tc_op``), compared bitwise
+against the same chain with the native gate, eagerly and under CUDA graph
+replay, and across ranks.
+"""
+
+import hashlib
+import os
+
+import pytest
+import torch
+from mpi4py import MPI
+from utils.llm_data import llm_models_root
+
+import tensorrt_llm  # noqa: F401
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.router_gemm import (  # noqa: E501
+    pre_moe_hop,
+    router_gemm,
+)
+from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams
+from tensorrt_llm.functional import AllReduceStrategy
+from tensorrt_llm.mapping import Mapping
+
+COMM = MPI.COMM_WORLD
+WORLD = COMM.Get_size()
+RANK = COMM.Get_rank()
+
+pytestmark = pytest.mark.skipif(WORLD < 2, reason="launch one process per GPU with srun --mpi=pmix")
+
+HIDDEN = 6144
+EXPERTS = 256
+HEAD_SLICE = 2048  # 64 heads * 256 v_head / TP8 per rank (o_proj in-features)
+EPS = 1e-5
+
+
+def _bits_equal(a, b):
+    return torch.equal(
+        a.contiguous().view(torch.int32 if a.element_size() == 4 else torch.int16),
+        b.contiguous().view(torch.int32 if b.element_size() == 4 else torch.int16),
+    )
+
+
+def _digest(t):
+    return hashlib.sha256(t.contiguous().cpu().view(torch.uint8).numpy().tobytes()).hexdigest()
+
+
+def _checkpoint_dir():
+    """The GLM-5.2 NVFP4 checkpoint: ``$GLM52_NVFP4_CKPT``, else under ``llm_models_root()``."""
+    root = llm_models_root()
+    default = os.path.join(str(root), "GLM-5.2-NVFP4") if root is not None else ""
+    return os.environ.get("GLM52_NVFP4_CKPT", default)
+
+
+@pytest.fixture(scope="module")
+def env():
+    ngpu = torch.cuda.device_count()
+    local = int(os.environ.get("SLURM_LOCALID", RANK % ngpu))
+    torch.cuda.set_device(local)
+    mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=ngpu)
+    strategy = AllReduceStrategy.AUTO if mapping.is_multi_node() else AllReduceStrategy.MNNVL
+    allreduce = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16)
+    assert allreduce.mnnvl_allreduce is not None, "MNNVL AR not engaged"
+    return mapping, allreduce
+
+
+@pytest.fixture(scope="module")
+def layer_weights(env):
+    """Real gate/norm/o_proj weights of one MoE layer if the checkpoint is set."""
+    from safetensors import safe_open
+
+    ckpt = _checkpoint_dir()
+    dev = torch.device("cuda", torch.cuda.current_device())
+    if ckpt and os.path.exists(os.path.join(ckpt, "model.safetensors.index.json")):
+        import json
+
+        with open(os.path.join(ckpt, "model.safetensors.index.json")) as f:
+            wmap = json.load(f)["weight_map"]
+
+        def load(name):
+            with safe_open(os.path.join(ckpt, wmap[name]), framework="pt") as f:
+                return f.get_tensor(name)
+
+        pre = "model.layers.40"
+        o_proj = load(f"{pre}.self_attn.o_proj.weight")
+        return {
+            "source": "real",
+            "gate": load(f"{pre}.mlp.gate.weight").to(dev),
+            "bias": load(f"{pre}.mlp.gate.e_score_correction_bias").to(dev).to(torch.bfloat16),
+            "norm": load(f"{pre}.post_attention_layernorm.weight").to(dev),
+            "o_proj_local": o_proj[:, RANK * HEAD_SLICE : (RANK + 1) * HEAD_SLICE]
+            .contiguous()
+            .to(dev),
+        }
+    g = torch.Generator(device=dev)
+    g.manual_seed(5)
+    return {
+        "source": "synthetic",
+        "gate": (torch.randn(EXPERTS, HIDDEN, generator=g, device=dev) * 0.02).to(torch.bfloat16),
+        "bias": torch.zeros(EXPERTS, dtype=torch.bfloat16, device=dev),
+        "norm": (1 + 0.1 * torch.randn(HIDDEN, generator=g, device=dev)).to(torch.bfloat16),
+        "o_proj_local": (torch.randn(HIDDEN, HEAD_SLICE, generator=g, device=dev) * 0.02).to(
+            torch.bfloat16
+        ),
+    }
+
+
+def _inputs(rows, seed):
+    """Per-rank attention output slice and the shared residual."""
+    dev = torch.device("cuda", torch.cuda.current_device())
+    g = torch.Generator(device=dev)
+    g.manual_seed(seed * 1000 + RANK)
+    attn = torch.randn(rows, HEAD_SLICE, generator=g, device=dev).to(torch.bfloat16)
+    g.manual_seed(seed * 1000 + 999)
+    residual = (torch.randn(rows, HIDDEN, generator=g, device=dev) * 4).to(torch.bfloat16)
+    return attn, residual
+
+
+def _native_chain(allreduce, w, attn, residual):
+    partial = torch.nn.functional.linear(attn, w["o_proj_local"])
+    hidden, res = allreduce(
+        partial,
+        all_reduce_params=AllReduceParams(
+            fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+            residual=residual,
+            norm_weight=w["norm"],
+            eps=EPS,
+            trigger_completion_at_end=False,
+        ),
+    )
+    logits = torch.ops.trtllm.dsv3_router_gemm_op(
+        hidden, w["gate"].t(), bias=None, out_dtype=torch.float32
+    )
+    return hidden, res, logits
+
+
+def _router_gemm_chain(allreduce, w, attn, residual):
+    partial = torch.nn.functional.linear(attn, w["o_proj_local"])
+    hidden, res, logits = pre_moe_hop(allreduce, partial, residual, w["norm"], EPS, w["gate"])
+    assert logits is not None
+    return hidden, res, logits
+
+
+def _route(logits, bias):
+    return torch.ops.trtllm.noaux_tc_op(logits, bias, 1, 1, 8, 2.5)
+
+
+def test_weight_source(layer_weights):
+    """Real checkpoint weights are used whenever the checkpoint is set."""
+    if os.path.exists(os.path.join(_checkpoint_dir(), "model.safetensors.index.json")):
+        assert layer_weights["source"] == "real"
+
+
+@pytest.mark.parametrize("rows", [1, 3, 6, 16])
+def test_hop_bitwise_and_rank_agreement(env, layer_weights, rows):
+    _, allreduce = env
+    attn, residual = _inputs(rows, seed=rows)
+    h_n, r_n, l_n = _native_chain(allreduce, layer_weights, attn, residual)
+    h_p, r_p, l_p = _router_gemm_chain(allreduce, layer_weights, attn, residual)
+    torch.cuda.synchronize()
+    assert _bits_equal(h_p, h_n) and _bits_equal(r_p, r_n)
+    assert _bits_equal(l_p, l_n)
+    v_n, i_n = _route(l_n, layer_weights["bias"])
+    v_p, i_p = _route(l_p, layer_weights["bias"])
+    assert torch.equal(i_p, i_n) and _bits_equal(v_p, v_n)
+    digests = COMM.allgather((_digest(l_p), _digest(i_p)))
+    assert all(d == digests[0] for d in digests), digests
+
+
+def test_hop_graph_replay(env, layer_weights):
+    """Capture native and router_gemm chains (with the MoE-style aux fork) and replay."""
+    _, allreduce = env
+    rows = 6
+    attn, residual = _inputs(rows, seed=77)
+    aux = torch.cuda.Stream()
+    ev0, ev1 = torch.cuda.Event(), torch.cuda.Event()
+    w_shared = torch.randn(512, HIDDEN, device="cuda").to(torch.bfloat16)
+
+    def body(chain):
+        hidden, res, logits = chain(allreduce, layer_weights, attn, residual)
+        ev0.record()
+        ids = _route(logits, layer_weights["bias"])[1]
+        with torch.cuda.stream(aux):
+            ev0.wait()
+            shared = torch.nn.functional.linear(hidden, w_shared)
+            ev1.record()
+        ev1.wait()
+        return hidden, res, logits, ids, shared
+
+    router_gemm(
+        torch.zeros(rows, HIDDEN, dtype=torch.bfloat16, device="cuda"), layer_weights["gate"]
+    )
+    outs = {}
+    graphs = {}
+    for name, chain in (("native", _native_chain), ("router_gemm", _router_gemm_chain)):
+        body(chain)
+        torch.cuda.synchronize()
+        COMM.Barrier()
+        graph = torch.cuda.CUDAGraph()
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s), torch.cuda.graph(graph, stream=s):
+            outs[name] = body(chain)
+        torch.cuda.current_stream().wait_stream(s)
+        graphs[name] = graph
+    torch.cuda.synchronize()
+    COMM.Barrier()
+    for step in range(30):
+        a, r = _inputs(rows, seed=1000 + step)
+        attn.copy_(a)
+        residual.copy_(r)
+        graphs["router_gemm"].replay()
+        graphs["native"].replay()
+        torch.cuda.synchronize()
+        for x, y in zip(outs["router_gemm"], outs["native"]):
+            assert _bits_equal(x, y), f"step {step}"
+        if step % 10 == 0:
+            digests = COMM.allgather(_digest(outs["router_gemm"][2]))
+            assert all(d == digests[0] for d in digests)
+    COMM.Barrier()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py
new file mode 100644
index 000000000000..37abab2187ff
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py
@@ -0,0 +1,446 @@
+# 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.
+"""Single-GPU correctness of the router_gemm against the native oracle.
+
+The oracle is the native ``torch.ops.trtllm.dsv3_router_gemm_op`` that
+``DeepseekV3Gate`` runs for this model. The router_gemm kernel reproduces its
+arithmetic, so every comparison below is bitwise (FP32 bit patterns), not a
+tolerance.
+"""
+
+import ast
+import json
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+from utils.llm_data import llm_models_root
+
+import tensorrt_llm  # noqa: F401  (loads the native torch ops)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.router_gemm import (  # noqa: E501
+    hop,
+    router_gemm,
+    router_gemm_supported,
+    router_logits,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or not 100
+    <= torch.cuda.get_device_capability()[0] * 10 + torch.cuda.get_device_capability()[1]
+    < 110,
+    reason="needs an SM100-family GPU",
+)
+
+HIDDEN = 6144
+EXPERTS = 256
+TOP_K = 8
+ROUTED_SCALING = 2.5
+
+# GLM-5.2 NVFP4 checkpoint directory. Router (mlp.gate) and
+# post_attention_layernorm weights are BF16 there and are read tensor by tensor;
+# tests that need them skip when it is absent. Activations are always synthetic
+# (real per-layer activations need a captured fixture).
+CKPT_ENV = "GLM52_NVFP4_CKPT"
+
+
+@pytest.fixture(scope="module")
+def real_layer_weights():
+    """Return ``fn(layer) -> dict`` of real gate/bias/norm weights, or skip."""
+    root = llm_models_root()
+    default = os.path.join(str(root), "GLM-5.2-NVFP4") if root is not None else ""
+    ckpt = os.environ.get(CKPT_ENV, default)
+    index = os.path.join(ckpt or "", "model.safetensors.index.json")
+    if not ckpt or not os.path.exists(index):
+        pytest.skip("GLM-5.2 checkpoint not available")
+    from safetensors import safe_open
+
+    with open(index) as f:
+        weight_map = json.load(f)["weight_map"]
+    cache = {}
+
+    def load(name):
+        with safe_open(os.path.join(ckpt, weight_map[name]), framework="pt") as f:
+            return f.get_tensor(name)
+
+    def get(layer):
+        if layer not in cache:
+            pre = f"model.layers.{layer}"
+            cache[layer] = {
+                "gate": load(f"{pre}.mlp.gate.weight"),
+                "bias": load(f"{pre}.mlp.gate.e_score_correction_bias"),
+                "post_norm": load(f"{pre}.post_attention_layernorm.weight"),
+            }
+        return cache[layer]
+
+    return get
+
+
+def native_logits(hidden, weight):
+    return torch.ops.trtllm.dsv3_router_gemm_op(
+        hidden, weight.t(), bias=None, out_dtype=torch.float32
+    )
+
+
+def assert_bitwise(actual, expected):
+    assert actual.shape == expected.shape and actual.dtype == expected.dtype
+    a_bits = actual.contiguous().view(torch.int32)
+    e_bits = expected.contiguous().view(torch.int32)
+    mismatch = (a_bits != e_bits).sum().item()
+    if mismatch:
+        diff = (actual.double() - expected.double()).abs().max().item()
+        raise AssertionError(
+            f"{mismatch}/{actual.numel()} logits differ in bits (max abs diff {diff:.3e})"
+        )
+
+
+def normed_activation(rows, hidden, gen, norm_weight=None, scale=1.0):
+    """Synthetic post-RMSNorm activation, BF16 (as the AR/norm emits)."""
+    x = torch.randn((rows, hidden), generator=gen, device="cuda") * scale
+    x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5)
+    if norm_weight is not None:
+        x = x * norm_weight.float()
+    return x.to(torch.bfloat16)
+
+
+def random_weight(hidden, gen, scale=0.02):
+    return (torch.randn((EXPERTS, hidden), generator=gen, device="cuda") * scale).to(torch.bfloat16)
+
+
+@pytest.fixture
+def gen():
+    g = torch.Generator(device="cuda")
+    g.manual_seed(1234)
+    return g
+
+
+@pytest.mark.parametrize("hidden_size", [6144, 7168, 4096])
+@pytest.mark.parametrize("rows", list(range(1, 17)))
+def test_bitwise_matches_native_random(rows, hidden_size, gen):
+    weight = random_weight(hidden_size, gen)
+    for scale in (1.0, 37.0):
+        hidden = normed_activation(rows, hidden_size, gen, scale=scale)
+        ok, reason = router_gemm_supported(hidden, weight)
+        if not ok:
+            # Only the staged activation can exceed shared memory (K=7168).
+            assert hidden_size == 7168 and "shared memory" in reason
+            assert router_logits(hidden, weight) is None
+            continue
+        assert_bitwise(router_gemm(hidden, weight), native_logits(hidden, weight))
+
+
+@pytest.mark.parametrize("experts_per_cta", [1, 2, 4, 8])
+@pytest.mark.parametrize("use_pdl", [True, False])
+@pytest.mark.parametrize("prefetch_weights", [True, False])
+@pytest.mark.parametrize("paired_fma", [True, False])
+@pytest.mark.parametrize("rows", [5, 6])
+def test_bitwise_schedule_variants(
+    experts_per_cta, use_pdl, prefetch_weights, paired_fma, rows, gen
+):
+    weight = random_weight(HIDDEN, gen)
+    hidden = normed_activation(rows, HIDDEN, gen)
+    out = router_gemm(
+        hidden,
+        weight,
+        experts_per_cta=experts_per_cta,
+        use_pdl=use_pdl,
+        prefetch_weights=prefetch_weights,
+        paired_fma=paired_fma,
+    )
+    assert_bitwise(out, native_logits(hidden, weight))
+
+
+@pytest.mark.parametrize("layer", [3, 4, 40, 77])
+@pytest.mark.parametrize("rows", [1, 6, 16])
+def test_bitwise_matches_native_real_weights(layer, rows, real_layer_weights, gen):
+    """Real GLM-5.2 router/norm weights; synthetic activations."""
+    w = real_layer_weights(layer)
+    weight = w["gate"].cuda()
+    norm_weight = w["post_norm"].cuda()
+    hidden = normed_activation(rows, HIDDEN, gen, norm_weight=norm_weight)
+    assert_bitwise(router_gemm(hidden, weight), native_logits(hidden, weight))
+
+
+def test_bitwise_special_values(gen):
+    """Zeros, signed zeros, cancellation, wide range and underflowing products."""
+    weight = random_weight(HIDDEN, gen)
+    rows = []
+    rows.append(torch.zeros(HIDDEN, device="cuda"))
+    rows.append(-torch.zeros(HIDDEN, device="cuda"))
+    alt = torch.ones(HIDDEN, device="cuda")
+    alt[1::2] = -1
+    rows.append(alt * 3.0)
+    wide = torch.randn(HIDDEN, generator=gen, device="cuda")
+    wide = wide * torch.exp2(torch.randint(-60, 60, (HIDDEN,), generator=gen, device="cuda"))
+    rows.append(wide)
+    rows.append(torch.randn(HIDDEN, generator=gen, device="cuda") * 1e-30)  # tiny
+    rows.append(torch.randn(HIDDEN, generator=gen, device="cuda") * 1e30)
+    hidden = torch.stack(rows).to(torch.bfloat16)
+    weight_tiny = weight.clone()
+    weight_tiny[:16] = (weight_tiny[:16].float() * 1e-20).to(torch.bfloat16)
+    for w in (weight, weight_tiny):
+        for epc, paired in ((1, True), (1, False), (2, True)):
+            assert_bitwise(
+                router_gemm(hidden, w, experts_per_cta=epc, paired_fma=paired),
+                native_logits(hidden, w),
+            )
+    neg_zero = torch.full((1, HIDDEN), -0.0, dtype=torch.bfloat16, device="cuda")
+    assert_bitwise(router_gemm(neg_zero, -weight.abs()), native_logits(neg_zero, -weight.abs()))
+
+
+def native_route(logits, bias):
+    return torch.ops.trtllm.noaux_tc_op(logits, bias, 1, 1, TOP_K, ROUTED_SCALING)
+
+
+@pytest.mark.parametrize("bias_dtype", [torch.float32, torch.bfloat16])
+def test_routing_decisions_identical_with_ties(bias_dtype, gen):
+    """Exact logit ties (duplicated expert rows) route identically."""
+    weight = random_weight(HIDDEN, gen)
+    weight[17] = weight[200]
+    weight[33] = weight[34]
+    bias = torch.zeros(EXPERTS, dtype=bias_dtype, device="cuda")
+    hidden = normed_activation(16, HIDDEN, gen)
+    ours = router_gemm(hidden, weight)
+    ref = native_logits(hidden, weight)
+    assert_bitwise(ours, ref)
+    assert torch.equal(ours[:, 17].view(torch.int32), ours[:, 200].view(torch.int32))
+    v_ours, i_ours = native_route(ours, bias)
+    v_ref, i_ref = native_route(ref, bias)
+    assert torch.equal(i_ours, i_ref)
+    assert torch.equal(v_ours.view(torch.int32), v_ref.view(torch.int32))
+
+
+def test_routing_decisions_real_weights(real_layer_weights, gen):
+    w = real_layer_weights(40)
+    weight = w["gate"].cuda()
+    bias = w["bias"].cuda().to(torch.bfloat16)
+    norm_weight = w["post_norm"].cuda()
+    for _ in range(8):
+        hidden = normed_activation(6, HIDDEN, gen, norm_weight=norm_weight)
+        ours = router_gemm(hidden, weight)
+        ref = native_logits(hidden, weight)
+        assert_bitwise(ours, ref)
+        assert torch.equal(native_route(ours, bias)[1], native_route(ref, bias)[1])
+
+
+def test_cuda_graph_replay(gen):
+    """Producer -> router inside one graph; replay with changing inputs."""
+    weight = random_weight(HIDDEN, gen)
+    source = normed_activation(6, HIDDEN, gen)
+    hidden = torch.empty_like(source)
+    out_native = torch.empty((6, EXPERTS), dtype=torch.float32, device="cuda")
+    router_gemm(source, weight)  # compile outside capture
+    native_logits(source, weight)
+    graph = torch.cuda.CUDAGraph()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        with torch.cuda.graph(graph, stream=stream):
+            hidden.copy_(source)
+            out_ours = router_gemm(hidden, weight)
+            out_native.copy_(native_logits(hidden, weight))
+    torch.cuda.current_stream().wait_stream(stream)
+    for step in range(20):
+        source.copy_(normed_activation(6, HIDDEN, gen, scale=1.0 + step))
+        graph.replay()
+        torch.cuda.synchronize()
+        assert_bitwise(out_ours, out_native)
+        assert_bitwise(out_ours, native_logits(source, weight))
+
+
+def test_padded_rows_do_not_affect_live_rows(gen):
+    weight = random_weight(HIDDEN, gen)
+    hidden = normed_activation(6, HIDDEN, gen)
+    live = 2
+    base = router_gemm(hidden, weight).clone()
+    hidden[live:] = normed_activation(6 - live, HIDDEN, gen, scale=1e3)
+    changed = router_gemm(hidden, weight)
+    assert_bitwise(changed[:live], base[:live])
+    hidden[live:] = 0
+    zero_rows = router_gemm(hidden, weight)[live:]
+    assert torch.all(zero_rows == 0)
+    assert_bitwise(router_gemm(hidden, weight), native_logits(hidden, weight))
+
+
+def test_ineligible_calls_return_none(gen):
+    weight = random_weight(HIDDEN, gen)
+    hidden = normed_activation(6, HIDDEN, gen)
+    cases = {
+        "fp16": (hidden.half(), weight.half()),
+        "rows_17": (normed_activation(17, HIDDEN, gen), weight),
+        "rows_0": (hidden[:0], weight),
+        "experts_128": (hidden, weight[:128]),
+        "hidden_5120": (hidden[:, :5120].contiguous(), weight[:, :5120].contiguous()),
+        "non_contiguous": (hidden.t().contiguous().t(), weight),
+        "misaligned": (
+            torch.empty(6 * HIDDEN + 1, dtype=torch.bfloat16, device="cuda")[1:].view(6, HIDDEN),
+            weight,
+        ),
+        "cpu": (hidden.cpu(), weight.cpu()),
+    }
+    before = hop.engagement_counts()["native_fallback"]
+    for name, (h, w) in cases.items():
+        ok, reason = router_gemm_supported(h, w)
+        assert not ok and reason, name
+        assert router_logits(h, w) is None, name
+    assert hop.engagement_counts()["native_fallback"] == before + len(cases)
+    assert router_logits(hidden, weight) is not None
+
+
+def _fake_model_config(tp_size=8, attention_dp=False, layers=78):
+    return SimpleNamespace(
+        mapping=SimpleNamespace(enable_attention_dp=attention_dp, tp_size=tp_size),
+        pretrained_config=SimpleNamespace(num_hidden_layers=layers),
+    )
+
+
+def _gate(use_cute_dsl_bf16_gemm=False):
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate
+
+    gate = DeepseekV3Gate(
+        HIDDEN,
+        EXPERTS,
+        top_k=TOP_K,
+        n_group=1,
+        topk_group=1,
+        routed_scaling_factor=ROUTED_SCALING,
+        dtype=torch.bfloat16,
+        moe_backend="TRTLLM",
+        use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm,
+    )
+    return gate.cuda()
+
+
+def test_make_moe_router_domain(gen):
+    cfg = _fake_model_config()
+    gate = _gate()
+    with torch.no_grad():
+        gate.weight.copy_(random_weight(HIDDEN, gen))
+    route = hop.make_moe_router(gate, 3, cfg)
+    assert route is not None
+    hidden = normed_activation(6, HIDDEN, gen)
+    assert_bitwise(route(hidden), gate(hidden))
+    assert route(hidden.half()) is None
+    assert hop.make_moe_router(gate, 78, cfg) is None  # MTP layer
+    assert hop.make_moe_router(gate, None, cfg) is None
+    assert hop.make_moe_router(gate, 3, _fake_model_config(attention_dp=True)) is None
+    assert hop.make_moe_router(gate, 3, _fake_model_config(tp_size=1)) is None
+    assert hop.make_moe_router(_gate(use_cute_dsl_bf16_gemm=True), 3, cfg) is None
+
+
+def test_registered_entry_points():
+    for name in ("router_logits", "pre_moe_hop", "make_moe_router"):
+        assert callable(dispatch.get_impl("router_gemm", name))
+
+
+def _enclosing_conditions(func, target):
+    """Conjuncts of the ``if`` tests whose body contains ``target``.
+
+    A ``target`` inside an ``else`` branch yields an ``<else>`` entry, so it
+    never matches a required condition.
+    """
+
+    def visit(node, conditions):
+        if node is target:
+            return conditions
+        if isinstance(node, ast.If):
+            test = node.test
+            parts = (
+                test.values
+                if isinstance(test, ast.BoolOp) and isinstance(test.op, ast.And)
+                else [test]
+            )
+            for branch, extra in (
+                (node.body, [ast.unparse(p) for p in parts]),
+                (node.orelse, ["<else>"]),
+            ):
+                for child in branch:
+                    found = visit(child, conditions + extra)
+                    if found is not None:
+                        return found
+            return None
+        for child in ast.iter_child_nodes(node):
+            found = visit(child, conditions)
+            if found is not None:
+                return found
+        return None
+
+    return visit(func, [])
+
+
+def _check_router_call_site(source, class_name):
+    """The call-site semantics of ``<class_name>.compute_routed_output``.
+
+    One router_gemm router call, guarded by the router's conditions, before one native
+    gate call that runs only when no logits exist (the fallback). The call site also
+    takes the post_attention_hop's logits as the ``router_logits`` argument; router_gemm
+    must then not overwrite them.
+    """
+    cls = next(
+        node
+        for node in ast.parse(source).body
+        if isinstance(node, ast.ClassDef) and node.name == class_name
+    )
+    func = next(
+        node
+        for node in cls.body
+        if isinstance(node, ast.FunctionDef) and node.name == "compute_routed_output"
+    )
+    calls = {}
+    for callee in ("self.router_gemm", "self.gate"):
+        found = [
+            node
+            for node in ast.walk(func)
+            if isinstance(node, ast.Call) and ast.unparse(node.func) == callee
+        ]
+        assert len(found) == 1, (callee, len(found))
+        calls[callee] = found[0]
+    router_gemm_call, gate_call = calls["self.router_gemm"], calls["self.gate"]
+    assert router_gemm_call.lineno < gate_call.lineno
+
+    router_guard = {"self.router_gemm is not None", "not use_dp_padding"}
+    router_gemm_guard = set(_enclosing_conditions(func, router_gemm_call))
+    assert router_guard <= router_gemm_guard <= router_guard | {"router_logits is None"}, (
+        router_gemm_guard
+    )
+    if "router_logits" in [arg.arg for arg in func.args.args]:
+        assert "router_logits is None" in router_gemm_guard, router_gemm_guard
+    gate_guard = _enclosing_conditions(func, gate_call)
+    assert gate_guard == ["router_logits is None"], gate_guard
+
+    assigned = {
+        ast.unparse(node.value)
+        for node in ast.walk(func)
+        if isinstance(node, ast.Assign)
+        and any(isinstance(t, ast.Name) and t.id == "router_logits" for t in node.targets)
+    }
+    assert assigned <= {"None", ast.unparse(router_gemm_call), ast.unparse(gate_call)}, assigned
+
+
+def test_router_call_site_in_target():
+    """The target MoE's router call site: one router_gemm selection (in ``GlmMoE.adopt``) and
+    the call-site semantics in ``GlmMoE.compute_routed_output``.
+    """
+    from pathlib import Path
+
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
+
+    source = Path(layers.__file__).read_text()
+    compile(source, "layers.py", "exec")
+    assert source.count('dispatch.is_enabled("router_gemm")') == 1
+    assert source.count("router_logits = self.gate(hidden_states)") == 1
+    _check_router_call_site(source, "GlmMoE")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py
new file mode 100644
index 000000000000..be9453677405
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py
@@ -0,0 +1,100 @@
+# 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.
+"""Test setup for the sparse_mla_decode package.
+
+When the installed TensorRT-LLM build does not contain the sparse_mla_decode kernel yet and
+``GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_SRC`` points at a source checkout, the kernel and its torch op
+are compiled from that checkout so that the tests can run against an existing
+image. Builds that compile the sparse_mla_decode sources into ``th_common`` skip this step.
+"""
+
+import os
+import sys
+import traceback
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+
+def _load_dev_op() -> None:
+    source_root = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_SRC")
+    if not source_root:
+        return
+    import torch
+
+    import tensorrt_llm
+
+    if hasattr(torch.ops.trtllm, "glm_moe_dsa_sparse_mla_decode"):
+        return
+    from torch.utils.cpp_extension import load
+
+    timing = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING") == "1"
+    # Development variants: GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_DEFINES="NAME_A+NAME_B" adds -DNAME_A
+    # -DNAME_B.
+    extra = [
+        d for d in os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_DEFINES", "").split("+") if d
+    ]
+    major, minor = torch.cuda.get_device_capability()
+    # GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_ARCH selects an architecture-specific target such as "103a".
+    arch_name = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_ARCH", f"{major}{minor}")
+    variant = "_".join([arch_name] + (["timing"] if timing else []) + [d.lower() for d in extra])
+    cpp = Path(source_root) / "cpp"
+    libs = Path(tensorrt_llm.__file__).resolve().parent / "libs"
+    build_dir = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_BUILD_DIR")
+    if build_dir:
+        build_dir = os.path.join(build_dir, variant)
+        os.makedirs(build_dir, exist_ok=True)
+    defines = ["-DENABLE_BF16", "-DENABLE_FP8"] + (
+        ["-DGLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING"] if timing else []
+    )
+    defines += [f"-D{d}" for d in extra]
+    # torch's extension defaults disable the half/bfloat16 operators that the
+    # TensorRT-LLM CUDA headers use; the TensorRT-LLM build does not set them.
+    undefines = [
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+        "-U__CUDA_NO_HALF2_OPERATORS__",
+    ]
+    arch = f"-gencode=arch=compute_{arch_name},code=sm_{arch_name}"
+    load(
+        name=f"glm_moe_dsa_sparse_mla_decode_dev_{variant}",
+        sources=[
+            str(cpp / "tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu"),
+            str(cpp / "tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp"),
+        ],
+        extra_include_paths=[str(cpp / "include"), str(cpp)],
+        extra_cflags=["-O3", "-std=c++17", *defines],
+        extra_cuda_cflags=[
+            "-O3",
+            "-std=c++17",
+            "-lineinfo",
+            "--expt-relaxed-constexpr",
+            arch,
+            *defines,
+            *undefines,
+        ],
+        extra_ldflags=[f"-L{libs}", "-ltensorrt_llm", f"-Wl,-rpath,{libs}"],
+        build_directory=build_dir,
+        is_python_module=False,
+        verbose=True,
+    )
+
+
+try:
+    _load_dev_op()
+except RuntimeError:
+    # Keep the native-only tests runnable; the sparse_mla_decode tests skip without the op.
+    traceback.print_exc()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py
new file mode 100644
index 000000000000..42103593ca4f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py
@@ -0,0 +1,397 @@
+# 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.
+"""One DSA ``MLA`` module through the target's call sites, for composition tests of sparse_mla_decode.
+
+Run as a script in a fresh interpreter: ``python sparse_mla_case.py <output.pt>`` with the fast
+paths to run in ``GLM_MOE_DSA_TEST_FAST_PATHS`` (``plan_testlib.FAST_PATHS_ENV``). It builds
+one GLM-shaped DSA MLA layer (8 local heads, TP1, FP8 latent cache, deterministic weights keyed
+by parameter name), adopts it as the target's ``GlmMLA``,
+runs a native context step and then MTP verify steps of 6 rows as decode steps through
+``GlmMLA.forward``, so that every package at the MLA / DSA call sites
+(absorbed_query/shared_indexer_absorbed_query
+projections, rope_kv_append RoPE/append and top-k bridge, attention_output_tail, sparse_mla_decode
+attention) runs as in the model. With ``GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK=fixed`` (default) the
+indexer's generation selection is replaced by a given local top-k, so that every process attends to
+the same rows; with ``GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK=indexer`` the indexer runs unmodified
+(needed when rope_kv_append's top-k bridge is a producer).
+
+Saved: per step the FMHA inputs (FP8 query bytes, bmm scales, global indices,
+KV lengths), the latent attention output and the module output, for an eager
+step, a step whose sparse_mla_decode launch is followed by a device synchronization (no
+programmatic overlap with the next kernel), and replays of a captured CUDA
+graph over alternating inputs (with the global indices each replay used);
+plus the per-package capability decisions.
+"""
+
+import json
+import math
+import os
+import sys
+import zlib
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import plan_testlib  # noqa: E402
+import torch  # noqa: E402
+
+import tensorrt_llm  # noqa: E402
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E402, E501
+    GlmMLA,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E402, E501
+    SparseMlaDecodeFmha,
+)
+from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha  # noqa: E402
+from tensorrt_llm._torch.attention.backends.interface import (  # noqa: E402
+    PositionalEmbeddingParams,
+    RopeParams,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager  # noqa: E402
+from tensorrt_llm._torch.attention.mla import MLA  # noqa: E402
+from tensorrt_llm._torch.metadata import KVCacheParams  # noqa: E402
+from tensorrt_llm._torch.model_config import ModelConfig  # noqa: E402
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention  # noqa: E402
+from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str  # noqa: E402
+from tensorrt_llm.bindings.executor import KvCacheConfig  # noqa: E402
+from tensorrt_llm.functional import PositionEmbeddingType  # noqa: E402
+from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig  # noqa: E402
+from tensorrt_llm.mapping import Mapping  # noqa: E402
+from tensorrt_llm.models.modeling_utils import QuantConfig  # noqa: E402
+from tensorrt_llm.quantization.mode import QuantAlgo  # noqa: E402
+
+HIDDEN = 6144
+HEADS = 8
+Q_LORA = 2048
+KV_LORA = 512
+NOPE = 192
+ROPE = 64
+V_HEAD = 256
+TOP_K = 2048
+ROWS = 6
+TOKENS_PER_BLOCK = 64
+
+
+def _rope_config(max_positions):
+    return SimpleNamespace(
+        max_position_embeddings=max_positions,
+        hidden_size=HIDDEN,
+        num_attention_heads=64,
+        qk_rope_head_dim=ROPE,
+        rope_theta=8.0e6,
+        rope_scaling=None,
+        model_type="deepseek_v32",
+    )
+
+
+def _fill(param: torch.Tensor, name: str) -> None:
+    """Deterministic values keyed by the parameter name (independent of creation order)."""
+    gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(name.encode()))
+    if param.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+        return
+    if param.dim() == 1:
+        values = 1.0 + 0.05 * torch.randn(param.shape, device="cuda", generator=gen)
+    else:
+        fan_in = param.shape[-1]
+        values = torch.randn(param.shape, device="cuda", generator=gen) / math.sqrt(fan_in)
+    with torch.no_grad():
+        param.copy_(values.to(param.dtype))
+
+
+def build(ctx_len: int):
+    max_seq_len = ctx_len + 4 * ROWS
+    max_positions = 1 << max(16, math.ceil(math.log2(max_seq_len + 1)))
+    mapping = Mapping(world_size=1, tp_size=1, rank=0)
+    sparse_config = DeepSeekSparseAttentionConfig(
+        index_n_heads=32, index_head_dim=128, index_topk=TOP_K, skip_indexer_for_short_seqs=False
+    )
+    model_config = ModelConfig(
+        mapping=mapping,
+        sparse_attention_config=sparse_config,
+        pretrained_config=SimpleNamespace(rms_norm_eps=1e-5),
+        quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value),
+    )
+    pos_embd_params = PositionalEmbeddingParams(
+        type=PositionEmbeddingType.yarn,
+        rope=RopeParams.from_config(_rope_config(max_positions)),
+        is_neox=False,
+    )
+    mla = MLA(
+        hidden_size=HIDDEN,
+        num_attention_heads=HEADS,
+        # GLM-5.2: num_key_value_heads == num_attention_heads (sizes o_proj's input).
+        num_key_value_heads=HEADS,
+        qk_nope_head_dim=NOPE,
+        qk_rope_head_dim=ROPE,
+        v_head_dim=V_HEAD,
+        q_lora_rank=Q_LORA,
+        kv_lora_rank=KV_LORA,
+        predicted_tokens_per_seq=1,
+        max_position_embeddings=max_positions,
+        bias=False,
+        pos_embd_params=pos_embd_params,
+        layer_idx=0,
+        dtype=torch.bfloat16,
+        config=model_config,
+        reduce_output=False,
+        rms_norm_eps=1e-5,
+    )
+    # The model's DeepseekV32Attention adds only its constructor to MLA; the case builds the MLA
+    # with GLM shapes directly and adopts it as the model's attention.
+    mla.__class__ = DeepseekV32Attention
+    GlmMLA.adopt(mla)
+    # The model builds the FMHA libraries in ``create_weights`` after adoption; a module built
+    # alone rebuilds them for the library adoption declared.
+    for backend in (mla.mha, mla.mqa):
+        if backend is not None and backend.prepended_fmha_lib_classes:
+            backend.update_quant_config(backend.quant_config)
+    mla = mla.cuda()
+    # Eager module path (the custom-op path resolves layers through the model engine).
+    mla.register_to_config = False
+    indexer = mla.mqa.indexer
+    indexer.cuda()
+    for name, param in list(mla.named_parameters()) + [
+        ("indexer." + n, p) for n, p in indexer.named_parameters()
+    ]:
+        _fill(param, name)
+    with torch.no_grad():
+        kv_b = mla.kv_b_proj.weight.view(HEADS, NOPE + V_HEAD, KV_LORA)
+        mla.kv_b_proj.weight.copy_(
+            torch.cat(
+                [kv_b[:, :NOPE].reshape(-1, KV_LORA), kv_b[:, NOPE:].reshape(-1, KV_LORA)], dim=0
+            )
+        )
+        k_nope = mla.kv_b_proj.weight[: HEADS * NOPE].view(HEADS, NOPE, KV_LORA)
+        v = mla.kv_b_proj.weight[HEADS * NOPE :].view(HEADS, V_HEAD, KV_LORA)
+        # In place, so that any packed view a package took at construction stays valid.
+        mla.k_b_proj_trans.copy_(k_nope.transpose(1, 2))
+        mla.v_b_proj.copy_(v)
+    indexer.post_load_weights()
+    mla.post_load_weights()
+
+    blocks = (max_seq_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
+    cache = DSACacheManager(
+        KvCacheConfig(max_tokens=blocks * TOKENS_PER_BLOCK, enable_block_reuse=False),
+        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=KV_LORA + ROPE,
+        tokens_per_block=TOKENS_PER_BLOCK,
+        max_seq_len=max_seq_len,
+        max_batch_size=1,
+        mapping=mapping,
+        dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
+        sparse_attn_config=sparse_config,
+        model_config=model_config,
+    )
+    cache.add_dummy_requests([0], [ctx_len])
+    return mla, cache, sparse_config, mapping
+
+
+def metadata(mla, cache, sparse_config, mapping, ctx_len, cached, rows, generation):
+    meta = type(mla.mqa).Metadata(
+        seq_lens=torch.tensor([rows], dtype=torch.int),
+        request_ids=[0],
+        max_num_requests=1,
+        num_contexts=0 if generation else 1,
+        prompt_lens=[ctx_len],
+        max_num_tokens=max(rows, ROWS),
+        kv_cache_manager=cache,
+        kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[cached]),
+        mapping=mapping,
+        sparse_attention_config=sparse_config,
+    )
+    if generation:
+        meta.update_spec_dec_param(
+            batch_size=1,
+            is_spec_decoding_enabled=True,
+            is_spec_dec_tree=False,
+            is_spec_dec_dynamic_tree=False,
+            max_draft_len=ROWS - 1,
+            max_total_draft_tokens=ROWS - 1,
+        )
+    meta.prepare()
+    return meta
+
+
+def local_topk(ctx_len: int, seed: int) -> torch.Tensor:
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    topk = torch.full((ROWS, TOP_K), -1, dtype=torch.int32)
+    for row in range(ROWS):
+        visible = ctx_len + row + 1
+        recent = torch.arange(visible - 32, visible)
+        rest = torch.randperm(visible - 32, generator=gen)[: TOP_K - 32]
+        topk[row] = torch.cat([recent, rest]).to(torch.int32)
+    return topk.cuda()
+
+
+class Capture:
+    """Records the FMHA inputs and the latent output of each attention call."""
+
+    def __init__(self):
+        self.calls = []
+        self.record_inputs = True
+        self.sync_after_sparse_decode = False
+        self._originals = {}
+
+    def install(self):
+        for cls in (SparseMlaDecodeFmha, FallbackFmha):
+            original = cls.forward
+            self._originals[cls] = original
+
+            def forward(fmha, q, k, v, md, args, _original=original, _cls=cls):
+                entry = {"library": _cls.__name__, "output": args.output}
+                if self.record_inputs:
+                    sparse = args.sparse_runtime_params
+                    entry.update(
+                        quant_q=args.quant_q_buffer.clone(),
+                        bmm1=args.mla_bmm1_scale.clone(),
+                        bmm2=args.mla_bmm2_scale.clone(),
+                        indices=sparse.sparse_attn_indices.clone(),
+                        kv_lens=md.kv_lens_cuda_runtime[md.num_contexts :].clone(),
+                    )
+                _original(fmha, q, k, v, md, args)
+                if self.sync_after_sparse_decode and _cls is SparseMlaDecodeFmha:
+                    torch.cuda.synchronize()
+                self.calls.append(entry)
+
+            cls.forward = forward
+
+
+def main(out_path: str) -> None:
+    plan_testlib.select_from_env(dispatch)
+    ctx_len = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_CTX", "32768"))
+    replays = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_REPLAYS", "8"))
+    topk_mode = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK", "fixed")
+    assert topk_mode in ("fixed", "indexer"), topk_mode
+    torch.manual_seed(0)
+    mla, cache, sparse_config, mapping = build(ctx_len)
+    gen = torch.Generator(device="cuda").manual_seed(1234)
+    hidden_ctx = torch.randn(ctx_len, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16)
+    hidden_gen = [
+        torch.randn(ROWS, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16) for _ in range(2)
+    ]
+    topks = [local_topk(ctx_len, 77), local_topk(ctx_len, 78)]
+    pool = cache.get_unique_primary_pool()
+
+    ctx_meta = metadata(mla, cache, sparse_config, mapping, ctx_len, 0, ctx_len, False)
+    with torch.no_grad():
+        mla.forward(torch.arange(ctx_len, device="cuda")[None], hidden_ctx, ctx_meta)
+    torch.cuda.synchronize()
+    for _ in range(ROWS):
+        cache.impl.add_token(0)
+    gen_meta = metadata(mla, cache, sparse_config, mapping, ctx_len, ctx_len, ROWS, True)
+    positions = torch.arange(ctx_len, ctx_len + ROWS, device="cuda")[None]
+
+    capture = Capture()
+    capture.install()
+    static_hidden = hidden_gen[0].clone()
+    static_topk = topks[0].clone()
+    if topk_mode == "fixed":
+        mla.mqa.indexer.forward_from_projected = Mock(return_value=static_topk)
+
+    def step(index):
+        static_hidden.copy_(hidden_gen[index])
+        static_topk.copy_(topks[index])
+        with torch.no_grad(), dispatch.step(dispatch.Phase.DECODE):
+            return mla.forward(positions, static_hidden, gen_meta)
+
+    result = {"eager": [], "serialized": [], "graph": []}
+    # Eager steps (inputs recorded), each input set once.
+    for index in (0, 1):
+        out = step(index).clone()
+        torch.cuda.synchronize()
+        call = capture.calls[-1]
+        result["eager"].append(
+            {
+                "input": index,
+                "module_output": out.cpu(),
+                "latent": call["output"].clone().cpu(),
+                "library": call["library"],
+                "quant_q": call["quant_q"].cpu(),
+                "bmm1": call["bmm1"].cpu(),
+                "bmm2": call["bmm2"].cpu(),
+                "indices": call["indices"].cpu(),
+                "kv_lens": call["kv_lens"].cpu(),
+            }
+        )
+    # The attended pool rows of each input, for an FP64 oracle in the parent.
+    pool_rows = pool.view(torch.uint8).reshape(-1, KV_LORA + ROPE)
+    for entry in result["eager"]:
+        idx = entry["indices"].long().clamp_min(0).cuda()
+        entry["rows"] = pool_rows[idx.reshape(-1)].reshape(ROWS, TOP_K, -1).cpu()
+    # Same inputs with the attention kernel completed before the next kernel is launched.
+    capture.sync_after_sparse_decode = True
+    for index in (0, 1):
+        out = step(index).clone()
+        torch.cuda.synchronize()
+        result["serialized"].append(
+            {
+                "input": index,
+                "module_output": out.cpu(),
+                "latent": capture.calls[-1]["output"].clone().cpu(),
+                "indices": capture.calls[-1]["indices"].cpu(),
+            }
+        )
+    capture.sync_after_sparse_decode = False
+    # A captured CUDA graph of one verify step, replayed over alternating inputs. The
+    # recorded indices are clones made inside the graph, refreshed by every replay.
+    step(0)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        with torch.no_grad(), dispatch.step(dispatch.Phase.DECODE):
+            graph_out = mla.forward(positions, static_hidden, gen_meta)
+    graph_call = capture.calls[-1]
+    for it in range(replays):
+        index = it % 2
+        static_hidden.copy_(hidden_gen[index])
+        static_topk.copy_(topks[index])
+        graph.replay()
+        torch.cuda.synchronize()
+        result["graph"].append(
+            {
+                "input": index,
+                "module_output": graph_out.clone().cpu(),
+                "latent": graph_call["output"].clone().cpu(),
+                "indices": graph_call["indices"].cpu(),
+            }
+        )
+    result["capability"] = {
+        f"{feature}/{key}": value for (feature, key), value in dispatch._CAPABILITY_LOGGED.items()
+    }
+    result["features"] = sorted(dispatch.enabled_features())
+    result["topk_mode"] = topk_mode
+    result["libraries"] = sorted({c["library"] for c in capture.calls})
+    torch.save(result, out_path)
+    print(
+        json.dumps(
+            {
+                "features": result["features"],
+                "topk_mode": topk_mode,
+                "libraries": result["libraries"],
+                "capability": result["capability"],
+            }
+        )
+    )
+    cache.shutdown()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py
new file mode 100644
index 000000000000..9aa3c4ad3a6d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py
@@ -0,0 +1,1026 @@
+# 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.
+"""Native DSA generation harness and FP64 oracles for the sparse_mla_decode tests.
+
+The harness drives the ``DSATrtllmAttention`` generation path exactly as
+``forward_absorption_generation`` does: ``mla_rope_generation`` (RoPE, FP8 query,
+latent append, device scales) followed by ``TrtllmAttention.forward``. Each layer and its
+indexer are adopted as the target's ``GlmDsaBackend`` / ``GlmIndexer`` and every generation
+forward is a decode step, so the predictor runs the target's (with no fast path selected, the
+built-in one). The
+indexer is either replaced by a fixed local top-k (``forward``) or runs
+unmodified on given indexer projections (``forward_native_indexer``). The FMHA
+library list of the layer decides whether the native library or sparse_mla_decode computes
+the attention.
+"""
+
+import contextlib
+import json
+import math
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Callable, Dict, Iterator, List, Optional, Sequence, Tuple
+from unittest.mock import Mock
+
+import torch
+
+import tensorrt_llm
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
+    GlmDsaBackend,
+    GlmIndexer,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
+    SparseMlaDecodeFmha,
+)
+from tensorrt_llm._torch.attention.backends.interface import (
+    AttentionInputType,
+    MLAParams,
+    PositionalEmbeddingParams,
+    RopeParams,
+)
+from tensorrt_llm._torch.attention.backends.sparse.dsa import DSABackendForwardArgs, DSACacheManager
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import (
+    transform_local_topk_and_prepare_pool_view,
+)
+from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
+from tensorrt_llm.bindings.executor import KvCacheConfig
+from tensorrt_llm.functional import PositionEmbeddingType
+from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+NUM_HEADS = 8
+KV_LORA_RANK = 512
+ROPE_DIM = 64
+NOPE_DIM = 192
+V_HEAD_DIM = 256
+HEAD_DIM = KV_LORA_RANK + ROPE_DIM
+Q_LORA_RANK = 2048
+HIDDEN_SIZE = 6144
+TOP_K = 2048
+TOKENS_PER_BLOCK = 64
+E4M3_MAX = 448.0
+
+
+@dataclass
+class _GlmRopeConfig:
+    """Rotary settings of GLM-5.2 (``rope_theta`` 8e6, no scaling)."""
+
+    max_position_embeddings: int
+    hidden_size: int = 6144
+    num_attention_heads: int = 64
+    qk_rope_head_dim: int = ROPE_DIM
+    rope_theta: float = 8.0e6
+    rope_scaling: Optional[dict] = None
+    model_type: str = "deepseek_v32"
+
+
+@dataclass
+class StepBuffers:
+    """Per-step tensors produced by the native ``mla_rope_generation``."""
+
+    fused_q: torch.Tensor
+    q_pe: torch.Tensor
+    latent: torch.Tensor
+    cu_q_seqlens: torch.Tensor
+    cu_kv_seqlens: torch.Tensor
+    scheduler_counter: torch.Tensor
+    bmm1_scale: torch.Tensor
+    bmm2_scale: torch.Tensor
+    quant_q: torch.Tensor
+
+
+class DsaGenHarness:
+    """One generation step of native DSA layers over an FP8 latent pool.
+
+    Args:
+        ctx_lens: Cached tokens per sequence before this step.
+        seq_len_q: Query rows per sequence (1 + draft tokens).
+        num_layers: Attention layers sharing the pool.
+        index_topk: DSA top-k of the sparse config.
+        kv_dtype: KV cache dtype (FP8 for the sparse_mla_decode domain).
+        seed: Seed of the random pool contents.
+        latent_std: Standard deviation of the random latent rows.
+        fill_random: Fill the whole pool with random rows (instead of running a context step).
+        defer_generation: Leave room for ``run_context`` before ``start_generation``.
+        sparse_overrides: Extra ``DeepSeekSparseAttentionConfig`` fields (indexer kernels).
+        mtp_draft_len: Draft length announced to the generation metadata, as the
+            speculative-decoding runner does (``update_spec_dec_param``); 0 leaves it unset.
+        ragged_verify_lens: Query rows per sequence of a ragged-verification step (one
+            entry per sequence; replaces ``seq_len_q``). The generation metadata then
+            enables ragged verification, and ``TrtllmAttention`` presents one row per
+            query token to the RoPE and FMHA calls.
+    """
+
+    def __init__(
+        self,
+        ctx_lens: Sequence[int],
+        seq_len_q: int,
+        num_layers: int = 1,
+        index_topk: int = TOP_K,
+        kv_dtype: torch.dtype = torch.float8_e4m3fn,
+        seed: int = 0,
+        latent_std: float = 1.0,
+        fill_random: bool = True,
+        defer_generation: bool = False,
+        sparse_overrides: Optional[dict] = None,
+        mtp_draft_len: int = 0,
+        ragged_verify_lens: Optional[Sequence[int]] = None,
+    ):
+        self.ctx_lens = [int(x) for x in ctx_lens]
+        self.num_seqs = len(self.ctx_lens)
+        self.ragged = ragged_verify_lens is not None
+        if self.ragged:
+            self.verify_lens = [int(x) for x in ragged_verify_lens]
+            assert len(self.verify_lens) == self.num_seqs and min(self.verify_lens) >= 1
+            self.seq_len_q = max(self.verify_lens)
+        else:
+            self.seq_len_q = int(seq_len_q)
+            self.verify_lens = [self.seq_len_q] * self.num_seqs
+        self.num_rows = sum(self.verify_lens)
+        # Sequence and position within the sequence's query rows, per generation row.
+        self.row_seq = [s for s, n in enumerate(self.verify_lens) for _ in range(n)]
+        self.row_pos = [p for n in self.verify_lens for p in range(n)]
+        self.device = torch.device("cuda")
+        self.mtp_draft_len = int(mtp_draft_len)
+        self.sparse_config = DeepSeekSparseAttentionConfig(
+            index_n_heads=32,
+            index_head_dim=128,
+            index_topk=index_topk,
+            skip_indexer_for_short_seqs=False,
+            **(sparse_overrides or {}),
+        )
+        self.attn_cls = get_attention_backend("TRTLLM", self.sparse_config)
+        max_seq_len = max(self.ctx_lens) + 2 * self.seq_len_q
+        max_positions = 1 << max(16, math.ceil(math.log2(max_seq_len + 1)))
+        rope_config = _GlmRopeConfig(max_position_embeddings=max_positions)
+        pos_embd_params = PositionalEmbeddingParams(
+            type=PositionEmbeddingType.yarn, rope=RopeParams.from_config(rope_config), is_neox=False
+        )
+        mla_params = MLAParams(
+            q_lora_rank=Q_LORA_RANK,
+            kv_lora_rank=KV_LORA_RANK,
+            qk_rope_head_dim=ROPE_DIM,
+            qk_nope_head_dim=NOPE_DIM,
+            v_head_dim=V_HEAD_DIM,
+            rope_append=True,
+            predicted_tokens_per_seq=1,
+            hidden_size=HIDDEN_SIZE,
+        )
+        quant_config = (
+            QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value)
+            if kv_dtype == torch.float8_e4m3fn
+            else None
+        )
+        self.layers = [
+            self.attn_cls(
+                layer_idx=i,
+                num_heads=NUM_HEADS,
+                head_dim=HEAD_DIM,
+                num_kv_heads=1,
+                quant_config=quant_config,
+                q_scaling=1.0,
+                pos_embd_params=pos_embd_params,
+                mla_params=mla_params,
+                sparse_attention_config=self.sparse_config,
+                dtype=torch.bfloat16,
+            )
+            for i in range(num_layers)
+        ]
+        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        model_config = ModelConfig(
+            mapping=self.mapping,
+            sparse_attention_config=self.sparse_config,
+            pretrained_config=_PretrainedStub(),
+        )
+        blocks_per_seq = (max_seq_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
+        self.kv_cache_manager = DSACacheManager(
+            KvCacheConfig(
+                max_tokens=blocks_per_seq * TOKENS_PER_BLOCK * self.num_seqs,
+                enable_block_reuse=False,
+            ),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
+            num_layers=num_layers,
+            num_kv_heads=1,
+            head_dim=HEAD_DIM,
+            tokens_per_block=TOKENS_PER_BLOCK,
+            max_seq_len=max_seq_len,
+            max_batch_size=self.num_seqs,
+            mapping=self.mapping,
+            dtype=str_dtype_to_binding(torch_dtype_to_str(kv_dtype)),
+            sparse_attn_config=self.sparse_config,
+            model_config=model_config,
+        )
+        for layer in self.layers:
+            GlmDsaBackend.adopt(layer)
+            if layer.indexer is not None:
+                GlmIndexer.adopt(layer.indexer)
+        self.request_ids = list(range(self.num_seqs))
+        self.kv_cache_manager.add_dummy_requests(self.request_ids, self.ctx_lens)
+        if fill_random:
+            self.fill_pool(seed, latent_std)
+        self.metadata = None
+        if not defer_generation:
+            self.start_generation()
+
+    def run_context(
+        self,
+        layer_idx: int,
+        fused_q: torch.Tensor,
+        q_pe: torch.Tensor,
+        latent: torch.Tensor,
+        topk: Optional[torch.Tensor],
+        indexer_k: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        chunk_tokens: Optional[int] = None,
+        attention: bool = True,
+    ) -> None:
+        """Native context step: RoPE and append the context latents of every sequence.
+
+        ``topk`` replaces the context-phase indexer selection (the context output is not
+        used). ``indexer_k`` = ``(k_fp8, k_scale)`` of the context tokens, when given, is
+        written to the indexer K cache with the native ``_update_k_cache``, as
+        ``_forward_dsa_attn`` does before the attention. With ``chunk_tokens`` (one
+        sequence only) the context runs as chunked prefill: one context step per chunk,
+        each with the previous chunks as cached tokens. With ``attention=False`` only the
+        indexer K cache is written (the latent pool keeps its contents).
+        """
+        if chunk_tokens is None or chunk_tokens >= max(self.ctx_lens):
+            chunks = [(None, None)]
+        else:
+            assert self.num_seqs == 1, "chunked context steps support one sequence"
+            ctx = self.ctx_lens[0]
+            chunks = [(s, min(s + chunk_tokens, ctx)) for s in range(0, ctx, chunk_tokens)]
+        layer = self.layers[layer_idx]
+        for start, end in chunks:
+            if start is None:
+                seq_lens, cached, rows = self.ctx_lens, [0] * self.num_seqs, slice(None)
+            else:
+                seq_lens, cached, rows = [end - start], [start], slice(start, end)
+            ctx_metadata = self.attn_cls.Metadata(
+                seq_lens=torch.tensor(seq_lens, dtype=torch.int),
+                request_ids=self.request_ids,
+                max_num_requests=self.num_seqs,
+                num_contexts=self.num_seqs,
+                prompt_lens=self.ctx_lens,
+                max_num_tokens=sum(seq_lens),
+                kv_cache_manager=self.kv_cache_manager,
+                kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached),
+                mapping=self.mapping,
+                sparse_attention_config=self.sparse_config,
+            )
+            ctx_metadata.prepare()
+            if indexer_k is not None:
+                layer.indexer._update_k_cache(indexer_k[0][rows], indexer_k[1][rows], ctx_metadata)
+            if not attention:
+                torch.cuda.synchronize()
+                continue
+            layer.indexer.forward_from_projected = Mock(return_value=topk[rows])
+            layer.forward(
+                fused_q[rows].clone(),
+                None,
+                None,
+                ctx_metadata,
+                attention_input_type=AttentionInputType.context_only,
+                latent_cache=latent[rows],
+                q_pe=q_pe[rows],
+                sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]),
+            )
+            torch.cuda.synchronize()
+
+    def start_generation(self) -> None:
+        """Allocate this step's rows and build the generation metadata."""
+        for request_id, rows in zip(self.request_ids, self.verify_lens):
+            for _ in range(rows):
+                self.kv_cache_manager.impl.add_token(request_id)
+        # The ragged row buffers are allocated at construction; the per-step widths are
+        # installed before ``prepare`` (they are not a constructor argument).
+        ragged_kwargs = dict(enable_ragged_verification=True) if self.ragged else {}
+        self.metadata = self.attn_cls.Metadata(
+            seq_lens=torch.tensor(self.verify_lens, dtype=torch.int),
+            request_ids=self.request_ids,
+            max_num_requests=self.num_seqs,
+            num_contexts=0,
+            prompt_lens=self.ctx_lens,
+            max_num_tokens=self.num_rows,
+            kv_cache_manager=self.kv_cache_manager,
+            kv_cache_params=KVCacheParams(
+                use_cache=True, num_cached_tokens_per_seq=list(self.ctx_lens)
+            ),
+            mapping=self.mapping,
+            sparse_attention_config=self.sparse_config,
+            **ragged_kwargs,
+        )
+        if self.ragged:
+            # Ragged verification sizes its row buffers from the draft ceiling and needs the
+            # dense spec-decoding metadata disabled.
+            self.metadata.update_spec_dec_param(
+                batch_size=self.num_seqs,
+                is_spec_decoding_enabled=False,
+                is_spec_dec_tree=False,
+                is_spec_dec_dynamic_tree=False,
+                max_draft_len=self.seq_len_q - 1,
+                max_total_draft_tokens=self.seq_len_q - 1,
+            )
+            self.metadata.ragged_verify_lens = list(self.verify_lens)
+        elif self.mtp_draft_len > 0:
+            self.metadata.update_spec_dec_param(
+                batch_size=self.num_seqs,
+                is_spec_decoding_enabled=True,
+                is_spec_dec_tree=False,
+                is_spec_dec_dynamic_tree=False,
+                max_draft_len=self.mtp_draft_len,
+                max_total_draft_tokens=self.mtp_draft_len,
+            )
+        self.metadata.prepare()
+
+    def shutdown(self) -> None:
+        self.kv_cache_manager.shutdown()
+
+    @property
+    def cu_len(self) -> int:
+        """Entries of the cumulative-length buffers: one per presented row plus one.
+
+        ``forward_absorption_generation`` sizes them the same way: per sequence, and
+        per query token under the ragged-verification presentation.
+        """
+        return (self.num_rows if self.ragged else self.num_seqs) + 1
+
+    # ----- pool access -----------------------------------------------------------------
+    def pool(self) -> torch.Tensor:
+        """Whole primary pool as an e4m3 tensor."""
+        pool = self.kv_cache_manager.get_unique_primary_pool()
+        if pool.dtype != torch.float8_e4m3fn and pool.element_size() == 1:
+            pool = pool.view(torch.float8_e4m3fn)
+        return pool
+
+    def pool_rows(self) -> torch.Tensor:
+        """``[num_pool_rows, 576]`` view addressed by global sparse indices."""
+        return self.pool().reshape(-1, HEAD_DIM)
+
+    def pool_base_pointer(self, layer_idx: int = 0) -> int:
+        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
+            pool_base_pointer,
+        )
+
+        return pool_base_pointer(
+            self.metadata, self.layers[layer_idx].get_local_layer_idx(self.metadata)
+        )
+
+    def fill_pool(self, seed: int, latent_std: float) -> None:
+        """Fill every pool row with random finite e4m3 latent/rope values."""
+        pool = self.pool()
+        gen = torch.Generator(device=self.device).manual_seed(seed)
+        flat = pool.view(-1)
+        chunk = 1 << 26
+        for start in range(0, flat.numel(), chunk):
+            end = min(start + chunk, flat.numel())
+            values = torch.randn(end - start, device=self.device, generator=gen) * latent_std
+            flat[start:end] = values.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn)
+
+    # ----- one step ---------------------------------------------------------------------
+    def make_step(self, layer_idx: int, seed: int, q_std: float = 1.0) -> StepBuffers:
+        """Random absorbed query and new latents, then the native ``mla_rope_generation``."""
+        gen = torch.Generator(device=self.device).manual_seed(seed)
+        rows = self.num_rows
+        fused_q = (
+            torch.randn(rows, NUM_HEADS * HEAD_DIM, device=self.device, generator=gen) * q_std
+        ).to(torch.bfloat16)
+        q_pe = (
+            torch.randn(rows, NUM_HEADS, ROPE_DIM, device=self.device, generator=gen) * q_std
+        ).to(torch.bfloat16)
+        latent = torch.randn(rows, HEAD_DIM, device=self.device, generator=gen).to(torch.bfloat16)
+        step = StepBuffers(
+            fused_q=fused_q,
+            q_pe=q_pe,
+            latent=latent,
+            cu_q_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
+            cu_kv_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
+            scheduler_counter=torch.empty(1, dtype=torch.uint32, device=self.device),
+            bmm1_scale=torch.empty(2, dtype=torch.float32, device=self.device),
+            bmm2_scale=torch.empty(1, dtype=torch.float32, device=self.device),
+            quant_q=torch.empty(rows, NUM_HEADS * HEAD_DIM, dtype=torch.uint8, device=self.device),
+        )
+        self.rope_and_quantize(layer_idx, step)
+        return step
+
+    def make_step_from(
+        self, layer_idx: int, fused_q: torch.Tensor, q_pe: torch.Tensor, latent: torch.Tensor
+    ) -> StepBuffers:
+        """Step buffers from given absorbed queries and latents, then ``mla_rope_generation``."""
+        step = StepBuffers(
+            fused_q=fused_q.clone(),
+            q_pe=q_pe.clone(),
+            latent=latent.clone(),
+            cu_q_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
+            cu_kv_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
+            scheduler_counter=torch.empty(1, dtype=torch.uint32, device=self.device),
+            bmm1_scale=torch.empty(2, dtype=torch.float32, device=self.device),
+            bmm2_scale=torch.empty(1, dtype=torch.float32, device=self.device),
+            quant_q=torch.empty(
+                self.num_rows, NUM_HEADS * HEAD_DIM, dtype=torch.uint8, device=self.device
+            ),
+        )
+        self.rope_and_quantize(layer_idx, step)
+        return step
+
+    def rope_and_quantize(self, layer_idx: int, step: StepBuffers) -> None:
+        self.layers[layer_idx].mla_rope_generation(
+            step.fused_q,
+            step.q_pe,
+            step.latent,
+            self.metadata,
+            step.cu_q_seqlens,
+            step.cu_kv_seqlens,
+            step.scheduler_counter,
+            step.bmm1_scale,
+            step.bmm2_scale,
+            step.quant_q,
+        )
+
+    def select_library(self, layer_idx: int, use_sparse_mla: bool) -> None:
+        """Put sparse_mla_decode first in (or remove it from) the layer's FMHA library list."""
+        layer = self.layers[layer_idx]
+        manager = layer._fmha_manager
+        libs = [lib for lib in manager.fmha_libs if not isinstance(lib, SparseMlaDecodeFmha)]
+        if use_sparse_mla:
+            libs.insert(0, SparseMlaDecodeFmha(layer))
+        manager.fmha_libs = libs
+        manager._cache.clear()
+
+    def selected_library(self, layer_idx: int) -> List[str]:
+        return [type(lib).__name__ for lib in self.layers[layer_idx]._fmha_manager._cache.values()]
+
+    def forward(self, layer_idx: int, step: StepBuffers, topk_local: torch.Tensor) -> torch.Tensor:
+        """``TrtllmAttention.forward`` for the generation rows; returns ``[R, 8*512]`` BF16."""
+        layer = self.layers[layer_idx]
+        layer.indexer.forward_from_projected = Mock(return_value=topk_local)
+        return self._forward_generation(layer_idx, step, [])
+
+    def update_indexer_cache(
+        self, layer_idx: int, k_fp8: torch.Tensor, k_scale: torch.Tensor
+    ) -> None:
+        """Native indexer K-cache append of this step's rows (``_update_k_cache``)."""
+        self.layers[layer_idx].indexer._update_k_cache(k_fp8, k_scale, self.metadata)
+
+    def forward_native_indexer(
+        self, layer_idx: int, step: StepBuffers, indexer_intermediates: Sequence[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """``TrtllmAttention.forward`` with the unmodified native indexer.
+
+        ``indexer_intermediates`` is the native ``pre_indexer_proj`` output of this step's
+        rows. Returns the attention output and a copy of the local top-k the native
+        ``Indexer.forward_from_projected`` produced (and the native remap consumed).
+        """
+        indexer = self.layers[layer_idx].indexer
+        indexer.__dict__.pop("forward_from_projected", None)
+        native_select = indexer.forward_from_projected
+        captured = []
+
+        def capture(*args, **kwargs):
+            topk = native_select(*args, **kwargs)
+            captured.append(topk.clone())
+            return topk
+
+        # Also keep the decode scores and lengths the native top-k module received.
+        top_k_module = indexer.top_k
+        top_k_module.__dict__.pop("forward", None)
+        native_top_k = top_k_module.forward
+        self.last_topk_inputs = None
+
+        def capture_top_k(scores, *args, **kwargs):
+            lengths = kwargs.get("sequence_lengths")
+            if isinstance(lengths, torch.Tensor):
+                self.last_topk_inputs = {
+                    "scores": scores.clone(),
+                    "sequence_lengths": lengths.clone(),
+                    "next_n": kwargs.get("next_n"),
+                }
+            return native_top_k(scores, *args, **kwargs)
+
+        indexer.forward_from_projected = capture
+        top_k_module.forward = capture_top_k
+        try:
+            out = self._forward_generation(layer_idx, step, list(indexer_intermediates))
+        finally:
+            indexer.__dict__.pop("forward_from_projected", None)
+            top_k_module.__dict__.pop("forward", None)
+        assert len(captured) == 1, len(captured)
+        return out, captured[0]
+
+    def _forward_generation(
+        self, layer_idx: int, step: StepBuffers, indexer_intermediates: List[torch.Tensor]
+    ) -> torch.Tensor:
+        layer = self.layers[layer_idx]
+        with dispatch.step(dispatch.Phase.DECODE):
+            return layer.forward(
+                step.fused_q,
+                None,
+                None,
+                self.metadata,
+                attention_input_type=AttentionInputType.generation_only,
+                latent_cache=step.latent,
+                q_pe=step.q_pe,
+                cu_q_seqlens=step.cu_q_seqlens,
+                cu_kv_seqlens=step.cu_kv_seqlens,
+                fmha_scheduler_counter=step.scheduler_counter,
+                mla_bmm1_scale=step.bmm1_scale,
+                mla_bmm2_scale=step.bmm2_scale,
+                quant_q_buffer=step.quant_q,
+                sparse_backend_args=DSABackendForwardArgs(
+                    indexer_intermediates=indexer_intermediates
+                ),
+            )
+
+    @contextlib.contextmanager
+    def capture_sparse_mla_inputs(self) -> Iterator[List[dict]]:
+        """Record the global indices, KV lengths and row split of every sparse_mla_decode call (eager only)."""
+        calls: List[dict] = []
+        original = SparseMlaDecodeFmha.forward
+
+        def forward(fmha, q, k, v, metadata, forward_args):
+            gen_lens = metadata.kv_lens_cuda_runtime[metadata.num_contexts :]
+            calls.append(
+                {
+                    "indices": forward_args.sparse_runtime_params.sparse_attn_indices.clone(),
+                    "seq_lens_kv": gen_lens.clone(),
+                    "seq_len_q": q.shape[0] // gen_lens.shape[0],
+                    "ragged": metadata.token_major_gen_view() is not None,
+                }
+            )
+            return original(fmha, q, k, v, metadata, forward_args)
+
+        SparseMlaDecodeFmha.forward = forward
+        try:
+            yield calls
+        finally:
+            SparseMlaDecodeFmha.forward = original
+
+    def global_indices(self, layer_idx: int, topk_local: torch.Tensor) -> torch.Tensor:
+        """Native local -> global pool-row conversion of a local top-k."""
+        layer = self.layers[layer_idx]
+        global_idx, _ = transform_local_topk_and_prepare_pool_view(
+            topk_local, self.metadata, layer.get_local_layer_idx(self.metadata), is_generation=True
+        )
+        return global_idx
+
+    def grouped_global_indices(self, layer_idx: int, topk_local: torch.Tensor) -> torch.Tensor:
+        """The same conversion through the native grouped (cross-layer fan-out) remap op."""
+        from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import (
+            transform_local_topk_and_prepare_pool_view_grouped,
+        )
+
+        local_layer = self.layers[layer_idx].get_local_layer_idx(self.metadata)
+        page_index_scale, layer_offset = (
+            self.metadata.kv_cache_manager.get_primary_pool_page_index_params(local_layer)
+        )
+        layer_ids = torch.tensor([int(layer_offset)], dtype=torch.int32, device=self.device)
+        grouped = transform_local_topk_and_prepare_pool_view_grouped(
+            topk_local, self.metadata, layer_ids, int(page_index_scale), is_generation=True
+        )
+        return grouped[0]
+
+    # ----- top-k builders ---------------------------------------------------------------
+    def causal_len(self, row: int) -> int:
+        """Tokens visible to generation row ``row`` (its own position included)."""
+        return self.ctx_lens[self.row_seq[row]] + self.row_pos[row] + 1
+
+    def list_limits(self) -> List[int]:
+        """Usable index-list positions per row (native rule: the causal length)."""
+        return [self.causal_len(r) for r in range(self.num_rows)]
+
+    def random_topk(self, seed: int, overlap: float = 0.0, recent: int = 0) -> torch.Tensor:
+        """Per-row distinct local token indices; ``-1`` fills rows with fewer tokens.
+
+        Args:
+            overlap: Fraction of each MTP row's selection copied from the previous row.
+            recent: Number of most recent tokens every row always selects.
+        """
+        gen = torch.Generator(device="cpu").manual_seed(seed)
+        topk = torch.full((self.num_rows, TOP_K), -1, dtype=torch.int32)
+        prev = None
+        for row in range(self.num_rows):
+            visible = self.causal_len(row)
+            if visible <= TOP_K:
+                sel = torch.arange(visible, dtype=torch.int64)
+                sel = sel[torch.randperm(visible, generator=gen)]
+            else:
+                must = torch.arange(visible - min(recent, TOP_K), visible, dtype=torch.int64)
+                chosen = set(must.tolist())
+                if prev is not None and overlap > 0 and self.row_pos[row] != 0:
+                    keep = prev[torch.randperm(prev.numel(), generator=gen)[: int(overlap * TOP_K)]]
+                    chosen.update(x for x in keep.tolist() if x < visible)
+                while len(chosen) < TOP_K:
+                    draw = torch.randint(0, visible, (2 * TOP_K,), generator=gen).tolist()
+                    for x in draw:
+                        chosen.add(x)
+                        if len(chosen) == TOP_K:
+                            break
+                sel = torch.tensor(sorted(chosen), dtype=torch.int64)
+                sel = sel[torch.randperm(sel.numel(), generator=gen)]
+            topk[row, : sel.numel()] = sel.to(torch.int32)
+            prev = sel
+        return topk.to(self.device)
+
+
+class _PretrainedStub:
+    rms_norm_eps = 1e-5
+
+
+def negative_inside_causal(
+    indices: torch.Tensor, seq_lens_kv: torch.Tensor, seq_len_q: int
+) -> torch.Tensor:
+    """Per row, the negative entries at index-list positions below the row's causal length.
+
+    This is the input class on which sparse_mla_decode deliberately differs from the native kernel
+    (sparse_mla_decode skips the entry, the native kernel reads the pool row ``-1`` addresses).
+    """
+    rows = indices.shape[0]
+    row = torch.arange(rows, device=indices.device)
+    seq = row // seq_len_q
+    usable = seq_lens_kv.to(indices.device).long()[seq] - (seq_len_q - 1 - row % seq_len_q)
+    positions = torch.arange(indices.shape[1], device=indices.device)
+    inside = positions[None, :] < usable[:, None]
+    return ((indices < 0) & inside).sum(dim=1)
+
+
+# ----- rope_kv_append's top-k bridge as the producer of sparse_mla_decode's indices
+# ------------------------------------
+
+
+def rope_append_selected() -> bool:
+    """Whether this process selects rope_kv_append (whose fused bridge then replaces the native remap)."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    return "rope_kv_append" in dispatch.SELECTOR_PATHS and dispatch.is_enabled("rope_kv_append")
+
+
+@contextlib.contextmanager
+def capture_rope_append_bridge() -> Iterator[List[dict]]:
+    """Record every generation call of rope_kv_append's fused top-k bridge (eager only).
+
+    The DSA predictor looks up rope_kv_append's ``remap_topk_to_global`` entry on every call, so
+    replacing the registered entry sees each call. Engagement is counted here, per call:
+    ``dispatch.log_capability`` keeps only the first decision per key, and the first
+    bridge decision of a process is usually a context batch, which the bridge declines.
+
+    Each record holds the layer, the local top-k the bridge remapped (cloned before the
+    call), the call's metadata, and the global indices it returned (the returned tensor).
+    Yields an empty list when rope_kv_append is not selected.
+    """
+    calls: List[dict] = []
+    if not rope_append_selected():
+        yield calls
+        return
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+
+    original = dispatch.get_impl("rope_kv_append", "remap_topk_to_global")
+
+    def bridge(backend, topk, metadata, packet):
+        if metadata.num_contexts != 0:
+            return original(backend, topk, metadata, packet)
+        source = topk if packet.source is None else packet.source
+        entry = {
+            "layer": backend.layer_idx,
+            "local_layer": backend.get_local_layer_idx(metadata),
+            "topk": source.clone(),
+            "metadata": metadata,
+        }
+        entry["global"] = original(backend, topk, metadata, packet)
+        calls.append(entry)
+        return entry["global"]
+
+    dispatch._REGISTRY["rope_kv_append"]["remap_topk_to_global"] = bridge
+    try:
+        yield calls
+    finally:
+        dispatch._REGISTRY["rope_kv_append"]["remap_topk_to_global"] = original
+
+
+def rope_append_bridge_report(call: dict, seq_lens_kv: torch.Tensor, seq_len_q: int) -> dict:
+    """Producer contract of one recorded bridge call (after the step completed).
+
+    ``equals_native_remap``: the bridge output equals the native local-to-global remap
+    of the same local top-k bit for bit. ``negative_below_causal``: per row, ``-1``
+    entries the bridge emitted at index-list positions below the row's causal length
+    (``seq_lens_kv`` = KV lengths of the generation sequences, this step included).
+    """
+    native, _ = transform_local_topk_and_prepare_pool_view(
+        call["topk"], call["metadata"], call["local_layer"], is_generation=True
+    )
+    out = call["global"]
+    return {
+        "layer": call["layer"],
+        "equals_native_remap": bool(out.shape == native.shape and torch.equal(out, native)),
+        "negative_below_causal": negative_inside_causal(out, seq_lens_kv, seq_len_q).tolist(),
+        "input_negative_below_causal": negative_inside_causal(
+            call["topk"], seq_lens_kv, seq_len_q
+        ).tolist(),
+        "negative_total": int((out < 0).sum().item()),
+    }
+
+
+# ----- CUPTI kernel records ------------------------------------------------------------------
+
+SPARSE_MLA_KERNEL = "sparseMlaDecodeKernel"
+
+
+def is_native_sparse_fmha(name: str) -> bool:
+    """The TRT-LLM-Gen sparse MLA generation kernel (``fmhaSm10x…Sparse…``)."""
+    return "fmhaSm1" in name and "Sparse" in name
+
+
+def kernel_trace(fn: Callable[[], None]) -> List[dict]:
+    """CUPTI records of the kernels ``fn`` enqueues (graph replays included), by start time.
+
+    Uses ``torch.profiler``'s CUDA activity tracing; each record has the kernel name, the
+    stream and the start / end time in microseconds.
+    """
+    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
+        fn()
+        torch.cuda.synchronize()
+    fd, path = tempfile.mkstemp(suffix=".json")
+    os.close(fd)
+    try:
+        prof.export_chrome_trace(path)
+        with open(path) as f:
+            events = json.load(f)["traceEvents"]
+    finally:
+        os.unlink(path)
+    kernels = [
+        {
+            "name": e["name"],
+            "stream": e.get("args", {}).get("stream"),
+            "start": float(e["ts"]),
+            "end": float(e["ts"]) + float(e.get("dur", 0.0)),
+        }
+        for e in events
+        if e.get("cat") == "kernel"
+    ]
+    return sorted(kernels, key=lambda k: k["start"])
+
+
+def attention_neighbours(kernels: List[dict]) -> List[dict]:
+    """For every sparse attention kernel (sparse_mla_decode or native): its neighbours on its stream.
+
+    ``successor_early``: the next kernel of the stream started before the attention
+    kernel ended (it was launched early as a programmatic dependent), by
+    ``successor_lead_us``; ``started_before_predecessor_end``: the attention kernel itself
+    started early, by ``predecessor_lead_us``.
+    """
+    by_stream: Dict[object, List[dict]] = {}
+    for k in kernels:
+        by_stream.setdefault(k["stream"], []).append(k)
+    out = []
+    for stream_kernels in by_stream.values():
+        for i, k in enumerate(stream_kernels):
+            is_sparse_mla = SPARSE_MLA_KERNEL in k["name"]
+            if not (is_sparse_mla or is_native_sparse_fmha(k["name"])):
+                continue
+            prev = stream_kernels[i - 1] if i > 0 else None
+            nxt = stream_kernels[i + 1] if i + 1 < len(stream_kernels) else None
+            out.append(
+                {
+                    "kernel": "sparse_mla_decode" if is_sparse_mla else "native",
+                    "start": k["start"],
+                    "duration_us": k["end"] - k["start"],
+                    "predecessor": prev["name"][:96] if prev else None,
+                    "started_before_predecessor_end": bool(prev and k["start"] < prev["end"]),
+                    "predecessor_lead_us": (prev["end"] - k["start"]) if prev else None,
+                    "successor": nxt["name"][:96] if nxt else None,
+                    "successor_early": bool(nxt and nxt["start"] < k["end"]),
+                    "successor_lead_us": (k["end"] - nxt["start"]) if nxt else None,
+                }
+            )
+    return sorted(out, key=lambda a: a["start"])
+
+
+# ----- FP64 oracles --------------------------------------------------------------------------
+
+
+def e4m3_round(scale: float) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Probability rounding to e4m3 after multiplying by ``scale``."""
+
+    def fn(p: torch.Tensor) -> torch.Tensor:
+        q = (p * scale).clamp(max=E4M3_MAX).to(torch.float32).to(torch.float8_e4m3fn)
+        return q.to(torch.float64) / scale
+
+    return fn
+
+
+def fp16_round(p: torch.Tensor) -> torch.Tensor:
+    return p.to(torch.float16).to(torch.float64)
+
+
+def bf16_round(p: torch.Tensor) -> torch.Tensor:
+    return p.to(torch.bfloat16).to(torch.float64)
+
+
+def oracle_attention(
+    q_bytes: torch.Tensor,
+    pool_rows: torch.Tensor,
+    global_idx: torch.Tensor,
+    bmm1_scale: float,
+    bmm2_scale: float,
+    *,
+    list_limit: Optional[Sequence[int]] = None,
+    prob_round: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+    split: Optional[int] = None,
+    step: Optional[int] = None,
+) -> torch.Tensor:
+    """FP64 sparse MLA reference over the e4m3 query and pool; returns ``[R, 8, 512]``.
+
+    Without ``prob_round`` this is the exact softmax. With it, the probabilities fed to
+    P*V are rounded the way an online-softmax kernel would round them: the index list
+    is cut into ``split``-position chunks merged at the end, each processed in
+    ``step``-position steps against the running maximum; row sums stay unrounded.
+    """
+    rows = q_bytes.shape[0]
+    q = q_bytes.reshape(rows, NUM_HEADS, HEAD_DIM).view(torch.float8_e4m3fn).to(torch.float64)
+    pool_bytes = pool_rows.view(torch.uint8)
+
+    def gather(sel: torch.Tensor) -> torch.Tensor:
+        return pool_bytes[sel].view(torch.float8_e4m3fn).to(torch.float64)
+
+    out = torch.zeros(rows, NUM_HEADS, KV_LORA_RANK, dtype=torch.float64, device=q.device)
+    positions = torch.arange(global_idx.shape[1], device=global_idx.device)
+    for r in range(rows):
+        idx = global_idx[r].to(torch.int64)
+        valid = idx >= 0
+        if list_limit is not None:
+            valid &= positions < list_limit[r]
+        if not bool(valid.any()):
+            continue
+        if prob_round is None:
+            k = gather(idx[valid])
+            s = (q[r] @ k.T) * bmm1_scale
+            p = torch.softmax(s, dim=-1)
+            out[r] = p @ k[:, :KV_LORA_RANK]
+            continue
+        chunk = split or global_idx.shape[1]
+        stride = step or chunk
+        maxes, sums, accs = [], [], []
+        for a in range(0, global_idx.shape[1], chunk):
+            m = torch.full((NUM_HEADS,), -math.inf, dtype=torch.float64, device=q.device)
+            row_sum = torch.zeros(NUM_HEADS, dtype=torch.float64, device=q.device)
+            acc = torch.zeros(NUM_HEADS, KV_LORA_RANK, dtype=torch.float64, device=q.device)
+            for b in range(a, min(a + chunk, global_idx.shape[1]), stride):
+                sel = valid[b : b + stride]
+                if not bool(sel.any()):
+                    continue
+                k = gather(idx[b : b + stride][sel])
+                s = (q[r] @ k.T) * bmm1_scale
+                m_new = torch.maximum(m, s.max(dim=-1).values)
+                alpha = torch.exp(m - m_new)
+                p = torch.exp(s - m_new[:, None])
+                row_sum = row_sum * alpha + p.sum(dim=-1)
+                acc = acc * alpha[:, None] + prob_round(p) @ k[:, :KV_LORA_RANK]
+                m = m_new
+            maxes.append(m)
+            sums.append(row_sum)
+            accs.append(acc)
+        m_all = torch.stack(maxes)
+        global_max = m_all.max(dim=0).values
+        weights = torch.where(
+            torch.isinf(m_all), torch.zeros_like(m_all), torch.exp(m_all - global_max)
+        )
+        total = (weights * torch.stack(sums)).sum(dim=0)
+        out[r] = (weights[:, :, None] * torch.stack(accs)).sum(dim=0) / total[:, None]
+    return out * bmm2_scale
+
+
+def error_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict:
+    """Normalized RMS and max-abs error of ``candidate`` against ``reference``."""
+    c = candidate.reshape(reference.shape).to(torch.float64)
+    diff = c - reference
+    ref_rms = reference.pow(2).mean().sqrt().item()
+    return {
+        "nrms": diff.pow(2).mean().sqrt().item() / max(ref_rms, 1e-30),
+        "max_abs": diff.abs().max().item(),
+        "ref_rms": ref_rms,
+        "nonfinite": int((~torch.isfinite(c)).sum().item()),
+    }
+
+
+# ----- native-versus-sparse_mla_decode comparison
+# -----------------------------------------------------------
+
+
+def record(filename: str, case: str, payload: dict) -> None:
+    """Print one JSON result line and append it to ``$GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS/<filename>``."""
+    import json
+    import os
+
+    line = json.dumps({"case": case, **payload}, sort_keys=True)
+    print(line)
+    results_dir = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS")
+    if results_dir:
+        with open(os.path.join(results_dir, filename), "a") as f:
+            f.write(line + "\n")
+
+
+def run_both(harness: DsaGenHarness, step: StepBuffers, topk: torch.Tensor):
+    """Native and sparse_mla_decode outputs for one step plus the FP64 oracle with the native rules."""
+    harness.select_library(0, use_sparse_mla=False)
+    native = harness.forward(0, step, topk).clone()
+    assert "SparseMlaDecodeFmha" not in harness.selected_library(0)
+    harness.select_library(0, use_sparse_mla=True)
+    sparse_out = harness.forward(0, step, topk).clone()
+    assert harness.selected_library(0) == ["SparseMlaDecodeFmha"], harness.selected_library(0)
+    global_idx = harness.global_indices(0, topk)
+    return native, sparse_out, global_idx, oracle_for(harness, step, global_idx)
+
+
+def usable_rows(harness: DsaGenHarness, global_idx: torch.Tensor) -> torch.Tensor:
+    positions = torch.arange(global_idx.shape[1], device=global_idx.device)
+    limits = torch.tensor(harness.list_limits(), device=global_idx.device)[:, None]
+    return ((global_idx >= 0) & (positions[None, :] < limits)).any(dim=1)
+
+
+def oracle_for(harness: DsaGenHarness, step: StepBuffers, global_idx: torch.Tensor) -> torch.Tensor:
+    """FP64 oracle with the native rules for this step's query and the current pool."""
+    return oracle_attention(
+        step.quant_q,
+        harness.pool_rows(),
+        global_idx,
+        float(step.bmm1_scale[0]),
+        float(step.bmm2_scale[0]),
+        list_limit=harness.list_limits(),
+    )
+
+
+def compare(
+    harness: DsaGenHarness,
+    step: StepBuffers,
+    topk: torch.Tensor,
+    case: str,
+    filename: str = "sparse_mla_correctness.jsonl",
+    **extra,
+) -> dict:
+    """Check sparse_mla_decode against native and the oracle with the frozen error budget."""
+    native, sparse_out, global_idx, oracle = run_both(harness, step, topk)
+    return check_budget(harness, native, sparse_out, global_idx, oracle, case, filename, **extra)
+
+
+def check_budget(
+    harness: DsaGenHarness,
+    native: torch.Tensor,
+    sparse_out: torch.Tensor,
+    global_idx: torch.Tensor,
+    oracle: torch.Tensor,
+    case: str,
+    filename: str = "sparse_mla_correctness.jsonl",
+    **extra,
+) -> dict:
+    """The frozen error budget for one native and one sparse_mla_decode output of the same inputs."""
+    native_err = error_stats(native.double(), oracle)
+    sparse_mla_err = error_stats(sparse_out.double(), oracle)
+    sparse_mla_vs_native = error_stats(sparse_out.double(), native.double().reshape(oracle.shape))
+    record(
+        filename,
+        case,
+        {
+            "native_vs_oracle": native_err,
+            "sparse_mla_vs_oracle": sparse_mla_err,
+            "sparse_mla_vs_native": sparse_mla_vs_native,
+            **extra,
+        },
+    )
+    native_rows = native.reshape(harness.num_rows, -1)
+    sparse_mla_rows = sparse_out.reshape(harness.num_rows, -1)
+    live = usable_rows(harness, global_idx)
+    if bool((~live).any()):
+        torch.testing.assert_close(
+            sparse_mla_rows[~live], native_rows[~live], rtol=0, atol=0, equal_nan=True
+        )
+    assert sparse_mla_err["nonfinite"] == native_err["nonfinite"]
+    assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (sparse_mla_err, native_err)
+    assert sparse_mla_err["max_abs"] <= 2.0 * native_err["max_abs"] + 1e-3, (
+        sparse_mla_err,
+        native_err,
+    )
+    oracle_rows = oracle.reshape(harness.num_rows, -1)
+    assert torch.allclose(sparse_mla_rows[live].double(), oracle_rows[live], atol=0.12, rtol=0.01)
+    return {
+        "native": native,
+        "sparse_mla_decode": sparse_out,
+        "oracle": oracle,
+        "global_idx": global_idx,
+    }
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py
new file mode 100644
index 000000000000..ecaac194346b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py
@@ -0,0 +1,204 @@
+# 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.
+"""Test-only PDL producer and consumer kernels for the sparse_mla_decode stress test.
+
+``delayed_scatter`` is a producer: it allows its dependents to launch at once
+(``griddepcontrol.launch_dependents``), waits for its own predecessor, spins for
+a given time and only then copies a list of segments into the buffers the next
+kernel reads. ``pdl_consumer`` is a consumer: it also allows its dependents to
+launch at once, waits for its predecessor and copies a buffer out, optionally
+after a spin. Both are launched with programmatic stream serialization, so a
+chain producer -> sparse_mla_decode -> consumer -> producer -> ... has no plain kernel
+boundary: every kernel may become resident while its predecessor still runs,
+and only ``griddepcontrol.wait`` orders the memory accesses.
+
+Both kernels record ``%globaltimer`` stamps so that the test can show that the
+early launches actually happened.
+"""
+
+import os
+
+import torch
+
+_CPP_DECLS = r"""
+#include <torch/extension.h>
+void delayed_scatter(torch::Tensor dsts, torch::Tensor srcs, torch::Tensor counts, int64_t spin_ns,
+    torch::Tensor stamps, bool use_pdl, int64_t num_blocks);
+void pdl_consumer(torch::Tensor dst, torch::Tensor src, int64_t spin_ns, torch::Tensor stamps, bool use_pdl);
+"""
+
+_CUDA_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+
+namespace
+{
+
+__device__ __forceinline__ unsigned long long globalTimer()
+{
+    unsigned long long t;
+    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+    return t;
+}
+
+__device__ __forceinline__ void spinFor(long long spinNs)
+{
+    if (spinNs <= 0)
+    {
+        return;
+    }
+    if (threadIdx.x == 0)
+    {
+        unsigned long long const start = globalTimer();
+        while (globalTimer() - start < static_cast<unsigned long long>(spinNs))
+        {
+        }
+    }
+    __syncthreads();
+}
+
+// stamps: [0] producer start (before waiting), [1] producer after its wait, [2] producer end.
+__global__ void delayedScatterKernel(long long const* dsts, long long const* srcs, long long const* counts,
+    int numSegments, long long spinNs, unsigned long long* stamps)
+{
+    unsigned long long const t0 = globalTimer();
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    unsigned long long const t1 = globalTimer();
+    if (blockIdx.x == 0 && threadIdx.x == 0)
+    {
+        stamps[0] = t0;
+        stamps[1] = t1;
+    }
+    spinFor(spinNs);
+    for (int s = 0; s < numSegments; ++s)
+    {
+        uint4* dst = reinterpret_cast<uint4*>(dsts[s]);
+        uint4 const* src = reinterpret_cast<uint4 const*>(srcs[s]);
+        long long const n = counts[s];
+        for (long long i = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x; i < n;
+             i += static_cast<long long>(gridDim.x) * blockDim.x)
+        {
+            dst[i] = src[i];
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0)
+    {
+        atomicMax(stamps + 2, globalTimer());
+    }
+}
+
+// stamps: [0] consumer start (before waiting), [1] consumer after its wait.
+__global__ void pdlConsumerKernel(
+    uint4 const* src, uint4* dst, long long n, long long spinNs, unsigned long long* stamps)
+{
+    unsigned long long const t0 = globalTimer();
+    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    unsigned long long const t1 = globalTimer();
+    if (blockIdx.x == 0 && threadIdx.x == 0)
+    {
+        stamps[0] = t0;
+        stamps[1] = t1;
+    }
+    spinFor(spinNs);
+    for (long long i = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x; i < n;
+         i += static_cast<long long>(gridDim.x) * blockDim.x)
+    {
+        dst[i] = src[i];
+    }
+}
+
+cudaLaunchAttribute pdlAttribute(bool usePdl)
+{
+    cudaLaunchAttribute attr;
+    attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attr.val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0;
+    return attr;
+}
+
+} // namespace
+
+void delayed_scatter(torch::Tensor dsts, torch::Tensor srcs, torch::Tensor counts, int64_t spin_ns,
+    torch::Tensor stamps, bool use_pdl, int64_t num_blocks)
+{
+    TORCH_CHECK(dsts.is_cuda() && srcs.is_cuda() && counts.is_cuda() && stamps.is_cuda());
+    TORCH_CHECK(dsts.scalar_type() == torch::kInt64 && srcs.scalar_type() == torch::kInt64
+        && counts.scalar_type() == torch::kInt64 && stamps.scalar_type() == torch::kInt64);
+    TORCH_CHECK(dsts.numel() == srcs.numel() && dsts.numel() == counts.numel() && stamps.numel() >= 3);
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(static_cast<unsigned>(num_blocks));
+    config.blockDim = dim3(256);
+    config.stream = at::cuda::getCurrentCUDAStream(dsts.get_device());
+    cudaLaunchAttribute attr = pdlAttribute(use_pdl);
+    config.attrs = &attr;
+    config.numAttrs = 1;
+    TORCH_CHECK(cudaLaunchKernelEx(&config, delayedScatterKernel, reinterpret_cast<long long const*>(dsts.data_ptr()),
+                    reinterpret_cast<long long const*>(srcs.data_ptr()),
+                    reinterpret_cast<long long const*>(counts.data_ptr()), static_cast<int>(dsts.numel()),
+                    static_cast<long long>(spin_ns), reinterpret_cast<unsigned long long*>(stamps.data_ptr()))
+        == cudaSuccess);
+}
+
+void pdl_consumer(torch::Tensor dst, torch::Tensor src, int64_t spin_ns, torch::Tensor stamps, bool use_pdl)
+{
+    TORCH_CHECK(dst.is_cuda() && src.is_cuda() && stamps.is_cuda() && stamps.scalar_type() == torch::kInt64);
+    TORCH_CHECK(dst.is_contiguous() && src.is_contiguous() && dst.nbytes() == src.nbytes() && src.nbytes() % 16 == 0);
+    TORCH_CHECK(stamps.numel() >= 2);
+    cudaLaunchConfig_t config{};
+    config.gridDim = dim3(16);
+    config.blockDim = dim3(256);
+    config.stream = at::cuda::getCurrentCUDAStream(dst.get_device());
+    cudaLaunchAttribute attr = pdlAttribute(use_pdl);
+    config.attrs = &attr;
+    config.numAttrs = 1;
+    TORCH_CHECK(cudaLaunchKernelEx(&config, pdlConsumerKernel, reinterpret_cast<uint4 const*>(src.data_ptr()),
+                    reinterpret_cast<uint4*>(dst.data_ptr()), static_cast<long long>(src.nbytes() / 16),
+                    static_cast<long long>(spin_ns), reinterpret_cast<unsigned long long*>(stamps.data_ptr()))
+        == cudaSuccess);
+}
+"""
+
+_MODULE = None
+
+
+def load_pdl_helpers():
+    """Build (once per process) and return the helper extension."""
+    global _MODULE
+    if _MODULE is None:
+        from torch.utils.cpp_extension import load_inline
+
+        major, minor = torch.cuda.get_device_capability()
+        arch = f"{major}{minor}"
+        build_dir = os.path.join(
+            os.environ.get(
+                "GLM_MOE_DSA_SPARSE_MLA_DECODE_HELPER_BUILD_DIR",
+                "/tmp/glm_moe_dsa_sparse_mla_decode_helpers",
+            ),
+            arch,
+        )
+        os.makedirs(build_dir, exist_ok=True)
+        _MODULE = load_inline(
+            name=f"glm_moe_dsa_sparse_mla_decode_pdl_helpers_{arch}",
+            cpp_sources=_CPP_DECLS,
+            cuda_sources=_CUDA_SRC,
+            functions=["delayed_scatter", "pdl_consumer"],
+            extra_cuda_cflags=["-O3", f"-gencode=arch=compute_{arch},code=sm_{arch}"],
+            build_directory=build_dir,
+            verbose=False,
+        )
+    return _MODULE
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py
new file mode 100644
index 000000000000..ca07bd36a9c7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py
@@ -0,0 +1,310 @@
+# 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.
+"""Real-weight sparse_attention contract inputs from GLM-5.2 layer 0 on real text.
+
+Layer 0 is the one target layer whose input is fully determined by the token
+ids (embedding + input RMSNorm), so its attention inputs are exactly the
+model's. The projections use the checkpoint's BF16 attention weights for the
+TP rank 0 head slice; RoPE, the FP8 latent cache and the FP8 query come from
+the native TRT-LLM ops driven by ``DsaGenHarness``. Two top-k sources exist:
+
+* ``indexer_topk``: the native indexer projections (``Indexer.pre_indexer_proj``)
+  with the checkpoint's indexer weights, followed by a *reference* FP32
+  evaluation of the indexer logits ``sum_h w_h * relu(q_h . k) * k_scale`` and an
+  exact top-k (not the native top-k kernels);
+* ``indexer_projections`` + ``DsaGenHarness.forward_native_indexer``: the native
+  indexer K cache written by the native ``_update_k_cache`` and the unmodified
+  native ``Indexer.forward_from_projected`` (paged MQA logits + decode top-k).
+"""
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterable, Optional
+
+import torch
+from utils.llm_data import llm_models_root
+
+import tensorrt_llm
+
+NUM_LOCAL_HEADS = 8
+QK_NOPE = 192
+QK_ROPE = 64
+V_HEAD = 256
+KV_LORA = 512
+NUM_HEADS = 64
+RMS_EPS = 1e-5
+CHECKPOINT_NAME = "GLM-5.2-NVFP4"
+
+
+def checkpoint_dir() -> Optional[Path]:
+    root = llm_models_root()
+    if root is None:
+        return None
+    path = Path(root) / CHECKPOINT_NAME
+    return path if (path / "config.json").is_file() else None
+
+
+def load_tensors(ckpt: Path, names: Iterable[str], device: str = "cuda") -> Dict[str, torch.Tensor]:
+    from safetensors import safe_open
+
+    weight_map = json.loads((ckpt / "model.safetensors.index.json").read_text())["weight_map"]
+    by_shard: Dict[str, list] = {}
+    for name in names:
+        by_shard.setdefault(weight_map[name], []).append(name)
+    tensors = {}
+    for shard, shard_names in by_shard.items():
+        with safe_open(str(ckpt / shard), framework="pt", device=device) as f:
+            for name in shard_names:
+                tensors[name] = f.get_tensor(name)
+    return tensors
+
+
+def real_token_ids(ckpt: Path, count: int) -> torch.Tensor:
+    """Token ids of real text: the installed TensorRT-LLM Python sources, in path order."""
+    from tokenizers import Tokenizer
+
+    tokenizer = Tokenizer.from_file(str(ckpt / "tokenizer.json"))
+    ids = []
+    for path in sorted(Path(tensorrt_llm.__file__).resolve().parent.rglob("*.py")):
+        ids.extend(tokenizer.encode(path.read_text(errors="ignore")).ids)
+        if len(ids) >= count:
+            break
+    if len(ids) < count:
+        raise ValueError(f"only {len(ids)} tokens of text available, {count} requested")
+    return torch.tensor(ids[:count], dtype=torch.long)
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = RMS_EPS) -> torch.Tensor:
+    xf = x.float()
+    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float()).to(x.dtype)
+
+
+@dataclass
+class LayerInputs:
+    """Token-wise projections of one layer for every token of the sequence."""
+
+    hidden: torch.Tensor  # [N, 6144] post input-RMSNorm
+    qr: torch.Tensor  # [N, 2048] post q_a RMSNorm
+    fused_q: torch.Tensor  # [N, 8 * 576] absorbed q_nope | raw q_pe
+    q_pe: torch.Tensor  # [N, 8, 64] raw q_pe
+    latent: torch.Tensor  # [N, 576] normalized latent | raw k_pe
+
+
+def layer_weight_names(layer: int) -> list:
+    prefix = f"model.layers.{layer}."
+    attn = prefix + "self_attn."
+    return [
+        "model.embed_tokens.weight",
+        prefix + "input_layernorm.weight",
+        attn + "q_a_proj.weight",
+        attn + "q_a_layernorm.weight",
+        attn + "q_b_proj.weight",
+        attn + "kv_a_proj_with_mqa.weight",
+        attn + "kv_a_layernorm.weight",
+        attn + "kv_b_proj.weight",
+        attn + "indexer.wq_b.weight",
+        attn + "indexer.wk.weight",
+        attn + "indexer.k_norm.weight",
+        attn + "indexer.k_norm.bias",
+        attn + "indexer.weights_proj.weight",
+    ]
+
+
+@torch.inference_mode()
+def project_layer0(weights: Dict[str, torch.Tensor], ids: torch.Tensor) -> LayerInputs:
+    """Layer-0 attention inputs of TP rank 0 (heads 0-7) for every token."""
+    attn = "model.layers.0.self_attn."
+    hidden = rms_norm(
+        weights["model.embed_tokens.weight"][ids.cuda()],
+        weights["model.layers.0.input_layernorm.weight"],
+    )
+    kv = hidden @ weights[attn + "kv_a_proj_with_mqa.weight"].T
+    compressed = rms_norm(kv[:, :KV_LORA], weights[attn + "kv_a_layernorm.weight"])
+    k_pe = kv[:, KV_LORA:]
+    qr = rms_norm(
+        hidden @ weights[attn + "q_a_proj.weight"].T, weights[attn + "q_a_layernorm.weight"]
+    )
+    q_b = weights[attn + "q_b_proj.weight"][: NUM_LOCAL_HEADS * (QK_NOPE + QK_ROPE)]
+    q = (qr @ q_b.T).view(-1, NUM_LOCAL_HEADS, QK_NOPE + QK_ROPE)
+    q_nope, q_pe = q.split([QK_NOPE, QK_ROPE], dim=-1)
+    kv_b = weights[attn + "kv_b_proj.weight"].view(NUM_HEADS, QK_NOPE + V_HEAD, KV_LORA)
+    w_uk = kv_b[:NUM_LOCAL_HEADS, :QK_NOPE, :]
+    q_absorbed = torch.einsum("nhd,hdk->nhk", q_nope.float(), w_uk.float()).to(torch.bfloat16)
+    fused_q = torch.cat([q_absorbed, q_pe], dim=-1).reshape(
+        -1, NUM_LOCAL_HEADS * (KV_LORA + QK_ROPE)
+    )
+    return LayerInputs(
+        hidden=hidden.contiguous(),
+        qr=qr.contiguous(),
+        fused_q=fused_q.contiguous(),
+        q_pe=q_pe.contiguous(),
+        latent=torch.cat([compressed, k_pe], dim=-1).contiguous(),
+    )
+
+
+def load_indexer_weights(indexer, weights: Dict[str, torch.Tensor], layer: int = 0) -> None:
+    prefix = f"model.layers.{layer}.self_attn.indexer."
+    with torch.no_grad():
+        for name, module_param in (
+            ("wq_b.weight", indexer.wq_b.weight),
+            ("wk.weight", indexer.wk.weight),
+            ("weights_proj.weight", indexer.weights_proj.weight),
+            ("k_norm.weight", indexer.k_norm.weight),
+            ("k_norm.bias", indexer.k_norm.bias),
+        ):
+            module_param.data = (
+                weights[prefix + name].to(device="cuda", dtype=module_param.dtype).contiguous()
+            )
+    indexer.cache_derived_state()
+
+
+@torch.inference_mode()
+def reference_indexer_logits(indexer, inputs: LayerInputs, ctx_len: int, seq_len_q: int) -> list:
+    """FP32 indexer logits ``sum_h w_h * relu(q_h . k) * k_scale`` of each generation row.
+
+    Uses the native indexer projections; row ``r`` covers its ``ctx_len + r + 1`` visible tokens.
+    """
+    num_tokens = ctx_len + seq_len_q
+    positions = torch.arange(num_tokens, dtype=torch.int32, device="cuda")
+    q_fp8, k_fp8, k_scale, weights, _ = indexer.pre_indexer_proj(
+        inputs.qr[:num_tokens], inputs.hidden[:num_tokens], positions
+    )
+    q = q_fp8[ctx_len:num_tokens].float().reshape(seq_len_q, indexer.n_heads, indexer.head_dim)
+    keys = (
+        k_fp8.float().reshape(num_tokens, indexer.head_dim)
+        * k_scale.float().reshape(num_tokens, -1)[:, :1]
+    )
+    head_weights = weights[ctx_len:num_tokens].float().reshape(seq_len_q, indexer.n_heads)
+    logits = []
+    for row in range(seq_len_q):
+        visible = ctx_len + row + 1
+        logits.append(
+            (head_weights[row][:, None] * torch.relu(q[row] @ keys[:visible].T)).sum(dim=0)
+        )
+    return logits
+
+
+@torch.inference_mode()
+def indexer_topk(
+    indexer, inputs: LayerInputs, ctx_len: int, seq_len_q: int, topk: int = 2048
+) -> torch.Tensor:
+    """Reference local top-k of the generation rows: exact top-k of the FP32 reference logits."""
+    out = torch.full((seq_len_q, topk), -1, dtype=torch.int32, device="cuda")
+    for row, logits in enumerate(reference_indexer_logits(indexer, inputs, ctx_len, seq_len_q)):
+        count = min(topk, logits.numel())
+        out[row, :count] = logits.topk(count).indices.to(torch.int32)
+    return out
+
+
+def selection_vs_reference(topk_local: torch.Tensor, logits: list, topk: int = 2048) -> list:
+    """How a selection ranks under the FP32 reference logits, per row.
+
+    ``kth`` is the reference top-k boundary value; ``ties_at_kth`` counts visible tokens whose
+    reference logit equals it; ``selected_at_or_above_kth`` is the fraction of the selection whose
+    reference logit is at least ``kth`` (1.0 for any exact top-k up to tie order), and
+    ``selected_within_1pct`` the same with a tolerance of 1 % of the logit range.
+    """
+    rows = []
+    for r, row_logits in enumerate(logits):
+        count = min(topk, row_logits.numel())
+        kth = row_logits.topk(count).values[-1]
+        ids = topk_local[r][topk_local[r] >= 0].long()
+        chosen = row_logits[ids]
+        span = (row_logits.max() - row_logits.min()).clamp_min(1e-30)
+        rows.append(
+            {
+                "kth": float(kth),
+                "max": float(row_logits.max()),
+                "ties_at_kth": int((row_logits == kth).sum()),
+                "zero_logits": int((row_logits == 0).sum()),
+                "selected_at_or_above_kth": float((chosen >= kth).float().mean()),
+                "selected_within_1pct": float((chosen >= kth - 0.01 * span).float().mean()),
+            }
+        )
+    return rows
+
+
+@torch.inference_mode()
+def indexer_projections(indexer, inputs: LayerInputs, start: int, end: int) -> list:
+    """Native ``pre_indexer_proj`` of tokens ``[start, end)``: ``[q_fp8, k_fp8, k_scale, w, q_scale]``."""
+    positions = torch.arange(start, end, dtype=torch.int32, device="cuda")
+    return list(indexer.pre_indexer_proj(inputs.qr[start:end], inputs.hidden[start:end], positions))
+
+
+def topk_contract(topk_local: torch.Tensor, ctx_len: int) -> list:
+    """Per generation row: where the local top-k holds negative, future and repeated ids.
+
+    Row ``r`` of a verify window sees ``ctx_len + r + 1`` tokens; ``negative_below_causal``
+    counts ``-1`` entries at list positions inside that causal length (where the native
+    sparse MLA kernel reads the row ``-1`` addresses and sparse_mla_decode skips the entry).
+    """
+    rows = []
+    for r in range(topk_local.shape[0]):
+        causal = ctx_len + r + 1
+        row = topk_local[r].long()
+        head = row[: min(causal, row.numel())]
+        valid = head[head >= 0]
+        rows.append(
+            {
+                "causal_len": causal,
+                "negative_below_causal": int((head < 0).sum()),
+                "negative_total": int((row < 0).sum()),
+                "ids_at_or_past_causal": int((row >= causal).sum()),
+                "duplicates": int(valid.numel() - valid.unique().numel()),
+                "valid": int((row >= 0).sum()),
+            }
+        )
+    return rows
+
+
+def selection_vs_kernel_scores(
+    topk_local: torch.Tensor, topk_inputs: dict, reference_logits: list, topk: int = 2048
+) -> list:
+    """The native decode scores the top-k kernel received, against its output and the reference.
+
+    Row ``r`` of a verify window of ``next_n`` rows scores ``seq_len - next_n + r + 1`` tokens
+    (the native decode top-k rule). ``exact_agreement``: overlap of the selection with the exact
+    top-k of those scores; ``scores_reference_corr``: Pearson correlation of the scores with the
+    FP32 reference logits of the same tokens.
+    """
+    scores = topk_inputs["scores"].float()
+    seq_len = int(topk_inputs["sequence_lengths"].reshape(-1)[0])
+    next_n = int(topk_inputs["next_n"])
+    rows = []
+    for r in range(topk_local.shape[0]):
+        visible = seq_len - next_n + r + 1
+        row_scores = scores[r, :visible]
+        count = min(topk, visible)
+        exact = set(row_scores.topk(count).indices.tolist())
+        chosen = set(x for x in topk_local[r].tolist() if x >= 0)
+        ref = reference_logits[r][:visible].float()
+        corr = torch.corrcoef(torch.stack([row_scores, ref]))[0, 1]
+        rows.append(
+            {
+                "scores_shape": list(scores.shape),
+                "visible": visible,
+                "exact_agreement": len(chosen & exact) / max(len(exact), 1),
+                "scores_reference_corr": float(corr),
+            }
+        )
+    return rows
+
+
+def recent_window_topk(num_tokens: int, topk: int = 2048) -> torch.Tensor:
+    """Causal context selection used only to run the native context step."""
+    rows = torch.arange(num_tokens, dtype=torch.int32, device="cuda")[:, None]
+    cols = torch.arange(topk, dtype=torch.int32, device="cuda")[None, :]
+    return (rows - cols).clamp_(min=-1)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py
new file mode 100644
index 000000000000..f590474a50c7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py
@@ -0,0 +1,480 @@
+# 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.
+"""Real GLM-5.2 attention layers through ``GlmMLA.forward``, for sparse_mla_decode composition tests.
+
+Run as a script in a fresh interpreter: ``python sparse_mla_stack_case.py <output.pt>`` with the
+fast paths to run in ``GLM_MOE_DSA_TEST_FAST_PATHS`` (``plan_testlib.FAST_PATHS_ENV``). It uses
+the attention_output_tail tests' ``real_mla_harness.AttentionStack``: the TP8 rank-0 shards of two
+consecutive real target layers (default 2, which runs its own
+indexer, and 3, which reuses layer 2's top-k), a native prefill of
+``GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_CTX`` real-token positions, and the MTP verify step of one request
+(6 rows) with the generation metadata the decoder runner builds (CUDA-graph
+metadata, ``update_spec_dec_param``, ``prepare``). The layers are the target's adopted ones and
+the verify step is a decode step, so ``GlmMLA.forward`` runs through ``register_to_config`` (the
+target's DSA attention op) as in the model and each package at the MLA / DSA call sites runs
+where the model runs it:
+absorbed_query (full-indexer layer) and shared_indexer_absorbed_query (shared-indexer layer) produce
+the query and latent row, rope_kv_append the RoPE / FP8 query / latent append, the indexer-K scatter
+and the top-k bridge, the attention_output_tail (attention_output_tail_prefetch /
+attention_output_into_hop schedules) consumes the attention output as its programmatic dependent,
+and sparse_mla_decode computes the attention. The process stands in for one of eight TP ranks
+(o_proj all-reduce off).
+
+``GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_TOPK=fixed`` (default) replaces the full layer's indexer
+selection by a given local top-k (deterministic indices; the shared layer reuses it through
+the shared top-k buffer); ``indexer`` runs the native indexer.
+
+Saved (``torch.save``):
+
+* ``inputs``: an eager step per input set; per layer the FMHA inputs (FP8
+  query bytes, bmm scales, global indices, KV lengths, attended pool rows),
+  cloned before the attention, plus the latent and module outputs;
+* ``pdl``: the same steps with no recording kernel anywhere, so the kernel
+  after the attention is the one the model runs there;
+* ``serialized``: the same steps with a device synchronization right after
+  each sparse_mla_decode launch (no programmatic overlap with the consumer);
+* ``graph``: replays of one captured step over alternating input sets;
+* ``kernels``: CUPTI records of two eager steps and of graph replays: kernel
+  counts and, per attention kernel, its predecessor and successor on the
+  stream and whether they overlapped it;
+* ``engagement``: per-layer counts of absorbed_query / shared_indexer_absorbed_query packets,
+  rope_kv_append fused RoPE calls and generation bridge calls, the attention_output_tail counters
+  and config, sparse_mla_decode's dispatch accounting, and the logged capability decisions;
+* ``rope_append_bridge``: per generation call of rope_kv_append's bridge in the ``inputs`` steps,
+  whether its output equals the native remap of the same local top-k, the
+  ``-1`` entries it emitted below a row's causal length, and whether the
+  attention received exactly that output.
+"""
+
+import dataclasses
+import json
+import os
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+from unittest.mock import Mock
+
+_HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(_HERE))
+sys.path.insert(0, str(_HERE.parent / "attention_output_tail"))
+sys.path.insert(0, str(_HERE.parent))
+
+import plan_testlib  # noqa: E402
+import real_mla_harness  # noqa: E402
+import torch  # noqa: E402
+from sparse_mla_harness import (  # noqa: E402
+    SPARSE_MLA_KERNEL,
+    attention_neighbours,
+    is_native_sparse_fmha,
+    kernel_trace,
+    rope_append_bridge_report,
+)
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
+    dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E402, E501
+    sparse_mla,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E402, E501
+    SparseMlaDecodeFmha,
+)
+from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha  # noqa: E402
+from tensorrt_llm._torch.distributed import AllReduceParams  # noqa: E402
+
+ROWS = 6
+TOP_K = 2048
+HEAD_DIM = 576
+
+
+def _selected(feature_id: str) -> bool:
+    return feature_id in dispatch.SELECTOR_PATHS and dispatch.is_enabled(feature_id)
+
+
+class Engagement:
+    """Per-call counters of the producer packages, read at their registered entry points."""
+
+    def __init__(self):
+        self.counts = defaultdict(int)
+        self.bridge = []
+        self.record_bridge = False
+        if _selected("absorbed_query"):
+            for name, label in (
+                ("dsa_proj", "absorbed_query"),
+                ("shared_dsa_proj", "shared_indexer_absorbed_query"),
+            ):
+                self._count_packets(name, label)
+        if _selected("rope_kv_append"):
+            rope = dispatch.get_impl("rope_kv_append", "mla_rope_generation")
+
+            def rope_hook(mqa, *args, **kwargs):
+                self.counts[f"rope_append_rope/layer{mqa.layer_idx}"] += 1
+                return rope(mqa, *args, **kwargs)
+
+            dispatch._REGISTRY["rope_kv_append"]["mla_rope_generation"] = rope_hook
+            bridge = dispatch.get_impl("rope_kv_append", "remap_topk_to_global")
+
+            def bridge_hook(backend, topk, metadata, packet):
+                if metadata.num_contexts != 0:
+                    return bridge(backend, topk, metadata, packet)
+                self.counts[f"rope_append_bridge/layer{backend.layer_idx}"] += 1
+                entry = None
+                if self.record_bridge:
+                    source = topk if packet.source is None else packet.source
+                    entry = {
+                        "layer": backend.layer_idx,
+                        "local_layer": backend.get_local_layer_idx(metadata),
+                        "topk": source.clone(),
+                        "metadata": metadata,
+                    }
+                out = bridge(backend, topk, metadata, packet)
+                if entry is not None:
+                    entry["global"] = out
+                    self.bridge.append(entry)
+                return out
+
+            dispatch._REGISTRY["rope_kv_append"]["remap_topk_to_global"] = bridge_hook
+
+    def _count_packets(self, name, label):
+        original = dispatch.get_impl("absorbed_query", name)
+
+        def hook(mla, *args, **kwargs):
+            packet = original(mla, *args, **kwargs)
+            if packet is not None:
+                self.counts[f"{label}/layer{mla.layer_idx}"] += 1
+            return packet
+
+        dispatch._REGISTRY["absorbed_query"][name] = hook
+
+
+class FmhaCapture:
+    """Records each attention call: library, layer, output and (optionally) input clones.
+
+    Only the input clones add kernels, and they run before the attention; the
+    output is kept by reference (which also keeps its memory from being reused)
+    and copied after the step, so nothing is inserted between the attention
+    kernel and its consumer.
+    """
+
+    def __init__(self):
+        self.calls = []
+        self.clone_inputs = False
+        self.sync_after_sparse_decode = False
+        for cls in (SparseMlaDecodeFmha, FallbackFmha):
+            original = cls.forward
+
+            def forward(fmha, q, k, v, md, args, _original=original, _cls=cls):
+                sparse = args.sparse_runtime_params
+                entry = {
+                    "library": _cls.__name__,
+                    "layer": fmha.attn.layer_idx,
+                    "generation": md.num_contexts == 0,
+                    "output": args.output,
+                    "indices_ref": sparse.sparse_attn_indices,
+                }
+                if self.clone_inputs:
+                    entry.update(
+                        quant_q=args.quant_q_buffer.clone(),
+                        bmm1=args.mla_bmm1_scale.clone(),
+                        bmm2=args.mla_bmm2_scale.clone(),
+                        indices=sparse.sparse_attn_indices.clone(),
+                        kv_lens=md.kv_lens_cuda_runtime[md.num_contexts :].clone(),
+                    )
+                _original(fmha, q, k, v, md, args)
+                if self.sync_after_sparse_decode and _cls is SparseMlaDecodeFmha:
+                    torch.cuda.synchronize()
+                self.calls.append(entry)
+
+            cls.forward = forward
+
+
+def local_topk(ctx_len: int, seed: int) -> torch.Tensor:
+    """Distinct local ids per row: the 32 most recent visible tokens plus random older ones."""
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    topk = torch.full((ROWS, TOP_K), -1, dtype=torch.int32)
+    for row in range(ROWS):
+        visible = ctx_len + row + 1
+        recent = torch.arange(visible - 32, visible)
+        rest = torch.randperm(visible - 32, generator=gen)[: TOP_K - 32]
+        topk[row] = torch.cat([recent, rest]).to(torch.int32)
+    return topk.cuda()
+
+
+def summarize_kernels(kernels, steps, layers):
+    names = Counter(k["name"] for k in kernels)
+    neighbours = attention_neighbours(kernels)
+    sparse_kernels = [a for a in neighbours if a["kernel"] == "sparse_mla_decode"]
+    return {
+        "steps": steps,
+        "layers": layers,
+        "sparse_mla_decode": sum(n for name, n in names.items() if SPARSE_MLA_KERNEL in name),
+        "native_sparse_fmha": sum(n for name, n in names.items() if is_native_sparse_fmha(name)),
+        "rope_append_bridge": sum(n for name, n in names.items() if "topkBridgeKernel" in name),
+        "rope_append_rope": sum(
+            n for name, n in names.items() if "ropeAppendScatterKernel" in name
+        ),
+        "native_remap": sum(n for name, n in names.items() if "convertReqIndexToGlobal" in name),
+        "top_kernels": [[name[:120], n] for name, n in names.most_common(48)],
+        "attention": neighbours,
+        "successors": Counter(a["successor"] for a in neighbours).most_common(),
+        "predecessors": Counter(a["predecessor"] for a in neighbours).most_common(),
+        "sparse_mla_successor_early": sum(a["successor_early"] for a in sparse_kernels),
+        "sparse_mla_started_before_predecessor_end": sum(
+            a["started_before_predecessor_end"] for a in sparse_kernels
+        ),
+        "sparse_mla_duration_us": sorted(a["duration_us"] for a in sparse_kernels),
+        "sparse_mla_successor_lead_us": sorted(
+            a["successor_lead_us"] for a in sparse_kernels if a["successor_lead_us"] is not None
+        ),
+        "w4_predecessor_lead_us": sorted(
+            a["predecessor_lead_us"] for a in sparse_kernels if a["predecessor_lead_us"] is not None
+        ),
+    }
+
+
+def drive(
+    stack, engagement, layers, ctx, replays, profiled_replays, topk_mode, features, out_path
+) -> None:
+    """Run, record and save every step of one process (see the module docstring)."""
+    # The prefill detached every tail; attach what the adoption built (None: native tail).
+    for layer, attn in stack.attn.items():
+        attn._attn_tail = stack.attn_tail_tails[layer]
+    full_layer = layers[0]
+    indexer = stack.attn[full_layer].mqa.indexer
+    assert indexer is not None, f"layer {full_layer} must run its own indexer"
+    shared_layers = [layer for layer in layers if stack.attn[layer].mqa.indexer is None]
+
+    hidden = [{layer: stack.gen_inputs(layer, step=s) for layer in layers} for s in (0, 1)]
+    topks = [local_topk(ctx, 77), local_topk(ctx, 78)]
+    static_hidden = {layer: hidden[0][layer].clone() for layer in layers}
+    static_topk = topks[0].clone()
+    if topk_mode == "fixed":
+        indexer.forward_from_projected = Mock(return_value=static_topk)
+    pos = stack.position_ids()
+    no_ar = AllReduceParams(enable_allreduce=False)
+    pool_rows = stack.kv_cache_manager.get_unique_primary_pool().view(torch.uint8)
+    pool_rows = pool_rows.reshape(-1, HEAD_DIM)
+    capture = FmhaCapture()
+
+    def load(index):
+        for layer in layers:
+            static_hidden[layer].copy_(hidden[index][layer])
+        static_topk.copy_(topks[index])
+
+    def step(index):
+        load(index)
+        with stack.scope():
+            return {
+                layer: stack.forward(layer, static_hidden[layer], pos, no_ar) for layer in layers
+            }
+
+    def generation_calls(start):
+        calls = [c for c in capture.calls[start:] if c["generation"]]
+        assert [c["layer"] for c in calls] == layers, [c["layer"] for c in calls]
+        return calls
+
+    # Warm-up (JIT compilation, autotuning, lazy buffers) with both input sets.
+    for index in (0, 1):
+        step(index)
+    torch.cuda.synchronize()
+
+    result = {"features": features, "topk_mode": topk_mode, "ctx": ctx, "layers": layers}
+    result["shared_layers"] = shared_layers
+    # Eager steps with the FMHA inputs recorded (clones before the attention).
+    capture.clone_inputs = True
+    engagement.record_bridge = True
+    result["inputs"] = []
+    bridge_pairs = []
+    for index in (0, 1):
+        start, bridge_start = len(capture.calls), len(engagement.bridge)
+        outs = step(index)
+        torch.cuda.synchronize()
+        per_layer = {}
+        for call in generation_calls(start):
+            idx = call["indices"]
+            per_layer[call["layer"]] = {
+                "library": call["library"],
+                "quant_q": call["quant_q"].cpu(),
+                "bmm1": call["bmm1"].cpu(),
+                "bmm2": call["bmm2"].cpu(),
+                "indices": idx.cpu(),
+                "kv_lens": call["kv_lens"].cpu(),
+                "latent": call["output"].clone().cpu(),
+                "module_output": outs[call["layer"]].clone().cpu(),
+                "rows": pool_rows[idx.long().clamp_min(0).reshape(-1)]
+                .reshape(ROWS, TOP_K, -1)
+                .cpu(),
+            }
+            for entry in engagement.bridge[bridge_start:]:
+                if entry["layer"] == call["layer"]:
+                    bridge_pairs.append((index, entry, idx, call["kv_lens"]))
+        result["inputs"].append({"input": index, "layers": per_layer})
+    capture.clone_inputs = False
+    engagement.record_bridge = False
+
+    def plain_steps(label):
+        records = []
+        for index in (0, 1):
+            start = len(capture.calls)
+            outs = step(index)
+            torch.cuda.synchronize()
+            records.append(
+                {
+                    "input": index,
+                    "layers": {
+                        c["layer"]: {
+                            "library": c["library"],
+                            "latent": c["output"].clone().cpu(),
+                            "indices": c["indices_ref"].clone().cpu(),
+                            "module_output": outs[c["layer"]].clone().cpu(),
+                        }
+                        for c in generation_calls(start)
+                    },
+                }
+            )
+        result[label] = records
+
+    plain_steps("pdl")
+    capture.sync_after_sparse_decode = True
+    plain_steps("serialized")
+    capture.sync_after_sparse_decode = False
+
+    # One captured verify step (both layers), replayed over alternating input sets.
+    step(0)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    start = len(capture.calls)
+    with stack.scope():
+        with torch.cuda.graph(graph):
+            graph_outs = {
+                layer: stack.forward(layer, static_hidden[layer], pos, no_ar) for layer in layers
+            }
+    graph_calls = generation_calls(start)
+
+    def replay(index):
+        load(index)
+        graph.replay()
+
+    result["graph"] = []
+    for it in range(replays):
+        index = it % 2
+        replay(index)
+        torch.cuda.synchronize()
+        result["graph"].append(
+            {
+                "input": index,
+                "layers": {
+                    c["layer"]: {
+                        "latent": c["output"].clone().cpu(),
+                        "indices": c["indices_ref"].clone().cpu(),
+                        "module_output": graph_outs[c["layer"]].clone().cpu(),
+                    }
+                    for c in graph_calls
+                },
+            }
+        )
+
+    # CUPTI kernel records: two eager steps without recording kernels, then graph replays.
+    eager_kernels = kernel_trace(lambda: [step(i) for i in (0, 1)])
+    graph_kernels = kernel_trace(lambda: [replay(i % 2) for i in range(profiled_replays)])
+    result["kernels"] = {
+        "eager": summarize_kernels(eager_kernels, 2, len(layers)),
+        "graph": summarize_kernels(graph_kernels, profiled_replays, len(layers)),
+    }
+
+    # rope_kv_append bridge contract, checked after every run so the native remap used as the
+    # reference cannot change the state of the steps above.
+    result["rope_append_bridge"] = []
+    for index, entry, received, kv_lens in bridge_pairs:
+        report = rope_append_bridge_report(entry, kv_lens, ROWS)
+        report["input"] = index
+        report["attention_received_bridge_output"] = bool(
+            torch.equal(received, entry["global"].to(received.device))
+        )
+        result["rope_append_bridge"].append(report)
+
+    tails = {}
+    for layer in layers:
+        tail = stack.attn_tail_tails[layer]
+        tails[layer] = (
+            None
+            if tail is None
+            else {
+                "engaged_calls": tail.engaged_calls,
+                "declined_calls": tail.declined_calls,
+                "static_reason": tail.static_reason,
+                "into_hop_published_calls": tail.into_hop_published_calls,
+                "config": dataclasses.asdict(tail.config),
+            }
+        )
+    sparse_decode_accounting = sparse_mla.engagement_snapshot()
+    result["engagement"] = {
+        "counts": dict(engagement.counts),
+        "attn_tail_tails": tails,
+        "sparse_mla_decode": {
+            "static": {str(k): v for k, v in sparse_decode_accounting["static"].items()},
+            "engaged_calls": {
+                str(k): v for k, v in sparse_decode_accounting["engaged_calls"].items()
+            },
+            "declined_selections": sparse_decode_accounting["declined_selections"],
+        },
+        "capability": {f"{f}/{key}": v for (f, key), v in dispatch._CAPABILITY_LOGGED.items()},
+    }
+    result["libraries"] = sorted({c["library"] for c in capture.calls if c["generation"]})
+    torch.save(result, out_path)
+    summary = {
+        "features": features,
+        "topk_mode": topk_mode,
+        "libraries": result["libraries"],
+        "engagement": result["engagement"],
+        "kernels": {
+            kind: {k: v for k, v in s.items() if k not in ("attention", "top_kernels")}
+            for kind, s in result["kernels"].items()
+        },
+        "rope_append_bridge": result["rope_append_bridge"],
+    }
+    print("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_SUMMARY " + json.dumps(summary, default=str))
+
+
+def main(out_path: str) -> None:
+    plan_testlib.select_from_env(dispatch)
+    ctx = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_CTX", "32768"))
+    layers = [
+        int(x)
+        for x in os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_LAYERS", "2,3").split(",")
+    ]
+    replays = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_REPLAYS", "8"))
+    profiled_replays = int(
+        os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_PROFILED_REPLAYS", "6")
+    )
+    topk_mode = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_TOPK", "fixed")
+    assert topk_mode in ("fixed", "indexer"), topk_mode
+    ckpt = os.environ["GLM52_NVFP4_CKPT"]
+    torch.manual_seed(0)
+    features = sorted(dispatch.enabled_features())
+    engagement = Engagement()
+
+    stack = real_mla_harness.AttentionStack(ckpt, layers, tp_rank=0, ctx=ctx, rows=ROWS)
+    # GEN steps run under inference mode, as the model engine runs them (absorbed_query's
+    # indexer projection updates views of its outputs in place).
+    with torch.inference_mode():
+        drive(
+            stack, engagement, layers, ctx, replays, profiled_replays, topk_mode, features, out_path
+        )
+    stack.shutdown()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py
new file mode 100644
index 000000000000..98dc9b455fc8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py
@@ -0,0 +1,100 @@
+# 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.
+"""Checks of the sparse-MLA decode library declaration on the DSA backend.
+
+These run on a GPU only on request (opt-in:
+``GLM_MOE_DSA_SPARSE_MLA_DECODE_ADAPTER_APPLIED=1``). Each check starts a fresh interpreter that
+installs the plan of the fast paths it is given (``plan_testlib.FAST_PATHS_ENV``). With
+sparse_mla_decode running, the layer declares the library ahead of the registry's, as ``attention.GlmMLA.adopt``
+does, and rebuilds its FMHA libraries, as ``create_weights`` does; the registry is unchanged.
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import plan_testlib  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_ADAPTER_APPLIED") != "1",
+    reason="opt-in GPU check",
+)
+
+_HERE = Path(__file__).resolve().parent
+
+_FORWARD_SCRIPT = r"""
+import json, os, sys
+sys.path.insert(0, sys.argv[1])
+sys.path.insert(0, os.path.dirname(sys.argv[1]))
+import conftest  # builds the dev op when needed
+import plan_testlib
+import torch
+from sparse_mla_harness import DsaGenHarness
+from tensorrt_llm._torch.attention.backends.fmha import registry
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
+    dispatch,
+)
+plan_testlib.select_from_env(dispatch)
+h = DsaGenHarness([20000], seq_len_q=6, seed=5)
+if dispatch.is_enabled("sparse_mla_decode"):
+    for layer in h.layers:
+        layer.prepended_fmha_lib_classes = (dispatch.get_impl("sparse_mla_decode", "fmha_lib"),)
+        layer.update_quant_config(layer.quant_config)
+try:
+    topk = h.random_topk(seed=6, overlap=0.7, recent=32)
+    step = h.make_step(0, seed=7)
+    out = h.forward(0, step, topk)
+    libs = [type(lib).__name__ for lib in h.layers[0]._fmha_manager.fmha_libs]
+    print(json.dumps({"registry": list(registry.FMHA_LIBS), "layer_libs": libs,
+                      "selected": h.selected_library(0),
+                      "out_finite": bool(torch.isfinite(out).all())}))
+finally:
+    h.shutdown()
+"""
+
+
+def _run(features: str) -> dict:
+    env = dict(os.environ)
+    env[plan_testlib.FAST_PATHS_ENV] = features
+    proc = subprocess.run(
+        [sys.executable, "-c", _FORWARD_SCRIPT, str(_HERE)],
+        env=env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert proc.returncode == 0, proc.stdout[-4000:] + proc.stderr[-4000:]
+    return json.loads(proc.stdout.strip().splitlines()[-1])
+
+
+def test_default_off_keeps_native_libraries():
+    result = _run("")
+    assert "glm_moe_dsa_sparse_mla_decode" not in result["registry"]
+    assert "SparseMlaDecodeFmha" not in result["layer_libs"]
+    assert result["selected"] == ["FallbackFmha"]
+
+
+def test_sparse_mla_selected_when_enabled():
+    result = _run("sparse_mla_decode")
+    assert "glm_moe_dsa_sparse_mla_decode" not in result["registry"]
+    assert result["layer_libs"][0] == "SparseMlaDecodeFmha"
+    assert result["layer_libs"].count("SparseMlaDecodeFmha") == 1
+    assert result["selected"] == ["SparseMlaDecodeFmha"]
+    assert result["out_finite"]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py
new file mode 100644
index 000000000000..7c9aca43c2b0
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py
@@ -0,0 +1,213 @@
+# 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.
+"""sparse_mla_decode composed with the packages around it at the MLA / DSA call sites.
+
+Each package set runs ``sparse_mla_case.py`` in a fresh interpreter, which installs the
+plan of that set, on the same deterministic layer and inputs:
+
+* producers (absorbed_query, absorbed_query + shared_indexer_absorbed_query, rope_kv_append) feed
+  the query and indices; for every set ``X`` the FMHA inputs of ``X`` and ``X + sparse_mla_decode``
+  must be bit-identical, and the latent output of each library is checked against an FP64 oracle
+  with the frozen error budget;
+* the consumer (attention_output_tail; attention_output_tail_prefetch where integrated) may start
+  before sparse_mla_decode finishes (sparse_mla_decode triggers its dependents before its gathers
+  complete): within each run the eager step, the same step with a device synchronization right after
+  the sparse_mla_decode launch, and replays of a captured graph over alternating inputs must give
+  bit-identical latent and module outputs whenever they attended the same
+  global indices.
+
+Every set runs with a fixed generation top-k (deterministic indices, so every
+replay is compared bit for bit); sets with rope_kv_append also run with the unmodified
+indexer, so that rope_kv_append's top-k bridge, where it engages, produces the indices
+sparse_mla_decode reads (the native indexer may return its selection in another order from run
+to run, so replays are compared only where the indices match). The capability
+decisions and reasons the packages log are recorded.
+
+Runs only where the installed build integrates the packages (an integration
+image); a package that declines the layer is reported, not failed.
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from sparse_mla_harness import error_stats, oracle_attention, record
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import plan_testlib  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION") != "1"
+    or not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available(),
+    reason="set GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION=1 on an SM10x integration build with the "
+    "sparse_mla_decode kernel",
+)
+
+_CASE = Path(__file__).resolve().parent / "sparse_mla_case.py"
+# Producers of the query/indices (absorbed_query, absorbed_query + shared_indexer_absorbed_query,
+# rope_kv_append RoPE/append and top-k bridge), the consumer tail (attention_output_tail,
+# attention_output_tail + attention_output_tail_prefetch) and one set with all of them; (set, top-k
+# mode).
+NEIGHBOUR_SETS = [
+    ("", "fixed"),
+    ("absorbed_query", "fixed"),
+    ("absorbed_query,shared_indexer_absorbed_query", "fixed"),
+    ("rope_kv_append", "fixed"),
+    ("rope_kv_append", "indexer"),
+    ("attention_output_tail", "fixed"),
+    ("attention_output_tail,attention_output_tail_prefetch", "fixed"),
+    ("absorbed_query,shared_indexer_absorbed_query,rope_kv_append,attention_output_tail", "fixed"),
+    (
+        "absorbed_query,shared_indexer_absorbed_query,rope_kv_append,attention_output_tail",
+        "indexer",
+    ),
+]
+ROWS = 6
+TOP_K = 2048
+
+
+def _run(features: str, topk_mode: str, out_dir: Path) -> dict:
+    out = out_dir / f"mla_{features.replace(',', '_') or 'native'}_{topk_mode}.pt"
+    env = dict(os.environ)
+    env[plan_testlib.FAST_PATHS_ENV] = features
+    env["GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK"] = topk_mode
+    env["TLLM_LOG_LEVEL"] = "INFO"
+    proc = subprocess.run(
+        [sys.executable, str(_CASE), str(out)], env=env, capture_output=True, text=True, check=False
+    )
+    assert proc.returncode == 0, proc.stdout[-6000:] + proc.stderr[-6000:]
+    result = torch.load(out, weights_only=False)
+    log = (proc.stdout + proc.stderr).splitlines()
+    result["capability_log"] = sorted(
+        {line[line.index("[glm_moe_dsa]") :] for line in log if "[glm_moe_dsa]" in line}
+    )
+    return result
+
+
+def _oracle(entry: dict) -> torch.Tensor:
+    rows = entry["rows"].cuda().reshape(ROWS * TOP_K, -1)
+    indices = entry["indices"].cuda()
+    local = torch.arange(ROWS * TOP_K, device="cuda", dtype=torch.int32).view(ROWS, TOP_K)
+    remapped = torch.where(indices >= 0, local, torch.full_like(local, -1))
+    kv_len = int(entry["kv_lens"][0])
+    limits = [kv_len - (ROWS - 1 - r) for r in range(ROWS)]
+    return oracle_attention(
+        entry["quant_q"].cuda(),
+        rows,
+        remapped,
+        float(entry["bmm1"][0]),
+        float(entry["bmm2"][0]),
+        list_limit=limits,
+    )
+
+
+def _same(a: torch.Tensor, b: torch.Tensor) -> bool:
+    return a.shape == b.shape and torch.equal(a.view(torch.uint8), b.view(torch.uint8))
+
+
+@pytest.mark.parametrize("neighbours,topk_mode", NEIGHBOUR_SETS)
+def test_sparse_mla_with_neighbours(neighbours, topk_mode, tmp_path):
+    out_dir = Path(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS", tmp_path))
+    with_sparse_mla = ",".join([f for f in neighbours.split(",") if f] + ["sparse_mla_decode"])
+    base = _run(neighbours, topk_mode, out_dir)
+    cand = _run(with_sparse_mla, topk_mode, out_dir)
+    assert base["libraries"] == ["FallbackFmha"], base["libraries"]
+    assert cand["libraries"] == ["SparseMlaDecodeFmha"], cand["libraries"]
+
+    summary = {
+        "neighbours": neighbours or "native",
+        "topk_mode": cand["topk_mode"],
+        "capability_without_sparse_mla": base["capability"],
+        "capability_with_sparse_mla": cand["capability"],
+        "capability_log_with_sparse_mla": cand["capability_log"],
+        "steps": [],
+    }
+    assert base["topk_mode"] == cand["topk_mode"]
+    for b, c in zip(base["eager"], cand["eager"]):
+        # The producers do not depend on the consumer: identical FMHA inputs.
+        for key in ("quant_q", "bmm1", "bmm2", "kv_lens"):
+            assert _same(b[key], c[key]), (neighbours, key)
+        same_order = _same(b["indices"], c["indices"])
+        if not same_order:
+            # Only an unmodified indexer may return the same selection in another order.
+            assert cand["topk_mode"] == "indexer", neighbours
+            assert torch.equal(b["indices"].sort(dim=-1).values, c["indices"].sort(dim=-1).values)
+        else:
+            assert _same(b["rows"], c["rows"]), (neighbours, "rows")
+        assert not bool((c["indices"] < 0).any()), "an index below the causal length is -1"
+        oracle = _oracle(c)
+        native_err = error_stats(b["latent"].cuda().double(), oracle)
+        sparse_mla_err = error_stats(c["latent"].cuda().double(), oracle)
+        module_diff = error_stats(c["module_output"].double(), b["module_output"].double())
+        summary["steps"].append(
+            {
+                "input": c["input"],
+                "indices_same_order": same_order,
+                "native_latent_vs_oracle": native_err,
+                "sparse_mla_latent_vs_oracle": sparse_mla_err,
+                "module_output_sparse_mla_vs_native": module_diff,
+            }
+        )
+        assert sparse_mla_err["nonfinite"] == native_err["nonfinite"] == 0
+        assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
+            sparse_mla_err,
+            native_err,
+        )
+        assert sparse_mla_err["max_abs"] <= 2.0 * native_err["max_abs"] + 1e-3, (
+            sparse_mla_err,
+            native_err,
+        )
+        assert module_diff["nonfinite"] == 0 and module_diff["nrms"] < 0.05, module_diff
+
+    # Programmatic overlap with the next kernel and graph replay change nothing.
+    for run, label in ((base, "without_sparse_mla"), (cand, "with_sparse_mla")):
+        eager = {e["input"]: e for e in run["eager"]}
+        mismatches = []
+        index_changes = []
+        for kind in ("serialized", "graph"):
+            for i, entry in enumerate(run[kind]):
+                ref = eager[entry["input"]]
+                if not _same(entry["indices"], ref["indices"]):
+                    index_changes.append((kind, i))
+                    continue
+                for key in ("latent", "module_output"):
+                    if not _same(entry[key], ref[key]):
+                        mismatches.append((kind, i, key))
+        summary[f"replay_mismatches_{label}"] = mismatches
+        summary[f"replay_index_changes_{label}"] = index_changes
+    record(
+        "sparse_mla_composition.jsonl",
+        "mla_composition",
+        json.loads(json.dumps(summary, default=str)),
+    )
+    # sparse_mla_decode is deterministic, so with it every replay and the serialized step that
+    # attended the same indices must match exactly; a fixed top-k never changes the indices.
+    assert not summary["replay_mismatches_with_sparse_mla"], summary[
+        "replay_mismatches_with_sparse_mla"
+    ]
+    if cand["topk_mode"] == "fixed":
+        assert not summary["replay_index_changes_with_sparse_mla"], summary[
+            "replay_index_changes_with_sparse_mla"
+        ]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py
new file mode 100644
index 000000000000..0231c119f646
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py
@@ -0,0 +1,123 @@
+# 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.
+"""Capability logging, error behaviour and engagement accounting of sparse_mla_decode.
+
+* The static decision is logged once per layer, not once per process.
+* A layer inside the sparse_mla_decode domain whose kernel is missing from the build or whose
+  device cannot launch it raises instead of falling back to native on one rank;
+  a layer outside the domain still takes the native library.
+* ``engagement_snapshot`` counts the sparse_mla_decode calls per layer and the declined
+  selections per reason.
+"""
+
+import pytest
+import torch
+from sparse_mla_harness import DsaGenHarness
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
+    SparseMlaDecodeFmha,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available(),
+    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
+)
+
+
+def test_capability_logged_per_layer():
+    harness = DsaGenHarness([256], seq_len_q=1, num_layers=3, seed=141)
+    try:
+        for layer in harness.layers:
+            assert SparseMlaDecodeFmha.is_available(layer)
+        snapshot = sparse_mla.engagement_snapshot()
+        for layer in harness.layers:
+            assert snapshot["static"][layer.layer_idx] == ""
+            assert (
+                dispatch._CAPABILITY_LOGGED[("sparse_mla_decode", f"layer{layer.layer_idx}")]
+                is True
+            )
+    finally:
+        harness.shutdown()
+
+
+def test_selected_sparse_mla_raises_instead_of_falling_back(monkeypatch):
+    harness = DsaGenHarness([256], seq_len_q=1, seed=142)
+    try:
+        layer = harness.layers[0]
+        monkeypatch.setattr(sparse_mla, "device_supported", lambda device: False)
+        with pytest.raises(RuntimeError, match="cannot launch its 16-CTA cluster kernel"):
+            SparseMlaDecodeFmha.is_available(layer)
+        monkeypatch.setattr(sparse_mla, "op_available", lambda: False)
+        with pytest.raises(RuntimeError, match="does not register"):
+            SparseMlaDecodeFmha.is_available(layer)
+    finally:
+        harness.shutdown()
+
+
+def test_layer_outside_domain_takes_native():
+    harness = DsaGenHarness([256], seq_len_q=1, kv_dtype=torch.bfloat16, seed=143)
+    try:
+        layer = harness.layers[0]
+        assert not SparseMlaDecodeFmha.is_available(layer)
+        assert sparse_mla.engagement_snapshot()["static"][layer.layer_idx] == "KV cache is not FP8"
+    finally:
+        harness.shutdown()
+
+
+def test_engagement_counts_per_layer_and_declines():
+    harness = DsaGenHarness([20000], seq_len_q=6, num_layers=2, seed=144)
+    try:
+        before = sparse_mla.engagement_snapshot()["engaged_calls"]
+        topk = harness.random_topk(seed=146)
+        for layer_idx in (0, 1):
+            harness.select_library(layer_idx, use_sparse_mla=True)
+            for _ in range(2):
+                harness.forward(layer_idx, harness.make_step(layer_idx, seed=147), topk)
+            assert harness.selected_library(layer_idx) == ["SparseMlaDecodeFmha"]
+        torch.cuda.synchronize()
+        after = sparse_mla.engagement_snapshot()["engaged_calls"]
+        for layer in harness.layers:
+            idx = layer.layer_idx
+            assert after.get(idx, 0) - before.get(idx, 0) == 2, (before, after)
+            assert dispatch._CAPABILITY_LOGGED[("sparse_mla_decode", f"call:layer{idx}")] is True
+    finally:
+        harness.shutdown()
+    other = DsaGenHarness([16384], seq_len_q=6, index_topk=1024, seed=148)
+    try:
+        declined_before = sparse_mla.engagement_snapshot()["declined_selections"]
+        gen = torch.Generator(device="cpu").manual_seed(149)
+        topk = (
+            torch.stack(
+                [
+                    torch.randperm(other.causal_len(r), generator=gen)[:1024]
+                    for r in range(other.num_rows)
+                ]
+            )
+            .to(torch.int32)
+            .cuda()
+        )
+        other.select_library(0, use_sparse_mla=True)
+        other.forward(0, other.make_step(0, seed=150), topk)
+        assert "SparseMlaDecodeFmha" not in other.selected_library(0)
+        declined_after = sparse_mla.engagement_snapshot()["declined_selections"]
+        assert declined_after.get("top-k 1024", 0) == declined_before.get("top-k 1024", 0) + 1
+    finally:
+        other.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py
new file mode 100644
index 000000000000..1d473fb57b00
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py
@@ -0,0 +1,171 @@
+# 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.
+"""Characterization of the native sparse MLA generation kernel (FP8 KV).
+
+The native trtllm-gen kernel ships as a binary, so its rounding and masking
+rules are measured here against FP64 models: which probability rounding it
+feeds to P*V, whether it ignores index-list positions past a row's causal KV
+length, what it writes for a row without valid entries, and whether repeated
+indices count once per occurrence. sparse_mla_decode must reproduce whatever these tests find.
+"""
+
+import json
+import os
+
+import pytest
+import torch
+from sparse_mla_harness import (
+    TOP_K,
+    DsaGenHarness,
+    bf16_round,
+    e4m3_round,
+    error_stats,
+    fp16_round,
+    oracle_attention,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="needs an SM10x GPU",
+)
+
+PROB_MODELS = {
+    "exact": {},
+    "fp16_split256_step128": dict(prob_round=fp16_round, split=256, step=128),
+    "bf16_split256_step128": dict(prob_round=bf16_round, split=256, step=128),
+    "e4m3x1_split256_step128": dict(prob_round=e4m3_round(1.0), split=256, step=128),
+    "e4m3x448_split256_step128": dict(prob_round=e4m3_round(448.0), split=256, step=128),
+    "e4m3x256_split256_step128": dict(prob_round=e4m3_round(256.0), split=256, step=128),
+    "e4m3x1_split2048_step128": dict(prob_round=e4m3_round(1.0), split=2048, step=128),
+    "e4m3x448_split2048_step128": dict(prob_round=e4m3_round(448.0), split=2048, step=128),
+    "e4m3x1_split128_step128": dict(prob_round=e4m3_round(1.0), split=128, step=128),
+    "e4m3x448_split128_step128": dict(prob_round=e4m3_round(448.0), split=128, step=128),
+}
+
+
+def record(case: str, payload: dict) -> None:
+    line = json.dumps({"case": case, **payload}, sort_keys=True)
+    print(line)
+    results_dir = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS")
+    if results_dir:
+        with open(os.path.join(results_dir, "native_semantics.jsonl"), "a") as f:
+            f.write(line + "\n")
+
+
+def run_native(harness: DsaGenHarness, topk: torch.Tensor, seed: int = 1, q_std: float = 1.0):
+    step = harness.make_step(0, seed=seed, q_std=q_std)
+    harness.select_library(0, use_sparse_mla=False)
+    out = harness.forward(0, step, topk).clone()
+    global_idx = harness.global_indices(0, topk)
+    return step, out, global_idx
+
+
+@pytest.mark.parametrize("ctx_len,q_std", [(32768, 1.0), (32768, 3.0), (32768, 0.3), (4096, 1.0)])
+def test_native_probability_rounding(ctx_len, q_std):
+    harness = DsaGenHarness([ctx_len], seq_len_q=6, seed=11)
+    try:
+        topk = harness.random_topk(seed=3, overlap=0.7, recent=32)
+        step, native, global_idx = run_native(harness, topk, seed=5, q_std=q_std)
+        b1 = float(step.bmm1_scale[0])
+        b2 = float(step.bmm2_scale[0])
+        native = native.double()
+        distances = {}
+        for name, kwargs in PROB_MODELS.items():
+            ref = oracle_attention(step.quant_q, harness.pool_rows(), global_idx, b1, b2, **kwargs)
+            ref_bf16 = ref.to(torch.bfloat16).double()
+            stats = error_stats(native, ref_bf16)
+            stats["exact_match_frac"] = float(
+                (native.reshape(ref_bf16.shape) == ref_bf16).double().mean().item()
+            )
+            distances[name] = stats
+        best = min(distances, key=lambda n: distances[n]["nrms"])
+        record(
+            "prob_rounding",
+            {
+                "ctx_len": ctx_len,
+                "q_std": q_std,
+                "bmm1_scale": [float(x) for x in step.bmm1_scale.tolist()],
+                "bmm2_scale": b2,
+                "best": best,
+                "distances": distances,
+            },
+        )
+        assert distances["exact"]["nonfinite"] == 0
+    finally:
+        harness.shutdown()
+
+
+def test_native_kv_length_clamp():
+    """Valid entries placed after a row's causal length in the index list."""
+    harness = DsaGenHarness([100], seq_len_q=6, seed=12)
+    try:
+        topk = harness.random_topk(seed=4)
+        extra_positions = [700, 1500, TOP_K - 1]
+        for row in range(harness.num_rows):
+            for pos in extra_positions:
+                topk[row, pos] = 3 + row
+        step, native, global_idx = run_native(harness, topk, seed=6)
+        b1 = float(step.bmm1_scale[0])
+        b2 = float(step.bmm2_scale[0])
+        limits = [harness.causal_len(r) for r in range(harness.num_rows)]
+        clamped = oracle_attention(
+            step.quant_q, harness.pool_rows(), global_idx, b1, b2, list_limit=limits
+        )
+        unclamped = oracle_attention(step.quant_q, harness.pool_rows(), global_idx, b1, b2)
+        record(
+            "kv_length_clamp",
+            {
+                "vs_clamped": error_stats(native.double(), clamped),
+                "vs_unclamped": error_stats(native.double(), unclamped),
+                "limits": limits,
+            },
+        )
+    finally:
+        harness.shutdown()
+
+
+def test_native_invalid_and_duplicate_entries():
+    harness = DsaGenHarness([8192], seq_len_q=6, seed=13)
+    try:
+        topk = harness.random_topk(seed=5)
+        # Row 1: no valid entry. Row 2: invalid entries interleaved with valid ones.
+        # Row 3: two tokens repeated with unequal multiplicity.
+        topk[1, :] = -1
+        topk[2, ::3] = -1
+        topk[3, :1500] = 17
+        topk[3, 1500:] = 4242
+        step, native, global_idx = run_native(harness, topk, seed=7)
+        b1 = float(step.bmm1_scale[0])
+        b2 = float(step.bmm2_scale[0])
+        exact = oracle_attention(step.quant_q, harness.pool_rows(), global_idx, b1, b2)
+        dedup_idx = global_idx.clone()
+        dedup_idx[3, 1:1500] = -1
+        dedup_idx[3, 1501:] = -1
+        dedup = oracle_attention(step.quant_q, harness.pool_rows(), dedup_idx, b1, b2)
+        row1 = native[1].double()
+        record(
+            "invalid_and_duplicates",
+            {
+                "all_invalid_row": {
+                    "nonfinite": int((~torch.isfinite(row1)).sum().item()),
+                    "max_abs": float(row1.nan_to_num().abs().max().item()),
+                },
+                "interleaved_invalid_row": error_stats(native[2].double(), exact[2]),
+                "duplicates_counted": error_stats(native[3].double(), exact[3]),
+                "duplicates_dedup": error_stats(native[3].double(), dedup[3]),
+            },
+        )
+    finally:
+        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py
new file mode 100644
index 000000000000..2604a240510e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py
@@ -0,0 +1,479 @@
+# 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.
+"""Delayed-producer and delayed-consumer stress of the sparse_mla_decode kernel under PDL.
+
+The sparse_mla_decode kernel initializes its mbarriers and arrives on the cluster barrier
+before ``griddepcontrol.wait``, issues ``griddepcontrol.launch_dependents``
+before its compute, gathers pool rows with bulk copies and pushes partial
+results to its cluster peers with ``st.async``. This test checks that none of
+that reads an input early or writes the output early.
+
+Every iteration is a chain without a plain kernel boundary:
+
+  producer (PDL)  -> sparse_mla_decode (PDL) -> consumer (PDL) -> next producer ...
+
+The producer lets its dependents launch at once, waits for its predecessor,
+spins, and only then writes everything sparse_mla_decode reads: the FP8 query, the index
+list, the KV lengths, the bmm scales and a block of pool rows the index list
+points into. The consumer lets its dependents launch at once, waits for sparse_mla_decode,
+optionally spins, and copies the output out. Consecutive iterations use
+different inputs, and each input set alternates between two versions of its
+pool rows, so an early read of any input, a late write of the output, or a
+stale cache line gives a different result. Each output must be bit-identical
+to the output of the same inputs computed in isolation (sparse_mla_decode is deterministic;
+this is checked too); the isolated outputs are checked against the FP64 oracle.
+
+Variants: eager chain with and without PDL on the helper kernels, the chain
+captured in a CUDA graph and replayed, and the eager chain while a second
+stream saturates HBM. ``%globaltimer`` stamps of the helper kernels and a
+CUPTI trace show whether the early launches happened.
+"""
+
+import os
+import statistics
+
+import pytest
+import torch
+from sparse_mla_harness import E4M3_MAX, error_stats, oracle_attention, record
+from sparse_mla_pdl_helpers import load_pdl_helpers
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available(),
+    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
+)
+
+NUM_HEADS = 8
+HEAD_DIM = 576
+KV_LORA_RANK = 512
+TOP_K = 2048
+STATIC_ROWS = 120_000
+OWNED_ROWS = 512
+NUM_SETS = 6
+ITERS = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STRESS_ITERS", "400"))
+GRAPH_REPLAYS = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STRESS_GRAPH_REPLAYS", "8"))
+PRODUCER_SPINS_NS = (0, 1000, 5000, 20000)
+CONSUMER_SPINS_NS = (0, 0, 3000, 10000)
+RESULTS = "sparse_mla_stress.jsonl"
+# compute-sanitizer serializes and slows the instrumented kernel; the early-launch evidence is
+# then not expected, only the correctness of every iteration.
+UNDER_SANITIZER = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_UNDER_SANITIZER") == "1"
+
+
+def _e4m3_bytes(shape, std: float, gen: torch.Generator) -> torch.Tensor:
+    values = torch.randn(*shape, device="cuda", generator=gen) * std
+    return values.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8)
+
+
+class StressCase:
+    """Live buffers read by sparse_mla_decode, staged input sets and their isolated reference outputs."""
+
+    def __init__(self, ctx_lens, seq_len_q: int, seed: int):
+        self.seq_len_q = seq_len_q
+        self.num_seqs = len(ctx_lens)
+        assert self.num_seqs <= 4
+        self.rows = self.num_seqs * seq_len_q
+        gen = torch.Generator(device="cuda").manual_seed(seed)
+        cpu_gen = torch.Generator(device="cpu").manual_seed(seed)
+        num_pool_rows = STATIC_ROWS + NUM_SETS * OWNED_ROWS
+        self.pool = _e4m3_bytes((num_pool_rows, HEAD_DIM), 1.0, gen)
+        self.q = torch.zeros(self.rows, NUM_HEADS * HEAD_DIM, dtype=torch.uint8, device="cuda")
+        self.indices = torch.zeros(self.rows, TOP_K, dtype=torch.int32, device="cuda")
+        self.kv_lens = torch.zeros(4, dtype=torch.int32, device="cuda")
+        self.scales = torch.zeros(4, dtype=torch.float32, device="cuda")
+        self.out = torch.zeros(
+            self.rows, NUM_HEADS * KV_LORA_RANK, dtype=torch.bfloat16, device="cuda"
+        )
+        self.sets = []
+        for k in range(NUM_SETS):
+            kv_lens = [c + seq_len_q + 3 * k for c in ctx_lens]
+            owned_base = STATIC_ROWS + k * OWNED_ROWS
+            limits = []
+            idx = torch.empty(self.rows, TOP_K, dtype=torch.int64)
+            for row in range(self.rows):
+                seq, pos = divmod(row, seq_len_q)
+                usable = kv_lens[seq] - (seq_len_q - 1 - pos)
+                limits.append(usable)
+                entries = torch.randint(0, STATIC_ROWS, (TOP_K,), generator=cpu_gen)
+                owned = torch.rand(TOP_K, generator=cpu_gen) < 0.125
+                entries[owned] = owned_base + torch.randint(
+                    0, OWNED_ROWS, (int(owned.sum()),), generator=cpu_gen
+                )
+                negative = torch.rand(TOP_K, generator=cpu_gen) < 0.02
+                entries[negative] = -1
+                entries[64:96] = entries[0:32]  # repeated rows count per occurrence
+                idx[row] = entries
+            kv_tensor = torch.zeros(4, dtype=torch.int32)
+            kv_tensor[: self.num_seqs] = torch.tensor(kv_lens, dtype=torch.int32)
+            softmax_scale = 0.0625 * (1.0 + 0.05 * k)
+            scales = torch.tensor(
+                [softmax_scale, softmax_scale * 1.4426950408889634, 1.0 + 0.01 * k, 0.0],
+                dtype=torch.float32,
+            )
+            variants = [
+                _e4m3_bytes((OWNED_ROWS, HEAD_DIM), 2.0, gen),
+                _e4m3_bytes((OWNED_ROWS, HEAD_DIM), 0.5, gen),
+            ]
+            self.sets.append(
+                {
+                    "q": _e4m3_bytes((self.rows, NUM_HEADS * HEAD_DIM), 1.0 + 0.25 * k, gen),
+                    "indices": idx.to(torch.int32).cuda(),
+                    "kv_lens": kv_tensor.cuda(),
+                    "scales": scales.cuda(),
+                    "owned_base": owned_base,
+                    "variants": variants,
+                    "limits": limits,
+                }
+            )
+        self.segments = {}
+        for k, s in enumerate(self.sets):
+            for v in range(2):
+                owned = self.pool[s["owned_base"] : s["owned_base"] + OWNED_ROWS]
+                pairs = [
+                    (self.q, s["q"]),
+                    (self.indices, s["indices"]),
+                    (self.kv_lens, s["kv_lens"]),
+                    (self.scales, s["scales"]),
+                    (owned, s["variants"][v]),
+                ]
+                for dst, src in pairs:
+                    assert dst.is_contiguous() and src.is_contiguous() and dst.nbytes == src.nbytes
+                    assert dst.nbytes % 16 == 0
+                self.segments[(k, v)] = tuple(
+                    torch.tensor(values, dtype=torch.int64, device="cuda")
+                    for values in (
+                        [d.data_ptr() for d, _ in pairs],
+                        [s_.data_ptr() for _, s_ in pairs],
+                        [d.nbytes // 16 for d, _ in pairs],
+                    )
+                )
+        self.references = {}
+
+    def apply_plain(self, k: int, v: int) -> None:
+        s = self.sets[k]
+        self.q.copy_(s["q"])
+        self.indices.copy_(s["indices"])
+        self.kv_lens.copy_(s["kv_lens"])
+        self.scales.copy_(s["scales"])
+        self.pool[s["owned_base"] : s["owned_base"] + OWNED_ROWS].copy_(s["variants"][v])
+
+    def launch_sparse_mla(self) -> None:
+        sparse_mla.sparse_mla_decode(
+            self.q.view(self.rows, NUM_HEADS, HEAD_DIM),
+            self.pool.data_ptr(),
+            self.indices,
+            self.scales[0:2],
+            self.scales[2:3],
+            self.kv_lens,
+            self.seq_len_q,
+            self.out,
+        )
+
+    def build_references(self, case: str) -> None:
+        """Isolated outputs of every (set, version); determinism and FP64-oracle checks."""
+        for k in range(NUM_SETS):
+            for v in range(2):
+                runs = []
+                for _ in range(2):
+                    self.apply_plain(k, v)
+                    torch.cuda.synchronize()
+                    self.out.fill_(float("nan"))
+                    self.launch_sparse_mla()
+                    torch.cuda.synchronize()
+                    runs.append(self.out.clone())
+                assert torch.equal(runs[0], runs[1]), (
+                    "sparse_mla_decode output is not deterministic"
+                )
+                s = self.sets[k]
+                oracle = oracle_attention(
+                    self.q,
+                    self.pool.view(torch.float8_e4m3fn),
+                    self.indices,
+                    float(s["scales"][0]),
+                    float(s["scales"][2]),
+                    list_limit=s["limits"],
+                )
+                err = error_stats(runs[0].double(), oracle)
+                record(
+                    RESULTS,
+                    f"{case}:reference",
+                    {"set": k, "version": v, "sparse_mla_vs_oracle": err},
+                )
+                assert err["nonfinite"] == 0
+                assert err["nrms"] < 0.05, err
+                assert torch.allclose(
+                    runs[0].double().reshape(oracle.shape), oracle, atol=0.12, rtol=0.01
+                )
+                self.references[(k, v)] = runs[0]
+        # The two versions of each set's pool rows must change the output, or a stale read of
+        # those rows would go unnoticed.
+        for k in range(NUM_SETS):
+            assert not torch.equal(self.references[(k, 0)], self.references[(k, 1)])
+
+
+def _schedule(iterations: int, seed: int):
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    plan = []
+    for i in range(iterations):
+        k = i % NUM_SETS
+        v = (i // NUM_SETS) % 2
+        p = PRODUCER_SPINS_NS[int(torch.randint(0, len(PRODUCER_SPINS_NS), (1,), generator=gen))]
+        c = CONSUMER_SPINS_NS[int(torch.randint(0, len(CONSUMER_SPINS_NS), (1,), generator=gen))]
+        plan.append((k, v, p, c))
+    return plan
+
+
+# A spin kernel ahead of an eager chain (about 0.2 s), so that the host enqueues the whole
+# chain before the GPU reaches it and the GPU, not the host, paces the chain.
+GATE_CYCLES = int(4e8)
+
+
+def _enqueue_chain(case: StressCase, helpers, plan, outs, stamps, use_pdl: bool) -> None:
+    for i, (k, v, producer_spin, consumer_spin) in enumerate(plan):
+        dsts, srcs, counts = case.segments[(k, v)]
+        helpers.delayed_scatter(dsts, srcs, counts, producer_spin, stamps[i, 0:3], use_pdl, 64)
+        case.launch_sparse_mla()
+        helpers.pdl_consumer(outs[i], case.out, consumer_spin, stamps[i, 3:5], use_pdl)
+
+
+def _mismatches(case: StressCase, plan, outs) -> list:
+    bad = []
+    for i, (k, v, _, _) in enumerate(plan):
+        if not torch.equal(outs[i], case.references[(k, v)]):
+            diff = (outs[i].float() - case.references[(k, v)].float()).abs()
+            bad.append(
+                {
+                    "iteration": i,
+                    "set": k,
+                    "version": v,
+                    "max_abs": float(diff.nan_to_num(1e30).max()),
+                }
+            )
+    return bad
+
+
+def _stamp_summary(stamps: torch.Tensor, plan) -> dict:
+    """Early-launch evidence from the helper kernels' ``%globaltimer`` stamps.
+
+    A consumer that started well before its ``griddepcontrol.wait`` returned was
+    launched while sparse_mla_decode was still running, i.e. through sparse_mla_decode's early
+    ``launch_dependents``.
+    """
+    s = stamps.cpu().tolist()
+    consumer_wait = [row[4] - row[3] for row in s]
+    producer_wait = [row[1] - row[0] for row in s]
+    return {
+        "consumer_wait_ns_median": statistics.median(consumer_wait),
+        "consumer_started_before_sparse_mla_done": sum(1 for x in consumer_wait if x > 500)
+        / len(s),
+        "producer_wait_ns_median": statistics.median(producer_wait),
+        "iterations": len(plan),
+    }
+
+
+@pytest.fixture(scope="module")
+def helpers():
+    return load_pdl_helpers()
+
+
+CASES = {
+    "one_seq_r6": ([40000], 6),
+    "three_seqs_r2": ([3000, 100, 9000], 2),
+}
+
+
+@pytest.mark.parametrize("use_pdl", [True, False], ids=["pdl", "no_pdl"])
+@pytest.mark.parametrize("case_name", list(CASES))
+def test_sparse_mla_delayed_producer_consumer_eager(helpers, case_name, use_pdl):
+    ctx_lens, seq_len_q = CASES[case_name]
+    case = StressCase(ctx_lens, seq_len_q, seed=1)
+    case.build_references(case_name)
+    plan = _schedule(ITERS, seed=2)
+    outs = torch.full(
+        (ITERS, case.rows, NUM_HEADS * KV_LORA_RANK),
+        float("nan"),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    stamps = torch.zeros(ITERS, 8, dtype=torch.int64, device="cuda")
+    # Start from a different set than the first iteration's, so iteration 0 also overwrites.
+    # Warm-up chain (module loading, launch attributes), then the measured chain.
+    _enqueue_chain(case, helpers, plan[:2], outs, stamps, use_pdl)
+    case.apply_plain(NUM_SETS - 1, 1)
+    outs.fill_(float("nan"))
+    stamps.zero_()
+    torch.cuda.synchronize()
+    torch.cuda._sleep(GATE_CYCLES)
+    _enqueue_chain(case, helpers, plan, outs, stamps, use_pdl)
+    torch.cuda.synchronize()
+    bad = _mismatches(case, plan, outs)
+    summary = _stamp_summary(stamps, plan)
+    record(
+        RESULTS,
+        f"{case_name}:eager",
+        {
+            "use_pdl_helpers": use_pdl,
+            "trtllm_enable_pdl_env": os.environ.get("TRTLLM_ENABLE_PDL"),
+            "mismatches": len(bad),
+            "first_mismatches": bad[:5],
+            **summary,
+        },
+    )
+    assert not bad, bad[:5]
+    if use_pdl and os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0" and not UNDER_SANITIZER:
+        # The chain must actually exercise the early launch, or this test proves nothing.
+        assert summary["consumer_started_before_sparse_mla_done"] > 0.25, summary
+
+
+@pytest.mark.parametrize("case_name", list(CASES))
+def test_sparse_mla_delayed_producer_consumer_graph(helpers, case_name):
+    ctx_lens, seq_len_q = CASES[case_name]
+    case = StressCase(ctx_lens, seq_len_q, seed=3)
+    case.build_references(f"{case_name}:graph")
+    length = 4 * NUM_SETS  # the version of each set alternates across replays too
+    plan = _schedule(length, seed=4)
+    outs = torch.full(
+        (length, case.rows, NUM_HEADS * KV_LORA_RANK),
+        float("nan"),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    stamps = torch.zeros(length, 8, dtype=torch.int64, device="cuda")
+    case.apply_plain(NUM_SETS - 1, 1)
+    # Warm up outside capture (attribute setup, lazy loading).
+    _enqueue_chain(case, helpers, plan[:2], outs, stamps, True)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        _enqueue_chain(case, helpers, plan, outs, stamps, True)
+    total_bad = 0
+    summaries = []
+    for replay in range(GRAPH_REPLAYS):
+        outs.fill_(float("nan"))
+        stamps.zero_()
+        graph.replay()
+        torch.cuda.synchronize()
+        bad = _mismatches(case, plan, outs)
+        total_bad += len(bad)
+        summaries.append(_stamp_summary(stamps, plan))
+        assert not bad, (replay, bad[:5])
+    record(
+        RESULTS,
+        f"{case_name}:graph",
+        {
+            "replays": GRAPH_REPLAYS,
+            "chain_length": length,
+            "mismatches": total_bad,
+            "consumer_started_before_sparse_mla_done": [
+                x["consumer_started_before_sparse_mla_done"] for x in summaries
+            ],
+            "consumer_wait_ns_median": [x["consumer_wait_ns_median"] for x in summaries],
+        },
+    )
+
+
+def test_sparse_mla_delayed_producer_consumer_under_hbm_contention(helpers):
+    ctx_lens, seq_len_q = CASES["one_seq_r6"]
+    case = StressCase(ctx_lens, seq_len_q, seed=5)
+    case.build_references("contention")
+    plan = _schedule(ITERS, seed=6)
+    outs = torch.full(
+        (ITERS, case.rows, NUM_HEADS * KV_LORA_RANK),
+        float("nan"),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    stamps = torch.zeros(ITERS, 8, dtype=torch.int64, device="cuda")
+    big_src = torch.ones(256 * 1024 * 1024, dtype=torch.float32, device="cuda")
+    big_dst = torch.empty_like(big_src)
+    case.apply_plain(NUM_SETS - 1, 1)
+    torch.cuda.synchronize()
+    side = torch.cuda.Stream()
+    with torch.cuda.stream(side):
+        for _ in range(160):
+            big_dst.copy_(big_src)
+    torch.cuda._sleep(GATE_CYCLES // 4)
+    _enqueue_chain(case, helpers, plan, outs, stamps, True)
+    torch.cuda.synchronize()
+    bad = _mismatches(case, plan, outs)
+    record(
+        RESULTS,
+        "contention:eager",
+        {"mismatches": len(bad), "first_mismatches": bad[:5], **_stamp_summary(stamps, plan)},
+    )
+    assert not bad, bad[:5]
+
+
+@pytest.mark.skipif(UNDER_SANITIZER, reason="CUPTI tracing is not combined with compute-sanitizer")
+def test_sparse_mla_early_launch_trace(helpers):
+    """CUPTI kernel intervals: sparse_mla_decode becomes resident before its producer finishes."""
+    ctx_lens, seq_len_q = CASES["one_seq_r6"]
+    case = StressCase(ctx_lens, seq_len_q, seed=7)
+    case.build_references("trace")
+    iterations = 48
+    plan = [(i % NUM_SETS, (i // NUM_SETS) % 2, 20000, 3000) for i in range(iterations)]
+    outs = torch.full(
+        (iterations, case.rows, NUM_HEADS * KV_LORA_RANK),
+        float("nan"),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    stamps = torch.zeros(iterations, 8, dtype=torch.int64, device="cuda")
+    case.apply_plain(NUM_SETS - 1, 1)
+    torch.cuda.synchronize()
+    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
+        torch.cuda._sleep(GATE_CYCLES // 4)
+        _enqueue_chain(case, helpers, plan, outs, stamps, True)
+        torch.cuda.synchronize()
+    assert not _mismatches(case, plan, outs)
+    kernels = sorted(
+        (
+            (e.time_range.start, e.time_range.end, e.name)
+            for e in prof.events()
+            if e.device_type == torch.autograd.DeviceType.CUDA
+        ),
+        key=lambda x: x[0],
+    )
+    producers = [k for k in kernels if "delayedScatter" in k[2]]
+    sparse_kernels = [k for k in kernels if "sparseMlaDecode" in k[2]]
+    consumers = [k for k in kernels if "pdlConsumer" in k[2]]
+    count = min(len(producers), len(sparse_kernels), len(consumers))
+    sparse_mla_early = sum(1 for i in range(count) if sparse_kernels[i][0] < producers[i][1])
+    consumer_early = sum(1 for i in range(count) if consumers[i][0] < sparse_kernels[i][1])
+    record(
+        RESULTS,
+        "early_launch_trace",
+        {
+            "kernels": {
+                "producer": len(producers),
+                "sparse_mla_decode": len(sparse_kernels),
+                "consumer": len(consumers),
+            },
+            "sparse_mla_started_before_producer_end": sparse_mla_early / max(count, 1),
+            "consumer_started_before_sparse_mla_end": consumer_early / max(count, 1),
+            "sparse_mla_start_minus_producer_end_us_median": statistics.median(
+                sparse_kernels[i][0] - producers[i][1] for i in range(count)
+            )
+            if count
+            else None,
+            **_stamp_summary(stamps, plan),
+        },
+    )
+    assert count == iterations
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py
new file mode 100644
index 000000000000..7f9ee560001d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py
@@ -0,0 +1,99 @@
+# 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.
+"""sparse_mla_decode under DSA ragged verification (per-request verify widths).
+
+With ragged verification enabled in the generation metadata, ``TrtllmAttention``
+selects and runs the FMHA library inside the token-major presentation: one row
+per query token, ``seq_len_q`` 1, and that row's causal KV length. sparse_mla_decode engages
+there; these tests check it against the native library and the FP64 oracle with
+the frozen error budget of ``test_sparse_mla_sparse_mla.py``, check the lengths sparse_mla_decode
+receives, and check that a ragged step whose widths are uniform gives the same
+sparse_mla_decode output as the uniform layout.
+"""
+
+import pytest
+import torch
+from sparse_mla_harness import DsaGenHarness, compare
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available(),
+    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
+)
+
+
+@pytest.mark.parametrize(
+    "ctx_lens,verify_lens",
+    [
+        ([8192, 300, 40000], [6, 2, 4]),  # mixed widths, one short request
+        ([391000], [3]),  # one request with a partial verify window
+        ([100, 2043], [5, 6]),  # rows straddling the top-k length
+        ([20000, 20000, 20000], [1, 6, 1]),
+    ],
+)
+def test_sparse_mla_ragged_verification_matches_native(ctx_lens, verify_lens):
+    harness = DsaGenHarness(
+        ctx_lens, seq_len_q=max(verify_lens), seed=121, ragged_verify_lens=verify_lens
+    )
+    try:
+        assert harness.metadata.token_major_gen_view() is not None
+        topk = harness.random_topk(seed=122, overlap=0.7, recent=32)
+        step = harness.make_step(0, seed=123)
+        with harness.capture_sparse_mla_inputs() as calls:
+            compare(
+                harness,
+                step,
+                topk,
+                "ragged_verify_matches_native",
+                ctx_lens=ctx_lens,
+                verify_lens=verify_lens,
+            )
+        assert len(calls) == 1, len(calls)
+        call = calls[0]
+        # One presented row per query token, each with its causal length.
+        assert call["ragged"] and call["seq_len_q"] == 1, call
+        expected = torch.tensor(harness.list_limits(), dtype=torch.int32)
+        assert torch.equal(call["seq_lens_kv"].cpu(), expected), (call["seq_lens_kv"], expected)
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_ragged_presentation_equals_uniform():
+    """Uniform widths through the ragged presentation give the uniform layout's sparse_mla_decode output."""
+    ctx_lens = [32768, 5000]
+    outputs = {}
+    queries = {}
+    for ragged in (False, True):
+        harness = DsaGenHarness(
+            ctx_lens, seq_len_q=6, seed=131, ragged_verify_lens=[6, 6] if ragged else None
+        )
+        try:
+            assert (harness.metadata.token_major_gen_view() is not None) == ragged
+            topk = harness.random_topk(seed=132, overlap=0.7, recent=32)
+            step = harness.make_step(0, seed=133)
+            harness.select_library(0, use_sparse_mla=True)
+            outputs[ragged] = harness.forward(0, step, topk).clone()
+            assert harness.selected_library(0) == ["SparseMlaDecodeFmha"]
+            queries[ragged] = (step.quant_q.clone(), step.bmm1_scale.clone())
+        finally:
+            harness.shutdown()
+    torch.testing.assert_close(queries[True][0], queries[False][0], rtol=0, atol=0)
+    torch.testing.assert_close(queries[True][1], queries[False][1], rtol=0, atol=0)
+    torch.testing.assert_close(outputs[True], outputs[False], rtol=0, atol=0)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py
new file mode 100644
index 000000000000..c9788cb648a4
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py
@@ -0,0 +1,404 @@
+# 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.
+"""sparse_mla_decode against native on real GLM-5.2 layer-0 attention inputs (see ``sparse_mla_real_fixture``).
+
+The context latents are written by the native context step and the generation
+rows by the native ``mla_rope_generation``. Two sources of the top-k:
+
+* ``test_sparse_mla_real_layer0``: native indexer projections + a reference FP32
+  logit evaluation and exact top-k (not the native top-k kernels);
+* ``test_sparse_mla_real_layer0_native_indexer``: the unmodified native indexer
+  (native indexer K cache from the context and generation rows, native paged
+  MQA logits and decode top-k) under the served GEN indexer configuration
+  and under the CUDA radix top-k; native and sparse_mla_decode consume the same native
+  selection. This test also checks the premise of sparse_mla_decode's one deliberate
+  difference (no ``-1`` inside a row's causal length) on the native top-k, on
+  the global indices sparse_mla_decode receives, and on the native grouped remap of the same
+  selection. With ``rope_kv_append`` selected, rope_kv_append's fused top-k bridge replaces the
+  native remap and produces the indices sparse_mla_decode reads: every generation call of the bridge
+  must then equal the native remap of the same local top-k bit for bit, emit no
+  ``-1`` below a row's causal length, and be what sparse_mla_decode received.
+
+The same error budget as ``test_sparse_mla_sparse_mla.py`` applies.
+"""
+
+import pytest
+import torch
+from sparse_mla_harness import (
+    DsaGenHarness,
+    capture_rope_append_bridge,
+    check_budget,
+    compare,
+    error_stats,
+    negative_inside_causal,
+    oracle_for,
+    record,
+    rope_append_bridge_report,
+    rope_append_selected,
+)
+from sparse_mla_real_fixture import (
+    checkpoint_dir,
+    indexer_projections,
+    indexer_topk,
+    layer_weight_names,
+    load_indexer_weights,
+    load_tensors,
+    project_layer0,
+    real_token_ids,
+    recent_window_topk,
+    reference_indexer_logits,
+    selection_vs_kernel_scores,
+    selection_vs_reference,
+    topk_contract,
+)
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available()
+    or checkpoint_dir() is None,
+    reason="needs an SM10x GPU, the sparse_mla_decode kernel and the GLM-5.2 NVFP4 checkpoint",
+)
+
+SEQ_LEN_Q = 6
+NATIVE_INDEXER_CONTEXTS = [100, 2043, 8192, 32768, 131072]
+# Long contexts run in their own test (and process). The native context attention step of this
+# harness faults at these lengths, before any generation-phase code runs (one 491K-token step,
+# or 64K-token chunks of a 391K context), so these cases write only the
+# native indexer K cache for the context tokens (native ``_update_k_cache``, 64K-token chunks)
+# and keep random latent rows in the pool: the selection is the native indexer's on real
+# indexer keys and queries, the attended latents are synthetic.
+NATIVE_INDEXER_LONG_CONTEXTS = [391000, 491061]
+CONTEXT_CHUNK_TOKENS = 65536
+# Indexer kernels of the served GEN role and the exact
+# CUDA radix decode top-k (indexerTopK.cu) that the default configuration selects.
+NATIVE_INDEXER_CONFIGS = {
+    "tp8_decode": dict(
+        enable_heuristic_topk=True, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=True
+    ),
+    "cuda_radix": dict(
+        enable_heuristic_topk=False, use_cute_dsl_paged_mqa_logits=False, use_cute_dsl_topk=False
+    ),
+    # Diagnostic only: the served CuTe-DSL paged MQA logits with the exact CUDA radix top-k.
+    "dsl_logits_cuda_radix": dict(
+        enable_heuristic_topk=False, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=False
+    ),
+}
+
+
+@pytest.fixture(scope="module")
+def layer0():
+    ckpt = checkpoint_dir()
+    max_tokens = max(NATIVE_INDEXER_CONTEXTS + NATIVE_INDEXER_LONG_CONTEXTS) + SEQ_LEN_Q
+    weights = load_tensors(ckpt, layer_weight_names(0))
+    ids = real_token_ids(ckpt, max_tokens)
+    inputs = project_layer0(weights, ids)
+    yield weights, inputs
+    del weights, inputs
+    torch.cuda.empty_cache()
+
+
+@pytest.mark.parametrize("ctx_len", [32768, 131072])
+def test_sparse_mla_real_layer0(layer0, ctx_len):
+    weights, inputs = layer0
+    harness = DsaGenHarness(
+        [ctx_len], seq_len_q=SEQ_LEN_Q, fill_random=False, defer_generation=True
+    )
+    try:
+        load_indexer_weights(harness.layers[0].indexer, weights)
+        harness.run_context(
+            0,
+            inputs.fused_q[:ctx_len],
+            inputs.q_pe[:ctx_len],
+            inputs.latent[:ctx_len],
+            recent_window_topk(ctx_len),
+        )
+        harness.start_generation()
+        rows = slice(ctx_len, ctx_len + SEQ_LEN_Q)
+        step = harness.make_step_from(
+            0, inputs.fused_q[rows], inputs.q_pe[rows], inputs.latent[rows]
+        )
+        topk = indexer_topk(harness.layers[0].indexer, inputs, ctx_len, SEQ_LEN_Q)
+        # Selection structure: overlap between consecutive MTP rows and spread over the context.
+        sets = [set(topk[r].tolist()) for r in range(SEQ_LEN_Q)]
+        overlap = [len(sets[r] & sets[r - 1]) / 2048 for r in range(1, SEQ_LEN_Q)]
+        recent = [sum(1 for x in sets[r] if x >= ctx_len - 2048) / 2048 for r in range(SEQ_LEN_Q)]
+        result = compare(
+            harness,
+            step,
+            topk,
+            "real_layer0_reference_topk",
+            ctx_len=ctx_len,
+            mtp_row_overlap=overlap,
+            recent_2048_fraction=recent,
+        )
+        # Softmax peakedness of the real scores, per row and head (exact FP64 model).
+        global_idx = result["global_idx"]
+        q = step.quant_q.view(SEQ_LEN_Q, 8, 576).view(torch.float8_e4m3fn).double()
+        pool = harness.pool_rows().view(torch.uint8)
+        top_prob = []
+        for r in range(SEQ_LEN_Q):
+            keys = pool[global_idx[r][global_idx[r] >= 0].long()].view(torch.float8_e4m3fn).double()
+            probs = torch.softmax(q[r] @ keys.T * float(step.bmm1_scale[0]), dim=-1)
+            top_prob.append(probs.max(dim=-1).values.mean().item())
+        record(
+            "sparse_mla_correctness.jsonl",
+            "real_layer0_scores",
+            {"ctx_len": ctx_len, "mean_max_prob_per_row": top_prob},
+        )
+        assert error_stats(result["sparse_mla_decode"].double(), result["oracle"])["nonfinite"] == 0
+    finally:
+        harness.shutdown()
+
+
+@pytest.mark.parametrize("ctx_len", NATIVE_INDEXER_CONTEXTS)
+@pytest.mark.parametrize("config", ["tp8_decode", "cuda_radix"])
+def test_sparse_mla_real_layer0_native_indexer(layer0, config, ctx_len):
+    _run_native_indexer_case(layer0, config, ctx_len)
+
+
+@pytest.mark.parametrize("ctx_len", NATIVE_INDEXER_LONG_CONTEXTS)
+@pytest.mark.parametrize("config", ["tp8_decode", "cuda_radix"])
+def test_sparse_mla_real_layer0_native_indexer_long(layer0, config, ctx_len):
+    _run_native_indexer_case(
+        layer0, config, ctx_len, chunk_tokens=CONTEXT_CHUNK_TOKENS, real_latents=False
+    )
+
+
+@pytest.mark.parametrize(
+    "config,seq_len_q,mtp",
+    [
+        ("tp8_decode", 1, 0),
+        ("tp8_decode", 4, 3),
+        ("dsl_logits_cuda_radix", 6, 5),
+        ("dsl_logits_cuda_radix", 1, 0),
+    ],
+)
+def test_sparse_mla_real_layer0_native_indexer_variants(layer0, config, seq_len_q, mtp):
+    """Diagnostic variants at 32K: verify-window width and the logits / top-k kernel pair."""
+    _run_native_indexer_case(layer0, config, 32768, seq_len_q=seq_len_q, mtp=mtp)
+
+
+def _run_native_indexer_case(
+    layer0, config, ctx_len, chunk_tokens=None, real_latents=True, seq_len_q=SEQ_LEN_Q, mtp=None
+):
+    weights, inputs = layer0
+    harness = DsaGenHarness(
+        [ctx_len],
+        seq_len_q=seq_len_q,
+        seed=17,
+        fill_random=not real_latents,
+        defer_generation=True,
+        sparse_overrides=NATIVE_INDEXER_CONFIGS[config],
+        mtp_draft_len=seq_len_q - 1 if mtp is None else mtp,
+    )
+    try:
+        indexer = harness.layers[0].indexer
+        load_indexer_weights(indexer, weights)
+        ctx_proj = indexer_projections(indexer, inputs, 0, ctx_len)
+        harness.run_context(
+            0,
+            inputs.fused_q[:ctx_len],
+            inputs.q_pe[:ctx_len],
+            inputs.latent[:ctx_len],
+            recent_window_topk(ctx_len) if real_latents else None,
+            indexer_k=(ctx_proj[1], ctx_proj[2]),
+            chunk_tokens=chunk_tokens,
+            attention=real_latents,
+        )
+        del ctx_proj
+        harness.start_generation()
+        rows = slice(ctx_len, ctx_len + seq_len_q)
+        step = harness.make_step_from(
+            0, inputs.fused_q[rows], inputs.q_pe[rows], inputs.latent[rows]
+        )
+        gen_proj = indexer_projections(indexer, inputs, ctx_len, ctx_len + seq_len_q)
+        harness.update_indexer_cache(0, gen_proj[1], gen_proj[2])
+
+        harness.select_library(0, use_sparse_mla=False)
+        with capture_rope_append_bridge() as bridge_native:
+            native, topk_native = harness.forward_native_indexer(0, step, gen_proj)
+        native = native.clone()
+        native_topk_inputs = harness.last_topk_inputs
+        assert "SparseMlaDecodeFmha" not in harness.selected_library(0)
+        harness.select_library(0, use_sparse_mla=True)
+        with (
+            harness.capture_sparse_mla_inputs() as sparse_mla_calls,
+            capture_rope_append_bridge() as bridge_sparse_mla,
+        ):
+            sparse_out, topk_sparse_mla = harness.forward_native_indexer(0, step, gen_proj)
+        sparse_out = sparse_out.clone()
+        torch.cuda.synchronize()
+        assert harness.selected_library(0) == ["SparseMlaDecodeFmha"], harness.selected_library(0)
+        assert len(sparse_mla_calls) == 1, len(sparse_mla_calls)
+        received = sparse_mla_calls[0]
+
+        contract = topk_contract(topk_native, ctx_len)
+        reference_logits = reference_indexer_logits(indexer, inputs, ctx_len, seq_len_q)
+        agreement = []
+        for r in range(seq_len_q):
+            native_ids = set(x for x in topk_native[r].tolist() if x >= 0)
+            count = min(2048, reference_logits[r].numel())
+            reference_ids = set(reference_logits[r].topk(count).indices.tolist())
+            agreement.append(len(native_ids & reference_ids) / max(len(reference_ids), 1))
+        same_selection = torch.equal(topk_native, topk_sparse_mla)
+        selection = {
+            "seq_len_q": seq_len_q,
+            "max_draft_tokens": harness.metadata.max_draft_tokens,
+            "dsl_expand": [
+                bool(harness.metadata.expand_for_dsl),
+                harness.metadata.dsl_expand_factor,
+                harness.metadata.dsl_atom,
+            ],
+            "decode_top_k": indexer.top_k.decode_implementation.value,
+            "self_sampling": bool(getattr(indexer, "_use_self_sampling_topk", False)),
+            "same_selection_in_both_runs": same_selection,
+            "rows": contract,
+            "agreement_with_reference_topk": agreement,
+            "native_selection_under_reference_logits": selection_vs_reference(
+                topk_native, reference_logits
+            ),
+            "native_selection_vs_kernel_scores": (
+                selection_vs_kernel_scores(topk_native, native_topk_inputs, reference_logits)
+                if native_topk_inputs is not None
+                else None
+            ),
+        }
+        global_native = harness.global_indices(0, topk_native)
+        oracle = oracle_for(harness, step, global_native)
+        # Producer contract at sparse_mla_decode's input: the global indices it received (the native
+        # remap of its run's selection) and the native grouped remap of the native run's selection.
+        grouped_native = harness.grouped_global_indices(0, topk_native)
+        received_negative = negative_inside_causal(
+            received["indices"], received["seq_lens_kv"], received["seq_len_q"]
+        )
+        grouped_negative = negative_inside_causal(
+            grouped_native, received["seq_lens_kv"], received["seq_len_q"]
+        )
+        selection["sparse_mla_input"] = {
+            "negative_below_causal": received_negative.tolist(),
+            "equals_remap_of_its_selection": torch.equal(
+                received["indices"], harness.global_indices(0, topk_sparse_mla)
+            ),
+            "seq_len_q": received["seq_len_q"],
+        }
+        selection["grouped_remap"] = {
+            "equals_per_layer_remap": torch.equal(grouped_native, global_native),
+            "negative_below_causal": grouped_negative.tolist(),
+        }
+        # rope_kv_append's bridge (when selected) is the producer of the indices of both runs.
+        rope_append_bridge = {
+            "selected": rope_append_selected(),
+            "calls_native_run": len(bridge_native),
+            "calls_sparse_mla_run": len(bridge_sparse_mla),
+            "calls": [
+                rope_append_bridge_report(call, received["seq_lens_kv"], received["seq_len_q"])
+                for call in bridge_native + bridge_sparse_mla
+            ],
+            "sparse_mla_received_bridge_output": (
+                bool(torch.equal(received["indices"], bridge_sparse_mla[0]["global"]))
+                if bridge_sparse_mla
+                else None
+            ),
+        }
+        selection["rope_append_bridge"] = rope_append_bridge
+        if same_selection:
+            check_budget(
+                harness,
+                native,
+                sparse_out,
+                global_native,
+                oracle,
+                "real_layer0_native_indexer",
+                config=config,
+                ctx_len=ctx_len,
+                context_chunk_tokens=chunk_tokens,
+                real_context_latents=real_latents,
+                selection=selection,
+            )
+        else:
+            # The native selection is not deterministic run to run. sparse_mla_decode is checked
+            # twice: on the exact selection the native run consumed (the native indexer output
+            # tensor fed to sparse_mla_decode through the fixed-top-k path) and, below, on the
+            # selection of its own run.
+            selection["sparse_mla_run_selection_overlap"] = [
+                len(set(topk_native[r].tolist()) & set(topk_sparse_mla[r].tolist()))
+                / max(len(set(topk_native[r].tolist())), 1)
+                for r in range(seq_len_q)
+            ]
+            sparse_mla_same = harness.forward(0, step, topk_native).clone()
+            assert harness.selected_library(0) == ["SparseMlaDecodeFmha"]
+            check_budget(
+                harness,
+                native,
+                sparse_mla_same,
+                global_native,
+                oracle,
+                "real_layer0_native_indexer_same_selection",
+                config=config,
+                ctx_len=ctx_len,
+                context_chunk_tokens=chunk_tokens,
+                real_context_latents=real_latents,
+                selection=selection,
+            )
+            global_sparse_mla = harness.global_indices(0, topk_sparse_mla)
+            native_err = error_stats(native.double(), oracle)
+            sparse_mla_err = error_stats(
+                sparse_out.double(), oracle_for(harness, step, global_sparse_mla)
+            )
+            record(
+                "sparse_mla_correctness.jsonl",
+                "real_layer0_native_indexer",
+                {
+                    "config": config,
+                    "ctx_len": ctx_len,
+                    "context_chunk_tokens": chunk_tokens,
+                    "real_context_latents": real_latents,
+                    "selection": selection,
+                    "native_vs_oracle": native_err,
+                    "sparse_mla_vs_oracle": sparse_mla_err,
+                },
+            )
+            assert sparse_mla_err["nonfinite"] == native_err["nonfinite"] == 0
+            assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
+                sparse_mla_err,
+                native_err,
+            )
+        # The premise of sparse_mla_decode's one deliberate difference from native: the native
+        # indexer emits no -1 inside a row's causal length, and neither remap creates one.
+        assert all(row["negative_below_causal"] == 0 for row in contract), contract
+        assert int(received_negative.sum()) == 0, selection["sparse_mla_input"]
+        assert selection["sparse_mla_input"]["equals_remap_of_its_selection"]
+        assert selection["grouped_remap"]["equals_per_layer_remap"]
+        assert int(grouped_negative.sum()) == 0, selection["grouped_remap"]
+        if rope_append_bridge["selected"]:
+            # The bridge served the generation call of both runs (counted per call: the
+            # logged rope_kv_append/bridge decision is the first one, the declined context step).
+            assert (
+                rope_append_bridge["calls_native_run"] == 1
+                and rope_append_bridge["calls_sparse_mla_run"] == 1
+            ), rope_append_bridge
+            for call in rope_append_bridge["calls"]:
+                assert call["equals_native_remap"], rope_append_bridge
+                assert sum(call["negative_below_causal"]) == 0, rope_append_bridge
+            assert rope_append_bridge["sparse_mla_received_bridge_output"], rope_append_bridge
+        else:
+            assert not bridge_native and not bridge_sparse_mla
+    finally:
+        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py
new file mode 100644
index 000000000000..c92aba441f83
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py
@@ -0,0 +1,195 @@
+# 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.
+"""rope_kv_append's fused top-k bridge as the producer of sparse_mla_decode's global indices.
+
+With the plan running ``sparse_mla_decode`` and ``rope_kv_append`` (installed by this test) the DSA
+predictor's local-to-global remap is ``remap_topk_to_global``
+(``trtllm::glm_moe_dsa_topk_index_bridge``), and its output is the index list sparse_mla_decode
+reads. sparse_mla_decode differs from the native kernel on one
+input class: a ``-1`` at an index-list position below the row's causal length (sparse_mla_decode
+skips it, the native kernel reads the pool row that ``-1`` addresses). For every
+generation call the bridge serves, this test asserts that
+
+* the bridge ran (counted per call, not read from the capability log, whose
+  first ``rope_kv_append/bridge`` decision is the declined context step),
+* its global indices equal the native remap of the same local top-k bit for
+  bit and contain no ``-1`` below the row's causal length,
+* sparse_mla_decode received exactly those indices, and sparse_mla_decode stays within the frozen
+  error budget against the FP64 oracle,
+* CUPTI shows rope_kv_append's ``topkBridgeKernel`` and sparse_mla_decode's
+  ``sparseMlaDecodeKernel`` once per generation call, and neither the native remap kernel nor the
+  native sparse FMHA kernel.
+
+A negative control puts ``-1`` inside the causal range of one row: the bridge
+passes it through exactly like the native remap, and the check above reports
+it, so the assertion is not vacuous.
+
+Runs only when both ops are in the build.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+from sparse_mla_harness import (
+    SPARSE_MLA_KERNEL,
+    DsaGenHarness,
+    capture_rope_append_bridge,
+    check_budget,
+    is_native_sparse_fmha,
+    kernel_trace,
+    negative_inside_causal,
+    oracle_for,
+    record,
+    rope_append_bridge_report,
+)
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import plan_testlib  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available()
+    or not hasattr(torch.ops.trtllm, "glm_moe_dsa_topk_index_bridge"),
+    reason="needs an SM10x build with the sparse_mla_decode and rope_kv_append kernels",
+)
+
+
+@pytest.fixture(autouse=True)
+def _plan(monkeypatch):
+    plan_testlib.select(dispatch, "sparse_mla_decode,rope_kv_append", monkeypatch)
+
+
+RESULTS = "sparse_mla_rope_append_bridge.jsonl"
+CASES = {
+    "ctx100_r6": ([100], 6),  # every row ends in -1 entries past its causal length
+    "ctx2042_r6": ([2042], 6),  # the last row sees exactly 2048 tokens
+    "ctx2043_r6": ([2043], 6),  # rows straddle 2048 visible tokens
+    "two_seqs_r6": ([5000, 300], 6),  # the second sequence has -1 tails
+    "four_seqs_r2": ([20000, 20000, 20000, 20000], 2),
+    "ctx32768_r6": ([32768], 6),  # target verify, MTP K = 5
+    "ctx131072_r1": ([131072], 1),  # draft steps 1-4
+}
+
+
+def _generation_kernels(harness, step, topk):
+    kernels = kernel_trace(lambda: harness.forward(0, step, topk))
+    names = [k["name"] for k in kernels]
+    return {
+        "rope_append_bridge": sum("topkBridgeKernel" in n for n in names),
+        "native_remap": sum("convertReqIndexToGlobal" in n for n in names),
+        "sparse_mla_decode": sum(SPARSE_MLA_KERNEL in n for n in names),
+        "native_sparse_fmha": sum(is_native_sparse_fmha(n) for n in names),
+    }
+
+
+@pytest.mark.parametrize("case", list(CASES))
+def test_rope_append_bridge_feeds_sparse_mla(case):
+    ctx_lens, seq_len_q = CASES[case]
+    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=31, mtp_draft_len=seq_len_q - 1)
+    try:
+        step = harness.make_step(0, seed=32)
+        topk = harness.random_topk(seed=33, overlap=0.7, recent=32)
+        seq_lens_kv = torch.tensor(
+            [c + seq_len_q for c in ctx_lens], dtype=torch.int32, device="cuda"
+        )
+        with (
+            capture_rope_append_bridge() as bridge,
+            harness.capture_sparse_mla_inputs() as sparse_mla_calls,
+        ):
+            harness.select_library(0, use_sparse_mla=False)
+            native = harness.forward(0, step, topk).clone()
+            harness.select_library(0, use_sparse_mla=True)
+            sparse_out = harness.forward(0, step, topk).clone()
+            torch.cuda.synchronize()
+        assert harness.selected_library(0) == ["SparseMlaDecodeFmha"]
+        # One bridge call per generation forward (native arm, then sparse_mla_decode arm).
+        assert len(bridge) == 2, len(bridge)
+        assert len(sparse_mla_calls) == 1, len(sparse_mla_calls)
+        received = sparse_mla_calls[0]
+        assert torch.equal(received["seq_lens_kv"], seq_lens_kv)
+        reports = [rope_append_bridge_report(call, seq_lens_kv, seq_len_q) for call in bridge]
+        same_as_bridge = torch.equal(received["indices"], bridge[1]["global"])
+        global_idx = harness.global_indices(0, topk)
+        oracle = oracle_for(harness, step, global_idx)
+        kernels = _generation_kernels(harness, step, topk)
+        check_budget(
+            harness,
+            native,
+            sparse_out,
+            global_idx,
+            oracle,
+            "rope_append_bridge_feeds_sparse_mla",
+            RESULTS,
+            shape=case,
+            bridge=reports,
+            sparse_mla_received_bridge_output=same_as_bridge,
+            generation_kernels=kernels,
+        )
+        for report in reports:
+            assert report["equals_native_remap"], report
+            assert sum(report["negative_below_causal"]) == 0, report
+        assert same_as_bridge
+        assert kernels == {
+            "rope_append_bridge": 1,
+            "native_remap": 0,
+            "sparse_mla_decode": 1,
+            "native_sparse_fmha": 0,
+        }, kernels
+    finally:
+        harness.shutdown()
+
+
+def test_rope_append_bridge_negative_control():
+    """``-1`` injected inside the causal range: passed through as by the native remap, detected."""
+    ctx_lens, seq_len_q = [32768], 6
+    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=41, mtp_draft_len=seq_len_q - 1)
+    try:
+        step = harness.make_step(0, seed=42)
+        topk = harness.random_topk(seed=43, overlap=0.7, recent=32)
+        topk[2, ::3] = -1
+        injected = int((topk[2] < 0).sum().item())
+        seq_lens_kv = torch.tensor([ctx_lens[0] + seq_len_q], dtype=torch.int32, device="cuda")
+        with (
+            capture_rope_append_bridge() as bridge,
+            harness.capture_sparse_mla_inputs() as sparse_mla_calls,
+        ):
+            harness.select_library(0, use_sparse_mla=True)
+            harness.forward(0, step, topk)
+            torch.cuda.synchronize()
+        assert len(bridge) == 1 and len(sparse_mla_calls) == 1
+        report = rope_append_bridge_report(bridge[0], seq_lens_kv, seq_len_q)
+        received = negative_inside_causal(
+            sparse_mla_calls[0]["indices"], seq_lens_kv, seq_len_q
+        ).tolist()
+        record(
+            RESULTS,
+            "rope_append_bridge_negative_control",
+            {"injected_row2": injected, "bridge": report, "sparse_mla_received_negative": received},
+        )
+        assert report["equals_native_remap"], report
+        expected = [0] * harness.num_rows
+        expected[2] = injected
+        assert report["negative_below_causal"] == expected, report
+        assert received == expected
+    finally:
+        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py
new file mode 100644
index 000000000000..df9ca44f0711
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py
@@ -0,0 +1,150 @@
+# 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.
+"""Small sparse_mla_decode cases for compute-sanitizer (memcheck, racecheck, synccheck, initcheck).
+
+Run under ``compute-sanitizer --tool <tool> --kernel-name kns=sparseMlaDecode``
+so that only the sparse_mla_decode kernel is instrumented. The cases cover every code path of
+the kernel at shapes small enough for instrumented execution: short rows with
+fewer usable positions than the top-k (zero-filled gather slots), rows that
+straddle 2048 usable positions, full rows, R = 6 and R = 1, several sequences,
+empty rows, negative and repeated ids, and CUDA-graph capture plus replay with
+rewritten inputs. Each case also checks the output against the native library
+and the FP64 oracle with the frozen error budget of ``test_sparse_mla_sparse_mla.py``.
+"""
+
+import pytest
+import torch
+from sparse_mla_harness import DsaGenHarness, compare, error_stats, oracle_attention, record
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available(),
+    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
+)
+
+RESULTS = "sparse_mla_sanitizer_cases.jsonl"
+
+
+@pytest.mark.parametrize(
+    "ctx_lens,seq_len_q",
+    [
+        ([100], 6),
+        ([100], 1),
+        ([2043], 6),
+        ([2043], 1),
+        ([8192], 6),
+        ([8192], 1),
+        ([3000, 100, 9000], 2),
+    ],
+)
+def test_sparse_mla_small_matches_native(ctx_lens, seq_len_q):
+    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=11)
+    try:
+        topk = harness.random_topk(seed=12, overlap=0.7, recent=32)
+        step = harness.make_step(0, seed=13)
+        compare(
+            harness,
+            step,
+            topk,
+            "small_matches_native",
+            RESULTS,
+            ctx_lens=ctx_lens,
+            seq_len_q=seq_len_q,
+        )
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_small_empty_negative_and_repeated():
+    harness = DsaGenHarness([4096, 40], seq_len_q=6, seed=14)
+    try:
+        topk = harness.random_topk(seed=15)
+        topk[1, :] = -1
+        topk[3, :1500] = 17
+        topk[3, 1500:] = 4242
+        for row in range(6, 12):
+            for pos in (300, 1500, 2047):
+                topk[row, pos] = 3
+        step = harness.make_step(0, seed=16)
+        compare(harness, step, topk, "small_empty_negative_repeated", RESULTS)
+        # Negative ids inside the causal range: sparse_mla_decode skips them. Only
+        # sparse_mla_decode's output is checked here (the native kernel reads the row that -1
+        # addresses).
+        topk = harness.random_topk(seed=17)
+        topk[2, ::3] = -1
+        topk[5, :1024] = -1
+        harness.select_library(0, use_sparse_mla=True)
+        sparse_out = harness.forward(0, step, topk).clone()
+        global_idx = harness.global_indices(0, topk)
+        oracle = oracle_attention(
+            step.quant_q,
+            harness.pool_rows(),
+            global_idx,
+            float(step.bmm1_scale[0]),
+            float(step.bmm2_scale[0]),
+            list_limit=harness.list_limits(),
+        )
+        err = error_stats(sparse_out.double(), oracle)
+        record(RESULTS, "small_negative_inside_range_sparse_mla", {"sparse_mla_vs_oracle": err})
+        assert err["nonfinite"] == 0 and err["nrms"] < 0.05, err
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_small_graph_replay():
+    harness = DsaGenHarness([8192], seq_len_q=6, seed=18)
+    try:
+        step = harness.make_step(0, seed=19)
+        topk = harness.random_topk(seed=20, overlap=0.7, recent=32)
+        harness.select_library(0, use_sparse_mla=True)
+        harness.forward(0, step, topk)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            out = harness.forward(0, step, topk)
+        pool_rows = harness.pool_rows()
+        for it in range(3):
+            topk.copy_(harness.random_topk(seed=300 + it, overlap=0.5, recent=16))
+            global_idx = harness.global_indices(0, topk)
+            rows = global_idx[global_idx >= 0].unique().long()
+            fresh = (
+                (torch.randn(rows.numel(), pool_rows.shape[1], device="cuda") * (1 + it))
+                .clamp(-448, 448)
+                .to(torch.float8_e4m3fn)
+            )
+            pool_rows.view(torch.uint8)[rows] = fresh.view(torch.uint8)
+            fresh_step = harness.make_step(0, seed=400 + it)
+            for name in ("fused_q", "q_pe", "latent", "quant_q", "bmm1_scale", "bmm2_scale"):
+                getattr(step, name).copy_(getattr(fresh_step, name))
+            graph.replay()
+            torch.cuda.synchronize()
+            oracle = oracle_attention(
+                step.quant_q,
+                pool_rows,
+                global_idx,
+                float(step.bmm1_scale[0]),
+                float(step.bmm2_scale[0]),
+                list_limit=harness.list_limits(),
+            )
+            err = error_stats(out.double(), oracle)
+            record(RESULTS, "small_graph_replay", {"iteration": it, "sparse_mla_vs_oracle": err})
+            assert err["nonfinite"] == 0 and err["nrms"] < 0.05, err
+    finally:
+        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py
new file mode 100644
index 000000000000..7a64a7e43bdf
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py
@@ -0,0 +1,285 @@
+# 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.
+"""Correctness of the sparse_mla_decode sparse MLA library against the native path.
+
+Every case runs the native ``mla_rope_generation`` once and then the same
+``TrtllmAttention.forward`` twice: with the native FMHA library and with sparse_mla_decode
+first in the layer's library list. Both outputs are compared with an FP64
+oracle over the same e4m3 query and pool rows that applies the native masking
+rules (index-list positions past a row's causal length are ignored, negative
+entries skipped, repeats counted).
+
+Error budget, fixed before any tuning:
+  * nrms(sparse_mla_decode - oracle) <= 1.25 * nrms(native - oracle) + 2e-4
+  * max|sparse_mla_decode - oracle| <= 2 * max|native - oracle| + 1e-3
+  * the existing FP8-KV sparse MLA tolerance: allclose(sparse_mla_decode, oracle, atol=0.12,
+    rtol=0.01)
+  * rows without a usable entry and non-finite values must match the native output exactly.
+"""
+
+import pytest
+import torch
+from sparse_mla_harness import (
+    DsaGenHarness,
+    compare,
+    error_stats,
+    oracle_attention,
+    record,
+    run_both,
+)
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
+    SparseMlaDecodeFmha,
+)
+from tensorrt_llm._torch.attention.backends.interface import (
+    AttentionForwardArgs,
+    AttentionInputType,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available(),
+    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
+)
+
+
+@pytest.mark.parametrize(
+    "ctx_lens,seq_len_q,q_std",
+    [
+        ([32768], 6, 1.0),  # target verify, MTP K=5
+        ([32768], 1, 1.0),  # draft steps 1-4
+        ([32768], 6, 4.0),  # peaky softmax
+        ([32768], 6, 0.1),  # nearly flat softmax
+        ([100], 6, 1.0),  # fewer visible tokens than top-k
+        ([2042], 6, 1.0),  # last row sees exactly 2048 tokens
+        ([2043], 6, 1.0),  # rows straddle 2048 visible tokens
+        ([5000, 300], 6, 1.0),  # two sequences, one short
+        ([20000, 20000, 20000, 20000], 2, 1.0),  # four sequences, MTP K=1
+        ([131072], 6, 1.0),
+        ([391000], 6, 1.0),  # served-length context
+        ([491061], 6, 1.0),  # controlled-decode context
+    ],
+)
+def test_sparse_mla_matches_native(ctx_lens, seq_len_q, q_std):
+    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=21)
+    try:
+        topk = harness.random_topk(seed=22, overlap=0.7, recent=32)
+        step = harness.make_step(0, seed=23, q_std=q_std)
+        compare(
+            harness,
+            step,
+            topk,
+            "matches_native",
+            ctx_lens=ctx_lens,
+            seq_len_q=seq_len_q,
+            q_std=q_std,
+        )
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_empty_rows_tails_and_duplicates():
+    """Index lists as the native indexer shapes them, plus repeated and empty rows."""
+    harness = DsaGenHarness([8192, 40], seq_len_q=6, seed=31)
+    try:
+        topk = harness.random_topk(seed=32)
+        topk[1, :] = -1  # no usable entry
+        topk[3, :1500] = 17  # one row repeated 1500 times ...
+        topk[3, 1500:] = 4242  # ... and another one 548 times
+        for row in range(6, 12):  # short sequence: valid ids past the causal length are ignored
+            for pos in (300, 1500, 2047):
+                topk[row, pos] = 3
+        step = harness.make_step(0, seed=33)
+        compare(harness, step, topk, "empty_rows_tails_and_duplicates")
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_causal_clamp_is_per_row():
+    """Valid ids at list positions between a row's causal length and the sequence length."""
+    harness = DsaGenHarness([100], seq_len_q=6, seed=41)
+    try:
+        topk = harness.random_topk(seed=42)
+        for row in range(harness.num_rows):
+            for pos in range(harness.causal_len(row), 106):
+                topk[row, pos] = 5
+        step = harness.make_step(0, seed=43)
+        result = compare(harness, step, topk, "causal_clamp_per_row")
+        global_idx = harness.global_indices(0, topk)
+        per_sequence = oracle_attention(
+            step.quant_q,
+            harness.pool_rows(),
+            global_idx,
+            float(step.bmm1_scale[0]),
+            float(step.bmm2_scale[0]),
+            list_limit=[106] * harness.num_rows,
+        )
+        native_rows = result["native"].double().reshape(result["oracle"].shape)
+        # Rows 0-4 differ between the two rules; the native kernel must follow the per-row one.
+        per_row_err = error_stats(native_rows[:5], result["oracle"][:5])
+        per_seq_err = error_stats(native_rows[:5], per_sequence[:5])
+        record(
+            "sparse_mla_correctness.jsonl",
+            "causal_clamp_native_rule",
+            {"per_row": per_row_err, "per_sequence": per_seq_err},
+        )
+        assert per_row_err["nrms"] < per_seq_err["nrms"], (per_row_err, per_seq_err)
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_skips_negative_entries_inside_range():
+    """Negative ids before the causal length are skipped (the native kernel does not check them)."""
+    harness = DsaGenHarness([8192], seq_len_q=6, seed=51)
+    try:
+        topk = harness.random_topk(seed=52)
+        topk[2, ::3] = -1
+        topk[4, 1000:] = -1
+        topk[5, :1024] = -1
+        step = harness.make_step(0, seed=53)
+        native, decode_out, global_idx, oracle = run_both(harness, step, topk)
+        per_row = {
+            r: {
+                "sparse_mla_decode": error_stats(decode_out[r].double(), oracle[r]),
+                "native": error_stats(native[r].double(), oracle[r]),
+            }
+            for r in range(harness.num_rows)
+        }
+        record("sparse_mla_correctness.jsonl", "negative_entries_inside_range", {"rows": per_row})
+        clean = max(per_row[r]["sparse_mla_decode"]["nrms"] for r in (0, 1, 3))
+        for r in (2, 4, 5):
+            assert per_row[r]["sparse_mla_decode"]["nrms"] <= 2.0 * clean, (r, per_row[r], clean)
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_cuda_graph_replay():
+    """Capture both libraries once; replay with new indices, queries and rewritten pool rows."""
+    harness = DsaGenHarness([65536], seq_len_q=6, seed=61)
+    try:
+        step = harness.make_step(0, seed=62)
+        topk = harness.random_topk(seed=63, overlap=0.7, recent=32)
+        graphs = {}
+        outputs = {}
+        for use_sparse_mla in (False, True):
+            harness.select_library(0, use_sparse_mla=use_sparse_mla)
+            harness.forward(0, step, topk)
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                outputs[use_sparse_mla] = harness.forward(0, step, topk)
+            graphs[use_sparse_mla] = graph
+        pool_rows = harness.pool_rows()
+        for it in range(6):
+            topk.copy_(harness.random_topk(seed=100 + it, overlap=0.5, recent=16))
+            if it % 2 == 1:
+                # Page reuse: overwrite the rows the next replay reads.
+                global_idx = harness.global_indices(0, topk)
+                rows = global_idx[global_idx >= 0].unique().long()
+                fresh = (
+                    (torch.randn(rows.numel(), pool_rows.shape[1], device="cuda") * (1 + it))
+                    .clamp(-448, 448)
+                    .to(torch.float8_e4m3fn)
+                )
+                pool_rows.view(torch.uint8)[rows] = fresh.view(torch.uint8)
+            fresh_step = harness.make_step(0, seed=200 + it)
+            for name in ("fused_q", "q_pe", "latent", "quant_q", "bmm1_scale", "bmm2_scale"):
+                getattr(step, name).copy_(getattr(fresh_step, name))
+            for use_sparse_mla in (False, True):
+                graphs[use_sparse_mla].replay()
+            torch.cuda.synchronize()
+            global_idx = harness.global_indices(0, topk)
+            oracle = oracle_attention(
+                step.quant_q,
+                pool_rows,
+                global_idx,
+                float(step.bmm1_scale[0]),
+                float(step.bmm2_scale[0]),
+                list_limit=harness.list_limits(),
+            )
+            native_err = error_stats(outputs[False].double(), oracle)
+            sparse_mla_err = error_stats(outputs[True].double(), oracle)
+            record(
+                "sparse_mla_correctness.jsonl",
+                "graph_replay",
+                {
+                    "iteration": it,
+                    "native_vs_oracle": native_err,
+                    "sparse_mla_vs_oracle": sparse_mla_err,
+                },
+            )
+            assert sparse_mla_err["nonfinite"] == 0
+            assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
+                it,
+                sparse_mla_err,
+                native_err,
+            )
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_falls_back_for_other_top_k():
+    """A request outside the domain (top-k 1024) must run the native library."""
+    harness = DsaGenHarness([16384], seq_len_q=6, index_topk=1024, seed=71)
+    try:
+        gen = torch.Generator(device="cpu").manual_seed(72)
+        topk = (
+            torch.stack(
+                [
+                    torch.randperm(harness.causal_len(r), generator=gen)[:1024]
+                    for r in range(harness.num_rows)
+                ]
+            )
+            .to(torch.int32)
+            .cuda()
+        )
+        step = harness.make_step(0, seed=73)
+        harness.select_library(0, use_sparse_mla=False)
+        native = harness.forward(0, step, topk).clone()
+        harness.select_library(0, use_sparse_mla=True)
+        fallback = harness.forward(0, step, topk).clone()
+        assert "SparseMlaDecodeFmha" not in harness.selected_library(0)
+        torch.testing.assert_close(fallback, native, rtol=0, atol=0)
+    finally:
+        harness.shutdown()
+
+
+def test_sparse_mla_static_eligibility():
+    fp8 = DsaGenHarness([256], seq_len_q=1, seed=81)
+    bf16 = DsaGenHarness([256], seq_len_q=1, kv_dtype=torch.bfloat16, seed=82)
+    try:
+        assert SparseMlaDecodeFmha.is_available(fp8.layers[0])
+        assert not SparseMlaDecodeFmha.is_available(bf16.layers[0])
+        assert sparse_mla.static_ineligibility(bf16.layers[0]) == "KV cache is not FP8"
+    finally:
+        fp8.shutdown()
+        bf16.shutdown()
+
+
+def test_sparse_mla_rejects_context_phase():
+    harness = DsaGenHarness([256], seq_len_q=1, seed=91)
+    try:
+        q = torch.empty(1, 8 * 576, dtype=torch.bfloat16, device="cuda")
+        args = AttentionForwardArgs(attention_input_type=AttentionInputType.context_only)
+        assert (
+            sparse_mla.request_ineligibility(q, None, None, harness.metadata, args)
+            == "not generation-only"
+        )
+    finally:
+        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py
new file mode 100644
index 000000000000..cebbefb99b37
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py
@@ -0,0 +1,350 @@
+# 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.
+"""sparse_mla_decode between engaged producers (absorbed_query, shared_indexer_absorbed_query,
+rope_kv_append) and its PDL consumer (the attention_output_tail).
+
+Each package set runs ``sparse_mla_stack_case.py`` in a fresh interpreter, without and
+with ``sparse_mla_decode``, on real GLM-5.2 layers at TP8 rank 0: layer 2 runs its own indexer
+(absorbed_query's layer), layer 3 reuses layer 2's top-k (shared_indexer_absorbed_query's layer).
+``MLA.forward`` runs the MTP verify step of one request (6 rows) as the model does. Asserted:
+
+* engagement, counted per call: absorbed_query returns its packet on the full layer and
+  shared_indexer_absorbed_query on the shared layer; rope_kv_append's fused RoPE and top-k bridge
+  serve every generation call; the attention_output_tail takes every call of both layers;
+  sparse_mla_decode runs on both layers;
+* the producers do not depend on the attention library: the FMHA inputs (FP8
+  query bytes, bmm scales, KV lengths, global indices and the pool rows they
+  address) are bit-identical without and with sparse_mla_decode;
+* sparse_mla_decode's latent output is within the frozen error budget against the FP64 oracle
+  on both layers, and the module outputs (after the attention_output_tail) agree with the
+  native-attention run;
+* sparse_mla_decode triggers its programmatic dependents before its gathers complete; the
+  attention_output_tail is that dependent. With sparse_mla_decode, the eager steps (no
+  recording kernels), the steps with a device synchronization right after each sparse_mla_decode
+  launch, and the replays of a captured step must give bit-identical latent and module outputs;
+* CUPTI: one ``sparseMlaDecodeKernel`` per DSA layer per step (eager and graph)
+  and no native sparse FMHA kernel with sparse_mla_decode; none without it; in graph replays
+  the kernel after sparse_mla_decode on its stream is attention_output_tail's W_UV kernel and it
+  becomes resident before sparse_mla_decode ends;
+* with rope_kv_append, every generation call of its bridge equals the native remap of the
+  same local top-k bit for bit, emits no ``-1`` below a row's causal length, and
+  is what the attention received.
+
+Neighbour sets: native, ``absorbed_query,shared_indexer_absorbed_query,attention_output_tail``
+(attention_output_tail default schedule) and the packages of the served selector that act at
+these call sites: absorbed_query, shared_indexer_absorbed_query, rope_kv_append,
+attention_output_tail, post_attention_hop, attention_output_tail_prefetch and
+attention_output_into_hop (with post_attention_hop and attention_output_into_hop the tail takes
+``INTO_HOP_SPARSE_DECODE_SCHEDULE`` behind sparse_mla_decode, which streams 32 ``w_o`` rows per
+CTA before its PDL wait). The last set also runs with the unmodified native indexer.
+
+Runs with ``GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION=1`` on an integration build that contains the
+packages, with the GLM-5.2 checkpoint.
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from sparse_mla_harness import error_stats, oracle_attention, record
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
+    sparse_mla,
+)
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import plan_testlib  # noqa: E402
+from utils.llm_data import llm_models_root
+
+_CKPT = os.environ.get(
+    "GLM52_NVFP4_CKPT",
+    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
+)
+_CASE = Path(__file__).resolve().parent / "sparse_mla_stack_case.py"
+_ATTN_TAIL_HARNESS = (
+    Path(__file__).resolve().parent.parent / "attention_output_tail" / "real_mla_harness.py"
+)
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION") != "1"
+    or not torch.cuda.is_available()
+    or torch.cuda.get_device_capability()[0] != 10
+    or not sparse_mla.op_available()
+    or not os.path.isfile(os.path.join(_CKPT, "config.json"))
+    or not _ATTN_TAIL_HARNESS.is_file(),
+    reason="set GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION=1 on an SM10x integration build with the checkpoint",
+)
+
+SERVED_SITE_SET = (
+    "absorbed_query,shared_indexer_absorbed_query,rope_kv_append,"
+    "attention_output_tail,post_attention_hop,attention_output_tail_prefetch,attention_output_into_hop"
+)
+STACK_SETS = [
+    ("", "fixed"),
+    ("absorbed_query,shared_indexer_absorbed_query,attention_output_tail", "fixed"),
+    (SERVED_SITE_SET, "fixed"),
+    (SERVED_SITE_SET, "indexer"),
+]
+ROWS = 6
+TOP_K = 2048
+FULL_LAYER, SHARED_LAYER = 2, 3
+TAIL_FEATURES = (
+    "attention_output_tail",
+    "attention_output_tail_prefetch",
+    "attention_output_into_hop",
+)
+
+
+def _run(features: str, topk_mode: str, out_dir: Path) -> dict:
+    out = out_dir / f"stack_{features.replace(',', '_') or 'native'}_{topk_mode}.pt"
+    env = dict(os.environ)
+    env[plan_testlib.FAST_PATHS_ENV] = features
+    env["GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_TOPK"] = topk_mode
+    env["GLM52_NVFP4_CKPT"] = _CKPT
+    env["TLLM_LOG_LEVEL"] = "INFO"
+    env.setdefault("TRTLLM_ENABLE_PDL", "1")
+    proc = subprocess.run(
+        [sys.executable, str(_CASE), str(out)], env=env, capture_output=True, text=True, check=False
+    )
+    assert proc.returncode == 0, proc.stdout[-8000:] + proc.stderr[-8000:]
+    result = torch.load(out, weights_only=False)
+    log = (proc.stdout + proc.stderr).splitlines()
+    result["capability_log"] = sorted(
+        {line[line.index("[glm_moe_dsa]") :] for line in log if "[glm_moe_dsa]" in line}
+    )
+    return result
+
+
+def _oracle(entry: dict) -> torch.Tensor:
+    rows = entry["rows"].cuda().reshape(ROWS * TOP_K, -1)
+    indices = entry["indices"].cuda()
+    local = torch.arange(ROWS * TOP_K, device="cuda", dtype=torch.int32).view(ROWS, TOP_K)
+    remapped = torch.where(indices >= 0, local, torch.full_like(local, -1))
+    kv_len = int(entry["kv_lens"][0])
+    limits = [kv_len - (ROWS - 1 - r) for r in range(ROWS)]
+    return oracle_attention(
+        entry["quant_q"].cuda(),
+        rows,
+        remapped,
+        float(entry["bmm1"][0]),
+        float(entry["bmm2"][0]),
+        list_limit=limits,
+    )
+
+
+def _same(a: torch.Tensor, b: torch.Tensor) -> bool:
+    return a.shape == b.shape and torch.equal(a.view(torch.uint8), b.view(torch.uint8))
+
+
+def _check_engagement(run: dict, selected: set, with_sparse_mla: bool) -> dict:
+    eng = run["engagement"]
+    counts = eng["counts"]
+    eager_steps = 2 + 2 + 2 + 1  # inputs, pdl, serialized, before capture (warm-up excluded)
+    checked = {}
+    if "absorbed_query" in selected:
+        checked["absorbed_query_full"] = counts.get(f"absorbed_query/layer{FULL_LAYER}", 0)
+        assert checked["absorbed_query_full"] >= eager_steps, counts
+    if "shared_indexer_absorbed_query" in selected:
+        checked["shared_query_shared"] = counts.get(
+            f"shared_indexer_absorbed_query/layer{SHARED_LAYER}", 0
+        )
+        assert checked["shared_query_shared"] >= eager_steps, counts
+        # The shared layer takes shared_indexer_absorbed_query's packet, not absorbed_query's.
+        assert counts.get(f"absorbed_query/layer{SHARED_LAYER}", 0) == 0, counts
+    if "rope_kv_append" in selected:
+        for layer in (FULL_LAYER, SHARED_LAYER):
+            assert counts.get(f"rope_append_rope/layer{layer}", 0) >= eager_steps, counts
+            assert counts.get(f"rope_append_bridge/layer{layer}", 0) >= eager_steps, counts
+    if selected & set(TAIL_FEATURES):
+        for layer in (FULL_LAYER, SHARED_LAYER):
+            tail = eng["attn_tail_tails"][layer]
+            assert tail is not None and tail["static_reason"] is None, tail
+            assert tail["engaged_calls"] >= eager_steps, tail
+            checked[f"attn_tail_layer{layer}"] = {
+                "engaged_calls": tail["engaged_calls"],
+                "early_rows": tail["config"]["early_rows"],
+                "oproj_grid": tail["config"]["oproj_grid"],
+            }
+    sparse_mla_calls = eng["sparse_mla_decode"]["engaged_calls"]
+    if with_sparse_mla:
+        for layer in (FULL_LAYER, SHARED_LAYER):
+            assert sparse_mla_calls.get(str(layer), 0) >= eager_steps, eng["sparse_mla_decode"]
+    else:
+        assert not sparse_mla_calls, eng["sparse_mla_decode"]
+    checked["sparse_mla_engaged_calls"] = sparse_mla_calls
+    return checked
+
+
+def _replay_check(run: dict, layers) -> dict:
+    """Compare every plain / serialized / graph step with the recorded eager step of its input."""
+    reference = {e["input"]: e["layers"] for e in run["inputs"]}
+    mismatches, index_changes, compared = [], [], 0
+    for kind in ("pdl", "serialized", "graph"):
+        for i, entry in enumerate(run[kind]):
+            ref = reference[entry["input"]]
+            for layer in layers:
+                got = entry["layers"][layer]
+                if not _same(got["indices"], ref[layer]["indices"]):
+                    index_changes.append((kind, i, layer))
+                for key in ("latent", "module_output"):
+                    compared += 1
+                    if not _same(got[key], ref[layer][key]):
+                        mismatches.append((kind, i, layer, key))
+    return {"compared": compared, "mismatches": mismatches, "index_changes": index_changes}
+
+
+@pytest.mark.parametrize("neighbours,topk_mode", STACK_SETS)
+def test_sparse_mla_stack_with_engaged_neighbours(neighbours, topk_mode, tmp_path):
+    out_dir = Path(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS", tmp_path))
+    selected = {f for f in neighbours.split(",") if f}
+    with_sparse_mla = ",".join(sorted(selected) + ["sparse_mla_decode"])
+    base = _run(neighbours, topk_mode, out_dir)
+    cand = _run(with_sparse_mla, topk_mode, out_dir)
+    layers = cand["layers"]
+    assert layers == [FULL_LAYER, SHARED_LAYER] and cand["shared_layers"] == [SHARED_LAYER]
+    assert base["libraries"] == ["FallbackFmha"], base["libraries"]
+    assert cand["libraries"] == ["SparseMlaDecodeFmha"], cand["libraries"]
+
+    summary = {
+        "neighbours": neighbours or "native",
+        "topk_mode": topk_mode,
+        "ctx": cand["ctx"],
+        "trtllm_enable_pdl": os.environ.get("TRTLLM_ENABLE_PDL", "1"),
+        "engagement_without_sparse_mla": _check_engagement(base, selected, False),
+        "engagement_with_sparse_mla": _check_engagement(cand, selected, True),
+        "capability_log_with_sparse_mla": cand["capability_log"],
+        "steps": [],
+    }
+    for b, c in zip(base["inputs"], cand["inputs"]):
+        for layer in layers:
+            bl, cl = b["layers"][layer], c["layers"][layer]
+            # The producers do not depend on the attention library.
+            for key in ("quant_q", "bmm1", "bmm2", "kv_lens"):
+                assert _same(bl[key], cl[key]), (neighbours, layer, key)
+            same_order = _same(bl["indices"], cl["indices"])
+            same_sets = torch.equal(
+                bl["indices"].sort(dim=-1).values, cl["indices"].sort(dim=-1).values
+            )
+            if topk_mode == "fixed":
+                assert same_order and _same(bl["rows"], cl["rows"]), (neighbours, layer)
+            assert not bool((cl["indices"] < 0).any()), "a -1 index below the causal length"
+            native_err = error_stats(bl["latent"].cuda().double(), _oracle(bl))
+            sparse_mla_err = error_stats(cl["latent"].cuda().double(), _oracle(cl))
+            module_diff = error_stats(cl["module_output"].double(), bl["module_output"].double())
+            summary["steps"].append(
+                {
+                    "input": c["input"],
+                    "layer": layer,
+                    "indices_same_order": same_order,
+                    "indices_same_sets": same_sets,
+                    "native_latent_vs_oracle": native_err,
+                    "sparse_mla_latent_vs_oracle": sparse_mla_err,
+                    "module_output_sparse_mla_vs_native": module_diff,
+                }
+            )
+            assert sparse_mla_err["nonfinite"] == native_err["nonfinite"] == 0
+            assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
+                layer,
+                sparse_mla_err,
+                native_err,
+            )
+            assert sparse_mla_err["max_abs"] <= 2.0 * native_err["max_abs"] + 1e-3, (
+                layer,
+                sparse_mla_err,
+            )
+            if same_sets:
+                assert module_diff["nonfinite"] == 0 and module_diff["nrms"] < 0.05, module_diff
+
+    summary["replays_without_sparse_mla"] = _replay_check(base, layers)
+    summary["replays_with_sparse_mla"] = _replay_check(cand, layers)
+    summary["kernels_without_sparse_mla"] = {
+        k: {f: v for f, v in s.items() if f != "attention"} for k, s in base["kernels"].items()
+    }
+    summary["kernels_with_sparse_mla"] = {
+        k: {f: v for f, v in s.items() if f != "attention"} for k, s in cand["kernels"].items()
+    }
+    summary["rope_append_bridge_without_sparse_mla"] = base["rope_append_bridge"]
+    summary["rope_append_bridge_with_sparse_mla"] = cand["rope_append_bridge"]
+    record(
+        "sparse_mla_stack_composition.jsonl", "stack_composition", json.loads(json.dumps(summary))
+    )
+
+    # sparse_mla_decode is deterministic: with it, every plain, serialized and replayed step
+    # reproduces the recorded eager step bit for bit when it attended the same indices (always with
+    # a fixed top-k; the native indexer may reorder or change its selection from step to step).
+    replays = summary["replays_with_sparse_mla"]
+    if topk_mode == "fixed":
+        assert not replays["mismatches"], replays
+    else:
+        changed = {(k, i, layer) for k, i, layer in replays["index_changes"]}
+        stable = [m for m in replays["mismatches"] if tuple(m[:3]) not in changed]
+        assert not stable, replays
+
+    # One attention kernel per DSA layer per step; the right library in each run.
+    for run, sparse_mla_expected in ((base, False), (cand, True)):
+        for kind, stats in run["kernels"].items():
+            calls = stats["steps"] * stats["layers"]
+            assert stats["sparse_mla_decode"] == (calls if sparse_mla_expected else 0), (
+                kind,
+                stats["sparse_mla_decode"],
+                calls,
+            )
+            assert stats["native_sparse_fmha"] == (0 if sparse_mla_expected else calls), (
+                kind,
+                stats["native_sparse_fmha"],
+                calls,
+            )
+            if "rope_kv_append" in selected:
+                assert stats["rope_append_bridge"] >= stats["steps"], (
+                    kind,
+                    stats["rope_append_bridge"],
+                )
+                if topk_mode == "fixed":
+                    assert stats["native_remap"] == 0, (kind, stats["native_remap"])
+    # The attention_output_tail is sparse_mla_decode's programmatic dependent in the captured graph
+    # and becomes resident well before sparse_mla_decode ends (sparse_mla_decode triggers before its
+    # gathers complete). Without PDL the tail starts only after sparse_mla_decode completes (CUPTI
+    # may still show a sub-microsecond overlap of the end and start stamps of back-to-back kernels,
+    # so the lead is compared, not its sign).
+    if selected & set(TAIL_FEATURES):
+        graph = cand["kernels"]["graph"]
+        sparse_mla_records = [a for a in graph["attention"] if a["kernel"] == "sparse_mla_decode"]
+        assert all("_wuv_kernel" in (a["successor"] or "") for a in sparse_mla_records), graph[
+            "successors"
+        ]
+        leads = graph["sparse_mla_successor_lead_us"]
+        median_lead = leads[len(leads) // 2]
+        if os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1":
+            assert graph["sparse_mla_successor_early"] == len(sparse_mla_records), graph[
+                "sparse_mla_successor_early"
+            ]
+            assert median_lead >= 1.0, leads
+        else:
+            assert median_lead < 0.5, leads
+
+    # rope_kv_append's bridge as the producer of the indices: equal to the native remap, no -1 below
+    # a row's causal length, and what the attention received.
+    if "rope_kv_append" in selected:
+        for run in (base, cand):
+            reports = run["rope_append_bridge"]
+            assert len(reports) == 2 * len(layers), reports
+            for report in reports:
+                assert report["equals_native_remap"], report
+                assert sum(report["negative_below_causal"]) == 0, report
+                assert report["attention_received_bridge_output"], report
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py
new file mode 100644
index 000000000000..5bc49b7e05b6
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py
@@ -0,0 +1,157 @@
+# 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.
+"""Contract table of the target's kernels (``kernels/contracts.py``).
+
+The module is loaded by file path so the test needs neither torch nor the
+compiled extensions; kernel sources are read as text.
+"""
+
+import importlib.util
+import re
+from pathlib import Path
+
+import pytest
+
+_ROOT = Path(__file__).resolve().parents[8]
+_TARGET = ("tensorrt_llm", "_torch", "_experimental", "modeling_v2", "models", "glm_moe_dsa")
+_KERNELS = (*_TARGET, "shared", "kernels")
+_TESTS = ("tests", "unittest", "_torch", "modeling_v2", "models", "glm_moe_dsa")
+_KERNEL_TESTS = (*_TESTS, "shared", "kernels")
+
+
+def _load(name):
+    path = _ROOT.joinpath(*_KERNELS, f"{name}.py")
+    spec = importlib.util.spec_from_file_location(f"glm_moe_dsa_kernels_{name}", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_contract_table():
+    contracts = _load("contracts")
+    owners = {k: c.owner for k, c in contracts.CONTRACTS.items()}
+    assert owners == {
+        "query_front": "absorbed_query",
+        "index_selection": "indexer_topk",
+        "sparse_attention": "sparse_mla_decode",
+        "attention_output_tail": "attention_output_tail",
+        "router": "router_gemm",
+        "routed_moe": "routed_moe",
+        "post_moe_hop": "post_moe_hop",
+        "draft_front": "draft_front",
+    }
+    # Rules are optional, so a four-field contract is still valid.
+    plain = contracts.Contract("x", "absorbed_query", ("a",), ("b",))
+    assert plain.rules == ()
+
+
+def test_attention_tail_rules():
+    contracts = _load("contracts")
+    tail = contracts.CONTRACTS["attention_output_tail"]
+    assert tail.outputs == ("partial[R,6144]",)
+    assert tail.rules == (
+        contracts.ATTENTION_TAIL_COMMUNICATION_RULE,
+        contracts.ATTENTION_TAIL_PDL_RULE,
+    )
+    comm = contracts.ATTENTION_TAIL_COMMUNICATION_RULE
+    assert "No communication is hidden inside attention_output_tail." in comm
+    assert "router_gemm/native all-reduce consumes the identical partial." in comm
+    pdl = contracts.ATTENTION_TAIL_PDL_RULE
+    for phrase in (
+        "attention_output_tail may be a PDL dependent of the latent's producer.",
+        "Before its griddepcontrol.wait a attention_output_tail kernel reads only immutable weights",
+        "writes only its own shared memory",
+        "no async-proxy global writes",
+        "Every latent read and every global write (v scratch, partial) follows the wait.",
+        "executes griddepcontrol.launch_dependents before it writes the partial",
+        "must execute griddepcontrol.wait before reading the partial",
+        "may be reused after attention_output_tail returns",
+    ):
+        assert phrase in pdl, phrase
+    switches = re.search(r"Kill switches: (.*)\.$", pdl).group(1)
+    assert switches == (
+        "TRTLLM_ENABLE_PDL=0 (all, entry included), TailConfig(trigger_dependents=False) (exit)"
+    )
+
+
+def _source(*parts):
+    path = _ROOT.joinpath(*parts)
+    if not path.is_file():
+        pytest.skip(f"{path} is not in this tree")
+    return path.read_text()
+
+
+def test_attention_tail_names_exist_in_the_tree():
+    """The kill switches and the into-hop protocol header the rules name exist."""
+    kernels = _source(*_KERNELS, "attention_output_tail", "kernels.py")
+    assert re.search(
+        r"def pdl_entry_enabled\(\) -> bool:\n(?:.*\n)*?    return pdl_enabled\(\)$", kernels, re.M
+    )
+    assert 'os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1"' in kernels
+    assert re.search(r"^\s+trigger_dependents: bool = True$", kernels, re.M)
+    rule = _load("contracts").ATTENTION_TAIL_COMMUNICATION_RULE
+    header = re.search(r"(cpp/\S+\.h)", rule).group(1)
+    assert "griddepcontrol.wait" in _source(*header.split("/"))
+
+
+def test_routed_moe_and_post_moe_hop_rules():
+    """The routed-input and deferred-norm seams and the token-major routed-MoE form."""
+    contracts = _load("contracts")
+    c = contracts.CONTRACTS
+    assert c["routed_moe"].rules == (
+        contracts.ROUTED_MOE_TOKEN_MAJOR_RULE,
+        contracts.POST_MOE_HOP_ROUTED_INPUT_RULE,
+    )
+    assert c["post_moe_hop"].rules == (
+        contracts.POST_MOE_HOP_ROUTED_INPUT_RULE,
+        contracts.POST_MOE_HOP_DEFERRED_NORM_RULE,
+        contracts.POST_MOE_HOP_SHARED_PROTOCOL_RULE,
+    )
+    assert c["query_front"].rules == (contracts.POST_MOE_HOP_DEFERRED_NORM_RULE,)
+    for phrase in ("identity t*8+k", "reads fc2_output only through the permutation"):
+        assert phrase in contracts.ROUTED_MOE_TOKEN_MAJOR_RULE, phrase
+    for phrase in (
+        "parsed only by dispatch.feature_variants",
+        "agreed by every rank",
+        "zero when routed_moe's kernel starts",
+        "the hop clears after reading them",
+        "routed_moe never reuses a buffer it wrote into",
+        "only after its own griddepcontrol.wait",
+    ):
+        assert phrase in contracts.POST_MOE_HOP_ROUTED_INPUT_RULE, phrase
+    for phrase in (
+        "The packet is not a tensor",
+        "materialize()",
+        "reads only immutable weights before its griddepcontrol.wait",
+    ):
+        assert phrase in contracts.POST_MOE_HOP_DEFERRED_NORM_RULE, phrase
+    assert "must change both in the same commit" in contracts.POST_MOE_HOP_SHARED_PROTOCOL_RULE
+
+
+def test_routed_moe_and_post_moe_hop_names_exist_in_the_tree():
+    """The names the routed-input and deferred-norm rules cite exist where the rules say."""
+    hop = _source(*_KERNELS, "post_moe_hop", "hop.py")
+    assert "def enable_routed_input(" in hop
+    assert 'ROUTED_MODES = ("row", "identity", "fp32")' in hop
+    assert 'dispatch.feature_variants(("routed_moe_into_hop",))["routed_moe_into_hop"]' in hop
+    packet = _source(*_KERNELS, "deferred_norm", "packet.py")
+    assert "class DeferredNormHidden:" in packet and "def materialize(self)" in packet
+    ptx = _source(*_KERNELS, "deferred_norm", "ptx.py")
+    assert "from ..post_moe_hop.kernels import (" in ptx
+    hop_kernel = _source(*_KERNELS, "deferred_norm", "hop_kernel.py")
+    assert "griddepcontrol_launch_dependents" in hop_kernel
+    assert "trigger_early" in hop_kernel
+    for name in ("test_deferred_norm_hop_mpi.py", "test_deferred_norm_routed_mpi.py"):
+        assert "mixed_with_post_moe_hop" in _source(*_KERNEL_TESTS, "deferred_norm", name)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py
new file mode 100644
index 000000000000..d60ed606d9c5
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py
@@ -0,0 +1,131 @@
+# 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.
+"""Names of the GLM-MoE-DSA torch ops and the layout of their C++ sources.
+
+The model-specific kernels live in ``cpp/tensorrt_llm/kernels/glmMoeDsa/`` and
+their torch ops in ``cpp/tensorrt_llm/thop/glmMoeDsa/``, one ``<stem>Op.cpp``
+per kernel ``<stem>.cu``; the kernels build as one OBJECT library for the
+SM100 family. Every op they register is
+``trtllm::glm_moe_dsa_<what>``, and every such name the Python sources use is
+registered there or by a ``torch.library.custom_op`` of the Python package
+under a name the C++ sources leave free. The sources are read as text, so
+the test needs neither torch nor the compiled extensions.
+"""
+
+import re
+from pathlib import Path
+
+_THIS = Path(__file__).resolve()
+_ROOT = _THIS.parents[8]
+_CPP = _ROOT / "cpp" / "tensorrt_llm"
+_KERNELS = _CPP / "kernels" / "glmMoeDsa"
+_THOP = _CPP / "thop" / "glmMoeDsa"
+_PYTHON = _ROOT / "tensorrt_llm"
+
+OP_PREFIX = "glm_moe_dsa_"
+
+_REGISTRATION = re.compile(r"\bm\.(def|impl)\(\s*(\"([A-Za-z0-9_.]+)|[A-Z][A-Z0-9_]*)")
+_DEFINE = re.compile(r"^\s*#\s*define\s+([A-Z][A-Z0-9_]*)\s+\"([A-Za-z0-9_.]+)\"", re.M)
+_PYTHON_OP = re.compile(
+    r"\btorch\.library\.custom_op\(\s*[\"']trtllm::(%s[a-z0-9_]+)[\"']" % OP_PREFIX
+)
+
+
+def _registered(path):
+    """(kind, op name) of every m.def / m.impl in ``path``.
+
+    A name given by a macro is resolved through the file's ``#define``s (every
+    definition counts); a macro without a string definition yields ``None``.
+    """
+    text = path.read_text(errors="replace")
+    macros = {}
+    for name, value in _DEFINE.findall(text):
+        macros.setdefault(name, []).append(value)
+    found = []
+    for match in _REGISTRATION.finditer(text):
+        kind = match.group(1)
+        if match.group(3) is not None:
+            found.append((kind, match.group(3)))
+        else:
+            found += [(kind, value) for value in macros.get(match.group(2), [None])]
+    return found
+
+
+def _thop_ops():
+    ops = {}
+    for path in sorted(_THOP.glob("*.cpp")):
+        for kind, name in _registered(path):
+            assert name is not None, f"{path.name}: m.{kind} names an op through an unknown macro"
+            ops.setdefault(name, set()).add(kind)
+    return ops
+
+
+def _python_ops():
+    """Every ``trtllm::glm_moe_dsa_*`` op a ``torch.library.custom_op`` of the package defines."""
+    ops = set()
+    for path in sorted(_PYTHON.rglob("*.py")):
+        ops.update(_PYTHON_OP.findall(path.read_text(errors="replace")))
+    return ops
+
+
+def test_sources_sit_in_the_glm_moe_dsa_folders():
+    kernels = sorted(p.stem for p in _KERNELS.glob("*.cu"))
+    headers = sorted(p.stem for p in _KERNELS.glob("*.h"))
+    ops = sorted(p.stem[: -len("Op")] for p in _THOP.glob("*Op.cpp"))
+    assert kernels, "no kernel sources under kernels/glmMoeDsa"
+    assert kernels == headers
+    assert kernels == ops
+    assert sorted(p.name for p in _THOP.glob("*.cpp")) == sorted(f"{s}Op.cpp" for s in ops)
+    cmake = (_CPP / "thop" / "CMakeLists.txt").read_text()
+    assert "add_th_common_family(th_common_glm_moe_dsa)" in cmake
+    assert "add_subdirectory(glmMoeDsa)" in cmake
+    assert "$<TARGET_OBJECTS:th_common_glm_moe_dsa>" in cmake
+    assert "th_common_glm_moe_dsa" in (_THOP / "CMakeLists.txt").read_text()
+
+
+def test_kernels_build_as_one_sm100_family_object_library():
+    cmake = (_KERNELS / "CMakeLists.txt").read_text()
+    listed = re.search(r"add_library\(\s*glm_moe_dsa_kernels_src\s+OBJECT\s+([^)]*)\)", cmake)
+    assert listed is not None
+    assert sorted(listed.group(1).split()) == sorted(p.name for p in _KERNELS.glob("*.cu"))
+    assert "set_cuda_architectures(glm_moe_dsa_kernels_src 100f)" in cmake
+    # The parent's globs leave the folder to its own target, which joins the
+    # shared library's link line.
+    parent = (_CPP / "kernels" / "CMakeLists.txt").read_text()
+    assert "add_subdirectory(glmMoeDsa)" in parent
+    assert 'list(FILTER SRC_CPP EXCLUDE REGEX "glmMoeDsa/.*")' in parent
+    assert 'list(FILTER SRC_CU EXCLUDE REGEX "glmMoeDsa/.*")' in parent
+    assert re.search(r"^\s*glm_moe_dsa_kernels_src$", (_CPP / "CMakeLists.txt").read_text(), re.M)
+
+
+def test_every_op_defined_in_thop_glm_moe_dsa_has_the_model_prefix():
+    ops = _thop_ops()
+    assert ops, "no op registrations under thop/glmMoeDsa"
+    assert sorted(n for n in ops if not n.startswith(OP_PREFIX)) == []
+    # Every implemented op is also defined (the schema lives next to the kernel).
+    assert sorted(n for n, kinds in ops.items() if "def" not in kinds) == []
+
+
+def test_python_sources_use_only_registered_glm_moe_dsa_ops():
+    used = re.compile(r"\b%s[a-z0-9_]+\b" % OP_PREFIX)
+    thop, python = set(_thop_ops()), _python_ops()
+    assert sorted(thop & python) == []
+    registered = thop | python
+    unknown = []
+    for path in sorted(_PYTHON.rglob("*.py")):
+        text = path.read_text(errors="replace")
+        rel = str(path.relative_to(_ROOT))
+        unknown += [(rel, m.group(0)) for m in used.finditer(text) if m.group(0) not in registered]
+    assert unknown == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py
new file mode 100644
index 000000000000..e67313ed9dcf
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py
@@ -0,0 +1,754 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The GLM-5.2 tp8 target's attention modules: their adoption, the FMHA libraries the attention
+backends declare, the phase gate of every override, the indexer's top-k region and bridge, and
+the target's DSA op 2.
+
+Host-only: the built-in modules are created without their constructors and carry only the
+attributes adoption and the overrides read. The kernel factories are stand-ins, and the built-in
+methods and the decode functions an override calls are recorders.
+"""
+
+from __future__ import annotations
+
+import ast
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    attention,
+    attention_forward,
+    dispatch,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
+    region,
+)
+from tensorrt_llm._torch.attention.backends.fmha import manager as fmha_manager
+from tensorrt_llm._torch.attention.backends.interface import AttentionInputType
+from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
+from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
+from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention
+from tensorrt_llm._torch.attention.mla import MLA
+from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
+
+pytestmark = pytest.mark.cpu_only
+
+PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
+
+
+# --- built-in modules without their constructors ------------------------------------------
+
+
+def _new(cls, **attributes):
+    obj = cls.__new__(cls)
+    if isinstance(obj, nn.Module):
+        nn.Module.__init__(obj)
+    for name, value in attributes.items():
+        setattr(obj, name, value)
+    return obj
+
+
+def _attention(layer_idx=3, full_indexer=True, dense_mha=False):
+    backend = _new(DSATrtllmAttention, indexer=_new(Indexer) if full_indexer else None)
+    return _new(
+        DeepseekV32Attention,
+        layer_idx=layer_idx,
+        sparse_attn_hooks=_new(DSAHooks, short_seq_mha_threshold=0, need_dense_mha=dense_mha),
+        mqa=backend,
+        mha=_new(TrtllmAttention) if dense_mha else None,
+        indexer=backend.indexer,
+    )
+
+
+class _SparseMlaDecodeFmha:
+    """Stands in for the sparse-MLA decode FMHA library."""
+
+
+class _Tail:
+    def __init__(self, mla) -> None:
+        self.mla = mla
+
+
+class _FastPaths:
+    """``dispatch.is_enabled`` / ``get_impl`` for the adoption tests."""
+
+    def __init__(self, enabled):
+        self.enabled = frozenset(enabled)
+        self.events = []
+        events = self.events
+
+        def make_attn_tail(mla):
+            events.append(("make_attn_tail", mla.layer_idx))
+            return _Tail(mla)
+
+        def wrap_attn_tail(mla, tail):
+            events.append(("wrap_attn_tail", tail))
+            return ("staged", tail)
+
+        self.entries = {
+            ("attention_output_tail", "make_attn_tail"): make_attn_tail,
+            ("oproj_weight_staging", "wrap_attn_tail"): wrap_attn_tail,
+            ("sparse_mla_decode", "fmha_lib"): _SparseMlaDecodeFmha,
+        }
+
+    def is_enabled(self, path):
+        assert path in dispatch.FAST_PATH_NAMES, path
+        return path in self.enabled
+
+    def get_impl(self, path, name):
+        return self.entries[(path, name)]
+
+
+@pytest.fixture
+def fast_paths(monkeypatch):
+    def install(enabled=dispatch.FAST_PATH_NAMES):
+        paths = _FastPaths(enabled)
+        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
+        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
+        return paths
+
+    return install
+
+
+def _keys(module):
+    return (
+        [name for name, _ in module.named_parameters()],
+        [name for name, _ in module.named_buffers()],
+        [name for name, _ in module.named_modules()],
+    )
+
+
+# --- adoption -----------------------------------------------------------------------------
+
+
+def test_adoption_reclasses_the_attention_modules_and_builds_the_tail(fast_paths):
+    paths = fast_paths()
+    mla = _attention()
+    hooks, backend, indexer = mla.sparse_attn_hooks, mla.mqa, mla.mqa.indexer
+    before = _keys(mla)
+
+    assert attention.GlmMLA.adopt(mla) is mla
+
+    assert type(mla) is attention.GlmMLA
+    assert type(hooks) is attention.GlmDsaHooks and mla.sparse_attn_hooks is hooks
+    assert type(backend) is attention.GlmDsaBackend and mla.mqa is backend
+    assert type(indexer) is attention.GlmIndexer and mla.indexer is indexer
+    assert hooks.short_seq_mha_threshold == 0 and hooks.need_dense_mha is False
+    # The tail is made for the layer, then wrapped by the o_proj weight staging.
+    staged, tail = mla._attn_tail
+    assert staged == "staged" and isinstance(tail, _Tail) and tail.mla is mla
+    assert paths.events == [("make_attn_tail", 3), ("wrap_attn_tail", tail)]
+    assert mla._tail_prefetch is True
+    assert _keys(mla) == before
+
+
+@pytest.mark.parametrize(
+    "path", ["attention_output_tail", "attention_output_tail_prefetch", "attention_output_into_hop"]
+)
+def test_any_tail_path_builds_the_tail(fast_paths, path):
+    paths = fast_paths({path})
+    mla = attention.GlmMLA.adopt(_attention())
+    assert isinstance(mla._attn_tail, _Tail) and paths.events == [("make_attn_tail", 3)]
+    assert mla._tail_prefetch is (path == "attention_output_tail_prefetch")
+
+
+def test_without_the_attention_paths_adoption_attaches_nothing(fast_paths):
+    paths = fast_paths(enabled=())
+    mla = attention.GlmMLA.adopt(_attention(full_indexer=False))
+    assert paths.events == []
+    assert mla._attn_tail is None and mla._tail_prefetch is False
+    assert mla.mqa.prepended_fmha_lib_classes == ()
+    assert "prepended_fmha_lib_classes" not in vars(mla.mqa)
+    assert mla.mqa.indexer is None and type(mla.mqa) is attention.GlmDsaBackend
+
+
+@pytest.mark.parametrize("dense_mha", [False, True], ids=["mqa only", "mha and mqa"])
+def test_the_sparse_mla_decode_library_is_declared_on_each_backend(fast_paths, dense_mha):
+    fast_paths({"sparse_mla_decode"})
+    mla = attention.GlmMLA.adopt(_attention(dense_mha=dense_mha))
+    backends = [mla.mqa] + ([mla.mha] if dense_mha else [])
+    for backend in backends:
+        assert backend.prepended_fmha_lib_classes == (_SparseMlaDecodeFmha,)
+    assert TrtllmAttention.prepended_fmha_lib_classes == ()
+
+
+def _other(cls):
+    return type(f"_Other{cls.__name__}", (cls,), {})
+
+
+def _refusal(case):
+    mla = _attention()
+    if case == "an attention subclass":
+        mla.__class__ = _other(DeepseekV32Attention)
+    elif case == "the MLA base class":
+        mla.__class__ = MLA
+    elif case == "a DSA hooks subclass":
+        mla.sparse_attn_hooks.__class__ = _other(DSAHooks)
+    elif case == "a DSA backend subclass":
+        mla.mqa.__class__ = _other(DSATrtllmAttention)
+    elif case == "an indexer subclass":
+        mla.mqa.indexer.__class__ = _other(Indexer)
+    elif case == "an adopted attention":
+        mla = attention.GlmMLA.adopt(mla)
+    return mla
+
+
+@pytest.mark.parametrize(
+    "case, got",
+    [
+        ("an attention subclass", "_OtherDeepseekV32Attention"),
+        ("the MLA base class", "MLA"),
+        ("a DSA hooks subclass", "_OtherDSAHooks"),
+        ("a DSA backend subclass", "_OtherDSATrtllmAttention"),
+        ("an indexer subclass", "_OtherIndexer"),
+        ("an adopted attention", "GlmMLA"),
+    ],
+)
+def test_adoption_refuses_anything_but_the_exact_builtin_class(fast_paths, case, got):
+    fast_paths(enabled=())
+    mla = _refusal(case)
+    with pytest.raises(TypeError, match=f"adopts the built-in .*, got {got}$"):
+        attention.GlmMLA.adopt(mla)
+
+
+def test_the_classes_declare_no_state_of_their_own():
+    for cls in (
+        attention.GlmMLA,
+        attention.GlmDsaHooks,
+        attention.GlmDsaBackend,
+        attention.GlmIndexer,
+    ):
+        assert "__slots__" not in vars(cls) and "__init__" not in vars(cls), cls
+    assert attention.GlmIndexer._topk_bridge is None
+    assert attention.GlmMLA._attn_tail is None and attention.GlmMLA._tail_prefetch is False
+
+
+# --- the FMHA libraries -------------------------------------------------------------------
+
+
+def _library(name, calls):
+    def is_available(cls, attn):
+        calls.append((name, attn.name))
+        return True
+
+    return type(
+        name,
+        (),
+        {"is_available": classmethod(is_available), "__init__": lambda self, attn: None},
+    )
+
+
+def _managers(monkeypatch, registry, declared):
+    """The libraries an mha and an mqa backend build, in the order ``create_weights`` builds
+    their managers, and the order their availability is asked in."""
+    calls = []
+    libraries = {name: _library(name, calls) for name in ("sparse", "a", "b", "fallback")}
+    monkeypatch.setattr(
+        fmha_manager, "get_enabled_fmha_lib_classes", lambda: [libraries[n] for n in registry]
+    )
+    built = []
+    for name in ("mha", "mqa"):
+        attn = types.SimpleNamespace(name=name)
+        if declared:
+            attn.prepended_fmha_lib_classes = tuple(libraries[n] for n in declared)
+        manager = fmha_manager.FmhaManager(attn)
+        built.append([type(lib).__name__ for lib in manager.fmha_libs])
+    return built, calls
+
+
+def test_the_declared_library_reproduces_a_registry_with_it_first(monkeypatch):
+    """A backend that declares the library builds the list and asks the libraries in the order a
+    registry with the library first would; one that declares nothing builds the registry's."""
+    first = _managers(monkeypatch, ["sparse", "a", "b", "fallback"], declared=())
+    declared = _managers(monkeypatch, ["a", "b", "fallback"], declared=("sparse",))
+    assert declared == first
+    assert declared[0] == [["sparse", "a", "b", "fallback"]] * 2
+    plain = _managers(monkeypatch, ["a", "b", "fallback"], declared=())
+    assert plain[0] == [["a", "b", "fallback"]] * 2
+
+
+# --- the phase gate -----------------------------------------------------------------------
+
+# Each override and the built-in method it falls back to.
+OVERRIDES = [
+    (attention.GlmMLA, DeepseekV32Attention, "forward"),
+    (attention.GlmMLA, DeepseekV32Attention, "forward_absorption_generation"),
+    (attention.GlmDsaHooks, DSAHooks, "forward"),
+    (attention.GlmDsaHooks, DSAHooks, "forward_custom_op"),
+    (attention.GlmDsaBackend, DSATrtllmAttention, "sparse_attn_predict"),
+    (attention.GlmIndexer, Indexer, "_fill_decode_topk"),
+]
+
+
+def _method(cls, name):
+    tree = ast.parse(Path(attention.__file__).read_text())
+    return next(
+        n
+        for c in tree.body
+        if isinstance(c, ast.ClassDef) and c.name == cls.__name__
+        for n in c.body
+        if isinstance(n, ast.FunctionDef) and n.name == name
+    )
+
+
+@pytest.mark.parametrize("cls, builtin, name", OVERRIDES, ids=lambda x: getattr(x, "__name__", x))
+def test_each_override_starts_with_the_phase_gate(cls, builtin, name):
+    """``if dispatch.phase() is not DECODE: return super().<name>(<every built-in parameter by
+    name>)``; the parameters only the decode composition takes are dropped."""
+    node = _method(cls, name)
+    docstring = isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant)
+    gate = node.body[1] if docstring else node.body[0]
+    assert ast.unparse(gate.test) == "dispatch.phase() is not dispatch.Phase.DECODE"
+    assert len(gate.body) == 1 and not gate.orelse
+    call = gate.body[0].value
+    assert ast.unparse(call.func) == f"super().{name}" and not call.args
+    parameters = inspect.signature(getattr(builtin, name)).parameters
+    expected = [
+        p.name if p.kind is not p.VAR_KEYWORD else None
+        for p in parameters.values()
+        if p.name != "self"
+    ]
+    assert [k.arg for k in call.keywords] == expected
+    # The override keeps the built-in parameters, in order, and may add keywords after them.
+    named = [p for p in parameters.values() if p.kind is not p.VAR_KEYWORD]
+    own = inspect.signature(getattr(cls, name)).parameters.values()
+    assert [p.name for p in own][: len(named)] == [p.name for p in named]
+    assert any(p.kind is p.VAR_KEYWORD for p in own) == (len(named) < len(parameters))
+
+
+def _record(monkeypatch, owner, name, calls, result=None):
+    def recorder(*args, **kwargs):
+        calls.append((name, args, kwargs))
+        return f"{getattr(owner, '__name__', owner)}.{name}" if result is None else result
+
+    monkeypatch.setattr(owner, name, recorder)
+
+
+def _adopted(monkeypatch):
+    monkeypatch.setattr(dispatch, "is_enabled", lambda path: False)
+    return attention.GlmMLA.adopt(_attention())
+
+
+@pytest.mark.parametrize("phase", [PREFILL, None], ids=["prefill", "outside a step"])
+def test_off_decode_every_override_runs_the_builtin_method(monkeypatch, phase):
+    mla = _adopted(monkeypatch)
+    calls = []
+    for owner, name in [
+        (MLA, "forward"),
+        (MLA, "forward_absorption_generation"),
+        (DSAHooks, "forward"),
+        (DSAHooks, "forward_custom_op"),
+        (DSATrtllmAttention, "sparse_attn_predict"),
+    ]:
+        _record(monkeypatch, owner, name, calls)
+    for name in (
+        "forward_absorption_generation",
+        "forward_sparse_attn",
+        "forward_sparse_attn_custom_op",
+    ):
+        monkeypatch.setattr(attention_forward, name, lambda *a, **k: pytest.fail())
+    monkeypatch.setattr(region, "cross_rank_topk_decode", lambda *a, **k: pytest.fail())
+
+    def run():
+        assert mla.forward("pos", "h", "md", "params", "latent") == "MLA.forward"
+        assert (
+            mla.forward_absorption_generation(
+                "q", "kv", "pe", "md", "out", "pos", "lc", "absorbed", "deferred", extra=1
+            )
+            == "MLA.forward_absorption_generation"
+        )
+        assert mla.sparse_attn_hooks.forward(mla, "pos", "h", "md", ["o"]) == "DSAHooks.forward"
+        assert (
+            mla.sparse_attn_hooks.forward_custom_op(mla, "h", "pos", ["o"], None)
+            == "DSAHooks.forward_custom_op"
+        )
+        assert (
+            mla.mqa.sparse_attn_predict("q", None, "md", "args")
+            == "DSATrtllmAttention.sparse_attn_predict"
+        )
+        assert mla.mqa.indexer._fill_decode_topk("md", "q", "w", "buf", 0, 6) is False
+
+    if phase is None:
+        run()
+    else:
+        with dispatch.step(phase):
+            run()
+    assert [(name, kwargs) for name, _, kwargs in calls] == [
+        (
+            "forward",
+            dict(
+                position_ids="pos",
+                hidden_states="h",
+                attn_metadata="md",
+                all_reduce_params="params",
+                latent_cache_gen="latent",
+            ),
+        ),
+        (
+            "forward_absorption_generation",
+            dict(
+                q="q",
+                compressed_kv="kv",
+                k_pe="pe",
+                attn_metadata="md",
+                output="out",
+                position_ids="pos",
+                latent_cache="lc",
+                extra=1,
+            ),
+        ),
+        (
+            "forward",
+            dict(
+                mla=mla,
+                position_ids="pos",
+                hidden_states="h",
+                attn_metadata="md",
+                attn_output=["o"],
+            ),
+        ),
+        (
+            "forward_custom_op",
+            dict(
+                mla=mla,
+                hidden_states="h",
+                position_ids="pos",
+                attn_output=["o"],
+                latent_cache_gen=None,
+            ),
+        ),
+        ("sparse_attn_predict", dict(q="q", k=None, metadata="md", forward_args="args")),
+    ]
+
+
+def test_on_a_decode_step_each_override_runs_the_decode_function(monkeypatch):
+    mla = _adopted(monkeypatch)
+    calls = []
+    for name in (
+        "forward_absorption_generation",
+        "forward_sparse_attn",
+        "forward_sparse_attn_custom_op",
+    ):
+        _record(monkeypatch, attention_forward, name, calls)
+    for owner, name in [
+        (MLA, "forward_absorption_generation"),
+        (DSAHooks, "forward"),
+        (DSAHooks, "forward_custom_op"),
+    ]:
+        monkeypatch.setattr(owner, name, lambda *a, **k: pytest.fail())
+
+    with dispatch.step(DECODE):
+        mla.forward_absorption_generation(
+            "q", "kv", "pe", "md", "out", "pos", "lc", "absorbed", "deferred", extra=1
+        )
+        assert mla.sparse_attn_hooks.forward(mla, "pos", "h", "md", ["o"]) is None
+        assert mla.sparse_attn_hooks.forward_custom_op(mla, "h", "pos", ["o"], None) is True
+
+    assert calls == [
+        (
+            "forward_absorption_generation",
+            (mla, "q", "kv", "pe", "md", "out"),
+            dict(
+                position_ids="pos",
+                latent_cache="lc",
+                absorbed_fused_q="absorbed",
+                deferred_indexer_k="deferred",
+                extra=1,
+            ),
+        ),
+        ("forward_sparse_attn", (mla, "pos", "h", "md", ["o"]), {}),
+        ("forward_sparse_attn_custom_op", (mla, "h", "pos", ["o"], None), {}),
+    ]
+
+
+# --- the indexer's top-k region and bridge ------------------------------------------------
+
+
+def _indexer(gvr_deterministic=False):
+    indexer = attention.GlmIndexer.adopt(_new(Indexer))
+    indexer.top_k = types.SimpleNamespace(gvr_deterministic=gvr_deterministic)
+    return indexer
+
+
+@pytest.mark.parametrize("cross_rank", [True, False])
+@pytest.mark.parametrize("gvr_deterministic", [False, True], ids=["default", "deterministic"])
+def test_the_decode_topk_region_tries_the_cross_rank_topk_first(
+    monkeypatch, cross_rank, gvr_deterministic
+):
+    calls = []
+    buffer = torch.zeros(8, 4, dtype=torch.int32)
+
+    def cross_rank_topk_decode(indexer, md, q, w, offset, rows, out, bridge):
+        calls.append(("cross_rank_topk", offset, rows, out.data_ptr(), bridge))
+        return cross_rank
+
+    def indexer_topk_decode(indexer, md, q, w, offset, rows, out):
+        calls.append(("indexer_topk", offset, rows, out.data_ptr()))
+        return True
+
+    monkeypatch.setattr(region, "cross_rank_topk_decode", cross_rank_topk_decode)
+    monkeypatch.setattr(region, "indexer_topk_decode", indexer_topk_decode)
+    indexer = _indexer(gvr_deterministic)
+    indexer._topk_bridge = "bridge"
+
+    with dispatch.step(DECODE):
+        assert indexer._fill_decode_topk("md", "q", "w", buffer, 2, 6) is True
+
+    rows = buffer[2:8].data_ptr()
+    expected = [] if gvr_deterministic else [("cross_rank_topk", 2, 6, rows, "bridge")]
+    if gvr_deterministic or not cross_rank:
+        expected.append(("indexer_topk", 2, 6, rows))
+    assert calls == expected
+
+
+def _rope_kv_append(monkeypatch, served, calls):
+    def entry(name):
+        def run(*args):
+            calls.append((name,) + args)
+            return served
+
+        return run
+
+    monkeypatch.setattr(dispatch, "get_impl", lambda path, name: entry((path, name)))
+
+
+@pytest.mark.parametrize(
+    "bridge, has_prefill, generations, served",
+    [
+        ("bridge", False, 3, True),
+        ("bridge", False, 3, False),
+        ("bridge", True, 3, True),
+        ("bridge", False, 1, True),
+        (None, False, 3, True),
+    ],
+)
+def test_the_reused_topk_copy_defers_to_the_bridge(
+    monkeypatch, bridge, has_prefill, generations, served
+):
+    """Three generation rows; the bridge is asked only for a pure decode call with one row per
+    request, and the built-in copy runs unless it takes the copy."""
+    calls = []
+    _rope_kv_append(monkeypatch, served, calls)
+    indexer = _indexer()
+    indexer._topk_bridge = bridge
+    shared = torch.arange(12, dtype=torch.int32).view(3, 4)
+    metadata = types.SimpleNamespace(shared_topk_indices=shared)
+    buffer = torch.full((4, 4), -1, dtype=torch.int32)
+
+    indexer._copy_reused_topk(metadata, buffer, 1, 3, generations, has_prefill, None)
+
+    asked = bridge is not None and not has_prefill and generations == 3
+    assert [c[0] for c in calls] == ([("rope_kv_append", "defer_reuse_copy")] if asked else [])
+    if asked:
+        _, bridge_arg, buffer_arg, offset, rows, has_prior = calls[0]
+        assert (bridge_arg, buffer_arg, offset, has_prior) == ("bridge", buffer, 1, False)
+        assert torch.equal(rows, shared)
+    if asked and served:
+        assert (buffer == -1).all()
+    else:
+        assert torch.equal(buffer[1:], shared[:generations].expand(3, 4))
+        assert (buffer[0] == -1).all()
+
+
+@pytest.mark.parametrize("served", [True, False])
+def test_the_mtp_row_stash_defers_to_the_bridge(monkeypatch, served):
+    calls, stock = [], []
+    _rope_kv_append(monkeypatch, served, calls)
+    monkeypatch.setattr(Indexer, "_store_mtp_shared_rows", lambda self, *a: stock.append(a))
+    indexer = _indexer()
+    indexer._topk_bridge = "bridge"
+    accepted = torch.tensor([9, 1, 2, 3])
+    metadata = types.SimpleNamespace(shared_topk_indices="shared", mtp_num_accepted=accepted)
+    args = ("buffer", metadata, 1, 3, 18, 0, False, True, True)
+
+    indexer._store_mtp_shared_rows(*args)
+
+    (call,) = calls
+    assert call[:6] == (("rope_kv_append", "defer_mtp_row_stash"), "bridge", "shared", 1, 3, 6)
+    assert torch.equal(call[6], accepted[1:4])
+    assert stock == ([] if served else [args])
+
+
+@pytest.mark.parametrize(
+    "bridge, has_decode, has_prefill, is_generation",
+    [
+        (None, True, False, True),
+        ("bridge", False, False, True),
+        ("bridge", True, True, True),
+        ("bridge", True, False, False),
+    ],
+)
+def test_the_mtp_row_stash_without_a_bridge_call_is_the_builtin_one(
+    monkeypatch, bridge, has_decode, has_prefill, is_generation
+):
+    calls, stock = [], []
+    _rope_kv_append(monkeypatch, True, calls)
+    monkeypatch.setattr(Indexer, "_store_mtp_shared_rows", lambda self, *a: stock.append(a))
+    indexer = _indexer()
+    indexer._topk_bridge = bridge
+    args = ("buffer", "md", 1, 3, 18, 0, has_prefill, has_decode, is_generation)
+    indexer._store_mtp_shared_rows(*args)
+    assert calls == [] and stock == [args]
+
+
+def test_the_builtin_seams_are_the_builtin_statements():
+    """The stock bodies of the three ``sparse_attn_indexer`` seams: the reuse copy, no decode
+    top-k of their own, and the MTP row stash."""
+    indexer = _new(Indexer)
+    shared = torch.arange(12, dtype=torch.int32).view(3, 4)
+    metadata = types.SimpleNamespace(shared_topk_indices=shared, mtp_num_accepted=None)
+    buffer = torch.full((4, 4), -1, dtype=torch.int32)
+    indexer._copy_reused_topk(metadata, buffer, 1, 3, 3, False, None)
+    assert torch.equal(buffer[1:], shared) and (buffer[0] == -1).all()
+    assert indexer._fill_decode_topk(metadata, None, None, buffer, 0, 4) is False
+
+    # Two generation requests at next_n = 2: the last row of each is stashed after row 1.
+    rows = torch.arange(16, dtype=torch.int32).view(4, 4)
+    metadata.shared_topk_indices = torch.zeros(3, 4, dtype=torch.int32)
+    indexer._store_mtp_shared_rows(rows, metadata, 1, 2, 4, 0, False, True, True)
+    assert torch.equal(metadata.shared_topk_indices[1:], rows[[1, 3]])
+    metadata.shared_topk_indices = None
+    indexer._store_mtp_shared_rows(rows, metadata, 0, 2, 4, 0, False, True, True)
+    assert torch.equal(metadata.shared_topk_indices, rows[[1, 3]])
+
+
+class _Bridged(Exception):
+    pass
+
+
+@pytest.mark.parametrize("fails", [False, True], ids=["returns", "raises"])
+def test_the_backend_hands_the_bridge_to_its_indexer_for_one_call(monkeypatch, fails):
+    seen = []
+    bridge_paths = {"rope_kv_append", "cross_rank_topk_global_emit"}
+    monkeypatch.setattr(dispatch, "is_enabled", lambda path: path in bridge_paths)
+
+    def get_impl(path, name):
+        assert (path, name) in {
+            ("cross_rank_topk_global_emit", "new_bridge_packet"),
+            ("cross_rank_topk_global_emit", "remap_topk_to_global"),
+        }, (path, name)
+        if name == "new_bridge_packet":
+            return lambda metadata, is_generation: ("bridge", is_generation)
+        return lambda backend, topk, metadata, bridge: ("global", topk, bridge)
+
+    monkeypatch.setattr(dispatch, "get_impl", get_impl)
+    backend = _attention().mqa
+    backend.__class__ = attention.GlmDsaBackend
+    indexer = backend.indexer
+    indexer.__class__ = attention.GlmIndexer
+    indexer.mtp_index_share = False
+    backend.get_local_layer_idx = lambda metadata: 0
+
+    def forward_from_projected(metadata, q, intermediates, is_generation):
+        seen.append((indexer._topk_bridge, is_generation))
+        if fails:
+            raise _Bridged
+        return "topk"
+
+    indexer.forward_from_projected = forward_from_projected
+    metadata = types.SimpleNamespace(
+        num_ctx_tokens=0,
+        num_tokens=6,
+        shared_topk_indices=None,
+        in_mtp_draft_loop=False,
+        kv_cache_manager=None,
+    )
+    args = types.SimpleNamespace(
+        attention_input_type=AttentionInputType.generation_only,
+        sparse_backend_args=types.SimpleNamespace(indexer_intermediates="intermediates"),
+    )
+
+    with dispatch.step(DECODE):
+        if fails:
+            with pytest.raises(_Bridged):
+                backend.sparse_attn_predict("q", None, metadata, args)
+        else:
+            out = backend.sparse_attn_predict("q", None, metadata, args)
+            assert out == (("global", "topk", ("bridge", True)), None)
+
+    assert seen == [(("bridge", True), True)]
+    assert indexer._topk_bridge is None
+
+
+# --- the attention-tail adoption and the absorbed-query module check ----------------------
+
+
+def test_the_absorbed_query_packet_takes_the_adopted_dsa_hooks(monkeypatch):
+    query_proj = pytest.importorskip(
+        f"{dispatch.__package__}.kernels.absorbed_query.query_proj",
+        reason="the absorbed_query kernels do not import here",
+    )
+    import tensorrt_llm._utils as utils
+
+    monkeypatch.setattr(utils, "get_sm_version", lambda: 103)
+    mla = types.SimpleNamespace(
+        sparse_attn_hooks=_new(attention.GlmDsaHooks), is_lite=True, fuse_qkv_a_proj=True
+    )
+    assert query_proj.module_eligibility(mla) == (False, "needs the fused [q_a | kv_a] projection")
+    mla.sparse_attn_hooks = _new(DSAHooks)
+    assert query_proj.module_eligibility(mla) == (False, "needs the fused [q_a | kv_a] projection")
+    mla.sparse_attn_hooks = _new(_other(DSAHooks))
+    assert query_proj.module_eligibility(mla) == (False, "not a DSA MLA module")
+
+
+# --- the target's DSA op 2 ----------------------------------------------------------------
+
+
+def _op_args():
+    t = torch.zeros(2, 4)
+    return (t, t, t, t, [t], None, "7", torch.zeros(2, 4))
+
+
+def test_the_op_runs_the_targets_op2_function_on_the_registered_layer(monkeypatch):
+    calls = []
+    monkeypatch.setattr(
+        attention_forward, "extract_extra_attrs", lambda layer, kind: ("md", f"mla{layer}{kind}")
+    )
+    monkeypatch.setattr(attention_forward, "_forward_dsa_attn", lambda *args: calls.append(args))
+    args = _op_args()
+
+    torch.ops.trtllm.glm_moe_dsa_mla_dsa_attn_inplace(*args)
+    attention_forward.maybe_bcg_glm_moe_dsa_mla_dsa_attn_inplace(*args)
+
+    assert len(calls) == 2
+    for call in calls:
+        assert call[0] == "mla7mla" and call[7] == "md" and call[6] is None
+        assert call[8] is args[-1]
+
+
+def test_the_op_mutates_only_its_output():
+    schema = torch.ops.trtllm.glm_moe_dsa_mla_dsa_attn_inplace.default._schema
+    assert [a.name for a in schema.arguments if a.alias_info and a.alias_info.is_write] == [
+        "output"
+    ]
+    assert len(schema.returns) == 0
+
+
+def test_the_ops_fake_runs_nothing(monkeypatch):
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    monkeypatch.setattr(attention_forward, "_forward_dsa_attn", lambda *a: pytest.fail())
+    with FakeTensorMode():
+        args = _op_args()
+        torch.ops.trtllm.glm_moe_dsa_mla_dsa_attn_inplace(*args)
+
+
+def test_only_the_target_registers_the_op():
+    """The op is this target's own: no other module of the package names it, and the target's
+    op-2 module registers it once, so it exists only after the target is imported.
+    """
+    name = "glm_moe_dsa_mla_dsa_attn_inplace"
+    module = Path(attention_forward.__file__).resolve()
+    package = next(parent for parent in module.parents if parent.name == "tensorrt_llm")
+    naming = [path for path in package.rglob("*.py") if name in path.read_text(encoding="utf-8")]
+    assert naming == [module]
+    source = module.read_text(encoding="utf-8")
+    assert source.count(f'@torch.library.custom_op("trtllm::{name}"') == 1
+    assert source.count(f"@{name}.register_fake") == 1
+    assert hasattr(torch.ops.trtllm, name)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py
new file mode 100644
index 000000000000..bebdfaeae377
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py
@@ -0,0 +1,534 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""How the GLM-5.2 tp8 target classifies a fast path's decline: the sites, the layer cells, the
+rules, and what ``log_capability`` does with each class.
+
+Host-only: the declines are synthetic calls with the keys and reasons the kernels log, the log
+goes into lists, and the sources are read as text, so nothing here needs a GPU.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+D = dispatch.Decline
+# The family: its shared code and every target of it.
+FAMILY = Path(dispatch.__file__).parents[1]
+SCOPE = "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE"
+LAYERS = "TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS"
+
+
+def _layers(spec: str) -> frozenset:
+    """``"0-2,6,10"`` as a set of layer indices."""
+    out = set()
+    for part in spec.split(","):
+        lo, _, hi = part.partition("-")
+        out |= set(range(int(lo), int(hi or lo) + 1))
+    return frozenset(out)
+
+
+SHARED = (
+    "3-5,7-9,11-13,15-17,19-21,23-25,27-29,31-33,35-37,39-41,43-45,47-49,51-53,55-57,59-61,"
+    "63-65,67-69,71-73,75-77"
+)
+FULL = "0-2,6,10,14,18,22,26,30,34,38,42,46,50,54,58,62,66,70,74"
+
+# The layers each per-layer key engages on, on every rank, when this target serves its
+# checkpoint: (path, key family, layers).
+CENSUS = (
+    ("absorbed_query", "layer", "0-77"),
+    ("attention_output_tail", "layer", "0-77"),
+    ("shared_indexer_absorbed_query", "layer", SHARED),
+    ("deferred_norm", "scope[shared]_into_layer", SHARED.replace("3-5", "4-5")),
+    ("sparse_mla_decode", "call:layer", "0-78"),
+    ("sparse_mla_decode", "layer", "0-78"),
+    ("draft_absorbed_query", "layer", "78"),
+    ("deferred_norm_evict_hints", "projection_layer", SHARED.replace("3-5", "4-5")),
+    ("indexer_layer_front", "draft_layer", "78"),
+    ("indexer_layer_front", "target_layer", FULL),
+    ("oproj_weight_staging", "layer", "0-77"),
+    ("fused_kv_indexer_projection", "layer", "3"),
+)
+
+# Every decline the served composition logs, on every rank: (path, key, reason). None of them
+# is a miss.
+SERVED_DECLINES = (
+    ("absorbed_query", "layer78", "target layers only"),
+    ("attention_output_tail", "layer78", "draft or unindexed layer"),
+    ("router_gemm", "router:not a target layer", "not a target layer"),
+    *(
+        (
+            "shared_indexer_absorbed_query",
+            f"layer{i}",
+            "full indexer layer (shared_indexer_absorbed_query serves shared layers)",
+        )
+        for i in sorted(_layers(FULL))
+    ),
+    ("shared_indexer_absorbed_query", "layer78", "target layers only"),
+    *(
+        (
+            "deferred_norm",
+            f"scope[shared]_into_layer{i}",
+            f"{SCOPE}=shared: layer {i} is a full-indexer layer; post_moe_hop hop",
+        )
+        for i in sorted(_layers(FULL) - {0, 1, 2})
+    ),
+    ("dense_mlp", "dense_mlp:row count outside [1, 8]", "row count outside [1, 8]"),
+    ("cross_layer_expert_prefetch", "spec_experts_declined", "128 rows outside [1, 8]"),
+    ("routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]"),
+    (
+        "sparse_mla_decode",
+        "request:not generation-only",
+        "not generation-only; first seen at layer 3 with 1 rows",
+    ),
+    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
+    ("draft_loop", "draft_loop:context rows present", "context rows present"),
+    ("draft_loop", "draft_step:dense top-k fallback active", "dense top-k fallback active"),
+    (
+        "deferred_input_prep",
+        "prep[no previous speculative device tensors]",
+        "no previous speculative device tensors",
+    ),
+    (
+        "deferred_input_prep",
+        "prep[not a steady decode-graph replay]",
+        "not a steady decode-graph replay",
+    ),
+    (
+        "decode_prologue",
+        "target_prologue[not a generation-only step]",
+        "not a generation-only step",
+    ),
+    *(
+        (
+            "indexer_layer_front",
+            f"target_layer{i}",
+            "shared-indexer layer (no q_a RMSNorm launch to remove; "
+            "shared_indexer_absorbed_query / absorbed_query serve it)",
+        )
+        for i in sorted(_layers(SHARED))
+    ),
+    ("fused_kv_indexer_projection", "layer78", f"draft layers not selected ({LAYERS})"),
+    *(
+        ("fused_kv_indexer_projection", f"layer{i}", f"full layers not selected ({LAYERS})")
+        for i in sorted(_layers(FULL))
+    ),
+    *(
+        ("fused_kv_indexer_projection", f"layer{i}", f"shared layers not selected ({LAYERS})")
+        for i in sorted(_layers(SHARED) - {3})
+    ),
+    *(
+        (
+            "draft_routed_moe",
+            f"rows={rows} unfinalized output requested",
+            "unfinalized output requested",
+        )
+        for rows in (1, 2, 6)
+    ),
+    ("draft_routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]"),
+    # Decided per object behind the first object's line: the deferred-norm hops into the
+    # full-indexer layers, and a dense attention instance of a sparse-attention layer.
+    (
+        "deferred_norm",
+        "deferred_hop",
+        f"{SCOPE}=shared: layer 6 is a full-indexer layer; post_moe_hop hop",
+    ),
+    (
+        "deferred_norm_evict_hints",
+        "deferred_hop",
+        f"{SCOPE}=shared: layer 6 is a full-indexer layer; post_moe_hop hop",
+    ),
+    ("sparse_mla_decode", "layer5", "not a DSA layer"),
+)
+
+# One decline per rule and site kind: (path, key, reason, class).
+RULES = (
+    # A path turned off.
+    ("sampler_store", "engaged", "turned off by TRTLLM_STAIRCASE_GLM_LL_DISABLE", D.DISABLED),
+    # A key no site names, and a path that does not exist.
+    ("routed_moe", "a_key_no_site_names", "workspace failed", D.UNCOVERED),
+    ("not_a_fast_path", "layer3", "workspace failed", D.UNCOVERED),
+    # Decided per call: a missing resource is a miss that is never raised.
+    (
+        "routed_moe",
+        "rows=6 not warmed up before CUDA graph capture",
+        "not warmed up before CUDA graph capture",
+        D.MISSED_PER_CALL,
+    ),
+    (
+        "rope_kv_append",
+        "scatter_defer",
+        "call site does not forward the deferred scatter",
+        D.MISSED_PER_CALL,
+    ),
+    (
+        "draft_front",
+        "draft_front_call_fallback",
+        "kernel not compiled before CUDA graph capture",
+        D.MISSED_PER_CALL,
+    ),
+    ("decode_prologue", "take_embedding", "packet does not match", D.MISSED_PER_CALL),
+    (
+        "indexer_layer_front",
+        "cp_indexer_decode_global_emit",
+        "workspace unavailable: peer rank",
+        D.MISSED_PER_CALL,
+    ),
+    (
+        "dense_mlp",
+        "dense_mlp:glm_moe_dsa_dense_mlp op not built",
+        "glm_moe_dsa_dense_mlp op not built",
+        D.MISSED_PER_CALL,
+    ),
+    ("indexer_topk", "indexer_decode", "CUTLASS DSL unavailable", D.MISSED_PER_CALL),
+    # Decided per call, outside the shapes (rows, metadata layout, step kind, KV band).
+    ("rope_kv_append", "rope", "outside the FP8 MLA decode domain", D.UNCOVERED),
+    ("sparse_mla_decode", "request:KV length layout", "KV length layout; first seen", D.UNCOVERED),
+    (
+        "decode_prologue",
+        "target_prologue[metadata is not DSAtrtllmAttentionMetadata]",
+        "metadata is not DSAtrtllmAttentionMetadata",
+        D.UNCOVERED,
+    ),
+    ("draft_tail", "draft_tail[graph]", "draft logits are not vocab-sharded", D.UNCOVERED),
+    (
+        "draft_tail",
+        "draft_tail[step: 9 rows outside [1, 8]]",
+        "identical on every TP rank",
+        D.UNCOVERED,
+    ),
+    ("sparse_mla_decode", "call:layer5", "first call", D.UNCOVERED),
+    # Decided once: a layer kind or a composition the path does not serve.
+    ("attention_output_into_hop", "layer3", "needs the post_attention_hop hop", D.UNCOVERED),
+    (
+        "attention_output_tail_prefetch",
+        "layer3",
+        "post_attention_hop without sparse_mla_decode: native tail is faster",
+        D.UNCOVERED,
+    ),
+    (
+        "attention_output_tail_prefetch",
+        "schedule",
+        "post_attention_hop without sparse_mla_decode",
+        D.UNCOVERED,
+    ),
+    (
+        "draft_live_rows",
+        "live_row_hop:no post_attention_hop",
+        "no post_attention_hop hop",
+        D.UNCOVERED,
+    ),
+    ("sparse_mla_decode", "layer40", "not MLA", D.UNCOVERED),
+    # Decided once, on a layer outside the cells.
+    ("fused_kv_indexer_projection", "layer4", "kernel compile failed", D.UNCOVERED),
+    ("attention_output_tail", "layer78", "staging buffer", D.UNCOVERED),
+    # Decided once, inside the cells or on a key without a layer: a miss.
+    ("fused_kv_indexer_projection", "layer3", "kernel compile failed", D.MISSED),
+    ("sparse_mla_decode", "layer12", "latent dims", D.MISSED),
+    ("sparse_mla_decode", "layer78", "KV cache is not FP8", D.MISSED),
+    ("sparse_mla_decode", "layer0", "rope_append disabled", D.MISSED),
+    ("indexer_layer_front", "rope", "indexer_layer_front rope op not loaded", D.MISSED),
+    (
+        "post_attention_hop",
+        "post_attn_hop:native all-reduce is not MNNVL",
+        "native all-reduce is not MNNVL",
+        D.MISSED,
+    ),
+    ("post_attention_hop", "post_attn_hop", "", D.MISSED),
+    ("router_gemm", "router:attention DP", "attention DP", D.MISSED),
+    ("post_moe_hop", "cluster_kernel", "the cluster kernel needs SM100", D.MISSED),
+    ("post_moe_hop", "post_moe_hop[cluster]", "workspace agreement failed", D.MISSED),
+    ("routed_moe_into_hop", "post_moe_routed_input_buffer", "no buffer", D.MISSED),
+    (
+        "cross_layer_expert_prefetch",
+        "link[k=8,start=moe]",
+        "0 MoE layers, 0 routed-kernel prefetch hooks",
+        D.MISSED,
+    ),
+    ("draft_front", "draft_front", "requires non-ADP tensor parallelism", D.MISSED),
+    ("draft_tail", "draft_tail[tp8]", "workspace setup failed: OSError()", D.MISSED),
+    ("host_page_retention", "host_policy", "libc10.so: not found", D.MISSED),
+    # The loop domain is final only once the overlap loop has declared it.
+    (
+        "deferred_input_prep",
+        "loop_domain",
+        "no overlap executor loop declared its domain; the store writer and the single-stream "
+        "loop are validated only there",
+        D.UNCOVERED,
+    ),
+    (
+        "deferred_input_prep",
+        "loop_domain",
+        "attention DP: the store writer and the single-stream loop are not validated with it",
+        D.MISSED,
+    ),
+)
+
+
+@pytest.fixture
+def capability(monkeypatch):
+    """A fresh capability log and fallback state; returns the info and warning lines."""
+    info, warnings = [], []
+
+    class _Logger:
+        @staticmethod
+        def info(message):
+            info.append(message)
+
+        @staticmethod
+        def warning(message):
+            warnings.append(message)
+
+    monkeypatch.setattr(dispatch, "_log_info", info.append)
+    monkeypatch.setattr(dispatch, "logger", _Logger)
+    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
+    dispatch.reset_fallback_state()
+    dispatch.set_require(False)
+    dispatch.classify_decline.cache_clear()
+    yield info, warnings
+    dispatch.reset_fallback_state()
+    dispatch.set_require(False)
+
+
+# --- the tables ---------------------------------------------------------------------------
+
+
+def test_the_layer_cells_are_the_census():
+    expected = {}
+    for path, family, layers in CENSUS:
+        expected.setdefault(path, {})[family] = _layers(layers)
+    assert {path: dict(families) for path, families in dispatch.LAYER_CELLS.items()} == expected
+
+
+def test_every_site_names_a_fast_path_and_every_fast_path_has_a_site():
+    assert {site.path for site in dispatch.DECLINE_KINDS} == set(dispatch.FAST_PATH_NAMES)
+
+
+def test_every_family_with_cells_is_a_site_of_its_path():
+    families = {(site.path, site.cells) for site in dispatch.DECLINE_KINDS if site.cells}
+    with_cells = {
+        (path, family) for path, cells in dispatch.LAYER_CELLS.items() for family in cells
+    }
+    assert with_cells <= families
+    # The per-layer keys the served composition never logs have no cells: every decline there
+    # is outside the path's shapes.
+    assert families - with_cells == {
+        ("deferred_norm", "projection_layer"),
+        ("attention_output_tail_prefetch", "layer"),
+        ("attention_output_into_hop", "layer"),
+    }
+
+
+def test_a_site_with_cells_captures_the_layer():
+    for site in dispatch.DECLINE_KINDS:
+        assert ("layer" in site.key.groupindex) == (site.cells is not None), site
+
+
+def _kernel_sources() -> list:
+    """The sources that log declines: every file of the family except the classifier itself."""
+    return [p for p in sorted(FAMILY.rglob("*.py")) if p != Path(dispatch.__file__)]
+
+
+def _source_text() -> str:
+    return "\n".join(p.read_text() for p in _kernel_sources())
+
+
+@pytest.mark.parametrize("fragment", dispatch.LAYER_KIND_REASONS + dispatch.DISABLED_REASONS)
+def test_every_reason_fragment_is_a_reason_the_kernels_give(fragment):
+    """A renamed reason must rename its fragment too, or its declines change class."""
+    assert fragment in _source_text()
+
+
+def _string_literals() -> set:
+    return {
+        node.value
+        for source in _kernel_sources()
+        for node in ast.walk(ast.parse(source.read_text()))
+        if isinstance(node, ast.Constant) and isinstance(node.value, str)
+    }
+
+
+def test_the_served_reasons_are_the_kernels_literals():
+    """Each constant reason of a served decline is a whole string literal of the kernels, so a
+    reworded reason fails here before its declines change class."""
+    literals = _string_literals()
+    # The reasons without a formatted field (a number or a knob name).
+    reasons = {
+        reason
+        for _, _, reason in SERVED_DECLINES
+        if not any(c.isdigit() for c in reason) and dispatch.KNOB_PREFIX not in reason
+    }
+    assert len(reasons) == 13
+    assert sorted(reasons - literals) == []
+
+
+def _render(node: ast.AST):
+    """A key the call can log: f-string fields become a sample value."""
+    if isinstance(node, ast.Constant):
+        return node.value
+    if not isinstance(node, ast.JoinedStr):
+        return None
+    out = []
+    for value in node.values:
+        if isinstance(value, ast.Constant):
+            out.append(value.value)
+        elif isinstance(value.value, ast.IfExp) and isinstance(value.value.body, ast.Constant):
+            out.append(value.value.body.value)
+        elif isinstance(value.value, ast.Attribute) and value.value.attr == "scope":
+            out.append("shared")
+        else:
+            out.append("0")
+    return "".join(out)
+
+
+def test_every_logged_key_has_a_site():
+    """Each ``log_capability`` call whose path and key can be read off the source logs a key
+    that a site of its path matches."""
+    checked, missing = 0, []
+    for source in _kernel_sources():
+        tree = ast.parse(source.read_text())
+        constants = {
+            node.targets[0].id: node.value.value
+            for node in tree.body
+            if isinstance(node, ast.Assign)
+            and len(node.targets) == 1
+            and isinstance(node.targets[0], ast.Name)
+            and isinstance(node.value, ast.Constant)
+            and isinstance(node.value.value, str)
+        }
+        for node in ast.walk(tree):
+            if not (isinstance(node, ast.Call) and len(node.args) >= 2):
+                continue
+            name = getattr(node.func, "attr", getattr(node.func, "id", None))
+            if name != "log_capability":
+                continue
+            first = node.args[0]
+            path = (
+                first.value
+                if isinstance(first, ast.Constant)
+                else constants.get(getattr(first, "id", None))
+            )
+            key = _render(node.args[1])
+            if path is None or key is None:
+                continue
+            checked += 1
+            if dispatch._decline_site(path, key) is None:
+                missing.append(f"{source.name}:{node.lineno} {path} {key!r}")
+    assert missing == []
+    assert checked >= 70
+
+
+# --- the rules ----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path, key, reason", SERVED_DECLINES)
+def test_no_decline_of_the_served_composition_is_a_miss(path, key, reason):
+    assert dispatch.classify_decline(path, key, reason) is D.UNCOVERED
+
+
+@pytest.mark.parametrize("path, key, reason, kind", RULES)
+def test_each_rule(path, key, reason, kind):
+    assert dispatch.classify_decline(path, key, reason) is kind
+
+
+def test_a_resource_reason_is_matched_in_any_case():
+    assert (
+        dispatch.classify_decline("routed_moe", "rows=4 x", "Workspace FAILED") is D.MISSED_PER_CALL
+    )
+
+
+def test_a_repeated_decline_is_one_lookup(capability):
+    args = ("routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]")
+    for _ in range(5):
+        dispatch.log_capability(*args[:2], False, args[2])
+    info = dispatch.classify_decline.cache_info()
+    assert (info.misses, info.hits) == (1, 4)
+
+
+# --- what log_capability does with each class ---------------------------------------------
+
+
+def test_an_engaged_decision_is_not_classified(capability):
+    assert dispatch.log_capability("post_moe_hop", "post_moe_hop[cluster]", True)
+    assert dispatch.classify_decline.cache_info().misses == 0
+
+
+def test_a_miss_decided_once_raises_under_require_after_its_line(capability):
+    info, warnings = capability
+    dispatch.set_require(True)
+    with pytest.raises(
+        dispatch.FastPathUnavailable, match=r"'post_moe_hop'.*post_moe_hop\[cluster\]: workspace"
+    ):
+        dispatch.log_capability("post_moe_hop", "post_moe_hop[cluster]", False, "workspace failed")
+    assert info == [
+        "[glm_moe_dsa] post_moe_hop/post_moe_hop[cluster]: native fallback (workspace failed)"
+    ]
+    assert warnings == []
+
+
+def test_a_miss_decided_once_is_counted_and_warned_once_under_auto(capability):
+    info, warnings = capability
+    for _ in range(2):
+        assert not dispatch.log_capability("sparse_mla_decode", "layer7", False, "latent dims")
+    assert dispatch.missed_counts() == {"sparse_mla_decode": 2}
+    assert warnings == [
+        "[glm_moe_dsa] sparse_mla_decode/missed: layer7: latent dims; falling back (logged once per reason)"
+    ]
+    assert len(info) == 1
+
+
+@pytest.mark.parametrize("require", [False, True], ids=["auto", "require"])
+def test_a_miss_decided_per_call_is_counted_and_never_raised(capability, require):
+    info, warnings = capability
+    dispatch.set_require(require)
+    key, reason = (
+        "rows=6 not warmed up before CUDA graph capture",
+        "not warmed up before CUDA graph capture",
+    )
+    for _ in range(3):
+        assert not dispatch.log_capability("routed_moe", key, False, reason)
+    assert dispatch.missed_counts() == {"routed_moe": 3}
+    assert warnings == [
+        f"[glm_moe_dsa] routed_moe/missed: {reason}; falling back (logged once per reason)"
+    ]
+    assert info == [f"[glm_moe_dsa] routed_moe/{key}: native fallback ({reason})"]
+
+
+@pytest.mark.parametrize(
+    "path, key, reason",
+    [
+        ("absorbed_query", "layer78", "target layers only"),
+        ("routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]"),
+        ("sampler_store", "engaged", "turned off by TRTLLM_STAIRCASE_GLM_LL_DISABLE"),
+        ("routed_moe", "a_key_no_site_names", "workspace failed"),
+    ],
+    ids=["uncovered-once", "uncovered-per-call", "disabled", "unknown-site"],
+)
+@pytest.mark.parametrize("require", [False, True], ids=["auto", "require"])
+def test_uncovered_and_disabled_declines_add_nothing_to_their_line(
+    capability, path, key, reason, require
+):
+    info, warnings = capability
+    dispatch.set_require(require)
+    for _ in range(2):
+        assert not dispatch.log_capability(path, key, False, reason)
+    assert info == [f"[glm_moe_dsa] {path}/{key}: native fallback ({reason})"]
+    assert warnings == [] and dispatch.missed_counts() == {}
+
+
+def test_the_served_declines_are_not_raised_under_require(capability):
+    """Under require, the served composition's declines run as before: logged, not raised."""
+    dispatch.set_require(True)
+    for path, key, reason in SERVED_DECLINES:
+        dispatch.log_capability(path, key, False, reason)
+    _, warnings = capability
+    assert warnings == [] and dispatch.missed_counts() == {}
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py
new file mode 100644
index 000000000000..518ffcd12525
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py
@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The composition of the GLM-5.2 tp8 target: the fast-path table, the DISABLE knob, the
+fallback policy and the rank-uniform checks.
+
+Host-only: the knobs are passed as an explicit environment and the collective is a fake
+two-rank group or a fake MPI communicator, so nothing here needs a GPU or MPI.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+import tensorrt_llm._utils as trtllm_utils
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode
+from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
+    dispatch,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+AUTO, REQUIRE = ModelingV2Mode.AUTO, ModelingV2Mode.REQUIRE
+DISABLE = dispatch.DISABLE_ENV
+ACCUMULATION = dispatch.KNOB_PREFIX + "ROUTED_MOE_INTO_HOP_ACCUMULATION"
+
+# The feature vocabulary, in this order: the fast-path names, the DISABLE tokens and the
+# knob infixes are the same words.
+VOCABULARY = (
+    "absorbed_query",
+    "draft_absorbed_query",
+    "shared_indexer_absorbed_query",
+    "fused_kv_indexer_projection",
+    "indexer_layer_front",
+    "deferred_norm",
+    "deferred_norm_evict_hints",
+    "rope_kv_append",
+    "indexer_topk",
+    "cross_rank_topk",
+    "cross_rank_topk_cluster_merge",
+    "cross_rank_topk_global_emit",
+    "sparse_mla_decode",
+    "attention_output_tail",
+    "attention_output_tail_prefetch",
+    "attention_output_into_hop",
+    "oproj_weight_staging",
+    "post_attention_hop",
+    "hop_routing",
+    "router_gemm",
+    "routed_moe",
+    "routed_moe_into_hop",
+    "shared_expert_prefetch",
+    "post_moe_hop",
+    "cross_layer_expert_prefetch",
+    "dense_mlp",
+    "draft_front",
+    "draft_routed_moe",
+    "draft_live_rows",
+    "draft_tail",
+    "draft_loop",
+    "decode_prologue",
+    "deferred_input_prep",
+    "sampler_store",
+    "execution_stream_loop",
+    "host_page_retention",
+)
+
+
+@pytest.fixture(autouse=True)
+def _fresh_fallback_state():
+    dispatch.set_require(False)
+    dispatch.reset_fallback_state()
+    yield
+    dispatch.set_require(False)
+    dispatch.reset_fallback_state()
+
+
+def _off(plan: dispatch.Plan) -> set:
+    return {name for name in dispatch.FAST_PATH_NAMES if not plan.enabled(name)}
+
+
+# --- the table --------------------------------------------------------------------------
+
+
+def test_the_fast_paths_are_the_vocabulary_in_order():
+    """The DISABLE tokens are these words; the log line lists them in this order."""
+    assert dispatch.FAST_PATH_NAMES == VOCABULARY
+
+
+def test_dependencies_name_fast_paths_and_have_no_cycle():
+    names = set(dispatch.FAST_PATH_NAMES)
+    for path in dispatch.FAST_PATHS:
+        assert set(path.depends_on) <= names, path
+        assert path.name not in path.depends_on, path
+        # A base that is also a (transitive) dependent of the path would be a cycle.
+        dependents = dispatch.disable_closure(frozenset({path.name})) - {path.name}
+        assert not set(path.depends_on) & dependents, path
+
+
+# --- DISABLE ----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("raw", [None, "", " , ,"], ids=["unset", "empty", "blank"])
+def test_every_fast_path_runs_by_default(raw):
+    environ = {} if raw is None else {DISABLE: raw}
+    plan = dispatch.resolve(AUTO, environ)
+    assert _off(plan) == set()
+    assert plan.disabled == plan.disabled_dependents == frozenset()
+
+
+def test_disabling_a_path_without_dependents_turns_off_only_it():
+    plan = dispatch.resolve(AUTO, {DISABLE: "draft_tail"})
+    assert _off(plan) == {"draft_tail"}
+    assert plan.disabled_dependents == frozenset()
+
+
+@pytest.mark.parametrize(
+    "disabled, dependents",
+    [
+        (
+            "routed_moe",
+            {
+                "routed_moe_into_hop",
+                "hop_routing",
+                "shared_expert_prefetch",
+                "cross_layer_expert_prefetch",
+            },
+        ),
+        (
+            "post_attention_hop",
+            {
+                "attention_output_into_hop",
+                "oproj_weight_staging",
+                "hop_routing",
+                "shared_expert_prefetch",
+                "draft_live_rows",
+                "host_page_retention",
+            },
+        ),
+        (
+            "cross_rank_topk",
+            {
+                "cross_rank_topk_cluster_merge",
+                "cross_rank_topk_global_emit",
+                "indexer_layer_front",
+            },
+        ),
+        ("draft_loop", {"draft_live_rows", "decode_prologue"}),
+    ],
+)
+def test_disabling_a_path_turns_off_every_path_that_depends_on_it(disabled, dependents):
+    plan = dispatch.resolve(AUTO, {DISABLE: disabled})
+    assert plan.disabled == {disabled}
+    assert plan.disabled_dependents == dependents
+    assert _off(plan) == {disabled} | dependents
+
+
+def test_disable_ignores_case_and_blanks():
+    plan = dispatch.resolve(AUTO, {DISABLE: " Draft_Tail ,, DENSE_MLP "})
+    assert _off(plan) == {"draft_tail", "dense_mlp"}
+
+
+def test_disable_naming_something_else_raises():
+    with pytest.raises(ValueError, match=r"not_a_fast_path") as excinfo:
+        dispatch.resolve(AUTO, {DISABLE: "draft_tail,not_a_fast_path"})
+    assert DISABLE in str(excinfo.value)
+
+
+def test_a_call_site_asking_for_a_misspelled_path_fails():
+    plan = dispatch.resolve(AUTO, {})
+    with pytest.raises(KeyError, match="draft_tails"):
+        plan.enabled("draft_tails")
+
+
+def test_only_this_targets_knobs_are_recorded():
+    plan = dispatch.resolve(
+        AUTO, {DISABLE: "draft_tail", ACCUMULATION: "identity", "UNRELATED": "1"}
+    )
+    assert plan.knobs == ((DISABLE, "draft_tail"), (ACCUMULATION, "identity"))
+
+
+# --- missed and uncovered ---------------------------------------------------------------
+
+
+def test_a_miss_raises_under_require():
+    dispatch.resolve(REQUIRE, {})
+    with pytest.raises(dispatch.FastPathUnavailable, match=r"'routed_moe'.*at 6 rows.*workspace"):
+        dispatch.missed("routed_moe", 6, "its workspace is not built")
+
+
+def test_a_miss_falls_back_and_is_counted_under_auto():
+    dispatch.resolve(AUTO, {})
+    dispatch.missed("routed_moe", None, "its workspace is not built")
+    dispatch.missed("routed_moe", 6, "its workspace is not built")
+    assert dispatch.missed_counts() == {"routed_moe": 2}
+
+
+@pytest.mark.parametrize("mode", [AUTO, REQUIRE])
+def test_an_uncovered_step_never_raises(mode):
+    dispatch.resolve(mode, {})
+    dispatch.uncovered("sparse_mla_decode", "more rows than the kernel takes")
+    assert dispatch.missed_counts() == {}
+
+
+@pytest.mark.parametrize("mode", [AUTO, REQUIRE])
+def test_a_miss_decided_per_call_never_raises_and_is_counted(mode):
+    dispatch.resolve(mode, {})
+    dispatch.missed("routed_moe", 6, "not warmed up before CUDA graph capture", per_call=True)
+    dispatch.missed("routed_moe", 6, "not warmed up before CUDA graph capture", per_call=True)
+    assert dispatch.missed_counts() == {"routed_moe": 2}
+
+
+# --- rank uniformity --------------------------------------------------------------------
+
+
+def _two_rank_group(peer_environ: dict, calls: list):
+    """A fake allgather over two ranks: this rank first, then a peer with ``peer_environ``."""
+
+    def allgather(knobs, world_size):
+        calls.append((knobs, world_size))
+        return [knobs, dispatch.read_knobs(peer_environ)]
+
+    return allgather
+
+
+def test_ranks_that_agree_pass_with_one_collective():
+    env = {DISABLE: "draft_tail", ACCUMULATION: "identity"}
+    calls = []
+    knobs = dispatch.read_knobs(env)
+    dispatch.check_rank_uniform(knobs, 2, _two_rank_group(dict(env), calls))
+    assert calls == [(knobs, 2)]
+
+
+def test_unrelated_environment_does_not_count():
+    dispatch.check_rank_uniform(
+        dispatch.read_knobs({"UNRELATED": "a"}), 2, _two_rank_group({"UNRELATED": "b"}, [])
+    )
+
+
+def test_ranks_that_would_run_different_fast_paths_raise_with_the_rank_table():
+    with pytest.raises(RuntimeError) as excinfo:
+        dispatch.check_rank_uniform(
+            dispatch.read_knobs({}), 2, _two_rank_group({DISABLE: "routed_moe"}, [])
+        )
+    message = str(excinfo.value)
+    rank0, rank1 = message.split("rank 1:")
+    assert "rank 0: knobs" in rank0 and f"{DISABLE}='<unset>'" in rank0
+    assert f"{DISABLE}='routed_moe'" in rank1
+
+
+def test_ranks_that_read_different_knob_values_raise():
+    with pytest.raises(RuntimeError, match=rf"{ACCUMULATION}='identity'"):
+        dispatch.check_rank_uniform(
+            dispatch.read_knobs({ACCUMULATION: "fp32"}),
+            2,
+            _two_rank_group({ACCUMULATION: "identity"}, []),
+        )
+
+
+def test_a_value_one_rank_cannot_parse_fails_the_check_on_every_rank():
+    """The check compares the raw values, so the rank with a valid value raises too."""
+    with pytest.raises(RuntimeError, match=rf"{DISABLE}='not_a_fast_path'"):
+        dispatch.check_rank_uniform(
+            dispatch.read_knobs({DISABLE: "draft_tail"}),
+            2,
+            _two_rank_group({DISABLE: "not_a_fast_path"}, []),
+        )
+
+
+def test_without_a_communicator_the_check_is_skipped():
+    dispatch.check_rank_uniform(dispatch.read_knobs({}), 2, lambda knobs, world_size: None)
+
+
+def test_the_ranks_agree_on_the_knobs_the_shared_switches_and_the_mode():
+    env = {
+        DISABLE: "draft_tail",
+        "TRTLLM_FORCE_MNNVL_AR": "1",
+        "TLLM_TEST_MNNVL": "0",
+        "TRTLLM_ENABLE_PDL": "1",
+        "TRTLLM_MODELING_V2_VALIDATE": "1",
+        "UNRELATED": "1",
+    }
+    inputs = dispatch.rank_uniform_inputs(dispatch.read_knobs(env), REQUIRE, env)
+    assert inputs == (
+        ("TLLM_TEST_MNNVL", "0"),
+        ("TRTLLM_ENABLE_PDL", "1"),
+        ("TRTLLM_FORCE_MNNVL_AR", "1"),
+        ("TRTLLM_MODELING_V2_VALIDATE", "1"),
+        (DISABLE, "draft_tail"),
+        (dispatch.MODE_KEY, "require"),
+    )
+    assert dispatch.rank_uniform_inputs((), AUTO, {}) == ((dispatch.MODE_KEY, "auto"),)
+
+
+@pytest.mark.parametrize(
+    "peer_env, peer_mode",
+    [({"TRTLLM_FORCE_MNNVL_AR": "1"}, REQUIRE), ({}, AUTO)],
+    ids=["shared-switch", "mode"],
+)
+def test_ranks_that_differ_in_a_shared_switch_or_the_mode_raise(peer_env, peer_mode):
+    mine = dispatch.rank_uniform_inputs((), REQUIRE, {})
+    peer = dispatch.rank_uniform_inputs((), peer_mode, peer_env)
+    with pytest.raises(RuntimeError, match="rank 1: knobs") as excinfo:
+        dispatch.check_rank_uniform(mine, 2, lambda value, world_size: [value, peer])
+    assert dispatch.MODE_KEY in str(excinfo.value)
+
+
+_STATE = (
+    ("model.layers.0", (("post_attn_hop", "PostAttentionHop"),)),
+    ("model.layers.1", (("post_moe_hop", "PostMoeHop", "engaged=True"),)),
+)
+
+
+def test_ranks_that_attached_the_same_state_pass_with_one_collective():
+    calls = []
+
+    def allgather(value, world_size):
+        calls.append(world_size)
+        return [value, value]
+
+    dispatch.check_rank_uniform_adopted(_STATE, 2, allgather)
+    assert calls == [2]
+
+
+def test_ranks_that_attached_different_state_raise_naming_the_first_module():
+    peer = (_STATE[0], ("model.layers.1", ()))
+    with pytest.raises(RuntimeError) as excinfo:
+        dispatch.check_rank_uniform_adopted(_STATE, 2, lambda value, world_size: [value, peer])
+    message = str(excinfo.value)
+    assert (
+        "rank 0: ('model.layers.1', (('post_moe_hop', 'PostMoeHop', 'engaged=True'),))" in message
+    )
+    assert "rank 1: ('model.layers.1', ())" in message
+    assert "model.layers.0" not in message
+
+
+def test_a_rank_with_fewer_modules_is_named():
+    with pytest.raises(RuntimeError, match=r"rank 1: <no module>"):
+        dispatch.check_rank_uniform_adopted(_STATE, 2, lambda value, world_size: [value, value[:1]])
+
+
+def test_without_a_communicator_the_adopted_check_is_skipped():
+    dispatch.check_rank_uniform_adopted(_STATE, 2, lambda value, world_size: None)
+
+
+class _Communicator:
+    def __init__(self, size: int) -> None:
+        self.size = size
+        self.gathered = []
+
+    def Get_size(self) -> int:  # noqa: N802 - the mpi4py name
+        return self.size
+
+    def allgather(self, knobs):
+        self.gathered.append(knobs)
+        return [knobs] * self.size
+
+
+def test_the_session_communicator_of_this_model_is_gathered_once(monkeypatch):
+    comm = _Communicator(8)
+    monkeypatch.setattr(trtllm_utils, "mpi_disabled", lambda: False)
+    monkeypatch.setattr(trtllm_utils, "mpi_comm", lambda: comm)
+    knobs = dispatch.read_knobs({DISABLE: "draft_tail"})
+    dispatch.check_rank_uniform(knobs, 8)
+    assert comm.gathered == [knobs]
+
+
+@pytest.mark.parametrize("size", [1, 16], ids=["singleton", "larger"])
+def test_a_communicator_that_is_not_this_models_ranks_is_not_gathered(monkeypatch, size):
+    """A singleton proves nothing; a larger one would wait on ranks that never build the model."""
+    comm = _Communicator(size)
+    monkeypatch.setattr(trtllm_utils, "mpi_disabled", lambda: False)
+    monkeypatch.setattr(trtllm_utils, "mpi_comm", lambda: comm)
+    dispatch.check_rank_uniform(dispatch.read_knobs({}), 8)
+    assert comm.gathered == []
+
+
+def test_without_mpi_nothing_is_gathered(monkeypatch):
+    monkeypatch.setattr(trtllm_utils, "mpi_disabled", lambda: True)
+    monkeypatch.setattr(trtllm_utils, "mpi_comm", lambda: pytest.fail("no communicator expected"))
+    dispatch.check_rank_uniform(dispatch.read_knobs({}), 8)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py
new file mode 100644
index 000000000000..b6e9335663b1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py
@@ -0,0 +1,297 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Nothing outside the glm_moe_dsa family reaches its targets, except the router line, and
+no target reaches another.
+
+With ``modeling_v2`` off the resolver returns before importing anything, so the
+target can change no other model's behaviour as long as no other module imports
+it. No shared module imports the target (``_S4_SHARED_IMPORTS`` is empty): the
+engine reaches its step-loop hooks through ``runtime_hooks``, and every other
+reference fails. Read off disk as text, like test_modeling_v2_claims.py: no
+module is imported.
+
+No shared Python or C++ file uses the family's names either (the C++ build adds the family's
+folders in a few CMake lines), and no model of the zoo declares step-loop hooks, so each runs
+with the engine's all-off default; that check imports the zoo's model modules.
+"""
+
+from __future__ import annotations
+
+import ast
+import re
+from collections import Counter
+from pathlib import Path
+
+import pytest
+
+import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2
+
+pytestmark = pytest.mark.cpu_only
+
+_MODELING_V2 = Path(_modeling_v2.__file__).resolve().parent
+_PACKAGE = _MODELING_V2.parents[2]
+_FAMILY = _MODELING_V2 / "models" / "glm_moe_dsa"
+_ROUTER_INDEX = _MODELING_V2 / "_router_index.py"
+
+# A reference to the family's modules or to the targets' classes.
+_REFERENCE = re.compile(
+    r"models[./]glm_moe_dsa\b|glm_5_2_nvfp4__sm_103__\w+|ModelingV2Glm52|GlmMoeDsa(Tp8|Tep|Dep)ModelingV2"
+)
+_ROUTER_LINE = '"GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing",'
+
+_TARGET = "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa"
+# The imports of the target from shared modules: (file, module, name, alias) -> count.
+# They are matched on the parsed import, so the formatter may wrap them.
+_S4_SHARED_IMPORTS: dict = {}
+
+
+def _target_imports(path: Path, rel: str):
+    """``(key, first line, last line)`` of every import of the target in ``path``."""
+    out = []
+    for node in ast.walk(ast.parse(path.read_text(errors="ignore"))):
+        if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module:
+            if node.module == _TARGET or node.module.startswith(_TARGET + "."):
+                for alias in node.names:
+                    key = (rel, node.module, alias.name, alias.asname)
+                    out.append((key, node.lineno, node.end_lineno))
+    return out
+
+
+def _shared_files():
+    for path in sorted(_PACKAGE.rglob("*.py")):
+        if _FAMILY not in path.parents:
+            yield path, str(path.relative_to(_PACKAGE.parent))
+
+
+def test_only_the_router_line_references_the_family():
+    offenders = []
+    for path, rel in _shared_files():
+        lines = path.read_text(errors="ignore").splitlines()
+        hits = [n for n, line in enumerate(lines, 1) if _REFERENCE.search(line)]
+        if not hits:
+            continue
+        allowed = set()
+        for key, first, last in _target_imports(path, rel):
+            if key in _S4_SHARED_IMPORTS:
+                allowed.update(range(first, last + 1))
+        for lineno in hits:
+            line = lines[lineno - 1]
+            if path == _ROUTER_INDEX and line.strip() == _ROUTER_LINE:
+                continue
+            if lineno in allowed:
+                continue
+            offenders.append(f"{rel}:{lineno}: {line.strip()}")
+    assert not offenders, "modules outside glm_moe_dsa reference it:\n" + "\n".join(offenders)
+
+
+def test_each_listed_shared_import_is_present():
+    found = Counter()
+    for path, rel in _shared_files():
+        if "glm_moe_dsa" in path.read_text(errors="ignore"):
+            found.update(key for key, _, _ in _target_imports(path, rel))
+    assert dict(found) == _S4_SHARED_IMPORTS
+
+
+def test_the_router_line_is_present():
+    lines = [line.strip() for line in _ROUTER_INDEX.read_text().splitlines()]
+    assert lines.count(_ROUTER_LINE) == 1
+
+
+# The family's names; no file outside the family uses them.
+_FORK_NAMES = re.compile(r"glm_ll_|\[glm_ll\]|GLM_LL|TRTLLM_STAIRCASE")
+_CPP_FORK_NAMES = re.compile(r"glm_ll|GLM_LL|glmMoeDsa")
+_REPO = _PACKAGE.parent
+_CPP = _REPO / "cpp" / "tensorrt_llm"
+_CPP_FAMILY_DIRS = ("kernels/glmMoeDsa", "thop/glmMoeDsa")
+# The lines each parent CMakeLists.txt may spend on the family's folder.
+_CMAKE_LINES = 3
+
+
+def test_no_shared_file_names_the_fork():
+    offenders = []
+    for path, rel in _shared_files():
+        for lineno, line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
+            if _FORK_NAMES.search(line):
+                offenders.append(f"{rel}:{lineno}: {line.strip()}")
+    assert not offenders, "modules outside glm_moe_dsa name the fork:\n" + "\n".join(offenders)
+
+
+def test_no_shared_cpp_file_names_the_family():
+    if not _CPP.is_dir():
+        pytest.skip("the C++ sources are not in this tree")
+    family = [_CPP / d for d in _CPP_FAMILY_DIRS]
+    offenders, cmake_lines = [], Counter()
+    for path in sorted(_CPP.rglob("*")):
+        if not path.is_file() or path.suffix not in (".cpp", ".cu", ".h", ".txt"):
+            continue
+        if any(d == path.parent or d in path.parents for d in family):
+            continue
+        rel = str(path.relative_to(_REPO))
+        for lineno, line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
+            if not _CPP_FORK_NAMES.search(line):
+                continue
+            if path.name == "CMakeLists.txt" and path.parent in {d.parent for d in family}:
+                cmake_lines[rel] += 1
+                continue
+            offenders.append(f"{rel}:{lineno}: {line.strip()}")
+    assert not offenders, "C++ files outside the family's folders name it:\n" + "\n".join(offenders)
+    assert all(count <= _CMAKE_LINES for count in cmake_lines.values()), dict(cmake_lines)
+
+
+def _zoo_model_classes():
+    """Class name -> providing module of every model of the zoo (imported lazily)."""
+    from tensorrt_llm._torch.models._arch_index import MODEL_CLASS_TO_MODULE
+
+    classes = dict(MODEL_CLASS_TO_MODULE)
+    classes.setdefault("DeepseekV3ForCausalLM", "modeling_deepseekv3")
+    classes.setdefault("GlmMoeDsaForCausalLM", "modeling_deepseekv3")
+    return classes
+
+
+_ZOO = _zoo_model_classes()
+
+
+def test_no_zoo_model_declares_runtime_hooks():
+    """Read off disk: no model of the zoo names the hooks, so each gets the all-off default."""
+    models = _PACKAGE / "_torch" / "models"
+    offenders = [
+        str(path.relative_to(_PACKAGE.parent))
+        for path in sorted(models.rglob("*.py"))
+        if "runtime_hooks" in path.read_text(errors="ignore")
+    ]
+    assert not offenders, offenders
+
+
+@pytest.mark.parametrize("name", sorted(_ZOO))
+def test_zoo_models_declare_no_runtime_hooks(name):
+    """Every model class of the zoo, without its constructor, reads as the all-off hooks."""
+    import importlib
+
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, runtime_hooks_of
+
+    try:
+        providing = importlib.import_module(f"tensorrt_llm._torch.models.{_ZOO[name]}")
+    except ModuleNotFoundError as err:
+        if err.name and err.name.split(".")[0] != "tensorrt_llm":
+            pytest.skip(f"{err.name} is not installed")
+        raise
+    cls = getattr(providing, name)
+    assert runtime_hooks_of(cls.__new__(cls)) is DEFAULT_RUNTIME_HOOKS
+
+
+# --- the targets ----------------------------------------------------------------------------
+
+_TARGET_DIRS = {
+    "tp8": _FAMILY / "glm_5_2_nvfp4__sm_103__tp8",
+    "tep": _FAMILY / "glm_5_2_nvfp4__sm_103__tep",
+    "dep": _FAMILY / "glm_5_2_nvfp4__sm_103__dep",
+}
+
+
+def _module_file(module: str):
+    """The file of a dotted module of the family, or None outside it."""
+    if not (module == _TARGET or module.startswith(_TARGET + ".")):
+        return None
+    base = (
+        _FAMILY.joinpath(*module[len(_TARGET) :].lstrip(".").split("."))
+        if module != _TARGET
+        else _FAMILY
+    )
+    for candidate in (base.with_suffix(".py"), base / "__init__.py"):
+        if candidate.is_file():
+            return candidate
+    return None
+
+
+def _package_of(path: Path) -> str:
+    parts = path.relative_to(_FAMILY).with_suffix("").parts
+    if parts[-1] == "__init__":
+        parts = parts[:-1]
+    return (
+        ".".join((_TARGET, *parts[:-1]))
+        if path.name != "__init__.py"
+        else ".".join((_TARGET, *parts))
+    )
+
+
+def _imports(path: Path):
+    """Every family module ``path`` imports, at module level or inside a function."""
+    found = set()
+    for node in ast.walk(ast.parse(path.read_text())):
+        if isinstance(node, ast.ImportFrom):
+            if node.level:
+                package = _package_of(path).split(".")
+                base = ".".join(package[: len(package) - node.level + 1])
+                module = f"{base}.{node.module}" if node.module else base
+            else:
+                module = node.module or ""
+            targets = [module] + [f"{module}.{alias.name}" for alias in node.names]
+        elif isinstance(node, ast.Import):
+            targets = [alias.name for alias in node.names]
+        else:
+            continue
+        for name in targets:
+            file = _module_file(name)
+            if file is not None:
+                found.add(file)
+    return found
+
+
+def _reachable(start: Path):
+    """Every family file ``start`` imports, transitively, with the parent packages Python
+    imports first."""
+    seen, todo = set(), [start]
+    while todo:
+        path = todo.pop()
+        if path in seen:
+            continue
+        seen.add(path)
+        for parent in path.parents:
+            if parent == _FAMILY.parent:
+                break
+            if (parent / "__init__.py").is_file() and parent / "__init__.py" not in seen:
+                todo.append(parent / "__init__.py")
+        todo.extend(_imports(path))
+    return seen
+
+
+@pytest.mark.parametrize(
+    "target, other",
+    [(t, o) for t in _TARGET_DIRS for o in _TARGET_DIRS if t != o],
+)
+def test_a_target_imports_nothing_of_the_other(target, other):
+    """Read off disk: the import closure of a target's modeling module stays in the target's
+    own directory, the family's shared package and the routing module."""
+    reached = _reachable(_TARGET_DIRS[target] / "modeling.py")
+    assert not [p for p in reached if _TARGET_DIRS[other] in p.parents], sorted(
+        str(p.relative_to(_FAMILY)) for p in reached if _TARGET_DIRS[other] in p.parents
+    )
+    allowed = (_TARGET_DIRS[target], _FAMILY / "shared")
+    stray = [p for p in reached if p.parent != _FAMILY and not any(d in p.parents for d in allowed)]
+    assert not stray, sorted(str(p.relative_to(_FAMILY)) for p in stray)
+
+
+def test_the_shared_package_imports_no_target():
+    shared = _FAMILY / "shared"
+    for path in sorted(shared.rglob("*.py")):
+        for imported in _imports(path):
+            assert shared in imported.parents or imported.parent == _FAMILY, (path, imported)
+
+
+def test_the_tep_target_loads_no_kernel_package_of_the_tp8_target():
+    """``get_impl`` imports a path's package from the installed target's ``kernel_modules``;
+    the tep target names none, so every package it loads is the family's."""
+    source = (_TARGET_DIRS["tep"] / "modeling.py").read_text()
+    assert "kernel_modules" not in source
+    assert "glm_5_2_nvfp4__sm_103__tp8" not in "".join(
+        p.read_text() for p in _TARGET_DIRS["tep"].glob("*.py")
+    )
+
+
+def test_the_dep_target_loads_no_kernel_package_of_another_target():
+    """The dep target names no ``kernel_modules`` either, and no other target's directory."""
+    source = (_TARGET_DIRS["dep"] / "modeling.py").read_text()
+    assert "kernel_modules" not in source
+    text = "".join(p.read_text() for p in _TARGET_DIRS["dep"].glob("*.py"))
+    assert "glm_5_2_nvfp4__sm_103__tp8" not in text
+    assert "glm_5_2_nvfp4__sm_103__tep" not in text
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py
new file mode 100644
index 000000000000..0ad3228f6c8e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py
@@ -0,0 +1,700 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Where a GlmMoeDsaForCausalLM deployment lands, driven through ``modeling_v2_resolve``.
+
+Synthetic configs, as in test_modeling_v2_routing.py: routing reads the config
+shape and layer pattern, the declared quantization, the mapping, the compute
+capability (patched here, so these run on any machine) and the
+speculative-decoding config.
+
+The three targets claim one checkpoint (GLM-5.2 NVFP4 with its bf16 MTP layer), on
+sm103, with one-model MTP-Eagle as a linear chain, at disjoint parallel layouts:
+tp8 (attention and experts tensor-parallel over eight ranks) and tep (attention
+tensor-parallel over four or eight ranks, experts expert-parallel over the same
+ranks), both without attention DP and at five draft tokens, and dep (attention DP
+with the experts expert-parallel over the same ranks: sixteen ranks at three draft
+tokens, eight ranks at three, or eight ranks with the LM head tensor-parallel in
+attention DP at five). Every other value of each criterion declines -- ``auto``
+falls back to the built-in implementation and ``require`` raises naming the fact
+that missed -- including quantization declarations the engine's own parser would
+reject.
+"""
+
+from __future__ import annotations
+
+import re
+
+import pytest
+import torch
+from transformers import PretrainedConfig
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
+    ModelingV2Mode,
+    modeling_v2_resolve,
+    resolved_mode,
+)
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm.llmapi.llm_args import (
+    MTPDecodingConfig,
+    NGramDecodingConfig,
+    SAEnhancerConfig,
+    TorchLlmArgs,
+)
+from tensorrt_llm.mapping import Mapping
+
+pytestmark = pytest.mark.cpu_only
+
+_SM103 = (10, 3)
+_DUMMY_MODEL = "/tmp/dummy_model"
+_TARGET = "models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.modeling.ModelingV2Glm52Nvfp4Sm103Tp8"
+_TEP_TARGET = "models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tep.modeling.ModelingV2Glm52Nvfp4Sm103Tep"
+_DEP_TARGET = "models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep.modeling.ModelingV2Glm52Nvfp4Sm103Dep"
+
+# The excluded-module patterns of the GLM-5.2 NVFP4 export, abridged: the head, the
+# embedding, the dense layers 0-2, the MTP layer 78, and per MoE layer the attention
+# and the shared expert.
+_NVFP4_IGNORE = [
+    "lm_head",
+    "model.embed_tokens",
+    "model.layers.0*",
+    "model.layers.1.*",
+    "model.layers.2.*",
+    "model.layers.78*",
+    "model.layers.3.self_attn*",
+    "model.layers.3.mlp.shared_experts*",
+]
+
+
+def _qualified(target) -> str:
+    """``models.<family>.<dir>.modeling.<Class>`` for a resolved target class."""
+    return f"{target.__module__.split('.modeling_v2.')[1]}.{target.__name__}"
+
+
+@pytest.fixture(autouse=True)
+def _on_sm103(monkeypatch):
+    """Route as if this were an sm_103 device, wherever the test actually runs."""
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: _SM103)
+
+
+def _nvfp4_quantization(ignore=_NVFP4_IGNORE) -> dict:
+    """The ``quantization_config`` the NVFP4 export writes into its config.json."""
+    return {
+        "quant_method": "modelopt",
+        "quant_algo": "NVFP4",
+        "producer": {"name": "modelopt"},
+        "kv_cache_scheme": {"dynamic": False, "num_bits": 8, "type": "float"},
+        "ignore": list(ignore),
+    }
+
+
+# The shape and layer pattern GLM-5.2-NVFP4's own config.json declares.
+_GLM52_SHAPE = dict(
+    num_hidden_layers=78,
+    hidden_size=6144,
+    n_routed_experts=256,
+    q_lora_rank=2048,
+    num_attention_heads=64,
+    index_topk=2048,
+    first_k_dense_replace=3,
+    index_topk_freq=4,
+    index_skip_topk_offset=3,
+)
+_SHAPE_FIELDS = (*_GLM52_SHAPE, "index_topk_pattern")
+
+
+def _glm52_fields(**overrides) -> dict:
+    fields = dict(
+        architectures=["GlmMoeDsaForCausalLM"],
+        model_type="glm_moe_dsa",
+        num_nextn_predict_layers=1,
+        quantization_config=_nvfp4_quantization(),
+        **_GLM52_SHAPE,
+    )
+    fields.update(overrides)
+    return fields
+
+
+def _glm52_config(**overrides) -> PretrainedConfig:
+    """The shape, layer pattern and quantization GLM-5.2-NVFP4's own config.json declares."""
+    return PretrainedConfig(**_glm52_fields(**overrides))
+
+
+def _shape_words(**overrides) -> str:
+    """The config shape a NoTarget message names for the GLM-5.2 shape with ``overrides``."""
+    values = {**_GLM52_SHAPE, "index_topk_pattern": None, **overrides}
+    return re.escape(f"config shape {tuple(values[name] for name in _SHAPE_FIELDS)}")
+
+
+def _llm_args(mode: str = "auto") -> TorchLlmArgs:
+    return TorchLlmArgs(model=_DUMMY_MODEL, modeling_v2=mode)
+
+
+def _model_config(pretrained_config, spec_config, **mapping_kwargs) -> ModelConfig:
+    return ModelConfig(
+        pretrained_config=pretrained_config,
+        mapping=Mapping(**mapping_kwargs),
+        spec_config=spec_config,
+    )
+
+
+_TP8 = dict(world_size=8, tp_size=8)
+_TEP4 = dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1)
+_TEP8 = dict(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1)
+_DEP16 = dict(world_size=16, tp_size=16, moe_ep_size=16, moe_tp_size=1, enable_attention_dp=True)
+_DEP8 = dict(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1, enable_attention_dp=True)
+_DEP8_LM_HEAD_TP = dict(_DEP8, enable_lm_head_tp_in_adp=True)
+# The deployments the dep target claims: (mapping, draft tokens per request).
+_DEP_CELLS = {
+    "dep16-mtp3": (_DEP16, 3),
+    "dep8-mtp3": (_DEP8, 3),
+    "dep8-lm-head-tp-mtp5": (_DEP8_LM_HEAD_TP, 5),
+}
+
+
+def _mtp(max_draft_len: int = 5, **kwargs) -> MTPDecodingConfig:
+    return MTPDecodingConfig(max_draft_len=max_draft_len, **kwargs)
+
+
+@pytest.mark.parametrize("mode", [None, "off"], ids=["default", "off"])
+def test_off_resolves_nothing(mode):
+    args = TorchLlmArgs(model=_DUMMY_MODEL) if mode is None else _llm_args(mode)
+    config = _model_config(_glm52_config(), _mtp(), **_TP8)
+    assert modeling_v2_resolve(config, args) is None
+
+
+def test_tp8_is_one_expert_group_without_attention_dp():
+    """The default tp8 mapping is the layout the target was written for."""
+    mapping = Mapping(**_TP8)
+    assert (mapping.moe_tp_size, mapping.moe_ep_size, mapping.enable_attention_dp) == (
+        8,
+        1,
+        False,
+    )
+
+
+@pytest.mark.parametrize("mode", ["auto", "require"])
+def test_glm52_nvfp4_tp8_mtp5_matches(mode):
+    config = _model_config(_glm52_config(), _mtp(), **_TP8)
+    target = modeling_v2_resolve(config, _llm_args(mode))
+    assert _qualified(target) == _TARGET
+    assert resolved_mode(target) is ModelingV2Mode(mode)
+
+
+@pytest.mark.parametrize("mode", ["auto", "require"])
+@pytest.mark.parametrize("mapping_kwargs", [_TEP4, _TEP8], ids=["tep4", "tep8"])
+def test_glm52_nvfp4_tep_mtp5_matches(mode, mapping_kwargs):
+    """Four or eight ranks of tensor-parallel attention with the experts expert-parallel over
+    the same ranks land on the tep target."""
+    config = _model_config(_glm52_config(), _mtp(), **mapping_kwargs)
+    target = modeling_v2_resolve(config, _llm_args(mode))
+    assert _qualified(target) == _TEP_TARGET
+    assert resolved_mode(target) is ModelingV2Mode(mode)
+
+
+def test_the_tep_mappings_are_one_expert_per_rank_group_without_attention_dp():
+    for kwargs, ranks in ((_TEP4, 4), (_TEP8, 8)):
+        mapping = Mapping(**kwargs)
+        assert (mapping.world_size, mapping.tp_size, mapping.moe_ep_size, mapping.moe_tp_size) == (
+            ranks,
+            ranks,
+            ranks,
+            1,
+        )
+        assert not mapping.enable_attention_dp
+
+
+def test_the_three_claims_are_disjoint():
+    """No mapping satisfies two parallel criteria: tp8 takes moe_ep 1, tep moe_ep = tp > 1, both
+    without attention DP, and dep takes attention DP."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa import routing
+
+    for kwargs in (_TP8, _TEP4, _TEP8, _DEP16, _DEP8, _DEP8_LM_HEAD_TP):
+        mapping = Mapping(**kwargs)
+        assert routing._tp8(mapping) + routing._tep(mapping) + routing._dep(mapping) == 1, kwargs
+
+
+@pytest.mark.parametrize("mode", ["auto", "require"])
+@pytest.mark.parametrize(
+    "mapping_kwargs, draft_len", list(_DEP_CELLS.values()), ids=list(_DEP_CELLS)
+)
+def test_glm52_nvfp4_dep_matches(mode, mapping_kwargs, draft_len):
+    """Attention DP with the experts expert-parallel over the same ranks lands on the dep target
+    at the draft length of its rank count and LM-head split."""
+    config = _model_config(_glm52_config(), _mtp(draft_len), **mapping_kwargs)
+    target = modeling_v2_resolve(config, _llm_args(mode))
+    assert _qualified(target) == _DEP_TARGET
+    assert resolved_mode(target) is ModelingV2Mode(mode)
+
+
+def test_the_dep_mappings_are_one_expert_per_rank_group_with_attention_dp():
+    for kwargs, ranks, lm_head_tp in ((_DEP16, 16, False), (_DEP8, 8, False)) + (
+        (_DEP8_LM_HEAD_TP, 8, True),
+    ):
+        mapping = Mapping(**kwargs)
+        assert (mapping.world_size, mapping.tp_size, mapping.moe_ep_size, mapping.moe_tp_size) == (
+            ranks,
+            ranks,
+            ranks,
+            1,
+        )
+        assert mapping.enable_attention_dp
+        assert mapping.enable_lm_head_tp_in_adp is lm_head_tp
+
+
+def test_the_dep_cells_are_the_routed_ones():
+    """The target's cell table and ``route``'s claim name the same deployments."""
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa import routing
+    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep import (  # noqa: E501
+        modeling,
+    )
+
+    assert set(modeling.CELLS) == {
+        (ranks, lm_head_tp, draft) for (ranks, lm_head_tp), draft in routing._DEP.items()
+    }
+    assert {
+        (Mapping(**kwargs).world_size, Mapping(**kwargs).enable_lm_head_tp_in_adp, draft)
+        for kwargs, draft in _DEP_CELLS.values()
+    } == set(modeling.CELLS)
+
+
+def test_the_dep_target_is_the_builtin_causal_lm():
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+
+    config = _model_config(_glm52_config(), _mtp(3), **_DEP8)
+    target = modeling_v2_resolve(config, _llm_args("require"))
+    assert issubclass(target, GlmMoeDsaForCausalLM)
+    assert (
+        target.get_preferred_kv_cache_manager_version()
+        == GlmMoeDsaForCausalLM.get_preferred_kv_cache_manager_version()
+    )
+    assert (
+        target.get_preferred_transceiver_runtime()
+        == GlmMoeDsaForCausalLM.get_preferred_transceiver_runtime()
+    )
+
+
+def test_the_tep_target_is_the_builtin_causal_lm():
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+
+    config = _model_config(_glm52_config(), _mtp(), **_TEP8)
+    target = modeling_v2_resolve(config, _llm_args("require"))
+    assert issubclass(target, GlmMoeDsaForCausalLM)
+    assert (
+        target.get_preferred_kv_cache_manager_version()
+        == GlmMoeDsaForCausalLM.get_preferred_kv_cache_manager_version()
+    )
+
+
+def test_the_engine_config_class_keeps_the_routed_fields():
+    """The loader builds glm_moe_dsa configs through the registered class; route reads the same fields."""
+    from tensorrt_llm._torch.pyexecutor.config_utils import _CONFIG_REGISTRY
+
+    pretrained_config = _CONFIG_REGISTRY["glm_moe_dsa"].from_dict(_glm52_fields())
+    config = _model_config(pretrained_config, _mtp(), **_TP8)
+    assert _qualified(modeling_v2_resolve(config, _llm_args("require"))) == _TARGET
+
+
+def test_the_target_is_the_builtin_causal_lm():
+    """The target keeps the built-in class's engine preferences, read before routing."""
+    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
+
+    config = _model_config(_glm52_config(), _mtp(), **_TP8)
+    target = modeling_v2_resolve(config, _llm_args("require"))
+    assert issubclass(target, GlmMoeDsaForCausalLM)
+    assert (
+        target.get_preferred_transceiver_runtime()
+        == GlmMoeDsaForCausalLM.get_preferred_transceiver_runtime()
+    )
+    assert (
+        target.get_preferred_kv_cache_manager_version()
+        == GlmMoeDsaForCausalLM.get_preferred_kv_cache_manager_version()
+    )
+
+
+# Each case changes one fact of the claimed deployment: (config overrides, mapping,
+# speculative config, compute capability, what the NoTarget message must name).
+_FOREIGN = {
+    "tp4": ({}, dict(world_size=4, tp_size=4), _mtp(), _SM103, r"ws=4 tp=4"),
+    "tp16": ({}, dict(world_size=16, tp_size=16), _mtp(), _SM103, r"ws=16 tp=16"),
+    "attention-dp": (
+        {},
+        dict(_TP8, enable_attention_dp=True),
+        _mtp(),
+        _SM103,
+        r"attention_dp=True",
+    ),
+    "tep2": (
+        {},
+        dict(world_size=2, tp_size=2, moe_ep_size=2, moe_tp_size=1),
+        _mtp(),
+        _SM103,
+        r"ws=2 tp=2 .* moe_tp=1 moe_ep=2",
+    ),
+    "tep16": (
+        {},
+        dict(world_size=16, tp_size=16, moe_ep_size=16, moe_tp_size=1),
+        _mtp(),
+        _SM103,
+        r"ws=16 tp=16 .* moe_tp=1 moe_ep=16",
+    ),
+    "tep8-attention-dp": (
+        {},
+        dict(_TEP8, enable_attention_dp=True),
+        _mtp(),
+        _SM103,
+        r"moe_ep=8 attention_dp=True",
+    ),
+    "tep4-attention-dp": (
+        {},
+        dict(_TEP4, enable_attention_dp=True),
+        _mtp(),
+        _SM103,
+        r"moe_ep=4 attention_dp=True",
+    ),
+    "tp8-ep4-tp2": (
+        {},
+        dict(_TP8, moe_ep_size=4, moe_tp_size=2),
+        _mtp(),
+        _SM103,
+        r"moe_tp=2 moe_ep=4",
+    ),
+    "tep4-pp2": (
+        {},
+        dict(world_size=8, tp_size=4, pp_size=2, moe_ep_size=4, moe_tp_size=1),
+        _mtp(),
+        _SM103,
+        r"ws=8 tp=4 pp=2",
+    ),
+    "tep8-sm100": ({}, _TEP8, _mtp(), (10, 0), r"compute capability \(10, 0\)"),
+    "tep4-mtp3": ({}, _TEP4, _mtp(3), _SM103, r"max_draft_len=3 "),
+    "tep8-bf16-checkpoint": (
+        dict(quantization_config=None),
+        _TEP8,
+        _mtp(),
+        _SM103,
+        r"quantization none",
+    ),
+    "sm100": ({}, _TP8, _mtp(), (10, 0), r"compute capability \(10, 0\)"),
+    "sm120": ({}, _TP8, _mtp(), (12, 0), r"compute capability \(12, 0\)"),
+    "sm90": ({}, _TP8, _mtp(), (9, 0), r"compute capability \(9, 0\)"),
+    "bf16-checkpoint": (
+        dict(quantization_config=None),
+        _TP8,
+        _mtp(),
+        _SM103,
+        r"quantization none",
+    ),
+    "fp8-checkpoint": (
+        dict(quantization_config={"quant_method": "fp8", "weight_block_size": [128, 128]}),
+        _TP8,
+        _mtp(),
+        _SM103,
+        r"quantization quant_method 'fp8' \(not a modelopt export\)",
+    ),
+    "fp8-modelopt-checkpoint": (
+        dict(quantization_config={**_nvfp4_quantization(), "quant_algo": "FP8"}),
+        _TP8,
+        _mtp(),
+        _SM103,
+        r"quantization FP8 with the MTP layer unquantized",
+    ),
+    "quantized-mtp-layer": (
+        dict(quantization_config=_nvfp4_quantization(["lm_head", "model.embed_tokens"])),
+        _TP8,
+        _mtp(),
+        _SM103,
+        r"quantization NVFP4 with the MTP layer quantized",
+    ),
+    **{
+        f"shape-{name}": (
+            {name: value},
+            _TP8,
+            _mtp(),
+            _SM103,
+            _shape_words(**{name: value}),
+        )
+        for name, value in (
+            ("num_hidden_layers", 92),
+            ("hidden_size", 7168),
+            ("n_routed_experts", 160),
+            ("q_lora_rank", 1536),
+            ("num_attention_heads", 128),
+            ("index_topk", 1024),
+            ("first_k_dense_replace", 1),
+            ("index_topk_freq", 1),
+            ("index_skip_topk_offset", 2),
+            ("index_topk_pattern", "FFFSSSF"),
+        )
+    },
+    "no-speculation": ({}, _TP8, None, _SM103, r"speculative decoding off"),
+    "mtp1": ({}, _TP8, _mtp(1), _SM103, r"max_draft_len=1 "),
+    "mtp2": ({}, _TP8, _mtp(2), _SM103, r"max_draft_len=2 "),
+    "mtp3": ({}, _TP8, _mtp(3), _SM103, r"max_draft_len=3 "),
+    "mtp4": ({}, _TP8, _mtp(4), _SM103, r"max_draft_len=4 "),
+    "mtp6": ({}, _TP8, _mtp(6), _SM103, r"max_draft_len=6 "),
+    "vanilla-mtp": ({}, _TP8, _mtp(use_mtp_vanilla=True), _SM103, r"speculative decoding MTP "),
+    "dynamic-tree": (
+        {},
+        _TP8,
+        _mtp(use_dynamic_tree=True, dynamic_tree_max_topK=1),
+        _SM103,
+        r"dynamic tree",
+    ),
+    "suffix-automaton-enhancer": (
+        {},
+        _TP8,
+        _mtp(sa_config=SAEnhancerConfig()),
+        _SM103,
+        r"with the suffix-automaton enhancer",
+    ),
+    "relaxed-acceptance": (
+        {},
+        _TP8,
+        _mtp(use_relaxed_acceptance_for_thinking=True),
+        _SM103,
+        r"with relaxed acceptance",
+    ),
+    "ngram": (
+        {},
+        _TP8,
+        NGramDecodingConfig(max_draft_len=5, max_matching_ngram_size=2),
+        _SM103,
+        r"speculative decoding NGRAM",
+    ),
+    # The dep claim: one fact of a claimed deployment changed.
+    "dep4-mtp3": (
+        {},
+        dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True),
+        _mtp(3),
+        _SM103,
+        r"ws=4 tp=4 .* attention_dp=True",
+    ),
+    "dep32-mtp3": (
+        {},
+        dict(world_size=32, tp_size=32, moe_ep_size=32, moe_tp_size=1, enable_attention_dp=True),
+        _mtp(3),
+        _SM103,
+        r"ws=32 tp=32",
+    ),
+    "dep16-mtp5": ({}, _DEP16, _mtp(), _SM103, r"ws=16 .* max_draft_len=5 "),
+    "dep16-lm-head-tp-mtp3": (
+        {},
+        dict(_DEP16, enable_lm_head_tp_in_adp=True),
+        _mtp(3),
+        _SM103,
+        r"ws=16 .* lm_head_tp_in_adp=True, .* max_draft_len=3 ",
+    ),
+    "dep16-lm-head-tp-mtp5": (
+        {},
+        dict(_DEP16, enable_lm_head_tp_in_adp=True),
+        _mtp(),
+        _SM103,
+        r"ws=16 .* lm_head_tp_in_adp=True, .* max_draft_len=5 ",
+    ),
+    "dep8-lm-head-tp-mtp3": (
+        {},
+        _DEP8_LM_HEAD_TP,
+        _mtp(3),
+        _SM103,
+        r"lm_head_tp_in_adp=True, .* max_draft_len=3 ",
+    ),
+    **{
+        f"dep8-mtp{draft}": ({}, _DEP8, _mtp(draft), _SM103, rf"max_draft_len={draft} ")
+        for draft in (1, 2, 4, 5)
+    },
+    **{
+        f"dep8-lm-head-tp-mtp{draft}": (
+            {},
+            _DEP8_LM_HEAD_TP,
+            _mtp(draft),
+            _SM103,
+            rf"max_draft_len={draft} ",
+        )
+        for draft in (1, 2, 4, 6)
+    },
+    "dep8-moe-tp2": (
+        {},
+        dict(_DEP8, moe_ep_size=4, moe_tp_size=2),
+        _mtp(3),
+        _SM103,
+        r"moe_tp=2 moe_ep=4 attention_dp=True",
+    ),
+    "dep8-pp2": (
+        {},
+        dict(
+            world_size=16,
+            tp_size=8,
+            pp_size=2,
+            moe_ep_size=8,
+            moe_tp_size=1,
+            enable_attention_dp=True,
+        ),
+        _mtp(3),
+        _SM103,
+        r"ws=16 tp=8 pp=2",
+    ),
+    "dep8-no-speculation": ({}, _DEP8, None, _SM103, r"speculative decoding off"),
+    "dep8-vanilla-mtp": (
+        {},
+        _DEP8,
+        _mtp(3, use_mtp_vanilla=True),
+        _SM103,
+        r"speculative decoding MTP ",
+    ),
+    "dep8-dynamic-tree": (
+        {},
+        _DEP8,
+        _mtp(3, use_dynamic_tree=True, dynamic_tree_max_topK=1),
+        _SM103,
+        r"dynamic tree",
+    ),
+    "dep8-suffix-automaton-enhancer": (
+        {},
+        _DEP8,
+        _mtp(3, sa_config=SAEnhancerConfig()),
+        _SM103,
+        r"with the suffix-automaton enhancer",
+    ),
+    "dep8-relaxed-acceptance": (
+        {},
+        _DEP8,
+        _mtp(3, use_relaxed_acceptance_for_thinking=True),
+        _SM103,
+        r"with relaxed acceptance",
+    ),
+    "dep8-sm100": ({}, _DEP8, _mtp(3), (10, 0), r"compute capability \(10, 0\)"),
+    "dep16-bf16-checkpoint": (
+        dict(quantization_config=None),
+        _DEP16,
+        _mtp(3),
+        _SM103,
+        r"quantization none",
+    ),
+    "dep8-lm-head-tp-quantized-mtp-layer": (
+        dict(quantization_config=_nvfp4_quantization(["lm_head", "model.embed_tokens"])),
+        _DEP8_LM_HEAD_TP,
+        _mtp(),
+        _SM103,
+        r"quantization NVFP4 with the MTP layer quantized",
+    ),
+    "dep8-shape-num_attention_heads": (
+        {"num_attention_heads": 128},
+        _DEP8,
+        _mtp(3),
+        _SM103,
+        _shape_words(num_attention_heads=128),
+    ),
+}
+
+
+@pytest.mark.parametrize(
+    "config_kwargs, mapping_kwargs, spec_config, capability, missed",
+    list(_FOREIGN.values()),
+    ids=list(_FOREIGN),
+)
+def test_a_foreign_deployment_does_not_match(
+    monkeypatch, config_kwargs, mapping_kwargs, spec_config, capability, missed
+):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
+    config = _model_config(_glm52_config(**config_kwargs), spec_config, **mapping_kwargs)
+
+    assert modeling_v2_resolve(config, _llm_args("auto")) is None
+    with pytest.raises(ValueError, match="GlmMoeDsaForCausalLM") as excinfo:
+        modeling_v2_resolve(config, _llm_args("require"))
+    message = str(excinfo.value)
+    assert "the glm_moe_dsa targets are" in message
+    assert "tp8 (8 ranks" in message and "tep (4 or 8 ranks" in message
+    assert "dep (attention DP" in message
+    assert re.search(missed, message), message
+
+
+def test_a_config_without_the_sparse_attention_fields_does_not_match():
+    """A GlmMoeDsa config that lacks a shape field declines instead of failing in route."""
+    pretrained_config = _glm52_config()
+    del pretrained_config.index_topk
+    config = _model_config(pretrained_config, _mtp(), **_TP8)
+
+    assert modeling_v2_resolve(config, _llm_args("auto")) is None
+    with pytest.raises(ValueError, match=_shape_words(index_topk=None)):
+        modeling_v2_resolve(config, _llm_args("require"))
+
+
+class _NoToDict:
+    """A quantization declaration that is neither a dict nor convertible to one."""
+
+    quant_method = "nvfp4"
+
+
+# Quantization declarations route must decline without parsing through the engine's
+# builders, which raise other errors than NoTarget on them: (declaration, what the
+# message names when the rest of the deployment matches).
+_UNREADABLE_QUANTIZATION = {
+    "modelopt-mixed-precision-flat": (
+        {
+            "quant_method": "modelopt",
+            "quant_algo": "MIXED_PRECISION",
+            "config_groups": {"group_0": {}},
+            "ignore": ["lm_head"],
+        },
+        r"quantization MIXED_PRECISION with the MTP layer quantized",
+    ),
+    "modelopt-mixed-precision-legacy": (
+        {"producer": {"name": "modelopt"}, "quantization": {"quant_algo": "MIXED_PRECISION"}},
+        r"quantization MIXED_PRECISION with the MTP layer quantized",
+    ),
+    "modelopt-unknown-algorithm": (
+        {"quant_method": "modelopt", "quant_algo": "NOT_AN_ALGORITHM"},
+        r"quantization a modelopt config without a readable global algorithm",
+    ),
+    "modelopt-no-algorithm": (
+        {"quant_method": "modelopt", "ignore": ["model.layers.78*"]},
+        r"quantization a modelopt config without a readable global algorithm",
+    ),
+    "modelopt-excluded-modules-not-a-list": (
+        {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": 5},
+        r"quantization modelopt NVFP4 with unreadable excluded modules",
+    ),
+    "fp8-64x64-blocks": (
+        {"quant_method": "fp8", "weight_block_size": [64, 64]},
+        r"quantization quant_method 'fp8'",
+    ),
+    "nvfp4-group-32": (
+        {"quant_method": "nvfp4", "group_size": 32},
+        r"quantization quant_method 'nvfp4'",
+    ),
+    "mxfp8-1x64-blocks": (
+        {"quant_method": "mxfp8", "weight_block_size": [1, 64]},
+        r"quantization quant_method 'mxfp8'",
+    ),
+    "compressed-tensors-without-strategy": (
+        {
+            "quant_method": "compressed-tensors",
+            "config_groups": {"group_0": {"weights": {"num_bits": 4, "type": "int"}}},
+        },
+        r"quantization quant_method 'compressed-tensors'",
+    ),
+    "not-a-dict": (_NoToDict(), r"quantization declared as a _NoToDict, not a dict"),
+}
+
+
+@pytest.mark.parametrize("capability", [_SM103, (9, 0)], ids=["sm103", "sm90"])
+@pytest.mark.parametrize(
+    "quantization, named",
+    list(_UNREADABLE_QUANTIZATION.values()),
+    ids=list(_UNREADABLE_QUANTIZATION),
+)
+def test_a_quantization_route_cannot_read_declines(monkeypatch, quantization, named, capability):
+    """Declined, never an exception other than NoTarget, so ``auto`` still falls back.
+
+    The quantization is read only for the deployment the target serves; on any
+    other deployment the message says it was not read.
+    """
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
+    config = _model_config(_glm52_config(quantization_config=quantization), _mtp(), **_TP8)
+
+    assert modeling_v2_resolve(config, _llm_args("auto")) is None
+    with pytest.raises(ValueError, match="GlmMoeDsaForCausalLM") as excinfo:
+        modeling_v2_resolve(config, _llm_args("require"))
+    expected = named if capability == _SM103 else r"quantization not read"
+    assert re.search(expected, str(excinfo.value)), str(excinfo.value)
diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py
index 52781ca6a355..97524438dee0 100644
--- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py
+++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py
@@ -129,6 +129,11 @@ def test_targets_do_not_share_files():
     Targets are allowed to import the catalog and nothing else of each
     other's. A shared helper between two targets is the first step back to
     the abstraction this package exists to avoid.
+
+    One declared exception: a target may import its own family's ``shared``
+    package (``models/<family>/shared/``), the code two or more targets of one
+    family call that cannot be a catalog entry (README, "Op ownership"). It is
+    never another target's directory.
     """
     for arch in _ARCHS:
         for target_dir, name in _routed_targets(arch):
@@ -138,6 +143,8 @@ def test_targets_do_not_share_files():
                     dots, tail = match.group(1), match.group(2)
                     if len(dots) == 1:
                         continue  # sibling within the target
+                    if len(dots) == 2 and tail.split(".")[0] == "shared":
+                        continue  # the family's shared package
                     assert tail.startswith("catalog"), (
                         f"{name}: {source_file.name} reaches outside its own "
                         f"directory for {tail!r}; targets may import the "

From a87e238391622bcc24abd6fc6478315504d314ca Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 22:28:11 -0700
Subject: [PATCH 05/19] [None][doc] modeling_v2: list the GLM-5.2 targets

Add the glm_moe_dsa family and its three targets to the modeling_v2 README.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_experimental/modeling_v2/README.md       | 31 ++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
index 481c86552981..8681fe5471c4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
@@ -110,9 +110,30 @@ models/<family>/
   <checkpoint>__<gpu arch>__<parallel>/
                       modeling.py  weights.py
                       ops/<category>/   entries only this target calls
+  shared/           code two or more targets of the family call that is not an entry
 catalog/<category>/ entries more than one target calls
 ```
 
+The targets, as `models/<family>/<target>`:
+
+* `deepseek_v3/r1_0528_nvfp4__sm_103__dep4`
+* `glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8`: GLM-5.2 NVFP4 low-latency decode with attention
+  and experts tensor-parallel over eight ranks. The fast paths shared by the GLM-5.2 targets are
+  under `glm_moe_dsa/shared/kernels/`, the routed-MoE ones only this target calls under its own
+  `kernels/`, with their C++ kernels and ops in `cpp/tensorrt_llm/kernels/glmMoeDsa` and
+  `cpp/tensorrt_llm/thop/glmMoeDsa`; `shared/dispatch.py` selects them and lists the knobs.
+  The tp8 and tep targets run the layer composition of `shared/layers.py`, each with its own plan.
+* `glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep`: the same checkpoint with attention tensor-parallel
+  over four or eight ranks and the experts expert-parallel over the same ranks. It runs the
+  shared fast paths that take its tensor-parallel size, in front of the built-in
+  expert-parallel MoE.
+* `glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep`: the same checkpoint with attention data parallelism
+  and the experts expert-parallel over the same ranks. Its built-in layers run the shared
+  attention, draft-loop and step-loop fast paths that are rank-local under attention DP; it
+  checks the KV-cache dtype and the MoE backend, which `route` does not read, when it is built.
+* `gpt_oss/gpt_oss_120b__sm_103__tp1`
+* `qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1`
+
 An entry, shared or owned, is one `OpWrapper` subclass: `__call__` is the op
 call a target makes, `reference` is what it is supposed to compute in plain
 torch, `is_valid` names the inputs it accepts and answers wrongly, and `CELLS`
@@ -281,6 +302,14 @@ names beside the code that reads it, only adds something to forget to update.
 An entry only one target calls lives under that target's `ops/`; one that two
 or more call lives in `catalog/`. A later target that needs an owned entry
 promotes it by moving the file and its test to the shared tree -- the test is
-the certification record, so nothing else has to travel with it. The five
+the certification record, so nothing else has to travel with it.
+
+Code that two or more targets of one family call and that cannot be an entry --
+it holds state across calls (collective workspaces, compiled kernels, a plan the
+kernels read) or registers torch ops when imported, so one copy per process is
+required -- lives in that family's `models/<family>/shared/`, moved there with its
+tests the same way. A target may import its own family's `shared` package and
+the catalog, nothing else (`test_modeling_v2_claims.py`); `shared` imports no
+target. The five
 other out-of-tree targets and the agent definitions are not migrated in this
 batch.

From c1a4f5456cd1a0f515cd0db079823f9a2883933d Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 22:46:16 -0700
Subject: [PATCH 06/19] [None][doc] modeling_v2: note that the GLM-5.2 TEP and
 DEP targets reuse the TP8 fast paths

Each of the two targets gets a short README: the fast paths were designed
and tuned for the TP8 target, these targets run the subset that applies at
their topology, and both should be revisited with fast paths of their own.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../glm_5_2_nvfp4__sm_103__dep/README.md          | 15 +++++++++++++++
 .../glm_5_2_nvfp4__sm_103__tep/README.md          | 14 ++++++++++++++
 2 files changed, 29 insertions(+)
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md
new file mode 100644
index 000000000000..6a8376b48935
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md
@@ -0,0 +1,15 @@
+# `glm_5_2_nvfp4__sm_103__dep`
+
+GLM-5.2 NVFP4 with attention data parallelism over eight or sixteen ranks and the experts
+expert-parallel over the same ranks.
+
+**This target is not separately optimized.** The family's fast paths were designed and tuned
+for low-latency decode on the TP8 target, `glm_5_2_nvfp4__sm_103__tp8`. Under attention data
+parallelism only the rank-local ones apply, and this target runs those from `../shared/`: the
+RoPE and KV append, the full-indexer layers' query front, the draft loop, the decode step's
+prologue and the host page retention policy. Its decoder layers, MoE modules and draft layer
+stay the built-in classes. No fast path was designed or tuned for DEP batch shapes.
+
+It is in the family because those rank-local paths also improve decode at this topology. DEP
+should be revisited in the future with fast paths of its own, in particular for the
+attention-DP communication and the expert-parallel MoE.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md
new file mode 100644
index 000000000000..eefb1402ee8d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md
@@ -0,0 +1,14 @@
+# `glm_5_2_nvfp4__sm_103__tep`
+
+GLM-5.2 NVFP4 with the attention heads tensor-parallel over four or eight ranks and the
+experts expert-parallel over the same ranks.
+
+**This target is not separately optimized.** The family's fast paths were designed and tuned
+for low-latency decode on the TP8 target, `glm_5_2_nvfp4__sm_103__tp8`. This target runs the
+fast paths in `../shared/` whose contracts also cover its tensor-parallel size, in front of the
+built-in expert-parallel MoE. No fast path was designed or tuned for TEP batch shapes, and the
+routed experts run the built-in MoE.
+
+It is in the family because those shared paths also improve decode at this topology. TEP should
+be revisited in the future with fast paths of its own, in particular for the expert-parallel
+MoE and the larger decode batches this topology serves.

From 4da9ca9dd3806fc531324f7a5fabbf45d4d33680 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Sat, 17 Oct 2026 22:54:22 -0700
Subject: [PATCH 07/19] [None][chore] modeling_v2 GLM-5.2: pre-commit fixes

clang-format on five C++ files and one test .cu (whitespace only), an
unused local and a blank line flagged by ruff, two CuTe tensor names that
codespell reads as misspellings, and the prep packet's pinned slots
allocated with prefer_pinned(), as the pinned-memory policy requires. The
stager declines when prefer_pinned() is False, so a packet is only built
when it is True.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../kernels/glmMoeDsa/denseMlp.cu             | 20 +++++++++----------
 .../kernels/glmMoeDsa/ropeKvAppend.cu         |  7 +++----
 .../kernels/glmMoeDsa/ropeKvAppend.h          |  2 +-
 .../thop/glmMoeDsa/denseMlpOp.cpp             |  3 +--
 .../thop/glmMoeDsa/draftTailOp.cpp            |  4 ++--
 .../kernels/deferred_input_prep/compat.py     |  1 -
 .../kernels/deferred_input_prep/packet.py     | 11 ++++++++--
 .../shared/kernels/draft_front/kernels.py     | 12 +++++------
 .../proj_kernel.py                            |  1 -
 .../shared/kernels/dense_mlp/pdl_delay.cu     |  3 ++-
 10 files changed, 34 insertions(+), 30 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu
index d2fe7ad69100..7837384be9b1 100644
--- a/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu
@@ -188,8 +188,8 @@ __device__ __forceinline__ void computeChunk(float (&accA)[4], float (&accB)[4],
 //! CTA-reduced accumulator fragment in warp 0 (other warps return false).
 //! acc[0..1] = tile row "group", activation rows 2 * tig + {0, 1}; acc[2..3] = tile row "group + 8".
 template <typename T>
-__device__ __forceinline__ bool streamTile(float (&acc)[4], Bf16 const* __restrict__ x,
-    Bf16 const* __restrict__ weight, int numRows, int kDim, int row0, int row1)
+__device__ __forceinline__ bool streamTile(float (&acc)[4], Bf16 const* __restrict__ x, Bf16 const* __restrict__ weight,
+    int numRows, int kDim, int row0, int row1)
 {
     constexpr int kWarps = T::kWarps;
     constexpr int kSteps = T::kSteps;
@@ -286,8 +286,8 @@ __device__ __forceinline__ float roundToBf16(float v)
 }
 
 template <typename T>
-__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) gateUpSwigluKernel(
-    Bf16 const* __restrict__ x, Bf16 const* __restrict__ weight, Bf16* __restrict__ h, int numRows, int hidden, int inter)
+__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) gateUpSwigluKernel(Bf16 const* __restrict__ x,
+    Bf16 const* __restrict__ weight, Bf16* __restrict__ h, int numRows, int hidden, int inter)
 {
     int const lane = threadIdx.x % kWarpSize;
     int const group = lane / kQuadThreads;
@@ -315,8 +315,8 @@ __global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) gateUpSwi
 }
 
 template <typename T>
-__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) downKernel(
-    Bf16 const* __restrict__ h, Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int inter, int hidden)
+__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) downKernel(Bf16 const* __restrict__ h,
+    Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int inter, int hidden)
 {
     int const lane = threadIdx.x % kWarpSize;
     int const group = lane / kQuadThreads;
@@ -381,16 +381,16 @@ bool isContractionSupported(int kDim, int variant)
 
 //! Dispatches a variant index to its kernel instance; kGateUp selects the gate/up or the down kernel.
 template <bool kGateUp>
-cudaError_t launchVariant(int variant, Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int kDim,
-    int outDim, bool enablePdl, cudaStream_t stream)
+cudaError_t launchVariant(int variant, Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int kDim, int outDim,
+    bool enablePdl, cudaStream_t stream)
 {
     auto launch = [&](auto tiling) -> cudaError_t
     {
         using T = decltype(tiling);
         if constexpr (kGateUp)
         {
-            return launchKernel<T>(gateUpSwigluKernel<T>, outDim / kGateUpCols, in, weight, out, numRows, kDim, outDim,
-                enablePdl, stream);
+            return launchKernel<T>(
+                gateUpSwigluKernel<T>, outDim / kGateUpCols, in, weight, out, numRows, kDim, outDim, enablePdl, stream);
         }
         else
         {
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu
index efe8a1e21e09..e13d352ba57a 100644
--- a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu
@@ -198,7 +198,7 @@ __global__ void __launch_bounds__(kRopeThreads) ropeAppendScatterKernel(RopeAppe
     {
         scale = p.quantScaleQ ? p.quantScaleQ[0] : 1.F;
     }
-    auto loadCoef = [&](int32_t pos, float2 (&c)[4])
+    auto loadCoef = [&](int32_t pos, float2(&c)[4])
     {
         float2 const* src = p.cosSin + static_cast<size_t>(kRopeDim) * pos + (vec * kEltsPerVec) / 2;
 #pragma unroll
@@ -415,9 +415,8 @@ __global__ void __launch_bounds__(kBridgeThreads) topkBridgeKernel(TopkBridgePar
         }
         else
         {
-            int64_t const acc = p.numAcceptedIs64
-                ? static_cast<int64_t const*>(p.numAccepted)[i]
-                : static_cast<int64_t>(static_cast<int32_t const*>(p.numAccepted)[i]);
+            int64_t const acc = p.numAcceptedIs64 ? static_cast<int64_t const*>(p.numAccepted)[i]
+                                                  : static_cast<int64_t>(static_cast<int32_t const*>(p.numAccepted)[i]);
             int64_t off = acc - 1;
             off = off < 0 ? 0 : (off > p.nextN - 1 ? p.nextN - 1 : off);
             srcRow = static_cast<int32_t>(static_cast<int64_t>(i) * p.nextN + off);
diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h
index 8f12c8b6cd61..8d4ea666f26b 100644
--- a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h
+++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h
@@ -95,7 +95,7 @@ struct TopkBridgeParams
 {
     int32_t const* topk{nullptr}; //!< Source rows [srcRows, numTopk] with row stride topkStride0, unit column stride.
     int64_t topkStride0{0};
-    int32_t numRows{0}; //!< Rows remapped to global indices.
+    int32_t numRows{0};           //!< Rows remapped to global indices.
     int32_t numTopk{0};
 
     int32_t const* reqIdx{nullptr}; //!< [numRows] request index per row, before subtracting reqIdxBias.
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp
index dd335256942d..c4843dcb6e44 100644
--- a/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp
@@ -55,8 +55,7 @@ void launchGateUp(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tenso
         static_cast<int>(variant), enablePdl, stream));
 }
 
-void launchDownProj(
-    at::Tensor const& h, at::Tensor const& downWeight, at::Tensor& out, int64_t variant, bool enablePdl)
+void launchDownProj(at::Tensor const& h, at::Tensor const& downWeight, at::Tensor& out, int64_t variant, bool enablePdl)
 {
     auto const stream = at::cuda::getCurrentCUDAStream(h.get_device());
     TLLM_CUDA_CHECK(tk::invokeDown(reinterpret_cast<__nv_bfloat16 const*>(h.data_ptr()),
diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp
index 62c53823f6ba..72cda4e7ab47 100644
--- a/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp
+++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp
@@ -116,8 +116,8 @@ void draftTail(at::Tensor const& hidden, at::Tensor const& weight, at::Tensor& t
     if (logitsOut.has_value())
     {
         TORCH_CHECK(logitsOut->device() == hidden.device() && logitsOut->scalar_type() == at::kBFloat16
-                && logitsOut->dim() == 2 && logitsOut->size(0) >= hidden.size(0)
-                && logitsOut->size(1) == weight.size(0) && logitsOut->stride(1) == 1,
+                && logitsOut->dim() == 2 && logitsOut->size(0) >= hidden.size(0) && logitsOut->size(1) == weight.size(0)
+                && logitsOut->stride(1) == 1,
             "glm_moe_dsa_draft_tail: logits_out must be a BF16 [rows, vocab_local] tensor with unit last stride");
         p.logitsOut = reinterpret_cast<__nv_bfloat16*>(logitsOut->data_ptr());
         p.logitsOutStride = logitsOut->stride(0);
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py
index d0ae536316f7..be05da8ad710 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/compat.py
@@ -42,7 +42,6 @@
 
 from typing import Optional
 
-
 # Executor-loop domain of this process: declared by the overlap loop at its start.
 _LOOP_STATE = {"declared": False, "reason": "", "logged": False}
 _LOOP_UNDECLARED = (
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py
index cd322a249581..f0c35b3e6276 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/packet.py
@@ -49,6 +49,8 @@
 import triton
 import triton.language as tl
 
+from tensorrt_llm._utils import prefer_pinned
+
 KIND_COPY = 0
 KIND_GATHER = 1
 KIND_BLOCK_OFFSETS = 2
@@ -143,9 +145,14 @@ def query() -> bool:
 class _PinnedSlot:
     """One packet's pinned storage and the event that marks its consumption."""
 
+    # The commit kernel reads this storage in place, so it must be pinned. The stager
+    # declines when prefer_pinned() is False, so a packet is only built when it is True.
+
     def __init__(self, max_desc: int, payload_words: int):
-        self.desc = torch.zeros((max_desc, DESC_WORDS), dtype=torch.int64, pin_memory=True)
-        self.payload = torch.zeros((payload_words,), dtype=torch.int32, pin_memory=True)
+        self.desc = torch.zeros(
+            (max_desc, DESC_WORDS), dtype=torch.int64, pin_memory=prefer_pinned()
+        )
+        self.payload = torch.zeros((payload_words,), dtype=torch.int32, pin_memory=prefer_pinned())
         self.desc_np = self.desc.numpy()
         self.payload_np = self.payload.numpy()
         self.event: Optional[torch.cuda.Event] = None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py
index 85077c66c76e..03cd2543f843 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/kernels.py
@@ -187,8 +187,8 @@ def __call__(
         mHs: cute.Tensor,
         mWe: cute.Tensor,
         mWh: cute.Tensor,
-        mYe: cute.Tensor,
-        mYh: cute.Tensor,
+        mOutE: cute.Tensor,
+        mOutH: cute.Tensor,
         M: Int64,
         eps_e: Float32,
         eps_h: Float32,
@@ -196,7 +196,7 @@ def __call__(
     ):
         tv_layout = self._tv_layout()
         tiler_mn = (self.rows_per_block, self.cols_per_tile)
-        self.kernel(mE, mHs, mWe, mWh, mYe, mYh, M, eps_e, eps_h, tv_layout, tiler_mn).launch(
+        self.kernel(mE, mHs, mWe, mWh, mOutE, mOutH, M, eps_e, eps_h, tv_layout, tiler_mn).launch(
             grid=[cute.ceil_div(M, self.rows_per_block), len(self.segments), 1],
             block=[self.num_threads, 1, 1],
             smem=self.smem_bytes(),
@@ -211,8 +211,8 @@ def kernel(
         mHs: cute.Tensor,
         mWe: cute.Tensor,
         mWh: cute.Tensor,
-        mYe: cute.Tensor,
-        mYh: cute.Tensor,
+        mOutE: cute.Tensor,
+        mOutH: cute.Tensor,
         M: Int64,
         eps_e: Float32,
         eps_h: Float32,
@@ -233,7 +233,7 @@ def kernel(
         )
 
         # Source 0 is the embedding half of the concatenation, source 1 the hidden half.
-        operands = ((mE, mWe, mYe, eps_e), (mHs, mWh, mYh, eps_h))
+        operands = ((mE, mWe, mOutE, eps_e), (mHs, mWh, mOutH, eps_h))
         seg0 = self.segments[0]
         if cutlass.const_expr(len(self.segments) == 1):
             self._normalize_segment(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py
index 9fe9cf2ea82c..361a72920252 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/proj_kernel.py
@@ -669,7 +669,6 @@ def kernel(
         done (last MMA warp), 6 main loop done (MMA warp 0), 4 rows stored,
         5 exit."""
         k = self.k
-        stride = self.stride
         tid, _, _ = cute.arch.thread_idx()
         bid, _, _ = cute.arch.block_idx()
         warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
index 19530f200648..4bf176903e38 100644
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
+++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
@@ -27,7 +27,8 @@
 namespace
 {
 
-__global__ void delayedCopyKernel(uint4 const* __restrict__ src, uint4* __restrict__ dst, int64_t numVec, int64_t delayNs)
+__global__ void delayedCopyKernel(
+    uint4 const* __restrict__ src, uint4* __restrict__ dst, int64_t numVec, int64_t delayNs)
 {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
     asm volatile("griddepcontrol.wait;" ::: "memory");

From d79d437e94ae285b503de03d7e7d0f035751d659 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Mon, 19 Oct 2026 14:19:08 -0700
Subject: [PATCH 08/19] [None][test] modeling_v2 GLM-5.2: replace the family's
 unit tests with end-to-end GSM8K gates

Remove the glm_moe_dsa unit and module tests, and gate each GLM-5.2 target
end to end instead: tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
serves GSM8K with modeling_v2="require" at tp8, tep4 and dep8. The tep4 case
runs in l0_gb300_multi_gpus, the tp8 and dep8 cases in
l0_gb300_multi_nodes_node2_gpu8. The engine-side runtime-hook tests stay.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_experimental/modeling_v2/README.md       |    2 +
 .../accuracy/modeling_v2/test_glm_moe_dsa.py  |  118 ++
 .../test-db/l0_gb300_multi_gpus.yml           |    2 +
 .../l0_gb300_multi_nodes_node2_gpu8.yml       |    4 +
 .../test_modeling_v2_glm_dep_plan.py          |  411 ------
 .../test_modeling_v2_glm_dep_target.py        |  501 -------
 .../test_modeling_v2_glm_tep_layers.py        |  328 -----
 .../test_modeling_v2_glm_tep_plan.py          |  355 -----
 .../test_expert_prefetch.py                   |  267 ----
 ...expert_prefetch_composition_hop_routing.py |  353 -----
 ...osition_shared_prefetch_expert_prefetch.py |  352 -----
 ...eferred_norm_routed_into_hop_wiring_mpi.py |  343 -----
 .../draft_live_rows/test_live_rows_oracle.py  |  270 ----
 .../draft_routed_moe/draft_moe_fixtures.py    |  315 -----
 .../kernels/draft_routed_moe/draft_moe_jit.py |  101 --
 .../test_draft_moe_adapter.py                 |  152 ---
 .../draft_routed_moe/test_draft_moe_oracle.py |  775 -----------
 .../kernels/hop_routing/hop_routing_jit.py    |  161 ---
 .../kernels/hop_routing/hop_routing_oracle.py |  214 ---
 .../hop_routing/test_hop_routing_adapter.py   |  245 ----
 .../test_hop_routing_hop_single_gpu.py        |  234 ----
 .../test_hop_routing_routed_moe_fork.py       |  200 ---
 .../test_attn_hop_integration_contract.py     |  161 ---
 .../routed_moe/routed_into_hop_checks.py      |  492 -------
 .../routed_moe/routed_moe_chain_bench.py      |  764 -----------
 .../kernels/routed_moe/routed_moe_fixtures.py | 1036 --------------
 .../routed_moe/test_routed_into_hop_mpi.py    |  464 -------
 .../routed_moe/test_routed_moe_integration.py |  400 ------
 .../test_routed_moe_native_consumer_mpi.py    |  176 ---
 .../routed_moe/test_routed_moe_routed_moe.py  |  722 ----------
 .../test_routed_moe_spec_prefetch.py          |  184 ---
 .../routed_moe/test_routed_moe_stress.py      |  812 -----------
 ...test_router_gemm_integration_invariants.py |  135 --
 .../shared_prefetch_fixtures.py               |  104 --
 .../shared_prefetch_se.py                     |  251 ----
 .../shared_prefetch_se_kernel.py              |  453 -------
 .../shared_prefetch_se_ptx.py                 |  120 --
 .../test_shared_prefetch_adapter.py           |  251 ----
 .../test_shared_prefetch_routed.py            |  454 -------
 .../test_shared_prefetch_shared.py            |  269 ----
 .../test_modeling_v2_glm_adopt.py             |  152 ---
 .../test_modeling_v2_glm_layers.py            |  888 ------------
 .../test_modeling_v2_glm_phase.py             |  213 ---
 .../test_modeling_v2_glm_selection.py         |  573 --------
 .../test_modeling_v2_glm_spec_worker.py       |  514 -------
 .../test_modeling_v2_glm_targets.py           |  597 ---------
 .../test_modeling_v2_glm_vocabulary.py        |  198 ---
 .../absorbed_query/absorbed_query_fixture.py  |  223 ----
 .../shared/kernels/absorbed_query/conftest.py |   22 -
 ...st_absorbed_query_attention_integration.py |  870 ------------
 .../test_absorbed_query_proj.py               |  869 ------------
 ...test_draft_query_draft_loop_integration.py | 1103 ---------------
 .../test_draft_query_fused_dsa_metadata.py    |  745 -----------
 .../test_draft_query_mtp_query_proj.py        |  439 ------
 ...test_shared_query_attention_integration.py |  447 -------
 .../test_shared_query_shared_absorb.py        |  599 ---------
 .../attention_output_tail/real_mla_harness.py |  485 -------
 .../attention_output_tail/tail_fixtures.py    |  191 ---
 .../attention_output_tail/tail_reference.py   |  101 --
 .../test_attn_tail_all_layers.py              |  130 --
 .../test_attn_tail_attn_tail.py               |  263 ----
 .../test_attn_tail_bridge_mpi.py              |  244 ----
 .../test_attn_tail_contract_o.py              |  290 ----
 .../test_attn_tail_contract_o_forward.py      |  276 ----
 .../test_attn_tail_gates.py                   |  132 --
 .../test_attn_tail_mla_adapter.py             |  359 -----
 .../test_attn_tail_mla_forward.py             |  269 ----
 .../test_attn_tail_mla_forward_mpi.py         |  423 ------
 .../test_attn_tail_pdl_stress.py              |  506 -------
 .../test_attn_tail_sanitizer_probe.py         |  115 --
 .../test_into_hop_producer.py                 |  358 -----
 .../test_tail_prefetch.py                     |  531 --------
 .../test_tail_prefetch_mla_forward.py         |  142 --
 .../test_cluster_merge_adapter.py             |  171 ---
 .../test_global_emit_adapter.py               |  535 --------
 ...t_composition_global_emit_indexer_front.py |  387 ------
 .../_decode_prologue_oracle.py                |  355 -----
 .../decode_prologue/test_decode_prologue.py   |  421 ------
 .../deferred_input_prep/input_prep_testlib.py |  230 ----
 .../test_input_prep_compat.py                 |  340 -----
 .../test_input_prep_kernels.py                |  872 ------------
 .../shared/kernels/deferred_norm/conftest.py  |   24 -
 .../deferred_norm/deferred_norm_utils.py      |  299 -----
 ...est_deferred_norm_attention_integration.py |  495 -------
 .../test_deferred_norm_controls.py            |  425 ------
 .../test_deferred_norm_hop_mpi.py             |  835 ------------
 .../test_deferred_norm_null_controls.py       |  559 --------
 .../test_deferred_norm_projection.py          |  293 ----
 .../test_deferred_norm_routed_mpi.py          |  946 -------------
 .../deferred_norm/test_deferred_norm_scope.py |  327 -----
 .../test_deferred_norm_scope_mpi.py           |  342 -----
 .../test_evict_hints_adapter.py               |  126 --
 .../test_evict_hints_config.py                |   85 --
 .../test_evict_hints_hop_mpi.py               |  424 ------
 .../test_evict_hints_projection.py            |  156 ---
 .../kernels/dense_mlp/dense_mlp_utils.py      |  190 ---
 .../shared/kernels/dense_mlp/pdl_delay.cu     |   74 -
 .../shared/kernels/dense_mlp/test_adapter.py  |  171 ---
 .../kernels/dense_mlp/test_dense_mlp.py       |  410 ------
 .../test_draft_front_draft_front.py           |  853 ------------
 .../draft_live_rows/test_live_rows_adapter.py |  472 -------
 .../draft_loop/_draft_loop_fixtures.py        |  311 -----
 .../draft_loop/test_draft_loop_draft_loop.py  |  614 ---------
 .../draft_loop/test_draft_loop_draft_step.py  |  314 -----
 .../draft_loop/test_draft_loop_embedding.py   |  177 ---
 .../kernels/draft_tail/draft_tail_nearties.py |  296 ----
 .../kernels/draft_tail/draft_tail_oracle.py   |  148 --
 .../draft_tail/test_draft_tail_draft_tail.py  |  736 ----------
 .../draft_tail/test_draft_tail_eligibility.py |  239 ----
 .../draft_tail/test_draft_tail_nearties.py    |   89 --
 .../fused_kv_indexer_projection/conftest.py   |   24 -
 .../fused_kv_utils.py                         |  288 ----
 .../test_fused_kv_attention_integration.py    |  275 ----
 .../test_fused_kv_entry.py                    |  294 ----
 .../test_fused_kv_projection.py               |  331 -----
 .../test_host_pages_host_policy.py            |  176 ---
 .../test_host_pages_selector.py               |   82 --
 .../test_indexer_front_adapter.py             |  246 ----
 .../indexer_topk/test_cross_rank_cp_policy.py |  195 ---
 .../test_indexer_topk_indexer_decode.py       | 1081 ---------------
 .../test_indexer_topk_pipeline.py             | 1188 -----------------
 .../indexer_topk/test_indexer_topk_policy.py  |  243 ----
 .../oproj_weight_staging/test_oproj_stage.py  |  839 ------------
 .../shared/kernels/plan_testlib.py            |   92 --
 .../post_attention_hop/into_hop_jit.py        |   81 --
 .../test_attn_hop_static.py                   |  230 ----
 .../test_into_hop_single_gpu.py               |  276 ----
 .../post_moe_hop/post_moe_hop_mpi_utils.py    |  261 ----
 .../post_moe_hop/post_moe_hop_real_fixture.py |  211 ---
 .../test_post_moe_hop_post_moe_mpi.py         |  796 -----------
 .../test_post_moe_hop_protocol_model.py       |  287 ----
 .../test_post_moe_hop_real_weights_mpi.py     |  218 ---
 .../shared/kernels/rope_kv_append/conftest.py |   22 -
 .../kernels/rope_kv_append/pdl_producer.py    |  194 ---
 .../kernels/rope_kv_append/rope_append_env.py |  624 ---------
 .../rope_kv_append/rope_append_stack.py       |  727 ----------
 .../test_rope_append_composition.py           |  226 ----
 .../test_rope_append_integration.py           |  501 -------
 .../rope_kv_append/test_rope_append_ops.py    |  597 ---------
 .../router_gemm/test_pre_moe_hop_mpi.py       |  241 ----
 .../kernels/router_gemm/test_router_gemm.py   |  446 -------
 .../kernels/sparse_mla_decode/conftest.py     |  100 --
 .../sparse_mla_decode/sparse_mla_case.py      |  397 ------
 .../sparse_mla_decode/sparse_mla_harness.py   | 1026 --------------
 .../sparse_mla_pdl_helpers.py                 |  204 ---
 .../sparse_mla_real_fixture.py                |  310 -----
 .../sparse_mla_stack_case.py                  |  480 -------
 .../test_sparse_mla_adapter.py                |  100 --
 .../test_sparse_mla_composition.py            |  213 ---
 .../test_sparse_mla_dispatch_accounting.py    |  123 --
 .../test_sparse_mla_native_semantics.py       |  171 ---
 .../test_sparse_mla_pdl_stress.py             |  479 -------
 .../test_sparse_mla_ragged_verify.py          |   99 --
 .../test_sparse_mla_real_weights.py           |  404 ------
 .../test_sparse_mla_rope_append_bridge.py     |  195 ---
 .../test_sparse_mla_sanitizer_cases.py        |  150 ---
 .../test_sparse_mla_sparse_mla.py             |  285 ----
 .../test_sparse_mla_stack_composition.py      |  350 -----
 .../shared/kernels/test_kernels_contracts.py  |  157 ---
 .../shared/kernels/test_kernels_op_names.py   |  131 --
 .../shared/test_modeling_v2_glm_attention.py  |  754 -----------
 .../test_modeling_v2_glm_decline_kinds.py     |  534 --------
 .../shared/test_modeling_v2_glm_dispatch.py   |  380 ------
 .../test_modeling_v2_glm_moe_dsa_isolation.py |  297 -----
 .../test_modeling_v2_glm_moe_dsa_routing.py   |  700 ----------
 165 files changed, 126 insertions(+), 59164 deletions(-)
 create mode 100644 tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
index 8681fe5471c4..d8fed3900e42 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
@@ -123,6 +123,8 @@ The targets, as `models/<family>/<target>`:
   `kernels/`, with their C++ kernels and ops in `cpp/tensorrt_llm/kernels/glmMoeDsa` and
   `cpp/tensorrt_llm/thop/glmMoeDsa`; `shared/dispatch.py` selects them and lists the knobs.
   The tp8 and tep targets run the layer composition of `shared/layers.py`, each with its own plan.
+  Each GLM-5.2 target is gated end to end by
+  `tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py` (GSM8K under `require`).
 * `glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep`: the same checkpoint with attention tensor-parallel
   over four or eight ranks and the experts expert-parallel over the same ranks. It runs the
   shared fast paths that take its tensor-parallel size, in front of the built-in
diff --git a/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py b/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
new file mode 100644
index 000000000000..7b6fc2d13e50
--- /dev/null
+++ b/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
@@ -0,0 +1,118 @@
+# 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.
+"""Whole-model gates for the glm_moe_dsa modeling_v2 targets (GLM-5.2 NVFP4, sm_103).
+
+One case per target, each serving GSM8K end to end with ``modeling_v2="require"``.
+Under ``"auto"`` a deployment that missed a target's route would fall back to the
+built-in model, pass, and report the built-in's accuracy as the target's; under
+``"require"`` it raises instead, and so does a step a fast path covers but cannot
+run. The mode and the environment below are LLM API arguments, so they reach every
+rank.
+
+Each deployment is one the family's ``route()`` claims: the GLM-5.2 modelopt NVFP4
+checkpoint (bf16 MTP layer) with one-model MTP-Eagle as a linear chain, at the
+target's topology and draft length; the DEP case also matches its cell's KV-cache
+dtype and MoE backend.
+"""
+
+import pytest
+
+from tensorrt_llm import LLM
+from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.llmapi import (
+    CudaGraphConfig,
+    DeepSeekSparseAttentionConfig,
+    KvCacheConfig,
+    MoeConfig,
+    MTPDecodingConfig,
+)
+
+from ...conftest import llm_models_root
+from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness
+
+# The targets assert their own SM at construction.
+skip_not_sm103 = pytest.mark.skipif(
+    get_sm_version() != 103, reason="the glm_moe_dsa targets are certified on sm_103 only"
+)
+
+# On every rank: check the target's step contract on the first forward, and run the
+# indexer key projection in BF16, the form the family's fused indexer paths take.
+# The deployments below are the ones the targets were validated at: per-target batch,
+# KV-cache dtype, MoE backend and DSA kernels.
+_RANK_ENV = {MODELING_V2_VALIDATE_ENV: "1", "TRTLLM_DSA_INDEXER_BF16": "1"}
+
+
+class TestModelingV2Glm52Nvfp4Sm103(LlmapiAccuracyTestHarness):
+    """glm-5.2-nvfp4 / sm_103 at tp8, tep4 and dep8."""
+
+    MODEL_NAME = "zai-org/GLM-5.2"
+    MODEL_PATH = f"{llm_models_root()}/GLM-5.2-NVFP4"
+
+    @staticmethod
+    def _llm_kwargs(*, draft_len: int, kv_dtype: str, moe_backend: str, batch: int) -> dict:
+        return dict(
+            modeling_v2="require",
+            env_overrides=dict(_RANK_ENV),
+            speculative_config=MTPDecodingConfig(max_draft_len=draft_len),
+            kv_cache_config=KvCacheConfig(
+                dtype=kv_dtype, enable_block_reuse=False, free_gpu_memory_fraction=0.75
+            ),
+            moe_config=MoeConfig(backend=moe_backend),
+            sparse_attention_config=DeepSeekSparseAttentionConfig(
+                enable_heuristic_topk=True,
+                use_cute_dsl_paged_mqa_logits=True,
+                use_cute_dsl_topk=True,
+            ),
+            cuda_graph_config=CudaGraphConfig(max_batch_size=batch, enable_padding=True),
+            max_batch_size=batch,
+            max_seq_len=8192,
+            enable_chunked_prefill=True,
+        )
+
+    def _evaluate(self, **llm_kwargs) -> None:
+        with LLM(self.MODEL_PATH, **llm_kwargs) as llm:
+            GSM8K(self.MODEL_NAME).evaluate(llm)
+
+    @skip_not_sm103
+    @pytest.mark.skip_less_mpi_world_size(8)
+    def test_gsm8k_tp8(self):
+        """The low-latency target: attention and experts tensor-parallel over 8 ranks."""
+        self._evaluate(
+            tensor_parallel_size=8,
+            moe_expert_parallel_size=1,
+            **self._llm_kwargs(draft_len=5, kv_dtype="fp8", moe_backend="TRTLLM", batch=1),
+        )
+
+    @skip_not_sm103
+    @pytest.mark.skip_less_device(4)
+    def test_gsm8k_tep4(self):
+        """Attention tensor-parallel over 4 ranks, experts expert-parallel over them."""
+        self._evaluate(
+            tensor_parallel_size=4,
+            moe_expert_parallel_size=4,
+            **self._llm_kwargs(draft_len=5, kv_dtype="fp8", moe_backend="TRTLLM", batch=2),
+        )
+
+    @skip_not_sm103
+    @pytest.mark.skip_less_mpi_world_size(8)
+    def test_gsm8k_dep8(self):
+        """Attention data parallel over 8 ranks: the FP8-KV, CuteDSL-MoE, MTP3 cell."""
+        self._evaluate(
+            tensor_parallel_size=8,
+            moe_expert_parallel_size=8,
+            enable_attention_dp=True,
+            **self._llm_kwargs(draft_len=3, kv_dtype="fp8", moe_backend="CUTEDSL", batch=16),
+        )
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index 1ae313f8fdad..b4ddbe3b50bb 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -20,6 +20,8 @@ l0_gb300_multi_gpus:
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock]
+  # GLM-5.2 NVFP4 modeling_v2 tep target, GSM8K under modeling_v2="require". Measured 17 min on 4 GB300.
+  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103::test_gsm8k_tep4 TIMEOUT (60)
   # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere
   # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry
   # is what picks them up:
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
index 0806efd6646a..26594bfbe349 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
@@ -17,3 +17,7 @@ l0_gb300_multi_nodes_node2_gpu8:
   # Kimi K3 GSM8K accuracy plus an acceptance-length gate on the TEP8 DSpark
   # recipe. Measured 25m45s on 2x4 GB300.
   - accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 TIMEOUT (60)
+  # GLM-5.2 NVFP4 modeling_v2 targets at tp8 and dep8, GSM8K under modeling_v2="require".
+  # Measured on 2x4 GB300: tp8 (batch 1) 56 min, dep8 10 min.
+  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103::test_gsm8k_tp8 TIMEOUT (90)
+  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103::test_gsm8k_dep8 TIMEOUT (60)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py
deleted file mode 100644
index 564c5b927364..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_plan.py
+++ /dev/null
@@ -1,411 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The fast paths the GLM-5.2 dep target serves in each of its cells, how it picks the cell, and
-what its plan does with the declines its deployments log.
-
-Host-only: the plan is resolved and installed from the cell's ``TargetPaths`` exactly as the
-shell does at construction. The engagement of each deployment is the per-rank decision set the
-fast paths log there (the same on every rank): the paths that engage, and the (path, key,
-reason) triples of the paths the cell serves that decline. None of the declines may be a miss:
-under ``modeling_v2='require'`` a miss raises.
-"""
-
-from __future__ import annotations
-
-import types
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-pytestmark = pytest.mark.cpu_only
-
-D = dispatch.Decline
-REQUIRE, AUTO = ModelingV2Mode.REQUIRE, ModelingV2Mode.AUTO
-
-# The deployments, by (ranks, LM head tensor-parallel in attention DP, draft tokens per request),
-# with their KV-cache dtype and MoE backend.
-DEP16 = (16, False, 3)
-DEP8 = (8, False, 3)
-DEP8_LM_HEAD_TP = (8, True, 5)
-DEPLOYMENTS = {
-    DEP16: (QuantAlgo.FP8, "TRTLLM"),
-    DEP8: (QuantAlgo.FP8, "CUTEDSL"),
-    DEP8_LM_HEAD_TP: (QuantAlgo.NVFP4, "CUTEDSL"),
-}
-IDS = ["dep16", "dep8", "dep8-lm-head-tp"]
-
-# The paths that engage at each deployment (device work on every rank).
-ENGAGED = {
-    DEP16: {
-        "rope_kv_append",
-        "indexer_layer_front",
-        "draft_loop",
-        "decode_prologue",
-        "host_page_retention",
-    },
-    DEP8: {
-        "rope_kv_append",
-        "indexer_layer_front",
-        "draft_loop",
-        "decode_prologue",
-        "host_page_retention",
-    },
-    DEP8_LM_HEAD_TP: {"draft_loop", "decode_prologue", "host_page_retention"},
-}
-# The bases the engaged paths extend; they decline every call at these deployments, or (the
-# post-attention hop) are never built.
-BASES = {
-    DEP16: {
-        "absorbed_query",
-        "draft_absorbed_query",
-        "indexer_topk",
-        "cross_rank_topk",
-        "post_attention_hop",
-    },
-    DEP8: {
-        "absorbed_query",
-        "draft_absorbed_query",
-        "indexer_topk",
-        "cross_rank_topk",
-        "post_attention_hop",
-    },
-    DEP8_LM_HEAD_TP: {"post_attention_hop"},
-}
-SERVED = {cell: ENGAGED[cell] | BASES[cell] for cell in DEPLOYMENTS}
-# Paths that log an engaged decision at these deployments without device work there: the
-# step-input stager's feature check (its stager, store writer and execution-stream loop all
-# decline attention DP) and the cross-layer expert prefetch links (the prefetch is issued only
-# by the routed-MoE launch, which does not run with expert-parallel experts).
-LOG_ONLY = {"deferred_input_prep", "sampler_store", "execution_stream_loop"} | {
-    "cross_layer_expert_prefetch"
-}
-
-ADP = "attention DP / CP not supported"
-ADP_TOPK = "attention data parallelism"
-
-# The declines the served paths log with an FP8 KV cache, on every rank.
-DECLINES_FP8_KV = (
-    *(("absorbed_query", f"layer{i}", ADP) for i in range(78)),
-    ("absorbed_query", "layer78", "target layers only"),
-    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
-    ("draft_absorbed_query", "layer78", ADP),
-    *(("indexer_layer_front", f"target_layer{i}", ADP) for i in range(78)),
-    ("indexer_layer_front", "draft_layer78", ADP),
-    ("indexer_layer_front", "cp_indexer_decode_cross_rank_topk", ADP_TOPK),
-    ("indexer_topk", "indexer_decode", ADP_TOPK),
-    ("indexer_topk", "indexer_decode", "MTP draft step (target layers only)"),
-    ("cross_rank_topk", "cp_indexer_decode", ADP_TOPK),
-    ("draft_loop", "draft_loop:context rows present", "context rows present"),
-    ("draft_loop", "draft_step:dense top-k fallback active", "dense top-k fallback active"),
-    (
-        "decode_prologue",
-        "target_prologue[not a generation-only step]",
-        "not a generation-only step",
-    ),
-)
-# The declines the served paths log with an NVFP4 KV cache and the LM head tensor-parallel in
-# attention DP, on every rank.
-DECLINES_NVFP4_KV = (
-    ("draft_loop", "draft_loop:LM-head TP in attention DP", "LM-head TP in attention DP"),
-    ("draft_loop", "draft_loop:context rows present", "context rows present"),
-    (
-        "decode_prologue",
-        "target_prologue[not a generation-only step]",
-        "not a generation-only step",
-    ),
-)
-DECLINES = {DEP16: DECLINES_FP8_KV, DEP8: DECLINES_FP8_KV, DEP8_LM_HEAD_TP: DECLINES_NVFP4_KV}
-# The engaged decisions the served paths log, as (path, key).
-ENGAGED_KEYS = {
-    DEP16: {
-        ("rope_kv_append", "rope"),
-        ("rope_kv_append", "bridge"),
-        ("indexer_layer_front", "rope"),
-        ("draft_loop", "draft_embedding"),
-        ("draft_loop", "draft_loop:engaged"),
-        ("draft_loop", "draft_step:engaged"),
-        ("decode_prologue", "target_prologue"),
-        ("decode_prologue", "target_embedding"),
-        ("host_page_retention", "host_policy"),
-    },
-    DEP8_LM_HEAD_TP: {
-        ("draft_loop", "draft_embedding"),
-        ("decode_prologue", "target_prologue"),
-        ("decode_prologue", "target_embedding"),
-        ("host_page_retention", "host_policy"),
-    },
-}
-ENGAGED_KEYS[DEP8] = ENGAGED_KEYS[DEP16]
-
-
-def _model_config(cell, spec=True, **overrides):
-    """What ``target_paths`` reads off the model config of a deployment; ``overrides`` replace
-    its ``kv_cache`` and ``moe_backend``."""
-    ranks, lm_head_tp, draft = cell
-    kv_cache, moe_backend = DEPLOYMENTS.get(cell, (QuantAlgo.FP8, "TRTLLM"))
-    facts = {"kv_cache": kv_cache, "moe_backend": moe_backend, **overrides}
-    return types.SimpleNamespace(
-        mapping=types.SimpleNamespace(
-            world_size=ranks, enable_attention_dp=True, enable_lm_head_tp_in_adp=lm_head_tp
-        ),
-        spec_config=types.SimpleNamespace(max_draft_len=draft) if spec else None,
-        quant_config=types.SimpleNamespace(kv_cache_quant_algo=facts["kv_cache"]),
-        moe_backend=facts["moe_backend"],
-    )
-
-
-@pytest.fixture
-def installed(monkeypatch):
-    """Resolve and install the plan of a deployment's cell, as the shell does."""
-    for name in ("_INSTALLED", "_VALIDATE", "_require", "_TARGET"):
-        monkeypatch.setattr(dispatch, name, list(getattr(dispatch, name)))
-    monkeypatch.setattr(dispatch, "_logged", set())
-    monkeypatch.setattr(dispatch, "_missed", {})
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-    dispatch.classify_decline.cache_clear()
-
-    def install(cell, environ=None, mode=REQUIRE):
-        paths = modeling.target_paths(_model_config(cell), mode)
-        plan = dispatch.resolve(mode, environ or {}, paths)
-        dispatch.install(plan, environ or {})
-        return plan
-
-    yield install
-    dispatch.classify_decline.cache_clear()
-
-
-# --- the cells ------------------------------------------------------------------------------
-
-
-def test_the_cells_are_the_three_deployments():
-    assert set(modeling.CELLS) == set(DEPLOYMENTS)
-    for cell, (kv_cache, moe_backend) in DEPLOYMENTS.items():
-        assert (modeling.CELLS[cell].kv_cache, modeling.CELLS[cell].moe_backend) == (
-            kv_cache,
-            moe_backend,
-        )
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-@pytest.mark.parametrize("mode", [REQUIRE, AUTO], ids=["require", "auto"])
-def test_each_deployment_selects_its_cells_paths(cell, mode):
-    assert modeling.target_paths(_model_config(cell), mode) is modeling.CELLS[cell].paths
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_the_kv_cache_dtype_may_be_the_enum_or_its_value(cell):
-    kv_cache = DEPLOYMENTS[cell][0]
-    config = _model_config(cell, kv_cache=kv_cache.value)
-    assert not isinstance(config.quant_config.kv_cache_quant_algo, QuantAlgo)
-    assert modeling.target_paths(config, REQUIRE) is modeling.CELLS[cell].paths
-
-
-# A deployment ``route`` claims whose KV-cache dtype or MoE backend is not its cell's, and one
-# outside the cells: (cell, the facts that differ, what the error names).
-_OUTSIDE = {
-    "dep16-nvfp4-kv": (DEP16, {"kv_cache": QuantAlgo.NVFP4}, "KV cache NVFP4, MoE backend TRTLLM"),
-    "dep16-cutedsl": (DEP16, {"moe_backend": "CUTEDSL"}, "KV cache FP8, MoE backend CUTEDSL"),
-    "dep8-trtllm": (DEP8, {"moe_backend": "TRTLLM"}, "KV cache FP8, MoE backend TRTLLM"),
-    "dep8-nvfp4-kv": (DEP8, {"kv_cache": QuantAlgo.NVFP4}, "KV cache NVFP4, MoE backend CUTEDSL"),
-    "dep8-unquantized-kv": (DEP8, {"kv_cache": None}, "KV cache None, MoE backend CUTEDSL"),
-    "dep8-lm-head-tp-fp8-kv": (
-        DEP8_LM_HEAD_TP,
-        {"kv_cache": QuantAlgo.FP8},
-        "KV cache FP8, MoE backend CUTEDSL",
-    ),
-    "dep8-lm-head-tp-cutlass": (
-        DEP8_LM_HEAD_TP,
-        {"moe_backend": "CUTLASS"},
-        "KV cache NVFP4, MoE backend CUTLASS",
-    ),
-    "dep16-mtp5": (
-        (16, False, 5),
-        {},
-        "16 ranks, LM head tensor-parallel in attention DP False, 5",
-    ),
-}
-
-
-@pytest.mark.parametrize("cell, overrides, named", list(_OUTSIDE.values()), ids=list(_OUTSIDE))
-def test_a_deployment_outside_the_cells_raises_under_require(cell, overrides, named):
-    config = _model_config(cell, **overrides)
-    with pytest.raises(ValueError, match="is not a deployment this target serves") as excinfo:
-        modeling.target_paths(config, REQUIRE)
-    message = str(excinfo.value)
-    assert named in message, message
-    # The error names the three cells.
-    assert message.count("draft tokens: KV cache") == 3, message
-
-
-@pytest.mark.parametrize("cell, overrides, named", list(_OUTSIDE.values()), ids=list(_OUTSIDE))
-def test_a_deployment_outside_the_cells_serves_nothing_under_auto(
-    monkeypatch, cell, overrides, named
-):
-    warnings = []
-    monkeypatch.setattr(modeling.logger, "warning", warnings.append)
-    config = _model_config(cell, **overrides)
-    assert modeling.target_paths(config, AUTO) is modeling.NO_PATHS
-    assert modeling.NO_PATHS.served == frozenset()
-    (warning,) = warnings
-    assert named in warning and warning.endswith("no fast path runs"), warning
-
-
-def test_without_speculation_no_cell_matches():
-    config = _model_config(DEP8, spec=False)
-    with pytest.raises(ValueError, match="None draft tokens"):
-        modeling.target_paths(config, REQUIRE)
-
-
-def test_the_auto_plan_outside_the_cells_runs_no_fast_path(installed, monkeypatch):
-    monkeypatch.setattr(modeling.logger, "warning", lambda message: None)
-    paths = modeling.target_paths(_model_config(DEP8, moe_backend="TRTLLM"), AUTO)
-    plan = dispatch.resolve(AUTO, {}, paths)
-    dispatch.install(plan, {})
-    assert dispatch.installed() == frozenset()
-    assert plan.describe().startswith("modeling_v2=auto, 0/0 fast paths")
-    assert not dispatch.require_mode()
-
-
-# --- the plans ------------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_the_cell_serves_the_engaged_paths_and_their_bases(installed, cell):
-    plan = installed(cell)
-    assert set(plan.served) == SERVED[cell]
-    assert dispatch.installed() == frozenset(SERVED[cell])
-    count = len(SERVED[cell])
-    assert plan.describe().startswith(f"modeling_v2=require, {count}/{count} fast paths")
-
-
-def test_the_cells_serve_ten_ten_and_four_paths():
-    assert [len(SERVED[cell]) for cell in DEPLOYMENTS] == [10, 10, 4]
-    assert modeling.CELLS[DEP16].paths is modeling.CELLS[DEP8].paths is modeling.FP8_KV_PATHS
-    assert modeling.CELLS[DEP8_LM_HEAD_TP].paths is modeling.NVFP4_KV_PATHS
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_every_base_of_a_served_path_is_served(cell):
-    served = SERVED[cell]
-    for path in served:
-        assert dispatch.REQUIRED_FEATURES.get(path, frozenset()) <= served, path
-    # The bases are exactly what the engaged paths extend.
-    required = set().union(*(dispatch.REQUIRED_FEATURES.get(p, frozenset()) for p in ENGAGED[cell]))
-    assert required - ENGAGED[cell] == BASES[cell]
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_no_log_only_path_is_served(cell):
-    """The step-input stager's parts and the expert prefetch do no device work under attention
-    DP; the execution-stream loop has no attention-DP gate of its own, so it must not be
-    served."""
-    assert not (SERVED[cell] & LOG_ONLY)
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_every_kernel_package_is_the_familys(cell):
-    assert dict(modeling.CELLS[cell].paths.kernel_modules) == {}
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_the_record_line_lists_the_served_paths(installed, monkeypatch, cell):
-    lines = []
-    monkeypatch.setattr(dispatch, "_log_info", lines.append)
-    installed(cell)
-    (record,) = [line for line in lines if "record" in line]
-    listed = record.split("paths=", 1)[1].split(" ", 1)[0].split(",")
-    assert {name.split("[", 1)[0] for name in listed} == SERVED[cell]
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_disable_may_name_only_a_served_path(installed, cell):
-    with pytest.raises(ValueError, match="not fast paths of this target"):
-        installed(cell, {dispatch.DISABLE_ENV: "execution_stream_loop"})
-    plan = installed(cell, {dispatch.DISABLE_ENV: "decode_prologue"})
-    assert not plan.enabled("decode_prologue") and plan.enabled("draft_loop")
-
-
-# Each path of the failure drop order with what it turns off (itself and its dependents).
-_DROP_ORDER = {
-    "decode_prologue": {"decode_prologue"},
-    "indexer_layer_front": {"indexer_layer_front"},
-    "draft_loop": {"draft_loop", "decode_prologue"},
-    "rope_kv_append": {"rope_kv_append", "indexer_layer_front"},
-}
-
-
-@pytest.mark.parametrize("path", list(_DROP_ORDER))
-def test_each_path_of_the_drop_order_can_be_disabled(installed, path):
-    plan = installed(DEP8, {dispatch.DISABLE_ENV: path})
-    off = {name for name in SERVED[DEP8] if not plan.enabled(name)}
-    assert off == _DROP_ORDER[path]
-    assert plan.describe().startswith(f"modeling_v2=require, {10 - len(off)}/10 fast paths")
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_the_routed_accumulation_knob_is_refused(installed, cell):
-    with pytest.raises(ValueError, match="does not serve routed_moe_into_hop"):
-        installed(cell, {dispatch.ACCUMULATION_ENV: "identity"})
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_the_cross_rank_wait_timeout_is_read(installed, cell):
-    plan = installed(cell, {dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV: "60000"})
-    assert plan.knobs == ((dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV, "60000"),)
-
-
-# --- the declines ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_no_logged_decline_is_a_miss(installed, cell):
-    """Under require, a decline decided once inside a path's cells would raise."""
-    installed(cell)
-    assert dispatch.require_mode()
-    for path, key, reason in DECLINES[cell]:
-        assert path in SERVED[cell], path
-        assert dispatch.classify_decline(path, key, reason) is D.UNCOVERED, (path, key, reason)
-        assert dispatch.log_capability(path, key, False, reason) is False
-    assert dispatch.missed_counts() == {}
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_the_engaged_decisions_are_of_served_paths(installed, cell):
-    installed(cell)
-    for path, key in ENGAGED_KEYS[cell]:
-        assert path in ENGAGED[cell], path
-        assert dispatch.log_capability(path, key, True) is True
-    assert {path for path, _ in ENGAGED_KEYS[cell]} == ENGAGED[cell]
-
-
-@pytest.mark.parametrize("cell", list(DEPLOYMENTS), ids=IDS)
-def test_every_base_declines_every_call(cell):
-    """A base the cell serves logs no engaged decision; every one but the never-built
-    post-attention hop declines."""
-    declining = {path for path, _, _ in DECLINES[cell]}
-    assert BASES[cell] - {"post_attention_hop"} <= declining
-    assert not ({path for path, _ in ENGAGED_KEYS[cell]} & BASES[cell])
-
-
-def test_the_family_cells_would_raise_on_the_query_fronts(installed):
-    """The family's cells hold every target layer for the query fronts; the dep table is what
-    keeps an attention-DP deployment from raising on their declines under require."""
-    plan = dispatch.resolve(REQUIRE, {}, dispatch.TargetPaths(served=SERVED[DEP8]))
-    dispatch.install(plan, {})
-    misses = {
-        (path, key)
-        for path, key, reason in DECLINES_FP8_KV
-        if dispatch.classify_decline(path, key, reason) is D.MISSED
-    }
-    assert ("absorbed_query", "layer0") in misses
-    assert ("draft_absorbed_query", "layer78") in misses
-    assert ("indexer_layer_front", "target_layer0") in misses
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py
deleted file mode 100644
index 3925082357b9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/test_modeling_v2_glm_dep_target.py
+++ /dev/null
@@ -1,501 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 dep target's construction, adoption, decode loop and step-loop hooks.
-
-Host-only. Adoption runs on the stand-in built-in modules and kernel factories of the tp8
-target's layer test (``test_modeling_v2_glm_layers.py``, loaded by path), with the fast paths
-of a dep cell: the attention modules and the speculative worker are re-classed, and the
-decoder layers, the MoE modules and the draft layer stay the built-in classes. The shell's
-construction runs with the built-in constructor stubbed and a recorder in place of
-``adopt_model``, as the tp8 target's shell test does.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import inspect
-import os
-import types
-from pathlib import Path
-
-import pytest
-import torch
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2 import _router_index
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    attention,
-    dispatch,
-    layers,
-    runtime,
-    spec_worker,
-)
-from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
-from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
-from tensorrt_llm._torch.models.modeling_deepseekv3 import (
-    DeepseekV3DecoderLayer,
-    DeepseekV3Model,
-    Deepseekv3MoE,
-    DeepseekV3MTP,
-    DeepseekV32Attention,
-    GlmMoeDsaForCausalLM,
-)
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of
-from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-pytestmark = pytest.mark.cpu_only
-
-PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
-REQUIRE, AUTO = _router_index.ModelingV2Mode.REQUIRE, _router_index.ModelingV2Mode.AUTO
-
-
-def _load_tp8_layer_test():
-    path = (
-        Path(__file__).resolve().parents[1]
-        / "glm_5_2_nvfp4__sm_103__tp8"
-        / "test_modeling_v2_glm_layers.py"
-    )
-    spec = importlib.util.spec_from_file_location("_glm_tp8_layer_test_for_dep", path)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-T = _load_tp8_layer_test()
-
-DEP16, DEP8, DEP8_LM_HEAD_TP = (16, False, 3), (8, False, 3), (8, True, 5)
-CELLS = [DEP16, DEP8, DEP8_LM_HEAD_TP]
-IDS = ["dep16", "dep8", "dep8-lm-head-tp"]
-
-
-def _served(cell):
-    return modeling.CELLS[cell].paths.served
-
-
-@pytest.fixture
-def fast_paths(monkeypatch):
-    def install(enabled):
-        paths = T._FastPaths(frozenset(enabled))
-        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-        monkeypatch.setattr(dispatch, "log_capability", paths.log_capability)
-        return paths
-
-    return install
-
-
-def _state(model):
-    return dict(modeling.adopted_state(types.SimpleNamespace(model=model.model)))
-
-
-# --- adoption -------------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("cell", [DEP16, DEP8], ids=IDS[:2])
-def test_with_an_fp8_kv_cache_the_attention_and_the_worker_are_adopted(fast_paths, cell):
-    paths = fast_paths(_served(cell))
-    model = T._built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-
-    # No fast-path object is built: the attention tail, the hops, the router GEMM and the draft
-    # front are not served.
-    assert paths.events == []
-    dense, moe, last, draft = model.model.layers
-    assert [type(m) for m in (dense, moe, last)] == [DeepseekV3DecoderLayer] * 3
-    assert type(draft) is DeepseekV3MTP
-    assert [type(m.mlp) for m in (moe, last, draft)] == [Deepseekv3MoE] * 3
-    for layer in (dense, moe, last, draft):
-        mla = layer.self_attn
-        assert type(mla) is attention.GlmMLA
-        assert type(mla.sparse_attn_hooks) is attention.GlmDsaHooks
-        assert type(mla.mqa) is attention.GlmDsaBackend
-        assert mla._attn_tail is None and mla._tail_prefetch is False
-        # The FMHA libraries stay the registry's.
-        assert "prepended_fmha_lib_classes" not in vars(mla.mqa)
-    assert [type(m.self_attn.mqa.indexer) for m in (dense, last)] == [attention.GlmIndexer] * 2
-    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
-
-
-def test_with_an_nvfp4_kv_cache_only_the_worker_is_adopted(fast_paths):
-    paths = fast_paths(_served(DEP8_LM_HEAD_TP))
-    model = T._built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-
-    assert paths.events == []
-    for layer in model.model.layers:
-        assert type(layer.self_attn) is DeepseekV32Attention
-        assert type(layer.self_attn.sparse_attn_hooks) is DSAHooks
-        assert type(layer.self_attn.mqa) is DSATrtllmAttention
-    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
-
-
-@pytest.mark.parametrize(
-    "disabled, attention_adopted, worker_adopted",
-    [
-        ({"rope_kv_append", "indexer_layer_front"}, False, True),
-        ({"indexer_layer_front"}, True, True),
-        ({"draft_loop", "decode_prologue"}, True, False),
-        (set(), True, True),
-    ],
-    ids=["no rope", "no indexer front", "no draft loop", "all"],
-)
-def test_each_adoption_follows_its_fast_path(
-    fast_paths, disabled, attention_adopted, worker_adopted
-):
-    fast_paths(_served(DEP8) - disabled)
-    model = T._built_in_model([])
-    modeling.adopt_model(model, plan=None)
-    expected = attention.GlmMLA if attention_adopted else DeepseekV32Attention
-    assert [type(layer.self_attn) for layer in model.model.layers] == [expected] * 4
-    worker = spec_worker.GlmMtpWorker if worker_adopted else MTPEagleWorker
-    assert type(model.spec_worker) is worker
-
-
-def test_without_a_plan_nothing_is_adopted(fast_paths):
-    fast_paths(modeling.NO_PATHS.served)
-    model = T._built_in_model([])
-    modeling.adopt_model(model, plan=None)
-    assert [type(layer.self_attn) for layer in model.model.layers] == [DeepseekV32Attention] * 4
-    assert type(model.spec_worker) is MTPEagleWorker
-
-
-def test_without_a_worker_the_attention_is_adopted_alone(fast_paths):
-    fast_paths(_served(DEP8))
-    model = T._built_in_model([])
-    model.spec_worker = None
-    modeling.adopt_model(model, plan=None)
-    assert type(model.model.layers[0].self_attn) is attention.GlmMLA
-
-
-@pytest.mark.parametrize("cell", CELLS, ids=IDS)
-def test_adoption_attaches_no_layer_state(fast_paths, cell):
-    """Every rank checks the same empty per-layer state: no fast-path object is attached to the
-    layers, their attention or their MoE."""
-    fast_paths(_served(cell))
-    model = T._built_in_model([])
-    modeling.adopt_model(model, plan=None)
-    state = _state(model)
-    assert list(state) == [f"model.layers.{i}" for i in range(4)]
-    assert all(attached == () for attached in state.values())
-
-
-@pytest.mark.parametrize("cell", CELLS, ids=IDS)
-def test_adoption_adds_no_parameter_buffer_or_module(fast_paths, cell):
-    fast_paths(_served(cell))
-    model = T._built_in_model([])
-    before = [T._keys(layer) for layer in model.model.layers]
-    modeling.adopt_model(model, plan=None)
-    assert [T._keys(layer) for layer in model.model.layers] == before
-
-
-def test_adoption_refuses_an_attention_module_that_is_not_the_builtin_class(fast_paths):
-    fast_paths(_served(DEP8))
-    model = T._built_in_model([])
-    attention.GlmMLA.adopt(model.model.layers[1].self_attn)
-    with pytest.raises(TypeError, match="GlmMLA adopts the built-in DeepseekV32Attention"):
-        modeling.adopt_model(model, plan=None)
-
-
-def test_the_target_uses_no_layer_class_of_the_family():
-    """The decoder layers, the MoE modules and the draft layer stay the built-in classes: the
-    target names none of the family's layer classes or their adoption."""
-    source = inspect.getsource(modeling)
-    for name in ("GlmDecoderLayer", "GlmMtpLayer", "GlmMoE", "layers.adopt_model"):
-        assert name not in source, name
-    assert modeling.adopted_state is layers.adopted_state
-
-
-# --- the core and its two targets -----------------------------------------------------------
-
-
-def _builtin_text_model() -> DeepseekV3Model:
-    model = DeepseekV3Model.__new__(DeepseekV3Model)
-    nn.Module.__init__(model)
-    return model
-
-
-def test_adopt_reclasses_the_builtin_text_model_in_place():
-    model = _builtin_text_model()
-    core = modeling.GlmMoeDsaDepModelingV2.adopt(model)
-    assert core is model and type(core) is modeling.GlmMoeDsaDepModelingV2
-    assert isinstance(core._prefill, modeling.PrefillTarget) and core._prefill.core is core
-    assert isinstance(core._decode, modeling.DecodeTarget) and core._decode.core is core
-    assert list(core.named_modules()) == [("", core)]
-    with pytest.raises(TypeError, match="GlmMoeDsaDepModelingV2"):
-        modeling.GlmMoeDsaDepModelingV2.adopt(core)
-
-
-@pytest.mark.parametrize("num_contexts, target", [(0, "_decode"), (1, "_prefill"), (4, "_prefill")])
-def test_each_step_goes_to_the_target_of_its_phase(num_contexts, target):
-    core = modeling.GlmMoeDsaDepModelingV2.adopt(_builtin_text_model())
-    md = types.SimpleNamespace(num_contexts=num_contexts)
-    assert core._select_target(md) is getattr(core, target)
-
-
-@pytest.fixture
-def decode_core():
-    calls = []
-    core = DeepseekV3Model.__new__(DeepseekV3Model)
-    nn.Module.__init__(core)
-    core.embed_tokens = T._Embedding(calls)
-    core.layers = nn.ModuleList([T._RecordingLayer(i, calls) for i in range(T.NUM_LAYERS + 1)])
-    core.num_hidden_layers = T.NUM_LAYERS
-    core.mapping_with_cp = None
-    core = modeling.GlmMoeDsaDepModelingV2.adopt(core)
-    for layer in core.layers:
-        layer.register_forward_pre_hook(
-            lambda m, a, k: calls.append(("hook", m.index)), with_kwargs=True
-        )
-    return core, calls
-
-
-@pytest.mark.parametrize("prologue_rows", [None, "rows"], ids=["no packet", "packet"])
-def test_the_decode_loop_runs_the_decoder_layers_as_modules(decode_core, prologue_rows):
-    core, calls = decode_core
-    rows = torch.full((2, T.HIDDEN), 7.0) if prologue_rows else None
-    seen = []
-
-    def take_embedding(attn_metadata, input_ids, embed_tokens):
-        seen.append((attn_metadata, input_ids, embed_tokens))
-        return rows
-
-    core.decode_prologue = types.SimpleNamespace(take_embedding=take_embedding)
-    md = T._NoContexts()
-    with dispatch.step(DECODE):
-        out = core._decode.forward(md, input_ids="ids", position_ids="pos", spec_metadata="spec")
-
-    assert seen == [(md, "ids", core.embed_tokens)]
-    embedded = [] if prologue_rows else [("embed", "ids")]
-    start = rows if prologue_rows else torch.zeros(2, T.HIDDEN)
-    assert torch.equal(out, start + T.NUM_LAYERS)
-    layer_calls = []
-    for i in range(T.NUM_LAYERS):
-        kwargs = calls[len(embedded) + 2 * i + 1][1]
-        assert kwargs["residual"] == (None if i == 0 else f"residual {i - 1}")
-        assert kwargs["attn_metadata"] is md and kwargs["spec_metadata"] == "spec"
-        layer_calls += [("hook", i), (i, kwargs, DECODE)]
-    # The draft layer is not run, and nothing follows the last decoder layer.
-    assert calls == embedded + layer_calls
-
-
-def test_without_the_decode_prologue_the_decode_loop_is_the_builtin_loop(decode_core):
-    core, calls = decode_core
-    with dispatch.step(DECODE):
-        out = core._decode.forward(T._NoContexts(), input_ids="ids")
-    assert torch.equal(out, torch.zeros(2, T.HIDDEN) + T.NUM_LAYERS)
-    assert [c[0] for c in calls] == ["embed"] + ["hook", 0, "hook", 1, "hook", 2]
-
-
-def test_the_decode_loop_reads_no_fast_path_and_never_the_step_kind():
-    source = inspect.getsource(modeling.DecodeTarget.forward)
-    assert "dispatch." not in source
-    assert "num_contexts" not in source and "num_ctx" not in source
-
-
-# --- construction ---------------------------------------------------------------------------
-
-
-def _model_config(cell, **facts):
-    ranks, lm_head_tp, draft = cell
-    defaults = {
-        DEP16: (QuantAlgo.FP8, "TRTLLM"),
-        DEP8: (QuantAlgo.FP8, "CUTEDSL"),
-        DEP8_LM_HEAD_TP: (QuantAlgo.NVFP4, "CUTEDSL"),
-    }[cell]
-    return types.SimpleNamespace(
-        mapping=types.SimpleNamespace(
-            world_size=ranks,
-            enable_attention_dp=True,
-            enable_lm_head_tp_in_adp=lm_head_tp,
-        ),
-        spec_config=types.SimpleNamespace(max_draft_len=draft),
-        quant_config=types.SimpleNamespace(kv_cache_quant_algo=facts.get("kv_cache", defaults[0])),
-        moe_backend=facts.get("moe_backend", defaults[1]),
-    )
-
-
-RANK_UNIFORM_ENV = dispatch.SHARED_ENV + (_router_index.MODELING_V2_VALIDATE_ENV,)
-
-
-@pytest.fixture
-def shell_construction(monkeypatch):
-    """The shell class with the built-in constructor stubbed, and the events of a construction."""
-    events = []
-
-    def adopt_model(model, plan):
-        events.append(("adopt", type(model.model).__name__, plan))
-
-    def builtin_init(self, model_config):
-        events.append("built-in model")
-        events.append(("installed", dispatch.installed()))
-        nn.Module.__init__(self)
-        self.model_config = model_config
-        self.model = _builtin_text_model()
-        self.model.layers = nn.ModuleList([nn.Module(), nn.Module()])
-
-    def allgather(value, world_size):
-        events.append(("allgather", value, world_size))
-        return [value] * world_size
-
-    shell_cls = modeling.ModelingV2Glm52Nvfp4Sm103Dep
-    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__init__", builtin_init)
-    monkeypatch.setattr(modeling, "adopt_model", adopt_model)
-    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__post_init__", lambda self: None)
-    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__pp_init__", lambda self: None)
-    monkeypatch.setattr(dispatch, "_require", [False])
-    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
-    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
-    monkeypatch.setattr(dispatch, "_TARGET", [dispatch.DEFAULT_TARGET])
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: events.append(message))
-    monkeypatch.setattr(modeling.logger, "warning", lambda message: events.append(message))
-    for name in list(os.environ):
-        if name.startswith(dispatch.KNOB_PREFIX) or name in RANK_UNIFORM_ENV:
-            monkeypatch.delenv(name)
-
-    def construct(mode, model_config):
-        monkeypatch.setitem(_router_index._RESOLVED_MODES, shell_cls, mode)
-        return shell_cls(model_config)
-
-    return construct, events
-
-
-@pytest.mark.parametrize("cell", CELLS, ids=IDS)
-def test_the_shell_checks_the_cell_the_ranks_and_the_plan_then_adopts(shell_construction, cell):
-    construct, events = shell_construction
-    shell = construct(REQUIRE, _model_config(cell))
-
-    allgather, record, built_in, installed, adopt, adopted = events
-    assert allgather == ("allgather", ((dispatch.MODE_KEY, "require"),), cell[0])
-    served = modeling.CELLS[cell].paths.served
-    assert shell.plan.describe().startswith(
-        f"modeling_v2=require, {len(served)}/{len(served)} fast paths"
-    )
-    assert record.startswith("[glm_moe_dsa] dispatch: record paths=")
-    assert built_in == "built-in model"
-    assert installed == ("installed", served)
-    assert adopt == ("adopt", "GlmMoeDsaDepModelingV2", shell.plan)
-    assert adopted == ("allgather", modeling.adopted_state(shell), cell[0])
-    assert shell.plan.target is modeling.CELLS[cell].paths
-    assert type(shell.model) is modeling.GlmMoeDsaDepModelingV2
-    assert "_runtime_hooks" not in vars(shell)
-
-
-@pytest.mark.parametrize(
-    "cell, facts",
-    [
-        (DEP16, {"moe_backend": "CUTEDSL"}),
-        (DEP8, {"kv_cache": QuantAlgo.NVFP4}),
-        (DEP8_LM_HEAD_TP, {"kv_cache": QuantAlgo.FP8}),
-    ],
-    ids=IDS,
-)
-def test_outside_the_cells_require_raises_before_any_collective(shell_construction, cell, facts):
-    construct, events = shell_construction
-    with pytest.raises(ValueError, match="is not a deployment this target serves"):
-        construct(REQUIRE, _model_config(cell, **facts))
-    assert events == []
-
-
-def test_outside_the_cells_auto_builds_the_builtin_model_without_fast_paths(shell_construction):
-    construct, events = shell_construction
-    shell = construct(AUTO, _model_config(DEP8, moe_backend="TRTLLM"))
-
-    warning, allgather, record = events[:3]
-    assert "is not a deployment this target serves" in warning
-    assert allgather == ("allgather", ((dispatch.MODE_KEY, "auto"),), 8)
-    assert record.startswith("[glm_moe_dsa] dispatch: record paths=<none> ")
-    assert shell.plan.describe().startswith("modeling_v2=auto, 0/0 fast paths")
-    assert ("installed", frozenset()) in events
-    assert shell.plan.target is modeling.NO_PATHS
-    assert not dispatch.require_mode()
-
-
-def test_a_disable_knob_turns_its_path_and_dependents_off(monkeypatch, shell_construction):
-    construct, _ = shell_construction
-    monkeypatch.setenv(dispatch.DISABLE_ENV, "draft_loop")
-    shell = construct(REQUIRE, _model_config(DEP8))
-    off = {name for name in dispatch.FAST_PATH_NAMES if not dispatch.is_enabled(name)}
-    assert off >= {"draft_loop", "decode_prologue"}
-    assert dispatch.installed() == modeling.CELLS[DEP8].paths.served - {
-        "draft_loop",
-        "decode_prologue",
-    }
-    assert shell.plan.describe().startswith("modeling_v2=require, 8/10 fast paths")
-
-
-# --- the step-loop hooks --------------------------------------------------------------------
-
-
-def _model_of_cell(cell):
-    paths = modeling.CELLS[cell].paths
-    return types.SimpleNamespace(
-        plan=dispatch.resolve(REQUIRE, {}, paths),
-        model=types.SimpleNamespace(decode_prologue=None),
-        model_config=types.SimpleNamespace(mapping=types.SimpleNamespace(enable_attention_dp=True)),
-    )
-
-
-@pytest.fixture
-def quiet_dispatch(monkeypatch):
-    monkeypatch.setattr(dispatch, "_require", [False])
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-
-
-@pytest.mark.usefixtures("quiet_dispatch")
-@pytest.mark.parametrize("cell", CELLS, ids=IDS)
-def test_the_hooks_declare_the_decode_prologue_and_no_step_loop_part(cell):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E501
-        DecodePrologueHook,
-    )
-
-    model = _model_of_cell(cell)
-    hooks = runtime.make_runtime_hooks(model)
-    assert isinstance(hooks, RuntimeHooks)
-    assert not hooks.loop_on_execution_stream
-    assert hooks.step_input_stager is None and hooks.sampler_store_writer is None
-    assert isinstance(hooks.decode_prologue, DecodePrologueHook)
-    assert model.model.decode_prologue is hooks.decode_prologue
-    assert isinstance(hooks.executor_events, runtime.GlmExecutorEvents)
-
-
-@pytest.mark.usefixtures("quiet_dispatch")
-@pytest.mark.parametrize("cell", CELLS, ids=IDS)
-def test_the_executor_events_install_only_the_host_policy(monkeypatch, cell):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
-        host_page_retention,
-    )
-
-    calls = []
-    monkeypatch.setattr(host_page_retention, "install_host_policy", calls.append)
-    events = runtime.make_runtime_hooks(_model_of_cell(cell)).executor_events
-    events.after_warmup()
-    events.before_overlap_loop()
-    events.at_shutdown()
-    assert calls == ["after_warmup"]
-
-
-def test_the_shell_builds_its_hooks_on_the_first_read(monkeypatch, shell_construction):
-    construct, _ = shell_construction
-    shell = construct(REQUIRE, _model_config(DEP8))
-    built = []
-
-    def make_runtime_hooks(model):
-        built.append(model)
-        return RuntimeHooks()
-
-    monkeypatch.setattr(runtime, "make_runtime_hooks", make_runtime_hooks)
-    hooks = runtime_hooks_of(shell)
-    assert built == [shell] and runtime_hooks_of(shell) is hooks
-
-    def broken(model):
-        raise AttributeError("plan")
-
-    shell2 = construct(REQUIRE, _model_config(DEP8))
-    monkeypatch.setattr(runtime, "make_runtime_hooks", broken)
-    with pytest.raises(RuntimeError, match="step-loop hooks"):
-        runtime_hooks_of(shell2)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py
deleted file mode 100644
index f22584d3245e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_layers.py
+++ /dev/null
@@ -1,328 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 tep target's adoption at each tensor-parallel size, the family's layer
-composition under its plans, its decode loop, and how its ``modeling.py`` differs from the tp8
-target's.
-
-Host-only, with the stand-in built-in modules and kernel factories of the tp8 target's layer
-test (``test_modeling_v2_glm_layers.py``, loaded by path): the fast paths the stand-ins see are
-the ones the tep plan serves at that size.
-"""
-
-from __future__ import annotations
-
-import ast
-import importlib.util
-import inspect
-import types
-from pathlib import Path
-
-import pytest
-import torch
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tep import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
-    modeling as tp8_modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    attention,
-    dispatch,
-    layers,
-    spec_worker,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-
-def _load_tp8_layer_test():
-    path = (
-        Path(__file__).resolve().parents[1]
-        / "glm_5_2_nvfp4__sm_103__tp8"
-        / "test_modeling_v2_glm_layers.py"
-    )
-    spec = importlib.util.spec_from_file_location("_glm_tp8_layer_test", path)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-T = _load_tp8_layer_test()
-
-
-def _served(tp_size):
-    return modeling.TARGET_PATHS[tp_size].served
-
-
-@pytest.fixture
-def fast_paths(monkeypatch):
-    def install(enabled, post_attn_hop=True):
-        paths = T._FastPaths(frozenset(enabled))
-        if not post_attn_hop:
-            # As at four ranks: the call site's native all-reduce is not multi-node NVLink, so
-            # the hop factory returns None.
-            paths.entries[("post_attention_hop", "make_post_attn_hop")] = lambda *a, **k: None
-        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-        monkeypatch.setattr(dispatch, "log_capability", paths.log_capability)
-        return paths
-
-    return install
-
-
-def _state(model):
-    return dict(modeling.adopted_state(types.SimpleNamespace(model=model.model)))
-
-
-def test_at_eight_ranks_adoption_builds_the_tp8_objects_but_the_routed_moe_ones(fast_paths):
-    paths = fast_paths(_served(8))
-    model = T._built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-
-    tail = [("make", "attn_tail"), ("make", "staged_attn_tail")]
-    pre_moe = tail + [("make", "post_attn_hop"), ("make", "into_hop")]
-    post_moe = [("make", "router_gemm"), ("make", "post_moe_hop")]
-    expected = (
-        pre_moe
-        + pre_moe
-        + post_moe
-        + [("make", "deferred_norm_evict_hints_hop"), ("deferred hop serves", 2, "model config")]
-        + pre_moe
-        + post_moe
-        + [
-            ("make", "post_attn_hop"),
-            ("make", "router_gemm"),
-            ("log_capability", "draft_front", "draft_front", True, ""),
-            ("make", "draft_live_row_hop"),
-        ]
-    )
-    assert [e for e in paths.events if e[0] != "set_successor_weights"] == expected
-    dense, moe, last, draft = model.model.layers
-    assert [type(m) for m in (dense, moe, last)] == [layers.GlmDecoderLayer] * 3
-    assert type(draft) is layers.GlmMtpLayer
-    assert [type(m.mlp) for m in (moe, last, draft)] == [layers.GlmMoE] * 3
-    assert type(moe.self_attn) is attention.GlmMLA
-    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
-    state = _state(model)
-    assert [entry[0] for entry in state["model.layers.1"]] == [
-        "post_attn_hop",
-        "_into_hop",
-        "post_moe_hop",
-        "deferred_norm_hop",
-        "self_attn._attn_tail",
-        "mlp.router_gemm",
-    ]
-    assert [entry[0] for entry in state["model.layers.3"]] == [
-        "post_attn_hop",
-        "draft_front",
-        "draft_live_row_hop",
-        "mlp.router_gemm",
-    ]
-
-
-def test_at_four_ranks_adoption_builds_the_router_gemm_and_the_draft_front(fast_paths):
-    paths = fast_paths(_served(4), post_attn_hop=False)
-    model = T._built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-
-    assert [e for e in paths.events if e[0] == "make"] == [("make", "router_gemm")] * 3
-    state = _state(model)
-    assert [entry[0] for entry in state["model.layers.0"]] == []
-    assert [entry[0] for entry in state["model.layers.1"]] == ["mlp.router_gemm"]
-    assert [entry[0] for entry in state["model.layers.3"]] == ["draft_front", "mlp.router_gemm"]
-    for layer in model.model.layers:
-        assert layer.post_attn_hop is None and layer._into_hop is None
-        assert layer.post_moe_hop is None and layer.deferred_norm_hop is None
-        assert layer.self_attn._attn_tail is None
-
-
-# The fast paths only the tp8 target serves: their kernels take tensor-parallel experts. The
-# objects only they attach, as (owner submodule, attribute, path).
-_ROUTED_MOE_PATHS = frozenset(
-    {
-        "routed_moe",
-        "routed_moe_into_hop",
-        "hop_routing",
-        "shared_expert_prefetch",
-        "draft_routed_moe",
-        "cross_layer_expert_prefetch",
-    }
-)
-_ROUTED_MOE_STATE = (
-    ("", "_hop_routing", "hop_routing"),
-    ("mlp", "_routed_into_hop", "routed_moe_into_hop"),
-)
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_adoption_attaches_no_routed_moe_object(fast_paths, tp_size):
-    paths = fast_paths(_served(tp_size), post_attn_hop=tp_size == 8)
-    model = T._built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-    for layer in model.model.layers:
-        for submodule, name, _ in _ROUTED_MOE_STATE:
-            owner = getattr(layer, submodule) if submodule else layer
-            assert getattr(owner, name, None) is None, (submodule, name)
-        assert layer._expert_prefetch_link is None
-
-
-def _enclosing_tests(func, target):
-    """The tests of the ``if`` statements whose body contains ``target``."""
-    found = []
-
-    def visit(node, stack):
-        if node is target:
-            found.extend(stack)
-            return True
-        if isinstance(node, ast.If):
-            return any(visit(c, stack + [ast.unparse(node.test)]) for c in node.body) or any(
-                visit(c, stack) for c in node.orelse
-            )
-        return any(visit(c, stack) for c in ast.iter_child_nodes(node))
-
-    visit(func, [])
-    return found
-
-
-def _paths_named(tree):
-    """Every fast path a ``dispatch.is_enabled`` / ``dispatch.get_impl`` call names."""
-    return {
-        call.args[0].value
-        for call in ast.walk(tree)
-        if isinstance(call, ast.Call)
-        and ast.unparse(call.func) in ("dispatch.is_enabled", "dispatch.get_impl")
-        and call.args
-        and isinstance(call.args[0], ast.Constant)
-    }
-
-
-def test_every_path_the_composition_names_is_served_or_routed_moe():
-    """The family's layers and this target's ``modeling.py`` name, besides the paths the target
-    serves at some size, only routed-MoE paths, which it serves at no size: their blocks stay
-    off under both plans."""
-    served = _served(4) | _served(8)
-    assert not served & _ROUTED_MOE_PATHS
-    for module in (layers, modeling):
-        named = _paths_named(ast.parse(Path(module.__file__).read_text()))
-        assert named <= set(dispatch.FAST_PATH_NAMES), module.__name__
-        unserved = named - served
-        assert unserved <= _ROUTED_MOE_PATHS, (module.__name__, sorted(unserved))
-    named = _paths_named(ast.parse(Path(layers.__file__).read_text()))
-    assert named & _ROUTED_MOE_PATHS == _ROUTED_MOE_PATHS - {"cross_layer_expert_prefetch"}
-
-
-def test_the_routed_moe_objects_are_attached_only_under_their_paths():
-    """``_hop_routing`` and the MoE's ``_routed_into_hop`` are set only inside a block gated on
-    their path, so under a plan without it they keep their class default None."""
-    tree = ast.parse(Path(layers.__file__).read_text())
-    seen = set()
-    for func in ast.walk(tree):
-        if not isinstance(func, ast.FunctionDef):
-            continue
-        for node in ast.walk(func):
-            if not isinstance(node, ast.Assign):
-                continue
-            for target in node.targets:
-                for submodule, name, path in _ROUTED_MOE_STATE:
-                    owner = f"self.{submodule}" if submodule else "self"
-                    if ast.unparse(target) == f"{owner}.{name}":
-                        seen.add(name)
-                        assert any(
-                            f"dispatch.is_enabled({path!r})" in test
-                            for test in _enclosing_tests(func, node)
-                        ), (func.name, name)
-    assert seen == {name for _, name, _ in _ROUTED_MOE_STATE}
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_the_moe_runs_the_builtin_experts_on_the_produced_logits(fast_paths, monkeypatch, tp_size):
-    fast_paths(_served(tp_size))
-    calls = []
-    moe = T._moe(calls)
-    moe.__class__ = layers.GlmMoE
-    moe.router_gemm = lambda hidden: "router gemm logits"
-    monkeypatch.setattr(dispatch, "phase", lambda: dispatch.Phase.DECODE)
-
-    hidden = torch.zeros(2, T.HIDDEN)
-    assert moe.compute_routed_output(hidden, None, [2], True) == "routed"
-    assert calls[-1][0][1] == "router gemm logits"
-    assert moe.compute_routed_output(hidden, None, [2], False, "hop logits") == "routed"
-    assert calls[-1][0][1] == "hop logits" and calls[-1][1]["do_finalize"] is False
-    moe.router_gemm = None
-    moe.compute_routed_output(hidden, None, [2], True)
-    assert calls[-1][0][1] == "gate logits"
-
-
-def test_the_decode_loop_has_no_expert_prefetch_join():
-    source = inspect.getsource(modeling.DecodeTarget.forward)
-    assert "cross_layer_expert_prefetch" not in source
-    assert "setup_aliases" not in vars(modeling.ModelingV2Glm52Nvfp4Sm103Tep)
-
-
-def _function_bodies(module):
-    """Every function of ``module`` as (qualified name, AST dump without docstrings)."""
-    tree = ast.parse(Path(module.__file__).read_text())
-    out = {}
-
-    def strip(node):
-        for child in ast.walk(node):
-            body = getattr(child, "body", None)
-            if (
-                isinstance(body, list)
-                and body
-                and isinstance(body[0], ast.Expr)
-                and isinstance(getattr(body[0], "value", None), ast.Constant)
-                and isinstance(body[0].value.value, str)
-            ):
-                child.body = body[1:] or [ast.Pass()]
-        return node
-
-    for cls in tree.body:
-        if isinstance(cls, ast.ClassDef):
-            for fn in cls.body:
-                if isinstance(fn, ast.FunctionDef):
-                    out[f"{cls.name}.{fn.name}"] = ast.dump(strip(fn))
-        elif isinstance(cls, ast.FunctionDef):
-            out[cls.name] = ast.dump(strip(cls))
-    return out
-
-
-# The tp8 target's class each tep class corresponds to, and the functions the tep target
-# changes: its plan per tensor-parallel size, and no cross-layer expert prefetch (a routed-MoE
-# path) in its decode loop and shell.
-_RENAMES = {
-    "GlmMoeDsaTepModelingV2": "GlmMoeDsaTp8ModelingV2",
-    "ModelingV2Glm52Nvfp4Sm103Tep": "ModelingV2Glm52Nvfp4Sm103Tp8",
-}
-_CHANGED = {"DecodeTarget.forward", "ModelingV2Glm52Nvfp4Sm103Tp8.__init__"}
-
-
-def test_the_targets_modeling_differs_only_in_the_declared_functions():
-    """Every function of the tep ``modeling.py`` equals the tp8 one's (with the tep class
-    names), except exactly the declared ones; the tp8 module has only the declared extra
-    function. Both adopt their layers with the family's ``adopt_model``."""
-    ours = {}
-    for name, body in _function_bodies(modeling).items():
-        for new, old in _RENAMES.items():
-            name = name.replace(new, old)
-            body = body.replace(f"'{new}'", f"'{old}'")
-        ours[name] = body
-    theirs = _function_bodies(tp8_modeling)
-    differ = {name for name in ours if name in theirs and ours[name] != theirs[name]}
-    assert differ == _CHANGED, (sorted(differ - _CHANGED), sorted(_CHANGED - differ))
-    only_tp8 = set(theirs) - set(ours)
-    assert only_tp8 <= {"ModelingV2Glm52Nvfp4Sm103Tp8.setup_aliases"}, sorted(only_tp8)
-    assert set(ours) - set(theirs) == set()
-    assert modeling.adopt_model is tp8_modeling.adopt_model is layers.adopt_model
-    assert modeling.adopted_state is tp8_modeling.adopted_state is layers.adopted_state
-
-
-def test_the_classes_add_no_parameter_buffer_or_module(fast_paths):
-    paths = fast_paths(_served(8))
-    model = T._built_in_model(paths.events)
-    before = [T._keys(layer) for layer in model.model.layers]
-    modeling.adopt_model(model, plan=None)
-    assert [T._keys(layer) for layer in model.model.layers] == before
-    assert all(isinstance(layer, nn.Module) for layer in model.model.layers)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py
deleted file mode 100644
index 9fea5839cfd3..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/test_modeling_v2_glm_tep_plan.py
+++ /dev/null
@@ -1,355 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The fast paths the GLM-5.2 tep target serves at each tensor-parallel size, and what its
-plan does with the declines its deployments log.
-
-Host-only: the plan is resolved and installed from the target's ``TARGET_PATHS`` exactly as
-the shell does at construction; the declines are the (path, key, reason) triples the kernels
-log at those deployments, decided at construction or in the warm-up, and none of them may be
-a miss: under ``modeling_v2='require'`` a miss raises.
-"""
-
-from __future__ import annotations
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tep import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
-    modeling as tp8_modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-D = dispatch.Decline
-REQUIRE = ModelingV2Mode.REQUIRE
-
-# The paths whose kernels take only tensor-parallel experts (every expert on every rank).
-ROUTED_MOE = {
-    "routed_moe",
-    "routed_moe_into_hop",
-    "hop_routing",
-    "shared_expert_prefetch",
-    "draft_routed_moe",
-    "cross_layer_expert_prefetch",
-}
-# At four ranks: the paths that engage there, and the bases they extend.
-TEP4_ENGAGED = {
-    "indexer_layer_front",
-    "rope_kv_append",
-    "indexer_topk",
-    "cross_rank_topk",
-    "cross_rank_topk_cluster_merge",
-    "cross_rank_topk_global_emit",
-    "router_gemm",
-    "dense_mlp",
-    "draft_front",
-    "draft_tail",
-    "draft_loop",
-    "decode_prologue",
-    "deferred_input_prep",
-    "sampler_store",
-    "execution_stream_loop",
-    "host_page_retention",
-}
-TEP4_BASES = {"absorbed_query", "draft_absorbed_query", "post_attention_hop"}
-SERVED = {
-    4: TEP4_ENGAGED | TEP4_BASES,
-    8: set(dispatch.FAST_PATH_NAMES) - ROUTED_MOE,
-}
-
-
-def _layers(spec: str) -> list:
-    out = []
-    for part in spec.split(","):
-        lo, _, hi = part.partition("-")
-        out += list(range(int(lo), int(hi or lo) + 1))
-    return out
-
-
-FULL = "0-2,6,10,14,18,22,26,30,34,38,42,46,50,54,58,62,66,70,74"
-SHARED = (
-    "3-5,7-9,11-13,15-17,19-21,23-25,27-29,31-33,35-37,39-41,43-45,47-49,51-53,55-57,59-61,"
-    "63-65,67-69,71-73,75-77"
-)
-TP8_DOMAIN = "shape outside the GLM-5.2 TP8 domain"
-
-# The declines the served paths log at four ranks (bs 2 and bs 5), on every rank.
-DECLINES_TP4 = (
-    *(("absorbed_query", f"layer{i}", TP8_DOMAIN) for i in range(78)),
-    ("absorbed_query", "layer78", "target layers only"),
-    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
-    ("draft_absorbed_query", "layer78", TP8_DOMAIN),
-    *(("indexer_layer_front", f"target_layer{i}", TP8_DOMAIN) for i in range(78)),
-    ("indexer_layer_front", "draft_layer78", TP8_DOMAIN),
-    (
-        "post_attention_hop",
-        "post_attn_hop:native all-reduce (AUTO) does not use MNNVL",
-        "native all-reduce (AUTO) does not use MNNVL",
-    ),
-    ("router_gemm", "router:not a target layer", "not a target layer"),
-    ("indexer_topk", "indexer_decode", "MTP draft step (target layers only)"),
-    ("indexer_topk", "indexer_decode", "5 generation requests > 4"),
-    ("cross_rank_topk", "cp_indexer_decode", "2 generation requests != 1"),
-    ("cross_rank_topk_cluster_merge", "cp_indexer_decode", "4 generation requests != 1"),
-    ("cross_rank_topk_global_emit", "cp_indexer_decode", "5 generation requests > 4"),
-    ("dense_mlp", "dense_mlp:row count outside [1, 8]", "row count outside [1, 8]"),
-    ("draft_loop", "draft_loop:context rows present", "context rows present"),
-    ("draft_loop", "draft_step:dense top-k fallback active", "dense top-k fallback active"),
-    (
-        "deferred_input_prep",
-        "prep[not a steady decode-graph replay]",
-        "not a steady decode-graph replay",
-    ),
-    (
-        "deferred_input_prep",
-        "prep[no previous speculative device tensors]",
-        "no previous speculative device tensors",
-    ),
-    (
-        "deferred_input_prep",
-        "prep[dummy row or row without previous tensors]",
-        "dummy row or row without previous tensors",
-    ),
-    (
-        "decode_prologue",
-        "target_prologue[not a generation-only step]",
-        "not a generation-only step",
-    ),
-)
-
-# The declines the served paths log at eight ranks (bs 5), on every rank.
-DECLINES_TP8 = (
-    ("absorbed_query", "layer78", "target layers only"),
-    ("attention_output_tail", "layer78", "draft or unindexed layer"),
-    *(
-        (
-            "shared_indexer_absorbed_query",
-            f"layer{i}",
-            "full indexer layer (shared_indexer_absorbed_query serves shared layers)",
-        )
-        for i in _layers(FULL)
-    ),
-    ("shared_indexer_absorbed_query", "layer78", "target layers only"),
-    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
-    *(
-        (
-            "deferred_norm",
-            f"scope[shared]_into_layer{i}",
-            "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE=shared: layer "
-            f"{i} is a full-indexer layer; post_moe_hop hop",
-        )
-        for i in _layers(FULL)[3:]
-    ),
-    *(
-        (
-            "indexer_layer_front",
-            f"target_layer{i}",
-            "shared-indexer layer (no q_a RMSNorm launch to remove; "
-            "shared_indexer_absorbed_query / absorbed_query serve it)",
-        )
-        for i in _layers(SHARED)
-    ),
-    *(
-        (
-            "fused_kv_indexer_projection",
-            f"layer{i}",
-            "full layers not selected (TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS)",
-        )
-        for i in _layers(FULL)
-    ),
-    (
-        "fused_kv_indexer_projection",
-        "layer78",
-        "draft layers not selected (TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS)",
-    ),
-    ("router_gemm", "router:not a target layer", "not a target layer"),
-    (
-        "sparse_mla_decode",
-        "request:not generation-only",
-        "not generation-only; first seen at layer 3 with 1 rows",
-    ),
-    ("indexer_topk", "indexer_decode", "MTP draft step (target layers only)"),
-    ("indexer_topk", "indexer_decode", "5 generation requests > 4"),
-    ("cross_rank_topk", "cp_indexer_decode", "2 generation requests != 1"),
-    ("cross_rank_topk_global_emit", "cp_indexer_decode", "5 generation requests > 4"),
-    ("dense_mlp", "dense_mlp:row count outside [1, 8]", "row count outside [1, 8]"),
-    ("draft_loop", "draft_loop:context rows present", "context rows present"),
-    (
-        "deferred_input_prep",
-        "prep[not a steady decode-graph replay]",
-        "not a steady decode-graph replay",
-    ),
-    (
-        "decode_prologue",
-        "target_prologue[not a generation-only step]",
-        "not a generation-only step",
-    ),
-)
-DECLINES = {4: DECLINES_TP4, 8: DECLINES_TP8}
-
-# The layer-keyed capabilities that engage at eight ranks: (path, key family, layers).
-ENGAGED_TP8 = (
-    ("absorbed_query", "layer", "0-77"),
-    ("attention_output_tail", "layer", "0-77"),
-    ("oproj_weight_staging", "layer", "0-77"),
-    ("shared_indexer_absorbed_query", "layer", SHARED),
-    ("draft_absorbed_query", "layer", "78"),
-    ("sparse_mla_decode", "layer", "0-78"),
-    ("sparse_mla_decode", "call:layer", "0-78"),
-    ("fused_kv_indexer_projection", "layer", "3"),
-    ("indexer_layer_front", "target_layer", FULL),
-    ("indexer_layer_front", "draft_layer", "78"),
-    ("deferred_norm", "scope[shared]_into_layer", SHARED.replace("3-5", "4-5")),
-    ("deferred_norm_evict_hints", "projection_layer", SHARED.replace("3-5", "4-5")),
-)
-
-
-@pytest.fixture
-def installed(monkeypatch):
-    """Resolve and install the target's plan at a tensor-parallel size, as the shell does."""
-    for name in ("_INSTALLED", "_VALIDATE", "_require", "_TARGET"):
-        monkeypatch.setattr(dispatch, name, list(getattr(dispatch, name)))
-    monkeypatch.setattr(dispatch, "_logged", set())
-    monkeypatch.setattr(dispatch, "_missed", {})
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-    dispatch.classify_decline.cache_clear()
-
-    def install(tp_size, environ=None):
-        plan = dispatch.resolve(REQUIRE, environ or {}, modeling.TARGET_PATHS[tp_size])
-        dispatch.install(plan, environ or {})
-        return plan
-
-    yield install
-    dispatch.classify_decline.cache_clear()
-
-
-def test_the_target_has_a_table_per_tensor_parallel_size():
-    assert sorted(modeling.TARGET_PATHS) == [4, 8]
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_the_target_serves_its_paths_at_each_size(installed, tp_size):
-    plan = installed(tp_size)
-    assert set(plan.served) == SERVED[tp_size]
-    assert dispatch.installed() == frozenset(SERVED[tp_size])
-    assert plan.describe().startswith(
-        f"modeling_v2=require, {len(SERVED[tp_size])}/{len(SERVED[tp_size])} fast paths"
-    )
-
-
-def test_the_sizes_serve_nineteen_and_thirty_paths():
-    assert len(SERVED[4]) == 19 and len(SERVED[8]) == 30
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_no_served_path_extends_one_the_target_does_not_serve(tp_size):
-    served = SERVED[tp_size]
-    for path in served:
-        assert dispatch.REQUIRED_FEATURES.get(path, frozenset()) <= served, path
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_the_target_never_serves_a_routed_moe_path(tp_size):
-    assert not (SERVED[tp_size] & ROUTED_MOE)
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_every_kernel_package_is_the_familys(tp_size):
-    """The target has no kernel package of its own: every served path imports from shared."""
-    assert dict(modeling.TARGET_PATHS[tp_size].kernel_modules) == {}
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_the_record_line_lists_the_served_paths(installed, monkeypatch, tp_size):
-    lines = []
-    monkeypatch.setattr(dispatch, "_log_info", lines.append)
-    installed(tp_size)
-    (record,) = [line for line in lines if "record" in line]
-    listed = record.split("paths=", 1)[1].split(" ", 1)[0].split(",")
-    assert {name.split("[", 1)[0] for name in listed} == SERVED[tp_size]
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_disable_may_name_only_a_served_path(installed, tp_size):
-    with pytest.raises(ValueError, match="not fast paths of this target"):
-        installed(tp_size, {dispatch.DISABLE_ENV: "routed_moe"})
-    plan = installed(tp_size, {dispatch.DISABLE_ENV: "draft_tail"})
-    assert not plan.enabled("draft_tail") and plan.enabled("draft_loop")
-
-
-def test_disable_turns_the_dependents_off_at_four_ranks(installed):
-    plan = installed(4, {dispatch.DISABLE_ENV: "absorbed_query"})
-    assert not plan.enabled("indexer_layer_front")
-    assert plan.describe().startswith("modeling_v2=require, 17/19 fast paths")
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_the_routed_accumulation_knob_is_refused(installed, tp_size):
-    with pytest.raises(ValueError, match="does not serve routed_moe_into_hop"):
-        installed(tp_size, {dispatch.ACCUMULATION_ENV: "identity"})
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_the_cross_rank_wait_timeout_is_read(installed, tp_size):
-    plan = installed(tp_size, {dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV: "60000"})
-    assert plan.knobs == ((dispatch.CROSS_RANK_TOPK_TIMEOUT_ENV, "60000"),)
-
-
-@pytest.mark.parametrize("tp_size", [4, 8])
-def test_no_logged_decline_is_a_miss(installed, tp_size):
-    """Under require, a decline decided once inside a path's cells would raise."""
-    installed(tp_size)
-    assert dispatch.require_mode()
-    for path, key, reason in DECLINES[tp_size]:
-        assert path in SERVED[tp_size], path
-        assert dispatch.classify_decline(path, key, reason) is not D.MISSED, (path, key, reason)
-        assert dispatch.log_capability(path, key, False, reason) is False
-    assert all(count >= 0 for count in dispatch.missed_counts().values())
-
-
-def test_at_four_ranks_the_bases_decline_as_uncovered(installed):
-    installed(4)
-    for path, key, reason in DECLINES_TP4:
-        if path in TEP4_BASES or key.startswith(("target_layer", "draft_layer")):
-            assert dispatch.classify_decline(path, key, reason) is D.UNCOVERED, (path, key)
-
-
-def test_the_four_rank_tables_would_raise_at_the_tp8_cells(installed):
-    """The tp8 cells hold the query fronts and the hop that four ranks decline on: the tep
-    table is what keeps a four-rank deployment from raising under require."""
-    installed(8)
-    misses = [
-        (path, key)
-        for path, key, reason in DECLINES_TP4
-        if dispatch.classify_decline(path, key, reason) is D.MISSED
-    ]
-    assert ("absorbed_query", "layer0") in misses
-    assert ("post_attention_hop", "post_attn_hop:native all-reduce (AUTO) does not use MNNVL") in (
-        misses
-    )
-
-
-def test_at_eight_ranks_every_engaged_layer_is_inside_the_cells(installed):
-    target = modeling.TARGET_PATHS[8]
-    for path, family, layers in ENGAGED_TP8:
-        cells = target.layer_cells[path][family]
-        assert set(_layers(layers)) <= cells, (path, family)
-
-
-def test_the_tp8_target_serves_every_path_with_the_family_tables():
-    target = tp8_modeling.TARGET_PATHS
-    assert target.served == frozenset(dispatch.FAST_PATH_NAMES)
-    assert target.layer_cells is dispatch.LAYER_CELLS
-    assert dict(target.uncovered) == {}
-    assert {path for path in target.kernel_modules} == ROUTED_MOE
-    assert all(
-        module.startswith(tp8_modeling.__package__ + ".kernels.")
-        for module in target.kernel_modules.values()
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py
deleted file mode 100644
index 724c6c33a153..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch.py
+++ /dev/null
@@ -1,267 +0,0 @@
-# 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.
-"""Tests of package cross_layer_expert_prefetch (cross-layer speculative routed-expert L2 prefetch)."""
-
-import pytest
-import torch
-from torch import nn
-
-cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-NE = 256
-HIDDEN = 6144
-KMAX = 32
-STAT_KS = (8, 12, 18, 24, 32)
-
-
-def _ref_union(logits, bias, rows):
-    s = torch.sigmoid(logits[:rows].float()) + bias.float()[None, :]
-    thr = s.sort(dim=1, descending=True).values[:, 7:8]
-    sel = s >= thr
-    return sel.sum(0).to(torch.int32), s.max(0).values
-
-
-def _ref_rank(logits, bias, rows):
-    cnt, mx = _ref_union(logits, bias, rows)
-    order = sorted(range(NE), key=lambda e: (-int(cnt[e]), -float(mx[e]), e))
-    return order[:KMAX]
-
-
-@pytest.fixture()
-def spec():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
-        spec_experts,
-    )
-
-    return spec_experts
-
-
-def test_the_served_config(spec):
-    cfg = spec.config()
-    assert cfg is spec.SERVED_CONFIG
-    assert (
-        cfg.k,
-        cfg.start,
-        cfg.pred,
-        cfg.fused,
-        cfg.delay_ns,
-        cfg.chunk,
-        cfg.ctas,
-        cfg.issue,
-        cfg.xc_ctas,
-    ) == (8, "moe", "self", True, 0, 16384, 16, "xc", 32)
-    assert cfg.start in spec.STARTS and cfg.pred in spec.PREDS and cfg.issue in spec.ISSUES
-    assert spec.ExpertPrefetchConfig(8, "moe", "self", True, 0, 16384, 16) == cfg
-
-
-@cuda
-@pytest.mark.parametrize("rows,ties", [(1, False), (6, False), (8, False), (1, True), (6, True)])
-def test_rank_and_stats_match_reference(rows, ties):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
-        kernels,
-    )
-
-    g = torch.Generator(device="cuda").manual_seed(rows)
-    dev = "cuda"
-    nxt = torch.randn(8, NE, device=dev, generator=g) * 3
-    nbias = (torch.randn(NE, device=dev, generator=g) * 0.5).bfloat16()
-    cur = torch.randn(8, NE, device=dev, generator=g) * 3
-    cbias = (torch.randn(NE, device=dev, generator=g) * 0.5).bfloat16()
-    if ties:
-        # 13 integer logits and no bias: many equal scores within and across
-        # rows (resolved by ID), distinct values far apart
-        nxt = nxt.round().clamp(-6, 6)
-        cur = cur.round().clamp(-6, 6)
-        nbias.zero_()
-        cbias.zero_()
-    cur_ids = torch.randperm(NE, device=dev, generator=g)[:KMAX].to(torch.int32)
-    out_ids = torch.zeros(KMAX, dtype=torch.int32, device=dev)
-    counters = torch.zeros(kernels.NUM_COUNTERS, dtype=torch.int64, device=dev)
-    kernels.rank_and_stats(nxt, nbias, out_ids, cur, cbias, cur_ids, counters, rows, True, True)
-    torch.cuda.synchronize()
-    assert out_ids.tolist() == _ref_rank(nxt, nbias, rows)
-    cnt, _ = _ref_union(cur, cbias, rows)
-    actual = cnt > 0
-    pos = {int(e): j for j, e in enumerate(cur_ids.tolist())}
-    exp = [1, int(actual.sum())]
-    for k in STAT_KS:
-        exp.append(sum(1 for e in range(NE) if actual[e] and pos.get(e, KMAX) < k))
-    assert counters.tolist() == exp
-
-
-def _exact_weight(g):
-    """Gate weights in multiples of 1/256: with hidden values in {-1, 0, 1} every
-    logit is exact in FP32 whatever the summation order."""
-    return (torch.randint(-2, 3, (NE, HIDDEN), device="cuda", generator=g) / 256).bfloat16()
-
-
-def _exact_hidden(rows, g):
-    return torch.randint(-1, 2, (rows, HIDDEN), device="cuda", generator=g).bfloat16()
-
-
-@cuda
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_gemm_rank_matches_reference(rows):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
-        kernels,
-    )
-
-    g = torch.Generator(device="cuda").manual_seed(100 + rows)
-    h = _exact_hidden(rows, g)
-    w = _exact_weight(g)
-    bias = (torch.randn(NE, device="cuda", generator=g) * 0.1).bfloat16()
-    scores = torch.zeros(kernels.MAX_ROWS, NE, device="cuda")
-    ticket = torch.zeros(1, dtype=torch.int32, device="cuda")
-    out = torch.zeros(KMAX, dtype=torch.int32, device="cuda")
-    ref = _ref_rank(h.float() @ w.float().t(), bias, rows)
-    for _ in range(3):  # every launch resets the ticket for the next one
-        out.zero_()
-        kernels.gemm_rank(h, w, bias, scores, ticket, out, rows)
-        torch.cuda.synchronize()
-        assert out.tolist() == ref
-        assert ticket.item() == 0
-    graph = torch.cuda.CUDAGraph()
-    s = torch.cuda.Stream()
-    s.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(s):
-        with torch.cuda.graph(graph, stream=s):
-            kernels.gemm_rank(h, w, bias, scores, ticket, out, rows)
-    torch.cuda.current_stream().wait_stream(s)
-    out.zero_()
-    for _ in range(3):
-        graph.replay()
-    torch.cuda.synchronize()
-    assert out.tolist() == ref and ticket.item() == 0
-
-
-class _Gate(nn.Module):
-    def __init__(self, seed):
-        super().__init__()
-        g = torch.Generator(device="cuda").manual_seed(seed)
-        self.weight = nn.Parameter(_exact_weight(g), requires_grad=False)
-        self.e_score_correction_bias = nn.Parameter(
-            (torch.randn(NE, device="cuda", generator=g) * 0.1).bfloat16(), requires_grad=False
-        )
-
-    def forward(self, h):
-        return h.float() @ self.weight.float().t()
-
-
-class _Experts(nn.Module):
-    def __init__(self):
-        super().__init__()
-        u8 = dict(dtype=torch.uint8, device="cuda")
-        self.w3_w1_weight = nn.Parameter(
-            torch.zeros(NE, 512, HIDDEN // 2, **u8), requires_grad=False
-        )
-        self.w3_w1_weight_scale = nn.Parameter(
-            torch.zeros(NE, 512, HIDDEN // 16, **u8), requires_grad=False
-        )
-        self.w2_weight = nn.Parameter(torch.zeros(NE, HIDDEN, 128, **u8), requires_grad=False)
-        self.w2_weight_scale = nn.Parameter(torch.zeros(NE, HIDDEN, 16, **u8), requires_grad=False)
-
-
-class _Mlp(nn.Module):
-    def __init__(self, seed, experts):
-        super().__init__()
-        self.gate = _Gate(seed)
-        self.experts = experts
-
-
-class _Layer(nn.Module):
-    def __init__(self, mlp):
-        super().__init__()
-        self.mlp = mlp
-
-
-@cuda
-def test_link_layers_binds_routed_kernel_hooks(spec, monkeypatch):
-    """At the fixed depth, predictor, issuing kernel and CTA count, every MoE layer's routed
-    kernel prefetches the experts its own prediction names."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe,
-    )
-
-    monkeypatch.setattr(routed_moe, "resolve_expert_owner", lambda m: (m, ""))
-    monkeypatch.setattr(spec.dispatch, "is_enabled", lambda f: True)
-    spec._STATES.clear()
-    gates = [_Gate(i) for i in range(3)]
-    leaves = [nn.Module() for _ in range(3)]
-    layers = [nn.Identity()]
-    for gate, leaf in zip(gates, leaves):
-        mlp = nn.Module()
-        mlp.gate, mlp.experts = gate, leaf
-        layers.append(_Layer(mlp))
-    assert spec.link_layers(layers, 4) == 3
-    bound = [hasattr(leaf, "_expert_prefetch_spec") for leaf in leaves]
-    # Every layer's prediction feeds its own routed kernel.
-    assert bound == [True, True, True]
-    assert leaves[2]._expert_prefetch_spec() is None  # before the device state exists
-    rows = 6
-    h = torch.zeros(rows, HIDDEN, device="cuda", dtype=torch.bfloat16)
-    logits = torch.zeros(rows, NE, device="cuda", dtype=torch.float32)
-    for i in (1, 2, 3):
-        link = layers[i]._expert_prefetch_link
-        link.on_hop(h, logits)
-        link.after_moe()
-        link.after_post_moe()
-    spec.join()
-    torch.cuda.synchronize()
-    st = spec._STATES[torch.cuda.current_device()]
-    ids, k, ctas, self_store = leaves[2]._expert_prefetch_spec()
-    assert ids.data_ptr() == st.pred_ids[3].data_ptr() and (k, ctas) == (8, 32)
-    assert self_store
-    # the routed kernel issues the prefetch: no prefetch kernel is launched
-    assert st.prefetch_calls == 0
-    assert not st.forked and not st.counters.any()
-    # without routed_moe nothing is bound
-    spec._STATES.clear()
-    for leaf in leaves:
-        leaf.__dict__.pop("_expert_prefetch_spec", None)
-    monkeypatch.setattr(spec.dispatch, "is_enabled", lambda f: f != "routed_moe")
-    spec.link_layers(layers, 4)
-    assert not any(hasattr(leaf, "_expert_prefetch_spec") for leaf in leaves)
-    spec._STATES.clear()
-
-
-@cuda
-def test_declines_outside_decode_shape(spec):
-    spec._STATES.clear()
-    experts = _Experts()
-    layers = [_Layer(_Mlp(1, experts)), _Layer(_Mlp(2, experts))]
-    spec.link_layers(layers, 2)
-    link = layers[0]._expert_prefetch_link
-    h = torch.zeros(16, HIDDEN, device="cuda", dtype=torch.bfloat16)
-    link.on_hop(h, None)
-    link.after_moe()
-    link.after_post_moe()
-    spec.join()
-    assert not spec._STATES
-
-
-@cuda
-def test_delay_kernel_waits():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.cross_layer_expert_prefetch import (  # noqa: E501
-        kernels,
-    )
-
-    kernels.delay(1000)  # compile
-    torch.cuda.synchronize()
-    t0, t1 = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
-    t0.record()
-    kernels.delay(50_000)
-    t1.record()
-    torch.cuda.synchronize()
-    assert t0.elapsed_time(t1) >= 0.045  # ms
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py
deleted file mode 100644
index 4070e5ca1c39..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py
+++ /dev/null
@@ -1,353 +0,0 @@
-# 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.
-"""Package cross_layer_expert_prefetch's speculative expert prefetch in hop_routing's routed_moe fork.
-
-* Every routed-MoE runner class of the package tree (routed_moe's runner and every
-  subclass, e.g. hop_routing's fork) compiles the prefetch when its key carries a
-  cross_layer_expert_prefetch spec (``key[6]``): ``spec_k`` / ``spec_ctas`` / ``spec_self`` of the
-  kernel equal the hook's values, ``spec_k == 0`` without a hook, and the
-  capability line names the runner's package. A runner that dropped the spec
-  would compile the prefetch out silently (the call still passes the spec
-  row's address). Kernel compilation is stubbed: no GPU needed.
-* On an SM100-family GPU with real GLM-5.2 weights (``GLM_MOE_DSA_REAL_MOE_INPUTS_DIR``), the
-  fork's outputs with the prefetch equal the fork's outputs without it bit
-  for bit (raw tuple, finalized output; the FP32 accumulator within the FP32
-  reorder bound of the exact sum), for valid and out-of-range expert IDs, R
-  = 1 / 6 / 8 and K = 8 / 12 with all or a subset of issuing CTAs; the
-  self-prediction store writes the call's distinct experts ranked by row
-  count; a CUDA graph with the prefetch replays bit-identically.
-"""
-
-import importlib
-import os
-import pkgutil
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-_ROOT = os.path.dirname(_HERE)
-sys.path.insert(0, os.path.join(_ROOT, "routed_moe"))
-sys.path.insert(0, os.path.join(_ROOT, "hop_routing"))
-
-KMAX = 32
-REORDER_REL = 2.0**-20
-_PKG = "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels"
-
-sm100 = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-
-# ---------------------------------------------------------------------------
-# Structural guard: the spec reaches every runner's kernel (no compilation)
-# ---------------------------------------------------------------------------
-def _runner_classes():
-    """routed_moe's routed-MoE runner and every subclass in the package tree."""
-    pkg = importlib.import_module(_PKG)
-    for mod in pkgutil.iter_modules(pkg.__path__):
-        if mod.ispkg:
-            try:
-                importlib.import_module(f"{_PKG}.{mod.name}.routed_moe")
-            except ModuleNotFoundError:
-                pass
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    out, stack = {RoutedMoe}, [RoutedMoe]
-    while stack:
-        for sub in stack.pop().__subclasses__():
-            if sub not in out:
-                out.add(sub)
-                stack.append(sub)
-    return sorted(out, key=lambda c: f"{c.__module__}.{c.__qualname__}")
-
-
-class _Owner:
-    """Stands in for the NVFP4 expert leaf: only what the compile path reads."""
-
-    def __init__(self):
-        for name in ("w3_w1_weight", "w3_w1_weight_scale", "w2_weight", "w2_weight_scale"):
-            setattr(self, name, torch.zeros(64, dtype=torch.uint8))
-
-
-class _Stream:
-    cuda_stream = 0
-
-
-def _compile_kernels(cls, monkeypatch, hook):
-    """The kernel objects ``cls._compiled`` constructs for an FP32-accumulator call (stubbed compile)."""
-    import cutlass.cute as cute
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe as routed_moe_mod,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    built = []
-
-    def fake_compile(fn, *args, **kwargs):
-        built.append(getattr(fn, "__self__", fn))
-        return object()
-
-    logged = []
-    monkeypatch.setattr(cute, "compile", fake_compile)
-    monkeypatch.setattr(torch.cuda, "current_stream", lambda *a, **k: _Stream())
-    monkeypatch.setattr(
-        dispatch, "log_capability", lambda f, k, e, r="": logged.append((f, k, e)) or e
-    )
-    # 24 co-resident four-CTA clusters: a 96-CTA expert-cluster grid, as on the served GPU.
-    monkeypatch.setattr(routed_moe_mod, "_max_active_clusters", lambda device, cluster: 24)
-    # Each class keeps its own compile cache: start from empty ones.
-    monkeypatch.setattr(routed_moe_mod, "_COMPILED", {})
-    mod = sys.modules[cls.__module__]
-    if hasattr(mod, "_COMPILED"):
-        monkeypatch.setattr(mod, "_COMPILED", {})
-    owner = _Owner()
-    if hook is not None:
-        owner._expert_prefetch_spec = hook
-    runner = cls.__new__(cls)
-    runner.module = owner
-    runner.bias = torch.zeros(256, dtype=torch.bfloat16)
-    runner.device = torch.device("cpu")
-    runner.num_ctas = 152
-    key = runner._key(False, 0, 6, True)
-    runner._compiled(False, 0, 6, True)
-    kernels = [k for k in built if hasattr(k, "spec_k")]
-    assert len(kernels) == 1, f"{cls.__qualname__}: expected one expert-cluster kernel, got {built}"
-    return key, kernels[0], logged
-
-
-@pytest.mark.parametrize("spec", [(12, 0, True), (8, 32, True), (8, 40, False), (0, 0, True)])
-def test_every_runner_compiles_the_expert_prefetch_spec(monkeypatch, spec):
-    ids = torch.zeros(KMAX, dtype=torch.int32)
-    classes = _runner_classes()
-    assert any(c.__module__.endswith(".hop_routing.routed_moe") for c in classes), classes
-    for cls in classes:
-        key, kern, logged = _compile_kernels(cls, monkeypatch, lambda: (ids,) + spec)
-        assert key[6] == spec, (cls.__qualname__, key)
-        got = (kern.spec_k, kern.spec_ctas, kern.spec_self)
-        assert got == spec, (
-            f"{cls.__qualname__} dropped the cross_layer_expert_prefetch spec: kernel {got}, key {spec}"
-        )
-        if spec[0]:
-            pkg = cls.__module__.split(".")[-2]
-            line = f"expert_prefetch[k={spec[0]},ctas={spec[1]},self={int(spec[2])}]"
-            assert (pkg, line, True) in logged, (cls.__qualname__, logged)
-
-
-def test_runner_without_expert_prefetch_has_no_prefetch(monkeypatch):
-    for cls in _runner_classes():
-        key, kern, logged = _compile_kernels(cls, monkeypatch, None)
-        assert key[6] is None
-        assert (kern.spec_k, kern.spec_ctas, kern.spec_self) == (0, 0, False), cls.__qualname__
-        assert not [entry for entry in logged if entry[1].startswith("expert_prefetch")], logged
-
-
-def test_expert_prefetch_requires_routed_moe():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    assert dispatch.REQUIRED_FEATURES["cross_layer_expert_prefetch"] == frozenset({"routed_moe"})
-    auto = dispatch.ModelingV2Mode.AUTO
-    assert dispatch.resolve(auto, {}).enabled("cross_layer_expert_prefetch")
-    plan = dispatch.resolve(auto, {dispatch.DISABLE_ENV: "routed_moe"})
-    assert not plan.enabled("cross_layer_expert_prefetch")
-
-
-# ---------------------------------------------------------------------------
-# Real weights: the fork with and without the prefetch
-# ---------------------------------------------------------------------------
-@pytest.fixture(scope="module")
-def layer():
-    import routed_moe_fixtures as fx
-
-    torch.manual_seed(0)
-    return fx.build_native_layer(
-        layer=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10")),
-        tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")),
-    )
-
-
-@pytest.fixture(scope="module")
-def fork(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing.routed_moe import (  # noqa: E501
-        HopRoutedMoe,
-    )
-
-    r = HopRoutedMoe(layer.module, layer.bias)
-    if r.schedule(8, False, True) != "xc":
-        pytest.skip("the speculative prefetch is part of the expert-cluster schedule")
-    return r
-
-
-class _Spec:
-    """Install a cross_layer_expert_prefetch hook on the fork's expert leaf."""
-
-    def __init__(self, runner, ids, k, ctas=0, self_store=False):
-        self.owner = runner.module
-        self.spec = (ids, k, ctas, self_store)
-
-    def __enter__(self):
-        self.owner._expert_prefetch_spec = lambda: self.spec
-        return self
-
-    def __exit__(self, *exc):
-        del self.owner._expert_prefetch_spec
-
-
-def _pack(logits, bias):
-    import hop_routing_oracle as orc
-
-    ids, w = orc.native_routing(logits, bias)
-    routing = torch.empty(2, ids.shape[0], 8, dtype=torch.int32, device="cuda")
-    routing[0].copy_(ids)
-    routing[1].copy_(w.float().view(torch.int32))
-    return routing
-
-
-def _windows(layer, rows, count, offset=0):
-    import routed_moe_fixtures as fx
-
-    win = [
-        (x, lg)
-        for x, lg in fx.real_windows(fx.load_real_capture(layer.layer), rows)
-        if x is not None
-    ]
-    if not win:
-        pytest.skip(
-            f"no captured hidden states of layer {layer.layer} (GLM_MOE_DSA_REAL_MOE_INPUTS_DIR)"
-        )
-    step = max(1, len(win) // count)
-    out = []
-    for i in range(count):
-        x, lg = win[(offset + i * step) % len(win)]
-        x = x.cuda().contiguous()
-        lg = lg.float().cuda().contiguous()
-        out.append((x, _pack(lg, layer.bias)))
-    return out
-
-
-def _ids(kind, seed):
-    g = torch.Generator().manual_seed(seed)
-    if kind == "valid":
-        ids = torch.randperm(256, generator=g)[:KMAX]
-    else:  # the kernel masks each ID into [0, 256)
-        ids = torch.randint(-(2**31), 2**31 - 1, (KMAX,), generator=g)
-    return ids.to(torch.int32).cuda()
-
-
-def _outputs(fork, x, routing):
-    fin = fork(x, routing)
-    fc2, scale, _ = fork(x, routing, do_finalize=False)
-    acc = torch.zeros(x.shape[0], 6144, dtype=torch.float32, device="cuda")
-    fork(x, routing, red_out=acc)
-    torch.cuda.synchronize()
-    return fin.clone(), fc2.clone(), scale.clone(), acc
-
-
-def _check_same(ref, got, what):
-    fin, fc2, scale, _ = ref
-    gfin, gfc2, gscale, gacc = got
-    assert torch.equal(gfin, fin), f"{what}: finalized output differs"
-    assert torch.equal(gfc2, fc2), f"{what}: raw FC2 rows differ"
-    assert torch.equal(gscale, scale), f"{what}: routing weights differ"
-    rows = gacc.shape[0]
-    terms = scale.double()[:, :, None] * fc2.view(rows, 8, -1).double()
-    exact = terms.sum(1)
-    bound = REORDER_REL * terms.abs().sum(1)
-    assert ((gacc.double() - exact).abs() <= bound).all(), f"{what}: accumulator"
-
-
-@sm100
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("k,ctas", [(12, 0), (8, 32), (8, 0), (12, 40)])
-@pytest.mark.parametrize("ids_kind", ["valid", "out_of_range"])
-def test_fork_prefetch_outputs_bit_identical(layer, fork, rows, k, ctas, ids_kind):
-    for i, (x, routing) in enumerate(_windows(layer, rows, 2, offset=rows)):
-        ref = _outputs(fork, x, routing)
-        ids = _ids(ids_kind, seed=rows * 100 + i)
-        with _Spec(fork, ids, k, ctas):
-            got = _outputs(fork, x, routing)
-        _check_same(ref, got, f"rows={rows} k={k} ctas={ctas} {ids_kind} window {i}")
-
-
-def _ranked_experts(routing):
-    ids = routing[0].flatten().long().cpu()
-    counts = torch.bincount(ids, minlength=256)
-    distinct = [e for e in range(256) if counts[e] > 0]
-    return sorted(distinct, key=lambda e: (-int(counts[e]), e))
-
-
-@sm100
-@pytest.mark.parametrize("rows", [1, 3, 6, 8])
-@pytest.mark.parametrize("k,ctas", [(12, 0), (8, 32)])
-def test_fork_self_store_ranks_by_row_count(layer, fork, rows, k, ctas):
-    for i, (x, routing) in enumerate(_windows(layer, rows, 3, offset=5 * rows)):
-        ref = _outputs(fork, x, routing)
-        ids = torch.full((KMAX,), 999, dtype=torch.int32, device="cuda")
-        with _Spec(fork, ids, k, ctas, self_store=True):
-            acc = torch.zeros(rows, 6144, dtype=torch.float32, device="cuda")
-            fork(x, routing, red_out=acc)
-            torch.cuda.synchronize()
-            got = ids.tolist()
-            # the store is idempotent for the same call (the prefetch reads the
-            # previous call's row first)
-            fc2, scale, _ = fork(x, routing, do_finalize=False)
-            gfin = fork(x, routing)
-            torch.cuda.synchronize()
-        assert torch.equal(fc2, ref[1]) and torch.equal(scale, ref[2]), (rows, i)
-        assert torch.equal(gfin, ref[0]), (rows, i)
-        want = _ranked_experts(routing)[:KMAX]
-        assert got[: len(want)] == want, (rows, i, got, want)
-        assert all(v == 999 for v in got[len(want) :]), (rows, i, got)
-
-
-@sm100
-def test_fork_prefetch_cuda_graph_replay(layer, fork):
-    rows = 6
-    wins = _windows(layer, rows, 6, offset=3)
-    # References: the fork without the prefetch, eager.
-    refs = []
-    for x, routing in wins:
-        fc2, scale, _ = fork(x, routing, do_finalize=False)
-        torch.cuda.synchronize()
-        refs.append((fc2.clone(), scale.clone()))
-    x_static = wins[0][0].clone()
-    r_static = wins[0][1].clone()
-    ids = torch.zeros(KMAX, dtype=torch.int32, device="cuda")
-    with _Spec(fork, ids, 8, 32, self_store=True):
-        stream = torch.cuda.Stream()
-        stream.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(stream):
-            fork(x_static, r_static, do_finalize=False)
-        torch.cuda.current_stream().wait_stream(stream)
-        # the eager call compiled the prefetch variant: the capture records it
-        assert fork.ready_for_capture(False, rows=rows)
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph, stream=stream):
-            fc2_s, scale_s, _ = fork(x_static, r_static, do_finalize=False)
-    for it, (x, routing) in enumerate(wins):
-        x_static.copy_(x)
-        r_static.copy_(routing)
-        graph.replay()
-        torch.cuda.synchronize()
-        assert torch.equal(fc2_s, refs[it][0]) and torch.equal(scale_s, refs[it][1]), it
-        # the replay stored this call's ranked experts for the next one
-        want = _ranked_experts(routing)[:KMAX]
-        assert ids.tolist()[: len(want)] == want, it
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py
deleted file mode 100644
index 73142cbf4ea7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/test_expert_prefetch_composition_shared_prefetch_expert_prefetch.py
+++ /dev/null
@@ -1,352 +0,0 @@
-# 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.
-"""The cross_layer_expert_prefetch expert prefetch in the shared_expert_prefetch fork.
-
-``cross_layer_expert_prefetch/test_expert_prefetch_composition_hop_routing.py`` checks that every
-routed-MoE runner class of the tree, shared_expert_prefetch's included, compiles the prefetch from a
-key with a cross_layer_expert_prefetch spec. This file checks what is specific to
-shared_expert_prefetch's port:
-
-* the merged pre-wait prologue of ``SharedPrefetchExpertClusterKernel``: the shared-expert chunks
-  are spread over cross_layer_expert_prefetch's issuing CTAs (``spec_ctas``) when
-  cross_layer_expert_prefetch restricts them, over the whole grid otherwise
-  (shared_expert_prefetch's mapping without cross_layer_expert_prefetch), and the chunks cover the
-  shared expert's bytes;
-* ``SharedPrefetchRoutedMoe.__call__`` passes cross_layer_expert_prefetch's row of predicted experts
-  as the 18th routed_moe address (routed_moe's dummy without a hook) before the two shared-expert
-  addresses, i.e. ``_NUM_XC_ADDRS + 2`` addresses in the kernel's order.
-  Both need no GPU (kernel compilation and launch are stubbed);
-* on an SM100-family GPU with real GLM-5.2 weights (``GLM_MOE_DSA_REAL_MOE_INPUTS_DIR``,
-  ``GLM_MOE_DSA_ROUTED_MOE_CKPT_DIR``): shared_expert_prefetch's FP32 accumulator with the prefetch (K = 8 / 12, all
-  or 32 issuing CTAs, valid and out-of-range IDs) is the exact sum of routed_moe's weighted rows within the
-  FP32 reorder bound (routed_moe_into_hop's fp32 class, as shared_expert_prefetch without
-  cross_layer_expert_prefetch); the self-store writes the call's distinct experts ranked by row
-  count; a CUDA graph with the prefetch replays within the bound and keeps storing.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-_ROOT = os.path.dirname(_HERE)
-for _sub in ("routed_moe", "hop_routing", "shared_expert_prefetch"):
-    sys.path.insert(0, os.path.join(_ROOT, _sub))
-
-KMAX = 32
-REORDER_REL = 2.0**-20
-LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
-
-sm100 = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-
-# ---------------------------------------------------------------------------
-# The merged prologue's issuers (no compilation)
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize(
-    "num_ctas,spec_k,spec_ctas,want",
-    [
-        (96, 0, 0, 96),  # shared_expert_prefetch without cross_layer_expert_prefetch: every CTA
-        (96, 0, 32, 96),  # spec_ctas, prefetch not compiled in: shared_expert_prefetch's mapping
-        (96, 8, 0, 96),  # cross_layer_expert_prefetch on every CTA
-        (96, 8, 32, 96 // 3),  # the prefetch on the early-resident CTAs: the same CTAs
-        (96, 12, 40, 40),
-        (96, 8, 200, 96),
-        (144, 8, 56, 56),
-    ],
-)
-def test_prologue_issuers_and_coverage(num_ctas, spec_k, spec_ctas, want):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch import (  # noqa: E501
-        kernel_xc as zk,
-    )
-
-    kern = zk.SharedPrefetchExpertClusterKernel(
-        num_ctas, spec_k=spec_k, spec_ctas=spec_ctas, spec_self=True, red_out=True
-    )
-    assert kern.se_ctas == want
-    assert (kern.spec_k, kern.spec_ctas) == (spec_k, spec_ctas)
-    for chunk, total in ((kern.se_chunk1, zk.SE_W1_BYTES), (kern.se_chunk2, zk.SE_W2_BYTES)):
-        assert chunk % 16 == 0
-        lanes = kern.se_ctas * zk.PF_LANES
-        assert lanes * chunk >= total, (chunk, total)
-        # every issuing CTA owns a chunk (no CTA past se_ctas is needed)
-        assert (lanes - zk.PF_LANES) * chunk < total
-    if spec_k == 0:
-        assert (kern.se_chunk1, kern.se_chunk2) == (
-            zk._chunk(zk.SE_W1_BYTES, num_ctas),
-            zk._chunk(zk.SE_W2_BYTES, num_ctas),
-        )
-
-
-# ---------------------------------------------------------------------------
-# SharedPrefetchRoutedMoe.__call__ passes cross_layer_expert_prefetch's row (no launch)
-# ---------------------------------------------------------------------------
-class _Stream:
-    cuda_stream = 0
-
-
-class _Owner:
-    """The NVFP4 expert leaf: only what the call path reads."""
-
-    def __init__(self):
-        for name in (
-            "w3_w1_weight",
-            "w3_w1_weight_scale",
-            "w2_weight",
-            "w2_weight_scale",
-            "fc31_alpha",
-            "fc31_scale_c",
-            "fc2_alpha",
-            "fc31_input_scale",
-        ):
-            setattr(self, name, torch.zeros(64, dtype=torch.uint8))
-
-
-class _Ws:
-    def __init__(self):
-        for name in ("rows", "scales", "dummy", "tile_count"):
-            setattr(self, name, torch.zeros(16, dtype=torch.uint8))
-        self.num_ctas = 152
-
-    def order_on(self, stream):
-        pass
-
-
-@pytest.mark.parametrize("hooked", [True, False])
-def test_call_passes_the_expert_prefetch_row(monkeypatch, hooked):
-    import cutlass
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe as routed_moe_mod,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
-        SharedPrefetchRoutedMoe,
-    )
-
-    calls = []
-    ws = _Ws()
-    monkeypatch.setattr(routed_moe_mod, "_bulk_copy_ready", lambda x: x)
-    monkeypatch.setattr(routed_moe_mod, "_workspace", lambda device, n: ws)
-    monkeypatch.setattr(torch.cuda, "current_stream", lambda *a, **k: _Stream())
-    # 24 co-resident four-CTA clusters: a 96-CTA expert-cluster grid, as on the served GPU.
-    monkeypatch.setattr(routed_moe_mod, "_max_active_clusters", lambda device, cluster: 24)
-    owner = _Owner()
-    ids = torch.zeros(KMAX, dtype=torch.int32)
-    if hooked:
-        owner._expert_prefetch_spec = lambda: (ids, 8, 32, True)
-    runner = SharedPrefetchRoutedMoe.__new__(SharedPrefetchRoutedMoe)
-    runner.module = owner
-    runner.bias = torch.zeros(256, dtype=torch.bfloat16)
-    runner.device = torch.device("cpu")
-    runner.num_ctas = 152
-    monkeypatch.setattr(runner, "_compiled", lambda *a, **k: (lambda *args: calls.append(args)))
-    x = torch.zeros(6, 16, dtype=torch.bfloat16)
-    routing = torch.zeros(2, 6, 8, dtype=torch.int32)
-    acc = torch.zeros(6, 16, dtype=torch.float32)
-    se = (torch.zeros(8, dtype=torch.bfloat16), torch.zeros(8, dtype=torch.bfloat16))
-    assert runner(x, routing, acc, se) == [acc]
-    assert len(calls) == 1
-    ints = [a for a in calls[0] if isinstance(a, cutlass.Int64)]
-    addrs = [int(a.value) if hasattr(a, "value") else int(a) for a in ints]
-    assert len(addrs) == routed_moe_mod._NUM_XC_ADDRS + 2 == 20
-    assert (
-        addrs[0] == x.data_ptr() and addrs[1] == routing.data_ptr() and addrs[10] == acc.data_ptr()
-    )
-    assert addrs[17] == (ids.data_ptr() if hooked else ws.dummy.data_ptr())
-    assert addrs[18:] == [se[0].data_ptr(), se[1].data_ptr()]
-
-
-# ---------------------------------------------------------------------------
-# Real weights: shared_expert_prefetch's fork with and without the prefetch
-# ---------------------------------------------------------------------------
-@pytest.fixture(scope="module")
-def layer():
-    import routed_moe_fixtures as fx
-
-    torch.manual_seed(0)
-    return fx.build_native_layer(
-        layer=LAYER, tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
-    )
-
-
-@pytest.fixture(scope="module")
-def runners(layer):
-    import shared_prefetch_fixtures as zf
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
-        SharedPrefetchRoutedMoe,
-        shared_weights,
-    )
-
-    mlp = zf.build_shared_mlp(LAYER, int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")))
-    se, reason = shared_weights(mlp)
-    assert se is not None, reason
-    shared_expert_prefetch = SharedPrefetchRoutedMoe(layer.module, layer.bias)
-    if shared_expert_prefetch.schedule(8, False, True) != "xc":
-        pytest.skip("the prefetch is part of the expert-cluster schedule")
-    return RoutedMoe.from_native(layer.module, layer.bias), shared_expert_prefetch, se, mlp
-
-
-class _Spec:
-    """Install a cross_layer_expert_prefetch hook on the fork's expert leaf."""
-
-    def __init__(self, runner, ids, k, ctas=0, self_store=False):
-        self.owner = runner.module
-        self.spec = (ids, k, ctas, self_store)
-
-    def __enter__(self):
-        self.owner._expert_prefetch_spec = lambda: self.spec
-        return self
-
-    def __exit__(self, *exc):
-        del self.owner._expert_prefetch_spec
-
-
-def _pack(logits, bias):
-    import hop_routing_oracle as orc
-
-    ids, w = orc.native_routing(logits, bias)
-    routing = torch.empty(2, ids.shape[0], 8, dtype=torch.int32, device="cuda")
-    routing[0].copy_(ids)
-    routing[1].copy_(w.float().view(torch.int32))
-    return routing
-
-
-def _windows(layer, rows, count, offset=0):
-    import routed_moe_fixtures as fx
-
-    win = [
-        (x, lg)
-        for x, lg in fx.real_windows(fx.load_real_capture(layer.layer), rows)
-        if x is not None
-    ]
-    if not win:
-        pytest.skip(
-            f"no captured hidden states of layer {layer.layer} (GLM_MOE_DSA_REAL_MOE_INPUTS_DIR)"
-        )
-    step = max(1, len(win) // count)
-    out = []
-    for i in range(count):
-        x, lg = win[(offset + i * step) % len(win)]
-        x = x.cuda().contiguous()
-        lg = lg.float().cuda().contiguous()
-        out.append((x, lg, _pack(lg, layer.bias)))
-    return out
-
-
-def _ids(kind, seed):
-    g = torch.Generator().manual_seed(seed)
-    if kind == "valid":
-        ids = torch.randperm(256, generator=g)[:KMAX]
-    else:  # the kernel masks each ID into [0, 256)
-        ids = torch.randint(-(2**31), 2**31 - 1, (KMAX,), generator=g)
-    return ids.to(torch.int32).cuda()
-
-
-def _exact(routed_runner, x, lg):
-    fc2, scale, idx = routed_runner(x, lg, do_finalize=False)
-    torch.cuda.synchronize()
-    terms = scale.double()[:, :, None] * fc2[idx.long()].double()
-    return terms.sum(1), REORDER_REL * terms.abs().sum(1)
-
-
-def _shared_prefetch_acc(shared_expert_prefetch, se, x, routing):
-    acc = torch.zeros(x.shape[0], 6144, dtype=torch.float32, device="cuda")
-    shared_expert_prefetch(x, routing, acc, se)
-    torch.cuda.synchronize()
-    return acc
-
-
-def _ranked_experts(routing):
-    ids = routing[0].flatten().long().cpu()
-    counts = torch.bincount(ids, minlength=256)
-    distinct = [e for e in range(256) if counts[e] > 0]
-    return sorted(distinct, key=lambda e: (-int(counts[e]), e))
-
-
-@sm100
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("k,ctas", [(8, 32), (12, 0), (8, 0), (12, 40)])
-@pytest.mark.parametrize("ids_kind", ["valid", "out_of_range"])
-def test_shared_prefetch_accumulator_within_bound(layer, runners, rows, k, ctas, ids_kind):
-    routed_runner, shared_expert_prefetch, se, _ = runners
-    for i, (x, lg, routing) in enumerate(_windows(layer, rows, 2, offset=rows)):
-        exact, bound = _exact(routed_runner, x, lg)
-        ref = _shared_prefetch_acc(shared_expert_prefetch, se, x, routing)
-        assert ((ref.double() - exact).abs() <= bound).all(), ("no prefetch", rows, i)
-        ids = _ids(ids_kind, seed=rows * 100 + i)
-        with _Spec(shared_expert_prefetch, ids, k, ctas):
-            got = _shared_prefetch_acc(shared_expert_prefetch, se, x, routing)
-        assert ((got.double() - exact).abs() <= bound).all(), (rows, k, ctas, ids_kind, i)
-
-
-@sm100
-@pytest.mark.parametrize("rows", [1, 3, 6, 8])
-@pytest.mark.parametrize("k,ctas", [(8, 32), (12, 0)])
-def test_shared_prefetch_self_store_ranks_by_row_count(layer, runners, rows, k, ctas):
-    routed_runner, shared_expert_prefetch, se, _ = runners
-    for i, (x, lg, routing) in enumerate(_windows(layer, rows, 3, offset=5 * rows)):
-        exact, bound = _exact(routed_runner, x, lg)
-        ids = torch.full((KMAX,), 999, dtype=torch.int32, device="cuda")
-        with _Spec(shared_expert_prefetch, ids, k, ctas, self_store=True):
-            acc = _shared_prefetch_acc(shared_expert_prefetch, se, x, routing)
-            got = ids.tolist()
-        assert ((acc.double() - exact).abs() <= bound).all(), (rows, i)
-        want = _ranked_experts(routing)[:KMAX]
-        assert got[: len(want)] == want, (rows, i, got, want)
-        assert all(v == 999 for v in got[len(want) :]), (rows, i, got)
-
-
-@sm100
-def test_shared_prefetch_cuda_graph_replay(layer, runners):
-    routed_runner, shared_expert_prefetch, se, _ = runners
-    rows = 6
-    wins = _windows(layer, rows, 6, offset=3)
-    refs = [_exact(routed_runner, x, lg) for x, lg, _ in wins]
-    x_static = wins[0][0].clone()
-    r_static = wins[0][2].clone()
-    acc = torch.zeros(rows, 6144, dtype=torch.float32, device="cuda")
-    ids = torch.zeros(KMAX, dtype=torch.int32, device="cuda")
-    with _Spec(shared_expert_prefetch, ids, 8, 32, self_store=True):
-        stream = torch.cuda.Stream()
-        stream.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(stream):
-            shared_expert_prefetch(x_static, r_static, acc, se)
-        torch.cuda.current_stream().wait_stream(stream)
-        torch.cuda.synchronize()
-        # the eager call compiled the prefetch variant: the capture records it
-        assert shared_expert_prefetch.ready_for_capture(False, rows=rows, red=True)
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph, stream=stream):
-            shared_expert_prefetch(x_static, r_static, acc, se)
-    for it, (x, _, routing) in enumerate(wins):
-        x_static.copy_(x)
-        r_static.copy_(routing)
-        acc.zero_()
-        graph.replay()
-        torch.cuda.synchronize()
-        exact, bound = refs[it]
-        assert ((acc.double() - exact).abs() <= bound).all(), it
-        want = _ranked_experts(routing)[:KMAX]
-        assert ids.tolist()[: len(want)] == want, it
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py
deleted file mode 100644
index d04624e97387..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/deferred_norm/test_deferred_norm_routed_into_hop_wiring_mpi.py
+++ /dev/null
@@ -1,343 +0,0 @@
-# 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.
-"""deferred_norm with routed_moe_into_hop through the target's decoder layer (8 ranks over 2 hosts).
-
-Runs the target's ``GlmDecoderLayer.forward_MoE`` as a decode step on a decoder-layer
-stand-in that carries a real MoE of layer 10 (``Deepseekv3MoE`` adopted by ``GlmMoE``;
-routed_moe's fixtures: real NVFP4 experts and router, captured real hidden-state windows),
-the layer's post_moe_hop with its routed_moe_into_hop routed input enabled, and the
-deferred_norm hop with its routed input enabled. Features
-``absorbed_query,post_moe_hop,routed_moe,routed_moe_into_hop,deferred_norm``: routed_moe writes the
-routed result into the post_moe_hop and ``forward_MoE`` takes the routed_moe_into_hop branch.
-
-Checks, per routed_moe_into_hop mode:
-
-* the routed_moe_into_hop branch returns a deferred_norm packet (the deferred_norm routed hop
-  engaged);
-* the packet's materialization and the residual equal
-  ``PostMoeHop.hop_routed`` on the routed content routed_moe wrote for that call
-  (captured before the hop consumed it) bit for bit;
-* without the deferred_norm hop (as with ``deferred_norm`` off) the same branch returns
-  post_moe_hop's normalized tensor, bit-identical to the materialized packet in the exact
-  routed_moe_into_hop modes (routed_moe's ``fp32`` mode accumulates with ``red.add``, so two
-  routed_moe calls can differ by FP32 reordering; there the two runs are compared
-  within routed_moe_into_hop's reorder class only through the captured-input check above);
-* a context-phase batch keeps post_moe_hop's ``hop_routed``;
-* a CUDA graph of the branch (multi-stream capture as served) replays with
-  fresh windows and matches the eager deferred_norm path.
-
-The decoder layer's adoption wiring is checked on the target's source: the routed enable is
-called only with ``routed_moe_into_hop`` selected and the post_moe_hop's routed input set,
-after the deferred_norm hop is built.
-"""
-
-import inspect
-import os
-import sys
-import types
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "..", "..", "shared", "kernels", "post_moe_hop"))
-sys.path.insert(0, os.path.join(_HERE, "..", "routed_moe"))
-sys.path.insert(0, os.path.join(_HERE, "..", "..", "..", "shared", "kernels"))
-
-import plan_testlib  # noqa: E402
-import post_moe_hop_mpi_utils as M  # noqa: E402
-
-LAYER = 10
-FEATURES = "absorbed_query,post_moe_hop,routed_moe,routed_moe_into_hop,deferred_norm"
-MODES = ("fp32", "identity", "row")
-
-
-def _set_features(value: str) -> None:
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    plan_testlib.select(dispatch, value)
-
-
-@pytest.fixture(autouse=True)
-def decode_step():
-    """Every test runs as a decode step, where the target's decoder layer takes its fast paths."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    with dispatch.step(dispatch.Phase.DECODE):
-        yield
-
-
-@pytest.fixture(scope="module")
-def env():
-    e = M.mpi_env()
-    if e.world != 8:
-        pytest.skip("needs the TP8 group")
-    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
-    return e
-
-
-@pytest.fixture(scope="module")
-def layer(env):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    if not hasattr(PostMoeHop, "enable_routed_input"):
-        pytest.skip("the installed post_moe_hop hop has no routed_moe_into_hop routed input")
-    import routed_moe_chain_bench as CB
-    import routed_moe_fixtures as fx
-
-    _set_features("")
-    moes, gammas = CB.build_layers(env, [LAYER])
-    moe = moes[LAYER]
-    hop = PostMoeHop(env.mapping, fx.HIDDEN, fx.TOP_K, torch.bfloat16, raw_producer=moe.experts)
-    assert M.all_ranks(hop.engaged), hop.reason
-    win = [x for x, _ in fx.real_windows(fx.load_real_capture(LAYER), 6)]
-    step = max(1, len(win) // 8)
-    windows = [win[i * step].cuda() for i in range(8)]
-    yield types.SimpleNamespace(moe=moe, gamma=gammas[LAYER], hop=hop, windows=windows)
-    _set_features("")
-
-
-def _stub(layer, deferred_norm, gamma):
-    """The decoder-layer attributes ``forward_MoE`` reads on the post_moe_hop path."""
-    return types.SimpleNamespace(
-        _post_attn_router_weight=lambda: None,
-        _into_hop=None,
-        fusion_config=types.SimpleNamespace(PRE_MOE_FUSION=False, POST_MOE_FUSION=True),
-        post_attention_layernorm=lambda h, r: (h, r),
-        allreduce=None,
-        moe_allreduce=None,
-        model_config=None,
-        is_p2p_supported=False,
-        mapping=M.mpi_env().mapping,
-        mlp=layer.moe,
-        post_moe_hop=layer.hop,
-        deferred_norm_hop=deferred_norm,
-        next_layer_layernorm=types.SimpleNamespace(weight=gamma, variance_epsilon=M.GLM_EPS),
-    )
-
-
-def _metadata(rows, context=False):
-    return types.SimpleNamespace(
-        all_rank_num_tokens=None,
-        num_contexts=1 if context else 0,
-        num_generations=0 if context else rows,
-    )
-
-
-def _forward(stub, x, residual, md):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmDecoderLayer,
-    )
-
-    return GlmDecoderLayer.forward_MoE(stub, x, md, residual)
-
-
-def _enable(layer, mode):
-    """Collective: routed_moe_into_hop routed input in ``mode`` and a deferred_norm hop serving it."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-
-    assert M.all_ranks(layer.hop.enable_routed_input(mode))
-    layer.moe._routed_into_hop = layer.hop
-    key = f"_deferred_norm_{id(layer.hop)}"
-    deferred_norm = getattr(layer, key, None)
-    if deferred_norm is None:
-        deferred_norm = DeferredNormHop(layer.hop)
-        setattr(layer, key, deferred_norm)
-    assert M.all_ranks(deferred_norm.engaged and deferred_norm.enable_routed_input()), (
-        deferred_norm.routed_reason
-    )
-    return deferred_norm
-
-
-class _Recorder:
-    """Wraps ``DeferredNormHop.hop_routed`` to keep a copy of the routed input."""
-
-    def __init__(self, deferred_norm):
-        self.deferred_norm = deferred_norm
-        self.inner = deferred_norm.hop_routed
-        self.calls = []
-
-    def __call__(self, routed, shared, residual, gamma, eps):
-        mode, _, scales = self.deferred_norm._routed_view(residual.shape[0])
-        self.calls.append(
-            (
-                routed.clone(),
-                None if scales is None else scales.clone(),
-                None if shared is None else shared.clone(),
-                residual.clone(),
-            )
-        )
-        return self.inner(routed, shared, residual, gamma, eps)
-
-
-def _replay_post_moe_hop(hop, mode, routed, scales, shared, residual, gamma):
-    rows = residual.shape[0]
-    if mode == "fp32":
-        buf = hop.routed_acc_buffer(rows)
-    elif mode == "identity":
-        buf, s = hop.routed_raw_buffers(rows)
-        s.copy_(scales)
-    else:
-        buf = hop.routed_input_buffer(rows)
-    buf.copy_(routed)
-    return hop.hop_routed(buf, shared, residual, gamma, M.GLM_EPS)
-
-
-def test_adoption_wiring_in_target_source(env, layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmDecoderLayer,
-    )
-
-    init = inspect.getsource(GlmDecoderLayer._adopt_post_moe)
-    body = inspect.getsource(GlmDecoderLayer.forward_MoE)
-    i_routed_into_hop = init.find("self.mlp._routed_into_hop = self.post_moe_hop")
-    i_deferred_norm = init.find('"deferred_hop",')
-    i_en = init.find("self.deferred_norm_hop.enable_routed_input()")
-    ok = 0 <= i_routed_into_hop < i_deferred_norm < i_en
-    ok &= 'dispatch.is_enabled("routed_moe_into_hop")' in init[i_deferred_norm:i_en]
-    ok &= "_routed_into_hop" in init[i_deferred_norm:i_en]
-    i_branch = body.find("elif use_post_moe_hop and len(hidden_states) == 2:")
-    i_raw = body.find("elif use_post_moe_hop:", i_branch + 1)
-    ok &= 0 <= i_branch < body.find("can_run_routed", i_branch) < i_raw
-    ok &= body.find("hop.hop_routed(", i_branch) < i_raw
-    assert ok
-
-
-@pytest.mark.parametrize("mode", MODES)
-def test_routed_into_hop_branch_returns_packet_equal_to_hop_routed(env, layer, mode):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        is_deferred,
-    )
-
-    deferred_norm = _enable(layer, mode)
-    rec = _Recorder(deferred_norm)
-    deferred_norm.hop_routed = rec
-    _set_features(FEATURES)
-    g = torch.Generator(device="cuda").manual_seed(300)
-    ok, n = True, 0
-    try:
-        with torch.inference_mode():
-            for i, x in enumerate(layer.windows):
-                rows = x.shape[0]
-                residual = (torch.randn(rows, 6144, device="cuda", generator=g) * 4.0).bfloat16()
-                out, res = _forward(
-                    _stub(layer, deferred_norm, layer.gamma), x, residual, _metadata(rows)
-                )
-                engaged = is_deferred(out) and len(rec.calls) == i + 1
-                ok &= engaged
-                if not engaged:
-                    print(
-                        f"GLM_MOE_DSA_DEFERRED_NORM_WIRING_NOT_ENGAGED mode={mode} call={i} rank={env.rank}",
-                        flush=True,
-                    )
-                    continue
-                h = out.materialize()
-                routed, scales, shared, res_in = rec.calls[-1]
-                h_ref, r_ref = _replay_post_moe_hop(
-                    layer.hop, mode, routed, scales, shared, res_in, layer.gamma
-                )
-                torch.cuda.synchronize()
-                same = torch.equal(h, h_ref) and torch.equal(res, r_ref)
-                ok &= same
-                # deferred_norm off: the same branch returns post_moe_hop's normalized tensor.
-                out_off, res_off = _forward(
-                    _stub(layer, None, layer.gamma), x, residual, _metadata(rows)
-                )
-                torch.cuda.synchronize()
-                ok &= isinstance(out_off, torch.Tensor) and not is_deferred(out_off)
-                if mode != "fp32":
-                    ok &= torch.equal(out_off, h) and torch.equal(res_off, res)
-                n += 1
-    finally:
-        del deferred_norm.hop_routed
-        _set_features("")
-    stats = env.comm.allgather((env.rank, n, ok))
-    if env.rank == 0:
-        print(
-            f"GLM_MOE_DSA_DEFERRED_NORM_WIRING mode={mode}: (rank, windows, ok) {stats}", flush=True
-        )
-    assert M.all_ranks(ok and n == len(layer.windows))
-
-
-def test_context_batch_keeps_hop_routed(env, layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        is_deferred,
-    )
-
-    deferred_norm = _enable(layer, "identity")
-    _set_features(FEATURES)
-    try:
-        with torch.inference_mode():
-            x = layer.windows[0]
-            residual = torch.zeros(x.shape[0], 6144, dtype=torch.bfloat16, device="cuda")
-            out, _ = _forward(
-                _stub(layer, deferred_norm, layer.gamma),
-                x,
-                residual,
-                _metadata(x.shape[0], context=True),
-            )
-            ok = isinstance(out, torch.Tensor) and not is_deferred(out)
-    finally:
-        _set_features("")
-    assert M.all_ranks(ok)
-
-
-def test_routed_into_hop_branch_cuda_graph_replay(env, layer):
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    mode = "identity"
-    deferred_norm = _enable(layer, mode)
-    _set_features(FEATURES)
-    try:
-        with torch.inference_mode():
-            rows = layer.windows[0].shape[0]
-            x = layer.windows[0].clone()
-            g = torch.Generator(device="cuda").manual_seed(7)
-            residual = (torch.randn(rows, 6144, device="cuda", generator=g) * 4.0).bfloat16()
-            md = _metadata(rows)
-            stub = _stub(layer, deferred_norm, layer.gamma)
-            stream = torch.cuda.Stream()
-            stream.wait_stream(torch.cuda.current_stream())
-            with torch.cuda.stream(stream):
-                for _ in range(2):
-                    _forward(stub, x, residual, md)
-            torch.cuda.current_stream().wait_stream(stream)
-            torch.cuda.synchronize()
-            env.comm.barrier()
-            graph = torch.cuda.CUDAGraph()
-            with with_multi_stream(True):
-                with torch.cuda.graph(graph, stream=stream):
-                    g_out, g_res = _forward(stub, x, residual, md)
-            torch.cuda.synchronize()
-            ok = True
-            for w in layer.windows[1:5]:
-                x.copy_(w)
-                env.comm.barrier()
-                graph.replay()
-                torch.cuda.synchronize()
-                h_graph = g_out.materialize()
-                e_out, e_res = _forward(stub, x, residual, md)
-                h_eager = e_out.materialize()
-                torch.cuda.synchronize()
-                ok &= torch.equal(h_graph, h_eager) and torch.equal(g_res, e_res)
-    finally:
-        _set_features("")
-    assert M.all_ranks(ok)
-    layer.hop.raise_if_error()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py
deleted file mode 100644
index 2da3c9dcc750..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_live_rows/test_live_rows_oracle.py
+++ /dev/null
@@ -1,270 +0,0 @@
-# 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.
-"""Gate C1(b, c): per-row R invariance of the MTP layer's MoE on real weights (one GPU, one TP8 rank shard).
-
-draft_live_rows runs the MTP layer's routed MoE, shared expert and their add on the live
-rows only. That is bit-identical to the served step 0 only if each of them
-computes a row the same way whatever the other rows of the call are. Checked
-here, bitwise, on rank ``GLM_MOE_DSA_DRAFT_LIVE_ROWS_TP_RANK``'s shard of ``model.layers.78``:
-
-* routed MoE, native FlashInfer TRT-LLM-Gen BF16 leaf and draft_routed_moe (each with its
-  own launch configuration per row count): row g of the R = 6 call equals the
-  R = 1 call on row g, for every g, with the six rows chosen to cover
-  D = 8..48 distinct experts (Gaussian rows and real token-embedding rows);
-  rows (g1, g2) of an R = 8 call equal the R = 2 call on them (B = 2 at K = 3);
-* shared expert (the served BF16 ``GatedMLP``, TP8 shard, real weights): row g
-  at M = 6 equals M = 1, rows (g1, g2) at M = 8 equal M = 2;
-* the same under CUDA-graph replay (graphs captured at R = 6 / R = 1, replayed
-  with new rows copied into their static inputs);
-* negative control: row g of the R = 6 call against the R = 1 call on another
-  row differs.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "draft_routed_moe"))
-import draft_moe_fixtures as fx  # noqa: E402
-
-TP_RANK = int(os.environ.get("GLM_MOE_DSA_DRAFT_LIVE_ROWS_TP_RANK", "0"))
-D_TARGETS = (8, 12, 16, 24, 31, 40, 48)
-
-
-def _bits(t):
-    return t.contiguous().view(torch.int16)
-
-
-def _eq(a, b):
-    return torch.equal(_bits(a), _bits(b))
-
-
-@pytest.fixture(scope="module")
-def layer():
-    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
-        pytest.skip("needs an SM100-family GPU")
-    import tensorrt_llm  # noqa: F401  (loads th_common: draft_routed_moe and the native MoE ops)
-
-    if not hasattr(torch.ops.trtllm, "glm_moe_dsa_draft_routed_moe"):
-        pytest.skip("draft_routed_moe op not in this build")
-    return fx.build_native(TP_RANK)
-
-
-@pytest.fixture(scope="module")
-def pools(layer):
-    gauss = fx.RowPool.gaussian(layer, num=1024, seed=10_000)
-    g = torch.Generator().manual_seed(77)
-    ids = torch.randint(0, 100_000, (1024,), generator=g).tolist()
-    xs = fx.embedding_rows(layer, ids)
-    emb = fx.RowPool(xs, fx.top8_sets(layer.bias, layer.logits(xs)))
-    return {"gauss": gauss, "emb": emb}
-
-
-def _routed(kind, layer, x, logits):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    if kind == "draft_routed_moe":
-        return rm.run(layer.module, layer.bias, x, logits)
-    return layer.native(x, logits)
-
-
-@pytest.mark.parametrize("kind", ["native", "draft_routed_moe"])
-@pytest.mark.parametrize("pool", ["gauss", "emb"])
-def test_routed_row_invariance_r6(layer, pools, kind, pool):
-    bad, cases, negatives = [], 0, 0
-    for t, target in enumerate(D_TARGETS):
-        for rep in range(3):
-            x6, d = pools[pool].take(6, target, seed=1000 * t + rep)
-            l6 = layer.logits(x6)
-            out6 = _routed(kind, layer, x6, l6)
-            for g in range(6):
-                out1 = _routed(kind, layer, x6[g : g + 1].contiguous(), l6[g : g + 1].contiguous())
-                cases += 1
-                if not _eq(out6[g : g + 1], out1):
-                    bad.append((target, d, rep, g))
-            # Negative control: another (different) row's R = 1 output differs from row 0's.
-            others = [h for h in range(1, 6) if not _eq(x6[h], x6[0])]
-            if others:
-                h = others[0]
-                out1o = _routed(kind, layer, x6[h : h + 1].contiguous(), l6[h : h + 1].contiguous())
-                assert not _eq(out6[0:1], out1o), "negative control not caught"
-                negatives += 1
-    torch.cuda.synchronize()
-    assert negatives >= len(D_TARGETS), negatives
-    print(
-        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE routed {kind} {pool} R=6 vs R=1 cases={cases} mismatching={len(bad)} "
-        f"negative_controls={negatives} {bad[:4]}",
-        flush=True,
-    )
-    assert not bad
-
-
-@pytest.mark.parametrize("kind", ["native", "draft_routed_moe"])
-def test_routed_pairs_r8(layer, pools, kind):
-    bad, cases = [], 0
-    for t, target in enumerate((16, 31, 48, 64)):
-        x8, d = pools["gauss"].take(8, target, seed=5000 + t)
-        l8 = layer.logits(x8)
-        out8 = _routed(kind, layer, x8, l8)
-        for g1, g2 in ((0, 7), (7, 0), (3, 4), (5, 2), (1, 6)):
-            idx = torch.tensor([g1, g2], device="cuda")
-            out2 = _routed(kind, layer, x8[idx].contiguous(), l8[idx].contiguous())
-            cases += 1
-            if not _eq(out8[idx], out2):
-                bad.append((target, d, g1, g2))
-    torch.cuda.synchronize()
-    print(
-        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE routed {kind} R=8 vs R=2 cases={cases} mismatching={len(bad)} {bad[:4]}",
-        flush=True,
-    )
-    assert not bad
-
-
-def _shared_expert(layer):
-    """The served BF16 shared expert of the MTP layer, TP rank shard, real weights (Deepseekv3MoE's construction)."""
-    import copy
-    import math
-
-    sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
-    import routed_moe_fixtures as wfx
-
-    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-    from tensorrt_llm.quantization.mode import QuantAlgo
-
-    mc = wfx._served_model_config(fx.MTP_LAYER, TP_RANK)
-    mc.quant_config.exclude_modules = [f"model.layers.{fx.MTP_LAYER}*"]
-    shared_mc = copy.copy(mc)
-    shared_mc.quant_config = QuantConfig(quant_algo=None, kv_cache_quant_algo=QuantAlgo.FP8)
-    tp = math.gcd(fx.INTER // 1, mc.mapping.tp_size)
-    se = GatedMLP(
-        hidden_size=fx.HIDDEN,
-        intermediate_size=fx.INTER,
-        bias=False,
-        dtype=torch.bfloat16,
-        config=shared_mc,
-        overridden_tp_size=tp,
-        reduce_output=False,
-    )
-    rd = fx.Reader()
-    pre = f"model.layers.{fx.MTP_LAYER}.mlp.shared_experts"
-    full = {p: rd.full(f"{pre}.{p}.weight") for p in ("gate_proj", "up_proj", "down_proj")}
-    if hasattr(se, "gate_up_proj"):
-        se.gate_up_proj.load_weights([{"weight": full["gate_proj"]}, {"weight": full["up_proj"]}])
-    else:
-        se.gate_proj.load_weights([{"weight": full["gate_proj"]}])
-        se.up_proj.load_weights([{"weight": full["up_proj"]}])
-    se.down_proj.load_weights([{"weight": full["down_proj"]}])
-    wfx._walk_post_load(se)
-    return se.cuda(), tp
-
-
-@pytest.fixture(scope="module")
-def shared(layer):
-    se, tp = _shared_expert(layer)
-    print(
-        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE shared expert TP{tp} rank {TP_RANK} {type(se).__name__}",
-        flush=True,
-    )
-    return se
-
-
-def test_shared_row_invariance(layer, pools, shared):
-    bad, cases = [], 0
-    with torch.inference_mode():
-        for seed in range(12):
-            x6, _ = pools["emb" if seed % 2 else "gauss"].take(6, 8 + 4 * seed, seed=seed)
-            out6 = shared(x6)
-            for g in range(6):
-                cases += 1
-                if not _eq(out6[g : g + 1], shared(x6[g : g + 1].contiguous())):
-                    bad.append((seed, g))
-            x8, _ = pools["gauss"].take(8, 40, seed=100 + seed)
-            out8 = shared(x8)
-            for g1, g2 in ((0, 7), (6, 1), (3, 4)):
-                idx = torch.tensor([g1, g2], device="cuda")
-                cases += 1
-                if not _eq(out8[idx], shared(x8[idx].contiguous())):
-                    bad.append((seed, g1, g2))
-        other = shared(x6[1:2].contiguous())
-        assert not _eq(out6[0:1], other) or _eq(x6[0], x6[1]), "negative control not caught"
-    torch.cuda.synchronize()
-    print(
-        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE shared M=6/8 vs M=1/2 cases={cases} mismatching={len(bad)} {bad[:4]}",
-        flush=True,
-    )
-    assert not bad
-
-
-def _capture(fn, *static):
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        fn(*static)
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        out = fn(*static)
-    torch.cuda.synchronize()
-    return graph, out
-
-
-@pytest.mark.parametrize("kind", ["native", "draft_routed_moe"])
-def test_graph_replay(layer, pools, shared, kind):
-    """R = 6 and R = 1 graphs of the routed MoE and of the shared expert, replayed with new rows."""
-    x6s = torch.empty(6, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
-    l6s = torch.empty(6, fx.NUM_EXPERTS, dtype=torch.float32, device="cuda")
-    x1s = torch.empty(1, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
-    l1s = torch.empty(1, fx.NUM_EXPERTS, dtype=torch.float32, device="cuda")
-    x0, _ = pools["gauss"].take(6, 31, seed=9)
-    x6s.copy_(x0)
-    l6s.copy_(layer.logits(x0))
-    x1s.copy_(x0[:1])
-    l1s.copy_(layer.logits(x0[:1]))
-    with torch.inference_mode():
-        g6, o6 = _capture(lambda x, lg: _routed(kind, layer, x, lg), x6s, l6s)
-        g1, o1 = _capture(lambda x, lg: _routed(kind, layer, x, lg), x1s, l1s)
-        s6, so6 = _capture(shared, x6s)
-        s1, so1 = _capture(shared, x1s)
-        bad, cases = [], 0
-        for trial, target in enumerate(D_TARGETS):
-            x, _ = pools["emb" if trial % 2 else "gauss"].take(6, target, seed=300 + trial)
-            lg = layer.logits(x)
-            x6s.copy_(x)
-            l6s.copy_(lg)
-            g6.replay()
-            s6.replay()
-            r6, sh6 = o6.clone(), so6.clone()
-            for g in range(6):
-                x1s.copy_(x[g : g + 1])
-                l1s.copy_(lg[g : g + 1])
-                g1.replay()
-                s1.replay()
-                cases += 1
-                if not (_eq(r6[g : g + 1], o1) and _eq(sh6[g : g + 1], so1)):
-                    bad.append((target, g))
-        torch.cuda.synchronize()
-    print(
-        f"GLM_MOE_DSA_DRAFT_LIVE_ROWS_ORACLE graph {kind} routed+shared R=6 vs R=1 cases={cases} "
-        f"mismatching={len(bad)} {bad[:4]}",
-        flush=True,
-    )
-    assert not bad
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py
deleted file mode 100644
index dae482dcc4ac..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_fixtures.py
+++ /dev/null
@@ -1,315 +0,0 @@
-# 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.
-"""Real-weight fixtures and the native oracle for the draft_routed_moe tests.
-
-The fixture is the MTP layer (``model.layers.78``, BF16 in the GLM-5.2 NVFP4
-checkpoint) as TP rank ``r`` of TP8/EP1 sees it: every routed expert's
-gate/up rows ``[256r, 256r+256)`` and down columns ``[256r, 256r+256)``, the
-router weight and the BF16 router bias. The native module is the stock
-FlashInfer TRT-LLM-Gen BF16 leaf (``trtllm_gen_leaf(None)``) loaded through its
-own ``load_weights``/``post_load_weights``, so its buffers carry exactly the
-shuffled BlockMajorK layouts of a served model. Slicing the checkpoint before
-loading a ``tp_size=1`` leaf is what the TP8 loader's ``load_weight_shard``
-does (COLUMN for w1/w3, ROW for w2).
-"""
-
-import json
-import os
-import tempfile
-from dataclasses import dataclass
-from typing import Dict, List, Tuple
-
-import torch
-from utils.llm_data import llm_models_root
-
-CKPT_DIR = os.environ.get(
-    "GLM_MOE_DSA_DRAFT_ROUTED_MOE_CKPT_DIR",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-CACHE_DIR = os.environ.get(
-    "GLM_MOE_DSA_DRAFT_ROUTED_MOE_FIXTURE_CACHE",
-    os.path.join(tempfile.gettempdir(), "glm_moe_dsa_draft_routed_moe_fixtures"),
-)
-
-MTP_LAYER = 78
-HIDDEN = 6144
-INTER = 2048
-NUM_EXPERTS = 256
-TOP_K = 8
-TP_SIZE = 8
-ROUTED_SCALING = 2.5
-RMS_EPS = 1e-5
-
-
-class Reader:
-    """Opens each safetensors shard once and slices tensors lazily."""
-
-    def __init__(self):
-        from safetensors import safe_open
-
-        self._safe_open = safe_open
-        with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
-            self._map = json.load(f)["weight_map"]
-        self._open = {}
-
-    def _handle(self, name: str):
-        shard = self._map[name]
-        if shard not in self._open:
-            self._open[shard] = self._safe_open(
-                os.path.join(CKPT_DIR, shard), framework="pt", device="cpu"
-            )
-        return self._open[shard]
-
-    def full(self, name: str) -> torch.Tensor:
-        return self._handle(name).get_tensor(name)
-
-    def rows(self, name: str, start: int, stop: int) -> torch.Tensor:
-        return self._handle(name).get_slice(name)[start:stop].clone()
-
-    def cols(self, name: str, start: int, stop: int) -> torch.Tensor:
-        return self._handle(name).get_slice(name)[:, start:stop].clone()
-
-
-def load_rank(tp_rank: int, layer: int = MTP_LAYER) -> Dict[str, torch.Tensor]:
-    """Rank-sliced raw tensors of the MTP layer's MoE (CPU, cached under ``CACHE_DIR``)."""
-    os.makedirs(CACHE_DIR, exist_ok=True)
-    path = os.path.join(CACHE_DIR, f"glm52_L{layer}_bf16_tp{TP_SIZE}r{tp_rank}.pt")
-    if os.path.exists(path):
-        return torch.load(path, map_location="cpu")
-    rd = Reader()
-    inter_rank = INTER // TP_SIZE
-    r0, r1 = tp_rank * inter_rank, (tp_rank + 1) * inter_rank
-    pre = f"model.layers.{layer}.mlp"
-    out: Dict[str, torch.Tensor] = {}
-    for e in range(NUM_EXPERTS):
-        out[f"{e}.w1.weight"] = rd.rows(f"{pre}.experts.{e}.gate_proj.weight", r0, r1)
-        out[f"{e}.w3.weight"] = rd.rows(f"{pre}.experts.{e}.up_proj.weight", r0, r1)
-        out[f"{e}.w2.weight"] = rd.cols(f"{pre}.experts.{e}.down_proj.weight", r0, r1)
-    out["gate.weight"] = rd.full(f"{pre}.gate.weight")
-    out["gate.e_score_correction_bias"] = rd.full(f"{pre}.gate.e_score_correction_bias")
-    out["post_attention_layernorm.weight"] = rd.full(
-        f"model.layers.{layer}.post_attention_layernorm.weight"
-    )
-    tmp = path + f".tmp{os.getpid()}"
-    torch.save(out, tmp)
-    os.replace(tmp, path)
-    return out
-
-
-@dataclass
-class NativeBf16Layer:
-    """The native FlashInfer TRT-LLM-Gen BF16 MoE leaf of one TP rank plus its router."""
-
-    module: torch.nn.Module
-    gate_weight: torch.Tensor  # [256, 6144] bf16
-    bias: torch.Tensor  # [256] bf16 (the TRTLLM gate dtype)
-    norm_weight: torch.Tensor  # [6144] bf16
-    tp_rank: int
-
-    def logits(self, x: torch.Tensor) -> torch.Tensor:
-        """FP32 router logits of ``x`` (FP32 GEMM; any FP32 logits are valid oracle inputs)."""
-        return (x.float() @ self.gate_weight.float().t()).contiguous()
-
-    def native(self, x: torch.Tensor, logits: torch.Tensor, do_finalize: bool = True):
-        """``trtllm_bf16_moe`` exactly as ``FlashinferTrtllmGenBf16Impl.run_moe`` calls it."""
-        from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.kernel_inputs import (
-            to_trtllm_gen_act_type,
-        )
-
-        m = self.module
-        return m.op_backend.run_bf16_moe(
-            router_logits=logits,
-            routing_bias=self.bias,
-            hidden_states=x,
-            gemm1_weights=m.w3_w1_weight,
-            gemm2_weights=m.w2_weight,
-            num_experts=m.num_slots,
-            top_k=TOP_K,
-            n_group=1,
-            topk_group=1,
-            intermediate_size=m.intermediate_size_per_partition,
-            local_expert_offset=m.slot_start,
-            local_num_experts=m.expert_size_per_partition,
-            routed_scaling_factor=ROUTED_SCALING,
-            routing_method_type=m.routing_method.routing_method_type,
-            topk_weights=None,
-            topk_ids=None,
-            gated_act_type=to_trtllm_gen_act_type(m.activation_type),
-            output=None,
-            use_shuffled_weight=m.quant_method.use_shuffled_weight,
-            weight_layout=m.quant_method.weight_layout,
-            do_finalize=do_finalize,
-        )
-
-
-def build_native(
-    tp_rank: int = 0, max_num_tokens: int = 128, identity_fc2: bool = False
-) -> NativeBf16Layer:
-    """The rank's native leaf.
-
-    ``identity_fc2`` replaces every expert's down projection by ``W[i, i % 256] = 1`` so that
-    ``gemm2_output[:, :256]`` is the native FC1 + SwiGLU intermediate (diagnostics only).
-    """
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
-    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
-    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import trtllm_gen_leaf
-    from tensorrt_llm.mapping import Mapping
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-
-    raw = load_rank(tp_rank)
-    bias = raw["gate.e_score_correction_bias"].to(torch.bfloat16).cuda()
-    routing = DeepSeekV3MoeRoutingMethod(
-        top_k=TOP_K,
-        n_group=1,
-        topk_group=1,
-        routed_scaling_factor=ROUTED_SCALING,
-        callable_e_score_correction_bias=lambda: bias,
-        is_fused=True,
-    )
-    model_config = ModelConfig(
-        quant_config=QuantConfig(),
-        mapping=Mapping(world_size=1, tp_size=1, rank=0),
-        moe_backend="TRTLLM",
-        max_num_tokens=max_num_tokens,
-    )
-    leaf_cls = trtllm_gen_leaf(None)
-    module = leaf_cls(
-        routing_method=routing,
-        num_experts=NUM_EXPERTS,
-        hidden_size=HIDDEN,
-        intermediate_size=INTER // TP_SIZE,
-        dtype=torch.bfloat16,
-        reduce_results=False,
-        model_config=model_config,
-        init_load_balancer=False,
-        weight_loading_mode=MoEWeightLoadingMode.VANILLA,
-    )
-    weights = {k: v for k, v in raw.items() if k[0].isdigit()}
-    if identity_fc2:
-        eye = torch.zeros(HIDDEN, INTER // TP_SIZE, dtype=torch.bfloat16)
-        eye[torch.arange(HIDDEN), torch.arange(HIDDEN) % (INTER // TP_SIZE)] = 1.0
-        weights = {k: (eye if k.endswith(".w2.weight") else v) for k, v in weights.items()}
-    module.load_weights([weights])
-    module.post_load_weights()
-    module = module.cuda()
-    return NativeBf16Layer(
-        module=module,
-        gate_weight=raw["gate.weight"].to(torch.bfloat16).cuda().contiguous(),
-        bias=bias,
-        norm_weight=raw["post_attention_layernorm.weight"].to(torch.bfloat16).cuda(),
-        tp_rank=tp_rank,
-    )
-
-
-def hidden_rows(layer: NativeBf16Layer, num: int, seed: int, scale: float = 1.0) -> torch.Tensor:
-    """``num`` RMS-normalized Gaussian rows times the layer's post-attention norm weight (BF16)."""
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    h = torch.randn(num, HIDDEN, device="cuda", generator=g) * scale
-    y = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + RMS_EPS)
-    return (y * layer.norm_weight.float()).to(torch.bfloat16).contiguous()
-
-
-def chain_like_rows(layer: NativeBf16Layer, num: int, seed: int) -> torch.Tensor:
-    """hidden_rows with the distinct-expert count of the draft chain at R = 6.
-
-    Independent rows select about 40 distinct experts, the chain's six draft rows about 31: rows 4 and 5 repeat
-    rows 0 and 1. Other row counts are hidden_rows unchanged.
-    """
-    x = hidden_rows(layer, num, seed)
-    if num == 6:
-        x[4:] = x[:2]
-    return x
-
-
-def distinct_experts(layer: NativeBf16Layer, logits: torch.Tensor) -> int:
-    """Distinct experts of the top-8 selection (sigmoid + correction bias), for reports."""
-    scores = torch.sigmoid(logits.float()) + layer.bias.float()
-    return int(scores.topk(8, dim=-1).indices.unique().numel())
-
-
-def top8_sets(bias: torch.Tensor, logits: torch.Tensor) -> List[frozenset]:
-    """Per row, the set of its top-8 experts (sigmoid + correction bias) as frozensets (for row selection)."""
-    scores = torch.sigmoid(logits.float()) + bias.float().to(logits.device)
-    return [frozenset(r) for r in scores.topk(TOP_K, dim=-1).indices.cpu().tolist()]
-
-
-def select_distinct(
-    sets: List[frozenset], rows: int, target: int, seed: int, steps: int = 600, sample: int = 256
-) -> Tuple[List[int], int]:
-    """Indices into ``sets`` (repeats allowed) of ``rows`` rows whose top-8 sets cover ``target`` distinct experts.
-
-    Seeded local search: start from random rows, then repeatedly replace one row by the candidate (a random sample
-    of the pool plus the selection's other rows, whose repeat lowers the count) that brings the union closest to the
-    target. Returns ``(indices, achieved count)``; the count differs from the target only when the search does not
-    reach it.
-    """
-    import random
-
-    rng = random.Random(seed)
-    n = len(sets)
-    pick = [rng.randrange(n) for _ in range(rows)]
-
-    def count(p):
-        u = set()
-        for i in p:
-            u |= sets[i]
-        return len(u)
-
-    best, best_d = list(pick), count(pick)
-    for _ in range(steps):
-        d = count(pick)
-        if abs(d - target) < abs(best_d - target):
-            best, best_d = list(pick), d
-        if d == target:
-            return pick, d
-        pos = rng.randrange(rows)
-        cands = rng.sample(range(n), min(n, sample)) + [pick[j] for j in range(rows) if j != pos]
-        scored = [
-            (abs(count(pick[:pos] + [c] + pick[pos + 1 :]) - target), rng.random(), c)
-            for c in cands
-        ]
-        pick[pos] = min(scored)[2]
-    d = count(pick)
-    if abs(d - target) < abs(best_d - target):
-        best, best_d = list(pick), d
-    return best, best_d
-
-
-class RowPool:
-    """Candidate rows and their top-8 sets, from which ``take`` draws rows with a chosen distinct-expert count."""
-
-    def __init__(self, xs: torch.Tensor, sets: List[frozenset]):
-        self.xs, self.sets = xs, sets
-
-    @classmethod
-    def gaussian(cls, layer: NativeBf16Layer, num: int = 1024, seed: int = 10_000) -> "RowPool":
-        xs = hidden_rows(layer, num, seed)
-        return cls(xs, top8_sets(layer.bias, layer.logits(xs)))
-
-    def take(self, rows: int, target: int, seed: int) -> Tuple[torch.Tensor, int]:
-        """``rows`` rows (repeats allowed) covering ``target`` distinct experts, and the count achieved."""
-        idx, d = select_distinct(self.sets, rows, target, seed)
-        return self.xs[torch.tensor(idx, device=self.xs.device)].contiguous(), d
-
-
-def embedding_rows(layer: NativeBf16Layer, token_ids: List[int]) -> torch.Tensor:
-    """RMSNorm(embedding) * post_attention_layernorm.weight for real token ids (BF16)."""
-    rd = Reader()
-    name = "model.embed_tokens.weight"
-    sl = rd._handle(name).get_slice(name)
-    emb = torch.stack([sl[t : t + 1][0] for t in token_ids]).cuda().float()
-    y = emb * torch.rsqrt(emb.pow(2).mean(-1, keepdim=True) + RMS_EPS)
-    return (y * layer.norm_weight.float()).to(torch.bfloat16).contiguous()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py
deleted file mode 100644
index 9e57868210e0..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/draft_moe_jit.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# 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.
-"""Load ``trtllm::glm_moe_dsa_draft_routed_moe`` for tests.
-
-A TensorRT-LLM build of this branch registers the op in ``th_common``. On an
-older image (development overlay) it is compiled from this checkout's C++
-sources with ``torch.utils.cpp_extension``, unmodified and with the same op
-schema; results obtained that way are labelled DEV_JIT.
-"""
-
-import hashlib
-import os
-from pathlib import Path
-from typing import Callable, Sequence
-
-import torch
-
-_REPO = Path(__file__).resolve().parents[9]
-_SOURCES = (
-    _REPO / "cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp",
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu",
-)
-_HEADERS = (_REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h",)
-OP_NAME = "glm_moe_dsa_draft_routed_moe"
-_NAME = "draft_moe_jit"
-
-
-def _digest() -> str:
-    h = hashlib.sha256()
-    for path in (*_SOURCES, *_HEADERS):
-        h.update(path.read_bytes())
-    return h.hexdigest()[:16]
-
-
-def ensure_op(
-    build_root: str,
-    is_builder: bool = True,
-    barrier: Callable[[], None] = lambda: None,
-    defines: Sequence[str] = (),
-) -> str:
-    """Make the op available in this process; returns where it came from.
-
-    ``is_builder`` must be true on exactly one process per shared build root;
-    ``barrier`` synchronizes the processes after that one has built it.
-    ``defines`` (e.g. ``("GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1",)``) builds a probe variant of
-    the op; it needs a process in which the op is not registered yet.
-    """
-    import tensorrt_llm  # noqa: F401  (loads th_common and the trtllm namespace)
-
-    if hasattr(torch.ops.trtllm, OP_NAME):
-        if defines:
-            raise RuntimeError(
-                f"{OP_NAME} is already registered; a probe build needs a fresh process"
-            )
-        barrier()
-        return "builtin"
-    tag = _digest()
-    name = _NAME
-    if defines:
-        tag += "-" + hashlib.sha256(",".join(defines).encode()).hexdigest()[:8]
-        name += "_probe"
-    build_dir = Path(build_root) / f"draft_routed_moe-{tag}"
-    library = build_dir / f"{name}.so"
-    if is_builder and not library.exists():
-        from torch.utils.cpp_extension import load
-
-        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
-        build_dir.mkdir(parents=True, exist_ok=True)
-        extra = [f"-D{d}" for d in defines]
-        load(
-            name=name,
-            sources=[str(s) for s in _SOURCES],
-            extra_include_paths=[str(_REPO / "cpp/include"), str(_REPO / "cpp")],
-            extra_cflags=["-O3", "-std=c++17", *extra],
-            extra_cuda_cflags=[
-                "-O3",
-                "-std=c++17",
-                "--expt-relaxed-constexpr",
-                "-Xptxas=-v",
-                *extra,
-            ],
-            build_directory=str(build_dir),
-            is_python_module=False,
-            verbose=True,
-        )
-    barrier()
-    if not hasattr(torch.ops.trtllm, OP_NAME):
-        torch.ops.load_library(str(library))
-    return f"DEV_JIT {library}"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py
deleted file mode 100644
index 115e9acfae19..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_adapter.py
+++ /dev/null
@@ -1,152 +0,0 @@
-# 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.
-"""The draft_routed_moe call site in the target's MoE.
-
-* a plan runs draft_routed_moe alone and next to F's GEN set, draft_routed_moe extends no
-  package and has a single numerical form; deferred_input_prep's gate
-  classifies draft_routed_moe's launch-configuration selectors as draft_routed_moe's;
-* ``GlmMoE.compute_routed_output`` dispatches to draft_routed_moe only when it is selected, before
-  routed_moe, and keeps the native path when draft_routed_moe declines;
-* an NVFP4 MoE layer (every target layer) declines draft_routed_moe before any device work,
-  so routed_moe / hop_routing see exactly the calls they see in F;
-* through ``GlmMoE.compute_routed_output`` on a decode step the MTP rank
-  shard returns draft_routed_moe's output, bit-identical to the native leaf.
-"""
-
-import inspect
-import os
-import re
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "..", "..", "shared", "kernels"))
-import plan_testlib  # noqa: E402
-
-F_GEN = (
-    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,router_gemm,"
-    "post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,rope_kv_append,"
-    "draft_loop,shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
-    "attention_output_tail_prefetch,attention_output_into_hop,routed_moe_into_hop,"
-    "deferred_norm,dense_mlp,deferred_input_prep,hop_routing,oproj_weight_staging,"
-    "fused_kv_indexer_projection,host_page_retention"
-)
-
-
-def _dispatch():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    return dispatch
-
-
-def test_selector_facts():
-    d = _dispatch()
-    assert "draft_routed_moe" in d.SELECTOR_PATHS
-    assert "draft_routed_moe" not in d.REQUIRED_FEATURES
-    assert "draft_routed_moe" not in d.NUMERICAL_VARIANTS
-    assert plan_testlib.plan_running(d, "draft_routed_moe").enabled("draft_routed_moe")
-    no_input_prep = frozenset(F_GEN.split(",")) - {"deferred_input_prep"}
-    assert plan_testlib.plan_running(d, no_input_prep | {"draft_routed_moe"}).enabled(
-        "draft_routed_moe"
-    )
-
-
-def test_call_site_order():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmMoE,
-    )
-
-    src = inspect.getsource(GlmMoE.compute_routed_output)
-    # The formatter wraps long statements: compare with the whitespace collapsed.
-    src = re.sub(r"(?<=[(\[])\s+|\s+(?=[)\]])", "", " ".join(src.split()))
-    draft_routed_moe = src.index('is_enabled("draft_routed_moe")')
-    assert draft_routed_moe < src.index('is_enabled("routed_moe")')
-    assert src.index('get_impl("draft_routed_moe", "routed_output")') > draft_routed_moe
-
-
-def test_nvfp4_layer_declines():
-    if not torch.cuda.is_available():
-        pytest.skip("needs a GPU")
-    import draft_moe_jit
-    import routed_moe_fixtures as wfx
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    # With the op loaded, the decline names the leaf type, not the missing op.
-    draft_moe_jit.ensure_op(
-        os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT", "/tmp/draft_routed_moe_jit")
-    )
-    layer = wfx.build_native_layer(10, 0)
-    x = torch.zeros(6, 6144, dtype=torch.bfloat16, device="cuda")
-    lg = torch.zeros(6, 256, dtype=torch.float32, device="cuda")
-    ok, reason = rm.eligibility(layer.module, layer.bias, x, lg, True)
-    print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_NVFP4_DECLINE {reason}", flush=True)
-    assert not ok and "BF16 leaf" in reason
-
-
-def test_patched_call_site_bit_exact(monkeypatch):
-    """``GlmMoE.compute_routed_output`` on a decode step, on a stand-in module that holds the MTP leaf."""
-    if not torch.cuda.is_available():
-        pytest.skip("needs a GPU")
-    import draft_moe_fixtures as fx
-    import draft_moe_jit
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmMoE,
-    )
-
-    draft_moe_jit.ensure_op(
-        os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT", "/tmp/draft_routed_moe_jit")
-    )
-    layer = fx.build_native(0)
-    plan_testlib.select(_dispatch(), "draft_routed_moe", monkeypatch)
-
-    class Gate:
-        e_score_correction_bias = layer.bias
-
-    class Stub:
-        use_dp = False
-        mapping = None
-        router_gemm = None
-        gate = Gate()
-        experts = layer.module
-
-    calls = {"n": 0}
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    inner = dispatch.get_impl("draft_routed_moe", "routed_output")
-
-    def counted(*a, **k):
-        out = inner(*a, **k)
-        calls["n"] += out is not None
-        return out
-
-    dispatch.register("draft_routed_moe", "routed_output", counted)
-    try:
-        for rows in (1, 6):
-            x = fx.hidden_rows(layer, rows, seed=40 + rows)
-            lg = layer.logits(x)
-            with dispatch.step(dispatch.Phase.DECODE):
-                out = GlmMoE.compute_routed_output(Stub(), x, None, None, True, lg)
-            assert torch.equal(out.view(torch.int16), layer.native(x, lg).view(torch.int16))
-    finally:
-        dispatch.register("draft_routed_moe", "routed_output", inner)
-    assert calls["n"] == 2
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py
deleted file mode 100644
index 68b5934c391d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/test_draft_moe_oracle.py
+++ /dev/null
@@ -1,775 +0,0 @@
-# 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.
-"""draft_routed_moe against the native ``trtllm_bf16_moe`` on the real MTP-layer rank shard (single GPU).
-
-Every comparison is bit for bit on the BF16 outputs. On a mismatch the stage
-diagnostics compare the routing weights with the native ``expert_weights`` and
-draft_routed_moe's per-(token, slot) FC2 rows (``debug_fc2``) with the native unfinalized
-``gemm2_output`` through ``expanded_idx_to_permuted_idx``.
-
-Environment: ``GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT`` (JIT build directory when the image has no
-op), ``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TP_RANK`` (rank shard, default 0),
-``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TRIALS`` (calls per row count).
-"""
-
-import contextlib
-import json
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-import draft_moe_fixtures as fx  # noqa: E402
-import draft_moe_jit  # noqa: E402
-
-TRIALS = int(os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_TRIALS", "24"))
-
-# Keys of a launch configuration (``_launch``): FC1 and FC2 "ctas,warps" and the flag bits, each
-# "<value>" or "<one row>+<more rows>".
-FC1, FC2, FLAGS = "fc1", "fc2", "flags"
-
-# Kernel families: the default selectors (TMA-fed kernels: 64 FC1 CTAs of 4 MMA warps at R = 1, 152 CTAs with two
-# chains per warp at R > 1 and the triples path above 38 distinct experts with its early trigger, FC1 loads L2
-# evict-first at R > 1, the 3-warp FC2 on 128 CTAs, FC2 weight prefetch at R = 1), the
-# cp.async kernels, the TMA-fed kernels on other grids (flag 64) and the TMA-fed kernels with two chains per FC1 warp,
-# two experts per FC2 warp step, FC1 L2 prefetch and the FC2 weight prefetch (flags 8, 16, 128, 2; the FC1 ring of
-# 32 KB stages is 6 deep at R <= 2, 4 deep at R = 3..7 and 3 deep at R = 8), and the FC1 triples path on a 120-CTA
-# grid (flag 4096: R > 1 calls whose 128-row tiles exceed the grid but whose 64-row sub-tiles fit three per CTA; the
-# others run the regular path, e.g. R = 8 with about 50 distinct experts; flag 8192: the early trigger on that path
-# only), and the default R > 1 kernels with the activation streamed through the FC1 ring (flag 65536). All must be
-# bit-exact.
-FAMILIES = {
-    "default": {},
-    "cpasync": {
-        FC1: "128,2",
-        FC2: "96,4",
-        FLAGS: "2",
-    },
-    "tma": {
-        FC1: "128,2+152,4",
-        FC2: "96,4",
-        FLAGS: "96",
-    },
-    "tmapairs": {
-        FC1: "64,2+152,4",
-        FC2: "96,4",
-        FLAGS: "250",
-    },
-    "triples": {
-        FC1: "64,4+120,4",
-        FC2: "128,3",
-        FLAGS: "66+12360",
-    },
-    "streamx": {
-        FC1: "64,4+152,4",
-        FC2: "128,3",
-        FLAGS: "66+77896",
-    },
-}
-DELAYED_FLAGS = {
-    "default": ("66+45128", "70+76"),
-    "cpasync": ("2", "6"),
-    "tma": ("96", "100"),
-    "tmapairs": ("250", "254"),
-    "triples": ("66+12360", "70+4172"),
-    "streamx": ("66+77896", "70+77900"),
-}
-
-
-def _per_rows(raw, parse):
-    """``(one row, more rows)`` of a "<value>" or "<one row>+<more rows>" entry."""
-    parts = raw.split("+")
-    assert len(parts) <= 2, raw
-    return parse(parts[0]), parse(parts[-1])
-
-
-def _ctas_warps(raw):
-    ctas, warps = (int(v) for v in raw.split(","))
-    return ctas, warps
-
-
-@contextlib.contextmanager
-def _launch(config):
-    """``config`` set over draft_routed_moe's launch defaults for the duration of the block."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    saved = rm.DEFAULT_FC1, rm.DEFAULT_FC2, rm.DEFAULT_FLAGS
-    config = config or {}
-    try:
-        if FC1 in config:
-            rm.DEFAULT_FC1 = _per_rows(config[FC1], _ctas_warps)
-        if FC2 in config:
-            rm.DEFAULT_FC2 = _per_rows(config[FC2], _ctas_warps)
-        if FLAGS in config:
-            rm.DEFAULT_FLAGS = _per_rows(config[FLAGS], int)
-        yield
-    finally:
-        rm.DEFAULT_FC1, rm.DEFAULT_FC2, rm.DEFAULT_FLAGS = saved
-
-
-@pytest.fixture(scope="module")
-def layer():
-    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
-        pytest.skip("needs an SM100-family GPU")
-    origin = draft_moe_jit.ensure_op(
-        os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_JIT_ROOT", "/tmp/draft_routed_moe_jit")
-    )
-    print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_OP_ORIGIN {origin}", flush=True)
-    return fx.build_native(int(os.environ.get("GLM_MOE_DSA_DRAFT_ROUTED_MOE_TP_RANK", "0")))
-
-
-def _draft_moe(layer, x, logits, debug=False, config=None):
-    """draft_routed_moe's output, its debug FC2 rows and whether the TMA-fed kernels ran."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    with _launch(config):
-        dbg = (
-            torch.full(
-                (x.shape[0], fx.TOP_K, fx.HIDDEN), float("nan"), dtype=torch.bfloat16, device="cuda"
-            )
-            if debug
-            else None
-        )
-        tma = bool(rm.launch_config(x.shape[0])[2] & rm.FLAG_TMA)
-        out = rm.run(layer.module, layer.bias, x, logits, debug_fc2=dbg)
-        return out, dbg, tma
-
-
-def _bits(t):
-    return t.contiguous().view(torch.int16)
-
-
-_IDENTITY = {}
-
-
-def _identity_layer(layer):
-    """The same rank's native leaf with identity-like down projections (built on the first mismatch only)."""
-    if layer.tp_rank not in _IDENTITY:
-        _IDENTITY[layer.tp_rank] = fx.build_native(layer.tp_rank, identity_fc2=True)
-    return _IDENTITY[layer.tp_rank]
-
-
-def _h_image_order(rows: int = 8, inter: int = 256) -> torch.Tensor:
-    """Index map of an intermediate entry stored as its shared-memory image (TMA family): plain [tok, j] -> image."""
-    tok = torch.arange(rows)[:, None]
-    j = torch.arange(inter)[None, :]
-    u = j // 8
-    return tok * inter + (((u & ~7) | ((u & 7) ^ (tok & 7))) * 8) + (j & 7)
-
-
-def _diagnose(layer, x, logits, out_draft_moe, dbg, tma=False):
-    """Stage-level comparison of one mismatching call (returns a dict of counts).
-
-    Reads draft_routed_moe's routing record and intermediate scratch, which still hold this call's
-    values.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    scratch = rm._scratch(x.device, create=False)
-    route = scratch.route.cpu()
-    h_draft_moe = scratch.h.view(64, 8 * (fx.INTER // fx.TP_SIZE)).clone()
-    if tma:
-        h_draft_moe = h_draft_moe[:, _h_image_order().flatten().to(h_draft_moe.device)]
-    h_draft_moe = h_draft_moe.view(64, 8, fx.INTER // fx.TP_SIZE)
-    gemm2, weights, idx = layer.native(x, logits, do_finalize=False)[:3]
-    rows = x.shape[0]
-    idx = idx.reshape(-1).long()
-    w_draft_moe = route[129 : 129 + rows * fx.TOP_K].view(torch.float32).reshape(rows, fx.TOP_K)
-    slot_d = route[65 : 65 + rows * fx.TOP_K].reshape(rows, fx.TOP_K).long()
-    w_nat = weights.float().cpu().reshape(rows, fx.TOP_K)
-    fc2_nat = gemm2[idx].reshape(rows, fx.TOP_K, fx.HIDDEN)
-    fc2_diff = (_bits(fc2_nat) != _bits(dbg)).sum(-1)
-    # Native FC1 + SwiGLU of every (token, slot): gemm2 output of the identity leaf, columns 0..255.
-    ident = _identity_layer(layer)
-    g_id, _, idx_id = ident.native(x, logits, do_finalize=False)[:3]
-    h_nat = g_id[idx_id.reshape(-1).long()].reshape(rows, fx.TOP_K, fx.HIDDEN)[
-        ..., : fx.INTER // fx.TP_SIZE
-    ]
-    h_mine = torch.stack(
-        [torch.stack([h_draft_moe[slot_d[t, k], t] for k in range(fx.TOP_K)]) for t in range(rows)]
-    )
-    h_diff = (h_nat.float() != h_mine.float()).sum(-1)
-    return {
-        "weights_differ": int((w_draft_moe != w_nat).sum()),
-        "fc1_h_differ_per_slot": h_diff.tolist(),
-        "fc2_rows_differ_per_slot": fc2_diff.tolist(),
-        "out_elems_differ": int((_bits(out_draft_moe) != _bits(layer.native(x, logits))).sum()),
-    }
-
-
-def _check(layer, x, logits, label, config=None):
-    native = layer.native(x, logits)
-    out, dbg, tma = _draft_moe(layer, x, logits, debug=True, config=config)
-    torch.cuda.synchronize()
-    differ = int((_bits(out) != _bits(native)).sum())
-    if differ:
-        info = _diagnose(layer, x, logits, out, dbg, tma)
-        print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_MISMATCH {label} " + json.dumps(info), flush=True)
-    return differ
-
-
-@pytest.mark.parametrize("family", list(FAMILIES))
-@pytest.mark.parametrize("rows", [1, 2, 3, 4, 5, 6, 7, 8])
-def test_bit_exact_gaussian(layer, rows, family):
-    bad = 0
-    for trial in range(TRIALS):
-        x = fx.hidden_rows(layer, rows, seed=1000 * rows + trial)
-        bad += (
-            _check(
-                layer,
-                x,
-                layer.logits(x),
-                f"gauss {family} R={rows} trial={trial}",
-                config=FAMILIES[family],
-            )
-            > 0
-        )
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE gaussian {family} R={rows} calls={TRIALS} mismatching_calls={bad}",
-        flush=True,
-    )
-    assert bad == 0
-
-
-@pytest.mark.parametrize("family", list(FAMILIES))
-@pytest.mark.parametrize("rows", [1, 6])
-def test_bit_exact_embeddings(layer, rows, family):
-    """Consecutive real token embeddings (normalized) as activation rows."""
-    ids = list(range(1000, 1000 + 64 * rows))
-    bad = 0
-    for trial in range(8):
-        x = fx.embedding_rows(layer, ids[trial * rows : (trial + 1) * rows])
-        bad += (
-            _check(
-                layer,
-                x,
-                layer.logits(x),
-                f"emb {family} R={rows} trial={trial}",
-                config=FAMILIES[family],
-            )
-            > 0
-        )
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE embeddings {family} R={rows} calls=8 mismatching_calls={bad}",
-        flush=True,
-    )
-    assert bad == 0
-
-
-@pytest.mark.parametrize("family", list(FAMILIES))
-def test_routing_ties(layer, family):
-    """Equal selection keys.
-
-    Equal logits for every expert (the bias decides, ties to the lower index), and pairs of experts with equal
-    bias and equal logits.
-    """
-    bad = 0
-    x = fx.hidden_rows(layer, 6, seed=7)
-    flat = torch.zeros(6, fx.NUM_EXPERTS, device="cuda")
-    flat[1] += 3.0
-    flat[2] -= 3.0
-    bias = layer.bias.float()
-    # rows 3..5: logits that make sigmoid + bias equal for groups of experts with equal bias
-    for r in range(3, 6):
-        flat[r] = torch.where(bias == bias[r * 7], 1.0, -2.0)
-    bad += _check(layer, x, flat.contiguous(), "ties", config=FAMILIES[family]) > 0
-    big = torch.randn(6, fx.NUM_EXPERTS, device="cuda") * 40.0
-    bad += _check(layer, x, big.contiguous(), "saturated logits", config=FAMILIES[family]) > 0
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE ties {family} calls=2 mismatching_calls={bad}",
-        flush=True,
-    )
-    assert bad == 0
-
-
-@pytest.mark.parametrize(
-    "config",
-    [
-        {
-            FC1: "148,1",
-            FC2: "148,1",
-            FLAGS: "0",
-        },
-        {
-            FC1: "64,4",
-            FC2: "48,2",
-            FLAGS: "6",
-        },
-        {
-            FC1: "1,1",
-            FC2: "1,1",
-            FLAGS: "2",
-        },
-        {
-            FC1: "128,2",
-            FC2: "96,4",
-            FLAGS: "26",
-        },
-        {
-            FC1: "37,4",
-            FC2: "5,2",
-            FLAGS: "24",
-        },
-        {
-            FC1: "1,1+148,4",
-            FC2: "148,1+5,2",
-            FLAGS: "0+30",
-        },
-        {
-            FC1: "148,2+64,4",
-            FC2: "96,4",
-            FLAGS: "34+58",
-        },
-        {
-            FC1: "128,4",
-            FC2: "37,1",
-            FLAGS: "62",
-        },
-        # Delayed producer of the FC2 grid: one FC1 warp (milliseconds) with the early trigger, so the FC2 CTAs
-        # are resident long before the routing record and the intermediate exist.
-        {
-            FC1: "1,1",
-            FC2: "148,4",
-            FLAGS: "6",
-        },
-        # TMA-fed kernels: one CTA running every tile (persistent loop, ring wrap-around across tiles), deep rings
-        # with the early trigger, other grid shapes, and a one-CTA FC1 producer ahead of a full FC2 grid.
-        {
-            FC1: "1,1",
-            FC2: "1,1",
-            FLAGS: "64",
-        },
-        {
-            FC1: "37,4",
-            FC2: "5,2",
-            FLAGS: "100",
-        },
-        {
-            FC1: "152,2",
-            FC2: "152,4",
-            FLAGS: "96",
-        },
-        {
-            FC1: "64,4+256,1",
-            FC2: "192,2+96,4",
-            FLAGS: "64+96",
-        },
-        {
-            FC1: "1,1",
-            FC2: "152,4",
-            FLAGS: "68",
-        },
-        # TMA-fed with two chains per FC1 warp: one CTA of one warp, an odd grid with the early trigger, L2 prefetch
-        # and the early FC2 prefetch, grids with more CTAs than tiles, and one 4-warp CTA ahead of a full FC2 grid.
-        {
-            FC1: "1,1",
-            FC2: "1,1",
-            FLAGS: "72",
-        },
-        {
-            FC1: "37,4",
-            FC2: "5,2",
-            FLAGS: "460",
-        },
-        {
-            FC1: "152,2+100,4",
-            FC2: "152,4",
-            FLAGS: "226+98",
-        },
-        {
-            FC1: "1,4",
-            FC2: "152,4",
-            FLAGS: "236",
-        },
-        # TMA-fed FC2 with 3 MMA warps (48-row groups): the full grid of 128 groups and an odd persistent grid.
-        {
-            FC1: "128,2+152,4",
-            FC2: "128,3",
-            FLAGS: "66+68",
-        },
-        {
-            FC1: "64,4",
-            FC2: "37,3",
-            FLAGS: "64",
-        },
-        # 3-warp FC2 with expert pairs, after FC1 tile pairs at R > 1: the full grid and an odd persistent grid.
-        {
-            FC1: "64,4+152,4",
-            FC2: "128,3",
-            FLAGS: "82+90",
-        },
-        {
-            FC1: "32,4+136,4",
-            FC2: "37,3",
-            FLAGS: "74+88",
-        },
-        # FC1 triples (flag 4096): the default grid with tile pairs, without tile pairs (64-row tiles), with the FC2
-        # weight prefetch after the last copy and right after routing, and a grid too small for one round of
-        # triples (regular path).
-        {
-            FC1: "64,4+152,4",
-            FC2: "128,3",
-            FLAGS: "66+4172",
-        },
-        {
-            FC1: "152,4",
-            FC2: "128,3",
-            FLAGS: "4164",
-        },
-        {
-            FC1: "112,4",
-            FC2: "96,4",
-            FLAGS: "4170+4426",
-        },
-        {
-            FC1: "20,4",
-            FC2: "37,3",
-            FLAGS: "4168",
-        },
-    ],
-)
-def test_launch_configs(layer, config):
-    bad = 0
-    for rows in (1, 3, 6, 8):
-        x = fx.hidden_rows(layer, rows, seed=99 + rows)
-        bad += _check(layer, x, layer.logits(x), f"config {config} R={rows}", config=config) > 0
-    assert bad == 0
-
-
-def _distinct(layer, logits):
-    """Distinct experts of a call (float32 restatement of the selection; only used to pick test cases)."""
-    keys = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5 + layer.bias.float()[None, :]
-    return len(set(keys.topk(fx.TOP_K, dim=-1).indices.flatten().tolist()))
-
-
-@pytest.mark.parametrize("flags", ["24", "58", "16", "80", "90"])
-def test_pairs_odd_distinct(layer, flags):
-    """Expert pairs in FC2 / tile pairs in FC1 with an odd number of distinct experts.
-
-    An unpaired expert per item makes the pair parity alternate across a warp's items; every warp has several
-    items.
-    """
-    config = {
-        FC1: "37,2",
-        FC2: "48,1",
-        FLAGS: flags,
-    }
-    bad, odd, seed = 0, 0, 0
-    while odd < 8 and seed < 400:
-        rows = (2, 3, 5, 7)[seed % 4]
-        x = fx.hidden_rows(layer, rows, seed=7000 + seed)
-        lg = layer.logits(x)
-        seed += 1
-        if _distinct(layer, lg) % 2 == 0:
-            continue
-        odd += 1
-        bad += (
-            _check(layer, x, lg, f"odd D flags={flags} R={rows} seed={seed - 1}", config=config) > 0
-        )
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE pairs flags={flags} odd-D calls={odd} mismatching_calls={bad}",
-        flush=True,
-    )
-    assert odd == 8 and bad == 0
-
-
-# Grids and flags of the triples test: the default grid with tile pairs (one round of triples up to 57 distinct
-# experts; early trigger always or on the triples path only), 64-row tiles (8 per expert), and a 120-CTA grid
-# (triples from 31 to 45 distinct experts).
-TRIPLES_CONFIGS = (
-    ("152,4", "128,3", "4172"),
-    ("152,4", "128,3", "12360"),
-    ("152,4", "128,3", "4160"),
-    ("120,4", "96,4", "4168"),
-    ("152,4", "128,3", "45128"),
-    ("152,4", "128,3", "77896"),
-    ("152,4", "128,3", "110664"),
-)
-
-
-@pytest.mark.parametrize("config", TRIPLES_CONFIGS, ids=lambda c: "/".join(c))
-def test_triples_distinct(layer, config):
-    """R = 5..8 with chosen distinct-expert counts around the FC1 path boundaries.
-
-    The boundaries are 38/39 and 57/58 distinct experts on the 152-CTA grid with tile pairs and 30/31 and 45/46 on
-    120 CTAs; every call is bit-exact whichever FC1 path it takes.
-    """
-    fc1, fc2, flags = config
-    pool = fx.RowPool.gaussian(layer)
-    cases = [
-        (5, 39),
-        (5, 40),
-        (6, 30),
-        (6, 31),
-        (6, 38),
-        (6, 39),
-        (6, 44),
-        (6, 45),
-        (6, 46),
-        (6, 48),
-    ]
-    cases += [(7, 41), (7, 50), (7, 56), (8, 57), (8, 58), (8, 64)]
-    bad, reached = 0, 0
-    for i, (rows, target) in enumerate(cases):
-        x, d = pool.take(rows, target, seed=3000 + i)
-        reached += d == target
-        bad += (
-            _check(
-                layer,
-                x,
-                layer.logits(x),
-                f"triples {config} R={rows} D={d}",
-                config={
-                    FC1: fc1,
-                    FC2: fc2,
-                    FLAGS: flags,
-                },
-            )
-            > 0
-        )
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE triples {config} calls={len(cases)} at target D={reached} "
-        f"mismatching_calls={bad}",
-        flush=True,
-    )
-    assert bad == 0 and reached >= len(cases) - 2
-
-
-def _route_weight_bits(layer, x, logits):
-    """(draft_routed_moe routing record weights, native expert_weights) of one call as int32 bit
-    patterns of their FP32 values."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    rows = x.shape[0]
-    out = rm.run(layer.module, layer.bias, x, logits)
-    scratch = rm._scratch(x.device, create=False)
-    fast_w = scratch.route[129 : 129 + rows * fx.TOP_K].clone()
-    nat_w = layer.native(x, logits, do_finalize=False)[1].reshape(-1).float().view(torch.int32)
-    return fast_w, nat_w, out
-
-
-def test_route_record(layer):
-    """draft_routed_moe's routing weights (its routing record) against the native ``expert_weights``, bit for bit.
-
-    102,400 rows (R = 8; logits of several scales, every eighth call rounded so that keys tie), plus the outputs of
-    every 64th call.
-    """
-    calls, rows = 12800, 8
-    x = fx.hidden_rows(layer, rows, seed=4)
-    g = torch.Generator(device="cuda").manual_seed(12)
-    differ = torch.zeros((), dtype=torch.int64, device="cuda")
-    out_bad = 0
-    for i in range(calls):
-        scale = (0.5, 2.0, 6.0, 30.0)[i % 4]
-        lg = torch.randn(rows, fx.NUM_EXPERTS, device="cuda", generator=g) * scale
-        if i % 8 == 7:
-            lg = torch.round(lg * 4.0) / 4.0
-        fast_w, nat_w, out = _route_weight_bits(layer, x, lg.contiguous())
-        differ += (fast_w != nat_w).sum()
-        if i % 64 == 0:
-            out_bad += int((_bits(out) != _bits(layer.native(x, lg.contiguous()))).sum()) > 0
-    differ = int(differ)
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE route record rows={calls * rows} differing weights={differ} "
-        f"mismatching outputs={out_bad}/{calls // 64}",
-        flush=True,
-    )
-    assert differ == 0 and out_bad == 0
-
-
-def test_route_record_nonfinite(layer):
-    """Logits with NaN, +Inf and -Inf entries.
-
-    Routing weights and outputs against native, bit for bit (NaN payloads included).
-    """
-    rows = 8
-    x = fx.hidden_rows(layer, rows, seed=4)
-    g = torch.Generator(device="cuda").manual_seed(13)
-    w_bad, out_bad, calls = 0, 0, 64
-    for i in range(calls):
-        lg = torch.randn(rows, fx.NUM_EXPERTS, device="cuda", generator=g) * 2.0
-        cols = torch.randint(0, fx.NUM_EXPERTS, (3,), device="cuda", generator=g)
-        lg[i % rows, cols[0]] = float("inf")
-        lg[(i + 1) % rows, cols[1]] = float("-inf")
-        if i % 2:
-            lg[(i + 2) % rows, cols[2]] = float("nan")
-        fast_w, nat_w, out = _route_weight_bits(layer, x, lg.contiguous())
-        nat_out = layer.native(x, lg.contiguous())
-        w_differ = (fast_w != nat_w).view(rows, fx.TOP_K)
-        o_differ = (_bits(out) != _bits(nat_out)).sum(-1)
-        if (bool(w_differ.any()) or bool(o_differ.any())) and w_bad + out_bad < 4:
-            from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-                routed_moe as rm,
-            )
-
-            route = rm._scratch(x.device, create=False).route.cpu()
-            distinct = route[1:65]
-            ids = distinct[route[65 : 65 + rows * fx.TOP_K].long()].view(rows, fx.TOP_K)
-            zw, nw = fast_w.view(rows, -1).tolist(), nat_w.view(rows, -1).tolist()
-            zo, no = _bits(out), _bits(nat_out)
-            for r in range(rows):
-                if bool(w_differ[r].any()) or int(o_differ[r]) > 0:
-                    print(
-                        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_NONFINITE call={i} row={r} "
-                        f"nan_in_row={bool(lg[r].isnan().any())} "
-                        f"draft_moe_ids={ids[r].tolist()} fast_w={[hex(v & 0xFFFFFFFF) for v in zw[r]]} "
-                        f"native_w={[hex(v & 0xFFFFFFFF) for v in nw[r]]} out_differ={int(o_differ[r])} "
-                        f"draft_moe_out_nan={int(out[r].isnan().sum())} native_out_nan={int(nat_out[r].isnan().sum())} "
-                        f"draft_moe_out0={hex(int(zo[r, 0]) & 0xFFFF)} native_out0={hex(int(no[r, 0]) & 0xFFFF)}",
-                        flush=True,
-                    )
-        w_bad += bool(w_differ.any())
-        out_bad += bool(o_differ.any())
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE route record non-finite calls={calls} mismatching "
-        f"weights={w_bad} outputs={out_bad}",
-        flush=True,
-    )
-    assert w_bad == 0 and out_bad == 0
-
-
-@pytest.mark.parametrize("family", list(FAMILIES))
-def test_delayed_producer(layer, family):
-    """Inputs written late by their producer.
-
-    A copy writes the inputs after a long device-side sleep on the same stream, eagerly and inside a CUDA graph
-    (PDL edges as captured), with the default and the early-trigger configurations.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    bad = 0
-    flag_sets = DELAYED_FLAGS[family]
-    for flags in flag_sets:
-        with _launch({**FAMILIES[family], FLAGS: flags}):
-            for rows in (1, 6):
-                x0 = fx.hidden_rows(layer, rows, seed=11 + rows)
-                lg0 = layer.logits(x0)
-                ref = layer.native(x0, lg0)
-                xs = torch.zeros_like(x0)
-                ls = torch.zeros_like(lg0)
-                torch.cuda.synchronize()
-                torch.cuda._sleep(20_000_000)
-                xs.copy_(x0)
-                ls.copy_(lg0)
-                out = rm.run(layer.module, layer.bias, xs, ls)
-                torch.cuda.synchronize()
-                bad += int((_bits(out) != _bits(ref)).sum()) > 0
-                xs.zero_()
-                ls.zero_()
-                graph = torch.cuda.CUDAGraph()
-                stream = torch.cuda.Stream()
-                with torch.cuda.stream(stream):
-                    with torch.cuda.graph(graph, stream=stream):
-                        torch.cuda._sleep(20_000_000)
-                        xs.copy_(x0)
-                        ls.copy_(lg0)
-                        out_g = rm.run(layer.module, layer.bias, xs, ls)
-                for _ in range(3):
-                    xs.zero_()
-                    ls.zero_()
-                    torch.cuda.synchronize()
-                    graph.replay()
-                    torch.cuda.synchronize()
-                    bad += int((_bits(out_g) != _bits(ref)).sum()) > 0
-    print(
-        f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE delayed producer {family} calls=16 mismatching={bad}",
-        flush=True,
-    )
-    assert bad == 0
-
-
-@pytest.mark.parametrize("family", list(FAMILIES))
-def test_graph_replay(layer, family):
-    """Captured once, replayed with new inputs copied into the static buffers."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    with _launch(FAMILIES[family]):
-        _graph_replay(layer, rm, family)
-
-
-def _graph_replay(layer, rm, family):
-    for rows in (1, 6):
-        xs = torch.empty(rows, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
-        ls = torch.empty(rows, fx.NUM_EXPERTS, dtype=torch.float32, device="cuda")
-        x0 = fx.hidden_rows(layer, rows, seed=5)
-        xs.copy_(x0)
-        ls.copy_(layer.logits(x0))
-        rm.run(layer.module, layer.bias, xs, ls)  # scratch exists before capture
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        stream = torch.cuda.Stream()
-        with torch.cuda.stream(stream):
-            with torch.cuda.graph(graph, stream=stream):
-                out = rm.run(layer.module, layer.bias, xs, ls)
-        bad = 0
-        for trial in range(6):
-            x = fx.hidden_rows(layer, rows, seed=500 + trial)
-            xs.copy_(x)
-            ls.copy_(layer.logits(x))
-            graph.replay()
-            torch.cuda.synchronize()
-            bad += int((_bits(out) != _bits(layer.native(xs, ls))).sum()) > 0
-        print(
-            f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_ORACLE graph {family} R={rows} replays=6 mismatching={bad}",
-            flush=True,
-        )
-        assert bad == 0
-
-
-def test_eligibility(layer):
-    """Every unsupported call declines before any device work."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.draft_routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-
-    x = fx.hidden_rows(layer, 6, seed=3)
-    lg = layer.logits(x)
-    assert rm.eligibility(layer.module, layer.bias, x, lg, True) == (True, "")
-    cases = {
-        "unfinalized": (layer.module, layer.bias, x, lg, False),
-        "rows 9": (
-            layer.module,
-            layer.bias,
-            fx.hidden_rows(layer, 9, seed=1),
-            layer.logits(fx.hidden_rows(layer, 9, seed=1)),
-            True,
-        ),
-        "fp16 activation": (layer.module, layer.bias, x.half(), lg, True),
-        "bf16 logits": (layer.module, layer.bias, x, lg.bfloat16(), True),
-        "strided logits": (layer.module, layer.bias, x, torch.cat([lg, lg], 1)[:, ::2], True),
-        "other bias": (layer.module, layer.bias.clone(), x, lg, True),
-        "fp32 bias": (layer.module, layer.bias.float(), x, lg, True),
-        "not a module": (object(), layer.bias, x, lg, True),
-    }
-    for name, args in cases.items():
-        ok, reason = rm.eligibility(*args)
-        print(f"GLM_MOE_DSA_DRAFT_ROUTED_MOE_DECLINE {name}: {reason}", flush=True)
-        assert not ok and reason, name
-    assert rm.routed_output(layer.module, layer.bias, x, lg, False) is None
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py
deleted file mode 100644
index 62c71bb72228..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_jit.py
+++ /dev/null
@@ -1,161 +0,0 @@
-# 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.
-"""Load ``trtllm::glm_moe_dsa_post_attention_hop_routing`` (and the test occupier) for tests.
-
-A TensorRT-LLM build of this branch registers the routing hop in
-``th_common``. On an older image (development overlay) it is compiled from
-this checkout's C++ sources with ``torch.utils.cpp_extension``, unmodified and
-with the same op schema. ``GLM_MOE_DSA_HOP_ROUTING_JIT_DEFINES`` adds compile definitions for
-diagnostic builds (``-DTLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS``); such a build is always
-compiled here and registers ``trtllm::glm_moe_dsa_post_attention_hop_routing_diag``.
-
-``ensure_occupier`` builds ``z2test::occupy(int blocks, int ns)``, a test-only
-kernel of ``blocks`` one-warp CTAs that each spin for ``ns`` nanoseconds; the
-progress tests run it on another stream to hold SMs while the hop waits.
-"""
-
-import hashlib
-import os
-from pathlib import Path
-from typing import Callable
-
-import torch
-
-_REPO = Path(__file__).resolve().parents[9]
-_SOURCES = (
-    _REPO / "cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp",
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu",
-)
-_HEADERS = (
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h",
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h",
-)
-_EXTRA = [f for f in os.environ.get("GLM_MOE_DSA_HOP_ROUTING_JIT_DEFINES", "").split() if f]
-DIAG = "-DTLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS" in _EXTRA
-# The diagnostic build registers ``glm_moe_dsa_post_attention_hop_routing_diag``.
-OP_NAME = (
-    "glm_moe_dsa_post_attention_hop_routing_diag"
-    if DIAG
-    else "glm_moe_dsa_post_attention_hop_routing"
-)
-_NAME = "hop_routing_diag_jit" if DIAG else "hop_routing_jit"
-
-_OCCUPIER_SRC = r"""
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-
-__global__ void z2OccupyKernel(long long ns)
-{
-    unsigned long long start;
-    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(start));
-    unsigned long long now = start;
-    while (now - start < static_cast<unsigned long long>(ns))
-    {
-        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now));
-    }
-}
-
-void occupy(int64_t blocks, int64_t ns)
-{
-    // One CTA per SM: all of an SM's shared memory, so no other CTA fits next to it.
-    int device = 0;
-    cudaGetDevice(&device);
-    int smem = 0;
-    cudaDeviceGetAttribute(&smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
-    cudaFuncSetAttribute(z2OccupyKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-    z2OccupyKernel<<<static_cast<unsigned>(blocks), 32, smem, at::cuda::getCurrentCUDAStream()>>>(ns);
-}
-
-// No tensor arguments: registered as a catch-all kernel.
-TORCH_LIBRARY(z2test, m)
-{
-    m.def("occupy(int blocks, int ns) -> ()", &occupy);
-}
-"""
-
-
-def _digest() -> str:
-    h = hashlib.sha256(" ".join(_EXTRA).encode())
-    for path in (*_SOURCES, *_HEADERS):
-        h.update(path.read_bytes())
-    return h.hexdigest()[:16]
-
-
-def ensure_op(build_root: str, is_builder: bool, barrier: Callable[[], None]) -> str:
-    """Make the routing hop available on every rank; returns where it came from.
-
-    ``is_builder`` must be true on exactly one rank per shared build root;
-    ``barrier`` synchronizes the ranks after that rank has built it.
-    """
-    import tensorrt_llm  # noqa: F401  (loads th_common and the trtllm namespace)
-
-    if hasattr(torch.ops.trtllm, OP_NAME) and not _EXTRA:
-        barrier()
-        return "builtin"
-    build_dir = Path(build_root) / f"hop_routing-{_digest()}"
-    library = build_dir / f"{_NAME}.so"
-    if is_builder and not library.exists():
-        from torch.utils.cpp_extension import load
-
-        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
-        build_dir.mkdir(parents=True, exist_ok=True)
-        load(
-            name=_NAME,
-            sources=[str(s) for s in _SOURCES],
-            extra_include_paths=[str(_REPO / "cpp/include"), str(_REPO / "cpp")],
-            extra_cflags=["-O3", "-std=c++17", *_EXTRA],
-            extra_cuda_cflags=[
-                "-O3",
-                "-std=c++17",
-                "--expt-relaxed-constexpr",
-                "-Xptxas=-v",
-                *_EXTRA,
-            ],
-            build_directory=str(build_dir),
-            is_python_module=False,
-            verbose=True,
-        )
-    barrier()
-    if not hasattr(torch.ops.trtllm, OP_NAME):
-        torch.ops.load_library(str(library))
-    return str(library)
-
-
-def ensure_occupier(build_root: str, is_builder: bool, barrier: Callable[[], None]) -> None:
-    """Make ``torch.ops.z2test.occupy`` available on every rank."""
-    if hasattr(torch.ops, "z2test") and hasattr(torch.ops.z2test, "occupy"):
-        barrier()
-        return
-    digest = hashlib.sha256(_OCCUPIER_SRC.encode()).hexdigest()[:16]
-    build_dir = Path(build_root) / f"hop_routing_occupier-{digest}"
-    library = build_dir / "z2test_occupy.so"
-    if is_builder and not library.exists():
-        from torch.utils.cpp_extension import load
-
-        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
-        build_dir.mkdir(parents=True, exist_ok=True)
-        src = build_dir / "z2test_occupy.cu"
-        src.write_text(_OCCUPIER_SRC)
-        load(
-            name="z2test_occupy",
-            sources=[str(src)],
-            extra_cuda_cflags=["-O3", "-std=c++17"],
-            build_directory=str(build_dir),
-            is_python_module=False,
-            verbose=False,
-        )
-    barrier()
-    if not (hasattr(torch.ops, "z2test") and hasattr(torch.ops.z2test, "occupy")):
-        torch.ops.load_library(str(library))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py
deleted file mode 100644
index 4a85a90d7f6e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop_routing_oracle.py
+++ /dev/null
@@ -1,214 +0,0 @@
-# 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.
-"""Native routing oracle and boundary logits for the hop_routing tests.
-
-``native_routing`` runs the served TRT-LLM-Gen routing runner (DeepSeekV3,
-n_group = topk_group = 1, top_k = 8, routed scaling 2.5, bf16 bias) through
-``trtllm::moe_topk_sort`` and reads the expert of every ``(t, k)`` back
-through the permutation: ``expanded_idx_to_permuted_idx[t, k]`` lies in
-tile ``p // TILE`` whose expert is ``tile_idx_to_expert_idx``. The weights
-are the runner's bf16 ``token_final_scales`` in top-k order.
-
-``boundary_rows`` builds FP32 logit rows whose native selection sits on an
-exact tie or a one-ulp key difference at the 8th/9th boundary (same warp,
-different warps, different clusters of the hop), plus NaN, +-Inf, -0.0,
-huge and subnormal logits; each row is labelled.
-"""
-
-import os
-import sys
-from typing import List, Optional, Tuple
-
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
-import routed_moe_fixtures as fx  # noqa: E402
-
-NUM_EXPERTS = 256
-TOP_K = 8
-TILE = 8
-ROUTED_SCALING = 2.5
-DEEPSEEK_V3 = 2  # RoutingMethodType.DeepSeekV3
-
-
-def native_routing(
-    logits: torch.Tensor, bias_bf16: torch.Tensor
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """``(ids int32 [R, 8], weights bf16 [R, 8])`` of the native runner."""
-    logits = logits.float().contiguous()
-    out = torch.ops.trtllm.moe_topk_sort(
-        logits,
-        bias_bf16.contiguous(),
-        NUM_EXPERTS,
-        TOP_K,
-        1,
-        1,
-        0,
-        NUM_EXPERTS,
-        ROUTED_SCALING,
-        TILE,
-        DEEPSEEK_V3,
-    )
-    tile_expert, tile_limit, exp_to_perm = out[0], out[1], out[2]
-    scales = out[6]
-    perm = exp_to_perm.long()
-    tile = perm // TILE
-    ids = tile_expert[tile]
-    assert bool((perm < tile_limit[tile]).all()), "permuted index outside its tile"
-    return ids.to(torch.int32), scales
-
-
-def reference_routing(logits: torch.Tensor, bias_bf16: torch.Tensor):
-    """The bit-faithful PyTorch model of routed_moe's routing (secondary check)."""
-    return fx.routing_reference(logits.float(), bias_bf16)
-
-
-def weight_bits(w_fp32: torch.Tensor) -> torch.Tensor:
-    return w_fp32.contiguous().view(torch.int32)
-
-
-def native_weight_bits(w_bf16: torch.Tensor) -> torch.Tensor:
-    return w_bf16.float().contiguous().view(torch.int32)
-
-
-def compare_routing(ids, weights, n_ids, n_w) -> Optional[str]:
-    """None when ids are equal and weights are bit-identical (NaN included)."""
-    if not torch.equal(ids, n_ids):
-        bad = (ids != n_ids).any(-1).nonzero().flatten().tolist()
-        return f"ids differ in rows {bad[:8]}: got {ids[bad[0]].tolist()} native {n_ids[bad[0]].tolist()}"
-    a, b = weight_bits(weights), native_weight_bits(n_w)
-    if not torch.equal(a, b):
-        bad = (a != b).any(-1).nonzero().flatten().tolist()
-        return f"weights differ in rows {bad[:8]}: got {weights[bad[0]].tolist()} native {n_w[bad[0]].float().tolist()}"
-    if not torch.equal(weight_bits(weights.to(torch.bfloat16).float()), a):
-        return "weights are not bf16-representable"
-    return None
-
-
-def _key(logit: float, bias_val: torch.Tensor) -> torch.Tensor:
-    lg = torch.tensor([logit], dtype=torch.float32, device=bias_val.device)
-    return fx._key_of(lg, bias_val)[0]
-
-
-def _logit_for_sigmoid(s: float) -> float:
-    return float(torch.logit(torch.tensor(s, dtype=torch.float64)))
-
-
-LOW_LOGIT = -12.0
-
-
-def boundary_rows(bias_bf16: torch.Tensor, seed: int = 0) -> List[Tuple[str, torch.Tensor]]:
-    """Labelled boundary rows (each ``[256]`` FP32 on the bias's device).
-
-    Every expert not placed explicitly gets a logit near ``LOW_LOGIT`` (key
-    ~ its bias). With ``k_low`` the largest such key, the placed experts aim
-    at keys ``k_low + 0.25`` (the 8th/9th pair) and ``k_low + 0.4 + ...``
-    (seven clear winners), so only experts whose bias lets their sigmoid
-    reach those keys are placed.
-    """
-    dev = bias_bf16.device
-    b = bias_bf16.float()
-    g = torch.Generator(device="cpu").manual_seed(seed)
-    rows: List[Tuple[str, torch.Tensor]] = []
-    low_sig = torch.sigmoid(torch.tensor(LOW_LOGIT + 0.5)).item()
-    k_low = float((b + low_sig).max().item())
-    pair_key = k_low + 0.25
-    win_key = k_low + 0.40
-    reach = [
-        e
-        for e in range(NUM_EXPERTS)
-        if 1e-3 < pair_key - float(b[e]) and win_key + 0.08 - float(b[e]) < 1 - 1e-3
-    ]
-
-    def base():
-        return (torch.rand(NUM_EXPERTS, generator=g) * 0.5 + LOW_LOGIT).to(dev)
-
-    def place_winners(row, avoid):
-        winners = [e for e in reach if e not in avoid]
-        perm = torch.randperm(len(winners), generator=g).tolist()
-        for i, j in enumerate(perm[:7]):
-            e = winners[j]
-            row[e] = _logit_for_sigmoid(win_key + 0.01 * i - float(b[e]))
-
-    if len(reach) >= 9:
-        # 8th/9th boundary pairs (e1 < e2): same warp, different warps, different hop clusters.
-        cand = sorted(reach)
-        pairs = [
-            (cand[0], cand[1]),
-            (cand[0], cand[-1]),
-            (cand[len(cand) // 2], cand[len(cand) // 2 + 1]),
-            (cand[1], cand[-2]),
-        ]
-        same_warp = [(a, c) for a, c in zip(cand, cand[1:]) if a // 32 == c // 32]
-        if same_warp:
-            pairs.append(same_warp[0])
-        cross_cluster = [(a, c) for a in cand for c in cand if a < 64 and c >= 192]
-        if cross_cluster:
-            pairs.append(cross_cluster[0])
-        for e1, e2 in pairs:
-            for offset, label in ((0, "tie"), (1, "e2+1ulp"), (-1, "e2-1ulp")):
-                row = base()
-                place_winners(row, (e1, e2))
-                l1 = _logit_for_sigmoid(pair_key - float(b[e1]))
-                k1 = _key(l1, b[e1])
-                l2 = fx.logit_for_key(float(k1.item()), b[e2], offset)
-                if l2 is None:
-                    continue
-                row[e1] = l1
-                row[e2] = l2
-                rows.append((f"{label}({e1},{e2})", row))
-        # Exact ties inside the top 8 (order among equal keys) across warps and clusters.
-        picks = [cand[(i * len(cand)) // 8] for i in range(8)]
-        if len(set(picks)) == 8:
-            row = base()
-            k = _key(_logit_for_sigmoid(pair_key - float(b[picks[0]])), b[picks[0]])
-            ok = True
-            for e in picks:
-                lg = fx.logit_for_key(float(k.item()), b[e])
-                if lg is None:
-                    ok = False
-                    break
-                row[e] = lg
-            if ok:
-                rows.append(("tie-top8", row))
-
-    row = base()
-    row[77] = float("nan")
-    rows.append(("nan1", row))
-    row = base()
-    row[200] = float("nan")
-    row[13] = float("nan")
-    rows.append(("nan2", row))
-    rows.append(("nan-all", torch.full((NUM_EXPERTS,), float("nan"), device=dev)))
-    row = base()
-    row[[1, 9, 100, 250]] = float("inf")
-    row[[2, 3]] = float("-inf")
-    rows.append(("inf", row))
-    rows.append(("neginf-all", torch.full((NUM_EXPERTS,), float("-inf"), device=dev)))
-    rows.append(("posinf-all", torch.full((NUM_EXPERTS,), float("inf"), device=dev)))
-    rows.append(("zero-all", torch.zeros(NUM_EXPERTS, device=dev)))
-    row = torch.zeros(NUM_EXPERTS, device=dev)
-    row[::2] = -0.0
-    rows.append(("negzero-mix", row))
-    row = base()
-    row[[4, 44, 144]] = 3.0e38
-    row[[5, 55]] = -3.0e38
-    row[[6, 66]] = 1.0e-42
-    row[[7, 77]] = -1.0e-42
-    rows.append(("huge-subnormal", row))
-    for scale in (0.5, 2.0, 8.0):
-        rows.append((f"randn*{scale}", (torch.randn(NUM_EXPERTS, generator=g) * scale).to(dev)))
-    return rows
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py
deleted file mode 100644
index fa011031dd45..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_adapter.py
+++ /dev/null
@@ -1,245 +0,0 @@
-# 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.
-"""Control flow of the hop_routing call sites in the target's decoder layer and MoE.
-
-With fake collaborators (no device work), on a decode step, checks that
-
-* the plan runs hop_routing with its closure, and not without it;
-* ``GlmDecoderLayer.forward_MoE`` runs the hop_routing hop instead of the into-hop entry's
-  exactly when the tail published and the hop_routing transport can run, hands the hop's
-  routing and logits to the MoE, and otherwise runs the into-hop entry's hop with no routing;
-* ``GlmMoE.compute_routed_output`` tries routed_moe's hop_routing fork first only when routing
-  is present (routed_moe_into_hop entry with the hop, the plain entry without), and falls back
-  to routed_moe on the same logits when the fork returns None.
-"""
-
-import os
-import sys
-import types
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
-    dispatch as _dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-    GlmDecoderLayer,
-    GlmMoE,
-)
-
-sys.path.insert(
-    0,
-    os.path.join(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "..", "shared", "kernels"
-    ),
-)
-
-import plan_testlib  # noqa: E402
-
-FEATURES = (
-    "attention_output_tail,post_moe_hop,post_attention_hop,attention_output_into_hop,"
-    "routed_moe,routed_moe_into_hop,hop_routing"
-)
-
-
-@pytest.fixture(autouse=True)
-def decode_step():
-    """The owned call sites run on decode steps."""
-    with _dispatch.step(_dispatch.Phase.DECODE):
-        yield
-
-
-@pytest.fixture
-def selector(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-    plan_testlib.select(dispatch, FEATURES, monkeypatch)
-    return dispatch
-
-
-def test_selector_accepts_hop_routing(selector):
-    dispatch = selector
-    assert dispatch.is_enabled("hop_routing")
-    for base in dispatch.REQUIRED_FEATURES["hop_routing"]:
-        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
-        assert not plan.enabled("hop_routing"), base
-
-
-class _Calls:
-    def __init__(self):
-        self.log = []
-
-    def fn(self, name, result):
-        def call(*args, **kwargs):
-            self.log.append((name, args, kwargs))
-            return result
-
-        return call
-
-
-def _decoder(hop_routing_can_run, calls):
-    dec = GlmDecoderLayer.__new__(GlmDecoderLayer)
-    torch.nn.Module.__init__(dec)
-    rows = 6
-    normed = torch.zeros(rows, 4)
-    res = torch.ones(rows, 4)
-    logits = torch.full((rows, 256), 2.0)
-    routing = torch.zeros(2, rows, 8, dtype=torch.int32)
-    bias = torch.zeros(256, dtype=torch.bfloat16)
-    weight = torch.zeros(256, 4)
-
-    class _HOP_ROUTING:
-        def can_run(self, *args, **kwargs):
-            calls.log.append(("hop_routing.can_run", args, kwargs))
-            return hop_routing_can_run
-
-        def __call__(self, *args, **kwargs):
-            calls.log.append(("hop_routing", args, kwargs))
-            return normed, res, logits, routing
-
-    class _INTO_HOP:
-        def __call__(self, *args, **kwargs):
-            calls.log.append(("into_hop", args, kwargs))
-            return normed, res, logits
-
-    class _MoE:
-        gate = types.SimpleNamespace(e_score_correction_bias=bias)
-        experts = types.SimpleNamespace(has_nvfp4=True)
-
-        def __call__(self, hidden_states, hidden_states_fp4, **kwargs):
-            calls.log.append(("mlp", (hidden_states,), kwargs))
-            return hidden_states
-
-    dec.__dict__.update(
-        _hop_routing=_HOP_ROUTING(),
-        _into_hop=_INTO_HOP(),
-        mlp=_MoE(),
-        _post_attn_router_weight=lambda: weight,
-        post_attention_layernorm=types.SimpleNamespace(
-            weight=torch.ones(4), variance_epsilon=1e-5, nvfp4_scale=None
-        ),
-        fusion_config=types.SimpleNamespace(PRE_MOE_FUSION=True, POST_MOE_FUSION=False),
-        mapping=types.SimpleNamespace(tp_size=8, is_multi_node=lambda: True),
-        post_moe_hop=None,
-        next_layer_layernorm=None,
-        layer_idx=3,
-        _apply_next_layer_layernorm=lambda h, r: (h, r),
-    )
-    return dec, (normed, res, logits, routing, bias, weight)
-
-
-@pytest.mark.parametrize("hop_routing_can_run", [True, False])
-def test_forward_moe_selects_the_hop(hop_routing_can_run):
-    calls = _Calls()
-    dec, (normed, res, logits, routing, bias, weight) = _decoder(hop_routing_can_run, calls)
-    meta = types.SimpleNamespace(all_rank_num_tokens=[6])
-    GlmDecoderLayer.forward_MoE(dec, torch.zeros(6, 4), meta, res, into_hop_published=True)
-    names = [c[0] for c in calls.log]
-    assert names[0] == "hop_routing.can_run"
-    can = calls.log[0]
-    assert can[1][2] is weight and can[1][3] is bias
-    mlp = [c for c in calls.log if c[0] == "mlp"][0]
-    assert mlp[2]["router_logits"] is logits
-    if hop_routing_can_run:
-        assert "hop_routing" in names and "into_hop" not in names
-        assert mlp[2]["hop_routing_out"] is routing
-    else:
-        assert "into_hop" in names and "hop_routing" not in names
-        assert mlp[2]["hop_routing_out"] is None
-
-
-def _moe(selector, calls, hop_routing_into, hop_routing_plain, hop):
-    dispatch = selector
-    for feature, name, result in (
-        ("hop_routing", "routed_output_into_hop", hop_routing_into),
-        ("hop_routing", "routed_output", hop_routing_plain),
-        ("routed_moe", "routed_output_into_hop", ["routed_moe_into"]),
-        ("routed_moe", "routed_output", ["routed_moe_plain"]),
-    ):
-        dispatch.register(feature, name, calls.fn(f"{feature}.{name}", result))
-    moe = GlmMoE.__new__(GlmMoE)
-    torch.nn.Module.__init__(moe)
-    moe.__dict__.update(
-        use_dp=False,
-        mapping=types.SimpleNamespace(tp_size=8),
-        router_gemm=None,
-        gate=types.SimpleNamespace(e_score_correction_bias=torch.zeros(256)),
-        experts=object(),
-        _routed_into_hop=hop,
-    )
-    return moe
-
-
-@pytest.fixture
-def restore_registry():
-    import importlib
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.modeling import (  # noqa: E501
-        TARGET_PATHS,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    # Import both packages (which registers their real entries) before the snapshot: a package registers only on
-    # its first import, so entries dropped by a restore would never come back.
-    for feature in ("hop_routing", "routed_moe"):
-        importlib.import_module(TARGET_PATHS.kernel_modules[feature])
-        dispatch.get_impl(feature, "routed_output")
-    saved = {k: dict(v) for k, v in dispatch._REGISTRY.items()}
-    yield
-    dispatch._REGISTRY.clear()
-    dispatch._REGISTRY.update(saved)
-
-
-@pytest.mark.parametrize(
-    "hop, routing_given, fork_result, want",
-    [
-        ("hop", True, ["fork_into"], ["hop_routing.routed_output_into_hop"]),
-        (
-            "hop",
-            True,
-            None,
-            ["hop_routing.routed_output_into_hop", "routed_moe.routed_output_into_hop"],
-        ),
-        ("hop", False, ["fork_into"], ["routed_moe.routed_output_into_hop"]),
-        (None, True, ["fork_plain"], ["hop_routing.routed_output"]),
-        (None, True, None, ["hop_routing.routed_output", "routed_moe.routed_output"]),
-        (None, False, ["fork_plain"], ["routed_moe.routed_output"]),
-    ],
-)
-def test_compute_routed_output_order(
-    selector, restore_registry, hop, routing_given, fork_result, want
-):
-    calls = _Calls()
-    moe = _moe(selector, calls, fork_result, fork_result, object() if hop else None)
-    x = torch.zeros(6, 4)
-    logits = torch.zeros(6, 256)
-    routing = torch.zeros(2, 6, 8, dtype=torch.int32) if routing_given else None
-    out = moe.compute_routed_output(
-        x, None, None, False, router_logits=logits, hop_routing_out=routing
-    )
-    assert [c[0] for c in calls.log] == want
-    for name, args, _ in calls.log:
-        assert args[3] is logits  # every entry receives the hop's logits
-        if name.startswith("hop_routing."):
-            assert args[4] is routing
-    results = {
-        "hop_routing.routed_output_into_hop": fork_result,
-        "hop_routing.routed_output": fork_result,
-        "routed_moe.routed_output_into_hop": ["routed_moe_into"],
-        "routed_moe.routed_output": ["routed_moe_plain"],
-    }
-    assert out == results[want[-1]]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py
deleted file mode 100644
index 9b0329202532..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_hop_single_gpu.py
+++ /dev/null
@@ -1,234 +0,0 @@
-# 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.
-"""The hop_routing hop on one GPU, over a simulated transport (hop_routing).
-
-Like the into-hop entry, the hop_routing hop reads every rank's partial from the local
-(unicast) view of the post_attention_hop workspace only, so a plain device buffer stands in
-for it: the test writes all eight ranks' sanitized partials into the current
-Lamport buffer, sets the protocol state, and runs the hop as rank 0 of 8.
-This needs no MNNVL and runs under every compute-sanitizer tool. Checks:
-
-* normed, residual and router logits are bit-identical to the into-hop op's on the
-  same transport (R = 1, 3, 6, 8, every ``router_load``, both poll modes);
-* ids and weights equal the native TRT-LLM-Gen routing runner on those
-  logits (ids equal, weight bits equal);
-* the gather protocol: the previous parity's buffer is reset to the sentinel
-  (it holds junk before the call), this call's buffer holds a score for
-  every live row and the sentinel elsewhere, the parity flips, the post_attention_hop state
-  advances as the into-hop entry's;
-* boundary rows through the ``logits_in`` instance;
-* CUDA-graph capture and replay with fresh partials and alternating parity.
-
-Partials, weights and bias are synthetic (seeded).
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(
-    0, os.path.join(os.path.dirname(_HERE), "..", "..", "shared", "kernels", "post_attention_hop")
-)
-import hop_routing_jit  # noqa: E402
-import hop_routing_oracle as orc  # noqa: E402
-import into_hop_jit  # noqa: E402
-
-H = 6144
-RANKS = 8
-MAX_ROWS = 8
-NUM_BUFFERS = 3
-BUFFER_BYTES = MAX_ROWS * RANKS * H * 2
-WORKSPACE_BYTES = BUFFER_BYTES * NUM_BUFFERS
-SENTINEL = -(2**31)  # 0x80000000 as int32
-GATHER_WORDS = 2 * MAX_ROWS * 256
-EPS = 1e-5
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="the hop_routing hop targets SM100/SM103",
-    ),
-    pytest.mark.threadleak(enabled=False),
-]
-
-_OPS = {}
-
-
-def _ops():
-    if not _OPS:
-        root = os.environ.get("GLM_MOE_DSA_HOP_ROUTING_JIT_ROOT", "/tmp/hop_routing_jit")
-        into_hop_jit.ensure_op(root, True, lambda: None)
-        hop_routing_jit.ensure_op(root, True, lambda: None)
-        _OPS["into_hop"] = torch.ops.trtllm.glm_moe_dsa_post_attention_hop_from_partials
-        _OPS["hop_routing"] = getattr(torch.ops.trtllm, hop_routing_jit.OP_NAME)
-    return _OPS
-
-
-def _sanitized(x):
-    bits = x.contiguous().view(torch.int16).clone()
-    bits[bits == -32768] = 0  # bf16 -0.0 -> +0.0
-    return bits.view(torch.bfloat16)
-
-
-class Transport:
-    """A device buffer laid out like the post_attention_hop workspace, its state and the hop_routing gather state."""
-
-    def __init__(self):
-        self.uc = torch.full((WORKSPACE_BYTES // 4,), SENTINEL, dtype=torch.int32, device="cuda")
-        self.flags = torch.zeros(4, dtype=torch.int32, device="cuda")
-        self.gather = torch.full((GATHER_WORDS,), SENTINEL, dtype=torch.int32, device="cuda")
-        self.zstate = torch.zeros(4, dtype=torch.int32, device="cuda")
-
-    def buffers(self):
-        return self.uc.view(torch.bfloat16).view(NUM_BUFFERS, MAX_ROWS, RANKS, H)
-
-    def publish(self, partials, cur, prev_rows):
-        prev = (cur + NUM_BUFFERS - 1) % NUM_BUFFERS
-        self.uc.fill_(SENTINEL)
-        bufs = self.buffers()
-        bufs[prev, :prev_rows].fill_(1.5)
-        for s, part in enumerate(partials):
-            bufs[cur, : part.shape[0], s] = _sanitized(part)
-        self.flags.copy_(torch.tensor([cur, prev_rows * RANKS * H * 2, 0, 0], dtype=torch.int32))
-
-
-def _inputs(rows, seed):
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    partials = [
-        (torch.randn(rows, H, generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
-        for _ in range(RANKS)
-    ]
-    residual = torch.randn(rows, H, generator=gen, device="cuda").to(torch.bfloat16)
-    gamma = (1 + 0.1 * torch.randn(H, generator=gen, device="cuda")).to(torch.bfloat16)
-    weight = (
-        (torch.randn(256, H, generator=gen, device="cuda") * 0.02).to(torch.bfloat16).contiguous()
-    )
-    bias = (torch.randn(256, generator=gen, device="cuda") * 0.05).to(torch.bfloat16).contiguous()
-    return partials, residual, gamma, weight, bias
-
-
-def _into_hop(tr, residual, gamma, weight, early_poll, router_load):
-    return _ops()["into_hop"](
-        residual, gamma, EPS, weight, tr.uc, tr.flags, 0, RANKS, early_poll, router_load
-    )
-
-
-def _hop_routing(
-    tr, residual, gamma, weight, bias, early_poll=False, router_load=1, logits_in=None
-):
-    return _ops()["hop_routing"](
-        residual,
-        gamma,
-        EPS,
-        weight,
-        bias,
-        tr.uc,
-        tr.flags,
-        tr.gather,
-        tr.zstate,
-        0,
-        RANKS,
-        early_poll,
-        router_load,
-        2.5,
-        logits_in,
-        None,
-    )
-
-
-def _check_routing(routing, logits, bias, rows):
-    ids, w = routing[0][:rows], routing[1][:rows].view(torch.float32)
-    n_ids, n_w = orc.native_routing(logits[:rows], bias)
-    err = orc.compare_routing(ids, w, n_ids, n_w)
-    assert err is None, err
-
-
-@pytest.mark.parametrize("rows", [1, 3, 6, 8])
-def test_identity_routing_and_protocol(rows):
-    partials, residual, gamma, weight, bias = _inputs(rows, 100 + rows)
-    tr = Transport()
-    for early_poll in (False, True):
-        for router_load in (0, 1, 2):
-            cur, prev_rows = (early_poll + 2 * router_load) % 3, 8 - rows + 1
-            tr.publish(partials, cur, prev_rows)
-            ref = _into_hop(tr, residual, gamma, weight, early_poll, router_load)
-            torch.cuda.synchronize()
-            tr.publish(partials, cur, prev_rows)
-            parity = int(tr.zstate[0].item()) & 1
-            gview = tr.gather.view(2, MAX_ROWS, 256)
-            gview[parity ^ 1].fill_(12345)  # junk in the previous call's buffer
-            got = _hop_routing(tr, residual, gamma, weight, bias, early_poll, router_load)
-            torch.cuda.synchronize()
-            for a, b in zip(got[:3], ref):
-                assert torch.equal(a, b)
-            _check_routing(got[3], got[2], bias, rows)
-            assert bool((gview[parity ^ 1] == SENTINEL).all()), "previous gather buffer not reset"
-            assert bool((gview[parity, :rows] != SENTINEL).all()), "missing scores"
-            assert bool((gview[parity, rows:] == SENTINEL).all()), "scores outside the call's rows"
-            assert int(tr.zstate[0].item()) == parity ^ 1
-            want = [(cur + 1) % NUM_BUFFERS, rows * RANKS * H * 2, 0, 0]
-            assert tr.flags.tolist() == want, (tr.flags.tolist(), want)
-            # This call's buffer must be all-sentinel again before its parity comes back.
-            gview[parity].fill_(SENTINEL)
-
-
-def test_boundary_rows_logits_in():
-    partials, residual, gamma, weight, bias = _inputs(8, 7)
-    rows_all = orc.boundary_rows(bias, seed=3)
-    mat = torch.stack([r for _, r in rows_all]).float()
-    tr = Transport()
-    cur = 0
-    for start in range(0, mat.shape[0], 8):
-        chunk = mat[start : start + 8].contiguous()
-        n = chunk.shape[0]
-        tr.publish([p[:n] for p in partials], cur, n)
-        got = _hop_routing(tr, residual[:n].contiguous(), gamma, weight, bias, logits_in=chunk)
-        torch.cuda.synchronize()
-        tr.gather.view(2, MAX_ROWS, 256)[int(tr.zstate[0].item()) ^ 1].fill_(SENTINEL)
-        _check_routing(got[3], chunk, bias, n)
-        cur = (cur + 1) % NUM_BUFFERS
-
-
-def test_graph_replay_alternates_parity():
-    rows = 6
-    partials, residual, gamma, weight, bias = _inputs(rows, 11)
-    tr = Transport()
-    tr.publish(partials, 0, rows)
-    _hop_routing(tr, residual, gamma, weight, bias)
-    torch.cuda.synchronize()
-    tr.publish(partials, 1, rows)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = _hop_routing(tr, residual, gamma, weight, bias)
-    for rep in range(4):
-        fresh, res_new, _, _, _ = _inputs(rows, 300 + rep)
-        residual.copy_(res_new)
-        parity = int(tr.zstate[0].item()) & 1
-        tr.publish(fresh, rep % NUM_BUFFERS, rows)
-        graph.replay()
-        torch.cuda.synchronize()
-        assert int(tr.zstate[0].item()) == parity ^ 1
-        got = [t.clone() for t in captured]
-        tr.publish(fresh, rep % NUM_BUFFERS, rows)
-        ref = _into_hop(tr, residual, gamma, weight, False, 1)
-        torch.cuda.synchronize()
-        for a, b in zip(got[:3], ref):
-            assert torch.equal(a, b)
-        _check_routing(got[3], got[2], bias, rows)
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py
deleted file mode 100644
index a5df11349cb7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/test_hop_routing_routed_moe_fork.py
+++ /dev/null
@@ -1,200 +0,0 @@
-# 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.
-"""The hop_routing routed_moe fork against routed_moe on one GPU (real GLM-5.2 NVFP4 weights).
-
-routed_moe computes the routing from the logits; the fork reads the same routing
-(the native runner's ids and weights for those logits, packed as the hop_routing hop
-writes them). Given that, every output must be routed_moe's bit for bit:
-
-* the routing tables the kernel builds (``DEBUG_ROUTING`` dump: top-k ids
-  and weights, slot experts, slot masks, slot count);
-* the raw tuple (FC2 rows and expert scales) and the finalized output;
-* the routed_moe_into_hop FP32 accumulator within the FP32 reorder bound (routed_moe's own output
-  class; the atomic adds are run-to-run nondeterministic), and identical to
-  routed_moe's from the same rows;
-* padded rows leave the live rows unchanged; CUDA-graph replay with fresh
-  routing; host-only declines (wrong routing shape or dtype) before any
-  device work.
-
-Inputs: real captured logits and post-norm hidden states of layer 10
-(``GLM_MOE_DSA_REAL_MOE_INPUTS_DIR``), plus random activations with the real router GEMM.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
-import hop_routing_oracle as orc  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-DEBUG_ROUTING = 1
-REORDER_REL = 2.0**-20
-LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
-
-
-@pytest.fixture(scope="module")
-def layer():
-    torch.manual_seed(0)
-    return fx.build_native_layer(
-        layer=LAYER, tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
-    )
-
-
-@pytest.fixture(scope="module")
-def runners(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing.routed_moe import (  # noqa: E501
-        HopRoutedMoe,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    return RoutedMoe.from_native(layer.module, layer.bias), HopRoutedMoe(layer.module, layer.bias)
-
-
-def _pack(logits, bias):
-    ids, w = orc.native_routing(logits, bias)
-    rows = ids.shape[0]
-    routing = torch.empty(2, rows, 8, dtype=torch.int32, device="cuda")
-    routing[0].copy_(ids)
-    routing[1].copy_(w.float().view(torch.int32))
-    return routing
-
-
-def _real(layer, rows, count, offset=0):
-    win = fx.real_windows(fx.load_real_capture(layer.layer), rows, stride=rows * 7)
-    out = []
-    for x, lg in win[offset : offset + count]:
-        out.append((x.cuda().contiguous(), lg.float().cuda().contiguous()))
-    return out
-
-
-def _random(layer, rows, seed):
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-    return x, layer.router_logits(x).float().contiguous()
-
-
-def _cases(layer, rows):
-    return _real(layer, rows, 3, offset=rows) + [_random(layer, rows, 40 + rows)]
-
-
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-def test_routing_tables_raw_and_finalized_equal_routed_moe(layer, runners, rows):
-    routed_runner, fork = runners
-    for i, (x, logits) in enumerate(_cases(layer, rows)):
-        routing = _pack(logits, layer.bias)
-        dbg_a = torch.zeros(2 * 1024 * 1024, dtype=torch.int32, device="cuda")
-        dbg_b = torch.zeros_like(dbg_a)
-        routed_runner(x, logits, debug=DEBUG_ROUTING, debug_buffer=dbg_a)
-        fork(x, routing, debug=DEBUG_ROUTING, debug_buffer=dbg_b)
-        raw_a = routed_runner(x, logits, do_finalize=False)
-        raw_b = fork(x, routing, do_finalize=False)
-        fin_a = routed_runner(x, logits)
-        fin_b = fork(x, routing)
-        torch.cuda.synchronize()
-        assert torch.equal(dbg_a[:2048], dbg_b[:2048]), f"case {i}: routing tables differ"
-        for a, b in zip(raw_a, raw_b):
-            assert torch.equal(a, b), f"case {i}: raw tuple differs"
-        assert torch.equal(fin_a, fin_b), f"case {i}: finalized output differs"
-
-
-def test_fp32_accumulator_matches_routed_moe_class(layer, runners):
-    routed_runner, fork = runners
-    rows = 6
-    if fork.schedule(rows, False, True) != "xc":
-        pytest.skip("FP32 accumulator output needs the expert-cluster schedule")
-    for x, logits in _real(layer, rows, 2, offset=3):
-        routing = _pack(logits, layer.bias)
-        fc2, scale, idx = routed_runner(x, logits, do_finalize=False)
-        terms = scale.double()[:, :, None] * fc2[idx.long()].double()
-        exact = terms.sum(1)
-        bound = REORDER_REL * terms.abs().sum(1)
-        for _ in range(4):
-            acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
-            fork(x, routing, red_out=acc)
-            torch.cuda.synchronize()
-            assert bool(((acc.double() - exact).abs() <= bound).all())
-
-
-@pytest.mark.parametrize("live", [1, 5])
-def test_padded_rows_do_not_change_live_rows(layer, runners, live):
-    _, fork = runners
-    ((x, logits),) = _real(layer, 8, 1, offset=1)
-    full = fork(x, _pack(logits, layer.bias), do_finalize=False)
-    part = fork(
-        x[:live].contiguous(), _pack(logits[:live].contiguous(), layer.bias), do_finalize=False
-    )
-    torch.cuda.synchronize()
-    assert torch.equal(full[0][: live * 8], part[0])
-    assert torch.equal(full[1][:live], part[1])
-
-
-def test_cuda_graph_replay(layer, runners):
-    _, fork = runners
-    rows = 6
-    cases = _real(layer, rows, 4, offset=2)
-    x = cases[0][0].clone()
-    routing = _pack(cases[0][1], layer.bias)
-    fork(x, routing, do_finalize=False)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = fork(x, routing, do_finalize=False)
-    for xn, ln in cases[1:]:
-        x.copy_(xn)
-        routing.copy_(_pack(ln, layer.bias))
-        graph.replay()
-        torch.cuda.synchronize()
-        got = [t.clone() for t in captured]
-        want = fork(x, routing, do_finalize=False)
-        torch.cuda.synchronize()
-        for a, b in zip(got, want):
-            assert torch.equal(a, b)
-    del graph
-
-
-def test_declines_before_device_work(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing import (  # noqa: E501
-        routed_moe as hop_routed_mod,
-    )
-
-    x, logits = _random(layer, 6, 3)
-    good = _pack(logits, layer.bias)
-    bad = [
-        None,
-        good[:, :5].contiguous(),
-        good.to(torch.int64),
-        good.transpose(1, 2).contiguous().transpose(1, 2),
-    ]
-    torch.cuda.synchronize()
-    for routing in bad:
-        before = torch.cuda.memory_allocated()
-        assert (
-            hop_routed_mod.routed_output(layer.module, layer.bias, x, logits, routing, False)
-            is None
-        )
-        assert torch.cuda.memory_allocated() == before
-    assert hop_routed_mod.routing_reason(good, x) == ""
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py
deleted file mode 100644
index 438ff7dba012..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/post_attention_hop/test_attn_hop_integration_contract.py
+++ /dev/null
@@ -1,161 +0,0 @@
-# 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.
-"""Call-site invariants of the post_attention_hop in the target's layers.
-
-The post_attention_hop keeps one workspace and one device-side buffer state per TP group
-for every layer, CUDA graph and eager call, as the native MNNVL all-reduce
-does. That is correct only while every rank issues the same sequence of hop
-calls on one stream, deciding each call from rank-uniform facts. These checks
-pin the preconditions on the target's layer module:
-
-* the hop and the into-hop transport (attention_output_into_hop) are created at one place each, at
-  layer adoption, through the dispatch entry points;
-* every hop call (post_attention_hop, into-hop entry) is issued directly on the layer's stream:
-  never inside a ``with`` block (e.g. ``torch.cuda.stream``) and never inside
-  a nested function or lambda (e.g. a closure handed to
-  ``maybe_execute_in_parallel``, which runs one branch on an auxiliary
-  stream);
-* the per-call decision reads only shapes, dtypes, contiguity and the
-  rank-agreed SM count: no device-to-host transfer or tensor value;
-* dense layers keep the native post-attention path: ``forward_mlp`` never
-  reaches the hop (scope of the served post_attention_hop: MoE and MTP layers).
-"""
-
-import ast
-import inspect
-import textwrap
-
-import pytest
-
-modeling = pytest.importorskip(
-    "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers"
-)
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop import (  # noqa: E402, E501
-    hop,
-)
-
-_HOP_CALLS = ("_post_attn_hop", "_into_hop", "post_attn_hop")
-
-
-@pytest.fixture(scope="module")
-def tree():
-    return ast.parse(inspect.getsource(modeling))
-
-
-def _self_attr_calls(node, attrs):
-    return [
-        n
-        for n in ast.walk(node)
-        if isinstance(n, ast.Call)
-        and isinstance(n.func, ast.Attribute)
-        and n.func.attr in attrs
-        and isinstance(n.func.value, ast.Name)
-        and n.func.value.id == "self"
-    ]
-
-
-def _parents(tree):
-    parent = {}
-    for node in ast.walk(tree):
-        for child in ast.iter_child_nodes(node):
-            parent[child] = node
-    return parent
-
-
-def _methods(tree, cls_name):
-    for node in ast.walk(tree):
-        if isinstance(node, ast.ClassDef) and node.name == cls_name:
-            return {n.name: n for n in node.body if isinstance(n, ast.FunctionDef)}
-    raise AssertionError(f"{cls_name} not found")
-
-
-def test_hop_created_once_through_dispatch(tree):
-    calls = [
-        n
-        for n in ast.walk(tree)
-        if isinstance(n, ast.Call)
-        and isinstance(n.func, ast.Call)
-        and isinstance(n.func.func, ast.Attribute)
-        and n.func.func.attr == "get_impl"
-        and len(n.func.args) == 2
-        and all(isinstance(a, ast.Constant) for a in n.func.args)
-    ]
-    entries = [(c.func.args[0].value, c.func.args[1].value) for c in calls]
-    assert entries.count(("post_attention_hop", "make_post_attn_hop")) == 1, entries
-    assert entries.count(("post_attention_hop", "make_into_hop_transport")) <= 1, entries
-    methods = _methods(tree, "_AdoptedLayerMixin")
-    init = methods["_adopt_pre_moe"]
-    in_init = [
-        c for c in ast.walk(init) if c in calls and c.func.args[0].value == "post_attention_hop"
-    ]
-    assert len(in_init) == len(
-        [e for e in entries if e[0] == "post_attention_hop" and e[1].startswith("make_")]
-    )
-
-
-def test_hop_calls_are_issued_on_the_layer_stream(tree):
-    parent = _parents(tree)
-    sites = []
-    for call in _self_attr_calls(tree, _HOP_CALLS):
-        node = call
-        while node in parent:
-            node = parent[node]
-            if isinstance(node, (ast.With, ast.AsyncWith)):
-                raise AssertionError(
-                    f"hop call at line {call.lineno} inside a with block: "
-                    f"{ast.unparse(node.items[0])}"
-                )
-            if isinstance(node, ast.Lambda):
-                raise AssertionError(f"hop call at line {call.lineno} inside a lambda")
-            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                outer = parent.get(node)
-                assert isinstance(outer, ast.ClassDef), (
-                    f"hop call at line {call.lineno} inside nested function {node.name}"
-                )
-                sites.append((outer.name, node.name))
-                break
-    # The served call sites: target MoE layers (forward_MoE) and the MTP
-    # layer; _post_attn_hop itself wraps the hop object.
-    assert ("GlmDecoderLayer", "forward_MoE") in sites, sites
-    assert ("GlmMtpLayer", "forward") in sites, sites
-
-
-def test_dense_layers_never_reach_the_hop(tree):
-    methods = _methods(tree, "GlmDecoderLayer")
-    forward_mlp = methods["forward_mlp"]
-    names = {n.attr for n in ast.walk(forward_mlp) if isinstance(n, ast.Attribute)}
-    assert not names & {"post_attn_hop", "_post_attn_hop", "_can_use_post_attn_hop", "_into_hop"}, (
-        names
-    )
-
-
-def test_call_decision_reads_no_tensor_values():
-    for fn in (hop.PostAttnHop.call_ineligibility, hop.PostAttnHop.can_run, hop.required_sms):
-        node = ast.parse(textwrap.dedent(inspect.getsource(fn)))
-        attrs = {n.attr for n in ast.walk(node) if isinstance(n, ast.Attribute)}
-        assert not attrs & {
-            "item",
-            "tolist",
-            "cpu",
-            "numpy",
-            "any",
-            "all",
-            "sum",
-            "max",
-            "min",
-            "nonzero",
-            "synchronize",
-        }, (fn.__name__, attrs)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py
deleted file mode 100644
index a7994e880c34..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_into_hop_checks.py
+++ /dev/null
@@ -1,492 +0,0 @@
-# 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.
-"""Collective checks of routed_moe_into_hop: routed_moe's finalize kernel -> post_moe_hop routed-input hop.
-
-Used by ``test_routed_into_hop_mpi.py``. Every function here is collective over the two-host TP8
-group (call it in the same order on every rank).
-
-A per-rank *producer* decides what a rank hands to the post_moe_hop for one call:
-
-* ``routed_moe_into_hop``: routed_moe writes into the hop's routed-input buffer ->
-  ``hop.hop_routed`` (the cluster kernel's routed-input variant of the hop's mode,
-  ``TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION``, default "fp32":
-  "row" = finalized by ``RoutedMoeFinalizeKernel``, "identity" = identity-layout raw
-  rows, "fp32" = FP32 accumulator added by the routed kernel);
-* ``routed_moe``: routed_moe raw tuple -> ``hop(...)`` (default variant);
-* ``native``: the native TRT-LLM-Gen raw tuple -> ``hop(...)``.
-
-The hop is constructed, and its routed input enabled, inside the served
-model loader's ``MetaInitMode`` (``served_construction``), as
-the target's decoder-layer adoption does in the server.
-
-Ranks may pick different producers for the same call (a rank whose routed_moe
-declines routed_moe_into_hop keeps the raw tuple); the collective order is one hop per call
-on every rank either way. The oracle of a call is the same per-rank choice
-with ``routed_moe_into_hop`` replaced by ``routed_moe``, i.e. every rank on the default (raw-tuple)
-variant. In modes "row" and "identity" routed_moe_into_hop's local partial must be
-bit-identical to the raw variant's, so the reduced outputs must be
-bit-identical too. In mode "fp32" the top-k sum order differs from the native
-finalize: the published partial must lie within the FP32 reorder bound of the
-raw variant's partial (``REORDER_REL * sum_k |w y| + ulp_bf16``), and the
-outputs must equal the native-arithmetic emulation of the published partials.
-"""
-
-import os
-import random
-import sys
-
-import torch
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(HERE), "routed_moe"))
-sys.path.insert(
-    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
-)
-import post_moe_hop_mpi_utils as U  # noqa: E402
-import routed_moe_chain_bench as CB  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-FEATURES = {
-    "routed_moe_into_hop": "post_moe_hop,routed_moe,routed_moe_into_hop",
-    "routed_moe": "routed_moe",
-    "native": "",
-}
-NORM_MAX_ULP = 2
-# eight exact FP32 products summed in another order: |delta| <= 7 * 2^-24 * sum|w y|
-REORDER_REL = 2.0**-20
-
-
-_WINDOWS = {}
-
-
-def windows(layer: int, rows: int, count: int, offset: int = 0):
-    """``count`` real hidden-state windows of the captured fixture (spread over it)."""
-    key = (layer, rows)
-    if key not in _WINDOWS:
-        _WINDOWS[key] = [x for x, _ in fx.real_windows(fx.load_real_capture(layer), rows)]
-    win = _WINDOWS[key]
-    step = max(1, len(win) // count)
-    return [win[(offset + i * step) % len(win)].cuda() for i in range(count)]
-
-
-def served_construction():
-    """The context in which the served model loader constructs the model.
-
-    ``PyTorchModelLoader`` builds the model inside ``MetaInitMode``: ``empty``
-    and ``empty_like`` produce meta tensors, and afterwards only module
-    parameters and buffers are materialized. The target's decoder-layer adoption
-    builds the post_moe_hop and enables its routed input inside that context, so
-    the harness does the same.
-    """
-    from tensorrt_llm._torch.models.modeling_utils import MetaInitMode
-
-    return MetaInitMode()
-
-
-_ROUTED_BUFFERS = {
-    "row": ("_routed_buffer",),
-    "identity": ("_routed_rows", "_routed_scales"),
-    "fp32": ("_routed_acc",),
-}
-
-
-def routed_buffers(hop):
-    """``{attribute: tensor or None}`` of the routed-input buffers of the hop's mode."""
-    return {name: getattr(hop, name, None) for name in _ROUTED_BUFFERS[hop.routed_mode]}
-
-
-def buffer_report(hop):
-    """Where each routed-input buffer of the hop's mode lives (for logs)."""
-    return {
-        name: None if t is None else f"{t.device} meta={t.is_meta} ptr={t.data_ptr():#x}"
-        for name, t in routed_buffers(hop).items()
-    }
-
-
-# MoE layers of GLM-5.2 that build a post_moe_hop (layers 3..77)
-MOE_LAYERS = 75
-
-
-def buffer_bytes(hop) -> int:
-    """Device bytes of the routed-input buffers of the hop's mode."""
-    return sum(t.numel() * t.element_size() for t in routed_buffers(hop).values() if t is not None)
-
-
-def buffer_problems(hop):
-    """Why the routed-input buffers are not usable device memory ([] when they are)."""
-    device = hop.workspace.state.device
-    problems = []
-    for name, t in routed_buffers(hop).items():
-        if t is None:
-            problems.append(f"{name} missing")
-        elif t.is_meta:
-            problems.append(f"{name} is a meta tensor (data_ptr {t.data_ptr():#x})")
-        elif t.device != device:
-            problems.append(f"{name} on {t.device}, the hop's workspace on {device}")
-        elif t.data_ptr() == 0 or t.data_ptr() % 16:
-            problems.append(f"{name} data_ptr {t.data_ptr():#x}")
-    return problems
-
-
-def build_hop(env, producer, mode=None):
-    """The post_moe_hop with routed input, constructed as the served model does.
-
-    Inside ``served_construction()``, with the layer's MoE module as the
-    declared raw-tuple producer, followed by the collective
-    ``enable_routed_input(mode)``. Returns ``(hop, enabled)``.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    with served_construction():
-        hop = PostMoeHop(env.mapping, fx.HIDDEN, fx.TOP_K, torch.bfloat16, raw_producer=producer)
-        enabled = hop.engaged and hop.enable_routed_input(mode)
-    return hop, enabled
-
-
-def build_stack(env, layers, check_buffers=True):
-    """Served MoE layers (real weights), the post_moe_hop with routed input on.
-
-    With ``check_buffers`` the routed-input buffers must be device memory on
-    every rank: the served construction must not leave them on the meta device.
-    """
-    CB.set_features("")
-    moes, gammas = CB.build_layers(env, layers)
-    hop, enabled = build_hop(env, next(iter(moes.values())).experts)
-    assert hop.engaged, hop.reason
-    assert enabled
-    if check_buffers:
-        problems = buffer_problems(hop)
-        assert U.all_ranks(not problems), f"rank {env.rank}: {problems}"
-    for moe in moes.values():
-        moe._routed_into_hop = hop
-    return moes, gammas, hop
-
-
-def produce(moe, x, producer: str):
-    """The MoE half of one call (a decode step); returns ``(kind, payload)`` without the hop."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    CB.set_features(FEATURES[producer])
-    with dispatch.step(dispatch.Phase.DECODE):
-        out = moe.forward(x, do_finalize=False)
-    if len(out) == 2:
-        assert producer == "routed_moe_into_hop", f"{producer}: unexpected [shared, routed] result"
-        return "routed", out
-    assert producer != "routed_moe_into_hop", (
-        "routed_moe_into_hop did not engage (raw tuple returned)"
-    )
-    return "raw", out
-
-
-def hop_call(hop, kind, payload, res, gamma, with_shared: bool):
-    """The post_moe_hop half of one call (collective)."""
-    if kind == "routed":
-        shared, routed = payload
-        return hop.hop_routed(routed, shared if with_shared else None, res, gamma, fx.RMS_EPS)
-    shared, fc2, scale, perm = payload
-    return hop(fc2, scale, perm, shared if with_shared else None, res, gamma, fx.RMS_EPS)
-
-
-def local_partial(kind, payload, with_shared: bool, hidden: int = fx.HIDDEN, hop=None):
-    """The rank partial the hop publishes: ``sanitize(bf16(routed + shared))``."""
-    if kind == "routed":
-        shared, routed = payload
-        if routed.dtype == torch.float32:
-            routed = routed.to(torch.bfloat16)
-        elif routed.shape[0] != shared.shape[0]:
-            # identity-layout raw rows; the routing weights are in the hop
-            rows = shared.shape[0]
-            ident = torch.arange(rows * fx.TOP_K, dtype=torch.int32, device=routed.device)
-            routed = U.reference_local_partial(
-                routed, hop._routed_scales[:rows], ident.view(rows, fx.TOP_K), None, hidden
-            )
-        local = routed if not with_shared else shared + routed
-    else:
-        shared, fc2, scale, perm = payload
-        local = U.reference_local_partial(fc2, scale, perm, shared if with_shared else None, hidden)
-    return U.sanitize_bf16(local)
-
-
-def call(moe, hop, x, res, gamma, producer: str, with_shared: bool):
-    """One call; returns ``(outputs, kind, payload)``.
-
-    In mode "fp32" the returned payload holds a copy of the accumulator taken
-    before the hop read and cleared the hop's own buffer.
-    """
-    kind, payload = produce(moe, x, producer)
-    seen = payload
-    if kind == "routed" and payload[1].dtype == torch.float32:
-        seen = [payload[0], payload[1].clone()]
-    out = hop_call(hop, kind, payload, res, gamma, with_shared)
-    return out, kind, seen
-
-
-def exact_mode(hop) -> bool:
-    """Whether routed_moe_into_hop's partial must equal the raw variant's bit for bit."""
-    return getattr(hop, "routed_mode", "row") != "fp32"
-
-
-def reorder_ok(local, ref_payload, with_shared: bool):
-    """FP32-reorder bound of a routed_moe_into_hop "fp32" partial against the raw variant's."""
-    shared, fc2, scale, perm = ref_payload
-    ref = local_partial("raw", ref_payload, with_shared)
-    routed_ref = local_partial("raw", ref_payload, False)
-    idx = perm.to(torch.long).clamp(min=0)
-    rows = fc2[idx.view(-1)].view(perm.shape[0], perm.shape[1], -1).float()
-    mag = (scale.float()[:, :, None] * rows).abs().sum(1)
-    # one BF16 rounding of the routed sum, one of the shared add (|x| 2^-7 >= ulp(x))
-    ulp = (ref.float().abs() + routed_ref.float().abs()).clamp(min=2.0**-126) * 2.0**-7
-    bound = REORDER_REL * mag + ulp
-    diff = (local.float() - ref.float()).abs()
-    return bool((diff <= bound).all()), int((diff > 0).sum())
-
-
-def oracle_producer(producer: str) -> str:
-    return "routed_moe" if producer == "routed_moe_into_hop" else producer
-
-
-def check_emulation(env, hop, kind, payload, with_shared, res, gamma, out, res_out):
-    """Published slots == gathered expected partials; outputs == native emulation.
-
-    Collective (gathers every rank's expected partial). Must directly follow
-    the checked call (reads the workspace's last buffer).
-    """
-    rows = res.shape[0]
-    mine = local_partial(kind, payload, with_shared, hop=hop)
-    expected = U.gather_bf16(mine).permute(1, 0, 2).cuda()
-    published = U.published_partials(hop, rows)
-    slots_exact = torch.equal(published.view(torch.int16), expected.view(torch.int16))
-    emu_norm, emu_res = U.emulate_post_reduction(expected, res, gamma, fx.RMS_EPS)
-    res_exact = torch.equal(res_out.view(torch.int16), emu_res.view(torch.int16))
-    norm_ulp = U.compare_stats(out, emu_norm)["max_ulp"]
-    return slots_exact and res_exact and norm_ulp <= NORM_MAX_ULP, {
-        "slots_exact": slots_exact,
-        "res_exact": res_exact,
-        "norm_max_ulp": norm_ulp,
-    }
-
-
-def same_or_close(hop, got, ref) -> bool:
-    """Bit-equal in the exact modes, close in mode "fp32".
-
-    In mode "fp32" the routed_moe_into_hop ranks' top-k sums are reordered, so the reduced
-    outputs differ by the BF16 rounding of the reordered partials (bounded
-    exactly per call by ``run_mixed``; here a small bound on residual and norm).
-    """
-    if exact_mode(hop):
-        return identical(got, ref)
-    ok = True
-    for a, b in zip(got, ref):
-        d = (a.float() - b.float()).abs()
-        ok &= bool((d <= 2.0**-6 * b.float().abs() + 2.0**-10).all())
-    return ok
-
-
-def identical(a, b) -> bool:
-    return torch.equal(a[0].view(torch.int16), b[0].view(torch.int16)) and torch.equal(
-        a[1].view(torch.int16), b[1].view(torch.int16)
-    )
-
-
-def warm(env, moes, rows_list, producers=("routed_moe_into_hop", "routed_moe", "native")):
-    """Compile every producer variant on every rank before any collective call.
-
-    A rank that compiles (seconds) while its peers already poll inside the post_moe_hop
-    kernel would make them hit the spin limit; this runs only the MoE half.
-    """
-    for moe in moes.values():
-        for rows in rows_list:
-            x = torch.zeros(rows, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
-            for producer in producers:
-                produce(moe, x, producer)
-    CB.set_features("")
-    torch.cuda.synchronize()
-    env.comm.barrier()
-
-
-def mixed_pattern(env, pattern: str):
-    """Per-rank producer of a mixed-variant call."""
-    if pattern == "odd-routed_moe":
-        return "routed_moe" if env.rank % 2 else "routed_moe_into_hop"
-    if pattern == "rank0-routed_moe":
-        return "routed_moe" if env.rank == 0 else "routed_moe_into_hop"
-    if pattern == "host1-routed_moe":
-        return "routed_moe" if env.rank >= env.world // 2 else "routed_moe_into_hop"
-    if pattern == "odd-native":
-        return "native" if env.rank % 2 else "routed_moe_into_hop"
-    if pattern == "all-routed_moe_into_hop":
-        return "routed_moe_into_hop"
-    raise ValueError(pattern)
-
-
-def run_mixed(env, moes, gammas, hop, pattern, with_shared, rows_list, per_rows=3, seed=0):
-    """Mixed-variant calls vs their all-raw-variant oracle; returns (ok, n, failures)."""
-    producer = mixed_pattern(env, pattern)
-    g = torch.Generator(device="cuda").manual_seed(900 + seed)
-    failures, n = [], 0
-    for layer in moes:
-        for rows in rows_list:
-            for i, x in enumerate(windows(layer, rows, per_rows, offset=seed)):
-                res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
-                got, kind, payload = call(
-                    moes[layer], hop, x, res, gammas[layer], producer, with_shared
-                )
-                torch.cuda.synchronize()
-                emu_ok, stats = check_emulation(
-                    env, hop, kind, payload, with_shared, res, gammas[layer], *got
-                )
-                mine = local_partial(kind, payload, with_shared, hop=hop)
-                ref, ref_kind, ref_payload = call(
-                    moes[layer], hop, x, res, gammas[layer], oracle_producer(producer), with_shared
-                )
-                torch.cuda.synchronize()
-                if kind == "routed" and not exact_mode(hop):
-                    same, n_diff = reorder_ok(mine, ref_payload, with_shared)
-                    stats["reorder_elems"] = n_diff
-                else:
-                    same = identical(got, ref)
-                if not (same and emu_ok):
-                    failures.append((layer, rows, i, stats, same))
-                n += 1
-    CB.set_features("")
-    return not failures, n, failures
-
-
-def run_stress(
-    env,
-    moes,
-    gammas,
-    hop,
-    batches: int,
-    batch: int,
-    max_delay_us: float,
-    rows_list=(1, 6, 8),
-    seed: int = 0,
-):
-    """Back-to-back batches with random per-rank producers and device delays.
-
-    No host synchronization separates the calls of a batch; every call is
-    checked bit for bit against its all-raw-variant oracle afterwards. The row
-    count, layer and window of a call are drawn from a generator seeded
-    identically on every rank (collective shapes must agree); the producer,
-    the shared-output choice and the delay are drawn per rank. Kernel-internal
-    delays come from ``routed_moe.DELAY_INJECTION`` if the caller set it (and compiled them
-    with ``warm``).
-    """
-    common = random.Random(31 + seed)
-    local = random.Random(1000 + 17 * seed + env.rank)
-    cycles_per_us = 2000.0  # approximate sm_103 SM clock; only the magnitude matters
-    layers = list(moes)
-    gen = torch.Generator(device="cuda").manual_seed(4000 + seed)
-    failures, calls, delayed, producers_seen = [], 0, 0, set()
-    for b in range(batches):
-        plan = []
-        for j in range(batch):
-            layer = layers[common.randrange(len(layers))]
-            rows = rows_list[common.randrange(len(rows_list))]
-            x = windows(layer, rows, 1, offset=common.randrange(1 << 16))[0]
-            res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=gen) * 4.0).bfloat16()
-            producer = ("routed_moe_into_hop", "routed_moe_into_hop", "routed_moe", "native")[
-                local.randrange(4)
-            ]
-            with_shared = local.random() < 0.75
-            delay = local.uniform(0.0, max_delay_us) if local.random() < 0.3 else 0.0
-            plan.append((layer, x, res, producer, with_shared, delay))
-            producers_seen.add(producer)
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        outs = []
-        for layer, x, res, producer, with_shared, delay in plan:
-            if delay:
-                torch.cuda._sleep(int(delay * cycles_per_us))
-                delayed += 1
-            got, _, _ = call(moes[layer], hop, x, res, gammas[layer], producer, with_shared)
-            outs.append(got)
-        torch.cuda.synchronize()
-        for j, ((layer, x, res, producer, with_shared, _), got) in enumerate(zip(plan, outs)):
-            ref, _, _ = call(
-                moes[layer], hop, x, res, gammas[layer], oracle_producer(producer), with_shared
-            )
-            torch.cuda.synchronize()
-            if not same_or_close(hop, got, ref):
-                failures.append((b, j, layer, x.shape[0], producer, with_shared))
-        calls += len(plan)
-    CB.set_features("")
-    bits = hop.error_bits()
-    return not failures and bits == 0, {
-        "calls": calls,
-        "delayed_on_rank": delayed,
-        "producers_on_rank": sorted(producers_seen),
-        "error_bits": bits,
-        "failures": failures[:6],
-    }
-
-
-def run_graph(env, moes, gammas, hop, pattern: str, replays: int = 4, rows: int = 6):
-    """A multi-stream CUDA graph (as served) of four mixed-producer calls.
-
-    Per-rank producers come from ``pattern``, the shared output alternates;
-    the graph is replayed on fresh real windows and every replay is checked
-    bit for bit against the eager all-raw-variant oracle. Returns
-    ``(ok, replays_checked)``.
-    """
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    layers = list(moes)
-    producer = mixed_pattern(env, pattern)
-    plan = [(layers[i % len(layers)], i % 2 == 0) for i in range(4)]
-    wins = {layer: windows(layer, rows, replays + 1, offset=11) for layer in layers}
-    xs = [wins[layer][0].clone() for layer, _ in plan]
-    g = torch.Generator(device="cuda").manual_seed(77)
-    ress = [
-        (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16() for _ in plan
-    ]
-
-    def body():
-        return [
-            call(moes[layer], hop, x, r, gammas[layer], producer, sh)[0]
-            for (layer, sh), x, r in zip(plan, xs, ress)
-        ]
-
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(2):
-            body()
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    graph = torch.cuda.CUDAGraph()
-    with with_multi_stream(True):
-        with torch.cuda.graph(graph, stream=stream):
-            outs = body()
-    torch.cuda.synchronize()
-    ok = True
-    for rep in range(1, replays + 1):
-        for (layer, _), x in zip(plan, xs):
-            x.copy_(wins[layer][rep])
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        graph.replay()
-        torch.cuda.synchronize()
-        for (layer, sh), x, r, got in zip(plan, xs, ress, outs):
-            ref, _, _ = call(moes[layer], hop, x, r, gammas[layer], oracle_producer(producer), sh)
-            torch.cuda.synchronize()
-            ok &= same_or_close(hop, got, ref)
-    CB.set_features("")
-    del graph
-    return ok and hop.error_bits() == 0, replays
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py
deleted file mode 100644
index 7b2677e9845e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_chain_bench.py
+++ /dev/null
@@ -1,764 +0,0 @@
-# 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.
-"""Real-neighbour benchmark of the routed MoE inside the served MoE block.
-
-Runs one process per GPU of the two-host TP8 group (``srun --mpi=pmix``, 8
-tasks). Every call is one target MoE layer of the served graph, built from the
-model's own modules with real GLM-5.2 NVFP4 weights:
-
-    router GEMM (main stream) || BF16 shared expert (aux stream)
-    -> routed MoE (native TRT-LLM-Gen runner, or routed_moe through the target's MoE)
-    -> stream join -> post-MoE hop
-
-The MoE is the built-in ``Deepseekv3MoE`` adopted by the target's ``GlmMoE``; its forward
-runs inside a decode step, with multi-stream on (as the CUDA-graph runner captures it) and
-``TRTLLM_ENABLE_PDL`` as served. The post-MoE hop is
-
-* ``fin`` (post_moe_hop off): ``shared + routed`` inside ``Deepseekv3MoE.forward``, then
-  the native MNNVL ``AllReduce(RESIDUAL_RMS_NORM)`` with the next layer's norm;
-* ``raw`` (post_moe_hop on): ``forward(do_finalize=False)`` and the post_moe_hop on the raw
-  tuple, exactly as the target's decoder layer calls it in ``forward_MoE``.
-
-Each call's hop is the predecessor of the next call's router GEMM. Calls
-rotate over several layers' weights (``--layers``), so a call never finds its
-experts in L2. One CUDA graph holds ``--calls`` calls; a trial replays it
-``--replays`` times after an MPI barrier; per-call time = trial time /
-(replays * calls); the summary is the median trial of the slowest rank.
-
-Inputs: ``--mode real`` uses the captured post-norm hidden states (real
-routing through the real router GEMM, consecutive-token windows of R rows,
-drawn from every position of the captured prompts);
-``--mode sweep`` replaces the router GEMM's result (the GEMM still runs) by
-constructed logits that route to exactly D distinct experts; ``--mode mix``
-draws every call's D from a served distribution (``--mix``, see ``MIXES``)
-and constructs its logits the same way, with fresh experts per call.
-
-Router-logit producer (``--logits``):
-
-* ``gate``: the MoE's own router GEMM runs on the main stream before the
-  routed MoE (native, router_gemm/W2 off). The shared expert starts during that GEMM.
-* ``supplied``: the router contract with logits present (post_attention_hop/router_gemm on): the
-  pre-MoE hop already produced the FP32 logits, the MoE's gate is skipped (it returns the
-  call's logits without launching anything), so the routed MoE and the
-  shared expert become ready at the same moment. Config labels get a
-  ``+sup`` suffix. The logits are the real router GEMM's output for the
-  call's hidden states (computed once before capture), or the sweep logits.
-
-Optional served neighbours before every call (identical in every arm):
-``--attn-mb N`` streams N MB of weights (a reduction over a rotating buffer,
-the stand-in for the attention block's weight reads, which also flushes L2
-as the served graph does), then ``--pre-hop ar`` runs the native two-host
-``AllReduce(RESIDUAL_RMS_NORM, trigger_completion_at_end=False)`` on a
-rank-local partial (the stand-in for the pre-MoE hop that precedes the MoE
-block; with ``--logits supplied`` it plays post_attention_hop's hop, which emits the logits).
-"""
-
-import argparse
-import json
-import os
-import statistics
-import sys
-import time
-
-import torch
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, HERE)
-sys.path.insert(
-    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
-)
-sys.path.insert(0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels"))
-import plan_testlib  # noqa: E402
-import post_moe_hop_mpi_utils as U  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-ACCUMULATION_ENV = "TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION"
-# routed_moe_into_hop routed-input modes (see routed_moe.routed_moe.routed_output_into_hop); the
-# routed kernels run at fixed launch parameters, so the routed-input mode is the only arm knob.
-ARM_ENV = {
-    "routed_moe_into_hop_I": {ACCUMULATION_ENV: "identity"},
-    "routed_moe_into_hop_F": {ACCUMULATION_ENV: "fp32"},
-    "routed_moe_into_hop_R": {ACCUMULATION_ENV: "row"},
-}
-ARM_KNOBS = (ACCUMULATION_ENV,)
-# Generic arm syntax ``routed_moe[K=V+K=V]`` / ``routed_moe_into_hop[...]`` with these short keys.
-ARM_KEYS = {"MODE": ACCUMULATION_ENV}
-# Served distinct-expert distributions at R = 6, from a server profile of
-# the routed_moe arms (fixture routed_moe-server-profile,
-# xc_end_histogram_rank0.csv: the share of two-round XC calls, i.e. D above
-# the cluster count, at 24 / 28 / 32 / 36 clusters). Bins (lo, hi, share);
-# inside a bin D is uniform, except the open top bin, which decays by 0.7
-# per expert from 37 (the real prefill windows reach p90 38, max 47).
-MIXES = {
-    # late in a decode segment: the distribution of D over its bins
-    "late": ((17, 24, 0.058), (25, 28, 0.077), (29, 32, 0.213), (33, 36, 0.195), (37, 48, 0.457)),
-    # early in a segment: fewer large D; the split below 33 follows the late
-    # window's proportions
-    "early": ((17, 24, 0.084), (25, 28, 0.111), (29, 32, 0.308), (33, 36, 0.380), (37, 48, 0.117)),
-}
-ARMS_REPORTED = tuple(["routed_moe", "routed_moe_into_hop"] + list(ARM_ENV))
-NUM_LAYERS = 78
-
-
-class _PassAllReduce(torch.nn.Module):
-    """Pass-through stand-in for the MoE module's own AllReduce.
-
-    The decoder layer disables that all-reduce (``enable_allreduce=False``
-    under POST_MOE_FUSION), so it only ever passes its input through.
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-
-    @property
-    def output_buffer_kind(self) -> int:
-        from tensorrt_llm.bindings.internal.thop import BufferKind
-
-        return int(BufferKind.DEFAULT)
-
-    def forward(self, input, *, all_reduce_params=None):
-        assert all_reduce_params is not None and not all_reduce_params.enable_allreduce
-        return input
-
-
-def arm_env(impl: str) -> dict:
-    """Environment of an arm: ``ARM_ENV[impl]`` or the generic ``base[K=V+...]``."""
-    if "[" not in impl:
-        return dict(ARM_ENV.get(impl, {}))
-    base, _, spec = impl.partition("[")
-    assert base in ("routed_moe", "routed_moe_into_hop") and spec.endswith("]"), impl
-    env = {}
-    for item in [v for v in spec[:-1].split("+") if v]:
-        key, _, val = item.partition("=")
-        env[ARM_KEYS[key]] = val
-    return env
-
-
-def mix_distinct(name: str, calls: int, rows: int, seed: int):
-    """Per-call distinct-expert counts drawn from ``MIXES[name]``."""
-    values, weights = [], []
-    for lo, hi, share in MIXES[name]:
-        if hi == fx.TOP_K * rows and lo > 32:
-            span = [0.7 ** (d - lo) for d in range(lo, hi + 1)]
-        else:
-            span = [1.0] * (hi - lo + 1)
-        total = sum(span)
-        for d, w in zip(range(lo, hi + 1), span):
-            values.append(d)
-            weights.append(share * w / total)
-    g = torch.Generator().manual_seed(seed)
-    idx = torch.multinomial(torch.tensor(weights), calls, replacement=True, generator=g)
-    return [values[i] for i in idx.tolist()]
-
-
-def mix_logits(moes, order, rows: int, distinct_per_call, seed: int):
-    """Constructed logits routing call ``i`` to ``distinct_per_call[i]`` experts,
-    drawn fresh per call (the attention stand-in evicts L2 between calls)."""
-    out = []
-    for i, (layer, d) in enumerate(zip(order, distinct_per_call)):
-        sets = fx.expert_sets_with_distinct(rows, d, seed=seed + i)
-        perm = torch.randperm(
-            fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(seed + 7919 * i)
-        ).tolist()
-        uniq = sorted(set(e for row in sets for e in row))
-        remap = {e: perm[j] for j, e in enumerate(uniq)}
-        sets = [[remap[e] for e in row] for row in sets]
-        bias = moes[layer].gate.e_score_correction_bias.to(torch.bfloat16)
-        lg = fx.logits_for_expert_sets(sets, bias, seed=seed + i)
-        ids, _ = fx.routing_reference(lg, bias)
-        assert len(set(ids.flatten().tolist())) == d
-        out.append(lg)
-    return out
-
-
-def distinct_stats(moes, order, logits_per_call):
-    """Distinct experts per call (native routing of the logits the MoE reads)."""
-    ds = []
-    for layer, lg in zip(order, logits_per_call):
-        bias = moes[layer].gate.e_score_correction_bias.to(torch.bfloat16)
-        ids, _ = fx.routing_reference(lg.float().cpu(), bias.cpu())
-        ds.append(len(set(ids.flatten().tolist())))
-    n = len(ds)
-    above = {c: round(100.0 * sum(d > c for d in ds) / n, 1) for c in (24, 26, 28, 32, 36)}
-    return {"mean": round(sum(ds) / n, 2), "median": sorted(ds)[n // 2], "pct_above": above}
-
-
-def set_features(ids: str) -> None:
-    """Install, for this process, the plan that runs exactly the fast paths ``ids`` (comma-separated)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    plan_testlib.select(dispatch, ids)
-
-
-def next_norm_weight(layer: int) -> torch.Tensor:
-    name = (
-        "model.norm.weight"
-        if layer + 1 >= NUM_LAYERS
-        else f"model.layers.{layer + 1}.input_layernorm.weight"
-    )
-    return fx._Reader().full(name).to(torch.bfloat16).cuda()
-
-
-class Engagement:
-    """Counts routed_moe adapter calls that engaged (returned a result)."""
-
-    def __init__(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-        dispatch.get_impl("routed_moe", "routed_output")  # import + register
-        self.counts = {}
-        for name in ("routed_output", "routed_output_into_hop"):
-            inner = dispatch._REGISTRY["routed_moe"].get(name)
-            if inner is None:
-                continue
-
-            def counted(*args, _inner=inner, _name=name, **kwargs):
-                out = _inner(*args, **kwargs)
-                key = (_name, out is not None)
-                self.counts[key] = self.counts.get(key, 0) + 1
-                return out
-
-            dispatch.register("routed_moe", name, counted)
-
-    def engaged(self, name: str) -> int:
-        return self.counts.get((name, True), 0)
-
-
-class Neighbours:
-    """Served kernels that precede each MoE block (identical in every arm).
-
-    ``attn_mb``: a weight-streaming stand-in for the attention block (one
-    reduction over ``attn_mb`` MB, two buffers used alternately so the reads
-    miss L2 as the served attention weights do). ``pre_hop``: the native
-    two-host fused all-reduce + residual + RMSNorm on a rank-local partial.
-    """
-
-    def __init__(self, env, rows: int, calls: int, attn_mb: int, pre_hop: str):
-        g = torch.Generator(device="cuda").manual_seed(500 + env.rank)
-        self.attn = []
-        if attn_mb > 0:
-            numel = attn_mb * (1 << 20) // 2
-            self.attn = [
-                torch.randn(numel, device="cuda", generator=g).bfloat16() for _ in range(2)
-            ]
-        self.pre = None
-        if pre_hop == "ar":
-            common = torch.Generator(device="cuda").manual_seed(600)
-            self.pre = [
-                (
-                    (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.05).bfloat16(),
-                    (torch.randn(rows, fx.HIDDEN, device="cuda", generator=common)).bfloat16(),
-                )
-                for _ in range(calls)
-            ]
-            self.gamma = (torch.rand(fx.HIDDEN, device="cuda", generator=common) + 0.5).bfloat16()
-        self.sink = []
-
-    def step(self, i: int):
-        if self.attn:
-            self.sink.append(self.attn[i % 2].amax())
-        if self.pre is not None:
-            partial, residual = self.pre[i]
-            self.sink.append(U.native_fused_allreduce(partial, residual, self.gamma, fx.RMS_EPS))
-
-
-class Chain:
-    """The per-rank call sequence of one arm, captured in one CUDA graph."""
-
-    def __init__(
-        self,
-        env,
-        moes,
-        gammas,
-        order,
-        xs,
-        residuals,
-        logits,
-        hop,
-        mode,
-        impl,
-        supplied=None,
-        neighbours=None,
-    ):
-        self.env, self.moes, self.gammas = env, moes, gammas
-        self.order, self.xs, self.residuals, self.logits = order, xs, residuals, logits
-        self.hop, self.mode, self.impl = hop, mode, impl
-        self.supplied, self.neighbours = supplied, neighbours
-        self.outs = []
-
-    def run(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-        with dispatch.step(dispatch.Phase.DECODE):
-            self._run()
-
-    def _run(self):
-        from tensorrt_llm._torch.distributed import AllReduceParams
-
-        self.outs = []
-        if self.neighbours is not None:
-            self.neighbours.sink = []
-        for i, layer in enumerate(self.order):
-            moe = self.moes[layer]
-            moe._routed_moe_bench_logits = None if self.logits is None else self.logits[i]
-            moe._routed_moe_bench_supplied = None if self.supplied is None else self.supplied[i]
-            if self.neighbours is not None:
-                self.neighbours.step(i)
-            x, res, gamma = self.xs[i], self.residuals[i], self.gammas[layer]
-            if self.mode == "fin":
-                h = moe.forward(
-                    x,
-                    do_finalize=True,
-                    final_all_reduce_params=AllReduceParams(enable_allreduce=False),
-                )
-                out = U.native_fused_allreduce(h, res, gamma, fx.RMS_EPS)
-            else:
-                res_moe = moe.forward(x, do_finalize=False)
-                if len(res_moe) == 2:
-                    # routed_moe_into_hop: [shared, routed] with routed finalized into the hop's
-                    # buffer
-                    shared, routed = res_moe
-                    out = self.hop.hop_routed(routed, shared, res, gamma, fx.RMS_EPS)
-                else:
-                    shared, fc2, scale, perm = res_moe
-                    out = self.hop(fc2, scale, perm, shared, res, gamma, fx.RMS_EPS)
-            self.outs.append(out)
-
-
-def build_layers(env, layers):
-    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe  # noqa: E501, F401  (registers routed_moe)
-
-    saved = fx._NoAllReduce
-    fx._NoAllReduce = _PassAllReduce
-    moes, gammas = {}, {}
-    try:
-        for layer in layers:
-            t0 = time.time()
-            full = fx.load_layer_full(layer)
-            moes[layer] = fx.build_deepseek_moe(layer, env.rank, full)
-            del full
-            gammas[layer] = next_norm_weight(layer)
-            if env.rank == 0:
-                print(f"[chain] layer {layer} built in {time.time() - t0:.1f}s", flush=True)
-    finally:
-        fx._NoAllReduce = saved
-    for moe in moes.values():
-        gate = moe.gate
-        orig = gate.forward
-
-        def forward(h, moe=moe, orig=orig):
-            if moe._routed_moe_bench_supplied is not None:
-                # the router contract with logits present: the gate is skipped
-                return moe._routed_moe_bench_supplied
-            out = orig(h)
-            if moe._routed_moe_bench_logits is not None:
-                # sweep: the router GEMM still runs as served; the routed MoE
-                # reads the constructed logits instead (no extra kernel)
-                out = moe._routed_moe_bench_logits
-            return out
-
-        gate.forward = forward
-        moe._routed_moe_bench_logits = None
-        moe._routed_moe_bench_supplied = None
-        moe._routed_moe_bench_router = orig
-    return moes, gammas
-
-
-def supplied_logits(moes, order, xs, logits):
-    """Per call, the FP32 logits a pre-MoE hop would hand over (static buffers)."""
-    out = []
-    with torch.no_grad():
-        for i, layer in enumerate(order):
-            if logits is not None:
-                lg = logits[i].float().contiguous().clone()
-            else:
-                lg = moes[layer]._routed_moe_bench_router(xs[i]).float().contiguous().clone()
-            out.append(lg)
-    torch.cuda.synchronize()
-    return out
-
-
-def real_inputs(layers, rows: int, calls: int, seed: int):
-    """Per call: (layer, x) from consecutive-token windows of the capture."""
-    per_layer = {}
-    for layer in layers:
-        win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer), rows)]
-        assert win and win[0] is not None, f"no hidden states captured for layer {layer}"
-        g = torch.Generator().manual_seed(seed + layer)
-        per_layer[layer] = [win[i] for i in torch.randperm(len(win), generator=g).tolist()]
-    order, xs = [], []
-    for i in range(calls):
-        layer = layers[i % len(layers)]
-        order.append(layer)
-        xs.append(per_layer[layer][(i // len(layers)) % len(per_layer[layer])].cuda())
-    return order, xs
-
-
-def sweep_logits(moes, order, rows: int, distinct: int, seed: int):
-    """Constructed logits routing each call to ``distinct`` experts.
-
-    Consecutive uses of one layer draw from disjoint expert pools.
-    """
-    pools = max(1, fx.NUM_EXPERTS // distinct)
-    perm = torch.randperm(fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(seed)).tolist()
-    out = []
-    for i, layer in enumerate(order):
-        sets = fx.expert_sets_with_distinct(rows, distinct, seed=seed + i)
-        p = (i // len(set(order))) % pools
-        pool = perm[p * distinct : (p + 1) * distinct]
-        uniq = sorted(set(e for row in sets for e in row))
-        remap = {e: pool[j] for j, e in enumerate(uniq)}
-        sets = [[remap[e] for e in row] for row in sets]
-        bias = moes[layer].gate.e_score_correction_bias.to(torch.bfloat16)
-        lg = fx.logits_for_expert_sets(sets, bias, seed=seed + i)
-        ids, _ = fx.routing_reference(lg, bias)
-        assert len(set(ids.flatten().tolist())) == distinct
-        out.append(lg)
-    return out
-
-
-def capture(env, chain):
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(2):
-            chain.run()
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    graph = torch.cuda.CUDAGraph()
-    with with_multi_stream(True):
-        with torch.cuda.graph(graph, stream=stream):
-            chain.run()
-    torch.cuda.synchronize()
-    for _ in range(2):
-        env.comm.barrier()
-        graph.replay()
-        torch.cuda.synchronize()
-    return graph
-
-
-def time_graph(env, graph, replays: int):
-    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    start.record()
-    for _ in range(replays):
-        graph.replay()
-    end.record()
-    end.synchronize()
-    return start.elapsed_time(end) * 1000.0
-
-
-def compare(env, a_outs, b_outs, label):
-    """Max |a - b| of the final hidden and residual over every call (sanity)."""
-    worst = [0.0, 0.0]
-    for a, b in zip(a_outs, b_outs):
-        for j in range(2):
-            worst[j] = max(worst[j], (a[j].float() - b[j].float()).abs().max().item())
-    allw = env.comm.allgather(worst)
-    if env.rank == 0:
-        print(
-            f"[chain] check {label}: max|d norm_out| {max(w[0] for w in allw):.3e} "
-            f"max|d residual| {max(w[1] for w in allw):.3e}",
-            flush=True,
-        )
-
-
-def run_config(
-    env, args, moes, gammas, hop, eng, label, order, xs, logits, supplied=None, neighbours=None
-):
-    torch.manual_seed(0)
-    g = torch.Generator(device="cuda").manual_seed(11)
-    residuals = [
-        (torch.randn(args.rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
-        for _ in order
-    ]
-    arms = [a for a in args.arms.split(",") if a]
-    graphs, chains = {}, {}
-    # Every arm is captured ``--captures`` times, in forward and then reverse
-    # arm order: an occasional graph replays with a per-call inter-kernel gap
-    # that is a property of that one instantiation (identical arms captured
-    # elsewhere do not show it), so an arm reports its fastest capture.
-    passes = [arms if p % 2 == 0 else arms[::-1] for p in range(args.captures)]
-    for p, seq in enumerate(passes):
-        for arm in seq:
-            impl, mode = arm.split(":")
-            # "routed_moe" is the routed_moe schedule, "routed_moe_into_hop" the same schedule
-            # writing into the post_moe_hop (raw mode)
-            feats = (
-                ""
-                if impl == "native"
-                else (
-                    "post_moe_hop,routed_moe,routed_moe_into_hop"
-                    if impl.startswith("routed_moe_into_hop")
-                    else "routed_moe"
-                )
-            )
-            if supplied is not None:
-                # the composition's logit producer (post_attention_hop's hop) is selected
-                feats = ",".join(f for f in (feats, "post_attention_hop") if f)
-            set_features(feats)
-            for knob in ARM_KNOBS:
-                os.environ.pop(knob, None)
-            os.environ.update(arm_env(impl))
-            if impl.startswith("routed_moe_into_hop"):
-                # collective, outside capture: the hop's routed-input mode of this arm
-                assert hop.enable_routed_input(os.environ.get(ACCUMULATION_ENV) or None)
-            entry = (
-                "routed_output_into_hop"
-                if impl.startswith("routed_moe_into_hop")
-                else "routed_output"
-            )
-            e0 = eng.engaged(entry)
-            chain = Chain(
-                env,
-                moes,
-                gammas,
-                order,
-                xs,
-                residuals,
-                None if supplied is not None else logits,
-                hop,
-                mode,
-                impl,
-                supplied=supplied,
-                neighbours=neighbours,
-            )
-            graphs[(arm, p)] = capture(env, chain)
-            if p == 0:
-                chains[arm] = chain
-            engaged = eng.engaged(entry) - e0
-            want = 0 if impl == "native" else 3 * len(order)  # two warmups + capture
-            ok = U.all_ranks(engaged == want)
-            if env.rank == 0:
-                print(
-                    f"[chain] {label} arm {arm} capture {p}: routed_moe engaged {engaged}/{want} "
-                    f"all_ranks={ok}"
-                )
-            assert ok, f"{arm}: unexpected routed_moe engagement"
-    set_features("")
-    for knob in ARM_KNOBS:
-        os.environ.pop(knob, None)
-    per = {key: [] for key in graphs}
-    for _ in range(args.trials):
-        for key in graphs:
-            per[key].append(time_graph(env, graphs[key], args.replays))
-    if args.profile_config == label:
-        # one replay per arm inside a cudaProfilerApi capture range, each in
-        # an NVTX range named after the arm (nsys --capture-range=cudaProfilerApi)
-        torch.cuda.synchronize()
-        torch.cuda.profiler.start()
-        for arm, p in graphs:
-            env.comm.barrier()
-            torch.cuda.nvtx.range_push(f"{label}:{arm}" + (f"#{p}" if p else ""))
-            graphs[(arm, p)].replay()
-            torch.cuda.synchronize()
-            torch.cuda.nvtx.range_pop()
-        torch.cuda.profiler.stop()
-    rows = []
-    lg_seen = supplied if supplied is not None else logits
-    dstat = None
-    if lg_seen is not None:
-        dstat = distinct_stats(moes, order, lg_seen)
-    elif env.rank == 0:
-        dstat = distinct_stats(
-            moes,
-            order,
-            [moes[layer]._routed_moe_bench_router(x).float() for layer, x in zip(order, xs)],
-        )
-    if env.rank == 0:
-        print(f"CHAIN_DISTINCT {label} {json.dumps(dstat)}", flush=True)
-    for arm in arms:
-        caps = []
-        for p in range(args.captures):
-            per_call = [t / (args.replays * len(order)) for t in per[(arm, p)]]
-            meds = env.comm.allgather(statistics.median(per_call))
-            caps.append(
-                {
-                    "median": max(meds),
-                    "spread": [min(meds), max(meds)],
-                    "best": max(env.comm.allgather(min(per_call))),
-                    "mean": max(env.comm.allgather(statistics.mean(per_call))),
-                }
-            )
-        best = min(caps, key=lambda c: c["median"])
-        rows.append(
-            {
-                "config": label,
-                "arm": arm,
-                "per_call_us_median_slowest_rank": best["median"],
-                "per_call_us_median_rank_spread": best["spread"],
-                "per_call_us_best_trial_slowest_rank": best["best"],
-                "per_call_us_mean_slowest_rank": best["mean"],
-                "per_capture_median_slowest_rank": [c["median"] for c in caps],
-                "distinct": dstat,
-                "calls_per_graph": len(order),
-                "replays": args.replays,
-                "trials": args.trials,
-            }
-        )
-        if env.rank == 0:
-            print("CHAIN_BENCH " + json.dumps(rows[-1]), flush=True)
-            if len(caps) > 1:
-                m = [c["median"] for c in caps]
-                flag = "  <-- capture spread > 1 us" if max(m) - min(m) > 1.0 else ""
-                print(
-                    f"CHAIN_CAPTURES {label} {arm}: " + " ".join(f"{v:.2f}" for v in m) + flag,
-                    flush=True,
-                )
-    by = {r["arm"]: r["per_call_us_median_slowest_rank"] for r in rows}
-    by_mean = {r["arm"]: r["per_call_us_mean_slowest_rank"] for r in rows}
-    impls = [a.split(":")[0] for a in arms if not a.startswith("native:")]
-    impls = list(dict.fromkeys(impls))
-    if env.rank == 0:
-        for mode in ("fin", "raw"):
-            for impl in impls:
-                if f"native:{mode}" in by and f"{impl}:{mode}" in by:
-                    d = by[f"native:{mode}"] - by[f"{impl}:{mode}"]
-                    dm = by_mean[f"native:{mode}"] - by_mean[f"{impl}:{mode}"]
-                    print(
-                        f"CHAIN_DELTA {label} {mode}: native - {impl} = {d:+.2f} us/call "
-                        f"(trial mean {dm:+.2f})",
-                        flush=True,
-                    )
-        for impl in ("routed_moe", "routed_moe_into_hop"):
-            if "native:fin" in by and f"{impl}:raw" in by:
-                d = by["native:fin"] - by[f"{impl}:raw"]
-                print(f"CHAIN_DELTA {label} native:fin - {impl}:raw = {d:+.2f} us/call", flush=True)
-        if "routed_moe:raw" in by and "routed_moe_into_hop:raw" in by:
-            d = by["routed_moe:raw"] - by["routed_moe_into_hop:raw"]
-            print(
-                f"CHAIN_DELTA {label} routed_moe:raw - routed_moe_into_hop:raw = {d:+.2f} us/call (routed input)",
-                flush=True,
-            )
-    for mode in ("fin", "raw"):
-        for impl in impls:
-            if f"native:{mode}" in chains and f"{impl}:{mode}" in chains:
-                compare(
-                    env,
-                    chains[f"native:{mode}"].outs,
-                    chains[f"{impl}:{mode}"].outs,
-                    f"{label} {mode} {impl}",
-                )
-    del graphs
-    torch.cuda.synchronize()
-    return rows
-
-
-def main(argv=None):
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--layers", default="10,30,50,77")
-    ap.add_argument("--rows", type=int, default=6)
-    ap.add_argument("--calls", type=int, default=48)
-    ap.add_argument("--replays", type=int, default=10)
-    ap.add_argument("--trials", type=int, default=9)
-    ap.add_argument("--captures", type=int, default=2, help="graphs per arm (best one reported)")
-    ap.add_argument("--arms", default="native:fin,routed_moe:fin,native:raw,routed_moe:raw")
-    ap.add_argument("--mode", default="real", help="real | sweep | mix (comma list)")
-    ap.add_argument("--mix", default="early,late", help="served D distributions (MIXES)")
-    ap.add_argument("--distinct", default="8,14,21,25,31,38,48")
-    ap.add_argument("--seed", type=int, default=0)
-    ap.add_argument(
-        "--logits", default="gate", help="gate | supplied | gate,supplied (router-logit producer)"
-    )
-    ap.add_argument("--attn-mb", type=int, default=0, help="attention weight-stream stand-in")
-    ap.add_argument("--pre-hop", default="none", choices=("none", "ar"))
-    ap.add_argument("--out", default="")
-    ap.add_argument(
-        "--profile-config", default="", help="config label replayed once per arm under nsys"
-    )
-    args = ap.parse_args(argv)
-
-    env = U.mpi_env()
-    assert env.world == 8 and env.multi_node, "needs the two-host TP8 group"
-    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
-    set_features("")
-    layers = [int(v) for v in args.layers.split(",")]
-    moes, gammas = build_layers(env, layers)
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    # as the target's decoder layer builds it: the layer's MoE module declares the
-    # raw-tuple contract the hop consumes
-    hop = PostMoeHop(
-        env.mapping,
-        fx.HIDDEN,
-        fx.TOP_K,
-        torch.bfloat16,
-        raw_producer=next(iter(moes.values())).experts,
-    )
-    assert hop.engaged, hop.reason
-    if any(a.startswith("routed_moe_into_hop") for a in args.arms.split(",")):
-        assert hop.enable_routed_input(), (
-            "routed_moe_into_hop routed input not enabled on every rank"
-        )
-        for moe in moes.values():
-            moe._routed_into_hop = hop
-    eng = Engagement()
-    if env.rank == 0:
-        print(
-            f"[chain] ranks={env.world} layers={layers} rows={args.rows} calls={args.calls} "
-            f"pdl={os.environ.get('TRTLLM_ENABLE_PDL')} post_moe_hop={hop.impl}",
-            flush=True,
-        )
-    results = []
-    modes = args.mode.split(",")
-    producers = [p for p in args.logits.split(",") if p]
-    assert set(producers) <= {"gate", "supplied"}, args.logits
-    order, xs = real_inputs(layers, args.rows, args.calls, args.seed)
-    neighbours = None
-    if args.attn_mb > 0 or args.pre_hop != "none":
-        neighbours = Neighbours(env, args.rows, len(order), args.attn_mb, args.pre_hop)
-    if env.rank == 0:
-        print(
-            f"[chain] logits={producers} attn_mb={args.attn_mb} pre_hop={args.pre_hop}",
-            flush=True,
-        )
-
-    def configs(label, lg):
-        out = []
-        for producer in producers:
-            sup = supplied_logits(moes, order, xs, lg) if producer == "supplied" else None
-            name = label + ("+sup" if producer == "supplied" else "")
-            out += run_config(
-                env, args, moes, gammas, hop, eng, name, order, xs, lg, sup, neighbours
-            )
-        return out
-
-    if "real" in modes:
-        results += configs("real", None)
-    if "mix" in modes:
-        for name in [v for v in args.mix.split(",") if v]:
-            ds = mix_distinct(name, len(order), args.rows, args.seed + 5000)
-            lg = mix_logits(moes, order, args.rows, ds, args.seed + 6000)
-            results += configs(f"mix{name}", lg)
-    if "sweep" in modes:
-        for d in [int(v) for v in args.distinct.split(",") if v]:
-            if d < fx.TOP_K or d > fx.TOP_K * args.rows:
-                continue
-            lg = sweep_logits(moes, order, args.rows, d, args.seed + 1000 * d)
-            results += configs(f"D{d}", lg)
-    if env.rank == 0 and args.out:
-        with open(args.out, "w") as f:
-            json.dump(results, f, indent=1)
-    hop.raise_if_error()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py
deleted file mode 100644
index f97e676e55e7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe_fixtures.py
+++ /dev/null
@@ -1,1036 +0,0 @@
-# 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.
-"""Real-weight fixtures and the native oracle for the routed-MoE tests.
-
-The fixture is one GLM-5.2 NVFP4 MoE layer as TP rank ``r`` of TP8/EP1 sees
-it: every routed expert's gate/up rows ``[256r, 256r+256)`` and down columns
-``[256r, 256r+256)``, with block and global scales, the router weight and the
-router bias. The native module is the stock TRT-LLM-Gen NVFP4 leaf loaded
-through its own ``load_weights``/``post_load_weights``, so its buffers carry
-exactly the shuffled/interleaved layouts and derived scalars of a served model.
-Slicing the checkpoint before loading a ``tp_size=1`` module is equivalent to
-the TP8 loader's ``load_weight_shard`` (COLUMN for w1/w3, ROW for w2).
-
-Hidden-state proxies are real token embeddings passed through the layer's
-``post_attention_layernorm`` (labelled semi-synthetic: the served MoE input is
-the normalized residual stream of that layer, which needs a full forward).
-"""
-
-import json
-import os
-import tempfile
-from dataclasses import dataclass
-from typing import Dict, List, Optional
-
-import torch
-from utils.llm_data import llm_models_root
-
-CKPT_DIR = os.environ.get(
-    "GLM_MOE_DSA_ROUTED_MOE_CKPT_DIR",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-CACHE_DIR = os.environ.get(
-    "GLM_MOE_DSA_ROUTED_MOE_FIXTURE_CACHE",
-    os.path.join(tempfile.gettempdir(), "glm_moe_dsa_routed_moe_fixtures"),
-)
-
-HIDDEN = 6144
-INTER = 2048
-NUM_EXPERTS = 256
-TOP_K = 8
-TP_SIZE = 8
-ROUTED_SCALING = 2.5
-RMS_EPS = 1e-5
-
-# A fixed English/code text; consecutive tokens of it stand in for the
-# consecutive verify rows of one MTP request.
-PROMPT_TEXT = (
-    "def merge_sorted(a, b):\n    i = j = 0\n    out = []\n"
-    "    while i < len(a) and j < len(b):\n        if a[i] <= b[j]:\n"
-    "            out.append(a[i]); i += 1\n        else:\n"
-    "            out.append(b[j]); j += 1\n    return out + a[i:] + b[j:]\n"
-    "The quick brown fox jumps over the lazy dog while the compiler "
-    "schedules tensor-core instructions across the streaming multiprocessors."
-)
-
-
-def _index() -> Dict[str, str]:
-    with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
-        return json.load(f)["weight_map"]
-
-
-class _Reader:
-    """Opens each safetensors shard once and slices tensors lazily."""
-
-    def __init__(self):
-        from safetensors import safe_open
-
-        self._safe_open = safe_open
-        self._map = _index()
-        self._open = {}
-
-    def _handle(self, name: str):
-        shard = self._map[name]
-        if shard not in self._open:
-            self._open[shard] = self._safe_open(
-                os.path.join(CKPT_DIR, shard), framework="pt", device="cpu"
-            )
-        return self._open[shard]
-
-    def full(self, name: str) -> torch.Tensor:
-        return self._handle(name).get_tensor(name)
-
-    def rows(self, name: str, start: int, stop: int) -> torch.Tensor:
-        # clone: a row slice may be a view of the whole tensor's storage,
-        # which torch.save would otherwise write out in full
-        return self._handle(name).get_slice(name)[start:stop].clone()
-
-    def cols(self, name: str, start: int, stop: int) -> torch.Tensor:
-        return self._handle(name).get_slice(name)[:, start:stop].clone()
-
-
-def load_layer_rank(layer: int, tp_rank: int) -> Dict[str, torch.Tensor]:
-    """Rank-sliced raw checkpoint tensors of one MoE layer (CPU, cached)."""
-    os.makedirs(CACHE_DIR, exist_ok=True)
-    path = os.path.join(CACHE_DIR, f"glm52_L{layer}_tp{TP_SIZE}r{tp_rank}.pt")
-    if os.path.exists(path):
-        return torch.load(path, map_location="cpu")
-    rd = _Reader()
-    inter_rank = INTER // TP_SIZE
-    r0, r1 = tp_rank * inter_rank, (tp_rank + 1) * inter_rank
-    pre = f"model.layers.{layer}.mlp"
-    out: Dict[str, torch.Tensor] = {}
-    for e in range(NUM_EXPERTS):
-        for proj, key in (("gate_proj", "w1"), ("up_proj", "w3")):
-            base = f"{pre}.experts.{e}.{proj}"
-            out[f"{e}.{key}.weight"] = rd.rows(f"{base}.weight", r0, r1)
-            out[f"{e}.{key}.weight_scale"] = rd.rows(f"{base}.weight_scale", r0, r1)
-            out[f"{e}.{key}.weight_scale_2"] = rd.full(f"{base}.weight_scale_2")
-            out[f"{e}.{key}.input_scale"] = rd.full(f"{base}.input_scale")
-        base = f"{pre}.experts.{e}.down_proj"
-        out[f"{e}.w2.weight"] = rd.cols(f"{base}.weight", r0 // 2, r1 // 2)
-        out[f"{e}.w2.weight_scale"] = rd.cols(f"{base}.weight_scale", r0 // 16, r1 // 16)
-        out[f"{e}.w2.weight_scale_2"] = rd.full(f"{base}.weight_scale_2")
-        out[f"{e}.w2.input_scale"] = rd.full(f"{base}.input_scale")
-    out["gate.weight"] = rd.full(f"{pre}.gate.weight")
-    out["gate.e_score_correction_bias"] = rd.full(f"{pre}.gate.e_score_correction_bias")
-    out["post_attention_layernorm.weight"] = rd.full(
-        f"model.layers.{layer}.post_attention_layernorm.weight"
-    )
-    tmp = path + f".tmp{os.getpid()}"
-    torch.save(out, tmp)
-    os.replace(tmp, path)
-    return out
-
-
-def load_layer_full(layer: int) -> Dict[str, torch.Tensor]:
-    """Unsliced checkpoint tensors of one MoE layer (CPU, not cached).
-
-    Keys follow the MoE loader's convention (``{e}.w1.weight`` ...), so the
-    TP8 module's own ``load_weights`` does the rank slicing.
-    """
-    rd = _Reader()
-    pre = f"model.layers.{layer}.mlp"
-    out: Dict[str, torch.Tensor] = {}
-    for e in range(NUM_EXPERTS):
-        for proj, key in (("gate_proj", "w1"), ("up_proj", "w3"), ("down_proj", "w2")):
-            base = f"{pre}.experts.{e}.{proj}"
-            for suffix in ("weight", "weight_scale", "weight_scale_2", "input_scale"):
-                out[f"{e}.{key}.{suffix}"] = rd.full(f"{base}.{suffix}")
-    out["gate.weight"] = rd.full(f"{pre}.gate.weight")
-    out["gate.e_score_correction_bias"] = rd.full(f"{pre}.gate.e_score_correction_bias")
-    out["post_attention_layernorm.weight"] = rd.full(
-        f"model.layers.{layer}.post_attention_layernorm.weight"
-    )
-    return out
-
-
-def _served_model_config(layer: int, tp_rank: int, max_num_tokens: int = 128):
-    """ModelConfig of one served GEN rank: TP8/EP1 over two 4-GPU hosts,
-    no attention DP, TRTLLM MoE backend, NVFP4 experts, BF16 shared expert.
-    """
-    from transformers import PretrainedConfig
-
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm.mapping import Mapping
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-    from tensorrt_llm.quantization.mode import QuantAlgo
-
-    quant_config = QuantConfig(
-        quant_algo=QuantAlgo.NVFP4,
-        kv_cache_quant_algo=QuantAlgo.FP8,
-        group_size=16,
-        exclude_modules=[f"model.layers.{layer}.mlp.shared_experts*"],
-    )
-    pretrained = PretrainedConfig(
-        hidden_size=HIDDEN,
-        n_routed_experts=NUM_EXPERTS,
-        num_experts_per_tok=TOP_K,
-        moe_intermediate_size=INTER,
-        n_shared_experts=1,
-        n_group=1,
-        topk_group=1,
-        routed_scaling_factor=ROUTED_SCALING,
-    )
-    mapping = Mapping(
-        world_size=TP_SIZE,
-        rank=tp_rank,
-        gpus_per_node=4,
-        tp_size=TP_SIZE,
-        moe_tp_size=TP_SIZE,
-        moe_ep_size=1,
-    )
-    return ModelConfig(
-        pretrained_config=pretrained,
-        mapping=mapping,
-        quant_config=quant_config,
-        moe_backend="TRTLLM",
-        max_num_tokens=max_num_tokens,
-    )
-
-
-def _walk_post_load(root: torch.nn.Module) -> None:
-    """The model loader's full post-load walk over one module tree."""
-    for mod in root.modules():
-        hook = getattr(mod, "post_load_weights", None)
-        if hook is not None and not getattr(mod, "_weights_removed", False):
-            hook()
-
-
-def _aux_streams():
-    from tensorrt_llm._torch.utils import AuxStreamType
-
-    return {t: torch.cuda.Stream() for t in AuxStreamType}
-
-
-def build_served_experts(layer: int, tp_rank: int, full: Dict[str, torch.Tensor]):
-    """The routed experts exactly as ``Deepseekv3MoE`` builds them: ``create_moe``
-    with the served TP8 ``ModelConfig`` (a ``ConfigurableMoE`` around the
-    resolved TRT-LLM-Gen leaf), loaded from unsliced checkpoint tensors.
-
-    Returns ``(wrapper, bias)`` with the BF16 router bias the gate routes with.
-    """
-    from tensorrt_llm._torch.moe.fused_moe import create_moe
-    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
-
-    bias = full["gate.e_score_correction_bias"].to(torch.bfloat16).cuda()
-    routing = DeepSeekV3MoeRoutingMethod(
-        top_k=TOP_K,
-        n_group=1,
-        topk_group=1,
-        routed_scaling_factor=ROUTED_SCALING,
-        callable_e_score_correction_bias=lambda: bias,
-        is_fused=True,
-    )
-    mc = _served_model_config(layer, tp_rank)
-    wrapper = create_moe(
-        num_experts=NUM_EXPERTS,
-        routing_method=routing,
-        hidden_size=HIDDEN,
-        intermediate_size=INTER,
-        dtype=torch.bfloat16,
-        reduce_results=False,
-        model_config=mc,
-        override_quant_config=mc.quant_config,
-        aux_stream_dict=_aux_streams(),
-        layer_idx=layer,
-    )
-    weights = {k: v for k, v in full.items() if k[0].isdigit()}
-    wrapper.load_weights([weights])
-    _walk_post_load(wrapper)
-    wrapper = wrapper.cuda()
-    return wrapper, bias
-
-
-class _NoAllReduce(torch.nn.Module):
-    """Stand-in for the cross-rank AllReduce of a single-GPU TP8-shard test;
-    the paths exercised here never reduce.
-    """
-
-    output_buffer_kind = 0
-
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-
-    def forward(self, *args, **kwargs):
-        raise RuntimeError("cross-rank reduction is not available in a single-GPU test")
-
-
-def build_deepseek_moe(layer: int, tp_rank: int, full: Dict[str, torch.Tensor]):
-    """One served MoE (TP8 rank shard): the built-in ``Deepseekv3MoE`` through its own
-    constructor, adopted by the target's ``GlmMoE`` as the target's decoder layer adopts it.
-
-    The NVFP4 experts and the gate carry the real checkpoint tensors; the BF16
-    shared expert gets random weights (it is identical in both arms of a
-    comparison). The cross-rank AllReduce is replaced by a stub, so only the
-    paths that do not reduce (``forward(do_finalize=False)`` and
-    ``compute_routed_output``) are usable. The target's routed fast paths run
-    only inside a decode step (``dispatch.step(dispatch.Phase.DECODE)``).
-    """
-    import tensorrt_llm._torch.distributed as trtllm_dist
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmMoE,
-    )
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import Deepseekv3MoE
-
-    mc = _served_model_config(layer, tp_rank)
-    saved = trtllm_dist.AllReduce
-    trtllm_dist.AllReduce = _NoAllReduce
-    try:
-        moe = Deepseekv3MoE(
-            num_experts=NUM_EXPERTS,
-            top_k=TOP_K,
-            hidden_size=HIDDEN,
-            intermediate_size=INTER,
-            shared_expert_intermediate_size=INTER,
-            aux_stream_dict=_aux_streams(),
-            dtype=torch.bfloat16,
-            model_config=mc,
-            override_quant_config=mc.quant_config,
-            layer_idx=layer,
-        )
-    finally:
-        trtllm_dist.AllReduce = saved
-    GlmMoE.adopt(moe, mc, layer)
-    moe.gate.load_weights(
-        [
-            {
-                "weight": full["gate.weight"],
-                "e_score_correction_bias": full["gate.e_score_correction_bias"],
-            }
-        ]
-    )
-    weights = {k: v for k, v in full.items() if k[0].isdigit()}
-    moe.experts.load_weights([weights])
-    g = torch.Generator().manual_seed(1234)
-    for p in moe.shared_experts.parameters():
-        if p.dtype.is_floating_point:
-            p.data.copy_(torch.randn(p.shape, generator=g) * 0.02)
-    _walk_post_load(moe)
-    return moe.cuda()
-
-
-def load_router_bias(layer: int) -> torch.Tensor:
-    """BF16 ``e_score_correction_bias`` of one MoE layer.
-
-    BF16 is what the TRT-LLM-Gen backend routes with.
-    """
-    return (
-        _Reader()
-        .full(f"model.layers.{layer}.mlp.gate.e_score_correction_bias")
-        .to(torch.bfloat16)
-        .cuda()
-    )
-
-
-def load_embedding_rows(token_ids: List[int]) -> torch.Tensor:
-    """Rows of ``model.embed_tokens.weight`` (BF16, CPU)."""
-    rd = _Reader()
-    name = "model.embed_tokens.weight"
-    sl = rd._handle(name).get_slice(name)
-    return torch.stack([sl[t : t + 1][0] for t in token_ids]).contiguous()
-
-
-def prompt_token_ids(max_tokens: int = 256) -> List[int]:
-    cache = os.path.join(CACHE_DIR, "prompt_ids.json")
-    if os.path.exists(cache):
-        with open(cache) as f:
-            return json.load(f)[:max_tokens]
-    # The checkpoint config does not validate under the image's transformers
-    # AutoConfig, so read the tokenizer file directly.
-    from tokenizers import Tokenizer
-
-    tok = Tokenizer.from_file(os.path.join(CKPT_DIR, "tokenizer.json"))
-    ids = tok.encode(PROMPT_TEXT).ids
-    os.makedirs(CACHE_DIR, exist_ok=True)
-    with open(cache, "w") as f:
-        json.dump(ids, f)
-    return ids[:max_tokens]
-
-
-@dataclass
-class NativeLayer:
-    """A native TRT-LLM-Gen NVFP4 MoE layer plus its router (one TP rank)."""
-
-    module: torch.nn.Module
-    gate_weight: torch.Tensor  # [256, 6144] bf16
-    bias: torch.Tensor  # [256] bf16, the native TRTLLM gate dtype
-    norm_weight: torch.Tensor  # [6144] bf16
-    layer: int
-    tp_rank: int
-
-    def router_logits(self, x: torch.Tensor) -> torch.Tensor:
-        """Native min-latency router GEMM (FP32 logits), as DeepseekV3Gate."""
-        return torch.ops.trtllm.dsv3_router_gemm_op(
-            x, self.gate_weight.t(), bias=None, out_dtype=torch.float32
-        )
-
-    def native_forward(
-        self,
-        x: torch.Tensor,
-        logits: torch.Tensor,
-        do_finalize: bool = True,
-        bias: Optional[torch.Tensor] = None,
-    ):
-        """The native region of the routed_moe contract: FP4 quant + routing/FC1/FC2/finalize.
-
-        Mirrors ``TRTLLMGenFp4BlockScaleBase.run_moe`` and
-        ``TRTLLMGenNvfp4Base.quantize_input`` argument for argument.
-        ``bias`` overrides the router bias (routing studies with another
-        layer's logits and bias).
-        """
-        m = self.module
-        x_fp4, x_sf = m.quantize_input(x)
-        inter_padded = m.w3_w1_weight.shape[-2] // 2
-        outputs = m.op_backend.run_fp4_block_scale_moe(
-            logits,
-            self.bias if bias is None else bias,
-            x_fp4,
-            x_sf,
-            m.w3_w1_weight,
-            m.w3_w1_weight_scale,
-            None,
-            m.act_alpha,
-            m.act_beta,
-            m.act_clamp,
-            m.w2_weight,
-            m.w2_weight_scale,
-            None,
-            m.fc31_scale_c.data,
-            m.fc31_alpha.data,
-            m.fc2_alpha.data,
-            m.num_slots,
-            TOP_K,
-            1,
-            1,
-            inter_padded,
-            m.slot_start,
-            m.expert_size_per_partition,
-            ROUTED_SCALING,
-            int(m.routing_method.routing_method_type),
-            do_finalize=do_finalize,
-            topk_weights=None,
-            topk_ids=None,
-            valid_hidden_size=m.hidden_size,
-            valid_intermediate_size=getattr(
-                m.quant_method, "intermediate_size_per_partition_lean", None
-            ),
-            gated_act_type=0,
-            output=None,
-            tune_max_num_tokens=m.max_num_tokens,
-            use_dp=False,
-        )
-        if do_finalize:
-            return outputs[0] if isinstance(outputs, (list, tuple)) else outputs
-        return outputs
-
-    def hidden_proxy(self, token_ids: List[int]) -> torch.Tensor:
-        """RMSNorm(embedding) * post_attention_layernorm.weight, BF16 on GPU."""
-        emb = load_embedding_rows(token_ids).cuda().float()
-        var = emb.pow(2).mean(-1, keepdim=True)
-        y = emb * torch.rsqrt(var + RMS_EPS)
-        return (y * self.norm_weight.float()).to(torch.bfloat16).contiguous()
-
-
-def build_native_layer(layer: int = 10, tp_rank: int = 0, max_num_tokens: int = 128) -> NativeLayer:
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
-    from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
-    from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import trtllm_gen_leaf
-    from tensorrt_llm.mapping import Mapping
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-    from tensorrt_llm.quantization.mode import QuantAlgo
-
-    raw = load_layer_rank(layer, tp_rank)
-    bias = raw["gate.e_score_correction_bias"].to(torch.bfloat16).cuda()
-    routing = DeepSeekV3MoeRoutingMethod(
-        top_k=TOP_K,
-        n_group=1,
-        topk_group=1,
-        routed_scaling_factor=ROUTED_SCALING,
-        callable_e_score_correction_bias=lambda: bias,
-        is_fused=True,
-    )
-    quant_config = QuantConfig(
-        quant_algo=QuantAlgo.NVFP4, kv_cache_quant_algo=QuantAlgo.FP8, group_size=16
-    )
-    model_config = ModelConfig(
-        quant_config=quant_config,
-        mapping=Mapping(world_size=1, tp_size=1, rank=0),
-        moe_backend="TRTLLM",
-        max_num_tokens=max_num_tokens,
-    )
-    leaf_cls = trtllm_gen_leaf(QuantAlgo.NVFP4)
-    module = leaf_cls(
-        routing_method=routing,
-        num_experts=NUM_EXPERTS,
-        hidden_size=HIDDEN,
-        intermediate_size=INTER // TP_SIZE,
-        dtype=torch.bfloat16,
-        reduce_results=False,
-        model_config=model_config,
-        init_load_balancer=False,
-        weight_loading_mode=MoEWeightLoadingMode.VANILLA,
-    )
-    weights = {k: v for k, v in raw.items() if k[0].isdigit()}
-    module.load_weights([weights])
-    module.post_load_weights()
-    module = module.cuda()
-    return NativeLayer(
-        module=module,
-        gate_weight=raw["gate.weight"].cuda().contiguous(),
-        bias=bias,
-        norm_weight=raw["post_attention_layernorm.weight"].cuda(),
-        layer=layer,
-        tp_rank=tp_rank,
-    )
-
-
-def _sum8(v: torch.Tensor, order: str) -> torch.Tensor:
-    """Sum of the 8 top-k values in a fixed FP32 association order."""
-    c = [v[:, i] for i in range(TOP_K)]
-    if order == "xor_desc":
-        # butterfly with offsets 16, 8, 4, 2, 1 over lanes 0..7 (zeros above)
-        return ((c[0] + c[4]) + (c[2] + c[6])) + ((c[1] + c[5]) + (c[3] + c[7]))
-    if order == "xor_asc":
-        return ((c[0] + c[1]) + (c[2] + c[3])) + ((c[4] + c[5]) + (c[6] + c[7]))
-    if order == "seq":
-        acc = c[0]
-        for x in c[1:]:
-            acc = acc + x
-        return acc
-    raise ValueError(order)
-
-
-ROUTING_SUM_ORDER = "xor_desc"
-
-
-def routing_reference(logits: torch.Tensor, bias_bf16: torch.Tensor, sum_order: str = None):
-    """Bit-faithful PyTorch model of the native SigmoidBias/ScaledSumNormalize.
-
-    Returns ``(ids[R,8] int32 in descending-key order, weights[R,8] bf16)``.
-    Selection key: FP32 ``0.5*tanh(0.5*x)+0.5 + float(bias)``; ties go to the
-    lower expert index (the native 64-bit key packs ``65535-idx`` below the
-    score bits).
-    """
-    s = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5
-    score = s + bias_bf16.float()[None, :]
-    # Stable descending sort keeps the lower index first among equal scores.
-    order = torch.sort(score, dim=-1, descending=True, stable=True).indices
-    ids = order[:, :TOP_K]
-    sel = torch.gather(score, 1, ids)
-    sig = sel - bias_bf16.float()[ids]
-    total = _sum8(sig, sum_order or ROUTING_SUM_ORDER)[:, None]
-    w = sig * ROUTED_SCALING / total
-    return ids.to(torch.int32), w.to(torch.bfloat16)
-
-
-def logits_for_expert_sets(
-    expert_sets: List[List[int]], bias_bf16: torch.Tensor, seed: int = 0
-) -> torch.Tensor:
-    """FP32 router logits whose native top-8 for row ``t`` is ``expert_sets[t]``.
-
-    Chosen experts get logits that make ``sigmoid + bias`` clear every
-    unchosen expert's maximum possible score by a margin; values within the
-    chosen set are distinct so the top-8 order is well defined.
-    """
-    g = torch.Generator().manual_seed(seed)
-    bias = bias_bf16.float().cpu()
-    num_rows = len(expert_sets)
-    logits = torch.full((num_rows, NUM_EXPERTS), -12.0)
-    logits += 0.5 * torch.rand((num_rows, NUM_EXPERTS), generator=g)
-    for t, chosen in enumerate(expert_sets):
-        assert len(chosen) == TOP_K and len(set(chosen)) == TOP_K
-        unchosen = [e for e in range(NUM_EXPERTS) if e not in set(chosen)]
-        ceiling = max(bias[e].item() for e in unchosen) + 0.05
-        for k, e in enumerate(chosen):
-            # need sigmoid(l) + bias[e] > ceiling + small spread
-            target = ceiling - bias[e].item() + 0.02 * (TOP_K - k)
-            target = min(max(target, 1e-3), 1.0 - 1e-6)
-            logit = torch.logit(torch.tensor(target, dtype=torch.float64))
-            logits[t, e] = float(logit) + 0.5
-    return logits.cuda()
-
-
-def expert_sets_with_distinct(
-    num_rows: int, num_distinct: int, seed: int = 0
-) -> Optional[List[List[int]]]:
-    """``num_rows`` top-8 sets whose union has exactly ``num_distinct`` experts."""
-    if num_distinct < TOP_K or num_distinct > TOP_K * num_rows:
-        return None
-    if num_distinct > NUM_EXPERTS:
-        return None
-    g = torch.Generator().manual_seed(seed)
-    pool = torch.randperm(NUM_EXPERTS, generator=g)[:num_distinct].tolist()
-    sets = []
-    # Cover the pool round-robin so every expert is used at least once.
-    cursor = 0
-    for _ in range(num_rows):
-        row = []
-        while len(row) < TOP_K:
-            e = pool[cursor % num_distinct]
-            cursor += 1
-            if e not in row:
-                row.append(e)
-        sets.append(row)
-    used = set(e for row in sets for e in row)
-    assert len(used) == num_distinct, (num_rows, num_distinct, len(used))
-    return sets
-
-
-# ---------------------------------------------------------------------------
-# Independent emulation of the native NVFP4 MoE arithmetic (reference only).
-# ---------------------------------------------------------------------------
-_E2M1 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
-
-
-def e2m1_lut(device) -> torch.Tensor:
-    vals = _E2M1 + [-v for v in _E2M1]
-    return torch.tensor(vals, dtype=torch.float64, device=device)
-
-
-def unpack_fp4(packed: torch.Tensor) -> torch.Tensor:
-    """uint8 [..., K/2] -> float64 [..., K]; element 0 in the low nibble."""
-    lut = e2m1_lut(packed.device)
-    lo = lut[(packed & 0xF).long()]
-    hi = lut[(packed >> 4).long()]
-    return torch.stack([lo, hi], dim=-1).flatten(-2)
-
-
-def unswizzle_sf(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
-    """128x4-interleaved scale bytes -> linear [rows, cols] (uint8)."""
-    flat = sf.contiguous().view(torch.uint8).flatten()
-    r = torch.arange(rows, device=flat.device)[:, None]
-    c = torch.arange(cols, device=flat.device)[None, :]
-    idx = ((r // 128) * (cols // 4) + c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4
-    return flat[idx]
-
-
-def e4m3_bytes_to_float(b: torch.Tensor) -> torch.Tensor:
-    return b.contiguous().view(torch.float8_e4m3fn).to(torch.float64)
-
-
-def quant_e2m1_rn(v: torch.Tensor) -> torch.Tensor:
-    """FP32 -> E2M1 value with round-to-nearest-even and saturation."""
-    a = v.abs().double()
-    grid = torch.tensor(_E2M1, dtype=torch.float64, device=v.device)
-    # nearest grid point; ties go to the even code
-    d = (a[..., None] - grid).abs()
-    best = d.min(-1, keepdim=True).values
-    cand = d == best
-    codes = torch.arange(8, device=v.device)
-    even = cand & (codes % 2 == 0)
-    pick = torch.where(even.any(-1, keepdim=True), even, cand)
-    code = (pick.int() * (8 - codes)).argmax(-1)  # lowest qualifying code
-    q = grid[code]
-    q = torch.where(a >= 6.0, torch.full_like(q, 6.0), q)
-    return torch.copysign(q, v.double())
-
-
-def quant_nvfp4_unit(h: torch.Tensor):
-    """NVFP4 quantization with unit global scale, block 16 along the last dim.
-
-    Returns the dequantized values (q * sf) as float64.
-    """
-    shp = h.shape
-    hb = h.reshape(*shp[:-1], shp[-1] // 16, 16).float()
-    amax = hb.abs().amax(-1, keepdim=True)
-    sf = (amax * (1.0 / 6.0)).to(torch.float8_e4m3fn).to(torch.float32)
-    osc = torch.where(amax > 0, 1.0 / sf, torch.zeros_like(sf))
-    q = quant_e2m1_rn(hb * osc)
-    return (q * sf.double()).reshape(shp)
-
-
-def stored_row_maps():
-    """Stored-row -> (is_gate, j) for w3_w1 and stored-row -> hidden for w2."""
-    from tensorrt_llm.quantization.utils.fp4_utils import (
-        get_reorder_rows_for_gated_act_gemm_row_indices,
-        get_shuffle_matrix_a_row_indices,
-    )
-
-    dummy13 = torch.empty(2 * (INTER // TP_SIZE), 1)
-    p0 = get_reorder_rows_for_gated_act_gemm_row_indices(dummy13)
-    p1 = get_shuffle_matrix_a_row_indices(dummy13, epilogue_tile_m=128)
-    orig13 = p0[p1]  # stored row -> row of cat([w3, w1])
-    inter = INTER // TP_SIZE
-    is_gate = orig13 >= inter
-    j13 = torch.where(is_gate, orig13 - inter, orig13)
-    dummy2 = torch.empty(HIDDEN, 1)
-    orig2 = get_shuffle_matrix_a_row_indices(dummy2, epilogue_tile_m=128)
-    return is_gate, j13, orig2
-
-
-def dequant_expert(module, e: int):
-    """Stored-row dequantized FC1/FC2 weights of expert ``e`` (no global)."""
-    w13 = unpack_fp4(module.w3_w1_weight[e])
-    sf13 = e4m3_bytes_to_float(
-        unswizzle_sf(module.w3_w1_weight_scale[e], 2 * (INTER // TP_SIZE), HIDDEN // 16)
-    )
-    w13 = w13 * sf13.repeat_interleave(16, dim=1)
-    w2 = unpack_fp4(module.w2_weight[e])
-    sf2 = e4m3_bytes_to_float(
-        unswizzle_sf(module.w2_weight_scale[e], HIDDEN, (INTER // TP_SIZE) // 16)
-    )
-    w2 = w2 * sf2.repeat_interleave(16, dim=1)
-    return w13, w2
-
-
-def _device_row_maps(device):
-    is_gate, j13, orig2 = stored_row_maps()
-    return is_gate.to(device), j13.to(device), orig2.to(device)
-
-
-def emulate_expert_row(module, e: int, xd_row: torch.Tensor, maps, cache=None) -> torch.Tensor:
-    """FC2 output row (BF16 values as float32) of expert ``e`` for one
-    dequantized activation row, with FP64 accumulation and the native
-    SwiGLU / NVFP4 / scaling steps.
-    """
-    is_gate, j13, orig2 = maps
-    inter = INTER // TP_SIZE
-    if cache is not None and e in cache:
-        w13, w2 = cache[e]
-    else:
-        w13, w2 = dequant_expert(module, e)
-        if cache is not None:
-            cache[e] = (w13, w2)
-    a13 = w13 @ xd_row  # stored-row accumulators
-    up = torch.zeros(inter, dtype=torch.float64, device=xd_row.device)
-    gate = torch.zeros(inter, dtype=torch.float64, device=xd_row.device)
-    up[j13[~is_gate]] = a13[~is_gate]
-    gate[j13[is_gate]] = a13[is_gate]
-    gr = (gate * float(module.fc31_alpha[e])).float()
-    us = (up * float(module.fc31_scale_c[e])).float()
-    h = us * (gr * torch.sigmoid(gr))
-    hd = quant_nvfp4_unit(h[None, :])[0]
-    a2 = w2 @ hd  # stored rows
-    y = torch.zeros(HIDDEN, dtype=torch.float64, device=xd_row.device)
-    y[orig2] = a2
-    return (y * float(module.fc2_alpha[e])).to(torch.bfloat16).float()
-
-
-def dequant_activations(module, x: torch.Tensor) -> torch.Tensor:
-    """The stock ``fp4_quantize`` of ``x``, dequantized to float64."""
-    rows = x.shape[0]
-    xq, xsf = module.quantize_input(x)
-    return unpack_fp4(xq.view(torch.uint8)) * e4m3_bytes_to_float(
-        xsf.view(torch.uint8).reshape(rows, HIDDEN // 16)
-    ).repeat_interleave(16, dim=1)
-
-
-def emulate_moe(layer: "NativeLayer", x: torch.Tensor, logits: torch.Tensor):
-    """FP64-accumulated emulation of the native region (finalized output).
-
-    Activations are quantized with the stock ``fp4_quantize`` op; weights are
-    the stock module buffers; routing is ``routing_reference``.
-    """
-    m = layer.module
-    rows = x.shape[0]
-    xd = dequant_activations(m, x)
-    ids, wts = routing_reference(logits, layer.bias)
-    maps = _device_row_maps(x.device)
-    out = torch.zeros(rows, HIDDEN, dtype=torch.float32, device=x.device)
-    cache = {}
-    for t in range(rows):
-        acc = torch.zeros(HIDDEN, dtype=torch.float32, device=x.device)
-        for k in range(TOP_K):
-            y = emulate_expert_row(m, int(ids[t, k]), xd[t], maps, cache)
-            acc = acc + wts[t, k].float() * y
-        out[t] = acc
-    return out.to(torch.bfloat16)
-
-
-def error_stats(a: torch.Tensor, b: torch.Tensor) -> dict:
-    a = a.float()
-    b = b.float()
-    diff = a - b
-    ref_rms = b.pow(2).mean().sqrt().item()
-    return {
-        "nrms": (diff.pow(2).mean().sqrt().item() / max(ref_rms, 1e-30)),
-        "max_abs": diff.abs().max().item(),
-        "ref_rms": ref_rms,
-        "cos": torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item(),
-    }
-
-
-# ---------------------------------------------------------------------------
-# Native capture of real routed-MoE inputs (read-only shared fixture).
-# ---------------------------------------------------------------------------
-REAL_DIR = os.environ.get("GLM_MOE_DSA_REAL_MOE_INPUTS_DIR", "")
-
-
-def load_real_capture(layer: int, real_dir: str = REAL_DIR):
-    """Per prompt, ``(hidden [T, 6144] bf16 or None, logits [T, 256] fp32)``
-    that the native ``compute_routed_output`` of ``layer`` received (CPU).
-    Hidden states were captured for layers 10, 30, 50 and 77 only.
-    """
-    import glob
-
-    out = []
-    for path in sorted(glob.glob(os.path.join(real_dir, f"rows*_layer{layer}.pt"))):
-        obj = torch.load(path, map_location="cpu")
-        out.append((obj.get("hidden"), obj["logits"]))
-    return out
-
-
-def real_windows(capture, rows: int, stride: Optional[int] = None):
-    """Consecutive-token windows ``(x, logits)`` (CPU); ``x`` is None when the
-    capture has no hidden states for this layer.
-    """
-    stride = stride or rows
-    win = []
-    for hidden, logits in capture:
-        for start in range(0, logits.shape[0] - rows + 1, stride):
-            x = None if hidden is None else hidden[start : start + rows]
-            win.append((x, logits[start : start + rows]))
-    return win
-
-
-def routing_keys(logits: torch.Tensor, bias_bf16: torch.Tensor) -> torch.Tensor:
-    """FP32 selection keys ``sigmoid + bias`` of the native routing."""
-    s = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5
-    return s + bias_bf16.float()[None, :].to(s.device)
-
-
-def boundary_gap(logits: torch.Tensor, bias_bf16: torch.Tensor) -> torch.Tensor:
-    """Per row, key(8th) - key(9th) in FP32 (0 means an exact boundary tie)."""
-    k = routing_keys(logits, bias_bf16)
-    top = torch.topk(k, TOP_K + 1, dim=-1).values
-    return top[:, TOP_K - 1] - top[:, TOP_K]
-
-
-def _key_of(logit: torch.Tensor, bias_val: torch.Tensor) -> torch.Tensor:
-    s = 0.5 * torch.tanh(0.5 * logit.float()) + 0.5
-    return s + bias_val.float()
-
-
-def logits_for_key(target_key: float, bias_val: torch.Tensor, offset_ulps: int = 0):
-    """All FP32 logits within +-2^16 logit ulps of the analytic solution whose
-    native key equals ``target_key`` shifted by ``offset_ulps`` key ulps (the
-    key arithmetic is FP32, evaluated on the device). May be empty.
-    """
-    dev = bias_val.device
-    want = torch.tensor(target_key, dtype=torch.float32, device=dev)
-    toward = torch.tensor(float("inf") if offset_ulps > 0 else -float("inf"), device=dev)
-    for _ in range(abs(offset_ulps)):
-        want = torch.nextafter(want, toward)
-    s_target = float(want.double() - bias_val.double())
-    s_target = min(max(s_target, 1e-4), 1 - 1e-4)
-    l0 = torch.logit(torch.tensor(s_target, dtype=torch.float64)).float().reshape(1).to(dev)
-    steps = torch.arange(-(1 << 16), 1 << 16, device=dev, dtype=torch.int32)
-    cand = (l0.view(torch.int32) + steps).view(torch.float32)
-    keys = _key_of(cand, bias_val)
-    return cand[keys == want]
-
-
-def logit_for_key(target_key: float, bias_val: torch.Tensor, offset_ulps: int = 0, avoid=None):
-    """One logit from ``logits_for_key`` (the one farthest from ``avoid`` when
-    given, else the middle hit) or ``None``.
-    """
-    hits = logits_for_key(target_key, bias_val, offset_ulps)
-    if hits.numel() == 0:
-        return None
-    if avoid is None:
-        return hits[hits.numel() // 2].item()
-    return hits[(hits - avoid).abs().argmax()].item()
-
-
-def raw_rows_by_topk(raw) -> torch.Tensor:
-    """``[R, 8, 6144]`` FC2 rows in top-k order from a raw tuple
-    ``(fc2_output, scales, expanded_idx_to_permuted_idx)``.
-    """
-    fc2, _, idx = raw
-    return fc2[idx.long()]
-
-
-def finalize_raw(raw) -> torch.Tensor:
-    """The stock finalize arithmetic: ``sum_k scale[t,k] * fc2[idx[t,k]]`` in
-    FP32, top-k order, rounded to BF16.
-    """
-    fc2, scale, idx = raw
-    rows = scale.shape[0]
-    fin = torch.zeros(rows, fc2.shape[1], device=fc2.device, dtype=torch.float32)
-    for k in range(TOP_K):
-        fin += scale[:, k : k + 1].float() * fc2[idx[:, k].long()].float()
-    return fin.to(torch.bfloat16)
-
-
-# ---------------------------------------------------------------------------
-# routed_moe-versus-native classification.
-#
-# Discrete invariants are exact: the expert scales of every (t, k) are equal
-# bit for bit and every (t, k) uses the same expert as native. Expert identity
-# is read from the FC2 rows: rows within 5 % relative L2 of native's row are
-# the same expert (rows of two different experts are far apart); any other
-# window is decoded exactly by ``ExpertDecoder`` (a probe activation shared by
-# all tokens turns every FC2 row into a per-expert fingerprint).
-# Floating-point results fall in one of three classes:
-#   exact    finalized outputs bitwise equal;
-#   reorder  every FC2 row bitwise equal to native's and the finalized output
-#            within the bound of the same eight exact products summed in a
-#            different FP32 order (routed_moe sums per expert slot and across split
-#            partials, the native finalize in top-k order):
-#            |a - b| <= 2^-20 * sum_k |w_k y_k| + ulp_bf16(max(|a|, |b|));
-#   budget   some FC2 rows differ (an FP4 rounding of the intermediate flipped
-#            because the split-K FC1 accumulation is associated differently);
-#            checked against the frozen error budget
-#            (nrms <= 1.5 * nrms(emul - native) + 2e-3, no row above 4x).
-# ---------------------------------------------------------------------------
-SAME_EXPERT_REL = 0.05
-
-
-def bf16_ulp(v: torch.Tensor) -> torch.Tensor:
-    _, e = torch.frexp(v.float().abs())
-    return torch.ldexp(torch.ones_like(v, dtype=torch.float32), e - 8)
-
-
-def raw_report(raw_a, raw_n) -> dict:
-    ra = raw_rows_by_topk(raw_a).float()
-    rn = raw_rows_by_topk(raw_n).float()
-    num = (ra - rn).pow(2).sum(-1).sqrt()
-    den = rn.pow(2).sum(-1).sqrt()
-    rel = torch.where(den > 0, num / den.clamp_min(1e-30), (num > 0).float() * 1e30)
-    rows_exact = (ra == rn).all(-1)
-    return {
-        "scale_exact": bool(torch.equal(raw_a[1], raw_n[1])),
-        "same_expert": bool((rel < SAME_EXPERT_REL).all()),
-        "rows_exact": int(rows_exact.sum()),
-        "rows": rows_exact.numel(),
-        "max_row_rel": float(rel.max()),
-    }
-
-
-def reorder_close(fin_a: torch.Tensor, fin_n: torch.Tensor, raw_n) -> bool:
-    fc2, scale, idx = raw_n
-    terms = (scale.float()[:, :, None] * fc2[idx.long()].float()).abs().sum(1)
-    a = fin_a.float()
-    b = fin_n.float()
-    bound = terms * 2.0**-20 + bf16_ulp(torch.maximum(a.abs(), b.abs()))
-    return bool(((a - b).abs() <= bound).all())
-
-
-def budget_check(out: torch.Tensor, native: torch.Tensor, emul: torch.Tensor):
-    """The frozen error budget; returns ``(ok, stats)``."""
-    e_native = error_stats(out, native)
-    e_emul = error_stats(emul, native)
-    budget = 1.5 * e_emul["nrms"] + 2e-3
-    row_ref = native.float().pow(2).mean(-1).sqrt()
-    row_err = (out.float() - native.float()).pow(2).mean(-1).sqrt()
-    ok = e_native["nrms"] <= budget and bool(torch.all(row_err <= 4 * budget * row_ref + 1e-6))
-    return ok, {"routed_moe_vs_native": e_native, "emul_vs_native": e_emul, "budget": budget}
-
-
-class ExpertDecoder:
-    """Exact expert IDs of a raw tuple, routed_moe's or native's.
-
-    With every token row equal to one probe activation, the FC2 row of
-    ``(t, k)`` depends only on the expert, so it is matched (nearest
-    neighbour, L2) against per-expert fingerprints: native FC2 rows of every
-    local expert for the probe, produced with constructed logits whose
-    expert of every ``(t, k)`` is known. Routing comes from the logits and
-    bias only, so a decoded call uses the window's logits with the probe as
-    activation.
-    """
-
-    def __init__(self, layer: "NativeLayer", probe: torch.Tensor):
-        self.layer = layer
-        self.probe = probe.reshape(1, HIDDEN).to(torch.bfloat16).cuda()
-        # Fingerprints: every expert's FC2 row for the probe, computed by the
-        # FP64 emulation of the native arithmetic from the native buffers
-        # (independent of any routing). A native or routed_moe row of the same
-        # expert differs from it at most by an FP4 rounding of the
-        # intermediate; rows of different experts are far apart.
-        m = layer.module
-        xd = dequant_activations(m, self.probe)[0]
-        maps = _device_row_maps(xd.device)
-        self.fps = torch.stack(
-            [emulate_expert_row(m, e, xd, maps) for e in range(NUM_EXPERTS)]
-        ).float()
-        norms = self.fps.norm(dim=-1)
-        d = torch.cdist(self.fps, self.fps) + torch.eye(NUM_EXPERTS, device="cuda") * 1e30
-        self.min_norm_ratio = float(norms.min() / norms.median())
-        self.min_separation = float(d.min() / norms.median())
-        assert self.min_norm_ratio > 0.05 and self.min_separation > 0.2, (
-            self.min_norm_ratio,
-            self.min_separation,
-        )
-
-    def _x(self, rows: int) -> torch.Tensor:
-        return self.probe.expand(rows, HIDDEN).contiguous()
-
-    def decode(self, raw):
-        rows = raw_rows_by_topk(raw).float()
-        flat = rows.reshape(-1, HIDDEN)
-        d = torch.cdist(flat, self.fps)
-        best = d.topk(2, dim=-1, largest=False)
-        margin = best.values[:, 1] / best.values[:, 0].clamp_min(1e-30)
-        return best.indices[:, 0].reshape(rows.shape[:2]), float(margin.min())
-
-    def compare(self, runner, logits, bias=None) -> dict:
-        """Decode native's and routed_moe's experts for ``logits`` (and ``bias``)."""
-        x = self._x(logits.shape[0])
-        n_raw = self.layer.native_forward(x, logits, do_finalize=False, bias=bias)
-        w_raw = runner(x, logits, do_finalize=False)
-        n_ids, n_margin = self.decode(n_raw)
-        w_ids, w_margin = self.decode(w_raw)
-        return {
-            "ids_equal": bool(torch.equal(n_ids, w_ids)),
-            "scale_exact": bool(torch.equal(n_raw[1], w_raw[1])),
-            "margin": min(n_margin, w_margin),
-            "native_ids": n_ids,
-            "routed_moe_ids": w_ids,
-        }
-
-
-def classify_vs_native(w_fin, w_raw, n_fin, n_raw, emul_fn=None, label="", id_check=None):
-    """Assert the discrete invariants and return ``(class, report)``.
-
-    ``emul_fn`` computes the FP64 emulation for the budget class; without it
-    a window of that class is returned as ``"budget-unchecked"``.
-    ``id_check`` (e.g. ``lambda: decoder.compare(runner, logits)``) decodes
-    the experts exactly when some FC2 row is not within 5 % of native's.
-    """
-    rep = raw_report(w_raw, n_raw)
-    assert rep["scale_exact"], f"{label}: expert scales differ from native"
-    if not rep["same_expert"]:
-        assert id_check is not None, f"{label}: FC2 rows far from native and no id check {rep}"
-        dec = id_check()
-        rep["decoded"] = {k: dec[k] for k in ("ids_equal", "scale_exact", "margin")}
-        assert dec["ids_equal"] and dec["scale_exact"] and dec["margin"] > 4, (
-            f"{label}: decoded experts differ from native {rep}"
-        )
-    if torch.equal(w_fin, n_fin):
-        return "exact", rep
-    if rep["rows_exact"] == rep["rows"]:
-        assert reorder_close(w_fin, n_fin, n_raw), f"{label}: beyond the FP32 reorder bound"
-        return "reorder", rep
-    if emul_fn is None:
-        return "budget-unchecked", rep
-    ok, stats = budget_check(w_fin, n_fin, emul_fn())
-    rep["budget"] = stats
-    assert ok, f"{label}: outside the frozen error budget {stats}"
-    return "budget", rep
-
-
-def check_call(layer, runner, x, logits, fin=None, raw=None, emul=True, label="", decoder=None):
-    """Run routed_moe (unless outputs are given) and native on one input; classify."""
-    if fin is None:
-        fin = runner(x, logits)
-    if raw is None:
-        raw = runner(x, logits, do_finalize=False)
-    n_fin = layer.native_forward(x, logits)
-    n_raw = layer.native_forward(x, logits, do_finalize=False)
-    emul_fn = (lambda: emulate_moe(layer, x, logits)) if emul else None
-    id_check = (lambda: decoder.compare(runner, logits)) if decoder is not None else None
-    return classify_vs_native(fin, raw, n_fin, n_raw, emul_fn, label, id_check)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py
deleted file mode 100644
index 6f773bad1aa6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_into_hop_mpi.py
+++ /dev/null
@@ -1,464 +0,0 @@
-# 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.
-"""routed_moe_into_hop on the two-host TP8 group.
-
-routed_moe writes its finalized routed output into the post_moe_hop, which adds the shared
-output, reduces across the group and normalizes.
-
-Run every test on all 8 ranks (``srun --mpi=pmix``, two hosts). The MoE is the built-in
-``Deepseekv3MoE`` adopted by the target's ``GlmMoE``, and every test runs as a decode step.
-Real GLM-5.2 weights (layers 10 and 77) and captured real hidden-state windows.
-
-Oracles, per call (identical inputs on every arm):
-
-* ``raw``: ``Deepseekv3MoE.forward(do_finalize=False)`` with routed_moe (raw tuple)
-  -> post_moe_hop on the raw tuple, i.e. the composition without routed_moe_into_hop;
-* ``native``: native MoE finalized -> native MNNVL ``AllReduce`` with
-  RESIDUAL_RMS_NORM.
-
-routed_moe_into_hop must equal ``raw`` bit for bit (same routed values, same post_moe_hop arithmetic
-after the finalize) and match ``native`` like the composition does.
-
-The hop is built, and its routed input enabled, inside the served model
-loader's ``MetaInitMode``, as the target's decoder-layer adoption does in the
-server (``routed_into_hop_checks.served_construction``).
-
-Tests of the routed-input variant and the routed_moe finalize -> post_moe_hop
-hand-off (helpers in ``routed_into_hop_checks.py``):
-
-* ``test_routed_into_hop_served_construction_buffers``: per mode, a hop built inside the
-  served loader's ``MetaInitMode`` has its routed-input buffers in device
-  memory, and its routed_moe_into_hop calls equal the raw variant;
-* ``test_routed_into_hop_shared_absent_variant``: the ``has_shared=False`` routed variant;
-* ``test_routed_into_hop_mixed_rank_variants``: some ranks keep the raw tuple (routed_moe raw or
-  the native runner) while the others use routed_moe_into_hop, on the shared post_moe_hop workspace;
-  no hang, bit-equal to every rank on the raw variant, published partials
-  and outputs equal to the native-arithmetic emulation;
-* ``test_routed_into_hop_mixed_rank_graph``: the same mix inside a multi-stream CUDA
-  graph, replayed on fresh windows;
-* ``test_routed_into_hop_delayed_stress``: back-to-back batches with random per-rank
-  producers, shared output and device delays, plus delay injection inside
-  the routed_moe kernels (``routed_moe.DELAY_INJECTION``: CTA entry of the
-  routed kernel and the finalize kernel's stores).
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(HERE), "routed_moe"))
-sys.path.insert(
-    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
-)
-sys.path.insert(0, HERE)
-import post_moe_hop_mpi_utils as U  # noqa: E402
-import routed_into_hop_checks as G  # noqa: E402
-import routed_moe_chain_bench as CB  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-LAYERS = [
-    int(v) for v in os.environ.get("GLM_MOE_DSA_ROUTED_MOE_INTO_HOP_LAYERS", "10,77").split(",")
-]
-
-
-@pytest.fixture(autouse=True)
-def decode_step():
-    """Every test runs as a decode step, where the target's MoE takes its routed fast paths."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    with dispatch.step(dispatch.Phase.DECODE):
-        yield
-
-
-@pytest.fixture(scope="module")
-def env():
-    e = U.mpi_env()
-    if e.world != 8 or not e.multi_node:
-        pytest.skip("needs the two-host TP8 group")
-    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
-    return e
-
-
-@pytest.fixture(scope="module")
-def stack(env):
-    # as the served model builds it: inside the loader's MetaInitMode, with the
-    # layer's MoE module declaring the raw-tuple contract the hop consumes
-    moes, gammas, hop = G.build_stack(env, LAYERS)
-    eng = CB.Engagement()
-    return moes, gammas, hop, eng
-
-
-@pytest.mark.parametrize("mode", ["row", "identity", "fp32"])
-def test_routed_into_hop_served_construction_buffers(env, stack, mode):
-    """Every mode's routed-input buffers are device memory when the hop is
-    built inside the served loader's ``MetaInitMode``, and a routed_moe_into_hop call in that
-    mode equals the raw variant bit for bit (modes "row", "identity") or within
-    the FP32 reorder bound ("fp32")."""
-    moes, gammas, _, _ = stack
-    with G.served_construction():
-        probe = torch.empty(4, device="cuda")
-    assert probe.is_meta, "the harness does not reproduce the served loader's MetaInitMode"
-    layer = LAYERS[0]
-    hop, enabled = G.build_hop(env, moes[layer].experts, mode)
-    assert hop.engaged and enabled and hop.routed_mode == mode
-    problems = G.buffer_problems(hop)
-    assert U.all_ranks(not problems), f"rank {env.rank}: {problems}"
-    saved = {m: m._routed_into_hop for m in moes.values()}
-    try:
-        for m in moes.values():
-            m._routed_into_hop = hop
-        G.warm(env, {layer: moes[layer]}, [6], producers=("routed_moe_into_hop", "routed_moe"))
-        ok, n, fails = G.run_mixed(
-            env, {layer: moes[layer]}, gammas, hop, "all-routed_moe_into_hop", True, [6], 2
-        )
-        hop.raise_if_error()
-    finally:
-        for m, h in saved.items():
-            m._routed_into_hop = h
-        CB.set_features("")
-    if env.rank == 0:
-        per_hop = G.buffer_bytes(hop)
-        print(
-            f"[routed_moe_into_hop] served construction mode={mode}: buffers {G.buffer_report(hop)}, {n} calls; "
-            f"{per_hop} bytes per hop, {per_hop * G.MOE_LAYERS / 2**20:.1f} MiB per rank over "
-            f"{G.MOE_LAYERS} MoE layers"
-        )
-    assert U.all_ranks(ok), f"rank {env.rank}: {fails}"
-
-
-@pytest.mark.parametrize("bad_rank", [0, 5])
-def test_routed_into_hop_buffer_problem_turns_routed_into_hop_off_on_every_rank(
-    env, stack, monkeypatch, bad_rank
-):
-    """The buffer check's failure path: if one rank's routed-input buffer is not
-    device memory (as an allocation under ``MetaInitMode`` leaves it),
-    ``enable_routed_input`` returns False on every rank and the MoE keeps the
-    raw tuple, before any device work touches the buffer."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        hop as hop_mod,
-    )
-
-    moes, gammas, _, _ = stack
-    layer = LAYERS[0]
-    if env.rank == bad_rank:
-        monkeypatch.setattr(
-            hop_mod.PostMoeHop,
-            "_routed_device_buffer",
-            lambda self, shape, dtype: torch.empty(shape, dtype=dtype, device="meta"),
-        )
-    hop, enabled = G.build_hop(env, moes[layer].experts, "identity")
-    monkeypatch.undo()
-    assert hop.engaged
-    assert U.all_ranks(not enabled and not hop._routed_enabled), (
-        "routed_moe_into_hop stayed on somewhere"
-    )
-    moe = moes[layer]
-    saved = moe._routed_into_hop
-    try:
-        moe._routed_into_hop = hop
-        CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
-        out = moe.forward(_windows(layer, 6, 1)[0], do_finalize=False)
-        torch.cuda.synchronize()
-        assert len(out) == 4, (
-            "routed_moe_into_hop must keep the raw tuple when its hop declined routed input"
-        )
-    finally:
-        moe._routed_into_hop = saved
-        CB.set_features("")
-    if env.rank == 0:
-        print(
-            f"[routed_moe_into_hop] buffer problem on rank {bad_rank}: routed_moe_into_hop off on all {env.world} ranks"
-        )
-
-
-@pytest.mark.parametrize("odd_rank", [0, 6])
-def test_routed_into_hop_modes_differing_across_ranks_turn_routed_into_hop_off_everywhere(
-    env, stack, monkeypatch, odd_rank
-):
-    """``TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION`` is read per rank.
-
-    It comes from each rank's environment. If one rank selects another mode
-    than its peers, ``enable_routed_input`` agrees on that and turns routed_moe_into_hop off on
-    every rank (the MoE keeps the raw tuple), instead of enabling a different
-    routed-input variant per rank.
-    """
-    moes, _, _, _ = stack
-    layer = LAYERS[0]
-    monkeypatch.setenv(
-        "TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION",
-        "identity" if env.rank == odd_rank else "fp32",
-    )
-    hop, enabled = G.build_hop(env, moes[layer].experts)
-    monkeypatch.delenv("TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION")
-    assert hop.engaged
-    assert U.all_ranks(not enabled and not hop._routed_enabled), (
-        "routed_moe_into_hop stayed on somewhere"
-    )
-    # The same environment on every rank, spelled as the selector accepts it,
-    # enables routed_moe_into_hop everywhere with the mode the feature record names.
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    monkeypatch.setenv("TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION", " Identity ")
-    hop2, enabled2 = G.build_hop(env, moes[layer].experts)
-    record = dispatch.feature_record({"post_moe_hop", "routed_moe", "routed_moe_into_hop"})
-    monkeypatch.delenv("TRTLLM_STAIRCASE_GLM_LL_ROUTED_MOE_INTO_HOP_ACCUMULATION")
-    assert U.all_ranks(enabled2 and hop2.routed_mode == "identity")
-    assert "routed_moe_into_hop[identity]" in record
-    if env.rank == 0:
-        print(
-            f"[routed_moe_into_hop] mode differs on rank {odd_rank}: routed_moe_into_hop off on all {env.world} ranks"
-        )
-
-
-def _windows(layer: int, rows: int, count: int):
-    win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer), rows)]
-    step = max(1, len(win) // count)
-    return [win[i * step].cuda() for i in range(count)]
-
-
-def _arm(moe, hop, x, res, gamma, features: str):
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    CB.set_features(features)
-    if features == "":
-        h = moe.forward(
-            x, do_finalize=True, final_all_reduce_params=AllReduceParams(enable_allreduce=False)
-        )
-        return U.native_fused_allreduce(h, res, gamma, fx.RMS_EPS), None
-    out = moe.forward(x, do_finalize=False)
-    if len(out) == 2:
-        shared, routed = out
-        return hop.hop_routed(routed, shared, res, gamma, fx.RMS_EPS), routed
-    shared, fc2, scale, perm = out
-    return hop(fc2, scale, perm, shared, res, gamma, fx.RMS_EPS), None
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_routed_into_hop_equals_composition_and_matches_native(env, stack, rows):
-    moes, gammas, hop, eng = stack
-    mode = hop.routed_mode
-    buffers = {
-        "row": lambda: hop.routed_input_buffer(rows),
-        "identity": lambda: hop.routed_raw_buffers(rows)[0],
-        "fp32": lambda: hop.routed_acc_buffer(rows),
-    }
-    g = torch.Generator(device="cuda").manual_seed(100 + rows)
-    worst_native = [0, 0.0]
-    exact_native = 0
-    total = 0
-    reordered = 0
-    for layer in LAYERS:
-        for x in _windows(layer, rows, 6):
-            res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
-            e0 = eng.engaged("routed_output_into_hop")
-            CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
-            out = moes[layer].forward(x, do_finalize=False)
-            assert eng.engaged("routed_output_into_hop") == e0 + 1, (
-                "routed_moe_into_hop did not engage"
-            )
-            assert len(out) == 2 and out[1].data_ptr() == buffers[mode]().data_ptr()
-            # the hop clears the FP32 accumulator after reading: keep a copy
-            seen = [out[0], out[1].clone() if mode == "fp32" else out[1]]
-            w_out, w_res = G.hop_call(hop, "routed", out, res, gammas[layer], True)
-            mine = G.local_partial("routed", seen, True, hop=hop)
-            (r_out, r_res), _, r_payload = G.call(
-                moes[layer], hop, x, res, gammas[layer], "routed_moe", True
-            )
-            (n_out, n_res), _ = _arm(moes[layer], hop, x, res, gammas[layer], "")
-            torch.cuda.synchronize()
-            if G.exact_mode(hop):
-                assert torch.equal(w_res, r_res), (
-                    f"L{layer} R={rows}: residual differs from routed_moe raw -> post_moe_hop"
-                )
-                assert torch.equal(w_out, r_out), (
-                    f"L{layer} R={rows}: output differs from routed_moe raw -> post_moe_hop"
-                )
-            else:
-                ok, n_diff = G.reorder_ok(mine, r_payload, True)
-                assert ok, f"L{layer} R={rows}: partial outside the FP32 reorder bound"
-                reordered += n_diff
-            st = U.compare_stats(w_res, n_res)
-            exact_native += int(st["max_ulp"] == 0)
-            worst_native[0] = max(worst_native[0], st["max_ulp"])
-            worst_native[1] = max(worst_native[1], st["max_abs"])
-            torch.testing.assert_close(w_out, n_out, rtol=0.05, atol=0.15)
-            total += 1
-    CB.set_features("")
-    allw = env.comm.allgather((exact_native, total, worst_native, reordered))
-    if env.rank == 0:
-        rel = "bit for bit" if G.exact_mode(hop) else "within the FP32 reorder bound"
-        print(
-            f"[routed_moe_into_hop] mode={mode} R={rows}: routed_moe_into_hop == composition (routed_moe "
-            f"raw -> post_moe_hop) {rel} on all ranks; "
-            f"residual vs native exact in {[a[0] for a in allw]} of {total} windows per rank, "
-            f"worst ulp/abs {[a[2] for a in allw]}, reordered partial elements {[a[3] for a in allw]}"
-        )
-
-
-def test_routed_into_hop_cuda_graph_replay(env, stack):
-    moes, gammas, hop, eng = stack
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    rows = 6
-    layer = LAYERS[0]
-    wins = _windows(layer, rows, 5)
-    x = wins[0].clone()
-    g = torch.Generator(device="cuda").manual_seed(7)
-    res = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 4.0).bfloat16()
-    CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(2):
-            _arm(
-                moes[layer],
-                hop,
-                x,
-                res,
-                gammas[layer],
-                "post_moe_hop,routed_moe,routed_moe_into_hop",
-            )
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    graph = torch.cuda.CUDAGraph()
-    with with_multi_stream(True):
-        with torch.cuda.graph(graph, stream=stream):
-            (g_out, g_res), _ = _arm(
-                moes[layer],
-                hop,
-                x,
-                res,
-                gammas[layer],
-                "post_moe_hop,routed_moe,routed_moe_into_hop",
-            )
-    torch.cuda.synchronize()
-    for w in wins[1:]:
-        x.copy_(w)
-        env.comm.barrier()
-        graph.replay()
-        torch.cuda.synchronize()
-        (e_out, e_res), _ = _arm(
-            moes[layer], hop, x, res, gammas[layer], "post_moe_hop,routed_moe,routed_moe_into_hop"
-        )
-        torch.cuda.synchronize()
-        assert G.same_or_close(hop, (g_out, g_res), (e_out, e_res)), "graph replay differs"
-    CB.set_features("")
-    hop.raise_if_error()
-
-
-def test_routed_into_hop_falls_back_to_raw_without_routed_input(env, stack):
-    moes, gammas, hop, eng = stack
-    layer = LAYERS[0]
-    x = _windows(layer, 6, 1)[0]
-    moe = moes[layer]
-    saved = moe._routed_into_hop
-    try:
-        moe._routed_into_hop = None
-        CB.set_features("post_moe_hop,routed_moe,routed_moe_into_hop")
-        out = moe.forward(x, do_finalize=False)
-        assert len(out) == 4, (
-            "without the hop's routed input routed_moe_into_hop must keep the raw tuple"
-        )
-    finally:
-        moe._routed_into_hop = saved
-        CB.set_features("")
-
-
-_WARM = set()
-
-
-def _warm(env, moes, rows_list):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe as routed_moe_mod,
-    )
-
-    key = (tuple(rows_list), routed_moe_mod.DELAY_INJECTION)
-    if key not in _WARM:
-        G.warm(env, moes, rows_list)
-        _WARM.add(key)
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_routed_into_hop_shared_absent_variant(env, stack, rows):
-    moes, gammas, hop, eng = stack
-    _warm(env, moes, [1, 6, 8])
-    ok, n, fails = G.run_mixed(
-        env, moes, gammas, hop, "all-routed_moe_into_hop", False, [rows], 3, seed=rows
-    )
-    hop.raise_if_error()
-    if env.rank == 0:
-        print(
-            f"[routed_moe_into_hop] shared absent R={rows}: {n} calls, routed variant == raw variant"
-        )
-    assert U.all_ranks(ok), f"rank {env.rank}: {fails}"
-
-
-@pytest.mark.parametrize("with_shared", [True, False], ids=["shared", "no_shared"])
-@pytest.mark.parametrize(
-    "pattern", ["odd-routed_moe", "rank0-routed_moe", "host1-routed_moe", "odd-native"]
-)
-def test_routed_into_hop_mixed_rank_variants(env, stack, pattern, with_shared):
-    moes, gammas, hop, eng = stack
-    _warm(env, moes, [1, 6, 8])
-    ok, n, fails = G.run_mixed(env, moes, gammas, hop, pattern, with_shared, [1, 6, 8], 2)
-    hop.raise_if_error()
-    producers = env.comm.allgather(G.mixed_pattern(env, pattern))
-    if env.rank == 0:
-        print(
-            f"[routed_moe_into_hop] mixed {pattern} shared={with_shared}: producers {producers}, {n} calls "
-            f"bit-equal to the all-raw variant"
-        )
-    assert U.all_ranks(ok), f"rank {env.rank}: {fails}"
-
-
-@pytest.mark.parametrize("pattern", ["odd-routed_moe", "odd-native"])
-def test_routed_into_hop_mixed_rank_graph(env, stack, pattern):
-    moes, gammas, hop, eng = stack
-    _warm(env, moes, [1, 6, 8])
-    ok, replays = G.run_graph(env, moes, gammas, hop, pattern)
-    hop.raise_if_error()
-    if env.rank == 0:
-        print(f"[routed_moe_into_hop] graph mixed {pattern}: {replays} replays x 4 calls bit-equal")
-    assert U.all_ranks(ok)
-
-
-@pytest.mark.parametrize("kernel_delay", [False, True], ids=["no_kernel_delay", "kernel_delay"])
-def test_routed_into_hop_delayed_stress(env, stack, monkeypatch, kernel_delay):
-    moes, gammas, hop, eng = stack
-    if kernel_delay:
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-            routed_moe as routed_moe_mod,
-        )
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel import (  # noqa: E501
-            DELAY_ENTRY,
-        )
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel_xc import (  # noqa: E501
-            DELAY_FINALIZE,
-        )
-
-        monkeypatch.setattr(routed_moe_mod, "DELAY_INJECTION", (DELAY_ENTRY | DELAY_FINALIZE, 3000))
-    _warm(env, moes, [1, 6, 8])
-    ok, stats = G.run_stress(
-        env, moes, gammas, hop, batches=6, batch=16, max_delay_us=1500.0, seed=int(kernel_delay)
-    )
-    allstats = env.comm.allgather(stats)
-    if env.rank == 0:
-        print(f"[routed_moe_into_hop] delayed stress kernel_delay={kernel_delay}: {allstats}")
-    hop.raise_if_error()
-    assert U.all_ranks(ok), f"rank {env.rank}: {stats}"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py
deleted file mode 100644
index 24abaaf27e14..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_integration.py
+++ /dev/null
@@ -1,400 +0,0 @@
-# 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.
-"""routed_moe on the served module types (sm_103, real GLM-5.2 NVFP4 weights).
-
-The routed experts are built the way ``Deepseekv3MoE`` builds them:
-``create_moe`` with the served TP8/EP1 ``ModelConfig`` (rank ``GLM_MOE_DSA_ROUTED_MOE_RANK``),
-which returns a ``ConfigurableMoE`` around the resolved TRT-LLM-Gen leaf, and
-loaded from unsliced checkpoint tensors so the native loader does the TP
-slicing. routed_moe is called through ``dispatch.get_impl("routed_moe", "routed_output")``
-with that wrapper and compared bit for bit with ``wrapper.forward`` (the
-served native path) and with the hand-mirrored ``native_forward`` oracle of
-the other tests.
-
-``test_adopted_deepseek_moe_*`` run the served MoE: the built-in
-``Deepseekv3MoE`` adopted by the target's ``GlmMoE``, called inside a decode
-step.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-sys.path.insert(
-    0,
-    os.path.join(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "..", "shared", "kernels"
-    ),
-)
-import plan_testlib  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
-RANK = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
-
-
-@pytest.fixture(scope="module")
-def full():
-    return fx.load_layer_full(LAYER)
-
-
-@pytest.fixture(scope="module")
-def served(full):
-    wrapper, bias = fx.build_served_experts(LAYER, RANK, full)
-    return wrapper, bias
-
-
-@pytest.fixture(scope="module")
-def leaf_layer():
-    return fx.build_native_layer(layer=LAYER, tp_rank=RANK)
-
-
-@pytest.fixture(scope="module")
-def leaf_runner(leaf_layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    return RoutedMoe.from_native(leaf_layer.module, leaf_layer.bias)
-
-
-@pytest.fixture(scope="module")
-def decoder(leaf_layer):
-    return fx.ExpertDecoder(leaf_layer, leaf_layer.hidden_proxy(fx.prompt_token_ids()[:1])[0])
-
-
-def _impl():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    return dispatch.get_impl("routed_moe", "routed_output")
-
-
-def _inputs(leaf_layer, rows, kind, seed):
-    ids = fx.prompt_token_ids()
-    if kind == "proxy":
-        start = (seed * 7) % (len(ids) - rows)
-        x = leaf_layer.hidden_proxy(ids[start : start + rows])
-        return x, leaf_layer.router_logits(x)
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-    if kind == "random":
-        return x, leaf_layer.router_logits(x)
-    d = int(kind[len("distinct") :])
-    sets = fx.expert_sets_with_distinct(rows, d, seed=seed)
-    return x, fx.logits_for_expert_sets(sets, leaf_layer.bias, seed=seed)
-
-
-def test_served_wrapper_type_and_resolution(served):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        eligibility,
-        resolve_expert_owner,
-    )
-    from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE
-
-    wrapper, _ = served
-    leaf = wrapper.backend
-    print(
-        f"[routed_moe] served experts: {type(wrapper).__name__} -> {type(leaf).__name__} "
-        f"provider={getattr(type(leaf), 'provider', None)} comm={wrapper.comm} "
-        f"tp={leaf.tp_size}/{leaf.tp_rank} ep={leaf.ep_size} "
-        f"w3_w1={tuple(leaf.w3_w1_weight.shape)}"
-    )
-    assert isinstance(wrapper, ConfigurableMoE)
-    # the wrapper only delegates: the tensors live on the leaf
-    assert "w3_w1_weight" not in wrapper._parameters
-    owner, reason = resolve_expert_owner(wrapper)
-    assert owner is leaf, reason
-    x = torch.zeros(6, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
-    lg = torch.zeros(6, fx.NUM_EXPERTS, device="cuda", dtype=torch.float32)
-    ok, reason = eligibility(wrapper, x, lg)
-    assert ok, reason
-
-
-def test_served_weights_equal_fixture_slices(served, leaf_layer):
-    """The native TP8 loader's rank shard equals the fixture's hand-sliced
-    shard byte for byte, so the leaf-level tests run on served buffers.
-    """
-    leaf = served[0].backend
-    ref = leaf_layer.module
-    for name in (
-        "w3_w1_weight",
-        "w3_w1_weight_scale",
-        "w2_weight",
-        "w2_weight_scale",
-        "fc31_alpha",
-        "fc31_scale_c",
-        "fc2_alpha",
-        "fc31_input_scale",
-    ):
-        a, b = getattr(leaf, name), getattr(ref, name)
-        assert a.shape == b.shape and a.dtype == b.dtype, name
-        assert torch.equal(a.reshape(-1).view(torch.uint8), b.reshape(-1).view(torch.uint8)), name
-    assert torch.equal(served[1], leaf_layer.bias)
-
-
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-@pytest.mark.parametrize("kind", ["proxy", "random", "distinct40"])
-def test_served_wrapper_matches_served_forward(
-    served, leaf_layer, leaf_runner, decoder, rows, kind
-):
-    wrapper, bias = served
-    if kind == "distinct40" and rows * fx.TOP_K < 40:
-        kind = "distinct8"
-    x, logits = _inputs(leaf_layer, rows, kind, seed=50 + rows)
-    out = _impl()(wrapper, bias, x, logits, True)
-    assert out is not None, "routed_moe did not engage on the served wrapper"
-    served_native = wrapper.forward(x, logits, do_finalize=True, output_dtype=torch.bfloat16)
-    mirrored = leaf_layer.native_forward(x, logits)
-    torch.cuda.synchronize()
-    # the hand-mirrored oracle is the served path, bit for bit
-    assert torch.equal(mirrored, served_native)
-    raw = _impl()(wrapper, bias, x, logits, False)
-    raw_native = wrapper.forward(x, logits, do_finalize=False, output_dtype=torch.bfloat16)
-    mirrored_raw = leaf_layer.native_forward(x, logits, do_finalize=False)
-    assert torch.equal(fx.raw_rows_by_topk(mirrored_raw), fx.raw_rows_by_topk(raw_native))
-    cls, rep = fx.classify_vs_native(
-        out,
-        raw,
-        served_native,
-        raw_native,
-        lambda: fx.emulate_moe(leaf_layer, x, logits),
-        f"rows={rows} {kind}",
-        lambda: decoder.compare(leaf_runner, logits),
-    )
-    print(f"[routed_moe] served wrapper rows={rows} {kind}: {cls} {rep}")
-
-
-def test_served_wrapper_fallbacks(served, leaf_layer):
-    """Wrapper-level behaviour routed_moe does not implement selects the native path
-    before any device work.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        _WORKSPACES,
-    )
-
-    wrapper, bias = served
-    x, logits = _inputs(leaf_layer, 6, "proxy", seed=3)
-    snapshot = {k: v.ws.clone() for k, v in _WORKSPACES.items()}
-    for attr, value in (
-        ("unfinalized_combine_fn", lambda *a, **k: None),
-        ("reduce_results", True),
-        ("apply_router_weight_on_input", True),
-        ("_enable_perfect_router", True),
-    ):
-        saved = getattr(wrapper, attr, None)
-        setattr(wrapper, attr, value)
-        try:
-            assert _impl()(wrapper, bias, x, logits, True) is None, attr
-        finally:
-            setattr(wrapper, attr, saved)
-    # FP32 activations and more than eight rows stay native
-    assert _impl()(wrapper, bias, x.float(), logits, True) is None
-    x9 = torch.cat([x, x[:3]])
-    assert _impl()(wrapper, bias, x9, torch.cat([logits, logits[:3]]), True) is None
-    torch.cuda.synchronize()
-    for k, v in _WORKSPACES.items():
-        assert torch.equal(v.ws, snapshot[k])
-    assert _impl()(wrapper, bias, x, logits, True) is not None
-
-
-# ---------------------------------------------------------------------------
-# The served MoE: Deepseekv3MoE adopted by the target's GlmMoE.
-# ---------------------------------------------------------------------------
-@pytest.fixture(scope="module")
-def deepseek_moe(full):
-    return fx.build_deepseek_moe(LAYER, RANK, full)
-
-
-class _Features:
-    """Run exactly the fast paths ``value`` in a block run as a decode step, and count
-    routed_moe engagements; the selector state before the block is restored after it.
-    """
-
-    def __init__(self, value: str):
-        self.value = value
-        self.engaged = 0
-        self.declined = 0
-
-    def __enter__(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-        self.dispatch = dispatch
-        self.step = dispatch.step(dispatch.Phase.DECODE)
-        self.step.__enter__()
-        self.saved_state = (dispatch._INSTALLED[0], dispatch._VALIDATE[0], dispatch._require[0])
-        plan_testlib.select(dispatch, self.value)
-        dispatch.get_impl("routed_moe", "routed_output")
-        self.saved_fn = dispatch._REGISTRY["routed_moe"]["routed_output"]
-
-        def counting(*args, **kwargs):
-            out = self.saved_fn(*args, **kwargs)
-            if out is None:
-                self.declined += 1
-            else:
-                self.engaged += 1
-            return out
-
-        dispatch._REGISTRY["routed_moe"]["routed_output"] = counting
-        return self
-
-    def __exit__(self, *exc):
-        self.dispatch._REGISTRY["routed_moe"]["routed_output"] = self.saved_fn
-        (
-            self.dispatch._INSTALLED[0],
-            self.dispatch._VALIDATE[0],
-            self.dispatch._require[0],
-        ) = self.saved_state
-        self.step.__exit__(*exc)
-        return False
-
-
-def _hidden(rows, seed):
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    return (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-
-
-def _assert_forward_equivalent(on, off, emul_fn, label, id_check=None):
-    """forward(do_finalize=False) results [shared, fc2, scales, idx]: the
-    shared expert output is identical; the routed raw tuples agree on every
-    expert and scale, and their stock finalization is exact or in budget.
-    """
-    assert torch.equal(on[0], off[0]), f"{label}: shared expert output differs"
-    raw_on, raw_off = on[1:], off[1:]
-    cls, rep = fx.classify_vs_native(
-        fx.finalize_raw(raw_on),
-        raw_on,
-        fx.finalize_raw(raw_off),
-        raw_off,
-        emul_fn,
-        label,
-        id_check,
-    )
-    return cls
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_adopted_deepseek_moe_forward(deepseek_moe, leaf_layer, leaf_runner, decoder, rows):
-    """The served MoE forward (shared expert on the aux stream in
-    parallel, gate, routed experts) with ``routed_moe`` on matches it with ``routed_moe``
-    off (same experts and scales, shared output identical, routed output exact
-    or in budget), and routed_moe actually engaged.
-    """
-    moe = deepseek_moe
-    classes = []
-    for it in range(4):
-        h = _hidden(rows, 1000 * rows + it)
-        with _Features("routed_moe") as f:
-            on = moe.forward(h, do_finalize=False)
-            fin_on = moe.compute_routed_output(h, None, None, True)
-        with _Features("") as f_off:
-            off = moe.forward(h, do_finalize=False)
-            fin_off = moe.compute_routed_output(h, None, None, True)
-        torch.cuda.synchronize()
-        assert f.engaged == 2 and f.declined == 0, (f.engaged, f.declined)
-        assert f_off.engaged == 0 and f_off.declined == 0
-        logits = moe.gate(h)
-
-        def emul(h=h, logits=logits):
-            return fx.emulate_moe(leaf_layer, h, logits)
-
-        def ids(logits=logits):
-            return decoder.compare(leaf_runner, logits)
-
-        label = f"rows={rows} it={it}"
-        classes.append(_assert_forward_equivalent(on, off, emul, label, ids))
-        cls, _ = fx.classify_vs_native(fin_on, on[1:], fin_off, off[1:], emul, label, ids)
-        classes.append(cls)
-    print(f"[routed_moe] adopted Deepseekv3MoE rows={rows}: classes {classes}")
-
-
-def test_adopted_deepseek_moe_graph(deepseek_moe, leaf_layer, leaf_runner, decoder):
-    """CUDA graph of the served forward with ``routed_moe`` on (eager warmup first,
-    as the model's graph runner does), replayed on new inputs and compared
-    with eager native forwards; overlapping shared-expert GEMMs on the aux
-    stream run in every replay.
-    """
-    moe = deepseek_moe
-    rows = 6
-    h_static = _hidden(rows, 1)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with _Features("routed_moe") as f:
-        with torch.cuda.stream(stream):
-            for _ in range(2):
-                moe.forward(h_static, do_finalize=False)
-        torch.cuda.current_stream().wait_stream(stream)
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph, stream=stream):
-            out_static = moe.forward(h_static, do_finalize=False)
-        torch.cuda.synchronize()
-    assert f.engaged == 3 and f.declined == 0, (f.engaged, f.declined)
-    for it in range(16):
-        h = _hidden(rows, 500 + it)
-        h_static.copy_(h)
-        graph.replay()
-        torch.cuda.synchronize()
-        with _Features("") as f_off:
-            ref = moe.forward(h, do_finalize=False)
-        torch.cuda.synchronize()
-        assert f_off.engaged == 0
-        logits = moe.gate(h)
-        _assert_forward_equivalent(
-            [t.clone() for t in out_static],
-            ref,
-            lambda h=h, logits=logits: fx.emulate_moe(leaf_layer, h, logits),
-            f"graph replay {it}",
-            lambda logits=logits: decoder.compare(leaf_runner, logits),
-        )
-
-
-def test_adopted_deepseek_moe_capture_without_warmup_falls_back(full):
-    """A routed_moe call whose first use is inside a CUDA graph capture takes the
-    native path (no compilation or workspace initialization is recorded).
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe,
-    )
-
-    moe = fx.build_deepseek_moe(LAYER, RANK, full)
-    h = _hidden(6, 7)
-    # warm the native path, which the capture below records
-    with _Features(""):
-        moe.compute_routed_output(h, None, None, True)
-    torch.cuda.synchronize()
-    saved_ws = dict(routed_moe._WORKSPACES)
-    saved_fn = dict(routed_moe._COMPILED)
-    routed_moe._WORKSPACES.clear()
-    routed_moe._COMPILED.clear()
-    try:
-        stream = torch.cuda.Stream()
-        stream.wait_stream(torch.cuda.current_stream())
-        with _Features("routed_moe") as f:
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph, stream=stream):
-                moe.compute_routed_output(h, None, None, True)
-        assert f.declined == 1 and f.engaged == 0
-        assert not routed_moe._WORKSPACES and not routed_moe._COMPILED
-    finally:
-        routed_moe._WORKSPACES.update(saved_ws)
-        routed_moe._COMPILED.update(saved_fn)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py
deleted file mode 100644
index 4840fd9ba137..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_native_consumer_mpi.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# 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.
-"""routed_moe's raw tuple through the native unfinalized consumer (one 4-GPU node).
-
-On a single-node TP group the served DeepSeek-V3 layer hands the
-``do_finalize=False`` MoE output to ``MoEAllReduce``, i.e. the native
-``moe_finalize_allreduce`` op: a local finalize that follows the
-permutation, the shared-expert add, the all-reduce, the residual add and the
-RMSNorm. Run every test on the 4 GPUs of one node::
-
-    srun -N1 --ntasks-per-node=4 --mpi=pmix python3 -m pytest ...
-    test_routed_moe_native_consumer_mpi.py
-
-Rank ``r`` holds TP8 shard ``r`` of a real GLM-5.2 NVFP4 layer (``GLM_MOE_DSA_ROUTED_MOE_LAYER``,
-default 10) and runs captured real hidden-state windows of that layer. Per call
-every rank feeds the same consumer three raw tuples:
-
-* ``routed_moe``: routed_moe's token-major rows (row ``t * 8 + k``) with the identity
-  permutation;
-* ``routed_moe_relocated``: the same rows written at the native runner's permuted
-  positions, with the native permutation and the native expert scales;
-* ``native``: the native runner's raw tuple.
-
-``routed_moe`` and ``routed_moe_relocated`` must give identical outputs bit for bit: the
-consumer reads the token-major layout exactly like the native expert-grouped
-layout. Where every rank's routed_moe rows equal the native rows bit for bit, ``routed_moe``
-and ``native`` must give identical outputs as well; the other calls (an FP4
-rounding flip inside routed_moe's frozen error budget) are counted and must stay
-close.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, HERE)
-sys.path.insert(
-    0, os.path.join(os.path.dirname(HERE), "..", "..", "shared", "kernels", "post_moe_hop")
-)
-import post_moe_hop_mpi_utils as U  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
-WINDOWS = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_CONSUMER_WINDOWS", "8"))
-
-
-@pytest.fixture(scope="module")
-def env():
-    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
-        pytest.skip("needs an SM100-family GPU")
-    e = U.mpi_env()
-    if e.world < 2 or e.multi_node:
-        pytest.skip("needs a single-node TP group of at least 2 ranks")
-    return e
-
-
-@pytest.fixture(scope="module")
-def setup(env):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-    from tensorrt_llm._torch.distributed import MoEAllReduce
-
-    layer = fx.build_native_layer(layer=LAYER, tp_rank=env.rank)
-    runner = RoutedMoe.from_native(layer.module, layer.bias)
-    consumer = MoEAllReduce(env.mapping)
-    # compile routed_moe's variants on every rank before the first collective call
-    for rows in (1, 6, 8):
-        x = torch.zeros(rows, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
-        logits = layer.router_logits(x)
-        runner(x, logits, do_finalize=False)
-        layer.native_forward(x, logits, do_finalize=False)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    return layer, runner, consumer
-
-
-def _windows(rows: int, count: int):
-    win = [x for x, _ in fx.real_windows(fx.load_real_capture(LAYER), rows)]
-    assert win and win[0] is not None, f"no captured hidden states for layer {LAYER}"
-    step = max(1, len(win) // count)
-    return [win[i * step].cuda().contiguous() for i in range(count)]
-
-
-def _consume(consumer, raw, shared, residual, gamma):
-    from tensorrt_llm._torch.distributed import MoEAllReduceParams
-
-    fc2, scale, idx = raw
-    params = MoEAllReduceParams(
-        expanded_idx_to_permuted_idx=idx,
-        expert_scale_factor=scale,
-        shared_expert_output=shared,
-        residual=residual,
-        norm_weight=gamma,
-        eps=fx.RMS_EPS,
-        is_cutlass_min_latency=False,
-    )
-    return consumer(fc2, all_reduce_params=params)
-
-
-def _relocate(routed_moe_raw, native_raw):
-    """routed_moe's rows at the native runner's permuted positions."""
-    fc2_w, _, idx_w = routed_moe_raw
-    fc2_n, scale_n, idx_n = native_raw
-    assert fc2_n.shape[1] == fc2_w.shape[1], (fc2_n.shape, fc2_w.shape)
-    out = torch.zeros_like(fc2_n)
-    out[idx_n.reshape(-1).long()] = fc2_w[idx_w.reshape(-1).long()]
-    return out, scale_n, idx_n
-
-
-def _same(a, b) -> bool:
-    return all(torch.equal(x.view(torch.int16), y.view(torch.int16)) for x, y in zip(a, b))
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_raw_tuple_through_moe_finalize_allreduce(env, setup, rows):
-    layer, runner, consumer = setup
-    g_rank = torch.Generator(device="cuda").manual_seed(300 + 17 * env.rank + rows)
-    g_common = torch.Generator(device="cuda").manual_seed(500 + rows)
-    gamma = layer.norm_weight
-    exact_calls, calls, worst = 0, 0, [0, 0.0]
-    for x in _windows(rows, WINDOWS):
-        logits = layer.router_logits(x)
-        shared = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g_rank) * 0.1).bfloat16()
-        residual = torch.randn(rows, fx.HIDDEN, device="cuda", generator=g_common) * 4.0
-        residual = residual.bfloat16()
-        native_raw = layer.native_forward(x, logits, do_finalize=False)
-        routed_moe_raw = runner(x, logits, do_finalize=False)
-        torch.cuda.synchronize()
-        assert routed_moe_raw[0].shape == (rows * fx.TOP_K, fx.HIDDEN)
-        ident = torch.arange(rows * fx.TOP_K, dtype=torch.int32)
-        assert torch.equal(routed_moe_raw[2].view(-1).cpu(), ident), (
-            "routed_moe permutation is not t * 8 + k"
-        )
-        assert (native_raw[2] >= 0).all(), "TP8/EP1: every (t, k) has a local native row"
-        assert torch.equal(routed_moe_raw[1], native_raw[1]), "expert scales differ from native"
-        rep = fx.raw_report(routed_moe_raw, native_raw)
-        assert rep["same_expert"], rep
-        rows_exact = U.all_ranks(rep["rows_exact"] == rep["rows"])
-
-        out_w = _consume(consumer, routed_moe_raw, shared, residual, gamma)
-        out_r = _consume(consumer, _relocate(routed_moe_raw, native_raw), shared, residual, gamma)
-        out_n = _consume(consumer, native_raw, shared, residual, gamma)
-        torch.cuda.synchronize()
-        assert _same(out_w, out_r), "token-major layout read differently from the native layout"
-        if rows_exact:
-            assert _same(out_w, out_n), "identical rows, different consumer output"
-            exact_calls += 1
-        st = U.compare_stats(out_w[1], out_n[1])
-        worst[0] = max(worst[0], st["max_ulp"])
-        worst[1] = max(worst[1], st["max_abs"])
-        torch.testing.assert_close(out_w[0], out_n[0], rtol=0.05, atol=0.15)
-        calls += 1
-    allw = env.comm.allgather((exact_calls, calls, worst))
-    if env.rank == 0:
-        print(
-            f"[routed_moe] native consumer R={rows}: routed_moe raw == routed_moe relocated to the native layout "
-            f"bit for bit on {env.world} ranks; routed_moe == native in {[a[0] for a in allw]} of "
-            f"{calls} calls (all-rank-exact rows); residual vs native worst ulp/abs "
-            f"{[a[2] for a in allw]}"
-        )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py
deleted file mode 100644
index 09abab9d414a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_routed_moe.py
+++ /dev/null
@@ -1,722 +0,0 @@
-# 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.
-"""sm_103 correctness tests of the routed_moe routed NVFP4 MoE against the native region.
-
-Real GLM-5.2 NVFP4 weights (layer ``GLM_MOE_DSA_ROUTED_MOE_LAYER``, default 10, of TP8 rank
-``GLM_MOE_DSA_ROUTED_MOE_RANK``, default 0) through the stock TRT-LLM-Gen leaf are the oracle.
-Error budget (frozen before any performance tuning):
-expert IDs and routing weights bit-exact; finalized output
-``nrms(routed_moe - native) <= 1.5 * nrms(emul - native) + 2e-3`` where ``emul`` is
-the FP64-accumulated emulation, and no row may exceed 4x the per-tensor
-budget.
-
-Native expert IDs and their top-k order are compared through the native raw
-tuple: for every ``(t, k)`` the expert scale equals native's bit for bit and
-the routed_moe FC2 row is the row of the same expert as native's row at
-``expanded_idx_to_permuted_idx[t, k]`` (see ``routed_moe_fixtures.classify_vs_native``
-for the exact / reorder / budget classes of the floating-point results).
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import routed_moe_fixtures as fx  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-DEBUG_ROUTING = 1
-DEBUG_QUANT = 2
-DEBUG_FC1 = 4
-
-
-@pytest.fixture(scope="module")
-def layer():
-    torch.manual_seed(0)
-    return fx.build_native_layer(
-        layer=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10")),
-        tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")),
-    )
-
-
-@pytest.fixture(scope="module")
-def runner(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    return RoutedMoe.from_native(layer.module, layer.bias)
-
-
-@pytest.fixture(scope="module")
-def decoder(layer):
-    return fx.ExpertDecoder(layer, layer.hidden_proxy(fx.prompt_token_ids()[:1])[0])
-
-
-def _inputs(layer, rows: int, kind: str, seed: int = 0):
-    ids = fx.prompt_token_ids()
-    if kind == "proxy":
-        start = (seed * 7) % (len(ids) - rows)
-        x = layer.hidden_proxy(ids[start : start + rows])
-        return x, layer.router_logits(x)
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-    if kind == "random":
-        return x, layer.router_logits(x)
-    if kind.startswith("distinct"):
-        d = int(kind[len("distinct") :])
-        sets = fx.expert_sets_with_distinct(rows, d, seed=seed)
-        return x, fx.logits_for_expert_sets(sets, layer.bias, seed=seed)
-    raise ValueError(kind)
-
-
-def _dbg_buffer():
-    return torch.zeros(2 * 1024 * 1024, dtype=torch.int32, device="cuda")
-
-
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-@pytest.mark.parametrize("kind", ["proxy", "random"])
-def test_routing_matches_native_bitwise(layer, runner, decoder, rows, kind):
-    x, logits = _inputs(layer, rows, kind, seed=rows)
-    dbg = _dbg_buffer()
-    runner(x, logits, debug=DEBUG_ROUTING, debug_buffer=dbg)
-    torch.cuda.synchronize()
-    ids_ref, w_ref = fx.routing_reference(logits, layer.bias)
-    ids = dbg[:64].view(8, 8)[:rows].cpu()
-    w = dbg[64:128].view(torch.float32).view(8, 8)[:rows].cpu()
-    assert torch.equal(ids, ids_ref.cpu()), (ids, ids_ref)
-    assert torch.equal(w, w_ref.float().cpu())
-    num_slots = int(dbg[256].item())
-    expect = sorted(set(ids_ref.flatten().tolist()))
-    assert num_slots == len(expect)
-    assert dbg[128 : 128 + num_slots].cpu().tolist() == expect
-    # native unfinalized expert weights, same top-k order
-    outs = layer.native_forward(x, logits, do_finalize=False)
-    assert torch.equal(outs[1].float().cpu(), w_ref.float().cpu())
-    # native expert IDs and order, through the raw rows
-    _assert_raw_matches_native(
-        runner(x, logits, do_finalize=False), outs, lambda: decoder.compare(runner, logits)
-    )
-    # and decoded exactly through the probe fingerprints
-    dec = decoder.compare(runner, logits)
-    assert dec["ids_equal"] and dec["scale_exact"], dec
-    assert torch.equal(dec["native_ids"].cpu().int(), ids_ref.cpu()), (dec["native_ids"], ids_ref)
-
-
-def _assert_raw_matches_native(raw_routed_moe, raw_native, id_check=None):
-    """Every (t, k): same expert scale bits and same expert as native."""
-    rep = fx.raw_report(raw_routed_moe, raw_native)
-    assert rep["scale_exact"], (raw_routed_moe[1], raw_native[1])
-    if not rep["same_expert"]:
-        assert id_check is not None, rep
-        dec = id_check()
-        assert dec["ids_equal"] and dec["scale_exact"] and dec["margin"] > 4, (rep, dec)
-    return rep
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_activation_quant_matches_native_bitwise(layer, runner, rows):
-    x, logits = _inputs(layer, rows, "proxy", seed=3)
-    dbg = _dbg_buffer()
-    runner(x, logits, debug=DEBUG_QUANT, debug_buffer=dbg)
-    torch.cuda.synchronize()
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel import (  # noqa: E501
-        SFB_BYTES,
-        SFB_CHUNK_STRIDE,
-        SFB_STAGE_STRIDE,
-    )
-
-    raw = dbg.view(torch.uint8)[4096 : 4096 + 24 * 1024 + SFB_BYTES]
-    sb = raw[: 24 * 1024].view(24, 1024)
-    sfb_flat = raw[24 * 1024 :]
-    xq, xsf = layer.module.quantize_input(x)
-    xq = xq.view(torch.uint8).view(rows, fx.HIDDEN // 2)
-    xsf = xsf.view(torch.uint8).view(rows, fx.HIDDEN // 16)
-    kb = torch.arange(24, device="cuda")[:, None, None]
-    n = torch.arange(rows, device="cuda")[None, :, None]
-    b = torch.arange(128, device="cuda")[None, None, :]
-    phys = n * 128 + (((b // 16) ^ n) * 16) + (b % 16)
-    got = torch.gather(sb[:, None, :].expand(24, rows, 1024), 2, phys.expand(24, rows, 128))
-    want = xq.view(rows, 24, 128).permute(1, 0, 2)
-    assert torch.equal(got, want)
-    f = torch.arange(16, device="cuda")[None, None, :]
-    sf_phys = kb * SFB_STAGE_STRIDE + (f // 4) * SFB_CHUNK_STRIDE + n * 16 + f % 4
-    got_sf = sfb_flat[sf_phys.expand(24, rows, 16)]
-    want_sf = xsf.view(rows, 24, 16).permute(1, 0, 2)
-    assert torch.equal(got_sf, want_sf)
-
-
-@pytest.mark.parametrize("rows,kind", [(1, "proxy"), (6, "proxy"), (8, "distinct40")])
-def test_fc1_accumulators_match_exact(layer, runner, rows, kind):
-    x, logits = _inputs(layer, rows, kind, seed=5)
-    dbg = _dbg_buffer()
-    runner(x, logits, debug=DEBUG_ROUTING | DEBUG_FC1, debug_buffer=dbg)
-    torch.cuda.synchronize()
-    num_slots = int(dbg[256].item())
-    slots = dbg[128 : 128 + num_slots].tolist()
-    acc = dbg.view(torch.float32)[65536 : 65536 + num_slots * 4 * 128 * 8].view(
-        num_slots, 4 * 128, 8
-    )
-    m = layer.module
-    xq, xsf = m.quantize_input(x)
-    xd = fx.unpack_fp4(xq.view(torch.uint8)) * fx.e4m3_bytes_to_float(
-        xsf.view(torch.uint8).reshape(rows, fx.HIDDEN // 16)
-    ).repeat_interleave(16, dim=1)
-    worst = 0.0
-    for s, e in enumerate(slots):
-        w13, _ = fx.dequant_expert(m, e)
-        ref = w13 @ xd.T  # [512, rows]
-        got = acc[s, :, :rows].double()
-        scale = ref.abs().max().item() + 1e-30
-        worst = max(worst, ((got - ref).abs().max().item()) / scale)
-    assert worst < 1e-5, worst
-
-
-def _check_vs_native(layer, out, x, logits, label):
-    native = layer.native_forward(x, logits)
-    emul = fx.emulate_moe(layer, x, logits)
-    e_native = fx.error_stats(out, native)
-    e_emul = fx.error_stats(emul, native)
-    e_self = fx.error_stats(out, emul)
-    print(
-        f"[routed_moe] {label}: routed_moe-vs-native {e_native} emul-vs-native {e_emul} routed_moe-vs-emul {e_self}"
-    )
-    budget = 1.5 * e_emul["nrms"] + 2e-3
-    assert e_native["nrms"] <= budget, (e_native, e_emul)
-    row_ref = native.float().pow(2).mean(-1).sqrt()
-    row_err = (out.float() - native.float()).pow(2).mean(-1).sqrt()
-    assert torch.all(row_err <= 4 * budget * row_ref + 1e-6), (row_err, row_ref)
-    return e_native
-
-
-@pytest.mark.parametrize("rows", [1, 2, 4, 6, 8])
-@pytest.mark.parametrize("kind", ["proxy", "random", "distinct8"])
-def test_finalized_matches_native(layer, runner, rows, kind):
-    x, logits = _inputs(layer, rows, kind, seed=11 + rows)
-    out = runner(x, logits)
-    torch.cuda.synchronize()
-    _check_vs_native(layer, out, x, logits, f"rows={rows} {kind}")
-
-
-def test_raw_tuple_consumable_by_stock_finalizer(layer, runner):
-    rows = 6
-    x, logits = _inputs(layer, rows, "proxy", seed=2)
-    fc2, scale, idx = runner(x, logits, do_finalize=False)
-    out_final = runner(x, logits)
-    torch.cuda.synchronize()
-    assert fc2.shape == (rows * 8, fx.HIDDEN) and fc2.dtype == torch.bfloat16
-    assert scale.shape == (rows, 8) and scale.dtype == torch.bfloat16
-    assert idx.shape == (rows, 8) and idx.dtype == torch.int32
-    _, w_ref = fx.routing_reference(logits, layer.bias)
-    assert torch.equal(scale, w_ref)
-    # stock finalize semantics: sum_k scale[t,k] * fc2[idx[t,k]] in FP32
-    fin = torch.zeros(rows, fx.HIDDEN, device="cuda", dtype=torch.float32)
-    for k in range(8):
-        fin += scale[:, k : k + 1].float() * fc2[idx[:, k].long()].float()
-    fin = fin.to(torch.bfloat16)
-    _check_vs_native(layer, fin, x, logits, "raw tuple")
-    e = fx.error_stats(fin, out_final)
-    assert e["nrms"] < 2e-3, e
-
-
-@pytest.mark.parametrize("offset", [1, 4, 7])
-def test_misaligned_activation_is_copied_not_faulted(layer, runner, offset):
-    """A contiguous activation view whose base is not 16-byte aligned (the bulk
-    copy's requirement) gives the same raw and finalized outputs as the aligned
-    tensor: the runner copies it instead of faulting or changing path."""
-    rows = 6
-    x, logits = _inputs(layer, rows, "proxy", seed=3)
-    store = torch.empty(rows * fx.HIDDEN + offset, device="cuda", dtype=torch.bfloat16)
-    x_mis = store[offset:].view(rows, fx.HIDDEN)
-    x_mis.copy_(x)
-    assert x_mis.is_contiguous() and x_mis.data_ptr() % 16 != 0
-    if runner.schedule(rows, True) != "xc" or runner.schedule(rows, False) != "xc":
-        pytest.skip("bit-exact repeat calls are a property of the expert-cluster schedule")
-    # Outputs may live in the runner's workspace: copy each before the next call.
-    ref = runner(x, logits).clone()
-    got = runner(x_mis, logits).clone()
-    raw_ref = [t.clone() for t in runner(x, logits, do_finalize=False)]
-    raw_got = [t.clone() for t in runner(x_mis, logits, do_finalize=False)]
-    torch.cuda.synchronize()
-    assert torch.equal(got, ref)
-    for a, b in zip(raw_got, raw_ref):
-        assert torch.equal(a, b)
-
-
-def test_cuda_graph_replay(layer, runner):
-    rows = 6
-    x0, l0 = _inputs(layer, rows, "proxy", seed=1)
-    x_static = x0.clone()
-    l_static = l0.clone()
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        runner(x_static, l_static)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        out_static = runner(x_static, l_static)
-    for it in range(12):
-        kind = ["proxy", "random", "distinct8", "distinct48"][it % 4]
-        x, logits = _inputs(layer, rows, kind, seed=100 + it)
-        x_static.copy_(x)
-        l_static.copy_(logits)
-        graph.replay()
-        torch.cuda.synchronize()
-        eager = runner(x, logits)
-        torch.cuda.synchronize()
-        e = fx.error_stats(out_static, eager)
-        assert e["nrms"] < 1e-6, (it, kind, e)
-        if it < 4:
-            _check_vs_native(layer, out_static.clone(), x, logits, f"graph replay {it} {kind}")
-
-
-def _real_inputs(layer, rows: int, count: int, offset: int = 0):
-    """``count`` windows of the captured hidden states of this layer and their
-    router logits (GPU)."""
-    win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer.layer), rows)]
-    assert win and win[0] is not None, f"no captured hidden states for layer {layer.layer}"
-    step = max(1, len(win) // count)
-    xs = [win[(offset + i * step) % len(win)].cuda().contiguous() for i in range(count)]
-    return [(x, layer.router_logits(x)) for x in xs]
-
-
-def _expert_ids(logits, bias):
-    ids, _ = fx.routing_reference(logits, bias)
-    return set(ids.flatten().tolist())
-
-
-def _paddings(layer, x, logits, pads: int, other, seed: int):
-    """Padded-row contents ``{name: (x_pad, logits_pad)}`` for ``pads`` rows."""
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    rand = (torch.randn(pads, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-    out = {
-        "zeros": torch.zeros(pads, fx.HIDDEN, dtype=torch.bfloat16, device="cuda"),
-        "scaled_copy": (x[:1].float() * 3.0).to(torch.bfloat16).repeat(pads, 1),
-        "real": other[:pads].contiguous(),
-        "random": rand,
-    }
-    out = {k: (v, layer.router_logits(v)) for k, v in out.items()}
-    # padded rows whose top-8 sets avoid every live expert: the call's
-    # distinct-expert count grows by 8 per padded row (two rounds when it
-    # exceeds the cluster count)
-    used = _expert_ids(logits, layer.bias)
-    free = [
-        e
-        for e in torch.randperm(
-            fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(seed)
-        ).tolist()
-        if e not in used
-    ]
-    sets = [free[8 * i : 8 * i + 8] for i in range(pads)]
-    out["new_experts"] = (rand, fx.logits_for_expert_sets(sets, layer.bias, seed=seed))
-    return out
-
-
-@pytest.mark.parametrize("live", [1, 2, 4, 6])
-def test_padded_rows_do_not_change_live_rows(layer, runner, live):
-    """Padding invariance of the default (expert-cluster) schedule, bit for bit.
-
-    Real hidden-state windows of ``live`` rows are padded to 8 rows with
-    zeros, a scaled copy of a live row, rows of another real window, random
-    rows, and rows routed to experts no live row uses (more distinct experts,
-    up to a two-round call). The live rows' raw FC2 rows, expert scales and
-    finalized output must equal the unpadded call's bit for bit.
-    """
-    if runner.schedule(8, True) != "xc" or runner.schedule(8, False) != "xc":
-        pytest.skip("padding invariance is a property of the expert-cluster schedule")
-    pads = 8 - live
-    clusters = runner._xc_ctas(8, True) // 4
-    others = _real_inputs(layer, 8, 3, offset=5)
-    seen = {}
-    for i, (x, logits) in enumerate(_real_inputs(layer, live, 3, offset=live)):
-        fin = runner(x, logits).clone()
-        raw = [t.clone() for t in runner(x, logits, do_finalize=False)]
-        torch.cuda.synchronize()
-        for name, (xp, lp) in _paddings(layer, x, logits, pads, others[i][0], 70 + i).items():
-            xa = torch.cat([x, xp]).contiguous()
-            la = torch.cat([logits, lp.float()]).contiguous()
-            fin_p = runner(xa, la)
-            raw_p = runner(xa, la, do_finalize=False)
-            torch.cuda.synchronize()
-            d = len(_expert_ids(la, layer.bias))
-            seen.setdefault(name, []).append(d)
-            tag = f"live={live} window {i} padding {name} D={d}"
-            assert torch.equal(fin_p[:live].view(torch.int16), fin.view(torch.int16)), tag
-            rows_p = raw_p[0][: live * 8].view(torch.int16)
-            assert torch.equal(rows_p, raw[0].view(torch.int16)), tag
-            assert torch.equal(raw_p[1][:live], raw[1]), tag
-            assert torch.isfinite(fin_p).all(), tag
-    two_round = [d for ds in seen.values() for d in ds if d > clusters]
-    print(
-        f"[routed_moe] padding invariance live={live}: bit-exact for paddings {seen} "
-        f"(distinct experts per call; {len(two_round)} two-round calls, {clusters} clusters)"
-    )
-    assert two_round, "no padded call reached the two-round schedule"
-
-
-REORDER_REL = 2.0**-20
-
-
-def test_fp32_accumulator_output_reorder_class(layer, runner):
-    """routed_moe_into_hop "fp32" output: the kernel adds the eight weighted rows of a token
-    with atomic FP32 adds. Every run, with or without padded rows, lies within
-    the FP32 reorder bound of the exact sum of the same products (the raw
-    rows' ``scale * row``, each exact in FP32); the run-to-run spread is
-    counted, not assumed."""
-    rows = 6
-    if runner.schedule(8, False, True) != "xc":
-        pytest.skip("FP32 accumulator output needs the expert-cluster schedule")
-    others = _real_inputs(layer, 8, 3, offset=9)
-    distinct_runs = []
-    for i, (x, logits) in enumerate(_real_inputs(layer, rows, 3, offset=2)):
-        fc2, scale, idx = runner(x, logits, do_finalize=False)
-        terms = scale.double()[:, :, None] * fc2[idx.long()].double()
-        exact = terms.sum(1)
-        bound = REORDER_REL * terms.abs().sum(1)
-        runs = []
-        for _ in range(8):
-            acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
-            runner(x, logits, red_out=acc)
-            runs.append(acc)
-        xp = torch.cat([x, others[i][0][:2]]).contiguous()
-        lp = torch.cat([logits, layer.router_logits(others[i][0][:2]).float()]).contiguous()
-        acc8 = torch.zeros(8, fx.HIDDEN, dtype=torch.float32, device="cuda")
-        runner(xp, lp, red_out=acc8)
-        torch.cuda.synchronize()
-        for j, acc in enumerate(runs + [acc8[:rows]]):
-            assert ((acc.double() - exact).abs() <= bound).all(), f"window {i} run {j}"
-        distinct_runs.append(len({r.cpu().numpy().tobytes() for r in runs}))
-    print(
-        f"[routed_moe] fp32 accumulator: within the FP32 reorder bound in every run; distinct "
-        f"results over 8 runs per window {distinct_runs}"
-    )
-
-
-def test_workspace_created_once_in_device_memory(layer, runner):
-    """The per-device workspace is device memory even when first requested
-    inside the served loader's ``MetaInitMode``, is never replaced, and a
-    request for more CTA slots than it holds raises."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe as rm,
-    )
-    from tensorrt_llm._torch.models.modeling_utils import MetaInitMode
-
-    dev = runner.device
-    key = dev.index if dev.index is not None else torch.cuda.current_device()
-    saved = rm._WORKSPACES.pop(key, None)
-    try:
-        with MetaInitMode():
-            assert torch.empty(1, device="cuda").is_meta, "MetaInitMode is not active"
-            fresh = rm.RoutedMoe.from_native(layer.module, layer.bias)
-            ws = rm._workspace(fresh.device, fresh.num_ctas)
-        for name, t in ws.tensors().items():
-            assert not t.is_meta and t.is_cuda and t.device.index == key, name
-            assert t.data_ptr() != 0, name
-        assert ws.num_ctas >= rm._sm_count(dev)
-        assert rm._workspace(dev, fresh.num_ctas) is ws
-        with pytest.raises(RuntimeError):
-            rm._workspace(dev, ws.num_ctas + 4)
-        assert rm._WORKSPACES[key] is ws
-        print(f"[routed_moe] workspace cuda:{key}: {ws.num_ctas} CTA slots, {ws.nbytes()} bytes")
-    finally:
-        if saved is not None:
-            rm._WORKSPACES[key] = saved
-
-
-def test_eager_calls_from_two_streams_are_ordered(layer, runner):
-    """A call issued on another stream than the previous call waits for the
-    previous stream's submitted work before it uses the shared workspace."""
-    rows = 6
-    x, logits = _inputs(layer, rows, "proxy", seed=4)
-    x2, l2 = _inputs(layer, rows, "distinct40", seed=6)
-    ref = runner(x, logits).clone()
-    ref2 = runner(x2, l2).clone()
-    torch.cuda.synchronize()
-    a, b = torch.cuda.Stream(), torch.cuda.Stream()
-    a.wait_stream(torch.cuda.current_stream())
-    b.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(a):
-        torch.cuda._sleep(100_000_000)  # tens of milliseconds of device time
-        out_a = runner(x, logits)
-    with torch.cuda.stream(b):
-        out_b = runner(x2, l2)
-        done_b = torch.cuda.Event()
-        done_b.record(b)
-    done_b.synchronize()
-    assert a.query(), "the call on the second stream did not wait for the first stream"
-    torch.cuda.synchronize()
-    assert torch.equal(out_a.view(torch.int16), ref.view(torch.int16))
-    assert torch.equal(out_b.view(torch.int16), ref2.view(torch.int16))
-
-
-def test_fallback_is_decided_before_device_work(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        _WORKSPACES,
-        routed_output,
-    )
-
-    x9 = torch.randn(9, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
-    l9 = torch.randn(9, 256, device="cuda")
-    snapshot = {k: v.ws.clone() for k, v in _WORKSPACES.items()}
-    assert routed_output(layer.module, layer.bias, x9, l9, True) is None
-    x1 = torch.randn(1, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
-    assert (
-        routed_output(
-            layer.module,
-            layer.bias,
-            x1,
-            torch.randn(1, 256, device="cuda").to(torch.bfloat16),
-            True,
-        )
-        is None
-    )
-    assert (
-        routed_output(
-            layer.module, layer.bias, x1.float(), torch.randn(1, 256, device="cuda"), True
-        )
-        is None
-    )
-    # FP32 router bias (non-TRTLLM gate), a bias tensor that is not the routing
-    # method's, and modules without registered expert parameters stay native
-    l1 = torch.randn(1, 256, device="cuda")
-    assert routed_output(layer.module, layer.bias.float(), x1, l1, True) is None
-    assert routed_output(layer.module, layer.bias.clone(), x1, l1, True) is None
-    assert routed_output(torch.nn.Linear(4, 4), layer.bias, x1, l1, True) is None
-
-    class _DelegatingOnly(torch.nn.Module):
-        """Exposes the expert tensors as properties but owns none of them."""
-
-        w3_w1_weight = property(lambda self: layer.module.w3_w1_weight)
-        w2_weight = property(lambda self: layer.module.w2_weight)
-
-    assert routed_output(_DelegatingOnly(), layer.bias, x1, l1, True) is None
-    torch.cuda.synchronize()
-    for k, v in _WORKSPACES.items():
-        assert torch.equal(v.ws, snapshot[k])
-
-
-def test_dispatch_registration():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    fn = dispatch.get_impl("routed_moe", "routed_output")
-    assert callable(fn)
-    assert callable(dispatch.get_impl("routed_moe", "eligibility"))
-
-
-# ---------------------------------------------------------------------------
-# Near-tie routing: FP32 logits built so that the 8th/9th selection keys tie
-# exactly, or differ by one key ulp, including experts with equal BF16 bias.
-# ---------------------------------------------------------------------------
-TIE_PATTERNS = (
-    "same_logit",  # equal bias, identical logits: exact key tie
-    "same_key",  # equal bias, different logits, exact key tie
-    "plus_ulp",  # equal bias, key(b) = key(a) + 1 ulp
-    "minus_ulp",  # equal bias, key(b) = key(a) - 1 ulp
-    "cross_bias_lo",  # bias(a) < bias(b), exact key tie
-    "cross_bias_hi",  # bias(a) > bias(b), exact key tie
-    "three_way",  # three equal keys for the last two slots
-    "inner_tie",  # two winners with equal keys (top-k order only)
-)
-_WIN_LOGIT = 6.0
-_CONT_LOGIT = 0.3
-_BOOST = 0.02
-
-
-def _tie_row(bias, pattern: str, gen: torch.Generator):
-    """One row of logits for ``pattern`` plus its resolution variants.
-
-    Returns ``(row, variants)``: ``variants`` are logit rows in which each
-    possible resolution of the tie wins clearly (by ``_BOOST`` in logit); the
-    native outcome of ``row`` must equal exactly one of them.
-    """
-    b = bias.float().cpu()
-    vals = sorted(set(b.tolist()))
-    groups = [[e for e in range(fx.NUM_EXPERTS) if b[e].item() == v] for v in vals]
-    big = [g for g in groups if len(g) >= 3]
-    perm = lambda n: torch.randperm(n, generator=gen).tolist()  # noqa: E731
-    row = -12.0 + 0.5 * torch.rand(fx.NUM_EXPERTS, generator=gen)
-    if pattern in ("cross_bias_lo", "cross_bias_hi"):
-        # cont[0] is the lower expert index; its bias is the lower one for
-        # "lo" and the higher one for "hi"
-        lo_g = groups[0] + (groups[1] if len(groups) > 1 else [])
-        hi_g = groups[-1] + (groups[-2] if len(groups) > 2 else [])
-        want_lo_first = pattern == "cross_bias_lo"
-        cands = [(x, y) for x in lo_g for y in hi_g if (x < y) == want_lo_first]
-        assert cands, pattern
-        x, y = cands[perm(len(cands))[0]]
-        cont = [x, y] if want_lo_first else [y, x]
-    else:
-        g = big[perm(len(big))[0]]
-        pick = [g[i] for i in perm(len(g))[:3]]
-        cont = sorted(pick[:3] if pattern == "three_way" else pick[:2])
-    n_win = {"three_way": 6, "inner_tie": 6}.get(pattern, 7)
-    others = [e for e in perm(fx.NUM_EXPERTS) if e not in cont]
-    winners = others[:n_win]
-    for i, e in enumerate(winners):
-        row[e] = _WIN_LOGIT + 0.01 * i
-    bdev = bias.float().cuda()
-    a = cont[0]
-    if pattern == "inner_tie":
-        # two contenders become winners with identical keys; no boundary tie
-        for e in cont[:2]:
-            row[e] = _WIN_LOGIT + 0.5
-    else:
-        for e in cont:
-            row[e] = _CONT_LOGIT
-    key_a = fx._key_of(torch.tensor(row[a].item(), device="cuda"), bdev[a]).item()
-    c = cont[1]
-    if pattern == "same_key" or pattern.startswith("cross_bias"):
-        hit = fx.logit_for_key(key_a, bdev[c], 0, avoid=row[a].item())
-        assert hit is not None, pattern
-        row[c] = hit
-    elif pattern in ("plus_ulp", "minus_ulp"):
-        hit = fx.logit_for_key(key_a, bdev[c], 1 if pattern == "plus_ulp" else -1)
-        assert hit is not None, pattern
-        row[c] = hit
-    variants = []
-    if pattern == "three_way":
-        for pair in ((0, 1), (0, 2), (1, 2)):
-            v = row.clone()
-            for i in pair:
-                v[cont[i]] += _BOOST
-            variants.append(v)
-    else:
-        for e in cont[:2]:
-            v = row.clone()
-            v[e] += _BOOST
-            variants.append(v)
-    return row, variants
-
-
-def _tie_batch(bias, patterns, seed):
-    gen = torch.Generator().manual_seed(seed)
-    rows, variants = [], []
-    for p in patterns:
-        r, v = _tie_row(bias, p, gen)
-        rows.append(r)
-        variants.append(v)
-    logits = torch.stack(rows).cuda().contiguous()
-    nvar = max(len(v) for v in variants)
-    var_batches = [
-        torch.stack([v[min(i, len(v) - 1)] for v in variants]).cuda().contiguous()
-        for i in range(nvar)
-    ]
-    return logits, var_batches, [len(v) for v in variants]
-
-
-def _check_ties(layer, runner, decoder, x, logits, var_batches, nvars, label):
-    native_raw = layer.native_forward(x, logits, do_finalize=False)
-    native_fin = layer.native_forward(x, logits)
-    routed_moe_raw = runner(x, logits, do_finalize=False)
-    routed_moe_fin = runner(x, logits)
-    torch.cuda.synchronize()
-    tie_rows = fx.raw_rows_by_topk(native_raw)
-    var_rows = [
-        fx.raw_rows_by_topk(layer.native_forward(x, v, do_finalize=False)) for v in var_batches
-    ]
-    resolution = []
-    for t in range(x.shape[0]):
-        match = [i for i in range(nvars[t]) if torch.equal(tie_rows[t], var_rows[i][t])]
-        distinct = len({tuple(var_rows[i][t].flatten()[::97].tolist()) for i in range(nvars[t])})
-        assert distinct == nvars[t], f"{label} row {t}: variants not distinguishable"
-        assert len(match) == 1, f"{label} row {t}: native outcome matches variants {match}"
-        resolution.append(match[0])
-    cls, rep = fx.classify_vs_native(
-        routed_moe_fin,
-        routed_moe_raw,
-        native_fin,
-        native_raw,
-        lambda: fx.emulate_moe(layer, x, logits),
-        label,
-        lambda: decoder.compare(runner, logits),
-    )
-    # the decoded native experts of every tie row, independent of the rows
-    dec = decoder.compare(runner, logits)
-    assert dec["ids_equal"] and dec["scale_exact"], (label, dec)
-    gaps = fx.boundary_gap(logits, layer.bias)
-    print(
-        f"[routed_moe] ties {label}: resolution={resolution} class={cls} "
-        f"rows_exact={rep['rows_exact']}/{rep['rows']} decode_margin={dec['margin']:.1f} "
-        f"boundary_gap={gaps.tolist()}"
-    )
-    return resolution
-
-
-@pytest.mark.parametrize("seed", [0, 1, 2, 3])
-def test_boundary_ties_match_native(layer, runner, decoder, seed):
-    """All eight tie patterns in one R=8 call; routed_moe equals native bit for bit
-    and the native outcome is identified as one clear resolution.
-    """
-    logits, var_batches, nvars = _tie_batch(layer.bias, TIE_PATTERNS, seed)
-    x, _ = _inputs(layer, len(TIE_PATTERNS), "proxy", seed=20 + seed)
-    res = _check_ties(layer, runner, decoder, x, logits, var_batches, nvars, f"R=8 seed={seed}")
-    # identical logits and bias give identical keys in any arithmetic; they
-    # resolve to the lower expert index (variant 0), as the native 64-bit key
-    # (score bits, 65535 - index) orders them
-    for p, r in zip(TIE_PATTERNS, res):
-        if p in ("same_logit", "three_way", "inner_tie"):
-            assert r == 0, (p, r)
-
-
-@pytest.mark.parametrize("pattern", TIE_PATTERNS)
-def test_boundary_tie_single_row(layer, runner, decoder, pattern):
-    logits, var_batches, nvars = _tie_batch(layer.bias, [pattern], 7)
-    x, _ = _inputs(layer, 1, "proxy", seed=31)
-    _check_ties(layer, runner, decoder, x, logits, var_batches, nvars, f"R=1 {pattern}")
-
-
-def test_near_tie_sweep_matches_native(layer, runner, decoder):
-    """Randomized near-boundary rows: the 8th/9th keys at -2..+2 key ulps
-    (random contender biases and logits); routed_moe equals native bit for bit.
-    """
-    gen = torch.Generator().manual_seed(77)
-    bdev = layer.bias.float().cuda()
-    batches = 48
-    exact_ties = 0
-    classes = {}
-    for i in range(batches):
-        rows = []
-        for t in range(8):
-            row = -12.0 + 0.5 * torch.rand(fx.NUM_EXPERTS, generator=gen)
-            order = torch.randperm(fx.NUM_EXPERTS, generator=gen).tolist()
-            for k, e in enumerate(order[:7]):
-                row[e] = _WIN_LOGIT + 0.01 * k
-            a, c = order[7], order[8]
-            row[a] = 0.1 + 0.6 * torch.rand(1, generator=gen).item()
-            key_a = fx._key_of(torch.tensor(row[a].item(), device="cuda"), bdev[a]).item()
-            off = int(torch.randint(-2, 3, (1,), generator=gen).item())
-            hit = fx.logit_for_key(key_a, bdev[c], off)
-            row[c] = hit if hit is not None else row[a]
-            rows.append(row)
-        logits = torch.stack(rows).cuda().contiguous()
-        exact_ties += int((fx.boundary_gap(logits, layer.bias) == 0).sum().item())
-        x, _ = _inputs(layer, 8, "random", seed=200 + i)
-        cls, _ = fx.check_call(
-            layer, runner, x, logits, label=f"near-tie batch {i}", decoder=decoder
-        )
-        dec = decoder.compare(runner, logits)
-        assert dec["ids_equal"] and dec["scale_exact"], (i, dec)
-        classes[cls] = classes.get(cls, 0) + 1
-    print(
-        f"[routed_moe] near-tie sweep: {batches * 8} rows, {exact_ties} exact boundary key ties, "
-        f"batch classes {classes}"
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py
deleted file mode 100644
index c3a19dd2c013..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_spec_prefetch.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# 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.
-"""routed_moe expert-cluster kernel with package cross_layer_expert_prefetch's speculative expert prefetch.
-
-The prefetch (``spec_k`` > 0) only issues ``cp.async.bulk.prefetch.L2``, so
-every output must equal the prefetch-off kernel bit for bit, for any content
-of the expert-ID row (including out-of-range IDs). The self-prediction store
-(``spec_self``) must write the call's distinct experts ranked by row count.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import routed_moe_fixtures as fx  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-KMAX = 32
-REORDER_REL = 2.0**-20
-
-
-@pytest.fixture(scope="module")
-def layer():
-    torch.manual_seed(0)
-    return fx.build_native_layer(
-        layer=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10")),
-        tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")),
-    )
-
-
-@pytest.fixture(scope="module")
-def runner(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    r = RoutedMoe.from_native(layer.module, layer.bias)
-    if r.schedule(8, False, True) != "xc":
-        pytest.skip("the speculative prefetch is part of the expert-cluster schedule")
-    return r
-
-
-class _Spec:
-    """Install a cross_layer_expert_prefetch hook on the runner's expert leaf."""
-
-    def __init__(self, runner, ids, k, ctas=0, self_store=False):
-        self.owner = runner.module
-        self.spec = (ids, k, ctas, self_store)
-
-    def __enter__(self):
-        self.owner._expert_prefetch_spec = lambda: self.spec
-        return self
-
-    def __exit__(self, *exc):
-        del self.owner._expert_prefetch_spec
-
-
-def _windows(layer, rows, count, offset=0):
-    win = [x for x, _ in fx.real_windows(fx.load_real_capture(layer.layer), rows)]
-    if not win:
-        pytest.skip(
-            f"no captured hidden states of layer {layer.layer} (GLM_MOE_DSA_REAL_MOE_INPUTS_DIR)"
-        )
-    step = max(1, len(win) // count)
-    xs = [win[(offset + i * step) % len(win)].cuda().contiguous() for i in range(count)]
-    return [(x, layer.router_logits(x)) for x in xs]
-
-
-def _ids(kind, seed):
-    g = torch.Generator().manual_seed(seed)
-    if kind == "valid":
-        ids = torch.randperm(fx.NUM_EXPERTS, generator=g)[:KMAX]
-    else:  # out of range: the kernel masks each ID into [0, 256)
-        ids = torch.randint(-(2**31), 2**31 - 1, (KMAX,), generator=g)
-    return ids.to(torch.int32).cuda()
-
-
-def _outputs(runner, x, logits):
-    fin = runner(x, logits)
-    fc2, scale, _ = runner(x, logits, do_finalize=False)
-    acc = torch.zeros(x.shape[0], fx.HIDDEN, dtype=torch.float32, device="cuda")
-    runner(x, logits, red_out=acc)
-    torch.cuda.synchronize()
-    return fin.clone(), fc2.clone(), scale.clone(), acc
-
-
-def _check_same(ref, got, what):
-    fin, fc2, scale, _ = ref
-    gfin, gfc2, gscale, gacc = got
-    assert torch.equal(gfin, fin), f"{what}: finalized output differs"
-    assert torch.equal(gfc2, fc2), f"{what}: raw FC2 rows differ"
-    assert torch.equal(gscale, scale), f"{what}: routing weights differ"
-    # FP32 accumulator: atomic adds, compare within the FP32 reorder bound of
-    # the exact sum of the same raw products
-    rows = gacc.shape[0]
-    terms = scale.double()[:, :, None] * fc2.view(rows, 8, -1).double()
-    exact = terms.sum(1)
-    bound = REORDER_REL * terms.abs().sum(1)
-    assert ((gacc.double() - exact).abs() <= bound).all(), f"{what}: accumulator"
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("k,ctas", [(12, 0), (32, 48), (1, 4)])
-@pytest.mark.parametrize("ids_kind", ["valid", "out_of_range"])
-def test_prefetch_outputs_bit_identical(layer, runner, rows, k, ctas, ids_kind):
-    for i, (x, logits) in enumerate(_windows(layer, rows, 2, offset=rows)):
-        ref = _outputs(runner, x, logits)
-        ids = _ids(ids_kind, seed=rows * 100 + i)
-        with _Spec(runner, ids, k, ctas):
-            got = _outputs(runner, x, logits)
-        _check_same(ref, got, f"rows={rows} k={k} ctas={ctas} {ids_kind} window {i}")
-
-
-def test_prefetch_cuda_graph_replay(layer, runner):
-    rows = 6
-    wins = _windows(layer, rows, 6, offset=3)
-    x_static, l_static = (t.clone() for t in wins[0])
-    ids = _ids("valid", seed=7)
-    with _Spec(runner, ids, 12):
-        stream = torch.cuda.Stream()
-        stream.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(stream):
-            runner(x_static, l_static, do_finalize=False)
-        torch.cuda.current_stream().wait_stream(stream)
-        # the eager call compiled the prefetch variant: the capture records it
-        assert runner.ready_for_capture(False, rows=rows)
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph, stream=stream):
-            fc2_s, scale_s, _ = runner(x_static, l_static, do_finalize=False)
-    for it, (x, logits) in enumerate(wins):
-        x_static.copy_(x)
-        l_static.copy_(logits)
-        ids.copy_(_ids("valid", seed=100 + it))
-        graph.replay()
-        torch.cuda.synchronize()
-        fc2, scale, _ = runner(x, logits, do_finalize=False)
-        torch.cuda.synchronize()
-        assert torch.equal(fc2_s, fc2) and torch.equal(scale_s, scale), it
-
-
-def _ranked_experts(logits, bias):
-    ids, _ = fx.routing_reference(logits, bias)
-    counts = torch.bincount(ids.flatten().long().cpu(), minlength=fx.NUM_EXPERTS)
-    distinct = [e for e in range(fx.NUM_EXPERTS) if counts[e] > 0]
-    return sorted(distinct, key=lambda e: (-int(counts[e]), e))
-
-
-@pytest.mark.parametrize("rows", [1, 3, 6, 8])
-def test_self_store_ranks_by_row_count(layer, runner, rows):
-    for i, (x, logits) in enumerate(_windows(layer, rows, 3, offset=5 * rows)):
-        ref = _outputs(runner, x, logits)
-        ids = torch.full((KMAX,), 999, dtype=torch.int32, device="cuda")
-        with _Spec(runner, ids, 12, self_store=True):
-            fc2, scale, _ = runner(x, logits, do_finalize=False)
-            torch.cuda.synchronize()
-            got = ids.tolist()
-            # the store is idempotent for the same call (the prefetch reads
-            # the previous call's row first)
-            gfin = runner(x, logits)
-            torch.cuda.synchronize()
-        assert torch.equal(fc2, ref[1]) and torch.equal(scale, ref[2]), (rows, i)
-        assert torch.equal(gfin, ref[0]), (rows, i)
-        want = _ranked_experts(logits, layer.bias)[:KMAX]
-        assert got[: len(want)] == want, (rows, i, got, want)
-        assert all(v == 999 for v in got[len(want) :]), (rows, i, got)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py
deleted file mode 100644
index afe1bf6302c9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/test_routed_moe_stress.py
+++ /dev/null
@@ -1,812 +0,0 @@
-# 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.
-"""Stress, scheduling and real-input tests of the routed_moe cross-CTA protocol (sm_103).
-
-* ``test_sanitizer_cells``: routed_moe-vs-native cells in both output modes plus a
-  short graph, meant to be run under ``compute-sanitizer`` (memcheck /
-  synccheck / racecheck, kernel filter ``kns=RoutedMoeKernel``) in both
-  FC1->FC2 hand-off modes.
-* ``test_delayed_*``: the kernel's delay injection (``routed_moe.DELAY_INJECTION``) holds a
-  pseudo-random third of the CTAs (re-drawn every launch) at CTA entry, before
-  publishing FC1 tiles and split partials, before split owners gather and
-  before FC2 operand fetches; outputs must stay bit-identical to undelayed
-  launches.
-* ``test_corun_*``: other kernels on a second stream hold SMs while routed_moe runs
-  (a shared-memory-heavy spinning kernel, cuBLAS GEMMs), so not all routed_moe CTAs
-  are co-resident; routed_moe must complete with bit-identical outputs.
-* ``test_graph_many_launches_real_layers``: one CUDA graph with dozens of
-  back-to-back routed_moe launches over several real layers (one shared workspace),
-  every output equal to eager routed_moe and classified against native.
-* ``test_real_*``: a native capture of real routed-MoE inputs.
-
-Comparisons with native use ``routed_moe_fixtures.classify_vs_native`` (exact
-expert IDs and scales; finalized output exact, FP32-reorder-close or within
-the frozen budget).
-
-Layers: ``GLM_MOE_DSA_ROUTED_MOE_REAL_LAYERS`` (default ``10,30,50,77``, the layers with captured
-hidden states) of TP8 rank ``GLM_MOE_DSA_ROUTED_MOE_RANK``.
-"""
-
-import faulthandler
-import os
-import sys
-import time
-
-import cutlass
-import cutlass.cute as cute
-import cutlass.utils as utils
-import pytest
-import torch
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import tcgen05
-from cutlass.cutlass_dsl import Int32, Int64, T, Uint32, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-    ptx,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import routed_moe_fixtures as fx  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-RANK = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
-LAYERS = [
-    int(v) for v in os.environ.get("GLM_MOE_DSA_ROUTED_MOE_REAL_LAYERS", "10,30,50,77").split(",")
-]
-WATCHDOG_S = 180
-
-
-@pytest.fixture(scope="module")
-def layers():
-    return {lay: fx.build_native_layer(layer=lay, tp_rank=RANK) for lay in LAYERS}
-
-
-@pytest.fixture(scope="module")
-def runners(layers):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    return {lay: RoutedMoe.from_native(m.module, m.bias) for lay, m in layers.items()}
-
-
-@pytest.fixture(scope="module")
-def decoders(layers):
-    probe = layers[LAYERS[0]].hidden_proxy(fx.prompt_token_ids()[:1])[0]
-    return {lay: fx.ExpertDecoder(m, probe) for lay, m in layers.items()}
-
-
-@pytest.fixture(scope="module")
-def captures():
-    return {
-        lay: [(None if h is None else h.cuda(), lg.cuda()) for h, lg in fx.load_real_capture(lay)]
-        for lay in LAYERS
-    }
-
-
-def _synthetic(layer, rows, distinct, seed):
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-    sets = fx.expert_sets_with_distinct(rows, distinct, seed=seed)
-    return x, fx.logits_for_expert_sets(sets, layer.bias, seed=seed)
-
-
-def _real(captures, lay, rows, index):
-    win = fx.real_windows(captures[lay], rows)
-    x, lg = win[(index * 131) % len(win)]
-    return x.cuda().contiguous(), lg.cuda().contiguous()
-
-
-CELLS = [(1, 8), (2, 14), (4, 25), (6, 21), (6, 38), (8, 48)]
-
-
-def _cells(layer, captures, lay, seed=0):
-    out = [(f"R{r}D{d}", *_synthetic(layer, r, d, seed + r * 100 + d)) for r, d in CELLS]
-    for i in range(3):
-        out.append((f"real{lay}R6#{i}", *_real(captures, lay, 6, seed + i)))
-    out.append((f"real{lay}R8", *_real(captures, lay, 8, seed + 7)))
-    return out
-
-
-def _identical(a, b, label):
-    """routed_moe against routed_moe: bit-identical finalized outputs or raw tuples."""
-    if isinstance(a, torch.Tensor):
-        assert torch.equal(a, b), label
-    else:
-        for i in range(3):
-            assert torch.equal(a[i], b[i]), f"{label} (raw part {i})"
-
-
-def _routed_moe_both(runner, x, lg):
-    return runner(x, lg), runner(x, lg, do_finalize=False)
-
-
-def _classify(layer, x, lg, fin, raw, label, emul=True):
-    n_fin = layer.native_forward(x, lg)
-    n_raw = layer.native_forward(x, lg, do_finalize=False)
-    emul_fn = (lambda: fx.emulate_moe(layer, x, lg)) if emul else None
-    return fx.classify_vs_native(fin, raw, n_fin, n_raw, emul_fn, label)
-
-
-# ---------------------------------------------------------------------------
-# (a) compute-sanitizer target
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize("mode", ["fin", "raw"])
-def test_sanitizer_cells(layers, runners, captures, mode):
-    """One output mode per run, so a sanitizer report names one compiled
-    kernel. Outputs are compared with native loosely here (the exact checks
-    are in the other tests). The graph segment needs the caching allocator
-    (``PYTORCH_NO_CUDA_MEMORY_CACHING=1`` would cudaMalloc during capture).
-    """
-    lay = LAYERS[0]
-    layer, runner = layers[lay], runners[lay]
-    fin = mode == "fin"
-    for label, x, lg in _cells(layer, captures, lay):
-        out = runner(x, lg, do_finalize=fin)
-        ref = layer.native_forward(x, lg, do_finalize=fin)
-        torch.cuda.synchronize()
-        if fin:
-            assert fx.error_stats(out, ref)["nrms"] < 2e-3, label
-        else:
-            rep = fx.raw_report(out, ref)
-            assert rep["scale_exact"] and rep["same_expert"], (label, rep)
-    if os.environ.get("PYTORCH_NO_CUDA_MEMORY_CACHING") == "1":
-        return
-    x, lg = _real(captures, lay, 6, 11)
-    eager = runner(x, lg, do_finalize=fin)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        outs = [runner(x, lg, do_finalize=fin) for _ in range(4)]
-    graph.replay()
-    torch.cuda.synchronize()
-    for i, o in enumerate(outs):
-        _identical(o, eager, f"graph launch {i}")
-
-
-class _BulkCopyControl:
-    """Racecheck control: the bulk-copy -> mbarrier (transaction count) ->
-    shared-load pattern routed_moe uses for its activation rows and FC2 operand
-    images, in isolation and correctly synchronized. Two phases on one
-    buffer: copy, wait, every thread reads, CTA barrier, copy again (the
-    write-after-read reuse), wait, read.
-    """
-
-    NBYTES = 2048
-
-    @cute.jit
-    def run(self, src: Int64, dst: Int64, stream):
-        @cute.struct
-        class Storage:
-            mbar: cute.struct.MemRange[cutlass.Int64, 2]
-            buf: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, 2048], 128]
-
-        self.storage_t = Storage
-        self.kernel(src, dst).launch(
-            grid=[1, 1, 1], block=[128, 1, 1], smem=Storage.size_in_bytes(), stream=stream
-        )
-
-    @cute.kernel
-    def kernel(self, src: Int64, dst: Int64):
-        tidx, _, _ = cute.arch.thread_idx()
-        smem = utils.SmemAllocator()
-        st = smem.allocate(self.storage_t)
-        bar = st.mbar.data_ptr()
-        buf = Int32(st.buf.data_ptr().toint())
-        if tidx == 0:
-            cute.arch.mbarrier_init(bar, 1)
-            cute.arch.mbarrier_init(bar + 1, 1)
-        cute.arch.mbarrier_init_fence()
-        cute.arch.sync_threads()
-        for phase in cutlass.range_constexpr(2):
-            if tidx == 0:
-                b = Int32((bar + phase).toint())
-                ptx.mbar_arrive_expect_tx(b, Int32(self.NBYTES))
-                ptx.bulk_g2s(buf, src + Int64(phase * self.NBYTES), Int32(self.NBYTES), b)
-            cute.arch.mbarrier_wait(bar + phase, 0)
-            for j in cutlass.range_constexpr(4):
-                v = ptx.lds_u32(buf + (j * 128 + tidx) * 4)
-                ptx.st_u32(dst + Int64(phase * self.NBYTES) + Int64((j * 128 + tidx) * 4), v)
-            # generic reads before the next async-proxy write of the buffer
-            ptx.fence_proxy_async_shared()
-            cute.arch.sync_threads()
-
-
-class _ChainControl:
-    """Racecheck control for routed_moe's shared-buffer reuse (FC1 B operand stages
-    overwritten by FC2 operands): warp 0 writes a buffer with generic stores,
-    proxy fence, mbarrier arrive; warp 1 waits and signals warp 2 through a
-    second mbarrier, with a plain arrive or with ``tcgen05.commit`` (routed_moe's
-    FC1-done signal); warp 2 waits, overwrites the buffer with a bulk copy,
-    waits for it and reads it back. Correctly ordered in both variants.
-    """
-
-    NBYTES = 1024
-
-    def __init__(self, use_commit: bool):
-        self.use_commit = use_commit
-
-    @cute.jit
-    def run(self, src: Int64, dst: Int64, stream):
-        @cute.struct
-        class Storage:
-            mbar: cute.struct.MemRange[cutlass.Int64, 3]
-            buf: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, 1024], 128]
-
-        self.storage_t = Storage
-        self.kernel(src, dst).launch(
-            grid=[1, 1, 1], block=[96, 1, 1], smem=Storage.size_in_bytes(), stream=stream
-        )
-
-    @cute.kernel
-    def kernel(self, src: Int64, dst: Int64):
-        tidx, _, _ = cute.arch.thread_idx()
-        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
-        lane = tidx % 32
-        smem = utils.SmemAllocator()
-        st = smem.allocate(self.storage_t)
-        bar = st.mbar.data_ptr()
-        buf = Int32(st.buf.data_ptr().toint())
-        if tidx == 0:
-            cute.arch.mbarrier_init(bar, 1)
-            cute.arch.mbarrier_init(bar + 1, 1)
-            cute.arch.mbarrier_init(bar + 2, 1)
-        cute.arch.mbarrier_init_fence()
-        cute.arch.sync_threads()
-        if warp == 0:
-            for j in cutlass.range_constexpr(8):
-                ptx.sts_u32(buf + (j * 32 + lane) * 4, Uint32(lane + 32 * j))
-            ptx.fence_proxy_async_shared()
-            cute.arch.sync_warp()
-            if lane == 0:
-                cute.arch.mbarrier_arrive(bar)
-        if warp == 1:
-            cute.arch.mbarrier_wait(bar, 0)
-            if cutlass.const_expr(self.use_commit):
-                with cute.arch.elect_one():
-                    tcgen05.commit(bar + 1)
-            else:
-                if lane == 0:
-                    cute.arch.mbarrier_arrive(bar + 1)
-        if warp == 2:
-            cute.arch.mbarrier_wait(bar + 1, 0)
-            if lane == 0:
-                b2 = Int32((bar + 2).toint())
-                ptx.mbar_arrive_expect_tx(b2, Int32(self.NBYTES))
-                ptx.bulk_g2s(buf, src, Int32(self.NBYTES), b2)
-            cute.arch.mbarrier_wait(bar + 2, 0)
-            for j in cutlass.range_constexpr(8):
-                v = ptx.lds_u32(buf + (j * 32 + lane) * 4)
-                ptx.st_u32(dst + Int64((j * 32 + lane) * 4), v)
-
-
-@pytest.mark.parametrize("use_commit", [False, True])
-def test_racecheck_control_chain(use_commit):
-    from cuda.bindings import driver as cuda
-
-    src = torch.randint(0, 255, (_ChainControl.NBYTES,), dtype=torch.uint8, device="cuda")
-    dst = torch.zeros_like(src)
-    cur = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-    fn = cute.compile(_ChainControl(use_commit).run, cutlass.Int64(0), cutlass.Int64(0), cur)
-    fn(cutlass.Int64(src.data_ptr()), cutlass.Int64(dst.data_ptr()), cur)
-    torch.cuda.synchronize()
-    assert torch.equal(src, dst)
-
-
-def test_racecheck_control_bulk_copy():
-    from cuda.bindings import driver as cuda
-
-    src = torch.randint(0, 255, (2 * _BulkCopyControl.NBYTES,), dtype=torch.uint8, device="cuda")
-    dst = torch.zeros_like(src)
-    cur = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-    ctl = _BulkCopyControl()
-    fn = cute.compile(ctl.run, cutlass.Int64(0), cutlass.Int64(0), cur)
-    fn(cutlass.Int64(src.data_ptr()), cutlass.Int64(dst.data_ptr()), cur)
-    torch.cuda.synchronize()
-    assert torch.equal(src, dst)
-
-
-# ---------------------------------------------------------------------------
-# (b) delayed producers / owners / consumers
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize("delay_ns", [2000, 9000])
-@pytest.mark.parametrize("finalize", [True, False])
-def test_delayed_ctas_exact(layers, runners, captures, monkeypatch, delay_ns, finalize):
-    """Delayed launches are bit-identical to undelayed ones (eager and in a
-    replayed graph); the undelayed outputs are classified against native.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe as routed_moe_mod,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.kernel import (  # noqa: E501
-        DELAY_CONSUMER,
-        DELAY_ENTRY,
-        DELAY_OWNER,
-        DELAY_PRODUCER,
-    )
-
-    lay = LAYERS[0]
-    layer, runner = layers[lay], runners[lay]
-    cells = _cells(layer, captures, lay, seed=delay_ns)
-    assert routed_moe_mod.DELAY_INJECTION[0] == 0
-    undelayed = []
-    classes = {}
-    for label, x, lg in cells:
-        fin, raw = _routed_moe_both(runner, x, lg)
-        cls, _ = _classify(layer, x, lg, fin, raw, label)
-        classes[cls] = classes.get(cls, 0) + 1
-        undelayed.append(fin if finalize else raw)
-    torch.cuda.synchronize()
-    sites = DELAY_ENTRY | DELAY_PRODUCER | DELAY_OWNER | DELAY_CONSUMER
-    monkeypatch.setattr(routed_moe_mod, "DELAY_INJECTION", (sites, delay_ns))
-    t0 = time.perf_counter()
-    for rep in range(3):  # every launch draws a new delayed set
-        for (label, x, lg), base in zip(cells, undelayed):
-            out = runner(x, lg, do_finalize=finalize)
-            torch.cuda.synchronize()
-            _identical(out, base, f"{label} rep{rep} delay={delay_ns}")
-    eager_s = time.perf_counter() - t0
-    # graph: all cells back to back, replayed (epochs advance per replay)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    static_inputs = [(x.clone(), lg.clone()) for _, x, lg in cells]
-    with torch.cuda.stream(stream):
-        runner(static_inputs[0][0], static_inputs[0][1], do_finalize=finalize)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        outs = [runner(x, lg, do_finalize=finalize) for x, lg in static_inputs]
-    for rep in range(4):
-        graph.replay()
-        torch.cuda.synchronize()
-        for (label, _, _), o, base in zip(cells, outs, undelayed):
-            _identical(o, base, f"graph {label} rep{rep}")
-    print(
-        f"[routed_moe] delayed CTAs bit-identical: delay_ns={delay_ns} finalize={finalize} "
-        f"eager {eager_s:.2f}s; undelayed vs native {classes}"
-    )
-
-
-# ---------------------------------------------------------------------------
-# (c) co-running kernels that hold SMs
-# ---------------------------------------------------------------------------
-@dsl_user_op
-def _sreg_u32(name: str, *, loc=None, ip=None) -> Uint32:
-    return Uint32(
-        llvm.inline_asm(
-            T.i32(),
-            [],
-            f"mov.u32 $0, %{name};",
-            "=r",
-            has_side_effects=True,
-            is_align_stack=False,
-            asm_dialect=llvm.AsmDialect.AD_ATT,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-class _SmHog:
-    """Test-only kernel: ``ctas`` CTAs that each reserve ``smem_bytes`` of
-    shared memory (more than half an SM, so one CTA per SM and no routed_moe CTA
-    beside it) and spin for ``spin_ns`` on the global timer. Each CTA records
-    ``started``, its SM id and its dynamic shared-memory size.
-    """
-
-    def __init__(self, ctas: int, smem_bytes: int, spin_ns: int):
-        self.ctas = ctas
-        self.smem_bytes = smem_bytes
-        self.spin_ns = spin_ns
-
-    @cute.jit
-    def run(self, info_addr: Int64, stream):
-        self.kernel(info_addr).launch(
-            grid=[self.ctas, 1, 1], block=[32, 1, 1], smem=self.smem_bytes, stream=stream
-        )
-
-    @cute.kernel
-    def kernel(self, info_addr: Int64):
-        tidx, _, _ = cute.arch.thread_idx()
-        bidx, _, _ = cute.arch.block_idx()
-        rec = info_addr + Int64(bidx) * Int64(16)
-        if tidx == 0:
-            ptx.st_u32(rec + Int64(4), _sreg_u32("smid"))
-            ptx.st_u32(rec + Int64(8), _sreg_u32("dynamic_smem_size"))
-            ptx.st_relaxed_gpu_u32(rec, Uint32(1))
-        t0 = ptx.globaltimer_u64()
-        t = t0
-        while t - t0 < Int64(self.spin_ns):
-            t = ptx.globaltimer_u64()
-        if tidx == 0:
-            ptx.st_u32(rec + Int64(12), Uint32(1))
-
-
-class _WaitStarted:
-    """Test-only kernel: one thread spins until ``n`` hog CTAs have started,
-    so work enqueued after it on the same stream runs while they are all
-    resident.
-    """
-
-    def __init__(self, n: int):
-        self.n = n
-
-    @cute.jit
-    def run(self, info_addr: Int64, stream):
-        self.kernel(info_addr).launch(grid=[1, 1, 1], block=[32, 1, 1], stream=stream)
-
-    @cute.kernel
-    def kernel(self, info_addr: Int64):
-        tidx, _, _ = cute.arch.thread_idx()
-        if tidx == 0:
-            i = Int32(0)
-            while i < self.n:
-                if ptx.ld_relaxed_gpu_u32(info_addr + Int64(i) * Int64(16)) != Uint32(0):
-                    i = i + 1
-
-
-def _make_hog(ctas: int, smem_bytes: int, spin_ns: int):
-    """Returns ``(launch_hog(stream), wait_started(stream), info)``; ``info``
-    is int32 ``[ctas, 4]``: started, SM id, dynamic smem bytes, done.
-    """
-    from cuda.bindings import driver as cuda
-
-    cur = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-    hog_fn = cute.compile(_SmHog(ctas, smem_bytes, spin_ns).run, cutlass.Int64(0), cur)
-    wait_fn = cute.compile(_WaitStarted(ctas).run, cutlass.Int64(0), cur)
-    info = torch.zeros(ctas, 4, dtype=torch.int32, device="cuda")
-
-    def launch(stream: torch.cuda.Stream):
-        hog_fn(cutlass.Int64(info.data_ptr()), cuda.CUstream(stream.cuda_stream))
-
-    def wait_started(stream: torch.cuda.Stream):
-        wait_fn(cutlass.Int64(info.data_ptr()), cuda.CUstream(stream.cuda_stream))
-
-    return launch, wait_started, info
-
-
-class _Watchdog:
-    def __enter__(self):
-        faulthandler.dump_traceback_later(WATCHDOG_S, exit=True)
-        return self
-
-    def __exit__(self, *exc):
-        faulthandler.cancel_dump_traceback_later()
-        return False
-
-
-def _refs(layer, runner, cells):
-    """Eager routed_moe outputs without co-runners (the reference the co-run
-    outputs must equal bit for bit), classified against native.
-    """
-    refs, classes = [], {}
-    for label, x, lg in cells:
-        fin, raw = _routed_moe_both(runner, x, lg)
-        cls, _ = _classify(layer, x, lg, fin, raw, label)
-        classes[cls] = classes.get(cls, 0) + 1
-        refs.append(fin)
-    torch.cuda.synchronize()
-    return refs, classes
-
-
-def _ev():
-    return torch.cuda.Event(enable_timing=True)
-
-
-def test_corun_sm_hog(layers, runners, captures):
-    """routed_moe while a spinning kernel holds SMs (one CTA per SM, more shared
-    memory than a routed_moe CTA leaves free). A and C enqueue routed_moe only after every
-    hog CTA is resident, so routed_moe starts with part (C) or none (A) of its grid
-    resident and must finish once the hog exits; B starts the hog at
-    staggered points of routed_moe's execution; D replays a routed_moe graph against hogs.
-    Outputs must equal the undisturbed routed_moe outputs bit for bit.
-    """
-    lay = LAYERS[0]
-    layer, runner = layers[lay], runners[lay]
-    sms = torch.cuda.get_device_properties(0).multi_processor_count
-    smem = 120 * 1024
-    hog_all, wait_all, info_all = _make_hog(sms, smem, 300_000)  # every SM, 300 us
-    hog_half, wait_half, info_half = _make_hog(sms // 2, smem, 2_000_000)  # half, 2 ms
-    cells = _cells(layer, captures, lay, seed=5)
-    refs, classes = _refs(layer, runner, cells)
-    side = torch.cuda.Stream()
-    main = torch.cuda.current_stream()
-    timing = {}
-
-    def check(outs, label):
-        for i, (o, r) in enumerate(zip(outs, refs)):
-            _identical(o, r, f"{label} #{i}")
-
-    def held(hog, wait, info, label):
-        """Hog on ``side``; routed_moe batch on ``main`` once every hog CTA runs."""
-        info.zero_()
-        e0, e_hog, e_first, e_all = _ev(), _ev(), _ev(), _ev()
-        side.wait_stream(main)
-        e0.record(side)
-        hog(side)
-        e_hog.record(side)
-        main.wait_event(e0)
-        wait(main)
-        outs = []
-        for i, (_, x, lg) in enumerate(cells):
-            outs.append(runner(x, lg))
-            if i == 0:
-                e_first.record(main)
-        e_all.record(main)
-        torch.cuda.synchronize()
-        check(outs, label)
-        assert int(info[:, 0].sum()) == info.shape[0] and int(info[:, 3].sum()) == info.shape[0]
-        timing[f"{label}_hog_end_ms"] = round(e0.elapsed_time(e_hog), 3)
-        timing[f"{label}_first_routed_moe_end_ms"] = round(e0.elapsed_time(e_first), 3)
-        timing[f"{label}_all_routed_moe_end_ms"] = round(e0.elapsed_time(e_all), 3)
-        timing[f"{label}_hog_sms"] = len(set(info[:, 1].tolist()))
-        timing[f"{label}_hog_dyn_smem"] = sorted(set(info[:, 2].tolist()))
-
-    with _Watchdog():
-        e0, e1 = _ev(), _ev()
-        e0.record(main)
-        outs = [runner(x, lg) for _, x, lg in cells]
-        e1.record(main)
-        torch.cuda.synchronize()
-        check(outs, "alone")
-        timing["alone_batch_ms"] = round(e0.elapsed_time(e1), 3)
-        held(hog_all, wait_all, info_all, "A")
-        held(hog_half, wait_half, info_half, "C")
-        # B: routed_moe first, the hog launched at staggered offsets
-        for delay_cycles in (0, 2000, 10000, 30000, 60000, 120000):
-            for rep in range(4):
-                info_all.zero_()
-                side.wait_stream(main)
-                outs = [runner(x, lg) for _, x, lg in cells[:4]]
-                with torch.cuda.stream(side):
-                    torch.cuda._sleep(delay_cycles)
-                hog_all(side)
-                outs += [runner(x, lg) for _, x, lg in cells[4:]]
-                torch.cuda.synchronize()
-                check(outs, f"hog after {delay_cycles} cycles rep{rep}")
-                assert int(info_all[:, 3].sum()) == sms
-        # D: a captured graph of routed_moe launches replayed against repeated hogs
-        static_inputs = [(x.clone(), lg.clone()) for _, x, lg in cells]
-        gstream = torch.cuda.Stream()
-        gstream.wait_stream(main)
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph, stream=gstream):
-            gouts = [runner(x, lg) for x, lg in static_inputs]
-        for rep in range(8):
-            side.wait_stream(main)
-            if rep % 2 == 0:
-                hog_all(side)
-                wait_all(main)
-            graph.replay()
-            if rep % 2 == 1:
-                hog_half(side)
-            torch.cuda.synchronize()
-            check(gouts, f"graph vs hog rep{rep}")
-    print(f"[routed_moe] co-run SM hog ok: timing {timing}; refs vs native {classes}")
-    # the hogs really held one SM per CTA with their full shared memory
-    assert timing["C_hog_sms"] == sms // 2 and timing["A_hog_sms"] == sms, timing
-    assert timing["A_hog_dyn_smem"] == [smem] and timing["C_hog_dyn_smem"] == [smem], timing
-
-
-def test_corun_gemm_side_stream(layers, runners, captures):
-    """CuBLAS GEMMs (shared-memory-heavy, many CTAs) on a second stream while
-    a graph of routed_moe launches replays, as the shared expert's GEMMs overlap the
-    routed experts in the served model.
-    """
-    lay = LAYERS[0]
-    layer, runner = layers[lay], runners[lay]
-    cells = _cells(layer, captures, lay, seed=9)
-    refs, classes = _refs(layer, runner, cells)
-    static_inputs = [(x.clone(), lg.clone()) for _, x, lg in cells]
-    main = torch.cuda.current_stream()
-    gstream = torch.cuda.Stream()
-    gstream.wait_stream(main)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=gstream):
-        gouts = [runner(x, lg) for x, lg in static_inputs]
-    side = torch.cuda.Stream()
-    a = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
-    b = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
-    small_w = torch.randn(512, fx.HIDDEN, device="cuda", dtype=torch.bfloat16)
-    with _Watchdog():
-        for rep in range(10):
-            side.wait_stream(main)
-            with torch.cuda.stream(side):
-                for _ in range(3):
-                    torch.matmul(a, b)
-                    # shared-expert-sized GEMMs as in the served overlap
-                    for x, _ in static_inputs[:4]:
-                        torch.nn.functional.linear(x, small_w)
-            graph.replay()
-            torch.cuda.synchronize()
-            for i, (o, r) in enumerate(zip(gouts, refs)):
-                _identical(o, r, f"gemm co-run rep{rep} #{i}")
-    print(f"[routed_moe] co-run GEMM ok; refs vs native {classes}")
-
-
-# ---------------------------------------------------------------------------
-# (d) many back-to-back launches in one graph, several real layers
-# ---------------------------------------------------------------------------
-def test_graph_many_launches_real_layers(layers, runners, captures):
-    """One graph, 12 launches per layer interleaved over the layers (one
-    shared workspace, every fourth launch in raw mode, an unrelated consumer
-    after each); every replayed output equals the eager routed_moe output for its
-    input bit for bit and is classified against native.
-    """
-    per_layer = 12
-    plan = []  # (layer, finalize)
-    for i in range(per_layer):
-        for lay in LAYERS:
-            plan.append((lay, (i % 4) != 3))
-    wins = {lay: fx.real_windows(captures[lay], 6) for lay in LAYERS}
-    static_inputs = [
-        (
-            torch.empty(6, fx.HIDDEN, device="cuda", dtype=torch.bfloat16),
-            torch.empty(6, fx.NUM_EXPERTS, device="cuda", dtype=torch.float32),
-        )
-        for _ in plan
-    ]
-
-    def fill(rep):
-        used = []
-        for j, ((lay, _), (xs, ls)) in enumerate(zip(plan, static_inputs)):
-            x, lg = wins[lay][(rep * 977 + j * 31) % len(wins[lay])]
-            xs.copy_(x)
-            ls.copy_(lg)
-            used.append((x, lg))
-        return used
-
-    used = fill(0)
-    main = torch.cuda.current_stream()
-    stream = torch.cuda.Stream()
-    stream.wait_stream(main)
-    with torch.cuda.stream(stream):
-        for (lay, fin), (xs, ls) in zip(plan, static_inputs):
-            runners[lay](xs, ls, do_finalize=fin)
-    main.wait_stream(stream)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    sink = torch.zeros(6, fx.HIDDEN, device="cuda", dtype=torch.float32)
-    with torch.cuda.graph(graph, stream=stream):
-        outs = []
-        for (lay, fin), (xs, ls) in zip(plan, static_inputs):
-            o = runners[lay](xs, ls, do_finalize=fin)
-            outs.append(o)
-            # an unrelated consumer between launches, as the model's layers
-            sink.add_(o.float() if fin else o[0][:6].float())
-    classes = {}
-    emul_left = 8
-    for rep in range(6):
-        if rep:
-            used = fill(rep)
-        graph.replay()
-        torch.cuda.synchronize()
-        for j, ((lay, fin), o, (x, lg)) in enumerate(zip(plan, outs, used)):
-            tag = f"rep{rep} launch{j} layer{lay}"
-            e_fin, e_raw = _routed_moe_both(runners[lay], x, lg)
-            _identical(o, e_fin if fin else e_raw, tag)
-            cls, _ = _classify(layers[lay], x, lg, e_fin, e_raw, tag, emul=emul_left > 0)
-            if cls == "budget":
-                emul_left -= 1
-            classes[cls] = classes.get(cls, 0) + 1
-    print(
-        f"[routed_moe] graph of {len(plan)} launches over layers {LAYERS}: 6 replays bit-identical "
-        f"to eager routed_moe; vs native {classes}"
-    )
-
-
-# ---------------------------------------------------------------------------
-# Real routed-MoE inputs (native capture): correctness, not only timing
-# ---------------------------------------------------------------------------
-EMUL_PER_CASE = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_EMUL_PER_CASE", "24"))
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_real_windows_match_native(layers, runners, decoders, captures, rows):
-    """Every real window of every captured layer: exact expert IDs and
-    scales; finalized output exact, FP32-reorder-close, or (sampled, FP64
-    emulation) within the frozen budget.
-    """
-    stride = {1: 7, 6: 6, 8: 8}[rows]
-    summary = {}
-    for lay in LAYERS:
-        layer, runner = layers[lay], runners[lay]
-        wins = fx.real_windows(captures[lay], rows, stride)
-        assert wins and wins[0][0] is not None, f"layer {lay}: no captured hidden states"
-        classes = {}
-        nrms = []
-        emul_left = EMUL_PER_CASE
-        for i, (x, lg) in enumerate(wins):
-            fin, raw = _routed_moe_both(runner, x, lg)
-            n_fin = layer.native_forward(x, lg)
-            n_raw = layer.native_forward(x, lg, do_finalize=False)
-            emul_fn = None
-            if emul_left > 0:
-
-                def emul_fn(x=x, lg=lg):
-                    return fx.emulate_moe(layer, x, lg)
-
-            cls, rep = fx.classify_vs_native(
-                fin,
-                raw,
-                n_fin,
-                n_raw,
-                emul_fn,
-                f"layer {lay} window {i}",
-                lambda lg=lg: decoders[lay].compare(runner, lg),
-            )
-            if cls == "budget":
-                emul_left -= 1
-            classes[cls] = classes.get(cls, 0) + 1
-            if cls != "exact":
-                nrms.append(fx.error_stats(fin, n_fin)["nrms"])
-        gaps = torch.cat([fx.boundary_gap(lg, layer.bias) for _, lg in wins])
-        summary[lay] = {
-            "windows": len(wins),
-            "classes": classes,
-            "max_nrms_non_exact": max(nrms) if nrms else 0.0,
-            "rows_exact_boundary_tie": int((gaps == 0).sum()),
-            "rows_gap_le_1ulp": int((gaps <= 4.8e-7).sum()),
-            "min_gap": float(gaps.min()),
-        }
-    print(f"[routed_moe] real windows R={rows} rank={RANK}: {summary}")
-
-
-def test_real_routing_all_layers(layers, decoders, captures):
-    """Routing on the real logits of every captured MoE layer with that
-    layer's own router bias: the experts of every (t, k), decoded exactly
-    from the FC2 rows of a probe activation (``ExpertDecoder``), and the
-    expert scale bits equal native's. The expert weights are layer
-    ``LAYERS[0]``'s; routing depends only on logits and bias.
-    """
-    import glob
-    import re
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-
-    base = LAYERS[0]
-    layer = layers[base]
-    decoder = decoders[base]
-    all_layers = sorted(
-        {
-            int(re.search(r"_layer(\d+)\.pt$", p).group(1))
-            for p in glob.glob(os.path.join(fx.REAL_DIR, "rows*_layer*.pt"))
-        }
-    )
-    totals = {"layers": len(all_layers), "windows": 0, "min_margin": float("inf")}
-    totals.update({"exact_ties": 0, "gap_le_1ulp": 0, "tie_windows_checked": 0})
-    for lay in all_layers:
-        cap = [(None, lg.cuda()) for _, lg in fx.load_real_capture(lay)]
-        bias = fx.load_router_bias(lay)
-        runner = RoutedMoe.from_native(layer.module, bias)
-        for i, (_, lg) in enumerate(fx.real_windows(cap, 6)):
-            dec = decoder.compare(runner, lg, bias=bias)
-            assert dec["ids_equal"] and dec["scale_exact"], (lay, i, dec)
-            totals["windows"] += 1
-            totals["min_margin"] = min(totals["min_margin"], dec["margin"])
-            if bool((fx.boundary_gap(lg, bias) <= 4.8e-7).any()):
-                totals["tie_windows_checked"] += 1
-        gaps = torch.cat([fx.boundary_gap(lg, bias) for _, lg in cap])
-        totals["exact_ties"] += int((gaps == 0).sum())
-        totals["gap_le_1ulp"] += int((gaps <= 4.8e-7).sum())
-    print(f"[routed_moe] real routing over all layers (decoded experts): {totals}")
-    assert totals["min_margin"] > 4, totals
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py
deleted file mode 100644
index 29835a49f106..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/router_gemm/test_router_gemm_integration_invariants.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# 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.
-"""Router-producer invariants of the target's MoE and layers (router_gemm, post_attention_hop).
-
-* one ``is_enabled("router_gemm")`` gate, at ``GlmMoE`` adoption, and the router GEMM is
-  None (the native gate) unless adoption attaches it;
-* ``GlmMoE.compute_routed_output`` has exactly one router_gemm router call and one native
-  gate call; router_gemm runs only when no logits were passed in, the gate only when neither
-  the caller (post_attention_hop) nor router_gemm produced logits, in that order;
-* the post_attention_hop's logits reach ``compute_routed_output`` only through
-  ``router_logits`` and only on the tensor-parallel path.
-"""
-
-import ast
-import inspect
-import textwrap
-
-import pytest
-
-modeling = pytest.importorskip(
-    "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers"
-)
-
-
-def _function(cls, name):
-    source = textwrap.dedent(inspect.getsource(getattr(cls, name)))
-    return ast.parse(source).body[0]
-
-
-def _self_calls(node, attr):
-    return [
-        n
-        for n in ast.walk(node)
-        if isinstance(n, ast.Call)
-        and isinstance(n.func, ast.Attribute)
-        and n.func.attr == attr
-        and isinstance(n.func.value, ast.Name)
-        and n.func.value.id == "self"
-    ]
-
-
-def _enclosing_ifs(func, target):
-    """Tests of the ``if`` statements that contain ``target`` (outermost first)."""
-    tests = []
-
-    def visit(node, stack):
-        if node is target:
-            tests.extend(stack)
-            return True
-        children = ast.iter_child_nodes(node)
-        if isinstance(node, ast.If):
-            for child in node.body:
-                if visit(child, stack + [node.test]):
-                    return True
-            for child in node.orelse:
-                if visit(child, stack):
-                    return True
-            return False
-        return any(visit(child, stack) for child in children)
-
-    visit(func, [])
-    return [ast.unparse(t) for t in tests]
-
-
-@pytest.fixture(scope="module")
-def integrated():
-    return inspect.getsource(modeling)
-
-
-def test_single_router_gemm_gate(integrated):
-    assert integrated.count('is_enabled("router_gemm")') == 1
-    adopt = inspect.getsource(modeling.GlmMoE.adopt)
-    assert 'is_enabled("router_gemm")' in adopt
-    assert modeling.GlmMoE.router_gemm is None
-
-
-def test_router_producer_order(integrated):
-    func = _function(modeling.GlmMoE, "compute_routed_output")
-    args = [a.arg for a in func.args.args]
-    assert "router_logits" in args
-    defaults = dict(zip(args[-len(func.args.defaults) :], func.args.defaults))
-    assert ast.unparse(defaults["router_logits"]) == "None"
-
-    router_gemm_calls = _self_calls(func, "router_gemm")
-    gate_calls = _self_calls(func, "gate")
-    assert len(router_gemm_calls) == 1 and len(gate_calls) == 1
-    assert router_gemm_calls[0].lineno < gate_calls[0].lineno
-
-    router_gemm_guard = " and ".join(_enclosing_ifs(func, router_gemm_calls[0]))
-    assert "router_logits is None" in router_gemm_guard
-    assert "self.router_gemm is not None" in router_gemm_guard
-    gate_guard = _enclosing_ifs(func, gate_calls[0])
-    assert gate_guard == ["router_logits is None"], gate_guard
-
-    # router_logits is assigned only by the two producers.
-    assigned = [
-        n.value
-        for n in ast.walk(func)
-        if isinstance(n, ast.Assign)
-        and any(isinstance(t, ast.Name) and t.id == "router_logits" for t in n.targets)
-    ]
-    assert sorted(ast.unparse(v) for v in assigned) == sorted(
-        [
-            ast.unparse(router_gemm_calls[0]),
-            ast.unparse(gate_calls[0]),
-        ]
-    )
-
-
-def test_attn_hop_logits_only_through_the_argument(integrated):
-    forward = inspect.getsource(modeling.GlmMoE.forward)
-    assert "assert router_logits is None or not self.use_dp" in forward
-    for cls, name in (
-        (modeling.GlmDecoderLayer, "forward_MoE"),
-        (modeling.GlmMtpLayer, "forward"),
-    ):
-        func = _function(cls, name)
-        hop_calls = _self_calls(func, "_post_attn_hop")
-        assert len(hop_calls) == 1, (cls.__name__, name)
-        guard = " and ".join(_enclosing_ifs(func, hop_calls[0]))
-        assert "PRE_MOE_FUSION" in guard and "_can_use_post_attn_hop" in guard
-        source = ast.unparse(func)
-        assert "router_logits=router_logits" in source
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py
deleted file mode 100644
index 0f9203385777..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_fixtures.py
+++ /dev/null
@@ -1,104 +0,0 @@
-# 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.
-"""Real GLM-5.2 shared-expert fixtures for the shared_expert_prefetch tests.
-
-* ``build_shared_mlp(layer, tp_rank)``: the served BF16 shared expert of one
-  MoE layer (one TP8 rank shard), built through ``GatedMLP``'s constructor
-  with the served model configuration and loaded from the checkpoint's
-  ``shared_experts.{gate,up,down}_proj`` tensors through its own loaders (the
-  rank slicing is the module's).
-* ``load_real_shared(mlp, layer)``: the same loading into an existing module
-  (the chain bench's MoE modules carry random shared-expert weights).
-* ``reference(mlp, x)``: the shared expert with the native rounding points and
-  exact (FP64) GEMM sums: ``h = bf16(x attention_output_into_hop^T)``, ``a`` = the native Triton
-  SiLU-and-mul of ``h``, ``y = bf16(a W2^T)``.
-"""
-
-import copy
-import os
-import sys
-
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
-import routed_moe_fixtures as fx  # noqa: E402
-
-TARGET_LAYERS = tuple(range(3, 78))
-
-
-def _shared_tensors(layer: int):
-    rd = fx._Reader()
-    pre = f"model.layers.{layer}.mlp.shared_experts"
-    return {
-        p: rd.full(f"{pre}.{p}.weight").to(torch.bfloat16)
-        for p in ("gate_proj", "up_proj", "down_proj")
-    }
-
-
-def load_real_shared(mlp, layer: int) -> None:
-    w = _shared_tensors(layer)
-    dev = mlp.gate_up_proj.weight.device
-    mlp.gate_up_proj.load_weights([{"weight": w["gate_proj"]}, {"weight": w["up_proj"]}])
-    mlp.down_proj.load_weights([{"weight": w["down_proj"]}])
-    fx._walk_post_load(mlp)
-    assert mlp.gate_up_proj.weight.device == dev
-
-
-def build_shared_mlp(layer: int, tp_rank: int):
-    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-
-    mc = fx._served_model_config(layer, tp_rank)
-    smc = copy.copy(mc)
-    smc.quant_config = QuantConfig()
-    mlp = GatedMLP(
-        hidden_size=fx.HIDDEN,
-        intermediate_size=fx.INTER,
-        bias=False,
-        dtype=torch.bfloat16,
-        config=smc,
-        overridden_tp_size=fx.TP_SIZE,
-        reduce_output=False,
-    ).cuda()
-    load_real_shared(mlp, layer)
-    return mlp
-
-
-def reference(mlp, x: torch.Tensor) -> torch.Tensor:
-    """FP64 GEMM sums with the native BF16 rounding points (see the module docstring)."""
-    w1 = mlp.gate_up_proj.weight.double()
-    w2 = mlp.down_proj.weight.double()
-    h = (x.double() @ w1.t()).to(torch.bfloat16)
-    a = torch.ops.trtllm.silu_and_mul(h)
-    return (a.double() @ w2.t()).to(torch.bfloat16)
-
-
-def native_in_batch(mlp, x: torch.Tensor, batch: int = 64) -> torch.Tensor:
-    """The native module on ``x`` placed in a ``batch``-row input (another cuBLAS
-    algorithm than at R rows): native algorithm-to-algorithm spread."""
-    big = torch.zeros(batch, x.shape[1], dtype=x.dtype, device=x.device)
-    big[: x.shape[0]] = x
-    return mlp(big)[: x.shape[0]].contiguous()
-
-
-def ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """|ordinal(a) - ordinal(b)| of BF16 values (sign-magnitude to ordinal)."""
-
-    def ordinal(t):
-        i = t.contiguous().view(torch.int16).to(torch.int32)
-        return torch.where(i < 0, -(i & 0x7FFF), i)
-
-    return (ordinal(a) - ordinal(b)).abs()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py
deleted file mode 100644
index 0299068cdba8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se.py
+++ /dev/null
@@ -1,251 +0,0 @@
-# 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.
-"""Host side of shared_expert_prefetch's two-kernel BF16 shared expert
-(``shared_prefetch_se_kernel.py``), an alternative.
-
-It is not part of package shared_expert_prefetch: in the two-host chain it ends later than the
-native shared expert beside the routed kernel,
-so the served path keeps the native ``GatedMLP``. The chain bench's ``se7``
-arms and ``test_shared_prefetch_shared.py`` run it from here.
-
-``shared_output(mlp, x)`` returns the shared expert's BF16 ``[R, 6144]``
-output for a served ``GatedMLP`` (the BF16 shared expert of a GLM-5.2 MoE
-layer, one TP8 rank shard), or ``None`` when the call is not eligible; the
-decision depends only on host-side state (module configuration, shapes,
-dtypes, device) and is taken before any device work, so ``None`` leaves the
-caller on the native ``GatedMLP`` path with nothing launched.
-
-The two kernels run on the current stream (the MoE's auxiliary stream in the
-served graph) with programmatic dependent launch. Numerical variant
-``se-reorder`` (see ``kernel_shared``). Side selectors (tests and
-measurements): ``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_UP_CTAS`` (up-kernel CTAs, 32 or 16),
-``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS`` (down-kernel CTAs, 48 or 24),
-``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_L2_PREFETCH`` (L2 prefetch bits, default 3),
-``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_SILU_VARIANT`` (SiLU instruction variant, tests only),
-``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS`` (1: weights staged in shared memory by bulk
-copies before the dependency wait, bit-identical outputs; default 0),
-``TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_TRIGGER`` (``end`` or ``early``: where the up kernel
-releases the down kernel; default ``end``).
-"""
-
-import os
-import threading
-from typing import Dict, Optional, Tuple
-
-import torch
-
-FEATURE = "shared_expert_prefetch"
-HIDDEN = 6144
-INTER = 256
-MAX_ROWS = 8
-
-_LOCK = threading.Lock()
-_COMPILED: Dict[tuple, tuple] = {}
-
-
-def _knobs() -> tuple:
-    import shared_prefetch_se_ptx as ptx
-
-    up = int(
-        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_UP_CTAS") or 32
-    )
-    down = int(
-        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS") or 48
-    )
-    pf = int(
-        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_L2_PREFETCH") or 3
-    )
-    silu = int(
-        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_SILU_VARIANT")
-        or ptx.SILU_DEFAULT
-    )
-    stage = (
-        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS")
-        or "0"
-    ) == "1"
-    trig = (
-        os.environ.get("TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_TRIGGER") or "end"
-    )
-    return up, down, pf, silu, stage, trig
-
-
-def _compiled(key: tuple):
-    with _LOCK:
-        fns = _COMPILED.get(key)
-        if fns is not None:
-            return fns
-        import cutlass
-        import cutlass.cute as cute
-        from cuda.bindings import driver as cuda
-        from shared_prefetch_se_kernel import (
-            PF_W2_OWN,
-            SharedExpertTwoKernelDown,
-            SharedExpertTwoKernelUp,
-        )
-
-        up, down, pf, silu, stage, trig = key
-        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-        z = cutlass.Int64(0)
-        up_fn = cute.compile(
-            SharedExpertTwoKernelUp(up, prefetch=pf & 3, silu=silu, stage=stage, trigger=trig).run,
-            z,
-            z,
-            z,
-            z,
-            cutlass.Int32(1),
-            stream,
-        )
-        down_fn = cute.compile(
-            SharedExpertTwoKernelDown(down, prefetch=pf & PF_W2_OWN, stage=stage).run,
-            z,
-            z,
-            z,
-            cutlass.Int32(1),
-            stream,
-        )
-        _COMPILED[key] = (up_fn, down_fn)
-        return up_fn, down_fn
-
-
-def eligibility(mlp, x) -> Tuple[bool, str]:
-    """Host-only check that the shared_expert_prefetch kernels compute ``mlp(x)`` (rank-identical inputs)."""
-    import torch.nn.functional as F
-
-    if not isinstance(x, torch.Tensor):
-        return False, "input is not a plain tensor"
-    if x.dtype != torch.bfloat16 or x.dim() != 2 or x.shape[1] != HIDDEN:
-        return False, f"input {x.dtype} {tuple(x.shape)}"
-    rows = x.shape[0]
-    if not 1 <= rows <= MAX_ROWS:
-        return False, f"rows {rows} outside 1..{MAX_ROWS}"
-    if not x.is_cuda or not x.is_contiguous() or x.data_ptr() % 16:
-        return False, "input is not a contiguous 16-byte aligned device tensor"
-    if torch.cuda.get_device_capability(x.device)[0] != 10:
-        return False, "not an SM100-family device"
-    if type(mlp).__name__ != "GatedMLP" or getattr(mlp, "split_gate_up", False):
-        return False, "not a fused gate/up GatedMLP"
-    if mlp.activation is not F.silu:
-        return False, "activation is not SiLU"
-    if (
-        mlp.swiglu_limit is not None
-        or mlp.swiglu_alpha not in (None, 1.0)
-        or mlp.swiglu_beta not in (None, 0.0)
-    ):
-        return False, "SwiGLU limit/alpha/beta"
-    gu, dn = mlp.gate_up_proj, mlp.down_proj
-    for name, lin, shape in (("gate_up", gu, (2 * INTER, HIDDEN)), ("down", dn, (HIDDEN, INTER))):
-        w = getattr(lin, "weight", None)
-        if (
-            w is None
-            or w.dtype != torch.bfloat16
-            or tuple(w.shape) != shape
-            or not w.is_contiguous()
-        ):
-            return False, f"{name} weight is not a contiguous BF16 {shape}"
-        if w.device != x.device or w.data_ptr() % 16:
-            return False, f"{name} weight is not a 16-byte aligned tensor on the input's device"
-        if getattr(lin, "bias", None) is not None:
-            return False, f"{name} has a bias"
-        if getattr(lin, "has_any_quant", False):
-            return False, f"{name} is quantized"
-    if getattr(dn, "reduce_output", False) and getattr(dn, "tp_size", 1) > 1:
-        return False, "down projection reduces across ranks"
-    for probe in (
-        "_can_fuse_gate_up_swiglu_fp4out",
-        "_can_fuse_gate_up_swiglu",
-        "_can_fuse_swiglu_fp8_quant",
-    ):
-        fn = getattr(mlp, probe, None)
-        if fn is not None and fn():
-            return False, f"native path is {probe[5:]}"
-    return True, ""
-
-
-class SharedExpertTwoKernel:
-    """The shared_expert_prefetch shared expert of one ``GatedMLP`` (one per module: its scratch is private)."""
-
-    def __init__(self, mlp):
-        self.w1 = mlp.gate_up_proj.weight
-        self.w2 = mlp.down_proj.weight
-        self.a = torch.zeros(MAX_ROWS, INTER, dtype=torch.bfloat16, device=self.w1.device)
-
-    @staticmethod
-    def ready_for_capture() -> bool:
-        return _knobs() in _COMPILED
-
-    def __call__(self, x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
-        import cutlass
-        from cuda.bindings import driver as cuda
-
-        up_fn, down_fn = _compiled(_knobs())
-        rows = x.shape[0]
-        if out is None:
-            out = torch.empty(rows, HIDDEN, dtype=torch.bfloat16, device=x.device)
-        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-        n = cutlass.Int32(rows)
-        up_fn(
-            cutlass.Int64(x.data_ptr()),
-            cutlass.Int64(self.w1.data_ptr()),
-            cutlass.Int64(self.w2.data_ptr()),
-            cutlass.Int64(self.a.data_ptr()),
-            n,
-            stream,
-        )
-        down_fn(
-            cutlass.Int64(self.a.data_ptr()),
-            cutlass.Int64(self.w2.data_ptr()),
-            cutlass.Int64(out.data_ptr()),
-            n,
-            stream,
-        )
-        return out
-
-
-def shared_eligible(mlp, x) -> bool:
-    """Whether ``shared_output(mlp, x)`` engages (host-only, same decision)."""
-    ok, _ = eligibility(mlp, x)
-    if ok and torch.cuda.is_current_stream_capturing():
-        runner = getattr(mlp, "_shared_expert_two_kernel", None)
-        ok = runner is not None and runner.ready_for_capture()
-    return ok
-
-
-def shared_output(mlp, x) -> Optional[torch.Tensor]:
-    """The shared_expert_prefetch shared expert's output for ``mlp(x)``, or ``None`` (native path; nothing launched)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    ok, reason = eligibility(mlp, x)
-    runner = getattr(mlp, "_shared_expert_two_kernel", None) if ok else None
-    if (
-        ok
-        and runner is not None
-        and (runner.w1 is not mlp.gate_up_proj.weight or runner.w2 is not mlp.down_proj.weight)
-    ):
-        runner = None
-    if (
-        ok
-        and torch.cuda.is_current_stream_capturing()
-        and (runner is None or not runner.ready_for_capture())
-    ):
-        ok, reason = False, "not warmed up before CUDA graph capture"
-    dispatch.log_capability(
-        FEATURE, f"shared_expert rows={x.shape[0] if ok else '-'} {reason or 'engaged'}", ok, reason
-    )
-    if not ok:
-        return None
-    if runner is None:
-        runner = SharedExpertTwoKernel(mlp)
-        mlp._shared_expert_two_kernel = runner
-    return runner(x)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py
deleted file mode 100644
index 8aac4f631df7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_kernel.py
+++ /dev/null
@@ -1,453 +0,0 @@
-# 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.
-"""BF16 shared expert of a GLM-5.2 MoE layer in two small kernels (CuTe DSL, SM100 family); an
-alternative to the native shared expert, used by ``shared_prefetch_se.py``.
-
-The per-rank shared expert is ``y = down(silu(gate(x)) * up(x))`` with
-``gate_up`` ``[512, 6144]`` (rows 0..255 gate, 256..511 up) and ``down``
-``[6144, 256]``, all BF16, for ``R <= 8`` rows. Rounding points are the
-native ones: gate and up rounded to BF16 after an FP32 accumulation, the
-SiLU-and-mul in FP32 from those BF16 values (the native Triton instruction
-sequence, ``ptx.silu_mul_triton``) rounded to BF16, the down projection
-accumulated in FP32 and rounded to BF16. Only the FP32 summation order inside
-each GEMM differs from cuBLAS (numerical variant ``se-reorder``).
-
-``SharedExpertTwoKernelUp`` (``ctas`` CTAs of 8 warps): CTA ``c`` owns the
-intermediate columns ``[c * cols, (c + 1) * cols)``; its gate and up rows form
-``cols / 8`` tiles of 16 rows (8 gate rows over the 8 up rows of the same
-columns), so the ``mma.sync`` m16n8k16 accumulator of a thread holds the gate
-and the up value of one column for two tokens. The 8 warps split K = 6144
-(768 each) and their FP32 partial tiles are summed in shared memory in warp
-order 0..7. It writes ``a`` (BF16 ``[R, 256]``, a per-module scratch buffer).
-
-``SharedExpertTwoKernelDown`` (``ctas`` CTAs of 8 warps): CTA ``c`` owns output
-features ``[c * F, (c + 1) * F)``, ``F = 6144 / ctas``; each warp computes
-16-feature tiles over the full K = 256 from ``a`` and writes ``y`` (BF16
-``[R, 6144]``).
-
-Both kernels read every weight with 16-byte non-coherent loads straight into
-``mma.sync`` fragments: thread ``(g, t)`` loads 8 consecutive K elements of
-its rows per 32-wide K chunk and the activation row ``g`` at the same 8 K
-positions, so A and B use the same permutation of K inside the chunk.
-Activation columns of tokens ``>= R`` read row ``R - 1`` and their outputs are
-never stored (MMA columns are independent).
-
-With ``stage`` the kernels read the same weight bytes from shared memory
-instead: before ``griddepcontrol.wait`` each warp bulk-copies
-(``cp.async.bulk``, completion on its own mbarrier) the weight rows it reads
-into a padded shared-memory image (row pitch + 16 bytes, so that the 16-byte
-fragment loads of a warp are bank-conflict free), waits for its mbarrier after
-the dependency wait, and runs the same ``mma.sync`` sequence on ``ld.shared``
-fragments. The arithmetic and its order are unchanged, so the outputs are
-bit-identical to the unstaged kernels. The up kernel stages its 16 rows
-(192 KB, ``ctas = 32`` only); the down kernel its ``F`` rows (64 KB at 48
-CTAs, 128 KB at 24). With ``trigger = "early"`` the up kernel releases its
-dependent (the down kernel) right after its dependency wait, so that the down
-kernel can stage its rows while the up kernel computes.
-
-Memory order and progress (programmatic dependent launch as served):
-
-* Before ``griddepcontrol.wait`` the up kernel reads nothing and writes
-  nothing; with ``prefetch`` it issues ``cp.async.bulk.prefetch.L2`` of its
-  own gate/up rows (bit 1) and of one 1/ctas slice of ``down`` (bit 2), all
-  immutable weights. The down kernel likewise prefetches its own ``down``
-  rows (bit 4).
-* After the wait, the up kernel reads ``x`` (written by the pre-MoE hop,
-  plain coherent loads) and writes ``a``; the down kernel reads ``a`` (written
-  by the up kernel, complete when its wait returns) and writes ``y``. ``a``'s
-  previous reader is the down kernel of the module's previous call, complete
-  before this call's predecessors ran; ``y`` is a fresh output tensor.
-* With ``stage``, the pre-wait bulk copies read only immutable weights and
-  write only the CTA's own shared memory; ``griddepcontrol.wait`` still
-  precedes every read of ``x`` or ``a``. The early trigger changes only when
-  the down kernel is scheduled: its own wait returns after the up kernel has
-  completed and its writes of ``a`` are visible.
-* No CTA waits for another CTA: the only cross-CTA dependency is the kernel
-  boundary between the two launches, so progress does not depend on
-  co-residency.
-"""
-
-import cutlass
-import cutlass.cute as cute
-import cutlass.utils as utils
-import shared_prefetch_se_ptx as ptx
-from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-    ptx as wptx,
-)
-
-HIDDEN = 6144
-INTER = 256  # per-rank shared intermediate size (2048 / TP8)
-NTOK = 8
-THREADS = 256
-WARPS = THREADS // 32
-W1_ROW_BYTES = HIDDEN * 2
-W2_ROW_BYTES = INTER * 2
-PF_W1 = 1
-PF_W2_SLICE = 2
-PF_W2_OWN = 4
-PAD = 16  # shared-memory row padding of the staged weight images
-ROWP1 = W1_ROW_BYTES + PAD
-ROWP2 = W2_ROW_BYTES + PAD
-SLICE1 = HIDDEN // WARPS * 2  # bytes of one gate/up row in one warp's K range
-
-
-class SharedExpertTwoKernelUp:
-    """gate/up GEMM + SiLU-and-mul (see the module docstring)."""
-
-    def __init__(
-        self,
-        ctas: int = 32,
-        prefetch: int = PF_W1 | PF_W2_SLICE,
-        unroll: int = 0,
-        silu: int = ptx.SILU_DEFAULT,
-        carveout: int = 100,
-        stage: bool = False,
-        trigger: str = "end",
-    ):
-        assert ctas in (16, 32), ctas
-        assert not stage or ctas == 32, "the staged up kernel holds 16 rows per CTA"
-        assert trigger in ("end", "early"), trigger
-        self.stage = stage
-        self.trigger = trigger
-        self.ctas = ctas
-        self.cols = INTER // ctas
-        self.tiles = self.cols // 8
-        self.chunks = HIDDEN // (WARPS * 32)
-        unroll = unroll or 8 // self.tiles
-        assert self.chunks % unroll == 0
-        self.unroll = unroll
-        self.prefetch = prefetch
-        self.silu = silu
-        self.carveout = carveout
-
-    @cute.jit
-    def run(
-        self,
-        x_addr: Int64,
-        w1_addr: Int64,
-        w2_addr: Int64,
-        a_addr: Int64,
-        num_tokens: Int32,
-        stream,
-    ):
-        @cute.struct
-        class SharedStorage:
-            part: cute.struct.Align[
-                cute.struct.MemRange[cutlass.Float32, WARPS * self.tiles * 128], 16
-            ]
-
-        @cute.struct
-        class StagedStorage:
-            part: cute.struct.Align[
-                cute.struct.MemRange[cutlass.Float32, WARPS * self.tiles * 128], 16
-            ]
-            mbar: cute.struct.MemRange[cutlass.Int64, WARPS]
-            wsm: cute.struct.Align[
-                cute.struct.MemRange[cutlass.Int64, 2 * self.cols * ROWP1 // 8], 128
-            ]
-
-        if cutlass.const_expr(self.stage):
-            self.shared_storage = StagedStorage
-        else:
-            self.shared_storage = SharedStorage
-        self.kernel(x_addr, w1_addr, w2_addr, a_addr, num_tokens).launch(
-            grid=[self.ctas, 1, 1],
-            block=[THREADS, 1, 1],
-            smem=self.shared_storage.size_in_bytes(),
-            stream=stream,
-            use_pdl=True,
-            preferred_smem_carveout=self.carveout,
-        )
-
-    @cute.kernel
-    def kernel(
-        self, x_addr: Int64, w1_addr: Int64, w2_addr: Int64, a_addr: Int64, num_tokens: Int32
-    ):
-        tidx, _, _ = cute.arch.thread_idx()
-        bidx, _, _ = cute.arch.block_idx()
-        warp = Int32(tidx) // 32
-        lane = Int32(tidx) % 32
-        g = lane // 4
-        t = lane % 4
-        c = Int32(bidx)
-        smem = utils.SmemAllocator()
-        storage = smem.allocate(self.shared_storage)
-        part = Int32(storage.part.data_ptr().toint())
-        if cutlass.const_expr(self.stage):
-            # Pre-wait: each warp bulk-copies its K range of the CTA's 2 * cols weight rows (immutable) into the
-            # padded image; row r < cols is gate row c * cols + r, row cols + r is up row INTER + c * cols + r.
-            mbar = storage.mbar.data_ptr()
-            wsm = Int32(storage.wsm.data_ptr().toint())
-            if tidx == 0:
-                for w in cutlass.range_constexpr(WARPS):
-                    cute.arch.mbarrier_init(mbar + w, 1)
-            cute.arch.mbarrier_init_fence()
-            cute.arch.sync_threads()
-            if lane == 0:
-                bar = Int32((mbar + warp).toint())
-                wptx.mbar_arrive_expect_tx(bar, Int32(2 * self.cols * SLICE1))
-                for r in cutlass.range_constexpr(2 * self.cols):
-                    src_row = c * self.cols + r
-                    if cutlass.const_expr(r >= self.cols):
-                        src_row = INTER + c * self.cols + (r - self.cols)
-                    wptx.bulk_g2s(
-                        wsm + r * ROWP1 + warp * SLICE1,
-                        w1_addr + Int64(src_row) * W1_ROW_BYTES + Int64(warp * SLICE1),
-                        Int32(SLICE1),
-                        bar,
-                    )
-
-        if cutlass.const_expr(self.prefetch & (PF_W1 | PF_W2_SLICE)):
-            if tidx == 0:
-                if cutlass.const_expr(self.prefetch & PF_W1):
-                    wptx.prefetch_l2_bulk(
-                        w1_addr + Int64(c * self.cols) * W1_ROW_BYTES,
-                        Int32(self.cols * W1_ROW_BYTES),
-                    )
-                    wptx.prefetch_l2_bulk(
-                        w1_addr + Int64(INTER + c * self.cols) * W1_ROW_BYTES,
-                        Int32(self.cols * W1_ROW_BYTES),
-                    )
-                if cutlass.const_expr(self.prefetch & PF_W2_SLICE):
-                    rows2 = HIDDEN // self.ctas
-                    wptx.prefetch_l2_bulk(
-                        w2_addr + Int64(c * rows2) * W2_ROW_BYTES, Int32(rows2 * W2_ROW_BYTES)
-                    )
-        cute.arch.griddepcontrol_wait()
-        if cutlass.const_expr(self.trigger == "early"):
-            cute.arch.griddepcontrol_launch_dependents()
-
-        xrow = g
-        if xrow >= num_tokens:
-            xrow = num_tokens - 1
-        xbase = x_addr + Int64(xrow * HIDDEN + warp * (HIDDEN // WARPS) + 8 * t) * 2
-        wbase = w1_addr + Int64((c * self.cols + g) * HIDDEN + warp * (HIDDEN // WARPS) + 8 * t) * 2
-        sbase = Int32(0)
-        if cutlass.const_expr(self.stage):
-            sbase = wsm + g * ROWP1 + warp * SLICE1 + 16 * t
-            cute.arch.mbarrier_wait(mbar + warp, 0)
-        acc = cute.make_rmem_tensor((self.tiles * 4,), Float32)
-        for i in cutlass.range_constexpr(self.tiles * 4):
-            acc[i] = Float32(0.0)
-        xv = cute.make_rmem_tensor((self.unroll * 4,), Uint32)
-        wv = cute.make_rmem_tensor((self.unroll * self.tiles * 8,), Uint32)
-        q0 = Int32(0)
-        while q0 < self.chunks:
-            for u in cutlass.range_constexpr(self.unroll):
-                off = Int64((q0 + u) * 64)
-                v = wptx.ld_v4_u32(xbase + off)
-                for i in cutlass.range_constexpr(4):
-                    xv[u * 4 + i] = v[i]
-                for j in cutlass.range_constexpr(self.tiles):
-                    if cutlass.const_expr(self.stage):
-                        soff = Int32((q0 + u) * 64)
-                        rg = wptx.lds_v4_u32(sbase + j * 8 * ROWP1 + soff)
-                        ru = wptx.lds_v4_u32(sbase + (self.cols + j * 8) * ROWP1 + soff)
-                    else:
-                        rg = wptx.ld_nc_v4_u32(wbase + Int64(j * 8 * W1_ROW_BYTES) + off)
-                        ru = wptx.ld_nc_v4_u32(wbase + Int64((INTER + j * 8) * W1_ROW_BYTES) + off)
-                    for i in cutlass.range_constexpr(4):
-                        wv[(u * self.tiles + j) * 8 + i] = rg[i]
-                        wv[(u * self.tiles + j) * 8 + 4 + i] = ru[i]
-            for u in cutlass.range_constexpr(self.unroll):
-                for j in cutlass.range_constexpr(self.tiles):
-                    b = (u * self.tiles + j) * 8
-                    for h in cutlass.range_constexpr(2):
-                        d = ptx.mma_bf16_16816(
-                            wv[b + 2 * h],
-                            wv[b + 4 + 2 * h],
-                            wv[b + 2 * h + 1],
-                            wv[b + 4 + 2 * h + 1],
-                            xv[u * 4 + 2 * h],
-                            xv[u * 4 + 2 * h + 1],
-                            acc[j * 4],
-                            acc[j * 4 + 1],
-                            acc[j * 4 + 2],
-                            acc[j * 4 + 3],
-                        )
-                        for i in cutlass.range_constexpr(4):
-                            acc[j * 4 + i] = d[i]
-            q0 = q0 + self.unroll
-
-        for j in cutlass.range_constexpr(self.tiles):
-            for i in cutlass.range_constexpr(4):
-                wptx.sts_u32(
-                    part + (((warp * self.tiles + j) * 32 + lane) * 4 + i) * 4,
-                    wptx.f32_bits(acc[j * 4 + i]),
-                )
-        cute.arch.sync_threads()
-        if warp < self.tiles:
-            j = warp
-            # warp order 0..7, starting from warp 0's partial (no +0.0 term)
-            s = cute.make_rmem_tensor((4,), Float32)
-            for i in cutlass.range_constexpr(4):
-                s[i] = wptx.bits_f32(wptx.lds_u32(part + ((j * 32 + lane) * 4 + i) * 4))
-            for w in cutlass.range_constexpr(WARPS - 1):
-                for i in cutlass.range_constexpr(4):
-                    s[i] = s[i] + wptx.bits_f32(
-                        wptx.lds_u32(part + ((((w + 1) * self.tiles + j) * 32 + lane) * 4 + i) * 4)
-                    )
-            col = c * self.cols + j * 8 + g
-            for h in cutlass.range_constexpr(2):
-                tok = 2 * t + h
-                gate = wptx.bf16_round(s[h])
-                up = wptx.bf16_round(s[2 + h])
-                av = ptx.silu_mul_triton(gate, up, self.silu)
-                if tok < num_tokens:
-                    ptx.st_u16(a_addr + Int64(tok * INTER + col) * 2, wptx.bf16_bits(av))
-        if cutlass.const_expr(self.trigger == "end"):
-            cute.arch.griddepcontrol_launch_dependents()
-
-
-class SharedExpertTwoKernelDown:
-    """down GEMM (see the module docstring)."""
-
-    def __init__(self, ctas: int = 48, prefetch: int = 0, carveout: int = 100, stage: bool = False):
-        assert ctas in (24, 48), ctas
-        self.ctas = ctas
-        self.feats = HIDDEN // ctas
-        self.tpw = self.feats // (WARPS * 16)
-        self.prefetch = prefetch
-        self.carveout = carveout
-        self.stage = stage
-
-    @cute.jit
-    def run(self, a_addr: Int64, w2_addr: Int64, y_addr: Int64, num_tokens: Int32, stream):
-        if cutlass.const_expr(self.stage):
-
-            @cute.struct
-            class StagedStorage:
-                mbar: cute.struct.MemRange[cutlass.Int64, WARPS]
-                wsm: cute.struct.Align[
-                    cute.struct.MemRange[cutlass.Int64, self.feats * ROWP2 // 8], 128
-                ]
-
-            self.shared_storage = StagedStorage
-            self.kernel(a_addr, w2_addr, y_addr, num_tokens).launch(
-                grid=[self.ctas, 1, 1],
-                block=[THREADS, 1, 1],
-                smem=self.shared_storage.size_in_bytes(),
-                stream=stream,
-                use_pdl=True,
-                preferred_smem_carveout=self.carveout,
-            )
-        else:
-            self.kernel(a_addr, w2_addr, y_addr, num_tokens).launch(
-                grid=[self.ctas, 1, 1],
-                block=[THREADS, 1, 1],
-                stream=stream,
-                use_pdl=True,
-                preferred_smem_carveout=self.carveout,
-            )
-
-    @cute.kernel
-    def kernel(self, a_addr: Int64, w2_addr: Int64, y_addr: Int64, num_tokens: Int32):
-        tidx, _, _ = cute.arch.thread_idx()
-        bidx, _, _ = cute.arch.block_idx()
-        warp = Int32(tidx) // 32
-        lane = Int32(tidx) % 32
-        g = lane // 4
-        t = lane % 4
-        c = Int32(bidx)
-        rows_w = 16 * self.tpw
-        if cutlass.const_expr(self.stage):
-            # Pre-wait: each warp bulk-copies its 16 * tpw down rows (immutable) into the padded image.
-            smem = utils.SmemAllocator()
-            storage = smem.allocate(self.shared_storage)
-            mbar = storage.mbar.data_ptr()
-            wsm = Int32(storage.wsm.data_ptr().toint())
-            if tidx == 0:
-                for w in cutlass.range_constexpr(WARPS):
-                    cute.arch.mbarrier_init(mbar + w, 1)
-            cute.arch.mbarrier_init_fence()
-            cute.arch.sync_threads()
-            if lane == 0:
-                bar = Int32((mbar + warp).toint())
-                wptx.mbar_arrive_expect_tx(bar, Int32(rows_w * W2_ROW_BYTES))
-                for r in cutlass.range_constexpr(rows_w):
-                    wptx.bulk_g2s(
-                        wsm + (warp * rows_w + r) * ROWP2,
-                        w2_addr + Int64(c * self.feats + warp * rows_w + r) * W2_ROW_BYTES,
-                        Int32(W2_ROW_BYTES),
-                        bar,
-                    )
-        if cutlass.const_expr(self.prefetch & PF_W2_OWN):
-            if tidx == 0:
-                wptx.prefetch_l2_bulk(
-                    w2_addr + Int64(c * self.feats) * W2_ROW_BYTES, Int32(self.feats * W2_ROW_BYTES)
-                )
-        cute.arch.griddepcontrol_wait()
-
-        arow = g
-        if arow >= num_tokens:
-            arow = num_tokens - 1
-        abase = a_addr + Int64(arow * INTER + 8 * t) * 2
-        f0 = c * self.feats + warp * (16 * self.tpw)
-        wbase = w2_addr + Int64((f0 + g) * INTER + 8 * t) * 2
-        sbase = Int32(0)
-        if cutlass.const_expr(self.stage):
-            sbase = wsm + (warp * rows_w + g) * ROWP2 + 16 * t
-            cute.arch.mbarrier_wait(mbar + warp, 0)
-        acc = cute.make_rmem_tensor((self.tpw * 4,), Float32)
-        for i in cutlass.range_constexpr(self.tpw * 4):
-            acc[i] = Float32(0.0)
-        half = 8 // self.tpw
-        av = cute.make_rmem_tensor((half * 4,), Uint32)
-        wv = cute.make_rmem_tensor((half * self.tpw * 8,), Uint32)
-        for p in cutlass.range_constexpr(self.tpw):
-            for u in cutlass.range_constexpr(half):
-                off = Int64((p * half + u) * 64)
-                v = wptx.ld_v4_u32(abase + off)
-                for i in cutlass.range_constexpr(4):
-                    av[u * 4 + i] = v[i]
-                for j in cutlass.range_constexpr(self.tpw):
-                    if cutlass.const_expr(self.stage):
-                        soff = Int32((p * half + u) * 64)
-                        r0 = wptx.lds_v4_u32(sbase + j * 16 * ROWP2 + soff)
-                        r1 = wptx.lds_v4_u32(sbase + (j * 16 + 8) * ROWP2 + soff)
-                    else:
-                        r0 = wptx.ld_nc_v4_u32(wbase + Int64(j * 16 * W2_ROW_BYTES) + off)
-                        r1 = wptx.ld_nc_v4_u32(wbase + Int64((j * 16 + 8) * W2_ROW_BYTES) + off)
-                    for i in cutlass.range_constexpr(4):
-                        wv[(u * self.tpw + j) * 8 + i] = r0[i]
-                        wv[(u * self.tpw + j) * 8 + 4 + i] = r1[i]
-            for u in cutlass.range_constexpr(half):
-                for j in cutlass.range_constexpr(self.tpw):
-                    b = (u * self.tpw + j) * 8
-                    for h in cutlass.range_constexpr(2):
-                        d = ptx.mma_bf16_16816(
-                            wv[b + 2 * h],
-                            wv[b + 4 + 2 * h],
-                            wv[b + 2 * h + 1],
-                            wv[b + 4 + 2 * h + 1],
-                            av[u * 4 + 2 * h],
-                            av[u * 4 + 2 * h + 1],
-                            acc[j * 4],
-                            acc[j * 4 + 1],
-                            acc[j * 4 + 2],
-                            acc[j * 4 + 3],
-                        )
-                        for i in cutlass.range_constexpr(4):
-                            acc[j * 4 + i] = d[i]
-        for j in cutlass.range_constexpr(self.tpw):
-            for h in cutlass.range_constexpr(2):
-                tok = 2 * t + h
-                if tok < num_tokens:
-                    row = y_addr + Int64(tok * HIDDEN + f0 + j * 16 + g) * 2
-                    ptx.st_u16(row, wptx.bf16_bits(acc[j * 4 + h]))
-                    ptx.st_u16(row + 16, wptx.bf16_bits(acc[j * 4 + 2 + h]))
-        cute.arch.griddepcontrol_launch_dependents()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py
deleted file mode 100644
index 78200b103e7e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/shared_prefetch_se_ptx.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# 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.
-"""Inline-PTX primitives of the shared_expert_prefetch shared-expert kernels
-(``shared_prefetch_se_kernel.py``), in addition to routed_moe's."""
-
-from cutlass._mlir.dialects import llvm
-from cutlass.cutlass_dsl import Float32, Int64, T, Uint32, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.ptx import (  # noqa: E501
-    _asm,
-)
-
-
-@dsl_user_op
-def mma_bf16_16816(
-    a0: Uint32,
-    a1: Uint32,
-    a2: Uint32,
-    a3: Uint32,
-    b0: Uint32,
-    b1: Uint32,
-    c0: Float32,
-    c1: Float32,
-    c2: Float32,
-    c3: Float32,
-    *,
-    loc=None,
-    ip=None,
-):
-    """``mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32``: D = A x B + C.
-
-    A fragment (row-major 16 x 16): ``a0`` = (row g, k 2t..2t+1), ``a1`` =
-    (row g+8, k 2t..), ``a2`` = (row g, k 2t+8..), ``a3`` = (row g+8, k
-    2t+8..); B fragment (16 x 8, column n = g): ``b0`` = (k 2t..2t+1),
-    ``b1`` = (k 2t+8..); C/D: (row g, n 2t / 2t+1), (row g+8, n 2t / 2t+1),
-    with g = lane / 4 and t = lane % 4.
-    """
-    ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in (a0, a1, a2, a3, b0, b1)]
-    ops += [Float32(v).ir_value(loc=loc, ip=ip) for v in (c0, c1, c2, c3)]
-    res = _asm(
-        llvm.StructType.get_literal([T.f32()] * 4),
-        ops,
-        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
-        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
-        loc=loc,
-        ip=ip,
-    )
-    return [Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-@dsl_user_op
-def st_u16(addr: Int64, value: Uint32, *, loc=None, ip=None):
-    """16-bit global store of the low half of ``value``."""
-    _asm(
-        None,
-        [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)],
-        "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.global.b16 [$0], h;\n}",
-        "l,r",
-        loc=loc,
-        ip=ip,
-    )
-
-
-# Candidate instruction sequences of the native Triton SiLU-and-mul for an FP32
-# gate ``$1`` and up ``$2`` (``a * sigmoid(a) * (b + 0.0)`` with ``sigmoid(a) =
-# 1 / (1 + exp(-a))``): the exponential as ``ex2.approx`` of ``-a * log2(e)``
-# (with or without flush-to-zero) and the reciprocal as an IEEE or an
-# approximate division. ``SILU_DEFAULT`` is Triton's lowering of FP32 ``exp``
-# (``ex2.approx.f32``) and FP32 ``/`` (``div.full.f32``); the shared_expert_prefetch tests check
-# every candidate against the Triton kernel over every BF16 gate value.
-_SILU_EX2 = ("ex2.approx.f32", "ex2.approx.ftz.f32")
-_SILU_DIV = ("div.rn.f32", "div.full.f32", "div.approx.f32")
-SILU_VARIANTS = tuple((e, d) for d in _SILU_DIV for e in _SILU_EX2)
-SILU_DEFAULT = 2
-
-
-def _silu_text(variant: int) -> str:
-    ex2, div = SILU_VARIANTS[variant]
-    return (
-        "{\n"
-        ".reg .f32 n, e, d, s, gs, ub;\n"
-        "mul.rn.f32 n, $1, 0fBFB8AA3B;\n"
-        f"{ex2} e, n;\n"
-        "add.rn.f32 d, e, 0f3F800000;\n"
-        f"{div} s, 0f3F800000, d;\n"
-        "mul.rn.f32 gs, $1, s;\n"
-        "add.rn.f32 ub, $2, 0f00000000;\n"
-        "mul.rn.f32 $0, gs, ub;\n"
-        "}"
-    )
-
-
-@dsl_user_op
-def silu_mul_triton(
-    g: Float32, u: Float32, variant: int = SILU_DEFAULT, *, loc=None, ip=None
-) -> Float32:
-    """``g * sigmoid(g) * (u + 0.0)`` in FP32 with one of ``SILU_VARIANTS``."""
-    return Float32(
-        _asm(
-            T.f32(),
-            [Float32(g).ir_value(loc=loc, ip=ip), Float32(u).ir_value(loc=loc, ip=ip)],
-            _silu_text(variant),
-            "=f,f,f",
-            loc=loc,
-            ip=ip,
-        )
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py
deleted file mode 100644
index 35121c45a059..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_adapter.py
+++ /dev/null
@@ -1,251 +0,0 @@
-# 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.
-"""Control flow of the shared_expert_prefetch call site in the target's MoE.
-
-With fake collaborators (no device work), on a decode step, checks that
-
-* the plan runs shared_expert_prefetch with its closure (hop_routing's fork into
-  routed_moe_into_hop's FP32 accumulator), not without it, and records no numerical variant
-  for it;
-* deferred_input_prep's tables cover the shared_expert_prefetch directory (with hop_routing's and
-  routed_moe's) and its side selectors;
-* ``GlmMoE.compute_routed_output`` tries ``shared_expert_prefetch.routed_output_into_hop`` (with the
-  module's shared expert) before hop_routing's fork, only with the hop's routing and
-  only when shared_expert_prefetch is selected; ``None`` falls through to hop_routing's fork; the
-  shared branch stays the native shared expert.
-"""
-
-import os
-import sys
-import types
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
-    dispatch as _dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
-
-sys.path.insert(
-    0,
-    os.path.join(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "..", "shared", "kernels"
-    ),
-)
-
-import plan_testlib  # noqa: E402
-
-F_FEATURES = (
-    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,"
-    "router_gemm,post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,"
-    "rope_kv_append,draft_loop,shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
-    "attention_output_tail_prefetch,attention_output_into_hop,routed_moe_into_hop,"
-    "deferred_norm,dense_mlp,deferred_input_prep,hop_routing,oproj_weight_staging,"
-    "fused_kv_indexer_projection,host_page_retention"
-)
-
-
-@pytest.fixture(autouse=True)
-def decode_step():
-    """The owned call site runs on decode steps."""
-    with _dispatch.step(_dispatch.Phase.DECODE):
-        yield
-
-
-def _selector(monkeypatch, features):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-    plan_testlib.select(dispatch, features, monkeypatch)
-    return dispatch
-
-
-@pytest.fixture
-def shared_prefetch_on(monkeypatch):
-    return _selector(monkeypatch, F_FEATURES + ",shared_expert_prefetch")
-
-
-@pytest.fixture
-def shared_prefetch_off(monkeypatch):
-    return _selector(monkeypatch, F_FEATURES)
-
-
-def test_selector(shared_prefetch_on, monkeypatch):
-    dispatch = shared_prefetch_on
-    assert dispatch.is_enabled("shared_expert_prefetch")
-    assert "shared_expert_prefetch" not in dispatch.NUMERICAL_VARIANTS
-    paths = dispatch.feature_record(dispatch.enabled_features()).split()[0]
-    assert paths.startswith("paths=") and "shared_expert_prefetch" in paths[6:].split(",")
-    closure = (
-        "attention_output_tail",
-        "post_moe_hop",
-        "routed_moe",
-        "post_attention_hop",
-        "attention_output_into_hop",
-        "routed_moe_into_hop",
-        "hop_routing",
-    )
-    assert set(dispatch.REQUIRED_FEATURES["shared_expert_prefetch"]) == set(closure)
-    _selector(monkeypatch, closure + ("shared_expert_prefetch",))
-    assert dispatch.installed() == frozenset(closure + ("shared_expert_prefetch",))
-    for missing in closure:
-        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: missing})
-        assert not plan.enabled("shared_expert_prefetch"), missing
-
-
-@pytest.fixture
-def restore_registry():
-    import importlib
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.modeling import (  # noqa: E501
-        TARGET_PATHS,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    for feature in ("hop_routing", "routed_moe", "shared_expert_prefetch"):
-        importlib.import_module(TARGET_PATHS.kernel_modules[feature])
-        dispatch.get_impl(feature, "routed_output_into_hop")
-    saved = {k: dict(v) for k, v in dispatch._REGISTRY.items()}
-    yield
-    dispatch._REGISTRY.clear()
-    dispatch._REGISTRY.update(saved)
-
-
-class _Log:
-    def __init__(self):
-        self.calls = []
-
-    def fn(self, name, result):
-        def call(*args, **kwargs):
-            self.calls.append((name, args, kwargs))
-            return result
-
-        return call
-
-    def names(self):
-        return [c[0] for c in self.calls]
-
-
-def _moe(dispatch, log, r7_result, monkeypatch):
-    from tensorrt_llm._torch.utils import EventType
-
-    for feature, result in (
-        ("shared_expert_prefetch", r7_result),
-        ("hop_routing", ["hop_routing_into"]),
-        ("routed_moe", ["routed_moe_into"]),
-    ):
-        dispatch.register(
-            feature, "routed_output_into_hop", log.fn(f"{feature}.routed_output_into_hop", result)
-        )
-
-    def sequential(fn0, fn1, *args, **kwargs):
-        log.calls.append(("parallel", (), {}))
-        return fn0(), fn1()
-
-    monkeypatch.setattr(layers, "maybe_execute_in_parallel", sequential)
-    moe = layers.GlmMoE.__new__(layers.GlmMoE)
-    torch.nn.Module.__init__(moe)
-
-    class _Shared:
-        def __call__(self, x):
-            log.calls.append(("native_shared", (x,), {}))
-            return torch.full((1,), 7.0)
-
-    moe.__dict__.update(
-        use_dp=False,
-        mapping=types.SimpleNamespace(tp_size=8),
-        router_gemm=None,
-        gate=types.SimpleNamespace(e_score_correction_bias=torch.zeros(256)),
-        experts=object(),
-        _routed_into_hop=object(),
-        shared_experts=_Shared(),
-        shared_output_scale=None,
-        shared_experts_use_fp4=False,
-        event_dict={k: None for k in EventType},
-        aux_stream=None,
-    )
-    return moe
-
-
-@pytest.mark.parametrize(
-    "routing_given, r7_result, want_routed, want_out",
-    [
-        (
-            True,
-            ["shared_prefetch_into"],
-            ["shared_expert_prefetch.routed_output_into_hop"],
-            ["shared_prefetch_into"],
-        ),
-        (
-            True,
-            None,
-            ["shared_expert_prefetch.routed_output_into_hop", "hop_routing.routed_output_into_hop"],
-            ["hop_routing_into"],
-        ),
-        (
-            False,
-            ["shared_prefetch_into"],
-            ["routed_moe.routed_output_into_hop"],
-            ["routed_moe_into"],
-        ),
-    ],
-)
-def test_routed_order(
-    shared_prefetch_on,
-    restore_registry,
-    monkeypatch,
-    routing_given,
-    r7_result,
-    want_routed,
-    want_out,
-):
-    log = _Log()
-    moe = _moe(shared_prefetch_on, log, r7_result, monkeypatch)
-    x = torch.zeros(6, 4)
-    logits = torch.zeros(6, 256)
-    routing = torch.zeros(2, 6, 8, dtype=torch.int32) if routing_given else None
-    out = moe.forward(x, do_finalize=False, router_logits=logits, hop_routing_out=routing)
-    names = log.names()
-    assert [n for n in names if n.endswith("routed_output_into_hop")] == want_routed
-    assert names.count("native_shared") == 1
-    assert float(out[0][0]) == 7.0 and out[1:] == want_out
-    for name, args, kwargs in log.calls:
-        if name == "shared_expert_prefetch.routed_output_into_hop":
-            want = (
-                moe.experts,
-                moe.gate.e_score_correction_bias,
-                x,
-                logits,
-                routing,
-                moe._routed_into_hop,
-                moe.shared_experts,
-            )
-            assert not kwargs and len(args) == len(want) and all(a is b for a, b in zip(args, want))
-
-
-def test_not_selected_never_calls_shared_prefetch(
-    shared_prefetch_off, restore_registry, monkeypatch
-):
-    log = _Log()
-    moe = _moe(shared_prefetch_off, log, ["shared_prefetch_into"], monkeypatch)
-    out = moe.forward(
-        torch.zeros(6, 4),
-        do_finalize=False,
-        router_logits=torch.zeros(6, 256),
-        hop_routing_out=torch.zeros(2, 6, 8, dtype=torch.int32),
-    )
-    assert not [n for n in log.names() if n.startswith("shared_expert_prefetch.")]
-    assert float(out[0][0]) == 7.0 and out[1:] == ["hop_routing_into"]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py
deleted file mode 100644
index 2f499ef2e4b5..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_routed.py
+++ /dev/null
@@ -1,454 +0,0 @@
-# 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.
-"""shared_expert_prefetch's routed fork (``SharedPrefetchExpertClusterKernel``: hop_routing's fork
-plus the pre-wait L2 prefetch of the shared expert's weights, on the shared_expert_prefetch grid) on
-one GPU, real GLM-5.2 NVFP4 weights.
-
-Given the same routing, neither the grid nor the prefetch may change any expert row:
-
-* raw FC2 rows and scales of hop_routing's fork at the shared_expert_prefetch grid (the fork's
-  grid set to 4 x shared_expert_prefetch's clusters) equal those at F's grid (24 clusters,
-  raised to ``ceil(R * 8 / 2)`` for R > 6) bit for bit, for every D bucket 8..48 at R = 6 and
-  for R = 1..8 (real windows);
-* shared_expert_prefetch's FP32 accumulator equals the exact sum of routed_moe's weighted rows
-  within the FP32 reorder bound (routed_moe_into_hop's fp32 class), for the same cases and in every
-  test below. The accumulator's atomic adds come in any order, so two runs
-  need not be bit-identical; every check is against that bound;
-* CUDA-graph replay with fresh routing and interleaved row counts (one graph
-  per R, replayed in alternating order);
-* a delayed producer: a kernel that releases its dependents first and writes
-  the routing and the activation only after a delay, so a read before the
-  dependency wait would see stale data;
-* the native shared expert co-running on a second stream (as served) changes
-  no routed result, and its own output stays bit-identical while the routed
-  kernel prefetches its weights;
-* host-only declines.
-
-Routings: constructed expert sets with exactly D distinct experts (fresh
-experts per case), routed through the native routing of logits built for
-them, and real captured windows of layer ``GLM_MOE_DSA_ROUTED_MOE_LAYER``.
-"""
-
-import contextlib
-import os
-import sys
-from unittest import mock
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-_ROOT = os.path.dirname(_HERE)
-sys.path.insert(0, _HERE)
-for sub in ("routed_moe", "hop_routing"):
-    sys.path.insert(0, os.path.join(_ROOT, sub))
-import hop_routing_oracle as orc  # noqa: E402
-import routed_moe_fixtures as fx  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM100-family GPU",
-)
-
-REORDER_REL = 2.0**-20
-WIDE_CLUSTERS = 36
-LAYER = int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_LAYER", "10"))
-F_CLUSTERS = 24
-D_CASES = (8, 12, 16, 20, 23, 24, 25, 26, 28, 30, 31, 32, 33, 35, 36, 37, 40, 44, 48)
-
-
-def _grid(runner, ctas):
-    """``runner``'s expert-cluster grid set to ``ctas`` CTAs inside the block."""
-    return mock.patch.object(runner, "_xc_ctas", lambda *args, **kwargs: ctas)
-
-
-def _clusters(n):
-    """shared_expert_prefetch's cluster count set to ``n`` inside the block (None: unchanged)."""
-    if n is None:
-        return contextlib.nullcontext()
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch import (  # noqa: E501
-        routed_moe,
-    )
-
-    return mock.patch.object(routed_moe, "DEFAULT_CLUSTERS", n)
-
-
-@pytest.fixture(scope="module")
-def layer():
-    torch.manual_seed(0)
-    return fx.build_native_layer(
-        layer=LAYER, tp_rank=int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0"))
-    )
-
-
-class _SHARED_PREFETCH:
-    """shared_expert_prefetch's runner with the layer's real shared expert ``mlp`` (its weights are
-    the prefetch target)."""
-
-    def __init__(self, runner, mlp, se):
-        self.runner, self.mlp, self.se = runner, mlp, se
-
-    def __call__(self, x, routing, red_out):
-        return self.runner(x, routing, red_out, self.se)
-
-    def __getattr__(self, name):
-        return getattr(self.runner, name)
-
-
-@pytest.fixture(scope="module")
-def runners(layer):
-    import shared_prefetch_fixtures as zf
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.hop_routing.routed_moe import (  # noqa: E501
-        HopRoutedMoe,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe import (  # noqa: E501
-        RoutedMoe,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
-        SharedPrefetchRoutedMoe,
-        shared_weights,
-    )
-
-    mlp = zf.build_shared_mlp(LAYER, int(os.environ.get("GLM_MOE_DSA_ROUTED_MOE_RANK", "0")))
-    se, reason = shared_weights(mlp)
-    assert se is not None, reason
-    return (
-        RoutedMoe.from_native(layer.module, layer.bias),
-        HopRoutedMoe(layer.module, layer.bias),
-        _SHARED_PREFETCH(SharedPrefetchRoutedMoe(layer.module, layer.bias), mlp, se),
-    )
-
-
-def _pack(logits, bias):
-    ids, w = orc.native_routing(logits, bias)
-    routing = torch.empty(2, ids.shape[0], 8, dtype=torch.int32, device="cuda")
-    routing[0].copy_(ids)
-    routing[1].copy_(w.float().view(torch.int32))
-    return routing
-
-
-def _constructed(layer, rows, d, seed):
-    sets = fx.expert_sets_with_distinct(rows, d, seed=seed)
-    perm = torch.randperm(
-        fx.NUM_EXPERTS, generator=torch.Generator().manual_seed(1000 + seed)
-    ).tolist()
-    uniq = sorted(set(e for row in sets for e in row))
-    remap = {e: perm[j] for j, e in enumerate(uniq)}
-    sets = [[remap[e] for e in row] for row in sets]
-    lg = fx.logits_for_expert_sets(sets, layer.bias.cpu(), seed=seed).float().cuda().contiguous()
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = (torch.randn(rows, fx.HIDDEN, device="cuda", generator=g) * 0.5).to(torch.bfloat16)
-    routing = _pack(lg, layer.bias)
-    assert len(set(routing[0].flatten().tolist())) == d
-    return x, lg, routing
-
-
-def _real(layer, rows, count, offset=0):
-    win = fx.real_windows(fx.load_real_capture(layer.layer), rows, stride=rows * 7)
-    out = []
-    for x, lg in win[offset : offset + count]:
-        x, lg = x.cuda().contiguous(), lg.float().cuda().contiguous()
-        out.append((x, lg, _pack(lg, layer.bias)))
-    return out
-
-
-def _cases(layer):
-    out = [("D%d" % d, 6) + _constructed(layer, 6, d, seed=d) for d in D_CASES]
-    for rows in range(1, 9):
-        out += [
-            ("real R%d #%d" % (rows, i), rows) + c
-            for i, c in enumerate(_real(layer, rows, 2, offset=rows))
-        ]
-    return out
-
-
-def _shared_prefetch_ctas(runner, rows):
-    return runner._xc_ctas(rows, False, True)
-
-
-def _f_ctas(runner, rows):
-    """F's FP32-accumulator grid for ``rows`` rows: 24 clusters, raised to cover rows x 8 experts in two rounds."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        routed_moe as routed_moe_mod,
-    )
-
-    co = routed_moe_mod._max_active_clusters(runner.device, 4)
-    return 4 * min(co, max(F_CLUSTERS, -(-rows * fx.TOP_K // 2)))
-
-
-def _acc_within_bound(routed_runner, x, lg, acc):
-    """Whether ``acc`` is the exact FP32-class sum of routed_moe's weighted raw rows for ``(x,
-    lg)``: the atomics add the top-k contributions in any order, so the check is the reorder bound
-    against the exact sum, not run-to-run equality."""
-    fc2, scale, idx = routed_runner(x, lg, do_finalize=False)
-    terms = scale.double()[:, :, None] * fc2[idx.long()].double()
-    exact = terms.sum(1)
-    bound = REORDER_REL * terms.abs().sum(1)
-    return bool(((acc.double() - exact).abs() <= bound).all())
-
-
-def test_grid(layer, runners):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch.routed_moe import (  # noqa: E501
-        DEFAULT_CLUSTERS,
-    )
-
-    _, _, shared_expert_prefetch = runners
-    co = __import__(
-        "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe.routed_moe",
-        fromlist=["x"],
-    )
-    clusters = co._max_active_clusters(shared_expert_prefetch.device, 4)
-    print(
-        f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_GRID co-resident 4-CTA clusters {clusters}; "
-        f"shared_expert_prefetch grid R6 {_shared_prefetch_ctas(shared_expert_prefetch, 6)} CTAs",
-        flush=True,
-    )
-    assert _shared_prefetch_ctas(shared_expert_prefetch, 6) == 4 * min(clusters, DEFAULT_CLUSTERS)
-    for want in (28, 32, 36):
-        with _clusters(want):
-            assert _shared_prefetch_ctas(shared_expert_prefetch, 6) == 4 * min(clusters, want)
-            assert _shared_prefetch_ctas(shared_expert_prefetch, 8) == 4 * min(
-                clusters, max(want, 32)
-            )  # two rounds cover 8 x 8 experts
-    assert _shared_prefetch_ctas(shared_expert_prefetch, 8) == 4 * min(
-        clusters, max(DEFAULT_CLUSTERS, 32)
-    )
-    assert shared_expert_prefetch.schedule(6, False, True) == "xc"
-
-
-def test_raw_rows_grid_independent_and_fp32_class(layer, runners):
-    routed_runner, hop_routing, shared_expert_prefetch = runners
-    bad = []
-    for name, rows, x, lg, routing in _cases(layer):
-        with _grid(hop_routing, _f_ctas(shared_expert_prefetch, rows)):
-            raw_f = hop_routing(x, routing, do_finalize=False)
-        with _clusters(WIDE_CLUSTERS):
-            wide = _shared_prefetch_ctas(shared_expert_prefetch, rows)
-        with _grid(hop_routing, wide):
-            raw_7 = hop_routing(x, routing, do_finalize=False)
-        fc2, scale, idx = routed_runner(x, lg, do_finalize=False)
-        torch.cuda.synchronize()
-        same = all(torch.equal(a, b) for a, b in zip(raw_f, raw_7))
-        same_routed_moe = all(torch.equal(a, b) for a, b in zip(raw_f, (fc2, scale, idx)))
-        terms = scale.double()[:, :, None] * fc2[idx.long()].double()
-        exact = terms.sum(1)
-        bound = REORDER_REL * terms.abs().sum(1)
-        ok_acc = True
-        for clusters in (None, WIDE_CLUSTERS):
-            with _clusters(clusters):
-                for _ in range(2):
-                    acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
-                    shared_expert_prefetch(x, routing, red_out=acc)
-                    torch.cuda.synchronize()
-                    ok_acc = ok_acc and bool(((acc.double() - exact).abs() <= bound).all())
-        print(
-            f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_ROUTED {name}: raw F-grid == {wide}-CTA grid {same}, == "
-            f"routed_moe {same_routed_moe}, "
-            f"fp32 within bound (shared_expert_prefetch default / wide) {ok_acc}",
-            flush=True,
-        )
-        if not (same and same_routed_moe and ok_acc):
-            bad.append(name)
-    assert not bad, bad
-
-
-def test_graph_replay_interleaved_rows(layer, runners):
-    routed_runner, _, shared_expert_prefetch = runners
-    graphs, state = {}, {}
-    for rows in (1, 3, 6, 8):
-        cases = _real(layer, rows, 3, offset=1)
-        x = cases[0][0].clone()
-        routing = cases[0][2].clone()
-        acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
-        shared_expert_prefetch(x, routing, red_out=acc)
-        torch.cuda.synchronize()
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            shared_expert_prefetch(x, routing, red_out=acc)
-        graphs[rows] = g
-        state[rows] = (x, routing, acc, cases)
-    for step in range(2):
-        for rows in (1, 3, 6, 8) if step == 0 else (8, 6, 3, 1):
-            x, routing, acc, cases = state[rows]
-            xn, ln, rn = cases[1 + step]
-            x.copy_(xn)
-            routing.copy_(rn)
-            acc.zero_()
-            graphs[rows].replay()
-            torch.cuda.synchronize()
-            assert _acc_within_bound(routed_runner, xn, ln, acc), (rows, step)
-    del graphs
-
-
-def _delay_kernel():
-    import cutlass
-    import cutlass.cute as cute
-    from cutlass.cutlass_dsl import Int32, Int64
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        ptx as wptx,
-    )
-
-    class DelayedCopy:
-        """Releases its dependents at entry, spins ``ns``, then copies ``words`` 16-byte words."""
-
-        @cute.jit
-        def run(self, src: Int64, dst: Int64, words: Int32, ns: Int64, stream):
-            self.kernel(src, dst, words, ns).launch(
-                grid=[1, 1, 1], block=[256, 1, 1], stream=stream, use_pdl=True
-            )
-
-        @cute.kernel
-        def kernel(self, src: Int64, dst: Int64, words: Int32, ns: Int64):
-            tidx, _, _ = cute.arch.thread_idx()
-            cute.arch.griddepcontrol_launch_dependents()
-            t0 = wptx.globaltimer_u64()
-            tn = t0
-            while tn - t0 < ns:
-                tn = wptx.globaltimer_u64()
-            i = Int32(tidx)
-            while i < words:
-                v = wptx.ld_v4_u32(src + Int64(i) * 16)
-                wptx.st_v4_u32(dst + Int64(i) * 16, v[0], v[1], v[2], v[3])
-                i = i + 256
-
-    from cuda.bindings import driver as cuda
-
-    z = cutlass.Int64(0)
-    fn = cute.compile(
-        DelayedCopy().run,
-        z,
-        z,
-        cutlass.Int32(1),
-        z,
-        cuda.CUstream(torch.cuda.current_stream().cuda_stream),
-    )
-
-    def launch(src, dst, ns):
-        assert src.nbytes % 16 == 0 and src.nbytes == dst.nbytes
-        fn(
-            cutlass.Int64(src.data_ptr()),
-            cutlass.Int64(dst.data_ptr()),
-            cutlass.Int32(src.nbytes // 16),
-            cutlass.Int64(ns),
-            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
-        )
-
-    return launch
-
-
-def test_delayed_producer(layer, runners):
-    """Routing and activation written 20 us after the producer releases the routed kernel."""
-    routed_runner, _, shared_expert_prefetch = runners
-    launch = _delay_kernel()
-    os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
-    for name, rows, x, lg, routing in _cases(layer)[::3]:
-        # packed source: [routing | x] so that one delayed copy writes both
-        src = torch.cat([routing.flatten().view(torch.uint8), x.flatten().view(torch.uint8)])
-        pad = (-src.numel()) % 16
-        src = torch.cat([src, torch.zeros(pad, dtype=torch.uint8, device="cuda")])
-        # stale bytes (huge activations, invalid expert ids), so that a read before the dependency wait changes
-        # the result
-        dst = torch.full_like(src, 0x7F)
-        r_view = dst[: routing.nbytes].view(torch.int32).view(2, rows, 8)
-        x_view = (
-            dst[routing.nbytes : routing.nbytes + x.nbytes]
-            .view(torch.bfloat16)
-            .view(rows, fx.HIDDEN)
-        )
-        acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
-        launch(src, dst, 20000)
-        shared_expert_prefetch(x_view, r_view, red_out=acc)
-        torch.cuda.synchronize()
-        assert _acc_within_bound(routed_runner, x, lg, acc), name
-
-
-def test_corun_shared_expert(layer, runners):
-    """shared_expert_prefetch's routed fork and the native shared expert (whose weights it
-    prefetches) on two streams, as served: same results as alone."""
-    routed_runner, _, shared_expert_prefetch = runners
-    mlp = shared_expert_prefetch.mlp
-    side = torch.cuda.Stream()
-    for name, rows, x, lg, routing in _cases(layer)[::2]:
-        want_acc = torch.zeros(rows, fx.HIDDEN, dtype=torch.float32, device="cuda")
-        shared_expert_prefetch(x, routing, red_out=want_acc)
-        want_se = mlp(x)
-        torch.cuda.synchronize()
-        acc = torch.zeros_like(want_acc)
-        side.wait_stream(torch.cuda.current_stream())
-        shared_expert_prefetch(x, routing, red_out=acc)
-        with torch.cuda.stream(side):
-            se = mlp(x)
-        torch.cuda.current_stream().wait_stream(side)
-        torch.cuda.synchronize()
-        assert _acc_within_bound(routed_runner, x, lg, acc), name
-        assert torch.equal(se.view(torch.int16), want_se.view(torch.int16)), name
-
-
-def test_declines_before_device_work(layer):
-    import shared_prefetch_fixtures as zf
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.shared_expert_prefetch import (  # noqa: E501
-        routed_moe as shared_prefetch_mod,
-    )
-
-    class _Hop:
-        _routed_enabled = True
-
-        def __init__(self, mode):
-            self.routed_mode = mode
-
-        def can_run(self, rows):
-            return rows <= 8
-
-    mlp = zf.build_shared_mlp(LAYER, 0)
-    x, lg, routing = _constructed(layer, 6, 30, seed=5)
-    torch.cuda.synchronize()
-    cases = (
-        (_Hop("identity"), routing, mlp),
-        (_Hop("row"), routing, mlp),
-        (_Hop("fp32"), None, mlp),
-        (_Hop("fp32"), routing[:, :5].contiguous(), mlp),
-        (None, routing, mlp),
-        (_Hop("fp32"), routing, None),
-        (_Hop("fp32"), routing, torch.nn.Linear(4, 4)),
-    )
-    for hop, r, shared in cases:
-        before = torch.cuda.memory_allocated()
-        assert (
-            shared_prefetch_mod.routed_output_into_hop(
-                layer.module, layer.bias, x, lg, r, hop, shared
-            )
-            is None
-        )
-        assert torch.cuda.memory_allocated() == before
-
-
-@pytest.mark.parametrize("live", [1, 5, 6])
-def test_padded_rows_do_not_change_live_rows(layer, runners, live):
-    """An 8-row call whose rows past ``live`` are padding: the live rows' FP32 accumulator equals the live-row call's
-    within the reorder bound."""
-    routed_runner, _, shared_expert_prefetch = runners
-    ((x, lg, routing),) = _real(layer, 8, 1, offset=4)
-    full = torch.zeros(8, fx.HIDDEN, dtype=torch.float32, device="cuda")
-    shared_expert_prefetch(x, routing, red_out=full)
-    xl = x[:live].contiguous()
-    rl = _pack(lg[:live].contiguous(), layer.bias)
-    part = torch.zeros(live, fx.HIDDEN, dtype=torch.float32, device="cuda")
-    shared_expert_prefetch(xl, rl, red_out=part)
-    torch.cuda.synchronize()
-    assert _acc_within_bound(routed_runner, xl, lg[:live].contiguous(), part)
-    assert _acc_within_bound(routed_runner, xl, lg[:live].contiguous(), full[:live])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py
deleted file mode 100644
index 315c577296e8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/test_shared_prefetch_shared.py
+++ /dev/null
@@ -1,269 +0,0 @@
-# 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.
-"""Single-GPU tests of shared_expert_prefetch's two-kernel BF16 shared expert
-(``shared_prefetch_se.py``, an alternative to the native shared expert) on real GLM-5.2 weights.
-
-* ``test_silu_variants``: every candidate SiLU-and-mul instruction sequence
-  (``shared_prefetch_se_ptx.SILU_VARIANTS``) against the native Triton ``silu_and_mul`` over
-  all 65,536 BF16 gate values times 24 up values; the default must be
-  bit-identical.
-* ``test_shared_expert``: the shared_expert_prefetch kernels against the native ``GatedMLP`` and
-  the FP64 reference with native rounding points (``shared_prefetch_fixtures.reference``)
-  on real hidden states (captured windows of layers 10, 30, 50, 77, all
-  rows R = 1..8) for the layers in ``GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_LAYERS`` (default
-  3,10,30,50,77; ``all`` = every target layer) and the TP ranks in
-  ``GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_RANKS`` (default 0). Prints the error statistics per layer
-  (``GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SE_STATS``) and checks the frozen error budget of
-  ``se-reorder``.
-* ``test_eligibility``: declines (host-only, nothing launched).
-"""
-
-import json
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "routed_moe"))
-import routed_moe_fixtures as fx  # noqa: E402
-import shared_prefetch_fixtures as zf  # noqa: E402
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
-
-
-def _silu_probe(variant: int):
-    import cutlass
-    import cutlass.cute as cute
-    import shared_prefetch_se_ptx as ptx
-    from cutlass.cutlass_dsl import Int32, Int64
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.kernels.routed_moe import (  # noqa: E501
-        ptx as wptx,
-    )
-
-    class SiluProbe:
-        @cute.jit
-        def run(self, g_addr: Int64, u_addr: Int64, o_addr: Int64, n: Int32, stream):
-            self.kernel(g_addr, u_addr, o_addr, n).launch(
-                grid=[(n + 255) // 256, 1, 1], block=[256, 1, 1], stream=stream
-            )
-
-        @cute.kernel
-        def kernel(self, g_addr: Int64, u_addr: Int64, o_addr: Int64, n: Int32):
-            tidx, _, _ = cute.arch.thread_idx()
-            bidx, _, _ = cute.arch.block_idx()
-            i = Int32(bidx) * 256 + Int32(tidx)
-            if i < n:
-                gv = wptx.bf16x2_to_f32(wptx.ld_nc_u16(g_addr + Int64(i) * 2), False)
-                uv = wptx.bf16x2_to_f32(wptx.ld_nc_u16(u_addr + Int64(i) * 2), False)
-                r = ptx.silu_mul_triton(gv, uv, variant)
-                ptx.st_u16(o_addr + Int64(i) * 2, wptx.bf16_bits(r))
-
-    from cuda.bindings import driver as cuda
-
-    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-    z = cutlass.Int64(0)
-    fn = cute.compile(SiluProbe().run, z, z, z, cutlass.Int32(1), stream)
-
-    def call(g, u):
-        out = torch.empty_like(g)
-        fn(
-            cutlass.Int64(g.data_ptr()),
-            cutlass.Int64(u.data_ptr()),
-            cutlass.Int64(out.data_ptr()),
-            cutlass.Int32(g.numel()),
-            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
-        )
-        return out
-
-    return call
-
-
-def test_silu_variants():
-    import shared_prefetch_se_ptx as ptx
-
-    import tensorrt_llm  # noqa: F401  (registers trtllm ops)
-
-    gates = torch.arange(-32768, 32768, dtype=torch.int32).to(torch.int16).view(torch.bfloat16)
-    gen = torch.Generator().manual_seed(7)
-    ups = [1.0, -1.0, 0.5, 2.0, 3.140625, -0.0078125, 1e-3, 1e-20, 300.0, -6e4] + list(
-        (torch.randn(14, generator=gen) * 4).tolist()
-    )
-    ups = torch.tensor(ups, dtype=torch.float32).to(torch.bfloat16)
-    rows = ups.numel()
-    d = gates.numel()
-    x = torch.empty(rows, 2 * d, dtype=torch.bfloat16)
-    x[:, :d] = gates[None, :]
-    x[:, d:] = ups[:, None]
-    x = x.cuda()
-    native = torch.ops.trtllm.silu_and_mul(x)
-    g = x[:, :d].contiguous().flatten()
-    u = x[:, d:].contiguous().flatten()
-    nat = native.flatten()
-    nat_nan = torch.isnan(nat)
-    results = {}
-    for v in range(len(ptx.SILU_VARIANTS)):
-        out = _silu_probe(v)(g, u)
-        same = (out.view(torch.int16) == nat.view(torch.int16)) | (torch.isnan(out) & nat_nan)
-        results[v] = int((~same).sum().item())
-        print(
-            f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SILU variant {v} {ptx.SILU_VARIANTS[v]}: {results[v]} of "
-            f"{nat.numel()} differ",
-            flush=True,
-        )
-    assert results[ptx.SILU_DEFAULT] == 0, results
-
-
-def _windows():
-    out = []
-    for layer in (10, 30, 50, 77):
-        cap = fx.load_real_capture(layer)
-        for rows in range(1, 9):
-            for x, _ in fx.real_windows(cap, rows, stride=97)[:12]:
-                out.append(x)
-    return out
-
-
-def _layers():
-    raw = os.environ.get("GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_LAYERS", "3,10,30,50,77")
-    return list(zf.TARGET_LAYERS) if raw == "all" else [int(v) for v in raw.split(",")]
-
-
-def _ranks():
-    return [
-        int(v) for v in os.environ.get("GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_RANKS", "0").split(",")
-    ]
-
-
-def _stats(y, ref):
-    d = y.float() - ref.float()
-    nrms = (d.pow(2).mean().sqrt() / ref.float().pow(2).mean().sqrt()).item()
-    ulp = zf.ulp_distance(y, ref)
-    return {
-        "nrms": nrms,
-        "max_ulp": int(ulp.max().item()),
-        "frac_diff": (ulp > 0).float().mean().item(),
-        "max_abs": d.abs().max().item(),
-    }
-
-
-def test_shared_expert():
-    from shared_prefetch_se import shared_output
-
-    wins = [w.cuda() for w in _windows()]
-    failures = []
-    for rank in _ranks():
-        for layer in _layers():
-            mlp = zf.build_shared_mlp(layer, rank)
-            acc = {k: ([], []) for k in ("shared_expert_prefetch", "native", "native64")}
-            repeat_ok = True
-            for x in wins:
-                ref = zf.reference(mlp, x)
-                y7 = shared_output(mlp, x)
-                assert y7 is not None
-                y7b = shared_output(mlp, x)
-                repeat_ok = repeat_ok and torch.equal(y7.view(torch.int16), y7b.view(torch.int16))
-                for k, y in (
-                    ("shared_expert_prefetch", y7),
-                    ("native", mlp(x)),
-                    ("native64", zf.native_in_batch(mlp, x)),
-                ):
-                    acc[k][0].append(y.flatten())
-                    acc[k][1].append(ref.flatten())
-            st = {k: _stats(torch.cat(v[0]), torch.cat(v[1])) for k, v in acc.items()}
-            shared_prefetch_vs_native = _stats(
-                torch.cat(acc["shared_expert_prefetch"][0]), torch.cat(acc["native"][0])
-            )
-            nat_spread = _stats(torch.cat(acc["native64"][0]), torch.cat(acc["native"][0]))
-            row = {
-                "layer": layer,
-                "rank": rank,
-                "windows": len(wins),
-                "deterministic": repeat_ok,
-                "vs_ref": st,
-                "shared_prefetch_vs_native": shared_prefetch_vs_native,
-                "native64_vs_native": nat_spread,
-            }
-            print("GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SE_STATS " + json.dumps(row), flush=True)
-            # Frozen se-reorder budget: against the FP64 reference with native rounding points,
-            # shared_expert_prefetch's normalized RMS error is at most 1.1x the larger of the two
-            # native algorithms' and its largest ulp distance at most theirs + 1; the output is
-            # run-to-run deterministic.
-            nat_nrms = max(st["native"]["nrms"], st["native64"]["nrms"])
-            nat_ulp = max(st["native"]["max_ulp"], st["native64"]["max_ulp"])
-            if not (
-                repeat_ok
-                and st["shared_expert_prefetch"]["nrms"] <= 1.1 * nat_nrms
-                and st["shared_expert_prefetch"]["max_ulp"] <= nat_ulp + 1
-            ):
-                failures.append(row)
-            del mlp
-    assert not failures, failures
-
-
-@pytest.mark.parametrize("trig", ["end", "early"])
-def test_staged_bit_identical(trig, monkeypatch):
-    """Weights staged in shared memory (bulk copies before the dependency wait) give the unstaged outputs bit for
-    bit: same fragments, same ``mma.sync`` order."""
-    from shared_prefetch_se import shared_output
-
-    wins = [w.cuda() for w in _windows()]
-    bad = []
-    for layer in _layers()[:2]:
-        mlp = zf.build_shared_mlp(layer, _ranks()[0])
-        for i, x in enumerate(wins):
-            monkeypatch.delenv(
-                "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS",
-                raising=False,
-            )
-            want = shared_output(mlp, x)
-            for down in ("48", "24"):
-                monkeypatch.setenv(
-                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_STAGE_WEIGHTS", "1"
-                )
-                monkeypatch.setenv(
-                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_TRIGGER", trig
-                )
-                monkeypatch.setenv(
-                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS", down
-                )
-                got = shared_output(mlp, x)
-                monkeypatch.delenv(
-                    "TRTLLM_STAIRCASE_GLM_LL_SHARED_EXPERT_PREFETCH_TWO_KERNEL_DOWN_CTAS"
-                )
-                torch.cuda.synchronize()
-                if not torch.equal(got.view(torch.int16), want.view(torch.int16)):
-                    bad.append((layer, i, x.shape[0], down))
-        del mlp
-    print(f"GLM_MOE_DSA_SHARED_EXPERT_PREFETCH_SE_STAGED trig={trig} mismatches {bad}", flush=True)
-    assert not bad
-
-
-def test_eligibility():
-    from shared_prefetch_se import eligibility
-
-    mlp = zf.build_shared_mlp(10, 0)
-    x = torch.zeros(6, fx.HIDDEN, dtype=torch.bfloat16, device="cuda")
-    assert eligibility(mlp, x) == (True, "")
-    assert not eligibility(mlp, torch.zeros(9, fx.HIDDEN, dtype=torch.bfloat16, device="cuda"))[0]
-    assert not eligibility(mlp, x.float())[0]
-    assert not eligibility(mlp, x[:, :-8])[0]
-    assert not eligibility(
-        mlp, torch.zeros(6, 2 * fx.HIDDEN, dtype=torch.bfloat16, device="cuda")[:, ::2]
-    )[0]
-    assert not eligibility(mlp, x.cpu())[0]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py
deleted file mode 100644
index 7e67cf291de1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_adopt.py
+++ /dev/null
@@ -1,152 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 tp8 target's class swaps side by side: the eight built-in classes the target
-adopts, what each adoption refuses, and the state the target classes may declare.
-
-Host-only: the objects are created without their constructors. What each adoption attaches is
-covered by the layer, attention and speculative-worker tests of this target.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    attention,
-    layers,
-    spec_worker,
-)
-from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.indexer import DeepseekV4Indexer
-from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
-from tensorrt_llm._torch.models.modeling_deepseekv3 import (
-    DeepseekV3DecoderLayer,
-    Deepseekv3MoE,
-    DeepseekV3MTP,
-    DeepseekV32Attention,
-)
-from tensorrt_llm._torch.models.modeling_exaone_moe import ExaoneMoeSparseMoEBlock
-from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-from tensorrt_llm._torch.speculative.mtp_dynamic_tree import MTPEagleDynamicTreeWorker
-
-pytestmark = pytest.mark.cpu_only
-
-# The target and the family code it adopts with (``shared``).
-TARGET_PACKAGE = layers.__name__.rpartition(".")[0].rpartition(".")[0]
-
-
-def _adopt_moe(moe):
-    return layers.GlmMoE.adopt(moe, None, 0)
-
-
-# (target class, the built-in class it adopts, the adoption)
-SWAPS = [
-    (attention.GlmMLA, DeepseekV32Attention, attention.GlmMLA.adopt),
-    (attention.GlmDsaHooks, DSAHooks, attention.GlmDsaHooks.adopt),
-    (attention.GlmDsaBackend, DSATrtllmAttention, attention.GlmDsaBackend.adopt),
-    (attention.GlmIndexer, Indexer, attention.GlmIndexer.adopt),
-    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, layers.GlmDecoderLayer.adopt),
-    (layers.GlmMtpLayer, DeepseekV3MTP, layers.GlmMtpLayer.adopt),
-    (layers.GlmMoE, Deepseekv3MoE, _adopt_moe),
-    (spec_worker.GlmMtpWorker, MTPEagleWorker, spec_worker.GlmMtpWorker.adopt),
-]
-SWAP_IDS = [cls.__name__ for cls, _, _ in SWAPS]
-
-# Every attribute the target's code reads that adoption may leave unset, with its default.
-CLASS_DEFAULTS = {
-    layers.GlmDecoderLayer: {
-        "post_attn_hop": None,
-        "_into_hop": None,
-        "_hop_routing": None,
-        "post_moe_hop": None,
-        "deferred_norm_hop": None,
-        "_expert_prefetch_link": None,
-    },
-    layers.GlmMtpLayer: {
-        "post_attn_hop": None,
-        "_into_hop": None,
-        "_hop_routing": None,
-        "post_moe_hop": None,
-        "deferred_norm_hop": None,
-        "_expert_prefetch_link": None,
-        "draft_front": None,
-        "draft_live_row_hop": None,
-        "_draft_live_rows": None,
-        "_draft_live_rows_compacted": False,
-    },
-    layers.GlmMoE: {"router_gemm": None, "_routed_into_hop": None},
-    attention.GlmMLA: {"_attn_tail": None, "_tail_prefetch": False},
-    attention.GlmIndexer: {"_topk_bridge": None},
-    spec_worker.GlmMtpWorker: {"_draft_loop_ops": None},
-}
-
-
-def _bare(cls):
-    obj = cls.__new__(cls)
-    if isinstance(obj, nn.Module):
-        nn.Module.__init__(obj)
-    return obj
-
-
-def _target_classes(cls):
-    return [klass for klass in cls.__mro__ if klass.__module__.startswith(TARGET_PACKAGE)]
-
-
-@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
-def test_each_target_class_extends_only_its_builtin_class(cls, builtin, adopt):
-    assert issubclass(cls, builtin)
-    added = [klass for klass in cls.__mro__ if klass not in builtin.__mro__]
-    assert added == _target_classes(cls), added
-
-
-@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
-def test_adoption_refuses_a_subclass_of_the_builtin_class(cls, builtin, adopt):
-    subclass = type(f"_Derived{builtin.__name__}", (builtin,), {})
-    obj = _bare(subclass)
-    with pytest.raises(TypeError, match=subclass.__name__):
-        adopt(obj)
-    assert type(obj) is subclass
-
-
-@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
-def test_adoption_refuses_an_adopted_object(cls, builtin, adopt):
-    obj = _bare(cls)
-    with pytest.raises(TypeError, match=cls.__name__):
-        adopt(obj)
-    assert type(obj) is cls
-
-
-@pytest.mark.parametrize(
-    "upstream, adopt",
-    [
-        (ExaoneMoeSparseMoEBlock, _adopt_moe),
-        (DeepseekV4Indexer, attention.GlmIndexer.adopt),
-        (MTPEagleDynamicTreeWorker, spec_worker.GlmMtpWorker.adopt),
-    ],
-    ids=lambda value: getattr(value, "__name__", None),
-)
-def test_adoption_refuses_the_upstream_subclasses(upstream, adopt):
-    obj = _bare(upstream)
-    with pytest.raises(TypeError, match=upstream.__name__):
-        adopt(obj)
-    assert type(obj) is upstream
-
-
-@pytest.mark.parametrize("cls, builtin, adopt", SWAPS, ids=SWAP_IDS)
-def test_the_target_classes_keep_the_builtin_layout(cls, builtin, adopt):
-    for klass in _target_classes(cls):
-        own = vars(klass)
-        assert "__init__" not in own and "__new__" not in own and "__slots__" not in own, klass
-        for name, value in own.items():
-            assert not isinstance(value, (torch.Tensor, nn.Module)), (klass, name)
-
-
-@pytest.mark.parametrize("cls", list(CLASS_DEFAULTS), ids=lambda cls: cls.__name__)
-def test_the_class_level_defaults(cls):
-    builtin = next(builtin for target, builtin, _ in SWAPS if target is cls)
-    for name, default in CLASS_DEFAULTS[cls].items():
-        assert getattr(cls, name) is default, name
-        assert not hasattr(builtin, name), name
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py
deleted file mode 100644
index a3c3df4e39bb..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_layers.py
+++ /dev/null
@@ -1,888 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 tp8 target's decoder layers, draft layer and MoE modules: their adoption, the
-phase gate of every forward override, the decode loop of ``DecodeTarget`` and the shell's
-aliases.
-
-Host-only: the built-in modules are created without their constructors and carry only the
-attributes that adoption and the phase checks read. The kernel factories are stand-ins that
-record their calls, and the built-in methods an override falls back to are recorders, so each
-test sees which call a step makes and in which order adoption builds the fast-path objects.
-"""
-
-from __future__ import annotations
-
-import ast
-import importlib.util
-import inspect
-import types
-from pathlib import Path
-
-import pytest
-import torch
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    attention,
-    dispatch,
-    layers,
-    spec_worker,
-)
-from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
-from tensorrt_llm._torch.models.modeling_deepseekv3 import (
-    DeepseekV3DecoderLayer,
-    DeepseekV3Gate,
-    DeepseekV3Model,
-    Deepseekv3MoE,
-    DeepseekV3MTP,
-    DeepseekV32Attention,
-    GlmMoeDsaForCausalLM,
-)
-from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-
-pytestmark = pytest.mark.cpu_only
-
-PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
-NUM_LAYERS = 3  # layer 0 dense, layers 1 and 2 MoE, layer 3 the draft layer
-HIDDEN = 16
-
-
-def _load_scope():
-    """``kernels/deferred_norm/scope.py`` without importing the deferred-norm kernels."""
-    path = Path(dispatch.__file__).parent / "kernels" / "deferred_norm" / "scope.py"
-    spec = importlib.util.spec_from_file_location(
-        f"{dispatch.__package__}.kernels.deferred_norm.scope", path
-    )
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-SCOPE = _load_scope()
-
-
-# --- built-in modules without their constructors ------------------------------------------
-
-
-def _bare(cls, **attributes):
-    module = cls.__new__(cls)
-    nn.Module.__init__(module)
-    for name, value in attributes.items():
-        setattr(module, name, value)
-    return module
-
-
-class _Gate(nn.Module):
-    """Stands in for the router gate: a weight, a bias and logits of its own."""
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.zeros(4, HIDDEN))
-        self.e_score_correction_bias = nn.Parameter(torch.zeros(4))
-        self.register_buffer("routing_table", torch.zeros(4))
-
-    def forward(self, hidden_states):
-        return "gate logits"
-
-
-class _Experts(nn.Module):
-    def __init__(self, calls=None) -> None:
-        super().__init__()
-        self.calls = calls if calls is not None else []
-        self.weight = nn.Parameter(torch.zeros(2, 2))
-
-    def forward(self, *args, **kwargs):
-        self.calls.append((args, kwargs))
-        return "routed"
-
-
-class _Tail:
-    """Stands in for an attention tail."""
-
-    def __init__(self, events) -> None:
-        self.events = events
-
-    def set_successor_weights(self, weights) -> None:
-        self.events.append(("set_successor_weights", tuple(weights)))
-
-
-class _SparseMlaDecodeFmha:
-    """Stands in for the sparse-MLA decode FMHA library."""
-
-
-def _attention(layer_idx):
-    """A layer's built-in attention modules: the MLA, its DSA hooks, its DSA backend and, on a
-    full-indexer layer (even index here), its indexer."""
-    backend = DSATrtllmAttention.__new__(DSATrtllmAttention)
-    backend.indexer = _bare(Indexer) if layer_idx % 2 == 0 else None
-    return _bare(
-        DeepseekV32Attention,
-        layer_idx=layer_idx,
-        sparse_attn_hooks=DSAHooks.__new__(DSAHooks),
-        mqa=backend,
-        mha=None,
-        indexer=backend.indexer,
-    )
-
-
-class _SparseParams:
-    def to_sparse_params(self, pretrained_config, layer_idx):
-        return types.SimpleNamespace(is_full_indexer_layer=layer_idx % 2 == 0)
-
-
-def _config():
-    return types.SimpleNamespace(
-        n_routed_experts=4,
-        first_k_dense_replace=1,
-        moe_layer_freq=1,
-        num_hidden_layers=NUM_LAYERS,
-        hidden_size=HIDDEN,
-        torch_dtype=torch.bfloat16,
-    )
-
-
-def _model_config(config):
-    mapping = types.SimpleNamespace(enable_attention_dp=False, tp_size=8, tp_rank=0)
-    return types.SimpleNamespace(
-        pretrained_config=config,
-        mapping=mapping,
-        moe_backend="TRTLLM",
-        sparse_attention_config=_SparseParams(),
-    )
-
-
-def _moe(experts_calls=None):
-    return _bare(
-        Deepseekv3MoE,
-        gate=_Gate(),
-        experts=_Experts(experts_calls),
-        use_dp=False,
-        mapping=types.SimpleNamespace(tp_size=8),
-        shared_experts=None,
-    )
-
-
-def _layer(layer_idx, events, model_config, cls=DeepseekV3DecoderLayer, nvfp4=True):
-    config = model_config.pretrained_config
-    moe = layer_idx >= config.first_k_dense_replace
-    return _bare(
-        cls,
-        model_config=model_config,
-        layer_idx=layer_idx,
-        mapping=model_config.mapping,
-        allreduce=object(),
-        self_attn=_attention(layer_idx),
-        mlp=_moe() if moe else nn.Linear(2, 2),
-        fusion_config=types.SimpleNamespace(PRE_MOE_FUSION=moe, POST_MOE_FUSION=moe),
-        is_nvfp4=nvfp4,
-        top_k=2,
-    )
-
-
-def _mtp(events, model_config):
-    layer = _layer(NUM_LAYERS, events, model_config, cls=DeepseekV3MTP, nvfp4=False)
-    layer.enorm = nn.LayerNorm(2)
-    layer.hnorm = nn.LayerNorm(2)
-    return layer
-
-
-def _keys(module):
-    return (
-        [name for name, _ in module.named_parameters()],
-        [name for name, _ in module.named_buffers()],
-        [name for name, _ in module.named_modules()],
-    )
-
-
-# --- stand-in kernel factories ------------------------------------------------------------
-
-
-class _Made:
-    """A fast-path object a factory made, with the call that made it."""
-
-    def __init__(self, events, name, *args, engaged=True, **kwargs):
-        self.events, self.name, self.args, self.kwargs = events, name, args, kwargs
-        self.engaged = engaged
-        events.append(("make", name))
-
-    def enable_routed_input(self) -> bool:
-        self.events.append(("enable_routed_input", self.name))
-        return True
-
-
-class _FastPaths:
-    """``dispatch.is_enabled`` / ``get_impl`` / ``log_capability`` for the adoption tests."""
-
-    def __init__(self, enabled, post_moe_engaged=True):
-        self.events = []
-        self.enabled = enabled
-        self.post_moe_engaged = post_moe_engaged
-        events = self.events
-
-        class DeferredNormHop(_Made):
-            def __init__(self, post_moe_hop, name="deferred_norm_hop"):
-                super().__init__(events, name, post_moe_hop)
-                # As the real hop: it names its post-MoE hop too, so the search for the
-                # constructing layer must skip it.
-                self.post_moe_hop = post_moe_hop
-                _, next_layer, how = SCOPE.next_layer_kind(post_moe_hop)
-                events.append(("deferred hop serves", next_layer, how))
-
-        class EvictHintsDeferredHop(DeferredNormHop):
-            def __init__(self, post_moe_hop):
-                super().__init__(post_moe_hop, name="deferred_norm_evict_hints_hop")
-
-        def draft_front(*args):
-            return None
-
-        def make_attn_tail(mla):
-            # As the real factory: the draft layer keeps the native tail.
-            if mla.layer_idx >= NUM_LAYERS:
-                return None
-            events.append(("make", "attn_tail"))
-            return _Tail(events)
-
-        def wrap_attn_tail(mla, tail):
-            if tail is not None:
-                events.append(("make", "staged_attn_tail"))
-            return tail
-
-        self.draft_front = draft_front
-        self.entries = {
-            ("attention_output_tail", "make_attn_tail"): make_attn_tail,
-            ("oproj_weight_staging", "wrap_attn_tail"): wrap_attn_tail,
-            ("sparse_mla_decode", "fmha_lib"): _SparseMlaDecodeFmha,
-            ("post_attention_hop", "make_post_attn_hop"): lambda *a, **k: _Made(
-                events, "post_attn_hop", *a, **k
-            ),
-            ("post_attention_hop", "make_into_hop_transport"): lambda hop: _Made(
-                events, "into_hop", hop
-            ),
-            ("hop_routing", "make_transport"): lambda hop: _Made(events, "hop_routing", hop),
-            ("router_gemm", "make_moe_router"): lambda **k: _Made(events, "router_gemm", **k),
-            ("post_moe_hop", "post_moe_hop"): lambda **k: _Made(
-                events, "post_moe_hop", engaged=self.post_moe_engaged, **k
-            ),
-            ("deferred_norm", "deferred_hop"): DeferredNormHop,
-            ("deferred_norm_evict_hints", "deferred_hop"): EvictHintsDeferredHop,
-            ("draft_front", "native_norm_reason"): lambda enorm, hnorm: "",
-            ("draft_front", "draft_front"): draft_front,
-            ("draft_live_rows", "make_live_row_hop"): lambda hop: _Made(
-                events, "draft_live_row_hop", hop
-            ),
-        }
-
-    def is_enabled(self, path):
-        assert path in dispatch.FAST_PATH_NAMES, path
-        return path in self.enabled
-
-    def get_impl(self, path, name):
-        assert path in self.enabled, path
-        return self.entries[(path, name)]
-
-    def log_capability(self, path, key, eligible, reason=""):
-        self.events.append(("log_capability", path, key, eligible, reason))
-        return eligible
-
-
-@pytest.fixture
-def fast_paths(monkeypatch):
-    def install(enabled=dispatch.FAST_PATH_NAMES, post_moe_engaged=True):
-        paths = _FastPaths(frozenset(enabled), post_moe_engaged)
-        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-        monkeypatch.setattr(dispatch, "log_capability", paths.log_capability)
-        return paths
-
-    return install
-
-
-def _built_in_model(events):
-    """The built-in model as ``adopt_model`` sees it: decoder layers, then the draft layer, and
-    the speculative worker."""
-    model_config = _model_config(_config())
-    core = DeepseekV3Model.__new__(DeepseekV3Model)
-    nn.Module.__init__(core)
-    core.num_hidden_layers = NUM_LAYERS
-    core.layers = nn.ModuleList(
-        [_layer(i, events, model_config) for i in range(NUM_LAYERS)] + [_mtp(events, model_config)]
-    )
-    return types.SimpleNamespace(model=core, spec_worker=_bare(MTPEagleWorker))
-
-
-# --- adoption -----------------------------------------------------------------------------
-
-
-def test_adoption_builds_each_object_in_the_builtin_constructor_order(fast_paths):
-    """Layer by layer: the attention tail, the post-attention hop and its transports, the router
-    GEMM, the post-MoE hops, the tail's successor weights; on the draft layer then its front and
-    live-row hop."""
-    paths = fast_paths()
-    model = _built_in_model(paths.events)
-
-    modeling.adopt_model(model, plan=None)
-
-    tail = [("make", "attn_tail"), ("make", "staged_attn_tail")]
-    pre_moe = tail + [("make", "post_attn_hop"), ("make", "into_hop"), ("make", "hop_routing")]
-    post_moe = [
-        ("make", "router_gemm"),
-        ("make", "post_moe_hop"),
-        ("enable_routed_input", "post_moe_hop"),
-    ]
-    layers_0 = pre_moe
-    layers_1 = (
-        pre_moe
-        + post_moe
-        + [
-            ("make", "deferred_norm_evict_hints_hop"),
-            ("deferred hop serves", 2, "model config"),
-            ("enable_routed_input", "deferred_norm_evict_hints_hop"),
-        ]
-    )
-    layers_2 = pre_moe + post_moe  # the last decoder layer: no next layer to defer into
-    tails = [e for e in paths.events if e[0] == "set_successor_weights"]
-    draft = [
-        ("make", "post_attn_hop"),
-        ("make", "router_gemm"),
-        ("log_capability", "draft_front", "draft_front", True, ""),
-        ("make", "draft_live_row_hop"),
-    ]
-    events = [e for e in paths.events if e[0] != "set_successor_weights"]
-    assert events == layers_0 + layers_1 + layers_2 + draft
-    # The tail stages its MoE layer's router weight, after that layer's post-MoE hops.
-    core = model.model
-    assert tails == [("set_successor_weights", (core.layers[i].mlp.gate.weight,)) for i in (1, 2)]
-    assert paths.events.index(tails[0]) == len(layers_0 + layers_1)
-
-
-def test_adoption_attaches_what_the_constructors_attached(fast_paths):
-    paths = fast_paths()
-    model = _built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-    dense, moe, last, draft = model.model.layers
-
-    assert [type(m) for m in (dense, moe, last)] == [layers.GlmDecoderLayer] * 3
-    assert type(draft) is layers.GlmMtpLayer
-    assert [type(m.mlp) for m in (moe, last, draft)] == [layers.GlmMoE] * 3
-    assert type(dense.mlp) is nn.Linear
-    for layer in (dense, moe, last, draft):
-        mla = layer.self_attn
-        assert type(mla) is attention.GlmMLA
-        assert type(mla.sparse_attn_hooks) is attention.GlmDsaHooks
-        assert type(mla.mqa) is attention.GlmDsaBackend
-        assert mla.mqa.prepended_fmha_lib_classes == (_SparseMlaDecodeFmha,)
-        assert mla._tail_prefetch is True
-    # Layers 0 and 2 run a full indexer, layers 1 and 3 reuse the shared top-k.
-    assert [type(m.self_attn.mqa.indexer) for m in (dense, last)] == [attention.GlmIndexer] * 2
-    assert moe.self_attn.mqa.indexer is None and draft.self_attn.mqa.indexer is None
-    assert isinstance(moe.self_attn._attn_tail, _Tail) and draft.self_attn._attn_tail is None
-
-    hop = moe.post_attn_hop
-    assert hop.args == (moe.mapping, HIDDEN, torch.bfloat16) and hop.kwargs == {
-        "native_allreduce": moe.allreduce
-    }
-    assert moe._into_hop.args == (hop,) and moe._hop_routing.args == (moe._into_hop,)
-    assert moe.mlp.router_gemm.kwargs == {
-        "gate": moe.mlp.gate,
-        "layer_idx": 1,
-        "model_config": moe.model_config,
-    }
-    assert moe.post_moe_hop.kwargs == {
-        "mapping": moe.mapping,
-        "hidden_size": HIDDEN,
-        "top_k": 2,
-        "dtype": torch.bfloat16,
-        "raw_producer": moe.mlp.experts,
-    }
-    assert moe.mlp._routed_into_hop is moe.post_moe_hop
-    assert moe.deferred_norm_hop.post_moe_hop is moe.post_moe_hop
-    assert last.deferred_norm_hop is None and last.post_moe_hop is not None
-    assert dense.post_moe_hop is None and dense.post_attn_hop is not None
-    assert draft._into_hop is None and draft.post_moe_hop is None
-    assert draft.draft_front is paths.draft_front
-    assert draft.draft_live_row_hop.args == (draft.post_attn_hop,)
-    assert draft._draft_live_rows is None and draft._draft_live_rows_compacted is False
-    assert type(model.spec_worker) is spec_worker.GlmMtpWorker
-    # The adopted state names every object, layer by layer.
-    state = dict(modeling.adopted_state(types.SimpleNamespace(model=model.model)))
-    assert [entry[0] for entry in state["model.layers.1"]] == [
-        "post_attn_hop",
-        "_into_hop",
-        "_hop_routing",
-        "post_moe_hop",
-        "deferred_norm_hop",
-        "self_attn._attn_tail",
-        "mlp.router_gemm",
-        "mlp._routed_into_hop",
-    ]
-    assert [entry[0] for entry in state["model.layers.3"]] == [
-        "post_attn_hop",
-        "draft_front",
-        "draft_live_row_hop",
-        "mlp.router_gemm",
-    ]
-
-
-def test_the_deferred_norm_variant_without_evict_hints(fast_paths):
-    paths = fast_paths(set(dispatch.FAST_PATH_NAMES) - {"deferred_norm_evict_hints"})
-    model = _built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-    assert ("make", "deferred_norm_hop") in paths.events
-    assert ("deferred hop serves", 2, "model config") in paths.events
-
-
-def test_without_the_fast_paths_adoption_attaches_nothing(fast_paths):
-    paths = fast_paths(enabled=())
-    model = _built_in_model(paths.events)
-    before = [_keys(layer) for layer in model.model.layers]
-
-    modeling.adopt_model(model, plan=None)
-
-    assert paths.events == []
-    for layer in model.model.layers:
-        assert type(layer.self_attn) is attention.GlmMLA and layer.self_attn._attn_tail is None
-        assert layer.self_attn.mqa.prepended_fmha_lib_classes == ()
-        assert "prepended_fmha_lib_classes" not in vars(layer.self_attn.mqa)
-        for name in ("post_attn_hop", "_into_hop", "_hop_routing", "post_moe_hop"):
-            assert getattr(layer, name) is None and name not in vars(layer)
-        assert layer.deferred_norm_hop is None and layer._expert_prefetch_link is None
-    draft = model.model.layers[NUM_LAYERS]
-    assert draft.draft_front is None and draft.draft_live_row_hop is None
-    assert model.model.layers[1].mlp.router_gemm is None
-    assert model.model.layers[1].mlp._routed_into_hop is None
-    assert [_keys(layer) for layer in model.model.layers] == before
-
-
-def test_adoption_adds_no_parameter_buffer_or_module(fast_paths):
-    paths = fast_paths()
-    model = _built_in_model(paths.events)
-    before = [_keys(layer) for layer in model.model.layers]
-    modeling.adopt_model(model, plan=None)
-    assert [_keys(layer) for layer in model.model.layers] == before
-    assert "mlp.gate.weight" in before[1][0] and before[1][1] == ["mlp.gate.routing_table"]
-
-
-def test_a_post_moe_hop_that_does_not_engage_is_not_attached(fast_paths):
-    paths = fast_paths(post_moe_engaged=False)
-    model = _built_in_model(paths.events)
-    modeling.adopt_model(model, plan=None)
-    moe = model.model.layers[1]
-    assert moe.post_moe_hop is None and moe.deferred_norm_hop is None
-    assert moe.mlp._routed_into_hop is None
-    assert ("enable_routed_input", "post_moe_hop") not in paths.events
-
-
-def test_the_draft_front_needs_non_adp_tensor_parallelism(fast_paths):
-    paths = fast_paths()
-    events = paths.events
-    model_config = _model_config(_config())
-    model_config.mapping.enable_attention_dp = True
-    draft = layers.GlmMtpLayer.adopt(_mtp(events, model_config))
-    assert draft.draft_front is None
-    assert (
-        "log_capability",
-        "draft_front",
-        "draft_front",
-        False,
-        "requires non-ADP tensor parallelism",
-    ) in events
-
-
-class _OtherDecoderLayer(DeepseekV3DecoderLayer):
-    pass
-
-
-class _OtherMoE(Deepseekv3MoE):
-    pass
-
-
-class _OtherAttention(DeepseekV32Attention):
-    pass
-
-
-def _refusal(case):
-    """``(adopt, module, the class the refusal names)``."""
-    events = []
-    model_config = _model_config(_config())
-    if case == "the draft layer as a decoder layer":
-        return layers.GlmDecoderLayer.adopt, _mtp(events, model_config), "DeepseekV3MTP"
-    if case == "a decoder layer subclass":
-        module = _layer(1, events, model_config, cls=_OtherDecoderLayer)
-        return layers.GlmDecoderLayer.adopt, module, "_OtherDecoderLayer"
-    if case == "a module as a decoder layer":
-        return layers.GlmDecoderLayer.adopt, nn.Module(), "Module"
-    if case == "a decoder layer as the draft layer":
-        return layers.GlmMtpLayer.adopt, _layer(0, events, model_config), "DeepseekV3DecoderLayer"
-    if case == "a decoder layer with an attention subclass":
-        module = _layer(1, events, model_config)
-        module.self_attn.__class__ = _OtherAttention
-        return layers.GlmDecoderLayer.adopt, module, "_OtherAttention"
-    module = _layer(1, events, model_config)
-    module.mlp = _bare(_OtherMoE, gate=_Gate())
-    return layers.GlmDecoderLayer.adopt, module, "_OtherMoE"
-
-
-@pytest.mark.parametrize(
-    "case",
-    [
-        "the draft layer as a decoder layer",
-        "a decoder layer subclass",
-        "a module as a decoder layer",
-        "a decoder layer as the draft layer",
-        "a decoder layer with an attention subclass",
-        "a decoder layer with a MoE subclass",
-    ],
-)
-def test_adoption_refuses_anything_but_the_exact_builtin_class(fast_paths, case):
-    fast_paths(enabled=())
-    adopt, module, got = _refusal(case)
-    with pytest.raises(TypeError, match=f"adopts the built-in .*, got {got}$"):
-        adopt(module)
-
-
-def test_the_moe_is_adopted_once(fast_paths):
-    fast_paths(enabled=())
-    model_config = _model_config(_config())
-    moe = layers.GlmMoE.adopt(_moe(), model_config, 1)
-    with pytest.raises(TypeError, match="GlmMoE adopts the built-in Deepseekv3MoE, got GlmMoE"):
-        layers.GlmMoE.adopt(moe, model_config, 1)
-
-
-def test_the_classes_declare_no_state_of_their_own():
-    for cls in (layers.GlmDecoderLayer, layers.GlmMtpLayer, layers.GlmMoE):
-        for base in cls.__mro__[: cls.__mro__.index(nn.Module)]:
-            if base.__module__ == layers.__name__:
-                assert "__slots__" not in vars(base) and "__init__" not in vars(base), base
-
-
-# --- the phase gate -----------------------------------------------------------------------
-
-# Each forward override and the built-in method it falls back to.
-OVERRIDES = [
-    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, "forward"),
-    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, "forward_MoE"),
-    (layers.GlmDecoderLayer, DeepseekV3DecoderLayer, "forward_mlp"),
-    (layers.GlmMtpLayer, DeepseekV3MTP, "forward"),
-    (layers.GlmMoE, Deepseekv3MoE, "forward"),
-    (layers.GlmMoE, Deepseekv3MoE, "compute_routed_output"),
-]
-
-
-def _gate_of(cls, name):
-    tree = ast.parse(Path(layers.__file__).read_text())
-    node = next(
-        n
-        for c in tree.body
-        if isinstance(c, ast.ClassDef) and c.name == cls.__name__
-        for n in c.body
-        if isinstance(n, ast.FunctionDef) and n.name == name
-    )
-    return node, node.body[0]
-
-
-@pytest.mark.parametrize("cls, builtin, name", OVERRIDES, ids=lambda x: getattr(x, "__name__", x))
-def test_each_override_starts_with_the_phase_gate(cls, builtin, name):
-    """``if dispatch.phase() is not DECODE: return super().<name>(<every built-in parameter by
-    name>)``; the parameters only the decode composition takes are dropped."""
-    node, gate = _gate_of(cls, name)
-    assert ast.unparse(gate.test) == "dispatch.phase() is not dispatch.Phase.DECODE"
-    assert len(gate.body) == 1 and not gate.orelse
-    call = gate.body[0].value
-    assert ast.unparse(call.func) == f"super().{name}" and not call.args
-    passed = [k.arg for k in call.keywords]
-    parameters = inspect.signature(getattr(builtin, name)).parameters
-    expected = [
-        p.name if p.kind is not p.VAR_KEYWORD else None
-        for p in parameters.values()
-        if p.name != "self"
-    ]
-    assert passed == expected
-    # The override keeps the built-in parameters, in order, and may add keywords after them.
-    own = list(inspect.signature(getattr(cls, name)).parameters)
-    assert own[: len(parameters)] == list(parameters)
-
-
-def _decode_layer():
-    model_config = _model_config(_config())
-    layer = _layer(1, [], model_config)
-    layer.__class__ = layers.GlmDecoderLayer
-    layer.mlp.__class__ = layers.GlmMoE
-    return layer
-
-
-def _calls_through(monkeypatch, owner, name, calls):
-    def recorder(self, *args, **kwargs):
-        calls.append((name, self, args, kwargs))
-        return f"{owner.__name__}.{name}"
-
-    monkeypatch.setattr(owner, name, recorder)
-
-
-LAYER_ARGS = dict(position_ids="pos", hidden_states="h", attn_metadata="md", residual="r")
-
-
-@pytest.mark.parametrize("phase", [PREFILL, None], ids=["prefill", "outside a step"])
-def test_off_decode_every_override_runs_the_builtin_method(monkeypatch, phase):
-    calls = []
-    for owner, name in [
-        (DeepseekV3DecoderLayer, "forward"),
-        (DeepseekV3DecoderLayer, "forward_MoE"),
-        (DeepseekV3DecoderLayer, "forward_mlp"),
-        (DeepseekV3MTP, "forward"),
-        (Deepseekv3MoE, "forward"),
-        (Deepseekv3MoE, "compute_routed_output"),
-    ]:
-        _calls_through(monkeypatch, owner, name, calls)
-    monkeypatch.setattr(layers.GlmDecoderLayer, "forward_decode", lambda *a, **k: pytest.fail())
-    layer = _decode_layer()
-    mtp = _bare(layers.GlmMtpLayer)
-
-    def run():
-        out = [
-            layer.forward(**LAYER_ARGS, spec_metadata="spec", lora_params="lora"),
-            layer.forward_MoE("h", "md", "r", "spec", into_hop_published=True),
-            layer.forward_mlp("h", "r"),
-            mtp.forward("ids", "pos", "h", "embed", "md", [2], "spec", lora_params="lora"),
-            layer.mlp.forward("h", None, [2], "params", False, "logits", "routing"),
-            layer.mlp.compute_routed_output("h", None, [2], True, "logits", "routing"),
-        ]
-        assert out == [
-            "DeepseekV3DecoderLayer.forward",
-            "DeepseekV3DecoderLayer.forward_MoE",
-            "DeepseekV3DecoderLayer.forward_mlp",
-            "DeepseekV3MTP.forward",
-            "Deepseekv3MoE.forward",
-            "Deepseekv3MoE.compute_routed_output",
-        ]
-
-    if phase is None:
-        run()
-    else:
-        with dispatch.step(phase):
-            run()
-    assert [(c[0], c[2], c[3]) for c in calls] == [
-        ("forward", (), {**LAYER_ARGS, "spec_metadata": "spec", "lora_params": "lora"}),
-        (
-            "forward_MoE",
-            (),
-            dict(hidden_states="h", attn_metadata="md", residual="r", spec_metadata="spec"),
-        ),
-        ("forward_mlp", (), dict(hidden_states="h", residual="r", spec_metadata=None)),
-        (
-            "forward",
-            (),
-            dict(
-                input_ids="ids",
-                position_ids="pos",
-                hidden_states="h",
-                embed_tokens="embed",
-                attn_metadata="md",
-                all_rank_num_tokens=[2],
-                spec_metadata="spec",
-                lora_params="lora",
-            ),
-        ),
-        (
-            "forward",
-            (),
-            dict(
-                hidden_states="h",
-                hidden_states_fp4=None,
-                all_rank_num_tokens=[2],
-                final_all_reduce_params="params",
-                do_finalize=False,
-            ),
-        ),
-        (
-            "compute_routed_output",
-            (),
-            dict(
-                hidden_states="h", hidden_states_fp4=None, all_rank_num_tokens=[2], do_finalize=True
-            ),
-        ),
-    ]
-
-
-def test_on_a_decode_step_the_layer_runs_its_decode_composition(monkeypatch):
-    calls = []
-    _calls_through(monkeypatch, DeepseekV3DecoderLayer, "forward", calls)
-    _calls_through(monkeypatch, layers.GlmDecoderLayer, "forward_decode", calls)
-    layer = _decode_layer()
-    with dispatch.step(DECODE):
-        assert layer.forward(**LAYER_ARGS, lora_params="lora") == "GlmDecoderLayer.forward_decode"
-    assert calls == [
-        ("forward_decode", layer, ("pos", "h", "md", "r", None), {"lora_params": "lora"})
-    ]
-
-
-@pytest.mark.parametrize("phase", [DECODE, PREFILL], ids=["decode", "prefill"])
-def test_the_router_gemm_replaces_the_gate_on_decode_steps_only(fast_paths, phase):
-    paths = fast_paths(enabled={"router_gemm"})
-    paths.entries[("router_gemm", "make_moe_router")] = lambda **k: (lambda h: "router logits")
-    experts = []
-    moe = layers.GlmMoE.adopt(_moe(experts), _model_config(_config()), 1)
-    hidden = torch.zeros(2, HIDDEN)
-
-    with dispatch.step(phase):
-        assert moe.compute_routed_output(hidden, None, [2], True) == "routed"
-
-    logits = "router logits" if phase is DECODE else "gate logits"
-    assert experts == [
-        (
-            (hidden, logits),
-            dict(
-                do_finalize=True,
-                output_dtype=hidden.dtype,
-                all_rank_num_tokens=[2],
-                use_dp_padding=False,
-            ),
-        )
-    ]
-
-
-def test_the_post_attention_hop_reads_the_builtin_gates_weight_only():
-    layer = _decode_layer()
-    assert layer._post_attn_router_weight() is None  # a stand-in gate
-    layer.mlp.gate = _bare(DeepseekV3Gate, weight=nn.Parameter(torch.zeros(4, HIDDEN)))
-    assert layer._post_attn_router_weight() is layer.mlp.gate.weight
-
-
-# --- DecodeTarget -------------------------------------------------------------------------
-
-
-class _NoContexts:
-    """Attention metadata a decode step passes; reading the context count fails the test."""
-
-    def __getattr__(self, name):
-        raise AssertionError(f"DecodeTarget read attn_metadata.{name}")
-
-
-class _RecordingLayer(nn.Module):
-    def __init__(self, index, calls) -> None:
-        super().__init__()
-        self.index, self.calls = index, calls
-
-    def forward(self, **kwargs):
-        self.calls.append((self.index, kwargs, dispatch.phase()))
-        return kwargs["hidden_states"] + 1, f"residual {self.index}"
-
-
-class _Embedding(nn.Module):
-    def __init__(self, calls) -> None:
-        super().__init__()
-        self.calls = calls
-
-    def forward(self, input_ids):
-        self.calls.append(("embed", input_ids))
-        return torch.zeros(2, HIDDEN)
-
-
-@pytest.fixture
-def decode_core(monkeypatch):
-    calls = []
-    core = DeepseekV3Model.__new__(DeepseekV3Model)
-    nn.Module.__init__(core)
-    core.embed_tokens = _Embedding(calls)
-    core.layers = nn.ModuleList([_RecordingLayer(i, calls) for i in range(NUM_LAYERS + 1)])
-    core.num_hidden_layers = NUM_LAYERS
-    core.mapping_with_cp = None
-    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(core)
-    for layer in core.layers:
-        layer.register_forward_pre_hook(
-            lambda m, a, k: calls.append(("hook", m.index)), with_kwargs=True
-        )
-    return core, calls
-
-
-@pytest.mark.parametrize("prologue_rows", [None, "rows"], ids=["no packet", "packet"])
-def test_the_decode_loop_runs_the_decoder_layers_as_modules(fast_paths, decode_core, prologue_rows):
-    core, calls = decode_core
-    paths = fast_paths(enabled={"decode_prologue", "cross_layer_expert_prefetch"})
-    rows = torch.full((2, HIDDEN), 7.0) if prologue_rows else None
-    seen = []
-
-    def take_embedding(attn_metadata, input_ids, embed_tokens):
-        seen.append((attn_metadata, input_ids, embed_tokens))
-        return rows
-
-    core.decode_prologue = types.SimpleNamespace(take_embedding=take_embedding)
-    paths.entries[("cross_layer_expert_prefetch", "join")] = lambda: calls.append(("join",))
-    md = _NoContexts()
-
-    with dispatch.step(DECODE):
-        out = core._decode.forward(
-            md, input_ids="ids", position_ids="pos", spec_metadata="spec", lora_params="x"
-        )
-
-    assert seen == [(md, "ids", core.embed_tokens)]
-    embedded = [] if prologue_rows else [("embed", "ids")]
-    expected_start = (rows if prologue_rows else torch.zeros(2, HIDDEN)) + NUM_LAYERS
-    assert torch.equal(out, expected_start)
-    layer_calls = []
-    for i in range(NUM_LAYERS):
-        kwargs = calls[len(embedded) + 2 * i + 1][1]
-        assert set(kwargs) == {
-            "position_ids",
-            "hidden_states",
-            "attn_metadata",
-            "residual",
-            "spec_metadata",
-        }
-        assert kwargs["residual"] == (None if i == 0 else f"residual {i - 1}")
-        assert kwargs["attn_metadata"] is md and kwargs["spec_metadata"] == "spec"
-        layer_calls += [("hook", i), (i, kwargs, DECODE)]
-    assert calls == embedded + layer_calls + [("join",)]
-
-
-def test_the_decode_loop_without_its_fast_paths_is_the_builtin_loop(fast_paths, decode_core):
-    core, calls = decode_core
-    fast_paths(enabled=())
-    with dispatch.step(DECODE):
-        out = core._decode.forward(_NoContexts(), input_ids="ids")
-    assert torch.equal(out, torch.zeros(2, HIDDEN) + NUM_LAYERS)
-    assert [c[0] for c in calls] == ["embed"] + ["hook", 0, "hook", 1, "hook", 2]
-
-
-def test_the_decode_loop_takes_ids_or_embeddings(fast_paths, decode_core):
-    core, _ = decode_core
-    fast_paths(enabled=())
-    with pytest.raises(ValueError, match="cannot specify both input_ids and inputs_embeds"):
-        core._decode.forward(_NoContexts(), input_ids="ids", inputs_embeds=torch.zeros(1))
-    with pytest.raises(ValueError, match="must specify either one"):
-        core._decode.forward(_NoContexts())
-
-
-def test_the_decode_loop_never_reads_the_step_kind():
-    source = inspect.getsource(modeling.DecodeTarget.forward)
-    assert "num_contexts" not in source and "num_ctx" not in source
-
-
-# --- the shell's aliases ------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("linked", [True, False], ids=["prefetch", "no prefetch"])
-def test_the_shell_links_the_expert_prefetch_after_the_builtin_aliases(
-    monkeypatch, fast_paths, linked
-):
-    paths = fast_paths(enabled={"cross_layer_expert_prefetch"} if linked else ())
-    events = paths.events
-    paths.entries[("cross_layer_expert_prefetch", "link_layers")] = (
-        lambda layers_, n: events.append(("link_layers", layers_, n))
-    )
-    monkeypatch.setattr(
-        GlmMoeDsaForCausalLM, "setup_aliases", lambda self: events.append("built-in aliases")
-    )
-    shell = _bare(modeling.ModelingV2Glm52Nvfp4Sm103Tp8)
-    shell.model = types.SimpleNamespace(layers=["layer"])
-    shell.model_config = types.SimpleNamespace(
-        pretrained_config=types.SimpleNamespace(num_hidden_layers=78)
-    )
-
-    shell.setup_aliases()
-
-    assert events == ["built-in aliases"] + ([("link_layers", ["layer"], 78)] if linked else [])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py
deleted file mode 100644
index 4bc0a870e4ff..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_phase.py
+++ /dev/null
@@ -1,213 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The phase of a model step of the GLM-5.2 tp8 target: the scope, the targets' phases and the
-shell's forward, which runs every step under the phase of the target that takes it.
-
-Host-only: the shell is created without its constructor around an adopted built-in text model
-without layers, and the built-in causal LM's forward is replaced by a recorder that also calls
-a stand-in speculative worker, so each test sees the phase every part of a step runs under.
-"""
-
-from __future__ import annotations
-
-import inspect
-
-import pytest
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model
-from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM
-
-pytestmark = pytest.mark.cpu_only
-
-PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
-
-
-class _Metadata:
-    def __init__(self, num_contexts: int) -> None:
-        self.num_contexts = num_contexts
-
-
-@pytest.fixture(autouse=True)
-def _no_phase():
-    assert dispatch.phase() is None
-    yield
-    assert dispatch.phase() is None
-
-
-# --- the scope ----------------------------------------------------------------------------
-
-
-def test_there_is_no_phase_outside_a_step():
-    assert dispatch.phase() is None
-
-
-@pytest.mark.parametrize("phase", [PREFILL, DECODE])
-def test_a_step_runs_under_its_phase(phase):
-    with dispatch.step(phase):
-        assert dispatch.phase() is phase
-    assert dispatch.phase() is None
-
-
-def test_steps_nest_and_restore_the_enclosing_phase():
-    with dispatch.step(PREFILL):
-        with dispatch.step(DECODE):
-            assert dispatch.phase() is DECODE
-            with dispatch.step(PREFILL):
-                assert dispatch.phase() is PREFILL
-            assert dispatch.phase() is DECODE
-        assert dispatch.phase() is PREFILL
-
-
-def test_a_step_that_raises_restores_the_enclosing_phase():
-    with dispatch.step(PREFILL):
-        with pytest.raises(ValueError, match="inside"):
-            with dispatch.step(DECODE):
-                raise ValueError("inside")
-        assert dispatch.phase() is PREFILL
-
-
-def test_the_phases():
-    assert [p.value for p in dispatch.Phase] == ["prefill", "decode"]
-
-
-# --- the targets --------------------------------------------------------------------------
-
-
-def test_each_target_names_its_phase():
-    assert modeling.PrefillTarget.phase is PREFILL
-    assert modeling.DecodeTarget.phase is DECODE
-
-
-@pytest.mark.parametrize(
-    "num_contexts, phase",
-    [(0, DECODE), (1, PREFILL), (3, PREFILL)],
-    ids=["decode", "prefill", "mixed"],
-)
-def test_the_selected_target_carries_the_phase_of_the_step(num_contexts, phase):
-    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
-    assert core._select_target(_Metadata(num_contexts)).phase is phase
-
-
-# --- the shell's forward ------------------------------------------------------------------
-
-
-def _builtin_text_model() -> DeepseekV3Model:
-    model = DeepseekV3Model.__new__(DeepseekV3Model)
-    nn.Module.__init__(model)
-    return model
-
-
-@pytest.fixture
-def shell_forward(monkeypatch):
-    """A shell around an adopted core, and the phases its step's parts ran under."""
-    seen = []
-
-    def builtin_forward(
-        self,
-        attn_metadata,
-        input_ids=None,
-        position_ids=None,
-        inputs_embeds=None,
-        return_context_logits=False,
-        spec_metadata=None,
-        resource_manager=None,
-        **kwargs,
-    ):
-        seen.append(("model", dispatch.phase()))
-        seen.append(
-            (
-                "arguments",
-                attn_metadata,
-                input_ids,
-                position_ids,
-                inputs_embeds,
-                return_context_logits,
-                spec_metadata,
-                resource_manager,
-                kwargs,
-            )
-        )
-        # The speculative worker and its draft loop run inside the built-in forward.
-        seen.append(("spec worker", dispatch.phase()))
-        if kwargs.get("fail"):
-            raise RuntimeError("forward failed")
-        return "logits"
-
-    monkeypatch.setattr(SpecDecOneEngineForCausalLM, "forward", builtin_forward)
-    shell = modeling.ModelingV2Glm52Nvfp4Sm103Tp8.__new__(modeling.ModelingV2Glm52Nvfp4Sm103Tp8)
-    nn.Module.__init__(shell)
-    shell.model = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
-    return shell, seen
-
-
-@pytest.mark.parametrize(
-    "num_contexts, phase",
-    [(0, DECODE), (1, PREFILL), (2, PREFILL)],
-    ids=["decode", "prefill", "mixed"],
-)
-def test_the_shell_runs_the_whole_step_under_the_phase_of_its_target(
-    shell_forward, num_contexts, phase
-):
-    shell, seen = shell_forward
-    md = _Metadata(num_contexts)
-
-    out = shell.forward(
-        attn_metadata=md,
-        input_ids="ids",
-        position_ids="positions",
-        inputs_embeds=None,
-        return_context_logits=True,
-        spec_metadata="spec",
-        resource_manager="resources",
-        lora_params="lora",
-    )
-
-    assert out == "logits"
-    assert seen == [
-        ("model", phase),
-        (
-            "arguments",
-            md,
-            "ids",
-            "positions",
-            None,
-            True,
-            "spec",
-            "resources",
-            {"lora_params": "lora"},
-        ),
-        ("spec worker", phase),
-    ]
-    assert dispatch.phase() is None
-
-
-def test_the_shell_takes_the_builtin_arguments_positionally(shell_forward):
-    shell, seen = shell_forward
-    md = _Metadata(0)
-
-    assert shell.forward(md, "ids", "positions", None, False, "spec", "resources") == "logits"
-    assert seen[1] == ("arguments", md, "ids", "positions", None, False, "spec", "resources", {})
-
-
-def test_a_step_that_fails_leaves_no_phase_behind(shell_forward):
-    shell, seen = shell_forward
-    with pytest.raises(RuntimeError, match="forward failed"):
-        shell.forward(attn_metadata=_Metadata(0), fail=True)
-    assert seen[-1] == ("spec worker", DECODE)
-    assert dispatch.phase() is None
-
-
-def test_the_shell_forward_takes_the_builtin_parameters():
-    builtin = inspect.signature(SpecDecOneEngineForCausalLM.forward)
-    shell = inspect.signature(modeling.ModelingV2Glm52Nvfp4Sm103Tp8.forward)
-    assert list(shell.parameters) == list(builtin.parameters)
-    for name, parameter in builtin.parameters.items():
-        assert shell.parameters[name].default == parameter.default, name
-        assert shell.parameters[name].kind == parameter.kind, name
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py
deleted file mode 100644
index 27afa09d82d3..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_selection.py
+++ /dev/null
@@ -1,573 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Which fast paths of the GLM-5.2 tp8 target run: the installed plan, the knobs the target
-refuses, the kernel registry, the capability log and the knobs folded into constants.
-
-Host-only: the knobs are set per test, the log goes into a list and the kernel imports are
-recorded, so nothing here needs a GPU or a kernel package.
-"""
-
-from __future__ import annotations
-
-import json
-import os
-import re
-import types
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
-    MODELING_V2_VALIDATE_ENV,
-    ModelingV2Mode,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-AUTO = ModelingV2Mode.AUTO
-ACCUMULATION = dispatch.KNOB_PREFIX + "ROUTED_MOE_INTO_HOP_ACCUMULATION"
-SCOPE = dispatch.KNOB_PREFIX + "DEFERRED_NORM_SCOPE"
-PARTS = ("sampler_store", "execution_stream_loop")
-
-# The record of every fast path but the parts that only deferred_input_prep brings along.
-SERVED_PATHS = (
-    "paths=absorbed_query,attention_output_into_hop,attention_output_tail,"
-    "attention_output_tail_prefetch,cross_layer_expert_prefetch,cross_rank_topk[rank-consistent],"
-    "cross_rank_topk_cluster_merge[rank-consistent],cross_rank_topk_global_emit[rank-consistent],"
-    "decode_prologue,deferred_input_prep,deferred_norm[shared],deferred_norm_evict_hints,dense_mlp,"
-    "draft_absorbed_query,draft_front,draft_live_rows,draft_loop,draft_routed_moe,draft_tail,"
-    "fused_kv_indexer_projection[kw-fused-a-order],hop_routing,host_page_retention,"
-    "indexer_layer_front,indexer_topk,oproj_weight_staging,post_attention_hop,post_moe_hop,"
-    "rope_kv_append,routed_moe,routed_moe_into_hop[fp32],router_gemm,shared_expert_prefetch,"
-    "shared_indexer_absorbed_query,sparse_mla_decode"
-)
-
-# The paths each path has no effect without: the closure of ``depends_on``.
-REQUIRED = {
-    "shared_indexer_absorbed_query": {"absorbed_query"},
-    "attention_output_tail_prefetch": {"attention_output_tail"},
-    "attention_output_into_hop": {"attention_output_tail", "post_attention_hop"},
-    "routed_moe_into_hop": {"post_moe_hop", "routed_moe"},
-    "deferred_norm": {"absorbed_query", "post_moe_hop"},
-    "cross_layer_expert_prefetch": {"routed_moe"},
-    "hop_routing": {
-        "attention_output_tail",
-        "post_attention_hop",
-        "attention_output_into_hop",
-        "routed_moe",
-    },
-    "oproj_weight_staging": {
-        "attention_output_tail",
-        "post_attention_hop",
-        "sparse_mla_decode",
-        "attention_output_into_hop",
-    },
-    "fused_kv_indexer_projection": {"absorbed_query"},
-    "host_page_retention": {"post_attention_hop"},
-    "shared_expert_prefetch": {
-        "attention_output_tail",
-        "post_moe_hop",
-        "routed_moe",
-        "post_attention_hop",
-        "attention_output_into_hop",
-        "routed_moe_into_hop",
-        "hop_routing",
-    },
-    "cross_rank_topk_cluster_merge": {"cross_rank_topk"},
-    "deferred_norm_evict_hints": {"absorbed_query", "post_moe_hop", "deferred_norm"},
-    "cross_rank_topk_global_emit": {
-        "rope_kv_append",
-        "cross_rank_topk",
-        "cross_rank_topk_cluster_merge",
-    },
-    "decode_prologue": {"draft_loop"},
-    "draft_live_rows": {"post_attention_hop", "draft_loop"},
-    "indexer_layer_front": {
-        "absorbed_query",
-        "indexer_topk",
-        "draft_absorbed_query",
-        "rope_kv_append",
-        "cross_rank_topk",
-    },
-}
-
-RECORD_PREFIX = "[glm_moe_dsa] dispatch: record "
-_RECORD = re.compile(r"^paths=(\S+) side_selectors=(\{.*\}) shared_env=(\{.*\})$")
-
-
-@pytest.fixture
-def lines(monkeypatch):
-    """A fresh selector state that logs into the returned list; no knob of the target is set."""
-    for name in list(os.environ):
-        if name.startswith(dispatch.SIDE_SELECTOR_PREFIXES) or name in dispatch.SHARED_ENV:
-            monkeypatch.delenv(name)
-    logged = []
-    monkeypatch.setattr(dispatch, "_log_info", logged.append)
-    monkeypatch.delenv(MODELING_V2_VALIDATE_ENV, raising=False)
-    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
-    monkeypatch.setattr(dispatch, "_VALIDATE", [False])
-    monkeypatch.setattr(dispatch, "_require", [False])
-    monkeypatch.setattr(dispatch, "_REGISTRY", {})
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    return logged
-
-
-def _parse(record):
-    m = _RECORD.match(record)
-    assert m, record
-    return m.group(1), json.loads(m.group(2)), json.loads(m.group(3))
-
-
-def _install(environ):
-    """Resolve and install the plan of ``environ``; return the paths it runs."""
-    dispatch.install(dispatch.resolve(AUTO, environ), environ)
-    return dispatch.installed()
-
-
-# --- the vocabulary -----------------------------------------------------------------------
-
-
-def test_the_selector_takes_every_fast_path_but_the_parts():
-    assert len(dispatch.FAST_PATH_NAMES) == 36
-    assert dispatch.SELECTOR_PATHS == tuple(n for n in dispatch.FAST_PATH_NAMES if n not in PARTS)
-    assert len(dispatch.SELECTOR_PATHS) == 34
-    assert dict(dispatch.PART_PATHS) == {
-        "prep": "deferred_input_prep",
-        "store": "sampler_store",
-        "stream": "execution_stream_loop",
-    }
-    assert dispatch.PARTS_ENV == "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_INPUT_PREP_PARTS"
-
-
-def test_the_prefetch_schedule_extends_the_attention_tail():
-    names = list(dispatch.FAST_PATH_NAMES)
-    assert names.index("attention_output_tail_prefetch") == names.index("attention_output_tail") + 1
-    prefetch = dispatch.FAST_PATHS[names.index("attention_output_tail_prefetch")]
-    assert prefetch.depends_on == ("attention_output_tail",)
-
-
-def test_every_fast_path_has_its_kernel_directory():
-    directories = set(dispatch.FAST_PATH_NAMES) - set(dispatch.KERNEL_DIRECTORY)
-    assert len(directories) == 28
-    assert len(dispatch.KERNEL_DIRECTORY) == 8
-    assert set(dispatch.KERNEL_DIRECTORY) <= set(dispatch.FAST_PATH_NAMES)
-    assert set(dispatch.KERNEL_DIRECTORY.values()) <= directories
-    assert dispatch.KERNELS_PACKAGE == dispatch.__name__.rpartition(".")[0] + ".kernels"
-
-
-def test_required_features_are_the_closure_of_the_table():
-    assert {k: set(v) for k, v in dispatch.REQUIRED_FEATURES.items()} == REQUIRED
-
-
-def test_the_variant_table_names_fast_paths_and_knobs():
-    for path, (name, choices) in dispatch.NUMERICAL_VARIANTS.items():
-        assert path in dispatch.FAST_PATH_NAMES
-        assert choices and len(set(choices)) == len(choices)
-        assert all(re.fullmatch(r"[a-z0-9][a-z0-9-]*", c) for c in choices)
-        assert name is None or name.startswith(dispatch.KNOB_PREFIX)
-    assert dispatch.NUMERICAL_VARIANTS["routed_moe_into_hop"] == (
-        ACCUMULATION,
-        ("fp32", "identity"),
-    )
-    # The deferred-norm scope is fixed in this target: no knob selects a variant.
-    assert dispatch.NUMERICAL_VARIANTS["deferred_norm"] == (None, ("shared",))
-
-
-# --- a path and its bases -----------------------------------------------------------------
-
-
-@pytest.mark.parametrize("path", dispatch.FAST_PATH_NAMES)
-def test_every_path_runs_on_its_minimal_base(lines, path):
-    """Turning off every path outside a path and its bases leaves exactly those running."""
-    kept = {path} | set(dispatch.REQUIRED_FEATURES.get(path, ()))
-    off = ",".join(sorted(set(dispatch.FAST_PATH_NAMES) - kept))
-    assert _install({dispatch.DISABLE_ENV: off} if off else {}) == kept
-
-
-# --- variants and the record --------------------------------------------------------------
-
-
-def test_the_served_record_names_every_variant(lines):
-    paths = frozenset(dispatch.SELECTOR_PATHS)
-    assert (
-        dispatch.feature_record(paths, {}) == f"{SERVED_PATHS} side_selectors={{}} shared_env={{}}"
-    )
-    assert dispatch.feature_variants(paths, {}) == {
-        "cross_rank_topk": "rank-consistent",
-        "cross_rank_topk_cluster_merge": "rank-consistent",
-        "cross_rank_topk_global_emit": "rank-consistent",
-        "deferred_norm": "shared",
-        "fused_kv_indexer_projection": "kw-fused-a-order",
-        "routed_moe_into_hop": "fp32",
-    }
-
-
-@pytest.mark.parametrize("mode", ["fp32", "identity", " Identity "])
-def test_the_accumulation_variant_follows_its_knob(mode):
-    paths = {"post_moe_hop", "routed_moe", "routed_moe_into_hop"}
-    env = {ACCUMULATION: mode}
-    listed, side, _ = _parse(dispatch.feature_record(paths, env))
-    assert listed == f"post_moe_hop,routed_moe,routed_moe_into_hop[{mode.strip().lower()}]"
-    assert side == env
-
-
-@pytest.mark.parametrize("scope", ["all", "shared", " Shared ", ""])
-def test_the_deferred_norm_variant_is_the_fixed_scope(scope):
-    """The record names the scope the hop runs; ``resolve`` refuses the scope variable."""
-    paths = {"absorbed_query", "post_moe_hop", "deferred_norm"}
-    env = {SCOPE: scope}
-    listed, side, _ = _parse(dispatch.feature_record(paths, env))
-    assert listed == "absorbed_query,deferred_norm[shared],post_moe_hop"
-    assert side == env
-
-
-@pytest.mark.parametrize(
-    "raw, env",
-    [
-        ("post_moe_hop,routed_moe,routed_moe_into_hop", {ACCUMULATION: "fp16"}),
-        # The finalized-row input stays a mode of the post-MoE hop, but no knob selects it.
-        ("post_moe_hop,routed_moe,routed_moe_into_hop", {ACCUMULATION: "row"}),
-    ],
-)
-def test_an_unknown_variant_of_a_selected_path_fails(raw, env):
-    with pytest.raises(ValueError, match="is not a .* variant"):
-        dispatch.feature_record(frozenset(raw.split(",")), env)
-
-
-def test_the_variant_knob_of_an_unselected_path_is_only_recorded():
-    env = {ACCUMULATION: "fp16", SCOPE: "x"}
-    listed, side, _ = _parse(dispatch.feature_record({"absorbed_query"}, env))
-    assert listed == "absorbed_query"
-    assert side == env
-
-
-def test_side_selectors_and_shared_env():
-    env = {
-        dispatch.DISABLE_ENV: "draft_tail",
-        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS": "100",
-        ACCUMULATION: "identity",
-        "TRTLLM_ENABLE_PDL": "1",
-        "TRTLLM_FORCE_MNNVL_AR": "0",
-        "PATH": "/usr/bin",
-    }
-    assert dispatch.side_selectors(env) == {
-        dispatch.DISABLE_ENV: "draft_tail",
-        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS": "100",
-        ACCUMULATION: "identity",
-    }
-    assert dispatch.shared_env(env) == {"TRTLLM_ENABLE_PDL": "1", "TRTLLM_FORCE_MNNVL_AR": "0"}
-    assert _parse(dispatch.feature_record({"absorbed_query"}, env)) == (
-        "absorbed_query",
-        dispatch.side_selectors(env),
-        dispatch.shared_env(env),
-    )
-
-
-# --- without a plan ---------------------------------------------------------------------
-
-
-def test_without_a_plan_no_fast_path_runs_and_nothing_is_logged(lines, monkeypatch):
-    """A process that does not build the target runs no fast path, whatever its environment."""
-    monkeypatch.setenv(dispatch.FEATURE_ENV, ",".join(dispatch.SELECTOR_PATHS))
-    monkeypatch.setenv(ACCUMULATION, "identity")
-    monkeypatch.setenv("TRTLLM_ENABLE_PDL", "1")
-    assert dispatch.installed() is None
-    assert dispatch.enabled_features() == frozenset()
-    assert not any(dispatch.is_enabled(path) for path in dispatch.FAST_PATH_NAMES)
-    assert not dispatch.validate_mode()
-    assert lines == []
-
-
-def test_asking_for_something_that_is_not_a_fast_path_fails(lines):
-    with pytest.raises(ValueError, match="'draft_tails' is not a fast path"):
-        dispatch.is_enabled("draft_tails")
-    assert lines == []
-
-
-# --- the installed plan -------------------------------------------------------------------
-
-
-def test_an_installed_plan_answers_every_query(lines):
-    plan = dispatch.resolve(AUTO, {})
-    lines.clear()
-    assert dispatch.installed() is None
-    dispatch.install(plan, {})
-    assert dispatch.installed() == frozenset(dispatch.FAST_PATH_NAMES)
-    assert all(dispatch.is_enabled(path) for path in dispatch.FAST_PATH_NAMES)
-    assert len(lines) == 1
-    listed, side, shared = _parse(lines[0].removeprefix(RECORD_PREFIX))
-    served = SERVED_PATHS.removeprefix("paths=").split(",")
-    assert listed.split(",") == sorted(served + list(PARTS))
-    assert (side, shared) == ({}, {})
-
-
-@pytest.mark.parametrize(
-    "part, other_parts",
-    [
-        ("deferred_input_prep", {"sampler_store", "execution_stream_loop"}),
-        ("sampler_store", {"deferred_input_prep", "execution_stream_loop"}),
-        ("execution_stream_loop", {"deferred_input_prep", "sampler_store"}),
-    ],
-)
-def test_each_part_of_the_input_preparation_is_a_path_of_its_own(lines, part, other_parts):
-    paths = _install({dispatch.DISABLE_ENV: part})
-    assert paths == frozenset(dispatch.FAST_PATH_NAMES) - {part}
-    assert other_parts <= paths
-
-
-@pytest.mark.parametrize("raw, validate", [("", False), ("0", False), ("1", True), (" On ", True)])
-def test_validate_follows_the_modeling_v2_switch(lines, monkeypatch, raw, validate):
-    monkeypatch.setenv(MODELING_V2_VALIDATE_ENV, raw)
-    plan = dispatch.resolve(AUTO, {})
-    assert plan.validate is validate
-    assert not dispatch.validate_mode()
-    dispatch.install(plan, {})
-    assert dispatch.validate_mode() is validate
-    dispatch.reset_fallback_state()
-    assert not dispatch.validate_mode()
-
-
-def test_disable_turns_off_a_path_and_its_dependents(lines):
-    dispatch.install(dispatch.resolve(AUTO, {dispatch.DISABLE_ENV: "post_attention_hop"}), {})
-    off = {path for path in dispatch.FAST_PATH_NAMES if not dispatch.is_enabled(path)}
-    assert off == {
-        "post_attention_hop",
-        "attention_output_into_hop",
-        "oproj_weight_staging",
-        "hop_routing",
-        "shared_expert_prefetch",
-        "draft_live_rows",
-        "host_page_retention",
-    }
-    listed, _, _ = _parse(lines[-1].removeprefix(RECORD_PREFIX))
-    assert set(listed.split(",")).isdisjoint(off)
-
-
-def test_a_bad_variant_of_an_enabled_path_fails_at_install(lines):
-    plan = dispatch.resolve(AUTO, {})
-    with pytest.raises(ValueError, match="is not a routed_moe_into_hop variant"):
-        dispatch.install(plan, {ACCUMULATION: "fp16"})
-    assert dispatch.installed() is None
-
-
-def test_the_variant_knob_of_a_disabled_path_is_only_recorded(lines):
-    env = {dispatch.DISABLE_ENV: "routed_moe_into_hop", ACCUMULATION: "fp16"}
-    dispatch.install(dispatch.resolve(AUTO, env), env)
-    assert not dispatch.is_enabled("routed_moe_into_hop")
-
-
-@pytest.mark.parametrize(
-    "knob", sorted(dispatch.REMOVED_KNOBS), ids=lambda k: k.removeprefix(dispatch.KNOB_PREFIX)
-)
-@pytest.mark.parametrize("value", ["deferred_input_prep", ""], ids=["set", "empty"])
-def test_a_removed_knob_is_refused_with_what_to_use_instead(lines, knob, value):
-    with pytest.raises(ValueError) as excinfo:
-        dispatch.resolve(AUTO, {knob: value})
-    assert str(excinfo.value) == (
-        f"this target does not read {knob} ({dispatch.REMOVED_KNOBS[knob]}); unset them"
-    )
-    assert dispatch.installed() is None
-    assert lines == []
-
-
-def test_every_removed_knob_set_is_named_in_one_refusal(lines):
-    env = {knob: "1" for knob in dispatch.REMOVED_KNOBS}
-    env[dispatch.DISABLE_ENV] = "draft_tail"
-    with pytest.raises(ValueError) as excinfo:
-        dispatch.resolve(AUTO, env)
-    message = str(excinfo.value)
-    assert all(f"{knob} (" in message for knob in dispatch.REMOVED_KNOBS)
-    assert f"{dispatch.DISABLE_ENV} (" not in message
-
-
-def test_the_removed_knobs_point_at_what_replaces_them():
-    hints = dispatch.REMOVED_KNOBS
-    assert set(hints) == {
-        dispatch.FEATURE_ENV,
-        dispatch.PARTS_ENV,
-        dispatch.KNOB_PREFIX + "DEFERRED_INPUT_PREP_DEBUG",
-        dispatch.KNOB_PREFIX + "DRAFT_LOOP_VERIFY",
-        dispatch.KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_STATS",
-        dispatch.KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_REPORT_EVERY",
-        dispatch.KNOB_PREFIX + "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE",
-    }
-    assert "DISABLE" in hints[dispatch.KNOB_PREFIX + "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE"]
-    assert dispatch.DISABLE_ENV in hints[dispatch.FEATURE_ENV]
-    assert dispatch.DISABLE_ENV in hints[dispatch.PARTS_ENV]
-    for part, path in dispatch.PART_PATHS.items():
-        assert f"{part} = {path}" in hints[dispatch.PARTS_ENV]
-    for knob in ("DEFERRED_INPUT_PREP_DEBUG", "DRAFT_LOOP_VERIFY"):
-        assert MODELING_V2_VALIDATE_ENV in hints[dispatch.KNOB_PREFIX + knob]
-
-
-def test_resetting_the_fallback_state_uninstalls_the_plan(lines):
-    dispatch.install(dispatch.resolve(AUTO, {dispatch.DISABLE_ENV: "draft_tail"}), {})
-    assert not dispatch.is_enabled("draft_tail")
-    dispatch.reset_fallback_state()
-    assert dispatch.installed() is None
-
-
-# --- the kernel registry ------------------------------------------------------------------
-
-
-def test_register_and_look_up(lines, monkeypatch):
-    monkeypatch.setattr(
-        dispatch, "importlib", types.SimpleNamespace(import_module=lambda name: pytest.fail(name))
-    )
-
-    def impl():
-        return 42
-
-    assert dispatch.register("draft_front", "draft_front", impl) is impl
-    assert dispatch.get_impl("draft_front", "draft_front")() == 42
-    with pytest.raises(KeyError, match="'draft_front' does not register 'missing'"):
-        dispatch.get_impl("draft_front", "missing")
-    with pytest.raises(ValueError, match="'not_a_path' is not a fast path"):
-        dispatch.register("not_a_path", "draft_front", impl)
-
-
-@pytest.mark.parametrize("path", dispatch.FAST_PATH_NAMES)
-def test_get_impl_imports_the_kernel_directory_of_the_path(lines, monkeypatch, path):
-    imported = []
-    monkeypatch.setattr(dispatch, "importlib", types.SimpleNamespace(import_module=imported.append))
-    with pytest.raises(KeyError, match=f"{path!r} does not register 'entry'"):
-        dispatch.get_impl(path, "entry")
-    directory = dispatch.KERNEL_DIRECTORY.get(path, path)
-    assert imported == [f"{dispatch.KERNELS_PACKAGE}.{directory}"]
-
-
-def test_capability_lines_are_logged_once_per_key(lines):
-    assert dispatch.log_capability("draft_front", "draft_front", True)
-    assert not dispatch.log_capability("attention_output_tail", "layer3", False, "dense")
-    # One line per (path, key) and process; the call still returns its own decision.
-    assert not dispatch.log_capability("draft_front", "draft_front", False)
-    assert lines == [
-        "[glm_moe_dsa] draft_front/draft_front: engaged",
-        "[glm_moe_dsa] attention_output_tail/layer3: native fallback (dense)",
-    ]
-
-
-# --- knobs folded into constants --------------------------------------------------------
-
-# Every variable with the knob prefix that the kernels of this target once read; they run at
-# fixed launch parameters now, so ``resolve`` refuses each of them when it is set.
-FOLDED_KNOBS = tuple(
-    dispatch.KNOB_PREFIX + name
-    for name in (
-        "DEFERRED_NORM_HOP_CARVEOUT",
-        "DEFERRED_NORM_PROJECTION_CONFIG",
-        "DEFERRED_NORM_EVICT_HINTS_HOP_CONFIG",
-        "DEFERRED_NORM_EVICT_HINTS_PROJECTION_CONFIG",
-        "ABSORBED_QUERY_FORM",
-        "FUSED_KV_INDEXER_PROJECTION_LAYERS",
-        "FUSED_KV_INDEXER_PROJECTION_CONFIG",
-        "INDEXER_LAYER_FRONT_SKIP_REUSED_INDEXER_Q",
-        "CROSS_RANK_TOPK_LOGITS_PDL",
-        "ATTENTION_OUTPUT_TAIL_PDL_ENTRY",
-        "OPROJ_WEIGHT_STAGING_CONFIG",
-        "ROUTER_GEMM_EXPERTS_PER_CTA",
-        "ROUTER_GEMM_PAIRED_FMA",
-        "ROUTED_MOE_SCHEDULE",
-        "ROUTED_MOE_EXPERT_CLUSTER_CTAS",
-        "ROUTED_MOE_FINALIZE_MODE",
-        "ROUTED_MOE_PDL_TRIGGER",
-        "ROUTED_MOE_MIN_CLUSTERS",
-        "ROUTED_MOE_STAGES",
-        "ROUTED_MOE_GRID_WIDE_L2_PREFETCH",
-        "ROUTED_MOE_MIN_FC1_KBLOCKS",
-        "ROUTED_MOE_ACCUMULATOR_STAGES",
-        "ROUTED_MOE_HANDOFF",
-        "ROUTED_MOE_DELAY_SITES",
-        "ROUTED_MOE_DELAY_NS",
-        "ROUTED_MOE_EXPERT_CLUSTER_L2_PREFETCH",
-        "ROUTED_MOE_TWO_ROUND_CLUSTERS",
-        "SHARED_EXPERT_PREFETCH_CLUSTERS",
-        "SHARED_EXPERT_PREFETCH_WEIGHTS",
-        "DRAFT_ROUTED_MOE_FC1_CONFIG",
-        "DRAFT_ROUTED_MOE_FC2_CONFIG",
-        "DRAFT_ROUTED_MOE_FLAGS",
-        "DRAFT_TAIL_MODE",
-        "DRAFT_TAIL_EXCHANGE",
-        "DRAFT_TAIL_VARIANT",
-        "DRAFT_TAIL_CTAS_PER_SM",
-        "DRAFT_TAIL_NUM_CTAS",
-        "DRAFT_TAIL_PDL",
-        "DRAFT_TAIL_TIMEOUT_CYCLES",
-        "DRAFT_TAIL_KEEP_PERMILLE",
-        "CROSS_LAYER_EXPERT_PREFETCH_START_POINT",
-        "CROSS_LAYER_EXPERT_PREFETCH_CHUNK",
-        "CROSS_LAYER_EXPERT_PREFETCH_CTAS",
-        "CROSS_LAYER_EXPERT_PREFETCH_DELAY_NS",
-        "CROSS_LAYER_EXPERT_PREFETCH_FUSED",
-        "DRAFT_LOOP_PARTS",
-        "DEFERRED_NORM_SCOPE",
-        "CROSS_LAYER_EXPERT_PREFETCH_DEPTH",
-        "CROSS_LAYER_EXPERT_PREFETCH_PREDICTOR",
-        "CROSS_LAYER_EXPERT_PREFETCH_ISSUE_KERNEL",
-        "CROSS_LAYER_EXPERT_PREFETCH_ROUTED_KERNEL_CTAS",
-        "DRAFT_ABSORBED_QUERY_SKIP_REUSED_INDEXER_Q",
-        "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE",
-        "ROPE_KV_APPEND_METADATA_HINTS",
-        "ROPE_KV_APPEND_METADATA_PREFETCH",
-    )
-)
-
-
-def test_the_knobs_of_the_target_are_the_plan_knobs():
-    assert dispatch.KNOBS == {
-        dispatch.DISABLE_ENV,
-        ACCUMULATION,
-        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS",
-    }
-    assert not set(FOLDED_KNOBS) & dispatch.KNOBS
-
-
-@pytest.mark.parametrize("knob", FOLDED_KNOBS, ids=lambda k: k.removeprefix(dispatch.KNOB_PREFIX))
-def test_a_folded_knob_is_refused_even_at_its_value(lines, knob):
-    with pytest.raises(ValueError) as excinfo:
-        dispatch.resolve(AUTO, {knob: "0"})
-    hint = dispatch.REMOVED_KNOBS.get(
-        knob, "the fast paths of this target run at fixed launch parameters"
-    )
-    assert str(excinfo.value) == f"this target does not read {knob} ({hint}); unset them"
-    assert dispatch.installed() is None
-    assert lines == []
-
-
-def test_the_knobs_of_the_target_resolve(lines):
-    env = {
-        dispatch.DISABLE_ENV: "draft_tail",
-        ACCUMULATION: "identity",
-        dispatch.KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS": "100",
-    }
-    plan = dispatch.resolve(AUTO, env)
-    assert plan.disabled == {"draft_tail"}
-    assert plan.knobs == tuple(sorted(env.items()))
-
-
-# Shared TensorRT-LLM switches a kernel of this target may read next to the plan.
-_SHARED_READS = {"TRTLLM_ENABLE_PDL", "TRTLLM_FORCE_MNNVL_AR", "TLLM_TEST_MNNVL"}
-
-
-def test_only_dispatch_reads_the_knobs():
-    """No module of the family but ``dispatch`` reads the environment, except the shared switches."""
-    root = os.path.dirname(os.path.dirname(dispatch.__file__))
-    reads = []
-    for directory, _, files in os.walk(root):
-        for name in sorted(files):
-            if not name.endswith(".py") or os.path.join(directory, name) == dispatch.__file__:
-                continue
-            path = os.path.join(directory, name)
-            with open(path) as f:
-                text = f.read()
-            for match in re.finditer(r"os\.(?:environ|getenv)\b", text):
-                read = re.match(
-                    r"os\.(?:environ(?:\.get)?|getenv)\s*[\[(]\s*[\"']([A-Z_]+)[\"']",
-                    text[match.start() :],
-                )
-                if read is None or read.group(1) not in _SHARED_READS:
-                    line = text.count("\n", 0, match.start()) + 1
-                    reads.append(f"{os.path.relpath(path, root)}:{line}")
-    assert reads == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py
deleted file mode 100644
index 951159049d65..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_spec_worker.py
+++ /dev/null
@@ -1,514 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 tp8 target's speculative worker: its adoption, the phase gate of every override,
-the draft loop it runs on a decode step, and the draft layer's live-row arming.
-
-Host-only: the built-in worker is created without its constructor and carries only the
-attributes the overrides read. The kernel entry points are stand-ins, and the built-in methods
-an override falls back to or calls are recorders.
-"""
-
-from __future__ import annotations
-
-import ast
-import contextlib
-import inspect
-import types
-from pathlib import Path
-
-import pytest
-import torch
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-    layers,
-    spec_worker,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows import (  # noqa: E501
-    step0,
-)
-from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker, MTPEagleWorker
-from tensorrt_llm._torch.speculative.mtp_dynamic_tree import MTPEagleDynamicTreeWorker
-
-pytestmark = pytest.mark.cpu_only
-
-PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
-GlmMtpWorker = spec_worker.GlmMtpWorker
-OVERRIDES = ["_forward_impl", "_forward_linear_draft_loop", "_run_draft_forward"]
-
-
-def _bare(cls, **attributes):
-    obj = cls.__new__(cls)
-    nn.Module.__init__(obj)
-    for name, value in attributes.items():
-        setattr(obj, name, value)
-    return obj
-
-
-def _keys(module):
-    return (
-        [name for name, _ in module.named_parameters()],
-        [name for name, _ in module.named_buffers()],
-        [name for name, _ in module.named_modules()],
-    )
-
-
-# --- adoption -----------------------------------------------------------------------------
-
-
-def test_adoption_reclasses_the_builtin_worker_in_place():
-    worker = _bare(MTPEagleWorker, is_mtp_eagle=True)
-    worker.register_buffer("probe", torch.zeros(1))
-    before = _keys(worker)
-
-    assert GlmMtpWorker.adopt(worker) is worker
-
-    assert type(worker) is GlmMtpWorker
-    assert worker._draft_loop_ops is None and "_draft_loop_ops" not in vars(worker)
-    assert _keys(worker) == before
-
-
-def _other_worker():
-    return _bare(type("_OtherMTPEagleWorker", (MTPEagleWorker,), {}))
-
-
-@pytest.mark.parametrize(
-    "make, got",
-    [
-        (lambda: _bare(Eagle3OneModelWorker), "Eagle3OneModelWorker"),
-        (lambda: _bare(MTPEagleDynamicTreeWorker), "MTPEagleDynamicTreeWorker"),
-        (_other_worker, "_OtherMTPEagleWorker"),
-        (lambda: GlmMtpWorker.adopt(_bare(MTPEagleWorker)), "GlmMtpWorker"),
-    ],
-    ids=["the Eagle3 worker", "the dynamic-tree worker", "a subclass", "an adopted worker"],
-)
-def test_adoption_refuses_anything_but_the_exact_builtin_worker(make, got):
-    with pytest.raises(TypeError, match=f"adopts the built-in MTPEagleWorker, got {got}$"):
-        GlmMtpWorker.adopt(make())
-
-
-def test_the_class_declares_no_state_of_its_own():
-    assert not {"__slots__", "__init__", "forward"} & set(vars(GlmMtpWorker))
-    assert GlmMtpWorker._draft_loop_ops is None
-
-
-@pytest.mark.parametrize("has_worker", [True, False], ids=["a worker", "no worker"])
-def test_the_worker_is_adopted_after_the_draft_layer(monkeypatch, has_worker):
-    order = []
-    for cls, label in (
-        (layers.GlmDecoderLayer, "decoder layer"),
-        (layers.GlmMtpLayer, "draft layer"),
-        (GlmMtpWorker, "worker"),
-    ):
-        monkeypatch.setattr(
-            cls, "adopt", classmethod(lambda c, module, label=label: order.append((label, module)))
-        )
-    core = types.SimpleNamespace(num_hidden_layers=2, layers=["layer 0", "layer 1", "draft"])
-    worker = object() if has_worker else None
-
-    modeling.adopt_model(types.SimpleNamespace(model=core, spec_worker=worker), plan=None)
-
-    expected = [
-        ("decoder layer", "layer 0"),
-        ("decoder layer", "layer 1"),
-        ("draft layer", "draft"),
-    ]
-    assert order == expected + ([("worker", worker)] if has_worker else [])
-
-
-# --- the phase gate -----------------------------------------------------------------------
-
-
-def _method(name):
-    tree = ast.parse(Path(spec_worker.__file__).read_text())
-    cls = next(c for c in tree.body if isinstance(c, ast.ClassDef) and c.name == "GlmMtpWorker")
-    return next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == name)
-
-
-@pytest.mark.parametrize("name", OVERRIDES)
-def test_each_override_starts_with_the_phase_gate(name):
-    """``if dispatch.phase() is not DECODE: return super().<name>(<every built-in parameter by
-    name>)``, after the built-in docstring if the method has one."""
-    node = _method(name)
-    docstring = isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant)
-    gate = node.body[1] if docstring else node.body[0]
-    assert ast.unparse(gate.test) == "dispatch.phase() is not dispatch.Phase.DECODE"
-    assert len(gate.body) == 1 and not gate.orelse
-    call = gate.body[0].value
-    assert ast.unparse(call.func) == f"super().{name}" and not call.args
-    parameters = inspect.signature(getattr(Eagle3OneModelWorker, name)).parameters
-    assert [k.arg for k in call.keywords] == [p for p in parameters if p != "self"]
-    assert all(k.arg == ast.unparse(k.value) for k in call.keywords)
-    assert list(inspect.signature(getattr(GlmMtpWorker, name)).parameters) == list(parameters)
-
-
-def _record(monkeypatch, owner, name, calls):
-    def recorder(*args, **kwargs):
-        calls.append((name, args, kwargs))
-        return f"{owner.__name__}.{name}"
-
-    monkeypatch.setattr(owner, name, recorder)
-
-
-@pytest.mark.parametrize("phase", [PREFILL, None], ids=["prefill", "outside a step"])
-def test_off_decode_every_override_runs_the_builtin_method(monkeypatch, phase):
-    worker = GlmMtpWorker.adopt(_bare(MTPEagleWorker, is_mtp_eagle=True))
-    calls = []
-    for name in OVERRIDES:
-        _record(monkeypatch, Eagle3OneModelWorker, name, calls)
-    monkeypatch.setattr(dispatch, "is_enabled", lambda path: pytest.fail(path))
-    monkeypatch.setattr(dispatch, "get_impl", lambda *a: pytest.fail(str(a)))
-
-    with dispatch.step(phase) if phase is not None else contextlib.nullcontext():
-        results = [
-            worker._forward_impl("ids", "pos", "h", "logits", "md", "spec", "draft", "rm"),
-            worker._forward_linear_draft_loop(
-                "inputs", "md", "spec", "draft", "kv", 0, 2, "accepted", "ranks"
-            ),
-            worker._run_draft_forward("draft", "inputs", "spec", 1),
-        ]
-
-    assert results == [f"Eagle3OneModelWorker.{name}" for name in OVERRIDES]
-    assert all(args == (worker,) for _, args, _ in calls)
-    assert [(name, kwargs) for name, _, kwargs in calls] == [
-        (
-            "_forward_impl",
-            dict(
-                input_ids="ids",
-                position_ids="pos",
-                hidden_states="h",
-                logits="logits",
-                attn_metadata="md",
-                spec_metadata="spec",
-                draft_model="draft",
-                resource_manager="rm",
-            ),
-        ),
-        (
-            "_forward_linear_draft_loop",
-            dict(
-                inputs="inputs",
-                attn_metadata="md",
-                spec_metadata="spec",
-                draft_model="draft",
-                draft_kv_cache_manager="kv",
-                num_contexts=0,
-                batch_size=2,
-                num_accepted_tokens="accepted",
-                original_all_rank_num_tokens="ranks",
-            ),
-        ),
-        (
-            "_run_draft_forward",
-            dict(draft_model="draft", inputs="inputs", spec_metadata="spec", step_idx=1),
-        ),
-    ]
-    assert "_draft_loop_ops" not in vars(worker)
-
-
-# --- the decode draft loop ----------------------------------------------------------------
-
-BATCH = 2
-HIDDEN = 4
-
-
-class _DraftLoopOps:
-    """Stands in for the draft-loop operations of one forward; records each call."""
-
-    def __init__(self, events, next_new_tokens="ops next new tokens"):
-        self.events = events
-        self.batch_size = BATCH
-        self.gather_ids = torch.arange(BATCH, dtype=torch.int64)
-        self.next_new_tokens = next_new_tokens
-
-    def accept(self, logits):
-        self.events.append(("accept", logits))
-        return "ops accepted", "ops num accepted"
-
-    def first_drafter_inputs(self, position_ids, hidden_states):
-        self.events.append(("first_drafter_inputs", tuple(position_ids.shape), hidden_states))
-        return "ops first inputs"
-
-    def step0_rows(self, hidden_states):
-        self.events.append(("step0_rows", hidden_states.shape[0]))
-        return self.gather_ids, hidden_states[:BATCH]
-
-    def step_update(self, attn_metadata, step, num_accepted, position_ids, gather_ids):
-        self.events.append(("step_update", step, position_ids is not None, gather_ids is not None))
-        return types.SimpleNamespace(position_ids=f"positions after step {step}")
-
-    def finalize(self, tokens):
-        self.events.append(("finalize", list(tokens)))
-        return "ops stacked"
-
-
-def _worker(events, **methods):
-    worker = GlmMtpWorker.adopt(
-        _bare(
-            MTPEagleWorker,
-            is_mtp_eagle=True,
-            sa_enhancer=None,
-            guided_decoder=None,
-            model_config=None,
-            _uses_external_shared_target_kv=False,
-            _is_mamba_hybrid_cache=False,
-            _auxiliary_state_handlers=[],
-        )
-    )
-    for name, result in methods.items():
-
-        def method(*args, _name=name, _result=result, **kwargs):
-            events.append((_name,))
-            return _result(worker) if callable(_result) else _result
-
-        setattr(worker, name, method)
-    return worker
-
-
-class _FastPaths:
-    """``dispatch.is_enabled`` / ``get_impl`` with stand-in draft-loop entry points."""
-
-    def __init__(self, events, enabled, entries):
-        self.events, self.enabled, self.entries = events, frozenset(enabled), entries
-
-    def is_enabled(self, path):
-        assert path in dispatch.FAST_PATH_NAMES, path
-        return path in self.enabled
-
-    def get_impl(self, path, name):
-        assert path in self.enabled, path
-        entry = self.entries[(path, name)]
-
-        def call(*args):
-            self.events.append((name,) + args)
-            return entry(*args)
-
-        return call
-
-
-@pytest.mark.parametrize("mode", ["engaged", "ops without next new tokens", "declined", "off"])
-def test_the_decode_forward_runs_the_draft_loop_operations(monkeypatch, mode):
-    events = []
-    ops = _DraftLoopOps(events, None if mode == "ops without next new tokens" else "ops tokens")
-    worker = _worker(
-        events,
-        _execute_guided_decoder_if_present=None,
-        sample_and_accept_draft_tokens=("accepted", "num accepted"),
-        _prepare_attn_metadata_for_spec_dec=None,
-        prepare_1st_drafter_inputs="first inputs",
-        get_draft_kv_cache_manager="draft kv",
-        _forward_draft_loop=lambda w: ("draft tokens", w._draft_loop_ops),
-        _restore_attn_metadata_from_spec_dec=None,
-        _prepare_next_new_tokens="next new tokens",
-    )
-    paths = _FastPaths(
-        events,
-        () if mode == "off" else {"draft_loop"},
-        {("draft_loop", "draft_loop_ops"): lambda *a: None if mode == "declined" else ops},
-    )
-    monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-    monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-    md = types.SimpleNamespace(num_seqs=BATCH, num_contexts=0, all_rank_num_tokens=None)
-    spec = types.SimpleNamespace(
-        runtime_draft_len=2, spec_resource_manager=None, batch_indices_cuda="batch indices"
-    )
-
-    with dispatch.step(DECODE):
-        out = worker._forward_impl(
-            "ids", torch.zeros(1, 6), "hidden", "logits", md, spec, "draft model", "rm"
-        )
-
-    engaged = mode in ("engaged", "ops without next new tokens")
-    start = [("_execute_guided_decoder_if_present",)]
-    if mode != "off":
-        start.append(("draft_loop_ops", worker, md, spec))
-    if engaged:
-        assert events == start + [
-            ("accept", "logits"),
-            ("_prepare_attn_metadata_for_spec_dec",),
-            ("first_drafter_inputs", (6,), "hidden"),
-            ("get_draft_kv_cache_manager",),
-            ("_forward_draft_loop",),
-            ("_restore_attn_metadata_from_spec_dec",),
-        ] + ([("_prepare_next_new_tokens",)] if ops.next_new_tokens is None else [])
-        assert out["new_tokens"] == "ops accepted"
-        assert out["new_tokens_lens"] == "ops num accepted"
-        assert out["next_draft_tokens"] == ("draft tokens", ops)
-        assert out["next_new_tokens"] == (ops.next_new_tokens or "next new tokens")
-    else:
-        assert events == start + [
-            ("sample_and_accept_draft_tokens",),
-            ("_prepare_attn_metadata_for_spec_dec",),
-            ("prepare_1st_drafter_inputs",),
-            ("get_draft_kv_cache_manager",),
-            ("_forward_draft_loop",),
-            ("_restore_attn_metadata_from_spec_dec",),
-            ("_prepare_next_new_tokens",),
-        ]
-        assert out["new_tokens"] == "accepted"
-        assert out["next_draft_tokens"] == ("draft tokens", None)
-        assert out["next_new_tokens"] == "next new tokens"
-    # The operations live for one forward only.
-    assert worker._draft_loop_ops is None and md.use_spec_decoding is True
-    assert out["logits"] == "logits"
-
-
-class _DraftLayerForward:
-    """The draft layer's decode forward as the draft loop sees it: it records the rows it was
-    armed with and, when armed, returns only the live rows, as the live-row hop does."""
-
-    def __init__(self, layer, compacts):
-        self.layer, self.compacts, self.calls = layer, compacts, []
-
-    def __call__(self, **kwargs):
-        rows = self.layer._draft_live_rows
-        self.calls.append((rows, kwargs["embed_tokens"], kwargs["all_rank_num_tokens"]))
-        if rows is not None and self.compacts:
-            self.layer._draft_live_rows_compacted = True
-            return torch.ones(rows.shape[0], HIDDEN)
-        return torch.zeros(3 * BATCH if len(self.calls) == 1 else BATCH, HIDDEN)
-
-
-@pytest.mark.parametrize("compacts", [True, False], ids=["compacted", "kept every row"])
-def test_the_decode_draft_loop_runs_the_draft_fast_paths(monkeypatch, compacts):
-    """Each draft-loop fast path is decided once before the steps; step 0 runs armed with the
-    draft-loop operations' live rows; the draft tail makes each step's tokens; the step updates
-    replace the native position advance and metadata update."""
-    events = []
-    ops = _DraftLoopOps(events)
-    draft_layer = _bare(layers.GlmMtpLayer)
-    forward = _DraftLayerForward(draft_layer, compacts)
-    draft_layer.forward = forward
-    draft_model = types.SimpleNamespace(embed_tokens="embedding table", mtp_layers=[draft_layer])
-    worker = _worker(events, draft_kv_cache_context=lambda w: contextlib.nullcontext(md))
-    worker._draft_loop_ops = ops
-    tail_tokens = iter(["tokens 0", "tokens 1"])
-    paths = _FastPaths(
-        events,
-        {"draft_loop", "draft_tail", "draft_live_rows"},
-        {
-            ("draft_tail", "select_draft_tail"): lambda *a: (
-                lambda rows: (events.append(("draft tail", rows.shape[0])), next(tail_tokens))[1]
-            ),
-            ("draft_live_rows", "select_step0"): lambda w, dm, s, o: step0.Step0LiveRows(
-                dm.mtp_layers[0], o.gather_ids
-            ),
-            ("draft_loop", "draft_embedding"): lambda table: ("draft embedding", table),
-        },
-    )
-    monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-    monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-    md = types.SimpleNamespace()
-    spec = types.SimpleNamespace(
-        runtime_draft_len=2,
-        all_rank_num_tokens=[3 * BATCH],
-        subseq_all_rank_num_tokens=[BATCH],
-        wants_advanced_draft_sampling=False,
-    )
-    inputs = {"input_ids": "ids", "position_ids": "positions", "hidden_states": "h"}
-
-    with dispatch.step(DECODE):
-        out = worker._forward_linear_draft_loop(
-            inputs, md, spec, draft_model, None, 0, BATCH, "num accepted", None
-        )
-
-    assert out == "ops stacked"
-    assert events[0] == ("select_draft_tail", worker, draft_model, spec)
-    assert events[1] == ("select_step0", worker, draft_model, spec, ops)
-    assert events[2] == ("draft_kv_cache_context",)
-    step_events = [e for e in events if e[0] in ("step0_rows", "draft tail", "step_update")]
-    assert step_events == ([] if compacts else [("step0_rows", 3 * BATCH)]) + [
-        ("draft tail", BATCH),
-        ("step_update", 0, True, True),
-        ("draft tail", BATCH),
-        ("step_update", 1, False, False),
-    ]
-    assert events[-1] == ("finalize", ["tokens 0", "tokens 1"])
-    # Step 0 ran armed with the live rows; step 1 did not; the layer is disarmed afterwards.
-    armed = [rows for rows, _, _ in forward.calls]
-    assert armed[0] is ops.gather_ids and armed[1] is None
-    assert [c[1:] for c in forward.calls] == [
-        (("draft embedding", "embedding table"), [3 * BATCH]),
-        (("draft embedding", "embedding table"), [BATCH]),
-    ]
-    assert draft_layer._draft_live_rows is None and draft_layer._draft_live_rows_compacted is False
-
-
-@pytest.mark.parametrize("draft_loop", [True, False], ids=["draft loop", "no draft loop"])
-def test_the_decode_draft_forward_takes_the_draft_loop_embedding(monkeypatch, draft_loop):
-    events = []
-    calls = []
-    layer = lambda **kwargs: calls.append(kwargs) or "draft layer output"  # noqa: E731
-    draft_model = types.SimpleNamespace(embed_tokens="embedding table", mtp_layers=[layer])
-    worker = _worker(events)
-    paths = _FastPaths(
-        events,
-        {"draft_loop"} if draft_loop else (),
-        {("draft_loop", "draft_embedding"): lambda table: ("draft embedding", table)},
-    )
-    monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-    monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-    spec = types.SimpleNamespace(all_rank_num_tokens=[6], subseq_all_rank_num_tokens=[2])
-
-    with dispatch.step(DECODE):
-        out = worker._run_draft_forward(draft_model, {"input_ids": "ids"}, spec, 0)
-
-    assert out == ("draft layer output", None)
-    table = ("draft embedding", "embedding table") if draft_loop else "embedding table"
-    assert calls == [dict(embed_tokens=table, all_rank_num_tokens=[6], input_ids="ids")]
-
-
-# --- the draft layer's live-row arming ----------------------------------------------------
-
-
-def _recording_draft_layer(writes):
-    class _Recording(layers.GlmMtpLayer):
-        def __setattr__(self, name, value):
-            if name.startswith("_draft_live_rows"):
-                writes.append((name, value))
-            super().__setattr__(name, value)
-
-    return _bare(_Recording)
-
-
-def test_the_draft_layer_is_armed_and_disarmed_through_its_api():
-    layer = _bare(layers.GlmMtpLayer)
-    rows = torch.arange(BATCH)
-    layer.arm_live_rows(rows)
-    assert layer._draft_live_rows is rows and layer._draft_live_rows_compacted is False
-    layer._draft_live_rows_compacted = True
-    assert layer.disarm_live_rows() is True
-    assert layer._draft_live_rows is None and layer._draft_live_rows_compacted is False
-    assert layer.disarm_live_rows() is False
-
-
-@pytest.mark.parametrize("outcome", ["compacted", "kept every row", "raised"])
-def test_step0_arming_writes_the_layer_state_in_the_builtin_order(outcome):
-    """Arm (rows, then not compacted), the draft layer's own write, then clear (rows, then
-    compacted), also when the step raises."""
-    writes = []
-    layer = _recording_draft_layer(writes)
-    rows = torch.arange(BATCH)
-    live = step0.Step0LiveRows(layer, rows)
-
-    with pytest.raises(RuntimeError) if outcome == "raised" else contextlib.nullcontext():
-        with live.armed() as armed:
-            assert armed is live and layer._draft_live_rows is rows
-            if outcome != "kept every row":
-                layer._draft_live_rows_compacted = True
-            if outcome == "raised":
-                raise RuntimeError("draft step failed")
-
-    assert live.compacted is (outcome == "compacted")
-    own = [("_draft_live_rows_compacted", True)] if outcome != "kept every row" else []
-    assert writes == [
-        ("_draft_live_rows", rows),
-        ("_draft_live_rows_compacted", False),
-        *own,
-        ("_draft_live_rows", None),
-        ("_draft_live_rows_compacted", False),
-    ]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py
deleted file mode 100644
index 97e626cf595b..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_targets.py
+++ /dev/null
@@ -1,597 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 tp8 core and its two targets, and what the shell keeps from the built-in model.
-
-Host-only: the core is adopted on a built-in text model created without its constructor
-(which builds CUDA streams and the decoder layers), and the built-in forward is replaced by
-a recorder, so each test sees exactly which call a step makes. The shell's construction runs
-with the built-in constructor stubbed, stand-in decoder layers, a recorder in place of
-``adopt_model`` (``test_modeling_v2_glm_layers.py`` covers it) and a stand-in for the
-sparse-MLA decode FMHA library.
-"""
-
-from __future__ import annotations
-
-import inspect
-import os
-import sys
-import types
-
-import pytest
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2 import _router_index
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8 import (  # noqa: E501
-    modeling,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-    runtime,
-)
-from tensorrt_llm._torch.attention.backends.fmha import manager as fmha_manager
-from tensorrt_llm._torch.attention.backends.fmha import registry as fmha_registry
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of
-
-pytestmark = pytest.mark.cpu_only
-
-
-def _builtin_text_model() -> DeepseekV3Model:
-    """A built-in text model as far as class and module state go, without its layers."""
-    model = DeepseekV3Model.__new__(DeepseekV3Model)
-    nn.Module.__init__(model)
-    return model
-
-
-class _Hop:
-    """Stands in for a post-MoE hop object a decoder layer carries."""
-
-    def __init__(self, engaged: bool = True) -> None:
-        self.engaged = engaged
-        self.impl = "cluster"
-
-
-class _Layer(nn.Module):
-    """Stands in for a decoder layer with the fast-path attributes the constructors set."""
-
-    def __init__(self, post_moe_hop=None, router_gemm=None) -> None:
-        super().__init__()
-        self.self_attn = nn.Module()
-        self.self_attn._attn_tail = None
-        self.mlp = nn.Module()
-        self.mlp.router_gemm = router_gemm
-        self.post_attn_hop = None
-        self.post_moe_hop = post_moe_hop
-
-
-class _Metadata:
-    def __init__(self, num_contexts: int) -> None:
-        self.num_contexts = num_contexts
-
-
-def test_adopt_reclasses_the_builtin_text_model_in_place():
-    model = _builtin_text_model()
-    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(model)
-    assert core is model
-    assert type(core) is modeling.GlmMoeDsaTp8ModelingV2
-    assert isinstance(core, DeepseekV3Model)
-    assert isinstance(core._prefill, modeling.PrefillTarget) and core._prefill.core is core
-    assert isinstance(core._decode, modeling.DecodeTarget) and core._decode.core is core
-    # Plain attributes only: no module or parameter joins the built-in tree.
-    assert list(core.named_modules()) == [("", core)]
-    assert list(core.named_parameters()) == []
-
-
-def test_adopt_refuses_anything_but_the_builtin_text_model():
-    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
-    with pytest.raises(TypeError, match="GlmMoeDsaTp8ModelingV2"):
-        modeling.GlmMoeDsaTp8ModelingV2.adopt(core)
-    with pytest.raises(TypeError, match="Module"):
-        modeling.GlmMoeDsaTp8ModelingV2.adopt(nn.Module())
-
-
-def _recorder(calls):
-    def forward(
-        self,
-        attn_metadata,
-        input_ids=None,
-        position_ids=None,
-        inputs_embeds=None,
-        spec_metadata=None,
-        **kwargs,
-    ):
-        calls.append((self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata))
-        calls.append(kwargs)
-        return "hidden states"
-
-    return forward
-
-
-def _run_step_both_ways(core, md, calls, receiver):
-    # Called the way the speculative causal LM calls its text model: keywords throughout.
-    out = core(
-        input_ids="ids",
-        attn_metadata=md,
-        position_ids="positions",
-        inputs_embeds=None,
-        spec_metadata="spec",
-        lora_params="lora",
-    )
-    assert out == "hidden states"
-    assert calls == [(receiver, md, "ids", "positions", None, "spec"), {"lora_params": "lora"}]
-
-    # And positionally, in the built-in forward's parameter order.
-    calls.clear()
-    assert core(md, "ids", "positions", None, "spec") == "hidden states"
-    assert calls == [(receiver, md, "ids", "positions", None, "spec"), {}]
-
-
-@pytest.mark.parametrize("num_contexts", [1, 4], ids=["prefill", "mixed"])
-def test_each_step_runs_the_builtin_forward_through_its_target(monkeypatch, num_contexts):
-    calls = []
-    monkeypatch.setattr(DeepseekV3Model, "forward", _recorder(calls))
-    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
-    md = _Metadata(num_contexts)
-
-    assert type(core._select_target(md)) is modeling.PrefillTarget
-    _run_step_both_ways(core, md, calls, core)
-
-
-def test_a_decode_step_runs_the_decode_targets_loop(monkeypatch):
-    calls, builtin = [], []
-    monkeypatch.setattr(DeepseekV3Model, "forward", _recorder(builtin))
-    monkeypatch.setattr(modeling.DecodeTarget, "forward", _recorder(calls))
-    core = modeling.GlmMoeDsaTp8ModelingV2.adopt(_builtin_text_model())
-    md = _Metadata(0)
-
-    assert type(core._select_target(md)) is modeling.DecodeTarget
-    _run_step_both_ways(core, md, calls, core._decode)
-    assert builtin == []
-
-
-def test_the_forwards_take_the_builtin_parameters():
-    builtin = inspect.signature(DeepseekV3Model.forward)
-    for forward in (
-        modeling.GlmMoeDsaTp8ModelingV2.forward,
-        modeling.PrefillTarget.forward,
-        modeling.DecodeTarget.forward,
-    ):
-        assert list(inspect.signature(forward).parameters) == list(builtin.parameters), forward
-
-
-def test_the_shell_is_the_builtin_causal_lm():
-    shell = modeling.ModelingV2Glm52Nvfp4Sm103Tp8
-    assert issubclass(shell, GlmMoeDsaForCausalLM)
-    # The model loader inspects load_weights' parameters to decide what it passes.
-    assert (
-        inspect.getfullargspec(shell.load_weights).args
-        == inspect.getfullargspec(GlmMoeDsaForCausalLM.load_weights).args
-    )
-
-
-def _hooks_of_plan(disable: str = "") -> RuntimeHooks:
-    """``make_runtime_hooks`` of a model whose plan disables ``disable``."""
-    env = {dispatch.DISABLE_ENV: disable} if disable else {}
-    model = types.SimpleNamespace(
-        plan=dispatch.resolve(_router_index.ModelingV2Mode.AUTO, env),
-        model=types.SimpleNamespace(decode_prologue=None),
-        model_config=types.SimpleNamespace(
-            mapping=types.SimpleNamespace(enable_attention_dp=False)
-        ),
-    )
-    return runtime.make_runtime_hooks(model)
-
-
-class _StoreWriter:
-    """Stands in for ``SamplerStoreWriter``, whose constructor allocates pinned host memory."""
-
-
-@pytest.fixture
-def host_only_store_writer(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
-        deferred_input_prep,
-    )
-
-    monkeypatch.setattr(deferred_input_prep, "SamplerStoreWriter", _StoreWriter)
-
-
-@pytest.mark.usefixtures("host_only_store_writer")
-def test_the_target_declares_the_step_input_stager_of_its_plan(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        PrepStager,
-    )
-
-    validate_env = _router_index.MODELING_V2_VALIDATE_ENV
-    monkeypatch.delenv(validate_env, raising=False)
-    hooks = _hooks_of_plan()
-    assert isinstance(hooks.step_input_stager, PrepStager)
-    assert hooks.step_input_stager.debug is None
-
-    monkeypatch.setenv(validate_env, "1")
-    assert _hooks_of_plan().step_input_stager.debug.level == 1
-    monkeypatch.delenv(validate_env)
-
-    assert _hooks_of_plan("deferred_input_prep").step_input_stager is None
-
-
-@pytest.mark.usefixtures("host_only_store_writer")
-def test_the_target_declares_the_decode_prologue_of_its_plan():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E501
-        DecodePrologueHook,
-    )
-
-    hooks = _hooks_of_plan()
-    assert isinstance(hooks.decode_prologue, DecodePrologueHook)
-    assert _hooks_of_plan("decode_prologue").decode_prologue is None
-
-
-@pytest.mark.usefixtures("host_only_store_writer")
-def test_the_target_declares_the_sampler_store_writer_of_its_plan():
-    assert isinstance(_hooks_of_plan().sampler_store_writer, _StoreWriter)
-    hooks = _hooks_of_plan("sampler_store")
-    assert hooks.sampler_store_writer is None
-    assert hooks.step_input_stager is not None
-
-
-@pytest.mark.usefixtures("host_only_store_writer")
-@pytest.mark.parametrize(
-    "disabled, off",
-    [
-        ("", ()),
-        ("deferred_input_prep", ("step_input_stager",)),
-        ("sampler_store", ("sampler_store_writer",)),
-        ("execution_stream_loop", ("loop_on_execution_stream",)),
-        (
-            "deferred_input_prep,sampler_store,execution_stream_loop",
-            ("step_input_stager", "sampler_store_writer", "loop_on_execution_stream"),
-        ),
-    ],
-)
-def test_each_part_of_deferred_input_prep_is_a_hook_of_its_own(disabled, off):
-    """The three parts are separate fast paths: disabling one leaves the other two declared."""
-    hooks = _hooks_of_plan(disabled)
-    for field in ("step_input_stager", "sampler_store_writer", "loop_on_execution_stream"):
-        assert bool(getattr(hooks, field)) is (field not in off), field
-    assert isinstance(hooks.executor_events, runtime.GlmExecutorEvents)
-
-
-@pytest.fixture
-def event_calls(monkeypatch):
-    """The kernel entry points the executor events call, recorded in order."""
-    calls = []
-    kernels = runtime.__name__.rsplit(".", 1)[0] + ".kernels"
-    monkeypatch.setitem(
-        sys.modules,
-        kernels + ".post_moe_hop",
-        types.SimpleNamespace(
-            check_health=lambda where, **kwargs: calls.append(("check_health", where, kwargs))
-        ),
-    )
-    monkeypatch.setitem(
-        sys.modules,
-        kernels + ".host_page_retention",
-        types.SimpleNamespace(
-            install_host_policy=lambda where: calls.append(("install_host_policy", where))
-        ),
-    )
-    compat = types.SimpleNamespace(
-        declare_overlap_loop=lambda adp: calls.append(("declare_overlap_loop", adp)),
-        loop_domain=lambda: calls.append(("loop_domain",)),
-    )
-    monkeypatch.setitem(
-        sys.modules,
-        kernels + ".deferred_input_prep",
-        types.SimpleNamespace(compat=compat),
-    )
-    return calls
-
-
-def _events(paths, enable_attention_dp=False):
-    plan = types.SimpleNamespace(enabled=lambda name: name in paths)
-    mapping = types.SimpleNamespace(enable_attention_dp=enable_attention_dp)
-    return runtime.GlmExecutorEvents(plan, mapping)
-
-
-def test_the_executor_events_act_for_the_fast_paths_of_the_plan(event_calls):
-    events = _events({"post_moe_hop", "host_page_retention", "sampler_store"})
-
-    events.after_warmup()
-    events.before_overlap_loop()
-    events.at_shutdown()
-    assert event_calls == [
-        ("check_health", "after_warmup", {"collective": True}),
-        ("install_host_policy", "after_warmup"),
-        ("declare_overlap_loop", False),
-        ("loop_domain",),
-        ("check_health", "end_of_serving", {"raise_on_error": False}),
-    ]
-
-
-@pytest.mark.parametrize("path", ["sampler_store", "execution_stream_loop"])
-def test_the_overlap_loop_declares_its_domain_for_either_part(event_calls, path):
-    _events({path}, enable_attention_dp=True).before_overlap_loop()
-    assert event_calls == [("declare_overlap_loop", True), ("loop_domain",)]
-
-
-def test_the_executor_events_do_nothing_without_their_fast_paths(event_calls):
-    events = _events(set())
-
-    events.after_warmup()
-    events.before_overlap_loop()
-    events.at_shutdown()
-    assert event_calls == []
-
-
-def test_the_stager_attaches_the_attribute_every_engine_call_site_reads():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        PrepStager,
-    )
-    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
-    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
-    from tensorrt_llm._torch.speculative.interface import SpecMetadata
-
-    for cls in (DecoderRunner, TrtllmAttentionMetadata, SpecMetadata, KVCacheManagerV2):
-        assert PrepStager._ATTACH_ATTR in vars(cls), cls
-        assert getattr(cls, PrepStager._ATTACH_ATTR) is None, cls
-
-
-class _ModelConfig:
-    """What the shell and its hooks read off the model config."""
-
-    class mapping:
-        world_size = 8
-        enable_attention_dp = False
-
-
-class _SparseMlaDecodeFmha:
-    """Stands in for the FMHA library the ``sparse_mla_decode`` kernel package registers."""
-
-
-SPARSE_MLA_LIB = "glm_moe_dsa_sparse_mla_decode"
-RANK_UNIFORM_ENV = dispatch.SHARED_ENV + (_router_index.MODELING_V2_VALIDATE_ENV,)
-
-
-@pytest.fixture
-def shell_construction(monkeypatch):
-    """The shell class with the built-in constructor stubbed, and the events of a construction.
-
-    The built-in constructor records the installed fast paths and the FMHA library list it
-    would build its attention layers from; ``adopt_model`` records the plan it attaches. The
-    selector state and the FMHA registry module are restored afterwards. The store writer the
-    hooks build is a host-only stand-in.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
-        deferred_input_prep,
-    )
-
-    monkeypatch.setattr(deferred_input_prep, "SamplerStoreWriter", _StoreWriter)
-    events = []
-
-    def adopt_model(model, plan):
-        events.append(("adopt", type(model.model).__name__, plan))
-
-    def builtin_init(self, model_config):
-        events.append("built-in model")
-        events.append(("installed", dispatch.installed()))
-        events.append(("fmha", list(fmha_registry.FMHA_LIBS)))
-        nn.Module.__init__(self)
-        self.model_config = model_config
-        self.model = _builtin_text_model()
-        self.model.layers = nn.ModuleList([_Layer(), _Layer(post_moe_hop=_Hop())])
-
-    def allgather(value, world_size):
-        events.append(("allgather", value, world_size))
-        return [value] * world_size
-
-    shell_cls = modeling.ModelingV2Glm52Nvfp4Sm103Tp8
-    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__init__", builtin_init)
-    monkeypatch.setattr(modeling, "adopt_model", adopt_model)
-    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__post_init__", lambda self: None)
-    monkeypatch.setattr(GlmMoeDsaForCausalLM, "__pp_init__", lambda self: None)
-    monkeypatch.setitem(
-        _router_index._RESOLVED_MODES, shell_cls, _router_index.ModelingV2Mode.REQUIRE
-    )
-    monkeypatch.setattr(dispatch, "_require", [False])
-    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
-    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: events.append(message))
-    monkeypatch.setattr(
-        dispatch, "_REGISTRY", {"sparse_mla_decode": {"fmha_lib": _SparseMlaDecodeFmha}}
-    )
-    for name in list(os.environ):
-        if name.startswith(dispatch.KNOB_PREFIX) or name in RANK_UNIFORM_ENV:
-            monkeypatch.delenv(name)
-    registry_namespace = dict(vars(fmha_registry))
-    yield shell_cls, events
-    vars(fmha_registry).update(registry_namespace)
-
-
-def test_the_shell_checks_the_ranks_resolves_the_plan_and_adopts_the_core(
-    monkeypatch, shell_construction
-):
-    """The knobs are checked across the ranks before any other construction step, and the
-    attached fast-path state after the model is built."""
-    shell_cls, events = shell_construction
-    monkeypatch.setenv(dispatch.DISABLE_ENV, "draft_tail")
-
-    shell = shell_cls(_ModelConfig())
-
-    allgather, record, built_in, installed, fmha, adopt, adopted = events
-    assert allgather == (
-        "allgather",
-        ((dispatch.DISABLE_ENV, "draft_tail"), (dispatch.MODE_KEY, "require")),
-        8,
-    )
-    # The layers are adopted after the core and before their state is checked.
-    assert adopt == ("adopt", "GlmMoeDsaTp8ModelingV2", shell.plan)
-    assert adopted == ("allgather", modeling.adopted_state(shell), 8)
-    assert record.startswith("[glm_moe_dsa] dispatch: record paths=")
-    assert built_in == "built-in model"
-    assert installed == ("installed", frozenset(dispatch.FAST_PATH_NAMES) - {"draft_tail"})
-    assert fmha[0] == "fmha"
-    assert type(shell.model) is modeling.GlmMoeDsaTp8ModelingV2
-    assert shell.plan.require and dispatch.require_mode()
-    assert not shell.plan.enabled("draft_tail") and shell.plan.enabled("routed_moe")
-    # The hooks are built on the engine's first read, not during construction.
-    assert "_runtime_hooks" not in vars(shell)
-    hooks = runtime_hooks_of(shell)
-    assert hooks.step_input_stager is not None
-    assert isinstance(hooks.sampler_store_writer, _StoreWriter)
-    assert hooks.loop_on_execution_stream
-    assert isinstance(hooks.executor_events, runtime.GlmExecutorEvents)
-    assert runtime_hooks_of(shell) is hooks
-
-
-def test_an_attribute_error_while_building_the_hooks_is_not_read_as_no_hooks(
-    monkeypatch, shell_construction
-):
-    shell_cls, _ = shell_construction
-    shell = shell_cls(_ModelConfig())
-
-    def broken(model):
-        raise AttributeError("plan")
-
-    monkeypatch.setattr(runtime, "make_runtime_hooks", broken)
-    with pytest.raises(RuntimeError, match="step-loop hooks"):
-        runtime_hooks_of(shell)
-
-
-def test_the_fast_paths_read_the_installed_plan(monkeypatch, shell_construction):
-    """The plan is installed before the built-in model is built; DISABLE turns paths off."""
-    shell_cls, events = shell_construction
-    monkeypatch.setenv(dispatch.DISABLE_ENV, "post_attention_hop")
-
-    shell = shell_cls(_ModelConfig())
-
-    off = {name for name in dispatch.FAST_PATH_NAMES if not dispatch.is_enabled(name)}
-    assert off == shell.plan.disabled | shell.plan.disabled_dependents
-    assert "post_attention_hop" in off and "hop_routing" in off
-    assert ("installed", frozenset(dispatch.FAST_PATH_NAMES) - off) in events
-    # A knob the target does not read, set after construction, changes nothing.
-    monkeypatch.setenv(dispatch.FEATURE_ENV, "post_attention_hop")
-    assert not dispatch.is_enabled("post_attention_hop")
-    record = next(e for e in events if isinstance(e, str) and "dispatch: record" in e)
-    assert "post_attention_hop" not in record.split(" side_selectors=")[0]
-
-
-def test_the_fmha_registry_stays_the_builtin_one(monkeypatch, shell_construction):
-    """sparse_mla_decode joins no process-wide FMHA library list: adoption declares it on the
-    GLM attention backends (``test_modeling_v2_glm_attention.py``)."""
-    shell_cls, events = shell_construction
-    builtin_libs = list(fmha_registry.FMHA_LIBS)
-    builtin_enabled = fmha_manager.get_enabled_fmha_lib_classes()
-
-    shell_cls(_ModelConfig())
-    monkeypatch.setenv(dispatch.DISABLE_ENV, "sparse_mla_decode")
-    shell_cls(_ModelConfig())
-
-    assert SPARSE_MLA_LIB not in builtin_libs
-    assert [e for e in events if isinstance(e, tuple) and e[0] == "fmha"] == [
-        ("fmha", builtin_libs)
-    ] * 2
-    assert list(fmha_registry.FMHA_LIBS) == builtin_libs
-    assert fmha_manager.get_enabled_fmha_lib_classes() == builtin_enabled
-
-
-@pytest.mark.parametrize(
-    "knob",
-    sorted(dispatch.REMOVED_KNOBS),
-    ids=lambda knob: knob.removeprefix(dispatch.KNOB_PREFIX),
-)
-def test_a_removed_knob_fails_before_the_builtin_model_is_built(
-    monkeypatch, shell_construction, knob
-):
-    shell_cls, events = shell_construction
-    monkeypatch.setenv(knob, "deferred_input_prep")
-
-    with pytest.raises(ValueError, match=f"this target does not read {knob} "):
-        shell_cls(_ModelConfig())
-    assert [event[0] for event in events] == ["allgather"]
-    assert dispatch.installed() is None
-
-
-def test_a_knob_value_no_rank_can_parse_fails_before_the_builtin_model_is_built(
-    monkeypatch, shell_construction
-):
-    shell_cls, events = shell_construction
-    monkeypatch.setenv(dispatch.DISABLE_ENV, "not_a_fast_path")
-
-    with pytest.raises(ValueError, match="not_a_fast_path"):
-        shell_cls(_ModelConfig())
-    assert [event[0] for event in events] == ["allgather"]
-
-
-def test_the_ranks_check_the_shared_switches_and_the_mode_too(monkeypatch, shell_construction):
-    shell_cls, events = shell_construction
-    monkeypatch.setenv("TRTLLM_ENABLE_PDL", "1")
-    monkeypatch.setenv(_router_index.MODELING_V2_VALIDATE_ENV, "1")
-
-    shell_cls(_ModelConfig())
-
-    assert events[0] == (
-        "allgather",
-        (
-            ("TRTLLM_ENABLE_PDL", "1"),
-            (_router_index.MODELING_V2_VALIDATE_ENV, "1"),
-            (dispatch.MODE_KEY, "require"),
-        ),
-        8,
-    )
-
-
-def test_ranks_that_disagree_fail_before_the_builtin_model_is_built(
-    monkeypatch, shell_construction
-):
-    shell_cls, events = shell_construction
-
-    def allgather(value, world_size):
-        events.append(("allgather", value, world_size))
-        return [value, value + (("TRTLLM_FORCE_MNNVL_AR", "1"),)]
-
-    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
-    with pytest.raises(RuntimeError, match="TRTLLM_FORCE_MNNVL_AR='1'"):
-        shell_cls(_ModelConfig())
-    assert [event[0] for event in events] == ["allgather"]
-    assert dispatch.installed() is None
-
-
-def test_the_adopted_state_names_each_attached_fast_path_object():
-    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[_Layer(), _Layer()]))
-    model.model.layers[1].post_moe_hop = _Hop()
-    model.model.layers[1].mlp.router_gemm = _Hop(engaged=False)
-    model.model.layers[1].self_attn._attn_tail = object()
-
-    assert modeling.adopted_state(model) == (
-        ("model.layers.0", ()),
-        (
-            "model.layers.1",
-            (
-                ("post_moe_hop", "_Hop", "engaged=True", "impl='cluster'"),
-                ("self_attn._attn_tail", "object"),
-                ("mlp.router_gemm", "_Hop", "engaged=False", "impl='cluster'"),
-            ),
-        ),
-    )
-
-
-def test_ranks_that_attached_different_state_fail_after_the_model_is_built(
-    monkeypatch, shell_construction
-):
-    shell_cls, events = shell_construction
-    hooks = []
-    monkeypatch.setattr(runtime, "make_runtime_hooks", lambda model: hooks.append(model))
-
-    def allgather(value, world_size):
-        events.append(("allgather", value, world_size))
-        if value and value[0][0] == "model.layers.0":
-            # A peer whose post-MoE hop on layer 1 did not engage.
-            return [value, (value[0], ("model.layers.1", ()))]
-        return [value] * world_size
-
-    monkeypatch.setattr(dispatch, "_world_allgather", allgather)
-    with pytest.raises(RuntimeError, match=r"rank 1: \('model.layers.1', \(\)\)") as excinfo:
-        shell_cls(_ModelConfig())
-    assert "rank 0: ('model.layers.1', (('post_moe_hop', '_Hop'" in str(excinfo.value)
-    assert "built-in model" in events and hooks == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py
deleted file mode 100644
index 5642cf41d4bb..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_modeling_v2_glm_vocabulary.py
+++ /dev/null
@@ -1,198 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""One name per fast path of the GLM-5.2 tp8 target, used the same way everywhere.
-
-The fast-path names of ``dispatch.FAST_PATHS`` are the ``DISABLE`` tokens, the kernel package
-directories (or a directory declared in ``KERNEL_DIRECTORY``), the log tags, the knob infixes,
-the kernel registry keys and the kernel test directories. Read off disk: no kernel module is
-imported.
-"""
-
-from __future__ import annotations
-
-import ast
-import re
-from pathlib import Path
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-# The family's code: its shared package and every target of it, each with its kernel packages.
-_TARGET = Path(dispatch.__file__).resolve().parents[1]
-_KERNEL_ROOTS = sorted(_TARGET.glob("*/kernels"))
-_TESTS = Path(__file__).resolve().parents[1]
-_PATHS = frozenset(dispatch.FAST_PATH_NAMES)
-
-_LOG_TAG = "[glm_moe_dsa] "
-# After the tag: a fast-path name and a separator, or the dispatcher's own lines.
-_TAGGED = re.compile(r"(?P<name>[a-z0-9_]+)(?P<sep>[/ :])")
-_KNOB = re.compile(re.escape(dispatch.KNOB_PREFIX) + r"(?P<infix>[A-Z0-9_]*)")
-# Knobs that select among the fast paths rather than tune one.
-_SELECTOR_KNOBS = ("FEATURES", "DISABLE")
-
-
-def _sources(root: Path):
-    return sorted(p for p in root.rglob("*.py") if "__pycache__" not in p.parts)
-
-
-def _kernel_directories() -> set:
-    """The kernel packages of the family, the shared ones and every target's own."""
-    return {
-        d.name
-        for kernels in _KERNEL_ROOTS
-        for d in kernels.iterdir()
-        if d.is_dir() and (d / "__init__.py").is_file()
-    }
-
-
-def _string_constants(tree: ast.AST):
-    """Every ``str`` constant, the constant parts of f-strings included."""
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Constant) and isinstance(node.value, str):
-            yield node
-
-
-def test_each_fast_path_has_one_name():
-    assert len(dispatch.FAST_PATH_NAMES) == 36
-    assert len(_PATHS) == 36
-
-
-def test_the_kernel_directories_are_the_fast_paths_without_a_declared_directory():
-    directories = _kernel_directories()
-    assert len(directories) == 28
-    assert directories == _PATHS - set(dispatch.KERNEL_DIRECTORY)
-    assert set(dispatch.KERNEL_DIRECTORY) <= _PATHS
-    assert set(dispatch.KERNEL_DIRECTORY.values()) <= directories
-
-
-def test_every_log_tag_names_a_fast_path_or_the_dispatcher():
-    offences = []
-    for path in _sources(_TARGET):
-        for node in _string_constants(ast.parse(path.read_text(encoding="utf-8"))):
-            text = node.value
-            start = text.find(_LOG_TAG)
-            while start >= 0:
-                match = _TAGGED.match(text, start + len(_LOG_TAG))
-                name = match.group("name") if match else None
-                if not (
-                    match and (name in _PATHS or (name == "dispatch" and match.group("sep") == ":"))
-                ):
-                    offences.append(f"{path.relative_to(_TARGET)}:{node.lineno}: {text[start:]!r}")
-                start = text.find(_LOG_TAG, start + 1)
-    assert not offences, "log lines must name a fast path after the tag:\n" + "\n".join(offences)
-
-
-def _knob_owner(infix: str):
-    """The fast path whose name starts the knob's infix (the longest such name), if any."""
-    owners = [p for p in _PATHS if infix.startswith(p.upper() + "_")]
-    return max(owners, key=len) if owners else None
-
-
-def test_every_knob_name_starts_with_its_fast_path():
-    offences = []
-    for path in _sources(_TARGET):
-        text = path.read_text(encoding="utf-8")
-        for match in _KNOB.finditer(text):
-            infix = match.group("infix")
-            if infix and infix not in _SELECTOR_KNOBS and _knob_owner(infix) is None:
-                line = text.count("\n", 0, match.start()) + 1
-                offences.append(f"{path.relative_to(_TARGET)}:{line}: {match.group(0)}")
-    assert not offences, "knob names must start with a fast-path name:\n" + "\n".join(offences)
-
-
-def _imports_register_from_dispatch(tree: ast.AST) -> bool:
-    return any(
-        isinstance(node, ast.ImportFrom)
-        and (node.module or "").split(".")[-1] == "dispatch"
-        and any(alias.name == "register" and alias.asname is None for alias in node.names)
-        for node in ast.walk(tree)
-    )
-
-
-def _module_strings(path: Path, tree: ast.Module) -> dict:
-    """Module-level ``NAME = "<str>"`` bindings, and those imported from a sibling module."""
-    out = {}
-    for stmt in tree.body:
-        if (
-            isinstance(stmt, ast.Assign)
-            and isinstance(stmt.value, ast.Constant)
-            and isinstance(stmt.value.value, str)
-        ):
-            out.update({t.id: stmt.value.value for t in stmt.targets if isinstance(t, ast.Name)})
-        elif isinstance(stmt, ast.ImportFrom) and stmt.level == 1 and stmt.module:
-            sibling = path.parent / f"{stmt.module}.py"
-            if sibling.is_file():
-                strings = _module_strings(sibling, ast.parse(sibling.read_text(encoding="utf-8")))
-                out.update(
-                    {
-                        alias.asname or alias.name: strings[alias.name]
-                        for alias in stmt.names
-                        if alias.name in strings
-                    }
-                )
-    return out
-
-
-def _registered_paths(path: Path):
-    """``(line, name)`` of every ``dispatch.register(<name>, …)`` call of one kernel module."""
-    tree = ast.parse(path.read_text(encoding="utf-8"))
-    bare = _imports_register_from_dispatch(tree)
-    names = _module_strings(path, tree)
-    for node in ast.walk(tree):
-        if not isinstance(node, ast.Call) or not node.args:
-            continue
-        func = node.func
-        if not (
-            (isinstance(func, ast.Name) and func.id == "register" and bare)
-            or (
-                isinstance(func, ast.Attribute)
-                and func.attr == "register"
-                and isinstance(func.value, ast.Name)
-                and func.value.id == "dispatch"
-            )
-        ):
-            continue
-        first = node.args[0]
-        if isinstance(first, ast.Constant) and isinstance(first.value, str):
-            yield node.lineno, first.value
-        elif isinstance(first, ast.Name) and first.id in names:
-            yield node.lineno, names[first.id]
-        else:
-            yield node.lineno, ast.unparse(first)
-
-
-def test_every_registry_entry_is_keyed_by_the_fast_path_of_its_directory():
-    offences, registering = [], set()
-    for kernels in _KERNEL_ROOTS:
-        for path in _sources(kernels):
-            relative = path.relative_to(kernels)
-            if len(relative.parts) < 2:
-                continue
-            directory = relative.parts[0]
-            for line, name in _registered_paths(path):
-                if name not in _PATHS or dispatch.KERNEL_DIRECTORY.get(name, name) != directory:
-                    offences.append(
-                        f"{kernels.parent.name}/kernels/{relative}:{line}: register({name!r}, ...)"
-                    )
-                else:
-                    registering.add(directory)
-    assert not offences, (
-        "get_impl imports kernels.<directory of the path>, so each package registers its "
-        "entry points under its own path:\n" + "\n".join(offences)
-    )
-    assert registering == _kernel_directories()
-
-
-def test_every_kernel_test_directory_is_a_kernel_directory():
-    """A target's tests of how it composes a shared package live under the package's name."""
-    for kernels in _KERNEL_ROOTS:
-        test_root = _TESTS / kernels.parent.name / "kernels"
-        if not test_root.is_dir():
-            continue
-        tests = {d.name for d in test_root.iterdir() if d.is_dir() and d.name != "__pycache__"}
-        assert tests <= _kernel_directories(), sorted(tests - _kernel_directories())
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py
deleted file mode 100644
index c63a5c287a14..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/absorbed_query_fixture.py
+++ /dev/null
@@ -1,223 +0,0 @@
-# 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.
-"""Real-weight native fixtures for the absorbed_query/projection package.
-
-``build_native_attention`` constructs the built-in ``DeepseekV32Attention``
-module that the GLM-5.2 decode engine uses for one target layer at TP8 rank
-``tp_rank``, adopts it as the target's ``GlmMLA`` right after construction (as the
-target's model does; on a decode step it runs the selected fast paths, and with none
-selected or on any other step it computes the built-in attention) and places the
-checkpoint tensors exactly as ``DeepseekV3WeightLoader`` does for the BF16 attention
-weights:
-
-* ``kv_a_proj_with_mqa.weight = cat([q_a_proj, kv_a_proj_with_mqa])``
-* ``q_b_proj`` through ``Linear.load_weights`` (column-parallel slicing)
-* ``k_b_proj_trans`` / ``v_b_proj`` / ``kv_b_proj`` from the TP slice of
-  ``kv_b_proj`` (heads ``[r*8, r*8+8)``)
-* indexer ``wq_b``/``wk``/``weights_proj``/``k_norm`` and
-  ``Indexer.post_load_weights`` (fused ``wk``+``weights_proj``).
-
-The attention linears of this checkpoint are excluded from NVFP4, so the
-module is built with a KV-cache-only ``QuantConfig`` (FP8 KV), which is the
-configuration ``apply_quant_config_exclude_modules`` leaves on them.
-
-Fixture labels: weights are **real** (checkpoint tensors); activations are
-**synthetic** unless a caller supplies captured ones.
-"""
-
-import json
-import os
-from typing import Dict, List, Optional
-
-import torch
-from utils.llm_data import llm_models_root
-
-CKPT_DIR = os.environ.get(
-    "GLM52_CKPT_DIR",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-TP_SIZE = 8
-
-# Full (own indexer) and shared (reused top-k) target layers of GLM-5.2.
-FULL_LAYER = 6
-SHARED_LAYER = 7
-
-
-# Feature names of the neighbour packages every arm of a comparison runs on
-# top of, the native reference included (e.g. "rope_kv_append" on an
-# integration image that contains rope_kv_append).
-BASE_FEATURES_ENV = "GLM_MOE_DSA_TEST_BASE_FEATURES"
-
-
-def with_base_features(features: str) -> str:
-    """Return the fast paths ``features`` as a comma-separated list.
-
-    The fast paths of ``GLM_MOE_DSA_TEST_BASE_FEATURES`` follow them.
-    """
-    names = [f.strip() for f in features.split(",") if f.strip()]
-    base = [f.strip() for f in os.environ.get(BASE_FEATURES_ENV, "").split(",") if f.strip()]
-    return ",".join(names + [b for b in base if b not in names])
-
-
-def checkpoint_available() -> bool:
-    return os.path.isfile(os.path.join(CKPT_DIR, "model.safetensors.index.json"))
-
-
-def load_layer_tensors(
-    layer_idx: int, names: Optional[List[str]] = None
-) -> Dict[str, torch.Tensor]:
-    """Load ``model.layers.<layer_idx>.self_attn.*`` tensors (CPU)."""
-    from safetensors import safe_open
-
-    with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
-        weight_map = json.load(f)["weight_map"]
-    prefix = f"model.layers.{layer_idx}.self_attn."
-    keys = [k for k in weight_map if k.startswith(prefix)]
-    if names is not None:
-        keys = [prefix + n for n in names if prefix + n in weight_map]
-    by_shard: Dict[str, List[str]] = {}
-    for k in keys:
-        by_shard.setdefault(weight_map[k], []).append(k)
-    out = {}
-    for shard, shard_keys in by_shard.items():
-        with safe_open(os.path.join(CKPT_DIR, shard), framework="pt") as f:
-            for k in shard_keys:
-                out[k[len(prefix) :]] = f.get_tensor(k)
-    return out
-
-
-def build_model_config(tp_rank: int = 0):
-    """Native ``ModelConfig`` for GLM-5.2 at TP8 with the served DSA flags."""
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, MTPDecodingConfig
-    from tensorrt_llm.mapping import Mapping
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-    from tensorrt_llm.quantization.mode import QuantAlgo
-
-    mapping = Mapping(world_size=TP_SIZE, tp_size=TP_SIZE, rank=tp_rank, gpus_per_node=4)
-    sparse = DeepSeekSparseAttentionConfig(
-        enable_heuristic_topk=True, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=True
-    )
-    model_config = ModelConfig.from_pretrained(
-        CKPT_DIR,
-        trust_remote_code=True,
-        mapping=mapping,
-        sparse_attention_config=sparse,
-        spec_config=MTPDecodingConfig(max_draft_len=5),
-        attn_backend="TRTLLM",
-    )
-    model_config._frozen = False
-    # Attention linears are excluded from NVFP4 in this checkpoint; the
-    # KV cache stays FP8 as in the served configuration.
-    model_config.quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8)
-    if model_config.pretrained_config.model_type == "glm_moe_dsa":
-        model_config.pretrained_config.model_type = "deepseek_v32"
-    model_config._frozen = True
-    return model_config
-
-
-def build_native_attention(
-    layer_idx: int,
-    tp_rank: int = 0,
-    model_config=None,
-    tensors: Optional[Dict[str, torch.Tensor]] = None,
-    aux_stream: Optional[torch.cuda.Stream] = None,
-):
-    """Build the built-in ``DeepseekV32Attention``, adopted as ``GlmMLA``, with real weights on GPU."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-        GlmMLA,
-    )
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
-    from tensorrt_llm._torch.utils import AuxStreamType
-
-    if model_config is None:
-        model_config = build_model_config(tp_rank)
-    if tensors is None:
-        tensors = load_layer_tensors(layer_idx)
-    if aux_stream is None:
-        aux_stream = torch.cuda.Stream()
-    attn = DeepseekV32Attention(
-        model_config, layer_idx=layer_idx, aux_stream_dict={AuxStreamType.Attention: aux_stream}
-    )
-    GlmMLA.adopt(attn)
-    # The model builds the FMHA libraries in ``create_weights`` after adoption; a module built
-    # alone rebuilds them for the library adoption declared.
-    for backend in (attn.mha, attn.mqa):
-        if backend is not None and backend.prepended_fmha_lib_classes:
-            backend.update_quant_config(backend.quant_config)
-    attn = attn.cuda()
-
-    cfg = model_config.pretrained_config
-    num_heads = cfg.num_attention_heads
-    nope, vdim = cfg.qk_nope_head_dim, cfg.v_head_dim
-    local_heads = num_heads // TP_SIZE
-
-    with torch.no_grad():
-        fused_a = torch.cat(
-            [tensors["q_a_proj.weight"], tensors["kv_a_proj_with_mqa.weight"]], dim=0
-        )
-        attn.kv_a_proj_with_mqa.weight.data.copy_(fused_a)
-        attn.q_a_layernorm.weight.data.copy_(tensors["q_a_layernorm.weight"])
-        attn.kv_a_layernorm.weight.data.copy_(tensors["kv_a_layernorm.weight"])
-        attn.q_b_proj.load_weights(weights=[{"weight": tensors["q_b_proj.weight"]}])
-
-        kv_b = tensors["kv_b_proj.weight"].unflatten(0, [num_heads, nope + vdim])
-        kv_b = torch.chunk(kv_b, TP_SIZE, dim=0)[tp_rank].contiguous()
-        k_nope, v_w = kv_b.split([nope, vdim], dim=1)
-        attn.k_b_proj_trans.data.copy_(k_nope.transpose(2, 1).contiguous())
-        attn.kv_b_proj.weight.data.copy_(
-            torch.cat(
-                [k_nope.reshape(local_heads * nope, -1), v_w.reshape(local_heads * vdim, -1)], dim=0
-            )
-        )
-        attn.v_b_proj = torch.nn.Parameter(v_w.contiguous().cuda(), requires_grad=False)
-
-        indexer = attn.mqa.indexer
-        if indexer is not None:
-            indexer.wq_b.load_weights(weights=[{"weight": tensors["indexer.wq_b.weight"]}])
-            indexer.wk.load_weights(weights=[{"weight": tensors["indexer.wk.weight"]}])
-            indexer.weights_proj.load_weights(
-                weights=[{"weight": tensors["indexer.weights_proj.weight"]}]
-            )
-            indexer.k_norm.weight.data.copy_(tensors["indexer.k_norm.weight"])
-            indexer.k_norm.bias.data.copy_(tensors["indexer.k_norm.bias"])
-            indexer.post_load_weights()
-    attn.post_load_weights()
-    return attn
-
-
-def make_hidden(
-    num_rows: int, hidden_size: int = 6144, seed: int = 0, gamma: Optional[torch.Tensor] = None
-) -> torch.Tensor:
-    """Synthetic post-RMSNorm hidden state ``[R, hidden]`` (BF16).
-
-    Unit-RMS Gaussian rows scaled by a real ``input_layernorm`` gain when one
-    is given, so activation magnitudes match what the layer sees.
-    """
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = torch.randn(num_rows, hidden_size, generator=g, device="cuda", dtype=torch.float32)
-    if gamma is not None:
-        x = x * gamma.float().cuda()
-    return x.to(torch.bfloat16)
-
-
-def load_input_layernorm(layer_idx: int) -> torch.Tensor:
-    from safetensors import safe_open
-
-    with open(os.path.join(CKPT_DIR, "model.safetensors.index.json")) as f:
-        weight_map = json.load(f)["weight_map"]
-    key = f"model.layers.{layer_idx}.input_layernorm.weight"
-    with safe_open(os.path.join(CKPT_DIR, weight_map[key]), framework="pt") as f:
-        return f.get_tensor(key)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py
deleted file mode 100644
index 24857f178128..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/conftest.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# 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.
-"""Make the absorbed_query fixture helpers importable under ``--import-mode=importlib``."""
-
-import sys
-from pathlib import Path
-
-_HERE = str(Path(__file__).resolve().parent)
-if _HERE not in sys.path:
-    sys.path.insert(0, _HERE)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py
deleted file mode 100644
index b65f038f957a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_attention_integration.py
+++ /dev/null
@@ -1,870 +0,0 @@
-# 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.
-"""Stock producer -> absorbed_query -> native DSA attention, through the target's attention.
-
-The attention modules are the target's adopted ones (``GlmMLA``, built by
-``absorbed_query_fixture.build_native_attention``); every generation step runs as a decode
-step (``dispatch.Phase.DECODE``) and the prefill as a prefill step. Real
-GLM-5.2 weights of the full-indexer layer 6 and the shared layer 7 at TP8
-rank 0; synthetic activations; FP8 KV cache (V1 ``DSACacheManager``).
-
-Six requests with a 2,600-token context are natively prefilled through both
-layers (the indexer then selects 2,048 of 2,600 keys). A generation step
-then presents six query rows in one of two layouts: six requests x one token,
-or one request x six tokens (the served MTP verify layout, K = 5).
-
-Both call paths are exercised: ``forward_impl`` (``forward_sparse_attn``)
-and ``custom_op`` (``MLA.forward``'s ``_forward_custom_op`` with the
-registered metadata, which the model uses). Every step, eager or captured,
-first calls ``on_update_kv_lens()`` on the metadata, as the decoder
-runner's ``_forward_step`` does through ``_preprocess_inputs``.
-
-Checks against the same step with absorbed_query off (all asserted unless noted):
-
-* only the new-token slots of every cache (latent caches of both layers,
-  indexer-key cache of layer 6) change, and every byte equals native;
-* the MLA query after native RoPE/quantization is bit-identical: the
-  absorbed half ``fused_q[..., :512]`` and the FP8 ``quant_q_buffer``
-  handed to the attention kernel, and the BMM scales written next to them
-  (recorded by wrapping the native ``mla_rope_generation``; in the FP8-KV
-  path the RoPE half of ``fused_q`` is not written, the rotated q-rope goes
-  only to ``quant_q_buffer``);
-* the indexer inputs (``q_fp8``, ``weights``) are bit-identical to native,
-  and the indexer's logical top-k (layer 6) and the converted sparse indices
-  of both layers select the same key sets (rows compared after sorting;
-  exact order and, on a failure, the per-row overlap reported);
-* the attention outputs: the native sparse attention is not bitwise
-  run-to-run deterministic, so the outputs are bounded by its measured
-  repeat difference over ``NATIVE_RUNS`` native runs of the same step:
-  ``max|absorbed_query - native| <= 2 * max(native repeat max diff, one BF16 ulp of
-  max|native|)`` and ``relRMS(absorbed_query - native) <= max(2 * native repeat relRMS,
-  2**-10)``. This is a secondary check: the negative control shows that a
-  zeroed absorbed query moves the shared layer's outputs by only a few
-  times that noise, while the bit-identity checks of the attention inputs
-  above flag it directly;
-* absorbed_query itself is deterministic: a repeated absorbed_query step, 20 more absorbed_query
-  steps and every CUDA-graph replay reproduce the caches and recorded attention inputs of
-  the eager absorbed_query step bit for bit.
-
-The graph test captures the adopted layer forward (absorbed_query, ``_forward_dsa_attn``
-and the native cache append) for CUDA-graph generation metadata, replays it
-once and 20 times with fresh hidden states, and compares each replay with an
-eager native step and with a captured native graph.
-"""
-
-import inspect
-import math
-import os
-import sys
-
-import absorbed_query_fixture
-import pytest
-import torch
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
-
-import plan_testlib  # noqa: E402
-
-CTX = 2600
-REQS = 6
-LAYERS = (absorbed_query_fixture.FULL_LAYER, absorbed_query_fixture.SHARED_LAYER)
-# layout name -> (requests, query tokens per request)
-LAYOUTS = {"6req_x1tok": (6, 1), "1req_x6tok": (1, 6)}
-REL_RMS_FLOOR = 2.0**-10
-
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-
-def _set_absorbed_query(enabled: bool, form: str = "preabsorbed"):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    query_proj.SERVED_FORM = form
-    plan_testlib.select(
-        dispatch, absorbed_query_fixture.with_base_features("absorbed_query" if enabled else "")
-    )
-
-
-class _Recorder:
-    """Records, per layer, what the native attention consumes.
-
-    Wraps the native ``mla_rope_generation`` (post-RoPE ``fused_q``, FP8
-    ``quant_q_buffer``, BMM scales), the indexer's ``sparse_attn_indexer``
-    (logical top-k) and ``sparse_attn_predict`` (converted indices). Inside
-    a CUDA-graph capture the clones are graph-pool tensors that every replay
-    refreshes. When the installed package registers rope_kv_append's fused replacement
-    of ``mla_rope_generation`` (integration images), that entry point is
-    wrapped as well and records the same tensors.
-    """
-
-    def __init__(self, attn_by_layer):
-        self.attn = attn_by_layer
-        self.data = {layer: {} for layer in attn_by_layer}
-        self._installed = []
-        self._rope_append = None
-
-    @staticmethod
-    def _rope_record(rec, kv_lora, a):
-        # The FP8-KV path writes the rotated q-rope only into
-        # quant_q_buffer; fused_q[..., kv_lora:] stays unwritten.
-        rec["absorbed_q"] = a["fused_q"][..., :kv_lora].clone()
-        for key, arg in (
-            ("quant_q", "quant_q_buffer"),
-            ("bmm1_scale", "mla_bmm1_scale"),
-            ("bmm2_scale", "mla_bmm2_scale"),
-        ):
-            if a.get(arg) is not None:
-                rec[key] = a[arg].clone()
-
-    def _wrap_rope_append_rope(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-        try:
-            impl = dispatch.get_impl("rope_kv_append", "mla_rope_generation")
-        except (ImportError, KeyError, ValueError):
-            return
-        sig = inspect.signature(impl)
-        by_mqa = {
-            id(attn.mqa): (self.data[layer], attn.kv_lora_rank) for layer, attn in self.attn.items()
-        }
-
-        def hook(*args, **kwargs):
-            out = impl(*args, **kwargs)
-            a = sig.bind(*args, **kwargs).arguments
-            hit = by_mqa.get(id(a["mqa"]))
-            if hit is not None:
-                self._rope_record(hit[0], hit[1], a)
-                # Marks the call as served by rope_kv_append (compared like any record).
-                hit[0]["rope_append_rope"] = torch.ones(
-                    1, dtype=torch.int32, device=a["fused_q"].device
-                )
-            return out
-
-        dispatch._REGISTRY["rope_kv_append"]["mla_rope_generation"] = hook
-        self._rope_append = (dispatch, impl)
-
-    def _wrap(self, obj, name, post):
-        orig = getattr(obj, name)
-        sig = inspect.signature(orig)
-
-        def hook(*args, **kwargs):
-            out = orig(*args, **kwargs)
-            post(sig.bind(*args, **kwargs).arguments, out)
-            return out
-
-        setattr(obj, name, hook)
-        self._installed.append((obj, name))
-
-    def __enter__(self):
-        for layer, attn in self.attn.items():
-            rec = self.data[layer]
-
-            def rope_post(a, _out, rec=rec, kv_lora=attn.kv_lora_rank):
-                self._rope_record(rec, kv_lora, a)
-
-            def predict_post(_a, out, rec=rec):
-                if isinstance(out, tuple) and out[0] is not None:
-                    rec["sparse_indices"] = out[0].clone()
-
-            self._wrap(attn.mqa, "mla_rope_generation", rope_post)
-            self._wrap(attn.mqa, "sparse_attn_predict", predict_post)
-            if attn.mqa.indexer is not None:
-
-                def topk_post(a, out, rec=rec):
-                    rec["topk"] = out.clone()
-                    rec["indexer_q_fp8"] = a["q_fp8"].clone()
-                    rec["indexer_weights"] = a["weights"].clone()
-
-                self._wrap(attn.mqa.indexer, "sparse_attn_indexer", topk_post)
-        self._wrap_rope_append_rope()
-        return self
-
-    def __exit__(self, *exc):
-        for obj, name in self._installed:
-            delattr(obj, name)
-        self._installed.clear()
-        if self._rope_append is not None:
-            dispatch, impl = self._rope_append
-            dispatch._REGISTRY["rope_kv_append"]["mla_rope_generation"] = impl
-            self._rope_append = None
-        return False
-
-    def snapshot(self):
-        return {layer: {k: v.clone() for k, v in rec.items()} for layer, rec in self.data.items()}
-
-
-class _Harness:
-    def __init__(self):
-        import tensorrt_llm.bindings
-        from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
-        from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
-        from tensorrt_llm._torch.metadata import KVCacheParams
-        from tensorrt_llm._torch.model_config import ModelConfig
-        from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
-        from tensorrt_llm.llmapi.llm_args import KvCacheConfig
-        from tensorrt_llm.mapping import Mapping
-        from tensorrt_llm.models.modeling_utils import QuantConfig
-        from tensorrt_llm.quantization.mode import QuantAlgo
-
-        self.KVCacheParams = KVCacheParams
-        self.model_config = absorbed_query_fixture.build_model_config(0)
-        pretrained = self.model_config.pretrained_config
-        sparse = self.model_config.sparse_attention_config
-        self.attn = {
-            layer: absorbed_query_fixture.build_native_attention(layer, 0, self.model_config)
-            for layer in LAYERS
-        }
-        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
-        kv_model_config = ModelConfig(
-            mapping=self.mapping,
-            sparse_attention_config=sparse,
-            pretrained_config=pretrained,
-            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8),
-        )
-        self.kvm = DSACacheManager(
-            KvCacheConfig(
-                dtype="fp8",
-                max_tokens=REQS * (CTX + 128),
-                enable_block_reuse=False,
-                event_buffer_max_size=0,
-            ),
-            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-            num_layers=max(LAYERS) + 1,
-            num_kv_heads=1,
-            head_dim=576,
-            tokens_per_block=64,
-            max_seq_len=CTX + 128,
-            max_batch_size=REQS,
-            mapping=self.mapping,
-            dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
-            sparse_attention_config=sparse,
-            model_config=kv_model_config,
-            pretrained_config=pretrained,
-            vocab_size=pretrained.vocab_size,
-            max_num_tokens=REQS * CTX,
-        )
-        self.req_ids = list(range(REQS))
-        for rid in self.req_ids:
-            # CTX + 6 new tokens fit the blocks of CTX + 1 tokens (64/block).
-            self.kvm.add_dummy_requests(
-                request_ids=[rid], token_nums=[CTX + 1], is_gen=False, prepare_resource=True
-            )
-        sparse_params = sparse.to_sparse_params(pretrained_config=pretrained)
-        self.Metadata = get_attention_backend("TRTLLM", sparse_params=sparse_params).Metadata
-        self.meta_params = sparse.to_sparse_metadata_params(pretrained_config=pretrained)
-        self.gamma = absorbed_query_fixture.load_input_layernorm(LAYERS[0])
-
-        # Native prefill through both layers fills the latent and indexer
-        # caches the way the stock producers do.
-        _set_absorbed_query(False)
-        hidden = absorbed_query_fixture.make_hidden(REQS * CTX, seed=11, gamma=self.gamma)
-        positions = torch.arange(CTX, dtype=torch.int32, device="cuda").repeat(REQS).unsqueeze(0)
-        ctx_meta = self.Metadata(
-            seq_lens=torch.tensor([CTX] * REQS, dtype=torch.int),
-            request_ids=self.req_ids,
-            max_num_requests=REQS,
-            num_contexts=REQS,
-            prompt_lens=[CTX] * REQS,
-            max_num_tokens=REQS * CTX,
-            kv_cache_manager=self.kvm,
-            kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0] * REQS),
-            mapping=self.mapping,
-            sparse_metadata_params=self.meta_params,
-        )
-        ctx_meta.prepare()
-        self.step(ctx_meta, positions, hidden, "forward_impl")
-        torch.cuda.synchronize()
-
-    # -- generation inputs -------------------------------------------------
-
-    def gen_inputs(self, layout, seed):
-        nreq, ntok = LAYOUTS[layout]
-        hidden = absorbed_query_fixture.make_hidden(nreq * ntok, seed=seed, gamma=self.gamma)
-        pos = (CTX + torch.arange(ntok, dtype=torch.int32, device="cuda")).repeat(nreq).unsqueeze(0)
-        return hidden, pos
-
-    def gen_metadata(self, layout, cuda_graph=False):
-        """Generation metadata as the runner builds it: graph metadata from
-        ``create_cuda_graph_metadata`` with a buffer registry (reserved,
-        graph-stable buffers, as the graph runner passes its own), MTP verify
-        width from ``update_spec_dec_param``, then ``prepare()``."""
-        nreq, ntok = LAYOUTS[layout]
-        req_ids = self.req_ids[:nreq]
-        kv_params = self.KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[CTX] * nreq)
-        meta = self.Metadata(
-            seq_lens=torch.tensor([ntok] * nreq, dtype=torch.int),
-            request_ids=req_ids,
-            max_num_requests=nreq,
-            num_contexts=0,
-            prompt_lens=[CTX] * nreq,
-            max_num_tokens=nreq * ntok,
-            kv_cache_manager=self.kvm,
-            kv_cache_params=kv_params,
-            mapping=self.mapping,
-            sparse_metadata_params=self.meta_params,
-        )
-        if cuda_graph:
-            from tensorrt_llm._torch.memory_buffer_utils import Buffers
-
-            meta = meta.create_cuda_graph_metadata(nreq, False, ntok - 1, Buffers())
-            assert meta.is_cuda_graph
-            meta.request_ids = req_ids
-            meta.prompt_lens = [CTX] * nreq
-            meta.num_contexts = 0
-            meta.kv_cache_params = kv_params
-        if ntok > 1:
-            meta.update_spec_dec_param(
-                batch_size=nreq,
-                is_spec_decoding_enabled=True,
-                is_spec_dec_tree=False,
-                is_spec_dec_dynamic_tree=False,
-                max_draft_len=ntok - 1,
-                max_total_draft_tokens=ntok - 1,
-            )
-        meta.prepare()
-        return meta
-
-    # -- forwards ----------------------------------------------------------
-
-    def forward(self, layer, positions, hidden, meta, path):
-        attn = self.attn[layer]
-        if path == "forward_impl":
-            out = attn._create_outputs(hidden, meta)
-            attn.forward_impl(positions, hidden, meta, attn_output=out)
-            return out[0]
-        # The model's path: MLA.forward with register_to_config goes through
-        # _forward_custom_op and the registered metadata.
-        import weakref
-
-        from tensorrt_llm._torch.utils import model_extra_attrs
-
-        extra = self.model_config.extra_attrs
-        extra["attention_metadata"] = weakref.ref(meta)
-        with model_extra_attrs(extra):
-            out = [torch.ops.trtllm.create_mla_outputs(hidden, attn.layer_idx_str)]
-            attn._forward_custom_op(hidden, positions, out, None)
-        return out[0]
-
-    def step(self, meta, positions, hidden, path):
-        """One forward step of both layers, as the decoder runner's
-        ``_forward_step`` runs it (eager or captured): ``_preprocess_inputs``
-        first refreshes the metadata derived from the KV lengths
-        (``on_update_kv_lens``), then the model forward, under the step's phase (a
-        generation-only step is a decode step)."""
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-            dispatch,
-        )
-
-        phase = dispatch.Phase.DECODE if meta.num_contexts == 0 else dispatch.Phase.PREFILL
-        meta.on_update_kv_lens()
-        with dispatch.step(phase):
-            return {layer: self.forward(layer, positions, hidden, meta, path) for layer in LAYERS}
-
-    def gen_step(self, layout, path, hidden, pos, meta=None):
-        meta = meta if meta is not None else self.gen_metadata(layout)
-        outs = {layer: out.clone() for layer, out in self.step(meta, pos, hidden, path).items()}
-        torch.cuda.synchronize()
-        return outs
-
-    def recorded_step(self, layout, path, hidden, pos):
-        with _Recorder(self.attn) as rec:
-            out = self.gen_step(layout, path, hidden, pos)
-        return out, rec.snapshot()
-
-    # -- caches ------------------------------------------------------------
-
-    def caches(self):
-        bufs = {f"latent{layer}": self.kvm.get_buffers(layer) for layer in LAYERS}
-        bufs[f"index_k{LAYERS[0]}"] = self.kvm.get_indexer_k_cache_buffers(LAYERS[0])
-        return bufs
-
-    def snapshot(self):
-        return {k: v.clone() for k, v in self.caches().items()}
-
-    def restore(self, snap):
-        for k, v in self.caches().items():
-            v.copy_(snap[k])
-        torch.cuda.synchronize()
-
-    def reset_absorbed_query_runners(self):
-        for layer in LAYERS:
-            self.attn[layer]._absorbed_query = None
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = _Harness()
-    yield h
-    _set_absorbed_query(False)
-
-
-# -- comparison helpers -----------------------------------------------------
-
-
-def _bytes(t):
-    return t.contiguous().view(torch.uint8)
-
-
-def _changed_bytes(before, after):
-    return _bytes(before) != _bytes(after)
-
-
-def _changed_latent_slots(before, after):
-    """(block, token) pairs whose bytes changed in a latent cache buffer
-    ``[blocks, kv_factor, tokens_per_block, heads, head_dim]``."""
-    diff = before.view(torch.uint8) != after.view(torch.uint8)
-    diff = diff.any(dim=4).any(dim=3).any(dim=1)
-    return {tuple(x) for x in diff.nonzero().tolist()}
-
-
-def _index_k_tokens(buf):
-    """Indexer-key cache blocks ``[64 x 128 FP8 | 64 x FP32 scale]`` as
-    (codes [blocks, 64, 128] uint8, scales [blocks, 64] float32)."""
-    raw = buf.contiguous().view(torch.uint8).reshape(buf.shape[0], -1)
-    tokens = buf.shape[1]
-    data = raw[:, : tokens * 128].reshape(-1, tokens, 128)
-    scale = raw[:, tokens * 128 : tokens * 132].contiguous().view(torch.float32).reshape(-1, tokens)
-    return data, scale
-
-
-def _changed_index_k_tokens(before, after):
-    db, sb = _index_k_tokens(before)
-    da, sa = _index_k_tokens(after)
-    changed = (db != da).any(dim=-1) | (_bytes(sb) != _bytes(sa)).view(
-        sb.shape[0], sb.shape[1], 4
-    ).any(dim=-1)
-    return {tuple(x) for x in changed.nonzero().tolist()}
-
-
-def _compare_caches(snap, ref_cache, cand_cache, n_new_tokens, report, tag):
-    """Only the new-token slots may differ from the pre-step snapshot, and
-    the candidate's bytes must equal the reference's."""
-    fails = []
-    for key in cand_cache:
-        ref, cand = ref_cache[key], cand_cache[key]
-        n_diff = int(_changed_bytes(ref, cand).sum())
-        report[f"{tag}:{key}_bytes_differing"] = n_diff
-        if key.startswith("latent"):
-            slots = _changed_latent_slots(snap[key], cand)
-            ref_slots = _changed_latent_slots(snap[key], ref)
-        else:
-            slots = _changed_index_k_tokens(snap[key], cand)
-            ref_slots = _changed_index_k_tokens(snap[key], ref)
-        if len(slots) != n_new_tokens or slots != ref_slots:
-            fails.append(
-                f"{tag}: {key} changed slots {sorted(slots)} vs reference {sorted(ref_slots)}"
-            )
-        if n_diff:
-            fails.append(f"{tag}: {key} differs in {n_diff} bytes")
-    return fails
-
-
-def _compare_records(ref, cand, report, tag):
-    """Bit equality of the recorded MLA attention inputs and indexer inputs;
-    top-k / sparse indices must select equal key sets (row-wise sorted; the
-    per-row overlap is reported when they do not)."""
-    fails = []
-    for layer in LAYERS:
-        r, c = ref[layer], cand[layer]
-        if set(r) != set(c) or "absorbed_q" not in r:
-            fails.append(f"{tag}: layer {layer} recorded {sorted(c)} vs {sorted(r)}")
-            return fails
-        for key in sorted(r):
-            if key in ("topk", "sparse_indices"):
-                rs = r[key].sort(dim=-1).values
-                cs = c[key].sort(dim=-1).values
-                report[f"{tag}:L{layer}.{key}_same_order"] = torch.equal(r[key], c[key])
-                if torch.equal(rs, cs):
-                    continue
-                overlap = min(
-                    len(set(x) & set(y)) / max(len(set(x)), 1)
-                    for x, y in zip(rs.tolist(), cs.tolist())
-                )
-                n_rows = int((rs != cs).any(dim=-1).sum())
-                report[f"{tag}:L{layer}.{key}_rows_differing"] = n_rows
-                report[f"{tag}:L{layer}.{key}_min_row_overlap"] = overlap
-                fails.append(f"{tag}: layer {layer} {key} selects different keys in {n_rows} rows")
-            else:
-                n_diff = int((_bytes(r[key]) != _bytes(c[key])).sum())
-                report[f"{tag}:L{layer}.{key}_bytes_differing"] = n_diff
-                if n_diff:
-                    fails.append(f"{tag}: layer {layer} {key} differs in {n_diff} bytes")
-    return fails
-
-
-def _rel_rms(x, ref):
-    return (
-        (x.double() - ref.double()).pow(2).mean().sqrt()
-        / ref.double().pow(2).mean().sqrt().clamp_min(1e-30)
-    ).item()
-
-
-def _bf16_ulp(value):
-    return 2.0 ** (math.floor(math.log2(max(value, 2.0**-126))) - 7)
-
-
-def _output_bounds(native_runs):
-    """Per-layer bounds from repeated native runs of the same step."""
-    ref = native_runs[0]
-    bounds = {}
-    for layer in LAYERS:
-        noise = max(
-            [
-                (run[layer].float() - ref[layer].float()).abs().max().item()
-                for run in native_runs[1:]
-            ]
-            + [0.0]
-        )
-        rel = max([_rel_rms(run[layer], ref[layer]) for run in native_runs[1:]] + [0.0])
-        ulp = _bf16_ulp(ref[layer].float().abs().max().item())
-        bounds[layer] = (2.0 * max(noise, ulp), max(2.0 * rel, REL_RMS_FLOOR), noise, rel)
-    return bounds
-
-
-def _compare_outputs(ref, cand, bounds, report, tag):
-    fails = []
-    for layer in LAYERS:
-        max_bound, rel_bound, noise, _ = bounds[layer]
-        diff = (cand[layer].float() - ref[layer].float()).abs().max().item()
-        rel = _rel_rms(cand[layer], ref[layer])
-        report[f"{tag}:L{layer}.out_max_diff"] = diff
-        report[f"{tag}:L{layer}.out_rel_rms"] = rel
-        report[f"L{layer}.native_repeat_max_diff"] = noise
-        if diff > max_bound or rel > rel_bound:
-            fails.append(
-                f"{tag}: layer {layer} output max diff {diff:.3e} "
-                f"(bound {max_bound:.3e}), relRMS {rel:.3e} (bound "
-                f"{rel_bound:.3e})"
-            )
-    return fails
-
-
-# Native runs per step used to estimate the attention's run-to-run noise
-# (one reference and four repeats).
-NATIVE_RUNS = 5
-
-
-def _native_runs(h, layout, path, hidden, pos, snap, repeats=NATIVE_RUNS):
-    """Repeated native steps from the same cache snapshot: outputs, the
-    recorded attention inputs and the caches after the first run. Native
-    caches and recorded inputs must themselves be repeatable."""
-    _set_absorbed_query(False)
-    outs, recs, cache, fails = [], [], None, []
-    for i in range(repeats):
-        h.restore(snap)
-        out, rec = h.recorded_step(layout, path, hidden, pos)
-        outs.append(out)
-        recs.append(rec)
-        if i == 0:
-            cache = h.snapshot()
-        else:
-            report = {}
-            fails += _compare_caches(
-                snap, cache, h.snapshot(), math.prod(LAYOUTS[layout]), report, f"native_repeat{i}"
-            )
-            fails += _compare_records(recs[0], rec, report, f"native_repeat{i}")
-    h.restore(snap)
-    assert not fails, f"native step is not repeatable: {fails}"
-    return outs, recs, cache
-
-
-# -- tests ------------------------------------------------------------------
-
-
-def _compare_absorbed_query_to_native(snap, native, candidate, n_new, bounds, report, tag):
-    """absorbed_query step vs native step: caches, MLA and indexer inputs bitwise, equal
-    top-k key sets, noise-bounded outputs."""
-    native_out, native_rec, native_cache = native
-    absorbed_query_out, absorbed_query_rec, absorbed_query_cache = candidate
-    fails = _compare_caches(snap, native_cache, absorbed_query_cache, n_new, report, tag)
-    fails += _compare_records(native_rec, absorbed_query_rec, report, tag)
-    fails += _compare_outputs(native_out, absorbed_query_out, bounds, report, tag)
-    return fails
-
-
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("form", ["auto", "preabsorbed", "native_q"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_absorbed_query_into_native_attention(harness, layout, form, path):
-    h = harness
-    n_new = math.prod(LAYOUTS[layout])
-    hidden, pos = h.gen_inputs(layout, seed=12)
-    snap = h.snapshot()
-    try:
-        native_outs, native_recs, native_cache = _native_runs(h, layout, path, hidden, pos, snap)
-        bounds = _output_bounds(native_outs)
-
-        _set_absorbed_query(True, form)
-        h.reset_absorbed_query_runners()
-        absorbed_query_out, absorbed_query_rec = h.recorded_step(layout, path, hidden, pos)
-        for layer in LAYERS:
-            runner = h.attn[layer]._absorbed_query
-            assert runner not in (None, False), f"absorbed_query did not engage on {layer}"
-        absorbed_query_cache = h.snapshot()
-
-        report, fails = {}, []
-        fails += _compare_absorbed_query_to_native(
-            snap,
-            (native_outs[0], native_recs[0], native_cache),
-            (absorbed_query_out, absorbed_query_rec, absorbed_query_cache),
-            n_new,
-            bounds,
-            report,
-            "absorbed_query",
-        )
-
-        # absorbed_query is deterministic: the same step again from the snapshot ...
-        h.restore(snap)
-        absorbed_query_again, absorbed_query_rec_again = h.recorded_step(layout, path, hidden, pos)
-        fails += _compare_caches(
-            snap, absorbed_query_cache, h.snapshot(), n_new, report, "absorbed_query_again"
-        )
-        fails += _compare_records(
-            absorbed_query_rec, absorbed_query_rec_again, report, "absorbed_query_again"
-        )
-        # ... and 20 more steps rewrite the same slots with the same bytes.
-        for _ in range(20):
-            again = h.gen_step(layout, path, hidden, pos)
-        fails += _compare_caches(
-            snap, absorbed_query_cache, h.snapshot(), n_new, report, "absorbed_query_repeat20"
-        )
-        fails += _compare_outputs(native_outs[0], again, bounds, report, "absorbed_query_repeat20")
-        print(f"layout={layout} form={form} path={path} {report}")
-        assert not fails, fails
-    finally:
-        h.restore(snap)
-        _set_absorbed_query(False)
-
-
-# Graphs captured by the CUDA-graph test: absorbed_query alone, native alone, or both in
-# one process sharing one memory pool (as the graph runner keeps several
-# graphs).
-GRAPH_SETS = {
-    "p1only": ("absorbed_query",),
-    "nativeonly": ("native",),
-    "both": ("absorbed_query", "native"),
-}
-
-
-@pytest.mark.parametrize("graph_set", list(GRAPH_SETS))
-@pytest.mark.parametrize("form", ["auto", "preabsorbed", "native_q"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_absorbed_query_cuda_graph_replay_into_native_attention(harness, layout, form, graph_set):
-    """Captured absorbed_query + native attention + native cache append, replayed once
-    and 20 times with fresh inputs. Each absorbed_query replay must reproduce the eager
-    absorbed_query step bit for bit (caches, recorded attention inputs) and compare with
-    the eager native step as the eager test does; a captured native graph
-    must reproduce the eager native step bit for bit.
-
-    Capture follows the graph runner: every configuration is warmed up on
-    the current stream first, then all are captured on ``torch.cuda.graph``'s
-    own capture stream into one shared memory pool."""
-    h = harness
-    n_new = math.prod(LAYOUTS[layout])
-    path = "custom_op"
-    snap = h.snapshot()
-    graphs = {}
-    try:
-        _graph_replay_checks(h, layout, form, n_new, path, snap, graphs, GRAPH_SETS[graph_set])
-    except BaseException as exc:
-        # Report before the graphs are destroyed: a CUDA error raised while
-        # tearing them down would otherwise hide the original failure.
-        print(f"graph test failure: {exc!r}", flush=True)
-        raise
-    finally:
-        graphs.clear()
-        h.reset_absorbed_query_runners()
-        h.restore(snap)
-        _set_absorbed_query(False)
-
-
-def _graph_replay_checks(h, layout, form, n_new, path, snap, graphs, wanted):
-    """Body of the CUDA-graph test; ``graphs`` is owned by the caller."""
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    configs = [
-        (name, name == "absorbed_query") for name in ("absorbed_query", "native") if name in wanted
-    ]
-    # Diagnostics only: GLM_MOE_DSA_ABSORBED_QUERY_IT_MULTI_STREAM=0 warms up and captures without
-    # the aux-stream overlap the graph runner enables.
-    multi_stream = os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_IT_MULTI_STREAM", "1") != "0"
-
-    def step(meta, static_pos, static_hidden):
-        # The runner's step: on_update_kv_lens() is captured too, so the
-        # graph recomputes the graph-temporary metadata (indexer scan
-        # lengths, slot mappings) instead of addressing tensors that a later
-        # prepare() rebinds.
-        return h.step(meta, static_pos, static_hidden, path)
-
-    # Pass 1: warm up every configuration (graph metadata, runners, JIT).
-    warm = {}
-    for name, enabled in configs:
-        _set_absorbed_query(enabled, form)
-        h.reset_absorbed_query_runners()
-        meta = h.gen_metadata(layout, cuda_graph=True)
-        static_hidden, static_pos = h.gen_inputs(layout, seed=200)
-        with with_multi_stream(multi_stream):
-            for _ in range(2):
-                step(meta, static_pos, static_hidden)
-        torch.cuda.synchronize()
-        print(f"[graph test] warmed up {name}", flush=True)
-        # The captured graph addresses its runners' counters: keep them.
-        runners = [h.attn[layer]._absorbed_query for layer in LAYERS]
-        if enabled:
-            for layer, runner in zip(LAYERS, runners):
-                assert runner not in (None, False), f"absorbed_query did not engage on {layer}"
-        warm[name] = (meta, static_hidden, static_pos, runners)
-    h.restore(snap)
-    # Pass 2: capture into one shared pool.
-    pool = None
-    for name, enabled in configs:
-        meta, static_hidden, static_pos, runners = warm[name]
-        _set_absorbed_query(enabled, form)
-        for layer, runner in zip(LAYERS, runners):
-            h.attn[layer]._absorbed_query = runner
-        graph = torch.cuda.CUDAGraph()
-        with with_multi_stream(multi_stream), _Recorder(h.attn) as rec:
-            with torch.cuda.graph(graph, pool=pool):
-                outs = step(meta, static_pos, static_hidden)
-        pool = graph.pool()
-        graphs[name] = (graph, meta, static_hidden, outs, rec, runners)
-        torch.cuda.synchronize()
-        print(f"[graph test] captured {name}", flush=True)
-    # Eager absorbed_query steps below use their own runners.
-    h.reset_absorbed_query_runners()
-
-    report, fails = {}, []
-    for trial, (seed, replays) in enumerate([(301, 1), (302, 1), (303, 20)]):
-        hidden, pos = h.gen_inputs(layout, seed=seed)
-        native_outs, native_recs, native_cache = _native_runs(h, layout, path, hidden, pos, snap)
-        bounds = _output_bounds(native_outs)
-        _set_absorbed_query(True, form)
-        h.restore(snap)
-        absorbed_query_eager = h.recorded_step(layout, path, hidden, pos)
-        absorbed_query_eager_cache = h.snapshot()
-        _set_absorbed_query(False)
-        print(f"[graph test] trial {trial}: eager steps done", flush=True)
-        replayed = {}
-        for name, (graph, meta, static_hidden, outs, rec, runners) in graphs.items():
-            h.restore(snap)
-            meta.prepare()  # host-side per-step preparation (no realloc)
-            static_hidden.copy_(hidden)
-            for _ in range(replays):
-                graph.replay()
-            torch.cuda.synchronize()
-            print(f"[graph test] trial {trial}: {name} x{replays} replayed", flush=True)
-            replayed[name] = (
-                {layer: outs[layer].clone() for layer in LAYERS},
-                rec.snapshot(),
-                h.snapshot(),
-            )
-        tag = f"t{trial}x{replays}"
-        if "absorbed_query" in replayed:
-            absorbed_query_out, absorbed_query_rec, absorbed_query_cache = replayed[
-                "absorbed_query"
-            ]
-            # Replays reproduce the eager absorbed_query step bit for bit ...
-            fails += _compare_caches(
-                snap,
-                absorbed_query_eager_cache,
-                absorbed_query_cache,
-                n_new,
-                report,
-                f"{tag}:absorbed_query_graph_vs_absorbed_query_eager",
-            )
-            fails += _compare_records(
-                absorbed_query_eager[1],
-                absorbed_query_rec,
-                report,
-                f"{tag}:absorbed_query_graph_vs_absorbed_query_eager",
-            )
-            # ... and compare with native as in the eager test.
-            fails += _compare_absorbed_query_to_native(
-                snap,
-                (native_outs[0], native_recs[0], native_cache),
-                (absorbed_query_out, absorbed_query_rec, absorbed_query_cache),
-                n_new,
-                bounds,
-                report,
-                f"{tag}:absorbed_query_graph_vs_native",
-            )
-        if "native" in replayed:
-            g_out, g_rec, g_cache = replayed["native"]
-            fails += _compare_caches(
-                snap, native_cache, g_cache, n_new, report, f"{tag}:native_graph_vs_eager"
-            )
-            fails += _compare_records(native_recs[0], g_rec, report, f"{tag}:native_graph_vs_eager")
-            fails += _compare_outputs(
-                native_outs[0], g_out, bounds, report, f"{tag}:native_graph_vs_eager"
-            )
-        absorbed_query_runners = graphs["absorbed_query"][5] if "absorbed_query" in graphs else []
-        for runner in absorbed_query_runners + [h.attn[layer]._absorbed_query for layer in LAYERS]:
-            if runner not in (None, False):
-                assert int(runner.sync.abs().sum()) == 0
-    print(f"graph layout={layout} form={form} graphs={wanted} {report}")
-    assert not fails, fails
-
-
-@torch.inference_mode()
-def test_checks_detect_a_wrong_preabsorbed_query(harness, monkeypatch):
-    """Negative control: an absorbed_query packet whose absorbed query is left unwritten
-    (zeroed) must fail both the absorbed-query / quant_q checks and the
-    output bound."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    h = harness
-    layout, path = "6req_x1tok", "forward_impl"
-    hidden, pos = h.gen_inputs(layout, seed=12)
-    snap = h.snapshot()
-    orig_project = query_proj.AbsorbedQueryProjection.project
-
-    def broken_project(self, *args, **kwargs):
-        packet = orig_project(self, *args, **kwargs)
-        if packet.fused_q is not None:
-            packet.fused_q[..., : self.mla.kv_lora_rank].zero_()
-        return packet
-
-    try:
-        native_outs, native_recs, _ = _native_runs(h, layout, path, hidden, pos, snap)
-        bounds = _output_bounds(native_outs)
-        monkeypatch.setattr(query_proj.AbsorbedQueryProjection, "project", broken_project)
-        _set_absorbed_query(True, "preabsorbed")
-        h.reset_absorbed_query_runners()
-        bad_out, bad_rec = h.recorded_step(layout, path, hidden, pos)
-        report = {}
-        rec_fails = _compare_records(native_recs[0], bad_rec, report, "bad")
-        out_fails = _compare_outputs(native_outs[0], bad_out, bounds, report, "bad")
-        print(f"negative control {report}")
-        assert any("absorbed_q" in f for f in rec_fails), rec_fails
-        assert any("quant_q" in f for f in rec_fails), rec_fails
-        assert len(out_fails) == len(LAYERS), out_fails
-    finally:
-        monkeypatch.undo()
-        h.restore(snap)
-        _set_absorbed_query(False)
-        h.reset_absorbed_query_runners()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py
deleted file mode 100644
index 58ec3fd16b63..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_absorbed_query_proj.py
+++ /dev/null
@@ -1,869 +0,0 @@
-# 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.
-"""absorbed_query: query_front packet correctness against the native TRT-LLM region.
-
-Fixtures: real GLM-5.2 checkpoint weights of one full-indexer layer (6) and
-one shared layer (7) at TP8 rank 0 (see ``absorbed_query_fixture``); synthetic
-activations (unit-RMS rows times the layer's real ``input_layernorm`` gain).
-
-Asserted properties:
-
-* The frozen error budget (fixed before tuning) for every floating output:
-  ``allclose(candidate, native, rtol=0.1, atol=1e-2)`` -- the tolerance of
-  the native fused-A GEMM test -- and normalized RMS error vs an FP64 oracle
-  (with the native BF16 intermediate roundings) ``<= 1.25 * native_error +
-  2**-12``.
-* Native agreement (``_native_agreement``), checked on every launch of the
-  default kernels that is compared with native: every output of the 4/9
-  list and the absorbed query ``fused_q[..., :512]`` is bitwise identical
-  to native -- ``q``, ``compressed_kv`` / ``latent[:, :512]`` (FlashInfer
-  RMSNorm fold order), ``k_pe`` / ``latent[:, 512:]``, the indexer
-  ``q_fp8`` / ``k_fp8`` codes, their UE8M0 scales and the head ``weights``.
-  The indexer key LayerNorm is the native ``k_norm`` call (its inductor
-  kernel is shared with the native region), the RoPE is FlashInfer's
-  ``fma(x, cos, rn(+-x_pair * sin))`` and the FP8 tail is the native
-  ``fused_cat_fp8`` arithmetic. The experimental prefetch kernel (own q_a
-  norm) is held to the frozen budget only.
-
-  ``test_native_agreement_statistics`` reports how often each output
-  differs over many inputs (with the first differing elements), and
-  whether the indexer's BF16 ``wq_b`` output already differs from the
-  native GEMM.
-* absorbed_query determinism: padded rows never change live rows, CUDA-graph replays
-  and every stress launch reproduce an undisturbed eager absorbed_query launch bit for
-  bit, and the per-layer sync counters are zero after every launch.
-* Stress (new inter-program synchronization): delayed producers/consumers,
-  many concurrent launches on several streams, and progress while a
-  concurrent kernel holds all but ``k`` SMs.
-
-``GLM_MOE_DSA_ABSORBED_QUERY_TEST_ROWS`` (comma list) and
-``GLM_MOE_DSA_ABSORBED_QUERY_TEST_PREFETCH`` (0/1) restrict the parametrization, e.g. for
-compute-sanitizer runs.
-"""
-
-import os
-import time
-
-import absorbed_query_fixture
-import pytest
-import torch
-import triton
-import triton.language as tl
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-_EPS_FLOOR = 2.0**-12
-
-
-def _env_list(name, default):
-    value = os.environ.get(name)
-    if not value:
-        return default
-    return [type(default[0])(int(x)) for x in value.split(",")]
-
-
-ROWS = _env_list("GLM_MOE_DSA_ABSORBED_QUERY_TEST_ROWS", [1, 2, 3, 4, 5, 6, 8, 16])
-PREFETCH = _env_list("GLM_MOE_DSA_ABSORBED_QUERY_TEST_PREFETCH", [True, False])
-STRESS_ROWS = _env_list("GLM_MOE_DSA_ABSORBED_QUERY_TEST_ROWS", [1, 6, 16])
-STRESS_LAUNCHES = int(os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_STRESS_LAUNCHES", "8"))
-LAYERS = [absorbed_query_fixture.FULL_LAYER, absorbed_query_fixture.SHARED_LAYER]
-
-# Kernel configurations whose inter-program waits are exercised by the
-# stress tests: the default (absorption of head h by the producers of head
-# h + 1), separate spinning absorption programs, the prefetching kernel, and
-# the legacy own-head-group layout (reference only).
-STRESS_CONFIGS = {
-    "default": {},
-    "separate_consumers": dict(absorb_in_producer=False),
-    "prefetch": dict(prefetch=True),
-    "legacy_own_head": dict(absorb_prev_head=False),
-}
-
-
-@pytest.fixture(scope="module")
-def model_config():
-    return absorbed_query_fixture.build_model_config(0)
-
-
-_ATTN_CACHE = {}
-
-
-def _attn(model_config, layer):
-    if layer not in _ATTN_CACHE:
-        _ATTN_CACHE[layer] = absorbed_query_fixture.build_native_attention(layer, 0, model_config)
-    return _ATTN_CACHE[layer]
-
-
-def _gamma(layer):
-    return absorbed_query_fixture.load_input_layernorm(layer)
-
-
-def _native(attn, hidden, position_ids):
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_dsa_proj
-
-    outs = forward_dsa_proj(attn, position_ids, hidden, None)
-    rows = hidden.shape[0]
-    fused_q = torch.zeros(
-        rows,
-        attn.num_heads_tp,
-        attn.kv_lora_rank + attn.qk_rope_head_dim,
-        dtype=hidden.dtype,
-        device=hidden.device,
-    )
-    q_nope = outs[0].view(rows, attn.num_heads_tp, attn.qk_head_dim)[..., : attn.qk_nope_head_dim]
-    attn._bmm_bf16_out(
-        q_nope.transpose(0, 1),
-        attn.k_b_proj_trans,
-        attn.k_b_proj_trans.transpose(1, 2),
-        fused_q[..., : attn.kv_lora_rank].transpose(0, 1),
-    )
-    return outs, fused_q
-
-
-def _oracle(attn, hidden):
-    """FP64 reference with the native BF16 intermediate roundings."""
-    kva = (hidden.double() @ attn.kv_a_proj_with_mqa.weight.double().t()).to(torch.bfloat16)
-    rows = hidden.shape[0]
-    q_a = kva[:, : attn.q_lora_rank].double()
-    eps = attn.q_a_layernorm.variance_epsilon
-    qn = q_a * torch.rsqrt(q_a.pow(2).mean(-1, keepdim=True) + eps)
-    qn = (qn * attn.q_a_layernorm.weight.double()).to(torch.bfloat16)
-    q = qn.double() @ attn.q_b_proj.weight.double().t()
-    q_b16 = q.to(torch.bfloat16).view(rows, attn.num_heads_tp, attn.qk_head_dim)
-    absorbed = torch.einsum(
-        "rhd,hcd->rhc", q_b16[..., : attn.qk_nope_head_dim].double(), attn.k_b_proj_trans.double()
-    )
-    ckv = kva[:, attn.q_lora_rank : attn.q_lora_rank + attn.kv_lora_rank].double()
-    ckv = ckv * torch.rsqrt(
-        ckv.pow(2).mean(-1, keepdim=True) + attn.kv_a_layernorm.variance_epsilon
-    )
-    ckv = ckv * attn.kv_a_layernorm.weight.double()
-    k_pe = kva[:, attn.q_lora_rank + attn.kv_lora_rank :]
-    return {"qr": qn.double(), "q": q, "absorbed": absorbed, "ckv": ckv, "k_pe": k_pe}
-
-
-def _rel_rms(x, ref):
-    x = x.double()
-    return ((x - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-30)).item()
-
-
-def _bytes(t):
-    return t.contiguous().view(torch.uint8)
-
-
-def _assert_bitwise(name, cand, native):
-    assert cand.shape == native.shape and cand.dtype == native.dtype, name
-    diff = _bytes(cand) != _bytes(native)
-    if bool(diff.any()):
-        max_abs = (cand.float() - native.float()).abs().max().item()
-        raise AssertionError(
-            f"{name}: {int(diff.sum())} of {diff.numel()} bytes differ "
-            f"(max abs difference {max_abs:.3e})"
-        )
-
-
-def _ordered(t):
-    """Order-preserving integer view: adjacent representable values differ
-    by 1 (+0 and -0 both map to 0)."""
-    t = t.contiguous()
-    if t.dtype == torch.float8_e4m3fn:
-        u = t.view(torch.uint8).to(torch.int32)
-        mag = u & 0x7F
-        return torch.where(u >= 0x80, -mag, mag)
-    if t.dtype == torch.bfloat16:
-        i = t.view(torch.int16).to(torch.int32)
-        return torch.where(i < 0, -(i & 0x7FFF), i)
-    if t.dtype == torch.float32:
-        i = t.view(torch.int32).to(torch.int64)
-        return torch.where(i < 0, -(i & 0x7FFFFFFF), i)
-    raise TypeError(t.dtype)
-
-
-def _steps(cand, native):
-    """Per-element distance in representable values (ulps / codes)."""
-    assert cand.shape == native.shape and cand.dtype == native.dtype
-    return (_ordered(cand) - _ordered(native)).abs()
-
-
-def _first_diffs(cand, ref, n=4):
-    """Up to ``n`` differing elements as ``(index, candidate, native)``."""
-    idx = (_ordered(cand) != _ordered(ref)).nonzero()[:n].tolist()
-    return [(tuple(i), cand[tuple(i)].float().item(), ref[tuple(i)].float().item()) for i in idx]
-
-
-def _native_agreement(outs, fused_q, native_outs, native_fq, kv_lora, rows=slice(None), check=True):
-    """Per-output (differing elements, elements, max steps) statistics of
-    absorbed_query vs native; with ``check`` asserts that every output is bitwise
-    identical to native (the module docstring)."""
-    out = dict(zip(_NAMES, outs))
-    nat = dict(zip(_NAMES, native_outs))
-    pairs = {name: (out[name][rows], nat[name][rows]) for name in ("q", "k_pe", "compressed_kv")}
-    latent_c, latent_n = out["latent_cache"][rows], nat["latent_cache"][rows]
-    pairs["latent_nope"] = (latent_c[:, :kv_lora], latent_n[:, :kv_lora])
-    pairs["latent_rope"] = (latent_c[:, kv_lora:], latent_n[:, kv_lora:])
-    if fused_q is not None:
-        pairs["absorbed"] = (fused_q[rows, :, :kv_lora], native_fq[rows, :, :kv_lora])
-    if "q_fp8" in out:
-        rope = 64
-        for side in ("q", "k"):
-            codes_c, codes_n = out[f"{side}_fp8"][rows], nat[f"{side}_fp8"][rows]
-            pairs[f"{side}_fp8_rope"] = (codes_c[..., :rope], codes_n[..., :rope])
-            pairs[f"{side}_fp8_nope"] = (codes_c[..., rope:], codes_n[..., rope:])
-            pairs[f"{side}_scale"] = (out[f"{side}_scale"][rows], nat[f"{side}_scale"][rows])
-        pairs["weights"] = (out["weights"][rows], nat["weights"][rows])
-    stats, fails = {}, []
-    for name, (cand, ref) in pairs.items():
-        steps = _steps(cand, ref)
-        n_diff = int((steps != 0).sum())
-        stats[name] = (n_diff, steps.numel(), int(steps.max()) if steps.numel() else 0)
-        if n_diff:
-            fails.append(
-                f"{name} differs from native: {stats[name]}, first "
-                f"(index, absorbed_query, native) {_first_diffs(cand, ref)}"
-            )
-    if check:
-        assert not fails, f"{fails} (stats {stats})"
-    return stats, fails
-
-
-def _assert_same_packet(a, b, name):
-    """Two absorbed_query launches on the same inputs must agree bit for bit."""
-    assert len(a.proj_outputs) == len(b.proj_outputs)
-    for out_name, x, y in zip(_NAMES, a.proj_outputs, b.proj_outputs):
-        _assert_bitwise(f"{name} {out_name}", x, y)
-    assert (a.fused_q is None) == (b.fused_q is None)
-    if a.fused_q is not None:
-        kv_lora = a.proj_outputs[1].shape[1]
-        _assert_bitwise(f"{name} absorbed", a.fused_q[..., :kv_lora], b.fused_q[..., :kv_lora])
-
-
-def _check_budget(name, cand, native, ref, report):
-    torch.testing.assert_close(
-        cand.float(), native.float(), rtol=0.1, atol=1e-2, msg=f"{name}: candidate vs native"
-    )
-    e_c, e_n = _rel_rms(cand, ref), _rel_rms(native, ref)
-    report[name] = (e_c, e_n)
-    assert e_c <= 1.25 * e_n + _EPS_FLOOR, (
-        f"{name}: candidate rel-RMS {e_c:.3e} vs native {e_n:.3e}"
-    )
-
-
-# Index of each output in the native 4/9-tensor list.
-_NAMES = (
-    "q",
-    "compressed_kv",
-    "k_pe",
-    "latent_cache",
-    "q_fp8",
-    "k_fp8",
-    "k_scale",
-    "weights",
-    "q_scale",
-)
-
-
-def _inputs(layer, rows, seed=0, padded_from=None):
-    hidden = absorbed_query_fixture.make_hidden(rows, seed=seed + layer, gamma=_gamma(layer))
-    if padded_from is not None:
-        hidden[padded_from:] = 0
-    position_ids = torch.arange(491061, 491061 + rows, dtype=torch.int32, device="cuda").unsqueeze(
-        0
-    )
-    return hidden, position_ids
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("rows", ROWS)
-@pytest.mark.parametrize("absorb", [True, False])
-@pytest.mark.parametrize("fuse_indexer", [True, False])
-@pytest.mark.parametrize("prefetch", PREFETCH)
-@torch.inference_mode()
-def test_contract_q_matches_native(model_config, layer, rows, absorb, fuse_indexer, prefetch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    if fuse_indexer and layer != absorbed_query_fixture.FULL_LAYER:
-        pytest.skip("indexer tail exists on full indexer layers only")
-    attn = _attn(model_config, layer)
-    hidden, position_ids = _inputs(layer, rows)
-    native_outs, native_fq = _native(attn, hidden, position_ids)
-    runner = AbsorbedQueryProjection(
-        attn, KernelConfig(fuse_indexer=fuse_indexer, prefetch=prefetch)
-    )
-    packet = runner.project(position_ids, hidden, absorb=absorb)
-    torch.cuda.synchronize()
-    ref = _oracle(attn, hidden)
-    outs = packet.proj_outputs
-    report = {}
-
-    assert len(outs) == len(native_outs) == (9 if attn.mqa.indexer else 4)
-    for cand, nat in zip(outs, native_outs):
-        assert cand.shape == nat.shape and cand.dtype == nat.dtype
-    assert (packet.fused_q is not None) == absorb
-    # Views of the latent row, as in the native list.
-    assert torch.equal(outs[3][:, : attn.kv_lora_rank], outs[1])
-    assert torch.equal(outs[3][:, attn.kv_lora_rank :], outs[2])
-    # Frozen budget on every floating output.
-    _check_budget("q", outs[0], native_outs[0], ref["q"], report)
-    _check_budget("compressed_kv", outs[1], native_outs[1], ref["ckv"], report)
-    if absorb:
-        _check_budget(
-            "absorbed",
-            packet.fused_q[..., : attn.kv_lora_rank],
-            native_fq[..., : attn.kv_lora_rank],
-            ref["absorbed"],
-            report,
-        )
-    if attn.mqa.indexer is not None:
-        q_fp8, k_fp8, k_scale, weights, q_scale = outs[4:]
-        nq_fp8, nk_fp8, nk_scale, nweights, nq_scale = native_outs[4:]
-        torch.testing.assert_close(
-            k_fp8.float() * k_scale, nk_fp8.float() * nk_scale, rtol=0.1, atol=1e-2
-        )
-        torch.testing.assert_close(
-            q_fp8.float() * q_scale, nq_fp8.float() * nq_scale, rtol=0.1, atol=1e-2
-        )
-        torch.testing.assert_close(weights, nweights, rtol=0.1, atol=1e-2)
-    # Bitwise native agreement; the experimental prefetch kernel normalizes
-    # q_a itself and is held to the frozen budget only.
-    stats, _ = _native_agreement(
-        outs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank, check=not prefetch
-    )
-    print(
-        f"layer={layer} rows={rows} absorb={absorb} "
-        f"fuse_indexer={fuse_indexer} prefetch={prefetch} "
-        f"budget(cand_relrms, native_relrms)={report} "
-        f"agreement(n_diff, n, max_steps)={stats}"
-    )
-    assert int(runner.sync.abs().sum()) == 0, "counters must reset"
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("rows", [6, 16])
-@torch.inference_mode()
-def test_native_agreement_statistics(model_config, layer, rows):
-    """Default kernel over many inputs: how often each output differs from
-    native (and by how many steps), and whether the fused indexer's BF16
-    wq_b output already differs from the native GEMM. Every input must
-    agree with native bit for bit."""
-    import json
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    attn = _attn(model_config, layer)
-    indexer = attn.mqa.indexer
-    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_keep_intermediates=True))
-    n_inputs = int(os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_AGREEMENT_INPUTS", "48"))
-    totals, inputs_with_diff, fails = {}, {}, []
-    for i in range(n_inputs):
-        hidden, position_ids = _inputs(layer, rows, seed=1000 + 17 * i)
-        native_outs, native_fq = _native(attn, hidden, position_ids)
-        packet = runner.project(position_ids, hidden, absorb=True)
-        stats, input_fails = _native_agreement(
-            packet.proj_outputs,
-            packet.fused_q,
-            native_outs,
-            native_fq,
-            attn.kv_lora_rank,
-            check=False,
-        )
-        if indexer is not None:
-            kva = attn.kv_a_proj_with_mqa(hidden)
-            qr = attn._q_a_layernorm_maybe_fused(kva[:, : attn.q_lora_rank])
-            steps = _steps(runner.debug["iq"], indexer.wq_b(qr))
-            stats["indexer_wq_b_bf16"] = (int((steps != 0).sum()), steps.numel(), int(steps.max()))
-        fails += [f"input {i}: {f}" for f in input_fails]
-        for name, (n_diff, n, max_steps) in stats.items():
-            tot = totals.setdefault(name, [0, 0, 0])
-            tot[0] += n_diff
-            tot[1] += n
-            tot[2] = max(tot[2], max_steps)
-            inputs_with_diff[name] = inputs_with_diff.get(name, 0) + (n_diff > 0)
-    summary = {
-        name: dict(
-            inputs_with_diff=inputs_with_diff[name],
-            inputs=n_inputs,
-            diff_elements=tot[0],
-            elements=tot[1],
-            max_steps=tot[2],
-        )
-        for name, tot in totals.items()
-    }
-    print(f"AGREEMENT layer={layer} rows={rows} {json.dumps(summary)}")
-    assert not fails, fails
-    assert int(runner.sync.abs().sum()) == 0
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("absorb", [True, False])
-@torch.inference_mode()
-def test_padded_rows_do_not_affect_live_rows(model_config, layer, absorb):
-    """Rows are independent: padded rows (zero or 1e3 garbage) never change
-    any live-row output (all 4/9 outputs and the absorbed query, bit for
-    bit), live rows agree with native bit for bit, and zero padded rows
-    stay finite."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-    )
-
-    attn = _attn(model_config, layer)
-    runner = AbsorbedQueryProjection(attn)
-    live = slice(0, 3)
-    hidden_zero, position_ids = _inputs(layer, 6, padded_from=3)
-    hidden_junk = hidden_zero.clone()
-    hidden_junk[3:] = 1e3 * torch.randn_like(hidden_junk[3:])
-    a = runner.project(position_ids, hidden_zero, absorb=absorb)
-    b = runner.project(position_ids, hidden_junk, absorb=absorb)
-    native_outs, native_fq = _native(attn, hidden_junk, position_ids)
-    torch.cuda.synchronize()
-    assert len(a.proj_outputs) == (9 if attn.mqa.indexer else 4)
-    for name, x, y in zip(_NAMES, a.proj_outputs, b.proj_outputs):
-        _assert_bitwise(f"padded {name}", x[live], y[live])
-    if absorb:
-        _assert_bitwise(
-            "padded absorbed",
-            a.fused_q[live, :, : attn.kv_lora_rank],
-            b.fused_q[live, :, : attn.kv_lora_rank],
-        )
-    _native_agreement(b.proj_outputs, b.fused_q, native_outs, native_fq, attn.kv_lora_rank, live)
-    for x in a.proj_outputs:
-        assert torch.isfinite(x[3:].float()).all()
-    assert int(runner.sync.abs().sum()) == 0
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("absorb", [True, False])
-@pytest.mark.parametrize("delay_ns", [0, 20000])
-@torch.inference_mode()
-def test_cuda_graph_replay(model_config, layer, absorb, delay_ns):
-    """One and many replays of a captured absorbed_query call (also with delayed
-    producers/consumers) reproduce an eager absorbed_query launch on the same inputs
-    bit for bit (all 4/9 outputs and the absorbed query), agree with
-    native bit for bit, and leave the per-layer counters at zero."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    attn = _attn(model_config, layer)
-    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_delay_ns=delay_ns))
-    eager = AbsorbedQueryProjection(attn)
-    rows = 6
-    static_hidden, position_ids = _inputs(layer, rows)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(2):
-            runner.project(position_ids, static_hidden, absorb=absorb)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        packet = runner.project(position_ids, static_hidden, absorb=absorb)
-    assert len(packet.proj_outputs) == (9 if attn.mqa.indexer else 4)
-    for step, replays in enumerate([1, 1, 50]):
-        hidden, _ = _inputs(layer, rows, seed=100 + step)
-        static_hidden.copy_(hidden)
-        for _ in range(replays):
-            graph.replay()
-        torch.cuda.synchronize()
-        reference = eager.project(position_ids, hidden, absorb=absorb)
-        native_outs, native_fq = _native(attn, hidden, position_ids)
-        torch.cuda.synchronize()
-        _assert_same_packet(packet, reference, f"replay x{replays}")
-        _native_agreement(
-            packet.proj_outputs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
-        )
-        assert int(runner.sync.abs().sum()) == 0
-        assert int(eager.sync.abs().sum()) == 0
-
-
-def _stress_runner(attn, config_name, delay_ns=0):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    return AbsorbedQueryProjection(
-        attn, KernelConfig(debug_delay_ns=delay_ns, **STRESS_CONFIGS[config_name])
-    )
-
-
-def _check_stressed(packet, reference, native, kv_lora, config_name, tag):
-    """A stressed launch must equal the undisturbed launch of the same
-    configuration bit for bit, and (except the experimental prefetch
-    kernel) agree with native bit for bit."""
-    _assert_same_packet(packet, reference, tag)
-    native_outs, native_fq = native
-    _native_agreement(
-        packet.proj_outputs,
-        packet.fused_q,
-        native_outs,
-        native_fq,
-        kv_lora,
-        check=config_name != "prefetch",
-    )
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("rows", STRESS_ROWS)
-@pytest.mark.parametrize("absorb", [True, False])
-@pytest.mark.parametrize("config_name", list(STRESS_CONFIGS))
-@torch.inference_mode()
-def test_delayed_producers_and_consumers(model_config, layer, rows, absorb, config_name):
-    """Producers delayed before they publish and consumers delayed before
-    they release/reset (up to 50 us, deterministic per program): many
-    back-to-back launches with fresh inputs reproduce the undelayed launch
-    bit for bit and leave the counters at zero."""
-    attn = _attn(model_config, layer)
-    runner = _stress_runner(attn, config_name, delay_ns=50000)
-    plain = _stress_runner(attn, config_name)
-    launches = []
-    for i in range(STRESS_LAUNCHES):
-        hidden, position_ids = _inputs(layer, rows, seed=500 + i)
-        launches.append(
-            (
-                runner.project(position_ids, hidden, absorb=absorb),
-                plain.project(position_ids, hidden, absorb=absorb),
-                _native(attn, hidden, position_ids),
-            )
-        )
-    torch.cuda.synchronize()
-    for i, (packet, reference, native) in enumerate(launches):
-        _check_stressed(
-            packet, reference, native, attn.kv_lora_rank, config_name, f"delayed launch {i}"
-        )
-    assert int(runner.sync.abs().sum()) == 0
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("config_name", list(STRESS_CONFIGS))
-@torch.inference_mode()
-def test_concurrent_launches_on_many_streams(model_config, layer, config_name):
-    """Six independent absorbed_query instances (own counters, like six layers) launched
-    concurrently on six streams, 30 rounds, with and without delays: every
-    stream's last result equals its undisturbed launch bit for bit, the
-    counters are zero, no hang."""
-    attn = _attn(model_config, layer)
-    n_streams, rounds, rows = 6, 30, 6
-    streams = [torch.cuda.Stream() for _ in range(n_streams)]
-    runners = [
-        _stress_runner(attn, config_name, delay_ns=(5000 if s % 2 else 0)) for s in range(n_streams)
-    ]
-    plain = _stress_runner(attn, config_name)
-    inputs = [_inputs(layer, rows, seed=700 + s) for s in range(n_streams)]
-    references = [plain.project(p, h, absorb=True) for h, p in inputs]
-    natives = [_native(attn, h, p) for h, p in inputs]
-    torch.cuda.synchronize()
-    last = [None] * n_streams
-    for _ in range(rounds):
-        for s in range(n_streams):
-            with torch.cuda.stream(streams[s]):
-                hidden, position_ids = inputs[s]
-                last[s] = runners[s].project(position_ids, hidden, absorb=True)
-    torch.cuda.synchronize()
-    for s in range(n_streams):
-        _check_stressed(
-            last[s], references[s], natives[s], attn.kv_lora_rank, config_name, f"stream {s}"
-        )
-        assert int(runners[s].sync.abs().sum()) == 0
-
-
-def _memcpy_copy(t):
-    """A new device tensor initialized by a host-to-device memcpy only."""
-    return t.cpu().to(t.device)
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("rows", STRESS_ROWS)
-@pytest.mark.parametrize("absorb", [True, False])
-@torch.inference_mode()
-def test_absorbed_query_kernel_with_memcpy_initialized_inputs(model_config, layer, rows, absorb):
-    """The absorbed_query kernel launch alone, with every tensor it reads placed by a
-    host-to-device copy: the kv_a output, the q_a_layernorm output, the
-    indexer wk/weights_proj output, the native k_norm output, positions,
-    RoPE-table rows, sync counters and the weights (swapped in for the
-    launch). The result must
-    equal the normal launch bit for bit.
-
-    For ``compute-sanitizer --tool initcheck --kernel-name kns=qproj
-    --check-api-memory-access no`` the suite sets ``GLM_MOE_DSA_ABSORBED_QUERY_INITCHECK_ONLY=1``:
-    then only the memcpy-fed launch runs, so the only instrumented kernel
-    reads only memory placed by memcpys or written by itself, and any report
-    is a read of memory nothing initialized."""
-    initcheck_only = os.environ.get("GLM_MOE_DSA_ABSORBED_QUERY_INITCHECK_ONLY") == "1"
-    attn = _attn(model_config, layer)
-    runner = _stress_runner(attn, "default")
-    hidden, _ = _inputs(layer, rows, seed=1200)
-    # Small positions so that a copy of the first RoPE-table rows suffices.
-    position_ids = torch.arange(rows, dtype=torch.int32).unsqueeze(0).cuda()
-    kva = attn.kv_a_proj_with_mqa(hidden)
-    fused = runner.indexer_fusable()
-    kw = torch.nn.functional.linear(hidden, attn.mqa.indexer._fused_wk_wp_weight) if fused else None
-    k_ln = attn.mqa.indexer.k_norm(kw[:, : attn.mqa.indexer.head_dim]) if fused else None
-    qn = runner.native_q_norm(kva)
-    reference = None
-    if not initcheck_only:
-        reference = (
-            runner.launch_indexer(kva, kw, k_ln, None, 0, position_ids, absorb)
-            if fused
-            else runner.launch(kva, need_qr=False, absorb=absorb)
-        )
-    torch.cuda.synchronize()
-
-    probe = _stress_runner(attn, "default")
-    probe.sync = _memcpy_copy(torch.zeros_like(probe.sync))
-    inputs = dict(kva=_memcpy_copy(kva), qn=_memcpy_copy(qn))
-    probe.native_q_norm = lambda _kva: inputs["qn"]
-    params = [
-        attn.q_b_proj.weight,
-        attn.k_b_proj_trans,
-        attn.kv_a_layernorm.weight,
-        attn.q_a_layernorm.weight,
-    ]
-    rope = attn.mqa.indexer.rotary_emb if fused else None
-    saved_table = rope.rotary_cos_sin if fused else None
-    if fused:
-        inputs.update(
-            kw=_memcpy_copy(kw),
-            k_ln=_memcpy_copy(k_ln),
-            pos=_memcpy_copy(position_ids),
-            rope=_memcpy_copy(saved_table[:rows].contiguous()),
-        )
-        indexer = attn.mqa.indexer
-        params += [indexer.wq_b.weight]
-    saved_params = [p.data for p in params]
-    try:
-        for p in params:
-            p.data = _memcpy_copy(p.data)
-        if fused:
-            rope.rotary_cos_sin = inputs["rope"]
-        torch.cuda.synchronize()
-        if fused:
-            result = probe.launch_indexer(
-                inputs["kva"], inputs["kw"], inputs["k_ln"], None, 0, inputs["pos"], absorb
-            )
-        else:
-            result = probe.launch(inputs["kva"], need_qr=False, absorb=absorb)
-        torch.cuda.synchronize()
-    finally:
-        for p, data in zip(params, saved_params):
-            p.data = data
-        if fused:
-            rope.rotary_cos_sin = saved_table
-    assert int(probe.sync.abs().sum()) == 0
-    if initcheck_only:
-        return
-    if fused:
-        _assert_same_packet(result, reference, "memcpy inputs")
-    else:
-        for name, x, y in zip(("q", "qr", "latent", "fused_q"), result, reference):
-            if x is None:
-                assert y is None
-                continue
-            if name == "fused_q":
-                x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
-            _assert_bitwise(f"memcpy inputs {name}", x, y)
-
-
-@triton.jit
-def _occupier_kernel(
-    a_ptr,
-    b_ptr,
-    out_ptr,
-    smid_ptr,
-    hold_ns,
-    ITERS: tl.constexpr,
-    BM: tl.constexpr,
-    BN: tl.constexpr,
-    BK: tl.constexpr,
-    NS: tl.constexpr,
-):
-    """Holds one SM per program for ``hold_ns``: its pipelined dot loop
-    needs ~192 KB of shared memory, so no absorbed_query program fits next to it."""
-    pid = tl.program_id(0)
-    tl.store(smid_ptr + pid, tl.extra.cuda.smid())
-    offs_m = tl.arange(0, BM)
-    offs_n = tl.arange(0, BN)
-    offs_k = tl.arange(0, BK)
-    acc = tl.zeros([BM, BN], dtype=tl.float32)
-    for k in tl.range(0, ITERS, num_stages=NS):
-        a = tl.load(a_ptr + offs_m[:, None] * (ITERS * BK) + (k * BK + offs_k)[None, :])
-        b = tl.load(b_ptr + (k * BK + offs_k)[:, None] * BN + offs_n[None, :])
-        acc = tl.dot(a, b, acc)
-    t0 = tl.extra.cuda.globaltimer()
-    while tl.extra.cuda.globaltimer() - t0 < hold_ns:
-        pass
-    tl.store(out_ptr + pid * BM * BN + offs_m[:, None] * BN + offs_n[None, :], acc)
-
-
-def _compiled_shared_bytes(jit_fn):
-    """Shared memory of the variants of a Triton kernel compiled so far."""
-    sizes = []
-    for caches in jit_fn.device_caches.values():
-        for kernel in caches[0].values():
-            sizes.append(kernel.metadata.shared)
-    return sizes
-
-
-def _launch_absorbed_query_kernel_only(runner, kva, kw, position_ids):
-    """The absorbed_query launch of ``project`` after its native bridges."""
-    if kw is not None:
-        k_ln = runner.mla.mqa.indexer.k_norm(kw[:, : runner.mla.mqa.indexer.head_dim])
-        return runner.launch_indexer(kva, kw, k_ln, None, 0, position_ids, True)
-    q, _, latent, fused_q = runner.launch(kva, need_qr=False, absorb=True)
-    return q, latent, fused_q
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("free_sms", [1, 2, 4, 8])
-@pytest.mark.parametrize("config_name", ["default", "separate_consumers", "legacy_own_head"])
-@torch.inference_mode()
-def test_progress_with_restricted_sms(model_config, layer, free_sms, config_name):
-    """A concurrent kernel holds all but ``free_sms`` SMs for 0.5 s (its
-    programs leave no room for an absorbed_query program on those SMs). The absorbed_query kernel
-    (preabsorbed form; on the full layer with the fused indexer tail) must
-    complete while the SMs are still held, with outputs equal to its
-    unconstrained launch and native bit for bit.
-    The legacy own-head-group layout is reported, not asserted: it needs
-    NOPE // BN co-resident programs of one head group."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        kernels,
-    )
-
-    attn = _attn(model_config, layer)
-    runner = _stress_runner(attn, config_name)
-    rows = 6
-    hidden, position_ids = _inputs(layer, rows, seed=900)
-    native_outs, native_fq = _native(attn, hidden, position_ids)
-    kva = attn.kv_a_proj_with_mqa(hidden)
-    kw = (
-        torch.nn.functional.linear(hidden, attn.mqa.indexer._fused_wk_wp_weight)
-        if runner.indexer_fusable()
-        else None
-    )
-    # Warm up (compile/load) before the timed launch; also the reference.
-    reference = _launch_absorbed_query_kernel_only(runner, kva, kw, position_ids)
-    torch.cuda.synchronize()
-
-    n_sm = torch.cuda.get_device_properties(0).multi_processor_count
-    n_occ = n_sm - free_sms
-    bm = bn = 128
-    bk, iters = 64, 32
-    a = torch.randn(bm, iters * bk, device="cuda").to(torch.bfloat16)
-    b = torch.randn(iters * bk, bn, device="cuda").to(torch.bfloat16)
-    out = torch.empty(n_occ, bm, bn, device="cuda")
-    smids = torch.full((n_occ,), -1, dtype=torch.int32, device="cuda")
-    occ_stream, absorbed_query_stream = torch.cuda.Stream(), torch.cuda.Stream()
-    ev = {k: torch.cuda.Event(enable_timing=True) for k in ("s", "o", "p")}
-    hold_ns = 500_000_000
-    torch.cuda.synchronize()
-    with torch.cuda.stream(occ_stream):
-        ev["s"].record()
-        occ = _occupier_kernel[(n_occ,)](
-            a, b, out, smids, hold_ns, iters, bm, bn, bk, 6, num_warps=4, num_stages=6
-        )
-        ev["o"].record()
-    time.sleep(0.05)  # the occupier is resident long before absorbed_query is launched
-    with torch.cuda.stream(absorbed_query_stream):
-        result = _launch_absorbed_query_kernel_only(runner, kva, kw, position_ids)
-        ev["p"].record()
-    torch.cuda.synchronize()
-    t_absorbed_query = ev["s"].elapsed_time(ev["p"])
-    t_occ = ev["s"].elapsed_time(ev["o"])
-    occ_smem = occ.metadata.shared
-    jit_fn = kernels.qproj_indexer_kernel if kw is not None else (kernels.qproj_kernel)
-    absorbed_query_smem = min(_compiled_shared_bytes(jit_fn) or [0])
-    held = int(torch.unique(smids).numel())
-    report = dict(
-        free_sms=free_sms,
-        held_sms=held,
-        occupier_smem=occ_smem,
-        absorbed_query_min_smem=absorbed_query_smem,
-        absorbed_query_done_ms=round(t_absorbed_query, 3),
-        occupier_done_ms=round(t_occ, 3),
-    )
-    print(f"layer={layer} config={config_name} {report}")
-    assert held == n_occ, report  # one occupier program per SM
-    assert occ_smem + absorbed_query_smem > 228 * 1024, report
-
-    if kw is not None:
-        _assert_same_packet(result, reference, "restricted")
-        _native_agreement(
-            result.proj_outputs, result.fused_q, native_outs, native_fq, attn.kv_lora_rank
-        )
-    else:
-        for name, x, y in zip(("q", "latent", "fused_q"), result, reference):
-            if name == "fused_q":
-                x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
-            _assert_bitwise(f"restricted {name}", x, y)
-        q, latent, fused_q = result
-        _assert_bitwise("q vs native", q, native_outs[0])
-        _assert_bitwise(
-            "absorbed vs native",
-            fused_q[..., : attn.kv_lora_rank],
-            native_fq[..., : attn.kv_lora_rank],
-        )
-        _assert_bitwise("latent vs native", latent, native_outs[3])
-    assert int(runner.sync.abs().sum()) == 0
-    if config_name != "legacy_own_head":
-        assert t_absorbed_query < t_occ - 100.0, (
-            f"absorbed_query did not complete while {n_occ} SMs were held: {report}"
-        )
-
-
-@torch.inference_mode()
-def test_eligibility_and_fallback(model_config):
-    """Unsupported calls return None before any GPU work."""
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    attn = _attn(model_config, absorbed_query_fixture.SHARED_LAYER)
-    ok, reason = query_proj.module_eligibility(attn)
-    assert ok, reason
-    gen = SimpleNamespace(num_contexts=0, num_generations=1)
-    ctx = SimpleNamespace(num_contexts=1, num_generations=0)
-    mixed = SimpleNamespace(num_contexts=1, num_generations=1)
-    hidden, position_ids = _inputs(absorbed_query_fixture.SHARED_LAYER, 6)
-    assert query_proj.call_eligibility(attn, hidden, gen)[0]
-    for meta in (None, ctx, mixed):
-        assert not query_proj.call_eligibility(attn, hidden, meta)[0]
-        attn._absorbed_query = None
-        assert query_proj.dsa_proj(attn, position_ids, hidden, meta) is None
-    big = torch.zeros(17, hidden.shape[1], dtype=hidden.dtype, device="cuda")
-    assert not query_proj.call_eligibility(attn, big, gen)[0]
-    assert not query_proj.call_eligibility(attn, hidden.float(), gen)[0]
-    # Draft (MTP) layers are outside the target-only domain.
-    saved = attn.layer_idx
-    try:
-        attn.layer_idx = attn.num_hidden_layers
-        assert not query_proj.module_eligibility(attn)[0]
-    finally:
-        attn.layer_idx = saved
-    attn._absorbed_query = None
-    packet = query_proj.dsa_proj(attn, position_ids, hidden, gen)
-    assert packet is not None and len(packet.proj_outputs) == 4
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py
deleted file mode 100644
index 37044e1331ac..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_draft_loop_integration.py
+++ /dev/null
@@ -1,1103 +0,0 @@
-# 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.
-"""draft_absorbed_query inside the MTP Eagle draft loop, through the target's attention and worker.
-
-Modules: the built-in ``DeepseekV32Attention`` (adopted as the target's ``GlmMLA``) of the
-full-indexer target
-layer 6 (stand-in for the target verify forward) and of the MTP layer 78,
-both with real GLM-5.2 weights at TP8 rank 0; FP8 KV cache (V1
-``DSACacheManager`` holding layers 6 and 78); synthetic activations. Two
-requests with a 2,600-token context are prefilled natively through both
-layers.
-
-One decode iteration runs as the GEN CUDA graph runs it, as a decode step:
-
-1. ``on_update_kv_lens()`` (the runner's ``_preprocess_inputs``);
-2. target verify: layer 6 on ``B x 6`` rows (MTP verify layout);
-3. ``_prepare_attn_metadata_for_spec_dec``, then the draft loop of the target's
-   adopted worker (``GlmMtpWorker._forward_linear_draft_loop`` on an ``MTPEagleWorker``),
-   then ``_restore_attn_metadata_from_spec_dec``. The loop is the real
-   one: it marks the draft loop and the accepted counts for the indexer,
-   calls ``set_skip_topk(i > 0)``, gathers each request's last accepted row,
-   advances positions, sets ``seq_lens`` to 1 and updates ``kv_lens_cuda``.
-   Its ``draft_model`` is a stub whose ``mtp_layers[0]`` runs the native
-   layer-78 attention (``MLA.forward``'s custom-op path, or ``forward_impl``)
-   on the rows the loop hands it and returns an attention-independent next
-   hidden state (a fixed column rotation), and whose ``shared_head`` returns
-   ``hidden[:, :1024]`` as logits. So draft step 0 runs layer 78 on
-   ``B x 6`` rows and steps 1-4 on ``B`` rows, exactly the draft-loop shapes.
-
-Asserted:
-
-* gating through the target's ``_absorbed_query_packet``: ``draft_absorbed_query`` engages only layer
-  78 (all five draft steps, rows ``[6B, B, B, B, B]``), ``absorbed_query`` only layer 6,
-  ``absorbed_query,draft_absorbed_query`` both, none of them nothing; layers that are not engaged
-  produce caches and attention inputs byte-identical to native;
-* indexer reuse semantics: in every run (native and draft_absorbed_query) draft step 0 scores
-  the top-k (``in_mtp_draft_loop`` set, ``indexer_skip_topk`` clear) and
-  steps 1-4 reuse it (both set), and the top-k of steps 1-4 equals, row for
-  row and in order, the step-0 top-k row of each request's last accepted
-  token;
-* cache writes: per draft step the indexer keys handed to
-  ``_update_k_cache`` (``k_fp8``, ``k_scale``) and, after the iteration,
-  every cache (latent and indexer-key caches of layers 6 and 78) are
-  bit-identical to native, and exactly the same slots changed;
-* MLA inputs per call (absorbed query, FP8 ``quant_q_buffer``, BMM scales)
-  and indexer inputs per call (FP8 query, head weights) bit-identical to
-  native; the top-k and sparse-index key sets of every call equal native's;
-* attention outputs per call within the native run-to-run noise bound (as
-  in the absorbed_query integration test);
-* determinism: a repeated draft_absorbed_query iteration reproduces caches and recorded inputs
-  bit for bit;
-* CUDA graph: the whole iteration (target layer + five draft steps) captured
-  in one graph, replayed 1, 1 and 20 times with fresh hidden states and
-  fresh accepted counts, reproduces the eager draft_absorbed_query iteration bit for bit and
-  compares with native as above; a captured native graph reproduces eager
-  native;
-* a negative control (zeroed absorbed query on layer 78) is detected.
-
-Every draft step launches the full fused-indexer kernel and writes every
-packet output, including the indexer query inputs that the reusing steps 1-4
-do not read.
-"""
-
-import inspect
-import os
-import sys
-from types import SimpleNamespace
-
-import absorbed_query_fixture
-import pytest
-import test_absorbed_query_attention_integration as aq_integration
-import torch
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
-
-import plan_testlib  # noqa: E402
-
-TARGET = absorbed_query_fixture.FULL_LAYER
-MTP = 78
-LAYERS = (TARGET, MTP)
-CTX = 2600
-MAX_REQS = 2
-DRAFT_LEN = 5
-VERIFY_ROWS = DRAFT_LEN + 1
-STUB_VOCAB = 1024
-# Native iterations per comparison. The output bound is twice the largest
-# native repeat difference, so it needs enough repeats to sample the native
-# sparse attention's run-to-run spread (graph replays included).
-NATIVE_RUNS = int(os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_IT_NATIVE_RUNS", "9"))
-# Accepted counts per request (including the target-sampled token, 1..6).
-ACCEPTED = {1: [3], 2: [6, 1]}
-
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-
-def _set_features(features: str, form: str = "auto"):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    query_proj.SERVED_FORM = form
-    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
-
-
-class _LaunchLog:
-    """Records ``(layer, rows)`` of every fused-indexer absorbed_query launch."""
-
-    def __enter__(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-            query_proj,
-        )
-
-        self.cls = query_proj.AbsorbedQueryProjection
-        self.orig = self.cls.launch_indexer
-        self.calls = []
-        log, orig = self.calls, self.orig
-
-        def launch_indexer(runner, kva, *args, **kwargs):
-            log.append((runner.mla.layer_idx, kva.shape[0]))
-            return orig(runner, kva, *args, **kwargs)
-
-        self.cls.launch_indexer = launch_indexer
-        return self
-
-    def __exit__(self, *exc):
-        self.cls.launch_indexer = self.orig
-        return False
-
-
-class _Recorder:
-    """Per-call records of what the native attention consumes, per layer:
-    lists indexed by call (one call for layer 6, five for layer 78).
-
-    Wraps ``mla_rope_generation`` (absorbed query, ``quant_q_buffer``, BMM
-    scales), ``sparse_attn_predict`` (converted indices), the indexer's
-    ``_update_k_cache`` (keys appended to the indexer cache) and
-    ``sparse_attn_indexer`` (logical top-k, indexer query inputs and the
-    draft-loop flags it sees). Inside a capture the clones are graph-pool
-    tensors that every replay refreshes; the flags are host values seen at
-    capture."""
-
-    def __init__(self, attn_by_layer):
-        self.attn = attn_by_layer
-        self.data = {layer: {} for layer in attn_by_layer}
-        self._installed = []
-
-    def _wrap(self, obj, name, post):
-        orig = getattr(obj, name)
-        sig = inspect.signature(orig)
-
-        def hook(*args, **kwargs):
-            out = orig(*args, **kwargs)
-            post(sig.bind(*args, **kwargs).arguments, out)
-            return out
-
-        setattr(obj, name, hook)
-        self._installed.append((obj, name))
-
-    def __enter__(self):
-        for layer, attn in self.attn.items():
-            rec = self.data[layer]
-
-            def add(key, value, rec=rec):
-                rec.setdefault(key, []).append(value)
-
-            def rope_post(a, _out, add=add, kv_lora=attn.kv_lora_rank):
-                add("absorbed_q", a["fused_q"][..., :kv_lora].clone())
-                for key, arg in (
-                    ("quant_q", "quant_q_buffer"),
-                    ("bmm1_scale", "mla_bmm1_scale"),
-                    ("bmm2_scale", "mla_bmm2_scale"),
-                ):
-                    if a.get(arg) is not None:
-                        add(key, a[arg].clone())
-
-            def predict_post(_a, out, add=add):
-                if isinstance(out, tuple) and out[0] is not None:
-                    add("sparse_indices", out[0].clone())
-
-            def kcache_post(a, _out, add=add):
-                add("index_k_fp8", a["k_fp8"].clone())
-                add("index_k_scale", a["k_scale"].clone())
-
-            def topk_post(a, out, add=add):
-                meta = a["metadata"]
-                add("topk", out.clone())
-                add("indexer_q_fp8", a["q_fp8"].clone())
-                add("indexer_weights", a["weights"].clone())
-                add("flags", (bool(meta.in_mtp_draft_loop), bool(meta.indexer_skip_topk)))
-
-            self._wrap(attn.mqa, "mla_rope_generation", rope_post)
-            self._wrap(attn.mqa, "sparse_attn_predict", predict_post)
-            self._wrap(attn.mqa.indexer, "_update_k_cache", kcache_post)
-            self._wrap(attn.mqa.indexer, "sparse_attn_indexer", topk_post)
-        return self
-
-    def __exit__(self, *exc):
-        for obj, name in self._installed:
-            delattr(obj, name)
-        self._installed.clear()
-        return False
-
-    def snapshot(self):
-        return {
-            layer: {
-                key: [v.clone() if torch.is_tensor(v) else v for v in vals]
-                for key, vals in rec.items()
-            }
-            for layer, rec in self.data.items()
-        }
-
-
-class _Engagement:
-    """Counts the absorbed_query package entry points the absorbed-query packet calls:
-    ``(entry, layer, rows, engaged)`` per call."""
-
-    def __enter__(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501, F401
-            absorbed_query,
-        )
-
-        self.registry = dispatch._REGISTRY["absorbed_query"]
-        self.saved = dict(self.registry)
-        self.calls = []
-        for name in ("dsa_proj", "mtp_dsa_proj"):
-            fn = self.saved[name]
-
-            def wrapped(mla, position_ids, hidden_states, attn_metadata, fn=fn, name=name):
-                packet = fn(mla, position_ids, hidden_states, attn_metadata)
-                self.calls.append((name, mla.layer_idx, hidden_states.shape[0], packet is not None))
-                return packet
-
-            self.registry[name] = wrapped
-        return self
-
-    def __exit__(self, *exc):
-        self.registry.update(self.saved)
-        return False
-
-
-class _StubShared:
-    def __call__(self, hidden_states, lm_head, attn_metadata, return_context_logits=False):
-        return hidden_states[:, :STUB_VOCAB].float()
-
-
-class _StubMTPLayer:
-    """``mtp_layers[0]``: the native layer-78 attention on the loop's rows;
-    the next hidden state does not depend on the attention output, so native
-    and draft_absorbed_query runs see identical inputs at every draft step."""
-
-    def __init__(self, harness, path, outputs):
-        self.h, self.path, self.outputs = harness, path, outputs
-        self.shared_head = _StubShared()
-
-    def __call__(
-        self,
-        input_ids,
-        position_ids,
-        hidden_states,
-        embed_tokens,
-        attn_metadata,
-        all_rank_num_tokens=None,
-        spec_metadata=None,
-        **kwargs,
-    ):
-        out = self.h.forward(MTP, position_ids, hidden_states, attn_metadata, self.path)
-        self.outputs.append(out.clone())
-        return torch.roll(hidden_states, shifts=97, dims=-1)
-
-
-class _StubDraftModel:
-    def __init__(self, mtp_layer):
-        self.mtp_layers = [mtp_layer]
-        self.embed_tokens = None
-        self.lm_head = None
-
-
-class _Harness:
-    def __init__(self):
-        import tensorrt_llm.bindings
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
-            GlmMtpWorker,
-        )
-        from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
-        from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
-        from tensorrt_llm._torch.metadata import KVCacheParams
-        from tensorrt_llm._torch.model_config import ModelConfig
-        from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-        from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
-        from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig
-        from tensorrt_llm.mapping import Mapping
-        from tensorrt_llm.models.modeling_utils import QuantConfig
-        from tensorrt_llm.quantization.mode import QuantAlgo
-
-        self.KVCacheParams = KVCacheParams
-        self.model_config = absorbed_query_fixture.build_model_config(0)
-        pretrained = self.model_config.pretrained_config
-        sparse = self.model_config.sparse_attention_config
-        self.attn = {
-            layer: absorbed_query_fixture.build_native_attention(layer, 0, self.model_config)
-            for layer in LAYERS
-        }
-        assert self.attn[MTP].mqa.indexer.mtp_index_share
-        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
-        kv_model_config = ModelConfig(
-            mapping=self.mapping,
-            sparse_attention_config=sparse,
-            pretrained_config=pretrained,
-            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8),
-        )
-        self.kvm = DSACacheManager(
-            KvCacheConfig(
-                dtype="fp8",
-                max_tokens=MAX_REQS * (CTX + 128),
-                enable_block_reuse=False,
-                event_buffer_max_size=0,
-            ),
-            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-            num_layers=len(LAYERS),
-            layer_mask=[i in LAYERS for i in range(MTP + 1)],
-            num_kv_heads=1,
-            head_dim=576,
-            tokens_per_block=64,
-            max_seq_len=CTX + 128,
-            max_batch_size=MAX_REQS,
-            mapping=self.mapping,
-            dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
-            sparse_attention_config=sparse,
-            model_config=kv_model_config,
-            pretrained_config=pretrained,
-            vocab_size=pretrained.vocab_size,
-            max_num_tokens=MAX_REQS * CTX,
-        )
-        self.req_ids = list(range(MAX_REQS))
-        for rid in self.req_ids:
-            # CTX + 10 tokens (6 verify slots, up to 4 more draft slots) fit
-            # the blocks of CTX + 1 tokens (64 per block).
-            self.kvm.add_dummy_requests(
-                request_ids=[rid], token_nums=[CTX + 1], is_gen=False, prepare_resource=True
-            )
-        sparse_params = sparse.to_sparse_params(pretrained_config=pretrained)
-        self.Metadata = get_attention_backend("TRTLLM", sparse_params=sparse_params).Metadata
-        self.meta_params = sparse.to_sparse_metadata_params(pretrained_config=pretrained)
-        assert self.meta_params.mtp_index_share
-        self.gamma = {layer: absorbed_query_fixture.load_input_layernorm(layer) for layer in LAYERS}
-        self.worker = GlmMtpWorker.adopt(
-            MTPEagleWorker(MTPDecodingConfig(max_draft_len=5), mapping=self.mapping)
-        )
-        assert self.worker.is_mtp_eagle
-        self.batch_indices = torch.arange(MAX_REQS, dtype=torch.int32, device="cuda")
-
-        # Native prefill of both requests through both layers.
-        _set_features("")
-        hidden = absorbed_query_fixture.make_hidden(
-            MAX_REQS * CTX, seed=11, gamma=self.gamma[TARGET]
-        )
-        positions = (
-            torch.arange(CTX, dtype=torch.int32, device="cuda").repeat(MAX_REQS).unsqueeze(0)
-        )
-        ctx_meta = self.Metadata(
-            seq_lens=torch.tensor([CTX] * MAX_REQS, dtype=torch.int),
-            request_ids=self.req_ids,
-            max_num_requests=MAX_REQS,
-            num_contexts=MAX_REQS,
-            prompt_lens=[CTX] * MAX_REQS,
-            max_num_tokens=MAX_REQS * CTX,
-            kv_cache_manager=self.kvm,
-            kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0] * MAX_REQS),
-            mapping=self.mapping,
-            sparse_metadata_params=self.meta_params,
-        )
-        ctx_meta.prepare()
-        ctx_meta.on_update_kv_lens()
-        for layer in LAYERS:
-            self.forward(layer, positions, hidden, ctx_meta, "forward_impl")
-        torch.cuda.synchronize()
-
-    # -- inputs and metadata -----------------------------------------------
-
-    def iteration_inputs(self, nreq, seed, accepted=None):
-        rows = nreq * VERIFY_ROWS
-        accepted = ACCEPTED[nreq] if accepted is None else accepted
-        return SimpleNamespace(
-            hidden_t=absorbed_query_fixture.make_hidden(rows, seed=seed, gamma=self.gamma[TARGET]),
-            pos_t=(CTX + torch.arange(VERIFY_ROWS, dtype=torch.int32, device="cuda"))
-            .repeat(nreq)
-            .unsqueeze(0),
-            hidden_d=absorbed_query_fixture.make_hidden(rows, seed=seed + 1, gamma=self.gamma[MTP]),
-            num_accepted=torch.tensor(accepted, dtype=torch.int32, device="cuda"),
-            input_ids=torch.arange(rows, dtype=torch.int32, device="cuda"),
-        )
-
-    def gen_metadata(self, nreq, cuda_graph=False):
-        """Verify-step metadata as the runner builds it (see the absorbed_query
-        integration test): ``nreq`` requests x 6 tokens, MTP verify width
-        from ``update_spec_dec_param``, then ``prepare()``."""
-        req_ids = self.req_ids[:nreq]
-        kv_params = self.KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[CTX] * nreq)
-        meta = self.Metadata(
-            seq_lens=torch.tensor([VERIFY_ROWS] * nreq, dtype=torch.int),
-            request_ids=req_ids,
-            max_num_requests=nreq,
-            num_contexts=0,
-            prompt_lens=[CTX] * nreq,
-            max_num_tokens=nreq * VERIFY_ROWS,
-            kv_cache_manager=self.kvm,
-            kv_cache_params=kv_params,
-            mapping=self.mapping,
-            sparse_metadata_params=self.meta_params,
-        )
-        if cuda_graph:
-            from tensorrt_llm._torch.memory_buffer_utils import Buffers
-
-            meta = meta.create_cuda_graph_metadata(nreq, False, DRAFT_LEN, Buffers())
-            assert meta.is_cuda_graph
-            meta.request_ids = req_ids
-            meta.prompt_lens = [CTX] * nreq
-            meta.num_contexts = 0
-            meta.kv_cache_params = kv_params
-        meta.update_spec_dec_param(
-            batch_size=nreq,
-            is_spec_decoding_enabled=True,
-            is_spec_dec_tree=False,
-            is_spec_dec_dynamic_tree=False,
-            max_draft_len=DRAFT_LEN,
-            max_total_draft_tokens=DRAFT_LEN,
-        )
-        meta.prepare()
-        return meta
-
-    # -- forwards ----------------------------------------------------------
-
-    def forward(self, layer, positions, hidden, meta, path):
-        attn = self.attn[layer]
-        if path == "forward_impl":
-            out = attn._create_outputs(hidden, meta)
-            attn.forward_impl(positions, hidden, meta, attn_output=out)
-            return out[0]
-        import weakref
-
-        from tensorrt_llm._torch.utils import model_extra_attrs
-
-        extra = self.model_config.extra_attrs
-        extra["attention_metadata"] = weakref.ref(meta)
-        with model_extra_attrs(extra):
-            out = [torch.ops.trtllm.create_mla_outputs(hidden, attn.layer_idx_str)]
-            attn._forward_custom_op(hidden, positions, out, None)
-        return out[0]
-
-    def iteration(self, meta, it, path):
-        """Target verify (layer 6) and the draft loop (layer 78), as the GEN graph
-        runs one decode iteration (a decode step). Returns the attention
-        outputs per layer and call."""
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-            dispatch,
-        )
-
-        with dispatch.step(dispatch.Phase.DECODE):
-            return self._iteration(meta, it, path)
-
-    def _iteration(self, meta, it, path):
-        nreq = it.num_accepted.shape[0]
-        meta.on_update_kv_lens()
-        outs = {TARGET: [], MTP: []}
-        outs[TARGET].append(self.forward(TARGET, it.pos_t, it.hidden_t, meta, path).clone())
-        draft_model = _StubDraftModel(_StubMTPLayer(self, path, outs[MTP]))
-        spec_metadata = SimpleNamespace(
-            runtime_draft_len=DRAFT_LEN,
-            batch_indices_cuda=self.batch_indices[:nreq],
-            all_rank_num_tokens=None,
-            subseq_all_rank_num_tokens=None,
-            wants_advanced_draft_sampling=False,
-            max_num_requests=nreq,
-        )
-        self.worker._prepare_attn_metadata_for_spec_dec(meta)
-        try:
-            inputs = {
-                "input_ids": it.input_ids,
-                "position_ids": it.pos_t.squeeze(0),
-                "hidden_states": it.hidden_d,
-                "attn_metadata": meta,
-                "spec_metadata": spec_metadata,
-            }
-            outs["tokens"] = self.worker._forward_linear_draft_loop(
-                inputs, meta, spec_metadata, draft_model, None, 0, nreq, it.num_accepted, None
-            )
-        finally:
-            self.worker._restore_attn_metadata_from_spec_dec(meta)
-            meta.use_spec_decoding = True
-        return outs
-
-    def run(self, nreq, it, path, meta=None):
-        """One eager iteration with the current features: outputs, records,
-        entry-point engagement."""
-        meta = meta if meta is not None else self.gen_metadata(nreq)
-        with _Recorder(self.attn) as rec, _Engagement() as eng:
-            outs = self.iteration(meta, it, path)
-        torch.cuda.synchronize()
-        return outs, rec.snapshot(), eng.calls
-
-    # -- caches ------------------------------------------------------------
-
-    def caches(self):
-        bufs = {}
-        for layer in LAYERS:
-            bufs[f"latent{layer}"] = self.kvm.get_buffers(layer)
-            bufs[f"index_k{layer}"] = self.kvm.get_indexer_k_cache_buffers(layer)
-        return bufs
-
-    def snapshot(self):
-        return {k: v.clone() for k, v in self.caches().items()}
-
-    def restore(self, snap):
-        for k, v in self.caches().items():
-            v.copy_(snap[k])
-        torch.cuda.synchronize()
-
-    def reset_runners(self):
-        for attn in self.attn.values():
-            attn._absorbed_query = None
-            attn._draft_absorbed_query = None
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = _Harness()
-    yield h
-    _set_features("")
-
-
-# -- comparison helpers -----------------------------------------------------
-
-
-def _expected_slots(nreq, accepted):
-    """Changed token slots per layer after one iteration: six verify slots
-    per request on both layers; on the MTP layer steps 1-4 append at
-    positions ``accepted .. accepted + 3`` after the context."""
-    mtp = sum(len(set(range(VERIFY_ROWS)) | set(range(a, a + DRAFT_LEN - 1))) for a in accepted)
-    return {TARGET: nreq * VERIFY_ROWS, MTP: mtp}
-
-
-def _compare_caches(snap, ref_cache, cand_cache, expected, report, tag):
-    """Changed slots equal the reference's (and the expected count) and
-    every byte equals the reference's."""
-    fails = []
-    for key in cand_cache:
-        ref, cand = ref_cache[key], cand_cache[key]
-        layer = int(key.replace("latent", "").replace("index_k", ""))
-        n_diff = int(aq_integration._changed_bytes(ref, cand).sum())
-        report[f"{tag}:{key}_bytes_differing"] = n_diff
-        if key.startswith("latent"):
-            slots = aq_integration._changed_latent_slots(snap[key], cand)
-            ref_slots = aq_integration._changed_latent_slots(snap[key], ref)
-        else:
-            slots = aq_integration._changed_index_k_tokens(snap[key], cand)
-            ref_slots = aq_integration._changed_index_k_tokens(snap[key], ref)
-        if slots != ref_slots or len(ref_slots) != expected[layer]:
-            fails.append(
-                f"{tag}: {key} changed {len(slots)} slots vs "
-                f"reference {len(ref_slots)} (expected "
-                f"{expected[layer]})"
-            )
-        if n_diff:
-            fails.append(f"{tag}: {key} differs in {n_diff} bytes")
-    return fails
-
-
-def _check_reuse(rec, accepted, report, tag):
-    """Indexer reuse semantics of one run: layer 6 outside the draft loop;
-    layer 78 scores at step 0 and reuses at steps 1-4, the reused rows being
-    the step-0 rows of each request's last accepted token, in order."""
-    fails = []
-    if rec[TARGET]["flags"] != [(False, False)]:
-        fails.append(f"{tag}: target flags {rec[TARGET]['flags']}")
-    flags = rec[MTP]["flags"]
-    want = [(True, False)] + [(True, True)] * (DRAFT_LEN - 1)
-    if flags != want:
-        fails.append(f"{tag}: draft-loop flags {flags}, expected {want}")
-    topk = rec[MTP]["topk"]
-    if len(topk) != DRAFT_LEN:
-        return fails + [f"{tag}: {len(topk)} layer-78 top-k calls"]
-    rows = torch.tensor(
-        [r * VERIFY_ROWS + min(max(a - 1, 0), DRAFT_LEN) for r, a in enumerate(accepted)],
-        device=topk[0].device,
-    )
-    stash = topk[0][rows]
-    for step in range(1, DRAFT_LEN):
-        if not torch.equal(topk[step], stash):
-            fails.append(f"{tag}: step {step} top-k is not the step-0 last-accepted row")
-    report[f"{tag}:reuse_checked_steps"] = DRAFT_LEN - 1
-    return fails
-
-
-def _same_sets(a, b):
-    return torch.equal(a.sort(dim=-1).values, b.sort(dim=-1).values)
-
-
-def _compare_records(ref, cand, report, tag, layers=LAYERS):
-    """Per call: MLA inputs (absorbed query, FP8 query, BMM scales), indexer
-    inputs (FP8 query, head weights) and appended indexer keys bitwise; top-k
-    and sparse-index key sets equal."""
-    fails = []
-    for layer in layers:
-        r, c = ref[layer], cand[layer]
-        if set(r) != set(c) or "absorbed_q" not in r:
-            return [f"{tag}: layer {layer} recorded {sorted(c)} vs {sorted(r)}"]
-        for key in sorted(r):
-            if len(r[key]) != len(c[key]):
-                fails.append(f"{tag}: layer {layer} {key}: {len(c[key])} calls vs {len(r[key])}")
-        if fails:
-            return fails
-        for i in range(len(r["absorbed_q"])):
-            name = f"{tag}:L{layer}.call{i}"
-            for key in (
-                "absorbed_q",
-                "quant_q",
-                "bmm1_scale",
-                "bmm2_scale",
-                "indexer_q_fp8",
-                "indexer_weights",
-                "index_k_fp8",
-                "index_k_scale",
-            ):
-                if key not in r:
-                    continue
-                n_diff = int(
-                    (aq_integration._bytes(r[key][i]) != aq_integration._bytes(c[key][i])).sum()
-                )
-                report[f"{name}.{key}_bytes_differing"] = n_diff
-                if n_diff:
-                    fails.append(f"{name}: {key} differs in {n_diff} bytes")
-            for key in ("topk", "sparse_indices"):
-                same = _same_sets(r[key][i], c[key][i])
-                report[f"{name}.{key}_same_sets"] = same
-                if not same:
-                    fails.append(f"{name}: {key} selects different keys")
-    return fails
-
-
-def _output_bounds(native_runs):
-    """Per (layer, call) bounds from repeated native iterations."""
-    ref = native_runs[0]
-    bounds = {}
-    for layer in LAYERS:
-        for i in range(len(ref[layer])):
-            base = ref[layer][i].float()
-            noise = max(
-                [(run[layer][i].float() - base).abs().max().item() for run in native_runs[1:]]
-                + [0.0]
-            )
-            rel = max(
-                [aq_integration._rel_rms(run[layer][i], ref[layer][i]) for run in native_runs[1:]]
-                + [0.0]
-            )
-            ulp = aq_integration._bf16_ulp(base.abs().max().item())
-            bounds[(layer, i)] = (
-                2.0 * max(noise, ulp),
-                max(2.0 * rel, aq_integration.REL_RMS_FLOOR),
-                noise,
-            )
-    return bounds
-
-
-def _compare_outputs(ref, cand, bounds, report, tag):
-    fails = []
-    for (layer, i), (max_bound, rel_bound, noise) in bounds.items():
-        diff = (cand[layer][i].float() - ref[layer][i].float()).abs().max().item()
-        rel = aq_integration._rel_rms(cand[layer][i], ref[layer][i])
-        report[f"{tag}:L{layer}.call{i}.out_max_diff"] = diff
-        report[f"{tag}:L{layer}.call{i}.out_rel_rms"] = rel
-        report[f"L{layer}.call{i}.native_repeat_max_diff"] = noise
-        if diff > max_bound or rel > rel_bound:
-            fails.append(
-                f"{tag}: L{layer} call {i} output max diff "
-                f"{diff:.3e} (bound {max_bound:.3e}), relRMS "
-                f"{rel:.3e} (bound {rel_bound:.3e})"
-            )
-    return fails
-
-
-def _native_runs(h, nreq, it, path, snap, repeats=NATIVE_RUNS):
-    """Repeated native iterations from the same snapshot; the native
-    iteration must itself be repeatable in caches and recorded inputs."""
-    _set_features("")
-    h.reset_runners()
-    outs, recs, cache, fails = [], [], None, []
-    accepted = it.num_accepted.tolist()
-    expected = _expected_slots(nreq, accepted)
-    for i in range(repeats):
-        h.restore(snap)
-        out, rec, calls = h.run(nreq, it, "custom_op" if path is None else path)
-        assert calls == [], f"native run called absorbed_query entry points: {calls}"
-        outs.append(out)
-        recs.append(rec)
-        report = {}
-        fails += _check_reuse(rec, accepted, report, f"native{i}")
-        if i == 0:
-            cache = h.snapshot()
-        else:
-            fails += _compare_caches(
-                snap, cache, h.snapshot(), expected, report, f"native_repeat{i}"
-            )
-            fails += _compare_records(recs[0], rec, report, f"native_repeat{i}")
-    h.restore(snap)
-    assert not fails, f"native iteration is not repeatable: {fails}"
-    return outs, recs, cache
-
-
-def _expected_calls(features, nreq):
-    """Entry-point calls the absorbed-query packet makes in one iteration."""
-    selected = set(features.split(","))
-    target_query, draft_query = "absorbed_query" in selected, "draft_absorbed_query" in selected
-    target_rows, steps = nreq * VERIFY_ROWS, [nreq * VERIFY_ROWS] + [nreq] * (DRAFT_LEN - 1)
-    calls = []
-    if target_query:
-        calls.append(("dsa_proj", TARGET, target_rows, True))
-    elif draft_query:
-        calls.append(("mtp_dsa_proj", TARGET, target_rows, False))
-    for rows in steps:
-        if target_query:
-            calls.append(("dsa_proj", MTP, rows, False))
-        if draft_query:
-            calls.append(("mtp_dsa_proj", MTP, rows, True))
-    return calls
-
-
-def _compare_to_native(snap, native, cand, accepted, nreq, bounds, report, tag):
-    """Caches and recorded inputs bit-identical to native, the indexer reuse
-    semantics, and outputs within the native run-to-run bound."""
-    native_out, native_rec, native_cache = native
-    out, rec, cache = cand
-    fails = _compare_caches(snap, native_cache, cache, _expected_slots(nreq, accepted), report, tag)
-    fails += _compare_records(native_rec, rec, report, tag)
-    fails += _check_reuse(rec, accepted, report, tag)
-    fails += _compare_outputs(native_out, out, bounds, report, tag)
-    return fails
-
-
-def _expected_launches(nreq):
-    """Fused-indexer launches of layer 78 in one iteration (draft_absorbed_query only)."""
-    return [(MTP, nreq * VERIFY_ROWS)] + [(MTP, nreq)] * (DRAFT_LEN - 1)
-
-
-# -- tests ------------------------------------------------------------------
-
-FEATURE_SETS = {
-    "none": "",
-    "absorbed_query": "absorbed_query",
-    "draft_absorbed_query": "draft_absorbed_query",
-    "absorbed_query_draft_query": "absorbed_query,draft_absorbed_query",
-}
-
-
-@pytest.mark.parametrize("features", list(FEATURE_SETS))
-@pytest.mark.parametrize("nreq", [1, 2])
-@torch.inference_mode()
-def test_feature_gating_through_adapter(harness, features, nreq):
-    """Which layers each selection engages, through the absorbed-query packet,
-    and native agreement of the whole iteration."""
-    h = harness
-    sel = FEATURE_SETS[features]
-    it = h.iteration_inputs(nreq, seed=12)
-    accepted = it.num_accepted.tolist()
-    snap = h.snapshot()
-    try:
-        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, "custom_op", snap)
-        bounds = _output_bounds(native_outs)
-        _set_features(sel)
-        h.reset_runners()
-        h.restore(snap)
-        out, rec, calls = h.run(nreq, it, "custom_op")
-        cache = h.snapshot()
-        expected = _expected_calls(sel, nreq)
-        assert calls == expected, f"engagement {calls} != {expected}"
-        engaged = {layer for _, layer, _, ok in calls if ok}
-        report = {}
-        fails = _compare_to_native(
-            snap,
-            (native_outs[0], native_recs[0], native_cache),
-            (out, rec, cache),
-            accepted,
-            nreq,
-            bounds,
-            report,
-            f"{features}",
-        )
-        print(f"features={sel!r} nreq={nreq} engaged={sorted(engaged)} {report}")
-        assert not fails, fails
-    finally:
-        h.restore(snap)
-        _set_features("")
-        h.reset_runners()
-
-
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("form", ["auto", "preabsorbed", "native_q"])
-@pytest.mark.parametrize("nreq", [1, 2])
-@torch.inference_mode()
-def test_draft_query_draft_loop_matches_native(harness, nreq, form, path):
-    """draft_absorbed_query alone (layer 78 on all five draft steps) against native: caches,
-    MLA inputs, indexer keys and reuse, outputs; then determinism (the same
-    iteration again reproduces caches and records bit for bit)."""
-    h = harness
-    it = h.iteration_inputs(nreq, seed=20 + nreq)
-    accepted = it.num_accepted.tolist()
-    snap = h.snapshot()
-    try:
-        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, path, snap)
-        bounds = _output_bounds(native_outs)
-        _set_features("draft_absorbed_query", form)
-        h.reset_runners()
-        h.restore(snap)
-        with _LaunchLog() as launches:
-            out, rec, calls = h.run(nreq, it, path)
-        cache = h.snapshot()
-        assert calls == _expected_calls("draft_absorbed_query", nreq), calls
-        assert launches.calls == _expected_launches(nreq), launches.calls
-        report = {}
-        fails = _compare_to_native(
-            snap,
-            (native_outs[0], native_recs[0], native_cache),
-            (out, rec, cache),
-            accepted,
-            nreq,
-            bounds,
-            report,
-            "draft_absorbed_query",
-        )
-        h.restore(snap)
-        _, rec_again, _ = h.run(nreq, it, path)
-        fails += _compare_caches(
-            snap, cache, h.snapshot(), _expected_slots(nreq, accepted), report, "draft_query_again"
-        )
-        fails += _compare_records(rec, rec_again, report, "draft_query_again")
-        print(f"nreq={nreq} form={form} path={path} {report}")
-        assert not fails, fails
-    finally:
-        h.restore(snap)
-        _set_features("")
-        h.reset_runners()
-
-
-GRAPH_SETS = {
-    "w5only": ("draft_absorbed_query",),
-    "nativeonly": ("native",),
-    "both": ("draft_absorbed_query", "native"),
-}
-# Diagnostics: restrict the graph cases (comma lists).
-GRAPH_NREQ = [
-    int(x)
-    for x in os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_IT_GRAPH_NREQ", "1,2").split(",")
-]
-GRAPH_SET_NAMES = os.environ.get(
-    "GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_IT_GRAPH_SETS", ",".join(GRAPH_SETS)
-).split(",")
-
-
-def _tensor_ptrs(meta):
-    """data_ptr of every CUDA tensor attribute of the metadata (what a
-    captured graph may have baked in), as the graph runner's strict buffer
-    check records them."""
-    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
-
-    return CUDAGraphRunner._snapshot_graph_tensor_ptrs(meta)
-
-
-def _moved(before, after):
-    return sorted(k for k, v in before.items() if after.get(k) != v)
-
-
-@pytest.mark.parametrize("graph_set", GRAPH_SET_NAMES)
-@pytest.mark.parametrize("form", ["auto", "native_q"])
-@pytest.mark.parametrize("nreq", GRAPH_NREQ)
-@torch.inference_mode()
-def test_draft_query_draft_loop_cuda_graph(harness, nreq, form, graph_set):
-    """The whole iteration (target layer + native draft loop over layer 78)
-    captured in one CUDA graph, as the GEN graph runner does: warm-up of
-    every configuration first, then capture on ``torch.cuda.graph``'s stream
-    into one shared pool. Replays 1, 1 and 20 times with fresh hidden states
-    and fresh accepted counts must reproduce the eager iteration of the same
-    selection bit for bit (caches, records) and compare with eager native as
-    the eager test does."""
-    h = harness
-    graphs = {}
-    snap = h.snapshot()
-    try:
-        _graph_checks(h, nreq, form, snap, graphs, GRAPH_SETS[graph_set])
-    except BaseException as exc:
-        print(f"graph test failure: {exc!r}", flush=True)
-        raise
-    finally:
-        graphs.clear()
-        h.reset_runners()
-        h.restore(snap)
-        _set_features("")
-
-
-def _graph_checks(h, nreq, form, snap, graphs, wanted):
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    path = "custom_op"
-    features = {"draft_absorbed_query": "draft_absorbed_query", "native": ""}
-    configs = [name for name in ("draft_absorbed_query", "native") if name in wanted]
-    warm = {}
-    for name in configs:
-        _set_features(features[name], form)
-        h.reset_runners()
-        meta = h.gen_metadata(nreq, cuda_graph=True)
-        static = h.iteration_inputs(nreq, seed=200)
-        with with_multi_stream(True):
-            for _ in range(2):
-                # Host-side per-step preparation, as before every replay
-                # (the draft loop moved kv_lens_cuda on the device).
-                meta.prepare()
-                h.iteration(meta, static, path)
-        torch.cuda.synchronize()
-        runners = {
-            layer: (h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query)
-            for layer in LAYERS
-        }
-        if name == "draft_absorbed_query":
-            assert runners[MTP][1] not in (None, False), "draft_absorbed_query did not engage"
-        warm[name] = (meta, static, runners)
-    h.restore(snap)
-    pool = None
-    for name in configs:
-        meta, static, runners = warm[name]
-        _set_features(features[name], form)
-        for layer, (r_absorbed_query, r_draft_query) in runners.items():
-            h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query = (
-                r_absorbed_query,
-                r_draft_query,
-            )
-        graph = torch.cuda.CUDAGraph()
-        ptrs_before = _tensor_ptrs(meta)
-        with (
-            with_multi_stream(True),
-            _Recorder(h.attn) as rec,
-            _Engagement() as eng,
-            _LaunchLog() as launches,
-        ):
-            with torch.cuda.graph(graph, pool=pool):
-                outs = h.iteration(meta, static, path)
-        pool = graph.pool()
-        ptrs_captured = _tensor_ptrs(meta)
-        print(
-            f"[graph test] {name}: metadata tensors rebound during "
-            f"capture: {_moved(ptrs_before, ptrs_captured)}; new: "
-            f"{sorted(set(ptrs_captured) - set(ptrs_before))}",
-            flush=True,
-        )
-        assert eng.calls == _expected_calls(features[name], nreq), eng.calls
-        if name == "draft_absorbed_query":
-            assert launches.calls == _expected_launches(nreq), launches.calls
-        graphs[name] = (graph, meta, static, outs, rec, runners, ptrs_captured)
-        torch.cuda.synchronize()
-        print(f"[graph test] captured {name}", flush=True)
-    h.reset_runners()
-
-    report, fails = {}, []
-    trials = [(301, 1, None), (302, 1, [1] * nreq), (303, 20, [VERIFY_ROWS] * nreq)]
-    for trial, (seed, replays, accepted) in enumerate(trials):
-        it = h.iteration_inputs(nreq, seed=seed, accepted=accepted)
-        accepted = it.num_accepted.tolist()
-        expected = _expected_slots(nreq, accepted)
-        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, path, snap)
-        bounds = _output_bounds(native_outs)
-        _set_features("draft_absorbed_query", form)
-        h.reset_runners()
-        h.restore(snap)
-        draft_query_eager = h.run(nreq, it, path)
-        draft_query_eager_cache = h.snapshot()
-        _set_features("")
-        h.reset_runners()
-        tag = f"t{trial}x{replays}"
-        for name, (graph, meta, static, outs, rec, runners, ptrs_captured) in graphs.items():
-            h.restore(snap)
-            meta.prepare()
-            moved = _moved(ptrs_captured, _tensor_ptrs(meta))
-            print(
-                f"[graph test] trial {trial} {name}: metadata tensors "
-                f"rebound by prepare() since capture: {moved}",
-                flush=True,
-            )
-            for key in ("hidden_t", "pos_t", "hidden_d", "num_accepted", "input_ids"):
-                getattr(static, key).copy_(getattr(it, key))
-            for i in range(replays):
-                if i:
-                    # The captured draft loop advances kv_lens_cuda on the
-                    # device; like the runner before every step, rebuild the
-                    # step's metadata from the host first (the replays
-                    # rewrite the same slots with the same bytes).
-                    meta.prepare()
-                graph.replay()
-            torch.cuda.synchronize()
-            print(f"[graph test] trial {trial} {name}: replayed x{replays}", flush=True)
-            g_out = {layer: [o.clone() for o in outs[layer]] for layer in LAYERS}
-            g_rec, g_cache = rec.snapshot(), h.snapshot()
-            fails += _check_reuse(g_rec, accepted, report, f"{tag}:{name}")
-            if name == "draft_absorbed_query":
-                fails += _compare_caches(
-                    snap,
-                    draft_query_eager_cache,
-                    g_cache,
-                    expected,
-                    report,
-                    f"{tag}:draft_query_graph_vs_draft_query_eager",
-                )
-                fails += _compare_records(
-                    draft_query_eager[1],
-                    g_rec,
-                    report,
-                    f"{tag}:draft_query_graph_vs_draft_query_eager",
-                )
-                fails += _compare_to_native(
-                    snap,
-                    (native_outs[0], native_recs[0], native_cache),
-                    (g_out, g_rec, g_cache),
-                    accepted,
-                    nreq,
-                    bounds,
-                    report,
-                    f"{tag}:draft_query_graph_vs_native",
-                )
-            else:
-                fails += _compare_caches(
-                    snap, native_cache, g_cache, expected, report, f"{tag}:native_graph_vs_eager"
-                )
-                fails += _compare_records(
-                    native_recs[0], g_rec, report, f"{tag}:native_graph_vs_eager"
-                )
-                fails += _compare_outputs(
-                    native_outs[0], g_out, bounds, report, f"{tag}:native_graph_vs_eager"
-                )
-            for r_absorbed_query, r_draft_query in runners.values():
-                for runner in (r_absorbed_query, r_draft_query):
-                    if runner not in (None, False):
-                        assert int(runner.sync.abs().sum()) == 0
-    print(f"graph nreq={nreq} form={form} graphs={wanted} {report}")
-    assert not fails, fails
-
-
-@torch.inference_mode()
-def test_checks_detect_a_wrong_mtp_query(harness, monkeypatch):
-    """Negative control: draft_absorbed_query packets whose absorbed query is zeroed must fail
-    the absorbed-query / quant_q checks on layer 78."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    h = harness
-    nreq, path = 1, "custom_op"
-    it = h.iteration_inputs(nreq, seed=12)
-    snap = h.snapshot()
-    orig_project = query_proj.AbsorbedQueryProjection.project
-
-    def broken_project(self, *args, **kwargs):
-        packet = orig_project(self, *args, **kwargs)
-        if packet.fused_q is not None:
-            packet.fused_q[..., : self.mla.kv_lora_rank].zero_()
-        return packet
-
-    try:
-        native_outs, native_recs, native_cache = _native_runs(h, nreq, it, path, snap)
-        monkeypatch.setattr(query_proj.AbsorbedQueryProjection, "project", broken_project)
-        _set_features("draft_absorbed_query", "preabsorbed")
-        h.reset_runners()
-        h.restore(snap)
-        _, bad_rec, calls = h.run(nreq, it, path)
-        assert calls == _expected_calls("draft_absorbed_query", nreq)
-        report = {}
-        fails = _compare_records(native_recs[0], bad_rec, report, "bad")
-        print(f"negative control {fails}")
-        for step in range(DRAFT_LEN):
-            assert any(f"L{MTP}.call{step}: absorbed_q" in f for f in fails), fails
-            assert any(f"L{MTP}.call{step}: quant_q" in f for f in fails), fails
-        assert not any(f"L{TARGET}." in f for f in fails), fails
-    finally:
-        monkeypatch.undo()
-        h.restore(snap)
-        _set_features("")
-        h.reset_runners()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py
deleted file mode 100644
index 7c7084e0eeb1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_fused_dsa_metadata.py
+++ /dev/null
@@ -1,745 +0,0 @@
-# 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.
-"""draft_absorbed_query in the MTP draft loop on the DSA decode-metadata path of the served env.
-
-The served GEN environment exports ``TRTLLM_FUSED_DSA_METADATA=1``. TensorRT-LLM
-does not read that name. The fused DSA decode-metadata kernel
-(``dsa/kernels.py::fused_dsa_decode_metadata``) is on by default and is
-turned off only by ``TRTLLM_DISABLE_FUSED_DSA_METADATA``
-(``dsa/metadata.py::_fused_dsa_meta_enabled``, cached per process). On a
-pure-decode step, ``DSAtrtllmAttentionMetadata.on_update_kv_lens`` then fills
-the request index per token, the indexer K-cache slot mappings and the two
-generation indptrs with one launch instead of the eager chain.
-
-One decode iteration (the harness of ``test_draft_query_draft_loop_integration.py``:
-target layer 6, then the native MTP Eagle draft loop over layer 78, real
-weights, FP8 KV cache, two requests with a 2,600-token context) updates the
-DSA decode metadata six times:
-
-* ``u0``: ``on_update_kv_lens()`` at the start of the step (B x 6 rows);
-* ``u1``..``u5``: after each of the five draft steps (B rows), through
-  ``update_for_spec_dec()`` -> ``on_update_kv_lens()`` (native step update),
-  or through draft_loop's one-launch ``draft_step_update`` when the step update of
-  the served draft_loop package is attached to the worker.
-
-``u0`` feeds layer 6 and draft step 0, and ``u<i>`` feeds draft step ``i``.
-Every update is recorded: which path ran (native or draft_loop), how many fused
-metadata kernels it launched, and the values it left in the request index,
-the slot mappings, the indptrs, the KV lengths (1-D and 2-D views) and the
-paged-MQA schedules.
-
-Arms per case (custom-op or ``forward_impl`` attention path; native or draft_loop
-draft-step update):
-
-* native, eager metadata chain (``TRTLLM_DISABLE_FUSED_DSA_METADATA=1``);
-* native, served chain (fused kernel), repeated for the output-noise bound;
-* draft_absorbed_query (layer 78 at every draft step), served chain, then repeated;
-* draft_absorbed_query, eager chain.
-
-Asserted:
-
-* served env: no TensorRT-LLM source or library reads
-  ``TRTLLM_FUSED_DSA_METADATA``; the fused path is on with the served env
-  and off only with the opt-out;
-* the update sequence: served chain -> one fused launch at each native update
-  (six per iteration, or one at ``u0`` with the draft_loop step update, whose five
-  steps then run draft_loop's kernel); eager chain -> no fused launch;
-* every recorded metadata value of every update is identical between the
-  eager and the fused chain, with and without draft_absorbed_query, and with the draft_loop step
-  update;
-* draft_absorbed_query under the served chain against native under the served chain: caches,
-  MLA inputs, indexer inputs and appended indexer keys bit-identical; top-k
-  sets equal; draft-loop top-k reuse at steps 1-4; outputs within the
-  native run-to-run bound; draft_absorbed_query repeat bit-identical; draft_absorbed_query under the
-  eager chain bit-identical to draft_absorbed_query under the served chain;
-* CUDA graph: the whole iteration captured with the served chain (draft_absorbed_query graph
-  and native graph in one pool), replayed 1, 1 and 20 times with fresh
-  inputs and accepted counts: the captured fused launches, the replayed
-  metadata, caches and records equal the eager runs as above;
-* a negative control: a slot mapping corrupted after the fused launch of one
-  draft step is caught in that update's metadata and in the layer-78 indexer
-  K cache.
-"""
-
-import contextlib
-import mmap
-import os
-import re
-from types import SimpleNamespace
-
-import pytest
-import test_draft_query_draft_loop_integration as dq_loop_integration
-import torch
-
-TARGET, MTP, LAYERS = (
-    dq_loop_integration.TARGET,
-    dq_loop_integration.MTP,
-    dq_loop_integration.LAYERS,
-)
-DRAFT_LEN, VERIFY_ROWS = dq_loop_integration.DRAFT_LEN, dq_loop_integration.VERIFY_ROWS
-NATIVE_RUNS = dq_loop_integration.NATIVE_RUNS
-SERVED_ENV = "TRTLLM_FUSED_DSA_METADATA"
-DISABLE_ENV = "TRTLLM_DISABLE_FUSED_DSA_METADATA"
-# Metadata outputs of one update: per token, per sequence.
-TOKEN_FIELDS = ("req_idx_per_token", "slot_mapping_fp8", "slot_mapping_scale")
-SEQ_FIELDS = ("gen_kv_indptr", "gen_cached_token_indptr")
-SCHEDULE_FIELDS = ("scheduler_metadata_buffer", "scheduler_metadata_buffer_full_next_n")
-# One indexer key in the FP8 slot mapping (bytes per token).
-INDEX_KEY_BYTES = 128
-
-pytestmark = dq_loop_integration.pytestmark
-harness = dq_loop_integration.harness
-
-
-def _select_chain(chain: str):
-    """Served env, with the fused metadata chain ("fused") or the opt-out
-    ("eager"). The gate is cached per process; it is reset here and never
-    flipped while a graph is being captured."""
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import _fused_dsa_meta_enabled
-
-    os.environ[SERVED_ENV] = "1"
-    if chain == "fused":
-        os.environ.pop(DISABLE_ENV, None)
-    else:
-        os.environ[DISABLE_ENV] = "1"
-    _fused_dsa_meta_enabled.cache_clear()
-    assert _fused_dsa_meta_enabled() == (chain == "fused")
-
-
-def _features(step: str):
-    """Feature selections of the native and the draft_absorbed_query arm. The draft_loop arm attaches
-    draft_loop's draft-loop operations to the worker (``_draft_loop_step_update``) instead
-    of selecting ``draft_loop``: the selection would also wrap the draft embedding,
-    which the harness's stub draft model does not have and which plays no
-    part in the decode metadata."""
-    del step
-    return {"draft_absorbed_query": "draft_absorbed_query", "native": ""}
-
-
-def _inputs(h, nreq, seed, accepted=None):
-    it = h.iteration_inputs(nreq, seed=seed, accepted=accepted)
-    # Step-0 rows of the last accepted tokens (the draft_loop step-0 gather).
-    it.gather_ids = (
-        torch.arange(nreq, device="cuda", dtype=torch.int64) * VERIFY_ROWS
-        + it.num_accepted.to(torch.int64)
-        - 1
-    )
-    return it
-
-
-class _MetaLog:
-    """Records every DSA decode-metadata update of the metadata object
-    ``meta``: the native ``on_update_kv_lens`` and draft_loop's step update. Inside a
-    capture the clones are graph-pool tensors that every replay refreshes;
-    the update kinds and launch counts are host values seen at capture.
-
-    ``corrupt_update``: index of the update after whose fused launch the FP8
-    slot mapping is moved by one indexer key (negative control)."""
-
-    def __init__(self, meta, corrupt_update=None):
-        self.meta = meta
-        self.corrupt_update = corrupt_update
-        self.updates = []
-        self.fused_rows = []
-
-    def __enter__(self):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E501
-            draft_loop as draft_loop_mod,
-        )
-        from tensorrt_llm._torch.attention.backends.sparse.dsa import metadata as dsa_metadata
-
-        log = self
-        self._mod = dsa_metadata
-        self._fused = dsa_metadata.fused_dsa_decode_metadata
-        self._ops_cls = draft_loop_mod.DraftLoopOps
-        self._step_update = draft_loop_mod.DraftLoopOps.step_update
-        native_update = self.meta.on_update_kv_lens
-
-        def fused(*args, **kwargs):
-            out = log._fused(*args, **kwargs)
-            log.fused_rows.append(kwargs["num_tokens"])
-            if log.corrupt_update == len(log.updates):
-                args[4].add_(INDEX_KEY_BYTES)  # slot_mapping_fp8[:num_tokens]
-            return out
-
-        def on_update_kv_lens():
-            before = len(log.fused_rows)
-            native_update()
-            log._record("native", len(log.fused_rows) - before)
-
-        def step_update(ops, attn_metadata, *args, **kwargs):
-            before = len(log.fused_rows)
-            result = log._step_update(ops, attn_metadata, *args, **kwargs)
-            if result is not None and attn_metadata is log.meta:
-                log._record("draft_loop", len(log.fused_rows) - before)
-            return result
-
-        dsa_metadata.fused_dsa_decode_metadata = fused
-        draft_loop_mod.DraftLoopOps.step_update = step_update
-        self.meta.on_update_kv_lens = on_update_kv_lens
-        return self
-
-    def __exit__(self, *exc):
-        self._mod.fused_dsa_decode_metadata = self._fused
-        self._ops_cls.step_update = self._step_update
-        del self.meta.on_update_kv_lens
-        return False
-
-    def _record(self, kind, fused):
-        m = self.meta
-        n, s, g = m.num_tokens, m.num_seqs, m.num_generations
-        fields = {k: getattr(m, k)[:n].clone() for k in TOKEN_FIELDS}
-        fields.update({k: getattr(m, k)[: s + 1].clone() for k in SEQ_FIELDS})
-        fields["kv_lens_cuda"] = m.kv_lens_cuda[:s].clone()
-        width = min(m.gen_token_stride, m.kv_lens_cuda_2d.shape[1])
-        fields["kv_lens_cuda_2d"] = m.kv_lens_cuda_2d[:g, :width].clone()
-        for k in SCHEDULE_FIELDS:
-            t = getattr(m, k, None)
-            if torch.is_tensor(t):
-                fields[k] = t.clone()
-        self.updates.append(SimpleNamespace(kind=kind, rows=n, seqs=s, fused=fused, fields=fields))
-
-    def snapshot(self):
-        return [
-            SimpleNamespace(
-                kind=u.kind,
-                rows=u.rows,
-                seqs=u.seqs,
-                fused=u.fused,
-                fields={k: v.clone() for k, v in u.fields.items()},
-            )
-            for u in self.updates
-        ]
-
-
-@contextlib.contextmanager
-def _draft_loop_step_update(h, meta, it, enabled):
-    """Attach draft_loop's draft-loop operations to the worker, as the served
-    ``_forward_impl`` does for an eligible forward with ``draft_loop`` selected. The
-    harness has no target sampling and no draft embedding, so only draft_loop's
-    step-0 row gather and its per-step metadata update run; acceptance, the
-    draft-token stack and the embedding stay native."""
-    if not enabled:
-        yield None
-        return
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E501
-        draft_loop as draft_loop_mod,
-    )
-
-    ops = draft_loop_mod.DraftLoopOps(h.worker, meta, SimpleNamespace(runtime_draft_len=DRAFT_LEN))
-    ops.gather_ids = it.gather_ids
-    h.worker._draft_loop_ops = ops
-    try:
-        yield ops
-    finally:
-        h.worker._draft_loop_ops = None
-
-
-def _run(h, nreq, it, features, step, chain, path="custom_op", corrupt_update=None):
-    """One eager iteration from the current caches."""
-    _select_chain(chain)
-    dq_loop_integration._set_features(features)
-    h.reset_runners()
-    meta = h.gen_metadata(nreq)
-    with (
-        _MetaLog(meta, corrupt_update) as ml,
-        _draft_loop_step_update(h, meta, it, step == "draft_loop") as ops,
-        dq_loop_integration._LaunchLog() as launches,
-    ):
-        out, rec, calls = h.run(nreq, it, path, meta=meta)
-    return SimpleNamespace(
-        out=out,
-        rec=rec,
-        calls=calls,
-        meta=ml.updates,
-        launches=launches.calls,
-        cache=h.snapshot(),
-        draft_loop_engaged=ops is not None and ops.step_engaged,
-    )
-
-
-def _expected_updates(nreq, step, fused):
-    """``(kind, rows, fused launches)`` of the six metadata updates."""
-    draft_loop_pkg = step == "draft_loop" and fused
-    first = ("native", nreq * VERIFY_ROWS, int(fused))
-    later = ("draft_loop", nreq, 0) if draft_loop_pkg else ("native", nreq, int(fused))
-    return [first] + [later] * DRAFT_LEN
-
-
-def _check_updates(updates, nreq, step, fused, tag):
-    got = [(u.kind, u.rows, u.fused) for u in updates]
-    want = _expected_updates(nreq, step, fused)
-    if got != want or any(u.seqs != nreq for u in updates):
-        return [f"{tag}: metadata updates {got}, expected {want}"]
-    return []
-
-
-def _compare_meta(ref, cand, report, tag):
-    """Every recorded value of every update bit-identical."""
-    if [(u.rows, u.seqs) for u in ref] != [(u.rows, u.seqs) for u in cand]:
-        return [
-            f"{tag}: metadata updates {[(u.rows, u.seqs) for u in cand]}"
-            f" vs {[(u.rows, u.seqs) for u in ref]}"
-        ]
-    fails = []
-    for i, (r, c) in enumerate(zip(ref, cand)):
-        if set(r.fields) != set(c.fields):
-            fails.append(f"{tag}: u{i} fields {sorted(c.fields)} vs {sorted(r.fields)}")
-            continue
-        for key, a in r.fields.items():
-            b = c.fields[key]
-            same = a.dtype == b.dtype and a.shape == b.shape and torch.equal(a, b)
-            report[f"{tag}:u{i}.{key}_same"] = same
-            if not same:
-                fails.append(f"{tag}: u{i} {key} differs")
-    return fails
-
-
-def _summary(report):
-    """Counters of a report: byte-difference counters (all must be 0),
-    key-set and metadata comparisons (all must hold)."""
-    nbytes = [v for k, v in report.items() if k.endswith("_bytes_differing")]
-    sets = [v for k, v in report.items() if k.endswith("_same_sets")]
-    meta = [v for k, v in report.items() if k.endswith("_same")]
-    return (
-        f"byte counters {len(nbytes)} (nonzero "
-        f"{sum(1 for v in nbytes if v)}), key-set comparisons "
-        f"{len(sets)} (unequal {sum(1 for v in sets if not v)}), "
-        f"metadata comparisons {len(meta)} (unequal "
-        f"{sum(1 for v in meta if not v)})"
-    )
-
-
-def _repeat_native(h, nreq, it, features, step, snap, path="custom_op"):
-    """Native iterations under the served chain from one snapshot; they must
-    repeat bit for bit in caches, recorded inputs and metadata."""
-    runs, fails = [], []
-    accepted = it.num_accepted.tolist()
-    expected = dq_loop_integration._expected_slots(nreq, accepted)
-    for i in range(NATIVE_RUNS):
-        h.restore(snap)
-        run = _run(h, nreq, it, features, step, "fused", path)
-        assert run.calls == [], f"native run called absorbed_query entry points: {run.calls}"
-        report = {}
-        fails += _check_updates(run.meta, nreq, step, True, f"native{i}")
-        fails += dq_loop_integration._check_reuse(run.rec, accepted, report, f"native{i}")
-        if i:
-            fails += dq_loop_integration._compare_caches(
-                snap, runs[0].cache, run.cache, expected, report, f"native_repeat{i}"
-            )
-            fails += dq_loop_integration._compare_records(
-                runs[0].rec, run.rec, report, f"native_repeat{i}"
-            )
-            fails += _compare_meta(runs[0].meta, run.meta, report, f"native_repeat{i}")
-        runs.append(run)
-    h.restore(snap)
-    assert not fails, f"native iteration is not repeatable: {fails}"
-    return runs
-
-
-def _output_bounds(native_outs):
-    """Per (layer, call) output bounds from native iterations of one input:
-    twice the largest difference over all pairs of native runs (at least
-    one BF16 ulp, relative RMS at least ``REL_RMS_FLOOR``).
-
-    The attention inputs of these runs are bit-identical (checked
-    separately), so their outputs differ only by the native layer-78
-    attention's run-to-run nondeterminism, which leaves some runs
-    identical and others about one BF16 ulp apart. Every pair samples it;
-    differences from one reference run alone underestimate it."""
-    aq_integration = dq_loop_integration.aq_integration
-    bounds = {}
-    for layer in LAYERS:
-        for i in range(len(native_outs[0][layer])):
-            outs = [run[layer][i] for run in native_outs]
-            noise, rel = 0.0, 0.0
-            for a in range(len(outs)):
-                for b in range(a + 1, len(outs)):
-                    noise = max(noise, (outs[a].float() - outs[b].float()).abs().max().item())
-                    rel = max(rel, aq_integration._rel_rms(outs[a], outs[b]))
-            ulp = aq_integration._bf16_ulp(outs[0].float().abs().max().item())
-            bounds[(layer, i)] = (
-                2.0 * max(noise, ulp),
-                max(2.0 * rel, aq_integration.REL_RMS_FLOOR),
-                noise,
-            )
-    return bounds
-
-
-def _restore_defaults(h, snap):
-    h.worker._draft_loop_ops = None
-    h.restore(snap)
-    dq_loop_integration._set_features("")
-    h.reset_runners()
-    _select_chain("fused")
-
-
-# -- tests ------------------------------------------------------------------
-
-
-def _names_served_env(data) -> bool:
-    """Whether ``data`` contains ``TRTLLM_FUSED_DSA_METADATA`` as a whole
-    name (not as the tail of ``TRTLLM_DISABLE_FUSED_DSA_METADATA``)."""
-    name = SERVED_ENV.encode()
-    start = 0
-    while True:
-        i = data.find(name, start)
-        if i < 0:
-            return False
-        if i == 0 or not re.match(rb"[A-Za-z0-9_]", data[i - 1 : i]):
-            return True
-        start = i + 1
-
-
-def test_served_env_selects_fused_metadata():
-    """The served ``TRTLLM_FUSED_DSA_METADATA=1`` is read nowhere in the
-    installed TensorRT-LLM (Python sources and shared libraries); the fused
-    chain is on with or without it and off only with the opt-out."""
-    import tensorrt_llm
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import _fused_dsa_meta_enabled
-
-    root = os.path.dirname(tensorrt_llm.__file__)
-    scanned, readers = {"py": 0, "so": 0}, []
-    for dirpath, _, files in os.walk(root):
-        for name in files:
-            kind = "py" if name.endswith(".py") else "so" if ".so" in name else None
-            if kind is None:
-                continue
-            path = os.path.join(dirpath, name)
-            with open(path, "rb") as fh:
-                if os.fstat(fh.fileno()).st_size == 0:
-                    continue
-                with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as data:
-                    scanned[kind] += 1
-                    if _names_served_env(data):
-                        readers.append(os.path.relpath(path, root))
-    print(f"scanned {scanned} under {root}; readers of {SERVED_ENV}: {readers}")
-    assert scanned["py"] > 1000 and scanned["so"] > 0, scanned
-    assert readers == [], readers
-
-    saved = {k: os.environ.get(k) for k in (SERVED_ENV, DISABLE_ENV)}
-    try:
-        os.environ.pop(DISABLE_ENV, None)
-        for served in ("1", "0", None):
-            if served is None:
-                os.environ.pop(SERVED_ENV, None)
-            else:
-                os.environ[SERVED_ENV] = served
-            _fused_dsa_meta_enabled.cache_clear()
-            assert _fused_dsa_meta_enabled(), served
-        os.environ[SERVED_ENV] = "1"
-        os.environ[DISABLE_ENV] = "1"
-        _fused_dsa_meta_enabled.cache_clear()
-        assert not _fused_dsa_meta_enabled()
-    finally:
-        for k, v in saved.items():
-            if v is None:
-                os.environ.pop(k, None)
-            else:
-                os.environ[k] = v
-        _fused_dsa_meta_enabled.cache_clear()
-
-
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("step", ["native", "draft_loop"])
-@pytest.mark.parametrize("nreq", [1, 2])
-@torch.inference_mode()
-def test_draft_query_draft_loop_on_fused_metadata(harness, nreq, step, path):
-    """draft_absorbed_query on layer 78 (all five draft steps) with the served metadata chain,
-    against native on the served chain and against the eager chain."""
-    h = harness
-    feats = _features(step)
-    it = _inputs(h, nreq, seed=40 + nreq)
-    accepted = it.num_accepted.tolist()
-    expected = dq_loop_integration._expected_slots(nreq, accepted)
-    snap = h.snapshot()
-    try:
-        natives = _repeat_native(h, nreq, it, feats["native"], step, snap, path)
-        native = natives[0]
-        bounds = _output_bounds([r.out for r in natives])
-        h.restore(snap)
-        chain = _run(h, nreq, it, feats["native"], step, "eager", path)
-        h.restore(snap)
-        fused_run = _run(h, nreq, it, feats["draft_absorbed_query"], step, "fused", path)
-        h.restore(snap)
-        draft_query_again = _run(h, nreq, it, feats["draft_absorbed_query"], step, "fused", path)
-        h.restore(snap)
-        draft_query_chain = _run(h, nreq, it, feats["draft_absorbed_query"], step, "eager", path)
-
-        assert fused_run.calls == dq_loop_integration._expected_calls(
-            "draft_absorbed_query", nreq
-        ), fused_run.calls
-        assert fused_run.launches == dq_loop_integration._expected_launches(nreq), (
-            fused_run.launches
-        )
-        assert draft_query_chain.calls == fused_run.calls and chain.calls == []
-        if step == "draft_loop":
-            assert native.draft_loop_engaged and fused_run.draft_loop_engaged
-
-        report, fails = {}, []
-        for tag, run, fused in (
-            ("native_fused", native, True),
-            ("native_eager", chain, False),
-            ("draft_query_fused", fused_run, True),
-            ("draft_query_again", draft_query_again, True),
-            ("draft_query_eager", draft_query_chain, False),
-        ):
-            fails += _check_updates(run.meta, nreq, step, fused, tag)
-        # Eager vs fused chain, native: metadata, caches and inputs.
-        fails += _compare_meta(chain.meta, native.meta, report, "native_fused_vs_eager")
-        fails += dq_loop_integration._compare_caches(
-            snap, chain.cache, native.cache, expected, report, "native_fused_vs_eager"
-        )
-        fails += dq_loop_integration._compare_records(
-            chain.rec, native.rec, report, "native_fused_vs_eager"
-        )
-        fails += dq_loop_integration._compare_outputs(
-            native.out, chain.out, bounds, report, "native_eager_vs_fused"
-        )
-        # draft_absorbed_query vs native, both on the served chain.
-        fails += _compare_meta(native.meta, fused_run.meta, report, "draft_query_vs_native")
-        fails += dq_loop_integration._compare_to_native(
-            snap,
-            (native.out, native.rec, native.cache),
-            (fused_run.out, fused_run.rec, fused_run.cache),
-            accepted,
-            nreq,
-            bounds,
-            report,
-            "draft_query_vs_native",
-        )
-        # draft_absorbed_query determinism and draft_absorbed_query eager chain vs served chain.
-        for tag, other in (
-            ("draft_query_again", draft_query_again),
-            ("draft_query_eager", draft_query_chain),
-        ):
-            fails += _compare_meta(fused_run.meta, other.meta, report, tag)
-            fails += dq_loop_integration._compare_caches(
-                snap, fused_run.cache, other.cache, expected, report, tag
-            )
-            fails += dq_loop_integration._compare_records(fused_run.rec, other.rec, report, tag)
-            fails += dq_loop_integration._check_reuse(other.rec, accepted, report, tag)
-        print(
-            f"nreq={nreq} step={step} path={path} accepted={accepted} "
-            f"updates={[(u.kind, u.rows, u.fused) for u in fused_run.meta]} "
-            f"summary: {_summary(report)}"
-        )
-        print(f"report {report}")
-        assert not fails, fails
-    finally:
-        _restore_defaults(h, snap)
-
-
-@pytest.mark.parametrize("step", ["native", "draft_loop"])
-@pytest.mark.parametrize("nreq", [1, 2])
-@torch.inference_mode()
-def test_draft_query_fused_metadata_cuda_graph(harness, nreq, step):
-    """The whole iteration with the served chain captured in one CUDA graph
-    per selection (draft_absorbed_query and native, one pool), replayed 1, 1 and 20 times with
-    fresh inputs and accepted counts (``prepare()`` before every replay)."""
-    h = harness
-    snap = h.snapshot()
-    graphs = {}
-    try:
-        _graph_checks(h, nreq, step, snap, graphs)
-    except BaseException as exc:
-        print(f"graph test failure: {exc!r}", flush=True)
-        raise
-    finally:
-        graphs.clear()
-        _restore_defaults(h, snap)
-
-
-def _graph_checks(h, nreq, step, snap, graphs):
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    path = "custom_op"
-    feats = _features(step)
-    _select_chain("fused")
-    warm = {}
-    for name, sel in feats.items():
-        dq_loop_integration._set_features(sel)
-        h.reset_runners()
-        meta = h.gen_metadata(nreq, cuda_graph=True)
-        static = _inputs(h, nreq, seed=200)
-        with with_multi_stream(True):
-            for _ in range(2):
-                meta.prepare()
-                with _draft_loop_step_update(h, meta, static, step == "draft_loop"):
-                    h.iteration(meta, static, path)
-        torch.cuda.synchronize()
-        runners = {
-            layer: (h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query)
-            for layer in LAYERS
-        }
-        if name == "draft_absorbed_query":
-            assert runners[MTP][1] not in (None, False), "draft_absorbed_query did not engage"
-        warm[name] = (meta, static, runners)
-    h.restore(snap)
-    pool = None
-    for name, sel in feats.items():
-        meta, static, runners = warm[name]
-        dq_loop_integration._set_features(sel)
-        for layer, (r_absorbed_query, r_draft_query) in runners.items():
-            h.attn[layer]._absorbed_query, h.attn[layer]._draft_absorbed_query = (
-                r_absorbed_query,
-                r_draft_query,
-            )
-        meta.prepare()
-        graph = torch.cuda.CUDAGraph()
-        with (
-            with_multi_stream(True),
-            dq_loop_integration._Recorder(h.attn) as rec,
-            dq_loop_integration._Engagement() as eng,
-            dq_loop_integration._LaunchLog() as launches,
-            _MetaLog(meta) as ml,
-            _draft_loop_step_update(h, meta, static, step == "draft_loop") as ops,
-        ):
-            with torch.cuda.graph(graph, pool=pool):
-                outs = h.iteration(meta, static, path)
-        pool = graph.pool()
-        assert eng.calls == dq_loop_integration._expected_calls(sel, nreq), eng.calls
-        if name == "draft_absorbed_query":
-            assert launches.calls == dq_loop_integration._expected_launches(nreq), launches.calls
-        if step == "draft_loop":
-            assert ops.step_engaged, "draft_loop step update declined at capture"
-        captured = _check_updates(ml.updates, nreq, step, True, f"capture:{name}")
-        assert not captured, captured
-        print(
-            f"[fused-meta graph] captured {name}: updates "
-            f"{[(u.kind, u.rows, u.fused) for u in ml.updates]}",
-            flush=True,
-        )
-        graphs[name] = (graph, meta, static, outs, rec, ml, runners)
-        torch.cuda.synchronize()
-    h.reset_runners()
-
-    report, fails = {}, []
-    trials = [(311, 1, None), (312, 1, [1] * nreq), (313, 20, [VERIFY_ROWS] * nreq)]
-    for trial, (seed, replays, accepted) in enumerate(trials):
-        it = _inputs(h, nreq, seed=seed, accepted=accepted)
-        accepted = it.num_accepted.tolist()
-        expected = dq_loop_integration._expected_slots(nreq, accepted)
-        natives = _repeat_native(h, nreq, it, feats["native"], step, snap, path)
-        native = natives[0]
-        bounds = _output_bounds([r.out for r in natives])
-        h.restore(snap)
-        chain = _run(h, nreq, it, feats["native"], step, "eager", path)
-        h.restore(snap)
-        draft_query_eager = _run(h, nreq, it, feats["draft_absorbed_query"], step, "fused", path)
-        _select_chain("fused")
-        dq_loop_integration._set_features("")
-        h.reset_runners()
-        tag = f"t{trial}x{replays}"
-        for name, (graph, meta, static, outs, rec, ml, runners) in graphs.items():
-            h.restore(snap)
-            meta.prepare()
-            for key in ("hidden_t", "pos_t", "hidden_d", "num_accepted", "input_ids", "gather_ids"):
-                getattr(static, key).copy_(getattr(it, key))
-            for i in range(replays):
-                if i:
-                    meta.prepare()
-                graph.replay()
-            torch.cuda.synchronize()
-            g_out = {layer: [o.clone() for o in outs[layer]] for layer in LAYERS}
-            g_rec, g_meta, g_cache = rec.snapshot(), ml.snapshot(), h.snapshot()
-            fails += dq_loop_integration._check_reuse(g_rec, accepted, report, f"{tag}:{name}")
-            fails += _compare_meta(
-                native.meta, g_meta, report, f"{tag}:{name}_graph_vs_native_fused"
-            )
-            fails += _compare_meta(
-                chain.meta, g_meta, report, f"{tag}:{name}_graph_vs_native_eager"
-            )
-            if name == "draft_absorbed_query":
-                fails += dq_loop_integration._compare_caches(
-                    snap,
-                    draft_query_eager.cache,
-                    g_cache,
-                    expected,
-                    report,
-                    f"{tag}:draft_query_graph_vs_draft_query_eager",
-                )
-                fails += dq_loop_integration._compare_records(
-                    draft_query_eager.rec,
-                    g_rec,
-                    report,
-                    f"{tag}:draft_query_graph_vs_draft_query_eager",
-                )
-                fails += dq_loop_integration._compare_to_native(
-                    snap,
-                    (native.out, native.rec, native.cache),
-                    (g_out, g_rec, g_cache),
-                    accepted,
-                    nreq,
-                    bounds,
-                    report,
-                    f"{tag}:draft_query_graph_vs_native",
-                )
-            else:
-                fails += dq_loop_integration._compare_caches(
-                    snap, native.cache, g_cache, expected, report, f"{tag}:native_graph_vs_eager"
-                )
-                fails += dq_loop_integration._compare_records(
-                    native.rec, g_rec, report, f"{tag}:native_graph_vs_eager"
-                )
-                fails += dq_loop_integration._compare_outputs(
-                    native.out, g_out, bounds, report, f"{tag}:native_graph_vs_eager"
-                )
-            for r_absorbed_query, r_draft_query in runners.values():
-                for runner in (r_absorbed_query, r_draft_query):
-                    if runner not in (None, False):
-                        assert int(runner.sync.abs().sum()) == 0
-            print(
-                f"[fused-meta graph] trial {trial} {name}: replayed "
-                f"x{replays}, accepted={accepted}",
-                flush=True,
-            )
-    print(f"graph nreq={nreq} step={step} summary: {_summary(report)}")
-    print(f"report {report}")
-    assert not fails, fails
-
-
-@torch.inference_mode()
-def test_checks_detect_corrupted_draft_step_metadata(harness):
-    """Negative control: the FP8 slot mapping of the update that feeds draft
-    step 2 (``u2``), moved by one indexer key after the fused launch, must be
-    flagged in that update only and in the layer-78 indexer K cache only."""
-    h = harness
-    nreq = 1
-    it = _inputs(h, nreq, seed=12)
-    accepted = it.num_accepted.tolist()
-    snap = h.snapshot()
-    try:
-        h.restore(snap)
-        native = _run(h, nreq, it, "", "native", "fused")
-        h.restore(snap)
-        bad = _run(h, nreq, it, "draft_absorbed_query", "native", "fused", corrupt_update=2)
-        assert bad.calls == dq_loop_integration._expected_calls("draft_absorbed_query", nreq)
-        report = {}
-        meta_fails = _compare_meta(native.meta, bad.meta, report, "bad")
-        cache_fails = dq_loop_integration._compare_caches(
-            snap,
-            native.cache,
-            bad.cache,
-            dq_loop_integration._expected_slots(nreq, accepted),
-            report,
-            "bad",
-        )
-        print(f"negative control: metadata {meta_fails}; caches {cache_fails}")
-        assert meta_fails == ["bad: u2 slot_mapping_fp8 differs"], meta_fails
-        assert any(f"index_k{MTP}" in f for f in cache_fails), cache_fails
-        assert not any(f"index_k{TARGET}" in f or "latent" in f for f in cache_fails), cache_fails
-    finally:
-        _restore_defaults(h, snap)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py
deleted file mode 100644
index a6bae5b8fc1d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_draft_query_mtp_query_proj.py
+++ /dev/null
@@ -1,439 +0,0 @@
-# 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.
-"""draft_absorbed_query: the absorbed_query query_front kernel on the MTP draft layer, against native.
-
-Fixtures: the real GLM-5.2 checkpoint weights of the one-model MTP layer
-(``model.layers.78.self_attn.*``, BF16, target-layer shapes, full indexer)
-at TP8 rank 0, placed exactly as for target layers (``absorbed_query_fixture``), and the
-layer's real ``input_layernorm`` gain; synthetic activations. The module is
-the native ``DeepseekV32Attention`` with ``layer_idx = 78``, as
-``DeepseekV3MTP`` builds it in the one-model engine.
-
-Draft-step shapes of the MTP Eagle loop at batch 1: step 0 presents the six
-verify rows (``R = 6``, positions ``p .. p + 5``), steps 1-4 one row each
-(``R = 1``, position ``p + accepted + i - 1``).
-
-Asserted, with the absorbed_query checks (``test_absorbed_query_proj``):
-
-* eligibility: absorbed_query's ``dsa_proj`` still declines layer 78 ("target layers
-  only"); ``mtp_dsa_proj`` engages on layer 78 only, and only for
-  generation-only calls inside the absorbed_query domain;
-* the frozen absorbed_query error budget and bit identity with native for every
-  output (``aq_proj._native_agreement``) for R in {1, 6} and the other absorbed_query row
-  counts, both forms, fused and native indexer tail;
-* padded rows (whole padded requests, as CUDA-graph padding adds them) never
-  change live rows;
-* the draft loop's five calls of one layer captured back to back in one CUDA
-  graph, sharing the layer's sync counters (also with delayed producers /
-  consumers), reproduce eager absorbed_query calls bit for bit on 1, 1 and 20 replays
-  and leave the counters at zero.
-"""
-
-import os
-
-import absorbed_query_fixture
-import pytest
-import test_absorbed_query_proj as aq_proj
-import torch
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-MTP_LAYER = 78
-# Context length at the start of the decode iteration (the controlled
-# decode workload) and draft-loop shape at batch 1: R_d = [6, 1, 1, 1, 1].
-CTX = 491061
-DRAFT_LEN = 5
-ROWS = aq_proj._env_list("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_TEST_ROWS", [1, 6, 2, 3, 4, 5, 8, 16])
-
-_ATTN = {}
-
-
-@pytest.fixture(scope="module")
-def model_config():
-    return absorbed_query_fixture.build_model_config(0)
-
-
-def _attn(model_config, layer=MTP_LAYER):
-    if layer not in _ATTN:
-        _ATTN[layer] = absorbed_query_fixture.build_native_attention(layer, 0, model_config)
-    return _ATTN[layer]
-
-
-def _gamma():
-    return absorbed_query_fixture.load_input_layernorm(MTP_LAYER)
-
-
-def _hidden(rows, seed):
-    return absorbed_query_fixture.make_hidden(rows, seed=seed, gamma=_gamma())
-
-
-def _draft_step_inputs(step, accepted=3, seed=0, batch=1):
-    """Hidden rows and positions of draft step ``step`` (batch requests)."""
-    if step == 0:
-        rows = batch * (DRAFT_LEN + 1)
-        pos = (CTX + torch.arange(DRAFT_LEN + 1, dtype=torch.int32)).repeat(batch)
-    else:
-        rows = batch
-        pos = torch.full((batch,), CTX + accepted + step - 1, dtype=torch.int32)
-    # Step >= 1 position ids are 1-D (one per request), as the draft loop
-    # builds them; step 0 reuses the target's [1, num_tokens] position ids.
-    pos = pos.cuda()
-    if step == 0:
-        pos = pos.unsqueeze(0)
-    return _hidden(rows, seed=seed + 31 * step), pos
-
-
-@torch.inference_mode()
-def test_mtp_layer_fixture(model_config):
-    """Layer 78 is the MTP layer: full indexer with draft-loop top-k reuse,
-    BF16 attention weights of the target-layer shapes."""
-    attn = _attn(model_config)
-    target = _attn(model_config, absorbed_query_fixture.FULL_LAYER)
-    assert attn.layer_idx == MTP_LAYER == attn.num_hidden_layers
-    indexer = attn.mqa.indexer
-    assert indexer is not None and indexer.mtp_index_share
-    for name in ("kv_a_proj_with_mqa", "q_b_proj"):
-        a, t = getattr(attn, name).weight, getattr(target, name).weight
-        assert a.dtype == t.dtype == torch.bfloat16 and a.shape == t.shape
-        assert not torch.equal(a, t), f"{name}: layer 78 weights not loaded"
-    assert attn.k_b_proj_trans.shape == target.k_b_proj_trans.shape
-    assert indexer.wq_b.weight.shape == target.mqa.indexer.wq_b.weight.shape
-
-
-@torch.inference_mode()
-def test_eligibility(model_config):
-    """absorbed_query declines layer 78 as before; draft_absorbed_query engages on layer 78 only, for
-    generation-only calls inside the absorbed_query domain, before any GPU work."""
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    mtp = _attn(model_config)
-    target = _attn(model_config, absorbed_query_fixture.FULL_LAYER)
-    assert query_proj.module_eligibility(mtp) == (False, "target layers only")
-    assert query_proj.module_eligibility(mtp, draft=True) == (True, "")
-    assert query_proj.module_eligibility(target) == (True, "")
-    assert query_proj.module_eligibility(target, draft=True) == (False, "draft (MTP) layers only")
-
-    gen = SimpleNamespace(num_contexts=0, num_generations=1)
-    ctx = SimpleNamespace(num_contexts=1, num_generations=0)
-    mixed = SimpleNamespace(num_contexts=1, num_generations=1)
-    for step in range(DRAFT_LEN):
-        hidden, pos = _draft_step_inputs(step)
-        for attr in ("_absorbed_query", "_draft_absorbed_query"):
-            for mod in (mtp, target):
-                setattr(mod, attr, None)
-        assert query_proj.dsa_proj(mtp, pos, hidden, gen) is None
-        assert mtp._absorbed_query is False
-        assert query_proj.mtp_dsa_proj(target, pos, hidden, gen) is None
-        assert target._draft_absorbed_query is False
-        packet = query_proj.mtp_dsa_proj(mtp, pos, hidden, gen)
-        assert packet is not None and len(packet.proj_outputs) == 9
-        assert packet.fused_q is not None  # auto form: preabsorbed (full)
-        for meta in (None, ctx, mixed):
-            assert query_proj.mtp_dsa_proj(mtp, pos, hidden, meta) is None
-    big = torch.zeros(17, 6144, dtype=torch.bfloat16, device="cuda")
-    assert query_proj.mtp_dsa_proj(mtp, None, big, gen) is None
-    assert query_proj.mtp_dsa_proj(mtp, None, big[:6].float(), gen) is None
-    for mod in (mtp, target):
-        mod._absorbed_query = mod._draft_absorbed_query = None
-
-
-@pytest.mark.parametrize("rows", ROWS)
-@pytest.mark.parametrize("absorb", [True, False])
-@pytest.mark.parametrize("fuse_indexer", [True, False])
-@torch.inference_mode()
-def test_mtp_contract_q_matches_native(model_config, rows, absorb, fuse_indexer):
-    """The query_front contract on layer 78 against the native region, absorbed_query budget and
-    bitwise agreement; R = 6 and R = 1 are the draft-loop shapes."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    attn = _attn(model_config)
-    hidden = _hidden(rows, seed=rows)
-    pos = (CTX + torch.arange(rows, dtype=torch.int32, device="cuda")).unsqueeze(0)
-    native_outs, native_fq = aq_proj._native(attn, hidden, pos)
-    runner = AbsorbedQueryProjection(attn, KernelConfig(fuse_indexer=fuse_indexer))
-    assert runner.indexer_fusable() == fuse_indexer
-    packet = runner.project(pos, hidden, absorb=absorb)
-    torch.cuda.synchronize()
-    ref = aq_proj._oracle(attn, hidden)
-    outs = packet.proj_outputs
-    report = {}
-    assert len(outs) == len(native_outs) == 9
-    for cand, nat in zip(outs, native_outs):
-        assert cand.shape == nat.shape and cand.dtype == nat.dtype
-    assert (packet.fused_q is not None) == absorb
-    assert torch.equal(outs[3][:, : attn.kv_lora_rank], outs[1])
-    assert torch.equal(outs[3][:, attn.kv_lora_rank :], outs[2])
-    aq_proj._check_budget("q", outs[0], native_outs[0], ref["q"], report)
-    aq_proj._check_budget("compressed_kv", outs[1], native_outs[1], ref["ckv"], report)
-    if absorb:
-        aq_proj._check_budget(
-            "absorbed",
-            packet.fused_q[..., : attn.kv_lora_rank],
-            native_fq[..., : attn.kv_lora_rank],
-            ref["absorbed"],
-            report,
-        )
-    q_fp8, k_fp8, k_scale, weights, q_scale = outs[4:]
-    nq_fp8, nk_fp8, nk_scale, nweights, nq_scale = native_outs[4:]
-    torch.testing.assert_close(
-        k_fp8.float() * k_scale, nk_fp8.float() * nk_scale, rtol=0.1, atol=1e-2
-    )
-    torch.testing.assert_close(
-        q_fp8.float() * q_scale, nq_fp8.float() * nq_scale, rtol=0.1, atol=1e-2
-    )
-    torch.testing.assert_close(weights, nweights, rtol=0.1, atol=1e-2)
-    stats, _ = aq_proj._native_agreement(
-        outs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
-    )
-    print(
-        f"layer={MTP_LAYER} rows={rows} absorb={absorb} "
-        f"fuse_indexer={fuse_indexer} budget={report} agreement={stats}"
-    )
-    assert int(runner.sync.abs().sum()) == 0
-
-
-@pytest.mark.parametrize("step", list(range(DRAFT_LEN)))
-@torch.inference_mode()
-def test_mtp_entry_point_per_draft_step(model_config, step):
-    """``mtp_dsa_proj`` at each draft step's shape and positions equals the
-    default absorbed_query runner and native bit for bit."""
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    attn = _attn(model_config)
-    attn._draft_absorbed_query = None
-    gen = SimpleNamespace(num_contexts=0, num_generations=1)
-    hidden, pos = _draft_step_inputs(step, seed=40)
-    packet = query_proj.mtp_dsa_proj(attn, pos, hidden, gen)
-    reference = query_proj.AbsorbedQueryProjection(attn).project(
-        pos, hidden, absorb=query_proj.absorb_for(attn)
-    )
-    native_outs, native_fq = aq_proj._native(attn, hidden, pos)
-    torch.cuda.synchronize()
-    assert packet.proj_outputs[0].shape[0] == (6 if step == 0 else 1)
-    aq_proj._assert_same_packet(packet, reference, f"step {step}")
-    aq_proj._native_agreement(
-        packet.proj_outputs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
-    )
-    assert int(attn._draft_absorbed_query.sync.abs().sum()) == 0
-    attn._draft_absorbed_query = None
-
-
-@pytest.mark.parametrize("rows", [6, 1])
-@torch.inference_mode()
-def test_mtp_native_agreement_statistics(model_config, rows):
-    """Default kernel on layer 78 over many inputs: per-output difference
-    statistics against native; every input must agree bit for bit."""
-    import json
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    attn = _attn(model_config)
-    indexer = attn.mqa.indexer
-    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_keep_intermediates=True))
-    n_inputs = int(os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_AGREEMENT_INPUTS", "48"))
-    totals, inputs_with_diff, fails = {}, {}, []
-    for i in range(n_inputs):
-        hidden = _hidden(rows, seed=2000 + 17 * i)
-        pos = (CTX + i + torch.arange(rows, dtype=torch.int32, device="cuda")).unsqueeze(0)
-        native_outs, native_fq = aq_proj._native(attn, hidden, pos)
-        packet = runner.project(pos, hidden, absorb=True)
-        stats, input_fails = aq_proj._native_agreement(
-            packet.proj_outputs,
-            packet.fused_q,
-            native_outs,
-            native_fq,
-            attn.kv_lora_rank,
-            check=False,
-        )
-        kva = attn.kv_a_proj_with_mqa(hidden)
-        qr = attn._q_a_layernorm_maybe_fused(kva[:, : attn.q_lora_rank])
-        steps = aq_proj._steps(runner.debug["iq"], indexer.wq_b(qr))
-        stats["indexer_wq_b_bf16"] = (int((steps != 0).sum()), steps.numel(), int(steps.max()))
-        fails += [f"input {i}: {f}" for f in input_fails]
-        for name, (n_diff, n, max_steps) in stats.items():
-            tot = totals.setdefault(name, [0, 0, 0])
-            tot[0] += n_diff
-            tot[1] += n
-            tot[2] = max(tot[2], max_steps)
-            inputs_with_diff[name] = inputs_with_diff.get(name, 0) + (n_diff > 0)
-    summary = {
-        name: dict(
-            inputs_with_diff=inputs_with_diff[name],
-            inputs=n_inputs,
-            diff_elements=tot[0],
-            elements=tot[1],
-            max_steps=tot[2],
-        )
-        for name, tot in totals.items()
-    }
-    print(f"AGREEMENT layer={MTP_LAYER} rows={rows} {json.dumps(summary)}")
-    assert not fails, fails
-    assert int(runner.sync.abs().sum()) == 0
-
-
-@pytest.mark.parametrize("rows,live", [(12, 6), (6, 3), (2, 1)])
-@pytest.mark.parametrize("absorb", [True, False])
-@torch.inference_mode()
-def test_mtp_padded_rows_do_not_affect_live_rows(model_config, rows, live, absorb):
-    """Graph padding adds whole dummy requests: 12 = 2 x 6 step-0 rows with
-    one live request, 2 step >= 1 rows with one live request. Padded rows
-    (zero or 1e3 garbage) never change any live-row output bit for bit, and
-    live rows agree with native bit for bit."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-    )
-
-    attn = _attn(model_config)
-    runner = AbsorbedQueryProjection(attn)
-    live_rows = slice(0, live)
-    hidden_zero = _hidden(rows, seed=77)
-    hidden_zero[live:] = 0
-    hidden_junk = hidden_zero.clone()
-    hidden_junk[live:] = 1e3 * torch.randn_like(hidden_junk[live:])
-    pos = (CTX + torch.arange(rows, dtype=torch.int32, device="cuda") % (DRAFT_LEN + 1)).unsqueeze(
-        0
-    )
-    a = runner.project(pos, hidden_zero, absorb=absorb)
-    b = runner.project(pos, hidden_junk, absorb=absorb)
-    native_outs, native_fq = aq_proj._native(attn, hidden_junk, pos)
-    torch.cuda.synchronize()
-    for name, x, y in zip(aq_proj._NAMES, a.proj_outputs, b.proj_outputs):
-        aq_proj._assert_bitwise(f"padded {name}", x[live_rows], y[live_rows])
-    if absorb:
-        aq_proj._assert_bitwise(
-            "padded absorbed",
-            a.fused_q[live_rows, :, : attn.kv_lora_rank],
-            b.fused_q[live_rows, :, : attn.kv_lora_rank],
-        )
-    aq_proj._native_agreement(
-        b.proj_outputs, b.fused_q, native_outs, native_fq, attn.kv_lora_rank, live_rows
-    )
-    for x in a.proj_outputs:
-        assert torch.isfinite(x[live:].float()).all()
-    assert int(runner.sync.abs().sum()) == 0
-
-
-def _draft_sequence(runner, static_inputs, absorb):
-    """The five draft-step calls of the MTP layer, back to back."""
-    return [runner.project(pos, hidden, absorb=absorb) for hidden, pos in static_inputs]
-
-
-def _check_step_vs_native(attn, packet, hidden, pos):
-    """Every output bit-identical to the native region."""
-    native_outs, native_fq = aq_proj._native(attn, hidden, pos)
-    torch.cuda.synchronize()
-    aq_proj._native_agreement(
-        packet.proj_outputs, packet.fused_q, native_outs, native_fq, attn.kv_lora_rank
-    )
-
-
-@pytest.mark.parametrize("absorb", [True, False])
-@pytest.mark.parametrize("delay_ns", [0, 20000])
-@torch.inference_mode()
-def test_mtp_draft_loop_graph_replay(model_config, absorb, delay_ns):
-    """One CUDA graph holding the draft loop's five calls of layer 78
-    (R = 6, then four R = 1 calls), all on the same runner and sync
-    counters, replayed 1, 1 and 20 times with fresh inputs: every call equals
-    an eager absorbed_query call and native bit for bit, and
-    the counters are zero afterwards (also with delayed producers and
-    consumers, which stresses counter reuse between adjacent launches)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    attn = _attn(model_config)
-    runner = AbsorbedQueryProjection(attn, KernelConfig(debug_delay_ns=delay_ns))
-    eager = AbsorbedQueryProjection(attn)
-    static_inputs = [_draft_step_inputs(step, seed=300) for step in range(DRAFT_LEN)]
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(2):
-            _draft_sequence(runner, static_inputs, absorb)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        packets = _draft_sequence(runner, static_inputs, absorb)
-    for trial, replays in enumerate([1, 1, 20]):
-        fresh = [
-            _draft_step_inputs(step, accepted=1 + trial, seed=400 + 10 * trial)
-            for step in range(DRAFT_LEN)
-        ]
-        for (s_hidden, s_pos), (hidden, pos) in zip(static_inputs, fresh):
-            s_hidden.copy_(hidden)
-            s_pos.copy_(pos)
-        for _ in range(replays):
-            graph.replay()
-        torch.cuda.synchronize()
-        for step, (packet, (hidden, pos)) in enumerate(zip(packets, fresh)):
-            reference = eager.project(pos, hidden, absorb=absorb)
-            torch.cuda.synchronize()
-            tag = f"trial {trial} x{replays} step {step}"
-            aq_proj._assert_same_packet(packet, reference, tag)
-            _check_step_vs_native(attn, packet, hidden, pos)
-        assert int(runner.sync.abs().sum()) == 0
-        assert int(eager.sync.abs().sum()) == 0
-    del graph
-
-
-@pytest.mark.parametrize("absorb", [True, False])
-@torch.inference_mode()
-def test_mtp_draft_loop_back_to_back_stress(model_config, absorb):
-    """Many eager draft-loop sequences (five back-to-back launches on one
-    runner) with delayed producers/consumers and fresh inputs reproduce the
-    undelayed launches bit for bit; counters stay zero."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        KernelConfig,
-    )
-
-    attn = _attn(model_config)
-    delayed = AbsorbedQueryProjection(attn, KernelConfig(debug_delay_ns=50000))
-    plain = AbsorbedQueryProjection(attn)
-    results = []
-    for it in range(int(os.environ.get("GLM_MOE_DSA_DRAFT_ABSORBED_QUERY_STRESS_ITERS", "6"))):
-        inputs = [
-            _draft_step_inputs(step, accepted=1 + it % 6, seed=600 + it)
-            for step in range(DRAFT_LEN)
-        ]
-        results.append(
-            (_draft_sequence(delayed, inputs, absorb), _draft_sequence(plain, inputs, absorb))
-        )
-    torch.cuda.synchronize()
-    for it, (got, want) in enumerate(results):
-        for step, (a, b) in enumerate(zip(got, want)):
-            aq_proj._assert_same_packet(a, b, f"iteration {it} step {step}")
-    assert int(delayed.sync.abs().sum()) == 0
-    assert int(plain.sync.abs().sum()) == 0
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py
deleted file mode 100644
index 2c3539d2c41d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_attention_integration.py
+++ /dev/null
@@ -1,447 +0,0 @@
-# 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.
-"""Stock producer -> shared_indexer_absorbed_query -> native DSA attention, through the target's
-attention (the absorbed-query packet of absorbed_query, draft_absorbed_query and
-shared_indexer_absorbed_query).
-
-Uses the absorbed_query integration harness
-(``absorbed_query/test_absorbed_query_attention_integration .py``): real weights of the full-indexer
-layer 6 and the shared layer 7 at TP8 rank 0, six requests natively prefilled with 2,600 tokens, FP8
-KV cache, generation steps in the 6 x 1 and the 1 x 6 (MTP verify) layouts, both call
-paths, every step preceded by ``on_update_kv_lens()`` as in the runner.
-
-With the fast paths ``absorbed_query,shared_indexer_absorbed_query`` running, layer 7 takes
-shared_indexer_absorbed_query (native kv_a GEMM + the shared_indexer_absorbed_query
-kernel, preabsorbed packet, no native BMM) and layer 6 takes absorbed_query.
-Checks against the same step with every fast path off:
-
-* engagement: shared_indexer_absorbed_query serves layer 7 only; absorbed_query serves layer 6 only;
-* layer 7 (shared_indexer_absorbed_query): its latent cache is **byte-identical** to native, and the
-  attention inputs recorded at the native ``mla_rope_generation`` (absorbed
-  query, FP8 ``quant_q_buffer``, BMM scales) and the converted sparse
-  indices are bit-identical / select the same keys;
-* layer 6 (absorbed_query): bit-identical caches, attention and indexer inputs and
-  equal top-k key sets (the absorbed_query harness's comparison);
-* outputs within the native-noise bound; a repeated shared_indexer_absorbed_query step and 20 more
-  steps reproduce the shared_indexer_absorbed_query caches and recorded inputs bit for bit;
-* CUDA graphs (shared_indexer_absorbed_query alone, native alone, both in one pool): replays 1, 1
-  and 20 times with fresh hidden states reproduce the eager shared_indexer_absorbed_query step bit
-  for bit and compare with native as above; counters are zero after each trial;
-* feature gating through the target's ``_absorbed_query_packet``: ``""``, ``absorbed_query``,
-  ``shared_indexer_absorbed_query``, ``absorbed_query,shared_indexer_absorbed_query``,
-  ``absorbed_query,draft_absorbed_query,shared_indexer_absorbed_query``;
-* negative control: a shared_indexer_absorbed_query packet whose absorbed query is zeroed is
-  flagged.
-"""
-
-import math
-import os
-import sys
-
-import absorbed_query_fixture
-import pytest
-import test_absorbed_query_attention_integration as aq_integration
-import torch
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
-
-import plan_testlib  # noqa: E402
-
-LAYERS = aq_integration.LAYERS
-SHARED = absorbed_query_fixture.SHARED_LAYER
-FULL = absorbed_query_fixture.FULL_LAYER
-LAYOUTS = aq_integration.LAYOUTS
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-
-def _set_features(features: str, form: str = "auto"):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    query_proj.SERVED_FORM = form
-    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
-
-
-def _reset_runners(h):
-    for layer in LAYERS:
-        h.attn[layer]._absorbed_query = None
-        h.attn[layer]._shared_indexer_query = None
-
-
-def _runners(h):
-    return {
-        layer: (
-            getattr(h.attn[layer], "_absorbed_query", None),
-            getattr(h.attn[layer], "_shared_indexer_query", None),
-        )
-        for layer in LAYERS
-    }
-
-
-def _assert_engaged(h):
-    r = _runners(h)
-    assert r[SHARED][1] not in (None, False), "shared_indexer_absorbed_query did not engage on L7"
-    assert r[SHARED][0] is None, (
-        "absorbed_query must not run on the shared_indexer_absorbed_query layer"
-    )
-    assert r[FULL][1] is False, "shared_indexer_absorbed_query must decline the full layer"
-    assert r[FULL][0] not in (None, False), "absorbed_query did not engage on L6"
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = aq_integration._Harness()
-    yield h
-    _set_features("")
-
-
-def _compare_shared_query_to_native(snap, native, cand, n_new, bounds, report, tag):
-    """absorbed_query harness comparison (bitwise on both layers) plus an explicit check
-    of the shared_indexer_absorbed_query layer's latent cache."""
-    fails = aq_integration._compare_absorbed_query_to_native(
-        snap, native, cand, n_new, bounds, report, tag
-    )
-    n_diff = report.get(f"{tag}:latent{SHARED}_bytes_differing")
-    if n_diff != 0:
-        fails.append(
-            f"{tag}: shared_indexer_absorbed_query layer latent cache differs from native in "
-            f"{n_diff} bytes"
-        )
-    return fails
-
-
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_shared_query_into_native_attention(harness, layout, path):
-    h = harness
-    n_new = math.prod(LAYOUTS[layout])
-    hidden, pos = h.gen_inputs(layout, seed=12)
-    snap = h.snapshot()
-    try:
-        native_outs, native_recs, native_cache = aq_integration._native_runs(
-            h, layout, path, hidden, pos, snap
-        )
-        bounds = aq_integration._output_bounds(native_outs)
-
-        _set_features("absorbed_query,shared_indexer_absorbed_query")
-        _reset_runners(h)
-        shared_query_out, shared_query_rec = h.recorded_step(layout, path, hidden, pos)
-        _assert_engaged(h)
-        shared_query_cache = h.snapshot()
-        rope_served = {layer: "rope_append_rope" in shared_query_rec[layer] for layer in LAYERS}
-        print(f"rope_kv_append rope served: {rope_served}", flush=True)
-        if path == "custom_op" and "rope_kv_append" in absorbed_query_fixture.with_base_features(
-            ""
-        ).split(","):
-            assert rope_served[SHARED], (
-                "rope_kv_append did not serve the shared_indexer_absorbed_query layer's rope"
-            )
-
-        report, fails = {}, []
-        fails += _compare_shared_query_to_native(
-            snap,
-            (native_outs[0], native_recs[0], native_cache),
-            (shared_query_out, shared_query_rec, shared_query_cache),
-            n_new,
-            bounds,
-            report,
-            "shared_indexer_absorbed_query",
-        )
-
-        h.restore(snap)
-        _, rec_again = h.recorded_step(layout, path, hidden, pos)
-        fails += aq_integration._compare_caches(
-            snap, shared_query_cache, h.snapshot(), n_new, report, "shared_query_again"
-        )
-        fails += aq_integration._compare_records(
-            shared_query_rec, rec_again, report, "shared_query_again"
-        )
-        for _ in range(20):
-            again = h.gen_step(layout, path, hidden, pos)
-        fails += aq_integration._compare_caches(
-            snap, shared_query_cache, h.snapshot(), n_new, report, "shared_query_repeat20"
-        )
-        fails += aq_integration._compare_outputs(
-            native_outs[0], again, bounds, report, "shared_query_repeat20"
-        )
-        assert h.attn[SHARED]._shared_indexer_query.counters_quiescent()
-        print(f"layout={layout} path={path} {report}")
-        assert not fails, fails
-    finally:
-        h.restore(snap)
-        _set_features("")
-        _reset_runners(h)
-
-
-@pytest.mark.parametrize(
-    "features,expect",
-    [
-        ("", {FULL: (None, None), SHARED: (None, None)}),
-        ("absorbed_query", {FULL: ("absorbed_query", None), SHARED: ("absorbed_query", None)}),
-        ("shared_indexer_absorbed_query", {FULL: (None, None), SHARED: (None, None)}),
-        (
-            "absorbed_query,shared_indexer_absorbed_query",
-            {FULL: ("absorbed_query", False), SHARED: (None, "shared_indexer_absorbed_query")},
-        ),
-        (
-            "absorbed_query,draft_absorbed_query,shared_indexer_absorbed_query",
-            {FULL: ("absorbed_query", False), SHARED: (None, "shared_indexer_absorbed_query")},
-        ),
-    ],
-)
-@torch.inference_mode()
-def test_feature_gating_through_adapter(harness, features, expect):
-    """Which runner serves which layer for each feature selection (the draft_absorbed_query
-    entry point declines target layers, so ``draft_absorbed_query`` changes nothing here).
-    ``shared_indexer_absorbed_query`` without ``absorbed_query`` is refused by the selection itself
-    (``_set_features``), before any runner is attached, when the installed selector has the
-    dependency rule (``REQUIRED_FEATURES``); without that rule it runs native."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-        shared_indexer,
-    )
-
-    h = harness
-    layout, path = "1req_x6tok", "custom_op"
-    hidden, pos = h.gen_inputs(layout, seed=21)
-    snap = h.snapshot()
-    ids = features.split(",")
-    rejected = (
-        "shared_indexer_absorbed_query" in ids
-        and "absorbed_query" not in ids
-        and "shared_indexer_absorbed_query" in getattr(dispatch, "REQUIRED_FEATURES", {})
-    )
-    try:
-        if rejected:
-            _reset_runners(h)
-            with pytest.raises(ValueError, match="which they extend"):
-                _set_features(features)
-            assert _runners(h) == {layer: (None, None) for layer in LAYERS}, _runners(h)
-            return
-        _set_features(features)
-        _reset_runners(h)
-        h.gen_step(layout, path, hidden, pos)
-        got = _runners(h)
-        for layer, (want_absorbed_query, want_shared_query) in expect.items():
-            absorbed_query_r, shared_query_r = got[layer]
-            if want_absorbed_query == "absorbed_query":
-                assert isinstance(absorbed_query_r, query_proj.AbsorbedQueryProjection), (
-                    features,
-                    layer,
-                    got,
-                )
-            else:
-                assert absorbed_query_r is want_absorbed_query, (features, layer, got)
-            if want_shared_query == "shared_indexer_absorbed_query":
-                assert isinstance(shared_query_r, shared_indexer.SharedIndexerProjection), (
-                    features,
-                    layer,
-                    got,
-                )
-            else:
-                assert shared_query_r is want_shared_query, (features, layer, got)
-    finally:
-        h.restore(snap)
-        _set_features("")
-        _reset_runners(h)
-
-
-GRAPH_SETS = {
-    "w10only": ("shared_indexer_absorbed_query",),
-    "nativeonly": ("native",),
-    "both": ("shared_indexer_absorbed_query", "native"),
-}
-
-
-@pytest.mark.parametrize("graph_set", list(GRAPH_SETS))
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_shared_query_cuda_graph_replay_into_native_attention(harness, layout, graph_set):
-    """Captured shared_indexer_absorbed_query (+ absorbed_query on layer 6) + native attention + native cache
-    append, replayed once and 20 times with fresh inputs (the absorbed_query harness's
-    graph protocol: warm-up pass, capture into one shared pool, the runner's
-    on_update_kv_lens() inside every step)."""
-    from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-
-    h = harness
-    n_new = math.prod(LAYOUTS[layout])
-    path = "custom_op"
-    snap = h.snapshot()
-    graphs = {}
-    wanted = GRAPH_SETS[graph_set]
-    configs = [
-        (
-            name,
-            "absorbed_query,shared_indexer_absorbed_query"
-            if name == "shared_indexer_absorbed_query"
-            else "",
-        )
-        for name in ("shared_indexer_absorbed_query", "native")
-        if name in wanted
-    ]
-    try:
-        warm = {}
-        for name, features in configs:
-            _set_features(features)
-            _reset_runners(h)
-            meta = h.gen_metadata(layout, cuda_graph=True)
-            static_hidden, static_pos = h.gen_inputs(layout, seed=200)
-            with with_multi_stream(True):
-                for _ in range(2):
-                    h.step(meta, static_pos, static_hidden, path)
-            torch.cuda.synchronize()
-            if name == "shared_indexer_absorbed_query":
-                _assert_engaged(h)
-            warm[name] = (meta, static_hidden, static_pos, _runners(h))
-        h.restore(snap)
-        pool = None
-        for name, features in configs:
-            meta, static_hidden, static_pos, runners = warm[name]
-            _set_features(features)
-            for layer, (absorbed_query_r, shared_query_r) in runners.items():
-                h.attn[layer]._absorbed_query = absorbed_query_r
-                h.attn[layer]._shared_indexer_query = shared_query_r
-            graph = torch.cuda.CUDAGraph()
-            with with_multi_stream(True), aq_integration._Recorder(h.attn) as rec:
-                with torch.cuda.graph(graph, pool=pool):
-                    outs = h.step(meta, static_pos, static_hidden, path)
-            pool = graph.pool()
-            graphs[name] = (graph, meta, static_hidden, outs, rec, runners)
-            torch.cuda.synchronize()
-        _reset_runners(h)
-
-        report, fails = {}, []
-        for trial, (seed, replays) in enumerate([(301, 1), (302, 1), (303, 20)]):
-            hidden, pos = h.gen_inputs(layout, seed=seed)
-            native_outs, native_recs, native_cache = aq_integration._native_runs(
-                h, layout, path, hidden, pos, snap
-            )
-            bounds = aq_integration._output_bounds(native_outs)
-            _set_features("absorbed_query,shared_indexer_absorbed_query")
-            h.restore(snap)
-            shared_query_eager = h.recorded_step(layout, path, hidden, pos)
-            shared_query_eager_cache = h.snapshot()
-            _set_features("")
-            replayed = {}
-            for name, (graph, meta, static_hidden, outs, rec, runners) in graphs.items():
-                h.restore(snap)
-                meta.prepare()
-                static_hidden.copy_(hidden)
-                for _ in range(replays):
-                    graph.replay()
-                torch.cuda.synchronize()
-                replayed[name] = (
-                    {layer: outs[layer].clone() for layer in LAYERS},
-                    rec.snapshot(),
-                    h.snapshot(),
-                )
-            tag = f"t{trial}x{replays}"
-            if "shared_indexer_absorbed_query" in replayed:
-                g_out, g_rec, g_cache = replayed["shared_indexer_absorbed_query"]
-                fails += aq_integration._compare_caches(
-                    snap,
-                    shared_query_eager_cache,
-                    g_cache,
-                    n_new,
-                    report,
-                    f"{tag}:shared_query_graph_vs_eager",
-                )
-                fails += aq_integration._compare_records(
-                    shared_query_eager[1], g_rec, report, f"{tag}:shared_query_graph_vs_eager"
-                )
-                fails += _compare_shared_query_to_native(
-                    snap,
-                    (native_outs[0], native_recs[0], native_cache),
-                    (g_out, g_rec, g_cache),
-                    n_new,
-                    bounds,
-                    report,
-                    f"{tag}:shared_query_graph_vs_native",
-                )
-            if "native" in replayed:
-                g_out, g_rec, g_cache = replayed["native"]
-                fails += aq_integration._compare_caches(
-                    snap, native_cache, g_cache, n_new, report, f"{tag}:native_graph_vs_eager"
-                )
-                fails += aq_integration._compare_records(
-                    native_recs[0], g_rec, report, f"{tag}:native_graph_vs_eager"
-                )
-            for _, (graph, meta, sh, outs, rec, runners) in graphs.items():
-                for absorbed_query_r, shared_query_r in runners.values():
-                    if absorbed_query_r not in (None, False):
-                        assert int(absorbed_query_r.sync.abs().sum()) == 0
-                    if shared_query_r not in (None, False):
-                        assert shared_query_r.counters_quiescent()
-        print(f"graph layout={layout} graphs={wanted} {report}")
-        assert not fails, fails
-    except BaseException as exc:
-        print(f"graph test failure: {exc!r}", flush=True)
-        raise
-    finally:
-        graphs.clear()
-        _reset_runners(h)
-        h.restore(snap)
-        _set_features("")
-
-
-@torch.inference_mode()
-def test_checks_detect_a_wrong_shared_query_absorbed_query(harness, monkeypatch):
-    """Negative control: zeroing the shared_indexer_absorbed_query absorbed query must be flagged by
-    the absorbed-query and quant_q checks of layer 7."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        shared_indexer,
-    )
-
-    h = harness
-    layout, path = "1req_x6tok", "custom_op"
-    hidden, pos = h.gen_inputs(layout, seed=12)
-    snap = h.snapshot()
-    orig = shared_indexer.SharedIndexerProjection.project
-
-    def broken(self, hidden_states):
-        packet = orig(self, hidden_states)
-        packet.fused_q[..., : self.mla.kv_lora_rank].zero_()
-        return packet
-
-    try:
-        native_outs, native_recs, _ = aq_integration._native_runs(
-            h, layout, path, hidden, pos, snap
-        )
-        monkeypatch.setattr(shared_indexer.SharedIndexerProjection, "project", broken)
-        _set_features("absorbed_query,shared_indexer_absorbed_query")
-        _reset_runners(h)
-        _, bad_rec = h.recorded_step(layout, path, hidden, pos)
-        report = {}
-        fails = aq_integration._compare_records(native_recs[0], bad_rec, report, "bad")
-        print(f"negative control {report}")
-        assert any(f"layer {SHARED} absorbed_q" in f for f in fails), fails
-        assert any(f"layer {SHARED} quant_q" in f for f in fails), fails
-        assert not any(f"layer {FULL} absorbed_q" in f for f in fails), fails
-    finally:
-        monkeypatch.undo()
-        h.restore(snap)
-        _set_features("")
-        _reset_runners(h)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py
deleted file mode 100644
index 7aaa78cc645c..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/absorbed_query/test_shared_query_shared_absorb.py
+++ /dev/null
@@ -1,599 +0,0 @@
-# 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.
-"""Contract tests of shared_indexer_absorbed_query (absorbed_query shared-indexer-layer path, preabsorbed form).
-
-Real GLM-5.2 weights of the shared layer 7 (and the full layer 6 for the
-fallback checks) at TP8 rank 0; synthetic activations (unit-RMS rows times
-the layer's real ``input_layernorm`` gain) unless stated.
-
-Acceptance (frozen before tuning): every shared_indexer_absorbed_query output is **bit-identical**
-to the native region it replaces -- ``forward_dsa_proj`` (``q``,
-``compressed_kv``, ``k_pe``, ``latent_cache``) followed by the native
-absorption BMM (``fused_q[..., :512]``) -- because the kernel reproduces the
-native RMSNorm arithmetic and uses absorbed_query's MMA sequences (absorbed_query is held to the
-same bitwise acceptance). The normalized q_a rows the kernel shares between its programs must equal the
-native ``q_a_layernorm`` output bit for bit as well.
-
-Also covered: every kernel configuration (the CuTe-DSL cluster kernel with
-both q exchanges, early PDL trigger on (default) / off and weights before /
-after the dependency wait; the Triton kernel's poll protocols, L2 prefetch,
-W_UK preload, early trigger, rows per norm program, tile shapes and
-orientations, counter protocols), R in {1..8, 16}, padded rows,
-CUDA-graph capture and replay (also with delayed programs), back-to-back and
-multi-stream stress, progress with only a few free SMs, memcpy-fed inputs
-(for compute-sanitizer initcheck), eligibility / fallback (returns
-``None`` before any GPU work), and the package selectors: ``shared_indexer_absorbed_query`` is off
-without ``absorbed_query``, and the absorbed_query form is logged.
-"""
-
-import os
-import sys
-import time
-from types import SimpleNamespace
-
-import absorbed_query_fixture
-import pytest
-import test_absorbed_query_proj as aq_proj
-import torch
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
-
-import plan_testlib  # noqa: E402
-
-SHARED = absorbed_query_fixture.SHARED_LAYER
-FULL = absorbed_query_fixture.FULL_LAYER
-ROWS = aq_proj._env_list(
-    "GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_TEST_ROWS", [1, 2, 3, 4, 5, 6, 7, 8, 16]
-)
-STRESS_ROWS = aq_proj._env_list("GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_TEST_ROWS", [1, 6, 16])
-STRESS_LAUNCHES = int(
-    os.environ.get("GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_STRESS_LAUNCHES", "8")
-)
-
-
-def _triton(**kwargs):
-    return dict(cute=False, **kwargs)
-
-
-# Kernel configurations that must all give the same (native) bits.
-# "default" is the production configuration (CuTe-DSL cluster kernel for
-# R <= 8, Triton kernel for R = 16); the Triton configurations set
-# cute=False.
-CONFIGS = {
-    "default": {},
-    "triton": _triton(),
-    "poll_atomic": _triton(poll_mode=0),
-    "poll_sleep": _triton(poll_mode=2),
-    "no_prefetch": _triton(prefetch_l2=False),
-    "no_preload": _triton(preload_wuk=False),
-    "no_trigger": _triton(early_trigger=False),
-    "norm_rows_2": _triton(norm_rows=2),
-    "norm_rows_8": _triton(norm_rows=8),
-    "norm_scalar": _triton(norm_vectorized=False),
-    "transposed_bn32": _triton(gemm_transposed=True, absorb_transposed=True),
-    "transposed_bn64": _triton(
-        bn=64, gemm_transposed=True, absorb_transposed=True, preload_wuk=False
-    ),
-    "bn64": _triton(bn=64, preload_wuk=False),
-    "epoch_bn16": _triton(bn=16, num_stages=8),
-    "absorbed_query_protocol": _triton(epoch_sync=False),
-    # CuTe-DSL cluster kernel (R <= 8; R = 16 takes the Triton kernel).
-    "cute": dict(cute=True),
-    "cute_no_trigger": dict(cute=True, early_trigger=False),
-    "cute_weights_after_wait": dict(cute=True, cute_weights_before_wait=False),
-    "cute_cluster_barrier": dict(cute=True, cute_push_async=False),
-}
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-
-@pytest.fixture(scope="module")
-def model_config():
-    return absorbed_query_fixture.build_model_config(0)
-
-
-def _runner(attn, **kwargs):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        shared_indexer,
-    )
-
-    return shared_indexer.SharedIndexerProjection(
-        attn, shared_indexer.SharedQueryKernelConfig(**kwargs)
-    )
-
-
-def _assert_packet_native(packet, native, kv_lora, tag, rows=slice(None)):
-    """shared_indexer_absorbed_query packet vs the native region: bit for bit."""
-    native_outs, native_fq = native
-    names = ("q", "compressed_kv", "k_pe", "latent_cache")
-    assert len(packet.proj_outputs) == 4, tag
-    for name, x, y in zip(names, packet.proj_outputs, native_outs):
-        assert x.shape == y.shape and x.dtype == y.dtype, (tag, name)
-        aq_proj._assert_bitwise(f"{tag} {name}", x[rows], y[rows])
-    assert packet.fused_q is not None
-    aq_proj._assert_bitwise(
-        f"{tag} absorbed", packet.fused_q[rows, :, :kv_lora], native_fq[rows, :, :kv_lora]
-    )
-
-
-def _assert_same(a, b, kv_lora, tag):
-    for i, (x, y) in enumerate(zip(a.proj_outputs, b.proj_outputs)):
-        aq_proj._assert_bitwise(f"{tag} out{i}", x, y)
-    aq_proj._assert_bitwise(f"{tag} absorbed", a.fused_q[..., :kv_lora], b.fused_q[..., :kv_lora])
-
-
-def _inputs(rows, seed=0, scale=1.0):
-    hidden, position_ids = aq_proj._inputs(SHARED, rows, seed=seed)
-    if scale != 1.0:
-        hidden = (hidden.float() * scale).to(torch.bfloat16)
-    return hidden, position_ids
-
-
-@pytest.mark.parametrize("kind", ["real_q_a", "real_ckv", "wide"])
-@pytest.mark.parametrize("rows", [1, 6, 64, 4096])
-@torch.inference_mode()
-def test_norm_reproduction_is_bitwise(model_config, kind, rows):
-    """The Triton reproduction of the native RMSNorm (sequential per-lane
-    fold) equals the native module call bit for bit, on strided kv_a views
-    and on wide-dynamic-range rows."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        shared_indexer,
-    )
-
-    attn = aq_proj._attn(model_config, SHARED)
-    gen = torch.Generator(device="cuda").manual_seed(rows)
-    if kind == "wide":
-        mag = torch.exp(torch.empty(rows, 2048, device="cuda").uniform_(-6, 6, generator=gen))
-        x = (torch.randn(rows, 2048, device="cuda", generator=gen) * mag).to(torch.bfloat16)
-        norm = attn.q_a_layernorm
-    else:
-        hidden = absorbed_query_fixture.make_hidden(rows, seed=rows, gamma=aq_proj._gamma(SHARED))
-        kva = attn.kv_a_proj_with_mqa(hidden)
-        if kind == "real_q_a":
-            x, norm = kva[:, :2048], attn.q_a_layernorm
-        else:
-            x, norm = kva[:, 2048:2560], attn.kv_a_layernorm
-    native = norm(x)
-    for rpp in (1, 4, 8):
-        for vec in (False, True):
-            rep = shared_indexer.fi_rmsnorm_reproduction(
-                x, norm.weight, norm.variance_epsilon, order=0, rows_per_program=rpp, vectorized=vec
-            )
-            aq_proj._assert_bitwise(f"{kind} rows={rows} rpp={rpp} vec={vec}", rep, native)
-
-
-@pytest.mark.parametrize("rows", ROWS)
-@pytest.mark.parametrize("config_name", list(CONFIGS))
-@torch.inference_mode()
-def test_contract_matches_native_bitwise(model_config, rows, config_name):
-    """Every output of the shared_indexer_absorbed_query packet equals the native region bit for bit
-    for every configuration; the shared normalized rows equal the native
-    q_a_layernorm output; ``q`` equals absorbed_query's; the counters return to zero."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-    )
-
-    if config_name not in ("default", "triton") and rows not in (1, 6, 16):
-        pytest.skip("non-default configurations run at R in {1, 6, 16}")
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn, debug_keep_qn=True, **CONFIGS[config_name])
-    hidden, position_ids = _inputs(rows, seed=rows)
-    packet = runner.project(hidden)
-    native = aq_proj._native(attn, hidden, position_ids)
-    kva = attn.kv_a_proj_with_mqa(hidden)
-    native_qn = attn.q_a_layernorm(kva[:, :2048])
-    aq_packet = AbsorbedQueryProjection(attn).project(position_ids, hidden, absorb=False)
-    torch.cuda.synchronize()
-    _assert_packet_native(packet, native, attn.kv_lora_rank, f"R={rows} {config_name}")
-    aq_proj._assert_bitwise("qn", runner.debug["qn"], native_qn)
-    aq_proj._assert_bitwise(
-        "q vs absorbed_query", packet.proj_outputs[0], aq_packet.proj_outputs[0]
-    )
-    # Packet layout: compressed_kv / k_pe are views of latent_cache.
-    q, ckv, k_pe, latent = packet.proj_outputs
-    assert ckv.data_ptr() == latent.data_ptr()
-    assert k_pe.data_ptr() == latent.data_ptr() + 512 * latent.element_size()
-    assert packet.fused_q.shape == (
-        rows,
-        attn.num_heads_tp,
-        attn.kv_lora_rank + attn.qk_rope_head_dim,
-    )
-    assert runner.counters_quiescent(), runner.epochs
-
-
-@pytest.mark.parametrize("rows", [1, 2, 3, 4, 5, 6, 7, 8])
-@torch.inference_mode()
-def test_cute_kernel_all_rows_back_to_back(model_config, rows):
-    """The CuTe-DSL cluster kernel at every R it serves: 8 back-to-back
-    launches with fresh inputs (also scaled), each bit-identical to the
-    native region, the shared normalized rows equal to the native
-    q_a_layernorm output."""
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn, cute=True, debug_keep_qn=True)
-    launches = []
-    for i in range(8):
-        hidden, position_ids = _inputs(rows, seed=4000 + i, scale=(1.0, 1e-3, 30.0)[i % 3])
-        packet = runner.project(hidden)
-        qn = runner.debug["qn"].clone()
-        kva = attn.kv_a_proj_with_mqa(hidden)
-        launches.append(
-            (
-                packet,
-                qn,
-                attn.q_a_layernorm(kva[:, :2048]),
-                aq_proj._native(attn, hidden, position_ids),
-            )
-        )
-    torch.cuda.synchronize()
-    for i, (packet, qn, native_qn, native) in enumerate(launches):
-        aq_proj._assert_bitwise(f"cute qn {i}", qn, native_qn)
-        _assert_packet_native(packet, native, attn.kv_lora_rank, f"cute launch {i}")
-
-
-@pytest.mark.parametrize("rows", [6, 16])
-@pytest.mark.parametrize("scale", [1.0, 1e-3, 30.0])
-@torch.inference_mode()
-def test_native_agreement_statistics(model_config, rows, scale):
-    """48 fresh inputs per case (also scaled hidden states, which move the
-    RMSNorm sums across many rounding boundaries): zero differing elements
-    in every output; absorbed_query's own compressed_kv norm is reported for
-    comparison."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-    )
-
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn)
-    aq_projection = AbsorbedQueryProjection(attn)
-    absorbed_query_ckv_diff = 0
-    for i in range(48):
-        hidden, position_ids = _inputs(rows, seed=3000 + i, scale=scale)
-        packet = runner.project(hidden)
-        native = aq_proj._native(attn, hidden, position_ids)
-        aq_packet = aq_projection.project(position_ids, hidden, absorb=False)
-        torch.cuda.synchronize()
-        _assert_packet_native(packet, native, attn.kv_lora_rank, f"input {i}")
-        absorbed_query_ckv_diff += int(
-            (aq_proj._steps(aq_packet.proj_outputs[1], native[0][1]) != 0).sum()
-        )
-    print(
-        f"rows={rows} scale={scale}: shared_indexer_absorbed_query 0 differing elements in 48 "
-        f"inputs; absorbed_query compressed_kv differing elements {absorbed_query_ckv_diff}"
-    )
-
-
-@pytest.mark.parametrize("cute", [False, True], ids=["triton", "cute"])
-@torch.inference_mode()
-def test_padded_rows_do_not_affect_live_rows(model_config, cute):
-    """Rows are independent: padded rows (zero or 1e3 garbage) never change a
-    live row, live rows equal native bit for bit, zero rows stay finite."""
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn, cute=cute)
-    live = slice(0, 3)
-    hidden_zero, position_ids = aq_proj._inputs(SHARED, 6, padded_from=3)
-    hidden_junk = hidden_zero.clone()
-    hidden_junk[3:] = 1e3 * torch.randn_like(hidden_junk[3:])
-    a = runner.project(hidden_zero)
-    b = runner.project(hidden_junk)
-    native = aq_proj._native(attn, hidden_junk, position_ids)
-    torch.cuda.synchronize()
-    for i, (x, y) in enumerate(zip(a.proj_outputs, b.proj_outputs)):
-        aq_proj._assert_bitwise(f"padded out{i}", x[live], y[live])
-    aq_proj._assert_bitwise(
-        "padded absorbed",
-        a.fused_q[live, :, : attn.kv_lora_rank],
-        b.fused_q[live, :, : attn.kv_lora_rank],
-    )
-    _assert_packet_native(b, native, attn.kv_lora_rank, "padded", live)
-    for x in a.proj_outputs:
-        assert torch.isfinite(x[3:].float()).all()
-    assert runner.counters_quiescent(), runner.epochs
-
-
-@pytest.mark.parametrize(
-    "delay_ns,cute",
-    [(0, False), (20000, False), (0, True)],
-    ids=["triton", "triton_delayed", "cute"],
-)
-@pytest.mark.parametrize("rows", [1, 6])
-@torch.inference_mode()
-def test_cuda_graph_replay(model_config, delay_ns, cute, rows):
-    """One and many replays of a captured shared_indexer_absorbed_query call (native kv_a GEMM +
-    shared_indexer_absorbed_query kernel, PDL-chained) with fresh inputs equal the native region bit
-    for bit and leave the counters at zero."""
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn, debug_delay_ns=delay_ns, cute=cute)
-    static_hidden, _ = _inputs(rows)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(2):
-            runner.project(static_hidden)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        packet = runner.project(static_hidden)
-    for step, replays in enumerate([1, 1, 50]):
-        hidden, position_ids = _inputs(rows, seed=100 + step)
-        static_hidden.copy_(hidden)
-        for _ in range(replays):
-            graph.replay()
-        torch.cuda.synchronize()
-        native = aq_proj._native(attn, hidden, position_ids)
-        torch.cuda.synchronize()
-        _assert_packet_native(packet, native, attn.kv_lora_rank, f"replay x{replays}")
-        assert runner.counters_quiescent(), runner.epochs
-
-
-@pytest.mark.parametrize("rows", STRESS_ROWS)
-@pytest.mark.parametrize(
-    "config_name", ["triton", "poll_atomic", "no_preload", "norm_rows_2", "absorbed_query_protocol"]
-)
-@torch.inference_mode()
-def test_delayed_programs_back_to_back(model_config, rows, config_name):
-    """Norm programs delayed before they release, producers before they
-    publish, absorption jobs before they count out (up to 50 us,
-    deterministic per program): back-to-back launches with fresh inputs
-    equal the native region bit for bit; counters end at zero."""
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn, debug_delay_ns=50000, **CONFIGS[config_name])
-    launches = []
-    for i in range(STRESS_LAUNCHES):
-        hidden, position_ids = _inputs(rows, seed=500 + i)
-        launches.append((runner.project(hidden), aq_proj._native(attn, hidden, position_ids)))
-    torch.cuda.synchronize()
-    for i, (packet, native) in enumerate(launches):
-        _assert_packet_native(packet, native, attn.kv_lora_rank, f"delayed launch {i}")
-    assert runner.counters_quiescent(), runner.epochs
-
-
-@pytest.mark.parametrize(
-    "config_name",
-    ["default", "triton", "poll_atomic", "absorbed_query_protocol", "cute_cluster_barrier"],
-)
-@torch.inference_mode()
-def test_concurrent_launches_on_many_streams(model_config, config_name):
-    """Six shared_indexer_absorbed_query instances (own counters, like six layers) concurrently on six
-    streams, 30 rounds, with and without delays: each stream's last result
-    equals the native region bit for bit; no hang; counters at zero."""
-    attn = aq_proj._attn(model_config, SHARED)
-    n_streams, rounds, rows = 6, 30, 6
-    streams = [torch.cuda.Stream() for _ in range(n_streams)]
-    cute = CONFIGS[config_name].get("cute", False)
-    runners = [
-        _runner(attn, debug_delay_ns=(5000 if s % 2 and not cute else 0), **CONFIGS[config_name])
-        for s in range(n_streams)
-    ]
-    inputs = [_inputs(rows, seed=700 + s) for s in range(n_streams)]
-    natives = [aq_proj._native(attn, h, p) for h, p in inputs]
-    torch.cuda.synchronize()
-    last = [None] * n_streams
-    for _ in range(rounds):
-        for s in range(n_streams):
-            with torch.cuda.stream(streams[s]):
-                last[s] = runners[s].project(inputs[s][0])
-    torch.cuda.synchronize()
-    for s in range(n_streams):
-        _assert_packet_native(last[s], natives[s], attn.kv_lora_rank, f"stream {s}")
-        assert runners[s].counters_quiescent()
-
-
-@pytest.mark.parametrize("free_sms", [1, 2, 4, 8])
-@torch.inference_mode()
-def test_progress_with_restricted_sms(model_config, free_sms):
-    """A concurrent kernel holds all but ``free_sms`` SMs for 0.5 s (no shared_indexer_absorbed_query
-    program fits beside one of its programs). The shared_indexer_absorbed_query kernel must complete
-    while those SMs are still held, with native bits: every spin-wait
-    targets lower program ids (norm programs, then producers, then the
-    extra absorption programs). Triton kernel only: the CuTe kernel's
-    clusters are co-scheduled by the hardware and do not wait on each
-    other."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        shared_indexer,
-    )
-
-    attn = aq_proj._attn(model_config, SHARED)
-    runner = _runner(attn, cute=False)
-    rows = 6
-    hidden, position_ids = _inputs(rows, seed=900)
-    native = aq_proj._native(attn, hidden, position_ids)
-    kva = attn.kv_a_proj_with_mqa(hidden)
-    reference = runner.launch(kva)
-    torch.cuda.synchronize()
-
-    n_sm = torch.cuda.get_device_properties(0).multi_processor_count
-    n_occ = n_sm - free_sms
-    bm = bn = 128
-    bk, iters = 64, 32
-    a = torch.randn(bm, iters * bk, device="cuda").to(torch.bfloat16)
-    b = torch.randn(iters * bk, bn, device="cuda").to(torch.bfloat16)
-    out = torch.empty(n_occ, bm, bn, device="cuda")
-    smids = torch.full((n_occ,), -1, dtype=torch.int32, device="cuda")
-    occ_stream, shared_query_stream = torch.cuda.Stream(), torch.cuda.Stream()
-    ev = {k: torch.cuda.Event(enable_timing=True) for k in ("s", "o", "p")}
-    torch.cuda.synchronize()
-    with torch.cuda.stream(occ_stream):
-        ev["s"].record()
-        occ = aq_proj._occupier_kernel[(n_occ,)](
-            a, b, out, smids, 500_000_000, iters, bm, bn, bk, 6, num_warps=4, num_stages=6
-        )
-        ev["o"].record()
-    time.sleep(0.05)
-    with torch.cuda.stream(shared_query_stream):
-        result = runner.launch(kva)
-        ev["p"].record()
-    torch.cuda.synchronize()
-    t_shared_query, t_occ = ev["s"].elapsed_time(ev["p"]), ev["s"].elapsed_time(ev["o"])
-    shared_query_smem = min(
-        aq_proj._compiled_shared_bytes(shared_indexer.shared_query_kernel) or [0]
-    )
-    held = int(torch.unique(smids).numel())
-    report = dict(
-        free_sms=free_sms,
-        held_sms=held,
-        occupier_smem=occ.metadata.shared,
-        shared_query_min_smem=shared_query_smem,
-        shared_query_done_ms=round(t_shared_query, 3),
-        occupier_done_ms=round(t_occ, 3),
-    )
-    print(f"restricted SMs {report}")
-    assert held == n_occ, report
-    assert occ.metadata.shared + shared_query_smem > 228 * 1024, report
-    for name, x, y in zip(("q", "latent", "fused_q"), result, reference):
-        if name == "fused_q":
-            x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
-        aq_proj._assert_bitwise(f"restricted {name}", x, y)
-    aq_proj._assert_bitwise("restricted q vs native", result[0], native[0][0])
-    assert runner.counters_quiescent(), runner.epochs
-    assert t_shared_query < t_occ - 100.0, (
-        f"shared_indexer_absorbed_query did not complete while {n_occ} SMs were held: {report}"
-    )
-
-
-@pytest.mark.parametrize("cute", [False, True], ids=["triton", "cute"])
-@pytest.mark.parametrize("rows", STRESS_ROWS)
-@torch.inference_mode()
-def test_kernel_with_memcpy_initialized_inputs(model_config, rows, cute):
-    """The shared_indexer_absorbed_query kernel alone with every tensor it reads placed by a
-    host-to-device copy (kv_a output, weights, norm gains, counters) equals
-    the normal launch bit for bit. ``GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_INITCHECK_ONLY=1``
-    (initcheck run) skips the reference launch so that only memcpy-fed memory is read by
-    instrumented kernels."""
-    initcheck_only = (
-        os.environ.get("GLM_MOE_DSA_SHARED_INDEXER_ABSORBED_QUERY_INITCHECK_ONLY") == "1"
-    )
-    attn = aq_proj._attn(model_config, SHARED)
-    hidden, _ = _inputs(rows, seed=1200)
-    kva = attn.kv_a_proj_with_mqa(hidden)
-    reference = None
-    if not initcheck_only:
-        reference = _runner(attn, cute=cute).launch(kva)
-    torch.cuda.synchronize()
-    probe = _runner(attn, cute=cute)
-    probe.sync = aq_proj._memcpy_copy(torch.zeros_like(probe.sync))
-    probe.epochs = aq_proj._memcpy_copy(torch.zeros_like(probe.epochs))
-    kva_copy = aq_proj._memcpy_copy(kva)
-    params = [
-        attn.q_b_proj.weight,
-        attn.k_b_proj_trans,
-        attn.kv_a_layernorm.weight,
-        attn.q_a_layernorm.weight,
-    ]
-    saved = [p.data for p in params]
-    try:
-        for p in params:
-            p.data = aq_proj._memcpy_copy(p.data)
-        torch.cuda.synchronize()
-        result = probe.launch(kva_copy)
-        torch.cuda.synchronize()
-    finally:
-        for p, data in zip(params, saved):
-            p.data = data
-    assert probe.counters_quiescent()
-    if initcheck_only:
-        return
-    for name, x, y in zip(("q", "latent", "fused_q"), result, reference):
-        if name == "fused_q":
-            x, y = x[..., : attn.kv_lora_rank], y[..., : attn.kv_lora_rank]
-        aq_proj._assert_bitwise(f"memcpy inputs {name}", x, y)
-
-
-@torch.inference_mode()
-def test_eligibility_and_fallback(model_config, monkeypatch):
-    """``shared_dsa_proj`` returns None before any GPU work outside its
-    domain (full-indexer and draft layers, context / mixed / missing
-    metadata, R > 16, non-BF16 input, force-native switch) and the packet
-    of a generation-only call on a shared layer otherwise."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        shared_indexer,
-    )
-
-    shared = aq_proj._attn(model_config, SHARED)
-    full = aq_proj._attn(model_config, FULL)
-    gen = SimpleNamespace(num_contexts=0, num_generations=1)
-    ctx = SimpleNamespace(num_contexts=1, num_generations=0)
-    mixed = SimpleNamespace(num_contexts=1, num_generations=1)
-    hidden, position_ids = _inputs(6)
-
-    assert shared_indexer.shared_module_eligibility(shared)[0]
-    ok, reason = shared_indexer.shared_module_eligibility(full)
-    assert not ok and "full indexer" in reason
-    full._shared_indexer_query = None
-    assert shared_indexer.shared_dsa_proj(full, position_ids, hidden, gen) is None
-    assert full._shared_indexer_query is False
-    for meta in (None, ctx, mixed):
-        shared._shared_indexer_query = None
-        assert shared_indexer.shared_dsa_proj(shared, position_ids, hidden, meta) is None
-    big = torch.zeros(17, hidden.shape[1], dtype=hidden.dtype, device="cuda")
-    assert shared_indexer.shared_dsa_proj(shared, None, big, gen) is None
-    assert shared_indexer.shared_dsa_proj(shared, None, hidden.float(), gen) is None
-    saved = shared.layer_idx
-    try:
-        shared.layer_idx = shared.num_hidden_layers
-        assert not shared_indexer.shared_module_eligibility(shared)[0]
-    finally:
-        shared.layer_idx = saved
-    shared._shared_indexer_query = None
-    packet = shared_indexer.shared_dsa_proj(shared, position_ids, hidden, gen)
-    assert packet is not None and len(packet.proj_outputs) == 4
-    assert packet.fused_q is not None
-    _assert_packet_native(
-        packet, aq_proj._native(shared, hidden, position_ids), shared.kv_lora_rank, "entry point"
-    )
-    shared._shared_indexer_query = None
-
-
-def test_package_selectors():
-    """``shared_indexer_absorbed_query`` needs ``absorbed_query`` (the plan turns it off with
-    ``absorbed_query``, so it never runs silently native); the absorbed_query form is
-    logged once next to the feature set."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    assert dispatch.REQUIRED_FEATURES["shared_indexer_absorbed_query"] == {"absorbed_query"}
-    plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: "absorbed_query"})
-    assert not plan.enabled("shared_indexer_absorbed_query")
-    assert plan_testlib.plan_running(
-        dispatch, "absorbed_query,shared_indexer_absorbed_query"
-    ).enabled("shared_indexer_absorbed_query")
-    try:
-        dispatch._CAPABILITY_LOGGED.pop(("absorbed_query", "form"), None)
-        assert query_proj.check_package_selectors() == query_proj.SERVED_FORM == "auto"
-        assert ("absorbed_query", "form") in dispatch._CAPABILITY_LOGGED
-    finally:
-        dispatch._CAPABILITY_LOGGED.pop(("absorbed_query", "form"), None)
-
-
-def test_registered_entry_point():
-    """``shared_dsa_proj`` is reachable through the dispatch registry under
-    the absorbed_query package, and the absorbed_query / draft_absorbed_query entries are
-    unchanged."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-        shared_indexer,
-    )
-
-    assert dispatch.get_impl("absorbed_query", "shared_dsa_proj") is shared_indexer.shared_dsa_proj
-    assert dispatch.get_impl("absorbed_query", "dsa_proj") is query_proj.dsa_proj
-    assert dispatch.get_impl("absorbed_query", "mtp_dsa_proj") is query_proj.mtp_dsa_proj
-    assert "shared_indexer_absorbed_query" in dispatch.SELECTOR_PATHS
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py
deleted file mode 100644
index 7f384ebceb21..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/real_mla_harness.py
+++ /dev/null
@@ -1,485 +0,0 @@
-# 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.
-"""Real GLM-5.2 DSA attention layers for attention_output_tail end-to-end tests and benchmarks.
-
-``AttentionStack`` builds, for one TP8 rank, the attention module the model
-builds for ``glm_moe_dsa`` checkpoints (``DeepseekV32Attention``) for a run of
-consecutive layers. It loads every attention weight of the rank's shard from
-the checkpoint the way ``DeepseekV3WeightLoader`` does, fills an FP8
-``DSACacheManager`` with a real prefill forward, and prepares GEN-only
-attention metadata for one request with ``rows`` query tokens (MTP target
-verification with ``rows - 1`` draft tokens) the way the decoder runner does:
-CUDA-graph metadata, ``update_spec_dec_param``, then ``prepare``.
-
-Each module is adopted as the target's ``GlmMLA`` right after construction, as the target's
-model adopts it, and the GEN forward runs as a decode step (``scope()``), so ``GlmMLA.forward``
-then runs exactly as in the model: ``register_to_config`` is set, so the GEN attention runs
-inside the target's DSA attention op and ``o_proj`` runs in ``GlmMLA.forward``. With a plan
-running ``attention_output_tail`` installed at construction (``fast_paths``, or ``select``),
-every target layer owns an attention_output_tail ``AttnTail``; ``set_attn_tail`` swaps it in and out to compare
-against the native tail on the same module, cache and inputs. The prefill runs as a prefill
-step (the built-in attention).
-
-Fixture labels: the weights (all attention projections, norms and indexer
-weights of the rank's shard) are real. The activations are synthetic: the
-prefill and GEN inputs are the layer's real ``input_layernorm`` applied to
-real token embeddings, not the residual stream of a full forward.
-
-The cache manager and the attention metadata use a one-rank mapping: an MLA
-cache holds one latent head, so it is not TP-sharded, and a one-rank mapping
-keeps their constructors free of collectives in a process that stands in for
-one of eight TP ranks.
-"""
-
-import contextlib
-import os
-import sys
-import weakref
-from typing import Dict, List, Optional
-
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-import tail_fixtures  # noqa: E402
-
-TP8 = 8
-REQUEST_ID = 7
-TOKEN_OFFSET = 1000
-
-
-def load_pretrained(ckpt: str):
-    """Load the checkpoint's config as the model does for ``glm_moe_dsa``."""
-    from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config
-
-    cfg = load_pretrained_config(ckpt, trust_remote_code=True)
-    # DeepseekV3ForCausalLM renames glm_moe_dsa to deepseek_v32, which makes
-    # the decoder layers build DeepseekV32Attention.
-    cfg.model_type = "deepseek_v32"
-    dtype = getattr(cfg, "torch_dtype", None)
-    if not isinstance(dtype, torch.dtype):
-        resolved = getattr(torch, str(dtype or "bfloat16").replace("torch.", ""))
-        try:
-            cfg.torch_dtype = resolved
-        except AttributeError:
-            cfg.dtype = resolved
-    return cfg
-
-
-def make_model_config(pretrained, tp_rank: int, max_draft_len: int):
-    """``ModelConfig`` of the served GEN role for TP8 rank ``tp_rank``.
-
-    Attention weights of this checkpoint are excluded from NVFP4, so the
-    module is built unquantized with an FP8 KV cache, as the model's
-    exclusion handling leaves it.
-    """
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, MTPDecodingConfig
-    from tensorrt_llm.mapping import Mapping
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-    from tensorrt_llm.quantization.mode import QuantAlgo
-
-    quant = QuantConfig()
-    quant.kv_cache_quant_algo = QuantAlgo.FP8
-    sparse = DeepSeekSparseAttentionConfig(
-        enable_heuristic_topk=True, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=True
-    )
-    spec = MTPDecodingConfig(max_draft_len=max_draft_len) if max_draft_len > 0 else None
-    model_config = ModelConfig(
-        pretrained_config=pretrained,
-        mapping=Mapping(world_size=TP8, tp_size=TP8, rank=tp_rank),
-        quant_config=quant,
-        sparse_attention_config=sparse,
-        spec_config=spec,
-    )
-    model_config.extra_attrs["kv_cache_dtype"] = "fp8"
-    return model_config
-
-
-def build_attention(model_config, layer_idx: int, aux_stream, reduce_output: bool):
-    """Build ``DeepseekV32Attention`` for one layer on the current device, adopted as ``GlmMLA``.
-
-    The module has TP8 shapes, but this process may run without the other
-    ranks, so the all-reduce modules its ``Linear``s would create are
-    replaced by ``None`` during construction; callers that reduce attach a
-    real ``AllReduce`` afterwards.
-    """
-    import tensorrt_llm._torch.distributed as distributed
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-        GlmMLA,
-    )
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
-    from tensorrt_llm._torch.utils import AuxStreamType
-
-    saved = distributed.AllReduce
-    distributed.AllReduce = lambda *args, **kwargs: None
-    try:
-        with torch.device("cuda"):
-            attn = DeepseekV32Attention(
-                model_config,
-                layer_idx=layer_idx,
-                aux_stream_dict={AuxStreamType.Attention: aux_stream},
-                reduce_output=reduce_output,
-            )
-            GlmMLA.adopt(attn)
-            # The model builds the FMHA libraries in ``create_weights`` after adoption; a module
-            # built alone rebuilds them for the library adoption declared.
-            for backend in (attn.mha, attn.mqa):
-                if backend is not None and backend.prepended_fmha_lib_classes:
-                    backend.update_quant_config(backend.quant_config)
-    finally:
-        distributed.AllReduce = saved
-    return attn.cuda()
-
-
-def load_attention_weights(attn, ckpt: str, layer: int, tp_rank: int, tp_size: int = TP8) -> None:
-    """Load the attention weights of ``layer`` like ``DeepseekV3WeightLoader``."""
-    from tensorrt_llm._torch.models.modeling_utils import filter_weights
-
-    prefix = f"model.layers.{layer}.self_attn"
-    weights = tail_fixtures.read_prefix(ckpt, prefix + ".")
-    heads = attn.num_heads
-    nope, v_dim, kv_lora = (attn.qk_nope_head_dim, attn.v_head_dim, attn.kv_lora_rank)
-    with torch.no_grad():
-        for name, module in attn.named_modules():
-            if not name or len(module._parameters) == 0:
-                continue
-            full = f"{prefix}.{name}"
-            leaf = name.split(".")[-1]
-            if leaf == "kv_b_proj":
-                kv_b = weights[f"{full}.weight"].unflatten(0, [heads, nope + v_dim])
-                kv_b = torch.chunk(kv_b, tp_size, dim=0)[tp_rank]
-                k_nope, v_w = kv_b.split([nope, v_dim], dim=1)
-                local = kv_b.shape[0]
-                module.weight.data.copy_(
-                    torch.cat(
-                        [k_nope.reshape(local * nope, kv_lora), v_w.reshape(local * v_dim, kv_lora)]
-                    )
-                )
-                attn.v_b_proj = torch.nn.Parameter(
-                    module.weight.data[local * nope :].view(local, v_dim, kv_lora),
-                    requires_grad=False,
-                )
-                attn.k_b_proj_trans.data.copy_(
-                    k_nope.transpose(2, 1).contiguous().reshape(attn.k_b_proj_trans.shape)
-                )
-            elif leaf == "kv_a_proj_with_mqa":
-                fused = torch.cat([weights[f"{prefix}.q_a_proj.weight"], weights[f"{full}.weight"]])
-                module.weight.data[: fused.shape[0]].copy_(fused)
-            else:
-                module_weights = filter_weights(full, weights)
-                if hasattr(module, "load_weights"):
-                    module.load_weights(weights=[module_weights])
-                else:
-                    for n, p in module.named_parameters():
-                        p.data.copy_(module_weights[n][:])
-        for m in attn.modules():
-            post = getattr(m, "post_load_weights", None)
-            if post is not None and not getattr(m, "_weights_removed", False):
-                post()
-
-
-def layer_inputs(ckpt: str, layer: int, first_token: int, num_tokens: int) -> torch.Tensor:
-    """Real ``input_layernorm`` of ``layer`` over real token embeddings."""
-    return tail_fixtures.normed_embeddings(ckpt, layer, first_token, num_tokens)
-
-
-def select(fast_paths: str) -> None:
-    """Install, for this process, the plan that runs exactly ``fast_paths`` (comma-separated)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-        dispatch,
-    )
-
-    plan_testlib.select(dispatch, fast_paths)
-
-
-class AttentionStack:
-    """Consecutive real DSA attention layers of one TP8 rank with a filled cache.
-
-    Args:
-        ckpt: checkpoint directory.
-        layers: consecutive layer indices; the first must run a full indexer.
-        tp_rank: TP8 rank whose weight shard is loaded.
-        ctx: cached context length per layer (filled by a real prefill).
-        rows: GEN query rows of the single request (``1 + max_draft_len``).
-        reduce_output: build ``o_proj`` with ``reduce_output`` (dense layers).
-        fast_paths: when given, the plan that runs exactly these fast paths is installed
-            (``select``) before the layers are built; otherwise the installed plan is used.
-    """
-
-    def __init__(
-        self,
-        ckpt: str,
-        layers: List[int],
-        tp_rank: int,
-        ctx: int = 4096,
-        rows: int = 6,
-        reduce_output: bool = False,
-        fast_paths: Optional[str] = None,
-    ):
-        from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
-
-        if fast_paths is not None:
-            select(fast_paths)
-
-        self.ckpt = ckpt
-        self.layers = list(layers)
-        self.tp_rank = tp_rank
-        self.ctx = ctx
-        self.rows = rows
-        self.max_draft_len = rows - 1
-        self.pretrained = load_pretrained(ckpt)
-        self.model_config = make_model_config(self.pretrained, tp_rank, self.max_draft_len)
-        sparse = self.model_config.sparse_attention_config
-        self.sparse_params = sparse.to_sparse_params(pretrained_config=self.pretrained)
-        self.sparse_metadata_params = sparse.to_sparse_metadata_params(
-            pretrained_config=self.pretrained
-        )
-        self.backend = get_attention_backend("TRTLLM", sparse_params=self.sparse_params)
-        self.aux_stream = torch.cuda.Stream()
-        self.attn: Dict[int, torch.nn.Module] = {}
-        for layer in self.layers:
-            attn = build_attention(self.model_config, layer, self.aux_stream, reduce_output)
-            load_attention_weights(attn, ckpt, layer, tp_rank)
-            self.attn[layer] = attn
-        self.attn_tail_tails = {
-            layer: getattr(attn, "_attn_tail", None) for layer, attn in self.attn.items()
-        }
-        self.kv_cache_manager = self._make_cache()
-        self.kv_cache_manager.add_dummy_requests(
-            request_ids=[REQUEST_ID], token_nums=[ctx + 64], is_gen=False, prepare_resource=True
-        )
-        self._prefill()
-        self.gen_metadata = self._gen_metadata()
-
-    # -- construction --------------------------------------------------------
-    def _one_rank_mapping(self):
-        from tensorrt_llm.mapping import Mapping
-
-        return Mapping(world_size=1, tp_size=1, rank=0)
-
-    def _make_cache(self):
-        import tensorrt_llm.bindings
-        from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
-        from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
-        from tensorrt_llm.llmapi.llm_args import KvCacheConfig
-
-        num_layers = max(self.layers) + 1
-        layer_mask = [i in self.layers for i in range(num_layers)]
-        max_tokens = self.ctx + 256
-        return DSACacheManager(
-            KvCacheConfig(
-                dtype="fp8",
-                max_tokens=max_tokens,
-                enable_block_reuse=False,
-                event_buffer_max_size=0,
-            ),
-            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-            num_layers=len(self.layers),
-            layer_mask=layer_mask,
-            num_kv_heads=1,
-            head_dim=self.pretrained.kv_lora_rank + self.pretrained.qk_rope_head_dim,
-            tokens_per_block=64,
-            max_seq_len=max_tokens,
-            max_batch_size=1,
-            mapping=self._one_rank_mapping(),
-            dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
-            max_num_tokens=max(8192, self.ctx),
-            sparse_attention_config=self.model_config.sparse_attention_config,
-            model_config=self.model_config,
-            pretrained_config=self.pretrained,
-            vocab_size=self.pretrained.vocab_size,
-        )
-
-    def _prefill(self) -> None:
-        from tensorrt_llm._torch.metadata import KVCacheParams
-
-        meta = self.backend.Metadata(
-            seq_lens=torch.tensor([self.ctx], dtype=torch.int),
-            request_ids=[REQUEST_ID],
-            max_num_requests=1,
-            num_contexts=1,
-            prompt_lens=[self.ctx],
-            max_num_tokens=self.ctx,
-            kv_cache_manager=self.kv_cache_manager,
-            kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]),
-            mapping=self._one_rank_mapping(),
-            sparse_metadata_params=self.sparse_metadata_params,
-        )
-        meta.prepare()
-        position_ids = torch.arange(self.ctx, dtype=torch.int32, device="cuda")[None]
-        from tensorrt_llm._torch.distributed import AllReduceParams
-
-        no_ar = AllReduceParams(enable_allreduce=False)
-        with self.scope(meta):
-            for layer in self.layers:
-                hidden = layer_inputs(self.ckpt, layer, TOKEN_OFFSET, self.ctx)
-                self.set_attn_tail(False)
-                self.attn[layer](
-                    position_ids=position_ids,
-                    hidden_states=hidden,
-                    attn_metadata=meta,
-                    all_reduce_params=no_ar,
-                )
-        torch.cuda.synchronize()
-
-    def _gen_metadata(self):
-        from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures
-        from tensorrt_llm._torch.metadata import KVCacheParams
-
-        base = self.backend.Metadata(
-            max_num_requests=1,
-            max_num_tokens=128,
-            max_num_sequences=1,
-            kv_cache_manager=self.kv_cache_manager,
-            mapping=self._one_rank_mapping(),
-            runtime_features=AttentionRuntimeFeatures(
-                has_speculative_draft_tokens=self.max_draft_len > 0
-            ),
-            enable_flash_mla=False,
-            enable_context_mla_with_cached_kv=False,
-            num_heads_per_kv=1,
-            sparse_metadata_params=self.sparse_metadata_params,
-        )
-        meta = base.create_cuda_graph_metadata(1, False, self.max_draft_len)
-        meta.beam_width = 1
-        meta.request_ids = [REQUEST_ID]
-        meta.prompt_lens = [self.ctx]
-        meta.num_contexts = 0
-        meta.num_chunked_ctx_requests = 0
-        meta.kv_cache_params = KVCacheParams(
-            use_cache=True,
-            num_cached_tokens_per_seq=[self.ctx],
-            num_extra_kv_tokens=max(self.max_draft_len - 1, 0),
-        )
-        meta.kv_cache_manager = self.kv_cache_manager
-        if self.max_draft_len > 0:
-            meta.update_spec_dec_param(
-                batch_size=1,
-                is_spec_decoding_enabled=True,
-                is_spec_dec_tree=False,
-                is_spec_dec_dynamic_tree=False,
-                max_draft_len=self.max_draft_len,
-                max_total_draft_tokens=self.max_draft_len,
-                num_contexts=0,
-            )
-        meta.prepare()
-        return meta
-
-    # -- running -------------------------------------------------------------
-    @contextlib.contextmanager
-    def scope(self, meta=None):
-        """Model extra attrs, multi-stream mode and the step's phase, as the engine forward
-        sets them: the GEN metadata (``meta`` None) is a decode step, ``meta`` a prefill."""
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-            dispatch,
-        )
-        from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
-        from tensorrt_llm._torch.utils import model_extra_attrs
-
-        phase = dispatch.Phase.DECODE if meta is None else dispatch.Phase.PREFILL
-        attrs = dict(self.model_config.extra_attrs)
-        attrs["attention_metadata"] = weakref.ref(meta or self.gen_metadata)
-        with model_extra_attrs(attrs), with_multi_stream(True), dispatch.step(phase):
-            yield
-
-    def set_attn_tail(self, enabled: bool, config=None) -> None:
-        """Attach (or detach) each layer's attention_output_tail; optionally set its config."""
-        for layer, attn in self.attn.items():
-            tail = self.attn_tail_tails[layer]
-            if enabled and tail is None:
-                raise RuntimeError(
-                    f"layer {layer} has no attention_output_tail tail "
-                    "(the installed plan does not run attention_output_tail)"
-                )
-            if tail is not None and config is not None:
-                tail.config = config
-            attn._attn_tail = tail if enabled else None
-
-    def gen_inputs(self, layer: int, step: int = 0) -> torch.Tensor:
-        """GEN hidden states ``[rows, hidden]`` for ``layer`` at ``step``."""
-        first = TOKEN_OFFSET + self.ctx + step * self.rows
-        return layer_inputs(self.ckpt, layer, first, self.rows)
-
-    def position_ids(self) -> torch.Tensor:
-        return torch.arange(self.ctx, self.ctx + self.rows, dtype=torch.int32, device="cuda")[None]
-
-    def forward(
-        self, layer: int, hidden: torch.Tensor, position_ids, all_reduce_params
-    ) -> torch.Tensor:
-        """One GEN ``GlmMLA.forward`` of ``layer``; call inside ``scope()``."""
-        return self.attn[layer](
-            position_ids=position_ids,
-            hidden_states=hidden,
-            attn_metadata=self.gen_metadata,
-            all_reduce_params=all_reduce_params,
-        )
-
-    def shutdown(self) -> None:
-        self.kv_cache_manager.shutdown()
-
-
-class LatentRecorder:
-    """Record the latent each layer's GEN tail receives (eager runs only).
-
-    Wraps the native W_UV entry (``MLA._bmm_bf16_out``) and the attention_output_tail entry
-    (``AttnTail.defer_w_uv``) on the given modules, so the test can build the
-    FP64 oracle from the latent the real attention produced.
-    """
-
-    def __init__(self, stack: AttentionStack):
-        self.stack = stack
-        self.latents: Dict[int, torch.Tensor] = {}
-        self._saved: List = []
-
-    def __enter__(self):
-        for layer, attn in self.stack.attn.items():
-            native = attn._bmm_bf16_out
-
-            def bmm(a, b, b_t, out, *rest, _layer=layer, _native=native):
-                # a: latent transposed to [H, R, C].
-                self.latents[_layer] = a.transpose(0, 1).contiguous().clone()
-                return _native(a, b, b_t, out, *rest)
-
-            attn._bmm_bf16_out = bmm
-            self._saved.append((attn, "_bmm_bf16_out", None))
-            tail = self.stack.attn_tail_tails[layer]
-            if tail is not None:
-                defer = tail.defer_w_uv
-
-                def wrapped(mla, meta, latent, output, _layer=layer, _defer=defer):
-                    self.latents[_layer] = latent.contiguous().clone()
-                    return _defer(mla, meta, latent, output)
-
-                tail.defer_w_uv = wrapped
-                self._saved.append((tail, "defer_w_uv", None))
-        return self
-
-    def __exit__(self, *exc):
-        for obj, name, _ in self._saved:
-            delattr(obj, name)
-        self._saved.clear()
-        return False
-
-
-def native_reference(attn, latent: torch.Tensor) -> Optional[torch.Tensor]:
-    """FP64 oracle of the tail on a recorded latent (native BF16 ``v``)."""
-    import tail_reference
-
-    ref, _ = tail_reference.oracle_attn_tail(
-        latent, attn.v_b_proj, attn.o_proj.weight, latent.shape[0]
-    )
-    return ref
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py
deleted file mode 100644
index 65eec6af6e4e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_fixtures.py
+++ /dev/null
@@ -1,191 +0,0 @@
-# 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.
-"""Real-weight fixtures for the attention_output_tail, read from a GLM-5.x checkpoint.
-
-The TP slicing mirrors ``DeepseekV3WeightLoader``: ``kv_b_proj`` is unflattened
-to ``[heads, qk_nope + v, kv_lora]``, chunked along heads per rank, and the
-``v`` rows form ``v_b_proj [heads / tp, v, kv_lora]``; the ROW-parallel
-``o_proj`` takes the rank's chunk of input features (columns).
-
-``compressed_kv_rows`` and ``mix_latent`` build a latent attention output from
-real layer weights (token embeddings → input RMSNorm → ``kv_a_proj_with_mqa``
-→ ``kv_a_layernorm``), mixed by random softmax weights. It is a real-weight, synthetic-activation
-fixture: the rows are convex combinations of real compressed-KV vectors, but
-the attention probabilities are random, not model activations.
-"""
-
-import json
-import os
-from typing import Dict, Iterable, Tuple
-
-import torch
-from utils.llm_data import llm_models_root
-
-GLM52_CKPT_ENV = "GLM52_NVFP4_CKPT"
-
-
-def checkpoint_dir() -> str:
-    """The GLM-5.2 NVFP4 checkpoint: ``$GLM52_NVFP4_CKPT``, else under ``llm_models_root()``."""
-    path = os.environ.get(GLM52_CKPT_ENV)
-    if path:
-        return path
-    root = llm_models_root()
-    return os.path.join(str(root), "GLM-5.2-NVFP4") if root is not None else ""
-
-
-def _index(ckpt_dir: str) -> Dict[str, str]:
-    with open(os.path.join(ckpt_dir, "model.safetensors.index.json")) as f:
-        return json.load(f)["weight_map"]
-
-
-def _read(
-    ckpt_dir: str, weight_map: Dict[str, str], name: str, rows: slice = slice(None)
-) -> torch.Tensor:
-    from safetensors import safe_open
-
-    path = os.path.join(ckpt_dir, weight_map[name])
-    with safe_open(path, framework="pt", device="cpu") as f:
-        return f.get_slice(name)[rows].contiguous()
-
-
-def read_tensor(ckpt_dir: str, name: str) -> torch.Tensor:
-    """Read one checkpoint tensor by its safetensors name (CPU)."""
-    return _read(ckpt_dir, _index(ckpt_dir), name)
-
-
-def read_prefix(ckpt_dir: str, prefix: str) -> Dict[str, torch.Tensor]:
-    """Read every checkpoint tensor whose name starts with ``prefix`` (CPU)."""
-    from safetensors import safe_open
-
-    weight_map = _index(ckpt_dir)
-    by_file: Dict[str, list] = {}
-    for name, fname in weight_map.items():
-        if name.startswith(prefix):
-            by_file.setdefault(fname, []).append(name)
-    out = {}
-    for fname, names in by_file.items():
-        with safe_open(os.path.join(ckpt_dir, fname), framework="pt", device="cpu") as f:
-            for name in names:
-                out[name] = f.get_tensor(name)
-    return out
-
-
-def normed_embeddings(
-    ckpt_dir: str, layer: int, first_token: int, num_tokens: int, device: str = "cuda"
-) -> torch.Tensor:
-    """``input_layernorm`` of ``layer`` over the embeddings of consecutive
-    token IDs ``first_token ... first_token + num_tokens - 1`` (BF16)."""
-    with open(os.path.join(ckpt_dir, "config.json")) as f:
-        eps = json.load(f).get("rms_norm_eps", 1e-6)
-    weight_map = _index(ckpt_dir)
-    emb = _read(
-        ckpt_dir,
-        weight_map,
-        "model.embed_tokens.weight",
-        slice(first_token, first_token + num_tokens),
-    ).to(device)
-    ln = _read(ckpt_dir, weight_map, f"model.layers.{layer}.input_layernorm.weight").to(device)
-    return _rms_norm(emb, ln, eps).contiguous()
-
-
-def load_tail_weights(
-    ckpt_dir: str, layer: int, ranks: Iterable[int], tp_size: int = 8
-) -> Dict[int, Tuple[torch.Tensor, torch.Tensor]]:
-    """Load ``(v_b_proj [H/tp, V, C], o_proj [N, H*V/tp])`` per TP rank (CPU)."""
-    with open(os.path.join(ckpt_dir, "config.json")) as f:
-        cfg = json.load(f)
-    heads = cfg["num_attention_heads"]
-    nope = cfg["qk_nope_head_dim"]
-    v_dim = cfg["v_head_dim"]
-    weight_map = _index(ckpt_dir)
-    prefix = f"model.layers.{layer}.self_attn"
-    kv_b = _read(ckpt_dir, weight_map, f"{prefix}.kv_b_proj.weight")
-    kv_b = kv_b.unflatten(0, [heads, nope + v_dim])
-    o_proj = _read(ckpt_dir, weight_map, f"{prefix}.o_proj.weight")
-    out = {}
-    for rank in ranks:
-        kv_b_rank = torch.chunk(kv_b, tp_size, dim=0)[rank]
-        _, v_rank = kv_b_rank.split([nope, v_dim], dim=1)
-        o_rank = torch.chunk(o_proj, tp_size, dim=1)[rank]
-        out[rank] = (v_rank.contiguous(), o_rank.contiguous())
-    return out
-
-
-def load_full_o_proj(ckpt_dir: str, layer: int) -> torch.Tensor:
-    """Return the unsharded ``o_proj`` weight ``[N, H*V]`` (CPU)."""
-    return _read(ckpt_dir, _index(ckpt_dir), f"model.layers.{layer}.self_attn.o_proj.weight")
-
-
-def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
-    xf = x.float()
-    y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
-    return (y * weight.float()).to(x.dtype)
-
-
-def compressed_kv_rows(
-    ckpt_dir: str, layer: int, context: int = 2048, token_offset: int = 1000, device: str = "cuda"
-) -> torch.Tensor:
-    """Return ``[context, kv_lora]`` normalized compressed-KV rows of ``layer``.
-
-    Computed from the embeddings of token IDs ``token_offset`` to
-    ``token_offset + context - 1``: input RMSNorm, the compressed part of
-    ``kv_a_proj_with_mqa``, then ``kv_a_layernorm``.
-    """
-    with open(os.path.join(ckpt_dir, "config.json")) as f:
-        cfg = json.load(f)
-    eps = cfg.get("rms_norm_eps", 1e-6)
-    kv_lora = cfg["kv_lora_rank"]
-    weight_map = _index(ckpt_dir)
-    prefix = f"model.layers.{layer}"
-    emb = _read(
-        ckpt_dir,
-        weight_map,
-        "model.embed_tokens.weight",
-        slice(token_offset, token_offset + context),
-    ).to(device)
-    ln = _read(ckpt_dir, weight_map, f"{prefix}.input_layernorm.weight").to(device)
-    kv_a = _read(ckpt_dir, weight_map, f"{prefix}.self_attn.kv_a_proj_with_mqa.weight").to(device)
-    kv_ln = _read(ckpt_dir, weight_map, f"{prefix}.self_attn.kv_a_layernorm.weight").to(device)
-    hidden = _rms_norm(emb, ln, eps)
-    compressed = torch.nn.functional.linear(hidden, kv_a)[:, :kv_lora]
-    return _rms_norm(compressed, kv_ln, eps)
-
-
-def mix_latent(
-    compressed: torch.Tensor,
-    num_rows: int,
-    num_local_heads: int = 8,
-    score_scale: float = 3.0,
-    seed: int = 0,
-) -> torch.Tensor:
-    """Return a ``[num_rows, num_local_heads, kv_lora]`` BF16 latent.
-
-    Each (row, head) is a softmax-weighted average of the ``compressed`` rows
-    with random scores, like a sparse-attention output over a real cache.
-    """
-    device = compressed.device
-    gen = torch.Generator(device=device).manual_seed(seed)
-    scores = (
-        torch.randn(
-            (num_rows, num_local_heads, compressed.shape[0]),
-            generator=gen,
-            device=device,
-            dtype=torch.float32,
-        )
-        * score_scale
-    )
-    probs = torch.softmax(scores, dim=-1)
-    latent = torch.einsum("rht,tc->rhc", probs, compressed.float())
-    return latent.to(torch.bfloat16).contiguous()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py
deleted file mode 100644
index c13f46af2f37..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/tail_reference.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# 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.
-"""Oracles for the attention_output_tail.
-
-``native_attn_tail`` runs the exact native TensorRT-LLM operations of the GEN
-tail: ``MLA._bmm_bf16_out`` for W_UV (the BF16 branch of
-``MLA.forward_absorption_generation``) followed by the ROW-parallel ``o_proj``
-GEMM (``UnquantizedLinearMethod.apply``). ``oracle_attn_tail`` is an
-independent FP64 reference that keeps the native BF16 rounding of the
-intermediate ``v``.
-"""
-
-from types import SimpleNamespace
-from typing import Optional, Tuple
-
-import torch
-
-
-def native_attn_tail(
-    latent: torch.Tensor,
-    w_uv: torch.Tensor,
-    w_o: torch.Tensor,
-    num_rows: int,
-    out_rows: Optional[int] = None,
-    o_proj: Optional[torch.nn.Module] = None,
-    return_v: bool = False,
-):
-    """Run the native W_UV BMM and o_proj GEMM.
-
-    Args:
-        latent: ``[>= num_rows, H * C]`` or ``[>= num_rows, H, C]`` BF16.
-        w_uv: ``[H, V, C]`` BF16 ``v_b_proj``.
-        w_o: ``[N, H * V]`` BF16 local o_proj weight (ignored when ``o_proj``
-            is given).
-        num_rows: live rows written by W_UV.
-        out_rows: rows of the o_proj input buffer (native ``create_output``
-            rows); rows past ``num_rows`` are zero-filled here so the oracle is
-            defined, whereas the native buffer leaves them uninitialized.
-        o_proj: optional native ``Linear`` (ROW parallel, no reduction) to run
-            instead of ``F.linear`` on ``w_o``.
-        return_v: also return the ``[out_rows, H * V]`` intermediate.
-    """
-    from tensorrt_llm._torch.attention.mla import MLA
-
-    heads, v_dim, c_dim = w_uv.shape
-    if out_rows is None:
-        out_rows = num_rows
-    attn_out = torch.zeros((out_rows, heads * v_dim), dtype=w_uv.dtype, device=latent.device)
-    lat = latent[:num_rows].reshape(num_rows, heads, c_dim)
-    v_view = attn_out[:num_rows].view(num_rows, heads, v_dim)
-    native_cfg = SimpleNamespace(use_cute_dsl_bf16_bmm=False)
-    MLA._bmm_bf16_out(
-        native_cfg, lat.transpose(0, 1), w_uv, w_uv.transpose(1, 2), v_view.transpose(0, 1)
-    )
-    if o_proj is not None:
-        partial = o_proj(attn_out)
-    else:
-        partial = torch.nn.functional.linear(attn_out, w_o)
-    if return_v:
-        return partial, attn_out
-    return partial
-
-
-def oracle_attn_tail(
-    latent: torch.Tensor, w_uv: torch.Tensor, w_o: torch.Tensor, num_rows: int
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """FP64 reference with the native BF16 rounding of the intermediate.
-
-    Returns:
-        ``(partial [num_rows, N] float64, v [num_rows, H * V] bfloat16)``.
-    """
-    heads, v_dim, c_dim = w_uv.shape
-    lat = latent[:num_rows].reshape(num_rows, heads, c_dim).double()
-    v = torch.einsum("rhc,hvc->rhv", lat, w_uv.double())
-    v_bf16 = v.reshape(num_rows, heads * v_dim).to(torch.bfloat16)
-    partial = v_bf16.double() @ w_o.double().t()
-    return partial, v_bf16
-
-
-def error_stats(x: torch.Tensor, ref: torch.Tensor) -> dict:
-    """Normalized RMS / max errors of ``x`` against ``ref`` (both cast to FP64)."""
-    diff = x.double() - ref.double()
-    ref_rms = ref.double().pow(2).mean().sqrt().item()
-    return {
-        "rel_rms": diff.pow(2).mean().sqrt().item() / max(ref_rms, 1e-300),
-        "max_abs": diff.abs().max().item(),
-        "ref_rms": ref_rms,
-        "ref_max": ref.double().abs().max().item(),
-    }
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py
deleted file mode 100644
index 95bac645109e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_all_layers.py
+++ /dev/null
@@ -1,130 +0,0 @@
-# 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.
-"""Default attention_output_tail engine on the real tail weights of all 78 target layers.
-
-For every target layer (3 dense destinations, 75 MoE destinations) the test
-loads the real TP8 ``v_b_proj`` / ``o_proj`` shards of one rank (rank
-``layer % 8``, so every rank is covered across the sweep) and of all eight
-ranks for layers 1, 41 and 76, and runs the default ``cute_stream`` engine
-with PDL at R ∈ {1, 6, 8} on real-weight, synthetic-activation latents (real
-compressed-KV rows of the same layer, random softmax mixtures).
-
-Checks per (layer, rank, rows):
-
-* the frozen partial-output budget of ``test_attn_tail_attn_tail.py`` against the
-  FP64 oracle with the native BF16 rounding of ``v``;
-* the strict intermediate check for the default engine: attention_output_tail's BF16 ``v``
-  equals native ``v`` (``MLA._bmm_bf16_out``) bit for bit, or differs by at
-  most one BF16 ulp per element; the count of non-identical elements is
-  printed. (The relaxed FP64-bound check of ``test_attn_tail_attn_tail.py`` applies
-  to the non-default engines.)
-
-Weights are re-read per layer from the checkpoint, so the device sees 78
-distinct weight sets (no warm reuse across layers).
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
-    DEFAULT_CONFIG,
-    attn_tail,
-    check_tail_shapes,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-TP_SIZE = 8
-LOCAL_HEADS, V_DIM, C_DIM = 8, 256, 512
-NUM_TARGET_LAYERS = 78
-ALL_RANK_LAYERS = (1, 41, 76)
-ROWS = (1, 6, 8)
-_CKPT = tail_fixtures.checkpoint_dir()
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="attention_output_tail kernels target SM100/SM103",
-    ),
-    pytest.mark.skipif(
-        not (_CKPT and os.path.isdir(_CKPT)), reason=f"{tail_fixtures.GLM52_CKPT_ENV} not available"
-    ),
-]
-
-
-def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """Per-element distance in BF16 ulps (+0 and -0 are equal)."""
-
-    def ordered(x):
-        bits = x.contiguous().view(torch.int16).to(torch.int32)
-        mag = bits & 0x7FFF
-        return torch.where(bits < 0, -mag, mag)
-
-    return (ordered(a) - ordered(b)).abs()
-
-
-@pytest.mark.parametrize("layer", range(NUM_TARGET_LAYERS))
-def test_default_engine_layer(layer):
-    ranks = list(range(TP_SIZE)) if layer in ALL_RANK_LAYERS else [layer % TP_SIZE]
-    weights = tail_fixtures.load_tail_weights(_CKPT, layer, ranks, TP_SIZE)
-    compressed = tail_fixtures.compressed_kv_rows(_CKPT, layer)
-    destination = "dense" if layer < 3 else "moe"
-    for rank in ranks:
-        w_uv, w_o = (t.cuda() for t in weights[rank])
-        for rows in ROWS:
-            latent = tail_fixtures.mix_latent(
-                compressed, rows, num_local_heads=LOCAL_HEADS, seed=layer * 100 + rank * 10 + rows
-            )
-            assert check_tail_shapes(latent, w_uv, w_o, rows, rows) is None
-            v_attn_tail = torch.empty(
-                (rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda"
-            )
-            attention_output_tail = attn_tail(
-                latent, w_uv, w_o, rows, v_scratch=v_attn_tail, config=DEFAULT_CONFIG
-            )
-            native, v_nat = tail_reference.native_attn_tail(latent, w_uv, w_o, rows, return_v=True)
-            ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, w_o, rows)
-            e_attn_tail = tail_reference.error_stats(attention_output_tail, ref)
-            e_nat = tail_reference.error_stats(native, ref)
-            label = (layer, destination, rank, rows)
-            assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (
-                label,
-                e_attn_tail,
-                e_nat,
-            )
-            assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_nat["ref_max"]), (
-                label,
-                e_attn_tail,
-                e_nat,
-            )
-            torch.testing.assert_close(
-                attention_output_tail, native, rtol=2**-6, atol=2**-6 * e_nat["ref_rms"]
-            )
-            ulps = bf16_ulp_distance(v_attn_tail, v_nat)
-            mismatches = int((ulps > 0).sum())
-            assert int(ulps.max()) <= 1, (label, int(ulps.max()), mismatches)
-            print(
-                f"[attention_output_tail-pdl-stress-layer] layer={layer} dest={destination} "
-                f"rank={rank} rows={rows} rel_rms attention_output_tail={e_attn_tail['rel_rms']:.3e} "
-                f"native={e_nat['rel_rms']:.3e} "
-                f"v_mismatch={mismatches}/{v_attn_tail.numel()} "
-                f"v_max_ulp={int(ulps.max())}"
-            )
-        del w_uv, w_o
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py
deleted file mode 100644
index b6fc4bc1bb36..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_attn_tail.py
+++ /dev/null
@@ -1,263 +0,0 @@
-# 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.
-"""Single-GPU correctness of the attention_output_tail (the attention_output_tail contract).
-
-Oracles: the native TensorRT-LLM ops (``MLA._bmm_bf16_out`` + ROW-parallel
-``Linear``) and an FP64 reference that keeps the native BF16 rounding of the
-W_UV intermediate.
-
-Error budget, frozen before kernel tuning (normalized against the FP64
-reference ``ref``; ``err(x) = ||x - ref||_2 / ||ref||_2``):
-
-* ``err(attention_output_tail) <= 1.25 * err(native) + 1e-6`` on every (layer, rank, rows) case;
-* ``max|attention_output_tail - ref| <= 2 * max|native - ref| + 2**-8 * max|ref|``;
-* ``torch.testing.assert_close(attention_output_tail, native, rtol=2**-6, atol=2**-6 * rms(ref))``.
-
-Fixtures: real GLM-5.2 ``kv_b_proj``/``o_proj`` weights sliced per TP8 rank
-(``GLM52_NVFP4_CKPT``) with two latent kinds: ``derived`` (real-weight,
-synthetic-activation: softmax mixtures of real compressed-KV rows) and
-``randn`` (synthetic). Tests that need the checkpoint skip without it.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
-    ROW_TILE,
-    TailConfig,
-    attn_tail,
-    check_tail_shapes,
-    check_weights,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-TP_SIZE = 8
-HEADS, V_DIM, C_DIM, HIDDEN = 64, 256, 512, 6144
-LOCAL_HEADS = HEADS // TP_SIZE
-# Dense layers 0 and 2, MoE layers 3, 40 and 77 (all TP8 ranks of each).
-LAYERS = (0, 2, 3, 40, 77)
-ROWS = (1, 2, 3, 4, 5, 6, 7, 8, 12, 16)
-
-_CKPT = tail_fixtures.checkpoint_dir()
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="attention_output_tail kernels target SM100/SM103",
-)
-
-needs_ckpt = pytest.mark.skipif(
-    not (_CKPT and os.path.isdir(_CKPT)),
-    reason=f"{tail_fixtures.GLM52_CKPT_ENV} checkpoint not available",
-)
-
-ENGINES = {
-    "fused": TailConfig(oproj_engine="fused"),
-    "cute_stream": TailConfig(oproj_engine="cute_stream"),
-    "triton": TailConfig(oproj_engine="triton"),
-}
-
-_WEIGHTS = {}
-_COMPRESSED = {}
-
-
-def _layer_weights(layer):
-    if layer not in _WEIGHTS:
-        _WEIGHTS[layer] = tail_fixtures.load_tail_weights(_CKPT, layer, range(TP_SIZE), TP_SIZE)
-    return _WEIGHTS[layer]
-
-
-def _latent(kind, layer, rows, seed):
-    if kind == "derived":
-        if layer not in _COMPRESSED:
-            _COMPRESSED[layer] = tail_fixtures.compressed_kv_rows(_CKPT, layer)
-        return tail_fixtures.mix_latent(
-            _COMPRESSED[layer], rows, num_local_heads=LOCAL_HEADS, seed=seed
-        )
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    return torch.randn(
-        (rows, LOCAL_HEADS, C_DIM), generator=gen, device="cuda", dtype=torch.float32
-    ).to(torch.bfloat16)
-
-
-def _check_budget(attention_output_tail, native, ref, label):
-    e_attn_tail = tail_reference.error_stats(attention_output_tail, ref)
-    e_nat = tail_reference.error_stats(native, ref)
-    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (label, e_attn_tail, e_nat)
-    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_nat["ref_max"]), (
-        label,
-        e_attn_tail,
-        e_nat,
-    )
-    torch.testing.assert_close(
-        attention_output_tail, native, rtol=2**-6, atol=2**-6 * e_nat["ref_rms"]
-    )
-    return e_attn_tail, e_nat
-
-
-@needs_ckpt
-@pytest.mark.parametrize("layer", LAYERS)
-def test_o_proj_slice_matches_native_row_linear(layer):
-    """The fixture's o_proj slice equals what a native ROW Linear loads."""
-    from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
-    from tensorrt_llm.mapping import Mapping
-
-    full = tail_fixtures.load_full_o_proj(_CKPT, layer)
-    weights = _layer_weights(layer)
-    for rank in (0, 3, 7):
-        linear = Linear(
-            HEADS * V_DIM,
-            HIDDEN,
-            bias=False,
-            dtype=torch.bfloat16,
-            mapping=Mapping(world_size=TP_SIZE, tp_size=TP_SIZE, rank=rank),
-            tensor_parallel_mode=TensorParallelMode.ROW,
-            reduce_output=False,
-        ).cuda()
-        linear.load_weights([{"weight": full}])
-        assert torch.equal(linear.weight.cpu(), weights[rank][1])
-
-
-@needs_ckpt
-@pytest.mark.parametrize("layer", LAYERS)
-@pytest.mark.parametrize("kind", ["derived", "randn"])
-@pytest.mark.parametrize("use_pdl", [True, False])
-@pytest.mark.parametrize("engine", list(ENGINES))
-def test_tail_matches_native_real_weights(layer, kind, use_pdl, engine):
-    """All TP8 ranks and row counts: attention_output_tail vs native vs FP64 within the budget."""
-    worst = {"attention_output_tail": 0.0, "native": 0.0}
-    for rank, (w_uv_cpu, w_o_cpu) in _layer_weights(layer).items():
-        w_uv, w_o = w_uv_cpu.cuda(), w_o_cpu.cuda()
-        for rows in ROWS:
-            latent = _latent(kind, layer, rows, seed=rank * 100 + rows)
-            assert check_tail_shapes(latent, w_uv, w_o, rows, rows) is None
-            attention_output_tail = attn_tail(
-                latent, w_uv, w_o, rows, use_pdl=use_pdl, config=ENGINES[engine]
-            )
-            native = tail_reference.native_attn_tail(latent, w_uv, w_o, rows)
-            ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, w_o, rows)
-            e_attn_tail, e_nat = _check_budget(
-                attention_output_tail, native, ref, (layer, rank, rows, kind)
-            )
-            worst["attention_output_tail"] = max(
-                worst["attention_output_tail"], e_attn_tail["rel_rms"]
-            )
-            worst["native"] = max(worst["native"], e_nat["rel_rms"])
-    print(
-        f"[attention_output_tail-err] layer={layer} kind={kind} pdl={use_pdl} engine={engine} "
-        f"worst_rel_rms attention_output_tail={worst['attention_output_tail']:.3e} native={worst['native']:.3e}"
-    )
-
-
-@needs_ckpt
-@pytest.mark.parametrize("layer", (2, 40))
-@pytest.mark.parametrize("engine", list(ENGINES))
-def test_intermediate_rounding_matches_native(layer, engine):
-    """The W_UV intermediate is a correctly rounded BF16 of the exact product.
-
-    Both attention_output_tail and native ``v`` must lie within one BF16 rounding step of the
-    FP64 product plus an FP32 accumulation-order term (``2**-20`` of the sum
-    of absolute products, which matters only under cancellation); elements
-    that differ between attention_output_tail and native must stay rare.
-    """
-    for rank in (0, 5):
-        w_uv_cpu, w_o_cpu = _layer_weights(layer)[rank]
-        w_uv, w_o = w_uv_cpu.cuda(), w_o_cpu.cuda()
-        rows = 6
-        latent = _latent("derived", layer, rows, seed=rank)
-        v_attn_tail = torch.empty((rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
-        attn_tail(latent, w_uv, w_o, rows, v_scratch=v_attn_tail, config=ENGINES[engine])
-        _, v_nat = tail_reference.native_attn_tail(latent, w_uv, w_o, rows, return_v=True)
-        lat64 = latent.reshape(rows, LOCAL_HEADS, C_DIM).double()
-        v64 = torch.einsum("rhc,hvc->rhv", lat64, w_uv.double()).reshape(rows, -1)
-        mag = torch.einsum("rhc,hvc->rhv", lat64.abs(), w_uv.double().abs()).reshape(rows, -1)
-        bound = torch.finfo(torch.bfloat16).eps * v64.abs() + 2**-20 * mag
-        for name, v in (("attention_output_tail", v_attn_tail), ("native", v_nat)):
-            err = (v.double() - v64).abs()
-            assert bool((err <= bound).all()), (name, (err - bound).max())
-        mismatch = (v_attn_tail != v_nat).float().mean().item()
-        assert mismatch < 0.01, mismatch
-        print(
-            f"[attention_output_tail-v] layer={layer} rank={rank} engine={engine} "
-            f"v_mismatch_fraction={mismatch:.2e}"
-        )
-
-
-@pytest.mark.parametrize(
-    "live,out_rows", [(1, 4), (6, 8), (6, 9), (6, 16), (6, 20), (9, 12), (16, 40)]
-)
-@pytest.mark.parametrize("engine", list(ENGINES))
-def test_padded_rows_are_zero(live, out_rows, engine):
-    torch.manual_seed(live)
-    w_uv = (torch.randn(LOCAL_HEADS, V_DIM, C_DIM, device="cuda") * 0.02).bfloat16()
-    w_o = (torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda") * 0.02).bfloat16()
-    latent = torch.randn(live, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
-    out = attn_tail(latent, w_uv, w_o, live, out_rows=out_rows, config=ENGINES[engine])
-    assert out.shape == (out_rows, HIDDEN)
-    assert bool((out[live:] == 0).all())
-    native = tail_reference.native_attn_tail(latent, w_uv, w_o, live)
-    ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, w_o, live)
-    _check_budget(out[:live], native, ref, (live, out_rows))
-
-
-@pytest.mark.parametrize("use_pdl", [True, False])
-@pytest.mark.parametrize("engine", list(ENGINES))
-def test_cuda_graph_replay(use_pdl, engine):
-    """Replay a captured tail with fresh latent contents and padded rows."""
-    torch.manual_seed(0)
-    rows, out_rows = 6, 8
-    w_uv = (torch.randn(LOCAL_HEADS, V_DIM, C_DIM, device="cuda") * 0.02).bfloat16()
-    w_o = (torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda") * 0.02).bfloat16()
-    # Latent rows live inside a wider buffer, like the backend's output.
-    latent_buf = torch.zeros(rows, LOCAL_HEADS * C_DIM, device="cuda").bfloat16()
-    cfg = ENGINES[engine]
-    attn_tail(latent_buf, w_uv, w_o, rows, out_rows=out_rows, use_pdl=use_pdl, config=cfg)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        out = attn_tail(latent_buf, w_uv, w_o, rows, out_rows=out_rows, use_pdl=use_pdl, config=cfg)
-    for step in range(64):
-        latent_buf.copy_(torch.randn(rows, LOCAL_HEADS * C_DIM, device="cuda").bfloat16())
-        graph.replay()
-        native = tail_reference.native_attn_tail(latent_buf, w_uv, w_o, rows)
-        ref, _ = tail_reference.oracle_attn_tail(latent_buf, w_uv, w_o, rows)
-        _check_budget(out[:rows].clone(), native, ref, ("graph", step))
-        assert bool((out[rows:] == 0).all())
-
-
-def test_rejects_unsupported_inputs():
-    w_uv = torch.zeros(LOCAL_HEADS, V_DIM, C_DIM, device="cuda").bfloat16()
-    w_o = torch.zeros(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda").bfloat16()
-    latent = torch.zeros(6, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
-    assert check_tail_shapes(latent, w_uv, w_o, 6, 6) is None
-    assert check_tail_shapes(latent, w_uv, w_o, ROW_TILE + 1, ROW_TILE + 1) is not None
-    assert check_tail_shapes(latent, w_uv, w_o, 0, 0) is not None
-    assert check_tail_shapes(latent.float(), w_uv, w_o, 6, 6) is not None
-    assert check_weights(w_uv.to(torch.float8_e4m3fn), w_o) is not None
-    assert check_weights(w_uv, w_o.t()) is not None
-    assert check_weights(w_uv, w_o[:, :1024]) is not None
-    assert check_tail_shapes(latent.transpose(0, 1), w_uv, w_o, 6, 6) is not None
-
-
-def test_registered_entry_points():
-    assert dispatch.get_impl("attention_output_tail", "attn_tail") is attn_tail
-    make_tail = dispatch.get_impl("attention_output_tail", "make_attn_tail")
-    assert callable(make_tail)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py
deleted file mode 100644
index 3ba4d02e55dc..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_bridge_mpi.py
+++ /dev/null
@@ -1,244 +0,0 @@
-# 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.
-"""attention_output_tail contract bridge with a synthetic latent, into the native reduction/norm.
-
-The latent is ``tail_fixtures.mix_latent`` (softmax mixes of real compressed-KV
-rows with random scores; no attention kernel runs), so this file checks only
-the consumer side of the bridge. The native attention producer → attention_output_tail → native
-reduction/norm chain runs in ``test_attn_tail_mla_forward_mpi.py``.
-
-Launch one process per GPU under MPI (``srun --mpi=pmix -n <tp>``); every
-rank runs the same tests in the same order. Each rank holds its own TP8 shard
-of real layer weights (rank ``r`` uses TP8 shard ``r``; with fewer than eight
-ranks the all-reduce sums a subset of shards, which still exercises the
-bridge). Both destinations are covered:
-
-* MoE layer (3): ``o_proj`` with ``enable_allreduce=False`` → decoder-level
-  ``AllReduce(RESIDUAL_RMS_NORM)`` with the real ``post_attention_layernorm``.
-* Dense layer (0): ``o_proj`` reduces inside ``Linear.forward``, then the real
-  ``post_attention_layernorm(hidden, residual)``.
-
-The native chain uses the native W_UV BMM and the native ROW ``Linear``; the
-attention_output_tail chain goes through ``AttnTail.defer_w_uv`` / ``AttnTail.project``, which
-returns the rank-local partial without communicating, and then reduces it as
-``MLA.forward`` does (``o_proj.all_reduce`` with the caller's params, only
-where ``o_proj`` reduces), with the same ``Linear`` and ``AllReduce``
-modules. ``project`` must not call the all-reduce. Outputs must match within
-the frozen attention_output_tail tolerance and be bit-identical across ranks.
-"""
-
-import hashlib
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
-    AttnTail,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-TP8 = 8
-HEADS, V_DIM, C_DIM, HIDDEN = 64, 256, 512, 6144
-LOCAL_HEADS = HEADS // TP8
-EPS = 1e-5
-
-_CKPT = tail_fixtures.checkpoint_dir()
-
-
-def _comm():
-    from mpi4py import MPI
-
-    return MPI.COMM_WORLD
-
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"),
-    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
-    pytest.mark.skipif(_comm().Get_size() < 2, reason="needs >= 2 MPI ranks"),
-]
-
-
-_SETUP = {}
-
-
-def _setup(layer):
-    if layer not in _SETUP:
-        _SETUP[layer] = _build(layer)
-    return _SETUP[layer]
-
-
-def _build(layer):
-    from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
-    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-    from tensorrt_llm.mapping import Mapping
-
-    comm = _comm()
-    rank, world = comm.Get_rank(), comm.Get_size()
-    torch.cuda.set_device(rank % torch.cuda.device_count())
-    mapping = Mapping(
-        world_size=world, tp_size=world, rank=rank, gpus_per_node=torch.cuda.device_count()
-    )
-    w_uv, w_o = tail_fixtures.load_tail_weights(_CKPT, layer, [rank], TP8)[rank]
-    w_uv = w_uv.cuda()
-    # A ROW Linear over `world` ranks whose local slice is this rank's TP8
-    # shard: load the [N, world * K_local] matrix built from TP8 shards.
-    o_proj = Linear(
-        world * LOCAL_HEADS * V_DIM,
-        HIDDEN,
-        bias=False,
-        dtype=torch.bfloat16,
-        mapping=mapping,
-        tensor_parallel_mode=TensorParallelMode.ROW,
-        reduce_output=True,
-    ).cuda()
-    o_proj.weight.data.copy_(w_o.cuda())
-    norm_w = tail_fixtures.read_tensor(
-        _CKPT, f"model.layers.{layer}.post_attention_layernorm.weight"
-    ).cuda()
-    norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda()
-    norm.weight.data.copy_(norm_w)
-    return rank, world, mapping, w_uv, o_proj, norm
-
-
-def _digest(t):
-    return hashlib.sha256(t.contiguous().view(torch.uint8).cpu().numpy().tobytes()).hexdigest()
-
-
-def _tail_ns(w_uv, o_proj):
-    tail = AttnTail("bridge")
-    tail._static_ok = True  # The weights are real; MLA-type checks need a module.
-    mla = SimpleNamespace(v_b_proj=w_uv, o_proj=o_proj)
-    return tail, mla
-
-
-def _run_tail(tail, mla, latent, out_rows, all_reduce_params):
-    """The attention_output_tail, then the reduction ``MLA.forward`` applies to it."""
-    rows = latent.shape[0]
-    attn_out = torch.empty((out_rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
-    meta = SimpleNamespace(num_contexts=0, num_tokens=rows)
-    o_proj = mla.o_proj
-    # Count calls of o_proj's AllReduce module while the tail runs.
-    reduce = o_proj.all_reduce
-    native_forward = reduce.forward
-    calls = []
-
-    def counting(*args, **kwargs):
-        calls.append(args)
-        return native_forward(*args, **kwargs)
-
-    reduce.forward = counting
-    try:
-        assert tail.arm(mla, meta)
-        assert tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
-        partial = tail.project(mla, attn_out)
-    finally:
-        del reduce.forward
-    assert not calls, "the attention_output_tail tail called the all-reduce"
-    if o_proj.reduce_output:
-        partial = o_proj.all_reduce(partial, all_reduce_params=all_reduce_params)
-    return partial
-
-
-@pytest.mark.parametrize("rows", [1, 6])
-def test_moe_layer_bridge(rows):
-    from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams
-
-    rank, world, mapping, w_uv, o_proj, norm = _setup(3)
-    allreduce = AllReduce(mapping=mapping, dtype=torch.bfloat16)
-    comp = tail_fixtures.compressed_kv_rows(_CKPT, 3)
-    latent = tail_fixtures.mix_latent(comp, rows, LOCAL_HEADS, seed=rank)
-    gen = torch.Generator(device="cuda").manual_seed(1234)
-    residual = torch.randn(
-        (rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
-    ).bfloat16()
-
-    def fused(partial):
-        return allreduce(
-            partial,
-            all_reduce_params=AllReduceParams(
-                fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
-                residual=residual,
-                norm_weight=norm.weight,
-                eps=EPS,
-                trigger_completion_at_end=False,
-            ),
-        )
-
-    no_ar = AllReduceParams(enable_allreduce=False)
-    native_partial = tail_reference.native_attn_tail(
-        latent, w_uv, None, rows, o_proj=lambda x: o_proj(x, all_reduce_params=no_ar)
-    )
-    hs_nat, res_nat = fused(native_partial)
-    tail, mla = _tail_ns(w_uv, o_proj)
-    attn_tail_partial = _run_tail(tail, mla, latent, rows, no_ar)
-    hs_attn_tail, res_attn_tail = fused(attn_tail_partial)
-    torch.cuda.synchronize()
-
-    ref, _ = tail_reference.oracle_attn_tail(latent, w_uv, o_proj.weight, rows)
-    e_attn_tail = tail_reference.error_stats(attn_tail_partial, ref)
-    e_nat = tail_reference.error_stats(native_partial, ref)
-    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (e_attn_tail, e_nat)
-    torch.testing.assert_close(res_attn_tail, res_nat, rtol=2**-6, atol=2**-6)
-    torch.testing.assert_close(hs_attn_tail, hs_nat, rtol=2**-6, atol=2**-5)
-    digests = _comm().allgather((_digest(hs_attn_tail), _digest(res_attn_tail)))
-    assert len(set(digests)) == 1, "post-norm outputs differ across ranks"
-    print(
-        f"[attention_output_tail-bridge] moe rank={rank}/{world} rows={rows} "
-        f"partial_err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e} "
-        f"hs_maxdiff={(hs_attn_tail.float() - hs_nat.float()).abs().max().item():.3e}"
-    )
-
-
-@pytest.mark.parametrize("rows", [1, 6])
-def test_dense_layer_bridge(rows):
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    rank, world, mapping, w_uv, o_proj, norm = _setup(0)
-    comp = tail_fixtures.compressed_kv_rows(_CKPT, 0)
-    latent = tail_fixtures.mix_latent(comp, rows, LOCAL_HEADS, seed=rank)
-    gen = torch.Generator(device="cuda").manual_seed(4321)
-    residual = torch.randn(
-        (rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
-    ).bfloat16()
-    with_ar = AllReduceParams(enable_allreduce=True)
-
-    native_reduced = tail_reference.native_attn_tail(
-        latent, w_uv, None, rows, o_proj=lambda x: o_proj(x, all_reduce_params=with_ar)
-    )
-    hs_nat, res_nat = norm(native_reduced, residual.clone())
-    tail, mla = _tail_ns(w_uv, o_proj)
-    attn_tail_reduced = _run_tail(tail, mla, latent, rows, with_ar)
-    hs_attn_tail, res_attn_tail = norm(attn_tail_reduced, residual.clone())
-    torch.cuda.synchronize()
-
-    torch.testing.assert_close(
-        attn_tail_reduced,
-        native_reduced,
-        rtol=2**-6,
-        atol=2**-6 * native_reduced.float().pow(2).mean().sqrt().item(),
-    )
-    torch.testing.assert_close(hs_attn_tail, hs_nat, rtol=2**-6, atol=2**-5)
-    digests = _comm().allgather((_digest(attn_tail_reduced), _digest(hs_attn_tail)))
-    assert len(set(digests)) == 1, "reduced outputs differ across ranks"
-    print(
-        f"[attention_output_tail-bridge] dense rank={rank}/{world} rows={rows} "
-        f"hs_maxdiff={(hs_attn_tail.float() - hs_nat.float()).abs().max().item():.3e}"
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py
deleted file mode 100644
index 5cba06b6e302..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o.py
+++ /dev/null
@@ -1,290 +0,0 @@
-# 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.
-"""The attention_output_tail contract on one GPU: no attention_output_tail entry point communicates.
-
-The attention_output_tail contract: "attention_output_tail owns W_UV + local o_proj.
-router_gemm/native all-reduce consumes the identical partial. No communication is hidden inside
-attention_output_tail."
-
-The real DSA ``MLA`` of TP8 rank 0 with real layer-3 weights gets an
-``o_proj`` that reduces its output, as in the model at TP8
-(``reduce_output=True``). A recording stand-in (``SpyAllReduce``) replaces
-the all-reduce module. It reports an NCCL symmetric-memory window, as the
-model's AUTO all-reduce does by default. The ``AttnTail`` entry points
-(``arm``, ``before_attention``, ``defer_w_uv``, ``project``) are driven with
-synthetic latents for every engine and named schedule the package ships:
-attention_output_tail's default, the ``attention_output_tail_prefetch`` and
-``attention_output_into_hop`` schedules (``INTO_HOP_SPARSE_DECODE_SCHEDULE`` is the served
-one), the into-hop publication, the ``fused`` and ``triton`` engines, and
-``attention_output_tail_prefetch``'s ``native`` engine with the cuBLAS and ring o_proj. The decoder
-layer's two ``AllReduceParams`` are used: enabled (dense layers, where ``Linear.forward`` reduces)
-and disabled (MoE layers, where the reduction is fused downstream).
-
-Checks:
-
-* no all-reduce call happens while any tail entry point runs;
-* ``project`` returns the rank-local partial: for the attention_output_tail kernels it equals
-  the functional ``attn_tail`` bit for bit; the ``native`` engine returns
-  the native ``o_proj`` GEMM output (bit-identical for cuBLAS, within the
-  frozen budget for the ring kernel) before any reduction;
-* captured in a CUDA graph, the tail still enqueues no all-reduce.
-
-``MLA.forward``'s reduction of this partial is tested at the real call site
-in ``test_attn_tail_contract_o_forward.py`` and ``test_attn_tail_mla_forward_mpi.py``.
-"""
-
-import dataclasses
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-import tail_reference  # noqa: E402
-import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
-    DEFAULT_CONFIG,
-    INTO_HOP_SCHEDULE,
-    INTO_HOP_SPARSE_DECODE_SCHEDULE,
-    TAIL_PREFETCH_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
-    TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE,
-    TailConfig,
-    attn_tail,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
-    AttnTail,
-)
-
-pytestmark = tail_adapter.pytestmark
-
-ROWS = 6
-_R = dataclasses.replace
-
-ATTN_TAIL_KERNEL_CONFIGS = {
-    "attention_output_tail": DEFAULT_CONFIG,
-    "attention_output_tail_prefetch": TAIL_PREFETCH_SCHEDULE,
-    "tail_prefetch_sparse_mla": TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
-    "tail_prefetch_sparse_mla_attn_hop": TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-    "attention_output_into_hop": INTO_HOP_SCHEDULE,
-    "into_hop_sparse_mla": INTO_HOP_SPARSE_DECODE_SCHEDULE,
-    "fused": TailConfig(oproj_engine="fused"),
-    "triton": TailConfig(oproj_engine="triton"),
-}
-
-
-TAIL_ENTRY_POINTS = ("arm", "before_attention", "defer_w_uv", "project")
-
-
-class SpyAllReduce:
-    """Stand-in for ``o_proj.all_reduce`` that records every call.
-
-    ``flag.in_tail`` is set while an attention_output_tail entry point runs. An enabled call
-    returns twice its input (two identical ranks), so a reduction that
-    happens is visible in the values; a disabled call returns its input, as
-    ``AllReduce.forward`` does.
-    """
-
-    def __init__(self, flag=None):
-        self.calls = []
-        self.flag = flag if flag is not None else SimpleNamespace(in_tail=False)
-
-    def uses_nccl_symmetric_memory_window(self) -> bool:
-        return True
-
-    def __call__(self, input, all_reduce_params=None):
-        self.calls.append(
-            SimpleNamespace(
-                in_tail=self.flag.in_tail, params=all_reduce_params, data_ptr=input.data_ptr()
-            )
-        )
-        if all_reduce_params is not None and not all_reduce_params.enable_allreduce:
-            return input
-        return input * 2
-
-
-def guard_tail(tail: AttnTail, spy: SpyAllReduce) -> None:
-    """Flag every all-reduce call made while a tail entry point runs."""
-    for name in TAIL_ENTRY_POINTS:
-        entry = getattr(tail, name)
-
-        def guarded(*args, _entry=entry, **kwargs):
-            spy.flag.in_tail = True
-            try:
-                return _entry(*args, **kwargs)
-            finally:
-                spy.flag.in_tail = False
-
-        setattr(tail, name, guarded)
-
-
-_MLA = {}
-
-
-def _reducing_mla():
-    """The real DSA MLA with an ``o_proj`` that reduces into a spy."""
-    if "m" not in _MLA:
-        mla = tail_adapter._build_mla()
-        mla.o_proj.reduce_output = True
-        _MLA["m"] = mla
-    mla = _MLA["m"]
-    spy = SpyAllReduce()
-    mla.o_proj.all_reduce = spy
-    return mla, spy
-
-
-def _params(enable: bool):
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    return AllReduceParams(enable_allreduce=enable)
-
-
-def _transport(rank=3):
-    """A post_attention_hop transport stand-in on one GPU (see attention_output_tail/test_into_hop_producer.py)."""
-    words = 8 * 8 * 6144 * 2 // 4
-    mc = torch.full((3 * words,), -(2**31), dtype=torch.int32, device="cuda")
-    flags = torch.tensor([1, 0, 0, 0], dtype=torch.int32, device="cuda")
-    return SimpleNamespace(mc_words=mc, flags=flags, rank=rank, cur=1)
-
-
-def _tail(cfg, mla, spy):
-    tail = AttnTail("layer3", cfg)
-    if cfg.successor_prefetch:
-        # The router weight an attention_output_tail_prefetch tail stages for the op after it.
-        tail.set_successor_weights([torch.zeros((256, 6144), dtype=torch.bfloat16, device="cuda")])
-    guard_tail(tail, spy)
-    return tail
-
-
-def _attn_tail_call(tail, mla, latent, params, hop_transport=None):
-    rows = latent.shape[0]
-    attn_out = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    if hop_transport is not None:
-        tail.request_into_hop(hop_transport, rows)
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point=tail.config.prefetch_fork)
-    assert tail.defer_w_uv(mla, meta, latent, attn_out)
-    return tail.project(mla, attn_out, params)
-
-
-@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
-@pytest.mark.parametrize("name", list(ATTN_TAIL_KERNEL_CONFIGS))
-def test_attn_tail_kernels_return_partial_without_communication(name, enable):
-    cfg = ATTN_TAIL_KERNEL_CONFIGS[name]
-    mla, spy = _reducing_mla()
-    latent = tail_adapter._latent(ROWS, seed=11)
-    want = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, ROWS, config=cfg)
-    tail = _tail(cfg, mla, spy)
-    got = _attn_tail_call(tail, mla, latent, _params(enable))
-    torch.cuda.synchronize()
-    assert spy.calls == [], "the attention_output_tail tail called the all-reduce"
-    # Engaged with a window all-reduce attached, as at TP8 in the model.
-    assert tail.static_reason is None
-    assert tail.engaged_calls == 1 and tail._pending is None
-    assert torch.equal(got, want), name
-    if cfg.successor_prefetch:
-        assert tail.prefetch_forks == 1 and not tail._prefetch.forked
-
-
-@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
-@pytest.mark.parametrize("name", ["attention_output_into_hop", "into_hop_sparse_mla"])
-def test_into_hop_publication_has_no_communication(name, enable):
-    cfg = ATTN_TAIL_KERNEL_CONFIGS[name]
-    mla, spy = _reducing_mla()
-    latent = tail_adapter._latent(ROWS, seed=12)
-    want = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, ROWS, config=cfg)
-    tail = _tail(cfg, mla, spy)
-    got = _attn_tail_call(tail, mla, latent, _params(enable), hop_transport=_transport())
-    torch.cuda.synchronize()
-    assert tail.take_into_hop_published()
-    assert spy.calls == []
-    assert torch.equal(got, want)
-
-
-def _native_engine_call(tail, mla, latent, params):
-    """What the op and MLA.forward do with attention_output_tail_prefetch's ``native`` engine."""
-    from tensorrt_llm._torch.attention.mla import MLA
-
-    rows = latent.shape[0]
-    attn_out = torch.zeros((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point="attn")
-    assert not tail.defer_w_uv(mla, meta, latent, attn_out)
-    v_view = attn_out.view(rows, mla.num_heads_tp, mla.v_head_dim)
-    MLA._bmm_bf16_out(
-        mla,
-        latent.transpose(0, 1),
-        mla.v_b_proj,
-        mla.v_b_proj.transpose(1, 2),
-        v_view.transpose(0, 1),
-    )
-    return tail.project(mla, attn_out, params), attn_out
-
-
-@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
-@pytest.mark.parametrize("join", ["before_oproj", "after_oproj"])
-@pytest.mark.parametrize("oproj", ["cublas", "ring"])
-def test_native_engine_returns_gemm_output_without_communication(oproj, join, enable):
-    mla, spy = _reducing_mla()
-    cfg = _R(TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE, native_oproj=oproj, prefetch_join=join)
-    tail = _tail(cfg, mla, spy)
-    latent = tail_adapter._latent(ROWS, seed=13)
-    got, attn_out = _native_engine_call(tail, mla, latent, _params(enable))
-    gemm = mla.o_proj.apply_linear(attn_out, None)
-    torch.cuda.synchronize()
-    assert spy.calls == []
-    assert tail.engaged_calls == 1 and tail.prefetch_forks == 1
-    assert not tail._prefetch.forked
-    if oproj == "cublas":
-        assert torch.equal(got, gemm)
-    else:
-        ref = attn_out.double() @ mla.o_proj.weight.double().t()
-        e_ring = tail_reference.error_stats(got, ref)
-        e_nat = tail_reference.error_stats(gemm, ref)
-        assert e_ring["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6
-        assert e_ring["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_ring["ref_max"])
-
-
-@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
-@pytest.mark.parametrize(
-    "name", ["attention_output_tail", "into_hop_sparse_mla", "attention_output_tail_prefetch"]
-)
-def test_graph_capture_enqueues_no_communication(name, enable):
-    cfg = ATTN_TAIL_KERNEL_CONFIGS[name]
-    mla, spy = _reducing_mla()
-    tail = _tail(cfg, mla, spy)
-    lat = tail_adapter._latent(ROWS, seed=14).contiguous()
-    params = _params(enable)
-    _attn_tail_call(tail, mla, lat, params)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = _attn_tail_call(tail, mla, lat, params)
-    assert spy.calls == []
-    for seed in range(4):
-        lat.copy_(tail_adapter._latent(ROWS, seed=140 + seed))
-        graph.replay()
-        torch.cuda.synchronize()
-        want = attn_tail(lat, mla.v_b_proj, mla.o_proj.weight, ROWS, config=cfg)
-        assert torch.equal(captured, want)
-    assert spy.calls == []
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py
deleted file mode 100644
index 4817785f65e8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_contract_o_forward.py
+++ /dev/null
@@ -1,276 +0,0 @@
-# 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.
-"""The attention_output_tail contract at the real call site: ``GlmMLA.forward`` reduces attention_output_tail's partial.
-
-Runs with the plan that runs only ``attention_output_tail`` (opt-in:
-``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED=1``).
-
-Builds TP8 rank 0 of real GLM-5.2 layers 2 and 3 (``real_mla_harness``: all
-attention weights real, FP8 KV cache filled by a real prefill) and runs the
-GEN step (R=6) as a decode step through ``GlmMLA.forward``, so the latent comes from the native
-DSA attention kernel. ``o_proj`` reduces its output, as in the model at TP8;
-on one GPU a recording stand-in (``test_attn_tail_contract_o.SpyAllReduce``)
-replaces the all-reduce. Both ``AllReduceParams`` of the decoder layer are
-used: enabled (dense destination, where ``Linear.forward`` reduces) and
-disabled (MoE destination, where the reduction is fused downstream). Tails:
-attention_output_tail's default schedule and ``INTO_HOP_SPARSE_DECODE_SCHEDULE``, the served
-one.
-
-Per layer call:
-
-* native: exactly one all-reduce call (from ``Linear.forward``) with the
-  call's ``AllReduceParams`` object;
-* attention_output_tail: the tail engages; no all-reduce call while a tail entry point runs;
-  exactly one call, from ``GlmMLA.forward``, with the same params object, on
-  the tensor ``project`` returned; the module output is that call's result,
-  and each arm is within the frozen attention_output_tail budget against the FP64 oracle of
-  the latent it received;
-* CUDA graph: the attention_output_tail step captured with both params records the same calls
-  and every replay with new inputs is within the frozen budget of the
-  native eager step on the same inputs.
-"""
-
-import os
-import sys
-from contextlib import nullcontext
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
-    DEFAULT_CONFIG,
-    INTO_HOP_SPARSE_DECODE_SCHEDULE,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
-    AttnTail,
-)
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-import real_mla_harness as harness  # noqa: E402
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-from test_attn_tail_contract_o import TAIL_ENTRY_POINTS, SpyAllReduce  # noqa: E402
-
-_CKPT = tail_fixtures.checkpoint_dir()
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="attention_output_tail targets SM100/SM103",
-    ),
-    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
-    pytest.mark.skipif(
-        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
-        reason="needs attention_output_tail selected (opt-in)",
-    ),
-    pytest.mark.threadleak(enabled=False),
-]
-
-LAYERS = [2, 3]
-CONFIGS = {
-    "attention_output_tail": DEFAULT_CONFIG,
-    "into_hop_sparse_mla": INTO_HOP_SPARSE_DECODE_SCHEDULE,
-}
-_STACK = {}
-
-
-@pytest.fixture
-def probe(monkeypatch):
-    """Instrument every ``AttnTail`` entry point (class level, so the
-    harness's per-instance latent recorder composes with it): set
-    ``in_tail`` while one runs and keep what ``project`` returned."""
-    state = SimpleNamespace(in_tail=False, returned={})
-    for name in TAIL_ENTRY_POINTS:
-        entry = getattr(AttnTail, name)
-
-        def run(self, *args, _entry=entry, _name=name, **kwargs):
-            state.in_tail = True
-            try:
-                out = _entry(self, *args, **kwargs)
-            finally:
-                state.in_tail = False
-            if _name == "project" and out is not None:
-                state.returned[self.layer_key] = out
-            return out
-
-        monkeypatch.setattr(AttnTail, name, run)
-    return state
-
-
-@pytest.fixture
-def setup(probe):
-    """The stack with a fresh spy all-reduce on every layer's ``o_proj``."""
-    if "s" not in _STACK:
-        stack = harness.AttentionStack(
-            _CKPT, LAYERS, tp_rank=0, ctx=4096, rows=6, fast_paths="attention_output_tail"
-        )
-        for layer in LAYERS:
-            assert stack.attn_tail_tails[layer] is not None, (
-                "no attention_output_tail tail: call sites or attention_output_tail missing"
-            )
-            stack.attn[layer].o_proj.reduce_output = True
-        _STACK["s"] = stack
-    stack = _STACK["s"]
-    spies = {}
-    for layer in LAYERS:
-        spies[layer] = SpyAllReduce(flag=probe)
-        stack.attn[layer].o_proj.all_reduce = spies[layer]
-    yield stack, spies, probe
-    stack.set_attn_tail(True, config=DEFAULT_CONFIG)
-
-
-def _params(enable):
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    return AllReduceParams(enable_allreduce=enable)
-
-
-def _step(stack, attention_output_tail, inputs, params, record=False, config=None):
-    """One GEN step over the stack; outputs are cloned as produced."""
-    stack.set_attn_tail(attention_output_tail, config=config)
-    pos = stack.position_ids()
-    outs = {}
-    recorder = harness.LatentRecorder(stack) if record else nullcontext()
-    with stack.scope(), recorder:
-        for layer in LAYERS:
-            outs[layer] = stack.forward(layer, inputs[layer], pos, params[layer]).clone()
-    torch.cuda.synchronize()
-    return outs, (dict(recorder.latents) if record else {})
-
-
-def _budget(stack, layer, attn_tail_out, lat_attn_tail, native_out, lat_native, scale):
-    """Each arm against the FP64 oracle of the latent it received."""
-    attn = stack.attn[layer]
-    e_attn_tail = tail_reference.error_stats(
-        attn_tail_out.float() / scale, harness.native_reference(attn, lat_attn_tail)
-    )
-    e_nat = tail_reference.error_stats(
-        native_out.float() / scale, harness.native_reference(attn, lat_native)
-    )
-    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (layer, e_attn_tail, e_nat)
-    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]), (
-        layer,
-        e_attn_tail,
-        e_nat,
-    )
-    return e_attn_tail, e_nat
-
-
-def _one_call_outside_tail(spy, params, label):
-    assert len(spy.calls) == 1, (label, spy.calls)
-    call = spy.calls[0]
-    assert not call.in_tail, f"{label}: all-reduce called inside attention_output_tail"
-    assert call.params is params, f"{label}: params object differs"
-    spy.calls.clear()
-    return call
-
-
-@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
-@pytest.mark.parametrize("name", list(CONFIGS))
-def test_call_site_reduces_like_native(setup, name, enable):
-    stack, spies, probe = setup
-    inputs = {layer: stack.gen_inputs(layer, step=0) for layer in LAYERS}
-    params = {layer: _params(enable) for layer in LAYERS}
-    scale = 2.0 if enable else 1.0
-
-    native, lat_native = _step(stack, False, inputs, params, record=True)
-    for layer in LAYERS:
-        _one_call_outside_tail(spies[layer], params[layer], f"native layer{layer}")
-    assert probe.returned == {}
-
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    attention_output_tail, lat_attn_tail = _step(
-        stack, True, inputs, params, record=True, config=CONFIGS[name]
-    )
-    for layer in LAYERS:
-        tail = stack.attn_tail_tails[layer]
-        assert tail.static_reason is None
-        assert tail.engaged_calls == engaged[layer] + 1, layer
-        call = _one_call_outside_tail(
-            spies[layer], params[layer], f"attention_output_tail layer{layer}"
-        )
-        partial = probe.returned[tail.layer_key]
-        assert call.data_ptr == partial.data_ptr(), (
-            "the all-reduce input is not the tensor attention_output_tail returned"
-        )
-        assert torch.equal(attention_output_tail[layer], partial * 2 if enable else partial)
-        e_attn_tail, e_nat = _budget(
-            stack,
-            layer,
-            attention_output_tail[layer],
-            lat_attn_tail[layer],
-            native[layer],
-            lat_native[layer],
-            scale,
-        )
-        print(
-            f"[attention_output_tail-contract-o] tail={name} allreduce={enable} layer={layer} "
-            f"err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e}"
-        )
-
-
-@pytest.mark.parametrize("enable", [True, False], ids=["allreduce_on", "allreduce_fused"])
-@pytest.mark.parametrize("name", list(CONFIGS))
-def test_call_site_graph_replay(setup, name, enable):
-    stack, spies, probe = setup
-    params = {layer: _params(enable) for layer in LAYERS}
-    scale = 2.0 if enable else 1.0
-    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
-    stack.set_attn_tail(True, config=CONFIGS[name])
-    pos = stack.position_ids()
-    with stack.scope():
-        for _ in range(2):
-            for layer in LAYERS:
-                stack.forward(layer, static_in[layer], pos, params[layer])
-        torch.cuda.synchronize()
-        for layer in LAYERS:
-            spies[layer].calls.clear()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph), harness.LatentRecorder(stack) as rec:
-            captured = {
-                layer: stack.forward(layer, static_in[layer], pos, params[layer])
-                for layer in LAYERS
-            }
-        captured_lat = dict(rec.latents)
-    for layer in LAYERS:
-        call = _one_call_outside_tail(spies[layer], params[layer], f"capture layer{layer}")
-        partial = probe.returned[stack.attn_tail_tails[layer].layer_key]
-        assert call.data_ptr == partial.data_ptr()
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    for step in range(1, 4):
-        for layer in LAYERS:
-            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
-        graph.replay()
-        torch.cuda.synchronize()
-        got = {layer: captured[layer].clone() for layer in LAYERS}
-        lat = {layer: captured_lat[layer].clone() for layer in LAYERS}
-        native, lat_native = _step(
-            stack, False, {layer: static_in[layer].clone() for layer in LAYERS}, params, record=True
-        )
-        for layer in LAYERS:
-            _one_call_outside_tail(spies[layer], params[layer], f"native step{step} layer{layer}")
-            e_attn_tail, e_nat = _budget(
-                stack, layer, got[layer], lat[layer], native[layer], lat_native[layer], scale
-            )
-            print(
-                f"[attention_output_tail-contract-o-graph] tail={name} allreduce={enable} "
-                f"step={step} layer={layer} err attention_output_tail={e_attn_tail['rel_rms']:.3e} "
-                f"native={e_nat['rel_rms']:.3e}"
-            )
-    # Replays run no Python: the tails' counters do not move.
-    assert engaged == {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py
deleted file mode 100644
index 03b999ef8aa1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_gates.py
+++ /dev/null
@@ -1,132 +0,0 @@
-# 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.
-"""attention_output_tail eligibility gates and registered-API state.
-
-* The static capability check declines attention-DP modules and o_proj
-  tensor-parallel widths other than 8 explicitly, before any other check or
-  GPU work, and the engaged decision logs its per-call domain.
-* The registered functional entry point never falls back to module-global
-  counters for the non-default fused engine: without a caller state it uses
-  zeroed counters private to the call, eagerly and under CUDA-graph replay.
-"""
-
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail import (  # noqa: E501
-    kernels,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
-    TailConfig,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
-    ENGAGED_DOMAIN,
-    SUPPORTED_TP_SIZE,
-    AttnTail,
-    _static_reason,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import tail_reference  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="attention_output_tail targets SM100/SM103",
-)
-
-LOCAL_HEADS, V_DIM, C_DIM, HIDDEN = 8, 256, 512, 6144
-
-
-def _module(tp_size, attention_dp):
-    from tensorrt_llm.mapping import Mapping
-
-    mapping = Mapping(world_size=8, tp_size=8, rank=0, enable_attention_dp=attention_dp)
-    mapping_o = Mapping(world_size=tp_size, tp_size=tp_size, rank=0)
-    # Anything past the gates would fail on this stub, so a returned reason
-    # proves the gate fired first.
-    return SimpleNamespace(mapping=mapping, mapping_o=mapping_o)
-
-
-def test_attention_dp_declines():
-    assert _static_reason(_module(8, True)) == "attention data parallelism"
-
-
-@pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
-def test_other_tp_sizes_decline(tp_size):
-    reason = _static_reason(_module(tp_size, False))
-    assert reason == (f"o_proj tensor-parallel size {tp_size} is not {SUPPORTED_TP_SIZE}")
-
-
-def test_declined_layer_takes_native_path_without_side_effects():
-    tail = AttnTail("adp")
-    meta = SimpleNamespace(num_contexts=0, num_tokens=6)
-    assert not tail.arm(_module(8, True), meta)
-    assert tail.static_reason == "attention data parallelism"
-    assert tail._pending is None and not tail._armed
-    assert tail._state is None
-    assert tail.engaged_calls == 0 and tail.declined_calls == 1
-    # The op-side hook then declines and touches nothing.
-    latent = torch.zeros(6, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
-    out = torch.zeros(6, LOCAL_HEADS * V_DIM, device="cuda").bfloat16()
-    assert not tail.defer_w_uv(None, meta, latent, out)
-    assert tail.project(None, out, None) is None
-
-
-def test_engaged_domain_text():
-    assert "GEN-only" in ENGAGED_DOMAIN and "TP8" in ENGAGED_DOMAIN
-    assert "1-8 token rows" in ENGAGED_DOMAIN
-
-
-def test_registered_fused_entry_has_no_module_state():
-    torch.manual_seed(5)
-    w_uv = (torch.randn(LOCAL_HEADS, V_DIM, C_DIM, device="cuda") * 0.02).bfloat16()
-    w_o = (torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, device="cuda") * 0.02).bfloat16()
-    latent = torch.randn(6, LOCAL_HEADS, C_DIM, device="cuda").bfloat16()
-    fused = TailConfig(oproj_engine="fused")
-    entry = dispatch.get_impl("attention_output_tail", "attn_tail")
-    assert entry is kernels.attn_tail
-    states_before = dict(kernels._DEFAULT_STATES)
-    first = entry(latent, w_uv, w_o, 6, config=fused)
-    second = entry(latent, w_uv, w_o, 6, config=fused)
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_fused import (  # noqa: E501
-        new_state,
-    )
-
-    own = new_state(latent.device)
-    explicit = entry(latent, w_uv, w_o, 6, config=fused, state=own)
-    torch.cuda.synchronize()
-    assert kernels._DEFAULT_STATES == states_before
-    assert torch.equal(first, second) and torch.equal(first, explicit)
-    native = tail_reference.native_attn_tail(latent, w_uv, w_o, 6)
-    torch.testing.assert_close(first, native, rtol=2**-6, atol=2**-6)
-
-    # Graph replay: the per-call counters are re-zeroed by every replay.
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = entry(latent, w_uv, w_o, 6, config=fused)
-    for step in range(16):
-        latent.copy_(torch.randn_like(latent, dtype=torch.float32).bfloat16())
-        graph.replay()
-        torch.cuda.synchronize()
-        expected = entry(latent, w_uv, w_o, 6, config=fused, state=own)
-        torch.cuda.synchronize()
-        assert torch.equal(captured, expected), step
-    assert kernels._DEFAULT_STATES == states_before
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py
deleted file mode 100644
index beb358aa9516..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_adapter.py
+++ /dev/null
@@ -1,359 +0,0 @@
-# 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.
-"""``AttnTail`` protocol on a real DSA ``MLA`` module with synthetic latents.
-
-Builds the GLM-5.2 attention module of TP8 rank 0 (DSA sparse attention, FP8
-KV cache, BF16 projections) with real ``v_b_proj``/``o_proj`` weights and
-drives the three ``AttnTail`` entry points directly (``arm``, ``defer_w_uv``,
-``project``) with synthetic latents (``tail_fixtures.mix_latent``) and
-``SimpleNamespace`` metadata. No attention kernel runs here; the real
-``MLA.forward`` path with the native attention producer is covered by
-``test_attn_tail_mla_forward.py`` and ``test_attn_tail_mla_forward_mpi.py``.
-
-* the static capability check accepts the module;
-* ``arm`` + ``defer_w_uv`` + ``project`` equal the native GEN tail
-  (``MLA._bmm_bf16_out`` into ``attn_output[0]`` followed by the native
-  ``MLA._project_output``) within the frozen attention_output_tail budget, for every row count;
-* ``defer_w_uv`` acts only on an armed call; ``arm`` declines (and touches no
-  state) under torch.compile tracing, TensorRT-LLM's compile policy and
-  breakable CUDA graphs; unsupported calls (context present, more than
-  ``MAX_ENGAGED_ROWS`` rows) are declined before any GPU work; a mismatched or
-  leftover packet raises; draft (MTP) layers get no tail;
-* the tail can be captured in a CUDA graph and replayed.
-"""
-
-import math
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
-    AttnTail,
-    make_attn_tail,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-_CKPT = tail_fixtures.checkpoint_dir()
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="attention_output_tail targets SM100/SM103",
-    ),
-    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
-]
-
-LAYER = 3
-
-
-def _rope_config():
-    return SimpleNamespace(
-        hidden_size=6144,
-        num_attention_heads=64,
-        rope_scaling={
-            "beta_fast": 32,
-            "beta_slow": 1,
-            "factor": 40,
-            "mscale": 1.0,
-            "mscale_all_dim": 1.0,
-            "original_max_position_embeddings": 4096,
-            "type": "yarn",
-        },
-        max_position_embeddings=4096,
-        rope_theta=10000.0,
-        qk_rope_head_dim=64,
-        model_type="deepseek_v2",
-    )
-
-
-def _build_mla(layer_idx=LAYER, num_hidden_layers=None):
-    from tensorrt_llm._torch.attention.mla import MLA
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
-    from tensorrt_llm.mapping import Mapping
-    from tensorrt_llm.models.modeling_utils import QuantConfig
-    from tensorrt_llm.quantization.mode import QuantAlgo
-
-    quant = QuantConfig()
-    quant.kv_cache_quant_algo = QuantAlgo.FP8
-    pretrained = SimpleNamespace(rms_norm_eps=1e-5)
-    if num_hidden_layers is not None:
-        pretrained.num_hidden_layers = num_hidden_layers
-    model_config = ModelConfig(
-        mapping=Mapping(world_size=8, tp_size=8, rank=0),
-        sparse_attention_config=DeepSeekSparseAttentionConfig(
-            index_n_heads=32, index_head_dim=128, index_topk=2048
-        ),
-        pretrained_config=pretrained,
-        quant_config=quant,
-    )
-    model_config.extra_attrs["kv_cache_dtype"] = "fp8"
-    # One process stands in for TP8 rank 0: the COLUMN Linears would build
-    # all-reduce workspaces over 8 local GPUs, which this test never uses.
-    import tensorrt_llm._torch.distributed as distributed
-
-    saved_all_reduce = distributed.AllReduce
-    distributed.AllReduce = lambda *args, **kwargs: None
-    try:
-        mla = _construct(MLA, model_config, layer_idx)
-    finally:
-        distributed.AllReduce = saved_all_reduce
-    w_uv, w_o = tail_fixtures.load_tail_weights(_CKPT, LAYER, [0])[0]
-    with torch.no_grad():
-        mla.v_b_proj.data = w_uv.cuda().contiguous()
-        mla.o_proj.weight.data.copy_(w_o.cuda())
-    return mla
-
-
-def _build_adopted_mla(layer_idx=LAYER, num_hidden_layers=None):
-    """``_build_mla``, adopted as the target's ``GlmMLA`` (which attaches the attention tail)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-        GlmMLA,
-    )
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
-
-    mla = _build_mla(layer_idx=layer_idx, num_hidden_layers=num_hidden_layers)
-    # The model's DeepseekV32Attention adds only its constructor to MLA.
-    mla.__class__ = DeepseekV32Attention
-    return GlmMLA.adopt(mla)
-
-
-def _construct(MLA, model_config, layer_idx):
-    from tensorrt_llm._torch.attention.backends.interface import (
-        PositionalEmbeddingParams,
-        RopeParams,
-    )
-    from tensorrt_llm.functional import PositionEmbeddingType
-
-    return MLA(
-        hidden_size=6144,
-        num_attention_heads=64,
-        num_key_value_heads=64,
-        qk_nope_head_dim=192,
-        qk_rope_head_dim=64,
-        v_head_dim=256,
-        q_lora_rank=2048,
-        kv_lora_rank=512,
-        predicted_tokens_per_seq=6,
-        max_position_embeddings=4096,
-        bias=False,
-        pos_embd_params=PositionalEmbeddingParams(
-            type=PositionEmbeddingType.yarn,
-            rope=RopeParams.from_config(_rope_config()),
-            is_neox=False,
-        ),
-        layer_idx=layer_idx,
-        dtype=torch.bfloat16,
-        config=model_config,
-        reduce_output=False,
-    ).cuda()
-
-
-_MLA = {}
-
-
-def _mla():
-    if "m" not in _MLA:
-        _MLA["m"] = _build_mla()
-    return _MLA["m"]
-
-
-def _native(mla, latent, out_rows):
-    from tensorrt_llm._torch.attention.mla import MLA
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    rows = latent.shape[0]
-    attn_out = torch.zeros(
-        (out_rows, mla.num_heads_tp * mla.v_head_dim), dtype=torch.bfloat16, device="cuda"
-    )
-    v_view = attn_out[:rows].view(rows, mla.num_heads_tp, mla.v_head_dim)
-    MLA._bmm_bf16_out(
-        mla,
-        latent.transpose(0, 1),
-        mla.v_b_proj,
-        mla.v_b_proj.transpose(1, 2),
-        v_view.transpose(0, 1),
-    )
-    meta = SimpleNamespace(num_contexts=0)
-    return mla._project_output([attn_out], None, meta, AllReduceParams(enable_allreduce=False))
-
-
-def _meta(rows, num_contexts=0):
-    return SimpleNamespace(num_contexts=num_contexts, num_tokens=rows)
-
-
-def _latent(rows, seed):
-    if "comp" not in _MLA:
-        _MLA["comp"] = tail_fixtures.compressed_kv_rows(_CKPT, LAYER)
-    return tail_fixtures.mix_latent(_MLA["comp"], rows, 8, seed=seed)
-
-
-def test_static_capability_accepts_dsa_mla():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
-        _static_reason,
-    )
-
-    mla = _mla()
-    assert mla.sparse_attn_hooks is not None
-    assert _static_reason(mla) is None
-
-
-@pytest.mark.parametrize("rows,out_rows", [(1, 1), (2, 2), (4, 4), (6, 6), (6, 8), (6, 20), (8, 8)])
-def test_defer_and_project_match_native(rows, out_rows):
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    mla = _mla()
-    latent = _latent(rows, seed=rows)
-    native = _native(mla, latent, out_rows)
-    tail = make_attn_tail(mla)
-    attn_out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = _meta(rows)
-    assert tail.arm(mla, meta)
-    assert tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
-    attention_output_tail = tail.project(mla, attn_out, AllReduceParams(enable_allreduce=False))
-    assert tail.project(mla, attn_out, None) is None  # packet consumed
-    ref, _ = tail_reference.oracle_attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows)
-    e_attn_tail = tail_reference.error_stats(attention_output_tail[:rows], ref)
-    e_nat = tail_reference.error_stats(native[:rows], ref)
-    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6
-    torch.testing.assert_close(
-        attention_output_tail[:rows], native[:rows], rtol=2**-6, atol=2**-6 * e_nat["ref_rms"]
-    )
-    assert bool((attention_output_tail[rows:] == 0).all())
-    assert tail.engaged_calls == 1 and tail.declined_calls == 0
-
-
-def test_declines_before_any_work():
-    mla = _mla()
-    tail = make_attn_tail(mla)
-    lat = _latent(6, seed=0)
-    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
-    # Not armed: the op-side entry never engages on its own.
-    assert not tail.defer_w_uv(mla, _meta(6), lat, out)
-    assert not tail.arm(mla, _meta(6, num_contexts=1))
-    assert not tail.defer_w_uv(mla, _meta(6, num_contexts=1), lat, out)
-    assert not tail.arm(mla, _meta(9))
-    # Armed, but the latent the op hands over does not fit the call.
-    assert tail.arm(mla, _meta(6))
-    big = torch.zeros((9, 8, 512), dtype=torch.bfloat16, device="cuda")
-    big_out = torch.empty((9, 2048), dtype=torch.bfloat16, device="cuda")
-    assert not tail.defer_w_uv(mla, _meta(6), big, big_out)
-    assert tail.declined_calls == 3 and tail.engaged_calls == 0
-    assert not tail._armed
-    assert tail.project(mla, out, None) is None
-
-
-def test_arm_declines_under_compile_and_breakable_graphs():
-    from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import enable_breakable_cuda_graph
-    from tensorrt_llm._torch.utils import torch_compiling
-
-    mla = _mla()
-    tail = make_attn_tail(mla)
-    with torch_compiling(True):
-        assert not tail.arm(mla, _meta(6))
-    with enable_breakable_cuda_graph():
-        assert not tail.arm(mla, _meta(6))
-    assert tail.declined_calls == 1  # the compile check touches no state
-
-    meta = _meta(6)
-
-    def traced(x):
-        return x + 1 if tail.arm(mla, meta) else x - 1
-
-    x = torch.zeros(4, device="cuda")
-    out = torch.compile(traced, backend="eager", fullgraph=True)(x)
-    assert torch.equal(out, x - 1)
-    assert not tail._armed and tail.declined_calls == 1
-
-
-def test_packet_errors_are_loud():
-    mla = _mla()
-    tail = make_attn_tail(mla)
-    meta = _meta(6)
-    lat = _latent(6, seed=1)
-    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
-    other = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
-    assert tail.arm(mla, meta)
-    assert tail.defer_w_uv(mla, meta, lat, out)
-    with pytest.raises(RuntimeError, match="not consumed"):
-        tail.arm(mla, meta)
-    with pytest.raises(RuntimeError, match="does not match"):
-        tail.project(mla, other, None)
-
-
-def test_draft_layers_keep_native_tail():
-    mla = _build_mla(layer_idx=78, num_hidden_layers=78)
-    assert make_attn_tail(mla) is None
-
-
-def test_graph_capture_and_replay():
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    mla = _mla()
-    tail = make_attn_tail(mla)
-    rows = 6
-    meta = _meta(rows)
-    lat = _latent(rows, seed=2).contiguous()
-    out = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    params = AllReduceParams(enable_allreduce=False)
-    # Eager warm-up (JIT, per-layer state) before capture, like the engine.
-    assert tail.arm(mla, meta)
-    assert tail.defer_w_uv(mla, meta, lat, out)
-    tail.project(mla, out, params)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        assert tail.arm(mla, meta)
-        assert tail.defer_w_uv(mla, meta, lat, out)
-        captured = tail.project(mla, out, params)
-    for seed in range(8):
-        lat.copy_(_latent(rows, seed=100 + seed))
-        graph.replay()
-        native = _native(mla, lat, rows)
-        torch.testing.assert_close(
-            captured.clone(),
-            native,
-            rtol=2**-6,
-            atol=2**-6 * native.float().pow(2).mean().sqrt().item(),
-        )
-    assert not math.isnan(captured.float().sum().item())
-
-
-@pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
-    reason="opt-in",
-)
-def test_adoption_attaches_tail_to_target_layers(monkeypatch):
-    """With attention_output_tail running, the adopted attention (``GlmMLA``) owns an attention_output_tail.
-
-    The call sites themselves run in ``test_attn_tail_mla_forward.py``.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-        dispatch,
-    )
-
-    plan_testlib.select(dispatch, "attention_output_tail", monkeypatch)
-    mla = _build_adopted_mla(layer_idx=LAYER, num_hidden_layers=78)
-    assert isinstance(mla._attn_tail, AttnTail)
-    draft = _build_adopted_mla(layer_idx=78, num_hidden_layers=78)
-    assert draft._attn_tail is None
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py
deleted file mode 100644
index 45a3230bbfb4..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward.py
+++ /dev/null
@@ -1,269 +0,0 @@
-# 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.
-"""attention_output_tail inside the target's ``GlmMLA.forward``: native DSA attention → attention_output_tail → output.
-
-Runs with the plan that runs only ``attention_output_tail`` (opt-in:
-``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED=1``).
-
-Builds TP8 rank 0 of real GLM-5.2 layers 2 (dense destination, full indexer)
-and 3 (MoE destination, shared indexer) as ``DeepseekV32Attention`` with all
-attention weights real, adopted as the target's ``GlmMLA``, fills an FP8 KV cache with a real
-4096-token prefill, and runs GEN-only ``GlmMLA.forward`` as a decode step through the
-``register_to_config`` path (the target's DSA attention op), as the model does. Every check compares
-attention_output_tail with the native tail on the same module, cache and inputs:
-
-* eager, R=6 (one request, MTP draft 5) and R=1: the latent the tail receives
-  comes from the native attention kernel and is identical in both arms; the
-  attention_output_tail output is within the frozen budget against the FP64 oracle of that
-  latent and matches native; attention_output_tail engages exactly once per layer and call;
-* CUDA-graph capture of the GEN step, replayed with new inputs;
-* under torch.compile (Dynamo traces ``MLA.forward``, the attention op stays
-  opaque) and under TensorRT-LLM's torch.compile policy, attention_output_tail declines in
-  ``MLA.forward`` before the op runs: the output equals native and no packet
-  or armed decision is left behind; a breakable CUDA graph also declines.
-
-Fixture labels: weights real, activations synthetic (real input_layernorm of
-real token embeddings).
-"""
-
-import os
-import sys
-from contextlib import nullcontext
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import real_mla_harness as harness  # noqa: E402
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-_CKPT = tail_fixtures.checkpoint_dir()
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="attention_output_tail targets SM100/SM103",
-    ),
-    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
-    pytest.mark.skipif(
-        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
-        reason="needs attention_output_tail selected (opt-in)",
-    ),
-    pytest.mark.threadleak(enabled=False),
-]
-
-LAYERS = [2, 3]
-CTX = 4096
-TAIL_PATHS = "attention_output_tail"
-_STACKS = {}
-
-
-def _stack(rows, fast_paths=TAIL_PATHS):
-    """The stack of ``rows`` GEN rows built (once) under the plan that runs ``fast_paths``."""
-    key = (rows, fast_paths)
-    if key not in _STACKS:
-        _STACKS[key] = harness.AttentionStack(
-            _CKPT, LAYERS, tp_rank=0, ctx=CTX, rows=rows, fast_paths=fast_paths
-        )
-    else:
-        harness.select(fast_paths)
-    return _STACKS[key]
-
-
-def _no_ar():
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    return AllReduceParams(enable_allreduce=False)
-
-
-def _run(stack, attention_output_tail, inputs, record=False, position_ids=None):
-    """One GEN step over the stack; returns per-layer outputs and latents."""
-    stack.set_attn_tail(attention_output_tail)
-    pos = stack.position_ids() if position_ids is None else position_ids
-    outs = {}
-    recorder = harness.LatentRecorder(stack) if record else nullcontext()
-    with stack.scope(), recorder:
-        for layer in stack.layers:
-            outs[layer] = stack.forward(layer, inputs[layer], pos, _no_ar()).clone()
-    torch.cuda.synchronize()
-    return outs, (dict(recorder.latents) if record else {})
-
-
-def _check_budget(stack, layer, attn_tail_out, native_out, latent):
-    ref = harness.native_reference(stack.attn[layer], latent)
-    e_attn_tail = tail_reference.error_stats(attn_tail_out, ref)
-    e_nat = tail_reference.error_stats(native_out, ref)
-    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (layer, e_attn_tail, e_nat)
-    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]), (
-        layer,
-        e_attn_tail,
-        e_nat,
-    )
-    torch.testing.assert_close(attn_tail_out, native_out, rtol=2**-6, atol=2**-6 * e_nat["ref_rms"])
-    return e_attn_tail, e_nat
-
-
-@pytest.mark.parametrize("rows", [6, 1])
-def test_eager_forward_matches_native(rows):
-    stack = _stack(rows)
-    inputs = {layer: stack.gen_inputs(layer, step=0) for layer in LAYERS}
-    native, lat_native = _run(stack, False, inputs, record=True)
-    # A second native run characterizes the run-to-run variation of the
-    # native attention latent itself.
-    _, lat_native2 = _run(stack, False, inputs, record=True)
-    before = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    attention_output_tail, lat_attn_tail = _run(stack, True, inputs, record=True)
-    for layer in LAYERS:
-        tail = stack.attn_tail_tails[layer]
-        assert tail.static_reason is None
-        assert tail.engaged_calls == before[layer] + 1, layer
-        assert tail._pending is None and not tail._armed
-        assert lat_native[layer].shape == (rows, 8, 512)
-        same = torch.equal(lat_native[layer], lat_attn_tail[layer])
-        native_rerun_same = torch.equal(lat_native[layer], lat_native2[layer])
-        native_rerun_diff = (
-            (lat_native[layer].float() - lat_native2[layer].float()).abs().max().item()
-        )
-        print(
-            f"[attention_output_tail-mla-latent] rows={rows} layer={layer} "
-            f"native_rerun_equal={native_rerun_same} "
-            f"native_rerun_maxdiff={native_rerun_diff:.3e} "
-            f"attn_tail_vs_native_equal={same}"
-        )
-        torch.testing.assert_close(lat_attn_tail[layer], lat_native[layer])
-        e_attn_tail, e_nat = _check_budget(
-            stack, layer, attention_output_tail[layer], native[layer], lat_attn_tail[layer]
-        )
-        print(
-            f"[attention_output_tail-mla-forward] rows={rows} layer={layer} latent_equal={same} "
-            f"err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e} "
-            f"maxdiff={(attention_output_tail[layer].float() - native[layer].float()).abs().max().item():.3e}"
-        )
-
-
-def test_cuda_graph_capture_and_replay():
-    stack = _stack(6)
-    pos = stack.position_ids()
-    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
-    stack.set_attn_tail(True)
-    with stack.scope():
-        for _ in range(2):
-            for layer in LAYERS:
-                stack.forward(layer, static_in[layer], pos, _no_ar())
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            captured = {
-                layer: stack.forward(layer, static_in[layer], pos, _no_ar()) for layer in LAYERS
-            }
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    for step in range(1, 5):
-        for layer in LAYERS:
-            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
-        graph.replay()
-        torch.cuda.synchronize()
-        got = {layer: captured[layer].clone() for layer in LAYERS}
-        native, lat = _run(
-            stack,
-            False,
-            {layer: static_in[layer].clone() for layer in LAYERS},
-            record=True,
-            position_ids=pos,
-        )
-        for layer in LAYERS:
-            e_attn_tail, e_nat = _check_budget(stack, layer, got[layer], native[layer], lat[layer])
-            print(
-                f"[attention_output_tail-mla-graph] step={step} layer={layer} "
-                f"err attention_output_tail={e_attn_tail['rel_rms']:.3e} native={e_nat['rel_rms']:.3e}"
-            )
-    # Replays run no Python: the engagement counters do not move.
-    assert engaged == {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    del graph
-
-
-def _assert_declined_and_native(stack, layer, out, native, engaged_before):
-    tail = stack.attn_tail_tails[layer]
-    assert tail.engaged_calls == engaged_before
-    assert tail._pending is None and not tail._armed
-    print(
-        f"[attention_output_tail-mla-declined] layer={layer} "
-        f"bitwise_equal_to_native={torch.equal(out, native)}"
-    )
-    torch.testing.assert_close(
-        out, native, rtol=2**-6, atol=2**-6 * native.float().pow(2).mean().sqrt().item()
-    )
-
-
-def test_torch_compile_takes_native_path():
-    stack = _stack(6)
-    inputs = {layer: stack.gen_inputs(layer, step=7) for layer in LAYERS}
-    native, _ = _run(stack, False, inputs)
-    pos = stack.position_ids()
-    stack.set_attn_tail(True)
-    compiled = {
-        layer: torch.compile(stack.attn[layer], backend="eager", dynamic=False) for layer in LAYERS
-    }
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    with stack.scope():
-        for _ in range(2):
-            outs = {
-                layer: compiled[layer](
-                    position_ids=pos,
-                    hidden_states=inputs[layer],
-                    attn_metadata=stack.gen_metadata,
-                    all_reduce_params=_no_ar(),
-                ).clone()
-                for layer in LAYERS
-            }
-            torch.cuda.synchronize()
-            for layer in LAYERS:
-                _assert_declined_and_native(
-                    stack, layer, outs[layer], native[layer], engaged[layer]
-                )
-    # The same module still takes attention_output_tail eagerly afterwards.
-    attention_output_tail, _ = _run(stack, True, inputs)
-    for layer in LAYERS:
-        assert stack.attn_tail_tails[layer].engaged_calls == engaged[layer] + 1
-        torch.testing.assert_close(
-            attention_output_tail[layer],
-            native[layer],
-            rtol=2**-6,
-            atol=2**-6 * native[layer].float().pow(2).mean().sqrt().item(),
-        )
-
-
-def test_trtllm_compile_policy_takes_native_path():
-    from tensorrt_llm._torch.utils import torch_compiling
-
-    stack = _stack(6)
-    inputs = {layer: stack.gen_inputs(layer, step=8) for layer in LAYERS}
-    native, _ = _run(stack, False, inputs)
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    with torch_compiling(True):
-        outs, _ = _run(stack, True, inputs)
-    for layer in LAYERS:
-        _assert_declined_and_native(stack, layer, outs[layer], native[layer], engaged[layer])
-
-
-def test_breakable_cuda_graph_declines():
-    from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import enable_breakable_cuda_graph
-
-    stack = _stack(6)
-    layer = LAYERS[0]
-    tail = stack.attn_tail_tails[layer]
-    with enable_breakable_cuda_graph():
-        assert not tail.arm(stack.attn[layer], stack.gen_metadata)
-    assert not tail._armed
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py
deleted file mode 100644
index 853abf9552b7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_mla_forward_mpi.py
+++ /dev/null
@@ -1,423 +0,0 @@
-# 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.
-"""Native attention → attention_output_tail → native reduction/norm, multi-rank, the target's GlmMLA.forward.
-
-Runs with the plan that runs only ``attention_output_tail`` (opt-in:
-``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED=1``),
-one process per GPU under
-``srun --mpi=pmix``. Rank ``r`` builds TP8 rank ``r`` of real layers 2 and 3
-(``real_mla_harness.AttentionStack``: all attention weights real, FP8 KV
-cache filled by a real prefill) and runs the GEN step (R=6, one request, MTP
-draft 5) as a decode step through ``GlmMLA.forward``, so the latent comes from the native DSA
-attention kernel inside the target's DSA attention op. Each rank's output
-feeds the native reduction exactly as the decoder layer does:
-
-* layer 2 (dense destination): ``o_proj`` reduces inside ``Linear.forward``
-  (``AllReduceParams(enable_allreduce=True)``), then the real
-  ``post_attention_layernorm(hidden, residual)``;
-* layer 3 (MoE destination): ``o_proj`` returns the rank-local partial
-  (``enable_allreduce=False``), then ``AllReduce(RESIDUAL_RMS_NORM)`` with the
-  real ``post_attention_layernorm`` weight.
-
-With eight ranks on two nodes (``-N2``) this is the GEN role's TP8 group;
-with fewer ranks the reduction sums the shards of ranks ``0 .. world-1``
-only, which still exercises the bridge.
-
-The native DSA attention output varies slightly between identical runs, so
-each arm is judged against its own oracle: every rank records the latent its
-tail received (``LatentRecorder``), the FP64 tail oracle of that latent (with
-the native BF16 rounding of ``v``) is summed over ranks in FP64, and the
-residual add and RMSNorm are applied in FP64. attention_output_tail's normalized RMS and max
-errors against its oracle must stay within the frozen attention_output_tail budget relative to
-native's errors against native's oracle. A second native run (A/A) is
-reported for scale. Results are gathered from all ranks before any assertion,
-so a failure on one rank cannot leave the others waiting in a collective.
-
-The attention_output_tail contract ("no communication is hidden inside attention_output_tail"): each
-``o_proj`` gets its own view of the shared all-reduce that counts calls and flags any call
-made while an ``AttnTail`` entry point runs. In both arms every layer call
-must reach the all-reduce exactly once, outside the tail, with the params
-object the caller passed: from ``Linear.forward`` natively, from
-``MLA.forward`` with attention_output_tail. A second test captures the attention_output_tail step
-(both destinations, real all-reduces) in a CUDA graph and replays it with new
-inputs against the same per-arm FP64 oracles.
-"""
-
-import hashlib
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import real_mla_harness as harness  # noqa: E402
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-_CKPT = tail_fixtures.checkpoint_dir()
-HIDDEN = 6144
-LAYERS = [2, 3]
-
-
-def _comm():
-    from mpi4py import MPI
-
-    return MPI.COMM_WORLD
-
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"),
-    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
-    pytest.mark.skipif(
-        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_ADAPTER_PATCH_APPLIED") != "1",
-        reason="needs attention_output_tail selected (opt-in)",
-    ),
-    pytest.mark.skipif(_comm().Get_size() < 2, reason="needs >= 2 MPI ranks"),
-    pytest.mark.threadleak(enabled=False),
-]
-
-_SETUP = {}
-TAIL_ENTRY_POINTS = ("arm", "before_attention", "defer_w_uv", "project")
-
-
-class _InTail:
-    """Set while an ``AttnTail`` entry point runs (class-level hooks, so
-    they compose with the harness's per-instance latent recorder)."""
-
-    active = False
-
-
-def _hook_tail_entry_points():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E501
-        AttnTail,
-    )
-
-    if getattr(AttnTail, "_contract_o_hooked", False):
-        return
-    for name in TAIL_ENTRY_POINTS:
-        entry = getattr(AttnTail, name)
-
-        def run(self, *args, _entry=entry, **kwargs):
-            _InTail.active = True
-            try:
-                return _entry(self, *args, **kwargs)
-            finally:
-                _InTail.active = False
-
-        setattr(AttnTail, name, run)
-    AttnTail._contract_o_hooked = True
-
-
-class _CountingAllReduce:
-    """One ``o_proj``'s view of the shared ``AllReduce``: forwards every
-    call and records whether it was made inside the attention_output_tail and with which
-    params object."""
-
-    def __init__(self, inner):
-        self.inner = inner
-        self.calls = []
-
-    def __call__(self, input, all_reduce_params=None):
-        self.calls.append((_InTail.active, all_reduce_params))
-        return self.inner(input, all_reduce_params=all_reduce_params)
-
-    def __getattr__(self, name):
-        return getattr(self.inner, name)
-
-
-def _setup():
-    if "s" in _SETUP:
-        return _SETUP["s"]
-    from tensorrt_llm._torch.distributed import AllReduce
-    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-    from tensorrt_llm.mapping import Mapping
-
-    comm = _comm()
-    rank, world = comm.Get_rank(), comm.Get_size()
-    torch.cuda.set_device(rank % torch.cuda.device_count())
-    mapping = Mapping(
-        world_size=world, tp_size=world, rank=rank, gpus_per_node=torch.cuda.device_count()
-    )
-    stack = harness.AttentionStack(
-        _CKPT, LAYERS, tp_rank=rank, ctx=4096, rows=6, fast_paths="attention_output_tail"
-    )
-    # As in the model: o_proj owns an all-reduce; the decoder layer enables
-    # it for dense layers and fuses it into the post-attention norm for MoE.
-    allreduce = AllReduce(mapping=mapping, dtype=torch.bfloat16)
-    for attn in stack.attn.values():
-        attn.o_proj.reduce_output = True
-        attn.o_proj.all_reduce = _CountingAllReduce(allreduce)
-    _hook_tail_entry_points()
-    eps = stack.pretrained.rms_norm_eps
-    norms = {}
-    for layer in LAYERS:
-        norm = RMSNorm(hidden_size=HIDDEN, eps=eps, dtype=torch.bfloat16).cuda()
-        norm.weight.data.copy_(
-            tail_fixtures.read_tensor(
-                _CKPT, f"model.layers.{layer}.post_attention_layernorm.weight"
-            ).cuda()
-        )
-        norms[layer] = norm
-    _SETUP["s"] = (rank, world, stack, norms, allreduce, eps)
-    return _SETUP["s"]
-
-
-def _digest(t):
-    return hashlib.sha256(t.contiguous().view(torch.uint8).cpu().numpy().tobytes()).hexdigest()
-
-
-def _chain(
-    stack, norms, allreduce, eps, attention_output_tail, inputs, residual, params=None, sync=True
-):
-    """GEN step of layers 2 (dense) and 3 (MoE) with their native reductions.
-
-    Returns the outputs and the latent each layer's tail received; ``params``
-    (if given) receives the ``AllReduceParams`` object passed to each layer.
-    """
-    from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams
-
-    stack.set_attn_tail(attention_output_tail)
-    pos = stack.position_ids()
-    params = {} if params is None else params
-    params[2] = AllReduceParams(enable_allreduce=True)
-    params[3] = AllReduceParams(enable_allreduce=False)
-    # Outputs are cloned right after they are produced: an all-reduce result
-    # may live in a workspace that the next all-reduce reuses.
-    out = {}
-    with stack.scope(), harness.LatentRecorder(stack) as recorder:
-        reduced = stack.forward(2, inputs[2], pos, params[2])
-        out["dense_reduced"] = reduced.clone()
-        hs2, res2 = norms[2](reduced, residual.clone())
-        out["dense_hidden"], out["dense_residual"] = hs2.clone(), res2.clone()
-        partial = stack.forward(3, inputs[3], pos, params[3])
-        out["moe_partial"] = partial.clone()
-        hs3, res3 = allreduce(
-            partial,
-            all_reduce_params=AllReduceParams(
-                fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
-                residual=residual.clone(),
-                norm_weight=norms[3].weight,
-                eps=eps,
-                trigger_completion_at_end=False,
-            ),
-        )
-        out["moe_hidden"], out["moe_residual"] = hs3.clone(), res3.clone()
-        latents = dict(recorder.latents)
-    if sync:
-        torch.cuda.synchronize()
-    return out, latents
-
-
-def _take_calls(stack, params, label):
-    """The attention_output_tail contract: one all-reduce call per layer, outside the tail, with
-    the caller's params object. Returns failure strings and clears."""
-    failures = []
-    for layer in LAYERS:
-        calls = stack.attn[layer].o_proj.all_reduce.calls
-        ok = len(calls) == 1 and not calls[0][0] and calls[0][1] is params[layer]
-        if not ok:
-            failures.append(
-                f"{label} layer{layer}: o_proj all-reduce calls "
-                f"(in_tail, same_params) = "
-                f"{[(c[0], c[1] is params[layer]) for c in calls]}"
-            )
-        calls.clear()
-    return failures
-
-
-def _oracle(stack, norms, eps, latents, residual):
-    """FP64 references of ``_chain``'s outputs from the recorded latents."""
-    import numpy as np
-    from mpi4py import MPI
-
-    def summed(local):
-        send = local.cpu().numpy()
-        recv = np.empty_like(send)
-        _comm().Allreduce(send, recv, op=MPI.SUM)
-        return torch.from_numpy(recv).cuda()
-
-    def norm(x, weight):
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight.double()
-
-    res_in = residual.double()
-    partial = {}
-    for layer in LAYERS:
-        attn = stack.attn[layer]
-        partial[layer], _ = tail_reference.oracle_attn_tail(
-            latents[layer], attn.v_b_proj, attn.o_proj.weight, stack.rows
-        )
-    dense = summed(partial[2])
-    moe = summed(partial[3])
-    return {
-        "dense_reduced": dense,
-        "dense_hidden": norm(dense + res_in, norms[2].weight),
-        "dense_residual": dense + res_in,
-        "moe_partial": partial[3],
-        "moe_hidden": norm(moe + res_in, norms[3].weight),
-        "moe_residual": moe + res_in,
-    }
-
-
-@pytest.mark.parametrize("step", [0, 1])
-def test_native_producer_attn_tail_native_reduction(step):
-    rank, world, stack, norms, allreduce, eps = _setup()
-    inputs = {layer: stack.gen_inputs(layer, step=step) for layer in LAYERS}
-    gen = torch.Generator(device="cuda").manual_seed(4321 + step)
-    residual = torch.randn(
-        (stack.rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
-    ).bfloat16()
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    for layer in LAYERS:
-        stack.attn[layer].o_proj.all_reduce.calls.clear()
-    failures, report = [], []
-    params = {}
-    native, lat_native = _chain(stack, norms, allreduce, eps, False, inputs, residual, params)
-    failures += _take_calls(stack, params, "native")
-    native_b, _ = _chain(stack, norms, allreduce, eps, False, inputs, residual, params)
-    failures += _take_calls(stack, params, "native_b")
-    attention_output_tail, lat_attn_tail = _chain(
-        stack, norms, allreduce, eps, True, inputs, residual, params
-    )
-    failures += _take_calls(stack, params, "attention_output_tail")
-    report.append("contract_o_calls=" + ("ok" if not failures else "BAD"))
-    ref_native = _oracle(stack, norms, eps, lat_native, residual)
-    ref_attn_tail = _oracle(stack, norms, eps, lat_attn_tail, residual)
-
-    for layer in LAYERS:
-        tail = stack.attn_tail_tails[layer]
-        # As in the model (default TLLM_NCCL_SYMMETRIC_ZERO_COPY, AUTO
-        # strategy, TP > 1), the o_proj all-reduce may use an NCCL window,
-        # but the native o_proj GEMM output is an ordinary F.linear buffer,
-        # so attention_output_tail still takes the call.
-        report.append(
-            f"layer{layer}: nccl_window_allreduce="
-            f"{stack.attn[layer].o_proj.all_reduce.uses_nccl_symmetric_memory_window()} "
-            f"static_reason={tail.static_reason} "
-            f"engaged={tail.engaged_calls - engaged[layer]}"
-        )
-        if tail.static_reason is not None or (tail.engaged_calls != engaged[layer] + 1):
-            failures.append(report[-1])
-    for key in native:
-        e_attn_tail = tail_reference.error_stats(attention_output_tail[key], ref_attn_tail[key])
-        e_nat = tail_reference.error_stats(native[key], ref_native[key])
-        aa = tail_reference.error_stats(native_b[key], native[key].double())
-        report.append(
-            f"{key}: rel_rms attention_output_tail={e_attn_tail['rel_rms']:.2e} "
-            f"native={e_nat['rel_rms']:.2e} "
-            f"max attention_output_tail={e_attn_tail['max_abs']:.2e} "
-            f"native={e_nat['max_abs']:.2e} "
-            f"(A/A rel_rms {aa['rel_rms']:.2e})"
-        )
-        if (
-            e_attn_tail["rel_rms"] > 1.25 * e_nat["rel_rms"] + 1e-6
-            or e_attn_tail["max_abs"] > 2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]
-        ):
-            failures.append(report[-1])
-    digests = tuple(
-        _digest(attention_output_tail[k])
-        for k in ("dense_reduced", "dense_hidden", "moe_hidden", "moe_residual")
-    )
-    gathered = _comm().allgather((rank, failures, report, digests))
-    if rank == 0:
-        for r, _, rep, _ in gathered:
-            print(
-                f"[attention_output_tail-mla-forward-mpi] rank={r}/{world} step={step} | "
-                + " | ".join(rep)
-            )
-    bad = [(r, f) for r, f, _, _ in gathered if f]
-    assert not bad, bad
-    assert len({d for _, _, _, d in gathered}) == 1, "reduced outputs differ across ranks"
-
-
-def test_graph_replay_native_reduction():
-    """The attention_output_tail step of both destinations, real all-reduces included, captured
-    in one CUDA graph and replayed with new inputs and residuals; each replay
-    is judged like the eager step, against native eager on the same inputs."""
-    rank, world, stack, norms, allreduce, eps = _setup()
-    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
-    static_res = torch.zeros((stack.rows, HIDDEN), dtype=torch.bfloat16, device="cuda")
-    for layer in LAYERS:
-        stack.attn[layer].o_proj.all_reduce.calls.clear()
-    failures, report = [], []
-    params = {}
-    for _ in range(2):  # eager warm-up (JIT, per-layer state), as the engine
-        _chain(stack, norms, allreduce, eps, True, static_in, static_res, params)
-        failures += _take_calls(stack, params, "warm-up")
-    _comm().Barrier()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured, captured_lat = _chain(
-            stack, norms, allreduce, eps, True, static_in, static_res, params, sync=False
-        )
-    failures += _take_calls(stack, params, "capture")
-    engaged = {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}
-    digests = []
-    for step in range(1, 4):
-        for layer in LAYERS:
-            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
-        gen = torch.Generator(device="cuda").manual_seed(4321 + step)
-        static_res.copy_(
-            torch.randn(
-                (stack.rows, HIDDEN), generator=gen, device="cuda", dtype=torch.float32
-            ).bfloat16()
-        )
-        _comm().Barrier()
-        graph.replay()
-        torch.cuda.synchronize()
-        got = {k: v.clone() for k, v in captured.items()}
-        lat = {k: v.clone() for k, v in captured_lat.items()}
-        native, lat_native = _chain(
-            stack,
-            norms,
-            allreduce,
-            eps,
-            False,
-            {layer: static_in[layer].clone() for layer in LAYERS},
-            static_res.clone(),
-            params,
-        )
-        failures += _take_calls(stack, params, f"native step{step}")
-        ref_attn_tail = _oracle(stack, norms, eps, lat, static_res)
-        ref_native = _oracle(stack, norms, eps, lat_native, static_res)
-        for key in got:
-            e_attn_tail = tail_reference.error_stats(got[key], ref_attn_tail[key])
-            e_nat = tail_reference.error_stats(native[key], ref_native[key])
-            report.append(
-                f"step{step} {key}: rel_rms "
-                f"graph={e_attn_tail['rel_rms']:.2e} "
-                f"native={e_nat['rel_rms']:.2e}"
-            )
-            if (
-                e_attn_tail["rel_rms"] > 1.25 * e_nat["rel_rms"] + 1e-6
-                or e_attn_tail["max_abs"] > 2 * e_nat["max_abs"] + 2**-8 * e_attn_tail["ref_max"]
-            ):
-                failures.append(report[-1])
-        digests.append(
-            tuple(_digest(got[k]) for k in ("dense_reduced", "dense_hidden", "moe_hidden"))
-        )
-    if engaged != {layer: stack.attn_tail_tails[layer].engaged_calls for layer in LAYERS}:
-        failures.append("replays ran tail Python")
-    gathered = _comm().allgather((rank, failures, report, tuple(digests)))
-    if rank == 0:
-        for r, _, rep, _ in gathered:
-            print(
-                f"[attention_output_tail-mla-forward-mpi-graph] rank={r}/{world} | "
-                + " | ".join(rep)
-            )
-    bad = [(r, f) for r, f, _, _ in gathered if f]
-    assert not bad, bad
-    assert len({d for _, _, _, d in gathered}) == 1, "replayed reduced outputs differ across ranks"
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py
deleted file mode 100644
index a2313306d43e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_pdl_stress.py
+++ /dev/null
@@ -1,506 +0,0 @@
-# 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.
-"""PDL boundary stress of the default attention_output_tail engine (``cute_stream``).
-
-The default attention_output_tail is two kernels with programmatic dependent launch (PDL) at
-both package boundaries:
-
-* ``_wuv_kernel`` is a PDL dependent of the latent's producer. It triggers its
-  own dependent at entry and loads ``w_uv`` before ``griddepcontrol.wait``;
-  it reads the latent and writes ``v`` only after the wait.
-* ``stream_oproj`` is a PDL dependent of ``_wuv_kernel``. It copies ``w_o``
-  rows into its own shared memory before its wait, reads ``v`` and writes the
-  partial after it, and triggers its own PDL consumer before it has written
-  the partial.
-
-The boundary contract is therefore: before its wait an attention_output_tail kernel touches only
-immutable weights and its own shared memory, and every consumer of the
-partial waits before reading it. These tests exercise that contract with
-synthetic neighbour kernels that stretch the overlap windows:
-
-* **Delayed producer.** A predecessor triggers its dependents at entry, spins
-  for ``delay`` µs, then writes the latent. The latent and ``v`` are poisoned
-  with NaN before the producer runs. An attention_output_tail read of the latent or of ``v`` that
-  is not ordered after the producer (or after ``_wuv_kernel``) shows up as
-  NaN or as another iteration's values.
-* **Delayed consumer with rapid buffer reuse.** Every kernel of the chain
-  producer → attention_output_tail → consumer is PDL-launched and triggers its dependents at
-  entry, so the next iteration's producer and attention_output_tail kernels may become resident
-  while the current consumer still runs. The consumer waits, spins, then
-  copies the partial. attention_output_tail's ``v`` and partial come from the CUDA-graph pool and
-  are reused across iterations. An attention_output_tail global write issued before its wait would
-  overwrite the partial the sleeping consumer has not read yet.
-
-Every captured iteration has its own latent, and the partial the consumer
-copies must equal, bit for bit, an eager attention_output_tail run on that latent
-(attention_output_tail is deterministic). The graphs are replayed many times with new latents. The
-consumer records ``%globaltimer`` at entry and after its wait, so the tests
-can report (and ``test_pdl_overlap_witness`` asserts) that the consumer really
-was resident before attention_output_tail finished.
-
-Weights: real GLM-5.2 ``v_b_proj`` / ``o_proj`` TP8 shards when
-``GLM52_NVFP4_CKPT`` is set, otherwise synthetic. Latents: real-weight,
-synthetic-activation mixtures of real compressed-KV rows when the checkpoint
-is available, otherwise random.
-
-``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_QUICK=1`` shrinks iteration and replay counts
-(compute-sanitizer runs); ``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SANITIZER=1`` skips the timing
-witness, because the sanitizer serializes kernels;
-``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SYNTHETIC=1`` uses synthetic weights and latents even when
-the checkpoint is available.
-"""
-
-import os
-import statistics
-import sys
-
-import pytest
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
-    DEFAULT_CONFIG,
-    TailConfig,
-    attn_tail,
-    check_tail_shapes,
-    pdl_enabled,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import tail_fixtures  # noqa: E402
-import tail_reference  # noqa: E402
-
-TP_SIZE = 8
-LOCAL_HEADS, V_DIM, C_DIM, HIDDEN = 8, 256, 512, 6144
-QUICK = os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_QUICK") == "1"
-UNDER_SANITIZER = os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SANITIZER") == "1"
-N_ITERS = 3 if QUICK else 16
-N_REPLAYS = 2 if QUICK else 24
-_CKPT = tail_fixtures.checkpoint_dir()
-# GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SYNTHETIC=1 forces synthetic weights and latents built by
-# elementwise kernels only (compute-sanitizer initcheck: the real-weight
-# fixture runs a cuBLAS GEMM whose output initcheck cannot see as written).
-_HAVE_CKPT = (
-    bool(_CKPT and os.path.isdir(_CKPT))
-    and os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PDL_STRESS_SYNTHETIC") != "1"
-)
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="attention_output_tail kernels target SM100/SM103",
-    ),
-    pytest.mark.skipif(not pdl_enabled(), reason="the stress needs TRTLLM_ENABLE_PDL=1"),
-]
-
-# A fully serialized attention_output_tail: no PDL at entry, inside, or at exit.
-SERIAL_CONFIG = TailConfig(pdl_entry=False, trigger_dependents=False)
-
-# ---------------------------------------------------------------------------
-# Synthetic neighbour kernels.
-# ---------------------------------------------------------------------------
-
-
-@triton.jit
-def _clock_fenced():
-    # The memory clobber keeps loads and stores from moving across the read.
-    return tl.inline_asm_elementwise(
-        "mov.u64 $0, %globaltimer;", "=l,~{memory}", [], dtype=tl.int64, is_pure=False, pack=1
-    )
-
-
-@triton.jit
-def _gdc_wait():
-    tl.inline_asm_elementwise(
-        "griddepcontrol.wait; // dummy $0",
-        "=r,~{memory}",
-        [],
-        dtype=tl.int32,
-        is_pure=False,
-        pack=1,
-    )
-
-
-@triton.jit
-def _gdc_launch_dependents():
-    tl.inline_asm_elementwise(
-        "griddepcontrol.launch_dependents; // dummy $0",
-        "=r,~{memory}",
-        [],
-        dtype=tl.int32,
-        is_pure=False,
-        pack=1,
-    )
-
-
-@triton.jit
-def _spin(delay_ns):
-    t0 = _clock_fenced()
-    t = t0
-    while t - t0 < delay_ns:
-        t = _clock_fenced()
-    return t
-
-
-@triton.jit
-def _delayed_producer(
-    src_ptr,
-    dst_ptr,
-    stamp_ptr,
-    numel,
-    delay_ns,
-    BLOCK: tl.constexpr,
-    TRIGGER: tl.constexpr,
-    WAIT: tl.constexpr,
-):
-    """Copy ``src`` to ``dst`` after spinning ``delay_ns``.
-
-    ``TRIGGER`` lets the PDL dependent (attention_output_tail's W_UV kernel) launch at entry;
-    ``WAIT`` (when itself PDL-launched) waits for its predecessor before the
-    spin. ``stamp[0]`` = time the copy starts.
-    """
-    pid = tl.program_id(0)
-    if TRIGGER:
-        _gdc_launch_dependents()
-    if WAIT:
-        _gdc_wait()
-    t = _spin(delay_ns)
-    offs = pid * BLOCK + tl.arange(0, BLOCK)
-    mask = offs < numel
-    x = tl.load(src_ptr + offs, mask=mask)
-    tl.store(dst_ptr + offs, x, mask=mask)
-    if pid == 0:
-        tl.store(stamp_ptr, t)
-
-
-@triton.jit
-def _delayed_consumer(
-    src_ptr, dst_ptr, stamp_ptr, numel, delay_ns, BLOCK: tl.constexpr, TRIGGER: tl.constexpr
-):
-    """PDL consumer of the attention_output_tail partial: wait, spin ``delay_ns``, copy.
-
-    ``stamp[0]`` = entry time (before the wait), ``stamp[1]`` = time the
-    wait returned, ``stamp[2]`` = time the copy starts.
-    """
-    pid = tl.program_id(0)
-    t_entry = _clock_fenced()
-    if TRIGGER:
-        _gdc_launch_dependents()
-    _gdc_wait()
-    t_wait = _clock_fenced()
-    t_copy = _spin(delay_ns)
-    offs = pid * BLOCK + tl.arange(0, BLOCK)
-    mask = offs < numel
-    x = tl.load(src_ptr + offs, mask=mask)
-    tl.store(dst_ptr + offs, x, mask=mask)
-    if pid == 0:
-        tl.store(stamp_ptr, t_entry)
-        tl.store(stamp_ptr + 1, t_wait)
-        tl.store(stamp_ptr + 2, t_copy)
-
-
-_BLOCK = 1024
-
-
-def launch_producer(src, dst, stamp, delay_us, trigger, wait, pdl):
-    numel = dst.numel()
-    _delayed_producer[(triton.cdiv(numel, _BLOCK),)](
-        src,
-        dst,
-        stamp,
-        numel,
-        int(delay_us * 1000),
-        BLOCK=_BLOCK,
-        TRIGGER=trigger,
-        WAIT=wait,
-        num_warps=4,
-        **({"launch_pdl": True} if pdl else {}),
-    )
-
-
-def launch_consumer(src, dst, stamp, delay_us, trigger):
-    numel = src.numel()
-    _delayed_consumer[(triton.cdiv(numel, _BLOCK),)](
-        src,
-        dst,
-        stamp,
-        numel,
-        int(delay_us * 1000),
-        BLOCK=_BLOCK,
-        TRIGGER=trigger,
-        num_warps=4,
-        launch_pdl=True,
-    )
-
-
-# ---------------------------------------------------------------------------
-# Fixtures.
-# ---------------------------------------------------------------------------
-
-_CACHE = {}
-
-
-def tail_weights(layer=3, rank=0):
-    key = ("w", layer, rank)
-    if key not in _CACHE:
-        if _HAVE_CKPT:
-            w_uv, w_o = tail_fixtures.load_tail_weights(_CKPT, layer, [rank], TP_SIZE)[rank]
-            _CACHE[key] = (w_uv.cuda(), w_o.cuda(), "real")
-        else:
-            gen = torch.Generator(device="cuda").manual_seed(layer * 8 + rank)
-            w_uv = (
-                torch.randn(LOCAL_HEADS, V_DIM, C_DIM, generator=gen, device="cuda") * 0.02
-            ).bfloat16()
-            w_o = (
-                torch.randn(HIDDEN, LOCAL_HEADS * V_DIM, generator=gen, device="cuda") * 0.02
-            ).bfloat16()
-            _CACHE[key] = (w_uv, w_o, "synthetic")
-    return _CACHE[key]
-
-
-def make_latents(count, rows, seed, layer=3):
-    """``[count, rows, 8, 512]`` BF16 latents, one per captured iteration."""
-    if _HAVE_CKPT:
-        key = ("kv", layer)
-        if key not in _CACHE:
-            _CACHE[key] = tail_fixtures.compressed_kv_rows(_CKPT, layer)
-        return torch.stack(
-            [
-                tail_fixtures.mix_latent(
-                    _CACHE[key], rows, num_local_heads=LOCAL_HEADS, seed=seed + i
-                )
-                for i in range(count)
-            ]
-        )
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    return torch.randn((count, rows, LOCAL_HEADS, C_DIM), generator=gen, device="cuda").bfloat16()
-
-
-def eager_expected(latents, w_uv, w_o, rows, out_rows, config):
-    """attention_output_tail on each latent, one isolated call at a time."""
-    outs = []
-    for i in range(latents.shape[0]):
-        lat = latents[i].clone()
-        torch.cuda.synchronize()
-        outs.append(attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows, config=config).clone())
-        torch.cuda.synchronize()
-    return torch.stack(outs)
-
-
-def check_against_native(latents, expected, w_uv, w_o, rows, label):
-    """Frozen attention_output_tail budget (test_attn_tail_attn_tail) on the first latent."""
-    native = tail_reference.native_attn_tail(latents[0], w_uv, w_o, rows)
-    ref, _ = tail_reference.oracle_attn_tail(latents[0], w_uv, w_o, rows)
-    e_attn_tail = tail_reference.error_stats(expected[0][:rows], ref)
-    e_nat = tail_reference.error_stats(native, ref)
-    assert e_attn_tail["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (label, e_attn_tail, e_nat)
-    assert e_attn_tail["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_nat["ref_max"]), (
-        label,
-        e_attn_tail,
-        e_nat,
-    )
-
-
-def assert_bitwise(caps, expected, label):
-    assert not bool(torch.isnan(caps).any()), (label, "NaN in a partial")
-    bad = [i for i in range(caps.shape[0]) if not torch.equal(caps[i], expected[i])]
-    if bad:
-        i = bad[0]
-        diff = (caps[i].float() - expected[i].float()).abs()
-        pytest.fail(
-            f"{label}: {len(bad)}/{caps.shape[0]} iterations differ "
-            f"from the eager attention_output_tail result (first {i}: "
-            f"{int((diff > 0).sum())} elements, max |diff| "
-            f"{diff.max().item():.3e})"
-        )
-
-
-def _median_us(values):
-    return statistics.median(values) / 1000.0 if values else float("nan")
-
-
-# ---------------------------------------------------------------------------
-# Tests.
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1), (8, 8)])
-def test_eager_exact_buffers(rows, out_rows):
-    """Exactly sized latent / v / partial (memcheck with the caching
-    allocator off sees every out-of-bounds access), default config vs the
-    fully serialized config, bit for bit, and the frozen budget vs native."""
-    w_uv, w_o, kind = tail_weights()
-    lat = make_latents(1, rows, seed=11)[0].contiguous()
-    assert check_tail_shapes(lat, w_uv, w_o, rows, out_rows) is None
-    v = torch.empty((rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
-    out = attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows, v_scratch=v)
-    serial = attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows, use_pdl=False, config=SERIAL_CONFIG)
-    torch.cuda.synchronize()
-    assert torch.equal(out, serial), "PDL and serialized attention_output_tail differ"
-    assert bool((out[rows:] == 0).all())
-    check_against_native(lat[None], out[None], w_uv, w_o, rows, ("eager", kind, rows, out_rows))
-    _, v_nat = tail_reference.native_attn_tail(lat, w_uv, w_o, rows, return_v=True)
-    print(
-        f"[attention_output_tail-pdl-stress-eager] weights={kind} rows={rows} out_rows={out_rows} "
-        f"pdl_equals_serial=True v_equals_native={torch.equal(v, v_nat)}"
-    )
-
-
-def _run_delayed_producer(rows, out_rows, delay_us, config):
-    w_uv, w_o, kind = tail_weights()
-    n = LOCAL_HEADS * V_DIM
-    latent = torch.empty((rows, LOCAL_HEADS, C_DIM), dtype=torch.bfloat16, device="cuda")
-    v = torch.empty((rows, n), dtype=torch.bfloat16, device="cuda")
-    srcs = make_latents(N_ITERS, rows, seed=1000)
-    caps = torch.empty((N_ITERS, out_rows, HIDDEN), dtype=torch.bfloat16, device="cuda")
-    stamps = torch.zeros((N_ITERS, 4), dtype=torch.int64, device="cuda")
-
-    def body():
-        for i in range(N_ITERS):
-            # Poison (ordinary, fully serialized kernels).
-            latent.fill_(float("nan"))
-            v.fill_(float("nan"))
-            # Not PDL-launched itself, but lets attention_output_tail's W_UV launch at its entry.
-            launch_producer(srcs[i], latent, stamps[i], delay_us, True, False, False)
-            out = attn_tail(latent, w_uv, w_o, rows, out_rows=out_rows, v_scratch=v, config=config)
-            launch_consumer(out, caps[i], stamps[i, 1:], 0, False)
-
-    body()  # eager pass: JIT compile + a first stressed run
-    torch.cuda.synchronize()
-    expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, config)
-    assert_bitwise(caps, expected, "delayed-producer eager")
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        body()
-    gaps = []
-    for rep in range(N_REPLAYS):
-        srcs.copy_(make_latents(N_ITERS, rows, seed=2000 + 100 * rep))
-        caps.fill_(float("nan"))
-        torch.cuda.synchronize()
-        graph.replay()
-        torch.cuda.synchronize()
-        expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, config)
-        assert_bitwise(caps, expected, f"delayed-producer replay {rep}")
-        st = stamps.cpu()
-        # Producer copy start -> consumer wait return: attention_output_tail's critical path
-        # after the latent becomes available (plus the copy itself).
-        gaps += [int(st[i, 2] - st[i, 0]) for i in range(N_ITERS)]
-    check_against_native(srcs, expected, w_uv, w_o, rows, ("delayed-producer", kind))
-    del graph
-    return kind, gaps
-
-
-@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1), (8, 8)])
-@pytest.mark.parametrize("delay_us", [0, 3, 20, 50])
-@pytest.mark.parametrize("entry", ["pdl_entry", "no_entry"])
-def test_delayed_producer(rows, out_rows, delay_us, entry):
-    config = DEFAULT_CONFIG if entry == "pdl_entry" else TailConfig(pdl_entry=False)
-    kind, gaps = _run_delayed_producer(rows, out_rows, delay_us, config)
-    print(
-        f"[attention_output_tail-pdl-stress-producer] weights={kind} rows={rows} "
-        f"out_rows={out_rows} delay_us={delay_us} entry={entry} "
-        f"iterations={len(gaps)} bitwise_ok=True "
-        f"producer_copy_to_consumer_wait_us_median={_median_us(gaps):.2f}"
-    )
-
-
-def _run_delayed_consumer(rows, out_rows, producer_us, consumer_us):
-    w_uv, w_o, kind = tail_weights()
-    latent = torch.empty((rows, LOCAL_HEADS, C_DIM), dtype=torch.bfloat16, device="cuda")
-    srcs = make_latents(N_ITERS, rows, seed=3000)
-    caps = torch.empty((N_ITERS, out_rows, HIDDEN), dtype=torch.bfloat16, device="cuda")
-    stamps = torch.zeros((N_ITERS, 4), dtype=torch.int64, device="cuda")
-
-    def body():
-        out_ptrs, v_ptrs = [], []
-        for i in range(N_ITERS):
-            # PDL-launched and triggering at entry: the chain may run ahead
-            # into this iteration while the previous consumer still spins.
-            launch_producer(srcs[i], latent, stamps[i], producer_us, True, True, True)
-            v = torch.empty((rows, LOCAL_HEADS * V_DIM), dtype=torch.bfloat16, device="cuda")
-            out = attn_tail(latent, w_uv, w_o, rows, out_rows=out_rows, v_scratch=v)
-            out_ptrs.append(out.data_ptr())
-            v_ptrs.append(v.data_ptr())
-            launch_consumer(out, caps[i], stamps[i, 1:], consumer_us, True)
-            del out, v  # back to the pool: the next iteration reuses them
-        return out_ptrs, v_ptrs
-
-    body()
-    torch.cuda.synchronize()
-    expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, DEFAULT_CONFIG)
-    assert_bitwise(caps, expected, "delayed-consumer eager")
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        out_ptrs, v_ptrs = body()
-    early = []
-    for rep in range(N_REPLAYS):
-        srcs.copy_(make_latents(N_ITERS, rows, seed=4000 + 100 * rep))
-        caps.fill_(float("nan"))
-        torch.cuda.synchronize()
-        graph.replay()
-        torch.cuda.synchronize()
-        expected = eager_expected(srcs, w_uv, w_o, rows, out_rows, DEFAULT_CONFIG)
-        assert_bitwise(caps, expected, f"delayed-consumer replay {rep}")
-        st = stamps.cpu()
-        # Consumer resident before attention_output_tail finished: time spent in its wait.
-        early += [int(st[i, 2] - st[i, 1]) for i in range(N_ITERS)]
-    check_against_native(srcs, expected, w_uv, w_o, rows, ("delayed-consumer", kind))
-    del graph
-    reuse = {
-        "distinct_out_ptrs": len(set(out_ptrs)),
-        "distinct_v_ptrs": len(set(v_ptrs)),
-        "v_overlaps_a_partial": any(
-            any(o <= p < o + out_rows * HIDDEN * 2 for o in out_ptrs) for p in v_ptrs
-        ),
-    }
-    return kind, early, reuse
-
-
-@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1)])
-@pytest.mark.parametrize("producer_us,consumer_us", [(0, 0), (0, 20), (20, 20), (5, 50)])
-def test_delayed_consumer_buffer_reuse(rows, out_rows, producer_us, consumer_us):
-    kind, early, reuse = _run_delayed_consumer(rows, out_rows, producer_us, consumer_us)
-    # The pool hands the same partial buffer to consecutive iterations, so
-    # the stress covers a reused output buffer.
-    assert reuse["distinct_out_ptrs"] < N_ITERS, reuse
-    frac = sum(e > 1000 for e in early) / max(len(early), 1)
-    print(
-        f"[attention_output_tail-pdl-stress-consumer] weights={kind} rows={rows} "
-        f"out_rows={out_rows} producer_us={producer_us} "
-        f"consumer_us={consumer_us} iterations={len(early)} bitwise_ok=True "
-        f"reuse={reuse} consumer_wait_us_median={_median_us(early):.2f} "
-        f"fraction_waited_over_1us={frac:.2f}"
-    )
-
-
-@pytest.mark.skipif(UNDER_SANITIZER, reason="the sanitizer serializes kernel launches")
-def test_pdl_overlap_witness():
-    """The stress is not vacuous: the PDL consumer of attention_output_tail is resident and
-    waiting before attention_output_tail has finished (attention_output_tail triggers its dependents
-    early), and attention_output_tail's PDL entry shortens the producer → consumer path (W_UV
-    launched under the producer)."""
-    _, early, _ = _run_delayed_consumer(6, 6, 0, 20)
-    frac = sum(e > 1000 for e in early) / len(early)
-    _, gaps_entry = _run_delayed_producer(6, 6, 20, DEFAULT_CONFIG)
-    _, gaps_noentry = _run_delayed_producer(6, 6, 20, TailConfig(pdl_entry=False))
-    print(
-        f"[attention_output_tail-pdl-stress-witness] consumer_wait_us_median={_median_us(early):.2f} "
-        f"fraction_waited_over_1us={frac:.2f} "
-        f"producer_to_consumer_us_median pdl_entry="
-        f"{_median_us(gaps_entry):.2f} no_entry={_median_us(gaps_noentry):.2f}"
-    )
-    assert frac >= 0.5, frac
-    assert _median_us(gaps_entry) < _median_us(gaps_noentry)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py
deleted file mode 100644
index 20f230b6e21b..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_attn_tail_sanitizer_probe.py
+++ /dev/null
@@ -1,115 +0,0 @@
-# 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.
-"""Attribution probes for compute-sanitizer initcheck reports.
-
-Run under ``compute-sanitizer --tool initcheck`` with
-``PYTORCH_NO_CUDA_MEMORY_CACHING=1``.
-Each probe produces one fully written tensor and then reads it once with an
-ordinary PyTorch conversion kernel, so pytest's per-test output attributes
-every report to one producer: native producers (H2D copy, elementwise,
-cuBLAS GEMMs, the native oracle tail) and the default attention_output_tail. A report under
-a native probe concerns how that producer's writes are tracked; a report
-under the attention_output_tail probe would mean attention_output_tail left part of its partial
-unwritten.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-
-def _read(y):
-    return float(y.float().abs().sum())
-
-
-def test_probe_h2d_copy_then_read():
-    x = torch.randn(2048, 512, dtype=torch.bfloat16).cuda()
-    torch.cuda.synchronize()
-    assert _read(x) > 0
-
-
-def test_probe_elementwise_then_read():
-    x = torch.randn(2048, 512, device="cuda").bfloat16()
-    torch.cuda.synchronize()
-    assert _read(x) > 0
-
-
-def test_probe_cublas_linear_then_read():
-    x = torch.randn(2048, 6144, device="cuda").bfloat16()
-    w = torch.randn(576, 6144, device="cuda").bfloat16()
-    y = torch.nn.functional.linear(x, w)
-    torch.cuda.synchronize()
-    assert _read(y) > 0
-
-
-def test_probe_cublas_linear_small_m_then_read():
-    x = torch.randn(6, 2048, device="cuda").bfloat16()
-    w = torch.randn(6144, 2048, device="cuda").bfloat16()
-    y = torch.nn.functional.linear(x, w)
-    torch.cuda.synchronize()
-    assert _read(y) > 0
-
-
-def test_probe_cublas_bmm_then_read():
-    a = torch.randn(8, 6, 512, device="cuda").bfloat16()
-    b = torch.randn(8, 512, 256, device="cuda").bfloat16()
-    y = torch.bmm(a, b)
-    torch.cuda.synchronize()
-    assert _read(y) > 0
-
-
-def _tail_inputs(rows):
-    gen = torch.Generator(device="cuda").manual_seed(rows)
-    w_uv = (torch.randn(8, 256, 512, generator=gen, device="cuda") * 0.02).bfloat16()
-    w_o = (torch.randn(6144, 2048, generator=gen, device="cuda") * 0.02).bfloat16()
-    lat = torch.randn(rows, 8, 512, generator=gen, device="cuda").bfloat16()
-    return lat, w_uv, w_o
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_probe_native_tail_then_read(rows):
-    """The native oracle tail (bmm into the o_proj input buffer, then
-    F.linear), read once by a conversion kernel."""
-    import tail_reference
-
-    lat, w_uv, w_o = _tail_inputs(rows)
-    native = tail_reference.native_attn_tail(lat, w_uv, w_o, rows)
-    torch.cuda.synchronize()
-    assert _read(native) > 0
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="attention_output_tail kernels target SM100/SM103",
-)
-@pytest.mark.parametrize("rows,out_rows", [(1, 1), (6, 6), (6, 8), (8, 8)])
-def test_probe_attn_tail_then_read(rows, out_rows):
-    """The default attention_output_tail (PDL on), its partial read once by a conversion
-    kernel: every element of the partial, padded rows included, must be
-    written by attention_output_tail."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E501
-        attn_tail,
-    )
-
-    lat, w_uv, w_o = _tail_inputs(rows)
-    out = attn_tail(lat, w_uv, w_o, rows, out_rows=out_rows)
-    torch.cuda.synchronize()
-    assert _read(out) > 0
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py
deleted file mode 100644
index 86f3015b779a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_into_hop_producer.py
+++ /dev/null
@@ -1,358 +0,0 @@
-# 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.
-"""attention_output_into_hop producer (single GPU): the publishing attention_output_tail o_proj and the tail protocol.
-
-A plain device buffer stands in for the post_attention_hop workspace (multicast view and
-local copy are the same memory on one GPU), filled with the Lamport sentinel,
-with a protocol state selecting the current buffer. On the real DSA ``MLA``
-of TP8 rank 0 with real layer-3 weights and synthetic latents:
-
-* ``request_into_hop`` → ``arm`` → ``defer_w_uv`` → ``project`` publishes exactly
-  slot ``[row][rank]`` of the current buffer for every output row (live rows
-  and the zero rows of padding the post_attention_hop also sums), as 32-bit words equal
-  to the BF16 partial with ``-0.0`` mapped to ``+0.0``; every other word
-  stays the sentinel; ``take_into_hop_published`` reports it once;
-* the ordinary partial is bit-identical to the attention_output_tail without publishing;
-* a request is consumed by the next ``arm`` even when the call declines;
-  non-attention_output_tail engines never publish; a request whose hop row count differs from
-  the output rows publishes nothing (the layer then runs the post_attention_hop);
-* CUDA-graph capture and replay publish into the same slots, also with
-  padded rows;
-* both publish widths (32-bit words on 152 CTAs or on request, 16-byte
-  stores on 128 CTAs) write the same words;
-* ``make_attn_tail`` engages the attention_output_into_hop schedules only next to the
-  post_attention_hop (``INTO_HOP_SPARSE_DECODE_SCHEDULE`` behind ``sparse_mla_decode``,
-  ``INTO_HOP_SCHEDULE`` behind the native FMHA); the plan turns attention_output_into_hop
-  off with either path it extends.
-"""
-
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
-sys.path.insert(0, os.path.dirname(_HERE))
-import plan_testlib  # noqa: E402
-import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
-    dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
-    DEFAULT_CONFIG,
-    INTO_HOP_SCHEDULE,
-    INTO_HOP_SPARSE_DECODE_SCHEDULE,
-    TAIL_PREFETCH_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
-    TailConfig,
-    attn_tail,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
-    AttnTail,
-    make_attn_tail,
-)
-
-pytestmark = tail_adapter.pytestmark
-
-H = 6144
-RANKS = 8
-MAX_ROWS = 8
-BUF_WORDS = MAX_ROWS * RANKS * H * 2 // 4
-SENTINEL = -(2**31)  # 0x80000000 as int32
-
-
-def _transport(cur=1, rank=3):
-    mc = torch.full((3 * BUF_WORDS,), SENTINEL, dtype=torch.int32, device="cuda")
-    flags = torch.tensor([cur, 0, 0, 0], dtype=torch.int32, device="cuda")
-    return SimpleNamespace(mc_words=mc, flags=flags, rank=rank, cur=cur)
-
-
-def _sanitized_words(partial):
-    bits = partial.contiguous().view(torch.int16).clone()
-    bits[bits == -32768] = 0  # bf16 -0.0 (0x8000) -> +0.0
-    return bits.view(torch.int32)
-
-
-def _check_slots(tr, partial, rows, live=None):
-    """Slots ``[:rows, rank]`` hold ``partial``'s words; rows past ``live``
-    are zero words; every other word of every buffer is the sentinel."""
-    buf = tr.mc_words[tr.cur * BUF_WORDS : (tr.cur + 1) * BUF_WORDS].view(MAX_ROWS, RANKS, H // 2)
-    want = _sanitized_words(partial[:rows])
-    assert torch.equal(buf[:rows, tr.rank], want)
-    assert not bool((buf[:rows, tr.rank] == SENTINEL).any())
-    if live is not None:
-        assert bool((buf[live:rows, tr.rank] == 0).all())
-    others = torch.ones(MAX_ROWS, RANKS, dtype=torch.bool, device="cuda")
-    others[:rows, tr.rank] = False
-    assert bool((buf[others] == SENTINEL).all())
-    for b in range(3):
-        if b != tr.cur:
-            assert bool((tr.mc_words[b * BUF_WORDS : (b + 1) * BUF_WORDS] == SENTINEL).all())
-
-
-def _step(tail, mla, latent, tr, out_rows=None, hop_rows=None):
-    """One armed call; the attention output has ``out_rows`` rows (live rows
-    first, as in ``MLA.forward``) and the request names ``hop_rows``."""
-    rows = latent.shape[0]
-    out_rows = rows if out_rows is None else out_rows
-    hop_rows = out_rows if hop_rows is None else hop_rows
-    out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    tail.request_into_hop(tr, hop_rows)
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point="attn")
-    assert tail.defer_w_uv(mla, meta, latent, out[:rows])
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    return tail.project(mla, out, AllReduceParams(enable_allreduce=False))
-
-
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-@pytest.mark.parametrize("grid,vector", [(0, True), (128, True), (128, False)])
-def test_publish_slots_and_partial(rows, grid, vector):
-    mla = tail_adapter._mla()
-    cfg = TailConfig(oproj_grid=grid, into_hop_vector_publish=vector)
-    tail = AttnTail("layer3", cfg)
-    tr = _transport()
-    latent = tail_adapter._latent(rows, seed=rows)
-    partial = _step(tail, mla, latent, tr)
-    torch.cuda.synchronize()
-    assert tail.take_into_hop_published() and not tail.take_into_hop_published()
-    ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows, config=cfg)
-    assert torch.equal(partial, ref)
-    _check_slots(tr, partial, rows)
-
-
-@pytest.mark.parametrize("live,out_rows", [(1, 4), (2, 3), (6, 8), (7, 8), (1, 8)])
-@pytest.mark.parametrize("grid,vector", [(0, True), (128, True), (128, False)])
-def test_publish_padded_rows(live, out_rows, grid, vector):
-    """Padded output rows (``hidden_states`` rows > ``num_tokens``): every
-    output row is published, the padding as zero words, so the hop, which
-    polls the residual's rows, sees all of them (and sums zeros, as the post_attention_hop
-    hop does from the zero rows of the ordinary partial)."""
-    mla = tail_adapter._mla()
-    cfg = TailConfig(oproj_grid=grid, into_hop_vector_publish=vector)
-    tail = AttnTail("layer3", cfg)
-    tr = _transport(cur=0, rank=6)
-    latent = tail_adapter._latent(live, seed=30 + live)
-    partial = _step(tail, mla, latent, tr, out_rows=out_rows)
-    torch.cuda.synchronize()
-    assert tail.take_into_hop_published()
-    assert tail.into_hop_published_calls == 1 and tail.into_hop_declined_calls == 0
-    ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, live, out_rows=out_rows, config=cfg)
-    assert partial.shape[0] == out_rows
-    assert torch.equal(partial, ref)
-    assert bool((partial[live:] == 0).all())
-    _check_slots(tr, partial, out_rows, live=live)
-
-
-@pytest.mark.parametrize("live,out_rows,hop_rows", [(6, 8, 6), (6, 6, 8), (4, 4, 3)])
-def test_hop_row_mismatch_publishes_nothing(live, out_rows, hop_rows):
-    """A request whose hop row count differs from the output rows is not
-    honoured: nothing is published (the hop resets exactly the rows it
-    polls, so extra or missing rows would corrupt the transport), the
-    ordinary partial is unchanged and the layer runs the post_attention_hop."""
-    mla = tail_adapter._mla()
-    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
-    tr = _transport()
-    latent = tail_adapter._latent(live, seed=70 + live)
-    partial = _step(tail, mla, latent, tr, out_rows=out_rows, hop_rows=hop_rows)
-    torch.cuda.synchronize()
-    assert not tail.take_into_hop_published()
-    assert tail.into_hop_published_calls == 0 and tail.into_hop_declined_calls == 1
-    assert tail.engaged_calls == 1
-    ref = attn_tail(
-        latent, mla.v_b_proj, mla.o_proj.weight, live, out_rows=out_rows, config=INTO_HOP_SCHEDULE
-    )
-    assert torch.equal(partial, ref)
-    assert bool((tr.mc_words == SENTINEL).all())
-
-
-@pytest.mark.parametrize("live,out_rows", [(6, 6), (6, 8), (1, 4), (8, 8)])
-@pytest.mark.parametrize("grid,vector", [(0, True), (128, True), (128, False)])
-def test_publisher_synthetic_inputs(live, out_rows, grid, vector):
-    """The publishing o_proj alone on fully initialized synthetic inputs
-    (no MLA module, so a compute-sanitizer run sees only this kernel and
-    its reference): ordinary partial bit-identical to ``stream_oproj``,
-    every output row published."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_oproj import (  # noqa: E501
-        stream_oproj,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.into_hop import (  # noqa: E501
-        stream_oproj_into_hop,
-    )
-
-    gen = torch.Generator(device="cuda").manual_seed(17 * live + out_rows)
-    v = (torch.randn(live, 2048, generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
-    w_o = (torch.randn(H, 2048, generator=gen, device="cuda") * 0.02).to(torch.bfloat16)
-    tr = _transport(cur=2, rank=1)
-    out = torch.empty((out_rows, H), dtype=torch.bfloat16, device="cuda")
-    stream_oproj_into_hop(
-        v,
-        w_o,
-        out,
-        True,
-        tr.mc_words,
-        tr.flags,
-        tr.rank,
-        grid=grid or None,
-        trigger_dependents=True,
-        vector_publish=vector,
-    )
-    ref = torch.empty((out_rows, H), dtype=torch.bfloat16, device="cuda")
-    stream_oproj(v, w_o, ref, use_pdl=True, grid=grid or None)
-    torch.cuda.synchronize()
-    assert torch.equal(out, ref)
-    _check_slots(tr, out, out_rows, live=live)
-
-
-def test_publisher_rejects_rows_outside_transport():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.into_hop import (  # noqa: E501
-        stream_oproj_into_hop,
-    )
-
-    mla = tail_adapter._mla()
-    tr = _transport()
-    v = torch.zeros((6, 2048), dtype=torch.bfloat16, device="cuda")
-    for out_rows in (9, 5):
-        out = torch.empty((out_rows, H), dtype=torch.bfloat16, device="cuda")
-        with pytest.raises(ValueError):
-            stream_oproj_into_hop(v, mla.o_proj.weight, out, True, tr.mc_words, tr.flags, tr.rank)
-
-
-def test_request_consumed_and_engines():
-    mla = tail_adapter._mla()
-    tr = _transport()
-    # Declined call: the request is consumed, nothing is published.
-    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
-    tail.request_into_hop(tr, 6)
-    assert not tail.arm(mla, tail_adapter._meta(6, num_contexts=1))
-    assert tail._into_hop_request is None and tail._into_hop is None
-    assert not tail.take_into_hop_published()
-    # Without a request the attention_output_tail does not publish.
-    latent = tail_adapter._latent(6, seed=9)
-    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(6)
-    assert tail.arm(mla, meta)
-    assert tail.defer_w_uv(mla, meta, latent, out)
-    tail.project(mla, out, None)
-    torch.cuda.synchronize()
-    assert not tail.take_into_hop_published()
-    assert bool((tr.mc_words == SENTINEL).all())
-    # A non-attention_output_tail engine never publishes.
-    native = AttnTail("layer3", TailConfig(oproj_engine="triton"))
-    native.request_into_hop(tr, 6)
-    assert native.arm(mla, meta)
-    assert native._into_hop is None
-
-
-def test_request_dropped_after_a_compile_declined_call():
-    """Under torch.compile ``arm`` returns before it reads the request; the
-    layer's ``take_into_hop_published`` drops it, so a later call without a request
-    never honours it (no publication into the transport)."""
-    from tensorrt_llm._torch.utils import torch_compiling
-
-    mla = tail_adapter._mla()
-    tr = _transport()
-    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
-    tail.request_into_hop(tr, 6)
-    with torch_compiling(True):
-        assert not tail.arm(mla, tail_adapter._meta(6))
-    assert tail._into_hop_request is not None  # arm touches no state under compile
-    assert not tail.take_into_hop_published()
-    assert tail._into_hop_request is None
-    latent = tail_adapter._latent(6, seed=9)
-    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(6)
-    assert tail.arm(mla, meta)
-    assert tail._into_hop is None
-    assert tail.defer_w_uv(mla, meta, latent, out)
-    tail.project(mla, out, None)
-    torch.cuda.synchronize()
-    assert not tail.take_into_hop_published()
-    assert bool((tr.mc_words == SENTINEL).all())
-
-
-@pytest.mark.parametrize("out_rows", [6, 8])
-def test_graph_capture_and_replay_publish(out_rows):
-    mla = tail_adapter._mla()
-    tail = AttnTail("layer3", INTO_HOP_SCHEDULE)
-    rows = 6
-    tr = _transport(cur=2, rank=5)
-    lat = tail_adapter._latent(rows, seed=21).contiguous()
-    _step(tail, mla, lat, tr, out_rows=out_rows)
-    tail.take_into_hop_published()
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = _step(tail, mla, lat, tr, out_rows=out_rows)
-    assert tail.take_into_hop_published()
-    for seed in range(4):
-        tr.mc_words.fill_(SENTINEL)
-        lat.copy_(tail_adapter._latent(rows, seed=800 + seed))
-        graph.replay()
-        torch.cuda.synchronize()
-        ref = attn_tail(
-            lat, mla.v_b_proj, mla.o_proj.weight, rows, out_rows=out_rows, config=INTO_HOP_SCHEDULE
-        )
-        assert torch.equal(captured, ref)
-        _check_slots(tr, captured, out_rows, live=rows)
-
-
-def test_make_attn_tail_selects_into_hop(monkeypatch):
-    mla = tail_adapter._build_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
-    cases = [
-        (
-            "attention_output_tail,attention_output_into_hop,post_attention_hop,sparse_mla_decode",
-            INTO_HOP_SPARSE_DECODE_SCHEDULE,
-        ),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,attention_output_into_hop,post_attention_hop,sparse_mla_decode",
-            INTO_HOP_SPARSE_DECODE_SCHEDULE,
-        ),
-        ("attention_output_tail,attention_output_into_hop,post_attention_hop", INTO_HOP_SCHEDULE),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,attention_output_into_hop,post_attention_hop",
-            INTO_HOP_SCHEDULE,
-        ),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,post_attention_hop,sparse_mla_decode",
-            TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-        ),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,sparse_mla_decode",
-            TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
-        ),
-        # attention_output_tail_prefetch has no schedule next to the post_attention_hop
-        # without sparse_mla_decode; the tail keeps its default schedule.
-        ("attention_output_tail,attention_output_tail_prefetch,post_attention_hop", DEFAULT_CONFIG),
-        ("attention_output_tail,attention_output_tail_prefetch", TAIL_PREFETCH_SCHEDULE),
-        ("attention_output_tail", DEFAULT_CONFIG),
-    ]
-    for features, want in cases:
-        plan_testlib.select(dispatch, features, monkeypatch)
-        tail = make_attn_tail(mla)
-        assert tail is not None and tail.config == want, features
-
-
-def test_attention_output_into_hop_is_off_without_the_paths_it_extends():
-    for base in ("attention_output_tail", "post_attention_hop"):
-        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
-        assert not plan.enabled("attention_output_into_hop"), base
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py
deleted file mode 100644
index c07963d4f391..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch.py
+++ /dev/null
@@ -1,531 +0,0 @@
-# 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.
-"""attention_output_tail_prefetch: off-path L2 prefetch of the attention_output_tail weights (protocol and kernel).
-
-Uses the real DSA ``MLA`` module of TP8 rank 0 with real layer-3
-``v_b_proj``/``o_proj`` weights and synthetic latents (the attention_output_tail adapter test's
-builders); the real ``MLA.forward`` path is in ``test_tail_prefetch_mla_forward.py``.
-
-* the prefetch kernel runs on every config the benchmarks use and changes no
-  tensor;
-* attention_output_tail_prefetch's default tail (native W_UV + native o_proj on prefetched weights)
-  is bit-identical to the native tail, and with the attention_output_tail kernels on prefetched
-  weights it is bit-identical to the attention_output_tail: the prefetch has no effect on
-  results;
-* fork/join protocol: ``before_attention`` forks only on an armed call and
-  only for layers the config selects; ``project`` always joins (engaged or
-  declined), at every configured join point; a leftover fork raises in the
-  next ``arm``;
-* CUDA-graph capture of arm → before_attention → defer → project replays
-  correctly with the auxiliary stream rejoined;
-* ``make_attn_tail`` picks the attention_output_tail_prefetch schedule of the running
-  neighbours exactly when ``attention_output_tail_prefetch`` runs;
-* the adopted attention (``GlmMLA``) carries the attention_output_tail_prefetch flag and its tail
-  uses ``TAIL_PREFETCH_SCHEDULE``; the layer adoption hands the MoE gate weight to the tail.
-"""
-
-import dataclasses
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
-sys.path.insert(0, os.path.dirname(_HERE))
-import plan_testlib  # noqa: E402
-import tail_reference  # noqa: E402
-import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
-    dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail import (  # noqa: E402, E501
-    prefetch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
-    DEFAULT_CONFIG,
-    TAIL_PREFETCH_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
-    TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE,
-    attn_tail,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
-    TailConfig as TailConfig_,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
-    AttnTail,
-    make_attn_tail,
-)
-
-pytestmark = tail_adapter.pytestmark
-
-_R = dataclasses.replace
-STAGED = TAIL_PREFETCH_STAGED_NATIVE_SCHEDULE
-# attention_output_tail_prefetch's prefetch protocol with the attention_output_tail kernels (the
-# default uses the native ones).
-TAIL_PREFETCH_KERNEL_CONFIGS = _R(STAGED, oproj_engine="cute_stream")
-
-
-def _no_ar():
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    return AllReduceParams(enable_allreduce=False)
-
-
-@pytest.mark.parametrize(
-    "cfg",
-    [
-        prefetch.PrefetchConfig(),
-        prefetch.PrefetchConfig(ctas=152),
-        prefetch.PrefetchConfig(chunk=16384, ctas=26),
-        prefetch.PrefetchConfig(gap_ns=500),
-    ],
-)
-def test_prefetch_kernel_changes_nothing(cfg):
-    mla = tail_adapter._mla()
-    weights = [mla.v_b_proj, mla.o_proj.weight]
-    assert prefetch.check_prefetch_tensors(weights, cfg.chunk) == ""
-    before = [w.clone() for w in weights]
-    probe = torch.randn(4096, device="cuda")
-    probe_before = probe.clone()
-    prefetch.launch_prefetch(weights, cfg)
-    torch.cuda.synchronize()
-    for w, b in zip(weights, before):
-        assert torch.equal(w, b)
-    assert torch.equal(probe, probe_before)
-
-
-def test_prefetch_rejects_bad_inputs():
-    a = torch.zeros(4096, dtype=torch.bfloat16, device="cuda")
-    assert prefetch.check_prefetch_tensors([a, a, a, a], 4096) != ""
-    assert prefetch.check_prefetch_tensors([a, a], 4000) != ""
-    odd = torch.zeros(1000, dtype=torch.bfloat16, device="cuda")
-    assert prefetch.check_prefetch_tensors([a, odd], 4096) != ""
-
-
-def _drive(tail, mla, latent, out_rows=None, decline=False):
-    rows = latent.shape[0]
-    out_rows = out_rows or rows
-    attn_out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point=tail.config.prefetch_fork)
-    if decline:
-        big = torch.zeros((9, 8, 512), dtype=torch.bfloat16, device="cuda")
-        assert not tail.defer_w_uv(mla, meta, big, attn_out[:rows])
-    else:
-        assert tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
-    return tail.project(mla, attn_out, _no_ar())
-
-
-@pytest.mark.parametrize("join", ["before_wuv", "before_oproj", "after_oproj"])
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_tail_prefetch_tail_bitwise_equals_attn_tail(join, rows):
-    mla = tail_adapter._mla()
-    latent = tail_adapter._latent(rows, seed=rows)
-    ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows, config=DEFAULT_CONFIG)
-    tail = AttnTail("layer3", _R(TAIL_PREFETCH_KERNEL_CONFIGS, prefetch_join=join))
-    got = _drive(tail, mla, latent)
-    torch.cuda.synchronize()
-    assert torch.equal(got, ref)
-    assert tail.prefetch_forks == 1
-    assert not tail._prefetch.forked and not tail._armed
-    assert tail._pending is None
-
-
-def test_fork_only_on_armed_calls_and_join_on_decline():
-    mla = tail_adapter._mla()
-    tail = AttnTail("layer3", TAIL_PREFETCH_KERNEL_CONFIGS)
-    # Not armed: no fork.
-    tail.before_attention(mla, point="attn")
-    assert tail.prefetch_forks == 0
-    assert not tail.arm(mla, tail_adapter._meta(6, num_contexts=1))
-    tail.before_attention(mla, point="attn")
-    assert tail.prefetch_forks == 0
-    # Armed, but called at the other fork point: no fork.
-    assert tail.arm(mla, tail_adapter._meta(6))
-    tail.before_attention(mla, point="start")
-    assert tail.prefetch_forks == 0
-    tail._armed = False
-    # Armed, forked, then declined at the W_UV step: project joins.
-    assert _drive(tail, mla, tail_adapter._latent(6, seed=3), decline=True) is None
-    assert tail.prefetch_forks == 1 and not tail._prefetch.forked
-    # An attention_output_tail-only tail never forks.
-    served_tail = AttnTail("layer3", DEFAULT_CONFIG)
-    out = _drive(served_tail, mla, tail_adapter._latent(6, seed=4))
-    assert out is not None and served_tail.prefetch_forks == 0
-    # A leftover fork is an error at the next arm.
-    assert tail.arm(mla, tail_adapter._meta(6))
-    tail.before_attention(mla, point="attn")
-    tail._armed = False
-    with pytest.raises(RuntimeError, match="not consumed"):
-        tail.arm(mla, tail_adapter._meta(6))
-    tail._prefetch.join()
-
-
-def test_layer_selection():
-    mla = tail_adapter._mla()
-    has_indexer = getattr(mla.mqa, "indexer", None) is not None
-    for layers in ("all", "shared", "full"):
-        tail = AttnTail("layer3", _R(TAIL_PREFETCH_KERNEL_CONFIGS, prefetch_layers=layers))
-        _drive(tail, mla, tail_adapter._latent(6, seed=5))
-        want = layers == "all" or (layers == "full") == has_indexer
-        assert tail.prefetch_forks == int(want), (layers, has_indexer)
-
-
-def _native_step(tail, mla, latent, out_rows=None):
-    """What the patched op and MLA.forward do with attention_output_tail_prefetch's native engine."""
-    from tensorrt_llm._torch.attention.mla import MLA
-
-    rows = latent.shape[0]
-    out_rows = out_rows or rows
-    attn_out = torch.zeros((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point="attn")
-    # The op's native W_UV (the tail declines the W_UV step).
-    assert not tail.defer_w_uv(mla, meta, latent, attn_out[:rows])
-    v_view = attn_out[:rows].view(rows, mla.num_heads_tp, mla.v_head_dim)
-    MLA._bmm_bf16_out(
-        mla,
-        latent.transpose(0, 1),
-        mla.v_b_proj,
-        mla.v_b_proj.transpose(1, 2),
-        v_view.transpose(0, 1),
-    )
-    return tail.project(mla, attn_out, _no_ar())
-
-
-@pytest.mark.parametrize("join", ["before_oproj", "after_oproj"])
-@pytest.mark.parametrize("rows,out_rows", [(1, 1), (6, 6), (6, 8), (8, 8)])
-def test_native_engine_bitwise_equals_native(join, rows, out_rows):
-    mla = tail_adapter._mla()
-    latent = tail_adapter._latent(rows, seed=10 + rows)
-    native = tail_adapter._native(mla, latent, out_rows)
-    tail = AttnTail("layer3", _R(STAGED, prefetch_join=join))
-    got = _native_step(tail, mla, latent, out_rows)
-    torch.cuda.synchronize()
-    assert torch.equal(got[:rows], native[:rows])
-    assert tail.prefetch_forks == 1 and tail.engaged_calls == 1
-    assert not tail._prefetch.forked and not tail._armed
-
-
-def test_native_engine_without_fork_takes_native_path():
-    mla = tail_adapter._mla()
-    # Layer selection excludes this layer: no fork, so the caller's native
-    # projection runs (project returns None).
-    has_indexer = getattr(mla.mqa, "indexer", None) is not None
-    other = "shared" if has_indexer else "full"
-    tail = AttnTail("layer3", _R(STAGED, prefetch_layers=other))
-    rows = 6
-    attn_out = torch.zeros((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point="attn")
-    assert tail.prefetch_forks == 0
-    assert not tail.defer_w_uv(mla, meta, tail_adapter._latent(rows, seed=6), attn_out)
-    assert tail.project(mla, attn_out, _no_ar()) is None
-
-
-def test_native_engine_graph_capture_and_replay():
-    mla = tail_adapter._mla()
-    tail = AttnTail("layer3", STAGED)
-    rows = 6
-    lat = tail_adapter._latent(rows, seed=17).contiguous()
-    _native_step(tail, mla, lat)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = _native_step(tail, mla, lat)
-    assert not tail._prefetch.forked
-    for seed in range(6):
-        lat.copy_(tail_adapter._latent(rows, seed=300 + seed))
-        graph.replay()
-        torch.cuda.synchronize()
-        assert torch.equal(captured, tail_adapter._native(mla, lat, rows))
-
-
-def test_graph_capture_and_replay():
-    mla = tail_adapter._mla()
-    tail = AttnTail("layer3", TAIL_PREFETCH_KERNEL_CONFIGS)
-    rows = 6
-    meta = tail_adapter._meta(rows)
-    lat = tail_adapter._latent(rows, seed=7).contiguous()
-    out = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
-
-    def step():
-        assert tail.arm(mla, meta)
-        tail.before_attention(mla, point="attn")
-        assert tail.defer_w_uv(mla, meta, lat, out)
-        return tail.project(mla, out, _no_ar())
-
-    step()
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = step()
-    assert not tail._prefetch.forked
-    for seed in range(6):
-        lat.copy_(tail_adapter._latent(rows, seed=200 + seed))
-        graph.replay()
-        torch.cuda.synchronize()
-        ref = attn_tail(lat, mla.v_b_proj, mla.o_proj.weight, rows, config=DEFAULT_CONFIG)
-        assert torch.equal(captured, ref)
-
-
-def test_make_attn_tail_selects_config(monkeypatch):
-    mla = tail_adapter._build_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
-    cases = [
-        ("attention_output_tail", DEFAULT_CONFIG),
-        ("attention_output_tail,attention_output_tail_prefetch", TAIL_PREFETCH_SCHEDULE),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,sparse_mla_decode",
-            TAIL_PREFETCH_SPARSE_DECODE_SCHEDULE,
-        ),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,sparse_mla_decode,post_attention_hop",
-            TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-        ),
-        (
-            "attention_output_tail,attention_output_tail_prefetch,post_attention_hop",
-            DEFAULT_CONFIG,
-        ),
-    ]
-    for features, want in cases:
-        plan_testlib.select(dispatch, features, monkeypatch)
-        tail = make_attn_tail(mla)
-        assert tail is not None and tail.config == want, features
-
-
-@pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED") != "1",
-    reason="opt-in",
-)
-def test_adoption_selects_tail_prefetch(monkeypatch):
-    import inspect
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-        attention_forward,
-    )
-
-    assert "before_attention" in inspect.getsource(attention_forward.forward_absorption_generation)
-    plan_testlib.select(
-        dispatch, "attention_output_tail,attention_output_tail_prefetch", monkeypatch
-    )
-    mla = tail_adapter._build_adopted_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
-    assert mla._tail_prefetch
-    assert mla._attn_tail.config == TAIL_PREFETCH_SCHEDULE
-    plan_testlib.select(dispatch, "attention_output_tail", monkeypatch)
-    assert not tail_adapter._build_adopted_mla(
-        layer_idx=tail_adapter.LAYER, num_hidden_layers=78
-    )._tail_prefetch
-    draft = tail_adapter._build_adopted_mla(layer_idx=78, num_hidden_layers=78)
-    assert draft._attn_tail is None
-
-
-@pytest.mark.parametrize("grid,slots,early", [(96, 3, 3), (96, 3, 0), (128, 3, 3), (64, 3, 2)])
-@pytest.mark.parametrize("rows,out_rows", [(1, 1), (6, 6), (6, 8), (8, 8), (8, 16)])
-def test_ring_oproj_within_budget(grid, slots, early, rows, out_rows):
-    """The small-footprint o_proj against the FP64 oracle of the same ``v``
-    (frozen attention_output_tail budget relative to the native GEMM), padded rows zero."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_oproj_ring import (  # noqa: E501
-        ring_oproj,
-    )
-
-    mla = tail_adapter._mla()
-    latent = tail_adapter._latent(rows, seed=40 + rows)
-    native, v = tail_reference.native_attn_tail(
-        latent, mla.v_b_proj, mla.o_proj.weight, rows, out_rows=out_rows, return_v=True
-    )
-    out = torch.full(
-        (out_rows, mla.o_proj.weight.shape[0]), 7.0, dtype=torch.bfloat16, device="cuda"
-    )
-    ring_oproj(v, mla.o_proj.weight, out, use_pdl=False, grid=grid, slots=slots, early=early)
-    torch.cuda.synchronize()
-    ref = v[:rows].double() @ mla.o_proj.weight.double().t()
-    e_ring = tail_reference.error_stats(out[:rows], ref)
-    e_nat = tail_reference.error_stats(native[:rows], ref)
-    assert e_ring["rel_rms"] <= 1.25 * e_nat["rel_rms"] + 1e-6, (e_ring, e_nat)
-    assert e_ring["max_abs"] <= (2 * e_nat["max_abs"] + 2**-8 * e_ring["ref_max"])
-    assert bool((out[rows:] == 0).all())
-
-
-def test_ring_engine_graph_replay_matches_native():
-    mla = tail_adapter._mla()
-    cfg = _R(STAGED, native_oproj="ring")
-    tail = AttnTail("layer3", cfg)
-    rows = 6
-    lat = tail_adapter._latent(rows, seed=51).contiguous()
-    _native_step(tail, mla, lat)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = _native_step(tail, mla, lat)
-    for seed in range(6):
-        lat.copy_(tail_adapter._latent(rows, seed=500 + seed))
-        graph.replay()
-        torch.cuda.synchronize()
-        native = tail_adapter._native(mla, lat, rows)
-        torch.testing.assert_close(
-            captured.clone(),
-            native,
-            rtol=2**-6,
-            atol=2**-6 * native.float().pow(2).mean().sqrt().item(),
-        )
-    assert tail.prefetch_forks == 2 and not tail._prefetch.forked
-
-
-def test_fork_ahead_protocol():
-    mla = tail_adapter._mla()
-    prev = _R(STAGED, prefetch_fork="prev")
-    # Other fork points ignore fork_ahead.
-    tail = AttnTail("layer3", STAGED)
-    tail.fork_ahead(mla)
-    assert tail.prefetch_forks == 0
-    # Forked ahead, then an armed call: before_attention does not fork
-    # again, project runs the native o_proj and joins; bit-identical.
-    tail = AttnTail("layer3", prev)
-    rows = 6
-    latent = tail_adapter._latent(rows, seed=60)
-    tail.fork_ahead(mla)
-    assert tail.prefetch_forks == 1 and tail._prefetch.forked
-    meta = tail_adapter._meta(rows)
-    attn_out = torch.zeros((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    assert tail.arm(mla, meta)
-    tail.before_attention(mla, point="attn")
-    tail.before_attention(mla, point="start")
-    assert tail.prefetch_forks == 1
-    from tensorrt_llm._torch.attention.mla import MLA
-
-    assert not tail.defer_w_uv(mla, meta, latent, attn_out)
-    v_view = attn_out.view(rows, mla.num_heads_tp, mla.v_head_dim)
-    MLA._bmm_bf16_out(
-        mla,
-        latent.transpose(0, 1),
-        mla.v_b_proj,
-        mla.v_b_proj.transpose(1, 2),
-        v_view.transpose(0, 1),
-    )
-    got = tail.project(mla, attn_out, _no_ar())
-    torch.cuda.synchronize()
-    assert torch.equal(got, tail_adapter._native(mla, latent, rows))
-    assert not tail._prefetch.forked and not tail._ahead
-    # Forked ahead, then a call that declines: arm joins at once.
-    tail.fork_ahead(mla)
-    assert not tail.arm(mla, tail_adapter._meta(rows, num_contexts=1))
-    assert not tail._prefetch.forked and not tail._ahead
-    # Twice ahead without a call in between is a no-op the second time.
-    tail.fork_ahead(mla)
-    tail.fork_ahead(mla)
-    assert tail.prefetch_forks == 3
-    assert not tail.arm(mla, tail_adapter._meta(9))
-    assert not tail._prefetch.forked
-
-
-def test_fork_ahead_graph_capture():
-    mla = tail_adapter._mla()
-    tail = AttnTail("layer3", _R(STAGED, prefetch_fork="prev"))
-    rows = 6
-    lat = tail_adapter._latent(rows, seed=61).contiguous()
-
-    def step():
-        tail.fork_ahead(mla)
-        return _native_step(tail, mla, lat)
-
-    step()
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = step()
-    assert not tail._prefetch.forked
-    for seed in range(4):
-        lat.copy_(tail_adapter._latent(rows, seed=700 + seed))
-        graph.replay()
-        torch.cuda.synchronize()
-        assert torch.equal(captured, tail_adapter._native(mla, lat, rows))
-
-
-def test_successor_prefetch_protocol_and_bitwise():
-    mla = tail_adapter._mla()
-    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
-    before = router.clone()
-    for staged in (False, True):
-        cfg = TailConfig_(successor_prefetch=True, prefetch_l2=staged, oproj_engine="cute_stream")
-        tail = AttnTail("layer3", cfg)
-        rows = 6
-        latent = tail_adapter._latent(rows, seed=80)
-        # Without successor weights and without tail staging: no fork.
-        if not staged:
-            assert _drive(tail, mla, latent) is not None
-            assert tail.prefetch_forks == 0
-        tail.set_successor_weights([router])
-        got = _drive(tail, mla, latent)
-        torch.cuda.synchronize()
-        ref = attn_tail(latent, mla.v_b_proj, mla.o_proj.weight, rows, config=DEFAULT_CONFIG)
-        assert torch.equal(got, ref)
-        assert tail.prefetch_forks == 1 and not tail._prefetch.forked
-    assert torch.equal(router, before)
-
-
-def test_prefetch_kernel_three_regions():
-    mla = tail_adapter._mla()
-    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
-    tensors = [mla.v_b_proj, mla.o_proj.weight, router]
-    assert prefetch.check_prefetch_tensors(tensors, 4096) == ""
-    copies = [t.clone() for t in tensors]
-    prefetch.launch_prefetch(tensors, prefetch.PrefetchConfig())
-    prefetch.launch_prefetch([router], prefetch.PrefetchConfig(ctas=8))
-    torch.cuda.synchronize()
-    for t, c in zip(tensors, copies):
-        assert torch.equal(t, c)
-
-
-@pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED") != "1",
-    reason="needs attention_output_tail and attention_output_tail_prefetch selected",
-)
-def test_adoption_registers_router_weight():
-    """The layer adoption hands the MoE gate weight to the tail."""
-    import inspect
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmDecoderLayer,
-    )
-
-    src = inspect.getsource(GlmDecoderLayer._adopt_tail_prefetch)
-    assert "set_successor_weights" in src and 'is_enabled("attention_output_tail_prefetch")' in src
-
-
-def test_ring_oproj_refuses_publishing():
-    """The ring o_proj has no live-rows-only publish path (the into-hop entry needs
-    every hop row); a caller that asks for it gets an error, not a partial
-    publication.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.cute_oproj_ring import (  # noqa: E501
-        ring_oproj,
-    )
-
-    mla = tail_adapter._mla()
-    v = torch.zeros((6, mla.o_proj.weight.shape[1]), dtype=torch.bfloat16, device="cuda")
-    out = torch.empty((8, mla.o_proj.weight.shape[0]), dtype=torch.bfloat16, device="cuda")
-    words = torch.zeros(16, dtype=torch.int32, device="cuda")
-    flags = torch.zeros(4, dtype=torch.int32, device="cuda")
-    with pytest.raises(ValueError, match="withdrawn"):
-        ring_oproj(
-            v, mla.o_proj.weight, out, use_pdl=False, mc_words=words, flags=flags, publish_rank=0
-        )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py
deleted file mode 100644
index 70d710a337a3..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/test_tail_prefetch_mla_forward.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# 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.
-"""attention_output_tail_prefetch inside the target's ``GlmMLA.forward`` (native DSA attention →
-attention_output_tail_prefetch tail).
-
-Runs with the plan that runs only ``attention_output_tail`` and ``attention_output_tail_prefetch``
-(opt-in: ``GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED=1``). Builds TP8 rank 0 of
-real GLM-5.2 layers 2 (full indexer) and 3 (shared indexer) with the attention_output_tail test
-harness (real weights, FP8 KV cache filled by a real prefill, GEN metadata as the decoder runner
-prepares it) and runs GEN-only ``GlmMLA.forward`` as a decode step through the target's DSA
-attention op:
-
-* eager R = 6 and R = 1: with a successor weight registered (a router-shaped
-  tensor, as the layer adoption registers the MoE gate weight), the
-  target's ``forward_absorption_generation`` forks the staging exactly once
-  per layer call (right before the attention backend call) and ``project``
-  joins it; the attention_output_tail_prefetch output is bit-identical to the attention_output_tail
-  kernels of ``TAIL_PREFETCH_SCHEDULE`` applied to the latent the attention produced (staging
-  changes no result) and within the frozen attention_output_tail budget against the FP64
-  oracle;
-* CUDA-graph capture of the GEN step with the fork/join inside, replayed
-  with new inputs, within budget of a native eager rerun.
-
-Fixture labels: weights real, activations synthetic.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
-import test_attn_tail_mla_forward as tail_forward  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
-    TAIL_PREFETCH_SCHEDULE,
-    attn_tail,
-)
-
-pytestmark = [
-    *tail_forward.pytestmark[:2],
-    pytest.mark.skipif(
-        os.environ.get("GLM_MOE_DSA_ATTENTION_OUTPUT_TAIL_PREFETCH_ADAPTER_PATCH_APPLIED") != "1",
-        reason="needs attention_output_tail and attention_output_tail_prefetch selected (opt-in)",
-    ),
-    pytest.mark.threadleak(enabled=False),
-]
-
-LAYERS = tail_forward.LAYERS
-PREFETCH_PATHS = "attention_output_tail,attention_output_tail_prefetch"
-
-
-@pytest.mark.parametrize("rows", [6, 1])
-def test_eager_forward_forks_joins_and_matches_attn_tail(rows):
-    stack = tail_forward._stack(rows, PREFETCH_PATHS)
-    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
-    for layer in LAYERS:
-        assert stack.attn[layer]._tail_prefetch
-        assert stack.attn_tail_tails[layer].config == TAIL_PREFETCH_SCHEDULE
-        stack.attn_tail_tails[layer].set_successor_weights([router])
-    inputs = {layer: stack.gen_inputs(layer, step=0) for layer in LAYERS}
-    native, lat_native = tail_forward._run(stack, False, inputs, record=True)
-    forks = {layer: stack.attn_tail_tails[layer].prefetch_forks for layer in LAYERS}
-    attention_output_tail, lat = tail_forward._run(stack, True, inputs, record=True)
-    for layer in LAYERS:
-        tail = stack.attn_tail_tails[layer]
-        attn = stack.attn[layer]
-        assert tail.prefetch_forks == forks[layer] + 1, layer
-        assert not tail._prefetch.forked and not tail._armed
-        ref = attn_tail(
-            lat[layer], attn.v_b_proj, attn.o_proj.weight, rows, config=TAIL_PREFETCH_SCHEDULE
-        )
-        assert torch.equal(attention_output_tail[layer], ref), layer
-        e_tail_prefetch, e_nat = tail_forward._check_budget(
-            stack, layer, attention_output_tail[layer], native[layer], lat[layer]
-        )
-        print(
-            f"[attention_output_tail_prefetch-mla-forward] rows={rows} layer={layer} "
-            f"err attention_output_tail_prefetch={e_tail_prefetch['rel_rms']:.3e} native={e_nat['rel_rms']:.3e}"
-        )
-
-
-def test_cuda_graph_capture_and_replay():
-    stack = tail_forward._stack(6, PREFETCH_PATHS)
-    router = torch.randn(256, 6144, device="cuda").to(torch.bfloat16)
-    for layer in LAYERS:
-        stack.attn_tail_tails[layer].set_successor_weights([router])
-    pos = stack.position_ids()
-    static_in = {layer: stack.gen_inputs(layer, step=0).clone() for layer in LAYERS}
-    stack.set_attn_tail(True)
-    with stack.scope():
-        for _ in range(2):
-            for layer in LAYERS:
-                stack.forward(layer, static_in[layer], pos, tail_forward._no_ar())
-        torch.cuda.synchronize()
-        forks = {layer: stack.attn_tail_tails[layer].prefetch_forks for layer in LAYERS}
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            captured = {
-                layer: stack.forward(layer, static_in[layer], pos, tail_forward._no_ar())
-                for layer in LAYERS
-            }
-    for layer in LAYERS:
-        assert stack.attn_tail_tails[layer].prefetch_forks == forks[layer] + 1
-        assert not stack.attn_tail_tails[layer]._prefetch.forked
-    for step in range(1, 5):
-        for layer in LAYERS:
-            static_in[layer].copy_(stack.gen_inputs(layer, step=step))
-        graph.replay()
-        torch.cuda.synchronize()
-        got = {layer: captured[layer].clone() for layer in LAYERS}
-        native, lat = tail_forward._run(
-            stack,
-            False,
-            {layer: static_in[layer].clone() for layer in LAYERS},
-            record=True,
-            position_ids=pos,
-        )
-        for layer in LAYERS:
-            e_tail_prefetch, e_nat = tail_forward._check_budget(
-                stack, layer, got[layer], native[layer], lat[layer]
-            )
-            print(
-                f"[attention_output_tail_prefetch-mla-graph] step={step} layer={layer} "
-                f"err attention_output_tail_prefetch={e_tail_prefetch['rel_rms']:.3e} "
-                f"native={e_nat['rel_rms']:.3e}"
-            )
-    del graph
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py
deleted file mode 100644
index c1ee2d65b758..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_cluster_merge/test_cluster_merge_adapter.py
+++ /dev/null
@@ -1,171 +0,0 @@
-# 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.
-"""Selector and call site of cross_rank_topk_cluster_merge.
-
-- ``cross_rank_topk_cluster_merge`` runs only with ``cross_rank_topk`` (disabling ``cross_rank_topk``
-  turns it off; disabling it keeps ``cross_rank_topk``) and is recorded as a
-  cross_rank_topk-identical ``rank-consistent`` variant.
-- The cross_rank_topk region of the indexer's decode top-k (``region.cross_rank_topk_decode``,
-  which ``GlmIndexer._fill_decode_topk`` asks first) tries cross_rank_topk_cluster_merge
-  first when ``cross_rank_topk_cluster_merge`` is selected and falls back to cross_rank_topk when
-  cross_rank_topk_cluster_merge declines; without ``cross_rank_topk_cluster_merge`` it calls
-  cross_rank_topk only (composition F unchanged), and without ``cross_rank_topk`` nothing.
-- The cross_rank_topk_cluster_merge package registers its entry points and declines a device that
-  cannot keep every row's cluster resident, and a group size its GPU tests
-  do not cover.
-"""
-
-import os
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-
-def _plan(disabled=()):
-    environ = {dispatch.DISABLE_ENV: ",".join(disabled)} if disabled else {}
-    return dispatch.resolve(dispatch.ModelingV2Mode.AUTO, environ)
-
-
-def test_selector_closure():
-    assert dispatch.REQUIRED_FEATURES["cross_rank_topk_cluster_merge"] == {"cross_rank_topk"}
-    assert _plan().enabled("cross_rank_topk_cluster_merge")
-    assert not _plan(["cross_rank_topk"]).enabled("cross_rank_topk_cluster_merge")
-    record = dispatch.feature_record(
-        ["cross_rank_topk", "cross_rank_topk_cluster_merge"], environ={}
-    )
-    assert (
-        "cross_rank_topk_cluster_merge[rank-consistent]" in record
-        and "cross_rank_topk[rank-consistent]" in record
-    )
-
-
-def test_disabling_cluster_merge_keeps_cross_rank_topk():
-    plan = _plan(["cross_rank_topk_cluster_merge"])
-    assert plan.enabled("cross_rank_topk")
-    assert not plan.enabled("cross_rank_topk_cluster_merge")
-
-
-class _Calls:
-    def __init__(self, cluster_merge_result):
-        self.cluster_merge_result = cluster_merge_result
-        self.log = []
-
-    def cross_rank_topk_cluster_merge(self, *args):
-        self.log.append("cross_rank_topk_cluster_merge")
-        return self.cluster_merge_result
-
-    def indexer_topk(self, *args):
-        self.log.append("indexer_topk")
-        return True
-
-
-@pytest.mark.parametrize(
-    "features, cluster_merge_result, expected, calls",
-    [
-        (
-            frozenset({"cross_rank_topk", "cross_rank_topk_cluster_merge"}),
-            True,
-            True,
-            ["cross_rank_topk_cluster_merge"],
-        ),
-        (
-            frozenset({"cross_rank_topk", "cross_rank_topk_cluster_merge"}),
-            False,
-            True,
-            ["cross_rank_topk_cluster_merge", "indexer_topk"],
-        ),
-        (frozenset({"cross_rank_topk"}), True, True, ["indexer_topk"]),
-        (frozenset(), True, False, []),
-    ],
-)
-def test_indexer_hook(monkeypatch, features, cluster_merge_result, expected, calls):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        region,
-    )
-
-    rec = _Calls(cluster_merge_result)
-    monkeypatch.setattr(dispatch, "enabled_features", lambda: features)
-    impls = {
-        ("cross_rank_topk_cluster_merge", "cp_indexer_decode"): rec.cross_rank_topk_cluster_merge,
-        ("indexer_topk", "cp_indexer_decode"): rec.indexer_topk,
-    }
-    monkeypatch.setattr(dispatch, "get_impl", lambda fid, name: impls[(fid, name)])
-    got = region.cross_rank_topk_decode(None, None, None, None, 0, 6, None)
-    assert got is expected
-    assert rec.log == calls
-
-
-def test_package_registration_and_cluster_reason():
-    impl = dispatch.get_impl("cross_rank_topk_cluster_merge", "cp_indexer_decode")
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_cluster_merge import (  # noqa: E501
-        cp_workspace,
-        try_cp_indexer_decode,
-    )
-
-    assert impl is try_cp_indexer_decode
-    assert dispatch.get_impl("cross_rank_topk_cluster_merge", "cp_decode_topk") is not None
-    info = [0] * 12
-    info[cp_workspace.INFO_MAX_ROWS] = 8
-    info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] = 16
-    assert cp_workspace.cluster_reason(info) == ""
-    info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] = 7
-    assert "7 co-resident" in cp_workspace.cluster_reason(info)
-    assert "not built" in cp_workspace.cluster_reason(None)
-
-
-def test_world_size_domain():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_cluster_merge import (  # noqa: E501
-        cp_indexer_decode,
-    )
-
-    assert cp_indexer_decode.SUPPORTED_WORLD_SIZES == (4, 8)
-    for world in (4, 8):
-        assert cp_indexer_decode.world_size_reason(world) == ""
-    for world in (1, 2, 3, 5, 6, 7, 16):
-        assert (
-            "not in cross_rank_topk_cluster_merge's verified sizes"
-            in cp_indexer_decode.world_size_reason(world)
-        )
-
-
-def test_op_info_when_built():
-    import torch
-
-    import tensorrt_llm  # noqa: F401  (the trtllm op namespace)
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_cluster_merge import (  # noqa: E501
-        cp_workspace,
-    )
-
-    if not torch.cuda.is_available():
-        pytest.skip("no GPU")
-    library = os.environ.get("GLM_MOE_DSA_CROSS_RANK_TOPK_CLUSTER_MERGE_JIT_LIB", "")
-    if not hasattr(
-        torch.ops.trtllm, "glm_moe_dsa_cross_rank_topk_cluster_merge_workspace_info"
-    ) and os.path.exists(library):
-        # A JIT build of the op on an image without it.
-        torch.ops.load_library(library)
-    info = cp_workspace.op_info(8)
-    if info is None:
-        pytest.skip(
-            "glm_moe_dsa_cross_rank_topk_cluster_merge is not built into this image (no development build given)"
-        )
-    assert info[cp_workspace.INFO_STATE_WORDS] == cp_workspace.EPOCH_BASE + 64
-    assert info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= info[cp_workspace.INFO_MAX_ROWS]
-    # Cached per group size and device; callers get their own copy.
-    again = cp_workspace.op_info(8)
-    assert again == info and again is not info
-    info4 = cp_workspace.op_info(4)
-    assert info4[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= info4[cp_workspace.INFO_MAX_ROWS]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py
deleted file mode 100644
index aa01442ce1b6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_adapter.py
+++ /dev/null
@@ -1,535 +0,0 @@
-# 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.
-"""Selector, call sites and bridge plan of cross_rank_topk_global_emit.
-
-- ``cross_rank_topk_global_emit`` runs only with ``rope_kv_append``, ``cross_rank_topk``
-  and ``cross_rank_topk_cluster_merge`` and is recorded as a ``rank-consistent`` variant; disabling it
-  keeps the paths it extends.
-- The cross_rank_topk region of the indexer's decode top-k (``region.cross_rank_topk_decode``)
-  tries cross_rank_topk_global_emit first (with the call's rope_kv_append packet, which
-  ``GlmIndexer._fill_decode_topk`` passes) when ``cross_rank_topk_global_emit`` is selected,
-  then cross_rank_topk_cluster_merge, then cross_rank_topk; without ``cross_rank_topk_global_emit``
-  it is B5's hook.
-- The backend's packet and remap (``GlmDsaBackend.sparse_attn_predict``) come from
-  cross_rank_topk_global_emit when it is selected.
-- ``bridge.predict_copy`` / ``predict_remap`` reproduce the call sites' rules
-  (rope_kv_append's stash, the backend's shared copy, rope_kv_append's grouping), and
-  ``remap_topk_to_global`` returns the planned pool rows, falls back to rope_kv_append's
-  remap without a plan, and raises on any difference from the plan.
-- The package registers its entry points; group sizes and op info.
-- deferred_input_prep's tables list cross_rank_topk_global_emit's package directories and its side
-  selector.
-"""
-
-import inspect
-import os
-import re
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-
-def _plan(disabled=()):
-    environ = {dispatch.DISABLE_ENV: ",".join(disabled)} if disabled else {}
-    return dispatch.resolve(dispatch.ModelingV2Mode.AUTO, environ)
-
-
-def test_selector_closure():
-    assert dispatch.REQUIRED_FEATURES["cross_rank_topk_global_emit"] == {
-        "rope_kv_append",
-        "cross_rank_topk",
-        "cross_rank_topk_cluster_merge",
-    }
-    assert _plan().enabled("cross_rank_topk_global_emit")
-    for base in ("rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"):
-        assert not _plan([base]).enabled("cross_rank_topk_global_emit"), base
-    record = dispatch.feature_record(
-        [
-            "rope_kv_append",
-            "cross_rank_topk",
-            "cross_rank_topk_cluster_merge",
-            "cross_rank_topk_global_emit",
-        ],
-        environ={},
-    )
-    assert (
-        "cross_rank_topk_global_emit[rank-consistent]" in record
-        and "cross_rank_topk_cluster_merge[rank-consistent]" in record
-    )
-
-
-def test_disabling_global_emit_keeps_cluster_merge():
-    plan = _plan(["cross_rank_topk_global_emit"])
-    assert not plan.enabled("cross_rank_topk_global_emit")
-    for base in ("rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"):
-        assert plan.enabled(base), base
-
-
-class _Calls:
-    def __init__(self, global_emit_result, cluster_merge_result):
-        self.global_emit_result = global_emit_result
-        self.cluster_merge_result = cluster_merge_result
-        self.log = []
-        self.packets = []
-
-    def cross_rank_topk_global_emit(self, *args):
-        self.log.append("cross_rank_topk_global_emit")
-        self.packets.append(args[-1])
-        assert len(args) == 8
-        return self.global_emit_result
-
-    def cross_rank_topk_cluster_merge(self, *args):
-        self.log.append("cross_rank_topk_cluster_merge")
-        assert len(args) == 7
-        return self.cluster_merge_result
-
-    def indexer_topk(self, *args):
-        self.log.append("indexer_topk")
-        return True
-
-
-@pytest.mark.parametrize(
-    "features, global_emit_result, cluster_merge_result, calls",
-    [
-        (
-            frozenset(
-                {
-                    "rope_kv_append",
-                    "cross_rank_topk",
-                    "cross_rank_topk_cluster_merge",
-                    "cross_rank_topk_global_emit",
-                }
-            ),
-            True,
-            True,
-            ["cross_rank_topk_global_emit"],
-        ),
-        (
-            frozenset(
-                {
-                    "rope_kv_append",
-                    "cross_rank_topk",
-                    "cross_rank_topk_cluster_merge",
-                    "cross_rank_topk_global_emit",
-                }
-            ),
-            False,
-            True,
-            ["cross_rank_topk_global_emit", "cross_rank_topk_cluster_merge"],
-        ),
-        (
-            frozenset(
-                {
-                    "rope_kv_append",
-                    "cross_rank_topk",
-                    "cross_rank_topk_cluster_merge",
-                    "cross_rank_topk_global_emit",
-                }
-            ),
-            False,
-            False,
-            ["cross_rank_topk_global_emit", "cross_rank_topk_cluster_merge", "indexer_topk"],
-        ),
-        (
-            frozenset({"rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"}),
-            True,
-            True,
-            ["cross_rank_topk_cluster_merge"],
-        ),
-        (frozenset({"rope_kv_append", "cross_rank_topk"}), True, True, ["indexer_topk"]),
-        (frozenset(), True, True, []),
-    ],
-)
-def test_indexer_hook(monkeypatch, features, global_emit_result, cluster_merge_result, calls):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        region,
-    )
-
-    rec = _Calls(global_emit_result, cluster_merge_result)
-    monkeypatch.setattr(dispatch, "enabled_features", lambda: features)
-    impls = {
-        ("cross_rank_topk_global_emit", "cp_indexer_decode"): rec.cross_rank_topk_global_emit,
-        ("cross_rank_topk_cluster_merge", "cp_indexer_decode"): rec.cross_rank_topk_cluster_merge,
-        ("indexer_topk", "cp_indexer_decode"): rec.indexer_topk,
-    }
-    monkeypatch.setattr(dispatch, "get_impl", lambda fid, name: impls[(fid, name)])
-    packet = object()
-    got = region.cross_rank_topk_decode(None, None, None, None, 0, 6, None, packet)
-    # Each tried region but the last declined; the last one served the call.
-    assert got is bool(calls)
-    assert rec.log == calls
-    assert all(p is packet for p in rec.packets)
-
-
-def test_indexer_call_site_passes_the_packet():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-        GlmIndexer,
-    )
-
-    src = inspect.getsource(GlmIndexer._fill_decode_topk)
-    call = src[src.index("region.cross_rank_topk_decode(") :]
-    call = call[: call.index(")\n")]
-    assert "self._topk_bridge" in call
-
-
-def test_backend_uses_global_emit_packet_and_remap_only_when_selected():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-        GlmDsaBackend,
-    )
-
-    src = inspect.getsource(GlmDsaBackend.sparse_attn_predict)
-    # The formatter wraps long statements: compare with the whitespace collapsed.
-    src = re.sub(r"(?<=[(\[])\s+|\s+(?=[)\]])", "", " ".join(src.split()))
-    assert (
-        'bridge_path = ("cross_rank_topk_global_emit" if '
-        'dispatch.is_enabled("cross_rank_topk_global_emit") else "rope_kv_append")' in src
-    )
-    assert 'get_impl(bridge_path, "new_bridge_packet")' in src
-    assert 'bridge_path, "remap_topk_to_global"' in src
-    # Every other rope_kv_append call site of the backend is rope_kv_append's own.
-    assert src.count('get_impl("rope_kv_append", "defer_shared_copy")') == 1
-
-
-def test_package_registration():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
-        cross_rank_topk_global_emit,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-        cp_indexer_decode,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        ops as rope_ops,
-    )
-
-    assert (
-        dispatch.get_impl("cross_rank_topk_global_emit", "cp_indexer_decode")
-        is cp_indexer_decode.try_cp_indexer_decode
-    )
-    assert (
-        dispatch.get_impl("cross_rank_topk_global_emit", "new_bridge_packet")
-        is bridge.new_bridge_packet
-    )
-    assert (
-        dispatch.get_impl("cross_rank_topk_global_emit", "remap_topk_to_global")
-        is bridge.remap_topk_to_global
-    )
-    assert issubclass(cross_rank_topk_global_emit.GlobalEmitBridgePacket, rope_ops.BridgePacket)
-    assert cross_rank_topk_global_emit.GlobalEmitBridgePacket().plan is None
-
-
-def test_world_size_domain():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        cp_indexer_decode,
-    )
-
-    assert cp_indexer_decode.SUPPORTED_WORLD_SIZES == (4, 8)
-    for world in (4, 8):
-        assert cp_indexer_decode.world_size_reason(world) == ""
-    for world in (1, 2, 3, 5, 6, 7, 16):
-        assert (
-            "not in cross_rank_topk_global_emit's verified sizes"
-            in cp_indexer_decode.world_size_reason(world)
-        )
-
-
-# ---- Bridge plan ------------------------------------------------------------------------------------------------
-
-
-class _Kvm:
-    def __init__(self, tokens_per_block=64):
-        self.tokens_per_block = tokens_per_block
-        self.layer_offsets = {layer: layer for layer in range(100)}
-
-    def get_primary_pool_page_index_params(self, local_layer):
-        return 79, local_layer
-
-
-def _struct(group_ids):
-    """Group struct of rope_kv_append's grouping: layers 0, 1 singletons, then groups of len(group_ids)."""
-    n = 2 + 2 * len(group_ids)
-    leader_of = [0, 1] + [2] * len(group_ids) + [2 + len(group_ids)] * len(group_ids)
-    slot_of = [0, 0] + list(range(len(group_ids))) * 2
-    return {
-        "leader_of": leader_of[:n],
-        "slot_of": slot_of[:n],
-        "group_active": {0: False, 1: False, 2: True, 2 + len(group_ids): True},
-        "group_scale": {2: 79, 2 + len(group_ids): 79},
-        "group_layer_ids": {
-            2: torch.tensor(group_ids, dtype=torch.int32),
-            2 + len(group_ids): torch.tensor(group_ids, dtype=torch.int32),
-        },
-    }
-
-
-def _metadata(*, draft=False, shared=True, accepted=torch.int32, group=4, rows=6, struct=None):
-    st = _struct(list(range(group))) if struct is None else struct
-    return SimpleNamespace(
-        num_contexts=0,
-        num_generations=1,
-        num_ctx_tokens=0,
-        num_tokens=rows,
-        in_mtp_draft_loop=draft,
-        shared_topk_indices=torch.zeros(8, 2048, dtype=torch.int32) if shared else None,
-        mtp_num_accepted=None if accepted is None else torch.tensor([3, 1], dtype=accepted),
-        kv_cache_manager=_Kvm(),
-        _ensure_group_remap_struct=lambda: st,
-        _group_remap_batched={},
-    )
-
-
-def _indexer(share=True, layer=2):
-    return SimpleNamespace(
-        mtp_index_share=share, layer_idx=layer, top_k=SimpleNamespace(top_k=2048)
-    )
-
-
-def test_predict_copy_rules():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-    )
-
-    # Target layer: the backend's shared copy of the generation rows.
-    md = _metadata()
-    c = bridge.predict_copy(_indexer(), md, 6)
-    assert (c.dst is md.shared_topk_indices, c.row_start, c.rows, c.mode) == (True, 0, 6, 0)
-    assert bridge.predict_copy(_indexer(), _metadata(shared=False), 6) is None
-    # MTP draft step 0 with index sharing: rope_kv_append's stash (mode 1, the request's accepted
-    # count).
-    md = _metadata(draft=True)
-    c = bridge.predict_copy(_indexer(), md, 6)
-    assert (c.dst is md.shared_topk_indices, c.row_start, c.rows, c.mode, c.next_n) == (
-        True,
-        0,
-        1,
-        1,
-        6,
-    )
-    assert c.num_accepted.dtype == torch.int32 and c.num_accepted.tolist() == [3]
-    c = bridge.predict_copy(_indexer(), _metadata(draft=True, accepted=None), 6)
-    assert c.mode == 1 and c.num_accepted is None
-    c = bridge.predict_copy(_indexer(), _metadata(draft=True, accepted=torch.int64), 6)
-    assert c.num_accepted.dtype == torch.int64
-    # The stash is not deferred (rope_kv_append declines): the native stash runs and the backend
-    # keeps the rows.
-    assert bridge.predict_copy(_indexer(), _metadata(draft=True, shared=False), 6) is None
-    assert bridge.predict_copy(_indexer(), _metadata(draft=True, accepted=torch.float32), 6) is None
-    # A draft step without index sharing takes the backend's shared copy.
-    c = bridge.predict_copy(_indexer(share=False), _metadata(draft=True), 6)
-    assert c.mode == 0 and c.rows == 6
-
-
-def test_predict_remap_rules(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-    )
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
-
-    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
-    md = _metadata(group=4)
-    r = bridge.predict_remap(md, 2)
-    assert (r.grouped, r.leader, r.slot, r.page_index_scale) == (True, 2, 0, 79)
-    assert r.layer_ids is md._ensure_group_remap_struct()["group_layer_ids"][2]
-    r = bridge.predict_remap(md, 1)
-    assert (r.grouped, r.page_index_scale, r.layer_offset) == (False, 79, 1)
-    assert bridge.predict_remap(_metadata(group=9), 2) is None
-    r = bridge.predict_remap(_metadata(draft=True), 2)
-    assert not r.grouped and r.layer_offset == 2
-    r = bridge.predict_remap(_metadata(struct={}), 2)
-    assert not r.grouped
-    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", False)
-    assert not bridge.predict_remap(_metadata(), 2).grouped
-
-
-def _backend(layer=2, indexer=True):
-    return SimpleNamespace(
-        indexer=object() if indexer else None, get_local_layer_idx=lambda metadata: layer
-    )
-
-
-def _planned(md, layer=2, grouped=True):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-    )
-
-    ids = torch.zeros(6, 2048, dtype=torch.int32)
-    remap = bridge.predict_remap(md, layer)
-    copy = bridge.predict_copy(_indexer(layer=layer), md, 6)
-    pool = (
-        torch.arange(4 * 6 * 2048, dtype=torch.int32).view(4, 6, 2048)
-        if remap.grouped
-        else torch.ones(6, 2048, dtype=torch.int32)
-    )
-    packet = bridge.GlobalEmitBridgePacket()
-    packet.plan = bridge.GlobalEmitPlan(ids=ids, remap=remap, copy=copy, pool=pool)
-    return packet, ids, pool, copy
-
-
-def test_remap_uses_the_plan_once(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        ops as rope_ops,
-    )
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
-
-    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
-    rope_append_calls = []
-    monkeypatch.setattr(
-        rope_ops, "remap_topk_to_global", lambda *a: rope_append_calls.append(a) or "rope_kv_append"
-    )
-    md = _metadata()
-    packet, ids, pool, copy = _planned(md, layer=3)
-    rope_ops.defer_shared_copy(packet, md.shared_topk_indices, 0, ids)
-    out = bridge.remap_topk_to_global(_backend(layer=3), ids, md, packet)
-    assert out.data_ptr() == pool[1].data_ptr() and md._group_remap_batched[2] is pool
-    assert packet.plan is None and not rope_append_calls
-    # A packet without a plan (draft reuse, declines) is rope_kv_append's remap.
-    assert (
-        bridge.remap_topk_to_global(_backend(layer=3), ids, md, packet) == "rope_kv_append"
-        and len(rope_append_calls) == 1
-    )
-    # Per-layer plan.
-    packet, ids, pool, copy = _planned(md, layer=1)
-    rope_ops.defer_shared_copy(packet, md.shared_topk_indices, 0, ids)
-    assert bridge.remap_topk_to_global(_backend(layer=1), ids, md, packet) is pool
-
-
-def test_remap_rejects_differences(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        ops as rope_ops,
-    )
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
-
-    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
-
-    def check(mutate, match):
-        md = _metadata()
-        packet, ids, pool, copy = _planned(md, layer=2)
-        rope_ops.defer_shared_copy(packet, md.shared_topk_indices, 0, ids)
-        args = mutate(md, packet, ids)
-        with pytest.raises(bridge.GlobalEmitContractError, match=match):
-            bridge.remap_topk_to_global(*args)
-
-    check(lambda md, p, ids: (_backend(), ids, md, setattr(p, "copy", None) or p), "copy")
-    check(lambda md, p, ids: (_backend(), ids, md, setattr(p.copy, "rows", 5) or p), "copy")
-    check(lambda md, p, ids: (_backend(), ids, md, setattr(p, "source", ids) or p), "source")
-    check(
-        lambda md, p, ids: (_backend(), ids.clone(), md, p),
-        "rows cross_rank_topk_global_emit wrote",
-    )
-    check(lambda md, p, ids: (_backend(layer=3), ids, md, p), "local layer")
-    check(lambda md, p, ids: (_backend(indexer=False), ids, md, p), "no indexer")
-
-    def regroup(md, p, ids):
-        md._ensure_group_remap_struct = lambda: {}
-        return _backend(), ids, md, p
-
-    check(regroup, "remap decision")
-
-
-def test_copy_reason_matches_the_bridge_domain():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        ops as rope_ops,
-    )
-
-    if not torch.cuda.is_available():
-        pytest.skip("no GPU")
-    dst = torch.zeros(8, 2048, dtype=torch.int32, device="cuda")
-    assert bridge.copy_reason(None, 6, 2048) == ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 6), 6, 2048) == ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 3, 6), 6, 2048) != ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 7), 6, 2048) != ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst.cpu(), 0, 6), 6, 2048) != ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst[:, :100], 0, 6), 6, 2048) != ""
-    acc = torch.tensor([2], dtype=torch.int32, device="cuda")
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 1, 1, 6, acc), 6, 2048) == ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 2, 1, 6, None), 6, 2048) != ""
-    assert bridge.copy_reason(rope_ops.BridgeCopy(dst, 0, 1, 1, 6, acc.float()), 6, 2048) != ""
-
-
-def test_bridge_reason_declines_before_device_work(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge,
-        cp_indexer_decode,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        ops as rope_ops,
-    )
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import backend as dsa_backend
-
-    monkeypatch.setattr(dsa_backend, "_GROUP_REMAP", True)
-    reason, _, _ = cp_indexer_decode.bridge_reason(
-        _indexer(), _metadata(), rope_ops.BridgePacket(), 6
-    )
-    assert "no cross_rank_topk_global_emit bridge packet" in reason
-    reason, _, _ = cp_indexer_decode.bridge_reason(_indexer(), _metadata(), None, 6)
-    assert "no cross_rank_topk_global_emit bridge packet" in reason
-    md = _metadata()
-    md.kv_cache_manager.tokens_per_block = 32
-    reason, _, _ = cp_indexer_decode.bridge_reason(
-        _indexer(), md, bridge.GlobalEmitBridgePacket(), 6
-    )
-    assert "block size" in reason
-    reason, _, _ = cp_indexer_decode.bridge_reason(
-        _indexer(), _metadata(group=9), bridge.GlobalEmitBridgePacket(), 6
-    )
-    assert "group" in reason
-    used = bridge.GlobalEmitBridgePacket()
-    used.copy = rope_ops.BridgeCopy(torch.zeros(1, 1), 0, 1)
-    reason, _, _ = cp_indexer_decode.bridge_reason(_indexer(), _metadata(), used, 6)
-    assert "already holds" in reason
-
-
-def test_op_info_when_built():
-    import tensorrt_llm  # noqa: F401  (the trtllm op namespace)
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        cp_workspace,
-    )
-
-    if not torch.cuda.is_available():
-        pytest.skip("no GPU")
-    library = os.environ.get("GLM_MOE_DSA_CROSS_RANK_TOPK_GLOBAL_EMIT_JIT_LIB", "")
-    if not hasattr(
-        torch.ops.trtllm, "glm_moe_dsa_cross_rank_topk_global_emit_workspace_info"
-    ) and os.path.exists(library):
-        # A JIT build of the op on an image without it.
-        torch.ops.load_library(library)
-    info = cp_workspace.op_info(8)
-    if info is None:
-        pytest.skip(
-            "glm_moe_dsa_cross_rank_topk_global_emit is not built into this image (no development build given)"
-        )
-    assert info[cp_workspace.INFO_STATE_WORDS] == cp_workspace.EPOCH_BASE + 64
-    assert info[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= info[cp_workspace.INFO_MAX_ROWS]
-    assert info[cp_workspace.INFO_MAX_GROUP_SIZE] == 8
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk.cross_rank_logits_runner import (  # noqa: E501
-        local_capacity,
-    )
-
-    for world in (4, 8):
-        winfo = cp_workspace.op_info(world)
-        assert winfo[cp_workspace.INFO_MAX_LOCAL_CAPACITY] >= local_capacity((1 << 20) + 14, world)
-        assert winfo[cp_workspace.INFO_MAX_ACTIVE_CLUSTERS] >= winfo[cp_workspace.INFO_MAX_ROWS]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py
deleted file mode 100644
index 14fc28efef99..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/cross_rank_topk_global_emit/test_global_emit_composition_global_emit_indexer_front.py
+++ /dev/null
@@ -1,387 +0,0 @@
-# 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.
-"""cross_rank_topk_global_emit and indexer_layer_front composed in the cross_rank_topk CP indexer
-region (host-side checks).
-
-cross_rank_topk_global_emit forks the region's top-k (its emission also writes rope_kv_append's
-bridge outputs) and indexer_layer_front forks the region's CP logits launch (entry wait, PDL
-launch). With both selected, cross_rank_topk_global_emit serves the region and takes
-indexer_layer_front's logits launcher; the indexer hook tries cross_rank_topk_global_emit, then
-indexer_layer_front (cross_rank_topk_cluster_merge's or cross_rank_topk's top-k), then
-cross_rank_topk_cluster_merge, then cross_rank_topk. These tests mock every device step: they check
-which entry serves a call, which logits launcher cross_rank_topk_global_emit uses, and the
-capability lines; the numerics are covered by the GPU tests of each package.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-
-class _Calls:
-    def __init__(self, results):
-        self.results = results
-        self.log = []
-
-    def make(self, name, nargs):
-        def entry(*args):
-            self.log.append(name)
-            assert len(args) == nargs, (name, len(args))
-            return self.results[name]
-
-        return entry
-
-
-_ALL = {
-    "cross_rank_topk_global_emit": True,
-    "indexer_layer_front": True,
-    "cross_rank_topk_cluster_merge": True,
-    "indexer_topk": True,
-}
-
-
-@pytest.mark.parametrize(
-    "features, results, calls",
-    [
-        # F4: cross_rank_topk_global_emit serves (with indexer_layer_front's logits).
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "cross_rank_topk_global_emit",
-                "indexer_layer_front",
-            },
-            _ALL,
-            ["cross_rank_topk_global_emit"],
-        ),
-        # cross_rank_topk_global_emit declines: indexer_layer_front's region
-        # (cross_rank_topk_cluster_merge's top-k) serves.
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "cross_rank_topk_global_emit",
-                "indexer_layer_front",
-            },
-            {**_ALL, "cross_rank_topk_global_emit": False},
-            ["cross_rank_topk_global_emit", "indexer_layer_front"],
-        ),
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "cross_rank_topk_global_emit",
-                "indexer_layer_front",
-            },
-            {**_ALL, "cross_rank_topk_global_emit": False, "indexer_layer_front": False},
-            ["cross_rank_topk_global_emit", "indexer_layer_front", "cross_rank_topk_cluster_merge"],
-        ),
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "cross_rank_topk_global_emit",
-                "indexer_layer_front",
-            },
-            {
-                "cross_rank_topk_global_emit": False,
-                "indexer_layer_front": False,
-                "cross_rank_topk_cluster_merge": False,
-                "indexer_topk": True,
-            },
-            [
-                "cross_rank_topk_global_emit",
-                "indexer_layer_front",
-                "cross_rank_topk_cluster_merge",
-                "indexer_topk",
-            ],
-        ),
-        # F4 - cross_rank_topk_global_emit: indexer_layer_front first.
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "indexer_layer_front",
-            },
-            _ALL,
-            ["indexer_layer_front"],
-        ),
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "indexer_layer_front",
-            },
-            {**_ALL, "indexer_layer_front": False},
-            ["indexer_layer_front", "cross_rank_topk_cluster_merge"],
-        ),
-        # Without indexer_layer_front: cross_rank_topk_global_emit first, then
-        # cross_rank_topk_cluster_merge.
-        (
-            {
-                "rope_kv_append",
-                "cross_rank_topk",
-                "cross_rank_topk_cluster_merge",
-                "cross_rank_topk_global_emit",
-            },
-            {**_ALL, "cross_rank_topk_global_emit": False},
-            ["cross_rank_topk_global_emit", "cross_rank_topk_cluster_merge"],
-        ),
-        # cross_rank_topk_cluster_merge only.
-        (
-            {"rope_kv_append", "cross_rank_topk", "cross_rank_topk_cluster_merge"},
-            _ALL,
-            ["cross_rank_topk_cluster_merge"],
-        ),
-        (set(), _ALL, []),
-    ],
-)
-def test_indexer_hook_order(monkeypatch, features, results, calls):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        region,
-    )
-
-    rec = _Calls(results)
-    impls = {
-        ("cross_rank_topk_global_emit", "cp_indexer_decode"): rec.make(
-            "cross_rank_topk_global_emit", 8
-        ),
-        ("indexer_layer_front", "cp_indexer_decode"): rec.make("indexer_layer_front", 7),
-        ("cross_rank_topk_cluster_merge", "cp_indexer_decode"): rec.make(
-            "cross_rank_topk_cluster_merge", 7
-        ),
-        ("indexer_topk", "cp_indexer_decode"): rec.make("indexer_topk", 7),
-    }
-    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(features))
-    monkeypatch.setattr(dispatch, "get_impl", lambda fid, name: impls[(fid, name)])
-    got = region.cross_rank_topk_decode(None, None, None, None, 0, 6, None, object())
-    assert rec.log == calls
-    assert got is bool(calls) and (not calls or results[calls[-1]])
-
-
-def _region_inputs(num_heads=64, head_dim=128, next_n=6):
-    top_k = SimpleNamespace(
-        needs_gvr_prior=False,
-        decode_implementation=SimpleNamespace(value="default"),
-        gvr_self_sampling=False,
-        top_k=2048,
-        compress_ratio=1,
-    )
-    indexer = SimpleNamespace(
-        top_k=top_k,
-        use_cute_dsl_paged_mqa_logits=False,
-        use_fp4=False,
-        use_gvr_emission=False,
-        n_heads=num_heads,
-        head_dim=head_dim,
-        layer_idx=3,
-    )
-    kvm = SimpleNamespace(
-        get_indexer_k_cache_buffers=lambda layer: torch.zeros(
-            4, 64, 1, head_dim + 4, dtype=torch.uint8
-        ),
-        tokens_per_block=64,
-        layer_offsets={3: 3},
-    )
-    metadata = SimpleNamespace(
-        num_contexts=0,
-        num_generations=1,
-        seq_lens=torch.tensor([next_n], dtype=torch.int32),
-        kv_cache_manager=kvm,
-        indexer_k_cache_block_offsets=torch.zeros(1, 16, dtype=torch.int32),
-        mapping=SimpleNamespace(tp_size=8, enable_attention_dp=False),
-        get_indexer_max_seq_len=lambda: 32768,
-        is_ragged_verify=False,
-        use_expanded_buffers_for_mtp=False,
-        in_mtp_draft_loop=False,
-        gen_indexer_kv_lens_cuda_runtime=torch.tensor([32768], dtype=torch.int32),
-        req_idx_per_token=torch.zeros(next_n, dtype=torch.int32),
-        num_ctx_tokens=0,
-        num_tokens=next_n,
-        block_table=torch.zeros(1, 512, dtype=torch.int32),
-        num_seqs=1,
-    )
-    q_fp8 = torch.zeros(next_n, num_heads, head_dim)
-    weights = torch.zeros(next_n, num_heads)
-    out = torch.zeros(next_n, 2048, dtype=torch.int32)
-    return indexer, metadata, q_fp8, weights, out
-
-
-@pytest.mark.parametrize("with_indexer_front", [True, False])
-def test_global_emit_takes_the_logits_launcher_of_indexer_front(monkeypatch, with_indexer_front):
-    """cross_rank_topk_global_emit's region launches indexer_layer_front's logits kernel exactly
-    when indexer_layer_front is selected."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        bridge as global_emit_bridge,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        cp_indexer_decode as global_emit_decode,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        cp_workspace as global_emit_ws,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
-        cp_logits_runner as indexer_front_runner,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        cross_rank_logits_runner as indexer_topk_runner,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        cross_rank_policy,
-    )
-
-    features = {
-        "absorbed_query",
-        "indexer_topk",
-        "draft_absorbed_query",
-        "rope_kv_append",
-        "cross_rank_topk",
-        "cross_rank_topk_cluster_merge",
-        "cross_rank_topk_global_emit",
-    } | ({"indexer_layer_front"} if with_indexer_front else set())
-    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(features))
-    launched = []
-
-    def runner(owner):
-        def ensure_compiled(*args, **kwargs):
-            launched.append((owner, "compile"))
-
-        def paged_mqa_logits_cp(
-            q, k_cache, weights, kv_lens, block_table, max_seq_len, rank, world
-        ):
-            launched.append((owner, "logits"))
-            assert (rank, world) == (2, 8)
-            return torch.zeros(q.shape[1], 16)
-
-        return ensure_compiled, paged_mqa_logits_cp
-
-    for owner, mod in (
-        ("indexer_layer_front", indexer_front_runner),
-        ("indexer_topk", indexer_topk_runner),
-    ):
-        ensure, logits = runner(owner)
-        monkeypatch.setattr(mod, "ensure_compiled", ensure)
-        monkeypatch.setattr(mod, "paged_mqa_logits_cp", logits)
-    monkeypatch.setattr(
-        cross_rank_policy,
-        "decide_cp_uniform",
-        lambda **kw: SimpleNamespace(eligible=True, reason=""),
-    )
-    monkeypatch.setattr(cross_rank_policy, "rank_local_reason", lambda **kw: "")
-    monkeypatch.setattr(global_emit_ws, "op_info", lambda world: (0,) * 16)
-    monkeypatch.setattr(global_emit_ws, "cluster_reason", lambda info: "")
-    workspace = SimpleNamespace(rank=2, world=8, state=None, mc_ptr=0, uc_ptr=0)
-
-    def get_workspace(mapping, prepare=None, local_reason=""):
-        prepare()
-        return workspace
-
-    monkeypatch.setattr(global_emit_ws, "get_workspace", get_workspace)
-    remap = SimpleNamespace(page_index_scale=1, grouped=False, layer_ids=None, layer_offset=3)
-    monkeypatch.setattr(global_emit_decode, "bridge_reason", lambda *a: ("", remap, None))
-    topk_calls = []
-
-    def topk(local_logits, kv_lens, out, ws, next_n, **kw):
-        topk_calls.append((tuple(local_logits.shape), next_n, kw["layer_id"]))
-        return torch.zeros(next_n, 2048, dtype=torch.int32)
-
-    monkeypatch.setattr(global_emit_decode, "cp_topk_from_local_logits", topk)
-    logged = {}
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", logged)
-
-    indexer, metadata, q_fp8, weights, out = _region_inputs()
-    packet = SimpleNamespace(plan=None)
-    assert global_emit_decode.try_cp_indexer_decode(
-        indexer, metadata, q_fp8, weights, 0, 6, out, packet
-    )
-    owner = "indexer_layer_front" if with_indexer_front else "indexer_topk"
-    assert launched == [(owner, "compile"), (owner, "logits")]
-    assert topk_calls == [((6, 16), 6, 3)]
-    assert isinstance(packet.plan, global_emit_bridge.GlobalEmitPlan) and packet.plan.ids is out
-    assert logged.get(("cross_rank_topk_global_emit", "cp_indexer_decode")) is True
-    assert (
-        ("indexer_layer_front", "cp_indexer_decode_global_emit") in logged
-    ) is with_indexer_front
-    if with_indexer_front:
-        assert logged[("indexer_layer_front", "cp_indexer_decode_global_emit")] is True
-
-
-def test_global_emit_decline_is_logged_for_both(monkeypatch):
-    """A cross_rank_topk_global_emit decline with indexer_layer_front selected is visible on both capability keys."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        cp_indexer_decode as global_emit_decode,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.cross_rank_topk_global_emit import (  # noqa: E501
-        cp_workspace as global_emit_ws,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        cross_rank_policy,
-    )
-
-    features = {
-        "absorbed_query",
-        "indexer_topk",
-        "draft_absorbed_query",
-        "rope_kv_append",
-        "cross_rank_topk",
-        "cross_rank_topk_cluster_merge",
-        "cross_rank_topk_global_emit",
-        "indexer_layer_front",
-    }
-    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(features))
-    monkeypatch.setattr(
-        cross_rank_policy,
-        "decide_cp_uniform",
-        lambda **kw: SimpleNamespace(eligible=False, reason="why"),
-    )
-    monkeypatch.setattr(cross_rank_policy, "rank_local_reason", lambda **kw: "")
-    monkeypatch.setattr(global_emit_ws, "op_info", lambda world: (0,) * 16)
-    monkeypatch.setattr(global_emit_ws, "cluster_reason", lambda info: "")
-    logged = {}
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", logged)
-    indexer, metadata, q_fp8, weights, out = _region_inputs()
-    packet = SimpleNamespace(plan=None)
-    assert not global_emit_decode.try_cp_indexer_decode(
-        indexer, metadata, q_fp8, weights, 0, 6, out, packet
-    )
-    assert packet.plan is None
-    assert logged == {
-        ("cross_rank_topk_global_emit", "cp_indexer_decode"): False,
-        ("indexer_layer_front", "cp_indexer_decode_global_emit"): False,
-    }
-
-
-def test_mla_rope_line():
-    """The rope_kv_append rope branch logs indexer_layer_front/rope and picks indexer_layer_front's
-    rope only when its op is loaded."""
-    import inspect
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-        attention_forward,
-    )
-
-    src = inspect.getsource(attention_forward.forward_absorption_generation)
-    assert '"indexer_layer_front", "rope_available"' in src
-    assert '"indexer_layer_front",\n' in src and '"rope",' in src
-    assert 'rope_owner = "rope_kv_append"' in src
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py
deleted file mode 100644
index 45eb2446a16a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/_decode_prologue_oracle.py
+++ /dev/null
@@ -1,355 +0,0 @@
-# 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.
-"""Native oracle of the decode_prologue target prologue (shared by the tests and the benchmark).
-
-The native reference is the literal ``DecoderRunner._preprocess_inputs``
-(called unbound on a runner stub that holds the overlap-scheduler offset
-buffers, the configuration and the mapping) followed by the target model's
-native ``Embedding``. It runs on a real CUDA-graph ``DSAtrtllmAttentionMetadata``
-prepared from a real ``DSACacheManager`` at the served GEN DSA
-configuration (draft_loop's fixtures), exactly as the decoder runner prepares a decode
-step: ``create_cuda_graph_metadata``, ``update_spec_dec_param``,
-``KVCacheParams``, ``prepare()``.
-"""
-
-import sys
-from pathlib import Path
-from types import SimpleNamespace
-from typing import Dict, List, Optional, Sequence
-
-import torch
-
-_DRAFT_LOOP_TESTS = Path(__file__).resolve().parent.parent / "draft_loop"
-sys.path.insert(0, str(_DRAFT_LOOP_TESTS))
-
-from _draft_loop_fixtures import (  # noqa: E402
-    MTP_K,
-    add_gen_requests,
-    load_checkpoint_tensor,
-    make_dsa_cache_manager,
-)
-
-VOCAB = 154880
-HIDDEN = 6144
-# Served GEN: max_num_tokens 128, max_batch_size 1 (2 for the padded case).
-MAX_NUM_TOKENS = 128
-POISON = 0xA5
-PACKET_ATTR = "_decode_prologue_embedding"
-
-# Every metadata buffer the native prologue writes; poisoned identically in
-# both arms before each run and compared in full afterwards.
-DEST_FIELDS = (
-    "kv_lens_cuda_2d",
-    "scheduler_metadata_buffer",
-    "scheduler_metadata_buffer_full_next_n",
-    "kv_lens_expanded_cuda",
-    "scheduler_metadata_buffer_expanded",
-    "req_idx_per_token",
-    "slot_mapping_fp8",
-    "slot_mapping_scale",
-    "gen_kv_indptr",
-    "gen_cached_token_indptr",
-    "topk_indices_buffer",
-)
-
-
-def make_graph_metadata(
-    cache_manager, mapping, sparse_config, request_ids, past_kv_lens, draft_len: int = MTP_K
-):
-    """A prepared CUDA-graph decode metadata (the runner's per-step preparation)."""
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata
-    from tensorrt_llm._torch.metadata import KVCacheParams
-
-    batch_size = len(request_ids)
-    base = DSAtrtllmAttentionMetadata(
-        seq_lens=torch.full((batch_size,), draft_len + 1, dtype=torch.int),
-        request_ids=list(request_ids),
-        max_num_requests=batch_size,
-        num_contexts=0,
-        prompt_lens=list(past_kv_lens),
-        max_num_tokens=MAX_NUM_TOKENS,
-        kv_cache_manager=cache_manager,
-        mapping=mapping,
-        sparse_attention_config=sparse_config,
-    )
-    md = base.create_cuda_graph_metadata(batch_size, False, draft_len)
-    md.request_ids = list(request_ids)
-    md.prompt_lens = list(past_kv_lens)
-    md.update_spec_dec_param(
-        batch_size=batch_size,
-        is_spec_decoding_enabled=True,
-        is_spec_dec_tree=False,
-        is_spec_dec_dynamic_tree=False,
-        max_draft_len=draft_len,
-        max_total_draft_tokens=draft_len,
-        spec_metadata=None,
-        spec_tree_manager=None,
-        num_contexts=0,
-    )
-    md.kv_cache_params = KVCacheParams(
-        use_cache=True,
-        num_cached_tokens_per_seq=list(past_kv_lens),
-        num_extra_kv_tokens=draft_len - 1,
-    )
-    md.prepare()
-    # Set by the runner's input preparation (no chunked-context rows in a decode step).
-    md.num_chunked_ctx_requests = 0
-    return md, base
-
-
-def make_runner_stub(model, num_tokens: int, batch_size: int):
-    """The attributes ``_preprocess_inputs`` and ``target_prologue`` read from a runner."""
-    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS
-    from tensorrt_llm.mapping import Mapping
-
-    return SimpleNamespace(
-        _config=SimpleNamespace(disable_overlap_scheduler=False),
-        previous_pos_id_offsets_cuda=torch.zeros(
-            (MAX_NUM_TOKENS,), dtype=torch.int32, device="cuda"
-        ),
-        previous_kv_lens_offsets_cuda=torch.zeros(
-            (max(batch_size, 1),), dtype=torch.int32, device="cuda"
-        ),
-        mapping=Mapping(world_size=1, tp_size=1, rank=0),
-        guided_decoder=None,
-        model=model,
-        runtime_hooks=DEFAULT_RUNTIME_HOOKS,
-    )
-
-
-def make_target_model(weight: torch.Tensor):
-    """A ``DeepseekV3Model`` shell holding only the native replicated ``Embedding``."""
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model
-    from tensorrt_llm._torch.modules.embedding import Embedding
-
-    embed = Embedding(weight.shape[0], weight.shape[1], dtype=weight.dtype).cuda()
-    embed.weight.data = weight
-    target = DeepseekV3Model.__new__(DeepseekV3Model)
-    torch.nn.Module.__init__(target)
-    target.embed_tokens = embed
-    return SimpleNamespace(model=target), embed
-
-
-def embedding_weight(synthetic: bool = False) -> torch.Tensor:
-    """The GLM-5.2 embedding table (synthetic only when asked or unavailable)."""
-    if not synthetic:
-        weight = load_checkpoint_tensor("model.embed_tokens.weight")
-        if weight is not None:
-            assert weight.shape == (VOCAB, HIDDEN) and weight.dtype == torch.bfloat16
-            return weight
-    gen = torch.Generator(device="cuda").manual_seed(12)
-    weight = torch.randn((VOCAB, HIDDEN), dtype=torch.bfloat16, device="cuda", generator=gen)
-    weight[3, :4] = torch.tensor([-0.0, float("inf"), float("-inf"), float("nan")])
-    return weight
-
-
-class Arm:
-    """One arm (native or decode_prologue): metadata, runner stub, static inputs."""
-
-    def __init__(self, name, past_kv_lens, weight, cache_manager, mapping, sparse_config):
-        self.name = name
-        self.batch_size = len(past_kv_lens)
-        ids = list(range(self.batch_size))
-        self.md, self._base = make_graph_metadata(
-            cache_manager, mapping, sparse_config, ids, past_kv_lens
-        )
-        self.num_tokens = self.md.num_tokens
-        self.model, self.embed = make_target_model(weight)
-        self.runner = make_runner_stub(self.model, self.num_tokens, self.batch_size)
-        self.input_ids = torch.zeros((MAX_NUM_TOKENS,), dtype=torch.int32, device="cuda")
-        self.position_ids = torch.zeros((1, MAX_NUM_TOKENS), dtype=torch.int32, device="cuda")
-        self.out = None
-
-    def inputs(self) -> dict:
-        n = self.num_tokens
-        return {
-            "attn_metadata": self.md,
-            "input_ids": self.input_ids[:n],
-            "position_ids": self.position_ids[:, :n],
-        }
-
-    def set_step(self, kv_lens, positions, ids, kv_off, pos_off) -> None:
-        """Install one step's inputs (eager copies outside any graph)."""
-        b, n = self.batch_size, self.num_tokens
-        self.md.kv_lens_cuda[:b].copy_(kv_lens)
-        self.position_ids[0, :n].copy_(positions)
-        self.input_ids[:n].copy_(ids)
-        self.runner.previous_kv_lens_offsets_cuda[:b].copy_(kv_off)
-        self.runner.previous_pos_id_offsets_cuda[:n].copy_(pos_off)
-
-    def poison(self) -> None:
-        for name in DEST_FIELDS:
-            t = getattr(self.md, name, None)
-            if isinstance(t, torch.Tensor):
-                t.view(torch.uint8).fill_(POISON) if t.is_contiguous() else t.fill_(-1)
-
-
-def run_native(
-    arm: Arm, inputs: Optional[dict] = None, enable_spec_decode: bool = True
-) -> torch.Tensor:
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
-
-    inputs = arm.inputs() if inputs is None else inputs
-    DecoderRunner._preprocess_inputs(
-        arm.runner, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=MTP_K
-    )
-    return arm.embed(inputs["input_ids"])
-
-
-def run_decode_prologue(
-    arm: Arm, inputs: Optional[dict] = None, enable_spec_decode: bool = True
-) -> torch.Tensor:
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E501
-        prologue,
-    )
-
-    inputs = arm.inputs() if inputs is None else inputs
-    done = prologue.target_prologue(
-        arm.runner, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=MTP_K
-    )
-    assert done, prologue.prologue_ineligibility(arm.runner, inputs, enable_spec_decode)
-    rows = prologue.take_embedding(inputs["attn_metadata"], inputs["input_ids"], arm.embed)
-    assert rows is not None
-    return rows
-
-
-def state(arm: Arm) -> Dict[str, object]:
-    """Every tensor (cloned) and scalar attribute of the metadata, plus the step tensors."""
-    out = {}
-    for name, value in vars(arm.md).items():
-        if isinstance(value, torch.Tensor):
-            out[name] = value.detach().clone()
-        elif isinstance(value, (bool, int, float, str, type(None))):
-            out[name] = value
-    out["@position_ids"] = arm.position_ids.clone()
-    out["@input_ids"] = arm.input_ids.clone()
-    out["@kv_off"] = arm.runner.previous_kv_lens_offsets_cuda.clone()
-    out["@pos_off"] = arm.runner.previous_pos_id_offsets_cuda.clone()
-    return out
-
-
-def _bits(t: torch.Tensor) -> torch.Tensor:
-    if t.dtype in (torch.bfloat16, torch.float16):
-        return t.view(torch.int16)
-    if t.dtype == torch.float32:
-        return t.view(torch.int32)
-    if t.dtype == torch.float64:
-        return t.view(torch.int64)
-    return t
-
-
-def _same(a, b) -> bool:
-    if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
-        if a.shape != b.shape or a.dtype != b.dtype:
-            return False
-        return bool(torch.equal(_bits(a).cpu(), _bits(b).cpu()))
-    return type(a) is type(b) and a == b
-
-
-def changed(before: dict, after: dict) -> set:
-    return {
-        k
-        for k in set(before) | set(after)
-        if not _same(before.get(k, "<missing>"), after.get(k, "<missing>"))
-    }
-
-
-def compare(native_before, native_after, z_before, z_after) -> List[str]:
-    """Names whose bytes (or host values) differ between the arms after the run.
-
-    Compared: every key either arm changed, every destination buffer in full,
-    and the step tensors.
-    """
-    keys = changed(native_before, native_after) | changed(z_before, z_after)
-    keys |= {k for k in DEST_FIELDS if k in native_after}
-    keys |= {"@position_ids", "@input_ids", "@kv_off", "@pos_off", "kv_lens_cuda"}
-    bad = []
-    for k in sorted(keys):
-        if not _same(native_after.get(k, "<missing>"), z_after.get(k, "<missing>")):
-            bad.append(k)
-    return bad
-
-
-def step_inputs(
-    gen: torch.Generator,
-    kv_base: Sequence[int],
-    accepted: Sequence[int],
-    num_tokens: int,
-    draft_len: int = MTP_K,
-    id_edges: bool = False,
-):
-    """One step's (kv_lens, positions, input ids, kv offsets, position offsets).
-
-    ``kv_base[r]`` is request r's KV length after the correction; the eager
-    preparation wrote ``kv_base[r] - offset`` (the overlap scheduler assumed
-    every draft accepted), and the correction adds the offset
-    ``accepted[r] - draft_len`` back. Positions follow the same rule per token.
-    """
-    b = len(kv_base)
-    kv_off = torch.tensor([a - draft_len for a in accepted], dtype=torch.int32)
-    kv_lens = torch.tensor(kv_base, dtype=torch.int32) - kv_off
-    per_req = num_tokens // b
-    pos = []
-    pos_off = []
-    for r in range(b):
-        start = kv_base[r] - per_req
-        for t in range(per_req):
-            pos.append(start + t - (accepted[r] + 1))
-            pos_off.append(accepted[r] + 1)
-    ids = torch.randint(0, VOCAB, (num_tokens,), generator=gen, dtype=torch.int32)
-    if id_edges:
-        edge = torch.tensor([-1, VOCAB, VOCAB + 7, 0, VOCAB - 1, 3], dtype=torch.int32)
-        ids[: min(num_tokens, edge.numel())] = edge[: min(num_tokens, edge.numel())]
-    return (
-        kv_lens.cuda(),
-        torch.tensor(pos, dtype=torch.int32).cuda(),
-        ids.cuda(),
-        kv_off.cuda(),
-        torch.tensor(pos_off, dtype=torch.int32).cuda(),
-    )
-
-
-def make_pair(
-    past_kv_lens, weight, max_kv: Optional[int] = None, sparse_update: Optional[dict] = None
-):
-    """A cache manager and two identically prepared arms on it.
-
-    ``sparse_update`` overrides fields of the arms' DSA configuration (the
-    metadata's parameters), e.g. the paged-MQA logits kernel choice.
-    """
-    batch_size = len(past_kv_lens)
-    cache_manager, mapping, sparse = make_dsa_cache_manager(
-        batch_size, max_kv or max(past_kv_lens) + 4096
-    )
-    if sparse_update:
-        sparse = sparse.model_copy(update=sparse_update)
-    add_gen_requests(cache_manager, list(range(batch_size)), past_kv_lens)
-    native = Arm("native", past_kv_lens, weight, cache_manager, mapping, sparse)
-    decode_prologue = Arm("decode_prologue", past_kv_lens, weight, cache_manager, mapping, sparse)
-    return cache_manager, native, decode_prologue
-
-
-def capture(arm: Arm, fn, debug_dot: Optional[str] = None):
-    """Warm up ``fn(arm)`` eagerly once, then capture it; returns (graph, output)."""
-    fn(arm)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    if debug_dot:
-        graph.enable_debug_mode()
-    with torch.cuda.graph(graph):
-        out = fn(arm)
-    if debug_dot:
-        graph.debug_dump(debug_dot)
-    return graph, out
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py
deleted file mode 100644
index b4218b33aa75..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/decode_prologue/test_decode_prologue.py
+++ /dev/null
@@ -1,421 +0,0 @@
-# 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.
-"""decode_prologue target prologue versus the native prologue, byte for byte (TP1 oracle).
-
-Native: the literal ``DecoderRunner._preprocess_inputs`` (two
-``on_update_kv_lens()`` calls around the overlap scheduler's accepted-count
-correction) and the target model's native ``Embedding``. Candidate: decode_prologue's one
-launch. Both run on identically prepared CUDA-graph DSA metadata of a real
-``DSACacheManager`` (served GEN DSA configuration) with the GLM-5.2
-embedding table. Before every run every destination buffer is poisoned
-identically (0xA5) in both arms; afterwards every metadata attribute either
-arm changed, every destination buffer in full, the position ids, the input ids,
-the offsets and the embedding rows must be bitwise equal.
-
-Cases: accepted counts 0-5 (KV offsets -5..0), contexts 2K-1M with lengths
-at k*256-1, k*256 and k*256+1 (DeepGEMM's SPLIT_KV), B in {1, 2}, a padded
-B = 2 graph whose second row is a short dummy, every DSL atom-split factor the
-picker yields (6, 3, 2), the dense top-k (skipped indexer) graph, out-of-range
-input ids, the single-call prologue without the correction (overlap scheduler
-off, or spec decode off for the step), the DeepGEMM paged-MQA path without the
-DSL atom split (factor 1); eager and CUDA graph with changing inputs at every
-replay.
-"""
-
-import os
-import re
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-sys.path.insert(0, str(Path(__file__).resolve().parent))
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-import _decode_prologue_oracle as orc  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.decode_prologue import (  # noqa: E402, E501
-    prologue,
-)
-
-SYNTHETIC = os.environ.get("GLM_MOE_DSA_DECODE_PROLOGUE_SYNTHETIC_EMBEDDING") == "1"
-
-
-@pytest.fixture(scope="module")
-def weight():
-    return orc.embedding_weight(synthetic=SYNTHETIC)
-
-
-def _boundary(k: int, d: int) -> int:
-    return k * 256 + d
-
-
-# (past KV lengths at prepare, list of per-replay post-correction KV lengths)
-CONTEXT_CASES = {
-    # B = 1: atom split factor 6 (2K), 3 (~12.8K), 2 (>= 32K); SPLIT_KV edges.
-    "b1_2k": ([2100], [[_boundary(9, d)] for d in (-1, 0, 1)]),
-    "b1_12k": ([12700], [[_boundary(50, d)] for d in (-1, 0, 1)]),
-    "b1_32k": ([32768], [[_boundary(128, d)] for d in (-1, 0, 1, 7)]),
-    "b1_131k": ([131072], [[_boundary(512, d)] for d in (-1, 0, 1)]),
-    "b1_491k": (
-        [491061],
-        [[491067], [_boundary(1918, 1)], [_boundary(1919, 0)], [_boundary(1919, -1)]],
-    ),
-    "b1_1m": ([1048000], [[_boundary(4094, d)] for d in (-1, 0, 1)]),
-    # B = 2 live rows (different contexts).
-    "b2_mixed": ([491061, 32768], [[491067, 32774], [_boundary(1919, 1), _boundary(128, -1)]]),
-    # Padded: a B = 2 graph whose second row is a short dummy request.
-    "b2_padded": ([131072, 1], [[131078, 7], [_boundary(512, 0), 7]]),
-}
-
-ACCEPTED_SWEEP = [0, 1, 2, 3, 4, 5]
-
-
-def _run_eager_pair(native, decode_prologue, step):
-    for arm in (native, decode_prologue):
-        arm.set_step(*step)
-        arm.poison()
-    torch.cuda.synchronize()
-    nb, zb = orc.state(native), orc.state(decode_prologue)
-    out_n = orc.run_native(native)
-    out_z = orc.run_decode_prologue(decode_prologue)
-    torch.cuda.synchronize()
-    na, za = orc.state(native), orc.state(decode_prologue)
-    bad = orc.compare(nb, na, zb, za)
-    assert out_n.shape == out_z.shape and out_n.dtype == out_z.dtype
-    assert torch.equal(orc._bits(out_n), orc._bits(out_z)), "embedding rows differ"
-    return bad
-
-
-@pytest.mark.parametrize("case", sorted(CONTEXT_CASES))
-def test_eager_matches_native(case, weight):
-    past, kv_steps = CONTEXT_CASES[case]
-    cm, native, decode_prologue = orc.make_pair(past, weight)
-    try:
-        md = decode_prologue.md
-        reason = prologue.prologue_ineligibility(
-            decode_prologue.runner, decode_prologue.inputs(), True
-        )
-        assert reason is None, reason
-        print(
-            f"[decode_prologue] {case}: B={md.num_seqs} N={md.num_tokens} expand_for_dsl={md.expand_for_dsl} "
-            f"factor={md.dsl_expand_factor} width={md.gen_token_stride} "
-            f"dense={md.skip_indexer_for_gen_reqs} num_sms={md.num_sms}"
-        )
-        gen = torch.Generator().manual_seed(len(case))
-        trials = 0
-        for kv_base in kv_steps:
-            for a in ACCEPTED_SWEEP:
-                accepted = [a] + [(a + 3) % 6 for _ in kv_base[1:]]
-                if case == "b2_padded":
-                    accepted[1] = 5  # a padding row has offset 0
-                step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=a == 0)
-                bad = _run_eager_pair(native, decode_prologue, step)
-                assert not bad, (case, kv_base, accepted, bad)
-                trials += 1
-        print(f"[decode_prologue] {case}: {trials} eager trials identical")
-    finally:
-        cm.shutdown()
-
-
-@pytest.mark.parametrize("case", ["b1_491k", "b1_32k", "b1_2k", "b2_mixed", "b2_padded"])
-def test_cuda_graph_replays_match_native(case, weight):
-    past, kv_steps = CONTEXT_CASES[case]
-    cm, native, decode_prologue = orc.make_pair(past, weight)
-    try:
-        g_n, out_n = orc.capture(native, orc.run_native)
-        g_z, out_z = orc.capture(decode_prologue, orc.run_decode_prologue)
-        gen = torch.Generator().manual_seed(7)
-        replays = 0
-        for rep in range(4):
-            for kv_base in kv_steps:
-                accepted = [int(x) for x in torch.randint(0, 6, (len(kv_base),), generator=gen)]
-                if case == "b2_padded":
-                    accepted[1] = 5
-                step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=rep == 1)
-                for arm in (native, decode_prologue):
-                    arm.set_step(*step)
-                    arm.poison()
-                torch.cuda.synchronize()
-                nb, zb = orc.state(native), orc.state(decode_prologue)
-                g_n.replay()
-                g_z.replay()
-                torch.cuda.synchronize()
-                bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
-                assert not bad, (case, rep, kv_base, accepted, bad)
-                assert torch.equal(orc._bits(out_n), orc._bits(out_z)), (case, rep, "embedding")
-                replays += 1
-        assert replays >= 8
-        print(f"[decode_prologue] {case}: {replays} graph replays identical")
-    finally:
-        cm.shutdown()
-
-
-def test_dense_topk_graph_matches_native(weight):
-    """Short context: the indexer is skipped and the dense top-k rows are written."""
-    cm, native, decode_prologue = orc.make_pair([1000], weight, max_kv=4096)
-    try:
-        assert decode_prologue.md.skip_indexer_for_gen_reqs
-        reason = prologue.prologue_ineligibility(
-            decode_prologue.runner, decode_prologue.inputs(), True
-        )
-        assert reason is None, reason
-        g_n, out_n = orc.capture(native, orc.run_native)
-        g_z, out_z = orc.capture(decode_prologue, orc.run_decode_prologue)
-        gen = torch.Generator().manual_seed(3)
-        for kv_base in ([1006], [1001], [2040], [7]):
-            for a in ACCEPTED_SWEEP:
-                step = orc.step_inputs(gen, kv_base, [a], native.num_tokens)
-                for arm in (native, decode_prologue):
-                    arm.set_step(*step)
-                    arm.poison()
-                torch.cuda.synchronize()
-                nb, zb = orc.state(native), orc.state(decode_prologue)
-                g_n.replay()
-                g_z.replay()
-                torch.cuda.synchronize()
-                bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
-                assert not bad, (kv_base, a, bad)
-                assert torch.equal(orc._bits(out_n), orc._bits(out_z))
-    finally:
-        cm.shutdown()
-
-
-def _eager_and_graph(
-    native, decode_prologue, kv_steps, run_n, run_z, label, sweep=(0, 2, 5), reps=2
-):
-    """Eager trials over ``sweep`` accepted counts, then CUDA-graph replays; returns the counts."""
-    gen = torch.Generator().manual_seed(11)
-    trials = 0
-    for kv_base in kv_steps:
-        for a in sweep:
-            accepted = [a] + [(a + 3) % 6 for _ in kv_base[1:]]
-            step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=a == 0)
-            for arm in (native, decode_prologue):
-                arm.set_step(*step)
-                arm.poison()
-            torch.cuda.synchronize()
-            nb, zb = orc.state(native), orc.state(decode_prologue)
-            out_n, out_z = run_n(native), run_z(decode_prologue)
-            torch.cuda.synchronize()
-            bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
-            assert not bad, (label, kv_base, accepted, bad)
-            assert torch.equal(orc._bits(out_n), orc._bits(out_z)), (label, "embedding")
-            trials += 1
-
-    g_n, out_n = orc.capture(native, run_n)
-    g_z, out_z = orc.capture(decode_prologue, run_z)
-    replays = 0
-    for rep in range(reps):
-        for kv_base in kv_steps:
-            accepted = [int(x) for x in torch.randint(0, 6, (len(kv_base),), generator=gen)]
-            step = orc.step_inputs(gen, kv_base, accepted, native.num_tokens, id_edges=rep == 1)
-            for arm in (native, decode_prologue):
-                arm.set_step(*step)
-                arm.poison()
-            torch.cuda.synchronize()
-            nb, zb = orc.state(native), orc.state(decode_prologue)
-            g_n.replay()
-            g_z.replay()
-            torch.cuda.synchronize()
-            bad = orc.compare(nb, orc.state(native), zb, orc.state(decode_prologue))
-            assert not bad, (label, rep, kv_base, accepted, bad)
-            assert torch.equal(orc._bits(out_n), orc._bits(out_z)), (label, rep, "embedding")
-            replays += 1
-    return trials, replays
-
-
-@pytest.mark.parametrize("mode", ["no_overlap", "no_spec_decode"])
-@pytest.mark.parametrize("case", ["b1_491k", "b1_2k", "b2_mixed"])
-def test_without_correction_matches_native(case, mode, weight):
-    """Without the overlap correction the native prologue is one ``on_update_kv_lens()``.
-
-    ``no_overlap``: spec decode with the overlap scheduler disabled;
-    ``no_spec_decode``: the step runs with ``enable_spec_decode=False``.
-    """
-    past, kv_steps = CONTEXT_CASES[case]
-    cm, native, decode_prologue = orc.make_pair(past, weight)
-    try:
-        spec = mode != "no_spec_decode"
-        if mode == "no_overlap":
-            for arm in (native, decode_prologue):
-                arm.runner._config.disable_overlap_scheduler = True
-        reason = prologue.prologue_ineligibility(
-            decode_prologue.runner, decode_prologue.inputs(), spec
-        )
-        assert reason is None, reason
-
-        def run_n(arm):
-            return orc.run_native(arm, enable_spec_decode=spec)
-
-        def run_z(arm):
-            return orc.run_decode_prologue(arm, enable_spec_decode=spec)
-
-        trials, replays = _eager_and_graph(
-            native, decode_prologue, kv_steps, run_n, run_z, (case, mode)
-        )
-        print(
-            f"[decode_prologue] {case} {mode}: {trials} eager trials and {replays} graph replays identical"
-        )
-    finally:
-        cm.shutdown()
-
-
-@pytest.mark.parametrize("case", ["b1_491k", "b1_2k", "b2_mixed"])
-def test_without_dsl_expansion_matches_native(case, weight):
-    """DeepGEMM paged-MQA logits without the DSL atom split: no expanded schedule (factor 1)."""
-    past, kv_steps = CONTEXT_CASES[case]
-    cm, native, decode_prologue = orc.make_pair(
-        past, weight, sparse_update={"use_cute_dsl_paged_mqa_logits": False}
-    )
-    try:
-        md = decode_prologue.md
-        assert not md.expand_for_dsl and not md.use_expanded_buffers_for_mtp
-        reason = prologue.prologue_ineligibility(
-            decode_prologue.runner, decode_prologue.inputs(), True
-        )
-        assert reason is None, reason
-        trials, replays = _eager_and_graph(
-            native,
-            decode_prologue,
-            kv_steps,
-            orc.run_native,
-            orc.run_decode_prologue,
-            (case, "no_dsl"),
-            sweep=tuple(ACCEPTED_SWEEP),
-        )
-        print(
-            f"[decode_prologue] {case} no_dsl: {trials} eager trials and {replays} graph replays identical"
-        )
-    finally:
-        cm.shutdown()
-
-
-class _HelixMapping:
-    @staticmethod
-    def has_cp_helix():
-        return True
-
-
-def test_declines_leave_native_untouched(weight):
-    """Out-of-domain steps: target_prologue returns False and writes nothing."""
-    cm, native, decode_prologue = orc.make_pair([32768], weight)
-    try:
-        md = decode_prologue.md
-        inputs = decode_prologue.inputs()
-        mrope = dict(inputs)
-        mrope["position_ids"] = torch.zeros((3, 1, md.num_tokens), dtype=torch.int32, device="cuda")
-        # (expected reason, attribute owner, attribute, value, inputs)
-        cases = [
-            ("expanded MTP buffers", md, "use_expanded_buffers_for_mtp", True, inputs),
-            ("ragged verification", md, "enable_ragged_verification", True, inputs),
-            (
-                "helix context parallelism",
-                decode_prologue.runner,
-                "mapping",
-                _HelixMapping(),
-                inputs,
-            ),
-            ("not a generation-only step", md, "_num_contexts", 1, inputs),
-            ("position ids are not a [1, N] int32 row", None, None, None, mrope),
-        ]
-        for expect, owner, name, value, case_inputs in cases:
-            saved = getattr(owner, name) if owner is not None else None
-            if owner is not None:
-                setattr(owner, name, value)
-            try:
-                reason = prologue.prologue_ineligibility(decode_prologue.runner, case_inputs, True)
-                before = orc.state(decode_prologue)
-                done = prologue.target_prologue(
-                    decode_prologue.runner,
-                    case_inputs,
-                    enable_spec_decode=True,
-                    runtime_draft_len=orc.MTP_K,
-                )
-                torch.cuda.synchronize()
-                after = orc.state(decode_prologue)
-            finally:
-                if owner is not None:
-                    setattr(owner, name, saved)
-            assert reason == expect, (expect, reason)
-            assert not done, expect
-            assert not orc.changed(before, after), (expect, orc.changed(before, after))
-            print(f"[decode_prologue] decline '{expect}' ok")
-    finally:
-        cm.shutdown()
-
-
-def test_take_embedding_requires_the_same_ids(weight):
-    cm, native, decode_prologue = orc.make_pair([32768], weight)
-    try:
-        inputs = decode_prologue.inputs()
-        assert prologue.target_prologue(
-            decode_prologue.runner, inputs, enable_spec_decode=True, runtime_draft_len=5
-        )
-        other_ids = inputs["input_ids"].clone()
-        assert prologue.take_embedding(decode_prologue.md, other_ids, decode_prologue.embed) is None
-        # The packet is consumed either way.
-        assert (
-            prologue.take_embedding(decode_prologue.md, inputs["input_ids"], decode_prologue.embed)
-            is None
-        )
-    finally:
-        cm.shutdown()
-
-
-_KINDS = ("KERNEL", "MEMCPY", "MEMSET", "EMPTY", "HOST", "EVENT", "MEM_ALLOC", "MEM_FREE")
-
-
-def _dot_nodes(path):
-    """(kind, label text) per node of a ``cudaGraphDebugDotPrint`` file, in file order."""
-    text = Path(path).read_text(errors="replace")
-    nodes = []
-    for label in re.findall(r'"graph_\d+_node_\d+"\[[^\]]*?label="(.*?)"\];', text, flags=re.S):
-        flat = re.sub(r"\\[<>]|[{}|]|\s+", " ", label).strip()
-        kind = next((k for k in _KINDS if flat.startswith(k)), "OTHER")
-        width = re.search(r"Width\s+(\d+)", flat)
-        name = re.search(r"topoId:\s*\d+\)\s+(\S+)", flat)
-        if kind == "MEMCPY" and width:
-            flat = f"MEMCPY {width.group(1)} B"
-        elif kind == "KERNEL" and name:
-            flat = name.group(1)
-        nodes.append((kind, flat[:160]))
-    return nodes
-
-
-def test_native_chain_census(weight, tmp_path):
-    """Print the captured native prologue's nodes (memcpy sizes included) and decode_prologue's."""
-    cm, native, decode_prologue = orc.make_pair([491061], weight)
-    try:
-        out_dir = Path(os.environ.get("GLM_MOE_DSA_DECODE_PROLOGUE_DOT_DIR", str(tmp_path)))
-        out_dir.mkdir(parents=True, exist_ok=True)
-        dot_n = str(out_dir / "native_prologue.dot")
-        dot_z = str(out_dir / "decode_prologue.dot")
-        orc.capture(native, orc.run_native, debug_dot=dot_n)
-        orc.capture(decode_prologue, orc.run_decode_prologue, debug_dot=dot_z)
-        for tag, path in (("native", dot_n), ("decode_prologue", dot_z)):
-            nodes = _dot_nodes(path)
-            kinds = [k for k, _ in nodes]
-            print(
-                f"[decode_prologue] census {tag}: {len(nodes)} nodes: "
-                f"{kinds.count('KERNEL')} kernels, {kinds.count('MEMCPY')} memcpy, "
-                f"{kinds.count('MEMSET')} memset; dot {path}"
-            )
-            for i, (kind, detail) in enumerate(nodes):
-                print(f"[decode_prologue]   {tag} {i:2d} {kind} {detail}")
-    finally:
-        cm.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py
deleted file mode 100644
index 90cf8c0922cc..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/input_prep_testlib.py
+++ /dev/null
@@ -1,230 +0,0 @@
-# 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.
-"""Helpers for the CPU tests of deferred_input_prep's host side (stager and loop domain).
-
-``load()`` returns deferred_input_prep's ``dispatch``, ``compat``, ``packet`` and ``stager``
-modules. Without a built ``tensorrt_llm`` it imports them from the source tree
-under a stand-in package, so ``PrepStager`` runs with CPU tensors. Logging and
-``prefer_pinned`` then need ``tensorrt_llm``: tests replace
-``dispatch._log_info`` and ``stager._host_memory_pinned``.
-
-``FakePacket`` replaces ``PrepPacket`` for CPU tensors. It performs the
-staged writes at its ``commit``, through NumPy views of the destinations.
-"""
-
-import importlib
-import importlib.util
-import sys
-import types
-from pathlib import Path
-from types import SimpleNamespace
-
-REPO = Path(__file__).resolve().parents[9]
-MODELING_V2 = REPO / "tensorrt_llm" / "_torch" / "_experimental" / "modeling_v2"
-TARGET = MODELING_V2 / "models" / "glm_moe_dsa" / "shared"
-PACKAGES = TARGET / "kernels"
-_STANDALONE = "_input_prep_testlib"
-
-
-def _module_from(name, path):
-    spec = importlib.util.spec_from_file_location(name, path)
-    module = importlib.util.module_from_spec(spec)
-    sys.modules[name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-def load_dispatch_standalone(name):
-    """The target's ``dispatch.py`` as module ``name``, without importing ``tensorrt_llm``.
-
-    Its absolute imports, ``tensorrt_llm.logger`` and the modeling_v2 router index, use
-    only the standard library: when they are not imported yet, they are loaded from their
-    source files and registered under their own names only while ``dispatch.py`` runs.
-    """
-    deps = (
-        ("tensorrt_llm.logger", REPO / "tensorrt_llm" / "logger.py"),
-        (
-            "tensorrt_llm._torch._experimental.modeling_v2._router_index",
-            MODELING_V2 / "_router_index.py",
-        ),
-    )
-    added = []
-    try:
-        for dep, path in deps:
-            if dep not in sys.modules:
-                added.append(dep)
-                _module_from(dep, path)
-        return _module_from(name, TARGET / "dispatch.py")
-    finally:
-        for dep in added:
-            sys.modules.pop(dep, None)
-
-
-def _standalone():
-    if _STANDALONE not in sys.modules:
-        for name, path in (
-            (_STANDALONE, TARGET),
-            (f"{_STANDALONE}.kernels", PACKAGES),
-            (f"{_STANDALONE}.kernels.deferred_input_prep", PACKAGES / "deferred_input_prep"),
-        ):
-            package = types.ModuleType(name)
-            package.__path__ = [str(path)]
-            sys.modules[name] = package
-        sys.modules[_STANDALONE].dispatch = load_dispatch_standalone(f"{_STANDALONE}.dispatch")
-    return SimpleNamespace(
-        dispatch=importlib.import_module(f"{_STANDALONE}.dispatch"),
-        compat=importlib.import_module(f"{_STANDALONE}.kernels.deferred_input_prep.compat"),
-        packet=importlib.import_module(f"{_STANDALONE}.kernels.deferred_input_prep.packet"),
-        stager=importlib.import_module(f"{_STANDALONE}.kernels.deferred_input_prep.stager"),
-    )
-
-
-def load():
-    """deferred_input_prep's host modules: the installed package, else the source tree standalone."""
-    try:
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-            compat,
-            packet,
-            stager,
-        )
-    except (ImportError, OSError):
-        return _standalone()
-    return SimpleNamespace(dispatch=dispatch, compat=compat, packet=packet, stager=stager)
-
-
-class FakePacket:
-    """``PrepPacket`` stand-in for CPU tensors (see the module docstring).
-
-    ``max_desc`` descriptors fit in one packet (one per write); a write whose
-    bytes overlap an earlier write of the packet raises ``PacketOverlap``.
-    """
-
-    def __init__(self, packet_module, max_desc=1 << 20):
-        self._full = packet_module.PacketFull
-        self._overlap = packet_module.PacketOverlap
-        self.max_desc = max_desc
-        self._ops = []
-        self._ranges = []
-        self.num_commits = 0
-        self.num_writes = 0
-
-    def begin(self):
-        self._ops, self._ranges = [], []
-
-    def abort(self):
-        self._ops, self._ranges = [], []
-
-    @property
-    def num_descriptors(self):
-        return len(self._ops)
-
-    def _claim(self, *dsts):
-        if len(self._ops) + 1 > self.max_desc:
-            raise self._full("fake packet is full")
-        new = [(t.data_ptr(), t.data_ptr() + t.numel() * t.element_size()) for t in dsts]
-        for lo, hi in new:
-            for a, b in self._ranges:
-                if lo < b and a < hi:
-                    raise self._overlap("overlapping destinations in one packet")
-        self._ranges.extend(new)
-
-    def _add(self, fn, *dsts):
-        self._claim(*dsts)
-        self._ops.append(fn)
-
-    def add_copy(self, dst, src, clamp_min=None):
-        def op():
-            values = src.reshape(dst.shape).numpy()
-            dst.numpy()[...] = values if clamp_min is None else values.clip(min=clamp_min)
-
-        self._add(op, dst)
-
-    def add_values(self, dst, values):
-        def op():
-            dst.numpy()[...] = list(values)
-
-        self._add(op, dst)
-
-    def add_gather(self, dst, src, indices, addend=0):
-        def op():
-            flat = src.reshape(-1).numpy()
-            out = dst.numpy()
-            out[...] = 0
-            out[: len(indices)] = flat[list(indices)] + addend
-
-        self._add(op, dst)
-
-    def add_zero(self, dst):
-        def op():
-            dst.numpy()[...] = 0
-
-        self._add(op, dst)
-
-    def add_clamp_min(self, dst, lo):
-        def op():
-            out = dst.numpy()
-            out[...] = out.clip(min=lo)
-
-        self._add(op, dst)
-
-    def commit(self, stream=None, record_event=True):
-        ops, self._ops, self._ranges = self._ops, [], []
-        for op in ops:
-            op()
-        self.num_commits += 1
-        self.num_writes += len(ops)
-        return len(ops)
-
-    def set_completion_event(self, event):
-        pass
-
-
-def eligibility_case(**overrides):
-    """Arguments of ``PrepStager.begin_if_eligible`` for a steady served decode step."""
-    mode = SimpleNamespace(use_one_engine=lambda: True)
-    cfg = SimpleNamespace(
-        spec_config=SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True),
-        disable_overlap_scheduler=overrides.pop("disable_overlap_scheduler", False),
-        enable_attention_dp=overrides.pop("enable_attention_dp", False),
-        use_mrope=False,
-    )
-    runner = SimpleNamespace(
-        _config=cfg,
-        use_beam_search=False,
-        mapping=SimpleNamespace(has_cp_helix=lambda: False, pp_size=overrides.pop("pp_size", 1)),
-        guided_decoder=None,
-        _lora=SimpleNamespace(cuda_graph_manager=None),
-    )
-    rows = overrides.pop("rows", [SimpleNamespace(is_dummy=False, py_batch_idx=0)])
-    requests = SimpleNamespace(num_context_requests=0, generation_requests=rows)
-    DSAtrtllmAttentionMetadata = type("DSAtrtllmAttentionMetadata", (), {})
-    manager = SimpleNamespace(index_mapper=object(), _use_per_layer_page_tables=False)
-    attn = DSAtrtllmAttentionMetadata()
-    attn.kv_cache_manager = manager
-    attn.draft_kv_cache_manager = None
-    args = dict(
-        runner=runner,
-        requests=requests,
-        attn_metadata=attn,
-        spec_metadata=SimpleNamespace(),
-        new_tensors_device=SimpleNamespace(next_draft_tokens=object()),
-        graph_replay=True,
-        is_dummy=False,
-        enable_spec_decode=True,
-        promoted=False,
-    )
-    args.update(overrides)
-    return args
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py
deleted file mode 100644
index a212ca16f519..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_compat.py
+++ /dev/null
@@ -1,340 +0,0 @@
-# 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.
-"""deferred_input_prep's executor-loop domain, its engagement rule, its debug level and its window (CPU only).
-
-The step-input stager engages when the target declares it (the plan runs
-``deferred_input_prep``) and the step is in its support domain; the store writer engages
-once a non-attention-DP overlap loop declared its domain. Nothing depends on which other
-packages are selected or on their sources. An engaged preparation opens the deferred-write window, which
-``PrepStager.flush`` can close early for the writes staged so far, also handing ``input_ids`` /
-``position_ids`` back to the graph replay's own copy.
-"""
-
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-import input_prep_testlib  # noqa: E402
-
-HOST = torch.arange(10, 14, dtype=torch.int32)
-
-
-@pytest.fixture
-def deferred_input_prep(monkeypatch):
-    """deferred_input_prep's host modules with logging captured in ``deferred_input_prep.lines``."""
-    modules = input_prep_testlib.load()
-    lines = []
-    monkeypatch.setattr(modules.dispatch, "_log_info", lines.append)
-    monkeypatch.setattr(modules.dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(modules.dispatch, "_INSTALLED", [None])
-    monkeypatch.setattr(
-        modules.compat, "_LOOP_STATE", {"declared": False, "reason": "", "logged": False}
-    )
-    modules.lines = lines
-    return modules
-
-
-@pytest.fixture
-def compat(deferred_input_prep):
-    return deferred_input_prep.compat
-
-
-def test_loop_domain_state(compat):
-    """Store writer: only after a non-ADP overlap loop declared."""
-    state = compat._LOOP_STATE
-    assert compat.loop_domain() is False
-    compat.declare_overlap_loop(False)
-    assert compat.loop_domain() is True
-    compat.declare_overlap_loop(True)
-    assert compat.loop_domain() is False and "attention DP" in state["reason"]
-    compat.declare_overlap_loop(False)
-    assert compat.loop_domain() is False
-    state.update(declared=False, reason="")
-    compat.declare_overlap_loop(None)
-    assert compat.loop_domain() is False and "did not pass" in state["reason"]
-
-
-def test_every_part_engages_without_any_record(deferred_input_prep, monkeypatch):
-    """A shared switch set to a value no oracle ran: deferred_input_prep engages.
-
-    Engagement reads no package source and no table of package versions, so
-    editing another package's sources cannot turn deferred_input_prep off.
-    """
-    monkeypatch.setenv("TRTLLM_ENABLE_PDL", "0")
-    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
-    st = deferred_input_prep.stager.PrepStager()
-    st.packet = input_prep_testlib.FakePacket(deferred_input_prep.packet)
-    case = input_prep_testlib.eligibility_case()
-    attached = (
-        case["runner"],
-        case["attn_metadata"],
-        case["spec_metadata"],
-        case["attn_metadata"].kv_cache_manager,
-    )
-    assert st.begin_if_eligible(**case) is st
-    # The attribute every engine call site reads (RuntimeHooks.step_input_stager).
-    assert all(obj.step_input_stager is st for obj in attached)
-    st.end()
-    assert all(obj.step_input_stager is None for obj in attached)
-    assert st.decisions == {"engaged": 1} and st.num_commits == 1
-    assert deferred_input_prep.lines[0].startswith(
-        "[glm_moe_dsa] deferred_input_prep/prep: engaged"
-    )
-    deferred_input_prep.compat.declare_overlap_loop(False)
-    assert deferred_input_prep.compat.loop_domain()
-
-
-def test_prep_declines_only_for_pageable_host_memory(deferred_input_prep, monkeypatch):
-    """Whether a stager exists is the plan's decision; the stager declines only pageable memory."""
-    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
-    st = deferred_input_prep.stager.PrepStager()
-    st.packet = input_prep_testlib.FakePacket(deferred_input_prep.packet)
-    assert st.begin_if_eligible(**input_prep_testlib.eligibility_case()) is st
-    st.end()
-    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: False)
-    st = deferred_input_prep.stager.PrepStager()
-    assert st.begin_if_eligible(**input_prep_testlib.eligibility_case()) is None
-    assert st.last_reason == "host page tables are not pinned (prefer_pinned)"
-    assert not st.active
-
-
-def test_debug_level_is_the_constructor_argument(deferred_input_prep, monkeypatch):
-    """The stager reads no environment: its debug level is the argument (0, 1 or 2)."""
-    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
-    monkeypatch.setenv("TRTLLM_MODELING_V2_VALIDATE", "1")
-    PrepStager = deferred_input_prep.stager.PrepStager
-    assert PrepStager().debug is None
-    assert PrepStager(debug_level=1).debug.level == 1
-    assert PrepStager(debug_level=2).debug.level == 2
-    with pytest.raises(ValueError, match="0, 1 or 2"):
-        PrepStager(debug_level=3)
-
-
-def _engaged(deferred_input_prep, monkeypatch, packet=None, debug_level=0):
-    """An engaged stager on a CPU case whose runner owns a 16-element ``input_ids_cuda``."""
-    monkeypatch.setattr(deferred_input_prep.stager, "_host_memory_pinned", lambda: True)
-    case = input_prep_testlib.eligibility_case()
-    case["runner"].input_ids_cuda = torch.zeros(16, dtype=torch.int32)
-    st = deferred_input_prep.stager.PrepStager(debug_level=debug_level)
-    st.packet = (
-        packet if packet is not None else input_prep_testlib.FakePacket(deferred_input_prep.packet)
-    )
-    assert st.begin_if_eligible(**case) is st
-    return st, case
-
-
-def test_engaged_line_states_the_window_contract(deferred_input_prep, monkeypatch):
-    st, _ = _engaged(deferred_input_prep, monkeypatch)
-    st.end()
-    assert deferred_input_prep.lines == [
-        "[glm_moe_dsa] deferred_input_prep/prep: engaged (deferred-write window from "
-        "DecoderRunner._forward_decoder's begin_if_eligible to stager.end(): the buffers written "
-        "through the deferred_input_prep stager are written later by the commit; do not read or "
-        "write them in between with other device ops; route writes through the stager API or "
-        "call stager.flush() first; for input_ids_cuda and position_ids_cuda (also staged into "
-        "the graph's static inputs) call stager.flush() and then write natively; "
-        "disable with TRTLLM_STAIRCASE_GLM_LL_DISABLE=deferred_input_prep)"
-    ]
-
-
-def test_flush_then_another_write_keeps_native_order(deferred_input_prep, monkeypatch):
-    """After flush, a write to a staged buffer lands after the staged write, as natively."""
-    st, case = _engaged(deferred_input_prep, monkeypatch)
-    buf = case["runner"].input_ids_cuda
-    st.copy(buf[:4], HOST)
-    st.flush()
-    assert buf[:4].tolist() == HOST.tolist() and st.num_flushes == 1
-    buf[:4].add_(100)
-    st.zero(buf[4:])
-    st.end()
-    assert st.packet.num_commits == 2 and st.num_commits == 1
-    assert buf.tolist() == (HOST + 100).tolist() + [0] * 12
-
-
-@pytest.mark.parametrize("max_desc", [0, 2, 1 << 20])
-def test_flush_issues_several_writes_in_staging_order(deferred_input_prep, monkeypatch, max_desc):
-    """Overlapping writes staged before one flush land in staging order.
-
-    ``max_desc`` 0: every write is native; 2: the packet fills up; otherwise
-    only the overlaps split the packet.
-    """
-    packet = input_prep_testlib.FakePacket(deferred_input_prep.packet, max_desc=max_desc)
-    st, case = _engaged(deferred_input_prep, monkeypatch, packet=packet)
-    buf = case["runner"].input_ids_cuda
-    buf.fill_(5)
-    st.copy(buf[:4], torch.tensor([-3, 1, 7, -9], dtype=torch.int32))
-    st.zero(buf[4:])
-    st.clamp_min(buf[:4], 0)
-    st.values(buf[:2], [-1, -2])
-    st.flush()
-    expected = [-1, -2, 7, 0] + [0] * 12
-    assert buf.tolist() == expected
-    st.end()
-    assert buf.tolist() == expected
-
-
-def test_flush_and_commit_fall_back_to_native_writes(deferred_input_prep, monkeypatch):
-    """With a packet that holds nothing, flush and the commit issue each write natively."""
-    full = input_prep_testlib.FakePacket(deferred_input_prep.packet, max_desc=0)
-    st, case = _engaged(deferred_input_prep, monkeypatch, packet=full)
-    buf = case["runner"].input_ids_cuda
-    st.copy(buf[:4], HOST)
-    st.flush()
-    assert buf[:4].tolist() == HOST.tolist() and st.num_native_copies == 1
-    st.values(buf[4:8], [1, 2, 3, 4])
-    st.zero(buf[8:])
-    st.end()
-    assert st.num_native_copies == 3 and full.num_writes == 0
-    assert buf.tolist() == HOST.tolist() + [1, 2, 3, 4] + [0] * 8
-    st.flush()  # outside a window: nothing to do
-    assert st.num_native_copies == 3
-
-
-def test_flush_ends_the_guarded_span_of_the_flushed_writes(deferred_input_prep, monkeypatch):
-    """Debug guard: a flushed buffer may be touched; a buffer staged later stays guarded."""
-    guard = sys.modules[deferred_input_prep.stager.DebugWindow.__module__]
-
-    def cpu_span(t):
-        lo = t.data_ptr()
-        return (lo, lo + t.numel() * t.element_size()) if t.numel() else None
-
-    monkeypatch.setattr(guard, "tensor_span", cpu_span)
-    st, case = _engaged(deferred_input_prep, monkeypatch, debug_level=1)
-    buf = case["runner"].input_ids_cuda
-    st.copy(buf[:4], HOST)
-    st.flush()
-    buf[:4].add_(100)
-    st.zero(buf[4:])
-    try:
-        with pytest.raises(guard.InputPrepHazardError, match="touches a destination staged"):
-            buf[4:].add_(1)
-    finally:
-        st.discard()
-    assert buf.tolist() == (HOST + 100).tolist() + [0] * 12
-
-
-STEP_TOKENS = 4
-STAGED_INPUT_IDS = [11, 21, 12, 22]
-STAGED_POSITIONS = [7, 8, 9, 10]
-
-
-def _stage_input_ids(st, runner):
-    """``overlap_inputs`` of a steady step: two rows with previous tensors, 2 tokens per row."""
-    new_tokens = torch.tensor([[11, 12, 13, 14], [21, 22, 23, 24]], dtype=torch.int32)
-    next_draft_tokens = torch.tensor([[31], [32], [33], [34]], dtype=torch.int32)
-    new_tokens_lens = torch.tensor([2, 1, 2, 2], dtype=torch.int32)
-    st.overlap_inputs(
-        runner, [0, 1], [0, 1, 2, 3], 0, 0, 2, 2, new_tokens, next_draft_tokens, new_tokens_lens
-    )
-
-
-def _stage_position_ids(st, runner):
-    st.position_ids(
-        runner,
-        runner.position_ids_cuda[:STEP_TOKENS],
-        torch.tensor(STAGED_POSITIONS, dtype=torch.int32),
-    )
-
-
-def _graph_step(deferred_input_prep, monkeypatch):
-    """An engaged step whose runner has the decode graph's static input buffers."""
-    st, case = _engaged(deferred_input_prep, monkeypatch)
-    runner = case["runner"]
-    for name in (
-        "position_ids_cuda",
-        "draft_tokens_cuda",
-        "previous_batch_indices_cuda",
-        "previous_pos_indices_cuda",
-        "previous_pos_id_offsets_cuda",
-        "previous_kv_lens_offsets_cuda",
-    ):
-        setattr(runner, name, torch.zeros(16, dtype=torch.int32))
-    static = {
-        "input_ids": torch.zeros(16, dtype=torch.int32),
-        "position_ids": torch.zeros((1, 16), dtype=torch.int32),
-    }
-    runner.cuda_graph_runner = SimpleNamespace(shared_static_tensors=static)
-    return st, runner, static
-
-
-def _graph_inputs(st, runner, static):
-    """``input_ids`` / ``position_ids`` as ``CUDAGraphRunner.replay`` hands them to the graph."""
-    n = STEP_TOKENS
-    if st.static_inputs_tokens() != n:
-        static["input_ids"][:n].copy_(runner.input_ids_cuda[:n])
-        static["position_ids"][0, :n].copy_(runner.position_ids_cuda[:n])
-    return static["input_ids"][:n].tolist(), static["position_ids"][0, :n].tolist()
-
-
-def test_staged_step_writes_the_graph_inputs(deferred_input_prep, monkeypatch):
-    """The commit writes the static inputs too, so the replay skips its own copy."""
-    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
-    _stage_input_ids(st, runner)
-    _stage_position_ids(st, runner)
-    st.end()
-    assert st.static_inputs_tokens() == STEP_TOKENS
-    assert runner.input_ids_cuda[:STEP_TOKENS].tolist() == STAGED_INPUT_IDS
-    assert _graph_inputs(st, runner, static) == (STAGED_INPUT_IDS, STAGED_POSITIONS)
-
-
-def test_native_writes_after_flush_reach_the_graph_inputs(deferred_input_prep, monkeypatch):
-    """flush, then native writes to input_ids_cuda / position_ids_cuda: the graph sees them."""
-    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
-    _stage_input_ids(st, runner)
-    _stage_position_ids(st, runner)
-    st.flush()
-    runner.input_ids_cuda[:STEP_TOKENS].copy_(torch.tensor([101, 102, 103, 104]))
-    runner.position_ids_cuda[:STEP_TOKENS].add_(100)
-    st.end()
-    assert st.static_inputs_tokens() == 0
-    assert _graph_inputs(st, runner, static) == (
-        [101, 102, 103, 104],
-        [p + 100 for p in STAGED_POSITIONS],
-    )
-
-
-def test_flush_between_input_ids_and_position_ids(deferred_input_prep, monkeypatch):
-    """Only position_ids is staged again after the flush: the replay copies both inputs."""
-    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
-    _stage_input_ids(st, runner)
-    st.flush()
-    runner.input_ids_cuda[:STEP_TOKENS].add_(100)
-    _stage_position_ids(st, runner)
-    st.end()
-    assert st.static_inputs_tokens() == 0
-    assert _graph_inputs(st, runner, static) == (
-        [t + 100 for t in STAGED_INPUT_IDS],
-        STAGED_POSITIONS,
-    )
-
-
-def test_staged_write_to_input_ids_does_not_reach_the_graph(deferred_input_prep, monkeypatch):
-    """Why the window contract sends new input_ids_cuda writes through flush.
-
-    The static inputs staged by ``overlap_inputs`` are what the graph replays;
-    a later staged write to ``input_ids_cuda`` changes only the runner buffer.
-    """
-    st, runner, static = _graph_step(deferred_input_prep, monkeypatch)
-    _stage_input_ids(st, runner)
-    _stage_position_ids(st, runner)
-    st.copy(runner.input_ids_cuda[:STEP_TOKENS], torch.tensor([101, 102, 103, 104]))
-    st.end()
-    assert runner.input_ids_cuda[:STEP_TOKENS].tolist() == [101, 102, 103, 104]
-    assert _graph_inputs(st, runner, static)[0] == STAGED_INPUT_IDS
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py
deleted file mode 100644
index 543ccdd38174..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_input_prep/test_input_prep_kernels.py
+++ /dev/null
@@ -1,872 +0,0 @@
-# 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.
-"""Byte-exact checks of the deferred_input_prep kernels against the native operations they replace.
-
-Also: ordering (last write wins), native fallbacks, delayed-producer stress of
-the pinned rings (GPU held back by a sleep kernel while many launches are
-staged) and the debug guard. Set ``GLM_MOE_DSA_DEFERRED_INPUT_PREP_UNDER_SANITIZER=1`` when running
-under compute-sanitizer: it slows the host enough that the "GPU still busy"
-precondition of the stress tests may not hold, so only that precondition is
-not asserted.
-"""
-
-import os
-from types import SimpleNamespace
-
-import pytest
-import torch
-import triton
-import triton.language as tl
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
-
-UNDER_SANITIZER = os.environ.get("GLM_MOE_DSA_DEFERRED_INPUT_PREP_UNDER_SANITIZER") == "1"
-# About one second of GPU spin on sm_103: far longer than staging the launches.
-HOLD_CYCLES = 2_000_000_000
-
-
-def _stager(debug_level=0):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.stager import (  # noqa: E501
-        PrepStager,
-    )
-
-    return PrepStager(debug_level=debug_level)
-
-
-@triton.jit
-def _copy_tensor_kernel(src, dst, n, BLOCK: tl.constexpr):
-    offs = tl.arange(0, BLOCK)
-    m = offs < n
-    tl.store(dst + offs, tl.load(src + offs, mask=m), mask=m)
-
-
-@triton.jit
-def _copy_addr_kernel(src_addr, dst_addr, n, BLOCK: tl.constexpr):
-    src = src_addr.to(tl.pointer_type(tl.int32))
-    dst = dst_addr.to(tl.pointer_type(tl.int32))
-    offs = tl.arange(0, BLOCK)
-    m = offs < n
-    tl.store(dst + offs, tl.load(src + offs, mask=m), mask=m)
-
-
-def _rand_i32(*shape, lo=-1000, hi=1 << 20):
-    return torch.randint(lo, hi, shape, dtype=torch.int32)
-
-
-def test_copy_values_gather_and_zero_tail():
-    torch.manual_seed(0)
-    dev = "cuda"
-    a = _rand_i32(37).to(dev)
-    b = torch.zeros(5000, dtype=torch.int32, device=dev)  # multi-chunk copy
-    c = _rand_i32(128).to(dev)  # gather target with zero tail (previous_pos_id_offsets)
-    d = torch.zeros(4, dtype=torch.int64, device=dev)  # int64 copy
-    f = torch.zeros(3, dtype=torch.float32, device=dev)
-    src = _rand_i32(7, 1, 3).to(dev)  # sampler store [W, S, B]
-    a_ref, b_ref, c_ref, d_ref, f_ref = a.clone(), b.clone(), c.clone(), d.clone(), f.clone()
-
-    host_b = _rand_i32(5000).pin_memory()
-    host_d = torch.tensor([1, -2, 1 << 40, 7], dtype=torch.int64).pin_memory()
-    host_f = torch.tensor([0.5, -1.25, 3.0], dtype=torch.float32)
-    values = list(range(10, 47))
-    idx = [5 * 3 + 2, 2 * 3 + 2, 0]
-
-    # Native reference.
-    a_ref.copy_(torch.tensor(values, dtype=torch.int32), non_blocking=True)
-    b_ref.copy_(host_b, non_blocking=True)
-    c_ref *= 0
-    c_ref[10:13] = src.flatten()[torch.tensor(idx, device=dev)] - 6
-    d_ref.copy_(host_d, non_blocking=True)
-    f_ref.copy_(host_f, non_blocking=True)
-
-    st = _stager()
-    st.begin()
-    st.values(a, values)
-    st.copy(b, host_b)
-    st.gather(c[:10], src, [], 0)
-    st.gather(c[10:13], src, idx, -6)
-    st.gather(c[13:], src, [], 0)
-    st.copy(d, host_d)
-    st.copy(f, host_f)
-    st.commit()
-    torch.cuda.synchronize()
-    for x, y in ((a, a_ref), (b, b_ref), (c, c_ref), (d, d_ref), (f, f_ref)):
-        assert torch.equal(x.view(torch.uint8), y.view(torch.uint8))
-
-
-def test_host_bytes_read_at_commit_time():
-    dev = "cuda"
-    dst = torch.zeros(8, dtype=torch.int32, device=dev)
-    host = torch.zeros(8, dtype=torch.int32).pin_memory()
-    st = _stager()
-    st.begin()
-    st.copy(dst, host)
-    host.fill_(42)  # a later in-place host update before commit is what the copy sees
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(dst.cpu(), torch.full((8,), 42, dtype=torch.int32))
-
-
-def test_block_offsets_matches_native_kernel():
-    from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import (
-        copy_batch_block_offsets_to_device,
-    )
-
-    torch.manual_seed(1)
-    pools, src_seqs, dst_seqs, blocks = 2, 3, 2, 16388
-    host = _rand_i32(pools, src_seqs, 2, blocks, lo=-1, hi=50000).pin_memory()
-    host[:, :, :, 7000:] = -1
-    scales = torch.tensor([3, 5], dtype=torch.int32).pin_memory()
-    kv_off = torch.tensor([11, 0], dtype=torch.int32).pin_memory()
-    copy_idx = torch.tensor([2, 0], dtype=torch.int32).pin_memory()
-    ref = _rand_i32(pools, dst_seqs + 1, 2, blocks).cuda()
-    out = ref.clone()
-    copy_batch_block_offsets_to_device(
-        host, ref, copy_idx, scales, kv_off, torch.cuda.current_stream().cuda_stream
-    )
-    st = _stager()
-    st.begin()
-    st.block_offsets(host, out, copy_idx.tolist(), scales.tolist(), kv_off.tolist())
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(out, ref)
-
-
-def test_copy_clamp_and_clamp_rest():
-    dev = "cuda"
-    buf = _rand_i32(4, 300).to(dev)
-    ref = buf.clone()
-    host = _rand_i32(1, 300, lo=-5, hi=5).pin_memory()
-    ref[:1].copy_(host, non_blocking=True)
-    ref.clamp_(min=0)
-    st = _stager()
-    st.begin()
-    st.copy(buf[:1], host, clamp_min=0)
-    st.clamp_min(buf[1:], 0)
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(buf, ref)
-
-
-def test_overlapping_writes_keep_native_order():
-    """Overlapping staged writes land in staging order (one more launch each)."""
-    torch.manual_seed(6)
-    dst = _rand_i32(64).cuda()
-    ref = dst.clone()
-    host = _rand_i32(64, lo=-20, hi=20).pin_memory()
-    ref.copy_(host, non_blocking=True)
-    ref[4:12] = torch.arange(8, dtype=torch.int32, device="cuda")
-    ref[:8].clamp_(min=5)
-    ref[60:].zero_()
-    st = _stager()
-    st.begin()
-    st.copy(dst, host)
-    st.values(dst[4:12], list(range(8)))
-    st.clamp_min(dst[:8], 5)
-    st.zero(dst[60:])
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(dst, ref)
-    assert st.num_flushes == 2 and st.num_native_copies == 0
-
-
-def test_flush_then_native_op_keeps_native_order():
-    """``flush`` in the window: a native op after it sees the staged writes."""
-    torch.manual_seed(8)
-    dst = _rand_i32(64).cuda()
-    ref = dst.clone()
-    host = _rand_i32(64, lo=-20, hi=20).pin_memory()
-    ref.copy_(host, non_blocking=True)
-    ref[:8].clamp_(min=5)
-    ref[4:12] += 7
-    ref[60:].zero_()
-    st = _stager()
-    st.begin()
-    st.copy(dst, host)
-    st.clamp_min(dst[:8], 5)
-    st.flush()
-    dst[4:12] += 7
-    st.zero(dst[60:])
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(dst, ref)
-    assert st.num_flushes == 2 and st.num_native_copies == 0 and st.num_commits == 1
-
-
-def test_every_kind_falls_back_to_native_in_order():
-    """Writes too large for an empty packet are issued natively, in order."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.packet import (  # noqa: E501
-        PrepPacket,
-    )
-
-    torch.manual_seed(7)
-    dev = "cuda"
-    a = _rand_i32(8).to(dev)
-    b = _rand_i32(5000).to(dev)
-    c = _rand_i32(32).to(dev)
-    e = _rand_i32(5000).to(dev)
-    f = _rand_i32(5000, lo=-9, hi=9).to(dev)
-    tail = _rand_i32(2).to(dev)
-    src = _rand_i32(6, 1, 1).to(dev)
-    host_b = _rand_i32(5000).pin_memory()
-    pools, blocks = 2, 9000
-    host_t = _rand_i32(pools, 3, 2, blocks, lo=-1, hi=50000).pin_memory()
-    bo = _rand_i32(pools, 2, 2, blocks).to(dev)
-    refs = [t.clone() for t in (a, b, c, e, f, tail, bo)]
-    ra, rb, rc, re_, rf, rtail, rbo = refs
-    ra.copy_(torch.arange(8, dtype=torch.int32), non_blocking=True)
-    rb.copy_(host_b, non_blocking=True)
-    rc.zero_()
-    rc[:6] = src.flatten() + 2
-    re_.zero_()
-    rf.clamp_(min=0)
-    for p in range(pools):
-        for s, row in enumerate([2, 0]):
-            v = host_t[p, row, 0].cuda()
-            rbo[p, s, 0] = torch.where(v == -1, 0, v * (p + 2))
-            rbo[p, s, 1] = torch.where(v == -1, 0, v * (p + 2) + 3)
-    rtail.copy_(torch.tensor([5, 6], dtype=torch.int32), non_blocking=True)
-
-    st = _stager()
-    st.packet = PrepPacket(max_desc=1, payload_words=4)
-    st.begin()
-    st.values(a, list(range(8)))  # payload 8 > 4 words
-    st.copy(b, host_b)  # 2 chunks > 1 descriptor
-    st.gather(c, src, list(range(6)), addend=2)  # payload 6 > 4
-    st.zero(e)  # 2 chunks
-    st.clamp_min(f, 0)  # 2 chunks
-    st.block_offsets(host_t, bo, [2, 0], [2, 3], [3, 3])  # 3 chunks per row
-    st.values(tail, [5, 6])  # fits
-    st.commit()
-    torch.cuda.synchronize()
-    for got, want in zip((a, b, c, e, f, tail, bo), refs):
-        assert torch.equal(got, want)
-    assert st.num_native_copies == 5 + 2 * pools  # a, b, c, e, f + every block row
-
-
-@pytest.mark.parametrize("caller_event", [False, True])
-def test_packet_ring_under_gpu_lag(caller_event):
-    """Delayed producer: many packets staged while the GPU is held back.
-
-    Every launch reads its descriptors, its payload and the caller's pinned
-    sources long after ``commit`` returned. The caller drops its pinned
-    temporary right after staging and the pinned allocator is churned with
-    junk of the same size, so a slot or source reused before its launch ran
-    would show up as wrong destination bytes.
-    """
-    torch.manual_seed(8)
-    dev = "cuda"
-    n_iter, words = 24, 3 * 4096 + 7
-    big = [torch.zeros(words, dtype=torch.int32, device=dev) for _ in range(n_iter)]
-    small = [torch.zeros(9, dtype=torch.int32, device=dev) for _ in range(n_iter)]
-    vals = [torch.zeros(5, dtype=torch.int32, device=dev) for _ in range(n_iter)]
-    st = _stager()
-    st.begin()
-    st.values(vals[0], [0] * 5)  # compile the kernel outside the lag window
-    st.commit()
-    torch.cuda.synchronize()
-
-    expected, junk = [], []
-    torch.cuda._sleep(HOLD_CYCLES)
-    held = torch.cuda.Event()
-    held.record()
-    for i in range(n_iter):
-        host_big = _rand_i32(words).pin_memory()  # read in place by the kernel
-        host_small = _rand_i32(9)  # pageable: goes through the payload
-        values = [100 * i + k for k in range(5)]
-        expected.append((host_big.clone(), host_small.clone(), values))
-        st.begin()
-        st.copy(big[i], host_big)
-        st.copy(small[i], host_small)
-        st.values(vals[i], values)
-        if caller_event:
-            st.commit(record_event=False)
-            ev = torch.cuda.Event()
-            ev.record()
-            st.set_completion_event(ev)
-        else:
-            st.commit()
-        del host_big, host_small
-        junk.append(torch.full((words,), -7, dtype=torch.int32).pin_memory())
-    lagging = not held.query()
-    slots_while_lagging = len(st.packet._slots)
-    torch.cuda.synchronize()
-    for i, (hb, hs, values) in enumerate(expected):
-        assert torch.equal(big[i].cpu(), hb), i
-        assert torch.equal(small[i].cpu(), hs), i
-        assert vals[i].tolist() == values, i
-    assert all(bool((j == -7).all()) for j in junk)
-    if not UNDER_SANITIZER:
-        assert lagging, "the GPU finished its hold before staging ended; inconclusive"
-        assert slots_while_lagging >= n_iter, slots_while_lagging
-
-
-def test_sampler_store_ring_under_gpu_lag():
-    """Delayed producer for the sampler store writer (slot ring + zero-copy host targets)."""
-    import random
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
-        SamplerStoreWriter,
-    )
-
-    torch.manual_seed(9)
-    rng = random.Random(9)
-    dev = "cuda"
-    num_slots, w, d, n_iter = 8, 6, 5, 24
-
-    def zeros(*shape):
-        return torch.zeros(shape, dtype=torch.int32, device=dev)
-
-    store = SimpleNamespace(
-        new_tokens=zeros(w, num_slots, 1),
-        next_new_tokens=zeros(w, num_slots, 1),
-        new_tokens_lens=zeros(num_slots),
-        next_draft_tokens=zeros(num_slots, d),
-    )
-    ref = SimpleNamespace(**{k: v.cpu() for k, v in vars(store).items()})
-    calls = []
-    for i in range(n_iter):
-        slots = rng.sample(range(num_slots), rng.randint(1, 3))
-        n = len(slots)
-        calls.append(
-            (
-                slots,
-                _rand_i32(n, w).to(dev),
-                _rand_i32(n, w).to(dev),
-                _rand_i32(n, lo=1, hi=7).to(dev),
-                _rand_i32(n, d).to(dev),
-            )
-        )
-    calls_cpu = [tuple(t.cpu() for t in c[1:]) for c in calls]
-    writer = SamplerStoreWriter()
-    warm = [
-        torch.empty_like(t, device="cpu", pin_memory=True)
-        for t in (store.new_tokens, store.new_tokens_lens, store.next_draft_tokens)
-    ]
-    scratch = SimpleNamespace(**{k: v.clone() for k, v in vars(store).items()})
-    # Compile every argument specialization (request count, alignments) of the
-    # kernel outside the lag window: a compile inside it outlasts the hold.
-    for call in calls:
-        writer.store_and_copy_out(scratch, *call, *warm)
-    torch.cuda.synchronize()
-
-    kept, junk = [], []
-    torch.cuda._sleep(HOLD_CYCLES)
-    held = torch.cuda.Event()
-    held.record()
-    for i, ((slots, o_new, o_next, o_lens, o_draft), cpu) in enumerate(zip(calls, calls_cpu)):
-        # Host reference from host copies made before the hold (a D2H copy here
-        # would wait for the GPU).
-        s = torch.tensor(slots)
-        ref.new_tokens[:, s, 0] = cpu[0].T
-        ref.next_new_tokens[:, s, 0] = cpu[1].T
-        ref.new_tokens_lens[s] = cpu[2]
-        ref.next_draft_tokens[s] = cpu[3]
-        want = (ref.new_tokens.clone(), ref.new_tokens_lens.clone(), ref.next_draft_tokens.clone())
-        host = [
-            torch.empty_like(t, device="cpu", pin_memory=True)
-            for t in (store.new_tokens, store.new_tokens_lens, store.next_draft_tokens)
-        ]
-        writer.store_and_copy_out(
-            store, slots, o_new, o_next, o_lens, o_draft, *host, record_event=False
-        )
-        ev = torch.cuda.Event()
-        ev.record()
-        writer.set_completion_event(ev)
-        if i % 2:
-            kept.append((host, want))
-        else:
-            # Dropped targets must stay owned by the writer until its launch ran:
-            # junk allocations of the same size must not receive the store bytes.
-            del host
-            junk.extend(torch.full_like(t, -9).pin_memory() for t in want)
-    lagging = not held.query()
-    ring_while_lagging = len(writer._slots)
-    torch.cuda.synchronize()
-    for host, want in kept:
-        for h, x in zip(host, want):
-            assert torch.equal(h, x)
-    assert all(bool((j == -9).all()) for j in junk)
-    for k in vars(store):
-        assert torch.equal(getattr(store, k).cpu(), getattr(ref, k)), k
-    if not UNDER_SANITIZER:
-        assert lagging, "the GPU finished its hold before staging ended; inconclusive"
-        assert ring_while_lagging > 4, ring_while_lagging
-
-
-@pytest.mark.parametrize("num_slots,slots", [(1, [0]), (4, [2, 0]), (8, [7, 3, 5])])
-def test_sampler_store_matches_index_copy_and_d2h(num_slots, slots):
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
-        SamplerStoreWriter,
-    )
-
-    torch.manual_seed(2)
-    dev = "cuda"
-    w_new, w_next, d = 6, 6, 5
-
-    def make_store():
-        return SimpleNamespace(
-            new_tokens=_rand_i32(w_new, num_slots, 1).to(dev),
-            next_new_tokens=_rand_i32(w_next, num_slots, 1).to(dev),
-            new_tokens_lens=_rand_i32(num_slots).to(dev),
-            next_draft_tokens=_rand_i32(num_slots, d).to(dev),
-        )
-
-    ref = make_store()
-    out = SimpleNamespace(**{k: v.clone() for k, v in vars(ref).items()})
-    n = len(slots)
-    big_new = _rand_i32(n + 2, w_new).to(dev)
-    o_new = big_new[1 : 1 + n]
-    o_next = _rand_i32(n, w_next).to(dev)
-    o_lens = _rand_i32(n, lo=1, hi=7).to(dev)
-    o_draft = _rand_i32(n, d).to(dev)
-
-    s = torch.tensor(slots, dtype=torch.long).to(dev, non_blocking=True)
-    ref.new_tokens.squeeze(-1).T.index_copy_(0, s, o_new)
-    ref.next_new_tokens.squeeze(-1).T.index_copy_(0, s, o_next)
-    ref.new_tokens_lens.index_copy_(0, s, o_lens)
-    ref.next_draft_tokens.index_copy_(0, s, o_draft)
-    h_ref = [
-        torch.empty_like(t, device="cpu", pin_memory=True)
-        for t in (ref.new_tokens, ref.new_tokens_lens, ref.next_draft_tokens)
-    ]
-    for h, t in zip(h_ref, (ref.new_tokens, ref.new_tokens_lens, ref.next_draft_tokens)):
-        h.copy_(t, non_blocking=True)
-
-    writer = SamplerStoreWriter()
-    assert writer.tensors_supported(out, o_new, o_next, o_lens, o_draft)
-    h_out = [
-        torch.empty_like(t, device="cpu", pin_memory=True)
-        for t in (out.new_tokens, out.new_tokens_lens, out.next_draft_tokens)
-    ]
-    writer.store_and_copy_out(out, slots, o_new, o_next, o_lens, o_draft, *h_out)
-    torch.cuda.synchronize()
-    for k in vars(ref):
-        assert torch.equal(getattr(out, k), getattr(ref, k)), k
-    for a, b in zip(h_out, h_ref):
-        assert torch.equal(a, b)
-
-
-def test_sampler_store_without_sampling_requests():
-    """No sampling request (a prefill chunk): empty outputs, the store is copied out unchanged.
-
-    The empty outputs have the null address; the launcher must accept it.
-    """
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
-        SamplerStoreWriter,
-    )
-
-    torch.manual_seed(3)
-    dev = "cuda"
-    num_slots, w, d = 4, 6, 5
-    store = SimpleNamespace(
-        new_tokens=_rand_i32(w, num_slots, 1).to(dev),
-        next_new_tokens=_rand_i32(w, num_slots, 1).to(dev),
-        new_tokens_lens=_rand_i32(num_slots).to(dev),
-        next_draft_tokens=_rand_i32(num_slots, d).to(dev),
-    )
-    ref = {k: v.clone() for k, v in vars(store).items()}
-    o_new = torch.empty((0, w), dtype=torch.int32, device=dev)
-    o_next = torch.empty((0, w), dtype=torch.int32, device=dev)
-    o_lens = torch.empty((0,), dtype=torch.int32, device=dev)
-    o_draft = torch.empty((0, d), dtype=torch.int32, device=dev)
-    assert o_new.data_ptr() == 0
-    writer = SamplerStoreWriter()
-    h_out = [
-        torch.full_like(t, -1, device="cpu").pin_memory()
-        for t in (store.new_tokens, store.new_tokens_lens, store.next_draft_tokens)
-    ]
-    writer.store_and_copy_out(store, [], o_new, o_next, o_lens, o_draft, *h_out)
-    torch.cuda.synchronize()
-    for k, v in ref.items():
-        assert torch.equal(getattr(store, k), v), k
-    for h, k in zip(h_out, ("new_tokens", "new_tokens_lens", "next_draft_tokens")):
-        assert torch.equal(h, ref[k].cpu()), k
-
-
-def test_pinned_in_place_and_pageable_payload_sources():
-    dev = "cuda"
-    big = torch.zeros(3 * 4096 + 17, dtype=torch.int32, device=dev)  # multi-chunk, in place
-    small = torch.zeros(9, dtype=torch.int32, device=dev)  # pageable -> payload
-    host_big = _rand_i32(big.numel()).pin_memory()
-    host_small = _rand_i32(9)
-    st = _stager()
-    st.begin()
-    st.copy(big, host_big)
-    st.copy(small, host_small)
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(big.cpu(), host_big) and torch.equal(small.cpu(), host_small)
-
-
-def test_packet_full_falls_back_to_native_writes():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.packet import (  # noqa: E501
-        PrepPacket,
-    )
-
-    torch.manual_seed(3)
-    dev = "cuda"
-    st = _stager()
-    st.packet = PrepPacket(max_desc=3, payload_words=64)
-    a = torch.zeros(5000, dtype=torch.int32, device=dev)
-    b = torch.zeros(100, dtype=torch.int32, device=dev)
-    c = _rand_i32(64).to(dev)
-    src = _rand_i32(6, 1, 1).to(dev)
-    host_a = _rand_i32(5000).pin_memory()
-    host_b = _rand_i32(100)
-    c_ref = torch.zeros_like(c)
-    c_ref[:6] = src.flatten() + 3
-    st.begin()
-    st.copy(a, host_a)  # 2 descriptors
-    st.copy(b, host_b)  # payload too small -> native
-    st.gather(c, src, list(range(6)), addend=3)  # 1 descriptor
-    st.zero(c[:0]) if c[:0].numel() else None
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(a.cpu(), host_a) and torch.equal(b.cpu(), host_b)
-    assert torch.equal(c, c_ref)
-    assert st.num_native_copies >= 1
-
-
-@pytest.mark.parametrize("next_n", [1, 6])
-def test_paged_mqa_schedule_host_replica_matches_deepgemm(next_n):
-    """The host replica equals ``get_paged_mqa_logits_metadata`` bit for bit."""
-    import random
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.paged_mqa_schedule import (  # noqa: E501
-        SPLIT_KV_FOR_BLOCK_KV,
-        paged_mqa_logits_schedule,
-    )
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _DG_SCHEDULE_BLOCK_KV
-    from tensorrt_llm.deep_gemm import get_num_sms, get_paged_mqa_logits_metadata
-
-    num_sms = get_num_sms()
-    split_kv = SPLIT_KV_FOR_BLOCK_KV[_DG_SCHEDULE_BLOCK_KV]
-    max_len = (1 << 20) + 14
-    rng = random.Random(next_n)
-    cases = [
-        [k * split_kv + d]
-        for k in range(0, max_len // split_kv + 1)
-        for d in (-1, 0, 1)
-        if 1 <= k * split_kv + d <= max_len
-    ]
-    cases += [[rng.randint(1, max_len) for _ in range(rng.randint(2, 8))] for _ in range(300)]
-    cases += [[rng.randint(1, 4096) for _ in range(rng.randint(2, 8))] for _ in range(300)]
-    mismatches = 0
-    for lens in cases:
-        t = torch.tensor(lens, dtype=torch.int32, device="cuda").view(-1, 1)
-        if next_n > 1:
-            t = t.expand(-1, next_n).contiguous()
-        native = get_paged_mqa_logits_metadata(t, _DG_SCHEDULE_BLOCK_KV, num_sms).cpu().numpy()
-        host = paged_mqa_logits_schedule(lens, next_n, split_kv, num_sms)
-        if native.shape != host.shape or not (native == host).all():
-            mismatches += 1
-            if mismatches <= 3:
-                print("mismatch", lens[:4], native.reshape(-1)[:12], host.reshape(-1)[:12])
-    assert mismatches == 0, f"{mismatches}/{len(cases)} schedules differ"
-
-
-# ---- debug guard ------------------------------------------------------------
-
-
-def test_debug_guard_flags_torch_op_on_staged_destination():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
-        InputPrepHazardError,
-    )
-
-    a = torch.zeros(8, dtype=torch.int32, device="cuda")
-    b = torch.zeros(8, dtype=torch.int32, device="cuda")
-    st = _stager(1)
-    st.begin()
-    a.add_(1)  # before staging: native order is preserved
-    st.values(a, list(range(8)))
-    b.add_(1)  # unrelated buffer
-    _ = a[2:4]  # a view reads no data
-    with pytest.raises(InputPrepHazardError, match="touches a destination staged"):
-        _ = a[2:4] * 2
-    st.discard()
-    # The guard is gone after discard.
-    _ = a * 2
-    torch.cuda.synchronize()
-
-
-def test_debug_guard_flags_triton_kernel_on_staged_destination():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
-        InputPrepHazardError,
-    )
-
-    a = torch.zeros(8, dtype=torch.int32, device="cuda")
-    out = torch.zeros(8, dtype=torch.int32, device="cuda")
-    st = _stager(1)
-    st.begin()
-    st.values(a, list(range(8)))
-    with pytest.raises(InputPrepHazardError, match="triton:_copy_tensor_kernel"):
-        _copy_tensor_kernel[(1,)](a, out, 8, BLOCK=8)
-    st.discard()
-
-
-def test_debug_guard_flags_hidden_writer_at_launch():
-    """A writer the hooks cannot see (direct CUDA launch) is caught before the launch."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
-        InputPrepHazardError,
-    )
-
-    a = torch.zeros(8, dtype=torch.int32, device="cuda")
-    junk = torch.full((8,), 3, dtype=torch.int32, device="cuda")
-    st = _stager(1)
-    st.begin()
-    st.zero(a)
-    _copy_addr_kernel[(1,)](junk.data_ptr(), a.data_ptr(), 8, BLOCK=8)
-    with pytest.raises(InputPrepHazardError, match="written by another operation"):
-        st.commit()
-    assert not st.active
-
-
-@pytest.mark.parametrize("level", [1, 2])
-def test_debug_levels_keep_final_bytes(level):
-    """The guard and the poison do not change what the launch writes."""
-    torch.manual_seed(10)
-    dst = _rand_i32(300).cuda()
-    ref = dst.clone()
-    host = _rand_i32(100).pin_memory()
-    ref[:100].copy_(host, non_blocking=True)
-    ref[100:200].clamp_(min=0)
-    ref[200:].zero_()
-    st = _stager(level)
-    st.begin()
-    st.copy(dst[:100], host)
-    st.clamp_min(dst[100:200], 0)
-    st.zero(dst[200:])
-    st.commit()
-    torch.cuda.synchronize()
-    assert torch.equal(dst, ref)
-
-
-def test_debug_poison_reaches_hidden_reader():
-    """Level 2: a reader the hooks cannot see consumes poison, not stale bytes."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.guard import (  # noqa: E501
-        POISON_BYTE,
-    )
-
-    a = torch.full((8,), 7, dtype=torch.int32, device="cuda")
-    out = torch.zeros(8, dtype=torch.int32, device="cuda")
-    st = _stager(2)
-    st.begin()
-    st.values(a, list(range(8)))
-    _copy_addr_kernel[(1,)](a.data_ptr(), out.data_ptr(), 8, BLOCK=8)
-    st.commit()
-    torch.cuda.synchronize()
-    assert a.tolist() == list(range(8))
-    assert bool((out.view(torch.uint8) == POISON_BYTE).all())
-
-
-# ---- support domain: executor loop and preparation eligibility ---------------
-
-
-@pytest.fixture
-def loop_state(monkeypatch):
-    """A fresh executor-loop domain state for the process."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        compat,
-    )
-
-    state = {"declared": False, "reason": "", "logged": False}
-    monkeypatch.setattr(compat, "_LOOP_STATE", state)
-    return state
-
-
-def _store_args(dev="cuda", num_slots=4, w=6, d=5):
-    store = SimpleNamespace(
-        new_tokens=torch.zeros(w, num_slots, 1, dtype=torch.int32, device=dev),
-        next_new_tokens=torch.zeros(w, num_slots, 1, dtype=torch.int32, device=dev),
-        new_tokens_lens=torch.zeros(num_slots, dtype=torch.int32, device=dev),
-        next_draft_tokens=torch.zeros(num_slots, d, dtype=torch.int32, device=dev),
-    )
-    outs = (
-        torch.zeros(1, w, dtype=torch.int32, device=dev),
-        torch.zeros(1, w, dtype=torch.int32, device=dev),
-        torch.zeros(1, dtype=torch.int32, device=dev),
-        torch.zeros(1, d, dtype=torch.int32, device=dev),
-    )
-    return store, outs
-
-
-def test_store_writer_needs_the_validated_loop_domain(loop_state):
-    """The writer keeps the native path until a non-ADP overlap loop declared its domain."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        compat,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.sampler_store import (  # noqa: E501
-        SamplerStoreWriter,
-    )
-
-    writer = SamplerStoreWriter()
-    store, outs = _store_args()
-    assert writer.tensors_supported(store, *outs)
-    # No overlap loop yet (also: the non-overlap and pipeline-parallel loops never declare).
-    assert not writer.supported(store, *outs)
-    compat.declare_overlap_loop(False)
-    assert writer.supported(store, *outs)
-    # A later loop with attention DP: final for the process.
-    compat.declare_overlap_loop(True)
-    assert not writer.supported(store, *outs)
-    compat.declare_overlap_loop(False)
-    assert not writer.supported(store, *outs)
-    assert writer.num_calls == 0
-
-
-@pytest.mark.parametrize(
-    "adp, expected",
-    [(False, True), (True, False), (None, False)],
-)
-def test_the_loop_domain_declines_under_attention_dp(loop_state, adp, expected):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        compat,
-    )
-
-    compat.declare_overlap_loop(adp)  # None: a caller without the setting
-    assert loop_state["declared"]
-    assert compat.loop_domain() is expected
-
-
-def _eligibility_case(**overrides):
-    """Arguments of ``PrepStager._ineligible_reason`` for a steady served decode step."""
-    mode = SimpleNamespace(use_one_engine=lambda: True)
-    cfg = SimpleNamespace(
-        spec_config=SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True),
-        disable_overlap_scheduler=overrides.pop("disable_overlap_scheduler", False),
-        enable_attention_dp=overrides.pop("enable_attention_dp", False),
-        use_mrope=False,
-    )
-    runner = SimpleNamespace(
-        _config=cfg,
-        use_beam_search=False,
-        mapping=SimpleNamespace(has_cp_helix=lambda: False, pp_size=overrides.pop("pp_size", 1)),
-        guided_decoder=None,
-        _lora=SimpleNamespace(cuda_graph_manager=None),
-    )
-    rows = overrides.pop("rows", [SimpleNamespace(is_dummy=False, py_batch_idx=0)])
-    requests = SimpleNamespace(num_context_requests=0, generation_requests=rows)
-    DSAtrtllmAttentionMetadata = type("DSAtrtllmAttentionMetadata", (), {})
-    manager = SimpleNamespace(index_mapper=object(), _use_per_layer_page_tables=False)
-    attn = DSAtrtllmAttentionMetadata()
-    attn.kv_cache_manager = manager
-    attn.draft_kv_cache_manager = None
-    args = dict(
-        runner=runner,
-        requests=requests,
-        attn_metadata=attn,
-        spec_metadata=object(),
-        new_tensors_device=SimpleNamespace(next_draft_tokens=object()),
-        graph_replay=True,
-        is_dummy=False,
-        enable_spec_decode=True,
-        promoted=False,
-    )
-    args.update(overrides)
-    return args
-
-
-def test_preparation_eligibility_declines_attention_dp_and_new_arrivals():
-    """Prep stays native under attention DP and for rows without previous device tensors.
-
-    A request that just arrived (also a disaggregated generation request whose
-    KV cache came from a context server) has no previous iteration's device
-    tensors (``py_batch_idx`` None), so its first preparation is native.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.stager import (  # noqa: E501
-        PrepStager,
-    )
-
-    reason = PrepStager._ineligible_reason
-    assert reason(**_eligibility_case()) is None
-    assert (
-        reason(**_eligibility_case(enable_attention_dp=True))
-        == "overlap off, attention DP or MRoPE"
-    )
-    assert reason(**_eligibility_case(disable_overlap_scheduler=True)) == (
-        "overlap off, attention DP or MRoPE"
-    )
-    arrival = [
-        SimpleNamespace(is_dummy=False, py_batch_idx=0),
-        SimpleNamespace(is_dummy=False, py_batch_idx=None),
-    ]
-    assert reason(**_eligibility_case(rows=arrival)) == "dummy row or row without previous tensors"
-    dummy = [SimpleNamespace(is_dummy=True, py_batch_idx=0)]
-    assert reason(**_eligibility_case(rows=dummy)) == "dummy row or row without previous tensors"
-    assert reason(**_eligibility_case(new_tensors_device=None)) == (
-        "no previous speculative device tensors"
-    )
-
-
-def test_preparation_eligibility_declines_pipeline_parallelism():
-    """deferred_input_prep is not validated with the pipeline-parallel executor loop: prep stays native."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep.stager import (  # noqa: E501
-        PrepStager,
-    )
-
-    reason = PrepStager._ineligible_reason
-    assert reason(**_eligibility_case(pp_size=1)) is None
-    assert reason(**_eligibility_case(pp_size=2)) == "pipeline parallelism"
-
-
-def test_prep_declines_without_pinned_host_memory(monkeypatch):
-    """Confidential Compute (``prefer_pinned()`` False) keeps every preparation native.
-
-    The staged page-table copy reads the pinned host table; with pageable host
-    tables it would raise in the middle of a preparation instead of declining.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        stager as stager_mod,
-    )
-
-    monkeypatch.setattr(stager_mod, "_host_memory_pinned", lambda: False)
-    st = stager_mod.PrepStager()
-    got = st.begin_if_eligible(
-        **_eligibility_case(),
-    )
-    assert got is None
-    assert st.last_reason == "host page tables are not pinned (prefer_pinned)"
-    assert not st.active
-
-
-def test_prep_decisions_are_logged_once_per_kind(monkeypatch):
-    """Every kind of preparation decision is logged once, later ones only counted."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_input_prep import (  # noqa: E501
-        stager as stager_mod,
-    )
-
-    lines = []
-    monkeypatch.setattr(dispatch, "_log_info", lines.append)
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(stager_mod, "_host_memory_pinned", lambda: True)
-    st = stager_mod.PrepStager()
-    for _ in range(3):
-        assert st.begin_if_eligible(**_eligibility_case(pp_size=2)) is None
-    for _ in range(2):
-        assert st.begin_if_eligible(**_eligibility_case(new_tensors_device=None)) is None
-    assert st.decisions == {
-        "pipeline parallelism": 3,
-        "no previous speculative device tensors": 2,
-    }
-    assert lines == [
-        "[glm_moe_dsa] deferred_input_prep/prep[pipeline parallelism]: native fallback (pipeline parallelism)",
-        "[glm_moe_dsa] deferred_input_prep/prep[no previous speculative device tensors]: native fallback "
-        "(no previous speculative device tensors)",
-    ]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py
deleted file mode 100644
index c48abebdbaa9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/conftest.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# 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.
-"""Make the deferred_norm helpers and the absorbed_query / post_moe_hop fixtures importable under
-``--import-mode=importlib``."""
-
-import sys
-from pathlib import Path
-
-_HERE = Path(__file__).resolve().parent
-for _path in (str(_HERE), str(_HERE.parent / "absorbed_query"), str(_HERE.parent / "post_moe_hop")):
-    if _path not in sys.path:
-        sys.path.insert(0, _path)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py
deleted file mode 100644
index 3a103b099ceb..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/deferred_norm_utils.py
+++ /dev/null
@@ -1,299 +0,0 @@
-# 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.
-"""Shared helpers of the deferred_norm tests and benchmarks.
-
-* ``hop_sumsq_partials``: torch emulation of the deferred_norm / post_moe_hop's square-sum
-  partials (per thread: 8 columns, sequential fp32 sum of bf16-rounded
-  squares; warp butterfly 16, 8, 4, 2, 1; one value per warp), bit-exact by
-  construction because every step is one IEEE fp32 operation.
-* ``fp64_oracle``: exact projection of the un-rounded RMSNorm output.
-* ``ProxyProducer``: a CuTe DSL stand-in for a hop kernel. It is launched
-  with programmatic dependent launch, releases its dependents at entry (or
-  at the end), spins for a given time and only then writes the packet
-  tensors; a dependent that reads them before its dependency wait would see
-  stale data.
-* error statistics (normalized RMS error, bf16 ulp distance).
-"""
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cuda.bindings.driver import CUstream
-from cutlass import Int32, Int64
-from cutlass._mlir.dialects import llvm
-from cutlass.cutlass_dsl import dsl_user_op
-
-H = 6144
-EPS = 1e-5
-HOP_CTAS = 8
-HOP_THREADS = H // (8 * HOP_CTAS)  # 96
-HOP_WARPS = HOP_THREADS // 32  # 3
-PARTIALS = HOP_CTAS * HOP_WARPS  # 24
-
-
-def hop_sumsq_partials(x: torch.Tensor) -> torch.Tensor:
-    """fp32 ``[R, 24]`` partials exactly as the hop computes them from bf16 ``x``."""
-    rows = x.shape[0]
-    xf = x.float().view(rows, HOP_CTAS, HOP_WARPS, 32, 8)
-    sq = (xf * xf).bfloat16().float()
-    s = torch.zeros(rows, HOP_CTAS, HOP_WARPS, 32, dtype=torch.float32, device=x.device)
-    for i in range(8):
-        s = s + sq[..., i]
-    lanes = torch.arange(32, device=x.device)
-    for mask in (16, 8, 4, 2, 1):
-        s = s + s[..., lanes ^ mask]
-    return s[..., 0].reshape(rows, PARTIALS).contiguous()
-
-
-def square_sum_in_order(partials: torch.Tensor) -> torch.Tensor:
-    """fp32 sum of the partials in index order (the consumer's order)."""
-    s = torch.zeros(partials.shape[0], dtype=torch.float32, device=partials.device)
-    for i in range(partials.shape[1]):
-        s = s + partials[:, i]
-    return s
-
-
-def make_frag_packet(x: torch.Tensor, gamma: torch.Tensor, eps: float = EPS):
-    """Packet of the hop computed in torch (bit-exact ``a`` and partials).
-
-    ``sumsq`` is a view of an 8-row buffer (``sumsq_full(packet)`` is the
-    full buffer) so a proxy producer can copy whole 8-row units.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHidden,
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        frag_from_rows,
-    )
-
-    rows = x.shape[0]
-    a = (x.float() * gamma.float()).bfloat16()
-    ss = torch.zeros(8, PARTIALS, dtype=torch.float32, device=x.device)
-    ss[:rows] = hop_sumsq_partials(x)
-    return DeferredNormHidden(frag_from_rows(a), ss[:rows], x, gamma, eps, materialize)
-
-
-def sumsq_full(p):
-    """The 8-row partial buffer behind a test packet's ``sumsq`` view (test
-    packets allocate 8 rows; inference-mode views do not keep ``_base``)."""
-    return torch.as_strided(p.sumsq, (8, PARTIALS), (PARTIALS, 1))
-
-
-def empty_like_packet(p):
-    """Same-shaped packet with fresh (uninitialized) a_frag / partials buffers."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHidden,
-    )
-
-    ss = torch.empty(8, PARTIALS, dtype=torch.float32, device=p.sumsq.device)
-    return DeferredNormHidden(
-        torch.empty_like(p.a_frag), ss[: p.num_rows], p.residual, p.gamma, p.eps, p.materializer
-    )
-
-
-def fp64_oracle(
-    x: torch.Tensor, gamma: torch.Tensor, weight: torch.Tensor, eps: float = EPS
-) -> torch.Tensor:
-    """``y* = r* · Σ_k x_k γ_k W_jk`` in fp64 (``r*`` from exact squares)."""
-    xd = x.double()
-    r = torch.rsqrt(xd.pow(2).mean(-1, keepdim=True) + eps)
-    return (xd * gamma.double() * r) @ weight.double().t()
-
-
-def norm_rms_err(y: torch.Tensor, ref: torch.Tensor) -> float:
-    ref = ref.double()
-    return float((y.double() - ref).norm() / ref.norm().clamp_min(1e-300))
-
-
-def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    def ordered(t):
-        v = t.contiguous().view(torch.int16).to(torch.int32)
-        return torch.where(v < 0, -(v + 32768), v)
-
-    return (ordered(a) - ordered(b)).abs()
-
-
-def compare(cand: torch.Tensor, native: torch.Tensor, oracle: torch.Tensor) -> dict:
-    ulp = bf16_ulp_distance(cand, native)
-    return {
-        "allclose_native": bool(torch.allclose(cand.float(), native.float(), rtol=0.1, atol=1e-2)),
-        "e_cand": norm_rms_err(cand, oracle),
-        "e_native": norm_rms_err(native, oracle),
-        "max_ulp_vs_native": int(ulp.max()),
-        "frac_diff_vs_native": float((ulp > 0).float().mean()),
-        "max_abs_vs_native": float((cand.float() - native.float()).abs().max()),
-    }
-
-
-def within_budget(stats: dict) -> bool:
-    return stats["allclose_native"] and stats["e_cand"] <= 1.25 * stats["e_native"] + 2.0**-12
-
-
-def _asm(res, args, text, cons, loc=None, ip=None):
-    return llvm.inline_asm(
-        res,
-        args,
-        text,
-        cons,
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def _spin_ns(ns: Int64, sleep_ns: int = 0, *, loc=None, ip=None):
-    nap = f"\t@p nanosleep.u32 {int(sleep_ns)};\n" if sleep_ns else ""
-    _asm(
-        None,
-        [ns.ir_value()],
-        "{\n\t.reg .u64 t0, t;\n\t.reg .pred p;\n\tmov.u64 t0, %globaltimer;\n\t"
-        "GLM_MOE_DSA_DEFERRED_NORM_SPIN:\n\tmov.u64 t, %globaltimer;\n\tsub.u64 t, t, t0;\n\t"
-        "setp.lt.u64 p, t, $0;\n" + nap + "\t@p bra GLM_MOE_DSA_DEFERRED_NORM_SPIN;\n\t}",
-        "l",
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def _copy16(dst: Int64, src: Int64, *, loc=None, ip=None):
-    _asm(
-        None,
-        [dst.ir_value(), src.ir_value()],
-        "{\n\t.reg .b32 a, b, c, d;\n\tld.global.v4.b32 {a, b, c, d}, [$1];\n\t"
-        "st.global.v4.b32 [$0], {a, b, c, d};\n\t}",
-        "l,l",
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def _copy4(dst: Int64, src: Int64, *, loc=None, ip=None):
-    _asm(
-        None,
-        [dst.ir_value(), src.ir_value()],
-        "{\n\t.reg .b32 a;\n\tld.global.b32 a, [$1];\n\tst.global.b32 [$0], a;\n\t}",
-        "l,l",
-        loc=loc,
-        ip=ip,
-    )
-
-
-class _ProxyKernel:
-    def __init__(self, trigger_at_start: bool, sleep_ns: int = 0, carveout: int = -1):
-        self.trigger_at_start = trigger_at_start
-        self.sleep_ns = sleep_ns
-        self.carveout = carveout
-
-    @cute.jit
-    def __call__(
-        self,
-        act_src: Int64,
-        act_dst: Int64,
-        ss_src: Int64,
-        ss_dst: Int64,
-        delay_ns: Int64,
-        anchor: cute.Tensor,
-        rows: Int32,
-        stream: CUstream,
-    ):
-        if cutlass.const_expr(self.carveout >= 0):
-            self.kernel(act_src, act_dst, ss_src, ss_dst, delay_ns, anchor).launch(
-                grid=[HOP_CTAS, rows, 1],
-                block=[HOP_THREADS, 1, 1],
-                stream=stream,
-                use_pdl=True,
-                preferred_smem_carveout=self.carveout,
-            )
-        else:
-            self.kernel(act_src, act_dst, ss_src, ss_dst, delay_ns, anchor).launch(
-                grid=[HOP_CTAS, rows, 1], block=[HOP_THREADS, 1, 1], stream=stream, use_pdl=True
-            )
-
-    @cute.kernel
-    def kernel(
-        self,
-        act_src: Int64,
-        act_dst: Int64,
-        ss_src: Int64,
-        ss_dst: Int64,
-        delay_ns: Int64,
-        anchor: cute.Tensor,
-    ):
-        tid, _, _ = cute.arch.thread_idx()
-        cta, row, _ = cute.arch.block_idx()
-        cute.arch.griddepcontrol_wait()
-        if cutlass.const_expr(self.trigger_at_start):
-            cute.arch.griddepcontrol_launch_dependents()
-        _spin_ns(delay_ns, self.sleep_ns)
-        off = Int64(row) * Int64(H * 2) + Int64((cta * Int32(HOP_THREADS) + tid) * Int32(16))
-        _copy16(act_dst + off, act_src + off)
-        if tid < Int32(HOP_WARPS):
-            o4 = Int64(row) * Int64(PARTIALS * 4) + Int64((cta * Int32(HOP_WARPS) + tid) * 4)
-            _copy4(ss_dst + o4, ss_src + o4)
-        if cutlass.const_expr(not self.trigger_at_start):
-            cute.arch.griddepcontrol_launch_dependents()
-
-
-_PROXY = {}
-
-
-def proxy_producer(
-    act_src, act_dst, ss_src, ss_dst, delay_ns: int, trigger_at_start: bool = True, rows: int = None
-):
-    """Launch the proxy producer on the current stream (PDL).
-
-    Copies ``rows`` x 12,288 bytes of ``act_src`` and ``rows`` x 24 fp32 of
-    ``ss_src`` (``rows`` defaults to ``act_src.shape[0]``; an ``a_frag``
-    buffer is 8 rows' worth).
-    """
-    import os
-
-    sleep_ns = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_PROXY_SLEEP_NS", "0"))
-    carveout = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_PROXY_CARVEOUT", "-1"))
-    key = (trigger_at_start, sleep_ns, carveout)
-    if key not in _PROXY:
-        from cutlass.cute import runtime as cute_runtime
-
-        z = Int64(0)
-        anchor = cute_runtime.make_fake_compact_tensor(
-            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
-        )
-        _PROXY[key] = cute.compile(
-            _ProxyKernel(trigger_at_start, sleep_ns, carveout),
-            z,
-            z,
-            z,
-            z,
-            z,
-            anchor,
-            Int32(1),
-            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
-            options="--enable-tvm-ffi",
-        )
-    _PROXY[key](
-        act_src.data_ptr(),
-        act_dst.data_ptr(),
-        ss_src.data_ptr(),
-        ss_dst.data_ptr(),
-        int(delay_ns),
-        ss_dst.view(torch.int32).view(-1),
-        rows if rows is not None else act_src.shape[0],
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py
deleted file mode 100644
index 10382e53ddbb..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_attention_integration.py
+++ /dev/null
@@ -1,495 +0,0 @@
-# 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.
-"""deferred_norm packet -> layer L+1 DSA attention through the target's adopted attention
-(absorbed_query, draft_absorbed_query, shared_indexer_absorbed_query, post_moe_hop and
-deferred_norm; single GPU; generation steps run as decode steps).
-
-Uses the absorbed_query integration harness (real weights of the full-indexer layer 6
-and the shared layer 7 at TP8 rank 0, six requests natively prefilled with
-2,600 tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6 (MTP verify), both call
-paths, ``on_update_kv_lens()`` before every step). Each layer gets its own
-input: residual-like rows ``x`` and that layer's real ``input_layernorm``
-gain; the deferred_norm packet carries ``a = bf16(x * gamma)`` and the hop-order
-square-sum partials, the native input is ``h = materialize(packet)`` (the post_moe_hop
-hop's output bit for bit). The target's deferred-norm scope is ``shared``: the hop defers
-the norm only into shared-indexer layers, so a deferred_norm step gives layer 7 the packet
-and layer 6 the normalized ``h`` (``SERVED_LAYERS``, ``_served``).
-
-* deviation (report + sanity bound): relative RMS vs the native
-  step of the attention inputs recorded at ``mla_rope_generation`` and the
-  indexer, top-k set agreement per row (identical rows, mean Jaccard), and
-  the attention output. Reported next to two native-only references: a
-  rerun of the native step (run-to-run noise) and the *rounding control*,
-  the native step on ``bf16(bf16(x * gamma) * r)`` (the normalized hidden
-  state with the gain applied first, a one-rounding perturbation of the
-  native input in native code only). Bound: output relRMS <= max(2**-7,
-  2 * control) and mean Jaccard >= 0.98 on every layer (2**-7 alone is
-  below the control: the attention output amplifies rounding-level input
-  differences through the FP8 query and the sparse softmax);
-* fallback exactness: with the deferred_norm projection declined (``_deferred_norm_projection =
-  False``) the packet is materialized and every cache, recorded input and
-  output equals the native step bit for bit;
-* engagement: the deferred_norm projection serves layer 7 and is never asked on layer 6,
-  shared_indexer_absorbed_query serves layer 7 under
-  ``absorbed_query,shared_indexer_absorbed_query,deferred_norm`` and absorbed_query serves layer 6;
-* CUDA graph: a captured deferred_norm step replays (1, 1, 20 times with fresh
-  inputs) bit-identically to the eager deferred_norm step;
-* a projection compile error on layer 7 propagates; on layer 6 the projection declines
-  before any compile;
-* negative control: a packet whose ``a_frag`` lacks the gain (``bf16(x)``
-  instead of ``bf16(x * gamma)``) violates the bound.
-
-Fallback and graph comparisons of attention outputs whose inputs and caches
-are bitwise equal allow 4 bf16 ulps of the output scale: the native sparse
-attention is not bitwise run-to-run stable.
-
-Every arm (native reference included) also runs the fast paths in
-``GLM_MOE_DSA_TEST_BASE_FEATURES`` (e.g. ``rope_kv_append``), as the absorbed_query harness
-does.
-"""
-
-import os
-import sys
-
-import absorbed_query_fixture
-import deferred_norm_utils as U
-import pytest
-import test_absorbed_query_attention_integration as aq_integration
-import torch
-
-LAYERS = aq_integration.LAYERS
-SHARED = absorbed_query_fixture.SHARED_LAYER
-FULL = absorbed_query_fixture.FULL_LAYER
-LAYOUTS = aq_integration.LAYOUTS
-# The layers the deferred_norm projection serves: the shared-indexer ones.
-SERVED_LAYERS = (SHARED,)
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
-
-import plan_testlib  # noqa: E402
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-
-def _set_features(features: str):
-    """Run exactly the fast paths ``features`` plus ``GLM_MOE_DSA_TEST_BASE_FEATURES``.
-
-    The latter are the neighbour paths every arm runs on, e.g.
-    ``rope_kv_append`` on an integration image (absorbed_query harness rule).
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query import (  # noqa: E501
-        query_proj,
-    )
-
-    query_proj.SERVED_FORM = "auto"
-    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
-
-
-def _served(packets):
-    """The inputs of a deferred_norm step: the packet on the layers the projection serves, the
-    normalized hidden state (the post-MoE hop's output) on the others."""
-    return {layer: p if layer in SERVED_LAYERS else p.materialize() for layer, p in packets.items()}
-
-
-def _reset_runners(h):
-    for layer in LAYERS:
-        for name in ("_absorbed_query", "_shared_indexer_query", "_deferred_norm_projection"):
-            setattr(h.attn[layer], name, None)
-
-
-class _DeferredNormHarness(aq_integration._Harness):
-    """absorbed_query harness whose layers take per-layer inputs (tensor or deferred_norm packet)."""
-
-    def __init__(self):
-        super().__init__()
-        self.gammas = {
-            layer: absorbed_query_fixture.load_input_layernorm(layer).cuda() for layer in LAYERS
-        }
-
-    def forward(self, layer, positions, hidden, meta, path):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-            is_deferred,
-        )
-
-        view = hidden.residual if is_deferred(hidden) else hidden
-        attn = self.attn[layer]
-        if path == "forward_impl":
-            out = attn._create_outputs(view, meta)
-            attn.forward_impl(positions, hidden, meta, attn_output=out)
-            return out[0]
-        import weakref
-
-        from tensorrt_llm._torch.utils import model_extra_attrs
-
-        extra = self.model_config.extra_attrs
-        extra["attention_metadata"] = weakref.ref(meta)
-        with model_extra_attrs(extra):
-            out = [torch.ops.trtllm.create_mla_outputs(view, attn.layer_idx_str)]
-            attn._forward_custom_op(hidden, positions, out, None)
-        return out[0]
-
-    def step(self, meta, positions, hidden, path):
-        """Per-layer inputs when ``hidden`` is a dict (layer -> tensor or
-        packet); otherwise the same tensor for every layer (native prefill). A
-        generation-only step runs as a decode step."""
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-            dispatch,
-        )
-
-        phase = dispatch.Phase.DECODE if meta.num_contexts == 0 else dispatch.Phase.PREFILL
-        meta.on_update_kv_lens()
-        with dispatch.step(phase):
-            return {
-                layer: self.forward(
-                    layer,
-                    positions,
-                    hidden[layer] if isinstance(hidden, dict) else hidden,
-                    meta,
-                    path,
-                )
-                for layer in LAYERS
-            }
-
-    def packets(self, layout, seed, scale=1.0):
-        nreq, ntok = LAYOUTS[layout]
-        out = {}
-        for layer in LAYERS:
-            g = torch.Generator(device="cuda").manual_seed(seed * 10 + layer)
-            x = torch.randn(nreq * ntok, U.H, generator=g, device="cuda") * 2.0
-            x = (x * scale).bfloat16()
-            out[layer] = U.make_frag_packet(x, self.gammas[layer])
-        pos = (
-            (aq_integration.CTX + torch.arange(ntok, dtype=torch.int32, device="cuda"))
-            .repeat(nreq)
-            .unsqueeze(0)
-        )
-        return out, pos
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = _DeferredNormHarness()
-    yield h
-    _set_features("")
-
-
-def _run(h, features, layout, path, inputs, pos, snap):
-    _set_features(features)
-    h.restore(snap)
-    out, rec = h.recorded_step(layout, path, inputs, pos)
-    cache = h.snapshot()
-    return out, rec, cache
-
-
-def _jaccard_rows(a, b):
-    rows = []
-    for ra, rb in zip(a.tolist(), b.tolist()):
-        sa = {v for v in ra if v >= 0}
-        sb = {v for v in rb if v >= 0}
-        union = len(sa | sb)
-        rows.append(1.0 if union == 0 else len(sa & sb) / union)
-    return rows
-
-
-def _output_stats(native, cand, layout):
-    """Deviation statistics per layer."""
-    n_out, n_rec, _ = native
-    c_out, c_rec, _ = cand
-    stats = {}
-    for layer in LAYERS:
-        s = {"out_rel_rms": aq_integration._rel_rms(c_out[layer], n_out[layer])}
-        nr, cr = n_rec[layer], c_rec[layer]
-        if "absorbed_q" in nr:
-            s["absorbed_q_rel_rms"] = aq_integration._rel_rms(cr["absorbed_q"], nr["absorbed_q"])
-        if "quant_q" in nr:
-            s["quant_q_bytes_diff_frac"] = float(
-                (cr["quant_q"].view(torch.uint8) != nr["quant_q"].view(torch.uint8)).float().mean()
-            )
-        if "indexer_weights" in nr:
-            s["indexer_weights_rel_rms"] = aq_integration._rel_rms(
-                cr["indexer_weights"], nr["indexer_weights"]
-            )
-            s["indexer_q_fp8_bytes_diff_frac"] = float(
-                (cr["indexer_q_fp8"].view(torch.uint8) != nr["indexer_q_fp8"].view(torch.uint8))
-                .float()
-                .mean()
-            )
-        if "topk" in nr:
-            jac = _jaccard_rows(cr["topk"], nr["topk"])
-            s["topk_rows_identical"] = sum(1 for j in jac if j == 1.0) / len(jac)
-            s["topk_mean_jaccard"] = sum(jac) / len(jac)
-        stats[layer] = s
-    return stats
-
-
-def _control_inputs(packets):
-    """Rounding control: ``bf16(bf16(x * gamma) * r)`` per layer (native code
-    consumes it; ``r`` in torch from the hop-order square sum)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    out = {}
-    for layer, p in packets.items():
-        rows, hidden = p.residual.shape
-        a = rows_from_frag(p.a_frag, rows, hidden)
-        r = torch.rsqrt(U.square_sum_in_order(p.sumsq) / hidden + p.eps)
-        out[layer] = (a.float() * r.unsqueeze(-1)).bfloat16()
-    return out
-
-
-def _check_output_bound(stats, control):
-    fails = []
-    for layer, s in stats.items():
-        bound = max(2.0**-7, 2.0 * control[layer]["out_rel_rms"])
-        if s["out_rel_rms"] > bound:
-            fails.append(f"L{layer} output relRMS {s['out_rel_rms']:.3e} > {bound:.3e}")
-        if s.get("topk_mean_jaccard", 1.0) < 0.98:
-            fails.append(f"L{layer} top-k mean Jaccard {s['topk_mean_jaccard']:.4f} < 0.98")
-    return fails
-
-
-@pytest.mark.parametrize(
-    "features",
-    [
-        "absorbed_query,post_moe_hop,deferred_norm",
-        "absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm",
-    ],
-)
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_deferred_norm_packet_into_attention(harness, layout, path, features):
-    h = harness
-    packets, pos = h.packets(layout, seed=3)
-    native_in = {layer: p.materialize() for layer, p in packets.items()}
-    snap = h.snapshot()
-    try:
-        base = ",".join(f for f in features.split(",") if f != "deferred_norm")
-        _reset_runners(h)
-        native = _run(h, base, layout, path, native_in, pos, snap)
-        _reset_runners(h)
-        cand = _run(h, features, layout, path, _served(packets), pos, snap)
-        for layer in LAYERS:
-            engaged = h.attn[layer]._deferred_norm_projection not in (None, False)
-            assert engaged == (layer in SERVED_LAYERS), f"deferred_norm on L{layer}: {engaged}"
-        if "shared_indexer_absorbed_query" in features:
-            assert h.attn[SHARED]._shared_indexer_query not in (None, False)
-        assert h.attn[FULL]._absorbed_query not in (None, False)
-        _reset_runners(h)
-        rerun = _run(h, base, layout, path, native_in, pos, snap)
-        _reset_runners(h)
-        control = _run(h, base, layout, path, _control_inputs(packets), pos, snap)
-        stats = _output_stats(native, cand, layout)
-        noise = _output_stats(native, rerun, layout)
-        ctrl = _output_stats(native, control, layout)
-        tag = f"layout={layout} path={path} features={features}"
-        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_BOUND {tag} {stats}", flush=True)
-        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_NOISE {tag} {noise}", flush=True)
-        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_CONTROL {tag} {ctrl}", flush=True)
-        fails = _check_output_bound(stats, ctrl)
-        assert not fails, fails
-    finally:
-        h.restore(snap)
-        _set_features("")
-
-
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_fallback_materialize_is_exact(harness, layout, path):
-    """deferred_norm projection declined -> materialized packet -> the absorbed_query path, bit for
-    bit equal to the native-input step."""
-    h = harness
-    packets, pos = h.packets(layout, seed=5)
-    native_in = {layer: p.materialize() for layer, p in packets.items()}
-    snap = h.snapshot()
-    try:
-        _reset_runners(h)
-        n_out, n_rec, n_cache = _run(h, "absorbed_query", layout, path, native_in, pos, snap)
-        _reset_runners(h)
-        for layer in LAYERS:
-            h.attn[layer]._deferred_norm_projection = False
-        _set_features("absorbed_query,post_moe_hop,deferred_norm")
-        h.restore(snap)
-        c_out, c_rec = h.recorded_step(layout, path, packets, pos)
-        c_cache = h.snapshot()
-        for key in n_cache:
-            assert torch.equal(c_cache[key], n_cache[key]), key
-        for layer in LAYERS:
-            for key, val in n_rec[layer].items():
-                if key in ("topk", "sparse_indices"):
-                    continue  # native top-k order is not run-to-run stable
-                assert torch.equal(c_rec[layer][key], val), (layer, key)
-            if "topk" in n_rec[layer]:
-                assert min(_jaccard_rows(c_rec[layer]["topk"], n_rec[layer]["topk"])) == 1.0
-        noise_runs = [n_out]
-        _reset_runners(h)
-        noise_runs.append(_run(h, "absorbed_query", layout, path, native_in, pos, snap)[0])
-        for layer in LAYERS:
-            diff = (c_out[layer].float() - n_out[layer].float()).abs().max().item()
-            noise = (noise_runs[1][layer].float() - n_out[layer].float()).abs().max().item()
-            ulp = aq_integration._bf16_ulp(n_out[layer].float().abs().max().item())
-            assert diff <= max(2 * noise, 4 * ulp), (layer, diff, noise, ulp)
-    finally:
-        h.restore(snap)
-        _set_features("")
-
-
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_cuda_graph_replay(harness, layout):
-    h = harness
-    snap = h.snapshot()
-    try:
-        _set_features("absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm")
-        _reset_runners(h)
-        packets, pos = h.packets(layout, seed=7)
-        static = {}
-        for layer, p in _served(packets).items():
-            if layer in SERVED_LAYERS:
-                st = U.empty_like_packet(p)
-                st.a_frag.copy_(p.a_frag)
-                st.sumsq.copy_(p.sumsq)
-            else:
-                st = p.clone()
-            static[layer] = st
-        static_pos = pos.clone()
-        meta = h.gen_metadata(layout, cuda_graph=True)
-        h.restore(snap)
-        h.step(meta, static_pos, static, "custom_op")
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        h.restore(snap)
-        with torch.cuda.graph(graph):
-            g_out = h.step(meta, static_pos, static, "custom_op")
-        for trial, replays in enumerate((1, 1, 20)):
-            fresh, _ = h.packets(layout, seed=100 + trial)
-            fresh = _served(fresh)
-            for layer in LAYERS:
-                if layer in SERVED_LAYERS:
-                    static[layer].a_frag.copy_(fresh[layer].a_frag)
-                    static[layer].sumsq.copy_(fresh[layer].sumsq)
-                else:
-                    static[layer].copy_(fresh[layer])
-            for _ in range(replays):
-                h.restore(snap)
-                graph.replay()
-            torch.cuda.synchronize()
-            got = {layer: g_out[layer].clone() for layer in LAYERS}
-            got_cache = h.snapshot()
-            h.restore(snap)
-            eager = h.gen_step(
-                layout, "custom_op", fresh, pos, meta=h.gen_metadata(layout, cuda_graph=True)
-            )
-            eager_cache = h.snapshot()
-            for key in eager_cache:
-                assert torch.equal(got_cache[key], eager_cache[key]), (trial, key)
-            for layer in LAYERS:
-                diff = (got[layer].float() - eager[layer].float()).abs().max().item()
-                ref = eager[layer].float().abs().max().item()
-                # Native sparse attention is not bitwise run-to-run stable:
-                # bound by 4 bf16 ulps of the output scale.
-                assert diff <= 4 * aq_integration._bf16_ulp(ref), (trial, layer, diff)
-    finally:
-        h.restore(snap)
-        _set_features("")
-
-
-@torch.inference_mode()
-def test_projection_compile_error_is_not_masked(harness, monkeypatch):
-    """A projection compile error in the consumer propagates (it invalidates
-    the run) instead of a rank-local decline to the materialized path."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        projection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.query_entry import (  # noqa: E501
-        deferred_norm_dsa_proj,
-    )
-
-    h = harness
-    packets, pos = h.packets("1req_x6tok", seed=11)
-
-    def _fail(self):
-        raise RuntimeError("injected projection compile failure")
-
-    monkeypatch.setattr(projection.DeferredNormProjection, "precompile", _fail)
-    try:
-        _set_features("absorbed_query,post_moe_hop,deferred_norm")
-        _reset_runners(h)
-        for layer in SERVED_LAYERS:
-            with pytest.raises(RuntimeError, match="injected projection compile failure"):
-                deferred_norm_dsa_proj(h.attn[layer], pos, packets[layer], None)
-            assert h.attn[layer]._deferred_norm_projection is None, (
-                "a failed runner must not be cached"
-            )
-        # A full-indexer layer keeps its native projection: declined before any compile.
-        assert deferred_norm_dsa_proj(h.attn[FULL], pos, packets[FULL], None) is None
-        assert h.attn[FULL]._deferred_norm_projection is False
-    finally:
-        _reset_runners(h)
-        _set_features("")
-
-
-@torch.inference_mode()
-def test_negative_control_missing_gain(harness):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHidden,
-        materialize,
-    )
-
-    h = harness
-    layout = "1req_x6tok"
-    packets, pos = h.packets(layout, seed=9)
-    native_in = {layer: p.materialize() for layer, p in packets.items()}
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        frag_from_rows,
-    )
-
-    bad = {
-        layer: DeferredNormHidden(
-            frag_from_rows(p.residual), p.sumsq, p.residual, p.gamma, U.EPS, materialize
-        )
-        for layer, p in packets.items()
-    }
-    snap = h.snapshot()
-    try:
-        _reset_runners(h)
-        native = _run(h, "absorbed_query", layout, "custom_op", native_in, pos, snap)
-        _reset_runners(h)
-        cand = _run(
-            h, "absorbed_query,post_moe_hop,deferred_norm", layout, "custom_op", bad, pos, snap
-        )
-        _reset_runners(h)
-        control = _run(
-            h, "absorbed_query", layout, "custom_op", _control_inputs(packets), pos, snap
-        )
-        stats = _output_stats(native, cand, layout)
-        ctrl = _output_stats(native, control, layout)
-        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_NEGATIVE {stats}", flush=True)
-        print(f"GLM_MOE_DSA_DEFERRED_NORM_OUTPUT_NEGATIVE_CONTROL {ctrl}", flush=True)
-        assert _check_output_bound(stats, ctrl), (
-            "the output bound must flag a packet without the gain"
-        )
-    finally:
-        h.restore(snap)
-        _set_features("")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py
deleted file mode 100644
index c22bc6a7eb39..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_controls.py
+++ /dev/null
@@ -1,425 +0,0 @@
-# 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.
-"""Diagnosis of deferred_norm's attention-output deviation next to native-only controls.
-
-Same single-GPU harness and metrics as ``test_deferred_norm_attention_integration.py``
-(real layer-6 / layer-7 weights at TP8 rank 0, six requests natively
-prefilled with 2,600 tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6, both
-call paths, ``absorbed_query`` and ``absorbed_query,shared_indexer_absorbed_query``). Every arm runs
-the same generation step from the same cache snapshot and is compared with the native reference step
-on ``h = materialize(packet)`` (the post_moe_hop's normalized output, bit for bit).
-
-Arms:
-
-* ``deferred_norm``: the deferred_norm packet through the deferred_norm projection (the candidate)
-  on the layers the projection serves (``dn_integration.SERVED_LAYERS``: the target's deferred-norm
-  scope is ``shared``), the normalized ``h`` on the full-indexer layer, as a served step
-  gives them.
-* ``rerun``: the native step again (run-to-run noise of native attention).
-* ``n1_rmsnorm``: native code on ``h'`` = TRT-LLM's ``RMSNorm`` module
-  (FlashInfer RMSNorm) applied to the same ``x`` and gain: a second in-tree
-  implementation of the same normalization.
-* ``n2_cublas``: native code with the ``kv_a_proj_with_mqa`` GEMM through
-  ``Linear.apply_linear`` (``cublas_mm``), the path ``DeepseekV3Linear``
-  takes for more than 16 tokens instead of ``dsv3_fused_a_gemm_op``; the
-  input ``h`` is unchanged.
-* ``n3_round``: native code on ``bf16(bf16(x * gamma) * r)`` (a single
-  extra rounding of the native input).
-* ``n4_dither1`` / ``n4_dither50``: native code on ``h`` with a random
-  +-1 bf16 ulp change in 1 % / 50 % of its elements.
-* ``ref``: native code whose projection outputs (``kva`` and, on the
-  full-indexer layer, the fused ``wk|weights_proj`` output ``kw``) are the
-  correctly rounded FP64 oracle ``bf16(r* * sum_k x_k gamma_k W_jk)``; every
-  arm is also reported as its distance to this step.
-* ``neg``: a packet without the gain through deferred_norm on the served layers (negative
-  control: a real bug).
-
-Reported per case, per arm and layer: the deviation metrics vs the native step
-(attention-output relRMS, absorbed-query relRMS, FP8 query bytes differing,
-indexer weights relRMS, FP8 indexer-query bytes differing, top-k rows
-identical, mean top-k Jaccard), the attention-output relRMS vs ``ref``, and
-whether the predeclared deviation bound (output relRMS <=
-2**-7 and mean Jaccard >= 0.98) holds. Per layout, per arm and served layer, the
-projection-level error ``e = ||y - y*|| / ||y*||`` of ``kva``.
-
-The numbers are evidence for the deviation bound, not a pass
-criterion: the assertions only check that every arm ran
-the intended path (deferred_norm engaged on the served layers only, hooks fired once, finite
-outputs) and that the experiment can see a real bug (on every served layer the negative
-control's output deviation exceeds every native-only control's).
-
-Fixture label: weights real (checkpoint), activations synthetic (Gaussian
-residual rows; the ``hot`` variant multiplies 12 channels by 30, as
-residual streams with outlier channels have).
-"""
-
-import contextlib
-import json
-import os
-
-import absorbed_query_fixture
-import deferred_norm_utils as U
-import pytest
-import test_absorbed_query_attention_integration as aq_integration
-import test_deferred_norm_attention_integration as dn_integration
-import torch
-
-LAYERS = aq_integration.LAYERS
-FULL = absorbed_query_fixture.FULL_LAYER
-LAYOUTS = aq_integration.LAYOUTS
-OUTPUT_BOUND = 2.0**-7
-NATIVE_CONTROLS = ("n1_rmsnorm", "n2_cublas", "n3_round", "n4_dither1", "n4_dither50")
-
-pytestmark = dn_integration.pytestmark
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = dn_integration._DeferredNormHarness()
-    yield h
-    dn_integration._set_features("")
-
-
-def _packets(h, layout, seed, hot):
-    """Per-layer packets of residual-like rows ``x`` (Gaussian x 2, optionally
-    12 channels x 30) with the layer's real gain."""
-    nreq, ntok = LAYOUTS[layout]
-    out = {}
-    for layer in LAYERS:
-        g = torch.Generator(device="cuda").manual_seed(seed * 10 + layer)
-        x = torch.randn(nreq * ntok, U.H, generator=g, device="cuda") * 2.0
-        if hot:
-            idx = torch.randint(0, U.H, (12,), generator=g, device="cuda")
-            x[:, idx] *= 30.0
-        out[layer] = U.make_frag_packet(x.bfloat16(), h.gammas[layer])
-    pos = (
-        (aq_integration.CTX + torch.arange(ntok, dtype=torch.int32, device="cuda"))
-        .repeat(nreq)
-        .unsqueeze(0)
-    )
-    return out, pos
-
-
-def _rmsnorm_inputs(packets):
-    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-
-    out = {}
-    for layer, p in packets.items():
-        norm = RMSNorm(hidden_size=U.H, eps=p.eps, dtype=torch.bfloat16, device="cuda")
-        norm.weight.data.copy_(p.gamma)
-        out[layer] = norm(p.residual)
-    return out
-
-
-def _dither(h, frac, seed):
-    """``h`` with a random +-1 ulp step in a fraction ``frac`` of its nonzero
-    elements (the bf16 bit pattern is sign-magnitude, so +-1 on the integer
-    view moves the magnitude by one ulp)."""
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    bits = h.contiguous().view(torch.int16)
-    pick = torch.rand(h.shape, generator=g, device="cuda") < frac
-    step = torch.where(torch.rand(h.shape, generator=g, device="cuda") < 0.5, -1, 1)
-    nonzero = (bits.to(torch.int32) & 0x7FFF) != 0
-    out = torch.where(pick & nonzero, bits.to(torch.int32) + step, bits.to(torch.int32))
-    return out.to(torch.int16).view(torch.bfloat16)
-
-
-@contextlib.contextmanager
-def _cublas_kv_a():
-    """``DeepseekV3Linear`` without its <= 16-token ``dsv3_fused_a_gemm_op``
-    branch: ``Linear.apply_linear`` (``cublas_mm``), as for larger batches."""
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Linear
-    from tensorrt_llm._torch.modules.linear import Linear
-
-    calls = [0]
-    own = DeepseekV3Linear.__dict__["apply_linear"]
-
-    def apply_linear(self, input, bias, lora_params=None, layer_idx=None):
-        calls[0] += 1
-        return Linear.apply_linear(self, input, bias, lora_params, layer_idx)
-
-    DeepseekV3Linear.apply_linear = apply_linear
-    try:
-        yield calls
-    finally:
-        DeepseekV3Linear.apply_linear = own
-
-
-def _oracle(p, weight):
-    return U.fp64_oracle(p.residual, p.gamma, weight, p.eps).to(torch.bfloat16)
-
-
-@contextlib.contextmanager
-def _oracle_projection(h, packets):
-    """The native step with ``kva`` / ``kw`` replaced by the correctly rounded
-    FP64 oracle of each layer's packet."""
-    f = torch.nn.functional
-    orig_linear = f.linear
-    kva = {
-        layer: _oracle(packets[layer], h.attn[layer].kv_a_proj_with_mqa.weight) for layer in LAYERS
-    }
-    kw = {}
-    for layer in LAYERS:
-        indexer = h.attn[layer].mqa.indexer
-        if indexer is not None:
-            w = indexer._fused_wk_wp_weight
-            kw[id(w)] = (layer, _oracle(packets[layer], w))
-    torch.cuda.synchronize()
-    calls = {f"kva{layer}": 0 for layer in LAYERS}
-    calls.update({f"kw{layer}": 0 for layer, _ in kw.values()})
-
-    def make_forward(layer):
-        def forward(input, *args, **kwargs):
-            ref = kva[layer]
-            assert input.shape[0] == ref.shape[0], (input.shape, ref.shape)
-            calls[f"kva{layer}"] += 1
-            return ref.clone()
-
-        return forward
-
-    def linear(input, weight, bias=None):
-        hit = kw.get(id(weight))
-        if hit is not None and bias is None:
-            layer, ref = hit
-            assert input.shape[0] == ref.shape[0], (input.shape, ref.shape)
-            calls[f"kw{layer}"] += 1
-            return ref.clone()
-        return orig_linear(input, weight, bias)
-
-    for layer in LAYERS:
-        h.attn[layer].kv_a_proj_with_mqa.forward = make_forward(layer)
-    f.linear = linear
-    try:
-        yield calls
-    finally:
-        f.linear = orig_linear
-        for layer in LAYERS:
-            del h.attn[layer].kv_a_proj_with_mqa.forward
-
-
-def _projection_errors(h, layer, p, inputs):
-    """``e`` of ``kva`` / ``kw`` per arm on one layer (projection level, vs
-    the FP64 oracle). ``inputs``: arm -> that arm's normalized hidden ``h``."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-    from tensorrt_llm._torch.modules.linear import Linear
-
-    attn = h.attn[layer]
-    lin = attn.kv_a_proj_with_mqa
-    indexer = attn.mqa.indexer
-    w_kw = indexer._fused_wk_wp_weight if indexer is not None else None
-    oracle = {"kva": U.fp64_oracle(p.residual, p.gamma, lin.weight, p.eps)}
-    if w_kw is not None:
-        oracle["kw"] = U.fp64_oracle(p.residual, p.gamma, w_kw, p.eps)
-    ok, reason, kw_weight = projection_eligibility(attn)
-    assert ok, reason
-    proj = DeferredNormProjection(attn, kw_weight=kw_weight)
-    assert proj.supported() is None, proj.supported()
-    kva_w, kw_w = proj.project(p)
-    arms = {"deferred_norm": {"kva": kva_w, "kw": kw_w}}
-    for arm, hid in inputs.items():
-        arms[arm] = {"kva": lin(hid)}
-        if w_kw is not None:
-            arms[arm]["kw"] = torch.nn.functional.linear(hid, w_kw)
-    arms["n2_cublas"] = {"kva": Linear.apply_linear(lin, inputs["native"], None)}
-    if w_kw is not None:
-        arms["n2_cublas"]["kw"] = torch.nn.functional.linear(inputs["native"], w_kw)
-    arms["ref"] = {k: v.to(torch.bfloat16) for k, v in oracle.items()}
-    torch.cuda.synchronize()
-    return {
-        arm: {
-            name: U.norm_rms_err(y, oracle[name])
-            for name, y in ys.items()
-            if name in oracle and y is not None
-        }
-        for arm, ys in arms.items()
-    }
-
-
-def _input_stats(inputs):
-    """Fraction of elements of each arm's ``h`` that differ from native and
-    their max bf16 ulp distance."""
-    out = {}
-    for arm, per_layer in inputs.items():
-        if arm == "native":
-            continue
-        out[arm] = {}
-        for layer, hid in per_layer.items():
-            ulp = U.bf16_ulp_distance(hid, inputs["native"][layer])
-            out[arm][layer] = {
-                "frac_diff": float((ulp > 0).float().mean()),
-                "max_ulp": int(ulp.max()),
-            }
-    return out
-
-
-def _vs_ref(arm_out, ref_out):
-    return {layer: aq_integration._rel_rms(arm_out[layer], ref_out[layer]) for layer in LAYERS}
-
-
-@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
-@pytest.mark.parametrize(
-    "features",
-    [
-        "absorbed_query,post_moe_hop,deferred_norm",
-        "absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm",
-    ],
-)
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_output_bound_with_native_controls(harness, layout, path, features, hot):
-    h = harness
-    base = ",".join(f for f in features.split(",") if f != "deferred_norm")
-    seed = 3 if not hot else 13
-    packets, pos = _packets(h, layout, seed, hot)
-    native_in = {layer: p.materialize() for layer, p in packets.items()}
-    inputs = {
-        "native": native_in,
-        "n1_rmsnorm": _rmsnorm_inputs(packets),
-        "n3_round": dn_integration._control_inputs(packets),
-        "n4_dither1": {layer: _dither(v, 0.01, 1000 + layer) for layer, v in native_in.items()},
-        "n4_dither50": {layer: _dither(v, 0.5, 2000 + layer) for layer, v in native_in.items()},
-    }
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHidden,
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        frag_from_rows,
-    )
-
-    neg = {
-        layer: DeferredNormHidden(
-            frag_from_rows(p.residual), p.sumsq, p.residual, p.gamma, p.eps, materialize
-        )
-        if layer in dn_integration.SERVED_LAYERS
-        else native_in[layer]
-        for layer, p in packets.items()
-    }
-    snap = h.snapshot()
-    runs = {}
-    try:
-
-        def native_run(hidden):
-            dn_integration._reset_runners(h)
-            return dn_integration._run(h, base, layout, path, hidden, pos, snap)
-
-        runs["native"] = native_run(native_in)
-        dn_integration._reset_runners(h)
-        runs["deferred_norm"] = dn_integration._run(
-            h, features, layout, path, dn_integration._served(packets), pos, snap
-        )
-        for layer in LAYERS:
-            engaged = h.attn[layer]._deferred_norm_projection not in (None, False)
-            assert engaged == (layer in dn_integration.SERVED_LAYERS), f"deferred_norm on L{layer}"
-        runs["rerun"] = native_run(native_in)
-        for arm in ("n1_rmsnorm", "n3_round", "n4_dither1", "n4_dither50"):
-            runs[arm] = native_run(inputs[arm])
-        with _cublas_kv_a() as calls:
-            runs["n2_cublas"] = native_run(native_in)
-        assert calls[0] == len(LAYERS), f"cuBLAS kv_a path ran {calls[0]} times"
-        with _oracle_projection(h, packets) as calls:
-            runs["ref"] = native_run(native_in)
-        assert all(v == 1 for v in calls.values()), calls
-        dn_integration._reset_runners(h)
-        runs["neg"] = dn_integration._run(h, features, layout, path, neg, pos, snap)
-
-        report = {
-            "layout": layout,
-            "path": path,
-            "features": features,
-            "data": "hot" if hot else "gauss",
-            "arms": {},
-        }
-        ref_out = runs["ref"][0]
-        for arm, run in runs.items():
-            for layer in LAYERS:
-                assert torch.isfinite(run[0][layer].float()).all(), (arm, layer)
-            if arm == "native":
-                continue
-            stats = dn_integration._output_stats(runs["native"], run, layout)
-            for layer, d in _vs_ref(run[0], ref_out).items():
-                stats[layer]["out_rel_rms_vs_ref"] = d
-                stats[layer]["original_output_bound_ok"] = bool(
-                    stats[layer]["out_rel_rms"] <= OUTPUT_BOUND
-                    and stats[layer].get("topk_mean_jaccard", 1.0) >= 0.98
-                )
-            report["arms"][arm] = stats
-        report["arms"]["native"] = {
-            layer: {"out_rel_rms_vs_ref": d}
-            for layer, d in _vs_ref(runs["native"][0], ref_out).items()
-        }
-        report["inputs"] = _input_stats(inputs)
-        print(
-            "GLM_MOE_DSA_DEFERRED_NORM_CONTROLS " + json.dumps(report, sort_keys=True), flush=True
-        )
-        out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_CONTROLS_OUT")
-        if out_dir:
-            os.makedirs(out_dir, exist_ok=True)
-            name = f"controls-{layout}-{path}-{features.replace(',', '_')}-{report['data']}.json"
-            with open(os.path.join(out_dir, name), "w") as fh:
-                json.dump(report, fh, indent=1, sort_keys=True)
-
-        for layer in dn_integration.SERVED_LAYERS:
-            native_max = max(report["arms"][a][layer]["out_rel_rms"] for a in NATIVE_CONTROLS)
-            assert report["arms"]["neg"][layer]["out_rel_rms"] > native_max, (
-                f"L{layer}: the negative control is not separated from the native controls"
-            )
-    finally:
-        h.restore(snap)
-        dn_integration._set_features("")
-
-
-@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_projection_error_of_every_arm(harness, layout, hot):
-    """The projection error (``e`` vs the FP64 oracle) of deferred_norm and of every
-    native-only input / GEMM variant of the diagnosis."""
-    h = harness
-    packets, _ = _packets(h, layout, 3 if not hot else 13, hot)
-    native_in = {layer: p.materialize() for layer, p in packets.items()}
-    per_arm_inputs = {
-        "native": native_in,
-        "n1_rmsnorm": _rmsnorm_inputs(packets),
-        "n3_round": dn_integration._control_inputs(packets),
-        "n4_dither1": {layer: _dither(v, 0.01, 1000 + layer) for layer, v in native_in.items()},
-        "n4_dither50": {layer: _dither(v, 0.5, 2000 + layer) for layer, v in native_in.items()},
-    }
-    errs = {
-        layer: _projection_errors(
-            h, layer, packets[layer], {arm: v[layer] for arm, v in per_arm_inputs.items()}
-        )
-        for layer in dn_integration.SERVED_LAYERS
-    }
-    report = {
-        "layout": layout,
-        "data": "hot" if hot else "gauss",
-        "e": errs,
-        "inputs": _input_stats(per_arm_inputs),
-    }
-    print(
-        "GLM_MOE_DSA_DEFERRED_NORM_CONTROLS_PROJ " + json.dumps(report, sort_keys=True), flush=True
-    )
-    for layer, arms in errs.items():
-        for arm, d in arms.items():
-            for name, e in d.items():
-                assert e == e and e < 0.05, (layer, arm, name, e)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py
deleted file mode 100644
index 22a4391d31da..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_hop_mpi.py
+++ /dev/null
@@ -1,835 +0,0 @@
-# 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.
-"""Collective tests of the deferred_norm deferred hop (8 ranks over 2 hosts, TP8).
-
-Run every test file process under ``srun --mpi=pmix`` (one process per GPU);
-every test is collective and runs in the same order on every rank.
-
-Exact invariants (the deferred_norm hop against the post_moe_hop on the same raw tuple and
-the same workspace):
-
-* ``residual_out`` bit-identical to post_moe_hop's;
-* the fragment-layout ``a_frag`` holds ``bf16(x * gamma)`` bitwise (torch);
-* ``sumsq`` bit-identical to the torch emulation of post_moe_hop's per-thread /
-  butterfly order;
-* ``materialize(packet)`` bit-identical to post_moe_hop's normalized output (this
-  checks that the deferred ``r`` equals post_moe_hop's ``r`` bit for bit).
-
-Workspace protocol shared with post_moe_hop (``post_moe_hop/kernels.py``; post_moe_hop's
-error handling), with deferred_norm and post_moe_hop hops on one workspace:
-
-* a rank that never joins makes every other rank's deferred_norm call time out after
-  the post_moe_hop's deadline, record the Lamport-timeout bit and poll count, and
-  return NaN (residual, activation, square sums, materialized output); post_moe_hop's
-  health check fails, and after a collective reset the next call is exact;
-* an error word left by an earlier call poisons every later deferred_norm and post_moe_hop call
-  of the group (NaN) and is carried from buffer to buffer until the reset;
-* a permutation entry at or beyond the FC2 row count is skipped (output as
-  post_moe_hop's) and reported (bit 4);
-* scrambled store order (fault-injected deferred_norm and post_moe_hop builds, chosen per rank
-  and per call, mixed with plain builds and random device delays) changes
-  no output bit;
-* contract violations raise ``PostMoeHopContractError`` before any device work;
-  layouts the kernel cannot read with 16-byte vectors are repaired on the
-  rank that has them (counted in the workspace), also in a CUDA graph;
-* the health check counts deferred_norm calls as hop calls.
-"""
-
-import functools
-import os
-import random
-import sys
-import time
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
-
-import deferred_norm_utils as U  # noqa: E402
-import post_moe_hop_mpi_utils as M  # noqa: E402
-
-
-def _env():
-    env = M.mpi_env()
-    if env.world != 8:
-        pytest.skip("the deferred_norm hop is specialized to TP8")
-    return env
-
-
-_HOPS = {}
-
-
-def _hops(trigger_early=True):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    if "post_moe_hop" not in _HOPS:
-        _HOPS["post_moe_hop"] = PostMoeHop(
-            env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False
-        )
-    post_moe_hop = _HOPS["post_moe_hop"]
-    key = ("deferred_norm", trigger_early)
-    if key not in _HOPS:
-        _HOPS[key] = DeferredNormHop(post_moe_hop, trigger_early=trigger_early)
-    return post_moe_hop, _HOPS[key]
-
-
-def _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
-    """This rank's packet invariants against post_moe_hop's outputs (not collective)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    torch.cuda.synchronize()
-    ok_res = torch.equal(res_deferred_norm, res_post_moe_hop)
-    scaled_ref = (res_deferred_norm.float() * gamma.float()).bfloat16()
-    ok_scaled = torch.equal(
-        rows_from_frag(packet.a_frag, res_deferred_norm.shape[0], res_deferred_norm.shape[1]),
-        scaled_ref,
-    )
-    ok_ss = torch.equal(packet.sumsq, U.hop_sumsq_partials(res_deferred_norm))
-    h = materialize(packet)
-    torch.cuda.synchronize()
-    ok_h = torch.equal(h, h_post_moe_hop)
-    ok = ok_res and ok_scaled and ok_ss and ok_h
-    if not ok:
-        print(
-            f"GLM_MOE_DSA_DEFERRED_NORM_HOP_FAIL {tag} rank={M.mpi_env().rank} res={ok_res} scaled={ok_scaled} "
-            f"sumsq={ok_ss} materialize={ok_h} "
-            f"max_ulp_h={int(U.bf16_ulp_distance(h, h_post_moe_hop).max())}",
-            flush=True,
-        )
-    return ok
-
-
-def _check(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
-    """Collective: ``_verdict`` on every rank."""
-    return M.all_ranks(
-        _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag)
-    )
-
-
-def _all_nan_packet(packet, res):
-    """Every value a consumer can derive from the packet is NaN."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    torch.cuda.synchronize()
-    a = rows_from_frag(packet.a_frag, res.shape[0], res.shape[1])
-    h = materialize(packet)
-    torch.cuda.synchronize()
-    return all(bool(torch.isnan(t.float()).all()) for t in (res, a, packet.sumsq, h))
-
-
-def _reset(post_moe_hop):
-    assert M.all_ranks(post_moe_hop.reset())
-    assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-def test_capability_engaged():
-    post_moe_hop, deferred_norm = _hops()
-    assert M.all_ranks(post_moe_hop.engaged and deferred_norm.engaged), deferred_norm.reason
-
-
-def test_hop_precompiles_served_projections():
-    """An engaged hop has compiled both served projection kernels on every
-    rank (the consumer's compile is a cache hit)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.proj_kernel import (  # noqa: E501
-        is_compiled,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        SERVED_N1,
-        served_kernel_configs,
-    )
-
-    _, deferred_norm = _hops()
-    cfgs = served_kernel_configs(torch.cuda.current_device())
-    ok = (
-        deferred_norm.engaged
-        and sorted(c["n1"] for c in cfgs) == sorted(SERVED_N1)
-        and all(is_compiled(**c) for c in cfgs)
-    )
-    assert M.all_ranks(ok)
-
-
-@pytest.mark.parametrize("failing_rank", [3, 7])
-def test_projection_compile_failure_disengages_every_rank(monkeypatch, failing_rank):
-    """A projection compile error on one rank disengages the hop on all
-    ranks (collective decision), so no rank produces a packet."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-        projection,
-    )
-
-    env = _env()
-    post_moe_hop, _ = _hops()
-    if env.rank == failing_rank:
-
-        def _fail(device):
-            raise RuntimeError("injected projection compile failure")
-
-        monkeypatch.setattr(projection, "precompile_served", _fail)
-    hop = DeferredNormHop(post_moe_hop)
-    if env.rank == failing_rank:
-        local = "injected projection compile failure" in hop.reason
-    else:
-        local = hop.reason == "another rank is not capable"
-    assert M.all_ranks(not hop.engaged and local), hop.reason
-
-
-def test_unsupported_projection_disengages_every_rank(monkeypatch):
-    """A served projection shape the kernel does not support (static,
-    rank-identical) disengages the hop on every rank."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501  # noqa: E501
-        DeferredNormHop,
-        projection,
-    )
-
-    post_moe_hop, _ = _hops()
-    monkeypatch.setattr(
-        projection,
-        "DeferredNormProjConfig",
-        functools.partial(projection.DeferredNormProjConfig, smem_rows=8),
-    )
-    hop = DeferredNormHop(post_moe_hop)
-    assert M.all_ranks(not hop.engaged and "staged rows" in hop.reason), hop.reason
-
-
-@pytest.mark.parametrize("trigger_early", [True, False])
-@pytest.mark.parametrize("scale_dtype", [torch.bfloat16, torch.float32])
-@pytest.mark.parametrize("with_shared", [True, False])
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-def test_deferred_matches_post_moe_hop(rows, with_shared, scale_dtype, trigger_early):
-    post_moe_hop, deferred_norm = _hops(trigger_early)
-    with torch.inference_mode():
-        fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(
-            rows, seed=40 + rows, scale_dtype=scale_dtype, with_shared=with_shared
-        )
-        h_post_moe_hop, res_post_moe_hop = post_moe_hop(
-            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
-        )
-        packet, res_deferred_norm = deferred_norm(
-            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
-        )
-        assert _check(
-            packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, f"rows={rows}"
-        )
-
-
-def test_empty_rows_and_large_residual():
-    post_moe_hop, deferred_norm = _hops()
-    with torch.inference_mode():
-        fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(6, seed=77, empty_rows=(1, 4))
-        residual = (residual.float() * 50.0).bfloat16()
-        h_post_moe_hop, res_post_moe_hop = post_moe_hop(
-            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
-        )
-        packet, res_deferred_norm = deferred_norm(
-            fc2, scales, perm, shared, residual, gamma, M.GLM_EPS
-        )
-        assert _check(
-            packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, "empty-rows"
-        )
-
-
-def test_back_to_back_mixed_with_post_moe_hop():
-    """deferred_norm and post_moe_hop hops interleaved on one workspace without host syncs."""
-    post_moe_hop, deferred_norm = _hops()
-    rows_seq = [6, 1, 8, 3, 6, 2, 6, 8, 5, 6, 1, 6]
-    with torch.inference_mode():
-        inputs = [M.make_raw_tuple(r, seed=900 + i) for i, r in enumerate(rows_seq)]
-        outs = []
-        for i, t in enumerate(inputs):
-            if i % 2 == 0:
-                outs.append(("deferred_norm", deferred_norm(*t, M.GLM_EPS)))
-            else:
-                outs.append(("post_moe_hop", post_moe_hop(*t, M.GLM_EPS)))
-        torch.cuda.synchronize()
-        # Reference: each call again, alone, through post_moe_hop.
-        failures = []
-        for i, (t, (kind, out)) in enumerate(zip(inputs, outs)):
-            h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-            if kind == "post_moe_hop":
-                ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
-            else:
-                ok = _check(out[0], out[1], h_ref, res_ref, t[5], f"b2b{i}")
-            if not ok:
-                failures.append(i)
-        assert M.all_ranks(not failures), failures
-        assert post_moe_hop.error_bits() == 0
-
-
-def test_cuda_graph_replay():
-    post_moe_hop, deferred_norm = _hops()
-    with torch.inference_mode():
-        static = [M.make_raw_tuple(6, seed=1200 + i) for i in range(3)]
-        bufs = [[x.clone() if x is not None else None for x in t] for t in static]
-
-        def body():
-            a = deferred_norm(*bufs[0], M.GLM_EPS)
-            b = post_moe_hop(*bufs[1], M.GLM_EPS)
-            c = deferred_norm(*bufs[2], M.GLM_EPS)
-            return a, b, c
-
-        body()
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            ga, gb, gc = body()
-        for trial, replays in enumerate((1, 1, 20)):
-            fresh = [M.make_raw_tuple(6, seed=1300 + 10 * trial + i) for i in range(3)]
-            for dst, src in zip(bufs, fresh):
-                for d, s in zip(dst, src):
-                    if d is not None:
-                        d.copy_(s)
-            for _ in range(replays):
-                M.mpi_env().comm.barrier()
-                graph.replay()
-            torch.cuda.synchronize()
-            ok = True
-            for (kind, got), t in zip(
-                (("deferred_norm", ga), ("post_moe_hop", gb), ("deferred_norm", gc)), fresh
-            ):
-                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-                if kind == "post_moe_hop":
-                    ok &= torch.equal(got[0], h_ref) and torch.equal(got[1], res_ref)
-                else:
-                    ok &= _check(got[0], got[1], h_ref, res_ref, t[5], f"graph{trial}")
-            assert M.all_ranks(ok), f"trial {trial}"
-
-
-def test_delayed_rank_stress():
-    """Random per-rank device delays inside back-to-back batches (no host sync)."""
-    post_moe_hop, deferred_norm = _hops()
-    env = M.mpi_env()
-    gen = torch.Generator().manual_seed(4242)
-    with torch.inference_mode():
-        for batch in range(4):
-            inputs = []
-            outs = []
-            for i in range(12):
-                rows = int(torch.randint(1, 9, (1,), generator=gen))
-                t = M.make_raw_tuple(rows, seed=5000 + 100 * batch + i)
-                local = torch.Generator().manual_seed(77 * batch + 13 * i + env.rank)
-                if float(torch.rand(1, generator=local)) < 0.3:
-                    torch.cuda._sleep(int(float(torch.rand(1, generator=local)) * 2_000_000))
-                inputs.append(t)
-                outs.append(deferred_norm(*t, M.GLM_EPS))
-            torch.cuda.synchronize()
-            ok = True
-            for t, (packet, res) in zip(inputs, outs):
-                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-                ok &= _check(packet, res, h_ref, res_ref, t[5], f"stress{batch}")
-            assert M.all_ranks(ok), f"batch {batch}"
-        assert post_moe_hop.error_bits() == 0
-
-
-@pytest.mark.parametrize("layer", [10, 76])
-def test_real_weights(layer):
-    """Real routed experts / router / shared expert / next-layer gain (synthetic activations)."""
-    from post_moe_hop_real_fixture import build_real_moe_layer
-
-    post_moe_hop, deferred_norm = _hops()
-    env = M.mpi_env()
-    with torch.inference_mode():
-        real = build_real_moe_layer(env.mapping, layer, max_num_tokens=64)
-        ok = True
-        for rows in (1, 6, 8):
-            x = real.hidden_states(rows, seed=31 + rows)
-            logits = real.gate(x)
-            shared = real.shared_output(x)
-            residual = real.residual(rows, seed=41 + rows)
-            fc2, scales, perm = real.run_moe(x, logits, False)
-            h_post_moe_hop, res_post_moe_hop = post_moe_hop(
-                fc2, scales, perm, shared, residual, real.next_gamma, M.GLM_EPS
-            )
-            packet, res_deferred_norm = deferred_norm(
-                fc2, scales, perm, shared, residual, real.next_gamma, M.GLM_EPS
-            )
-            ok &= _check(
-                packet,
-                res_deferred_norm,
-                h_post_moe_hop,
-                res_post_moe_hop,
-                real.next_gamma,
-                f"L{layer}R{rows}",
-            )
-        real.moe.destroy()
-        assert M.all_ranks(ok)
-
-
-# ---------------------------------------------------------------------------
-# Workspace protocol shared with post_moe_hop: bounded errors, sticky poison, permutation
-# range, scrambled store order, raw-tuple contract and layout repairs.
-# ---------------------------------------------------------------------------
-
-_MISSING_TIMEOUT_S = 0.2
-
-
-def _impatient():
-    """A post_moe_hop with a short poll deadline on the same workspace, and the deferred_norm
-    hop built on it (collective; the deferred_norm kernels are compile-cache hits)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    if "impatient" not in _HOPS:
-        post_moe_hop = PostMoeHop(
-            env.mapping,
-            M.GLM_HIDDEN,
-            M.GLM_TOP_K,
-            require_multi_node=False,
-            timeout_s=_MISSING_TIMEOUT_S,
-        )
-        _HOPS["impatient"] = (post_moe_hop, DeferredNormHop(post_moe_hop))
-    return _HOPS["impatient"]
-
-
-@pytest.mark.parametrize("missing", [0, 5, 7])
-def test_missing_rank_times_out_with_nan_and_reset(missing):
-    """A rank that never joins makes every other rank's deferred_norm call time out,
-    report and return NaN instead of hanging; a collective reset recovers."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-
-    post_moe_hop, deferred_norm = _hops()
-    imp_post_moe_hop, imp_deferred_norm = _impatient()
-    env = M.mpi_env()
-    assert M.all_ranks(
-        imp_deferred_norm.engaged and imp_post_moe_hop.workspace is post_moe_hop.workspace
-    )
-    assert M.all_ranks(imp_deferred_norm.timeout_ns == int(_MISSING_TIMEOUT_S * 1e9))
-    rows = 6
-    with torch.inference_mode():
-        tensors = M.make_raw_tuple(rows, seed=600 + missing)
-        before = M.lamport_snapshot(post_moe_hop)
-        assert M.all_ranks(before["ok"]), before
-        env.comm.barrier()
-        start = time.monotonic()
-        packet = res = None
-        if env.rank != missing:
-            packet, res = imp_deferred_norm(*tensors, M.GLM_EPS)
-        torch.cuda.synchronize()
-        elapsed = time.monotonic() - start
-        counters = post_moe_hop.workspace.counters()
-        bits, polls = counters["error_bits"], counters["timeout_polls"]
-        if env.rank == missing:
-            reported = bits == 0 and polls == 0
-        else:
-            reported = (
-                bool(bits & 1)
-                and polls > 0
-                and _MISSING_TIMEOUT_S <= elapsed < 30.0
-                and _all_nan_packet(packet, res)
-            )
-        stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
-        if env.rank == 0:
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_MISSING rank {missing} skipped: (rank, elapsed_s, error_bits, "
-                f"timeout_polls, ok) {stats}",
-                flush=True,
-            )
-        assert M.all_ranks(reported), stats
-        try:
-            check_health(f"deferred_norm-missing-rank{missing}-local")
-            local_raised = False
-        except PostMoeHopError:
-            local_raised = True
-        try:
-            check_health(f"deferred_norm-missing-rank{missing}-collective", collective=True)
-            group_raised = False
-        except PostMoeHopError:
-            group_raised = True
-        assert M.all_ranks(local_raised == (env.rank != missing) and group_raised)
-        _reset(post_moe_hop)
-        t = M.make_raw_tuple(rows, seed=650 + missing)
-        packet, res = deferred_norm(*t, M.GLM_EPS)
-        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-        assert _check(packet, res, h_ref, res_ref, t[5], f"after-reset{missing}")
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize(
-    "sequence",
-    [
-        "deferred_norm,deferred_norm,deferred_norm",
-        "deferred_norm,post_moe_hop,deferred_norm",
-        "post_moe_hop,deferred_norm,post_moe_hop",
-    ],
-)
-def test_error_poisons_every_rank_until_reset(sequence):
-    """An error word left by an earlier call on one rank makes the next calls
-    of every rank NaN, deferred_norm and post_moe_hop alike, and is carried through all three
-    Lamport buffers until the collective reset."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
-        STATE_EPOCH,
-        STATE_ERROR_WORDS,
-    )
-
-    post_moe_hop, deferred_norm = _hops()
-    env = M.mpi_env()
-    victim = env.world - 2
-    with torch.inference_mode():
-        torch.cuda.synchronize()
-        if env.rank == victim:
-            # As if this rank's previous call had timed out.
-            epoch = int(post_moe_hop.workspace.state[STATE_EPOCH].item())
-            post_moe_hop.workspace.state[STATE_ERROR_WORDS[(epoch + 2) % 3]] = 1
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        poisoned = True
-        for call, kind in enumerate(sequence.split(",")):
-            t = M.make_raw_tuple(6, seed=860 + call)
-            if kind == "deferred_norm":
-                packet, res = deferred_norm(*t, M.GLM_EPS)
-                poisoned &= _all_nan_packet(packet, res)
-            else:
-                out, res = post_moe_hop(*t, M.GLM_EPS)
-                torch.cuda.synchronize()
-                poisoned &= bool(torch.isnan(out.float()).all() and torch.isnan(res.float()).all())
-        words = post_moe_hop.workspace.state.cpu().tolist()
-        carried = all(words[w] == 1 for w in STATE_ERROR_WORDS)
-        clean = not any(words[w] for w in STATE_ERROR_WORDS)
-        if env.rank == 0:
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_POISON {sequence}: per-rank error words "
-                f"{env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])}",
-                flush=True,
-            )
-        else:
-            env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])
-        assert M.all_ranks(poisoned and (carried if env.rank == victim else clean))
-        with pytest.raises(PostMoeHopError):
-            check_health("deferred_norm-poisoned", collective=True)
-        _reset(post_moe_hop)
-        t = M.make_raw_tuple(6, seed=870)
-        packet, res = deferred_norm(*t, M.GLM_EPS)
-        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-        assert _check(packet, res, h_ref, res_ref, t[5], "after-poison-reset")
-
-
-def test_permutation_out_of_range_is_skipped_and_reported():
-    """Entries at or beyond the FC2 row count are skipped (as by post_moe_hop and the
-    torch oracle) and set error bit 4 on the rank that has them. FC2 holds
-    exactly the referenced rows, so under memcheck without the caching
-    allocator a read of a skipped entry would be an invalid access."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
-        ERROR_PERM_RANGE,
-    )
-
-    post_moe_hop, deferred_norm = _hops()
-    env = M.mpi_env()
-    culprit = 2 % env.world
-    rows = 6
-    with torch.inference_mode():
-        t = list(M.make_raw_tuple(rows, seed=880, padded_rows=rows * M.GLM_TOP_K))
-        if env.rank == culprit:
-            t[2] = t[2].clone()
-            t[2][1, 3] = t[0].shape[0]
-            t[2][4, 0] = t[0].shape[0] + 1000
-        packet, res = deferred_norm(*t, M.GLM_EPS)
-        torch.cuda.synchronize()
-        bits = post_moe_hop.error_bits()
-        reported = bits == (ERROR_PERM_RANGE if env.rank == culprit else 0)
-        with pytest.raises(PostMoeHopError):
-            check_health("deferred_norm-perm-range", collective=True)
-        _reset(post_moe_hop)
-        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-        torch.cuda.synchronize()
-        ok = _verdict(packet, res, h_ref, res_ref, t[5], "perm-range") and reported
-        # The oracle skips entries outside [0, N_padded) like the kernels.
-        ref_local = M.sanitize_bf16(M.reference_local_partial(t[0], t[1], t[2], t[3]))
-        slots = M.gather_bf16(ref_local).permute(1, 0, 2).cuda()
-        _, emu_res = M.emulate_post_reduction(slots, t[4], t[5], M.GLM_EPS)
-        ok &= torch.equal(res.view(torch.int16), emu_res.view(torch.int16))
-        assert M.all_ranks(ok), bits
-        _reset(post_moe_hop)
-
-
-_FAULT_DELAYS_NS = (3000, 60000)
-
-
-def _fault_hops():
-    """post_moe_hop and deferred_norm hops with fault-injected kernel builds on the same workspace."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    post_moe_hop, deferred_norm = _hops()
-    if "fault" not in _HOPS:
-        hops = {"post_moe_hop": post_moe_hop, "deferred_norm": deferred_norm}
-        for d in _FAULT_DELAYS_NS:
-            fp6 = PostMoeHop(
-                env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False, fault_delay_ns=d
-            )
-            hops[f"post_moe_hop_fault{d}"] = fp6
-            hops[f"deferred_norm_fault{d}"] = DeferredNormHop(fp6)
-        _HOPS["fault"] = hops
-    return _HOPS["fault"]
-
-
-def test_fault_injected_store_order_mixed_with_post_moe_hop():
-    """Scrambled store order inside the kernels (fault-injected builds), a
-    per-rank choice of deferred_norm or post_moe_hop kernel for the same collective call, and
-    random per-rank device delays: every output bit as in an undisturbed post_moe_hop
-    call, the call counter advances once per call, no error bit."""
-    hops = _fault_hops()
-    post_moe_hop = hops["post_moe_hop"]
-    env = M.mpi_env()
-    names = sorted(hops)
-    assert M.all_ranks(all(h.engaged for h in hops.values()))
-    common = random.Random(4711)
-    local = random.Random(1000 + env.rank)
-    calls_before = post_moe_hop.workspace.calls()
-    failed, used, total = [], {}, 0
-    with torch.inference_mode():
-        for batch in range(6):
-            plan = [
-                (
-                    common.randint(1, 8),
-                    common.random() < 0.75,
-                    torch.float32 if common.random() < 0.2 else torch.bfloat16,
-                )
-                for _ in range(12)
-            ]
-            inputs = [
-                M.make_raw_tuple(r, seed=9000 + 100 * batch + j, with_shared=sh, scale_dtype=sd)
-                for j, (r, sh, sd) in enumerate(plan)
-            ]
-            picks = [local.choice(names) for _ in plan]
-            torch.cuda.synchronize()
-            env.comm.barrier()
-            outs = []
-            for name, t in zip(picks, inputs):
-                if local.random() < 0.3:
-                    torch.cuda._sleep(int(local.uniform(0.0, 50.0) * 2000))
-                used[name] = used.get(name, 0) + 1
-                outs.append(hops[name](*t, M.GLM_EPS))
-            torch.cuda.synchronize()
-            for j, (name, t, out) in enumerate(zip(picks, inputs, outs)):
-                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-                if name.startswith("deferred_norm"):
-                    ok = _verdict(out[0], out[1], h_ref, res_ref, t[5], f"fault{batch}.{j}")
-                else:
-                    torch.cuda.synchronize()
-                    ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
-                if not ok:
-                    failed.append((batch, j, name))
-            total += len(plan)
-        counted = post_moe_hop.workspace.calls() - calls_before
-    bits = post_moe_hop.error_bits()
-    if env.rank == 0:
-        print(
-            f"GLM_MOE_DSA_DEFERRED_NORM_FAULT calls={total} (+{total} reference calls) used_on_rank0={used}",
-            flush=True,
-        )
-    assert M.all_ranks(not failed and counted == 2 * total and bits == 0), (
-        failed[:4],
-        counted,
-        bits,
-    )
-
-
-@pytest.mark.parametrize(
-    "case",
-    ("fc2_dtype", "scale_shape", "perm_dtype", "residual_shape", "shared_dtype", "gamma_shape"),
-)
-def test_contract_violation_raises_before_any_device_work(case):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopContractError,
-    )
-
-    post_moe_hop, deferred_norm = _hops()
-    rows = 6
-    with torch.inference_mode():
-        fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(rows, seed=920)
-        if case == "fc2_dtype":
-            fc2 = fc2.half()
-        elif case == "scale_shape":
-            scales = scales[:, : M.GLM_TOP_K - 1]
-        elif case == "perm_dtype":
-            perm = perm.long()
-        elif case == "residual_shape":
-            residual = residual[:, : M.GLM_HIDDEN // 2]
-        elif case == "shared_dtype":
-            shared = shared.float()
-        else:
-            gamma = gamma[:-8]
-        calls_before = post_moe_hop.workspace.calls()
-        with pytest.raises(PostMoeHopContractError):
-            deferred_norm(fc2, scales, perm, shared, residual, gamma, M.GLM_EPS)
-        assert M.all_ranks(post_moe_hop.workspace.calls() == calls_before)
-        t = M.make_raw_tuple(rows, seed=921)
-        packet, res = deferred_norm(*t, M.GLM_EPS)
-        h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-        assert _check(packet, res, h_ref, res_ref, t[5], f"after-contract-{case}")
-
-
-def _misaligned(t, offset_elems):
-    storage = torch.empty(t.numel() + 16, dtype=t.dtype, device=t.device)
-    view = storage[offset_elems : offset_elems + t.numel()].view(t.shape)
-    view.copy_(t)
-    return view
-
-
-def _transposed(t):
-    view = torch.empty(tuple(reversed(t.shape)), dtype=t.dtype, device=t.device).t()
-    view.copy_(t)
-    return view
-
-
-def _padded_rows(t, extra_cols):
-    storage = torch.zeros((t.shape[0], t.shape[1] + extra_cols), dtype=t.dtype, device=t.device)
-    view = storage[:, : t.shape[1]]
-    view.copy_(t)
-    return view
-
-
-def _needs_copy(t, row_major_2d=False):
-    if t.data_ptr() % 16 or t.stride(-1) != 1:
-        return True
-    if row_major_2d:
-        return (t.stride(0) * t.element_size()) % 16 != 0
-    return not t.is_contiguous()
-
-
-def _relayout(rank, tensors):
-    """Rank-dependent layouts; returns (tensors, expected repairs on this rank)."""
-    fc2, scales, perm, shared, residual, gamma = tensors
-    kind = rank % 4
-    if kind == 0:
-        fc2 = _misaligned(fc2, 1)
-        perm = _misaligned(perm, 1)
-        gamma = _misaligned(gamma, 2)
-    elif kind == 1:
-        fc2 = _padded_rows(fc2, 4)
-        scales = _transposed(scales)
-        residual = _misaligned(residual, 4)
-    elif kind == 2:
-        fc2 = _padded_rows(fc2, 8)  # aligned row pitch: read in place
-        if shared is not None:
-            shared = _transposed(shared)
-    repairs = int(_needs_copy(fc2, row_major_2d=True)) + sum(
-        int(_needs_copy(t)) for t in (scales, perm, shared, residual, gamma) if t is not None
-    )
-    return (fc2, scales, perm, shared, residual, gamma), repairs
-
-
-def test_unaligned_layouts_are_repaired_locally():
-    """Per-rank layout repairs keep the collective and every output bit."""
-    post_moe_hop, deferred_norm = _hops()
-    env = M.mpi_env()
-    ok = True
-    with torch.inference_mode():
-        for i, rows in enumerate((6, 1, 8)):
-            reference = M.make_raw_tuple(rows, seed=880 + i)
-            tensors, expected = _relayout(env.rank, reference)
-            before = post_moe_hop.workspace.layout_repairs
-            packet, res = deferred_norm(*tensors, M.GLM_EPS)
-            repaired = post_moe_hop.workspace.layout_repairs - before
-            h_ref, res_ref = post_moe_hop(*reference, M.GLM_EPS)
-            ok &= _verdict(packet, res, h_ref, res_ref, reference[5], f"relayout{rows}")
-            ok &= repaired == expected
-        # The repaired call inside a CUDA graph, replayed with fresh values.
-        rows = 6
-        static_ref = M.make_raw_tuple(rows, seed=890)
-        static, _ = _relayout(env.rank, static_ref)
-        deferred_norm(*static, M.GLM_EPS)
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            g_packet, g_res = deferred_norm(*static, M.GLM_EPS)
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        for replay in range(4):
-            fresh = M.make_raw_tuple(rows, seed=900 + replay, padded_rows=static_ref[0].shape[0])
-            for dst, src in zip(static, fresh):
-                if dst is not None:
-                    dst.copy_(src)
-            graph.replay()
-            torch.cuda.synchronize()
-            h_ref, res_ref = post_moe_hop(*fresh, M.GLM_EPS)
-            ok &= _verdict(g_packet, g_res, h_ref, res_ref, fresh[5], f"relayout-graph{replay}")
-    assert M.all_ranks(ok and post_moe_hop.error_bits() == 0)
-
-
-def test_health_check_counts_deferred_norm_calls():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        check_health,
-    )
-
-    post_moe_hop, deferred_norm = _hops()
-    with torch.inference_mode():
-        before = check_health("deferred_norm-engagement-before", collective=True)
-        for i in range(5):
-            deferred_norm(*M.make_raw_tuple(6, seed=940 + i), M.GLM_EPS)
-        after = check_health("deferred_norm-engagement-after", collective=True)
-    assert M.all_ranks(
-        after.calls - before.calls == 5
-        and after.error_bits == 0
-        and not after.group_error
-        and after.calls == post_moe_hop.workspace.calls()
-    )
-
-
-def test_memcheck_positive_control():
-    """Tests only, run under compute-sanitizer memcheck by the sanitizer
-    suite (``GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL=1``): the kernel is told FC2 has one more
-    row than it does and the permutation references that row, so memcheck
-    must report an invalid global read in the deferred_norm hop."""
-    if os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL") != "1":
-        pytest.skip("memcheck positive control (sanitizer suite only)")
-    post_moe_hop, deferred_norm = _hops()
-    rows = 1
-    with torch.inference_mode():
-        t = list(M.make_raw_tuple(rows, seed=990, padded_rows=rows * M.GLM_TOP_K))
-        t[2] = t[2].clone()
-        t[2][0, 0] = t[0].shape[0]
-        deferred_norm.launch(*t, M.GLM_EPS, fc2_rows=t[0].shape[0] + 1)
-        torch.cuda.synchronize()
-    _reset(post_moe_hop)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py
deleted file mode 100644
index c7471662e07c..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_null_controls.py
+++ /dev/null
@@ -1,559 +0,0 @@
-# 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.
-"""Null distribution of deferred_norm's attention-output distances among correct
-rounding variants of the same projection.
-
-Same harness, cases and inputs as ``test_deferred_norm_controls.py`` (real layer-6 /
-layer-7 weights at TP8 rank 0, six requests natively prefilled with 2,600
-tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6, both call paths, ``absorbed_query`` and
-``absorbed_query,shared_indexer_absorbed_query``, Gaussian and hot-channel residual rows).
-
-The null family ``S_k`` (``k = 1..K``) models every correct implementation
-of "RMSNorm, then the projection" with one bf16 rounding of the normalized
-input and one of the output: native code whose projection outputs (``kva``
-and, on the full-indexer layer, ``kw``) are ``bf16(W . h_k)`` with ``h_k =
-h* + u_k * ulp(h*)``, where ``h* = x * gamma * r*`` is the exact normalized
-hidden state (FP64), ``ulp`` the bf16 ulp at ``h*`` and ``u_k`` i.i.d.
-uniform on [-1/2, 1/2]. Native (``h = bf16((x * r) * gamma)``, then the
-GEMM) and deferred_norm (``a = bf16(x * gamma)``, the GEMM, then ``r``) are both of
-this form; their rounding errors are deterministic draws, the family's are
-random draws.
-
-Per case, all on the same cache snapshot:
-
-* ``native``, ``rerun``, ``deferred_norm``: as in ``test_deferred_norm_controls.py`` (the
-  deferred_norm projection on the layers it serves, ``dn_integration.SERVED_LAYERS``: the target's
-  deferred-norm scope is ``shared``).
-* ``ref``: native code on the correctly rounded FP64 oracle projection.
-* ``swap``: native code whose projection outputs on the served layers are deferred_norm's
-  projection kernel outputs (``DeferredNormProjection.project`` on the same packet). It
-  differs from the ``deferred_norm`` arm only in the code path around the projection.
-* ``S_1 .. S_K``.
-
-Reported per case, arm and layer: attention-output relRMS vs ``ref`` and vs
-``native`` (overall and per token row), FP8 query bytes and mean top-k
-Jaccard vs ``ref``; the null family's order statistics and the counts of
-draws at or above ``deferred_norm`` / ``native``.
-
-``test_projection_null`` reports, per data set, served layer and projection output, the
-projection-level error ``e = ||y - y*|| / ||y*||`` of native, deferred_norm and
-``K_PROJ`` draws of the family; per-segment errors; the bias and the
-largest error in units of the bf16 ulp at ``y*``; and each kernel against
-its own arithmetic specification evaluated in FP64 (native: ``bf16(W .
-h)``; deferred_norm: ``bf16(r . (W . a))`` with the hop-order ``r``).
-
-Assertions are harness facts only (hooks fired once per projection and run,
-deferred_norm engaged on the served layers only, finite outputs). The decision rule that uses these numbers was
-written before the data.
-"""
-
-import contextlib
-import json
-import os
-
-import deferred_norm_utils as U
-import pytest
-import test_absorbed_query_attention_integration as aq_integration
-import test_deferred_norm_attention_integration as dn_integration
-import test_deferred_norm_controls as controls
-import torch
-
-LAYERS = aq_integration.LAYERS
-LAYOUTS = aq_integration.LAYOUTS
-K_ATTN = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_K", "48"))
-K_PROJ = int(os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_K_PROJ", "1000"))
-SEED_BASE = 7000
-
-pytestmark = dn_integration.pytestmark
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = dn_integration._DeferredNormHarness()
-    yield h
-    dn_integration._set_features("")
-
-
-def _kw_weight(h, layer):
-    indexer = h.attn[layer].mqa.indexer
-    return None if indexer is None else indexer._fused_wk_wp_weight
-
-
-def _weights64(h):
-    """Per layer: name -> FP64 copy of the projection weight."""
-    out = {}
-    for layer in LAYERS:
-        out[layer] = {"kva": h.attn[layer].kv_a_proj_with_mqa.weight.double()}
-        w = _kw_weight(h, layer)
-        if w is not None:
-            out[layer]["kw"] = w.double()
-    return out
-
-
-def _h_star(p):
-    """Exact normalized hidden state ``x * gamma * r*`` (FP64)."""
-    xd = p.residual.double()
-    r = torch.rsqrt(xd.pow(2).mean(-1, keepdim=True) + p.eps)
-    return xd * p.gamma.double() * r
-
-
-def _bf16_ulp64(t):
-    """bf16 ulp at ``|t|`` (FP64; 0 where ``t == 0``)."""
-    _, e = torch.frexp(t)
-    ulp = torch.exp2((e - 8).to(torch.float64))
-    return torch.where(t == 0, torch.zeros_like(t), ulp)
-
-
-def _draw(hstar, ulp, w64, seed):
-    """One member of the null family: name -> ``bf16(W . (h* + u ulp))``."""
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    u = torch.rand(hstar.shape, generator=g, device="cuda", dtype=torch.float64) - 0.5
-    hk = hstar + u * ulp
-    return {name: (hk @ w.t()).to(torch.bfloat16) for name, w in w64.items()}
-
-
-def _seed(k, layer):
-    return SEED_BASE + 97 * k + layer
-
-
-@contextlib.contextmanager
-def _projection_override(h, kva, kw):
-    """Native step with ``kv_a_proj_with_mqa`` returning ``kva[layer]`` and the
-    indexer's fused ``wk|weights_proj`` GEMM returning ``kw[layer]``."""
-    f = torch.nn.functional
-    orig_linear = f.linear
-    by_weight = {id(_kw_weight(h, layer)): (layer, t) for layer, t in kw.items()}
-    calls = {f"kva{layer}": 0 for layer in kva}
-    calls.update({f"kw{layer}": 0 for layer in kw})
-
-    def make_forward(layer):
-        def forward(input, *args, **kwargs):
-            t = kva[layer]
-            assert input.shape[0] == t.shape[0], (input.shape, t.shape)
-            calls[f"kva{layer}"] += 1
-            return t.clone()
-
-        return forward
-
-    def linear(input, weight, bias=None):
-        hit = by_weight.get(id(weight))
-        if hit is not None and bias is None:
-            layer, t = hit
-            assert input.shape[0] == t.shape[0], (input.shape, t.shape)
-            calls[f"kw{layer}"] += 1
-            return t.clone()
-        return orig_linear(input, weight, bias)
-
-    for layer in kva:
-        h.attn[layer].kv_a_proj_with_mqa.forward = make_forward(layer)
-    f.linear = linear
-    try:
-        yield calls
-    finally:
-        f.linear = orig_linear
-        for layer in kva:
-            del h.attn[layer].kv_a_proj_with_mqa.forward
-
-
-def _split(per_layer):
-    kva = {layer: d["kva"] for layer, d in per_layer.items()}
-    kw = {layer: d["kw"] for layer, d in per_layer.items() if d.get("kw") is not None}
-    return kva, kw
-
-
-def _deferred_norm_projection_outputs(h, packets):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-
-    out = {}
-    for layer in dn_integration.SERVED_LAYERS:
-        attn = h.attn[layer]
-        ok, reason, kw_weight = projection_eligibility(attn)
-        assert ok, reason
-        proj = DeferredNormProjection(attn, kw_weight=kw_weight)
-        assert proj.supported() is None, proj.supported()
-        kva, kw = proj.project(packets[layer])
-        out[layer] = {"kva": kva, "kw": kw}
-    torch.cuda.synchronize()
-    return out
-
-
-def _row_rel(x, ref):
-    xd = x.double().reshape(x.shape[0], -1)
-    rd = ref.double().reshape(ref.shape[0], -1)
-    return ((xd - rd).pow(2).mean(1).sqrt() / rd.pow(2).mean(1).sqrt().clamp_min(1e-30)).tolist()
-
-
-def _order_stats(v):
-    s = sorted(v)
-    n = len(s)
-    return {
-        "min": s[0],
-        "p10": s[int(0.1 * (n - 1))],
-        "median": s[(n - 1) // 2],
-        "p90": s[int(0.9 * (n - 1))],
-        "max": s[-1],
-        "mean": sum(s) / n,
-    }
-
-
-@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
-@pytest.mark.parametrize(
-    "features",
-    [
-        "absorbed_query,post_moe_hop,deferred_norm",
-        "absorbed_query,post_moe_hop,shared_indexer_absorbed_query,deferred_norm",
-    ],
-)
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_attention_null(harness, layout, path, features, hot):
-    h = harness
-    base = ",".join(f for f in features.split(",") if f != "deferred_norm")
-    packets, pos = controls._packets(h, layout, 3 if not hot else 13, hot)
-    native_in = {layer: p.materialize() for layer, p in packets.items()}
-    w64 = _weights64(h)
-    hstar = {layer: _h_star(p) for layer, p in packets.items()}
-    ulp = {layer: _bf16_ulp64(v) for layer, v in hstar.items()}
-    oracle = {
-        layer: {
-            name: U.fp64_oracle(p.residual, p.gamma, w, p.eps).to(torch.bfloat16)
-            for name, w in w64[layer].items()
-        }
-        for layer, p in packets.items()
-    }
-    swap = _deferred_norm_projection_outputs(h, packets)
-    snap = h.snapshot()
-    runs = {}
-    try:
-
-        def native_run(hidden=native_in):
-            dn_integration._reset_runners(h)
-            return dn_integration._run(h, base, layout, path, hidden, pos, snap)
-
-        def override_run(per_layer):
-            kva, kw = _split(per_layer)
-            with _projection_override(h, kva, kw) as calls:
-                run = native_run()
-            assert all(v == 1 for v in calls.values()), calls
-            return run
-
-        runs["native"] = native_run()
-        dn_integration._reset_runners(h)
-        runs["deferred_norm"] = dn_integration._run(
-            h, features, layout, path, dn_integration._served(packets), pos, snap
-        )
-        for layer in LAYERS:
-            engaged = h.attn[layer]._deferred_norm_projection not in (None, False)
-            assert engaged == (layer in dn_integration.SERVED_LAYERS), f"deferred_norm on L{layer}"
-        runs["rerun"] = native_run()
-        runs["ref"] = override_run(oracle)
-        runs["swap"] = override_run(swap)
-        null = []
-        for k in range(1, K_ATTN + 1):
-            draw = {
-                layer: _draw(hstar[layer], ulp[layer], w64[layer], _seed(k, layer))
-                for layer in LAYERS
-            }
-            null.append(override_run(draw))
-
-        ref_out, nat_out = runs["ref"][0], runs["native"][0]
-        report = {
-            "layout": layout,
-            "path": path,
-            "features": features,
-            "data": "hot" if hot else "gauss",
-            "k": K_ATTN,
-            "layers": {},
-        }
-        for layer in LAYERS:
-            for run in list(runs.values()) + null:
-                assert torch.isfinite(run[0][layer].float()).all(), layer
-            arms = {}
-            for arm, run in runs.items():
-                vs_ref = dn_integration._output_stats(runs["ref"], run, layout)[layer]
-                arms[arm] = {
-                    "d_ref": aq_integration._rel_rms(run[0][layer], ref_out[layer]),
-                    "d_nat": aq_integration._rel_rms(run[0][layer], nat_out[layer]),
-                    "rows_d_ref": _row_rel(run[0][layer], ref_out[layer]),
-                    "quant_q_bytes_diff_vs_ref": vs_ref.get("quant_q_bytes_diff_frac"),
-                    "topk_jaccard_vs_ref": vs_ref.get("topk_mean_jaccard"),
-                }
-            d_ref = [aq_integration._rel_rms(r[0][layer], ref_out[layer]) for r in null]
-            d_nat = [aq_integration._rel_rms(r[0][layer], nat_out[layer]) for r in null]
-            rows = [_row_rel(r[0][layer], ref_out[layer]) for r in null]
-            vs_ref = [dn_integration._output_stats(runs["ref"], r, layout)[layer] for r in null]
-            q_bytes = [s.get("quant_q_bytes_diff_frac") for s in vs_ref]
-            jac = [s.get("topk_mean_jaccard") for s in vs_ref if "topk_mean_jaccard" in s]
-            nrows = len(rows[0])
-            null_stats = {
-                "d_ref": _order_stats(d_ref),
-                "d_nat": _order_stats(d_nat),
-                "d_ref_values": d_ref,
-                "d_nat_values": d_nat,
-                "rows_d_ref": [_order_stats([r[i] for r in rows]) for i in range(nrows)],
-                "quant_q_bytes_diff_vs_ref": _order_stats(q_bytes),
-            }
-            if jac:
-                null_stats["topk_jaccard_vs_ref"] = _order_stats(jac)
-            for arm in ("deferred_norm", "native", "swap"):
-                arms[arm]["null_ge_d_ref"] = sum(1 for v in d_ref if v >= arms[arm]["d_ref"])
-                arms[arm]["null_ge_d_nat"] = sum(1 for v in d_nat if v >= arms[arm]["d_nat"])
-            arms["deferred_norm"]["d_swap"] = aq_integration._rel_rms(
-                runs["deferred_norm"][0][layer], runs["swap"][0][layer]
-            )
-            report["layers"][str(layer)] = {"arms": arms, "null": null_stats}
-        print(
-            "GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS " + json.dumps(report, sort_keys=True),
-            flush=True,
-        )
-        out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_OUT")
-        if out_dir:
-            os.makedirs(out_dir, exist_ok=True)
-            name = (
-                f"null-controls-{layout}-{path}-{features.replace(',', '_')}-{report['data']}.json"
-            )
-            with open(os.path.join(out_dir, name), "w") as fh:
-                json.dump(report, fh, indent=1, sort_keys=True)
-    finally:
-        h.restore(snap)
-        dn_integration._set_features("")
-
-
-def _segments(h, layer, name, n):
-    attn = h.attn[layer]
-    if name == "kva":
-        kv = attn.kv_lora_rank
-        rope = getattr(attn, "qk_rope_head_dim", 64)
-        q = n - kv - rope
-        return {"q_a": (0, q), "kv_a": (q, q + kv), "k_rope": (q + kv, n)}
-    wk = attn.mqa.indexer.head_dim
-    return {"wk": (0, wk), "weights_proj": (wk, n)}
-
-
-def _audit(y, ystar, segments):
-    """Error statistics of ``y`` vs the FP64 oracle ``y*``."""
-    yd = y.double()
-    ulp = _bf16_ulp64(ystar)
-    err_ulp = torch.where(ulp > 0, (yd - ystar) / ulp.clamp_min(1e-300), torch.zeros_like(yd))
-    nearest = ystar.to(torch.bfloat16)
-    out = {
-        "e": U.norm_rms_err(y, ystar),
-        "frac_not_nearest": float((y != nearest).float().mean()),
-        "max_err_ulp": float(err_ulp.abs().max()),
-        "mean_err_ulp": float(err_ulp.mean()),
-        "frac_err_ge_1ulp": float((err_ulp.abs() >= 1.0).float().mean()),
-        "seg_e": {s: U.norm_rms_err(y[:, a:b], ystar[:, a:b]) for s, (a, b) in segments.items()},
-        "row_e": [U.norm_rms_err(y[i], ystar[i]) for i in range(y.shape[0])],
-    }
-    return out
-
-
-def _vs_spec(y, spec):
-    ulp = U.bf16_ulp_distance(y, spec)
-    return {"max_ulp": int(ulp.max()), "frac_diff": float((ulp > 0).float().mean())}
-
-
-@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
-@torch.inference_mode()
-def test_projection_null(harness, hot):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    h = harness
-    packets, _ = controls._packets(h, "6req_x1tok", 3 if not hot else 13, hot)
-    w64 = _weights64(h)
-    swap = _deferred_norm_projection_outputs(h, packets)
-    report = {"data": "hot" if hot else "gauss", "k": K_PROJ, "layers": {}}
-    for layer in dn_integration.SERVED_LAYERS:
-        p = packets[layer]
-        rows, hidden = p.residual.shape
-        hstar = _h_star(p)
-        ulp = _bf16_ulp64(hstar)
-        h_nat = p.materialize()
-        a = rows_from_frag(p.a_frag, rows, hidden)
-        r_w = torch.rsqrt(U.square_sum_in_order(p.sumsq) / hidden + p.eps)
-        layer_rep = {}
-        for name, w in w64[layer].items():
-            ystar = U.fp64_oracle(p.residual, p.gamma, w, p.eps)
-            segs = _segments(h, layer, name, w.shape[0])
-            if name == "kva":
-                y_nat = h.attn[layer].kv_a_proj_with_mqa(h_nat)
-            else:
-                y_nat = torch.nn.functional.linear(h_nat, _kw_weight(h, layer))
-            y_deferred_norm = swap[layer][name]
-            spec_nat = (h_nat.double() @ w.t()).to(torch.bfloat16)
-            spec_deferred_norm = (r_w.double().unsqueeze(-1) * (a.double() @ w.t())).to(
-                torch.bfloat16
-            )
-            e_null = []
-            for k in range(1, K_PROJ + 1):
-                g = torch.Generator(device="cuda").manual_seed(_seed(k, layer))
-                u = torch.rand(hstar.shape, generator=g, device="cuda", dtype=torch.float64) - 0.5
-                y = ((hstar + u * ulp) @ w.t()).to(torch.bfloat16)
-                e_null.append(U.norm_rms_err(y, ystar))
-            draw1 = _draw(hstar, ulp, {name: w}, _seed(1, layer))[name]
-            e_deferred_norm = U.norm_rms_err(y_deferred_norm, ystar)
-            e_nat = U.norm_rms_err(y_nat, ystar)
-            layer_rep[name] = {
-                "native": _audit(y_nat, ystar, segs),
-                "deferred_norm": _audit(y_deferred_norm, ystar, segs),
-                "ref": _audit(ystar.to(torch.bfloat16), ystar, segs),
-                "null_draw1": _audit(draw1, ystar, segs),
-                "null_e": _order_stats(e_null),
-                "null_ge_e_deferred_norm": sum(1 for v in e_null if v >= e_deferred_norm),
-                "null_ge_e_native": sum(1 for v in e_null if v >= e_nat),
-                "native_vs_spec": _vs_spec(y_nat, spec_nat),
-                "deferred_norm_vs_spec": _vs_spec(y_deferred_norm, spec_deferred_norm),
-                "deferred_norm_vs_native": _vs_spec(y_deferred_norm, y_nat),
-            }
-        report["layers"][str(layer)] = layer_rep
-    torch.cuda.synchronize()
-    print(
-        "GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_PROJ " + json.dumps(report, sort_keys=True),
-        flush=True,
-    )
-    out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_OUT")
-    if out_dir:
-        os.makedirs(out_dir, exist_ok=True)
-        with open(os.path.join(out_dir, f"null-controls-proj-{report['data']}.json"), "w") as fh:
-            json.dump(report, fh, indent=1, sort_keys=True)
-    for layer, d in report["layers"].items():
-        for name, s in d.items():
-            assert s["native"]["e"] == s["native"]["e"] and s["null_e"]["max"] < 0.05, (layer, name)
-
-
-def _fp32_reference(op, w64, scale=None):
-    """The same specification through an FP32 cuBLAS GEMM (TF32 off): a
-    third accumulation order of the same products."""
-    prev = torch.backends.cuda.matmul.allow_tf32
-    torch.backends.cuda.matmul.allow_tf32 = False
-    try:
-        y = op.float() @ w64.float().t()
-        if scale is not None:
-            y = y * scale.unsqueeze(-1)
-        return y.to(torch.bfloat16)
-    finally:
-        torch.backends.cuda.matmul.allow_tf32 = prev
-
-
-def _accumulation_audit(y, s, t_abs, ystar, max_list=16):
-    """``y`` (bf16) against its exact pre-rounding specification ``s`` (FP64).
-
-    ``rho`` is the error beyond the output rounding (half a bf16 ulp at
-    ``max(|s|, |y|)``) and beyond a ``2**-21 |s|`` allowance for the scale
-    ``r`` (``rsqrt.approx``), in units of ``2**-24 * sum_k |t_jk|``: the
-    scale of FP32 accumulation error of the dot product's terms ``t_jk``.
-    """
-    yd = y.double()
-    err = (yd - s).abs()
-    half = 0.5 * _bf16_ulp64(torch.maximum(s.abs(), yd.abs()))
-    excess = (err - half - 2.0**-21 * s.abs()).clamp_min(0.0)
-    rho = excess / (2.0**-24 * t_abs).clamp_min(1e-300)
-    ulp = U.bf16_ulp_distance(y, s.to(torch.bfloat16))
-    big = (ulp > 1).nonzero().tolist()
-    flat = rho.flatten()
-    k999 = max(1, int(0.999 * flat.numel()))
-    out = {
-        "frac_diff": float((ulp > 0).float().mean()),
-        "max_ulp": int(ulp.max()),
-        "n_gt1ulp": len(big),
-        "rho_max": float(rho.max()),
-        "rho_p999": float(flat.kthvalue(k999).values),
-        "gt1ulp": [
-            {
-                "row": i,
-                "col": j,
-                "ulp": int(ulp[i, j]),
-                "cancel": float(s[i, j].abs() / t_abs[i, j]),
-                "rho": float(rho[i, j]),
-            }
-            for i, j in big[:max_list]
-        ],
-        "bias_rel": float((yd - ystar).mean() / ystar.pow(2).mean().sqrt()),
-        "frac_above_oracle": float(
-            ((yd - ystar) > 0).double().sum() / ((yd - ystar) != 0).double().sum().clamp_min(1)
-        ),
-    }
-    return out
-
-
-@pytest.mark.parametrize("hot", [False, True], ids=["gauss", "hot"])
-@torch.inference_mode()
-def test_projection_spec_accumulation(harness, hot):
-    """Each projection kernel (deferred_norm's and native's) against its own exact
-    arithmetic specification, with the error measured in units of the FP32
-    accumulation scale of its terms; next to an FP32 cuBLAS reference of the
-    same specification."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    h = harness
-    packets, _ = controls._packets(h, "6req_x1tok", 3 if not hot else 13, hot)
-    w64 = _weights64(h)
-    deferred_norm_out = _deferred_norm_projection_outputs(h, packets)
-    report = {"data": "hot" if hot else "gauss", "layers": {}}
-    for layer in dn_integration.SERVED_LAYERS:
-        p = packets[layer]
-        rows, hidden = p.residual.shape
-        h_nat = p.materialize()
-        a = rows_from_frag(p.a_frag, rows, hidden)
-        r_w = torch.rsqrt(U.square_sum_in_order(p.sumsq) / hidden + p.eps)
-        layer_rep = {}
-        for name, w in w64[layer].items():
-            ystar = U.fp64_oracle(p.residual, p.gamma, w, p.eps)
-            if name == "kva":
-                y_nat = h.attn[layer].kv_a_proj_with_mqa(h_nat)
-            else:
-                y_nat = torch.nn.functional.linear(h_nat, _kw_weight(h, layer))
-            s_deferred_norm = r_w.double().unsqueeze(-1) * (a.double() @ w.t())
-            t_deferred_norm = r_w.double().unsqueeze(-1) * (a.double().abs() @ w.abs().t())
-            s_nat = h_nat.double() @ w.t()
-            t_nat = h_nat.double().abs() @ w.abs().t()
-            layer_rep[name] = {
-                "deferred_norm_kernel": _accumulation_audit(
-                    deferred_norm_out[layer][name], s_deferred_norm, t_deferred_norm, ystar
-                ),
-                "deferred_norm_fp32_ref": _accumulation_audit(
-                    _fp32_reference(a, w, r_w), s_deferred_norm, t_deferred_norm, ystar
-                ),
-                "native_kernel": _accumulation_audit(y_nat, s_nat, t_nat, ystar),
-                "native_fp32_ref": _accumulation_audit(
-                    _fp32_reference(h_nat, w), s_nat, t_nat, ystar
-                ),
-            }
-        report["layers"][str(layer)] = layer_rep
-    torch.cuda.synchronize()
-    print(
-        "GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_SPEC " + json.dumps(report, sort_keys=True),
-        flush=True,
-    )
-    out_dir = os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_NULL_CONTROLS_OUT")
-    if out_dir:
-        os.makedirs(out_dir, exist_ok=True)
-        with open(os.path.join(out_dir, f"null-controls-spec-{report['data']}.json"), "w") as fh:
-            json.dump(report, fh, indent=1, sort_keys=True)
-    for layer, d in report["layers"].items():
-        for name, s in d.items():
-            for impl, v in s.items():
-                assert v["rho_max"] == v["rho_max"], (layer, name, impl)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py
deleted file mode 100644
index da9db76e5168..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_projection.py
+++ /dev/null
@@ -1,293 +0,0 @@
-# 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.
-"""Single-GPU tests of the deferred_norm staged projection (real GLM-5.2 weights).
-
-Native chain: ``h = bf16((x * r) * gamma)`` (the post_moe_hop's output, produced
-here by the deferred_norm materialization kernel, which runs post_moe_hop's arithmetic) followed
-by the native ``kv_a_proj_with_mqa`` (``dsv3_fused_a_gemm_op``).
-deferred_norm: the staged projection on ``a = bf16(x * gamma)`` and the hop's
-square-sum partials (emulated bit-exactly in torch).
-
-The target's deferred-norm scope is ``shared``: the projection serves the shared-indexer
-layers (``LAYERS``); a full-indexer layer (``FULL_INDEXER_LAYERS``) keeps its native
-projection, and the projection's eligibility check declines it.
-
-Frozen error budget: per output tensor
-``allclose(y_deferred_norm, y_native, rtol=0.1, atol=1e-2)`` and normalized RMS error
-vs the FP64 oracle ``e_deferred_norm <= 1.25 * e_native + 2**-12``.
-
-Fixture label: weights real (checkpoint), activations synthetic.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
-
-import absorbed_query_fixture  # noqa: E402
-import deferred_norm_utils as U  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
-    reason="needs a GPU and the GLM-5.2 checkpoint",
-)
-
-LAYERS = [int(x) for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_TEST_LAYERS", "7").split(",")]
-FULL_INDEXER_LAYERS = [
-    int(x) for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_TEST_FULL_LAYERS", "6").split(",")
-]
-ROWS = [int(x) for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_TEST_ROWS", "1,6,8").split(",")]
-_MODS = {}
-
-
-def _module(layer):
-    if layer not in _MODS:
-        with torch.inference_mode():
-            cfg = absorbed_query_fixture.build_model_config(0)
-            _MODS[layer] = (
-                absorbed_query_fixture.build_native_attention(layer, 0, cfg),
-                absorbed_query_fixture.load_input_layernorm(layer).cuda(),
-            )
-    return _MODS[layer]
-
-
-def make_residual(rows, seed, scale=1.0):
-    """Residual-stream-like rows: unit Gaussian plus a few large channels."""
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = torch.randn(rows, U.H, generator=g, device="cuda") * 2.0
-    hot = torch.randint(0, U.H, (12,), generator=g, device="cuda")
-    x[:, hot] *= 30.0
-    return (x * scale).bfloat16()
-
-
-def make_packet(x, gamma):
-    return U.make_frag_packet(x, gamma)
-
-
-def _projection(attn):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-
-    ok, reason, kw = projection_eligibility(attn)
-    assert ok, reason
-    proj = DeferredNormProjection(attn, kw_weight=kw)
-    assert proj.supported() is None, proj.supported()
-    return proj
-
-
-def _native(attn, h):
-    kva = attn.kv_a_proj_with_mqa(h)
-    kw = None
-    if attn.mqa.indexer is not None:
-        kw = torch.nn.functional.linear(h, attn.mqa.indexer._fused_wk_wp_weight)
-    return kva, kw
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-def test_served_configs_cover_the_consumer(layer):
-    """``precompile_served`` (run by the producing hop before its cross-rank
-    agreement) compiles exactly the kernel the consumer launches on this
-    layer kind, so the consumer's ``precompile`` is a cache hit."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.proj_kernel import (  # noqa: E501
-        is_compiled,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        precompile_served,
-        served_kernel_configs,
-    )
-
-    device = torch.cuda.current_device()
-    assert precompile_served(device) is None
-    proj = _projection(_module(layer)[0])
-    assert proj.served_config() in served_kernel_configs(device)
-    assert is_compiled(**proj.served_config())
-
-
-@pytest.mark.parametrize("scale", [1.0, 1e-3, 30.0])
-@pytest.mark.parametrize("rows", ROWS)
-@pytest.mark.parametrize("layer", LAYERS)
-def test_projection_within_frozen_budget(layer, rows, scale):
-    attn, gamma = _module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        x = make_residual(rows, seed=100 + rows + layer, scale=scale)
-        packet = make_packet(x, gamma)
-        h = packet.materialize()
-        kva_n, kw_n = _native(attn, h)
-        kva_w, kw_w = proj.project(packet)
-        torch.cuda.synchronize()
-        outs = [("kva", kva_w, kva_n, attn.kv_a_proj_with_mqa.weight)]
-        if kw_n is not None:
-            outs.append(("kw", kw_w, kw_n, attn.mqa.indexer._fused_wk_wp_weight))
-        for name, cand, nat, w in outs:
-            oracle = U.fp64_oracle(x, gamma, w)
-            stats = U.compare(cand, nat, oracle)
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_BUDGET layer={layer} rows={rows} scale={scale} out={name} {stats}"
-            )
-            assert U.within_budget(stats), (name, stats)
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-def test_native_input_mode(layer):
-    """Same kernel on the normalized hidden with r = 1 (attribution arm)."""
-    attn, gamma = _module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        x = make_residual(6, seed=7)
-        h = make_packet(x, gamma).materialize()
-        kva_n, kw_n = _native(attn, h)
-        kva_w, kw_w = proj.run(h, None, U.EPS)
-        torch.cuda.synchronize()
-        for name, cand, nat in (("kva", kva_w, kva_n), ("kw", kw_w, kw_n)):
-            if nat is None:
-                continue
-            ulp = U.bf16_ulp_distance(cand, nat)
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_NATIVE_INPUT layer={layer} out={name} max_ulp={int(ulp.max())} "
-                f"frac_diff={float((ulp > 0).float().mean()):.3e}"
-            )
-            assert torch.allclose(cand.float(), nat.float(), rtol=0.1, atol=1e-2)
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-def test_rows_match_single_row_runs(layer):
-    """Every token row equals a launch of that row alone (no cross-row mixing)."""
-    attn, gamma = _module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        x = make_residual(8, seed=11)
-        kva, kw = proj.project(make_packet(x, gamma))
-        for t in range(8):
-            kva1, kw1 = proj.project(make_packet(x[t : t + 1].contiguous(), gamma))
-            assert torch.equal(kva1[0], kva[t])
-            if kw is not None:
-                assert torch.equal(kw1[0], kw[t])
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-def test_graph_replay_matches_eager(layer):
-    attn, gamma = _module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        x0 = make_residual(6, seed=21)
-        static = U.empty_like_packet(make_packet(x0, gamma))
-        ref0 = make_packet(x0, gamma)
-        static.a_frag.copy_(ref0.a_frag)
-        static.sumsq.copy_(ref0.sumsq)
-        proj.project(static)
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            g_kva, g_kw = proj.project(static)
-        for trial, replays in enumerate((1, 1, 50)):
-            x = make_residual(6, seed=300 + trial)
-            ref = make_packet(x, gamma)
-            static.a_frag.copy_(ref.a_frag)
-            static.sumsq.copy_(ref.sumsq)
-            for _ in range(replays):
-                graph.replay()
-            e_kva, e_kw = proj.project(ref)
-            torch.cuda.synchronize()
-            assert torch.equal(g_kva, e_kva)
-            if e_kw is not None:
-                assert torch.equal(g_kw, e_kw)
-
-
-@pytest.mark.parametrize("stage_before_wait", [True, False])
-@pytest.mark.parametrize("delay_ns", [0, 5000, 50000])
-@pytest.mark.parametrize("layer", LAYERS)
-def test_pdl_delayed_producer(layer, delay_ns, stage_before_wait):
-    """The producer releases its dependents at entry and writes the packet only
-    after ``delay_ns``; the projection (launched with PDL, staging its weights
-    before its dependency wait) must still read the final packet."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormProjConfig,
-        DeferredNormProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-
-    attn, gamma = _module(layer)
-    _, _, kwt = projection_eligibility(attn)
-    proj = DeferredNormProjection(
-        attn, DeferredNormProjConfig(stage_before_wait=stage_before_wait), kw_weight=kwt
-    )
-    if not proj.use_pdl:
-        pytest.skip("PDL disabled (TRTLLM_ENABLE_PDL=0)")
-    with torch.inference_mode():
-        for trial in range(4):
-            x = make_residual(6, seed=500 + trial)
-            src = make_packet(x, gamma)
-            dst = U.empty_like_packet(src)
-            dst.a_frag.fill_(float("nan"))
-            U.sumsq_full(dst).fill_(float("nan"))
-            torch.cuda.synchronize()
-            U.proxy_producer(
-                src.a_frag, dst.a_frag, U.sumsq_full(src), U.sumsq_full(dst), delay_ns, rows=8
-            )
-            kva, kw = proj.project(dst)
-            ref_kva, ref_kw = proj.project(src)
-            torch.cuda.synchronize()
-            assert torch.equal(kva, ref_kva), f"trial {trial}"
-            if kw is not None:
-                assert torch.equal(kw, ref_kw)
-
-
-def test_materialize_close_to_torch():
-    """The fallback's normalized output vs a torch RMSNorm of the same x
-    (bit-exactness against the post_moe_hop is asserted by the MPI tests)."""
-    _, gamma = _module(LAYERS[0])
-    with torch.inference_mode():
-        x = make_residual(6, seed=3)
-        h = make_packet(x, gamma).materialize()
-        s = U.square_sum_in_order(U.hop_sumsq_partials(x))
-        r = torch.rsqrt(s / U.H + U.EPS).unsqueeze(-1)
-        ref = ((x.float() * r) * gamma.float()).bfloat16()
-        assert int(U.bf16_ulp_distance(h, ref).max()) <= 1
-
-
-@pytest.mark.parametrize("layer", FULL_INDEXER_LAYERS)
-def test_full_indexer_layer_keeps_its_native_projection(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-
-    attn, _ = _module(layer)
-    assert attn.mqa.indexer is not None
-    ok, reason, kw = projection_eligibility(attn)
-    assert not ok and kw is None
-    assert reason.endswith("=shared: full-indexer layer keeps its native projection"), reason
-
-
-def test_eligibility_rejects_out_of_domain():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.proj_kernel import (  # noqa: E501
-        projection_supported,
-    )
-
-    assert projection_supported(6144, 2624, 152, 512, 18, 4) is None
-    assert projection_supported(6144, 2784, 152, 512, 18, 4) is None
-    assert projection_supported(6144, 2624, 64, 512, 18, 4) is not None  # 41 rows per CTA
-    assert projection_supported(6144, 2784, 152, 512, 17, 4) is not None  # 2 rows over staged
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py
deleted file mode 100644
index 05718956b17c..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_routed_mpi.py
+++ /dev/null
@@ -1,946 +0,0 @@
-# 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.
-"""Collective tests of the deferred_norm deferred hop on routed_moe_into_hop's routed input (8 ranks, 2 hosts).
-
-With ``routed_moe_into_hop`` selected, routed_moe writes its per-rank routed result into the
-post_moe_hop hop's routed-input buffer (``PostMoeHop.enable_routed_input``; modes
-``fp32`` (default), ``identity``, ``row``) and the layer calls
-``PostMoeHop.hop_routed``. ``DeferredNormHop.hop_routed`` takes the same
-inputs and returns the deferred packet. The routed content is written here
-as routed_moe would write it (per rank), from a synthetic raw tuple.
-
-Exact invariants, per mode, against ``PostMoeHop.hop_routed`` on the same
-routed content and the same workspace:
-
-* ``residual_out`` bit-identical;
-* ``a_frag`` holds ``bf16(x * gamma)`` bitwise; ``sumsq`` equals the torch
-  emulation of the hop's order; ``materialize(packet)`` equals
-  ``hop_routed``'s normalized output bit for bit;
-* ``fp32`` mode: the hop clears the accumulator rows it read (the buffer is
-  zero after every call, as after post_moe_hop's).
-
-post_moe_hop's error protocol on the routed variant, mixed with post_moe_hop routed, post_moe_hop
-raw and deferred_norm raw calls on one workspace: missing rank (deadline, NaN, health check,
-reset), sticky poison carried through every Lamport buffer, fault-injected
-store order with per-rank kernel choice, contract violations before any
-device work, per-rank layout repairs (also in a CUDA graph), graph replay,
-the health check's call count, and the collective enable (post_moe_hop routed input
-off, compile failure on one rank, modes that differ across ranks, mode
-switch).
-"""
-
-import os
-import random
-import sys
-import time
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
-
-import deferred_norm_utils as U  # noqa: E402
-import post_moe_hop_mpi_utils as M  # noqa: E402
-
-MODES = ("fp32", "identity", "row")
-
-
-def _env():
-    env = M.mpi_env()
-    if env.world != 8:
-        pytest.skip("the deferred_norm hop is specialized to TP8")
-    return env
-
-
-def _routed_api():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    if not hasattr(PostMoeHop, "enable_routed_input"):
-        pytest.skip("the installed post_moe_hop hop has no routed_moe_into_hop routed input")
-
-
-_HOPS = {}
-
-
-def _hops(mode, trigger_early=True, timeout_s=None, fault_delay_ns=0):
-    """A post_moe_hop with its routed input in ``mode`` and a deferred_norm hop on it (collective)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    _routed_api()
-    key = (mode, timeout_s, fault_delay_ns)
-    if key not in _HOPS:
-        kw = {}
-        if timeout_s is not None:
-            kw["timeout_s"] = timeout_s
-        if fault_delay_ns:
-            kw["fault_delay_ns"] = fault_delay_ns
-        post_moe_hop = PostMoeHop(
-            env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False, **kw
-        )
-        assert M.all_ranks(post_moe_hop.engaged and post_moe_hop.enable_routed_input(mode)), (
-            post_moe_hop.reason
-        )
-        _HOPS[key] = {"post_moe_hop": post_moe_hop}
-    entry = _HOPS[key]
-    if trigger_early not in entry:
-        deferred_norm = DeferredNormHop(entry["post_moe_hop"], trigger_early=trigger_early)
-        assert M.all_ranks(deferred_norm.engaged and deferred_norm.enable_routed_input()), (
-            deferred_norm.reason,
-            deferred_norm.routed_reason,
-        )
-        entry[trigger_early] = deferred_norm
-    return entry["post_moe_hop"], entry[trigger_early]
-
-
-def _source(mode, rows, seed, with_shared=True):
-    """Per-rank routed content as routed_moe writes it in ``mode``, plus shared /
-    residual / gamma (the residual and gamma are replicated)."""
-    fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(
-        rows, seed=seed, with_shared=with_shared
-    )
-    if mode == "row":
-        src = M.reference_local_partial(fc2, scales, perm, None)
-    elif mode == "identity":
-        valid = perm >= 0
-        safe = torch.where(valid, perm, torch.zeros_like(perm)).long()
-        src = (
-            fc2[safe.view(-1)].contiguous(),
-            torch.where(valid, scales, torch.zeros_like(scales)).bfloat16().contiguous(),
-        )
-    else:
-        valid = (perm >= 0).unsqueeze(-1)
-        safe = torch.where(perm >= 0, perm, torch.zeros_like(perm)).long()
-        prod = fc2[safe.view(-1)].view(rows, M.GLM_TOP_K, -1).float() * scales.float().unsqueeze(-1)
-        src = torch.where(valid, prod, torch.zeros_like(prod)).sum(1).contiguous()
-    return src, shared, residual, gamma
-
-
-def _fill(post_moe_hop, mode, src, rows):
-    """Write ``src`` into the post_moe_hop's routed buffer (stream-ordered); return the view."""
-    if mode == "row":
-        buf = post_moe_hop.routed_input_buffer(rows)
-        buf.copy_(src)
-        return buf
-    if mode == "identity":
-        buf, scales = post_moe_hop.routed_raw_buffers(rows)
-        buf.copy_(src[0])
-        scales.copy_(src[1])
-        return buf
-    buf = post_moe_hop.routed_acc_buffer(rows)
-    buf.copy_(src)
-    return buf
-
-
-def _post_moe_hop_reference(post_moe_hop, mode, src, rows, shared, residual, gamma):
-    routed = _fill(post_moe_hop, mode, src, rows)
-    return post_moe_hop.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
-
-
-def _deferred_norm_call(post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma):
-    routed = _fill(post_moe_hop, mode, src, rows)
-    return deferred_norm.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
-
-
-def _acc_clear(post_moe_hop, mode):
-    if mode != "fp32":
-        return True
-    torch.cuda.synchronize()
-    return bool((post_moe_hop.routed_acc_buffer(post_moe_hop.max_rows) == 0).all())
-
-
-def _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
-    """This rank's packet invariants against post_moe_hop's outputs (not collective)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    torch.cuda.synchronize()
-    ok_res = torch.equal(res_deferred_norm, res_post_moe_hop)
-    scaled_ref = (res_deferred_norm.float() * gamma.float()).bfloat16()
-    ok_scaled = torch.equal(
-        rows_from_frag(packet.a_frag, res_deferred_norm.shape[0], res_deferred_norm.shape[1]),
-        scaled_ref,
-    )
-    ok_ss = torch.equal(packet.sumsq, U.hop_sumsq_partials(res_deferred_norm))
-    h = materialize(packet)
-    torch.cuda.synchronize()
-    ok_h = torch.equal(h, h_post_moe_hop)
-    ok = ok_res and ok_scaled and ok_ss and ok_h
-    if not ok:
-        print(
-            f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_FAIL {tag} rank={M.mpi_env().rank} res={ok_res} scaled={ok_scaled} "
-            f"sumsq={ok_ss} materialize={ok_h} "
-            f"max_ulp_h={int(U.bf16_ulp_distance(h, h_post_moe_hop).max())}",
-            flush=True,
-        )
-    return ok
-
-
-def _check(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag):
-    return M.all_ranks(
-        _verdict(packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, tag)
-    )
-
-
-def _all_nan_packet(packet, res):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    torch.cuda.synchronize()
-    a = rows_from_frag(packet.a_frag, res.shape[0], res.shape[1])
-    h = materialize(packet)
-    torch.cuda.synchronize()
-    return all(bool(torch.isnan(t.float()).all()) for t in (res, a, packet.sumsq, h))
-
-
-def _reset(post_moe_hop):
-    assert M.all_ranks(post_moe_hop.reset())
-    assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-# ---------------------------------------------------------------------------
-# Collective enable
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("mode", MODES)
-def test_routed_engaged(mode):
-    post_moe_hop, deferred_norm = _hops(mode)
-    assert M.all_ranks(deferred_norm.routed_engaged and deferred_norm.routed_mode == mode)
-    md = type("Md", (), {"num_contexts": 0, "num_generations": 6})()
-    ctx = type("Md", (), {"num_contexts": 1, "num_generations": 0})()
-    assert M.all_ranks(
-        deferred_norm.can_run_routed(6, md)
-        and not deferred_norm.can_run_routed(6, ctx)
-        and not deferred_norm.can_run_routed(9, md)
-        and not deferred_norm.can_run_routed(0, md)
-    )
-
-
-def test_routed_requires_post_moe_hop_routed_input():
-    """Without the post_moe_hop's routed input the routed variant stays off on every rank."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    _routed_api()
-    post_moe_hop = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
-    deferred_norm = DeferredNormHop(post_moe_hop)
-    ok = (
-        deferred_norm.engaged
-        and not deferred_norm.enable_routed_input()
-        and not deferred_norm.routed_engaged
-    )
-    ok &= "routed input is not enabled" in deferred_norm.routed_reason
-    md = type("Md", (), {"num_contexts": 0, "num_generations": 6})()
-    ok &= not deferred_norm.can_run_routed(6, md)
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source("fp32", 6, seed=31)
-        with pytest.raises(ValueError):
-            deferred_norm.hop_routed(src, shared, residual, gamma, M.GLM_EPS)
-    assert M.all_ranks(ok), deferred_norm.routed_reason
-
-
-@pytest.mark.parametrize("failing_rank", [2, 7])
-def test_routed_compile_failure_disengages_every_rank(monkeypatch, failing_rank):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-        hop_kernel,
-    )
-
-    env = _env()
-    post_moe_hop, _ = _hops("fp32")
-    deferred_norm = DeferredNormHop(post_moe_hop)
-    if env.rank == failing_rank:
-        real = hop_kernel.compiled_hop
-
-        def _fail(*args, **kwargs):
-            if kwargs.get("routed_input"):
-                raise RuntimeError("injected routed compile failure")
-            return real(*args, **kwargs)
-
-        monkeypatch.setattr(hop_kernel, "compiled_hop", _fail)
-    engaged = deferred_norm.enable_routed_input()
-    if env.rank == failing_rank:
-        local = "injected routed compile failure" in deferred_norm.routed_reason
-    else:
-        local = deferred_norm.routed_reason == "another rank is not capable"
-    assert M.all_ranks(
-        deferred_norm.engaged and not engaged and not deferred_norm.routed_engaged and local
-    ), deferred_norm.routed_reason
-
-
-def test_routed_modes_must_agree_across_ranks():
-    """post_moe_hop hops whose routed-input modes differ across ranks (environment)
-    leave the deferred_norm routed variant off on every rank."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    _routed_api()
-    post_moe_hop = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
-    assert M.all_ranks(post_moe_hop.enable_routed_input("fp32" if env.rank % 2 else "identity"))
-    deferred_norm = DeferredNormHop(post_moe_hop)
-    engaged = deferred_norm.enable_routed_input()
-    assert M.all_ranks(not engaged and "modes differ" in deferred_norm.routed_reason), (
-        deferred_norm.routed_reason
-    )
-
-
-def test_routed_mode_switch_needs_reenable():
-    """After the post_moe_hop switches modes (collective) the routed variant
-    declines until it is enabled again, then serves the new mode exactly."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    _routed_api()
-    post_moe_hop = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
-    assert M.all_ranks(post_moe_hop.enable_routed_input("fp32"))
-    deferred_norm = DeferredNormHop(post_moe_hop)
-    assert M.all_ranks(deferred_norm.enable_routed_input() and deferred_norm.routed_mode == "fp32")
-    md = type("Md", (), {"num_contexts": 0, "num_generations": 6})()
-    assert M.all_ranks(post_moe_hop.enable_routed_input("identity"))
-    ok = not deferred_norm.can_run_routed(6, md)
-    ok &= deferred_norm.enable_routed_input() and deferred_norm.routed_mode == "identity"
-    ok &= deferred_norm.can_run_routed(6, md)
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source("identity", 6, seed=33)
-        h_post_moe_hop, res_post_moe_hop = _post_moe_hop_reference(
-            post_moe_hop, "identity", src, 6, shared, residual, gamma
-        )
-        packet, res = _deferred_norm_call(
-            post_moe_hop, deferred_norm, "identity", src, 6, shared, residual, gamma
-        )
-        ok &= _verdict(packet, res, h_post_moe_hop, res_post_moe_hop, gamma, "mode-switch")
-    assert M.all_ranks(ok)
-
-
-# ---------------------------------------------------------------------------
-# Exactness against post_moe_hop's hop_routed
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("trigger_early", [True, False])
-@pytest.mark.parametrize("with_shared", [True, False])
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-@pytest.mark.parametrize("mode", MODES)
-def test_routed_packet_matches_hop_routed(mode, rows, with_shared, trigger_early):
-    post_moe_hop, deferred_norm = _hops(mode, trigger_early)
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, rows, seed=50 + rows, with_shared=with_shared)
-        h_post_moe_hop, res_post_moe_hop = _post_moe_hop_reference(
-            post_moe_hop, mode, src, rows, shared, residual, gamma
-        )
-        cleared_post_moe_hop = _acc_clear(post_moe_hop, mode)
-        packet, res_deferred_norm = _deferred_norm_call(
-            post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma
-        )
-        cleared_deferred_norm = _acc_clear(post_moe_hop, mode)
-        ok = _verdict(
-            packet, res_deferred_norm, h_post_moe_hop, res_post_moe_hop, gamma, f"{mode} R={rows}"
-        )
-        assert M.all_ranks(ok and cleared_post_moe_hop and cleared_deferred_norm), (
-            cleared_post_moe_hop,
-            cleared_deferred_norm,
-        )
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize("mode", MODES)
-def test_routed_large_residual_and_zero_routed(mode):
-    post_moe_hop, deferred_norm = _hops(mode)
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, 6, seed=77)
-        residual = (residual.float() * 50.0).bfloat16()
-        if mode == "identity":
-            src = (src[0], torch.zeros_like(src[1]))
-        else:
-            src = torch.zeros_like(src)
-        h_post_moe_hop, res_post_moe_hop = _post_moe_hop_reference(
-            post_moe_hop, mode, src, 6, shared, residual, gamma
-        )
-        packet, res = _deferred_norm_call(
-            post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma
-        )
-        assert _check(packet, res, h_post_moe_hop, res_post_moe_hop, gamma, f"{mode}-large-res")
-
-
-@pytest.mark.parametrize("mode", MODES)
-def test_back_to_back_routed_and_raw_mixed_with_post_moe_hop(mode):
-    """deferred_norm routed, post_moe_hop routed, deferred_norm raw and post_moe_hop raw calls interleaved on one
-    workspace without host syncs (each routed call refills the buffer on
-    the stream, as routed_moe does)."""
-    post_moe_hop, deferred_norm = _hops(mode)
-    kinds = [
-        "routed_deferred_norm",
-        "routed_post_moe_hop",
-        "deferred_norm",
-        "post_moe_hop",
-        "routed_deferred_norm",
-        "routed_deferred_norm",
-        "post_moe_hop",
-        "deferred_norm",
-        "routed_post_moe_hop",
-        "routed_deferred_norm",
-        "deferred_norm",
-        "routed_post_moe_hop",
-    ]
-    rows_seq = [6, 1, 8, 3, 6, 2, 6, 8, 5, 6, 1, 6]
-    with torch.inference_mode():
-        inputs = []
-        for i, (kind, r) in enumerate(zip(kinds, rows_seq)):
-            if kind.startswith("routed_"):
-                inputs.append(_source(mode, r, seed=900 + i))
-            else:
-                inputs.append(M.make_raw_tuple(r, seed=900 + i))
-        outs = []
-        for kind, r, t in zip(kinds, rows_seq, inputs):
-            if kind == "routed_deferred_norm":
-                outs.append(_deferred_norm_call(post_moe_hop, deferred_norm, mode, t[0], r, *t[1:]))
-            elif kind == "routed_post_moe_hop":
-                outs.append(_post_moe_hop_reference(post_moe_hop, mode, t[0], r, *t[1:]))
-            elif kind == "deferred_norm":
-                outs.append(deferred_norm(*t, M.GLM_EPS))
-            else:
-                outs.append(post_moe_hop(*t, M.GLM_EPS))
-        torch.cuda.synchronize()
-        failures = []
-        for i, (kind, r, t, out) in enumerate(zip(kinds, rows_seq, inputs, outs)):
-            if kind.startswith("routed_"):
-                h_ref, res_ref = _post_moe_hop_reference(post_moe_hop, mode, t[0], r, *t[1:])
-                gamma = t[3]
-            else:
-                h_ref, res_ref = post_moe_hop(*t, M.GLM_EPS)
-                gamma = t[5]
-            torch.cuda.synchronize()
-            if kind.endswith("post_moe_hop"):
-                ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
-            else:
-                ok = _verdict(out[0], out[1], h_ref, res_ref, gamma, f"b2b{i}-{kind}")
-            if not ok:
-                failures.append((i, kind))
-        assert M.all_ranks(not failures and _acc_clear(post_moe_hop, mode)), failures
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize("mode", MODES)
-def test_routed_cuda_graph_replay(mode):
-    """Fill + deferred_norm routed, fill + post_moe_hop routed and a deferred_norm raw call in one CUDA
-    graph, replayed with fresh routed content."""
-    post_moe_hop, deferred_norm = _hops(mode)
-    rows = 6
-    with torch.inference_mode():
-        static = [_source(mode, rows, seed=1200 + i) for i in range(2)]
-        raw = [x.clone() if x is not None else None for x in M.make_raw_tuple(rows, seed=1250)]
-
-        def clone_src(s):
-            return tuple(x.clone() for x in s) if isinstance(s, tuple) else s.clone()
-
-        bufs = [(clone_src(s[0]),) + tuple(s[1:]) for s in static]
-
-        def body():
-            a = _deferred_norm_call(
-                post_moe_hop, deferred_norm, mode, bufs[0][0], rows, *bufs[0][1:]
-            )
-            b = _post_moe_hop_reference(post_moe_hop, mode, bufs[1][0], rows, *bufs[1][1:])
-            c = deferred_norm(*raw, M.GLM_EPS)
-            return a, b, c
-
-        body()
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            ga, gb, gc = body()
-        for trial, replays in enumerate((1, 1, 20)):
-            fresh = [_source(mode, rows, seed=1300 + 10 * trial + i) for i in range(2)]
-            fresh_raw = M.make_raw_tuple(rows, seed=1350 + trial)
-            for dst, src in zip(bufs, fresh):
-                if isinstance(dst[0], tuple):
-                    for d, s in zip(dst[0], src[0]):
-                        d.copy_(s)
-                else:
-                    dst[0].copy_(src[0])
-                for d, s in zip(dst[1:], src[1:]):
-                    if d is not None:
-                        d.copy_(s)
-            for d, s in zip(raw, fresh_raw):
-                if d is not None:
-                    d.copy_(s)
-            for _ in range(replays):
-                M.mpi_env().comm.barrier()
-                graph.replay()
-            torch.cuda.synchronize()
-            ok = True
-            h0, r0 = _post_moe_hop_reference(post_moe_hop, mode, fresh[0][0], rows, *fresh[0][1:])
-            ok &= _verdict(ga[0], ga[1], h0, r0, fresh[0][3], f"graph{trial}-routed_deferred_norm")
-            h1, r1 = _post_moe_hop_reference(post_moe_hop, mode, fresh[1][0], rows, *fresh[1][1:])
-            torch.cuda.synchronize()
-            ok &= torch.equal(gb[0], h1) and torch.equal(gb[1], r1)
-            h2, r2 = post_moe_hop(*fresh_raw, M.GLM_EPS)
-            ok &= _verdict(gc[0], gc[1], h2, r2, fresh_raw[5], f"graph{trial}-deferred_norm")
-            assert M.all_ranks(ok and _acc_clear(post_moe_hop, mode)), f"trial {trial}"
-
-
-def test_routed_delayed_rank_stress():
-    """Random per-rank device delays inside back-to-back routed batches."""
-    mode = "fp32"
-    post_moe_hop, deferred_norm = _hops(mode)
-    env = M.mpi_env()
-    gen = torch.Generator().manual_seed(4243)
-    with torch.inference_mode():
-        for batch in range(3):
-            inputs, outs = [], []
-            for i in range(12):
-                rows = int(torch.randint(1, 9, (1,), generator=gen))
-                t = _source(mode, rows, seed=5000 + 100 * batch + i)
-                local = torch.Generator().manual_seed(77 * batch + 13 * i + env.rank)
-                if float(torch.rand(1, generator=local)) < 0.3:
-                    torch.cuda._sleep(int(float(torch.rand(1, generator=local)) * 2_000_000))
-                inputs.append((rows, t))
-                outs.append(
-                    _deferred_norm_call(post_moe_hop, deferred_norm, mode, t[0], rows, *t[1:])
-                )
-            torch.cuda.synchronize()
-            ok = True
-            for (rows, t), (packet, res) in zip(inputs, outs):
-                h_ref, res_ref = _post_moe_hop_reference(post_moe_hop, mode, t[0], rows, *t[1:])
-                ok &= _verdict(packet, res, h_ref, res_ref, t[3], f"stress{batch}")
-            assert M.all_ranks(ok), f"batch {batch}"
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-# ---------------------------------------------------------------------------
-# post_moe_hop's error protocol on the routed variant
-# ---------------------------------------------------------------------------
-
-_MISSING_TIMEOUT_S = 0.2
-
-
-@pytest.mark.parametrize("missing", [0, 5])
-@pytest.mark.parametrize("mode", ["fp32", "identity"])
-def test_routed_missing_rank_times_out_with_nan_and_reset(mode, missing):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-
-    post_moe_hop, deferred_norm = _hops(mode)
-    imp_post_moe_hop, imp_deferred_norm = _hops(mode, timeout_s=_MISSING_TIMEOUT_S)
-    env = M.mpi_env()
-    assert M.all_ranks(
-        imp_deferred_norm.routed_engaged and imp_post_moe_hop.workspace is post_moe_hop.workspace
-    )
-    rows = 6
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, rows, seed=600 + missing)
-        before = M.lamport_snapshot(post_moe_hop)
-        assert M.all_ranks(before["ok"]), before
-        env.comm.barrier()
-        start = time.monotonic()
-        packet = res = None
-        if env.rank != missing:
-            packet, res = _deferred_norm_call(
-                imp_post_moe_hop, imp_deferred_norm, mode, src, rows, shared, residual, gamma
-            )
-        torch.cuda.synchronize()
-        elapsed = time.monotonic() - start
-        counters = post_moe_hop.workspace.counters()
-        bits, polls = counters["error_bits"], counters["timeout_polls"]
-        if env.rank == missing:
-            reported = bits == 0 and polls == 0
-        else:
-            reported = (
-                bool(bits & 1)
-                and polls > 0
-                and _MISSING_TIMEOUT_S <= elapsed < 30.0
-                and _all_nan_packet(packet, res)
-                and _acc_clear(imp_post_moe_hop, mode)
-            )
-        stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
-        if env.rank == 0:
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_MISSING mode={mode} rank {missing} skipped: (rank, elapsed_s, "
-                f"error_bits, timeout_polls, ok) {stats}",
-                flush=True,
-            )
-        assert M.all_ranks(reported), stats
-        try:
-            check_health(f"routed_deferred_norm-missing-rank{missing}-local")
-            local_raised = False
-        except PostMoeHopError:
-            local_raised = True
-        try:
-            check_health(f"routed_deferred_norm-missing-rank{missing}-collective", collective=True)
-            group_raised = False
-        except PostMoeHopError:
-            group_raised = True
-        assert M.all_ranks(local_raised == (env.rank != missing) and group_raised)
-        _reset(post_moe_hop)
-        src, shared, residual, gamma = _source(mode, rows, seed=650 + missing)
-        packet, res = _deferred_norm_call(
-            post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma
-        )
-        h_ref, res_ref = _post_moe_hop_reference(
-            post_moe_hop, mode, src, rows, shared, residual, gamma
-        )
-        assert _check(packet, res, h_ref, res_ref, gamma, f"after-reset{missing}")
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize(
-    "sequence",
-    [
-        "routed_deferred_norm,routed_deferred_norm,routed_deferred_norm",
-        "routed_deferred_norm,routed_post_moe_hop,deferred_norm",
-        "routed_post_moe_hop,routed_deferred_norm,post_moe_hop",
-        "deferred_norm,routed_deferred_norm,routed_post_moe_hop",
-    ],
-)
-def test_routed_error_poisons_every_rank_until_reset(sequence):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
-        STATE_EPOCH,
-        STATE_ERROR_WORDS,
-    )
-
-    mode = "fp32"
-    post_moe_hop, deferred_norm = _hops(mode)
-    env = M.mpi_env()
-    victim = env.world - 3
-    with torch.inference_mode():
-        torch.cuda.synchronize()
-        if env.rank == victim:
-            epoch = int(post_moe_hop.workspace.state[STATE_EPOCH].item())
-            post_moe_hop.workspace.state[STATE_ERROR_WORDS[(epoch + 2) % 3]] = 1
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        poisoned = True
-        for call, kind in enumerate(sequence.split(",")):
-            if kind.startswith("routed_"):
-                src, shared, residual, gamma = _source(mode, 6, seed=860 + call)
-                if kind == "routed_deferred_norm":
-                    packet, res = _deferred_norm_call(
-                        post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma
-                    )
-                    poisoned &= _all_nan_packet(packet, res)
-                else:
-                    out, res = _post_moe_hop_reference(
-                        post_moe_hop, mode, src, 6, shared, residual, gamma
-                    )
-                    torch.cuda.synchronize()
-                    poisoned &= bool(
-                        torch.isnan(out.float()).all() and torch.isnan(res.float()).all()
-                    )
-            else:
-                t = M.make_raw_tuple(6, seed=860 + call)
-                if kind == "deferred_norm":
-                    packet, res = deferred_norm(*t, M.GLM_EPS)
-                    poisoned &= _all_nan_packet(packet, res)
-                else:
-                    out, res = post_moe_hop(*t, M.GLM_EPS)
-                    torch.cuda.synchronize()
-                    poisoned &= bool(
-                        torch.isnan(out.float()).all() and torch.isnan(res.float()).all()
-                    )
-        words = post_moe_hop.workspace.state.cpu().tolist()
-        carried = all(words[w] == 1 for w in STATE_ERROR_WORDS)
-        clean = not any(words[w] for w in STATE_ERROR_WORDS)
-        per_rank = env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])
-        if env.rank == 0:
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_POISON {sequence}: per-rank error words {per_rank}",
-                flush=True,
-            )
-        assert M.all_ranks(
-            poisoned
-            and (carried if env.rank == victim else clean)
-            and _acc_clear(post_moe_hop, mode)
-        )
-        with pytest.raises(PostMoeHopError):
-            check_health("routed_deferred_norm-poisoned", collective=True)
-        _reset(post_moe_hop)
-        src, shared, residual, gamma = _source(mode, 6, seed=870)
-        packet, res = _deferred_norm_call(
-            post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma
-        )
-        h_ref, res_ref = _post_moe_hop_reference(
-            post_moe_hop, mode, src, 6, shared, residual, gamma
-        )
-        assert _check(packet, res, h_ref, res_ref, gamma, "after-poison-reset")
-
-
-_FAULT_DELAYS_NS = (3000, 60000)
-
-
-@pytest.mark.parametrize("mode", ["fp32", "identity"])
-def test_routed_fault_injected_store_order_mixed_with_post_moe_hop(mode):
-    """Fault-injected deferred_norm / post_moe_hop routed builds chosen per rank and per call for
-    the same collective call, plus random device delays: every output bit as
-    in an undisturbed post_moe_hop routed call, one call count per call, no error."""
-    post_moe_hop, deferred_norm = _hops(mode)
-    hops = {
-        "post_moe_hop": (post_moe_hop, post_moe_hop),
-        "deferred_norm": (post_moe_hop, deferred_norm),
-    }
-    for d in _FAULT_DELAYS_NS:
-        faulty_post_moe, faulty_deferred_norm = _hops(mode, fault_delay_ns=d)
-        hops[f"post_moe_hop_fault{d}"] = (faulty_post_moe, faulty_post_moe)
-        hops[f"deferred_norm_fault{d}"] = (faulty_post_moe, faulty_deferred_norm)
-    env = M.mpi_env()
-    names = sorted(hops)
-    common = random.Random(4712)
-    local = random.Random(2000 + env.rank)
-    calls_before = post_moe_hop.workspace.calls()
-    failed, used, total = [], {}, 0
-    with torch.inference_mode():
-        for batch in range(5):
-            plan = [(common.randint(1, 8), common.random() < 0.75) for _ in range(12)]
-            inputs = [
-                _source(mode, r, seed=9100 + 100 * batch + j, with_shared=sh)
-                for j, (r, sh) in enumerate(plan)
-            ]
-            picks = [local.choice(names) for _ in plan]
-            torch.cuda.synchronize()
-            env.comm.barrier()
-            outs = []
-            for name, (r, _), t in zip(picks, plan, inputs):
-                if local.random() < 0.3:
-                    torch.cuda._sleep(int(local.uniform(0.0, 50.0) * 2000))
-                used[name] = used.get(name, 0) + 1
-                owner, hop = hops[name]
-                routed = _fill(owner, mode, t[0], r)
-                outs.append(hop.hop_routed(routed, t[1], t[2], t[3], M.GLM_EPS))
-            torch.cuda.synchronize()
-            for j, (name, (r, _), t, out) in enumerate(zip(picks, plan, inputs, outs)):
-                h_ref, res_ref = _post_moe_hop_reference(post_moe_hop, mode, t[0], r, *t[1:])
-                if name.startswith("deferred_norm"):
-                    ok = _verdict(out[0], out[1], h_ref, res_ref, t[3], f"fault{batch}.{j}")
-                else:
-                    torch.cuda.synchronize()
-                    ok = torch.equal(out[0], h_ref) and torch.equal(out[1], res_ref)
-                if not ok:
-                    failed.append((batch, j, name))
-            total += len(plan)
-        counted = post_moe_hop.workspace.calls() - calls_before
-    bits = post_moe_hop.error_bits()
-    if env.rank == 0:
-        print(
-            f"GLM_MOE_DSA_DEFERRED_NORM_ROUTED_FAULT mode={mode} calls={total} (+{total} reference calls) "
-            f"used_on_rank0={used}",
-            flush=True,
-        )
-    assert M.all_ranks(not failed and counted == 2 * total and bits == 0), (
-        failed[:4],
-        counted,
-        bits,
-    )
-
-
-@pytest.mark.parametrize(
-    "case", ("foreign_buffer", "row_count", "residual_dtype", "shared_shape", "gamma_shape")
-)
-def test_routed_contract_violation_raises_before_any_device_work(case):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopContractError,
-    )
-
-    mode = "fp32"
-    post_moe_hop, deferred_norm = _hops(mode)
-    rows = 6
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, rows, seed=920)
-        routed = _fill(post_moe_hop, mode, src, rows)
-        if case == "foreign_buffer":
-            routed = routed.clone()
-        elif case == "row_count":
-            routed = post_moe_hop.routed_acc_buffer(rows + 1)
-        elif case == "residual_dtype":
-            residual = residual.half()
-        elif case == "shared_shape":
-            shared = shared[:, : M.GLM_HIDDEN // 2]
-        else:
-            gamma = gamma[:-8]
-        calls_before = post_moe_hop.workspace.calls()
-        with pytest.raises(PostMoeHopContractError):
-            deferred_norm.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
-        assert M.all_ranks(post_moe_hop.workspace.calls() == calls_before)
-        post_moe_hop.routed_acc_buffer(post_moe_hop.max_rows).zero_()
-        src, shared, residual, gamma = _source(mode, rows, seed=921)
-        packet, res = _deferred_norm_call(
-            post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma
-        )
-        h_ref, res_ref = _post_moe_hop_reference(
-            post_moe_hop, mode, src, rows, shared, residual, gamma
-        )
-        assert _check(packet, res, h_ref, res_ref, gamma, f"after-contract-{case}")
-
-
-def _misaligned(t, offset_elems):
-    storage = torch.empty(t.numel() + 16, dtype=t.dtype, device=t.device)
-    view = storage[offset_elems : offset_elems + t.numel()].view(t.shape)
-    view.copy_(t)
-    return view
-
-
-def _transposed(t):
-    view = torch.empty(tuple(reversed(t.shape)), dtype=t.dtype, device=t.device).t()
-    view.copy_(t)
-    return view
-
-
-def _needs_copy(t):
-    """The hop's 16-byte vector rule for a 1-D / row-major operand."""
-    return bool(t.data_ptr() % 16 or t.stride(-1) != 1 or not t.is_contiguous())
-
-
-def _relayout(rank, shared, residual, gamma):
-    """Rank-dependent layouts; returns them and the repairs expected on this
-    rank (a one-row "transposed" tensor is still contiguous)."""
-    kind = rank % 4
-    if kind == 0:
-        gamma = _misaligned(gamma, 2)
-    elif kind == 1:
-        residual = _misaligned(residual, 4)
-    elif kind == 2 and shared is not None:
-        shared = _transposed(shared)
-    repairs = sum(int(_needs_copy(t)) for t in (shared, residual, gamma) if t is not None)
-    return shared, residual, gamma, repairs
-
-
-def test_routed_unaligned_layouts_are_repaired_locally():
-    mode = "fp32"
-    post_moe_hop, deferred_norm = _hops(mode)
-    env = M.mpi_env()
-    ok = True
-    with torch.inference_mode():
-        for i, rows in enumerate((6, 1, 8)):
-            src, shared, residual, gamma = _source(mode, rows, seed=880 + i)
-            r_shared, r_residual, r_gamma, expected = _relayout(env.rank, shared, residual, gamma)
-            before = post_moe_hop.workspace.layout_repairs
-            packet, res = _deferred_norm_call(
-                post_moe_hop, deferred_norm, mode, src, rows, r_shared, r_residual, r_gamma
-            )
-            repaired = post_moe_hop.workspace.layout_repairs - before
-            h_ref, res_ref = _post_moe_hop_reference(
-                post_moe_hop, mode, src, rows, shared, residual, gamma
-            )
-            ok &= _verdict(packet, res, h_ref, res_ref, gamma, f"relayout{rows}")
-            ok &= repaired == expected
-        rows = 6
-        src, shared, residual, gamma = _source(mode, rows, seed=890)
-        s_src = src.clone()
-        s_shared, s_residual, s_gamma, _ = _relayout(env.rank, shared, residual, gamma)
-        _deferred_norm_call(
-            post_moe_hop, deferred_norm, mode, s_src, rows, s_shared, s_residual, s_gamma
-        )
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            g_packet, g_res = _deferred_norm_call(
-                post_moe_hop, deferred_norm, mode, s_src, rows, s_shared, s_residual, s_gamma
-            )
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        for replay in range(4):
-            f_src, f_shared, f_residual, _ = _source(mode, rows, seed=900 + replay)
-            s_src.copy_(f_src)
-            s_shared.copy_(f_shared)
-            s_residual.copy_(f_residual)
-            graph.replay()
-            torch.cuda.synchronize()
-            h_ref, res_ref = _post_moe_hop_reference(
-                post_moe_hop, mode, f_src, rows, f_shared, f_residual, gamma
-            )
-            ok &= _verdict(g_packet, g_res, h_ref, res_ref, gamma, f"relayout-graph{replay}")
-    assert M.all_ranks(ok and post_moe_hop.error_bits() == 0)
-
-
-def test_health_check_counts_deferred_norm_routed_calls():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        check_health,
-    )
-
-    mode = "fp32"
-    post_moe_hop, deferred_norm = _hops(mode)
-    with torch.inference_mode():
-        before = check_health("routed_deferred_norm-engagement-before", collective=True)
-        for i in range(5):
-            src, shared, residual, gamma = _source(mode, 6, seed=940 + i)
-            _deferred_norm_call(post_moe_hop, deferred_norm, mode, src, 6, shared, residual, gamma)
-        after = check_health("routed_deferred_norm-engagement-after", collective=True)
-    assert M.all_ranks(
-        after.calls - before.calls == 5
-        and after.error_bits == 0
-        and not after.group_error
-        and after.calls == post_moe_hop.workspace.calls()
-    )
-
-
-def test_routed_memcheck_positive_control():
-    """Tests only, run under compute-sanitizer memcheck by the sanitizer
-    suite (``GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL=1``): the routed variant reads (and, in
-    fp32 mode, clears) an exactly sized accumulator with twice its row
-    pitch, so row 1 is out of bounds and memcheck must report it."""
-    if os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_MEMCHECK_CONTROL") != "1":
-        pytest.skip("memcheck positive control (sanitizer suite only)")
-    mode = "fp32"
-    post_moe_hop, deferred_norm = _hops(mode)
-    rows = 2
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, rows, seed=990)
-        acc = src.clone()
-        deferred_norm.launch_routed(
-            acc,
-            shared,
-            residual,
-            gamma,
-            M.GLM_EPS,
-            row_bytes=2 * acc.stride(0) * acc.element_size(),
-        )
-        torch.cuda.synchronize()
-    _reset(post_moe_hop)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py
deleted file mode 100644
index a1b2b6506bfd..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope.py
+++ /dev/null
@@ -1,327 +0,0 @@
-# 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.
-"""Single-process tests of the deferred_norm scope (``deferred_norm/scope.py``).
-
-The scope logic (this target's ``shared`` scope, the next layer's indexer type, the constructing
-decoder layer found from the calling frames) needs only the standard library and the target's
-``dispatch`` and runs without a built ``tensorrt_llm``. The remaining tests need the installed
-package and the GLM-5.2 checkpoint:
-
-* the GLM-5.2 indexer layout seen by the scope (the model config's
-  ``to_sparse_params``) has 18 full-indexer layers among the 74 next layers of
-  the deferred_norm boundaries (6, 10, ..., 74) and agrees with the native attention
-  module (an ``Indexer`` exactly on full layers);
-* the consumer guard: under ``shared``, this target's scope, the deferred_norm projection
-  declines a full-indexer layer (so a packet there is materialized for the
-  native path) and accepts shared layers.
-"""
-
-import importlib
-import importlib.util
-import sys
-import types
-from pathlib import Path
-
-import pytest
-
-REPO = Path(__file__).resolve().parents[9]
-TARGET = (
-    REPO
-    / "tensorrt_llm"
-    / "_torch"
-    / "_experimental"
-    / "modeling_v2"
-    / "models"
-    / "glm_moe_dsa"
-    / "shared"
-)
-_STANDALONE = "_deferred_norm_scope_standalone"
-# GLM-5.2: 78 target layers, full-indexer layers where max(i - 2, 0) % 4 == 0.
-GLM52 = dict(
-    num_hidden_layers=78, index_topk_freq=4, index_skip_topk_offset=3, index_topk_pattern=None
-)
-FULL_NEXT = list(range(6, 78, 4))
-BOUNDARIES = list(range(3, 77))  # L -> L+1 with a deferred_norm hop (first_k_dense_replace = 3)
-
-
-def _load_dispatch_standalone(name):
-    """The target's ``dispatch.py`` as module ``name``, without importing ``tensorrt_llm``."""
-    spec = importlib.util.spec_from_file_location(
-        "deferred_norm_scope_input_prep_testlib",
-        Path(__file__).resolve().parents[1] / "deferred_input_prep" / "input_prep_testlib.py",
-    )
-    testlib = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(testlib)
-    return testlib.load_dispatch_standalone(name)
-
-
-def _load_scope():
-    """``scope.py``: the installed package, else the source tree under a stand-in package."""
-    try:
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-            scope,
-        )
-
-        return scope
-    except (ImportError, OSError):
-        pass
-    if _STANDALONE not in sys.modules:
-        for name, path in (
-            (_STANDALONE, TARGET),
-            (f"{_STANDALONE}.kernels", TARGET / "kernels"),
-            (f"{_STANDALONE}.kernels.deferred_norm", TARGET / "kernels" / "deferred_norm"),
-        ):
-            package = types.ModuleType(name)
-            package.__path__ = [str(path)]
-            sys.modules[name] = package
-        sys.modules[f"{_STANDALONE}.dispatch"] = _load_dispatch_standalone(
-            f"{_STANDALONE}.dispatch"
-        )
-    return importlib.import_module(f"{_STANDALONE}.kernels.deferred_norm.scope")
-
-
-@pytest.fixture
-def scope():
-    return _load_scope()
-
-
-class _Sparse:
-    """``to_sparse_params`` with GLM-5.2's frequency rule; records its calls."""
-
-    def __init__(self):
-        self.calls = []
-
-    def to_sparse_params(self, pretrained_config=None, layer_idx=None):
-        self.calls.append((pretrained_config, layer_idx))
-        offset = pretrained_config.index_skip_topk_offset
-        full = max(layer_idx - offset + 1, 0) % pretrained_config.index_topk_freq == 0
-        return types.SimpleNamespace(is_full_indexer_layer=full)
-
-
-def _model_config(sparse=None):
-    return types.SimpleNamespace(
-        sparse_attention_config=sparse or _Sparse(),
-        pretrained_config=types.SimpleNamespace(**GLM52),
-    )
-
-
-class _Hop:
-    """Stands in for ``DeferredNormHop.__init__``: resolves the next layer like it."""
-
-    def __init__(self, scope, post_moe_hop, arg=None):
-        self.kind = scope.next_layer_kind(post_moe_hop, arg)
-
-
-class _DecoderLayer:
-    """The order of ``GlmDecoderLayer._adopt_post_moe``: post_moe_hop first, then deferred_norm."""
-
-    def __init__(self, scope, layer_idx, model_config, post_moe_hop):
-        self.model_config = model_config
-        self.layer_idx = layer_idx
-        self.post_moe_hop = post_moe_hop
-        self.deferred_norm = _Hop(scope, self.post_moe_hop)
-
-
-def test_the_scope_of_this_target(scope):
-    assert scope.SCOPES == ("all", "shared") and scope.DEFAULT_SCOPE == "shared"
-    assert scope.UNRESTRICTED_SCOPE == "all"
-    assert scope.SCOPE_LABEL == "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE"
-    assert scope.selected_scope() == "shared"
-
-
-def test_the_record_names_the_scope_of_this_target(scope):
-    """The record's deferred_norm variant is the scope the hop runs, whatever the environment
-    holds (``resolve`` refuses the variable)."""
-    try:
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    except (ImportError, OSError):
-        dispatch = _load_dispatch_standalone("dispatch_standalone")
-    assert dispatch.NUMERICAL_VARIANTS["deferred_norm"] == (None, (scope.DEFAULT_SCOPE,))
-    for raw in (None, "", " ", "SHARED", " shared ", "all", "bogus"):
-        env = {} if raw is None else {scope.SCOPE_LABEL: raw}
-        assert dispatch.feature_variants(["deferred_norm"], env) == {
-            "deferred_norm": scope.DEFAULT_SCOPE
-        }
-
-
-def test_glm52_frequency_rule_boundaries(scope):
-    cfg = _model_config()
-    full = {i for i in range(78) if scope.is_full_indexer_layer(cfg, i)}
-    assert full == {0, 1, 2} | set(FULL_NEXT)
-    into_full = [b for b in BOUNDARIES if scope.is_full_indexer_layer(cfg, b + 1)]
-    assert into_full == [n - 1 for n in FULL_NEXT] and len(into_full) == 18
-    assert len(BOUNDARIES) - len(into_full) == 56
-
-
-def test_unknown_layout(scope):
-    assert (
-        scope.is_full_indexer_layer(types.SimpleNamespace(sparse_attention_config=None), 6) is None
-    )
-    params = types.SimpleNamespace(to_sparse_params=lambda **kw: types.SimpleNamespace())
-    assert (
-        scope.is_full_indexer_layer(
-            types.SimpleNamespace(sparse_attention_config=params, pretrained_config=None), 6
-        )
-        is None
-    )
-
-
-@pytest.mark.parametrize("layer_idx", [3, 5, 6, 9, 73, 75, 76])
-def test_next_layer_from_the_constructing_layer(scope, layer_idx):
-    hop_stub = object()
-    cfg = _model_config()
-    layer = _DecoderLayer(scope, layer_idx, cfg, hop_stub)
-    full, nxt, how = layer.deferred_norm.kind
-    assert nxt == layer_idx + 1 and how == "model config"
-    assert full is ((layer_idx + 1) in FULL_NEXT)
-    # The same call as MLA.__init__ of the next layer.
-    assert cfg.sparse_attention_config.calls == [(cfg.pretrained_config, layer_idx + 1)]
-
-
-def test_constructing_layer_must_hold_this_post_moe_hop(scope):
-    hop_stub, other = object(), object()
-    unresolved = (None, None, "no constructing decoder layer found")
-    assert _Hop(scope, hop_stub).kind == unresolved
-
-    class _OtherHop:
-        def __init__(self):
-            self.model_config = _model_config()
-            self.layer_idx = 5
-            self.post_moe_hop = other
-            self.deferred_norm = _Hop(scope, hop_stub)
-
-    assert _OtherHop().deferred_norm.kind == unresolved
-
-    class _NoIndex:
-        def __init__(self):
-            self.post_moe_hop = hop_stub
-            self.layer_idx = None
-            self.model_config = _model_config()
-            self.deferred_norm = _Hop(scope, hop_stub)
-
-    assert _NoIndex().deferred_norm.kind == unresolved
-
-
-def test_explicit_argument_wins(scope):
-    hop_stub = object()
-
-    class _Layer:
-        def __init__(self, arg):
-            self.model_config = _model_config()
-            self.layer_idx = 5  # next layer 6 is full
-            self.post_moe_hop = hop_stub
-            self.deferred_norm = _Hop(scope, hop_stub, arg)
-
-    assert _Layer(False).deferred_norm.kind == (False, None, "argument")
-    assert _Layer(True).deferred_norm.kind == (True, None, "argument")
-    assert _Layer(None).deferred_norm.kind == (True, 6, "model config")
-
-
-def test_unresolved_config(scope):
-    hop_stub = object()
-    cfg = types.SimpleNamespace(sparse_attention_config=None, pretrained_config=None)
-    layer = _DecoderLayer(scope, 5, cfg, hop_stub)
-    assert layer.deferred_norm.kind[0] is None and layer.deferred_norm.kind[1] == 6
-
-
-def test_decline_reason(scope):
-    assert scope.decline_reason("all", True, 6, "model config") is None
-    assert scope.decline_reason("all", None, None, "x") is None
-    assert scope.decline_reason("shared", False, 7, "model config") is None
-    why = scope.decline_reason("shared", True, 6, "model config")
-    assert "layer 6 is a full-indexer layer" in why and "post_moe_hop hop" in why
-    why = scope.decline_reason("shared", None, None, "no constructing decoder layer found")
-    assert "unknown" in why and "no constructing decoder layer found" in why
-
-
-# ---- installed package and checkpoint ------------------------------------------------
-
-
-def _installed():
-    try:
-        import torch  # noqa: F401
-
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501, F401
-            projection,
-        )
-    except (ImportError, OSError) as exc:
-        pytest.skip(f"needs the installed tensorrt_llm ({exc!r})")
-    import absorbed_query_fixture
-
-    if not absorbed_query_fixture.checkpoint_available():
-        pytest.skip("needs the GLM-5.2 checkpoint")
-    return absorbed_query_fixture
-
-
-_ATTN = {}
-
-
-def _attention(absorbed_query_fixture, layer):
-    import torch
-
-    if not torch.cuda.is_available():
-        pytest.skip("needs a GPU")
-    if "cfg" not in _ATTN:
-        _ATTN["cfg"] = absorbed_query_fixture.build_model_config(0)
-    if layer not in _ATTN:
-        with torch.inference_mode():
-            _ATTN[layer] = absorbed_query_fixture.build_native_attention(layer, 0, _ATTN["cfg"])
-    return _ATTN["cfg"], _ATTN[layer]
-
-
-def test_glm52_model_config_layout(scope):
-    """The scope's view of the real GLM-5.2 model config (no GPU work)."""
-    absorbed_query_fixture = _installed()
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        scope as installed,
-    )
-
-    cfg = absorbed_query_fixture.build_model_config(0)
-    full = [b + 1 for b in BOUNDARIES if installed.is_full_indexer_layer(cfg, b + 1)]
-    assert full == FULL_NEXT, full
-    assert len(BOUNDARIES) - len(full) == 56
-
-
-@pytest.mark.parametrize("layer", [6, 7])
-def test_scope_layout_matches_the_attention_module(layer):
-    """``is_full_indexer_layer`` is true exactly where the native module has an indexer."""
-    absorbed_query_fixture = _installed()
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        scope as installed,
-    )
-
-    cfg, attn = _attention(absorbed_query_fixture, layer)
-    has_indexer = attn.mqa.indexer is not None
-    assert installed.is_full_indexer_layer(cfg, layer) is has_indexer
-    assert has_indexer is (layer in FULL_NEXT)
-
-
-@pytest.mark.parametrize("layer", [6, 7])
-def test_consumer_guard(layer):
-    """Under ``shared`` the deferred_norm projection declines full-indexer layers only."""
-    absorbed_query_fixture = _installed()
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.scope import (  # noqa: E501
-        SCOPE_LABEL,
-    )
-
-    _, attn = _attention(absorbed_query_fixture, layer)
-    ok, reason, kw = projection_eligibility(attn)
-    if layer in FULL_NEXT:
-        assert not ok and kw is None and f"{SCOPE_LABEL}=shared" in reason, reason
-    else:
-        assert ok and kw is None, reason
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py
deleted file mode 100644
index e265c090fda3..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm/test_deferred_norm_scope_mpi.py
+++ /dev/null
@@ -1,342 +0,0 @@
-# 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.
-"""Collective tests of the deferred_norm scope (8 ranks over 2 hosts, TP8).
-
-The ``shared`` scope, this target's, keeps the deferred-norm hop only on boundaries whose next
-layer is a shared-indexer layer. The tests that compare against the hop's unrestricted scope
-``all`` build the hop with that scope directly. Checks:
-
-* unrestricted scope: the hop is the default deferred_norm hop (engaged, no next-layer
-  lookup, no per-boundary log key);
-* ``shared``, next layer full: the hop is disengaged on every rank and the
-  layer's choice (``can_run`` / ``can_run_routed`` false; the decoder layer
-  keeps ``None``) is post_moe_hop's hop, whose output is post_moe_hop's by construction; the
-  routed input is never enabled on such a hop;
-* ``shared``, next layer shared: the hop is engaged and its packets (raw
-  tuple and routed_moe_into_hop routed input in ``fp32`` and ``identity``) are bit-identical
-  to the unrestricted hop's packets on the same inputs and workspace, and
-  their materialization equals post_moe_hop's normalized output (``a_frag``, ``sumsq``,
-  residual, ``materialize``);
-* rank-differing scopes: one rank with another scope or an unknown
-  next-layer type can only disengage a boundary on every rank (one agreement,
-  the same number of collectives in every scope);
-* the decoder-layer path: a stand-in with the target decoder layer's adoption order
-  (``GlmDecoderLayer._adopt_post_moe``: post_moe_hop assigned, then the deferred_norm factory
-  called with only the post_moe_hop) and the real GLM-5.2 model config resolves layers 5, 9, 73 ->
-  full (post_moe_hop) and 3, 6, 76 -> shared (deferred_norm hop) on every rank.
-"""
-
-import importlib
-import os
-import sys
-import types
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
-sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
-
-import post_moe_hop_mpi_utils as M  # noqa: E402
-
-FULL_NEXT = set(range(6, 78, 4))
-
-
-def _env():
-    env = M.mpi_env()
-    if env.world != 8:
-        pytest.skip("the deferred_norm hop is specialized to TP8")
-    return env
-
-
-_POST_MOE_HOP = {}
-
-
-def _post_moe_hop(mode=None):
-    """One post_moe_hop per routed-input mode (``None``: raw tuple only), collective."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    if mode not in _POST_MOE_HOP:
-        hop_stub = PostMoeHop(env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False)
-        ok = hop_stub.engaged
-        if mode is not None:
-            if not hasattr(hop_stub, "enable_routed_input"):
-                pytest.skip(
-                    "the installed post_moe_hop hop has no routed_moe_into_hop routed input"
-                )
-            ok = ok and hop_stub.enable_routed_input(mode)
-        assert M.all_ranks(ok), hop_stub.reason
-        _POST_MOE_HOP[mode] = hop_stub
-    return _POST_MOE_HOP[mode]
-
-
-UNRESTRICTED = "all"
-_HOP_MODULE = (
-    "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa."
-    "shared.kernels.deferred_norm.hop"
-)
-
-
-def _hop(monkeypatch, hop_stub, scope_value, **kwargs):
-    """A deferred_norm hop built under this target's scope (``scope_value`` None or ``shared``).
-
-    ``UNRESTRICTED``: the hop's unrestricted scope, which this target does not run; the hop
-    module's scope reader returns it for this construction.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-
-    with monkeypatch.context() as patch:
-        if scope_value == UNRESTRICTED:
-            hop_module = importlib.import_module(_HOP_MODULE)
-            patch.setattr(hop_module, "selected_scope", lambda: UNRESTRICTED)
-        return DeferredNormHop(hop_stub, **kwargs)
-
-
-def _metadata(rows):
-    return types.SimpleNamespace(all_rank_num_tokens=None, num_contexts=0, num_generations=rows)
-
-
-def _same_packet(pa, ra, pb, rb):
-    """Same packet content (the fragment buffer holds 8 rows; rows >= R are not written)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    torch.cuda.synchronize()
-    rows, k = ra.shape
-    return (
-        torch.equal(rows_from_frag(pa.a_frag, rows, k), rows_from_frag(pb.a_frag, rows, k))
-        and torch.equal(pa.sumsq, pb.sumsq)
-        and torch.equal(ra, rb)
-        and pa.eps == pb.eps
-    )
-
-
-def _materialized_is_post_moe_hop(packet, h_post_moe_hop):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        materialize,
-    )
-
-    h = materialize(packet)
-    torch.cuda.synchronize()
-    return torch.equal(h, h_post_moe_hop)
-
-
-def test_default_scope_is_the_b11_hop(monkeypatch):
-    hop_stub = _post_moe_hop()
-    # The next-layer argument is ignored by the unrestricted scope.
-    hop = _hop(monkeypatch, hop_stub, UNRESTRICTED, next_layer_full_indexer=True)
-    ok = (
-        hop.engaged
-        and hop.scope == "all"
-        and hop.next_layer_full is None
-        and hop.next_layer_idx is None
-        and hop.reason == ""
-    )
-    assert M.all_ranks(ok), (hop.scope, hop.reason)
-
-
-def test_shared_scope_full_next_layer_keeps_post_moe_hop(monkeypatch):
-    hop_stub = _post_moe_hop()
-    hop = _hop(monkeypatch, hop_stub, None, next_layer_full_indexer=True)
-    ok = (
-        not hop.engaged
-        and hop.scope == "shared"
-        and hop.next_layer_full is True
-        and "full-indexer layer" in hop.reason
-        and not hop.can_run(6, _metadata(6))
-        and not hop.can_run_routed(6, _metadata(6))
-    )
-    assert M.all_ranks(ok), hop.reason
-    # The decoder layer enables the routed input only on a kept (engaged) hop;
-    # called anyway, the enable declines on every rank without a collective.
-    assert M.all_ranks(not hop.enable_routed_input() and not hop.routed_engaged)
-
-
-@pytest.mark.parametrize("with_shared", [True, False])
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-def test_shared_scope_shared_next_layer_equals_default(monkeypatch, rows, with_shared):
-    hop_stub = _post_moe_hop()
-    default = _hop(monkeypatch, hop_stub, UNRESTRICTED)
-    shared = _hop(monkeypatch, hop_stub, None, next_layer_full_indexer=False)
-    assert M.all_ranks(
-        default.engaged
-        and shared.engaged
-        and shared.next_layer_full is False
-        and shared.can_run(rows, _metadata(rows))
-    ), shared.reason
-    with torch.inference_mode():
-        fc2, scales, perm, sh, residual, gamma = M.make_raw_tuple(
-            rows, seed=900 + rows, with_shared=with_shared
-        )
-        h_post_moe_hop, res_post_moe_hop = hop_stub(
-            fc2, scales, perm, sh, residual, gamma, M.GLM_EPS
-        )
-        pa, ra = default(fc2, scales, perm, sh, residual, gamma, M.GLM_EPS)
-        pb, rb = shared(fc2, scales, perm, sh, residual, gamma, M.GLM_EPS)
-        ok = (
-            _same_packet(pa, ra, pb, rb)
-            and torch.equal(rb, res_post_moe_hop)
-            and _materialized_is_post_moe_hop(pb, h_post_moe_hop)
-        )
-    assert M.all_ranks(ok)
-
-
-def _routed_source(mode, rows, seed):
-    fc2, scales, perm, shared, residual, gamma = M.make_raw_tuple(rows, seed=seed)
-    if mode == "identity":
-        valid = perm >= 0
-        safe = torch.where(valid, perm, torch.zeros_like(perm)).long()
-        src = (
-            fc2[safe.view(-1)].contiguous(),
-            torch.where(valid, scales, torch.zeros_like(scales)).bfloat16().contiguous(),
-        )
-    else:
-        valid = (perm >= 0).unsqueeze(-1)
-        safe = torch.where(perm >= 0, perm, torch.zeros_like(perm)).long()
-        prod = fc2[safe.view(-1)].view(rows, M.GLM_TOP_K, -1).float() * scales.float().unsqueeze(-1)
-        src = torch.where(valid, prod, torch.zeros_like(prod)).sum(1).contiguous()
-    return src, shared, residual, gamma
-
-
-def _fill(hop_stub, mode, src, rows):
-    if mode == "identity":
-        buf, scales = hop_stub.routed_raw_buffers(rows)
-        buf.copy_(src[0])
-        scales.copy_(src[1])
-        return buf
-    buf = hop_stub.routed_acc_buffer(rows)
-    buf.copy_(src)
-    return buf
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("mode", ["fp32", "identity"])
-def test_shared_scope_routed_equals_default(monkeypatch, mode, rows):
-    hop_stub = _post_moe_hop(mode)
-    default = _hop(monkeypatch, hop_stub, UNRESTRICTED)
-    shared = _hop(monkeypatch, hop_stub, None, next_layer_full_indexer=False)
-    ok = default.enable_routed_input() and shared.enable_routed_input()
-    assert M.all_ranks(ok and shared.can_run_routed(rows, _metadata(rows))), shared.routed_reason
-    with torch.inference_mode():
-        src, sh, residual, gamma = _routed_source(mode, rows, seed=950 + rows)
-        h_post_moe_hop, res_post_moe_hop = hop_stub.hop_routed(
-            _fill(hop_stub, mode, src, rows), sh, residual, gamma, M.GLM_EPS
-        )
-        pa, ra = default.hop_routed(
-            _fill(hop_stub, mode, src, rows), sh, residual, gamma, M.GLM_EPS
-        )
-        pb, rb = shared.hop_routed(_fill(hop_stub, mode, src, rows), sh, residual, gamma, M.GLM_EPS)
-        ok = (
-            _same_packet(pa, ra, pb, rb)
-            and torch.equal(rb, res_post_moe_hop)
-            and _materialized_is_post_moe_hop(pb, h_post_moe_hop)
-        )
-        if mode == "fp32":
-            torch.cuda.synchronize()
-            ok = ok and bool((hop_stub.routed_acc_buffer(hop_stub.max_rows) == 0).all())
-    assert M.all_ranks(ok)
-
-
-@pytest.mark.parametrize("next_full", [True, False])
-@pytest.mark.parametrize("odd_rank", [0, 3, 7])
-def test_rank_differing_scope_agrees(monkeypatch, odd_rank, next_full):
-    """One rank with the unrestricted scope: a full boundary is post_moe_hop everywhere, a
-    shared boundary is deferred_norm everywhere."""
-    env = _env()
-    hop_stub = _post_moe_hop()
-    value = UNRESTRICTED if env.rank == odd_rank else "shared"
-    hop = _hop(monkeypatch, hop_stub, value, next_layer_full_indexer=next_full)
-    if next_full:
-        local = (
-            "another rank is not capable" == hop.reason
-            if env.rank == odd_rank
-            else "full-indexer layer" in hop.reason
-        )
-        assert M.all_ranks(not hop.engaged and local), hop.reason
-    else:
-        assert M.all_ranks(hop.engaged), hop.reason
-
-
-@pytest.mark.parametrize("odd_rank", [2, 5])
-def test_unresolved_next_layer_on_one_rank(monkeypatch, odd_rank):
-    """A rank that cannot resolve the next layer keeps post_moe_hop on every rank."""
-    env = _env()
-    hop_stub = _post_moe_hop()
-    arg = None if env.rank == odd_rank else False
-    hop = _hop(monkeypatch, hop_stub, "shared", next_layer_full_indexer=arg)
-    if env.rank == odd_rank:
-        local = "unknown" in hop.reason and "no constructing decoder layer found" in hop.reason
-    else:
-        local = hop.reason == "another rank is not capable"
-    assert M.all_ranks(not hop.engaged and local), hop.reason
-
-
-def _glm52_model_config():
-    import absorbed_query_fixture
-
-    if not absorbed_query_fixture.checkpoint_available():
-        pytest.skip("needs the GLM-5.2 checkpoint config")
-    return absorbed_query_fixture.build_model_config(M.mpi_env().rank)
-
-
-class _DecoderLayerStandIn:
-    """``GlmDecoderLayer._adopt_post_moe``'s deferred_norm wiring, in its order."""
-
-    def __init__(self, model_config, layer_idx, post_moe_hop):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-        self.model_config = model_config
-        self.layer_idx = layer_idx
-        self.post_moe_hop = post_moe_hop
-        hop = dispatch.get_impl("deferred_norm", "deferred_hop")(self.post_moe_hop)
-        self.deferred_norm_hop = hop if hop.engaged else None
-        self.deferred_norm_reason = hop.reason
-
-
-def test_decoder_layer_resolves_the_next_layer():
-    """Under ``shared``, this target's scope, full next layers keep the post-MoE hop."""
-    hop_stub = _post_moe_hop()
-    cfg = _glm52_model_config()
-    ok = True
-    for layer_idx in (3, 5, 6, 9, 73, 76):
-        layer = _DecoderLayerStandIn(cfg, layer_idx, hop_stub)
-        hop = layer.deferred_norm_hop
-        if (layer_idx + 1) not in FULL_NEXT:
-            good = (
-                hop is not None
-                and hop.next_layer_idx == layer_idx + 1
-                and hop.next_layer_full is False
-            )
-        else:
-            good = (
-                hop is None
-                and f"layer {layer_idx + 1} is a full-indexer layer" in layer.deferred_norm_reason
-            )
-        if not good:
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_SCOPE_LAYER_FAIL rank={M.mpi_env().rank} layer={layer_idx} "
-                f"reason={layer.deferred_norm_reason}",
-                flush=True,
-            )
-        ok = ok and good
-    assert M.all_ranks(ok)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py
deleted file mode 100644
index f0b0b59ecdf9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_adapter.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# 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.
-"""Control flow of the deferred_norm_evict_hints call sites in the target.
-
-Without device work, checks that
-
-* the plan runs deferred_norm_evict_hints with its closure ``(absorbed_query, post_moe_hop,
-  deferred_norm)`` and turns it off with any of them;
-* the DSA projection entry (``attention_forward._absorbed_query_packet``) hands a deferred_norm
-  packet to deferred_norm_evict_hints's ``dsa_proj``
-  exactly when deferred_norm_evict_hints is selected (deferred_norm's otherwise), and leaves other
-  hidden states to the next entry;
-* the layer adoption (``_adopt_post_moe``) asks deferred_norm_evict_hints's ``deferred_hop`` factory when
-  deferred_norm_evict_hints is selected (deferred_norm's otherwise), with the decision taken on the
-  host before the factory runs;
-* deferred_input_prep knows deferred_norm_evict_hints's package directories and side selectors.
-"""
-
-import inspect
-import os
-import re
-import sys
-
-import pytest
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import plan_testlib  # noqa: E402
-
-BASE = "absorbed_query,post_moe_hop,routed_moe_into_hop,deferred_norm,routed_moe"
-
-
-@pytest.fixture
-def select(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-
-    def _select(features):
-        plan_testlib.select(dispatch, features, monkeypatch)
-        return dispatch
-
-    return _select
-
-
-def test_selector_closure(select):
-    dispatch = select(BASE + ",deferred_norm_evict_hints")
-    assert dispatch.is_enabled("deferred_norm_evict_hints") and dispatch.is_enabled("deferred_norm")
-    assert dispatch.REQUIRED_FEATURES["deferred_norm_evict_hints"] == frozenset(
-        {"absorbed_query", "post_moe_hop", "deferred_norm"}
-    )
-    for base in ("absorbed_query", "post_moe_hop", "deferred_norm"):
-        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
-        assert not plan.enabled("deferred_norm_evict_hints"), base
-
-
-@pytest.mark.parametrize("with_evict_hints", [True, False])
-def test_dsa_entry_routes_deferred_norm_packets(select, monkeypatch, with_evict_hints):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-        attention_forward,
-    )
-
-    dispatch = select(BASE + (",deferred_norm_evict_hints" if with_evict_hints else ""))
-    calls = []
-    packet = object()
-    real_get_impl = dispatch.get_impl
-
-    def fake_get_impl(feature, name):
-        if name == "is_deferred":
-            return lambda h: h is packet
-        if name == "dsa_proj" and feature in ("deferred_norm", "deferred_norm_evict_hints"):
-            return lambda *args: calls.append(feature) or f"{feature}-packet"
-        return real_get_impl(feature, name)
-
-    monkeypatch.setattr(dispatch, "get_impl", fake_get_impl)
-    out = attention_forward._absorbed_query_packet(object(), None, packet, None)
-    expect = "deferred_norm_evict_hints" if with_evict_hints else "deferred_norm"
-    assert out == f"{expect}-packet" and calls == [expect]
-
-
-def test_decoder_layer_selects_the_hop_factory():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmDecoderLayer,
-    )
-
-    src = inspect.getsource(GlmDecoderLayer._adopt_post_moe)
-    # The formatter wraps long statements: compare with the whitespace collapsed.
-    src = re.sub(r"(?<=[(\[])\s+|\s+(?=[)\]])", "", " ".join(src.split()))
-    assert (
-        '"deferred_norm_evict_hints" if dispatch.is_enabled("deferred_norm_evict_hints") else "deferred_norm"'
-        in src
-    )
-    # The choice is made before the factory is called (host-side selector).
-    head, tail = src.split(
-        '"deferred_norm_evict_hints" if dispatch.is_enabled("deferred_norm_evict_hints") else "deferred_norm"',
-        1,
-    )
-    assert head.rstrip().endswith("get_impl(")
-    assert re.match(r', "deferred_hop",?\)\(self\.post_moe_hop\)', tail), tail[:60]
-
-
-def test_evict_hints_registers_its_entry_points(select):
-    dispatch = select(BASE + ",deferred_norm_evict_hints")
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        is_deferred,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
-        EvictHintsDeferredHop,
-        evict_hints_dsa_proj,
-    )
-
-    assert dispatch.get_impl("deferred_norm_evict_hints", "deferred_hop") is EvictHintsDeferredHop
-    assert dispatch.get_impl("deferred_norm_evict_hints", "dsa_proj") is evict_hints_dsa_proj
-    assert dispatch.get_impl("deferred_norm_evict_hints", "is_deferred") is is_deferred
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py
deleted file mode 100644
index 79ad8ea35702..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_config.py
+++ /dev/null
@@ -1,85 +0,0 @@
-# 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.
-"""Host-only tests of the deferred_norm_evict_hints configuration parsing (no GPU work)."""
-
-import dataclasses
-
-
-def test_pre_wait_prefetch_is_a_test_arm():
-    """The served hop configuration prefetches nothing before its dependency wait; tests build
-    the pre-wait arms explicitly."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.hop import (  # noqa: E501
-        PRE_WAIT_PREFETCH,
-        EvictHintsHopConfig,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
-        parse_cfg,
-    )
-
-    assert EvictHintsHopConfig().poll_prefetch == "none"
-    assert EvictHintsHopConfig().poll_prefetch not in PRE_WAIT_PREFETCH
-    for mode in PRE_WAIT_PREFETCH + ("none", "wait3", "t3"):
-        assert parse_cfg(f"poll_prefetch={mode}", EvictHintsHopConfig).poll_prefetch == mode
-
-
-def test_defaults_are_deferred_norm_launch():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormProjConfig,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
-        EvictHintsHopConfig,
-        EvictHintsProjConfig,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
-        EVICT_HINTS_CONFIG_KEYS,
-    )
-
-    z = EvictHintsProjConfig()
-    w = DeferredNormProjConfig()
-    for f in dataclasses.fields(DeferredNormProjConfig):
-        assert getattr(z, f.name) == getattr(w, f.name), f.name
-    extra = [
-        f.name
-        for f in dataclasses.fields(EvictHintsProjConfig)
-        if f.name not in {g.name for g in dataclasses.fields(DeferredNormProjConfig)}
-    ]
-    assert tuple(extra) == EVICT_HINTS_CONFIG_KEYS
-    # Served default: deferred_norm's launch plus the evict_first policy on the weight copies.
-    assert not any(getattr(z, k) for k in EVICT_HINTS_CONFIG_KEYS if k != "copy_hint")
-    assert z.copy_hint == "evict_first"
-    h = EvictHintsHopConfig()
-    assert (h.carveout, h.poll_prefetch, h.clear_hint, h.ctas, h.ffi, h.dyn_smem) == (
-        -1,
-        "none",
-        "none",
-        8,
-        True,
-        0,
-    )
-
-
-def test_full_indexer_layers_keep_deferred_norm_launch():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
-        EvictHintsProjConfig,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
-        config_for,
-        parse_cfg,
-    )
-
-    cfg = parse_cfg("prod_warp=1,split_chunk=1,split_delay_ns=3000", EvictHintsProjConfig)
-    assert config_for(cfg, 0) is cfg
-    assert config_for(cfg, 160) == EvictHintsProjConfig(copy_hint="none")
-    assert config_for(EvictHintsProjConfig(), 160).copy_hint == "none"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py
deleted file mode 100644
index 26bc7fda6289..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_hop_mpi.py
+++ /dev/null
@@ -1,424 +0,0 @@
-# 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.
-"""Collective tests of the deferred_norm_evict_hints hop against F's deferred_norm hop (8 ranks, 2 hosts).
-
-For every deferred_norm_evict_hints hop configuration (``CONFIGS``: four experiment arms, the
-defaults ``carveout=-1,poll_prefetch=none`` through TVM-FFI, and the
-co-resident executor hop ``carveout=100,poll_prefetch=none,ffi=0``), on the
-same post_moe_hop workspace and the same inputs as deferred_norm's hop (routed_moe_into_hop routed
-input in each mode, and the raw tuple), every test below runs per configuration:
-
-* ``residual_out``, ``a_frag``, ``sumsq`` and ``materialize()`` are bit-identical
-  to deferred_norm's (and so to post_moe_hop's ``hop_routed`` / ``__call__``, which deferred_norm's
-  own suite checks); the fp32 accumulator is zero after the call;
-* deferred_norm_evict_hints, deferred_norm and post_moe_hop calls interleave on one workspace
-  without host syncs;
-* a CUDA graph of deferred_norm_evict_hints / deferred_norm / post_moe_hop calls replays with fresh
-  inputs;
-* random per-rank device delays and fault-injected store order;
-* a missing rank times out with NaN, error bits and the post_moe_hop health check,
-  and a collective reset restores exact results.
-"""
-
-import os
-import sys
-import time
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "post_moe_hop"))
-sys.path.insert(0, os.path.join(_HERE, "..", "deferred_norm"))
-
-import post_moe_hop_mpi_utils as M  # noqa: E402
-from test_deferred_norm_routed_mpi import _fill, _source  # noqa: E402
-
-CONFIGS = os.environ.get(
-    "GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_HOP_CFGS",
-    "carveout=100,poll_prefetch=wait3,ffi=0;"
-    "carveout=100,poll_prefetch=cur,ffi=0;carveout=100,poll_prefetch=pre3,ffi=0;"
-    "carveout=100,poll_prefetch=wait3,clear_hint=evict_last,ffi=0;"
-    "carveout=-1,poll_prefetch=none;carveout=100,poll_prefetch=none,ffi=0",
-).split(";")
-MODES = ("fp32", "identity", "row")
-_HOPS = {}
-
-
-def _env():
-    env = M.mpi_env()
-    if env.world != 8:
-        pytest.skip("the hop is specialized to TP8")
-    return env
-
-
-def _hops(mode, cfg, timeout_s=None, fault_delay_ns=0):
-    """post_moe_hop (routed input ``mode``), deferred_norm hop and deferred_norm_evict_hints hop
-    ``cfg`` on one workspace."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormHop,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
-        EvictHintsDeferredHop,
-        EvictHintsHopConfig,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
-        parse_cfg,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    env = _env()
-    key = (mode, timeout_s, fault_delay_ns)
-    if key not in _HOPS:
-        kw = {}
-        if timeout_s is not None:
-            kw["timeout_s"] = timeout_s
-        if fault_delay_ns:
-            kw["fault_delay_ns"] = fault_delay_ns
-        post_moe_hop = PostMoeHop(
-            env.mapping, M.GLM_HIDDEN, M.GLM_TOP_K, require_multi_node=False, **kw
-        )
-        assert M.all_ranks(post_moe_hop.engaged and post_moe_hop.enable_routed_input(mode)), (
-            post_moe_hop.reason
-        )
-        deferred_norm = DeferredNormHop(post_moe_hop, trigger_early=True)
-        assert M.all_ranks(deferred_norm.engaged and deferred_norm.enable_routed_input())
-        _HOPS[key] = {"post_moe_hop": post_moe_hop, "deferred_norm": deferred_norm}
-    entry = _HOPS[key]
-    if cfg not in entry:
-        hop_config = EvictHintsDeferredHop(
-            entry["post_moe_hop"], trigger_early=True, config=parse_cfg(cfg, EvictHintsHopConfig)
-        )
-        assert M.all_ranks(hop_config.engaged and hop_config.enable_routed_input()), (
-            hop_config.reason,
-            hop_config.routed_reason,
-        )
-        entry[cfg] = hop_config
-    return entry["post_moe_hop"], entry["deferred_norm"], entry[cfg]
-
-
-def _bits(t):
-    return t.view(torch.int16) if t.element_size() == 2 else t.view(torch.int32)
-
-
-def _outs(packet, res):
-    """Residual, the R written rows of ``a_frag``, ``sumsq`` and the materialization."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        materialize,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.packet import (  # noqa: E501
-        rows_from_frag,
-    )
-
-    h = materialize(packet)
-    torch.cuda.synchronize()
-    a = rows_from_frag(packet.a_frag, res.shape[0], res.shape[1])
-    return [res.clone(), a.clone(), packet.sumsq.clone(), h.clone()]
-
-
-def _same(a, b, tag):
-    same = [torch.equal(_bits(x), _bits(y)) for x, y in zip(a, b)]
-    if not all(same):
-        print(
-            f"GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_HOP_FAIL {tag} rank={M.mpi_env().rank} "
-            f"residual/a_frag/sumsq/h={same}",
-            flush=True,
-        )
-    return all(same)
-
-
-def _acc_clear(post_moe_hop, mode):
-    if mode != "fp32":
-        return True
-    torch.cuda.synchronize()
-    return bool((post_moe_hop.routed_acc_buffer(post_moe_hop.max_rows) == 0).all())
-
-
-def _routed(post_moe_hop, hop, mode, src, rows, shared, residual, gamma):
-    routed = _fill(post_moe_hop, mode, src, rows)
-    return hop.hop_routed(routed, shared, residual, gamma, M.GLM_EPS)
-
-
-@pytest.mark.parametrize("with_shared", [True, False])
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-@pytest.mark.parametrize("mode", MODES)
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_routed_matches_deferred_norm(cfg, mode, rows, with_shared):
-    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, rows, seed=50 + rows, with_shared=with_shared)
-        ref = _outs(*_routed(post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma))
-        got = _outs(*_routed(post_moe_hop, hop_config, mode, src, rows, shared, residual, gamma))
-        ok = _same(got, ref, f"{cfg} {mode} R={rows}") and _acc_clear(post_moe_hop, mode)
-        assert M.all_ranks(ok)
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_raw_tuple_matches_deferred_norm(cfg, rows):
-    post_moe_hop, deferred_norm, hop_config = _hops("fp32", cfg)
-    with torch.inference_mode():
-        t = M.make_raw_tuple(rows, seed=300 + rows)
-        ref = _outs(*deferred_norm(*t, M.GLM_EPS))
-        got = _outs(*hop_config(*t, M.GLM_EPS))
-        assert M.all_ranks(_same(got, ref, f"{cfg} raw R={rows}"))
-
-
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_interleaved_with_deferred_norm_and_post_moe_hop(cfg):
-    """deferred_norm_evict_hints / deferred_norm / post_moe_hop calls (routed and raw) back to back on one workspace."""
-    mode = "fp32"
-    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
-    kinds = [
-        "routed_deferred_norm_evict_hints",
-        "routed_deferred_norm",
-        "routed_post_moe_hop",
-        "deferred_norm_evict_hints",
-        "post_moe_hop",
-        "routed_deferred_norm_evict_hints",
-        "routed_deferred_norm_evict_hints",
-        "deferred_norm",
-        "routed_post_moe_hop",
-        "routed_deferred_norm_evict_hints",
-        "deferred_norm_evict_hints",
-        "routed_deferred_norm",
-    ]
-    rows_seq = [6, 1, 8, 3, 6, 2, 6, 8, 5, 6, 1, 6]
-    with torch.inference_mode():
-        inputs = [
-            (
-                _source(mode, r, seed=900 + i)
-                if k.startswith("routed_")
-                else M.make_raw_tuple(r, seed=900 + i)
-            )
-            for i, (k, r) in enumerate(zip(kinds, rows_seq))
-        ]
-        outs = []
-        for kind, r, t in zip(kinds, rows_seq, inputs):
-            if kind.startswith("routed_"):
-                hop = {
-                    "routed_deferred_norm_evict_hints": hop_config,
-                    "routed_deferred_norm": deferred_norm,
-                }.get(kind)
-                if hop is None:
-                    outs.append(
-                        post_moe_hop.hop_routed(
-                            _fill(post_moe_hop, mode, t[0], r), *t[1:], M.GLM_EPS
-                        )
-                    )
-                else:
-                    outs.append(_routed(post_moe_hop, hop, mode, t[0], r, *t[1:]))
-            else:
-                hop = {
-                    "deferred_norm_evict_hints": hop_config,
-                    "deferred_norm": deferred_norm,
-                    "post_moe_hop": post_moe_hop,
-                }[kind]
-                outs.append(hop(*t, M.GLM_EPS))
-        torch.cuda.synchronize()
-        snap = [
-            [x.clone() for x in o] if kind.endswith("post_moe_hop") else _outs(*o)
-            for kind, o in zip(kinds, outs)
-        ]
-        failures = []
-        for i, (kind, r, t, got) in enumerate(zip(kinds, rows_seq, inputs, snap)):
-            if kind.endswith("post_moe_hop"):
-                if kind.startswith("routed_"):
-                    ref = post_moe_hop.hop_routed(
-                        _fill(post_moe_hop, mode, t[0], r), *t[1:], M.GLM_EPS
-                    )
-                else:
-                    ref = post_moe_hop(*t, M.GLM_EPS)
-                torch.cuda.synchronize()
-                ok = torch.equal(got[0], ref[0]) and torch.equal(got[1], ref[1])
-            else:
-                if kind.startswith("routed_"):
-                    ref = _outs(*_routed(post_moe_hop, deferred_norm, mode, t[0], r, *t[1:]))
-                else:
-                    ref = _outs(*deferred_norm(*t, M.GLM_EPS))
-                ok = _same(got, ref, f"interleaved{i}-{kind}")
-            if not ok:
-                failures.append((i, kind))
-        assert M.all_ranks(not failures and _acc_clear(post_moe_hop, mode)), failures
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_cuda_graph_replay(cfg):
-    """Fill + deferred_norm_evict_hints routed, fill + deferred_norm routed, a
-    deferred_norm_evict_hints raw call in one graph, replayed."""
-    mode = "fp32"
-    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
-    rows = 6
-    with torch.inference_mode():
-        bufs = [list(_source(mode, rows, seed=1200 + i)) for i in range(2)]
-        raw = [x.clone() if x is not None else None for x in M.make_raw_tuple(rows, seed=1250)]
-
-        def body():
-            a = _routed(post_moe_hop, hop_config, mode, bufs[0][0], rows, *bufs[0][1:])
-            b = _routed(post_moe_hop, deferred_norm, mode, bufs[1][0], rows, *bufs[1][1:])
-            c = hop_config(*raw, M.GLM_EPS)
-            return a, b, c
-
-        body()
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            ga, gb, gc = body()
-        for trial, replays in enumerate((1, 1, 20)):
-            fresh = [_source(mode, rows, seed=1300 + 10 * trial + i) for i in range(2)]
-            fresh_raw = M.make_raw_tuple(rows, seed=1350 + trial)
-            for dst, src in zip(bufs, fresh):
-                for d, s in zip(dst, src):
-                    if d is not None:
-                        d.copy_(s)
-            for d, s in zip(raw, fresh_raw):
-                if d is not None:
-                    d.copy_(s)
-            for _ in range(replays):
-                M.mpi_env().comm.barrier()
-                graph.replay()
-            torch.cuda.synchronize()
-            got = [_outs(*ga), _outs(*gb), _outs(*gc)]
-            refs = [
-                _outs(
-                    *_routed(post_moe_hop, deferred_norm, mode, fresh[0][0], rows, *fresh[0][1:])
-                ),
-                _outs(
-                    *_routed(post_moe_hop, deferred_norm, mode, fresh[1][0], rows, *fresh[1][1:])
-                ),
-                _outs(*deferred_norm(*fresh_raw, M.GLM_EPS)),
-            ]
-            ok = all(_same(g, r, f"graph{trial}-{j}") for j, (g, r) in enumerate(zip(got, refs)))
-            assert M.all_ranks(ok and _acc_clear(post_moe_hop, mode)), f"trial {trial}"
-
-
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_delayed_rank_stress(cfg):
-    """Random per-rank device delays inside back-to-back deferred_norm_evict_hints routed batches."""
-    mode = "fp32"
-    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg)
-    env = M.mpi_env()
-    gen = torch.Generator().manual_seed(4243)
-    with torch.inference_mode():
-        for batch in range(3):
-            inputs, outs = [], []
-            for i in range(12):
-                rows = int(torch.randint(1, 9, (1,), generator=gen))
-                t = _source(mode, rows, seed=5000 + 100 * batch + i)
-                local = torch.Generator().manual_seed(77 * batch + 13 * i + env.rank)
-                if float(torch.rand(1, generator=local)) < 0.3:
-                    torch.cuda._sleep(int(float(torch.rand(1, generator=local)) * 2_000_000))
-                inputs.append((rows, t))
-                outs.append(_routed(post_moe_hop, hop_config, mode, t[0], rows, *t[1:]))
-            torch.cuda.synchronize()
-            snap = [_outs(*o) for o in outs]
-            ok = True
-            for (rows, t), got in zip(inputs, snap):
-                ok &= _same(
-                    got,
-                    _outs(*_routed(post_moe_hop, deferred_norm, mode, t[0], rows, *t[1:])),
-                    f"stress{batch}",
-                )
-            assert M.all_ranks(ok), f"batch {batch}"
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_fault_injected_store_order(cfg):
-    """Scrambled store order (post_moe_hop's fault injection) in the deferred_norm_evict_hints kernel."""
-    mode = "fp32"
-    post_moe_hop, deferred_norm, hop_config = _hops(mode, cfg, fault_delay_ns=4000)
-    with torch.inference_mode():
-        inputs = [(r, _source(mode, r, seed=7000 + i)) for i, r in enumerate([6, 1, 8, 6, 3, 6])]
-        outs = [_routed(post_moe_hop, hop_config, mode, t[0], r, *t[1:]) for r, t in inputs]
-        torch.cuda.synchronize()
-        snap = [_outs(*o) for o in outs]
-        ok = True
-        for (r, t), got in zip(inputs, snap):
-            ok &= _same(
-                got, _outs(*_routed(post_moe_hop, deferred_norm, mode, t[0], r, *t[1:])), "fault"
-            )
-        assert M.all_ranks(ok)
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-
-
-_MISSING_TIMEOUT_S = 0.2
-
-
-@pytest.mark.parametrize("missing", [0, 5])
-@pytest.mark.parametrize("cfg", CONFIGS)
-def test_missing_rank_times_out_with_nan_and_reset(cfg, missing):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-
-    mode = "fp32"
-    post_moe_hop, deferred_norm, _ = _hops(mode, cfg)
-    imp_post_moe_hop, _, imp_evict_hints = _hops(mode, cfg, timeout_s=_MISSING_TIMEOUT_S)
-    env = M.mpi_env()
-    assert M.all_ranks(imp_post_moe_hop.workspace is post_moe_hop.workspace)
-    rows = 6
-    with torch.inference_mode():
-        src, shared, residual, gamma = _source(mode, rows, seed=600 + missing)
-        env.comm.barrier()
-        start = time.monotonic()
-        snap = None
-        if env.rank != missing:
-            snap = _outs(
-                *_routed(
-                    imp_post_moe_hop, imp_evict_hints, mode, src, rows, shared, residual, gamma
-                )
-            )
-        torch.cuda.synchronize()
-        elapsed = time.monotonic() - start
-        counters = post_moe_hop.workspace.counters()
-        bits, polls = counters["error_bits"], counters["timeout_polls"]
-        if env.rank == missing:
-            reported = bits == 0 and polls == 0
-        else:
-            nan = all(bool(torch.isnan(t.float()).all()) for t in (snap[0], snap[2], snap[3]))
-            reported = (
-                bool(bits & 1)
-                and polls > 0
-                and _MISSING_TIMEOUT_S <= elapsed < 30.0
-                and nan
-                and _acc_clear(imp_post_moe_hop, mode)
-            )
-        stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
-        if env.rank == 0:
-            print(
-                f"GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_MISSING cfg={cfg} rank {missing} skipped: {stats}",
-                flush=True,
-            )
-        assert M.all_ranks(reported), stats
-        try:
-            check_health(f"deferred_norm_evict_hints-missing-rank{missing}", collective=True)
-            raised = False
-        except PostMoeHopError:
-            raised = True
-        assert M.all_ranks(raised)
-        assert M.all_ranks(post_moe_hop.reset())
-        assert M.all_ranks(post_moe_hop.error_bits() == 0)
-        src, shared, residual, gamma = _source(mode, rows, seed=650 + missing)
-        _, _, hop_config = _hops(mode, cfg)
-        got = _outs(*_routed(post_moe_hop, hop_config, mode, src, rows, shared, residual, gamma))
-        ref = _outs(*_routed(post_moe_hop, deferred_norm, mode, src, rows, shared, residual, gamma))
-        assert M.all_ranks(_same(got, ref, f"after-reset{missing}"))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py
deleted file mode 100644
index 7ee030e1624f..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/deferred_norm_evict_hints/test_evict_hints_projection.py
+++ /dev/null
@@ -1,156 +0,0 @@
-# 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.
-"""Single-GPU tests of the deferred_norm_evict_hints projection against F's deferred_norm projection (real weights).
-
-For every deferred_norm_evict_hints projection configuration (``CONFIGS``) and every layer
-the projection serves (``GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_LAYERS``: a shared-indexer
-target layer; the target's deferred-norm scope is ``shared``, so full-indexer layers keep their
-native projection), the fused-A output
-is bit-identical to ``DeferredNormProjection``'s on the same deferred packet, for R in
-1..8 and residual scales from 1e-3 to 30, in eager mode and in a CUDA graph
-with PDL. ``precompile_served`` covers the consumer's launch configuration.
-
-Fixture label: weights real (checkpoint), activations synthetic.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
-sys.path.insert(0, os.path.join(_HERE, "..", "deferred_norm"))
-
-import absorbed_query_fixture  # noqa: E402
-from test_deferred_norm_projection import _module, make_packet, make_residual  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
-    reason="needs a GPU and the GLM-5.2 checkpoint",
-)
-
-LAYERS = [
-    int(x)
-    for x in os.environ.get("GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_LAYERS", "7").split(",")
-]
-CONFIGS = os.environ.get(
-    "GLM_MOE_DSA_DEFERRED_NORM_EVICT_HINTS_TEST_PROJ_CFGS",
-    "maxrreg=112;maxrreg=0,maxntid=576;maxrreg=112,copy_hint=evict_first;"
-    "maxrreg=112,split_chunk=2,split_delay_ns=2500;maxrreg=0;copy_warps=1;prod_warp=1;"
-    "prod_warp=1,split_chunk=1,split_delay_ns=3000;"
-    "prod_warp=1,split_chunk=2,split_delay_ns=2500,copy_hint=evict_first;"
-    "prod_warp=1,defer_tail=8,defer_wait=1;prod_warp=1,split_chunk=1,defer_tail=3,defer_wait=1;"
-    "prod_warp=1,defer_tail=5,split_delay_ns=2000;",
-).split(";")
-
-
-def _pair(layer, cfg_text):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm import (  # noqa: E501
-        DeferredNormProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints import (  # noqa: E501
-        EvictHintsProjConfig,
-        EvictHintsProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
-        parse_cfg,
-    )
-
-    attn, gamma = _module(layer)
-    ok, reason, kw = projection_eligibility(attn)
-    assert ok, reason
-    deferred_norm = DeferredNormProjection(attn, kw_weight=kw)
-    hop_config = EvictHintsProjection(attn, parse_cfg(cfg_text, EvictHintsProjConfig), kw_weight=kw)
-    assert hop_config.supported() is None
-    return attn, gamma, deferred_norm, hop_config
-
-
-def _bits(t):
-    return None if t is None else t.view(torch.int16)
-
-
-def _equal(a, b):
-    return all((x is None and y is None) or torch.equal(_bits(x), _bits(y)) for x, y in zip(a, b))
-
-
-@pytest.mark.parametrize("cfg", CONFIGS)
-@pytest.mark.parametrize("layer", LAYERS)
-def test_served_configs_cover_the_consumer(layer, cfg):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.proj_kernel import (  # noqa: E501
-        is_compiled,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.deferred_norm_evict_hints.projection import (  # noqa: E501
-        EvictHintsProjConfig,
-        parse_cfg,
-        precompile_served,
-        served_kernel_configs,
-    )
-
-    config = parse_cfg(cfg, EvictHintsProjConfig)
-    device = torch.cuda.current_device()
-    assert precompile_served(device, config) is None
-    _, _, _, hop_config = _pair(layer, cfg)
-    assert hop_config.served_config() in served_kernel_configs(device, config)
-    assert is_compiled(**hop_config.served_config())
-
-
-@pytest.mark.parametrize("scale", [1.0, 1e-3, 30.0])
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-@pytest.mark.parametrize("cfg", CONFIGS)
-@pytest.mark.parametrize("layer", LAYERS)
-def test_bit_identical_to_deferred_norm(layer, cfg, rows, scale):
-    _, gamma, deferred_norm, hop_config = _pair(layer, cfg)
-    with torch.inference_mode():
-        x = make_residual(rows, seed=100 + rows + layer, scale=scale)
-        packet = make_packet(x, gamma)
-        ref = deferred_norm.project(packet)
-        got = hop_config.project(packet)
-        torch.cuda.synchronize()
-        assert _equal(got, ref)
-
-
-@pytest.mark.parametrize("cfg", CONFIGS)
-@pytest.mark.parametrize("layer", LAYERS)
-def test_graph_with_pdl_bit_identical(layer, cfg):
-    """deferred_norm_evict_hints then deferred_norm projection of 4 packets in one CUDA graph, replayed with
-    fresh packets (PDL chain between the projections)."""
-    _, gamma, deferred_norm, hop_config = _pair(layer, cfg)
-    with torch.inference_mode():
-        packets = [make_packet(make_residual(6, seed=500 + i), gamma) for i in range(4)]
-
-        def body():
-            return [(hop_config.project(p), deferred_norm.project(p)) for p in packets]
-
-        body()
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            outs = body()
-        for trial in range(3):
-            for i, p in enumerate(packets):
-                fresh = make_packet(make_residual(6, seed=600 + 10 * trial + i), gamma)
-                p.a_frag.copy_(fresh.a_frag)
-                p.sumsq.copy_(fresh.sumsq)
-            graph.replay()
-            torch.cuda.synchronize()
-            for (got, ref), p in zip(outs, packets):
-                assert _equal(got, ref)
-                assert _equal(got, deferred_norm.project(p))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py
deleted file mode 100644
index b8c481449425..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/dense_mlp_utils.py
+++ /dev/null
@@ -1,190 +0,0 @@
-# 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.
-"""Shared helpers for the dense_mlp tests and microbenchmarks.
-
-* ``ensure_op``: make ``trtllm::glm_moe_dsa_dense_mlp*`` available. An image built from
-  a commit that contains the dense_mlp thop sources already registers the ops; for
-  in-job development on an older image the same sources are JIT-compiled.
-* Native oracle: the exact native BF16 ``GatedMLP`` compute (``F.linear`` ->
-  TRT-LLM ``swiglu`` -> ``F.linear``) and an FP64 reference with the native
-  BF16 intermediate rounding.
-* Real data: GLM-5.2 dense-layer weight shards as the GEN loader lays them out
-  (TP4 sub-group, fused [gate; up]) and real post-RMSNorm hidden rows.
-"""
-
-import json
-import os
-from pathlib import Path
-from typing import Dict, Optional, Tuple
-
-import torch
-import torch.nn.functional as F
-from utils.llm_data import llm_models_root
-
-REPO_ROOT = Path(__file__).resolve().parents[9]
-DEFAULT_CKPT = os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4")
-DEFAULT_REAL_HIDDEN = os.environ.get("GLM_MOE_DSA_REAL_MOE_INPUTS_DIR", "")
-GLM52_HIDDEN = 6144
-GLM52_INTER = 12288
-GLM52_DENSE_TP = 4
-
-
-def ensure_op(build_dir: Optional[str] = None, verbose: bool = False) -> str:
-    """Return "builtin" or "jit" after making the dense_mlp ops available."""
-    import tensorrt_llm  # noqa: F401  (loads th_common and its op fragments)
-
-    if hasattr(torch.ops.trtllm, "glm_moe_dsa_dense_mlp"):
-        return "builtin"
-    from torch.utils.cpp_extension import load
-
-    libs_dir = Path(tensorrt_llm.__file__).resolve().parent / "libs"
-    build_dir = build_dir or os.environ.get("GLM_MOE_DSA_DENSE_MLP_JIT_BUILD_DIR")
-    if build_dir:
-        Path(build_dir).mkdir(parents=True, exist_ok=True)
-    load(
-        name="dense_mlp_jit",
-        sources=[
-            str(REPO_ROOT / "cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu"),
-            str(REPO_ROOT / "cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp"),
-        ],
-        extra_include_paths=[
-            str(REPO_ROOT / "cpp/include"),
-            str(REPO_ROOT / "cpp"),
-        ],
-        extra_cflags=["-O3", "-DENABLE_BF16", "-DENABLE_FP8"],
-        extra_cuda_cflags=[
-            "-O3",
-            "-DENABLE_BF16",
-            "-DENABLE_FP8",
-            "-gencode=arch=compute_103a,code=sm_103a",
-            "-lineinfo",
-        ],
-        extra_ldflags=[f"-L{libs_dir}", "-ltensorrt_llm", f"-Wl,-rpath,{libs_dir}"],
-        build_directory=build_dir,
-        is_python_module=False,
-        verbose=verbose,
-    )
-    assert hasattr(torch.ops.trtllm, "glm_moe_dsa_dense_mlp")
-    return "jit"
-
-
-def native_dense_mlp(
-    x: torch.Tensor, gate_up_w: torch.Tensor, down_w: torch.Tensor
-) -> torch.Tensor:
-    """Native BF16 GatedMLP compute without the output all-reduce."""
-    from tensorrt_llm._torch.modules.swiglu import swiglu
-
-    h1 = F.linear(x, gate_up_w)
-    h2 = swiglu(h1)
-    return F.linear(h2, down_w)
-
-
-def native_intermediate(x: torch.Tensor, gate_up_w: torch.Tensor) -> torch.Tensor:
-    from tensorrt_llm._torch.modules.swiglu import swiglu
-
-    return swiglu(F.linear(x, gate_up_w))
-
-
-def oracle_dense_mlp(
-    x: torch.Tensor, gate_up_w: torch.Tensor, down_w: torch.Tensor
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """FP64 sums with the native BF16 rounding points.
-
-    Returns ``(h, out)`` as FP64 tensors whose values are BF16-representable
-    where the native path rounds (gate, up, h); ``out`` is unrounded.
-    """
-    inter = down_w.shape[1]
-    gu = x.double() @ gate_up_w.double().t()
-    gate = gu[:, :inter].to(torch.bfloat16).double()
-    up = gu[:, inter:].to(torch.bfloat16).double()
-    h = (gate * torch.sigmoid(gate) * up).to(torch.bfloat16).double()
-    out = h @ down_w.double().t()
-    return h, out
-
-
-def error_stats(value: torch.Tensor, ref: torch.Tensor) -> Dict[str, float]:
-    diff = value.double() - ref.double()
-    scale = ref.double().pow(2).mean().sqrt().item()
-    return {
-        "rms": diff.pow(2).mean().sqrt().item(),
-        "max": diff.abs().max().item(),
-        "nrms": diff.pow(2).mean().sqrt().item() / max(scale, 1e-30),
-        "ref_rms": scale,
-    }
-
-
-def _safetensors_index(ckpt: str) -> Dict[str, str]:
-    with open(os.path.join(ckpt, "model.safetensors.index.json")) as f:
-        return json.load(f)["weight_map"]
-
-
-def load_dense_shard(
-    layer: int, shard: int, ckpt: str = DEFAULT_CKPT, tp: int = GLM52_DENSE_TP, device: str = "cuda"
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Return GEN-layout ``(gate_up [2I/tp, H], down [H, I/tp])`` of one rank.
-
-    Matches the native loader: gate rows ``[I/tp * shard, ...)`` then the same
-    up rows; down input columns of that shard.
-    """
-    from safetensors import safe_open
-
-    index = _safetensors_index(ckpt)
-    prefix = f"model.layers.{layer}.mlp."
-    tensors = {}
-    for name in ("gate_proj", "up_proj", "down_proj"):
-        key = prefix + name + ".weight"
-        with safe_open(os.path.join(ckpt, index[key]), framework="pt") as f:
-            sl = f.get_slice(key)
-            rows, cols = sl.get_shape()
-            if name == "down_proj":
-                width = cols // tp
-                tensors[name] = sl[:, width * shard : width * (shard + 1)]
-            else:
-                width = rows // tp
-                tensors[name] = sl[width * shard : width * (shard + 1), :]
-    gate_up = torch.cat([tensors["gate_proj"], tensors["up_proj"]], dim=0)
-    down = tensors["down_proj"].contiguous()
-    assert gate_up.dtype == torch.bfloat16 and down.dtype == torch.bfloat16
-    return gate_up.to(device).contiguous(), down.to(device).contiguous()
-
-
-def load_real_hidden(
-    fixture_dir: str = DEFAULT_REAL_HIDDEN, layer: int = 10, tokens: int = 2311
-) -> Optional[torch.Tensor]:
-    """Real post-RMSNorm hidden rows [tokens, 6144] (BF16) or None if absent.
-
-    The fixture captured the routed-MoE input (``post_attention_layernorm``
-    output) of the native model; it serves as a realistic activation for the
-    dense MLP, which consumes the same kind of tensor in layers 0-2.
-    """
-    path = os.path.join(fixture_dir, f"rows{tokens}_layer{layer}.pt")
-    if not os.path.exists(path):
-        return None
-    obj = torch.load(path, map_location="cpu", weights_only=False)
-    if isinstance(obj, dict):
-        for key in ("hidden_states", "hidden", "x", "input"):
-            if key in obj and isinstance(obj[key], torch.Tensor):
-                return obj[key].to(torch.bfloat16)
-        for value in obj.values():
-            if (
-                isinstance(value, torch.Tensor)
-                and value.dim() == 2
-                and value.shape[-1] == GLM52_HIDDEN
-            ):
-                return value.to(torch.bfloat16)
-        return None
-    if isinstance(obj, torch.Tensor):
-        return obj.to(torch.bfloat16)
-    return None
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
deleted file mode 100644
index 4bf176903e38..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/pdl_delay.cu
+++ /dev/null
@@ -1,74 +0,0 @@
-/*
- * 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.
- */
-
-// Test-only delayed producer for programmatic dependent launch: releases its dependents immediately,
-// spins for a given number of nanoseconds, then copies src into dst. A PDL consumer that reads dst
-// before griddepcontrol.wait would observe stale data.
-
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/extension.h>
-
-#include <cstdint>
-
-namespace
-{
-
-__global__ void delayedCopyKernel(
-    uint4 const* __restrict__ src, uint4* __restrict__ dst, int64_t numVec, int64_t delayNs)
-{
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    asm volatile("griddepcontrol.wait;" ::: "memory");
-    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
-#endif
-    uint64_t start;
-    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(start));
-    uint64_t now = start;
-    while (now - start < static_cast<uint64_t>(delayNs))
-    {
-        asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now));
-    }
-    for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < numVec;
-         i += static_cast<int64_t>(gridDim.x) * blockDim.x)
-    {
-        dst[i] = src[i];
-    }
-}
-
-} // namespace
-
-void delayedCopy(torch::Tensor src, torch::Tensor dst, int64_t delayNs)
-{
-    TORCH_CHECK(src.is_contiguous() && dst.is_contiguous() && src.nbytes() == dst.nbytes() && src.nbytes() % 16 == 0);
-    auto const stream = at::cuda::getCurrentCUDAStream(src.get_device());
-    cudaLaunchConfig_t config{};
-    config.gridDim = dim3(4);
-    config.blockDim = dim3(128);
-    config.stream = stream;
-    cudaLaunchAttribute attrs[1];
-    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-    attrs[0].val.programmaticStreamSerializationAllowed = 1;
-    config.attrs = attrs;
-    config.numAttrs = 1;
-    TORCH_CHECK(cudaLaunchKernelEx(&config, delayedCopyKernel, reinterpret_cast<uint4 const*>(src.data_ptr()),
-                    reinterpret_cast<uint4*>(dst.data_ptr()), static_cast<int64_t>(src.nbytes() / 16), delayNs)
-        == cudaSuccess);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
-{
-    m.def("delayed_copy", &delayedCopy);
-}
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py
deleted file mode 100644
index 60af1ceb683d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_adapter.py
+++ /dev/null
@@ -1,171 +0,0 @@
-# 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.
-"""Tests of the dense_mlp call site in the target's decoder layer.
-
-``GlmDecoderLayer.forward_mlp`` is exercised on a decode step on a minimal layer object
-(dense GatedMLP, single rank, no fusion), with and without ``dense_mlp`` selected.
-"""
-
-import os
-import sys
-import types
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import plan_testlib  # noqa: E402
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
-
-HIDDEN = 2048
-INTER = 4096
-
-
-@pytest.fixture(scope="module", autouse=True)
-def dense_mlp_op():
-    import dense_mlp_utils as wu
-
-    wu.ensure_op()
-
-
-@pytest.fixture(autouse=True)
-def decode_step():
-    """The owned call site runs on decode steps."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    with dispatch.step(dispatch.Phase.DECODE):
-        yield
-
-
-def _select(monkeypatch, value):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    plan_testlib.select(dispatch, value, monkeypatch)
-
-
-@pytest.fixture
-def spy(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    dispatch.get_impl("dense_mlp", "dense_mlp_forward")
-    real = dispatch._REGISTRY["dense_mlp"]["dense_mlp_forward"]
-
-    class Calls(list):
-        pass
-
-    calls = Calls()
-
-    def wrapped(*args, **kwargs):
-        result = real(*args, **kwargs)
-        calls.append(result is not None)
-        if result is not None:
-            wrapped.last = (args[0], args[1].clone(), args[2], result.clone())
-        return result
-
-    wrapped.last = None
-
-    monkeypatch.setitem(dispatch._REGISTRY["dense_mlp"], "dense_mlp_forward", wrapped)
-    calls.wrapped = wrapped
-    return calls
-
-
-def _layer():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmDecoderLayer,
-    )
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.models.modeling_utils import EagerFusionConfig
-    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
-    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-
-    torch.manual_seed(0)
-    mlp = GatedMLP(
-        hidden_size=HIDDEN,
-        intermediate_size=INTER,
-        bias=False,
-        dtype=torch.bfloat16,
-        config=ModelConfig(),
-        reduce_output=False,
-    ).cuda()
-    with torch.no_grad():
-        mlp.gate_up_proj.weight.normal_(0, 0.02)
-        mlp.down_proj.weight.normal_(0, 0.02)
-    layer = types.SimpleNamespace(
-        mlp=mlp,
-        fusion_config=EagerFusionConfig(),
-        mlp_tp_size=1,
-        allreduce=None,
-        layer_idx=0,
-        post_attention_layernorm=RMSNorm(hidden_size=HIDDEN, eps=1e-5, dtype=torch.bfloat16).cuda(),
-        next_layer_layernorm=RMSNorm(hidden_size=HIDDEN, eps=1e-5, dtype=torch.bfloat16).cuda(),
-    )
-    layer._apply_next_layer_layernorm = types.MethodType(
-        GlmDecoderLayer._apply_next_layer_layernorm, layer
-    )
-    return layer, GlmDecoderLayer.forward_mlp
-
-
-def _run(layer, forward_mlp, rows, seed=1):
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    hidden = torch.randn(rows, HIDDEN, device="cuda", generator=g).to(torch.bfloat16)
-    residual = torch.randn(rows, HIDDEN, device="cuda", generator=g).to(torch.bfloat16)
-    return forward_mlp(layer, hidden, residual)
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_dense_mlp_engages_and_matches_native(monkeypatch, spy, rows):
-    layer, forward_mlp = _layer()
-    _select(monkeypatch, "")
-    _run(layer, forward_mlp, rows)
-    assert spy == []
-    _select(monkeypatch, "dense_mlp")
-    hidden, residual = _run(layer, forward_mlp, rows)
-    assert spy == [True]
-    # The MLP output handed back to forward_mlp matches the native module on
-    # the same input within the BF16 GEMM tolerance.
-    mlp, x_in, params, dense_mlp_out = spy.wrapped.last
-    torch.testing.assert_close(
-        dense_mlp_out, mlp(x_in, final_all_reduce_params=params), rtol=1e-2, atol=5e-3
-    )
-    # Everything after the MLP stays native: the native forward_mlp fed with
-    # the same MLP output reproduces the dense_mlp layer outputs bit for bit.
-    _select(monkeypatch, "")
-    native_mlp = layer.mlp
-    layer.mlp = lambda x, final_all_reduce_params=None: dense_mlp_out.clone()
-    ref_hidden, ref_residual = _run(layer, forward_mlp, rows)
-    layer.mlp = native_mlp
-    assert torch.equal(hidden, ref_hidden)
-    assert torch.equal(residual, ref_residual)
-
-
-def test_ineligible_rows_take_native_path(monkeypatch, spy):
-    layer, forward_mlp = _layer()
-    _select(monkeypatch, "")
-    nat_hidden, nat_residual = _run(layer, forward_mlp, 12)
-    _select(monkeypatch, "dense_mlp")
-    hidden, residual = _run(layer, forward_mlp, 12)
-    assert spy == [False]
-    assert torch.equal(hidden, nat_hidden)
-    assert torch.equal(residual, nat_residual)
-
-
-def test_other_features_do_not_engage_dense_mlp(monkeypatch, spy):
-    layer, forward_mlp = _layer()
-    _select(monkeypatch, "absorbed_query")
-    _run(layer, forward_mlp, 6)
-    assert spy == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py
deleted file mode 100644
index 6e266ee0e1c1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/dense_mlp/test_dense_mlp.py
+++ /dev/null
@@ -1,410 +0,0 @@
-# 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.
-"""GPU tests of package dense_mlp (small-M BF16 dense MLP).
-
-Oracle: the native BF16 GatedMLP compute (cuBLAS ``F.linear`` -> TRT-LLM
-``swiglu`` -> ``F.linear``) and an FP64 reference with native BF16 rounding.
-Tolerance against native: the existing TRT-LLM small-M BF16 GEMM tolerance
-(``tests/unittest/_torch/modules/test_low_m_gemm.py``: rtol 1e-2, atol 5e-3)
-for each GEMM stage. Error budget against the FP64 oracle, frozen before
-tuning: candidate normalized RMS error <= 1.25 x native + 1e-4.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-from utils.llm_data import llm_models_root
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
-
-NRMS_FACTOR = 1.25
-NRMS_SLACK = 1e-4
-
-
-@pytest.fixture(scope="module", autouse=True)
-def dense_mlp_op():
-    import dense_mlp_utils as wu
-
-    return wu.ensure_op()
-
-
-@pytest.fixture(autouse=True)
-def _seed_inputs():
-    """Seed the global generator so each case draws the same inputs in any order."""
-    torch.manual_seed(0)
-
-
-def _random_weights(hidden, inter, scale=0.02, seed=0):
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    gate_up = (scale * torch.randn(2 * inter, hidden, device="cuda", generator=g)).to(
-        torch.bfloat16
-    )
-    down = (scale * torch.randn(hidden, inter, device="cuda", generator=g)).to(torch.bfloat16)
-    return gate_up, down
-
-
-def assert_within_bf16_ulps(actual, expected, ulps, abs_floor_rel=1e-5):
-    """|actual - expected| <= ulps BF16 ulps of the larger magnitude + floor.
-
-    The absolute floor (``abs_floor_rel`` x RMS of ``expected``) covers values
-    near zero, where a different FP32 summation order of a cancelling sum
-    moves the result by far more than its own (tiny) ulp.
-    """
-    a = actual.float()
-    e = expected.float()
-    mag = torch.maximum(a.abs(), e.abs()).clamp_min(2.0**-126)
-    ulp = torch.exp2(torch.floor(torch.log2(mag)) - 7)
-    floor = abs_floor_rel * e.pow(2).mean().sqrt().item()
-    excess = ((a - e).abs() - ulps * ulp - floor).max().item()
-    assert excess <= 0, f"difference exceeds {ulps} BF16 ulps by {excess}"
-
-
-# SwiGLU of BF16-rounded gate and up, rounded to BF16: when another FP32
-# summation order flips the rounding of gate and of up by one ulp each, the
-# product moves by up to two ulps, plus one ulp from rounding h.
-SWIGLU_ULPS = 3
-
-
-def _check_against_native_and_oracle(x, gate_up, down, variants=(0, 0)):
-    """Per-stage native comparison plus the end-to-end FP64-oracle budget.
-
-    The down stage (one BF16 rounding) uses the existing small-M BF16 GEMM
-    tolerance against the native down GEMM on the candidate's own SwiGLU
-    output. The gate/up + SwiGLU stage has three rounding points (gate, up,
-    h) and is bounded by SWIGLU_ULPS against native. End to end, the frozen
-    FP64-oracle error budget applies.
-    """
-    import dense_mlp_utils as wu
-
-    h = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up, variants[0])
-    assert_within_bf16_ulps(h, wu.native_intermediate(x, gate_up), SWIGLU_ULPS)
-    out = torch.ops.trtllm.glm_moe_dsa_dense_mlp_down(h, down, variants[1])
-    torch.testing.assert_close(out, F.linear(h, down), rtol=1e-2, atol=5e-3)
-    cand = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
-    assert torch.equal(cand, out)
-    nat = wu.native_dense_mlp(x, gate_up, down)
-    _, ref = wu.oracle_dense_mlp(x, gate_up, down)
-    cand_err = wu.error_stats(cand, ref)["nrms"]
-    nat_err = wu.error_stats(nat, ref)["nrms"]
-    assert cand_err <= NRMS_FACTOR * nat_err + NRMS_SLACK, (cand_err, nat_err)
-    return cand
-
-
-# Contraction multiple (warps x 32 x steps) of each variant in
-# denseMlp.cu.
-VARIANT_K_MULTIPLE = [4 * 32 * 8, 8 * 32 * 4, 8 * 32 * 4, 4 * 32 * 8, 8 * 32 * 8]
-
-# (gate/up variant, down variant) pairs exercised by the shape-generic tests.
-VARIANT_PAIRS = [(0, 0), (1, 0), (2, 1), (3, 3), (4, 3)]
-
-
-def _supported(variants, hidden, inter):
-    return (
-        hidden % VARIANT_K_MULTIPLE[variants[0]] == 0
-        and inter % VARIANT_K_MULTIPLE[variants[1]] == 0
-    )
-
-
-@pytest.mark.parametrize("num_rows", [1, 2, 3, 4, 5, 6, 7, 8])
-@pytest.mark.parametrize("hidden,inter", [(6144, 3072), (1024, 1024), (2048, 4096)])
-@pytest.mark.parametrize("variants", VARIANT_PAIRS)
-def test_matches_native_random(num_rows, hidden, inter, variants):
-    if not _supported(variants, hidden, inter):
-        pytest.skip("tiling does not divide the shape")
-    gate_up, down = _random_weights(hidden, inter)
-    x = torch.randn(num_rows, hidden, device="cuda", dtype=torch.bfloat16)
-    _check_against_native_and_oracle(x, gate_up, down, variants)
-
-
-@pytest.mark.parametrize("num_rows", [1, 6, 8])
-def test_intermediate_rounding_matches_native(num_rows):
-    """gate/up/h rounding follows native; h differs only by summation order."""
-    import dense_mlp_utils as wu
-
-    gate_up, down = _random_weights(6144, 3072, seed=3)
-    x = torch.randn(num_rows, 6144, device="cuda", dtype=torch.bfloat16)
-    h = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up)
-    nat_h = wu.native_intermediate(x, gate_up)
-    ref_h, _ = wu.oracle_dense_mlp(x, gate_up, down)
-    assert_within_bf16_ulps(h, nat_h, SWIGLU_ULPS)
-    # Values that native reproduces exactly from the FP64 oracle must mostly
-    # match exactly here too (only BF16 rounding-boundary cases may differ).
-    assert (h == nat_h).float().mean().item() > 0.97
-    assert wu.error_stats(h, ref_h)["nrms"] <= (
-        NRMS_FACTOR * wu.error_stats(nat_h, ref_h)["nrms"] + NRMS_SLACK
-    )
-
-
-@pytest.mark.parametrize("variant", range(len(VARIANT_K_MULTIPLE)))
-def test_all_variants(variant):
-    gate_up, down = _random_weights(6144, 3072, seed=5)
-    x = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
-    ref = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, 0, 0)
-    h = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up, variant)
-    h0 = torch.ops.trtllm.glm_moe_dsa_dense_mlp_gate_up_swiglu(x, gate_up, 0)
-    assert_within_bf16_ulps(h, h0, SWIGLU_ULPS)
-    if 3072 % VARIANT_K_MULTIPLE[variant] != 0:
-        with pytest.raises(RuntimeError):
-            torch.ops.trtllm.glm_moe_dsa_dense_mlp_down(h0, down, variant)
-        return
-    out = torch.ops.trtllm.glm_moe_dsa_dense_mlp_down(h0, down, variant)
-    torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-3)
-
-
-@pytest.mark.parametrize("variants", VARIANT_PAIRS)
-def test_deterministic(variants):
-    gate_up, down = _random_weights(6144, 3072, seed=7)
-    x = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
-    first = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
-    for _ in range(5):
-        assert torch.equal(
-            torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants), first
-        )
-
-
-@pytest.mark.parametrize("variants", VARIANT_PAIRS)
-def test_cuda_graph_replay_tracks_inputs(variants):
-    gate_up, down = _random_weights(6144, 3072, seed=11)
-    x_static = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_static, gate_up, down, *variants)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        out_static = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_static, gate_up, down, *variants)
-    for seed in range(4):
-        torch.manual_seed(100 + seed)
-        new_x = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
-        x_static.copy_(new_x)
-        graph.replay()
-        eager = torch.ops.trtllm.glm_moe_dsa_dense_mlp(new_x, gate_up, down, *variants)
-        torch.cuda.synchronize()
-        assert torch.equal(out_static, eager)
-
-
-def test_rejects_unsupported_inputs():
-    gate_up, down = _random_weights(6144, 3072)
-    with pytest.raises(RuntimeError):
-        torch.ops.trtllm.glm_moe_dsa_dense_mlp(
-            torch.randn(9, 6144, device="cuda", dtype=torch.bfloat16), gate_up, down
-        )
-    with pytest.raises(RuntimeError):
-        torch.ops.trtllm.glm_moe_dsa_dense_mlp(
-            torch.randn(6, 6144, device="cuda", dtype=torch.float16), gate_up, down
-        )
-    with pytest.raises(RuntimeError):
-        torch.ops.trtllm.glm_moe_dsa_dense_mlp(
-            torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16), gate_up, down, 99, 0
-        )
-
-
-def _make_gated_mlp(hidden, inter, tp_size=1, overridden_tp_size=None):
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
-    from tensorrt_llm.mapping import Mapping
-
-    mapping = Mapping(world_size=tp_size, rank=0, tp_size=tp_size, gpus_per_node=max(tp_size, 1))
-    mlp = GatedMLP(
-        hidden_size=hidden,
-        intermediate_size=inter,
-        bias=False,
-        dtype=torch.bfloat16,
-        config=ModelConfig(mapping=mapping),
-        overridden_tp_size=overridden_tp_size,
-        reduce_output=tp_size > 1,
-    ).cuda()
-    gate_up, down = _random_weights(hidden, inter // (overridden_tp_size or tp_size), seed=13)
-    with torch.no_grad():
-        mlp.gate_up_proj.weight.copy_(gate_up)
-        mlp.down_proj.weight.copy_(down)
-    return mlp
-
-
-@pytest.mark.parametrize("num_rows", [1, 6, 8])
-def test_forward_matches_gated_mlp_module(num_rows):
-    import dense_mlp_utils as wu
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
-        dense_mlp,
-        dense_mlp_forward,
-    )
-
-    mlp = _make_gated_mlp(2048, 4096)
-    x = torch.randn(num_rows, 2048, device="cuda", dtype=torch.bfloat16)
-    cand = dense_mlp_forward(mlp, x)
-    assert cand is not None
-    # The entry point runs exactly the op with the module's weights and the
-    # package's default tilings; the module's native forward is the
-    # reference the op is checked against.
-    gate_up, down = mlp.gate_up_proj.weight, mlp.down_proj.weight
-    variants = (dense_mlp.GATE_UP_VARIANT, dense_mlp.DOWN_VARIANT)
-    assert torch.equal(cand, _check_against_native_and_oracle(x, gate_up, down, variants))
-    _, ref = wu.oracle_dense_mlp(x, gate_up, down)
-    assert wu.error_stats(cand, ref)["nrms"] <= (
-        NRMS_FACTOR * wu.error_stats(mlp(x), ref)["nrms"] + NRMS_SLACK
-    )
-
-
-def test_eligibility_falls_back_before_device_work():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
-        check_eligibility,
-        dense_mlp_forward,
-    )
-    from tensorrt_llm._torch.distributed import AllReduceParams
-    from tensorrt_llm._torch.utils import Fp4QuantizedTensor
-
-    mlp = _make_gated_mlp(2048, 4096)
-    x = torch.randn(6, 2048, device="cuda", dtype=torch.bfloat16)
-    assert check_eligibility(mlp, x)[0]
-    cases = {
-        "rows": torch.randn(9, 2048, device="cuda", dtype=torch.bfloat16),
-        "dtype": x.float(),
-        "noncontig": torch.randn(2048, 6, device="cuda", dtype=torch.bfloat16).t(),
-        "3d": x.view(1, 6, 2048),
-        "fp4": Fp4QuantizedTensor(x, x),
-    }
-    for name, bad in cases.items():
-        eligible, reason = check_eligibility(mlp, bad)
-        assert not eligible, name
-        assert dense_mlp_forward(mlp, bad) is None, name
-    assert not check_eligibility(mlp, x, lora_params={"x": 1})[0]
-    mlp.swiglu_limit = 7.0
-    assert not check_eligibility(mlp, x)[0]
-    mlp.swiglu_limit = None
-    # Output all-reduce requested on a TP>1 down_proj -> native path. The
-    # mapping is swapped in place so no communicator is needed.
-    from tensorrt_llm.mapping import Mapping
-
-    mlp.down_proj.reduce_output = True
-    mlp.down_proj.mapping = Mapping(world_size=2, rank=0, tp_size=2, gpus_per_node=2)
-    assert not check_eligibility(mlp, x)[0]
-    assert not check_eligibility(mlp, x, AllReduceParams(enable_allreduce=True))[0]
-    assert check_eligibility(mlp, x, AllReduceParams(enable_allreduce=False))[0]
-    mlp.activation = F.gelu
-    assert not check_eligibility(mlp, x, AllReduceParams(enable_allreduce=False))[0]
-
-
-def test_attention_dp_takes_native_path():
-    """Attention DP runs the dense MLP unsharded: outside dense_mlp's TP domain."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
-        check_eligibility,
-        dense_mlp_forward,
-    )
-    from tensorrt_llm.mapping import Mapping
-
-    mlp = _make_gated_mlp(2048, 4096)
-    x = torch.randn(6, 2048, device="cuda", dtype=torch.bfloat16)
-    assert check_eligibility(mlp, x)[0]
-    mlp.mapping = Mapping(
-        world_size=1, rank=0, tp_size=1, gpus_per_node=1, enable_attention_dp=True
-    )
-    assert check_eligibility(mlp, x) == (False, "attention data parallelism")
-    assert dense_mlp_forward(mlp, x) is None
-
-
-def test_registered_with_dispatch():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.dense_mlp import (  # noqa: E501
-        dense_mlp_forward,
-    )
-
-    assert dispatch.get_impl("dense_mlp", "dense_mlp_forward") is dense_mlp_forward
-
-
-@pytest.mark.skipif(
-    not os.path.isdir(os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4")),
-    reason="GLM-5.2 checkpoint not mounted",
-)
-@pytest.mark.parametrize("layer", [0, 1, 2])
-def test_real_weights_real_hidden(layer):
-    import dense_mlp_utils as wu
-
-    gate_up, down = wu.load_dense_shard(layer, shard=layer % 4)
-    real = wu.load_real_hidden()
-    rows = [real[i : i + 6] for i in (0, 777, 1500)] if real is not None else []
-    rows.append(torch.randn(6, 6144, dtype=torch.bfloat16))
-    for x in rows:
-        for variants in VARIANT_PAIRS:
-            _check_against_native_and_oracle(x.cuda().contiguous(), gate_up, down, variants)
-
-
-def _pdl_delay_module():
-    from torch.utils.cpp_extension import load
-
-    build_dir = os.environ.get("GLM_MOE_DSA_DENSE_MLP_JIT_BUILD_DIR")
-    if build_dir:
-        os.makedirs(build_dir, exist_ok=True)
-    return load(
-        name="dense_mlp_pdl_delay",
-        sources=[os.path.join(os.path.dirname(os.path.abspath(__file__)), "pdl_delay.cu")],
-        extra_cuda_cflags=["-O3", "-gencode=arch=compute_103a,code=sm_103a"]
-        if torch.cuda.get_device_capability() == (10, 3)
-        else ["-O3"],
-        build_directory=build_dir,
-        verbose=False,
-    )
-
-
-@pytest.mark.skipif(
-    os.environ.get("TRTLLM_ENABLE_PDL") != "1",
-    reason="the dense_mlp kernels launch with PDL only when TRTLLM_ENABLE_PDL=1",
-)
-@pytest.mark.parametrize("variants", [(1, 0), (0, 0)])
-@pytest.mark.parametrize("delay_ns", [0, 20_000, 200_000])
-def test_pdl_delayed_producer(variants, delay_ns):
-    """Activations written late by a PDL predecessor are the ones consumed.
-
-    The predecessor releases its dependents immediately and writes x only
-    after a delay; the dense_mlp kernels issue weight loads before
-    griddepcontrol.wait but must read x after it, in eager mode and in a
-    replayed CUDA graph.
-    """
-    mod = _pdl_delay_module()
-    gate_up, down = _random_weights(6144, 3072, seed=17)
-    x_old = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
-    x_new = torch.randn(6, 6144, device="cuda", dtype=torch.bfloat16)
-    expected = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_new, gate_up, down, *variants)
-    stale = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x_old, gate_up, down, *variants)
-    assert not torch.equal(expected, stale)
-    x = x_old.clone()
-    mod.delayed_copy(x_new, x, delay_ns)
-    out = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
-    torch.cuda.synchronize()
-    assert torch.equal(out, expected)
-
-    x.copy_(x_old)
-    graph = torch.cuda.CUDAGraph()
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        mod.delayed_copy(x_new, x, delay_ns)
-        torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    with torch.cuda.graph(graph):
-        mod.delayed_copy(x_new, x, delay_ns)
-        out_g = torch.ops.trtllm.glm_moe_dsa_dense_mlp(x, gate_up, down, *variants)
-    for _ in range(3):
-        x.copy_(x_old)
-        graph.replay()
-        torch.cuda.synchronize()
-        assert torch.equal(out_g, expected)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py
deleted file mode 100644
index 4fe078ecd6c1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_front/test_draft_front_draft_front.py
+++ /dev/null
@@ -1,853 +0,0 @@
-# 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.
-"""draft_front (the draft_front contract) against the native ``DeepseekV3MTP`` front end.
-
-The oracle runs the native ops of ``DeepseekV3MTP.forward``: TensorRT-LLM
-``RMSNorm`` modules for ``enorm``/``hnorm`` (run in parallel through
-``maybe_execute_in_parallel``), ``torch.concat`` and ``torch.chunk``. The
-candidate must reproduce chunk ``tp_rank`` bit for bit, with the native view's
-shape and strides.
-
-Fixtures:
-* synthetic: random activations and weights (labelled ``synthetic``).
-* real: ``enorm``/``hnorm``/``eh_proj``/``model.norm``/``shared_head.norm``
-  weights and embedding rows read from the GLM-5.2 NVFP4 checkpoint
-  (``GLM52_NVFP4_CKPT``); the hidden states entering the draft layer are
-  synthetic pre-norm activations passed through the real ``model.norm`` (first
-  draft pass) or ``shared_head.norm`` (later passes) weights.
-
-The native-norm-path tests check that draft_front serves a call only when the native
-``enorm``/``hnorm`` run the FlashInfer CuTe DSL RMSNorm that it mirrors, and
-leaves every other native path (gemma, cuda.tile, NVFP4, torch fallback,
-FlashInfer CUDA JIT norm, unverified FlashInfer) to the native front.
-"""
-
-import functools
-import json
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn.functional as F
-from utils.llm_data import llm_models_root
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_front import (  # noqa: E501
-    native_norm,
-    native_norm_reason,
-    pack_draft_front,
-    plan_segments,
-    unsupported_reason,
-)
-from tensorrt_llm._torch.modules import rms_norm as rms_norm_module
-from tensorrt_llm._torch.modules.multi_stream_utils import (
-    maybe_execute_in_parallel,
-    with_multi_stream,
-)
-from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-
-GLM_H = 6144
-GLM_EPS = 1e-5
-GLM_TP = 8
-CKPT = os.environ.get(
-    "GLM52_NVFP4_CKPT",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-# Draft passes at K=5 for one request: first pass verifies 6 rows, then 1 row each.
-DRAFT_ROWS = (6, 1)
-
-requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import plan_testlib  # noqa: E402
-
-
-class NativeDraftFront:
-    """Native ``DeepseekV3MTP`` front end: parallel enorm/hnorm, concat, chunk."""
-
-    def __init__(
-        self,
-        enorm_weight: torch.Tensor,
-        hnorm_weight: torch.Tensor,
-        enorm_eps: float = GLM_EPS,
-        hnorm_eps: float = GLM_EPS,
-    ):
-        hidden_size = enorm_weight.shape[0]
-        self.enorm = RMSNorm(hidden_size=hidden_size, eps=enorm_eps, dtype=torch.bfloat16).cuda()
-        self.hnorm = RMSNorm(hidden_size=hidden_size, eps=hnorm_eps, dtype=torch.bfloat16).cuda()
-        with torch.no_grad():
-            self.enorm.weight.copy_(enorm_weight)
-            self.hnorm.weight.copy_(hnorm_weight)
-        self.aux_stream = torch.cuda.Stream()
-        self.events = (torch.cuda.Event(), torch.cuda.Event())
-
-    def __call__(
-        self, embeds: torch.Tensor, hidden: torch.Tensor, tp_size: int, tp_rank: int
-    ) -> torch.Tensor:
-        e, h = maybe_execute_in_parallel(
-            lambda: self.enorm(embeds),
-            lambda: self.hnorm(hidden),
-            self.events[0],
-            self.events[1],
-            self.aux_stream,
-            disable_on_compile=True,
-        )
-        x = torch.concat([e, h], dim=-1)
-        if tp_size > 1:
-            x = torch.chunk(x, tp_size, dim=-1)[tp_rank]
-        return x
-
-    def candidate(
-        self, embeds: torch.Tensor, hidden: torch.Tensor, tp_size: int, tp_rank: int
-    ) -> torch.Tensor | None:
-        return pack_draft_front(
-            embeds,
-            hidden,
-            self.enorm.weight,
-            self.enorm.variance_epsilon,
-            self.hnorm.weight,
-            self.hnorm.variance_epsilon,
-            tp_rank,
-            tp_size,
-        )
-
-
-def assert_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor, what: str = "") -> None:
-    assert actual is not None, f"{what}: candidate fell back to native"
-    assert actual.shape == expected.shape, (what, actual.shape, expected.shape)
-    assert actual.stride() == expected.stride(), (what, actual.stride(), expected.stride())
-    assert actual.dtype == expected.dtype
-    a = actual.contiguous().view(torch.int16)
-    b = expected.contiguous().view(torch.int16)
-    mismatch = a != b
-    if bool(mismatch.any()):
-        diff = (actual.float() - expected.float()).abs()
-        pytest.fail(
-            f"{what}: {int(mismatch.sum())}/{mismatch.numel()} elements differ, "
-            f"max |diff| {float(diff.max()):.3e}"
-        )
-
-
-def make_inputs(
-    rows: int, hidden_size: int = GLM_H, seed: int = 0, scale_spread: bool = True
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Synthetic BF16 embeddings / hidden states with per-row scale spread and outliers."""
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    e = torch.randn(rows, hidden_size, device="cuda", generator=gen)
-    h = torch.randn(rows, hidden_size, device="cuda", generator=gen)
-    if scale_spread:
-        scales = torch.logspace(-2, 2, steps=max(rows, 1), device="cuda")[:rows, None]
-        e = e * scales.flip(0) * 0.05
-        h = h * scales
-        h[:, 7] *= 50.0  # an outlier channel, as in residual streams
-    return e.bfloat16(), h.bfloat16()
-
-
-def make_weights(hidden_size: int = GLM_H, seed: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    we = (1.0 + 0.5 * torch.randn(hidden_size, device="cuda", generator=gen)).bfloat16()
-    wh = (0.5 + 0.25 * torch.randn(hidden_size, device="cuda", generator=gen)).bfloat16()
-    return we, wh
-
-
-def fp64_reference(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
-    x64 = x.double()
-    return x64 * torch.rsqrt(x64.pow(2).mean(-1, keepdim=True) + eps) * w.double()
-
-
-# --------------------------------------------------------------------------- planning (CPU)
-
-
-@pytest.mark.parametrize("hidden_size", [GLM_H, 7168, 2048])
-@pytest.mark.parametrize("tp_size", [1, 2, 3, 4, 8])
-def test_plan_segments_cover_native_chunks(hidden_size, tp_size):
-    if (2 * hidden_size) % tp_size:
-        pytest.skip("not a native chunking")
-    width = 2 * hidden_size // tp_size
-    for rank in range(tp_size):
-        cols = []
-        for seg in plan_segments(hidden_size, tp_size, rank):
-            assert 0 <= seg.begin < seg.end <= hidden_size
-            cols.extend(
-                range(seg.source * hidden_size + seg.begin, seg.source * hidden_size + seg.end)
-            )
-        assert cols == list(range(rank * width, (rank + 1) * width))
-
-
-def test_plan_segments_glm_tp8_one_source_per_rank():
-    for rank in range(GLM_TP):
-        (seg,) = plan_segments(GLM_H, GLM_TP, rank)
-        assert seg.source == (0 if rank < 4 else 1)
-        assert seg.end - seg.begin == 1536
-
-
-def test_plan_segments_rejects_bad_tp():
-    with pytest.raises(ValueError):
-        plan_segments(GLM_H, 5, 0)
-    with pytest.raises(ValueError):
-        plan_segments(GLM_H, 8, 8)
-
-
-# --------------------------------------------------------------------------- equivalence
-
-
-@requires_cuda
-@pytest.mark.parametrize("rows", [1, 2, 3, 6, 7, 64, 128, 1000])
-@pytest.mark.parametrize("tp_size", [1, 2, 3, 4, 8])
-def test_matches_native_all_ranks_synthetic(rows, tp_size):
-    e, h = make_inputs(rows, seed=rows)
-    native = NativeDraftFront(*make_weights())
-    for rank in range(tp_size):
-        assert_bitwise_equal(
-            native.candidate(e, h, tp_size, rank),
-            native(e, h, tp_size, rank),
-            f"synthetic rows={rows} tp={tp_size} rank={rank}",
-        )
-
-
-@requires_cuda
-@pytest.mark.parametrize("hidden_size", [2048, 4096, 7168])
-@pytest.mark.parametrize("rows", [1, 6])
-def test_matches_native_other_hidden_sizes(hidden_size, rows):
-    e, h = make_inputs(rows, hidden_size, seed=3)
-    native = NativeDraftFront(*make_weights(hidden_size))
-    for rank in range(GLM_TP):
-        assert_bitwise_equal(
-            native.candidate(e, h, GLM_TP, rank),
-            native(e, h, GLM_TP, rank),
-            f"H={hidden_size} rows={rows} rank={rank}",
-        )
-
-
-@requires_cuda
-@pytest.mark.parametrize("enorm_eps, hnorm_eps", [(1e-5, 1e-5), (1e-6, 1e-3), (1e-2, 1e-6)])
-def test_epsilon_per_source(enorm_eps, hnorm_eps):
-    # Small magnitudes so that eps dominates mean(x^2) for some rows.
-    e, h = make_inputs(6, seed=11)
-    e = (e.float() * 1e-3).bfloat16()
-    h = (h.float() * 1e-3).bfloat16()
-    native = NativeDraftFront(*make_weights(), enorm_eps=enorm_eps, hnorm_eps=hnorm_eps)
-    for tp_size in (1, 3, 8):
-        for rank in range(tp_size):
-            assert_bitwise_equal(
-                native.candidate(e, h, tp_size, rank),
-                native(e, h, tp_size, rank),
-                f"eps=({enorm_eps},{hnorm_eps}) tp={tp_size} rank={rank}",
-            )
-
-
-@requires_cuda
-def test_extreme_rows():
-    rows = 8
-    e, h = make_inputs(rows, seed=5, scale_spread=False)
-    e[0].zero_()
-    h[1].zero_()
-    e[2] = 3.0e4
-    h[3] = -2.0e4
-    e[4] = 1.0e-30
-    h[5].copy_(torch.linspace(-1.0e3, 1.0e3, GLM_H, device="cuda").bfloat16())
-    h[6, ::2] = 0.0
-    e[7, 100] = 6.0e4
-    native = NativeDraftFront(*make_weights())
-    for rank in range(GLM_TP):
-        assert_bitwise_equal(
-            native.candidate(e, h, GLM_TP, rank), native(e, h, GLM_TP, rank), f"extreme rank={rank}"
-        )
-
-
-@requires_cuda
-def test_strided_and_gathered_inputs():
-    """Draft-loop operands: row slices of a larger buffer and gathered rows."""
-    native = NativeDraftFront(*make_weights())
-    big_e, big_h = make_inputs(16, seed=9)
-    wide = torch.empty(16, 2 * GLM_H, dtype=torch.bfloat16, device="cuda")
-    wide[:, GLM_H:] = big_h
-    cases = {
-        "row_slice": (big_e[3:9], big_h[3:9]),
-        "wide_row_stride": (big_e[:6], wide[:6, GLM_H:]),
-        "gathered": (
-            big_e[torch.tensor([5, 0, 9], device="cuda")],
-            big_h[torch.tensor([2, 2, 15], device="cuda")],
-        ),
-    }
-    for name, (e, h) in cases.items():
-        for rank in range(GLM_TP):
-            assert_bitwise_equal(
-                native.candidate(e, h, GLM_TP, rank),
-                native(e, h, GLM_TP, rank),
-                f"{name} rank={rank}",
-            )
-
-
-@requires_cuda
-def test_padded_rows_do_not_leak():
-    """Rows past the live count hold stale data; live rows must not depend on them."""
-    native = NativeDraftFront(*make_weights())
-    live = 2
-    e, h = make_inputs(6, seed=21)
-    e2, h2 = e.clone(), h.clone()
-    e2[live:] = 1.0e4
-    h2[live:] = -7.0
-    for rank in range(GLM_TP):
-        out_a = native.candidate(e, h, GLM_TP, rank)[:live].clone()
-        out_b = native.candidate(e2, h2, GLM_TP, rank)
-        assert torch.equal(out_a, out_b[:live])
-        assert_bitwise_equal(out_b, native(e2, h2, GLM_TP, rank), f"padded rank={rank}")
-
-
-@requires_cuda
-@pytest.mark.parametrize("pdl", ["0", "1"])
-def test_cuda_graph_replay(monkeypatch, pdl):
-    """Capture once per rank and draft-row count; replay with fresh inputs."""
-    monkeypatch.setenv("TRTLLM_ENABLE_PDL", pdl)
-    native = NativeDraftFront(*make_weights())
-    for rows in DRAFT_ROWS:
-        for rank in range(GLM_TP):
-            static_e, static_h = make_inputs(rows, seed=100)
-            native.candidate(static_e, static_h, GLM_TP, rank)  # compile outside capture
-            torch.cuda.synchronize()
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph):
-                static_out = native.candidate(static_e, static_h, GLM_TP, rank)
-            for step in range(8):
-                e, h = make_inputs(rows, seed=200 + step)
-                static_e.copy_(e)
-                static_h.copy_(h)
-                graph.replay()
-                torch.cuda.synchronize()
-                assert_bitwise_equal(
-                    static_out,
-                    native(static_e, static_h, GLM_TP, rank),
-                    f"graph pdl={pdl} rows={rows} rank={rank} step={step}",
-                )
-
-
-@requires_cuda
-def test_graph_with_native_neighbors():
-    """Embedding lookup -> candidate -> eh_proj GEMM captured in one multi-stream graph."""
-    vocab = 4096
-    gen = torch.Generator(device="cuda").manual_seed(7)
-    table = torch.randn(vocab, GLM_H, device="cuda", generator=gen).bfloat16()
-    w_eh = (torch.randn(GLM_H, 2 * GLM_H, device="cuda", generator=gen) * 0.02).bfloat16()
-    native = NativeDraftFront(*make_weights())
-    width = 2 * GLM_H // GLM_TP
-    for rows in DRAFT_ROWS:
-        ids = torch.randint(0, vocab, (rows,), device="cuda", generator=gen)
-        _, h = make_inputs(rows, seed=31)
-        for rank in range(GLM_TP):
-            w_r = w_eh[:, rank * width : (rank + 1) * width].contiguous()
-
-            def native_step():
-                return F.linear(native(F.embedding(ids, table), h, GLM_TP, rank), w_r)
-
-            def candidate_step():
-                return F.linear(native.candidate(F.embedding(ids, table), h, GLM_TP, rank), w_r)
-
-            expected = native_step()
-            candidate_step()
-            torch.cuda.synchronize()
-            graph = torch.cuda.CUDAGraph()
-            with with_multi_stream(True), torch.cuda.graph(graph):
-                out_native = native_step()
-                out_cand = candidate_step()
-            for step in range(4):
-                ids.copy_(torch.randint(0, vocab, (rows,), device="cuda", generator=gen))
-                graph.replay()
-                torch.cuda.synchronize()
-                expected = native_step()
-                assert_bitwise_equal(out_native, expected, f"native-in-graph rank={rank}")
-                assert_bitwise_equal(out_cand, expected, f"neighbors rows={rows} rank={rank}")
-
-
-@requires_cuda
-def test_capture_before_compile_falls_back():
-    """A configuration first seen inside CUDA-graph capture takes the native path."""
-    hidden_size, rank = 1024, 5
-    e, h = make_inputs(6, hidden_size, seed=55)
-    native = NativeDraftFront(*make_weights(hidden_size))
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        assert native.candidate(e, h, GLM_TP, rank) is None
-    assert_bitwise_equal(
-        native.candidate(e, h, GLM_TP, rank),
-        native(e, h, GLM_TP, rank),
-        "eager after capture fallback",
-    )
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        static_out = native.candidate(e, h, GLM_TP, rank)
-    assert static_out is not None
-    graph.replay()
-    torch.cuda.synchronize()
-    assert_bitwise_equal(static_out, native(e, h, GLM_TP, rank), "capture after compile")
-
-
-# --------------------------------------------------------------------------- served neighbours
-
-
-def make_eh_shard(tp_size: int = GLM_TP, seed: int = 9) -> torch.Tensor:
-    """A ``[H, 2H / tp]`` BF16 stand-in for the local ROW-parallel ``eh_proj`` weight shard."""
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    return (
-        torch.randn(GLM_H, 2 * GLM_H // tp_size, device="cuda", generator=gen) * 0.02
-    ).bfloat16()
-
-
-@requires_cuda
-@pytest.mark.parametrize("pdl", ["0", "1"])
-def test_graph_with_served_neighbors(monkeypatch, pdl):
-    """LM-head-sized L2 eviction -> embedding -> front -> eh_proj GEMM -> input RMSNorm."""
-    monkeypatch.setenv("TRTLLM_ENABLE_PDL", pdl)
-    vocab = 4096
-    gen = torch.Generator(device="cuda").manual_seed(17)
-    table = torch.randn(vocab, GLM_H, device="cuda", generator=gen).bfloat16()
-    flush_w = torch.randn(19360, GLM_H, device="cuda", generator=gen).bfloat16()
-    flush_x = torch.randn(1, GLM_H, device="cuda", generator=gen).bfloat16()
-    native = NativeDraftFront(*make_weights())
-    w_eh = make_eh_shard()
-    input_norm = RMSNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16).cuda()
-    for rows in DRAFT_ROWS:
-        ids = torch.randint(0, vocab, (rows,), device="cuda", generator=gen)
-        _, h = make_inputs(rows, seed=33)
-        for rank in (0, 5):
-
-            def native_step():
-                F.linear(flush_x, flush_w)
-                return input_norm(F.linear(native(F.embedding(ids, table), h, GLM_TP, rank), w_eh))
-
-            def candidate_step():
-                F.linear(flush_x, flush_w)
-                front = native.candidate(F.embedding(ids, table), h, GLM_TP, rank)
-                return input_norm(F.linear(front, w_eh))
-
-            native_step()
-            candidate_step()
-            torch.cuda.synchronize()
-            graph = torch.cuda.CUDAGraph()
-            with with_multi_stream(True), torch.cuda.graph(graph):
-                out_native = native_step()
-                out_cand = candidate_step()
-            for step in range(4):
-                ids.copy_(torch.randint(0, vocab, (rows,), device="cuda", generator=gen))
-                graph.replay()
-                torch.cuda.synchronize()
-                assert_bitwise_equal(
-                    out_cand, out_native, f"served pdl={pdl} rows={rows} rank={rank}"
-                )
-
-
-# --------------------------------------------------------------------------- fallback
-
-
-@requires_cuda
-def test_unsupported_operands_take_native_path():
-    e, h = make_inputs(4)
-    we, wh = make_weights()
-    flat = torch.empty(4 * GLM_H + 8, dtype=torch.bfloat16, device="cuda")
-    misaligned = flat[1 : 1 + 4 * GLM_H].view(4, GLM_H)
-    wide = torch.empty(4, 2 * GLM_H, dtype=torch.bfloat16, device="cuda")
-    cases = {
-        "fp32": (e.float(), h.float(), we.float(), wh.float(), 8, 0),
-        "fp16_mixed": (e.half(), h, we, wh, 8, 0),
-        "misaligned": (misaligned, h, we, wh, 8, 0),
-        "column_stride": (e, wide[:, ::2], we, wh, 8, 0),
-        "rows_mismatch": (e, h[:3], we, wh, 8, 0),
-        "tp_not_dividing": (e, h, we, wh, 5, 0),
-        "rank_out_of_range": (e, h, we, wh, 8, 8),
-        "weight_shape": (e, h, we[:-8], wh, 8, 0),
-        "cpu": (e.cpu(), h.cpu(), we.cpu(), wh.cpu(), 8, 0),
-    }
-    for name, (ce, ch, cwe, cwh, tp_size, rank) in cases.items():
-        assert unsupported_reason(ce, ch, cwe, cwh, tp_size, rank) is not None, name
-        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
-            out = pack_draft_front(ce, ch, cwe, GLM_EPS, cwh, GLM_EPS, rank, tp_size)
-            torch.cuda.synchronize()
-        assert out is None, name
-        kernels = [
-            ev.name for ev in prof.events() if ev.device_type == torch.autograd.DeviceType.CUDA
-        ]
-        assert not kernels, (name, kernels)
-
-
-@requires_cuda
-def test_empty_rows():
-    e, h = make_inputs(1)
-    we, wh = make_weights()
-    out = pack_draft_front(e[:0], h[:0], we, GLM_EPS, wh, GLM_EPS, 3, GLM_TP)
-    assert out is not None and out.shape == (0, 2 * GLM_H // GLM_TP)
-
-
-# --------------------------------------------------------------------------- native norm path
-
-
-@pytest.fixture
-def fresh_norm_path():
-    """Recompute the per-device FlashInfer path decision inside the test and after it."""
-    native_norm._flashinfer_path_reason.cache_clear()
-    yield
-    native_norm._flashinfer_path_reason.cache_clear()
-
-
-def glm_norm(**kwargs) -> RMSNorm:
-    """``enorm``/``hnorm`` as ``DeepseekV3MTP`` builds them, on the CPU."""
-    return RMSNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16, **kwargs)
-
-
-def _with_attr(name: str, value) -> RMSNorm:
-    norm = glm_norm()
-    setattr(norm, name, value)
-    return norm
-
-
-class _OverridingNorm(RMSNorm):
-    def forward(self, hidden_states, *args, **kwargs):
-        return super().forward(hidden_states, *args, **kwargs)
-
-
-class _InheritingNorm(RMSNorm):
-    pass
-
-
-def _module_cases() -> dict:
-    """Norm modules whose native forward is not the plain FlashInfer RMSNorm."""
-    return {
-        "gemma": glm_norm(use_gemma=True),
-        "cuda_tile": _with_attr("use_cuda_tile", True),
-        "nvfp4": _with_attr("is_nvfp4", True),
-        "hp_output": _with_attr("return_hp_output", True),
-        "forward_override": _OverridingNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16),
-        "foreign_module": torch.nn.RMSNorm(GLM_H, eps=GLM_EPS),
-    }
-
-
-def test_native_norm_module_paths_decline():
-    plain = glm_norm()
-    assert native_norm.norm_module_reason("enorm", plain) is None
-    inheriting = _InheritingNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16)
-    assert native_norm.norm_module_reason("hnorm", inheriting) is None
-    for case, norm in _module_cases().items():
-        assert native_norm.norm_module_reason("enorm", norm) is not None, case
-        reason = native_norm_reason(norm, plain)
-        assert reason is not None and reason.startswith("enorm"), (case, reason)
-        reason = native_norm_reason(plain, norm)
-        assert reason is not None and reason.startswith("hnorm"), (case, reason)
-
-
-@requires_cuda
-def test_native_norm_glm_config_takes_mirrored_path(monkeypatch, fresh_norm_path):
-    """GLM's enorm/hnorm reach FlashInfer's CuTe DSL RMSNorm, so draft_front serves and is bitwise."""
-    import flashinfer.norm as flashinfer_norm
-
-    native = NativeDraftFront(*make_weights())
-    assert native_norm_reason(native.enorm, native.hnorm) is None
-    calls = {"cute": [], "cuda_jit": 0}
-    cute = flashinfer_norm.rmsnorm_cute
-    jit = flashinfer_norm.get_norm_module
-
-    def cute_spy(input, *args, **kwargs):
-        calls["cute"].append(input.dim())
-        return cute(input, *args, **kwargs)
-
-    def jit_spy():
-        calls["cuda_jit"] += 1
-        return jit()
-
-    monkeypatch.setattr(flashinfer_norm, "rmsnorm_cute", cute_spy)
-    monkeypatch.setattr(flashinfer_norm, "get_norm_module", jit_spy)
-    for rows in DRAFT_ROWS:
-        e, h = make_inputs(rows, seed=rows + 70)
-        assert unsupported_reason(e, h, native.enorm.weight, native.hnorm.weight, GLM_TP, 0) is None
-        for rank in (0, 4):
-            calls["cute"].clear()
-            expected = native(e, h, GLM_TP, rank)
-            assert calls["cute"] == [2, 2], (rows, rank, calls)
-            assert_bitwise_equal(
-                native.candidate(e, h, GLM_TP, rank),
-                expected,
-                f"mirrored path rows={rows} rank={rank}",
-            )
-    assert calls["cuda_jit"] == 0
-
-
-@requires_cuda
-def test_gemma_native_differs_from_candidate():
-    """A gemma enorm computes x * rstd * (w + 1): draft_front would change the result."""
-    we, wh = make_weights()
-    native = NativeDraftFront(we, wh)
-    gemma = glm_norm(use_gemma=True).cuda()
-    with torch.no_grad():
-        gemma.weight.copy_(we)
-    e, h = make_inputs(6, seed=81)
-    rank = 0  # an embedding-source rank
-    width = 2 * GLM_H // GLM_TP
-    native_gemma = gemma(e)[:, :width]
-    candidate = native.candidate(e, h, GLM_TP, rank)
-    assert not torch.equal(candidate, native_gemma)
-    assert native_norm_reason(gemma, native.hnorm) is not None
-
-
-def _process_case(monkeypatch, case: str) -> None:
-    import flashinfer
-    import flashinfer.norm as flashinfer_norm
-
-    if case == "flashinfer_unavailable":
-        monkeypatch.setattr(rms_norm_module, "IS_FLASHINFER_AVAILABLE", False)
-    elif case == "cuda_jit_norm":
-        monkeypatch.setattr(flashinfer_norm, "_use_cuda_norm", lambda device: True)
-    elif case == "unverified_flashinfer":
-        monkeypatch.setattr(flashinfer, "__version__", "0.0.0")
-    elif case == "no_norm_dispatch":
-        monkeypatch.delattr(flashinfer_norm, "_use_cuda_norm")
-    else:
-        raise AssertionError(case)
-
-
-@requires_cuda
-@pytest.mark.parametrize(
-    "case", ["flashinfer_unavailable", "cuda_jit_norm", "unverified_flashinfer", "no_norm_dispatch"]
-)
-def test_native_norm_process_paths_fall_back(monkeypatch, fresh_norm_path, case):
-    """Per-process conditions under which native RMSNorm leaves FlashInfer's CuTe kernel."""
-    native = NativeDraftFront(*make_weights())
-    e, h = make_inputs(6, seed=91)
-    rank = 4  # a hidden-source rank
-    candidate = native.candidate(e, h, GLM_TP, rank)
-    assert candidate is not None
-    _process_case(monkeypatch, case)
-    native_norm._flashinfer_path_reason.cache_clear()
-    assert native_norm.flashinfer_path_reason(e.device) is not None, case
-    assert native_norm_reason(native.enorm, native.hnorm) is not None, case
-    weights = (native.enorm.weight, native.hnorm.weight)
-    assert unsupported_reason(e, h, *weights, GLM_TP, rank) is not None, case
-    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
-        out = pack_draft_front(e, h, weights[0], GLM_EPS, weights[1], GLM_EPS, rank, GLM_TP)
-        torch.cuda.synchronize()
-    assert out is None, case
-    kernels = [ev.name for ev in prof.events() if ev.device_type == torch.autograd.DeviceType.CUDA]
-    assert not kernels, (case, kernels)
-    if case == "flashinfer_unavailable":
-        # The torch path rounds x * rstd to BF16 before the weight multiply.
-        torch_path = native(e, h, GLM_TP, rank)
-        assert not torch.equal(torch_path, candidate)
-
-
-@requires_cuda
-def test_adapter_native_norm_gate(monkeypatch, fresh_norm_path):
-    """The draft layer's adoption binds draft_front only when enorm/hnorm take the mirrored path.
-
-    Builds the built-in ``DeepseekV3MTP`` with the decoder-layer base, ``eh_proj`` and the
-    shared head stubbed out, then runs the draft step of its adoption as ``GlmMtpLayer``.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-        reclass,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmMtpLayer,
-    )
-    from tensorrt_llm._torch.models import modeling_deepseekv3 as model_module
-
-    plan_testlib.select(dispatch, "draft_front", monkeypatch)
-    monkeypatch.setattr(
-        model_module.DeepseekV3DecoderLayer,
-        "__init__",
-        lambda self, *args, **kwargs: torch.nn.Module.__init__(self),
-    )
-    monkeypatch.setattr(model_module, "Linear", lambda *args, **kwargs: torch.nn.Identity())
-    monkeypatch.setattr(
-        model_module, "DeepseekV3MTPHead", lambda *args, **kwargs: torch.nn.Identity()
-    )
-    aux_streams = {model_module.AuxStreamType.MoeShared: torch.cuda.Stream()}
-
-    def build(norm_factory=RMSNorm, attention_dp: bool = False, tp_size: int = GLM_TP):
-        monkeypatch.setattr(model_module, "RMSNorm", norm_factory)
-        config = SimpleNamespace(
-            pretrained_config=SimpleNamespace(
-                hidden_size=GLM_H,
-                moe_intermediate_size=2048,
-                n_routed_experts=256,
-                n_shared_experts=1,
-                num_experts_per_tok=8,
-                rms_norm_eps=GLM_EPS,
-                torch_dtype=torch.bfloat16,
-            ),
-            mapping=SimpleNamespace(enable_attention_dp=attention_dp, tp_size=tp_size, tp_rank=0),
-            skip_create_weights_in_init=True,
-            use_cute_dsl_blockscaling_mm=False,
-        )
-        layer = model_module.DeepseekV3MTP(config, 78, aux_streams)
-        # The stubbed decoder-layer constructor does not keep the model config.
-        layer.model_config = config
-        reclass(layer, model_module.DeepseekV3MTP, GlmMtpLayer)
-        layer._adopt_draft()
-        return layer.draft_front
-
-    def factory(name, value):
-        def make(**kwargs):
-            norm = RMSNorm(**kwargs)
-            setattr(norm, name, value)
-            return norm
-
-        return make
-
-    assert build() is pack_draft_front
-    assert build(functools.partial(RMSNorm, use_gemma=True)) is None
-    assert build(factory("use_cuda_tile", True)) is None
-    assert build(factory("is_nvfp4", True)) is None
-    assert build(attention_dp=True) is None
-    assert build(tp_size=1) is None
-    monkeypatch.setattr(rms_norm_module, "IS_FLASHINFER_AVAILABLE", False)
-    native_norm._flashinfer_path_reason.cache_clear()
-    assert build() is None
-
-
-def test_dispatch_registration(monkeypatch):
-    plan_testlib.select(dispatch, "draft_front", monkeypatch)
-    assert dispatch.is_enabled("draft_front")
-    assert dispatch.get_impl("draft_front", "draft_front") is pack_draft_front
-    assert dispatch.get_impl("draft_front", "unsupported_reason") is unsupported_reason
-    assert dispatch.get_impl("draft_front", "native_norm_reason") is native_norm_reason
-    # draft_front is the draft_front contract only: nothing that changes the eh_proj reduction
-    # or the draft layer's numerics is registered under this path.
-    with pytest.raises(KeyError):
-        dispatch.get_impl("draft_front", "input_norm_fold")
-
-
-def test_dispatch_default_off(monkeypatch):
-    """Without an installed plan (a process that does not build the target) draft_front is off."""
-    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
-    assert not dispatch.is_enabled("draft_front")
-
-
-# --------------------------------------------------------------------------- real checkpoint
-
-
-@functools.cache
-def _real_fixture() -> dict | None:
-    index_path = os.path.join(CKPT, "model.safetensors.index.json")
-    if not os.path.isfile(index_path):
-        return None
-    from safetensors import safe_open
-
-    with open(index_path) as f:
-        weight_map = json.load(f)["weight_map"]
-
-    def load(name: str, rows: list[int] | None = None) -> torch.Tensor:
-        with safe_open(os.path.join(CKPT, weight_map[name]), framework="pt", device="cpu") as f:
-            if rows is None:
-                return f.get_tensor(name)
-            sl = f.get_slice(name)
-            return torch.cat([sl[r : r + 1] for r in rows])
-
-    token_ids = [
-        0,
-        1,
-        2,
-        3,
-        13,
-        198,
-        271,
-        1000,
-        9999,
-        12345,
-        50000,
-        77777,
-        100000,
-        151329,
-        154000,
-        154879,
-    ]
-    return {
-        "enorm": load("model.layers.78.enorm.weight").cuda(),
-        "hnorm": load("model.layers.78.hnorm.weight").cuda(),
-        "eh_proj": load("model.layers.78.eh_proj.weight").cuda(),
-        "model_norm": load("model.norm.weight").cuda(),
-        "shared_head_norm": load("model.layers.78.shared_head.norm.weight").cuda(),
-        "embeds": load("model.embed_tokens.weight", token_ids).cuda(),
-        "token_ids": token_ids,
-    }
-
-
-def _real_hidden(rows: int, norm_weight: torch.Tensor, seed: int) -> torch.Tensor:
-    """Synthetic pre-norm residual activations through a real output RMSNorm weight."""
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    pre = torch.randn(rows, GLM_H, device="cuda", generator=gen) * 3.0
-    pre[:, 11] *= 40.0
-    pre = pre.bfloat16()
-    norm = RMSNorm(hidden_size=GLM_H, eps=GLM_EPS, dtype=torch.bfloat16).cuda()
-    with torch.no_grad():
-        norm.weight.copy_(norm_weight)
-    return norm(pre)
-
-
-@requires_cuda
-@pytest.mark.parametrize("draft_pass", [0, 1])
-def test_real_checkpoint_all_ranks(draft_pass):
-    fx = _real_fixture()
-    if fx is None:
-        pytest.skip(f"checkpoint not found at {CKPT}")
-    rows = DRAFT_ROWS[0] if draft_pass == 0 else DRAFT_ROWS[1]
-    e = fx["embeds"][draft_pass : draft_pass + rows]
-    h = _real_hidden(
-        rows, fx["model_norm"] if draft_pass == 0 else fx["shared_head_norm"], seed=40 + draft_pass
-    )
-    native = NativeDraftFront(fx["enorm"], fx["hnorm"])
-    width = 2 * GLM_H // GLM_TP
-    full_ref = torch.cat(
-        [fp64_reference(e, fx["enorm"], GLM_EPS), fp64_reference(h, fx["hnorm"], GLM_EPS)], dim=-1
-    )
-    for rank in range(GLM_TP):
-        out_native = native(e, h, GLM_TP, rank)
-        out_cand = native.candidate(e, h, GLM_TP, rank)
-        assert_bitwise_equal(out_cand, out_native, f"real pass={draft_pass} rank={rank}")
-        ref = full_ref[:, rank * width : (rank + 1) * width]
-        for label, out in (("native", out_native), ("candidate", out_cand)):
-            err = out.double() - ref
-            rel_rms = float(err.pow(2).mean().sqrt() / ref.pow(2).mean().sqrt())
-            print(
-                f"GLM_MOE_DSA_DRAFT_FRONT_STAT real pass={draft_pass} rank={rank} {label} "
-                f"max_abs={float(err.abs().max()):.3e} rel_rms={rel_rms:.3e}"
-            )
-        # eh_proj consumer: the local ROW-parallel GEMM sees an identical operand.
-        w_r = fx["eh_proj"][:, rank * width : (rank + 1) * width].contiguous()
-        assert_bitwise_equal(
-            F.linear(out_cand, w_r),
-            F.linear(out_native, w_r),
-            f"real eh_proj partial pass={draft_pass} rank={rank}",
-        )
-
-
-@requires_cuda
-def test_real_checkpoint_draft_loop_rows():
-    """All 16 real embedding rows through one first pass and later single-row passes."""
-    fx = _real_fixture()
-    if fx is None:
-        pytest.skip(f"checkpoint not found at {CKPT}")
-    native = NativeDraftFront(fx["enorm"], fx["hnorm"])
-    e_all = fx["embeds"]
-    h_all = _real_hidden(e_all.shape[0], fx["shared_head_norm"], seed=77)
-    for start in range(0, e_all.shape[0]):
-        for rows in (1, 6):
-            if start + rows > e_all.shape[0]:
-                continue
-            e, h = e_all[start : start + rows], h_all[start : start + rows]
-            for rank in range(GLM_TP):
-                assert_bitwise_equal(
-                    native.candidate(e, h, GLM_TP, rank),
-                    native(e, h, GLM_TP, rank),
-                    f"real loop start={start} rows={rows} rank={rank}",
-                )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py
deleted file mode 100644
index 530dd3a0785e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_live_rows/test_live_rows_adapter.py
+++ /dev/null
@@ -1,472 +0,0 @@
-# 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.
-"""The draft_live_rows call sites in the target and their host-side decisions.
-
-* the plan runs draft_live_rows with its closure (post_attention_hop, draft_loop), next to
-  the served GEN set without deferred_input_prep, and turns it off with either base; draft_live_rows has
-  one numerical form; the package reads no environment variable and
-  imports no other package (no deferred_input_prep side-selector or package-directory entry);
-* the call sites: ``GlmMtpLayer`` adoption builds the live-row hop only when draft_live_rows
-  is selected and its forward tries it only on an armed call that the post_attention_hop
-  serves, before the post_attention_hop; the draft loop (``GlmMtpWorker``) decides once per
-  forward, arms step 0 only, and makes draft_loop's step-0 gather the identity only when the
-  layer compacted;
-* the host-side eligibility of the live-row hop and of the step-0 arming,
-  case by case, with no device work; the arming is cleared even when the
-  step raises;
-* the live-row hop is built on every rank of the post_attention_hop's TP group or on none:
-  whether the operator is loaded is agreed over the group's communicator
-  (mocked here; one collective per post_attention_hop, also on a rank without the
-  operator), with a negative control showing that the rank-local rule
-  disagrees whenever one rank lacks the operator.
-"""
-
-import ast
-import inspect
-import os
-import re
-import sys
-import threading
-from pathlib import Path
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import plan_testlib  # noqa: E402
-
-GEN_FAST_PATHS_WITHOUT_INPUT_PREP = (
-    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,"
-    "router_gemm,post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,"
-    "rope_kv_append,draft_loop,shared_indexer_absorbed_query,routed_moe,"
-    "cross_rank_topk,attention_output_tail_prefetch,attention_output_into_hop,"
-    "routed_moe_into_hop,deferred_norm,dense_mlp,hop_routing,oproj_weight_staging,"
-    "fused_kv_indexer_projection,host_page_retention,shared_expert_prefetch,"
-    "draft_routed_moe,cross_layer_expert_prefetch"
-)
-
-
-def _dispatch():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    return dispatch
-
-
-def test_selector_facts(monkeypatch):
-    d = _dispatch()
-    monkeypatch.setattr(d, "_log_info", lambda message: None)
-    assert "draft_live_rows" in d.SELECTOR_PATHS
-    assert set(d.REQUIRED_FEATURES["draft_live_rows"]) == {"post_attention_hop", "draft_loop"}
-    assert "draft_live_rows" not in d.NUMERICAL_VARIANTS
-    assert plan_testlib.select(
-        d, "post_attention_hop,draft_loop,draft_live_rows", monkeypatch
-    ) == frozenset({"post_attention_hop", "draft_loop", "draft_live_rows"})
-    no_input_prep = frozenset(GEN_FAST_PATHS_WITHOUT_INPUT_PREP.split(","))
-    assert plan_testlib.select(
-        d, GEN_FAST_PATHS_WITHOUT_INPUT_PREP + ",draft_live_rows", monkeypatch
-    ) == no_input_prep | {"draft_live_rows"}
-    assert plan_testlib.select(
-        d,
-        GEN_FAST_PATHS_WITHOUT_INPUT_PREP + ",cross_rank_topk_cluster_merge,draft_live_rows",
-        monkeypatch,
-    ) >= {
-        "cross_rank_topk_cluster_merge",
-        "draft_live_rows",
-    }
-    for base in ("post_attention_hop", "draft_loop"):
-        plan = d.resolve(d.ModelingV2Mode.AUTO, {d.DISABLE_ENV: base})
-        assert not plan.enabled("draft_live_rows"), base
-
-
-def _package_dir() -> Path:
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
-        draft_live_rows,
-    )
-
-    return Path(draft_live_rows.__file__).parent
-
-
-def test_package_is_self_contained():
-    """No side selector and no import of another package's code."""
-    for path in sorted(_package_dir().glob("*.py")):
-        src = path.read_text()
-        assert "environ" not in src and "getenv" not in src, path.name
-        for node in ast.walk(ast.parse(src)):
-            if isinstance(node, ast.ImportFrom) and node.level >= 2:
-                assert node.module == "dispatch", (path.name, node.module)
-
-
-def test_mtp_call_site():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmMtpLayer,
-    )
-
-    init = inspect.getsource(GlmMtpLayer._adopt_draft)
-    assert 'is_enabled("draft_live_rows")' in init and re.search(
-        r'get_impl\(\s*"draft_live_rows",\s*"make_live_row_hop"\s*\)', init
-    )
-    fwd = inspect.getsource(GlmMtpLayer.forward)
-    draft_live_rows = fwd.index("self.draft_live_row_hop.can_run(")
-    # The post_attention_hop decision is made once, before either hop; the live-row hop is tried
-    # first and only when armed.
-    assert (
-        fwd.index("use_post_attn_hop = ") < draft_live_rows < fwd.index("elif use_post_attn_hop:")
-    )
-    assert fwd.index("live_rows is not None") < draft_live_rows
-    assert fwd.count("self._can_use_post_attn_hop(") == 1
-    assert "self._draft_live_rows_compacted = True" in fwd
-
-
-def test_draft_loop_call_site():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
-        GlmMtpWorker,
-    )
-
-    src = inspect.getsource(GlmMtpWorker._forward_linear_draft_loop)
-    match = re.search(r'get_impl\(\s*"draft_live_rows",\s*"select_step0"\s*\)', src)
-    assert match is not None
-    sel = match.start()
-    assert sel < src.index("for i in range(runtime_draft_len):")
-    assert "if i == 0 and live_rows_step0 is not None:" in src
-    ident = src.index("if live_rows_compacted:")
-    assert ident < src.index("draft_loop_ops.step0_rows(hidden_states)")
-    assert "gather_ids, draft_rows = draft_loop_ops.gather_ids, hidden_states" in src
-
-
-class _AttnHopStub:
-    enabled = True
-    rank = 2
-    uc_buffer = None
-    mc_ptr = 1
-    flags = None
-    min_sm_count = 148
-
-
-def test_live_row_hop_eligibility():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows.hop import (  # noqa: E501
-        LiveRowHop,
-        required_sms,
-    )
-
-    hop = LiveRowHop(_AttnHopStub())
-    h6 = torch.zeros(6, 6144, dtype=torch.bfloat16)
-    w = torch.zeros(256, 6144, dtype=torch.bfloat16)
-    ids = torch.zeros(1, dtype=torch.int64)
-    assert hop.call_ineligibility(h6, ids, w) is None
-    assert hop.call_ineligibility(h6, torch.zeros(2, dtype=torch.int64), w) is None
-    assert hop.call_ineligibility(h6, ids, None) == "no router weight"
-    assert hop.call_ineligibility(h6, None, w) == "no live rows"
-    assert hop.call_ineligibility(h6, ids.int(), w) == "live-row ids layout"
-    assert hop.call_ineligibility(h6, ids.view(1, 1), w) == "live-row ids layout"
-    assert (
-        hop.call_ineligibility(h6, torch.zeros(4, dtype=torch.int64)[::2], w)
-        == "live-row ids layout"
-    )
-    assert hop.call_ineligibility(h6, torch.zeros(7, dtype=torch.int64), w).startswith(
-        "live rows 7 of 6"
-    )
-    assert hop.call_ineligibility(h6, torch.zeros(0, dtype=torch.int64), w).startswith(
-        "live rows 0"
-    )
-    h12 = torch.zeros(12, 6144, dtype=torch.bfloat16)
-    assert hop.call_ineligibility(h12, ids, w).startswith("live rows 1 of 12")
-    small = _AttnHopStub()
-    small.min_sm_count = 24
-    assert LiveRowHop(small).call_ineligibility(h6, ids, w) == "SM count"
-    assert [required_sms(r) for r in (1, 2, 3, 8)] == [32, 64, 96, 96]
-
-
-class _Comm:
-    """Mocked TP communicator (MPI style): this rank's value plus fixed peer values."""
-
-    def __init__(self, peers):
-        self.peers = [int(v) for v in peers]
-        self.calls = []
-
-    def allreduce(self, value, op=None):
-        self.calls.append(int(value))
-        return int(value) + sum(self.peers)
-
-    def Get_size(self):
-        return len(self.peers) + 1
-
-
-class _FakeAttnHop(_AttnHopStub):
-    def __init__(self, comm, rank=2):
-        self.comm = comm
-        self.rank = rank
-
-
-@pytest.fixture
-def live_row_hop(monkeypatch):
-    """The hop module with an empty decision table and the MPI agreement path."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows import (  # noqa: E501
-        hop,
-    )
-
-    monkeypatch.setattr(hop.LiveRowHop, "_instances", {})
-    monkeypatch.delenv("TLLM_DISABLE_MPI", raising=False)
-    return hop
-
-
-@pytest.mark.parametrize(
-    "local, peers, built, reason",
-    [
-        (True, [1] * 7, True, ""),
-        (True, [1] * 6 + [0], False, "another rank lacks trtllm::glm_moe_dsa_draft_live_row_hop"),
-        (
-            True,
-            [1] * 3 + [0] * 4,
-            False,
-            "another rank lacks trtllm::glm_moe_dsa_draft_live_row_hop",
-        ),
-        (False, [1] * 7, False, "trtllm::glm_moe_dsa_draft_live_row_hop is not loaded"),
-        (False, [0] * 7, False, "trtllm::glm_moe_dsa_draft_live_row_hop is not loaded"),
-    ],
-)
-def test_live_row_hop_agreed_over_tp_group(live_row_hop, monkeypatch, local, peers, built, reason):
-    monkeypatch.setattr(live_row_hop, "op_available", lambda: local)
-    comm = _Comm(peers)
-    w2 = _FakeAttnHop(comm)
-    hop = live_row_hop.LiveRowHop.for_post_attn_hop(w2)
-    assert (hop is not None) == built
-    # This rank takes part in the agreement whatever its own answer: a rank that
-    # skipped it would leave its peers waiting in the collective.
-    assert comm.calls == [int(local)]
-    assert live_row_hop.LiveRowHop._instances[id(w2)] == (w2, hop, reason)
-    if built:
-        assert hop.w2 is w2 and hop.rank == w2.rank
-    # Later calls for the same post_attention_hop return the recorded decision without
-    # communicating.
-    assert live_row_hop.LiveRowHop.for_post_attn_hop(w2) is hop
-    assert comm.calls == [int(local)]
-
-
-def test_live_row_hop_without_attn_hop_group(live_row_hop, monkeypatch):
-    """No enabled post_attention_hop (identical on all ranks) or no communicator: no hop, no collective."""
-    monkeypatch.setattr(live_row_hop, "op_available", lambda: True)
-    assert live_row_hop.LiveRowHop.for_post_attn_hop(None) is None
-    comm = _Comm([1] * 7)
-    off = _FakeAttnHop(comm)
-    off.enabled = False
-    assert live_row_hop.LiveRowHop.for_post_attn_hop(off) is None
-    lone = _FakeAttnHop(None)
-    assert live_row_hop.LiveRowHop.for_post_attn_hop(lone) is None
-    assert (
-        live_row_hop.LiveRowHop._instances[id(lone)][2]
-        == "post_attention_hop hop has no communicator"
-    )
-    assert comm.calls == []
-
-
-class _ThreadComm:
-    """One in-process TP group of ``size`` threads; a single integer-sum allreduce."""
-
-    def __init__(self, size):
-        self.size = size
-        self.total = 0
-        self.lock = threading.Lock()
-        self.barrier = threading.Barrier(size, timeout=60)
-
-    def allreduce(self, value, op=None):
-        with self.lock:
-            self.total += int(value)
-        self.barrier.wait()
-        return self.total
-
-    def Get_size(self):
-        return self.size
-
-
-@pytest.mark.parametrize("lacking", [(), (0,), (3,), (5,), (7,), (4, 5, 6, 7), tuple(range(8))])
-def test_live_row_hop_rank_agreement_eight_ranks(live_row_hop, monkeypatch, lacking):
-    """Eight ranks (threads) build the hop concurrently: every rank gets it, or none does."""
-    size = 8
-    tls = threading.local()
-    monkeypatch.setattr(live_row_hop, "op_available", lambda: tls.has_op)
-    comm = _ThreadComm(size)
-    built, local_rule, errors = [None] * size, [None] * size, []
-
-    def rank_main(r):
-        tls.has_op = r not in lacking
-        try:
-            built[r] = (
-                live_row_hop.LiveRowHop.for_post_attn_hop(_FakeAttnHop(comm, rank=r)) is not None
-            )
-            # The construction rule before the agreement: this rank's own operator only.
-            local_rule[r] = live_row_hop.op_available()
-        except Exception as error:  # noqa: BLE001  (reported below with the rank)
-            errors.append((r, repr(error)))
-
-    threads = [threading.Thread(target=rank_main, args=(r,)) for r in range(size)]
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join(timeout=120)
-    assert not any(t.is_alive() for t in threads) and not errors, errors
-    assert built == [not lacking] * size
-    # Negative control: the rank-local rule splits the group whenever some but not all ranks lack the op.
-    split = 0 < len(lacking) < size
-    assert (len(set(local_rule)) == 2) == split
-
-
-class _Layer:
-    """The draft layer's live-row state and its arming API (``GlmMtpLayer``)."""
-
-    def __init__(self, hop):
-        self.draft_live_row_hop = hop
-        self._draft_live_rows = None
-        self._draft_live_rows_compacted = False
-
-    def arm_live_rows(self, rows):
-        self._draft_live_rows = rows
-        self._draft_live_rows_compacted = False
-
-    def disarm_live_rows(self):
-        compacted = bool(self._draft_live_rows_compacted)
-        self._draft_live_rows = None
-        self._draft_live_rows_compacted = False
-        return compacted
-
-
-class _Draft:
-    def __init__(self, hop):
-        self.mtp_layers = [_Layer(hop)]
-
-
-class _Spec:
-    layers_to_capture = ()
-
-
-class _DRAFT_LOOP:
-    def __init__(self, b):
-        self.batch_size = b
-        self.gather_ids = torch.zeros(b, dtype=torch.int64)
-
-
-class _Worker:
-    is_mtp_eagle = True
-
-
-def test_step0_eligibility():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows.step0 import (  # noqa: E501
-        select_step0,
-        step0_ineligibility,
-    )
-
-    hop = object()
-    draft_loop = _DRAFT_LOOP(1)
-    assert step0_ineligibility(_Worker(), _Draft(hop), _Spec(), draft_loop) is None
-    other = _Worker()
-    other.is_mtp_eagle = False
-    assert step0_ineligibility(other, _Draft(hop), _Spec(), draft_loop) == "not MTP Eagle one-model"
-    assert (
-        step0_ineligibility(_Worker(), _Draft(hop), _Spec(), None) == "no draft_loop draft-loop ops"
-    )
-    cap = _Spec()
-    cap.layers_to_capture = (0,)
-    assert (
-        step0_ineligibility(_Worker(), _Draft(hop), cap, draft_loop) == "MTP layer output captured"
-    )
-    none = _Spec()
-    none.layers_to_capture = None
-    assert step0_ineligibility(_Worker(), _Draft(hop), none, draft_loop) is None
-    assert (
-        step0_ineligibility(_Worker(), _Draft(None), _Spec(), draft_loop)
-        == "MTP layer has no live-row hop"
-    )
-    bad = _DRAFT_LOOP(2)
-    bad.gather_ids = torch.zeros(3, dtype=torch.int64)
-    assert step0_ineligibility(_Worker(), _Draft(hop), _Spec(), bad) == "draft_loop gather ids"
-    bad.gather_ids = None
-    assert step0_ineligibility(_Worker(), _Draft(hop), _Spec(), bad) == "draft_loop gather ids"
-    assert select_step0(_Worker(), _Draft(None), _Spec(), draft_loop) is None
-    assert select_step0(_Worker(), _Draft(hop), _Spec(), draft_loop) is not None
-
-
-def test_arming_is_cleared():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows.step0 import (  # noqa: E501
-        select_step0,
-    )
-
-    draft = _Draft(object())
-    draft_loop = _DRAFT_LOOP(1)
-    layer = draft.mtp_layers[0]
-    s0 = select_step0(_Worker(), draft, _Spec(), draft_loop)
-    with s0.armed() as live:
-        assert layer._draft_live_rows is draft_loop.gather_ids
-        layer._draft_live_rows_compacted = True
-    assert (
-        live.compacted and layer._draft_live_rows is None and not layer._draft_live_rows_compacted
-    )
-    with s0.armed() as live:
-        pass
-    assert not live.compacted
-    with pytest.raises(RuntimeError):
-        with s0.armed():
-            layer._draft_live_rows_compacted = True
-            raise RuntimeError("step failed")
-    assert layer._draft_live_rows is None and not layer._draft_live_rows_compacted
-
-
-def test_mtp_layer_default_off(monkeypatch):
-    """Without draft_live_rows the MTP layer's draft_live_rows attributes stay inert (no hop, never armed)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.layers import (  # noqa: E501
-        GlmMtpLayer,
-    )
-
-    d = _dispatch()
-    monkeypatch.setattr(d, "_log_info", lambda message: None)
-    plan_testlib.select(d, GEN_FAST_PATHS_WITHOUT_INPUT_PREP, monkeypatch)
-    assert not d.is_enabled("draft_live_rows")
-    src = inspect.getsource(GlmMtpLayer._adopt_draft)
-    # The hop is created only under is_enabled("draft_live_rows"); the class defaults keep
-    # the attributes inert either way.
-    assert 'if self.post_attn_hop is not None and dispatch.is_enabled("draft_live_rows"):' in src
-    assert GlmMtpLayer.draft_live_row_hop is None and GlmMtpLayer._draft_live_rows is None
-
-
-@pytest.mark.skipif(
-    not os.environ.get("GLM_MOE_DSA_DRAFT_LIVE_ROWS_OP_LIB") and not torch.cuda.is_available(),
-    reason="needs the draft_live_rows op",
-)
-def test_op_registered_with_fake():
-    if not torch.cuda.is_available():
-        pytest.skip("needs a GPU")
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_live_rows import (  # noqa: E501
-        hop as live_row_hop,
-    )
-
-    if not live_row_hop.op_available():
-        pytest.skip("draft_live_rows op not loaded in this process")
-    live_row_hop._register_fake()  # idempotent
-    with torch._subclasses.fake_tensor.FakeTensorMode():
-        x = torch.empty(6, 6144, dtype=torch.bfloat16, device="cuda")
-        ids = torch.empty(2, dtype=torch.int64, device="cuda")
-        out = torch.ops.trtllm.glm_moe_dsa_draft_live_row_hop(
-            x,
-            x,
-            ids,
-            torch.empty(6144, dtype=torch.bfloat16, device="cuda"),
-            1e-5,
-            torch.empty(256, 6144, dtype=torch.bfloat16, device="cuda"),
-            torch.empty(10, device="cuda"),
-            1,
-            torch.empty(4, dtype=torch.int32, device="cuda"),
-            0,
-            8,
-        )
-        assert [tuple(t.shape) for t in out] == [(2, 6144), (2, 6144), (2, 256)]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py
deleted file mode 100644
index 6d5aac009ff9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/_draft_loop_fixtures.py
+++ /dev/null
@@ -1,311 +0,0 @@
-# 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.
-"""Shared helpers for the draft_loop tests and microbenchmarks."""
-
-import json
-import os
-from pathlib import Path
-from typing import Optional
-
-import torch
-from utils.llm_data import llm_models_root
-
-GLM52_CKPT_ENV = "GLM52_CKPT"
-_DEFAULT_CKPT = os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4")
-
-
-def checkpoint_dir() -> Optional[Path]:
-    """Return the GLM-5.2 NVFP4 checkpoint directory if it is readable."""
-    path = Path(os.environ.get(GLM52_CKPT_ENV, _DEFAULT_CKPT))
-    if (path / "model.safetensors.index.json").is_file():
-        return path
-    return None
-
-
-def load_checkpoint_tensor(name: str, device: str = "cuda") -> Optional[torch.Tensor]:
-    """Load one tensor of the GLM-5.2 checkpoint, or None if unavailable."""
-    ckpt = checkpoint_dir()
-    if ckpt is None:
-        return None
-    from safetensors import safe_open
-
-    weight_map = json.loads((ckpt / "model.safetensors.index.json").read_text())["weight_map"]
-    with safe_open(str(ckpt / weight_map[name]), framework="pt", device="cpu") as f:
-        return f.get_tensor(name).to(device)
-
-
-def real_token_ids(num_tokens: int) -> Optional[torch.Tensor]:
-    """Tokenize a fixed tool-use style text with the checkpoint tokenizer."""
-    ckpt = checkpoint_dir()
-    if ckpt is None:
-        return None
-    # The raw tokenizer avoids loading the model config through transformers.
-    from tokenizers import Tokenizer
-
-    tokenizer = Tokenizer.from_file(str(ckpt / "tokenizer.json"))
-    text = (
-        "<tool_call>search_files(path='/repo/src', pattern='def forward')</tool_call>\n"
-        "The draft loop embeds each proposed token, normalizes it, and projects the "
-        "concatenation with the previous hidden state; 你好，世界。 0x7f3a 3.14159 "
-    ) * 64
-    ids = tokenizer.encode(text, add_special_tokens=True).ids
-    assert len(ids) >= num_tokens
-    return torch.tensor(ids[:num_tokens], dtype=torch.int32)
-
-
-def bits(t: torch.Tensor) -> torch.Tensor:
-    """Bit view of a floating tensor (for bitwise equality incl. -0.0/NaN)."""
-    if t.dtype in (torch.bfloat16, torch.float16):
-        return t.view(torch.int16)
-    if t.dtype == torch.float32:
-        return t.view(torch.int32)
-    return t
-
-
-def graph_time_us(fn, reps: int = 50, replays: int = 20, warmup: int = 3) -> float:
-    """Mean device time of ``fn`` per call, measured as a CUDA graph of ``reps`` calls.
-
-    The graph holds ``reps`` back-to-back invocations so the result includes
-    the inter-kernel launch gaps a captured decode graph pays.
-    """
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(warmup):
-            fn()
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        for _ in range(reps):
-            fn()
-    for _ in range(3):
-        graph.replay()
-    torch.cuda.synchronize()
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(replays):
-        graph.replay()
-    end.record()
-    torch.cuda.synchronize()
-    elapsed_ms = start.elapsed_time(end)
-    del graph
-    return elapsed_ms * 1000.0 / (replays * reps)
-
-
-# ----------------------------------------------------------------------------
-# Real DSA generation metadata at the served GEN configuration
-# ----------------------------------------------------------------------------
-
-MTP_K = 5
-TOKENS_PER_BLOCK = 64
-
-
-def served_tp8_sparse_config():
-    """DSA config of the served GEN role (GLM-5.2 indexer geometry)."""
-    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
-
-    return DeepSeekSparseAttentionConfig(
-        index_n_heads=32,
-        index_head_dim=128,
-        index_topk=2048,
-        enable_heuristic_topk=True,
-        use_cute_dsl_paged_mqa_logits=True,
-        use_cute_dsl_topk=True,
-        indexer_k_dtype="fp8",
-    )
-
-
-def make_dsa_cache_manager(batch_size: int, max_kv_len: int, draft_len: int = MTP_K):
-    """A one-layer FP8 DSA cache manager sized for ``batch_size`` long requests."""
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm.bindings import DataType
-    from tensorrt_llm.bindings.executor import KvCacheConfig
-    from tensorrt_llm.bindings.internal.batch_manager import CacheType
-    from tensorrt_llm.mapping import Mapping
-
-    mapping = Mapping(world_size=1, tp_size=1, rank=0)
-    sparse_config = served_tp8_sparse_config()
-    model_config = ModelConfig(
-        mapping=mapping,
-        sparse_attention_config=sparse_config,
-        pretrained_config=SimpleNamespace(rms_norm_eps=1e-5),
-    )
-    max_seq_len = max_kv_len + 2 * (draft_len + 1) + TOKENS_PER_BLOCK
-    blocks = (max_seq_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK + 2
-    return (
-        DSACacheManager(
-            KvCacheConfig(
-                max_tokens=batch_size * blocks * TOKENS_PER_BLOCK, enable_block_reuse=False
-            ),
-            CacheType.SELFKONLY,
-            num_layers=1,
-            num_kv_heads=1,
-            head_dim=576,
-            tokens_per_block=TOKENS_PER_BLOCK,
-            max_seq_len=max_seq_len,
-            max_batch_size=batch_size,
-            mapping=mapping,
-            dtype=DataType.FP8,
-            sparse_attn_config=sparse_config,
-            model_config=model_config,
-        ),
-        mapping,
-        sparse_config,
-    )
-
-
-def make_gen_dsa_metadata(
-    cache_manager,
-    mapping,
-    sparse_config,
-    request_ids,
-    past_kv_lens,
-    draft_len: int = MTP_K,
-    draft_cache_manager=None,
-):
-    """Prepared generation metadata for one MTP verify step (K+1 query tokens each).
-
-    Follows the decoder runner's order: ``update_spec_dec_param`` (one-model MTP,
-    linear tree), then ``KVCacheParams`` with the one-model extra KV tokens and
-    ``prepare()``.
-    """
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata
-    from tensorrt_llm._torch.metadata import KVCacheParams
-
-    batch_size = len(request_ids)
-    md = DSAtrtllmAttentionMetadata(
-        seq_lens=torch.full((batch_size,), draft_len + 1, dtype=torch.int),
-        request_ids=list(request_ids),
-        max_num_requests=batch_size,
-        num_contexts=0,
-        prompt_lens=list(past_kv_lens),
-        max_num_tokens=batch_size * (draft_len + 1),
-        kv_cache_manager=cache_manager,
-        draft_kv_cache_manager=draft_cache_manager,
-        mapping=mapping,
-        sparse_attention_config=sparse_config,
-    )
-    md.update_spec_dec_param(
-        batch_size=batch_size,
-        is_spec_decoding_enabled=True,
-        is_spec_dec_tree=False,
-        is_spec_dec_dynamic_tree=False,
-        max_draft_len=draft_len,
-        max_total_draft_tokens=draft_len,
-        spec_metadata=None,
-        spec_tree_manager=None,
-        num_contexts=0,
-    )
-    md.kv_cache_params = KVCacheParams(
-        use_cache=True,
-        num_cached_tokens_per_seq=list(past_kv_lens),
-        num_extra_kv_tokens=draft_len - 1,
-    )
-    md.prepare()
-    return md
-
-
-def add_gen_requests(cache_manager, request_ids, past_kv_lens, draft_len: int = MTP_K):
-    """Allocate KV blocks for generation requests with the given past lengths."""
-    cache_manager.add_dummy_requests(
-        list(request_ids),
-        token_nums=[n + 1 for n in past_kv_lens],
-        is_gen=True,
-        prepare_resource=True,
-        max_num_draft_tokens=draft_len,
-    )
-
-
-def tensor_state(obj) -> dict:
-    """Clone every tensor and scalar attribute of ``obj`` (for state comparison)."""
-    state = {}
-    for name, value in vars(obj).items():
-        if isinstance(value, torch.Tensor):
-            state[name] = value.detach().clone()
-        elif isinstance(value, (bool, int, float, str, type(None))):
-            state[name] = value
-    return state
-
-
-def changed_keys(before: dict, after: dict) -> set:
-    """Attributes whose value differs between two ``tensor_state`` snapshots."""
-    keys = set()
-    for name in set(before) | set(after):
-        a, b = before.get(name, "<missing>"), after.get(name, "<missing>")
-        if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
-            if a.shape != b.shape or a.dtype != b.dtype or not torch.equal(a.cpu(), b.cpu()):
-                keys.add(name)
-        elif type(a) is not type(b) or a != b:
-            keys.add(name)
-    return keys
-
-
-def assert_same_state(native: dict, candidate: dict, keys, what: str) -> None:
-    """Assert that ``keys`` hold identical values in both snapshots."""
-    for name in sorted(keys):
-        a, b = native.get(name, "<missing>"), candidate.get(name, "<missing>")
-        if isinstance(a, torch.Tensor):
-            assert isinstance(b, torch.Tensor), f"{what}: {name} type differs"
-            assert a.shape == b.shape and a.dtype == b.dtype, f"{what}: {name} shape/dtype"
-            assert torch.equal(a.cpu(), b.cpu()), f"{what}: {name} differs"
-        else:
-            assert a == b, f"{what}: {name} differs ({a!r} vs {b!r})"
-
-
-def native_draft_step_tail(
-    md,
-    i,
-    batch_size,
-    num_contexts,
-    runtime_draft_len,
-    num_accepted_tokens,
-    position_ids,
-    gather_ids,
-):
-    """Literal copy of the native loop tail for draft step ``i`` (MTP Eagle, DSA)."""
-    from tensorrt_llm._torch.speculative.mtp import _select_mtp_position_ids
-
-    new_position_ids = None
-    if i + 1 < runtime_draft_len:
-        new_position_ids = _select_mtp_position_ids(position_ids, gather_ids) + 1
-    attn_metadata = md
-    if i == 0:
-        attn_metadata._seq_lens[:batch_size].fill_(1)
-        attn_metadata._seq_lens_cuda[:batch_size].fill_(1)
-        attn_metadata.on_update()
-        has_kv_cache = attn_metadata.kv_cache_manager is not None
-        if has_kv_cache and hasattr(attn_metadata, "host_request_types"):
-            attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1)
-        if has_kv_cache:
-            attn_metadata.num_contexts = 0
-        if hasattr(attn_metadata, "kv_lens_cuda"):
-            attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (
-                runtime_draft_len - num_accepted_tokens[num_contexts:]
-            )
-            attn_metadata.kv_lens_cuda[:num_contexts] += 1
-        # _prepare_flash_mla_generation_layout is a no-op without context rows.
-        if hasattr(attn_metadata, "kv_lens_cuda"):
-            attn_metadata.update_for_spec_dec()
-        attn_metadata.use_spec_decoding = False
-    else:
-        if hasattr(attn_metadata, "kv_lens_cuda"):
-            attn_metadata.kv_lens_cuda[:batch_size] += 1
-            attn_metadata.update_for_spec_dec()
-    return new_position_ids
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py
deleted file mode 100644
index 77905e327352..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_loop.py
+++ /dev/null
@@ -1,614 +0,0 @@
-# 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.
-"""draft-loop bookkeeping versus the native MTP one-model worker.
-
-Unit level: fused acceptance (strict + forced-acceptance override, including
-the draw counter) against ``_sample_and_accept_draft_tokens_base``; draft-0
-inputs against ``prepare_1st_drafter_inputs``; the row gather and the
-post-loop finalize against the native indexing.
-
-Worker level: the real ``MTPEagleWorker.forward`` of the target's adopted worker
-(``GlmMtpWorker``) on decode steps -- the native draft loop -- on real DSA metadata,
-with draft_loop off and on. A stub MTP layer (consuming the real embedding callable)
-records what every draft step sees; outputs, per-step inputs, per-step
-metadata and the final metadata state must be identical, eagerly and under
-CUDA-graph replay.
-"""
-
-import sys
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-sys.path.insert(0, str(Path(__file__).resolve().parent))
-
-from _draft_loop_fixtures import (  # noqa: E402
-    MTP_K,
-    add_gen_requests,
-    assert_same_state,
-    changed_keys,
-    make_dsa_cache_manager,
-    make_gen_dsa_metadata,
-    tensor_state,
-)
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
-    dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E402, E501
-    draft_loop as draft_loop_mod,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E402, E501
-    draft_step as draft_step_mod,
-)
-
-VOCAB = 1024
-HIDDEN = 6144
-
-
-def _make_worker(force: float):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
-        GlmMtpWorker,
-    )
-    from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-    from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
-
-    worker = GlmMtpWorker.adopt(MTPEagleWorker(MTPDecodingConfig(max_draft_len=MTP_K)))
-    worker.force_num_accepted_tokens = force
-    return worker
-
-
-def _spec_metadata(batch_size, draft_tokens, is_cuda_graph=False):
-    return SimpleNamespace(
-        runtime_draft_len=MTP_K,
-        draft_tokens=draft_tokens,
-        batch_indices_cuda=torch.arange(max(batch_size, 8), dtype=torch.int32, device="cuda"),
-        is_cuda_graph=is_cuda_graph,
-        is_all_greedy_sample=True,
-        use_rejection_sampling=False,
-        enable_penalty=False,
-        wants_advanced_draft_sampling=False,
-        all_rank_num_tokens=None,
-        subseq_all_rank_num_tokens=None,
-        spec_resource_manager=None,
-        max_num_requests=batch_size,
-        gather_ids=None,
-        context_prompt_lookahead_tokens=None,
-        is_ragged_verify=False,
-    )
-
-
-def _target_logits(targets: torch.Tensor) -> torch.Tensor:
-    """Float logits whose argmax per row is ``targets``."""
-    logits = torch.randn((targets.numel(), VOCAB), device="cuda")
-    logits[torch.arange(targets.numel(), device="cuda"), targets.long()] = 100.0
-    return logits
-
-
-def _random_case(gen, batch_size):
-    """Draft tokens and target tokens with a random leading-match length per row."""
-    draft = torch.randint(0, VOCAB, (batch_size, MTP_K), generator=gen, dtype=torch.int32)
-    target = torch.randint(0, VOCAB, (batch_size, MTP_K + 1), generator=gen, dtype=torch.int32)
-    for b in range(batch_size):
-        n = int(torch.randint(0, MTP_K + 1, (1,), generator=gen))
-        target[b, :n] = draft[b, :n]
-        if n < MTP_K and target[b, n] == draft[b, n]:
-            target[b, n] = (draft[b, n] + 1) % VOCAB
-    return draft.cuda().reshape(-1), target.cuda()
-
-
-@pytest.mark.parametrize("force", [0.0, 2.61, 0.5, 3.0, 5.0, 7.0])
-@pytest.mark.parametrize("batch_size", [1, 3, 8])
-def test_accept_matches_native(force, batch_size):
-    native, fused = _make_worker(force), _make_worker(force)
-    gen = torch.Generator().manual_seed(int(force * 100) + batch_size)
-    md_stub = SimpleNamespace(num_seqs=batch_size, num_contexts=0, num_ctx_tokens=0)
-    for call in range(40):
-        draft, target = _random_case(gen, batch_size)
-        logits = _target_logits(target.reshape(-1))
-        spec = _spec_metadata(batch_size, draft)
-        acc_n, num_n = native._sample_and_accept_draft_tokens_base(
-            logits, draft.reshape(batch_size, MTP_K), 0, batch_size, spec
-        )
-        ops = draft_loop_mod.DraftLoopOps(fused, md_stub, spec)
-        assert draft_loop_mod.DraftLoopOps.ineligibility(fused, md_stub, spec) is None
-        acc_w, num_w = ops.accept(logits)
-        assert acc_w.dtype == acc_n.dtype and torch.equal(acc_w, acc_n), call
-        assert num_w.dtype == num_n.dtype and torch.equal(num_w, num_n), call
-        if native._force_accept_rng_counter is not None:
-            assert torch.equal(native._force_accept_rng_counter, fused._force_accept_rng_counter)
-        # Draft-0 inputs and step-0 rows.
-        inputs = native.prepare_1st_drafter_inputs(
-            input_ids=torch.zeros((batch_size * (MTP_K + 1),), dtype=torch.int32, device="cuda"),
-            position_ids=None,
-            hidden_states=None,
-            accepted_tokens=acc_n,
-            attn_metadata=md_stub,
-            spec_metadata=spec,
-            draft_model=None,
-        )
-        assert torch.equal(ops.input_ids0, inputs["input_ids"])
-        assert ops.input_ids0.data_ptr() != acc_w.data_ptr()
-        rows = (spec.batch_indices_cuda[:batch_size] * (MTP_K + 1)).long() + num_n - 1 + 0
-        assert torch.equal(ops.gather_ids, rows)
-
-
-def test_gather_rows_and_finalize_match_native():
-    worker = _make_worker(0.0)
-    gen = torch.Generator(device="cuda").manual_seed(3)
-    for batch_size in (1, 2, 6):
-        hidden = torch.randn((batch_size * (MTP_K + 1), HIDDEN), generator=gen, device="cuda")
-        hidden = hidden.to(torch.bfloat16)
-        rows = torch.randint(0, hidden.shape[0], (batch_size,), generator=gen, device="cuda")
-        assert torch.equal(draft_loop_mod.gather_rows(hidden, rows), hidden[rows])
-
-        draft = [
-            torch.randint(0, VOCAB, (batch_size,), generator=gen, device="cuda", dtype=torch.int32)
-            for _ in range(MTP_K)
-        ]
-        accepted = torch.randint(
-            0, VOCAB, (batch_size, MTP_K + 1), generator=gen, device="cuda", dtype=torch.int32
-        )
-        num_acc = torch.randint(
-            1, MTP_K + 2, (batch_size,), generator=gen, device="cuda", dtype=torch.int32
-        )
-        spec = _spec_metadata(batch_size, None)
-        ops = draft_loop_mod.DraftLoopOps(
-            worker, SimpleNamespace(num_seqs=batch_size, num_contexts=0), spec
-        )
-        ops.accepted_tokens, ops.num_accepted_tokens = accepted, num_acc
-        stacked = torch.stack(draft, dim=1)
-        expected_new = worker._prepare_next_new_tokens(
-            accepted, stacked, spec.batch_indices_cuda, batch_size, num_acc
-        )
-        got = ops.finalize(draft)
-        assert torch.equal(got, stacked)
-        assert torch.equal(ops.next_new_tokens, expected_new)
-
-
-# ----------------------------------------------------------------------------
-# Worker level: the adopted worker's draft loop on decode steps
-# ----------------------------------------------------------------------------
-
-_STEP_FIELDS = (
-    "kv_lens_cuda",
-    "kv_lens_cuda_2d",
-    "scheduler_metadata_buffer",
-    "scheduler_metadata_buffer_full_next_n",
-    "kv_lens_expanded_cuda",
-    "scheduler_metadata_buffer_expanded",
-    "req_idx_per_token",
-    "slot_mapping_fp8",
-    "slot_mapping_scale",
-    "gen_kv_indptr",
-    "gen_cached_token_indptr",
-)
-
-
-class _Recorder:
-    """Per-step copies of what the MTP layer sees (static buffers under capture)."""
-
-    def __init__(self):
-        self.steps = []
-        self.static = None
-        self.index = 0
-
-    def record(self, md, input_ids, position_ids, hidden, emb):
-        items = {
-            "input_ids": input_ids,
-            "position_ids": position_ids,
-            "hidden": hidden,
-            "emb": emb,
-            "seq_lens_cuda": md.seq_lens_cuda,
-        }
-        items.update({f: getattr(md, f) for f in _STEP_FIELDS})
-        if torch.cuda.is_current_stream_capturing():
-            if self.static is None:
-                self.static = []
-            if self.index == len(self.static):
-                self.static.append({k: torch.empty_like(v) for k, v in items.items()})
-            for k, v in items.items():
-                self.static[self.index][k].copy_(v)
-            self.index += 1
-        else:
-            self.steps.append({k: v.clone() for k, v in items.items()})
-            self.steps[-1]["num_tokens"] = md.num_tokens
-
-
-class _StubHead:
-    def __call__(self, hidden, lm_head, attn_metadata, return_context_logits):
-        return hidden.float() @ lm_head.weight.float().t()
-
-
-class _StubMTPLayer(torch.nn.Module):
-    """Deterministic draft layer: embeds the ids with the given callable."""
-
-    def __init__(self, recorder):
-        super().__init__()
-        self.recorder = recorder
-        self.shared_head = _StubHead()
-
-    def forward(
-        self,
-        input_ids,
-        position_ids,
-        hidden_states,
-        embed_tokens,
-        attn_metadata,
-        all_rank_num_tokens=None,
-        spec_metadata=None,
-        **kwargs,
-    ):
-        emb = embed_tokens(input_ids)
-        self.recorder.record(attn_metadata, input_ids, position_ids, hidden_states, emb)
-        mix = emb.float() + 0.5 * hidden_states.float()
-        mix = mix + (position_ids.float() % 97).unsqueeze(-1) * 1e-2
-        return mix.to(hidden_states.dtype)
-
-
-def _draft_model(recorder, seed=11):
-    from tensorrt_llm._torch.modules.embedding import Embedding
-
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    embed = Embedding(VOCAB, HIDDEN, dtype=torch.bfloat16).cuda()
-    embed.weight.data = torch.randn((VOCAB, HIDDEN), generator=gen, device="cuda").to(
-        torch.bfloat16
-    )
-    lm_head = SimpleNamespace(
-        weight=torch.randn((VOCAB, HIDDEN), generator=gen, device="cuda").to(torch.bfloat16)
-    )
-    return SimpleNamespace(
-        mtp_layers=[_StubMTPLayer(recorder)], embed_tokens=embed, lm_head=lm_head
-    )
-
-
-def _forward_inputs(batch_size, past, gen):
-    draft, target = _random_case(gen, batch_size)
-    n = batch_size * (MTP_K + 1)
-    return {
-        "input_ids": torch.randint(0, VOCAB, (n,), generator=gen, dtype=torch.int32).cuda(),
-        "position_ids": torch.cat(
-            [torch.arange(p, p + MTP_K + 1, dtype=torch.int32) for p in past]
-        ).cuda()[None],
-        "hidden_states": torch.randn((n, HIDDEN), generator=gen).to(torch.bfloat16).cuda(),
-        "logits": _target_logits(target.reshape(-1)),
-        "draft_tokens": draft,
-    }
-
-
-class _ResourceManagerStub:
-    """Resource manager exposing only the separate draft KV cache manager."""
-
-    def __init__(self, draft_manager):
-        self.draft_manager = draft_manager
-
-    def get_resource_manager(self, kind):
-        from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
-
-        return self.draft_manager if kind == ResourceManagerType.DRAFT_KV_CACHE_MANAGER else None
-
-
-def _run_forward(worker, md, spec, draft_model, inputs):
-    resource_manager = None
-    if md.draft_kv_cache_manager is not None:
-        # One-model MTP with a separate draft KV cache, as the GEN runner sets up.
-        worker.use_separate_draft_kv_cache = True
-        resource_manager = _ResourceManagerStub(md.draft_kv_cache_manager)
-    with dispatch.step(dispatch.Phase.DECODE):
-        return worker.forward(
-            input_ids=inputs["input_ids"],
-            position_ids=inputs["position_ids"],
-            hidden_states=inputs["hidden_states"],
-            logits=inputs["logits"],
-            attn_metadata=md,
-            spec_metadata=spec,
-            draft_model=draft_model,
-            resource_manager=resource_manager,
-        )
-
-
-_OUTPUT_KEYS = ("new_tokens", "new_tokens_lens", "next_draft_tokens", "next_new_tokens")
-
-
-@pytest.mark.parametrize("force", [0.0, 2.61])
-@pytest.mark.parametrize("past", [[391061], [391061, 120000, 4100]])
-@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
-def test_worker_forward_matches_native(monkeypatch, force, past, separate_draft):
-    batch_size = len(past)
-    cache_manager, mapping, sparse = make_dsa_cache_manager(batch_size, max(past))
-    draft_manager = None
-    calls = {"accept": 0, "step": 0}
-    real_accept, real_step = draft_loop_mod.DraftLoopOps.accept, draft_step_mod.draft_step_update
-
-    def spy_accept(self, logits):
-        calls["accept"] += 1
-        return real_accept(self, logits)
-
-    def spy_step(*args, **kwargs):
-        calls["step"] += 1
-        return real_step(*args, **kwargs)
-
-    monkeypatch.setattr(draft_loop_mod.DraftLoopOps, "accept", spy_accept)
-    monkeypatch.setattr(draft_step_mod, "draft_step_update", spy_step)
-    try:
-        ids = list(range(batch_size))
-        add_gen_requests(cache_manager, ids, past)
-        if separate_draft:
-            draft_manager, _, _ = make_dsa_cache_manager(batch_size, max(past))
-            add_gen_requests(draft_manager, ids, past)
-        gen = torch.Generator().manual_seed(len(past) + int(force * 10))
-        for it in range(3):
-            inputs = _forward_inputs(batch_size, past, gen)
-            results = {}
-            for name, features in (
-                ("native", frozenset()),
-                ("draft_loop", frozenset({"draft_loop"})),
-            ):
-                monkeypatch.setattr(dispatch, "enabled_features", lambda f=features: f)
-                recorder = _Recorder()
-                worker = _make_worker(force)
-                md = make_gen_dsa_metadata(
-                    cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
-                )
-                base = tensor_state(md)
-                spec = _spec_metadata(batch_size, inputs["draft_tokens"].clone())
-                out = _run_forward(worker, md, spec, _draft_model(recorder), inputs)
-                torch.cuda.synchronize()
-                results[name] = (out, recorder.steps, base, tensor_state(md), worker)
-            assert calls["accept"] == it + 1 and calls["step"] == (it + 1) * MTP_K
-            (o_n, s_n, b_n, f_n, w_n), (o_w, s_w, b_w, f_w, w_w) = (
-                results["native"],
-                results["draft_loop"],
-            )
-            for k in _OUTPUT_KEYS:
-                assert o_n[k].dtype == o_w[k].dtype and torch.equal(o_n[k], o_w[k]), (it, k)
-            assert len(s_n) == len(s_w) == MTP_K
-            for i, (a, b) in enumerate(zip(s_n, s_w)):
-                for k in a:
-                    if isinstance(a[k], torch.Tensor):
-                        assert a[k].dtype == b[k].dtype and torch.equal(a[k], b[k]), (it, i, k)
-                    else:
-                        assert a[k] == b[k], (it, i, k)
-            keys = changed_keys(b_n, f_n) | changed_keys(b_w, f_w)
-            assert_same_state(f_n, f_w, keys, f"iteration {it} final metadata")
-            if w_n._force_accept_rng_counter is not None:
-                assert torch.equal(w_n._force_accept_rng_counter, w_w._force_accept_rng_counter)
-    finally:
-        cache_manager.shutdown()
-        if draft_manager is not None:
-            draft_manager.shutdown()
-
-
-def _capture_and_replay(
-    monkeypatch,
-    features,
-    cache_manager,
-    mapping,
-    sparse,
-    ids,
-    past,
-    replay_inputs,
-    between=None,
-    keep_graph=False,
-    draft_manager=None,
-):
-    """Capture one worker forward for ``features``, replay it per input set.
-
-    Returns per-replay clones of the outputs, the per-step records and the
-    draw counter. ``between`` runs after capture and before the replays.
-    """
-    batch_size = len(ids)
-    monkeypatch.setattr(dispatch, "enabled_features", lambda f=features: f)
-    recorder = _Recorder()
-    worker = _make_worker(2.61)
-    md = make_gen_dsa_metadata(
-        cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
-    )
-    kv_init = md.kv_lens_cuda[:batch_size].clone()
-    static = {k: v.clone() for k, v in replay_inputs[0].items()}
-    spec = _spec_metadata(batch_size, static["draft_tokens"], is_cuda_graph=True)
-    draft_model = _draft_model(recorder)
-    host_seq = md._seq_lens.clone()
-    # The forced-acceptance pool is built eagerly, never under capture.
-    worker._ensure_force_accept_rng_state(torch.device("cuda"))
-    for _ in range(2):  # eager warmup (forced acceptance is skipped here)
-        md.kv_lens_cuda[:batch_size].copy_(kv_init)
-        _run_forward(worker, md, spec, draft_model, static)
-        md._seq_lens.copy_(host_seq)
-        md.num_contexts = 0
-        md.on_update()
-    md.kv_lens_cuda[:batch_size].copy_(kv_init)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        out = _run_forward(worker, md, spec, draft_model, static)
-    md._seq_lens.copy_(host_seq)
-    if between is not None:
-        between()
-    results = []
-    for replay, inputs in enumerate(replay_inputs):
-        for k, v in inputs.items():
-            static[k].copy_(v)
-        md.kv_lens_cuda[:batch_size].copy_(kv_init)
-        graph.replay()
-        try:
-            torch.cuda.synchronize()
-        except (RuntimeError, torch.AcceleratorError) as err:
-            raise AssertionError(f"{sorted(features)} graph replay {replay} faulted") from err
-        results.append(
-            (
-                {k: out[k].clone() for k in _OUTPUT_KEYS},
-                [{k: v.clone() for k, v in step.items()} for step in recorder.static],
-                worker._force_accept_rng_counter.clone(),
-            )
-        )
-    if keep_graph:
-        return results, graph, md, static
-    del graph
-    torch.cuda.synchronize()
-    return results
-
-
-def _graph_setup():
-    past = [391061, 200000]
-    cache_manager, mapping, sparse = make_dsa_cache_manager(len(past), max(past) + 64)
-    ids = list(range(len(past)))
-    add_gen_requests(cache_manager, ids, past)
-    gen = torch.Generator().manual_seed(21)
-    replay_inputs = [_forward_inputs(len(past), past, gen) for _ in range(10)]
-    return cache_manager, mapping, sparse, ids, past, replay_inputs
-
-
-@pytest.mark.parametrize(
-    "scenario",
-    [
-        "native",
-        "native_then_self_check",
-        "draft_loop",
-        "native_then_new_metadata",
-        "native_then_native_eager",
-        "native_then_draft_loop_eager",
-        "native_then_native_capture",
-        "native_then_draft_loop_capture",
-    ],
-)
-def test_graph_isolation_diagnostics(monkeypatch, scenario):
-    """A captured graph replays cleanly after other work (fault localization)."""
-    cache_manager, mapping, sparse, ids, past, replay_inputs = _graph_setup()
-    try:
-        between = None
-        if scenario.startswith("native_then_") and scenario != "native_then_self_check":
-            eager_features = {
-                "native_then_native_eager": frozenset(),
-                "native_then_draft_loop_eager": frozenset({"draft_loop"}),
-                "native_then_native_capture": frozenset(),
-                "native_then_draft_loop_capture": frozenset({"draft_loop"}),
-            }.get(scenario)
-            second_graphs = []
-
-            def between():
-                if scenario.endswith("_capture"):
-                    # A second captured graph that stays alive during the replays.
-                    second_graphs.append(
-                        _capture_and_replay(
-                            monkeypatch,
-                            eager_features,
-                            cache_manager,
-                            mapping,
-                            sparse,
-                            ids,
-                            past,
-                            replay_inputs[:1],
-                            keep_graph=True,
-                        )
-                    )
-                    monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset())
-                    return
-                md2 = make_gen_dsa_metadata(cache_manager, mapping, sparse, ids, past)
-                if eager_features is None:
-                    return
-                monkeypatch.setattr(dispatch, "enabled_features", lambda: eager_features)
-                worker2 = _make_worker(2.61)
-                spec2 = _spec_metadata(len(ids), replay_inputs[0]["draft_tokens"].clone())
-                for _ in range(2):
-                    _run_forward(worker2, md2, spec2, _draft_model(_Recorder()), replay_inputs[0])
-                torch.cuda.synchronize()
-                monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset())
-
-        if scenario == "native_then_self_check":
-
-            def between():
-                draft_step_mod._SPLIT_KV_BY_SMS.clear()
-                assert (
-                    draft_step_mod.schedule_self_check(
-                        torch.cuda.get_device_properties(0).multi_processor_count,
-                        torch.device("cuda"),
-                    )
-                    is not None
-                )
-
-        features = frozenset({"draft_loop"}) if scenario == "draft_loop" else frozenset()
-        results = _capture_and_replay(
-            monkeypatch,
-            features,
-            cache_manager,
-            mapping,
-            sparse,
-            ids,
-            past,
-            replay_inputs[:3],
-            between,
-        )
-        assert len(results) == 3
-    finally:
-        cache_manager.shutdown()
-
-
-@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
-def test_worker_forward_cuda_graph(monkeypatch, separate_draft):
-    """Captured native and draft_loop forwards replay identically (forced acceptance 2.61).
-
-    Outputs, per-step records and draw counters are compared over 10 replays
-    with fresh inputs.
-    """
-    # The draft_loop schedule self-check runs eagerly before any capture, as in a
-    # serving warmup.
-    assert draft_step_mod.schedule_self_check(
-        torch.cuda.get_device_properties(0).multi_processor_count, torch.device("cuda")
-    )
-    cache_manager, mapping, sparse, ids, past, replay_inputs = _graph_setup()
-    draft_manager = None
-    try:
-        if separate_draft:
-            draft_manager, _, _ = make_dsa_cache_manager(len(past), max(past) + 64)
-            add_gen_requests(draft_manager, ids, past)
-        native = _capture_and_replay(
-            monkeypatch,
-            frozenset(),
-            cache_manager,
-            mapping,
-            sparse,
-            ids,
-            past,
-            replay_inputs,
-            draft_manager=draft_manager,
-        )
-        fused = _capture_and_replay(
-            monkeypatch,
-            frozenset({"draft_loop"}),
-            cache_manager,
-            mapping,
-            sparse,
-            ids,
-            past,
-            replay_inputs,
-            draft_manager=draft_manager,
-        )
-        for replay, ((o_n, r_n, c_n), (o_w, r_w, c_w)) in enumerate(zip(native, fused)):
-            for k in _OUTPUT_KEYS:
-                assert torch.equal(o_n[k], o_w[k]), (replay, k)
-            assert len(r_n) == len(r_w) == MTP_K
-            for i in range(MTP_K):
-                for k in r_n[i]:
-                    assert torch.equal(r_n[i][k], r_w[i][k]), (replay, i, k)
-            assert torch.equal(c_n, c_w), replay
-    finally:
-        cache_manager.shutdown()
-        if draft_manager is not None:
-            draft_manager.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py
deleted file mode 100644
index a1a6147b6ec2..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_draft_step.py
+++ /dev/null
@@ -1,314 +0,0 @@
-# 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.
-"""draft_loop fused draft-step metadata update versus the native draft-loop tail.
-
-The native reference is the literal per-step tail of
-``Eagle3OneModelWorker._forward_linear_draft_loop`` (position advance, step-0
-switch to one query token, KV-length update, ``update_for_spec_dec``) running
-on a real ``DSAtrtllmAttentionMetadata`` prepared from a real ``DSACacheManager``
-at the served GEN configuration. After every step every tensor and scalar
-attribute either path changed must be identical.
-"""
-
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-sys.path.insert(0, str(Path(__file__).resolve().parent))
-
-from _draft_loop_fixtures import (  # noqa: E402
-    MTP_K,
-    add_gen_requests,
-    assert_same_state,
-    changed_keys,
-    make_dsa_cache_manager,
-    make_gen_dsa_metadata,
-    native_draft_step_tail,
-    tensor_state,
-)
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop import (  # noqa: E402, E501
-    draft_step as draft_step_mod,
-)
-
-
-def _step(md, use_draft_loop, i, batch_size, num_accepted, position_ids, gather_ids):
-    pos_arg = position_ids if i + 1 < MTP_K else None
-    gather = gather_ids if i == 0 else slice(0, batch_size)
-    if use_draft_loop:
-        result = draft_step_mod.draft_step_update(
-            md, i, batch_size, 0, MTP_K, num_accepted, pos_arg, gather_ids if i == 0 else None
-        )
-        assert result is not None, draft_step_mod.step_ineligibility(md, batch_size, 0)
-        return result.position_ids
-    return native_draft_step_tail(md, i, batch_size, 0, MTP_K, num_accepted, position_ids, gather)
-
-
-def _enter_draft(md):
-    """The native loop's draft-KV entry: rebind to the draft buffers, refresh."""
-    if md.draft_kv_cache_manager is None:
-        return None
-    from tensorrt_llm._torch.speculative.interface import prepare_attn_metadata_for_draft_replay
-
-    saved = prepare_attn_metadata_for_draft_replay(md, md.draft_kv_cache_manager)
-    md.on_update_kv_lens()
-    return saved
-
-
-def _exit_draft(md, saved):
-    if saved is not None:
-        from tensorrt_llm._torch.speculative.interface import (
-            restore_attn_metadata_after_draft_replay,
-        )
-
-        restore_attn_metadata_after_draft_replay(md, saved)
-
-
-def _run_loop(md, use_draft_loop, num_accepted, position_ids, gather_ids):
-    batch_size = md.num_seqs
-    md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
-    saved = _enter_draft(md)
-    states = []
-    pos = position_ids
-    for i in range(MTP_K):
-        pos = _step(md, use_draft_loop, i, batch_size, num_accepted, pos, gather_ids)
-        torch.cuda.synchronize()
-        states.append((tensor_state(md), None if pos is None else pos.clone()))
-    _exit_draft(md, saved)
-    md.restore_from_spec_dec()
-    md.on_update()
-    return states, tensor_state(md)
-
-
-def _inputs(batch_size, past_kv_lens, accepted):
-    num_accepted = torch.tensor(accepted, dtype=torch.int32, device="cuda")
-    # Target verify positions: past_kv_len .. past_kv_len + K for each request.
-    pos = torch.cat(
-        [torch.arange(n, n + MTP_K + 1, dtype=torch.int32) for n in past_kv_lens]
-    ).cuda()
-    gather_ids = (
-        torch.arange(batch_size, dtype=torch.int64, device="cuda") * (MTP_K + 1)
-        + num_accepted.to(torch.int64)
-        - 1
-    )
-    return num_accepted, pos, gather_ids
-
-
-CASES = [
-    # (past KV lengths, accepted counts)
-    ([391061], [4]),
-    ([391061], [1]),
-    ([391061], [6]),
-    ([2100], [3]),
-    ([65536 - 3], [6]),
-    ([(1 << 20) - 64], [2]),
-    ([491061, 32768, 131072], [6, 1, 3]),
-    ([4096, 4097, 4098, 4099, 4100, 4101, 4102, 300000], [1, 2, 3, 4, 5, 6, 4, 3]),
-]
-
-
-@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
-@pytest.mark.parametrize("past_kv_lens, accepted", CASES)
-def test_step_update_matches_native(past_kv_lens, accepted, separate_draft):
-    batch_size = len(past_kv_lens)
-    cache_manager, mapping, sparse = make_dsa_cache_manager(batch_size, max(past_kv_lens))
-    draft_manager = None
-    try:
-        ids = list(range(batch_size))
-        add_gen_requests(cache_manager, ids, past_kv_lens)
-        if separate_draft:
-            draft_manager, _, _ = make_dsa_cache_manager(batch_size, max(past_kv_lens))
-            add_gen_requests(draft_manager, ids, past_kv_lens)
-        md_native = make_gen_dsa_metadata(
-            cache_manager, mapping, sparse, ids, past_kv_lens, draft_cache_manager=draft_manager
-        )
-        md_draft_loop = make_gen_dsa_metadata(
-            cache_manager, mapping, sparse, ids, past_kv_lens, draft_cache_manager=draft_manager
-        )
-        reason = draft_step_mod.step_ineligibility(md_draft_loop, batch_size, 0)
-        assert reason is None, reason
-        print(
-            f"expand_for_dsl={md_draft_loop.expand_for_dsl} factor={md_draft_loop.dsl_expand_factor} "
-            f"num_sms={md_draft_loop.num_sms} split_kv={draft_step_mod._SPLIT_KV_BY_SMS[md_draft_loop.num_sms]} "
-            f"width={md_draft_loop.gen_token_stride}"
-        )
-        base_native, base_draft_loop = tensor_state(md_native), tensor_state(md_draft_loop)
-        inputs = _inputs(batch_size, past_kv_lens, accepted)
-        steps_native, final_native = _run_loop(md_native, False, *inputs)
-        steps_draft_loop, final_draft_loop = _run_loop(md_draft_loop, True, *inputs)
-        for i, ((s_n, p_n), (s_w, p_w)) in enumerate(zip(steps_native, steps_draft_loop)):
-            keys = changed_keys(base_native, s_n) | changed_keys(base_draft_loop, s_w)
-            # The query lengths always change (K + 1 -> 1); the step buffers are
-            # compared whether or not they changed in this step.
-            assert "_seq_lens_cuda" in keys
-            assert_same_state(s_n, s_w, keys | set(_SNAPSHOT_FIELDS), f"step {i}")
-            if p_n is None:
-                assert p_w is None
-            else:
-                assert p_w is not None and p_w.dtype == p_n.dtype and torch.equal(p_w, p_n)
-        keys = changed_keys(base_native, final_native) | changed_keys(
-            base_draft_loop, final_draft_loop
-        )
-        assert_same_state(final_native, final_draft_loop, keys, "after restore")
-    finally:
-        cache_manager.shutdown()
-        if draft_manager is not None:
-            draft_manager.shutdown()
-
-
-def test_ineligible_short_context_is_untouched():
-    """Dense top-k (short context) keeps the native chain; draft_loop changes nothing."""
-    cache_manager, mapping, sparse = make_dsa_cache_manager(1, 1000)
-    try:
-        add_gen_requests(cache_manager, [0], [1000])
-        md = make_gen_dsa_metadata(cache_manager, mapping, sparse, [0], [1000])
-        assert draft_step_mod.step_ineligibility(md, 1, 0) == "dense top-k fallback active"
-        md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
-        before = tensor_state(md)
-        num_accepted, pos, gather_ids = _inputs(1, [1000], [3])
-        assert (
-            draft_step_mod.draft_step_update(md, 0, 1, 0, MTP_K, num_accepted, pos, gather_ids)
-            is None
-        )
-        torch.cuda.synchronize()
-        assert not changed_keys(before, tensor_state(md))
-        md.restore_from_spec_dec()
-    finally:
-        cache_manager.shutdown()
-
-
-def test_schedule_matches_deepgemm():
-    """The draft_loop schedule equals DeepGEMM's on random 1-D/2-D context lengths."""
-    from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata
-
-    num_sms = torch.cuda.get_device_properties(0).multi_processor_count
-    split_kv = draft_step_mod.schedule_self_check(num_sms, torch.device("cuda"))
-    assert split_kv is not None
-    gen = torch.Generator().manual_seed(80)
-    for trial in range(300):
-        num_requests = int(torch.randint(1, 17, (1,), generator=gen))
-        next_n = int(torch.randint(1, 9, (1,), generator=gen))
-        hi = [64, 4096, 1 << 20][trial % 3]
-        lens = torch.randint(0, hi + 15, (num_requests, 1), generator=gen, dtype=torch.int32)
-        lens = lens.expand(-1, next_n).contiguous().cuda()
-        native = get_paged_mqa_logits_metadata(lens, 64, num_sms)
-        ours = draft_step_mod.paged_mqa_schedule(lens, num_sms, split_kv)
-        assert torch.equal(native, ours), (lens.tolist(), next_n)
-
-
-_SNAPSHOT_FIELDS = (
-    "kv_lens_cuda",
-    "kv_lens_cuda_2d",
-    "scheduler_metadata_buffer",
-    "scheduler_metadata_buffer_full_next_n",
-    "kv_lens_expanded_cuda",
-    "scheduler_metadata_buffer_expanded",
-    "req_idx_per_token",
-    "slot_mapping_fp8",
-    "slot_mapping_scale",
-    "gen_kv_indptr",
-    "gen_cached_token_indptr",
-)
-
-
-def _capture_loop(md, use_draft_loop, kv_init, num_accepted, pos0, gather_ids):
-    """Capture prepare_for_spec_dec + 5 steps + restore, with in-graph snapshots."""
-    batch_size = md.num_seqs
-    snaps = [{f: torch.empty_like(getattr(md, f)) for f in _SNAPSHOT_FIELDS} for _ in range(MTP_K)]
-    seq_snaps = [torch.empty((batch_size,), dtype=torch.int32, device="cuda") for _ in range(MTP_K)]
-    pos_snaps = [
-        torch.empty((batch_size,), dtype=pos0.dtype, device="cuda") for _ in range(MTP_K - 1)
-    ]
-
-    def body():
-        md.kv_lens_cuda[:batch_size].copy_(kv_init)
-        md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
-        saved = _enter_draft(md)
-        pos = pos0
-        for i in range(MTP_K):
-            pos = _step(md, use_draft_loop, i, batch_size, num_accepted, pos, gather_ids)
-            for f in _SNAPSHOT_FIELDS:
-                snaps[i][f].copy_(getattr(md, f))
-            seq_snaps[i].copy_(md.seq_lens_cuda[:batch_size])
-            if pos is not None:
-                pos_snaps[i].copy_(pos)
-        _exit_draft(md, saved)
-        md.restore_from_spec_dec()
-        md.on_update()
-
-    # Warm up eagerly (JIT and the schedule self-check), then reset host state.
-    host_seq = md._seq_lens.clone()
-    num_contexts = md.num_contexts
-    body()
-    md._seq_lens.copy_(host_seq)
-    md.num_contexts = num_contexts
-    md.use_spec_decoding = True
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        body()
-    md._seq_lens.copy_(host_seq)
-    md.num_contexts = num_contexts
-    return graph, snaps, seq_snaps, pos_snaps
-
-
-@pytest.mark.parametrize("separate_draft", [False, True], ids=["shared_kv", "draft_kv"])
-def test_cuda_graph_replays_match_native(separate_draft):
-    past = [391061, 120000]
-    batch_size = len(past)
-    cache_manager, mapping, sparse = make_dsa_cache_manager(batch_size, max(past) + 64)
-    draft_manager = None
-    try:
-        ids = list(range(batch_size))
-        add_gen_requests(cache_manager, ids, past)
-        if separate_draft:
-            draft_manager, _, _ = make_dsa_cache_manager(batch_size, max(past) + 64)
-            add_gen_requests(draft_manager, ids, past)
-        md_native = make_gen_dsa_metadata(
-            cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
-        )
-        md_draft_loop = make_gen_dsa_metadata(
-            cache_manager, mapping, sparse, ids, past, draft_cache_manager=draft_manager
-        )
-        kv_init = md_native.kv_lens_cuda[:batch_size].clone()
-        num_accepted, pos0, gather_ids = _inputs(batch_size, past, [3, 5])
-        g_native = _capture_loop(md_native, False, kv_init, num_accepted, pos0, gather_ids)
-        g_draft_loop = _capture_loop(md_draft_loop, True, kv_init, num_accepted, pos0, gather_ids)
-        gen = torch.Generator().manual_seed(5)
-        for replay in range(12):
-            shift = int(torch.randint(-40, 40, (1,), generator=gen))
-            kv_init.copy_(torch.tensor([n + MTP_K + 1 + shift for n in past], dtype=torch.int32))
-            num_accepted.copy_(
-                torch.randint(1, MTP_K + 2, (batch_size,), generator=gen, dtype=torch.int32)
-            )
-            gather_ids.copy_(torch.arange(batch_size) * (MTP_K + 1) + num_accepted.cpu() - 1)
-            pos0.copy_(torch.cat([torch.arange(n + shift, n + shift + MTP_K + 1) for n in past]))
-            g_native[0].replay()
-            g_draft_loop[0].replay()
-            torch.cuda.synchronize()
-            for i in range(MTP_K):
-                for f in _SNAPSHOT_FIELDS:
-                    assert torch.equal(g_native[1][i][f], g_draft_loop[1][i][f]), (replay, i, f)
-                assert torch.equal(g_native[2][i], g_draft_loop[2][i]), (replay, i, "seq_lens")
-                if i < MTP_K - 1:
-                    assert torch.equal(g_native[3][i], g_draft_loop[3][i]), (replay, i, "pos")
-    finally:
-        cache_manager.shutdown()
-        if draft_manager is not None:
-            draft_manager.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py
deleted file mode 100644
index 913e918cb7e1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_loop/test_draft_loop_embedding.py
+++ /dev/null
@@ -1,177 +0,0 @@
-# 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.
-"""draft_loop draft embedding versus the native ``Embedding`` (bitwise)."""
-
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-sys.path.insert(0, str(Path(__file__).resolve().parent))
-
-from _draft_loop_fixtures import bits, load_checkpoint_tensor, real_token_ids  # noqa: E402
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop.embedding import (  # noqa: E402, E501
-    DraftEmbedding,
-    draft_embedding,
-    wrap_draft_embedding,
-)
-from tensorrt_llm._torch.modules.embedding import Embedding  # noqa: E402
-
-VOCAB = 154880
-HIDDEN = 6144
-
-
-def _native_embedding(weight: torch.Tensor) -> Embedding:
-    module = Embedding(weight.shape[0], weight.shape[1], dtype=weight.dtype).cuda()
-    module.weight.data = weight
-    return module
-
-
-def _edge_ids(vocab: int, dtype: torch.dtype) -> torch.Tensor:
-    info = torch.iinfo(dtype)
-    values = [0, 1, vocab - 1, -1, vocab, vocab + 7, info.min, info.max, 12345]
-    return torch.tensor(values, dtype=dtype, device="cuda")
-
-
-@pytest.fixture(scope="module")
-def synthetic_weight():
-    gen = torch.Generator(device="cuda").manual_seed(8)
-    weight = torch.randn((VOCAB, HIDDEN), dtype=torch.bfloat16, device="cuda", generator=gen)
-    # Keep special bit patterns in the table so a bit-copy is required.
-    weight[3, :4] = torch.tensor([-0.0, float("inf"), float("-inf"), float("nan")])
-    return weight
-
-
-@pytest.fixture(scope="module")
-def real_weight():
-    weight = load_checkpoint_tensor("model.embed_tokens.weight")
-    if weight is None:
-        pytest.skip("GLM-5.2 checkpoint not available")
-    assert weight.shape == (VOCAB, HIDDEN) and weight.dtype == torch.bfloat16
-    return weight
-
-
-@pytest.mark.parametrize("num_rows", [1, 2, 5, 6, 7, 64])
-@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
-def test_matches_native_synthetic(synthetic_weight, num_rows, id_dtype):
-    native = _native_embedding(synthetic_weight)
-    gen = torch.Generator(device="cuda").manual_seed(num_rows)
-    ids = torch.randint(0, VOCAB, (num_rows,), generator=gen, device="cuda", dtype=id_dtype)
-    ids[0] = 3
-    expected = native(ids)
-    got = draft_embedding(ids, synthetic_weight)
-    assert got.shape == expected.shape and got.dtype == expected.dtype
-    assert torch.equal(bits(got), bits(expected))
-
-
-@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
-def test_out_of_range_ids_give_zero_rows(synthetic_weight, id_dtype):
-    native = _native_embedding(synthetic_weight)
-    ids = _edge_ids(VOCAB, id_dtype)
-    expected = native(ids)
-    got = draft_embedding(ids, synthetic_weight)
-    assert torch.equal(bits(got), bits(expected))
-    # +0.0, not -0.0, for every invalid row.
-    invalid = (ids < 0) | (ids >= VOCAB)
-    assert torch.all(bits(got[invalid]) == 0)
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
-def test_other_dtypes(dtype):
-    weight = torch.randn((1000, 384), dtype=dtype, device="cuda")
-    native = _native_embedding(weight)
-    ids = torch.tensor([0, 999, 1000, -5, 17], dtype=torch.int32, device="cuda")
-    assert torch.equal(bits(draft_embedding(ids, weight)), bits(native(ids)))
-
-
-def test_real_weights_real_ids(real_weight):
-    native = _native_embedding(real_weight)
-    ids_host = real_token_ids(512)
-    assert ids_host is not None
-    for chunk in (1, 6):
-        for start in range(0, 512, 37):
-            ids = ids_host[start : start + chunk].cuda()
-            assert torch.equal(bits(draft_embedding(ids, real_weight)), bits(native(ids)))
-    # Every vocabulary row once.
-    all_ids = torch.arange(VOCAB, dtype=torch.int32, device="cuda")
-    assert torch.equal(bits(draft_embedding(all_ids, real_weight)), bits(native(all_ids)))
-
-
-def test_cuda_graph_replay(synthetic_weight):
-    """Replays read the static id buffer; each replay matches the native result."""
-    native = _native_embedding(synthetic_weight)
-    static_ids = torch.zeros((6,), dtype=torch.int32, device="cuda")
-    draft_embedding(static_ids, synthetic_weight)  # compile outside capture
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        static_out = draft_embedding(static_ids, synthetic_weight)
-    gen = torch.Generator(device="cuda").manual_seed(1)
-    for replay in range(20):
-        new_ids = torch.randint(
-            -3, VOCAB + 3, (6,), generator=gen, device="cuda", dtype=torch.int32
-        )
-        static_ids.copy_(new_ids)
-        graph.replay()
-        torch.cuda.synchronize()
-        assert torch.equal(bits(static_out), bits(native(new_ids))), f"replay {replay}"
-
-
-def test_wrapper_uses_fused_path_and_proxies_attributes(synthetic_weight, monkeypatch):
-    native = _native_embedding(synthetic_weight)
-    wrapper = wrap_draft_embedding(native)
-    assert wrap_draft_embedding(native) is wrapper
-    assert wrapper.ineligible_reason is None
-    assert wrapper.weight is native.weight
-    assert wrapper.num_embeddings == VOCAB
-
-    calls = []
-    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop.embedding as emb  # noqa: E501
-
-    real = emb.draft_embedding
-    monkeypatch.setattr(emb, "draft_embedding", lambda *a: calls.append(1) or real(*a))
-    ids = torch.tensor([5, 6], dtype=torch.int32, device="cuda")
-    assert torch.equal(bits(wrapper(ids)), bits(native(ids)))
-    assert calls == [1]
-
-
-def test_wrapper_native_fallbacks(synthetic_weight, monkeypatch):
-    native = _native_embedding(synthetic_weight)
-    wrapper = DraftEmbedding(native)
-    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_loop.embedding as emb  # noqa: E501
-
-    def _fail(*_):
-        raise AssertionError("fused path taken for an ineligible input")
-
-    monkeypatch.setattr(emb, "draft_embedding", _fail)
-    # 2-D ids and non-contiguous ids take the native module.
-    ids_2d = torch.tensor([[1, 2], [3, 4]], dtype=torch.int32, device="cuda")
-    assert torch.equal(bits(wrapper(ids_2d)), bits(native(ids_2d)))
-    ids_strided = torch.arange(10, dtype=torch.int32, device="cuda")[::2]
-    assert torch.equal(bits(wrapper(ids_strided)), bits(native(ids_strided)))
-    # A tensor-parallel module is ineligible as a whole.
-    native.tp_size = 2
-    assert DraftEmbedding(native).ineligible_reason == "embedding is tensor-parallel"
-    native.tp_size = 1
-
-
-def test_empty_ids(synthetic_weight):
-    ids = torch.empty((0,), dtype=torch.int32, device="cuda")
-    out = draft_embedding(ids, synthetic_weight)
-    assert out.shape == (0, HIDDEN)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py
deleted file mode 100644
index fe2499054729..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_nearties.py
+++ /dev/null
@@ -1,296 +0,0 @@
-# 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.
-"""BF16 near-tie tooling for the draft_tail decision-exactness study.
-
-* Native decisions: the 8-shard native oracle (``draft_tail_oracle``) on BF16 logits
-  computed by ``F.linear`` (nvjet) at a given GEMM row count M.
-* Exact-rounding decisions: native semantics (largest BF16 logit, lowest id
-  on ties) applied to logits computed in FP64 and rounded once to BF16.
-* Adversarial near-ties: real hidden rows changed in a few BF16 elements so
-  the FP64 logits of their two best tokens straddle one BF16 rounding
-  midpoint at a chosen distance (in FP32 ulps at the midpoint). There the
-  BF16 decision depends on the GEMM's FP32 accumulation order.
-"""
-
-import math
-import os
-import sys
-from typing import Dict, List, Sequence, Tuple
-
-import torch
-import torch.nn.functional as F
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import draft_tail_oracle as oracle  # noqa: E402
-
-# ---------------------------------------------------------------------------
-# BF16 grid helpers (finite values only)
-# ---------------------------------------------------------------------------
-
-
-def bf16_step(x: torch.Tensor, k) -> torch.Tensor:
-    """BF16 values ``k`` grid steps above ``x`` (BF16), along the real line;
-    -0 and +0 are one point. ``k`` is an int or a broadcastable int tensor."""
-    bits = x.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF
-    mag = bits & 0x7FFF
-    ordv = torch.where((bits & 0x8000) != 0, -mag, mag) + k
-    out = torch.where(ordv >= 0, ordv, 0x8000 | (-ordv))
-    out = torch.where(out >= 0x8000, out - 0x10000, out)
-    return out.to(torch.int16).view(torch.bfloat16)
-
-
-def bf16_round_exact(x: torch.Tensor) -> torch.Tensor:
-    """Round FP64 ``x`` to BF16 once (nearest, ties to even), avoiding the
-    double rounding of FP64 -> FP32 -> BF16."""
-    x = x.double()
-    trunc = (
-        (x.float().contiguous().view(torch.int32) & -65536).view(torch.float32).to(torch.bfloat16)
-    )  # toward zero, from FP32
-    # FP32 rounding of x may have crossed a BF16 grid point: step back.
-    lo = torch.where(
-        trunc.double().abs() > x.abs(), bf16_step(trunc, torch.where(x > 0, -1, 1)), trunc
-    )
-    hi = bf16_step(lo, torch.where(x >= 0, 1, -1))
-    lo_d, hi_d = lo.double(), hi.double()
-    mid = (lo_d + hi_d) / 2
-    away = x.abs() > mid.abs()
-    tie = x.abs() == mid.abs()
-    lo_even = (lo.view(torch.int16) & 1) == 0
-    pick_hi = away | (tie & ~lo_even)
-    return torch.where(pick_hi, hi, lo)
-
-
-def bf16_ulp(x: torch.Tensor) -> torch.Tensor:
-    """BF16 ulp at |x| (FP64), for normal values."""
-    e = torch.floor(torch.log2(x.abs().double().clamp_min(1e-30)))
-    return torch.pow(2.0, e - 7)
-
-
-def fp32_ulp(x: torch.Tensor) -> torch.Tensor:
-    e = torch.floor(torch.log2(x.abs().double().clamp_min(1e-30)))
-    return torch.pow(2.0, e - 23)
-
-
-# ---------------------------------------------------------------------------
-# Decisions
-# ---------------------------------------------------------------------------
-
-
-def native_bf16_logits(
-    hidden: torch.Tensor, shards: Sequence[torch.Tensor], m: int
-) -> List[torch.Tensor]:
-    """Per-shard native BF16 logits of ``hidden`` [N, H] computed with GEMMs
-    of ``m`` rows (the last chunk is padded with copies of its first row so
-    every GEMM has exactly ``m`` rows)."""
-    out = [[] for _ in shards]
-    n = hidden.shape[0]
-    for b in range(0, n, m):
-        chunk = hidden[b : b + m]
-        live = chunk.shape[0]
-        if live < m:
-            chunk = torch.cat([chunk, chunk[:1].expand(m - live, -1)])
-        for r, w in enumerate(shards):
-            out[r].append(F.linear(chunk, w)[:live])
-    return [torch.cat(o) for o in out]
-
-
-def native_tokens(shard_logits: Sequence[torch.Tensor], chunk: int = 512) -> torch.Tensor:
-    """8-shard native oracle (``_get_local_max_and_combined`` + gather +
-    ``_get_draft_tokens_from_gathered``) over many rows."""
-    n = shard_logits[0].shape[0]
-    toks = []
-    for b in range(0, n, chunk):
-        toks.append(oracle.native_tokens_from_logits([lg[b : b + chunk] for lg in shard_logits]))
-    return torch.cat(toks)
-
-
-def bf16_key(values_bf16: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
-    """draft_tail's packed key (ordered BF16 << 32 | ~id) as int64 (no NaN here)."""
-    bits = values_bf16.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF
-    ordered = torch.where((bits & 0x8000) != 0, 0x7FFF - (bits & 0x7FFF), bits | 0x8000)
-    # Same mapping as orderedBf16 for non-NaN values, up to a constant that
-    # does not change the order (-0 == +0 below).
-    ordered = torch.where((bits & 0x7FFF) == 0, torch.full_like(bits, 0x7FFF), ordered)
-    return (ordered << 32) | (0xFFFFFFFF - ids.to(torch.int64))
-
-
-def topk_fp64(
-    hidden: torch.Tensor, head: torch.Tensor, shard_logits: Sequence[torch.Tensor], k: int = 16
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Ids of the ``k`` largest native BF16 logits per row and their FP64
-    logits (sorted descending by FP64 value)."""
-    full = torch.cat(shard_logits, dim=-1).float()
-    ids = full.topk(k, dim=-1).indices
-    w = head[ids].double()  # [N, k, H]
-    vals = torch.einsum("nh,nkh->nk", hidden.double(), w)
-    order = vals.argsort(dim=-1, descending=True)
-    return ids.gather(1, order), vals.gather(1, order)
-
-
-def exact_tokens(ids: torch.Tensor, vals: torch.Tensor) -> torch.Tensor:
-    """Native semantics on once-rounded FP64 logits of the candidates."""
-    rounded = bf16_round_exact(vals)
-    keys = bf16_key(rounded, ids)
-    return ids.gather(1, keys.argmax(dim=1, keepdim=True)).squeeze(1).to(torch.int32)
-
-
-def near_tie_stats(ids: torch.Tensor, vals: torch.Tensor) -> Dict:
-    """Per-row near-tie measures from the FP64 candidates.
-
-    gap_bf16_ulps: (top1 - top2) / BF16 ulp at top1.
-    edge_fp32_ulps: for the candidates within one BF16 ulp of the top, the
-    smallest distance of any of them to a BF16 rounding midpoint, in FP32
-    ulps; small values mean a different FP32 accumulation order can change
-    that candidate's BF16 value.
-    """
-    top = vals[:, :1]
-    ulp_b = bf16_ulp(top)
-    gap = (vals[:, 0] - vals[:, 1]) / ulp_b[:, 0]
-    rounded = bf16_round_exact(vals).double()
-    # nearest midpoint: between rounded and its neighbour on x's side
-    side = torch.where(vals >= rounded, 1, -1)
-    neigh = bf16_step(bf16_round_exact(vals), side).double()
-    mid = (rounded + neigh) / 2
-    dist = (vals - mid).abs() / fp32_ulp(vals)
-    close = (top - vals) <= ulp_b
-    edge = torch.where(close, dist, torch.full_like(dist, float("inf")))
-    return {
-        "gap_bf16_ulps": gap,
-        "edge_fp32_ulps": edge.min(dim=1).values,
-    }
-
-
-def draft_tail_full_tokens(
-    runner, hidden: torch.Tensor, head: torch.Tensor, variant: int, chunk: int = 8
-) -> torch.Tensor:
-    """draft_tail fused-kernel decision over the whole vocabulary (one tp_size 1
-    shard; equals the 8-shard global decision, see test_draft_tail_draft_tail)."""
-    out = []
-    for b in range(0, hidden.shape[0], chunk):
-        h = hidden[b : b + chunk].contiguous()
-        keys = torch.zeros(h.shape[0], dtype=torch.int64, device="cuda")
-        runner.run_fused(h, head, 0, keys_out=keys, variant=variant)
-        out.append(oracle.keys_to_tokens(keys))
-    return torch.cat(out)
-
-
-def draft_tail_logits_tokens(
-    runner, shard_logits: Sequence[torch.Tensor], chunk: int = 8
-) -> torch.Tensor:
-    """draft_tail logits mode (native BF16 logits + draft_tail), 8 shards emulated."""
-    n = shard_logits[0].shape[0]
-    out = []
-    for b in range(0, n, chunk):
-        keys = []
-        for r, lg in enumerate(shard_logits):
-            k = torch.zeros(min(chunk, n - b), dtype=torch.int64, device="cuda")
-            runner.run_from_logits(lg[b : b + chunk].contiguous(), r * lg.shape[1], keys_out=k)
-            keys.append(k)
-        out.append(oracle.keys_to_tokens(oracle.global_key(keys)))
-    return torch.cat(out)
-
-
-# ---------------------------------------------------------------------------
-# Adversarial near-ties
-# ---------------------------------------------------------------------------
-
-_STEPS = (1, 2, 4, 8, 16, 32)
-
-
-def construct_straddle(
-    hidden: torch.Tensor,
-    head: torch.Tensor,
-    tok_hi: torch.Tensor,
-    tok_lo: torch.Tensor,
-    tau_fp32_ulps: torch.Tensor,
-    iters: int = 3000,
-) -> Dict[str, torch.Tensor]:
-    """Move rows ``hidden`` [B, H] (BF16) by BF16 grid steps of single
-    elements until the FP64 logit of ``tok_hi`` sits ``tau`` FP32 ulps above
-    a BF16 rounding midpoint m and that of ``tok_lo`` ``tau`` below it (tau
-    per row, times a random factor in [0.5, 1]); m is the midpoint nearest
-    to the mean of the two starting logits. Greedy: each iteration applies,
-    per row, the single (element, step) change that most reduces the
-    distance to the targets.
-    """
-    b, hdim = hidden.shape
-    h = hidden.clone()
-    w_hi = head[tok_hi].double()
-    w_lo = head[tok_lo].double()
-
-    def logits(hh):
-        hd = hh.double()
-        return (hd * w_hi).sum(-1), (hd * w_lo).sum(-1)
-
-    l_hi, l_lo = logits(h)
-    mean = (l_hi + l_lo) / 2
-    lo_grid = bf16_round_exact(mean)
-    lo_grid = torch.where(lo_grid.double() > mean, bf16_step(lo_grid, -1), lo_grid)
-    mid = (lo_grid.double() + bf16_step(lo_grid, 1).double()) / 2
-    gen = torch.Generator(device=hidden.device).manual_seed(7)
-    jitter = 0.5 + 0.5 * torch.rand(2, b, generator=gen, device=hidden.device, dtype=torch.float64)
-    ulp = fp32_ulp(mid)
-    t_hi = mid + tau_fp32_ulps.double() * ulp * jitter[0]
-    t_lo = mid - tau_fp32_ulps.double() * ulp * jitter[1]
-    rows = torch.arange(b, device=hidden.device)
-    for _ in range(iters):
-        r_hi, r_lo = t_hi - l_hi, t_lo - l_lo
-        cur = torch.maximum(r_hi.abs(), r_lo.abs())
-        done = cur <= 0.05 * tau_fp32_ulps.double() * ulp
-        if bool(done.all()):
-            break
-        best_cost = cur.clone()
-        best_idx = torch.full((b,), -1, dtype=torch.int64, device=hidden.device)
-        best_step = torch.zeros(b, dtype=torch.int64, device=hidden.device)
-        hd = h.double()
-        for s in _STEPS:
-            for sign in (1, -1):
-                moved = bf16_step(h, sign * s).double()
-                delta = moved - hd
-                c = torch.maximum(
-                    (r_hi[:, None] - delta * w_hi).abs(), (r_lo[:, None] - delta * w_lo).abs()
-                )
-                c = torch.where(torch.isfinite(moved), c, torch.full_like(c, float("inf")))
-                val, idx = c.min(dim=1)
-                better = val < best_cost
-                best_cost = torch.where(better, val, best_cost)
-                best_idx = torch.where(better, idx, best_idx)
-                best_step = torch.where(better, torch.full_like(best_step, sign * s), best_step)
-        move = (best_idx >= 0) & ~done
-        if not bool(move.any()):
-            break
-        sel = rows[move]
-        col = best_idx[move]
-        h[sel, col] = bf16_step(h[sel, col], best_step[move])
-        l_hi, l_lo = logits(h)
-    return {
-        "hidden": h,
-        "mid": mid,
-        "l_hi": l_hi,
-        "l_lo": l_lo,
-        "dist_hi_fp32_ulps": (l_hi - mid) / ulp,
-        "dist_lo_fp32_ulps": (l_lo - mid) / ulp,
-        "changed_elements": (h != hidden).sum(-1),
-    }
-
-
-def wilson_upper(k: int, n: int, z: float = 1.96) -> float:
-    if n == 0:
-        return float("nan")
-    p = k / n
-    den = 1 + z * z / n
-    c = p + z * z / (2 * n)
-    r = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
-    return (c + r) / den
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py
deleted file mode 100644
index a844bdfbd158..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/draft_tail_oracle.py
+++ /dev/null
@@ -1,148 +0,0 @@
-# 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.
-"""Native oracle and fixtures shared by the draft_tail tests and benchmarks.
-
-The oracle runs the native TensorRT-LLM code itself: ``F.linear`` (what
-``UnquantizedLinearMethod`` runs for the BF16 LM head) per vocab shard, then
-``SpecWorkerBase._get_local_max_and_combined`` per rank, the all-gather along
-the last dimension (a concatenation in rank order, as ``allgather(dim=-1)``
-returns) and the compiled ``SpecWorkerBase._get_draft_tokens_from_gathered``.
-On one GPU the eight ranks are emulated shard by shard.
-"""
-
-import json
-import os
-from types import SimpleNamespace
-from typing import List, Optional
-
-import torch
-import torch.nn.functional as F
-from utils.llm_data import llm_models_root
-
-CKPT = os.environ.get(
-    "GLM_MOE_DSA_DRAFT_TAIL_CKPT",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-VOCAB = 154880
-HIDDEN = 6144
-TP = 8
-SHARD = VOCAB // TP
-
-
-def checkpoint_available() -> bool:
-    return os.path.exists(os.path.join(CKPT, "model.safetensors.index.json"))
-
-
-def _tensor_file(name: str) -> str:
-    index = json.load(open(os.path.join(CKPT, "model.safetensors.index.json")))
-    return os.path.join(CKPT, index["weight_map"][name])
-
-
-def load_lm_head_rows(begin: int, end: int, device="cuda") -> torch.Tensor:
-    from safetensors import safe_open
-
-    with safe_open(_tensor_file("lm_head.weight"), framework="pt", device="cpu") as f:
-        return f.get_slice("lm_head.weight")[begin:end].to(device).contiguous()
-
-
-def load_lm_head_shards(tp: int = TP, device="cuda") -> List[torch.Tensor]:
-    rows = VOCAB // tp
-    full = load_lm_head_rows(0, VOCAB, device=device)
-    return [full[r * rows : (r + 1) * rows].contiguous() for r in range(tp)]
-
-
-def load_tensor(name: str, device="cuda") -> torch.Tensor:
-    from safetensors import safe_open
-
-    with safe_open(_tensor_file(name), framework="pt", device="cpu") as f:
-        return f.get_tensor(name).to(device)
-
-
-def mtp_norm_weight(device="cuda") -> torch.Tensor:
-    return load_tensor("model.layers.78.shared_head.norm.weight", device)
-
-
-def rms_normalize(x: torch.Tensor, gamma: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
-    xf = x.float()
-    y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
-    return (y * gamma.float()).to(torch.bfloat16)
-
-
-def synthetic_hidden(
-    rows: int,
-    gamma: torch.Tensor,
-    generator: torch.Generator,
-    peaked_from: Optional[torch.Tensor] = None,
-    peak_scale: float = 0.0,
-) -> torch.Tensor:
-    """Normed hidden rows as the MTP shared-head RMSNorm emits them.
-
-    ``peaked_from`` (rows of the LM head) adds a component along chosen token
-    directions so the logits have a clear winner, like real decoding states;
-    without it the logits are flat and near-ties are frequent.
-    """
-    x = torch.randn(rows, HIDDEN, generator=generator, device="cuda", dtype=torch.float32)
-    if peaked_from is not None:
-        d = peaked_from.float()
-        d = d / d.norm(dim=-1, keepdim=True)
-        x = x + peak_scale * d * x.norm(dim=-1, keepdim=True)
-    return rms_normalize(x, gamma)
-
-
-class _Rank(SimpleNamespace):
-    pass
-
-
-def native_combined(logits: torch.Tensor, tp_rank: int) -> torch.Tensor:
-    from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
-
-    return SpecWorkerBase._get_local_max_and_combined(
-        None, logits, mapping_lm_tp=_Rank(tp_rank=tp_rank)
-    )
-
-
-def native_from_gathered(gathered: torch.Tensor) -> torch.Tensor:
-    from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
-
-    return SpecWorkerBase._get_draft_tokens_from_gathered(None, gathered)
-
-
-def native_tokens_from_logits(shard_logits: List[torch.Tensor]) -> torch.Tensor:
-    """Native greedy TP draft tokens from per-rank BF16 vocab-shard logits."""
-    combined = [native_combined(lg, r) for r, lg in enumerate(shard_logits)]
-    return native_from_gathered(torch.cat(combined, dim=-1))
-
-
-def native_tokens(hidden: torch.Tensor, shards: List[torch.Tensor]) -> torch.Tensor:
-    return native_tokens_from_logits([F.linear(hidden, w) for w in shards])
-
-
-def keys_to_tokens(keys: torch.Tensor) -> torch.Tensor:
-    low = keys & 0xFFFFFFFF
-    return (0xFFFFFFFF - low).to(torch.int32)
-
-
-def global_key(keys_per_rank: List[torch.Tensor]) -> torch.Tensor:
-    """Unsigned max of packed keys (all keys are < 2**63 here)."""
-    stacked = torch.stack(keys_per_rank, dim=0)
-    assert bool((stacked >= 0).all()), "keys must fit in int64 as non-negative"
-    return stacked.max(dim=0).values
-
-
-def fp64_logits(
-    hidden: torch.Tensor, rows: torch.Tensor, weight_full_rows: torch.Tensor
-) -> torch.Tensor:
-    """Exact-ish logits of selected vocab rows (FP64 accumulation)."""
-    return hidden.double() @ weight_full_rows.double().t()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py
deleted file mode 100644
index da0a918a2417..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_draft_tail.py
+++ /dev/null
@@ -1,736 +0,0 @@
-# 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.
-"""Single-GPU correctness of the draft_tail against the native oracle.
-
-The eight TP ranks of GLM-5.2 are emulated shard by shard on one GPU: each
-shard runs the draft_tail kernel without exchange (``tp_size=1`` runner, explicit
-vocab offset) and returns its packed key; the global decision is the unsigned
-max of the eight keys, exactly what the exchange computes on every rank. The
-cross-rank exchange itself needs 8 GPUs and is not tested here.
-"""
-
-import json
-import os
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import draft_tail_nearties as nt  # noqa: E402
-import draft_tail_oracle as oracle  # noqa: E402
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
-
-NUM_VARIANTS = 15
-FIRST_MMA_VARIANT = 7
-NATIVE_ORDER_VARIANTS = (11, 12, 13, 14)
-
-
-def _family(variant):
-    """Kernels of one family accumulate a row in the same order: CUDA-core
-    (0-6), tensor-core permuted K (7-10), tensor-core native K order
-    (11-14)."""
-    if variant < FIRST_MMA_VARIANT:
-        return 0
-    return 11 if variant in NATIVE_ORDER_VARIANTS else FIRST_MMA_VARIANT
-
-
-def _draft_tail():
-    import tensorrt_llm  # noqa: F401  (loads the trtllm torch ops)
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail import (  # noqa: E501
-        draft_tail,
-    )
-
-    return draft_tail
-
-
-@pytest.fixture(scope="module")
-def runner():
-    dt = _draft_tail()
-    r = dt.DraftTailRunner(0, 1)
-    assert r.prepare()
-    return r
-
-
-@pytest.fixture(scope="module")
-def shards():
-    if not oracle.checkpoint_available():
-        pytest.skip("GLM-5.2 checkpoint not mounted")
-    return oracle.load_lm_head_shards()
-
-
-@pytest.fixture(scope="module")
-def gamma():
-    if not oracle.checkpoint_available():
-        pytest.skip("GLM-5.2 checkpoint not mounted")
-    return oracle.mtp_norm_weight()
-
-
-class _Head:
-    """Minimal stand-in exposing ``weight`` like LMHead."""
-
-    def __init__(self, weight):
-        self.weight = weight
-
-
-@pytest.fixture(scope="module")
-def full_head(shards):
-    """The whole vocabulary as one shard: with a tp_size 1 runner the draft_tail
-    decision over it is the global decision of the 8 shards (the packed-key
-    max is associative and every row's dot product is computed alike)."""
-    return _Head(torch.cat(shards, dim=0))
-
-
-def _full_keys(runner, hidden, head, variant):
-    keys = torch.zeros(hidden.shape[0], dtype=torch.int64, device="cuda")
-    runner.run_fused(hidden, head.weight, 0, keys_out=keys, variant=variant)
-    return keys
-
-
-def _e2e_near_tie_rows():
-    """Draft LM-head input rows of the real BF16 near-ties found end to end.
-
-    ``GLM_MOE_DSA_DRAFT_TAIL_NEAR_TIE_DIR`` is a directory with the native run's captured draft LM-head
-    inputs (with each draft-loop call's first row) and ``tie_analysis.json``.
-    """
-    d = os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_NEAR_TIE_DIR")
-    if not d:
-        return []
-    ties = [t for t in json.load(open(os.path.join(d, "tie_analysis.json"))) if t.get("near_tie")]
-    starts = torch.load(os.path.join(d, "native_eager_draft.pt"))["start_rows"]
-    rows = torch.load(os.path.join(d, "native_eager_hidden.pt"))["rows"]
-    return [
-        (f"e2e call {t['call']} step {t['step']}", rows[starts[t["call"]] + t["step"]].to("cuda"))
-        for t in ties
-    ]
-
-
-def _family_split_rows(runner, head, gamma, want=4):
-    """Constructed BF16 near-ties: seeded flat MTP-normed rows whose winner
-    the CUDA-core family (variant 4) and the tensor-core family (variant 10)
-    round differently. ``GLM_MOE_DSA_DRAFT_TAIL_FAMILY_SPLIT_BATCHES`` bounds the search."""
-    batches = int(os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_FAMILY_SPLIT_BATCHES", "512"))
-    gen = torch.Generator(device="cuda").manual_seed(2026)
-    found = []
-    for b in range(batches):
-        h = oracle.synthetic_hidden(8, gamma, gen)
-        a = oracle.keys_to_tokens(_full_keys(runner, h, head, 4))
-        c = oracle.keys_to_tokens(_full_keys(runner, h, head, 10))
-        for i in torch.nonzero(a != c).flatten().tolist():
-            found.append((f"family split, search batch {b} row {i}", h[i]))
-        if len(found) >= want:
-            break
-    return found[:want]
-
-
-@pytest.fixture(scope="module")
-def live_rows(runner, full_head, gamma):
-    """(label, hidden row) pairs used as live rows: real and constructed
-    BF16 near-ties on which the two kernel families disagree, plus random
-    flat and peaked rows."""
-    gen = torch.Generator(device="cuda").manual_seed(31)
-    rows = _e2e_near_tie_rows()
-    rows += _family_split_rows(runner, full_head, gamma)
-    rows.append(("flat", oracle.synthetic_hidden(1, gamma, gen)[0]))
-    peak = full_head.weight[123456:123457]
-    rows.append(("peaked", oracle.synthetic_hidden(1, gamma, gen, peak, 0.3)[0]))
-    return rows
-
-
-def _draft_tail_keys_fused(runner, hidden, shards, **kw):
-    keys = []
-    for r, w in enumerate(shards):
-        k = torch.zeros(hidden.shape[0], dtype=torch.int64, device="cuda")
-        runner.run_fused(hidden, w, r * w.shape[0], keys_out=k, **kw)
-        keys.append(k)
-    return keys
-
-
-def _draft_tail_keys_logits(runner, shard_logits):
-    keys = []
-    for r, lg in enumerate(shard_logits):
-        k = torch.zeros(lg.shape[0], dtype=torch.int64, device="cuda")
-        runner.run_from_logits(lg, r * lg.shape[1], keys_out=k)
-        keys.append(k)
-    return keys
-
-
-# ---------------------------------------------------------------------------
-# Selection semantics on given BF16 logits (logits kernel): exact equality.
-# ---------------------------------------------------------------------------
-
-
-def _random_logits(rows, gen, tp=oracle.TP, shard=oracle.SHARD):
-    return [
-        (torch.randn(rows, shard, generator=gen, device="cuda") * 3).to(torch.bfloat16)
-        for _ in range(tp)
-    ]
-
-
-@pytest.mark.parametrize("rows", [1, 2, 6, 8])
-def test_logits_selection_matches_native_random(runner, rows):
-    gen = torch.Generator(device="cuda").manual_seed(rows)
-    for _ in range(20):
-        lg = _random_logits(rows, gen)
-        ref = oracle.native_tokens_from_logits(lg)
-        got = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, lg)))
-        torch.testing.assert_close(got, ref, rtol=0, atol=0)
-
-
-def test_logits_selection_ties_and_specials(runner):
-    """Cross-shard and in-shard ties, -0/+0, infinities, NaN."""
-    gen = torch.Generator(device="cuda").manual_seed(7)
-    rows = 8
-    cases = []
-    lg = _random_logits(rows, gen)
-    # row 0: same max value in shards 5 and 2 -> lower global id (shard 2)
-    for r, c in ((5, 100), (2, 19000)):
-        lg[r][0, c] = 1000.0
-    # row 1: tie inside one shard at two columns -> lower column
-    lg[3][1, 777] = 1000.0
-    lg[3][1, 12] = 1000.0
-    # row 2: every logit equal -> global id 0
-    for r in range(oracle.TP):
-        lg[r][2].fill_(-1.5)
-    # row 3: max is +0 vs -0 at different ids -> treated equal -> lower id
-    for r in range(oracle.TP):
-        lg[r][3].fill_(-5.0)
-    lg[6][3, 3] = 0.0
-    lg[1][3, 9] = -0.0
-    # row 4: +inf in shard 4, large finite elsewhere
-    lg[4][4, 55] = float("inf")
-    lg[0][4, 1] = 3e38
-    # row 5: all -inf except one -inf-adjacent value
-    for r in range(oracle.TP):
-        lg[r][5].fill_(float("-inf"))
-    lg[7][5, 19359] = -3e38
-    # row 6: tie across the last column of shard 0 and first of shard 1
-    lg[0][6, oracle.SHARD - 1] = 500.0
-    lg[1][6, 0] = 500.0
-    # row 7: NaN in shard 3 (native: NaN ranks above every number)
-    lg[3][7, 4242] = float("nan")
-    cases.append(lg)
-    for lg in cases:
-        ref = oracle.native_tokens_from_logits(lg)
-        got = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, lg)))
-        torch.testing.assert_close(got, ref, rtol=0, atol=0)
-    assert got[0].item() == 2 * oracle.SHARD + 19000
-    assert got[1].item() == 3 * oracle.SHARD + 12
-    assert got[2].item() == 0
-
-
-def test_logits_selection_padded_vocab(runner):
-    """Vocab not divisible by TP: the last shard carries zero padding rows.
-
-    The native path includes the padding columns of the last shard in its
-    local max (they produce logit 0); draft_tail must make the same decision.
-    """
-    gen = torch.Generator(device="cuda").manual_seed(11)
-    tp, shard = 8, 1000
-    vocab = tp * shard - 37
-    for rows in (1, 3):
-        lg = [
-            (torch.randn(rows, shard, generator=gen, device="cuda") - 4).to(torch.bfloat16)
-            for _ in range(tp)
-        ]
-        lg[-1][:, vocab - (tp - 1) * shard :] = 0.0  # zero weight rows
-        ref = oracle.native_tokens_from_logits(lg)
-        got = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, lg)))
-        torch.testing.assert_close(got, ref, rtol=0, atol=0)
-
-
-# ---------------------------------------------------------------------------
-# Fused GEMV path on the real GLM-5.2 LM head.
-# ---------------------------------------------------------------------------
-
-
-def _assert_near_ties(reports):
-    """A token mismatch is acceptable only when the two winners' exact logits
-    are within one BF16 ulp of the maximum: both kernels then round FP32 sums
-    of different accumulation order across the same BF16 boundary."""
-    for rep in reports:
-        top = max(abs(v) for v in rep["fp64"])
-        ulp = 2.0 ** (int(torch.tensor(top).log2().floor().item()) - 7)
-        assert abs(rep["fp64_gap"]) <= ulp, rep
-
-
-def _decision_report(hidden, shards, got, ref):
-    """For each mismatch, the FP64 logit gap between the two winners."""
-    full = torch.cat(shards, dim=0)
-    out = []
-    for i in torch.nonzero(got != ref).flatten().tolist():
-        ids = torch.tensor([ref[i].item(), got[i].item()], device="cuda")
-        exact = hidden[i : i + 1].double() @ full[ids].double().t()
-        native_bf16 = F.linear(hidden[i : i + 1], full[ids])
-        out.append(
-            {
-                "row": i,
-                "native": ref[i].item(),
-                "draft_tail": got[i].item(),
-                "fp64_gap": (exact[0, 0] - exact[0, 1]).item(),
-                "fp64": exact[0].tolist(),
-                "native_bf16": native_bf16[0].float().tolist(),
-            }
-        )
-    return out
-
-
-@pytest.mark.parametrize("peak", [0.0, 0.3, 1.0])
-def test_fused_matches_native_real_lm_head(runner, shards, gamma, peak):
-    """Exact token equality vs native on the real head; any mismatch must be
-    a BF16 rounding near-tie (FP64 gap below one BF16 ulp of the max)."""
-    gen = torch.Generator(device="cuda").manual_seed(int(peak * 10) + 1)
-    full_rows = torch.cat(shards, dim=0)
-    rows_total, mismatches, reports = 0, 0, []
-    for it in range(16):
-        rows = 8
-        target = torch.randint(0, oracle.VOCAB, (rows,), generator=gen, device="cuda")
-        hidden = oracle.synthetic_hidden(
-            rows, gamma, gen, peaked_from=full_rows[target] if peak > 0 else None, peak_scale=peak
-        )
-        ref = oracle.native_tokens(hidden, shards)
-        got = oracle.keys_to_tokens(
-            oracle.global_key(_draft_tail_keys_fused(runner, hidden, shards))
-        )
-        rows_total += rows
-        bad = got != ref
-        if bool(bad.any()):
-            mismatches += int(bad.sum())
-            reports += _decision_report(hidden, shards, got, ref)
-    print(f"[draft_tail] peak={peak}: {mismatches}/{rows_total} token mismatches; {reports}")
-    _assert_near_ties(reports)
-
-
-@pytest.mark.parametrize("variant", list(range(NUM_VARIANTS)))
-def test_fused_variants_bitwise_identical(runner, shards, gamma, variant):
-    """Within a kernel family (CUDA-core 0-6, tensor-core 7-10, tensor-core
-    native order 11-14) every shape, CTA count and L2 keep hint accumulates
-    a row in the same order, so the keys are identical."""
-    base_variant = _family(variant)
-    gen = torch.Generator(device="cuda").manual_seed(3)
-    for rows in (1, 2, 5, 8):
-        hidden = oracle.synthetic_hidden(rows, gamma, gen)
-        base = _draft_tail_keys_fused(runner, hidden, shards[:2], variant=base_variant)
-        for num_ctas, keep in ((0, 0), (37, 0), (300, 0), (0, 500), (304, 1000)):
-            got = _draft_tail_keys_fused(
-                runner, hidden, shards[:2], variant=variant, num_ctas=num_ctas, keep_permille=keep
-            )
-            for a, b in zip(base, got):
-                torch.testing.assert_close(a, b, rtol=0, atol=0)
-
-
-@pytest.mark.parametrize("variant", [0, FIRST_MMA_VARIANT, 11, 13])
-@pytest.mark.parametrize("rows", [1, 3, 6, 8])
-def test_fused_families_match_native(runner, shards, gamma, variant, rows):
-    """Every kernel family makes the native decision on the real head for
-    every row count (mismatches only at BF16 near-ties; none at all for the
-    native-order family)."""
-    gen = torch.Generator(device="cuda").manual_seed(100 + rows)
-    reports = []
-    for it in range(8):
-        peak = (0.0, 0.3, 1.0)[it % 3]
-        hidden = oracle.synthetic_hidden(
-            rows, gamma, gen, torch.cat(shards[:1])[it * 8 : it * 8 + rows] if peak else None, peak
-        )
-        ref = oracle.native_tokens(hidden, shards)
-        got = oracle.keys_to_tokens(
-            oracle.global_key(_draft_tail_keys_fused(runner, hidden, shards, variant=variant))
-        )
-        reports += _decision_report(hidden, shards, got, ref)
-    print(f"[draft_tail] variant={variant} rows={rows}: mismatches {reports}")
-    _assert_near_ties(reports)
-    if variant in NATIVE_ORDER_VARIANTS:
-        assert not reports, reports
-
-
-@pytest.mark.parametrize("variant", [4, 10, 11, 13])
-def test_fused_family_rows_independent(runner, shards, gamma, variant):
-    """Within one kernel, row r's key does not depend on the other rows of
-    the call or on the row count."""
-    gen = torch.Generator(device="cuda").manual_seed(5)
-    live = oracle.synthetic_hidden(3, gamma, gen)
-    w = shards[4]
-    single = [
-        _draft_tail_keys_fused(runner, live[i : i + 1], [w], variant=variant)[0] for i in range(3)
-    ]
-    for rows in range(3, 9):
-        padded = torch.zeros(rows, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda")
-        padded[:3] = live
-        if rows > 3:
-            padded[3:] = oracle.synthetic_hidden(rows - 3, gamma, gen) * 7
-        keys = _draft_tail_keys_fused(runner, padded, [w], variant=variant)[0]
-        for i in range(3):
-            assert keys[i].item() == single[i].item()
-
-
-def test_select_live_row_key_independent_of_row_count(runner, full_head, gamma, live_rows):
-    """The shipped dispatch (``runner.select``, as the model adapter calls
-    it) gives a live row the same packed key, hence the same draft token,
-    for every row count R = 1..8 and position, so CUDA-graph padding and
-    batch composition cannot change a live row's decision. The live rows
-    include real and constructed BF16 near-ties on which the two kernel
-    families round differently: a dispatch that picked the family by row
-    count would fail here."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
-        DraftTailOptions,
-    )
-
-    assert runner.options == DraftTailOptions()  # the shipped defaults
-    gen = torch.Generator(device="cuda").manual_seed(17)
-    pad = torch.cat(
-        [
-            oracle.synthetic_hidden(4, gamma, gen) * 7,
-            torch.zeros(4, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda"),
-        ]
-    )
-    split = 0
-    for label, row in live_rows:
-        v4 = int(_full_keys(runner, row[None], full_head, 4)[0])
-        v10 = int(_full_keys(runner, row[None], full_head, 10)[0])
-        vdef = int(_full_keys(runner, row[None], full_head, DraftTailOptions().variant)[0])
-        logits = F.linear(row[None], full_head.weight)
-        native_key = int(
-            nt.bf16_key(logits, torch.arange(logits.shape[1], device="cuda")[None]).max()
-        )
-        split += int(v4 != v10)
-        seen = set()
-        for rows in range(1, 9):
-            for pos in sorted({0, rows - 1}):
-                batch = pad[:rows].clone()
-                batch[pos] = row
-                keys = torch.zeros(rows, dtype=torch.int64, device="cuda")
-                tokens = runner.select(batch, full_head, keys_out=keys)
-                assert int(tokens[pos]) == int(oracle.keys_to_tokens(keys[pos : pos + 1])[0])
-                seen.add(int(keys[pos]))
-        print(
-            f"[draft_tail-rowcount] {label}: select keys over R=1..8 {seen}; "
-            f"v4 token {int(oracle.keys_to_tokens(torch.tensor([v4])))} "
-            f"v10 token {int(oracle.keys_to_tokens(torch.tensor([v10])))}"
-        )
-        assert len(seen) == 1, (label, seen)
-        # The shipped kernel for every R is the default variant, and its key
-        # (winning BF16 value and id) is native's.
-        assert seen == {vdef}, (label, seen, vdef)
-        assert vdef == native_key, (label, vdef, native_key)
-    print(
-        f"[draft_tail-rowcount] live rows {len(live_rows)}, rows where the two "
-        f"families disagree: {split}"
-    )
-
-
-def test_row_count_characterization(runner, full_head, shards, gamma, live_rows):
-    """For the record (no assertion on native): the native decision on each
-    live row when it shares the LM-head GEMM with M - 1 other rows (8-shard
-    native oracle; cuBLAS picks its kernel by M), next to both draft_tail kernel
-    families and the FP64 logits of every candidate."""
-    gen = torch.Generator(device="cuda").manual_seed(23)
-    pad = oracle.synthetic_hidden(64, gamma, gen)
-    full = full_head.weight
-    for label, row in live_rows:
-        native = {}
-        bf16 = {}
-        for m in (1, 2, 3, 4, 5, 6, 7, 8, 16, 32, 64):
-            batch = pad[:m].clone()
-            batch[0] = row
-            logits = [F.linear(batch, w) for w in shards]
-            tok = int(oracle.native_tokens_from_logits(logits)[0])
-            native[m] = tok
-            bf16[m] = logits[tok // oracle.SHARD][0, tok % oracle.SHARD].item()
-        fam = {
-            v: int(oracle.keys_to_tokens(_full_keys(runner, row[None], full_head, v))[0])
-            for v in (4, 10)
-        }
-        cands = sorted(set(native.values()) | set(fam.values()))
-        exact = (row[None].double() @ full[cands].double().t())[0].tolist()
-        print(
-            "[draft_tail-rowcount-native] "
-            + json.dumps(
-                {
-                    "row": label,
-                    "native_token_by_M": native,
-                    "native_bf16_winner_logit_by_M": bf16,
-                    "draft_tail_v4": fam[4],
-                    "draft_tail_v10": fam[10],
-                    "fp64_logits": dict(zip(cands, exact)),
-                }
-            )
-        )
-
-
-def test_fused_cuda_graph_replay(runner, shards, gamma):
-    """Scratch keys and the arrival counter are reset by every launch:
-    repeated replays with changing inputs keep matching native."""
-    gen = torch.Generator(device="cuda").manual_seed(9)
-    rows = 2
-    static_h = torch.zeros(rows, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda")
-    static_keys = [torch.zeros(rows, dtype=torch.int64, device="cuda") for _ in shards]
-
-    def step():
-        for r, w in enumerate(shards):
-            runner.run_fused(static_h, w, r * w.shape[0], keys_out=static_keys[r])
-
-    s = torch.cuda.Stream()
-    s.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(s):
-        step()
-    torch.cuda.current_stream().wait_stream(s)
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        step()
-    for it in range(50):
-        static_h.copy_(oracle.synthetic_hidden(rows, gamma, gen))
-        g.replay()
-        torch.cuda.synchronize()
-        got = oracle.keys_to_tokens(oracle.global_key(static_keys))
-        eager = oracle.keys_to_tokens(
-            oracle.global_key(_draft_tail_keys_fused(runner, static_h, shards))
-        )
-        torch.testing.assert_close(got, eager, rtol=0, atol=0)
-        ref = oracle.native_tokens(static_h, shards)
-        _assert_near_ties(_decision_report(static_h, shards, got, ref))
-    assert runner.scratch[:8].abs().sum().item() == 0
-    assert runner.scratch[8].item() & 0xFFFFFFFF == 0
-
-
-def test_runner_tp1_tokens_match_native_full_vocab(runner, shards, gamma):
-    """tp_size == 1 (no exchange): one shard is the whole local vocab; the
-    native path is a plain argmax over the logits."""
-    gen = torch.Generator(device="cuda").manual_seed(13)
-    w = shards[1]
-    for rows in (1, 4, 8):
-        h = oracle.synthetic_hidden(rows, gamma, gen, w[:rows], 1.0)
-        tokens = runner.run_fused(h, w, 0)
-        ref = torch.argmax(F.linear(h, w), dim=-1).to(torch.int32)
-        torch.testing.assert_close(tokens, ref, rtol=0, atol=0)
-
-
-@pytest.mark.skipif(
-    not os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE"),
-    reason="GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE (captured draft LM-head inputs) not set",
-)
-@pytest.mark.parametrize("variant", [0, FIRST_MMA_VARIANT, 11, 13])
-def test_real_hidden_states_exact_tokens(runner, shards, variant):
-    """Real MTP draft LM-head inputs captured from a native GLM-5.2 run:
-    token-exact vs native for every row, through
-    the fused kernel (both families) and the logits kernel; any mismatch must
-    be a BF16 near-tie."""
-    rows = torch.load(os.environ["GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE"])["rows"].to("cuda")
-    assert rows.dtype == torch.bfloat16 and rows.shape[1] == oracle.HIDDEN
-    total = fused_bad = logits_bad = 0
-    reports = []
-    for i in range(0, rows.shape[0], 8):
-        h = rows[i : i + 8].contiguous()
-        logits = [F.linear(h, w) for w in shards]
-        ref = oracle.native_tokens_from_logits(logits)
-        got_f = oracle.keys_to_tokens(
-            oracle.global_key(_draft_tail_keys_fused(runner, h, shards, variant=variant))
-        )
-        got_l = oracle.keys_to_tokens(oracle.global_key(_draft_tail_keys_logits(runner, logits)))
-        # The logits kernel consumes the native BF16 logits: always exact.
-        torch.testing.assert_close(got_l, ref, rtol=0, atol=0)
-        total += h.shape[0]
-        fused_bad += int((got_f != ref).sum())
-        logits_bad += int((got_l != ref).sum())
-        reports += _decision_report(h, shards, got_f, ref)
-    print(
-        f"[draft_tail] real hidden rows={total} variant={variant}: fused "
-        f"mismatches={fused_bad}, logits mismatches={logits_bad}; "
-        f"{reports}"
-    )
-    _assert_near_ties(reports)
-    if variant in NATIVE_ORDER_VARIANTS:
-        assert fused_bad == 0
-
-
-def _exactness_rows(full_head, gamma):
-    """Rows for the bit-level comparison: real draft LM-head inputs
-    (``GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE``, first ``GLM_MOE_DSA_DRAFT_TAIL_EXACT_ROWS``),
-    knife-edge rows built from the real rows with the smallest top-2 gaps (their two best logits
-    straddle a BF16 rounding midpoint within a fraction of an FP32 ulp, so
-    any difference in their FP32 sums changes the BF16 values), and
-    synthetic flat / peaked MTP-normed rows."""
-    gen = torch.Generator(device="cuda").manual_seed(41)
-    rows = [oracle.synthetic_hidden(16, gamma, gen)]
-    rows.append(oracle.synthetic_hidden(16, gamma, gen, full_head.weight[1000:1016], 0.5))
-    path = os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_HIDDEN_FIXTURE")
-    if path:
-        n = int(os.environ.get("GLM_MOE_DSA_DRAFT_TAIL_EXACT_ROWS", "256"))
-        real = torch.load(path)["rows"][:n].to("cuda")
-        rows.append(real)
-        shards = [
-            full_head.weight[r * oracle.SHARD : (r + 1) * oracle.SHARD] for r in range(oracle.TP)
-        ]
-        logits = nt.native_bf16_logits(real, shards, 8)
-        ids, vals = nt.topk_fp64(real, full_head.weight, logits)
-        gaps = (vals[:, 0] - vals[:, 1]) / nt.bf16_ulp(vals[:, 0])
-        base = gaps.argsort()[:16]
-        lo = torch.minimum(ids[base, 0], ids[base, 1])
-        hi = torch.maximum(ids[base, 0], ids[base, 1])
-        for up, down in ((lo, hi), (hi, lo)):
-            res = nt.construct_straddle(
-                real[base], full_head.weight, up, down, torch.full((16,), 0.5, device="cuda")
-            )
-            rows.append(res["hidden"])
-    return torch.cat(rows)
-
-
-@pytest.mark.parametrize("variant", list(NATIVE_ORDER_VARIANTS))
-def test_native_order_logits_bitwise_equal_native(runner, shards, full_head, gamma, variant):
-    """The native-order fused variants compute every BF16 logit bit for bit
-    like the native LM-head GEMM (``F.linear``), for every row count, so
-    their draft tokens are native ones. Variant 10 (permuted K order) is
-    measured on the same rows for comparison."""
-    rows = _exactness_rows(full_head, gamma)
-    diffs = {variant: 0, 10: 0}
-    total = 0
-    for i in range(0, rows.shape[0], 8):
-        for r in (1, 3, 8):
-            h = rows[i : i + r].contiguous()
-            for s, w in enumerate(shards):
-                ref = F.linear(h, w)
-                for v in diffs:
-                    out = torch.empty_like(ref)
-                    runner.run_fused(h, w, s * w.shape[0], variant=v, logits_out=out)
-                    diffs[v] += int((out.view(torch.int16) != ref.view(torch.int16)).sum())
-                total += ref.numel()
-    print(
-        f"[draft_tail-exact] rows={rows.shape[0]} logits compared={total} "
-        f"BF16 logits != native: {diffs}"
-    )
-    assert diffs[variant] == 0
-
-
-def _loaded_cublas_versions():
-    """cuBLAS / cuBLASLt libraries mapped into this process and their
-    versions (the native LM-head GEMM's kernel is chosen by these)."""
-    import ctypes
-
-    paths = set()
-    with open("/proc/self/maps") as f:
-        for line in f:
-            p = line.split()[-1]
-            if "libcublas" in os.path.basename(p) and ".so" in p:
-                paths.add(p)
-    out = {}
-    for p in sorted(paths):
-        lib = ctypes.CDLL(p)
-        fn = "cublasLtGetProperty" if "cublasLt" in p else "cublasGetProperty"
-        try:
-            ver = []
-            for prop in range(3):  # MAJOR_VERSION, MINOR_VERSION, PATCH_LEVEL
-                v = ctypes.c_int(-1)
-                getattr(lib, fn)(prop, ctypes.byref(v))
-                ver.append(v.value)
-            out[p] = ".".join(str(v) for v in ver)
-        except AttributeError:
-            out[p] = None
-    return out
-
-
-def test_native_lm_head_kernel_record(shards):
-    """For the record (no assertion): the GEMM kernel this image's cuBLAS
-    picks for the native draft LM head (``F.linear``, one vocab shard) at
-    M = 1, 3, 8, and the library versions. The native-order bitwise claim
-    is tied to this kernel; a different name on another image is the cue to
-    re-run ``test_native_order_logits_bitwise_equal_native``."""
-    from torch.profiler import ProfilerActivity, profile
-
-    h = torch.randn(8, oracle.HIDDEN, device="cuda").to(torch.bfloat16)
-    w = shards[0]
-    kernels = {}
-    for m in (1, 3, 8):
-        F.linear(h[:m], w)
-        torch.cuda.synchronize()
-        with profile(activities=[ProfilerActivity.CUDA]) as prof:
-            F.linear(h[:m], w)
-            torch.cuda.synchronize()
-        kernels[m] = sorted(
-            {e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA}
-        )
-    try:
-        blas = str(torch.backends.cuda.preferred_blas_library())
-    except (AttributeError, RuntimeError):
-        blas = None
-    print(
-        "[draft_tail-nvjet-record] "
-        + json.dumps(
-            {
-                "torch": torch.__version__,
-                "cuda": torch.version.cuda,
-                "device": torch.cuda.get_device_name(0),
-                "capability": list(torch.cuda.get_device_capability(0)),
-                "preferred_blas": blas,
-                "cublas_libs": _loaded_cublas_versions(),
-                "env": {
-                    k: v
-                    for k, v in os.environ.items()
-                    if k.startswith(("CUBLAS", "TORCH_BLAS", "NVIDIA_TF32"))
-                },
-                "kernels_by_M": kernels,
-            }
-        )
-    )
-
-
-@pytest.mark.parametrize("m", [1, 8])
-def test_native_order_logits_bitwise_equal_native_in_graph(runner, shards, full_head, gamma, m):
-    """The served draft head runs inside a CUDA graph. Capture the native
-    LM-head GEMM and the shipped default draft_tail kernel (``DraftTailOptions().variant``)
-    per shard in graphs with M rows, replay them over the exactness rows
-    (real, knife-edge, synthetic) and require the draft_tail BF16 logits to equal
-    the graph-captured native logits bit for bit; the graph-captured native
-    logits are also compared with eager ``F.linear`` (printed)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
-        DraftTailOptions,
-    )
-
-    variant = DraftTailOptions().variant
-    rows = _exactness_rows(full_head, gamma)
-    rows = rows[: rows.shape[0] // m * m]
-    static_h = torch.zeros(m, oracle.HIDDEN, dtype=torch.bfloat16, device="cuda")
-    nat = [None] * len(shards)
-    w3 = [torch.empty(m, w.shape[0], dtype=torch.bfloat16, device="cuda") for w in shards]
-
-    def step():
-        for s, w in enumerate(shards):
-            nat[s] = F.linear(static_h, w)
-            runner.run_fused(static_h, w, s * w.shape[0], variant=variant, logits_out=w3[s])
-
-    st = torch.cuda.Stream()
-    st.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(st):
-        step()
-    torch.cuda.current_stream().wait_stream(st)
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        step()
-    draft_tail_vs_graph = graph_vs_eager = total = 0
-    for i in range(0, rows.shape[0], m):
-        static_h.copy_(rows[i : i + m])
-        g.replay()
-        torch.cuda.synchronize()
-        for s, w in enumerate(shards):
-            eager = F.linear(rows[i : i + m].contiguous(), w)
-            draft_tail_vs_graph += int((w3[s].view(torch.int16) != nat[s].view(torch.int16)).sum())
-            graph_vs_eager += int((nat[s].view(torch.int16) != eager.view(torch.int16)).sum())
-            total += eager.numel()
-    print(
-        f"[draft_tail-exact-graph] M={m} variant={variant} rows={rows.shape[0]} "
-        f"logits compared={total} draft_tail != graph native: {draft_tail_vs_graph}; "
-        f"graph native != eager native: {graph_vs_eager}"
-    )
-    assert draft_tail_vs_graph == 0
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py
deleted file mode 100644
index 876cca23c360..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_eligibility.py
+++ /dev/null
@@ -1,239 +0,0 @@
-# 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.
-"""draft_tail dispatch: native fallback is chosen before any device work.
-
-Uses the real ``LMHead`` and ``DeepseekV3MTPHead`` classes with stand-in
-worker / metadata objects; tp_size 1, so no exchange workspace is needed.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
-
-VOCAB = 4096
-
-
-def _objects(**over):
-    import tensorrt_llm  # noqa: F401
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3MTPHead
-    from tensorrt_llm._torch.modules.embedding import LMHead
-    from tensorrt_llm._torch.modules.linear import TensorParallelMode
-    from tensorrt_llm.mapping import Mapping
-
-    mapping = over.pop("mapping", Mapping(world_size=1, rank=0, tp_size=1))
-    head = LMHead(
-        VOCAB,
-        6144,
-        dtype=over.pop("head_dtype", torch.bfloat16),
-        mapping=mapping,
-        tensor_parallel_mode=TensorParallelMode.COLUMN,
-    )
-    head.weight.data = (torch.randn(head.weight.shape, device="cuda") * 0.02).to(head.weight.dtype)
-    shared_head = over.pop("shared_head", object.__new__(DeepseekV3MTPHead))
-    worker = SimpleNamespace(
-        is_mtp_eagle=over.pop("is_mtp_eagle", True),
-        guided_decoder=over.pop("guided_decoder", None),
-        _d2t=over.pop("d2t", None),
-        mapping=mapping,
-    )
-    spec = SimpleNamespace(
-        wants_advanced_draft_sampling=over.pop("advanced", False),
-        draft_vocab_size=0,
-        vocab_size=VOCAB,
-        num_seq_slots=over.pop("num_seq_slots", 8),
-        max_num_requests=over.pop("max_num_requests", 8),
-        is_cuda_graph=over.pop("is_cuda_graph", False),
-    )
-    draft = SimpleNamespace(mtp_layers=[SimpleNamespace(shared_head=shared_head)], lm_head=head)
-    assert not over, over
-    return worker, draft, spec
-
-
-def _select(worker, draft, spec):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    return dispatch.get_impl("draft_tail", "select_draft_tail")(worker, draft, spec)
-
-
-def test_eligible_path_matches_native_argmax():
-    worker, draft, spec = _objects()
-    tail = _select(worker, draft, spec)
-    assert tail is not None
-    gen = torch.Generator(device="cuda").manual_seed(0)
-    for rows in range(1, 9):
-        h = torch.randn(rows, 6144, generator=gen, device="cuda").to(torch.bfloat16)
-        tokens = tail(h)
-        ref = torch.argmax(F.linear(h, draft.lm_head.weight), dim=-1).to(torch.int32)
-        assert tokens is not None and tokens.dtype == torch.int32
-        torch.testing.assert_close(tokens, ref, rtol=0, atol=0)
-
-
-def test_per_step_fallback_only_for_rank_shared_properties():
-    """A step leaves draft_tail only for properties every TP rank shares: the row
-    count (same requests on every rank), dtype and hidden width (same
-    model)."""
-    worker, draft, spec = _objects()
-    tail = _select(worker, draft, spec)
-    h = torch.randn(9, 6144, device="cuda").to(torch.bfloat16)
-    assert tail(h) is None  # more than 8 rows
-    assert tail(h[:0]) is None  # no rows
-    assert tail(h[:2].half()) is None  # dtype
-    assert tail(h[:2, :4096].contiguous()) is None  # hidden size
-    assert tail(h[0]) is None  # not a row block
-
-
-def _misaligned_layouts(rows, gen):
-    """Row blocks whose layout the kernels cannot read directly: non-unit
-    last stride, 8-byte-aligned base, row stride not a multiple of 8."""
-    src = torch.randn(rows, 6144, generator=gen, device="cuda").to(torch.bfloat16)
-    strided = src.t().contiguous().t()
-    off = torch.empty(rows * 6144 + 4, dtype=torch.bfloat16, device="cuda")
-    shifted = off[4:].view(rows, 6144)
-    shifted.copy_(src)
-    wide = torch.empty(rows, 6148, dtype=torch.bfloat16, device="cuda")
-    pitched = wide[:, :6144]
-    pitched.copy_(src)
-    layouts = [shifted, pitched]
-    if rows > 1:  # one row has no meaningful last-dim transpose
-        layouts.append(strided)
-    return src, layouts
-
-
-def test_per_step_layout_is_staged_never_native():
-    """Strides and alignment are rank-local (allocator placement): they must
-    not move one rank to the native collective while its peers run draft_tail's
-    exchange. Such rows are copied into draft_tail's aligned staging rows."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
-        _kernel_layout_ok,
-    )
-
-    worker, draft, spec = _objects()
-    tail = _select(worker, draft, spec)
-    runner = dispatch.get_impl("draft_tail", "get_runner")(worker.mapping)
-    gen = torch.Generator(device="cuda").manual_seed(3)
-    for rows in (1, 2, 5, 8):
-        src, layouts = _misaligned_layouts(rows, gen)
-        ref = torch.argmax(F.linear(src, draft.lm_head.weight), dim=-1).to(torch.int32)
-        for x in layouts:
-            assert not _kernel_layout_ok(x)
-            before = runner.staged_steps
-            tokens = tail(x)
-            assert tokens is not None
-            assert runner.staged_steps == before + 1
-            torch.testing.assert_close(tokens, ref, rtol=0, atol=0)
-        before = runner.staged_steps
-        torch.testing.assert_close(tail(src), ref, rtol=0, atol=0)
-        assert runner.staged_steps == before  # aligned rows: no copy
-
-
-def test_per_step_staging_under_cuda_graph():
-    """The staging copy is captured with the step and replays with new
-    contents of the same (misaligned) source buffer."""
-    worker, draft, spec = _objects()
-    tail = _select(worker, draft, spec)
-    gen = torch.Generator(device="cuda").manual_seed(4)
-    off = torch.zeros(3 * 6144 + 4, dtype=torch.bfloat16, device="cuda")
-    x = off[4:].view(3, 6144)
-    x.copy_(torch.randn(3, 6144, generator=gen, device="cuda"))
-    tail(x)  # warm-up outside capture
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        out = tail(x)
-    for _ in range(5):
-        x.copy_(torch.randn(3, 6144, generator=gen, device="cuda"))
-        graph.replay()
-        torch.cuda.synchronize()
-        ref = torch.argmax(F.linear(x.contiguous(), draft.lm_head.weight), dim=-1).to(torch.int32)
-        torch.testing.assert_close(out, ref, rtol=0, atol=0)
-
-
-def test_per_step_device_mismatch_raises():
-    worker, draft, spec = _objects()
-    tail = _select(worker, draft, spec)
-    with pytest.raises(RuntimeError):
-        tail(torch.zeros(1, 6144, dtype=torch.bfloat16))
-
-
-@pytest.mark.parametrize(
-    "change",
-    [
-        dict(is_mtp_eagle=False),
-        dict(guided_decoder=object()),
-        dict(d2t=torch.zeros(4, dtype=torch.int64)),
-        dict(advanced=True),
-        dict(shared_head=object()),
-        dict(head_dtype=torch.float16),
-    ],
-)
-def test_configuration_fallbacks(change):
-    worker, draft, spec = _objects(**change)
-    assert _select(worker, draft, spec) is None
-
-
-@pytest.mark.parametrize(
-    "change,eligible",
-    [
-        (dict(num_seq_slots=9), False),
-        (dict(num_seq_slots=16, max_num_requests=4, is_cuda_graph=True), False),
-        (dict(num_seq_slots=0, max_num_requests=16), False),
-        (dict(num_seq_slots=0, max_num_requests=4, is_cuda_graph=True), False),
-        (dict(num_seq_slots=0, max_num_requests=4), True),
-        (dict(num_seq_slots=2, max_num_requests=1, is_cuda_graph=True), True),
-    ],
-)
-def test_row_bound_decides_per_configuration(change, eligible):
-    """draft_tail engages only if no draft step of the configuration can exceed 8
-    rows, judged from the engine's sequence-slot pool, which CUDA-graph
-    metadata keeps: a step can then never switch between draft_tail and native
-    rounding because of its (padded) row count."""
-    worker, draft, spec = _objects(**change)
-    assert (_select(worker, draft, spec) is not None) == eligible
-
-
-def test_one_kernel_for_every_row_count():
-    """The shipped options use one kernel shape for every row count, in the
-    native accumulation order (native draft tokens, near-ties included)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.draft_tail.draft_tail import (  # noqa: E501
-        FIRST_MMA_VARIANT,
-        NATIVE_ORDER_VARIANTS,
-        DraftTailOptions,
-    )
-
-    opts = DraftTailOptions()
-    assert opts.mode == "fused" and opts.variant >= FIRST_MMA_VARIANT
-    assert opts.variant in NATIVE_ORDER_VARIANTS
-    assert opts.shape(152) == (opts.variant, 0)
-
-
-def test_attention_dp_falls_back():
-    from tensorrt_llm.mapping import Mapping
-
-    mapping = Mapping(world_size=2, rank=0, tp_size=2, enable_attention_dp=True)
-    worker, draft, spec = _objects(mapping=mapping)
-    assert _select(worker, draft, spec) is None
-
-
-def test_feature_off_by_default(monkeypatch):
-    """Without an installed plan (a process that does not build the target) draft_tail is off."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
-    assert not dispatch.is_enabled("draft_tail")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py
deleted file mode 100644
index b2a773ca0212..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/draft_tail/test_draft_tail_nearties.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# 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.
-"""Self-checks of the near-tie tooling used by the flip-rate study."""
-
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import draft_tail_nearties as nt  # noqa: E402
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
-
-
-def test_bf16_step_walks_the_grid():
-    x = torch.tensor([1.0, -1.0, 0.0, 3.0, 17.875], dtype=torch.bfloat16)
-    up = nt.bf16_step(x, 1).double()
-    down = nt.bf16_step(x, -1).double()
-    assert bool((up > x.double()).all()) and bool((down < x.double()).all())
-    # No BF16 value lies strictly between x and its neighbours.
-    mid_up = ((x.double() + up) / 2).float().to(torch.bfloat16).double()
-    assert bool(((mid_up == x.double()) | (mid_up == up)).all())
-    assert nt.bf16_step(x, 1)[4].item() == 18.0
-    assert nt.bf16_step(nt.bf16_step(x, 5), -5).tolist() == x.tolist()
-
-
-def test_bf16_round_exact_is_nearest_without_double_rounding():
-    gen = torch.Generator().manual_seed(0)
-    v = torch.randn(200000, generator=gen, dtype=torch.float64) * 20
-    r = nt.bf16_round_exact(v)
-    lo = nt.bf16_step(r, -1).double()
-    hi = nt.bf16_step(r, 1).double()
-    rd = r.double()
-    assert bool(((v - rd).abs() <= (v - lo).abs()).all())
-    assert bool(((v - rd).abs() <= (v - hi).abs()).all())
-    # Just above / below a midpoint, closer than FP32 can resolve.
-    b = torch.tensor([17.875], dtype=torch.bfloat16)
-    n = nt.bf16_step(b, 1)
-    m = (b.double() + n.double()) / 2
-    assert nt.bf16_round_exact(m + 1e-9).item() == n.item()
-    assert nt.bf16_round_exact(m - 1e-9).item() == b.item()
-    # Exactly at the midpoint: ties to even (17.875 has an odd mantissa).
-    assert nt.bf16_round_exact(m).item() == n.item()
-
-
-def test_bf16_key_orders_value_then_lower_id():
-    vals = torch.tensor([[1.0, 2.0, 2.0, -0.0, 0.0, -3.0]], dtype=torch.bfloat16)
-    ids = torch.arange(6)[None]
-    assert int(nt.bf16_key(vals, ids).argmax()) == 1
-    z = torch.tensor([[-0.0, 0.0]], dtype=torch.bfloat16)
-    k = nt.bf16_key(z, torch.arange(2)[None])
-    assert int(k.argmax()) == 0  # -0 == +0: lower id
-
-
-def test_construct_straddle_reaches_targets():
-    gen = torch.Generator(device="cuda").manual_seed(1)
-    head = (torch.randn(64, 6144, generator=gen, device="cuda") * 0.02).to(torch.bfloat16)
-    hidden = torch.randn(16, 6144, generator=gen, device="cuda").to(torch.bfloat16)
-    up = torch.full((16,), 3, device="cuda")
-    down = torch.full((16,), 7, device="cuda")
-    for tau in (1.0, 64.0):
-        res = nt.construct_straddle(hidden, head, up, down, torch.full((16,), tau, device="cuda"))
-        ok = (
-            (res["dist_hi_fp32_ulps"] > 0)
-            & (res["dist_lo_fp32_ulps"] < 0)
-            & (res["dist_hi_fp32_ulps"] <= 1.1 * tau)
-            & (res["dist_lo_fp32_ulps"] >= -1.1 * tau)
-        )
-        # Greedy single-element steps from large toy gaps (about 2): most
-        # rows reach the band (real near-tie rows start close); callers
-        # keep only rows that did.
-        assert int(ok.sum()) >= 8, (tau, res["dist_hi_fp32_ulps"], res["dist_lo_fp32_ulps"])
-        # The returned logits are those of the returned rows.
-        h = res["hidden"].double()
-        torch.testing.assert_close((h * head[3].double()).sum(-1), res["l_hi"], rtol=0, atol=1e-9)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py
deleted file mode 100644
index 9cc26219a55b..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/conftest.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# 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.
-"""Make the fused_kv_indexer_projection helpers and the absorbed_query / post_moe_hop fixtures importable under
-``--import-mode=importlib``."""
-
-import sys
-from pathlib import Path
-
-_HERE = Path(__file__).resolve().parent
-for _path in (str(_HERE), str(_HERE.parent / "absorbed_query"), str(_HERE.parent / "post_moe_hop")):
-    if _path not in sys.path:
-        sys.path.insert(0, _path)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py
deleted file mode 100644
index a4be67eb41aa..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/fused_kv_utils.py
+++ /dev/null
@@ -1,288 +0,0 @@
-# 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.
-"""Shared helpers of the fused_kv_indexer_projection tests and benchmarks.
-
-* ``module(layer)``: the native ``DeepseekV32Attention`` of a GLM-5.2 layer
-  at TP8 rank 0 with real checkpoint weights (absorbed_query's fixture; layer 78 is the
-  MTP layer), with its real ``input_layernorm`` gain.
-* ``native_hidden``: the tensor the native GEMMs consume: the native
-  ``RMSNorm`` module (FlashInfer kernel) applied with the real gain to
-  synthetic residual-stream rows; ``real_hidden``: rows of the real
-  normalized MoE-input fixture (real-text prefill, GLM-5.2 NVFP4).
-* ``native(attn, h)``: the native ``kv_a_proj_with_mqa`` call
-  (``dsv3_fused_a_gemm_op``) and ``F.linear`` on the indexer's fused
-  ``wk | weights_proj`` weight.
-* error statistics against an FP64 oracle and the frozen budget (the deferred_norm
-  budget for the same GEMM).
-* ``proxy_producer``: a CuTe DSL stand-in for the preceding kernel. It is
-  launched with programmatic dependent launch, releases its dependents at
-  entry, spins for a given time and only then writes the activation; a
-  dependent that read it before its dependency wait would see stale data.
-"""
-
-import glob
-import os
-import sys
-
-import cutlass.cute as cute
-import torch
-from cuda.bindings.driver import CUstream
-from cutlass import Int32, Int64
-from cutlass._mlir.dialects import llvm
-from cutlass.cutlass_dsl import dsl_user_op
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-for _p in (_HERE, os.path.join(_HERE, "..", "absorbed_query")):
-    if _p not in sys.path:
-        sys.path.insert(0, _p)
-
-import absorbed_query_fixture  # noqa: E402
-
-H = 6144
-EPS = 1e-5
-MTP_LAYER = 78
-REAL_HIDDEN_DIR = os.environ.get("GLM_MOE_DSA_REAL_MOE_INPUTS_DIR", "")
-
-_MODS = {}
-_MODEL_CONFIG = {}
-
-
-def model_config(tp_rank: int = 0):
-    if tp_rank not in _MODEL_CONFIG:
-        _MODEL_CONFIG[tp_rank] = absorbed_query_fixture.build_model_config(tp_rank)
-    return _MODEL_CONFIG[tp_rank]
-
-
-def module(layer: int):
-    """``(attn, input_layernorm gain)`` of ``layer`` (cached)."""
-    if layer not in _MODS:
-        with torch.inference_mode():
-            _MODS[layer] = (
-                absorbed_query_fixture.build_native_attention(layer, 0, model_config(0)),
-                absorbed_query_fixture.load_input_layernorm(layer).cuda(),
-            )
-    return _MODS[layer]
-
-
-def make_residual(rows: int, seed: int, scale: float = 1.0) -> torch.Tensor:
-    """Residual-stream-like rows: Gaussian plus a few large channels."""
-    g = torch.Generator(device="cuda").manual_seed(seed)
-    x = torch.randn(rows, H, generator=g, device="cuda") * 2.0
-    hot = torch.randint(0, H, (12,), generator=g, device="cuda")
-    x[:, hot] *= 30.0
-    return (x * scale).bfloat16()
-
-
-_NORMS = {}
-
-
-def native_rmsnorm(gamma: torch.Tensor):
-    """The native ``RMSNorm`` module (as the decoder layers build it)."""
-    key = gamma.data_ptr()
-    if key not in _NORMS:
-        from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-
-        norm = RMSNorm(hidden_size=H, eps=EPS, dtype=torch.bfloat16).cuda()
-        norm.weight.data.copy_(gamma)
-        _NORMS[key] = norm
-    return _NORMS[key]
-
-
-def native_hidden(rows: int, seed: int, gamma: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
-    """Native RMSNorm output of synthetic residual rows, times ``scale``."""
-    h = native_rmsnorm(gamma)(make_residual(rows, seed))
-    if scale != 1.0:
-        h = (h.float() * scale).bfloat16()
-    return h.contiguous()
-
-
-_REAL = []
-
-
-def real_hidden(rows: int, offset: int = 0):
-    """``rows`` rows of a real normalized hidden state, or ``None`` if the
-    fixture is not mounted (first 2-D BF16 ``[N, 6144]`` tensor of the first
-    file)."""
-    if not _REAL:
-        files = sorted(glob.glob(os.path.join(REAL_HIDDEN_DIR, "*.pt")))
-        if not files:
-            return None
-        obj = torch.load(files[0], map_location="cpu", weights_only=False)
-        cands = (
-            obj.values()
-            if isinstance(obj, dict)
-            else (obj if isinstance(obj, (list, tuple)) else [obj])
-        )
-        found = None
-        for t in cands:
-            if (
-                isinstance(t, torch.Tensor)
-                and t.dim() == 2
-                and t.shape[1] == H
-                and t.dtype == torch.bfloat16
-            ):
-                found = t
-                break
-        if found is None:
-            return None
-        _REAL.append(found)
-    t = _REAL[0]
-    return t[offset : offset + rows].contiguous().cuda()
-
-
-def native(attn, h: torch.Tensor):
-    kva = attn.kv_a_proj_with_mqa(h)
-    kw = None
-    if attn.mqa.indexer is not None:
-        kw = torch.nn.functional.linear(h, attn.mqa.indexer._fused_wk_wp_weight)
-    return kva, kw
-
-
-def fp64_oracle(h: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
-    return h.double() @ weight.double().t()
-
-
-def norm_rms_err(y: torch.Tensor, ref: torch.Tensor) -> float:
-    ref = ref.double()
-    return float((y.double() - ref).norm() / ref.norm().clamp_min(1e-300))
-
-
-def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    def ordered(t):
-        v = t.contiguous().view(torch.int16).to(torch.int32)
-        return torch.where(v < 0, -(v + 32768), v)
-
-    return (ordered(a) - ordered(b)).abs()
-
-
-def compare(cand: torch.Tensor, nat: torch.Tensor, oracle: torch.Tensor) -> dict:
-    ulp = bf16_ulp_distance(cand, nat)
-    return {
-        "bit_identical": bool(torch.equal(cand.view(torch.int16), nat.view(torch.int16))),
-        "allclose_native": bool(torch.allclose(cand.float(), nat.float(), rtol=0.1, atol=1e-2)),
-        "e_cand": norm_rms_err(cand, oracle),
-        "e_native": norm_rms_err(nat, oracle),
-        "max_ulp_vs_native": int(ulp.max()),
-        "frac_diff_vs_native": float((ulp > 0).float().mean()),
-        "max_abs_vs_native": float((cand.float() - nat.float()).abs().max()),
-    }
-
-
-def within_budget(stats: dict) -> bool:
-    """Frozen budget of a non-bit-identical output."""
-    return stats["allclose_native"] and stats["e_cand"] <= 1.25 * stats["e_native"] + 2.0**-12
-
-
-# --------------------------------------------------------------------------
-# Proxy producer (PDL): releases dependents at entry, writes after a delay.
-# --------------------------------------------------------------------------
-
-
-def _asm(res, args, text, cons, loc=None, ip=None):
-    return llvm.inline_asm(
-        res,
-        args,
-        text,
-        cons,
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def _spin_ns(ns: Int64, *, loc=None, ip=None):
-    _asm(
-        None,
-        [ns.ir_value()],
-        "{\n\t.reg .u64 t0, t;\n\t.reg .pred p;\n\tmov.u64 t0, %globaltimer;\n\t"
-        "GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_SPIN:\n\tmov.u64 t, %globaltimer;\n\tsub.u64 t, t, t0;\n\t"
-        "setp.lt.u64 p, t, $0;\n\t@p bra GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_SPIN;\n\t}",
-        "l",
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def _copy16(dst: Int64, src: Int64, *, loc=None, ip=None):
-    _asm(
-        None,
-        [dst.ir_value(), src.ir_value()],
-        "{\n\t.reg .b32 a, b, c, d;\n\tld.global.v4.b32 {a, b, c, d}, [$1];\n\t"
-        "st.global.v4.b32 [$0], {a, b, c, d};\n\t}",
-        "l,l",
-        loc=loc,
-        ip=ip,
-    )
-
-
-class _ProxyKernel:
-    """Grid ``[8, rows]`` x 96 threads; CTA ``(c, r)`` copies 1536 B of row ``r``."""
-
-    @cute.jit
-    def __call__(
-        self,
-        src: Int64,
-        dst: Int64,
-        delay_ns: Int64,
-        anchor: cute.Tensor,
-        rows: Int32,
-        stream: CUstream,
-    ):
-        self.kernel(src, dst, delay_ns, anchor).launch(
-            grid=[8, rows, 1], block=[96, 1, 1], stream=stream, use_pdl=True
-        )
-
-    @cute.kernel
-    def kernel(self, src: Int64, dst: Int64, delay_ns: Int64, anchor: cute.Tensor):
-        tid, _, _ = cute.arch.thread_idx()
-        cta, row, _ = cute.arch.block_idx()
-        cute.arch.griddepcontrol_wait()
-        cute.arch.griddepcontrol_launch_dependents()
-        _spin_ns(delay_ns)
-        off = Int64(row) * Int64(H * 2) + Int64((cta * Int32(96) + tid) * Int32(16))
-        _copy16(dst + off, src + off)
-
-
-_PROXY = {}
-
-
-def proxy_producer(src: torch.Tensor, dst: torch.Tensor, delay_ns: int) -> None:
-    """Launch the proxy producer on the current stream (PDL): ``dst <- src``
-    (``[R, 6144]`` BF16) after spinning ``delay_ns``."""
-    if "k" not in _PROXY:
-        from cutlass.cute import runtime as cute_runtime
-
-        z = Int64(0)
-        anchor = cute_runtime.make_fake_compact_tensor(
-            Int32, (cute.sym_int(),), stride_order=(0,), assumed_align=4
-        )
-        _PROXY["k"] = cute.compile(
-            _ProxyKernel(),
-            z,
-            z,
-            z,
-            anchor,
-            Int32(1),
-            stream=cute_runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
-            options="--enable-tvm-ffi",
-        )
-    assert src.shape == dst.shape and src.shape[1] == H and src.is_contiguous()
-    _PROXY["k"](
-        src.data_ptr(), dst.data_ptr(), int(delay_ns), dst.view(torch.int32).view(-1), src.shape[0]
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py
deleted file mode 100644
index 76fe1fe14450..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_attention_integration.py
+++ /dev/null
@@ -1,275 +0,0 @@
-# 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.
-"""fused_kv_indexer_projection through the target's absorbed-query packet into the native DSA
-attention (single GPU; the target's adopted attention, generation steps as decode steps).
-
-Uses the absorbed_query integration harness (real weights of the full-indexer layer 6
-and the shared layer 7 at TP8 rank 0, six requests natively prefilled with
-2,600 tokens, FP8 KV cache, layouts 6 x 1 and 1 x 6 (MTP verify), both call
-paths, ``on_update_kv_lens()`` before every step). The same native
-normalized hidden state goes through ``absorbed_query,shared_indexer_absorbed_query``
-(absorbed_query / shared_indexer_absorbed_query with their native GEMM bridges) and
-``absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection``
-(fused_kv_indexer_projection serves both layers: ``entry.DEFAULT_LAYERS`` set to
-``full,shared``).
-
-Checks (the fused_kv_indexer_projection contract and the attention-level sanity bound
-declared with it):
-
-* engagement: fused_kv_indexer_projection serves both layers;
-* every latent-cache byte and every recorded tensor derived only from
-  ``kva`` (absorbed / FP8 query, BMM scales, indexer FP8 query) is
-  bit-identical; the indexer-K cache, indexer weights and top-k sets are
-  bit-identical when ``kw`` is, else reported and bounded (top-k mean
-  Jaccard >= 0.98, output relRMS <= 2**-7);
-* attention outputs: within ``max(2 * native rerun noise, 4 bf16 ulps)``
-  (the native sparse attention is not bitwise run-to-run stable);
-* CUDA graph: a captured fused_kv_indexer_projection step replays (1, 1, 20 times, fresh inputs)
-  equal to the eager fused_kv_indexer_projection step;
-* with no layer kind selected, ``fused_kv_indexer_projection`` selected runs the absorbed_query /
-  shared_indexer_absorbed_query path bit for bit.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
-sys.path.insert(0, os.path.join(_HERE, ".."))
-
-import absorbed_query_fixture  # noqa: E402
-import plan_testlib  # noqa: E402
-import test_absorbed_query_attention_integration as aq_integration  # noqa: E402
-
-LAYERS = aq_integration.LAYERS
-LAYOUTS = aq_integration.LAYOUTS
-KW_EXACT = os.environ.get("GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW_EXACT", "0") == "1"
-KW_RECORDS = ("indexer_weights",)
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"),
-    pytest.mark.skipif(
-        not absorbed_query_fixture.checkpoint_available(), reason="needs the GLM-5.2 checkpoint"
-    ),
-]
-
-
-def _set_features(features: str, fused_kv_layers: str = "full,shared"):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        entry,
-    )
-
-    entry.DEFAULT_LAYERS = tuple(fused_kv_layers.split(","))
-    plan_testlib.select(dispatch, absorbed_query_fixture.with_base_features(features))
-
-
-def _reset_runners(h):
-    for layer in LAYERS:
-        for name in ("_absorbed_query", "_shared_indexer_query", "_fused_kv_projection"):
-            setattr(h.attn[layer], name, None)
-
-
-@pytest.fixture(scope="module")
-def harness():
-    with torch.inference_mode():
-        h = aq_integration._Harness()
-    yield h
-    _set_features("")
-
-
-def _run(h, features, layout, path, hidden, pos, snap, fused_kv_layers="full,shared"):
-    _set_features(features, fused_kv_layers)
-    _reset_runners(h)
-    h.restore(snap)
-    out, rec = h.recorded_step(layout, path, hidden, pos)
-    return out, rec, h.snapshot()
-
-
-def _jaccard_rows(a, b):
-    rows = []
-    for ra, rb in zip(a.tolist(), b.tolist()):
-        sa = {v for v in ra if v >= 0}
-        sb = {v for v in rb if v >= 0}
-        union = len(sa | sb)
-        rows.append(1.0 if union == 0 else len(sa & sb) / union)
-    return rows
-
-
-def _compare(native, rerun, cand, report):
-    n_out, n_rec, n_cache = native
-    c_out, c_rec, c_cache = cand
-    fails = []
-    for key in n_cache:
-        same = torch.equal(c_cache[key], n_cache[key])
-        report[f"cache_{key}_identical"] = same
-        if not same and (key.startswith("latent") or KW_EXACT):
-            fails.append(f"cache {key} differs")
-    for layer in LAYERS:
-        for key, val in n_rec[layer].items():
-            if key in ("topk", "sparse_indices"):
-                continue
-            same = torch.equal(c_rec[layer][key], val)
-            report[f"L{layer}_{key}_identical"] = same
-            if not same and (key not in KW_RECORDS or KW_EXACT):
-                fails.append(f"L{layer} record {key} differs")
-        if "topk" in n_rec[layer]:
-            jac = _jaccard_rows(c_rec[layer]["topk"], n_rec[layer]["topk"])
-            report[f"L{layer}_topk_mean_jaccard"] = sum(jac) / len(jac)
-            if min(jac) < 1.0 and KW_EXACT:
-                fails.append(f"L{layer} top-k sets differ")
-            if sum(jac) / len(jac) < 0.98:
-                fails.append(f"L{layer} top-k mean Jaccard {sum(jac) / len(jac):.4f} < 0.98")
-        diff = (c_out[layer].float() - n_out[layer].float()).abs().max().item()
-        noise = (rerun[0][layer].float() - n_out[layer].float()).abs().max().item()
-        ulp = aq_integration._bf16_ulp(n_out[layer].float().abs().max().item())
-        rel = aq_integration._rel_rms(c_out[layer], n_out[layer])
-        report[f"L{layer}_out_maxdiff"] = diff
-        report[f"L{layer}_out_noise"] = noise
-        report[f"L{layer}_out_rel_rms"] = rel
-        if KW_EXACT or LAYERS.index(layer) == 1:
-            # kva-only layer (shared) or exact kw: native run-to-run bound.
-            if diff > max(2 * noise, 4 * ulp):
-                fails.append(f"L{layer} output diff {diff:.3e} > max(2*{noise:.3e}, 4*{ulp:.3e})")
-        elif rel > 2.0**-7:
-            fails.append(f"L{layer} output relRMS {rel:.3e} > 2**-7")
-    return fails
-
-
-@pytest.mark.parametrize("path", ["custom_op", "forward_impl"])
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_fused_kv_into_attention(harness, layout, path):
-    h = harness
-    hidden, pos = h.gen_inputs(layout, seed=3)
-    snap = h.snapshot()
-    try:
-        native = _run(
-            h, "absorbed_query,shared_indexer_absorbed_query", layout, path, hidden, pos, snap
-        )
-        cand = _run(
-            h,
-            "absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection",
-            layout,
-            path,
-            hidden,
-            pos,
-            snap,
-        )
-        for layer in LAYERS:
-            assert h.attn[layer]._fused_kv_projection not in (None, False), (
-                f"fused_kv_indexer_projection not engaged on L{layer}"
-            )
-        rerun = _run(
-            h, "absorbed_query,shared_indexer_absorbed_query", layout, path, hidden, pos, snap
-        )
-        report = {}
-        fails = _compare(native, rerun, cand, report)
-        print(
-            f"GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_ATTN layout={layout} path={path} {report}",
-            flush=True,
-        )
-        assert not fails, fails
-    finally:
-        h.restore(snap)
-        _set_features("")
-
-
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_no_layer_kind_selected_is_native_path(harness, layout):
-    h = harness
-    hidden, pos = h.gen_inputs(layout, seed=5)
-    snap = h.snapshot()
-    try:
-        n_out, n_rec, n_cache = _run(
-            h,
-            "absorbed_query,shared_indexer_absorbed_query",
-            layout,
-            "custom_op",
-            hidden,
-            pos,
-            snap,
-        )
-        c_out, c_rec, c_cache = _run(
-            h,
-            "absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection",
-            layout,
-            "custom_op",
-            hidden,
-            pos,
-            snap,
-            fused_kv_layers="",
-        )
-        for layer in LAYERS:
-            assert h.attn[layer]._fused_kv_projection is False
-        for key in n_cache:
-            assert torch.equal(c_cache[key], n_cache[key]), key
-        for layer in LAYERS:
-            for key, val in n_rec[layer].items():
-                if key in ("topk", "sparse_indices"):
-                    continue
-                assert torch.equal(c_rec[layer][key], val), (layer, key)
-    finally:
-        h.restore(snap)
-        _set_features("")
-
-
-@pytest.mark.parametrize("layout", list(LAYOUTS))
-@torch.inference_mode()
-def test_fused_kv_cuda_graph_replay(harness, layout):
-    h = harness
-    snap = h.snapshot()
-    try:
-        _set_features("absorbed_query,shared_indexer_absorbed_query,fused_kv_indexer_projection")
-        _reset_runners(h)
-        hidden, pos = h.gen_inputs(layout, seed=7)
-        static = hidden.clone()
-        static_pos = pos.clone()
-        meta = h.gen_metadata(layout, cuda_graph=True)
-        h.restore(snap)
-        h.step(meta, static_pos, static, "custom_op")
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        h.restore(snap)
-        with torch.cuda.graph(graph):
-            g_out = h.step(meta, static_pos, static, "custom_op")
-        for trial, replays in enumerate((1, 1, 20)):
-            fresh, _ = h.gen_inputs(layout, seed=100 + trial)
-            static.copy_(fresh)
-            for _ in range(replays):
-                h.restore(snap)
-                graph.replay()
-            torch.cuda.synchronize()
-            got = {layer: g_out[layer].clone() for layer in LAYERS}
-            got_cache = h.snapshot()
-            h.restore(snap)
-            eager = h.gen_step(
-                layout, "custom_op", fresh, pos, meta=h.gen_metadata(layout, cuda_graph=True)
-            )
-            eager_cache = h.snapshot()
-            for key in eager_cache:
-                assert torch.equal(got_cache[key], eager_cache[key]), (trial, key)
-            for layer in LAYERS:
-                diff = (got[layer].float() - eager[layer].float()).abs().max().item()
-                ref = eager[layer].float().abs().max().item()
-                assert diff <= 4 * aq_integration._bf16_ulp(ref), (trial, layer, diff)
-    finally:
-        h.restore(snap)
-        _set_features("")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py
deleted file mode 100644
index fffdd1d29299..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_entry.py
+++ /dev/null
@@ -1,294 +0,0 @@
-# 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.
-"""The query_front packet through the fused_kv_indexer_projection entry (``fused_kv_dsa_proj``)
-against the host package's own path on the same hidden state (real weights, single GPU).
-
-* full-indexer target layer 0 and the MTP layer 78 (draft_absorbed_query): absorbed_query's
-  ``project`` with its native bridges (``fused_a`` + q_a norm next to the
-  ``wk | weights_proj`` GEMM + k_norm on the aux stream) vs fused_kv_indexer_projection +
-  the same norms + the same absorbed_query kernel. Every packet tensor that depends only
-  on ``kva`` must be bit-identical; the tensors derived from ``kw`` (indexer
-  key FP8 + scale, indexer weights) are bit-identical when ``kw`` is, and
-  otherwise held to the frozen budget.
-* shared-indexer target layer 3: shared_indexer_absorbed_query's ``project`` (native ``fused_a`` +
-  shared_indexer_absorbed_query kernel) and absorbed_query's shared-layer path vs
-  fused_kv_indexer_projection + the same kernels: bit-identical.
-* static declines (layer kinds not selected, ineligible module) return
-  ``None`` before any GPU work; a compile error propagates and is not
-  cached.
-
-The feature selection is set by replacing ``dispatch.enabled_features``
-(test-only), so these tests do not need the target's model.
-"""
-
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
-
-import absorbed_query_fixture  # noqa: E402
-import fused_kv_utils as U  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
-    reason="needs a GPU and the GLM-5.2 checkpoint",
-)
-
-KW_EXACT = os.environ.get("GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW_EXACT", "0") == "1"
-CTX = 491061
-GEN_META = SimpleNamespace(num_contexts=0, num_generations=1)
-# Packet slots (AbsorbedQueryPacket.proj_outputs) of a full-indexer layer.
-KVA_SLOTS = (0, 1, 2, 3, 4, 8)  # q, compressed_kv, k_pe, latent, indexer q_fp8, q_scale
-KW_SLOTS = (5, 6, 7)  # indexer k_fp8, k_scale, weights
-
-
-@pytest.fixture
-def features(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        entry,
-    )
-
-    def select(ids):
-        monkeypatch.setattr(dispatch, "enabled_features", lambda: frozenset(ids))
-
-    select(
-        {
-            "absorbed_query",
-            "draft_absorbed_query",
-            "shared_indexer_absorbed_query",
-            "fused_kv_indexer_projection",
-        }
-    )
-    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("full", "shared", "draft"))
-    return select
-
-
-def _reset(attn):
-    for name in (
-        "_absorbed_query",
-        "_draft_absorbed_query",
-        "_shared_indexer_query",
-        "_fused_kv_projection",
-    ):
-        setattr(attn, name, None)
-
-
-def _pos(rows):
-    return (CTX + torch.arange(rows, dtype=torch.int32, device="cuda")).unsqueeze(0)
-
-
-def _assert_same(a, b, tag):
-    assert a.shape == b.shape and a.dtype == b.dtype, tag
-    assert torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)), tag
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("layer", [0, U.MTP_LAYER])
-@torch.inference_mode()
-def test_full_layer_packet(features, layer, rows):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        absorb_for,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
-        fused_kv_dsa_proj,
-    )
-
-    attn, gamma = U.module(layer)
-    _reset(attn)
-    h = U.native_hidden(rows, seed=31 + rows, gamma=gamma)
-    pos = _pos(rows)
-    ref = AbsorbedQueryProjection(attn).project(pos, h, absorb=absorb_for(attn))
-    got = fused_kv_dsa_proj(attn, pos, h, GEN_META)
-    torch.cuda.synchronize()
-    assert got is not None and attn._fused_kv_projection not in (None, False)
-    assert len(got.proj_outputs) == len(ref.proj_outputs) == 9
-    for i in KVA_SLOTS:
-        _assert_same(got.proj_outputs[i], ref.proj_outputs[i], f"slot {i}")
-    assert (got.fused_q is None) == (ref.fused_q is None)
-    if ref.fused_q is not None:
-        _assert_same(got.fused_q[..., :512], ref.fused_q[..., :512], "fused_q")
-    k_fp8, k_scale, weights = (got.proj_outputs[i] for i in KW_SLOTS)
-    nk_fp8, nk_scale, nweights = (ref.proj_outputs[i] for i in KW_SLOTS)
-    if KW_EXACT:
-        for i in KW_SLOTS:
-            _assert_same(got.proj_outputs[i], ref.proj_outputs[i], f"slot {i}")
-    else:
-        torch.testing.assert_close(
-            k_fp8.float() * k_scale, nk_fp8.float() * nk_scale, rtol=0.1, atol=1e-2
-        )
-        torch.testing.assert_close(weights, nweights, rtol=0.1, atol=1e-2)
-    diff = {
-        f"slot{i}": float(
-            (
-                got.proj_outputs[i].contiguous().view(torch.uint8)
-                != ref.proj_outputs[i].contiguous().view(torch.uint8)
-            )
-            .float()
-            .mean()
-        )
-        for i in KW_SLOTS
-    }
-    print(
-        f"GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_ENTRY layer={layer} rows={rows} kw_derived_byte_diff={diff}"
-    )
-
-
-@pytest.mark.parametrize("host", ["shared_indexer_absorbed_query", "absorbed_query"])
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@torch.inference_mode()
-def test_shared_layer_packet(features, host, rows):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.query_proj import (  # noqa: E501
-        AbsorbedQueryProjection,
-        absorb_for,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.absorbed_query.shared_indexer import (  # noqa: E501
-        SharedIndexerProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
-        fused_kv_dsa_proj,
-    )
-
-    if host == "absorbed_query":
-        features({"absorbed_query", "fused_kv_indexer_projection"})
-    attn, gamma = U.module(3)
-    _reset(attn)
-    h = U.native_hidden(rows, seed=71 + rows, gamma=gamma)
-    pos = _pos(rows)
-    if host == "shared_indexer_absorbed_query":
-        ref = SharedIndexerProjection(attn).project(h)
-    else:
-        ref = AbsorbedQueryProjection(attn).project(pos, h, absorb=absorb_for(attn))
-    got = fused_kv_dsa_proj(attn, pos, h, GEN_META)
-    torch.cuda.synchronize()
-    assert got is not None
-    assert len(got.proj_outputs) == len(ref.proj_outputs) == 4
-    for i in range(4):
-        _assert_same(got.proj_outputs[i], ref.proj_outputs[i], f"slot {i}")
-    assert (got.fused_q is None) == (ref.fused_q is None)
-    if ref.fused_q is not None:
-        _assert_same(got.fused_q[..., :512], ref.fused_q[..., :512], "fused_q")
-
-
-@torch.inference_mode()
-def test_static_declines_before_gpu_work(features, monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        entry,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
-        fused_kv_dsa_proj,
-    )
-
-    attn, gamma = U.module(0)
-    h = U.native_hidden(6, seed=3, gamma=gamma)
-    # Layer kind not selected.
-    _reset(attn)
-    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("shared",))
-    assert fused_kv_dsa_proj(attn, _pos(6), h, GEN_META) is None
-    assert attn._fused_kv_projection is False
-    # fused_kv_indexer_projection not selected.
-    _reset(attn)
-    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("full",))
-    features({"absorbed_query"})
-    assert fused_kv_dsa_proj(attn, _pos(6), h, GEN_META) is None
-    # Per call: context batch, too many rows, non-contiguous input.
-    features({"absorbed_query", "fused_kv_indexer_projection"})
-    _reset(attn)
-    assert (
-        fused_kv_dsa_proj(attn, _pos(6), h, SimpleNamespace(num_contexts=1, num_generations=0))
-        is None
-    )
-    h9 = U.native_hidden(9, seed=4, gamma=gamma)
-    assert fused_kv_dsa_proj(attn, _pos(9), h9, GEN_META) is None
-    assert fused_kv_dsa_proj(attn, _pos(3), h9[::3], GEN_META) is None
-    _reset(attn)
-
-
-def test_layer_kinds_and_default_selection(monkeypatch):
-    """Layer 3 is the dense-fed shared layer; the default serves only it."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        entry,
-    )
-
-    def mla(idx, full):
-        return SimpleNamespace(
-            layer_idx=idx,
-            num_hidden_layers=78,
-            mqa=SimpleNamespace(indexer=object() if full else None),
-        )
-
-    assert entry.layer_kind(mla(0, True)) == "full"
-    assert entry.layer_kind(mla(3, False)) == "dense_shared"
-    assert entry.layer_kind(mla(7, False)) == "shared"
-    assert entry.layer_kind(mla(78, True)) == "draft"
-    served = entry.served_layer_kinds()
-    assert served == ("dense_shared",)
-    assert entry.kind_selected("dense_shared", served)
-    for kind in ("full", "shared", "draft"):
-        assert not entry.kind_selected(kind, served)
-    assert entry.kind_selected("dense_shared", ("shared",))
-    assert not entry.kind_selected("shared", ("dense_shared",))
-
-
-@torch.inference_mode()
-def test_default_selection_serves_layer3_only(features, monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        entry,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
-        fused_kv_dsa_proj,
-    )
-
-    monkeypatch.setattr(entry, "DEFAULT_LAYERS", ("dense_shared",))
-    full, gamma0 = U.module(0)
-    _reset(full)
-    assert (
-        fused_kv_dsa_proj(full, _pos(6), U.native_hidden(6, seed=8, gamma=gamma0), GEN_META) is None
-    )
-    assert full._fused_kv_projection is False
-    shared, gamma3 = U.module(3)
-    _reset(shared)
-    got = fused_kv_dsa_proj(shared, _pos(6), U.native_hidden(6, seed=9, gamma=gamma3), GEN_META)
-    assert got is not None
-    _reset(full)
-    _reset(shared)
-
-
-@torch.inference_mode()
-def test_compile_error_is_not_masked(features, monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        projection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.entry import (  # noqa: E501
-        fused_kv_dsa_proj,
-    )
-
-    def _fail(self):
-        raise RuntimeError("injected projection compile failure")
-
-    attn, gamma = U.module(0)
-    _reset(attn)
-    monkeypatch.setattr(projection.FusedKvProjection, "precompile", _fail)
-    with pytest.raises(RuntimeError, match="injected projection compile failure"):
-        fused_kv_dsa_proj(attn, _pos(6), U.native_hidden(6, seed=5, gamma=gamma), GEN_META)
-    assert attn._fused_kv_projection is None, "a failed runner must not be cached"
-    _reset(attn)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py
deleted file mode 100644
index 1dd8e71e8a00..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/fused_kv_indexer_projection/test_fused_kv_projection.py
+++ /dev/null
@@ -1,331 +0,0 @@
-# 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.
-"""Single-GPU tests of the fused_kv_indexer_projection merged projection (real GLM-5.2 weights).
-
-Oracle: the native ``kv_a_proj_with_mqa`` call (``dsv3_fused_a_gemm_op``)
-and ``F.linear`` on the indexer's fused ``wk | weights_proj`` weight, on the
-same normalized hidden state.
-
-Contract (frozen before measurement): ``kva`` is
-bit-identical to native for every R = 1..8; ``kw`` is bit-identical or, as
-the named variant, within the frozen budget ``allclose(rtol=0.1,
-atol=1e-2)`` and ``e_fused_kv <= 1.25 e_native + 2^-12`` against the FP64 oracle.
-
-Fixture labels: weights real (checkpoint); activations: native RMSNorm of
-synthetic residual rows with the real gain, scaled copies, plain Gaussian
-rows, and (when mounted) real normalized hidden rows of the captured fixture.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-sys.path.insert(0, os.path.join(_HERE, "..", "absorbed_query"))
-
-import absorbed_query_fixture  # noqa: E402
-import fused_kv_utils as U  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or not absorbed_query_fixture.checkpoint_available(),
-    reason="needs a GPU and the GLM-5.2 checkpoint",
-)
-
-LAYERS = [
-    int(x)
-    for x in os.environ.get(
-        "GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_TEST_LAYERS", "0,1,2,3,78"
-    ).split(",")
-]
-ROWS = [
-    int(x)
-    for x in os.environ.get(
-        "GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_TEST_ROWS", "1,2,3,4,5,6,7,8"
-    ).split(",")
-]
-# Whether kw must be bit-identical (set once the order probe found the native order).
-KW_EXACT = os.environ.get("GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW_EXACT", "0") == "1"
-INPUTS = ("norm", "norm_x1e-3", "norm_x30", "gauss", "real")
-
-
-def _projection(attn, config=None):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        FusedKvProjection,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.projection import (  # noqa: E501
-        projection_eligibility,
-    )
-
-    ok, reason, kw = projection_eligibility(attn)
-    assert ok, reason
-    proj = FusedKvProjection(attn, config, kw_weight=kw)
-    assert proj.supported() is None, proj.supported()
-    return proj
-
-
-def _parse_cfg(text: str, **extra):
-    from dataclasses import replace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection import (  # noqa: E501
-        FusedKvProjConfig,
-    )
-
-    return replace(FusedKvProjConfig.parse(text), **extra)
-
-
-def _input(kind: str, rows: int, seed: int, gamma: torch.Tensor):
-    if kind == "norm":
-        return U.native_hidden(rows, seed, gamma)
-    if kind == "norm_x1e-3":
-        return U.native_hidden(rows, seed, gamma, scale=1e-3)
-    if kind == "norm_x30":
-        return U.native_hidden(rows, seed, gamma, scale=30.0)
-    if kind == "gauss":
-        g = torch.Generator(device="cuda").manual_seed(seed)
-        return torch.randn(rows, U.H, generator=g, device="cuda").bfloat16()
-    return U.real_hidden(rows, offset=97 * seed % 1024)
-
-
-def _check_outputs(attn, h, kva, kw, tag):
-    kva_n, kw_n = U.native(attn, h)
-    torch.cuda.synchronize()
-    assert kva.shape == kva_n.shape and kva.dtype == kva_n.dtype and kva.stride() == kva_n.stride()
-    ok = torch.equal(kva.view(torch.int16), kva_n.view(torch.int16))
-    if not ok:
-        stats = U.compare(kva, kva_n, U.fp64_oracle(h, attn.kv_a_proj_with_mqa.weight))
-        pytest.fail(f"kva not bit-identical {tag}: {stats}")
-    if kw_n is None:
-        assert kw is None
-        return None
-    assert kw.shape == kw_n.shape and kw.dtype == kw_n.dtype and kw.stride() == kw_n.stride()
-    stats = U.compare(kw, kw_n, U.fp64_oracle(h, attn.mqa.indexer._fused_wk_wp_weight))
-    print(f"GLM_MOE_DSA_FUSED_KV_INDEXER_PROJECTION_KW {tag} {stats}")
-    if KW_EXACT:
-        assert stats["bit_identical"], (tag, stats)
-    else:
-        assert U.within_budget(stats), (tag, stats)
-    return stats
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-def test_served_kernel_supported_and_compiled(layer):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.proj_kernel import (  # noqa: E501
-        is_compiled,
-    )
-
-    attn, _ = U.module(layer)
-    proj = _projection(attn)
-    proj.precompile()
-    assert is_compiled(**proj.kernel_config())
-    assert proj.n1 == (0 if attn.mqa.indexer is None else 160)
-
-
-@pytest.mark.parametrize("kind", INPUTS)
-@pytest.mark.parametrize("rows", ROWS)
-@pytest.mark.parametrize("layer", LAYERS)
-def test_kva_bit_identical_kw_contract(layer, rows, kind):
-    attn, gamma = U.module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        h = _input(kind, rows, seed=1000 + 10 * layer + rows, gamma=gamma)
-        if h is None:
-            pytest.skip("real hidden-state fixture not mounted")
-        kva, kw = proj.run(h)
-        _check_outputs(attn, h, kva, kw, f"layer={layer} rows={rows} input={kind}")
-
-
-@pytest.mark.parametrize("layer", LAYERS)
-def test_rows_match_single_row_runs(layer):
-    """Every token row equals a launch of that row alone (padded or other
-    rows of the batch never change a live row)."""
-    attn, gamma = U.module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        h = U.native_hidden(8, seed=11, gamma=gamma)
-        kva, kw = proj.run(h)
-        for t in range(8):
-            kva1, kw1 = proj.run(h[t : t + 1].contiguous())
-            assert torch.equal(kva1[0], kva[t])
-            if kw is not None:
-                assert torch.equal(kw1[0], kw[t])
-
-
-@pytest.mark.parametrize("layer", LAYERS[:2] + LAYERS[-1:])
-def test_graph_replay_changing_rows(layer):
-    """Graphs captured at R = 6 and R = 1 (the draft-loop shapes) replayed
-    interleaved with new inputs equal eager launches."""
-    attn, gamma = U.module(layer)
-    proj = _projection(attn)
-    with torch.inference_mode():
-        static_inputs, graphs, outs = {}, {}, {}
-        for rows in (6, 1):
-            static_inputs[rows] = U.native_hidden(rows, seed=rows, gamma=gamma)
-            proj.run(static_inputs[rows])
-            torch.cuda.synchronize()
-            graphs[rows] = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graphs[rows]):
-                outs[rows] = proj.run(static_inputs[rows])
-        for trial in range(6):
-            rows = (6, 1)[trial % 2]
-            static_inputs[rows].copy_(U.native_hidden(rows, seed=300 + trial, gamma=gamma))
-            for _ in range(1 + 20 * (trial == 4)):
-                graphs[rows].replay()
-            e_kva, e_kw = proj.run(static_inputs[rows])
-            torch.cuda.synchronize()
-            g_kva, g_kw = outs[rows]
-            assert torch.equal(g_kva, e_kva)
-            if e_kw is not None:
-                assert torch.equal(g_kw, e_kw)
-            _check_outputs(
-                attn, static_inputs[rows], g_kva, g_kw, f"graph layer={layer} rows={rows}"
-            )
-
-
-@pytest.mark.parametrize(
-    "cfg", ["b_load=cg", "b_load=ca", "copy_method=tma,b_load=ca", "prewait_chunks=1"]
-)
-@pytest.mark.parametrize("stage_before_wait", [True, False])
-@pytest.mark.parametrize("delay_ns", [0, 5000, 50000])
-@pytest.mark.parametrize("layer", LAYERS[:1] + LAYERS[-1:])
-def test_pdl_delayed_producer(layer, delay_ns, stage_before_wait, cfg):
-    """The producer releases its dependents at entry and writes the
-    activation only after ``delay_ns``; the projection (PDL, staging weights
-    before its dependency wait) must read the final activation. The same
-    activation buffer is rewritten every trial, so an activation line kept in
-    L1 from the previous trial (``b_load=ca``) would show up as a mismatch."""
-    attn, gamma = U.module(layer)
-    proj = _projection(attn, _parse_cfg(cfg, stage_before_wait=stage_before_wait))
-    if not proj.use_pdl:
-        pytest.skip("PDL disabled (TRTLLM_ENABLE_PDL=0)")
-    with torch.inference_mode():
-        dst = torch.empty(6, U.H, dtype=torch.bfloat16, device="cuda")
-        for trial in range(6):
-            src = U.native_hidden(6, seed=500 + trial, gamma=gamma)
-            dst.fill_(float("nan") if trial % 2 else 0.0)
-            torch.cuda.synchronize()
-            U.proxy_producer(src, dst, delay_ns)
-            kva, kw = proj.run(dst)
-            ref_kva, ref_kw = U.native(attn, src)
-            torch.cuda.synchronize()
-            assert torch.equal(kva, ref_kva), f"trial {trial}"
-            if kw is not None:
-                r_kva, r_kw = proj.run(src)
-                torch.cuda.synchronize()
-                assert torch.equal(kw, r_kw), f"trial {trial}"
-
-
-@pytest.mark.parametrize(
-    "cfg",
-    [
-        "stage_before_wait=0",
-        "trigger=end",
-        "trigger=none",
-        "nsub=8",
-        "nsub=2",
-        "b_dist=4",
-        "b_dist=24",
-        "b_dist=48",
-        "a_dist=4",
-        "ctas=148",
-        "ctas=150",
-        "l2_prefetch=0",
-        "producer_warp=0",
-        "copy_method=cpasync,producer_warps=4",
-        "copy_method=cpasync,producer_warps=2,nsub=8",
-        "b_load=ca",
-        "chunks=48-24-12-12",
-        "chunks=40-32-16-8,b_load=ca",
-        "chunks=64-32,copy_method=cpasync,producer_warps=4",
-        "copy_method=tma",
-        "copy_method=tma,b_load=ca,nsub=8",
-        "copy_method=tma,nsub=24",
-        "copy_method=tma,tma_groups=3,chunks=48-24-12-12",
-        "copy_method=tma,producer_warp=0",
-        "copy_method=tma,stage_before_wait=0,b_load=ca",
-        "copy_method=tma,ctas=148",
-        "b_load=cg",
-        "nsub=4,b_load=cg",
-        "prewait_chunks=1",
-        "prewait_chunks=2,copy_method=tma",
-        "prewait_chunks=0",
-        "prewait_chunks=3,copy_method=cpasync,producer_warps=4",
-    ],
-)
-@pytest.mark.parametrize("layer", LAYERS[:1] + LAYERS[3:4])
-def test_launch_configs_are_numerically_identical(cfg, layer):
-    """Launch-configuration knobs change the schedule only, never a bit."""
-    attn, gamma = U.module(layer)
-    proj = _projection(attn, _parse_cfg(cfg))
-    ref = _projection(attn)
-    with torch.inference_mode():
-        for rows in (1, 6, 8):
-            h = U.native_hidden(rows, seed=900 + rows, gamma=gamma)
-            a = proj.run(h)
-            b = ref.run(h)
-            torch.cuda.synchronize()
-            assert torch.equal(a[0], b[0])
-            if b[1] is not None:
-                assert torch.equal(a[1], b[1])
-
-
-def test_corunning_kernels_progress():
-    """The projection on one stream while another stream keeps native GEMMs
-    of the same weights in flight: every launch completes and is correct (no
-    inter-CTA dependency, no co-residency assumption)."""
-    attn, gamma = U.module(LAYERS[0])
-    proj = _projection(attn)
-    side = torch.cuda.Stream()
-    with torch.inference_mode():
-        hs = [U.native_hidden(6, seed=700 + i, gamma=gamma) for i in range(8)]
-        refs = [U.native(attn, h) for h in hs]
-        torch.cuda.synchronize()
-        side.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(side):
-            for _ in range(40):
-                for h in hs:
-                    U.native(attn, h)
-        outs = [proj.run(h) for h in hs for _ in range(5)]
-        torch.cuda.synchronize()
-        for i, (kva, kw) in enumerate(outs):
-            r_kva, r_kw = refs[i // 5]
-            assert torch.equal(kva, r_kva)
-            if r_kw is not None:
-                assert U.within_budget(
-                    U.compare(
-                        kw, r_kw, U.fp64_oracle(hs[i // 5], attn.mqa.indexer._fused_wk_wp_weight)
-                    )
-                )
-
-
-def test_supported_rejects_out_of_domain():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.fused_kv_indexer_projection.proj_kernel import (  # noqa: E501
-        projection_supported,
-    )
-
-    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4) is None
-    assert projection_supported(6144, 2624, 152, 4, 18, 4, 4) is None
-    assert projection_supported(6144, 2784, 152, 4, 17, 4, 4) is not None  # 2 direct rows
-    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4, (48, 24, 12, 12)) is None
-    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4, (48, 24, 12)) is not None
-    assert projection_supported(6144, 2784, 152, 4, 18, 4, 4, (), 160, "tma") is None
-    assert projection_supported(6144, 2624, 152, 4, 18, 4, 4, (), 0, "tma") is None
-    assert projection_supported(6144, 2784, 152, 4, 18, 4, 24, (), 160, "tma", 3) is not None
-    assert projection_supported(6144, 2784, 152, 4, 18, 4, 8, (), 160, "tma", 3) is None
-    assert projection_supported(6144, 2784, 100, 4, 18, 4, 4, (), 160, "tma") is not None
-    assert projection_supported(6144, 2784, 100, 4, 18, 4, 4) is not None  # 28 rows per CTA
-    assert projection_supported(6144, 2784, 152, 4, 18, 3, 4) is not None  # slices vs warps
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py
deleted file mode 100644
index 45bfe5fbdbed..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_host_policy.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# 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.
-"""host_page_retention host policy: locating and setting mimalloc's purge delay in libc10.so.
-
-CPU only. Tests that need PyTorch's statically linked mimalloc skip on builds
-without it.
-"""
-
-import ctypes
-import json
-import os
-import subprocess
-import sys
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.host_page_retention import (  # noqa: E501
-    host_policy as hp,
-)
-
-
-def _options_or_skip():
-    try:
-        return hp.MimallocOptions(hp.torch_libc10_path())
-    except hp.HostPolicyUnavailable as exc:
-        pytest.skip(f"no mimalloc option table in this PyTorch build: {exc}")
-
-
-def _run(code, env_extra):
-    env = {k: v for k, v in os.environ.items() if not k.startswith("MIMALLOC_")}
-    env.update(env_extra)
-    out = subprocess.run(
-        [sys.executable, "-c", code],
-        env=env,
-        check=True,
-        capture_output=True,
-        text=True,
-        cwd="/tmp",
-        timeout=300,
-    )
-    return json.loads(out.stdout.strip().splitlines()[-1])
-
-
-_READ = """
-import json, torch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.host_page_retention import host_policy as hp
-opts = hp.MimallocOptions(hp.torch_libc10_path())
-before = opts.get("purge_delay")
-result = None
-if {apply}:
-    result = hp.apply_purge_policy()
-print(json.dumps({{"before": before, "after": opts.get("purge_delay"), "result": result}}))
-"""  # noqa: E501
-
-
-def test_option_table_is_found_and_consistent():
-    opts = _options_or_skip()
-    assert opts.names.count(hp.PURGE_DELAY_OPTION) == 1
-    assert len(opts.names) == len(set(opts.names))
-    for name in opts.names:
-        opts.get(name)  # mi_option_get agrees with the table for every option
-
-
-def test_set_takes_effect_and_restores():
-    opts = _options_or_skip()
-    old = opts.get(hp.PURGE_DELAY_OPTION)
-    try:
-        opts.set(hp.PURGE_DELAY_OPTION, 1234)
-        assert opts.get(hp.PURGE_DELAY_OPTION) == 1234
-    finally:
-        opts.set(hp.PURGE_DELAY_OPTION, old)
-    assert opts.get(hp.PURGE_DELAY_OPTION) == old
-
-
-def test_runtime_policy_equals_environment_setting():
-    _options_or_skip()
-    env_policy = _run(_READ.format(apply=False), {"MIMALLOC_PURGE_DELAY": str(hp.PURGE_DELAY_MS)})
-    served = _run(_READ.format(apply=False), {"MIMALLOC_PURGE_DELAY": "0"})
-    applied = _run(_READ.format(apply=True), {"MIMALLOC_PURGE_DELAY": "0"})
-    assert env_policy["before"] == hp.PURGE_DELAY_MS
-    assert served["before"] == 0 and served["after"] == 0
-    assert applied["result"] == [0, hp.PURGE_DELAY_MS]
-    assert applied["after"] == env_policy["before"]
-
-
-def test_already_set_is_left_alone():
-    _options_or_skip()
-    out = _run(_READ.format(apply=True), {"MIMALLOC_PURGE_DELAY": str(hp.PURGE_DELAY_MS)})
-    assert out["result"] == [hp.PURGE_DELAY_MS, hp.PURGE_DELAY_MS]
-
-
-def test_bad_descriptor_table_is_rejected():
-    opts = _options_or_skip()
-    good = ctypes.string_at(opts._table, opts._count * hp._DESC.size)
-    buf = ctypes.create_string_buffer(good, len(good))
-    value, init, option, name, legacy = hp._DESC.unpack_from(buf.raw, hp._DESC.size)
-    hp._DESC.pack_into(buf, hp._DESC.size, value, init, option + 7, name, legacy)
-    fake = object.__new__(hp.MimallocOptions)
-    fake._table = ctypes.addressof(buf)
-    fake._count = opts._count
-    with pytest.raises(hp.HostPolicyUnavailable, match="not a mimalloc option table"):
-        fake._check_table()
-
-
-@pytest.mark.parametrize("content", [b"", b"not an elf file at all" * 4])
-def test_non_elf_library_is_rejected(tmp_path, content):
-    path = tmp_path / "libc10.so"
-    path.write_bytes(content)
-    with pytest.raises(hp.HostPolicyUnavailable):
-        hp.MimallocOptions(str(path))
-
-
-def test_missing_library_is_rejected():
-    with pytest.raises(hp.HostPolicyUnavailable, match="not found"):
-        hp.MimallocOptions("/nonexistent/libc10.so")
-
-
-def test_library_not_loaded_is_rejected(tmp_path):
-    # A copy of libc10.so is an ELF file with the symbols but is not loaded
-    # under that path: the policy must not dlopen a second copy.
-    _options_or_skip()
-    copy = tmp_path / "libc10.so"
-    copy.write_bytes(open(hp.torch_libc10_path(), "rb").read())
-    with pytest.raises(hp.HostPolicyUnavailable, match="not loaded"):
-        hp.MimallocOptions(str(copy))
-
-
-def test_install_declines_without_side_effect(monkeypatch):
-    opts = _options_or_skip()
-    old = opts.get(hp.PURGE_DELAY_OPTION)
-    logged = []
-    monkeypatch.setattr(hp, "torch_libc10_path", lambda: "/nonexistent/libc10.so")
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(dispatch, "_log_info", logged.append)
-    assert hp.install_host_policy("test") is False
-    assert opts.get(hp.PURGE_DELAY_OPTION) == old
-    assert logged and "host_page_retention/host_policy: native fallback" in logged[0]
-
-
-def test_install_engages_and_logs(monkeypatch):
-    opts = _options_or_skip()
-    old = opts.get(hp.PURGE_DELAY_OPTION)
-    logged = []
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    monkeypatch.setattr(dispatch, "_log_info", logged.append)
-    try:
-        assert hp.install_host_policy("test") is True
-        assert opts.get(hp.PURGE_DELAY_OPTION) == hp.PURGE_DELAY_MS
-        assert "host_page_retention/host_policy: engaged" in logged[0]
-        assert f"{old} -> {hp.PURGE_DELAY_MS} ms" in logged[0]
-    finally:
-        opts.set(hp.PURGE_DELAY_OPTION, old)
-
-
-def test_package_registers_entry_point():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
-        host_page_retention,
-    )
-
-    assert (
-        dispatch.get_impl("host_page_retention", "host_policy")
-        is host_page_retention.install_host_policy
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py
deleted file mode 100644
index a4e3d588ca78..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/host_page_retention/test_host_pages_selector.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# 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.
-"""host_page_retention selection and call site.
-
-A plan runs ``host_page_retention`` with its closure, ``post_attention_hop``, and turning that
-off turns it off; the target's executor events install the host policy once after warmup,
-after the post-MoE hop health check.
-"""
-
-import os
-import sys
-import types
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
-    dispatch,
-    runtime,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import plan_testlib  # noqa: E402
-
-B1 = (
-    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,draft_front,router_gemm,"
-    "post_attention_hop,draft_tail,sparse_mla_decode,draft_absorbed_query,rope_kv_append,"
-    "draft_loop,shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
-    "attention_output_tail_prefetch,attention_output_into_hop,routed_moe_into_hop,deferred_norm,"
-    "dense_mlp,deferred_input_prep"
-)
-
-
-def test_closure_is_attn_hop():
-    assert dispatch.REQUIRED_FEATURES["host_page_retention"] == frozenset({"post_attention_hop"})
-    assert "host_page_retention" not in dispatch.NUMERICAL_VARIANTS  # allocator policy only
-
-
-def test_selected_with_adapter(monkeypatch):
-    paths = frozenset(B1.split(",")) | {"host_page_retention"}
-    monkeypatch.setattr(dispatch, "_log_info", lambda message: None)
-    assert plan_testlib.select(dispatch, paths, monkeypatch) == paths
-    assert dispatch.is_enabled("host_page_retention")
-    plan = dispatch.resolve(
-        dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: "post_attention_hop"}
-    )
-    assert not plan.enabled("host_page_retention")
-
-
-def test_the_host_policy_is_installed_after_warmup_after_the_hop_health_check(monkeypatch):
-    calls = []
-    kernels = runtime.__name__.rsplit(".", 1)[0] + ".kernels"
-    monkeypatch.setitem(
-        sys.modules,
-        kernels + ".post_moe_hop",
-        types.SimpleNamespace(check_health=lambda *args, **kwargs: calls.append(("health", args))),
-    )
-    monkeypatch.setitem(
-        sys.modules,
-        kernels + ".host_page_retention",
-        types.SimpleNamespace(install_host_policy=lambda where: calls.append(("policy", where))),
-    )
-    paths = {"post_moe_hop", "host_page_retention"}
-    plan = types.SimpleNamespace(enabled=lambda name: name in paths)
-    events = runtime.GlmExecutorEvents(plan, types.SimpleNamespace(enable_attention_dp=False))
-
-    events.after_warmup()
-    assert calls == [("health", ("after_warmup",)), ("policy", "after_warmup")]
-    paths.discard("post_moe_hop")
-    calls.clear()
-    events.after_warmup()
-    assert calls == [("policy", "after_warmup")]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py
deleted file mode 100644
index adcb534b1114..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/test_indexer_front_adapter.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# 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.
-"""indexer_layer_front selector and call sites.
-
-* the plan: indexer_layer_front runs only with its closure (absorbed_query, indexer_topk,
-  draft_absorbed_query, rope_kv_append, cross_rank_topk) and has no numerical variant;
-* the registered entry points and their call sites in the target (the absorbed-query packet,
-  the cross-rank top-k region and the owned ``forward_absorption_generation``);
-* the ``skip_reused_iq`` selector (0 / 1, anything else fails before any GPU
-  work) and the reuse condition it acts on, checked against the indexer's
-  own source;
-* deferred_input_prep's source map lists indexer_layer_front's directories and side selector.
-"""
-
-import inspect
-import os
-import sys
-
-import pytest
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import plan_testlib  # noqa: E402
-
-CLOSURE = (
-    "absorbed_query",
-    "indexer_topk",
-    "draft_absorbed_query",
-    "rope_kv_append",
-    "cross_rank_topk",
-)
-
-
-@pytest.fixture
-def dispatch(monkeypatch):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
-        dispatch as d,
-    )
-
-    return d
-
-
-def _select(monkeypatch, d, features):
-    monkeypatch.setattr(d, "_log_info", lambda message: None)
-    return plan_testlib.select(d, features, monkeypatch)
-
-
-def test_selector_tables(dispatch):
-    assert "indexer_layer_front" in dispatch.SELECTOR_PATHS
-    assert dispatch.REQUIRED_FEATURES["indexer_layer_front"] == frozenset(CLOSURE)
-    assert "indexer_layer_front" not in dispatch.NUMERICAL_VARIANTS
-
-
-def test_closure_enforced(dispatch, monkeypatch):
-    """Turning off any path of the closure turns indexer_layer_front off with it."""
-    monkeypatch.setattr(dispatch, "_require", [False])
-    for base in CLOSURE:
-        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
-        assert not plan.enabled("indexer_layer_front"), base
-    assert "indexer_layer_front" in _select(
-        monkeypatch, dispatch, CLOSURE + ("indexer_layer_front",)
-    )
-    assert "indexer_layer_front" in _select(
-        monkeypatch, dispatch, CLOSURE + ("cross_rank_topk_cluster_merge", "indexer_layer_front")
-    )
-
-
-def test_entry_points_registered(dispatch):
-    import tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front  # noqa: E501, F401
-
-    for name in ("dsa_proj", "rope_available", "mla_rope_generation", "cp_indexer_decode"):
-        assert callable(dispatch.get_impl("indexer_layer_front", name)), name
-
-
-def test_call_sites_patched():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-        attention_forward,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        region,
-    )
-
-    assert 'get_impl("indexer_layer_front", "dsa_proj")' in inspect.getsource(
-        attention_forward._absorbed_query_packet
-    )
-    src = inspect.getsource(region.cross_rank_topk_decode)
-    assert src.index('"indexer_layer_front", "cp_indexer_decode"') < src.index(
-        '"cross_rank_topk_cluster_merge", "cp_indexer_decode"'
-    )
-    assert '"indexer_layer_front", "rope_available"' in inspect.getsource(
-        attention_forward.forward_absorption_generation
-    )
-
-
-def test_skip_reused_iq_is_off_in_this_target():
-    """Every runner starts with the option off and the entry never turns it on."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
-        query_proj,
-    )
-
-    init = inspect.getsource(query_proj.IndexerFrontQueryProjection.__init__)
-    assert "self.skip_reused_iq = False" in init
-    assert "skip_reused_iq" not in inspect.getsource(query_proj._runner)
-
-
-def test_reuses_shared_topk_condition():
-    """The option fires only under the indexer's own ``reuse_topk`` condition."""
-    from types import SimpleNamespace
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
-        query_proj,
-    )
-
-    def mla(share):
-        return SimpleNamespace(mqa=SimpleNamespace(indexer=SimpleNamespace(mtp_index_share=share)))
-
-    def md(loop, skip, rows):
-        return SimpleNamespace(
-            in_mtp_draft_loop=loop, indexer_skip_topk=skip, shared_topk_indices=rows
-        )
-
-    assert query_proj.reuses_shared_topk(mla(True), md(True, True, object()))
-    assert not query_proj.reuses_shared_topk(mla(False), md(True, True, object()))
-    assert not query_proj.reuses_shared_topk(mla(True), md(False, True, object()))
-    assert not query_proj.reuses_shared_topk(mla(True), md(True, False, object()))
-    assert not query_proj.reuses_shared_topk(mla(True), md(True, True, None))
-
-
-def _conjuncts(node):
-    import ast
-
-    if isinstance(node, ast.BoolOp) and isinstance(node.op, ast.And):
-        return [c for v in node.values for c in _conjuncts(v)]
-    return [node]
-
-
-def _normalized(node):
-    """``getattr(obj, "x", default)`` -> ``obj.x``; the indexer and metadata names unified."""
-    import ast
-
-    roles = {"self": "indexer", "metadata": "metadata", "attn_metadata": "metadata"}
-
-    class Norm(ast.NodeTransformer):
-        def visit_Call(self, call):
-            self.generic_visit(call)
-            if isinstance(call.func, ast.Name) and call.func.id == "getattr":
-                return ast.Attribute(value=call.args[0], attr=call.args[1].value, ctx=ast.Load())
-            return call
-
-        def visit_Name(self, name):
-            return ast.Name(id=roles.get(name.id, name.id), ctx=ast.Load())
-
-    return ast.unparse(Norm().visit(node))
-
-
-def test_reuses_shared_topk_matches_indexer_source():
-    """``reuses_shared_topk`` is term for term the indexer's ``reuse_topk``.
-
-    The branch the indexer takes under it reads none of ``q_fp8`` / ``q_scale`` / ``weights``
-    (the outputs the ``skip_reused_iq`` option leaves unwritten).
-    """
-    import ast
-    import textwrap
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
-        query_proj,
-    )
-    from tensorrt_llm._torch.attention.backends.sparse.dsa import indexer as indexer_mod
-
-    tree = ast.parse(textwrap.dedent(inspect.getsource(indexer_mod.Indexer.sparse_attn_indexer)))
-    assigns = [
-        n
-        for n in ast.walk(tree)
-        if isinstance(n, ast.Assign)
-        and any(isinstance(t, ast.Name) and t.id == "reuse_topk" for t in n.targets)
-    ]
-    assert len(assigns) == 1, "the indexer assigns reuse_topk once"
-    indexer_terms = sorted(_normalized(c) for c in _conjuncts(assigns[0].value))
-
-    fn = ast.parse(textwrap.dedent(inspect.getsource(query_proj.reuses_shared_topk)))
-    aliases = {
-        ast.unparse(s.targets[0]): ast.unparse(s.value)
-        for s in fn.body[0].body
-        if isinstance(s, ast.Assign)
-    }
-    assert aliases == {"indexer": "mla.mqa.indexer"}, aliases
-    ret = next(n for n in ast.walk(fn) if isinstance(n, ast.Return))
-    value = ret.value.args[0] if isinstance(ret.value, ast.Call) else ret.value  # bool(...)
-    indexer_front_terms = sorted(_normalized(c) for c in _conjuncts(value))
-    assert indexer_front_terms == indexer_terms, (indexer_front_terms, indexer_terms)
-
-    branches = [
-        n
-        for n in ast.walk(tree)
-        if isinstance(n, ast.If)
-        and isinstance(n.test, ast.BoolOp)
-        and ast.unparse(n.test.values[-1]) == "reuse_topk"
-    ]
-    assert len(branches) == 1
-    names = {n.id for s in branches[0].body for n in ast.walk(s) if isinstance(n, ast.Name)}
-    assert not names & {"q_fp8", "q_scale", "weights"}, names
-
-
-def test_entry_declines_without_feature(dispatch, monkeypatch):
-    """With indexer_layer_front off the entry returns None before touching the module."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
-        query_proj,
-    )
-
-    _select(monkeypatch, dispatch, [])
-
-    class Boom:
-        def __getattr__(self, name):
-            raise AssertionError(f"module touched: {name}")
-
-    assert not dispatch.is_enabled("indexer_layer_front")
-    assert query_proj.dsa_proj(Boom(), None, None, None) is None
-
-
-@pytest.mark.skipif(
-    not os.environ.get("GLM_MOE_DSA_INDEXER_LAYER_FRONT_JIT_LIB"),
-    reason="needs the development build of the indexer_layer_front op",
-)
-def test_rope_op_available():
-    import torch
-
-    import tensorrt_llm  # noqa: F401
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_layer_front import (  # noqa: E501
-        rope,
-    )
-
-    if not rope.rope_available():
-        torch.ops.load_library(os.environ["GLM_MOE_DSA_INDEXER_LAYER_FRONT_JIT_LIB"])
-    assert rope.rope_available()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py
deleted file mode 100644
index 57926839977b..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_cross_rank_cp_policy.py
+++ /dev/null
@@ -1,195 +0,0 @@
-# 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.
-"""Eligibility of the cross_rank_topk context-partitioned indexer region (no GPU needed).
-
-``cross_rank_policy.py`` and ``policy.py`` are loaded from the indexer_topk package directory
-without executing the package ``__init__`` (which imports torch).
-"""
-
-import importlib
-import sys
-import types
-from pathlib import Path
-
-import pytest
-
-_INDEXER_TOPK_DIR = (
-    Path(__file__).resolve().parents[9]
-    / "tensorrt_llm"
-    / "_torch"
-    / "_experimental"
-    / "modeling_v2"
-    / "models"
-    / "glm_moe_dsa"
-    / "shared"
-    / "kernels"
-    / "indexer_topk"
-)
-
-
-def _load_cp_policy():
-    name = "_cross_rank_policy_pkg"
-    if name not in sys.modules:
-        pkg = types.ModuleType(name)
-        pkg.__path__ = [str(_INDEXER_TOPK_DIR)]
-        sys.modules[name] = pkg
-    return importlib.import_module(f"{name}.cross_rank_policy")
-
-
-cp_policy = _load_cp_policy()
-
-SERVED_MAX_SEQ_LEN = (1 << 20) + 14
-# Largest per-rank share of the served envelope at world 8, and the kernel limit on sm_103.
-LOCAL_CAPACITY_DRAFT_LOOP = 131136
-MAX_LOCAL_CAPACITY = 179712
-
-_SERVED_TP8 = dict(
-    world_size=8,
-    op_available=True,
-    local_capacity=LOCAL_CAPACITY_DRAFT_LOOP,
-    max_local_capacity=MAX_LOCAL_CAPACITY,
-    allow_draft_step=False,
-    sm_version=103,
-    cute_dsl_available=True,
-    native_uses_cute_dsl_logits=True,
-    use_fp4_indexer=False,
-    use_gvr_emission=False,
-    needs_gvr_prior=False,
-    is_ragged_verify=False,
-    expanded_mtp_buffers=False,
-    in_mtp_draft_loop=False,
-    attention_dp=False,
-    topk_decode_impl="cute_dsl_gvr",
-    topk_self_sampling=True,
-    index_topk=2048,
-    compress_ratio=1,
-    num_heads=32,
-    head_dim=128,
-    query_dtype="torch.float8_e4m3fn",
-    query_shape=(6, 32, 128),
-    query_contiguous=True,
-    weights_dtype="torch.float32",
-    weights_shape=(6, 32),
-    weights_contiguous=True,
-    k_cache_dtype="torch.uint8",
-    k_cache_shape=(16384, 64, 1, 132),
-    block_table_cols=16388,
-    max_seq_len=SERVED_MAX_SEQ_LEN,
-    num_contexts=0,
-    token_offset=0,
-    num_gen_tokens=6,
-    num_generations=1,
-    uniform_verify_width=True,
-)
-
-
-def test_tp8_decode_is_eligible():
-    d = cp_policy.decide_cp(**_SERVED_TP8)
-    assert d.eligible, d.reason
-
-
-@pytest.mark.parametrize("world", [4, 8])
-def test_supported_world_sizes(world):
-    assert cp_policy.decide_cp(**dict(_SERVED_TP8, world_size=world)).eligible
-
-
-FALLBACKS = {
-    "two_requests": dict(
-        num_generations=2, num_gen_tokens=12, query_shape=(12, 32, 128), weights_shape=(12, 32)
-    ),
-    "world_2": dict(world_size=2),
-    "world_16": dict(world_size=16),
-    "op_missing": dict(op_available=False),
-    "local_width_above_kernel_limit": dict(local_capacity=MAX_LOCAL_CAPACITY + 64),
-    "mtp_draft_step": dict(in_mtp_draft_loop=True),
-    "attention_dp": dict(attention_dp=True),
-    "indexer_topk_envelope_page32": dict(k_cache_shape=(32768, 32, 1, 132)),
-    "indexer_topk_envelope_next_n4": dict(
-        num_gen_tokens=4, query_shape=(4, 32, 128), weights_shape=(4, 32)
-    ),
-    "indexer_topk_envelope_max_seq_len": dict(max_seq_len=SERVED_MAX_SEQ_LEN + 1),
-    "indexer_topk_envelope_sm100": dict(sm_version=100),
-    "indexer_topk_envelope_fp4": dict(use_fp4_indexer=True),
-    "indexer_topk_envelope_ragged": dict(is_ragged_verify=True),
-    "indexer_topk_envelope_mixed_batch": dict(num_contexts=1),
-}
-
-
-@pytest.mark.parametrize("name", sorted(FALLBACKS))
-def test_native_fallback_cases(name):
-    d = cp_policy.decide_cp(**dict(_SERVED_TP8, **FALLBACKS[name]))
-    assert not d.eligible
-    assert d.reason
-
-
-def test_draft_step_only_when_allowed():
-    inputs = dict(_SERVED_TP8, in_mtp_draft_loop=True)
-    assert not cp_policy.decide_cp(**inputs).eligible
-    assert cp_policy.decide_cp(**dict(inputs, allow_draft_step=True)).eligible
-
-
-# The rank-local facts of a decision: a rank's build and device.
-_LOCAL_KEYS = ("op_available", "max_local_capacity", "sm_version", "cute_dsl_available")
-
-
-def _uniform(inputs):
-    return {k: v for k, v in inputs.items() if k not in _LOCAL_KEYS}
-
-
-def _local(inputs):
-    return dict(
-        op_available=inputs["op_available"],
-        local_capacity=inputs["local_capacity"],
-        max_local_capacity=inputs["max_local_capacity"],
-        sm_version=inputs["sm_version"],
-        cute_dsl_available=inputs["cute_dsl_available"],
-    )
-
-
-@pytest.mark.parametrize("name", sorted(FALLBACKS))
-def test_uniform_and_local_parts_reproduce_decide_cp(name):
-    """``decide_cp`` = rank-uniform decision AND no rank-local reason.
-
-    The adapter decides a forward with ``decide_cp_uniform`` (identical on
-    every rank) and hands ``rank_local_reason`` to the collective workspace
-    creation, so a fact of one rank's build or device keeps every rank native
-    instead of moving that rank alone out of the exchange.
-    """
-    inputs = dict(_SERVED_TP8, **FALLBACKS[name])
-    full = cp_policy.decide_cp(**inputs)
-    uniform = cp_policy.decide_cp_uniform(**_uniform(inputs))
-    local = cp_policy.rank_local_reason(**_local(inputs))
-    assert full.eligible == (uniform.eligible and not local), (full, uniform, local)
-
-
-@pytest.mark.parametrize(
-    "fact, value",
-    [
-        ("op_available", False),
-        ("sm_version", 100),
-        ("cute_dsl_available", False),
-        ("max_local_capacity", LOCAL_CAPACITY_DRAFT_LOOP - 64),
-    ],
-)
-def test_rank_local_facts_never_decide_a_forward(fact, value):
-    """A rank-local fact never changes the per-forward decision.
-
-    A rank whose own build or device differs makes the same decision as its
-    peers; only its local reason differs.
-    """
-    inputs = dict(_SERVED_TP8, **{fact: value})
-    assert cp_policy.decide_cp_uniform(**_uniform(inputs)).eligible
-    assert cp_policy.rank_local_reason(**_local(inputs))
-    assert not cp_policy.rank_local_reason(**_local(_SERVED_TP8))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py
deleted file mode 100644
index c1472647f66c..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_indexer_decode.py
+++ /dev/null
@@ -1,1081 +0,0 @@
-# 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.
-"""GPU correctness of the indexer_topk indexer decode package against native TRT-LLM.
-
-Native oracle: the CuTe-DSL FP8 paged-MQA logits launch exactly as
-``Indexer.sparse_attn_indexer`` issues it for the GEN DSL path (wave-aware
-row-group split from ``_pick_dsl_expand`` with expanded lengths, page tables
-and schedule), followed by the same ``TopK`` module. Most fixtures are
-synthetic (random FP8 keys with UE8M0 scales, random queries and weights).
-``test_real_weight_fixture`` uses a real GLM-5.2 layer-0 fixture when
-``GLM_MOE_DSA_INDEXER_TOPK_REAL_FIXTURE`` points at it.
-"""
-
-import os
-import sys
-from types import SimpleNamespace
-from unittest.mock import Mock, patch
-
-import pytest
-import torch
-
-from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
-from tensorrt_llm._utils import get_sm_version
-
-if not torch.cuda.is_available() or not IS_CUTLASS_DSL_AVAILABLE:
-    pytest.skip("CUDA and the CUTLASS DSL are required", allow_module_level=True)
-if get_sm_version() not in (100, 103):
-    pytest.skip("indexer_topk targets SM100/SM103", allow_module_level=True)
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-    GlmIndexer,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-    indexer_decode as indexer_topk,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk.logits_runner import (  # noqa: E501
-    DEFAULT_CONFIG,
-    LogitsConfig,
-)
-from tensorrt_llm._torch.attention.backends.sparse.dsa import indexer as native_indexer_mod
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _pick_dsl_expand
-from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation
-from tensorrt_llm.deep_gemm import get_num_sms, get_paged_mqa_logits_metadata
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-
-HEADS = 32
-HEAD_DIM = 128
-BLOCK_KV = 64
-TOP_K = 2048
-NEXT_N = 6
-SCHEDULE_BLOCK_KV = 64
-FP8_ATOMS = (1, 2, 3, 4)
-_DEV = torch.device("cuda")
-
-
-def _ue8m0(x):
-    return torch.pow(2.0, torch.ceil(torch.log2(x.clamp_min(1e-12))))
-
-
-def make_pool(num_blocks, gen, k_values=None):
-    """Fused FP8 indexer K pages ``[num_blocks, 64, 1, 132]`` plus their parts."""
-    if k_values is None:
-        k_values = torch.randn(num_blocks, BLOCK_KV, HEAD_DIM, device=_DEV, generator=gen)
-    scale = _ue8m0(k_values.abs().amax(dim=-1) / 448.0)
-    k_fp8 = (k_values / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
-    data = k_fp8.view(torch.uint8).reshape(num_blocks, BLOCK_KV * HEAD_DIM)
-    scales = scale.float().contiguous().view(torch.uint8).reshape(num_blocks, BLOCK_KV * 4)
-    fused = (
-        torch.cat([data, scales], dim=1).contiguous().view(num_blocks, BLOCK_KV, 1, HEAD_DIM + 4)
-    )
-    return fused, k_fp8, scale.float()
-
-
-def make_top_k():
-    return TopK(
-        TOP_K, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=True
-    )
-
-
-class Case:
-    """B requests x NEXT_N verify rows over one layer's paged FP8 K cache."""
-
-    def __init__(
-        self,
-        kv_lens,
-        seed=0,
-        max_seq_len=1 << 20,
-        shuffle=True,
-        q_scale=1.0,
-        weight_sign=None,
-        distinct_pages=None,
-        block_table_cols=None,
-        boost_from_token=None,
-    ):
-        """Build the case.
-
-        Args:
-            block_table_cols: Block-table width; defaults to
-                ``ceil(max_seq_len / 64)``.
-            boost_from_token: When set, the pages holding logical tokens at or
-                beyond this position get 64x larger keys (so 64x larger
-                per-token scales and identical FP8 values), which puts those
-                tokens into the top-k when the weights are positive.
-        """
-        gen = torch.Generator(device=_DEV).manual_seed(seed)
-        self.B = len(kv_lens)
-        self.kv_lens_list = list(kv_lens)
-        self.max_seq_len = max_seq_len
-        blocks = [max(1, -(-n // BLOCK_KV)) for n in kv_lens]
-        max_blocks = block_table_cols or -(-max_seq_len // BLOCK_KV)
-        assert max(blocks) <= max_blocks
-        num_blocks = (distinct_pages or sum(blocks)) + 4
-        self.pool, self.k_fp8, self.k_scale = make_pool(num_blocks, gen)
-        bt = torch.zeros(self.B, max_blocks, dtype=torch.int32, device=_DEV)
-        base = 0
-        for b, nb in enumerate(blocks):
-            if distinct_pages:
-                pages = torch.randint(0, distinct_pages, (nb,), device=_DEV, generator=gen)
-            else:
-                pages = torch.arange(base, base + nb, device=_DEV)
-                if shuffle:
-                    pages = pages[torch.randperm(nb, device=_DEV, generator=gen)]
-                base += nb
-            bt[b, :nb] = pages.to(torch.int32)
-        self.block_table = bt
-        if boost_from_token is not None:
-            k_values = torch.randn(num_blocks, BLOCK_KV, HEAD_DIM, device=_DEV, generator=gen)
-            first_page = boost_from_token // BLOCK_KV
-            for b, nb in enumerate(blocks):
-                if nb > first_page:
-                    k_values[bt[b, first_page:nb].long()] *= 64.0
-            self.pool, self.k_fp8, self.k_scale = make_pool(num_blocks, gen, k_values)
-        self.kv_lens = torch.tensor(kv_lens, dtype=torch.int32, device=_DEV)
-        q = torch.randn(self.B, NEXT_N, HEADS, HEAD_DIM, device=_DEV, generator=gen) * q_scale
-        self.q = q.to(torch.float8_e4m3fn)
-        w = torch.randn(self.B * NEXT_N, HEADS, device=_DEV, generator=gen)
-        if weight_sign is not None:
-            w = w.abs() * weight_sign
-        self.weights = w.float()
-        self.num_sms = get_num_sms()
-        self.schedule = get_paged_mqa_logits_metadata(
-            self.kv_lens.view(-1, 1), SCHEDULE_BLOCK_KV, self.num_sms
-        )
-        self.factor, self.atom = _pick_dsl_expand(
-            NEXT_N,
-            batch_size=self.B,
-            max_ctx=max(kv_lens),
-            num_sms=self.num_sms,
-            kernel_atoms=FP8_ATOMS,
-        )
-        self.kv_lens_exp = self.kv_lens.repeat_interleave(self.factor)
-        self.block_table_exp = self.block_table.repeat_interleave(self.factor, dim=0)
-        self.schedule_exp = get_paged_mqa_logits_metadata(
-            self.kv_lens_exp.view(-1, 1), SCHEDULE_BLOCK_KV, self.num_sms
-        )
-        rows = self.B * NEXT_N
-        self.radix_idx = torch.empty(rows, 10, TOP_K, dtype=torch.int32, device=_DEV)
-        self.radix_val = torch.empty(rows, 10, TOP_K, dtype=torch.float32, device=_DEV)
-
-    @classmethod
-    def from_fixture(cls, fixture, kv_len, shuffle, seed=0):
-        """Real-weight case: the fixture's pages, queries and weights."""
-        case = cls([kv_len], seed=seed)
-        gen = torch.Generator(device=_DEV).manual_seed(seed)
-        nb = -(-kv_len // BLOCK_KV)
-        pages = fixture["pool"][:nb].to(_DEV)
-        perm = torch.arange(nb, device=_DEV)
-        if shuffle:
-            perm = perm[torch.randperm(nb, device=_DEV, generator=gen)]
-        pool = torch.zeros(nb + 4, BLOCK_KV, 1, HEAD_DIM + 4, dtype=torch.uint8, device=_DEV)
-        pool[perm] = pages
-        case.pool = pool
-        flat = pool.view(nb + 4, -1)
-        case.k_fp8 = flat[:, : BLOCK_KV * HEAD_DIM].contiguous().view(torch.float8_e4m3fn)
-        case.k_fp8 = case.k_fp8.view(nb + 4, BLOCK_KV, HEAD_DIM)
-        case.k_scale = flat[:, BLOCK_KV * HEAD_DIM :].contiguous().view(torch.float32)
-        case.block_table.zero_()
-        case.block_table[0, :nb] = perm.to(torch.int32)
-        case.block_table_exp = case.block_table.repeat_interleave(case.factor, dim=0)
-        case.q = fixture["q"][kv_len].to(_DEV).view(1, NEXT_N, HEADS, HEAD_DIM)
-        case.weights = fixture["weights"][kv_len].to(_DEV).float()
-        return case
-
-    @property
-    def row_lens(self):
-        return [
-            max(0, self.kv_lens_list[r // NEXT_N] - NEXT_N + r % NEXT_N + 1)
-            for r in range(self.B * NEXT_N)
-        ]
-
-    def native_logits(self):
-        """Native GEN DSL launch, including the wave-aware row-group split."""
-        if self.factor > 1:
-            q = self.q.reshape(self.B * self.factor, self.atom, HEADS, HEAD_DIM)
-            return torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits(
-                q,
-                self.pool,
-                self.weights,
-                self.kv_lens_exp,
-                self.block_table_exp,
-                self.schedule_exp,
-                self.max_seq_len,
-            )
-        return torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits(
-            self.q,
-            self.pool,
-            self.weights,
-            self.kv_lens,
-            self.block_table,
-            self.schedule,
-            self.max_seq_len,
-        )
-
-    def indexer_topk_logits(self, config=None):
-        return indexer_topk.paged_mqa_logits(
-            self.q,
-            self.pool,
-            self.weights,
-            self.kv_lens,
-            self.block_table,
-            self.schedule,
-            self.max_seq_len,
-            config=config,
-        )
-
-    def native_region(self, top_k, out):
-        top_k(
-            self.native_logits(),
-            out,
-            is_prefill=False,
-            sequence_lengths=self.kv_lens,
-            scan_lengths=self.kv_lens,
-            row_kv_lens=None,
-            next_n=NEXT_N,
-            max_seq_len=self.max_seq_len,
-            gvr_ext_kwargs=None,
-            radix_aux_indices=self.radix_idx,
-            radix_aux_logits=self.radix_val,
-        )
-        return out
-
-    def indexer_topk_region(self, top_k, out, selector=None, config=None):
-        return indexer_topk.decode_topk(
-            q=self.q,
-            weights=self.weights,
-            k_cache=self.pool,
-            block_table=self.block_table,
-            context_lens=self.kv_lens,
-            kv_lens=self.kv_lens,
-            schedule_meta=self.schedule,
-            max_seq_len=self.max_seq_len,
-            top_k=top_k,
-            out=out,
-            radix_aux_indices=self.radix_idx,
-            radix_aux_logits=self.radix_val,
-            logits_config=config,
-            selector=selector,
-        )
-
-    def new_out(self, fill=-7):
-        return torch.full((self.B * NEXT_N, TOP_K), fill, dtype=torch.int32, device=_DEV)
-
-
-def assert_logits_equal(case, native, cand):
-    for r, n in enumerate(case.row_lens):
-        if n > 0:
-            assert torch.equal(native[r, :n], cand[r, :n]), f"row {r}: logits differ"
-
-
-def assert_selection_contract(case, logits, sel):
-    """Native the index_selection contract semantics, tie-aware (value multiset vs torch.topk)."""
-    for r, n in enumerate(case.row_lens):
-        idx = sel[r].long()
-        if n <= 0:
-            assert bool((idx == -1).all()), f"row {r}: empty row must be all -1"
-        elif n <= TOP_K:
-            assert torch.equal(torch.sort(idx[:n]).values, torch.arange(n, device=_DEV)), r
-            assert bool((idx[n:] == -1).all()), f"row {r}: short row tail must be -1"
-        else:
-            assert int(idx.min()) >= 0 and int(idx.max()) < n, f"row {r}: index out of range"
-            assert int(torch.unique(idx).numel()) == TOP_K, f"row {r}: duplicate index"
-            ref = torch.topk(logits[r, :n], TOP_K).values
-            got = torch.sort(logits[r].gather(0, idx) + 0.0, descending=True).values
-            assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), (
-                f"row {r}: selected value multiset is not the exact top-k"
-            )
-
-
-def reference_logits(case):
-    """Independent fp32 oracle: scale_j * sum_h w_h * relu(q_h . k_j)."""
-    k_all = case.k_fp8.float()
-    out = torch.full(
-        (case.B * NEXT_N, max(case.kv_lens_list)), float("nan"), device=_DEV, dtype=torch.float32
-    )
-    for b, n in enumerate(case.kv_lens_list):
-        nb = -(-n // BLOCK_KV)
-        pages = case.block_table[b, :nb].long()
-        k = k_all[pages].reshape(-1, HEAD_DIM)[:n]
-        s = case.k_scale[pages].reshape(-1)[:n]
-        qk = torch.matmul(case.q[b].float().reshape(NEXT_N * HEADS, HEAD_DIM), k.t())
-        qk = torch.relu(qk).view(NEXT_N, HEADS, n)
-        w = case.weights[b * NEXT_N : (b + 1) * NEXT_N]
-        out[b * NEXT_N : (b + 1) * NEXT_N, :n] = (w.unsqueeze(-1) * qk).sum(dim=1) * s
-    return out
-
-
-# --- logits ---------------------------------------------------------------
-
-LENGTHS = [3, 6, 7, 2047, 2048, 2049, 2054, 32768, 131072, 491061, (1 << 20) - 1]
-
-
-# None runs DEFAULT_CONFIG (rows, 8 register-cached weights), the kernel the
-# adapter hook ships; the rowsplit variants are opt-in alternatives.
-VARIANTS = {
-    "default_rows_draft_loop": None,
-    "rowsplit": LogitsConfig(variant="rowsplit"),
-    "rowsplit_wg3": LogitsConfig(variant="rowsplit", num_math_wg=3),
-}
-
-
-def test_default_config_is_rows_draft_loop():
-    assert DEFAULT_CONFIG == LogitsConfig(variant="rows", num_w_in_reg=8)
-
-
-@pytest.mark.parametrize("variant", sorted(VARIANTS))
-@pytest.mark.parametrize("kv_len", LENGTHS)
-def test_logits_bit_exact_vs_native(kv_len, variant):
-    case = Case([kv_len], seed=kv_len)
-    native = case.native_logits()[:, :kv_len].clone()
-    cand = case.indexer_topk_logits(VARIANTS[variant])[:, :kv_len]
-    assert_logits_equal(case, native, cand)
-
-
-@pytest.mark.parametrize("variant", sorted(VARIANTS))
-@pytest.mark.parametrize(
-    "kv_lens",
-    [[4224, 64], [491061, 300], [7, 131072, 2049], [1, 2, 3, 4]],
-    ids=lambda v: "-".join(map(str, v)),
-)
-def test_logits_bit_exact_multi_request(kv_lens, variant):
-    # Odd and even tile counts per request and schedule splits that end at
-    # request boundaries.
-    case = Case([max(n, NEXT_N) for n in kv_lens], seed=sum(kv_lens))
-    width = max(case.kv_lens_list)
-    native = case.native_logits()[:, :width].clone()
-    cand = case.indexer_topk_logits(VARIANTS[variant])[:, :width]
-    assert_logits_equal(case, native, cand)
-
-
-@pytest.mark.parametrize("kv_len", [2049, 32768, 131072])
-def test_logits_vs_reference(kv_len):
-    case = Case([kv_len], seed=7 * kv_len)
-    cand = case.indexer_topk_logits()[:, :kv_len]
-    ref = reference_logits(case)
-    for r, n in enumerate(case.row_lens):
-        # Same tolerances as the native CuTe-DSL FP8 paged-MQA logits test.
-        torch.testing.assert_close(cand[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
-
-
-# --- selection --------------------------------------------------------------
-
-
-SELECTORS = ["native", "radix"]
-
-
-@pytest.mark.parametrize("selector", SELECTORS)
-@pytest.mark.parametrize("kv_len", LENGTHS)
-def test_selection_matches_native(kv_len, selector):
-    case = Case([kv_len], seed=11 + kv_len)
-    top_k = make_top_k()
-    native = case.native_region(top_k, case.new_out()).clone()
-    logits = case.native_logits()[:, :kv_len].clone()
-    cand = case.indexer_topk_region(top_k, case.new_out(), selector=selector)
-    assert_selection_contract(case, logits, cand)
-    # Continuous random logits: the top-k set is unique, so any exact
-    # selection engine must return the native set.
-    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
-
-
-@pytest.mark.parametrize("selector", SELECTORS)
-@pytest.mark.parametrize(
-    "kv_lens",
-    [[30000, 64], [491061, 7, 2048, 2055], [131072, 131073]],
-    ids=lambda v: "-".join(map(str, v)),
-)
-def test_selection_multi_request(kv_lens, selector):
-    case = Case(kv_lens, seed=sum(kv_lens) + 1)
-    top_k = make_top_k()
-    width = max(kv_lens)
-    logits = case.native_logits()[:, :width].clone()
-    native = case.native_region(top_k, case.new_out()).clone()
-    cand = case.indexer_topk_region(top_k, case.new_out(), selector=selector)
-    assert_selection_contract(case, logits, cand)
-    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
-
-
-@pytest.mark.parametrize("selector", SELECTORS)
-@pytest.mark.parametrize("kind", ["duplicate_pages", "all_zero", "negative_weights"])
-def test_ties(kind, selector):
-    kv_len = 40000
-    if kind == "duplicate_pages":
-        # 8 distinct pages repeated across the context: every logit value
-        # occurs ~78 times, so the top-k boundary falls inside a tie class.
-        case = Case([kv_len], seed=3, distinct_pages=8)
-    elif kind == "all_zero":
-        case = Case([kv_len], seed=4, q_scale=0.0)
-    else:
-        case = Case([kv_len], seed=5, weight_sign=-1.0)
-    top_k = make_top_k()
-    native = case.native_region(top_k, case.new_out()).clone()
-    logits = case.native_logits()[:, :kv_len].clone()
-    cand = case.indexer_topk_region(top_k, case.new_out(), selector=selector)
-    assert_logits_equal(case, logits, case.indexer_topk_logits()[:, :kv_len])
-    assert_selection_contract(case, logits, cand)
-    # Characterize tie handling (all-zero logits overflow the radix
-    # candidate buffer and exercise its full-row fallback).
-    same_ids = torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
-    print(f"[ties:{kind}:{selector}] identical index sets vs native: {same_ids}")
-
-
-def test_padded_batch_rows():
-    # Request 1 stands in for a CUDA-graph padding request with a short cache.
-    case = Case([30000, 64], seed=21)
-    top_k = make_top_k()
-    native = case.native_region(top_k, case.new_out()).clone()
-    logits = case.native_logits()[:, :30000].clone()
-    cand = case.indexer_topk_region(top_k, case.new_out())
-    assert_logits_equal(case, logits, case.indexer_topk_logits()[:, :30000])
-    assert_selection_contract(case, logits, cand)
-    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
-
-
-_FIXTURE_PATH = os.environ.get("GLM_MOE_DSA_INDEXER_TOPK_REAL_FIXTURE", "")
-
-
-@pytest.mark.skipif(
-    not os.path.isfile(_FIXTURE_PATH), reason="GLM_MOE_DSA_INDEXER_TOPK_REAL_FIXTURE not provided"
-)
-@pytest.mark.parametrize("shuffle", [False, True], ids=["contiguous", "shuffled"])
-def test_real_weight_fixture(shuffle):
-    fixture = torch.load(_FIXTURE_PATH, map_location="cpu", weights_only=False)
-    print(f"[real] {fixture['label']} sha={fixture['token_ids_sha256'][:16]}")
-    top_k = make_top_k()
-    for kv_len in fixture["kv_lens"]:
-        case = Case.from_fixture(fixture, kv_len, shuffle)
-        logits = case.native_logits()[:, :kv_len].clone()
-        assert_logits_equal(case, logits, case.indexer_topk_logits()[:, :kv_len])
-        native = case.native_region(top_k, case.new_out()).clone()
-        cand = case.indexer_topk_region(top_k, case.new_out())
-        assert_selection_contract(case, logits, cand)
-        assert torch.equal(torch.sort(native, dim=1).values, torch.sort(cand, dim=1).values)
-        split = LogitsConfig(variant="rowsplit")
-        assert_logits_equal(case, logits, case.indexer_topk_logits(split)[:, :kv_len])
-        radix = case.indexer_topk_region(top_k, case.new_out(), selector="radix", config=split)
-        assert_selection_contract(case, logits, radix)
-        assert torch.equal(torch.sort(native, dim=1).values, torch.sort(radix, dim=1).values)
-        if kv_len <= 131072:
-            ref = reference_logits(case)
-            for r, n in enumerate(case.row_lens):
-                torch.testing.assert_close(logits[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
-        # Boundary sensitivity of the real distribution: gap between the k-th
-        # and (k+1)-th largest logits, and overlap of adjacent verify rows.
-        for r, n in enumerate(case.row_lens):
-            vals = torch.topk(logits[r, :n], TOP_K + 1).values
-            gap = (vals[TOP_K - 1] - vals[TOP_K]).item()
-            overlap = (
-                torch.isin(cand[r], cand[r - 1]).float().mean().item() if r > 0 else float("nan")
-            )
-            print(
-                f"[real] kv={kv_len} row={r} kth={vals[TOP_K - 1].item():.6g} "
-                f"gap={gap:.3g} overlap_prev_row={overlap:.4f}"
-            )
-
-
-# --- CUDA graph replay --------------------------------------------------------
-
-
-GRAPH_CONFIGS = {"default_rows_draft_loop": None, "rowsplit": LogitsConfig(variant="rowsplit")}
-GRAPH_STEPS = [20000, 20006, 20012, 35000, 9000, 65530]
-
-
-def _graph_step_lens(kv_len, batch):
-    # Request 0 follows the step length; later requests are shorter, down to
-    # a 6-token request (verify rows of 1..6 keys, padding-like).
-    return [max(NEXT_N, kv_len - 4099 * b) for b in range(batch)]
-
-
-def _native_oracle(lens, max_seq_len, pool, k_fp8, k_scale, block_table, q, weights, kv_lens):
-    """Eager native launch state for the given tensors (no new pool)."""
-    batch = len(lens)
-    ref = Case([NEXT_N] * batch, seed=0, max_seq_len=max_seq_len)
-    ref.kv_lens_list = list(lens)
-    ref.pool, ref.k_fp8, ref.k_scale = pool, k_fp8, k_scale
-    ref.block_table, ref.q, ref.weights, ref.kv_lens = block_table, q, weights, kv_lens
-    ref.schedule = get_paged_mqa_logits_metadata(kv_lens.view(-1, 1), 64, ref.num_sms)
-    ref.factor, ref.atom = _pick_dsl_expand(
-        NEXT_N, batch_size=batch, max_ctx=max(lens), num_sms=ref.num_sms, kernel_atoms=FP8_ATOMS
-    )
-    ref.kv_lens_exp = kv_lens.repeat_interleave(ref.factor)
-    ref.block_table_exp = block_table.repeat_interleave(ref.factor, dim=0)
-    ref.schedule_exp = get_paged_mqa_logits_metadata(ref.kv_lens_exp.view(-1, 1), 64, ref.num_sms)
-    return ref
-
-
-def _run_graph_replay(batch, logits_config, selector, max_seq_len, steps, block_table_cols=None):
-    """Capture indexer_topk once, then replay with new lengths, pages, tables and queries.
-
-    After every replay the indexer_topk logits (read from the graph-stable arena the
-    captured launch wrote) must be bit-exact to an eager native launch on the
-    same state, and the selection must match the native region.
-    """
-    case = Case(
-        _graph_step_lens(steps[0], batch),
-        seed=31,
-        max_seq_len=max_seq_len,
-        block_table_cols=block_table_cols,
-    )
-    # Persistent pool sized for the largest step; pages are rewritten between
-    # replays to model page reuse by later requests.
-    total_blocks = batch * -(-max_seq_len // BLOCK_KV) + 4
-    gen = torch.Generator(device=_DEV).manual_seed(99 + batch)
-    case.pool, case.k_fp8, case.k_scale = make_pool(total_blocks, gen)
-    top_k = make_top_k()
-    out = case.new_out()
-    q, weights = case.q.clone(), case.weights.clone()
-    kv_lens, block_table = case.kv_lens.clone(), case.block_table.clone()
-    schedule, pool = case.schedule.clone(), case.pool
-    arena = {}
-    indexer_topk_logits_fn = indexer_topk.paged_mqa_logits
-
-    def recording_logits(*args, **kwargs):
-        arena["logits"] = indexer_topk_logits_fn(*args, **kwargs)
-        return arena["logits"]
-
-    def indexer_topk_call():
-        return indexer_topk.decode_topk(
-            q=q,
-            weights=weights,
-            k_cache=pool,
-            block_table=block_table,
-            context_lens=kv_lens,
-            kv_lens=kv_lens,
-            schedule_meta=schedule,
-            max_seq_len=max_seq_len,
-            top_k=top_k,
-            out=out,
-            radix_aux_indices=case.radix_idx,
-            radix_aux_logits=case.radix_val,
-            logits_config=logits_config,
-            selector=selector,
-        )
-
-    with patch.object(indexer_topk, "paged_mqa_logits", side_effect=recording_logits):
-        stream = torch.cuda.Stream()
-        stream.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(stream):
-            indexer_topk_call()
-        torch.cuda.current_stream().wait_stream(stream)
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            indexer_topk_call()
-    captured_logits = arena["logits"]
-    assert captured_logits.shape == (batch * NEXT_N, max_seq_len)
-
-    for step, kv_len in enumerate(steps):
-        lens = _graph_step_lens(kv_len, batch)
-        new_pool, new_k, new_s = make_pool(total_blocks, gen)
-        pool.copy_(new_pool)
-        del new_pool
-        perm = torch.randperm(total_blocks, device=_DEV, generator=gen).to(torch.int32)
-        block_table.zero_()
-        base = 0
-        for b, n in enumerate(lens):
-            nb = -(-n // BLOCK_KV)
-            block_table[b, :nb] = perm[base : base + nb]
-            base += nb
-        kv_lens.copy_(torch.tensor(lens, dtype=torch.int32))
-        schedule.copy_(get_paged_mqa_logits_metadata(kv_lens.view(-1, 1), 64, case.num_sms))
-        q.copy_((torch.randn(q.shape, device=_DEV, generator=gen)).to(torch.float8_e4m3fn))
-        weights.copy_(torch.randn(weights.shape, device=_DEV, generator=gen))
-        out.fill_(-7)
-        graph.replay()
-        torch.cuda.synchronize()
-        ref = _native_oracle(
-            lens, max_seq_len, pool, new_k, new_s, block_table, q, weights, kv_lens
-        )
-        width = max(lens)
-        logits = ref.native_logits()[:, :width].clone()
-        assert_logits_equal(ref, logits, captured_logits[:, :width])
-        native = ref.native_region(make_top_k(), ref.new_out())
-        assert_selection_contract(ref, logits, out)
-        assert torch.equal(torch.sort(native, dim=1).values, torch.sort(out, dim=1).values), step
-
-
-@pytest.mark.parametrize("selector", SELECTORS)
-@pytest.mark.parametrize("config_name", sorted(GRAPH_CONFIGS))
-@pytest.mark.parametrize("batch", [1, 4])
-def test_cuda_graph_replay_growing_context_and_page_reuse(batch, config_name, selector):
-    _run_graph_replay(batch, GRAPH_CONFIGS[config_name], selector, 65536, GRAPH_STEPS)
-
-
-# --- served indexer max_seq_len (2**20 + 14) --------------------------------------
-
-# Served GEN: net max_seq_len 2**20, MTP K = 5, overlap scheduler on. The
-# indexer max_seq_len is 2**20 + 5 + 4 + 5 and KV cache manager V2 sizes the
-# block table at round_up(ceil((2**20 + 14 + 10) / 64), 4) columns; manager V1
-# would use ceil((2**20 + 14) / 64) columns.
-SERVED_MAX_SEQ_LEN = (1 << 20) + 14
-V2_BLOCK_TABLE_COLS = 16388
-V1_BLOCK_TABLE_COLS = 16385
-# Live indexer lengths around 2**20 up to the served bound (target layers
-# reach 2**20 + 10 with the overlap scheduler; the bound itself is included).
-BOUND_LENGTHS = [
-    (1 << 20) - 6,
-    (1 << 20) - 1,
-    1 << 20,
-    (1 << 20) + 1,
-    (1 << 20) + 5,
-    (1 << 20) + 10,
-    (1 << 20) + 13,
-    SERVED_MAX_SEQ_LEN,
-]
-
-
-def _bound_case(kv_lens, seed, **kwargs):
-    return Case(
-        kv_lens,
-        seed=seed,
-        max_seq_len=SERVED_MAX_SEQ_LEN,
-        block_table_cols=V2_BLOCK_TABLE_COLS,
-        **kwargs,
-    )
-
-
-def test_policy_envelope_is_the_served_bound():
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-        policy,
-    )
-
-    assert policy.MAX_INDEXER_SEQ_LEN == SERVED_MAX_SEQ_LEN
-    assert V2_BLOCK_TABLE_COLS % 2 == 0 and V1_BLOCK_TABLE_COLS % 2 == 1
-
-
-@pytest.mark.parametrize("kv_len", BOUND_LENGTHS)
-def test_served_bound_logits_and_selection(kv_len):
-    case = _bound_case([kv_len], seed=kv_len % 9973)
-    native = case.native_logits()[:, :kv_len].clone()
-    cand = case.indexer_topk_logits()
-    assert cand.shape == (NEXT_N, SERVED_MAX_SEQ_LEN)
-    assert cand.stride(0) % 256 == 0 and cand.stride(0) >= SERVED_MAX_SEQ_LEN
-    assert_logits_equal(case, native, cand[:, :kv_len])
-    top_k = make_top_k()
-    native_sel = case.native_region(top_k, case.new_out()).clone()
-    sel = case.indexer_topk_region(top_k, case.new_out())
-    assert_selection_contract(case, native, sel)
-    assert torch.equal(torch.sort(native_sel, dim=1).values, torch.sort(sel, dim=1).values)
-
-
-@pytest.mark.parametrize("kv_len", [(1 << 20) + 1, (1 << 20) + 10, SERVED_MAX_SEQ_LEN])
-def test_served_bound_selects_ids_beyond_2pow20(kv_len):
-    """Keys at the tail dominate, so IDs >= 2**20 must be selected exactly."""
-    boost_from = (1 << 20) - 64
-    case = _bound_case([kv_len], seed=kv_len % 7919, weight_sign=1.0, boost_from_token=boost_from)
-    native = case.native_logits()[:, :kv_len].clone()
-    assert_logits_equal(case, native, case.indexer_topk_logits()[:, :kv_len])
-    top_k = make_top_k()
-    native_sel = case.native_region(top_k, case.new_out()).clone()
-    sel = case.indexer_topk_region(top_k, case.new_out())
-    assert_selection_contract(case, native, sel)
-    assert torch.equal(torch.sort(native_sel, dim=1).values, torch.sort(sel, dim=1).values)
-    beyond = 0
-    for r, n in enumerate(case.row_lens):
-        tail = torch.arange(boost_from, n, device=_DEV, dtype=torch.int32)
-        assert bool(torch.isin(tail, sel[r]).all()), f"row {r}: boosted tail not selected"
-        beyond += int((sel[r] >= (1 << 20)).sum())
-    assert beyond > 0, "no ID >= 2**20 was selected"
-
-
-@pytest.mark.parametrize(
-    "kv_lens",
-    [
-        [SERVED_MAX_SEQ_LEN, (1 << 20) + 1],
-        [SERVED_MAX_SEQ_LEN, (1 << 20) + 10, 1 << 20, (1 << 20) - 1],
-        [(1 << 20) + 7, 491061, NEXT_N, 2049],
-    ],
-    ids=lambda v: "-".join(map(str, v)),
-)
-def test_served_bound_multi_request(kv_lens):
-    case = _bound_case(kv_lens, seed=sum(kv_lens) % 10007)
-    width = max(kv_lens)
-    native = case.native_logits()[:, :width].clone()
-    assert_logits_equal(case, native, case.indexer_topk_logits()[:, :width])
-    top_k = make_top_k()
-    native_sel = case.native_region(top_k, case.new_out()).clone()
-    sel = case.indexer_topk_region(top_k, case.new_out())
-    assert_selection_contract(case, native, sel)
-    assert torch.equal(torch.sort(native_sel, dim=1).values, torch.sort(sel, dim=1).values)
-
-
-def test_served_bound_vs_reference():
-    case = _bound_case([SERVED_MAX_SEQ_LEN], seed=17)
-    cand = case.indexer_topk_logits()[:, :SERVED_MAX_SEQ_LEN]
-    ref = reference_logits(case)
-    for r, n in enumerate(case.row_lens):
-        torch.testing.assert_close(cand[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
-
-
-@pytest.mark.parametrize("launch", ["indexer_topk", "native"])
-@pytest.mark.parametrize(
-    "cols", [V2_BLOCK_TABLE_COLS, V1_BLOCK_TABLE_COLS], ids=["v2_cols", "v1_cols"]
-)
-def test_served_bound_single_launch(launch, cols):
-    """One logits launch at the bound, checked against the fp32 oracle.
-
-    Each parameter runs exactly one kind of kernel launch so the test can run
-    under compute-sanitizer memcheck (with PYTORCH_NO_CUDA_MEMORY_CACHING=1)
-    one parameter at a time. With the odd V1 width the last 128-token tile's
-    second page ID is one entry past the row; the policy declines that width.
-    """
-    case = Case(
-        [SERVED_MAX_SEQ_LEN], seed=23, max_seq_len=SERVED_MAX_SEQ_LEN, block_table_cols=cols
-    )
-    logits = case.indexer_topk_logits() if launch == "indexer_topk" else case.native_logits()
-    torch.cuda.synchronize()
-    ref = reference_logits(case)
-    for r, n in enumerate(case.row_lens):
-        torch.testing.assert_close(logits[r, :n], ref[r, :n], atol=5e-5, rtol=1e-5)
-
-
-SERVED_GRAPH_STEPS = [
-    (1 << 20) - 600,
-    1 << 20,
-    (1 << 20) + 1,
-    SERVED_MAX_SEQ_LEN,
-    491061,
-    (1 << 20) + 10,
-]
-
-
-@pytest.mark.parametrize("batch", [1, 4])
-def test_cuda_graph_replay_across_2pow20(batch):
-    """Shipped kernel + native top-k, captured at the served max_seq_len."""
-    _run_graph_replay(
-        batch,
-        None,
-        "native",
-        SERVED_MAX_SEQ_LEN,
-        SERVED_GRAPH_STEPS,
-        block_table_cols=V2_BLOCK_TABLE_COLS,
-    )
-
-
-# --- adapter hook ---------------------------------------------------------------
-
-
-def _hook_objects(case, **overrides):
-    top_k = make_top_k()
-    indexer = SimpleNamespace(
-        use_cute_dsl_paged_mqa_logits=True,
-        use_fp4=False,
-        use_gvr_emission=False,
-        top_k=top_k,
-        n_heads=HEADS,
-        head_dim=HEAD_DIM,
-        layer_idx=0,
-    )
-    metadata = SimpleNamespace(
-        num_contexts=0,
-        num_generations=case.B,
-        seq_lens=torch.full((case.B,), NEXT_N, dtype=torch.int32),
-        is_ragged_verify=False,
-        in_mtp_draft_loop=False,
-        gen_indexer_kv_lens_cuda_runtime=case.kv_lens,
-        kv_lens_cuda_runtime=case.kv_lens,
-        use_expanded_buffers_for_mtp=False,
-        mapping=SimpleNamespace(enable_attention_dp=False),
-        kv_cache_manager=SimpleNamespace(get_indexer_k_cache_buffers=lambda idx: case.pool),
-        indexer_k_cache_block_offsets=case.block_table,
-        scheduler_metadata_buffer=case.schedule,
-        get_indexer_max_seq_len=lambda: case.max_seq_len,
-        radix_aux_indices=case.radix_idx,
-        radix_aux_logits=case.radix_val,
-    )
-    for key, value in overrides.items():
-        target = indexer if hasattr(indexer, key) else metadata
-        setattr(target, key, value)
-    return indexer, metadata
-
-
-HOOK_LENS = [50000, 4224, 64, 131072]
-
-
-@pytest.mark.parametrize("batch", [1, 2, 3, 4])
-def test_adapter_hook_engages_and_matches_native(batch):
-    case = Case(HOOK_LENS[:batch], seed=41 + batch)
-    indexer, metadata = _hook_objects(case)
-    rows = batch * NEXT_N
-    q_tok = case.q.reshape(rows, HEADS, HEAD_DIM)
-    out = case.new_out()
-    assert indexer_topk.try_indexer_decode(indexer, metadata, q_tok, case.weights, 0, rows, out)
-    native = case.native_region(make_top_k(), case.new_out())
-    assert_selection_contract(case, case.native_logits()[:, : max(HOOK_LENS[:batch])], out)
-    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(out, dim=1).values)
-
-
-@pytest.mark.parametrize("batch", [1, 4])
-def test_adapter_hook_engages_at_served_bound(batch):
-    """Served GEN values: indexer max_seq_len 2**20 + 14, V2 table width."""
-    lens = [SERVED_MAX_SEQ_LEN, (1 << 20) + 1, 1 << 20, (1 << 20) - 1][:batch]
-    case = _bound_case(lens, seed=61 + batch)
-    indexer, metadata = _hook_objects(case)
-    rows = batch * NEXT_N
-    q_tok = case.q.reshape(rows, HEADS, HEAD_DIM)
-    out = case.new_out()
-    engaged_before = indexer_topk.DECISION_COUNTS.get("engaged", 0)
-    assert indexer_topk.try_indexer_decode(indexer, metadata, q_tok, case.weights, 0, rows, out)
-    assert indexer_topk.DECISION_COUNTS["engaged"] == engaged_before + 1
-    logits = case.native_logits()[:, : max(lens)].clone()
-    native = case.native_region(make_top_k(), case.new_out())
-    assert_selection_contract(case, logits, out)
-    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(out, dim=1).values)
-
-
-# Each case: indexer/metadata overrides plus the call shape. Every one must be
-# declined by the policy before any device work.
-_I32 = torch.int32
-FALLBACK_CASES = {
-    "fp4_keys": dict(overrides={"use_fp4": True}),
-    "non_dsl_native_logits": dict(overrides={"use_cute_dsl_paged_mqa_logits": False}),
-    "gvr_emission": dict(overrides={"use_gvr_emission": True}),
-    "ragged_verify": dict(overrides={"is_ragged_verify": True}),
-    "expanded_mtp_buffers": dict(overrides={"use_expanded_buffers_for_mtp": True}),
-    "mtp_draft_step": dict(overrides={"in_mtp_draft_loop": True}),
-    "attention_dp": dict(overrides={"mapping": SimpleNamespace(enable_attention_dp=True)}),
-    "radix_topk_engine": dict(radix_topk=True),
-    "sm100": dict(sm=100),
-    "page16": dict(page=16),
-    "page32": dict(page=32),
-    "page128": dict(page=128),
-    "non_uniform_seq_lens": dict(overrides={"seq_lens": torch.tensor([4], dtype=_I32)}),
-    "next_n4": dict(overrides={"seq_lens": torch.tensor([4], dtype=_I32)}, next_n=4),
-    # Constructor-invalid for the rows kernel (8 cached weights > 4 heads).
-    "heads4_next_n4": dict(
-        overrides={"n_heads": 4, "seq_lens": torch.tensor([4], dtype=_I32)}, heads=4, next_n=4
-    ),
-    "batch5": dict(
-        overrides={"num_generations": 5, "seq_lens": torch.full((5,), NEXT_N, dtype=_I32)},
-        batch=5,
-    ),
-    "mixed_context_batch": dict(
-        overrides={"num_contexts": 1, "seq_lens": torch.tensor([100, NEXT_N], dtype=_I32)}
-    ),
-    "max_seq_len_2m": dict(overrides={"get_indexer_max_seq_len": lambda: 1 << 21}),
-    "max_seq_len_above_served_bound": dict(
-        overrides={"get_indexer_max_seq_len": lambda: SERVED_MAX_SEQ_LEN + 1}
-    ),
-    # KV cache manager V1 at the served max_seq_len: odd block-table width.
-    "block_table_v1_odd_width": dict(
-        overrides={"get_indexer_max_seq_len": lambda: SERVED_MAX_SEQ_LEN},
-        block_table_cols=V1_BLOCK_TABLE_COLS,
-    ),
-    "bf16_weights": dict(weights_dtype=torch.bfloat16),
-    "non_contiguous_query": dict(non_contiguous_q=True),
-}
-
-
-@pytest.mark.parametrize("name", sorted(FALLBACK_CASES))
-def test_adapter_hook_falls_back_before_side_effects(name):
-    spec = FALLBACK_CASES[name]
-    case = Case([50000], seed=43)
-    overrides = dict(spec.get("overrides", {}))
-    if "page" in spec:
-        view = case.pool.view(-1, spec["page"], 1, HEAD_DIM + 4)
-        overrides["kv_cache_manager"] = SimpleNamespace(
-            get_indexer_k_cache_buffers=lambda idx: view
-        )
-    if spec.get("radix_topk"):
-        overrides["top_k"] = TopK(TOP_K, decode_implementation=TopKImplementation.CUDA_RADIX)
-    if "block_table_cols" in spec:
-        overrides["indexer_k_cache_block_offsets"] = torch.zeros(
-            1, spec["block_table_cols"], dtype=torch.int32, device=_DEV
-        )
-    indexer, metadata = _hook_objects(case, **overrides)
-    heads = spec.get("heads", HEADS)
-    rows = spec.get("next_n", NEXT_N) * spec.get("batch", 1)
-    q_wide = torch.zeros(rows, heads, 2 * HEAD_DIM, device=_DEV).to(torch.float8_e4m3fn)
-    q_tok = q_wide[..., :HEAD_DIM]
-    if not spec.get("non_contiguous_q"):
-        q_tok = q_tok.contiguous()
-    weights = torch.zeros(rows, heads, device=_DEV, dtype=spec.get("weights_dtype", torch.float32))
-    out = torch.full((rows, TOP_K), -7, dtype=torch.int32, device=_DEV)
-    torch.cuda.synchronize()
-    decisions_before = sum(indexer_topk.DECISION_COUNTS.values())
-    engaged_before = indexer_topk.DECISION_COUNTS.get("engaged", 0)
-    with patch.object(
-        indexer_topk, "get_sm_version", return_value=spec.get("sm", get_sm_version())
-    ):
-        with patch.object(indexer_topk, "decode_topk") as launched:
-            assert not indexer_topk.try_indexer_decode(
-                indexer, metadata, q_tok, weights, 0, rows, out
-            )
-    torch.cuda.synchronize()
-    assert not launched.called, "indexer_topk must decline before building or launching a kernel"
-    assert bool((out == -7).all()), "fallback must leave the output untouched"
-    assert sum(indexer_topk.DECISION_COUNTS.values()) == decisions_before + 1
-    assert indexer_topk.DECISION_COUNTS.get("engaged", 0) == engaged_before
-
-
-def test_adapter_hook_temporal_gvr_prior_falls_back():
-    case = Case([50000], seed=45)
-    indexer, metadata = _hook_objects(case)
-    indexer.top_k = TopK(
-        TOP_K, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=False
-    )
-    out = case.new_out(fill=-7)
-    q_tok = case.q.reshape(NEXT_N, HEADS, HEAD_DIM)
-    assert not indexer_topk.try_indexer_decode(
-        indexer, metadata, q_tok, case.weights, 0, NEXT_N, out
-    )
-    assert bool((out == -7).all())
-
-
-# --- the target's indexer against the built-in Indexer ------------------------------
-
-
-def _make_indexer(module):
-    from tensorrt_llm._torch.attention.backends.interface import (
-        PositionalEmbeddingParams,
-        RopeParams,
-    )
-    from tensorrt_llm.functional import PositionEmbeddingType
-    from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
-
-    sparse_params = DeepSeekSparseAttentionConfig(
-        index_head_dim=HEAD_DIM,
-        index_n_heads=HEADS,
-        index_topk=TOP_K,
-        use_cute_dsl_topk=True,
-        use_cute_dsl_paged_mqa_logits=True,
-        enable_heuristic_topk=True,
-        index_share_for_mtp_iteration=True,
-    ).to_sparse_params()
-    pos = PositionalEmbeddingParams(
-        type=PositionEmbeddingType.rope_gpt_neox,
-        rope=RopeParams(dim=64, theta=10000.0, max_positions=4096),
-        is_neox=True,
-    )
-    mla = SimpleNamespace(hidden_size=6144, q_lora_rank=2048, qk_rope_head_dim=64)
-    with patch.object(module, "RotaryEmbedding", return_value=Mock()):
-        return module.Indexer(
-            quant_config=None,
-            pos_embd_params=pos,
-            mla_params=mla,
-            skip_create_weights_in_init=True,
-            sparse_params=sparse_params,
-            dtype=torch.bfloat16,
-            layer_idx=0,
-        )
-
-
-def _decode_metadata(case, in_mtp_draft_loop, num_accepted=None):
-    rows = case.B * NEXT_N
-    return SimpleNamespace(
-        kv_cache_manager=SimpleNamespace(
-            quant_block_size=128,
-            get_indexer_k_cache_buffers=lambda idx: case.pool,
-            layer_offsets={0: 0},
-        ),
-        num_contexts=0,
-        num_generations=case.B,
-        num_ctx_tokens=0,
-        num_tokens=rows,
-        get_empty=lambda buffers, shape, cache_name, dtype, capture_graph: torch.empty(
-            shape, dtype=dtype, device=_DEV
-        ),
-        cuda_graph_buffers=None,
-        is_cuda_graph=False,
-        skip_indexer_for_ctx_reqs=False,
-        skip_indexer_for_gen_reqs=False,
-        gvr_prior_indices=None,
-        in_mtp_draft_loop=in_mtp_draft_loop,
-        indexer_skip_topk=False,
-        shared_topk_indices=None,
-        mtp_num_accepted=num_accepted,
-        seq_lens=torch.full((case.B,), NEXT_N, dtype=torch.int32),
-        is_ragged_verify=False,
-        use_expanded_buffers_for_mtp=False,
-        kv_lens_cuda_2d=case.kv_lens.view(-1, 1).expand(-1, NEXT_N).contiguous(),
-        indexer_k_cache_block_offsets=case.block_table,
-        scheduler_metadata_buffer=case.schedule,
-        scheduler_metadata_buffer_full_next_n=case.schedule,
-        gen_indexer_kv_lens_cuda_runtime=case.kv_lens,
-        get_indexer_max_seq_len=lambda: case.max_seq_len,
-        dsl_expand_factor=case.factor,
-        dsl_atom=case.atom,
-        kv_lens_expanded_cuda=case.kv_lens_exp,
-        block_table_expanded=case.block_table_exp,
-        scheduler_metadata_buffer_expanded=case.schedule_exp,
-        kv_lens_row_reorder=None,
-        kv_lens_cuda_runtime=case.kv_lens,
-        ragged_row_kv_lens=lambda n: None,
-        radix_aux_indices=case.radix_idx,
-        radix_aux_logits=case.radix_val,
-        _tokens_per_block=BLOCK_KV,
-    )
-
-
-@pytest.mark.parametrize("in_mtp_draft_loop", [False, True], ids=["target", "draft_loop"])
-def test_adopted_indexer_preserves_native_indexer_semantics(monkeypatch, in_mtp_draft_loop):
-    """Stock projections -> the target's ``GlmIndexer.sparse_attn_indexer`` in a decode step
-    (indexer_topk on) == the built-in indexer.
-
-    Covers the target forward (indexer_topk engaged) and an MTP draft step (indexer_topk declines;
-    the native path computes the selection and stashes each request's last
-    accepted row for the reuse steps).
-    """
-    case = Case([60000, 45000], seed=51)
-    rows = case.B * NEXT_N
-    q_tok = case.q.reshape(rows, HEADS, HEAD_DIM)
-    k_dummy = torch.zeros(rows, HEAD_DIM, device=_DEV).to(torch.float8_e4m3fn)
-    k_scale_dummy = torch.ones(rows, 1, device=_DEV)
-    hidden = torch.zeros(rows, 6144, device=_DEV, dtype=torch.bfloat16)
-    accepted = torch.tensor([3, 6], dtype=torch.int32, device=_DEV)
-
-    native_indexer = _make_indexer(native_indexer_mod)
-    meta_native = _decode_metadata(case, in_mtp_draft_loop, accepted)
-    native = native_indexer.sparse_attn_indexer(
-        meta_native, hidden, q_tok, k_dummy, k_scale_dummy, case.weights, is_generation=True
-    ).clone()
-
-    with monkeypatch.context() as mp:
-        plan_testlib.select(dispatch, "indexer_topk", mp)
-        indexer_topk_indexer = GlmIndexer.adopt(_make_indexer(native_indexer_mod))
-        meta_indexer_topk = _decode_metadata(case, in_mtp_draft_loop, accepted)
-        with patch.object(
-            indexer_topk, "decode_topk", wraps=indexer_topk.decode_topk
-        ) as indexer_topk_launch:
-            with dispatch.step(dispatch.Phase.DECODE):
-                got = indexer_topk_indexer.sparse_attn_indexer(
-                    meta_indexer_topk,
-                    hidden,
-                    q_tok,
-                    k_dummy,
-                    k_scale_dummy,
-                    case.weights,
-                    is_generation=True,
-                )
-
-    # indexer_topk engages for the target forward and declines in the MTP draft step.
-    assert indexer_topk_launch.call_count == (0 if in_mtp_draft_loop else 1)
-    assert torch.equal(torch.sort(native, dim=1).values, torch.sort(got, dim=1).values)
-    if in_mtp_draft_loop:
-        assert meta_native.shared_topk_indices is not None
-        # The native engine's output order within a row varies run to run.
-        assert torch.equal(
-            torch.sort(meta_native.shared_topk_indices, dim=1).values,
-            torch.sort(meta_indexer_topk.shared_topk_indices, dim=1).values,
-        )
-    else:
-        assert meta_indexer_topk.shared_topk_indices is None
-
-
-def test_without_a_plan_indexer_topk_is_off(monkeypatch):
-    monkeypatch.setattr(dispatch, "_INSTALLED", [None])
-    assert not dispatch.is_enabled("indexer_topk")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py
deleted file mode 100644
index c5300f6c1a1f..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_pipeline.py
+++ /dev/null
@@ -1,1188 +0,0 @@
-# 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.
-"""indexer_topk inside the real DSA decode pipeline.
-
-Every object on the path to and from indexer_topk is the native one:
-
-- ``DSACacheManagerV2`` holding three GLM-5.2-shaped MLA layers: layer 0 owns
-  a full indexer, layer 1 is a shared-indexer layer, layer 2 is the MTP layer
-  (full indexer). FP8 MLA cache, FP8 indexer K, 64-token pages and the served
-  ``max_seq_len`` (2**20 + 14).
-- ``MLA`` modules with the per-rank GLM-5.2 TP8 attention shapes (8 heads,
-  q-nope 192, q-rope 64, v 256, kv-LoRA 512, q-LoRA 2048; indexer 32 x 128,
-  top-k 2048), adopted as the target's ``GlmMLA`` (with its DSA backend and
-  indexer), running the target's ``_forward_dsa_attn`` as a decode step:
-  indexer K append, the indexer
-  (native or indexer_topk), the backend's shared-layer copy, the local-to-global remap,
-  RoPE/FP8-query/KV append and sparse MLA (native, or sparse_mla_decode/rope_kv_append when
-  selected).
-- ``DSAtrtllmAttentionMetadata`` with ``update_spec_dec_param(max_draft_len=5)``
-  and ``prepare()`` per decode step, followed by the decode runner's
-  ``_preprocess_inputs`` sequence: ``on_update_kv_lens()``, the overlap
-  scheduler's accepted-token correction of ``kv_lens_cuda``, and
-  ``on_update_kv_lens()`` again.
-- The one-model MTP-Eagle draft loop of the target's worker,
-  ``GlmMtpWorker._forward_linear_draft_loop``, itself (index share on, accepted-row
-  stash, per-step metadata updates), with a stub draft model whose MTP layer runs the
-  real layer-2 attention. With
-  ``draft_loop`` selected the loop takes draft_loop's per-step metadata update.
-
-The context is written by the native producers (indexer K quantize + scatter,
-MLA RoPE-append of random latents). A decode sequence is four target verify
-steps (next_n = 6) separated by three draft loops with accepted counts
-``ACCEPTED``. Every arm replays the same sequence on the same caches; arms
-differ only in whether indexer_topk runs. ``GLM_MOE_DSA_TEST_FAST_PATHS``
-(``plan_testlib.FAST_PATHS_ENV``) is the base set of fast paths (for example
-``sparse_mla_decode,rope_kv_append`` or ``draft_loop``), whose plan is installed before any layer
-is built; the tests add ``indexer_topk`` on top of it for the indexer_topk arms.
-
-``test_rank_agreement`` runs the same sequence on every rank of an MPI job
-(8 ranks over 2 hosts) and compares digests across ranks.
-``test_indexer_topk_region_timing`` (``GLM_MOE_DSA_INDEXER_TOPK_PIPELINE_TIMING=1``) times the attention region of a
-full + shared layer pair in CUDA graphs with indexer_topk off and on.
-"""
-
-import hashlib
-import json
-import os
-import socket
-import sys
-import types
-from dataclasses import dataclass
-from types import SimpleNamespace
-from typing import Dict, List
-
-import pytest
-import torch
-
-from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
-from tensorrt_llm._utils import get_sm_version
-
-if not torch.cuda.is_available() or not IS_CUTLASS_DSL_AVAILABLE:
-    pytest.skip("CUDA and the CUTLASS DSL are required", allow_module_level=True)
-if get_sm_version() != 103:
-    pytest.skip("indexer_topk is verified on GB300 (sm103) only", allow_module_level=True)
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-    GlmMLA,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention_forward import (  # noqa: E501
-    _forward_dsa_attn,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-    indexer_decode as indexer_topk,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.spec_worker import (  # noqa: E501
-    GlmMtpWorker,
-)
-from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams
-from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManagerV2
-from tensorrt_llm._torch.attention.backends.sparse.dsa import indexer as indexer_mod
-from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
-from tensorrt_llm._torch.attention.mla import MLA
-from tensorrt_llm._torch.metadata import KVCacheParams
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
-from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
-from tensorrt_llm.bindings import DataType
-from tensorrt_llm.bindings.internal.batch_manager import CacheType
-from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata
-from tensorrt_llm.functional import PositionEmbeddingType
-from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig
-from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.models.modeling_utils import QuantConfig
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-# DSACacheManagerV2 owns background threads.
-pytestmark = pytest.mark.threadleak(enabled=False)
-
-LOCAL_HEADS = 8
-QK_NOPE = 192
-QK_ROPE = 64
-V_HEAD = 256
-KV_LORA = 512
-Q_LORA = 2048
-HIDDEN = 6144
-INDEX_HEADS = 32
-INDEX_DIM = 128
-TOP_K = 2048
-DRAFT_LEN = 5
-NEXT_N = DRAFT_LEN + 1
-TOKENS_PER_BLOCK = 64
-SCHEDULE_BLOCK_KV = indexer_mod._DG_SCHEDULE_BLOCK_KV
-# GLM-5.2: max_position_embeddings 2**20, rope theta 8e6, no rope scaling.
-MAX_POSITIONS = 1 << 20
-ROPE_THETA = 8.0e6
-# Served GEN: max_seq_len 2**20 plus the K = 5 one-model MTP slack.
-SERVED_MAX_SEQ_LEN = (1 << 20) + 14
-NUM_EXTRA_KV_TOKENS = DRAFT_LEN - 1
-LAYER_FULL, LAYER_SHARED, LAYER_MTP = 0, 1, 2
-DRAFT_VOCAB = 64
-# Accepted tokens (bonus included, 1..K+1) of the three draft loops.
-ACCEPTED = (3, 1, 6)
-_DEV = torch.device("cuda")
-BASE_FEATURES = os.environ.get(plan_testlib.FAST_PATHS_ENV, "")
-
-
-def _base_feature_set():
-    return set(plan_testlib.names(BASE_FEATURES))
-
-
-@pytest.fixture(scope="module", autouse=True)
-def base_plan():
-    """The plan of the base fast paths, installed before any layer is adopted."""
-    saved = (dispatch._INSTALLED[0], dispatch._VALIDATE[0], dispatch._require[0])
-    plan_testlib.select(dispatch, BASE_FEATURES)
-    yield
-    dispatch._INSTALLED[0], dispatch._VALIDATE[0], dispatch._require[0] = saved
-
-
-@pytest.fixture(autouse=True)
-def decode_step():
-    """Every test runs the target's decode steps."""
-    with dispatch.step(dispatch.Phase.DECODE):
-        yield
-
-
-@dataclass
-class _RopeConfig:
-    hidden_size: int
-    num_attention_heads: int
-    rope_scaling: object
-    max_position_embeddings: int
-    rope_theta: float
-    qk_rope_head_dim: int
-    model_type: str
-
-
-def _pos_embd_params():
-    rope_config = _RopeConfig(
-        hidden_size=HIDDEN,
-        num_attention_heads=LOCAL_HEADS,
-        rope_scaling=None,
-        max_position_embeddings=MAX_POSITIONS,
-        rope_theta=ROPE_THETA,
-        qk_rope_head_dim=QK_ROPE,
-        model_type="glm_moe_dsa",
-    )
-    # Same construction as the DeepSeek-V3-family attention layers.
-    return PositionalEmbeddingParams(
-        type=PositionEmbeddingType.yarn, rope=RopeParams.from_config(rope_config), is_neox=False
-    )
-
-
-def _expected_target_kv_lens(context_len: int) -> List[int]:
-    lens = [context_len]
-    for accepted in ACCEPTED:
-        lens.append(lens[-1] + accepted)
-    return lens
-
-
-@dataclass
-class StepInputs:
-    """Token-wise MLA/indexer inputs of one layer call."""
-
-    q: torch.Tensor
-    compressed_kv: torch.Tensor
-    k_pe: torch.Tensor
-    hidden: torch.Tensor
-    qr: torch.Tensor
-    position_ids: torch.Tensor
-
-
-class Recorder:
-    """Ordered record of one arm: indexer calls, shared rows and layer outputs."""
-
-    def __init__(self, name: str, canonical: bool):
-        self.name = name
-        self.canonical = canonical
-        self.events: List[dict] = []
-        self.capability: List[tuple] = []
-        self.sparse_mla_calls: List[tuple] = []
-        self.tag = None
-
-    def add(self, kind: str, **fields):
-        fields.update(kind=kind, tag=self.tag)
-        self.events.append(fields)
-
-
-def _canonicalize_(rows: torch.Tensor):
-    """Sort each top-k row in place (ascending local token ids)."""
-    assert bool((rows >= 0).all()), "tested lengths fill every top-k slot"
-    rows.copy_(torch.sort(rows, dim=1).values)
-
-
-class IndexerTap:
-    """Wraps ``Indexer.sparse_attn_indexer`` to record each call of an arm.
-
-    With ``canonical`` set, the rows every consumer reads (the returned rows,
-    and on a draft-loop reuse step the stashed row the reuse copies) are
-    sorted in place, so the attention input order does not depend on the
-    native top-k engine's in-row order, which varies run to run.
-    """
-
-    def __init__(self, recorder: Recorder):
-        self.rec = recorder
-
-    def wrap(self, original):
-        tap = self
-
-        def sparse_attn_indexer(indexer, metadata, *args, **kwargs):
-            in_draft = bool(metadata.in_mtp_draft_loop)
-            reuse = in_draft and bool(metadata.indexer_skip_topk)
-            pre = {}
-            if not in_draft:
-                # The per-request schedule and lengths indexer_topk consumes, as they are
-                # when the indexer runs.
-                pre["sched"] = metadata.scheduler_metadata_buffer.clone()
-                pre["kv"] = metadata.gen_indexer_kv_lens_cuda_runtime.clone()
-                pre["kv_lens_cuda"] = metadata.kv_lens_cuda[: metadata.num_seqs].clone()
-                pre["num_sms"] = metadata.num_sms
-                pre["block_table_cols"] = metadata.indexer_k_cache_block_offsets.shape[1]
-            if reuse:
-                stash = metadata.shared_topk_indices[: metadata.num_generations]
-                pre["stash"] = stash.clone()
-                if tap.rec.canonical:
-                    _canonicalize_(stash)
-            engaged = indexer_topk.DECISION_COUNTS.get("engaged", 0)
-            decided = sum(indexer_topk.DECISION_COUNTS.values())
-            out = original(indexer, metadata, *args, **kwargs)
-            rows = out[: metadata.num_tokens - metadata.num_ctx_tokens]
-            if tap.rec.canonical and not reuse:
-                _canonicalize_(rows)
-            tap.rec.add(
-                "indexer",
-                layer=indexer.layer_idx,
-                in_draft=in_draft,
-                reuse=reuse,
-                # On a reuse step rope_kv_append may defer the copy into its bridge, so the
-                # returned buffer is only meaningful when the indexer selected.
-                rows=None if reuse else torch.sort(rows, dim=1).values.clone(),
-                indexer_topk_engaged=indexer_topk.DECISION_COUNTS.get("engaged", 0) - engaged,
-                indexer_topk_decisions=sum(indexer_topk.DECISION_COUNTS.values()) - decided,
-                **pre,
-            )
-            return out
-
-        return sparse_attn_indexer
-
-
-class _StubMtpLayer:
-    """MTP layer of the stub draft model: runs the real layer-2 attention."""
-
-    def __init__(self, pipe: "GlmDsaPipeline", rec: Recorder, loop_idx: int):
-        self.pipe = pipe
-        self.rec = rec
-        self.loop_idx = loop_idx
-        self.step = 0
-
-    def __call__(
-        self,
-        *,
-        embed_tokens,
-        all_rank_num_tokens,
-        input_ids,
-        position_ids,
-        hidden_states,
-        attn_metadata,
-        spec_metadata,
-        **kwargs,
-    ):
-        step = self.step
-        self.step += 1
-        self.rec.tag = ("draft", self.loop_idx, step)
-        n = attn_metadata.num_tokens
-        positions = position_ids.reshape(-1)[:n].to(torch.int32).contiguous()
-        seed = 1000 + 37 * self.loop_idx + step
-        inputs = self.pipe.make_inputs(seed=seed, n=n, positions=positions)
-        out = self.pipe.layer_forward(LAYER_MTP, attn_metadata, inputs)
-        self.rec.add("output", layer=LAYER_MTP, out=out.clone())
-        # Hidden states that depend on the attention output, so the drafted
-        # tokens do as well.
-        hidden = torch.zeros(n, HIDDEN, dtype=torch.bfloat16, device=_DEV)
-        hidden[:, : out.shape[1]] = out
-        return hidden
-
-    @staticmethod
-    def shared_head(hidden_states, lm_head, attn_metadata, return_context_logits):
-        return hidden_states[:, :DRAFT_VOCAB].float()
-
-
-def _sample_draft_tokens(logits, spec_metadata, batch_size, draft_step=0, mapping_lm_head_tp=None):
-    return torch.argmax(logits, dim=-1).to(torch.int32)
-
-
-def _make_worker():
-    """The attributes of an MTP-Eagle one-model worker the draft loop reads."""
-    worker = SimpleNamespace(
-        is_mtp_eagle=True,
-        guided_decoder=None,
-        sa_enhancer=None,
-        model_config=None,
-        spec_config=SimpleNamespace(use_relaxed_acceptance_for_thinking=False),
-        _auxiliary_state_handlers=None,
-        _uses_external_shared_target_kv=False,
-        use_dynamic_tree=False,
-        _draft_loop_ops=None,
-    )
-    worker.draft_kv_cache_context = types.MethodType(SpecWorkerBase.draft_kv_cache_context, worker)
-    for name in (
-        "_run_draft_forward",
-        "_get_step_all_rank_num_tokens",
-        "_prepare_flash_mla_generation_layout",
-    ):
-        setattr(worker, name, types.MethodType(getattr(GlmMtpWorker, name), worker))
-    worker.sample_draft_tokens = _sample_draft_tokens
-    worker._can_use_rejection_sampling = lambda spec_metadata: False
-    return worker
-
-
-def _make_spec_metadata():
-    return SimpleNamespace(
-        runtime_draft_len=DRAFT_LEN,
-        batch_indices_cuda=torch.zeros(1, dtype=torch.int32, device=_DEV),
-        all_rank_num_tokens=None,
-        subseq_all_rank_num_tokens=None,
-        wants_advanced_draft_sampling=False,
-        max_num_requests=1,
-        draft_tokens=torch.zeros(DRAFT_LEN, dtype=torch.int32, device=_DEV),
-        enable_penalty=False,
-    )
-
-
-class GlmDsaPipeline:
-    """Three GLM-5.2-shaped DSA MLA layers over one ``DSACacheManagerV2``."""
-
-    def __init__(self, context_len: int, seed: int = 0):
-        # ``context_len`` is the KV length of the first target verify step
-        # (its six verify tokens included).
-        self.context_len = context_len
-        self.prompt_len = context_len - NEXT_N
-        self.capacity = context_len + sum(ACCEPTED) + 2 * NEXT_N + NUM_EXTRA_KV_TOKENS
-        assert self.capacity + NUM_EXTRA_KV_TOKENS <= SERVED_MAX_SEQ_LEN
-        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
-        self.pretrained = SimpleNamespace(
-            rms_norm_eps=1e-5,
-            num_hidden_layers=2,
-            index_topk_pattern=["F", "S"],
-            qk_rope_head_dim=QK_ROPE,
-        )
-        self.sparse_config = DeepSeekSparseAttentionConfig(
-            index_n_heads=INDEX_HEADS,
-            index_head_dim=INDEX_DIM,
-            index_topk=TOP_K,
-            use_cute_dsl_topk=True,
-            use_cute_dsl_paged_mqa_logits=True,
-            enable_heuristic_topk=True,
-            index_share_for_mtp_iteration=True,
-        )
-        quant_config = QuantConfig()
-        quant_config.kv_cache_quant_algo = QuantAlgo.FP8
-        self.model_config = ModelConfig(
-            mapping=self.mapping,
-            sparse_attention_config=self.sparse_config,
-            pretrained_config=self.pretrained,
-            quant_config=quant_config,
-        )
-        self.model_config.extra_attrs["kv_cache_dtype"] = "fp8"
-        pos = _pos_embd_params()
-        self.layers = []
-        for layer_idx in range(3):
-            mla = MLA(
-                hidden_size=HIDDEN,
-                num_attention_heads=LOCAL_HEADS,
-                num_key_value_heads=1,
-                qk_nope_head_dim=QK_NOPE,
-                qk_rope_head_dim=QK_ROPE,
-                v_head_dim=V_HEAD,
-                q_lora_rank=Q_LORA,
-                kv_lora_rank=KV_LORA,
-                predicted_tokens_per_seq=NEXT_N,
-                max_position_embeddings=MAX_POSITIONS,
-                bias=False,
-                pos_embd_params=pos,
-                layer_idx=layer_idx,
-                dtype=torch.bfloat16,
-                config=self.model_config,
-            )
-            # The model's DeepseekV32Attention adds only its constructor to MLA; the pipeline
-            # builds the MLA at GLM shapes directly and adopts it as the model's attention. The
-            # model builds the FMHA libraries in ``create_weights`` after adoption; a module built
-            # alone rebuilds them for the library adoption declared.
-            mla.__class__ = DeepseekV32Attention
-            GlmMLA.adopt(mla)
-            for backend in (mla.mha, mla.mqa):
-                if backend is not None and backend.prepended_fmha_lib_classes:
-                    backend.update_quant_config(backend.quant_config)
-            mla = mla.to(_DEV)
-            if mla.mqa.indexer is not None:
-                mla.indexer = mla.mqa.indexer.to(_DEV)
-            self.layers.append(mla)
-        assert self.layers[LAYER_FULL].mqa.indexer is not None
-        assert self.layers[LAYER_SHARED].mqa.indexer is None, "layer 1 must be shared"
-        assert self.layers[LAYER_MTP].mqa.indexer is not None
-        self._init_weights(seed)
-
-        self.kv_cache_manager = DSACacheManagerV2(
-            KvCacheConfig(
-                dtype="fp8",
-                # At least max_seq_len, or the manager clamps max_seq_len.
-                max_tokens=SERVED_MAX_SEQ_LEN + 64 * TOKENS_PER_BLOCK,
-                enable_block_reuse=False,
-                host_cache_size=0,
-            ),
-            CacheType.SELFKONLY,
-            num_layers=3,
-            num_kv_heads=1,
-            head_dim=KV_LORA + QK_ROPE,
-            tokens_per_block=TOKENS_PER_BLOCK,
-            max_seq_len=SERVED_MAX_SEQ_LEN,
-            max_batch_size=1,
-            mapping=self.mapping,
-            dtype=DataType.FP8,
-            sparse_attention_config=self.sparse_config,
-            pretrained_config=self.pretrained,
-        )
-        assert self.kv_cache_manager.max_seq_len == SERVED_MAX_SEQ_LEN
-        assert self.kv_cache_manager.indexer_k_cache_local_layer_mask == [True, False, True]
-        added = self.kv_cache_manager.add_dummy_requests(
-            request_ids=[0], token_nums=[self.capacity], is_gen=False, prepare_resource=True
-        )
-        assert added is not None, "KV cache allocation failed"
-        sparse_params = self.sparse_config.to_sparse_params(pretrained_config=self.pretrained)
-        self.attn_cls = get_attention_backend("TRTLLM", sparse_params=sparse_params)
-        self.metadata_params = self.sparse_config.to_sparse_metadata_params(
-            pretrained_config=self.pretrained
-        )
-        assert self.metadata_params.has_shared_indexer_layers
-        assert self.metadata_params.mtp_index_share
-        self._populate_context(seed)
-        self.md = self._new_metadata(
-            num_contexts=0,
-            seq_len=NEXT_N,
-            cached=self.prompt_len,
-            max_num_tokens=64,
-            extra_kv=NUM_EXTRA_KV_TOKENS,
-        )
-        with torch.no_grad():
-            self.embed_tokens = torch.nn.Embedding(DRAFT_VOCAB, HIDDEN).to(_DEV, torch.bfloat16)
-
-    # ------------------------------------------------------------------ setup
-    def _init_weights(self, seed: int):
-        gen = torch.Generator(device=_DEV).manual_seed(seed)
-
-        def randn_like(t):
-            return (torch.randn(t.shape, device=_DEV, generator=gen) * 0.02).to(t.dtype)
-
-        with torch.no_grad():
-            for mla in self.layers:
-                mla.kv_b_proj.weight.copy_(randn_like(mla.kv_b_proj.weight))
-                kv_b = mla.kv_b_proj.weight.data.view(LOCAL_HEADS, QK_NOPE + V_HEAD, KV_LORA)
-                mla.v_b_proj.data = kv_b[:, QK_NOPE:, :].contiguous()
-                mla.k_b_proj_trans.data = kv_b[:, :QK_NOPE, :].transpose(1, 2).contiguous()
-                indexer = mla.mqa.indexer
-                if indexer is not None:
-                    for linear in (indexer.wq_b, indexer.wk, indexer.weights_proj):
-                        linear.weight.copy_(randn_like(linear.weight))
-                    indexer.post_load_weights()
-
-    def _new_metadata(self, *, num_contexts, seq_len, cached, max_num_tokens, extra_kv):
-        return self.attn_cls.Metadata(
-            seq_lens=torch.tensor([seq_len], dtype=torch.int),
-            request_ids=[0],
-            max_num_requests=1,
-            num_contexts=num_contexts,
-            prompt_lens=[self.prompt_len],
-            max_num_tokens=max_num_tokens,
-            kv_cache_manager=self.kv_cache_manager,
-            kv_cache_params=KVCacheParams(
-                use_cache=True, num_cached_tokens_per_seq=[cached], num_extra_kv_tokens=extra_kv
-            ),
-            mapping=self.mapping,
-            sparse_metadata_params=self.metadata_params,
-        )
-
-    @torch.inference_mode()
-    def _populate_context(self, seed: int):
-        """Write the prompt with the native producers (one prefill-shaped call)."""
-        n = self.prompt_len
-        md = self._new_metadata(num_contexts=1, seq_len=n, cached=0, max_num_tokens=n, extra_kv=0)
-        md.prepare()
-        gen = torch.Generator(device=_DEV).manual_seed(seed + 1)
-        for layer_idx in (LAYER_FULL, LAYER_MTP):
-            indexer = self.layers[layer_idx].mqa.indexer
-            k = torch.randn(n, INDEX_DIM, device=_DEV, generator=gen).to(torch.bfloat16)
-            rope = indexer.rope_dim
-            # The indexer's own K quantizer and K-cache scatter.
-            k_fp8, k_scale = indexer._prep_q_or_k(
-                k[:, :rope].contiguous(), k[:, rope:].contiguous()
-            )
-            indexer._update_k_cache(k_fp8, k_scale, md)
-            del k, k_fp8, k_scale
-        for mla in self.layers:
-            latent = torch.randn(n, KV_LORA + QK_ROPE, device=_DEV, generator=gen)
-            latent = latent.to(torch.bfloat16)
-            dummy_q = torch.zeros(
-                n, LOCAL_HEADS * (QK_NOPE + QK_ROPE), device=_DEV, dtype=torch.bfloat16
-            )
-            mla.mqa.mla_rope_append_paged_kv_assign_q(dummy_q, latent, md)
-            del latent, dummy_q
-        torch.cuda.synchronize()
-        del md
-        torch.cuda.empty_cache()
-
-    def shutdown(self):
-        self.kv_cache_manager.shutdown()
-
-    # ------------------------------------------------------------- execution
-    def make_inputs(self, seed: int, n: int, positions: torch.Tensor) -> StepInputs:
-        gen = torch.Generator(device=_DEV).manual_seed(seed)
-
-        def rnd(*shape):
-            return torch.randn(*shape, device=_DEV, generator=gen).to(torch.bfloat16)
-
-        return StepInputs(
-            q=rnd(n, LOCAL_HEADS * (QK_NOPE + QK_ROPE)),
-            compressed_kv=rnd(n, KV_LORA),
-            k_pe=rnd(n, QK_ROPE),
-            hidden=rnd(n, HIDDEN),
-            qr=rnd(n, Q_LORA),
-            position_ids=positions,
-        )
-
-    def project(self, layer_idx: int, inputs: StepInputs):
-        """Graph-capturable indexer projections of a full-indexer layer (Op 1)."""
-        indexer = self.layers[layer_idx].mqa.indexer
-        if indexer is None:
-            return []
-        return list(indexer.pre_indexer_proj(inputs.qr, inputs.hidden, inputs.position_ids))
-
-    def attend(self, layer_idx: int, md, inputs: StepInputs, intermediates) -> torch.Tensor:
-        """``_forward_dsa_attn`` of one layer (Op 2) on copies of the inputs."""
-        mla = self.layers[layer_idx]
-        n = md.num_tokens
-        latent = torch.cat([inputs.compressed_kv, inputs.k_pe], dim=-1)
-        output = torch.empty(n, LOCAL_HEADS * V_HEAD, dtype=torch.bfloat16, device=_DEV)
-        _forward_dsa_attn(
-            mla,
-            inputs.q.clone(),
-            inputs.compressed_kv.clone(),
-            inputs.k_pe.clone(),
-            latent,
-            intermediates,
-            inputs.position_ids,
-            md,
-            output,
-        )
-        return output
-
-    def layer_forward(self, layer_idx: int, md, inputs: StepInputs) -> torch.Tensor:
-        return self.attend(layer_idx, md, inputs, self.project(layer_idx, inputs))
-
-    def prepare_target_step(self, host_cached: int, kv_offset: int):
-        """Runner order: host metadata + ``prepare()``, then ``_preprocess_inputs``."""
-        md = self.md
-        md.seq_lens = torch.tensor([NEXT_N], dtype=torch.int)
-        md.num_contexts = 0
-        md.request_ids = [0]
-        md.prompt_lens = [self.prompt_len]
-        md.kv_cache_params = KVCacheParams(
-            use_cache=True,
-            num_cached_tokens_per_seq=[host_cached],
-            num_extra_kv_tokens=NUM_EXTRA_KV_TOKENS,
-        )
-        md.update_spec_dec_param(
-            batch_size=1,
-            is_spec_decoding_enabled=True,
-            is_spec_dec_tree=False,
-            is_spec_dec_dynamic_tree=False,
-            max_draft_len=DRAFT_LEN,
-            max_total_draft_tokens=DRAFT_LEN,
-            spec_metadata=None,
-            spec_tree_manager=None,
-            num_contexts=0,
-        )
-        md.prepare()
-        # DecoderRunner._preprocess_inputs (overlap scheduler + spec decode).
-        md.on_update_kv_lens()
-        md.kv_lens_cuda[0:1] += kv_offset
-        md.on_update_kv_lens()
-        return md
-
-    @torch.inference_mode()
-    def run_sequence(self, rec: Recorder, draft_loop: bool):
-        """Four target steps separated by three draft loops."""
-        md = self.md
-        kv = self.context_len
-        cached = kv - NEXT_N
-        host_cached = cached
-        kv_offset = 0
-        for step in range(len(ACCEPTED) + 1):
-            rec.tag = ("target", step)
-            self.prepare_target_step(host_cached, kv_offset)
-            positions = torch.arange(cached, cached + NEXT_N, dtype=torch.int32, device=_DEV)
-            rec.add("target_kv", kv=kv)
-            for layer_idx in (LAYER_FULL, LAYER_SHARED):
-                if layer_idx == LAYER_SHARED:
-                    rec.add(
-                        "shared_rows",
-                        rows=torch.sort(md.shared_topk_indices[:NEXT_N], dim=1).values.clone(),
-                    )
-                inputs = self.make_inputs(seed=10 * step + layer_idx, n=NEXT_N, positions=positions)
-                out = self.layer_forward(layer_idx, md, inputs)
-                rec.add("output", layer=layer_idx, out=out.clone())
-            if step == len(ACCEPTED):
-                break
-            accepted = ACCEPTED[step]
-            self.run_draft_loop(rec, step, accepted, positions, draft_loop)
-            rec.tag = ("after_loop", step)
-            rec.add("stale_sched", sched=md.scheduler_metadata_buffer.clone())
-            # Overlap scheduler: the host assumes every draft was accepted and
-            # the next forward corrects kv_lens_cuda by (accepted - next_n).
-            host_cached = cached + NEXT_N
-            kv_offset = accepted - NEXT_N
-            cached += accepted
-            kv += accepted
-        torch.cuda.synchronize()
-        return rec
-
-    def run_draft_loop(
-        self, rec: Recorder, loop_idx: int, accepted: int, positions, draft_loop: bool
-    ):
-        md = self.md
-        worker = _make_worker()
-        spec_md = _make_spec_metadata()
-        num_accepted = torch.tensor([accepted], dtype=torch.int32, device=_DEV)
-        draft_model = SimpleNamespace(
-            embed_tokens=self.embed_tokens,
-            lm_head=None,
-            mtp_layers=[_StubMtpLayer(self, rec, loop_idx)],
-        )
-        if draft_loop:
-            ops = dispatch.get_impl("draft_loop", "draft_loop_ops")(worker, md, spec_md)
-            assert ops is not None, "draft_loop declined the draft loop (see the capability log)"
-            # Acceptance is not run here (there are no target logits); give draft_loop
-            # the step-0 rows the native loop computes for one generation request.
-            ops.gather_ids = (num_accepted.to(torch.int64) - 1).contiguous()
-            worker._draft_loop_ops = ops
-        gen = torch.Generator(device=_DEV).manual_seed(7000 + loop_idx)
-        inputs = {
-            "input_ids": torch.randint(0, DRAFT_VOCAB, (NEXT_N,), device=_DEV, generator=gen),
-            "position_ids": positions.clone(),
-            "hidden_states": torch.randn(NEXT_N, HIDDEN, device=_DEV, generator=gen).to(
-                torch.bfloat16
-            ),
-            "attn_metadata": md,
-            "spec_metadata": spec_md,
-        }
-        # The worker's _forward_impl around the loop.
-        md.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
-        try:
-            tokens = GlmMtpWorker._forward_linear_draft_loop(
-                worker,
-                inputs,
-                md,
-                spec_md,
-                draft_model,
-                None,
-                0,
-                1,
-                num_accepted,
-                None,
-            )
-        finally:
-            md.restore_from_spec_dec()
-            md.on_update()
-            md.use_spec_decoding = True
-        rec.tag = ("draft_tokens", loop_idx)
-        rec.add("draft_tokens", tokens=tokens.clone())
-
-
-# --------------------------------------------------------------------- arms
-def _install_recorders(mp, rec: Recorder):
-    original_log = dispatch.log_capability
-
-    def log_capability(feature_id, key, eligible, reason=""):
-        rec.capability.append((feature_id, key, bool(eligible), reason, rec.tag))
-        return original_log(feature_id, key, eligible, reason)
-
-    mp.setattr(dispatch, "log_capability", log_capability)
-    tap = IndexerTap(rec)
-    original_indexer = indexer_mod.Indexer.sparse_attn_indexer
-    mp.setattr(indexer_mod.Indexer, "sparse_attn_indexer", tap.wrap(original_indexer))
-    if "sparse_mla_decode" in _base_feature_set():
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-            sparse_mla as sparse_mla_sparse_mla,
-        )
-
-        sparse_mla_cls = sparse_mla_sparse_mla.SparseMlaDecodeFmha
-        original_forward = sparse_mla_cls.forward
-
-        def forward(self, *args, **kwargs):
-            rec.sparse_mla_calls.append(rec.tag)
-            return original_forward(self, *args, **kwargs)
-
-        mp.setattr(sparse_mla_cls, "forward", forward)
-
-
-def _select(mp, with_indexer_topk: bool):
-    features = _base_feature_set()
-    if with_indexer_topk:
-        features.add("indexer_topk")
-    else:
-        features.discard("indexer_topk")
-    plan_testlib.select(dispatch, features, mp)
-    assert dispatch.is_enabled("indexer_topk") == with_indexer_topk
-
-
-def run_arm(pipe: GlmDsaPipeline, monkeypatch, name: str, with_indexer_topk: bool, canonical: bool):
-    rec = Recorder(name, canonical)
-    with monkeypatch.context() as mp:
-        _select(mp, with_indexer_topk)
-        _install_recorders(mp, rec)
-        pipe.run_sequence(rec, draft_loop="draft_loop" in _base_feature_set())
-    return rec
-
-
-# -------------------------------------------------------------- comparisons
-def _fresh_schedule(kv: torch.Tensor, num_sms: int) -> torch.Tensor:
-    return get_paged_mqa_logits_metadata(kv.view(-1, 1), SCHEDULE_BLOCK_KV, num_sms)
-
-
-def _indexer_calls(rec: Recorder, phase: str):
-    return [e for e in rec.events if e["kind"] == "indexer" and e["tag"][0] == phase]
-
-
-def check_arm_invariants(rec: Recorder, pipe: GlmDsaPipeline, with_indexer_topk: bool):
-    """Per-arm checks that do not need another arm."""
-    expected_kv = _expected_target_kv_lens(pipe.context_len)
-    target_calls = _indexer_calls(rec, "target")
-    assert len(target_calls) == len(expected_kv), "one full-layer indexer call per target step"
-    for call, kv in zip(target_calls, expected_kv):
-        assert call["layer"] == LAYER_FULL and not call["in_draft"]
-        # The lengths and the per-request schedule indexer_topk reads are this step's,
-        # not a draft step's leftovers.
-        assert call["kv"].tolist() == [kv], (rec.name, call["tag"], call["kv"].tolist(), kv)
-        assert call["kv_lens_cuda"].tolist() == [kv]
-        fresh = _fresh_schedule(call["kv"], call["num_sms"])
-        assert torch.equal(call["sched"], fresh), (rec.name, call["tag"], "stale schedule")
-        assert call["block_table_cols"] % 2 == 0
-        assert call["indexer_topk_decisions"] == (1 if with_indexer_topk else 0)
-        assert call["indexer_topk_engaged"] == (1 if with_indexer_topk else 0), (
-            rec.name,
-            call["tag"],
-        )
-    draft_calls = _indexer_calls(rec, "draft")
-    assert len(draft_calls) == DRAFT_LEN * len(ACCEPTED)
-    for call in draft_calls:
-        assert call["layer"] == LAYER_MTP and call["in_draft"]
-        assert call["reuse"] == (call["tag"][2] > 0)
-        # indexer_topk is target-only: it declines every draft selection before device
-        # work, and the reuse steps never reach it.
-        assert call["indexer_topk_engaged"] == 0
-        assert call["indexer_topk_decisions"] == (
-            1 if with_indexer_topk and not call["reuse"] else 0
-        )
-    # The shared layer reads exactly the full layer's rows of the same step.
-    for step in range(len(expected_kv)):
-        full = next(e for e in target_calls if e["tag"] == ("target", step))
-        shared = next(
-            e for e in rec.events if e["kind"] == "shared_rows" and e["tag"] == ("target", step)
-        )
-        assert torch.equal(full["rows"], shared["rows"]), (rec.name, step)
-    features = _base_feature_set()
-    caps = rec.capability
-    num_attention_calls = 2 * len(expected_kv) + DRAFT_LEN * len(ACCEPTED)
-    # The consumers of indexer_topk's rows are the target layers; draft-step decisions
-    # of the neighbour packages are reported, not asserted.
-    if "sparse_mla_decode" in features:
-        declined = [c for c in caps if c[0] == "sparse_mla_decode" and not c[2]]
-        assert not [c for c in declined if c[4] and c[4][0] == "target"], declined
-        target_sparse_mla = [t for t in rec.sparse_mla_calls if t[0] == "target"]
-        # Both target layers (full and shared) of every step attend through sparse_mla_decode.
-        assert len(target_sparse_mla) == 2 * len(expected_kv), rec.sparse_mla_calls
-        if len(rec.sparse_mla_calls) != num_attention_calls or declined:
-            print(
-                f"[indexer_topk] {rec.name}: sparse_mla_decode calls {len(rec.sparse_mla_calls)}, declines {declined}"
-            )
-    if "rope_kv_append" in features:
-        bridge = [c for c in caps if c[0] == "rope_kv_append" and c[1] == "bridge"]
-        target_bridge = [c for c in bridge if c[4][0] == "target"]
-        assert len(target_bridge) == 2 * len(expected_kv), bridge
-        assert all(c[2] for c in target_bridge), bridge
-        if len(bridge) != num_attention_calls or not all(c[2] for c in bridge):
-            print(
-                f"[indexer_topk] {rec.name}: rope_kv_append bridge decisions {[c[2] for c in bridge]}"
-            )
-    if "draft_loop" in features:
-        steps = [c for c in caps if c[0] == "draft_loop" and c[1].startswith("draft_step:")]
-        assert len(steps) == DRAFT_LEN * len(ACCEPTED), steps
-        assert all(c[1] == "draft_step:engaged" for c in steps), steps
-
-
-def _pairs(rec_a: Recorder, rec_b: Recorder):
-    assert [(e["kind"], e["tag"]) for e in rec_a.events] == [
-        (e["kind"], e["tag"]) for e in rec_b.events
-    ]
-    return zip(rec_a.events, rec_b.events)
-
-
-def compare_selection(rec_a: Recorder, rec_b: Recorder, draft_tokens: bool = True):
-    """Discrete invariants: identical selections, shared rows and stashes."""
-    for a, b in _pairs(rec_a, rec_b):
-        if a["kind"] == "indexer":
-            if a["rows"] is not None:
-                assert torch.equal(a["rows"], b["rows"]), (rec_a.name, rec_b.name, a["tag"])
-            if "stash" in a:
-                sa = torch.sort(a["stash"], dim=1).values
-                sb = torch.sort(b["stash"], dim=1).values
-                assert torch.equal(sa, sb), (rec_a.name, rec_b.name, a["tag"], "stash")
-        elif a["kind"] == "shared_rows":
-            assert torch.equal(a["rows"], b["rows"]), (rec_a.name, rec_b.name, a["tag"])
-        elif a["kind"] == "draft_tokens" and draft_tokens:
-            assert torch.equal(a["tokens"], b["tokens"]), (rec_a.name, rec_b.name, a["tag"])
-
-
-def output_diffs(rec_a: Recorder, rec_b: Recorder) -> Dict[tuple, float]:
-    diffs = {}
-    for a, b in _pairs(rec_a, rec_b):
-        if a["kind"] == "output":
-            fa, fb = a["out"].float(), b["out"].float()
-            assert torch.isfinite(fa).all() and torch.isfinite(fb).all()
-            diffs[(a["tag"], a["layer"])] = (fa - fb).abs().max().item()
-    return diffs
-
-
-def _report(title: str, diffs: Dict[tuple, float]):
-    worst = max(diffs.values()) if diffs else 0.0
-    nonzero = sum(1 for v in diffs.values() if v != 0.0)
-    print(
-        f"[indexer_topk] {title}: {len(diffs)} layer outputs, {nonzero} differ, max |d| = {worst:.3e}"
-    )
-
-
-# --------------------------------------------------------------------- tests
-# DeepGEMM's decode schedule balances ceil(kv / SPLIT_KV) work units per
-# request over the SMs (SPLIT_KV is one of 64..512 tokens), so it changes only
-# where kv crosses a multiple of SPLIT_KV. 66,046 and 1,048,062 put the KV
-# length the first draft loop leaves behind (L + a1 - 1) on a multiple of 512
-# and the next target step's (L + a1) one past it, so a leftover draft-step
-# schedule differs from the target schedule for every SPLIT_KV candidate.
-SCHEDULE_SPLIT_LCM = 512
-CONTEXT_LENS = [66046, 491061, 1048062]
-
-
-def _is_schedule_boundary(context_len: int) -> bool:
-    return (context_len + ACCEPTED[0] - 1) % SCHEDULE_SPLIT_LCM == 0
-
-
-@pytest.fixture(scope="module", params=CONTEXT_LENS, ids=lambda n: f"kv{n}")
-def pipeline(request):
-    pipe = GlmDsaPipeline(request.param)
-    yield pipe
-    pipe.shutdown()
-
-
-def _served_order_report(title: str, rec_a: Recorder, rec_b: Recorder) -> float:
-    diffs = output_diffs(rec_a, rec_b)
-    outputs_a = {(e["tag"], e["layer"]): e["out"] for e in rec_a.events if e["kind"] == "output"}
-    outputs_b = {(e["tag"], e["layer"]): e["out"] for e in rec_b.events if e["kind"] == "output"}
-    worst_cos = 1.0
-    for key in diffs:
-        cos = torch.nn.functional.cosine_similarity(
-            outputs_a[key].float().flatten(), outputs_b[key].float().flatten(), dim=0
-        ).item()
-        worst_cos = min(worst_cos, cos)
-    _report(title, diffs)
-    print(f"[indexer_topk] {title}: min cosine {worst_cos:.7f}")
-    return worst_cos
-
-
-def test_indexer_topk_in_real_decode_pipeline(pipeline, monkeypatch):
-    """indexer_topk on vs off in the real metadata/draft-loop pipeline (sets and outputs)."""
-    print(f"[indexer_topk] base features '{BASE_FEATURES}', kv {pipeline.context_len}")
-    native_a = run_arm(pipeline, monkeypatch, "native_a", with_indexer_topk=False, canonical=True)
-    native_b = run_arm(pipeline, monkeypatch, "native_b", with_indexer_topk=False, canonical=True)
-    with_indexer_topk = run_arm(
-        pipeline, monkeypatch, "indexer_topk", with_indexer_topk=True, canonical=True
-    )
-    for rec, on in ((native_a, False), (native_b, False), (with_indexer_topk, True)):
-        check_arm_invariants(rec, pipeline, on)
-    compare_selection(native_a, native_b)
-    compare_selection(native_a, with_indexer_topk)
-    native_env = output_diffs(native_a, native_b)
-    indexer_topk_diff = output_diffs(native_a, with_indexer_topk)
-    _report("native vs native (canonical order)", native_env)
-    _report("indexer_topk vs native (canonical order)", indexer_topk_diff)
-    for key, value in indexer_topk_diff.items():
-        # Identical index sets in the same order: the attention is the same
-        # computation, so allow only what native shows run to run.
-        assert value <= native_env[key], (key, value, native_env[key])
-
-    # Served order (no canonicalization): the sets still match. The outputs
-    # then depend on the native top-k's run-to-run in-row order (the FP8
-    # attention tiles the selected tokens in that order), for native and indexer_topk
-    # alike; this is characterized, not bounded by a tolerance.
-    native_raw = run_arm(
-        pipeline, monkeypatch, "native_raw", with_indexer_topk=False, canonical=False
-    )
-    native_raw_b = run_arm(
-        pipeline, monkeypatch, "native_raw_b", with_indexer_topk=False, canonical=False
-    )
-    indexer_topk_raw = run_arm(
-        pipeline, monkeypatch, "indexer_topk_raw", with_indexer_topk=True, canonical=False
-    )
-    check_arm_invariants(native_raw, pipeline, False)
-    check_arm_invariants(native_raw_b, pipeline, False)
-    check_arm_invariants(indexer_topk_raw, pipeline, True)
-    compare_selection(native_raw, native_raw_b, draft_tokens=False)
-    compare_selection(native_raw, indexer_topk_raw, draft_tokens=False)
-    compare_selection(native_a, indexer_topk_raw, draft_tokens=False)
-    native_cos = _served_order_report("native vs native (served order)", native_raw, native_raw_b)
-    indexer_topk_cos = _served_order_report(
-        "indexer_topk vs native (served order)", native_raw, indexer_topk_raw
-    )
-    assert native_cos > 0.999 and indexer_topk_cos > 0.999, (native_cos, indexer_topk_cos)
-
-    # The freshness check is not vacuous at the boundary lengths: the schedule
-    # the first draft loop leaves behind is that of its last step and differs
-    # from the next target step's.
-    stale = [e for e in native_a.events if e["kind"] == "stale_sched"]
-    targets = _indexer_calls(native_a, "target")
-    differs = [not torch.equal(s["sched"], t["sched"]) for s, t in zip(stale, targets[1:])]
-    print(
-        f"[indexer_topk] leftover draft schedule differs from the next target schedule: {differs}"
-    )
-    if _is_schedule_boundary(pipeline.context_len):
-        leftover_kv = torch.tensor(
-            [pipeline.context_len + ACCEPTED[0] - 1], dtype=torch.int32, device=_DEV
-        )
-        leftover = _fresh_schedule(leftover_kv, targets[1]["num_sms"])
-        assert torch.equal(stale[0]["sched"], leftover), "leftover is the last draft step's"
-        assert differs[0], "boundary case must exercise a schedule change"
-    caps = sorted({c[:3] for c in with_indexer_topk.capability}, key=str)
-    print(f"[indexer_topk] capability decisions in the indexer_topk arm: {caps}")
-    print(
-        f"[indexer_topk] sparse_mla_decode sparse-MLA calls per arm: {len(with_indexer_topk.sparse_mla_calls)}"
-    )
-    print(f"[indexer_topk] indexer block-table width: {targets[0]['block_table_cols']}")
-    print(f"[indexer_topk] indexer_topk decisions: {dict(indexer_topk.DECISION_COUNTS)}")
-
-
-def _mpi_comm():
-    try:
-        from mpi4py import MPI
-    except ImportError:
-        return None
-    comm = MPI.COMM_WORLD
-    return comm if comm.Get_size() > 1 else None
-
-
-def _digest(t: torch.Tensor) -> str:
-    t = t.detach().contiguous().cpu()
-    if t.dtype in (torch.bfloat16, torch.float16):
-        t = t.view(torch.int16)
-    elif t.element_size() == 1:
-        t = t.view(torch.uint8)
-    return hashlib.sha256(t.numpy().tobytes()).hexdigest()[:16]
-
-
-@pytest.mark.parametrize("context_len", [491061])
-def test_rank_agreement(context_len, monkeypatch):
-    """Every rank of the MPI job selects identical sets with indexer_topk on and off."""
-    comm = _mpi_comm()
-    if comm is None:
-        pytest.skip("needs an MPI job with more than one rank")
-    rank, size = comm.Get_rank(), comm.Get_size()
-    pipe = GlmDsaPipeline(context_len)
-    try:
-        native = run_arm(pipe, monkeypatch, "native", with_indexer_topk=False, canonical=True)
-        with_indexer_topk = run_arm(
-            pipe, monkeypatch, "indexer_topk", with_indexer_topk=True, canonical=True
-        )
-        check_arm_invariants(native, pipe, False)
-        check_arm_invariants(with_indexer_topk, pipe, True)
-        compare_selection(native, with_indexer_topk)
-        kvm = pipe.kv_cache_manager
-        first_page = (
-            int(kvm.get_pool_block_indices(1, request_ids=[0])[0, 0])
-            * kvm.indexer_k_cache_page_scale
-        )
-        props = torch.cuda.get_device_properties(torch.cuda.current_device())
-
-        def selections(rec):
-            return [
-                (str(e["tag"]), e["layer"], _digest(e["rows"]))
-                for e in rec.events
-                if e["kind"] == "indexer" and e["rows"] is not None
-            ]
-
-        digest = {
-            "host": socket.gethostname(),
-            "gpu": str(getattr(props, "uuid", props.name)),
-            "indexer_k_first_page": _digest(
-                kvm.get_indexer_k_cache_buffers(LAYER_FULL)[first_page]
-            ),
-            "selections": selections(with_indexer_topk),
-            "native_selections": selections(native),
-            "outputs": [
-                (str(e["tag"]), e["layer"], _digest(e["out"]))
-                for e in with_indexer_topk.events
-                if e["kind"] == "output"
-            ],
-            "indexer_topk_decisions": dict(indexer_topk.DECISION_COUNTS),
-        }
-    finally:
-        pipe.shutdown()
-    everyone = comm.allgather(digest)
-    if rank == 0:
-        hosts = sorted({d["host"] for d in everyone})
-        print(f"[indexer_topk] {size} ranks on {len(hosts)} hosts: {hosts}")
-        for r, d in enumerate(everyone):
-            print(
-                f"[indexer_topk] rank {r} {d['host']} {d['gpu']} indexer_topk decisions {d['indexer_topk_decisions']}"
-            )
-        print(f"[indexer_topk] rank 0 selection digests: {everyone[0]['selections']}")
-    ref = everyone[0]
-    for r, d in enumerate(everyone):
-        assert d["indexer_k_first_page"] == ref["indexer_k_first_page"], r
-        assert d["selections"] == ref["selections"], r
-        assert d["native_selections"] == ref["native_selections"], r
-        assert d["selections"] == d["native_selections"], r
-        assert d["indexer_topk_decisions"] == ref["indexer_topk_decisions"], r
-        # The attention is replicated per rank on identical inputs as well.
-        assert d["outputs"] == ref["outputs"], r
-
-
-# ------------------------------------------------------------------- timing
-TIMING_LENS = [4224, 16384, 32768, 65536, 131072, 314721, 391000, 491061]
-_TIMING_REPS = 16
-_TIMING_ROUNDS = 15
-_TIMING_REPLAYS = 5
-
-
-def _replay_us(graph: torch.cuda.CUDAGraph, replays: int) -> float:
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(replays):
-        graph.replay()
-    end.record()
-    end.synchronize()
-    return start.elapsed_time(end) * 1000.0 / replays
-
-
-def _median(values):
-    values = sorted(values)
-    return values[len(values) // 2]
-
-
-def _capture_region(pipe, monkeypatch, kv: int, with_indexer_topk: bool, flush: torch.Tensor):
-    """CUDA graph of ``_TIMING_REPS`` x (L2 flush, full layer + shared layer attention)."""
-    with monkeypatch.context() as mp, torch.inference_mode():
-        _select(mp, with_indexer_topk)
-        md = pipe.prepare_target_step(host_cached=kv - NEXT_N, kv_offset=0)
-        positions = torch.arange(kv - NEXT_N, kv, dtype=torch.int32, device=_DEV)
-        work = []
-        for layer_idx in (LAYER_FULL, LAYER_SHARED):
-            inputs = pipe.make_inputs(seed=500 + layer_idx, n=NEXT_N, positions=positions)
-            work.append((layer_idx, inputs, pipe.project(layer_idx, inputs)))
-
-        def region():
-            for layer_idx, inputs, intermediates in work:
-                pipe.attend(layer_idx, md, inputs, intermediates)
-
-        engaged = indexer_topk.DECISION_COUNTS.get("engaged", 0)
-        region()  # compiles the kernels of this shape outside the capture
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph, capture_error_mode="thread_local"):
-            for _ in range(_TIMING_REPS):
-                flush.zero_()
-                region()
-        torch.cuda.synchronize()
-        engaged = indexer_topk.DECISION_COUNTS.get("engaged", 0) - engaged
-    assert engaged == ((_TIMING_REPS + 1) if with_indexer_topk else 0), engaged
-    return graph, work
-
-
-def _time_arm(pipe, monkeypatch, kv: int, with_indexer_topk: bool, flush, flush_graph):
-    """Per-region microseconds (median, spread) of one arm at one KV length.
-
-    The graph is replayed before the next ``prepare()``: the non-graph metadata
-    re-derives per-step tensors (such as the cached pool view) there, which a
-    graph captured earlier would still point at.
-    """
-    graph, work = _capture_region(pipe, monkeypatch, kv, with_indexer_topk, flush)
-    flush_samples, samples = [], []
-    for _ in range(_TIMING_ROUNDS):
-        flush_samples.append(_replay_us(flush_graph, _TIMING_REPLAYS))
-        samples.append(_replay_us(graph, _TIMING_REPLAYS))
-    region_us = (_median(samples) - _median(flush_samples)) / _TIMING_REPS
-    spread_us = (max(samples) - min(samples)) / _TIMING_REPS
-    del graph, work
-    torch.cuda.synchronize()
-    torch.cuda.empty_cache()
-    return region_us, spread_us
-
-
-@pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_INDEXER_TOPK_PIPELINE_TIMING") != "1",
-    reason="set GLM_MOE_DSA_INDEXER_TOPK_PIPELINE_TIMING=1 to time",
-)
-def test_indexer_topk_region_timing(monkeypatch):
-    """Attention region of a full + shared layer pair, indexer_topk off vs on, in CUDA graphs.
-
-    The region is ``_forward_dsa_attn`` of both layers with the served
-    neighbours of the indexer (indexer K scatter, shared-layer copy, remap,
-    RoPE/FP8-query/KV append, sparse MLA). Each capture flushes L2 before
-    every region (the served K caches of 21 full layers are not L2 resident);
-    the flush alone is timed in its own graph and subtracted. Native is timed
-    before and after indexer_topk at each length; their difference is the drift.
-    """
-    pipe = GlmDsaPipeline(max(TIMING_LENS) + TOKENS_PER_BLOCK)
-    flush = torch.empty(1 << 29, dtype=torch.uint8, device=_DEV)
-    results = []
-    try:
-        flush_graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(flush_graph, capture_error_mode="thread_local"):
-            for _ in range(_TIMING_REPS):
-                flush.zero_()
-        for kv in TIMING_LENS:
-            native_1, spread_1 = _time_arm(pipe, monkeypatch, kv, False, flush, flush_graph)
-            indexer_topk_us, indexer_topk_spread = _time_arm(
-                pipe, monkeypatch, kv, True, flush, flush_graph
-            )
-            native_2, spread_2 = _time_arm(pipe, monkeypatch, kv, False, flush, flush_graph)
-            native_us = (native_1 + native_2) / 2
-            row = {
-                "kv": kv,
-                "base_features": BASE_FEATURES,
-                "native_region_us": round(native_us, 3),
-                "indexer_topk_region_us": round(indexer_topk_us, 3),
-                "delta_us": round(indexer_topk_us - native_us, 3),
-                "native_drift_us": round(native_2 - native_1, 3),
-                "native_spread_us": round(max(spread_1, spread_2), 3),
-                "indexer_topk_spread_us": round(indexer_topk_spread, 3),
-            }
-            results.append(row)
-            print(f"[indexer_topk-timing] {json.dumps(row)}")
-    finally:
-        pipe.shutdown()
-    assert len(results) == len(TIMING_LENS)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py
deleted file mode 100644
index 3e737bd68d21..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/test_indexer_topk_policy.py
+++ /dev/null
@@ -1,243 +0,0 @@
-# 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.
-"""Eligibility table of the indexer_topk indexer decode package (no GPU needed).
-
-``policy.py`` is loaded by file path so this test needs neither torch nor the
-compiled extensions.
-"""
-
-import importlib.util
-from pathlib import Path
-
-import pytest
-
-_POLICY = (
-    Path(__file__).resolve().parents[9]
-    / "tensorrt_llm"
-    / "_torch"
-    / "_experimental"
-    / "modeling_v2"
-    / "models"
-    / "glm_moe_dsa"
-    / "shared"
-    / "kernels"
-    / "indexer_topk"
-    / "policy.py"
-)
-
-
-def _load_policy():
-    spec = importlib.util.spec_from_file_location("indexer_topk_policy", _POLICY)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-# Served GEN: max_seq_len 2**20 with MTP K = 5 and the overlap
-# scheduler gives an indexer max_seq_len of 2**20 + 14, and KV cache manager V2
-# sizes the block table at round_up(ceil((2**20 + 14 + 10) / 64), 4) columns.
-SERVED_MAX_SEQ_LEN = (1 << 20) + 14
-V2_BLOCK_TABLE_COLS = 16388
-# KV cache manager V1 at the same max_seq_len: ceil((2**20 + 14) / 64) columns.
-V1_BLOCK_TABLE_COLS = 16385
-
-# Served GEN: sm103, CuTe-DSL logits, FP8 indexer keys, self-sampling GVR,
-# one request with 1 + 5 verify rows, 32 heads of dim 128, 64-token pages,
-# TP8 without attention DP, the served indexer max_seq_len.
-_SERVED_TP8 = dict(
-    sm_version=103,
-    cute_dsl_available=True,
-    native_uses_cute_dsl_logits=True,
-    use_fp4_indexer=False,
-    use_gvr_emission=False,
-    needs_gvr_prior=False,
-    is_ragged_verify=False,
-    expanded_mtp_buffers=False,
-    in_mtp_draft_loop=False,
-    attention_dp=False,
-    topk_decode_impl="cute_dsl_gvr",
-    topk_self_sampling=True,
-    index_topk=2048,
-    compress_ratio=1,
-    num_heads=32,
-    head_dim=128,
-    query_dtype="torch.float8_e4m3fn",
-    query_shape=(6, 32, 128),
-    query_contiguous=True,
-    weights_dtype="torch.float32",
-    weights_shape=(6, 32),
-    weights_contiguous=True,
-    k_cache_dtype="torch.uint8",
-    k_cache_shape=(16384, 64, 1, 132),
-    block_table_cols=V2_BLOCK_TABLE_COLS,
-    max_seq_len=SERVED_MAX_SEQ_LEN,
-    num_contexts=0,
-    token_offset=0,
-    num_gen_tokens=6,
-    num_generations=1,
-    uniform_verify_width=True,
-)
-
-
-def _with_shape(heads=32, head_dim=128, page=64, next_n=6, batch=1, **extra):
-    """Served decision inputs with a different indexer/verify shape."""
-    rows = next_n * batch
-    return {
-        **_SERVED_TP8,
-        "num_heads": heads,
-        "head_dim": head_dim,
-        "query_shape": (rows, heads, head_dim),
-        "weights_shape": (rows, heads),
-        "k_cache_shape": (16384, page, 1, head_dim + 4),
-        "num_gen_tokens": rows,
-        "num_generations": batch,
-        **extra,
-    }
-
-
-def test_the_served_tp8_shape_is_eligible():
-    decision = _load_policy().decide(**_SERVED_TP8)
-    assert decision.eligible
-    assert decision.reason == ""
-
-
-@pytest.mark.parametrize(
-    "override, reason_part",
-    [
-        ({"sm_version": 90}, "sm90"),
-        # SM100 has the same kernels but has not been validated on a GPU.
-        ({"sm_version": 100}, "sm100"),
-        ({"sm_version": 120}, "sm120"),
-        ({"cute_dsl_available": False}, "DSL"),
-        ({"native_uses_cute_dsl_logits": False}, "CuTe-DSL"),
-        ({"use_fp4_indexer": True}, "FP4"),
-        ({"use_gvr_emission": True}, "emission"),
-        ({"needs_gvr_prior": True}, "prior"),
-        ({"is_ragged_verify": True}, "ragged"),
-        ({"expanded_mtp_buffers": True}, "expanded"),
-        ({"in_mtp_draft_loop": True}, "draft"),
-        ({"attention_dp": True}, "data parallelism"),
-        ({"topk_decode_impl": "cuda_radix"}, "top-k engine"),
-        ({"topk_decode_impl": "cute_dsl_radix"}, "top-k engine"),
-        ({"topk_self_sampling": False}, "top-k engine"),
-        ({"index_topk": 512}, "index_topk"),
-        ({"compress_ratio": 4}, "compress_ratio"),
-        ({"query_dtype": "torch.bfloat16"}, "query dtype"),
-        ({"query_contiguous": False}, "layout"),
-        ({"query_shape": (6, 32, 64)}, "layout"),
-        ({"weights_dtype": "torch.bfloat16"}, "weights dtype"),
-        ({"weights_contiguous": False}, "layout"),
-        ({"weights_shape": (6, 16)}, "layout"),
-        ({"k_cache_dtype": "torch.float8_e4m3fn"}, "K cache layout"),
-        ({"k_cache_shape": (16384, 64, 1, 68)}, "K cache layout"),
-        ({"k_cache_shape": (16384, 64, 132)}, "K cache view"),
-        ({"max_seq_len": SERVED_MAX_SEQ_LEN + 1}, "max_seq_len"),
-        ({"max_seq_len": 1 << 21}, "max_seq_len"),
-        ({"max_seq_len": 0}, "max_seq_len"),
-        # V1 table at the served max_seq_len: the last 128-token tile of a
-        # 2**20 + 1 .. 2**20 + 14 token request would read one entry past the row.
-        ({"block_table_cols": V1_BLOCK_TABLE_COLS}, "block table width"),
-        ({"block_table_cols": 0}, "block table width"),
-        ({"num_contexts": 1}, "context requests"),
-        ({"token_offset": 1}, "exceed"),
-        ({"num_generations": 0}, "no generation"),
-        (
-            {
-                "num_gen_tokens": 7,
-                "num_generations": 2,
-                "query_shape": (7, 32, 128),
-                "weights_shape": (7, 32),
-            },
-            "non-uniform",
-        ),
-        ({"uniform_verify_width": False}, "non-uniform"),
-    ],
-)
-def test_native_fallback_cases(override, reason_part):
-    decision = _load_policy().decide(**{**_SERVED_TP8, **override})
-    assert not decision.eligible
-    assert reason_part in decision.reason
-
-
-@pytest.mark.parametrize(
-    "shape, reason_part",
-    [
-        # Constructor-invalid for the rows kernel (8 cached weights > 4 heads):
-        # must fall back before any kernel is built.
-        (dict(heads=4, next_n=4), "num_heads"),
-        (dict(heads=16), "num_heads"),
-        (dict(heads=64), "num_heads"),
-        (dict(head_dim=64), "head_dim"),
-        (dict(page=16), "tokens_per_block"),
-        (dict(page=32), "tokens_per_block"),
-        (dict(page=128), "tokens_per_block"),
-        (dict(next_n=1), "next_n"),
-        (dict(next_n=4), "next_n"),
-        (dict(next_n=8), "next_n"),
-        (dict(batch=5), "generation requests"),
-    ],
-    ids=lambda v: "-".join(f"{k}{w}" for k, w in v.items()) if isinstance(v, dict) else v,
-)
-def test_unverified_shapes_fall_back(shape, reason_part):
-    decision = _load_policy().decide(**_with_shape(**shape))
-    assert not decision.eligible
-    assert reason_part in decision.reason
-
-
-def test_envelope_is_the_served_bound():
-    policy = _load_policy()
-    assert policy.MAX_INDEXER_SEQ_LEN == SERVED_MAX_SEQ_LEN
-    # net max_seq_len + (K) overlap lead + (K - 1) extra KV tokens + (K) tokens
-    # per generation step beyond the golden token, for K = 5.
-    assert policy.MAX_INDEXER_SEQ_LEN == (1 << 20) + 5 + 4 + 5
-
-
-@pytest.mark.parametrize(
-    "max_seq_len, block_table_cols",
-    [
-        (SERVED_MAX_SEQ_LEN, V2_BLOCK_TABLE_COLS),
-        # V1 and V2 at the 2**20 net envelope without MTP slack.
-        (1 << 20, 16384),
-        (1 << 20, 16388),
-        (131072, 2048),
-    ],
-)
-def test_served_lengths_and_table_widths_are_eligible(max_seq_len, block_table_cols):
-    decision = _load_policy().decide(
-        **{**_SERVED_TP8, "max_seq_len": max_seq_len, "block_table_cols": block_table_cols}
-    )
-    assert decision.eligible, decision.reason
-
-
-@pytest.mark.parametrize("batch", [1, 2, 3, 4])
-def test_verified_batch_sizes(batch):
-    assert _load_policy().decide(**_with_shape(batch=batch)).eligible
-
-
-def test_eligible_set_is_the_verified_envelope():
-    """No shape outside the GPU-verified one is eligible."""
-    policy = _load_policy()
-    eligible = set()
-    for sm in (90, 100, 103, 120):
-        for heads in (4, 8, 16, 32, 64):
-            for head_dim in (64, 128, 256):
-                for page in (16, 32, 64, 128):
-                    for next_n in range(1, 10):
-                        for batch in range(1, 9):
-                            inputs = _with_shape(heads, head_dim, page, next_n, batch)
-                            inputs["sm_version"] = sm
-                            if policy.decide(**inputs).eligible:
-                                eligible.add((sm, heads, head_dim, page, next_n, batch))
-    assert eligible == {(103, 32, 128, 64, 6, b) for b in range(1, 5)}
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py
deleted file mode 100644
index 81e6f86dfff8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/test_oproj_stage.py
+++ /dev/null
@@ -1,839 +0,0 @@
-# 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.
-"""oproj_weight_staging ``oproj_weight_staging`` (single GPU): the staging o_proj against the served
-attention_output_tail.
-
-Oracle: the served attention_output_tail/attention_output_into_hop kernels
-(``attention_output_tail.kernels.attn_tail``, its ``stream_oproj`` and ``stream_oproj_into_hop``) on
-the same inputs. oproj_weight_staging changes no arithmetic, so every check is bit-for-bit
-(``torch.equal``), on real GLM-5.2 TP8 weight slices (``tail_fixtures.load_tail_weights``) and
-real-weight latents (``tail_fixtures.mix_latent``):
-
-* every staging configuration (L2 prefetch or shared-memory staging, at
-  entry or once the early rows landed, from the o_proj or from the W_UV
-  kernel, cache-policy hints, partial staging), on the served schedule
-  (``INTO_HOP_SPARSE_DECODE_SCHEDULE``: 128 CTAs, 32 early rows) and on other ``early_rows``
-  / grids; the BF16 ``v`` intermediate as well;
-* the attention_output_into_hop publish (a device buffer stands in for the post_attention_hop
-  transport, as in ``attention_output_tail/test_into_hop_producer.py``): same words, same untouched
-  sentinels;
-* padded rows (live < out rows, also out rows 9..16);
-* CUDA-graph capture and replay with fresh contents, alternating graphs of
-  different row counts;
-* a delayed producer of the latent (a kernel that releases its PDL
-  dependents at entry and writes the latent only after a ``%globaltimer``
-  spin; the tail's W_UV kernel waits for it) and a delayed producer of the
-  o_proj's own input ``v`` (the o_proj is its direct PDL dependent, so every
-  o_proj CTA is resident and done staging long before its wait returns), and
-  a co-running kernel on another stream;
-* ``%globaltimer`` stamps of the diagnostic build are ordered;
-* ``StagedAttnTail`` on the real DSA ``MLA`` module: the attention_output_tail protocol (arm,
-  defer_w_uv, project, into-hop request) is unchanged and the partial equals
-  attention_output_tail's; schedules oproj_weight_staging does not take keep attention_output_tail's
-  tail; geometry declines run attention_output_tail.
-
-Opt-in (``GLM_MOE_DSA_OPROJ_WEIGHT_STAGING_ADAPTER_PATCH_APPLIED=1``): the adopted attention
-(``GlmMLA.adopt``) wraps the tail;
-every set of fast paths containing ``oproj_weight_staging`` that a plan can run engages it (the
-closure is the engaged domain, so no runnable set leaves it inert); deferred_input_prep's gate
-maps ``oproj_weight_staging`` to its directories and side selector.
-"""
-
-import dataclasses
-import os
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-import triton
-import triton.language as tl
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "attention_output_tail"))
-import tail_fixtures  # noqa: E402
-import test_attn_tail_mla_adapter as tail_adapter  # noqa: E402
-
-sys.path.insert(0, os.path.dirname(_HERE))
-import plan_testlib  # noqa: E402
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
-    dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.kernels import (  # noqa: E402, E501
-    INTO_HOP_SCHEDULE,
-    INTO_HOP_SPARSE_DECODE_SCHEDULE,
-    TAIL_PREFETCH_SCHEDULE,
-    TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE,
-    TailConfig,
-    attn_tail,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.attention_output_tail.tail import (  # noqa: E402, E501
-    AttnTail,
-    make_attn_tail,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging import (  # noqa: E402, E501
-    DEFAULT_STAGE,
-    StageConfig,
-    StagedAttnTail,
-    stage_tail,
-    wrap_attn_tail,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging.cute_oproj_stage import (  # noqa: E402, E501
-    TRACE_POINTS,
-    stream_oproj_stage,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging.kernels import (  # noqa: E402, E501
-    check_stage_config,
-    wuv_stage_geometry,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.oproj_weight_staging.tail import (  # noqa: E402, E501
-    config_reason,
-    geometry_reason,
-)
-
-_CKPT = tail_fixtures.checkpoint_dir()
-_PATCHED = os.environ.get("GLM_MOE_DSA_OPROJ_WEIGHT_STAGING_ADAPTER_PATCH_APPLIED") == "1"
-needs_patch = pytest.mark.skipif(not _PATCHED, reason="opt-in test of the adopted attention")
-# Served GEN selector (without deferred_input_prep, which the eligibility test adds).
-B1_GEN = (
-    "absorbed_query,indexer_topk,attention_output_tail,post_moe_hop,"
-    "draft_front,router_gemm,post_attention_hop,draft_tail,"
-    "sparse_mla_decode,draft_absorbed_query,rope_kv_append,draft_loop,"
-    "shared_indexer_absorbed_query,routed_moe,cross_rank_topk,"
-    "attention_output_tail_prefetch,attention_output_into_hop,"
-    "routed_moe_into_hop,deferred_norm,dense_mlp"
-)
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="oproj_weight_staging targets SM100/SM103",
-    ),
-    pytest.mark.skipif(not (_CKPT and os.path.isdir(_CKPT)), reason="checkpoint not available"),
-]
-
-TP_SIZE = 8
-LOCAL_HEADS = 8
-H = 6144
-RANKS = 8
-MAX_ROWS = 8
-BUF_WORDS = MAX_ROWS * RANKS * H * 2 // 4
-SENTINEL = -(2**31)
-
-S = StageConfig
-ARMS = {
-    "none": S(stage="none"),
-    "l2_entry": S(stage="l2", when="entry"),
-    "l2_landed": S(stage="l2", when="landed"),
-    "l2_landed_el": S(stage="l2", when="landed", prefetch_policy="evict_last"),
-    "l2_half": S(stage="l2", when="landed", rows=8),
-    "smem_entry": S(stage="smem", when="entry"),
-    "smem_landed": S(stage="smem", when="landed"),
-    "smem_half": S(stage="smem", when="landed", rows=8),
-    "wuv": S(stage="l2", host="wuv"),
-    "wuv_el": S(stage="l2", host="wuv", prefetch_policy="evict_last"),
-    "ef": S(stage="none", copy_policy="evict_first"),
-    "l2_landed_ef": S(stage="l2", when="landed", copy_policy="evict_first"),
-    "smem_landed_ef": S(stage="smem", when="landed", copy_policy="evict_first"),
-    "smem_chain": S(stage="smem", when="chain"),
-    "smem_chain_e16": S(stage="smem", when="chain", early_rows=16),
-    "smem_chain_e8_ef": S(stage="smem", when="chain", early_rows=8, copy_policy="evict_first"),
-    "default": DEFAULT_STAGE,
-}
-PUBLISH_ARMS = (
-    "none",
-    "l2_landed",
-    "smem_landed",
-    "wuv",
-    "l2_landed_ef",
-    "smem_chain_e16",
-    "default",
-)
-
-_W = {}
-_COMP = {}
-
-
-def _weights(layer, rank):
-    if layer not in _W:
-        _W[layer] = {
-            r: (w_uv.cuda().contiguous(), w_o.cuda().contiguous())
-            for r, (w_uv, w_o) in tail_fixtures.load_tail_weights(
-                _CKPT, layer, range(TP_SIZE), TP_SIZE
-            ).items()
-        }
-    return _W[layer][rank]
-
-
-def _latent(layer, rows, seed):
-    if layer not in _COMP:
-        _COMP[layer] = tail_fixtures.compressed_kv_rows(_CKPT, layer)
-    return tail_fixtures.mix_latent(_COMP[layer], rows, LOCAL_HEADS, seed=seed)
-
-
-def _transport(cur=1, rank=3):
-    mc = torch.full((3 * BUF_WORDS,), SENTINEL, dtype=torch.int32, device="cuda")
-    flags = torch.tensor([cur, 0, 0, 0], dtype=torch.int32, device="cuda")
-    return SimpleNamespace(mc_words=mc, flags=flags, rank=rank, cur=cur)
-
-
-def _pair(layer, rank, rows, out_rows, cfg, stage, seed, into_hop_pair=False):
-    """Run attention_output_tail's tail and oproj_weight_staging's tail on the same inputs; return both outputs
-    (and the v intermediates, and the two transports when ``into_hop_pair``)."""
-    w_uv, w_o = _weights(layer, rank)
-    lat = _latent(layer, rows, seed)
-    v_attn_tail = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    v_oproj_stage = torch.empty_like(v_attn_tail)
-    tr_attn_tail = _transport() if into_hop_pair else None
-    tr_oproj_stage = _transport() if into_hop_pair else None
-    ref = attn_tail(
-        lat,
-        w_uv,
-        w_o,
-        rows,
-        out_rows=out_rows,
-        v_scratch=v_attn_tail,
-        config=cfg,
-        hop_transport=tr_attn_tail,
-    )
-    got = stage_tail(
-        lat,
-        w_uv,
-        w_o,
-        rows,
-        out_rows=out_rows,
-        v_scratch=v_oproj_stage,
-        config=cfg,
-        stage=stage,
-        hop_transport=tr_oproj_stage,
-    )
-    torch.cuda.synchronize()
-    return ref, got, v_attn_tail, v_oproj_stage, tr_attn_tail, tr_oproj_stage
-
-
-@pytest.mark.parametrize("arm", list(ARMS))
-@pytest.mark.parametrize("layer,rank", [(3, 0), (77, 7)])
-@pytest.mark.parametrize("rows", [1, 6, 8])
-def test_served_schedule_bit_identical(arm, layer, rank, rows):
-    ref, got, v_attn_tail, v_oproj_stage, _, _ = _pair(
-        layer, rank, rows, rows, INTO_HOP_SPARSE_DECODE_SCHEDULE, ARMS[arm], seed=10 * rows + rank
-    )
-    assert torch.equal(v_attn_tail, v_oproj_stage)
-    assert torch.equal(ref, got), (arm, layer, rank, rows)
-
-
-@pytest.mark.parametrize("arm", PUBLISH_ARMS)
-@pytest.mark.parametrize("live,out_rows", [(1, 1), (6, 6), (8, 8), (6, 8), (1, 4)])
-def test_into_hop_publish_bit_identical(arm, live, out_rows):
-    ref, got, _, _, tr_attn_tail, tr_oproj_stage = _pair(
-        40,
-        3,
-        live,
-        out_rows,
-        INTO_HOP_SPARSE_DECODE_SCHEDULE,
-        ARMS[arm],
-        seed=live + out_rows,
-        into_hop_pair=True,
-    )
-    assert torch.equal(ref, got)
-    assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words), arm
-    buf = tr_oproj_stage.mc_words[
-        tr_oproj_stage.cur * BUF_WORDS : (tr_oproj_stage.cur + 1) * BUF_WORDS
-    ].view(MAX_ROWS, RANKS, H // 2)
-    assert not bool((buf[:out_rows, tr_oproj_stage.rank] == SENTINEL).any())
-
-
-@pytest.mark.parametrize(
-    "arm", ["none", "l2_entry", "l2_landed", "smem_landed", "smem_chain_e16", "wuv", "default"]
-)
-@pytest.mark.parametrize("live,out_rows,publish", [(6, 6, True), (1, 4, True), (6, 12, False)])
-def test_synthetic_inputs_bit_identical(arm, live, out_rows, publish):
-    """Every input written by plain stores (``torch.randn``/``full``), for
-    compute-sanitizer initcheck: the real-weight fixtures run cuBLAS GEMMs
-    whose outputs initcheck reports as uninitialized when torch reads them."""
-    gen = torch.Generator(device="cuda").manual_seed(live + out_rows)
-    w_uv = (torch.randn((8, 256, 512), generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
-    w_o = (torch.randn((H, 2048), generator=gen, device="cuda") * 0.02).to(torch.bfloat16)
-    lat = torch.randn((live, 8, 512), generator=gen, device="cuda").to(torch.bfloat16)
-    tr_attn_tail = _transport() if publish else None
-    tr_oproj_stage = _transport() if publish else None
-    ref = attn_tail(
-        lat,
-        w_uv,
-        w_o,
-        live,
-        out_rows=out_rows,
-        config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
-        hop_transport=tr_attn_tail,
-    )
-    got = stage_tail(
-        lat,
-        w_uv,
-        w_o,
-        live,
-        out_rows=out_rows,
-        config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
-        stage=ARMS[arm],
-        hop_transport=tr_oproj_stage,
-    )
-    torch.cuda.synchronize()
-    assert torch.equal(ref, got)
-    if publish:
-        assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words)
-
-
-@pytest.mark.parametrize(
-    "cfg_name,cfg",
-    [
-        ("attn_tail_default", TailConfig()),
-        ("g128_e24", TailConfig(oproj_grid=128, early_rows=24)),
-        ("g128_e40", TailConfig(oproj_grid=128, early_rows=40)),
-        ("g128_e48", TailConfig(oproj_grid=128, early_rows=48)),
-        ("attention_output_into_hop", INTO_HOP_SCHEDULE),
-    ],
-)
-@pytest.mark.parametrize(
-    "arm", ["l2_entry", "l2_landed", "smem_landed", "smem_chain", "default", "ef"]
-)
-def test_other_schedules_bit_identical(cfg_name, cfg, arm):
-    ref, got, v_attn_tail, v_oproj_stage, _, _ = _pair(3, 5, 6, 6, cfg, ARMS[arm], seed=7)
-    assert torch.equal(v_attn_tail, v_oproj_stage)
-    assert torch.equal(ref, got), (cfg_name, arm)
-
-
-@pytest.mark.parametrize("early", [16, 24, 40])
-def test_early_rows_override(early):
-    stage = S(stage="l2", when="landed", early_rows=early)
-    cfg = dataclasses.replace(INTO_HOP_SPARSE_DECODE_SCHEDULE, early_rows=early)
-    w_uv, w_o = _weights(3, 2)
-    lat = _latent(3, 6, 99)
-    ref = attn_tail(lat, w_uv, w_o, 6, config=cfg)
-    got = stage_tail(lat, w_uv, w_o, 6, config=INTO_HOP_SPARSE_DECODE_SCHEDULE, stage=stage)
-    torch.cuda.synchronize()
-    assert torch.equal(ref, got)
-
-
-@pytest.mark.parametrize("live,out_rows", [(6, 12), (8, 16), (1, 9)])
-@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
-def test_wide_padded_rows(live, out_rows, arm):
-    ref, got, _, _, _, _ = _pair(
-        2, 1, live, out_rows, INTO_HOP_SPARSE_DECODE_SCHEDULE, ARMS[arm], seed=live
-    )
-    assert torch.equal(ref, got)
-    assert bool((got[live:] == 0).all())
-
-
-@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
-def test_graph_replay_changing_rows(arm):
-    """Graphs for 1 live row (4 out rows) and 6 live rows (8 out rows),
-    replayed alternately with fresh latents; each replay equals attention_output_tail."""
-    w_uv, w_o = _weights(3, 0)
-    stage = ARMS[arm]
-    graphs = {}
-    for rows, out_rows in ((1, 4), (6, 8)):
-        lat = _latent(3, rows, 0).clone()
-        stage_tail(
-            lat,
-            w_uv,
-            w_o,
-            rows,
-            out_rows=out_rows,
-            config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
-            stage=stage,
-        )
-        torch.cuda.synchronize()
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            out = stage_tail(
-                lat,
-                w_uv,
-                w_o,
-                rows,
-                out_rows=out_rows,
-                config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
-                stage=stage,
-            )
-        graphs[rows] = (g, lat, out, out_rows)
-    for step in range(12):
-        rows = (1, 6)[step % 2]
-        g, lat, out, out_rows = graphs[rows]
-        lat.copy_(_latent(3, rows, 100 + step))
-        g.replay()
-        ref = attn_tail(
-            lat, w_uv, w_o, rows, out_rows=out_rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE
-        )
-        torch.cuda.synchronize()
-        assert torch.equal(out, ref), (arm, step)
-
-
-@triton.jit
-def _globaltimer(like):
-    return tl.inline_asm_elementwise(
-        "mov.u64 $0, %globaltimer;", "=l,r", [like], dtype=tl.int64, is_pure=False, pack=1
-    )
-
-
-@triton.jit
-def _slow_copy_kernel(src, dst, n, delay_ns, stamps, BLOCK: tl.constexpr):
-    """One CTA, launched with PDL (Triton emits the ``griddepcontrol``
-    instructions only for such launches): waits for its predecessor,
-    releases its own dependents, spins ``delay_ns`` on ``%globaltimer``, then
-    copies ``src`` into ``dst`` (the latent). One CTA leaves SMs free, so
-    o_proj CTAs (197,400 B of shared memory each) can become resident next
-    to the parked W_UV grid."""
-    tl.extra.cuda.gdc_wait()
-    tl.extra.cuda.gdc_launch_dependents()
-    pid = tl.program_id(0)
-    start = _globaltimer(pid)
-    now = start
-    while now - start < delay_ns:
-        now = _globaltimer(pid)
-    for base in range(0, n, BLOCK):
-        offs = base + tl.arange(0, BLOCK)
-        mask = offs < n
-        tl.store(dst + offs, tl.load(src + offs, mask=mask), mask=mask)
-    # Start of the spin and end of the latent's stores (``%globaltimer``).
-    tl.store(stamps, start)
-    tl.store(stamps + 1, _globaltimer(pid))
-
-
-@triton.jit
-def _hold_kernel(delay_ns, BLOCK: tl.constexpr):
-    pid = tl.program_id(0)
-    start = _globaltimer(pid)
-    now = start
-    while now - start < delay_ns:
-        now = _globaltimer(pid)
-
-
-def _hold_stream(delay_ns=2_000_000):
-    """Occupy the stream with one spinning CTA (no PDL) while the host
-    enqueues the launches that follow. Without it the host's launch latency
-    (tens of us per Triton / CuTe launch) exceeds the producer's delay, the
-    producer has completed before its PDL dependents are even enqueued, and
-    they never park in ``griddepcontrol.wait``."""
-    _hold_kernel[(1,)](delay_ns, BLOCK=1)
-
-
-def _slow_produce(src, dst, delay_ns, stamps=None):
-    if stamps is None:
-        stamps = torch.zeros(2, dtype=torch.int64, device=src.device)
-    _slow_copy_kernel[(1,)](src, dst, src.numel(), delay_ns, stamps, BLOCK=1024, launch_pdl=True)
-    return stamps
-
-
-@pytest.mark.parametrize(
-    "arm",
-    ["l2_entry", "l2_landed", "smem_landed", "wuv", "l2_landed_ef", "smem_chain_e16", "default"],
-)
-@pytest.mark.parametrize("delay_ns", [0, 3000, 20000])
-def test_delayed_producer(arm, delay_ns):
-    """The latent's producer is late (it releases its PDL dependents, then
-    spins ``delay_ns`` before writing the latent; a held stream lets the
-    host enqueue the tail before the producer starts). Outputs equal attention_output_tail's
-    under the same producer, and no o_proj CTA leaves its wait before the
-    producer's last store (``%globaltimer`` stamps of both kernels)."""
-    w_uv, w_o = _weights(40, 6)
-    rows = 6
-    src = _latent(40, rows, 5).contiguous()
-    lat_attn_tail = torch.zeros_like(src)
-    lat_oproj_stage = torch.zeros_like(src)
-    tr = torch.zeros(128 * TRACE_POINTS, dtype=torch.int64, device="cuda")
-    for _ in range(2):
-        _hold_stream()
-        _slow_produce(src, lat_attn_tail, delay_ns)
-        ref = attn_tail(lat_attn_tail, w_uv, w_o, rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE)
-        _hold_stream()
-        ps = _slow_produce(src, lat_oproj_stage, delay_ns)
-        got = stage_tail(
-            lat_oproj_stage,
-            w_uv,
-            w_o,
-            rows,
-            config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
-            stage=ARMS[arm],
-            trace=tr,
-        )
-        torch.cuda.synchronize()
-        assert torch.equal(ref, got), (arm, delay_ns)
-        st = tr.view(128, TRACE_POINTS).cpu()
-        p_start, p_end = (int(x) for x in ps.cpu())
-        assert p_end - p_start >= delay_ns
-        # Every o_proj CTA passed its wait only after the producer's stores
-        # (through W_UV's completion), however early it became resident.
-        assert int(st[:, 1].min()) >= p_end, (arm, delay_ns)
-    parked = (st[:, 1] - st[:, 0]).float()
-    print(
-        f"delayed producer {delay_ns} ns, {arm}: o_proj resident "
-        f"{float(parked.mean()):.0f} ns (mean) / {float(parked.max()):.0f} "
-        f"ns (max) before its wait returned"
-    )
-
-
-@pytest.mark.parametrize(
-    "arm",
-    ["none", "l2_entry", "l2_landed", "smem_landed", "l2_landed_ef", "smem_chain_e16", "default"],
-)
-@pytest.mark.parametrize("publish", [False, True])
-def test_parked_oproj_bit_identical(arm, publish):
-    """The o_proj is the direct PDL dependent of a one-CTA producer that
-    releases its dependents at entry and writes ``v`` only after a 20 us
-    spin (behind a held stream, so the host has enqueued the o_proj before
-    the producer starts). All 128 o_proj CTAs fit next to it, so each one
-    issues its pre-wait copies and staging and then stays parked in
-    ``griddepcontrol.wait`` for most of the spin (the case where the staging
-    has long completed when the wait returns). Outputs (and the into-hop words) equal the served
-    tail's on the same latent, for the plain and the stamped build; no CTA leaves its wait
-    before the producer's last store; the median CTA was parked at least
-    half the delay."""
-    delay_ns = 20000
-    w_uv, w_o = _weights(40, 6)
-    rows = 6
-    stage = ARMS[arm]
-    early = (
-        INTO_HOP_SPARSE_DECODE_SCHEDULE.early_rows if stage.early_rows is None else stage.early_rows
-    )
-    lat = _latent(40, rows, 21)
-    v_src = torch.empty((rows, 2048), dtype=torch.bfloat16, device="cuda")
-    tr_attn_tail = _transport() if publish else None
-    ref = attn_tail(
-        lat,
-        w_uv,
-        w_o,
-        rows,
-        v_scratch=v_src,
-        config=INTO_HOP_SPARSE_DECODE_SCHEDULE,
-        hop_transport=tr_attn_tail,
-    )
-    torch.cuda.synchronize()
-    for trace in (False, True, True):
-        # The first stamped run compiles that build (host time beyond the
-        # held stream); the parked time is asserted on the second.
-        tr = torch.zeros(128 * TRACE_POINTS, dtype=torch.int64, device="cuda") if trace else None
-        v = torch.zeros_like(v_src)
-        got = torch.empty_like(ref)
-        tr_oproj_stage = _transport() if publish else None
-        torch.cuda.synchronize()
-        _hold_stream()
-        ps = _slow_produce(v_src, v, delay_ns)
-        stream_oproj_stage(
-            v,
-            w_o,
-            got,
-            use_pdl=True,
-            early_rows=early,
-            grid=128,
-            trigger_dependents=True,
-            stage=stage.stage,
-            when=stage.when,
-            stage_rows=stage.rows,
-            prefetch_policy=stage.prefetch_policy,
-            copy_policy=stage.copy_policy,
-            hop_transport=tr_oproj_stage,
-            trace=tr,
-        )
-        torch.cuda.synchronize()
-        assert torch.equal(ref, got), (arm, publish, trace)
-        if publish:
-            assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words), (arm, trace)
-        if not trace:
-            continue
-        st = tr.view(128, TRACE_POINTS).cpu()
-        p_start, p_end = (int(x) for x in ps.cpu())
-        assert p_end - p_start >= delay_ns
-        assert int(st[:, 1].min()) >= p_end, arm
-        parked = (st[:, 1] - st[:, 0]).float()
-        print(
-            f"parked o_proj {arm} publish={publish}: resident "
-            f"{float(parked.median()):.0f} ns (median) / "
-            f"{float(parked.min()):.0f} ns (min) before its wait returned"
-        )
-    assert float(parked.median()) >= delay_ns / 2, arm
-
-
-@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
-def test_co_running_kernel_progress(arm):
-    """A long GEMM on another stream occupies SMs while the tail runs; the
-    tail completes with attention_output_tail's result (no wait in oproj_weight_staging depends on
-    co-residency)."""
-    w_uv, w_o = _weights(3, 4)
-    rows = 6
-    lat = _latent(3, rows, 17)
-    ref = attn_tail(lat, w_uv, w_o, rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE)
-    a = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
-    side = torch.cuda.Stream()
-    torch.cuda.synchronize()
-    outs = []
-    for _ in range(3):
-        with torch.cuda.stream(side):
-            c = a @ a
-        outs.append(
-            stage_tail(
-                lat, w_uv, w_o, rows, config=INTO_HOP_SPARSE_DECODE_SCHEDULE, stage=ARMS[arm]
-            )
-        )
-    torch.cuda.synchronize()
-    del c
-    for got in outs:
-        assert torch.equal(ref, got)
-
-
-def test_trace_stamps_ordered():
-    w_uv, w_o = _weights(3, 0)
-    lat = _latent(3, 6, 3)
-    tr = torch.zeros(128 * TRACE_POINTS, dtype=torch.int64, device="cuda")
-    out = stage_tail(
-        lat, w_uv, w_o, 6, config=INTO_HOP_SPARSE_DECODE_SCHEDULE, stage=ARMS["l2_landed"], trace=tr
-    )
-    torch.cuda.synchronize()
-    ref = attn_tail(lat, w_uv, w_o, 6, config=INTO_HOP_SPARSE_DECODE_SCHEDULE)
-    assert torch.equal(out, ref)
-    st = tr.view(128, TRACE_POINTS).cpu()
-    assert bool((st > 0).all())
-    for k in range(TRACE_POINTS - 1):
-        assert bool((st[:, k + 1] >= st[:, k]).all()), k
-
-
-def test_stage_rejects_bad_config():
-    w_uv, w_o = _weights(3, 0)
-    v = torch.zeros(6, 2048, dtype=torch.bfloat16, device="cuda")
-    out = torch.empty(6, H, dtype=torch.bfloat16, device="cuda")
-    for bad in (
-        dict(stage="tma"),
-        dict(when="late"),
-        dict(prefetch_policy="evict_first"),
-        dict(copy_policy="evict_last"),
-        dict(when="chain"),
-    ):
-        kwargs = dict(stage="l2", when="landed", prefetch_policy="normal", copy_policy="none")
-        kwargs.update(bad)
-        with pytest.raises(ValueError):
-            stream_oproj_stage(v, w_o, out, use_pdl=True, early_rows=32, grid=128, **kwargs)
-
-
-def test_wuv_geometry():
-    assert wuv_stage_geometry(H, 2048, 128, 32, -1) == (32, 16 * 4096)
-    assert wuv_stage_geometry(H, 2048, 128, 32, 8) == (32, 8 * 4096)
-    # 152 CTAs: 40 or 42 rows per CTA, not covered evenly.
-    assert wuv_stage_geometry(H, 2048, 152, 16, -1) is None
-    assert wuv_stage_geometry(H, 2048, 128, 48, -1) is None
-
-
-def test_the_served_stage_is_valid():
-    assert check_stage_config(DEFAULT_STAGE) is None
-    assert (DEFAULT_STAGE.stage, DEFAULT_STAGE.when, DEFAULT_STAGE.early_rows) == (
-        "smem",
-        "landed",
-        24,
-    )
-    assert check_stage_config(dataclasses.replace(DEFAULT_STAGE, stage="l2", when="chain"))
-
-
-# ---- StagedAttnTail on the real DSA MLA module (attention_output_tail protocol unchanged) -------
-
-
-def _step(tail, mla, latent, tr=None, out_rows=None):
-    from tensorrt_llm._torch.distributed import AllReduceParams
-
-    rows = latent.shape[0]
-    out_rows = rows if out_rows is None else out_rows
-    out = torch.empty((out_rows, 2048), dtype=torch.bfloat16, device="cuda")
-    meta = tail_adapter._meta(rows)
-    if tr is not None:
-        tail.request_into_hop(tr, out_rows)
-    assert tail.arm(mla, meta)
-    assert tail.defer_w_uv(mla, meta, latent, out[:rows])
-    return tail.project(mla, out, AllReduceParams(enable_allreduce=False))
-
-
-@pytest.mark.parametrize("arm", ["l2_landed", "smem_landed", "wuv", "default"])
-@pytest.mark.parametrize("rows,out_rows", [(6, 6), (6, 8), (1, 1)])
-def test_oproj_stage_tail_matches_attn_tail(arm, rows, out_rows):
-    mla = tail_adapter._mla()
-    served_tail = AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE)
-    oproj_stage_tail = wrap_attn_tail(
-        mla, AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE), ARMS[arm]
-    )
-    assert isinstance(oproj_stage_tail, StagedAttnTail)
-    lat = tail_adapter._latent(rows, seed=rows + out_rows)
-    tr_attn_tail, tr_oproj_stage = _transport(cur=0, rank=2), _transport(cur=0, rank=2)
-    ref = _step(served_tail, mla, lat, tr_attn_tail, out_rows)
-    got = _step(oproj_stage_tail, mla, lat, tr_oproj_stage, out_rows)
-    torch.cuda.synchronize()
-    assert torch.equal(ref, got)
-    assert torch.equal(tr_attn_tail.mc_words, tr_oproj_stage.mc_words)
-    assert served_tail.take_into_hop_published() and oproj_stage_tail.take_into_hop_published()
-    assert oproj_stage_tail.staging_calls == 1 and oproj_stage_tail.engaged_calls == 1
-    assert oproj_stage_tail.staging_reason is None
-    # No packet left: the native projection would follow.
-    assert oproj_stage_tail.project(mla, torch.empty(rows, 2048, device="cuda"), None) is None
-
-
-def test_oproj_stage_tail_declines_like_attn_tail():
-    mla = tail_adapter._mla()
-    tail = wrap_attn_tail(
-        mla, AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE), ARMS["l2_landed"]
-    )
-    lat = tail_adapter._latent(6, seed=0)
-    out = torch.empty((6, 2048), dtype=torch.bfloat16, device="cuda")
-    assert not tail.defer_w_uv(mla, tail_adapter._meta(6), lat, out)
-    assert not tail.arm(mla, tail_adapter._meta(6, num_contexts=1))
-    assert not tail.arm(mla, tail_adapter._meta(9))
-    assert tail.declined_calls == 2 and tail.staging_calls == 0
-    # torch.compile policy: arm declines without touching state.
-    from tensorrt_llm._torch.utils import torch_compiling
-
-    with torch_compiling(True):
-        assert not tail.arm(mla, tail_adapter._meta(6))
-    assert tail.declined_calls == 2
-
-
-def test_wrap_keeps_attn_tail_for_other_schedules():
-    mla = tail_adapter._mla()
-    for cfg in (
-        TAIL_PREFETCH_SCHEDULE,
-        TailConfig(oproj_engine="native"),
-        TailConfig(oproj_engine="fused"),
-        TailConfig(trigger_late=True),
-        TailConfig(),
-        INTO_HOP_SCHEDULE,
-    ):
-        tail = AttnTail("layer3", cfg)
-        assert config_reason(cfg, StageConfig()) is not None
-        assert wrap_attn_tail(mla, tail, StageConfig()) is tail
-    assert wrap_attn_tail(mla, None, StageConfig()) is None
-    served = AttnTail("layer3", INTO_HOP_SPARSE_DECODE_SCHEDULE)
-    assert wrap_attn_tail(mla, served, StageConfig(host="sparse_mla_decode")) is served
-    wrapped = wrap_attn_tail(mla, served, None)
-    assert isinstance(wrapped, StagedAttnTail)
-    assert wrapped.stage == DEFAULT_STAGE
-
-
-def test_geometry_decline_runs_attn_tail():
-    """W_UV host on the 152-CTA attention_output_tail default schedule cannot cover the rows:
-    the layer logs an oproj_weight_staging decline and runs attention_output_tail's own project."""
-    mla = tail_adapter._mla()
-    cfg = TailConfig()
-    assert geometry_reason(mla, cfg, ARMS["wuv"]) is not None
-    assert geometry_reason(mla, INTO_HOP_SPARSE_DECODE_SCHEDULE, ARMS["wuv"]) is None
-    tail = StagedAttnTail("layer3", cfg, ARMS["wuv"])
-    lat = tail_adapter._latent(6, seed=4)
-    got = _step(tail, mla, lat)
-    ref = _step(AttnTail("layer3", cfg), mla, lat)
-    torch.cuda.synchronize()
-    assert torch.equal(ref, got)
-    assert tail.staging_calls == 0 and tail.staging_reason is not None
-
-
-@needs_patch
-def test_adoption_attaches_oproj_stage_tail(monkeypatch):
-    """With ``oproj_weight_staging`` running on its closure (alone and with
-    attention_output_tail_prefetch, as served), the attention adoption wraps the
-    attention_output_tail of a target layer; the MTP layer keeps no tail; without
-    ``oproj_weight_staging`` the plain attention_output_tail ``AttnTail``.
-    """
-    for feats in (
-        "attention_output_tail,post_attention_hop,"
-        "sparse_mla_decode,attention_output_into_hop,"
-        "oproj_weight_staging",
-        "attention_output_tail,post_attention_hop,"
-        "sparse_mla_decode,attention_output_tail_prefetch,"
-        "attention_output_into_hop,oproj_weight_staging",
-    ):
-        plan_testlib.select(dispatch, feats, monkeypatch)
-        mla = tail_adapter._build_adopted_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
-        assert isinstance(mla._attn_tail, StagedAttnTail), feats
-        assert mla._attn_tail.config == INTO_HOP_SPARSE_DECODE_SCHEDULE
-        assert mla._attn_tail.stage == DEFAULT_STAGE
-        mtp = tail_adapter._build_adopted_mla(layer_idx=78, num_hidden_layers=78)
-        assert mtp._attn_tail is None
-    plan_testlib.select(
-        dispatch,
-        "attention_output_tail,post_attention_hop,sparse_mla_decode,"
-        "attention_output_tail_prefetch,attention_output_into_hop",
-        monkeypatch,
-    )
-    mla = tail_adapter._build_adopted_mla(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
-    assert type(mla._attn_tail) is AttnTail
-    assert mla._attn_tail.config == INTO_HOP_SPARSE_DECODE_SCHEDULE
-
-
-HOST_SCHEDULE_PATHS = (
-    "attention_output_tail",
-    "post_attention_hop",
-    "sparse_mla_decode",
-    "attention_output_into_hop",
-)
-
-
-@needs_patch
-def test_oproj_stage_runs_only_with_its_host_schedule():
-    """``oproj_weight_staging`` stages only in ``INTO_HOP_SPARSE_DECODE_SCHEDULE``
-    (attention_output_tail's tail with attention_output_into_hop and post_attention_hop behind
-    sparse_mla_decode): the plan turns it off with any of attention_output_tail,
-    post_attention_hop, sparse_mla_decode, attention_output_into_hop instead of running
-    attention_output_tail under oproj_weight_staging's name. This includes disabling only
-    ``attention_output_into_hop``: the attention_output_tail_prefetch schedule that then runs has
-    the same value as ``INTO_HOP_SPARSE_DECODE_SCHEDULE`` but no into-hop publication."""
-    assert TAIL_PREFETCH_SPARSE_DECODE_HOP_SCHEDULE == INTO_HOP_SPARSE_DECODE_SCHEDULE
-    assert dispatch.REQUIRED_FEATURES["oproj_weight_staging"] == set(HOST_SCHEDULE_PATHS)
-    for base in HOST_SCHEDULE_PATHS:
-        plan = dispatch.resolve(dispatch.ModelingV2Mode.AUTO, {dispatch.DISABLE_ENV: base})
-        assert not plan.enabled("oproj_weight_staging"), base
-    assert plan_testlib.plan_running(
-        dispatch, HOST_SCHEDULE_PATHS + ("oproj_weight_staging",)
-    ).enabled("oproj_weight_staging")
-    b1 = frozenset(B1_GEN.split(","))
-    for paths in (b1, b1 | {"deferred_input_prep"}):
-        assert plan_testlib.plan_running(dispatch, paths | {"oproj_weight_staging"}).enabled(
-            "oproj_weight_staging"
-        )
-
-
-@needs_patch
-def test_every_runnable_selection_engages_oproj_stage(monkeypatch):
-    """Over every subset of the paths ``make_attn_tail`` reads (attention_output_tail, post_attention_hop,
-    sparse_mla_decode, attention_output_tail_prefetch, attention_output_into_hop) plus
-    ``oproj_weight_staging``: each set a plan can run gives a target layer an
-    attention_output_tail with ``INTO_HOP_SPARSE_DECODE_SCHEDULE`` that ``wrap_attn_tail`` wraps, so
-    ``oproj_weight_staging`` never runs inert."""
-    import itertools
-
-    mla = SimpleNamespace(layer_idx=tail_adapter.LAYER, num_hidden_layers=78)
-    others = (
-        "attention_output_tail",
-        "post_attention_hop",
-        "sparse_mla_decode",
-        "attention_output_tail_prefetch",
-        "attention_output_into_hop",
-    )
-    runnable = []
-    for n in range(len(others) + 1):
-        for combo in itertools.combinations(others, n):
-            ids = combo + ("oproj_weight_staging",)
-            if not dispatch.REQUIRED_FEATURES["oproj_weight_staging"] <= set(ids):
-                continue
-            plan_testlib.select(dispatch, ids, monkeypatch)
-            tail = wrap_attn_tail(mla, make_attn_tail(mla))
-            assert isinstance(tail, StagedAttnTail), ids
-            assert tail.config == INTO_HOP_SPARSE_DECODE_SCHEDULE, ids
-            runnable.append(",".join(ids))
-    assert sorted(runnable) == sorted(
-        [
-            "attention_output_tail,post_attention_hop,sparse_mla_decode,attention_output_tail_prefetch,attention_output_into_hop,oproj_weight_staging",
-            "attention_output_tail,post_attention_hop,sparse_mla_decode,attention_output_into_hop,oproj_weight_staging",
-        ]
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py
deleted file mode 100644
index bcb6ffa7a21c..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/plan_testlib.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# 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.
-"""Run a chosen set of the target's fast paths in a test.
-
-``select(dispatch, paths)`` resolves and installs the plan that runs exactly ``paths``: every
-other fast path is listed in ``TRTLLM_STAIRCASE_GLM_LL_DISABLE``, so the plan is built by the
-same ``resolve`` and ``install`` as the target's at construction. ``paths`` must hold every
-path that one of them extends (``depends_on``), since ``DISABLE`` would turn such a path off;
-``select`` raises otherwise. ``dispatch`` is the selector module the code under test reads
-(some harnesses load the target's modules under a stand-in package).
-
-With ``monkeypatch`` the selector state of the process (the installed plan, its validation
-switch and the ``require`` mode) is restored when the test ends; without it the selection
-stays installed, as a process-level harness wants.
-
-A test that runs its case in a child interpreter passes the paths in ``FAST_PATHS_ENV``; the
-child installs them with ``select_from_env``.
-"""
-
-import os
-from typing import FrozenSet, Iterable, Mapping, Optional, Union
-
-Paths = Union[str, Iterable[str]]
-
-FAST_PATHS_ENV = "GLM_MOE_DSA_TEST_FAST_PATHS"
-
-
-def names(paths: Paths) -> FrozenSet[str]:
-    """The fast-path names of a comma-separated string or an iterable of names."""
-    if isinstance(paths, str):
-        return frozenset(name.strip() for name in paths.split(",") if name.strip())
-    return frozenset(paths)
-
-
-def plan_running(dispatch, paths: Paths, mode=None):
-    """The resolved plan that runs exactly ``paths``.
-
-    Raises:
-        ValueError: if a name is not a fast path, or if ``paths`` lacks a path one of them
-            extends.
-    """
-    wanted = names(paths)
-    unknown = sorted(wanted - set(dispatch.FAST_PATH_NAMES))
-    if unknown:
-        raise ValueError(f"{unknown} are not fast paths of the target")
-    bases = {base for path in wanted for base in dispatch.REQUIRED_FEATURES.get(path, ())}
-    missing = sorted(bases - wanted)
-    if missing:
-        raise ValueError(f"{sorted(wanted)} lacks {missing}, which they extend")
-    off = sorted(set(dispatch.FAST_PATH_NAMES) - wanted)
-    environ = {dispatch.DISABLE_ENV: ",".join(off)} if off else {}
-    plan = dispatch.resolve(mode if mode is not None else dispatch.ModelingV2Mode.AUTO, environ)
-    running = frozenset(name for name in dispatch.FAST_PATH_NAMES if plan.enabled(name))
-    assert running == wanted, (sorted(running), sorted(wanted))
-    return plan
-
-
-def select(
-    dispatch,
-    paths: Paths,
-    monkeypatch=None,
-    environ: Optional[Mapping[str, str]] = None,
-    mode=None,
-) -> FrozenSet[str]:
-    """Install the plan that runs exactly ``paths`` and return the installed paths.
-
-    ``environ`` is what the record line reads (the process environment by default), so a
-    variant knob set there selects the numerical variant of a selected path.
-    """
-    if monkeypatch is not None:
-        for name in ("_INSTALLED", "_VALIDATE", "_require"):
-            monkeypatch.setattr(dispatch, name, list(getattr(dispatch, name)))
-    dispatch.install(plan_running(dispatch, paths, mode), environ)
-    return dispatch.installed()
-
-
-def select_from_env(dispatch, environ: Optional[Mapping[str, str]] = None) -> FrozenSet[str]:
-    """Install, for this process, the plan that runs exactly the paths in ``FAST_PATHS_ENV``."""
-    environ = os.environ if environ is None else environ
-    return select(dispatch, environ.get(FAST_PATHS_ENV, ""), environ=environ)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py
deleted file mode 100644
index 62710ef7ae15..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/into_hop_jit.py
+++ /dev/null
@@ -1,81 +0,0 @@
-# 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.
-"""Load ``trtllm::glm_moe_dsa_post_attention_hop_from_partials`` for tests.
-
-A TensorRT-LLM build of this branch registers the op in ``th_common``. On an
-older image (development overlay) the op is compiled from this checkout's
-C++ sources with ``torch.utils.cpp_extension`` instead, unmodified and with
-the same op schema (``-DTLLM_GLM_MOE_DSA_ATTENTION_OUTPUT_INTO_HOP_JIT`` only names the library).
-"""
-
-import hashlib
-import os
-from pathlib import Path
-from typing import Callable
-
-import torch
-
-_REPO = Path(__file__).resolve().parents[9]
-_SOURCES = (
-    _REPO / "cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp",
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu",
-)
-_HEADERS = (
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h",
-    _REPO / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h",
-)
-_NAME = "into_hop_jit"
-_OP = "glm_moe_dsa_post_attention_hop_from_partials"
-
-
-def _digest() -> str:
-    h = hashlib.sha256()
-    for path in (*_SOURCES, *_HEADERS):
-        h.update(path.read_bytes())
-    return h.hexdigest()[:16]
-
-
-def ensure_op(build_root: str, is_builder: bool, barrier: Callable[[], None]) -> str:
-    """Make the op available on every rank; returns where it came from.
-
-    ``is_builder`` must be true on exactly one rank per shared build root;
-    ``barrier`` synchronizes the ranks after that rank has built it.
-    """
-    import tensorrt_llm  # noqa: F401  (loads th_common and the trtllm namespace)
-
-    if hasattr(torch.ops.trtllm, _OP):
-        barrier()
-        return "builtin"
-    build_dir = Path(build_root) / f"attention_output_into_hop-{_digest()}"
-    library = build_dir / f"{_NAME}.so"
-    if is_builder and not library.exists():
-        from torch.utils.cpp_extension import load
-
-        os.environ["TORCH_CUDA_ARCH_LIST"] = "10.3"
-        build_dir.mkdir(parents=True, exist_ok=True)
-        load(
-            name=_NAME,
-            sources=[str(s) for s in _SOURCES],
-            extra_include_paths=[str(_REPO / "cpp/include"), str(_REPO / "cpp")],
-            extra_cflags=["-O3", "-std=c++17"],
-            extra_cuda_cflags=["-O3", "-std=c++17", "--expt-relaxed-constexpr", "-Xptxas=-v"],
-            build_directory=str(build_dir),
-            is_python_module=False,
-            verbose=True,
-        )
-    barrier()
-    if not hasattr(torch.ops.trtllm, _OP):
-        torch.ops.load_library(str(library))
-    return str(library)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py
deleted file mode 100644
index dae079e8d365..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_attn_hop_static.py
+++ /dev/null
@@ -1,230 +0,0 @@
-# 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.
-"""Single-process checks of post_attention_hop (no collective, no GPU kernel).
-
-The eight-rank correctness checks need two 4-GPU nodes of one NVLink domain
-and run as collective tests.
-"""
-
-import re
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop import (  # noqa: E501
-    hop,
-)
-from tensorrt_llm.mapping import Mapping
-
-_HEADER = (
-    Path(__file__).resolve().parents[9] / "cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h"
-)
-
-
-def _header_constant(name: str) -> int:
-    match = re.search(rf"constexpr int {name} = (\d+);", _HEADER.read_text())
-    assert match, name
-    return int(match.group(1))
-
-
-def test_constants_match_kernel_header():
-    assert hop.HIDDEN_SIZE == _header_constant("kHiddenDim")
-    assert hop.NUM_RANKS == _header_constant("kNumRanks")
-    assert hop.NUM_EXPERTS == _header_constant("kNumExperts")
-    assert hop.MAX_ROWS == _header_constant("kMaxRows")
-    assert hop.NUM_LAMPORT_BUFFERS == _header_constant("kNumLamportBuffers")
-
-
-def test_entry_point_registered():
-    assert (
-        dispatch.get_impl("post_attention_hop", "make_post_attn_hop")
-        == hop.PostAttnHop.get_or_create
-    )
-    assert dispatch.get_impl("post_attention_hop", "engagement_counts") == hop.engagement_counts
-
-
-def test_router_layout_constants_match_kernel():
-    """required_sms mirrors the kernel's launch configuration."""
-    source = _HEADER.with_suffix(".cu").read_text()
-    assert f"constexpr int kSlices = {hop.CLUSTER_CTAS};" in source
-    for rows in range(1, hop.MAX_ROWS + 1):
-        groups = min(rows, hop.ROUTER_MAX_ROW_GROUPS)
-        assert (
-            f"case {rows}: return launchTyped<{rows}, "
-            f"{hop.ROUTER_EXPERT_GROUPS}, {groups}>(params);"
-        ) in source
-
-
-@pytest.mark.parametrize(
-    "rows, router, expected",
-    [
-        (1, False, 8),
-        (6, False, 48),
-        (8, False, 64),
-        (1, True, 32),
-        (2, True, 64),
-        (6, True, 96),
-        (8, True, 96),
-    ],
-)
-def test_required_sms(rows, router, expected):
-    assert hop.required_sms(rows, router) == expected
-
-
-def _mapping():
-    return Mapping(world_size=8, tp_size=8, gpus_per_node=4, rank=0)
-
-
-def test_native_allreduce_must_use_mnnvl():
-    from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce
-    from tensorrt_llm.functional import AllReduceStrategy
-
-    mapping = _mapping()
-    nccl = SimpleNamespace(strategy=AllReduceStrategy.NCCL, mnnvl_allreduce=None)
-    reason = hop.native_allreduce_ineligibility(nccl, mapping)
-    assert reason is not None and "NCCL" in reason
-    auto_single_node = SimpleNamespace(strategy=AllReduceStrategy.AUTO, mnnvl_allreduce=None)
-    assert hop.native_allreduce_ineligibility(auto_single_node, mapping) is not None
-    mnnvl = SimpleNamespace(strategy=AllReduceStrategy.AUTO, mnnvl_allreduce=object())
-    assert hop.native_allreduce_ineligibility(mnnvl, mapping) is None
-    # Call sites that do not pass their all-reduce: the TP group must have a
-    # native MNNVL workspace.
-    workspaces = MNNVLAllReduce.allreduce_mnnvl_workspaces
-    assert mapping not in workspaces
-    assert hop.native_allreduce_ineligibility(None, mapping) is not None
-    workspaces[mapping] = {}
-    try:
-        assert hop.native_allreduce_ineligibility(None, mapping) is None
-    finally:
-        del workspaces[mapping]
-
-
-def test_factory_declines_non_mnnvl_call_site_before_any_collective():
-    """A non-MNNVL call site returns None without creating the workspace."""
-    from tensorrt_llm.functional import AllReduceStrategy
-
-    mapping = _mapping()
-    nccl = SimpleNamespace(strategy=AllReduceStrategy.NCCL, mnnvl_allreduce=None)
-    assert (
-        hop.PostAttnHop.get_or_create(
-            mapping, hop.HIDDEN_SIZE, torch.bfloat16, native_allreduce=nccl
-        )
-        is None
-    )
-    assert mapping not in hop.PostAttnHop._instances
-
-
-def _enabled_hop(min_sm_count):
-    h = hop.PostAttnHop.__new__(hop.PostAttnHop)
-    h.enabled = True
-    h.min_sm_count = min_sm_count
-    return h
-
-
-def test_call_eligibility_and_fallback_counters():
-    h = _enabled_hop(152)
-    x6 = torch.zeros(6, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
-    x9 = torch.zeros(9, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
-    gamma = torch.zeros(hop.HIDDEN_SIZE, dtype=torch.bfloat16)
-    router = torch.zeros(hop.NUM_EXPERTS, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
-    hop.reset_engagement_counts()
-    assert h.can_run(x6, x6, gamma, router)
-    assert h.can_run(x6[:1], x6[:1], gamma)
-    assert not h.can_run(x9, x9, gamma, router)
-    assert not h.can_run(x6.half(), x6.half(), gamma)
-    assert not h.can_run(x6, None, gamma)
-    assert not h.can_run(x6, x6, gamma, router.float())
-    counts = hop.engagement_counts()
-    assert counts["native_fallback"] == 4
-    assert counts[f"native_fallback:rows outside [1, {hop.MAX_ROWS}]"] == 1
-    assert counts["native_fallback:dtype"] == 1
-    assert counts["native_fallback:no residual"] == 1
-    assert counts["native_fallback:router weight"] == 1
-    assert "engaged" not in counts
-    # A group whose smallest GPU has 64 SMs: the router layout needs 96 SMs
-    # from 3 rows on; the router-less layout fits up to 8 rows.
-    small = _enabled_hop(64)
-    assert small.call_ineligibility(x6[:2], x6[:2], gamma, router) is None
-    assert small.call_ineligibility(x6[:3], x6[:3], gamma, router) == "SM count"
-    x8 = torch.zeros(8, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
-    assert small.call_ineligibility(x8, x8, gamma) is None
-    hop.reset_engagement_counts()
-    assert hop.engagement_counts() == {}
-
-
-@pytest.mark.parametrize(
-    "kwargs, hidden, dtype, expected",
-    [
-        (dict(world_size=8, tp_size=8, gpus_per_node=4), 6144, torch.bfloat16, None),
-        (dict(world_size=4, tp_size=4, gpus_per_node=4), 6144, torch.bfloat16, "tp_size"),
-        (
-            dict(world_size=8, tp_size=8, gpus_per_node=4, enable_attention_dp=True),
-            6144,
-            torch.bfloat16,
-            "attention DP",
-        ),
-        (dict(world_size=8, tp_size=8, gpus_per_node=4), 7168, torch.bfloat16, "hidden size"),
-        (dict(world_size=8, tp_size=8, gpus_per_node=4), 6144, torch.float16, "dtype"),
-    ],
-)
-def test_static_ineligibility(kwargs, hidden, dtype, expected):
-    reason = hop.static_ineligibility(Mapping(rank=0, **kwargs), hidden, dtype)
-    if expected is None:
-        assert reason is None
-    else:
-        assert reason is not None and expected in reason
-
-
-def test_disabled_hop_never_runs():
-    disabled = hop.PostAttnHop.__new__(hop.PostAttnHop)
-    disabled.enabled = False
-    x = torch.zeros(6, hop.HIDDEN_SIZE, dtype=torch.bfloat16)
-    assert not disabled.can_run(x, x, torch.zeros(hop.HIDDEN_SIZE, dtype=torch.bfloat16))
-
-
-def test_hop_call_sites_keep_dense_layers_native():
-    """Only PRE_MOE_FUSION sites (MoE and MTP layers) call the hop.
-
-    The dense layers' native chain (plain all-reduce in o_proj, then a separate
-    add+RMSNorm) normalizes the FP32 sum, which the hop does not reproduce, so
-    the target must not route ``forward_mlp`` through the hop, and the hop's adoption
-    must not touch o_proj's all-reduce.
-    """
-    import inspect
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
-
-    source = Path(layers.__file__).read_text()
-    assert source.count("self._post_attn_hop(\n") == 2
-    for method in (layers.GlmDecoderLayer.forward_MoE, layers.GlmMtpLayer.forward):
-        assert "self._post_attn_hop(\n" in inspect.getsource(method)
-    assert "post_attn_hop" not in inspect.getsource(layers.GlmDecoderLayer.forward_mlp)
-    adoption = inspect.getsource(layers.GlmDecoderLayer._adopt_pre_moe)
-    assert "enable_allreduce" not in adoption
-    assert "o_proj" not in adoption
-
-
-def test_adoption_passes_the_native_allreduce():
-    """The hop serves a layer only if that layer's native AllReduce is MNNVL."""
-    import inspect
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
-
-    adoption = inspect.getsource(layers.GlmDecoderLayer._adopt_pre_moe)
-    assert '"make_post_attn_hop")' in adoption
-    assert "native_allreduce=self.allreduce" in adoption
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py
deleted file mode 100644
index 5a6d75c3428a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_attention_hop/test_into_hop_single_gpu.py
+++ /dev/null
@@ -1,276 +0,0 @@
-# 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.
-"""The into-hop entry on one GPU, over a simulated transport (attention_output_into_hop).
-
-The hop from partials (``trtllm::glm_moe_dsa_post_attention_hop_from_partials``) reads every rank's
-partial from the local (unicast) view of the post_attention_hop workspace only, so on one GPU
-a plain device buffer can stand in for it: the test writes all eight ranks'
-sanitized partials into slot ``[row][rank]`` of the current Lamport buffer
-(what the eight producers' multicast stores leave there), sets the protocol
-state, and runs the hop as rank 0 of 8. This needs no MNNVL and runs under
-every compute-sanitizer tool, the transport included. Checks:
-
-* ``residual_out`` equals the rank-order FP32 sum rounded to BF16 plus the
-  residual, exactly; ``normed`` and the router logits agree with references
-  of the same math within one BF16 ulp / FP32 summation-order tolerance;
-* every ``early_poll`` mode and ``router_load`` point gives bit-identical
-  outputs;
-* the protocol: the previous call's buffer (``clearBytes`` of it) is reset
-  to the sentinel, the current buffer is left as published, the state
-  advances to the next buffer with this call's byte count;
-* rows 1, 6, 8, router on/off; CUDA-graph capture and replay with fresh
-  partials.
-
-The partial values are synthetic (seeded); the real transport needs eight
-ranks.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-_HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, _HERE)
-import into_hop_jit  # noqa: E402
-
-H = 6144
-RANKS = 8
-MAX_ROWS = 8
-NUM_BUFFERS = 3
-BUFFER_BYTES = MAX_ROWS * RANKS * H * 2
-WORKSPACE_BYTES = BUFFER_BYTES * NUM_BUFFERS
-SENTINEL = -(2**31)  # 0x80000000 as int32
-EPS = 1e-5
-
-pytestmark = [
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-        reason="the attention_output_into_hop hop targets SM100/SM103",
-    ),
-    pytest.mark.threadleak(enabled=False),
-]
-
-_OP = {}
-
-
-def _op():
-    if "op" not in _OP:
-        into_hop_jit.ensure_op(
-            os.environ.get(
-                "GLM_MOE_DSA_ATTENTION_OUTPUT_INTO_HOP_JIT_ROOT",
-                "/tmp/attention_output_into_hop_jit",
-            ),
-            True,
-            lambda: None,
-        )
-        _OP["op"] = torch.ops.trtllm.glm_moe_dsa_post_attention_hop_from_partials
-    return _OP["op"]
-
-
-def _sanitized(x):
-    bits = x.contiguous().view(torch.int16).clone()
-    bits[bits == -32768] = 0  # bf16 -0.0 -> +0.0
-    return bits.view(torch.bfloat16)
-
-
-class Transport:
-    """A device buffer laid out like the post_attention_hop workspace, plus its state."""
-
-    def __init__(self):
-        self.uc = torch.full((WORKSPACE_BYTES // 4,), SENTINEL, dtype=torch.int32, device="cuda")
-        self.flags = torch.zeros(4, dtype=torch.int32, device="cuda")
-
-    def buffers(self):
-        return self.uc.view(torch.bfloat16).view(NUM_BUFFERS, MAX_ROWS, RANKS, H)
-
-    def publish(self, partials, cur, prev_rows):
-        """State as after the previous hop (prev_rows rows) and the eight
-        producers of this call; the previous buffer holds stale data."""
-        prev = (cur + NUM_BUFFERS - 1) % NUM_BUFFERS
-        self.uc.fill_(SENTINEL)
-        bufs = self.buffers()
-        bufs[prev, :prev_rows].fill_(1.5)  # stale words to be reset
-        for s, part in enumerate(partials):
-            bufs[cur, : part.shape[0], s] = _sanitized(part)
-        self.flags.copy_(torch.tensor([cur, prev_rows * RANKS * H * 2, 0, 0], dtype=torch.int32))
-
-
-def _inputs(rows, seed, router):
-    gen = torch.Generator(device="cuda").manual_seed(seed)
-    partials = [
-        (torch.randn(rows, H, generator=gen, device="cuda") * 0.05).to(torch.bfloat16)
-        for _ in range(RANKS)
-    ]
-    residual = torch.randn(rows, H, generator=gen, device="cuda").to(torch.bfloat16)
-    gamma = (1 + 0.1 * torch.randn(H, generator=gen, device="cuda")).to(torch.bfloat16)
-    weight = (
-        (torch.randn(256, H, generator=gen, device="cuda") * 0.02).to(torch.bfloat16).contiguous()
-        if router
-        else None
-    )
-    return partials, residual, gamma, weight
-
-
-def _run(tr, residual, gamma, weight, early_poll=False, router_load=0):
-    normed, res_out, logits = _op()(
-        residual, gamma, EPS, weight, tr.uc, tr.flags, 0, RANKS, early_poll, router_load
-    )
-    return normed, res_out, (logits if weight is not None else None)
-
-
-def _reference(partials, residual, gamma, weight, normed_k):
-    acc = torch.zeros_like(partials[0], dtype=torch.float32)
-    for part in partials:  # rank order, FP32
-        acc = acc + _sanitized(part).float()
-    res_out = acc.to(torch.bfloat16) + residual  # bf16 add, as __hadd
-    y = res_out.float()
-    sq = (res_out * res_out).float().sum(-1, keepdim=True)
-    normed = (y * torch.rsqrt(sq / H + EPS) * gamma.float()).to(torch.bfloat16)
-    logits = normed_k.float() @ weight.float().t() if weight is not None else None
-    return normed, res_out, logits
-
-
-@pytest.mark.parametrize("rows", [1, 6, 8])
-@pytest.mark.parametrize("router", [True, False])
-def test_hop_math_variants_and_protocol(rows, router):
-    partials, residual, gamma, weight = _inputs(rows, 100 + rows, router)
-    tr = Transport()
-    outs = []
-    for early_poll in (False, True):
-        for router_load in (0, 1, 2) if router else (0,):
-            cur, prev_rows = (early_poll + 2 * router_load) % 3, 8 - rows + 1
-            tr.publish(partials, cur, prev_rows)
-            out = _run(tr, residual, gamma, weight, early_poll, router_load)
-            torch.cuda.synchronize()
-            bufs = tr.buffers()
-            prev = (cur + NUM_BUFFERS - 1) % NUM_BUFFERS
-            # The previous call's rows are reset to the sentinel ...
-            assert bool((tr.uc.view(NUM_BUFFERS, -1)[prev] == SENTINEL).all())
-            # ... the published buffer is left as written ...
-            for s, part in enumerate(partials):
-                assert torch.equal(bufs[cur, :rows, s], _sanitized(part))
-            # ... and the state advances with this call's byte count.
-            want = [(cur + 1) % NUM_BUFFERS, rows * RANKS * H * 2, 0, 0]
-            assert tr.flags.tolist() == want, (tr.flags.tolist(), want)
-            outs.append(out)
-    for out in outs[1:]:
-        for a, b in zip(out, outs[0]):
-            assert (a is None and b is None) or torch.equal(a, b)
-    normed, res_out, logits = outs[0]
-    ref_n, ref_r, ref_l = _reference(partials, residual, gamma, weight, normed)
-    assert torch.equal(res_out, ref_r)
-    torch.testing.assert_close(normed.float(), ref_n.float(), rtol=2**-7, atol=2**-7)
-    if router:
-        torch.testing.assert_close(logits, ref_l, rtol=1e-3, atol=1e-3)
-
-
-@pytest.mark.parametrize("router_load", [0, 1, 2])
-def test_hop_graph_replay(router_load):
-    rows = 6
-    partials, residual, gamma, weight = _inputs(rows, 7, True)
-    tr = Transport()
-    tr.publish(partials, 0, rows)
-    _run(tr, residual, gamma, weight, router_load=router_load)
-    torch.cuda.synchronize()
-    tr.publish(partials, 1, rows)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        captured = _run(tr, residual, gamma, weight, router_load=router_load)
-    for rep in range(4):
-        fresh, res_new, _, _ = _inputs(rows, 300 + rep, True)
-        residual.copy_(res_new)
-        tr.publish(fresh, rep % NUM_BUFFERS, rows)
-        graph.replay()
-        torch.cuda.synchronize()
-        got = [t.clone() for t in captured]
-        tr.publish(fresh, rep % NUM_BUFFERS, rows)
-        want = _run(tr, residual, gamma, weight, router_load=router_load)
-        torch.cuda.synchronize()
-        for a, b in zip(got, want):
-            assert torch.equal(a, b)
-    del graph
-
-
-class _FakeHop:
-    """The parts of the post_attention_hop ``PostAttnHop`` that ``IntoHopTransport`` reads.
-
-    ``can_run`` (the post_attention_hop's own, fallback-counting check) must not be
-    called by the into-hop transport; ``call_ineligibility`` is the rank-uniform
-    per-call answer.
-    """
-
-    def __init__(self, uc, reason=None):
-        self.uc_buffer = uc
-        self.reason = reason
-        self.checked = []
-
-    def call_ineligibility(self, hidden_states, residual, norm_weight, router_weight=None):
-        self.checked.append(router_weight is not None)
-        return self.reason
-
-    def can_run(self, *args, **kwargs):
-        raise AssertionError("IntoHopTransport must not count post_attention_hop fallbacks")
-
-
-def _transport(tr, reason=None):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop.from_partials import (  # noqa: E501
-        IntoHopTransport,
-    )
-
-    into_hop = object.__new__(IntoHopTransport)
-    into_hop.hop = _FakeHop(tr.uc, reason)
-    into_hop.rank = 0
-    into_hop.flags = tr.flags
-    into_hop.early_poll = False
-    into_hop.router_load = 1
-    return into_hop
-
-
-def test_transport_eligibility_does_not_count_attn_hop_fallbacks():
-    tr = Transport()
-    residual = torch.zeros(6, H, dtype=torch.bfloat16, device="cuda")
-    gamma = torch.ones(H, dtype=torch.bfloat16, device="cuda")
-    assert _transport(tr).can_run(residual, gamma, None)
-    into_hop = _transport(tr, reason="SM count")
-    assert not into_hop.can_run(residual, gamma, gamma)
-    assert into_hop.hop.checked == [True]
-
-
-@pytest.mark.parametrize("router", [True, False])
-def test_transport_call_counts_and_matches_op(router):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_attention_hop import (  # noqa: E501
-        from_partials as from_partials_mod,
-    )
-
-    rows = 6
-    partials, residual, gamma, weight = _inputs(rows, 500 + router, router)
-    tr = Transport()
-    tr.publish(partials, 2, rows)
-    want = _run(tr, residual, gamma, weight, router_load=1)
-    torch.cuda.synchronize()
-    from_partials_mod.reset_into_hop_engagement_counts()
-    tr.publish(partials, 2, rows)
-    got = _transport(tr)(residual, gamma, EPS, weight)
-    torch.cuda.synchronize()
-    for a, b in zip(got, want):
-        assert (a is None and b is None) or torch.equal(a, b)
-    counts = dispatch.get_impl("post_attention_hop", "into_hop_engagement_counts")()
-    assert counts == (
-        {"into_hop_engaged": 1, "into_hop_engaged_router": 1} if router else {"into_hop_engaged": 1}
-    )
-    from_partials_mod.reset_into_hop_engagement_counts()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py
deleted file mode 100644
index b88e35074ae6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_mpi_utils.py
+++ /dev/null
@@ -1,261 +0,0 @@
-# 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.
-"""Shared helpers for the post_moe_hop collective tests and microbenchmarks.
-
-Every process of the TP group runs the same code under ``srun --mpi=pmix``;
-helpers here are collective where noted and must be called in the same order
-on every rank.
-"""
-
-import os
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import torch
-from mpi4py import MPI
-
-GLM_HIDDEN = 6144
-GLM_TOP_K = 8
-GLM_EPS = 1e-5
-
-
-@dataclass
-class MpiEnv:
-    comm: object
-    rank: int
-    world: int
-    gpus_per_node: int
-    mapping: object
-
-    @property
-    def multi_node(self) -> bool:
-        return self.mapping.is_multi_node()
-
-
-_ENV: Optional[MpiEnv] = None
-
-
-def mpi_env() -> MpiEnv:
-    """Per-process TP mapping over every MPI rank (collective-free)."""
-    global _ENV
-    if _ENV is None:
-        from tensorrt_llm.mapping import Mapping
-
-        comm = MPI.COMM_WORLD
-        rank, world = comm.Get_rank(), comm.Get_size()
-        gpus_per_node = int(
-            os.environ.get("GLM_MOE_DSA_POST_MOE_HOP_GPUS_PER_NODE", torch.cuda.device_count())
-        )
-        torch.cuda.set_device(rank % gpus_per_node)
-        mapping = Mapping(world_size=world, rank=rank, gpus_per_node=gpus_per_node, tp_size=world)
-        _ENV = MpiEnv(comm, rank, world, gpus_per_node, mapping)
-    return _ENV
-
-
-def all_ranks(ok: bool) -> bool:
-    """Collective AND of a per-rank verdict."""
-    return bool(mpi_env().comm.allreduce(int(bool(ok)), op=MPI.MIN))
-
-
-def gather_bf16(tensor: torch.Tensor) -> torch.Tensor:
-    """Collective all-gather of a bf16 tensor into a CPU ``[world, ...]`` tensor."""
-    host = tensor.detach().contiguous().view(torch.int16).cpu().numpy()
-    parts = mpi_env().comm.allgather(host)
-    return torch.stack([torch.from_numpy(p).view(torch.bfloat16) for p in parts])
-
-
-def make_raw_tuple(
-    rows: int,
-    hidden: int = GLM_HIDDEN,
-    top_k: int = GLM_TOP_K,
-    padded_rows: Optional[int] = None,
-    seed: int = 0,
-    scale_dtype: torch.dtype = torch.bfloat16,
-    invalid_fraction: float = 0.1,
-    with_shared: bool = True,
-    empty_rows: Tuple[int, ...] = (),
-):
-    """Synthetic rank-local raw MoE tuple plus replicated residual/gamma.
-
-    The FC2 rows, routing weights and permutation are rank-specific, the
-    residual and norm weight are identical on every rank (TP replication).
-    Rows listed in ``empty_rows`` have no valid expert slot.
-    """
-    env = mpi_env()
-    device = torch.device("cuda")
-    padded_rows = padded_rows or max(128, rows * top_k * 2)
-    local = torch.Generator(device=device).manual_seed(10_000 * seed + env.rank)
-    common = torch.Generator(device=device).manual_seed(20_000 * seed + 17)
-    fc2 = (torch.randn(padded_rows, hidden, device=device, generator=local) * 0.05).bfloat16()
-    perm = torch.randperm(padded_rows, device=device, generator=local)[: rows * top_k]
-    perm = perm.view(rows, top_k).to(torch.int32)
-    invalid = torch.rand(rows, top_k, device=device, generator=local) < invalid_fraction
-    perm = torch.where(invalid, torch.full_like(perm, -1), perm)
-    for row in empty_rows:
-        perm[row] = -1
-    scales = torch.rand(rows, top_k, device=device, generator=local) * 0.6
-    scales = scales.to(scale_dtype)
-    shared = None
-    if with_shared:
-        shared = (torch.randn(rows, hidden, device=device, generator=local) * 0.1).bfloat16()
-    residual = torch.randn(rows, hidden, device=device, generator=common).bfloat16()
-    gamma = (torch.rand(hidden, device=device, generator=common) + 0.5).bfloat16()
-    return fc2, scales, perm, shared, residual, gamma
-
-
-def reference_local_partial(fc2, scales, perm, shared, hidden: int = GLM_HIDDEN) -> torch.Tensor:
-    """Native multi-node local arithmetic: finalize (fp32, k order) -> bf16 -> + shared."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        native_finalize_reference,
-    )
-
-    routed = native_finalize_reference(fc2, scales, perm, hidden)
-    return routed if shared is None else shared + routed
-
-
-def published_partials(hop, rows: int) -> torch.Tensor:
-    """Read back ``[rows, world, H]`` partials of the hop's most recent call.
-
-    The buffer of the last completed call stays intact until the next call
-    clears it, so this reads exactly what every rank published.
-    """
-    ws = hop.workspace
-    torch.cuda.synchronize()
-    epoch = int(ws.state[0].item())
-    last = (epoch + 2) % 3
-    view = ws.uc.view(3, ws.max_rows, ws.world, ws.hidden_size)
-    return view[last, :rows].clone()
-
-
-def lamport_snapshot(hop) -> dict:
-    """Fill state of this rank's local Lamport copy (synchronizes the device).
-
-    At a quiescent point, when every rank's last call completed, the protocol
-    keeps exactly the rows of the last call's buffer filled in every slot and
-    every other element at the sentinel; ``ok`` reports that invariant.
-
-    Returns:
-        ``epoch``, ``prev_rows``, ``filled`` (non-sentinel element count per
-        ``[buffer][row][slot]``, int64 CPU), ``ok``, and ``partials_filled``
-        (non-sentinel row-sum partials per buffer).
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
-        MAX_CHUNKS,
-    )
-
-    ws = hop.workspace
-    torch.cuda.synchronize()
-    epoch, prev_rows = int(ws.state[0].item()), int(ws.state[1].item())
-    view = ws.uc.view(torch.int16).view(3, ws.max_rows, ws.world, ws.hidden_size)
-    filled = (view != -0x8000).sum(dim=-1).cpu()
-    expected = torch.zeros_like(filled)
-    expected[(epoch + 2) % 3, :prev_rows] = ws.hidden_size
-    parts = ws.partials.view(torch.int32).view(3, ws.max_rows * MAX_CHUNKS)
-    partials_filled = (parts != -0x80000000).sum(dim=-1).cpu().tolist()
-    return dict(
-        epoch=epoch,
-        prev_rows=prev_rows,
-        filled=filled,
-        ok=torch.equal(filled, expected),
-        partials_filled=partials_filled,
-    )
-
-
-def emulate_post_reduction(
-    partials: torch.Tensor, residual: torch.Tensor, gamma: torch.Tensor, eps: float
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Exact emulation of the native MNNVL one-shot arithmetic on given partials.
-
-    Args:
-        partials: bf16 ``[R, world, H]`` sanitized rank partials.
-
-    Returns:
-        ``(norm_out, residual_out)`` bf16 ``[R, H]``.
-    """
-    acc = torch.zeros(
-        partials.shape[0], partials.shape[2], dtype=torch.float32, device=partials.device
-    )
-    for r in range(partials.shape[1]):
-        acc = acc + partials[:, r].float()
-    reduced = acc.bfloat16()
-    new_residual = (reduced.float() + residual.float()).bfloat16()
-    x = new_residual.float()
-    square_sum = (x * x).bfloat16().float().sum(-1, keepdim=True)
-    inv_rms = torch.rsqrt(square_sum / x.shape[-1] + eps)
-    norm = ((x * inv_rms) * gamma.float()).bfloat16()
-    return norm, new_residual
-
-
-def sanitize_bf16(x: torch.Tensor) -> torch.Tensor:
-    bits = x.view(torch.int16)
-    return torch.where(bits == -32768, torch.zeros_like(bits), bits).view(torch.bfloat16)
-
-
-def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """Elementwise distance in bf16 units in the last place."""
-
-    def ordered(x):
-        v = x.contiguous().view(torch.int16).to(torch.int32)
-        return torch.where(v < 0, -(v + 32768), v)
-
-    return (ordered(a) - ordered(b)).abs()
-
-
-def compare_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict:
-    ulps = bf16_ulp_distance(candidate, reference)
-    diff = (candidate.float() - reference.float()).abs()
-    rel = diff / reference.float().abs().clamp_min(1e-6)
-    return {
-        "max_abs": float(diff.max()),
-        "max_rel": float(rel.max()),
-        "max_ulp": int(ulps.max()),
-        "frac_ulp_gt0": float((ulps > 0).float().mean()),
-        "frac_ulp_gt1": float((ulps > 1).float().mean()),
-        "rms_err": float(diff.pow(2).mean().sqrt()),
-    }
-
-
-_NATIVE_AR = None
-
-
-def native_allreduce():
-    """Cached native ``AllReduce(AUTO, bf16)`` of the TP group (MNNVL across hosts)."""
-    global _NATIVE_AR
-    if _NATIVE_AR is None:
-        from tensorrt_llm._torch.distributed import AllReduce
-        from tensorrt_llm.functional import AllReduceStrategy
-
-        _NATIVE_AR = AllReduce(
-            mapping=mpi_env().mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16
-        )
-    return _NATIVE_AR
-
-
-def native_fused_allreduce(
-    local: torch.Tensor, residual: torch.Tensor, gamma: torch.Tensor, eps: float
-):
-    """The native post-MoE hop: ``AllReduce`` with RESIDUAL_RMS_NORM."""
-    from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams
-
-    return native_allreduce()(
-        local,
-        all_reduce_params=AllReduceParams(
-            fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
-            residual=residual,
-            norm_weight=gamma,
-            eps=eps,
-            trigger_completion_at_end=False,
-        ),
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py
deleted file mode 100644
index a7c518c09f86..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/post_moe_hop_real_fixture.py
+++ /dev/null
@@ -1,211 +0,0 @@
-# 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.
-"""Real-weight native MoE fixture for one GLM-5.2 NVFP4 target layer.
-
-Builds, on every TP rank, the native TRT-LLM pieces that feed the post-MoE
-hop of the served GEN role (TP8/EP1, ``moe_backend="TRTLLM"``):
-
-* ``DeepseekV3Gate`` with the layer's real router weight and correction bias;
-* the native NVFP4 MoE (``create_moe`` -> TRTLLM-gen NVFP4 leaf) with the
-  layer's real 256 routed experts (TP-sliced by the module's own loader);
-* the rank's TP slice of the real BF16 shared expert (evaluated with torch);
-* the real RMSNorm weight of the next layer (``input_layernorm`` of
-  ``layer + 1``, or ``model.norm`` after the last layer).
-
-Weights are real; activations are synthetic (a unit-RMS random hidden state
-scaled by the layer's real ``post_attention_layernorm`` weight, and a random
-residual), identical on every rank as TP requires.
-"""
-
-import json
-import os
-from dataclasses import dataclass
-from typing import Dict, List
-
-import torch
-from utils.llm_data import llm_models_root
-
-CKPT = os.environ.get(
-    "GLM_MOE_DSA_POST_MOE_HOP_GLM52_CKPT",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-NUM_EXPERTS = 256
-TOP_K = 8
-HIDDEN = 6144
-MOE_INTER = 2048
-NUM_LAYERS = 78
-ROUTED_SCALING = 2.5
-EPS = 1e-5
-
-
-def _load(names: List[str]) -> Dict[str, torch.Tensor]:
-    from safetensors import safe_open
-
-    with open(os.path.join(CKPT, "model.safetensors.index.json")) as f:
-        weight_map = json.load(f)["weight_map"]
-    by_shard: Dict[str, List[str]] = {}
-    for name in names:
-        by_shard.setdefault(weight_map[name], []).append(name)
-    out = {}
-    for shard, shard_names in sorted(by_shard.items()):
-        with safe_open(os.path.join(CKPT, shard), framework="pt", device="cpu") as f:
-            for name in shard_names:
-                out[name] = f.get_tensor(name)
-    return out
-
-
-@dataclass
-class RealMoELayer:
-    layer: int
-    gate: torch.nn.Module
-    moe: torch.nn.Module
-    shared_gate: torch.Tensor
-    shared_up: torch.Tensor
-    shared_down: torch.Tensor
-    post_attn_gamma: torch.Tensor
-    next_gamma: torch.Tensor
-    expert_input_scale: float
-
-    def shared_output(self, x: torch.Tensor) -> torch.Tensor:
-        """Rank-local TP partial of the BF16 shared expert (no reduction)."""
-        gate = torch.nn.functional.linear(x, self.shared_gate)
-        up = torch.nn.functional.linear(x, self.shared_up)
-        return torch.nn.functional.linear(torch.nn.functional.silu(gate) * up, self.shared_down)
-
-    def hidden_states(self, rows: int, seed: int) -> torch.Tensor:
-        gen = torch.Generator(device="cuda").manual_seed(seed)
-        x = torch.randn(rows, HIDDEN, device="cuda", generator=gen)
-        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS)
-        return (x * self.post_attn_gamma.float()).bfloat16()
-
-    @staticmethod
-    def residual(rows: int, seed: int) -> torch.Tensor:
-        gen = torch.Generator(device="cuda").manual_seed(seed + 1)
-        return (torch.randn(rows, HIDDEN, device="cuda", generator=gen) * 4.0).bfloat16()
-
-    def run_moe(self, x: torch.Tensor, router_logits: torch.Tensor, do_finalize: bool):
-        return self.moe.forward(
-            x, router_logits, do_finalize=do_finalize, output_dtype=torch.bfloat16
-        )
-
-
-def build_real_moe_layer(mapping, layer: int, max_num_tokens: int = 64) -> RealMoELayer:
-    """Load layer ``layer`` of the checkpoint into the native modules on this rank."""
-    from transformers.configuration_utils import PretrainedConfig
-
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate
-    from tensorrt_llm._torch.moe.fused_moe import create_moe
-    from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
-
-    prefix = f"model.layers.{layer}"
-    next_norm = (
-        "model.norm.weight"
-        if layer + 1 >= NUM_LAYERS
-        else f"model.layers.{layer + 1}.input_layernorm.weight"
-    )
-    names = [
-        f"{prefix}.mlp.gate.weight",
-        f"{prefix}.mlp.gate.e_score_correction_bias",
-        f"{prefix}.post_attention_layernorm.weight",
-        next_norm,
-    ] + [f"{prefix}.mlp.shared_experts.{p}.weight" for p in ("gate_proj", "up_proj", "down_proj")]
-    suffixes = ("weight", "weight_scale", "weight_scale_2", "input_scale")
-    hf_to_native = {"gate_proj": "w1", "up_proj": "w3", "down_proj": "w2"}
-    for expert in range(NUM_EXPERTS):
-        for proj in hf_to_native:
-            names += [f"{prefix}.mlp.experts.{expert}.{proj}.{s}" for s in suffixes]
-    tensors = _load(names)
-
-    pretrained = PretrainedConfig()
-    pretrained.num_experts = NUM_EXPERTS
-    pretrained.n_routed_experts = NUM_EXPERTS
-    pretrained.hidden_size = HIDDEN
-    pretrained.intermediate_size = MOE_INTER
-    pretrained.moe_intermediate_size = MOE_INTER
-    pretrained.torch_dtype = torch.bfloat16
-    model_config = ModelConfig(
-        pretrained_config=pretrained,
-        mapping=mapping,
-        quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16),
-        moe_backend="TRTLLM",
-        max_num_tokens=max_num_tokens,
-    )
-
-    with torch.device("cuda"):
-        gate = DeepseekV3Gate(
-            HIDDEN,
-            NUM_EXPERTS,
-            top_k=TOP_K,
-            n_group=1,
-            topk_group=1,
-            routed_scaling_factor=ROUTED_SCALING,
-            dtype=torch.bfloat16,
-            fuse_routing_kernel=True,
-            apply_routing=False,
-            moe_backend="TRTLLM",
-        )
-        gate.load_weights(
-            [
-                {
-                    "weight": tensors[f"{prefix}.mlp.gate.weight"],
-                    "e_score_correction_bias": tensors[
-                        f"{prefix}.mlp.gate.e_score_correction_bias"
-                    ],
-                }
-            ]
-        )
-        moe = create_moe(
-            routing_method=gate.routing_method,
-            num_experts=NUM_EXPERTS,
-            hidden_size=HIDDEN,
-            intermediate_size=MOE_INTER,
-            dtype=torch.bfloat16,
-            reduce_results=False,
-            model_config=model_config,
-            layer_idx=layer,
-            allow_backend_degradation=False,
-        )
-        weights = {}
-        for expert in range(NUM_EXPERTS):
-            for proj, native in hf_to_native.items():
-                for s in suffixes:
-                    weights[f"{expert}.{native}.{s}"] = tensors[
-                        f"{prefix}.mlp.experts.{expert}.{proj}.{s}"
-                    ]
-        moe.load_weights([weights])
-        moe.post_load_weights()
-    moe.cuda()
-
-    tp, rank = mapping.tp_size, mapping.tp_rank
-    width = MOE_INTER // tp
-    rows = slice(rank * width, (rank + 1) * width)
-    shared_gate = tensors[f"{prefix}.mlp.shared_experts.gate_proj.weight"][rows].cuda()
-    shared_up = tensors[f"{prefix}.mlp.shared_experts.up_proj.weight"][rows].cuda()
-    shared_down = (
-        tensors[f"{prefix}.mlp.shared_experts.down_proj.weight"][:, rows].contiguous().cuda()
-    )
-    input_scale = float(tensors[f"{prefix}.mlp.experts.0.gate_proj.input_scale"])
-    return RealMoELayer(
-        layer,
-        gate,
-        moe,
-        shared_gate,
-        shared_up,
-        shared_down,
-        tensors[f"{prefix}.post_attention_layernorm.weight"].cuda(),
-        tensors[next_norm].cuda(),
-        input_scale,
-    )
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
deleted file mode 100644
index 69641ad86883..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
+++ /dev/null
@@ -1,796 +0,0 @@
-# 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.
-"""Collective correctness tests of the post_moe_hop (the post_moe_hop contract).
-
-Run one pytest process per GPU of the TP group, e.g. two 4-GPU hosts:
-
-    srun --mpi=pmix -N2 --ntasks-per-node=4 ... \
-        python3 -m pytest --import-mode=importlib -p no:cacheprovider -x -q \
-        tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_post_moe_mpi.py
-
-Every test is collective: all ranks run the same tests in the same order and
-each verdict is AND-reduced across ranks before it is asserted, so a failure
-on one rank fails the test everywhere instead of deadlocking the next test.
-Inputs are synthetic raw tuples (rank-specific FC2 rows/weights/permutation,
-replicated residual and norm weight); the real-weight native-runner check is
-in ``test_post_moe_hop_real_weights_mpi.py``.
-
-Oracles:
-* the native finalize arithmetic in torch (``native_finalize_reference``) for
-  the rank-local partial, which must equal what the kernel published;
-* an exact emulation of the native MNNVL one-shot arithmetic applied to the
-  partials every rank published (residual output must be bit-exact);
-* the native TRT-LLM ``AllReduce`` (MNNVL across hosts) with RESIDUAL_RMS_NORM
-  on the same partial (residual output bit-exact, normalized output compared
-  in bf16 ulps and with the existing MNNVL test tolerance rtol=0.05/atol=0.15).
-"""
-
-import os
-import sys
-import time
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import post_moe_hop_mpi_utils as U  # noqa: E402
-
-ROW_CASES = (1, 2, 6, 16)
-NORM_MAX_ULP = 2
-
-
-def _log(env, message):
-    if env.rank == 0:
-        print(f"[post_moe_hop-test] {message}", flush=True)
-
-
-@pytest.fixture(scope="module")
-def env():
-    e = U.mpi_env()
-    if e.world < 2:
-        pytest.skip("needs a TP group of at least two ranks under srun --mpi=pmix")
-    _log(e, f"world={e.world} gpus_per_node={e.gpus_per_node} multi_node={e.multi_node}")
-    return e
-
-
-@pytest.fixture(scope="module", params=("triton", "cluster"))
-def hop(env, request):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    U.native_allreduce()
-    hop = PostMoeHop(
-        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, impl=request.param
-    )
-    assert U.all_ranks(hop.impl == request.param), f"{request.param} kernel unavailable"
-    _log(env, f"--- hop impl={hop.impl}")
-    return hop
-
-
-def _check_call(env, hop, fc2, scales, perm, shared, residual, gamma, out, res_out, published=None):
-    """Return (ok, stats) for one completed call against all three oracles."""
-    ref_local = U.sanitize_bf16(U.reference_local_partial(fc2, scales, perm, shared))
-    expected_slots = U.gather_bf16(ref_local).permute(1, 0, 2).cuda()  # [R, world, H]
-    stats = {}
-    ok = True
-    if published is not None:
-        slots_equal = torch.equal(published.view(torch.int16), expected_slots.view(torch.int16))
-        stats["slots_exact"] = slots_equal
-        ok &= slots_equal
-    emu_norm, emu_res = U.emulate_post_reduction(expected_slots, residual, gamma, U.GLM_EPS)
-    res_exact = torch.equal(res_out.view(torch.int16), emu_res.view(torch.int16))
-    norm_stats = U.compare_stats(out, emu_norm)
-    stats.update(res_exact=res_exact, norm_vs_emulation=norm_stats)
-    ok &= res_exact and norm_stats["max_ulp"] <= NORM_MAX_ULP
-    return ok, stats
-
-
-class _Backend:
-    def __init__(self, declares: bool):
-        self._declares = declares
-
-    def supports_unfinalized_output(self) -> bool:
-        return self._declares
-
-
-class _Routing:
-    def __init__(self, top_k: int):
-        self.experts_per_token = top_k
-
-
-class _Producer:
-    """Stand-in for a ``ConfigurableMoE``: the attributes ``producer_capability`` reads."""
-
-    def __init__(
-        self,
-        declares=True,
-        hidden=U.GLM_HIDDEN,
-        top_k=U.GLM_TOP_K,
-        dtype=torch.bfloat16,
-        reduce_results=False,
-    ):
-        self.backend = _Backend(declares)
-        self.routing_method = _Routing(top_k)
-        self.hidden_size = hidden
-        self.dtype = dtype
-        self.reduce_results = reduce_results
-
-
-def _reset_all(env, hop):
-    assert U.all_ranks(hop.reset())
-    assert U.all_ranks(hop.error_bits() == 0)
-
-
-def test_capability_and_dispatch(env, hop):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-        check_health,
-    )
-    from tensorrt_llm.mapping import Mapping
-
-    assert dispatch.get_impl("post_moe_hop", "post_moe_hop") is PostMoeHop
-    assert dispatch.get_impl("post_moe_hop", "health_check") is check_health
-    assert U.all_ranks(hop.engaged), f"hop not engaged: {hop.reason}"
-    assert U.all_ranks(
-        not hop.can_run(0)
-        and hop.can_run(6)
-        and hop.can_run(hop.max_rows)
-        and not hop.can_run(hop.max_rows + 1)
-    )
-
-    # One vetoing rank must disengage every rank, without a hang.
-    veto = "test veto" if env.rank == env.world - 1 else None
-    vetoed = PostMoeHop(
-        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, local_veto=veto
-    )
-    assert U.all_ranks(not vetoed.engaged and not vetoed.can_run(6))
-
-    # A TP1 mapping never negotiates (no collective) and stays native.
-    single = PostMoeHop(
-        Mapping(world_size=1, rank=0, tp_size=1),
-        U.GLM_HIDDEN,
-        U.GLM_TOP_K,
-        require_multi_node=False,
-    )
-    assert not single.engaged
-
-    # The default domain is multi-node only; the default kernel is the cluster kernel.
-    default = PostMoeHop(env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K)
-    assert U.all_ranks(default.engaged == env.multi_node)
-    assert U.all_ranks(not default.engaged or default.impl == "cluster")
-
-    # A cluster width that does not split H into whole warps falls back to Triton.
-    fallback = PostMoeHop(
-        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, impl="cluster", cluster=5
-    )
-    assert U.all_ranks(fallback.engaged and fallback.impl == "triton")
-    _log(env, f"capability ok; multi_node={env.multi_node}")
-
-
-@pytest.mark.parametrize(
-    "case", ("declares", "no_declaration", "reduces", "top_k_one_rank", "hidden", "dtype")
-)
-def test_producer_contract_decided_before_runner(env, hop, case):
-    """The producer's declared raw-tuple contract is part of the collective decision."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    producer = {
-        "declares": _Producer(),
-        "no_declaration": _Producer(declares=False),
-        "reduces": _Producer(reduce_results=True),
-        # A mismatch on one rank only must disengage every rank.
-        "top_k_one_rank": _Producer(top_k=6 if env.rank == 1 else U.GLM_TOP_K),
-        "hidden": _Producer(hidden=U.GLM_HIDDEN // 2),
-        "dtype": _Producer(dtype=torch.float16),
-    }[case]
-    candidate = PostMoeHop(
-        env.mapping,
-        U.GLM_HIDDEN,
-        U.GLM_TOP_K,
-        require_multi_node=False,
-        impl=hop.impl,
-        raw_producer=producer,
-    )
-    expected = case == "declares"
-    _log(
-        env,
-        f"[{hop.impl}] producer {case}: engaged={candidate.engaged} reason={candidate.reason!r}",
-    )
-    assert U.all_ranks(candidate.engaged == expected and candidate.can_run(6) == expected)
-
-
-@pytest.mark.parametrize("scale_dtype", (torch.bfloat16, torch.float32), ids=("w_bf16", "w_fp32"))
-@pytest.mark.parametrize("with_shared", (True, False), ids=("shared", "no_shared"))
-@pytest.mark.parametrize("rows", ROW_CASES)
-def test_matches_oracles(env, hop, rows, with_shared, scale_dtype):
-    seed = 100 + rows * 4 + 2 * int(with_shared) + int(scale_dtype == torch.float32)
-    fc2, scales, perm, shared, residual, gamma = U.make_raw_tuple(
-        rows, seed=seed, scale_dtype=scale_dtype, with_shared=with_shared
-    )
-    out, res_out = hop(fc2, scales, perm, shared, residual, gamma, U.GLM_EPS)
-    published = U.published_partials(hop, rows)
-    ok, stats = _check_call(
-        env, hop, fc2, scales, perm, shared, residual, gamma, out, res_out, published
-    )
-
-    # Native TRT-LLM hop on the identical rank partial.
-    native_norm, native_res = U.native_fused_allreduce(
-        published[:, env.rank].contiguous(), residual, gamma, U.GLM_EPS
-    )
-    torch.cuda.synchronize()
-    native_res_exact = torch.equal(native_res.view(torch.int16), res_out.view(torch.int16))
-    native_stats = U.compare_stats(out, native_norm)
-    try:
-        torch.testing.assert_close(out, native_norm, rtol=0.05, atol=0.15)
-        torch.testing.assert_close(res_out, native_res, rtol=0.05, atol=0.15)
-        within_tolerance = True
-    except AssertionError:
-        within_tolerance = False
-    ok &= native_res_exact and within_tolerance and native_stats["max_ulp"] <= NORM_MAX_ULP
-    ok &= hop.error_bits() == 0
-    _log(
-        env,
-        f"[{hop.impl}] rows={rows} shared={with_shared} w={scale_dtype}: {stats} "
-        f"native_res_exact={native_res_exact} norm_vs_native={native_stats}",
-    )
-    assert U.all_ranks(ok), f"rank {env.rank}: {stats} native={native_stats}"
-
-
-@pytest.mark.parametrize("with_shared", (True, False), ids=("shared", "no_shared"))
-def test_empty_and_padded_rows(env, hop, with_shared):
-    rows = 6
-    fc2, scales, perm, shared, residual, gamma = U.make_raw_tuple(
-        rows, seed=300 + int(with_shared), with_shared=with_shared, empty_rows=(1, 4)
-    )
-    residual[2].zero_()  # a padded token row
-    if shared is not None:
-        shared[2].zero_()
-    out, res_out = hop(fc2, scales, perm, shared, residual, gamma, U.GLM_EPS)
-    published = U.published_partials(hop, rows)
-    ok, stats = _check_call(
-        env, hop, fc2, scales, perm, shared, residual, gamma, out, res_out, published
-    )
-    if not with_shared:
-        ok &= bool((published[1].view(torch.int16) == 0).all())
-        ok &= bool((published[4].view(torch.int16) == 0).all())
-    ok &= bool(torch.isfinite(out.float()).all())
-    _log(env, f"empty/padded rows shared={with_shared}: {stats}")
-    assert U.all_ranks(ok), stats
-
-
-@pytest.mark.parametrize(
-    "block,warps,pdl", ((512, 2, False), (2048, 8, True), (8192, 16, True), (8192, 32, False))
-)
-def test_launch_configs(env, hop, block, warps, pdl):
-    """Chunked and single-CTA-per-row grids, with and without PDL, share one workspace."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    variant = PostMoeHop(
-        env.mapping,
-        U.GLM_HIDDEN,
-        U.GLM_TOP_K,
-        require_multi_node=False,
-        block=block,
-        num_warps=warps,
-        use_pdl=pdl,
-        impl="triton",
-    )
-    ok = variant.engaged and variant.workspace is hop.workspace
-    for rows in (1, 6):
-        tensors = U.make_raw_tuple(rows, seed=350 + rows + block)
-        out, res_out = variant(*tensors, U.GLM_EPS)
-        call_ok, stats = _check_call(
-            env, variant, *tensors, out, res_out, U.published_partials(variant, rows)
-        )
-        ok &= call_ok
-    ok &= variant.error_bits() == 0
-    _log(env, f"launch config block={block} warps={warps} pdl={pdl}: ok={ok}")
-    assert U.all_ranks(ok), stats
-
-
-def test_narrow_launch_config_disengages(env, hop):
-    """A Triton program whose column vector does not span every warp is refused."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    variant = PostMoeHop(
-        env.mapping,
-        U.GLM_HIDDEN,
-        U.GLM_TOP_K,
-        require_multi_node=False,
-        block=1024,
-        num_warps=8,
-        impl="triton",
-    )
-    _log(env, f"block=1024 num_warps=8: engaged={variant.engaged} reason={variant.reason!r}")
-    assert U.all_ranks(not variant.engaged and "256 * num_warps" in variant.reason)
-
-
-def _run_back_to_back(env, hops, sequence, seed):
-    """Launch every call without a host sync, then check each one."""
-    calls_before = hops[0].workspace.calls()
-    inputs, outputs = [], []
-    for i, rows in enumerate(sequence):
-        tensors = U.make_raw_tuple(
-            rows,
-            seed=seed + i,
-            with_shared=(i % 4 != 3),
-            scale_dtype=torch.float32 if i % 5 == 0 else torch.bfloat16,
-        )
-        inputs.append(tensors)
-        outputs.append(hops[i % len(hops)](*tensors, U.GLM_EPS))
-    torch.cuda.synchronize()
-    ok = hops[0].workspace.calls() - calls_before == len(sequence)
-    failed = []
-    for i, (tensors, (out, res_out)) in enumerate(zip(inputs, outputs)):
-        call_ok, stats = _check_call(env, hops[0], *tensors, out, res_out)
-        ok &= call_ok
-        if not call_ok:
-            failed.append(
-                (i, sequence[i], stats["res_exact"], stats["norm_vs_emulation"]["max_ulp"])
-            )
-    ok &= hops[0].error_bits() == 0
-    return ok, failed
-
-
-def test_rapid_reuse_varying_rows(env, hop):
-    sequence = [6, 1, 16, 3, 6, 6, 2, 16, 5, 6] * 3
-    ok, failed = _run_back_to_back(env, [hop], sequence, seed=400)
-    _log(env, f"[{hop.impl}] rapid reuse over {len(sequence)} calls ok={ok} failed={failed[:6]}")
-    assert U.all_ranks(ok), failed
-
-
-def test_mixed_launch_configs_back_to_back(env, hop):
-    """The fixture's kernel interleaved with chunked and single-program Triton hops."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    variants = [hop] + [
-        PostMoeHop(
-            env.mapping,
-            U.GLM_HIDDEN,
-            U.GLM_TOP_K,
-            require_multi_node=False,
-            block=block,
-            num_warps=warps,
-            use_pdl=pdl,
-            impl="triton",
-        )
-        for block, warps, pdl in ((8192, 16, True), (512, 2, False))
-    ]
-    sequence = [6, 16, 1, 6, 3, 16, 6, 2, 6, 5, 16, 6] * 2
-    ok, failed = _run_back_to_back(env, variants, sequence, seed=450)
-    _log(env, f"mixed configs over {len(sequence)} calls ok={ok} failed={failed[:6]}")
-    assert U.all_ranks(ok), failed
-
-
-def test_cuda_graph_replay_with_native_interleave(env, hop):
-    num_layers, rows = 4, 6
-    static = [U.make_raw_tuple(rows, seed=500 + layer) for layer in range(num_layers)]
-    native_in = torch.randn(rows, U.GLM_HIDDEN, device="cuda").bfloat16()
-    native_residual, native_gamma = static[0][4].clone(), static[0][5]
-
-    def body():
-        results = []
-        native_out = None
-        for layer, tensors in enumerate(static):
-            results.append(hop(*tensors, U.GLM_EPS))
-            if layer == 1:
-                native_out = U.native_fused_allreduce(
-                    native_in, native_residual, native_gamma, U.GLM_EPS
-                )
-        return results, native_out
-
-    body()  # eager warmup: Triton compile and native workspaces
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        graph_results, graph_native = body()
-    torch.cuda.synchronize()
-    env.comm.barrier()
-
-    ok = True
-    calls_before = hop.workspace.calls()
-    for replay in range(20):
-        for layer, tensors in enumerate(static):
-            fresh = U.make_raw_tuple(
-                rows, seed=600 + 10 * replay + layer, padded_rows=tensors[0].shape[0]
-            )
-            for dst, src in zip(tensors, fresh):
-                if dst is not None:
-                    dst.copy_(src)
-        native_in.copy_(torch.randn_like(native_in.float()).bfloat16())
-        graph.replay()
-        torch.cuda.synchronize()
-        for tensors, (out, res_out) in zip(static, graph_results):
-            call_ok, stats = _check_call(env, hop, *tensors, out, res_out)
-            ok &= call_ok
-        gathered = U.gather_bf16(native_in).permute(1, 0, 2).cuda()
-        _, native_res_ref = U.emulate_post_reduction(
-            gathered, native_residual, native_gamma, U.GLM_EPS
-        )
-        ok &= torch.equal(graph_native[1].view(torch.int16), native_res_ref.view(torch.int16))
-    ok &= hop.workspace.calls() - calls_before == 20 * num_layers
-    ok &= hop.error_bits() == 0
-    _log(env, f"[{hop.impl}] graph replay x20 with native interleave ok={ok}")
-    assert U.all_ranks(ok)
-
-
-def test_delayed_rank_stress(env, hop):
-    rows = 6
-    ok = True
-    for it in range(12):
-        victim = it % env.world
-        tensors = U.make_raw_tuple(rows, seed=700 + it)
-        env.comm.barrier()
-        if env.rank == victim:
-            if it % 3 == 0:
-                time.sleep(0.02)  # host-side launch delay
-            else:
-                torch.cuda._sleep(int(2e6 * (1 + it % 4)))  # ~1-4 ms of device delay
-        out, res_out = hop(*tensors, U.GLM_EPS)
-        torch.cuda.synchronize()
-        call_ok, stats = _check_call(env, hop, *tensors, out, res_out)
-        ok &= call_ok
-    ok &= hop.error_bits() == 0
-    _log(env, f"[{hop.impl}] delayed-rank stress ok={ok}")
-    assert U.all_ranks(ok)
-
-
-def _missing_rank_trial(env, hop, impatient, missing, seed):
-    """One call that rank ``missing`` skips; returns this rank's verdict."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-
-    rows, timeout_s = 6, impatient.timeout_ns / 1e9
-    tensors = U.make_raw_tuple(rows, seed=seed)
-    before = U.lamport_snapshot(hop)
-    _log(
-        env,
-        f"[{hop.impl}] Lamport invariant before rank {missing} skips, on every rank "
-        f"(rank, epoch, prev_rows, ok): "
-        f"{env.comm.allgather((env.rank, before['epoch'], before['prev_rows'], before['ok']))}",
-    )
-    assert U.all_ranks(before["ok"]), before
-    cur = before["epoch"] % 3
-    env.comm.barrier()
-    start = time.monotonic()
-    out = res_out = None
-    if env.rank != missing:
-        out, res_out = impatient(*tensors, U.GLM_EPS)
-    torch.cuda.synchronize()
-    elapsed = time.monotonic() - start
-    counters = hop.workspace.counters()
-    bits, polls = counters["error_bits"], counters["timeout_polls"]
-    nan_chunks = slot_filled = None
-    if env.rank == missing:
-        reported = bits == 0 and polls == 0
-    else:
-        # Every program of the call must time out (the skipped rank's slot
-        # stays empty) and therefore write NaN for all of its columns.
-        slot_filled = U.lamport_snapshot(hop)["filled"][cur, :rows, missing].tolist()
-        nan_chunks = [
-            torch.isnan(t.float()).view(rows, U.GLM_HIDDEN // 1024, 1024).all(-1).sum(-1).tolist()
-            for t in (out, res_out)
-        ]
-        reported = (
-            bool(bits & 1)
-            and polls > 0
-            and timeout_s <= elapsed < 30.0
-            and bool(torch.isnan(out.float()).all())
-            and bool(torch.isnan(res_out.float()).all())
-        )
-    stats = env.comm.allgather((env.rank, round(elapsed, 3), bits, polls, reported))
-    _log(
-        env,
-        f"[{hop.impl}] rank {missing} skipped; per rank (rank, elapsed_s, error_bits, "
-        f"timeout_polls, ok): {stats}",
-    )
-    if not U.all_ranks(reported):
-        detail = env.comm.allgather((env.rank, nan_chunks, slot_filled))
-        _log(
-            env,
-            f"[{hop.impl}] per rank (rank, all-NaN 1024-column chunks per row of "
-            f"[out, residual_out], filled elements per row of the skipped rank's slot): {detail}",
-        )
-        return False
-
-    try:
-        hop.raise_if_error()
-        raised = False
-    except RuntimeError:
-        raised = True
-    try:
-        check_health(f"missing-rank{missing}-local")
-        local_raised = False
-    except PostMoeHopError:
-        local_raised = True
-    try:
-        check_health(f"missing-rank{missing}-collective", collective=True)
-        group_raised = False
-    except PostMoeHopError:
-        group_raised = True
-    expected = env.rank != missing
-    return raised == expected and local_raised == expected and group_raised
-
-
-def test_bounded_error_and_collective_reset(env, hop):
-    """A rank that never joins makes its peers time out, poison and report, not hang."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-        check_health,
-    )
-
-    impatient = PostMoeHop(
-        env.mapping,
-        U.GLM_HIDDEN,
-        U.GLM_TOP_K,
-        require_multi_node=False,
-        timeout_s=0.2,
-        impl=hop.impl,
-    )
-    assert U.all_ranks(impatient.workspace is hop.workspace)
-    for trial, missing in enumerate(sorted({0, env.world // 2 + 1, env.world - 1})):
-        ok = _missing_rank_trial(env, hop, impatient, missing, seed=800 + trial)
-        assert U.all_ranks(ok), f"trial {trial}: rank {missing} skipped"
-        _reset_all(env, hop)
-
-    rows = 6
-    tensors = U.make_raw_tuple(rows, seed=810)
-    out, res_out = hop(*tensors, U.GLM_EPS)
-    torch.cuda.synchronize()
-    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, U.published_partials(hop, rows))
-    assert U.all_ranks(call_ok and hop.error_bits() == 0), stats
-    health = check_health("after-reset", collective=True)
-    assert U.all_ranks(not health.group_error and health.workspaces == 1)
-
-
-def test_error_poisons_every_rank_until_reset(env, hop):
-    """After an error recorded on one rank, every rank's next outputs are NaN."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop.kernels import (  # noqa: E501
-        STATE_EPOCH,
-        STATE_ERROR_WORDS,
-    )
-
-    victim = env.world - 2
-    torch.cuda.synchronize()
-    if env.rank == victim:
-        # As if the previous call of this rank had timed out.
-        epoch = int(hop.workspace.state[STATE_EPOCH].item())
-        hop.workspace.state[STATE_ERROR_WORDS[(epoch + 2) % 3]] = 1
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    poisoned = True
-    # Three calls cover every Lamport buffer, so the bits must be carried from
-    # each buffer's error word into the next one.
-    for call in range(3):
-        tensors = U.make_raw_tuple(6, seed=850 + call)
-        out, res_out = hop(*tensors, U.GLM_EPS)
-        torch.cuda.synchronize()
-        poisoned &= bool(torch.isnan(out.float()).all()) and bool(
-            torch.isnan(res_out.float()).all()
-        )
-    words = hop.workspace.state.cpu().tolist()
-    carried = all(words[w] == 1 for w in STATE_ERROR_WORDS)
-    _log(
-        env,
-        f"[{hop.impl}] victim rank {victim} poisoned; three calls all NaN on this rank: "
-        f"{poisoned}; per-rank error words: "
-        f"{env.comm.allgather([words[w] for w in STATE_ERROR_WORDS])}",
-    )
-    assert U.all_ranks(
-        poisoned
-        and (carried if env.rank == victim else not any(words[w] for w in STATE_ERROR_WORDS))
-    )
-    with pytest.raises(PostMoeHopError):
-        check_health("poisoned", collective=True)
-    _reset_all(env, hop)
-    tensors = U.make_raw_tuple(6, seed=860)
-    out, res_out = hop(*tensors, U.GLM_EPS)
-    torch.cuda.synchronize()
-    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, U.published_partials(hop, 6))
-    assert U.all_ranks(call_ok), stats
-
-
-def test_permutation_out_of_range_is_skipped_and_reported(env, hop):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopError,
-        check_health,
-    )
-
-    culprit = 2 % env.world
-    rows = 6
-    tensors = U.make_raw_tuple(rows, seed=870)
-    if env.rank == culprit:
-        tensors[2][1, 3] = tensors[0].shape[0]
-        tensors[2][4, 0] = tensors[0].shape[0] + 1000
-    out, res_out = hop(*tensors, U.GLM_EPS)
-    published = U.published_partials(hop, rows)
-    # The oracle skips entries outside [0, N_padded) like the kernels.
-    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, published)
-    bits = hop.error_bits()
-    call_ok &= bits == (4 if env.rank == culprit else 0)
-    _log(env, f"[{hop.impl}] out-of-range permutation on rank {culprit}: {stats} bits={bits}")
-    assert U.all_ranks(call_ok), (stats, bits)
-    with pytest.raises(PostMoeHopError):
-        check_health("perm-range", collective=True)
-    _reset_all(env, hop)
-
-
-def _misaligned(t: torch.Tensor, offset_elems: int) -> torch.Tensor:
-    """Same values in a view whose base is ``offset_elems`` past a fresh allocation."""
-    storage = torch.empty(t.numel() + 16, dtype=t.dtype, device=t.device)
-    view = storage[offset_elems : offset_elems + t.numel()].view(t.shape)
-    view.copy_(t)
-    return view
-
-
-def _transposed(t: torch.Tensor) -> torch.Tensor:
-    """Same values, non-contiguous (column-major) layout."""
-    view = torch.empty(tuple(reversed(t.shape)), dtype=t.dtype, device=t.device).t()
-    view.copy_(t)
-    return view
-
-
-def _padded_rows(t: torch.Tensor, extra_cols: int) -> torch.Tensor:
-    """Same values in a ``[:, :H]`` view of a wider row-major allocation."""
-    storage = torch.zeros((t.shape[0], t.shape[1] + extra_cols), dtype=t.dtype, device=t.device)
-    view = storage[:, : t.shape[1]]
-    view.copy_(t)
-    return view
-
-
-def _needs_copy(t: torch.Tensor, row_major_2d: bool = False) -> bool:
-    """The documented 16-byte vector-access contract of the post_moe_hop kernels."""
-    if t.data_ptr() % 16 or t.stride(-1) != 1:
-        return True
-    if row_major_2d:
-        return (t.stride(0) * t.element_size()) % 16 != 0
-    return not t.is_contiguous()
-
-
-def _relayout(env, tensors):
-    """Rank-dependent layouts; returns (tensors, expected repairs on this rank)."""
-    fc2, scales, perm, shared, residual, gamma = tensors
-    kind = env.rank % 4
-    if kind == 0:
-        fc2 = _misaligned(fc2, 1)  # base % 16 == 2
-        perm = _misaligned(perm, 1)  # base % 16 == 4
-        gamma = _misaligned(gamma, 2)  # base % 16 == 4
-    elif kind == 1:
-        fc2 = _padded_rows(fc2, 4)  # row pitch % 16 == 8
-        scales = _transposed(scales)  # not contiguous (for R > 1)
-        residual = _misaligned(residual, 4)  # base % 16 == 8
-    elif kind == 2:
-        fc2 = _padded_rows(fc2, 8)  # aligned row pitch: read in place
-        if shared is not None:
-            shared = _transposed(shared)  # inner stride != 1 (for R > 1)
-    repairs = int(_needs_copy(fc2, row_major_2d=True)) + sum(
-        int(_needs_copy(t)) for t in (scales, perm, shared, residual, gamma) if t is not None
-    )
-    return (fc2, scales, perm, shared, residual, gamma), repairs
-
-
-def test_unaligned_layouts_are_repaired_locally(env, hop):
-    """Per-rank layout repairs keep the collective and the arithmetic unchanged."""
-    ok = True
-    for i, rows in enumerate((6, 1, 16)):
-        reference = U.make_raw_tuple(rows, seed=880 + i)
-        tensors, expected = _relayout(env, reference)
-        before = hop.workspace.layout_repairs
-        out, res_out = hop(*tensors, U.GLM_EPS)
-        call_ok, stats = _check_call(
-            env, hop, *reference, out, res_out, U.published_partials(hop, rows)
-        )
-        repaired = hop.workspace.layout_repairs - before
-        call_ok &= repaired == expected and (expected > 0) == (
-            env.rank % 4 in (0, 1) or (rows > 1 and env.rank % 4 == 2)
-        )
-        ok &= call_ok
-        _log(
-            env,
-            f"[{hop.impl}] relayout rows={rows}: {stats} repairs={repaired} expected={expected}",
-        )
-
-    # The repaired call inside a CUDA graph, replayed with fresh values.
-    rows = 6
-    static_ref = U.make_raw_tuple(rows, seed=890)
-    static, expected = _relayout(env, static_ref)
-    hop(*static, U.GLM_EPS)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        graph_out, graph_res = hop(*static, U.GLM_EPS)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    for replay in range(4):
-        fresh = U.make_raw_tuple(rows, seed=900 + replay, padded_rows=static_ref[0].shape[0])
-        for dst, src in zip(static, fresh):
-            if dst is not None:
-                dst.copy_(src)
-        graph.replay()
-        torch.cuda.synchronize()
-        call_ok, stats = _check_call(env, hop, *fresh, graph_out, graph_res)
-        ok &= call_ok
-    ok &= hop.error_bits() == 0
-    _log(env, f"[{hop.impl}] relayout graph replay x4 ok={ok}")
-    assert U.all_ranks(ok)
-
-
-@pytest.mark.parametrize(
-    "case",
-    ("fc2_dtype", "scale_shape", "perm_dtype", "residual_shape", "shared_dtype", "gamma_shape"),
-)
-def test_contract_violation_raises_before_any_device_work(env, hop, case):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHopContractError,
-    )
-
-    rows = 6
-    fc2, scales, perm, shared, residual, gamma = U.make_raw_tuple(rows, seed=920)
-    if case == "fc2_dtype":
-        fc2 = fc2.half()
-    elif case == "scale_shape":
-        scales = scales[:, : U.GLM_TOP_K - 1]
-    elif case == "perm_dtype":
-        perm = perm.long()
-    elif case == "residual_shape":
-        residual = residual[:, : U.GLM_HIDDEN // 2]
-    elif case == "shared_dtype":
-        shared = shared.float()
-    else:
-        gamma = gamma[:-8]
-    calls_before = hop.workspace.calls()
-    with pytest.raises(PostMoeHopContractError):
-        hop(fc2, scales, perm, shared, residual, gamma, U.GLM_EPS)
-    assert U.all_ranks(hop.workspace.calls() == calls_before)
-    tensors = U.make_raw_tuple(rows, seed=921)
-    out, res_out = hop(*tensors, U.GLM_EPS)
-    call_ok, stats = _check_call(env, hop, *tensors, out, res_out, U.published_partials(hop, rows))
-    assert U.all_ranks(call_ok and hop.error_bits() == 0), stats
-
-
-def test_health_check_reports_engagement(env, hop):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        check_health,
-    )
-
-    before = check_health("engagement-before", collective=True)
-    for i in range(5):
-        hop(*U.make_raw_tuple(6, seed=940 + i), U.GLM_EPS)
-    after = check_health("engagement-after", collective=True)
-    ok = (
-        after.calls - before.calls == 5
-        and after.error_bits == 0
-        and not after.group_error
-        and after.workspaces == 1
-        and after.calls == hop.workspace.calls()
-    )
-    _log(env, f"[{hop.impl}] health {before} -> {after}")
-    assert U.all_ranks(ok)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py
deleted file mode 100644
index f2d3e482bf3a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_protocol_model.py
+++ /dev/null
@@ -1,287 +0,0 @@
-# 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.
-"""Explicit-state model check of the post_moe_hop Lamport epoch/clear protocol (CPU only).
-
-compute-sanitizer has no race detector for global memory, so the safety of
-the post_moe_hop exchange protocol is checked here on an abstract model, exhaustively
-over every interleaving of its individual memory events for small bounds.
-
-Model of the cross-rank exchange (``kernels.py`` / ``cute_kernel.py``): ``W``
-ranks run the same sequence of calls; call ``n`` has ``rows[n]`` token rows
-and uses Lamport buffer ``n % 3`` (the epoch is the call index; the last
-arriving program advances it, and the next call reads it after the kernel
-boundary, so it is not modelled further). Every rank owns a local copy of the
-three buffers, each ``[row][slot]``; a slot is either the sentinel or the
-value ``(call, publisher)``. Within call ``n`` rank ``r`` performs, in any
-order (the kernel's threads are concurrent):
-
-* ``PUB(r, n, row, d)`` for every row and destination rank ``d``: one
-  multicast store into ``d``'s copy of slot ``[row][r]``. Each store is
-  issued when the call starts and may land at any later time, also after the
-  call ended (no completion guarantee is assumed for remote stores);
-* ``CLR(r, n, row, s)`` for every row of the previous call and every slot:
-  a local store of the sentinel into buffer ``(n - 1) % 3`` of ``r``'s copy;
-* ``POLL(r, n, row)``: enabled once no slot of the row in ``r``'s copy of
-  buffer ``n % 3`` holds the sentinel; it reads every slot.
-
-A call ends when its clears and polls are done; the next call of that rank
-starts only after that (stream order, kernel boundary). Safety: every
-completed poll reads exactly ``(n, s)`` in slot ``s``, every run reaches the
-end without a stuck poll, and every remote store lands. The Triton kernel's
-intra-GPU row-sum exchange is the same scheme with programs as "ranks", one
-shared copy, one program clearing, and a kernel boundary between calls
-(``shared_memory=True``), which is checked separately.
-
-Mutated protocols (two buffers, clearing the current buffer, no clear,
-clearing the current call's row count, polling only one slot) and a timed-out
-poll must be rejected, which shows the checker can find the failures it is
-meant to exclude.
-"""
-
-from collections import deque
-from dataclasses import dataclass, replace
-from typing import Optional, Tuple
-
-import pytest
-
-SENT = None
-
-
-@dataclass(frozen=True)
-class Protocol:
-    buffers: int = 3
-    clear: str = "previous"  # "previous", "current" or "none"
-    clear_rows: str = "previous_call"  # or "current_call"
-    poll_slots: str = "all"  # or "first"
-    shared_memory: bool = False  # Triton row-sum exchange: one copy, one clearing program
-    allow_timeout: bool = False  # a poll may give up once (error path)
-
-
-def _events(proto: Protocol, world: int, rows: Tuple[int, ...], rank: int, n: int):
-    """Local (clear, poll) events of call ``n`` of ``rank``; remote stores are separate."""
-    clears = []
-    if proto.clear != "none" and (not proto.shared_memory or rank == 0):
-        if proto.clear_rows == "previous_call":
-            clear_count = rows[n - 1] if n > 0 else 0
-        else:
-            clear_count = rows[n]
-        clears = [("CLR", r, s) for r in range(clear_count) for s in range(world)]
-    polls = [("POLL", r) for r in range(rows[n])]
-    return frozenset(clears + polls)
-
-
-class _Model:
-    def __init__(self, proto: Protocol, world: int, rows: Tuple[int, ...]):
-        self.proto = proto
-        self.world = world
-        self.rows = rows
-        self.calls = len(rows)
-        self.max_rows = max(rows)
-        self.copies = 1 if proto.shared_memory else world
-
-    def initial(self):
-        mem = tuple(
-            SENT for _ in range(self.copies * self.proto.buffers * self.max_rows * self.world)
-        )
-        # per rank: (call index, remaining local events, call started)
-        ranks = tuple((0, None) for _ in range(self.world))
-        return (ranks, frozenset(), mem, 0)  # (ranks, pending stores, memory, timeouts used)
-
-    def _addr(self, copy: int, buf: int, row: int, slot: int) -> int:
-        return ((copy * self.proto.buffers + buf) * self.max_rows + row) * self.world + slot
-
-    def _cur(self, n: int) -> int:
-        return n % self.proto.buffers
-
-    def _dirty(self, n: int) -> int:
-        if self.proto.clear == "current":
-            return self._cur(n)
-        return (n - 1) % self.proto.buffers
-
-    def successors(self, state):
-        """Yield ``(label, next_state, violation)`` for every enabled event."""
-        ranks, pending, mem, timeouts = state
-        for rank, (n, remaining) in enumerate(ranks):
-            if n >= self.calls:
-                continue
-            if remaining is None:
-                # Shared memory models one kernel's programs: call n starts
-                # only after every program ended call n - 1 (kernel boundary).
-                if self.proto.shared_memory and any(m < n for m, _ in ranks):
-                    continue
-                new_ranks = list(ranks)
-                new_ranks[rank] = (n, _events(self.proto, self.world, self.rows, rank, n))
-                stores = pending | frozenset(
-                    (rank, n, row, dest)
-                    for row in range(self.rows[n])
-                    for dest in ((0,) if self.proto.shared_memory else range(self.world))
-                )
-                yield f"START r{rank} n{n}", (tuple(new_ranks), stores, mem, timeouts), None
-                continue
-            for event in remaining:
-                kind = event[0]
-                if kind == "CLR":
-                    _, row, slot = event
-                    copy = 0 if self.proto.shared_memory else rank
-                    new_mem = list(mem)
-                    new_mem[self._addr(copy, self._dirty(n), row, slot)] = SENT
-                    nxt = self._advance(ranks, rank, n, remaining - {event})
-                    yield (
-                        f"CLR r{rank} n{n} row{row} s{slot}",
-                        (nxt, pending, tuple(new_mem), timeouts),
-                        None,
-                    )
-                else:
-                    _, row = event
-                    copy = 0 if self.proto.shared_memory else rank
-                    slots = range(self.world) if self.proto.poll_slots == "all" else range(1)
-                    values = [mem[self._addr(copy, self._cur(n), row, s)] for s in slots]
-                    timed_out = False
-                    if any(v is SENT for v in values):
-                        if not self.proto.allow_timeout or timeouts:
-                            continue
-                        timed_out = True
-                    violation = None
-                    if not timed_out:
-                        read = [
-                            mem[self._addr(copy, self._cur(n), row, s)] for s in range(self.world)
-                        ]
-                        expected = [(n, s) for s in range(self.world)]
-                        if read != expected:
-                            violation = (
-                                f"r{rank} call {n} row {row} read {read}, expected {expected}"
-                            )
-                    nxt = self._advance(ranks, rank, n, remaining - {event})
-                    label = f"{'TIMEOUT' if timed_out else 'POLL'} r{rank} n{n} row{row}"
-                    yield label, (nxt, pending, mem, timeouts + int(timed_out)), violation
-        for store in pending:
-            publisher, n, row, dest = store
-            new_mem = list(mem)
-            new_mem[self._addr(dest, self._cur(n), row, publisher)] = (n, publisher)
-            yield (
-                f"LAND r{publisher} n{n} row{row} -> r{dest}",
-                (ranks, pending - {store}, tuple(new_mem), timeouts),
-                None,
-            )
-
-    @staticmethod
-    def _advance(ranks, rank, n, remaining):
-        new_ranks = list(ranks)
-        new_ranks[rank] = (n + 1, None) if not remaining else (n, remaining)
-        return tuple(new_ranks)
-
-    def done(self, state) -> bool:
-        ranks, pending, _, _ = state
-        return not pending and all(n >= self.calls for n, _ in ranks)
-
-
-@dataclass
-class CheckResult:
-    states: int
-    violation: Optional[str]
-    trace: Tuple[str, ...]
-
-
-def explore(proto: Protocol, world: int, rows: Tuple[int, ...], max_states: int = 3_000_000):
-    """Breadth-first exploration; returns the first violation with its event trace."""
-    model = _Model(proto, world, rows)
-    start = model.initial()
-    parent = {start: None}
-    queue = deque([start])
-
-    def trace_to(state, last_label):
-        labels = [last_label]
-        while parent[state] is not None:
-            state, label = parent[state]
-            labels.append(label)
-        return tuple(reversed(labels))
-
-    while queue:
-        state = queue.popleft()
-        moved = False
-        for label, nxt, violation in model.successors(state):
-            moved = True
-            if violation is not None:
-                return CheckResult(len(parent), violation, trace_to(state, label))
-            if nxt not in parent:
-                parent[nxt] = (state, label)
-                if len(parent) > max_states:
-                    raise RuntimeError(f"state bound {max_states} exceeded")
-                queue.append(nxt)
-        if not moved and not model.done(state):
-            return CheckResult(
-                len(parent), "stuck: no event enabled before the end", trace_to(state, "<stuck>")
-            )
-    return CheckResult(len(parent), None, ())
-
-
-GOOD = Protocol()
-ROW_PATTERNS = ((1, 1, 1, 1, 1, 1, 1), (2, 1, 2, 2, 1, 1, 2))
-
-
-@pytest.mark.parametrize("rows", ROW_PATTERNS, ids=("rows1", "rows_varying"))
-def test_cross_rank_protocol_is_safe_two_ranks(rows):
-    result = explore(GOOD, world=2, rows=rows)
-    print(
-        f"[post_moe_hop-model] cross-rank W=2 rows={rows}: {result.states} states, "
-        f"violation={result.violation}"
-    )
-    assert result.violation is None, result
-
-
-def test_cross_rank_protocol_is_safe_three_ranks():
-    # Four calls are the fewest that reuse a buffer (call 3 reuses call 0's),
-    # after the clear of call 1 and with calls 1 and 2 in flight.
-    rows = (1, 1, 1, 1)
-    result = explore(GOOD, world=3, rows=rows, max_states=8_000_000)
-    print(
-        f"[post_moe_hop-model] cross-rank W=3 rows={rows}: {result.states} states, "
-        f"violation={result.violation}"
-    )
-    assert result.violation is None, result
-
-
-@pytest.mark.parametrize(
-    "rows", ((1, 1, 1, 1, 1, 1), (2, 1, 2, 1, 2, 2)), ids=("rows1", "rows_varying")
-)
-def test_intra_gpu_row_sum_exchange_is_safe(rows):
-    proto = replace(GOOD, shared_memory=True)
-    result = explore(proto, world=3, rows=rows)
-    print(
-        f"[post_moe_hop-model] row-sum exchange 3 programs rows={rows}: {result.states} states, "
-        f"violation={result.violation}"
-    )
-    assert result.violation is None, result
-
-
-MUTANTS = {
-    "two_buffers": (replace(GOOD, buffers=2), 2, (1, 1, 1, 1, 1)),
-    "clear_current_buffer": (replace(GOOD, clear="current"), 2, (1, 1, 1, 1)),
-    "no_clear": (replace(GOOD, clear="none"), 2, (1, 1, 1, 1, 1)),
-    "clear_current_row_count": (replace(GOOD, clear_rows="current_call"), 2, (2, 1, 1, 2, 2)),
-    "poll_first_slot_only": (replace(GOOD, poll_slots="first"), 2, (1, 1, 1)),
-    "timeout_without_reset": (replace(GOOD, allow_timeout=True), 2, (1, 1, 1, 1, 1)),
-}
-
-
-@pytest.mark.parametrize("name", sorted(MUTANTS))
-def test_mutated_protocol_is_rejected(name):
-    proto, world, rows = MUTANTS[name]
-    result = explore(proto, world=world, rows=rows)
-    print(
-        f"[post_moe_hop-model] mutant {name}: {result.states} states, violation={result.violation}; "
-        f"trace tail {result.trace[-6:]}"
-    )
-    assert result.violation is not None, f"mutant {name} was not rejected"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py
deleted file mode 100644
index ccb4b23d64a1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/post_moe_hop/test_post_moe_hop_real_weights_mpi.py
+++ /dev/null
@@ -1,218 +0,0 @@
-# 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.
-"""post_moe_hop fed by the native NVFP4 runner with real GLM-5.2 weights (collective).
-
-Fixture: real weights of target layers (router, 256 routed NVFP4 experts,
-BF16 shared expert, next-layer RMSNorm weight) loaded on every TP rank through
-the native modules; synthetic activations (see ``post_moe_hop_real_fixture``).
-
-For each case the same input goes through
-
-* native chain: runner ``do_finalize=True`` -> ``shared + routed`` -> native
-  ``AllReduce`` RESIDUAL_RMS_NORM (MNNVL one-shot across hosts), i.e. the
-  served GEN multi-node path;
-* candidate chain: runner ``do_finalize=False`` -> ``PostMoeHop``;
-
-and must agree: the runner's own finalize equals the torch finalize oracle on
-its raw tuple, the residual outputs are bit-exact, the normalized outputs are
-within the error budget frozen in ``test_post_moe_hop_post_moe_mpi.py``. A CUDA-graph
-case replays the captured candidate chain with fresh inputs.
-
-Set ``GLM_MOE_DSA_POST_MOE_HOP_REAL_LAYERS`` (comma-separated layer ids, default ``10,77``).
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import post_moe_hop_mpi_utils as U  # noqa: E402
-
-LAYERS = tuple(
-    int(x) for x in os.environ.get("GLM_MOE_DSA_POST_MOE_HOP_REAL_LAYERS", "10,77").split(",")
-)
-NORM_MAX_ULP = 2
-
-
-def _log(env, message):
-    if env.rank == 0:
-        print(f"[post_moe_hop-real] {message}", flush=True)
-
-
-@pytest.fixture(scope="module")
-def env():
-    e = U.mpi_env()
-    if e.world < 2:
-        pytest.skip("needs a TP group of at least two ranks under srun --mpi=pmix")
-    return e
-
-
-@pytest.fixture(scope="module", params=("triton", "cluster"))
-def hop(env, request):
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-    )
-
-    U.native_allreduce()
-    hop = PostMoeHop(
-        env.mapping, U.GLM_HIDDEN, U.GLM_TOP_K, require_multi_node=False, impl=request.param
-    )
-    assert U.all_ranks(hop.impl == request.param), f"{request.param} kernel unavailable"
-    return hop
-
-
-@pytest.fixture(scope="module", params=LAYERS, ids=lambda layer: f"layer{layer}")
-def real_layer(env, request):
-    from post_moe_hop_real_fixture import build_real_moe_layer
-
-    layer = build_real_moe_layer(env.mapping, request.param)
-    torch.cuda.synchronize()
-    env.comm.barrier()
-    yield layer
-    layer.moe.destroy()
-
-
-def test_native_producer_declares_raw_contract(env, hop, real_layer):
-    """The native NVFP4 MoE module passes the pre-runner producer check."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.post_moe_hop import (  # noqa: E501
-        PostMoeHop,
-        producer_capability,
-    )
-
-    declared, reason = producer_capability(
-        real_layer.moe, U.GLM_HIDDEN, U.GLM_TOP_K, torch.bfloat16
-    )
-    bound = PostMoeHop(
-        env.mapping,
-        U.GLM_HIDDEN,
-        U.GLM_TOP_K,
-        require_multi_node=False,
-        impl=hop.impl,
-        raw_producer=real_layer.moe,
-    )
-    backend = type(getattr(real_layer.moe, "backend", real_layer.moe)).__name__
-    _log(
-        env,
-        f"[{hop.impl}] layer {real_layer.layer} producer {type(real_layer.moe).__name__}/"
-        f"{backend}: declared={declared} {reason!r} engaged={bound.engaged}",
-    )
-    assert U.all_ranks(declared and bound.engaged and bound.workspace is hop.workspace)
-
-
-def _native_chain(layer, x, logits, shared, residual):
-    routed = layer.run_moe(x, logits, do_finalize=True)
-    local = shared + routed
-    norm, new_residual = U.native_fused_allreduce(local, residual, layer.next_gamma, U.GLM_EPS)
-    return routed, norm, new_residual
-
-
-def _candidate_chain(layer, hop, x, logits, shared, residual):
-    fc2, scales, perm = layer.run_moe(x, logits, do_finalize=False)
-    out, new_residual = hop(fc2, scales, perm, shared, residual, layer.next_gamma, U.GLM_EPS)
-    return (fc2, scales, perm), out, new_residual
-
-
-@pytest.mark.parametrize("rows", (1, 6, 16))
-def test_real_runner_native_vs_post_moe_hop(env, hop, real_layer, rows):
-    with torch.inference_mode():
-        x = real_layer.hidden_states(rows, seed=1000 + rows)
-        residual = real_layer.residual(rows, seed=2000 + rows)
-        logits = real_layer.gate(x)
-        shared = real_layer.shared_output(x)
-        routed, native_norm, native_res = _native_chain(real_layer, x, logits, shared, residual)
-        raw, out, new_residual = _candidate_chain(real_layer, hop, x, logits, shared, residual)
-        torch.cuda.synchronize()
-
-    fc2, scales, perm = raw
-    layout = {
-        "fc2_base_mod16": fc2.data_ptr() % 16,
-        "fc2_strides": tuple(fc2.stride()),
-        "scales_contig_base_mod16": (scales.is_contiguous(), scales.data_ptr() % 16),
-        "perm_contig_base_mod16": (perm.is_contiguous(), perm.data_ptr() % 16),
-    }
-    finalize_oracle = U.reference_local_partial(fc2, scales, perm, None)
-    finalize_ulps = U.bf16_ulp_distance(finalize_oracle, routed)
-    res_exact = torch.equal(new_residual.view(torch.int16), native_res.view(torch.int16))
-    norm_stats = U.compare_stats(out, native_norm)
-    res_stats = U.compare_stats(new_residual, native_res)
-    try:
-        torch.testing.assert_close(out, native_norm, rtol=0.05, atol=0.15)
-        torch.testing.assert_close(new_residual, native_res, rtol=0.05, atol=0.15)
-        within_tolerance = True
-    except AssertionError:
-        within_tolerance = False
-    distinct = int(perm[perm >= 0].numel())
-    fixture = {
-        "layer": real_layer.layer,
-        "rows": rows,
-        "scale_dtype": str(scales.dtype),
-        "fc2_shape": list(fc2.shape),
-        "valid_slots": distinct,
-        "finalize_max_ulp": int(finalize_ulps.max()),
-        "finalize_frac_ulp_gt0": float((finalize_ulps > 0).float().mean()),
-    }
-    _log(
-        env,
-        f"[{hop.impl}] {fixture} raw_layout={layout} res_exact={res_exact} "
-        f"res_vs_native={res_stats} norm_vs_native={norm_stats} tol={within_tolerance}",
-    )
-    # The native producer needs no layout repair.
-    ok = (
-        int(finalize_ulps.max()) == 0
-        and res_exact
-        and within_tolerance
-        and norm_stats["max_ulp"] <= NORM_MAX_ULP
-        and hop.error_bits() == 0
-        and hop.workspace.layout_repairs == 0
-    )
-    assert U.all_ranks(ok), f"rank {env.rank}: {fixture} {res_stats} {norm_stats}"
-
-
-def test_real_runner_graph_replay(env, hop, real_layer):
-    rows = 6
-    with torch.inference_mode():
-        x = real_layer.hidden_states(rows, seed=3000)
-        residual = real_layer.residual(rows, seed=3001)
-        logits = real_layer.gate(x)
-        shared = real_layer.shared_output(x)
-        _candidate_chain(real_layer, hop, x, logits, shared, residual)
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            _, graph_out, graph_res = _candidate_chain(real_layer, hop, x, logits, shared, residual)
-        torch.cuda.synchronize()
-        env.comm.barrier()
-        ok = True
-        worst = {}
-        for replay in range(8):
-            x.copy_(real_layer.hidden_states(rows, seed=3100 + replay))
-            residual.copy_(real_layer.residual(rows, seed=3200 + replay))
-            logits.copy_(real_layer.gate(x))
-            shared.copy_(real_layer.shared_output(x))
-            graph.replay()
-            torch.cuda.synchronize()
-            _, native_norm, native_res = _native_chain(real_layer, x, logits, shared, residual)
-            torch.cuda.synchronize()
-            res_exact = torch.equal(graph_res.view(torch.int16), native_res.view(torch.int16))
-            stats = U.compare_stats(graph_out, native_norm)
-            call_ok = res_exact and stats["max_ulp"] <= NORM_MAX_ULP
-            ok &= call_ok
-            if not call_ok and not worst:
-                worst = {"replay": replay, "res_exact": res_exact, **stats}
-        ok &= hop.error_bits() == 0
-    _log(env, f"[{hop.impl}] layer {real_layer.layer} graph replay x8 ok={ok} {worst}")
-    assert U.all_ranks(ok), worst
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py
deleted file mode 100644
index 12e7dfa1b62d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/conftest.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# 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.
-"""Make the shared rope_kv_append test helpers importable under ``--import-mode=importlib``."""
-
-import sys
-from pathlib import Path
-
-_HERE = str(Path(__file__).resolve().parent)
-if _HERE not in sys.path:
-    sys.path.insert(0, _HERE)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py
deleted file mode 100644
index 4463d6a6953b..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/pdl_producer.py
+++ /dev/null
@@ -1,194 +0,0 @@
-# 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.
-"""Test-only delayed metadata producer launched with programmatic dependent launch.
-
-``delayed_write(srcs, dsts, delay_cycles)`` enqueues one kernel on the current
-stream (capturable into a CUDA graph) that
-
-1. executes ``griddepcontrol.launch_dependents`` first, so a following kernel
-   launched with programmatic stream serialization may start right away;
-2. waits for its own predecessor (``griddepcontrol.wait``);
-3. spins for ``delay_cycles`` SM cycles;
-4. copies every ``srcs[i]`` into ``dsts[i]``.
-
-It stands for an in-graph metadata writer (for example a fused DSA metadata
-kernel) that is PDL-chained to the next kernel: a consumer is correct only if
-it reads ``dsts`` after its own ``griddepcontrol.wait``.
-
-``early_read(srcs, dsts)`` is the sensitivity control of that stress: a
-PDL-launched kernel that copies ``srcs`` into ``dsts`` *before* its own
-``griddepcontrol.wait`` (then waits and releases its dependents), i.e. a
-consumer that violates the rule. Placed right after ``delayed_write`` it must
-observe the stale values, which shows the stress opens the window.
-"""
-
-import functools
-from typing import List
-
-import torch
-
-_CPP_DECL = (
-    "void delayed_write(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts, "
-    "int64_t delay_cycles);\n"
-    "void early_read(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts);"
-)
-
-_CUDA_SRC = r"""
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAException.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <cstdint>
-#include <vector>
-
-namespace
-{
-constexpr int kMaxCopies = 8;
-
-struct Copies
-{
-    uint32_t const* src[kMaxCopies];
-    uint32_t* dst[kMaxCopies];
-    int64_t words[kMaxCopies];
-    int32_t n;
-};
-
-__global__ void delayedWriterKernel(Copies c, int64_t delayCycles)
-{
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
-    asm volatile("griddepcontrol.wait;" ::: "memory");
-#endif
-    if (threadIdx.x == 0)
-    {
-        long long const start = clock64();
-        while (clock64() - start < delayCycles)
-        {
-        }
-    }
-    __syncthreads();
-    for (int32_t i = 0; i < c.n; ++i)
-    {
-        for (int64_t w = threadIdx.x; w < c.words[i]; w += blockDim.x)
-        {
-            c.dst[i][w] = c.src[i][w];
-        }
-    }
-}
-
-__global__ void earlyReaderKernel(Copies c)
-{
-    for (int32_t i = 0; i < c.n; ++i)
-    {
-        for (int64_t w = threadIdx.x; w < c.words[i]; w += blockDim.x)
-        {
-            c.dst[i][w] = c.src[i][w];
-        }
-    }
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    asm volatile("griddepcontrol.wait;" ::: "memory");
-    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
-#endif
-}
-
-Copies makeCopies(std::vector<torch::Tensor> const& srcs, std::vector<torch::Tensor> const& dsts)
-{
-    TORCH_CHECK(srcs.size() == dsts.size() && srcs.size() <= kMaxCopies, "up to 8 (src, dst) pairs");
-    Copies c{};
-    c.n = static_cast<int32_t>(srcs.size());
-    for (size_t i = 0; i < srcs.size(); ++i)
-    {
-        TORCH_CHECK(srcs[i].is_cuda() && dsts[i].is_cuda() && srcs[i].is_contiguous() && dsts[i].is_contiguous(),
-            "contiguous CUDA tensors expected");
-        TORCH_CHECK(srcs[i].nbytes() == dsts[i].nbytes() && srcs[i].nbytes() % 4 == 0, "equal sizes, 4-byte words");
-        c.src[i] = static_cast<uint32_t const*>(srcs[i].data_ptr());
-        c.dst[i] = static_cast<uint32_t*>(dsts[i].data_ptr());
-        c.words[i] = static_cast<int64_t>(srcs[i].nbytes() / 4);
-    }
-    return c;
-}
-
-template <typename Kernel, typename... Args>
-void launchPdl(Kernel kernel, Args... args)
-{
-    cudaLaunchConfig_t cfg{};
-    cfg.gridDim = dim3(1);
-    cfg.blockDim = dim3(256);
-    cfg.dynamicSmemBytes = 0;
-    cfg.stream = at::cuda::getCurrentCUDAStream();
-    cudaLaunchAttribute attr[1];
-    attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-    attr[0].val.programmaticStreamSerializationAllowed = 1;
-    cfg.attrs = attr;
-    cfg.numAttrs = 1;
-    C10_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kernel, args...));
-}
-} // namespace
-
-void delayed_write(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts, int64_t delay_cycles)
-{
-    launchPdl(delayedWriterKernel, makeCopies(srcs, dsts), delay_cycles);
-}
-
-void early_read(std::vector<torch::Tensor> srcs, std::vector<torch::Tensor> dsts)
-{
-    launchPdl(earlyReaderKernel, makeCopies(srcs, dsts));
-}
-"""
-
-
-@functools.lru_cache(maxsize=1)
-def _module():
-    import os
-
-    from torch.utils.cpp_extension import load_inline
-
-    # Build for the visible GPU only (the image's TORCH_CUDA_ARCH_LIST includes pre-sm_90 targets).
-    major, minor = torch.cuda.get_device_capability()
-    saved = os.environ.get("TORCH_CUDA_ARCH_LIST")
-    os.environ["TORCH_CUDA_ARCH_LIST"] = f"{major}.{minor}"
-    try:
-        return load_inline(
-            name="rope_append_pdl_delayed_producer",
-            cpp_sources=[_CPP_DECL],
-            cuda_sources=[_CUDA_SRC],
-            functions=["delayed_write", "early_read"],
-            extra_cuda_cflags=["-O2"],
-            verbose=False,
-        )
-    finally:
-        if saved is None:
-            os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
-        else:
-            os.environ["TORCH_CUDA_ARCH_LIST"] = saved
-
-
-def available() -> bool:
-    """Whether the producer can be built (needs nvcc and a compiler at test time)."""
-    try:
-        _module()
-    except (OSError, RuntimeError, ImportError) as err:  # no toolchain in this environment
-        print(f"pdl_producer unavailable: {err}")
-        return False
-    return True
-
-
-def delayed_write(srcs: List[torch.Tensor], dsts: List[torch.Tensor], delay_cycles: int) -> None:
-    _module().delayed_write(list(srcs), list(dsts), int(delay_cycles))
-
-
-def early_read(srcs: List[torch.Tensor], dsts: List[torch.Tensor]) -> None:
-    _module().early_read(list(srcs), list(dsts))
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py
deleted file mode 100644
index ea47ff28bd87..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_env.py
+++ /dev/null
@@ -1,624 +0,0 @@
-# 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.
-"""Native-state environment shared by the rope_kv_append tests and microbenchmark.
-
-It allocates a real ``DSACacheManager`` (FP8 MLA latent pool + FP8 indexer
-K-cache pool, tokens_per_block 64) at GLM-5.2 TP8 decode geometry and builds
-the per-step tensors the native ops read: cache lengths, KV block offsets,
-indexer slot mappings, the DSA block table and request indices. Both the
-native ops (``mla_rope_generation``, ``indexer_k_cache_scatter_op``,
-``convert_req_index_to_global[_grouped]``, the top-k reuse ``copy_``) and the
-rope_kv_append ops are called on the same state so that results can be compared byte for
-byte.
-"""
-
-import os
-from dataclasses import dataclass
-from types import SimpleNamespace
-from typing import List, Optional
-
-import torch
-
-import tensorrt_llm
-import tensorrt_llm.bindings
-from tensorrt_llm._torch.attention.backends.interface import RopeParams
-from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import _compute_slot_mappings
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm.bindings.executor import KvCacheConfig
-from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
-from tensorrt_llm.mapping import Mapping
-
-NUM_HEADS = 8  # GLM-5.2 64 heads at TP8
-Q_LORA_RANK = 2048
-KV_LORA_RANK = 512
-QK_NOPE = 192
-QK_ROPE = 64
-V_HEAD = 256
-LATENT = KV_LORA_RANK + QK_ROPE
-INDEX_HEAD_DIM = 128
-TOPK = 2048
-TOKENS_PER_BLOCK = 64
-QUANT_MODE_FP8_KV = 128
-ROPE_THETA = 8_000_000.0
-MAX_POSITIONS = 1 << 20
-
-
-def load_dev_library() -> None:
-    """Load a development build of the rope_kv_append ops when ``GLM_MOE_DSA_ROPE_KV_APPEND_DEV_LIB`` points to one."""
-    path = os.environ.get("GLM_MOE_DSA_ROPE_KV_APPEND_DEV_LIB")
-    if path and not hasattr(torch.ops.trtllm, "glm_moe_dsa_topk_index_bridge"):
-        torch.ops.load_library(path)
-
-
-CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1  # tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
-PAD = None  # batch entry standing for the CUDA-graph padding dummy request
-PAD_DRAFT_LEN = 5  # served MTP depth; the padding dummy is created per draft length
-
-
-@dataclass
-class Step:
-    """One generation step of a batch of requests (``rows`` query tokens each).
-
-    Request-major like the native metadata: request ``i`` owns token rows
-    ``[i * rows, (i + 1) * rows)`` and row ``i`` of the block offsets, the
-    block table and the indexer block offsets.
-    """
-
-    request_ids: List[int]
-    ctx_lens: List[int]  # cached tokens before this step, per batch entry
-    rows: int
-    is_pad: List[bool]  # entries that are the CUDA-graph padding dummy
-    kv_lens_cuda: torch.Tensor  # [N] int32, cache length after the append
-    kv_lens_cpu: torch.Tensor
-    prompt_lens_cpu: torch.Tensor
-    block_offsets: torch.Tensor  # [num_pools, N, 2, max_blocks]
-    block_table: torch.Tensor  # [N, max_blocks] pool slots, -1 padding
-    req_idx: torch.Tensor  # [N * rows] int32 request index of each token
-    slot_fp8: torch.Tensor  # [N * rows] int64
-    slot_scale: torch.Tensor  # [N * rows] int64
-    idx_block_offsets: torch.Tensor  # [N, max_blocks] indexer K-cache pages
-
-    @property
-    def num_seqs(self) -> int:
-        return len(self.request_ids)
-
-    @property
-    def num_tokens(self) -> int:
-        return self.num_seqs * self.rows
-
-    @property
-    def request_id(self) -> int:
-        return self.request_ids[0]
-
-    @property
-    def ctx_len(self) -> int:
-        return self.ctx_lens[0]
-
-    def kv_lens(self) -> List[int]:
-        return [c + self.rows for c in self.ctx_lens]
-
-    def slot_mappings(self, kv_lens):
-        """Native ``_compute_slot_mappings`` for each request's ``rows`` tokens ending at ``kv_lens``."""
-        if isinstance(kv_lens, int):
-            kv_lens = [kv_lens] * self.num_seqs
-        dev = self.idx_block_offsets.device
-        positions = torch.cat(
-            [torch.arange(k - self.rows, k, dtype=torch.int64, device=dev) for k in kv_lens]
-        )
-        req_ind = torch.arange(self.num_seqs, dtype=torch.int64, device=dev).repeat_interleave(
-            self.rows
-        )
-        return _compute_slot_mappings(
-            positions,
-            self.idx_block_offsets,
-            req_ind,
-            INDEX_HEAD_DIM,
-            TOKENS_PER_BLOCK,
-            128,
-            data_bytes_per_token=INDEX_HEAD_DIM,
-        )
-
-    def set_kv_lens(self, kv_lens: List[int]) -> None:
-        """Advance the step in place (static addresses, as across CUDA-graph replays)."""
-        self.kv_lens_cuda.copy_(torch.tensor(kv_lens, dtype=torch.int32))
-        self.kv_lens_cpu.copy_(torch.tensor(kv_lens, dtype=torch.int32))
-        fp8, scl = self.slot_mappings(kv_lens)
-        self.slot_fp8.copy_(fp8)
-        self.slot_scale.copy_(scl)
-
-
-class RopeAppendEnv:
-    """Real DSA cache manager plus helpers to call native and rope_kv_append ops."""
-
-    def __init__(
-        self,
-        num_layers: int = 4,
-        max_tokens: int = (1 << 20) + 8192,
-        max_seq_len: int = (1 << 20) + 64,
-        max_batch_size: int = 40,
-    ):
-        self.device = torch.device("cuda")
-        self.num_layers = num_layers
-        self.max_seq_len = max_seq_len
-        mapping = Mapping(world_size=1, tp_size=1, rank=0)
-        self.sparse_config = DeepSeekSparseAttentionConfig(
-            index_n_heads=32, index_head_dim=INDEX_HEAD_DIM, index_topk=TOPK
-        )
-        model_config = ModelConfig(
-            mapping=mapping,
-            sparse_attention_config=self.sparse_config,
-            pretrained_config=SimpleNamespace(rms_norm_eps=1e-5),
-        )
-        self.kvm = DSACacheManager(
-            KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False),
-            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-            num_layers=num_layers,
-            num_kv_heads=1,
-            head_dim=LATENT,
-            tokens_per_block=TOKENS_PER_BLOCK,
-            max_seq_len=max_seq_len,
-            max_batch_size=max_batch_size,
-            mapping=mapping,
-            dtype=tensorrt_llm.bindings.DataType.FP8,
-            sparse_attention_config=self.sparse_config,
-            model_config=model_config,
-        )
-        rope = RopeParams(
-            dim=QK_ROPE, theta=ROPE_THETA, max_positions=max_seq_len, duplicate_data=True
-        )
-        _, self.cos_sin = rope.create_rope_const_params()
-        self.kv_scale_orig_quant = torch.ones(1, dtype=torch.float32, device=self.device)
-        self._next_request = 1000
-        self._live: List = []
-        self._pad = {}
-        self._layer_ids = {}
-
-    # ----------------------------------------------------------------- state
-    def pools(self) -> List[torch.Tensor]:
-        """Every cache byte the ops may write: the MLA pool and the indexer pools."""
-        pools = [self.kvm.get_unique_primary_pool().view(torch.uint8)]
-        for pool in self.kvm.indexer_k_cache_pool_per_layer:
-            if pool is not None:
-                pools.append(pool.view(torch.uint8))
-        return pools
-
-    def snapshot(self) -> List[torch.Tensor]:
-        return [p.clone() for p in self.pools()]
-
-    def restore(self, snap: List[torch.Tensor]) -> None:
-        for p, s in zip(self.pools(), snap):
-            p.copy_(s)
-
-    def _padding_dummy(self, rows: int):
-        """The CUDA-graph padding request, created as ``CUDAGraphRunner._get_or_create_padding_dummy``."""
-        rid = CUDA_GRAPH_DUMMY_REQUEST_ID - PAD_DRAFT_LEN
-        if rid not in self._pad:
-            reqs = self.kvm.add_dummy_requests(
-                [rid], token_nums=None, is_gen=True, max_num_draft_tokens=PAD_DRAFT_LEN
-            )
-            assert reqs, "KV pool exhausted"
-            self._live.extend(reqs)
-            self._pad[rid] = reqs[0]
-        # Engine view of a generation dummy: past = prompt_len = token_num - 1 = draft_len.
-        assert rows <= 1 + PAD_DRAFT_LEN
-        return rid, PAD_DRAFT_LEN
-
-    def new_step(self, ctx_len: int, rows: int) -> Step:
-        """One generation request whose cache holds ``ctx_len`` tokens plus this step."""
-        return self.new_batch([ctx_len], rows)
-
-    def new_batch(self, ctx_lens: List[Optional[int]], rows: int) -> Step:
-        """A generation batch; ``PAD`` entries are the CUDA-graph padding dummy (repeated as padding is)."""
-        rids, ctxs, is_pad = [], [], []
-        for ctx in ctx_lens:
-            if ctx is PAD:
-                rid, past = self._padding_dummy(rows)
-                rids.append(rid)
-                ctxs.append(past)
-                is_pad.append(True)
-                continue
-            rid = self._next_request
-            self._next_request += 1
-            reqs = self.kvm.add_dummy_requests(
-                [rid], token_nums=[ctx + 1], is_gen=True, max_num_draft_tokens=rows - 1
-            )
-            assert reqs, "KV pool exhausted"
-            self._live.extend(reqs)
-            rids.append(rid)
-            ctxs.append(ctx)
-            is_pad.append(False)
-        n = len(rids)
-        kv_lens = [c + rows for c in ctxs]
-        dev = self.device
-        offsets = torch.zeros(
-            (self.kvm.num_pools, n, 2, self.kvm.max_blocks_per_seq), dtype=torch.int32, device=dev
-        )
-        self.kvm.copy_batch_block_offsets(offsets, rids, 1, 0, n)
-        pool_idx = self.kvm.get_pool_block_indices(n)  # [n, max_blocks], host
-        bt = pool_idx.clone()
-        for i, kv in enumerate(kv_lens):
-            bt[i, (kv + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK :] = -1
-        idx_offsets = (pool_idx * self.kvm.indexer_k_cache_page_scale).to(dev)
-        step = Step(
-            request_ids=rids,
-            ctx_lens=ctxs,
-            rows=rows,
-            is_pad=is_pad,
-            kv_lens_cuda=torch.tensor(kv_lens, dtype=torch.int32, device=dev),
-            kv_lens_cpu=torch.tensor(kv_lens, dtype=torch.int32),
-            prompt_lens_cpu=torch.tensor(ctxs, dtype=torch.int32),
-            block_offsets=offsets,
-            block_table=bt.to(dev),
-            req_idx=torch.arange(n, dtype=torch.int32, device=dev).repeat_interleave(rows),
-            slot_fp8=torch.empty(n * rows, dtype=torch.int64, device=dev),
-            slot_scale=torch.empty(n * rows, dtype=torch.int64, device=dev),
-            idx_block_offsets=idx_offsets,
-        )
-        fp8, scl = step.slot_mappings(kv_lens)
-        step.slot_fp8.copy_(fp8)
-        step.slot_scale.copy_(scl)
-        return step
-
-    def release(self) -> None:
-        for req in self._live:
-            self.kvm.free_resources(req)
-        self._live = []
-        self._pad = {}
-
-    def shutdown(self) -> None:
-        self.release()
-        self.kvm.shutdown()
-
-    # ------------------------------------------------------------ native ops
-    def native_rope(
-        self, step: Step, layer: int, inp: "RopeInputs", out: "RopeOutputs", q_scaling: float = 1.0
-    ) -> None:
-        """``TrtllmAttention._mla_rope_generation_impl`` with the DSA/MLA decode arguments."""
-        torch.ops.trtllm.mla_rope_generation(
-            inp.fused_q,
-            inp.q_pe,
-            inp.latent,
-            self.cos_sin,
-            out.cu_q,
-            out.cu_kv,
-            out.counter,
-            out.bmm1,
-            out.bmm2,
-            out.quant_q,
-            step.kv_lens_cuda,
-            step.kv_lens_cpu,
-            step.prompt_lens_cpu,
-            0,
-            step.block_offsets,
-            self.kvm.kv_cache_pool_pointers,
-            self.kvm.kv_cache_pool_mapping,
-            None,
-            None,
-            self.kv_scale_orig_quant,
-            None,
-            None,
-            [None, None],
-            1,
-            layer,
-            NUM_HEADS,
-            1,
-            LATENT,
-            0,
-            TOKENS_PER_BLOCK,
-            self.max_seq_len,
-            1,
-            QUANT_MODE_FP8_KV,
-            q_scaling,
-            Q_LORA_RANK,
-            KV_LORA_RANK,
-            QK_NOPE,
-            QK_ROPE,
-            V_HEAD,
-            True,
-            None,
-            1e-6,
-            False,
-            False,
-            False,
-            False,
-            None,
-        )
-
-    def native_scatter(self, step: Step, layer: int, inp: "RopeInputs") -> None:
-        """``Indexer._update_k_cache``."""
-        k_cache = self.kvm.get_indexer_k_cache_buffers(layer)
-        torch.ops.trtllm.indexer_k_cache_scatter_op(
-            inp.k_fp8, inp.k_scale, k_cache, step.slot_fp8, step.slot_scale, inp.k_fp8.shape[0]
-        )
-
-    def layer_ids(self, layers: List[int]) -> torch.Tensor:
-        """Device layer-offset tensor of a group, built once (as the native group struct does)."""
-        key = tuple(layers)
-        if key not in self._layer_ids:
-            self._layer_ids[key] = torch.tensor(
-                [self.kvm.get_primary_pool_page_index_params(m)[1] for m in layers],
-                dtype=torch.int32,
-                device=self.device,
-            )
-        return self._layer_ids[key]
-
-    def native_bridge(
-        self,
-        step: Step,
-        topk: torch.Tensor,
-        layers: List[int],
-        shared: Optional[torch.Tensor],
-        grouped: bool,
-    ) -> torch.Tensor:
-        """Top-k reuse copy + ``transform_local_topk_and_prepare_pool_view[_grouped]``."""
-        if shared is not None:
-            shared[0 : topk.shape[0], : topk.shape[1]].copy_(topk)
-        scale, _ = self.kvm.get_primary_pool_page_index_params(layers[0])
-        stride = scale * TOKENS_PER_BLOCK
-        if grouped:
-            layer_ids = self.layer_ids(layers)
-            return torch.ops.trtllm.convert_req_index_to_global_grouped(
-                step.req_idx,
-                step.block_table,
-                topk,
-                TOKENS_PER_BLOCK,
-                topk.shape[1],
-                stride,
-                layer_ids,
-            )
-        _, off = self.kvm.get_primary_pool_page_index_params(layers[0])
-        return torch.ops.trtllm.convert_req_index_to_global(
-            step.req_idx, step.block_table, topk, TOKENS_PER_BLOCK, topk.shape[1], stride, off
-        )
-
-    # --------------------------------------------------------------- rope_kv_append ops
-    def rope_append_rope(
-        self,
-        step: Step,
-        layer: int,
-        inp: "RopeInputs",
-        out: "RopeOutputs",
-        with_scatter: bool,
-        q_scaling: float = 1.0,
-    ) -> None:
-        idx = [None, None, None, None, None]
-        n_idx = 0
-        if with_scatter:
-            idx = [
-                inp.k_fp8,
-                inp.k_scale,
-                self.kvm.get_indexer_k_cache_buffers(layer),
-                step.slot_fp8,
-                step.slot_scale,
-            ]
-            n_idx = inp.k_fp8.shape[0]
-        torch.ops.trtllm.glm_moe_dsa_rope_kv_append(
-            inp.fused_q,
-            inp.q_pe,
-            inp.latent,
-            self.cos_sin,
-            out.cu_q,
-            out.cu_kv,
-            out.counter,
-            out.bmm1,
-            out.bmm2,
-            out.quant_q,
-            step.kv_lens_cuda,
-            step.prompt_lens_cpu,
-            0,
-            step.block_offsets,
-            self.kvm.kv_cache_pool_pointers,
-            self.kvm.kv_cache_pool_mapping,
-            self.kv_scale_orig_quant,
-            layer,
-            NUM_HEADS,
-            TOKENS_PER_BLOCK,
-            self.max_seq_len,
-            QUANT_MODE_FP8_KV,
-            q_scaling,
-            QK_NOPE,
-            QK_ROPE,
-            KV_LORA_RANK,
-            *idx,
-            n_idx,
-        )
-
-    def rope_append_bridge(
-        self,
-        step: Step,
-        topk: torch.Tensor,
-        layers: List[int],
-        shared: Optional[torch.Tensor],
-        grouped: bool,
-        copy_mode: int = 0,
-        next_n: int = 1,
-        num_accepted=None,
-        copy_rows: Optional[int] = None,
-        num_rows: Optional[int] = None,
-    ) -> torch.Tensor:
-        scale, off = self.kvm.get_primary_pool_page_index_params(layers[0])
-        layer_ids = None
-        if grouped:
-            layer_ids = self.layer_ids(layers)
-        rows = topk.shape[0] if num_rows is None else num_rows
-        crows = 0 if shared is None else (rows if copy_rows is None else copy_rows)
-        return torch.ops.trtllm.glm_moe_dsa_topk_index_bridge(
-            topk,
-            rows,
-            step.req_idx,
-            0,
-            step.block_table,
-            TOKENS_PER_BLOCK,
-            scale * TOKENS_PER_BLOCK,
-            layer_ids,
-            off,
-            shared,
-            0,
-            crows,
-            copy_mode,
-            next_n,
-            num_accepted,
-        )
-
-
-@dataclass
-class RopeInputs:
-    fused_q: torch.Tensor  # [R, H, 576] bf16
-    q: torch.Tensor  # [R, H * 256] bf16, q_pe is a strided view of it
-    q_pe: torch.Tensor
-    latent: torch.Tensor  # [R, 576] bf16
-    k_fp8: torch.Tensor  # [R, 128] fp8
-    k_scale: torch.Tensor  # [R, 1] fp32
-
-
-@dataclass
-class RopeOutputs:
-    quant_q: torch.Tensor
-    cu_q: torch.Tensor
-    cu_kv: torch.Tensor
-    counter: torch.Tensor
-    bmm1: torch.Tensor
-    bmm2: torch.Tensor
-
-    @staticmethod
-    def make(rows: int, device, fill: int = 0xA5, num_seqs: int = 1) -> "RopeOutputs":
-        """Outputs pre-filled with sentinels; the cu_seqlens buffers have 3 slack entries."""
-
-        def i32(n: int) -> torch.Tensor:
-            return torch.full((n,), -7, dtype=torch.int32, device=device)
-
-        return RopeOutputs(
-            quant_q=torch.full((rows, NUM_HEADS, LATENT), fill, dtype=torch.uint8, device=device),
-            cu_q=i32(num_seqs + 4),
-            cu_kv=i32(num_seqs + 4),
-            counter=torch.full((1,), 99, dtype=torch.int32, device=device).view(torch.uint32),
-            bmm1=torch.full((2,), -3.0, dtype=torch.float32, device=device),
-            bmm2=torch.full((1,), -3.0, dtype=torch.float32, device=device),
-        )
-
-    def tensors(self):
-        return [
-            self.quant_q,
-            self.cu_q,
-            self.cu_kv,
-            self.counter.view(torch.int32),
-            self.bmm1,
-            self.bmm2,
-        ]
-
-
-def make_inputs(
-    rows: int, device, scale: float = 1.0, seed: int = 0, specials: bool = False
-) -> RopeInputs:
-    g = torch.Generator(device=device).manual_seed(seed)
-    fused_q = (torch.randn((rows, NUM_HEADS, LATENT), generator=g, device=device) * scale).to(
-        torch.bfloat16
-    )
-    q = (
-        torch.randn((rows, NUM_HEADS * (QK_NOPE + QK_ROPE)), generator=g, device=device) * scale
-    ).to(torch.bfloat16)
-    latent = (torch.randn((rows, LATENT), generator=g, device=device) * scale).to(torch.bfloat16)
-    if specials:
-        vals = torch.tensor(
-            [
-                0.0,
-                -0.0,
-                1e-40,
-                -1e-38,
-                448.0,
-                449.0,
-                1e4,
-                -1e4,
-                float("inf"),
-                float("-inf"),
-                float("nan"),
-                2.0**-10,
-                3.0e-3,
-            ],
-            device=device,
-        ).to(torch.bfloat16)
-        for t in (fused_q.view(-1), q.view(-1), latent.view(-1)):
-            pos = torch.randint(0, t.numel(), (vals.numel() * 8,), generator=g, device=device)
-            t[pos] = vals.repeat(8)
-    q_pe = q.view(rows, NUM_HEADS, QK_NOPE + QK_ROPE)[..., QK_NOPE:]
-    k_fp8 = (
-        torch.randint(0, 256, (rows, INDEX_HEAD_DIM), generator=g, device=device, dtype=torch.int32)
-        .to(torch.uint8)
-        .view(torch.float8_e4m3fn)
-    )
-    k_scale = torch.rand((rows, 1), generator=g, device=device, dtype=torch.float32) + 0.01
-    return RopeInputs(fused_q=fused_q, q=q, q_pe=q_pe, latent=latent, k_fp8=k_fp8, k_scale=k_scale)
-
-
-def make_topk(
-    rows: int, kv_len: int, max_blocks: int, device, seed: int = 0, invalid_frac: float = 0.05
-) -> torch.Tensor:
-    """Top-k ids in [0, kv_len) plus -1 padding and out-of-table ids."""
-    g = torch.Generator(device=device).manual_seed(seed)
-    topk = torch.randint(
-        0, max(kv_len, 1), (rows, TOPK), generator=g, device=device, dtype=torch.int32
-    )
-    mask = torch.rand((rows, TOPK), generator=g, device=device) < invalid_frac
-    topk[mask] = -1
-    beyond = torch.rand((rows, TOPK), generator=g, device=device) < invalid_frac / 2
-    topk[beyond] = torch.randint(
-        kv_len,
-        kv_len + 4 * TOKENS_PER_BLOCK,
-        (int(beyond.sum()),),
-        generator=g,
-        device=device,
-        dtype=torch.int32,
-    )
-    far = torch.rand((rows, TOPK), generator=g, device=device) < invalid_frac / 4
-    topk[far] = max_blocks * TOKENS_PER_BLOCK + 5
-    return topk.contiguous()
-
-
-def make_batch_inputs(
-    step: Step, device, scale: float = 1.0, seed: int = 0, specials: bool = False
-) -> RopeInputs:
-    """Inputs for every token of ``step``.
-
-    Repeated padding entries are the same request writing the same cache slots,
-    so they carry identical rows (as padded graph replays feed identical dummy
-    tokens); otherwise the final cache bytes would depend on write order for
-    native and rope_kv_append alike.
-    """
-    inp = make_inputs(step.num_tokens, device, scale=scale, seed=seed, specials=specials)
-    pads = [i for i, p in enumerate(step.is_pad) if p]
-    r = step.rows
-    for i in pads[1:]:
-        src, dst = slice(pads[0] * r, pads[0] * r + r), slice(i * r, i * r + r)
-        for t in (inp.fused_q, inp.q, inp.latent, inp.k_fp8.view(torch.uint8), inp.k_scale):
-            t[dst] = t[src]
-    return inp
-
-
-def make_batch_topk(
-    step: Step,
-    device,
-    seed: int = 0,
-    rows: Optional[int] = None,
-    kv_lens: Optional[List[int]] = None,
-) -> torch.Tensor:
-    """Per-request top-k rows (``rows`` per request, default ``step.rows``) inside each cache length."""
-    rows = step.rows if rows is None else rows
-    kv_lens = step.kv_lens() if kv_lens is None else kv_lens
-    parts = [
-        make_topk(rows, kv, step.block_table.shape[1], device, seed=seed * 131 + i)
-        for i, kv in enumerate(kv_lens)
-    ]
-    return torch.cat(parts).contiguous()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py
deleted file mode 100644
index fa41b7101dd3..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/rope_append_stack.py
+++ /dev/null
@@ -1,727 +0,0 @@
-# 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.
-"""Real DSA layer stack for the rope_kv_append call-site, composition and region tests.
-
-``Stack`` builds real ``MLA`` modules at GLM-5.2 TP8 attention geometry (8
-local heads, q-LoRA 2048, kv-LoRA 512, nope 192, rope 64, v 256, 32 x 128
-indexer, top-k 2048, FP8 KV cache, block 64, heuristic GVR top-k, CuTe DSL
-logits) in a full/shared indexer pattern plus the MTP layer, adopted as the target's
-``GlmMLA``, a real ``DSACacheManager`` filled with FP8 context, and real generation metadata.
-On decode steps, ``run_seg`` drives the target's ``_forward_dsa_attn`` with given projections;
-``run_proj`` drives the target's ``forward_sparse_attn`` from hidden states, so the
-projection producer (native, or absorbed_query/shared_indexer_absorbed_query/draft_absorbed_query
-when they run) is the real one. ``Engagement`` counts native, rope_kv_append and neighbour-package
-entry points and can record each layer's FP8 query/scheduler/scale buffers and global indices.
-"""
-
-import os
-import sys
-
-import torch
-from rope_append_env import load_dev_library  # noqa: F401  (re-exported for the tests)
-from transformers import PretrainedConfig
-
-import tensorrt_llm
-import tensorrt_llm.bindings
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-    GlmMLA,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention_forward import (  # noqa: E501
-    _forward_dsa_attn,
-    forward_sparse_attn,
-)
-from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams
-from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager
-from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
-from tensorrt_llm._torch.attention.mla import MLA
-from tensorrt_llm._torch.metadata import KVCacheParams
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
-from tensorrt_llm.bindings.executor import KvCacheConfig
-from tensorrt_llm.functional import PositionEmbeddingType
-from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
-from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.models.modeling_utils import QuantConfig
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-
-GLM_HIDDEN = 6144
-
-HIDDEN = 1024  # run_seg does not exercise the projections
-HEADS = 8
-NOPE, ROPE, V, Q_LORA, KV_LORA = 192, 64, 256, 2048, 512
-PATTERN = ["F", "S", "S", "S"]
-MTP_LAYER = len(PATTERN)
-TOKENS_PER_BLOCK = 64
-CTX_REQ_ID = 9  # a request in its context phase (fallback cases)
-CTX_REQ_TOKENS = 256
-
-
-def set_features(ids: str) -> None:
-    """Install, for this process, the plan that runs exactly the fast paths ``ids`` (comma-separated)."""
-    plan_testlib.select(dispatch, ids)
-
-
-class Engagement:
-    """Counts native and rope_kv_append entry points so a silent native fallback cannot pass a comparison."""
-
-    def __init__(self, capture: bool = False):
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (  # noqa: E501
-            rope_kv_append as rope_append_pkg,
-        )
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-            ops as rope_ops,
-        )
-        from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
-        from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-        from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention
-
-        self.counts = {}
-        self._undo = []
-        self.capture = capture
-        # Per rope call: (quant_q, cu_q, cu_kv, scheduler counter, bmm1, bmm2); per remap: indices.
-        self.rope_outputs = []
-        self.global_indices = []
-        self._wrap(Indexer, "_update_k_cache", "native_scatter")
-        self._wrap(TrtllmAttention, "mla_rope_generation", "native_rope")
-        self._wrap(DSATrtllmAttention, "_remap_topk_to_global", "native_remap")
-        self._wrap_neighbours()
-        self.bridge_checks = 0
-        orig_bridge = rope_ops.topk_bridge
-        counted_bridge = self._counted(orig_bridge, "rope_append_bridge_kernel")
-
-        def checked_bridge(topk, metadata, num_rows, **kwargs):
-            out = counted_bridge(topk, metadata, num_rows, **kwargs)
-            self._check_bridge(topk, metadata, num_rows, out, **kwargs)
-            return out
-
-        rope_ops.topk_bridge = checked_bridge
-        self._undo.append(lambda: setattr(rope_ops, "topk_bridge", orig_bridge))
-        del rope_append_pkg
-        registry = dispatch._REGISTRY["rope_kv_append"]
-        for name in (
-            "mla_rope_generation",
-            "remap_topk_to_global",
-            "run_indexer_k_scatter_native",
-            "check_indexer_k_scatter_consumed",
-        ):
-            orig = registry[name]
-            registry[name] = self._counted(orig, "rope_append_" + name)
-            self._undo.append(lambda n=name, o=orig: registry.__setitem__(n, o))
-        orig_defer = registry["defer_indexer_k_scatter"]
-
-        def counted_defer(*args, **kwargs):
-            pkt = orig_defer(*args, **kwargs)
-            key = (
-                "rope_append_deferred_scatter"
-                if pkt is not None
-                else "rope_append_scatter_not_deferred"
-            )
-            self.counts[key] = self.counts.get(key, 0) + 1
-            return pkt
-
-        registry["defer_indexer_k_scatter"] = counted_defer
-        self._undo.append(lambda: registry.__setitem__("defer_indexer_k_scatter", orig_defer))
-
-    def _counted(self, fn, key):
-        def wrapper(*args, **kwargs):
-            self.counts[key] = self.counts.get(key, 0) + 1
-            if key == "rope_append_mla_rope_generation" and kwargs.get("indexer_k") is not None:
-                self.counts["rope_append_fused_scatter"] = (
-                    self.counts.get("rope_append_fused_scatter", 0) + 1
-                )
-            out = fn(*args, **kwargs)
-            if self.capture and key in ("native_rope", "rope_append_mla_rope_generation"):
-                # (mqa, fused_q, q_pe, latent, metadata, cu_q, cu_kv, counter, bmm1, bmm2, quant_q)
-                md = args[4]
-                n = md.num_seqs + 1
-                cu_q, cu_kv, counter, bmm1, bmm2, quant_q = args[5:11]
-                self.rope_outputs.append(
-                    tuple(t.clone() for t in (quant_q, cu_q[:n], cu_kv[:n], counter, bmm1, bmm2))
-                )
-            if self.capture and key in ("native_remap", "rope_append_remap_topk_to_global"):
-                self.global_indices.append(out.clone())
-            return out
-
-        return wrapper
-
-    def _wrap_neighbours(self):
-        """Count engagement of the neighbour packages present in this tree."""
-
-        def wrap_registry(pkg, name, key, engaged):
-            try:
-                orig = dispatch.get_impl(pkg, name)
-            except (ImportError, KeyError):
-                return
-            registry = dispatch._REGISTRY[pkg]
-
-            def wrapper(*args, **kwargs):
-                out = orig(*args, **kwargs)
-                if engaged(out):
-                    self.counts[key] = self.counts.get(key, 0) + 1
-                return out
-
-            registry[name] = wrapper
-            self._undo.append(lambda: registry.__setitem__(name, orig))
-
-        wrap_registry("absorbed_query", "dsa_proj", "aq_packet", lambda o: o is not None)
-        wrap_registry(
-            "absorbed_query", "shared_dsa_proj", "shared_query_packet", lambda o: o is not None
-        )
-        wrap_registry(
-            "absorbed_query", "mtp_dsa_proj", "draft_query_packet", lambda o: o is not None
-        )
-        wrap_registry("indexer_topk", "indexer_decode", "indexer_topk_decode", lambda o: bool(o))
-        try:
-            fmha = dispatch.get_impl("sparse_mla_decode", "fmha_lib")
-        except (ImportError, KeyError):
-            fmha = None
-        if fmha is not None:
-            self._wrap(fmha, "forward", "sparse_mla_fmha")
-
-    def _check_bridge(
-        self,
-        topk,
-        metadata,
-        num_rows,
-        out,
-        *,
-        layer_offset=None,
-        layer_ids=None,
-        page_index_scale,
-        copy=None,
-    ):
-        """Per-call oracle at the real call site: native remap and native copy of the same rows.
-
-        Independent of which rows the (run-to-run nondeterministic) native
-        top-k selected in this run.
-        """
-        torch.cuda.synchronize()
-        tpb = metadata.kv_cache_manager.tokens_per_block
-        nc = metadata.num_contexts
-        block_table = metadata.block_table[nc : metadata.num_seqs]
-        req = (metadata.req_idx_per_token[metadata.num_ctx_tokens : metadata.num_tokens] - nc)[
-            :num_rows
-        ]
-        src = topk[:num_rows]
-        stride = page_index_scale * tpb
-        if layer_ids is not None:
-            want = torch.ops.trtllm.convert_req_index_to_global_grouped(
-                req.to(torch.int32).contiguous(),
-                block_table,
-                src,
-                tpb,
-                src.shape[1],
-                stride,
-                layer_ids,
-            )
-        else:
-            want = torch.ops.trtllm.convert_req_index_to_global(
-                req.to(torch.int32).contiguous(),
-                block_table,
-                src,
-                tpb,
-                src.shape[1],
-                stride,
-                layer_offset,
-            )
-        assert torch.equal(out, want), (
-            "rope_kv_append remap differs from the native remap of the same rows"
-        )
-        if copy is not None:
-            k = topk.shape[1]
-            got = copy.dst[copy.row_start : copy.row_start + copy.rows, :k]
-            if copy.mode == 0:
-                exp = topk[: copy.rows]
-            else:
-                base = torch.arange(copy.rows, device=topk.device, dtype=torch.long) * copy.next_n
-                if copy.num_accepted is None:
-                    off = torch.full_like(base, copy.next_n - 1)
-                else:
-                    off = (copy.num_accepted[: copy.rows].long() - 1).clamp(0, copy.next_n - 1)
-                exp = topk[base + off]
-            assert torch.equal(got, exp), (
-                f"rope_kv_append top-k copy (mode {copy.mode}) differs from native"
-            )
-        self.bridge_checks += 1
-
-    def _wrap(self, owner, attr, key):
-        orig = getattr(owner, attr)
-        setattr(owner, attr, self._counted(orig, key))
-        self._undo.append(lambda: setattr(owner, attr, orig))
-
-    def close(self):
-        for undo in reversed(self._undo):
-            undo()
-
-
-class Stack:
-    def __init__(self, ctx_len: int, pattern=None, real_proj: bool = False):
-        """``real_proj``: GLM hidden size 6144 and random projection weights for ``run_proj``."""
-        pattern = list(PATTERN if pattern is None else pattern)
-        self.pattern = pattern
-        self.mtp_layer = len(pattern)
-        self.hidden = GLM_HIDDEN if real_proj else HIDDEN
-        dev = torch.device("cuda")
-        self.dev = dev
-        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
-        self.pretrained = PretrainedConfig(
-            hidden_size=self.hidden,
-            num_attention_heads=HEADS,
-            num_key_value_heads=HEADS,
-            max_position_embeddings=1 << 20,
-            rope_parameters={"rope_theta": 8_000_000, "rope_type": "default"},
-            qk_rope_head_dim=ROPE,
-            qk_nope_head_dim=NOPE,
-            kv_lora_rank=KV_LORA,
-            q_lora_rank=Q_LORA,
-            v_head_dim=V,
-            model_type="glm_moe_dsa",
-            rms_norm_eps=1e-5,
-            num_hidden_layers=len(pattern),
-            index_topk_pattern=pattern,
-            index_share_for_mtp_iteration=True,
-            index_n_heads=32,
-            index_head_dim=128,
-            index_topk=2048,
-        )
-        self.sparse = DeepSeekSparseAttentionConfig(
-            index_n_heads=32,
-            index_head_dim=128,
-            index_topk=2048,
-            enable_heuristic_topk=True,
-            use_cute_dsl_paged_mqa_logits=True,
-            use_cute_dsl_topk=True,
-            skip_indexer_for_short_seqs=False,
-        )
-        self.model_config = ModelConfig(
-            mapping=self.mapping,
-            sparse_attention_config=self.sparse,
-            pretrained_config=self.pretrained,
-            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8),
-        )
-        g = torch.Generator(device=dev).manual_seed(1234)
-        self.layers = []
-        for idx in range(len(pattern) + 1):
-            mla = MLA(
-                hidden_size=self.hidden,
-                num_attention_heads=HEADS,
-                num_key_value_heads=1,
-                qk_nope_head_dim=NOPE,
-                qk_rope_head_dim=ROPE,
-                v_head_dim=V,
-                q_lora_rank=Q_LORA,
-                kv_lora_rank=KV_LORA,
-                predicted_tokens_per_seq=6,
-                max_position_embeddings=1 << 20,
-                bias=False,
-                pos_embd_params=PositionalEmbeddingParams(
-                    type=PositionEmbeddingType.yarn,
-                    rope=RopeParams.from_config(self.pretrained),
-                    is_neox=False,
-                ),
-                layer_idx=idx,
-                dtype=torch.bfloat16,
-                config=self.model_config,
-            )
-            # The model's DeepseekV32Attention adds only its constructor to MLA; the stack builds
-            # the MLA at GLM geometry directly and adopts it as the model's attention.
-            mla.__class__ = DeepseekV32Attention
-            GlmMLA.adopt(mla)
-            for backend in (mla.mha, mla.mqa):
-                if backend is not None and backend.prepended_fmha_lib_classes:
-                    backend.update_quant_config(backend.quant_config)
-            mla = mla.to(dev)
-            with torch.no_grad():
-                if real_proj:
-                    for name, prm in mla.named_parameters():
-                        if not prm.dtype.is_floating_point:
-                            continue
-                        if name.endswith("bias"):
-                            prm.zero_()
-                        elif "norm" in name:
-                            prm.fill_(1.0)
-                        else:
-                            prm.normal_(0.0, 0.02, generator=g)
-                k_nope = torch.randn(HEADS, NOPE, KV_LORA, generator=g, device=dev) * 0.05
-                v = torch.randn(HEADS, V, KV_LORA, generator=g, device=dev) * 0.05
-                mla.k_b_proj_trans.data = k_nope.transpose(1, 2).contiguous().to(torch.bfloat16)
-                mla.v_b_proj.data = v.contiguous().to(torch.bfloat16)
-                if mla.mqa.indexer is not None:
-                    mla.mqa.indexer.to(dev)
-                    for w in (
-                        mla.mqa.indexer.wq_b.weight,
-                        mla.mqa.indexer.wk.weight,
-                        mla.mqa.indexer.weights_proj.weight,
-                    ):
-                        w.normal_(0.0, 0.02, generator=g)
-                    mla.mqa.indexer.post_load_weights()
-            self.layers.append(mla)
-        assert [m.mqa.indexer is not None for m in self.layers] == [p == "F" for p in pattern] + [
-            True
-        ]
-
-        # A second request with a different cache length for the two-request cases.
-        self.ctx_b = ctx_len // 3 + 4099
-        # Block-table width = max_seq_len / 64 pages; keep it even so indexer_topk (whose logits
-        # tiles span two pages) stays inside its envelope at every context length.
-        pages_per_tile = 2
-        max_seq_len = -(-(ctx_len + 1024) // (pages_per_tile * TOKENS_PER_BLOCK)) * (
-            pages_per_tile * TOKENS_PER_BLOCK
-        )
-        self.kvm = DSACacheManager(
-            KvCacheConfig(
-                max_tokens=ctx_len + self.ctx_b + 8192 + CTX_REQ_TOKENS + 256,
-                enable_block_reuse=False,
-            ),
-            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-            num_layers=len(pattern) + 1,
-            num_kv_heads=1,
-            head_dim=KV_LORA + ROPE,
-            tokens_per_block=TOKENS_PER_BLOCK,
-            max_seq_len=max_seq_len,
-            max_batch_size=3,
-            mapping=self.mapping,
-            dtype=tensorrt_llm.bindings.DataType.FP8,
-            sparse_attention_config=self.sparse,
-            model_config=self.model_config,
-            pretrained_config=self.pretrained,
-        )
-        self.ctx_len = ctx_len
-        self.rid = 7
-        self.rids = [7, 8]
-        self.ctxs = [ctx_len, self.ctx_b]
-        self.reqs = self.kvm.add_dummy_requests(
-            self.rids, token_nums=[c + 1 for c in self.ctxs], is_gen=True, max_num_draft_tokens=5
-        )
-        # FP8 context: valid E4M3 latents and indexer keys, positive indexer scales.
-        pool = self.kvm.get_unique_primary_pool()
-        flat = pool.view(torch.uint8).view(-1)
-        chunk = flat.numel() if flat.numel() <= (1 << 31) else (1 << 30)
-        for start in range(0, flat.numel(), chunk):
-            n = min(chunk, flat.numel() - start)
-            flat[start : start + n].copy_(
-                (torch.randn(n, generator=g, device=dev) * 0.5)
-                .to(torch.float8_e4m3fn)
-                .view(torch.uint8)
-            )
-        for ipool in self.kvm.indexer_k_cache_pool_per_layer:
-            if ipool is None:
-                continue
-            blocks = ipool.view(torch.uint8).view(self.kvm.num_blocks, -1)
-            data = TOKENS_PER_BLOCK * 128
-            blocks[:, :data].copy_(
-                (torch.randn(blocks.shape[0], data, generator=g, device=dev))
-                .to(torch.float8_e4m3fn)
-                .view(torch.uint8)
-            )
-            blocks[:, data:].view(torch.float32).copy_(
-                torch.rand(blocks.shape[0], TOKENS_PER_BLOCK, generator=g, device=dev) + 0.05
-            )
-        self.attn_cls = get_attention_backend(
-            "TRTLLM", sparse_params=self.sparse.to_sparse_params(pretrained_config=self.pretrained)
-        )
-
-    def cached(self, cached: int, nreq: int):
-        """Cached lengths of the first ``nreq`` requests; ``cached`` is request 0's."""
-        return [c + (cached - self.ctx_len) for c in self.ctxs[:nreq]]
-
-    def metadata(
-        self, rows: int, cached: int, nreq: int = 1, ctx_tokens: int = 0, ragged_lens=None
-    ):
-        """Generation metadata for ``nreq`` requests (``rows`` query tokens each).
-
-        ``ctx_tokens`` > 0 prepends the context request ``CTX_REQ_ID`` with that
-        many new tokens (mixed batch; ``nreq`` = 0 gives a pure context batch).
-        ``ragged_lens`` installs ragged verification with these per-request
-        query lengths.
-        """
-        nctx = 1 if ctx_tokens else 0
-        gen_lens = list(ragged_lens) if ragged_lens is not None else [rows] * nreq
-        kwargs = {}
-        if ragged_lens is not None:
-            kwargs.update(enable_ragged_verification=True)
-        md = self.attn_cls.Metadata(
-            seq_lens=torch.tensor([ctx_tokens] * nctx + gen_lens, dtype=torch.int),
-            request_ids=[CTX_REQ_ID] * nctx + self.rids[:nreq],
-            max_num_requests=nctx + nreq,
-            num_contexts=nctx,
-            prompt_lens=[ctx_tokens] * nctx + self.ctxs[:nreq],
-            max_num_tokens=128 + ctx_tokens,
-            kv_cache_manager=self.kvm,
-            kv_cache_params=KVCacheParams(
-                use_cache=True, num_cached_tokens_per_seq=[0] * nctx + self.cached(cached, nreq)
-            ),
-            mapping=self.mapping,
-            sparse_metadata_params=self.sparse.to_sparse_metadata_params(
-                pretrained_config=self.pretrained
-            ),
-            **kwargs,
-        )
-        if ragged_lens is not None:
-            md.ragged_verify_lens = list(ragged_lens)
-        md.update_spec_dec_param(
-            batch_size=nctx + nreq,
-            is_spec_decoding_enabled=True,
-            is_spec_dec_tree=False,
-            is_spec_dec_dynamic_tree=False,
-            max_draft_len=5,
-            max_total_draft_tokens=5,
-        )
-        md.prepare()
-        return md
-
-    def add_context_request(self):
-        """Allocate the context-phase request used by the mixed / context fallback cases."""
-        if not hasattr(self, "ctx_req"):
-            self.ctx_req = self.kvm.add_dummy_requests(
-                [CTX_REQ_ID], token_nums=[CTX_REQ_TOKENS], is_gen=False
-            )
-            self.reqs = list(self.reqs) + list(self.ctx_req)
-
-    def inputs(self, rows: int, cached: int, seed: int, nreq: int = 1):
-        return self.inputs_seg([(rows, c) for c in self.cached(cached, nreq)], seed)
-
-    def inputs_seg(self, segs, seed: int):
-        """Random layer inputs for request-major segments ``(rows, first_position)``."""
-        g = torch.Generator(device=self.dev).manual_seed(seed)
-        n = sum(r for r, _ in segs)
-        q = (torch.randn(n, HEADS * (NOPE + ROPE), generator=g, device=self.dev)).to(torch.bfloat16)
-        ckv = torch.randn(n, KV_LORA, generator=g, device=self.dev).to(torch.bfloat16)
-        k_pe = torch.randn(n, ROPE, generator=g, device=self.dev).to(torch.bfloat16)
-        qr = torch.randn(n, Q_LORA, generator=g, device=self.dev).to(torch.bfloat16)
-        hidden = torch.randn(n, self.hidden, generator=g, device=self.dev).to(torch.bfloat16)
-        pos = torch.cat(
-            [torch.arange(p, p + r, dtype=torch.int32, device=self.dev) for r, p in segs]
-        ).unsqueeze(0)
-        return q, ckv, k_pe, qr, hidden, pos
-
-    def state(self, md):
-        out = [p.view(torch.uint8).clone() for p in [self.kvm.get_unique_primary_pool()]]
-        out += [
-            p.view(torch.uint8).clone()
-            for p in self.kvm.indexer_k_cache_pool_per_layer
-            if p is not None
-        ]
-        if md.shared_topk_indices is not None:
-            out.append(md.shared_topk_indices.clone())
-        return out
-
-    def restore(self, md, state):
-        pools = [self.kvm.get_unique_primary_pool()] + [
-            p for p in self.kvm.indexer_k_cache_pool_per_layer if p is not None
-        ]
-        for p, s in zip(pools, state):
-            p.view(torch.uint8).copy_(s)
-        if md.shared_topk_indices is not None:
-            md.shared_topk_indices.copy_(state[-1])
-        md._group_remap_batched.clear()
-
-    @torch.inference_mode()
-    def run_layers(self, md, layer_ids, rows, cached, seed, nreq: int = 1):
-        return self.run_seg(md, layer_ids, [(rows, c) for c in self.cached(cached, nreq)], seed)
-
-    @torch.inference_mode()
-    def run_seg(self, md, layer_ids, segs, seed, latent_row_stride: int = 0):
-        """Run ``_forward_dsa_attn`` of ``layer_ids`` on request-major segments.
-
-        ``latent_row_stride`` > 576 passes ``latent_cache`` as a row-strided
-        view (outside the fused rope domain).
-        """
-        q, ckv, k_pe, qr, hidden, pos = self.inputs_seg(segs, seed)
-        n = q.shape[0]
-        outs = []
-        for idx in layer_ids:
-            mla = self.layers[idx]
-            inter = []
-            if mla.mqa.indexer is not None:
-                inter = list(mla.mqa.indexer.pre_indexer_proj(qr, hidden, pos))
-            latent = torch.cat([ckv, k_pe], dim=-1)
-            if latent_row_stride:
-                buf = torch.zeros(n, latent_row_stride, dtype=latent.dtype, device=self.dev)
-                buf[:, : KV_LORA + ROPE].copy_(latent)
-                latent = buf[:, : KV_LORA + ROPE]
-            out = torch.empty(n, HEADS * V, dtype=torch.bfloat16, device=self.dev)
-            with dispatch.step(dispatch.Phase.DECODE):
-                _forward_dsa_attn(
-                    mla, q.clone(), ckv.clone(), k_pe.clone(), latent, inter, pos, md, out
-                )
-            outs.append(out)
-        torch.cuda.synchronize()
-        return outs
-
-    @torch.inference_mode()
-    def run_proj(self, md, layer_ids, segs, seed):
-        """Run ``forward_sparse_attn`` (projections + attention) of ``layer_ids`` from hidden states."""
-        _, _, _, _, hidden, pos = self.inputs_seg(segs, seed)
-        outs = []
-        for idx in layer_ids:
-            out = torch.empty(hidden.shape[0], HEADS * V, dtype=torch.bfloat16, device=self.dev)
-            with dispatch.step(dispatch.Phase.DECODE):
-                forward_sparse_attn(self.layers[idx], pos, hidden, md, [out])
-            outs.append(out)
-        torch.cuda.synchronize()
-        return outs
-
-    @torch.inference_mode()
-    def indexer_k(self, mla, rows, cached, seed, nreq=1):
-        """Indexer K inputs of one layer call, as ``_forward_dsa_attn`` slices them."""
-        _, _, _, qr, hidden, pos = self.inputs(rows, cached, seed, nreq)
-        _, k_fp8, k_scale, _, _ = mla.mqa.indexer.pre_indexer_proj(qr, hidden, pos)
-        return k_fp8[: rows * nreq], k_scale[: rows * nreq]
-
-    @staticmethod
-    def _topk_set_diff(a, b):
-        """Per-row number of entries in one selection but not the other (order ignored)."""
-        diffs = []
-        for ra, rb in zip(a, b):
-            sa, sb = set(ra.tolist()), set(rb.tolist())
-            diffs.append(len(sa ^ sb))
-        return diffs
-
-    @staticmethod
-    def _state_equal(i, a, b, is_topk):
-        # The native GVR top-k emits its selection in a run-to-run varying order,
-        # so the top-k reuse buffer is compared as a per-row set.
-        if is_topk:
-            return torch.equal(torch.sort(a, dim=-1).values, torch.sort(b, dim=-1).values)
-        return torch.equal(a, b)
-
-    def _compare_captures(self):
-        """Per-layer FP8 query / scheduler / scale buffers and global indices vs the rope_kv_append-off runs.
-
-        A buffer whose two rope_kv_append-off runs agree must be byte-identical with rope_kv_append;
-        buffers that differ between the rope_kv_append-off runs (selection-dependent
-        global indices under a nondeterministic top-k) are counted and
-        reported instead.
-        """
-        (rope_a, idx_a), (rope_b, idx_b), (rope_w, idx_w) = self.captures
-        assert len(rope_a) == len(rope_b) == len(rope_w) > 0, (len(rope_a), len(rope_w))
-        assert len(idx_a) == len(idx_b) == len(idx_w) > 0, (len(idx_a), len(idx_w))
-        names = ("quant_q", "cu_q", "cu_kv", "counter", "bmm1", "bmm2")
-        report = {"rope_exact": 0, "rope_native_nondet": 0, "idx_exact": 0, "idx_native_nondet": 0}
-        for i, (a, b, w) in enumerate(zip(rope_a, rope_b, rope_w)):
-            for name, x, y, z in zip(names, a, b, w):
-                xb, yb, zb = (t.view(torch.uint8) for t in (x, y, z))
-                if torch.equal(xb, yb):
-                    assert torch.equal(xb, zb), (
-                        f"rope call {i}: {name} differs from the rope_kv_append-off run"
-                    )
-                    report["rope_exact"] += 1
-                else:
-                    report["rope_native_nondet"] += 1
-        # The native top-k emits each row's selection in a run-to-run varying order, so global
-        # indices are compared per row as sets when the two rope_kv_append-off runs differ only in
-        # order.
-        report.update(idx_set_exact=0)
-        for i, (a, b, w) in enumerate(zip(idx_a, idx_b, idx_w)):
-            if torch.equal(a, b):
-                assert torch.equal(a, w), (
-                    f"remap call {i}: global indices differ from the rope_kv_append-off run"
-                )
-                report["idx_exact"] += 1
-                continue
-            sa, sb, sw = (torch.sort(t, dim=-1).values for t in (a, b, w))
-            if torch.equal(sa, sb):
-                assert torch.equal(sa, sw), (
-                    f"remap call {i}: global index sets differ from rope_kv_append-off"
-                )
-                report["idx_set_exact"] += 1
-            else:
-                report["idx_native_nondet"] += 1
-        print("capture comparison:", report)
-        self.capture_report = report
-
-    def compare(self, md, step, base: str = "", capture: bool = False):
-        """Run ``step`` twice with features ``base`` and once with ``base`` + rope_kv_append from the same state.
-
-        Always: cache pools byte-identical to native, and every rope_kv_append bridge call
-        equal to the native remap/copy of the rows it was given (per-call
-        oracle). When the two native runs select the same top-k sets, rope_kv_append must
-        select the same sets too and its sparse-MLA output must be exact (or
-        within the native run-to-run spread). When the native selection itself
-        differs between runs (characterized and printed), the selection-
-        dependent outputs cannot be compared run to run; the per-call oracle
-        and the pools still are.
-        """
-        state0 = self.state(md)
-        runs = []
-        self.engagement = []
-        self.bridge_checks = []
-        self.captures = []
-        rope_append_features = ",".join([f for f in base.split(",") if f] + ["rope_kv_append"])
-        for features in (base, base, rope_append_features):
-            self.restore(md, state0)
-            set_features(features)
-            eng = Engagement(capture=capture)
-            try:
-                outs = step()
-            finally:
-                eng.close()
-            self.engagement.append(eng.counts)
-            self.bridge_checks.append(eng.bridge_checks)
-            self.captures.append((eng.rope_outputs, eng.global_indices))
-            runs.append((outs, self.state(md)))
-        set_features("")
-        if capture:
-            self._compare_captures()
-        print("engagement native/native/rope_kv_append:", self.engagement)
-        assert self.bridge_checks[2] == self.engagement[2].get("rope_append_bridge_kernel", 0)
-        (out_n, st_n), (out_n2, st_n2), (out_w, st_w) = runs
-        topk_idx = len(st_n) - 1 if md.shared_topk_indices is not None else -1
-        for i, (a, b, c) in enumerate(zip(st_n, st_n2, st_w)):
-            if i == topk_idx:
-                continue
-            assert torch.equal(a, b), f"native cache state {i} is not deterministic"
-            if not torch.equal(a, c):
-                raise AssertionError(f"cache state {i}: {(a != c).sum().item()} bytes differ")
-        assert any(not torch.equal(a, b) for a, b in zip(state0, st_n)), "native step wrote nothing"
-        self.native_topk_diff = None
-        if topk_idx >= 0:
-            self.native_topk_diff = self._topk_set_diff(st_n[topk_idx], st_n2[topk_idx])
-            if any(self.native_topk_diff):
-                nz = [(r, d) for r, d in enumerate(self.native_topk_diff) if d]
-                print(
-                    f"native top-k selection differs between native runs (row, |symmetric diff|): {nz}"
-                )
-                return None
-            rope_append_diff = self._topk_set_diff(st_n[topk_idx], st_w[topk_idx])
-            assert not any(rope_append_diff), (
-                f"rope_kv_append run selected different top-k sets: {rope_append_diff}"
-            )
-        report = []
-        for i, (a, a2, b) in enumerate(zip(out_n, out_n2, out_w)):
-            native_spread = (a.float() - a2.float()).abs().max().item()
-            rope_append_diff = (a.float() - b.float()).abs().max().item()
-            rope_append_diff2 = (a2.float() - b.float()).abs().max().item()
-            report.append((native_spread, min(rope_append_diff, rope_append_diff2)))
-            if native_spread == 0.0:
-                assert torch.equal(a.view(torch.int16), b.view(torch.int16)), f"output {i} differs"
-            else:
-                assert min(rope_append_diff, rope_append_diff2) <= 2 * native_spread, (
-                    i,
-                    native_spread,
-                    rope_append_diff,
-                )
-        print("output native-spread / rope_kv_append-diff per layer:", report)
-        return report
-
-    def shutdown(self):
-        for r in self.reqs:
-            self.kvm.free_resources(r)
-        self.kvm.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py
deleted file mode 100644
index 6108700c5caf..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_composition.py
+++ /dev/null
@@ -1,226 +0,0 @@
-# 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.
-"""rope_kv_append composed with its neighbour packages (needs a tree that contains them).
-
-Producers of rope_kv_append's inputs: the native projection + absorption BMM, absorbed_query (target
-layers), absorbed_query + shared_indexer_absorbed_query (shared-indexer target layers) and
-draft_absorbed_query (the MTP layer); consumers of its outputs: the native trtllm-gen sparse MLA or
-sparse_mla_decode, with the native or indexer_topk indexer logits in between. Every case runs
-``forward_sparse_attn`` (projections + attention) from hidden states with the neighbour set alone,
-twice, and with rope_kv_append added, from identical state. Compared byte for byte
-against the rope_kv_append-off runs: every KV and indexer-K pool byte, each layer's FP8
-query buffer, ``cu_q``/``cu_kv`` scheduler offsets, scheduler counter, BMM1/BMM2
-scales and global KV indices (a buffer that differs between the two rope_kv_append-off
-runs is reported, not compared). rope_kv_append and every selected neighbour must engage,
-and rope_kv_append must not change the neighbours' engagement.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-from rope_append_stack import MTP_LAYER, PATTERN, Stack, load_dev_library
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-
-
-def _have(pkg: str, name: str) -> bool:
-    try:
-        dispatch.get_impl(pkg, name)
-    except (ImportError, KeyError):
-        return False
-    return True
-
-
-NEIGHBOURS_PRESENT = all(
-    _have(p, n)
-    for p, n in (
-        ("absorbed_query", "dsa_proj"),
-        ("absorbed_query", "shared_dsa_proj"),
-        ("absorbed_query", "mtp_dsa_proj"),
-        ("sparse_mla_decode", "fmha_lib"),
-        ("indexer_topk", "indexer_decode"),
-    )
-)
-
-
-pytestmark = [
-    pytest.mark.threadleak(enabled=False),
-    pytest.mark.skipif(
-        not NEIGHBOURS_PRESENT,
-        reason="absorbed_query/draft_absorbed_query/shared_indexer_absorbed_query/"
-        "sparse_mla_decode/indexer_topk are not in this tree",
-    ),
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
-        reason="rope_kv_append targets Blackwell decode",
-    ),
-]
-
-CTX = 32768
-TARGET_CASES = {
-    "native": "",
-    "absorbed_query": "absorbed_query",
-    "absorbed_query_shared_query": "absorbed_query,shared_indexer_absorbed_query",
-    "indexer_topk": "indexer_topk",
-    "sparse_mla_decode": "sparse_mla_decode",
-    "all": "absorbed_query,shared_indexer_absorbed_query,draft_absorbed_query,indexer_topk,sparse_mla_decode",
-}
-DRAFT_CASES = {
-    "native": "",
-    "draft_absorbed_query": "draft_absorbed_query",
-    "sparse_mla_decode": "sparse_mla_decode",
-    "all": "absorbed_query,shared_indexer_absorbed_query,draft_absorbed_query,indexer_topk,sparse_mla_decode",
-}
-NEIGHBOUR_KEYS = (
-    "aq_packet",
-    "shared_query_packet",
-    "draft_query_packet",
-    "indexer_topk_decode",
-    "sparse_mla_fmha",
-)
-
-
-@pytest.fixture(scope="module")
-def stack():
-    load_dev_library()
-    # The adopted layers declare the sparse-MLA decode library only under a plan that runs
-    # sparse_mla_decode: the stack is built under the plan of ``plan_testlib.FAST_PATHS_ENV``.
-    plan_testlib.select_from_env(dispatch)
-    s = Stack(CTX, real_proj=True)
-    yield s
-    s.shutdown()
-
-
-def _check_neighbours(features: str, base_counts, rope_append_counts, layers: int, draft: bool):
-    feats = {f for f in features.split(",") if f}
-    for key in NEIGHBOUR_KEYS:
-        assert base_counts.get(key, 0) == rope_append_counts.get(key, 0), (
-            key,
-            base_counts,
-            rope_append_counts,
-        )
-    if "absorbed_query" in feats and not draft:
-        shared = sum(1 for p in PATTERN[:layers] if p == "S")
-        assert (
-            rope_append_counts.get("aq_packet", 0)
-            + rope_append_counts.get("shared_query_packet", 0)
-            == layers
-        )
-        if "shared_indexer_absorbed_query" in feats:
-            assert rope_append_counts.get("shared_query_packet", 0) == shared
-    if "draft_absorbed_query" in feats and draft:
-        assert rope_append_counts.get("draft_query_packet", 0) == 1
-    if "sparse_mla_decode" in feats:
-        assert rope_append_counts.get("sparse_mla_fmha", 0) == layers
-
-
-def _sparse_mla_declared(stack) -> bool:
-    """Whether the stack's layers put the sparse-MLA decode library ahead of the registry's.
-
-    Once declared, sparse_mla_decode serves every eligible layer, so sparse_mla_decode cases run
-    in a process whose stack was built with ``sparse_mla_decode`` in
-    ``plan_testlib.FAST_PATHS_ENV``, the other cases in one without.
-    """
-    return any(getattr(mla.mqa, "prepended_fmha_lib_classes", ()) for mla in stack.layers)
-
-
-def _skip_unless_sparse_mla_matches(stack, features: str):
-    wants_sparse_mla = "sparse_mla_decode" in features.split(",")
-    if wants_sparse_mla != _sparse_mla_declared(stack):
-        pytest.skip(
-            "sparse_mla_decode cases need the stack built with sparse_mla_decode running; others one without"
-            if wants_sparse_mla
-            else "the stack declares sparse_mla_decode, which would serve this case"
-        )
-
-
-@pytest.mark.parametrize("case", list(TARGET_CASES))
-def test_target_layers_composed(stack, case):
-    """Verify step (M = 6) through full + 3 shared target layers."""
-    features = TARGET_CASES[case]
-    _skip_unless_sparse_mla_matches(stack, features)
-    md = stack.metadata(6, stack.ctx_len)
-    stack.compare(
-        md,
-        lambda: stack.run_proj(md, range(len(PATTERN)), [(6, stack.ctx_len)], seed=51),
-        base=features,
-        capture=True,
-    )
-    base_counts, _, rope_append_counts = stack.engagement
-    print(f"{case}: base {base_counts}\n{case}: rope_kv_append {rope_append_counts}")
-    assert (
-        rope_append_counts.get("rope_append_mla_rope_generation") == 4
-        and rope_append_counts.get("rope_append_fused_scatter") == 1
-    )
-    assert (
-        rope_append_counts.get("rope_append_remap_topk_to_global") == 4
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert rope_append_counts.get("rope_append_check_indexer_k_scatter_consumed") == 1
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-    _check_neighbours(features, base_counts, rope_append_counts, len(PATTERN), draft=False)
-    report = stack.capture_report
-    # quant_q may differ between runs without rope_kv_append; the rest may not.
-    assert report["rope_exact"] >= 4 * 5
-    assert report["idx_exact"] + report["idx_set_exact"] + report["idx_native_nondet"] == 4
-
-
-@pytest.mark.parametrize("case", list(DRAFT_CASES))
-def test_mtp_draft_steps_composed(stack, case):
-    """MTP layer: draft step 0 (6 rows, last-accepted-row stash) and a reuse step (1 row)."""
-    features = DRAFT_CASES[case]
-    _skip_unless_sparse_mla_matches(stack, features)
-    md0 = stack.metadata(6, stack.ctx_len)
-    md0.set_in_mtp_draft_loop(True)
-    md0.set_mtp_num_accepted(torch.tensor([4], dtype=torch.int32, device=stack.dev))
-    md0.set_skip_topk(False)
-    stack.compare(
-        md0,
-        lambda: stack.run_proj(md0, [MTP_LAYER], [(6, stack.ctx_len)], seed=52),
-        base=features,
-        capture=True,
-    )
-    base_counts, _, rope_append_counts = stack.engagement
-    print(f"{case} step0: base {base_counts}\n{case} step0: rope_kv_append {rope_append_counts}")
-    assert (
-        rope_append_counts.get("rope_append_fused_scatter") == 1
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-    _check_neighbours(features, base_counts, rope_append_counts, 1, draft=True)
-
-    md1 = stack.metadata(1, stack.ctx_len + 5)
-    md1.set_in_mtp_draft_loop(True)
-    md1.set_skip_topk(True)
-    md1.shared_topk_indices.copy_(md0.shared_topk_indices)
-    stack.compare(
-        md1,
-        lambda: stack.run_proj(md1, [MTP_LAYER], [(1, stack.ctx_len + 5)], seed=53),
-        base=features,
-        capture=True,
-    )
-    base_counts, _, rope_append_counts = stack.engagement
-    print(f"{case} step1: base {base_counts}\n{case} step1: rope_kv_append {rope_append_counts}")
-    assert (
-        rope_append_counts.get("rope_append_fused_scatter") == 1
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-    _check_neighbours(features, base_counts, rope_append_counts, 1, draft=True)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py
deleted file mode 100644
index 6d21dc14199d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_integration.py
+++ /dev/null
@@ -1,501 +0,0 @@
-# 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.
-"""rope_kv_append through the target's DSA call sites (the adopted attention, on decode steps).
-
-Real MLA/DSA modules at GLM-5.2 TP8 attention geometry (8 local heads,
-q-LoRA 2048, kv-LoRA 512, nope 192, rope 64, v 256, 32 x 128 indexer, top-k
-2048, FP8 KV cache, block 64, heuristic GVR top-k, CuTe DSL logits), four
-target layers in the GLM pattern full/shared/shared/shared plus the MTP layer,
-a real ``DSACacheManager`` filled with FP8 context and real generation
-metadata (verify M = 6, MTP draft rows 6 then 1), for one request and for
-two requests with different cache lengths. Each step runs the production
-``_forward_dsa_attn`` natively (no fast path running) and with rope_kv_append
-from identical state. Every KV and indexer-K pool byte must be identical, the
-shared top-k buffer must hold the same per-row selection and the attention
-outputs must be within the native run-to-run spread.
-"""
-
-import pytest
-import torch
-from rope_append_stack import (
-    CTX_REQ_TOKENS,
-    HEADS,
-    KV_LORA,
-    MTP_LAYER,
-    PATTERN,
-    ROPE,
-    Stack,
-    load_dev_library,
-    set_features,
-)
-
-import tensorrt_llm
-import tensorrt_llm.bindings
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-pytestmark = [
-    pytest.mark.threadleak(enabled=False),
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
-        reason="rope_kv_append targets Blackwell decode",
-    ),
-]
-
-
-@pytest.fixture(scope="module", params=[32768, 491_061], ids=["ctx32k", "ctx491k"])
-def stack(request):
-    load_dev_library()
-    s = Stack(request.param)
-    yield s
-    s.shutdown()
-
-
-@pytest.fixture(scope="module")
-def stack32k():
-    """32K context stack for the call-site fallback cases."""
-    load_dev_library()
-    s = Stack(32768)
-    s.add_context_request()
-    yield s
-    s.shutdown()
-
-
-def test_native_state_characterization(stack):
-    """Native run-to-run behaviour: cache bytes are deterministic; the top-k selection is reported.
-
-    The heuristic GVR top-k emits its rows in a varying order and, on some
-    inputs, selects different sets in two runs from identical state; the
-    per-row symmetric difference is printed so it is characterized, not hidden.
-    """
-    md = stack.metadata(6, stack.ctx_len)
-    state0 = stack.state(md)
-    set_features("")
-    stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=1)
-    st_a = stack.state(md)
-    stack.restore(md, state0)
-    stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=1)
-    st_b = stack.state(md)
-    topk_idx = len(st_a) - 1 if md.shared_topk_indices is not None else -1
-    for i, (x, y) in enumerate(zip(st_a, st_b)):
-        if i != topk_idx:
-            assert torch.equal(x, y), f"native cache state {i} differs"
-    if topk_idx >= 0:
-        diff = Stack._topk_set_diff(st_a[topk_idx][:6], st_b[topk_idx][:6])
-        print(
-            f"ctx {stack.ctx_len}: native top-k per-row symmetric difference between runs: {diff}"
-        )
-
-
-def test_target_verify_layers(stack):
-    """Full layer (deferred scatter, grouped bridge with the reuse copy) + 3 shared layers."""
-    md = stack.metadata(6, stack.ctx_len)
-    stack.compare(md, lambda: stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=2))
-    native, _, rope_append_counts = stack.engagement
-    assert native.get("native_scatter") == 1 and native.get("native_rope") == 4
-    assert native.get("native_remap") == 4 and "rope_append_mla_rope_generation" not in native
-    # rope_kv_append: one fused rope per layer (the full layer carries the scatter), one grouped
-    # bridge launch for the whole group, no native rope/scatter/remap.
-    assert (
-        rope_append_counts.get("rope_append_mla_rope_generation") == 4
-        and rope_append_counts.get("rope_append_fused_scatter") == 1
-    )
-    assert (
-        rope_append_counts.get("rope_append_remap_topk_to_global") == 4
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-
-
-@pytest.mark.parametrize("accepted", [1, 4, 6])
-def test_mtp_draft_steps(stack, accepted):
-    """Draft step 0 (row stash of the last accepted row) and a reuse step (rows = 1)."""
-    md0 = stack.metadata(6, stack.ctx_len)
-    md0.set_in_mtp_draft_loop(True)
-    md0.set_mtp_num_accepted(torch.tensor([accepted], dtype=torch.int32, device=stack.dev))
-    md0.set_skip_topk(False)
-    stack.compare(md0, lambda: stack.run_layers(md0, [MTP_LAYER], 6, stack.ctx_len, seed=3))
-    rope_append_counts = stack.engagement[2]
-    assert (
-        rope_append_counts.get("rope_append_fused_scatter") == 1
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-
-    md1 = stack.metadata(1, stack.ctx_len + 5)
-    md1.set_in_mtp_draft_loop(True)
-    md1.set_skip_topk(True)
-    md1.shared_topk_indices.copy_(md0.shared_topk_indices)
-    stack.compare(md1, lambda: stack.run_layers(md1, [MTP_LAYER], 1, stack.ctx_len + 5, seed=4))
-    rope_append_counts = stack.engagement[2]
-    assert (
-        rope_append_counts.get("rope_append_fused_scatter") == 1
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-
-
-def test_target_verify_layers_two_requests(stack):
-    """Two generation requests with different cache lengths (request-major rows 0..5, 6..11)."""
-    md = stack.metadata(6, stack.ctx_len, nreq=2)
-    assert md.num_generations == 2 and md.num_tokens == 12
-    stack.compare(
-        md, lambda: stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=5, nreq=2)
-    )
-    rope_append_counts = stack.engagement[2]
-    assert (
-        rope_append_counts.get("rope_append_mla_rope_generation") == 4
-        and rope_append_counts.get("rope_append_fused_scatter") == 1
-    )
-    assert (
-        rope_append_counts.get("rope_append_remap_topk_to_global") == 4
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-
-
-@pytest.mark.parametrize("accepted", [(1, 6), (4, 2)])
-def test_mtp_draft_steps_two_requests(stack, accepted):
-    """Draft step 0 stashes each request's own last accepted row; reuse step with 2 x 1 rows."""
-    md0 = stack.metadata(6, stack.ctx_len, nreq=2)
-    md0.set_in_mtp_draft_loop(True)
-    md0.set_mtp_num_accepted(torch.tensor(list(accepted), dtype=torch.int32, device=stack.dev))
-    md0.set_skip_topk(False)
-    stack.compare(md0, lambda: stack.run_layers(md0, [MTP_LAYER], 6, stack.ctx_len, seed=6, nreq=2))
-    rope_append_counts = stack.engagement[2]
-    assert (
-        rope_append_counts.get("rope_append_fused_scatter") == 1
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-
-    md1 = stack.metadata(1, stack.ctx_len + 5, nreq=2)
-    md1.set_in_mtp_draft_loop(True)
-    md1.set_skip_topk(True)
-    md1.shared_topk_indices.copy_(md0.shared_topk_indices)
-    stack.compare(
-        md1, lambda: stack.run_layers(md1, [MTP_LAYER], 1, stack.ctx_len + 5, seed=7, nreq=2)
-    )
-    rope_append_counts = stack.engagement[2]
-    assert (
-        rope_append_counts.get("rope_append_fused_scatter") == 1
-        and rope_append_counts.get("rope_append_bridge_kernel") == 1
-    )
-    assert not any(k.startswith("native") for k in rope_append_counts), rope_append_counts
-
-
-def test_lost_deferred_scatter_raises(stack, monkeypatch):
-    """A call site that defers the indexer scatter but drops the packet fails at once.
-
-    Reproduces a call-site placement defect (the packet never reached the
-    generation rope, so nothing wrote the step's indexer keys): the deferring
-    call site's consumed check must raise during that forward.
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
-        attention_forward,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        RopeKvAppendContractError,
-    )
-
-    orig = attention_forward.forward_generation_sparse_attn
-
-    def dropping(*args, deferred_indexer_k=None, **kwargs):
-        return orig(*args, **kwargs)
-
-    monkeypatch.setattr(attention_forward, "forward_generation_sparse_attn", dropping)
-    md = stack.metadata(6, stack.ctx_len)
-    state0 = stack.state(md)
-    set_features("rope_kv_append")
-    try:
-        with pytest.raises(RopeKvAppendContractError, match="not consumed"):
-            stack.run_layers(md, [0], 6, stack.ctx_len, seed=8)
-    finally:
-        set_features("")
-        # The raise skipped the layer's join of the forked prev-top-k copy.
-        stack.layers[0].mqa.indexer.maybe_join_prev_topk_copy()
-        torch.cuda.synchronize()
-        stack.restore(md, state0)
-
-
-def test_deferred_scatter_contract(stack):
-    """The packet is single use, the check sees a missing consumer, undeclared callers get None."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
-        rope_kv_append,
-    )
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.rope_kv_append import (  # noqa: E501
-        RopeKvAppendContractError,
-    )
-
-    md = stack.metadata(6, stack.ctx_len)
-    state0 = stack.state(md)
-    mla = stack.layers[0]
-    k_fp8, k_scale = stack.indexer_k(mla, 6, stack.ctx_len, seed=9)
-    set_features("rope_kv_append")
-    try:
-        assert rope_kv_append.defer_indexer_k_scatter(mla, k_fp8, k_scale, md) is None
-        pkt = rope_kv_append.defer_indexer_k_scatter(mla, k_fp8, k_scale, md, forwards_packet=True)
-        assert pkt is not None and pkt.consumer is None and pkt.layer_idx == 0
-        with pytest.raises(RopeKvAppendContractError, match="not consumed"):
-            rope_kv_append.check_indexer_k_scatter_consumed(pkt)
-        rope_kv_append.run_indexer_k_scatter_native(pkt)
-        rope_kv_append.check_indexer_k_scatter_consumed(pkt)
-        with pytest.raises(RopeKvAppendContractError, match="consumed twice"):
-            rope_kv_append.run_indexer_k_scatter_native(pkt)
-        rope_kv_append.check_indexer_k_scatter_consumed(None)
-        torch.cuda.synchronize()
-        # The native fallback wrote exactly what the native scatter writes.
-        after_rope_append = stack.state(md)
-        stack.restore(md, state0)
-        set_features("")
-        mla.mqa.indexer._update_k_cache(k_fp8, k_scale, md)
-        torch.cuda.synchronize()
-        native = stack.state(md)
-        for i, (a, b) in enumerate(zip(after_rope_append, native)):
-            assert torch.equal(a, b), f"state {i} differs"
-    finally:
-        set_features("")
-        stack.restore(md, state0)
-
-
-def test_pool_view_memo_is_lazy(stack):
-    """The fused bridge leaves the pool-view memo unpopulated; a later reader populates it.
-
-    Accepted deviation from the native remap's side effect: every reader of
-    ``_cached_*`` calls ``_ensure_pool_view_cached()`` first, so values are
-    unchanged and only the memo's ``req_idx - num_contexts`` kernel moves to
-    the first such reader of the step (none on the rope_kv_append/sparse_mla_decode/native-FMHA
-    decode path).
-    """
-    md = stack.metadata(6, stack.ctx_len)
-    state0 = stack.state(md)
-    set_features("rope_kv_append")
-    try:
-        stack.run_layers(md, range(len(PATTERN)), 6, stack.ctx_len, seed=10)
-        assert not md._pool_cache_valid
-        md._ensure_pool_view_cached()
-        nc = md.num_contexts
-        want_req = md.req_idx_per_token[md.num_ctx_tokens : md.num_tokens] - nc
-        assert torch.equal(md._cached_req_idx_gen, want_req)
-        assert torch.equal(md._cached_block_table_gen, md.block_table[nc : md.num_seqs])
-    finally:
-        set_features("")
-        stack.restore(md, state0)
-
-
-def test_default_off_dispatch(stack):
-    """With the selector empty, no rope_kv_append helper is consulted (native call graph)."""
-    set_features("")
-    md = stack.metadata(6, stack.ctx_len)
-    assert dispatch.enabled_features() == frozenset()
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
-        rope_kv_append,
-    )
-
-    assert rope_kv_append.defer_indexer_k_scatter(stack.layers[0], None, None, md) is None
-    assert rope_kv_append.new_bridge_packet(md, True) is None
-
-
-# ------------------------------------------------------------------------------------------
-# Call-site fallbacks: every rejection selects the native path before any rope_kv_append side
-# effect, the pools stay byte-identical to the rope_kv_append-off run and the engagement counters
-# show which path ran.
-
-ROPE_APPEND_ENTRY_KEYS = (
-    "rope_append_mla_rope_generation",
-    "rope_append_remap_topk_to_global",
-    "rope_append_bridge_kernel",
-    "rope_append_deferred_scatter",
-    "rope_append_fused_scatter",
-    "rope_append_run_indexer_k_scatter_native",
-)
-
-
-def _no_rope_append_engaged(counts):
-    return not any(counts.get(k) for k in ROPE_APPEND_ENTRY_KEYS)
-
-
-def test_fallback_rope_ineligible_after_deferral(stack32k):
-    """Rope outside the fused domain (row-strided latent): the deferred scatter runs natively."""
-    st = stack32k
-    md = st.metadata(6, st.ctx_len)
-    segs = [(6, st.ctx_len)]
-    st.compare(
-        md, lambda: st.run_seg(md, range(len(PATTERN)), segs, seed=31, latent_row_stride=640)
-    )
-    native, _, scatter_counts = st.engagement
-    assert native.get("native_scatter") == 1 and native.get("native_rope") == 4
-    assert (
-        scatter_counts.get("rope_append_deferred_scatter") == 1
-        and scatter_counts.get("rope_append_run_indexer_k_scatter_native") == 1
-    )
-    assert scatter_counts.get("rope_append_check_indexer_k_scatter_consumed") == 1
-    assert not scatter_counts.get("rope_append_mla_rope_generation") and not scatter_counts.get(
-        "native_scatter"
-    )
-    assert scatter_counts.get("native_rope") == 4
-    # The bridge does not depend on the latent layout and stays engaged.
-    assert (
-        scatter_counts.get("rope_append_remap_topk_to_global") == 4
-        and scatter_counts.get("rope_append_bridge_kernel") == 1
-    )
-
-
-@pytest.fixture(scope="module")
-def stack_g9():
-    """One full layer leading 8 shared layers: a group larger than the fused bridge supports."""
-    load_dev_library()
-    s = Stack(32768, pattern=["F"] + ["S"] * 8)
-    yield s
-    s.shutdown()
-
-
-def test_fallback_group_larger_than_8(stack_g9):
-    """Group of 9: the leader copies and remaps natively, followers slice the native batch."""
-    st = stack_g9
-    md = st.metadata(6, st.ctx_len)
-    st.compare(md, lambda: st.run_layers(md, range(len(st.pattern)), 6, st.ctx_len, seed=32))
-    _, _, scatter_counts = st.engagement
-    assert scatter_counts.get("native_remap") == 1 and not scatter_counts.get(
-        "rope_append_bridge_kernel"
-    )
-    assert scatter_counts.get("rope_append_remap_topk_to_global") == 9
-    assert (
-        scatter_counts.get("rope_append_mla_rope_generation") == 9
-        and scatter_counts.get("rope_append_fused_scatter") == 1
-    )
-
-
-def test_fallback_context_batch(stack32k):
-    """Pure context batch: no rope_kv_append entry engages."""
-    st = stack32k
-    md = st.metadata(0, st.ctx_len, nreq=0, ctx_tokens=CTX_REQ_TOKENS)
-    assert md.num_contexts == 1 and md.num_generations == 0
-    st.compare(md, lambda: st.run_seg(md, range(len(PATTERN)), [(CTX_REQ_TOKENS, 0)], seed=33))
-    native, _, scatter_counts = st.engagement
-    assert _no_rope_append_engaged(scatter_counts), scatter_counts
-    assert native.get("native_scatter") == 1 and scatter_counts.get("native_scatter") == 1
-
-
-def test_fallback_mixed_batch(stack32k):
-    """Context + generation in one batch: no rope_kv_append entry engages."""
-    st = stack32k
-    md = st.metadata(6, st.ctx_len, nreq=1, ctx_tokens=CTX_REQ_TOKENS)
-    assert md.num_contexts == 1 and md.num_generations == 1
-    segs = [(CTX_REQ_TOKENS, 0), (6, st.ctx_len)]
-    st.compare(md, lambda: st.run_seg(md, range(len(PATTERN)), segs, seed=34))
-    native, _, scatter_counts = st.engagement
-    assert _no_rope_append_engaged(scatter_counts), scatter_counts
-    assert scatter_counts.get("native_scatter") == native.get("native_scatter") == 1
-    assert scatter_counts.get("native_rope") == native.get("native_rope")
-
-
-def test_fallback_ragged_verify(stack32k):
-    """Ragged verification (token-major rows): rope and scatter native; bridge per-call exact."""
-    st = stack32k
-    md = st.metadata(6, st.ctx_len, nreq=2, ragged_lens=[6, 3])
-    assert md.is_ragged_verify and md.token_major_gen_view() is not None
-    c0, c1 = st.cached(st.ctx_len, 2)
-    st.compare(md, lambda: st.run_seg(md, range(len(PATTERN)), [(6, c0), (3, c1)], seed=35))
-    native, _, scatter_counts = st.engagement
-    assert not scatter_counts.get("rope_append_deferred_scatter") and not scatter_counts.get(
-        "rope_append_mla_rope_generation"
-    )
-    assert scatter_counts.get("native_scatter") == native.get("native_scatter") == 1
-    assert scatter_counts.get("native_rope") == native.get("native_rope")
-
-
-class _Over:
-    """Attribute overlay over a real object (unset attributes forward to it)."""
-
-    def __init__(self, base, **overrides):
-        self._base = base
-        self.__dict__.update(overrides)
-
-    def __getattr__(self, name):
-        return getattr(self._base, name)
-
-
-def test_eligibility_rejections(stack32k):
-    """Host-side predicates reject every unsupported mode before device work.
-
-    Starts from a real eligible GLM decode call and changes one condition at
-    a time (NVFP4 / FP4 cache or attention, beam > 1, helix, ragged, strided
-    latent, non-packed fused_q, residual-dim cache, missing KV manager).
-    """
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels import (
-        rope_kv_append,
-    )
-
-    st = stack32k
-    md = st.metadata(6, st.ctx_len)
-    mla = st.layers[0]
-    mqa = mla.mqa
-    n = 6
-    fused_q = torch.zeros(n, HEADS, KV_LORA + ROPE, dtype=torch.bfloat16, device=st.dev)
-    q_pe = torch.zeros(n, HEADS, ROPE, dtype=torch.bfloat16, device=st.dev)
-    latent = torch.zeros(n, KV_LORA + ROPE, dtype=torch.bfloat16, device=st.dev)
-    quant_q = torch.zeros(n, HEADS, KV_LORA + ROPE, dtype=torch.uint8, device=st.dev)
-    s1 = torch.zeros(2, dtype=torch.float32, device=st.dev)
-    s2 = torch.zeros(1, dtype=torch.float32, device=st.dev)
-    k_fp8, k_scale = st.indexer_k(mla, 6, st.ctx_len, seed=36)
-
-    def rope(mqa_=mqa, md_=md, latent_=latent, fused_q_=fused_q):
-        return rope_kv_append.rope_eligible(mqa_, md_, fused_q_, q_pe, latent_, quant_q, s1, s2)
-
-    def defer(mla_=mla, md_=md):
-        return rope_kv_append.defer_indexer_k_scatter(
-            mla_, k_fp8, k_scale, md_, forwards_packet=True
-        )
-
-    set_features("rope_kv_append")
-    try:
-        assert rope() and rope_kv_append.bridge_eligible(md, True) and defer() is not None
-        nvfp4 = tensorrt_llm.bindings.DataType.NVFP4
-        cases = {
-            "fp4_kv_cache": dict(mqa_=_Over(mqa, has_fp4_kv_cache=True)),
-            "no_fp8_kv_cache": dict(mqa_=_Over(mqa, has_fp8_kv_cache=False)),
-            "fp4_mla_attention": dict(mqa_=_Over(mqa, uses_fp4_mla_attention=True)),
-            "beam2": dict(md_=_Over(md, beam_width=2)),
-            "helix": dict(md_=_Over(md, helix_position_offsets=torch.zeros(1))),
-            "ragged": dict(md_=_Over(md, is_ragged_verify=True)),
-            "token_major": dict(md_=_Over(md, token_major_gen_view=lambda: object())),
-            "context": dict(md_=_Over(md, num_contexts=1)),
-            "no_kv_manager": dict(md_=_Over(md, kv_cache_manager=None)),
-            "residual_dim_cache": dict(
-                md_=_Over(
-                    md, kv_cache_manager=_Over(md.kv_cache_manager, mla_kv_cache_residual_dim=64)
-                )
-            ),
-        }
-        for name, kw in cases.items():
-            assert not rope(**kw), name
-        strided = torch.zeros(n, 640, dtype=torch.bfloat16, device=st.dev)[:, : KV_LORA + ROPE]
-        assert not rope(latent_=strided)
-        assert not rope(fused_q_=fused_q.transpose(0, 1).contiguous().transpose(0, 1))
-        for name in ("fp4_kv_cache", "beam2", "helix", "ragged", "token_major", "context"):
-            kw = cases[name]
-            assert (
-                defer(
-                    mla_=_Over(mla, mqa=kw["mqa_"]) if "mqa_" in kw else mla, md_=kw.get("md_", md)
-                )
-                is None
-            ), name
-        assert not rope_kv_append.bridge_eligible(
-            _Over(md, kv_cache_manager=_Over(md.kv_cache_manager, dtype=nvfp4)), True
-        )
-        assert not rope_kv_append.bridge_eligible(_Over(md, num_contexts=1), True)
-        assert not rope_kv_append.bridge_eligible(md, False)
-    finally:
-        set_features("")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py
deleted file mode 100644
index 818ff85ba7b1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/rope_kv_append/test_rope_append_ops.py
+++ /dev/null
@@ -1,597 +0,0 @@
-# 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.
-"""rope_kv_append fused ops versus the native TensorRT-LLM ops, byte for byte.
-
-The oracle is the native op sequence on a real ``DSACacheManager``:
-``indexer_k_cache_scatter_op`` + ``mla_rope_generation`` for the fused rope,
-and the top-k reuse ``copy_`` + ``convert_req_index_to_global[_grouped]`` for
-the bridge. Every cache pool byte, the FP8 query buffer, the FMHA scheduler
-buffers, the bmm scales, the global index tensor and the reuse buffer must be
-identical. Batches cover one request, several live requests with distinct
-cache lengths and the CUDA-graph padding dummy request; a captured-graph
-stress test puts a delayed PDL-launched metadata producer right before each
-rope_kv_append kernel.
-"""
-
-import os
-
-import pdl_producer
-import pytest
-import torch
-from rope_append_env import (
-    NUM_HEADS,
-    PAD,
-    TOKENS_PER_BLOCK,
-    TOPK,
-    RopeAppendEnv,
-    RopeOutputs,
-    load_dev_library,
-    make_batch_inputs,
-    make_batch_topk,
-    make_inputs,
-    make_topk,
-)
-
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-
-pytestmark = [
-    pytest.mark.threadleak(enabled=False),
-    pytest.mark.skipif(
-        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
-        reason="rope_kv_append targets Blackwell decode",
-    ),
-]
-
-LONG_CTX = 491_061
-NEAR_MAX_CTX = (1 << 20) + 8
-
-
-@pytest.fixture(scope="module")
-def env():
-    load_dev_library()
-    e = RopeAppendEnv()
-    yield e
-    e.shutdown()
-
-
-@pytest.fixture(autouse=True)
-def _release(env):
-    yield
-    env.release()
-
-
-def _assert_pools_equal(env, expected):
-    for i, (got, want) in enumerate(zip(env.pools(), expected)):
-        if not torch.equal(got, want):
-            diff = (got != want).nonzero()
-            raise AssertionError(
-                f"pool {i}: {diff.shape[0]} bytes differ, first at {diff[:4].tolist()}"
-            )
-
-
-def _assert_outputs_equal(out_w, out_n):
-    names = ["quant_q", "cu_q", "cu_kv", "counter", "bmm1", "bmm2"]
-    for name, got, want in zip(names, out_w.tensors(), out_n.tensors()):
-        g = got.view(torch.uint8) if got.dtype != torch.uint8 else got
-        w = want.view(torch.uint8) if want.dtype != torch.uint8 else want
-        assert torch.equal(g, w), f"{name} differs: {(g != w).sum().item()} bytes"
-
-
-def _run_rope_pair(env, step, layer, inp, with_scatter, q_scaling=1.0):
-    snap0 = env.snapshot()
-    out_n = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
-    if with_scatter:
-        env.native_scatter(step, layer, inp)
-    env.native_rope(step, layer, inp, out_n, q_scaling=q_scaling)
-    torch.cuda.synchronize()
-    native_pools = env.snapshot()
-    # The native ops must have written the cache (guards against a vacuous comparison).
-    assert any(not torch.equal(a, b) for a, b in zip(native_pools, snap0))
-
-    env.restore(snap0)
-    out_w = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
-    env.rope_append_rope(step, layer, inp, out_w, with_scatter, q_scaling=q_scaling)
-    torch.cuda.synchronize()
-    _assert_pools_equal(env, native_pools)
-    _assert_outputs_equal(out_w, out_n)
-    del snap0, native_pools
-    return out_n
-
-
-@pytest.mark.parametrize("with_scatter", [True, False], ids=["full_layer", "shared_layer"])
-@pytest.mark.parametrize("ctx", [1, 2047, 2048, 2049, 32768, 131072, LONG_CTX, NEAR_MAX_CTX])
-@pytest.mark.parametrize("rows", [6, 1])
-def test_rope_append_scatter_matches_native(env, rows, ctx, with_scatter):
-    step = env.new_step(ctx, rows)
-    # Distinct values per case: requests reuse the same pool blocks across tests.
-    inp = make_inputs(rows, env.device, scale=2.0, seed=ctx * 7 + rows * 3 + int(with_scatter) * 2)
-    layer = 3 if ctx % 2 else 0
-    _run_rope_pair(env, step, layer, inp, with_scatter)
-
-
-@pytest.mark.parametrize(
-    "scale,specials",
-    [(1e-3, False), (40.0, False), (300.0, True)],
-    ids=["tiny", "wide", "saturating+specials"],
-)
-@pytest.mark.parametrize("q_scaling", [1.0, 0.53366])
-def test_rope_value_ranges_match_native(env, scale, specials, q_scaling):
-    step = env.new_step(LONG_CTX, 6)
-    inp = make_inputs(
-        6, env.device, scale=scale, seed=int(scale * 1000) + int(q_scaling * 100), specials=specials
-    )
-    _run_rope_pair(env, step, 2, inp, True, q_scaling=q_scaling)
-
-
-def test_rope_noncontiguous_pages(env):
-    """Requests interleaved in the pool so the written block is not the first slot."""
-    for ctx in (64 * 37 + 5, 1000, 64 * 3 - 1):
-        env.new_step(ctx, 6)
-    step = env.new_step(LONG_CTX - 7, 6)
-    inp = make_inputs(6, env.device, seed=11)
-    _run_rope_pair(env, step, 1, inp, True)
-
-
-@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "per_layer"])
-@pytest.mark.parametrize("ctx", [1, 2047, 2049, 131072, LONG_CTX, NEAR_MAX_CTX])
-@pytest.mark.parametrize("rows", [6, 1])
-def test_bridge_matches_native(env, rows, ctx, grouped):
-    step = env.new_step(ctx, rows)
-    topk = make_topk(rows, ctx + rows, step.block_table.shape[1], env.device, seed=ctx)
-    layers = [0, 1, 2, 3] if grouped else [2]
-    shared_n = torch.full((16, TOPK), -5, dtype=torch.int32, device=env.device)
-    shared_w = shared_n.clone()
-    want = env.native_bridge(step, topk, layers, shared_n, grouped)
-    got = env.rope_append_bridge(step, topk, layers, shared_w, grouped)
-    torch.cuda.synchronize()
-    assert got.shape == want.shape
-    assert torch.equal(got, want)
-    assert torch.equal(shared_w, shared_n)
-
-
-def test_bridge_permuted_page_map(env):
-    """Non-contiguous page map and padding entries in the block table."""
-    step = env.new_step(LONG_CTX, 6)
-    nblk = (LONG_CTX + 6 + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
-    g = torch.Generator(device=env.device).manual_seed(3)
-    perm = torch.randperm(env.kvm.num_blocks, generator=g, device=env.device)[:nblk].to(torch.int32)
-    bt = torch.full_like(step.block_table, -1)
-    bt[0, :nblk] = perm
-    bt[0, 17] = -1  # a hole inside the used range
-    step.block_table = bt
-    topk = make_topk(6, LONG_CTX + 6, bt.shape[1], env.device, seed=5)
-    topk[0, :8] = torch.arange(17 * TOKENS_PER_BLOCK, 17 * TOKENS_PER_BLOCK + 8, device=env.device)
-    for grouped, layers in ((True, [0, 1, 2, 3]), (False, [1])):
-        shared_n = torch.full((8, TOPK), -5, dtype=torch.int32, device=env.device)
-        shared_w = shared_n.clone()
-        want = env.native_bridge(step, topk, layers, shared_n, grouped)
-        got = env.rope_append_bridge(step, topk, layers, shared_w, grouped)
-        assert torch.equal(got, want)
-        assert torch.equal(shared_w, shared_n)
-        assert (want == -1).any()
-
-
-@pytest.mark.parametrize("acc_dtype", [torch.int32, torch.int64, None])
-def test_bridge_mtp_last_accepted_copy(env, acc_dtype):
-    """Draft step 0: MTP last-accepted row stash (native ``_mtp_last_accepted_rows`` + copy)."""
-    next_n, num_gen = 6, 1
-    step = env.new_step(LONG_CTX, next_n)
-    topk = make_topk(next_n, LONG_CTX + next_n, step.block_table.shape[1], env.device, seed=9)
-    for accepted in (1, 3, 6, 0, 9):
-        acc = (
-            None
-            if acc_dtype is None
-            else torch.tensor([accepted], dtype=acc_dtype, device=env.device)
-        )
-        md = type("Md", (), {"mtp_num_accepted": acc})()
-        rows = Indexer._mtp_last_accepted_rows(None, topk, md, 0, num_gen, next_n)
-        shared_n = torch.full((4, TOPK), -5, dtype=torch.int32, device=env.device)
-        shared_n[0 : rows.shape[0], : rows.shape[1]].copy_(rows)
-        want = env.native_bridge(step, topk, [3], None, False)
-
-        shared_w = torch.full((4, TOPK), -5, dtype=torch.int32, device=env.device)
-        got = env.rope_append_bridge(
-            step,
-            topk,
-            [3],
-            shared_w,
-            False,
-            copy_mode=1,
-            next_n=next_n,
-            num_accepted=acc,
-            copy_rows=num_gen,
-        )
-        assert torch.equal(got, want)
-        assert torch.equal(shared_w, shared_n), f"accepted={accepted}"
-
-
-def test_bridge_draft_reuse_copy(env):
-    """Draft steps 1..4: reuse copy shared -> top-k buffer, remap from the shared rows."""
-    step = env.new_step(LONG_CTX + 3, 1)
-    shared = make_topk(4, LONG_CTX + 4, step.block_table.shape[1], env.device, seed=13)
-    buf_n = torch.full((8, TOPK), -5, dtype=torch.int32, device=env.device)
-    buf_n[0:1, :] = shared[:1, :]
-    want = env.native_bridge(step, buf_n[0:1], [3], None, False)
-    buf_w = torch.full((8, TOPK), -5, dtype=torch.int32, device=env.device)
-    got = env.rope_append_bridge(step, shared[:1], [3], buf_w, False, num_rows=1, copy_rows=1)
-    assert torch.equal(got, want)
-    assert torch.equal(buf_w, buf_n)
-
-
-# Multi-request batches. ``PAD`` is the CUDA-graph padding dummy request, created and repeated
-# in the batch exactly as ``CUDAGraphRunner`` pads a generation batch. b33 crosses the 32-request
-# chunk of the cu_kv_seqlens scan.
-BATCHES = {
-    "b2": [LONG_CTX, 2049],
-    "b3pad": [131072 + 13, 7, PAD],
-    "b33": [LONG_CTX] + [37 * i * i + 5 * i + 1 for i in range(1, 31)] + [PAD, PAD],
-}
-assert len(BATCHES["b33"]) == 33
-
-
-def _expected_cu(step):
-    kv = step.kv_lens()
-    cu_q = [NUM_HEADS * step.rows * i for i in range(step.num_seqs + 1)]
-    cu_kv = [sum(kv[:i]) for i in range(step.num_seqs + 1)]
-    return cu_q, cu_kv
-
-
-@pytest.mark.parametrize("with_scatter", [True, False], ids=["full_layer", "shared_layer"])
-@pytest.mark.parametrize("rows", [6, 1])
-@pytest.mark.parametrize("batch", list(BATCHES))
-def test_multi_request_rope_matches_native(env, batch, rows, with_scatter):
-    step = env.new_batch(BATCHES[batch], rows)
-    assert step.num_seqs == len(BATCHES[batch]) and any(step.is_pad) == (PAD in BATCHES[batch])
-    inp = make_batch_inputs(
-        step,
-        env.device,
-        scale=2.0,
-        seed=len(BATCHES[batch]) * 17 + rows * 3 + int(with_scatter) * 2,
-    )
-    out_n = _run_rope_pair(env, step, 3 if rows == 6 else 0, inp, with_scatter)
-    # Oracle sanity: per-request scheduler offsets, the scan carries across 32-request chunks.
-    cu_q, cu_kv = _expected_cu(step)
-    n = step.num_seqs + 1
-    assert out_n.cu_q[:n].tolist() == cu_q and out_n.cu_kv[:n].tolist() == cu_kv
-    assert (out_n.cu_q[n:] == -7).all() and (out_n.cu_kv[n:] == -7).all()
-
-
-@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "per_layer"])
-@pytest.mark.parametrize("rows", [6, 1])
-@pytest.mark.parametrize("batch", list(BATCHES))
-def test_multi_request_bridge_matches_native(env, batch, rows, grouped):
-    step = env.new_batch(BATCHES[batch], rows)
-    topk = make_batch_topk(step, env.device, seed=rows * 7 + int(grouped))
-    layers = [0, 1, 2, 3] if grouped else [2]
-    shared_n = torch.full((step.num_tokens + 8, TOPK), -5, dtype=torch.int32, device=env.device)
-    shared_w = shared_n.clone()
-    want = env.native_bridge(step, topk, layers, shared_n, grouped)
-    got = env.rope_append_bridge(step, topk, layers, shared_w, grouped)
-    torch.cuda.synchronize()
-    assert got.shape == want.shape
-    assert torch.equal(got, want)
-    assert torch.equal(shared_w, shared_n)
-    # Non-vacuous: the result depends on each row's request (different block-table rows).
-    req0 = torch.zeros_like(step.req_idx)
-    scale, off = env.kvm.get_primary_pool_page_index_params(layers[0])
-    as_req0 = torch.ops.trtllm.convert_req_index_to_global(
-        req0, step.block_table, topk, TOKENS_PER_BLOCK, TOPK, scale * TOKENS_PER_BLOCK, off
-    )
-    first = want[0] if grouped else want
-    assert not torch.equal(first[step.rows :], as_req0[step.rows :])
-
-
-@pytest.mark.parametrize("acc_dtype", [torch.int32, torch.int64, None])
-@pytest.mark.parametrize("batch", list(BATCHES))
-def test_multi_request_mtp_last_accepted_copy(env, batch, acc_dtype):
-    """Draft step 0 with several generation requests: stash row i * next_n + clamp(acc_i - 1)."""
-    next_n = 6
-    step = env.new_batch(BATCHES[batch], next_n)
-    num_gen = step.num_seqs
-    topk = make_batch_topk(step, env.device, seed=41)
-    pattern = [1, 3, 6, 0, 9, 2, 5, 4]
-    accepted = [pattern[(i * 5 + 1) % len(pattern)] for i in range(num_gen)]
-    acc = None if acc_dtype is None else torch.tensor(accepted, dtype=acc_dtype, device=env.device)
-    md = type("Md", (), {"mtp_num_accepted": acc})()
-    rows = Indexer._mtp_last_accepted_rows(None, topk, md, 0, num_gen, next_n)
-    assert rows.shape[0] == num_gen
-    shared_n = torch.full((num_gen + 4, TOPK), -5, dtype=torch.int32, device=env.device)
-    shared_n[0:num_gen, :].copy_(rows)
-    want = env.native_bridge(step, topk, [3], None, False)
-    shared_w = torch.full((num_gen + 4, TOPK), -5, dtype=torch.int32, device=env.device)
-    got = env.rope_append_bridge(
-        step,
-        topk,
-        [3],
-        shared_w,
-        False,
-        copy_mode=1,
-        next_n=next_n,
-        num_accepted=acc,
-        copy_rows=num_gen,
-    )
-    torch.cuda.synchronize()
-    assert torch.equal(got, want)
-    assert torch.equal(shared_w, shared_n), f"accepted={accepted}"
-
-
-@pytest.mark.parametrize("batch", list(BATCHES))
-def test_multi_request_draft_reuse_copy(env, batch):
-    """Draft steps 1..4 with several requests: one shared row per request."""
-    step = env.new_batch(BATCHES[batch], 1)
-    n = step.num_seqs
-    shared = make_batch_topk(step, env.device, seed=43, rows=1)
-    buf_n = torch.full((n + 8, TOPK), -5, dtype=torch.int32, device=env.device)
-    buf_n[0:n, :] = shared[:n, :]
-    want = env.native_bridge(step, buf_n[0:n], [3], None, False)
-    buf_w = torch.full((n + 8, TOPK), -5, dtype=torch.int32, device=env.device)
-    got = env.rope_append_bridge(step, shared[:n], [3], buf_w, False, num_rows=n, copy_rows=n)
-    torch.cuda.synchronize()
-    assert torch.equal(got, want)
-    assert torch.equal(buf_w, buf_n)
-
-
-GRAPH_BATCHES = {"single": [LONG_CTX], "b3pad": BATCHES["b3pad"], "b33": BATCHES["b33"]}
-
-
-@pytest.mark.parametrize("batch", list(GRAPH_BATCHES))
-def test_cuda_graph_replay(env, batch):
-    """Capture both rope_kv_append ops once and replay with new values and advancing cache lengths.
-
-    Live requests advance by ``rows`` per replay as in successive verify steps;
-    padding entries keep the dummy's fixed length, as padded replays do.
-    """
-    rows, spec = 6, GRAPH_BATCHES[batch]
-    alloc = [c if c is PAD else c + 4 * rows for c in spec]
-    step = env.new_batch(alloc, rows)
-    base = [c if pad else c - 4 * rows for c, pad in zip(step.ctx_lens, step.is_pad)]
-
-    def kv_at(i):
-        return [b + rows if pad else b + (i + 1) * rows for b, pad in zip(base, step.is_pad)]
-
-    inp = make_batch_inputs(step, env.device, seed=21)
-    out_w = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
-    topk = make_batch_topk(step, env.device, seed=22, kv_lens=kv_at(0))
-    shared_w = torch.full((step.num_tokens + 8, TOPK), -5, dtype=torch.int32, device=env.device)
-
-    def set_step(i):
-        step.set_kv_lens(kv_at(i))
-        new = make_batch_inputs(step, env.device, seed=100 + i)
-        for dst, src in (
-            (inp.fused_q, new.fused_q),
-            (inp.q, new.q),
-            (inp.latent, new.latent),
-            (inp.k_fp8, new.k_fp8),
-            (inp.k_scale, new.k_scale),
-        ):
-            dst.copy_(src)
-        topk.copy_(make_batch_topk(step, env.device, seed=200 + i, kv_lens=kv_at(i)))
-
-    set_step(0)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):  # warm up outside capture
-        env.rope_append_rope(step, 2, inp, out_w, True)
-        env.rope_append_bridge(step, topk, [0, 1, 2, 3], shared_w, True)
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        env.rope_append_rope(step, 2, inp, out_w, True)
-        out_static = env.rope_append_bridge(step, topk, [0, 1, 2, 3], shared_w, True)
-
-    for i in range(1, 4):
-        set_step(i)
-        torch.cuda.synchronize()
-        snap0 = env.snapshot()
-        out_n = RopeOutputs.make(step.num_tokens, env.device, num_seqs=step.num_seqs)
-        env.native_scatter(step, 2, inp)
-        env.native_rope(step, 2, inp, out_n)
-        shared_n = shared_w.clone()
-        want = env.native_bridge(step, topk, [0, 1, 2, 3], shared_n, True)
-        torch.cuda.synchronize()
-        native_pools = env.snapshot()
-        env.restore(snap0)
-        graph.replay()
-        torch.cuda.synchronize()
-        _assert_pools_equal(env, native_pools)
-        _assert_outputs_equal(out_w, out_n)
-        assert torch.equal(out_static, want)
-        assert torch.equal(shared_w, shared_n)
-    assert NUM_HEADS == 8
-
-
-# --------------------------------------------------------------------------------------------
-# Delayed-producer stress: in the captured graph a PDL-launched producer that releases its
-# dependents at entry and writes the per-step metadata ~40 us later runs immediately before each
-# rope_kv_append kernel. Correct rope_kv_append kernels read that metadata only after
-# griddepcontrol.wait.
-
-DELAY_CYCLES = 80_000  # ~40 us at the sm_103 SM clock
-STRESS_REPLAYS = 5
-
-
-def _pdl_enabled() -> bool:
-    return os.environ.get("TRTLLM_ENABLE_PDL", "0") == "1"
-
-
-def _delayed_producer_run(env):
-    """Return, per replay, the names of rope_kv_append outputs that differ from the native oracle.
-
-    Also returns, per replay, whether a control consumer that reads the cache
-    lengths *before* its own wait (``pdl_producer.early_read``) right after the
-    first producer saw the stale values, i.e. whether the stress opened the
-    window a pre-wait reader would fall into.
-    """
-    rows = 6
-    live = [LONG_CTX, 2049 + 64 * 5]
-    step = env.new_batch([c + rows for c in live] + [PAD], rows)
-    kv_new = step.kv_lens()  # this verify step (the allocated capacity) ...
-    kv_old = [k if pad else k - rows for k, pad in zip(kv_new, step.is_pad)]  # ... and the previous
-    dev = env.device
-
-    # "Old" (stale) and "new" values of every per-step metadata tensor rope_kv_append reads.
-    step.set_kv_lens(kv_old)
-    old = {
-        "kv": step.kv_lens_cuda.clone(),
-        "fp8": step.slot_fp8.clone(),
-        "scale": step.slot_scale.clone(),
-        "offsets": step.block_offsets.clone(),
-    }
-    step.set_kv_lens(kv_new)
-    new = {
-        "kv": step.kv_lens_cuda.clone(),
-        "fp8": step.slot_fp8.clone(),
-        "scale": step.slot_scale.clone(),
-        "offsets": step.block_offsets.clone(),
-    }
-    old["offsets"][:, [0, 1]] = new["offsets"][
-        :, [1, 0]
-    ]  # stale page map: the two live rows swapped
-    new["req"] = step.req_idx.clone()
-    old["req"] = new["req"].clone()
-    old["req"][:rows], old["req"][rows : 2 * rows] = 1, 0  # stale request index of the live rows
-    new["acc"] = torch.tensor([2, 5, 1], dtype=torch.int32, device=dev)
-    old["acc"] = torch.tensor([6, 1, 4], dtype=torch.int32, device=dev)
-    num_acc = old["acc"].clone()
-
-    live_rope = [step.kv_lens_cuda, step.slot_fp8, step.slot_scale, step.block_offsets]
-    new_rope = [new["kv"], new["fp8"], new["scale"], new["offsets"]]
-    old_rope = [old["kv"], old["fp8"], old["scale"], old["offsets"]]
-
-    def set_meta(src_rope, src_req, src_acc):
-        for dst, src in zip(live_rope, src_rope):
-            dst.copy_(src)
-        step.req_idx.copy_(src_req)
-        num_acc.copy_(src_acc)
-
-    inp = make_batch_inputs(step, dev, seed=71)
-    topk = make_batch_topk(step, dev, seed=72, kv_lens=kv_new)
-    layers = [0, 1, 2, 3]
-
-    # Native oracle on the new metadata.
-    set_meta(new_rope, new["req"], new["acc"])
-    snap0 = env.snapshot()
-    out_n = RopeOutputs.make(step.num_tokens, dev, num_seqs=step.num_seqs)
-    env.native_scatter(step, 2, inp)
-    env.native_rope(step, 2, inp, out_n)
-    shared_n = torch.full((step.num_tokens + 8, TOPK), -5, dtype=torch.int32, device=dev)
-    want_grouped = env.native_bridge(step, topk, layers, shared_n, True)
-    md = type("Md", (), {"mtp_num_accepted": num_acc})()
-    stash_rows = Indexer._mtp_last_accepted_rows(None, topk, md, 0, step.num_seqs, rows)
-    stash_n = torch.full((8, TOPK), -5, dtype=torch.int32, device=dev)
-    stash_n[0 : step.num_seqs].copy_(stash_rows)
-    want_stash = env.native_bridge(step, topk, [3], None, False)
-    torch.cuda.synchronize()
-    native_pools = env.snapshot()
-    env.restore(snap0)
-
-    out_w = RopeOutputs.make(step.num_tokens, dev, num_seqs=step.num_seqs)
-    shared_w = torch.full_like(shared_n, -5)
-    stash_w = torch.full_like(stash_n, -5)
-    seen_early = torch.zeros_like(step.kv_lens_cuda)
-    kv_copy = torch.empty_like(step.kv_lens_cuda)
-
-    def body():
-        # Control: a pre-wait reader in the same position as a rope_kv_append kernel.
-        pdl_producer.delayed_write([new["kv"]], [kv_copy], DELAY_CYCLES)
-        pdl_producer.early_read([kv_copy], [seen_early])
-        pdl_producer.delayed_write(new_rope, live_rope, DELAY_CYCLES)
-        env.rope_append_rope(step, 2, inp, out_w, True)
-        pdl_producer.delayed_write([new["req"]], [step.req_idx], DELAY_CYCLES)
-        grouped = env.rope_append_bridge(step, topk, layers, shared_w, True)
-        pdl_producer.delayed_write([new["acc"]], [num_acc], DELAY_CYCLES)
-        stash = env.rope_append_bridge(
-            step,
-            topk,
-            [3],
-            stash_w,
-            False,
-            copy_mode=1,
-            next_n=rows,
-            num_accepted=num_acc,
-            copy_rows=step.num_seqs,
-        )
-        return grouped, stash
-
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):  # warm up (and build the producer) outside capture
-        body()
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        got_grouped, got_stash = body()
-
-    report = []
-    stale_seen = []
-    for _ in range(STRESS_REPLAYS):
-        env.restore(snap0)
-        set_meta(old_rope, old["req"], old["acc"])
-        kv_copy.copy_(old["kv"])
-        seen_early.fill_(-1)
-        shared_w.fill_(-5)
-        stash_w.fill_(-5)
-        torch.cuda.synchronize()
-        graph.replay()
-        torch.cuda.synchronize()
-        bad = [
-            f"pool{i}"
-            for i, (a, b) in enumerate(zip(env.pools(), native_pools))
-            if not torch.equal(a, b)
-        ]
-        names = ["quant_q", "cu_q", "cu_kv", "counter", "bmm1", "bmm2"]
-        bad += [
-            n
-            for n, a, b in zip(names, out_w.tensors(), out_n.tensors())
-            if not torch.equal(a.view(torch.uint8), b.view(torch.uint8))
-        ]
-        bad += [
-            n
-            for n, a, b in (
-                ("grouped_remap", got_grouped, want_grouped),
-                ("shared_copy", shared_w, shared_n),
-                ("stash_remap", got_stash, want_stash),
-                ("stash_copy", stash_w, stash_n),
-            )
-            if not torch.equal(a, b)
-        ]
-        report.append(bad)
-        stale_seen.append(torch.equal(seen_early, old["kv"]))
-    print(
-        f"delayed producer: rope_kv_append outputs differing from native per replay {report}; "
-        f"control pre-wait reader saw stale cache lengths per replay {stale_seen}"
-    )
-    return report, stale_seen
-
-
-@pytest.fixture(scope="module")
-def pdl_producer_ok():
-    if not pdl_producer.available():
-        pytest.skip("cannot build the test PDL producer")
-    return True
-
-
-def test_pdl_delayed_producer(env, pdl_producer_ok):
-    """rope_kv_append is correct with a PDL-chained metadata writer right before each of its kernels.
-
-    The control consumer that reads before its wait must see the stale values
-    in at least one replay, so the stress does open the window (with PDL on).
-    """
-    report, stale_seen = _delayed_producer_run(env)
-    assert all(not bad for bad in report), report
-    if _pdl_enabled():
-        assert any(stale_seen), "the control pre-wait reader never saw stale metadata"
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py
deleted file mode 100644
index 13f7293c7e36..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_pre_moe_hop_mpi.py
+++ /dev/null
@@ -1,241 +0,0 @@
-# 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.
-"""Multi-rank tests of the router_gemm pre-MoE hop with the native MNNVL all-reduce.
-
-Run one pytest process per GPU, every rank executing the same tests:
-
-    srun --mpi=pmix -N2 -n8 python3 -m pytest --import-mode=importlib <this>
-
-With 4 ranks on one node the MNNVL strategy is requested explicitly; with 8
-ranks over two nodes ``AUTO`` selects it as in the served model. The chain
-under test is: stock row-parallel o_proj -> native fused AR + residual +
-RMSNorm -> router_gemm router -> native routing (``noaux_tc_op``), compared bitwise
-against the same chain with the native gate, eagerly and under CUDA graph
-replay, and across ranks.
-"""
-
-import hashlib
-import os
-
-import pytest
-import torch
-from mpi4py import MPI
-from utils.llm_data import llm_models_root
-
-import tensorrt_llm  # noqa: F401
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.router_gemm import (  # noqa: E501
-    pre_moe_hop,
-    router_gemm,
-)
-from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams
-from tensorrt_llm.functional import AllReduceStrategy
-from tensorrt_llm.mapping import Mapping
-
-COMM = MPI.COMM_WORLD
-WORLD = COMM.Get_size()
-RANK = COMM.Get_rank()
-
-pytestmark = pytest.mark.skipif(WORLD < 2, reason="launch one process per GPU with srun --mpi=pmix")
-
-HIDDEN = 6144
-EXPERTS = 256
-HEAD_SLICE = 2048  # 64 heads * 256 v_head / TP8 per rank (o_proj in-features)
-EPS = 1e-5
-
-
-def _bits_equal(a, b):
-    return torch.equal(
-        a.contiguous().view(torch.int32 if a.element_size() == 4 else torch.int16),
-        b.contiguous().view(torch.int32 if b.element_size() == 4 else torch.int16),
-    )
-
-
-def _digest(t):
-    return hashlib.sha256(t.contiguous().cpu().view(torch.uint8).numpy().tobytes()).hexdigest()
-
-
-def _checkpoint_dir():
-    """The GLM-5.2 NVFP4 checkpoint: ``$GLM52_NVFP4_CKPT``, else under ``llm_models_root()``."""
-    root = llm_models_root()
-    default = os.path.join(str(root), "GLM-5.2-NVFP4") if root is not None else ""
-    return os.environ.get("GLM52_NVFP4_CKPT", default)
-
-
-@pytest.fixture(scope="module")
-def env():
-    ngpu = torch.cuda.device_count()
-    local = int(os.environ.get("SLURM_LOCALID", RANK % ngpu))
-    torch.cuda.set_device(local)
-    mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=ngpu)
-    strategy = AllReduceStrategy.AUTO if mapping.is_multi_node() else AllReduceStrategy.MNNVL
-    allreduce = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16)
-    assert allreduce.mnnvl_allreduce is not None, "MNNVL AR not engaged"
-    return mapping, allreduce
-
-
-@pytest.fixture(scope="module")
-def layer_weights(env):
-    """Real gate/norm/o_proj weights of one MoE layer if the checkpoint is set."""
-    from safetensors import safe_open
-
-    ckpt = _checkpoint_dir()
-    dev = torch.device("cuda", torch.cuda.current_device())
-    if ckpt and os.path.exists(os.path.join(ckpt, "model.safetensors.index.json")):
-        import json
-
-        with open(os.path.join(ckpt, "model.safetensors.index.json")) as f:
-            wmap = json.load(f)["weight_map"]
-
-        def load(name):
-            with safe_open(os.path.join(ckpt, wmap[name]), framework="pt") as f:
-                return f.get_tensor(name)
-
-        pre = "model.layers.40"
-        o_proj = load(f"{pre}.self_attn.o_proj.weight")
-        return {
-            "source": "real",
-            "gate": load(f"{pre}.mlp.gate.weight").to(dev),
-            "bias": load(f"{pre}.mlp.gate.e_score_correction_bias").to(dev).to(torch.bfloat16),
-            "norm": load(f"{pre}.post_attention_layernorm.weight").to(dev),
-            "o_proj_local": o_proj[:, RANK * HEAD_SLICE : (RANK + 1) * HEAD_SLICE]
-            .contiguous()
-            .to(dev),
-        }
-    g = torch.Generator(device=dev)
-    g.manual_seed(5)
-    return {
-        "source": "synthetic",
-        "gate": (torch.randn(EXPERTS, HIDDEN, generator=g, device=dev) * 0.02).to(torch.bfloat16),
-        "bias": torch.zeros(EXPERTS, dtype=torch.bfloat16, device=dev),
-        "norm": (1 + 0.1 * torch.randn(HIDDEN, generator=g, device=dev)).to(torch.bfloat16),
-        "o_proj_local": (torch.randn(HIDDEN, HEAD_SLICE, generator=g, device=dev) * 0.02).to(
-            torch.bfloat16
-        ),
-    }
-
-
-def _inputs(rows, seed):
-    """Per-rank attention output slice and the shared residual."""
-    dev = torch.device("cuda", torch.cuda.current_device())
-    g = torch.Generator(device=dev)
-    g.manual_seed(seed * 1000 + RANK)
-    attn = torch.randn(rows, HEAD_SLICE, generator=g, device=dev).to(torch.bfloat16)
-    g.manual_seed(seed * 1000 + 999)
-    residual = (torch.randn(rows, HIDDEN, generator=g, device=dev) * 4).to(torch.bfloat16)
-    return attn, residual
-
-
-def _native_chain(allreduce, w, attn, residual):
-    partial = torch.nn.functional.linear(attn, w["o_proj_local"])
-    hidden, res = allreduce(
-        partial,
-        all_reduce_params=AllReduceParams(
-            fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
-            residual=residual,
-            norm_weight=w["norm"],
-            eps=EPS,
-            trigger_completion_at_end=False,
-        ),
-    )
-    logits = torch.ops.trtllm.dsv3_router_gemm_op(
-        hidden, w["gate"].t(), bias=None, out_dtype=torch.float32
-    )
-    return hidden, res, logits
-
-
-def _router_gemm_chain(allreduce, w, attn, residual):
-    partial = torch.nn.functional.linear(attn, w["o_proj_local"])
-    hidden, res, logits = pre_moe_hop(allreduce, partial, residual, w["norm"], EPS, w["gate"])
-    assert logits is not None
-    return hidden, res, logits
-
-
-def _route(logits, bias):
-    return torch.ops.trtllm.noaux_tc_op(logits, bias, 1, 1, 8, 2.5)
-
-
-def test_weight_source(layer_weights):
-    """Real checkpoint weights are used whenever the checkpoint is set."""
-    if os.path.exists(os.path.join(_checkpoint_dir(), "model.safetensors.index.json")):
-        assert layer_weights["source"] == "real"
-
-
-@pytest.mark.parametrize("rows", [1, 3, 6, 16])
-def test_hop_bitwise_and_rank_agreement(env, layer_weights, rows):
-    _, allreduce = env
-    attn, residual = _inputs(rows, seed=rows)
-    h_n, r_n, l_n = _native_chain(allreduce, layer_weights, attn, residual)
-    h_p, r_p, l_p = _router_gemm_chain(allreduce, layer_weights, attn, residual)
-    torch.cuda.synchronize()
-    assert _bits_equal(h_p, h_n) and _bits_equal(r_p, r_n)
-    assert _bits_equal(l_p, l_n)
-    v_n, i_n = _route(l_n, layer_weights["bias"])
-    v_p, i_p = _route(l_p, layer_weights["bias"])
-    assert torch.equal(i_p, i_n) and _bits_equal(v_p, v_n)
-    digests = COMM.allgather((_digest(l_p), _digest(i_p)))
-    assert all(d == digests[0] for d in digests), digests
-
-
-def test_hop_graph_replay(env, layer_weights):
-    """Capture native and router_gemm chains (with the MoE-style aux fork) and replay."""
-    _, allreduce = env
-    rows = 6
-    attn, residual = _inputs(rows, seed=77)
-    aux = torch.cuda.Stream()
-    ev0, ev1 = torch.cuda.Event(), torch.cuda.Event()
-    w_shared = torch.randn(512, HIDDEN, device="cuda").to(torch.bfloat16)
-
-    def body(chain):
-        hidden, res, logits = chain(allreduce, layer_weights, attn, residual)
-        ev0.record()
-        ids = _route(logits, layer_weights["bias"])[1]
-        with torch.cuda.stream(aux):
-            ev0.wait()
-            shared = torch.nn.functional.linear(hidden, w_shared)
-            ev1.record()
-        ev1.wait()
-        return hidden, res, logits, ids, shared
-
-    router_gemm(
-        torch.zeros(rows, HIDDEN, dtype=torch.bfloat16, device="cuda"), layer_weights["gate"]
-    )
-    outs = {}
-    graphs = {}
-    for name, chain in (("native", _native_chain), ("router_gemm", _router_gemm_chain)):
-        body(chain)
-        torch.cuda.synchronize()
-        COMM.Barrier()
-        graph = torch.cuda.CUDAGraph()
-        s = torch.cuda.Stream()
-        s.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(s), torch.cuda.graph(graph, stream=s):
-            outs[name] = body(chain)
-        torch.cuda.current_stream().wait_stream(s)
-        graphs[name] = graph
-    torch.cuda.synchronize()
-    COMM.Barrier()
-    for step in range(30):
-        a, r = _inputs(rows, seed=1000 + step)
-        attn.copy_(a)
-        residual.copy_(r)
-        graphs["router_gemm"].replay()
-        graphs["native"].replay()
-        torch.cuda.synchronize()
-        for x, y in zip(outs["router_gemm"], outs["native"]):
-            assert _bits_equal(x, y), f"step {step}"
-        if step % 10 == 0:
-            digests = COMM.allgather(_digest(outs["router_gemm"][2]))
-            assert all(d == digests[0] for d in digests)
-    COMM.Barrier()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py
deleted file mode 100644
index 37abab2187ff..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/router_gemm/test_router_gemm.py
+++ /dev/null
@@ -1,446 +0,0 @@
-# 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.
-"""Single-GPU correctness of the router_gemm against the native oracle.
-
-The oracle is the native ``torch.ops.trtllm.dsv3_router_gemm_op`` that
-``DeepseekV3Gate`` runs for this model. The router_gemm kernel reproduces its
-arithmetic, so every comparison below is bitwise (FP32 bit patterns), not a
-tolerance.
-"""
-
-import ast
-import json
-import os
-from types import SimpleNamespace
-
-import pytest
-import torch
-from utils.llm_data import llm_models_root
-
-import tensorrt_llm  # noqa: F401  (loads the native torch ops)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.router_gemm import (  # noqa: E501
-    hop,
-    router_gemm,
-    router_gemm_supported,
-    router_logits,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or not 100
-    <= torch.cuda.get_device_capability()[0] * 10 + torch.cuda.get_device_capability()[1]
-    < 110,
-    reason="needs an SM100-family GPU",
-)
-
-HIDDEN = 6144
-EXPERTS = 256
-TOP_K = 8
-ROUTED_SCALING = 2.5
-
-# GLM-5.2 NVFP4 checkpoint directory. Router (mlp.gate) and
-# post_attention_layernorm weights are BF16 there and are read tensor by tensor;
-# tests that need them skip when it is absent. Activations are always synthetic
-# (real per-layer activations need a captured fixture).
-CKPT_ENV = "GLM52_NVFP4_CKPT"
-
-
-@pytest.fixture(scope="module")
-def real_layer_weights():
-    """Return ``fn(layer) -> dict`` of real gate/bias/norm weights, or skip."""
-    root = llm_models_root()
-    default = os.path.join(str(root), "GLM-5.2-NVFP4") if root is not None else ""
-    ckpt = os.environ.get(CKPT_ENV, default)
-    index = os.path.join(ckpt or "", "model.safetensors.index.json")
-    if not ckpt or not os.path.exists(index):
-        pytest.skip("GLM-5.2 checkpoint not available")
-    from safetensors import safe_open
-
-    with open(index) as f:
-        weight_map = json.load(f)["weight_map"]
-    cache = {}
-
-    def load(name):
-        with safe_open(os.path.join(ckpt, weight_map[name]), framework="pt") as f:
-            return f.get_tensor(name)
-
-    def get(layer):
-        if layer not in cache:
-            pre = f"model.layers.{layer}"
-            cache[layer] = {
-                "gate": load(f"{pre}.mlp.gate.weight"),
-                "bias": load(f"{pre}.mlp.gate.e_score_correction_bias"),
-                "post_norm": load(f"{pre}.post_attention_layernorm.weight"),
-            }
-        return cache[layer]
-
-    return get
-
-
-def native_logits(hidden, weight):
-    return torch.ops.trtllm.dsv3_router_gemm_op(
-        hidden, weight.t(), bias=None, out_dtype=torch.float32
-    )
-
-
-def assert_bitwise(actual, expected):
-    assert actual.shape == expected.shape and actual.dtype == expected.dtype
-    a_bits = actual.contiguous().view(torch.int32)
-    e_bits = expected.contiguous().view(torch.int32)
-    mismatch = (a_bits != e_bits).sum().item()
-    if mismatch:
-        diff = (actual.double() - expected.double()).abs().max().item()
-        raise AssertionError(
-            f"{mismatch}/{actual.numel()} logits differ in bits (max abs diff {diff:.3e})"
-        )
-
-
-def normed_activation(rows, hidden, gen, norm_weight=None, scale=1.0):
-    """Synthetic post-RMSNorm activation, BF16 (as the AR/norm emits)."""
-    x = torch.randn((rows, hidden), generator=gen, device="cuda") * scale
-    x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5)
-    if norm_weight is not None:
-        x = x * norm_weight.float()
-    return x.to(torch.bfloat16)
-
-
-def random_weight(hidden, gen, scale=0.02):
-    return (torch.randn((EXPERTS, hidden), generator=gen, device="cuda") * scale).to(torch.bfloat16)
-
-
-@pytest.fixture
-def gen():
-    g = torch.Generator(device="cuda")
-    g.manual_seed(1234)
-    return g
-
-
-@pytest.mark.parametrize("hidden_size", [6144, 7168, 4096])
-@pytest.mark.parametrize("rows", list(range(1, 17)))
-def test_bitwise_matches_native_random(rows, hidden_size, gen):
-    weight = random_weight(hidden_size, gen)
-    for scale in (1.0, 37.0):
-        hidden = normed_activation(rows, hidden_size, gen, scale=scale)
-        ok, reason = router_gemm_supported(hidden, weight)
-        if not ok:
-            # Only the staged activation can exceed shared memory (K=7168).
-            assert hidden_size == 7168 and "shared memory" in reason
-            assert router_logits(hidden, weight) is None
-            continue
-        assert_bitwise(router_gemm(hidden, weight), native_logits(hidden, weight))
-
-
-@pytest.mark.parametrize("experts_per_cta", [1, 2, 4, 8])
-@pytest.mark.parametrize("use_pdl", [True, False])
-@pytest.mark.parametrize("prefetch_weights", [True, False])
-@pytest.mark.parametrize("paired_fma", [True, False])
-@pytest.mark.parametrize("rows", [5, 6])
-def test_bitwise_schedule_variants(
-    experts_per_cta, use_pdl, prefetch_weights, paired_fma, rows, gen
-):
-    weight = random_weight(HIDDEN, gen)
-    hidden = normed_activation(rows, HIDDEN, gen)
-    out = router_gemm(
-        hidden,
-        weight,
-        experts_per_cta=experts_per_cta,
-        use_pdl=use_pdl,
-        prefetch_weights=prefetch_weights,
-        paired_fma=paired_fma,
-    )
-    assert_bitwise(out, native_logits(hidden, weight))
-
-
-@pytest.mark.parametrize("layer", [3, 4, 40, 77])
-@pytest.mark.parametrize("rows", [1, 6, 16])
-def test_bitwise_matches_native_real_weights(layer, rows, real_layer_weights, gen):
-    """Real GLM-5.2 router/norm weights; synthetic activations."""
-    w = real_layer_weights(layer)
-    weight = w["gate"].cuda()
-    norm_weight = w["post_norm"].cuda()
-    hidden = normed_activation(rows, HIDDEN, gen, norm_weight=norm_weight)
-    assert_bitwise(router_gemm(hidden, weight), native_logits(hidden, weight))
-
-
-def test_bitwise_special_values(gen):
-    """Zeros, signed zeros, cancellation, wide range and underflowing products."""
-    weight = random_weight(HIDDEN, gen)
-    rows = []
-    rows.append(torch.zeros(HIDDEN, device="cuda"))
-    rows.append(-torch.zeros(HIDDEN, device="cuda"))
-    alt = torch.ones(HIDDEN, device="cuda")
-    alt[1::2] = -1
-    rows.append(alt * 3.0)
-    wide = torch.randn(HIDDEN, generator=gen, device="cuda")
-    wide = wide * torch.exp2(torch.randint(-60, 60, (HIDDEN,), generator=gen, device="cuda"))
-    rows.append(wide)
-    rows.append(torch.randn(HIDDEN, generator=gen, device="cuda") * 1e-30)  # tiny
-    rows.append(torch.randn(HIDDEN, generator=gen, device="cuda") * 1e30)
-    hidden = torch.stack(rows).to(torch.bfloat16)
-    weight_tiny = weight.clone()
-    weight_tiny[:16] = (weight_tiny[:16].float() * 1e-20).to(torch.bfloat16)
-    for w in (weight, weight_tiny):
-        for epc, paired in ((1, True), (1, False), (2, True)):
-            assert_bitwise(
-                router_gemm(hidden, w, experts_per_cta=epc, paired_fma=paired),
-                native_logits(hidden, w),
-            )
-    neg_zero = torch.full((1, HIDDEN), -0.0, dtype=torch.bfloat16, device="cuda")
-    assert_bitwise(router_gemm(neg_zero, -weight.abs()), native_logits(neg_zero, -weight.abs()))
-
-
-def native_route(logits, bias):
-    return torch.ops.trtllm.noaux_tc_op(logits, bias, 1, 1, TOP_K, ROUTED_SCALING)
-
-
-@pytest.mark.parametrize("bias_dtype", [torch.float32, torch.bfloat16])
-def test_routing_decisions_identical_with_ties(bias_dtype, gen):
-    """Exact logit ties (duplicated expert rows) route identically."""
-    weight = random_weight(HIDDEN, gen)
-    weight[17] = weight[200]
-    weight[33] = weight[34]
-    bias = torch.zeros(EXPERTS, dtype=bias_dtype, device="cuda")
-    hidden = normed_activation(16, HIDDEN, gen)
-    ours = router_gemm(hidden, weight)
-    ref = native_logits(hidden, weight)
-    assert_bitwise(ours, ref)
-    assert torch.equal(ours[:, 17].view(torch.int32), ours[:, 200].view(torch.int32))
-    v_ours, i_ours = native_route(ours, bias)
-    v_ref, i_ref = native_route(ref, bias)
-    assert torch.equal(i_ours, i_ref)
-    assert torch.equal(v_ours.view(torch.int32), v_ref.view(torch.int32))
-
-
-def test_routing_decisions_real_weights(real_layer_weights, gen):
-    w = real_layer_weights(40)
-    weight = w["gate"].cuda()
-    bias = w["bias"].cuda().to(torch.bfloat16)
-    norm_weight = w["post_norm"].cuda()
-    for _ in range(8):
-        hidden = normed_activation(6, HIDDEN, gen, norm_weight=norm_weight)
-        ours = router_gemm(hidden, weight)
-        ref = native_logits(hidden, weight)
-        assert_bitwise(ours, ref)
-        assert torch.equal(native_route(ours, bias)[1], native_route(ref, bias)[1])
-
-
-def test_cuda_graph_replay(gen):
-    """Producer -> router inside one graph; replay with changing inputs."""
-    weight = random_weight(HIDDEN, gen)
-    source = normed_activation(6, HIDDEN, gen)
-    hidden = torch.empty_like(source)
-    out_native = torch.empty((6, EXPERTS), dtype=torch.float32, device="cuda")
-    router_gemm(source, weight)  # compile outside capture
-    native_logits(source, weight)
-    graph = torch.cuda.CUDAGraph()
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        with torch.cuda.graph(graph, stream=stream):
-            hidden.copy_(source)
-            out_ours = router_gemm(hidden, weight)
-            out_native.copy_(native_logits(hidden, weight))
-    torch.cuda.current_stream().wait_stream(stream)
-    for step in range(20):
-        source.copy_(normed_activation(6, HIDDEN, gen, scale=1.0 + step))
-        graph.replay()
-        torch.cuda.synchronize()
-        assert_bitwise(out_ours, out_native)
-        assert_bitwise(out_ours, native_logits(source, weight))
-
-
-def test_padded_rows_do_not_affect_live_rows(gen):
-    weight = random_weight(HIDDEN, gen)
-    hidden = normed_activation(6, HIDDEN, gen)
-    live = 2
-    base = router_gemm(hidden, weight).clone()
-    hidden[live:] = normed_activation(6 - live, HIDDEN, gen, scale=1e3)
-    changed = router_gemm(hidden, weight)
-    assert_bitwise(changed[:live], base[:live])
-    hidden[live:] = 0
-    zero_rows = router_gemm(hidden, weight)[live:]
-    assert torch.all(zero_rows == 0)
-    assert_bitwise(router_gemm(hidden, weight), native_logits(hidden, weight))
-
-
-def test_ineligible_calls_return_none(gen):
-    weight = random_weight(HIDDEN, gen)
-    hidden = normed_activation(6, HIDDEN, gen)
-    cases = {
-        "fp16": (hidden.half(), weight.half()),
-        "rows_17": (normed_activation(17, HIDDEN, gen), weight),
-        "rows_0": (hidden[:0], weight),
-        "experts_128": (hidden, weight[:128]),
-        "hidden_5120": (hidden[:, :5120].contiguous(), weight[:, :5120].contiguous()),
-        "non_contiguous": (hidden.t().contiguous().t(), weight),
-        "misaligned": (
-            torch.empty(6 * HIDDEN + 1, dtype=torch.bfloat16, device="cuda")[1:].view(6, HIDDEN),
-            weight,
-        ),
-        "cpu": (hidden.cpu(), weight.cpu()),
-    }
-    before = hop.engagement_counts()["native_fallback"]
-    for name, (h, w) in cases.items():
-        ok, reason = router_gemm_supported(h, w)
-        assert not ok and reason, name
-        assert router_logits(h, w) is None, name
-    assert hop.engagement_counts()["native_fallback"] == before + len(cases)
-    assert router_logits(hidden, weight) is not None
-
-
-def _fake_model_config(tp_size=8, attention_dp=False, layers=78):
-    return SimpleNamespace(
-        mapping=SimpleNamespace(enable_attention_dp=attention_dp, tp_size=tp_size),
-        pretrained_config=SimpleNamespace(num_hidden_layers=layers),
-    )
-
-
-def _gate(use_cute_dsl_bf16_gemm=False):
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate
-
-    gate = DeepseekV3Gate(
-        HIDDEN,
-        EXPERTS,
-        top_k=TOP_K,
-        n_group=1,
-        topk_group=1,
-        routed_scaling_factor=ROUTED_SCALING,
-        dtype=torch.bfloat16,
-        moe_backend="TRTLLM",
-        use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm,
-    )
-    return gate.cuda()
-
-
-def test_make_moe_router_domain(gen):
-    cfg = _fake_model_config()
-    gate = _gate()
-    with torch.no_grad():
-        gate.weight.copy_(random_weight(HIDDEN, gen))
-    route = hop.make_moe_router(gate, 3, cfg)
-    assert route is not None
-    hidden = normed_activation(6, HIDDEN, gen)
-    assert_bitwise(route(hidden), gate(hidden))
-    assert route(hidden.half()) is None
-    assert hop.make_moe_router(gate, 78, cfg) is None  # MTP layer
-    assert hop.make_moe_router(gate, None, cfg) is None
-    assert hop.make_moe_router(gate, 3, _fake_model_config(attention_dp=True)) is None
-    assert hop.make_moe_router(gate, 3, _fake_model_config(tp_size=1)) is None
-    assert hop.make_moe_router(_gate(use_cute_dsl_bf16_gemm=True), 3, cfg) is None
-
-
-def test_registered_entry_points():
-    for name in ("router_logits", "pre_moe_hop", "make_moe_router"):
-        assert callable(dispatch.get_impl("router_gemm", name))
-
-
-def _enclosing_conditions(func, target):
-    """Conjuncts of the ``if`` tests whose body contains ``target``.
-
-    A ``target`` inside an ``else`` branch yields an ``<else>`` entry, so it
-    never matches a required condition.
-    """
-
-    def visit(node, conditions):
-        if node is target:
-            return conditions
-        if isinstance(node, ast.If):
-            test = node.test
-            parts = (
-                test.values
-                if isinstance(test, ast.BoolOp) and isinstance(test.op, ast.And)
-                else [test]
-            )
-            for branch, extra in (
-                (node.body, [ast.unparse(p) for p in parts]),
-                (node.orelse, ["<else>"]),
-            ):
-                for child in branch:
-                    found = visit(child, conditions + extra)
-                    if found is not None:
-                        return found
-            return None
-        for child in ast.iter_child_nodes(node):
-            found = visit(child, conditions)
-            if found is not None:
-                return found
-        return None
-
-    return visit(func, [])
-
-
-def _check_router_call_site(source, class_name):
-    """The call-site semantics of ``<class_name>.compute_routed_output``.
-
-    One router_gemm router call, guarded by the router's conditions, before one native
-    gate call that runs only when no logits exist (the fallback). The call site also
-    takes the post_attention_hop's logits as the ``router_logits`` argument; router_gemm
-    must then not overwrite them.
-    """
-    cls = next(
-        node
-        for node in ast.parse(source).body
-        if isinstance(node, ast.ClassDef) and node.name == class_name
-    )
-    func = next(
-        node
-        for node in cls.body
-        if isinstance(node, ast.FunctionDef) and node.name == "compute_routed_output"
-    )
-    calls = {}
-    for callee in ("self.router_gemm", "self.gate"):
-        found = [
-            node
-            for node in ast.walk(func)
-            if isinstance(node, ast.Call) and ast.unparse(node.func) == callee
-        ]
-        assert len(found) == 1, (callee, len(found))
-        calls[callee] = found[0]
-    router_gemm_call, gate_call = calls["self.router_gemm"], calls["self.gate"]
-    assert router_gemm_call.lineno < gate_call.lineno
-
-    router_guard = {"self.router_gemm is not None", "not use_dp_padding"}
-    router_gemm_guard = set(_enclosing_conditions(func, router_gemm_call))
-    assert router_guard <= router_gemm_guard <= router_guard | {"router_logits is None"}, (
-        router_gemm_guard
-    )
-    if "router_logits" in [arg.arg for arg in func.args.args]:
-        assert "router_logits is None" in router_gemm_guard, router_gemm_guard
-    gate_guard = _enclosing_conditions(func, gate_call)
-    assert gate_guard == ["router_logits is None"], gate_guard
-
-    assigned = {
-        ast.unparse(node.value)
-        for node in ast.walk(func)
-        if isinstance(node, ast.Assign)
-        and any(isinstance(t, ast.Name) and t.id == "router_logits" for t in node.targets)
-    }
-    assert assigned <= {"None", ast.unparse(router_gemm_call), ast.unparse(gate_call)}, assigned
-
-
-def test_router_call_site_in_target():
-    """The target MoE's router call site: one router_gemm selection (in ``GlmMoE.adopt``) and
-    the call-site semantics in ``GlmMoE.compute_routed_output``.
-    """
-    from pathlib import Path
-
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import layers
-
-    source = Path(layers.__file__).read_text()
-    compile(source, "layers.py", "exec")
-    assert source.count('dispatch.is_enabled("router_gemm")') == 1
-    assert source.count("router_logits = self.gate(hidden_states)") == 1
-    _check_router_call_site(source, "GlmMoE")
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py
deleted file mode 100644
index be9453677405..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/conftest.py
+++ /dev/null
@@ -1,100 +0,0 @@
-# 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.
-"""Test setup for the sparse_mla_decode package.
-
-When the installed TensorRT-LLM build does not contain the sparse_mla_decode kernel yet and
-``GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_SRC`` points at a source checkout, the kernel and its torch op
-are compiled from that checkout so that the tests can run against an existing
-image. Builds that compile the sparse_mla_decode sources into ``th_common`` skip this step.
-"""
-
-import os
-import sys
-import traceback
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parent))
-
-
-def _load_dev_op() -> None:
-    source_root = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_SRC")
-    if not source_root:
-        return
-    import torch
-
-    import tensorrt_llm
-
-    if hasattr(torch.ops.trtllm, "glm_moe_dsa_sparse_mla_decode"):
-        return
-    from torch.utils.cpp_extension import load
-
-    timing = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING") == "1"
-    # Development variants: GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_DEFINES="NAME_A+NAME_B" adds -DNAME_A
-    # -DNAME_B.
-    extra = [
-        d for d in os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_DEFINES", "").split("+") if d
-    ]
-    major, minor = torch.cuda.get_device_capability()
-    # GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_ARCH selects an architecture-specific target such as "103a".
-    arch_name = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_ARCH", f"{major}{minor}")
-    variant = "_".join([arch_name] + (["timing"] if timing else []) + [d.lower() for d in extra])
-    cpp = Path(source_root) / "cpp"
-    libs = Path(tensorrt_llm.__file__).resolve().parent / "libs"
-    build_dir = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_DEV_BUILD_DIR")
-    if build_dir:
-        build_dir = os.path.join(build_dir, variant)
-        os.makedirs(build_dir, exist_ok=True)
-    defines = ["-DENABLE_BF16", "-DENABLE_FP8"] + (
-        ["-DGLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING"] if timing else []
-    )
-    defines += [f"-D{d}" for d in extra]
-    # torch's extension defaults disable the half/bfloat16 operators that the
-    # TensorRT-LLM CUDA headers use; the TensorRT-LLM build does not set them.
-    undefines = [
-        "-U__CUDA_NO_HALF_OPERATORS__",
-        "-U__CUDA_NO_HALF_CONVERSIONS__",
-        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
-        "-U__CUDA_NO_HALF2_OPERATORS__",
-    ]
-    arch = f"-gencode=arch=compute_{arch_name},code=sm_{arch_name}"
-    load(
-        name=f"glm_moe_dsa_sparse_mla_decode_dev_{variant}",
-        sources=[
-            str(cpp / "tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu"),
-            str(cpp / "tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp"),
-        ],
-        extra_include_paths=[str(cpp / "include"), str(cpp)],
-        extra_cflags=["-O3", "-std=c++17", *defines],
-        extra_cuda_cflags=[
-            "-O3",
-            "-std=c++17",
-            "-lineinfo",
-            "--expt-relaxed-constexpr",
-            arch,
-            *defines,
-            *undefines,
-        ],
-        extra_ldflags=[f"-L{libs}", "-ltensorrt_llm", f"-Wl,-rpath,{libs}"],
-        build_directory=build_dir,
-        is_python_module=False,
-        verbose=True,
-    )
-
-
-try:
-    _load_dev_op()
-except RuntimeError:
-    # Keep the native-only tests runnable; the sparse_mla_decode tests skip without the op.
-    traceback.print_exc()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py
deleted file mode 100644
index 42103593ca4f..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_case.py
+++ /dev/null
@@ -1,397 +0,0 @@
-# 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.
-"""One DSA ``MLA`` module through the target's call sites, for composition tests of sparse_mla_decode.
-
-Run as a script in a fresh interpreter: ``python sparse_mla_case.py <output.pt>`` with the fast
-paths to run in ``GLM_MOE_DSA_TEST_FAST_PATHS`` (``plan_testlib.FAST_PATHS_ENV``). It builds
-one GLM-shaped DSA MLA layer (8 local heads, TP1, FP8 latent cache, deterministic weights keyed
-by parameter name), adopts it as the target's ``GlmMLA``,
-runs a native context step and then MTP verify steps of 6 rows as decode steps through
-``GlmMLA.forward``, so that every package at the MLA / DSA call sites
-(absorbed_query/shared_indexer_absorbed_query
-projections, rope_kv_append RoPE/append and top-k bridge, attention_output_tail, sparse_mla_decode
-attention) runs as in the model. With ``GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK=fixed`` (default) the
-indexer's generation selection is replaced by a given local top-k, so that every process attends to
-the same rows; with ``GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK=indexer`` the indexer runs unmodified
-(needed when rope_kv_append's top-k bridge is a producer).
-
-Saved: per step the FMHA inputs (FP8 query bytes, bmm scales, global indices,
-KV lengths), the latent attention output and the module output, for an eager
-step, a step whose sparse_mla_decode launch is followed by a device synchronization (no
-programmatic overlap with the next kernel), and replays of a captured CUDA
-graph over alternating inputs (with the global indices each replay used);
-plus the per-package capability decisions.
-"""
-
-import json
-import math
-import os
-import sys
-import zlib
-from pathlib import Path
-from types import SimpleNamespace
-from unittest.mock import Mock
-
-sys.path.insert(0, str(Path(__file__).resolve().parent))
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-
-import plan_testlib  # noqa: E402
-import torch  # noqa: E402
-
-import tensorrt_llm  # noqa: E402
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E402, E501
-    GlmMLA,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E402, E501
-    SparseMlaDecodeFmha,
-)
-from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha  # noqa: E402
-from tensorrt_llm._torch.attention.backends.interface import (  # noqa: E402
-    PositionalEmbeddingParams,
-    RopeParams,
-)
-from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager  # noqa: E402
-from tensorrt_llm._torch.attention.mla import MLA  # noqa: E402
-from tensorrt_llm._torch.metadata import KVCacheParams  # noqa: E402
-from tensorrt_llm._torch.model_config import ModelConfig  # noqa: E402
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention  # noqa: E402
-from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str  # noqa: E402
-from tensorrt_llm.bindings.executor import KvCacheConfig  # noqa: E402
-from tensorrt_llm.functional import PositionEmbeddingType  # noqa: E402
-from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig  # noqa: E402
-from tensorrt_llm.mapping import Mapping  # noqa: E402
-from tensorrt_llm.models.modeling_utils import QuantConfig  # noqa: E402
-from tensorrt_llm.quantization.mode import QuantAlgo  # noqa: E402
-
-HIDDEN = 6144
-HEADS = 8
-Q_LORA = 2048
-KV_LORA = 512
-NOPE = 192
-ROPE = 64
-V_HEAD = 256
-TOP_K = 2048
-ROWS = 6
-TOKENS_PER_BLOCK = 64
-
-
-def _rope_config(max_positions):
-    return SimpleNamespace(
-        max_position_embeddings=max_positions,
-        hidden_size=HIDDEN,
-        num_attention_heads=64,
-        qk_rope_head_dim=ROPE,
-        rope_theta=8.0e6,
-        rope_scaling=None,
-        model_type="deepseek_v32",
-    )
-
-
-def _fill(param: torch.Tensor, name: str) -> None:
-    """Deterministic values keyed by the parameter name (independent of creation order)."""
-    gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(name.encode()))
-    if param.dtype not in (torch.bfloat16, torch.float16, torch.float32):
-        return
-    if param.dim() == 1:
-        values = 1.0 + 0.05 * torch.randn(param.shape, device="cuda", generator=gen)
-    else:
-        fan_in = param.shape[-1]
-        values = torch.randn(param.shape, device="cuda", generator=gen) / math.sqrt(fan_in)
-    with torch.no_grad():
-        param.copy_(values.to(param.dtype))
-
-
-def build(ctx_len: int):
-    max_seq_len = ctx_len + 4 * ROWS
-    max_positions = 1 << max(16, math.ceil(math.log2(max_seq_len + 1)))
-    mapping = Mapping(world_size=1, tp_size=1, rank=0)
-    sparse_config = DeepSeekSparseAttentionConfig(
-        index_n_heads=32, index_head_dim=128, index_topk=TOP_K, skip_indexer_for_short_seqs=False
-    )
-    model_config = ModelConfig(
-        mapping=mapping,
-        sparse_attention_config=sparse_config,
-        pretrained_config=SimpleNamespace(rms_norm_eps=1e-5),
-        quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value),
-    )
-    pos_embd_params = PositionalEmbeddingParams(
-        type=PositionEmbeddingType.yarn,
-        rope=RopeParams.from_config(_rope_config(max_positions)),
-        is_neox=False,
-    )
-    mla = MLA(
-        hidden_size=HIDDEN,
-        num_attention_heads=HEADS,
-        # GLM-5.2: num_key_value_heads == num_attention_heads (sizes o_proj's input).
-        num_key_value_heads=HEADS,
-        qk_nope_head_dim=NOPE,
-        qk_rope_head_dim=ROPE,
-        v_head_dim=V_HEAD,
-        q_lora_rank=Q_LORA,
-        kv_lora_rank=KV_LORA,
-        predicted_tokens_per_seq=1,
-        max_position_embeddings=max_positions,
-        bias=False,
-        pos_embd_params=pos_embd_params,
-        layer_idx=0,
-        dtype=torch.bfloat16,
-        config=model_config,
-        reduce_output=False,
-        rms_norm_eps=1e-5,
-    )
-    # The model's DeepseekV32Attention adds only its constructor to MLA; the case builds the MLA
-    # with GLM shapes directly and adopts it as the model's attention.
-    mla.__class__ = DeepseekV32Attention
-    GlmMLA.adopt(mla)
-    # The model builds the FMHA libraries in ``create_weights`` after adoption; a module built
-    # alone rebuilds them for the library adoption declared.
-    for backend in (mla.mha, mla.mqa):
-        if backend is not None and backend.prepended_fmha_lib_classes:
-            backend.update_quant_config(backend.quant_config)
-    mla = mla.cuda()
-    # Eager module path (the custom-op path resolves layers through the model engine).
-    mla.register_to_config = False
-    indexer = mla.mqa.indexer
-    indexer.cuda()
-    for name, param in list(mla.named_parameters()) + [
-        ("indexer." + n, p) for n, p in indexer.named_parameters()
-    ]:
-        _fill(param, name)
-    with torch.no_grad():
-        kv_b = mla.kv_b_proj.weight.view(HEADS, NOPE + V_HEAD, KV_LORA)
-        mla.kv_b_proj.weight.copy_(
-            torch.cat(
-                [kv_b[:, :NOPE].reshape(-1, KV_LORA), kv_b[:, NOPE:].reshape(-1, KV_LORA)], dim=0
-            )
-        )
-        k_nope = mla.kv_b_proj.weight[: HEADS * NOPE].view(HEADS, NOPE, KV_LORA)
-        v = mla.kv_b_proj.weight[HEADS * NOPE :].view(HEADS, V_HEAD, KV_LORA)
-        # In place, so that any packed view a package took at construction stays valid.
-        mla.k_b_proj_trans.copy_(k_nope.transpose(1, 2))
-        mla.v_b_proj.copy_(v)
-    indexer.post_load_weights()
-    mla.post_load_weights()
-
-    blocks = (max_seq_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
-    cache = DSACacheManager(
-        KvCacheConfig(max_tokens=blocks * TOKENS_PER_BLOCK, enable_block_reuse=False),
-        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-        num_layers=1,
-        num_kv_heads=1,
-        head_dim=KV_LORA + ROPE,
-        tokens_per_block=TOKENS_PER_BLOCK,
-        max_seq_len=max_seq_len,
-        max_batch_size=1,
-        mapping=mapping,
-        dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)),
-        sparse_attn_config=sparse_config,
-        model_config=model_config,
-    )
-    cache.add_dummy_requests([0], [ctx_len])
-    return mla, cache, sparse_config, mapping
-
-
-def metadata(mla, cache, sparse_config, mapping, ctx_len, cached, rows, generation):
-    meta = type(mla.mqa).Metadata(
-        seq_lens=torch.tensor([rows], dtype=torch.int),
-        request_ids=[0],
-        max_num_requests=1,
-        num_contexts=0 if generation else 1,
-        prompt_lens=[ctx_len],
-        max_num_tokens=max(rows, ROWS),
-        kv_cache_manager=cache,
-        kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[cached]),
-        mapping=mapping,
-        sparse_attention_config=sparse_config,
-    )
-    if generation:
-        meta.update_spec_dec_param(
-            batch_size=1,
-            is_spec_decoding_enabled=True,
-            is_spec_dec_tree=False,
-            is_spec_dec_dynamic_tree=False,
-            max_draft_len=ROWS - 1,
-            max_total_draft_tokens=ROWS - 1,
-        )
-    meta.prepare()
-    return meta
-
-
-def local_topk(ctx_len: int, seed: int) -> torch.Tensor:
-    gen = torch.Generator(device="cpu").manual_seed(seed)
-    topk = torch.full((ROWS, TOP_K), -1, dtype=torch.int32)
-    for row in range(ROWS):
-        visible = ctx_len + row + 1
-        recent = torch.arange(visible - 32, visible)
-        rest = torch.randperm(visible - 32, generator=gen)[: TOP_K - 32]
-        topk[row] = torch.cat([recent, rest]).to(torch.int32)
-    return topk.cuda()
-
-
-class Capture:
-    """Records the FMHA inputs and the latent output of each attention call."""
-
-    def __init__(self):
-        self.calls = []
-        self.record_inputs = True
-        self.sync_after_sparse_decode = False
-        self._originals = {}
-
-    def install(self):
-        for cls in (SparseMlaDecodeFmha, FallbackFmha):
-            original = cls.forward
-            self._originals[cls] = original
-
-            def forward(fmha, q, k, v, md, args, _original=original, _cls=cls):
-                entry = {"library": _cls.__name__, "output": args.output}
-                if self.record_inputs:
-                    sparse = args.sparse_runtime_params
-                    entry.update(
-                        quant_q=args.quant_q_buffer.clone(),
-                        bmm1=args.mla_bmm1_scale.clone(),
-                        bmm2=args.mla_bmm2_scale.clone(),
-                        indices=sparse.sparse_attn_indices.clone(),
-                        kv_lens=md.kv_lens_cuda_runtime[md.num_contexts :].clone(),
-                    )
-                _original(fmha, q, k, v, md, args)
-                if self.sync_after_sparse_decode and _cls is SparseMlaDecodeFmha:
-                    torch.cuda.synchronize()
-                self.calls.append(entry)
-
-            cls.forward = forward
-
-
-def main(out_path: str) -> None:
-    plan_testlib.select_from_env(dispatch)
-    ctx_len = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_CTX", "32768"))
-    replays = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_REPLAYS", "8"))
-    topk_mode = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK", "fixed")
-    assert topk_mode in ("fixed", "indexer"), topk_mode
-    torch.manual_seed(0)
-    mla, cache, sparse_config, mapping = build(ctx_len)
-    gen = torch.Generator(device="cuda").manual_seed(1234)
-    hidden_ctx = torch.randn(ctx_len, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16)
-    hidden_gen = [
-        torch.randn(ROWS, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16) for _ in range(2)
-    ]
-    topks = [local_topk(ctx_len, 77), local_topk(ctx_len, 78)]
-    pool = cache.get_unique_primary_pool()
-
-    ctx_meta = metadata(mla, cache, sparse_config, mapping, ctx_len, 0, ctx_len, False)
-    with torch.no_grad():
-        mla.forward(torch.arange(ctx_len, device="cuda")[None], hidden_ctx, ctx_meta)
-    torch.cuda.synchronize()
-    for _ in range(ROWS):
-        cache.impl.add_token(0)
-    gen_meta = metadata(mla, cache, sparse_config, mapping, ctx_len, ctx_len, ROWS, True)
-    positions = torch.arange(ctx_len, ctx_len + ROWS, device="cuda")[None]
-
-    capture = Capture()
-    capture.install()
-    static_hidden = hidden_gen[0].clone()
-    static_topk = topks[0].clone()
-    if topk_mode == "fixed":
-        mla.mqa.indexer.forward_from_projected = Mock(return_value=static_topk)
-
-    def step(index):
-        static_hidden.copy_(hidden_gen[index])
-        static_topk.copy_(topks[index])
-        with torch.no_grad(), dispatch.step(dispatch.Phase.DECODE):
-            return mla.forward(positions, static_hidden, gen_meta)
-
-    result = {"eager": [], "serialized": [], "graph": []}
-    # Eager steps (inputs recorded), each input set once.
-    for index in (0, 1):
-        out = step(index).clone()
-        torch.cuda.synchronize()
-        call = capture.calls[-1]
-        result["eager"].append(
-            {
-                "input": index,
-                "module_output": out.cpu(),
-                "latent": call["output"].clone().cpu(),
-                "library": call["library"],
-                "quant_q": call["quant_q"].cpu(),
-                "bmm1": call["bmm1"].cpu(),
-                "bmm2": call["bmm2"].cpu(),
-                "indices": call["indices"].cpu(),
-                "kv_lens": call["kv_lens"].cpu(),
-            }
-        )
-    # The attended pool rows of each input, for an FP64 oracle in the parent.
-    pool_rows = pool.view(torch.uint8).reshape(-1, KV_LORA + ROPE)
-    for entry in result["eager"]:
-        idx = entry["indices"].long().clamp_min(0).cuda()
-        entry["rows"] = pool_rows[idx.reshape(-1)].reshape(ROWS, TOP_K, -1).cpu()
-    # Same inputs with the attention kernel completed before the next kernel is launched.
-    capture.sync_after_sparse_decode = True
-    for index in (0, 1):
-        out = step(index).clone()
-        torch.cuda.synchronize()
-        result["serialized"].append(
-            {
-                "input": index,
-                "module_output": out.cpu(),
-                "latent": capture.calls[-1]["output"].clone().cpu(),
-                "indices": capture.calls[-1]["indices"].cpu(),
-            }
-        )
-    capture.sync_after_sparse_decode = False
-    # A captured CUDA graph of one verify step, replayed over alternating inputs. The
-    # recorded indices are clones made inside the graph, refreshed by every replay.
-    step(0)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        with torch.no_grad(), dispatch.step(dispatch.Phase.DECODE):
-            graph_out = mla.forward(positions, static_hidden, gen_meta)
-    graph_call = capture.calls[-1]
-    for it in range(replays):
-        index = it % 2
-        static_hidden.copy_(hidden_gen[index])
-        static_topk.copy_(topks[index])
-        graph.replay()
-        torch.cuda.synchronize()
-        result["graph"].append(
-            {
-                "input": index,
-                "module_output": graph_out.clone().cpu(),
-                "latent": graph_call["output"].clone().cpu(),
-                "indices": graph_call["indices"].cpu(),
-            }
-        )
-    result["capability"] = {
-        f"{feature}/{key}": value for (feature, key), value in dispatch._CAPABILITY_LOGGED.items()
-    }
-    result["features"] = sorted(dispatch.enabled_features())
-    result["topk_mode"] = topk_mode
-    result["libraries"] = sorted({c["library"] for c in capture.calls})
-    torch.save(result, out_path)
-    print(
-        json.dumps(
-            {
-                "features": result["features"],
-                "topk_mode": topk_mode,
-                "libraries": result["libraries"],
-                "capability": result["capability"],
-            }
-        )
-    )
-    cache.shutdown()
-
-
-if __name__ == "__main__":
-    main(sys.argv[1])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py
deleted file mode 100644
index 9aa3c4ad3a6d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_harness.py
+++ /dev/null
@@ -1,1026 +0,0 @@
-# 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.
-"""Native DSA generation harness and FP64 oracles for the sparse_mla_decode tests.
-
-The harness drives the ``DSATrtllmAttention`` generation path exactly as
-``forward_absorption_generation`` does: ``mla_rope_generation`` (RoPE, FP8 query,
-latent append, device scales) followed by ``TrtllmAttention.forward``. Each layer and its
-indexer are adopted as the target's ``GlmDsaBackend`` / ``GlmIndexer`` and every generation
-forward is a decode step, so the predictor runs the target's (with no fast path selected, the
-built-in one). The
-indexer is either replaced by a fixed local top-k (``forward``) or runs
-unmodified on given indexer projections (``forward_native_indexer``). The FMHA
-library list of the layer decides whether the native library or sparse_mla_decode computes
-the attention.
-"""
-
-import contextlib
-import json
-import math
-import os
-import tempfile
-from dataclasses import dataclass
-from typing import Callable, Dict, Iterator, List, Optional, Sequence, Tuple
-from unittest.mock import Mock
-
-import torch
-
-import tensorrt_llm
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.attention import (  # noqa: E501
-    GlmDsaBackend,
-    GlmIndexer,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
-    SparseMlaDecodeFmha,
-)
-from tensorrt_llm._torch.attention.backends.interface import (
-    AttentionInputType,
-    MLAParams,
-    PositionalEmbeddingParams,
-    RopeParams,
-)
-from tensorrt_llm._torch.attention.backends.sparse.dsa import DSABackendForwardArgs, DSACacheManager
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import (
-    transform_local_topk_and_prepare_pool_view,
-)
-from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
-from tensorrt_llm._torch.metadata import KVCacheParams
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
-from tensorrt_llm.bindings.executor import KvCacheConfig
-from tensorrt_llm.functional import PositionEmbeddingType
-from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig
-from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.models.modeling_utils import QuantConfig
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-NUM_HEADS = 8
-KV_LORA_RANK = 512
-ROPE_DIM = 64
-NOPE_DIM = 192
-V_HEAD_DIM = 256
-HEAD_DIM = KV_LORA_RANK + ROPE_DIM
-Q_LORA_RANK = 2048
-HIDDEN_SIZE = 6144
-TOP_K = 2048
-TOKENS_PER_BLOCK = 64
-E4M3_MAX = 448.0
-
-
-@dataclass
-class _GlmRopeConfig:
-    """Rotary settings of GLM-5.2 (``rope_theta`` 8e6, no scaling)."""
-
-    max_position_embeddings: int
-    hidden_size: int = 6144
-    num_attention_heads: int = 64
-    qk_rope_head_dim: int = ROPE_DIM
-    rope_theta: float = 8.0e6
-    rope_scaling: Optional[dict] = None
-    model_type: str = "deepseek_v32"
-
-
-@dataclass
-class StepBuffers:
-    """Per-step tensors produced by the native ``mla_rope_generation``."""
-
-    fused_q: torch.Tensor
-    q_pe: torch.Tensor
-    latent: torch.Tensor
-    cu_q_seqlens: torch.Tensor
-    cu_kv_seqlens: torch.Tensor
-    scheduler_counter: torch.Tensor
-    bmm1_scale: torch.Tensor
-    bmm2_scale: torch.Tensor
-    quant_q: torch.Tensor
-
-
-class DsaGenHarness:
-    """One generation step of native DSA layers over an FP8 latent pool.
-
-    Args:
-        ctx_lens: Cached tokens per sequence before this step.
-        seq_len_q: Query rows per sequence (1 + draft tokens).
-        num_layers: Attention layers sharing the pool.
-        index_topk: DSA top-k of the sparse config.
-        kv_dtype: KV cache dtype (FP8 for the sparse_mla_decode domain).
-        seed: Seed of the random pool contents.
-        latent_std: Standard deviation of the random latent rows.
-        fill_random: Fill the whole pool with random rows (instead of running a context step).
-        defer_generation: Leave room for ``run_context`` before ``start_generation``.
-        sparse_overrides: Extra ``DeepSeekSparseAttentionConfig`` fields (indexer kernels).
-        mtp_draft_len: Draft length announced to the generation metadata, as the
-            speculative-decoding runner does (``update_spec_dec_param``); 0 leaves it unset.
-        ragged_verify_lens: Query rows per sequence of a ragged-verification step (one
-            entry per sequence; replaces ``seq_len_q``). The generation metadata then
-            enables ragged verification, and ``TrtllmAttention`` presents one row per
-            query token to the RoPE and FMHA calls.
-    """
-
-    def __init__(
-        self,
-        ctx_lens: Sequence[int],
-        seq_len_q: int,
-        num_layers: int = 1,
-        index_topk: int = TOP_K,
-        kv_dtype: torch.dtype = torch.float8_e4m3fn,
-        seed: int = 0,
-        latent_std: float = 1.0,
-        fill_random: bool = True,
-        defer_generation: bool = False,
-        sparse_overrides: Optional[dict] = None,
-        mtp_draft_len: int = 0,
-        ragged_verify_lens: Optional[Sequence[int]] = None,
-    ):
-        self.ctx_lens = [int(x) for x in ctx_lens]
-        self.num_seqs = len(self.ctx_lens)
-        self.ragged = ragged_verify_lens is not None
-        if self.ragged:
-            self.verify_lens = [int(x) for x in ragged_verify_lens]
-            assert len(self.verify_lens) == self.num_seqs and min(self.verify_lens) >= 1
-            self.seq_len_q = max(self.verify_lens)
-        else:
-            self.seq_len_q = int(seq_len_q)
-            self.verify_lens = [self.seq_len_q] * self.num_seqs
-        self.num_rows = sum(self.verify_lens)
-        # Sequence and position within the sequence's query rows, per generation row.
-        self.row_seq = [s for s, n in enumerate(self.verify_lens) for _ in range(n)]
-        self.row_pos = [p for n in self.verify_lens for p in range(n)]
-        self.device = torch.device("cuda")
-        self.mtp_draft_len = int(mtp_draft_len)
-        self.sparse_config = DeepSeekSparseAttentionConfig(
-            index_n_heads=32,
-            index_head_dim=128,
-            index_topk=index_topk,
-            skip_indexer_for_short_seqs=False,
-            **(sparse_overrides or {}),
-        )
-        self.attn_cls = get_attention_backend("TRTLLM", self.sparse_config)
-        max_seq_len = max(self.ctx_lens) + 2 * self.seq_len_q
-        max_positions = 1 << max(16, math.ceil(math.log2(max_seq_len + 1)))
-        rope_config = _GlmRopeConfig(max_position_embeddings=max_positions)
-        pos_embd_params = PositionalEmbeddingParams(
-            type=PositionEmbeddingType.yarn, rope=RopeParams.from_config(rope_config), is_neox=False
-        )
-        mla_params = MLAParams(
-            q_lora_rank=Q_LORA_RANK,
-            kv_lora_rank=KV_LORA_RANK,
-            qk_rope_head_dim=ROPE_DIM,
-            qk_nope_head_dim=NOPE_DIM,
-            v_head_dim=V_HEAD_DIM,
-            rope_append=True,
-            predicted_tokens_per_seq=1,
-            hidden_size=HIDDEN_SIZE,
-        )
-        quant_config = (
-            QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value)
-            if kv_dtype == torch.float8_e4m3fn
-            else None
-        )
-        self.layers = [
-            self.attn_cls(
-                layer_idx=i,
-                num_heads=NUM_HEADS,
-                head_dim=HEAD_DIM,
-                num_kv_heads=1,
-                quant_config=quant_config,
-                q_scaling=1.0,
-                pos_embd_params=pos_embd_params,
-                mla_params=mla_params,
-                sparse_attention_config=self.sparse_config,
-                dtype=torch.bfloat16,
-            )
-            for i in range(num_layers)
-        ]
-        self.mapping = Mapping(world_size=1, tp_size=1, rank=0)
-        model_config = ModelConfig(
-            mapping=self.mapping,
-            sparse_attention_config=self.sparse_config,
-            pretrained_config=_PretrainedStub(),
-        )
-        blocks_per_seq = (max_seq_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
-        self.kv_cache_manager = DSACacheManager(
-            KvCacheConfig(
-                max_tokens=blocks_per_seq * TOKENS_PER_BLOCK * self.num_seqs,
-                enable_block_reuse=False,
-            ),
-            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
-            num_layers=num_layers,
-            num_kv_heads=1,
-            head_dim=HEAD_DIM,
-            tokens_per_block=TOKENS_PER_BLOCK,
-            max_seq_len=max_seq_len,
-            max_batch_size=self.num_seqs,
-            mapping=self.mapping,
-            dtype=str_dtype_to_binding(torch_dtype_to_str(kv_dtype)),
-            sparse_attn_config=self.sparse_config,
-            model_config=model_config,
-        )
-        for layer in self.layers:
-            GlmDsaBackend.adopt(layer)
-            if layer.indexer is not None:
-                GlmIndexer.adopt(layer.indexer)
-        self.request_ids = list(range(self.num_seqs))
-        self.kv_cache_manager.add_dummy_requests(self.request_ids, self.ctx_lens)
-        if fill_random:
-            self.fill_pool(seed, latent_std)
-        self.metadata = None
-        if not defer_generation:
-            self.start_generation()
-
-    def run_context(
-        self,
-        layer_idx: int,
-        fused_q: torch.Tensor,
-        q_pe: torch.Tensor,
-        latent: torch.Tensor,
-        topk: Optional[torch.Tensor],
-        indexer_k: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        chunk_tokens: Optional[int] = None,
-        attention: bool = True,
-    ) -> None:
-        """Native context step: RoPE and append the context latents of every sequence.
-
-        ``topk`` replaces the context-phase indexer selection (the context output is not
-        used). ``indexer_k`` = ``(k_fp8, k_scale)`` of the context tokens, when given, is
-        written to the indexer K cache with the native ``_update_k_cache``, as
-        ``_forward_dsa_attn`` does before the attention. With ``chunk_tokens`` (one
-        sequence only) the context runs as chunked prefill: one context step per chunk,
-        each with the previous chunks as cached tokens. With ``attention=False`` only the
-        indexer K cache is written (the latent pool keeps its contents).
-        """
-        if chunk_tokens is None or chunk_tokens >= max(self.ctx_lens):
-            chunks = [(None, None)]
-        else:
-            assert self.num_seqs == 1, "chunked context steps support one sequence"
-            ctx = self.ctx_lens[0]
-            chunks = [(s, min(s + chunk_tokens, ctx)) for s in range(0, ctx, chunk_tokens)]
-        layer = self.layers[layer_idx]
-        for start, end in chunks:
-            if start is None:
-                seq_lens, cached, rows = self.ctx_lens, [0] * self.num_seqs, slice(None)
-            else:
-                seq_lens, cached, rows = [end - start], [start], slice(start, end)
-            ctx_metadata = self.attn_cls.Metadata(
-                seq_lens=torch.tensor(seq_lens, dtype=torch.int),
-                request_ids=self.request_ids,
-                max_num_requests=self.num_seqs,
-                num_contexts=self.num_seqs,
-                prompt_lens=self.ctx_lens,
-                max_num_tokens=sum(seq_lens),
-                kv_cache_manager=self.kv_cache_manager,
-                kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached),
-                mapping=self.mapping,
-                sparse_attention_config=self.sparse_config,
-            )
-            ctx_metadata.prepare()
-            if indexer_k is not None:
-                layer.indexer._update_k_cache(indexer_k[0][rows], indexer_k[1][rows], ctx_metadata)
-            if not attention:
-                torch.cuda.synchronize()
-                continue
-            layer.indexer.forward_from_projected = Mock(return_value=topk[rows])
-            layer.forward(
-                fused_q[rows].clone(),
-                None,
-                None,
-                ctx_metadata,
-                attention_input_type=AttentionInputType.context_only,
-                latent_cache=latent[rows],
-                q_pe=q_pe[rows],
-                sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]),
-            )
-            torch.cuda.synchronize()
-
-    def start_generation(self) -> None:
-        """Allocate this step's rows and build the generation metadata."""
-        for request_id, rows in zip(self.request_ids, self.verify_lens):
-            for _ in range(rows):
-                self.kv_cache_manager.impl.add_token(request_id)
-        # The ragged row buffers are allocated at construction; the per-step widths are
-        # installed before ``prepare`` (they are not a constructor argument).
-        ragged_kwargs = dict(enable_ragged_verification=True) if self.ragged else {}
-        self.metadata = self.attn_cls.Metadata(
-            seq_lens=torch.tensor(self.verify_lens, dtype=torch.int),
-            request_ids=self.request_ids,
-            max_num_requests=self.num_seqs,
-            num_contexts=0,
-            prompt_lens=self.ctx_lens,
-            max_num_tokens=self.num_rows,
-            kv_cache_manager=self.kv_cache_manager,
-            kv_cache_params=KVCacheParams(
-                use_cache=True, num_cached_tokens_per_seq=list(self.ctx_lens)
-            ),
-            mapping=self.mapping,
-            sparse_attention_config=self.sparse_config,
-            **ragged_kwargs,
-        )
-        if self.ragged:
-            # Ragged verification sizes its row buffers from the draft ceiling and needs the
-            # dense spec-decoding metadata disabled.
-            self.metadata.update_spec_dec_param(
-                batch_size=self.num_seqs,
-                is_spec_decoding_enabled=False,
-                is_spec_dec_tree=False,
-                is_spec_dec_dynamic_tree=False,
-                max_draft_len=self.seq_len_q - 1,
-                max_total_draft_tokens=self.seq_len_q - 1,
-            )
-            self.metadata.ragged_verify_lens = list(self.verify_lens)
-        elif self.mtp_draft_len > 0:
-            self.metadata.update_spec_dec_param(
-                batch_size=self.num_seqs,
-                is_spec_decoding_enabled=True,
-                is_spec_dec_tree=False,
-                is_spec_dec_dynamic_tree=False,
-                max_draft_len=self.mtp_draft_len,
-                max_total_draft_tokens=self.mtp_draft_len,
-            )
-        self.metadata.prepare()
-
-    def shutdown(self) -> None:
-        self.kv_cache_manager.shutdown()
-
-    @property
-    def cu_len(self) -> int:
-        """Entries of the cumulative-length buffers: one per presented row plus one.
-
-        ``forward_absorption_generation`` sizes them the same way: per sequence, and
-        per query token under the ragged-verification presentation.
-        """
-        return (self.num_rows if self.ragged else self.num_seqs) + 1
-
-    # ----- pool access -----------------------------------------------------------------
-    def pool(self) -> torch.Tensor:
-        """Whole primary pool as an e4m3 tensor."""
-        pool = self.kv_cache_manager.get_unique_primary_pool()
-        if pool.dtype != torch.float8_e4m3fn and pool.element_size() == 1:
-            pool = pool.view(torch.float8_e4m3fn)
-        return pool
-
-    def pool_rows(self) -> torch.Tensor:
-        """``[num_pool_rows, 576]`` view addressed by global sparse indices."""
-        return self.pool().reshape(-1, HEAD_DIM)
-
-    def pool_base_pointer(self, layer_idx: int = 0) -> int:
-        from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
-            pool_base_pointer,
-        )
-
-        return pool_base_pointer(
-            self.metadata, self.layers[layer_idx].get_local_layer_idx(self.metadata)
-        )
-
-    def fill_pool(self, seed: int, latent_std: float) -> None:
-        """Fill every pool row with random finite e4m3 latent/rope values."""
-        pool = self.pool()
-        gen = torch.Generator(device=self.device).manual_seed(seed)
-        flat = pool.view(-1)
-        chunk = 1 << 26
-        for start in range(0, flat.numel(), chunk):
-            end = min(start + chunk, flat.numel())
-            values = torch.randn(end - start, device=self.device, generator=gen) * latent_std
-            flat[start:end] = values.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn)
-
-    # ----- one step ---------------------------------------------------------------------
-    def make_step(self, layer_idx: int, seed: int, q_std: float = 1.0) -> StepBuffers:
-        """Random absorbed query and new latents, then the native ``mla_rope_generation``."""
-        gen = torch.Generator(device=self.device).manual_seed(seed)
-        rows = self.num_rows
-        fused_q = (
-            torch.randn(rows, NUM_HEADS * HEAD_DIM, device=self.device, generator=gen) * q_std
-        ).to(torch.bfloat16)
-        q_pe = (
-            torch.randn(rows, NUM_HEADS, ROPE_DIM, device=self.device, generator=gen) * q_std
-        ).to(torch.bfloat16)
-        latent = torch.randn(rows, HEAD_DIM, device=self.device, generator=gen).to(torch.bfloat16)
-        step = StepBuffers(
-            fused_q=fused_q,
-            q_pe=q_pe,
-            latent=latent,
-            cu_q_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
-            cu_kv_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
-            scheduler_counter=torch.empty(1, dtype=torch.uint32, device=self.device),
-            bmm1_scale=torch.empty(2, dtype=torch.float32, device=self.device),
-            bmm2_scale=torch.empty(1, dtype=torch.float32, device=self.device),
-            quant_q=torch.empty(rows, NUM_HEADS * HEAD_DIM, dtype=torch.uint8, device=self.device),
-        )
-        self.rope_and_quantize(layer_idx, step)
-        return step
-
-    def make_step_from(
-        self, layer_idx: int, fused_q: torch.Tensor, q_pe: torch.Tensor, latent: torch.Tensor
-    ) -> StepBuffers:
-        """Step buffers from given absorbed queries and latents, then ``mla_rope_generation``."""
-        step = StepBuffers(
-            fused_q=fused_q.clone(),
-            q_pe=q_pe.clone(),
-            latent=latent.clone(),
-            cu_q_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
-            cu_kv_seqlens=torch.empty(self.cu_len, dtype=torch.int32, device=self.device),
-            scheduler_counter=torch.empty(1, dtype=torch.uint32, device=self.device),
-            bmm1_scale=torch.empty(2, dtype=torch.float32, device=self.device),
-            bmm2_scale=torch.empty(1, dtype=torch.float32, device=self.device),
-            quant_q=torch.empty(
-                self.num_rows, NUM_HEADS * HEAD_DIM, dtype=torch.uint8, device=self.device
-            ),
-        )
-        self.rope_and_quantize(layer_idx, step)
-        return step
-
-    def rope_and_quantize(self, layer_idx: int, step: StepBuffers) -> None:
-        self.layers[layer_idx].mla_rope_generation(
-            step.fused_q,
-            step.q_pe,
-            step.latent,
-            self.metadata,
-            step.cu_q_seqlens,
-            step.cu_kv_seqlens,
-            step.scheduler_counter,
-            step.bmm1_scale,
-            step.bmm2_scale,
-            step.quant_q,
-        )
-
-    def select_library(self, layer_idx: int, use_sparse_mla: bool) -> None:
-        """Put sparse_mla_decode first in (or remove it from) the layer's FMHA library list."""
-        layer = self.layers[layer_idx]
-        manager = layer._fmha_manager
-        libs = [lib for lib in manager.fmha_libs if not isinstance(lib, SparseMlaDecodeFmha)]
-        if use_sparse_mla:
-            libs.insert(0, SparseMlaDecodeFmha(layer))
-        manager.fmha_libs = libs
-        manager._cache.clear()
-
-    def selected_library(self, layer_idx: int) -> List[str]:
-        return [type(lib).__name__ for lib in self.layers[layer_idx]._fmha_manager._cache.values()]
-
-    def forward(self, layer_idx: int, step: StepBuffers, topk_local: torch.Tensor) -> torch.Tensor:
-        """``TrtllmAttention.forward`` for the generation rows; returns ``[R, 8*512]`` BF16."""
-        layer = self.layers[layer_idx]
-        layer.indexer.forward_from_projected = Mock(return_value=topk_local)
-        return self._forward_generation(layer_idx, step, [])
-
-    def update_indexer_cache(
-        self, layer_idx: int, k_fp8: torch.Tensor, k_scale: torch.Tensor
-    ) -> None:
-        """Native indexer K-cache append of this step's rows (``_update_k_cache``)."""
-        self.layers[layer_idx].indexer._update_k_cache(k_fp8, k_scale, self.metadata)
-
-    def forward_native_indexer(
-        self, layer_idx: int, step: StepBuffers, indexer_intermediates: Sequence[torch.Tensor]
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """``TrtllmAttention.forward`` with the unmodified native indexer.
-
-        ``indexer_intermediates`` is the native ``pre_indexer_proj`` output of this step's
-        rows. Returns the attention output and a copy of the local top-k the native
-        ``Indexer.forward_from_projected`` produced (and the native remap consumed).
-        """
-        indexer = self.layers[layer_idx].indexer
-        indexer.__dict__.pop("forward_from_projected", None)
-        native_select = indexer.forward_from_projected
-        captured = []
-
-        def capture(*args, **kwargs):
-            topk = native_select(*args, **kwargs)
-            captured.append(topk.clone())
-            return topk
-
-        # Also keep the decode scores and lengths the native top-k module received.
-        top_k_module = indexer.top_k
-        top_k_module.__dict__.pop("forward", None)
-        native_top_k = top_k_module.forward
-        self.last_topk_inputs = None
-
-        def capture_top_k(scores, *args, **kwargs):
-            lengths = kwargs.get("sequence_lengths")
-            if isinstance(lengths, torch.Tensor):
-                self.last_topk_inputs = {
-                    "scores": scores.clone(),
-                    "sequence_lengths": lengths.clone(),
-                    "next_n": kwargs.get("next_n"),
-                }
-            return native_top_k(scores, *args, **kwargs)
-
-        indexer.forward_from_projected = capture
-        top_k_module.forward = capture_top_k
-        try:
-            out = self._forward_generation(layer_idx, step, list(indexer_intermediates))
-        finally:
-            indexer.__dict__.pop("forward_from_projected", None)
-            top_k_module.__dict__.pop("forward", None)
-        assert len(captured) == 1, len(captured)
-        return out, captured[0]
-
-    def _forward_generation(
-        self, layer_idx: int, step: StepBuffers, indexer_intermediates: List[torch.Tensor]
-    ) -> torch.Tensor:
-        layer = self.layers[layer_idx]
-        with dispatch.step(dispatch.Phase.DECODE):
-            return layer.forward(
-                step.fused_q,
-                None,
-                None,
-                self.metadata,
-                attention_input_type=AttentionInputType.generation_only,
-                latent_cache=step.latent,
-                q_pe=step.q_pe,
-                cu_q_seqlens=step.cu_q_seqlens,
-                cu_kv_seqlens=step.cu_kv_seqlens,
-                fmha_scheduler_counter=step.scheduler_counter,
-                mla_bmm1_scale=step.bmm1_scale,
-                mla_bmm2_scale=step.bmm2_scale,
-                quant_q_buffer=step.quant_q,
-                sparse_backend_args=DSABackendForwardArgs(
-                    indexer_intermediates=indexer_intermediates
-                ),
-            )
-
-    @contextlib.contextmanager
-    def capture_sparse_mla_inputs(self) -> Iterator[List[dict]]:
-        """Record the global indices, KV lengths and row split of every sparse_mla_decode call (eager only)."""
-        calls: List[dict] = []
-        original = SparseMlaDecodeFmha.forward
-
-        def forward(fmha, q, k, v, metadata, forward_args):
-            gen_lens = metadata.kv_lens_cuda_runtime[metadata.num_contexts :]
-            calls.append(
-                {
-                    "indices": forward_args.sparse_runtime_params.sparse_attn_indices.clone(),
-                    "seq_lens_kv": gen_lens.clone(),
-                    "seq_len_q": q.shape[0] // gen_lens.shape[0],
-                    "ragged": metadata.token_major_gen_view() is not None,
-                }
-            )
-            return original(fmha, q, k, v, metadata, forward_args)
-
-        SparseMlaDecodeFmha.forward = forward
-        try:
-            yield calls
-        finally:
-            SparseMlaDecodeFmha.forward = original
-
-    def global_indices(self, layer_idx: int, topk_local: torch.Tensor) -> torch.Tensor:
-        """Native local -> global pool-row conversion of a local top-k."""
-        layer = self.layers[layer_idx]
-        global_idx, _ = transform_local_topk_and_prepare_pool_view(
-            topk_local, self.metadata, layer.get_local_layer_idx(self.metadata), is_generation=True
-        )
-        return global_idx
-
-    def grouped_global_indices(self, layer_idx: int, topk_local: torch.Tensor) -> torch.Tensor:
-        """The same conversion through the native grouped (cross-layer fan-out) remap op."""
-        from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import (
-            transform_local_topk_and_prepare_pool_view_grouped,
-        )
-
-        local_layer = self.layers[layer_idx].get_local_layer_idx(self.metadata)
-        page_index_scale, layer_offset = (
-            self.metadata.kv_cache_manager.get_primary_pool_page_index_params(local_layer)
-        )
-        layer_ids = torch.tensor([int(layer_offset)], dtype=torch.int32, device=self.device)
-        grouped = transform_local_topk_and_prepare_pool_view_grouped(
-            topk_local, self.metadata, layer_ids, int(page_index_scale), is_generation=True
-        )
-        return grouped[0]
-
-    # ----- top-k builders ---------------------------------------------------------------
-    def causal_len(self, row: int) -> int:
-        """Tokens visible to generation row ``row`` (its own position included)."""
-        return self.ctx_lens[self.row_seq[row]] + self.row_pos[row] + 1
-
-    def list_limits(self) -> List[int]:
-        """Usable index-list positions per row (native rule: the causal length)."""
-        return [self.causal_len(r) for r in range(self.num_rows)]
-
-    def random_topk(self, seed: int, overlap: float = 0.0, recent: int = 0) -> torch.Tensor:
-        """Per-row distinct local token indices; ``-1`` fills rows with fewer tokens.
-
-        Args:
-            overlap: Fraction of each MTP row's selection copied from the previous row.
-            recent: Number of most recent tokens every row always selects.
-        """
-        gen = torch.Generator(device="cpu").manual_seed(seed)
-        topk = torch.full((self.num_rows, TOP_K), -1, dtype=torch.int32)
-        prev = None
-        for row in range(self.num_rows):
-            visible = self.causal_len(row)
-            if visible <= TOP_K:
-                sel = torch.arange(visible, dtype=torch.int64)
-                sel = sel[torch.randperm(visible, generator=gen)]
-            else:
-                must = torch.arange(visible - min(recent, TOP_K), visible, dtype=torch.int64)
-                chosen = set(must.tolist())
-                if prev is not None and overlap > 0 and self.row_pos[row] != 0:
-                    keep = prev[torch.randperm(prev.numel(), generator=gen)[: int(overlap * TOP_K)]]
-                    chosen.update(x for x in keep.tolist() if x < visible)
-                while len(chosen) < TOP_K:
-                    draw = torch.randint(0, visible, (2 * TOP_K,), generator=gen).tolist()
-                    for x in draw:
-                        chosen.add(x)
-                        if len(chosen) == TOP_K:
-                            break
-                sel = torch.tensor(sorted(chosen), dtype=torch.int64)
-                sel = sel[torch.randperm(sel.numel(), generator=gen)]
-            topk[row, : sel.numel()] = sel.to(torch.int32)
-            prev = sel
-        return topk.to(self.device)
-
-
-class _PretrainedStub:
-    rms_norm_eps = 1e-5
-
-
-def negative_inside_causal(
-    indices: torch.Tensor, seq_lens_kv: torch.Tensor, seq_len_q: int
-) -> torch.Tensor:
-    """Per row, the negative entries at index-list positions below the row's causal length.
-
-    This is the input class on which sparse_mla_decode deliberately differs from the native kernel
-    (sparse_mla_decode skips the entry, the native kernel reads the pool row ``-1`` addresses).
-    """
-    rows = indices.shape[0]
-    row = torch.arange(rows, device=indices.device)
-    seq = row // seq_len_q
-    usable = seq_lens_kv.to(indices.device).long()[seq] - (seq_len_q - 1 - row % seq_len_q)
-    positions = torch.arange(indices.shape[1], device=indices.device)
-    inside = positions[None, :] < usable[:, None]
-    return ((indices < 0) & inside).sum(dim=1)
-
-
-# ----- rope_kv_append's top-k bridge as the producer of sparse_mla_decode's indices
-# ------------------------------------
-
-
-def rope_append_selected() -> bool:
-    """Whether this process selects rope_kv_append (whose fused bridge then replaces the native remap)."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    return "rope_kv_append" in dispatch.SELECTOR_PATHS and dispatch.is_enabled("rope_kv_append")
-
-
-@contextlib.contextmanager
-def capture_rope_append_bridge() -> Iterator[List[dict]]:
-    """Record every generation call of rope_kv_append's fused top-k bridge (eager only).
-
-    The DSA predictor looks up rope_kv_append's ``remap_topk_to_global`` entry on every call, so
-    replacing the registered entry sees each call. Engagement is counted here, per call:
-    ``dispatch.log_capability`` keeps only the first decision per key, and the first
-    bridge decision of a process is usually a context batch, which the bridge declines.
-
-    Each record holds the layer, the local top-k the bridge remapped (cloned before the
-    call), the call's metadata, and the global indices it returned (the returned tensor).
-    Yields an empty list when rope_kv_append is not selected.
-    """
-    calls: List[dict] = []
-    if not rope_append_selected():
-        yield calls
-        return
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-
-    original = dispatch.get_impl("rope_kv_append", "remap_topk_to_global")
-
-    def bridge(backend, topk, metadata, packet):
-        if metadata.num_contexts != 0:
-            return original(backend, topk, metadata, packet)
-        source = topk if packet.source is None else packet.source
-        entry = {
-            "layer": backend.layer_idx,
-            "local_layer": backend.get_local_layer_idx(metadata),
-            "topk": source.clone(),
-            "metadata": metadata,
-        }
-        entry["global"] = original(backend, topk, metadata, packet)
-        calls.append(entry)
-        return entry["global"]
-
-    dispatch._REGISTRY["rope_kv_append"]["remap_topk_to_global"] = bridge
-    try:
-        yield calls
-    finally:
-        dispatch._REGISTRY["rope_kv_append"]["remap_topk_to_global"] = original
-
-
-def rope_append_bridge_report(call: dict, seq_lens_kv: torch.Tensor, seq_len_q: int) -> dict:
-    """Producer contract of one recorded bridge call (after the step completed).
-
-    ``equals_native_remap``: the bridge output equals the native local-to-global remap
-    of the same local top-k bit for bit. ``negative_below_causal``: per row, ``-1``
-    entries the bridge emitted at index-list positions below the row's causal length
-    (``seq_lens_kv`` = KV lengths of the generation sequences, this step included).
-    """
-    native, _ = transform_local_topk_and_prepare_pool_view(
-        call["topk"], call["metadata"], call["local_layer"], is_generation=True
-    )
-    out = call["global"]
-    return {
-        "layer": call["layer"],
-        "equals_native_remap": bool(out.shape == native.shape and torch.equal(out, native)),
-        "negative_below_causal": negative_inside_causal(out, seq_lens_kv, seq_len_q).tolist(),
-        "input_negative_below_causal": negative_inside_causal(
-            call["topk"], seq_lens_kv, seq_len_q
-        ).tolist(),
-        "negative_total": int((out < 0).sum().item()),
-    }
-
-
-# ----- CUPTI kernel records ------------------------------------------------------------------
-
-SPARSE_MLA_KERNEL = "sparseMlaDecodeKernel"
-
-
-def is_native_sparse_fmha(name: str) -> bool:
-    """The TRT-LLM-Gen sparse MLA generation kernel (``fmhaSm10x…Sparse…``)."""
-    return "fmhaSm1" in name and "Sparse" in name
-
-
-def kernel_trace(fn: Callable[[], None]) -> List[dict]:
-    """CUPTI records of the kernels ``fn`` enqueues (graph replays included), by start time.
-
-    Uses ``torch.profiler``'s CUDA activity tracing; each record has the kernel name, the
-    stream and the start / end time in microseconds.
-    """
-    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
-        fn()
-        torch.cuda.synchronize()
-    fd, path = tempfile.mkstemp(suffix=".json")
-    os.close(fd)
-    try:
-        prof.export_chrome_trace(path)
-        with open(path) as f:
-            events = json.load(f)["traceEvents"]
-    finally:
-        os.unlink(path)
-    kernels = [
-        {
-            "name": e["name"],
-            "stream": e.get("args", {}).get("stream"),
-            "start": float(e["ts"]),
-            "end": float(e["ts"]) + float(e.get("dur", 0.0)),
-        }
-        for e in events
-        if e.get("cat") == "kernel"
-    ]
-    return sorted(kernels, key=lambda k: k["start"])
-
-
-def attention_neighbours(kernels: List[dict]) -> List[dict]:
-    """For every sparse attention kernel (sparse_mla_decode or native): its neighbours on its stream.
-
-    ``successor_early``: the next kernel of the stream started before the attention
-    kernel ended (it was launched early as a programmatic dependent), by
-    ``successor_lead_us``; ``started_before_predecessor_end``: the attention kernel itself
-    started early, by ``predecessor_lead_us``.
-    """
-    by_stream: Dict[object, List[dict]] = {}
-    for k in kernels:
-        by_stream.setdefault(k["stream"], []).append(k)
-    out = []
-    for stream_kernels in by_stream.values():
-        for i, k in enumerate(stream_kernels):
-            is_sparse_mla = SPARSE_MLA_KERNEL in k["name"]
-            if not (is_sparse_mla or is_native_sparse_fmha(k["name"])):
-                continue
-            prev = stream_kernels[i - 1] if i > 0 else None
-            nxt = stream_kernels[i + 1] if i + 1 < len(stream_kernels) else None
-            out.append(
-                {
-                    "kernel": "sparse_mla_decode" if is_sparse_mla else "native",
-                    "start": k["start"],
-                    "duration_us": k["end"] - k["start"],
-                    "predecessor": prev["name"][:96] if prev else None,
-                    "started_before_predecessor_end": bool(prev and k["start"] < prev["end"]),
-                    "predecessor_lead_us": (prev["end"] - k["start"]) if prev else None,
-                    "successor": nxt["name"][:96] if nxt else None,
-                    "successor_early": bool(nxt and nxt["start"] < k["end"]),
-                    "successor_lead_us": (k["end"] - nxt["start"]) if nxt else None,
-                }
-            )
-    return sorted(out, key=lambda a: a["start"])
-
-
-# ----- FP64 oracles --------------------------------------------------------------------------
-
-
-def e4m3_round(scale: float) -> Callable[[torch.Tensor], torch.Tensor]:
-    """Probability rounding to e4m3 after multiplying by ``scale``."""
-
-    def fn(p: torch.Tensor) -> torch.Tensor:
-        q = (p * scale).clamp(max=E4M3_MAX).to(torch.float32).to(torch.float8_e4m3fn)
-        return q.to(torch.float64) / scale
-
-    return fn
-
-
-def fp16_round(p: torch.Tensor) -> torch.Tensor:
-    return p.to(torch.float16).to(torch.float64)
-
-
-def bf16_round(p: torch.Tensor) -> torch.Tensor:
-    return p.to(torch.bfloat16).to(torch.float64)
-
-
-def oracle_attention(
-    q_bytes: torch.Tensor,
-    pool_rows: torch.Tensor,
-    global_idx: torch.Tensor,
-    bmm1_scale: float,
-    bmm2_scale: float,
-    *,
-    list_limit: Optional[Sequence[int]] = None,
-    prob_round: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
-    split: Optional[int] = None,
-    step: Optional[int] = None,
-) -> torch.Tensor:
-    """FP64 sparse MLA reference over the e4m3 query and pool; returns ``[R, 8, 512]``.
-
-    Without ``prob_round`` this is the exact softmax. With it, the probabilities fed to
-    P*V are rounded the way an online-softmax kernel would round them: the index list
-    is cut into ``split``-position chunks merged at the end, each processed in
-    ``step``-position steps against the running maximum; row sums stay unrounded.
-    """
-    rows = q_bytes.shape[0]
-    q = q_bytes.reshape(rows, NUM_HEADS, HEAD_DIM).view(torch.float8_e4m3fn).to(torch.float64)
-    pool_bytes = pool_rows.view(torch.uint8)
-
-    def gather(sel: torch.Tensor) -> torch.Tensor:
-        return pool_bytes[sel].view(torch.float8_e4m3fn).to(torch.float64)
-
-    out = torch.zeros(rows, NUM_HEADS, KV_LORA_RANK, dtype=torch.float64, device=q.device)
-    positions = torch.arange(global_idx.shape[1], device=global_idx.device)
-    for r in range(rows):
-        idx = global_idx[r].to(torch.int64)
-        valid = idx >= 0
-        if list_limit is not None:
-            valid &= positions < list_limit[r]
-        if not bool(valid.any()):
-            continue
-        if prob_round is None:
-            k = gather(idx[valid])
-            s = (q[r] @ k.T) * bmm1_scale
-            p = torch.softmax(s, dim=-1)
-            out[r] = p @ k[:, :KV_LORA_RANK]
-            continue
-        chunk = split or global_idx.shape[1]
-        stride = step or chunk
-        maxes, sums, accs = [], [], []
-        for a in range(0, global_idx.shape[1], chunk):
-            m = torch.full((NUM_HEADS,), -math.inf, dtype=torch.float64, device=q.device)
-            row_sum = torch.zeros(NUM_HEADS, dtype=torch.float64, device=q.device)
-            acc = torch.zeros(NUM_HEADS, KV_LORA_RANK, dtype=torch.float64, device=q.device)
-            for b in range(a, min(a + chunk, global_idx.shape[1]), stride):
-                sel = valid[b : b + stride]
-                if not bool(sel.any()):
-                    continue
-                k = gather(idx[b : b + stride][sel])
-                s = (q[r] @ k.T) * bmm1_scale
-                m_new = torch.maximum(m, s.max(dim=-1).values)
-                alpha = torch.exp(m - m_new)
-                p = torch.exp(s - m_new[:, None])
-                row_sum = row_sum * alpha + p.sum(dim=-1)
-                acc = acc * alpha[:, None] + prob_round(p) @ k[:, :KV_LORA_RANK]
-                m = m_new
-            maxes.append(m)
-            sums.append(row_sum)
-            accs.append(acc)
-        m_all = torch.stack(maxes)
-        global_max = m_all.max(dim=0).values
-        weights = torch.where(
-            torch.isinf(m_all), torch.zeros_like(m_all), torch.exp(m_all - global_max)
-        )
-        total = (weights * torch.stack(sums)).sum(dim=0)
-        out[r] = (weights[:, :, None] * torch.stack(accs)).sum(dim=0) / total[:, None]
-    return out * bmm2_scale
-
-
-def error_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict:
-    """Normalized RMS and max-abs error of ``candidate`` against ``reference``."""
-    c = candidate.reshape(reference.shape).to(torch.float64)
-    diff = c - reference
-    ref_rms = reference.pow(2).mean().sqrt().item()
-    return {
-        "nrms": diff.pow(2).mean().sqrt().item() / max(ref_rms, 1e-30),
-        "max_abs": diff.abs().max().item(),
-        "ref_rms": ref_rms,
-        "nonfinite": int((~torch.isfinite(c)).sum().item()),
-    }
-
-
-# ----- native-versus-sparse_mla_decode comparison
-# -----------------------------------------------------------
-
-
-def record(filename: str, case: str, payload: dict) -> None:
-    """Print one JSON result line and append it to ``$GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS/<filename>``."""
-    import json
-    import os
-
-    line = json.dumps({"case": case, **payload}, sort_keys=True)
-    print(line)
-    results_dir = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS")
-    if results_dir:
-        with open(os.path.join(results_dir, filename), "a") as f:
-            f.write(line + "\n")
-
-
-def run_both(harness: DsaGenHarness, step: StepBuffers, topk: torch.Tensor):
-    """Native and sparse_mla_decode outputs for one step plus the FP64 oracle with the native rules."""
-    harness.select_library(0, use_sparse_mla=False)
-    native = harness.forward(0, step, topk).clone()
-    assert "SparseMlaDecodeFmha" not in harness.selected_library(0)
-    harness.select_library(0, use_sparse_mla=True)
-    sparse_out = harness.forward(0, step, topk).clone()
-    assert harness.selected_library(0) == ["SparseMlaDecodeFmha"], harness.selected_library(0)
-    global_idx = harness.global_indices(0, topk)
-    return native, sparse_out, global_idx, oracle_for(harness, step, global_idx)
-
-
-def usable_rows(harness: DsaGenHarness, global_idx: torch.Tensor) -> torch.Tensor:
-    positions = torch.arange(global_idx.shape[1], device=global_idx.device)
-    limits = torch.tensor(harness.list_limits(), device=global_idx.device)[:, None]
-    return ((global_idx >= 0) & (positions[None, :] < limits)).any(dim=1)
-
-
-def oracle_for(harness: DsaGenHarness, step: StepBuffers, global_idx: torch.Tensor) -> torch.Tensor:
-    """FP64 oracle with the native rules for this step's query and the current pool."""
-    return oracle_attention(
-        step.quant_q,
-        harness.pool_rows(),
-        global_idx,
-        float(step.bmm1_scale[0]),
-        float(step.bmm2_scale[0]),
-        list_limit=harness.list_limits(),
-    )
-
-
-def compare(
-    harness: DsaGenHarness,
-    step: StepBuffers,
-    topk: torch.Tensor,
-    case: str,
-    filename: str = "sparse_mla_correctness.jsonl",
-    **extra,
-) -> dict:
-    """Check sparse_mla_decode against native and the oracle with the frozen error budget."""
-    native, sparse_out, global_idx, oracle = run_both(harness, step, topk)
-    return check_budget(harness, native, sparse_out, global_idx, oracle, case, filename, **extra)
-
-
-def check_budget(
-    harness: DsaGenHarness,
-    native: torch.Tensor,
-    sparse_out: torch.Tensor,
-    global_idx: torch.Tensor,
-    oracle: torch.Tensor,
-    case: str,
-    filename: str = "sparse_mla_correctness.jsonl",
-    **extra,
-) -> dict:
-    """The frozen error budget for one native and one sparse_mla_decode output of the same inputs."""
-    native_err = error_stats(native.double(), oracle)
-    sparse_mla_err = error_stats(sparse_out.double(), oracle)
-    sparse_mla_vs_native = error_stats(sparse_out.double(), native.double().reshape(oracle.shape))
-    record(
-        filename,
-        case,
-        {
-            "native_vs_oracle": native_err,
-            "sparse_mla_vs_oracle": sparse_mla_err,
-            "sparse_mla_vs_native": sparse_mla_vs_native,
-            **extra,
-        },
-    )
-    native_rows = native.reshape(harness.num_rows, -1)
-    sparse_mla_rows = sparse_out.reshape(harness.num_rows, -1)
-    live = usable_rows(harness, global_idx)
-    if bool((~live).any()):
-        torch.testing.assert_close(
-            sparse_mla_rows[~live], native_rows[~live], rtol=0, atol=0, equal_nan=True
-        )
-    assert sparse_mla_err["nonfinite"] == native_err["nonfinite"]
-    assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (sparse_mla_err, native_err)
-    assert sparse_mla_err["max_abs"] <= 2.0 * native_err["max_abs"] + 1e-3, (
-        sparse_mla_err,
-        native_err,
-    )
-    oracle_rows = oracle.reshape(harness.num_rows, -1)
-    assert torch.allclose(sparse_mla_rows[live].double(), oracle_rows[live], atol=0.12, rtol=0.01)
-    return {
-        "native": native,
-        "sparse_mla_decode": sparse_out,
-        "oracle": oracle,
-        "global_idx": global_idx,
-    }
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py
deleted file mode 100644
index ecaac194346b..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_pdl_helpers.py
+++ /dev/null
@@ -1,204 +0,0 @@
-# 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.
-"""Test-only PDL producer and consumer kernels for the sparse_mla_decode stress test.
-
-``delayed_scatter`` is a producer: it allows its dependents to launch at once
-(``griddepcontrol.launch_dependents``), waits for its own predecessor, spins for
-a given time and only then copies a list of segments into the buffers the next
-kernel reads. ``pdl_consumer`` is a consumer: it also allows its dependents to
-launch at once, waits for its predecessor and copies a buffer out, optionally
-after a spin. Both are launched with programmatic stream serialization, so a
-chain producer -> sparse_mla_decode -> consumer -> producer -> ... has no plain kernel
-boundary: every kernel may become resident while its predecessor still runs,
-and only ``griddepcontrol.wait`` orders the memory accesses.
-
-Both kernels record ``%globaltimer`` stamps so that the test can show that the
-early launches actually happened.
-"""
-
-import os
-
-import torch
-
-_CPP_DECLS = r"""
-#include <torch/extension.h>
-void delayed_scatter(torch::Tensor dsts, torch::Tensor srcs, torch::Tensor counts, int64_t spin_ns,
-    torch::Tensor stamps, bool use_pdl, int64_t num_blocks);
-void pdl_consumer(torch::Tensor dst, torch::Tensor src, int64_t spin_ns, torch::Tensor stamps, bool use_pdl);
-"""
-
-_CUDA_SRC = r"""
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda_runtime.h>
-
-namespace
-{
-
-__device__ __forceinline__ unsigned long long globalTimer()
-{
-    unsigned long long t;
-    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
-    return t;
-}
-
-__device__ __forceinline__ void spinFor(long long spinNs)
-{
-    if (spinNs <= 0)
-    {
-        return;
-    }
-    if (threadIdx.x == 0)
-    {
-        unsigned long long const start = globalTimer();
-        while (globalTimer() - start < static_cast<unsigned long long>(spinNs))
-        {
-        }
-    }
-    __syncthreads();
-}
-
-// stamps: [0] producer start (before waiting), [1] producer after its wait, [2] producer end.
-__global__ void delayedScatterKernel(long long const* dsts, long long const* srcs, long long const* counts,
-    int numSegments, long long spinNs, unsigned long long* stamps)
-{
-    unsigned long long const t0 = globalTimer();
-    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
-    asm volatile("griddepcontrol.wait;" ::: "memory");
-    unsigned long long const t1 = globalTimer();
-    if (blockIdx.x == 0 && threadIdx.x == 0)
-    {
-        stamps[0] = t0;
-        stamps[1] = t1;
-    }
-    spinFor(spinNs);
-    for (int s = 0; s < numSegments; ++s)
-    {
-        uint4* dst = reinterpret_cast<uint4*>(dsts[s]);
-        uint4 const* src = reinterpret_cast<uint4 const*>(srcs[s]);
-        long long const n = counts[s];
-        for (long long i = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x; i < n;
-             i += static_cast<long long>(gridDim.x) * blockDim.x)
-        {
-            dst[i] = src[i];
-        }
-    }
-    __syncthreads();
-    if (threadIdx.x == 0)
-    {
-        atomicMax(stamps + 2, globalTimer());
-    }
-}
-
-// stamps: [0] consumer start (before waiting), [1] consumer after its wait.
-__global__ void pdlConsumerKernel(
-    uint4 const* src, uint4* dst, long long n, long long spinNs, unsigned long long* stamps)
-{
-    unsigned long long const t0 = globalTimer();
-    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
-    asm volatile("griddepcontrol.wait;" ::: "memory");
-    unsigned long long const t1 = globalTimer();
-    if (blockIdx.x == 0 && threadIdx.x == 0)
-    {
-        stamps[0] = t0;
-        stamps[1] = t1;
-    }
-    spinFor(spinNs);
-    for (long long i = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x; i < n;
-         i += static_cast<long long>(gridDim.x) * blockDim.x)
-    {
-        dst[i] = src[i];
-    }
-}
-
-cudaLaunchAttribute pdlAttribute(bool usePdl)
-{
-    cudaLaunchAttribute attr;
-    attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
-    attr.val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0;
-    return attr;
-}
-
-} // namespace
-
-void delayed_scatter(torch::Tensor dsts, torch::Tensor srcs, torch::Tensor counts, int64_t spin_ns,
-    torch::Tensor stamps, bool use_pdl, int64_t num_blocks)
-{
-    TORCH_CHECK(dsts.is_cuda() && srcs.is_cuda() && counts.is_cuda() && stamps.is_cuda());
-    TORCH_CHECK(dsts.scalar_type() == torch::kInt64 && srcs.scalar_type() == torch::kInt64
-        && counts.scalar_type() == torch::kInt64 && stamps.scalar_type() == torch::kInt64);
-    TORCH_CHECK(dsts.numel() == srcs.numel() && dsts.numel() == counts.numel() && stamps.numel() >= 3);
-    cudaLaunchConfig_t config{};
-    config.gridDim = dim3(static_cast<unsigned>(num_blocks));
-    config.blockDim = dim3(256);
-    config.stream = at::cuda::getCurrentCUDAStream(dsts.get_device());
-    cudaLaunchAttribute attr = pdlAttribute(use_pdl);
-    config.attrs = &attr;
-    config.numAttrs = 1;
-    TORCH_CHECK(cudaLaunchKernelEx(&config, delayedScatterKernel, reinterpret_cast<long long const*>(dsts.data_ptr()),
-                    reinterpret_cast<long long const*>(srcs.data_ptr()),
-                    reinterpret_cast<long long const*>(counts.data_ptr()), static_cast<int>(dsts.numel()),
-                    static_cast<long long>(spin_ns), reinterpret_cast<unsigned long long*>(stamps.data_ptr()))
-        == cudaSuccess);
-}
-
-void pdl_consumer(torch::Tensor dst, torch::Tensor src, int64_t spin_ns, torch::Tensor stamps, bool use_pdl)
-{
-    TORCH_CHECK(dst.is_cuda() && src.is_cuda() && stamps.is_cuda() && stamps.scalar_type() == torch::kInt64);
-    TORCH_CHECK(dst.is_contiguous() && src.is_contiguous() && dst.nbytes() == src.nbytes() && src.nbytes() % 16 == 0);
-    TORCH_CHECK(stamps.numel() >= 2);
-    cudaLaunchConfig_t config{};
-    config.gridDim = dim3(16);
-    config.blockDim = dim3(256);
-    config.stream = at::cuda::getCurrentCUDAStream(dst.get_device());
-    cudaLaunchAttribute attr = pdlAttribute(use_pdl);
-    config.attrs = &attr;
-    config.numAttrs = 1;
-    TORCH_CHECK(cudaLaunchKernelEx(&config, pdlConsumerKernel, reinterpret_cast<uint4 const*>(src.data_ptr()),
-                    reinterpret_cast<uint4*>(dst.data_ptr()), static_cast<long long>(src.nbytes() / 16),
-                    static_cast<long long>(spin_ns), reinterpret_cast<unsigned long long*>(stamps.data_ptr()))
-        == cudaSuccess);
-}
-"""
-
-_MODULE = None
-
-
-def load_pdl_helpers():
-    """Build (once per process) and return the helper extension."""
-    global _MODULE
-    if _MODULE is None:
-        from torch.utils.cpp_extension import load_inline
-
-        major, minor = torch.cuda.get_device_capability()
-        arch = f"{major}{minor}"
-        build_dir = os.path.join(
-            os.environ.get(
-                "GLM_MOE_DSA_SPARSE_MLA_DECODE_HELPER_BUILD_DIR",
-                "/tmp/glm_moe_dsa_sparse_mla_decode_helpers",
-            ),
-            arch,
-        )
-        os.makedirs(build_dir, exist_ok=True)
-        _MODULE = load_inline(
-            name=f"glm_moe_dsa_sparse_mla_decode_pdl_helpers_{arch}",
-            cpp_sources=_CPP_DECLS,
-            cuda_sources=_CUDA_SRC,
-            functions=["delayed_scatter", "pdl_consumer"],
-            extra_cuda_cflags=["-O3", f"-gencode=arch=compute_{arch},code=sm_{arch}"],
-            build_directory=build_dir,
-            verbose=False,
-        )
-    return _MODULE
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py
deleted file mode 100644
index ca07bd36a9c7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_real_fixture.py
+++ /dev/null
@@ -1,310 +0,0 @@
-# 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.
-"""Real-weight sparse_attention contract inputs from GLM-5.2 layer 0 on real text.
-
-Layer 0 is the one target layer whose input is fully determined by the token
-ids (embedding + input RMSNorm), so its attention inputs are exactly the
-model's. The projections use the checkpoint's BF16 attention weights for the
-TP rank 0 head slice; RoPE, the FP8 latent cache and the FP8 query come from
-the native TRT-LLM ops driven by ``DsaGenHarness``. Two top-k sources exist:
-
-* ``indexer_topk``: the native indexer projections (``Indexer.pre_indexer_proj``)
-  with the checkpoint's indexer weights, followed by a *reference* FP32
-  evaluation of the indexer logits ``sum_h w_h * relu(q_h . k) * k_scale`` and an
-  exact top-k (not the native top-k kernels);
-* ``indexer_projections`` + ``DsaGenHarness.forward_native_indexer``: the native
-  indexer K cache written by the native ``_update_k_cache`` and the unmodified
-  native ``Indexer.forward_from_projected`` (paged MQA logits + decode top-k).
-"""
-
-import json
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Dict, Iterable, Optional
-
-import torch
-from utils.llm_data import llm_models_root
-
-import tensorrt_llm
-
-NUM_LOCAL_HEADS = 8
-QK_NOPE = 192
-QK_ROPE = 64
-V_HEAD = 256
-KV_LORA = 512
-NUM_HEADS = 64
-RMS_EPS = 1e-5
-CHECKPOINT_NAME = "GLM-5.2-NVFP4"
-
-
-def checkpoint_dir() -> Optional[Path]:
-    root = llm_models_root()
-    if root is None:
-        return None
-    path = Path(root) / CHECKPOINT_NAME
-    return path if (path / "config.json").is_file() else None
-
-
-def load_tensors(ckpt: Path, names: Iterable[str], device: str = "cuda") -> Dict[str, torch.Tensor]:
-    from safetensors import safe_open
-
-    weight_map = json.loads((ckpt / "model.safetensors.index.json").read_text())["weight_map"]
-    by_shard: Dict[str, list] = {}
-    for name in names:
-        by_shard.setdefault(weight_map[name], []).append(name)
-    tensors = {}
-    for shard, shard_names in by_shard.items():
-        with safe_open(str(ckpt / shard), framework="pt", device=device) as f:
-            for name in shard_names:
-                tensors[name] = f.get_tensor(name)
-    return tensors
-
-
-def real_token_ids(ckpt: Path, count: int) -> torch.Tensor:
-    """Token ids of real text: the installed TensorRT-LLM Python sources, in path order."""
-    from tokenizers import Tokenizer
-
-    tokenizer = Tokenizer.from_file(str(ckpt / "tokenizer.json"))
-    ids = []
-    for path in sorted(Path(tensorrt_llm.__file__).resolve().parent.rglob("*.py")):
-        ids.extend(tokenizer.encode(path.read_text(errors="ignore")).ids)
-        if len(ids) >= count:
-            break
-    if len(ids) < count:
-        raise ValueError(f"only {len(ids)} tokens of text available, {count} requested")
-    return torch.tensor(ids[:count], dtype=torch.long)
-
-
-def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = RMS_EPS) -> torch.Tensor:
-    xf = x.float()
-    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float()).to(x.dtype)
-
-
-@dataclass
-class LayerInputs:
-    """Token-wise projections of one layer for every token of the sequence."""
-
-    hidden: torch.Tensor  # [N, 6144] post input-RMSNorm
-    qr: torch.Tensor  # [N, 2048] post q_a RMSNorm
-    fused_q: torch.Tensor  # [N, 8 * 576] absorbed q_nope | raw q_pe
-    q_pe: torch.Tensor  # [N, 8, 64] raw q_pe
-    latent: torch.Tensor  # [N, 576] normalized latent | raw k_pe
-
-
-def layer_weight_names(layer: int) -> list:
-    prefix = f"model.layers.{layer}."
-    attn = prefix + "self_attn."
-    return [
-        "model.embed_tokens.weight",
-        prefix + "input_layernorm.weight",
-        attn + "q_a_proj.weight",
-        attn + "q_a_layernorm.weight",
-        attn + "q_b_proj.weight",
-        attn + "kv_a_proj_with_mqa.weight",
-        attn + "kv_a_layernorm.weight",
-        attn + "kv_b_proj.weight",
-        attn + "indexer.wq_b.weight",
-        attn + "indexer.wk.weight",
-        attn + "indexer.k_norm.weight",
-        attn + "indexer.k_norm.bias",
-        attn + "indexer.weights_proj.weight",
-    ]
-
-
-@torch.inference_mode()
-def project_layer0(weights: Dict[str, torch.Tensor], ids: torch.Tensor) -> LayerInputs:
-    """Layer-0 attention inputs of TP rank 0 (heads 0-7) for every token."""
-    attn = "model.layers.0.self_attn."
-    hidden = rms_norm(
-        weights["model.embed_tokens.weight"][ids.cuda()],
-        weights["model.layers.0.input_layernorm.weight"],
-    )
-    kv = hidden @ weights[attn + "kv_a_proj_with_mqa.weight"].T
-    compressed = rms_norm(kv[:, :KV_LORA], weights[attn + "kv_a_layernorm.weight"])
-    k_pe = kv[:, KV_LORA:]
-    qr = rms_norm(
-        hidden @ weights[attn + "q_a_proj.weight"].T, weights[attn + "q_a_layernorm.weight"]
-    )
-    q_b = weights[attn + "q_b_proj.weight"][: NUM_LOCAL_HEADS * (QK_NOPE + QK_ROPE)]
-    q = (qr @ q_b.T).view(-1, NUM_LOCAL_HEADS, QK_NOPE + QK_ROPE)
-    q_nope, q_pe = q.split([QK_NOPE, QK_ROPE], dim=-1)
-    kv_b = weights[attn + "kv_b_proj.weight"].view(NUM_HEADS, QK_NOPE + V_HEAD, KV_LORA)
-    w_uk = kv_b[:NUM_LOCAL_HEADS, :QK_NOPE, :]
-    q_absorbed = torch.einsum("nhd,hdk->nhk", q_nope.float(), w_uk.float()).to(torch.bfloat16)
-    fused_q = torch.cat([q_absorbed, q_pe], dim=-1).reshape(
-        -1, NUM_LOCAL_HEADS * (KV_LORA + QK_ROPE)
-    )
-    return LayerInputs(
-        hidden=hidden.contiguous(),
-        qr=qr.contiguous(),
-        fused_q=fused_q.contiguous(),
-        q_pe=q_pe.contiguous(),
-        latent=torch.cat([compressed, k_pe], dim=-1).contiguous(),
-    )
-
-
-def load_indexer_weights(indexer, weights: Dict[str, torch.Tensor], layer: int = 0) -> None:
-    prefix = f"model.layers.{layer}.self_attn.indexer."
-    with torch.no_grad():
-        for name, module_param in (
-            ("wq_b.weight", indexer.wq_b.weight),
-            ("wk.weight", indexer.wk.weight),
-            ("weights_proj.weight", indexer.weights_proj.weight),
-            ("k_norm.weight", indexer.k_norm.weight),
-            ("k_norm.bias", indexer.k_norm.bias),
-        ):
-            module_param.data = (
-                weights[prefix + name].to(device="cuda", dtype=module_param.dtype).contiguous()
-            )
-    indexer.cache_derived_state()
-
-
-@torch.inference_mode()
-def reference_indexer_logits(indexer, inputs: LayerInputs, ctx_len: int, seq_len_q: int) -> list:
-    """FP32 indexer logits ``sum_h w_h * relu(q_h . k) * k_scale`` of each generation row.
-
-    Uses the native indexer projections; row ``r`` covers its ``ctx_len + r + 1`` visible tokens.
-    """
-    num_tokens = ctx_len + seq_len_q
-    positions = torch.arange(num_tokens, dtype=torch.int32, device="cuda")
-    q_fp8, k_fp8, k_scale, weights, _ = indexer.pre_indexer_proj(
-        inputs.qr[:num_tokens], inputs.hidden[:num_tokens], positions
-    )
-    q = q_fp8[ctx_len:num_tokens].float().reshape(seq_len_q, indexer.n_heads, indexer.head_dim)
-    keys = (
-        k_fp8.float().reshape(num_tokens, indexer.head_dim)
-        * k_scale.float().reshape(num_tokens, -1)[:, :1]
-    )
-    head_weights = weights[ctx_len:num_tokens].float().reshape(seq_len_q, indexer.n_heads)
-    logits = []
-    for row in range(seq_len_q):
-        visible = ctx_len + row + 1
-        logits.append(
-            (head_weights[row][:, None] * torch.relu(q[row] @ keys[:visible].T)).sum(dim=0)
-        )
-    return logits
-
-
-@torch.inference_mode()
-def indexer_topk(
-    indexer, inputs: LayerInputs, ctx_len: int, seq_len_q: int, topk: int = 2048
-) -> torch.Tensor:
-    """Reference local top-k of the generation rows: exact top-k of the FP32 reference logits."""
-    out = torch.full((seq_len_q, topk), -1, dtype=torch.int32, device="cuda")
-    for row, logits in enumerate(reference_indexer_logits(indexer, inputs, ctx_len, seq_len_q)):
-        count = min(topk, logits.numel())
-        out[row, :count] = logits.topk(count).indices.to(torch.int32)
-    return out
-
-
-def selection_vs_reference(topk_local: torch.Tensor, logits: list, topk: int = 2048) -> list:
-    """How a selection ranks under the FP32 reference logits, per row.
-
-    ``kth`` is the reference top-k boundary value; ``ties_at_kth`` counts visible tokens whose
-    reference logit equals it; ``selected_at_or_above_kth`` is the fraction of the selection whose
-    reference logit is at least ``kth`` (1.0 for any exact top-k up to tie order), and
-    ``selected_within_1pct`` the same with a tolerance of 1 % of the logit range.
-    """
-    rows = []
-    for r, row_logits in enumerate(logits):
-        count = min(topk, row_logits.numel())
-        kth = row_logits.topk(count).values[-1]
-        ids = topk_local[r][topk_local[r] >= 0].long()
-        chosen = row_logits[ids]
-        span = (row_logits.max() - row_logits.min()).clamp_min(1e-30)
-        rows.append(
-            {
-                "kth": float(kth),
-                "max": float(row_logits.max()),
-                "ties_at_kth": int((row_logits == kth).sum()),
-                "zero_logits": int((row_logits == 0).sum()),
-                "selected_at_or_above_kth": float((chosen >= kth).float().mean()),
-                "selected_within_1pct": float((chosen >= kth - 0.01 * span).float().mean()),
-            }
-        )
-    return rows
-
-
-@torch.inference_mode()
-def indexer_projections(indexer, inputs: LayerInputs, start: int, end: int) -> list:
-    """Native ``pre_indexer_proj`` of tokens ``[start, end)``: ``[q_fp8, k_fp8, k_scale, w, q_scale]``."""
-    positions = torch.arange(start, end, dtype=torch.int32, device="cuda")
-    return list(indexer.pre_indexer_proj(inputs.qr[start:end], inputs.hidden[start:end], positions))
-
-
-def topk_contract(topk_local: torch.Tensor, ctx_len: int) -> list:
-    """Per generation row: where the local top-k holds negative, future and repeated ids.
-
-    Row ``r`` of a verify window sees ``ctx_len + r + 1`` tokens; ``negative_below_causal``
-    counts ``-1`` entries at list positions inside that causal length (where the native
-    sparse MLA kernel reads the row ``-1`` addresses and sparse_mla_decode skips the entry).
-    """
-    rows = []
-    for r in range(topk_local.shape[0]):
-        causal = ctx_len + r + 1
-        row = topk_local[r].long()
-        head = row[: min(causal, row.numel())]
-        valid = head[head >= 0]
-        rows.append(
-            {
-                "causal_len": causal,
-                "negative_below_causal": int((head < 0).sum()),
-                "negative_total": int((row < 0).sum()),
-                "ids_at_or_past_causal": int((row >= causal).sum()),
-                "duplicates": int(valid.numel() - valid.unique().numel()),
-                "valid": int((row >= 0).sum()),
-            }
-        )
-    return rows
-
-
-def selection_vs_kernel_scores(
-    topk_local: torch.Tensor, topk_inputs: dict, reference_logits: list, topk: int = 2048
-) -> list:
-    """The native decode scores the top-k kernel received, against its output and the reference.
-
-    Row ``r`` of a verify window of ``next_n`` rows scores ``seq_len - next_n + r + 1`` tokens
-    (the native decode top-k rule). ``exact_agreement``: overlap of the selection with the exact
-    top-k of those scores; ``scores_reference_corr``: Pearson correlation of the scores with the
-    FP32 reference logits of the same tokens.
-    """
-    scores = topk_inputs["scores"].float()
-    seq_len = int(topk_inputs["sequence_lengths"].reshape(-1)[0])
-    next_n = int(topk_inputs["next_n"])
-    rows = []
-    for r in range(topk_local.shape[0]):
-        visible = seq_len - next_n + r + 1
-        row_scores = scores[r, :visible]
-        count = min(topk, visible)
-        exact = set(row_scores.topk(count).indices.tolist())
-        chosen = set(x for x in topk_local[r].tolist() if x >= 0)
-        ref = reference_logits[r][:visible].float()
-        corr = torch.corrcoef(torch.stack([row_scores, ref]))[0, 1]
-        rows.append(
-            {
-                "scores_shape": list(scores.shape),
-                "visible": visible,
-                "exact_agreement": len(chosen & exact) / max(len(exact), 1),
-                "scores_reference_corr": float(corr),
-            }
-        )
-    return rows
-
-
-def recent_window_topk(num_tokens: int, topk: int = 2048) -> torch.Tensor:
-    """Causal context selection used only to run the native context step."""
-    rows = torch.arange(num_tokens, dtype=torch.int32, device="cuda")[:, None]
-    cols = torch.arange(topk, dtype=torch.int32, device="cuda")[None, :]
-    return (rows - cols).clamp_(min=-1)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py
deleted file mode 100644
index f590474a50c7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/sparse_mla_stack_case.py
+++ /dev/null
@@ -1,480 +0,0 @@
-# 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.
-"""Real GLM-5.2 attention layers through ``GlmMLA.forward``, for sparse_mla_decode composition tests.
-
-Run as a script in a fresh interpreter: ``python sparse_mla_stack_case.py <output.pt>`` with the
-fast paths to run in ``GLM_MOE_DSA_TEST_FAST_PATHS`` (``plan_testlib.FAST_PATHS_ENV``). It uses
-the attention_output_tail tests' ``real_mla_harness.AttentionStack``: the TP8 rank-0 shards of two
-consecutive real target layers (default 2, which runs its own
-indexer, and 3, which reuses layer 2's top-k), a native prefill of
-``GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_CTX`` real-token positions, and the MTP verify step of one request
-(6 rows) with the generation metadata the decoder runner builds (CUDA-graph
-metadata, ``update_spec_dec_param``, ``prepare``). The layers are the target's adopted ones and
-the verify step is a decode step, so ``GlmMLA.forward`` runs through ``register_to_config`` (the
-target's DSA attention op) as in the model and each package at the MLA / DSA call sites runs
-where the model runs it:
-absorbed_query (full-indexer layer) and shared_indexer_absorbed_query (shared-indexer layer) produce
-the query and latent row, rope_kv_append the RoPE / FP8 query / latent append, the indexer-K scatter
-and the top-k bridge, the attention_output_tail (attention_output_tail_prefetch /
-attention_output_into_hop schedules) consumes the attention output as its programmatic dependent,
-and sparse_mla_decode computes the attention. The process stands in for one of eight TP ranks
-(o_proj all-reduce off).
-
-``GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_TOPK=fixed`` (default) replaces the full layer's indexer
-selection by a given local top-k (deterministic indices; the shared layer reuses it through
-the shared top-k buffer); ``indexer`` runs the native indexer.
-
-Saved (``torch.save``):
-
-* ``inputs``: an eager step per input set; per layer the FMHA inputs (FP8
-  query bytes, bmm scales, global indices, KV lengths, attended pool rows),
-  cloned before the attention, plus the latent and module outputs;
-* ``pdl``: the same steps with no recording kernel anywhere, so the kernel
-  after the attention is the one the model runs there;
-* ``serialized``: the same steps with a device synchronization right after
-  each sparse_mla_decode launch (no programmatic overlap with the consumer);
-* ``graph``: replays of one captured step over alternating input sets;
-* ``kernels``: CUPTI records of two eager steps and of graph replays: kernel
-  counts and, per attention kernel, its predecessor and successor on the
-  stream and whether they overlapped it;
-* ``engagement``: per-layer counts of absorbed_query / shared_indexer_absorbed_query packets,
-  rope_kv_append fused RoPE calls and generation bridge calls, the attention_output_tail counters
-  and config, sparse_mla_decode's dispatch accounting, and the logged capability decisions;
-* ``rope_append_bridge``: per generation call of rope_kv_append's bridge in the ``inputs`` steps,
-  whether its output equals the native remap of the same local top-k, the
-  ``-1`` entries it emitted below a row's causal length, and whether the
-  attention received exactly that output.
-"""
-
-import dataclasses
-import json
-import os
-import sys
-from collections import Counter, defaultdict
-from pathlib import Path
-from unittest.mock import Mock
-
-_HERE = Path(__file__).resolve().parent
-sys.path.insert(0, str(_HERE))
-sys.path.insert(0, str(_HERE.parent / "attention_output_tail"))
-sys.path.insert(0, str(_HERE.parent))
-
-import plan_testlib  # noqa: E402
-import real_mla_harness  # noqa: E402
-import torch  # noqa: E402
-from sparse_mla_harness import (  # noqa: E402
-    SPARSE_MLA_KERNEL,
-    attention_neighbours,
-    is_native_sparse_fmha,
-    kernel_trace,
-    rope_append_bridge_report,
-)
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E402
-    dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E402, E501
-    sparse_mla,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E402, E501
-    SparseMlaDecodeFmha,
-)
-from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha  # noqa: E402
-from tensorrt_llm._torch.distributed import AllReduceParams  # noqa: E402
-
-ROWS = 6
-TOP_K = 2048
-HEAD_DIM = 576
-
-
-def _selected(feature_id: str) -> bool:
-    return feature_id in dispatch.SELECTOR_PATHS and dispatch.is_enabled(feature_id)
-
-
-class Engagement:
-    """Per-call counters of the producer packages, read at their registered entry points."""
-
-    def __init__(self):
-        self.counts = defaultdict(int)
-        self.bridge = []
-        self.record_bridge = False
-        if _selected("absorbed_query"):
-            for name, label in (
-                ("dsa_proj", "absorbed_query"),
-                ("shared_dsa_proj", "shared_indexer_absorbed_query"),
-            ):
-                self._count_packets(name, label)
-        if _selected("rope_kv_append"):
-            rope = dispatch.get_impl("rope_kv_append", "mla_rope_generation")
-
-            def rope_hook(mqa, *args, **kwargs):
-                self.counts[f"rope_append_rope/layer{mqa.layer_idx}"] += 1
-                return rope(mqa, *args, **kwargs)
-
-            dispatch._REGISTRY["rope_kv_append"]["mla_rope_generation"] = rope_hook
-            bridge = dispatch.get_impl("rope_kv_append", "remap_topk_to_global")
-
-            def bridge_hook(backend, topk, metadata, packet):
-                if metadata.num_contexts != 0:
-                    return bridge(backend, topk, metadata, packet)
-                self.counts[f"rope_append_bridge/layer{backend.layer_idx}"] += 1
-                entry = None
-                if self.record_bridge:
-                    source = topk if packet.source is None else packet.source
-                    entry = {
-                        "layer": backend.layer_idx,
-                        "local_layer": backend.get_local_layer_idx(metadata),
-                        "topk": source.clone(),
-                        "metadata": metadata,
-                    }
-                out = bridge(backend, topk, metadata, packet)
-                if entry is not None:
-                    entry["global"] = out
-                    self.bridge.append(entry)
-                return out
-
-            dispatch._REGISTRY["rope_kv_append"]["remap_topk_to_global"] = bridge_hook
-
-    def _count_packets(self, name, label):
-        original = dispatch.get_impl("absorbed_query", name)
-
-        def hook(mla, *args, **kwargs):
-            packet = original(mla, *args, **kwargs)
-            if packet is not None:
-                self.counts[f"{label}/layer{mla.layer_idx}"] += 1
-            return packet
-
-        dispatch._REGISTRY["absorbed_query"][name] = hook
-
-
-class FmhaCapture:
-    """Records each attention call: library, layer, output and (optionally) input clones.
-
-    Only the input clones add kernels, and they run before the attention; the
-    output is kept by reference (which also keeps its memory from being reused)
-    and copied after the step, so nothing is inserted between the attention
-    kernel and its consumer.
-    """
-
-    def __init__(self):
-        self.calls = []
-        self.clone_inputs = False
-        self.sync_after_sparse_decode = False
-        for cls in (SparseMlaDecodeFmha, FallbackFmha):
-            original = cls.forward
-
-            def forward(fmha, q, k, v, md, args, _original=original, _cls=cls):
-                sparse = args.sparse_runtime_params
-                entry = {
-                    "library": _cls.__name__,
-                    "layer": fmha.attn.layer_idx,
-                    "generation": md.num_contexts == 0,
-                    "output": args.output,
-                    "indices_ref": sparse.sparse_attn_indices,
-                }
-                if self.clone_inputs:
-                    entry.update(
-                        quant_q=args.quant_q_buffer.clone(),
-                        bmm1=args.mla_bmm1_scale.clone(),
-                        bmm2=args.mla_bmm2_scale.clone(),
-                        indices=sparse.sparse_attn_indices.clone(),
-                        kv_lens=md.kv_lens_cuda_runtime[md.num_contexts :].clone(),
-                    )
-                _original(fmha, q, k, v, md, args)
-                if self.sync_after_sparse_decode and _cls is SparseMlaDecodeFmha:
-                    torch.cuda.synchronize()
-                self.calls.append(entry)
-
-            cls.forward = forward
-
-
-def local_topk(ctx_len: int, seed: int) -> torch.Tensor:
-    """Distinct local ids per row: the 32 most recent visible tokens plus random older ones."""
-    gen = torch.Generator(device="cpu").manual_seed(seed)
-    topk = torch.full((ROWS, TOP_K), -1, dtype=torch.int32)
-    for row in range(ROWS):
-        visible = ctx_len + row + 1
-        recent = torch.arange(visible - 32, visible)
-        rest = torch.randperm(visible - 32, generator=gen)[: TOP_K - 32]
-        topk[row] = torch.cat([recent, rest]).to(torch.int32)
-    return topk.cuda()
-
-
-def summarize_kernels(kernels, steps, layers):
-    names = Counter(k["name"] for k in kernels)
-    neighbours = attention_neighbours(kernels)
-    sparse_kernels = [a for a in neighbours if a["kernel"] == "sparse_mla_decode"]
-    return {
-        "steps": steps,
-        "layers": layers,
-        "sparse_mla_decode": sum(n for name, n in names.items() if SPARSE_MLA_KERNEL in name),
-        "native_sparse_fmha": sum(n for name, n in names.items() if is_native_sparse_fmha(name)),
-        "rope_append_bridge": sum(n for name, n in names.items() if "topkBridgeKernel" in name),
-        "rope_append_rope": sum(
-            n for name, n in names.items() if "ropeAppendScatterKernel" in name
-        ),
-        "native_remap": sum(n for name, n in names.items() if "convertReqIndexToGlobal" in name),
-        "top_kernels": [[name[:120], n] for name, n in names.most_common(48)],
-        "attention": neighbours,
-        "successors": Counter(a["successor"] for a in neighbours).most_common(),
-        "predecessors": Counter(a["predecessor"] for a in neighbours).most_common(),
-        "sparse_mla_successor_early": sum(a["successor_early"] for a in sparse_kernels),
-        "sparse_mla_started_before_predecessor_end": sum(
-            a["started_before_predecessor_end"] for a in sparse_kernels
-        ),
-        "sparse_mla_duration_us": sorted(a["duration_us"] for a in sparse_kernels),
-        "sparse_mla_successor_lead_us": sorted(
-            a["successor_lead_us"] for a in sparse_kernels if a["successor_lead_us"] is not None
-        ),
-        "w4_predecessor_lead_us": sorted(
-            a["predecessor_lead_us"] for a in sparse_kernels if a["predecessor_lead_us"] is not None
-        ),
-    }
-
-
-def drive(
-    stack, engagement, layers, ctx, replays, profiled_replays, topk_mode, features, out_path
-) -> None:
-    """Run, record and save every step of one process (see the module docstring)."""
-    # The prefill detached every tail; attach what the adoption built (None: native tail).
-    for layer, attn in stack.attn.items():
-        attn._attn_tail = stack.attn_tail_tails[layer]
-    full_layer = layers[0]
-    indexer = stack.attn[full_layer].mqa.indexer
-    assert indexer is not None, f"layer {full_layer} must run its own indexer"
-    shared_layers = [layer for layer in layers if stack.attn[layer].mqa.indexer is None]
-
-    hidden = [{layer: stack.gen_inputs(layer, step=s) for layer in layers} for s in (0, 1)]
-    topks = [local_topk(ctx, 77), local_topk(ctx, 78)]
-    static_hidden = {layer: hidden[0][layer].clone() for layer in layers}
-    static_topk = topks[0].clone()
-    if topk_mode == "fixed":
-        indexer.forward_from_projected = Mock(return_value=static_topk)
-    pos = stack.position_ids()
-    no_ar = AllReduceParams(enable_allreduce=False)
-    pool_rows = stack.kv_cache_manager.get_unique_primary_pool().view(torch.uint8)
-    pool_rows = pool_rows.reshape(-1, HEAD_DIM)
-    capture = FmhaCapture()
-
-    def load(index):
-        for layer in layers:
-            static_hidden[layer].copy_(hidden[index][layer])
-        static_topk.copy_(topks[index])
-
-    def step(index):
-        load(index)
-        with stack.scope():
-            return {
-                layer: stack.forward(layer, static_hidden[layer], pos, no_ar) for layer in layers
-            }
-
-    def generation_calls(start):
-        calls = [c for c in capture.calls[start:] if c["generation"]]
-        assert [c["layer"] for c in calls] == layers, [c["layer"] for c in calls]
-        return calls
-
-    # Warm-up (JIT compilation, autotuning, lazy buffers) with both input sets.
-    for index in (0, 1):
-        step(index)
-    torch.cuda.synchronize()
-
-    result = {"features": features, "topk_mode": topk_mode, "ctx": ctx, "layers": layers}
-    result["shared_layers"] = shared_layers
-    # Eager steps with the FMHA inputs recorded (clones before the attention).
-    capture.clone_inputs = True
-    engagement.record_bridge = True
-    result["inputs"] = []
-    bridge_pairs = []
-    for index in (0, 1):
-        start, bridge_start = len(capture.calls), len(engagement.bridge)
-        outs = step(index)
-        torch.cuda.synchronize()
-        per_layer = {}
-        for call in generation_calls(start):
-            idx = call["indices"]
-            per_layer[call["layer"]] = {
-                "library": call["library"],
-                "quant_q": call["quant_q"].cpu(),
-                "bmm1": call["bmm1"].cpu(),
-                "bmm2": call["bmm2"].cpu(),
-                "indices": idx.cpu(),
-                "kv_lens": call["kv_lens"].cpu(),
-                "latent": call["output"].clone().cpu(),
-                "module_output": outs[call["layer"]].clone().cpu(),
-                "rows": pool_rows[idx.long().clamp_min(0).reshape(-1)]
-                .reshape(ROWS, TOP_K, -1)
-                .cpu(),
-            }
-            for entry in engagement.bridge[bridge_start:]:
-                if entry["layer"] == call["layer"]:
-                    bridge_pairs.append((index, entry, idx, call["kv_lens"]))
-        result["inputs"].append({"input": index, "layers": per_layer})
-    capture.clone_inputs = False
-    engagement.record_bridge = False
-
-    def plain_steps(label):
-        records = []
-        for index in (0, 1):
-            start = len(capture.calls)
-            outs = step(index)
-            torch.cuda.synchronize()
-            records.append(
-                {
-                    "input": index,
-                    "layers": {
-                        c["layer"]: {
-                            "library": c["library"],
-                            "latent": c["output"].clone().cpu(),
-                            "indices": c["indices_ref"].clone().cpu(),
-                            "module_output": outs[c["layer"]].clone().cpu(),
-                        }
-                        for c in generation_calls(start)
-                    },
-                }
-            )
-        result[label] = records
-
-    plain_steps("pdl")
-    capture.sync_after_sparse_decode = True
-    plain_steps("serialized")
-    capture.sync_after_sparse_decode = False
-
-    # One captured verify step (both layers), replayed over alternating input sets.
-    step(0)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    start = len(capture.calls)
-    with stack.scope():
-        with torch.cuda.graph(graph):
-            graph_outs = {
-                layer: stack.forward(layer, static_hidden[layer], pos, no_ar) for layer in layers
-            }
-    graph_calls = generation_calls(start)
-
-    def replay(index):
-        load(index)
-        graph.replay()
-
-    result["graph"] = []
-    for it in range(replays):
-        index = it % 2
-        replay(index)
-        torch.cuda.synchronize()
-        result["graph"].append(
-            {
-                "input": index,
-                "layers": {
-                    c["layer"]: {
-                        "latent": c["output"].clone().cpu(),
-                        "indices": c["indices_ref"].clone().cpu(),
-                        "module_output": graph_outs[c["layer"]].clone().cpu(),
-                    }
-                    for c in graph_calls
-                },
-            }
-        )
-
-    # CUPTI kernel records: two eager steps without recording kernels, then graph replays.
-    eager_kernels = kernel_trace(lambda: [step(i) for i in (0, 1)])
-    graph_kernels = kernel_trace(lambda: [replay(i % 2) for i in range(profiled_replays)])
-    result["kernels"] = {
-        "eager": summarize_kernels(eager_kernels, 2, len(layers)),
-        "graph": summarize_kernels(graph_kernels, profiled_replays, len(layers)),
-    }
-
-    # rope_kv_append bridge contract, checked after every run so the native remap used as the
-    # reference cannot change the state of the steps above.
-    result["rope_append_bridge"] = []
-    for index, entry, received, kv_lens in bridge_pairs:
-        report = rope_append_bridge_report(entry, kv_lens, ROWS)
-        report["input"] = index
-        report["attention_received_bridge_output"] = bool(
-            torch.equal(received, entry["global"].to(received.device))
-        )
-        result["rope_append_bridge"].append(report)
-
-    tails = {}
-    for layer in layers:
-        tail = stack.attn_tail_tails[layer]
-        tails[layer] = (
-            None
-            if tail is None
-            else {
-                "engaged_calls": tail.engaged_calls,
-                "declined_calls": tail.declined_calls,
-                "static_reason": tail.static_reason,
-                "into_hop_published_calls": tail.into_hop_published_calls,
-                "config": dataclasses.asdict(tail.config),
-            }
-        )
-    sparse_decode_accounting = sparse_mla.engagement_snapshot()
-    result["engagement"] = {
-        "counts": dict(engagement.counts),
-        "attn_tail_tails": tails,
-        "sparse_mla_decode": {
-            "static": {str(k): v for k, v in sparse_decode_accounting["static"].items()},
-            "engaged_calls": {
-                str(k): v for k, v in sparse_decode_accounting["engaged_calls"].items()
-            },
-            "declined_selections": sparse_decode_accounting["declined_selections"],
-        },
-        "capability": {f"{f}/{key}": v for (f, key), v in dispatch._CAPABILITY_LOGGED.items()},
-    }
-    result["libraries"] = sorted({c["library"] for c in capture.calls if c["generation"]})
-    torch.save(result, out_path)
-    summary = {
-        "features": features,
-        "topk_mode": topk_mode,
-        "libraries": result["libraries"],
-        "engagement": result["engagement"],
-        "kernels": {
-            kind: {k: v for k, v in s.items() if k not in ("attention", "top_kernels")}
-            for kind, s in result["kernels"].items()
-        },
-        "rope_append_bridge": result["rope_append_bridge"],
-    }
-    print("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_SUMMARY " + json.dumps(summary, default=str))
-
-
-def main(out_path: str) -> None:
-    plan_testlib.select_from_env(dispatch)
-    ctx = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_CTX", "32768"))
-    layers = [
-        int(x)
-        for x in os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_LAYERS", "2,3").split(",")
-    ]
-    replays = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_REPLAYS", "8"))
-    profiled_replays = int(
-        os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_PROFILED_REPLAYS", "6")
-    )
-    topk_mode = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_TOPK", "fixed")
-    assert topk_mode in ("fixed", "indexer"), topk_mode
-    ckpt = os.environ["GLM52_NVFP4_CKPT"]
-    torch.manual_seed(0)
-    features = sorted(dispatch.enabled_features())
-    engagement = Engagement()
-
-    stack = real_mla_harness.AttentionStack(ckpt, layers, tp_rank=0, ctx=ctx, rows=ROWS)
-    # GEN steps run under inference mode, as the model engine runs them (absorbed_query's
-    # indexer projection updates views of its outputs in place).
-    with torch.inference_mode():
-        drive(
-            stack, engagement, layers, ctx, replays, profiled_replays, topk_mode, features, out_path
-        )
-    stack.shutdown()
-
-
-if __name__ == "__main__":
-    main(sys.argv[1])
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py
deleted file mode 100644
index 98dc9b455fc8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_adapter.py
+++ /dev/null
@@ -1,100 +0,0 @@
-# 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.
-"""Checks of the sparse-MLA decode library declaration on the DSA backend.
-
-These run on a GPU only on request (opt-in:
-``GLM_MOE_DSA_SPARSE_MLA_DECODE_ADAPTER_APPLIED=1``). Each check starts a fresh interpreter that
-installs the plan of the fast paths it is given (``plan_testlib.FAST_PATHS_ENV``). With
-sparse_mla_decode running, the layer declares the library ahead of the registry's, as ``attention.GlmMLA.adopt``
-does, and rebuilds its FMHA libraries, as ``create_weights`` does; the registry is unchanged.
-"""
-
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-import plan_testlib  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_ADAPTER_APPLIED") != "1",
-    reason="opt-in GPU check",
-)
-
-_HERE = Path(__file__).resolve().parent
-
-_FORWARD_SCRIPT = r"""
-import json, os, sys
-sys.path.insert(0, sys.argv[1])
-sys.path.insert(0, os.path.dirname(sys.argv[1]))
-import conftest  # builds the dev op when needed
-import plan_testlib
-import torch
-from sparse_mla_harness import DsaGenHarness
-from tensorrt_llm._torch.attention.backends.fmha import registry
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (
-    dispatch,
-)
-plan_testlib.select_from_env(dispatch)
-h = DsaGenHarness([20000], seq_len_q=6, seed=5)
-if dispatch.is_enabled("sparse_mla_decode"):
-    for layer in h.layers:
-        layer.prepended_fmha_lib_classes = (dispatch.get_impl("sparse_mla_decode", "fmha_lib"),)
-        layer.update_quant_config(layer.quant_config)
-try:
-    topk = h.random_topk(seed=6, overlap=0.7, recent=32)
-    step = h.make_step(0, seed=7)
-    out = h.forward(0, step, topk)
-    libs = [type(lib).__name__ for lib in h.layers[0]._fmha_manager.fmha_libs]
-    print(json.dumps({"registry": list(registry.FMHA_LIBS), "layer_libs": libs,
-                      "selected": h.selected_library(0),
-                      "out_finite": bool(torch.isfinite(out).all())}))
-finally:
-    h.shutdown()
-"""
-
-
-def _run(features: str) -> dict:
-    env = dict(os.environ)
-    env[plan_testlib.FAST_PATHS_ENV] = features
-    proc = subprocess.run(
-        [sys.executable, "-c", _FORWARD_SCRIPT, str(_HERE)],
-        env=env,
-        capture_output=True,
-        text=True,
-        check=False,
-    )
-    assert proc.returncode == 0, proc.stdout[-4000:] + proc.stderr[-4000:]
-    return json.loads(proc.stdout.strip().splitlines()[-1])
-
-
-def test_default_off_keeps_native_libraries():
-    result = _run("")
-    assert "glm_moe_dsa_sparse_mla_decode" not in result["registry"]
-    assert "SparseMlaDecodeFmha" not in result["layer_libs"]
-    assert result["selected"] == ["FallbackFmha"]
-
-
-def test_sparse_mla_selected_when_enabled():
-    result = _run("sparse_mla_decode")
-    assert "glm_moe_dsa_sparse_mla_decode" not in result["registry"]
-    assert result["layer_libs"][0] == "SparseMlaDecodeFmha"
-    assert result["layer_libs"].count("SparseMlaDecodeFmha") == 1
-    assert result["selected"] == ["SparseMlaDecodeFmha"]
-    assert result["out_finite"]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py
deleted file mode 100644
index 7c9aca43c2b0..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_composition.py
+++ /dev/null
@@ -1,213 +0,0 @@
-# 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.
-"""sparse_mla_decode composed with the packages around it at the MLA / DSA call sites.
-
-Each package set runs ``sparse_mla_case.py`` in a fresh interpreter, which installs the
-plan of that set, on the same deterministic layer and inputs:
-
-* producers (absorbed_query, absorbed_query + shared_indexer_absorbed_query, rope_kv_append) feed
-  the query and indices; for every set ``X`` the FMHA inputs of ``X`` and ``X + sparse_mla_decode``
-  must be bit-identical, and the latent output of each library is checked against an FP64 oracle
-  with the frozen error budget;
-* the consumer (attention_output_tail; attention_output_tail_prefetch where integrated) may start
-  before sparse_mla_decode finishes (sparse_mla_decode triggers its dependents before its gathers
-  complete): within each run the eager step, the same step with a device synchronization right after
-  the sparse_mla_decode launch, and replays of a captured graph over alternating inputs must give
-  bit-identical latent and module outputs whenever they attended the same
-  global indices.
-
-Every set runs with a fixed generation top-k (deterministic indices, so every
-replay is compared bit for bit); sets with rope_kv_append also run with the unmodified
-indexer, so that rope_kv_append's top-k bridge, where it engages, produces the indices
-sparse_mla_decode reads (the native indexer may return its selection in another order from run
-to run, so replays are compared only where the indices match). The capability
-decisions and reasons the packages log are recorded.
-
-Runs only where the installed build integrates the packages (an integration
-image); a package that declines the layer is reported, not failed.
-"""
-
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-from sparse_mla_harness import error_stats, oracle_attention, record
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-import plan_testlib  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION") != "1"
-    or not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available(),
-    reason="set GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION=1 on an SM10x integration build with the "
-    "sparse_mla_decode kernel",
-)
-
-_CASE = Path(__file__).resolve().parent / "sparse_mla_case.py"
-# Producers of the query/indices (absorbed_query, absorbed_query + shared_indexer_absorbed_query,
-# rope_kv_append RoPE/append and top-k bridge), the consumer tail (attention_output_tail,
-# attention_output_tail + attention_output_tail_prefetch) and one set with all of them; (set, top-k
-# mode).
-NEIGHBOUR_SETS = [
-    ("", "fixed"),
-    ("absorbed_query", "fixed"),
-    ("absorbed_query,shared_indexer_absorbed_query", "fixed"),
-    ("rope_kv_append", "fixed"),
-    ("rope_kv_append", "indexer"),
-    ("attention_output_tail", "fixed"),
-    ("attention_output_tail,attention_output_tail_prefetch", "fixed"),
-    ("absorbed_query,shared_indexer_absorbed_query,rope_kv_append,attention_output_tail", "fixed"),
-    (
-        "absorbed_query,shared_indexer_absorbed_query,rope_kv_append,attention_output_tail",
-        "indexer",
-    ),
-]
-ROWS = 6
-TOP_K = 2048
-
-
-def _run(features: str, topk_mode: str, out_dir: Path) -> dict:
-    out = out_dir / f"mla_{features.replace(',', '_') or 'native'}_{topk_mode}.pt"
-    env = dict(os.environ)
-    env[plan_testlib.FAST_PATHS_ENV] = features
-    env["GLM_MOE_DSA_SPARSE_MLA_DECODE_MLA_TOPK"] = topk_mode
-    env["TLLM_LOG_LEVEL"] = "INFO"
-    proc = subprocess.run(
-        [sys.executable, str(_CASE), str(out)], env=env, capture_output=True, text=True, check=False
-    )
-    assert proc.returncode == 0, proc.stdout[-6000:] + proc.stderr[-6000:]
-    result = torch.load(out, weights_only=False)
-    log = (proc.stdout + proc.stderr).splitlines()
-    result["capability_log"] = sorted(
-        {line[line.index("[glm_moe_dsa]") :] for line in log if "[glm_moe_dsa]" in line}
-    )
-    return result
-
-
-def _oracle(entry: dict) -> torch.Tensor:
-    rows = entry["rows"].cuda().reshape(ROWS * TOP_K, -1)
-    indices = entry["indices"].cuda()
-    local = torch.arange(ROWS * TOP_K, device="cuda", dtype=torch.int32).view(ROWS, TOP_K)
-    remapped = torch.where(indices >= 0, local, torch.full_like(local, -1))
-    kv_len = int(entry["kv_lens"][0])
-    limits = [kv_len - (ROWS - 1 - r) for r in range(ROWS)]
-    return oracle_attention(
-        entry["quant_q"].cuda(),
-        rows,
-        remapped,
-        float(entry["bmm1"][0]),
-        float(entry["bmm2"][0]),
-        list_limit=limits,
-    )
-
-
-def _same(a: torch.Tensor, b: torch.Tensor) -> bool:
-    return a.shape == b.shape and torch.equal(a.view(torch.uint8), b.view(torch.uint8))
-
-
-@pytest.mark.parametrize("neighbours,topk_mode", NEIGHBOUR_SETS)
-def test_sparse_mla_with_neighbours(neighbours, topk_mode, tmp_path):
-    out_dir = Path(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS", tmp_path))
-    with_sparse_mla = ",".join([f for f in neighbours.split(",") if f] + ["sparse_mla_decode"])
-    base = _run(neighbours, topk_mode, out_dir)
-    cand = _run(with_sparse_mla, topk_mode, out_dir)
-    assert base["libraries"] == ["FallbackFmha"], base["libraries"]
-    assert cand["libraries"] == ["SparseMlaDecodeFmha"], cand["libraries"]
-
-    summary = {
-        "neighbours": neighbours or "native",
-        "topk_mode": cand["topk_mode"],
-        "capability_without_sparse_mla": base["capability"],
-        "capability_with_sparse_mla": cand["capability"],
-        "capability_log_with_sparse_mla": cand["capability_log"],
-        "steps": [],
-    }
-    assert base["topk_mode"] == cand["topk_mode"]
-    for b, c in zip(base["eager"], cand["eager"]):
-        # The producers do not depend on the consumer: identical FMHA inputs.
-        for key in ("quant_q", "bmm1", "bmm2", "kv_lens"):
-            assert _same(b[key], c[key]), (neighbours, key)
-        same_order = _same(b["indices"], c["indices"])
-        if not same_order:
-            # Only an unmodified indexer may return the same selection in another order.
-            assert cand["topk_mode"] == "indexer", neighbours
-            assert torch.equal(b["indices"].sort(dim=-1).values, c["indices"].sort(dim=-1).values)
-        else:
-            assert _same(b["rows"], c["rows"]), (neighbours, "rows")
-        assert not bool((c["indices"] < 0).any()), "an index below the causal length is -1"
-        oracle = _oracle(c)
-        native_err = error_stats(b["latent"].cuda().double(), oracle)
-        sparse_mla_err = error_stats(c["latent"].cuda().double(), oracle)
-        module_diff = error_stats(c["module_output"].double(), b["module_output"].double())
-        summary["steps"].append(
-            {
-                "input": c["input"],
-                "indices_same_order": same_order,
-                "native_latent_vs_oracle": native_err,
-                "sparse_mla_latent_vs_oracle": sparse_mla_err,
-                "module_output_sparse_mla_vs_native": module_diff,
-            }
-        )
-        assert sparse_mla_err["nonfinite"] == native_err["nonfinite"] == 0
-        assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
-            sparse_mla_err,
-            native_err,
-        )
-        assert sparse_mla_err["max_abs"] <= 2.0 * native_err["max_abs"] + 1e-3, (
-            sparse_mla_err,
-            native_err,
-        )
-        assert module_diff["nonfinite"] == 0 and module_diff["nrms"] < 0.05, module_diff
-
-    # Programmatic overlap with the next kernel and graph replay change nothing.
-    for run, label in ((base, "without_sparse_mla"), (cand, "with_sparse_mla")):
-        eager = {e["input"]: e for e in run["eager"]}
-        mismatches = []
-        index_changes = []
-        for kind in ("serialized", "graph"):
-            for i, entry in enumerate(run[kind]):
-                ref = eager[entry["input"]]
-                if not _same(entry["indices"], ref["indices"]):
-                    index_changes.append((kind, i))
-                    continue
-                for key in ("latent", "module_output"):
-                    if not _same(entry[key], ref[key]):
-                        mismatches.append((kind, i, key))
-        summary[f"replay_mismatches_{label}"] = mismatches
-        summary[f"replay_index_changes_{label}"] = index_changes
-    record(
-        "sparse_mla_composition.jsonl",
-        "mla_composition",
-        json.loads(json.dumps(summary, default=str)),
-    )
-    # sparse_mla_decode is deterministic, so with it every replay and the serialized step that
-    # attended the same indices must match exactly; a fixed top-k never changes the indices.
-    assert not summary["replay_mismatches_with_sparse_mla"], summary[
-        "replay_mismatches_with_sparse_mla"
-    ]
-    if cand["topk_mode"] == "fixed":
-        assert not summary["replay_index_changes_with_sparse_mla"], summary[
-            "replay_index_changes_with_sparse_mla"
-        ]
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py
deleted file mode 100644
index 0231c119f646..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_dispatch_accounting.py
+++ /dev/null
@@ -1,123 +0,0 @@
-# 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.
-"""Capability logging, error behaviour and engagement accounting of sparse_mla_decode.
-
-* The static decision is logged once per layer, not once per process.
-* A layer inside the sparse_mla_decode domain whose kernel is missing from the build or whose
-  device cannot launch it raises instead of falling back to native on one rank;
-  a layer outside the domain still takes the native library.
-* ``engagement_snapshot`` counts the sparse_mla_decode calls per layer and the declined
-  selections per reason.
-"""
-
-import pytest
-import torch
-from sparse_mla_harness import DsaGenHarness
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
-    SparseMlaDecodeFmha,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available(),
-    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
-)
-
-
-def test_capability_logged_per_layer():
-    harness = DsaGenHarness([256], seq_len_q=1, num_layers=3, seed=141)
-    try:
-        for layer in harness.layers:
-            assert SparseMlaDecodeFmha.is_available(layer)
-        snapshot = sparse_mla.engagement_snapshot()
-        for layer in harness.layers:
-            assert snapshot["static"][layer.layer_idx] == ""
-            assert (
-                dispatch._CAPABILITY_LOGGED[("sparse_mla_decode", f"layer{layer.layer_idx}")]
-                is True
-            )
-    finally:
-        harness.shutdown()
-
-
-def test_selected_sparse_mla_raises_instead_of_falling_back(monkeypatch):
-    harness = DsaGenHarness([256], seq_len_q=1, seed=142)
-    try:
-        layer = harness.layers[0]
-        monkeypatch.setattr(sparse_mla, "device_supported", lambda device: False)
-        with pytest.raises(RuntimeError, match="cannot launch its 16-CTA cluster kernel"):
-            SparseMlaDecodeFmha.is_available(layer)
-        monkeypatch.setattr(sparse_mla, "op_available", lambda: False)
-        with pytest.raises(RuntimeError, match="does not register"):
-            SparseMlaDecodeFmha.is_available(layer)
-    finally:
-        harness.shutdown()
-
-
-def test_layer_outside_domain_takes_native():
-    harness = DsaGenHarness([256], seq_len_q=1, kv_dtype=torch.bfloat16, seed=143)
-    try:
-        layer = harness.layers[0]
-        assert not SparseMlaDecodeFmha.is_available(layer)
-        assert sparse_mla.engagement_snapshot()["static"][layer.layer_idx] == "KV cache is not FP8"
-    finally:
-        harness.shutdown()
-
-
-def test_engagement_counts_per_layer_and_declines():
-    harness = DsaGenHarness([20000], seq_len_q=6, num_layers=2, seed=144)
-    try:
-        before = sparse_mla.engagement_snapshot()["engaged_calls"]
-        topk = harness.random_topk(seed=146)
-        for layer_idx in (0, 1):
-            harness.select_library(layer_idx, use_sparse_mla=True)
-            for _ in range(2):
-                harness.forward(layer_idx, harness.make_step(layer_idx, seed=147), topk)
-            assert harness.selected_library(layer_idx) == ["SparseMlaDecodeFmha"]
-        torch.cuda.synchronize()
-        after = sparse_mla.engagement_snapshot()["engaged_calls"]
-        for layer in harness.layers:
-            idx = layer.layer_idx
-            assert after.get(idx, 0) - before.get(idx, 0) == 2, (before, after)
-            assert dispatch._CAPABILITY_LOGGED[("sparse_mla_decode", f"call:layer{idx}")] is True
-    finally:
-        harness.shutdown()
-    other = DsaGenHarness([16384], seq_len_q=6, index_topk=1024, seed=148)
-    try:
-        declined_before = sparse_mla.engagement_snapshot()["declined_selections"]
-        gen = torch.Generator(device="cpu").manual_seed(149)
-        topk = (
-            torch.stack(
-                [
-                    torch.randperm(other.causal_len(r), generator=gen)[:1024]
-                    for r in range(other.num_rows)
-                ]
-            )
-            .to(torch.int32)
-            .cuda()
-        )
-        other.select_library(0, use_sparse_mla=True)
-        other.forward(0, other.make_step(0, seed=150), topk)
-        assert "SparseMlaDecodeFmha" not in other.selected_library(0)
-        declined_after = sparse_mla.engagement_snapshot()["declined_selections"]
-        assert declined_after.get("top-k 1024", 0) == declined_before.get("top-k 1024", 0) + 1
-    finally:
-        other.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py
deleted file mode 100644
index 1d473fb57b00..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_native_semantics.py
+++ /dev/null
@@ -1,171 +0,0 @@
-# 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.
-"""Characterization of the native sparse MLA generation kernel (FP8 KV).
-
-The native trtllm-gen kernel ships as a binary, so its rounding and masking
-rules are measured here against FP64 models: which probability rounding it
-feeds to P*V, whether it ignores index-list positions past a row's causal KV
-length, what it writes for a row without valid entries, and whether repeated
-indices count once per occurrence. sparse_mla_decode must reproduce whatever these tests find.
-"""
-
-import json
-import os
-
-import pytest
-import torch
-from sparse_mla_harness import (
-    TOP_K,
-    DsaGenHarness,
-    bf16_round,
-    e4m3_round,
-    error_stats,
-    fp16_round,
-    oracle_attention,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="needs an SM10x GPU",
-)
-
-PROB_MODELS = {
-    "exact": {},
-    "fp16_split256_step128": dict(prob_round=fp16_round, split=256, step=128),
-    "bf16_split256_step128": dict(prob_round=bf16_round, split=256, step=128),
-    "e4m3x1_split256_step128": dict(prob_round=e4m3_round(1.0), split=256, step=128),
-    "e4m3x448_split256_step128": dict(prob_round=e4m3_round(448.0), split=256, step=128),
-    "e4m3x256_split256_step128": dict(prob_round=e4m3_round(256.0), split=256, step=128),
-    "e4m3x1_split2048_step128": dict(prob_round=e4m3_round(1.0), split=2048, step=128),
-    "e4m3x448_split2048_step128": dict(prob_round=e4m3_round(448.0), split=2048, step=128),
-    "e4m3x1_split128_step128": dict(prob_round=e4m3_round(1.0), split=128, step=128),
-    "e4m3x448_split128_step128": dict(prob_round=e4m3_round(448.0), split=128, step=128),
-}
-
-
-def record(case: str, payload: dict) -> None:
-    line = json.dumps({"case": case, **payload}, sort_keys=True)
-    print(line)
-    results_dir = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS")
-    if results_dir:
-        with open(os.path.join(results_dir, "native_semantics.jsonl"), "a") as f:
-            f.write(line + "\n")
-
-
-def run_native(harness: DsaGenHarness, topk: torch.Tensor, seed: int = 1, q_std: float = 1.0):
-    step = harness.make_step(0, seed=seed, q_std=q_std)
-    harness.select_library(0, use_sparse_mla=False)
-    out = harness.forward(0, step, topk).clone()
-    global_idx = harness.global_indices(0, topk)
-    return step, out, global_idx
-
-
-@pytest.mark.parametrize("ctx_len,q_std", [(32768, 1.0), (32768, 3.0), (32768, 0.3), (4096, 1.0)])
-def test_native_probability_rounding(ctx_len, q_std):
-    harness = DsaGenHarness([ctx_len], seq_len_q=6, seed=11)
-    try:
-        topk = harness.random_topk(seed=3, overlap=0.7, recent=32)
-        step, native, global_idx = run_native(harness, topk, seed=5, q_std=q_std)
-        b1 = float(step.bmm1_scale[0])
-        b2 = float(step.bmm2_scale[0])
-        native = native.double()
-        distances = {}
-        for name, kwargs in PROB_MODELS.items():
-            ref = oracle_attention(step.quant_q, harness.pool_rows(), global_idx, b1, b2, **kwargs)
-            ref_bf16 = ref.to(torch.bfloat16).double()
-            stats = error_stats(native, ref_bf16)
-            stats["exact_match_frac"] = float(
-                (native.reshape(ref_bf16.shape) == ref_bf16).double().mean().item()
-            )
-            distances[name] = stats
-        best = min(distances, key=lambda n: distances[n]["nrms"])
-        record(
-            "prob_rounding",
-            {
-                "ctx_len": ctx_len,
-                "q_std": q_std,
-                "bmm1_scale": [float(x) for x in step.bmm1_scale.tolist()],
-                "bmm2_scale": b2,
-                "best": best,
-                "distances": distances,
-            },
-        )
-        assert distances["exact"]["nonfinite"] == 0
-    finally:
-        harness.shutdown()
-
-
-def test_native_kv_length_clamp():
-    """Valid entries placed after a row's causal length in the index list."""
-    harness = DsaGenHarness([100], seq_len_q=6, seed=12)
-    try:
-        topk = harness.random_topk(seed=4)
-        extra_positions = [700, 1500, TOP_K - 1]
-        for row in range(harness.num_rows):
-            for pos in extra_positions:
-                topk[row, pos] = 3 + row
-        step, native, global_idx = run_native(harness, topk, seed=6)
-        b1 = float(step.bmm1_scale[0])
-        b2 = float(step.bmm2_scale[0])
-        limits = [harness.causal_len(r) for r in range(harness.num_rows)]
-        clamped = oracle_attention(
-            step.quant_q, harness.pool_rows(), global_idx, b1, b2, list_limit=limits
-        )
-        unclamped = oracle_attention(step.quant_q, harness.pool_rows(), global_idx, b1, b2)
-        record(
-            "kv_length_clamp",
-            {
-                "vs_clamped": error_stats(native.double(), clamped),
-                "vs_unclamped": error_stats(native.double(), unclamped),
-                "limits": limits,
-            },
-        )
-    finally:
-        harness.shutdown()
-
-
-def test_native_invalid_and_duplicate_entries():
-    harness = DsaGenHarness([8192], seq_len_q=6, seed=13)
-    try:
-        topk = harness.random_topk(seed=5)
-        # Row 1: no valid entry. Row 2: invalid entries interleaved with valid ones.
-        # Row 3: two tokens repeated with unequal multiplicity.
-        topk[1, :] = -1
-        topk[2, ::3] = -1
-        topk[3, :1500] = 17
-        topk[3, 1500:] = 4242
-        step, native, global_idx = run_native(harness, topk, seed=7)
-        b1 = float(step.bmm1_scale[0])
-        b2 = float(step.bmm2_scale[0])
-        exact = oracle_attention(step.quant_q, harness.pool_rows(), global_idx, b1, b2)
-        dedup_idx = global_idx.clone()
-        dedup_idx[3, 1:1500] = -1
-        dedup_idx[3, 1501:] = -1
-        dedup = oracle_attention(step.quant_q, harness.pool_rows(), dedup_idx, b1, b2)
-        row1 = native[1].double()
-        record(
-            "invalid_and_duplicates",
-            {
-                "all_invalid_row": {
-                    "nonfinite": int((~torch.isfinite(row1)).sum().item()),
-                    "max_abs": float(row1.nan_to_num().abs().max().item()),
-                },
-                "interleaved_invalid_row": error_stats(native[2].double(), exact[2]),
-                "duplicates_counted": error_stats(native[3].double(), exact[3]),
-                "duplicates_dedup": error_stats(native[3].double(), dedup[3]),
-            },
-        )
-    finally:
-        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py
deleted file mode 100644
index 2604a240510e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_pdl_stress.py
+++ /dev/null
@@ -1,479 +0,0 @@
-# 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.
-"""Delayed-producer and delayed-consumer stress of the sparse_mla_decode kernel under PDL.
-
-The sparse_mla_decode kernel initializes its mbarriers and arrives on the cluster barrier
-before ``griddepcontrol.wait``, issues ``griddepcontrol.launch_dependents``
-before its compute, gathers pool rows with bulk copies and pushes partial
-results to its cluster peers with ``st.async``. This test checks that none of
-that reads an input early or writes the output early.
-
-Every iteration is a chain without a plain kernel boundary:
-
-  producer (PDL)  -> sparse_mla_decode (PDL) -> consumer (PDL) -> next producer ...
-
-The producer lets its dependents launch at once, waits for its predecessor,
-spins, and only then writes everything sparse_mla_decode reads: the FP8 query, the index
-list, the KV lengths, the bmm scales and a block of pool rows the index list
-points into. The consumer lets its dependents launch at once, waits for sparse_mla_decode,
-optionally spins, and copies the output out. Consecutive iterations use
-different inputs, and each input set alternates between two versions of its
-pool rows, so an early read of any input, a late write of the output, or a
-stale cache line gives a different result. Each output must be bit-identical
-to the output of the same inputs computed in isolation (sparse_mla_decode is deterministic;
-this is checked too); the isolated outputs are checked against the FP64 oracle.
-
-Variants: eager chain with and without PDL on the helper kernels, the chain
-captured in a CUDA graph and replayed, and the eager chain while a second
-stream saturates HBM. ``%globaltimer`` stamps of the helper kernels and a
-CUPTI trace show whether the early launches happened.
-"""
-
-import os
-import statistics
-
-import pytest
-import torch
-from sparse_mla_harness import E4M3_MAX, error_stats, oracle_attention, record
-from sparse_mla_pdl_helpers import load_pdl_helpers
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available(),
-    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
-)
-
-NUM_HEADS = 8
-HEAD_DIM = 576
-KV_LORA_RANK = 512
-TOP_K = 2048
-STATIC_ROWS = 120_000
-OWNED_ROWS = 512
-NUM_SETS = 6
-ITERS = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STRESS_ITERS", "400"))
-GRAPH_REPLAYS = int(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_STRESS_GRAPH_REPLAYS", "8"))
-PRODUCER_SPINS_NS = (0, 1000, 5000, 20000)
-CONSUMER_SPINS_NS = (0, 0, 3000, 10000)
-RESULTS = "sparse_mla_stress.jsonl"
-# compute-sanitizer serializes and slows the instrumented kernel; the early-launch evidence is
-# then not expected, only the correctness of every iteration.
-UNDER_SANITIZER = os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_UNDER_SANITIZER") == "1"
-
-
-def _e4m3_bytes(shape, std: float, gen: torch.Generator) -> torch.Tensor:
-    values = torch.randn(*shape, device="cuda", generator=gen) * std
-    return values.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8)
-
-
-class StressCase:
-    """Live buffers read by sparse_mla_decode, staged input sets and their isolated reference outputs."""
-
-    def __init__(self, ctx_lens, seq_len_q: int, seed: int):
-        self.seq_len_q = seq_len_q
-        self.num_seqs = len(ctx_lens)
-        assert self.num_seqs <= 4
-        self.rows = self.num_seqs * seq_len_q
-        gen = torch.Generator(device="cuda").manual_seed(seed)
-        cpu_gen = torch.Generator(device="cpu").manual_seed(seed)
-        num_pool_rows = STATIC_ROWS + NUM_SETS * OWNED_ROWS
-        self.pool = _e4m3_bytes((num_pool_rows, HEAD_DIM), 1.0, gen)
-        self.q = torch.zeros(self.rows, NUM_HEADS * HEAD_DIM, dtype=torch.uint8, device="cuda")
-        self.indices = torch.zeros(self.rows, TOP_K, dtype=torch.int32, device="cuda")
-        self.kv_lens = torch.zeros(4, dtype=torch.int32, device="cuda")
-        self.scales = torch.zeros(4, dtype=torch.float32, device="cuda")
-        self.out = torch.zeros(
-            self.rows, NUM_HEADS * KV_LORA_RANK, dtype=torch.bfloat16, device="cuda"
-        )
-        self.sets = []
-        for k in range(NUM_SETS):
-            kv_lens = [c + seq_len_q + 3 * k for c in ctx_lens]
-            owned_base = STATIC_ROWS + k * OWNED_ROWS
-            limits = []
-            idx = torch.empty(self.rows, TOP_K, dtype=torch.int64)
-            for row in range(self.rows):
-                seq, pos = divmod(row, seq_len_q)
-                usable = kv_lens[seq] - (seq_len_q - 1 - pos)
-                limits.append(usable)
-                entries = torch.randint(0, STATIC_ROWS, (TOP_K,), generator=cpu_gen)
-                owned = torch.rand(TOP_K, generator=cpu_gen) < 0.125
-                entries[owned] = owned_base + torch.randint(
-                    0, OWNED_ROWS, (int(owned.sum()),), generator=cpu_gen
-                )
-                negative = torch.rand(TOP_K, generator=cpu_gen) < 0.02
-                entries[negative] = -1
-                entries[64:96] = entries[0:32]  # repeated rows count per occurrence
-                idx[row] = entries
-            kv_tensor = torch.zeros(4, dtype=torch.int32)
-            kv_tensor[: self.num_seqs] = torch.tensor(kv_lens, dtype=torch.int32)
-            softmax_scale = 0.0625 * (1.0 + 0.05 * k)
-            scales = torch.tensor(
-                [softmax_scale, softmax_scale * 1.4426950408889634, 1.0 + 0.01 * k, 0.0],
-                dtype=torch.float32,
-            )
-            variants = [
-                _e4m3_bytes((OWNED_ROWS, HEAD_DIM), 2.0, gen),
-                _e4m3_bytes((OWNED_ROWS, HEAD_DIM), 0.5, gen),
-            ]
-            self.sets.append(
-                {
-                    "q": _e4m3_bytes((self.rows, NUM_HEADS * HEAD_DIM), 1.0 + 0.25 * k, gen),
-                    "indices": idx.to(torch.int32).cuda(),
-                    "kv_lens": kv_tensor.cuda(),
-                    "scales": scales.cuda(),
-                    "owned_base": owned_base,
-                    "variants": variants,
-                    "limits": limits,
-                }
-            )
-        self.segments = {}
-        for k, s in enumerate(self.sets):
-            for v in range(2):
-                owned = self.pool[s["owned_base"] : s["owned_base"] + OWNED_ROWS]
-                pairs = [
-                    (self.q, s["q"]),
-                    (self.indices, s["indices"]),
-                    (self.kv_lens, s["kv_lens"]),
-                    (self.scales, s["scales"]),
-                    (owned, s["variants"][v]),
-                ]
-                for dst, src in pairs:
-                    assert dst.is_contiguous() and src.is_contiguous() and dst.nbytes == src.nbytes
-                    assert dst.nbytes % 16 == 0
-                self.segments[(k, v)] = tuple(
-                    torch.tensor(values, dtype=torch.int64, device="cuda")
-                    for values in (
-                        [d.data_ptr() for d, _ in pairs],
-                        [s_.data_ptr() for _, s_ in pairs],
-                        [d.nbytes // 16 for d, _ in pairs],
-                    )
-                )
-        self.references = {}
-
-    def apply_plain(self, k: int, v: int) -> None:
-        s = self.sets[k]
-        self.q.copy_(s["q"])
-        self.indices.copy_(s["indices"])
-        self.kv_lens.copy_(s["kv_lens"])
-        self.scales.copy_(s["scales"])
-        self.pool[s["owned_base"] : s["owned_base"] + OWNED_ROWS].copy_(s["variants"][v])
-
-    def launch_sparse_mla(self) -> None:
-        sparse_mla.sparse_mla_decode(
-            self.q.view(self.rows, NUM_HEADS, HEAD_DIM),
-            self.pool.data_ptr(),
-            self.indices,
-            self.scales[0:2],
-            self.scales[2:3],
-            self.kv_lens,
-            self.seq_len_q,
-            self.out,
-        )
-
-    def build_references(self, case: str) -> None:
-        """Isolated outputs of every (set, version); determinism and FP64-oracle checks."""
-        for k in range(NUM_SETS):
-            for v in range(2):
-                runs = []
-                for _ in range(2):
-                    self.apply_plain(k, v)
-                    torch.cuda.synchronize()
-                    self.out.fill_(float("nan"))
-                    self.launch_sparse_mla()
-                    torch.cuda.synchronize()
-                    runs.append(self.out.clone())
-                assert torch.equal(runs[0], runs[1]), (
-                    "sparse_mla_decode output is not deterministic"
-                )
-                s = self.sets[k]
-                oracle = oracle_attention(
-                    self.q,
-                    self.pool.view(torch.float8_e4m3fn),
-                    self.indices,
-                    float(s["scales"][0]),
-                    float(s["scales"][2]),
-                    list_limit=s["limits"],
-                )
-                err = error_stats(runs[0].double(), oracle)
-                record(
-                    RESULTS,
-                    f"{case}:reference",
-                    {"set": k, "version": v, "sparse_mla_vs_oracle": err},
-                )
-                assert err["nonfinite"] == 0
-                assert err["nrms"] < 0.05, err
-                assert torch.allclose(
-                    runs[0].double().reshape(oracle.shape), oracle, atol=0.12, rtol=0.01
-                )
-                self.references[(k, v)] = runs[0]
-        # The two versions of each set's pool rows must change the output, or a stale read of
-        # those rows would go unnoticed.
-        for k in range(NUM_SETS):
-            assert not torch.equal(self.references[(k, 0)], self.references[(k, 1)])
-
-
-def _schedule(iterations: int, seed: int):
-    gen = torch.Generator(device="cpu").manual_seed(seed)
-    plan = []
-    for i in range(iterations):
-        k = i % NUM_SETS
-        v = (i // NUM_SETS) % 2
-        p = PRODUCER_SPINS_NS[int(torch.randint(0, len(PRODUCER_SPINS_NS), (1,), generator=gen))]
-        c = CONSUMER_SPINS_NS[int(torch.randint(0, len(CONSUMER_SPINS_NS), (1,), generator=gen))]
-        plan.append((k, v, p, c))
-    return plan
-
-
-# A spin kernel ahead of an eager chain (about 0.2 s), so that the host enqueues the whole
-# chain before the GPU reaches it and the GPU, not the host, paces the chain.
-GATE_CYCLES = int(4e8)
-
-
-def _enqueue_chain(case: StressCase, helpers, plan, outs, stamps, use_pdl: bool) -> None:
-    for i, (k, v, producer_spin, consumer_spin) in enumerate(plan):
-        dsts, srcs, counts = case.segments[(k, v)]
-        helpers.delayed_scatter(dsts, srcs, counts, producer_spin, stamps[i, 0:3], use_pdl, 64)
-        case.launch_sparse_mla()
-        helpers.pdl_consumer(outs[i], case.out, consumer_spin, stamps[i, 3:5], use_pdl)
-
-
-def _mismatches(case: StressCase, plan, outs) -> list:
-    bad = []
-    for i, (k, v, _, _) in enumerate(plan):
-        if not torch.equal(outs[i], case.references[(k, v)]):
-            diff = (outs[i].float() - case.references[(k, v)].float()).abs()
-            bad.append(
-                {
-                    "iteration": i,
-                    "set": k,
-                    "version": v,
-                    "max_abs": float(diff.nan_to_num(1e30).max()),
-                }
-            )
-    return bad
-
-
-def _stamp_summary(stamps: torch.Tensor, plan) -> dict:
-    """Early-launch evidence from the helper kernels' ``%globaltimer`` stamps.
-
-    A consumer that started well before its ``griddepcontrol.wait`` returned was
-    launched while sparse_mla_decode was still running, i.e. through sparse_mla_decode's early
-    ``launch_dependents``.
-    """
-    s = stamps.cpu().tolist()
-    consumer_wait = [row[4] - row[3] for row in s]
-    producer_wait = [row[1] - row[0] for row in s]
-    return {
-        "consumer_wait_ns_median": statistics.median(consumer_wait),
-        "consumer_started_before_sparse_mla_done": sum(1 for x in consumer_wait if x > 500)
-        / len(s),
-        "producer_wait_ns_median": statistics.median(producer_wait),
-        "iterations": len(plan),
-    }
-
-
-@pytest.fixture(scope="module")
-def helpers():
-    return load_pdl_helpers()
-
-
-CASES = {
-    "one_seq_r6": ([40000], 6),
-    "three_seqs_r2": ([3000, 100, 9000], 2),
-}
-
-
-@pytest.mark.parametrize("use_pdl", [True, False], ids=["pdl", "no_pdl"])
-@pytest.mark.parametrize("case_name", list(CASES))
-def test_sparse_mla_delayed_producer_consumer_eager(helpers, case_name, use_pdl):
-    ctx_lens, seq_len_q = CASES[case_name]
-    case = StressCase(ctx_lens, seq_len_q, seed=1)
-    case.build_references(case_name)
-    plan = _schedule(ITERS, seed=2)
-    outs = torch.full(
-        (ITERS, case.rows, NUM_HEADS * KV_LORA_RANK),
-        float("nan"),
-        dtype=torch.bfloat16,
-        device="cuda",
-    )
-    stamps = torch.zeros(ITERS, 8, dtype=torch.int64, device="cuda")
-    # Start from a different set than the first iteration's, so iteration 0 also overwrites.
-    # Warm-up chain (module loading, launch attributes), then the measured chain.
-    _enqueue_chain(case, helpers, plan[:2], outs, stamps, use_pdl)
-    case.apply_plain(NUM_SETS - 1, 1)
-    outs.fill_(float("nan"))
-    stamps.zero_()
-    torch.cuda.synchronize()
-    torch.cuda._sleep(GATE_CYCLES)
-    _enqueue_chain(case, helpers, plan, outs, stamps, use_pdl)
-    torch.cuda.synchronize()
-    bad = _mismatches(case, plan, outs)
-    summary = _stamp_summary(stamps, plan)
-    record(
-        RESULTS,
-        f"{case_name}:eager",
-        {
-            "use_pdl_helpers": use_pdl,
-            "trtllm_enable_pdl_env": os.environ.get("TRTLLM_ENABLE_PDL"),
-            "mismatches": len(bad),
-            "first_mismatches": bad[:5],
-            **summary,
-        },
-    )
-    assert not bad, bad[:5]
-    if use_pdl and os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0" and not UNDER_SANITIZER:
-        # The chain must actually exercise the early launch, or this test proves nothing.
-        assert summary["consumer_started_before_sparse_mla_done"] > 0.25, summary
-
-
-@pytest.mark.parametrize("case_name", list(CASES))
-def test_sparse_mla_delayed_producer_consumer_graph(helpers, case_name):
-    ctx_lens, seq_len_q = CASES[case_name]
-    case = StressCase(ctx_lens, seq_len_q, seed=3)
-    case.build_references(f"{case_name}:graph")
-    length = 4 * NUM_SETS  # the version of each set alternates across replays too
-    plan = _schedule(length, seed=4)
-    outs = torch.full(
-        (length, case.rows, NUM_HEADS * KV_LORA_RANK),
-        float("nan"),
-        dtype=torch.bfloat16,
-        device="cuda",
-    )
-    stamps = torch.zeros(length, 8, dtype=torch.int64, device="cuda")
-    case.apply_plain(NUM_SETS - 1, 1)
-    # Warm up outside capture (attribute setup, lazy loading).
-    _enqueue_chain(case, helpers, plan[:2], outs, stamps, True)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        _enqueue_chain(case, helpers, plan, outs, stamps, True)
-    total_bad = 0
-    summaries = []
-    for replay in range(GRAPH_REPLAYS):
-        outs.fill_(float("nan"))
-        stamps.zero_()
-        graph.replay()
-        torch.cuda.synchronize()
-        bad = _mismatches(case, plan, outs)
-        total_bad += len(bad)
-        summaries.append(_stamp_summary(stamps, plan))
-        assert not bad, (replay, bad[:5])
-    record(
-        RESULTS,
-        f"{case_name}:graph",
-        {
-            "replays": GRAPH_REPLAYS,
-            "chain_length": length,
-            "mismatches": total_bad,
-            "consumer_started_before_sparse_mla_done": [
-                x["consumer_started_before_sparse_mla_done"] for x in summaries
-            ],
-            "consumer_wait_ns_median": [x["consumer_wait_ns_median"] for x in summaries],
-        },
-    )
-
-
-def test_sparse_mla_delayed_producer_consumer_under_hbm_contention(helpers):
-    ctx_lens, seq_len_q = CASES["one_seq_r6"]
-    case = StressCase(ctx_lens, seq_len_q, seed=5)
-    case.build_references("contention")
-    plan = _schedule(ITERS, seed=6)
-    outs = torch.full(
-        (ITERS, case.rows, NUM_HEADS * KV_LORA_RANK),
-        float("nan"),
-        dtype=torch.bfloat16,
-        device="cuda",
-    )
-    stamps = torch.zeros(ITERS, 8, dtype=torch.int64, device="cuda")
-    big_src = torch.ones(256 * 1024 * 1024, dtype=torch.float32, device="cuda")
-    big_dst = torch.empty_like(big_src)
-    case.apply_plain(NUM_SETS - 1, 1)
-    torch.cuda.synchronize()
-    side = torch.cuda.Stream()
-    with torch.cuda.stream(side):
-        for _ in range(160):
-            big_dst.copy_(big_src)
-    torch.cuda._sleep(GATE_CYCLES // 4)
-    _enqueue_chain(case, helpers, plan, outs, stamps, True)
-    torch.cuda.synchronize()
-    bad = _mismatches(case, plan, outs)
-    record(
-        RESULTS,
-        "contention:eager",
-        {"mismatches": len(bad), "first_mismatches": bad[:5], **_stamp_summary(stamps, plan)},
-    )
-    assert not bad, bad[:5]
-
-
-@pytest.mark.skipif(UNDER_SANITIZER, reason="CUPTI tracing is not combined with compute-sanitizer")
-def test_sparse_mla_early_launch_trace(helpers):
-    """CUPTI kernel intervals: sparse_mla_decode becomes resident before its producer finishes."""
-    ctx_lens, seq_len_q = CASES["one_seq_r6"]
-    case = StressCase(ctx_lens, seq_len_q, seed=7)
-    case.build_references("trace")
-    iterations = 48
-    plan = [(i % NUM_SETS, (i // NUM_SETS) % 2, 20000, 3000) for i in range(iterations)]
-    outs = torch.full(
-        (iterations, case.rows, NUM_HEADS * KV_LORA_RANK),
-        float("nan"),
-        dtype=torch.bfloat16,
-        device="cuda",
-    )
-    stamps = torch.zeros(iterations, 8, dtype=torch.int64, device="cuda")
-    case.apply_plain(NUM_SETS - 1, 1)
-    torch.cuda.synchronize()
-    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
-        torch.cuda._sleep(GATE_CYCLES // 4)
-        _enqueue_chain(case, helpers, plan, outs, stamps, True)
-        torch.cuda.synchronize()
-    assert not _mismatches(case, plan, outs)
-    kernels = sorted(
-        (
-            (e.time_range.start, e.time_range.end, e.name)
-            for e in prof.events()
-            if e.device_type == torch.autograd.DeviceType.CUDA
-        ),
-        key=lambda x: x[0],
-    )
-    producers = [k for k in kernels if "delayedScatter" in k[2]]
-    sparse_kernels = [k for k in kernels if "sparseMlaDecode" in k[2]]
-    consumers = [k for k in kernels if "pdlConsumer" in k[2]]
-    count = min(len(producers), len(sparse_kernels), len(consumers))
-    sparse_mla_early = sum(1 for i in range(count) if sparse_kernels[i][0] < producers[i][1])
-    consumer_early = sum(1 for i in range(count) if consumers[i][0] < sparse_kernels[i][1])
-    record(
-        RESULTS,
-        "early_launch_trace",
-        {
-            "kernels": {
-                "producer": len(producers),
-                "sparse_mla_decode": len(sparse_kernels),
-                "consumer": len(consumers),
-            },
-            "sparse_mla_started_before_producer_end": sparse_mla_early / max(count, 1),
-            "consumer_started_before_sparse_mla_end": consumer_early / max(count, 1),
-            "sparse_mla_start_minus_producer_end_us_median": statistics.median(
-                sparse_kernels[i][0] - producers[i][1] for i in range(count)
-            )
-            if count
-            else None,
-            **_stamp_summary(stamps, plan),
-        },
-    )
-    assert count == iterations
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py
deleted file mode 100644
index 7f9ee560001d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_ragged_verify.py
+++ /dev/null
@@ -1,99 +0,0 @@
-# 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.
-"""sparse_mla_decode under DSA ragged verification (per-request verify widths).
-
-With ragged verification enabled in the generation metadata, ``TrtllmAttention``
-selects and runs the FMHA library inside the token-major presentation: one row
-per query token, ``seq_len_q`` 1, and that row's causal KV length. sparse_mla_decode engages
-there; these tests check it against the native library and the FP64 oracle with
-the frozen error budget of ``test_sparse_mla_sparse_mla.py``, check the lengths sparse_mla_decode
-receives, and check that a ragged step whose widths are uniform gives the same
-sparse_mla_decode output as the uniform layout.
-"""
-
-import pytest
-import torch
-from sparse_mla_harness import DsaGenHarness, compare
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available(),
-    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
-)
-
-
-@pytest.mark.parametrize(
-    "ctx_lens,verify_lens",
-    [
-        ([8192, 300, 40000], [6, 2, 4]),  # mixed widths, one short request
-        ([391000], [3]),  # one request with a partial verify window
-        ([100, 2043], [5, 6]),  # rows straddling the top-k length
-        ([20000, 20000, 20000], [1, 6, 1]),
-    ],
-)
-def test_sparse_mla_ragged_verification_matches_native(ctx_lens, verify_lens):
-    harness = DsaGenHarness(
-        ctx_lens, seq_len_q=max(verify_lens), seed=121, ragged_verify_lens=verify_lens
-    )
-    try:
-        assert harness.metadata.token_major_gen_view() is not None
-        topk = harness.random_topk(seed=122, overlap=0.7, recent=32)
-        step = harness.make_step(0, seed=123)
-        with harness.capture_sparse_mla_inputs() as calls:
-            compare(
-                harness,
-                step,
-                topk,
-                "ragged_verify_matches_native",
-                ctx_lens=ctx_lens,
-                verify_lens=verify_lens,
-            )
-        assert len(calls) == 1, len(calls)
-        call = calls[0]
-        # One presented row per query token, each with its causal length.
-        assert call["ragged"] and call["seq_len_q"] == 1, call
-        expected = torch.tensor(harness.list_limits(), dtype=torch.int32)
-        assert torch.equal(call["seq_lens_kv"].cpu(), expected), (call["seq_lens_kv"], expected)
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_ragged_presentation_equals_uniform():
-    """Uniform widths through the ragged presentation give the uniform layout's sparse_mla_decode output."""
-    ctx_lens = [32768, 5000]
-    outputs = {}
-    queries = {}
-    for ragged in (False, True):
-        harness = DsaGenHarness(
-            ctx_lens, seq_len_q=6, seed=131, ragged_verify_lens=[6, 6] if ragged else None
-        )
-        try:
-            assert (harness.metadata.token_major_gen_view() is not None) == ragged
-            topk = harness.random_topk(seed=132, overlap=0.7, recent=32)
-            step = harness.make_step(0, seed=133)
-            harness.select_library(0, use_sparse_mla=True)
-            outputs[ragged] = harness.forward(0, step, topk).clone()
-            assert harness.selected_library(0) == ["SparseMlaDecodeFmha"]
-            queries[ragged] = (step.quant_q.clone(), step.bmm1_scale.clone())
-        finally:
-            harness.shutdown()
-    torch.testing.assert_close(queries[True][0], queries[False][0], rtol=0, atol=0)
-    torch.testing.assert_close(queries[True][1], queries[False][1], rtol=0, atol=0)
-    torch.testing.assert_close(outputs[True], outputs[False], rtol=0, atol=0)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py
deleted file mode 100644
index c9788cb648a4..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_real_weights.py
+++ /dev/null
@@ -1,404 +0,0 @@
-# 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.
-"""sparse_mla_decode against native on real GLM-5.2 layer-0 attention inputs (see ``sparse_mla_real_fixture``).
-
-The context latents are written by the native context step and the generation
-rows by the native ``mla_rope_generation``. Two sources of the top-k:
-
-* ``test_sparse_mla_real_layer0``: native indexer projections + a reference FP32
-  logit evaluation and exact top-k (not the native top-k kernels);
-* ``test_sparse_mla_real_layer0_native_indexer``: the unmodified native indexer
-  (native indexer K cache from the context and generation rows, native paged
-  MQA logits and decode top-k) under the served GEN indexer configuration
-  and under the CUDA radix top-k; native and sparse_mla_decode consume the same native
-  selection. This test also checks the premise of sparse_mla_decode's one deliberate
-  difference (no ``-1`` inside a row's causal length) on the native top-k, on
-  the global indices sparse_mla_decode receives, and on the native grouped remap of the same
-  selection. With ``rope_kv_append`` selected, rope_kv_append's fused top-k bridge replaces the
-  native remap and produces the indices sparse_mla_decode reads: every generation call of the bridge
-  must then equal the native remap of the same local top-k bit for bit, emit no
-  ``-1`` below a row's causal length, and be what sparse_mla_decode received.
-
-The same error budget as ``test_sparse_mla_sparse_mla.py`` applies.
-"""
-
-import pytest
-import torch
-from sparse_mla_harness import (
-    DsaGenHarness,
-    capture_rope_append_bridge,
-    check_budget,
-    compare,
-    error_stats,
-    negative_inside_causal,
-    oracle_for,
-    record,
-    rope_append_bridge_report,
-    rope_append_selected,
-)
-from sparse_mla_real_fixture import (
-    checkpoint_dir,
-    indexer_projections,
-    indexer_topk,
-    layer_weight_names,
-    load_indexer_weights,
-    load_tensors,
-    project_layer0,
-    real_token_ids,
-    recent_window_topk,
-    reference_indexer_logits,
-    selection_vs_kernel_scores,
-    selection_vs_reference,
-    topk_contract,
-)
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available()
-    or checkpoint_dir() is None,
-    reason="needs an SM10x GPU, the sparse_mla_decode kernel and the GLM-5.2 NVFP4 checkpoint",
-)
-
-SEQ_LEN_Q = 6
-NATIVE_INDEXER_CONTEXTS = [100, 2043, 8192, 32768, 131072]
-# Long contexts run in their own test (and process). The native context attention step of this
-# harness faults at these lengths, before any generation-phase code runs (one 491K-token step,
-# or 64K-token chunks of a 391K context), so these cases write only the
-# native indexer K cache for the context tokens (native ``_update_k_cache``, 64K-token chunks)
-# and keep random latent rows in the pool: the selection is the native indexer's on real
-# indexer keys and queries, the attended latents are synthetic.
-NATIVE_INDEXER_LONG_CONTEXTS = [391000, 491061]
-CONTEXT_CHUNK_TOKENS = 65536
-# Indexer kernels of the served GEN role and the exact
-# CUDA radix decode top-k (indexerTopK.cu) that the default configuration selects.
-NATIVE_INDEXER_CONFIGS = {
-    "tp8_decode": dict(
-        enable_heuristic_topk=True, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=True
-    ),
-    "cuda_radix": dict(
-        enable_heuristic_topk=False, use_cute_dsl_paged_mqa_logits=False, use_cute_dsl_topk=False
-    ),
-    # Diagnostic only: the served CuTe-DSL paged MQA logits with the exact CUDA radix top-k.
-    "dsl_logits_cuda_radix": dict(
-        enable_heuristic_topk=False, use_cute_dsl_paged_mqa_logits=True, use_cute_dsl_topk=False
-    ),
-}
-
-
-@pytest.fixture(scope="module")
-def layer0():
-    ckpt = checkpoint_dir()
-    max_tokens = max(NATIVE_INDEXER_CONTEXTS + NATIVE_INDEXER_LONG_CONTEXTS) + SEQ_LEN_Q
-    weights = load_tensors(ckpt, layer_weight_names(0))
-    ids = real_token_ids(ckpt, max_tokens)
-    inputs = project_layer0(weights, ids)
-    yield weights, inputs
-    del weights, inputs
-    torch.cuda.empty_cache()
-
-
-@pytest.mark.parametrize("ctx_len", [32768, 131072])
-def test_sparse_mla_real_layer0(layer0, ctx_len):
-    weights, inputs = layer0
-    harness = DsaGenHarness(
-        [ctx_len], seq_len_q=SEQ_LEN_Q, fill_random=False, defer_generation=True
-    )
-    try:
-        load_indexer_weights(harness.layers[0].indexer, weights)
-        harness.run_context(
-            0,
-            inputs.fused_q[:ctx_len],
-            inputs.q_pe[:ctx_len],
-            inputs.latent[:ctx_len],
-            recent_window_topk(ctx_len),
-        )
-        harness.start_generation()
-        rows = slice(ctx_len, ctx_len + SEQ_LEN_Q)
-        step = harness.make_step_from(
-            0, inputs.fused_q[rows], inputs.q_pe[rows], inputs.latent[rows]
-        )
-        topk = indexer_topk(harness.layers[0].indexer, inputs, ctx_len, SEQ_LEN_Q)
-        # Selection structure: overlap between consecutive MTP rows and spread over the context.
-        sets = [set(topk[r].tolist()) for r in range(SEQ_LEN_Q)]
-        overlap = [len(sets[r] & sets[r - 1]) / 2048 for r in range(1, SEQ_LEN_Q)]
-        recent = [sum(1 for x in sets[r] if x >= ctx_len - 2048) / 2048 for r in range(SEQ_LEN_Q)]
-        result = compare(
-            harness,
-            step,
-            topk,
-            "real_layer0_reference_topk",
-            ctx_len=ctx_len,
-            mtp_row_overlap=overlap,
-            recent_2048_fraction=recent,
-        )
-        # Softmax peakedness of the real scores, per row and head (exact FP64 model).
-        global_idx = result["global_idx"]
-        q = step.quant_q.view(SEQ_LEN_Q, 8, 576).view(torch.float8_e4m3fn).double()
-        pool = harness.pool_rows().view(torch.uint8)
-        top_prob = []
-        for r in range(SEQ_LEN_Q):
-            keys = pool[global_idx[r][global_idx[r] >= 0].long()].view(torch.float8_e4m3fn).double()
-            probs = torch.softmax(q[r] @ keys.T * float(step.bmm1_scale[0]), dim=-1)
-            top_prob.append(probs.max(dim=-1).values.mean().item())
-        record(
-            "sparse_mla_correctness.jsonl",
-            "real_layer0_scores",
-            {"ctx_len": ctx_len, "mean_max_prob_per_row": top_prob},
-        )
-        assert error_stats(result["sparse_mla_decode"].double(), result["oracle"])["nonfinite"] == 0
-    finally:
-        harness.shutdown()
-
-
-@pytest.mark.parametrize("ctx_len", NATIVE_INDEXER_CONTEXTS)
-@pytest.mark.parametrize("config", ["tp8_decode", "cuda_radix"])
-def test_sparse_mla_real_layer0_native_indexer(layer0, config, ctx_len):
-    _run_native_indexer_case(layer0, config, ctx_len)
-
-
-@pytest.mark.parametrize("ctx_len", NATIVE_INDEXER_LONG_CONTEXTS)
-@pytest.mark.parametrize("config", ["tp8_decode", "cuda_radix"])
-def test_sparse_mla_real_layer0_native_indexer_long(layer0, config, ctx_len):
-    _run_native_indexer_case(
-        layer0, config, ctx_len, chunk_tokens=CONTEXT_CHUNK_TOKENS, real_latents=False
-    )
-
-
-@pytest.mark.parametrize(
-    "config,seq_len_q,mtp",
-    [
-        ("tp8_decode", 1, 0),
-        ("tp8_decode", 4, 3),
-        ("dsl_logits_cuda_radix", 6, 5),
-        ("dsl_logits_cuda_radix", 1, 0),
-    ],
-)
-def test_sparse_mla_real_layer0_native_indexer_variants(layer0, config, seq_len_q, mtp):
-    """Diagnostic variants at 32K: verify-window width and the logits / top-k kernel pair."""
-    _run_native_indexer_case(layer0, config, 32768, seq_len_q=seq_len_q, mtp=mtp)
-
-
-def _run_native_indexer_case(
-    layer0, config, ctx_len, chunk_tokens=None, real_latents=True, seq_len_q=SEQ_LEN_Q, mtp=None
-):
-    weights, inputs = layer0
-    harness = DsaGenHarness(
-        [ctx_len],
-        seq_len_q=seq_len_q,
-        seed=17,
-        fill_random=not real_latents,
-        defer_generation=True,
-        sparse_overrides=NATIVE_INDEXER_CONFIGS[config],
-        mtp_draft_len=seq_len_q - 1 if mtp is None else mtp,
-    )
-    try:
-        indexer = harness.layers[0].indexer
-        load_indexer_weights(indexer, weights)
-        ctx_proj = indexer_projections(indexer, inputs, 0, ctx_len)
-        harness.run_context(
-            0,
-            inputs.fused_q[:ctx_len],
-            inputs.q_pe[:ctx_len],
-            inputs.latent[:ctx_len],
-            recent_window_topk(ctx_len) if real_latents else None,
-            indexer_k=(ctx_proj[1], ctx_proj[2]),
-            chunk_tokens=chunk_tokens,
-            attention=real_latents,
-        )
-        del ctx_proj
-        harness.start_generation()
-        rows = slice(ctx_len, ctx_len + seq_len_q)
-        step = harness.make_step_from(
-            0, inputs.fused_q[rows], inputs.q_pe[rows], inputs.latent[rows]
-        )
-        gen_proj = indexer_projections(indexer, inputs, ctx_len, ctx_len + seq_len_q)
-        harness.update_indexer_cache(0, gen_proj[1], gen_proj[2])
-
-        harness.select_library(0, use_sparse_mla=False)
-        with capture_rope_append_bridge() as bridge_native:
-            native, topk_native = harness.forward_native_indexer(0, step, gen_proj)
-        native = native.clone()
-        native_topk_inputs = harness.last_topk_inputs
-        assert "SparseMlaDecodeFmha" not in harness.selected_library(0)
-        harness.select_library(0, use_sparse_mla=True)
-        with (
-            harness.capture_sparse_mla_inputs() as sparse_mla_calls,
-            capture_rope_append_bridge() as bridge_sparse_mla,
-        ):
-            sparse_out, topk_sparse_mla = harness.forward_native_indexer(0, step, gen_proj)
-        sparse_out = sparse_out.clone()
-        torch.cuda.synchronize()
-        assert harness.selected_library(0) == ["SparseMlaDecodeFmha"], harness.selected_library(0)
-        assert len(sparse_mla_calls) == 1, len(sparse_mla_calls)
-        received = sparse_mla_calls[0]
-
-        contract = topk_contract(topk_native, ctx_len)
-        reference_logits = reference_indexer_logits(indexer, inputs, ctx_len, seq_len_q)
-        agreement = []
-        for r in range(seq_len_q):
-            native_ids = set(x for x in topk_native[r].tolist() if x >= 0)
-            count = min(2048, reference_logits[r].numel())
-            reference_ids = set(reference_logits[r].topk(count).indices.tolist())
-            agreement.append(len(native_ids & reference_ids) / max(len(reference_ids), 1))
-        same_selection = torch.equal(topk_native, topk_sparse_mla)
-        selection = {
-            "seq_len_q": seq_len_q,
-            "max_draft_tokens": harness.metadata.max_draft_tokens,
-            "dsl_expand": [
-                bool(harness.metadata.expand_for_dsl),
-                harness.metadata.dsl_expand_factor,
-                harness.metadata.dsl_atom,
-            ],
-            "decode_top_k": indexer.top_k.decode_implementation.value,
-            "self_sampling": bool(getattr(indexer, "_use_self_sampling_topk", False)),
-            "same_selection_in_both_runs": same_selection,
-            "rows": contract,
-            "agreement_with_reference_topk": agreement,
-            "native_selection_under_reference_logits": selection_vs_reference(
-                topk_native, reference_logits
-            ),
-            "native_selection_vs_kernel_scores": (
-                selection_vs_kernel_scores(topk_native, native_topk_inputs, reference_logits)
-                if native_topk_inputs is not None
-                else None
-            ),
-        }
-        global_native = harness.global_indices(0, topk_native)
-        oracle = oracle_for(harness, step, global_native)
-        # Producer contract at sparse_mla_decode's input: the global indices it received (the native
-        # remap of its run's selection) and the native grouped remap of the native run's selection.
-        grouped_native = harness.grouped_global_indices(0, topk_native)
-        received_negative = negative_inside_causal(
-            received["indices"], received["seq_lens_kv"], received["seq_len_q"]
-        )
-        grouped_negative = negative_inside_causal(
-            grouped_native, received["seq_lens_kv"], received["seq_len_q"]
-        )
-        selection["sparse_mla_input"] = {
-            "negative_below_causal": received_negative.tolist(),
-            "equals_remap_of_its_selection": torch.equal(
-                received["indices"], harness.global_indices(0, topk_sparse_mla)
-            ),
-            "seq_len_q": received["seq_len_q"],
-        }
-        selection["grouped_remap"] = {
-            "equals_per_layer_remap": torch.equal(grouped_native, global_native),
-            "negative_below_causal": grouped_negative.tolist(),
-        }
-        # rope_kv_append's bridge (when selected) is the producer of the indices of both runs.
-        rope_append_bridge = {
-            "selected": rope_append_selected(),
-            "calls_native_run": len(bridge_native),
-            "calls_sparse_mla_run": len(bridge_sparse_mla),
-            "calls": [
-                rope_append_bridge_report(call, received["seq_lens_kv"], received["seq_len_q"])
-                for call in bridge_native + bridge_sparse_mla
-            ],
-            "sparse_mla_received_bridge_output": (
-                bool(torch.equal(received["indices"], bridge_sparse_mla[0]["global"]))
-                if bridge_sparse_mla
-                else None
-            ),
-        }
-        selection["rope_append_bridge"] = rope_append_bridge
-        if same_selection:
-            check_budget(
-                harness,
-                native,
-                sparse_out,
-                global_native,
-                oracle,
-                "real_layer0_native_indexer",
-                config=config,
-                ctx_len=ctx_len,
-                context_chunk_tokens=chunk_tokens,
-                real_context_latents=real_latents,
-                selection=selection,
-            )
-        else:
-            # The native selection is not deterministic run to run. sparse_mla_decode is checked
-            # twice: on the exact selection the native run consumed (the native indexer output
-            # tensor fed to sparse_mla_decode through the fixed-top-k path) and, below, on the
-            # selection of its own run.
-            selection["sparse_mla_run_selection_overlap"] = [
-                len(set(topk_native[r].tolist()) & set(topk_sparse_mla[r].tolist()))
-                / max(len(set(topk_native[r].tolist())), 1)
-                for r in range(seq_len_q)
-            ]
-            sparse_mla_same = harness.forward(0, step, topk_native).clone()
-            assert harness.selected_library(0) == ["SparseMlaDecodeFmha"]
-            check_budget(
-                harness,
-                native,
-                sparse_mla_same,
-                global_native,
-                oracle,
-                "real_layer0_native_indexer_same_selection",
-                config=config,
-                ctx_len=ctx_len,
-                context_chunk_tokens=chunk_tokens,
-                real_context_latents=real_latents,
-                selection=selection,
-            )
-            global_sparse_mla = harness.global_indices(0, topk_sparse_mla)
-            native_err = error_stats(native.double(), oracle)
-            sparse_mla_err = error_stats(
-                sparse_out.double(), oracle_for(harness, step, global_sparse_mla)
-            )
-            record(
-                "sparse_mla_correctness.jsonl",
-                "real_layer0_native_indexer",
-                {
-                    "config": config,
-                    "ctx_len": ctx_len,
-                    "context_chunk_tokens": chunk_tokens,
-                    "real_context_latents": real_latents,
-                    "selection": selection,
-                    "native_vs_oracle": native_err,
-                    "sparse_mla_vs_oracle": sparse_mla_err,
-                },
-            )
-            assert sparse_mla_err["nonfinite"] == native_err["nonfinite"] == 0
-            assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
-                sparse_mla_err,
-                native_err,
-            )
-        # The premise of sparse_mla_decode's one deliberate difference from native: the native
-        # indexer emits no -1 inside a row's causal length, and neither remap creates one.
-        assert all(row["negative_below_causal"] == 0 for row in contract), contract
-        assert int(received_negative.sum()) == 0, selection["sparse_mla_input"]
-        assert selection["sparse_mla_input"]["equals_remap_of_its_selection"]
-        assert selection["grouped_remap"]["equals_per_layer_remap"]
-        assert int(grouped_negative.sum()) == 0, selection["grouped_remap"]
-        if rope_append_bridge["selected"]:
-            # The bridge served the generation call of both runs (counted per call: the
-            # logged rope_kv_append/bridge decision is the first one, the declined context step).
-            assert (
-                rope_append_bridge["calls_native_run"] == 1
-                and rope_append_bridge["calls_sparse_mla_run"] == 1
-            ), rope_append_bridge
-            for call in rope_append_bridge["calls"]:
-                assert call["equals_native_remap"], rope_append_bridge
-                assert sum(call["negative_below_causal"]) == 0, rope_append_bridge
-            assert rope_append_bridge["sparse_mla_received_bridge_output"], rope_append_bridge
-        else:
-            assert not bridge_native and not bridge_sparse_mla
-    finally:
-        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py
deleted file mode 100644
index c92aba441f83..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_rope_append_bridge.py
+++ /dev/null
@@ -1,195 +0,0 @@
-# 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.
-"""rope_kv_append's fused top-k bridge as the producer of sparse_mla_decode's global indices.
-
-With the plan running ``sparse_mla_decode`` and ``rope_kv_append`` (installed by this test) the DSA
-predictor's local-to-global remap is ``remap_topk_to_global``
-(``trtllm::glm_moe_dsa_topk_index_bridge``), and its output is the index list sparse_mla_decode
-reads. sparse_mla_decode differs from the native kernel on one
-input class: a ``-1`` at an index-list position below the row's causal length (sparse_mla_decode
-skips it, the native kernel reads the pool row that ``-1`` addresses). For every
-generation call the bridge serves, this test asserts that
-
-* the bridge ran (counted per call, not read from the capability log, whose
-  first ``rope_kv_append/bridge`` decision is the declined context step),
-* its global indices equal the native remap of the same local top-k bit for
-  bit and contain no ``-1`` below the row's causal length,
-* sparse_mla_decode received exactly those indices, and sparse_mla_decode stays within the frozen
-  error budget against the FP64 oracle,
-* CUPTI shows rope_kv_append's ``topkBridgeKernel`` and sparse_mla_decode's
-  ``sparseMlaDecodeKernel`` once per generation call, and neither the native remap kernel nor the
-  native sparse FMHA kernel.
-
-A negative control puts ``-1`` inside the causal range of one row: the bridge
-passes it through exactly like the native remap, and the check above reports
-it, so the assertion is not vacuous.
-
-Runs only when both ops are in the build.
-"""
-
-import os
-import sys
-
-import pytest
-import torch
-from sparse_mla_harness import (
-    SPARSE_MLA_KERNEL,
-    DsaGenHarness,
-    capture_rope_append_bridge,
-    check_budget,
-    is_native_sparse_fmha,
-    kernel_trace,
-    negative_inside_causal,
-    oracle_for,
-    record,
-    rope_append_bridge_report,
-)
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import dispatch
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-import plan_testlib  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available()
-    or not hasattr(torch.ops.trtllm, "glm_moe_dsa_topk_index_bridge"),
-    reason="needs an SM10x build with the sparse_mla_decode and rope_kv_append kernels",
-)
-
-
-@pytest.fixture(autouse=True)
-def _plan(monkeypatch):
-    plan_testlib.select(dispatch, "sparse_mla_decode,rope_kv_append", monkeypatch)
-
-
-RESULTS = "sparse_mla_rope_append_bridge.jsonl"
-CASES = {
-    "ctx100_r6": ([100], 6),  # every row ends in -1 entries past its causal length
-    "ctx2042_r6": ([2042], 6),  # the last row sees exactly 2048 tokens
-    "ctx2043_r6": ([2043], 6),  # rows straddle 2048 visible tokens
-    "two_seqs_r6": ([5000, 300], 6),  # the second sequence has -1 tails
-    "four_seqs_r2": ([20000, 20000, 20000, 20000], 2),
-    "ctx32768_r6": ([32768], 6),  # target verify, MTP K = 5
-    "ctx131072_r1": ([131072], 1),  # draft steps 1-4
-}
-
-
-def _generation_kernels(harness, step, topk):
-    kernels = kernel_trace(lambda: harness.forward(0, step, topk))
-    names = [k["name"] for k in kernels]
-    return {
-        "rope_append_bridge": sum("topkBridgeKernel" in n for n in names),
-        "native_remap": sum("convertReqIndexToGlobal" in n for n in names),
-        "sparse_mla_decode": sum(SPARSE_MLA_KERNEL in n for n in names),
-        "native_sparse_fmha": sum(is_native_sparse_fmha(n) for n in names),
-    }
-
-
-@pytest.mark.parametrize("case", list(CASES))
-def test_rope_append_bridge_feeds_sparse_mla(case):
-    ctx_lens, seq_len_q = CASES[case]
-    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=31, mtp_draft_len=seq_len_q - 1)
-    try:
-        step = harness.make_step(0, seed=32)
-        topk = harness.random_topk(seed=33, overlap=0.7, recent=32)
-        seq_lens_kv = torch.tensor(
-            [c + seq_len_q for c in ctx_lens], dtype=torch.int32, device="cuda"
-        )
-        with (
-            capture_rope_append_bridge() as bridge,
-            harness.capture_sparse_mla_inputs() as sparse_mla_calls,
-        ):
-            harness.select_library(0, use_sparse_mla=False)
-            native = harness.forward(0, step, topk).clone()
-            harness.select_library(0, use_sparse_mla=True)
-            sparse_out = harness.forward(0, step, topk).clone()
-            torch.cuda.synchronize()
-        assert harness.selected_library(0) == ["SparseMlaDecodeFmha"]
-        # One bridge call per generation forward (native arm, then sparse_mla_decode arm).
-        assert len(bridge) == 2, len(bridge)
-        assert len(sparse_mla_calls) == 1, len(sparse_mla_calls)
-        received = sparse_mla_calls[0]
-        assert torch.equal(received["seq_lens_kv"], seq_lens_kv)
-        reports = [rope_append_bridge_report(call, seq_lens_kv, seq_len_q) for call in bridge]
-        same_as_bridge = torch.equal(received["indices"], bridge[1]["global"])
-        global_idx = harness.global_indices(0, topk)
-        oracle = oracle_for(harness, step, global_idx)
-        kernels = _generation_kernels(harness, step, topk)
-        check_budget(
-            harness,
-            native,
-            sparse_out,
-            global_idx,
-            oracle,
-            "rope_append_bridge_feeds_sparse_mla",
-            RESULTS,
-            shape=case,
-            bridge=reports,
-            sparse_mla_received_bridge_output=same_as_bridge,
-            generation_kernels=kernels,
-        )
-        for report in reports:
-            assert report["equals_native_remap"], report
-            assert sum(report["negative_below_causal"]) == 0, report
-        assert same_as_bridge
-        assert kernels == {
-            "rope_append_bridge": 1,
-            "native_remap": 0,
-            "sparse_mla_decode": 1,
-            "native_sparse_fmha": 0,
-        }, kernels
-    finally:
-        harness.shutdown()
-
-
-def test_rope_append_bridge_negative_control():
-    """``-1`` injected inside the causal range: passed through as by the native remap, detected."""
-    ctx_lens, seq_len_q = [32768], 6
-    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=41, mtp_draft_len=seq_len_q - 1)
-    try:
-        step = harness.make_step(0, seed=42)
-        topk = harness.random_topk(seed=43, overlap=0.7, recent=32)
-        topk[2, ::3] = -1
-        injected = int((topk[2] < 0).sum().item())
-        seq_lens_kv = torch.tensor([ctx_lens[0] + seq_len_q], dtype=torch.int32, device="cuda")
-        with (
-            capture_rope_append_bridge() as bridge,
-            harness.capture_sparse_mla_inputs() as sparse_mla_calls,
-        ):
-            harness.select_library(0, use_sparse_mla=True)
-            harness.forward(0, step, topk)
-            torch.cuda.synchronize()
-        assert len(bridge) == 1 and len(sparse_mla_calls) == 1
-        report = rope_append_bridge_report(bridge[0], seq_lens_kv, seq_len_q)
-        received = negative_inside_causal(
-            sparse_mla_calls[0]["indices"], seq_lens_kv, seq_len_q
-        ).tolist()
-        record(
-            RESULTS,
-            "rope_append_bridge_negative_control",
-            {"injected_row2": injected, "bridge": report, "sparse_mla_received_negative": received},
-        )
-        assert report["equals_native_remap"], report
-        expected = [0] * harness.num_rows
-        expected[2] = injected
-        assert report["negative_below_causal"] == expected, report
-        assert received == expected
-    finally:
-        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py
deleted file mode 100644
index df9ca44f0711..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sanitizer_cases.py
+++ /dev/null
@@ -1,150 +0,0 @@
-# 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.
-"""Small sparse_mla_decode cases for compute-sanitizer (memcheck, racecheck, synccheck, initcheck).
-
-Run under ``compute-sanitizer --tool <tool> --kernel-name kns=sparseMlaDecode``
-so that only the sparse_mla_decode kernel is instrumented. The cases cover every code path of
-the kernel at shapes small enough for instrumented execution: short rows with
-fewer usable positions than the top-k (zero-filled gather slots), rows that
-straddle 2048 usable positions, full rows, R = 6 and R = 1, several sequences,
-empty rows, negative and repeated ids, and CUDA-graph capture plus replay with
-rewritten inputs. Each case also checks the output against the native library
-and the FP64 oracle with the frozen error budget of ``test_sparse_mla_sparse_mla.py``.
-"""
-
-import pytest
-import torch
-from sparse_mla_harness import DsaGenHarness, compare, error_stats, oracle_attention, record
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available(),
-    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
-)
-
-RESULTS = "sparse_mla_sanitizer_cases.jsonl"
-
-
-@pytest.mark.parametrize(
-    "ctx_lens,seq_len_q",
-    [
-        ([100], 6),
-        ([100], 1),
-        ([2043], 6),
-        ([2043], 1),
-        ([8192], 6),
-        ([8192], 1),
-        ([3000, 100, 9000], 2),
-    ],
-)
-def test_sparse_mla_small_matches_native(ctx_lens, seq_len_q):
-    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=11)
-    try:
-        topk = harness.random_topk(seed=12, overlap=0.7, recent=32)
-        step = harness.make_step(0, seed=13)
-        compare(
-            harness,
-            step,
-            topk,
-            "small_matches_native",
-            RESULTS,
-            ctx_lens=ctx_lens,
-            seq_len_q=seq_len_q,
-        )
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_small_empty_negative_and_repeated():
-    harness = DsaGenHarness([4096, 40], seq_len_q=6, seed=14)
-    try:
-        topk = harness.random_topk(seed=15)
-        topk[1, :] = -1
-        topk[3, :1500] = 17
-        topk[3, 1500:] = 4242
-        for row in range(6, 12):
-            for pos in (300, 1500, 2047):
-                topk[row, pos] = 3
-        step = harness.make_step(0, seed=16)
-        compare(harness, step, topk, "small_empty_negative_repeated", RESULTS)
-        # Negative ids inside the causal range: sparse_mla_decode skips them. Only
-        # sparse_mla_decode's output is checked here (the native kernel reads the row that -1
-        # addresses).
-        topk = harness.random_topk(seed=17)
-        topk[2, ::3] = -1
-        topk[5, :1024] = -1
-        harness.select_library(0, use_sparse_mla=True)
-        sparse_out = harness.forward(0, step, topk).clone()
-        global_idx = harness.global_indices(0, topk)
-        oracle = oracle_attention(
-            step.quant_q,
-            harness.pool_rows(),
-            global_idx,
-            float(step.bmm1_scale[0]),
-            float(step.bmm2_scale[0]),
-            list_limit=harness.list_limits(),
-        )
-        err = error_stats(sparse_out.double(), oracle)
-        record(RESULTS, "small_negative_inside_range_sparse_mla", {"sparse_mla_vs_oracle": err})
-        assert err["nonfinite"] == 0 and err["nrms"] < 0.05, err
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_small_graph_replay():
-    harness = DsaGenHarness([8192], seq_len_q=6, seed=18)
-    try:
-        step = harness.make_step(0, seed=19)
-        topk = harness.random_topk(seed=20, overlap=0.7, recent=32)
-        harness.select_library(0, use_sparse_mla=True)
-        harness.forward(0, step, topk)
-        torch.cuda.synchronize()
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            out = harness.forward(0, step, topk)
-        pool_rows = harness.pool_rows()
-        for it in range(3):
-            topk.copy_(harness.random_topk(seed=300 + it, overlap=0.5, recent=16))
-            global_idx = harness.global_indices(0, topk)
-            rows = global_idx[global_idx >= 0].unique().long()
-            fresh = (
-                (torch.randn(rows.numel(), pool_rows.shape[1], device="cuda") * (1 + it))
-                .clamp(-448, 448)
-                .to(torch.float8_e4m3fn)
-            )
-            pool_rows.view(torch.uint8)[rows] = fresh.view(torch.uint8)
-            fresh_step = harness.make_step(0, seed=400 + it)
-            for name in ("fused_q", "q_pe", "latent", "quant_q", "bmm1_scale", "bmm2_scale"):
-                getattr(step, name).copy_(getattr(fresh_step, name))
-            graph.replay()
-            torch.cuda.synchronize()
-            oracle = oracle_attention(
-                step.quant_q,
-                pool_rows,
-                global_idx,
-                float(step.bmm1_scale[0]),
-                float(step.bmm2_scale[0]),
-                list_limit=harness.list_limits(),
-            )
-            err = error_stats(out.double(), oracle)
-            record(RESULTS, "small_graph_replay", {"iteration": it, "sparse_mla_vs_oracle": err})
-            assert err["nonfinite"] == 0 and err["nrms"] < 0.05, err
-    finally:
-        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py
deleted file mode 100644
index 7a64a7e43bdf..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_sparse_mla.py
+++ /dev/null
@@ -1,285 +0,0 @@
-# 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.
-"""Correctness of the sparse_mla_decode sparse MLA library against the native path.
-
-Every case runs the native ``mla_rope_generation`` once and then the same
-``TrtllmAttention.forward`` twice: with the native FMHA library and with sparse_mla_decode
-first in the layer's library list. Both outputs are compared with an FP64
-oracle over the same e4m3 query and pool rows that applies the native masking
-rules (index-list positions past a row's causal length are ignored, negative
-entries skipped, repeats counted).
-
-Error budget, fixed before any tuning:
-  * nrms(sparse_mla_decode - oracle) <= 1.25 * nrms(native - oracle) + 2e-4
-  * max|sparse_mla_decode - oracle| <= 2 * max|native - oracle| + 1e-3
-  * the existing FP8-KV sparse MLA tolerance: allclose(sparse_mla_decode, oracle, atol=0.12,
-    rtol=0.01)
-  * rows without a usable entry and non-finite values must match the native output exactly.
-"""
-
-import pytest
-import torch
-from sparse_mla_harness import (
-    DsaGenHarness,
-    compare,
-    error_stats,
-    oracle_attention,
-    record,
-    run_both,
-)
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode.sparse_mla import (  # noqa: E501
-    SparseMlaDecodeFmha,
-)
-from tensorrt_llm._torch.attention.backends.interface import (
-    AttentionForwardArgs,
-    AttentionInputType,
-)
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available(),
-    reason="needs an SM10x GPU and the sparse_mla_decode kernel",
-)
-
-
-@pytest.mark.parametrize(
-    "ctx_lens,seq_len_q,q_std",
-    [
-        ([32768], 6, 1.0),  # target verify, MTP K=5
-        ([32768], 1, 1.0),  # draft steps 1-4
-        ([32768], 6, 4.0),  # peaky softmax
-        ([32768], 6, 0.1),  # nearly flat softmax
-        ([100], 6, 1.0),  # fewer visible tokens than top-k
-        ([2042], 6, 1.0),  # last row sees exactly 2048 tokens
-        ([2043], 6, 1.0),  # rows straddle 2048 visible tokens
-        ([5000, 300], 6, 1.0),  # two sequences, one short
-        ([20000, 20000, 20000, 20000], 2, 1.0),  # four sequences, MTP K=1
-        ([131072], 6, 1.0),
-        ([391000], 6, 1.0),  # served-length context
-        ([491061], 6, 1.0),  # controlled-decode context
-    ],
-)
-def test_sparse_mla_matches_native(ctx_lens, seq_len_q, q_std):
-    harness = DsaGenHarness(ctx_lens, seq_len_q=seq_len_q, seed=21)
-    try:
-        topk = harness.random_topk(seed=22, overlap=0.7, recent=32)
-        step = harness.make_step(0, seed=23, q_std=q_std)
-        compare(
-            harness,
-            step,
-            topk,
-            "matches_native",
-            ctx_lens=ctx_lens,
-            seq_len_q=seq_len_q,
-            q_std=q_std,
-        )
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_empty_rows_tails_and_duplicates():
-    """Index lists as the native indexer shapes them, plus repeated and empty rows."""
-    harness = DsaGenHarness([8192, 40], seq_len_q=6, seed=31)
-    try:
-        topk = harness.random_topk(seed=32)
-        topk[1, :] = -1  # no usable entry
-        topk[3, :1500] = 17  # one row repeated 1500 times ...
-        topk[3, 1500:] = 4242  # ... and another one 548 times
-        for row in range(6, 12):  # short sequence: valid ids past the causal length are ignored
-            for pos in (300, 1500, 2047):
-                topk[row, pos] = 3
-        step = harness.make_step(0, seed=33)
-        compare(harness, step, topk, "empty_rows_tails_and_duplicates")
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_causal_clamp_is_per_row():
-    """Valid ids at list positions between a row's causal length and the sequence length."""
-    harness = DsaGenHarness([100], seq_len_q=6, seed=41)
-    try:
-        topk = harness.random_topk(seed=42)
-        for row in range(harness.num_rows):
-            for pos in range(harness.causal_len(row), 106):
-                topk[row, pos] = 5
-        step = harness.make_step(0, seed=43)
-        result = compare(harness, step, topk, "causal_clamp_per_row")
-        global_idx = harness.global_indices(0, topk)
-        per_sequence = oracle_attention(
-            step.quant_q,
-            harness.pool_rows(),
-            global_idx,
-            float(step.bmm1_scale[0]),
-            float(step.bmm2_scale[0]),
-            list_limit=[106] * harness.num_rows,
-        )
-        native_rows = result["native"].double().reshape(result["oracle"].shape)
-        # Rows 0-4 differ between the two rules; the native kernel must follow the per-row one.
-        per_row_err = error_stats(native_rows[:5], result["oracle"][:5])
-        per_seq_err = error_stats(native_rows[:5], per_sequence[:5])
-        record(
-            "sparse_mla_correctness.jsonl",
-            "causal_clamp_native_rule",
-            {"per_row": per_row_err, "per_sequence": per_seq_err},
-        )
-        assert per_row_err["nrms"] < per_seq_err["nrms"], (per_row_err, per_seq_err)
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_skips_negative_entries_inside_range():
-    """Negative ids before the causal length are skipped (the native kernel does not check them)."""
-    harness = DsaGenHarness([8192], seq_len_q=6, seed=51)
-    try:
-        topk = harness.random_topk(seed=52)
-        topk[2, ::3] = -1
-        topk[4, 1000:] = -1
-        topk[5, :1024] = -1
-        step = harness.make_step(0, seed=53)
-        native, decode_out, global_idx, oracle = run_both(harness, step, topk)
-        per_row = {
-            r: {
-                "sparse_mla_decode": error_stats(decode_out[r].double(), oracle[r]),
-                "native": error_stats(native[r].double(), oracle[r]),
-            }
-            for r in range(harness.num_rows)
-        }
-        record("sparse_mla_correctness.jsonl", "negative_entries_inside_range", {"rows": per_row})
-        clean = max(per_row[r]["sparse_mla_decode"]["nrms"] for r in (0, 1, 3))
-        for r in (2, 4, 5):
-            assert per_row[r]["sparse_mla_decode"]["nrms"] <= 2.0 * clean, (r, per_row[r], clean)
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_cuda_graph_replay():
-    """Capture both libraries once; replay with new indices, queries and rewritten pool rows."""
-    harness = DsaGenHarness([65536], seq_len_q=6, seed=61)
-    try:
-        step = harness.make_step(0, seed=62)
-        topk = harness.random_topk(seed=63, overlap=0.7, recent=32)
-        graphs = {}
-        outputs = {}
-        for use_sparse_mla in (False, True):
-            harness.select_library(0, use_sparse_mla=use_sparse_mla)
-            harness.forward(0, step, topk)
-            torch.cuda.synchronize()
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph):
-                outputs[use_sparse_mla] = harness.forward(0, step, topk)
-            graphs[use_sparse_mla] = graph
-        pool_rows = harness.pool_rows()
-        for it in range(6):
-            topk.copy_(harness.random_topk(seed=100 + it, overlap=0.5, recent=16))
-            if it % 2 == 1:
-                # Page reuse: overwrite the rows the next replay reads.
-                global_idx = harness.global_indices(0, topk)
-                rows = global_idx[global_idx >= 0].unique().long()
-                fresh = (
-                    (torch.randn(rows.numel(), pool_rows.shape[1], device="cuda") * (1 + it))
-                    .clamp(-448, 448)
-                    .to(torch.float8_e4m3fn)
-                )
-                pool_rows.view(torch.uint8)[rows] = fresh.view(torch.uint8)
-            fresh_step = harness.make_step(0, seed=200 + it)
-            for name in ("fused_q", "q_pe", "latent", "quant_q", "bmm1_scale", "bmm2_scale"):
-                getattr(step, name).copy_(getattr(fresh_step, name))
-            for use_sparse_mla in (False, True):
-                graphs[use_sparse_mla].replay()
-            torch.cuda.synchronize()
-            global_idx = harness.global_indices(0, topk)
-            oracle = oracle_attention(
-                step.quant_q,
-                pool_rows,
-                global_idx,
-                float(step.bmm1_scale[0]),
-                float(step.bmm2_scale[0]),
-                list_limit=harness.list_limits(),
-            )
-            native_err = error_stats(outputs[False].double(), oracle)
-            sparse_mla_err = error_stats(outputs[True].double(), oracle)
-            record(
-                "sparse_mla_correctness.jsonl",
-                "graph_replay",
-                {
-                    "iteration": it,
-                    "native_vs_oracle": native_err,
-                    "sparse_mla_vs_oracle": sparse_mla_err,
-                },
-            )
-            assert sparse_mla_err["nonfinite"] == 0
-            assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
-                it,
-                sparse_mla_err,
-                native_err,
-            )
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_falls_back_for_other_top_k():
-    """A request outside the domain (top-k 1024) must run the native library."""
-    harness = DsaGenHarness([16384], seq_len_q=6, index_topk=1024, seed=71)
-    try:
-        gen = torch.Generator(device="cpu").manual_seed(72)
-        topk = (
-            torch.stack(
-                [
-                    torch.randperm(harness.causal_len(r), generator=gen)[:1024]
-                    for r in range(harness.num_rows)
-                ]
-            )
-            .to(torch.int32)
-            .cuda()
-        )
-        step = harness.make_step(0, seed=73)
-        harness.select_library(0, use_sparse_mla=False)
-        native = harness.forward(0, step, topk).clone()
-        harness.select_library(0, use_sparse_mla=True)
-        fallback = harness.forward(0, step, topk).clone()
-        assert "SparseMlaDecodeFmha" not in harness.selected_library(0)
-        torch.testing.assert_close(fallback, native, rtol=0, atol=0)
-    finally:
-        harness.shutdown()
-
-
-def test_sparse_mla_static_eligibility():
-    fp8 = DsaGenHarness([256], seq_len_q=1, seed=81)
-    bf16 = DsaGenHarness([256], seq_len_q=1, kv_dtype=torch.bfloat16, seed=82)
-    try:
-        assert SparseMlaDecodeFmha.is_available(fp8.layers[0])
-        assert not SparseMlaDecodeFmha.is_available(bf16.layers[0])
-        assert sparse_mla.static_ineligibility(bf16.layers[0]) == "KV cache is not FP8"
-    finally:
-        fp8.shutdown()
-        bf16.shutdown()
-
-
-def test_sparse_mla_rejects_context_phase():
-    harness = DsaGenHarness([256], seq_len_q=1, seed=91)
-    try:
-        q = torch.empty(1, 8 * 576, dtype=torch.bfloat16, device="cuda")
-        args = AttentionForwardArgs(attention_input_type=AttentionInputType.context_only)
-        assert (
-            sparse_mla.request_ineligibility(q, None, None, harness.metadata, args)
-            == "not generation-only"
-        )
-    finally:
-        harness.shutdown()
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py
deleted file mode 100644
index cebbefb99b37..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/sparse_mla_decode/test_sparse_mla_stack_composition.py
+++ /dev/null
@@ -1,350 +0,0 @@
-# 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.
-"""sparse_mla_decode between engaged producers (absorbed_query, shared_indexer_absorbed_query,
-rope_kv_append) and its PDL consumer (the attention_output_tail).
-
-Each package set runs ``sparse_mla_stack_case.py`` in a fresh interpreter, without and
-with ``sparse_mla_decode``, on real GLM-5.2 layers at TP8 rank 0: layer 2 runs its own indexer
-(absorbed_query's layer), layer 3 reuses layer 2's top-k (shared_indexer_absorbed_query's layer).
-``MLA.forward`` runs the MTP verify step of one request (6 rows) as the model does. Asserted:
-
-* engagement, counted per call: absorbed_query returns its packet on the full layer and
-  shared_indexer_absorbed_query on the shared layer; rope_kv_append's fused RoPE and top-k bridge
-  serve every generation call; the attention_output_tail takes every call of both layers;
-  sparse_mla_decode runs on both layers;
-* the producers do not depend on the attention library: the FMHA inputs (FP8
-  query bytes, bmm scales, KV lengths, global indices and the pool rows they
-  address) are bit-identical without and with sparse_mla_decode;
-* sparse_mla_decode's latent output is within the frozen error budget against the FP64 oracle
-  on both layers, and the module outputs (after the attention_output_tail) agree with the
-  native-attention run;
-* sparse_mla_decode triggers its programmatic dependents before its gathers complete; the
-  attention_output_tail is that dependent. With sparse_mla_decode, the eager steps (no
-  recording kernels), the steps with a device synchronization right after each sparse_mla_decode
-  launch, and the replays of a captured step must give bit-identical latent and module outputs;
-* CUPTI: one ``sparseMlaDecodeKernel`` per DSA layer per step (eager and graph)
-  and no native sparse FMHA kernel with sparse_mla_decode; none without it; in graph replays
-  the kernel after sparse_mla_decode on its stream is attention_output_tail's W_UV kernel and it
-  becomes resident before sparse_mla_decode ends;
-* with rope_kv_append, every generation call of its bridge equals the native remap of the
-  same local top-k bit for bit, emits no ``-1`` below a row's causal length, and
-  is what the attention received.
-
-Neighbour sets: native, ``absorbed_query,shared_indexer_absorbed_query,attention_output_tail``
-(attention_output_tail default schedule) and the packages of the served selector that act at
-these call sites: absorbed_query, shared_indexer_absorbed_query, rope_kv_append,
-attention_output_tail, post_attention_hop, attention_output_tail_prefetch and
-attention_output_into_hop (with post_attention_hop and attention_output_into_hop the tail takes
-``INTO_HOP_SPARSE_DECODE_SCHEDULE`` behind sparse_mla_decode, which streams 32 ``w_o`` rows per
-CTA before its PDL wait). The last set also runs with the unmodified native indexer.
-
-Runs with ``GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION=1`` on an integration build that contains the
-packages, with the GLM-5.2 checkpoint.
-"""
-
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-from sparse_mla_harness import error_stats, oracle_attention, record
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.sparse_mla_decode import (  # noqa: E501
-    sparse_mla,
-)
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-import plan_testlib  # noqa: E402
-from utils.llm_data import llm_models_root
-
-_CKPT = os.environ.get(
-    "GLM52_NVFP4_CKPT",
-    os.path.join(str(llm_models_root()), "GLM-5.2-NVFP4"),
-)
-_CASE = Path(__file__).resolve().parent / "sparse_mla_stack_case.py"
-_ATTN_TAIL_HARNESS = (
-    Path(__file__).resolve().parent.parent / "attention_output_tail" / "real_mla_harness.py"
-)
-
-pytestmark = pytest.mark.skipif(
-    os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION") != "1"
-    or not torch.cuda.is_available()
-    or torch.cuda.get_device_capability()[0] != 10
-    or not sparse_mla.op_available()
-    or not os.path.isfile(os.path.join(_CKPT, "config.json"))
-    or not _ATTN_TAIL_HARNESS.is_file(),
-    reason="set GLM_MOE_DSA_SPARSE_MLA_DECODE_COMPOSITION=1 on an SM10x integration build with the checkpoint",
-)
-
-SERVED_SITE_SET = (
-    "absorbed_query,shared_indexer_absorbed_query,rope_kv_append,"
-    "attention_output_tail,post_attention_hop,attention_output_tail_prefetch,attention_output_into_hop"
-)
-STACK_SETS = [
-    ("", "fixed"),
-    ("absorbed_query,shared_indexer_absorbed_query,attention_output_tail", "fixed"),
-    (SERVED_SITE_SET, "fixed"),
-    (SERVED_SITE_SET, "indexer"),
-]
-ROWS = 6
-TOP_K = 2048
-FULL_LAYER, SHARED_LAYER = 2, 3
-TAIL_FEATURES = (
-    "attention_output_tail",
-    "attention_output_tail_prefetch",
-    "attention_output_into_hop",
-)
-
-
-def _run(features: str, topk_mode: str, out_dir: Path) -> dict:
-    out = out_dir / f"stack_{features.replace(',', '_') or 'native'}_{topk_mode}.pt"
-    env = dict(os.environ)
-    env[plan_testlib.FAST_PATHS_ENV] = features
-    env["GLM_MOE_DSA_SPARSE_MLA_DECODE_STACK_TOPK"] = topk_mode
-    env["GLM52_NVFP4_CKPT"] = _CKPT
-    env["TLLM_LOG_LEVEL"] = "INFO"
-    env.setdefault("TRTLLM_ENABLE_PDL", "1")
-    proc = subprocess.run(
-        [sys.executable, str(_CASE), str(out)], env=env, capture_output=True, text=True, check=False
-    )
-    assert proc.returncode == 0, proc.stdout[-8000:] + proc.stderr[-8000:]
-    result = torch.load(out, weights_only=False)
-    log = (proc.stdout + proc.stderr).splitlines()
-    result["capability_log"] = sorted(
-        {line[line.index("[glm_moe_dsa]") :] for line in log if "[glm_moe_dsa]" in line}
-    )
-    return result
-
-
-def _oracle(entry: dict) -> torch.Tensor:
-    rows = entry["rows"].cuda().reshape(ROWS * TOP_K, -1)
-    indices = entry["indices"].cuda()
-    local = torch.arange(ROWS * TOP_K, device="cuda", dtype=torch.int32).view(ROWS, TOP_K)
-    remapped = torch.where(indices >= 0, local, torch.full_like(local, -1))
-    kv_len = int(entry["kv_lens"][0])
-    limits = [kv_len - (ROWS - 1 - r) for r in range(ROWS)]
-    return oracle_attention(
-        entry["quant_q"].cuda(),
-        rows,
-        remapped,
-        float(entry["bmm1"][0]),
-        float(entry["bmm2"][0]),
-        list_limit=limits,
-    )
-
-
-def _same(a: torch.Tensor, b: torch.Tensor) -> bool:
-    return a.shape == b.shape and torch.equal(a.view(torch.uint8), b.view(torch.uint8))
-
-
-def _check_engagement(run: dict, selected: set, with_sparse_mla: bool) -> dict:
-    eng = run["engagement"]
-    counts = eng["counts"]
-    eager_steps = 2 + 2 + 2 + 1  # inputs, pdl, serialized, before capture (warm-up excluded)
-    checked = {}
-    if "absorbed_query" in selected:
-        checked["absorbed_query_full"] = counts.get(f"absorbed_query/layer{FULL_LAYER}", 0)
-        assert checked["absorbed_query_full"] >= eager_steps, counts
-    if "shared_indexer_absorbed_query" in selected:
-        checked["shared_query_shared"] = counts.get(
-            f"shared_indexer_absorbed_query/layer{SHARED_LAYER}", 0
-        )
-        assert checked["shared_query_shared"] >= eager_steps, counts
-        # The shared layer takes shared_indexer_absorbed_query's packet, not absorbed_query's.
-        assert counts.get(f"absorbed_query/layer{SHARED_LAYER}", 0) == 0, counts
-    if "rope_kv_append" in selected:
-        for layer in (FULL_LAYER, SHARED_LAYER):
-            assert counts.get(f"rope_append_rope/layer{layer}", 0) >= eager_steps, counts
-            assert counts.get(f"rope_append_bridge/layer{layer}", 0) >= eager_steps, counts
-    if selected & set(TAIL_FEATURES):
-        for layer in (FULL_LAYER, SHARED_LAYER):
-            tail = eng["attn_tail_tails"][layer]
-            assert tail is not None and tail["static_reason"] is None, tail
-            assert tail["engaged_calls"] >= eager_steps, tail
-            checked[f"attn_tail_layer{layer}"] = {
-                "engaged_calls": tail["engaged_calls"],
-                "early_rows": tail["config"]["early_rows"],
-                "oproj_grid": tail["config"]["oproj_grid"],
-            }
-    sparse_mla_calls = eng["sparse_mla_decode"]["engaged_calls"]
-    if with_sparse_mla:
-        for layer in (FULL_LAYER, SHARED_LAYER):
-            assert sparse_mla_calls.get(str(layer), 0) >= eager_steps, eng["sparse_mla_decode"]
-    else:
-        assert not sparse_mla_calls, eng["sparse_mla_decode"]
-    checked["sparse_mla_engaged_calls"] = sparse_mla_calls
-    return checked
-
-
-def _replay_check(run: dict, layers) -> dict:
-    """Compare every plain / serialized / graph step with the recorded eager step of its input."""
-    reference = {e["input"]: e["layers"] for e in run["inputs"]}
-    mismatches, index_changes, compared = [], [], 0
-    for kind in ("pdl", "serialized", "graph"):
-        for i, entry in enumerate(run[kind]):
-            ref = reference[entry["input"]]
-            for layer in layers:
-                got = entry["layers"][layer]
-                if not _same(got["indices"], ref[layer]["indices"]):
-                    index_changes.append((kind, i, layer))
-                for key in ("latent", "module_output"):
-                    compared += 1
-                    if not _same(got[key], ref[layer][key]):
-                        mismatches.append((kind, i, layer, key))
-    return {"compared": compared, "mismatches": mismatches, "index_changes": index_changes}
-
-
-@pytest.mark.parametrize("neighbours,topk_mode", STACK_SETS)
-def test_sparse_mla_stack_with_engaged_neighbours(neighbours, topk_mode, tmp_path):
-    out_dir = Path(os.environ.get("GLM_MOE_DSA_SPARSE_MLA_DECODE_RESULTS", tmp_path))
-    selected = {f for f in neighbours.split(",") if f}
-    with_sparse_mla = ",".join(sorted(selected) + ["sparse_mla_decode"])
-    base = _run(neighbours, topk_mode, out_dir)
-    cand = _run(with_sparse_mla, topk_mode, out_dir)
-    layers = cand["layers"]
-    assert layers == [FULL_LAYER, SHARED_LAYER] and cand["shared_layers"] == [SHARED_LAYER]
-    assert base["libraries"] == ["FallbackFmha"], base["libraries"]
-    assert cand["libraries"] == ["SparseMlaDecodeFmha"], cand["libraries"]
-
-    summary = {
-        "neighbours": neighbours or "native",
-        "topk_mode": topk_mode,
-        "ctx": cand["ctx"],
-        "trtllm_enable_pdl": os.environ.get("TRTLLM_ENABLE_PDL", "1"),
-        "engagement_without_sparse_mla": _check_engagement(base, selected, False),
-        "engagement_with_sparse_mla": _check_engagement(cand, selected, True),
-        "capability_log_with_sparse_mla": cand["capability_log"],
-        "steps": [],
-    }
-    for b, c in zip(base["inputs"], cand["inputs"]):
-        for layer in layers:
-            bl, cl = b["layers"][layer], c["layers"][layer]
-            # The producers do not depend on the attention library.
-            for key in ("quant_q", "bmm1", "bmm2", "kv_lens"):
-                assert _same(bl[key], cl[key]), (neighbours, layer, key)
-            same_order = _same(bl["indices"], cl["indices"])
-            same_sets = torch.equal(
-                bl["indices"].sort(dim=-1).values, cl["indices"].sort(dim=-1).values
-            )
-            if topk_mode == "fixed":
-                assert same_order and _same(bl["rows"], cl["rows"]), (neighbours, layer)
-            assert not bool((cl["indices"] < 0).any()), "a -1 index below the causal length"
-            native_err = error_stats(bl["latent"].cuda().double(), _oracle(bl))
-            sparse_mla_err = error_stats(cl["latent"].cuda().double(), _oracle(cl))
-            module_diff = error_stats(cl["module_output"].double(), bl["module_output"].double())
-            summary["steps"].append(
-                {
-                    "input": c["input"],
-                    "layer": layer,
-                    "indices_same_order": same_order,
-                    "indices_same_sets": same_sets,
-                    "native_latent_vs_oracle": native_err,
-                    "sparse_mla_latent_vs_oracle": sparse_mla_err,
-                    "module_output_sparse_mla_vs_native": module_diff,
-                }
-            )
-            assert sparse_mla_err["nonfinite"] == native_err["nonfinite"] == 0
-            assert sparse_mla_err["nrms"] <= 1.25 * native_err["nrms"] + 2e-4, (
-                layer,
-                sparse_mla_err,
-                native_err,
-            )
-            assert sparse_mla_err["max_abs"] <= 2.0 * native_err["max_abs"] + 1e-3, (
-                layer,
-                sparse_mla_err,
-            )
-            if same_sets:
-                assert module_diff["nonfinite"] == 0 and module_diff["nrms"] < 0.05, module_diff
-
-    summary["replays_without_sparse_mla"] = _replay_check(base, layers)
-    summary["replays_with_sparse_mla"] = _replay_check(cand, layers)
-    summary["kernels_without_sparse_mla"] = {
-        k: {f: v for f, v in s.items() if f != "attention"} for k, s in base["kernels"].items()
-    }
-    summary["kernels_with_sparse_mla"] = {
-        k: {f: v for f, v in s.items() if f != "attention"} for k, s in cand["kernels"].items()
-    }
-    summary["rope_append_bridge_without_sparse_mla"] = base["rope_append_bridge"]
-    summary["rope_append_bridge_with_sparse_mla"] = cand["rope_append_bridge"]
-    record(
-        "sparse_mla_stack_composition.jsonl", "stack_composition", json.loads(json.dumps(summary))
-    )
-
-    # sparse_mla_decode is deterministic: with it, every plain, serialized and replayed step
-    # reproduces the recorded eager step bit for bit when it attended the same indices (always with
-    # a fixed top-k; the native indexer may reorder or change its selection from step to step).
-    replays = summary["replays_with_sparse_mla"]
-    if topk_mode == "fixed":
-        assert not replays["mismatches"], replays
-    else:
-        changed = {(k, i, layer) for k, i, layer in replays["index_changes"]}
-        stable = [m for m in replays["mismatches"] if tuple(m[:3]) not in changed]
-        assert not stable, replays
-
-    # One attention kernel per DSA layer per step; the right library in each run.
-    for run, sparse_mla_expected in ((base, False), (cand, True)):
-        for kind, stats in run["kernels"].items():
-            calls = stats["steps"] * stats["layers"]
-            assert stats["sparse_mla_decode"] == (calls if sparse_mla_expected else 0), (
-                kind,
-                stats["sparse_mla_decode"],
-                calls,
-            )
-            assert stats["native_sparse_fmha"] == (0 if sparse_mla_expected else calls), (
-                kind,
-                stats["native_sparse_fmha"],
-                calls,
-            )
-            if "rope_kv_append" in selected:
-                assert stats["rope_append_bridge"] >= stats["steps"], (
-                    kind,
-                    stats["rope_append_bridge"],
-                )
-                if topk_mode == "fixed":
-                    assert stats["native_remap"] == 0, (kind, stats["native_remap"])
-    # The attention_output_tail is sparse_mla_decode's programmatic dependent in the captured graph
-    # and becomes resident well before sparse_mla_decode ends (sparse_mla_decode triggers before its
-    # gathers complete). Without PDL the tail starts only after sparse_mla_decode completes (CUPTI
-    # may still show a sub-microsecond overlap of the end and start stamps of back-to-back kernels,
-    # so the lead is compared, not its sign).
-    if selected & set(TAIL_FEATURES):
-        graph = cand["kernels"]["graph"]
-        sparse_mla_records = [a for a in graph["attention"] if a["kernel"] == "sparse_mla_decode"]
-        assert all("_wuv_kernel" in (a["successor"] or "") for a in sparse_mla_records), graph[
-            "successors"
-        ]
-        leads = graph["sparse_mla_successor_lead_us"]
-        median_lead = leads[len(leads) // 2]
-        if os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1":
-            assert graph["sparse_mla_successor_early"] == len(sparse_mla_records), graph[
-                "sparse_mla_successor_early"
-            ]
-            assert median_lead >= 1.0, leads
-        else:
-            assert median_lead < 0.5, leads
-
-    # rope_kv_append's bridge as the producer of the indices: equal to the native remap, no -1 below
-    # a row's causal length, and what the attention received.
-    if "rope_kv_append" in selected:
-        for run in (base, cand):
-            reports = run["rope_append_bridge"]
-            assert len(reports) == 2 * len(layers), reports
-            for report in reports:
-                assert report["equals_native_remap"], report
-                assert sum(report["negative_below_causal"]) == 0, report
-                assert report["attention_received_bridge_output"], report
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py
deleted file mode 100644
index 5bc49b7e05b6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_contracts.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# 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.
-"""Contract table of the target's kernels (``kernels/contracts.py``).
-
-The module is loaded by file path so the test needs neither torch nor the
-compiled extensions; kernel sources are read as text.
-"""
-
-import importlib.util
-import re
-from pathlib import Path
-
-import pytest
-
-_ROOT = Path(__file__).resolve().parents[8]
-_TARGET = ("tensorrt_llm", "_torch", "_experimental", "modeling_v2", "models", "glm_moe_dsa")
-_KERNELS = (*_TARGET, "shared", "kernels")
-_TESTS = ("tests", "unittest", "_torch", "modeling_v2", "models", "glm_moe_dsa")
-_KERNEL_TESTS = (*_TESTS, "shared", "kernels")
-
-
-def _load(name):
-    path = _ROOT.joinpath(*_KERNELS, f"{name}.py")
-    spec = importlib.util.spec_from_file_location(f"glm_moe_dsa_kernels_{name}", path)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-def test_contract_table():
-    contracts = _load("contracts")
-    owners = {k: c.owner for k, c in contracts.CONTRACTS.items()}
-    assert owners == {
-        "query_front": "absorbed_query",
-        "index_selection": "indexer_topk",
-        "sparse_attention": "sparse_mla_decode",
-        "attention_output_tail": "attention_output_tail",
-        "router": "router_gemm",
-        "routed_moe": "routed_moe",
-        "post_moe_hop": "post_moe_hop",
-        "draft_front": "draft_front",
-    }
-    # Rules are optional, so a four-field contract is still valid.
-    plain = contracts.Contract("x", "absorbed_query", ("a",), ("b",))
-    assert plain.rules == ()
-
-
-def test_attention_tail_rules():
-    contracts = _load("contracts")
-    tail = contracts.CONTRACTS["attention_output_tail"]
-    assert tail.outputs == ("partial[R,6144]",)
-    assert tail.rules == (
-        contracts.ATTENTION_TAIL_COMMUNICATION_RULE,
-        contracts.ATTENTION_TAIL_PDL_RULE,
-    )
-    comm = contracts.ATTENTION_TAIL_COMMUNICATION_RULE
-    assert "No communication is hidden inside attention_output_tail." in comm
-    assert "router_gemm/native all-reduce consumes the identical partial." in comm
-    pdl = contracts.ATTENTION_TAIL_PDL_RULE
-    for phrase in (
-        "attention_output_tail may be a PDL dependent of the latent's producer.",
-        "Before its griddepcontrol.wait a attention_output_tail kernel reads only immutable weights",
-        "writes only its own shared memory",
-        "no async-proxy global writes",
-        "Every latent read and every global write (v scratch, partial) follows the wait.",
-        "executes griddepcontrol.launch_dependents before it writes the partial",
-        "must execute griddepcontrol.wait before reading the partial",
-        "may be reused after attention_output_tail returns",
-    ):
-        assert phrase in pdl, phrase
-    switches = re.search(r"Kill switches: (.*)\.$", pdl).group(1)
-    assert switches == (
-        "TRTLLM_ENABLE_PDL=0 (all, entry included), TailConfig(trigger_dependents=False) (exit)"
-    )
-
-
-def _source(*parts):
-    path = _ROOT.joinpath(*parts)
-    if not path.is_file():
-        pytest.skip(f"{path} is not in this tree")
-    return path.read_text()
-
-
-def test_attention_tail_names_exist_in_the_tree():
-    """The kill switches and the into-hop protocol header the rules name exist."""
-    kernels = _source(*_KERNELS, "attention_output_tail", "kernels.py")
-    assert re.search(
-        r"def pdl_entry_enabled\(\) -> bool:\n(?:.*\n)*?    return pdl_enabled\(\)$", kernels, re.M
-    )
-    assert 'os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1"' in kernels
-    assert re.search(r"^\s+trigger_dependents: bool = True$", kernels, re.M)
-    rule = _load("contracts").ATTENTION_TAIL_COMMUNICATION_RULE
-    header = re.search(r"(cpp/\S+\.h)", rule).group(1)
-    assert "griddepcontrol.wait" in _source(*header.split("/"))
-
-
-def test_routed_moe_and_post_moe_hop_rules():
-    """The routed-input and deferred-norm seams and the token-major routed-MoE form."""
-    contracts = _load("contracts")
-    c = contracts.CONTRACTS
-    assert c["routed_moe"].rules == (
-        contracts.ROUTED_MOE_TOKEN_MAJOR_RULE,
-        contracts.POST_MOE_HOP_ROUTED_INPUT_RULE,
-    )
-    assert c["post_moe_hop"].rules == (
-        contracts.POST_MOE_HOP_ROUTED_INPUT_RULE,
-        contracts.POST_MOE_HOP_DEFERRED_NORM_RULE,
-        contracts.POST_MOE_HOP_SHARED_PROTOCOL_RULE,
-    )
-    assert c["query_front"].rules == (contracts.POST_MOE_HOP_DEFERRED_NORM_RULE,)
-    for phrase in ("identity t*8+k", "reads fc2_output only through the permutation"):
-        assert phrase in contracts.ROUTED_MOE_TOKEN_MAJOR_RULE, phrase
-    for phrase in (
-        "parsed only by dispatch.feature_variants",
-        "agreed by every rank",
-        "zero when routed_moe's kernel starts",
-        "the hop clears after reading them",
-        "routed_moe never reuses a buffer it wrote into",
-        "only after its own griddepcontrol.wait",
-    ):
-        assert phrase in contracts.POST_MOE_HOP_ROUTED_INPUT_RULE, phrase
-    for phrase in (
-        "The packet is not a tensor",
-        "materialize()",
-        "reads only immutable weights before its griddepcontrol.wait",
-    ):
-        assert phrase in contracts.POST_MOE_HOP_DEFERRED_NORM_RULE, phrase
-    assert "must change both in the same commit" in contracts.POST_MOE_HOP_SHARED_PROTOCOL_RULE
-
-
-def test_routed_moe_and_post_moe_hop_names_exist_in_the_tree():
-    """The names the routed-input and deferred-norm rules cite exist where the rules say."""
-    hop = _source(*_KERNELS, "post_moe_hop", "hop.py")
-    assert "def enable_routed_input(" in hop
-    assert 'ROUTED_MODES = ("row", "identity", "fp32")' in hop
-    assert 'dispatch.feature_variants(("routed_moe_into_hop",))["routed_moe_into_hop"]' in hop
-    packet = _source(*_KERNELS, "deferred_norm", "packet.py")
-    assert "class DeferredNormHidden:" in packet and "def materialize(self)" in packet
-    ptx = _source(*_KERNELS, "deferred_norm", "ptx.py")
-    assert "from ..post_moe_hop.kernels import (" in ptx
-    hop_kernel = _source(*_KERNELS, "deferred_norm", "hop_kernel.py")
-    assert "griddepcontrol_launch_dependents" in hop_kernel
-    assert "trigger_early" in hop_kernel
-    for name in ("test_deferred_norm_hop_mpi.py", "test_deferred_norm_routed_mpi.py"):
-        assert "mixed_with_post_moe_hop" in _source(*_KERNEL_TESTS, "deferred_norm", name)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py
deleted file mode 100644
index d60ed606d9c5..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/kernels/test_kernels_op_names.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# 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.
-"""Names of the GLM-MoE-DSA torch ops and the layout of their C++ sources.
-
-The model-specific kernels live in ``cpp/tensorrt_llm/kernels/glmMoeDsa/`` and
-their torch ops in ``cpp/tensorrt_llm/thop/glmMoeDsa/``, one ``<stem>Op.cpp``
-per kernel ``<stem>.cu``; the kernels build as one OBJECT library for the
-SM100 family. Every op they register is
-``trtllm::glm_moe_dsa_<what>``, and every such name the Python sources use is
-registered there or by a ``torch.library.custom_op`` of the Python package
-under a name the C++ sources leave free. The sources are read as text, so
-the test needs neither torch nor the compiled extensions.
-"""
-
-import re
-from pathlib import Path
-
-_THIS = Path(__file__).resolve()
-_ROOT = _THIS.parents[8]
-_CPP = _ROOT / "cpp" / "tensorrt_llm"
-_KERNELS = _CPP / "kernels" / "glmMoeDsa"
-_THOP = _CPP / "thop" / "glmMoeDsa"
-_PYTHON = _ROOT / "tensorrt_llm"
-
-OP_PREFIX = "glm_moe_dsa_"
-
-_REGISTRATION = re.compile(r"\bm\.(def|impl)\(\s*(\"([A-Za-z0-9_.]+)|[A-Z][A-Z0-9_]*)")
-_DEFINE = re.compile(r"^\s*#\s*define\s+([A-Z][A-Z0-9_]*)\s+\"([A-Za-z0-9_.]+)\"", re.M)
-_PYTHON_OP = re.compile(
-    r"\btorch\.library\.custom_op\(\s*[\"']trtllm::(%s[a-z0-9_]+)[\"']" % OP_PREFIX
-)
-
-
-def _registered(path):
-    """(kind, op name) of every m.def / m.impl in ``path``.
-
-    A name given by a macro is resolved through the file's ``#define``s (every
-    definition counts); a macro without a string definition yields ``None``.
-    """
-    text = path.read_text(errors="replace")
-    macros = {}
-    for name, value in _DEFINE.findall(text):
-        macros.setdefault(name, []).append(value)
-    found = []
-    for match in _REGISTRATION.finditer(text):
-        kind = match.group(1)
-        if match.group(3) is not None:
-            found.append((kind, match.group(3)))
-        else:
-            found += [(kind, value) for value in macros.get(match.group(2), [None])]
-    return found
-
-
-def _thop_ops():
-    ops = {}
-    for path in sorted(_THOP.glob("*.cpp")):
-        for kind, name in _registered(path):
-            assert name is not None, f"{path.name}: m.{kind} names an op through an unknown macro"
-            ops.setdefault(name, set()).add(kind)
-    return ops
-
-
-def _python_ops():
-    """Every ``trtllm::glm_moe_dsa_*`` op a ``torch.library.custom_op`` of the package defines."""
-    ops = set()
-    for path in sorted(_PYTHON.rglob("*.py")):
-        ops.update(_PYTHON_OP.findall(path.read_text(errors="replace")))
-    return ops
-
-
-def test_sources_sit_in_the_glm_moe_dsa_folders():
-    kernels = sorted(p.stem for p in _KERNELS.glob("*.cu"))
-    headers = sorted(p.stem for p in _KERNELS.glob("*.h"))
-    ops = sorted(p.stem[: -len("Op")] for p in _THOP.glob("*Op.cpp"))
-    assert kernels, "no kernel sources under kernels/glmMoeDsa"
-    assert kernels == headers
-    assert kernels == ops
-    assert sorted(p.name for p in _THOP.glob("*.cpp")) == sorted(f"{s}Op.cpp" for s in ops)
-    cmake = (_CPP / "thop" / "CMakeLists.txt").read_text()
-    assert "add_th_common_family(th_common_glm_moe_dsa)" in cmake
-    assert "add_subdirectory(glmMoeDsa)" in cmake
-    assert "$<TARGET_OBJECTS:th_common_glm_moe_dsa>" in cmake
-    assert "th_common_glm_moe_dsa" in (_THOP / "CMakeLists.txt").read_text()
-
-
-def test_kernels_build_as_one_sm100_family_object_library():
-    cmake = (_KERNELS / "CMakeLists.txt").read_text()
-    listed = re.search(r"add_library\(\s*glm_moe_dsa_kernels_src\s+OBJECT\s+([^)]*)\)", cmake)
-    assert listed is not None
-    assert sorted(listed.group(1).split()) == sorted(p.name for p in _KERNELS.glob("*.cu"))
-    assert "set_cuda_architectures(glm_moe_dsa_kernels_src 100f)" in cmake
-    # The parent's globs leave the folder to its own target, which joins the
-    # shared library's link line.
-    parent = (_CPP / "kernels" / "CMakeLists.txt").read_text()
-    assert "add_subdirectory(glmMoeDsa)" in parent
-    assert 'list(FILTER SRC_CPP EXCLUDE REGEX "glmMoeDsa/.*")' in parent
-    assert 'list(FILTER SRC_CU EXCLUDE REGEX "glmMoeDsa/.*")' in parent
-    assert re.search(r"^\s*glm_moe_dsa_kernels_src$", (_CPP / "CMakeLists.txt").read_text(), re.M)
-
-
-def test_every_op_defined_in_thop_glm_moe_dsa_has_the_model_prefix():
-    ops = _thop_ops()
-    assert ops, "no op registrations under thop/glmMoeDsa"
-    assert sorted(n for n in ops if not n.startswith(OP_PREFIX)) == []
-    # Every implemented op is also defined (the schema lives next to the kernel).
-    assert sorted(n for n, kinds in ops.items() if "def" not in kinds) == []
-
-
-def test_python_sources_use_only_registered_glm_moe_dsa_ops():
-    used = re.compile(r"\b%s[a-z0-9_]+\b" % OP_PREFIX)
-    thop, python = set(_thop_ops()), _python_ops()
-    assert sorted(thop & python) == []
-    registered = thop | python
-    unknown = []
-    for path in sorted(_PYTHON.rglob("*.py")):
-        text = path.read_text(errors="replace")
-        rel = str(path.relative_to(_ROOT))
-        unknown += [(rel, m.group(0)) for m in used.finditer(text) if m.group(0) not in registered]
-    assert unknown == []
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py
deleted file mode 100644
index e67313ed9dcf..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_attention.py
+++ /dev/null
@@ -1,754 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The GLM-5.2 tp8 target's attention modules: their adoption, the FMHA libraries the attention
-backends declare, the phase gate of every override, the indexer's top-k region and bridge, and
-the target's DSA op 2.
-
-Host-only: the built-in modules are created without their constructors and carry only the
-attributes adoption and the overrides read. The kernel factories are stand-ins, and the built-in
-methods and the decode functions an override calls are recorders.
-"""
-
-from __future__ import annotations
-
-import ast
-import inspect
-import types
-from pathlib import Path
-
-import pytest
-import torch
-from torch import nn
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    attention,
-    attention_forward,
-    dispatch,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared.kernels.indexer_topk import (  # noqa: E501
-    region,
-)
-from tensorrt_llm._torch.attention.backends.fmha import manager as fmha_manager
-from tensorrt_llm._torch.attention.backends.interface import AttentionInputType
-from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention
-from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks
-from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention
-from tensorrt_llm._torch.attention.mla import MLA
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention
-
-pytestmark = pytest.mark.cpu_only
-
-PREFILL, DECODE = dispatch.Phase.PREFILL, dispatch.Phase.DECODE
-
-
-# --- built-in modules without their constructors ------------------------------------------
-
-
-def _new(cls, **attributes):
-    obj = cls.__new__(cls)
-    if isinstance(obj, nn.Module):
-        nn.Module.__init__(obj)
-    for name, value in attributes.items():
-        setattr(obj, name, value)
-    return obj
-
-
-def _attention(layer_idx=3, full_indexer=True, dense_mha=False):
-    backend = _new(DSATrtllmAttention, indexer=_new(Indexer) if full_indexer else None)
-    return _new(
-        DeepseekV32Attention,
-        layer_idx=layer_idx,
-        sparse_attn_hooks=_new(DSAHooks, short_seq_mha_threshold=0, need_dense_mha=dense_mha),
-        mqa=backend,
-        mha=_new(TrtllmAttention) if dense_mha else None,
-        indexer=backend.indexer,
-    )
-
-
-class _SparseMlaDecodeFmha:
-    """Stands in for the sparse-MLA decode FMHA library."""
-
-
-class _Tail:
-    def __init__(self, mla) -> None:
-        self.mla = mla
-
-
-class _FastPaths:
-    """``dispatch.is_enabled`` / ``get_impl`` for the adoption tests."""
-
-    def __init__(self, enabled):
-        self.enabled = frozenset(enabled)
-        self.events = []
-        events = self.events
-
-        def make_attn_tail(mla):
-            events.append(("make_attn_tail", mla.layer_idx))
-            return _Tail(mla)
-
-        def wrap_attn_tail(mla, tail):
-            events.append(("wrap_attn_tail", tail))
-            return ("staged", tail)
-
-        self.entries = {
-            ("attention_output_tail", "make_attn_tail"): make_attn_tail,
-            ("oproj_weight_staging", "wrap_attn_tail"): wrap_attn_tail,
-            ("sparse_mla_decode", "fmha_lib"): _SparseMlaDecodeFmha,
-        }
-
-    def is_enabled(self, path):
-        assert path in dispatch.FAST_PATH_NAMES, path
-        return path in self.enabled
-
-    def get_impl(self, path, name):
-        return self.entries[(path, name)]
-
-
-@pytest.fixture
-def fast_paths(monkeypatch):
-    def install(enabled=dispatch.FAST_PATH_NAMES):
-        paths = _FastPaths(enabled)
-        monkeypatch.setattr(dispatch, "is_enabled", paths.is_enabled)
-        monkeypatch.setattr(dispatch, "get_impl", paths.get_impl)
-        return paths
-
-    return install
-
-
-def _keys(module):
-    return (
-        [name for name, _ in module.named_parameters()],
-        [name for name, _ in module.named_buffers()],
-        [name for name, _ in module.named_modules()],
-    )
-
-
-# --- adoption -----------------------------------------------------------------------------
-
-
-def test_adoption_reclasses_the_attention_modules_and_builds_the_tail(fast_paths):
-    paths = fast_paths()
-    mla = _attention()
-    hooks, backend, indexer = mla.sparse_attn_hooks, mla.mqa, mla.mqa.indexer
-    before = _keys(mla)
-
-    assert attention.GlmMLA.adopt(mla) is mla
-
-    assert type(mla) is attention.GlmMLA
-    assert type(hooks) is attention.GlmDsaHooks and mla.sparse_attn_hooks is hooks
-    assert type(backend) is attention.GlmDsaBackend and mla.mqa is backend
-    assert type(indexer) is attention.GlmIndexer and mla.indexer is indexer
-    assert hooks.short_seq_mha_threshold == 0 and hooks.need_dense_mha is False
-    # The tail is made for the layer, then wrapped by the o_proj weight staging.
-    staged, tail = mla._attn_tail
-    assert staged == "staged" and isinstance(tail, _Tail) and tail.mla is mla
-    assert paths.events == [("make_attn_tail", 3), ("wrap_attn_tail", tail)]
-    assert mla._tail_prefetch is True
-    assert _keys(mla) == before
-
-
-@pytest.mark.parametrize(
-    "path", ["attention_output_tail", "attention_output_tail_prefetch", "attention_output_into_hop"]
-)
-def test_any_tail_path_builds_the_tail(fast_paths, path):
-    paths = fast_paths({path})
-    mla = attention.GlmMLA.adopt(_attention())
-    assert isinstance(mla._attn_tail, _Tail) and paths.events == [("make_attn_tail", 3)]
-    assert mla._tail_prefetch is (path == "attention_output_tail_prefetch")
-
-
-def test_without_the_attention_paths_adoption_attaches_nothing(fast_paths):
-    paths = fast_paths(enabled=())
-    mla = attention.GlmMLA.adopt(_attention(full_indexer=False))
-    assert paths.events == []
-    assert mla._attn_tail is None and mla._tail_prefetch is False
-    assert mla.mqa.prepended_fmha_lib_classes == ()
-    assert "prepended_fmha_lib_classes" not in vars(mla.mqa)
-    assert mla.mqa.indexer is None and type(mla.mqa) is attention.GlmDsaBackend
-
-
-@pytest.mark.parametrize("dense_mha", [False, True], ids=["mqa only", "mha and mqa"])
-def test_the_sparse_mla_decode_library_is_declared_on_each_backend(fast_paths, dense_mha):
-    fast_paths({"sparse_mla_decode"})
-    mla = attention.GlmMLA.adopt(_attention(dense_mha=dense_mha))
-    backends = [mla.mqa] + ([mla.mha] if dense_mha else [])
-    for backend in backends:
-        assert backend.prepended_fmha_lib_classes == (_SparseMlaDecodeFmha,)
-    assert TrtllmAttention.prepended_fmha_lib_classes == ()
-
-
-def _other(cls):
-    return type(f"_Other{cls.__name__}", (cls,), {})
-
-
-def _refusal(case):
-    mla = _attention()
-    if case == "an attention subclass":
-        mla.__class__ = _other(DeepseekV32Attention)
-    elif case == "the MLA base class":
-        mla.__class__ = MLA
-    elif case == "a DSA hooks subclass":
-        mla.sparse_attn_hooks.__class__ = _other(DSAHooks)
-    elif case == "a DSA backend subclass":
-        mla.mqa.__class__ = _other(DSATrtllmAttention)
-    elif case == "an indexer subclass":
-        mla.mqa.indexer.__class__ = _other(Indexer)
-    elif case == "an adopted attention":
-        mla = attention.GlmMLA.adopt(mla)
-    return mla
-
-
-@pytest.mark.parametrize(
-    "case, got",
-    [
-        ("an attention subclass", "_OtherDeepseekV32Attention"),
-        ("the MLA base class", "MLA"),
-        ("a DSA hooks subclass", "_OtherDSAHooks"),
-        ("a DSA backend subclass", "_OtherDSATrtllmAttention"),
-        ("an indexer subclass", "_OtherIndexer"),
-        ("an adopted attention", "GlmMLA"),
-    ],
-)
-def test_adoption_refuses_anything_but_the_exact_builtin_class(fast_paths, case, got):
-    fast_paths(enabled=())
-    mla = _refusal(case)
-    with pytest.raises(TypeError, match=f"adopts the built-in .*, got {got}$"):
-        attention.GlmMLA.adopt(mla)
-
-
-def test_the_classes_declare_no_state_of_their_own():
-    for cls in (
-        attention.GlmMLA,
-        attention.GlmDsaHooks,
-        attention.GlmDsaBackend,
-        attention.GlmIndexer,
-    ):
-        assert "__slots__" not in vars(cls) and "__init__" not in vars(cls), cls
-    assert attention.GlmIndexer._topk_bridge is None
-    assert attention.GlmMLA._attn_tail is None and attention.GlmMLA._tail_prefetch is False
-
-
-# --- the FMHA libraries -------------------------------------------------------------------
-
-
-def _library(name, calls):
-    def is_available(cls, attn):
-        calls.append((name, attn.name))
-        return True
-
-    return type(
-        name,
-        (),
-        {"is_available": classmethod(is_available), "__init__": lambda self, attn: None},
-    )
-
-
-def _managers(monkeypatch, registry, declared):
-    """The libraries an mha and an mqa backend build, in the order ``create_weights`` builds
-    their managers, and the order their availability is asked in."""
-    calls = []
-    libraries = {name: _library(name, calls) for name in ("sparse", "a", "b", "fallback")}
-    monkeypatch.setattr(
-        fmha_manager, "get_enabled_fmha_lib_classes", lambda: [libraries[n] for n in registry]
-    )
-    built = []
-    for name in ("mha", "mqa"):
-        attn = types.SimpleNamespace(name=name)
-        if declared:
-            attn.prepended_fmha_lib_classes = tuple(libraries[n] for n in declared)
-        manager = fmha_manager.FmhaManager(attn)
-        built.append([type(lib).__name__ for lib in manager.fmha_libs])
-    return built, calls
-
-
-def test_the_declared_library_reproduces_a_registry_with_it_first(monkeypatch):
-    """A backend that declares the library builds the list and asks the libraries in the order a
-    registry with the library first would; one that declares nothing builds the registry's."""
-    first = _managers(monkeypatch, ["sparse", "a", "b", "fallback"], declared=())
-    declared = _managers(monkeypatch, ["a", "b", "fallback"], declared=("sparse",))
-    assert declared == first
-    assert declared[0] == [["sparse", "a", "b", "fallback"]] * 2
-    plain = _managers(monkeypatch, ["a", "b", "fallback"], declared=())
-    assert plain[0] == [["a", "b", "fallback"]] * 2
-
-
-# --- the phase gate -----------------------------------------------------------------------
-
-# Each override and the built-in method it falls back to.
-OVERRIDES = [
-    (attention.GlmMLA, DeepseekV32Attention, "forward"),
-    (attention.GlmMLA, DeepseekV32Attention, "forward_absorption_generation"),
-    (attention.GlmDsaHooks, DSAHooks, "forward"),
-    (attention.GlmDsaHooks, DSAHooks, "forward_custom_op"),
-    (attention.GlmDsaBackend, DSATrtllmAttention, "sparse_attn_predict"),
-    (attention.GlmIndexer, Indexer, "_fill_decode_topk"),
-]
-
-
-def _method(cls, name):
-    tree = ast.parse(Path(attention.__file__).read_text())
-    return next(
-        n
-        for c in tree.body
-        if isinstance(c, ast.ClassDef) and c.name == cls.__name__
-        for n in c.body
-        if isinstance(n, ast.FunctionDef) and n.name == name
-    )
-
-
-@pytest.mark.parametrize("cls, builtin, name", OVERRIDES, ids=lambda x: getattr(x, "__name__", x))
-def test_each_override_starts_with_the_phase_gate(cls, builtin, name):
-    """``if dispatch.phase() is not DECODE: return super().<name>(<every built-in parameter by
-    name>)``; the parameters only the decode composition takes are dropped."""
-    node = _method(cls, name)
-    docstring = isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant)
-    gate = node.body[1] if docstring else node.body[0]
-    assert ast.unparse(gate.test) == "dispatch.phase() is not dispatch.Phase.DECODE"
-    assert len(gate.body) == 1 and not gate.orelse
-    call = gate.body[0].value
-    assert ast.unparse(call.func) == f"super().{name}" and not call.args
-    parameters = inspect.signature(getattr(builtin, name)).parameters
-    expected = [
-        p.name if p.kind is not p.VAR_KEYWORD else None
-        for p in parameters.values()
-        if p.name != "self"
-    ]
-    assert [k.arg for k in call.keywords] == expected
-    # The override keeps the built-in parameters, in order, and may add keywords after them.
-    named = [p for p in parameters.values() if p.kind is not p.VAR_KEYWORD]
-    own = inspect.signature(getattr(cls, name)).parameters.values()
-    assert [p.name for p in own][: len(named)] == [p.name for p in named]
-    assert any(p.kind is p.VAR_KEYWORD for p in own) == (len(named) < len(parameters))
-
-
-def _record(monkeypatch, owner, name, calls, result=None):
-    def recorder(*args, **kwargs):
-        calls.append((name, args, kwargs))
-        return f"{getattr(owner, '__name__', owner)}.{name}" if result is None else result
-
-    monkeypatch.setattr(owner, name, recorder)
-
-
-def _adopted(monkeypatch):
-    monkeypatch.setattr(dispatch, "is_enabled", lambda path: False)
-    return attention.GlmMLA.adopt(_attention())
-
-
-@pytest.mark.parametrize("phase", [PREFILL, None], ids=["prefill", "outside a step"])
-def test_off_decode_every_override_runs_the_builtin_method(monkeypatch, phase):
-    mla = _adopted(monkeypatch)
-    calls = []
-    for owner, name in [
-        (MLA, "forward"),
-        (MLA, "forward_absorption_generation"),
-        (DSAHooks, "forward"),
-        (DSAHooks, "forward_custom_op"),
-        (DSATrtllmAttention, "sparse_attn_predict"),
-    ]:
-        _record(monkeypatch, owner, name, calls)
-    for name in (
-        "forward_absorption_generation",
-        "forward_sparse_attn",
-        "forward_sparse_attn_custom_op",
-    ):
-        monkeypatch.setattr(attention_forward, name, lambda *a, **k: pytest.fail())
-    monkeypatch.setattr(region, "cross_rank_topk_decode", lambda *a, **k: pytest.fail())
-
-    def run():
-        assert mla.forward("pos", "h", "md", "params", "latent") == "MLA.forward"
-        assert (
-            mla.forward_absorption_generation(
-                "q", "kv", "pe", "md", "out", "pos", "lc", "absorbed", "deferred", extra=1
-            )
-            == "MLA.forward_absorption_generation"
-        )
-        assert mla.sparse_attn_hooks.forward(mla, "pos", "h", "md", ["o"]) == "DSAHooks.forward"
-        assert (
-            mla.sparse_attn_hooks.forward_custom_op(mla, "h", "pos", ["o"], None)
-            == "DSAHooks.forward_custom_op"
-        )
-        assert (
-            mla.mqa.sparse_attn_predict("q", None, "md", "args")
-            == "DSATrtllmAttention.sparse_attn_predict"
-        )
-        assert mla.mqa.indexer._fill_decode_topk("md", "q", "w", "buf", 0, 6) is False
-
-    if phase is None:
-        run()
-    else:
-        with dispatch.step(phase):
-            run()
-    assert [(name, kwargs) for name, _, kwargs in calls] == [
-        (
-            "forward",
-            dict(
-                position_ids="pos",
-                hidden_states="h",
-                attn_metadata="md",
-                all_reduce_params="params",
-                latent_cache_gen="latent",
-            ),
-        ),
-        (
-            "forward_absorption_generation",
-            dict(
-                q="q",
-                compressed_kv="kv",
-                k_pe="pe",
-                attn_metadata="md",
-                output="out",
-                position_ids="pos",
-                latent_cache="lc",
-                extra=1,
-            ),
-        ),
-        (
-            "forward",
-            dict(
-                mla=mla,
-                position_ids="pos",
-                hidden_states="h",
-                attn_metadata="md",
-                attn_output=["o"],
-            ),
-        ),
-        (
-            "forward_custom_op",
-            dict(
-                mla=mla,
-                hidden_states="h",
-                position_ids="pos",
-                attn_output=["o"],
-                latent_cache_gen=None,
-            ),
-        ),
-        ("sparse_attn_predict", dict(q="q", k=None, metadata="md", forward_args="args")),
-    ]
-
-
-def test_on_a_decode_step_each_override_runs_the_decode_function(monkeypatch):
-    mla = _adopted(monkeypatch)
-    calls = []
-    for name in (
-        "forward_absorption_generation",
-        "forward_sparse_attn",
-        "forward_sparse_attn_custom_op",
-    ):
-        _record(monkeypatch, attention_forward, name, calls)
-    for owner, name in [
-        (MLA, "forward_absorption_generation"),
-        (DSAHooks, "forward"),
-        (DSAHooks, "forward_custom_op"),
-    ]:
-        monkeypatch.setattr(owner, name, lambda *a, **k: pytest.fail())
-
-    with dispatch.step(DECODE):
-        mla.forward_absorption_generation(
-            "q", "kv", "pe", "md", "out", "pos", "lc", "absorbed", "deferred", extra=1
-        )
-        assert mla.sparse_attn_hooks.forward(mla, "pos", "h", "md", ["o"]) is None
-        assert mla.sparse_attn_hooks.forward_custom_op(mla, "h", "pos", ["o"], None) is True
-
-    assert calls == [
-        (
-            "forward_absorption_generation",
-            (mla, "q", "kv", "pe", "md", "out"),
-            dict(
-                position_ids="pos",
-                latent_cache="lc",
-                absorbed_fused_q="absorbed",
-                deferred_indexer_k="deferred",
-                extra=1,
-            ),
-        ),
-        ("forward_sparse_attn", (mla, "pos", "h", "md", ["o"]), {}),
-        ("forward_sparse_attn_custom_op", (mla, "h", "pos", ["o"], None), {}),
-    ]
-
-
-# --- the indexer's top-k region and bridge ------------------------------------------------
-
-
-def _indexer(gvr_deterministic=False):
-    indexer = attention.GlmIndexer.adopt(_new(Indexer))
-    indexer.top_k = types.SimpleNamespace(gvr_deterministic=gvr_deterministic)
-    return indexer
-
-
-@pytest.mark.parametrize("cross_rank", [True, False])
-@pytest.mark.parametrize("gvr_deterministic", [False, True], ids=["default", "deterministic"])
-def test_the_decode_topk_region_tries_the_cross_rank_topk_first(
-    monkeypatch, cross_rank, gvr_deterministic
-):
-    calls = []
-    buffer = torch.zeros(8, 4, dtype=torch.int32)
-
-    def cross_rank_topk_decode(indexer, md, q, w, offset, rows, out, bridge):
-        calls.append(("cross_rank_topk", offset, rows, out.data_ptr(), bridge))
-        return cross_rank
-
-    def indexer_topk_decode(indexer, md, q, w, offset, rows, out):
-        calls.append(("indexer_topk", offset, rows, out.data_ptr()))
-        return True
-
-    monkeypatch.setattr(region, "cross_rank_topk_decode", cross_rank_topk_decode)
-    monkeypatch.setattr(region, "indexer_topk_decode", indexer_topk_decode)
-    indexer = _indexer(gvr_deterministic)
-    indexer._topk_bridge = "bridge"
-
-    with dispatch.step(DECODE):
-        assert indexer._fill_decode_topk("md", "q", "w", buffer, 2, 6) is True
-
-    rows = buffer[2:8].data_ptr()
-    expected = [] if gvr_deterministic else [("cross_rank_topk", 2, 6, rows, "bridge")]
-    if gvr_deterministic or not cross_rank:
-        expected.append(("indexer_topk", 2, 6, rows))
-    assert calls == expected
-
-
-def _rope_kv_append(monkeypatch, served, calls):
-    def entry(name):
-        def run(*args):
-            calls.append((name,) + args)
-            return served
-
-        return run
-
-    monkeypatch.setattr(dispatch, "get_impl", lambda path, name: entry((path, name)))
-
-
-@pytest.mark.parametrize(
-    "bridge, has_prefill, generations, served",
-    [
-        ("bridge", False, 3, True),
-        ("bridge", False, 3, False),
-        ("bridge", True, 3, True),
-        ("bridge", False, 1, True),
-        (None, False, 3, True),
-    ],
-)
-def test_the_reused_topk_copy_defers_to_the_bridge(
-    monkeypatch, bridge, has_prefill, generations, served
-):
-    """Three generation rows; the bridge is asked only for a pure decode call with one row per
-    request, and the built-in copy runs unless it takes the copy."""
-    calls = []
-    _rope_kv_append(monkeypatch, served, calls)
-    indexer = _indexer()
-    indexer._topk_bridge = bridge
-    shared = torch.arange(12, dtype=torch.int32).view(3, 4)
-    metadata = types.SimpleNamespace(shared_topk_indices=shared)
-    buffer = torch.full((4, 4), -1, dtype=torch.int32)
-
-    indexer._copy_reused_topk(metadata, buffer, 1, 3, generations, has_prefill, None)
-
-    asked = bridge is not None and not has_prefill and generations == 3
-    assert [c[0] for c in calls] == ([("rope_kv_append", "defer_reuse_copy")] if asked else [])
-    if asked:
-        _, bridge_arg, buffer_arg, offset, rows, has_prior = calls[0]
-        assert (bridge_arg, buffer_arg, offset, has_prior) == ("bridge", buffer, 1, False)
-        assert torch.equal(rows, shared)
-    if asked and served:
-        assert (buffer == -1).all()
-    else:
-        assert torch.equal(buffer[1:], shared[:generations].expand(3, 4))
-        assert (buffer[0] == -1).all()
-
-
-@pytest.mark.parametrize("served", [True, False])
-def test_the_mtp_row_stash_defers_to_the_bridge(monkeypatch, served):
-    calls, stock = [], []
-    _rope_kv_append(monkeypatch, served, calls)
-    monkeypatch.setattr(Indexer, "_store_mtp_shared_rows", lambda self, *a: stock.append(a))
-    indexer = _indexer()
-    indexer._topk_bridge = "bridge"
-    accepted = torch.tensor([9, 1, 2, 3])
-    metadata = types.SimpleNamespace(shared_topk_indices="shared", mtp_num_accepted=accepted)
-    args = ("buffer", metadata, 1, 3, 18, 0, False, True, True)
-
-    indexer._store_mtp_shared_rows(*args)
-
-    (call,) = calls
-    assert call[:6] == (("rope_kv_append", "defer_mtp_row_stash"), "bridge", "shared", 1, 3, 6)
-    assert torch.equal(call[6], accepted[1:4])
-    assert stock == ([] if served else [args])
-
-
-@pytest.mark.parametrize(
-    "bridge, has_decode, has_prefill, is_generation",
-    [
-        (None, True, False, True),
-        ("bridge", False, False, True),
-        ("bridge", True, True, True),
-        ("bridge", True, False, False),
-    ],
-)
-def test_the_mtp_row_stash_without_a_bridge_call_is_the_builtin_one(
-    monkeypatch, bridge, has_decode, has_prefill, is_generation
-):
-    calls, stock = [], []
-    _rope_kv_append(monkeypatch, True, calls)
-    monkeypatch.setattr(Indexer, "_store_mtp_shared_rows", lambda self, *a: stock.append(a))
-    indexer = _indexer()
-    indexer._topk_bridge = bridge
-    args = ("buffer", "md", 1, 3, 18, 0, has_prefill, has_decode, is_generation)
-    indexer._store_mtp_shared_rows(*args)
-    assert calls == [] and stock == [args]
-
-
-def test_the_builtin_seams_are_the_builtin_statements():
-    """The stock bodies of the three ``sparse_attn_indexer`` seams: the reuse copy, no decode
-    top-k of their own, and the MTP row stash."""
-    indexer = _new(Indexer)
-    shared = torch.arange(12, dtype=torch.int32).view(3, 4)
-    metadata = types.SimpleNamespace(shared_topk_indices=shared, mtp_num_accepted=None)
-    buffer = torch.full((4, 4), -1, dtype=torch.int32)
-    indexer._copy_reused_topk(metadata, buffer, 1, 3, 3, False, None)
-    assert torch.equal(buffer[1:], shared) and (buffer[0] == -1).all()
-    assert indexer._fill_decode_topk(metadata, None, None, buffer, 0, 4) is False
-
-    # Two generation requests at next_n = 2: the last row of each is stashed after row 1.
-    rows = torch.arange(16, dtype=torch.int32).view(4, 4)
-    metadata.shared_topk_indices = torch.zeros(3, 4, dtype=torch.int32)
-    indexer._store_mtp_shared_rows(rows, metadata, 1, 2, 4, 0, False, True, True)
-    assert torch.equal(metadata.shared_topk_indices[1:], rows[[1, 3]])
-    metadata.shared_topk_indices = None
-    indexer._store_mtp_shared_rows(rows, metadata, 0, 2, 4, 0, False, True, True)
-    assert torch.equal(metadata.shared_topk_indices, rows[[1, 3]])
-
-
-class _Bridged(Exception):
-    pass
-
-
-@pytest.mark.parametrize("fails", [False, True], ids=["returns", "raises"])
-def test_the_backend_hands_the_bridge_to_its_indexer_for_one_call(monkeypatch, fails):
-    seen = []
-    bridge_paths = {"rope_kv_append", "cross_rank_topk_global_emit"}
-    monkeypatch.setattr(dispatch, "is_enabled", lambda path: path in bridge_paths)
-
-    def get_impl(path, name):
-        assert (path, name) in {
-            ("cross_rank_topk_global_emit", "new_bridge_packet"),
-            ("cross_rank_topk_global_emit", "remap_topk_to_global"),
-        }, (path, name)
-        if name == "new_bridge_packet":
-            return lambda metadata, is_generation: ("bridge", is_generation)
-        return lambda backend, topk, metadata, bridge: ("global", topk, bridge)
-
-    monkeypatch.setattr(dispatch, "get_impl", get_impl)
-    backend = _attention().mqa
-    backend.__class__ = attention.GlmDsaBackend
-    indexer = backend.indexer
-    indexer.__class__ = attention.GlmIndexer
-    indexer.mtp_index_share = False
-    backend.get_local_layer_idx = lambda metadata: 0
-
-    def forward_from_projected(metadata, q, intermediates, is_generation):
-        seen.append((indexer._topk_bridge, is_generation))
-        if fails:
-            raise _Bridged
-        return "topk"
-
-    indexer.forward_from_projected = forward_from_projected
-    metadata = types.SimpleNamespace(
-        num_ctx_tokens=0,
-        num_tokens=6,
-        shared_topk_indices=None,
-        in_mtp_draft_loop=False,
-        kv_cache_manager=None,
-    )
-    args = types.SimpleNamespace(
-        attention_input_type=AttentionInputType.generation_only,
-        sparse_backend_args=types.SimpleNamespace(indexer_intermediates="intermediates"),
-    )
-
-    with dispatch.step(DECODE):
-        if fails:
-            with pytest.raises(_Bridged):
-                backend.sparse_attn_predict("q", None, metadata, args)
-        else:
-            out = backend.sparse_attn_predict("q", None, metadata, args)
-            assert out == (("global", "topk", ("bridge", True)), None)
-
-    assert seen == [(("bridge", True), True)]
-    assert indexer._topk_bridge is None
-
-
-# --- the attention-tail adoption and the absorbed-query module check ----------------------
-
-
-def test_the_absorbed_query_packet_takes_the_adopted_dsa_hooks(monkeypatch):
-    query_proj = pytest.importorskip(
-        f"{dispatch.__package__}.kernels.absorbed_query.query_proj",
-        reason="the absorbed_query kernels do not import here",
-    )
-    import tensorrt_llm._utils as utils
-
-    monkeypatch.setattr(utils, "get_sm_version", lambda: 103)
-    mla = types.SimpleNamespace(
-        sparse_attn_hooks=_new(attention.GlmDsaHooks), is_lite=True, fuse_qkv_a_proj=True
-    )
-    assert query_proj.module_eligibility(mla) == (False, "needs the fused [q_a | kv_a] projection")
-    mla.sparse_attn_hooks = _new(DSAHooks)
-    assert query_proj.module_eligibility(mla) == (False, "needs the fused [q_a | kv_a] projection")
-    mla.sparse_attn_hooks = _new(_other(DSAHooks))
-    assert query_proj.module_eligibility(mla) == (False, "not a DSA MLA module")
-
-
-# --- the target's DSA op 2 ----------------------------------------------------------------
-
-
-def _op_args():
-    t = torch.zeros(2, 4)
-    return (t, t, t, t, [t], None, "7", torch.zeros(2, 4))
-
-
-def test_the_op_runs_the_targets_op2_function_on_the_registered_layer(monkeypatch):
-    calls = []
-    monkeypatch.setattr(
-        attention_forward, "extract_extra_attrs", lambda layer, kind: ("md", f"mla{layer}{kind}")
-    )
-    monkeypatch.setattr(attention_forward, "_forward_dsa_attn", lambda *args: calls.append(args))
-    args = _op_args()
-
-    torch.ops.trtllm.glm_moe_dsa_mla_dsa_attn_inplace(*args)
-    attention_forward.maybe_bcg_glm_moe_dsa_mla_dsa_attn_inplace(*args)
-
-    assert len(calls) == 2
-    for call in calls:
-        assert call[0] == "mla7mla" and call[7] == "md" and call[6] is None
-        assert call[8] is args[-1]
-
-
-def test_the_op_mutates_only_its_output():
-    schema = torch.ops.trtllm.glm_moe_dsa_mla_dsa_attn_inplace.default._schema
-    assert [a.name for a in schema.arguments if a.alias_info and a.alias_info.is_write] == [
-        "output"
-    ]
-    assert len(schema.returns) == 0
-
-
-def test_the_ops_fake_runs_nothing(monkeypatch):
-    from torch._subclasses.fake_tensor import FakeTensorMode
-
-    monkeypatch.setattr(attention_forward, "_forward_dsa_attn", lambda *a: pytest.fail())
-    with FakeTensorMode():
-        args = _op_args()
-        torch.ops.trtllm.glm_moe_dsa_mla_dsa_attn_inplace(*args)
-
-
-def test_only_the_target_registers_the_op():
-    """The op is this target's own: no other module of the package names it, and the target's
-    op-2 module registers it once, so it exists only after the target is imported.
-    """
-    name = "glm_moe_dsa_mla_dsa_attn_inplace"
-    module = Path(attention_forward.__file__).resolve()
-    package = next(parent for parent in module.parents if parent.name == "tensorrt_llm")
-    naming = [path for path in package.rglob("*.py") if name in path.read_text(encoding="utf-8")]
-    assert naming == [module]
-    source = module.read_text(encoding="utf-8")
-    assert source.count(f'@torch.library.custom_op("trtllm::{name}"') == 1
-    assert source.count(f"@{name}.register_fake") == 1
-    assert hasattr(torch.ops.trtllm, name)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py
deleted file mode 100644
index bebdfaeae377..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_decline_kinds.py
+++ /dev/null
@@ -1,534 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""How the GLM-5.2 tp8 target classifies a fast path's decline: the sites, the layer cells, the
-rules, and what ``log_capability`` does with each class.
-
-Host-only: the declines are synthetic calls with the keys and reasons the kernels log, the log
-goes into lists, and the sources are read as text, so nothing here needs a GPU.
-"""
-
-from __future__ import annotations
-
-import ast
-from pathlib import Path
-
-import pytest
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-D = dispatch.Decline
-# The family: its shared code and every target of it.
-FAMILY = Path(dispatch.__file__).parents[1]
-SCOPE = "TRTLLM_STAIRCASE_GLM_LL_DEFERRED_NORM_SCOPE"
-LAYERS = "TRTLLM_STAIRCASE_GLM_LL_FUSED_KV_INDEXER_PROJECTION_LAYERS"
-
-
-def _layers(spec: str) -> frozenset:
-    """``"0-2,6,10"`` as a set of layer indices."""
-    out = set()
-    for part in spec.split(","):
-        lo, _, hi = part.partition("-")
-        out |= set(range(int(lo), int(hi or lo) + 1))
-    return frozenset(out)
-
-
-SHARED = (
-    "3-5,7-9,11-13,15-17,19-21,23-25,27-29,31-33,35-37,39-41,43-45,47-49,51-53,55-57,59-61,"
-    "63-65,67-69,71-73,75-77"
-)
-FULL = "0-2,6,10,14,18,22,26,30,34,38,42,46,50,54,58,62,66,70,74"
-
-# The layers each per-layer key engages on, on every rank, when this target serves its
-# checkpoint: (path, key family, layers).
-CENSUS = (
-    ("absorbed_query", "layer", "0-77"),
-    ("attention_output_tail", "layer", "0-77"),
-    ("shared_indexer_absorbed_query", "layer", SHARED),
-    ("deferred_norm", "scope[shared]_into_layer", SHARED.replace("3-5", "4-5")),
-    ("sparse_mla_decode", "call:layer", "0-78"),
-    ("sparse_mla_decode", "layer", "0-78"),
-    ("draft_absorbed_query", "layer", "78"),
-    ("deferred_norm_evict_hints", "projection_layer", SHARED.replace("3-5", "4-5")),
-    ("indexer_layer_front", "draft_layer", "78"),
-    ("indexer_layer_front", "target_layer", FULL),
-    ("oproj_weight_staging", "layer", "0-77"),
-    ("fused_kv_indexer_projection", "layer", "3"),
-)
-
-# Every decline the served composition logs, on every rank: (path, key, reason). None of them
-# is a miss.
-SERVED_DECLINES = (
-    ("absorbed_query", "layer78", "target layers only"),
-    ("attention_output_tail", "layer78", "draft or unindexed layer"),
-    ("router_gemm", "router:not a target layer", "not a target layer"),
-    *(
-        (
-            "shared_indexer_absorbed_query",
-            f"layer{i}",
-            "full indexer layer (shared_indexer_absorbed_query serves shared layers)",
-        )
-        for i in sorted(_layers(FULL))
-    ),
-    ("shared_indexer_absorbed_query", "layer78", "target layers only"),
-    *(
-        (
-            "deferred_norm",
-            f"scope[shared]_into_layer{i}",
-            f"{SCOPE}=shared: layer {i} is a full-indexer layer; post_moe_hop hop",
-        )
-        for i in sorted(_layers(FULL) - {0, 1, 2})
-    ),
-    ("dense_mlp", "dense_mlp:row count outside [1, 8]", "row count outside [1, 8]"),
-    ("cross_layer_expert_prefetch", "spec_experts_declined", "128 rows outside [1, 8]"),
-    ("routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]"),
-    (
-        "sparse_mla_decode",
-        "request:not generation-only",
-        "not generation-only; first seen at layer 3 with 1 rows",
-    ),
-    *(("draft_absorbed_query", f"layer{i}", "draft (MTP) layers only") for i in range(78)),
-    ("draft_loop", "draft_loop:context rows present", "context rows present"),
-    ("draft_loop", "draft_step:dense top-k fallback active", "dense top-k fallback active"),
-    (
-        "deferred_input_prep",
-        "prep[no previous speculative device tensors]",
-        "no previous speculative device tensors",
-    ),
-    (
-        "deferred_input_prep",
-        "prep[not a steady decode-graph replay]",
-        "not a steady decode-graph replay",
-    ),
-    (
-        "decode_prologue",
-        "target_prologue[not a generation-only step]",
-        "not a generation-only step",
-    ),
-    *(
-        (
-            "indexer_layer_front",
-            f"target_layer{i}",
-            "shared-indexer layer (no q_a RMSNorm launch to remove; "
-            "shared_indexer_absorbed_query / absorbed_query serve it)",
-        )
-        for i in sorted(_layers(SHARED))
-    ),
-    ("fused_kv_indexer_projection", "layer78", f"draft layers not selected ({LAYERS})"),
-    *(
-        ("fused_kv_indexer_projection", f"layer{i}", f"full layers not selected ({LAYERS})")
-        for i in sorted(_layers(FULL))
-    ),
-    *(
-        ("fused_kv_indexer_projection", f"layer{i}", f"shared layers not selected ({LAYERS})")
-        for i in sorted(_layers(SHARED) - {3})
-    ),
-    *(
-        (
-            "draft_routed_moe",
-            f"rows={rows} unfinalized output requested",
-            "unfinalized output requested",
-        )
-        for rows in (1, 2, 6)
-    ),
-    ("draft_routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]"),
-    # Decided per object behind the first object's line: the deferred-norm hops into the
-    # full-indexer layers, and a dense attention instance of a sparse-attention layer.
-    (
-        "deferred_norm",
-        "deferred_hop",
-        f"{SCOPE}=shared: layer 6 is a full-indexer layer; post_moe_hop hop",
-    ),
-    (
-        "deferred_norm_evict_hints",
-        "deferred_hop",
-        f"{SCOPE}=shared: layer 6 is a full-indexer layer; post_moe_hop hop",
-    ),
-    ("sparse_mla_decode", "layer5", "not a DSA layer"),
-)
-
-# One decline per rule and site kind: (path, key, reason, class).
-RULES = (
-    # A path turned off.
-    ("sampler_store", "engaged", "turned off by TRTLLM_STAIRCASE_GLM_LL_DISABLE", D.DISABLED),
-    # A key no site names, and a path that does not exist.
-    ("routed_moe", "a_key_no_site_names", "workspace failed", D.UNCOVERED),
-    ("not_a_fast_path", "layer3", "workspace failed", D.UNCOVERED),
-    # Decided per call: a missing resource is a miss that is never raised.
-    (
-        "routed_moe",
-        "rows=6 not warmed up before CUDA graph capture",
-        "not warmed up before CUDA graph capture",
-        D.MISSED_PER_CALL,
-    ),
-    (
-        "rope_kv_append",
-        "scatter_defer",
-        "call site does not forward the deferred scatter",
-        D.MISSED_PER_CALL,
-    ),
-    (
-        "draft_front",
-        "draft_front_call_fallback",
-        "kernel not compiled before CUDA graph capture",
-        D.MISSED_PER_CALL,
-    ),
-    ("decode_prologue", "take_embedding", "packet does not match", D.MISSED_PER_CALL),
-    (
-        "indexer_layer_front",
-        "cp_indexer_decode_global_emit",
-        "workspace unavailable: peer rank",
-        D.MISSED_PER_CALL,
-    ),
-    (
-        "dense_mlp",
-        "dense_mlp:glm_moe_dsa_dense_mlp op not built",
-        "glm_moe_dsa_dense_mlp op not built",
-        D.MISSED_PER_CALL,
-    ),
-    ("indexer_topk", "indexer_decode", "CUTLASS DSL unavailable", D.MISSED_PER_CALL),
-    # Decided per call, outside the shapes (rows, metadata layout, step kind, KV band).
-    ("rope_kv_append", "rope", "outside the FP8 MLA decode domain", D.UNCOVERED),
-    ("sparse_mla_decode", "request:KV length layout", "KV length layout; first seen", D.UNCOVERED),
-    (
-        "decode_prologue",
-        "target_prologue[metadata is not DSAtrtllmAttentionMetadata]",
-        "metadata is not DSAtrtllmAttentionMetadata",
-        D.UNCOVERED,
-    ),
-    ("draft_tail", "draft_tail[graph]", "draft logits are not vocab-sharded", D.UNCOVERED),
-    (
-        "draft_tail",
-        "draft_tail[step: 9 rows outside [1, 8]]",
-        "identical on every TP rank",
-        D.UNCOVERED,
-    ),
-    ("sparse_mla_decode", "call:layer5", "first call", D.UNCOVERED),
-    # Decided once: a layer kind or a composition the path does not serve.
-    ("attention_output_into_hop", "layer3", "needs the post_attention_hop hop", D.UNCOVERED),
-    (
-        "attention_output_tail_prefetch",
-        "layer3",
-        "post_attention_hop without sparse_mla_decode: native tail is faster",
-        D.UNCOVERED,
-    ),
-    (
-        "attention_output_tail_prefetch",
-        "schedule",
-        "post_attention_hop without sparse_mla_decode",
-        D.UNCOVERED,
-    ),
-    (
-        "draft_live_rows",
-        "live_row_hop:no post_attention_hop",
-        "no post_attention_hop hop",
-        D.UNCOVERED,
-    ),
-    ("sparse_mla_decode", "layer40", "not MLA", D.UNCOVERED),
-    # Decided once, on a layer outside the cells.
-    ("fused_kv_indexer_projection", "layer4", "kernel compile failed", D.UNCOVERED),
-    ("attention_output_tail", "layer78", "staging buffer", D.UNCOVERED),
-    # Decided once, inside the cells or on a key without a layer: a miss.
-    ("fused_kv_indexer_projection", "layer3", "kernel compile failed", D.MISSED),
-    ("sparse_mla_decode", "layer12", "latent dims", D.MISSED),
-    ("sparse_mla_decode", "layer78", "KV cache is not FP8", D.MISSED),
-    ("sparse_mla_decode", "layer0", "rope_append disabled", D.MISSED),
-    ("indexer_layer_front", "rope", "indexer_layer_front rope op not loaded", D.MISSED),
-    (
-        "post_attention_hop",
-        "post_attn_hop:native all-reduce is not MNNVL",
-        "native all-reduce is not MNNVL",
-        D.MISSED,
-    ),
-    ("post_attention_hop", "post_attn_hop", "", D.MISSED),
-    ("router_gemm", "router:attention DP", "attention DP", D.MISSED),
-    ("post_moe_hop", "cluster_kernel", "the cluster kernel needs SM100", D.MISSED),
-    ("post_moe_hop", "post_moe_hop[cluster]", "workspace agreement failed", D.MISSED),
-    ("routed_moe_into_hop", "post_moe_routed_input_buffer", "no buffer", D.MISSED),
-    (
-        "cross_layer_expert_prefetch",
-        "link[k=8,start=moe]",
-        "0 MoE layers, 0 routed-kernel prefetch hooks",
-        D.MISSED,
-    ),
-    ("draft_front", "draft_front", "requires non-ADP tensor parallelism", D.MISSED),
-    ("draft_tail", "draft_tail[tp8]", "workspace setup failed: OSError()", D.MISSED),
-    ("host_page_retention", "host_policy", "libc10.so: not found", D.MISSED),
-    # The loop domain is final only once the overlap loop has declared it.
-    (
-        "deferred_input_prep",
-        "loop_domain",
-        "no overlap executor loop declared its domain; the store writer and the single-stream "
-        "loop are validated only there",
-        D.UNCOVERED,
-    ),
-    (
-        "deferred_input_prep",
-        "loop_domain",
-        "attention DP: the store writer and the single-stream loop are not validated with it",
-        D.MISSED,
-    ),
-)
-
-
-@pytest.fixture
-def capability(monkeypatch):
-    """A fresh capability log and fallback state; returns the info and warning lines."""
-    info, warnings = [], []
-
-    class _Logger:
-        @staticmethod
-        def info(message):
-            info.append(message)
-
-        @staticmethod
-        def warning(message):
-            warnings.append(message)
-
-    monkeypatch.setattr(dispatch, "_log_info", info.append)
-    monkeypatch.setattr(dispatch, "logger", _Logger)
-    monkeypatch.setattr(dispatch, "_CAPABILITY_LOGGED", {})
-    dispatch.reset_fallback_state()
-    dispatch.set_require(False)
-    dispatch.classify_decline.cache_clear()
-    yield info, warnings
-    dispatch.reset_fallback_state()
-    dispatch.set_require(False)
-
-
-# --- the tables ---------------------------------------------------------------------------
-
-
-def test_the_layer_cells_are_the_census():
-    expected = {}
-    for path, family, layers in CENSUS:
-        expected.setdefault(path, {})[family] = _layers(layers)
-    assert {path: dict(families) for path, families in dispatch.LAYER_CELLS.items()} == expected
-
-
-def test_every_site_names_a_fast_path_and_every_fast_path_has_a_site():
-    assert {site.path for site in dispatch.DECLINE_KINDS} == set(dispatch.FAST_PATH_NAMES)
-
-
-def test_every_family_with_cells_is_a_site_of_its_path():
-    families = {(site.path, site.cells) for site in dispatch.DECLINE_KINDS if site.cells}
-    with_cells = {
-        (path, family) for path, cells in dispatch.LAYER_CELLS.items() for family in cells
-    }
-    assert with_cells <= families
-    # The per-layer keys the served composition never logs have no cells: every decline there
-    # is outside the path's shapes.
-    assert families - with_cells == {
-        ("deferred_norm", "projection_layer"),
-        ("attention_output_tail_prefetch", "layer"),
-        ("attention_output_into_hop", "layer"),
-    }
-
-
-def test_a_site_with_cells_captures_the_layer():
-    for site in dispatch.DECLINE_KINDS:
-        assert ("layer" in site.key.groupindex) == (site.cells is not None), site
-
-
-def _kernel_sources() -> list:
-    """The sources that log declines: every file of the family except the classifier itself."""
-    return [p for p in sorted(FAMILY.rglob("*.py")) if p != Path(dispatch.__file__)]
-
-
-def _source_text() -> str:
-    return "\n".join(p.read_text() for p in _kernel_sources())
-
-
-@pytest.mark.parametrize("fragment", dispatch.LAYER_KIND_REASONS + dispatch.DISABLED_REASONS)
-def test_every_reason_fragment_is_a_reason_the_kernels_give(fragment):
-    """A renamed reason must rename its fragment too, or its declines change class."""
-    assert fragment in _source_text()
-
-
-def _string_literals() -> set:
-    return {
-        node.value
-        for source in _kernel_sources()
-        for node in ast.walk(ast.parse(source.read_text()))
-        if isinstance(node, ast.Constant) and isinstance(node.value, str)
-    }
-
-
-def test_the_served_reasons_are_the_kernels_literals():
-    """Each constant reason of a served decline is a whole string literal of the kernels, so a
-    reworded reason fails here before its declines change class."""
-    literals = _string_literals()
-    # The reasons without a formatted field (a number or a knob name).
-    reasons = {
-        reason
-        for _, _, reason in SERVED_DECLINES
-        if not any(c.isdigit() for c in reason) and dispatch.KNOB_PREFIX not in reason
-    }
-    assert len(reasons) == 13
-    assert sorted(reasons - literals) == []
-
-
-def _render(node: ast.AST):
-    """A key the call can log: f-string fields become a sample value."""
-    if isinstance(node, ast.Constant):
-        return node.value
-    if not isinstance(node, ast.JoinedStr):
-        return None
-    out = []
-    for value in node.values:
-        if isinstance(value, ast.Constant):
-            out.append(value.value)
-        elif isinstance(value.value, ast.IfExp) and isinstance(value.value.body, ast.Constant):
-            out.append(value.value.body.value)
-        elif isinstance(value.value, ast.Attribute) and value.value.attr == "scope":
-            out.append("shared")
-        else:
-            out.append("0")
-    return "".join(out)
-
-
-def test_every_logged_key_has_a_site():
-    """Each ``log_capability`` call whose path and key can be read off the source logs a key
-    that a site of its path matches."""
-    checked, missing = 0, []
-    for source in _kernel_sources():
-        tree = ast.parse(source.read_text())
-        constants = {
-            node.targets[0].id: node.value.value
-            for node in tree.body
-            if isinstance(node, ast.Assign)
-            and len(node.targets) == 1
-            and isinstance(node.targets[0], ast.Name)
-            and isinstance(node.value, ast.Constant)
-            and isinstance(node.value.value, str)
-        }
-        for node in ast.walk(tree):
-            if not (isinstance(node, ast.Call) and len(node.args) >= 2):
-                continue
-            name = getattr(node.func, "attr", getattr(node.func, "id", None))
-            if name != "log_capability":
-                continue
-            first = node.args[0]
-            path = (
-                first.value
-                if isinstance(first, ast.Constant)
-                else constants.get(getattr(first, "id", None))
-            )
-            key = _render(node.args[1])
-            if path is None or key is None:
-                continue
-            checked += 1
-            if dispatch._decline_site(path, key) is None:
-                missing.append(f"{source.name}:{node.lineno} {path} {key!r}")
-    assert missing == []
-    assert checked >= 70
-
-
-# --- the rules ----------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("path, key, reason", SERVED_DECLINES)
-def test_no_decline_of_the_served_composition_is_a_miss(path, key, reason):
-    assert dispatch.classify_decline(path, key, reason) is D.UNCOVERED
-
-
-@pytest.mark.parametrize("path, key, reason, kind", RULES)
-def test_each_rule(path, key, reason, kind):
-    assert dispatch.classify_decline(path, key, reason) is kind
-
-
-def test_a_resource_reason_is_matched_in_any_case():
-    assert (
-        dispatch.classify_decline("routed_moe", "rows=4 x", "Workspace FAILED") is D.MISSED_PER_CALL
-    )
-
-
-def test_a_repeated_decline_is_one_lookup(capability):
-    args = ("routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]")
-    for _ in range(5):
-        dispatch.log_capability(*args[:2], False, args[2])
-    info = dispatch.classify_decline.cache_info()
-    assert (info.misses, info.hits) == (1, 4)
-
-
-# --- what log_capability does with each class ---------------------------------------------
-
-
-def test_an_engaged_decision_is_not_classified(capability):
-    assert dispatch.log_capability("post_moe_hop", "post_moe_hop[cluster]", True)
-    assert dispatch.classify_decline.cache_info().misses == 0
-
-
-def test_a_miss_decided_once_raises_under_require_after_its_line(capability):
-    info, warnings = capability
-    dispatch.set_require(True)
-    with pytest.raises(
-        dispatch.FastPathUnavailable, match=r"'post_moe_hop'.*post_moe_hop\[cluster\]: workspace"
-    ):
-        dispatch.log_capability("post_moe_hop", "post_moe_hop[cluster]", False, "workspace failed")
-    assert info == [
-        "[glm_moe_dsa] post_moe_hop/post_moe_hop[cluster]: native fallback (workspace failed)"
-    ]
-    assert warnings == []
-
-
-def test_a_miss_decided_once_is_counted_and_warned_once_under_auto(capability):
-    info, warnings = capability
-    for _ in range(2):
-        assert not dispatch.log_capability("sparse_mla_decode", "layer7", False, "latent dims")
-    assert dispatch.missed_counts() == {"sparse_mla_decode": 2}
-    assert warnings == [
-        "[glm_moe_dsa] sparse_mla_decode/missed: layer7: latent dims; falling back (logged once per reason)"
-    ]
-    assert len(info) == 1
-
-
-@pytest.mark.parametrize("require", [False, True], ids=["auto", "require"])
-def test_a_miss_decided_per_call_is_counted_and_never_raised(capability, require):
-    info, warnings = capability
-    dispatch.set_require(require)
-    key, reason = (
-        "rows=6 not warmed up before CUDA graph capture",
-        "not warmed up before CUDA graph capture",
-    )
-    for _ in range(3):
-        assert not dispatch.log_capability("routed_moe", key, False, reason)
-    assert dispatch.missed_counts() == {"routed_moe": 3}
-    assert warnings == [
-        f"[glm_moe_dsa] routed_moe/missed: {reason}; falling back (logged once per reason)"
-    ]
-    assert info == [f"[glm_moe_dsa] routed_moe/{key}: native fallback ({reason})"]
-
-
-@pytest.mark.parametrize(
-    "path, key, reason",
-    [
-        ("absorbed_query", "layer78", "target layers only"),
-        ("routed_moe", "rows=128 128 rows outside [1, 8]", "128 rows outside [1, 8]"),
-        ("sampler_store", "engaged", "turned off by TRTLLM_STAIRCASE_GLM_LL_DISABLE"),
-        ("routed_moe", "a_key_no_site_names", "workspace failed"),
-    ],
-    ids=["uncovered-once", "uncovered-per-call", "disabled", "unknown-site"],
-)
-@pytest.mark.parametrize("require", [False, True], ids=["auto", "require"])
-def test_uncovered_and_disabled_declines_add_nothing_to_their_line(
-    capability, path, key, reason, require
-):
-    info, warnings = capability
-    dispatch.set_require(require)
-    for _ in range(2):
-        assert not dispatch.log_capability(path, key, False, reason)
-    assert info == [f"[glm_moe_dsa] {path}/{key}: native fallback ({reason})"]
-    assert warnings == [] and dispatch.missed_counts() == {}
-
-
-def test_the_served_declines_are_not_raised_under_require(capability):
-    """Under require, the served composition's declines run as before: logged, not raised."""
-    dispatch.set_require(True)
-    for path, key, reason in SERVED_DECLINES:
-        dispatch.log_capability(path, key, False, reason)
-    _, warnings = capability
-    assert warnings == [] and dispatch.missed_counts() == {}
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py
deleted file mode 100644
index 518ffcd12525..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/shared/test_modeling_v2_glm_dispatch.py
+++ /dev/null
@@ -1,380 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The composition of the GLM-5.2 tp8 target: the fast-path table, the DISABLE knob, the
-fallback policy and the rank-uniform checks.
-
-Host-only: the knobs are passed as an explicit environment and the collective is a fake
-two-rank group or a fake MPI communicator, so nothing here needs a GPU or MPI.
-"""
-
-from __future__ import annotations
-
-import pytest
-
-import tensorrt_llm._utils as trtllm_utils
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode
-from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.shared import (  # noqa: E501
-    dispatch,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-AUTO, REQUIRE = ModelingV2Mode.AUTO, ModelingV2Mode.REQUIRE
-DISABLE = dispatch.DISABLE_ENV
-ACCUMULATION = dispatch.KNOB_PREFIX + "ROUTED_MOE_INTO_HOP_ACCUMULATION"
-
-# The feature vocabulary, in this order: the fast-path names, the DISABLE tokens and the
-# knob infixes are the same words.
-VOCABULARY = (
-    "absorbed_query",
-    "draft_absorbed_query",
-    "shared_indexer_absorbed_query",
-    "fused_kv_indexer_projection",
-    "indexer_layer_front",
-    "deferred_norm",
-    "deferred_norm_evict_hints",
-    "rope_kv_append",
-    "indexer_topk",
-    "cross_rank_topk",
-    "cross_rank_topk_cluster_merge",
-    "cross_rank_topk_global_emit",
-    "sparse_mla_decode",
-    "attention_output_tail",
-    "attention_output_tail_prefetch",
-    "attention_output_into_hop",
-    "oproj_weight_staging",
-    "post_attention_hop",
-    "hop_routing",
-    "router_gemm",
-    "routed_moe",
-    "routed_moe_into_hop",
-    "shared_expert_prefetch",
-    "post_moe_hop",
-    "cross_layer_expert_prefetch",
-    "dense_mlp",
-    "draft_front",
-    "draft_routed_moe",
-    "draft_live_rows",
-    "draft_tail",
-    "draft_loop",
-    "decode_prologue",
-    "deferred_input_prep",
-    "sampler_store",
-    "execution_stream_loop",
-    "host_page_retention",
-)
-
-
-@pytest.fixture(autouse=True)
-def _fresh_fallback_state():
-    dispatch.set_require(False)
-    dispatch.reset_fallback_state()
-    yield
-    dispatch.set_require(False)
-    dispatch.reset_fallback_state()
-
-
-def _off(plan: dispatch.Plan) -> set:
-    return {name for name in dispatch.FAST_PATH_NAMES if not plan.enabled(name)}
-
-
-# --- the table --------------------------------------------------------------------------
-
-
-def test_the_fast_paths_are_the_vocabulary_in_order():
-    """The DISABLE tokens are these words; the log line lists them in this order."""
-    assert dispatch.FAST_PATH_NAMES == VOCABULARY
-
-
-def test_dependencies_name_fast_paths_and_have_no_cycle():
-    names = set(dispatch.FAST_PATH_NAMES)
-    for path in dispatch.FAST_PATHS:
-        assert set(path.depends_on) <= names, path
-        assert path.name not in path.depends_on, path
-        # A base that is also a (transitive) dependent of the path would be a cycle.
-        dependents = dispatch.disable_closure(frozenset({path.name})) - {path.name}
-        assert not set(path.depends_on) & dependents, path
-
-
-# --- DISABLE ----------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("raw", [None, "", " , ,"], ids=["unset", "empty", "blank"])
-def test_every_fast_path_runs_by_default(raw):
-    environ = {} if raw is None else {DISABLE: raw}
-    plan = dispatch.resolve(AUTO, environ)
-    assert _off(plan) == set()
-    assert plan.disabled == plan.disabled_dependents == frozenset()
-
-
-def test_disabling_a_path_without_dependents_turns_off_only_it():
-    plan = dispatch.resolve(AUTO, {DISABLE: "draft_tail"})
-    assert _off(plan) == {"draft_tail"}
-    assert plan.disabled_dependents == frozenset()
-
-
-@pytest.mark.parametrize(
-    "disabled, dependents",
-    [
-        (
-            "routed_moe",
-            {
-                "routed_moe_into_hop",
-                "hop_routing",
-                "shared_expert_prefetch",
-                "cross_layer_expert_prefetch",
-            },
-        ),
-        (
-            "post_attention_hop",
-            {
-                "attention_output_into_hop",
-                "oproj_weight_staging",
-                "hop_routing",
-                "shared_expert_prefetch",
-                "draft_live_rows",
-                "host_page_retention",
-            },
-        ),
-        (
-            "cross_rank_topk",
-            {
-                "cross_rank_topk_cluster_merge",
-                "cross_rank_topk_global_emit",
-                "indexer_layer_front",
-            },
-        ),
-        ("draft_loop", {"draft_live_rows", "decode_prologue"}),
-    ],
-)
-def test_disabling_a_path_turns_off_every_path_that_depends_on_it(disabled, dependents):
-    plan = dispatch.resolve(AUTO, {DISABLE: disabled})
-    assert plan.disabled == {disabled}
-    assert plan.disabled_dependents == dependents
-    assert _off(plan) == {disabled} | dependents
-
-
-def test_disable_ignores_case_and_blanks():
-    plan = dispatch.resolve(AUTO, {DISABLE: " Draft_Tail ,, DENSE_MLP "})
-    assert _off(plan) == {"draft_tail", "dense_mlp"}
-
-
-def test_disable_naming_something_else_raises():
-    with pytest.raises(ValueError, match=r"not_a_fast_path") as excinfo:
-        dispatch.resolve(AUTO, {DISABLE: "draft_tail,not_a_fast_path"})
-    assert DISABLE in str(excinfo.value)
-
-
-def test_a_call_site_asking_for_a_misspelled_path_fails():
-    plan = dispatch.resolve(AUTO, {})
-    with pytest.raises(KeyError, match="draft_tails"):
-        plan.enabled("draft_tails")
-
-
-def test_only_this_targets_knobs_are_recorded():
-    plan = dispatch.resolve(
-        AUTO, {DISABLE: "draft_tail", ACCUMULATION: "identity", "UNRELATED": "1"}
-    )
-    assert plan.knobs == ((DISABLE, "draft_tail"), (ACCUMULATION, "identity"))
-
-
-# --- missed and uncovered ---------------------------------------------------------------
-
-
-def test_a_miss_raises_under_require():
-    dispatch.resolve(REQUIRE, {})
-    with pytest.raises(dispatch.FastPathUnavailable, match=r"'routed_moe'.*at 6 rows.*workspace"):
-        dispatch.missed("routed_moe", 6, "its workspace is not built")
-
-
-def test_a_miss_falls_back_and_is_counted_under_auto():
-    dispatch.resolve(AUTO, {})
-    dispatch.missed("routed_moe", None, "its workspace is not built")
-    dispatch.missed("routed_moe", 6, "its workspace is not built")
-    assert dispatch.missed_counts() == {"routed_moe": 2}
-
-
-@pytest.mark.parametrize("mode", [AUTO, REQUIRE])
-def test_an_uncovered_step_never_raises(mode):
-    dispatch.resolve(mode, {})
-    dispatch.uncovered("sparse_mla_decode", "more rows than the kernel takes")
-    assert dispatch.missed_counts() == {}
-
-
-@pytest.mark.parametrize("mode", [AUTO, REQUIRE])
-def test_a_miss_decided_per_call_never_raises_and_is_counted(mode):
-    dispatch.resolve(mode, {})
-    dispatch.missed("routed_moe", 6, "not warmed up before CUDA graph capture", per_call=True)
-    dispatch.missed("routed_moe", 6, "not warmed up before CUDA graph capture", per_call=True)
-    assert dispatch.missed_counts() == {"routed_moe": 2}
-
-
-# --- rank uniformity --------------------------------------------------------------------
-
-
-def _two_rank_group(peer_environ: dict, calls: list):
-    """A fake allgather over two ranks: this rank first, then a peer with ``peer_environ``."""
-
-    def allgather(knobs, world_size):
-        calls.append((knobs, world_size))
-        return [knobs, dispatch.read_knobs(peer_environ)]
-
-    return allgather
-
-
-def test_ranks_that_agree_pass_with_one_collective():
-    env = {DISABLE: "draft_tail", ACCUMULATION: "identity"}
-    calls = []
-    knobs = dispatch.read_knobs(env)
-    dispatch.check_rank_uniform(knobs, 2, _two_rank_group(dict(env), calls))
-    assert calls == [(knobs, 2)]
-
-
-def test_unrelated_environment_does_not_count():
-    dispatch.check_rank_uniform(
-        dispatch.read_knobs({"UNRELATED": "a"}), 2, _two_rank_group({"UNRELATED": "b"}, [])
-    )
-
-
-def test_ranks_that_would_run_different_fast_paths_raise_with_the_rank_table():
-    with pytest.raises(RuntimeError) as excinfo:
-        dispatch.check_rank_uniform(
-            dispatch.read_knobs({}), 2, _two_rank_group({DISABLE: "routed_moe"}, [])
-        )
-    message = str(excinfo.value)
-    rank0, rank1 = message.split("rank 1:")
-    assert "rank 0: knobs" in rank0 and f"{DISABLE}='<unset>'" in rank0
-    assert f"{DISABLE}='routed_moe'" in rank1
-
-
-def test_ranks_that_read_different_knob_values_raise():
-    with pytest.raises(RuntimeError, match=rf"{ACCUMULATION}='identity'"):
-        dispatch.check_rank_uniform(
-            dispatch.read_knobs({ACCUMULATION: "fp32"}),
-            2,
-            _two_rank_group({ACCUMULATION: "identity"}, []),
-        )
-
-
-def test_a_value_one_rank_cannot_parse_fails_the_check_on_every_rank():
-    """The check compares the raw values, so the rank with a valid value raises too."""
-    with pytest.raises(RuntimeError, match=rf"{DISABLE}='not_a_fast_path'"):
-        dispatch.check_rank_uniform(
-            dispatch.read_knobs({DISABLE: "draft_tail"}),
-            2,
-            _two_rank_group({DISABLE: "not_a_fast_path"}, []),
-        )
-
-
-def test_without_a_communicator_the_check_is_skipped():
-    dispatch.check_rank_uniform(dispatch.read_knobs({}), 2, lambda knobs, world_size: None)
-
-
-def test_the_ranks_agree_on_the_knobs_the_shared_switches_and_the_mode():
-    env = {
-        DISABLE: "draft_tail",
-        "TRTLLM_FORCE_MNNVL_AR": "1",
-        "TLLM_TEST_MNNVL": "0",
-        "TRTLLM_ENABLE_PDL": "1",
-        "TRTLLM_MODELING_V2_VALIDATE": "1",
-        "UNRELATED": "1",
-    }
-    inputs = dispatch.rank_uniform_inputs(dispatch.read_knobs(env), REQUIRE, env)
-    assert inputs == (
-        ("TLLM_TEST_MNNVL", "0"),
-        ("TRTLLM_ENABLE_PDL", "1"),
-        ("TRTLLM_FORCE_MNNVL_AR", "1"),
-        ("TRTLLM_MODELING_V2_VALIDATE", "1"),
-        (DISABLE, "draft_tail"),
-        (dispatch.MODE_KEY, "require"),
-    )
-    assert dispatch.rank_uniform_inputs((), AUTO, {}) == ((dispatch.MODE_KEY, "auto"),)
-
-
-@pytest.mark.parametrize(
-    "peer_env, peer_mode",
-    [({"TRTLLM_FORCE_MNNVL_AR": "1"}, REQUIRE), ({}, AUTO)],
-    ids=["shared-switch", "mode"],
-)
-def test_ranks_that_differ_in_a_shared_switch_or_the_mode_raise(peer_env, peer_mode):
-    mine = dispatch.rank_uniform_inputs((), REQUIRE, {})
-    peer = dispatch.rank_uniform_inputs((), peer_mode, peer_env)
-    with pytest.raises(RuntimeError, match="rank 1: knobs") as excinfo:
-        dispatch.check_rank_uniform(mine, 2, lambda value, world_size: [value, peer])
-    assert dispatch.MODE_KEY in str(excinfo.value)
-
-
-_STATE = (
-    ("model.layers.0", (("post_attn_hop", "PostAttentionHop"),)),
-    ("model.layers.1", (("post_moe_hop", "PostMoeHop", "engaged=True"),)),
-)
-
-
-def test_ranks_that_attached_the_same_state_pass_with_one_collective():
-    calls = []
-
-    def allgather(value, world_size):
-        calls.append(world_size)
-        return [value, value]
-
-    dispatch.check_rank_uniform_adopted(_STATE, 2, allgather)
-    assert calls == [2]
-
-
-def test_ranks_that_attached_different_state_raise_naming_the_first_module():
-    peer = (_STATE[0], ("model.layers.1", ()))
-    with pytest.raises(RuntimeError) as excinfo:
-        dispatch.check_rank_uniform_adopted(_STATE, 2, lambda value, world_size: [value, peer])
-    message = str(excinfo.value)
-    assert (
-        "rank 0: ('model.layers.1', (('post_moe_hop', 'PostMoeHop', 'engaged=True'),))" in message
-    )
-    assert "rank 1: ('model.layers.1', ())" in message
-    assert "model.layers.0" not in message
-
-
-def test_a_rank_with_fewer_modules_is_named():
-    with pytest.raises(RuntimeError, match=r"rank 1: <no module>"):
-        dispatch.check_rank_uniform_adopted(_STATE, 2, lambda value, world_size: [value, value[:1]])
-
-
-def test_without_a_communicator_the_adopted_check_is_skipped():
-    dispatch.check_rank_uniform_adopted(_STATE, 2, lambda value, world_size: None)
-
-
-class _Communicator:
-    def __init__(self, size: int) -> None:
-        self.size = size
-        self.gathered = []
-
-    def Get_size(self) -> int:  # noqa: N802 - the mpi4py name
-        return self.size
-
-    def allgather(self, knobs):
-        self.gathered.append(knobs)
-        return [knobs] * self.size
-
-
-def test_the_session_communicator_of_this_model_is_gathered_once(monkeypatch):
-    comm = _Communicator(8)
-    monkeypatch.setattr(trtllm_utils, "mpi_disabled", lambda: False)
-    monkeypatch.setattr(trtllm_utils, "mpi_comm", lambda: comm)
-    knobs = dispatch.read_knobs({DISABLE: "draft_tail"})
-    dispatch.check_rank_uniform(knobs, 8)
-    assert comm.gathered == [knobs]
-
-
-@pytest.mark.parametrize("size", [1, 16], ids=["singleton", "larger"])
-def test_a_communicator_that_is_not_this_models_ranks_is_not_gathered(monkeypatch, size):
-    """A singleton proves nothing; a larger one would wait on ranks that never build the model."""
-    comm = _Communicator(size)
-    monkeypatch.setattr(trtllm_utils, "mpi_disabled", lambda: False)
-    monkeypatch.setattr(trtllm_utils, "mpi_comm", lambda: comm)
-    dispatch.check_rank_uniform(dispatch.read_knobs({}), 8)
-    assert comm.gathered == []
-
-
-def test_without_mpi_nothing_is_gathered(monkeypatch):
-    monkeypatch.setattr(trtllm_utils, "mpi_disabled", lambda: True)
-    monkeypatch.setattr(trtllm_utils, "mpi_comm", lambda: pytest.fail("no communicator expected"))
-    dispatch.check_rank_uniform(dispatch.read_knobs({}), 8)
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py
deleted file mode 100644
index b6e9335663b1..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_isolation.py
+++ /dev/null
@@ -1,297 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Nothing outside the glm_moe_dsa family reaches its targets, except the router line, and
-no target reaches another.
-
-With ``modeling_v2`` off the resolver returns before importing anything, so the
-target can change no other model's behaviour as long as no other module imports
-it. No shared module imports the target (``_S4_SHARED_IMPORTS`` is empty): the
-engine reaches its step-loop hooks through ``runtime_hooks``, and every other
-reference fails. Read off disk as text, like test_modeling_v2_claims.py: no
-module is imported.
-
-No shared Python or C++ file uses the family's names either (the C++ build adds the family's
-folders in a few CMake lines), and no model of the zoo declares step-loop hooks, so each runs
-with the engine's all-off default; that check imports the zoo's model modules.
-"""
-
-from __future__ import annotations
-
-import ast
-import re
-from collections import Counter
-from pathlib import Path
-
-import pytest
-
-import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2
-
-pytestmark = pytest.mark.cpu_only
-
-_MODELING_V2 = Path(_modeling_v2.__file__).resolve().parent
-_PACKAGE = _MODELING_V2.parents[2]
-_FAMILY = _MODELING_V2 / "models" / "glm_moe_dsa"
-_ROUTER_INDEX = _MODELING_V2 / "_router_index.py"
-
-# A reference to the family's modules or to the targets' classes.
-_REFERENCE = re.compile(
-    r"models[./]glm_moe_dsa\b|glm_5_2_nvfp4__sm_103__\w+|ModelingV2Glm52|GlmMoeDsa(Tp8|Tep|Dep)ModelingV2"
-)
-_ROUTER_LINE = '"GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing",'
-
-_TARGET = "tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa"
-# The imports of the target from shared modules: (file, module, name, alias) -> count.
-# They are matched on the parsed import, so the formatter may wrap them.
-_S4_SHARED_IMPORTS: dict = {}
-
-
-def _target_imports(path: Path, rel: str):
-    """``(key, first line, last line)`` of every import of the target in ``path``."""
-    out = []
-    for node in ast.walk(ast.parse(path.read_text(errors="ignore"))):
-        if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module:
-            if node.module == _TARGET or node.module.startswith(_TARGET + "."):
-                for alias in node.names:
-                    key = (rel, node.module, alias.name, alias.asname)
-                    out.append((key, node.lineno, node.end_lineno))
-    return out
-
-
-def _shared_files():
-    for path in sorted(_PACKAGE.rglob("*.py")):
-        if _FAMILY not in path.parents:
-            yield path, str(path.relative_to(_PACKAGE.parent))
-
-
-def test_only_the_router_line_references_the_family():
-    offenders = []
-    for path, rel in _shared_files():
-        lines = path.read_text(errors="ignore").splitlines()
-        hits = [n for n, line in enumerate(lines, 1) if _REFERENCE.search(line)]
-        if not hits:
-            continue
-        allowed = set()
-        for key, first, last in _target_imports(path, rel):
-            if key in _S4_SHARED_IMPORTS:
-                allowed.update(range(first, last + 1))
-        for lineno in hits:
-            line = lines[lineno - 1]
-            if path == _ROUTER_INDEX and line.strip() == _ROUTER_LINE:
-                continue
-            if lineno in allowed:
-                continue
-            offenders.append(f"{rel}:{lineno}: {line.strip()}")
-    assert not offenders, "modules outside glm_moe_dsa reference it:\n" + "\n".join(offenders)
-
-
-def test_each_listed_shared_import_is_present():
-    found = Counter()
-    for path, rel in _shared_files():
-        if "glm_moe_dsa" in path.read_text(errors="ignore"):
-            found.update(key for key, _, _ in _target_imports(path, rel))
-    assert dict(found) == _S4_SHARED_IMPORTS
-
-
-def test_the_router_line_is_present():
-    lines = [line.strip() for line in _ROUTER_INDEX.read_text().splitlines()]
-    assert lines.count(_ROUTER_LINE) == 1
-
-
-# The family's names; no file outside the family uses them.
-_FORK_NAMES = re.compile(r"glm_ll_|\[glm_ll\]|GLM_LL|TRTLLM_STAIRCASE")
-_CPP_FORK_NAMES = re.compile(r"glm_ll|GLM_LL|glmMoeDsa")
-_REPO = _PACKAGE.parent
-_CPP = _REPO / "cpp" / "tensorrt_llm"
-_CPP_FAMILY_DIRS = ("kernels/glmMoeDsa", "thop/glmMoeDsa")
-# The lines each parent CMakeLists.txt may spend on the family's folder.
-_CMAKE_LINES = 3
-
-
-def test_no_shared_file_names_the_fork():
-    offenders = []
-    for path, rel in _shared_files():
-        for lineno, line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
-            if _FORK_NAMES.search(line):
-                offenders.append(f"{rel}:{lineno}: {line.strip()}")
-    assert not offenders, "modules outside glm_moe_dsa name the fork:\n" + "\n".join(offenders)
-
-
-def test_no_shared_cpp_file_names_the_family():
-    if not _CPP.is_dir():
-        pytest.skip("the C++ sources are not in this tree")
-    family = [_CPP / d for d in _CPP_FAMILY_DIRS]
-    offenders, cmake_lines = [], Counter()
-    for path in sorted(_CPP.rglob("*")):
-        if not path.is_file() or path.suffix not in (".cpp", ".cu", ".h", ".txt"):
-            continue
-        if any(d == path.parent or d in path.parents for d in family):
-            continue
-        rel = str(path.relative_to(_REPO))
-        for lineno, line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
-            if not _CPP_FORK_NAMES.search(line):
-                continue
-            if path.name == "CMakeLists.txt" and path.parent in {d.parent for d in family}:
-                cmake_lines[rel] += 1
-                continue
-            offenders.append(f"{rel}:{lineno}: {line.strip()}")
-    assert not offenders, "C++ files outside the family's folders name it:\n" + "\n".join(offenders)
-    assert all(count <= _CMAKE_LINES for count in cmake_lines.values()), dict(cmake_lines)
-
-
-def _zoo_model_classes():
-    """Class name -> providing module of every model of the zoo (imported lazily)."""
-    from tensorrt_llm._torch.models._arch_index import MODEL_CLASS_TO_MODULE
-
-    classes = dict(MODEL_CLASS_TO_MODULE)
-    classes.setdefault("DeepseekV3ForCausalLM", "modeling_deepseekv3")
-    classes.setdefault("GlmMoeDsaForCausalLM", "modeling_deepseekv3")
-    return classes
-
-
-_ZOO = _zoo_model_classes()
-
-
-def test_no_zoo_model_declares_runtime_hooks():
-    """Read off disk: no model of the zoo names the hooks, so each gets the all-off default."""
-    models = _PACKAGE / "_torch" / "models"
-    offenders = [
-        str(path.relative_to(_PACKAGE.parent))
-        for path in sorted(models.rglob("*.py"))
-        if "runtime_hooks" in path.read_text(errors="ignore")
-    ]
-    assert not offenders, offenders
-
-
-@pytest.mark.parametrize("name", sorted(_ZOO))
-def test_zoo_models_declare_no_runtime_hooks(name):
-    """Every model class of the zoo, without its constructor, reads as the all-off hooks."""
-    import importlib
-
-    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, runtime_hooks_of
-
-    try:
-        providing = importlib.import_module(f"tensorrt_llm._torch.models.{_ZOO[name]}")
-    except ModuleNotFoundError as err:
-        if err.name and err.name.split(".")[0] != "tensorrt_llm":
-            pytest.skip(f"{err.name} is not installed")
-        raise
-    cls = getattr(providing, name)
-    assert runtime_hooks_of(cls.__new__(cls)) is DEFAULT_RUNTIME_HOOKS
-
-
-# --- the targets ----------------------------------------------------------------------------
-
-_TARGET_DIRS = {
-    "tp8": _FAMILY / "glm_5_2_nvfp4__sm_103__tp8",
-    "tep": _FAMILY / "glm_5_2_nvfp4__sm_103__tep",
-    "dep": _FAMILY / "glm_5_2_nvfp4__sm_103__dep",
-}
-
-
-def _module_file(module: str):
-    """The file of a dotted module of the family, or None outside it."""
-    if not (module == _TARGET or module.startswith(_TARGET + ".")):
-        return None
-    base = (
-        _FAMILY.joinpath(*module[len(_TARGET) :].lstrip(".").split("."))
-        if module != _TARGET
-        else _FAMILY
-    )
-    for candidate in (base.with_suffix(".py"), base / "__init__.py"):
-        if candidate.is_file():
-            return candidate
-    return None
-
-
-def _package_of(path: Path) -> str:
-    parts = path.relative_to(_FAMILY).with_suffix("").parts
-    if parts[-1] == "__init__":
-        parts = parts[:-1]
-    return (
-        ".".join((_TARGET, *parts[:-1]))
-        if path.name != "__init__.py"
-        else ".".join((_TARGET, *parts))
-    )
-
-
-def _imports(path: Path):
-    """Every family module ``path`` imports, at module level or inside a function."""
-    found = set()
-    for node in ast.walk(ast.parse(path.read_text())):
-        if isinstance(node, ast.ImportFrom):
-            if node.level:
-                package = _package_of(path).split(".")
-                base = ".".join(package[: len(package) - node.level + 1])
-                module = f"{base}.{node.module}" if node.module else base
-            else:
-                module = node.module or ""
-            targets = [module] + [f"{module}.{alias.name}" for alias in node.names]
-        elif isinstance(node, ast.Import):
-            targets = [alias.name for alias in node.names]
-        else:
-            continue
-        for name in targets:
-            file = _module_file(name)
-            if file is not None:
-                found.add(file)
-    return found
-
-
-def _reachable(start: Path):
-    """Every family file ``start`` imports, transitively, with the parent packages Python
-    imports first."""
-    seen, todo = set(), [start]
-    while todo:
-        path = todo.pop()
-        if path in seen:
-            continue
-        seen.add(path)
-        for parent in path.parents:
-            if parent == _FAMILY.parent:
-                break
-            if (parent / "__init__.py").is_file() and parent / "__init__.py" not in seen:
-                todo.append(parent / "__init__.py")
-        todo.extend(_imports(path))
-    return seen
-
-
-@pytest.mark.parametrize(
-    "target, other",
-    [(t, o) for t in _TARGET_DIRS for o in _TARGET_DIRS if t != o],
-)
-def test_a_target_imports_nothing_of_the_other(target, other):
-    """Read off disk: the import closure of a target's modeling module stays in the target's
-    own directory, the family's shared package and the routing module."""
-    reached = _reachable(_TARGET_DIRS[target] / "modeling.py")
-    assert not [p for p in reached if _TARGET_DIRS[other] in p.parents], sorted(
-        str(p.relative_to(_FAMILY)) for p in reached if _TARGET_DIRS[other] in p.parents
-    )
-    allowed = (_TARGET_DIRS[target], _FAMILY / "shared")
-    stray = [p for p in reached if p.parent != _FAMILY and not any(d in p.parents for d in allowed)]
-    assert not stray, sorted(str(p.relative_to(_FAMILY)) for p in stray)
-
-
-def test_the_shared_package_imports_no_target():
-    shared = _FAMILY / "shared"
-    for path in sorted(shared.rglob("*.py")):
-        for imported in _imports(path):
-            assert shared in imported.parents or imported.parent == _FAMILY, (path, imported)
-
-
-def test_the_tep_target_loads_no_kernel_package_of_the_tp8_target():
-    """``get_impl`` imports a path's package from the installed target's ``kernel_modules``;
-    the tep target names none, so every package it loads is the family's."""
-    source = (_TARGET_DIRS["tep"] / "modeling.py").read_text()
-    assert "kernel_modules" not in source
-    assert "glm_5_2_nvfp4__sm_103__tp8" not in "".join(
-        p.read_text() for p in _TARGET_DIRS["tep"].glob("*.py")
-    )
-
-
-def test_the_dep_target_loads_no_kernel_package_of_another_target():
-    """The dep target names no ``kernel_modules`` either, and no other target's directory."""
-    source = (_TARGET_DIRS["dep"] / "modeling.py").read_text()
-    assert "kernel_modules" not in source
-    text = "".join(p.read_text() for p in _TARGET_DIRS["dep"].glob("*.py"))
-    assert "glm_5_2_nvfp4__sm_103__tp8" not in text
-    assert "glm_5_2_nvfp4__sm_103__tep" not in text
diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py
deleted file mode 100644
index 0ad3228f6c8e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_moe_dsa_routing.py
+++ /dev/null
@@ -1,700 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Where a GlmMoeDsaForCausalLM deployment lands, driven through ``modeling_v2_resolve``.
-
-Synthetic configs, as in test_modeling_v2_routing.py: routing reads the config
-shape and layer pattern, the declared quantization, the mapping, the compute
-capability (patched here, so these run on any machine) and the
-speculative-decoding config.
-
-The three targets claim one checkpoint (GLM-5.2 NVFP4 with its bf16 MTP layer), on
-sm103, with one-model MTP-Eagle as a linear chain, at disjoint parallel layouts:
-tp8 (attention and experts tensor-parallel over eight ranks) and tep (attention
-tensor-parallel over four or eight ranks, experts expert-parallel over the same
-ranks), both without attention DP and at five draft tokens, and dep (attention DP
-with the experts expert-parallel over the same ranks: sixteen ranks at three draft
-tokens, eight ranks at three, or eight ranks with the LM head tensor-parallel in
-attention DP at five). Every other value of each criterion declines -- ``auto``
-falls back to the built-in implementation and ``require`` raises naming the fact
-that missed -- including quantization declarations the engine's own parser would
-reject.
-"""
-
-from __future__ import annotations
-
-import re
-
-import pytest
-import torch
-from transformers import PretrainedConfig
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
-    ModelingV2Mode,
-    modeling_v2_resolve,
-    resolved_mode,
-)
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm.llmapi.llm_args import (
-    MTPDecodingConfig,
-    NGramDecodingConfig,
-    SAEnhancerConfig,
-    TorchLlmArgs,
-)
-from tensorrt_llm.mapping import Mapping
-
-pytestmark = pytest.mark.cpu_only
-
-_SM103 = (10, 3)
-_DUMMY_MODEL = "/tmp/dummy_model"
-_TARGET = "models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tp8.modeling.ModelingV2Glm52Nvfp4Sm103Tp8"
-_TEP_TARGET = "models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__tep.modeling.ModelingV2Glm52Nvfp4Sm103Tep"
-_DEP_TARGET = "models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep.modeling.ModelingV2Glm52Nvfp4Sm103Dep"
-
-# The excluded-module patterns of the GLM-5.2 NVFP4 export, abridged: the head, the
-# embedding, the dense layers 0-2, the MTP layer 78, and per MoE layer the attention
-# and the shared expert.
-_NVFP4_IGNORE = [
-    "lm_head",
-    "model.embed_tokens",
-    "model.layers.0*",
-    "model.layers.1.*",
-    "model.layers.2.*",
-    "model.layers.78*",
-    "model.layers.3.self_attn*",
-    "model.layers.3.mlp.shared_experts*",
-]
-
-
-def _qualified(target) -> str:
-    """``models.<family>.<dir>.modeling.<Class>`` for a resolved target class."""
-    return f"{target.__module__.split('.modeling_v2.')[1]}.{target.__name__}"
-
-
-@pytest.fixture(autouse=True)
-def _on_sm103(monkeypatch):
-    """Route as if this were an sm_103 device, wherever the test actually runs."""
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: _SM103)
-
-
-def _nvfp4_quantization(ignore=_NVFP4_IGNORE) -> dict:
-    """The ``quantization_config`` the NVFP4 export writes into its config.json."""
-    return {
-        "quant_method": "modelopt",
-        "quant_algo": "NVFP4",
-        "producer": {"name": "modelopt"},
-        "kv_cache_scheme": {"dynamic": False, "num_bits": 8, "type": "float"},
-        "ignore": list(ignore),
-    }
-
-
-# The shape and layer pattern GLM-5.2-NVFP4's own config.json declares.
-_GLM52_SHAPE = dict(
-    num_hidden_layers=78,
-    hidden_size=6144,
-    n_routed_experts=256,
-    q_lora_rank=2048,
-    num_attention_heads=64,
-    index_topk=2048,
-    first_k_dense_replace=3,
-    index_topk_freq=4,
-    index_skip_topk_offset=3,
-)
-_SHAPE_FIELDS = (*_GLM52_SHAPE, "index_topk_pattern")
-
-
-def _glm52_fields(**overrides) -> dict:
-    fields = dict(
-        architectures=["GlmMoeDsaForCausalLM"],
-        model_type="glm_moe_dsa",
-        num_nextn_predict_layers=1,
-        quantization_config=_nvfp4_quantization(),
-        **_GLM52_SHAPE,
-    )
-    fields.update(overrides)
-    return fields
-
-
-def _glm52_config(**overrides) -> PretrainedConfig:
-    """The shape, layer pattern and quantization GLM-5.2-NVFP4's own config.json declares."""
-    return PretrainedConfig(**_glm52_fields(**overrides))
-
-
-def _shape_words(**overrides) -> str:
-    """The config shape a NoTarget message names for the GLM-5.2 shape with ``overrides``."""
-    values = {**_GLM52_SHAPE, "index_topk_pattern": None, **overrides}
-    return re.escape(f"config shape {tuple(values[name] for name in _SHAPE_FIELDS)}")
-
-
-def _llm_args(mode: str = "auto") -> TorchLlmArgs:
-    return TorchLlmArgs(model=_DUMMY_MODEL, modeling_v2=mode)
-
-
-def _model_config(pretrained_config, spec_config, **mapping_kwargs) -> ModelConfig:
-    return ModelConfig(
-        pretrained_config=pretrained_config,
-        mapping=Mapping(**mapping_kwargs),
-        spec_config=spec_config,
-    )
-
-
-_TP8 = dict(world_size=8, tp_size=8)
-_TEP4 = dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1)
-_TEP8 = dict(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1)
-_DEP16 = dict(world_size=16, tp_size=16, moe_ep_size=16, moe_tp_size=1, enable_attention_dp=True)
-_DEP8 = dict(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1, enable_attention_dp=True)
-_DEP8_LM_HEAD_TP = dict(_DEP8, enable_lm_head_tp_in_adp=True)
-# The deployments the dep target claims: (mapping, draft tokens per request).
-_DEP_CELLS = {
-    "dep16-mtp3": (_DEP16, 3),
-    "dep8-mtp3": (_DEP8, 3),
-    "dep8-lm-head-tp-mtp5": (_DEP8_LM_HEAD_TP, 5),
-}
-
-
-def _mtp(max_draft_len: int = 5, **kwargs) -> MTPDecodingConfig:
-    return MTPDecodingConfig(max_draft_len=max_draft_len, **kwargs)
-
-
-@pytest.mark.parametrize("mode", [None, "off"], ids=["default", "off"])
-def test_off_resolves_nothing(mode):
-    args = TorchLlmArgs(model=_DUMMY_MODEL) if mode is None else _llm_args(mode)
-    config = _model_config(_glm52_config(), _mtp(), **_TP8)
-    assert modeling_v2_resolve(config, args) is None
-
-
-def test_tp8_is_one_expert_group_without_attention_dp():
-    """The default tp8 mapping is the layout the target was written for."""
-    mapping = Mapping(**_TP8)
-    assert (mapping.moe_tp_size, mapping.moe_ep_size, mapping.enable_attention_dp) == (
-        8,
-        1,
-        False,
-    )
-
-
-@pytest.mark.parametrize("mode", ["auto", "require"])
-def test_glm52_nvfp4_tp8_mtp5_matches(mode):
-    config = _model_config(_glm52_config(), _mtp(), **_TP8)
-    target = modeling_v2_resolve(config, _llm_args(mode))
-    assert _qualified(target) == _TARGET
-    assert resolved_mode(target) is ModelingV2Mode(mode)
-
-
-@pytest.mark.parametrize("mode", ["auto", "require"])
-@pytest.mark.parametrize("mapping_kwargs", [_TEP4, _TEP8], ids=["tep4", "tep8"])
-def test_glm52_nvfp4_tep_mtp5_matches(mode, mapping_kwargs):
-    """Four or eight ranks of tensor-parallel attention with the experts expert-parallel over
-    the same ranks land on the tep target."""
-    config = _model_config(_glm52_config(), _mtp(), **mapping_kwargs)
-    target = modeling_v2_resolve(config, _llm_args(mode))
-    assert _qualified(target) == _TEP_TARGET
-    assert resolved_mode(target) is ModelingV2Mode(mode)
-
-
-def test_the_tep_mappings_are_one_expert_per_rank_group_without_attention_dp():
-    for kwargs, ranks in ((_TEP4, 4), (_TEP8, 8)):
-        mapping = Mapping(**kwargs)
-        assert (mapping.world_size, mapping.tp_size, mapping.moe_ep_size, mapping.moe_tp_size) == (
-            ranks,
-            ranks,
-            ranks,
-            1,
-        )
-        assert not mapping.enable_attention_dp
-
-
-def test_the_three_claims_are_disjoint():
-    """No mapping satisfies two parallel criteria: tp8 takes moe_ep 1, tep moe_ep = tp > 1, both
-    without attention DP, and dep takes attention DP."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa import routing
-
-    for kwargs in (_TP8, _TEP4, _TEP8, _DEP16, _DEP8, _DEP8_LM_HEAD_TP):
-        mapping = Mapping(**kwargs)
-        assert routing._tp8(mapping) + routing._tep(mapping) + routing._dep(mapping) == 1, kwargs
-
-
-@pytest.mark.parametrize("mode", ["auto", "require"])
-@pytest.mark.parametrize(
-    "mapping_kwargs, draft_len", list(_DEP_CELLS.values()), ids=list(_DEP_CELLS)
-)
-def test_glm52_nvfp4_dep_matches(mode, mapping_kwargs, draft_len):
-    """Attention DP with the experts expert-parallel over the same ranks lands on the dep target
-    at the draft length of its rank count and LM-head split."""
-    config = _model_config(_glm52_config(), _mtp(draft_len), **mapping_kwargs)
-    target = modeling_v2_resolve(config, _llm_args(mode))
-    assert _qualified(target) == _DEP_TARGET
-    assert resolved_mode(target) is ModelingV2Mode(mode)
-
-
-def test_the_dep_mappings_are_one_expert_per_rank_group_with_attention_dp():
-    for kwargs, ranks, lm_head_tp in ((_DEP16, 16, False), (_DEP8, 8, False)) + (
-        (_DEP8_LM_HEAD_TP, 8, True),
-    ):
-        mapping = Mapping(**kwargs)
-        assert (mapping.world_size, mapping.tp_size, mapping.moe_ep_size, mapping.moe_tp_size) == (
-            ranks,
-            ranks,
-            ranks,
-            1,
-        )
-        assert mapping.enable_attention_dp
-        assert mapping.enable_lm_head_tp_in_adp is lm_head_tp
-
-
-def test_the_dep_cells_are_the_routed_ones():
-    """The target's cell table and ``route``'s claim name the same deployments."""
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa import routing
-    from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_103__dep import (  # noqa: E501
-        modeling,
-    )
-
-    assert set(modeling.CELLS) == {
-        (ranks, lm_head_tp, draft) for (ranks, lm_head_tp), draft in routing._DEP.items()
-    }
-    assert {
-        (Mapping(**kwargs).world_size, Mapping(**kwargs).enable_lm_head_tp_in_adp, draft)
-        for kwargs, draft in _DEP_CELLS.values()
-    } == set(modeling.CELLS)
-
-
-def test_the_dep_target_is_the_builtin_causal_lm():
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
-
-    config = _model_config(_glm52_config(), _mtp(3), **_DEP8)
-    target = modeling_v2_resolve(config, _llm_args("require"))
-    assert issubclass(target, GlmMoeDsaForCausalLM)
-    assert (
-        target.get_preferred_kv_cache_manager_version()
-        == GlmMoeDsaForCausalLM.get_preferred_kv_cache_manager_version()
-    )
-    assert (
-        target.get_preferred_transceiver_runtime()
-        == GlmMoeDsaForCausalLM.get_preferred_transceiver_runtime()
-    )
-
-
-def test_the_tep_target_is_the_builtin_causal_lm():
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
-
-    config = _model_config(_glm52_config(), _mtp(), **_TEP8)
-    target = modeling_v2_resolve(config, _llm_args("require"))
-    assert issubclass(target, GlmMoeDsaForCausalLM)
-    assert (
-        target.get_preferred_kv_cache_manager_version()
-        == GlmMoeDsaForCausalLM.get_preferred_kv_cache_manager_version()
-    )
-
-
-def test_the_engine_config_class_keeps_the_routed_fields():
-    """The loader builds glm_moe_dsa configs through the registered class; route reads the same fields."""
-    from tensorrt_llm._torch.pyexecutor.config_utils import _CONFIG_REGISTRY
-
-    pretrained_config = _CONFIG_REGISTRY["glm_moe_dsa"].from_dict(_glm52_fields())
-    config = _model_config(pretrained_config, _mtp(), **_TP8)
-    assert _qualified(modeling_v2_resolve(config, _llm_args("require"))) == _TARGET
-
-
-def test_the_target_is_the_builtin_causal_lm():
-    """The target keeps the built-in class's engine preferences, read before routing."""
-    from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
-
-    config = _model_config(_glm52_config(), _mtp(), **_TP8)
-    target = modeling_v2_resolve(config, _llm_args("require"))
-    assert issubclass(target, GlmMoeDsaForCausalLM)
-    assert (
-        target.get_preferred_transceiver_runtime()
-        == GlmMoeDsaForCausalLM.get_preferred_transceiver_runtime()
-    )
-    assert (
-        target.get_preferred_kv_cache_manager_version()
-        == GlmMoeDsaForCausalLM.get_preferred_kv_cache_manager_version()
-    )
-
-
-# Each case changes one fact of the claimed deployment: (config overrides, mapping,
-# speculative config, compute capability, what the NoTarget message must name).
-_FOREIGN = {
-    "tp4": ({}, dict(world_size=4, tp_size=4), _mtp(), _SM103, r"ws=4 tp=4"),
-    "tp16": ({}, dict(world_size=16, tp_size=16), _mtp(), _SM103, r"ws=16 tp=16"),
-    "attention-dp": (
-        {},
-        dict(_TP8, enable_attention_dp=True),
-        _mtp(),
-        _SM103,
-        r"attention_dp=True",
-    ),
-    "tep2": (
-        {},
-        dict(world_size=2, tp_size=2, moe_ep_size=2, moe_tp_size=1),
-        _mtp(),
-        _SM103,
-        r"ws=2 tp=2 .* moe_tp=1 moe_ep=2",
-    ),
-    "tep16": (
-        {},
-        dict(world_size=16, tp_size=16, moe_ep_size=16, moe_tp_size=1),
-        _mtp(),
-        _SM103,
-        r"ws=16 tp=16 .* moe_tp=1 moe_ep=16",
-    ),
-    "tep8-attention-dp": (
-        {},
-        dict(_TEP8, enable_attention_dp=True),
-        _mtp(),
-        _SM103,
-        r"moe_ep=8 attention_dp=True",
-    ),
-    "tep4-attention-dp": (
-        {},
-        dict(_TEP4, enable_attention_dp=True),
-        _mtp(),
-        _SM103,
-        r"moe_ep=4 attention_dp=True",
-    ),
-    "tp8-ep4-tp2": (
-        {},
-        dict(_TP8, moe_ep_size=4, moe_tp_size=2),
-        _mtp(),
-        _SM103,
-        r"moe_tp=2 moe_ep=4",
-    ),
-    "tep4-pp2": (
-        {},
-        dict(world_size=8, tp_size=4, pp_size=2, moe_ep_size=4, moe_tp_size=1),
-        _mtp(),
-        _SM103,
-        r"ws=8 tp=4 pp=2",
-    ),
-    "tep8-sm100": ({}, _TEP8, _mtp(), (10, 0), r"compute capability \(10, 0\)"),
-    "tep4-mtp3": ({}, _TEP4, _mtp(3), _SM103, r"max_draft_len=3 "),
-    "tep8-bf16-checkpoint": (
-        dict(quantization_config=None),
-        _TEP8,
-        _mtp(),
-        _SM103,
-        r"quantization none",
-    ),
-    "sm100": ({}, _TP8, _mtp(), (10, 0), r"compute capability \(10, 0\)"),
-    "sm120": ({}, _TP8, _mtp(), (12, 0), r"compute capability \(12, 0\)"),
-    "sm90": ({}, _TP8, _mtp(), (9, 0), r"compute capability \(9, 0\)"),
-    "bf16-checkpoint": (
-        dict(quantization_config=None),
-        _TP8,
-        _mtp(),
-        _SM103,
-        r"quantization none",
-    ),
-    "fp8-checkpoint": (
-        dict(quantization_config={"quant_method": "fp8", "weight_block_size": [128, 128]}),
-        _TP8,
-        _mtp(),
-        _SM103,
-        r"quantization quant_method 'fp8' \(not a modelopt export\)",
-    ),
-    "fp8-modelopt-checkpoint": (
-        dict(quantization_config={**_nvfp4_quantization(), "quant_algo": "FP8"}),
-        _TP8,
-        _mtp(),
-        _SM103,
-        r"quantization FP8 with the MTP layer unquantized",
-    ),
-    "quantized-mtp-layer": (
-        dict(quantization_config=_nvfp4_quantization(["lm_head", "model.embed_tokens"])),
-        _TP8,
-        _mtp(),
-        _SM103,
-        r"quantization NVFP4 with the MTP layer quantized",
-    ),
-    **{
-        f"shape-{name}": (
-            {name: value},
-            _TP8,
-            _mtp(),
-            _SM103,
-            _shape_words(**{name: value}),
-        )
-        for name, value in (
-            ("num_hidden_layers", 92),
-            ("hidden_size", 7168),
-            ("n_routed_experts", 160),
-            ("q_lora_rank", 1536),
-            ("num_attention_heads", 128),
-            ("index_topk", 1024),
-            ("first_k_dense_replace", 1),
-            ("index_topk_freq", 1),
-            ("index_skip_topk_offset", 2),
-            ("index_topk_pattern", "FFFSSSF"),
-        )
-    },
-    "no-speculation": ({}, _TP8, None, _SM103, r"speculative decoding off"),
-    "mtp1": ({}, _TP8, _mtp(1), _SM103, r"max_draft_len=1 "),
-    "mtp2": ({}, _TP8, _mtp(2), _SM103, r"max_draft_len=2 "),
-    "mtp3": ({}, _TP8, _mtp(3), _SM103, r"max_draft_len=3 "),
-    "mtp4": ({}, _TP8, _mtp(4), _SM103, r"max_draft_len=4 "),
-    "mtp6": ({}, _TP8, _mtp(6), _SM103, r"max_draft_len=6 "),
-    "vanilla-mtp": ({}, _TP8, _mtp(use_mtp_vanilla=True), _SM103, r"speculative decoding MTP "),
-    "dynamic-tree": (
-        {},
-        _TP8,
-        _mtp(use_dynamic_tree=True, dynamic_tree_max_topK=1),
-        _SM103,
-        r"dynamic tree",
-    ),
-    "suffix-automaton-enhancer": (
-        {},
-        _TP8,
-        _mtp(sa_config=SAEnhancerConfig()),
-        _SM103,
-        r"with the suffix-automaton enhancer",
-    ),
-    "relaxed-acceptance": (
-        {},
-        _TP8,
-        _mtp(use_relaxed_acceptance_for_thinking=True),
-        _SM103,
-        r"with relaxed acceptance",
-    ),
-    "ngram": (
-        {},
-        _TP8,
-        NGramDecodingConfig(max_draft_len=5, max_matching_ngram_size=2),
-        _SM103,
-        r"speculative decoding NGRAM",
-    ),
-    # The dep claim: one fact of a claimed deployment changed.
-    "dep4-mtp3": (
-        {},
-        dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True),
-        _mtp(3),
-        _SM103,
-        r"ws=4 tp=4 .* attention_dp=True",
-    ),
-    "dep32-mtp3": (
-        {},
-        dict(world_size=32, tp_size=32, moe_ep_size=32, moe_tp_size=1, enable_attention_dp=True),
-        _mtp(3),
-        _SM103,
-        r"ws=32 tp=32",
-    ),
-    "dep16-mtp5": ({}, _DEP16, _mtp(), _SM103, r"ws=16 .* max_draft_len=5 "),
-    "dep16-lm-head-tp-mtp3": (
-        {},
-        dict(_DEP16, enable_lm_head_tp_in_adp=True),
-        _mtp(3),
-        _SM103,
-        r"ws=16 .* lm_head_tp_in_adp=True, .* max_draft_len=3 ",
-    ),
-    "dep16-lm-head-tp-mtp5": (
-        {},
-        dict(_DEP16, enable_lm_head_tp_in_adp=True),
-        _mtp(),
-        _SM103,
-        r"ws=16 .* lm_head_tp_in_adp=True, .* max_draft_len=5 ",
-    ),
-    "dep8-lm-head-tp-mtp3": (
-        {},
-        _DEP8_LM_HEAD_TP,
-        _mtp(3),
-        _SM103,
-        r"lm_head_tp_in_adp=True, .* max_draft_len=3 ",
-    ),
-    **{
-        f"dep8-mtp{draft}": ({}, _DEP8, _mtp(draft), _SM103, rf"max_draft_len={draft} ")
-        for draft in (1, 2, 4, 5)
-    },
-    **{
-        f"dep8-lm-head-tp-mtp{draft}": (
-            {},
-            _DEP8_LM_HEAD_TP,
-            _mtp(draft),
-            _SM103,
-            rf"max_draft_len={draft} ",
-        )
-        for draft in (1, 2, 4, 6)
-    },
-    "dep8-moe-tp2": (
-        {},
-        dict(_DEP8, moe_ep_size=4, moe_tp_size=2),
-        _mtp(3),
-        _SM103,
-        r"moe_tp=2 moe_ep=4 attention_dp=True",
-    ),
-    "dep8-pp2": (
-        {},
-        dict(
-            world_size=16,
-            tp_size=8,
-            pp_size=2,
-            moe_ep_size=8,
-            moe_tp_size=1,
-            enable_attention_dp=True,
-        ),
-        _mtp(3),
-        _SM103,
-        r"ws=16 tp=8 pp=2",
-    ),
-    "dep8-no-speculation": ({}, _DEP8, None, _SM103, r"speculative decoding off"),
-    "dep8-vanilla-mtp": (
-        {},
-        _DEP8,
-        _mtp(3, use_mtp_vanilla=True),
-        _SM103,
-        r"speculative decoding MTP ",
-    ),
-    "dep8-dynamic-tree": (
-        {},
-        _DEP8,
-        _mtp(3, use_dynamic_tree=True, dynamic_tree_max_topK=1),
-        _SM103,
-        r"dynamic tree",
-    ),
-    "dep8-suffix-automaton-enhancer": (
-        {},
-        _DEP8,
-        _mtp(3, sa_config=SAEnhancerConfig()),
-        _SM103,
-        r"with the suffix-automaton enhancer",
-    ),
-    "dep8-relaxed-acceptance": (
-        {},
-        _DEP8,
-        _mtp(3, use_relaxed_acceptance_for_thinking=True),
-        _SM103,
-        r"with relaxed acceptance",
-    ),
-    "dep8-sm100": ({}, _DEP8, _mtp(3), (10, 0), r"compute capability \(10, 0\)"),
-    "dep16-bf16-checkpoint": (
-        dict(quantization_config=None),
-        _DEP16,
-        _mtp(3),
-        _SM103,
-        r"quantization none",
-    ),
-    "dep8-lm-head-tp-quantized-mtp-layer": (
-        dict(quantization_config=_nvfp4_quantization(["lm_head", "model.embed_tokens"])),
-        _DEP8_LM_HEAD_TP,
-        _mtp(),
-        _SM103,
-        r"quantization NVFP4 with the MTP layer quantized",
-    ),
-    "dep8-shape-num_attention_heads": (
-        {"num_attention_heads": 128},
-        _DEP8,
-        _mtp(3),
-        _SM103,
-        _shape_words(num_attention_heads=128),
-    ),
-}
-
-
-@pytest.mark.parametrize(
-    "config_kwargs, mapping_kwargs, spec_config, capability, missed",
-    list(_FOREIGN.values()),
-    ids=list(_FOREIGN),
-)
-def test_a_foreign_deployment_does_not_match(
-    monkeypatch, config_kwargs, mapping_kwargs, spec_config, capability, missed
-):
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
-    config = _model_config(_glm52_config(**config_kwargs), spec_config, **mapping_kwargs)
-
-    assert modeling_v2_resolve(config, _llm_args("auto")) is None
-    with pytest.raises(ValueError, match="GlmMoeDsaForCausalLM") as excinfo:
-        modeling_v2_resolve(config, _llm_args("require"))
-    message = str(excinfo.value)
-    assert "the glm_moe_dsa targets are" in message
-    assert "tp8 (8 ranks" in message and "tep (4 or 8 ranks" in message
-    assert "dep (attention DP" in message
-    assert re.search(missed, message), message
-
-
-def test_a_config_without_the_sparse_attention_fields_does_not_match():
-    """A GlmMoeDsa config that lacks a shape field declines instead of failing in route."""
-    pretrained_config = _glm52_config()
-    del pretrained_config.index_topk
-    config = _model_config(pretrained_config, _mtp(), **_TP8)
-
-    assert modeling_v2_resolve(config, _llm_args("auto")) is None
-    with pytest.raises(ValueError, match=_shape_words(index_topk=None)):
-        modeling_v2_resolve(config, _llm_args("require"))
-
-
-class _NoToDict:
-    """A quantization declaration that is neither a dict nor convertible to one."""
-
-    quant_method = "nvfp4"
-
-
-# Quantization declarations route must decline without parsing through the engine's
-# builders, which raise other errors than NoTarget on them: (declaration, what the
-# message names when the rest of the deployment matches).
-_UNREADABLE_QUANTIZATION = {
-    "modelopt-mixed-precision-flat": (
-        {
-            "quant_method": "modelopt",
-            "quant_algo": "MIXED_PRECISION",
-            "config_groups": {"group_0": {}},
-            "ignore": ["lm_head"],
-        },
-        r"quantization MIXED_PRECISION with the MTP layer quantized",
-    ),
-    "modelopt-mixed-precision-legacy": (
-        {"producer": {"name": "modelopt"}, "quantization": {"quant_algo": "MIXED_PRECISION"}},
-        r"quantization MIXED_PRECISION with the MTP layer quantized",
-    ),
-    "modelopt-unknown-algorithm": (
-        {"quant_method": "modelopt", "quant_algo": "NOT_AN_ALGORITHM"},
-        r"quantization a modelopt config without a readable global algorithm",
-    ),
-    "modelopt-no-algorithm": (
-        {"quant_method": "modelopt", "ignore": ["model.layers.78*"]},
-        r"quantization a modelopt config without a readable global algorithm",
-    ),
-    "modelopt-excluded-modules-not-a-list": (
-        {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": 5},
-        r"quantization modelopt NVFP4 with unreadable excluded modules",
-    ),
-    "fp8-64x64-blocks": (
-        {"quant_method": "fp8", "weight_block_size": [64, 64]},
-        r"quantization quant_method 'fp8'",
-    ),
-    "nvfp4-group-32": (
-        {"quant_method": "nvfp4", "group_size": 32},
-        r"quantization quant_method 'nvfp4'",
-    ),
-    "mxfp8-1x64-blocks": (
-        {"quant_method": "mxfp8", "weight_block_size": [1, 64]},
-        r"quantization quant_method 'mxfp8'",
-    ),
-    "compressed-tensors-without-strategy": (
-        {
-            "quant_method": "compressed-tensors",
-            "config_groups": {"group_0": {"weights": {"num_bits": 4, "type": "int"}}},
-        },
-        r"quantization quant_method 'compressed-tensors'",
-    ),
-    "not-a-dict": (_NoToDict(), r"quantization declared as a _NoToDict, not a dict"),
-}
-
-
-@pytest.mark.parametrize("capability", [_SM103, (9, 0)], ids=["sm103", "sm90"])
-@pytest.mark.parametrize(
-    "quantization, named",
-    list(_UNREADABLE_QUANTIZATION.values()),
-    ids=list(_UNREADABLE_QUANTIZATION),
-)
-def test_a_quantization_route_cannot_read_declines(monkeypatch, quantization, named, capability):
-    """Declined, never an exception other than NoTarget, so ``auto`` still falls back.
-
-    The quantization is read only for the deployment the target serves; on any
-    other deployment the message says it was not read.
-    """
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
-    config = _model_config(_glm52_config(quantization_config=quantization), _mtp(), **_TP8)
-
-    assert modeling_v2_resolve(config, _llm_args("auto")) is None
-    with pytest.raises(ValueError, match="GlmMoeDsaForCausalLM") as excinfo:
-        modeling_v2_resolve(config, _llm_args("require"))
-    expected = named if capability == _SM103 else r"quantization not read"
-    assert re.search(expected, str(excinfo.value)), str(excinfo.value)

From b209b9f9cee1299a88213dee3cc19c8fb854955c Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Wed, 21 Oct 2026 09:52:22 -0700
Subject: [PATCH 09/19] [None][refactor] modeling_v2 GLM-5.2: one default
 target for the tp8, tep and dep layouts

Following review, the three GLM-5.2 targets become one,
glm_5_2_nvfp4__sm_103__default, as the gpt-oss and DeepSeek targets did.
route() admits the same layouts and speculation as before and returns
that target for all of them; the target picks the fast paths of the
layout when it is built (target_paths): every path at tp8, the tep
tables per tensor-parallel size, and the dep cells, whose KV-cache dtype
and MoE backend it still checks at construction. Under attention DP it
adopts only the attention modules and the speculative worker, as the dep
target did.

The forward, the fast-path tables and the adoption are unchanged per
layout. The target's README states that only tp8 is optimized and that
tep and dep should be revisited. The accuracy test class follows the
target name.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_experimental/modeling_v2/README.md       |  32 +-
 .../glm_5_2_nvfp4__sm_103__default/README.md  |  24 ++
 .../__init__.py                               |   2 +-
 .../kernels}/__init__.py                      |   2 +-
 .../cross_layer_expert_prefetch/__init__.py   |   0
 .../cross_layer_expert_prefetch/kernels.py    |   0
 .../spec_experts.py                           |   0
 .../kernels/draft_routed_moe/__init__.py      |   0
 .../kernels/draft_routed_moe/routed_moe.py    |   0
 .../kernels/hop_routing/__init__.py           |   0
 .../kernels/hop_routing/hop.py                |   0
 .../kernels/hop_routing/routed_moe.py         |   0
 .../kernels/routed_moe/__init__.py            |   0
 .../kernels/routed_moe/kernel.py              |   0
 .../kernels/routed_moe/kernel_xc.py           |   0
 .../kernels/routed_moe/ptx.py                 |   0
 .../kernels/routed_moe/routed_moe.py          |   0
 .../shared_expert_prefetch/__init__.py        |   0
 .../shared_expert_prefetch/kernel_xc.py       |   0
 .../shared_expert_prefetch/routed_moe.py      |   0
 .../modeling.py                               | 276 +++++++++-----
 .../weights.py                                |   4 +-
 .../glm_5_2_nvfp4__sm_103__dep/README.md      |  15 -
 .../glm_5_2_nvfp4__sm_103__dep/weights.py     |  21 --
 .../glm_5_2_nvfp4__sm_103__tep/README.md      |  14 -
 .../glm_5_2_nvfp4__sm_103__tep/modeling.py    | 357 ------------------
 .../glm_5_2_nvfp4__sm_103__tep/weights.py     |  21 --
 .../glm_5_2_nvfp4__sm_103__tp8/__init__.py    |   3 -
 .../kernels/__init__.py                       |   3 -
 .../glm_5_2_nvfp4__sm_103__tp8/modeling.py    | 318 ----------------
 .../modeling_v2/models/glm_moe_dsa/routing.py | 110 +++---
 .../accuracy/modeling_v2/test_glm_moe_dsa.py  |  30 +-
 .../test-db/l0_gb300_multi_gpus.yml           |   4 +-
 .../l0_gb300_multi_nodes_node2_gpu8.yml       |   6 +-
 34 files changed, 286 insertions(+), 956 deletions(-)
 create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/README.md
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tep => glm_5_2_nvfp4__sm_103__default}/__init__.py (77%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__dep => glm_5_2_nvfp4__sm_103__default/kernels}/__init__.py (59%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/cross_layer_expert_prefetch/__init__.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/cross_layer_expert_prefetch/kernels.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/cross_layer_expert_prefetch/spec_experts.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/draft_routed_moe/__init__.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/draft_routed_moe/routed_moe.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/hop_routing/__init__.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/hop_routing/hop.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/hop_routing/routed_moe.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/routed_moe/__init__.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/routed_moe/kernel.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/routed_moe/kernel_xc.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/routed_moe/ptx.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/routed_moe/routed_moe.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/shared_expert_prefetch/__init__.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/shared_expert_prefetch/kernel_xc.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/kernels/shared_expert_prefetch/routed_moe.py (100%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__dep => glm_5_2_nvfp4__sm_103__default}/modeling.py (58%)
 rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_103__tp8 => glm_5_2_nvfp4__sm_103__default}/weights.py (84%)
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py
 delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
index 212889f5024f..a765e47c5ff0 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md
@@ -115,27 +115,17 @@ models/<family>/
 catalog/<category>/ entries more than one target calls
 ```
 
-The targets, as `models/<family>/<target>`:
-
-* `deepseek_v3/r1_0528_nvfp4__sm_10x__default`
-* `glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8`: GLM-5.2 NVFP4 low-latency decode with attention
-  and experts tensor-parallel over eight ranks. The fast paths shared by the GLM-5.2 targets are
-  under `glm_moe_dsa/shared/kernels/`, the routed-MoE ones only this target calls under its own
-  `kernels/`, with their C++ kernels and ops in `cpp/tensorrt_llm/kernels/glmMoeDsa` and
-  `cpp/tensorrt_llm/thop/glmMoeDsa`; `shared/dispatch.py` selects them and lists the knobs.
-  The tp8 and tep targets run the layer composition of `shared/layers.py`, each with its own plan.
-  Each GLM-5.2 target is gated end to end by
-  `tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py` (GSM8K under `require`).
-* `glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep`: the same checkpoint with attention tensor-parallel
-  over four or eight ranks and the experts expert-parallel over the same ranks. It runs the
-  shared fast paths that take its tensor-parallel size, in front of the built-in
-  expert-parallel MoE.
-* `glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep`: the same checkpoint with attention data parallelism
-  and the experts expert-parallel over the same ranks. Its built-in layers run the shared
-  attention, draft-loop and step-loop fast paths that are rank-local under attention DP; it
-  checks the KV-cache dtype and the MoE backend, which `route` does not read, when it is built.
-* `gpt_oss/gpt_oss_120b__sm_10x__default`
-* `qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1`
+GLM-5.2's target, `glm_moe_dsa/glm_5_2_nvfp4__sm_103__default`, is a default target (see
+"Default and specialized targets") for GLM-5.2 NVFP4 low-latency decode. It admits tp8, tep
+at four or eight ranks, and dep at the cells its routing module lists, and picks the fast
+paths of the layout when it is built; the fast paths were designed and tuned for tp8, and
+tep and dep run the subset that applies to them. The fast paths are under
+`glm_moe_dsa/shared/kernels/`, the routed-MoE ones under the target's own `kernels/`, with
+their C++ kernels and ops in `cpp/tensorrt_llm/kernels/glmMoeDsa` and
+`cpp/tensorrt_llm/thop/glmMoeDsa`; `shared/dispatch.py` selects them and lists the knobs.
+With tensor-parallel attention the layers run the composition of `shared/layers.py`. The
+target is gated end to end by `tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py`
+(GSM8K under `require` at tp8, tep4 and dep8).
 
 An entry, shared or owned, is one `OpWrapper` subclass: `__call__` is the op
 call a target makes, `reference` is what it is supposed to compute in plain
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/README.md
new file mode 100644
index 000000000000..1d1137f191e4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/README.md
@@ -0,0 +1,24 @@
+# `glm_5_2_nvfp4__sm_103__default`
+
+GLM-5.2 NVFP4 on sm_103, at the layouts `../routing.py` admits. The target picks the fast
+paths of the layout when it is built (`target_paths` in `modeling.py`).
+
+**Only TP8 is optimized.** The family's fast paths were designed and tuned for low-latency
+decode with the attention heads and the experts tensor-parallel over eight ranks, and TP8 runs
+all of them, including the routed-MoE kernels under `kernels/`.
+
+TEP and DEP are not separately optimized. They run the fast paths in `../shared/` whose
+contracts also cover them:
+
+* TEP (attention tensor-parallel over four or eight ranks, experts expert-parallel over the same
+  ranks) runs the shared paths that take its tensor-parallel size, in front of the built-in
+  expert-parallel MoE.
+* DEP (attention data parallelism, experts expert-parallel over the same ranks) runs only the
+  rank-local ones: the RoPE and KV append, the full-indexer layers' query front, the draft loop,
+  the decode step's prologue and the host page retention policy. Its decoder layers, MoE modules
+  and draft layer stay the built-in classes.
+
+No fast path was designed or tuned for TEP or DEP batch shapes. They are admitted because those
+shared paths also improve decode there. TEP and DEP should be revisited in the future with fast
+paths of their own, in particular for the expert-parallel MoE and, under DEP, the attention-DP
+communication.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/__init__.py
similarity index 77%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/__init__.py
index bfe54ffe7404..1130d5a82f64 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/__init__.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/__init__.py
@@ -1,3 +1,3 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""GLM-5.2-NVFP4 / sm_103 / tep."""
+"""GLM-5.2-NVFP4 / sm_103 / default."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/__init__.py
similarity index 59%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/__init__.py
index dd69779ecd5e..f5a3eac8c769 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/__init__.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/__init__.py
@@ -1,3 +1,3 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""GLM-5.2-NVFP4 / sm_103 / dep."""
+"""Kernels of the target's routed-MoE fast paths (tp8 layout), one package per fast path."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/cross_layer_expert_prefetch/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/cross_layer_expert_prefetch/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/cross_layer_expert_prefetch/kernels.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/kernels.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/cross_layer_expert_prefetch/kernels.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/spec_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/cross_layer_expert_prefetch/spec_experts.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/cross_layer_expert_prefetch/spec_experts.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/cross_layer_expert_prefetch/spec_experts.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/draft_routed_moe/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/draft_routed_moe/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/draft_routed_moe/routed_moe.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/draft_routed_moe/routed_moe.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/draft_routed_moe/routed_moe.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/hop_routing/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/hop_routing/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/hop_routing/hop.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/hop.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/hop_routing/hop.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/hop_routing/routed_moe.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/hop_routing/routed_moe.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/hop_routing/routed_moe.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/kernel.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/kernel.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel_xc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/kernel_xc.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/kernel_xc.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/kernel_xc.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/ptx.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/ptx.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/ptx.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/routed_moe.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/routed_moe/routed_moe.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/routed_moe/routed_moe.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/shared_expert_prefetch/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/shared_expert_prefetch/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/kernel_xc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/shared_expert_prefetch/kernel_xc.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/kernel_xc.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/shared_expert_prefetch/kernel_xc.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/routed_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/shared_expert_prefetch/routed_moe.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/shared_expert_prefetch/routed_moe.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/kernels/shared_expert_prefetch/routed_moe.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/modeling.py
similarity index 58%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/modeling.py
index 5b2625da55cb..e97274ee00da 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/modeling.py
@@ -1,55 +1,53 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / dep.
+"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / default.
 
 GLM-5.2, modelopt NVFP4 export: 78 layers, hidden 6144, 64 query heads, MLA with
 DSA sparse attention (index top-k 2048), layers 0-2 dense and 3-77 MoE (256
 routed experts at top-8 plus a shared expert). Layer 78 is the checkpoint's bf16
-MTP module, driven by the one-model MTP-Eagle loop. `dep` is attention data
-parallelism over eight or sixteen ranks (each rank runs the whole attention on its
-own requests) and the experts expert-parallel over the same ranks.
-
-``route`` claims three deployments by the facts it reads: the rank count, whether
-the LM head is tensor-parallel in attention DP, and the draft tokens per request.
-Two more facts fix what the fast paths can serve there, the KV-cache dtype and the
-MoE backend, and only the model config carries them, so this target checks them
-first, at construction (``CELLS``, ``target_paths``). Under ``require`` a
-deployment outside the cells raises; under ``auto`` it serves no fast path and
-runs the built-in forward.
-
-Under attention DP the fast paths that take this target's shapes are rank-local:
-the attention's rope, KV append and top-k bridge with an FP8 KV cache
-(``rope_kv_append``), the rope of the full-indexer layers' query front
-(``indexer_layer_front``), the draft loop (``draft_loop``), the decode step's
-prologue (``decode_prologue``) and the host page retention policy
-(``host_page_retention``). The plan of a cell serves those and the bases they
-extend; the bases decline every call under attention DP. The other fast paths of
-the family take tensor-parallel attention, so this target serves none of them: its
-decoder layers, MoE modules and draft layer stay the built-in classes, which run
-the built-in model's attention-DP communication on every step.
-
-The shell is the built-in ``GlmMoeDsaForCausalLM``: the same module tree, the
-same weights and the same engine preferences (KV cache manager, KV transceiver,
-sparse-attention defaults), so a built-in context server and this target's
-generation server exchange the same KV and indexer-K layout. Construction runs,
-in order:
-
-1. the cell check (``target_paths``), from the model config alone, so every rank
-   decides alike before any collective;
+MTP module, driven by the one-model MTP-Eagle loop.
+
+One target serves every layout ``route`` admits; what differs between them is
+which fast paths run (``target_paths``):
+
+* tp8, the attention heads and the experts tensor-parallel over eight ranks:
+  every fast path of the family. The fast paths were designed and tuned for this
+  layout's low-latency decode.
+* tep, the attention heads tensor-parallel over four or eight ranks and the
+  experts expert-parallel over the same ranks: at eight ranks every fast path
+  but the routed-MoE ones, which take only tensor-parallel experts; at four
+  ranks the paths whose kernels take four ranks and the bases they extend. The
+  experts run the built-in expert-parallel MoE.
+* dep, attention data parallelism and the experts expert-parallel over the
+  same ranks: the rank-local paths of the cell the deployment falls in
+  (``DEP_CELLS``). Two facts of a cell, the KV-cache dtype and the MoE backend,
+  are not read by ``route``, so this target checks them at construction.
+
+Construction runs, in order:
+
+1. the fast paths of the layout (``target_paths``), from the model config
+   alone, so every rank decides alike before any collective;
 2. the rank-uniform check of the knobs (``dispatch.check_rank_uniform``), before
    every collective the construction below makes;
-3. the dispatch plan (``dispatch.resolve``) of the paths the cell serves, before
-   the built-in model is built, so a knob value no rank can parse fails early;
+3. the dispatch plan (``dispatch.resolve``) of those paths, before the built-in
+   model is built, so a knob value no rank can parse fails early;
 4. the plan is installed (``dispatch.install``): from here on the fast paths read
    it;
 5. the built-in model, then the core: its text model adopted as
-   ``GlmMoeDsaDepModelingV2``, whose forward hands each step to
-   ``PrefillTarget`` or ``DecodeTarget``;
-6. ``adopt_model``, which re-classes the attention modules of every layer
-   (``shared.attention``) when the plan runs an attention path, and the
-   speculative worker (``shared.spec_worker``) when it runs the draft loop;
-7. the rank-uniform check of the attached fast-path state
-   (``dispatch.check_rank_uniform_adopted``, ``shared.layers.adopted_state``).
+   ``GlmMoeDsaModelingV2``, whose forward hands each step to ``PrefillTarget``
+   or ``DecodeTarget``;
+6. the adoption of the layout. With tensor-parallel attention, ``adopt_model``
+   (``shared.layers``) re-classes the decoder layers, the draft layer, their
+   attention modules (``shared.attention``) and their MoE modules as the
+   family's layer classes and attaches the fast-path objects the plan enables
+   to them, then re-classes the speculative worker (``shared.spec_worker``).
+   Under attention DP, ``_adopt_attention_dp`` re-classes only the attention
+   modules and the speculative worker, as far as the plan runs their paths; the
+   layers, their MoE modules and the draft layer stay the built-in classes,
+   which run the built-in attention-DP communication. The attention backends
+   build their FMHA libraries after construction, in ``create_weights``, each
+   with the libraries adoption declared on it first;
+7. the rank-uniform check of that state (``dispatch.check_rank_uniform_adopted``).
 
 The step-loop hooks the target declares to the engine (``runtime_hooks``,
 ``shared.runtime.make_runtime_hooks``) are built later, on the engine's first
@@ -61,11 +59,10 @@
 Every step runs under the phase of the target that takes it
 (``dispatch.step``): the shell's forward enters it around the built-in forward,
 so the model, the speculative worker and its draft loop all run in it. A
-prefill step runs the built-in decoder forward and the built-in draft loop; a
-decode step runs ``DecodeTarget``'s loop over the built-in layers, the attention
-modules' decode composition and the worker's draft loop. Each rank picks its
-phase from its own requests; the layers' communication is the built-in one in
-both phases.
+prefill step runs the built-in decoder forward, the built-in layers and the
+built-in draft loop; a decode step runs ``DecodeTarget``'s loop, the layers'
+decode composition and the worker's draft loop. Under attention DP each rank
+picks its phase from its own requests.
 """
 
 from __future__ import annotations
@@ -90,7 +87,7 @@
 
 from ..shared import dispatch, runtime
 from ..shared.attention import GlmMLA
-from ..shared.layers import adopted_state
+from ..shared.layers import adopt_model, adopted_state
 from ..shared.spec_worker import GlmMtpWorker
 from . import weights as _weights
 
@@ -98,10 +95,73 @@
     from ..shared.kernels.decode_prologue import DecodePrologueHook
 
 
-# The paths that engage under attention DP with an FP8 KV cache: the attention's rope, KV append
-# and top-k bridge, the full-indexer layers' query-front rope, the draft loop, the decode
-# prologue and the host page retention policy.
-_FP8_KV_ENGAGED = frozenset(
+# The routed-MoE paths, which run from this target's own kernel packages. They take
+# tensor-parallel experts (every expert on every rank).
+_ROUTED_MOE_PACKAGES = {
+    "routed_moe": "routed_moe",
+    "routed_moe_into_hop": "routed_moe",
+    "hop_routing": "hop_routing",
+    "shared_expert_prefetch": "shared_expert_prefetch",
+    "draft_routed_moe": "draft_routed_moe",
+    "cross_layer_expert_prefetch": "cross_layer_expert_prefetch",
+}
+_KERNEL_MODULES = {
+    path: f"{__package__}.kernels.{package}" for path, package in _ROUTED_MOE_PACKAGES.items()
+}
+_NO_LAYERS: frozenset = frozenset()
+
+#: tp8: every fast path of the family, with the family's tables.
+TP8_PATHS = dispatch.TargetPaths(kernel_modules=_KERNEL_MODULES)
+
+# tep at four ranks: the attention operands, the attention output and the post-MoE hop are
+# outside the kernels' shapes (sixteen local heads, a four-rank output projection, a
+# single-node group); these paths take four ranks.
+_TEP4_PATHS = frozenset(
+    {
+        "indexer_layer_front",
+        "rope_kv_append",
+        "indexer_topk",
+        "cross_rank_topk",
+        "cross_rank_topk_cluster_merge",
+        "cross_rank_topk_global_emit",
+        "router_gemm",
+        "dense_mlp",
+        "draft_front",
+        "draft_tail",
+        "draft_loop",
+        "decode_prologue",
+        "deferred_input_prep",
+        "sampler_store",
+        "execution_stream_loop",
+        "host_page_retention",
+    }
+)
+# The bases these paths extend at four ranks. Their kernels decline every layer there: the query
+# fronts on their layer cells, the post-attention hop on its native all-reduce, which a
+# single-node group runs without multi-node NVLink.
+_TEP4_BASES = frozenset({"absorbed_query", "draft_absorbed_query", "post_attention_hop"})
+
+#: tep, per tensor-parallel size.
+TEP_PATHS = {
+    4: dispatch.TargetPaths(
+        served=_TEP4_PATHS | _TEP4_BASES,
+        layer_cells={
+            **dispatch.LAYER_CELLS,
+            "absorbed_query": {"layer": _NO_LAYERS},
+            "draft_absorbed_query": {"layer": _NO_LAYERS},
+            "indexer_layer_front": {"target_layer": _NO_LAYERS, "draft_layer": _NO_LAYERS},
+        },
+        uncovered={"post_attention_hop": ("does not use MNNVL",)},
+    ),
+    8: dispatch.TargetPaths(
+        served=frozenset(dispatch.FAST_PATH_NAMES) - frozenset(_ROUTED_MOE_PACKAGES)
+    ),
+}
+
+# dep with an FP8 KV cache: the attention's rope, KV append and top-k bridge, the full-indexer
+# layers' query-front rope, the draft loop, the decode prologue and the host page retention
+# policy engage.
+_DEP_FP8_KV_ENGAGED = frozenset(
     {
         "rope_kv_append",
         "indexer_layer_front",
@@ -113,7 +173,7 @@
 # The bases these paths extend. Under attention DP their kernels decline every call: the query
 # fronts on every layer, the top-k selection on every step, and the post-attention hop is never
 # built (the built-in layers run no post-attention all-reduce under attention DP).
-_FP8_KV_BASES = frozenset(
+_DEP_FP8_KV_BASES = frozenset(
     {
         "absorbed_query",
         "draft_absorbed_query",
@@ -122,16 +182,15 @@
         "post_attention_hop",
     }
 )
-# With an NVFP4 KV cache the attention paths are outside the shapes their kernels take; the
+# dep with an NVFP4 KV cache: the attention paths are outside the shapes their kernels take; the
 # draft loop, the decode prologue and the host page retention policy engage.
-_NVFP4_KV_ENGAGED = frozenset({"draft_loop", "decode_prologue", "host_page_retention"})
-_NVFP4_KV_BASES = frozenset({"post_attention_hop"})
-_NO_LAYERS: frozenset = frozenset()
+_DEP_NVFP4_KV_ENGAGED = frozenset({"draft_loop", "decode_prologue", "host_page_retention"})
+_DEP_NVFP4_KV_BASES = frozenset({"post_attention_hop"})
 
-#: What this target serves with an FP8 KV cache. The query fronts decline attention DP on every
-#: layer, which is outside their shapes here.
-FP8_KV_PATHS = dispatch.TargetPaths(
-    served=_FP8_KV_ENGAGED | _FP8_KV_BASES,
+#: dep with an FP8 KV cache. The query fronts decline attention DP on every layer, which is
+#: outside their shapes here.
+DEP_FP8_KV_PATHS = dispatch.TargetPaths(
+    served=_DEP_FP8_KV_ENGAGED | _DEP_FP8_KV_BASES,
     layer_cells={
         **dispatch.LAYER_CELLS,
         "absorbed_query": {"layer": _NO_LAYERS},
@@ -139,38 +198,36 @@
         "indexer_layer_front": {"target_layer": _NO_LAYERS, "draft_layer": _NO_LAYERS},
     },
 )
-#: What this target serves with an NVFP4 KV cache.
-NVFP4_KV_PATHS = dispatch.TargetPaths(served=_NVFP4_KV_ENGAGED | _NVFP4_KV_BASES)
-#: What a deployment outside the cells serves under ``auto``: no fast path.
+#: dep with an NVFP4 KV cache.
+DEP_NVFP4_KV_PATHS = dispatch.TargetPaths(served=_DEP_NVFP4_KV_ENGAGED | _DEP_NVFP4_KV_BASES)
+#: What a dep deployment outside the cells serves under ``auto``: no fast path.
 NO_PATHS = dispatch.TargetPaths(served=frozenset())
 
 
 @dataclass(frozen=True)
 class Cell:
-    """One deployment this target serves: the facts ``route`` does not read, and the plan."""
+    """One dep deployment this target serves: the facts ``route`` does not read, and the plan."""
 
     kv_cache: QuantAlgo
     moe_backend: str
     paths: dispatch.TargetPaths
 
 
-#: The deployments this target serves, per (ranks, LM head tensor-parallel in attention DP,
+#: The dep deployments this target serves, per (ranks, LM head tensor-parallel in attention DP,
 #: draft tokens per request), the facts ``route`` claims them by.
-CELLS: Mapping[Tuple[int, bool, int], Cell] = {
-    (16, False, 3): Cell(QuantAlgo.FP8, "TRTLLM", FP8_KV_PATHS),
-    (8, False, 3): Cell(QuantAlgo.FP8, "CUTEDSL", FP8_KV_PATHS),
-    (8, True, 5): Cell(QuantAlgo.NVFP4, "CUTEDSL", NVFP4_KV_PATHS),
+DEP_CELLS: Mapping[Tuple[int, bool, int], Cell] = {
+    (16, False, 3): Cell(QuantAlgo.FP8, "TRTLLM", DEP_FP8_KV_PATHS),
+    (8, False, 3): Cell(QuantAlgo.FP8, "CUTEDSL", DEP_FP8_KV_PATHS),
+    (8, True, 5): Cell(QuantAlgo.NVFP4, "CUTEDSL", DEP_NVFP4_KV_PATHS),
 }
 
 
-def target_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.TargetPaths:
-    """What this target serves on ``model_config``'s deployment: its cell's paths, or none.
-
-    Reads only the model config, so every rank decides alike.
+def _dep_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.TargetPaths:
+    """What this target serves on a dep deployment: its cell's paths, or none.
 
     Raises:
-        ValueError: under ``require``, if the deployment is not one of ``CELLS`` (its KV-cache
-            dtype or MoE backend is not its cell's).
+        ValueError: under ``require``, if the deployment is not one of ``DEP_CELLS`` (its
+            KV-cache dtype or MoE backend is not its cell's).
     """
     mapping = model_config.mapping
     spec_config = model_config.spec_config
@@ -180,7 +237,7 @@ def target_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.Ta
         getattr(spec_config, "max_draft_len", None),
     )
     kv_cache = model_config.quant_config.kv_cache_quant_algo
-    cell = CELLS.get(key)
+    cell = DEP_CELLS.get(key)
     if cell is not None and (kv_cache, model_config.moe_backend) == (
         cell.kv_cache,
         cell.moe_backend,
@@ -189,12 +246,12 @@ def target_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.Ta
     served = "; ".join(
         f"{ranks} ranks, LM head tensor-parallel in attention DP {lm_head_tp}, {draft} draft "
         f"tokens: KV cache {cell.kv_cache.value}, MoE backend {cell.moe_backend}"
-        for (ranks, lm_head_tp, draft), cell in CELLS.items()
+        for (ranks, lm_head_tp, draft), cell in DEP_CELLS.items()
     )
     words = (
         f"{key[0]} ranks, LM head tensor-parallel in attention DP {key[1]}, {key[2]} draft "
         f"tokens, KV cache {kv_cache}, MoE backend {model_config.moe_backend} is not a "
-        f"deployment this target serves ({served})"
+        f"dep deployment this target serves ({served})"
     )
     if mode is ModelingV2Mode.REQUIRE:
         raise ValueError(f"modeling_v2='require': {words}")
@@ -202,15 +259,32 @@ def target_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.Ta
     return NO_PATHS
 
 
-def adopt_model(model: "ModelingV2Glm52Nvfp4Sm103Dep", plan: dispatch.Plan) -> None:
+def target_paths(model_config: ModelConfig, mode: ModelingV2Mode) -> dispatch.TargetPaths:
+    """What this target serves on ``model_config``'s layout.
+
+    ``route`` admits only the tp8, tep and dep layouts, so attention DP and expert
+    parallelism tell them apart. Reads only the model config, so every rank decides alike.
+
+    Raises:
+        ValueError: under ``require``, if a dep deployment is not one of ``DEP_CELLS``.
+    """
+    mapping = model_config.mapping
+    if mapping.enable_attention_dp:
+        return _dep_paths(model_config, mode)
+    if mapping.moe_ep_size > 1:
+        return TEP_PATHS[mapping.tp_size]
+    return TP8_PATHS
+
+
+def _adopt_attention_dp(model: "ModelingV2Glm52Nvfp4Sm103Default") -> None:
     """Re-class the attention modules of every layer and the speculative worker, as far as
-    ``plan`` runs their fast paths.
+    the installed plan runs their fast paths.
 
     The attention modules of the decoder layers and of the draft layer, in index order, when
     the plan runs the rope and KV append (the attention paths that engage under attention DP
     extend it); then the speculative worker, which the built-in model builds after the draft
     layer, when the plan runs the draft loop. The layers, their MoE modules and the draft
-    layer stay the built-in classes. The fast paths read the installed plan, which is ``plan``.
+    layer stay the built-in classes.
 
     Raises:
         TypeError: if an attention module or the speculative worker is not exactly the
@@ -223,7 +297,7 @@ def adopt_model(model: "ModelingV2Glm52Nvfp4Sm103Dep", plan: dispatch.Plan) -> N
         GlmMtpWorker.adopt(model.spec_worker)
 
 
-class GlmMoeDsaDepModelingV2(DeepseekV3Model):
+class GlmMoeDsaModelingV2(DeepseekV3Model):
     """The text model: the built-in ``DeepseekV3Model``, whose forward picks a target per step.
 
     Never constructed directly: ``adopt`` re-classes the built-in model the shell
@@ -238,7 +312,7 @@ class GlmMoeDsaDepModelingV2(DeepseekV3Model):
     decode_prologue: Optional["DecodePrologueHook"] = None
 
     @classmethod
-    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaDepModelingV2":
+    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaModelingV2":
         """Re-class the built-in text model as this core and build its two targets.
 
         Raises:
@@ -290,7 +364,7 @@ class PrefillTarget:
 
     phase = dispatch.Phase.PREFILL
 
-    def __init__(self, core: GlmMoeDsaDepModelingV2) -> None:
+    def __init__(self, core: GlmMoeDsaModelingV2) -> None:
         self.core = core
 
     def forward(
@@ -321,7 +395,7 @@ class DecodeTarget:
 
     phase = dispatch.Phase.DECODE
 
-    def __init__(self, core: GlmMoeDsaDepModelingV2) -> None:
+    def __init__(self, core: GlmMoeDsaModelingV2) -> None:
         self.core = core
 
     def forward(
@@ -335,10 +409,10 @@ def forward(
     ) -> torch.Tensor:
         """The decoder layers of the text model on a decode step.
 
-        The built-in loop with the decode prologue's embedding. Each layer is the built-in
-        decoder layer, called as a module, so its forward hooks run; its attention runs the
-        attention modules' decode composition under the step's phase. The draft layer is not
-        run here; the last decoder layer applies the final norm.
+        The built-in loop with the decode prologue's embedding and the expert-prefetch join.
+        Each layer is called as a module, so its forward hooks run, and runs its decode
+        composition under the step's phase. The draft layer is not run here; the last decoder
+        layer applies the final norm.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError(
@@ -366,6 +440,9 @@ def forward(
                 residual=residual,
                 spec_metadata=spec_metadata,
             )
+        if dispatch.is_enabled("cross_layer_expert_prefetch"):
+            # The expert-prefetch stream rejoins before the drafts.
+            dispatch.get_impl("cross_layer_expert_prefetch", "join")()
 
         hidden_states = maybe_allgather_for_helix_cp(
             hidden_states, attn_metadata, self.core.mapping_with_cp
@@ -373,8 +450,8 @@ def forward(
         return hidden_states
 
 
-class ModelingV2Glm52Nvfp4Sm103Dep(GlmMoeDsaForCausalLM):
-    """GLM-5.2-NVFP4 on sm_103 at dep: the built-in causal LM with this target's core and dispatch."""
+class ModelingV2Glm52Nvfp4Sm103Default(GlmMoeDsaForCausalLM):
+    """GLM-5.2-NVFP4 on sm_103: the built-in causal LM with this target's core and dispatch."""
 
     def __init__(self, model_config: ModelConfig):
         world_size = model_config.mapping.world_size
@@ -385,9 +462,12 @@ def __init__(self, model_config: ModelConfig):
         plan = dispatch.resolve(mode, dict(knobs), paths)
         dispatch.install(plan)
         super().__init__(model_config)
-        GlmMoeDsaDepModelingV2.adopt(self.model)
+        GlmMoeDsaModelingV2.adopt(self.model)
         self.plan = plan
-        adopt_model(self, self.plan)
+        if model_config.mapping.enable_attention_dp:
+            _adopt_attention_dp(self)
+        else:
+            adopt_model(self, self.plan)
         dispatch.check_rank_uniform_adopted(adopted_state(self), world_size)
 
     @property
@@ -439,5 +519,15 @@ def forward(
                 **kwargs,
             )
 
+    def setup_aliases(self) -> None:
+        """The built-in aliases, then the cross-layer expert prefetch links."""
+        super().setup_aliases()
+        # Bind each decoder MoE layer to the next MoE layer's gate and experts (structural
+        # references only).
+        if dispatch.is_enabled("cross_layer_expert_prefetch"):
+            dispatch.get_impl("cross_layer_expert_prefetch", "link_layers")(
+                self.model.layers, self.config.num_hidden_layers
+            )
+
     def load_weights(self, weights):
         _weights.load(self, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/weights.py
similarity index 84%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/weights.py
index b77cd3c04a62..6310bbbd6cab 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/weights.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__default/weights.py
@@ -13,9 +13,9 @@
 from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
 
 if TYPE_CHECKING:
-    from .modeling import ModelingV2Glm52Nvfp4Sm103Tp8
+    from .modeling import ModelingV2Glm52Nvfp4Sm103Default
 
 
-def load(model: "ModelingV2Glm52Nvfp4Sm103Tp8", weights: Any) -> None:
+def load(model: "ModelingV2Glm52Nvfp4Sm103Default", weights: Any) -> None:
     """Load the checkpoint ``weights`` into ``model`` with the built-in loader."""
     GlmMoeDsaForCausalLM.load_weights(model, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md
deleted file mode 100644
index 6a8376b48935..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/README.md
+++ /dev/null
@@ -1,15 +0,0 @@
-# `glm_5_2_nvfp4__sm_103__dep`
-
-GLM-5.2 NVFP4 with attention data parallelism over eight or sixteen ranks and the experts
-expert-parallel over the same ranks.
-
-**This target is not separately optimized.** The family's fast paths were designed and tuned
-for low-latency decode on the TP8 target, `glm_5_2_nvfp4__sm_103__tp8`. Under attention data
-parallelism only the rank-local ones apply, and this target runs those from `../shared/`: the
-RoPE and KV append, the full-indexer layers' query front, the draft loop, the decode step's
-prologue and the host page retention policy. Its decoder layers, MoE modules and draft layer
-stay the built-in classes. No fast path was designed or tuned for DEP batch shapes.
-
-It is in the family because those rank-local paths also improve decode at this topology. DEP
-should be revisited in the future with fast paths of its own, in particular for the
-attention-DP communication and the expert-parallel MoE.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py
deleted file mode 100644
index b9046fba42dd..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__dep/weights.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The weights this target expects: the GLM-5.2 NVFP4 checkpoint as the built-in model loads it.
-
-The target's module tree is the built-in model's, so the built-in loader maps
-every checkpoint tensor, including the bf16 MTP layer and the NVFP4 scales.
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Any
-
-from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
-
-if TYPE_CHECKING:
-    from .modeling import ModelingV2Glm52Nvfp4Sm103Dep
-
-
-def load(model: "ModelingV2Glm52Nvfp4Sm103Dep", weights: Any) -> None:
-    """Load the checkpoint ``weights`` into ``model`` with the built-in loader."""
-    GlmMoeDsaForCausalLM.load_weights(model, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md
deleted file mode 100644
index eefb1402ee8d..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/README.md
+++ /dev/null
@@ -1,14 +0,0 @@
-# `glm_5_2_nvfp4__sm_103__tep`
-
-GLM-5.2 NVFP4 with the attention heads tensor-parallel over four or eight ranks and the
-experts expert-parallel over the same ranks.
-
-**This target is not separately optimized.** The family's fast paths were designed and tuned
-for low-latency decode on the TP8 target, `glm_5_2_nvfp4__sm_103__tp8`. This target runs the
-fast paths in `../shared/` whose contracts also cover its tensor-parallel size, in front of the
-built-in expert-parallel MoE. No fast path was designed or tuned for TEP batch shapes, and the
-routed experts run the built-in MoE.
-
-It is in the family because those shared paths also improve decode at this topology. TEP should
-be revisited in the future with fast paths of its own, in particular for the expert-parallel
-MoE and the larger decode batches this topology serves.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py
deleted file mode 100644
index d4c83cbaec3c..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/modeling.py
+++ /dev/null
@@ -1,357 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / tep.
-
-GLM-5.2, modelopt NVFP4 export: 78 layers, hidden 6144, 64 query heads, MLA with
-DSA sparse attention (index top-k 2048), layers 0-2 dense and 3-77 MoE (256
-routed experts at top-8 plus a shared expert). Layer 78 is the checkpoint's bf16
-MTP module, driven by the one-model MTP-Eagle loop at five draft tokens. `tep`
-is the attention heads tensor-parallel over four or eight ranks and the experts
-expert-parallel over the same ranks: each rank holds whole experts and runs the
-built-in expert-parallel MoE on them.
-
-The fast paths this target serves depend on its tensor-parallel size
-(``TARGET_PATHS``): at eight ranks the attention runs at the same shapes as on
-the tp8 target, so it serves every fast path of the family except the routed-MoE
-ones, which take only tensor-parallel experts; at four ranks it serves the paths
-whose kernels take four ranks (the indexer, the top-k selection, the router GEMM,
-the dense MLP, the draft loop and the step-loop paths) and the bases they extend.
-The layers run the family's composition (``shared.layers``), in which the block of
-a path this target does not serve never runs.
-
-The shell is the built-in ``GlmMoeDsaForCausalLM``: the same module tree, the
-same weights and the same engine preferences (KV cache manager, KV transceiver,
-sparse-attention defaults), so a built-in context server and this target's
-generation server exchange the same KV and indexer-K layout. Construction runs,
-in order:
-
-1. the rank-uniform check of the knobs (``dispatch.check_rank_uniform``), first,
-   so it runs before every collective the construction below makes;
-2. the dispatch plan (``dispatch.resolve``) of the paths this target serves at its
-   tensor-parallel size, before the built-in model is built, so a knob value no
-   rank can parse fails early;
-3. the plan is installed (``dispatch.install``): from here on the fast paths read
-   it;
-4. the built-in model, then the core: its text model adopted as
-   ``GlmMoeDsaTepModelingV2``, whose forward hands each step to
-   ``PrefillTarget`` or ``DecodeTarget``;
-5. ``adopt_model`` (``shared.layers``), which re-classes the decoder layers,
-   the draft layer, their attention modules (``shared.attention``) and their MoE
-   modules as the family's layer classes and attaches the fast-path objects the
-   plan enables to them, then re-classes the speculative worker
-   (``shared.spec_worker``). The attention backends build their FMHA libraries
-   after construction, in ``create_weights``, each with the libraries adoption
-   declared on it first;
-6. the rank-uniform check of that state (``dispatch.check_rank_uniform_adopted``).
-
-The step-loop hooks the target declares to the engine (``runtime_hooks``,
-``shared.runtime.make_runtime_hooks``) are built later, on the engine's first
-read, after the model is loaded.
-
-The rank-uniform check of step 1 also covers the shared switches the kernels
-read and the modeling_v2 mode (``dispatch.rank_uniform_inputs``).
-
-Every step runs under the phase of the target that takes it
-(``dispatch.step``): the shell's forward enters it around the built-in forward,
-so the model, the speculative worker and its draft loop all run in it. A
-prefill step runs the built-in decoder forward, the built-in layers and the
-built-in draft loop; a decode step runs ``DecodeTarget``'s loop, the layers'
-decode composition and the worker's draft loop.
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import resolved_mode
-from tensorrt_llm._torch.attention.attention import maybe_allgather_for_helix_cp
-from tensorrt_llm._torch.attention.backends import AttentionMetadata
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
-from tensorrt_llm._torch.speculative import SpecMetadata
-
-from ..shared import dispatch, runtime
-from ..shared.layers import adopt_model, adopted_state
-from . import weights as _weights
-
-if TYPE_CHECKING:
-    from ..shared.kernels.decode_prologue import DecodePrologueHook
-
-
-# The routed-MoE paths take tensor-parallel experts (every expert on every rank); with the experts
-# expert-parallel their kernels decline every call, so this target serves none of them.
-_ROUTED_MOE_PATHS = frozenset(
-    {
-        "routed_moe",
-        "routed_moe_into_hop",
-        "hop_routing",
-        "shared_expert_prefetch",
-        "draft_routed_moe",
-        "cross_layer_expert_prefetch",
-    }
-)
-# At four ranks the attention operands, the attention output and the post-MoE hop are outside
-# the kernels' shapes (sixteen local heads, a four-rank output projection, a single-node group);
-# these paths take four ranks.
-_TEP4_PATHS = frozenset(
-    {
-        "indexer_layer_front",
-        "rope_kv_append",
-        "indexer_topk",
-        "cross_rank_topk",
-        "cross_rank_topk_cluster_merge",
-        "cross_rank_topk_global_emit",
-        "router_gemm",
-        "dense_mlp",
-        "draft_front",
-        "draft_tail",
-        "draft_loop",
-        "decode_prologue",
-        "deferred_input_prep",
-        "sampler_store",
-        "execution_stream_loop",
-        "host_page_retention",
-    }
-)
-# The bases these paths extend at four ranks. Their kernels decline every layer there: the query
-# fronts on their layer cells, the post-attention hop on its native all-reduce, which a
-# single-node group runs without multi-node NVLink.
-_TEP4_BASES = frozenset({"absorbed_query", "draft_absorbed_query", "post_attention_hop"})
-_NO_LAYERS: frozenset = frozenset()
-
-#: What this target serves at each tensor-parallel size.
-TARGET_PATHS = {
-    4: dispatch.TargetPaths(
-        served=_TEP4_PATHS | _TEP4_BASES,
-        layer_cells={
-            **dispatch.LAYER_CELLS,
-            "absorbed_query": {"layer": _NO_LAYERS},
-            "draft_absorbed_query": {"layer": _NO_LAYERS},
-            "indexer_layer_front": {"target_layer": _NO_LAYERS, "draft_layer": _NO_LAYERS},
-        },
-        uncovered={"post_attention_hop": ("does not use MNNVL",)},
-    ),
-    8: dispatch.TargetPaths(served=frozenset(dispatch.FAST_PATH_NAMES) - _ROUTED_MOE_PATHS),
-}
-
-
-class GlmMoeDsaTepModelingV2(DeepseekV3Model):
-    """The text model: the built-in ``DeepseekV3Model``, whose forward picks a target per step.
-
-    Never constructed directly: ``adopt`` re-classes the built-in model the shell
-    has already built, so the module tree, the parameters and their names stay
-    the built-in model's. The class adds methods and plain attributes only.
-    """
-
-    _prefill: "PrefillTarget"
-    _decode: "DecodeTarget"
-    # The decode prologue the step loop runs, set when the target's runtime hooks are built;
-    # the decode target takes the embedding rows of its launch.
-    decode_prologue: Optional["DecodePrologueHook"] = None
-
-    @classmethod
-    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaTepModelingV2":
-        """Re-class the built-in text model as this core and build its two targets.
-
-        Raises:
-            TypeError: if ``model`` is not exactly the built-in ``DeepseekV3Model``.
-        """
-        if type(model) is not DeepseekV3Model:
-            raise TypeError(
-                f"{cls.__name__} adopts the built-in DeepseekV3Model, got {type(model).__name__}"
-            )
-        model.__class__ = cls
-        model._prefill = PrefillTarget(model)
-        model._decode = DecodeTarget(model)
-        return model
-
-    def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
-        """Which target runs this step.
-
-        One host-side integer read and nothing else: a decode-only CUDA graph is a
-        per-capture constant environment. A mixed batch -- context and generation
-        rows in one step -- goes to prefill, so prefill is the general case and
-        decode the specialization that may assume it has no context rows.
-        """
-        return self._decode if attn_metadata.num_contexts == 0 else self._prefill
-
-    def forward(
-        self,
-        attn_metadata,
-        input_ids: torch.IntTensor | None = None,
-        position_ids: torch.IntTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        spec_metadata=None,
-        **kwargs,
-    ) -> torch.Tensor:
-        return self._select_target(attn_metadata).forward(
-            attn_metadata,
-            input_ids=input_ids,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            spec_metadata=spec_metadata,
-            **kwargs,
-        )
-
-
-class PrefillTarget:
-    """The general case: context rows, and possibly generation rows beside them.
-
-    Runs the built-in decoder forward.
-    """
-
-    phase = dispatch.Phase.PREFILL
-
-    def __init__(self, core: GlmMoeDsaTepModelingV2) -> None:
-        self.core = core
-
-    def forward(
-        self,
-        attn_metadata,
-        input_ids: torch.IntTensor | None = None,
-        position_ids: torch.IntTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        spec_metadata=None,
-        **kwargs,
-    ) -> torch.Tensor:
-        return DeepseekV3Model.forward(
-            self.core,
-            attn_metadata,
-            input_ids=input_ids,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            spec_metadata=spec_metadata,
-            **kwargs,
-        )
-
-
-class DecodeTarget:
-    """The specialization: routed to only when there are no context rows.
-
-    Runs this target's decoder loop.
-    """
-
-    phase = dispatch.Phase.DECODE
-
-    def __init__(self, core: GlmMoeDsaTepModelingV2) -> None:
-        self.core = core
-
-    def forward(
-        self,
-        attn_metadata: AttentionMetadata,
-        input_ids: Optional[torch.IntTensor] = None,
-        position_ids: Optional[torch.IntTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        spec_metadata: Optional[SpecMetadata] = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        """The decoder layers of the text model on a decode step.
-
-        The built-in loop with the decode prologue's embedding.
-        Each layer is called as a module, so its forward hooks run, and runs its decode
-        composition under the step's phase. The draft layer is not run here; the last decoder
-        layer applies the final norm.
-        """
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
-            )
-
-        prologue = self.core.decode_prologue
-        if inputs_embeds is None and prologue is not None:
-            # The rows the decode prologue launch embedded (None unless that launch ran for
-            # exactly these ids and this table).
-            inputs_embeds = prologue.take_embedding(
-                attn_metadata, input_ids, self.core.embed_tokens
-            )
-        if inputs_embeds is None:
-            inputs_embeds = self.core.embed_tokens(input_ids)
-
-        hidden_states = inputs_embeds
-        residual = None
-
-        for idx, decoder_layer in enumerate(self.core.layers[: self.core.num_hidden_layers]):
-            hidden_states, residual = decoder_layer(
-                position_ids=position_ids,
-                hidden_states=hidden_states,
-                attn_metadata=attn_metadata,
-                residual=residual,
-                spec_metadata=spec_metadata,
-            )
-
-        hidden_states = maybe_allgather_for_helix_cp(
-            hidden_states, attn_metadata, self.core.mapping_with_cp
-        )
-        return hidden_states
-
-
-class ModelingV2Glm52Nvfp4Sm103Tep(GlmMoeDsaForCausalLM):
-    """GLM-5.2-NVFP4 on sm_103 at tep: the built-in causal LM with this target's core and dispatch."""
-
-    def __init__(self, model_config: ModelConfig):
-        world_size = model_config.mapping.world_size
-        knobs = dispatch.read_knobs()
-        mode = resolved_mode(type(self))
-        dispatch.check_rank_uniform(dispatch.rank_uniform_inputs(knobs, mode), world_size)
-        plan = dispatch.resolve(mode, dict(knobs), TARGET_PATHS[model_config.mapping.tp_size])
-        dispatch.install(plan)
-        super().__init__(model_config)
-        GlmMoeDsaTepModelingV2.adopt(self.model)
-        self.plan = plan
-        adopt_model(self, self.plan)
-        dispatch.check_rank_uniform_adopted(adopted_state(self), world_size)
-
-    @property
-    def runtime_hooks(self) -> RuntimeHooks:
-        """The step-loop hooks this target declares, built on the engine's first read.
-
-        Built after the model is loaded: a hook may hold pinned host memory, which model
-        construction (under meta-device initialization) must not allocate.
-
-        Raises:
-            RuntimeError: if building the hooks raises ``AttributeError``, which
-                ``runtime_hooks_of`` would otherwise read as "no hooks declared".
-        """
-        hooks = self.__dict__.get("_runtime_hooks")
-        if hooks is None:
-            try:
-                hooks = runtime.make_runtime_hooks(self)
-            except AttributeError as exc:
-                raise RuntimeError(
-                    f"building the step-loop hooks of {type(self).__name__} failed: {exc}"
-                ) from exc
-            self._runtime_hooks = hooks
-        return hooks
-
-    def forward(
-        self,
-        attn_metadata,
-        input_ids: torch.IntTensor | None = None,
-        position_ids: torch.IntTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        return_context_logits: bool = False,
-        spec_metadata=None,
-        resource_manager=None,
-        **kwargs,
-    ):
-        """The built-in forward under the phase of the target that takes this step.
-
-        The scope covers the text model, the speculative worker and its draft loop.
-        """
-        with dispatch.step(self.model._select_target(attn_metadata).phase):
-            return super().forward(
-                attn_metadata,
-                input_ids=input_ids,
-                position_ids=position_ids,
-                inputs_embeds=inputs_embeds,
-                return_context_logits=return_context_logits,
-                spec_metadata=spec_metadata,
-                resource_manager=resource_manager,
-                **kwargs,
-            )
-
-    def load_weights(self, weights):
-        _weights.load(self, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py
deleted file mode 100644
index 09b1849fe5f4..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep/weights.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The weights this target expects: the GLM-5.2 NVFP4 checkpoint as the built-in model loads it.
-
-The target's module tree is the built-in model's, so the built-in loader maps
-every checkpoint tensor, including the bf16 MTP layer and the NVFP4 scales.
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Any
-
-from tensorrt_llm._torch.models.modeling_deepseekv3 import GlmMoeDsaForCausalLM
-
-if TYPE_CHECKING:
-    from .modeling import ModelingV2Glm52Nvfp4Sm103Tep
-
-
-def load(model: "ModelingV2Glm52Nvfp4Sm103Tep", weights: Any) -> None:
-    """Load the checkpoint ``weights`` into ``model`` with the built-in loader."""
-    GlmMoeDsaForCausalLM.load_weights(model, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py
deleted file mode 100644
index 87156adf61ff..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GLM-5.2-NVFP4 / sm_103 / tp8."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py
deleted file mode 100644
index 94b44234a92b..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Kernels of the GLM-5.2-NVFP4 sm_103 tp8 target, one package per fast path."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py
deleted file mode 100644
index 7abbfb3c7d0d..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/modeling.py
+++ /dev/null
@@ -1,318 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""ModelingV2 target: glm-5.2-nvfp4 / sm_103 / tp8.
-
-GLM-5.2, modelopt NVFP4 export: 78 layers, hidden 6144, 64 query heads, MLA with
-DSA sparse attention (index top-k 2048), layers 0-2 dense and 3-77 MoE (256
-routed experts at top-8 plus a shared expert). Layer 78 is the checkpoint's bf16
-MTP module, driven by the one-model MTP-Eagle loop at five draft tokens. `tp8`
-is the attention heads and the experts tensor-parallel over eight ranks.
-
-The shell is the built-in ``GlmMoeDsaForCausalLM``: the same module tree, the
-same weights and the same engine preferences (KV cache manager, KV transceiver,
-sparse-attention defaults), so a built-in context server and this target's
-generation server exchange the same KV and indexer-K layout. Construction runs,
-in order:
-
-1. the rank-uniform check of the knobs (``dispatch.check_rank_uniform``), first,
-   so it runs before every collective the construction below makes;
-2. the dispatch plan (``dispatch.resolve``), before the built-in model is built,
-   so a knob value no rank can parse fails early;
-3. the plan is installed (``dispatch.install``): from here on the fast paths read
-   it;
-4. the built-in model, then the core: its text model adopted as
-   ``GlmMoeDsaTp8ModelingV2``, whose forward hands each step to
-   ``PrefillTarget`` or ``DecodeTarget``;
-5. ``adopt_model`` (``shared.layers``), which re-classes the decoder layers,
-   the draft layer, their attention modules (``shared.attention``) and their MoE
-   modules as the family's layer classes and attaches the fast-path objects the
-   plan enables to them, then re-classes the speculative worker
-   (``shared.spec_worker``). The attention backends build their FMHA libraries
-   after construction, in ``create_weights``, each with the libraries adoption
-   declared on it first;
-6. the rank-uniform check of that state (``dispatch.check_rank_uniform_adopted``).
-
-The step-loop hooks the target declares to the engine (``runtime_hooks``,
-``shared.runtime.make_runtime_hooks``) are built later, on the engine's first
-read, after the model is loaded.
-
-The rank-uniform check of step 1 also covers the shared switches the kernels
-read and the modeling_v2 mode (``dispatch.rank_uniform_inputs``).
-
-Every step runs under the phase of the target that takes it
-(``dispatch.step``): the shell's forward enters it around the built-in forward,
-so the model, the speculative worker and its draft loop all run in it. A
-prefill step runs the built-in decoder forward, the built-in layers and the
-built-in draft loop; a decode step runs ``DecodeTarget``'s loop, the layers'
-decode composition and the worker's draft loop.
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import resolved_mode
-from tensorrt_llm._torch.attention.attention import maybe_allgather_for_helix_cp
-from tensorrt_llm._torch.attention.backends import AttentionMetadata
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Model, GlmMoeDsaForCausalLM
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
-from tensorrt_llm._torch.speculative import SpecMetadata
-
-from ..shared import dispatch, runtime
-from ..shared.layers import adopt_model, adopted_state
-from . import weights as _weights
-
-if TYPE_CHECKING:
-    from ..shared.kernels.decode_prologue import DecodePrologueHook
-
-
-# Every fast path of the family, with the family's tables; the routed-MoE paths run from this
-# target's own kernel packages.
-TARGET_PATHS = dispatch.TargetPaths(
-    kernel_modules={
-        path: f"{__package__}.kernels.{package}"
-        for path, package in (
-            ("routed_moe", "routed_moe"),
-            ("routed_moe_into_hop", "routed_moe"),
-            ("hop_routing", "hop_routing"),
-            ("shared_expert_prefetch", "shared_expert_prefetch"),
-            ("draft_routed_moe", "draft_routed_moe"),
-            ("cross_layer_expert_prefetch", "cross_layer_expert_prefetch"),
-        )
-    }
-)
-
-
-class GlmMoeDsaTp8ModelingV2(DeepseekV3Model):
-    """The text model: the built-in ``DeepseekV3Model``, whose forward picks a target per step.
-
-    Never constructed directly: ``adopt`` re-classes the built-in model the shell
-    has already built, so the module tree, the parameters and their names stay
-    the built-in model's. The class adds methods and plain attributes only.
-    """
-
-    _prefill: "PrefillTarget"
-    _decode: "DecodeTarget"
-    # The decode prologue the step loop runs, set when the target's runtime hooks are built;
-    # the decode target takes the embedding rows of its launch.
-    decode_prologue: Optional["DecodePrologueHook"] = None
-
-    @classmethod
-    def adopt(cls, model: DeepseekV3Model) -> "GlmMoeDsaTp8ModelingV2":
-        """Re-class the built-in text model as this core and build its two targets.
-
-        Raises:
-            TypeError: if ``model`` is not exactly the built-in ``DeepseekV3Model``.
-        """
-        if type(model) is not DeepseekV3Model:
-            raise TypeError(
-                f"{cls.__name__} adopts the built-in DeepseekV3Model, got {type(model).__name__}"
-            )
-        model.__class__ = cls
-        model._prefill = PrefillTarget(model)
-        model._decode = DecodeTarget(model)
-        return model
-
-    def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
-        """Which target runs this step.
-
-        One host-side integer read and nothing else: a decode-only CUDA graph is a
-        per-capture constant environment. A mixed batch -- context and generation
-        rows in one step -- goes to prefill, so prefill is the general case and
-        decode the specialization that may assume it has no context rows.
-        """
-        return self._decode if attn_metadata.num_contexts == 0 else self._prefill
-
-    def forward(
-        self,
-        attn_metadata,
-        input_ids: torch.IntTensor | None = None,
-        position_ids: torch.IntTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        spec_metadata=None,
-        **kwargs,
-    ) -> torch.Tensor:
-        return self._select_target(attn_metadata).forward(
-            attn_metadata,
-            input_ids=input_ids,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            spec_metadata=spec_metadata,
-            **kwargs,
-        )
-
-
-class PrefillTarget:
-    """The general case: context rows, and possibly generation rows beside them.
-
-    Runs the built-in decoder forward.
-    """
-
-    phase = dispatch.Phase.PREFILL
-
-    def __init__(self, core: GlmMoeDsaTp8ModelingV2) -> None:
-        self.core = core
-
-    def forward(
-        self,
-        attn_metadata,
-        input_ids: torch.IntTensor | None = None,
-        position_ids: torch.IntTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        spec_metadata=None,
-        **kwargs,
-    ) -> torch.Tensor:
-        return DeepseekV3Model.forward(
-            self.core,
-            attn_metadata,
-            input_ids=input_ids,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            spec_metadata=spec_metadata,
-            **kwargs,
-        )
-
-
-class DecodeTarget:
-    """The specialization: routed to only when there are no context rows.
-
-    Runs this target's decoder loop.
-    """
-
-    phase = dispatch.Phase.DECODE
-
-    def __init__(self, core: GlmMoeDsaTp8ModelingV2) -> None:
-        self.core = core
-
-    def forward(
-        self,
-        attn_metadata: AttentionMetadata,
-        input_ids: Optional[torch.IntTensor] = None,
-        position_ids: Optional[torch.IntTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        spec_metadata: Optional[SpecMetadata] = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        """The decoder layers of the text model on a decode step.
-
-        The built-in loop with the decode prologue's embedding and the expert-prefetch join.
-        Each layer is called as a module, so its forward hooks run, and runs its decode
-        composition under the step's phase. The draft layer is not run here; the last decoder
-        layer applies the final norm.
-        """
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
-            )
-
-        prologue = self.core.decode_prologue
-        if inputs_embeds is None and prologue is not None:
-            # The rows the decode prologue launch embedded (None unless that launch ran for
-            # exactly these ids and this table).
-            inputs_embeds = prologue.take_embedding(
-                attn_metadata, input_ids, self.core.embed_tokens
-            )
-        if inputs_embeds is None:
-            inputs_embeds = self.core.embed_tokens(input_ids)
-
-        hidden_states = inputs_embeds
-        residual = None
-
-        for idx, decoder_layer in enumerate(self.core.layers[: self.core.num_hidden_layers]):
-            hidden_states, residual = decoder_layer(
-                position_ids=position_ids,
-                hidden_states=hidden_states,
-                attn_metadata=attn_metadata,
-                residual=residual,
-                spec_metadata=spec_metadata,
-            )
-        if dispatch.is_enabled("cross_layer_expert_prefetch"):
-            # The expert-prefetch stream rejoins before the drafts.
-            dispatch.get_impl("cross_layer_expert_prefetch", "join")()
-
-        hidden_states = maybe_allgather_for_helix_cp(
-            hidden_states, attn_metadata, self.core.mapping_with_cp
-        )
-        return hidden_states
-
-
-class ModelingV2Glm52Nvfp4Sm103Tp8(GlmMoeDsaForCausalLM):
-    """GLM-5.2-NVFP4 on sm_103 at tp8: the built-in causal LM with this target's core and dispatch."""
-
-    def __init__(self, model_config: ModelConfig):
-        world_size = model_config.mapping.world_size
-        knobs = dispatch.read_knobs()
-        mode = resolved_mode(type(self))
-        dispatch.check_rank_uniform(dispatch.rank_uniform_inputs(knobs, mode), world_size)
-        plan = dispatch.resolve(mode, dict(knobs), TARGET_PATHS)
-        dispatch.install(plan)
-        super().__init__(model_config)
-        GlmMoeDsaTp8ModelingV2.adopt(self.model)
-        self.plan = plan
-        adopt_model(self, self.plan)
-        dispatch.check_rank_uniform_adopted(adopted_state(self), world_size)
-
-    @property
-    def runtime_hooks(self) -> RuntimeHooks:
-        """The step-loop hooks this target declares, built on the engine's first read.
-
-        Built after the model is loaded: a hook may hold pinned host memory, which model
-        construction (under meta-device initialization) must not allocate.
-
-        Raises:
-            RuntimeError: if building the hooks raises ``AttributeError``, which
-                ``runtime_hooks_of`` would otherwise read as "no hooks declared".
-        """
-        hooks = self.__dict__.get("_runtime_hooks")
-        if hooks is None:
-            try:
-                hooks = runtime.make_runtime_hooks(self)
-            except AttributeError as exc:
-                raise RuntimeError(
-                    f"building the step-loop hooks of {type(self).__name__} failed: {exc}"
-                ) from exc
-            self._runtime_hooks = hooks
-        return hooks
-
-    def forward(
-        self,
-        attn_metadata,
-        input_ids: torch.IntTensor | None = None,
-        position_ids: torch.IntTensor | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        return_context_logits: bool = False,
-        spec_metadata=None,
-        resource_manager=None,
-        **kwargs,
-    ):
-        """The built-in forward under the phase of the target that takes this step.
-
-        The scope covers the text model, the speculative worker and its draft loop.
-        """
-        with dispatch.step(self.model._select_target(attn_metadata).phase):
-            return super().forward(
-                attn_metadata,
-                input_ids=input_ids,
-                position_ids=position_ids,
-                inputs_embeds=inputs_embeds,
-                return_context_logits=return_context_logits,
-                spec_metadata=spec_metadata,
-                resource_manager=resource_manager,
-                **kwargs,
-            )
-
-    def setup_aliases(self) -> None:
-        """The built-in aliases, then the cross-layer expert prefetch links."""
-        super().setup_aliases()
-        # Bind each decoder MoE layer to the next MoE layer's gate and experts (structural
-        # references only).
-        if dispatch.is_enabled("cross_layer_expert_prefetch"):
-            dispatch.get_impl("cross_layer_expert_prefetch", "link_layers")(
-                self.model.layers, self.config.num_hidden_layers
-            )
-
-    def load_weights(self, weights):
-        _weights.load(self, weights)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py
index 8a40d97a8102..86294131b70e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py
@@ -3,26 +3,27 @@
 """Where a GlmMoeDsaForCausalLM config lands. Read this file and you know.
 
 ``route`` is one forward-reading decision: the facts that fix which code runs,
-compared against the deployment the target below was written for. When it does
-not fit, ``route`` raises ``NoTarget`` stating those facts. In ``auto`` the
-engine then uses the built-in GlmMoeDsa implementation; in ``require`` it raises
-that error.
-
-The three targets are decode specializations of one checkpoint on one GPU
-architecture, and differ in how the attention and the experts are split. Their
-fast paths are built for the checkpoint's operand shapes and layer pattern, their
-ranks and the rows per request (the target token and its draft tokens), so each of
-those facts is a routing criterion rather than a runtime branch:
-
-* ``glm_5_2_nvfp4__sm_103__tp8``: attention and experts tensor-parallel over
-  eight ranks, five draft tokens;
-* ``glm_5_2_nvfp4__sm_103__tep``: attention tensor-parallel over four or eight
-  ranks and the experts expert-parallel over the same ranks, five draft tokens;
-* ``glm_5_2_nvfp4__sm_103__dep``: attention data parallelism and the experts
-  expert-parallel over the same ranks, at the rank counts, LM-head splits and
-  draft lengths of ``_DEP``. Two more facts fix what its fast paths serve, the
-  KV-cache dtype and the MoE backend; ``route`` does not read them, so the
-  target checks them when it is built.
+compared against the deployments the target below was written for. When none
+fits, ``route`` raises ``NoTarget`` stating those facts. In ``auto`` the engine
+then uses the built-in GlmMoeDsa implementation; in ``require`` it raises that
+error.
+
+One target, ``glm_5_2_nvfp4__sm_103__default``, serves one checkpoint on one GPU
+architecture at the layouts below; it picks the fast paths of each layout when
+it is built. The fast paths are built for the checkpoint's operand shapes and
+layer pattern, the ranks and the rows per request (the target token and its
+draft tokens), so each of those facts is a routing criterion rather than a
+runtime branch. A layout is admitted only once it has passed its gate:
+
+* tp8: attention and experts tensor-parallel over eight ranks, five draft
+  tokens. The fast paths were designed and tuned for this layout;
+* tep: attention tensor-parallel over four or eight ranks and the experts
+  expert-parallel over the same ranks, five draft tokens;
+* dep: attention data parallelism and the experts expert-parallel over the same
+  ranks, at the rank counts, LM-head splits and draft lengths of ``_DEP``. Two
+  more facts fix what its fast paths serve, the KV-cache dtype and the MoE
+  backend; the target checks both when it is built, since ``route`` does not see
+  the MoE backend.
 """
 
 from __future__ import annotations
@@ -123,8 +124,8 @@ def _tp8(m: Any) -> bool:
     """Attention heads and experts both tensor-parallel over eight ranks, with every
     expert on every rank and no attention data parallelism.
 
-    The split decides which weights each rank loads and into which shapes, so it
-    selects a target rather than a runtime branch.
+    The split decides which weights each rank loads and into which shapes, so it is
+    a routing criterion; the target picks the layout's fast paths when it is built.
     """
     return (
         m.world_size == _TP
@@ -141,8 +142,8 @@ def _tep(m: Any) -> bool:
     """Attention heads tensor-parallel over four or eight ranks and the experts expert-parallel
     over the same ranks (each rank holds whole experts), with no attention data parallelism.
 
-    The split decides which weights each rank loads and into which shapes, so it
-    selects a target rather than a runtime branch.
+    The split decides which weights each rank loads and into which shapes, so it is
+    a routing criterion; the target picks the layout's fast paths when it is built.
     """
     return (
         m.world_size == m.tp_size
@@ -160,8 +161,8 @@ def _dep(m: Any) -> bool:
     and the experts expert-parallel over the same ranks, at a rank count and LM-head split
     of ``_DEP``.
 
-    The split decides which weights each rank loads and into which shapes, so it
-    selects a target rather than a runtime branch.
+    The split decides which weights each rank loads and into which shapes, so it is
+    a routing criterion; the target picks the layout's fast paths when it is built.
     """
     return (
         m.world_size == m.tp_size
@@ -210,6 +211,16 @@ def _speculation(spec_config: Optional[Any]) -> str:
     return words
 
 
+def _layout_and_speculation(mapping: Any, spec_config: Optional[Any]) -> bool:
+    """One of the admitted layouts, with the speculation its fast paths are built for."""
+    if _tp8(mapping) or _tep(mapping):
+        return _mtp_eagle_linear(spec_config, _DRAFT_LEN)
+    if _dep(mapping):
+        draft_len = _DEP[(mapping.world_size, bool(mapping.enable_lm_head_tp_in_adp))]
+        return _mtp_eagle_linear(spec_config, draft_len)
+    return False
+
+
 def route(
     pretrained_config: Any,
     mapping: Any,
@@ -220,7 +231,7 @@ def route(
     """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
 
     ``kv_cache_quant_algo`` is the KV-cache algorithm the engine will run with. It
-    selects the fast paths of a dep deployment, which the dep target checks together
+    selects the fast paths of a dep deployment, which the target checks together
     with the MoE backend when it is built, so ``route`` only reports it.
 
     The target is imported where it is returned, so declining costs no import
@@ -230,52 +241,19 @@ def route(
     # None and the shape does not match.
     shape = tuple(getattr(pretrained_config, name, None) for name in _SHAPE_FIELDS)
 
-    # GLM-5.2-NVFP4 (bf16 MTP layer) on sm103, 8 ranks of tensor-parallel attention
-    # and experts, one-model MTP-Eagle at five draft tokens. The quantization is
-    # read last, and only for this deployment.
+    # GLM-5.2-NVFP4 (bf16 MTP layer) on sm103 at an admitted layout with its speculation.
+    # The quantization is read last, and only for such a deployment.
     quantization = "not read"
     if (
         compute_capability == _SM
         and shape == _SHAPE
-        and _tp8(mapping)
-        and _mtp_eagle_linear(spec_config, _DRAFT_LEN)
-    ):
-        served, quantization = _nvfp4_with_bf16_mtp_layer(pretrained_config)
-        if served:
-            from .glm_5_2_nvfp4__sm_103__tp8.modeling import ModelingV2Glm52Nvfp4Sm103Tp8
-
-            return ModelingV2Glm52Nvfp4Sm103Tp8
-
-    # The same checkpoint, architecture and speculation, with 4 or 8 ranks of
-    # tensor-parallel attention and the experts expert-parallel over them.
-    if (
-        compute_capability == _SM
-        and shape == _SHAPE
-        and _tep(mapping)
-        and _mtp_eagle_linear(spec_config, _DRAFT_LEN)
-    ):
-        served, quantization = _nvfp4_with_bf16_mtp_layer(pretrained_config)
-        if served:
-            from .glm_5_2_nvfp4__sm_103__tep.modeling import ModelingV2Glm52Nvfp4Sm103Tep
-
-            return ModelingV2Glm52Nvfp4Sm103Tep
-
-    # The same checkpoint and architecture with attention data parallelism and the experts
-    # expert-parallel over the same ranks, at the draft length of the rank count and LM-head
-    # split.
-    if (
-        compute_capability == _SM
-        and shape == _SHAPE
-        and _dep(mapping)
-        and _mtp_eagle_linear(
-            spec_config, _DEP[(mapping.world_size, bool(mapping.enable_lm_head_tp_in_adp))]
-        )
+        and _layout_and_speculation(mapping, spec_config)
     ):
         served, quantization = _nvfp4_with_bf16_mtp_layer(pretrained_config)
         if served:
-            from .glm_5_2_nvfp4__sm_103__dep.modeling import ModelingV2Glm52Nvfp4Sm103Dep
+            from .glm_5_2_nvfp4__sm_103__default.modeling import ModelingV2Glm52Nvfp4Sm103Default
 
-            return ModelingV2Glm52Nvfp4Sm103Dep
+            return ModelingV2Glm52Nvfp4Sm103Default
 
     raise NoTarget(
         f"compute capability {compute_capability}, config shape {shape}, quantization "
@@ -284,7 +262,7 @@ def route(
         f"moe_ep={mapping.moe_ep_size} attention_dp={mapping.enable_attention_dp} "
         f"lm_head_tp_in_adp={mapping.enable_lm_head_tp_in_adp}, "
         f"speculative decoding {_speculation(spec_config)}, KV cache "
-        f"{kv_cache_quant_algo or 'bf16'}; the glm_moe_dsa targets are "
+        f"{kv_cache_quant_algo or 'bf16'}; the glm_moe_dsa target is "
         "the GLM-5.2 modelopt NVFP4 checkpoint (bf16 MTP layer) on sm103, with one-model "
         "MTP-Eagle as a linear chain, without the suffix-automaton enhancer or relaxed "
         "acceptance, at tp8 (8 ranks of tensor-parallel attention and experts) or tep (4 or 8 "
diff --git a/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py b/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
index 7b6fc2d13e50..cff94506b7cb 100644
--- a/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
+++ b/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py
@@ -12,19 +12,19 @@
 # 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.
-"""Whole-model gates for the glm_moe_dsa modeling_v2 targets (GLM-5.2 NVFP4, sm_103).
+"""Whole-model gates for the glm_moe_dsa modeling_v2 target (GLM-5.2 NVFP4, sm_103).
 
-One case per target, each serving GSM8K end to end with ``modeling_v2="require"``.
-Under ``"auto"`` a deployment that missed a target's route would fall back to the
-built-in model, pass, and report the built-in's accuracy as the target's; under
-``"require"`` it raises instead, and so does a step a fast path covers but cannot
-run. The mode and the environment below are LLM API arguments, so they reach every
-rank.
+One case per layout family the target admits (tp8, tep, dep), each serving GSM8K end
+to end with ``modeling_v2="require"``. Under ``"auto"`` a deployment that missed the
+route would fall back to the built-in model, pass, and report the built-in's accuracy
+as the target's; under ``"require"`` it raises instead, and so does a step a fast
+path covers but cannot run. The mode and the environment below are LLM API
+arguments, so they reach every rank.
 
 Each deployment is one the family's ``route()`` claims: the GLM-5.2 modelopt NVFP4
 checkpoint (bf16 MTP layer) with one-model MTP-Eagle as a linear chain, at the
-target's topology and draft length; the DEP case also matches its cell's KV-cache
-dtype and MoE backend.
+layout's draft length; the dep case also matches its cell's KV-cache dtype and MoE
+backend.
 """
 
 import pytest
@@ -43,20 +43,20 @@
 from ...conftest import llm_models_root
 from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness
 
-# The targets assert their own SM at construction.
+# The target asserts its own SM at construction.
 skip_not_sm103 = pytest.mark.skipif(
-    get_sm_version() != 103, reason="the glm_moe_dsa targets are certified on sm_103 only"
+    get_sm_version() != 103, reason="the glm_moe_dsa target is certified on sm_103 only"
 )
 
 # On every rank: check the target's step contract on the first forward, and run the
 # indexer key projection in BF16, the form the family's fused indexer paths take.
-# The deployments below are the ones the targets were validated at: per-target batch,
+# The deployments below are the ones the target was validated at: per-layout batch,
 # KV-cache dtype, MoE backend and DSA kernels.
 _RANK_ENV = {MODELING_V2_VALIDATE_ENV: "1", "TRTLLM_DSA_INDEXER_BF16": "1"}
 
 
-class TestModelingV2Glm52Nvfp4Sm103(LlmapiAccuracyTestHarness):
-    """glm-5.2-nvfp4 / sm_103 at tp8, tep4 and dep8."""
+class TestModelingV2Glm52Nvfp4Sm103Default(LlmapiAccuracyTestHarness):
+    """glm-5.2-nvfp4 / sm_103 / default at tp8, tep4 and dep8."""
 
     MODEL_NAME = "zai-org/GLM-5.2"
     MODEL_PATH = f"{llm_models_root()}/GLM-5.2-NVFP4"
@@ -89,7 +89,7 @@ def _evaluate(self, **llm_kwargs) -> None:
     @skip_not_sm103
     @pytest.mark.skip_less_mpi_world_size(8)
     def test_gsm8k_tp8(self):
-        """The low-latency target: attention and experts tensor-parallel over 8 ranks."""
+        """tp8, the layout the fast paths are tuned for: attention and experts TP over 8 ranks."""
         self._evaluate(
             tensor_parallel_size=8,
             moe_expert_parallel_size=1,
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index 6f20557c6c8c..f500ac76f678 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -28,8 +28,8 @@ l0_gb300_multi_gpus:
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[dep4_moetp4]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[tp4_ep4]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4]
-  # GLM-5.2 NVFP4 modeling_v2 tep target, GSM8K under modeling_v2="require". Measured 17 min on 4 GB300.
-  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103::test_gsm8k_tep4 TIMEOUT (60)
+  # GLM-5.2 NVFP4 modeling_v2 target at tep4, GSM8K under modeling_v2="require". Measured 17 min on 4 GB300.
+  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103Default::test_gsm8k_tep4 TIMEOUT (60)
   # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere
   # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry
   # is what picks them up:
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
index 84c4cbfc8179..df5108775837 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
@@ -21,7 +21,7 @@ l0_gb300_multi_nodes_node2_gpu8:
   - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[dep8]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[tp8_moetp2_ep4]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[dep8]
-  # GLM-5.2 NVFP4 modeling_v2 targets at tp8 and dep8, GSM8K under modeling_v2="require".
+  # GLM-5.2 NVFP4 modeling_v2 target at tp8 and dep8, GSM8K under modeling_v2="require".
   # Measured on 2x4 GB300: tp8 (batch 1) 56 min, dep8 10 min.
-  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103::test_gsm8k_tp8 TIMEOUT (90)
-  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103::test_gsm8k_dep8 TIMEOUT (60)
+  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103Default::test_gsm8k_tp8 TIMEOUT (90)
+  - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103Default::test_gsm8k_dep8 TIMEOUT (60)

From 4183b67de505725763d3dcce672a3fc4f8ce7189 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Wed, 21 Oct 2026 10:37:58 -0700
Subject: [PATCH 10/19] [None][fix] modeling_v2 GLM-5.2: import the Blackwell
 CuTe helpers from their new location

main split tensorrt_llm/_torch/cute_dsl_kernels/ into the domain kernel
trees (#19275); cute_ptx_utils and utils now live under
tensorrt_llm/_torch/kernels/blackwell/.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../shared/kernels/attention_output_tail/cute_fused.py        | 2 +-
 .../shared/kernels/attention_output_tail/cute_oproj.py        | 2 +-
 .../shared/kernels/attention_output_tail/cute_oproj_ring.py   | 2 +-
 .../shared/kernels/attention_output_tail/into_hop.py          | 2 +-
 .../kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py    | 4 ++--
 .../kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py   | 4 ++--
 .../shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py  | 2 +-
 .../kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py     | 2 +-
 .../shared/kernels/oproj_weight_staging/cute_oproj_stage.py   | 2 +-
 9 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py
index 5102e384c037..d3290bc895d7 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_fused.py
@@ -103,7 +103,7 @@ def _build_fused(
     from cutlass.cutlass_dsl import T, dsl_user_op
     from cutlass.utils.smem_allocator import SmemAllocator
 
-    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+    from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import mma_sync
 
     k_dim = heads * v_dim
     geo = plan(n_dim, k_dim, grid)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py
index 6ab70ca5ab18..cba8c6527c39 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj.py
@@ -104,7 +104,7 @@ def _build(
     from cutlass.cutlass_dsl import dsl_user_op
     from cutlass.utils.smem_allocator import SmemAllocator
 
-    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+    from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import mma_sync
 
     geo = plan(n_dim, k_dim, grid)
     rows_max = geo["rows_max"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py
index 521fe4810d70..429c7de0c99b 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/cute_oproj_ring.py
@@ -103,7 +103,7 @@ def _build(
     from cutlass.cutlass_dsl import T, dsl_user_op
     from cutlass.utils.smem_allocator import SmemAllocator
 
-    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+    from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import mma_sync
 
     geo = plan(n_dim, k_dim, grid, slots)
     rows = geo["rows"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py
index 7a6dec5e0681..ec5ca1b57814 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/attention_output_tail/into_hop.py
@@ -114,7 +114,7 @@ def _build(
     from cutlass.cutlass_dsl import T, dsl_user_op
     from cutlass.utils.smem_allocator import SmemAllocator
 
-    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+    from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import mma_sync
 
     geo = plan(n_dim, k_dim, grid)
     rows_max = geo["rows_max"]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py
index bdf3d8013d38..16608edd3104 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_layer_front/fp8_paged_mqa_logits_cp.py
@@ -58,7 +58,7 @@
 
 indexer_topk variant of the CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100).
 
-Derived from ``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
+Derived from ``attention/kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
 for the GLM-5.x indexer shape, where all verify rows of a request fit one UMMA
 tile (``N = next_n * num_heads <= 256``). The only functional change is the
 size of the per-thread register cache of epilogue weights
@@ -132,7 +132,7 @@
 from cutlass.cutlass_dsl import dsl_user_op
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 
-from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import griddepcontrol_launch_dependents
+from tensorrt_llm._torch.kernels.blackwell.utils import griddepcontrol_launch_dependents
 
 # CuTe DSL CUDA 13 validates rounding modes as string literals. The string
 # form is also accepted by older wrappers, so keep it version-independent.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py
index f064a7aa5871..b10102353b79 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_cross_rank.py
@@ -35,7 +35,7 @@
 
 indexer_topk variant of the CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100).
 
-Derived from ``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
+Derived from ``attention/kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
 for the GLM-5.x indexer shape, where all verify rows of a request fit one UMMA
 tile (``N = next_n * num_heads <= 256``). The only functional change is the
 size of the per-thread register cache of epilogue weights
@@ -109,7 +109,7 @@
 from cutlass.cutlass_dsl import dsl_user_op
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 
-from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import griddepcontrol_launch_dependents
+from tensorrt_llm._torch.kernels.blackwell.utils import griddepcontrol_launch_dependents
 
 # CuTe DSL CUDA 13 validates rounding modes as string literals. The string
 # form is also accepted by older wrappers, so keep it version-independent.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py
index 32c115e2a7da..6727372ce230 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rows.py
@@ -15,7 +15,7 @@
 """
 indexer_topk variant of the CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100).
 
-Derived from ``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
+Derived from ``attention/kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``
 for the GLM-5.x indexer shape, where all verify rows of a request fit one UMMA
 tile (``N = next_n * num_heads <= 256``). The only functional change is the
 size of the per-thread register cache of epilogue weights
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py
index 4cbd26c8b892..c00c567c7010 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/indexer_topk/fp8_paged_mqa_logits_rowsplit.py
@@ -16,7 +16,7 @@
 
 Scores all ``next_n`` verify rows of a request against one stream of the
 paged FP8 indexer K cache, derived from the native CuTe DSL kernel
-``cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``.
+``attention/kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits``.
 
 Differences from the native kernel (same TMA/UMMA building blocks, same
 epilogue arithmetic per element):
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py
index 1ee494a105f2..a9becc975382 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/oproj_weight_staging/cute_oproj_stage.py
@@ -154,7 +154,7 @@ def _build(
     from cutlass.cutlass_dsl import T, dsl_user_op
     from cutlass.utils.smem_allocator import SmemAllocator
 
-    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import mma_sync
+    from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import mma_sync
 
     geo = plan(n_dim, k_dim, grid)
     rows_max = geo["rows_max"]

From 8c7260d44ddd18d62c1a6609a0bdaf25051210f6 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Wed, 21 Oct 2026 11:15:49 -0700
Subject: [PATCH 11/19] [None][test] modeling_v2 GLM-5.2: update the measured
 GSM8K gate times in the test-db comments

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml    | 2 +-
 .../test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml      | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index f500ac76f678..a41eff4652b5 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -28,7 +28,7 @@ l0_gb300_multi_gpus:
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[dep4_moetp4]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[tp4_ep4]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4]
-  # GLM-5.2 NVFP4 modeling_v2 target at tep4, GSM8K under modeling_v2="require". Measured 17 min on 4 GB300.
+  # GLM-5.2 NVFP4 modeling_v2 target at tep4, GSM8K under modeling_v2="require". Measured 25 min on 4 GB300.
   - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103Default::test_gsm8k_tep4 TIMEOUT (60)
   # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere
   # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
index df5108775837..6f4d36acb422 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml
@@ -22,6 +22,6 @@ l0_gb300_multi_nodes_node2_gpu8:
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[tp8_moetp2_ep4]
   - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[dep8]
   # GLM-5.2 NVFP4 modeling_v2 target at tp8 and dep8, GSM8K under modeling_v2="require".
-  # Measured on 2x4 GB300: tp8 (batch 1) 56 min, dep8 10 min.
+  # Measured on 2x4 GB300: tp8 (batch 1) 34 min, dep8 6 min.
   - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103Default::test_gsm8k_tp8 TIMEOUT (90)
   - accuracy/modeling_v2/test_glm_moe_dsa.py::TestModelingV2Glm52Nvfp4Sm103Default::test_gsm8k_dep8 TIMEOUT (60)

From 0bb9782fa865b35184cbf6ac40a62d219656c83a Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Wed, 21 Oct 2026 11:56:02 -0700
Subject: [PATCH 12/19] [None][chore] modeling_v2 GLM-5.2: rename a contract
 string bandit reads as a hardcoded password

bandit B105 matches "token" in an assigned name, so it flagged the
routed-MoE layout rule ROUTED_MOE_TOKEN_MAJOR_RULE. It is now
ROUTED_MOE_UNFINALIZED_LAYOUT_RULE; the text is unchanged.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../models/glm_moe_dsa/shared/kernels/contracts.py          | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py
index c90328e8a5b3..ed2bfbfe087e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/shared/kernels/contracts.py
@@ -103,7 +103,7 @@ class Contract:
 )
 
 # Physical form of the routed_moe contract's unfinalized tuple as routed_moe returns it.
-ROUTED_MOE_TOKEN_MAJOR_RULE = (
+ROUTED_MOE_UNFINALIZED_LAYOUT_RULE = (
     "A routed-MoE package may return the unfinalized tuple in token-major "
     "form: fc2_output[R*8,6144] BF16 whose row t*8+k holds token t's k-th "
     "expert output (row pitch 6144 elements), expert scales [R,8] BF16 with "
@@ -126,7 +126,7 @@ class Contract:
     "rank that disagrees on the mode, the buffers or the compile keeps every "
     "rank on the native seam. Modes: row = finalized rows [R,6144] BF16; "
     "identity = token-major raw rows [R*8,6144] BF16 and scales [R,8] "
-    "(ROUTED_MOE_TOKEN_MAJOR_RULE); fp32 = an FP32 accumulator [R,6144] that is "
+    "(ROUTED_MOE_UNFINALIZED_LAYOUT_RULE); fp32 = an FP32 accumulator [R,6144] that is "
     "zero when routed_moe's kernel starts, to which routed_moe adds the eight weighted rows "
     "of a token with atomic FP32 adds (order run-to-run nondeterministic), "
     "and whose rows the hop clears after reading them. routed_moe never reuses a "
@@ -217,7 +217,7 @@ class Contract:
             "or (fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx)",
             "or the post_moe_hop hop's routed-input buffer (routed_moe_into_hop)",
         ),
-        (ROUTED_MOE_TOKEN_MAJOR_RULE, POST_MOE_HOP_ROUTED_INPUT_RULE),
+        (ROUTED_MOE_UNFINALIZED_LAYOUT_RULE, POST_MOE_HOP_ROUTED_INPUT_RULE),
     ),
     "post_moe_hop": Contract(
         "post-MoE hop",

From e52b5e2ee29e1d04d53323077f5c2769f5e92f9d Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Wed, 21 Oct 2026 21:55:47 -0700
Subject: [PATCH 13/19] [TRTLLM-17096][feat] runtime hooks: model opt-in engine
 fast paths

Add RuntimeHooks, a per-model declaration of engine-side fast paths the
model provides itself: a step-input stager, a speculative sampler store
writer, a decode prologue, executor events, and running the overlap loop
on the execution stream. Every hook defaults to off, and a model without
a declaration runs the native path unchanged.

Also add two extension points a model can use: FMHA libraries a backend
puts ahead of the enabled ones, and the DSA indexer's reused top-k copy
as an overridable method.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_torch/attention/backends/fmha/manager.py |   5 +-
 .../attention/backends/sparse/dsa/indexer.py  | 136 ++-
 .../attention/backends/sparse/dsa/metadata.py |  98 +-
 .../_torch/attention/backends/trtllm.py       |  61 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   8 +-
 .../_torch/pyexecutor/cuda_graph_runner.py    |  14 +-
 .../pyexecutor/engine/runners/common.py       |   8 +
 .../engine/runners/decoder/graphs.py          |   7 +
 .../engine/runners/decoder/inputs.py          |  98 +-
 .../engine/runners/decoder/runner.py          | 154 ++-
 .../kv_cache/kv_cache_manager_v2.py           |  13 +
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  65 +-
 .../_torch/pyexecutor/runtime_hooks.py        | 221 +++++
 tensorrt_llm/_torch/speculative/eagle3.py     |   9 +-
 tensorrt_llm/_torch/speculative/interface.py  |  10 +-
 .../_torch/speculative/spec_sampler_base.py   |  82 +-
 .../attention/sparse/dsa/test_dsa_indexer.py  | 100 ++
 .../_torch/attention/test_fmha_manager.py     |  57 ++
 .../_torch/executor/test_gpu_keepalive.py     |   1 +
 .../_torch/executor/test_py_executor.py       |   3 +
 .../_torch/executor/test_runtime_hooks.py     | 924 ++++++++++++++++++
 21 files changed, 1942 insertions(+), 132 deletions(-)
 create mode 100644 tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
 create mode 100644 tests/unittest/_torch/executor/test_runtime_hooks.py

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py
index c9afff66e1d8..a4770ad3dbba 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py
@@ -308,7 +308,10 @@ def __init__(self, attn: TrtllmAttention) -> None:
         # normal forwarding never pays for an environment lookup.
         self._cache_sanity_check_enabled = _is_fmha_cache_sanity_check_enabled()
         self.fmha_libs: list[Fmha] = []
-        for fmha_cls in get_enabled_fmha_lib_classes():
+        for fmha_cls in (
+            *getattr(attn, "prepended_fmha_lib_classes", ()),
+            *get_enabled_fmha_lib_classes(),
+        ):
             if fmha_cls.is_available(attn):
                 self.fmha_libs.append(fmha_cls(attn))
 
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
index a4d3c25b7bce..9e1ff5b6f902 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
@@ -1240,6 +1240,18 @@ def recompute_slot_mappings(
         metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices
         metadata.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices
 
+        stager = getattr(metadata, "step_input_stager", None)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                metadata.slot_mapping_fp8[:total_new_kv_tokens],
+                metadata.host_slot_mapping_fp8[:total_new_kv_tokens],
+            )
+            stager.copy(
+                metadata.slot_mapping_scale[:total_new_kv_tokens],
+                metadata.host_slot_mapping_scale[:total_new_kv_tokens],
+            )
+            return
         metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_(
             metadata.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True
         )
@@ -1514,7 +1526,13 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None:
         if num_generations > 0:
             # Prepare schedule metadata for fp8_paged_mqa_logits
             # This is a preprocessing step that computes scheduling information for the kernel
-            Indexer.prepare_scheduler_metadata(metadata)
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager = getattr(metadata, "step_input_stager", None)
+            if stager is None:
+                Indexer.prepare_scheduler_metadata(metadata)
+            elif not stager.scheduler_metadata(metadata):
+                # Reads staged KV lengths: run right after the staged launch.
+                stager.defer(lambda: Indexer.prepare_scheduler_metadata(metadata))
 
     def _update_k_cache(
         self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata
@@ -1987,9 +2005,24 @@ def sparse_attn_indexer(
         )
 
         if has_decode and not metadata.skip_indexer_for_gen_reqs and reuse_topk:
-            topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = (
-                metadata.shared_topk_indices[:num_generations, :]
+            self._copy_reused_topk(
+                metadata,
+                topk_indices_buffer,
+                token_offset,
+                num_gen_tokens,
+                num_generations,
+                has_prefill,
+                gvr_prior_indices,
             )
+        elif (
+            has_decode
+            and not metadata.skip_indexer_for_gen_reqs
+            and self._fill_decode_topk(
+                metadata, q_fp8, weights, topk_indices_buffer, token_offset, num_gen_tokens
+            )
+        ):
+            # A subclass filled the generation rows.
+            pass
         elif has_decode and not metadata.skip_indexer_for_gen_reqs:
             # Get decode lengths per request (from seq_lens) for validation
             gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
@@ -2268,30 +2301,83 @@ def sparse_attn_indexer(
                 prev_topk_dst.copy_(last_mtp_topk)
 
         if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk:
-            rows = None
-            if has_decode:
-                next_n = num_gen_tokens // num_generations
-                gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
-                rows = self._mtp_last_accepted_rows(
-                    gen_topk, metadata, num_contexts, num_generations, next_n
-                )
-            if has_prefill:
-                ctx_last = (
-                    torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
-                )
-                ctx_rows = topk_indices_buffer[ctx_last]
-                rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
-            if rows is not None:
-                shared_topk_indices = metadata.shared_topk_indices
-                if shared_topk_indices is None:
-                    metadata.shared_topk_indices = rows.contiguous()
-                else:
-                    row_start = num_contexts if is_generation else 0
-                    shared_topk_indices[
-                        row_start : row_start + rows.shape[0], : rows.shape[1]
-                    ].copy_(rows)
+            self._store_mtp_shared_rows(
+                topk_indices_buffer,
+                metadata,
+                num_contexts,
+                num_generations,
+                num_gen_tokens,
+                token_offset,
+                has_prefill,
+                has_decode,
+                is_generation,
+            )
         return topk_indices_buffer
 
+    def _copy_reused_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+        num_generations: int,
+        has_prefill: bool,
+        gvr_prior_indices: Optional[torch.Tensor],
+    ) -> None:
+        """Copy the shared TopK rows an MTP draft step reuses into its generation rows."""
+        topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = (
+            metadata.shared_topk_indices[:num_generations, :]
+        )
+
+    def _fill_decode_topk(
+        self,
+        metadata: DSAtrtllmAttentionMetadata,
+        q_fp8: torch.Tensor,
+        weights: torch.Tensor,
+        topk_indices_buffer: torch.Tensor,
+        token_offset: int,
+        num_gen_tokens: int,
+    ) -> bool:
+        """Fill the generation TopK rows ahead of the decode branch; True if it filled them.
+
+        This class leaves them to the decode branch of ``sparse_attn_indexer``.
+        """
+        return False
+
+    def _store_mtp_shared_rows(
+        self,
+        topk_indices_buffer: torch.Tensor,
+        metadata: DSAtrtllmAttentionMetadata,
+        num_contexts: int,
+        num_generations: int,
+        num_gen_tokens: int,
+        token_offset: int,
+        has_prefill: bool,
+        has_decode: bool,
+        is_generation: Optional[bool],
+    ) -> None:
+        """Store each request's last accepted TopK row for the next MTP draft step to reuse."""
+        rows = None
+        if has_decode:
+            next_n = num_gen_tokens // num_generations
+            gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+            rows = self._mtp_last_accepted_rows(
+                gen_topk, metadata, num_contexts, num_generations, next_n
+            )
+        if has_prefill:
+            ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
+            ctx_rows = topk_indices_buffer[ctx_last]
+            rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
+        if rows is not None:
+            shared_topk_indices = metadata.shared_topk_indices
+            if shared_topk_indices is None:
+                metadata.shared_topk_indices = rows.contiguous()
+            else:
+                row_start = num_contexts if is_generation else 0
+                shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(
+                    rows
+                )
+
     def _mtp_last_accepted_rows(
         self,
         gen_topk: torch.Tensor,
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
index e0773630dcb3..5c9c4686375c 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
@@ -1486,6 +1486,15 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale
         num_blocks = page_indices.shape[1]
         self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices)
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.indexer_k_cache_block_offsets[: self.num_seqs],
+                self.host_indexer_k_cache_block_offsets[: self.num_seqs],
+                clamp_min=0,
+            )
+            self.step_input_stager.clamp_min(self.indexer_k_cache_block_offsets[self.num_seqs :], 0)
+            return pool_indices
         self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
             self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
         )
@@ -1702,9 +1711,16 @@ def prepare_dense_topk_indices(self, kv_lens, device=False):  # device=False mea
                     kv_lens[self.num_contexts : self.num_seqs],
                     self.num_tokens - self.num_ctx_tokens,
                 )
-                self.topk_indices_buffer[gen_range, :].copy_(
-                    self.host_topk_indices_buffer[gen_range, :], non_blocking=True
-                )
+                if self.step_input_stager is not None:
+                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                    self.step_input_stager.copy(
+                        self.topk_indices_buffer[gen_range, :],
+                        self.host_topk_indices_buffer[gen_range, :],
+                    )
+                else:
+                    self.topk_indices_buffer[gen_range, :].copy_(
+                        self.host_topk_indices_buffer[gen_range, :], non_blocking=True
+                    )
 
     def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
         # The DeepGEMM paged-MQA kernel (fp8_paged_mqa_logits) runs a native
@@ -1804,9 +1820,16 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 num_tokens = self.num_generations * expand_factor
                 gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(expand_factor)
                 self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded)
-                self.kv_lens_expanded_cuda[:num_tokens].copy_(
-                    self.kv_lens_expanded_host[:num_tokens], non_blocking=True
-                )
+                if self.step_input_stager is not None:
+                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                    self.step_input_stager.copy(
+                        self.kv_lens_expanded_cuda[:num_tokens],
+                        self.kv_lens_expanded_host[:num_tokens],
+                    )
+                else:
+                    self.kv_lens_expanded_cuda[:num_tokens].copy_(
+                        self.kv_lens_expanded_host[:num_tokens], non_blocking=True
+                    )
                 self._refresh_expanded_block_table(expand_factor)
         else:
             self.dsl_expand_factor = 1
@@ -2036,6 +2059,15 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None):
         self.host_block_table_expanded[:num_tokens, :max_len].copy_(
             expanded_blocks, non_blocking=True
         )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.block_table_expanded[:num_tokens],
+                self.host_block_table_expanded[:num_tokens],
+                clamp_min=0,
+            )
+            self.step_input_stager.clamp_min(self.block_table_expanded[num_tokens:], 0)
+            return
         self.block_table_expanded[:num_tokens].copy_(
             self.host_block_table_expanded[:num_tokens], non_blocking=True
         )
@@ -2076,9 +2108,15 @@ def prepare_for_indexer_k_cache(self):
             pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1
         )
         # Copy to GPU
-        self.block_table[: self.num_seqs, :max_blocks_used].copy_(
-            host_block_table, non_blocking=True
-        )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.block_table[: self.num_seqs, :max_blocks_used], host_block_table
+            )
+        else:
+            self.block_table[: self.num_seqs, :max_blocks_used].copy_(
+                host_block_table, non_blocking=True
+            )
 
     def prepare_for_skip_indexer(self, kv_lens: torch.Tensor):
         num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens
@@ -2144,9 +2182,17 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_cached_token_indptr[1 : self.num_generations + 1],
             )
-            self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
-                self.host_gen_cached_token_indptr[: self.num_generations + 1], non_blocking=True
-            )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self.gen_cached_token_indptr[: self.num_generations + 1],
+                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
+                )
+            else:
+                self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
+                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
+                    non_blocking=True,
+                )
             # generation kv indptr
             torch.cumsum(
                 kv_lens[self.num_contexts : self.num_seqs],
@@ -2154,9 +2200,16 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_kv_indptr[1 : self.num_generations + 1],
             )
-            self.gen_kv_indptr[: self.num_generations + 1].copy_(
-                self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
-            )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self.gen_kv_indptr[: self.num_generations + 1],
+                    self.host_gen_kv_indptr[: self.num_generations + 1],
+                )
+            else:
+                self.gen_kv_indptr[: self.num_generations + 1].copy_(
+                    self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
+                )
         else:
             self.max_gen_seq_len = 0
 
@@ -2168,7 +2221,14 @@ def prepare_for_indices_conversion(self):
             self.seq_lens,
             dim=0,
         )
-        self.req_idx_per_token[: self.num_tokens].copy_(
-            self.host_req_idx_per_token[: self.num_tokens],
-            non_blocking=True,
-        )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(
+                self.req_idx_per_token[: self.num_tokens],
+                self.host_req_idx_per_token[: self.num_tokens],
+            )
+        else:
+            self.req_idx_per_token[: self.num_tokens].copy_(
+                self.host_req_idx_per_token[: self.num_tokens],
+                non_blocking=True,
+            )
diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py
index 1d7f183e6e85..64b29917880b 100644
--- a/tensorrt_llm/_torch/attention/backends/trtllm.py
+++ b/tensorrt_llm/_torch/attention/backends/trtllm.py
@@ -28,6 +28,7 @@
     from tensorrt_llm.mapping import Mapping
 
     from ...model_config import ModelConfig
+    from ...pyexecutor.runtime_hooks import StepInputStager
     from ...speculative.interface import SpecMetadata
     from ...speculative.spec_tree_manager import SpecTreeManager
     from .sparse.dsa.metadata import TokenMajorGenView
@@ -123,6 +124,8 @@ def generate_spec_decoding_packed_mask(max_num_requests: int,
 
 @dataclass(kw_only=True)
 class TrtllmAttentionMetadata(AttentionMetadata):
+    # A step-input stager while it records this metadata's writes, else None.
+    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     workspace: Optional[torch.Tensor] = None
     cuda_graph_workspace: Optional[torch.Tensor] = None
     workspace_reclaimable: bool = field(default=True, init=False)
@@ -1006,8 +1009,14 @@ def prepare(self) -> None:
             device='cpu',
         )
         self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens)
-        self.prompt_lens_cuda[:self.num_seqs].copy_(
-            self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
+        stager = self.step_input_stager
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.prompt_lens_cuda[:self.num_seqs],
+                        self.prompt_lens_cpu[:self.num_seqs])
+        else:
+            self.prompt_lens_cuda[:self.num_seqs].copy_(
+                self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
 
         # number of tokens in the kv cache for each sequence in the batch
         cached_token_lens = torch.tensor(
@@ -1060,9 +1069,14 @@ def prepare(self) -> None:
         # the sequence length including the cached tokens and the input tokens.
         self.kv_lens[:self.num_seqs].copy_(
             kv_lens + self.kv_cache_params.num_extra_kv_tokens)
-        self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
-            kv_lens[:self.num_seqs]),
-                                                non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.kv_lens_cuda[:self.num_seqs],
+                        kv_lens[:self.num_seqs])
+        else:
+            self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
+                kv_lens[:self.num_seqs]),
+                                                    non_blocking=True)
         # total kv lens for context requests and generation requests, without extra tokens
         self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item()
         self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self.
@@ -1429,23 +1443,42 @@ def prepare_context_mla_with_cached_kv(self,
                      dtype=torch.int64,
                      out=self.host_ctx_cached_token_indptr[1:self.num_contexts +
                                                            1])
-        self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
-            non_blocking=True)
+        stager = self.step_input_stager
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.ctx_cached_token_indptr[:self.num_contexts + 1],
+                self.host_ctx_cached_token_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
+                non_blocking=True)
         torch.cumsum(
             self.seq_lens[:self.num_contexts],
             dim=0,
             dtype=torch.int64,
             out=self.host_ctx_uncached_token_indptr[1:self.num_contexts + 1])
-        self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
-            non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(
+                self.ctx_uncached_token_indptr[:self.num_contexts + 1],
+                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
+                non_blocking=True)
         torch.cumsum(kv_lens[:self.num_contexts],
                      dim=0,
                      dtype=torch.int64,
                      out=self.host_ctx_kv_indptr[1:self.num_contexts + 1])
-        self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
-            self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
+        if stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.copy(self.ctx_kv_indptr[:self.num_contexts + 1],
+                        self.host_ctx_kv_indptr[:self.num_contexts + 1])
+        else:
+            self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
+                self.host_ctx_kv_indptr[:self.num_contexts + 1],
+                non_blocking=True)
 
     def spec_decoding_param_prepare_for_blackwell(self) -> None:
         """
@@ -1654,6 +1687,8 @@ def is_sm_version_trtllm_gen_kernel(self, sm):
 class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]):
 
     Metadata = TrtllmAttentionMetadata
+    # FMHA library classes this instance tries before the registry's libraries.
+    prepended_fmha_lib_classes: tuple = ()
 
     @staticmethod
     def is_sm_version_trtllm_gen_kernel(sm):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 4d3b4af850e8..2090179caa6c 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -53,6 +53,7 @@
 from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens,
                            get_num_spec_layers, get_spec_decoder,
                            should_use_separate_draft_kv_cache)
+from ..speculative.spec_sampler_base import SpecSampler
 from ..utils import is_gdn_replay_enabled
 from . import config_utils
 from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer,
@@ -77,6 +78,7 @@
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
+from .runtime_hooks import runtime_hooks_of
 from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler
 from .sampler.penalties import PenaltyStore
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
@@ -4587,7 +4589,11 @@ def instantiate_sampler(
     )
     if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder(
     ):
-        return get_spec_decoder(sampler_args, engine.spec_config)
+        sampler = get_spec_decoder(sampler_args, engine.spec_config)
+        store_writer = runtime_hooks_of(engine.model).sampler_store_writer
+        if store_writer is not None and isinstance(sampler, SpecSampler):
+            sampler.store_writer = store_writer
+        return sampler
 
     if mm_encoder_only:
         # NOTE: handle model outputs specially for mm encoder executor/engine.
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index e0d5abd61b35..118499f61732 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -160,6 +160,10 @@ class CUDAGraphRunner:
     multiple graphs, keyed by batch shape and execution-path specializations.
     """
     WARMUP_STEPS = 1
+    # Tokens of input_ids / position_ids that a step-input stager already
+    # wrote into the static inputs for the next replay (0: the replay copies
+    # them).
+    staged_static_input_tokens = 0
 
     def __init__(self, config: CUDAGraphRunnerConfig):
         self.config = config
@@ -804,7 +808,12 @@ def replay(
                 f"replay() got {seqlen} tokens for key {key}, but the graph "
                 f"was captured for {expected_num_tokens} tokens. A shorter "
                 "input_ids leaves the tail of the static input buffer stale.")
-        static_tensors["input_ids"][:seqlen].copy_(input_ids)
+        # A step-input stager already wrote input_ids / position_ids[:seqlen]
+        # into the static buffers.
+        staged = self.staged_static_input_tokens == seqlen
+        self.staged_static_input_tokens = 0
+        if not staged:
+            static_tensors["input_ids"][:seqlen].copy_(input_ids)
 
         position_ids = current_inputs["position_ids"]
         if self.config.use_mrope:
@@ -851,7 +860,8 @@ def replay(
                     f"for key {key}, but expected {expected_position_ids_shape}. "
                     "torch.Tensor.copy_() silently broadcasts mismatched shapes, "
                     "which would corrupt the static input buffer.")
-            static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
+            if not staged:
+                static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
 
         num_encoder_tokens = key.num_encoder_tokens
         if num_encoder_tokens:
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
index 98dce6b41550..cb24f68deef7 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
@@ -22,6 +22,7 @@
 from .interface import ScheduledInputs
 
 if TYPE_CHECKING:
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm._torch.pyexecutor.sampler.sampler import SampleStateTensors
 
 
@@ -137,11 +138,13 @@ def pad_input_tokens(
     enable_attention_dp: bool,
     prefill_cuda_graph_backend: PrefillCudaGraphBackend,
     prefill_cuda_graph_num_tokens: list[int],
+    step_input_stager: "StepInputStager | None" = None,
 ) -> int:
     """Pad the batch to the token count its prefill graph or attention-DP peers need.
 
     Zeroes the padded rows of ``input_ids`` and ``position_ids``, whose last
     dimension is the token one, and returns the token count the model runs on.
+    With ``step_input_stager`` the zeroing is staged on it.
     """
     padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = get_padding_params(
         num_tokens,
@@ -158,6 +161,11 @@ def pad_input_tokens(
         attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens
     if attn_metadata.padded_num_tokens is None:
         return num_tokens
+    if step_input_stager is not None:
+        # Staged write; see the window banner in DecoderRunner._forward_decoder.
+        step_input_stager.zero(input_ids[num_tokens:padded_num_tokens])
+        step_input_stager.zero(position_ids[..., num_tokens:padded_num_tokens])
+        return padded_num_tokens
     input_ids[num_tokens:padded_num_tokens].fill_(0)
     position_ids[..., num_tokens:padded_num_tokens].fill_(0)
     return padded_num_tokens
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
index 4046fd5ca8ea..1dc3fca7e80f 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
@@ -487,11 +487,14 @@ def run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
+        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         """Run the model by replaying the graph for ``key``, capturing it first if needed.
 
         Without a key, record the model body while a breakable prefill graph
         captures, run through one when it covers the tokens, otherwise eagerly.
+        ``before_replay`` runs right before the replay of an already captured
+        graph, after the draft-cache metadata refresh.
         """
         with with_shared_pool(self.cuda_graph_runner.get_graph_pool()):
             if self.captures_model_body(key):
@@ -509,6 +512,7 @@ def run(
                 enable_spec_decode=enable_spec_decode,
                 gather_context_logits=gather_context_logits,
                 position_id_offsets=position_id_offsets,
+                before_replay=before_replay,
             )
 
     def _run(
@@ -524,6 +528,7 @@ def _run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
+        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         breakable_runner = self.breakable_cuda_graph_runner
         if key is None:
@@ -574,6 +579,8 @@ def capture_model(model_inputs: dict[str, Any]) -> Any:
                     attn_metadata, draft_kv_cache_manager
                 )
                 try:
+                    if before_replay is not None:
+                        before_replay()
                     with MoeLoadBalancerIterContext(moe_load_balancer):
                         outputs = self.cuda_graph_runner.replay(
                             key, inputs, position_id_offsets=position_id_offsets
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
index ac57b3ef45eb..0aa1d1c35c76 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
@@ -25,6 +25,7 @@
     ResourceManager,
     ResourceManagerType,
 )
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue, StepInputStager
 from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase
@@ -311,6 +312,11 @@ class DecoderInputPreparer:
     path overrides ``_new_extra_inputs`` and ``_prepare_inputs_fast_path``.
     """
 
+    # The model's decode prologue (RuntimeHooks.decode_prologue), set by the runner.
+    decode_prologue: DecodePrologue | None = None
+    # The model's step-input stager while it records this preparer's writes, else None.
+    step_input_stager: StepInputStager | None = None
+
     def __init__(
         self,
         *,
@@ -518,11 +524,7 @@ def pad_tokens(
                 mapping=self.mapping,
                 dist=self.dist,
             )
-        return pad_input_tokens(
-            attn_metadata,
-            num_tokens,
-            num_contexts,
-            attn_all_rank_num_tokens,
+        padding = dict(
             input_ids=self._buffers.input_ids_cuda,
             position_ids=self._position_ids_buffer(),
             dist=self.dist,
@@ -534,6 +536,12 @@ def pad_tokens(
             ),
             prefill_cuda_graph_num_tokens=self._config.prefill_cuda_graph_num_tokens,
         )
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            padding["step_input_stager"] = self.step_input_stager
+        return pad_input_tokens(
+            attn_metadata, num_tokens, num_contexts, attn_all_rank_num_tokens, **padding
+        )
 
     def model_inputs(
         self,
@@ -836,10 +844,18 @@ def prepare_inputs(
         self._copy_position_ids_to_device(host)
 
         if enable_spec_decode:
-            self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
-                torch.tensor(host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()),
-                non_blocking=True,
+            gather_ids = torch.tensor(
+                host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()
             )
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)], gather_ids
+                )
+            else:
+                self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
+                    gather_ids, non_blocking=True
+                )
 
         self._prepare_attn_metadata(
             host,
@@ -1705,24 +1721,52 @@ def previous_seq_slots_device() -> torch.Tensor:
         assert total_num_tokens <= self._config.max_num_tokens, (
             f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self._config.max_num_tokens})"  # noqa: E501
         )
+        stager = self.step_input_stager
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids, dtype=torch.int, pin_memory=prefer_pinned())
-            self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(self._buffers.input_ids_cuda[:num_tokens], input_ids)
+            else:
+                self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
 
         if num_draft_tokens > 0:
             draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned())
-            self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
-                draft_tokens, non_blocking=True
-            )
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(self._buffers.draft_tokens_cuda[: len(draft_tokens)], draft_tokens)
+            else:
+                self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
+                    draft_tokens, non_blocking=True
+                )
         if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0:
             num_accepted_draft_tokens = torch.tensor(
                 num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()
             )
-            self._buffers.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_(
-                num_accepted_draft_tokens, non_blocking=True
+            accepted = self._buffers.num_accepted_draft_tokens_cuda[
+                : len(num_accepted_draft_tokens)
+            ]
+            if stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                stager.copy(accepted, num_accepted_draft_tokens)
+            else:
+                accepted.copy_(num_accepted_draft_tokens, non_blocking=True)
+        if next_draft_tokens_device is not None and stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            stager.overlap_inputs(
+                self._buffers,
+                previous_batch_indices,
+                previous_pos_indices,
+                num_tokens,
+                num_draft_tokens,
+                self.get_runtime_tokens_per_gen_step(runtime_draft_len),
+                len(extend_requests) - len(extend_dummy_requests),
+                new_tokens_device,
+                next_draft_tokens_device,
+                new_tokens_lens_device,
             )
-        if next_draft_tokens_device is not None:
+        elif next_draft_tokens_device is not None:
             # Initialize these two values to zeros
             self._buffers.previous_pos_id_offsets_cuda.zero_()
             self._buffers.previous_kv_lens_offsets_cuda.zero_()
@@ -1822,6 +1866,10 @@ def previous_seq_slots_device() -> torch.Tensor:
         # cover the previous-batch rows; padding-request rows are covered by
         # neither, so without this the embedding would gather stale contents
         # of the uninitialized buffer, which may exceed the vocabulary size.
+        if stager is not None and host.lists.unwritten_input_id_rows:
+            # A native write to input_ids_cuda inside the window needs a flush
+            # first; see the window banner in DecoderRunner._forward_decoder.
+            stager.flush()
         _zero_fill_row_spans(self._buffers.input_ids_cuda, host.lists.unwritten_input_id_rows)
 
         host.num_tokens = num_tokens
@@ -1870,6 +1918,11 @@ def _copy_position_ids_to_device(self, host: HostInputs) -> None:
                 self._buffers.mrope_position_ids_cuda[:, :, start_idx:end_idx].copy_(
                     segment[:, :, : end_idx - start_idx], non_blocking=True
                 )
+        elif self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.position_ids(
+                self._buffers.position_ids_cuda[:total_num_tokens], host_position_ids
+            )
         else:
             self._buffers.position_ids_cuda[:total_num_tokens].copy_(
                 host_position_ids, non_blocking=True
@@ -2178,6 +2231,11 @@ def record_input_copy(self) -> None:
         self._prepare_inputs_event = torch.cuda.Event()
         self._prepare_inputs_event.record()
 
+    @property
+    def input_copy_event(self) -> torch.cuda.Event | None:
+        """The event ``record_input_copy`` recorded last."""
+        return self._prepare_inputs_event
+
     def wait_for_input_copy(self) -> None:
         if self._prepare_inputs_event is not None:
             self._prepare_inputs_event.synchronize()
@@ -2259,6 +2317,16 @@ def preprocess_inputs(
         Make some changes to the device inputs and avoid blocking the async data transfer
         """
         attn_meta = inputs.get("attn_metadata")
+        # A model's decode prologue may replace the device work below with one
+        # launch; it decides on the host before any device work and returns
+        # False to keep the code below.
+        prologue = self.decode_prologue
+        if prologue is not None and prologue(
+            self, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=runtime_draft_len
+        ):
+            if self.guided_decoder is not None:
+                self.guided_decoder.token_event.record()
+            return inputs
         # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step.
         if attn_meta is not None:
             attn_meta.on_update_kv_lens()
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 2480fecb6ce2..060f30b6b6a5 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -32,6 +32,11 @@
 )
 from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder
 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import (
+    DEFAULT_RUNTIME_HOOKS,
+    RuntimeHooks,
+    runtime_hooks_of,
+)
 from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer
@@ -123,6 +128,9 @@ def __getattr__(self, name: str) -> Any:
 class DecoderRunner(ScheduledModelRunner):
     """Run scheduled decoder generation with its own buffers, graphs and metadata."""
 
+    # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise.
+    runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS
+
     # Optional execution paths that a subclass can turn off.
     _context_warmups_supported = True
     _context_graph_promotion_supported = True
@@ -195,6 +203,8 @@ def __init__(
         self._dummy_requests = self._make_dummy_requests()
         self._init_decoder_state()
         self._inputs = self._make_input_preparer(mm_item_scheduler)
+        self.runtime_hooks = runtime_hooks_of(self.model)
+        self._inputs.decode_prologue = self.runtime_hooks.decode_prologue
         self._graphs = self._make_cuda_graphs()
         self._initialize_moe_workspace_reclaimer()
         self._warmup = DecoderWarmup(
@@ -612,34 +622,100 @@ def _forward_decoder(
                     execution_requests.generation_requests,
                 )
 
-            with self._cuda_graph_metadata_prepare_scope(
-                attn_metadata,
-                key,
-                execution_requests,
-                execution_promoted_context_ids,
-                enable_spec_decode=enable_spec_decode,
-            ):
-                prepared = self._inputs.prepare_inputs(
+            # ==== Step-input window opens here ====
+            # When the model declares a step-input stager and begin_if_eligible
+            # returns it, every device write that prepare_inputs (and the
+            # metadata and KV-cache-manager code it calls) routes through
+            # ``step_input_stager`` (call sites marked "Staged write") is only
+            # recorded; close_input_window below, right before the graph
+            # replay, issues them all. Until then, do not read or write those
+            # buffers with any other device op: a read sees stale bytes and a
+            # write is overwritten. Route new writes through the stager or
+            # call stager.flush() first. Exception: input_ids_cuda and
+            # position_ids_cuda are also staged into the graph's static
+            # inputs, which the replay then does not refresh from them; to
+            # change them here, call stager.flush() and then write them
+            # natively. Nothing checks this contract in a production run.
+            # Contract: StepInputStager in pyexecutor/runtime_hooks.py.
+            stager = None
+            stager_source = self.runtime_hooks.step_input_stager
+            if stager_source is not None:
+                graph_replay = (
+                    can_run_graph
+                    and not self.cuda_graph_runner.needs_capture(key)
+                    and not self.cuda_graph_runner.is_warmup_only
+                )
+                stager = stager_source.begin_if_eligible(
+                    self._inputs,
                     execution_requests,
-                    kv_cache_manager,
                     attn_metadata,
                     spec_metadata,
                     new_tensors_device,
-                    cache_indirection_buffer,
-                    resource_manager,
-                    can_run_graph,
-                    execution_promoted_context_ids,
-                    use_lora_graph=use_lora_graph,
-                    enable_spec_decode=enable_spec_decode,
-                    runtime_draft_len=runtime_draft_len,
+                    graph_replay=graph_replay,
+                    static_inputs=(
+                        self.cuda_graph_runner.shared_static_tensors if graph_replay else None
+                    ),
                     is_dummy=is_dummy,
-                    gather_context_logits=gather_context_logits,
+                    enable_spec_decode=enable_spec_decode,
+                    promoted=bool(execution_promoted_context_ids),
                 )
+            try:
+                # The window spans the scope: the scope is a no-op whenever
+                # speculative decoding is on, which every stager requires.
+                with self._cuda_graph_metadata_prepare_scope(
+                    attn_metadata,
+                    key,
+                    execution_requests,
+                    execution_promoted_context_ids,
+                    enable_spec_decode=enable_spec_decode,
+                ):
+                    prepared = self._inputs.prepare_inputs(
+                        execution_requests,
+                        kv_cache_manager,
+                        attn_metadata,
+                        spec_metadata,
+                        new_tensors_device,
+                        cache_indirection_buffer,
+                        resource_manager,
+                        can_run_graph,
+                        execution_promoted_context_ids,
+                        use_lora_graph=use_lora_graph,
+                        enable_spec_decode=enable_spec_decode,
+                        runtime_draft_len=runtime_draft_len,
+                        is_dummy=is_dummy,
+                        gather_context_logits=gather_context_logits,
+                    )
+            except BaseException:
+                if stager is not None:
+                    stager.discard()
+                raise
             inputs = prepared.kwargs
             runtime_draft_len = prepared.runtime_draft_len
             if execution_promoted_context_ids:
                 self._inputs.record_promoted_iter_states(scheduled_requests)
-            self._inputs.record_input_copy()
+            close_input_window = None
+            window_open = stager is not None
+            if stager is None:
+                self._inputs.record_input_copy()
+            else:
+
+                def close_input_window() -> None:
+                    nonlocal window_open
+                    window_open = False
+                    # ==== Step-input window closes here ====
+                    # end() issues every staged write; nothing between
+                    # begin_if_eligible above and this call may have touched
+                    # the staged buffers with other device ops. The input-copy
+                    # event must follow the launch (host page-table reads
+                    # happen at execution time). A nonzero
+                    # static_inputs_tokens() makes the replay skip its
+                    # input_ids / position_ids copy.
+                    stager.end(record_event=False)
+                    self._inputs.record_input_copy()
+                    stager.set_completion_event(self._inputs.input_copy_event)
+                    self.cuda_graph_runner.staged_static_input_tokens = (
+                        stager.static_inputs_tokens()
+                    )
 
             def run_model(model_inputs: dict[str, Any]) -> Any:
                 with MoeLoadBalancerIterContext(moe_load_balancer):
@@ -661,24 +737,30 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None:
 
             graph_key = key if can_run_graph else None
             captures_model_body = self._graphs.captures_model_body(graph_key)
-            outputs = self._graphs.run(
-                graph_key,
-                inputs,
-                run_model,
-                postprocess_inputs,
-                attn_metadata=attn_metadata,
-                draft_kv_cache_manager=draft_kv_cache_manager,
-                moe_load_balancer=moe_load_balancer,
-                enable_spec_decode=enable_spec_decode,
-                gather_context_logits=gather_context_logits,
-                position_id_offsets=(
-                    self._buffers.previous_pos_id_offsets_cuda
-                    if enable_spec_decode
-                    and not self._config.disable_overlap_scheduler
-                    and attn_metadata.kv_cache_manager is not None
-                    else None
-                ),
-            )
+            try:
+                outputs = self._graphs.run(
+                    graph_key,
+                    inputs,
+                    run_model,
+                    postprocess_inputs,
+                    attn_metadata=attn_metadata,
+                    draft_kv_cache_manager=draft_kv_cache_manager,
+                    moe_load_balancer=moe_load_balancer,
+                    enable_spec_decode=enable_spec_decode,
+                    gather_context_logits=gather_context_logits,
+                    position_id_offsets=(
+                        self._buffers.previous_pos_id_offsets_cuda
+                        if enable_spec_decode
+                        and not self._config.disable_overlap_scheduler
+                        and attn_metadata.kv_cache_manager is not None
+                        else None
+                    ),
+                    before_replay=close_input_window,
+                )
+            except BaseException:
+                if window_open:
+                    stager.discard()
+                raise
             if captures_model_body:
                 return outputs
 
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 07a5e93717aa..c825736a372c 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
@@ -134,6 +134,7 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS
@@ -1249,6 +1250,8 @@ def _spec_recompute_claim_limit(
 class KVCacheManagerV2(BaseResourceManager):
     # Sparse managers attach a metadata batch during initialization.
     sparse_metadata_batch: Batch | None = None
+    # A step-input stager while it records this manager's writes, else None.
+    step_input_stager: Optional["StepInputStager"] = None
     # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction.
     # Declared on the class so it is present even when an instance is built without running __init__.
     _cold_pool_group_membership_cache: Optional[tuple[tuple[int, frozenset[int]], ...]] = None
@@ -6754,6 +6757,16 @@ def copy_batch_block_offsets(
             )
             return
 
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.block_offsets(
+                self.host_kv_cache_block_offsets,
+                dst_tensor,
+                copy_idx.tolist(),
+                self._index_scale_ints,
+                self.kv_offset.tolist(),
+            )
+            return
         copy_batch_block_offsets_to_device(
             self.host_kv_cache_block_offsets,
             dst_tensor,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 9d326ab0c6e4..f06c73f28213 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -119,6 +119,7 @@
                             get_from_waiting_queue, merge_requests)
 from .resource_manager import (NoFreeSlotsError, ResourceManager,
                                ResourceManagerType, request_context)
+from .runtime_hooks import ExecutorEvents, runtime_hooks_of
 from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState,
                       SampleStateTensors)
 from .scheduler import (RequestScheduler, ScheduledRequests,
@@ -409,6 +410,16 @@ class PendingEncoderStep:
 
 
 class PyExecutor:
+    # True while the overlap loop runs with the execution stream current
+    # (_overlap_loop_stream_scope).
+    _loop_on_execution_stream = False
+    # Whether the model opted into running the overlap loop on the execution
+    # stream (RuntimeHooks.loop_on_execution_stream); read once at construction.
+    _execution_stream_loop_armed = False
+    # The model's executor events (RuntimeHooks.executor_events); read once at
+    # construction.
+    _executor_events: Optional[ExecutorEvents] = None
+
     # Minimum number of async micro batches for async PP execution.
     # This is a trade-off between memory usage and performance.
     # If the number of micro batches is too small, the executor will spend too much time in synchronization.
@@ -526,6 +537,9 @@ def __init__(
         # related modules
         self.resource_manager = resource_manager
         self.model_engine = model_engine
+        hooks = runtime_hooks_of(model_engine.model)
+        self._execution_stream_loop_armed = hooks.loop_on_execution_stream
+        self._executor_events = hooks.executor_events
         self._enable_adp_dummy_fixes = getattr(model_engine,
                                                "_enable_adp_dummy_fixes", False)
         self._enable_scheduler_aware_adp_dummy = getattr(
@@ -920,6 +934,10 @@ def __init__(
         # Ensure the default stream waits for execution_stream to complete
         # before subsequent operations.
         torch.cuda.current_stream().wait_stream(self.execution_stream)
+        if self._executor_events is not None:
+            # The model's actions after warmup and graph capture, on every rank
+            # (some are collective, so every rank reaches this point).
+            self._executor_events.after_warmup()
         self.encoder_launch_executor = (ThreadPoolExecutor(
             max_workers=1, thread_name_prefix="encoder-launch")
                                         if self.is_encoder_decoder else None)
@@ -3025,6 +3043,9 @@ def _executor_loop_cleanup(self):
                     f"Error waiting on PP send handles during cleanup: "
                     f"{traceback.format_exc()}")
 
+        if self._executor_events is not None and self._event_loop_completed:
+            self._executor_events.at_shutdown()
+
     def _pp_schedule_and_propagate(self, microbatch_id: int):
         """The first PP rank schedules the requests and propagates the result to all other PP ranks."""
 
@@ -5725,11 +5746,44 @@ def _wait_for_model_engine_input_copy(self):
         if wait_for_input_copy is not None:
             wait_for_input_copy()
 
+    @contextmanager
+    def _overlap_loop_stream_scope(self):
+        """Run the overlap loop with the execution stream current, when the model opted in.
+
+        _forward_step makes the execution stream wait for the current stream
+        before every forward and the current stream wait for the execution
+        stream after it. On the thread's default stream, which in steady decode
+        carries only these waits and timing events, step N+1's inputs wait for
+        decode graph N through two cross-stream event hops. With the execution
+        stream current the loop's stream work is queued in the same order on one
+        stream and _forward_step skips both waits.
+
+        Not taken with a two-model drafter, a guided decoder, a KV connector,
+        an encoder-decoder model or multimodal encoder item scheduling, whose
+        side streams synchronize with the current stream.
+        """
+        if not (self._execution_stream_loop_armed and self.drafter is None and
+                self.draft_model_engine is None and self.guided_decoder is None
+                and self.kv_connector_manager is None
+                and not self.is_encoder_decoder
+                and not self._mm_encoder_item_scheduling_enabled):
+            yield
+            return
+        with torch.cuda.stream(self.execution_stream):
+            self._loop_on_execution_stream = True
+            try:
+                yield
+            finally:
+                self._loop_on_execution_stream = False
+
     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))
-        with self._profiler() as profile_step, self.hang_detector:
+        if self._executor_events is not None:
+            self._executor_events.before_overlap_loop()
+        with self._overlap_loop_stream_scope(), self._profiler() as profile_step, \
+                self.hang_detector:
             iter_start_time = time.time()
             iter_stats = None
             target_inputs = None
@@ -8484,8 +8538,10 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
             cache_indirection_buffer = self.sampler.get_cache_indirection()
 
             # Run model forward on the execution stream for proper synchronization
-            # with KVCacheTransferManager's onboard/offload operations.
-            self.execution_stream.wait_stream(torch.cuda.current_stream())
+            # with KVCacheTransferManager's onboard/offload operations. With the
+            # loop on the execution stream both waits are program order.
+            if not self._loop_on_execution_stream:
+                self.execution_stream.wait_stream(torch.cuda.current_stream())
             self._attach_encoder_output_to_execution_stream(scheduled_requests)
             with torch.cuda.stream(self.execution_stream):
                 outputs = forward(scheduled_requests, self.resource_manager,
@@ -8499,7 +8555,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
 
             # Ensure the default stream waits for execution_stream to complete
             # before downstream operations use the outputs.
-            torch.cuda.current_stream().wait_stream(self.execution_stream)
+            if not self._loop_on_execution_stream:
+                torch.cuda.current_stream().wait_stream(self.execution_stream)
 
             self._kv_connector_wait_for_save()
             if route_capture is not None:
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
new file mode 100644
index 000000000000..24ffbccfbb8b
--- /dev/null
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Step-loop behaviour a model may opt into, read through ``runtime_hooks_of``.
+
+A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class
+attribute or a property). Every field defaults to off, and a model that declares
+nothing gets ``RuntimeHooks()``: the step loop, the samplers and the attention
+metadata then run exactly as they do without this module. Every engine call site
+of a hook runs its original statements when the hook is off.
+
+The hooks change *where* and *how* the step loop moves data, never what a
+forward computes. The protocols below state only the methods an engine call
+site needs; the objects that implement them are owned by the model.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Protocol, Sequence
+
+if TYPE_CHECKING:
+    import torch
+
+
+class StepInputStager(Protocol):
+    """Records a decode step's input preparation and issues it as one launch.
+
+    ``begin_if_eligible`` opens a deferred-write window on a step it supports
+    and returns the stager, or returns None and the engine prepares the step
+    natively. While the window is open, the engine's call sites route their
+    device writes to the staging methods instead of issuing them. ``end``
+    issues every staged write in staging order, then the deferred native
+    operations; ``discard`` drops them after a failed preparation.
+
+    While its window is open the stager sets itself as ``step_input_stager`` on
+    the decoder's input preparer, the attention metadata, the speculative
+    metadata and the KV cache managers, and resets that attribute to None at
+    ``end`` and ``discard``. Every call site reads the attribute and runs its
+    original statement when it is None.
+    """
+
+    def begin_if_eligible(
+        self,
+        preparer: Any,
+        requests: Any,
+        attn_metadata: Any,
+        spec_metadata: Any,
+        new_tensors_device: Any,
+        *,
+        graph_replay: bool,
+        static_inputs: Optional[Dict[str, "torch.Tensor"]],
+        is_dummy: bool,
+        enable_spec_decode: bool,
+        promoted: bool,
+    ) -> Optional["StepInputStager"]:
+        """Open the window for this step, or return None to keep the native preparation.
+
+        ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
+        records. ``static_inputs`` are the static input tensors of the CUDA
+        graph the step replays (None when it replays none).
+        """
+
+    def copy(
+        self, dst: "torch.Tensor", src: "torch.Tensor", clamp_min: Optional[int] = None
+    ) -> None:
+        """Stage ``dst.copy_(src, non_blocking=True)`` for a host ``src``."""
+
+    def zero(self, dst: "torch.Tensor") -> None:
+        """Stage ``dst.zero_()``."""
+
+    def clamp_min(self, dst: "torch.Tensor", lo: int) -> None:
+        """Stage ``dst.clamp_(min=lo)``."""
+
+    def block_offsets(
+        self,
+        host_table: "torch.Tensor",
+        dst: "torch.Tensor",
+        copy_index: Sequence[int],
+        index_scales: Sequence[int],
+        kv_offsets: Sequence[int],
+    ) -> None:
+        """Stage the KV cache manager's batch copy of block offsets to the device."""
+
+    def position_ids(self, dst: "torch.Tensor", host_position_ids: "torch.Tensor") -> None:
+        """Stage the position-id upload and the CUDA graph's static copy of it."""
+
+    def overlap_inputs(
+        self,
+        buffers: Any,
+        previous_batch_indices: Sequence[int],
+        previous_pos_indices: Sequence[int],
+        num_tokens: int,
+        num_draft_tokens: int,
+        tokens_per_gen_step: int,
+        num_extend_requests: int,
+        new_tokens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+    ) -> None:
+        """Stage the overlap scheduler's gathers of the previous step's outputs into ``buffers``."""
+
+    def scheduler_metadata(self, metadata: Any) -> bool:
+        """Stage the sparse-attention indexer's scheduler metadata.
+
+        Returns False, having staged nothing, when the caller must defer the
+        native computation instead.
+        """
+
+    def defer(self, fn: Callable[[], None]) -> None:
+        """Run ``fn`` (native device work that reads staged values) after the launch."""
+
+    def flush(self) -> None:
+        """Issue every write staged so far, before the window closes."""
+
+    def end(self, record_event: bool = True) -> None:
+        """Close the window and issue the staged preparation."""
+
+    def discard(self) -> None:
+        """Close the window after a failed preparation; the staged writes are dropped."""
+
+    def set_completion_event(self, event: "torch.cuda.Event") -> None:
+        """Event the caller recorded right after ``end(record_event=False)``."""
+
+    def static_inputs_tokens(self) -> int:
+        """Tokens of the CUDA graph's static inputs the last launch wrote (0: none)."""
+
+
+class SamplerStoreWriter(Protocol):
+    """Stores a speculative sampler step's outputs and copies the store to the host in one launch."""
+
+    def supported(
+        self,
+        store: Any,
+        new_tokens: "torch.Tensor",
+        next_new_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+    ) -> bool:
+        """Whether this step's output tensors fit the writer; False keeps the native stores."""
+
+    def store_and_copy_out(
+        self,
+        store: Any,
+        slots: Sequence[int],
+        new_tokens: "torch.Tensor",
+        next_new_tokens: "torch.Tensor",
+        new_tokens_lens: "torch.Tensor",
+        next_draft_tokens: "torch.Tensor",
+        host_new_tokens: "torch.Tensor",
+        host_new_tokens_lens: "torch.Tensor",
+        host_next_draft_tokens: "torch.Tensor",
+        record_event: bool = True,
+    ) -> None:
+        """Scatter the outputs into the store rows ``slots`` and copy the store to the host tensors."""
+
+    def set_completion_event(self, event: "torch.cuda.Event") -> None:
+        """Event the caller recorded right after ``store_and_copy_out(record_event=False)``."""
+
+
+class DecodePrologue(Protocol):
+    """Replaces the decode step's input prologue with one launch.
+
+    Called first in the step's input preprocessing, inside CUDA graph capture
+    as well. It decides on the host before any device work; returning False
+    keeps the native prologue, and returning True means every engine-visible
+    tensor the native prologue writes holds the value it would have written.
+    """
+
+    def __call__(
+        self, preparer: Any, inputs: Any, *, enable_spec_decode: bool, runtime_draft_len: int
+    ) -> bool:
+        """``preparer`` is the ``DecoderInputPreparer`` running the step's preprocessing."""
+
+
+class ExecutorEvents(Protocol):
+    """Points in the executor's life a model may act at."""
+
+    def after_warmup(self) -> None:
+        """After warmup and CUDA graph capture, before the first request is served."""
+
+    def before_overlap_loop(self) -> None:
+        """Right before the overlap executor loop starts, on every rank; must not raise."""
+
+    def at_shutdown(self) -> None:
+        """When the executor loop has completed, during cleanup; must not raise."""
+
+
+@dataclass(frozen=True)
+class RuntimeHooks:
+    """The step-loop hooks a model opts into. All off by default."""
+
+    # Run the overlap executor loop with the execution stream current, so the
+    # loop's own stream work and the forward are ordered on one stream.
+    loop_on_execution_stream: bool = False
+    # Decode input preparation recorded and issued as one launch (see StepInputStager).
+    step_input_stager: Optional[StepInputStager] = None
+    # The speculative sampler's store update and host copy in one launch (see SamplerStoreWriter).
+    sampler_store_writer: Optional[SamplerStoreWriter] = None
+    # The decode step's input prologue in one launch (see DecodePrologue).
+    decode_prologue: Optional[DecodePrologue] = None
+    # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents).
+    executor_events: Optional[ExecutorEvents] = None
+
+
+DEFAULT_RUNTIME_HOOKS = RuntimeHooks()
+
+
+def runtime_hooks_of(model: object) -> RuntimeHooks:
+    """The hooks ``model`` declares, or the all-off default when it declares none.
+
+    Declaring hooks is optional for every model class, so absence is the common
+    case and means the default.
+    """
+    hooks = getattr(model, "runtime_hooks", None)
+    if hooks is None:
+        return DEFAULT_RUNTIME_HOOKS
+    if not isinstance(hooks, RuntimeHooks):
+        raise TypeError(
+            f"{type(model).__name__}.runtime_hooks must be a RuntimeHooks, got {type(hooks).__name__}"
+        )
+    return hooks
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 15993e5c26ba..0328e29746cc 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -348,8 +348,13 @@ def prepare(self):
                                      dtype=torch.int,
                                      device='cpu',
                                      pin_memory=prefer_pinned())
-        self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
-                                                 non_blocking=True)
+        if self.step_input_stager is not None:
+            # Staged write; see the window banner in DecoderRunner._forward_decoder.
+            self.step_input_stager.copy(self.batch_indices_cuda[:num_seqs],
+                                        batch_indices)
+        else:
+            self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
+                                                     non_blocking=True)
         # Snapshot the real (unpadded) scheduled token count before the decrement
         # below. maybe_capture_hidden_states runs during the target forward and
         # must bound the capture by this count; self.num_tokens is about to be
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 8edccc3e3f8a..a8bdaf4d4589 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import TYPE_CHECKING, List, Optional, Protocol, Type
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Protocol, Type
 
 import torch
 from packaging.version import Version
@@ -37,6 +37,7 @@
 if TYPE_CHECKING:
     from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
     from ..pyexecutor.llm_request import LlmRequest
+    from ..pyexecutor.runtime_hooks import StepInputStager
 
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer
@@ -459,6 +460,8 @@ class SpecMetadata:
     """
     Metadata for speculative decoding.
     """
+    # A step-input stager while it records this metadata's writes, else None.
+    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     # The max number of requests in a single batch.
     max_num_requests: int
     # The number of draft layers. (Also the number of draft tokens for the linear tree.)
@@ -881,6 +884,11 @@ def _populate_request_rng_state(
                                        device='cuda')
 
         def _upload(dst: torch.Tensor, values: list[int]) -> None:
+            if self.step_input_stager is not None:
+                # Staged write; see the window banner in DecoderRunner._forward_decoder.
+                self.step_input_stager.copy(
+                    dst[:len(values)], torch.tensor(values, dtype=torch.int64))
+                return
             dst[:len(values)].copy_(torch.tensor(values,
                                                  dtype=torch.int64,
                                                  pin_memory=prefer_pinned()),
diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
index 5fb5b185dc0b..dc40cca3a7c5 100644
--- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py
+++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
@@ -27,9 +27,11 @@
 
 import torch
 
+from ..._utils import prefer_pinned
 from ...sampling_params import SamplingParams
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from ..pyexecutor.resource_manager import BaseResourceManager
+from ..pyexecutor.runtime_hooks import SamplerStoreWriter
 from ..pyexecutor.sampler import (
     DEFAULT_BEAM_IDX,
     AsyncWorkerMixin,
@@ -167,6 +169,10 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin):
 
     SampleState = SampleStateSpec
 
+    # The model's one-launch store writer, bound once by instantiate_sampler; None keeps
+    # the native stores and host copies.
+    store_writer: Optional[SamplerStoreWriter] = None
+
     def is_generation_model(self) -> bool:
         return True
 
@@ -442,8 +448,11 @@ def sample_async(
             for r in sampling_requests
         ]
 
-        slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
-        slots = slots.to(device="cuda", non_blocking=True)
+        store_writer = self.store_writer
+        slots = None
+        if store_writer is None:
+            slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
+            slots = slots.to(device="cuda", non_blocking=True)
 
         o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests]
         o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests]
@@ -514,11 +523,28 @@ def sample_async(
         elif o_next_new_tokens.shape[1] > next_new_tokens_width:
             o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width]
 
-        # Use index_copy_ for efficient copying (slots are unique)
-        self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
-        self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
-        self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
-        self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
+        # The writer stores the four outputs and copies the three store tensors
+        # only; a step that publishes verify_lens keeps the native path.
+        staged_store = (
+            store_writer is not None
+            and o_verify_lens is None
+            and prefer_pinned()
+            and not self._async_worker_active()
+            and store_writer.supported(
+                self.store, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens
+            )
+        )
+        if not staged_store:
+            if slots is None:
+                slots = torch.as_tensor(
+                    [r.py_seq_slot for r in sampling_requests], dtype=torch.long
+                )
+                slots = slots.to(device="cuda", non_blocking=True)
+            # Use index_copy_ for efficient copying (slots are unique)
+            self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
+            self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
+            self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
+            self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
 
         # Create sample state with async D2H copy
         device_tensors = SampleStateTensorsSpec(
@@ -527,13 +553,43 @@ def sample_async(
             next_draft_tokens=self.store.next_draft_tokens,
         )
 
-        host_tensors = SampleStateTensorsSpec(
-            new_tokens=self._copy_to_host(self.store.new_tokens),
-            new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
-            next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
-            verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None),
-        )
+        if staged_store:
+            # staged_store implies prefer_pinned(), so these buffers are pinned.
+            host_tensors = SampleStateTensorsSpec(
+                new_tokens=torch.empty_like(
+                    self.store.new_tokens, device="cpu", pin_memory=prefer_pinned()
+                ),
+                new_tokens_lens=torch.empty_like(
+                    self.store.new_tokens_lens, device="cpu", pin_memory=prefer_pinned()
+                ),
+                next_draft_tokens=torch.empty_like(
+                    self.store.next_draft_tokens, device="cpu", pin_memory=prefer_pinned()
+                ),
+            )
+            store_writer.store_and_copy_out(
+                self.store,
+                [r.py_seq_slot for r in sampling_requests],
+                o_new_tokens,
+                o_next_new_tokens,
+                o_new_tokens_lens,
+                o_next_draft_tokens,
+                host_tensors.new_tokens,
+                host_tensors.new_tokens_lens,
+                host_tensors.next_draft_tokens,
+                record_event=False,
+            )
+        else:
+            host_tensors = SampleStateTensorsSpec(
+                new_tokens=self._copy_to_host(self.store.new_tokens),
+                new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
+                next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
+                verify_lens=(
+                    self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None
+                ),
+            )
         sampler_event = self._record_sampler_event()
+        if staged_store:
+            store_writer.set_completion_event(sampler_event.cuda_event)
 
         # Add dummy draft tokens to context requests for KV cache preparation
         for request in finished_context_requests:
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index 1aad3d7d79c4..3fbef1289654 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -5327,3 +5327,103 @@ def _run_indexer():
     assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, (
         "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena"
     )
+
+
+# --- The subclass seams of Indexer.sparse_attn_indexer ----------------------------------------
+
+
+def _seam_indexer():
+    indexer = Indexer.__new__(Indexer)
+    torch.nn.Module.__init__(indexer)
+    return indexer
+
+
+def test_reused_topk_copy_is_the_row_copy():
+    gen = torch.Generator().manual_seed(0)
+    shared = torch.randint(0, 1000, (4, 16), generator=gen, dtype=torch.int32)
+    buffer = torch.full((12, 16), -1, dtype=torch.int32)
+    expected = buffer.clone()
+    expected[2:5, :] = shared[:3, :]
+
+    _seam_indexer()._copy_reused_topk(
+        SimpleNamespace(shared_topk_indices=shared), buffer, 2, 3, 3, False, None
+    )
+
+    assert torch.equal(buffer, expected)
+
+
+def test_fill_decode_topk_leaves_the_rows_to_the_decode_branch():
+    buffer = torch.full((6, 8), -1, dtype=torch.int32)
+
+    filled = _seam_indexer()._fill_decode_topk(SimpleNamespace(), None, None, buffer, 0, 6)
+
+    assert filled is False
+    assert bool((buffer == -1).all())
+
+
+def _stored_rows_inline(
+    indexer,
+    topk_indices_buffer,
+    metadata,
+    num_contexts,
+    num_generations,
+    num_gen_tokens,
+    token_offset,
+    has_prefill,
+    has_decode,
+    is_generation,
+):
+    """The statements ``sparse_attn_indexer`` ran in place of ``_store_mtp_shared_rows``."""
+    rows = None
+    if has_decode:
+        next_n = num_gen_tokens // num_generations
+        gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens]
+        rows = indexer._mtp_last_accepted_rows(
+            gen_topk, metadata, num_contexts, num_generations, next_n
+        )
+    if has_prefill:
+        ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1
+        ctx_rows = topk_indices_buffer[ctx_last]
+        rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows])
+    if rows is not None:
+        shared_topk_indices = metadata.shared_topk_indices
+        if shared_topk_indices is None:
+            metadata.shared_topk_indices = rows.contiguous()
+        else:
+            row_start = num_contexts if is_generation else 0
+            shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(rows)
+
+
+@pytest.mark.parametrize("has_prefill, has_decode", [(False, True), (True, True), (True, False)])
+@pytest.mark.parametrize("is_generation", [True, False])
+@pytest.mark.parametrize("preset", [False, True], ids=["unset", "preset"])
+@pytest.mark.parametrize("accepted", [None, (0, 0, 1, 4, 2)], ids=["no_counts", "counts"])
+def test_mtp_shared_rows_are_the_inline_statements(
+    has_prefill, has_decode, is_generation, preset, accepted
+):
+    num_contexts, num_generations, next_n = 2, 3, 4
+    seq_lens = torch.tensor([5, 3] + [next_n] * num_generations, dtype=torch.int32)
+    token_offset = int(seq_lens[:num_contexts].sum())
+    num_gen_tokens = num_generations * next_n
+    gen = torch.Generator().manual_seed(1)
+    buffer = torch.randint(
+        0, 1000, (token_offset + num_gen_tokens, 6), generator=gen, dtype=torch.int32
+    )
+
+    def metadata():
+        return SimpleNamespace(
+            seq_lens_cuda=seq_lens.clone(),
+            mtp_num_accepted=None if accepted is None else torch.tensor(accepted),
+            shared_topk_indices=(torch.full((8, 8), -7, dtype=torch.int32) if preset else None),
+        )
+
+    args = (num_contexts, num_generations, num_gen_tokens, token_offset)
+    flags = (has_prefill, has_decode, is_generation)
+    expected, got = metadata(), metadata()
+    indexer = _seam_indexer()
+    _stored_rows_inline(indexer, buffer.clone(), expected, *args, *flags)
+    seam_buffer = buffer.clone()
+    indexer._store_mtp_shared_rows(seam_buffer, got, *args, *flags)
+
+    assert torch.equal(seam_buffer, buffer)
+    assert torch.equal(got.shared_topk_indices, expected.shared_topk_indices)
diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py
index 7cfb59eee093..397bfd93a4b5 100644
--- a/tests/unittest/_torch/attention/test_fmha_manager.py
+++ b/tests/unittest/_torch/attention/test_fmha_manager.py
@@ -974,3 +974,60 @@ def __init__(self, attn: TrtllmAttention) -> None:
     assert isinstance(attn._fmha_manager.fmha_libs[0], _NewFmha)
     assert quant_states_during_construction == [True]
     assert events == [("support", "old", None)]
+
+
+def _listed_fmha(name: str, available: bool, checked: list) -> type:
+    """An FMHA library class that records its availability check."""
+
+    def is_available(cls, attn) -> bool:
+        checked.append((name, attn))
+        return available
+
+    return type(
+        name,
+        (),
+        {"is_available": classmethod(is_available), "__init__": lambda self, attn: None},
+    )
+
+
+@pytest.mark.parametrize("declared", [False, True], ids=["undeclared", "empty"])
+def test_manager_without_prepended_libraries_builds_the_registry_list(declared: bool) -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)]
+    attn = FakeAttention()
+    if declared:
+        attn.prepended_fmha_lib_classes = ()
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == registry
+    assert TrtllmAttention.prepended_fmha_lib_classes == ()
+
+
+def test_manager_tries_a_prepended_library_first() -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)]
+    prepended = _listed_fmha("Prepended", True, checked)
+    attn = FakeAttention()
+    attn.prepended_fmha_lib_classes = (prepended,)
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == [prepended, *registry]
+    assert [name for name, _ in checked] == ["Prepended", "First", "Second"]
+    assert all(checked_attn is attn for _, checked_attn in checked)
+
+
+def test_manager_skips_an_unavailable_prepended_library() -> None:
+    checked: list = []
+    registry = [_listed_fmha("First", True, checked)]
+    attn = FakeAttention()
+    attn.prepended_fmha_lib_classes = (_listed_fmha("Prepended", False, checked),)
+
+    with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry):
+        manager = FmhaManager(attn)
+
+    assert [type(lib) for lib in manager.fmha_libs] == registry
+    assert [name for name, _ in checked] == ["Prepended", "First"]
diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive.py b/tests/unittest/_torch/executor/test_gpu_keepalive.py
index 5c3318c00690..51d9b63b277e 100644
--- a/tests/unittest/_torch/executor/test_gpu_keepalive.py
+++ b/tests/unittest/_torch/executor/test_gpu_keepalive.py
@@ -315,6 +315,7 @@ class Stub:
     # notified first, then the keepalive is closed, then the PP handles are awaited.
     class CleanupStub:
         _executor_loop_cleanup = pe.PyExecutor._executor_loop_cleanup
+        _executor_events = None
 
         def __init__(self, keepalive):
             self.events = []
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 20c64eaad57d..613e6b7ddd4e 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -102,6 +102,7 @@ def test_forward_step_carries_context_logits_request_to_runner(request_flags, ex
         iter_counter=0,
         sampler=Mock(get_cache_indirection=Mock(return_value=cache_indirection)),
         execution_stream=Mock(),
+        _loop_on_execution_stream=False,
         _iter_adp_dummy_ctx_tokens=0,
         _iter_adp_dummy_gen_tokens=0,
         _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens,
@@ -2245,6 +2246,8 @@ class _CleanupStub:
     """Stub for _executor_loop_cleanup: records the order in which the
     shutdown notification and PP-handle wait happen."""
 
+    _executor_events = None
+
     def __init__(self, pp_handles_raise=False):
         self.response_lock = threading.Lock()
         self.response_cv = threading.Condition(self.response_lock)
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
new file mode 100644
index 000000000000..3993cfa7b42a
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -0,0 +1,924 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).
+
+The engine's call sites run their original statements when a hook is off: every
+object a hook attaches to declares it off at class level, every call a call
+site makes on a hook is a method of the hook's protocol, every read of a hook in
+the engine files is decided by its all-off value toward the original statements
+(checked on the source), and call sites of each hook run on stand-ins issue
+their original work with the hook off and hand it to the hook when it is on.
+"""
+
+import ast
+import contextlib
+import dataclasses
+import inspect
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+import tensorrt_llm
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import (
+    DEFAULT_RUNTIME_HOOKS,
+    ExecutorEvents,
+    RuntimeHooks,
+    SamplerStoreWriter,
+    StepInputStager,
+    runtime_hooks_of,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+
+class _PlainModel(torch.nn.Module):
+    """A model that declares no hooks, like every model outside the opt-in."""
+
+
+class _ClassHooksModel(torch.nn.Module):
+    runtime_hooks = RuntimeHooks(loop_on_execution_stream=True)
+
+
+class _Events:
+    def after_warmup(self) -> None:
+        pass
+
+    def before_overlap_loop(self) -> None:
+        pass
+
+    def at_shutdown(self) -> None:
+        pass
+
+
+@pytest.mark.parametrize("model", [_PlainModel(), object()], ids=["module", "object"])
+def test_a_model_without_hooks_gets_the_all_off_default(model):
+    hooks = runtime_hooks_of(model)
+
+    assert hooks is DEFAULT_RUNTIME_HOOKS
+    assert hooks == RuntimeHooks()
+    assert not hooks.loop_on_execution_stream
+    assert hooks.step_input_stager is None
+    assert hooks.sampler_store_writer is None
+    assert hooks.decode_prologue is None
+    assert hooks.executor_events is None
+
+
+def test_every_field_defaults_to_off():
+    for field in dataclasses.fields(RuntimeHooks):
+        assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) in (False, None), field.name
+
+
+def test_hooks_declared_as_none_mean_the_default():
+    model = _PlainModel()
+    model.runtime_hooks = None
+
+    assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS
+
+
+def test_class_attribute_hooks_are_returned_as_declared():
+    assert runtime_hooks_of(_ClassHooksModel()) is _ClassHooksModel.runtime_hooks
+
+
+def test_instance_attribute_hooks_are_returned_as_declared():
+    hooks = RuntimeHooks(executor_events=_Events())
+    model = _PlainModel()
+    model.runtime_hooks = hooks
+
+    assert runtime_hooks_of(model) is hooks
+
+
+def test_instance_attribute_overrides_class_attribute():
+    hooks = RuntimeHooks()
+    model = _ClassHooksModel()
+    model.runtime_hooks = hooks
+
+    assert runtime_hooks_of(model) is hooks
+
+
+@pytest.mark.parametrize(
+    "declared",
+    [True, {"loop_on_execution_stream": True}, "RuntimeHooks()"],
+    ids=["bool", "dict", "str"],
+)
+def test_hooks_of_the_wrong_type_raise(declared):
+    model = _PlainModel()
+    model.runtime_hooks = declared
+
+    with pytest.raises(TypeError, match=r"_PlainModel\.runtime_hooks must be a RuntimeHooks"):
+        runtime_hooks_of(model)
+
+
+def test_hooks_are_immutable_once_declared():
+    hooks = RuntimeHooks()
+
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        hooks.loop_on_execution_stream = True
+
+
+def test_engine_classes_declare_the_stager_off():
+    """Every object a hook attaches to declares the attribute off at class level."""
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+        DSAtrtllmAttentionMetadata,
+    )
+    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+    from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata
+    from tensorrt_llm._torch.speculative.interface import SpecMetadata
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    for cls in (
+        TrtllmAttentionMetadata,
+        DSAtrtllmAttentionMetadata,
+        SpecMetadata,
+        Eagle3OneModelSpecMetadata,
+    ):
+        assert cls.step_input_stager is None, cls
+        # A class attribute, not a dataclass field: the constructors do not take it.
+        assert "step_input_stager" not in {f.name for f in dataclasses.fields(cls)}, cls
+    assert DecoderInputPreparer.step_input_stager is None
+    assert DecoderInputPreparer.decode_prologue is None
+    assert DecoderRunner.runtime_hooks is DEFAULT_RUNTIME_HOOKS
+    assert KVCacheManagerV2.step_input_stager is None
+    assert CUDAGraphRunner.staged_static_input_tokens == 0
+    assert SpecSampler.store_writer is None
+    assert PyExecutor._execution_stream_loop_armed is False
+    assert PyExecutor._loop_on_execution_stream is False
+    assert PyExecutor._executor_events is None
+
+
+# Per hook protocol: the local names and attribute names a call site holds the hook in, and
+# the engine files with call sites.
+_HOOK_CONSUMERS = {
+    StepInputStager: (
+        ("stager", "stager_source", "step_input_stager"),
+        ("step_input_stager",),
+        (
+            "pyexecutor/engine/runners/decoder/runner.py",
+            "pyexecutor/engine/runners/decoder/inputs.py",
+            "pyexecutor/engine/runners/common.py",
+            "pyexecutor/kv_cache/kv_cache_manager_v2.py",
+            "attention/backends/trtllm.py",
+            "attention/backends/sparse/dsa/metadata.py",
+            "attention/backends/sparse/dsa/indexer.py",
+            "speculative/eagle3.py",
+            "speculative/interface.py",
+        ),
+    ),
+    SamplerStoreWriter: (
+        ("store_writer",),
+        ("store_writer",),
+        ("speculative/spec_sampler_base.py",),
+    ),
+    ExecutorEvents: ((), ("_executor_events",), ("pyexecutor/py_executor.py",)),
+}
+
+
+def _protocol_methods(protocol) -> set:
+    return {
+        name
+        for name, value in vars(protocol).items()
+        if inspect.isfunction(value) and not name.startswith("_")
+    }
+
+
+def _hook_calls(tree: ast.AST, local_names, attr_names):
+    """``(method, line)`` of every call on a hook held in one of the locals or attributes."""
+    for node in ast.walk(tree):
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
+            continue
+        receiver = node.func.value
+        if (isinstance(receiver, ast.Name) and receiver.id in local_names) or (
+            isinstance(receiver, ast.Attribute) and receiver.attr in attr_names
+        ):
+            yield node.func.attr, node.lineno
+
+
+@pytest.mark.parametrize("protocol", list(_HOOK_CONSUMERS), ids=lambda p: p.__name__)
+def test_every_hook_call_site_uses_its_protocol(protocol):
+    """A call site calls only its hook's protocol methods, and the protocol has no unused method."""
+    local_names, attr_names, consumers = _HOOK_CONSUMERS[protocol]
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    methods = _protocol_methods(protocol)
+    used, unknown = set(), []
+    for rel in consumers:
+        tree = ast.parse((root / rel).read_text())
+        calls = list(_hook_calls(tree, local_names, attr_names))
+        assert calls, f"{rel} has no {protocol.__name__} call site"
+        for method, line in calls:
+            used.add(method)
+            if method not in methods:
+                unknown.append(f"{rel}:{line}: {method}")
+    assert not unknown, f"calls outside the {protocol.__name__} protocol:\n" + "\n".join(unknown)
+    assert used == methods
+
+
+class _StoreWriterModel(torch.nn.Module):
+    def __init__(self, writer) -> None:
+        super().__init__()
+        self.runtime_hooks = RuntimeHooks(sampler_store_writer=writer)
+
+
+def _instantiate_sampler(monkeypatch, model, sampler):
+    """``instantiate_sampler`` for a speculative engine whose spec decoder is ``sampler``."""
+    from tensorrt_llm._torch.pyexecutor import _util
+
+    monkeypatch.setattr(_util, "confidential_compute_enabled", lambda: False)
+    monkeypatch.setattr(_util, "resolve_max_num_sequences", lambda *args, **kwargs: 8)
+    monkeypatch.setattr(_util, "create_torch_sampler_args", lambda **kwargs: "sampler args")
+    monkeypatch.setattr(_util, "get_spec_decoder", lambda args, spec_config: sampler)
+    engine = types.SimpleNamespace(
+        max_seq_len=64,
+        spec_config=types.SimpleNamespace(
+            spec_dec_mode=types.SimpleNamespace(has_spec_decoder=lambda: True)
+        ),
+        model=model,
+    )
+    llm_args = types.SimpleNamespace(
+        sampler_force_async_worker=False,
+        disable_overlap_scheduler=False,
+        enable_speculative_beam_history_d2h=False,
+    )
+    return _util.instantiate_sampler(
+        engine,
+        llm_args,
+        None,
+        max_batch_size=8,
+        max_beam_width=1,
+        mm_encoder_only=False,
+        speculative_config=None,
+    )
+
+
+def test_instantiate_sampler_binds_no_writer_without_hooks(monkeypatch):
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    sampler = SpecSampler.__new__(SpecSampler)
+
+    assert _instantiate_sampler(monkeypatch, _PlainModel(), sampler) is sampler
+    assert "store_writer" not in vars(sampler)
+    assert sampler.store_writer is None
+
+
+def test_instantiate_sampler_binds_the_declared_writer(monkeypatch):
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    writer = object()
+    sampler = SpecSampler.__new__(SpecSampler)
+
+    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(writer), sampler) is sampler
+    assert sampler.store_writer is writer
+
+
+def test_executor_events_protocol():
+    assert _protocol_methods(ExecutorEvents) == {
+        "after_warmup",
+        "before_overlap_loop",
+        "at_shutdown",
+    }
+    events = _Events()
+    assert RuntimeHooks(executor_events=events).executor_events is events
+
+
+class _StreamRecorder:
+    """Stands in for ``torch.cuda.stream``: records the streams it is entered with."""
+
+    def __init__(self) -> None:
+        self.entered = []
+
+    @contextlib.contextmanager
+    def __call__(self, stream):
+        self.entered.append(stream)
+        yield
+
+
+def _scope_executor(armed: bool, **overrides):
+    """A ``PyExecutor`` with only the attributes ``_overlap_loop_stream_scope`` reads."""
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor._execution_stream_loop_armed = armed
+    executor.execution_stream = object()
+    attrs = dict(
+        drafter=None,
+        draft_model_engine=None,
+        guided_decoder=None,
+        kv_connector_manager=None,
+        is_encoder_decoder=False,
+        _mm_encoder_item_scheduling_enabled=False,
+    )
+    attrs.update(overrides)
+    for name, value in attrs.items():
+        setattr(executor, name, value)
+    return executor
+
+
+def test_overlap_loop_scope_is_a_no_op_when_not_armed(monkeypatch):
+    executor = _scope_executor(armed=False)
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is False
+    assert recorder.entered == []
+    assert "_loop_on_execution_stream" not in vars(executor)
+
+
+def test_overlap_loop_scope_engages_when_armed(monkeypatch):
+    executor = _scope_executor(armed=True)
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is True
+    assert recorder.entered == [executor.execution_stream]
+    assert executor._loop_on_execution_stream is False
+
+
+def test_overlap_loop_scope_resets_the_flag_when_the_loop_raises(monkeypatch):
+    executor = _scope_executor(armed=True)
+    monkeypatch.setattr(torch.cuda, "stream", _StreamRecorder())
+
+    with pytest.raises(RuntimeError, match="loop failed"):
+        with executor._overlap_loop_stream_scope():
+            raise RuntimeError("loop failed")
+    assert executor._loop_on_execution_stream is False
+
+
+@pytest.mark.parametrize(
+    "name, value",
+    [
+        ("drafter", object()),
+        ("draft_model_engine", object()),
+        ("guided_decoder", object()),
+        ("kv_connector_manager", object()),
+        ("is_encoder_decoder", True),
+        ("_mm_encoder_item_scheduling_enabled", True),
+    ],
+)
+def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value):
+    executor = _scope_executor(armed=True, **{name: value})
+    recorder = _StreamRecorder()
+    monkeypatch.setattr(torch.cuda, "stream", recorder)
+
+    with executor._overlap_loop_stream_scope():
+        assert executor._loop_on_execution_stream is False
+    assert recorder.entered == []
+
+
+def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
+    sampler = types.SimpleNamespace()
+
+    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(object()), sampler) is sampler
+    assert "store_writer" not in vars(sampler)
+
+
+# ---------------------------------------------------------------------------------------------
+# With every hook off, every call site runs its original statements.
+#
+# The check below reads the engine files that consume a hook and evaluates each read of a hook
+# at its all-off value. An ``if`` that reads a hook must be decided by that value; the branch it
+# takes is checked like any other code and the other branch is not run, so it may use the hook
+# freely. Outside such tests, live code may only bind a hook value to a local (or to an
+# attribute that declares the hook) at its all-off value. Any other read is reported with its
+# file and line.
+# ---------------------------------------------------------------------------------------------
+
+_SEAM_CONSUMERS = (
+    "pyexecutor/engine/runners/decoder/runner.py",
+    "pyexecutor/engine/runners/decoder/inputs.py",
+    "pyexecutor/cuda_graph_runner.py",
+    "pyexecutor/kv_cache/kv_cache_manager_v2.py",
+    "pyexecutor/py_executor.py",
+    "pyexecutor/_util.py",
+    "attention/backends/trtllm.py",
+    "attention/backends/sparse/dsa/metadata.py",
+    "attention/backends/sparse/dsa/indexer.py",
+    "speculative/eagle3.py",
+    "speculative/interface.py",
+    "speculative/spec_sampler_base.py",
+)
+
+# Attributes that hold a hook, and their value when every hook is off.
+_SEAM_ATTRS = {
+    "step_input_stager": None,
+    "store_writer": None,
+    "_executor_events": None,
+    "_loop_on_execution_stream": False,
+    "_execution_stream_loop_armed": False,
+    "staged_static_input_tokens": 0,
+    "runtime_hooks": DEFAULT_RUNTIME_HOOKS,
+    "decode_prologue": None,
+}
+# Locals the call sites hold a hook (or a decision on one) in.
+_SEAM_LOCALS = ("stager", "stager_source", "store_writer", "prologue", "staged", "staged_store")
+_HOOK_FIELDS = {field.name for field in dataclasses.fields(RuntimeHooks)}
+
+
+class _Unknown:
+    """A value the all-off evaluation cannot decide."""
+
+
+_UNKNOWN = _Unknown()
+
+
+class _AllOffWalk:
+    """Walks one engine file with every hook at its all-off value."""
+
+    def __init__(self, rel: str) -> None:
+        self.rel = rel
+        self.problems = []
+        self.decided = 0
+
+    def report(self, node: ast.AST, what: str) -> None:
+        self.problems.append(f"{self.rel}:{node.lineno}: {what}")
+
+    # -- reads of a hook ------------------------------------------------------------------
+    @staticmethod
+    def _is_seam_call(node: ast.AST) -> bool:
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)):
+            return False
+        if node.func.id == "runtime_hooks_of":
+            return True
+        return (
+            node.func.id == "getattr"
+            and len(node.args) >= 2
+            and isinstance(node.args[1], ast.Constant)
+            and node.args[1].value == "step_input_stager"
+        )
+
+    def reads_seam(self, node: ast.AST, env: dict) -> bool:
+        for sub in ast.walk(node):
+            if isinstance(sub, ast.Attribute) and sub.attr in _SEAM_ATTRS:
+                return True
+            if isinstance(sub, ast.Name) and sub.id in env:
+                return True
+            if self._is_seam_call(sub):
+                return True
+        return False
+
+    def value(self, node: ast.AST, env: dict):
+        """The all-off value of ``node``, or ``_UNKNOWN``."""
+        if isinstance(node, ast.Constant):
+            return node.value
+        if isinstance(node, ast.Name):
+            return env.get(node.id, _UNKNOWN)
+        if isinstance(node, ast.Attribute):
+            receiver = self.value(node.value, env)
+            if receiver is DEFAULT_RUNTIME_HOOKS and node.attr in _HOOK_FIELDS:
+                return getattr(DEFAULT_RUNTIME_HOOKS, node.attr)
+            if node.attr in _SEAM_ATTRS:
+                return _SEAM_ATTRS[node.attr]
+            return _UNKNOWN
+        if self._is_seam_call(node):
+            return DEFAULT_RUNTIME_HOOKS if node.func.id == "runtime_hooks_of" else None
+        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
+            operand = self.value(node.operand, env)
+            return _UNKNOWN if operand is _UNKNOWN else not operand
+        if isinstance(node, ast.BoolOp):
+            values = [self.value(v, env) for v in node.values]
+            known = [v for v in values if v is not _UNKNOWN]
+            if isinstance(node.op, ast.And):
+                if any(not v for v in known):
+                    return False
+                return True if len(known) == len(values) else _UNKNOWN
+            if any(known):
+                return True
+            return False if len(known) == len(values) else _UNKNOWN
+        if isinstance(node, ast.Compare) and len(node.ops) == 1:
+            left, right = self.value(node.left, env), self.value(node.comparators[0], env)
+            op = node.ops[0]
+            if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN:
+                if isinstance(node.comparators[0], ast.Constant) and right is None:
+                    return (left is None) == isinstance(op, ast.Is)
+            if (
+                isinstance(op, ast.Eq)
+                and isinstance(node.left, ast.Attribute)
+                and node.left.attr == "staged_static_input_tokens"
+            ):
+                # No tokens were staged; a graph replay has at least one token.
+                return False
+        return _UNKNOWN
+
+    # -- statements -------------------------------------------------------------------------
+    def block(self, body, env: dict) -> bool:
+        """Walk live statements; True when the block always ends in return or raise."""
+        for stmt in body:
+            if self.stmt(stmt, env):
+                return True
+        return False
+
+    def bind(self, stmt, targets, value, env: dict) -> None:
+        reads = self.reads_seam(value, env)
+        names = [t.id for t in targets if isinstance(t, ast.Name)]
+        attrs = [t for t in targets if isinstance(t, ast.Attribute) and t.attr in _SEAM_ATTRS]
+        if not (reads or attrs or any(name in env or name in _SEAM_LOCALS for name in names)):
+            return
+        if len(names) + len(attrs) != len(targets):
+            self.report(stmt, "a hook value is bound to something other than a local")
+            return
+        off = self.value(value, env)
+        if off is _UNKNOWN:
+            if reads or attrs:
+                self.report(stmt, "a hook value with every hook off is not decided")
+            for name in names:
+                env.pop(name, None)
+            return
+        for name in names:
+            env[name] = off
+        for attr in attrs:
+            if off != _SEAM_ATTRS[attr.attr]:
+                self.report(stmt, f"{attr.attr} is set to a value other than its all-off value")
+
+    def stmt(self, stmt: ast.stmt, env: dict) -> bool:
+        if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            self.block(stmt.body, dict(env))
+            return False
+        if isinstance(stmt, ast.ClassDef):
+            self.block(stmt.body, {})
+            return False
+        if isinstance(stmt, (ast.If, ast.While)):
+            if not self.reads_seam(stmt.test, env):
+                body_ends = self.block(stmt.body, env)
+                else_ends = self.block(stmt.orelse, env)
+                return isinstance(stmt, ast.If) and body_ends and else_ends
+            taken = self.value(stmt.test, env)
+            if taken is _UNKNOWN:
+                self.report(stmt, "a test of a hook is not decided with every hook off")
+                return False
+            self.decided += 1
+            if isinstance(stmt, ast.While) and taken:
+                self.report(stmt, "a loop on a hook runs with every hook off")
+                return False
+            return self.block(stmt.body if taken else stmt.orelse, env)
+        if isinstance(stmt, ast.Assign):
+            self.bind(stmt, stmt.targets, stmt.value, env)
+            return False
+        if isinstance(stmt, ast.AnnAssign):
+            if stmt.value is not None:
+                self.bind(stmt, [stmt.target], stmt.value, env)
+            return False
+        if isinstance(stmt, (ast.With, ast.AsyncWith)):
+            for item in stmt.items:
+                if self.reads_seam(item.context_expr, env):
+                    self.report(stmt, "a hook is entered as a context")
+            return self.block(stmt.body, env)
+        if isinstance(stmt, ast.Try):
+            ends = self.block(stmt.body, env)
+            for handler in stmt.handlers:
+                self.block(handler.body, env)
+            self.block(stmt.orelse, env)
+            return self.block(stmt.finalbody, env) or (ends and not stmt.handlers)
+        if isinstance(stmt, (ast.For, ast.AsyncFor)):
+            if self.reads_seam(stmt.iter, env):
+                self.report(stmt, "a loop iterates over a hook")
+            self.block(stmt.body, env)
+            self.block(stmt.orelse, env)
+            return False
+        if self.reads_seam(stmt, env):
+            self.report(stmt, "a hook is used outside a test decided with every hook off")
+        return isinstance(stmt, (ast.Return, ast.Raise))
+
+
+@pytest.mark.parametrize("rel", _SEAM_CONSUMERS)
+def test_every_seam_read_has_a_native_branch(rel):
+    """With every hook off, each call site takes the branch with its original statements."""
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    walk = _AllOffWalk(rel)
+    walk.block(ast.parse((root / rel).read_text()).body, {})
+    assert not walk.problems, "\n".join(walk.problems)
+    assert walk.decided, f"{rel} has no call site of a hook"
+
+
+def test_the_seam_check_reports_undecided_reads():
+    source = """
+def step(self, x):
+    stager = self.step_input_stager
+    if stager is not None:
+        stager.copy(x)
+    else:
+        x.copy_()
+    stager.flush()
+    if x or stager:
+        pass
+    self.step_input_stager = x
+"""
+    walk = _AllOffWalk("sample.py")
+    walk.block(ast.parse(source).body, {})
+    assert walk.decided == 1
+    assert walk.problems == [
+        "sample.py:8: a hook is used outside a test decided with every hook off",
+        "sample.py:9: a test of a hook is not decided with every hook off",
+        "sample.py:11: a hook value with every hook off is not decided",
+    ]
+
+
+def test_the_decode_prologue_call_site_uses_its_protocol():
+    """The input preparer calls the prologue with the protocol's positional and keyword parameters."""
+    from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue
+
+    signature = inspect.signature(DecodePrologue.__call__)
+    positional = [
+        p.name
+        for p in signature.parameters.values()
+        if p.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD and p.name != "self"
+    ]
+    keywords = {
+        p.name for p in signature.parameters.values() if p.kind is inspect.Parameter.KEYWORD_ONLY
+    }
+    root = Path(tensorrt_llm.__file__).resolve().parent / "_torch"
+    tree = ast.parse((root / "pyexecutor/engine/runners/decoder/inputs.py").read_text())
+    calls = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == "prologue"
+    ]
+    assert len(calls) == 1
+    assert len(calls[0].args) == len(positional)
+    assert {keyword.arg for keyword in calls[0].keywords} == keywords
+
+
+# -- the decode prologue ------------------------------------------------------------------------
+
+
+class _AttnMetadata:
+    def __init__(self) -> None:
+        self.kv_lens_updates = 0
+
+    def on_update_kv_lens(self) -> None:
+        self.kv_lens_updates += 1
+
+
+def _prologue_preparer(prologue):
+    """A ``DecoderInputPreparer`` with only what ``preprocess_inputs`` reads without spec decoding."""
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer
+
+    preparer = DecoderInputPreparer.__new__(DecoderInputPreparer)
+    if prologue is not None:
+        preparer.decode_prologue = prologue
+    preparer.guided_decoder = None
+    preparer._config = types.SimpleNamespace(disable_overlap_scheduler=True)
+    preparer.mapping = types.SimpleNamespace(has_cp_helix=lambda: False)
+    return preparer
+
+
+def test_preprocess_inputs_runs_the_native_prologue_without_a_hook():
+    preparer = _prologue_preparer(None)
+    inputs = {"attn_metadata": _AttnMetadata()}
+
+    assert preparer.decode_prologue is None
+    assert preparer.preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=0) is (
+        inputs
+    )
+    assert inputs["attn_metadata"].kv_lens_updates == 1
+
+
+@pytest.mark.parametrize("engaged", [False, True], ids=["declined", "engaged"])
+def test_preprocess_inputs_keeps_the_native_prologue_when_the_hook_declines(engaged):
+    calls = []
+
+    def prologue(preparer, inputs, *, enable_spec_decode, runtime_draft_len):
+        calls.append((preparer, inputs, enable_spec_decode, runtime_draft_len))
+        return engaged
+
+    preparer = _prologue_preparer(prologue)
+    inputs = {"attn_metadata": _AttnMetadata()}
+
+    assert preparer.preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=3) is (
+        inputs
+    )
+    assert calls == [(preparer, inputs, False, 3)]
+    assert inputs["attn_metadata"].kv_lens_updates == (0 if engaged else 1)
+
+
+# -- the execution-stream loop --------------------------------------------------------------------
+
+
+class _Stream:
+    def __init__(self, name: str, log: list) -> None:
+        self.name, self.log = name, log
+
+    def wait_stream(self, other: "_Stream") -> None:
+        self.log.append(f"{self.name} waits for {other.name}")
+
+
+def _forward_step_executor(monkeypatch, log: list):
+    """A ``PyExecutor`` with only what ``_forward_step`` reads, logging its stream work."""
+    from tensorrt_llm._torch.pyexecutor import py_executor
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    current = _Stream("current", log)
+
+    @contextlib.contextmanager
+    def stream(s):
+        log.append(f"enter {s.name}")
+        yield
+        log.append(f"leave {s.name}")
+
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda: current)
+    monkeypatch.setattr(torch.cuda, "stream", stream)
+    monkeypatch.setattr(py_executor, "nvtx_range", lambda message: (lambda fn: fn))
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor.iter_counter = 0
+    executor.execution_stream = _Stream("execution", log)
+    executor.resource_manager = object()
+    executor.sampler = types.SimpleNamespace(get_cache_indirection=lambda: None)
+    executor.model_engine = types.SimpleNamespace(
+        forward=lambda *args, **kwargs: log.append("forward") or "outputs"
+    )
+    executor._iter_adp_dummy_ctx_tokens = 0
+    executor._iter_adp_dummy_gen_tokens = 0
+    executor._compute_adp_dummy_tokens = lambda requests: (0, 0)
+    executor._maybe_record_hang_diagnostic_phase = lambda *args, **kwargs: None
+    executor._attach_encoder_output_to_execution_stream = lambda requests: None
+    executor._mark_cross_kv_projection_consumed = lambda requests: None
+    executor._kv_connector_wait_for_save = lambda: None
+    return executor
+
+
+_NO_REQUESTS = types.SimpleNamespace(
+    context_requests=[], num_context_requests=0, num_generation_requests=0
+)
+
+
+def test_forward_step_orders_the_streams_off_the_execution_stream_loop(monkeypatch):
+    log = []
+    executor = _forward_step_executor(monkeypatch, log)
+
+    assert executor._forward_step(_NO_REQUESTS) == "outputs"
+    assert log == [
+        "execution waits for current",
+        "enter execution",
+        "forward",
+        "leave execution",
+        "current waits for execution",
+    ]
+
+
+def test_forward_step_skips_the_waits_on_the_execution_stream_loop(monkeypatch):
+    log = []
+    executor = _forward_step_executor(monkeypatch, log)
+    executor._loop_on_execution_stream = True
+
+    assert executor._forward_step(_NO_REQUESTS) == "outputs"
+    assert log == ["enter execution", "forward", "leave execution"]
+
+
+# -- executor events --------------------------------------------------------------------------
+
+
+def _cleanup_executor(events, completed: bool):
+    """A ``PyExecutor`` with only what ``_executor_loop_cleanup`` reads, without pipeline stages."""
+    import threading
+
+    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+    executor = PyExecutor.__new__(PyExecutor)
+    executor.response_cv = threading.Condition()
+    executor.shutdown_event = threading.Event()
+    executor.num_micro_batches = 0
+    executor._event_loop_completed = completed
+    if events is not None:
+        executor._executor_events = events
+    return executor
+
+
+class _RecordingEvents(_Events):
+    def __init__(self) -> None:
+        self.calls = []
+
+    def at_shutdown(self) -> None:
+        self.calls.append("at_shutdown")
+
+
+@pytest.mark.parametrize("completed", [False, True])
+def test_executor_loop_cleanup_without_events(completed):
+    executor = _cleanup_executor(None, completed)
+
+    executor._executor_loop_cleanup()
+    assert executor.is_shutdown and executor.shutdown_event.is_set()
+
+
+@pytest.mark.parametrize("completed", [False, True])
+def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(completed):
+    events = _RecordingEvents()
+    executor = _cleanup_executor(events, completed)
+
+    executor._executor_loop_cleanup()
+    assert executor.is_shutdown
+    assert events.calls == (["at_shutdown"] if completed else [])
+
+
+# -- the step-input stager ------------------------------------------------------------------
+
+
+class _RecordingStager:
+    def __init__(self) -> None:
+        self.calls = []
+
+    def __getattr__(self, name):
+        return lambda *args, **kwargs: self.calls.append((name, args, kwargs))
+
+
+def _block_offsets_manager(monkeypatch, log: list):
+    """A ``KVCacheManagerV2`` with only what ``copy_batch_block_offsets`` reads (one page table)."""
+    from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2
+    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+
+    monkeypatch.setattr(
+        kv_cache_manager_v2,
+        "copy_batch_block_offsets_to_device",
+        lambda *args: log.append(("native", args)),
+    )
+    manager = KVCacheManagerV2.__new__(KVCacheManagerV2)
+    manager._publish_sparse_metadata = lambda: None
+    manager.max_copy_beam_width = 1
+    manager.kv_cache_type = None
+    manager._use_per_layer_page_tables = False
+    manager.index_mapper = types.SimpleNamespace(
+        get_copy_index=lambda request_ids, num_contexts, beam_width, cross: torch.tensor([1, 0])
+    )
+    manager.host_kv_cache_block_offsets = torch.zeros(2, 4, dtype=torch.int32)
+    manager.index_scales = torch.ones(1, dtype=torch.int32)
+    manager._index_scale_ints = [1]
+    manager.kv_offset = torch.zeros(1, dtype=torch.int32)
+    manager._stream = types.SimpleNamespace(cuda_stream=7)
+    return manager
+
+
+def test_kv_cache_block_offsets_copy_natively_without_a_stager(monkeypatch):
+    log = []
+    manager = _block_offsets_manager(monkeypatch, log)
+    dst = torch.empty(2, 4, dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
+    assert len(log) == 1
+    kind, args = log[0]
+    assert kind == "native"
+    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
+    assert args[2].tolist() == [1, 0]
+    assert args[3] is manager.index_scales and args[4] is manager.kv_offset and args[5] == 7
+
+
+def test_kv_cache_block_offsets_are_staged_with_a_stager(monkeypatch):
+    log = []
+    manager = _block_offsets_manager(monkeypatch, log)
+    manager.step_input_stager = _RecordingStager()
+    dst = torch.empty(2, 4, dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
+    assert log == []
+    ((name, args, _),) = manager.step_input_stager.calls
+    assert name == "block_offsets"
+    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
+    assert args[2:] == ([1, 0], [1], [0])
+
+
+def _slot_mapping_inputs(tokens: int):
+    metadata = types.SimpleNamespace(
+        host_indexer_k_cache_block_offsets=torch.tensor([[0, 1]], dtype=torch.int32),
+        host_slot_mapping_fp8=torch.zeros(tokens, dtype=torch.int64),
+        host_slot_mapping_scale=torch.zeros(tokens, dtype=torch.int64),
+        slot_mapping_fp8=torch.full((tokens,), -1, dtype=torch.int64),
+        slot_mapping_scale=torch.full((tokens,), -1, dtype=torch.int64),
+    )
+    params = types.SimpleNamespace(
+        batch_size=1,
+        tokens_per_block=4,
+        head_dim=8,
+        new_kv_tokens=torch.tensor([tokens]),
+        data_bytes_per_token=8,
+        cached_kv_tokens=torch.tensor([1]),
+        quant_block_size=8,
+    )
+    return metadata, params
+
+
+def test_indexer_slot_mappings_copy_natively_without_a_stager():
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+    metadata, params = _slot_mapping_inputs(3)
+
+    Indexer.recompute_slot_mappings(metadata, params)
+    assert torch.equal(metadata.slot_mapping_fp8, metadata.host_slot_mapping_fp8)
+    assert torch.equal(metadata.slot_mapping_scale, metadata.host_slot_mapping_scale)
+
+
+def test_indexer_slot_mappings_are_staged_with_a_stager():
+    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
+
+    metadata, params = _slot_mapping_inputs(3)
+    metadata.step_input_stager = _RecordingStager()
+
+    Indexer.recompute_slot_mappings(metadata, params)
+    assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
+    assert bool((metadata.slot_mapping_fp8 == -1).all())
+    assert bool((metadata.slot_mapping_scale == -1).all())

From 37013e6a17fa932c75ba1c2bfd627031a6b4530b Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Thu, 22 Oct 2026 09:34:08 -0700
Subject: [PATCH 14/19] [TRTLLM-17096][fix] runtime hooks: run
 ExecutorEvents.after_warmup inside the distributed warmup guard

after_warmup may run collectives, so a rank-local failure there must arm
the rank-crash watchdog like a warmup failure does.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index f06c73f28213..2bcac7d79198 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -936,8 +936,10 @@ def __init__(
         torch.cuda.current_stream().wait_stream(self.execution_stream)
         if self._executor_events is not None:
             # The model's actions after warmup and graph capture, on every rank
-            # (some are collective, so every rank reaches this point).
-            self._executor_events.after_warmup()
+            # (some are collective, so every rank reaches this point, and a
+            # rank-local failure arms the rank-crash watchdog).
+            with _distributed_warmup_guard(self.dist, self.dist.mapping):
+                self._executor_events.after_warmup()
         self.encoder_launch_executor = (ThreadPoolExecutor(
             max_workers=1, thread_name_prefix="encoder-launch")
                                         if self.is_encoder_decoder else None)

From 005b339a596b8df723bdf0ee8b3c508f9d2e5da1 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Thu, 22 Oct 2026 09:34:08 -0700
Subject: [PATCH 15/19] [TRTLLM-17096][fix] runtime hooks: run
 ExecutorEvents.after_warmup inside the distributed warmup guard

after_warmup may run collectives, so a rank-local failure there must arm
the rank-crash watchdog like a warmup failure does.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index f06c73f28213..2bcac7d79198 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -936,8 +936,10 @@ def __init__(
         torch.cuda.current_stream().wait_stream(self.execution_stream)
         if self._executor_events is not None:
             # The model's actions after warmup and graph capture, on every rank
-            # (some are collective, so every rank reaches this point).
-            self._executor_events.after_warmup()
+            # (some are collective, so every rank reaches this point, and a
+            # rank-local failure arms the rank-crash watchdog).
+            with _distributed_warmup_guard(self.dist, self.dist.mapping):
+                self._executor_events.after_warmup()
         self.encoder_launch_executor = (ThreadPoolExecutor(
             max_workers=1, thread_name_prefix="encoder-launch")
                                         if self.is_encoder_decoder else None)

From 247c83ec9c05cfc6bc4bf6b72f54b421f3412111 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Thu, 22 Oct 2026 11:51:34 -0700
Subject: [PATCH 16/19] [TRTLLM-17096][fix] runtime hooks: stage step inputs
 only on the replay of a captured graph

The step-input window closes right before the replay of an already captured
CUDA graph. An eager step, an initial capture or a warmup would run the model
before the staged writes are issued, so the decoder runner now offers the
window only on such a replay, and StepInputStager documents the restriction.

Add behavioral tests with reference hooks: the runner window lifecycle
(staged writes issued before the completion event and the replay, discard on
a failure while the window is open), steps that are not a replay never open a
window, before_replay runs only before the replay of a captured graph, and
the speculative sampler's store writer against the native stores on
nonconsecutive slots, with its native fallbacks.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../engine/runners/decoder/runner.py          |  31 +-
 .../_torch/pyexecutor/runtime_hooks.py        |   8 +-
 .../_torch/executor/test_runtime_hooks.py     | 474 ++++++++++++++++++
 3 files changed, 498 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 060f30b6b6a5..e31f20da8e40 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -640,25 +640,28 @@ def _forward_decoder(
             stager = None
             stager_source = self.runtime_hooks.step_input_stager
             if stager_source is not None:
+                # Only the replay of an already captured graph closes the window
+                # (before_replay); an eager step, an initial capture or a warmup
+                # would run the model before the staged writes are issued, so
+                # those steps always prepare natively.
                 graph_replay = (
                     can_run_graph
                     and not self.cuda_graph_runner.needs_capture(key)
                     and not self.cuda_graph_runner.is_warmup_only
                 )
-                stager = stager_source.begin_if_eligible(
-                    self._inputs,
-                    execution_requests,
-                    attn_metadata,
-                    spec_metadata,
-                    new_tensors_device,
-                    graph_replay=graph_replay,
-                    static_inputs=(
-                        self.cuda_graph_runner.shared_static_tensors if graph_replay else None
-                    ),
-                    is_dummy=is_dummy,
-                    enable_spec_decode=enable_spec_decode,
-                    promoted=bool(execution_promoted_context_ids),
-                )
+                if graph_replay:
+                    stager = stager_source.begin_if_eligible(
+                        self._inputs,
+                        execution_requests,
+                        attn_metadata,
+                        spec_metadata,
+                        new_tensors_device,
+                        graph_replay=graph_replay,
+                        static_inputs=self.cuda_graph_runner.shared_static_tensors,
+                        is_dummy=is_dummy,
+                        enable_spec_decode=enable_spec_decode,
+                        promoted=bool(execution_promoted_context_ids),
+                    )
             try:
                 # The window spans the scope: the scope is a no-op whenever
                 # speculative decoding is on, which every stager requires.
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
index 24ffbccfbb8b..72987309443f 100644
--- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -55,9 +55,15 @@ def begin_if_eligible(
     ) -> Optional["StepInputStager"]:
         """Open the window for this step, or return None to keep the native preparation.
 
+        The engine calls this only for the replay of an already captured CUDA
+        graph, so ``graph_replay`` is always True: the window closes right
+        before the replay, and an eager step, an initial capture or a warmup
+        would run the model before the staged writes are issued. Those steps
+        always prepare natively.
+
         ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
         records. ``static_inputs`` are the static input tensors of the CUDA
-        graph the step replays (None when it replays none).
+        graph the step replays.
         """
 
     def copy(
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index 3993cfa7b42a..95745dd901db 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -922,3 +922,477 @@ def test_indexer_slot_mappings_are_staged_with_a_stager():
     assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
     assert bool((metadata.slot_mapping_fp8 == -1).all())
     assert bool((metadata.slot_mapping_scale == -1).all())
+
+
+# -- the step-input window in the decoder runner -------------------------------------------------
+
+
+class _ReferenceStager:
+    """A step-input stager that accepts every step it is offered and logs the window's lifecycle."""
+
+    def __init__(self, log: list, accept: bool = True, tokens: int = 3) -> None:
+        self.log, self.accept, self.tokens = log, accept, tokens
+
+    def begin_if_eligible(
+        self,
+        preparer,
+        requests,
+        attn_metadata,
+        spec_metadata,
+        new_tensors_device,
+        *,
+        graph_replay,
+        static_inputs,
+        is_dummy,
+        enable_spec_decode,
+        promoted,
+    ):
+        self.log.append(("begin", graph_replay, static_inputs))
+        return self if self.accept else None
+
+    def end(self, record_event: bool = True) -> None:
+        self.log.append(("end", record_event))
+
+    def discard(self) -> None:
+        self.log.append("discard")
+
+    def set_completion_event(self, event) -> None:
+        self.log.append(("completion_event", event))
+
+    def static_inputs_tokens(self) -> int:
+        return self.tokens
+
+
+class _WindowGraphs:
+    """The ``DecoderCudaGraphs`` calls ``_forward_decoder`` makes, with one selected graph key.
+
+    ``run`` follows ``DecoderCudaGraphs._run``: ``before_replay`` runs only right before the
+    replay of an already captured graph, and the replay consumes the staged static inputs.
+    """
+
+    def __init__(self, log: list, key, *, needs_capture=False, warmup_only=False, fail=None):
+        self.log, self.key, self.fail = log, key, fail
+        self.cuda_graph_runner = types.SimpleNamespace(
+            needs_capture=lambda key: needs_capture,
+            is_warmup_only=warmup_only,
+            shared_static_tensors={"input_ids": torch.zeros(4, dtype=torch.int64)},
+            staged_static_input_tokens=0,
+        )
+
+    def attention_dp_execution_batch(self, requests):
+        return requests
+
+    def promote_final_contexts(self, requests, kv_cache_manager, **kwargs):
+        return requests, frozenset()
+
+    def pad_batch(self, requests, resource_manager, runtime_draft_len):
+        return contextlib.nullcontext(requests)
+
+    def select(self, requests, resource_manager, **kwargs):
+        return types.SimpleNamespace(
+            key=self.key,
+            use_lora_graph=False,
+            attn_metadata=kwargs["attn_metadata"],
+            spec_metadata=None,
+        )
+
+    def captures_model_body(self, key) -> bool:
+        return False
+
+    def run(self, key, inputs, run_model, postprocess_inputs, *, before_replay=None, **kwargs):
+        runner = self.cuda_graph_runner
+        if self.fail == "before_replay":
+            raise RuntimeError("graph run failed")
+        if key is None:
+            self.log.append("eager")
+        elif runner.needs_capture(key) or runner.is_warmup_only:
+            self.log.append("capture")
+        else:
+            if before_replay is not None:
+                before_replay()
+            self.log.append(("replay", runner.staged_static_input_tokens))
+            runner.staged_static_input_tokens = 0
+            if self.fail == "replay":
+                raise RuntimeError("replay failed")
+        return "outputs"
+
+
+class _WindowPreparer:
+    """The ``DecoderInputPreparer`` calls ``_forward_decoder`` makes."""
+
+    def __init__(self, log: list, fail: bool = False) -> None:
+        self.log, self.fail = log, fail
+        self.input_copy_event = "input-copy-event"
+
+    def prepare_inputs(self, *args, **kwargs):
+        self.log.append("prepare")
+        if self.fail:
+            raise RuntimeError("preparation failed")
+        return types.SimpleNamespace(kwargs={}, runtime_draft_len=0, gather_ids=None)
+
+    def record_input_copy(self) -> None:
+        self.log.append("record_input_copy")
+
+    def postprocess_inputs(self, *args, **kwargs) -> None:
+        pass
+
+
+def _window_runner(monkeypatch, graphs, preparer, stager):
+    """A ``DecoderRunner`` with only what ``_forward_decoder`` reads (no speculation, no CED)."""
+    from tensorrt_llm._torch.pyexecutor import ced_replay
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+
+    monkeypatch.setattr(runner_module, "get_draft_kv_cache_manager", lambda *args: None)
+    monkeypatch.setattr(ced_replay, "uses_private_decoder_cache", lambda *args: False)
+    attn_metadata = types.SimpleNamespace()
+    runner = DecoderRunner.__new__(DecoderRunner)
+    runner._config = types.SimpleNamespace(spec_config=None, use_mrope=False, is_spec_decode=False)
+    runner._metadata = types.SimpleNamespace(
+        setup_attention=lambda *args: attn_metadata,
+        attn_metadata=attn_metadata,
+        spec_metadata=None,
+    )
+    runner._graphs = graphs
+    runner._inputs = preparer
+    runner._stage_in_graph_sampling = None
+    runner._cuda_graph_metadata_prepare_scope = lambda *args, **kwargs: contextlib.nullcontext()
+    runner._execute_logit_post_processors = lambda *args: None
+    runner.forward_pass_callable = None
+    runner.moe_load_balancer = None
+    runner.model = types.SimpleNamespace()
+    runner.kv_cache_manager_key = "kv_cache_manager"
+    runner.runtime_hooks = RuntimeHooks(step_input_stager=stager)
+    return runner
+
+
+def _decode_step(runner):
+    forward_inputs = types.SimpleNamespace(
+        batch=types.SimpleNamespace(num_context_requests=0),
+        new_tensors_device=None,
+        cache_indirection_buffer=None,
+        gather_context_logits=False,
+        enable_spec_decode=False,
+        runtime_draft_len=0,
+    )
+    resource_manager = types.SimpleNamespace(get_resource_manager=lambda key: object())
+    return runner._forward_decoder(forward_inputs, resource_manager, is_dummy=False)
+
+
+def _names(log: list) -> list:
+    return [entry if isinstance(entry, str) else entry[0] for entry in log]
+
+
+_GRAPH_KEY = (1, 0, False)
+_STAGED_STEP = ["begin", "prepare", "end", "record_input_copy", "completion_event", "replay"]
+
+
+def test_a_replay_issues_the_staged_preparation_before_the_replay(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    assert _decode_step(runner) == "outputs"
+    assert log == [
+        ("begin", True, graphs.cuda_graph_runner.shared_static_tensors),
+        "prepare",
+        ("end", False),
+        "record_input_copy",
+        ("completion_event", "input-copy-event"),
+        ("replay", 3),
+    ]
+
+
+@pytest.mark.parametrize(
+    "key, needs_capture, warmup_only, ran",
+    [
+        (None, False, False, "eager"),
+        (_GRAPH_KEY, True, False, "capture"),
+        (_GRAPH_KEY, False, True, "capture"),
+    ],
+    ids=["eager", "initial_capture", "warmup_only"],
+)
+def test_a_step_that_is_not_a_replay_never_opens_a_window(
+    monkeypatch, key, needs_capture, warmup_only, ran
+):
+    log = []
+    graphs = _WindowGraphs(log, key, needs_capture=needs_capture, warmup_only=warmup_only)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    _decode_step(runner)
+    assert log == ["prepare", "record_input_copy", ran]
+    assert graphs.cuda_graph_runner.staged_static_input_tokens == 0
+
+
+def test_a_declined_replay_prepares_natively(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    stager = _ReferenceStager(log, accept=False)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), stager)
+
+    _decode_step(runner)
+    assert _names(log) == ["begin", "prepare", "record_input_copy", "replay"]
+    assert log[-1] == ("replay", 0)
+
+
+def test_a_failed_preparation_discards_the_window(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(
+        monkeypatch, graphs, _WindowPreparer(log, fail=True), _ReferenceStager(log)
+    )
+
+    with pytest.raises(RuntimeError, match="preparation failed"):
+        _decode_step(runner)
+    assert _names(log) == ["begin", "prepare", "discard"]
+
+
+@pytest.mark.parametrize(
+    "fail, expected",
+    [
+        ("before_replay", ["begin", "prepare", "discard"]),
+        ("replay", _STAGED_STEP),
+    ],
+    ids=["window_open", "window_closed"],
+)
+def test_a_failed_run_discards_only_a_window_still_open(monkeypatch, fail, expected):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY, fail=fail)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    with pytest.raises(RuntimeError):
+        _decode_step(runner)
+    assert _names(log) == expected
+
+
+def test_replay_eager_replay_opens_one_window_per_replay(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    _decode_step(runner)
+    graphs.key = None
+    _decode_step(runner)
+    graphs.key = _GRAPH_KEY
+    _decode_step(runner)
+    assert _names(log) == [
+        *_STAGED_STEP,
+        "prepare",
+        "record_input_copy",
+        "eager",
+        *_STAGED_STEP,
+    ]
+    assert [entry for entry in log if entry[0] == "replay"] == [("replay", 3), ("replay", 3)]
+
+
+def _replay_graphs(monkeypatch, log: list, *, needs_capture: bool, warmup_only: bool):
+    """A ``DecoderCudaGraphs`` with only what ``_run`` reads, logging capture and replay."""
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs
+
+    monkeypatch.setattr(graphs_module, "prepare_attn_metadata_for_draft_replay", lambda *args: None)
+    monkeypatch.setattr(
+        graphs_module, "restore_attn_metadata_after_draft_replay", lambda *args: None
+    )
+    monkeypatch.setattr(
+        graphs_module, "MoeLoadBalancerIterContext", lambda balancer: contextlib.nullcontext()
+    )
+    graphs = DecoderCudaGraphs.__new__(DecoderCudaGraphs)
+    graphs.breakable_cuda_graph_runner = None
+    graphs._torch_compile_bypass = contextlib.nullcontext
+    graphs.cuda_graph_runner = types.SimpleNamespace(
+        needs_capture=lambda key: needs_capture,
+        is_warmup_only=warmup_only,
+        capture=lambda key, fn, inputs, **kwargs: log.append("capture"),
+        replay=lambda key, inputs, **kwargs: log.append("replay"),
+    )
+    return graphs
+
+
+@pytest.mark.parametrize(
+    "key, needs_capture, warmup_only, expected",
+    [
+        (_GRAPH_KEY, False, False, ["before_replay", "replay"]),
+        (_GRAPH_KEY, True, False, ["capture", "replay"]),
+        (_GRAPH_KEY, True, True, ["capture"]),
+        (None, False, False, ["eager"]),
+    ],
+    ids=["replay", "capture_then_replay", "warmup_only", "eager"],
+)
+def test_before_replay_runs_only_before_the_replay_of_a_captured_graph(
+    monkeypatch, key, needs_capture, warmup_only, expected
+):
+    log = []
+    graphs = _replay_graphs(monkeypatch, log, needs_capture=needs_capture, warmup_only=warmup_only)
+
+    graphs._run(
+        key,
+        {"input_ids": torch.zeros(2, dtype=torch.int64)},
+        lambda inputs: log.append("eager"),
+        lambda inputs: None,
+        attn_metadata=None,
+        draft_kv_cache_manager=None,
+        moe_load_balancer=None,
+        enable_spec_decode=True,
+        gather_context_logits=False,
+        position_id_offsets=None,
+        before_replay=lambda: log.append("before_replay"),
+    )
+    assert log == expected
+
+
+# -- the speculative sampler's store writer -----------------------------------------------------
+
+
+class _ReferenceStoreWriter:
+    """Stores like the native ``index_copy_`` path, then copies the store to the host tensors."""
+
+    def __init__(self, supported: bool = True) -> None:
+        self._supported = supported
+        self.calls = []
+
+    def supported(self, store, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens):
+        self.calls.append("supported")
+        return self._supported
+
+    def store_and_copy_out(
+        self,
+        store,
+        slots,
+        new_tokens,
+        next_new_tokens,
+        new_tokens_lens,
+        next_draft_tokens,
+        host_new_tokens,
+        host_new_tokens_lens,
+        host_next_draft_tokens,
+        record_event=True,
+    ) -> None:
+        self.calls.append(("store_and_copy_out", list(slots), record_event))
+        index = torch.tensor(slots, dtype=torch.long)
+        store.new_tokens.squeeze(-1).T.index_copy_(0, index, new_tokens)
+        store.next_new_tokens.squeeze(-1).T.index_copy_(0, index, next_new_tokens)
+        store.new_tokens_lens.index_copy_(0, index, new_tokens_lens)
+        store.next_draft_tokens.index_copy_(0, index, next_draft_tokens)
+        host_new_tokens.copy_(store.new_tokens)
+        host_new_tokens_lens.copy_(store.new_tokens_lens)
+        host_next_draft_tokens.copy_(store.next_draft_tokens)
+
+    def set_completion_event(self, event) -> None:
+        self.calls.append(("completion_event", event))
+
+
+def _cpu_sampling(monkeypatch, *, pinned: bool) -> None:
+    """Run ``SpecSampler.sample_async`` on CPU: device transfers stay on the host."""
+    from tensorrt_llm._torch.speculative import spec_sampler_base
+
+    original_to = torch.Tensor.to
+    original_empty_like = torch.empty_like
+
+    def to_cpu(tensor, *args, **kwargs):
+        if kwargs.get("device") == "cuda":
+            kwargs = {**kwargs, "device": "cpu"}
+        return original_to(tensor, *args, **kwargs)
+
+    def empty_like(tensor, *args, pin_memory=False, **kwargs):
+        # Pinning needs a CUDA context; the writer's host buffers are plain host tensors here.
+        return original_empty_like(tensor, *args, **kwargs)
+
+    monkeypatch.setattr(torch.Tensor, "to", to_cpu)
+    monkeypatch.setattr(torch, "empty_like", empty_like)
+    monkeypatch.setattr(spec_sampler_base, "prefer_pinned", lambda: pinned)
+    monkeypatch.setattr(spec_sampler_base, "handle_stop_criteria", lambda *args, **kwargs: False)
+
+
+def _store_sampler(writer, *, async_worker: bool = False):
+    """A ``SpecSampler`` with a six-slot store, draft length 3 and the given store writer."""
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    sampler = SpecSampler.__new__(SpecSampler)
+    sampler.draft_len = 3
+    sampler.max_accepted_path_len = 4
+    sampler.max_seq_len = 128
+    sampler.store = types.SimpleNamespace(
+        new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
+        next_new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
+        next_draft_tokens=torch.zeros(6, 3, dtype=torch.int32),
+        new_tokens_lens=torch.zeros(6, dtype=torch.int32),
+    )
+    sampler.store_writer = writer
+    sampler._async_worker_active = lambda: async_worker
+    sampler._copy_to_host = lambda tensor: tensor.clone()
+    sampler._record_sampler_event = lambda: types.SimpleNamespace(cuda_event="sampler-event")
+    return sampler
+
+
+def _sample_generation_step(sampler, *, verify_lens: bool = False):
+    """Three generation requests in nonconsecutive slots 4, 1 and 3."""
+    from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
+
+    def request(request_id: int, slot: int):
+        return types.SimpleNamespace(
+            py_request_id=request_id,
+            py_seq_slot=slot,
+            state=LlmRequestState.GENERATION_IN_PROGRESS,
+            py_draft_tokens=[10, 11, 12],
+            py_draft_tokens_effective_len=None,
+            py_verify_len=None,
+            py_decoding_iter=0,
+        )
+
+    scheduled = types.SimpleNamespace(
+        context_requests_chunking=[],
+        context_requests_last_chunk=[],
+        generation_requests=[request(7, 4), request(8, 1), request(9, 3)],
+    )
+    outputs = dict(
+        new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 100,
+        new_tokens_lens=torch.tensor([1, 3, 2], dtype=torch.int32),
+        next_draft_tokens=torch.arange(9, dtype=torch.int32).reshape(3, 3) + 200,
+        next_new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 300,
+    )
+    if verify_lens:
+        outputs["verify_lens"] = torch.tensor([2, 3, 1], dtype=torch.int32)
+        outputs["verify_lens_in_output_order"] = True
+    return sampler.sample_async(scheduled, outputs, [])
+
+
+def _assert_same_stores(sampler, state, native, native_state) -> None:
+    for name in ("new_tokens", "next_new_tokens", "next_draft_tokens", "new_tokens_lens"):
+        assert torch.equal(getattr(sampler.store, name), getattr(native.store, name)), name
+    for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"):
+        assert torch.equal(getattr(state.host, name), getattr(native_state.host, name)), name
+
+
+def test_the_store_writer_matches_the_native_stores_on_nonconsecutive_slots(monkeypatch):
+    _cpu_sampling(monkeypatch, pinned=True)
+    native = _store_sampler(None)
+    writer = _ReferenceStoreWriter()
+    sampler = _store_sampler(writer)
+
+    native_state = _sample_generation_step(native)
+    state = _sample_generation_step(sampler)
+    assert writer.calls == [
+        "supported",
+        ("store_and_copy_out", [4, 1, 3], False),
+        ("completion_event", "sampler-event"),
+    ]
+    assert native.store.new_tokens_lens.tolist() == [0, 3, 0, 2, 1, 0]
+    _assert_same_stores(sampler, state, native, native_state)
+    assert state.host.verify_lens is None
+    assert state.device.new_tokens is sampler.store.next_new_tokens
+
+
+@pytest.mark.parametrize("case", ["declined", "verify_lens", "async_worker", "not_pinned"])
+def test_the_store_writer_falls_back_to_the_native_stores(monkeypatch, case):
+    _cpu_sampling(monkeypatch, pinned=case != "not_pinned")
+    native = _store_sampler(None)
+    writer = _ReferenceStoreWriter(supported=case != "declined")
+    sampler = _store_sampler(writer, async_worker=case == "async_worker")
+
+    native_state = _sample_generation_step(native, verify_lens=case == "verify_lens")
+    state = _sample_generation_step(sampler, verify_lens=case == "verify_lens")
+    assert writer.calls == (["supported"] if case == "declined" else [])
+    _assert_same_stores(sampler, state, native, native_state)
+    if case == "verify_lens":
+        assert state.host.verify_lens.tolist() == [2, 3, 1]

From f6214d7723931b3247ea78665e619c03fe023c22 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Thu, 22 Oct 2026 11:51:34 -0700
Subject: [PATCH 17/19] [TRTLLM-17096][fix] runtime hooks: stage step inputs
 only on the replay of a captured graph

The step-input window closes right before the replay of an already captured
CUDA graph. An eager step, an initial capture or a warmup would run the model
before the staged writes are issued, so the decoder runner now offers the
window only on such a replay, and StepInputStager documents the restriction.

Add behavioral tests with reference hooks: the runner window lifecycle
(staged writes issued before the completion event and the replay, discard on
a failure while the window is open), steps that are not a replay never open a
window, before_replay runs only before the replay of a captured graph, and
the speculative sampler's store writer against the native stores on
nonconsecutive slots, with its native fallbacks.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../engine/runners/decoder/runner.py          |  31 +-
 .../_torch/pyexecutor/runtime_hooks.py        |   8 +-
 .../_torch/executor/test_runtime_hooks.py     | 474 ++++++++++++++++++
 3 files changed, 498 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 060f30b6b6a5..e31f20da8e40 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -640,25 +640,28 @@ def _forward_decoder(
             stager = None
             stager_source = self.runtime_hooks.step_input_stager
             if stager_source is not None:
+                # Only the replay of an already captured graph closes the window
+                # (before_replay); an eager step, an initial capture or a warmup
+                # would run the model before the staged writes are issued, so
+                # those steps always prepare natively.
                 graph_replay = (
                     can_run_graph
                     and not self.cuda_graph_runner.needs_capture(key)
                     and not self.cuda_graph_runner.is_warmup_only
                 )
-                stager = stager_source.begin_if_eligible(
-                    self._inputs,
-                    execution_requests,
-                    attn_metadata,
-                    spec_metadata,
-                    new_tensors_device,
-                    graph_replay=graph_replay,
-                    static_inputs=(
-                        self.cuda_graph_runner.shared_static_tensors if graph_replay else None
-                    ),
-                    is_dummy=is_dummy,
-                    enable_spec_decode=enable_spec_decode,
-                    promoted=bool(execution_promoted_context_ids),
-                )
+                if graph_replay:
+                    stager = stager_source.begin_if_eligible(
+                        self._inputs,
+                        execution_requests,
+                        attn_metadata,
+                        spec_metadata,
+                        new_tensors_device,
+                        graph_replay=graph_replay,
+                        static_inputs=self.cuda_graph_runner.shared_static_tensors,
+                        is_dummy=is_dummy,
+                        enable_spec_decode=enable_spec_decode,
+                        promoted=bool(execution_promoted_context_ids),
+                    )
             try:
                 # The window spans the scope: the scope is a no-op whenever
                 # speculative decoding is on, which every stager requires.
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
index 24ffbccfbb8b..72987309443f 100644
--- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -55,9 +55,15 @@ def begin_if_eligible(
     ) -> Optional["StepInputStager"]:
         """Open the window for this step, or return None to keep the native preparation.
 
+        The engine calls this only for the replay of an already captured CUDA
+        graph, so ``graph_replay`` is always True: the window closes right
+        before the replay, and an eager step, an initial capture or a warmup
+        would run the model before the staged writes are issued. Those steps
+        always prepare natively.
+
         ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
         records. ``static_inputs`` are the static input tensors of the CUDA
-        graph the step replays (None when it replays none).
+        graph the step replays.
         """
 
     def copy(
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index 3993cfa7b42a..95745dd901db 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -922,3 +922,477 @@ def test_indexer_slot_mappings_are_staged_with_a_stager():
     assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
     assert bool((metadata.slot_mapping_fp8 == -1).all())
     assert bool((metadata.slot_mapping_scale == -1).all())
+
+
+# -- the step-input window in the decoder runner -------------------------------------------------
+
+
+class _ReferenceStager:
+    """A step-input stager that accepts every step it is offered and logs the window's lifecycle."""
+
+    def __init__(self, log: list, accept: bool = True, tokens: int = 3) -> None:
+        self.log, self.accept, self.tokens = log, accept, tokens
+
+    def begin_if_eligible(
+        self,
+        preparer,
+        requests,
+        attn_metadata,
+        spec_metadata,
+        new_tensors_device,
+        *,
+        graph_replay,
+        static_inputs,
+        is_dummy,
+        enable_spec_decode,
+        promoted,
+    ):
+        self.log.append(("begin", graph_replay, static_inputs))
+        return self if self.accept else None
+
+    def end(self, record_event: bool = True) -> None:
+        self.log.append(("end", record_event))
+
+    def discard(self) -> None:
+        self.log.append("discard")
+
+    def set_completion_event(self, event) -> None:
+        self.log.append(("completion_event", event))
+
+    def static_inputs_tokens(self) -> int:
+        return self.tokens
+
+
+class _WindowGraphs:
+    """The ``DecoderCudaGraphs`` calls ``_forward_decoder`` makes, with one selected graph key.
+
+    ``run`` follows ``DecoderCudaGraphs._run``: ``before_replay`` runs only right before the
+    replay of an already captured graph, and the replay consumes the staged static inputs.
+    """
+
+    def __init__(self, log: list, key, *, needs_capture=False, warmup_only=False, fail=None):
+        self.log, self.key, self.fail = log, key, fail
+        self.cuda_graph_runner = types.SimpleNamespace(
+            needs_capture=lambda key: needs_capture,
+            is_warmup_only=warmup_only,
+            shared_static_tensors={"input_ids": torch.zeros(4, dtype=torch.int64)},
+            staged_static_input_tokens=0,
+        )
+
+    def attention_dp_execution_batch(self, requests):
+        return requests
+
+    def promote_final_contexts(self, requests, kv_cache_manager, **kwargs):
+        return requests, frozenset()
+
+    def pad_batch(self, requests, resource_manager, runtime_draft_len):
+        return contextlib.nullcontext(requests)
+
+    def select(self, requests, resource_manager, **kwargs):
+        return types.SimpleNamespace(
+            key=self.key,
+            use_lora_graph=False,
+            attn_metadata=kwargs["attn_metadata"],
+            spec_metadata=None,
+        )
+
+    def captures_model_body(self, key) -> bool:
+        return False
+
+    def run(self, key, inputs, run_model, postprocess_inputs, *, before_replay=None, **kwargs):
+        runner = self.cuda_graph_runner
+        if self.fail == "before_replay":
+            raise RuntimeError("graph run failed")
+        if key is None:
+            self.log.append("eager")
+        elif runner.needs_capture(key) or runner.is_warmup_only:
+            self.log.append("capture")
+        else:
+            if before_replay is not None:
+                before_replay()
+            self.log.append(("replay", runner.staged_static_input_tokens))
+            runner.staged_static_input_tokens = 0
+            if self.fail == "replay":
+                raise RuntimeError("replay failed")
+        return "outputs"
+
+
+class _WindowPreparer:
+    """The ``DecoderInputPreparer`` calls ``_forward_decoder`` makes."""
+
+    def __init__(self, log: list, fail: bool = False) -> None:
+        self.log, self.fail = log, fail
+        self.input_copy_event = "input-copy-event"
+
+    def prepare_inputs(self, *args, **kwargs):
+        self.log.append("prepare")
+        if self.fail:
+            raise RuntimeError("preparation failed")
+        return types.SimpleNamespace(kwargs={}, runtime_draft_len=0, gather_ids=None)
+
+    def record_input_copy(self) -> None:
+        self.log.append("record_input_copy")
+
+    def postprocess_inputs(self, *args, **kwargs) -> None:
+        pass
+
+
+def _window_runner(monkeypatch, graphs, preparer, stager):
+    """A ``DecoderRunner`` with only what ``_forward_decoder`` reads (no speculation, no CED)."""
+    from tensorrt_llm._torch.pyexecutor import ced_replay
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
+
+    monkeypatch.setattr(runner_module, "get_draft_kv_cache_manager", lambda *args: None)
+    monkeypatch.setattr(ced_replay, "uses_private_decoder_cache", lambda *args: False)
+    attn_metadata = types.SimpleNamespace()
+    runner = DecoderRunner.__new__(DecoderRunner)
+    runner._config = types.SimpleNamespace(spec_config=None, use_mrope=False, is_spec_decode=False)
+    runner._metadata = types.SimpleNamespace(
+        setup_attention=lambda *args: attn_metadata,
+        attn_metadata=attn_metadata,
+        spec_metadata=None,
+    )
+    runner._graphs = graphs
+    runner._inputs = preparer
+    runner._stage_in_graph_sampling = None
+    runner._cuda_graph_metadata_prepare_scope = lambda *args, **kwargs: contextlib.nullcontext()
+    runner._execute_logit_post_processors = lambda *args: None
+    runner.forward_pass_callable = None
+    runner.moe_load_balancer = None
+    runner.model = types.SimpleNamespace()
+    runner.kv_cache_manager_key = "kv_cache_manager"
+    runner.runtime_hooks = RuntimeHooks(step_input_stager=stager)
+    return runner
+
+
+def _decode_step(runner):
+    forward_inputs = types.SimpleNamespace(
+        batch=types.SimpleNamespace(num_context_requests=0),
+        new_tensors_device=None,
+        cache_indirection_buffer=None,
+        gather_context_logits=False,
+        enable_spec_decode=False,
+        runtime_draft_len=0,
+    )
+    resource_manager = types.SimpleNamespace(get_resource_manager=lambda key: object())
+    return runner._forward_decoder(forward_inputs, resource_manager, is_dummy=False)
+
+
+def _names(log: list) -> list:
+    return [entry if isinstance(entry, str) else entry[0] for entry in log]
+
+
+_GRAPH_KEY = (1, 0, False)
+_STAGED_STEP = ["begin", "prepare", "end", "record_input_copy", "completion_event", "replay"]
+
+
+def test_a_replay_issues_the_staged_preparation_before_the_replay(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    assert _decode_step(runner) == "outputs"
+    assert log == [
+        ("begin", True, graphs.cuda_graph_runner.shared_static_tensors),
+        "prepare",
+        ("end", False),
+        "record_input_copy",
+        ("completion_event", "input-copy-event"),
+        ("replay", 3),
+    ]
+
+
+@pytest.mark.parametrize(
+    "key, needs_capture, warmup_only, ran",
+    [
+        (None, False, False, "eager"),
+        (_GRAPH_KEY, True, False, "capture"),
+        (_GRAPH_KEY, False, True, "capture"),
+    ],
+    ids=["eager", "initial_capture", "warmup_only"],
+)
+def test_a_step_that_is_not_a_replay_never_opens_a_window(
+    monkeypatch, key, needs_capture, warmup_only, ran
+):
+    log = []
+    graphs = _WindowGraphs(log, key, needs_capture=needs_capture, warmup_only=warmup_only)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    _decode_step(runner)
+    assert log == ["prepare", "record_input_copy", ran]
+    assert graphs.cuda_graph_runner.staged_static_input_tokens == 0
+
+
+def test_a_declined_replay_prepares_natively(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    stager = _ReferenceStager(log, accept=False)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), stager)
+
+    _decode_step(runner)
+    assert _names(log) == ["begin", "prepare", "record_input_copy", "replay"]
+    assert log[-1] == ("replay", 0)
+
+
+def test_a_failed_preparation_discards_the_window(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(
+        monkeypatch, graphs, _WindowPreparer(log, fail=True), _ReferenceStager(log)
+    )
+
+    with pytest.raises(RuntimeError, match="preparation failed"):
+        _decode_step(runner)
+    assert _names(log) == ["begin", "prepare", "discard"]
+
+
+@pytest.mark.parametrize(
+    "fail, expected",
+    [
+        ("before_replay", ["begin", "prepare", "discard"]),
+        ("replay", _STAGED_STEP),
+    ],
+    ids=["window_open", "window_closed"],
+)
+def test_a_failed_run_discards_only_a_window_still_open(monkeypatch, fail, expected):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY, fail=fail)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    with pytest.raises(RuntimeError):
+        _decode_step(runner)
+    assert _names(log) == expected
+
+
+def test_replay_eager_replay_opens_one_window_per_replay(monkeypatch):
+    log = []
+    graphs = _WindowGraphs(log, _GRAPH_KEY)
+    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
+
+    _decode_step(runner)
+    graphs.key = None
+    _decode_step(runner)
+    graphs.key = _GRAPH_KEY
+    _decode_step(runner)
+    assert _names(log) == [
+        *_STAGED_STEP,
+        "prepare",
+        "record_input_copy",
+        "eager",
+        *_STAGED_STEP,
+    ]
+    assert [entry for entry in log if entry[0] == "replay"] == [("replay", 3), ("replay", 3)]
+
+
+def _replay_graphs(monkeypatch, log: list, *, needs_capture: bool, warmup_only: bool):
+    """A ``DecoderCudaGraphs`` with only what ``_run`` reads, logging capture and replay."""
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module
+    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs
+
+    monkeypatch.setattr(graphs_module, "prepare_attn_metadata_for_draft_replay", lambda *args: None)
+    monkeypatch.setattr(
+        graphs_module, "restore_attn_metadata_after_draft_replay", lambda *args: None
+    )
+    monkeypatch.setattr(
+        graphs_module, "MoeLoadBalancerIterContext", lambda balancer: contextlib.nullcontext()
+    )
+    graphs = DecoderCudaGraphs.__new__(DecoderCudaGraphs)
+    graphs.breakable_cuda_graph_runner = None
+    graphs._torch_compile_bypass = contextlib.nullcontext
+    graphs.cuda_graph_runner = types.SimpleNamespace(
+        needs_capture=lambda key: needs_capture,
+        is_warmup_only=warmup_only,
+        capture=lambda key, fn, inputs, **kwargs: log.append("capture"),
+        replay=lambda key, inputs, **kwargs: log.append("replay"),
+    )
+    return graphs
+
+
+@pytest.mark.parametrize(
+    "key, needs_capture, warmup_only, expected",
+    [
+        (_GRAPH_KEY, False, False, ["before_replay", "replay"]),
+        (_GRAPH_KEY, True, False, ["capture", "replay"]),
+        (_GRAPH_KEY, True, True, ["capture"]),
+        (None, False, False, ["eager"]),
+    ],
+    ids=["replay", "capture_then_replay", "warmup_only", "eager"],
+)
+def test_before_replay_runs_only_before_the_replay_of_a_captured_graph(
+    monkeypatch, key, needs_capture, warmup_only, expected
+):
+    log = []
+    graphs = _replay_graphs(monkeypatch, log, needs_capture=needs_capture, warmup_only=warmup_only)
+
+    graphs._run(
+        key,
+        {"input_ids": torch.zeros(2, dtype=torch.int64)},
+        lambda inputs: log.append("eager"),
+        lambda inputs: None,
+        attn_metadata=None,
+        draft_kv_cache_manager=None,
+        moe_load_balancer=None,
+        enable_spec_decode=True,
+        gather_context_logits=False,
+        position_id_offsets=None,
+        before_replay=lambda: log.append("before_replay"),
+    )
+    assert log == expected
+
+
+# -- the speculative sampler's store writer -----------------------------------------------------
+
+
+class _ReferenceStoreWriter:
+    """Stores like the native ``index_copy_`` path, then copies the store to the host tensors."""
+
+    def __init__(self, supported: bool = True) -> None:
+        self._supported = supported
+        self.calls = []
+
+    def supported(self, store, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens):
+        self.calls.append("supported")
+        return self._supported
+
+    def store_and_copy_out(
+        self,
+        store,
+        slots,
+        new_tokens,
+        next_new_tokens,
+        new_tokens_lens,
+        next_draft_tokens,
+        host_new_tokens,
+        host_new_tokens_lens,
+        host_next_draft_tokens,
+        record_event=True,
+    ) -> None:
+        self.calls.append(("store_and_copy_out", list(slots), record_event))
+        index = torch.tensor(slots, dtype=torch.long)
+        store.new_tokens.squeeze(-1).T.index_copy_(0, index, new_tokens)
+        store.next_new_tokens.squeeze(-1).T.index_copy_(0, index, next_new_tokens)
+        store.new_tokens_lens.index_copy_(0, index, new_tokens_lens)
+        store.next_draft_tokens.index_copy_(0, index, next_draft_tokens)
+        host_new_tokens.copy_(store.new_tokens)
+        host_new_tokens_lens.copy_(store.new_tokens_lens)
+        host_next_draft_tokens.copy_(store.next_draft_tokens)
+
+    def set_completion_event(self, event) -> None:
+        self.calls.append(("completion_event", event))
+
+
+def _cpu_sampling(monkeypatch, *, pinned: bool) -> None:
+    """Run ``SpecSampler.sample_async`` on CPU: device transfers stay on the host."""
+    from tensorrt_llm._torch.speculative import spec_sampler_base
+
+    original_to = torch.Tensor.to
+    original_empty_like = torch.empty_like
+
+    def to_cpu(tensor, *args, **kwargs):
+        if kwargs.get("device") == "cuda":
+            kwargs = {**kwargs, "device": "cpu"}
+        return original_to(tensor, *args, **kwargs)
+
+    def empty_like(tensor, *args, pin_memory=False, **kwargs):
+        # Pinning needs a CUDA context; the writer's host buffers are plain host tensors here.
+        return original_empty_like(tensor, *args, **kwargs)
+
+    monkeypatch.setattr(torch.Tensor, "to", to_cpu)
+    monkeypatch.setattr(torch, "empty_like", empty_like)
+    monkeypatch.setattr(spec_sampler_base, "prefer_pinned", lambda: pinned)
+    monkeypatch.setattr(spec_sampler_base, "handle_stop_criteria", lambda *args, **kwargs: False)
+
+
+def _store_sampler(writer, *, async_worker: bool = False):
+    """A ``SpecSampler`` with a six-slot store, draft length 3 and the given store writer."""
+    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
+
+    sampler = SpecSampler.__new__(SpecSampler)
+    sampler.draft_len = 3
+    sampler.max_accepted_path_len = 4
+    sampler.max_seq_len = 128
+    sampler.store = types.SimpleNamespace(
+        new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
+        next_new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
+        next_draft_tokens=torch.zeros(6, 3, dtype=torch.int32),
+        new_tokens_lens=torch.zeros(6, dtype=torch.int32),
+    )
+    sampler.store_writer = writer
+    sampler._async_worker_active = lambda: async_worker
+    sampler._copy_to_host = lambda tensor: tensor.clone()
+    sampler._record_sampler_event = lambda: types.SimpleNamespace(cuda_event="sampler-event")
+    return sampler
+
+
+def _sample_generation_step(sampler, *, verify_lens: bool = False):
+    """Three generation requests in nonconsecutive slots 4, 1 and 3."""
+    from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
+
+    def request(request_id: int, slot: int):
+        return types.SimpleNamespace(
+            py_request_id=request_id,
+            py_seq_slot=slot,
+            state=LlmRequestState.GENERATION_IN_PROGRESS,
+            py_draft_tokens=[10, 11, 12],
+            py_draft_tokens_effective_len=None,
+            py_verify_len=None,
+            py_decoding_iter=0,
+        )
+
+    scheduled = types.SimpleNamespace(
+        context_requests_chunking=[],
+        context_requests_last_chunk=[],
+        generation_requests=[request(7, 4), request(8, 1), request(9, 3)],
+    )
+    outputs = dict(
+        new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 100,
+        new_tokens_lens=torch.tensor([1, 3, 2], dtype=torch.int32),
+        next_draft_tokens=torch.arange(9, dtype=torch.int32).reshape(3, 3) + 200,
+        next_new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 300,
+    )
+    if verify_lens:
+        outputs["verify_lens"] = torch.tensor([2, 3, 1], dtype=torch.int32)
+        outputs["verify_lens_in_output_order"] = True
+    return sampler.sample_async(scheduled, outputs, [])
+
+
+def _assert_same_stores(sampler, state, native, native_state) -> None:
+    for name in ("new_tokens", "next_new_tokens", "next_draft_tokens", "new_tokens_lens"):
+        assert torch.equal(getattr(sampler.store, name), getattr(native.store, name)), name
+    for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"):
+        assert torch.equal(getattr(state.host, name), getattr(native_state.host, name)), name
+
+
+def test_the_store_writer_matches_the_native_stores_on_nonconsecutive_slots(monkeypatch):
+    _cpu_sampling(monkeypatch, pinned=True)
+    native = _store_sampler(None)
+    writer = _ReferenceStoreWriter()
+    sampler = _store_sampler(writer)
+
+    native_state = _sample_generation_step(native)
+    state = _sample_generation_step(sampler)
+    assert writer.calls == [
+        "supported",
+        ("store_and_copy_out", [4, 1, 3], False),
+        ("completion_event", "sampler-event"),
+    ]
+    assert native.store.new_tokens_lens.tolist() == [0, 3, 0, 2, 1, 0]
+    _assert_same_stores(sampler, state, native, native_state)
+    assert state.host.verify_lens is None
+    assert state.device.new_tokens is sampler.store.next_new_tokens
+
+
+@pytest.mark.parametrize("case", ["declined", "verify_lens", "async_worker", "not_pinned"])
+def test_the_store_writer_falls_back_to_the_native_stores(monkeypatch, case):
+    _cpu_sampling(monkeypatch, pinned=case != "not_pinned")
+    native = _store_sampler(None)
+    writer = _ReferenceStoreWriter(supported=case != "declined")
+    sampler = _store_sampler(writer, async_worker=case == "async_worker")
+
+    native_state = _sample_generation_step(native, verify_lens=case == "verify_lens")
+    state = _sample_generation_step(sampler, verify_lens=case == "verify_lens")
+    assert writer.calls == (["supported"] if case == "declined" else [])
+    _assert_same_stores(sampler, state, native, native_state)
+    if case == "verify_lens":
+        assert state.host.verify_lens.tolist() == [2, 3, 1]

From ce1d3f397fe4613a3207827a74b81fb2485c0a46 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Mon, 26 Oct 2026 21:49:27 -0700
Subject: [PATCH 18/19] [TRTLLM-17096][chore] runtime hooks: drop the
 step-input stager and the sampler store writer

Main's one-model speculative decoding work now batches the overlap input
gathers and moves the sampler's host copies off the critical path natively,
so these two hooks no longer fit the code they replaced. Remove them and
their call sites; they can be ported onto the new native paths later.

Keep the decode prologue, the executor events, the execution-stream loop and
the two extension points (prepended FMHA libraries, the DSA indexer's
overridable top-k copy).

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../attention/backends/sparse/dsa/indexer.py  |  20 +-
 .../attention/backends/sparse/dsa/metadata.py |  98 +--
 .../_torch/attention/backends/trtllm.py       |  59 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   8 +-
 .../_torch/pyexecutor/cuda_graph_runner.py    |  14 +-
 .../pyexecutor/engine/runners/common.py       |   8 -
 .../engine/runners/decoder/graphs.py          |   7 -
 .../engine/runners/decoder/inputs.py          |  86 +-
 .../engine/runners/decoder/runner.py          | 153 +---
 .../kv_cache/kv_cache_manager_v2.py           |  13 -
 .../_torch/pyexecutor/runtime_hooks.py        | 154 +---
 tensorrt_llm/_torch/speculative/eagle3.py     |   9 +-
 tensorrt_llm/_torch/speculative/interface.py  |  10 +-
 .../_torch/speculative/spec_sampler_base.py   |  82 +-
 .../_torch/executor/test_runtime_hooks.py     | 754 +-----------------
 15 files changed, 129 insertions(+), 1346 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
index 9e1ff5b6f902..d69ea5be0128 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py
@@ -1240,18 +1240,6 @@ def recompute_slot_mappings(
         metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices
         metadata.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices
 
-        stager = getattr(metadata, "step_input_stager", None)
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(
-                metadata.slot_mapping_fp8[:total_new_kv_tokens],
-                metadata.host_slot_mapping_fp8[:total_new_kv_tokens],
-            )
-            stager.copy(
-                metadata.slot_mapping_scale[:total_new_kv_tokens],
-                metadata.host_slot_mapping_scale[:total_new_kv_tokens],
-            )
-            return
         metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_(
             metadata.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True
         )
@@ -1526,13 +1514,7 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None:
         if num_generations > 0:
             # Prepare schedule metadata for fp8_paged_mqa_logits
             # This is a preprocessing step that computes scheduling information for the kernel
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager = getattr(metadata, "step_input_stager", None)
-            if stager is None:
-                Indexer.prepare_scheduler_metadata(metadata)
-            elif not stager.scheduler_metadata(metadata):
-                # Reads staged KV lengths: run right after the staged launch.
-                stager.defer(lambda: Indexer.prepare_scheduler_metadata(metadata))
+            Indexer.prepare_scheduler_metadata(metadata)
 
     def _update_k_cache(
         self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
index 5c9c4686375c..e0773630dcb3 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py
@@ -1486,15 +1486,6 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale
         num_blocks = page_indices.shape[1]
         self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices)
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.indexer_k_cache_block_offsets[: self.num_seqs],
-                self.host_indexer_k_cache_block_offsets[: self.num_seqs],
-                clamp_min=0,
-            )
-            self.step_input_stager.clamp_min(self.indexer_k_cache_block_offsets[self.num_seqs :], 0)
-            return pool_indices
         self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
             self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
         )
@@ -1711,16 +1702,9 @@ def prepare_dense_topk_indices(self, kv_lens, device=False):  # device=False mea
                     kv_lens[self.num_contexts : self.num_seqs],
                     self.num_tokens - self.num_ctx_tokens,
                 )
-                if self.step_input_stager is not None:
-                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                    self.step_input_stager.copy(
-                        self.topk_indices_buffer[gen_range, :],
-                        self.host_topk_indices_buffer[gen_range, :],
-                    )
-                else:
-                    self.topk_indices_buffer[gen_range, :].copy_(
-                        self.host_topk_indices_buffer[gen_range, :], non_blocking=True
-                    )
+                self.topk_indices_buffer[gen_range, :].copy_(
+                    self.host_topk_indices_buffer[gen_range, :], non_blocking=True
+                )
 
     def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
         # The DeepGEMM paged-MQA kernel (fp8_paged_mqa_logits) runs a native
@@ -1820,16 +1804,9 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 num_tokens = self.num_generations * expand_factor
                 gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(expand_factor)
                 self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded)
-                if self.step_input_stager is not None:
-                    # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                    self.step_input_stager.copy(
-                        self.kv_lens_expanded_cuda[:num_tokens],
-                        self.kv_lens_expanded_host[:num_tokens],
-                    )
-                else:
-                    self.kv_lens_expanded_cuda[:num_tokens].copy_(
-                        self.kv_lens_expanded_host[:num_tokens], non_blocking=True
-                    )
+                self.kv_lens_expanded_cuda[:num_tokens].copy_(
+                    self.kv_lens_expanded_host[:num_tokens], non_blocking=True
+                )
                 self._refresh_expanded_block_table(expand_factor)
         else:
             self.dsl_expand_factor = 1
@@ -2059,15 +2036,6 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None):
         self.host_block_table_expanded[:num_tokens, :max_len].copy_(
             expanded_blocks, non_blocking=True
         )
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.block_table_expanded[:num_tokens],
-                self.host_block_table_expanded[:num_tokens],
-                clamp_min=0,
-            )
-            self.step_input_stager.clamp_min(self.block_table_expanded[num_tokens:], 0)
-            return
         self.block_table_expanded[:num_tokens].copy_(
             self.host_block_table_expanded[:num_tokens], non_blocking=True
         )
@@ -2108,15 +2076,9 @@ def prepare_for_indexer_k_cache(self):
             pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1
         )
         # Copy to GPU
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.block_table[: self.num_seqs, :max_blocks_used], host_block_table
-            )
-        else:
-            self.block_table[: self.num_seqs, :max_blocks_used].copy_(
-                host_block_table, non_blocking=True
-            )
+        self.block_table[: self.num_seqs, :max_blocks_used].copy_(
+            host_block_table, non_blocking=True
+        )
 
     def prepare_for_skip_indexer(self, kv_lens: torch.Tensor):
         num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens
@@ -2182,17 +2144,9 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_cached_token_indptr[1 : self.num_generations + 1],
             )
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    self.gen_cached_token_indptr[: self.num_generations + 1],
-                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
-                )
-            else:
-                self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
-                    self.host_gen_cached_token_indptr[: self.num_generations + 1],
-                    non_blocking=True,
-                )
+            self.gen_cached_token_indptr[: self.num_generations + 1].copy_(
+                self.host_gen_cached_token_indptr[: self.num_generations + 1], non_blocking=True
+            )
             # generation kv indptr
             torch.cumsum(
                 kv_lens[self.num_contexts : self.num_seqs],
@@ -2200,16 +2154,9 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens:
                 dtype=torch.int64,
                 out=self.host_gen_kv_indptr[1 : self.num_generations + 1],
             )
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    self.gen_kv_indptr[: self.num_generations + 1],
-                    self.host_gen_kv_indptr[: self.num_generations + 1],
-                )
-            else:
-                self.gen_kv_indptr[: self.num_generations + 1].copy_(
-                    self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
-                )
+            self.gen_kv_indptr[: self.num_generations + 1].copy_(
+                self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True
+            )
         else:
             self.max_gen_seq_len = 0
 
@@ -2221,14 +2168,7 @@ def prepare_for_indices_conversion(self):
             self.seq_lens,
             dim=0,
         )
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(
-                self.req_idx_per_token[: self.num_tokens],
-                self.host_req_idx_per_token[: self.num_tokens],
-            )
-        else:
-            self.req_idx_per_token[: self.num_tokens].copy_(
-                self.host_req_idx_per_token[: self.num_tokens],
-                non_blocking=True,
-            )
+        self.req_idx_per_token[: self.num_tokens].copy_(
+            self.host_req_idx_per_token[: self.num_tokens],
+            non_blocking=True,
+        )
diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py
index 64b29917880b..c3d757fcd071 100644
--- a/tensorrt_llm/_torch/attention/backends/trtllm.py
+++ b/tensorrt_llm/_torch/attention/backends/trtllm.py
@@ -28,7 +28,6 @@
     from tensorrt_llm.mapping import Mapping
 
     from ...model_config import ModelConfig
-    from ...pyexecutor.runtime_hooks import StepInputStager
     from ...speculative.interface import SpecMetadata
     from ...speculative.spec_tree_manager import SpecTreeManager
     from .sparse.dsa.metadata import TokenMajorGenView
@@ -124,8 +123,6 @@ def generate_spec_decoding_packed_mask(max_num_requests: int,
 
 @dataclass(kw_only=True)
 class TrtllmAttentionMetadata(AttentionMetadata):
-    # A step-input stager while it records this metadata's writes, else None.
-    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     workspace: Optional[torch.Tensor] = None
     cuda_graph_workspace: Optional[torch.Tensor] = None
     workspace_reclaimable: bool = field(default=True, init=False)
@@ -1009,14 +1006,8 @@ def prepare(self) -> None:
             device='cpu',
         )
         self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens)
-        stager = self.step_input_stager
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(self.prompt_lens_cuda[:self.num_seqs],
-                        self.prompt_lens_cpu[:self.num_seqs])
-        else:
-            self.prompt_lens_cuda[:self.num_seqs].copy_(
-                self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
+        self.prompt_lens_cuda[:self.num_seqs].copy_(
+            self.prompt_lens_cpu[:self.num_seqs], non_blocking=True)
 
         # number of tokens in the kv cache for each sequence in the batch
         cached_token_lens = torch.tensor(
@@ -1069,14 +1060,9 @@ def prepare(self) -> None:
         # the sequence length including the cached tokens and the input tokens.
         self.kv_lens[:self.num_seqs].copy_(
             kv_lens + self.kv_cache_params.num_extra_kv_tokens)
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(self.kv_lens_cuda[:self.num_seqs],
-                        kv_lens[:self.num_seqs])
-        else:
-            self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
-                kv_lens[:self.num_seqs]),
-                                                    non_blocking=True)
+        self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory(
+            kv_lens[:self.num_seqs]),
+                                                non_blocking=True)
         # total kv lens for context requests and generation requests, without extra tokens
         self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item()
         self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self.
@@ -1443,42 +1429,23 @@ def prepare_context_mla_with_cached_kv(self,
                      dtype=torch.int64,
                      out=self.host_ctx_cached_token_indptr[1:self.num_contexts +
                                                            1])
-        stager = self.step_input_stager
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(
-                self.ctx_cached_token_indptr[:self.num_contexts + 1],
-                self.host_ctx_cached_token_indptr[:self.num_contexts + 1])
-        else:
-            self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
-                self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
-                non_blocking=True)
+        self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_(
+            self.host_ctx_cached_token_indptr[:self.num_contexts + 1],
+            non_blocking=True)
         torch.cumsum(
             self.seq_lens[:self.num_contexts],
             dim=0,
             dtype=torch.int64,
             out=self.host_ctx_uncached_token_indptr[1:self.num_contexts + 1])
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(
-                self.ctx_uncached_token_indptr[:self.num_contexts + 1],
-                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1])
-        else:
-            self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
-                self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
-                non_blocking=True)
+        self.ctx_uncached_token_indptr[:self.num_contexts + 1].copy_(
+            self.host_ctx_uncached_token_indptr[:self.num_contexts + 1],
+            non_blocking=True)
         torch.cumsum(kv_lens[:self.num_contexts],
                      dim=0,
                      dtype=torch.int64,
                      out=self.host_ctx_kv_indptr[1:self.num_contexts + 1])
-        if stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.copy(self.ctx_kv_indptr[:self.num_contexts + 1],
-                        self.host_ctx_kv_indptr[:self.num_contexts + 1])
-        else:
-            self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
-                self.host_ctx_kv_indptr[:self.num_contexts + 1],
-                non_blocking=True)
+        self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
+            self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
 
     def spec_decoding_param_prepare_for_blackwell(self) -> None:
         """
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 2090179caa6c..4d3b4af850e8 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -53,7 +53,6 @@
 from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens,
                            get_num_spec_layers, get_spec_decoder,
                            should_use_separate_draft_kv_cache)
-from ..speculative.spec_sampler_base import SpecSampler
 from ..utils import is_gdn_replay_enabled
 from . import config_utils
 from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer,
@@ -78,7 +77,6 @@
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
-from .runtime_hooks import runtime_hooks_of
 from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler
 from .sampler.penalties import PenaltyStore
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
@@ -4589,11 +4587,7 @@ def instantiate_sampler(
     )
     if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder(
     ):
-        sampler = get_spec_decoder(sampler_args, engine.spec_config)
-        store_writer = runtime_hooks_of(engine.model).sampler_store_writer
-        if store_writer is not None and isinstance(sampler, SpecSampler):
-            sampler.store_writer = store_writer
-        return sampler
+        return get_spec_decoder(sampler_args, engine.spec_config)
 
     if mm_encoder_only:
         # NOTE: handle model outputs specially for mm encoder executor/engine.
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index 118499f61732..e0d5abd61b35 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -160,10 +160,6 @@ class CUDAGraphRunner:
     multiple graphs, keyed by batch shape and execution-path specializations.
     """
     WARMUP_STEPS = 1
-    # Tokens of input_ids / position_ids that a step-input stager already
-    # wrote into the static inputs for the next replay (0: the replay copies
-    # them).
-    staged_static_input_tokens = 0
 
     def __init__(self, config: CUDAGraphRunnerConfig):
         self.config = config
@@ -808,12 +804,7 @@ def replay(
                 f"replay() got {seqlen} tokens for key {key}, but the graph "
                 f"was captured for {expected_num_tokens} tokens. A shorter "
                 "input_ids leaves the tail of the static input buffer stale.")
-        # A step-input stager already wrote input_ids / position_ids[:seqlen]
-        # into the static buffers.
-        staged = self.staged_static_input_tokens == seqlen
-        self.staged_static_input_tokens = 0
-        if not staged:
-            static_tensors["input_ids"][:seqlen].copy_(input_ids)
+        static_tensors["input_ids"][:seqlen].copy_(input_ids)
 
         position_ids = current_inputs["position_ids"]
         if self.config.use_mrope:
@@ -860,8 +851,7 @@ def replay(
                     f"for key {key}, but expected {expected_position_ids_shape}. "
                     "torch.Tensor.copy_() silently broadcasts mismatched shapes, "
                     "which would corrupt the static input buffer.")
-            if not staged:
-                static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
+            static_tensors["position_ids"][:, :seqlen].copy_(position_ids)
 
         num_encoder_tokens = key.num_encoder_tokens
         if num_encoder_tokens:
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
index cb24f68deef7..98dce6b41550 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py
@@ -22,7 +22,6 @@
 from .interface import ScheduledInputs
 
 if TYPE_CHECKING:
-    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm._torch.pyexecutor.sampler.sampler import SampleStateTensors
 
 
@@ -138,13 +137,11 @@ def pad_input_tokens(
     enable_attention_dp: bool,
     prefill_cuda_graph_backend: PrefillCudaGraphBackend,
     prefill_cuda_graph_num_tokens: list[int],
-    step_input_stager: "StepInputStager | None" = None,
 ) -> int:
     """Pad the batch to the token count its prefill graph or attention-DP peers need.
 
     Zeroes the padded rows of ``input_ids`` and ``position_ids``, whose last
     dimension is the token one, and returns the token count the model runs on.
-    With ``step_input_stager`` the zeroing is staged on it.
     """
     padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens = get_padding_params(
         num_tokens,
@@ -161,11 +158,6 @@ def pad_input_tokens(
         attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens
     if attn_metadata.padded_num_tokens is None:
         return num_tokens
-    if step_input_stager is not None:
-        # Staged write; see the window banner in DecoderRunner._forward_decoder.
-        step_input_stager.zero(input_ids[num_tokens:padded_num_tokens])
-        step_input_stager.zero(position_ids[..., num_tokens:padded_num_tokens])
-        return padded_num_tokens
     input_ids[num_tokens:padded_num_tokens].fill_(0)
     position_ids[..., num_tokens:padded_num_tokens].fill_(0)
     return padded_num_tokens
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
index 1dc3fca7e80f..4046fd5ca8ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py
@@ -487,14 +487,11 @@ def run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
-        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         """Run the model by replaying the graph for ``key``, capturing it first if needed.
 
         Without a key, record the model body while a breakable prefill graph
         captures, run through one when it covers the tokens, otherwise eagerly.
-        ``before_replay`` runs right before the replay of an already captured
-        graph, after the draft-cache metadata refresh.
         """
         with with_shared_pool(self.cuda_graph_runner.get_graph_pool()):
             if self.captures_model_body(key):
@@ -512,7 +509,6 @@ def run(
                 enable_spec_decode=enable_spec_decode,
                 gather_context_logits=gather_context_logits,
                 position_id_offsets=position_id_offsets,
-                before_replay=before_replay,
             )
 
     def _run(
@@ -528,7 +524,6 @@ def _run(
         enable_spec_decode: bool,
         gather_context_logits: bool,
         position_id_offsets: torch.Tensor | None,
-        before_replay: Callable[[], None] | None = None,
     ) -> Any:
         breakable_runner = self.breakable_cuda_graph_runner
         if key is None:
@@ -579,8 +574,6 @@ def capture_model(model_inputs: dict[str, Any]) -> Any:
                     attn_metadata, draft_kv_cache_manager
                 )
                 try:
-                    if before_replay is not None:
-                        before_replay()
                     with MoeLoadBalancerIterContext(moe_load_balancer):
                         outputs = self.cuda_graph_runner.replay(
                             key, inputs, position_id_offsets=position_id_offsets
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
index 0aa1d1c35c76..694fa2d122da 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
@@ -25,7 +25,7 @@
     ResourceManager,
     ResourceManagerType,
 )
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue, StepInputStager
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue
 from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase
@@ -314,8 +314,6 @@ class DecoderInputPreparer:
 
     # The model's decode prologue (RuntimeHooks.decode_prologue), set by the runner.
     decode_prologue: DecodePrologue | None = None
-    # The model's step-input stager while it records this preparer's writes, else None.
-    step_input_stager: StepInputStager | None = None
 
     def __init__(
         self,
@@ -524,7 +522,11 @@ def pad_tokens(
                 mapping=self.mapping,
                 dist=self.dist,
             )
-        padding = dict(
+        return pad_input_tokens(
+            attn_metadata,
+            num_tokens,
+            num_contexts,
+            attn_all_rank_num_tokens,
             input_ids=self._buffers.input_ids_cuda,
             position_ids=self._position_ids_buffer(),
             dist=self.dist,
@@ -536,12 +538,6 @@ def pad_tokens(
             ),
             prefill_cuda_graph_num_tokens=self._config.prefill_cuda_graph_num_tokens,
         )
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            padding["step_input_stager"] = self.step_input_stager
-        return pad_input_tokens(
-            attn_metadata, num_tokens, num_contexts, attn_all_rank_num_tokens, **padding
-        )
 
     def model_inputs(
         self,
@@ -844,18 +840,10 @@ def prepare_inputs(
         self._copy_position_ids_to_device(host)
 
         if enable_spec_decode:
-            gather_ids = torch.tensor(
-                host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()
+            self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
+                torch.tensor(host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()),
+                non_blocking=True,
             )
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)], gather_ids
-                )
-            else:
-                self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_(
-                    gather_ids, non_blocking=True
-                )
 
         self._prepare_attn_metadata(
             host,
@@ -1721,52 +1709,24 @@ def previous_seq_slots_device() -> torch.Tensor:
         assert total_num_tokens <= self._config.max_num_tokens, (
             f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self._config.max_num_tokens})"  # noqa: E501
         )
-        stager = self.step_input_stager
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids, dtype=torch.int, pin_memory=prefer_pinned())
-            if stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                stager.copy(self._buffers.input_ids_cuda[:num_tokens], input_ids)
-            else:
-                self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
+            self._buffers.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True)
 
         if num_draft_tokens > 0:
             draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned())
-            if stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                stager.copy(self._buffers.draft_tokens_cuda[: len(draft_tokens)], draft_tokens)
-            else:
-                self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
-                    draft_tokens, non_blocking=True
-                )
+            self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_(
+                draft_tokens, non_blocking=True
+            )
         if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0:
             num_accepted_draft_tokens = torch.tensor(
                 num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()
             )
-            accepted = self._buffers.num_accepted_draft_tokens_cuda[
-                : len(num_accepted_draft_tokens)
-            ]
-            if stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                stager.copy(accepted, num_accepted_draft_tokens)
-            else:
-                accepted.copy_(num_accepted_draft_tokens, non_blocking=True)
-        if next_draft_tokens_device is not None and stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            stager.overlap_inputs(
-                self._buffers,
-                previous_batch_indices,
-                previous_pos_indices,
-                num_tokens,
-                num_draft_tokens,
-                self.get_runtime_tokens_per_gen_step(runtime_draft_len),
-                len(extend_requests) - len(extend_dummy_requests),
-                new_tokens_device,
-                next_draft_tokens_device,
-                new_tokens_lens_device,
+            self._buffers.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_(
+                num_accepted_draft_tokens, non_blocking=True
             )
-        elif next_draft_tokens_device is not None:
+        if next_draft_tokens_device is not None:
             # Initialize these two values to zeros
             self._buffers.previous_pos_id_offsets_cuda.zero_()
             self._buffers.previous_kv_lens_offsets_cuda.zero_()
@@ -1866,10 +1826,6 @@ def previous_seq_slots_device() -> torch.Tensor:
         # cover the previous-batch rows; padding-request rows are covered by
         # neither, so without this the embedding would gather stale contents
         # of the uninitialized buffer, which may exceed the vocabulary size.
-        if stager is not None and host.lists.unwritten_input_id_rows:
-            # A native write to input_ids_cuda inside the window needs a flush
-            # first; see the window banner in DecoderRunner._forward_decoder.
-            stager.flush()
         _zero_fill_row_spans(self._buffers.input_ids_cuda, host.lists.unwritten_input_id_rows)
 
         host.num_tokens = num_tokens
@@ -1918,11 +1874,6 @@ def _copy_position_ids_to_device(self, host: HostInputs) -> None:
                 self._buffers.mrope_position_ids_cuda[:, :, start_idx:end_idx].copy_(
                     segment[:, :, : end_idx - start_idx], non_blocking=True
                 )
-        elif self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.position_ids(
-                self._buffers.position_ids_cuda[:total_num_tokens], host_position_ids
-            )
         else:
             self._buffers.position_ids_cuda[:total_num_tokens].copy_(
                 host_position_ids, non_blocking=True
@@ -2231,11 +2182,6 @@ def record_input_copy(self) -> None:
         self._prepare_inputs_event = torch.cuda.Event()
         self._prepare_inputs_event.record()
 
-    @property
-    def input_copy_event(self) -> torch.cuda.Event | None:
-        """The event ``record_input_copy`` recorded last."""
-        return self._prepare_inputs_event
-
     def wait_for_input_copy(self) -> None:
         if self._prepare_inputs_event is not None:
             self._prepare_inputs_event.synchronize()
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index e31f20da8e40..baacf2684be7 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -622,103 +622,34 @@ def _forward_decoder(
                     execution_requests.generation_requests,
                 )
 
-            # ==== Step-input window opens here ====
-            # When the model declares a step-input stager and begin_if_eligible
-            # returns it, every device write that prepare_inputs (and the
-            # metadata and KV-cache-manager code it calls) routes through
-            # ``step_input_stager`` (call sites marked "Staged write") is only
-            # recorded; close_input_window below, right before the graph
-            # replay, issues them all. Until then, do not read or write those
-            # buffers with any other device op: a read sees stale bytes and a
-            # write is overwritten. Route new writes through the stager or
-            # call stager.flush() first. Exception: input_ids_cuda and
-            # position_ids_cuda are also staged into the graph's static
-            # inputs, which the replay then does not refresh from them; to
-            # change them here, call stager.flush() and then write them
-            # natively. Nothing checks this contract in a production run.
-            # Contract: StepInputStager in pyexecutor/runtime_hooks.py.
-            stager = None
-            stager_source = self.runtime_hooks.step_input_stager
-            if stager_source is not None:
-                # Only the replay of an already captured graph closes the window
-                # (before_replay); an eager step, an initial capture or a warmup
-                # would run the model before the staged writes are issued, so
-                # those steps always prepare natively.
-                graph_replay = (
-                    can_run_graph
-                    and not self.cuda_graph_runner.needs_capture(key)
-                    and not self.cuda_graph_runner.is_warmup_only
-                )
-                if graph_replay:
-                    stager = stager_source.begin_if_eligible(
-                        self._inputs,
-                        execution_requests,
-                        attn_metadata,
-                        spec_metadata,
-                        new_tensors_device,
-                        graph_replay=graph_replay,
-                        static_inputs=self.cuda_graph_runner.shared_static_tensors,
-                        is_dummy=is_dummy,
-                        enable_spec_decode=enable_spec_decode,
-                        promoted=bool(execution_promoted_context_ids),
-                    )
-            try:
-                # The window spans the scope: the scope is a no-op whenever
-                # speculative decoding is on, which every stager requires.
-                with self._cuda_graph_metadata_prepare_scope(
-                    attn_metadata,
-                    key,
+            with self._cuda_graph_metadata_prepare_scope(
+                attn_metadata,
+                key,
+                execution_requests,
+                execution_promoted_context_ids,
+                enable_spec_decode=enable_spec_decode,
+            ):
+                prepared = self._inputs.prepare_inputs(
                     execution_requests,
+                    kv_cache_manager,
+                    attn_metadata,
+                    spec_metadata,
+                    new_tensors_device,
+                    cache_indirection_buffer,
+                    resource_manager,
+                    can_run_graph,
                     execution_promoted_context_ids,
+                    use_lora_graph=use_lora_graph,
                     enable_spec_decode=enable_spec_decode,
-                ):
-                    prepared = self._inputs.prepare_inputs(
-                        execution_requests,
-                        kv_cache_manager,
-                        attn_metadata,
-                        spec_metadata,
-                        new_tensors_device,
-                        cache_indirection_buffer,
-                        resource_manager,
-                        can_run_graph,
-                        execution_promoted_context_ids,
-                        use_lora_graph=use_lora_graph,
-                        enable_spec_decode=enable_spec_decode,
-                        runtime_draft_len=runtime_draft_len,
-                        is_dummy=is_dummy,
-                        gather_context_logits=gather_context_logits,
-                    )
-            except BaseException:
-                if stager is not None:
-                    stager.discard()
-                raise
+                    runtime_draft_len=runtime_draft_len,
+                    is_dummy=is_dummy,
+                    gather_context_logits=gather_context_logits,
+                )
             inputs = prepared.kwargs
             runtime_draft_len = prepared.runtime_draft_len
             if execution_promoted_context_ids:
                 self._inputs.record_promoted_iter_states(scheduled_requests)
-            close_input_window = None
-            window_open = stager is not None
-            if stager is None:
-                self._inputs.record_input_copy()
-            else:
-
-                def close_input_window() -> None:
-                    nonlocal window_open
-                    window_open = False
-                    # ==== Step-input window closes here ====
-                    # end() issues every staged write; nothing between
-                    # begin_if_eligible above and this call may have touched
-                    # the staged buffers with other device ops. The input-copy
-                    # event must follow the launch (host page-table reads
-                    # happen at execution time). A nonzero
-                    # static_inputs_tokens() makes the replay skip its
-                    # input_ids / position_ids copy.
-                    stager.end(record_event=False)
-                    self._inputs.record_input_copy()
-                    stager.set_completion_event(self._inputs.input_copy_event)
-                    self.cuda_graph_runner.staged_static_input_tokens = (
-                        stager.static_inputs_tokens()
-                    )
+            self._inputs.record_input_copy()
 
             def run_model(model_inputs: dict[str, Any]) -> Any:
                 with MoeLoadBalancerIterContext(moe_load_balancer):
@@ -740,30 +671,24 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None:
 
             graph_key = key if can_run_graph else None
             captures_model_body = self._graphs.captures_model_body(graph_key)
-            try:
-                outputs = self._graphs.run(
-                    graph_key,
-                    inputs,
-                    run_model,
-                    postprocess_inputs,
-                    attn_metadata=attn_metadata,
-                    draft_kv_cache_manager=draft_kv_cache_manager,
-                    moe_load_balancer=moe_load_balancer,
-                    enable_spec_decode=enable_spec_decode,
-                    gather_context_logits=gather_context_logits,
-                    position_id_offsets=(
-                        self._buffers.previous_pos_id_offsets_cuda
-                        if enable_spec_decode
-                        and not self._config.disable_overlap_scheduler
-                        and attn_metadata.kv_cache_manager is not None
-                        else None
-                    ),
-                    before_replay=close_input_window,
-                )
-            except BaseException:
-                if window_open:
-                    stager.discard()
-                raise
+            outputs = self._graphs.run(
+                graph_key,
+                inputs,
+                run_model,
+                postprocess_inputs,
+                attn_metadata=attn_metadata,
+                draft_kv_cache_manager=draft_kv_cache_manager,
+                moe_load_balancer=moe_load_balancer,
+                enable_spec_decode=enable_spec_decode,
+                gather_context_logits=gather_context_logits,
+                position_id_offsets=(
+                    self._buffers.previous_pos_id_offsets_cuda
+                    if enable_spec_decode
+                    and not self._config.disable_overlap_scheduler
+                    and attn_metadata.kv_cache_manager is not None
+                    else None
+                ),
+            )
             if captures_model_body:
                 return outputs
 
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 c825736a372c..07a5e93717aa 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
@@ -134,7 +134,6 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
-    from tensorrt_llm._torch.pyexecutor.runtime_hooks import StepInputStager
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS
@@ -1250,8 +1249,6 @@ def _spec_recompute_claim_limit(
 class KVCacheManagerV2(BaseResourceManager):
     # Sparse managers attach a metadata batch during initialization.
     sparse_metadata_batch: Batch | None = None
-    # A step-input stager while it records this manager's writes, else None.
-    step_input_stager: Optional["StepInputStager"] = None
     # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction.
     # Declared on the class so it is present even when an instance is built without running __init__.
     _cold_pool_group_membership_cache: Optional[tuple[tuple[int, frozenset[int]], ...]] = None
@@ -6757,16 +6754,6 @@ def copy_batch_block_offsets(
             )
             return
 
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.block_offsets(
-                self.host_kv_cache_block_offsets,
-                dst_tensor,
-                copy_idx.tolist(),
-                self._index_scale_ints,
-                self.kv_offset.tolist(),
-            )
-            return
         copy_batch_block_offsets_to_device(
             self.host_kv_cache_block_offsets,
             dst_tensor,
diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
index 72987309443f..772845919949 100644
--- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
+++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py
@@ -4,8 +4,8 @@
 
 A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class
 attribute or a property). Every field defaults to off, and a model that declares
-nothing gets ``RuntimeHooks()``: the step loop, the samplers and the attention
-metadata then run exactly as they do without this module. Every engine call site
+nothing gets ``RuntimeHooks()``: the step loop then runs exactly as it does
+without this module. Every engine call site
 of a hook runs its original statements when the hook is off.
 
 The hooks change *where* and *how* the step loop moves data, never what a
@@ -16,151 +16,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Protocol, Sequence
-
-if TYPE_CHECKING:
-    import torch
-
-
-class StepInputStager(Protocol):
-    """Records a decode step's input preparation and issues it as one launch.
-
-    ``begin_if_eligible`` opens a deferred-write window on a step it supports
-    and returns the stager, or returns None and the engine prepares the step
-    natively. While the window is open, the engine's call sites route their
-    device writes to the staging methods instead of issuing them. ``end``
-    issues every staged write in staging order, then the deferred native
-    operations; ``discard`` drops them after a failed preparation.
-
-    While its window is open the stager sets itself as ``step_input_stager`` on
-    the decoder's input preparer, the attention metadata, the speculative
-    metadata and the KV cache managers, and resets that attribute to None at
-    ``end`` and ``discard``. Every call site reads the attribute and runs its
-    original statement when it is None.
-    """
-
-    def begin_if_eligible(
-        self,
-        preparer: Any,
-        requests: Any,
-        attn_metadata: Any,
-        spec_metadata: Any,
-        new_tensors_device: Any,
-        *,
-        graph_replay: bool,
-        static_inputs: Optional[Dict[str, "torch.Tensor"]],
-        is_dummy: bool,
-        enable_spec_decode: bool,
-        promoted: bool,
-    ) -> Optional["StepInputStager"]:
-        """Open the window for this step, or return None to keep the native preparation.
-
-        The engine calls this only for the replay of an already captured CUDA
-        graph, so ``graph_replay`` is always True: the window closes right
-        before the replay, and an eager step, an initial capture or a warmup
-        would run the model before the staged writes are issued. Those steps
-        always prepare natively.
-
-        ``preparer`` is the ``DecoderInputPreparer`` whose writes the window
-        records. ``static_inputs`` are the static input tensors of the CUDA
-        graph the step replays.
-        """
-
-    def copy(
-        self, dst: "torch.Tensor", src: "torch.Tensor", clamp_min: Optional[int] = None
-    ) -> None:
-        """Stage ``dst.copy_(src, non_blocking=True)`` for a host ``src``."""
-
-    def zero(self, dst: "torch.Tensor") -> None:
-        """Stage ``dst.zero_()``."""
-
-    def clamp_min(self, dst: "torch.Tensor", lo: int) -> None:
-        """Stage ``dst.clamp_(min=lo)``."""
-
-    def block_offsets(
-        self,
-        host_table: "torch.Tensor",
-        dst: "torch.Tensor",
-        copy_index: Sequence[int],
-        index_scales: Sequence[int],
-        kv_offsets: Sequence[int],
-    ) -> None:
-        """Stage the KV cache manager's batch copy of block offsets to the device."""
-
-    def position_ids(self, dst: "torch.Tensor", host_position_ids: "torch.Tensor") -> None:
-        """Stage the position-id upload and the CUDA graph's static copy of it."""
-
-    def overlap_inputs(
-        self,
-        buffers: Any,
-        previous_batch_indices: Sequence[int],
-        previous_pos_indices: Sequence[int],
-        num_tokens: int,
-        num_draft_tokens: int,
-        tokens_per_gen_step: int,
-        num_extend_requests: int,
-        new_tokens: "torch.Tensor",
-        next_draft_tokens: "torch.Tensor",
-        new_tokens_lens: "torch.Tensor",
-    ) -> None:
-        """Stage the overlap scheduler's gathers of the previous step's outputs into ``buffers``."""
-
-    def scheduler_metadata(self, metadata: Any) -> bool:
-        """Stage the sparse-attention indexer's scheduler metadata.
-
-        Returns False, having staged nothing, when the caller must defer the
-        native computation instead.
-        """
-
-    def defer(self, fn: Callable[[], None]) -> None:
-        """Run ``fn`` (native device work that reads staged values) after the launch."""
-
-    def flush(self) -> None:
-        """Issue every write staged so far, before the window closes."""
-
-    def end(self, record_event: bool = True) -> None:
-        """Close the window and issue the staged preparation."""
-
-    def discard(self) -> None:
-        """Close the window after a failed preparation; the staged writes are dropped."""
-
-    def set_completion_event(self, event: "torch.cuda.Event") -> None:
-        """Event the caller recorded right after ``end(record_event=False)``."""
-
-    def static_inputs_tokens(self) -> int:
-        """Tokens of the CUDA graph's static inputs the last launch wrote (0: none)."""
-
-
-class SamplerStoreWriter(Protocol):
-    """Stores a speculative sampler step's outputs and copies the store to the host in one launch."""
-
-    def supported(
-        self,
-        store: Any,
-        new_tokens: "torch.Tensor",
-        next_new_tokens: "torch.Tensor",
-        new_tokens_lens: "torch.Tensor",
-        next_draft_tokens: "torch.Tensor",
-    ) -> bool:
-        """Whether this step's output tensors fit the writer; False keeps the native stores."""
-
-    def store_and_copy_out(
-        self,
-        store: Any,
-        slots: Sequence[int],
-        new_tokens: "torch.Tensor",
-        next_new_tokens: "torch.Tensor",
-        new_tokens_lens: "torch.Tensor",
-        next_draft_tokens: "torch.Tensor",
-        host_new_tokens: "torch.Tensor",
-        host_new_tokens_lens: "torch.Tensor",
-        host_next_draft_tokens: "torch.Tensor",
-        record_event: bool = True,
-    ) -> None:
-        """Scatter the outputs into the store rows ``slots`` and copy the store to the host tensors."""
-
-    def set_completion_event(self, event: "torch.cuda.Event") -> None:
-        """Event the caller recorded right after ``store_and_copy_out(record_event=False)``."""
+from typing import Any, Optional, Protocol
 
 
 class DecodePrologue(Protocol):
@@ -198,10 +54,6 @@ class RuntimeHooks:
     # Run the overlap executor loop with the execution stream current, so the
     # loop's own stream work and the forward are ordered on one stream.
     loop_on_execution_stream: bool = False
-    # Decode input preparation recorded and issued as one launch (see StepInputStager).
-    step_input_stager: Optional[StepInputStager] = None
-    # The speculative sampler's store update and host copy in one launch (see SamplerStoreWriter).
-    sampler_store_writer: Optional[SamplerStoreWriter] = None
     # The decode step's input prologue in one launch (see DecodePrologue).
     decode_prologue: Optional[DecodePrologue] = None
     # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents).
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 0328e29746cc..15993e5c26ba 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -348,13 +348,8 @@ def prepare(self):
                                      dtype=torch.int,
                                      device='cpu',
                                      pin_memory=prefer_pinned())
-        if self.step_input_stager is not None:
-            # Staged write; see the window banner in DecoderRunner._forward_decoder.
-            self.step_input_stager.copy(self.batch_indices_cuda[:num_seqs],
-                                        batch_indices)
-        else:
-            self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
-                                                     non_blocking=True)
+        self.batch_indices_cuda[:num_seqs].copy_(batch_indices,
+                                                 non_blocking=True)
         # Snapshot the real (unpadded) scheduled token count before the decrement
         # below. maybe_capture_hidden_states runs during the target forward and
         # must bound the capture by this count; self.num_tokens is about to be
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index a8bdaf4d4589..8edccc3e3f8a 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import TYPE_CHECKING, ClassVar, List, Optional, Protocol, Type
+from typing import TYPE_CHECKING, List, Optional, Protocol, Type
 
 import torch
 from packaging.version import Version
@@ -37,7 +37,6 @@
 if TYPE_CHECKING:
     from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
     from ..pyexecutor.llm_request import LlmRequest
-    from ..pyexecutor.runtime_hooks import StepInputStager
 
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer
@@ -460,8 +459,6 @@ class SpecMetadata:
     """
     Metadata for speculative decoding.
     """
-    # A step-input stager while it records this metadata's writes, else None.
-    step_input_stager: ClassVar[Optional["StepInputStager"]] = None
     # The max number of requests in a single batch.
     max_num_requests: int
     # The number of draft layers. (Also the number of draft tokens for the linear tree.)
@@ -884,11 +881,6 @@ def _populate_request_rng_state(
                                        device='cuda')
 
         def _upload(dst: torch.Tensor, values: list[int]) -> None:
-            if self.step_input_stager is not None:
-                # Staged write; see the window banner in DecoderRunner._forward_decoder.
-                self.step_input_stager.copy(
-                    dst[:len(values)], torch.tensor(values, dtype=torch.int64))
-                return
             dst[:len(values)].copy_(torch.tensor(values,
                                                  dtype=torch.int64,
                                                  pin_memory=prefer_pinned()),
diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
index dc40cca3a7c5..5fb5b185dc0b 100644
--- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py
+++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py
@@ -27,11 +27,9 @@
 
 import torch
 
-from ..._utils import prefer_pinned
 from ...sampling_params import SamplingParams
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from ..pyexecutor.resource_manager import BaseResourceManager
-from ..pyexecutor.runtime_hooks import SamplerStoreWriter
 from ..pyexecutor.sampler import (
     DEFAULT_BEAM_IDX,
     AsyncWorkerMixin,
@@ -169,10 +167,6 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin):
 
     SampleState = SampleStateSpec
 
-    # The model's one-launch store writer, bound once by instantiate_sampler; None keeps
-    # the native stores and host copies.
-    store_writer: Optional[SamplerStoreWriter] = None
-
     def is_generation_model(self) -> bool:
         return True
 
@@ -448,11 +442,8 @@ def sample_async(
             for r in sampling_requests
         ]
 
-        store_writer = self.store_writer
-        slots = None
-        if store_writer is None:
-            slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
-            slots = slots.to(device="cuda", non_blocking=True)
+        slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
+        slots = slots.to(device="cuda", non_blocking=True)
 
         o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests]
         o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests]
@@ -523,28 +514,11 @@ def sample_async(
         elif o_next_new_tokens.shape[1] > next_new_tokens_width:
             o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width]
 
-        # The writer stores the four outputs and copies the three store tensors
-        # only; a step that publishes verify_lens keeps the native path.
-        staged_store = (
-            store_writer is not None
-            and o_verify_lens is None
-            and prefer_pinned()
-            and not self._async_worker_active()
-            and store_writer.supported(
-                self.store, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens
-            )
-        )
-        if not staged_store:
-            if slots is None:
-                slots = torch.as_tensor(
-                    [r.py_seq_slot for r in sampling_requests], dtype=torch.long
-                )
-                slots = slots.to(device="cuda", non_blocking=True)
-            # Use index_copy_ for efficient copying (slots are unique)
-            self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
-            self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
-            self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
-            self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
+        # Use index_copy_ for efficient copying (slots are unique)
+        self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
+        self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
+        self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
+        self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
 
         # Create sample state with async D2H copy
         device_tensors = SampleStateTensorsSpec(
@@ -553,43 +527,13 @@ def sample_async(
             next_draft_tokens=self.store.next_draft_tokens,
         )
 
-        if staged_store:
-            # staged_store implies prefer_pinned(), so these buffers are pinned.
-            host_tensors = SampleStateTensorsSpec(
-                new_tokens=torch.empty_like(
-                    self.store.new_tokens, device="cpu", pin_memory=prefer_pinned()
-                ),
-                new_tokens_lens=torch.empty_like(
-                    self.store.new_tokens_lens, device="cpu", pin_memory=prefer_pinned()
-                ),
-                next_draft_tokens=torch.empty_like(
-                    self.store.next_draft_tokens, device="cpu", pin_memory=prefer_pinned()
-                ),
-            )
-            store_writer.store_and_copy_out(
-                self.store,
-                [r.py_seq_slot for r in sampling_requests],
-                o_new_tokens,
-                o_next_new_tokens,
-                o_new_tokens_lens,
-                o_next_draft_tokens,
-                host_tensors.new_tokens,
-                host_tensors.new_tokens_lens,
-                host_tensors.next_draft_tokens,
-                record_event=False,
-            )
-        else:
-            host_tensors = SampleStateTensorsSpec(
-                new_tokens=self._copy_to_host(self.store.new_tokens),
-                new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
-                next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
-                verify_lens=(
-                    self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None
-                ),
-            )
+        host_tensors = SampleStateTensorsSpec(
+            new_tokens=self._copy_to_host(self.store.new_tokens),
+            new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens),
+            next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens),
+            verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None),
+        )
         sampler_event = self._record_sampler_event()
-        if staged_store:
-            store_writer.set_completion_event(sampler_event.cuda_event)
 
         # Add dummy draft tokens to context requests for KV cache preparation
         for request in finished_context_requests:
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index 95745dd901db..f3a79f0e1180 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -3,11 +3,11 @@
 """How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).
 
 The engine's call sites run their original statements when a hook is off: every
-object a hook attaches to declares it off at class level, every call a call
-site makes on a hook is a method of the hook's protocol, every read of a hook in
-the engine files is decided by its all-off value toward the original statements
-(checked on the source), and call sites of each hook run on stand-ins issue
-their original work with the hook off and hand it to the hook when it is on.
+object that holds a hook declares it off at class level, every call a call site
+makes on a hook is a method of the hook's protocol, every read of a hook in the
+engine files is decided by its all-off value toward the original statements
+(checked on the source), and call sites of each hook run on stand-ins with the
+hook off and on.
 """
 
 import ast
@@ -25,8 +25,6 @@
     DEFAULT_RUNTIME_HOOKS,
     ExecutorEvents,
     RuntimeHooks,
-    SamplerStoreWriter,
-    StepInputStager,
     runtime_hooks_of,
 )
 
@@ -59,8 +57,6 @@ def test_a_model_without_hooks_gets_the_all_off_default(model):
     assert hooks is DEFAULT_RUNTIME_HOOKS
     assert hooks == RuntimeHooks()
     assert not hooks.loop_on_execution_stream
-    assert hooks.step_input_stager is None
-    assert hooks.sampler_store_writer is None
     assert hooks.decode_prologue is None
     assert hooks.executor_events is None
 
@@ -117,36 +113,14 @@ def test_hooks_are_immutable_once_declared():
         hooks.loop_on_execution_stream = True
 
 
-def test_engine_classes_declare_the_stager_off():
-    """Every object a hook attaches to declares the attribute off at class level."""
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
-        DSAtrtllmAttentionMetadata,
-    )
-    from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
-    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner
+def test_engine_classes_declare_the_hooks_off():
+    """Every object that holds a hook declares it off at class level."""
     from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer
     from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
-    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
     from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
-    from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata
-    from tensorrt_llm._torch.speculative.interface import SpecMetadata
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    for cls in (
-        TrtllmAttentionMetadata,
-        DSAtrtllmAttentionMetadata,
-        SpecMetadata,
-        Eagle3OneModelSpecMetadata,
-    ):
-        assert cls.step_input_stager is None, cls
-        # A class attribute, not a dataclass field: the constructors do not take it.
-        assert "step_input_stager" not in {f.name for f in dataclasses.fields(cls)}, cls
-    assert DecoderInputPreparer.step_input_stager is None
+
     assert DecoderInputPreparer.decode_prologue is None
     assert DecoderRunner.runtime_hooks is DEFAULT_RUNTIME_HOOKS
-    assert KVCacheManagerV2.step_input_stager is None
-    assert CUDAGraphRunner.staged_static_input_tokens == 0
-    assert SpecSampler.store_writer is None
     assert PyExecutor._execution_stream_loop_armed is False
     assert PyExecutor._loop_on_execution_stream is False
     assert PyExecutor._executor_events is None
@@ -155,26 +129,6 @@ def test_engine_classes_declare_the_stager_off():
 # Per hook protocol: the local names and attribute names a call site holds the hook in, and
 # the engine files with call sites.
 _HOOK_CONSUMERS = {
-    StepInputStager: (
-        ("stager", "stager_source", "step_input_stager"),
-        ("step_input_stager",),
-        (
-            "pyexecutor/engine/runners/decoder/runner.py",
-            "pyexecutor/engine/runners/decoder/inputs.py",
-            "pyexecutor/engine/runners/common.py",
-            "pyexecutor/kv_cache/kv_cache_manager_v2.py",
-            "attention/backends/trtllm.py",
-            "attention/backends/sparse/dsa/metadata.py",
-            "attention/backends/sparse/dsa/indexer.py",
-            "speculative/eagle3.py",
-            "speculative/interface.py",
-        ),
-    ),
-    SamplerStoreWriter: (
-        ("store_writer",),
-        ("store_writer",),
-        ("speculative/spec_sampler_base.py",),
-    ),
     ExecutorEvents: ((), ("_executor_events",), ("pyexecutor/py_executor.py",)),
 }
 
@@ -218,63 +172,6 @@ def test_every_hook_call_site_uses_its_protocol(protocol):
     assert used == methods
 
 
-class _StoreWriterModel(torch.nn.Module):
-    def __init__(self, writer) -> None:
-        super().__init__()
-        self.runtime_hooks = RuntimeHooks(sampler_store_writer=writer)
-
-
-def _instantiate_sampler(monkeypatch, model, sampler):
-    """``instantiate_sampler`` for a speculative engine whose spec decoder is ``sampler``."""
-    from tensorrt_llm._torch.pyexecutor import _util
-
-    monkeypatch.setattr(_util, "confidential_compute_enabled", lambda: False)
-    monkeypatch.setattr(_util, "resolve_max_num_sequences", lambda *args, **kwargs: 8)
-    monkeypatch.setattr(_util, "create_torch_sampler_args", lambda **kwargs: "sampler args")
-    monkeypatch.setattr(_util, "get_spec_decoder", lambda args, spec_config: sampler)
-    engine = types.SimpleNamespace(
-        max_seq_len=64,
-        spec_config=types.SimpleNamespace(
-            spec_dec_mode=types.SimpleNamespace(has_spec_decoder=lambda: True)
-        ),
-        model=model,
-    )
-    llm_args = types.SimpleNamespace(
-        sampler_force_async_worker=False,
-        disable_overlap_scheduler=False,
-        enable_speculative_beam_history_d2h=False,
-    )
-    return _util.instantiate_sampler(
-        engine,
-        llm_args,
-        None,
-        max_batch_size=8,
-        max_beam_width=1,
-        mm_encoder_only=False,
-        speculative_config=None,
-    )
-
-
-def test_instantiate_sampler_binds_no_writer_without_hooks(monkeypatch):
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    sampler = SpecSampler.__new__(SpecSampler)
-
-    assert _instantiate_sampler(monkeypatch, _PlainModel(), sampler) is sampler
-    assert "store_writer" not in vars(sampler)
-    assert sampler.store_writer is None
-
-
-def test_instantiate_sampler_binds_the_declared_writer(monkeypatch):
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    writer = object()
-    sampler = SpecSampler.__new__(SpecSampler)
-
-    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(writer), sampler) is sampler
-    assert sampler.store_writer is writer
-
-
 def test_executor_events_protocol():
     assert _protocol_methods(ExecutorEvents) == {
         "after_warmup",
@@ -371,13 +268,6 @@ def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value
     assert recorder.entered == []
 
 
-def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
-    sampler = types.SimpleNamespace()
-
-    assert _instantiate_sampler(monkeypatch, _StoreWriterModel(object()), sampler) is sampler
-    assert "store_writer" not in vars(sampler)
-
-
 # ---------------------------------------------------------------------------------------------
 # With every hook off, every call site runs its original statements.
 #
@@ -390,33 +280,20 @@ def test_instantiate_sampler_leaves_other_spec_decoders_alone(monkeypatch):
 # ---------------------------------------------------------------------------------------------
 
 _SEAM_CONSUMERS = (
-    "pyexecutor/engine/runners/decoder/runner.py",
     "pyexecutor/engine/runners/decoder/inputs.py",
-    "pyexecutor/cuda_graph_runner.py",
-    "pyexecutor/kv_cache/kv_cache_manager_v2.py",
     "pyexecutor/py_executor.py",
-    "pyexecutor/_util.py",
-    "attention/backends/trtllm.py",
-    "attention/backends/sparse/dsa/metadata.py",
-    "attention/backends/sparse/dsa/indexer.py",
-    "speculative/eagle3.py",
-    "speculative/interface.py",
-    "speculative/spec_sampler_base.py",
 )
 
 # Attributes that hold a hook, and their value when every hook is off.
 _SEAM_ATTRS = {
-    "step_input_stager": None,
-    "store_writer": None,
     "_executor_events": None,
     "_loop_on_execution_stream": False,
     "_execution_stream_loop_armed": False,
-    "staged_static_input_tokens": 0,
     "runtime_hooks": DEFAULT_RUNTIME_HOOKS,
     "decode_prologue": None,
 }
 # Locals the call sites hold a hook (or a decision on one) in.
-_SEAM_LOCALS = ("stager", "stager_source", "store_writer", "prologue", "staged", "staged_store")
+_SEAM_LOCALS = ("prologue",)
 _HOOK_FIELDS = {field.name for field in dataclasses.fields(RuntimeHooks)}
 
 
@@ -441,15 +318,10 @@ def report(self, node: ast.AST, what: str) -> None:
     # -- reads of a hook ------------------------------------------------------------------
     @staticmethod
     def _is_seam_call(node: ast.AST) -> bool:
-        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)):
-            return False
-        if node.func.id == "runtime_hooks_of":
-            return True
         return (
-            node.func.id == "getattr"
-            and len(node.args) >= 2
-            and isinstance(node.args[1], ast.Constant)
-            and node.args[1].value == "step_input_stager"
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "runtime_hooks_of"
         )
 
     def reads_seam(self, node: ast.AST, env: dict) -> bool:
@@ -476,7 +348,7 @@ def value(self, node: ast.AST, env: dict):
                 return _SEAM_ATTRS[node.attr]
             return _UNKNOWN
         if self._is_seam_call(node):
-            return DEFAULT_RUNTIME_HOOKS if node.func.id == "runtime_hooks_of" else None
+            return DEFAULT_RUNTIME_HOOKS
         if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
             operand = self.value(node.operand, env)
             return _UNKNOWN if operand is _UNKNOWN else not operand
@@ -496,13 +368,6 @@ def value(self, node: ast.AST, env: dict):
             if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN:
                 if isinstance(node.comparators[0], ast.Constant) and right is None:
                     return (left is None) == isinstance(op, ast.Is)
-            if (
-                isinstance(op, ast.Eq)
-                and isinstance(node.left, ast.Attribute)
-                and node.left.attr == "staged_static_input_tokens"
-            ):
-                # No tokens were staged; a graph replay has at least one token.
-                return False
         return _UNKNOWN
 
     # -- statements -------------------------------------------------------------------------
@@ -598,15 +463,15 @@ def test_every_seam_read_has_a_native_branch(rel):
 def test_the_seam_check_reports_undecided_reads():
     source = """
 def step(self, x):
-    stager = self.step_input_stager
-    if stager is not None:
-        stager.copy(x)
+    events = self._executor_events
+    if events is not None:
+        events.after_warmup()
     else:
         x.copy_()
-    stager.flush()
-    if x or stager:
+    events.at_shutdown()
+    if x or events:
         pass
-    self.step_input_stager = x
+    self._executor_events = x
 """
     walk = _AllOffWalk("sample.py")
     walk.block(ast.parse(source).body, {})
@@ -815,584 +680,3 @@ def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(complete
     executor._executor_loop_cleanup()
     assert executor.is_shutdown
     assert events.calls == (["at_shutdown"] if completed else [])
-
-
-# -- the step-input stager ------------------------------------------------------------------
-
-
-class _RecordingStager:
-    def __init__(self) -> None:
-        self.calls = []
-
-    def __getattr__(self, name):
-        return lambda *args, **kwargs: self.calls.append((name, args, kwargs))
-
-
-def _block_offsets_manager(monkeypatch, log: list):
-    """A ``KVCacheManagerV2`` with only what ``copy_batch_block_offsets`` reads (one page table)."""
-    from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2
-    from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
-
-    monkeypatch.setattr(
-        kv_cache_manager_v2,
-        "copy_batch_block_offsets_to_device",
-        lambda *args: log.append(("native", args)),
-    )
-    manager = KVCacheManagerV2.__new__(KVCacheManagerV2)
-    manager._publish_sparse_metadata = lambda: None
-    manager.max_copy_beam_width = 1
-    manager.kv_cache_type = None
-    manager._use_per_layer_page_tables = False
-    manager.index_mapper = types.SimpleNamespace(
-        get_copy_index=lambda request_ids, num_contexts, beam_width, cross: torch.tensor([1, 0])
-    )
-    manager.host_kv_cache_block_offsets = torch.zeros(2, 4, dtype=torch.int32)
-    manager.index_scales = torch.ones(1, dtype=torch.int32)
-    manager._index_scale_ints = [1]
-    manager.kv_offset = torch.zeros(1, dtype=torch.int32)
-    manager._stream = types.SimpleNamespace(cuda_stream=7)
-    return manager
-
-
-def test_kv_cache_block_offsets_copy_natively_without_a_stager(monkeypatch):
-    log = []
-    manager = _block_offsets_manager(monkeypatch, log)
-    dst = torch.empty(2, 4, dtype=torch.int32)
-
-    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
-    assert len(log) == 1
-    kind, args = log[0]
-    assert kind == "native"
-    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
-    assert args[2].tolist() == [1, 0]
-    assert args[3] is manager.index_scales and args[4] is manager.kv_offset and args[5] == 7
-
-
-def test_kv_cache_block_offsets_are_staged_with_a_stager(monkeypatch):
-    log = []
-    manager = _block_offsets_manager(monkeypatch, log)
-    manager.step_input_stager = _RecordingStager()
-    dst = torch.empty(2, 4, dtype=torch.int32)
-
-    manager.copy_batch_block_offsets(dst, [3, 4], 1, 0, 2)
-    assert log == []
-    ((name, args, _),) = manager.step_input_stager.calls
-    assert name == "block_offsets"
-    assert args[0] is manager.host_kv_cache_block_offsets and args[1] is dst
-    assert args[2:] == ([1, 0], [1], [0])
-
-
-def _slot_mapping_inputs(tokens: int):
-    metadata = types.SimpleNamespace(
-        host_indexer_k_cache_block_offsets=torch.tensor([[0, 1]], dtype=torch.int32),
-        host_slot_mapping_fp8=torch.zeros(tokens, dtype=torch.int64),
-        host_slot_mapping_scale=torch.zeros(tokens, dtype=torch.int64),
-        slot_mapping_fp8=torch.full((tokens,), -1, dtype=torch.int64),
-        slot_mapping_scale=torch.full((tokens,), -1, dtype=torch.int64),
-    )
-    params = types.SimpleNamespace(
-        batch_size=1,
-        tokens_per_block=4,
-        head_dim=8,
-        new_kv_tokens=torch.tensor([tokens]),
-        data_bytes_per_token=8,
-        cached_kv_tokens=torch.tensor([1]),
-        quant_block_size=8,
-    )
-    return metadata, params
-
-
-def test_indexer_slot_mappings_copy_natively_without_a_stager():
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-
-    metadata, params = _slot_mapping_inputs(3)
-
-    Indexer.recompute_slot_mappings(metadata, params)
-    assert torch.equal(metadata.slot_mapping_fp8, metadata.host_slot_mapping_fp8)
-    assert torch.equal(metadata.slot_mapping_scale, metadata.host_slot_mapping_scale)
-
-
-def test_indexer_slot_mappings_are_staged_with_a_stager():
-    from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer
-
-    metadata, params = _slot_mapping_inputs(3)
-    metadata.step_input_stager = _RecordingStager()
-
-    Indexer.recompute_slot_mappings(metadata, params)
-    assert [name for name, _, _ in metadata.step_input_stager.calls] == ["copy", "copy"]
-    assert bool((metadata.slot_mapping_fp8 == -1).all())
-    assert bool((metadata.slot_mapping_scale == -1).all())
-
-
-# -- the step-input window in the decoder runner -------------------------------------------------
-
-
-class _ReferenceStager:
-    """A step-input stager that accepts every step it is offered and logs the window's lifecycle."""
-
-    def __init__(self, log: list, accept: bool = True, tokens: int = 3) -> None:
-        self.log, self.accept, self.tokens = log, accept, tokens
-
-    def begin_if_eligible(
-        self,
-        preparer,
-        requests,
-        attn_metadata,
-        spec_metadata,
-        new_tensors_device,
-        *,
-        graph_replay,
-        static_inputs,
-        is_dummy,
-        enable_spec_decode,
-        promoted,
-    ):
-        self.log.append(("begin", graph_replay, static_inputs))
-        return self if self.accept else None
-
-    def end(self, record_event: bool = True) -> None:
-        self.log.append(("end", record_event))
-
-    def discard(self) -> None:
-        self.log.append("discard")
-
-    def set_completion_event(self, event) -> None:
-        self.log.append(("completion_event", event))
-
-    def static_inputs_tokens(self) -> int:
-        return self.tokens
-
-
-class _WindowGraphs:
-    """The ``DecoderCudaGraphs`` calls ``_forward_decoder`` makes, with one selected graph key.
-
-    ``run`` follows ``DecoderCudaGraphs._run``: ``before_replay`` runs only right before the
-    replay of an already captured graph, and the replay consumes the staged static inputs.
-    """
-
-    def __init__(self, log: list, key, *, needs_capture=False, warmup_only=False, fail=None):
-        self.log, self.key, self.fail = log, key, fail
-        self.cuda_graph_runner = types.SimpleNamespace(
-            needs_capture=lambda key: needs_capture,
-            is_warmup_only=warmup_only,
-            shared_static_tensors={"input_ids": torch.zeros(4, dtype=torch.int64)},
-            staged_static_input_tokens=0,
-        )
-
-    def attention_dp_execution_batch(self, requests):
-        return requests
-
-    def promote_final_contexts(self, requests, kv_cache_manager, **kwargs):
-        return requests, frozenset()
-
-    def pad_batch(self, requests, resource_manager, runtime_draft_len):
-        return contextlib.nullcontext(requests)
-
-    def select(self, requests, resource_manager, **kwargs):
-        return types.SimpleNamespace(
-            key=self.key,
-            use_lora_graph=False,
-            attn_metadata=kwargs["attn_metadata"],
-            spec_metadata=None,
-        )
-
-    def captures_model_body(self, key) -> bool:
-        return False
-
-    def run(self, key, inputs, run_model, postprocess_inputs, *, before_replay=None, **kwargs):
-        runner = self.cuda_graph_runner
-        if self.fail == "before_replay":
-            raise RuntimeError("graph run failed")
-        if key is None:
-            self.log.append("eager")
-        elif runner.needs_capture(key) or runner.is_warmup_only:
-            self.log.append("capture")
-        else:
-            if before_replay is not None:
-                before_replay()
-            self.log.append(("replay", runner.staged_static_input_tokens))
-            runner.staged_static_input_tokens = 0
-            if self.fail == "replay":
-                raise RuntimeError("replay failed")
-        return "outputs"
-
-
-class _WindowPreparer:
-    """The ``DecoderInputPreparer`` calls ``_forward_decoder`` makes."""
-
-    def __init__(self, log: list, fail: bool = False) -> None:
-        self.log, self.fail = log, fail
-        self.input_copy_event = "input-copy-event"
-
-    def prepare_inputs(self, *args, **kwargs):
-        self.log.append("prepare")
-        if self.fail:
-            raise RuntimeError("preparation failed")
-        return types.SimpleNamespace(kwargs={}, runtime_draft_len=0, gather_ids=None)
-
-    def record_input_copy(self) -> None:
-        self.log.append("record_input_copy")
-
-    def postprocess_inputs(self, *args, **kwargs) -> None:
-        pass
-
-
-def _window_runner(monkeypatch, graphs, preparer, stager):
-    """A ``DecoderRunner`` with only what ``_forward_decoder`` reads (no speculation, no CED)."""
-    from tensorrt_llm._torch.pyexecutor import ced_replay
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner
-
-    monkeypatch.setattr(runner_module, "get_draft_kv_cache_manager", lambda *args: None)
-    monkeypatch.setattr(ced_replay, "uses_private_decoder_cache", lambda *args: False)
-    attn_metadata = types.SimpleNamespace()
-    runner = DecoderRunner.__new__(DecoderRunner)
-    runner._config = types.SimpleNamespace(spec_config=None, use_mrope=False, is_spec_decode=False)
-    runner._metadata = types.SimpleNamespace(
-        setup_attention=lambda *args: attn_metadata,
-        attn_metadata=attn_metadata,
-        spec_metadata=None,
-    )
-    runner._graphs = graphs
-    runner._inputs = preparer
-    runner._stage_in_graph_sampling = None
-    runner._cuda_graph_metadata_prepare_scope = lambda *args, **kwargs: contextlib.nullcontext()
-    runner._execute_logit_post_processors = lambda *args: None
-    runner.forward_pass_callable = None
-    runner.moe_load_balancer = None
-    runner.model = types.SimpleNamespace()
-    runner.kv_cache_manager_key = "kv_cache_manager"
-    runner.runtime_hooks = RuntimeHooks(step_input_stager=stager)
-    return runner
-
-
-def _decode_step(runner):
-    forward_inputs = types.SimpleNamespace(
-        batch=types.SimpleNamespace(num_context_requests=0),
-        new_tensors_device=None,
-        cache_indirection_buffer=None,
-        gather_context_logits=False,
-        enable_spec_decode=False,
-        runtime_draft_len=0,
-    )
-    resource_manager = types.SimpleNamespace(get_resource_manager=lambda key: object())
-    return runner._forward_decoder(forward_inputs, resource_manager, is_dummy=False)
-
-
-def _names(log: list) -> list:
-    return [entry if isinstance(entry, str) else entry[0] for entry in log]
-
-
-_GRAPH_KEY = (1, 0, False)
-_STAGED_STEP = ["begin", "prepare", "end", "record_input_copy", "completion_event", "replay"]
-
-
-def test_a_replay_issues_the_staged_preparation_before_the_replay(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    assert _decode_step(runner) == "outputs"
-    assert log == [
-        ("begin", True, graphs.cuda_graph_runner.shared_static_tensors),
-        "prepare",
-        ("end", False),
-        "record_input_copy",
-        ("completion_event", "input-copy-event"),
-        ("replay", 3),
-    ]
-
-
-@pytest.mark.parametrize(
-    "key, needs_capture, warmup_only, ran",
-    [
-        (None, False, False, "eager"),
-        (_GRAPH_KEY, True, False, "capture"),
-        (_GRAPH_KEY, False, True, "capture"),
-    ],
-    ids=["eager", "initial_capture", "warmup_only"],
-)
-def test_a_step_that_is_not_a_replay_never_opens_a_window(
-    monkeypatch, key, needs_capture, warmup_only, ran
-):
-    log = []
-    graphs = _WindowGraphs(log, key, needs_capture=needs_capture, warmup_only=warmup_only)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    _decode_step(runner)
-    assert log == ["prepare", "record_input_copy", ran]
-    assert graphs.cuda_graph_runner.staged_static_input_tokens == 0
-
-
-def test_a_declined_replay_prepares_natively(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    stager = _ReferenceStager(log, accept=False)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), stager)
-
-    _decode_step(runner)
-    assert _names(log) == ["begin", "prepare", "record_input_copy", "replay"]
-    assert log[-1] == ("replay", 0)
-
-
-def test_a_failed_preparation_discards_the_window(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    runner = _window_runner(
-        monkeypatch, graphs, _WindowPreparer(log, fail=True), _ReferenceStager(log)
-    )
-
-    with pytest.raises(RuntimeError, match="preparation failed"):
-        _decode_step(runner)
-    assert _names(log) == ["begin", "prepare", "discard"]
-
-
-@pytest.mark.parametrize(
-    "fail, expected",
-    [
-        ("before_replay", ["begin", "prepare", "discard"]),
-        ("replay", _STAGED_STEP),
-    ],
-    ids=["window_open", "window_closed"],
-)
-def test_a_failed_run_discards_only_a_window_still_open(monkeypatch, fail, expected):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY, fail=fail)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    with pytest.raises(RuntimeError):
-        _decode_step(runner)
-    assert _names(log) == expected
-
-
-def test_replay_eager_replay_opens_one_window_per_replay(monkeypatch):
-    log = []
-    graphs = _WindowGraphs(log, _GRAPH_KEY)
-    runner = _window_runner(monkeypatch, graphs, _WindowPreparer(log), _ReferenceStager(log))
-
-    _decode_step(runner)
-    graphs.key = None
-    _decode_step(runner)
-    graphs.key = _GRAPH_KEY
-    _decode_step(runner)
-    assert _names(log) == [
-        *_STAGED_STEP,
-        "prepare",
-        "record_input_copy",
-        "eager",
-        *_STAGED_STEP,
-    ]
-    assert [entry for entry in log if entry[0] == "replay"] == [("replay", 3), ("replay", 3)]
-
-
-def _replay_graphs(monkeypatch, log: list, *, needs_capture: bool, warmup_only: bool):
-    """A ``DecoderCudaGraphs`` with only what ``_run`` reads, logging capture and replay."""
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module
-    from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs
-
-    monkeypatch.setattr(graphs_module, "prepare_attn_metadata_for_draft_replay", lambda *args: None)
-    monkeypatch.setattr(
-        graphs_module, "restore_attn_metadata_after_draft_replay", lambda *args: None
-    )
-    monkeypatch.setattr(
-        graphs_module, "MoeLoadBalancerIterContext", lambda balancer: contextlib.nullcontext()
-    )
-    graphs = DecoderCudaGraphs.__new__(DecoderCudaGraphs)
-    graphs.breakable_cuda_graph_runner = None
-    graphs._torch_compile_bypass = contextlib.nullcontext
-    graphs.cuda_graph_runner = types.SimpleNamespace(
-        needs_capture=lambda key: needs_capture,
-        is_warmup_only=warmup_only,
-        capture=lambda key, fn, inputs, **kwargs: log.append("capture"),
-        replay=lambda key, inputs, **kwargs: log.append("replay"),
-    )
-    return graphs
-
-
-@pytest.mark.parametrize(
-    "key, needs_capture, warmup_only, expected",
-    [
-        (_GRAPH_KEY, False, False, ["before_replay", "replay"]),
-        (_GRAPH_KEY, True, False, ["capture", "replay"]),
-        (_GRAPH_KEY, True, True, ["capture"]),
-        (None, False, False, ["eager"]),
-    ],
-    ids=["replay", "capture_then_replay", "warmup_only", "eager"],
-)
-def test_before_replay_runs_only_before_the_replay_of_a_captured_graph(
-    monkeypatch, key, needs_capture, warmup_only, expected
-):
-    log = []
-    graphs = _replay_graphs(monkeypatch, log, needs_capture=needs_capture, warmup_only=warmup_only)
-
-    graphs._run(
-        key,
-        {"input_ids": torch.zeros(2, dtype=torch.int64)},
-        lambda inputs: log.append("eager"),
-        lambda inputs: None,
-        attn_metadata=None,
-        draft_kv_cache_manager=None,
-        moe_load_balancer=None,
-        enable_spec_decode=True,
-        gather_context_logits=False,
-        position_id_offsets=None,
-        before_replay=lambda: log.append("before_replay"),
-    )
-    assert log == expected
-
-
-# -- the speculative sampler's store writer -----------------------------------------------------
-
-
-class _ReferenceStoreWriter:
-    """Stores like the native ``index_copy_`` path, then copies the store to the host tensors."""
-
-    def __init__(self, supported: bool = True) -> None:
-        self._supported = supported
-        self.calls = []
-
-    def supported(self, store, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens):
-        self.calls.append("supported")
-        return self._supported
-
-    def store_and_copy_out(
-        self,
-        store,
-        slots,
-        new_tokens,
-        next_new_tokens,
-        new_tokens_lens,
-        next_draft_tokens,
-        host_new_tokens,
-        host_new_tokens_lens,
-        host_next_draft_tokens,
-        record_event=True,
-    ) -> None:
-        self.calls.append(("store_and_copy_out", list(slots), record_event))
-        index = torch.tensor(slots, dtype=torch.long)
-        store.new_tokens.squeeze(-1).T.index_copy_(0, index, new_tokens)
-        store.next_new_tokens.squeeze(-1).T.index_copy_(0, index, next_new_tokens)
-        store.new_tokens_lens.index_copy_(0, index, new_tokens_lens)
-        store.next_draft_tokens.index_copy_(0, index, next_draft_tokens)
-        host_new_tokens.copy_(store.new_tokens)
-        host_new_tokens_lens.copy_(store.new_tokens_lens)
-        host_next_draft_tokens.copy_(store.next_draft_tokens)
-
-    def set_completion_event(self, event) -> None:
-        self.calls.append(("completion_event", event))
-
-
-def _cpu_sampling(monkeypatch, *, pinned: bool) -> None:
-    """Run ``SpecSampler.sample_async`` on CPU: device transfers stay on the host."""
-    from tensorrt_llm._torch.speculative import spec_sampler_base
-
-    original_to = torch.Tensor.to
-    original_empty_like = torch.empty_like
-
-    def to_cpu(tensor, *args, **kwargs):
-        if kwargs.get("device") == "cuda":
-            kwargs = {**kwargs, "device": "cpu"}
-        return original_to(tensor, *args, **kwargs)
-
-    def empty_like(tensor, *args, pin_memory=False, **kwargs):
-        # Pinning needs a CUDA context; the writer's host buffers are plain host tensors here.
-        return original_empty_like(tensor, *args, **kwargs)
-
-    monkeypatch.setattr(torch.Tensor, "to", to_cpu)
-    monkeypatch.setattr(torch, "empty_like", empty_like)
-    monkeypatch.setattr(spec_sampler_base, "prefer_pinned", lambda: pinned)
-    monkeypatch.setattr(spec_sampler_base, "handle_stop_criteria", lambda *args, **kwargs: False)
-
-
-def _store_sampler(writer, *, async_worker: bool = False):
-    """A ``SpecSampler`` with a six-slot store, draft length 3 and the given store writer."""
-    from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler
-
-    sampler = SpecSampler.__new__(SpecSampler)
-    sampler.draft_len = 3
-    sampler.max_accepted_path_len = 4
-    sampler.max_seq_len = 128
-    sampler.store = types.SimpleNamespace(
-        new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
-        next_new_tokens=torch.zeros(4, 6, 1, dtype=torch.int32),
-        next_draft_tokens=torch.zeros(6, 3, dtype=torch.int32),
-        new_tokens_lens=torch.zeros(6, dtype=torch.int32),
-    )
-    sampler.store_writer = writer
-    sampler._async_worker_active = lambda: async_worker
-    sampler._copy_to_host = lambda tensor: tensor.clone()
-    sampler._record_sampler_event = lambda: types.SimpleNamespace(cuda_event="sampler-event")
-    return sampler
-
-
-def _sample_generation_step(sampler, *, verify_lens: bool = False):
-    """Three generation requests in nonconsecutive slots 4, 1 and 3."""
-    from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
-
-    def request(request_id: int, slot: int):
-        return types.SimpleNamespace(
-            py_request_id=request_id,
-            py_seq_slot=slot,
-            state=LlmRequestState.GENERATION_IN_PROGRESS,
-            py_draft_tokens=[10, 11, 12],
-            py_draft_tokens_effective_len=None,
-            py_verify_len=None,
-            py_decoding_iter=0,
-        )
-
-    scheduled = types.SimpleNamespace(
-        context_requests_chunking=[],
-        context_requests_last_chunk=[],
-        generation_requests=[request(7, 4), request(8, 1), request(9, 3)],
-    )
-    outputs = dict(
-        new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 100,
-        new_tokens_lens=torch.tensor([1, 3, 2], dtype=torch.int32),
-        next_draft_tokens=torch.arange(9, dtype=torch.int32).reshape(3, 3) + 200,
-        next_new_tokens=torch.arange(12, dtype=torch.int32).reshape(3, 4) + 300,
-    )
-    if verify_lens:
-        outputs["verify_lens"] = torch.tensor([2, 3, 1], dtype=torch.int32)
-        outputs["verify_lens_in_output_order"] = True
-    return sampler.sample_async(scheduled, outputs, [])
-
-
-def _assert_same_stores(sampler, state, native, native_state) -> None:
-    for name in ("new_tokens", "next_new_tokens", "next_draft_tokens", "new_tokens_lens"):
-        assert torch.equal(getattr(sampler.store, name), getattr(native.store, name)), name
-    for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"):
-        assert torch.equal(getattr(state.host, name), getattr(native_state.host, name)), name
-
-
-def test_the_store_writer_matches_the_native_stores_on_nonconsecutive_slots(monkeypatch):
-    _cpu_sampling(monkeypatch, pinned=True)
-    native = _store_sampler(None)
-    writer = _ReferenceStoreWriter()
-    sampler = _store_sampler(writer)
-
-    native_state = _sample_generation_step(native)
-    state = _sample_generation_step(sampler)
-    assert writer.calls == [
-        "supported",
-        ("store_and_copy_out", [4, 1, 3], False),
-        ("completion_event", "sampler-event"),
-    ]
-    assert native.store.new_tokens_lens.tolist() == [0, 3, 0, 2, 1, 0]
-    _assert_same_stores(sampler, state, native, native_state)
-    assert state.host.verify_lens is None
-    assert state.device.new_tokens is sampler.store.next_new_tokens
-
-
-@pytest.mark.parametrize("case", ["declined", "verify_lens", "async_worker", "not_pinned"])
-def test_the_store_writer_falls_back_to_the_native_stores(monkeypatch, case):
-    _cpu_sampling(monkeypatch, pinned=case != "not_pinned")
-    native = _store_sampler(None)
-    writer = _ReferenceStoreWriter(supported=case != "declined")
-    sampler = _store_sampler(writer, async_worker=case == "async_worker")
-
-    native_state = _sample_generation_step(native, verify_lens=case == "verify_lens")
-    state = _sample_generation_step(sampler, verify_lens=case == "verify_lens")
-    assert writer.calls == (["supported"] if case == "declined" else [])
-    _assert_same_stores(sampler, state, native, native_state)
-    if case == "verify_lens":
-        assert state.host.verify_lens.tolist() == [2, 3, 1]

From 6b322651209aac397cd6c172f8c9c21d721c8b43 Mon Sep 17 00:00:00 2001
From: Xiao Wang <24860335+xwang233@users.noreply.github.com>
Date: Tue, 27 Oct 2026 05:41:53 -0700
Subject: [PATCH 19/19] [TRTLLM-17096][fix] runtime hooks: read the model from
 DecoderRunner._model

main renamed the decoder runner's model attribute to _model, so reading
self.model failed at runner construction for every model.

Update two tests to main's code: the GVR prior test calls
sparse_attn_indexer on a SimpleNamespace, which now needs the indexer
seams bound; the forward-step stub needs route_capture,
enable_spec_decode and runtime_draft_len.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
---
 .../_torch/pyexecutor/engine/runners/decoder/runner.py       | 2 +-
 .../_torch/attention/sparse/dsa/test_indexer_gvr_prior.py    | 3 +++
 tests/unittest/_torch/executor/test_runtime_hooks.py         | 5 ++++-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
index 07041d032d56..0f7586cf2c27 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py
@@ -211,7 +211,7 @@ def __init__(
         self._dummy_requests = self._make_dummy_requests()
         self._init_decoder_state()
         self._inputs = self._make_input_preparer(mm_item_scheduler)
-        self.runtime_hooks = runtime_hooks_of(self.model)
+        self.runtime_hooks = runtime_hooks_of(self._model)
         self._inputs.decode_prologue = self.runtime_hooks.decode_prologue
         self._graphs = self._make_cuda_graphs()
         self._initialize_moe_workspace_reclaimer()
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py
index 623a33a5e536..04aa0c8da075 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py
@@ -177,6 +177,9 @@ def prefill(
         aux_stream=None,
     )
     indexer.forward_prepared = MethodType(Indexer.forward_prepared, indexer)
+    indexer._copy_reused_topk = MethodType(Indexer._copy_reused_topk, indexer)
+    indexer._fill_decode_topk = MethodType(Indexer._fill_decode_topk, indexer)
+    indexer._store_mtp_shared_rows = MethodType(Indexer._store_mtp_shared_rows, indexer)
     with (
         patch.object(
             top_k, "update_gvr_prior_from_prefill", wraps=top_k.update_gvr_prior_from_prefill
diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py
index f3a79f0e1180..d2459e6804de 100644
--- a/tests/unittest/_torch/executor/test_runtime_hooks.py
+++ b/tests/unittest/_torch/executor/test_runtime_hooks.py
@@ -593,11 +593,14 @@ def stream(s):
 
     executor = PyExecutor.__new__(PyExecutor)
     executor.iter_counter = 0
+    executor.enable_spec_decode = False
+    executor.runtime_draft_len = 0
     executor.execution_stream = _Stream("execution", log)
     executor.resource_manager = object()
     executor.sampler = types.SimpleNamespace(get_cache_indirection=lambda: None)
     executor.model_engine = types.SimpleNamespace(
-        forward=lambda *args, **kwargs: log.append("forward") or "outputs"
+        forward=lambda *args, **kwargs: log.append("forward") or "outputs",
+        route_capture=None,
     )
     executor._iter_adp_dummy_ctx_tokens = 0
     executor._iter_adp_dummy_gen_tokens = 0